[ { "path": ".git-blame-ignore-revs", "content": "# Exclude these commits from git blame (e.g. mass reformatting).\n# These are ignored by GitHub automatically.\n# To enable this locally, run:\n#\n# git config blame.ignoreRevsFile .git-blame-ignore-revs\n\n3134e5f840c12c8f32613ce520101a047c89dcc2 # refactor(whitespace): rm temporary react fragments (#7161)\ned3f72bc75f3e3a9ae9e4d8cd38278f9c97e78b4 # refactor(whitespace): rm react fragment #7190\n7b927e79c25f4ddfd18a067f489e122acd2c89de # chore(format): format files where `ruff` and `black` agree (#9339)\n" }, { "path": ".github/CODEOWNERS", "content": "* @onyx-dot-app/onyx-core-team\n# Helm charts Owners\n/helm/ @justin-tahara\n\n# Web standards updates\n/web/STANDARDS.md @raunakab @Weves\n\n# Agent context files\n/CLAUDE.md @Weves\n/AGENTS.md @Weves\n\n# Beta cherry-pick workflow owners\n/.github/workflows/post-merge-beta-cherry-pick.yml @justin-tahara @jmelahman\n" }, { "path": ".github/actionlint.yml", "content": "self-hosted-runner:\n # Labels of self-hosted runner in array of strings.\n labels:\n - extras=ecr-cache\n - extras=s3-cache\n - hdd=256\n - runs-on\n - runner=1cpu-linux-arm64\n - runner=1cpu-linux-x64\n - runner=2cpu-linux-arm64\n - runner=2cpu-linux-x64\n - runner=4cpu-linux-arm64\n - runner=4cpu-linux-x64\n - runner=8cpu-linux-arm64\n - runner=8cpu-linux-x64\n - runner=16cpu-linux-arm64\n - runner=16cpu-linux-x64\n - ubuntu-slim # Currently in public preview\n - volume=40gb\n - volume=50gb\n\n# Configuration variables in array of strings defined in your repository or\n# organization. `null` means disabling configuration variables check.\n# Empty array means no configuration variable is allowed.\nconfig-variables: null\n\n# Configuration for file paths. The keys are glob patterns to match to file\n# paths relative to the repository root. The values are the configurations for\n# the file paths. Note that the path separator is always '/'.\n# The following configurations are available.\n#\n# \"ignore\" is an array of regular expression patterns. Matched error messages\n# are ignored. This is similar to the \"-ignore\" command line option.\npaths:\n # Glob pattern relative to the repository root for matching files. The path separator is always '/'.\n # This example configures any YAML file under the '.github/workflows/' directory.\n .github/workflows/**/*.{yml,yaml}:\n # TODO: These are real and should be fixed eventually.\n ignore:\n - 'shellcheck reported issue in this script: SC2038:.+'\n - 'shellcheck reported issue in this script: SC2046:.+'\n - 'shellcheck reported issue in this script: SC2086:.+'\n - 'shellcheck reported issue in this script: SC2193:.+'\n" }, { "path": ".github/actions/build-backend-image/action.yml", "content": "name: \"Build Backend Image\"\ndescription: \"Builds and pushes the backend Docker image with cache reuse\"\ninputs:\n runs-on-ecr-cache:\n description: \"ECR cache registry from runs-on/action\"\n required: true\n ref-name:\n description: \"Git ref name used for cache suffix fallback\"\n required: true\n pr-number:\n description: \"Optional PR number for cache suffix\"\n required: false\n default: \"\"\n github-sha:\n description: \"Commit SHA used for cache keys\"\n required: true\n run-id:\n description: \"GitHub run ID used in output image tag\"\n required: true\n docker-username:\n description: \"Docker Hub username\"\n required: true\n docker-token:\n description: \"Docker Hub token\"\n required: true\n docker-no-cache:\n description: \"Set to 'true' to disable docker build cache\"\n required: false\n default: \"false\"\nruns:\n using: \"composite\"\n steps:\n - name: Format branch name for cache\n id: format-branch\n shell: bash\n env:\n PR_NUMBER: ${{ inputs.pr-number }}\n REF_NAME: ${{ inputs.ref-name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> \"$GITHUB_OUTPUT\"\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ inputs.docker-username }}\n password: ${{ inputs.docker-token }}\n\n - name: Build and push Backend Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n push: true\n tags: ${{ inputs.runs-on-ecr-cache }}:nightly-llm-it-backend-${{ inputs.run-id }}\n cache-from: |\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:backend-cache-${{ inputs.github-sha }}\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:backend-cache\n type=registry,ref=onyxdotapp/onyx-backend:latest\n cache-to: |\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:backend-cache-${{ inputs.github-sha }},mode=max\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:backend-cache,mode=max\n no-cache: ${{ inputs.docker-no-cache == 'true' }}\n" }, { "path": ".github/actions/build-integration-image/action.yml", "content": "name: \"Build Integration Image\"\ndescription: \"Builds and pushes the integration test image with docker bake\"\ninputs:\n runs-on-ecr-cache:\n description: \"ECR cache registry from runs-on/action\"\n required: true\n ref-name:\n description: \"Git ref name used for cache suffix fallback\"\n required: true\n pr-number:\n description: \"Optional PR number for cache suffix\"\n required: false\n default: \"\"\n github-sha:\n description: \"Commit SHA used for cache keys\"\n required: true\n run-id:\n description: \"GitHub run ID used in output image tag\"\n required: true\n docker-username:\n description: \"Docker Hub username\"\n required: true\n docker-token:\n description: \"Docker Hub token\"\n required: true\nruns:\n using: \"composite\"\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ inputs.docker-username }}\n password: ${{ inputs.docker-token }}\n\n - name: Format branch name for cache\n id: format-branch\n shell: bash\n env:\n PR_NUMBER: ${{ inputs.pr-number }}\n REF_NAME: ${{ inputs.ref-name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> \"$GITHUB_OUTPUT\"\n\n - name: Build and push integration test image with Docker Bake\n shell: bash\n env:\n RUNS_ON_ECR_CACHE: ${{ inputs.runs-on-ecr-cache }}\n INTEGRATION_REPOSITORY: ${{ inputs.runs-on-ecr-cache }}\n TAG: nightly-llm-it-${{ inputs.run-id }}\n CACHE_SUFFIX: ${{ steps.format-branch.outputs.cache-suffix }}\n HEAD_SHA: ${{ inputs.github-sha }}\n run: |\n docker buildx bake --push \\\n --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA} \\\n --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX} \\\n --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache \\\n --set backend.cache-from=type=registry,ref=onyxdotapp/onyx-backend:latest \\\n --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA},mode=max \\\n --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX},mode=max \\\n --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache,mode=max \\\n --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA} \\\n --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX} \\\n --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache \\\n --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA},mode=max \\\n --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX},mode=max \\\n --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache,mode=max \\\n integration\n" }, { "path": ".github/actions/build-model-server-image/action.yml", "content": "name: \"Build Model Server Image\"\ndescription: \"Builds and pushes the model server Docker image with cache reuse\"\ninputs:\n runs-on-ecr-cache:\n description: \"ECR cache registry from runs-on/action\"\n required: true\n ref-name:\n description: \"Git ref name used for cache suffix fallback\"\n required: true\n pr-number:\n description: \"Optional PR number for cache suffix\"\n required: false\n default: \"\"\n github-sha:\n description: \"Commit SHA used for cache keys\"\n required: true\n run-id:\n description: \"GitHub run ID used in output image tag\"\n required: true\n docker-username:\n description: \"Docker Hub username\"\n required: true\n docker-token:\n description: \"Docker Hub token\"\n required: true\nruns:\n using: \"composite\"\n steps:\n - name: Format branch name for cache\n id: format-branch\n shell: bash\n env:\n PR_NUMBER: ${{ inputs.pr-number }}\n REF_NAME: ${{ inputs.ref-name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> \"$GITHUB_OUTPUT\"\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ inputs.docker-username }}\n password: ${{ inputs.docker-token }}\n\n - name: Build and push Model Server Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile.model_server\n push: true\n tags: ${{ inputs.runs-on-ecr-cache }}:nightly-llm-it-model-server-${{ inputs.run-id }}\n cache-from: |\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:model-server-cache-${{ inputs.github-sha }}\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:model-server-cache\n type=registry,ref=onyxdotapp/onyx-model-server:latest\n cache-to: |\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:model-server-cache-${{ inputs.github-sha }},mode=max\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ inputs.runs-on-ecr-cache }}:model-server-cache,mode=max\n" }, { "path": ".github/actions/run-nightly-provider-chat-test/action.yml", "content": "name: \"Run Nightly Provider Chat Test\"\ndescription: \"Starts required compose services and runs nightly provider integration test\"\ninputs:\n provider:\n description: \"Provider slug for NIGHTLY_LLM_PROVIDER\"\n required: true\n models:\n description: \"Comma-separated model list for NIGHTLY_LLM_MODELS\"\n required: true\n provider-api-key:\n description: \"API key for NIGHTLY_LLM_API_KEY\"\n required: false\n default: \"\"\n strict:\n description: \"String true/false for NIGHTLY_LLM_STRICT\"\n required: true\n api-base:\n description: \"Optional NIGHTLY_LLM_API_BASE\"\n required: false\n default: \"\"\n api-version:\n description: \"Optional NIGHTLY_LLM_API_VERSION\"\n required: false\n default: \"\"\n deployment-name:\n description: \"Optional NIGHTLY_LLM_DEPLOYMENT_NAME\"\n required: false\n default: \"\"\n custom-config-json:\n description: \"Optional NIGHTLY_LLM_CUSTOM_CONFIG_JSON\"\n required: false\n default: \"\"\n runs-on-ecr-cache:\n description: \"ECR cache registry from runs-on/action\"\n required: true\n run-id:\n description: \"GitHub run ID used in image tags\"\n required: true\n docker-username:\n description: \"Docker Hub username\"\n required: true\n docker-token:\n description: \"Docker Hub token\"\n required: true\nruns:\n using: \"composite\"\n steps:\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ inputs.docker-username }}\n password: ${{ inputs.docker-token }}\n\n - name: Create .env file for Docker Compose\n shell: bash\n env:\n ECR_CACHE: ${{ inputs.runs-on-ecr-cache }}\n RUN_ID: ${{ inputs.run-id }}\n run: |\n cat < deployment/docker_compose/.env\n COMPOSE_PROFILES=s3-filestore\n ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true\n LICENSE_ENFORCEMENT_ENABLED=false\n AUTH_TYPE=basic\n POSTGRES_POOL_PRE_PING=true\n POSTGRES_USE_NULL_POOL=true\n REQUIRE_EMAIL_VERIFICATION=false\n DISABLE_TELEMETRY=true\n INTEGRATION_TESTS_MODE=true\n AUTO_LLM_UPDATE_INTERVAL_SECONDS=10\n AWS_REGION_NAME=us-west-2\n ONYX_BACKEND_IMAGE=${ECR_CACHE}:nightly-llm-it-backend-${RUN_ID}\n ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:nightly-llm-it-model-server-${RUN_ID}\n EOF2\n\n - name: Start Docker containers\n shell: bash\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.dev.yml up -d --wait \\\n relational_db \\\n index \\\n cache \\\n minio \\\n api_server \\\n inference_model_server\n\n - name: Run nightly provider integration test\n uses: nick-fields/retry@ce71cc2ab81d554ebbe88c79ab5975992d79ba08 # ratchet:nick-fields/retry@v3\n env:\n MODELS: ${{ inputs.models }}\n NIGHTLY_LLM_PROVIDER: ${{ inputs.provider }}\n NIGHTLY_LLM_API_KEY: ${{ inputs.provider-api-key }}\n NIGHTLY_LLM_API_BASE: ${{ inputs.api-base }}\n NIGHTLY_LLM_API_VERSION: ${{ inputs.api-version }}\n NIGHTLY_LLM_DEPLOYMENT_NAME: ${{ inputs.deployment-name }}\n NIGHTLY_LLM_CUSTOM_CONFIG_JSON: ${{ inputs.custom-config-json }}\n NIGHTLY_LLM_STRICT: ${{ inputs.strict }}\n RUNS_ON_ECR_CACHE: ${{ inputs.runs-on-ecr-cache }}\n RUN_ID: ${{ inputs.run-id }}\n with:\n timeout_minutes: 20\n max_attempts: 2\n retry_wait_seconds: 10\n command: |\n docker run --rm --network onyx_default \\\n --name test-runner \\\n -e POSTGRES_HOST=relational_db \\\n -e POSTGRES_USER=postgres \\\n -e POSTGRES_PASSWORD=password \\\n -e POSTGRES_DB=postgres \\\n -e DB_READONLY_USER=db_readonly_user \\\n -e DB_READONLY_PASSWORD=password \\\n -e POSTGRES_POOL_PRE_PING=true \\\n -e POSTGRES_USE_NULL_POOL=true \\\n -e VESPA_HOST=index \\\n -e REDIS_HOST=cache \\\n -e API_SERVER_HOST=api_server \\\n -e TEST_WEB_HOSTNAME=test-runner \\\n -e AWS_REGION_NAME=us-west-2 \\\n -e NIGHTLY_LLM_PROVIDER=\"${NIGHTLY_LLM_PROVIDER}\" \\\n -e NIGHTLY_LLM_MODELS=\"${MODELS}\" \\\n -e NIGHTLY_LLM_API_KEY=\"${NIGHTLY_LLM_API_KEY}\" \\\n -e NIGHTLY_LLM_API_BASE=\"${NIGHTLY_LLM_API_BASE}\" \\\n -e NIGHTLY_LLM_API_VERSION=\"${NIGHTLY_LLM_API_VERSION}\" \\\n -e NIGHTLY_LLM_DEPLOYMENT_NAME=\"${NIGHTLY_LLM_DEPLOYMENT_NAME}\" \\\n -e NIGHTLY_LLM_CUSTOM_CONFIG_JSON=\"${NIGHTLY_LLM_CUSTOM_CONFIG_JSON}\" \\\n -e NIGHTLY_LLM_STRICT=\"${NIGHTLY_LLM_STRICT}\" \\\n ${RUNS_ON_ECR_CACHE}:nightly-llm-it-${RUN_ID} \\\n /app/tests/integration/tests/llm_workflows/test_nightly_provider_chat_workflow.py\n" }, { "path": ".github/actions/setup-playwright/action.yml", "content": "name: \"Setup Playwright\"\ndescription: \"Sets up Playwright and system deps (assumes Python and Playwright are installed)\"\nruns:\n using: \"composite\"\n steps:\n - name: Cache playwright cache\n uses: runs-on/cache@50350ad4242587b6c8c2baa2e740b1bc11285ff4 # ratchet:runs-on/cache@v4\n with:\n path: ~/.cache/ms-playwright\n key: ${{ runner.os }}-${{ runner.arch }}-playwright-${{ hashFiles('backend/requirements/default.txt') }}\n restore-keys: |\n ${{ runner.os }}-${{ runner.arch }}-playwright-\n\n - name: Install playwright\n shell: bash\n run: |\n playwright install chromium --with-deps\n" }, { "path": ".github/actions/setup-python-and-install-dependencies/action.yml", "content": "name: \"Setup Python and Install Dependencies\"\ndescription: \"Sets up Python with uv and installs deps\"\ninputs:\n requirements:\n description: \"Newline-separated list of requirement files to install (relative to repo root)\"\n required: true\nruns:\n using: \"composite\"\n steps:\n - name: Compute requirements hash\n id: req-hash\n shell: bash\n env:\n REQUIREMENTS: ${{ inputs.requirements }}\n run: |\n # Hash the contents of the specified requirement files\n hash=\"\"\n while IFS= read -r req; do\n if [ -n \"$req\" ] && [ -f \"$req\" ]; then\n hash=\"$hash$(sha256sum \"$req\")\"\n fi\n done <<< \"$REQUIREMENTS\"\n echo \"hash=$(echo \"$hash\" | sha256sum | cut -d' ' -f1)\" >> \"$GITHUB_OUTPUT\"\n\n # NOTE: This comes before Setup uv since clean-ups run in reverse chronological order\n # such that Setup uv's prune-cache is able to prune the cache before we upload.\n - name: Cache uv cache directory\n uses: runs-on/cache@50350ad4242587b6c8c2baa2e740b1bc11285ff4 # ratchet:runs-on/cache@v4\n with:\n path: ~/.cache/uv\n key: ${{ runner.os }}-uv-${{ steps.req-hash.outputs.hash }}\n restore-keys: |\n ${{ runner.os }}-uv-\n\n - name: Setup uv\n uses: astral-sh/setup-uv@ed21f2f24f8dd64503750218de024bcf64c7250a # ratchet:astral-sh/setup-uv@v7\n with:\n version: \"0.9.9\"\n # TODO: Enable caching once there is a uv.lock file checked in.\n # with:\n # enable-cache: true\n\n - name: Setup Python\n uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # ratchet:actions/setup-python@v5\n with:\n python-version: \"3.11\"\n\n - name: Create virtual environment\n shell: bash\n env:\n VENV_DIR: ${{ runner.temp }}/venv\n run: | # zizmor: ignore[github-env]\n uv venv \"$VENV_DIR\"\n # Validate path before adding to GITHUB_PATH to prevent code injection\n if [ -d \"$VENV_DIR/bin\" ]; then\n realpath \"$VENV_DIR/bin\" >> \"$GITHUB_PATH\"\n else\n echo \"Error: $VENV_DIR/bin does not exist\"\n exit 1\n fi\n\n - name: Install Python dependencies with uv\n shell: bash\n env:\n REQUIREMENTS: ${{ inputs.requirements }}\n run: |\n # Build the uv pip install command with each requirement file as array elements\n cmd=(\"uv\" \"pip\" \"install\")\n while IFS= read -r req; do\n # Skip empty lines\n if [ -n \"$req\" ]; then\n cmd+=(\"-r\" \"$req\")\n fi\n done <<< \"$REQUIREMENTS\"\n echo \"Running: ${cmd[*]}\"\n \"${cmd[@]}\"\n" }, { "path": ".github/actions/slack-notify/action.yml", "content": "name: \"Slack Notify\"\ndescription: \"Sends a Slack notification for workflow events\"\ninputs:\n webhook-url:\n description: \"Slack webhook URL (can also use SLACK_WEBHOOK_URL env var)\"\n required: false\n details:\n description: \"Additional message body content\"\n required: false\n failed-jobs:\n description: \"Deprecated alias for details\"\n required: false\n mention:\n description: \"GitHub username to resolve to a Slack @-mention. Replaces {mention} in details.\"\n required: false\n title:\n description: \"Title for the notification\"\n required: false\n default: \"🚨 Workflow Failed\"\n ref-name:\n description: \"Git ref name (tag/branch)\"\n required: false\nruns:\n using: \"composite\"\n steps:\n - name: Send Slack notification\n shell: bash\n env:\n SLACK_WEBHOOK_URL: ${{ inputs.webhook-url }}\n DETAILS: ${{ inputs.details }}\n FAILED_JOBS: ${{ inputs.failed-jobs }}\n MENTION_USER: ${{ inputs.mention }}\n TITLE: ${{ inputs.title }}\n REF_NAME: ${{ inputs.ref-name }}\n REPO: ${{ github.repository }}\n WORKFLOW: ${{ github.workflow }}\n RUN_NUMBER: ${{ github.run_number }}\n RUN_ID: ${{ github.run_id }}\n SERVER_URL: ${{ github.server_url }}\n GITHUB_REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -z \"$SLACK_WEBHOOK_URL\" ]; then\n echo \"webhook-url input or SLACK_WEBHOOK_URL env var is not set, skipping notification\"\n exit 0\n fi\n\n # Build workflow URL\n WORKFLOW_URL=\"${SERVER_URL}/${REPO}/actions/runs/${RUN_ID}\"\n\n # Use ref_name from input or fall back to github.ref_name\n if [ -z \"$REF_NAME\" ]; then\n REF_NAME=\"$GITHUB_REF_NAME\"\n fi\n\n if [ -z \"$DETAILS\" ]; then\n DETAILS=\"$FAILED_JOBS\"\n fi\n\n # Resolve {mention} placeholder if a GitHub username was provided.\n # Looks up the username in user-mappings.json (co-located with this action)\n # and replaces {mention} with <@SLACK_ID> for a Slack @-mention.\n # Falls back to the plain GitHub username if not found in the mapping.\n if [ -n \"$MENTION_USER\" ]; then\n MAPPINGS_FILE=\"${GITHUB_ACTION_PATH}/user-mappings.json\"\n slack_id=\"$(jq -r --arg gh \"$MENTION_USER\" 'to_entries[] | select(.value | ascii_downcase == ($gh | ascii_downcase)) | .key' \"$MAPPINGS_FILE\" 2>/dev/null | head -1)\"\n\n if [ -n \"$slack_id\" ]; then\n mention_text=\"<@${slack_id}>\"\n else\n mention_text=\"${MENTION_USER}\"\n fi\n\n DETAILS=\"${DETAILS//\\{mention\\}/$mention_text}\"\n TITLE=\"${TITLE//\\{mention\\}/}\"\n else\n DETAILS=\"${DETAILS//\\{mention\\}/}\"\n TITLE=\"${TITLE//\\{mention\\}/}\"\n fi\n\n normalize_multiline() {\n printf '%s' \"$1\" | awk 'BEGIN { ORS=\"\"; first=1 } { if (!first) printf \"\\\\n\"; printf \"%s\", $0; first=0 }'\n }\n\n DETAILS=\"$(normalize_multiline \"$DETAILS\")\"\n REF_NAME=\"$(normalize_multiline \"$REF_NAME\")\"\n TITLE=\"$(normalize_multiline \"$TITLE\")\"\n\n # Escape JSON special characters\n escape_json() {\n local input=\"$1\"\n # Escape backslashes first (but preserve \\n sequences)\n # Protect \\n sequences temporarily\n input=$(printf '%s' \"$input\" | sed 's/\\\\n/\\x01NL\\x01/g')\n # Escape remaining backslashes\n input=$(printf '%s' \"$input\" | sed 's/\\\\/\\\\\\\\/g')\n # Restore \\n sequences (single backslash, will be correct in JSON)\n input=$(printf '%s' \"$input\" | sed 's/\\x01NL\\x01/\\\\n/g')\n # Escape quotes\n printf '%s' \"$input\" | sed 's/\"/\\\\\"/g'\n }\n\n REF_NAME_ESC=$(escape_json \"$REF_NAME\")\n DETAILS_ESC=$(escape_json \"$DETAILS\")\n WORKFLOW_URL_ESC=$(escape_json \"$WORKFLOW_URL\")\n TITLE_ESC=$(escape_json \"$TITLE\")\n\n # Build JSON payload piece by piece\n # Note: DETAILS_ESC already contains \\n sequences that should remain as \\n in JSON\n PAYLOAD=\"{\"\n PAYLOAD=\"${PAYLOAD}\\\"text\\\":\\\"${TITLE_ESC}\\\",\"\n PAYLOAD=\"${PAYLOAD}\\\"blocks\\\":[{\"\n PAYLOAD=\"${PAYLOAD}\\\"type\\\":\\\"header\\\",\"\n PAYLOAD=\"${PAYLOAD}\\\"text\\\":{\\\"type\\\":\\\"plain_text\\\",\\\"text\\\":\\\"${TITLE_ESC}\\\"}\"\n PAYLOAD=\"${PAYLOAD}},{\"\n PAYLOAD=\"${PAYLOAD}\\\"type\\\":\\\"section\\\",\"\n PAYLOAD=\"${PAYLOAD}\\\"fields\\\":[\"\n if [ -n \"$REF_NAME\" ]; then\n PAYLOAD=\"${PAYLOAD}{\\\"type\\\":\\\"mrkdwn\\\",\\\"text\\\":\\\"*Ref:*\\\\n${REF_NAME_ESC}\\\"},\"\n fi\n PAYLOAD=\"${PAYLOAD}{\\\"type\\\":\\\"mrkdwn\\\",\\\"text\\\":\\\"*Run ID:*\\\\n#${RUN_NUMBER}\\\"}\"\n PAYLOAD=\"${PAYLOAD}]\"\n PAYLOAD=\"${PAYLOAD}}\"\n if [ -n \"$DETAILS\" ]; then\n PAYLOAD=\"${PAYLOAD},{\"\n PAYLOAD=\"${PAYLOAD}\\\"type\\\":\\\"section\\\",\"\n PAYLOAD=\"${PAYLOAD}\\\"text\\\":{\\\"type\\\":\\\"mrkdwn\\\",\\\"text\\\":\\\"${DETAILS_ESC}\\\"}\"\n PAYLOAD=\"${PAYLOAD}}\"\n fi\n PAYLOAD=\"${PAYLOAD},{\"\n PAYLOAD=\"${PAYLOAD}\\\"type\\\":\\\"actions\\\",\"\n PAYLOAD=\"${PAYLOAD}\\\"elements\\\":[{\"\n PAYLOAD=\"${PAYLOAD}\\\"type\\\":\\\"button\\\",\"\n PAYLOAD=\"${PAYLOAD}\\\"text\\\":{\\\"type\\\":\\\"plain_text\\\",\\\"text\\\":\\\"View Workflow Run\\\"},\"\n PAYLOAD=\"${PAYLOAD}\\\"url\\\":\\\"${WORKFLOW_URL_ESC}\\\"\"\n PAYLOAD=\"${PAYLOAD}}]\"\n PAYLOAD=\"${PAYLOAD}}\"\n PAYLOAD=\"${PAYLOAD}]\"\n PAYLOAD=\"${PAYLOAD}}\"\n\n curl -X POST -H 'Content-type: application/json' \\\n --data \"$PAYLOAD\" \\\n \"$SLACK_WEBHOOK_URL\"\n" }, { "path": ".github/actions/slack-notify/user-mappings.json", "content": "{\n \"U05SAGZPEA1\": \"yuhongsun96\",\n \"U05SAH6UGUD\": \"Weves\",\n \"U07PWEQB7A5\": \"evan-onyx\",\n \"U07V1SM68KF\": \"joachim-danswer\",\n \"U08JZ9N3QNN\": \"raunakab\",\n \"U08L24NCLJE\": \"Subash-Mohan\",\n \"U090B9M07B2\": \"wenxi-onyx\",\n \"U094RASDP0Q\": \"duo-onyx\",\n \"U096L8ZQ85B\": \"justin-tahara\",\n \"U09AHV8UBQX\": \"jessicasingh7\",\n \"U09KAL5T3C2\": \"nmgarza5\",\n \"U09KPGVQ70R\": \"acaprau\",\n \"U09QR8KTSJH\": \"rohoswagger\",\n \"U09RB4NTXA4\": \"jmelahman\",\n \"U0A6K9VCY6A\": \"Danelegend\",\n \"U0AGC4KH71A\": \"Bo-Onyx\"\n}\n" }, { "path": ".github/dependabot.yml", "content": "version: 2\nupdates:\n - package-ecosystem: \"github-actions\"\n directory: \"/\"\n schedule:\n interval: \"weekly\"\n cooldown:\n default-days: 7\n open-pull-requests-limit: 3\n assignees:\n - \"jmelahman\"\n labels:\n - \"dependabot:actions\"\n - package-ecosystem: \"pip\"\n directory: \"/backend\"\n schedule:\n interval: \"weekly\"\n cooldown:\n default-days: 7\n open-pull-requests-limit: 3\n assignees:\n - \"jmelahman\"\n labels:\n - \"dependabot:python\"\n" }, { "path": ".github/pull_request_template.md", "content": "## Description\n\n\n\n## How Has This Been Tested?\n\n\n\n## Additional Options\n\n- [ ] [Optional] Please cherry-pick this PR to the latest release version.\n- [ ] [Optional] Override Linear Check\n" }, { "path": ".github/runs-on.yml", "content": "_extend: .github-private\n" }, { "path": ".github/workflows/deployment.yml", "content": "name: Build and Push Docker Images on Tag\n\non:\n push:\n tags:\n - \"*\"\n workflow_dispatch:\n\n# Set restrictive default permissions for all jobs. Jobs that need more permissions\n# should explicitly declare them.\npermissions:\n # Required for OIDC authentication with AWS\n id-token: write # zizmor: ignore[excessive-permissions]\n\nenv:\n EDGE_TAG: ${{ startsWith(github.ref_name, 'nightly-latest') }}\n\njobs:\n # Determine which components to build based on the tag\n determine-builds:\n # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.\n runs-on: ubuntu-slim\n timeout-minutes: 90\n outputs:\n build-desktop: ${{ steps.check.outputs.build-desktop }}\n build-web: ${{ steps.check.outputs.build-web }}\n build-web-cloud: ${{ steps.check.outputs.build-web-cloud }}\n build-backend: ${{ steps.check.outputs.build-backend }}\n build-backend-craft: ${{ steps.check.outputs.build-backend-craft }}\n build-model-server: ${{ steps.check.outputs.build-model-server }}\n is-cloud-tag: ${{ steps.check.outputs.is-cloud-tag }}\n is-beta: ${{ steps.check.outputs.is-beta }}\n is-beta-standalone: ${{ steps.check.outputs.is-beta-standalone }}\n is-latest: ${{ steps.check.outputs.is-latest }}\n is-test-run: ${{ steps.check.outputs.is-test-run }}\n sanitized-tag: ${{ steps.check.outputs.sanitized-tag }}\n short-sha: ${{ steps.check.outputs.short-sha }}\n steps:\n - name: Checkout (for git tags)\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n fetch-depth: 0\n fetch-tags: true\n\n - name: Setup uv\n uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n version: \"0.9.9\"\n enable-cache: false\n\n - name: Check which components to build and version info\n id: check\n env:\n EVENT_NAME: ${{ github.event_name }}\n run: |\n set -eo pipefail\n TAG=\"${GITHUB_REF_NAME}\"\n # Sanitize tag name by replacing slashes with hyphens (for Docker tag compatibility)\n SANITIZED_TAG=$(echo \"$TAG\" | tr '/' '-')\n SHORT_SHA=\"${GITHUB_SHA::7}\"\n\n # Initialize all flags to false\n IS_CLOUD=false\n IS_NIGHTLY=false\n IS_VERSION_TAG=false\n IS_STABLE=false\n IS_BETA=false\n IS_BETA_STANDALONE=false\n IS_LATEST=false\n IS_PROD_TAG=false\n IS_TEST_RUN=false\n BUILD_DESKTOP=false\n BUILD_WEB=false\n BUILD_WEB_CLOUD=false\n BUILD_BACKEND=true\n BUILD_BACKEND_CRAFT=false\n BUILD_MODEL_SERVER=true\n\n # Determine tag type based on pattern matching (do regex checks once)\n if [[ \"$TAG\" == *cloud* ]]; then\n IS_CLOUD=true\n fi\n if [[ \"$TAG\" == nightly* ]]; then\n IS_NIGHTLY=true\n fi\n if [[ \"$TAG\" =~ ^v[0-9]+\\.[0-9]+\\.[0-9]+ ]]; then\n IS_VERSION_TAG=true\n fi\n if [[ \"$TAG\" =~ ^v[0-9]+\\.[0-9]+\\.[0-9]+$ ]]; then\n IS_STABLE=true\n fi\n if [[ \"$TAG\" =~ ^v[0-9]+\\.[0-9]+\\.[0-9]+-beta(\\.[0-9]+)?$ ]]; then\n IS_BETA=true\n fi\n\n # Determine what to build based on tag type\n if [[ \"$IS_CLOUD\" == \"true\" ]]; then\n BUILD_WEB_CLOUD=true\n else\n BUILD_WEB=true\n # Only build desktop for semver tags (excluding beta)\n if [[ \"$IS_VERSION_TAG\" == \"true\" ]] && [[ \"$IS_BETA\" != \"true\" ]]; then\n BUILD_DESKTOP=true\n fi\n fi\n\n # Standalone version checks (for backend/model-server - version excluding cloud tags)\n if [[ \"$IS_BETA\" == \"true\" ]] && [[ \"$IS_CLOUD\" != \"true\" ]]; then\n IS_BETA_STANDALONE=true\n fi\n\n # Determine if this tag should get the \"latest\" Docker tag.\n # Only the highest semver stable tag (vX.Y.Z exactly) gets \"latest\".\n if [[ \"$IS_STABLE\" == \"true\" ]]; then\n HIGHEST_STABLE=$(uv run --no-sync --with onyx-devtools ods latest-stable-tag) || {\n echo \"::error::Failed to determine highest stable tag via 'ods latest-stable-tag'\"\n exit 1\n }\n if [[ \"$TAG\" == \"$HIGHEST_STABLE\" ]]; then\n IS_LATEST=true\n fi\n fi\n\n # Build craft-latest backend alongside the regular latest.\n if [[ \"$IS_LATEST\" == \"true\" ]]; then\n BUILD_BACKEND_CRAFT=true\n fi\n\n # Determine if this is a production tag\n # Production tags are: version tags (v1.2.3*) or nightly tags\n if [[ \"$IS_VERSION_TAG\" == \"true\" ]] || [[ \"$IS_NIGHTLY\" == \"true\" ]]; then\n IS_PROD_TAG=true\n fi\n\n # Determine if this is a test run (workflow_dispatch on non-production ref)\n if [[ \"$EVENT_NAME\" == \"workflow_dispatch\" ]] && [[ \"$IS_PROD_TAG\" != \"true\" ]]; then\n IS_TEST_RUN=true\n fi\n {\n echo \"build-desktop=$BUILD_DESKTOP\"\n echo \"build-web=$BUILD_WEB\"\n echo \"build-web-cloud=$BUILD_WEB_CLOUD\"\n echo \"build-backend=$BUILD_BACKEND\"\n echo \"build-backend-craft=$BUILD_BACKEND_CRAFT\"\n echo \"build-model-server=$BUILD_MODEL_SERVER\"\n echo \"is-cloud-tag=$IS_CLOUD\"\n echo \"is-beta=$IS_BETA\"\n echo \"is-beta-standalone=$IS_BETA_STANDALONE\"\n echo \"is-latest=$IS_LATEST\"\n echo \"is-test-run=$IS_TEST_RUN\"\n echo \"sanitized-tag=$SANITIZED_TAG\"\n echo \"short-sha=$SHORT_SHA\"\n } >> \"$GITHUB_OUTPUT\"\n\n check-version-tag:\n runs-on: ubuntu-slim\n timeout-minutes: 10\n if: ${{ !startsWith(github.ref_name, 'nightly-latest') && github.event_name != 'workflow_dispatch' }}\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n fetch-depth: 0\n\n - name: Setup uv\n uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n version: \"0.9.9\"\n # NOTE: This isn't caching much and zizmor suggests this could be poisoned, so disable.\n enable-cache: false\n\n - name: Validate tag is versioned correctly\n run: |\n uv run --no-sync --with release-tag tag --check\n\n notify-slack-on-tag-check-failure:\n needs:\n - check-version-tag\n if: always() && needs.check-version-tag.result == 'failure' && github.event_name != 'workflow_dispatch'\n runs-on: ubuntu-slim\n timeout-minutes: 10\n environment: release\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Send Slack notification\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.MONITOR_DEPLOYMENTS_WEBHOOK }}\n failed-jobs: \"• check-version-tag\"\n title: \"🚨 Version Tag Check Failed\"\n ref-name: ${{ github.ref_name }}\n\n # Create GitHub release first, before desktop builds start.\n # This ensures all desktop matrix jobs upload to the same release instead of\n # racing to create duplicate releases.\n create-release:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-desktop == 'true'\n runs-on: ubuntu-slim\n timeout-minutes: 10\n permissions:\n contents: write\n outputs:\n release-id: ${{ steps.create-release.outputs.id }}\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Determine release tag\n id: release-tag\n env:\n IS_TEST_RUN: ${{ needs.determine-builds.outputs.is-test-run }}\n SHORT_SHA: ${{ needs.determine-builds.outputs.short-sha }}\n run: |\n if [ \"${IS_TEST_RUN}\" == \"true\" ]; then\n echo \"tag=v0.0.0-dev+${SHORT_SHA}\" >> \"$GITHUB_OUTPUT\"\n else\n echo \"tag=${GITHUB_REF_NAME}\" >> \"$GITHUB_OUTPUT\"\n fi\n\n - name: Create GitHub Release\n id: create-release\n uses: softprops/action-gh-release@da05d552573ad5aba039eaac05058a918a7bf631 # ratchet:softprops/action-gh-release@v2\n with:\n tag_name: ${{ steps.release-tag.outputs.tag }}\n name: ${{ steps.release-tag.outputs.tag }}\n body: \"See the assets to download this version and install.\"\n draft: true\n prerelease: false\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\n build-desktop:\n needs:\n - determine-builds\n - create-release\n if: needs.determine-builds.outputs.build-desktop == 'true'\n permissions:\n id-token: write\n contents: write\n actions: read\n strategy:\n fail-fast: false\n matrix:\n include:\n - platform: \"macos-latest\" # Build a universal image for macOS.\n args: \"--target universal-apple-darwin\"\n - platform: \"ubuntu-24.04\"\n args: \"--bundles deb,rpm\"\n - platform: \"ubuntu-24.04-arm\" # Only available in public repos.\n args: \"--bundles deb,rpm\"\n - platform: \"windows-latest\"\n args: \"\"\n\n runs-on: ${{ matrix.platform }}\n timeout-minutes: 90\n environment: release\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6.0.2\n with:\n # NOTE: persist-credentials is needed for tauri-action to upload assets to GitHub releases.\n persist-credentials: true # zizmor: ignore[artipacked]\n\n - name: Configure AWS credentials\n if: startsWith(matrix.platform, 'macos-')\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n if: startsWith(matrix.platform, 'macos-')\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n APPLE_ID, deploy/apple-id\n APPLE_PASSWORD, deploy/apple-password\n APPLE_CERTIFICATE, deploy/apple-certificate\n APPLE_CERTIFICATE_PASSWORD, deploy/apple-certificate-password\n KEYCHAIN_PASSWORD, deploy/keychain-password\n APPLE_TEAM_ID, deploy/apple-team-id\n parse-json-secrets: true\n\n - name: install dependencies (ubuntu only)\n if: startsWith(matrix.platform, 'ubuntu-')\n run: |\n sudo apt-get update\n sudo apt-get install -y \\\n build-essential \\\n libglib2.0-dev \\\n libgirepository1.0-dev \\\n libgtk-3-dev \\\n libjavascriptcoregtk-4.1-dev \\\n libwebkit2gtk-4.1-dev \\\n libayatana-appindicator3-dev \\\n gobject-introspection \\\n pkg-config \\\n curl \\\n xdg-utils\n\n - name: setup node\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v6.3.0\n with:\n node-version: 24\n package-manager-cache: false\n\n - name: install Rust stable\n uses: dtolnay/rust-toolchain@6d9817901c499d6b02debbb57edb38d33daa680b # zizmor: ignore[impostor-commit]\n with:\n # Those targets are only used on macos runners so it's in an `if` to slightly speed up windows and linux builds.\n targets: ${{ matrix.platform == 'macos-latest' && 'aarch64-apple-darwin,x86_64-apple-darwin' || '' }}\n\n - name: install frontend dependencies\n working-directory: ./desktop\n run: npm install\n\n - name: Inject version (Unix)\n if: runner.os != 'Windows'\n working-directory: ./desktop\n env:\n SHORT_SHA: ${{ needs.determine-builds.outputs.short-sha }}\n IS_TEST_RUN: ${{ needs.determine-builds.outputs.is-test-run }}\n run: |\n if [ \"${IS_TEST_RUN}\" == \"true\" ]; then\n VERSION=\"0.0.0-dev+${SHORT_SHA}\"\n else\n VERSION=\"${GITHUB_REF_NAME#v}\"\n fi\n echo \"Injecting version: $VERSION\"\n\n # Update Cargo.toml\n sed \"s/^version = .*/version = \\\"$VERSION\\\"/\" src-tauri/Cargo.toml > src-tauri/Cargo.toml.tmp\n mv src-tauri/Cargo.toml.tmp src-tauri/Cargo.toml\n\n # Update tauri.conf.json\n jq --arg v \"$VERSION\" '.version = $v' src-tauri/tauri.conf.json > src-tauri/tauri.conf.json.tmp\n mv src-tauri/tauri.conf.json.tmp src-tauri/tauri.conf.json\n\n # Update package.json\n jq --arg v \"$VERSION\" '.version = $v' package.json > package.json.tmp\n mv package.json.tmp package.json\n\n echo \"Versions set to: $VERSION\"\n\n - name: Inject version (Windows)\n if: runner.os == 'Windows'\n working-directory: ./desktop\n shell: pwsh\n env:\n IS_TEST_RUN: ${{ needs.determine-builds.outputs.is-test-run }}\n run: |\n # Windows MSI requires numeric-only build metadata, so we skip the SHA suffix\n if ($env:IS_TEST_RUN -eq \"true\") {\n $VERSION = \"0.0.0\"\n } else {\n # Strip 'v' prefix and any pre-release suffix (e.g., -beta.13) for MSI compatibility\n $VERSION = \"$env:GITHUB_REF_NAME\" -replace '^v', '' -replace '-.*$', ''\n }\n Write-Host \"Injecting version: $VERSION\"\n\n # Update Cargo.toml\n $cargo = Get-Content src-tauri/Cargo.toml -Raw\n $cargo = $cargo -replace '(?m)^version = .*', \"version = `\"$VERSION`\"\"\n Set-Content src-tauri/Cargo.toml $cargo -NoNewline\n\n # Update tauri.conf.json\n $json = Get-Content src-tauri/tauri.conf.json | ConvertFrom-Json\n $json.version = $VERSION\n $json | ConvertTo-Json -Depth 100 | Set-Content src-tauri/tauri.conf.json\n\n # Update package.json\n $pkg = Get-Content package.json | ConvertFrom-Json\n $pkg.version = $VERSION\n $pkg | ConvertTo-Json -Depth 100 | Set-Content package.json\n\n Write-Host \"Versions set to: $VERSION\"\n\n - name: Import Apple Developer Certificate\n if: startsWith(matrix.platform, 'macos-')\n run: |\n echo $APPLE_CERTIFICATE | base64 --decode > certificate.p12\n security create-keychain -p \"$KEYCHAIN_PASSWORD\" build.keychain\n security default-keychain -s build.keychain\n security unlock-keychain -p \"$KEYCHAIN_PASSWORD\" build.keychain\n security set-keychain-settings -t 3600 -u build.keychain\n security import certificate.p12 -k build.keychain -P \"$APPLE_CERTIFICATE_PASSWORD\" -T /usr/bin/codesign\n security set-key-partition-list -S apple-tool:,apple:,codesign: -s -k \"$KEYCHAIN_PASSWORD\" build.keychain\n security find-identity -v -p codesigning build.keychain\n\n - name: Verify Certificate\n if: startsWith(matrix.platform, 'macos-')\n run: |\n CERT_INFO=$(security find-identity -v -p codesigning build.keychain | grep -E \"(Developer ID Application|Apple Distribution|Apple Development)\" | head -n 1)\n CERT_ID=$(echo \"$CERT_INFO\" | awk -F'\"' '{print $2}')\n echo \"CERT_ID=$CERT_ID\" >> $GITHUB_ENV\n echo \"Certificate imported.\"\n\n - uses: tauri-apps/tauri-action@73fb865345c54760d875b94642314f8c0c894afa # ratchet:tauri-apps/tauri-action@action-v0.6.1\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n APPLE_ID: ${{ env.APPLE_ID }}\n APPLE_PASSWORD: ${{ env.APPLE_PASSWORD }}\n APPLE_SIGNING_IDENTITY: ${{ env.CERT_ID }}\n APPLE_TEAM_ID: ${{ env.APPLE_TEAM_ID }}\n with:\n # Use the release created by the create-release job to avoid race conditions\n # when multiple matrix jobs try to create/update the same release simultaneously\n releaseId: ${{ needs.create-release.outputs.release-id }}\n assetNamePattern: \"[name]_[arch][ext]\"\n args: ${{ matrix.args }}\n\n build-web-amd64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-web == 'true'\n runs-on:\n - runs-on\n - runner=4cpu-linux-x64\n - run-id=${{ github.run_id }}-web-amd64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-web-server\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./web\n file: ./web/Dockerfile\n platforms: linux/amd64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n NODE_OPTIONS=--max-old-space-size=8192\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-amd64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:edge\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-amd64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-web-arm64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-web == 'true'\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - run-id=${{ github.run_id }}-web-arm64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-web-server\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./web\n file: ./web/Dockerfile\n platforms: linux/arm64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n NODE_OPTIONS=--max-old-space-size=8192\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-arm64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:edge\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-arm64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n merge-web:\n needs:\n - determine-builds\n - build-web-amd64\n - build-web-arm64\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-merge-web\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-web-server\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n tags: |\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run == 'true' && format('web-{0}', needs.determine-builds.outputs.sanitized-tag) || github.ref_name }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-latest == 'true' && 'latest' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-latest == 'true' && 'craft-latest' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && env.EDGE_TAG == 'true' && 'edge' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-beta == 'true' && 'beta' || '' }}\n\n - name: Create and push manifest\n env:\n IMAGE_REPO: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n AMD64_DIGEST: ${{ needs.build-web-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.build-web-arm64.outputs.digest }}\n META_TAGS: ${{ steps.meta.outputs.tags }}\n run: |\n IMAGES=\"${IMAGE_REPO}@${AMD64_DIGEST} ${IMAGE_REPO}@${ARM64_DIGEST}\"\n docker buildx imagetools create \\\n $(printf '%s\\n' \"${META_TAGS}\" | xargs -I {} echo -t {}) \\\n $IMAGES\n\n build-web-cloud-amd64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-web-cloud == 'true'\n runs-on:\n - runs-on\n - runner=4cpu-linux-x64\n - run-id=${{ github.run_id }}-web-cloud-amd64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-web-server-cloud\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./web\n file: ./web/Dockerfile\n platforms: linux/amd64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n NEXT_PUBLIC_CLOUD_ENABLED=true\n NEXT_PUBLIC_POSTHOG_KEY=${{ secrets.POSTHOG_KEY }}\n NEXT_PUBLIC_POSTHOG_HOST=${{ secrets.POSTHOG_HOST }}\n NEXT_PUBLIC_SENTRY_DSN=${{ secrets.SENTRY_DSN }}\n NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY=${{ secrets.STRIPE_PUBLISHABLE_KEY }}\n NEXT_PUBLIC_RECAPTCHA_SITE_KEY=${{ vars.NEXT_PUBLIC_RECAPTCHA_SITE_KEY }}\n NEXT_PUBLIC_GTM_ENABLED=true\n NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED=true\n NEXT_PUBLIC_INCLUDE_ERROR_POPUP_SUPPORT_LINK=true\n NODE_OPTIONS=--max-old-space-size=8192\n SENTRY_RELEASE=${{ github.sha }}\n secrets: |\n sentry_auth_token=${{ secrets.SENTRY_AUTH_TOKEN }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:cloudweb-cache-amd64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:cloudweb-cache-amd64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-web-cloud-arm64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-web-cloud == 'true'\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - run-id=${{ github.run_id }}-web-cloud-arm64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-web-server-cloud\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./web\n file: ./web/Dockerfile\n platforms: linux/arm64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n NEXT_PUBLIC_CLOUD_ENABLED=true\n NEXT_PUBLIC_POSTHOG_KEY=${{ secrets.POSTHOG_KEY }}\n NEXT_PUBLIC_POSTHOG_HOST=${{ secrets.POSTHOG_HOST }}\n NEXT_PUBLIC_SENTRY_DSN=${{ secrets.SENTRY_DSN }}\n NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY=${{ secrets.STRIPE_PUBLISHABLE_KEY }}\n NEXT_PUBLIC_RECAPTCHA_SITE_KEY=${{ vars.NEXT_PUBLIC_RECAPTCHA_SITE_KEY }}\n NEXT_PUBLIC_GTM_ENABLED=true\n NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED=true\n NEXT_PUBLIC_INCLUDE_ERROR_POPUP_SUPPORT_LINK=true\n NODE_OPTIONS=--max-old-space-size=8192\n SENTRY_RELEASE=${{ github.sha }}\n secrets: |\n sentry_auth_token=${{ secrets.SENTRY_AUTH_TOKEN }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:cloudweb-cache-arm64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:cloudweb-cache-arm64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n merge-web-cloud:\n needs:\n - determine-builds\n - build-web-cloud-amd64\n - build-web-cloud-arm64\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-merge-web-cloud\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-web-server-cloud\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n tags: |\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run == 'true' && format('web-cloud-{0}', needs.determine-builds.outputs.sanitized-tag) || github.ref_name }}\n\n - name: Create and push manifest\n env:\n IMAGE_REPO: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n AMD64_DIGEST: ${{ needs.build-web-cloud-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.build-web-cloud-arm64.outputs.digest }}\n META_TAGS: ${{ steps.meta.outputs.tags }}\n run: |\n IMAGES=\"${IMAGE_REPO}@${AMD64_DIGEST} ${IMAGE_REPO}@${ARM64_DIGEST}\"\n docker buildx imagetools create \\\n $(printf '%s\\n' \"${META_TAGS}\" | xargs -I {} echo -t {}) \\\n $IMAGES\n\n build-backend-amd64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-backend == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-backend-amd64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-backend-cloud' || 'onyxdotapp/onyx-backend' }}\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n platforms: linux/amd64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-amd64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:edge\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-amd64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-backend-arm64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-backend == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - run-id=${{ github.run_id }}-backend-arm64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-backend-cloud' || 'onyxdotapp/onyx-backend' }}\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n platforms: linux/arm64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-arm64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:edge\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-arm64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n merge-backend:\n needs:\n - determine-builds\n - build-backend-amd64\n - build-backend-arm64\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-merge-backend\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n env:\n REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-backend-cloud' || 'onyxdotapp/onyx-backend' }}\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n tags: |\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run == 'true' && format('backend-{0}', needs.determine-builds.outputs.sanitized-tag) || github.ref_name }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-latest == 'true' && 'latest' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && env.EDGE_TAG == 'true' && 'edge' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-beta-standalone == 'true' && 'beta' || '' }}\n\n - name: Create and push manifest\n env:\n IMAGE_REPO: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n AMD64_DIGEST: ${{ needs.build-backend-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.build-backend-arm64.outputs.digest }}\n META_TAGS: ${{ steps.meta.outputs.tags }}\n run: |\n IMAGES=\"${IMAGE_REPO}@${AMD64_DIGEST} ${IMAGE_REPO}@${ARM64_DIGEST}\"\n docker buildx imagetools create \\\n $(printf '%s\\n' \"${META_TAGS}\" | xargs -I {} echo -t {}) \\\n $IMAGES\n\n build-backend-craft-amd64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-backend-craft == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-backend-craft-amd64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-backend\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n platforms: linux/amd64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n ENABLE_CRAFT=true\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-craft-cache-amd64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-craft-cache-amd64,mode=max\n outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-backend-craft-arm64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-backend-craft == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - run-id=${{ github.run_id }}-backend-craft-arm64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-backend\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n platforms: linux/arm64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n ENABLE_CRAFT=true\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-craft-cache-arm64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-craft-cache-arm64,mode=max\n outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n merge-backend-craft:\n needs:\n - determine-builds\n - build-backend-craft-amd64\n - build-backend-craft-arm64\n if: needs.determine-builds.outputs.build-backend-craft == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-merge-backend-craft\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-backend\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n tags: |\n type=raw,value=craft-latest\n\n - name: Create and push manifest\n env:\n IMAGE_REPO: ${{ env.REGISTRY_IMAGE }}\n AMD64_DIGEST: ${{ needs.build-backend-craft-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.build-backend-craft-arm64.outputs.digest }}\n META_TAGS: ${{ steps.meta.outputs.tags }}\n run: |\n IMAGES=\"${IMAGE_REPO}@${AMD64_DIGEST} ${IMAGE_REPO}@${ARM64_DIGEST}\"\n docker buildx imagetools create \\\n $(printf '%s\\n' \"${META_TAGS}\" | xargs -I {} echo -t {}) \\\n $IMAGES\n\n build-model-server-amd64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-model-server == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-model-server-amd64\n - volume=40gb\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-model-server-cloud' || 'onyxdotapp/onyx-model-server' }}\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n with:\n buildkitd-flags: ${{ vars.DOCKER_DEBUG == 'true' && '--debug' || '' }}\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n env:\n DEBUG: ${{ vars.DOCKER_DEBUG == 'true' && 1 || 0 }}\n with:\n context: ./backend\n file: ./backend/Dockerfile.model_server\n platforms: linux/amd64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-amd64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:edge\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-amd64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ env.EDGE_TAG != 'true' && vars.MODEL_SERVER_NO_CACHE == 'true' }}\n provenance: false\n sbom: false\n\n build-model-server-arm64:\n needs: determine-builds\n if: needs.determine-builds.outputs.build-model-server == 'true'\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - run-id=${{ github.run_id }}-model-server-arm64\n - volume=40gb\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-model-server-cloud' || 'onyxdotapp/onyx-model-server' }}\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n with:\n buildkitd-flags: ${{ vars.DOCKER_DEBUG == 'true' && '--debug' || '' }}\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n env:\n DEBUG: ${{ vars.DOCKER_DEBUG == 'true' && 1 || 0 }}\n with:\n context: ./backend\n file: ./backend/Dockerfile.model_server\n platforms: linux/arm64\n labels: ${{ steps.meta.outputs.labels }}\n build-args: |\n ONYX_VERSION=${{ github.ref_name }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-arm64\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:edge\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-arm64,mode=max\n outputs: type=image,name=${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n no-cache: ${{ env.EDGE_TAG != 'true' && vars.MODEL_SERVER_NO_CACHE == 'true' }}\n provenance: false\n sbom: false\n\n merge-model-server:\n needs:\n - determine-builds\n - build-model-server-amd64\n - build-model-server-arm64\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-merge-model-server\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n env:\n REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-model-server-cloud' || 'onyxdotapp/onyx-model-server' }}\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n tags: |\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run == 'true' && format('model-server-{0}', needs.determine-builds.outputs.sanitized-tag) || github.ref_name }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-latest == 'true' && 'latest' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-latest == 'true' && 'craft-latest' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && env.EDGE_TAG == 'true' && 'edge' || '' }}\n type=raw,value=${{ needs.determine-builds.outputs.is-test-run != 'true' && needs.determine-builds.outputs.is-beta-standalone == 'true' && 'beta' || '' }}\n\n - name: Create and push manifest\n env:\n IMAGE_REPO: ${{ needs.determine-builds.outputs.is-test-run == 'true' && env.RUNS_ON_ECR_CACHE || env.REGISTRY_IMAGE }}\n AMD64_DIGEST: ${{ needs.build-model-server-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.build-model-server-arm64.outputs.digest }}\n META_TAGS: ${{ steps.meta.outputs.tags }}\n run: |\n IMAGES=\"${IMAGE_REPO}@${AMD64_DIGEST} ${IMAGE_REPO}@${ARM64_DIGEST}\"\n docker buildx imagetools create \\\n $(printf '%s\\n' \"${META_TAGS}\" | xargs -I {} echo -t {}) \\\n $IMAGES\n\n trivy-scan:\n needs:\n - determine-builds\n - merge-web\n - merge-web-cloud\n - merge-backend\n - merge-model-server\n if: >-\n always() && !cancelled() &&\n (needs.merge-web.result == 'success' ||\n needs.merge-web-cloud.result == 'success' ||\n needs.merge-backend.result == 'success' ||\n needs.merge-model-server.result == 'success')\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - run-id=${{ github.run_id }}-trivy-scan-${{ matrix.component }}\n - extras=ecr-cache\n permissions:\n security-events: write # needed for SARIF uploads\n timeout-minutes: 10\n strategy:\n fail-fast: false\n matrix:\n include:\n - component: web\n registry-image: onyxdotapp/onyx-web-server\n - component: web-cloud\n registry-image: onyxdotapp/onyx-web-server-cloud\n - component: backend\n registry-image: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-backend-cloud' || 'onyxdotapp/onyx-backend' }}\n trivyignore: backend/.trivyignore\n - component: model-server\n registry-image: ${{ contains(github.ref_name, 'cloud') && 'onyxdotapp/onyx-model-server-cloud' || 'onyxdotapp/onyx-model-server' }}\n steps:\n - name: Check if this scan should run\n id: should-run\n run: |\n case \"$COMPONENT\" in\n web) RESULT=\"$MERGE_WEB\" ;;\n web-cloud) RESULT=\"$MERGE_WEB_CLOUD\" ;;\n backend) RESULT=\"$MERGE_BACKEND\" ;;\n model-server) RESULT=\"$MERGE_MODEL_SERVER\" ;;\n esac\n if [ \"$RESULT\" == \"success\" ]; then\n echo \"run=true\" >> \"$GITHUB_OUTPUT\"\n else\n echo \"run=false\" >> \"$GITHUB_OUTPUT\"\n fi\n env:\n COMPONENT: ${{ matrix.component }}\n MERGE_WEB: ${{ needs.merge-web.result }}\n MERGE_WEB_CLOUD: ${{ needs.merge-web-cloud.result }}\n MERGE_BACKEND: ${{ needs.merge-backend.result }}\n MERGE_MODEL_SERVER: ${{ needs.merge-model-server.result }}\n\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n if: steps.should-run.outputs.run == 'true'\n\n - name: Checkout\n if: steps.should-run.outputs.run == 'true' && matrix.trivyignore != ''\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Determine scan image\n if: steps.should-run.outputs.run == 'true'\n id: scan-image\n run: |\n if [ \"$IS_TEST_RUN\" == \"true\" ]; then\n echo \"image=${RUNS_ON_ECR_CACHE}:${TAG_PREFIX}-${SANITIZED_TAG}\" >> \"$GITHUB_OUTPUT\"\n else\n echo \"image=docker.io/${REGISTRY_IMAGE}:${REF_NAME}\" >> \"$GITHUB_OUTPUT\"\n fi\n env:\n IS_TEST_RUN: ${{ needs.determine-builds.outputs.is-test-run }}\n TAG_PREFIX: ${{ matrix.component }}\n SANITIZED_TAG: ${{ needs.determine-builds.outputs.sanitized-tag }}\n REGISTRY_IMAGE: ${{ matrix.registry-image }}\n REF_NAME: ${{ github.ref_name }}\n\n - name: Run Trivy vulnerability scanner\n if: steps.should-run.outputs.run == 'true'\n uses: aquasecurity/trivy-action@57a97c7e7821a5776cebc9bb87c984fa69cba8f1 # ratchet:aquasecurity/trivy-action@v0.35.0\n with:\n image-ref: ${{ steps.scan-image.outputs.image }}\n severity: CRITICAL,HIGH\n format: \"sarif\"\n output: \"trivy-results.sarif\"\n trivyignores: ${{ matrix.trivyignore }}\n env:\n TRIVY_USERNAME: ${{ secrets.DOCKER_USERNAME }}\n TRIVY_PASSWORD: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Upload Trivy scan results to GitHub Security tab\n if: steps.should-run.outputs.run == 'true'\n uses: github/codeql-action/upload-sarif@ba454b8ab46733eb6145342877cd148270bb77ab\n with:\n sarif_file: \"trivy-results.sarif\"\n\n notify-slack-on-failure:\n needs:\n - determine-builds\n - build-desktop\n - build-web-amd64\n - build-web-arm64\n - merge-web\n - build-web-cloud-amd64\n - build-web-cloud-arm64\n - merge-web-cloud\n - build-backend-amd64\n - build-backend-arm64\n - merge-backend\n - build-backend-craft-amd64\n - build-backend-craft-arm64\n - merge-backend-craft\n - build-model-server-amd64\n - build-model-server-arm64\n - merge-model-server\n if: always() && (needs.build-desktop.result == 'failure' || needs.build-web-amd64.result == 'failure' || needs.build-web-arm64.result == 'failure' || needs.merge-web.result == 'failure' || needs.build-web-cloud-amd64.result == 'failure' || needs.build-web-cloud-arm64.result == 'failure' || needs.merge-web-cloud.result == 'failure' || needs.build-backend-amd64.result == 'failure' || needs.build-backend-arm64.result == 'failure' || needs.merge-backend.result == 'failure' || (needs.determine-builds.outputs.build-backend-craft == 'true' && (needs.build-backend-craft-amd64.result == 'failure' || needs.build-backend-craft-arm64.result == 'failure' || needs.merge-backend-craft.result == 'failure')) || needs.build-model-server-amd64.result == 'failure' || needs.build-model-server-arm64.result == 'failure' || needs.merge-model-server.result == 'failure') && needs.determine-builds.outputs.is-test-run != 'true'\n # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.\n runs-on: ubuntu-slim\n timeout-minutes: 90\n environment: release\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Determine failed jobs\n id: failed-jobs\n shell: bash\n run: |\n FAILED_JOBS=\"\"\n if [ \"${NEEDS_BUILD_DESKTOP_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-desktop\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_WEB_AMD64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-web-amd64\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_WEB_ARM64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-web-arm64\\\\n\"\n fi\n if [ \"${NEEDS_MERGE_WEB_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• merge-web\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_WEB_CLOUD_AMD64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-web-cloud-amd64\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_WEB_CLOUD_ARM64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-web-cloud-arm64\\\\n\"\n fi\n if [ \"${NEEDS_MERGE_WEB_CLOUD_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• merge-web-cloud\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_BACKEND_AMD64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-backend-amd64\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_BACKEND_ARM64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-backend-arm64\\\\n\"\n fi\n if [ \"${NEEDS_MERGE_BACKEND_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• merge-backend\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_MODEL_SERVER_AMD64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-model-server-amd64\\\\n\"\n fi\n if [ \"${NEEDS_BUILD_MODEL_SERVER_ARM64_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• build-model-server-arm64\\\\n\"\n fi\n if [ \"${NEEDS_MERGE_MODEL_SERVER_RESULT}\" == \"failure\" ]; then\n FAILED_JOBS=\"${FAILED_JOBS}• merge-model-server\\\\n\"\n fi\n # Remove trailing \\n and set output\n FAILED_JOBS=$(printf '%s' \"$FAILED_JOBS\" | sed 's/\\\\n$//')\n echo \"jobs=$FAILED_JOBS\" >> \"$GITHUB_OUTPUT\"\n env:\n NEEDS_BUILD_DESKTOP_RESULT: ${{ needs.build-desktop.result }}\n NEEDS_BUILD_WEB_AMD64_RESULT: ${{ needs.build-web-amd64.result }}\n NEEDS_BUILD_WEB_ARM64_RESULT: ${{ needs.build-web-arm64.result }}\n NEEDS_MERGE_WEB_RESULT: ${{ needs.merge-web.result }}\n NEEDS_BUILD_WEB_CLOUD_AMD64_RESULT: ${{ needs.build-web-cloud-amd64.result }}\n NEEDS_BUILD_WEB_CLOUD_ARM64_RESULT: ${{ needs.build-web-cloud-arm64.result }}\n NEEDS_MERGE_WEB_CLOUD_RESULT: ${{ needs.merge-web-cloud.result }}\n NEEDS_BUILD_BACKEND_AMD64_RESULT: ${{ needs.build-backend-amd64.result }}\n NEEDS_BUILD_BACKEND_ARM64_RESULT: ${{ needs.build-backend-arm64.result }}\n NEEDS_MERGE_BACKEND_RESULT: ${{ needs.merge-backend.result }}\n NEEDS_BUILD_MODEL_SERVER_AMD64_RESULT: ${{ needs.build-model-server-amd64.result }}\n NEEDS_BUILD_MODEL_SERVER_ARM64_RESULT: ${{ needs.build-model-server-arm64.result }}\n NEEDS_MERGE_MODEL_SERVER_RESULT: ${{ needs.merge-model-server.result }}\n\n - name: Send Slack notification\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.MONITOR_DEPLOYMENTS_WEBHOOK }}\n failed-jobs: ${{ steps.failed-jobs.outputs.jobs }}\n title: \"🚨 Deployment Workflow Failed\"\n ref-name: ${{ github.ref_name }}\n" }, { "path": ".github/workflows/docker-tag-beta.yml", "content": "# This workflow is set up to be manually triggered via the GitHub Action tab.\n# Given a version, it will tag those backend and webserver images as \"beta\".\n\nname: Tag Beta Version\n\non:\n workflow_dispatch:\n inputs:\n version:\n description: \"The version (ie v1.0.0-beta.0) to tag as beta\"\n required: true\n\npermissions:\n contents: read\n\njobs:\n tag:\n # See https://runs-on.com/runners/linux/\n # use a lower powered instance since this just does i/o to docker hub\n runs-on: [runs-on, runner=2cpu-linux-x64, \"run-id=${{ github.run_id }}-tag\"]\n timeout-minutes: 45\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Enable Docker CLI experimental features\n run: echo \"DOCKER_CLI_EXPERIMENTAL=enabled\" >> $GITHUB_ENV\n\n - name: Pull, Tag and Push Web Server Image\n env:\n VERSION: ${{ github.event.inputs.version }}\n run: |\n docker buildx imagetools create -t onyxdotapp/onyx-web-server:beta onyxdotapp/onyx-web-server:${VERSION}\n\n - name: Pull, Tag and Push API Server Image\n env:\n VERSION: ${{ github.event.inputs.version }}\n run: |\n docker buildx imagetools create -t onyxdotapp/onyx-backend:beta onyxdotapp/onyx-backend:${VERSION}\n\n - name: Pull, Tag and Push Model Server Image\n env:\n VERSION: ${{ github.event.inputs.version }}\n run: |\n docker buildx imagetools create -t onyxdotapp/onyx-model-server:beta onyxdotapp/onyx-model-server:${VERSION}\n" }, { "path": ".github/workflows/docker-tag-latest.yml", "content": "# This workflow is set up to be manually triggered via the GitHub Action tab.\n# Given a version, it will tag those backend and webserver images as \"latest\".\n\nname: Tag Latest Version\n\non:\n workflow_dispatch:\n inputs:\n version:\n description: \"The version (ie v0.0.1) to tag as latest\"\n required: true\n\npermissions:\n contents: read\n\njobs:\n tag:\n # See https://runs-on.com/runners/linux/\n # use a lower powered instance since this just does i/o to docker hub\n runs-on: [runs-on, runner=2cpu-linux-x64, \"run-id=${{ github.run_id }}-tag\"]\n timeout-minutes: 45\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Enable Docker CLI experimental features\n run: echo \"DOCKER_CLI_EXPERIMENTAL=enabled\" >> $GITHUB_ENV\n\n - name: Pull, Tag and Push Web Server Image\n env:\n VERSION: ${{ github.event.inputs.version }}\n run: |\n docker buildx imagetools create -t onyxdotapp/onyx-web-server:latest onyxdotapp/onyx-web-server:${VERSION}\n\n - name: Pull, Tag and Push API Server Image\n env:\n VERSION: ${{ github.event.inputs.version }}\n run: |\n docker buildx imagetools create -t onyxdotapp/onyx-backend:latest onyxdotapp/onyx-backend:${VERSION}\n\n - name: Pull, Tag and Push Model Server Image\n env:\n VERSION: ${{ github.event.inputs.version }}\n run: |\n docker buildx imagetools create -t onyxdotapp/onyx-model-server:latest onyxdotapp/onyx-model-server:${VERSION}\n" }, { "path": ".github/workflows/helm-chart-releases.yml", "content": "name: Release Onyx Helm Charts\n\non:\n push:\n branches:\n - main\n\npermissions: write-all\n\njobs:\n release:\n permissions:\n contents: write\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n fetch-depth: 0\n persist-credentials: false\n\n - name: Install Helm CLI\n uses: azure/setup-helm@1a275c3b69536ee54be43f2070a358922e12c8d4 # ratchet:azure/setup-helm@v4\n with:\n version: v3.12.1\n\n - name: Add required Helm repositories\n run: |\n helm repo add ingress-nginx https://kubernetes.github.io/ingress-nginx\n helm repo add onyx-vespa https://onyx-dot-app.github.io/vespa-helm-charts\n helm repo add opensearch https://opensearch-project.github.io/helm-charts\n helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts\n helm repo add ot-container-kit https://ot-container-kit.github.io/helm-charts\n helm repo add minio https://charts.min.io/\n helm repo add code-interpreter https://onyx-dot-app.github.io/python-sandbox/\n helm repo update\n\n - name: Build chart dependencies\n run: |\n set -euo pipefail\n for chart_dir in deployment/helm/charts/*; do\n if [ -f \"$chart_dir/Chart.yaml\" ]; then\n echo \"Building dependencies for $chart_dir\"\n helm dependency build \"$chart_dir\"\n fi\n done\n\n - name: Publish Helm charts to gh-pages\n # NOTE: HEAD of https://github.com/stefanprodan/helm-gh-pages/pull/43\n uses: stefanprodan/helm-gh-pages@ad32ad3b8720abfeaac83532fd1e9bdfca5bbe27 # zizmor: ignore[impostor-commit]\n with:\n token: ${{ secrets.GITHUB_TOKEN }}\n charts_dir: deployment/helm/charts\n branch: gh-pages\n commit_username: ${{ github.actor }}\n commit_email: ${{ github.actor }}@users.noreply.github.com\n" }, { "path": ".github/workflows/merge-group.yml", "content": "name: Merge Group-Specific\n\non:\n merge_group:\n\npermissions:\n contents: read\n\njobs:\n # This job immediately succeeds to satisfy branch protection rules on merge_group events.\n # There is a similarly named \"required\" job in pr-integration-tests.yml which runs the actual\n # integration tests. That job runs on both pull_request and merge_group events, and this job\n # exists solely to provide a fast-passing check with the same name for branch protection.\n # The actual tests remain enforced on presubmit (pull_request events).\n required:\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Success\n run: echo \"Success\"\n # This job immediately succeeds to satisfy branch protection rules on merge_group events.\n # There is a similarly named \"playwright-required\" job in pr-playwright-tests.yml which runs\n # the actual playwright tests. That job runs on both pull_request and merge_group events, and\n # this job exists solely to provide a fast-passing check with the same name for branch protection.\n # The actual tests remain enforced on presubmit (pull_request events).\n playwright-required:\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Success\n run: echo \"Success\"\n" }, { "path": ".github/workflows/nightly-close-stale-issues.yml", "content": "name: 'Nightly - Close stale issues and PRs'\non:\n schedule:\n - cron: '0 11 * * *' # Runs every day at 3 AM PST / 4 AM PDT / 11 AM UTC\n\npermissions:\n # contents: write # only for delete-branch option\n issues: write\n pull-requests: write\n\njobs:\n stale:\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # ratchet:actions/stale@v10\n with:\n stale-issue-message: 'This issue is stale because it has been open 75 days with no activity. Remove stale label or comment or this will be closed in 15 days.'\n stale-pr-message: 'This PR is stale because it has been open 75 days with no activity. Remove stale label or comment or this will be closed in 15 days.'\n close-issue-message: 'This issue was closed because it has been stalled for 90 days with no activity.'\n close-pr-message: 'This PR was closed because it has been stalled for 90 days with no activity.'\n days-before-stale: 75\n# days-before-close: 90 # uncomment after we test stale behavior\n" }, { "path": ".github/workflows/nightly-llm-provider-chat.yml", "content": "name: Nightly LLM Provider Chat Tests\nconcurrency:\n group: Nightly-LLM-Provider-Chat-${{ github.workflow }}-${{ github.ref_name }}\n cancel-in-progress: true\n\non:\n schedule:\n # Runs daily at 10:30 UTC (2:30 AM PST / 3:30 AM PDT)\n - cron: \"30 10 * * *\"\n workflow_dispatch:\n\npermissions:\n contents: read\n\njobs:\n provider-chat-test:\n uses: ./.github/workflows/reusable-nightly-llm-provider-chat.yml\n secrets:\n AWS_OIDC_ROLE_ARN: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n permissions:\n contents: read\n id-token: write\n with:\n openai_models: ${{ vars.NIGHTLY_LLM_OPENAI_MODELS }}\n anthropic_models: ${{ vars.NIGHTLY_LLM_ANTHROPIC_MODELS }}\n bedrock_models: ${{ vars.NIGHTLY_LLM_BEDROCK_MODELS }}\n vertex_ai_models: ${{ vars.NIGHTLY_LLM_VERTEX_AI_MODELS }}\n azure_models: ${{ vars.NIGHTLY_LLM_AZURE_MODELS }}\n azure_api_base: ${{ vars.NIGHTLY_LLM_AZURE_API_BASE }}\n ollama_models: ${{ vars.NIGHTLY_LLM_OLLAMA_MODELS }}\n openrouter_models: ${{ vars.NIGHTLY_LLM_OPENROUTER_MODELS }}\n strict: true\n\n notify-slack-on-failure:\n needs: [provider-chat-test]\n if: failure() && github.event_name == 'schedule'\n runs-on: ubuntu-slim\n environment: ci-protected\n timeout-minutes: 5\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Send Slack notification\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.SLACK_WEBHOOK }}\n failed-jobs: provider-chat-test\n title: \"🚨 Scheduled LLM Provider Chat Tests failed!\"\n ref-name: ${{ github.ref_name }}\n" }, { "path": ".github/workflows/post-merge-beta-cherry-pick.yml", "content": "name: Post-Merge Beta Cherry-Pick\n\non:\n pull_request_target:\n types:\n - closed\n\n# SECURITY NOTE:\n# This workflow intentionally uses pull_request_target so post-merge automation can\n# use base-repo credentials. Do not checkout PR head refs in this workflow\n# (e.g. github.event.pull_request.head.sha). Only trusted base refs are allowed.\npermissions:\n contents: read\n\njobs:\n resolve-cherry-pick-request:\n if: >-\n github.event.pull_request.merged == true\n && github.event.pull_request.base.ref == 'main'\n && github.event.pull_request.head.repo.full_name == github.repository\n outputs:\n should_cherrypick: ${{ steps.gate.outputs.should_cherrypick }}\n pr_number: ${{ steps.gate.outputs.pr_number }}\n merge_commit_sha: ${{ steps.gate.outputs.merge_commit_sha }}\n merged_by: ${{ steps.gate.outputs.merged_by }}\n gate_error: ${{ steps.gate.outputs.gate_error }}\n runs-on: ubuntu-latest\n timeout-minutes: 10\n steps:\n - name: Resolve merged PR and checkbox state\n id: gate\n env:\n GH_TOKEN: ${{ github.token }}\n PR_NUMBER: ${{ github.event.pull_request.number }}\n # SECURITY: keep PR body in env/plain-text handling; avoid directly\n # inlining github.event.pull_request.body into shell commands.\n PR_BODY: ${{ github.event.pull_request.body }}\n MERGE_COMMIT_SHA: ${{ github.event.pull_request.merge_commit_sha }}\n MERGED_BY: ${{ github.event.pull_request.merged_by.login }}\n # Explicit merger allowlist used because pull_request_target runs with\n # the default GITHUB_TOKEN, which cannot reliably read org/team\n # membership for this repository context.\n ALLOWED_MERGERS: |\n acaprau\n bo-onyx\n danelegend\n duo-onyx\n evan-onyx\n jessicasingh7\n jmelahman\n joachim-danswer\n justin-tahara\n nmgarza5\n raunakab\n rohoswagger\n subash-mohan\n trial2onyx\n wenxi-onyx\n weves\n yuhongsun96\n run: |\n echo \"pr_number=${PR_NUMBER}\" >> \"$GITHUB_OUTPUT\"\n echo \"merged_by=${MERGED_BY}\" >> \"$GITHUB_OUTPUT\"\n\n if ! echo \"${PR_BODY}\" | grep -qiE \"\\\\[x\\\\][[:space:]]*(\\\\[[^]]+\\\\][[:space:]]*)?Please cherry-pick this PR to the latest release version\"; then\n echo \"should_cherrypick=false\" >> \"$GITHUB_OUTPUT\"\n echo \"Cherry-pick checkbox not checked for PR #${PR_NUMBER}. Skipping.\"\n exit 0\n fi\n\n # Keep should_cherrypick output before any possible exit 1 below so\n # notify-slack can still gate on this output even if this job fails.\n echo \"should_cherrypick=true\" >> \"$GITHUB_OUTPUT\"\n echo \"Cherry-pick checkbox checked for PR #${PR_NUMBER}.\"\n\n if [ -z \"${MERGE_COMMIT_SHA}\" ] || [ \"${MERGE_COMMIT_SHA}\" = \"null\" ]; then\n echo \"gate_error=missing-merge-commit-sha\" >> \"$GITHUB_OUTPUT\"\n echo \"::error::PR #${PR_NUMBER} requested cherry-pick, but merge_commit_sha is missing.\"\n exit 1\n fi\n\n echo \"merge_commit_sha=${MERGE_COMMIT_SHA}\" >> \"$GITHUB_OUTPUT\"\n\n normalized_merged_by=\"$(printf '%s' \"${MERGED_BY}\" | tr '[:upper:]' '[:lower:]')\"\n normalized_allowed_mergers=\"$(printf '%s\\n' \"${ALLOWED_MERGERS}\" | tr '[:upper:]' '[:lower:]')\"\n if ! printf '%s\\n' \"${normalized_allowed_mergers}\" | grep -Fxq \"${normalized_merged_by}\"; then\n echo \"gate_error=not-allowed-merger\" >> \"$GITHUB_OUTPUT\"\n echo \"::error::${MERGED_BY} is not in the explicit cherry-pick merger allowlist. Failing cherry-pick gate.\"\n exit 1\n fi\n\n exit 0\n\n cherry-pick-to-latest-release:\n needs:\n - resolve-cherry-pick-request\n if: needs.resolve-cherry-pick-request.outputs.should_cherrypick == 'true' && needs.resolve-cherry-pick-request.result == 'success'\n permissions:\n contents: write\n pull-requests: write\n outputs:\n cherry_pick_pr_url: ${{ steps.run_cherry_pick.outputs.pr_url }}\n cherry_pick_reason: ${{ steps.run_cherry_pick.outputs.reason }}\n cherry_pick_details: ${{ steps.run_cherry_pick.outputs.details }}\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Checkout repository\n # SECURITY: keep checkout pinned to trusted base branch; do not switch to PR head refs.\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n fetch-depth: 0\n persist-credentials: true\n ref: main\n\n - name: Install the latest version of uv\n uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n enable-cache: false\n version: \"0.9.9\"\n\n - name: Configure git identity\n run: |\n git config user.name \"github-actions[bot]\"\n git config user.email \"github-actions[bot]@users.noreply.github.com\"\n\n - name: Create cherry-pick PR to latest release\n id: run_cherry_pick\n env:\n GH_TOKEN: ${{ github.token }}\n GITHUB_TOKEN: ${{ github.token }}\n CHERRY_PICK_ASSIGNEE: ${{ needs.resolve-cherry-pick-request.outputs.merged_by }}\n MERGE_COMMIT_SHA: ${{ needs.resolve-cherry-pick-request.outputs.merge_commit_sha }}\n run: |\n output_file=\"$(mktemp)\"\n set +e\n uv run --no-sync --with onyx-devtools ods cherry-pick \"${MERGE_COMMIT_SHA}\" --yes --no-verify 2>&1 | tee \"$output_file\"\n pipe_statuses=(\"${PIPESTATUS[@]}\")\n exit_code=\"${pipe_statuses[0]}\"\n tee_exit=\"${pipe_statuses[1]:-0}\"\n set -e\n if [ \"${tee_exit}\" -ne 0 ]; then\n echo \"status=failure\" >> \"$GITHUB_OUTPUT\"\n echo \"reason=output-capture-failed\" >> \"$GITHUB_OUTPUT\"\n echo \"::error::tee failed to capture cherry-pick output (exit ${tee_exit}); cannot classify result.\"\n exit 1\n fi\n\n if [ \"${exit_code}\" -eq 0 ]; then\n pr_url=\"$(sed -n 's/^.*PR created successfully: \\(https:\\/\\/github\\.com\\/[^[:space:]]\\+\\/pull\\/[0-9]\\+\\).*$/\\1/p' \"$output_file\" | tail -n 1)\"\n echo \"status=success\" >> \"$GITHUB_OUTPUT\"\n if [ -n \"${pr_url}\" ]; then\n echo \"pr_url=${pr_url}\" >> \"$GITHUB_OUTPUT\"\n fi\n exit 0\n fi\n\n echo \"status=failure\" >> \"$GITHUB_OUTPUT\"\n\n reason=\"command-failed\"\n if grep -qiE \"merge conflict during cherry-pick|CONFLICT|could not apply|cherry-pick in progress with staged changes\" \"$output_file\"; then\n reason=\"merge-conflict\"\n fi\n echo \"reason=${reason}\" >> \"$GITHUB_OUTPUT\"\n\n {\n echo \"details<> \"$GITHUB_OUTPUT\"\n\n - name: Mark workflow as failed if cherry-pick failed\n if: steps.run_cherry_pick.outputs.status == 'failure'\n env:\n CHERRY_PICK_REASON: ${{ steps.run_cherry_pick.outputs.reason }}\n run: |\n echo \"::error::Automated cherry-pick failed (${CHERRY_PICK_REASON}).\"\n exit 1\n\n notify-slack-on-cherry-pick-success:\n needs:\n - resolve-cherry-pick-request\n - cherry-pick-to-latest-release\n if: needs.resolve-cherry-pick-request.outputs.should_cherrypick == 'true' && needs.resolve-cherry-pick-request.result == 'success' && needs.cherry-pick-to-latest-release.result == 'success'\n runs-on: ubuntu-slim\n environment: ci-protected\n timeout-minutes: 10\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Fail if Slack webhook secret is missing\n env:\n CHERRY_PICK_PRS_WEBHOOK: ${{ secrets.CHERRY_PICK_PRS_WEBHOOK }}\n run: |\n if [ -z \"${CHERRY_PICK_PRS_WEBHOOK}\" ]; then\n echo \"::error::CHERRY_PICK_PRS_WEBHOOK is not configured.\"\n exit 1\n fi\n\n - name: Build cherry-pick success summary\n id: success-summary\n env:\n SOURCE_PR_NUMBER: ${{ needs.resolve-cherry-pick-request.outputs.pr_number }}\n MERGE_COMMIT_SHA: ${{ needs.resolve-cherry-pick-request.outputs.merge_commit_sha }}\n CHERRY_PICK_PR_URL: ${{ needs.cherry-pick-to-latest-release.outputs.cherry_pick_pr_url }}\n run: |\n source_pr_url=\"https://github.com/${GITHUB_REPOSITORY}/pull/${SOURCE_PR_NUMBER}\"\n details=\"*Cherry-pick PR opened successfully.*\\\\n• author: {mention}\\\\n• source PR: ${source_pr_url}\"\n if [ -n \"${CHERRY_PICK_PR_URL}\" ]; then\n details=\"${details}\\\\n• cherry-pick PR: ${CHERRY_PICK_PR_URL}\"\n fi\n if [ -n \"${MERGE_COMMIT_SHA}\" ]; then\n details=\"${details}\\\\n• merge SHA: ${MERGE_COMMIT_SHA}\"\n fi\n\n echo \"details=${details}\" >> \"$GITHUB_OUTPUT\"\n\n - name: Notify #cherry-pick-prs about cherry-pick success\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.CHERRY_PICK_PRS_WEBHOOK }}\n mention: ${{ needs.resolve-cherry-pick-request.outputs.merged_by }}\n details: ${{ steps.success-summary.outputs.details }}\n title: \"✅ Automated Cherry-Pick PR Opened\"\n ref-name: ${{ github.event.pull_request.base.ref }}\n\n notify-slack-on-cherry-pick-failure:\n needs:\n - resolve-cherry-pick-request\n - cherry-pick-to-latest-release\n if: always() && needs.resolve-cherry-pick-request.outputs.should_cherrypick == 'true' && (needs.resolve-cherry-pick-request.result == 'failure' || needs.cherry-pick-to-latest-release.result == 'failure')\n runs-on: ubuntu-slim\n environment: ci-protected\n timeout-minutes: 10\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Fail if Slack webhook secret is missing\n env:\n CHERRY_PICK_PRS_WEBHOOK: ${{ secrets.CHERRY_PICK_PRS_WEBHOOK }}\n run: |\n if [ -z \"${CHERRY_PICK_PRS_WEBHOOK}\" ]; then\n echo \"::error::CHERRY_PICK_PRS_WEBHOOK is not configured.\"\n exit 1\n fi\n\n - name: Build cherry-pick failure summary\n id: failure-summary\n env:\n SOURCE_PR_NUMBER: ${{ needs.resolve-cherry-pick-request.outputs.pr_number }}\n MERGE_COMMIT_SHA: ${{ needs.resolve-cherry-pick-request.outputs.merge_commit_sha }}\n GATE_ERROR: ${{ needs.resolve-cherry-pick-request.outputs.gate_error }}\n CHERRY_PICK_REASON: ${{ needs.cherry-pick-to-latest-release.outputs.cherry_pick_reason }}\n CHERRY_PICK_DETAILS: ${{ needs.cherry-pick-to-latest-release.outputs.cherry_pick_details }}\n run: |\n source_pr_url=\"https://github.com/${GITHUB_REPOSITORY}/pull/${SOURCE_PR_NUMBER}\"\n\n reason_text=\"cherry-pick command failed\"\n if [ \"${GATE_ERROR}\" = \"missing-merge-commit-sha\" ]; then\n reason_text=\"requested cherry-pick but merge commit SHA was missing\"\n elif [ \"${GATE_ERROR}\" = \"not-allowed-merger\" ]; then\n reason_text=\"merger is not in the explicit cherry-pick allowlist\"\n elif [ \"${CHERRY_PICK_REASON}\" = \"output-capture-failed\" ]; then\n reason_text=\"failed to capture cherry-pick output for classification\"\n elif [ \"${CHERRY_PICK_REASON}\" = \"merge-conflict\" ]; then\n reason_text=\"merge conflict during cherry-pick\"\n fi\n\n details_excerpt=\"$(printf '%s' \"${CHERRY_PICK_DETAILS}\" | tail -n 8 | tr '\\n' ' ' | sed \"s/[[:space:]]\\\\+/ /g\" | sed \"s/\\\"/'/g\" | cut -c1-350)\"\n if [ -n \"${GATE_ERROR}\" ]; then\n failed_job_label=\"resolve-cherry-pick-request\"\n else\n failed_job_label=\"cherry-pick-to-latest-release\"\n fi\n details=\"• author: {mention}\\\\n• ${failed_job_label}\\\\n• source PR: ${source_pr_url}\\\\n• reason: ${reason_text}\"\n if [ -n \"${MERGE_COMMIT_SHA}\" ]; then\n details=\"${details}\\\\n• merge SHA: ${MERGE_COMMIT_SHA}\"\n fi\n if [ -n \"${details_excerpt}\" ]; then\n details=\"${details}\\\\n• excerpt: ${details_excerpt}\"\n fi\n\n echo \"details=${details}\" >> \"$GITHUB_OUTPUT\"\n\n - name: Notify #cherry-pick-prs about cherry-pick failure\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.CHERRY_PICK_PRS_WEBHOOK }}\n mention: ${{ needs.resolve-cherry-pick-request.outputs.merged_by }}\n details: ${{ steps.failure-summary.outputs.details }}\n title: \"🚨 Automated Cherry-Pick Failed\"\n ref-name: ${{ github.event.pull_request.base.ref }}\n" }, { "path": ".github/workflows/pr-database-tests.yml", "content": "name: Database Tests\nconcurrency:\n group: Database-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\njobs:\n database-tests:\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - \"run-id=${{ github.run_id }}-database-tests\"\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup Python and Install Dependencies\n uses: ./.github/actions/setup-python-and-install-dependencies\n with:\n requirements: |\n backend/requirements/default.txt\n backend/requirements/dev.txt\n\n - name: Generate OpenAPI schema and Python client\n shell: bash\n # TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\n env:\n LICENSE_ENFORCEMENT_ENABLED: \"false\"\n run: |\n ods openapi all\n\n # needed for pulling external images otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Start Docker containers\n working-directory: ./deployment/docker_compose\n run: |\n docker compose -f docker-compose.yml -f docker-compose.dev.yml up -d \\\n relational_db\n\n - name: Run Database Tests\n working-directory: ./backend\n run: pytest -m alembic tests/integration/tests/migrations/\n" }, { "path": ".github/workflows/pr-desktop-build.yml", "content": "name: Build Desktop App\nconcurrency:\n group: Build-Desktop-App-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n paths:\n - \"desktop/**\"\n - \".github/workflows/pr-desktop-build.yml\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\njobs:\n build-desktop:\n name: Build Desktop (${{ matrix.platform }})\n runs-on: ${{ matrix.os }}\n timeout-minutes: 60\n strategy:\n fail-fast: false\n matrix:\n include:\n - platform: linux\n os: ubuntu-latest\n target: x86_64-unknown-linux-gnu\n args: \"--bundles deb,rpm\"\n # TODO: Fix and enable the macOS build.\n #- platform: macos\n # os: macos-latest\n # target: universal-apple-darwin\n # args: \"--target universal-apple-darwin\"\n # TODO: Fix and enable the Windows build.\n #- platform: windows\n # os: windows-latest\n # target: x86_64-pc-windows-msvc\n # args: \"\"\n\n steps:\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd\n with:\n persist-credentials: false\n\n - name: Setup node\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f\n with:\n node-version: 24\n cache: \"npm\" # zizmor: ignore[cache-poisoning]\n cache-dependency-path: ./desktop/package-lock.json\n\n - name: Setup Rust\n uses: dtolnay/rust-toolchain@efa25f7f19611383d5b0ccf2d1c8914531636bf9\n with:\n toolchain: stable\n targets: ${{ matrix.target }}\n\n - name: Cache Cargo registry and build\n uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # zizmor: ignore[cache-poisoning]\n with:\n path: |\n ~/.cargo/bin/\n ~/.cargo/registry/index/\n ~/.cargo/registry/cache/\n ~/.cargo/git/db/\n desktop/src-tauri/target/\n key: ${{ runner.os }}-cargo-${{ hashFiles('desktop/src-tauri/Cargo.lock') }}\n restore-keys: |\n ${{ runner.os }}-cargo-\n\n - name: Install Linux dependencies\n if: matrix.platform == 'linux'\n run: |\n sudo apt-get update\n sudo apt-get install -y \\\n build-essential \\\n libglib2.0-dev \\\n libgirepository1.0-dev \\\n libgtk-3-dev \\\n libjavascriptcoregtk-4.1-dev \\\n libwebkit2gtk-4.1-dev \\\n libayatana-appindicator3-dev \\\n gobject-introspection \\\n pkg-config \\\n curl \\\n xdg-utils\n\n - name: Install npm dependencies\n working-directory: ./desktop\n run: npm ci\n\n - name: Build desktop app\n working-directory: ./desktop\n run: npx tauri build ${{ matrix.args }}\n env:\n TAURI_SIGNING_PRIVATE_KEY: \"\"\n TAURI_SIGNING_PRIVATE_KEY_PASSWORD: \"\"\n\n - name: Upload build artifacts\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: desktop-build-${{ matrix.platform }}-${{ github.run_id }}\n path: |\n desktop/src-tauri/target/release/bundle/\n retention-days: 7\n if-no-files-found: ignore\n" }, { "path": ".github/workflows/pr-external-dependency-unit-tests.yml", "content": "name: External Dependency Unit Tests\nconcurrency:\n group: External-Dependency-Unit-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches: [main]\n paths:\n - \"backend/**\"\n - \"pyproject.toml\"\n - \"uv.lock\"\n - \".github/workflows/pr-external-dependency-unit-tests.yml\"\n - \".github/actions/setup-python-and-install-dependencies/**\"\n - \".github/actions/setup-playwright/**\"\n - \"deployment/docker_compose/docker-compose.yml\"\n - \"deployment/docker_compose/docker-compose.dev.yml\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\nenv:\n # AWS credentials for S3-specific test\n S3_AWS_ACCESS_KEY_ID_FOR_TEST: ${{ secrets.S3_AWS_ACCESS_KEY_ID }}\n S3_AWS_SECRET_ACCESS_KEY_FOR_TEST: ${{ secrets.S3_AWS_SECRET_ACCESS_KEY }}\n\n # MinIO\n S3_ENDPOINT_URL: \"http://localhost:9004\"\n S3_AWS_ACCESS_KEY_ID: \"minioadmin\"\n S3_AWS_SECRET_ACCESS_KEY: \"minioadmin\"\n\n # Confluence\n CONFLUENCE_TEST_SPACE_URL: ${{ vars.CONFLUENCE_TEST_SPACE_URL }}\n CONFLUENCE_TEST_SPACE: ${{ vars.CONFLUENCE_TEST_SPACE }}\n CONFLUENCE_TEST_PAGE_ID: ${{ secrets.CONFLUENCE_TEST_PAGE_ID }}\n CONFLUENCE_USER_NAME: ${{ vars.CONFLUENCE_USER_NAME }}\n CONFLUENCE_ACCESS_TOKEN: ${{ secrets.CONFLUENCE_ACCESS_TOKEN }}\n CONFLUENCE_ACCESS_TOKEN_SCOPED: ${{ secrets.CONFLUENCE_ACCESS_TOKEN_SCOPED }}\n\n # Jira\n JIRA_ADMIN_API_TOKEN: ${{ secrets.JIRA_ADMIN_API_TOKEN }}\n\n # LLMs\n OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}\n ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}\n VERTEX_CREDENTIALS: ${{ secrets.VERTEX_CREDENTIALS }}\n VERTEX_LOCATION: ${{ vars.VERTEX_LOCATION }}\n\n # Code Interpreter\n # TODO: debug why this is failing and enable\n CODE_INTERPRETER_BASE_URL: http://localhost:8000\n\njobs:\n discover-test-dirs:\n # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.\n runs-on: ubuntu-slim\n timeout-minutes: 45\n outputs:\n test-dirs: ${{ steps.set-matrix.outputs.test-dirs }}\n steps:\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Discover test directories\n id: set-matrix\n run: |\n # Find all subdirectories in backend/tests/external_dependency_unit\n dirs=$(find backend/tests/external_dependency_unit -mindepth 1 -maxdepth 1 -type d -exec basename {} \\; | sort | jq -R -s -c 'split(\"\\n\")[:-1]')\n echo \"test-dirs=$dirs\" >> $GITHUB_OUTPUT\n\n external-dependency-unit-tests:\n needs: discover-test-dirs\n # Use larger runner with more resources for Vespa\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - ${{ format('run-id={0}-external-dependency-unit-tests-job-{1}', github.run_id, strategy['job-index']) }}\n - extras=s3-cache\n timeout-minutes: 45\n strategy:\n fail-fast: false\n matrix:\n test-dir: ${{ fromJson(needs.discover-test-dirs.outputs.test-dirs) }}\n\n env:\n PYTHONPATH: ./backend\n MODEL_SERVER_HOST: \"disabled\"\n DISABLE_TELEMETRY: \"true\"\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup Python and Install Dependencies\n uses: ./.github/actions/setup-python-and-install-dependencies\n with:\n requirements: |\n backend/requirements/default.txt\n backend/requirements/dev.txt\n backend/requirements/ee.txt\n\n - name: Setup Playwright\n uses: ./.github/actions/setup-playwright\n\n # needed for pulling Vespa, Redis, Postgres, and Minio images\n # otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Create .env file for Docker Compose\n run: |\n cat < deployment/docker_compose/.env\n COMPOSE_PROFILES=s3-filestore,opensearch-enabled\n DISABLE_TELEMETRY=true\n OPENSEARCH_FOR_ONYX_ENABLED=true\n EOF\n\n - name: Set up Standard Dependencies\n run: |\n cd deployment/docker_compose\n docker compose \\\n -f docker-compose.yml \\\n -f docker-compose.dev.yml \\\n up -d \\\n minio \\\n relational_db \\\n cache \\\n index \\\n opensearch \\\n code-interpreter\n\n - name: Run migrations\n run: |\n cd backend\n # Run migrations to head\n alembic upgrade head\n alembic heads --verbose\n\n - name: Run Tests for ${{ matrix.test-dir }}\n shell: script -q -e -c \"bash --noprofile --norc -eo pipefail {0}\"\n env:\n TEST_DIR: ${{ matrix.test-dir }}\n run: |\n py.test \\\n --durations=8 \\\n -o junit_family=xunit2 \\\n -xv \\\n --ff \\\n backend/tests/external_dependency_unit/${TEST_DIR}\n\n - name: Collect Docker logs on failure\n if: failure()\n run: |\n mkdir -p docker-logs\n cd deployment/docker_compose\n\n # Get list of running containers\n containers=$(docker compose -f docker-compose.yml -f docker-compose.dev.yml ps -q)\n\n # Collect logs from each container\n for container in $containers; do\n container_name=$(docker inspect --format='{{.Name}}' $container | sed 's/^\\///')\n echo \"Collecting logs from $container_name...\"\n docker logs $container > ../../docker-logs/${container_name}.log 2>&1\n done\n\n cd ../..\n echo \"Docker logs collected in docker-logs directory\"\n\n - name: Upload Docker logs\n if: failure()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-logs-${{ matrix.test-dir }}\n path: docker-logs/\n retention-days: 7\n" }, { "path": ".github/workflows/pr-golang-tests.yml", "content": "name: Golang Tests\nconcurrency:\n group: Golang-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions: {}\n\nenv:\n GO_VERSION: \"1.26\"\n\njobs:\n detect-modules:\n runs-on: ubuntu-latest\n timeout-minutes: 10\n outputs:\n modules: ${{ steps.set-modules.outputs.modules }}\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd\n with:\n persist-credentials: false\n - id: set-modules\n run: echo \"modules=$(find . -name 'go.mod' -exec dirname {} \\; | jq -Rc '[.,inputs]')\" >> \"$GITHUB_OUTPUT\"\n\n golang:\n needs: detect-modules\n runs-on: ubuntu-latest\n timeout-minutes: 10\n strategy:\n matrix:\n modules: ${{ fromJSON(needs.detect-modules.outputs.modules) }}\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n - uses: actions/setup-go@4dc6199c7b1a012772edbd06daecab0f50c9053c # zizmor: ignore[cache-poisoning]\n with:\n go-version: ${{ env.GO_VERSION }}\n cache-dependency-path: \"**/go.sum\"\n\n - run: go mod tidy\n working-directory: ${{ matrix.modules }}\n - run: git diff --exit-code go.mod go.sum\n working-directory: ${{ matrix.modules }}\n\n - run: go test ./...\n working-directory: ${{ matrix.modules }}\n" }, { "path": ".github/workflows/pr-helm-chart-testing.yml", "content": "name: Helm - Lint and Test Charts\nconcurrency:\n group: Helm-Lint-and-Test-Charts-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches: [main]\n push:\n tags:\n - \"v*.*.*\"\n workflow_dispatch: # Allows manual triggering\n\npermissions:\n contents: read\n\njobs:\n helm-chart-check:\n # See https://runs-on.com/runners/linux/\n runs-on:\n [\n runs-on,\n runner=8cpu-linux-x64,\n hdd=256,\n \"run-id=${{ github.run_id }}-helm-chart-check\",\n ]\n timeout-minutes: 45\n\n # fetch-depth 0 is required for helm/chart-testing-action\n steps:\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n fetch-depth: 0\n persist-credentials: false\n\n - name: Set up Helm\n uses: azure/setup-helm@1a275c3b69536ee54be43f2070a358922e12c8d4 # ratchet:azure/setup-helm@v4.3.1\n with:\n version: v3.19.0\n\n - name: Set up chart-testing\n uses: helm/chart-testing-action@2e2940618cb426dce2999631d543b53cdcfc8527\n with:\n uv_version: \"0.9.9\"\n\n # even though we specify chart-dirs in ct.yaml, it isn't used by ct for the list-changed command...\n - name: Run chart-testing (list-changed)\n id: list-changed\n env:\n DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}\n run: |\n echo \"default_branch: ${DEFAULT_BRANCH}\"\n changed=$(ct list-changed --remote origin --target-branch ${DEFAULT_BRANCH} --chart-dirs deployment/helm/charts)\n echo \"list-changed output: $changed\"\n if [[ -n \"$changed\" ]]; then\n echo \"changed=true\" >> \"$GITHUB_OUTPUT\"\n fi\n\n # uncomment to force run chart-testing\n # - name: Force run chart-testing (list-changed)\n # id: list-changed\n # run: echo \"changed=true\" >> $GITHUB_OUTPUT\n # lint all charts if any changes were detected\n - name: Run chart-testing (lint)\n if: steps.list-changed.outputs.changed == 'true'\n run: ct lint --config ct.yaml --all\n # the following would lint only changed charts, but linting isn't expensive\n # run: ct lint --config ct.yaml --target-branch ${{ github.event.repository.default_branch }}\n\n - name: Create kind cluster\n if: steps.list-changed.outputs.changed == 'true'\n uses: helm/kind-action@ef37e7f390d99f746eb8b610417061a60e82a6cc # ratchet:helm/kind-action@v1.14.0\n\n - name: Pre-install cluster status check\n if: steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Pre-install Cluster Status ===\"\n kubectl get nodes -o wide\n kubectl get pods --all-namespaces\n kubectl get storageclass\n\n - name: Add Helm repositories and update\n if: steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Adding Helm repositories ===\"\n helm repo add ingress-nginx https://kubernetes.github.io/ingress-nginx\n helm repo add vespa https://onyx-dot-app.github.io/vespa-helm-charts\n helm repo add opensearch https://opensearch-project.github.io/helm-charts\n helm repo add cloudnative-pg https://cloudnative-pg.github.io/charts\n helm repo add ot-container-kit https://ot-container-kit.github.io/helm-charts\n helm repo add minio https://charts.min.io/\n helm repo add code-interpreter https://onyx-dot-app.github.io/python-sandbox/\n helm repo update\n\n - name: Install Redis operator\n if: steps.list-changed.outputs.changed == 'true'\n shell: bash\n run: |\n echo \"=== Installing redis-operator CRDs ===\"\n helm upgrade --install redis-operator ot-container-kit/redis-operator \\\n --namespace redis-operator --create-namespace --wait --timeout 300s\n\n - name: Pre-pull required images\n if: steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Pre-pulling required images to avoid timeout ===\"\n KIND_CLUSTER=$(kubectl config current-context | sed 's/kind-//')\n echo \"Kind cluster: $KIND_CLUSTER\"\n\n IMAGES=(\n \"ghcr.io/cloudnative-pg/cloudnative-pg:1.27.0\"\n \"quay.io/opstree/redis:v7.0.15\"\n \"docker.io/onyxdotapp/onyx-web-server:latest\"\n )\n\n for image in \"${IMAGES[@]}\"; do\n echo \"Pre-pulling $image\"\n if docker pull \"$image\"; then\n kind load docker-image \"$image\" --name \"$KIND_CLUSTER\" || echo \"Failed to load $image into kind\"\n else\n echo \"Failed to pull $image\"\n fi\n done\n\n echo \"=== Images loaded into Kind cluster ===\"\n docker exec \"$KIND_CLUSTER\"-control-plane crictl images | grep -E \"(cloudnative-pg|redis|onyx)\" || echo \"Some images may still be loading...\"\n\n - name: Validate chart dependencies\n if: steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Validating chart dependencies ===\"\n cd deployment/helm/charts/onyx\n helm dependency update\n helm lint . --set auth.userauth.values.user_auth_secret=placeholder\n\n - name: Run chart-testing (install) with enhanced monitoring\n timeout-minutes: 25\n if: steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Starting chart installation with monitoring ===\"\n\n # Function to monitor cluster state\n monitor_cluster() {\n while true; do\n echo \"=== Cluster Status Check at $(date) ===\"\n # Only show non-running pods to reduce noise\n NON_RUNNING_PODS=$(kubectl get pods --all-namespaces --field-selector=status.phase!=Running,status.phase!=Succeeded --no-headers 2>/dev/null | wc -l)\n if [ \"$NON_RUNNING_PODS\" -gt 0 ]; then\n echo \"Non-running pods:\"\n kubectl get pods --all-namespaces --field-selector=status.phase!=Running,status.phase!=Succeeded\n else\n echo \"All pods running successfully\"\n fi\n # Only show recent events if there are issues\n RECENT_EVENTS=$(kubectl get events --sort-by=.lastTimestamp --all-namespaces --field-selector=type!=Normal 2>/dev/null | tail -5)\n if [ -n \"$RECENT_EVENTS\" ]; then\n echo \"Recent warnings/errors:\"\n echo \"$RECENT_EVENTS\"\n fi\n sleep 60\n done\n }\n\n # Start monitoring in background\n monitor_cluster &\n MONITOR_PID=$!\n\n # Set up cleanup\n cleanup() {\n echo \"=== Cleaning up monitoring process ===\"\n kill $MONITOR_PID 2>/dev/null || true\n echo \"=== Final cluster state ===\"\n kubectl get pods --all-namespaces\n kubectl get events --all-namespaces --sort-by=.lastTimestamp | tail -20\n }\n\n # Trap cleanup on exit\n trap cleanup EXIT\n\n # Run the actual installation with detailed logging\n # Note that opensearch.enabled is true whereas others in this install\n # are false. There is some work that needs to be done to get this\n # entire step working in CI, enabling opensearch here is a small step\n # in that direction. If this is causing issues, disabling it in this\n # step should be ok in the short term.\n echo \"=== Starting ct install ===\"\n set +e\n ct install --all \\\n --helm-extra-set-args=\"\\\n --set=nginx.enabled=false \\\n --set=minio.enabled=false \\\n --set=vespa.enabled=false \\\n --set=opensearch.enabled=true \\\n --set=auth.opensearch.enabled=true \\\n --set=auth.userauth.values.user_auth_secret=test-secret \\\n --set=slackbot.enabled=false \\\n --set=postgresql.enabled=true \\\n --set=postgresql.cluster.storage.storageClass=standard \\\n --set=redis.enabled=true \\\n --set=redis.storageSpec.volumeClaimTemplate.spec.storageClassName=standard \\\n --set=webserver.replicaCount=1 \\\n --set=api.replicaCount=0 \\\n --set=inferenceCapability.replicaCount=0 \\\n --set=indexCapability.replicaCount=0 \\\n --set=celery_beat.replicaCount=0 \\\n --set=celery_worker_heavy.replicaCount=0 \\\n --set=celery_worker_docfetching.replicaCount=0 \\\n --set=celery_worker_docprocessing.replicaCount=0 \\\n --set=celery_worker_light.replicaCount=0 \\\n --set=celery_worker_monitoring.replicaCount=0 \\\n --set=celery_worker_primary.replicaCount=0 \\\n --set=celery_worker_user_file_processing.replicaCount=0 \\\n --set=celery_worker_user_files_indexing.replicaCount=0\" \\\n --helm-extra-args=\"--timeout 900s --debug\" \\\n --debug --config ct.yaml\n CT_EXIT=$?\n set -e\n\n if [[ $CT_EXIT -ne 0 ]]; then\n echo \"ct install failed with exit code $CT_EXIT\"\n exit $CT_EXIT\n else\n echo \"=== Installation completed successfully ===\"\n fi\n\n kubectl get pods --all-namespaces\n\n - name: Post-install verification\n if: steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Post-install verification ===\"\n if ! kubectl cluster-info >/dev/null 2>&1; then\n echo \"ERROR: Kubernetes cluster is not reachable after install\"\n exit 1\n fi\n kubectl get pods --all-namespaces\n kubectl get services --all-namespaces\n # Only show issues if they exist\n kubectl describe pods --all-namespaces | grep -A 5 -B 2 \"Failed\\|Error\\|Warning\" || echo \"No pod issues found\"\n\n - name: Cleanup on failure\n if: failure() && steps.list-changed.outputs.changed == 'true'\n run: |\n echo \"=== Cleanup on failure ===\"\n if ! kubectl cluster-info >/dev/null 2>&1; then\n echo \"Skipping failure cleanup: Kubernetes cluster is not reachable\"\n exit 0\n fi\n echo \"=== Final cluster state ===\"\n kubectl get pods --all-namespaces\n kubectl get events --all-namespaces --sort-by=.lastTimestamp | tail -10\n\n echo \"=== Pod descriptions for debugging ===\"\n kubectl describe pods --all-namespaces | grep -A 10 -B 3 \"Failed\\|Error\\|Warning\\|Pending\" || echo \"No problematic pods found\"\n\n echo \"=== Recent logs for debugging ===\"\n kubectl logs --all-namespaces --tail=50 | grep -i \"error\\|timeout\\|failed\\|pull\" || echo \"No error logs found\"\n\n echo \"=== Helm releases ===\"\n helm list --all-namespaces\n # the following would install only changed charts, but we only have one chart so\n # don't worry about that for now\n # run: ct install --target-branch ${{ github.event.repository.default_branch }}\n" }, { "path": ".github/workflows/pr-integration-tests.yml", "content": "name: Run Integration Tests v2\nconcurrency:\n group: Run-Integration-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\nenv:\n # Test Environment Variables\n OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}\n SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}\n SLACK_BOT_TOKEN_TEST_SPACE: ${{ secrets.SLACK_BOT_TOKEN_TEST_SPACE }}\n CONFLUENCE_TEST_SPACE_URL: ${{ vars.CONFLUENCE_TEST_SPACE_URL }}\n CONFLUENCE_USER_NAME: ${{ vars.CONFLUENCE_USER_NAME }}\n CONFLUENCE_ACCESS_TOKEN: ${{ secrets.CONFLUENCE_ACCESS_TOKEN }}\n CONFLUENCE_ACCESS_TOKEN_SCOPED: ${{ secrets.CONFLUENCE_ACCESS_TOKEN_SCOPED }}\n JIRA_BASE_URL: ${{ secrets.JIRA_BASE_URL }}\n JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}\n JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}\n JIRA_API_TOKEN_SCOPED: ${{ secrets.JIRA_API_TOKEN_SCOPED }}\n PERM_SYNC_SHAREPOINT_CLIENT_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_CLIENT_ID }}\n PERM_SYNC_SHAREPOINT_PRIVATE_KEY: ${{ secrets.PERM_SYNC_SHAREPOINT_PRIVATE_KEY }}\n PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD: ${{ secrets.PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD }}\n PERM_SYNC_SHAREPOINT_DIRECTORY_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_DIRECTORY_ID }}\n EXA_API_KEY: ${{ secrets.EXA_API_KEY }}\n GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN: ${{ secrets.ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN }}\n GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC: ${{ secrets.ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC }}\n GITHUB_ADMIN_EMAIL: ${{ secrets.ONYX_GITHUB_ADMIN_EMAIL }}\n GITHUB_TEST_USER_1_EMAIL: ${{ secrets.ONYX_GITHUB_TEST_USER_1_EMAIL }}\n GITHUB_TEST_USER_2_EMAIL: ${{ secrets.ONYX_GITHUB_TEST_USER_2_EMAIL }}\n\njobs:\n discover-test-dirs:\n # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.\n runs-on: ubuntu-slim\n timeout-minutes: 45\n outputs:\n test-dirs: ${{ steps.set-matrix.outputs.test-dirs }}\n editions: ${{ steps.set-editions.outputs.editions }}\n steps:\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Discover test directories\n id: set-matrix\n run: |\n # Find all leaf-level directories in both test directories\n tests_dirs=$(find backend/tests/integration/tests -mindepth 1 -maxdepth 1 -type d ! -name \"__pycache__\" ! -name \"mcp\" ! -name \"no_vectordb\" -exec basename {} \\; | sort)\n connector_dirs=$(find backend/tests/integration/connector_job_tests -mindepth 1 -maxdepth 1 -type d ! -name \"__pycache__\" -exec basename {} \\; | sort)\n\n # Create JSON array with directory info\n all_dirs=\"\"\n for dir in $tests_dirs; do\n all_dirs=\"$all_dirs{\\\"path\\\":\\\"tests/$dir\\\",\\\"name\\\":\\\"tests-$dir\\\"},\"\n done\n for dir in $connector_dirs; do\n all_dirs=\"$all_dirs{\\\"path\\\":\\\"connector_job_tests/$dir\\\",\\\"name\\\":\\\"connector-$dir\\\"},\"\n done\n\n # Remove trailing comma and wrap in array\n all_dirs=\"[${all_dirs%,}]\"\n echo \"test-dirs=$all_dirs\" >> $GITHUB_OUTPUT\n\n - name: Determine editions to test\n id: set-editions\n run: |\n # On PRs, only run EE tests. On merge_group and tags, run both EE and MIT.\n if [ \"${{ github.event_name }}\" = \"pull_request\" ]; then\n echo 'editions=[\"ee\"]' >> $GITHUB_OUTPUT\n else\n echo 'editions=[\"ee\",\"mit\"]' >> $GITHUB_OUTPUT\n fi\n\n build-backend-image:\n runs-on:\n [\n runs-on,\n runner=1cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-backend-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n # needed for pulling Vespa, Redis, Postgres, and Minio images\n # otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Build and push Backend Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n push: true\n tags: ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-backend-test-${{ github.run_id }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache\n type=registry,ref=onyxdotapp/onyx-backend:latest\n cache-to: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache,mode=max\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-model-server-image:\n runs-on:\n [\n runs-on,\n runner=1cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-model-server-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n # needed for pulling Vespa, Redis, Postgres, and Minio images\n # otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Build and push Model Server Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile.model_server\n push: true\n tags: ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-model-server-test-${{ github.run_id }}\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache\n type=registry,ref=onyxdotapp/onyx-model-server:latest\n cache-to: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max\n\n build-integration-image:\n runs-on:\n [\n runs-on,\n runner=2cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-integration-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n # needed for pulling openapitools/openapi-generator-cli\n # otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Build and push integration test image with Docker Bake\n env:\n INTEGRATION_REPOSITORY: ${{ env.RUNS_ON_ECR_CACHE }}\n TAG: integration-test-${{ github.run_id }}\n CACHE_SUFFIX: ${{ steps.format-branch.outputs.cache-suffix }}\n HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }}\n run: |\n docker buildx bake --push \\\n --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA} \\\n --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX} \\\n --set backend.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache \\\n --set backend.cache-from=type=registry,ref=onyxdotapp/onyx-backend:latest \\\n --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${HEAD_SHA},mode=max \\\n --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache-${CACHE_SUFFIX},mode=max \\\n --set backend.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:backend-cache,mode=max \\\n --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA} \\\n --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX} \\\n --set integration.cache-from=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache \\\n --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${HEAD_SHA},mode=max \\\n --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache-${CACHE_SUFFIX},mode=max \\\n --set integration.cache-to=type=registry,ref=${RUNS_ON_ECR_CACHE}:integration-cache,mode=max \\\n integration\n\n integration-tests:\n needs:\n [\n discover-test-dirs,\n build-backend-image,\n build-model-server-image,\n build-integration-image,\n ]\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - ${{ format('run-id={0}-integration-tests-{1}-job-{2}', github.run_id, matrix.edition, strategy['job-index']) }}\n - extras=ecr-cache\n timeout-minutes: 45\n\n strategy:\n fail-fast: false\n matrix:\n test-dir: ${{ fromJson(needs.discover-test-dirs.outputs.test-dirs) }}\n edition: ${{ fromJson(needs.discover-test-dirs.outputs.editions) }}\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n # needed for pulling Vespa, Redis, Postgres, and Minio images\n # otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n # NOTE: Use pre-ping/null pool to reduce flakiness due to dropped connections\n # NOTE: don't need web server for integration tests\n - name: Create .env file for Docker Compose\n env:\n ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}\n RUN_ID: ${{ github.run_id }}\n EDITION: ${{ matrix.edition }}\n run: |\n # Base config shared by both editions\n cat < deployment/docker_compose/.env\n COMPOSE_PROFILES=s3-filestore\n OPENSEARCH_FOR_ONYX_ENABLED=false\n AUTH_TYPE=basic\n POSTGRES_POOL_PRE_PING=true\n POSTGRES_USE_NULL_POOL=true\n REQUIRE_EMAIL_VERIFICATION=false\n DISABLE_TELEMETRY=true\n ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID}\n ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID}\n INTEGRATION_TESTS_MODE=true\n MCP_SERVER_ENABLED=true\n AUTO_LLM_UPDATE_INTERVAL_SECONDS=10\n EOF\n\n # EE-only config\n if [ \"$EDITION\" = \"ee\" ]; then\n cat <> deployment/docker_compose/.env\n ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true\n # TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\n LICENSE_ENFORCEMENT_ENABLED=false\n CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS=0.001\n EOF\n fi\n\n - name: Start Docker containers\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.dev.yml up \\\n relational_db \\\n index \\\n cache \\\n minio \\\n api_server \\\n inference_model_server \\\n indexing_model_server \\\n background \\\n -d\n id: start_docker\n\n - name: Wait for services to be ready\n run: |\n echo \"Starting wait-for-service script...\"\n\n wait_for_service() {\n local url=$1\n local label=$2\n local timeout=${3:-300} # default 5 minutes\n local start_time\n start_time=$(date +%s)\n\n while true; do\n local current_time\n current_time=$(date +%s)\n local elapsed_time=$((current_time - start_time))\n\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. ${label} did not become ready in $timeout seconds.\"\n exit 1\n fi\n\n local response\n response=$(curl -s -o /dev/null -w \"%{http_code}\" \"$url\" || echo \"curl_error\")\n\n if [ \"$response\" = \"200\" ]; then\n echo \"${label} is ready!\"\n break\n elif [ \"$response\" = \"curl_error\" ]; then\n echo \"Curl encountered an error while checking ${label}. Retrying in 5 seconds...\"\n else\n echo \"${label} not ready yet (HTTP status $response). Retrying in 5 seconds...\"\n fi\n\n sleep 5\n done\n }\n\n wait_for_service \"http://localhost:8080/health\" \"API server\"\n echo \"Finished waiting for services.\"\n\n - name: Start Mock Services\n run: |\n cd backend/tests/integration/mock_services\n docker compose -f docker-compose.mock-it-services.yml \\\n -p mock-it-services-stack up -d\n\n - name: Run Integration Tests (${{ matrix.edition }}) for ${{ matrix.test-dir.name }}\n uses: nick-fields/retry@ce71cc2ab81d554ebbe88c79ab5975992d79ba08 # ratchet:nick-fields/retry@v3\n with:\n timeout_minutes: 20\n max_attempts: 3\n retry_wait_seconds: 10\n command: |\n echo \"Running ${{ matrix.edition }} integration tests for ${{ matrix.test-dir.path }}...\"\n docker run --rm --network onyx_default \\\n --name test-runner \\\n -e POSTGRES_HOST=relational_db \\\n -e POSTGRES_USER=postgres \\\n -e POSTGRES_PASSWORD=password \\\n -e POSTGRES_DB=postgres \\\n -e DB_READONLY_USER=db_readonly_user \\\n -e DB_READONLY_PASSWORD=password \\\n -e POSTGRES_POOL_PRE_PING=true \\\n -e POSTGRES_USE_NULL_POOL=true \\\n -e VESPA_HOST=index \\\n -e ENABLE_OPENSEARCH_INDEXING_FOR_ONYX=false \\\n -e REDIS_HOST=cache \\\n -e API_SERVER_HOST=api_server \\\n -e OPENAI_API_KEY=${OPENAI_API_KEY} \\\n -e EXA_API_KEY=${EXA_API_KEY} \\\n -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \\\n -e SLACK_BOT_TOKEN_TEST_SPACE=${SLACK_BOT_TOKEN_TEST_SPACE} \\\n -e CONFLUENCE_TEST_SPACE_URL=${CONFLUENCE_TEST_SPACE_URL} \\\n -e CONFLUENCE_USER_NAME=${CONFLUENCE_USER_NAME} \\\n -e CONFLUENCE_ACCESS_TOKEN=${CONFLUENCE_ACCESS_TOKEN} \\\n -e CONFLUENCE_ACCESS_TOKEN_SCOPED=${CONFLUENCE_ACCESS_TOKEN_SCOPED} \\\n -e JIRA_BASE_URL=${JIRA_BASE_URL} \\\n -e JIRA_USER_EMAIL=${JIRA_USER_EMAIL} \\\n -e JIRA_API_TOKEN=${JIRA_API_TOKEN} \\\n -e JIRA_API_TOKEN_SCOPED=${JIRA_API_TOKEN_SCOPED} \\\n -e PERM_SYNC_SHAREPOINT_CLIENT_ID=${PERM_SYNC_SHAREPOINT_CLIENT_ID} \\\n -e PERM_SYNC_SHAREPOINT_PRIVATE_KEY=\"${PERM_SYNC_SHAREPOINT_PRIVATE_KEY}\" \\\n -e PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD=${PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD} \\\n -e PERM_SYNC_SHAREPOINT_DIRECTORY_ID=${PERM_SYNC_SHAREPOINT_DIRECTORY_ID} \\\n -e GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN=${GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN} \\\n -e GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC=${GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC} \\\n -e GITHUB_ADMIN_EMAIL=${GITHUB_ADMIN_EMAIL} \\\n -e GITHUB_TEST_USER_1_EMAIL=${GITHUB_TEST_USER_1_EMAIL} \\\n -e GITHUB_TEST_USER_2_EMAIL=${GITHUB_TEST_USER_2_EMAIL} \\\n -e TEST_WEB_HOSTNAME=test-runner \\\n -e MOCK_CONNECTOR_SERVER_HOST=mock_connector_server \\\n -e MOCK_CONNECTOR_SERVER_PORT=8001 \\\n -e ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=${{ matrix.edition == 'ee' && 'true' || 'false' }} \\\n ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-${{ github.run_id }} \\\n /app/tests/integration/${{ matrix.test-dir.path }}\n\n # ------------------------------------------------------------\n # Always gather logs BEFORE \"down\":\n - name: Dump API server logs\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose logs --no-color api_server > $GITHUB_WORKSPACE/api_server.log || true\n\n - name: Dump all-container logs (optional)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose logs --no-color > $GITHUB_WORKSPACE/docker-compose.log || true\n\n - name: Upload logs\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-all-logs-${{ matrix.edition }}-${{ matrix.test-dir.name }}\n path: ${{ github.workspace }}/docker-compose.log\n # ------------------------------------------------------------\n\n onyx-lite-tests:\n needs: [build-backend-image, build-integration-image]\n runs-on:\n [\n runs-on,\n runner=4cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-onyx-lite-tests\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Create .env file for Onyx Lite Docker Compose\n env:\n ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}\n RUN_ID: ${{ github.run_id }}\n run: |\n cat < deployment/docker_compose/.env\n ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true\n LICENSE_ENFORCEMENT_ENABLED=false\n AUTH_TYPE=basic\n POSTGRES_POOL_PRE_PING=true\n POSTGRES_USE_NULL_POOL=true\n REQUIRE_EMAIL_VERIFICATION=false\n DISABLE_TELEMETRY=true\n ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID}\n INTEGRATION_TESTS_MODE=true\n EOF\n\n # Start only the services needed for Onyx Lite (Postgres + API server)\n - name: Start Docker containers (onyx-lite)\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.onyx-lite.yml -f docker-compose.dev.yml up \\\n relational_db \\\n api_server \\\n -d\n id: start_docker_onyx_lite\n\n - name: Wait for services to be ready\n run: |\n echo \"Starting wait-for-service script (onyx-lite)...\"\n start_time=$(date +%s)\n timeout=300\n while true; do\n current_time=$(date +%s)\n elapsed_time=$((current_time - start_time))\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. Service did not become ready in $timeout seconds.\"\n exit 1\n fi\n response=$(curl -s -o /dev/null -w \"%{http_code}\" http://localhost:8080/health || echo \"curl_error\")\n if [ \"$response\" = \"200\" ]; then\n echo \"API server is ready!\"\n break\n elif [ \"$response\" = \"curl_error\" ]; then\n echo \"Curl encountered an error; retrying...\"\n else\n echo \"Service not ready yet (HTTP $response). Retrying in 5 seconds...\"\n fi\n sleep 5\n done\n\n - name: Run Onyx Lite Integration Tests\n uses: nick-fields/retry@ce71cc2ab81d554ebbe88c79ab5975992d79ba08 # ratchet:nick-fields/retry@v3\n with:\n timeout_minutes: 20\n max_attempts: 3\n retry_wait_seconds: 10\n command: |\n echo \"Running onyx-lite integration tests...\"\n docker run --rm --network onyx_default \\\n --name test-runner \\\n -e POSTGRES_HOST=relational_db \\\n -e POSTGRES_USER=postgres \\\n -e POSTGRES_PASSWORD=password \\\n -e POSTGRES_DB=postgres \\\n -e DB_READONLY_USER=db_readonly_user \\\n -e DB_READONLY_PASSWORD=password \\\n -e POSTGRES_POOL_PRE_PING=true \\\n -e POSTGRES_USE_NULL_POOL=true \\\n -e API_SERVER_HOST=api_server \\\n -e OPENAI_API_KEY=${OPENAI_API_KEY} \\\n -e TEST_WEB_HOSTNAME=test-runner \\\n ${{ env.RUNS_ON_ECR_CACHE }}:integration-test-${{ github.run_id }} \\\n /app/tests/integration/tests/no_vectordb\n\n - name: Dump API server logs (onyx-lite)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.onyx-lite.yml -f docker-compose.dev.yml \\\n logs --no-color api_server > $GITHUB_WORKSPACE/api_server_onyx_lite.log || true\n\n - name: Dump all-container logs (onyx-lite)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.onyx-lite.yml -f docker-compose.dev.yml \\\n logs --no-color > $GITHUB_WORKSPACE/docker-compose-onyx-lite.log || true\n\n - name: Upload logs (onyx-lite)\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-all-logs-onyx-lite\n path: ${{ github.workspace }}/docker-compose-onyx-lite.log\n\n - name: Stop Docker containers (onyx-lite)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.onyx-lite.yml -f docker-compose.dev.yml down -v\n\n multitenant-tests:\n needs:\n [build-backend-image, build-model-server-image, build-integration-image]\n runs-on:\n [\n runs-on,\n runner=8cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-multitenant-tests\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Start Docker containers for multi-tenant tests\n env:\n ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}\n RUN_ID: ${{ github.run_id }}\n run: |\n cd deployment/docker_compose\n ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \\\n LICENSE_ENFORCEMENT_ENABLED=false \\\n MULTI_TENANT=true \\\n AUTH_TYPE=cloud \\\n REQUIRE_EMAIL_VERIFICATION=false \\\n DISABLE_TELEMETRY=true \\\n OPENAI_DEFAULT_API_KEY=${OPENAI_API_KEY} \\\n ONYX_BACKEND_IMAGE=${ECR_CACHE}:integration-test-backend-test-${RUN_ID} \\\n ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:integration-test-model-server-test-${RUN_ID} \\\n DEV_MODE=true \\\n OPENSEARCH_FOR_ONYX_ENABLED=false \\\n docker compose -f docker-compose.multitenant-dev.yml up \\\n relational_db \\\n index \\\n cache \\\n minio \\\n api_server \\\n inference_model_server \\\n indexing_model_server \\\n background \\\n -d\n id: start_docker_multi_tenant\n\n - name: Wait for service to be ready (multi-tenant)\n run: |\n echo \"Starting wait-for-service script for multi-tenant...\"\n docker logs -f onyx-api_server-1 &\n start_time=$(date +%s)\n timeout=300\n while true; do\n current_time=$(date +%s)\n elapsed_time=$((current_time - start_time))\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. Service did not become ready in 5 minutes.\"\n exit 1\n fi\n response=$(curl -s -o /dev/null -w \"%{http_code}\" http://localhost:8080/health || echo \"curl_error\")\n if [ \"$response\" = \"200\" ]; then\n echo \"Service is ready!\"\n break\n elif [ \"$response\" = \"curl_error\" ]; then\n echo \"Curl encountered an error; retrying...\"\n else\n echo \"Service not ready yet (HTTP $response). Retrying in 5 seconds...\"\n fi\n sleep 5\n done\n echo \"Finished waiting for service.\"\n\n - name: Run Multi-Tenant Integration Tests\n env:\n ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}\n RUN_ID: ${{ github.run_id }}\n run: |\n echo \"Running multi-tenant integration tests...\"\n docker run --rm --network onyx_default \\\n --name test-runner \\\n -e POSTGRES_HOST=relational_db \\\n -e POSTGRES_USER=postgres \\\n -e POSTGRES_PASSWORD=password \\\n -e DB_READONLY_USER=db_readonly_user \\\n -e DB_READONLY_PASSWORD=password \\\n -e POSTGRES_DB=postgres \\\n -e POSTGRES_USE_NULL_POOL=true \\\n -e VESPA_HOST=index \\\n -e ENABLE_OPENSEARCH_INDEXING_FOR_ONYX=false \\\n -e REDIS_HOST=cache \\\n -e API_SERVER_HOST=api_server \\\n -e OPENAI_API_KEY=${OPENAI_API_KEY} \\\n -e EXA_API_KEY=${EXA_API_KEY} \\\n -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \\\n -e SLACK_BOT_TOKEN_TEST_SPACE=${SLACK_BOT_TOKEN_TEST_SPACE} \\\n -e TEST_WEB_HOSTNAME=test-runner \\\n -e AUTH_TYPE=cloud \\\n -e MULTI_TENANT=true \\\n -e SKIP_RESET=true \\\n -e REQUIRE_EMAIL_VERIFICATION=false \\\n -e DISABLE_TELEMETRY=true \\\n -e DEV_MODE=true \\\n ${ECR_CACHE}:integration-test-${RUN_ID} \\\n /app/tests/integration/multitenant_tests\n\n - name: Dump API server logs (multi-tenant)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.multitenant-dev.yml logs --no-color api_server > $GITHUB_WORKSPACE/api_server_multitenant.log || true\n\n - name: Dump all-container logs (multi-tenant)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.multitenant-dev.yml logs --no-color > $GITHUB_WORKSPACE/docker-compose-multitenant.log || true\n\n - name: Upload logs (multi-tenant)\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-all-logs-multitenant\n path: ${{ github.workspace }}/docker-compose-multitenant.log\n\n - name: Stop multi-tenant Docker containers\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.multitenant-dev.yml down -v\n\n required:\n # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.\n runs-on: ubuntu-slim\n timeout-minutes: 45\n needs: [integration-tests, onyx-lite-tests, multitenant-tests]\n if: ${{ always() }}\n steps:\n - name: Check job status\n if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') || contains(needs.*.result, 'skipped') }}\n run: exit 1\n" }, { "path": ".github/workflows/pr-jest-tests.yml", "content": "name: Run Jest Tests\nconcurrency:\n group: Run-Jest-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\njobs:\n jest-tests:\n name: Jest Tests\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup node\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v4\n with:\n node-version: 22\n cache: \"npm\" # zizmor: ignore[cache-poisoning] test-only workflow; no deploy artifacts\n cache-dependency-path: ./web/package-lock.json\n\n - name: Install node dependencies\n working-directory: ./web\n run: npm ci\n\n - name: Run Jest tests\n working-directory: ./web\n run: npm test -- --ci --coverage --maxWorkers=50%\n\n - name: Upload coverage reports\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: jest-coverage-${{ github.run_id }}\n path: ./web/coverage\n retention-days: 7\n" }, { "path": ".github/workflows/pr-labeler.yml", "content": "name: PR Labeler\n\non:\n pull_request:\n branches:\n - main\n types:\n - opened\n - reopened\n - synchronize\n - edited\n\npermissions:\n contents: read\n\njobs:\n validate_pr_title:\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Check PR title for Conventional Commits\n env:\n PR_TITLE: ${{ github.event.pull_request.title }}\n run: |\n echo \"PR Title: $PR_TITLE\"\n if [[ ! \"$PR_TITLE\" =~ ^(feat|fix|docs|test|ci|refactor|perf|chore|revert|build)(\\(.+\\))?:\\ .+ ]]; then\n echo \"::error::❌ Your PR title does not follow the Conventional Commits format.\n This check ensures that all pull requests use clear, consistent titles that help automate changelogs and improve project history.\n\n Please update your PR title to follow the Conventional Commits style.\n Here is a link to a blog explaining the reason why we've included the Conventional Commits style into our PR titles: https://xfuture-blog.com/working-with-conventional-commits\n\n **Here are some examples of valid PR titles:**\n - feat: add user authentication\n - fix(login): handle null password error\n - docs(readme): update installation instructions\"\n exit 1\n fi\n" }, { "path": ".github/workflows/pr-linear-check.yml", "content": "name: Ensure PR references Linear\nconcurrency:\n group: Ensure-PR-references-Linear-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n pull_request:\n types: [opened, edited, reopened, synchronize]\n\npermissions:\n contents: read\n\njobs:\n linear-check:\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - name: Check PR body for Linear link or override\n env:\n PR_BODY: ${{ github.event.pull_request.body }}\n run: |\n # Looking for \"https://linear.app\" in the body\n if echo \"$PR_BODY\" | grep -qE \"https://linear\\.app\"; then\n echo \"Found a Linear link. Check passed.\"\n exit 0\n fi\n\n # Looking for a checked override: \"[x] Override Linear Check\"\n if echo \"$PR_BODY\" | grep -q \"\\[x\\].*Override Linear Check\"; then\n echo \"Override box is checked. Check passed.\"\n exit 0\n fi\n\n # Otherwise, fail the run\n echo \"No Linear link or override found in the PR description.\"\n exit 1\n" }, { "path": ".github/workflows/pr-playwright-tests.yml", "content": "name: Run Playwright Tests\nconcurrency:\n group: Run-Playwright-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n push:\n tags:\n - \"v*.*.*\"\n # TODO: Remove this if we enable merge-queues for release branches.\n branches:\n - \"release/**\"\n\npermissions:\n contents: read\n\nenv:\n # Test Environment Variables\n OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}\n SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}\n GEN_AI_API_KEY: ${{ secrets.OPENAI_API_KEY }}\n EXA_API_KEY: ${{ secrets.EXA_API_KEY }}\n FIRECRAWL_API_KEY: ${{ secrets.FIRECRAWL_API_KEY }}\n GOOGLE_PSE_API_KEY: ${{ secrets.GOOGLE_PSE_API_KEY }}\n GOOGLE_PSE_SEARCH_ENGINE_ID: ${{ secrets.GOOGLE_PSE_SEARCH_ENGINE_ID }}\n\n # for federated slack tests\n SLACK_CLIENT_ID: ${{ secrets.SLACK_CLIENT_ID }}\n SLACK_CLIENT_SECRET: ${{ secrets.SLACK_CLIENT_SECRET }}\n\n # for MCP Oauth tests\n MCP_OAUTH_CLIENT_ID: ${{ secrets.MCP_OAUTH_CLIENT_ID }}\n MCP_OAUTH_CLIENT_SECRET: ${{ secrets.MCP_OAUTH_CLIENT_SECRET }}\n MCP_OAUTH_ISSUER: ${{ secrets.MCP_OAUTH_ISSUER }}\n MCP_OAUTH_JWKS_URI: ${{ secrets.MCP_OAUTH_JWKS_URI }}\n MCP_OAUTH_USERNAME: ${{ vars.MCP_OAUTH_USERNAME }}\n MCP_OAUTH_PASSWORD: ${{ secrets.MCP_OAUTH_PASSWORD }}\n\n # for MCP API Key tests\n MCP_API_KEY: test-api-key-12345\n MCP_API_KEY_TEST_PORT: 8005\n MCP_API_KEY_TEST_URL: http://host.docker.internal:8005/mcp\n MCP_API_KEY_SERVER_HOST: 0.0.0.0\n MCP_API_KEY_SERVER_PUBLIC_HOST: host.docker.internal\n\n MOCK_LLM_RESPONSE: true\n MCP_TEST_SERVER_PORT: 8004\n MCP_TEST_SERVER_URL: http://host.docker.internal:8004/mcp\n MCP_TEST_SERVER_PUBLIC_URL: http://host.docker.internal:8004/mcp\n MCP_TEST_SERVER_BIND_HOST: 0.0.0.0\n MCP_TEST_SERVER_PUBLIC_HOST: host.docker.internal\n MCP_SERVER_HOST: 0.0.0.0\n MCP_SERVER_PUBLIC_HOST: host.docker.internal\n MCP_SERVER_PUBLIC_URL: http://host.docker.internal:8004/mcp\n\n # Visual regression S3 bucket (shared across all jobs)\n PLAYWRIGHT_S3_BUCKET: onyx-playwright-artifacts\n\njobs:\n build-web-image:\n runs-on:\n [\n runs-on,\n runner=4cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-web-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n # needed for pulling external images otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Build and push Web Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./web\n file: ./web/Dockerfile\n platforms: linux/arm64\n tags: ${{ env.RUNS_ON_ECR_CACHE }}:playwright-test-web-${{ github.run_id }}\n push: true\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ github.event.pull_request.head.sha || github.sha }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache\n type=registry,ref=onyxdotapp/onyx-web-server:latest\n cache-to: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:web-cache,mode=max\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-backend-image:\n runs-on:\n [\n runs-on,\n runner=1cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-backend-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n # needed for pulling external images otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Build and push Backend Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile\n platforms: linux/arm64\n tags: ${{ env.RUNS_ON_ECR_CACHE }}:playwright-test-backend-${{ github.run_id }}\n push: true\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache\n type=registry,ref=onyxdotapp/onyx-backend:latest\n cache-to: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:backend-cache,mode=max\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n build-model-server-image:\n runs-on:\n [\n runs-on,\n runner=1cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-model-server-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n # needed for pulling external images otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Build and push Model Server Docker image\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend\n file: ./backend/Dockerfile.model_server\n platforms: linux/arm64\n tags: ${{ env.RUNS_ON_ECR_CACHE }}:playwright-test-model-server-${{ github.run_id }}\n push: true\n cache-from: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache\n type=registry,ref=onyxdotapp/onyx-model-server:latest\n cache-to: |\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max\n no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' }}\n\n playwright-tests:\n needs: [build-web-image, build-backend-image, build-model-server-image]\n name: Playwright Tests (${{ matrix.project }})\n permissions:\n id-token: write # Required for OIDC-based AWS credential exchange (S3 access)\n contents: read\n runs-on:\n - runs-on\n - runner=8cpu-linux-arm64\n - \"run-id=${{ github.run_id }}-playwright-tests-${{ matrix.project }}\"\n - \"extras=ecr-cache\"\n - volume=50gb\n timeout-minutes: 45\n strategy:\n fail-fast: false\n matrix:\n project: [admin, exclusive]\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup node\n # zizmor: ignore[cache-poisoning] ephemeral runners; no release artifacts\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v4\n with:\n node-version: 22\n cache: \"npm\" # zizmor: ignore[cache-poisoning]\n cache-dependency-path: ./web/package-lock.json\n\n - name: Install node dependencies\n working-directory: ./web\n run: npm ci\n\n - name: Cache playwright cache\n # zizmor: ignore[cache-poisoning] ephemeral runners; no release artifacts\n uses: runs-on/cache@a5f51d6f3fece787d03b7b4e981c82538a0654ed # ratchet:runs-on/cache@v4\n with:\n path: ~/.cache/ms-playwright\n key: ${{ runner.os }}-playwright-npm-${{ hashFiles('web/package-lock.json') }}\n restore-keys: |\n ${{ runner.os }}-playwright-npm-\n\n - name: Install playwright browsers\n working-directory: ./web\n run: npx playwright install --with-deps\n\n - name: Create .env file for Docker Compose\n env:\n OPENAI_API_KEY_VALUE: ${{ env.OPENAI_API_KEY }}\n EXA_API_KEY_VALUE: ${{ env.EXA_API_KEY }}\n ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}\n RUN_ID: ${{ github.run_id }}\n run: |\n cat < deployment/docker_compose/.env\n COMPOSE_PROFILES=s3-filestore\n ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true\n # TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\n LICENSE_ENFORCEMENT_ENABLED=false\n AUTH_TYPE=basic\n INTEGRATION_TESTS_MODE=true\n GEN_AI_API_KEY=${OPENAI_API_KEY_VALUE}\n EXA_API_KEY=${EXA_API_KEY_VALUE}\n REQUIRE_EMAIL_VERIFICATION=false\n DISABLE_TELEMETRY=true\n ONYX_BACKEND_IMAGE=${ECR_CACHE}:playwright-test-backend-${RUN_ID}\n ONYX_MODEL_SERVER_IMAGE=${ECR_CACHE}:playwright-test-model-server-${RUN_ID}\n ONYX_WEB_SERVER_IMAGE=${ECR_CACHE}:playwright-test-web-${RUN_ID}\n EOF\n\n # needed for pulling Vespa, Redis, Postgres, and Minio images\n # otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Start Docker containers\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.dev.yml -f docker-compose.mcp-oauth-test.yml -f docker-compose.mcp-api-key-test.yml up -d\n id: start_docker\n\n - name: Wait for service to be ready\n run: |\n echo \"Starting wait-for-service script...\"\n\n docker logs -f onyx-api_server-1 &\n\n start_time=$(date +%s)\n timeout=300 # 5 minutes in seconds\n\n while true; do\n current_time=$(date +%s)\n elapsed_time=$((current_time - start_time))\n\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. Service did not become ready in 5 minutes.\"\n exit 1\n fi\n\n # Use curl with error handling to ignore specific exit code 56\n response=$(curl -s -o /dev/null -w \"%{http_code}\" http://localhost:8080/health || echo \"curl_error\")\n\n if [ \"$response\" = \"200\" ]; then\n echo \"Service is ready!\"\n break\n elif [ \"$response\" = \"curl_error\" ]; then\n echo \"Curl encountered an error, possibly exit code 56. Continuing to retry...\"\n else\n echo \"Service not ready yet (HTTP status $response). Retrying in 5 seconds...\"\n fi\n\n sleep 5\n done\n echo \"Finished waiting for service.\"\n\n - name: Wait for MCP OAuth mock server\n run: |\n echo \"Waiting for MCP OAuth mock server on port ${MCP_TEST_SERVER_PORT:-8004}...\"\n start_time=$(date +%s)\n timeout=120\n\n while true; do\n current_time=$(date +%s)\n elapsed_time=$((current_time - start_time))\n\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. MCP OAuth mock server did not become ready in ${timeout}s.\"\n exit 1\n fi\n\n if curl -sf \"http://localhost:${MCP_TEST_SERVER_PORT:-8004}/healthz\" > /dev/null; then\n echo \"MCP OAuth mock server is ready!\"\n break\n fi\n\n sleep 3\n done\n\n - name: Wait for MCP API Key mock server\n run: |\n echo \"Waiting for MCP API Key mock server on port ${MCP_API_KEY_TEST_PORT:-8005}...\"\n start_time=$(date +%s)\n timeout=120\n\n while true; do\n current_time=$(date +%s)\n elapsed_time=$((current_time - start_time))\n\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. MCP API Key mock server did not become ready in ${timeout}s.\"\n exit 1\n fi\n\n if curl -sf \"http://localhost:${MCP_API_KEY_TEST_PORT:-8005}/healthz\" > /dev/null; then\n echo \"MCP API Key mock server is ready!\"\n break\n fi\n\n sleep 3\n done\n\n - name: Wait for web server to be ready\n run: |\n echo \"Waiting for web server on port 3000...\"\n start_time=$(date +%s)\n timeout=120\n\n while true; do\n current_time=$(date +%s)\n elapsed_time=$((current_time - start_time))\n\n if [ $elapsed_time -ge $timeout ]; then\n echo \"Timeout reached. Web server did not become ready in ${timeout}s.\"\n exit 1\n fi\n\n if curl -sf \"http://localhost:3000/api/health\" > /dev/null 2>&1 || \\\n curl -sf \"http://localhost:3000/\" > /dev/null 2>&1; then\n echo \"Web server is ready!\"\n break\n fi\n\n echo \"Web server not ready yet. Retrying in 3 seconds...\"\n sleep 3\n done\n\n - name: Run Playwright tests\n working-directory: ./web\n env:\n PROJECT: ${{ matrix.project }}\n run: |\n npx playwright test --project ${PROJECT}\n\n - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n if: always()\n with:\n # Includes test results and trace.zip files\n name: playwright-test-results-${{ matrix.project }}-${{ github.run_id }}\n path: ./web/output/playwright/\n retention-days: 30\n\n - name: Upload screenshots\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n if: always()\n with:\n name: playwright-screenshots-${{ matrix.project }}-${{ github.run_id }}\n path: ./web/output/screenshots/\n retention-days: 30\n\n # --- Visual Regression Diff ---\n - name: Configure AWS credentials\n if: always()\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Install the latest version of uv\n if: always()\n uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n enable-cache: false\n version: \"0.9.9\"\n\n - name: Determine baseline revision\n if: always()\n id: baseline-rev\n env:\n EVENT_NAME: ${{ github.event_name }}\n BASE_REF: ${{ github.event.pull_request.base.ref }}\n MERGE_GROUP_BASE_REF: ${{ github.event.merge_group.base_ref }}\n GH_REF: ${{ github.ref }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ \"${EVENT_NAME}\" = \"pull_request\" ]; then\n # PRs compare against the base branch (e.g. main, release/2.5)\n echo \"rev=${BASE_REF}\" >> \"$GITHUB_OUTPUT\"\n elif [ \"${EVENT_NAME}\" = \"merge_group\" ]; then\n # Merge queue compares against the target branch (e.g. refs/heads/main -> main)\n echo \"rev=${MERGE_GROUP_BASE_REF#refs/heads/}\" >> \"$GITHUB_OUTPUT\"\n elif [[ \"${GH_REF}\" == refs/tags/* ]]; then\n # Tag builds compare against the tag name\n echo \"rev=${REF_NAME}\" >> \"$GITHUB_OUTPUT\"\n else\n # Push builds (main, release/*) compare against the branch name\n echo \"rev=${REF_NAME}\" >> \"$GITHUB_OUTPUT\"\n fi\n\n - name: Generate screenshot diff report\n if: always()\n env:\n PROJECT: ${{ matrix.project }}\n PLAYWRIGHT_S3_BUCKET: ${{ env.PLAYWRIGHT_S3_BUCKET }}\n BASELINE_REV: ${{ steps.baseline-rev.outputs.rev }}\n run: |\n uv run --no-sync --with onyx-devtools ods screenshot-diff compare \\\n --project \"${PROJECT}\" \\\n --rev \"${BASELINE_REV}\"\n\n - name: Upload visual diff report to S3\n if: always()\n env:\n PROJECT: ${{ matrix.project }}\n PR_NUMBER: ${{ github.event.pull_request.number }}\n RUN_ID: ${{ github.run_id }}\n run: |\n SUMMARY_FILE=\"web/output/screenshot-diff/${PROJECT}/summary.json\"\n if [ ! -f \"${SUMMARY_FILE}\" ]; then\n echo \"No summary file found — skipping S3 upload.\"\n exit 0\n fi\n\n HAS_DIFF=$(jq -r '.has_differences' \"${SUMMARY_FILE}\")\n if [ \"${HAS_DIFF}\" != \"true\" ]; then\n echo \"No visual differences for ${PROJECT} — skipping S3 upload.\"\n exit 0\n fi\n\n aws s3 sync \"web/output/screenshot-diff/${PROJECT}/\" \\\n \"s3://${PLAYWRIGHT_S3_BUCKET}/reports/pr-${PR_NUMBER}/${RUN_ID}/${PROJECT}/\"\n\n - name: Upload visual diff summary\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n if: always()\n with:\n name: screenshot-diff-summary-${{ matrix.project }}\n path: ./web/output/screenshot-diff/${{ matrix.project }}/summary.json\n if-no-files-found: ignore\n retention-days: 5\n\n - name: Upload visual diff report artifact\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n if: always()\n with:\n name: screenshot-diff-report-${{ matrix.project }}-${{ github.run_id }}\n path: ./web/output/screenshot-diff/${{ matrix.project }}/\n if-no-files-found: ignore\n retention-days: 30\n\n - name: Update S3 baselines\n if: >-\n success() && (\n github.ref == 'refs/heads/main' ||\n startsWith(github.ref, 'refs/heads/release/') ||\n startsWith(github.ref, 'refs/tags/v') ||\n (\n github.event_name == 'merge_group' && (\n github.event.merge_group.base_ref == 'refs/heads/main' ||\n startsWith(github.event.merge_group.base_ref, 'refs/heads/release/')\n )\n )\n )\n env:\n PROJECT: ${{ matrix.project }}\n PLAYWRIGHT_S3_BUCKET: ${{ env.PLAYWRIGHT_S3_BUCKET }}\n BASELINE_REV: ${{ steps.baseline-rev.outputs.rev }}\n run: |\n if [ -d \"web/output/screenshots/\" ] && [ \"$(ls -A web/output/screenshots/)\" ]; then\n uv run --no-sync --with onyx-devtools ods screenshot-diff upload-baselines \\\n --project \"${PROJECT}\" \\\n --rev \"${BASELINE_REV}\" \\\n --delete\n else\n echo \"No screenshots to upload for ${PROJECT} — skipping baseline update.\"\n fi\n\n # save before stopping the containers so the logs can be captured\n - name: Save Docker logs\n if: success() || failure()\n env:\n WORKSPACE: ${{ github.workspace }}\n run: |\n cd deployment/docker_compose\n docker compose logs > docker-compose.log\n mv docker-compose.log ${WORKSPACE}/docker-compose.log\n\n - name: Upload logs\n if: success() || failure()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-logs-${{ matrix.project }}-${{ github.run_id }}\n path: ${{ github.workspace }}/docker-compose.log\n\n playwright-tests-lite:\n needs: [build-web-image, build-backend-image]\n name: Playwright Tests (lite)\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - \"run-id=${{ github.run_id }}-playwright-tests-lite\"\n - \"extras=ecr-cache\"\n timeout-minutes: 30\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup node\n # zizmor: ignore[cache-poisoning] ephemeral runners; no release artifacts\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v4\n with:\n node-version: 22\n cache: \"npm\" # zizmor: ignore[cache-poisoning]\n cache-dependency-path: ./web/package-lock.json\n\n - name: Install node dependencies\n working-directory: ./web\n run: npm ci\n\n - name: Cache playwright cache\n # zizmor: ignore[cache-poisoning] ephemeral runners; no release artifacts\n uses: runs-on/cache@a5f51d6f3fece787d03b7b4e981c82538a0654ed # ratchet:runs-on/cache@v4\n with:\n path: ~/.cache/ms-playwright\n key: ${{ runner.os }}-playwright-npm-${{ hashFiles('web/package-lock.json') }}\n restore-keys: |\n ${{ runner.os }}-playwright-npm-\n\n - name: Install playwright browsers\n working-directory: ./web\n run: npx playwright install --with-deps\n\n - name: Create .env file for Docker Compose\n env:\n OPENAI_API_KEY_VALUE: ${{ env.OPENAI_API_KEY }}\n ECR_CACHE: ${{ env.RUNS_ON_ECR_CACHE }}\n RUN_ID: ${{ github.run_id }}\n run: |\n cat < deployment/docker_compose/.env\n ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true\n LICENSE_ENFORCEMENT_ENABLED=false\n AUTH_TYPE=basic\n INTEGRATION_TESTS_MODE=true\n GEN_AI_API_KEY=${OPENAI_API_KEY_VALUE}\n MOCK_LLM_RESPONSE=true\n REQUIRE_EMAIL_VERIFICATION=false\n DISABLE_TELEMETRY=true\n ONYX_BACKEND_IMAGE=${ECR_CACHE}:playwright-test-backend-${RUN_ID}\n ONYX_WEB_SERVER_IMAGE=${ECR_CACHE}:playwright-test-web-${RUN_ID}\n EOF\n\n # needed for pulling external images otherwise, we hit the \"Unauthenticated users\" limit\n # https://docs.docker.com/docker-hub/usage/\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Start Docker containers (lite)\n run: |\n cd deployment/docker_compose\n docker compose -f docker-compose.yml -f docker-compose.onyx-lite.yml -f docker-compose.dev.yml up -d\n id: start_docker\n\n - name: Run Playwright tests (lite)\n working-directory: ./web\n run: npx playwright test --project lite\n\n - uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n if: always()\n with:\n name: playwright-test-results-lite-${{ github.run_id }}\n path: ./web/output/playwright/\n retention-days: 30\n\n - name: Save Docker logs\n if: success() || failure()\n env:\n WORKSPACE: ${{ github.workspace }}\n run: |\n cd deployment/docker_compose\n docker compose logs > docker-compose.log\n mv docker-compose.log ${WORKSPACE}/docker-compose.log\n\n - name: Upload logs\n if: success() || failure()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-logs-lite-${{ github.run_id }}\n path: ${{ github.workspace }}/docker-compose.log\n\n # Post a single combined visual regression comment after all matrix jobs finish\n visual-regression-comment:\n needs: [playwright-tests]\n if: >-\n always() &&\n github.event_name == 'pull_request' &&\n needs.playwright-tests.result != 'cancelled'\n runs-on: ubuntu-slim\n timeout-minutes: 5\n permissions:\n pull-requests: write\n steps:\n - name: Download visual diff summaries\n uses: actions/download-artifact@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3\n with:\n pattern: screenshot-diff-summary-*\n path: summaries/\n\n - name: Post combined PR comment\n env:\n GH_TOKEN: ${{ github.token }}\n PR_NUMBER: ${{ github.event.pull_request.number }}\n RUN_ID: ${{ github.run_id }}\n REPO: ${{ github.repository }}\n S3_BUCKET: ${{ env.PLAYWRIGHT_S3_BUCKET }}\n run: |\n MARKER=\"\"\n\n # Build the markdown table from all summary files\n TABLE_HEADER=\"| Project | Changed | Added | Removed | Unchanged | Report |\"\n TABLE_DIVIDER=\"|---------|---------|-------|---------|-----------|--------|\"\n TABLE_ROWS=\"\"\n HAS_ANY_SUMMARY=false\n\n for SUMMARY_DIR in summaries/screenshot-diff-summary-*/; do\n SUMMARY_FILE=\"${SUMMARY_DIR}summary.json\"\n if [ ! -f \"${SUMMARY_FILE}\" ]; then\n continue\n fi\n\n HAS_ANY_SUMMARY=true\n PROJECT=$(jq -r '.project' \"${SUMMARY_FILE}\")\n CHANGED=$(jq -r '.changed' \"${SUMMARY_FILE}\")\n ADDED=$(jq -r '.added' \"${SUMMARY_FILE}\")\n REMOVED=$(jq -r '.removed' \"${SUMMARY_FILE}\")\n UNCHANGED=$(jq -r '.unchanged' \"${SUMMARY_FILE}\")\n TOTAL=$(jq -r '.total' \"${SUMMARY_FILE}\")\n HAS_DIFF=$(jq -r '.has_differences' \"${SUMMARY_FILE}\")\n\n if [ \"${TOTAL}\" = \"0\" ]; then\n REPORT_LINK=\"_No screenshots_\"\n elif [ \"${HAS_DIFF}\" = \"true\" ]; then\n REPORT_URL=\"https://${S3_BUCKET}.s3.us-east-2.amazonaws.com/reports/pr-${PR_NUMBER}/${RUN_ID}/${PROJECT}/index.html\"\n REPORT_LINK=\"[View Report](${REPORT_URL})\"\n else\n REPORT_LINK=\"✅ No changes\"\n fi\n\n TABLE_ROWS=\"${TABLE_ROWS}| \\`${PROJECT}\\` | ${CHANGED} | ${ADDED} | ${REMOVED} | ${UNCHANGED} | ${REPORT_LINK} |\\n\"\n done\n\n if [ \"${HAS_ANY_SUMMARY}\" = \"false\" ]; then\n echo \"No visual diff summaries found — skipping PR comment.\"\n exit 0\n fi\n\n BODY=$(printf '%s\\n' \\\n \"${MARKER}\" \\\n \"### 🖼️ Visual Regression Report\" \\\n \"\" \\\n \"${TABLE_HEADER}\" \\\n \"${TABLE_DIVIDER}\" \\\n \"$(printf '%b' \"${TABLE_ROWS}\")\")\n\n # Upsert: find existing comment with the marker, or create a new one\n EXISTING_COMMENT_ID=$(gh api \\\n \"repos/${REPO}/issues/${PR_NUMBER}/comments\" \\\n --jq \".[] | select(.body | startswith(\\\"${MARKER}\\\")) | .id\" \\\n 2>/dev/null | head -1)\n\n if [ -n \"${EXISTING_COMMENT_ID}\" ]; then\n gh api \\\n --method PATCH \\\n \"repos/${REPO}/issues/comments/${EXISTING_COMMENT_ID}\" \\\n -f body=\"${BODY}\"\n else\n gh api \\\n --method POST \\\n \"repos/${REPO}/issues/${PR_NUMBER}/comments\" \\\n -f body=\"${BODY}\"\n fi\n\n playwright-required:\n # NOTE: Github-hosted runners have about 20s faster queue times and are preferred here.\n runs-on: ubuntu-slim\n timeout-minutes: 45\n needs: [playwright-tests, playwright-tests-lite]\n if: ${{ always() }}\n steps:\n - name: Check job status\n if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') || contains(needs.*.result, 'skipped') }}\n run: exit 1\n" }, { "path": ".github/workflows/pr-python-checks.yml", "content": "name: Python Checks\nconcurrency:\n group: Python-Checks-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - \"release/**\"\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\njobs:\n mypy-check:\n # See https://runs-on.com/runners/linux/\n # Note: Mypy seems quite optimized for x64 compared to arm64.\n # Similarly, mypy is single-threaded and incremental, so 2cpu is sufficient.\n runs-on:\n [\n runs-on,\n runner=2cpu-linux-x64,\n \"run-id=${{ github.run_id }}-mypy-check\",\n \"extras=s3-cache\",\n ]\n timeout-minutes: 45\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup Python and Install Dependencies\n uses: ./.github/actions/setup-python-and-install-dependencies\n with:\n requirements: |\n backend/requirements/default.txt\n backend/requirements/dev.txt\n backend/requirements/model_server.txt\n backend/requirements/ee.txt\n\n - name: Generate OpenAPI schema and Python client\n shell: bash\n # TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\n env:\n LICENSE_ENFORCEMENT_ENABLED: \"false\"\n run: |\n ods openapi all\n\n - name: Cache mypy cache\n if: ${{ vars.DISABLE_MYPY_CACHE != 'true' }}\n uses: runs-on/cache@a5f51d6f3fece787d03b7b4e981c82538a0654ed # ratchet:runs-on/cache@v4\n with:\n path: .mypy_cache\n key: mypy-${{ runner.os }}-${{ github.base_ref || github.event.merge_group.base_ref || 'main' }}-${{ hashFiles('**/*.py', '**/*.pyi', 'pyproject.toml') }}\n restore-keys: |\n mypy-${{ runner.os }}-${{ github.base_ref || github.event.merge_group.base_ref || 'main' }}-\n mypy-${{ runner.os }}-\n\n - name: Run MyPy\n env:\n MYPY_FORCE_COLOR: 1\n TERM: xterm-256color\n run: mypy .\n" }, { "path": ".github/workflows/pr-python-connector-tests.yml", "content": "name: Connector Tests\nconcurrency:\n group: Connector-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches: [main]\n paths:\n - \"backend/**\"\n - \"pyproject.toml\"\n - \"uv.lock\"\n - \".github/workflows/pr-python-connector-tests.yml\"\n - \".github/actions/setup-python-and-install-dependencies/**\"\n - \".github/actions/setup-playwright/**\"\n push:\n tags:\n - \"v*.*.*\"\n schedule:\n # This cron expression runs the job daily at 16:00 UTC (9am PT)\n - cron: \"0 16 * * *\"\n\npermissions:\n id-token: write # Required for OIDC-based AWS credential exchange\n contents: read\n\nenv:\n PYTHONPATH: ./backend\n DISABLE_TELEMETRY: \"true\"\n R2_ACCOUNT_ID_DAILY_CONNECTOR_TESTS: ${{ vars.R2_ACCOUNT_ID_DAILY_CONNECTOR_TESTS }}\n CONFLUENCE_TEST_SPACE_URL: ${{ vars.CONFLUENCE_TEST_SPACE_URL }}\n CONFLUENCE_TEST_SPACE: ${{ vars.CONFLUENCE_TEST_SPACE }}\n CONFLUENCE_USER_NAME: ${{ vars.CONFLUENCE_USER_NAME }}\n SF_USERNAME: ${{ vars.SF_USERNAME }}\n IMAP_HOST: ${{ vars.IMAP_HOST }}\n IMAP_USERNAME: ${{ vars.IMAP_USERNAME }}\n IMAP_MAILBOXES: ${{ vars.IMAP_MAILBOXES }}\n AIRTABLE_TEST_BASE_ID: ${{ vars.AIRTABLE_TEST_BASE_ID }}\n AIRTABLE_TEST_TABLE_ID: ${{ vars.AIRTABLE_TEST_TABLE_ID }}\n AIRTABLE_TEST_TABLE_NAME: ${{ vars.AIRTABLE_TEST_TABLE_NAME }}\n SHAREPOINT_CLIENT_ID: ${{ vars.SHAREPOINT_CLIENT_ID }}\n SHAREPOINT_CLIENT_DIRECTORY_ID: ${{ vars.SHAREPOINT_CLIENT_DIRECTORY_ID }}\n SHAREPOINT_SITE: ${{ vars.SHAREPOINT_SITE }}\n BITBUCKET_EMAIL: ${{ vars.BITBUCKET_EMAIL }}\n\njobs:\n connectors-check:\n # See https://runs-on.com/runners/linux/\n runs-on:\n [\n runs-on,\n runner=8cpu-linux-x64,\n \"run-id=${{ github.run_id }}-connectors-check\",\n \"extras=s3-cache\",\n ]\n timeout-minutes: 45\n environment: ci-protected\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup Python and Install Dependencies\n uses: ./.github/actions/setup-python-and-install-dependencies\n with:\n requirements: |\n backend/requirements/default.txt\n backend/requirements/dev.txt\n\n - name: Setup Playwright\n uses: ./.github/actions/setup-playwright\n\n - name: Detect Connector changes\n id: changes\n uses: dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36 # ratchet:dorny/paths-filter@v3\n with:\n filters: |\n hubspot:\n - 'backend/onyx/connectors/hubspot/**'\n - 'backend/tests/daily/connectors/hubspot/**'\n - 'uv.lock'\n salesforce:\n - 'backend/onyx/connectors/salesforce/**'\n - 'backend/tests/daily/connectors/salesforce/**'\n - 'uv.lock'\n github:\n - 'backend/onyx/connectors/github/**'\n - 'backend/tests/daily/connectors/github/**'\n - 'uv.lock'\n file_processing:\n - 'backend/onyx/file_processing/**'\n - 'uv.lock'\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v4\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get connector test secrets from AWS Secrets Manager\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2\n with:\n parse-json-secrets: false\n secret-ids: |\n AWS_ACCESS_KEY_ID_DAILY_CONNECTOR_TESTS, test/aws-access-key-id\n AWS_SECRET_ACCESS_KEY_DAILY_CONNECTOR_TESTS, test/aws-secret-access-key\n R2_ACCESS_KEY_ID_DAILY_CONNECTOR_TESTS, test/r2-access-key-id\n R2_SECRET_ACCESS_KEY_DAILY_CONNECTOR_TESTS, test/r2-secret-access-key\n GCS_ACCESS_KEY_ID_DAILY_CONNECTOR_TESTS, test/gcs-access-key-id\n GCS_SECRET_ACCESS_KEY_DAILY_CONNECTOR_TESTS, test/gcs-secret-access-key\n CONFLUENCE_ACCESS_TOKEN, test/confluence-access-token\n CONFLUENCE_ACCESS_TOKEN_SCOPED, test/confluence-access-token-scoped\n JIRA_BASE_URL, test/jira-base-url\n JIRA_USER_EMAIL, test/jira-user-email\n JIRA_API_TOKEN, test/jira-api-token\n JIRA_API_TOKEN_SCOPED, test/jira-api-token-scoped\n GONG_ACCESS_KEY, test/gong-access-key\n GONG_ACCESS_KEY_SECRET, test/gong-access-key-secret\n GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR, test/google-drive-service-account-json\n GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR_TEST_USER_1, test/google-drive-oauth-creds-test-user-1\n GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR, test/google-drive-oauth-creds\n GOOGLE_GMAIL_SERVICE_ACCOUNT_JSON_STR, test/google-gmail-service-account-json\n GOOGLE_GMAIL_OAUTH_CREDENTIALS_JSON_STR, test/google-gmail-oauth-creds\n SLAB_BOT_TOKEN, test/slab-bot-token\n ZENDESK_SUBDOMAIN, test/zendesk-subdomain\n ZENDESK_EMAIL, test/zendesk-email\n ZENDESK_TOKEN, test/zendesk-token\n SF_PASSWORD, test/sf-password\n SF_SECURITY_TOKEN, test/sf-security-token\n HUBSPOT_ACCESS_TOKEN, test/hubspot-access-token\n IMAP_PASSWORD, test/imap-password\n AIRTABLE_ACCESS_TOKEN, test/airtable-access-token\n SHAREPOINT_CLIENT_SECRET, test/sharepoint-client-secret\n PERM_SYNC_SHAREPOINT_CLIENT_ID, test/perm-sync-sharepoint-client-id\n PERM_SYNC_SHAREPOINT_PRIVATE_KEY, test/perm-sync-sharepoint-private-key\n PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD, test/perm-sync-sharepoint-cert-password\n PERM_SYNC_SHAREPOINT_DIRECTORY_ID, test/perm-sync-sharepoint-directory-id\n ACCESS_TOKEN_GITHUB, test/github-access-token\n GITLAB_ACCESS_TOKEN, test/gitlab-access-token\n GITBOOK_SPACE_ID, test/gitbook-space-id\n GITBOOK_API_KEY, test/gitbook-api-key\n NOTION_INTEGRATION_TOKEN, test/notion-integration-token\n HIGHSPOT_KEY, test/highspot-key\n HIGHSPOT_SECRET, test/highspot-secret\n SLACK_BOT_TOKEN, test/slack-bot-token\n DISCORD_CONNECTOR_BOT_TOKEN, test/discord-bot-token\n TEAMS_APPLICATION_ID, test/teams-application-id\n TEAMS_DIRECTORY_ID, test/teams-directory-id\n TEAMS_SECRET, test/teams-secret\n BITBUCKET_WORKSPACE, test/bitbucket-workspace\n BITBUCKET_API_TOKEN, test/bitbucket-api-token\n FIREFLIES_API_KEY, test/fireflies-api-key\n\n - name: Run Tests (excluding HubSpot, Salesforce, GitHub, and Coda)\n shell: script -q -e -c \"bash --noprofile --norc -eo pipefail {0}\"\n run: |\n py.test \\\n -n 8 \\\n --dist loadfile \\\n --durations=8 \\\n -o junit_family=xunit2 \\\n -xv \\\n --ff \\\n backend/tests/daily/connectors \\\n --ignore backend/tests/daily/connectors/hubspot \\\n --ignore backend/tests/daily/connectors/salesforce \\\n --ignore backend/tests/daily/connectors/github \\\n --ignore backend/tests/daily/connectors/coda\n\n - name: Run HubSpot Connector Tests\n if: ${{ github.event_name == 'schedule' || steps.changes.outputs.hubspot == 'true' || steps.changes.outputs.file_processing == 'true' }}\n shell: script -q -e -c \"bash --noprofile --norc -eo pipefail {0}\"\n run: |\n py.test \\\n -n 8 \\\n --dist loadfile \\\n --durations=8 \\\n -o junit_family=xunit2 \\\n -xv \\\n --ff \\\n backend/tests/daily/connectors/hubspot\n\n - name: Run Salesforce Connector Tests\n if: ${{ github.event_name == 'schedule' || steps.changes.outputs.salesforce == 'true' || steps.changes.outputs.file_processing == 'true' }}\n shell: script -q -e -c \"bash --noprofile --norc -eo pipefail {0}\"\n run: |\n py.test \\\n -n 8 \\\n --dist loadfile \\\n --durations=8 \\\n -o junit_family=xunit2 \\\n -xv \\\n --ff \\\n backend/tests/daily/connectors/salesforce\n\n - name: Run GitHub Connector Tests\n if: ${{ github.event_name == 'schedule' || steps.changes.outputs.github == 'true' || steps.changes.outputs.file_processing == 'true' }}\n shell: script -q -e -c \"bash --noprofile --norc -eo pipefail {0}\"\n run: |\n py.test \\\n -n 8 \\\n --dist loadfile \\\n --durations=8 \\\n -o junit_family=xunit2 \\\n -xv \\\n --ff \\\n backend/tests/daily/connectors/github\n\n - name: Alert on Failure\n if: failure() && github.event_name == 'schedule'\n env:\n SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }}\n REPO: ${{ github.repository }}\n RUN_ID: ${{ github.run_id }}\n run: |\n curl -X POST \\\n -H 'Content-type: application/json' \\\n --data \"{\\\"text\\\":\\\"Scheduled Connector Tests failed! Check the run at: https://github.com/${REPO}/actions/runs/${RUN_ID}\\\"}\" \\\n $SLACK_WEBHOOK\n" }, { "path": ".github/workflows/pr-python-model-tests.yml", "content": "name: Model Server Tests\n\non:\n schedule:\n # This cron expression runs the job daily at 16:00 UTC (9am PT)\n - cron: \"0 16 * * *\"\n workflow_dispatch:\n\npermissions:\n contents: read\n\nenv:\n # Bedrock\n AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}\n AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}\n AWS_REGION_NAME: ${{ vars.AWS_REGION_NAME }}\n\n # API keys for testing\n COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}\n LITELLM_API_KEY: ${{ secrets.LITELLM_API_KEY }}\n LITELLM_API_URL: ${{ secrets.LITELLM_API_URL }}\n OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}\n AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}\n AZURE_API_URL: ${{ vars.AZURE_API_URL }}\n\njobs:\n model-check:\n # See https://runs-on.com/runners/linux/\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - \"run-id=${{ github.run_id }}-model-check\"\n - \"extras=ecr-cache\"\n environment: ci-protected\n timeout-minutes: 45\n\n env:\n PYTHONPATH: ./backend\n\n steps:\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup Python and Install Dependencies\n uses: ./.github/actions/setup-python-and-install-dependencies\n with:\n requirements: |\n backend/requirements/default.txt\n backend/requirements/dev.txt\n\n - name: Format branch name for cache\n id: format-branch\n env:\n PR_NUMBER: ${{ github.event.pull_request.number }}\n REF_NAME: ${{ github.ref_name }}\n run: |\n if [ -n \"${PR_NUMBER}\" ]; then\n CACHE_SUFFIX=\"${PR_NUMBER}\"\n else\n # shellcheck disable=SC2001\n CACHE_SUFFIX=$(echo \"${REF_NAME}\" | sed 's/[^A-Za-z0-9._-]/-/g')\n fi\n echo \"cache-suffix=${CACHE_SUFFIX}\" >> $GITHUB_OUTPUT\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9\n with:\n username: ${{ secrets.DOCKER_USERNAME }}\n password: ${{ secrets.DOCKER_TOKEN }}\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f\n\n - name: Build and load\n uses: docker/bake-action@82490499d2e5613fcead7e128237ef0b0ea210f7 # ratchet:docker/bake-action@v7.0.0\n env:\n TAG: model-server-${{ github.run_id }}\n with:\n load: true\n targets: model-server\n set: |\n model-server.cache-from=type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }}\n model-server.cache-from=type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }}\n model-server.cache-from=type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache\n model-server.cache-from=type=registry,ref=onyxdotapp/onyx-model-server:latest\n model-server.cache-to=type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ github.event.pull_request.head.sha || github.sha }},mode=max\n model-server.cache-to=type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache-${{ steps.format-branch.outputs.cache-suffix }},mode=max\n model-server.cache-to=type=registry,ref=${{ env.RUNS_ON_ECR_CACHE }}:model-server-cache,mode=max\n\n - name: Start Docker containers\n id: start_docker\n env:\n IMAGE_TAG: model-server-${{ github.run_id }}\n run: |\n cd deployment/docker_compose\n docker compose \\\n -f docker-compose.yml \\\n -f docker-compose.dev.yml \\\n up -d --wait \\\n inference_model_server\n\n - name: Run Tests\n run: |\n py.test -o junit_family=xunit2 -xv --ff backend/tests/daily/llm\n py.test -o junit_family=xunit2 -xv --ff backend/tests/daily/embedding\n\n - name: Alert on Failure\n if: failure() && github.event_name == 'schedule'\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.SLACK_WEBHOOK }}\n failed-jobs: model-check\n title: \"🚨 Scheduled Model Tests failed!\"\n ref-name: ${{ github.ref_name }}\n\n - name: Dump all-container logs (optional)\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose logs --no-color > $GITHUB_WORKSPACE/docker-compose.log || true\n\n - name: Upload logs\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-all-logs\n path: ${{ github.workspace }}/docker-compose.log\n" }, { "path": ".github/workflows/pr-python-tests.yml", "content": "name: Python Unit Tests\nconcurrency:\n group: Python-Unit-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request:\n branches:\n - main\n - 'release/**'\n push:\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\njobs:\n backend-check:\n # See https://runs-on.com/runners/linux/\n runs-on: [runs-on, runner=2cpu-linux-arm64, \"run-id=${{ github.run_id }}-backend-check\"]\n timeout-minutes: 45\n\n\n env:\n PYTHONPATH: ./backend\n REDIS_CLOUD_PYTEST_PASSWORD: ${{ secrets.REDIS_CLOUD_PYTEST_PASSWORD }}\n DISABLE_TELEMETRY: \"true\"\n # TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\n LICENSE_ENFORCEMENT_ENABLED: \"false\"\n\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Setup Python and Install Dependencies\n uses: ./.github/actions/setup-python-and-install-dependencies\n with:\n requirements: |\n backend/requirements/default.txt\n backend/requirements/dev.txt\n backend/requirements/model_server.txt\n backend/requirements/ee.txt\n\n - name: Run Tests\n shell: script -q -e -c \"bash --noprofile --norc -eo pipefail {0}\"\n run: py.test -o junit_family=xunit2 -xv --ff backend/tests/unit\n" }, { "path": ".github/workflows/pr-quality-checks.yml", "content": "name: Quality Checks PR\nconcurrency:\n group: Quality-Checks-PR-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}\n cancel-in-progress: true\n\non:\n merge_group:\n pull_request: null\n push:\n branches:\n - main\n tags:\n - \"v*.*.*\"\n\npermissions:\n contents: read\n\njobs:\n quality-checks:\n runs-on: ubuntu-latest\n timeout-minutes: 45\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n fetch-depth: 0\n persist-credentials: false\n - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # ratchet:actions/setup-python@v6\n with:\n python-version: \"3.11\"\n - name: Setup Terraform\n uses: hashicorp/setup-terraform@5e8dbf3c6d9deaf4193ca7a8fb23f2ac83bb6c85 # ratchet:hashicorp/setup-terraform@v4.0.0\n - name: Setup node\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v6\n with: # zizmor: ignore[cache-poisoning]\n node-version: 22\n cache: \"npm\"\n cache-dependency-path: ./web/package-lock.json\n - name: Install node dependencies\n working-directory: ./web\n run: npm ci\n - uses: j178/prek-action@0bb87d7f00b0c99306c8bcb8b8beba1eb581c037 # ratchet:j178/prek-action@v1\n with:\n prek-version: '0.3.4'\n extra-args: ${{ github.event_name == 'pull_request' && format('--from-ref {0} --to-ref {1}', github.event.pull_request.base.sha, github.event.pull_request.head.sha) || github.event_name == 'merge_group' && format('--from-ref {0} --to-ref {1}', github.event.merge_group.base_sha, github.event.merge_group.head_sha) || github.ref_name == 'main' && '--all-files' || '' }}\n - name: Check Actions\n uses: giner/check-actions@28d366c7cbbe235f9624a88aa31a628167eee28c # ratchet:giner/check-actions@v1.0.1\n with:\n check_permissions: false\n check_versions: false\n" }, { "path": ".github/workflows/preview.yml", "content": "name: Preview Deployment\nenv:\n VERCEL_ORG_ID: ${{ secrets.VERCEL_ORG_ID }}\n VERCEL_PROJECT_ID: ${{ secrets.VERCEL_PROJECT_ID }}\n VERCEL_CLI: vercel@50.14.1\non:\n push:\n branches-ignore:\n - main\n paths:\n - \"web/**\"\npermissions:\n contents: read\n pull-requests: write\njobs:\n Deploy-Preview:\n runs-on: ubuntu-latest\n timeout-minutes: 30\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd\n with:\n persist-credentials: false\n\n - name: Setup node\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v4\n with:\n node-version: 22\n cache: \"npm\"\n cache-dependency-path: ./web/package-lock.json\n\n - name: Pull Vercel Environment Information\n run: npx --yes ${{ env.VERCEL_CLI }} pull --yes --environment=preview --token=${{ secrets.VERCEL_TOKEN }}\n\n - name: Build Project Artifacts\n run: npx --yes ${{ env.VERCEL_CLI }} build --token=${{ secrets.VERCEL_TOKEN }}\n\n - name: Deploy Project Artifacts to Vercel\n id: deploy\n run: |\n DEPLOYMENT_URL=$(npx --yes ${{ env.VERCEL_CLI }} deploy --prebuilt --token=${{ secrets.VERCEL_TOKEN }})\n echo \"url=$DEPLOYMENT_URL\" >> \"$GITHUB_OUTPUT\"\n\n - name: Update PR comment with deployment URL\n if: always() && steps.deploy.outputs.url\n env:\n GH_TOKEN: ${{ github.token }}\n DEPLOYMENT_URL: ${{ steps.deploy.outputs.url }}\n run: |\n # Find the PR for this branch\n PR_NUMBER=$(gh pr list --head \"$GITHUB_REF_NAME\" --json number --jq '.[0].number')\n if [ -z \"$PR_NUMBER\" ]; then\n echo \"No open PR found for branch $GITHUB_REF_NAME, skipping comment.\"\n exit 0\n fi\n\n COMMENT_MARKER=\"\"\n COMMENT_BODY=\"$COMMENT_MARKER\n **Preview Deployment**\n\n | Status | Preview | Commit | Updated |\n | --- | --- | --- | --- |\n | ✅ | $DEPLOYMENT_URL | \\`${GITHUB_SHA::7}\\` | $(date -u '+%Y-%m-%d %H:%M:%S UTC') |\"\n\n # Find existing comment by marker\n EXISTING_COMMENT_ID=$(gh api \"repos/$GITHUB_REPOSITORY/issues/$PR_NUMBER/comments\" \\\n --jq \".[] | select(.body | startswith(\\\"$COMMENT_MARKER\\\")) | .id\" | head -1)\n\n if [ -n \"$EXISTING_COMMENT_ID\" ]; then\n gh api \"repos/$GITHUB_REPOSITORY/issues/comments/$EXISTING_COMMENT_ID\" \\\n --method PATCH --field body=\"$COMMENT_BODY\"\n else\n gh pr comment \"$PR_NUMBER\" --body \"$COMMENT_BODY\"\n fi\n" }, { "path": ".github/workflows/release-cli.yml", "content": "name: Release CLI\n\non:\n push:\n tags:\n - \"cli/v*.*.*\"\n\njobs:\n pypi:\n runs-on: ubuntu-latest\n environment:\n name: release-cli\n permissions:\n id-token: write\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n enable-cache: false\n version: \"0.9.9\"\n - run: |\n for goos in linux windows darwin; do\n for goarch in amd64 arm64; do\n GOOS=\"$goos\" GOARCH=\"$goarch\" uv build --wheel\n done\n done\n working-directory: cli\n - run: uv publish\n working-directory: cli\n\n docker-amd64:\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-cli-amd64\n - extras=ecr-cache\n environment: deploy\n permissions:\n id-token: write\n timeout-minutes: 30\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-cli\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v6.0.0\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2.0.10\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # ratchet:docker/setup-buildx-action@v4\n\n - name: Login to Docker Hub\n uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # ratchet:docker/login-action@v4\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@d08e5c354a6adb9ed34480a06d141179aa583294 # ratchet:docker/build-push-action@v7\n with:\n context: ./cli\n file: ./cli/Dockerfile\n platforms: linux/amd64\n cache-from: type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: type=inline\n outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n\n docker-arm64:\n runs-on:\n - runs-on\n - runner=2cpu-linux-arm64\n - run-id=${{ github.run_id }}-cli-arm64\n - extras=ecr-cache\n environment: deploy\n permissions:\n id-token: write\n timeout-minutes: 30\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-cli\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v6.0.0\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2.0.10\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # ratchet:docker/setup-buildx-action@v4\n\n - name: Login to Docker Hub\n uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # ratchet:docker/login-action@v4\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@d08e5c354a6adb9ed34480a06d141179aa583294 # ratchet:docker/build-push-action@v7\n with:\n context: ./cli\n file: ./cli/Dockerfile\n platforms: linux/arm64\n cache-from: type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: type=inline\n outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n\n merge-docker:\n needs:\n - docker-amd64\n - docker-arm64\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-cli-merge\n environment: deploy\n permissions:\n id-token: write\n timeout-minutes: 10\n env:\n REGISTRY_IMAGE: onyxdotapp/onyx-cli\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # ratchet:aws-actions/configure-aws-credentials@v6.0.0\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802 # ratchet:aws-actions/aws-secretsmanager-get-secrets@v2.0.10\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # ratchet:docker/setup-buildx-action@v4\n\n - name: Login to Docker Hub\n uses: docker/login-action@b45d80f862d83dbcd57f89517bcf500b2ab88fb2 # ratchet:docker/login-action@v4\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Create and push manifest\n env:\n AMD64_DIGEST: ${{ needs.docker-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.docker-arm64.outputs.digest }}\n TAG: ${{ github.ref_name }}\n run: |\n SANITIZED_TAG=\"${TAG#cli/}\"\n IMAGES=(\n \"${REGISTRY_IMAGE}@${AMD64_DIGEST}\"\n \"${REGISTRY_IMAGE}@${ARM64_DIGEST}\"\n )\n\n if [[ \"$TAG\" =~ ^cli/v[0-9]+\\.[0-9]+\\.[0-9]+$ ]]; then\n docker buildx imagetools create \\\n -t \"${REGISTRY_IMAGE}:${SANITIZED_TAG}\" \\\n -t \"${REGISTRY_IMAGE}:latest\" \\\n \"${IMAGES[@]}\"\n else\n docker buildx imagetools create \\\n -t \"${REGISTRY_IMAGE}:${SANITIZED_TAG}\" \\\n \"${IMAGES[@]}\"\n fi\n" }, { "path": ".github/workflows/release-devtools.yml", "content": "name: Release Devtools\n\non:\n push:\n tags:\n - \"ods/v*.*.*\"\n\njobs:\n pypi:\n runs-on: ubuntu-latest\n environment:\n name: release-devtools\n permissions:\n id-token: write\n timeout-minutes: 10\n strategy:\n matrix:\n os-arch:\n - { goos: \"linux\", goarch: \"amd64\" }\n - { goos: \"linux\", goarch: \"arm64\" }\n - { goos: \"windows\", goarch: \"amd64\" }\n - { goos: \"windows\", goarch: \"arm64\" }\n - { goos: \"darwin\", goarch: \"amd64\" }\n - { goos: \"darwin\", goarch: \"arm64\" }\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n enable-cache: false\n version: \"0.9.9\"\n - run: |\n GOOS=\"${{ matrix.os-arch.goos }}\" \\\n GOARCH=\"${{ matrix.os-arch.goarch }}\" \\\n uv build --wheel\n working-directory: tools/ods\n - run: uv publish\n working-directory: tools/ods\n" }, { "path": ".github/workflows/reusable-nightly-llm-provider-chat.yml", "content": "name: Reusable Nightly LLM Provider Chat Tests\n\non:\n workflow_call:\n inputs:\n openai_models:\n description: \"Comma-separated models for openai\"\n required: false\n default: \"\"\n type: string\n anthropic_models:\n description: \"Comma-separated models for anthropic\"\n required: false\n default: \"\"\n type: string\n bedrock_models:\n description: \"Comma-separated models for bedrock\"\n required: false\n default: \"\"\n type: string\n vertex_ai_models:\n description: \"Comma-separated models for vertex_ai\"\n required: false\n default: \"\"\n type: string\n azure_models:\n description: \"Comma-separated models for azure\"\n required: false\n default: \"\"\n type: string\n ollama_models:\n description: \"Comma-separated models for ollama_chat\"\n required: false\n default: \"\"\n type: string\n openrouter_models:\n description: \"Comma-separated models for openrouter\"\n required: false\n default: \"\"\n type: string\n azure_api_base:\n description: \"API base for azure provider\"\n required: false\n default: \"\"\n type: string\n strict:\n description: \"Default NIGHTLY_LLM_STRICT passed to tests\"\n required: false\n default: true\n type: boolean\n secrets:\n AWS_OIDC_ROLE_ARN:\n description: \"AWS role ARN for OIDC auth\"\n required: true\n\npermissions:\n contents: read\n id-token: write\n\njobs:\n build-backend-image:\n runs-on:\n [\n runs-on,\n runner=1cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-backend-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n environment: ci-protected\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, test/docker-username\n DOCKER_TOKEN, test/docker-token\n\n - name: Build backend image\n uses: ./.github/actions/build-backend-image\n with:\n runs-on-ecr-cache: ${{ env.RUNS_ON_ECR_CACHE }}\n ref-name: ${{ github.ref_name }}\n pr-number: ${{ github.event.pull_request.number }}\n github-sha: ${{ github.sha }}\n run-id: ${{ github.run_id }}\n docker-username: ${{ env.DOCKER_USERNAME }}\n docker-token: ${{ env.DOCKER_TOKEN }}\n docker-no-cache: ${{ vars.DOCKER_NO_CACHE == 'true' && 'true' || 'false' }}\n\n build-model-server-image:\n runs-on:\n [\n runs-on,\n runner=1cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-model-server-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n environment: ci-protected\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, test/docker-username\n DOCKER_TOKEN, test/docker-token\n\n - name: Build model server image\n uses: ./.github/actions/build-model-server-image\n with:\n runs-on-ecr-cache: ${{ env.RUNS_ON_ECR_CACHE }}\n ref-name: ${{ github.ref_name }}\n pr-number: ${{ github.event.pull_request.number }}\n github-sha: ${{ github.sha }}\n run-id: ${{ github.run_id }}\n docker-username: ${{ env.DOCKER_USERNAME }}\n docker-token: ${{ env.DOCKER_TOKEN }}\n\n build-integration-image:\n runs-on:\n [\n runs-on,\n runner=2cpu-linux-arm64,\n \"run-id=${{ github.run_id }}-build-integration-image\",\n \"extras=ecr-cache\",\n ]\n timeout-minutes: 45\n environment: ci-protected\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, test/docker-username\n DOCKER_TOKEN, test/docker-token\n\n - name: Build integration image\n uses: ./.github/actions/build-integration-image\n with:\n runs-on-ecr-cache: ${{ env.RUNS_ON_ECR_CACHE }}\n ref-name: ${{ github.ref_name }}\n pr-number: ${{ github.event.pull_request.number }}\n github-sha: ${{ github.sha }}\n run-id: ${{ github.run_id }}\n docker-username: ${{ env.DOCKER_USERNAME }}\n docker-token: ${{ env.DOCKER_TOKEN }}\n\n provider-chat-test:\n needs:\n [\n build-backend-image,\n build-model-server-image,\n build-integration-image,\n ]\n strategy:\n fail-fast: false\n matrix:\n include:\n - provider: openai\n models: ${{ inputs.openai_models }}\n api_key_env: OPENAI_API_KEY\n custom_config_env: \"\"\n api_base: \"\"\n api_version: \"\"\n deployment_name: \"\"\n required: true\n - provider: anthropic\n models: ${{ inputs.anthropic_models }}\n api_key_env: ANTHROPIC_API_KEY\n custom_config_env: \"\"\n api_base: \"\"\n api_version: \"\"\n deployment_name: \"\"\n required: true\n - provider: bedrock\n models: ${{ inputs.bedrock_models }}\n api_key_env: BEDROCK_API_KEY\n custom_config_env: \"\"\n api_base: \"\"\n api_version: \"\"\n deployment_name: \"\"\n required: false\n - provider: vertex_ai\n models: ${{ inputs.vertex_ai_models }}\n api_key_env: \"\"\n custom_config_env: NIGHTLY_LLM_VERTEX_AI_CUSTOM_CONFIG_JSON\n api_base: \"\"\n api_version: \"\"\n deployment_name: \"\"\n required: false\n - provider: azure\n models: ${{ inputs.azure_models }}\n api_key_env: AZURE_API_KEY\n custom_config_env: \"\"\n api_base: ${{ inputs.azure_api_base }}\n api_version: \"2025-04-01-preview\"\n deployment_name: \"\"\n required: false\n - provider: ollama_chat\n models: ${{ inputs.ollama_models }}\n api_key_env: OLLAMA_API_KEY\n custom_config_env: \"\"\n api_base: \"https://ollama.com\"\n api_version: \"\"\n deployment_name: \"\"\n required: false\n - provider: openrouter\n models: ${{ inputs.openrouter_models }}\n api_key_env: OPENROUTER_API_KEY\n custom_config_env: \"\"\n api_base: \"https://openrouter.ai/api/v1\"\n api_version: \"\"\n deployment_name: \"\"\n required: false\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - \"run-id=${{ github.run_id }}-nightly-${{ matrix.provider }}-provider-chat-test\"\n - extras=ecr-cache\n timeout-minutes: 45\n environment: ci-protected\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n # Keep JSON values unparsed so vertex custom config is passed as raw JSON.\n parse-json-secrets: false\n secret-ids: |\n DOCKER_USERNAME, test/docker-username\n DOCKER_TOKEN, test/docker-token\n OPENAI_API_KEY, test/openai-api-key\n ANTHROPIC_API_KEY, test/anthropic-api-key\n BEDROCK_API_KEY, test/bedrock-api-key\n NIGHTLY_LLM_VERTEX_AI_CUSTOM_CONFIG_JSON, test/nightly-llm-vertex-ai-custom-config-json\n AZURE_API_KEY, test/azure-api-key\n OLLAMA_API_KEY, test/ollama-api-key\n OPENROUTER_API_KEY, test/openrouter-api-key\n\n - name: Run nightly provider chat test\n uses: ./.github/actions/run-nightly-provider-chat-test\n with:\n provider: ${{ matrix.provider }}\n models: ${{ matrix.models }}\n provider-api-key: ${{ matrix.api_key_env && env[matrix.api_key_env] || '' }}\n strict: ${{ inputs.strict && 'true' || 'false' }}\n api-base: ${{ matrix.api_base }}\n api-version: ${{ matrix.api_version }}\n deployment-name: ${{ matrix.deployment_name }}\n custom-config-json: ${{ matrix.custom_config_env && env[matrix.custom_config_env] || '' }}\n runs-on-ecr-cache: ${{ env.RUNS_ON_ECR_CACHE }}\n run-id: ${{ github.run_id }}\n docker-username: ${{ env.DOCKER_USERNAME }}\n docker-token: ${{ env.DOCKER_TOKEN }}\n\n - name: Dump API server logs\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose logs --no-color api_server > $GITHUB_WORKSPACE/api_server.log || true\n\n - name: Dump all-container logs\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose logs --no-color > $GITHUB_WORKSPACE/docker-compose.log || true\n\n - name: Upload logs\n if: always()\n uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f\n with:\n name: docker-all-logs-nightly-${{ matrix.provider }}-llm-provider\n path: |\n ${{ github.workspace }}/api_server.log\n ${{ github.workspace }}/docker-compose.log\n\n - name: Stop Docker containers\n if: always()\n run: |\n cd deployment/docker_compose\n docker compose down -v\n" }, { "path": ".github/workflows/sandbox-deployment.yml", "content": "name: Build and Push Sandbox Image on Tag\n\non:\n push:\n tags:\n - \"experimental-cc4a.*\"\n\n# Restrictive defaults; jobs declare what they need.\npermissions: {}\n\njobs:\n check-sandbox-changes:\n runs-on: ubuntu-slim\n timeout-minutes: 10\n permissions:\n contents: read\n outputs:\n sandbox-changed: ${{ steps.check.outputs.sandbox-changed }}\n new-version: ${{ steps.version.outputs.new-version }}\n steps:\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n fetch-depth: 0\n\n - name: Check for sandbox-relevant file changes\n id: check\n run: |\n # Get the previous tag to diff against\n CURRENT_TAG=\"${GITHUB_REF_NAME}\"\n PREVIOUS_TAG=$(git tag --sort=-creatordate | grep '^experimental-cc4a\\.' | grep -v \"^${CURRENT_TAG}$\" | head -n 1)\n\n if [ -z \"$PREVIOUS_TAG\" ]; then\n echo \"No previous experimental-cc4a tag found, building unconditionally\"\n echo \"sandbox-changed=true\" >> \"$GITHUB_OUTPUT\"\n exit 0\n fi\n\n echo \"Comparing ${PREVIOUS_TAG}..${CURRENT_TAG}\"\n\n # Check if any sandbox-relevant files changed\n SANDBOX_PATHS=(\n \"backend/onyx/server/features/build/sandbox/\"\n )\n\n CHANGED=false\n for path in \"${SANDBOX_PATHS[@]}\"; do\n if git diff --name-only \"${PREVIOUS_TAG}..${CURRENT_TAG}\" -- \"$path\" | grep -q .; then\n echo \"Changes detected in: $path\"\n CHANGED=true\n break\n fi\n done\n\n echo \"sandbox-changed=$CHANGED\" >> \"$GITHUB_OUTPUT\"\n\n - name: Determine new sandbox version\n id: version\n if: steps.check.outputs.sandbox-changed == 'true'\n run: |\n # Query Docker Hub for the latest versioned tag\n LATEST_TAG=$(curl -s \"https://hub.docker.com/v2/repositories/onyxdotapp/sandbox/tags?page_size=100\" \\\n | jq -r '.results[].name' \\\n | grep -E '^v[0-9]+\\.[0-9]+\\.[0-9]+$' \\\n | sort -V \\\n | tail -n 1)\n\n if [ -z \"$LATEST_TAG\" ]; then\n echo \"No existing version tags found on Docker Hub, starting at 0.1.1\"\n NEW_VERSION=\"0.1.1\"\n else\n CURRENT_VERSION=\"${LATEST_TAG#v}\"\n echo \"Latest version on Docker Hub: $CURRENT_VERSION\"\n\n # Increment patch version\n MAJOR=$(echo \"$CURRENT_VERSION\" | cut -d. -f1)\n MINOR=$(echo \"$CURRENT_VERSION\" | cut -d. -f2)\n PATCH=$(echo \"$CURRENT_VERSION\" | cut -d. -f3)\n NEW_PATCH=$((PATCH + 1))\n NEW_VERSION=\"${MAJOR}.${MINOR}.${NEW_PATCH}\"\n fi\n\n echo \"New version: $NEW_VERSION\"\n echo \"new-version=$NEW_VERSION\" >> \"$GITHUB_OUTPUT\"\n\n build-sandbox-amd64:\n needs: check-sandbox-changes\n if: needs.check-sandbox-changes.outputs.sandbox-changed == 'true'\n runs-on:\n - runs-on\n - runner=4cpu-linux-x64\n - run-id=${{ github.run_id }}-sandbox-amd64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n permissions:\n contents: read\n id-token: write\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/sandbox\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push AMD64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend/onyx/server/features/build/sandbox/kubernetes/docker\n file: ./backend/onyx/server/features/build/sandbox/kubernetes/docker/Dockerfile\n platforms: linux/amd64\n labels: ${{ steps.meta.outputs.labels }}\n cache-from: |\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n\n build-sandbox-arm64:\n needs: check-sandbox-changes\n if: needs.check-sandbox-changes.outputs.sandbox-changed == 'true'\n runs-on:\n - runs-on\n - runner=4cpu-linux-arm64\n - run-id=${{ github.run_id }}-sandbox-arm64\n - extras=ecr-cache\n timeout-minutes: 90\n environment: release\n permissions:\n contents: read\n id-token: write\n outputs:\n digest: ${{ steps.build.outputs.digest }}\n env:\n REGISTRY_IMAGE: onyxdotapp/sandbox\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Checkout\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n persist-credentials: false\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Build and push ARM64\n id: build\n uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 # ratchet:docker/build-push-action@v6\n with:\n context: ./backend/onyx/server/features/build/sandbox/kubernetes/docker\n file: ./backend/onyx/server/features/build/sandbox/kubernetes/docker/Dockerfile\n platforms: linux/arm64\n labels: ${{ steps.meta.outputs.labels }}\n cache-from: |\n type=registry,ref=${{ env.REGISTRY_IMAGE }}:latest\n cache-to: |\n type=inline\n outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true\n\n merge-sandbox:\n needs:\n - check-sandbox-changes\n - build-sandbox-amd64\n - build-sandbox-arm64\n runs-on:\n - runs-on\n - runner=2cpu-linux-x64\n - run-id=${{ github.run_id }}-merge-sandbox\n - extras=ecr-cache\n timeout-minutes: 30\n environment: release\n permissions:\n id-token: write\n env:\n REGISTRY_IMAGE: onyxdotapp/sandbox\n steps:\n - uses: runs-on/action@cd2b598b0515d39d78c38a02d529db87d2196d1e # ratchet:runs-on/action@v2\n\n - name: Configure AWS credentials\n uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7\n with:\n role-to-assume: ${{ secrets.AWS_OIDC_ROLE_ARN }}\n aws-region: us-east-2\n\n - name: Get AWS Secrets\n uses: aws-actions/aws-secretsmanager-get-secrets@a9a7eb4e2f2871d30dc5b892576fde60a2ecc802\n with:\n secret-ids: |\n DOCKER_USERNAME, deploy/docker-username\n DOCKER_TOKEN, deploy/docker-token\n parse-json-secrets: true\n\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # ratchet:docker/setup-buildx-action@v3\n\n - name: Login to Docker Hub\n uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # ratchet:docker/login-action@v3\n with:\n username: ${{ env.DOCKER_USERNAME }}\n password: ${{ env.DOCKER_TOKEN }}\n\n - name: Docker meta\n id: meta\n uses: docker/metadata-action@030e881283bb7a6894de51c315a6bfe6a94e05cf # ratchet:docker/metadata-action@v6.0.0\n with:\n images: ${{ env.REGISTRY_IMAGE }}\n flavor: |\n latest=false\n tags: |\n type=raw,value=v${{ needs.check-sandbox-changes.outputs.new-version }}\n type=raw,value=latest\n\n - name: Create and push manifest\n env:\n IMAGE_REPO: ${{ env.REGISTRY_IMAGE }}\n AMD64_DIGEST: ${{ needs.build-sandbox-amd64.outputs.digest }}\n ARM64_DIGEST: ${{ needs.build-sandbox-arm64.outputs.digest }}\n META_TAGS: ${{ steps.meta.outputs.tags }}\n run: |\n IMAGES=\"${IMAGE_REPO}@${AMD64_DIGEST} ${IMAGE_REPO}@${ARM64_DIGEST}\"\n docker buildx imagetools create \\\n $(printf '%s\\n' \"${META_TAGS}\" | xargs -I {} echo -t {}) \\\n $IMAGES\n" }, { "path": ".github/workflows/storybook-deploy.yml", "content": "name: Storybook Deploy\nenv:\n VERCEL_ORG_ID: ${{ secrets.VERCEL_ORG_ID }}\n VERCEL_PROJECT_ID: prj_sG49mVsA25UsxIPhN2pmBJlikJZM\n VERCEL_CLI: vercel@50.14.1\n VERCEL_TOKEN: ${{ secrets.VERCEL_TOKEN }}\n\nconcurrency:\n group: storybook-deploy-production\n cancel-in-progress: true\n\non:\n workflow_dispatch:\n push:\n branches:\n - main\n paths:\n - \"web/lib/opal/**\"\n - \"web/src/refresh-components/**\"\n - \"web/.storybook/**\"\n - \"web/package.json\"\n - \"web/package-lock.json\"\npermissions:\n contents: read\njobs:\n Deploy-Storybook:\n runs-on: ubuntu-latest\n environment: ci-protected\n timeout-minutes: 30\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v4\n with:\n persist-credentials: false\n\n - name: Setup node\n uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # ratchet:actions/setup-node@v4\n with:\n node-version: 22\n cache: \"npm\"\n cache-dependency-path: ./web/package-lock.json\n\n - name: Install dependencies\n working-directory: web\n run: npm ci\n\n - name: Build Storybook\n working-directory: web\n run: npm run storybook:build\n\n - name: Deploy to Vercel (Production)\n working-directory: web\n run: npx --yes \"$VERCEL_CLI\" deploy storybook-static/ --prod --yes --token=\"$VERCEL_TOKEN\"\n\n notify-slack-on-failure:\n needs: Deploy-Storybook\n if: always() && needs.Deploy-Storybook.result == 'failure'\n runs-on: ubuntu-latest\n environment: ci-protected\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v4\n with:\n persist-credentials: false\n sparse-checkout: .github/actions/slack-notify\n\n - name: Send Slack notification\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.MONITOR_DEPLOYMENTS_WEBHOOK }}\n failed-jobs: \"• Deploy-Storybook\"\n title: \"🚨 Storybook Deploy Failed\"\n" }, { "path": ".github/workflows/sync_foss.yml", "content": "name: Sync FOSS Repo\n\non:\n schedule:\n # Run daily at 3am PT (11am UTC during PST)\n - cron: '0 11 * * *'\n workflow_dispatch:\n\njobs:\n sync-foss:\n runs-on: ubuntu-latest\n environment: ci-protected\n timeout-minutes: 45\n permissions:\n contents: read\n steps:\n - name: Checkout main Onyx repo\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n fetch-depth: 0\n persist-credentials: false\n\n - name: Install git-filter-repo\n run: |\n sudo apt-get update && sudo apt-get install -y git-filter-repo\n\n - name: Configure SSH for deploy key\n env:\n FOSS_REPO_DEPLOY_KEY: ${{ secrets.FOSS_REPO_DEPLOY_KEY }}\n run: |\n mkdir -p ~/.ssh\n echo \"$FOSS_REPO_DEPLOY_KEY\" > ~/.ssh/id_ed25519\n chmod 600 ~/.ssh/id_ed25519\n ssh-keyscan github.com >> ~/.ssh/known_hosts\n\n - name: Set Git config\n run: |\n git config --global user.name \"onyx-bot\"\n git config --global user.email \"bot@onyx.app\"\n\n - name: Build FOSS version\n run: bash backend/scripts/make_foss_repo.sh\n\n - name: Push to FOSS repo\n env:\n FOSS_REPO_URL: git@github.com:onyx-dot-app/onyx-foss.git\n run: |\n cd /tmp/foss_repo\n git remote add public \"$FOSS_REPO_URL\"\n git push --force public main\n" }, { "path": ".github/workflows/tag-nightly.yml", "content": "name: Nightly Tag Push\n\non:\n schedule:\n - cron: \"0 10 * * *\" # Runs every day at 2 AM PST / 3 AM PDT / 10 AM UTC\n workflow_dispatch:\n\npermissions:\n contents: write # Allows pushing tags to the repository\n\njobs:\n create-and-push-tag:\n runs-on: ubuntu-slim\n environment: ci-protected\n timeout-minutes: 45\n\n steps:\n # actions using GITHUB_TOKEN cannot trigger another workflow, but we do want this to trigger docker pushes\n # see https://github.com/orgs/community/discussions/27028#discussioncomment-3254367 for the workaround we\n # implement here which needs an actual user's deploy key\n - name: Checkout code\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6\n with:\n ssh-key: \"${{ secrets.DEPLOY_KEY }}\"\n persist-credentials: true\n\n - name: Set up Git user\n run: |\n git config user.name \"Onyx Bot [bot]\"\n git config user.email \"onyx-bot[bot]@onyx.app\"\n\n - name: Check for existing nightly tag\n id: check_tag\n run: |\n if git tag --points-at HEAD --list \"nightly-latest*\" | grep -q .; then\n echo \"A tag starting with 'nightly-latest' already exists on HEAD.\"\n echo \"tag_exists=true\" >> $GITHUB_OUTPUT\n else\n echo \"No tag starting with 'nightly-latest' exists on HEAD.\"\n echo \"tag_exists=false\" >> $GITHUB_OUTPUT\n fi\n\n # don't tag again if HEAD already has a nightly-latest tag on it\n - name: Create Nightly Tag\n if: steps.check_tag.outputs.tag_exists == 'false'\n env:\n DATE: ${{ github.run_id }}\n run: |\n TAG_NAME=\"nightly-latest-$(date +'%Y%m%d')\"\n echo \"Creating tag: $TAG_NAME\"\n git tag $TAG_NAME\n\n - name: Push Tag\n if: steps.check_tag.outputs.tag_exists == 'false'\n run: |\n TAG_NAME=\"nightly-latest-$(date +'%Y%m%d')\"\n git push origin $TAG_NAME\n\n - name: Send Slack notification\n if: failure()\n uses: ./.github/actions/slack-notify\n with:\n webhook-url: ${{ secrets.MONITOR_DEPLOYMENTS_WEBHOOK }}\n title: \"🚨 Nightly Tag Push Failed\"\n ref-name: ${{ github.ref_name }}\n failed-jobs: \"create-and-push-tag\"\n" }, { "path": ".github/workflows/zizmor.yml", "content": "name: Run Zizmor\n\non:\n push:\n branches: [\"main\"]\n pull_request:\n branches: [\"**\"]\n paths:\n - \".github/**\"\n\npermissions: {}\n\njobs:\n zizmor:\n name: zizmor\n runs-on: ubuntu-slim\n timeout-minutes: 45\n permissions:\n security-events: write # needed for SARIF uploads\n steps:\n - name: Checkout repository\n uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # ratchet:actions/checkout@v6.0.2\n with:\n persist-credentials: false\n\n - name: Install the latest version of uv\n uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # ratchet:astral-sh/setup-uv@v7\n with:\n enable-cache: false\n version: \"0.9.9\"\n\n - name: Run zizmor\n run: uv run --no-sync --with zizmor zizmor --format=sarif . > results.sarif\n env:\n GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\n - name: Upload SARIF file\n uses: github/codeql-action/upload-sarif@ba454b8ab46733eb6145342877cd148270bb77ab # ratchet:github/codeql-action/upload-sarif@codeql-bundle-v2.23.5\n with:\n sarif_file: results.sarif\n category: zizmor\n" }, { "path": ".gitignore", "content": "# editors\n.vscode/*\n!/.vscode/env_template.txt\n!/.vscode/env.web_template.txt\n!/.vscode/launch.json\n!/.vscode/tasks.template.jsonc\n.zed\n.cursor\n!/.cursor/mcp.json\n!/.cursor/skills/\n\n# macos\n.DS_store\n\n# python\n.venv\n.mypy_cache\n.idea\n\n# testing\n/web/test-results/\nbackend/onyx/agent_search/main/test_data.json\nbackend/tests/regression/answer_quality/test_data.json\nbackend/tests/regression/search_quality/eval-*\nbackend/tests/regression/search_quality/search_eval_config.yaml\nbackend/tests/regression/search_quality/*.json\nbackend/onyx/evals/data/\nbackend/onyx/evals/one_off/*.json\n*.log\n*.csv\n\n# secret files\n.env\njira_test_env\nsettings.json\n\n# others\n/deployment/data/nginx/app.conf\n/deployment/data/nginx/mcp.conf.inc\n/deployment/data/nginx/mcp_upstream.conf.inc\n*.sw?\n/backend/tests/regression/answer_quality/search_test_config.yaml\n*.egg-info\n\n# Local .terraform directories\n**/.terraform/*\n\n# Local .tfstate files\n*.tfstate\n*.tfstate.*\n\n# Local .terraform.lock.hcl file\n.terraform.lock.hcl\n\nnode_modules\n\n# MCP configs\n.playwright-mcp\n\n# plans\nplans/\n" }, { "path": ".greptile/config.json", "content": "{\n \"labels\": [],\n \"comment\": \"\",\n \"fixWithAI\": true,\n \"hideFooter\": false,\n \"strictness\": 3,\n \"statusCheck\": true,\n \"commentTypes\": [\n \"logic\",\n \"syntax\",\n \"style\"\n ],\n \"instructions\": \"\",\n \"disabledLabels\": [],\n \"excludeAuthors\": [\n \"dependabot[bot]\",\n \"renovate[bot]\"\n ],\n \"ignoreKeywords\": \"\",\n \"ignorePatterns\": \"\",\n \"includeAuthors\": [],\n \"summarySection\": {\n \"included\": true,\n \"collapsible\": false,\n \"defaultOpen\": false\n },\n \"excludeBranches\": [],\n \"fileChangeLimit\": 300,\n \"includeBranches\": [],\n \"includeKeywords\": \"\",\n \"triggerOnUpdates\": true,\n \"updateExistingSummaryComment\": true,\n \"updateSummaryOnly\": false,\n \"issuesTableSection\": {\n \"included\": true,\n \"collapsible\": false,\n \"defaultOpen\": false\n },\n \"statusCommentsEnabled\": true,\n \"confidenceScoreSection\": {\n \"included\": true,\n \"collapsible\": false\n },\n \"sequenceDiagramSection\": {\n \"included\": true,\n \"collapsible\": false,\n \"defaultOpen\": false\n },\n \"shouldUpdateDescription\": false,\n \"rules\": [\n {\n \"scope\": [\"web/**\"],\n \"rule\": \"In Onyx's Next.js app, the `app/ee/admin/` directory is a filesystem convention for Enterprise Edition route overrides — it does NOT add an `/ee/` prefix to the URL. Both `app/admin/groups/page.tsx` and `app/ee/admin/groups/page.tsx` serve the same URL `/admin/groups`. Hardcoded `/admin/...` paths in router.push() calls are correct and do NOT break EE deployments. Do not flag hardcoded admin paths as bugs.\"\n },\n {\n \"scope\": [\"web/**\"],\n \"rule\": \"In Onyx, each API key creates a unique user row in the database with a unique `user_id` (UUID). There is a 1:1 mapping between API keys and their backing user records. Multiple API keys do NOT share the same `user_id`. Do not flag potential duplicate row IDs when using `user_id` from API key descriptors.\"\n },\n {\n \"scope\": [\"backend/**/*.py\"],\n \"rule\": \"Never raise HTTPException directly in business code. Use `raise OnyxError(OnyxErrorCode.XXX, \\\"message\\\")` from `onyx.error_handling.exceptions`. A global FastAPI exception handler converts OnyxError into structured JSON responses with {\\\"error_code\\\": \\\"...\\\", \\\"detail\\\": \\\"...\\\"}. Error codes are defined in `onyx.error_handling.error_codes.OnyxErrorCode`. For upstream errors with dynamic HTTP status codes, use `status_code_override`: `raise OnyxError(OnyxErrorCode.BAD_GATEWAY, detail, status_code_override=upstream_status)`.\"\n }\n ]\n}\n" }, { "path": ".greptile/files.json", "content": "[\n {\n \"scope\": [],\n \"path\": \"contributing_guides/best_practices.md\",\n \"description\": \"Best practices for contributing to the codebase\"\n },\n {\n \"scope\": [\"web/**\"],\n \"path\": \"web/AGENTS.md\",\n \"description\": \"Frontend coding standards for the web directory\"\n },\n {\n \"scope\": [\"web/**\"],\n \"path\": \"web/tests/README.md\",\n \"description\": \"Frontend testing guide and conventions\"\n },\n {\n \"scope\": [\"web/**\"],\n \"path\": \"web/CLAUDE.md\",\n \"description\": \"Single source of truth for frontend coding standards\"\n },\n {\n \"scope\": [\"web/**\"],\n \"path\": \"web/lib/opal/README.md\",\n \"description\": \"Opal component library usage guide\"\n },\n {\n \"scope\": [\"backend/**\"],\n \"path\": \"backend/tests/README.md\",\n \"description\": \"Backend testing guide covering all 4 test types, fixtures, and conventions\"\n },\n {\n \"scope\": [\"backend/onyx/connectors/**\"],\n \"path\": \"backend/onyx/connectors/README.md\",\n \"description\": \"Connector development guide covering design, interfaces, and required changes\"\n },\n {\n \"scope\": [],\n \"path\": \"CLAUDE.md\",\n \"description\": \"Project instructions and coding standards\"\n },\n {\n \"scope\": [],\n \"path\": \"backend/alembic/README.md\",\n \"description\": \"Migration guidance, including multi-tenant migration behavior\"\n },\n {\n \"scope\": [],\n \"path\": \"deployment/helm/charts/onyx/values-lite.yaml\",\n \"description\": \"Lite deployment Helm values and service assumptions\"\n },\n {\n \"scope\": [],\n \"path\": \"deployment/docker_compose/docker-compose.onyx-lite.yml\",\n \"description\": \"Lite deployment Docker Compose overlay and disabled service behavior\"\n }\n]\n" }, { "path": ".greptile/rules.md", "content": "# Greptile Review Rules\n\n## Type Annotations\n\nUse explicit type annotations for variables to enhance code clarity, especially when moving type hints around in the code.\n\n## Best Practices\n\nUse the \"Engineering Best Practices\" section of `CONTRIBUTING.md` as core review context. Prefer consistency with existing patterns, fix issues in code you touch, avoid tacking new features onto muddy interfaces, fail loudly instead of silently swallowing errors, keep code strictly typed, preserve clear state boundaries, remove duplicate or dead logic, break up overly long functions, avoid hidden import-time side effects, respect module boundaries, and favor correctness-by-construction over relying on callers to use an API correctly.\n\n## TODOs\n\nWhenever a TODO is added, there must always be an associated name or ticket with that TODO in the style of `TODO(name): ...` or `TODO(1234): ...`\n\n## Debugging Code\n\nRemove temporary debugging code before merging to production, especially tenant-specific debugging logs.\n\n## Hardcoded Booleans\n\nWhen hardcoding a boolean variable to a constant value, remove the variable entirely and clean up all places where it's used rather than just setting it to a constant.\n\n## Multi-tenant vs Single-tenant\n\nCode changes must consider both multi-tenant and single-tenant deployments. In multi-tenant mode, preserve tenant isolation, ensure tenant context is propagated correctly, and avoid assumptions that only hold for a single shared schema or globally shared state. In single-tenant mode, avoid introducing unnecessary tenant-specific requirements or cloud-only control-plane dependencies.\n\n## Nginx Routing — New Backend Routes\n\nWhenever a new backend route is added that does NOT start with `/api`, it must also be explicitly added to ALL nginx configs:\n\n- `deployment/helm/charts/onyx/templates/nginx-conf.yaml` (Helm/k8s)\n- `deployment/data/nginx/app.conf.template` (docker-compose dev)\n- `deployment/data/nginx/app.conf.template.prod` (docker-compose prod)\n- `deployment/data/nginx/app.conf.template.no-letsencrypt` (docker-compose no-letsencrypt)\n\nRoutes not starting with `/api` are not caught by the existing `^/(api|openapi\\.json)` location block and will fall through to `location /`, which proxies to the Next.js web server and returns an HTML 404. The new location block must be placed before the `/api` block. Examples of routes that need this treatment: `/scim`, `/mcp`.\n\n## Full vs Lite Deployments\n\nCode changes must consider both regular Onyx deployments and Onyx lite deployments. Lite deployments disable the vector DB, Redis, model servers, and background workers by default, use PostgreSQL-backed cache/auth/file storage, and rely on the API server to handle background work. Do not assume those services are available unless the code path is explicitly limited to full deployments.\n\n## SWR Cache Keys — Always Use SWR_KEYS Registry\n\nAll `useSWR()` calls and `mutate()` calls in the frontend must reference the centralized `SWR_KEYS` registry in `web/src/lib/swr-keys.ts` instead of inline endpoint strings or local string constants. Never write `useSWR(\"/api/some/endpoint\", ...)` or `mutate(\"/api/some/endpoint\")` — always use the corresponding `SWR_KEYS.someEndpoint` constant. If the endpoint does not yet exist in the registry, add it there first. This applies to all variants of an endpoint (e.g. query-string variants like `?get_editable=true` must also be registered as their own key).\n" }, { "path": ".pre-commit-config.yaml", "content": "default_install_hook_types:\n - pre-commit\n - post-checkout\n - post-merge\n - post-rewrite\nrepos:\n - repo: https://github.com/astral-sh/uv-pre-commit\n # From: https://github.com/astral-sh/uv-pre-commit/pull/53/commits/d30b4298e4fb63ce8609e29acdbcf4c9018a483c\n rev: d30b4298e4fb63ce8609e29acdbcf4c9018a483c\n hooks:\n - id: uv-sync\n args: [\"--locked\", \"--all-extras\"]\n - id: uv-lock\n - id: uv-export\n name: uv-export default.txt\n args:\n [\n \"--no-emit-project\",\n \"--no-default-groups\",\n \"--no-hashes\",\n \"--extra\",\n \"backend\",\n \"-o\",\n \"backend/requirements/default.txt\",\n ]\n files: ^(pyproject\\.toml|uv\\.lock|backend/requirements/.*\\.txt)$\n - id: uv-export\n name: uv-export dev.txt\n args:\n [\n \"--no-emit-project\",\n \"--no-default-groups\",\n \"--no-hashes\",\n \"--extra\",\n \"dev\",\n \"-o\",\n \"backend/requirements/dev.txt\",\n ]\n files: ^(pyproject\\.toml|uv\\.lock|backend/requirements/.*\\.txt)$\n - id: uv-export\n name: uv-export ee.txt\n args:\n [\n \"--no-emit-project\",\n \"--no-default-groups\",\n \"--no-hashes\",\n \"--extra\",\n \"ee\",\n \"-o\",\n \"backend/requirements/ee.txt\",\n ]\n files: ^(pyproject\\.toml|uv\\.lock|backend/requirements/.*\\.txt)$\n - id: uv-export\n name: uv-export model_server.txt\n args:\n [\n \"--no-emit-project\",\n \"--no-default-groups\",\n \"--no-hashes\",\n \"--extra\",\n \"model_server\",\n \"-o\",\n \"backend/requirements/model_server.txt\",\n ]\n files: ^(pyproject\\.toml|uv\\.lock|backend/requirements/.*\\.txt)$\n - id: uv-run\n name: Check lazy imports\n args: [\"--active\", \"--with=onyx-devtools\", \"ods\", \"check-lazy-imports\"]\n pass_filenames: true\n files: ^backend/(?!\\.venv/|scripts/).*\\.py$\n # NOTE: This takes ~6s on a single, large module which is prohibitively slow.\n # - id: uv-run\n # name: mypy\n # args: [\"--all-extras\", \"mypy\"]\n # pass_filenames: true\n # files: ^backend/.*\\.py$\n\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c # frozen: v6.0.0\n hooks:\n - id: check-added-large-files\n name: Check for added large files\n args: [\"--maxkb=1500\"]\n\n - repo: https://github.com/rhysd/actionlint\n rev: a443f344ff32813837fa49f7aa6cbc478d770e62 # frozen: v1.7.9\n hooks:\n - id: actionlint\n\n - repo: https://github.com/psf/black\n rev: 8a737e727ac5ab2f1d4cf5876720ed276dc8dc4b # frozen: 25.1.0\n hooks:\n - id: black\n language_version: python3.11\n\n # this is a fork which keeps compatibility with black\n - repo: https://github.com/wimglenn/reorder-python-imports-black\n rev: f55cd27f90f0cf0ee775002c2383ce1c7820013d # frozen: v3.14.0\n hooks:\n - id: reorder-python-imports\n args: [\"--py311-plus\", \"--application-directories=backend/\"]\n # need to ignore alembic files, since reorder-python-imports gets confused\n # and thinks that alembic is a local package since there is a folder\n # in the backend directory called `alembic`\n exclude: ^backend/alembic/\n\n # These settings will remove unused imports with side effects\n # Note: The repo currently does not and should not have imports with side effects\n - repo: https://github.com/PyCQA/autoflake\n rev: 0544741e2b4a22b472d9d93e37d4ea9153820bb1 # frozen: v2.3.1\n hooks:\n - id: autoflake\n args:\n [\n \"--remove-all-unused-imports\",\n \"--remove-unused-variables\",\n \"--in-place\",\n \"--recursive\",\n ]\n\n - repo: https://github.com/golangci/golangci-lint\n rev: 5d1e709b7be35cb2025444e19de266b056b7b7ee # frozen: v2.10.1\n hooks:\n - id: golangci-lint\n language_version: \"1.26.1\"\n entry: bash -c \"find . -name go.mod -not -path './.venv/*' -print0 | xargs -0 -I{} bash -c 'cd \\\"$(dirname {})\\\" && golangci-lint run ./...'\"\n\n - repo: https://github.com/astral-sh/ruff-pre-commit\n # Ruff version.\n rev: 971923581912ef60a6b70dbf0c3e9a39563c9d47 # frozen: v0.11.4\n hooks:\n - id: ruff\n\n - repo: https://github.com/pre-commit/mirrors-prettier\n rev: ffb6a759a979008c0e6dff86e39f4745a2d9eac4 # frozen: v3.1.0\n hooks:\n - id: prettier\n types_or: [html, css, javascript, ts, tsx]\n language_version: system\n\n - repo: https://github.com/sirwart/ripsecrets\n rev: 7d94620933e79b8acaa0cd9e60e9864b07673d86 # frozen: v0.1.11\n hooks:\n - id: ripsecrets\n args:\n - --additional-pattern\n - ^sk-[A-Za-z0-9_\\-]{20,}$\n\n - repo: local\n hooks:\n - id: terraform-fmt\n name: terraform fmt\n entry: terraform fmt -recursive\n language: system\n pass_filenames: false\n files: \\.tf$\n\n - id: npm-install\n name: npm install\n description: \"Automatically run 'npm install' after a checkout, pull or rebase\"\n language: system\n entry: bash -c 'cd web && npm install --no-save'\n pass_filenames: false\n files: ^web/package(-lock)?\\.json$\n stages: [post-checkout, post-merge, post-rewrite]\n - id: npm-install-check\n name: npm install --package-lock-only\n description: \"Check the 'web/package-lock.json' is updated\"\n language: system\n entry: bash -c 'cd web && npm install --package-lock-only'\n pass_filenames: false\n files: ^web/package(-lock)?\\.json$\n\n # Uses tsgo (TypeScript's native Go compiler) for ~10x faster type checking.\n # This is a preview package - if it breaks:\n # 1. Try updating: cd web && npm update @typescript/native-preview\n # 2. Or fallback to tsc: replace 'tsgo' with 'tsc' below\n - id: typescript-check\n name: TypeScript type check\n entry: bash -c 'cd web && npx tsgo --noEmit --project tsconfig.types.json'\n language: system\n pass_filenames: false\n files: ^web/.*\\.(ts|tsx)$\n" }, { "path": ".prettierignore", "content": "backend/tests/integration/tests/pruning/website\n" }, { "path": ".vscode/env.web_template.txt", "content": "# Copy this file to .env.web in the .vscode folder.\n# Fill in the values as needed\n# Web Server specific environment variables\n# Minimal set needed for Next.js dev server\n\n# Auth\nAUTH_TYPE=basic\nDEV_MODE=true\n\n# Enable the full set of Danswer Enterprise Edition features.\n# NOTE: DO NOT ENABLE THIS UNLESS YOU HAVE A PAID ENTERPRISE LICENSE (or if you\n# are using this for local testing/development).\nENABLE_PAID_ENTERPRISE_EDITION_FEATURES=false\n\n# Enable Onyx Craft\nENABLE_CRAFT=true\n" }, { "path": ".vscode/env_template.txt", "content": "# Copy this file to .env in the .vscode folder.\n# Fill in the values as needed; it is recommended to set the\n# GEN_AI_API_KEY value to avoid having to set up an LLM in the UI.\n# Also check out onyx/backend/scripts/restart_containers.sh for a script to\n# restart the containers which Onyx relies on outside of VSCode/Cursor\n# processes.\n\n\nAUTH_TYPE=basic\n# Recommended for basic auth - used for signing password reset and verification tokens\n# Generate a secure value with: openssl rand -hex 32\nUSER_AUTH_SECRET=\"\"\nDEV_MODE=true\n\n\n# Always keep these on for Dev.\n# Logs model prompts, reasoning, and answer to stdout.\nLOG_ONYX_MODEL_INTERACTIONS=False\n# More verbose logging\nLOG_LEVEL=debug\n\n\n# Useful if you want to toggle auth on/off (google_oauth/OIDC specifically).\nOAUTH_CLIENT_ID=\nOAUTH_CLIENT_SECRET=\nOPENID_CONFIG_URL=\nSAML_CONF_DIR=//onyx/backend/ee/onyx/configs/saml_config\n\n\n# Generally not useful for dev, we don't generally want to set up an SMTP server\n# for dev.\nREQUIRE_EMAIL_VERIFICATION=False\n\n\n# Set these so if you wipe the DB, you don't end up having to go through the UI\n# every time.\nGEN_AI_API_KEY=\nOPENAI_API_KEY=\n# If answer quality isn't important for dev, use gpt-4o-mini since it's cheaper.\nGEN_AI_MODEL_VERSION=gpt-4o\n\n\n# Python stuff\nPYTHONPATH=../backend\nPYTHONUNBUFFERED=1\n\n\n# Enable the full set of Danswer Enterprise Edition features.\n# NOTE: DO NOT ENABLE THIS UNLESS YOU HAVE A PAID ENTERPRISE LICENSE (or if you\n# are using this for local testing/development).\nENABLE_PAID_ENTERPRISE_EDITION_FEATURES=False\n\n\n# S3 File Store Configuration (MinIO for local development)\nS3_ENDPOINT_URL=http://localhost:9004\nS3_FILE_STORE_BUCKET_NAME=onyx-file-store-bucket\nS3_AWS_ACCESS_KEY_ID=minioadmin\nS3_AWS_SECRET_ACCESS_KEY=minioadmin\n\n\n# Show extra/uncommon connectors.\nSHOW_EXTRA_CONNECTORS=True\n\n\n# Local langsmith tracing\nLANGSMITH_TRACING=\"true\"\nLANGSMITH_ENDPOINT=\"https://api.smith.langchain.com\"\nLANGSMITH_API_KEY=\nLANGSMITH_PROJECT=\n\n\n# Local Confluence OAuth testing\n# OAUTH_CONFLUENCE_CLOUD_CLIENT_ID=\n# OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET=\n# NEXT_PUBLIC_TEST_ENV=True\n\n\n# OpenSearch\n# Arbitrary password is fine for local development.\nOPENSEARCH_INITIAL_ADMIN_PASSWORD=\n" }, { "path": ".vscode/launch.json", "content": "{\n // Use IntelliSense to learn about possible attributes.\n // Hover to view descriptions of existing attributes.\n // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387\n \"version\": \"0.2.0\",\n \"compounds\": [\n {\n // Dummy entry used to label the group\n \"name\": \"--- Compound ---\",\n \"configurations\": [\"--- Individual ---\"],\n \"presentation\": {\n \"group\": \"1\"\n }\n },\n {\n \"name\": \"Run All Onyx Services\",\n \"configurations\": [\n \"Web Server\",\n \"Model Server\",\n \"API Server\",\n \"MCP Server\",\n \"Slack Bot\",\n \"Celery primary\",\n \"Celery light\",\n \"Celery heavy\",\n \"Celery docfetching\",\n \"Celery docprocessing\",\n \"Celery user_file_processing\",\n \"Celery beat\"\n ],\n \"presentation\": {\n \"group\": \"1\"\n }\n },\n {\n \"name\": \"Web / Model / API\",\n \"configurations\": [\"Web Server\", \"Model Server\", \"API Server\"],\n \"presentation\": {\n \"group\": \"1\"\n }\n },\n {\n \"name\": \"Celery\",\n \"configurations\": [\n \"Celery primary\",\n \"Celery light\",\n \"Celery heavy\",\n \"Celery kg_processing\",\n \"Celery monitoring\",\n \"Celery user_file_processing\",\n \"Celery docfetching\",\n \"Celery docprocessing\",\n \"Celery beat\"\n ],\n \"presentation\": {\n \"group\": \"1\"\n },\n \"stopAll\": true\n }\n ],\n \"configurations\": [\n {\n // Dummy entry used to label the group\n \"name\": \"--- Individual ---\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"presentation\": {\n \"group\": \"2\",\n \"order\": 0\n }\n },\n {\n \"name\": \"Web Server\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"cwd\": \"${workspaceRoot}/web\",\n \"runtimeExecutable\": \"npm\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env.web\",\n \"runtimeArgs\": [\"run\", \"dev\"],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"console\": \"integratedTerminal\",\n \"consoleTitle\": \"Web Server Console\"\n },\n {\n \"name\": \"Model Server\",\n \"consoleName\": \"Model Server\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"uvicorn\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\"\n },\n \"args\": [\"model_server.main:app\", \"--reload\", \"--port\", \"9000\"],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Model Server Console\"\n },\n {\n \"name\": \"API Server\",\n \"consoleName\": \"API Server\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"uvicorn\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\"\n },\n \"args\": [\"onyx.main:app\", \"--reload\", \"--port\", \"8080\"],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"API Server Console\",\n \"justMyCode\": false\n },\n {\n \"name\": \"Slack Bot\",\n \"consoleName\": \"Slack Bot\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"program\": \"onyx/onyxbot/slack/listener.py\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Slack Bot Console\"\n },\n {\n \"name\": \"Discord Bot\",\n \"consoleName\": \"Discord Bot\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"program\": \"onyx/onyxbot/discord/client.py\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Discord Bot Console\"\n },\n {\n \"name\": \"MCP Server\",\n \"consoleName\": \"MCP Server\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"uvicorn\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"MCP_SERVER_ENABLED\": \"true\",\n \"MCP_SERVER_PORT\": \"8090\",\n \"MCP_SERVER_CORS_ORIGINS\": \"http://localhost:*\",\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\"\n },\n \"args\": [\n \"onyx.mcp_server.api:mcp_app\",\n \"--reload\",\n \"--port\",\n \"8090\",\n \"--timeout-graceful-shutdown\",\n \"0\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"MCP Server Console\"\n },\n {\n \"name\": \"Celery primary\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.primary\",\n \"worker\",\n \"--pool=threads\",\n \"--concurrency=4\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=primary@%n\",\n \"-Q\",\n \"celery\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery primary Console\"\n },\n {\n \"name\": \"Celery light\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.light\",\n \"worker\",\n \"--pool=threads\",\n \"--concurrency=64\",\n \"--prefetch-multiplier=8\",\n \"--loglevel=INFO\",\n \"--hostname=light@%n\",\n \"-Q\",\n \"vespa_metadata_sync,connector_deletion,doc_permissions_upsert,index_attempt_cleanup,opensearch_migration\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery light Console\"\n },\n {\n \"name\": \"Celery heavy\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.heavy\",\n \"worker\",\n \"--pool=threads\",\n \"--concurrency=4\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=heavy@%n\",\n \"-Q\",\n \"connector_pruning,connector_doc_permissions_sync,connector_external_group_sync,csv_generation\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery heavy Console\",\n \"justMyCode\": false\n },\n {\n \"name\": \"Celery kg_processing\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.kg_processing\",\n \"worker\",\n \"--pool=threads\",\n \"--concurrency=2\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=kg_processing@%n\",\n \"-Q\",\n \"kg_processing\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery kg_processing Console\"\n },\n {\n \"name\": \"Celery monitoring\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.monitoring\",\n \"worker\",\n \"--pool=threads\",\n \"--concurrency=1\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=monitoring@%n\",\n \"-Q\",\n \"monitoring\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery monitoring Console\"\n },\n {\n \"name\": \"Celery user_file_processing\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.user_file_processing\",\n \"worker\",\n \"--pool=threads\",\n \"--concurrency=2\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=user_file_processing@%n\",\n \"-Q\",\n \"user_file_processing,user_file_project_sync,user_file_delete\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery user_file_processing Console\",\n \"justMyCode\": false\n },\n {\n \"name\": \"Celery docfetching\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.docfetching\",\n \"worker\",\n \"--pool=threads\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=docfetching@%n\",\n \"-Q\",\n \"connector_doc_fetching\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery docfetching Console\",\n \"justMyCode\": false\n },\n {\n \"name\": \"Celery docprocessing\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"ENABLE_MULTIPASS_INDEXING\": \"false\",\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.docprocessing\",\n \"worker\",\n \"--pool=threads\",\n \"--prefetch-multiplier=1\",\n \"--loglevel=INFO\",\n \"--hostname=docprocessing@%n\",\n \"-Q\",\n \"docprocessing\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery docprocessing Console\",\n \"justMyCode\": false\n },\n {\n \"name\": \"Celery beat\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"celery\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-A\",\n \"onyx.background.celery.versioned_apps.beat\",\n \"beat\",\n \"--loglevel=INFO\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Celery beat Console\"\n },\n {\n \"name\": \"Pytest\",\n \"consoleName\": \"Pytest\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"pytest\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"-v\"\n // Specify a specific module/test to run or provide nothing to run all tests\n // \"tests/unit/onyx/llm/answering/test_prune_and_merge.py\"\n ],\n \"presentation\": {\n \"group\": \"2\"\n },\n \"consoleTitle\": \"Pytest Console\"\n },\n {\n // Dummy entry used to label the group\n \"name\": \"--- Tasks ---\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"presentation\": {\n \"group\": \"3\",\n \"order\": 0\n }\n },\n {\n \"name\": \"Clear and Restart External Volumes and Containers\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"bash\",\n \"runtimeArgs\": [\n \"${workspaceFolder}/backend/scripts/restart_containers.sh\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"3\"\n }\n },\n {\n \"name\": \"Eval CLI\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"program\": \"${workspaceFolder}/backend/onyx/evals/eval_cli.py\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"console\": \"integratedTerminal\",\n \"justMyCode\": false,\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"presentation\": {\n \"group\": \"3\"\n },\n \"env\": {\n \"LOG_LEVEL\": \"INFO\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\"--verbose\"],\n \"consoleTitle\": \"Eval CLI Console\"\n },\n {\n // Celery jobs launched through a single background script (legacy)\n // Recommend using the \"Celery (all)\" compound launch instead.\n \"name\": \"Background Jobs\",\n \"consoleName\": \"Background Jobs\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"program\": \"scripts/dev_run_background_jobs.py\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"LOG_LEVEL\": \"DEBUG\",\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n }\n },\n {\n \"name\": \"Install Python Requirements\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"uv\",\n \"runtimeArgs\": [\n \"sync\",\n \"--all-extras\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"3\"\n }\n },\n {\n \"name\": \"Build Sandbox Templates\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"module\": \"onyx.server.features.build.sandbox.build_templates\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.vscode/.env\",\n \"env\": {\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"3\"\n },\n \"consoleTitle\": \"Build Sandbox Templates\"\n },\n {\n // Dummy entry used to label the group\n \"name\": \"--- Database ---\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"presentation\": {\n \"group\": \"4\",\n \"order\": 0\n }\n },\n {\n \"name\": \"Restore seeded database dump\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"uv\",\n \"runtimeArgs\": [\n \"run\",\n \"--with\",\n \"onyx-devtools\",\n \"ods\",\n \"db\",\n \"restore\",\n \"--fetch-seeded\",\n \"--yes\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"4\"\n }\n },\n {\n \"name\": \"Clean restore seeded database dump (destructive)\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"uv\",\n \"runtimeArgs\": [\n \"run\",\n \"--with\",\n \"onyx-devtools\",\n \"ods\",\n \"db\",\n \"restore\",\n \"--fetch-seeded\",\n \"--clean\",\n \"--yes\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"4\"\n }\n },\n {\n \"name\": \"Create database snapshot\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"uv\",\n \"runtimeArgs\": [\n \"run\",\n \"--with\",\n \"onyx-devtools\",\n \"ods\",\n \"db\",\n \"dump\",\n \"backup.dump\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"4\"\n }\n },\n {\n \"name\": \"Clean restore database snapshot (destructive)\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"uv\",\n \"runtimeArgs\": [\n \"run\",\n \"--with\",\n \"onyx-devtools\",\n \"ods\",\n \"db\",\n \"restore\",\n \"--clean\",\n \"--yes\",\n \"backup.dump\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"4\"\n }\n },\n {\n \"name\": \"Upgrade database to head revision\",\n \"type\": \"node\",\n \"request\": \"launch\",\n \"runtimeExecutable\": \"uv\",\n \"runtimeArgs\": [\n \"run\",\n \"--with\",\n \"onyx-devtools\",\n \"ods\",\n \"db\",\n \"upgrade\"\n ],\n \"cwd\": \"${workspaceFolder}\",\n \"console\": \"integratedTerminal\",\n \"presentation\": {\n \"group\": \"4\"\n }\n },\n {\n // script to generate the openapi schema\n \"name\": \"Onyx OpenAPI Schema Generator\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"program\": \"backend/scripts/onyx_openapi_schema.py\",\n \"cwd\": \"${workspaceFolder}\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"env\": {\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \"backend\"\n },\n \"args\": [\"--filename\", \"backend/generated/openapi.json\", \"--generate-python-client\"]\n },\n {\n // script to debug multi tenant db issues\n \"name\": \"Onyx DB Manager (Top Chunks)\",\n \"type\": \"debugpy\",\n \"request\": \"launch\",\n \"program\": \"scripts/debugging/onyx_db.py\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"env\": {\n \"PYTHONUNBUFFERED\": \"1\",\n \"PYTHONPATH\": \".\"\n },\n \"args\": [\n \"--password\",\n \"your_password_here\",\n \"--port\",\n \"5433\",\n \"--report\",\n \"top-chunks\",\n \"--filename\",\n \"generated/tenants_by_num_docs.csv\"\n ]\n },\n {\n \"name\": \"Debug React Web App in Chrome\",\n \"type\": \"chrome\",\n \"request\": \"launch\",\n \"url\": \"http://localhost:3000\",\n \"webRoot\": \"${workspaceFolder}/web\"\n }\n ]\n}\n" }, { "path": ".vscode/tasks.template.jsonc", "content": "{\n \"version\": \"2.0.0\",\n \"tasks\": [\n {\n \"type\": \"austin\",\n \"label\": \"Profile celery beat\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"options\": {\n \"cwd\": \"${workspaceFolder}/backend\"\n },\n \"command\": [\n \"sudo\",\n \"-E\"\n ],\n \"args\": [\n \"celery\",\n \"-A\",\n \"onyx.background.celery.versioned_apps.beat\",\n \"beat\",\n \"--loglevel=INFO\"\n ]\n },\n {\n \"type\": \"shell\",\n \"label\": \"Generate Onyx OpenAPI Python client\",\n \"cwd\": \"${workspaceFolder}/backend\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"options\": {\n \"cwd\": \"${workspaceFolder}/backend\"\n },\n \"command\": [\n \"openapi-generator\"\n ],\n \"args\": [\n \"generate\",\n \"-i\",\n \"generated/openapi.json\",\n \"-g\",\n \"python\",\n \"-o\",\n \"generated/onyx_openapi_client\",\n \"--package-name\",\n \"onyx_openapi_client\",\n ]\n },\n {\n \"type\": \"shell\",\n \"label\": \"Generate Typescript Fetch client (openapi-generator)\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"options\": {\n \"cwd\": \"${workspaceFolder}\"\n },\n \"command\": [\n \"openapi-generator\"\n ],\n \"args\": [\n \"generate\",\n \"-i\",\n \"backend/generated/openapi.json\",\n \"-g\",\n \"typescript-fetch\",\n \"-o\",\n \"${workspaceFolder}/web/src/lib/generated/onyx_api\",\n \"--additional-properties=disallowAdditionalPropertiesIfNotPresent=false,legacyDiscriminatorBehavior=false,supportsES6=true\",\n ]\n },\n {\n \"type\": \"shell\",\n \"label\": \"Generate TypeScript Client (openapi-ts)\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"options\": {\n \"cwd\": \"${workspaceFolder}/web\"\n },\n \"command\": [\n \"npx\"\n ],\n \"args\": [\n \"openapi-typescript\",\n \"../backend/generated/openapi.json\",\n \"--output\",\n \"./src/lib/generated/onyx-schema.ts\",\n ]\n },\n {\n \"type\": \"shell\",\n \"label\": \"Generate TypeScript Client (orval)\",\n \"envFile\": \"${workspaceFolder}/.env\",\n \"options\": {\n \"cwd\": \"${workspaceFolder}/web\"\n },\n \"command\": [\n \"npx\"\n ],\n \"args\": [\n \t\"orval\",\n \"--config\",\n \"orval.config.js\",\n ]\n }\n ]\n}\n" }, { "path": "AGENTS.md", "content": "# PROJECT KNOWLEDGE BASE\n\nThis file provides guidance to AI agents when working with code in this repository.\n\n## KEY NOTES\n\n- If you run into any missing python dependency errors, try running your command with `source .venv/bin/activate` \\\n to assume the python venv.\n- To make tests work, check the `.env` file at the root of the project to find an OpenAI key.\n- If using `playwright` to explore the frontend, you can usually log in with username `a@example.com` and password\n `a`. The app can be accessed at `http://localhost:3000`.\n- You should assume that all Onyx services are running. To verify, you can check the `backend/log` directory to\n make sure we see logs coming out from the relevant service.\n- To connect to the Postgres database, use: `docker exec -it onyx-relational_db-1 psql -U postgres -c \"\"`\n- When making calls to the backend, always go through the frontend. E.g. make a call to `http://localhost:3000/api/persona` not `http://localhost:8080/api/persona`\n- Put ALL db operations under the `backend/onyx/db` / `backend/ee/onyx/db` directories. Don't run queries\n outside of those directories.\n\n## Project Overview\n\n**Onyx** (formerly Danswer) is an open-source Gen-AI and Enterprise Search platform that connects to company documents, apps, and people. It features a modular architecture with both Community Edition (MIT licensed) and Enterprise Edition offerings.\n\n### Background Workers (Celery)\n\nOnyx uses Celery for asynchronous task processing with multiple specialized workers:\n\n#### Worker Types\n\n1. **Primary Worker** (`celery_app.py`)\n - Coordinates core background tasks and system-wide operations\n - Handles connector management, document sync, pruning, and periodic checks\n - Runs with 4 threads concurrency\n - Tasks: connector deletion, vespa sync, pruning, LLM model updates, user file sync\n\n2. **Docfetching Worker** (`docfetching`)\n - Fetches documents from external data sources (connectors)\n - Spawns docprocessing tasks for each document batch\n - Implements watchdog monitoring for stuck connectors\n - Configurable concurrency (default from env)\n\n3. **Docprocessing Worker** (`docprocessing`)\n - Processes fetched documents through the indexing pipeline:\n - Upserts documents to PostgreSQL\n - Chunks documents and adds contextual information\n - Embeds chunks via model server\n - Writes chunks to Vespa vector database\n - Updates document metadata\n - Configurable concurrency (default from env)\n\n4. **Light Worker** (`light`)\n - Handles lightweight, fast operations\n - Tasks: vespa operations, document permissions sync, external group sync\n - Higher concurrency for quick tasks\n\n5. **Heavy Worker** (`heavy`)\n - Handles resource-intensive operations\n - Primary task: document pruning operations\n - Runs with 4 threads concurrency\n\n6. **KG Processing Worker** (`kg_processing`)\n - Handles Knowledge Graph processing and clustering\n - Builds relationships between documents\n - Runs clustering algorithms\n - Configurable concurrency\n\n7. **Monitoring Worker** (`monitoring`)\n - System health monitoring and metrics collection\n - Monitors Celery queues, process memory, and system status\n - Single thread (monitoring doesn't need parallelism)\n - Cloud-specific monitoring tasks\n\n8. **User File Processing Worker** (`user_file_processing`)\n - Processes user-uploaded files\n - Handles user file indexing and project synchronization\n - Configurable concurrency\n\n9. **Beat Worker** (`beat`)\n - Celery's scheduler for periodic tasks\n - Uses DynamicTenantScheduler for multi-tenant support\n - Schedules tasks like:\n - Indexing checks (every 15 seconds)\n - Connector deletion checks (every 20 seconds)\n - Vespa sync checks (every 20 seconds)\n - Pruning checks (every 20 seconds)\n - KG processing (every 60 seconds)\n - Monitoring tasks (every 5 minutes)\n - Cleanup tasks (hourly)\n\n#### Key Features\n\n- **Thread-based Workers**: All workers use thread pools (not processes) for stability\n- **Tenant Awareness**: Multi-tenant support with per-tenant task isolation. There is a\n middleware layer that automatically finds the appropriate tenant ID when sending tasks\n via Celery Beat.\n- **Task Prioritization**: High, Medium, Low priority queues\n- **Monitoring**: Built-in heartbeat and liveness checking\n- **Failure Handling**: Automatic retry and failure recovery mechanisms\n- **Redis Coordination**: Inter-process communication via Redis\n- **PostgreSQL State**: Task state and metadata stored in PostgreSQL\n\n#### Important Notes\n\n**Defining Tasks**:\n\n- Always use `@shared_task` rather than `@celery_app`\n- Put tasks under `background/celery/tasks/` or `ee/background/celery/tasks`\n- Never enqueue a task without an expiration. Always supply `expires=` when\n sending tasks, either from the beat schedule or directly from another task. It\n should never be acceptable to submit code which enqueues tasks without an\n expiration, as doing so can lead to unbounded task queue growth.\n\n**Defining APIs**:\nWhen creating new FastAPI APIs, do NOT use the `response_model` field. Instead, just type the\nfunction.\n\n**Testing Updates**:\nIf you make any updates to a celery worker and you want to test these changes, you will need\nto ask me to restart the celery worker. There is no auto-restart on code-change mechanism.\n\n**Task Time Limits**:\nSince all tasks are executed in thread pools, the time limit features of Celery are silently \ndisabled and won't work. Timeout logic must be implemented within the task itself.\n\n### Code Quality\n\n```bash\n# Install and run pre-commit hooks\npre-commit install\npre-commit run --all-files\n```\n\nNOTE: Always make sure everything is strictly typed (both in Python and Typescript).\n\n## Architecture Overview\n\n### Technology Stack\n\n- **Backend**: Python 3.11, FastAPI, SQLAlchemy, Alembic, Celery\n- **Frontend**: Next.js 15+, React 18, TypeScript, Tailwind CSS\n- **Database**: PostgreSQL with Redis caching\n- **Search**: Vespa vector database\n- **Auth**: OAuth2, SAML, multi-provider support\n- **AI/ML**: LangChain, LiteLLM, multiple embedding models\n\n### Directory Structure\n\n```\nbackend/\n├── onyx/\n│ ├── auth/ # Authentication & authorization\n│ ├── chat/ # Chat functionality & LLM interactions\n│ ├── connectors/ # Data source connectors\n│ ├── db/ # Database models & operations\n│ ├── document_index/ # Vespa integration\n│ ├── federated_connectors/ # External search connectors\n│ ├── llm/ # LLM provider integrations\n│ └── server/ # API endpoints & routers\n├── ee/ # Enterprise Edition features\n├── alembic/ # Database migrations\n└── tests/ # Test suites\n\nweb/\n├── src/app/ # Next.js app router pages\n├── src/components/ # Reusable React components\n└── src/lib/ # Utilities & business logic\n```\n\n## Frontend Standards\n\nFrontend standards for the `web/` and `desktop/` projects live in `web/AGENTS.md`.\n\n## Database & Migrations\n\n### Running Migrations\n\n```bash\n# Standard migrations\nalembic upgrade head\n\n# Multi-tenant (Enterprise)\nalembic -n schema_private upgrade head\n```\n\n### Creating Migrations\n\n```bash\n# Create migration\nalembic revision -m \"description\"\n\n# Multi-tenant migration\nalembic -n schema_private revision -m \"description\"\n```\n\nWrite the migration manually and place it in the file that alembic creates when running the above command.\n\n## Testing Strategy\n\nFirst, you must activate the virtual environment with `source .venv/bin/activate`.\n\nThere are 4 main types of tests within Onyx:\n\n### Unit Tests\n\nThese should not assume any Onyx/external services are available to be called.\nInteractions with the outside world should be mocked using `unittest.mock`. Generally, only\nwrite these for complex, isolated modules e.g. `citation_processing.py`.\n\nTo run them:\n\n```bash\npytest -xv backend/tests/unit\n```\n\n### External Dependency Unit Tests\n\nThese tests assume that all external dependencies of Onyx are available and callable (e.g. Postgres, Redis,\nMinIO/S3, Vespa are running + OpenAI can be called + any request to the internet is fine + etc.).\n\nHowever, the actual Onyx containers are not running and with these tests we call the function to test directly.\nWe can also mock components/calls at will.\n\nThe goal with these tests are to minimize mocking while giving some flexibility to mock things that are flakey,\nneed strictly controlled behavior, or need to have their internal behavior validated (e.g. verify a function is called\nwith certain args, something that would be impossible with proper integration tests).\n\nA great example of this type of test is `backend/tests/external_dependency_unit/connectors/confluence/test_confluence_group_sync.py`.\n\nTo run them:\n\n```bash\npython -m dotenv -f .vscode/.env run -- pytest backend/tests/external_dependency_unit\n```\n\n### Integration Tests\n\nStandard integration tests. Every test in `backend/tests/integration` runs against a real Onyx deployment. We cannot\nmock anything in these tests. Prefer writing integration tests (or External Dependency Unit Tests if mocking/internal\nverification is necessary) over any other type of test.\n\nTests are parallelized at a directory level.\n\nWhen writing integration tests, make sure to check the root `conftest.py` for useful fixtures + the `backend/tests/integration/common_utils` directory for utilities. Prefer (if one exists), calling the appropriate Manager\nclass in the utils over directly calling the APIs with a library like `requests`. Prefer using fixtures rather than\ncalling the utilities directly (e.g. do NOT create admin users with\n`admin_user = UserManager.create(name=\"admin_user\")`, instead use the `admin_user` fixture).\n\nA great example of this type of test is `backend/tests/integration/tests/streaming_endpoints/test_chat_stream.py`.\n\nTo run them:\n\n```bash\npython -m dotenv -f .vscode/.env run -- pytest backend/tests/integration\n```\n\n### Playwright (E2E) Tests\n\nThese tests are an even more complete version of the Integration Tests mentioned above. Has all services of Onyx\nrunning, _including_ the Web Server.\n\nUse these tests for anything that requires significant frontend <-> backend coordination.\n\nTests are located at `web/tests/e2e`. Tests are written in TypeScript.\n\nTo run them:\n\n```bash\nnpx playwright test \n```\n\nFor shared fixtures, best practices, and detailed guidance, see `backend/tests/README.md`.\n\n## Logs\n\nWhen (1) writing integration tests or (2) doing live tests (e.g. curl / playwright) you can get access\nto logs via the `backend/log/_debug.log` file. All Onyx services (api_server, web_server, celery_X)\nwill be tailing their logs to this file.\n\n## Security Considerations\n\n- Never commit API keys or secrets to repository\n- Use encrypted credential storage for connector credentials\n- Follow RBAC patterns for new features\n- Implement proper input validation with Pydantic models\n- Use parameterized queries to prevent SQL injection\n\n## AI/LLM Integration\n\n- Multiple LLM providers supported via LiteLLM\n- Configurable models per feature (chat, search, embeddings)\n- Streaming support for real-time responses\n- Token management and rate limiting\n- Custom prompts and agent actions\n\n## Creating a Plan\n\nWhen creating a plan in the `plans` directory, make sure to include at least these elements:\n\n**Issues to Address**\nWhat the change is meant to do.\n\n**Important Notes**\nThings you come across in your research that are important to the implementation.\n\n**Implementation strategy**\nHow you are going to make the changes happen. High level approach.\n\n**Tests**\nWhat unit (use rarely), external dependency unit, integration, and playwright tests you plan to write to\nverify the correct behavior. Don't overtest. Usually, a given change only needs one type of test.\n\nDo NOT include these: _Timeline_, _Rollback plan_\n\nThis is a minimal list - feel free to include more. Do NOT write code as part of your plan.\nKeep it high level. You can reference certain files or functions though.\n\nBefore writing your plan, make sure to do research. Explore the relevant sections in the codebase.\n\n## Error Handling\n\n**Always raise `OnyxError` from `onyx.error_handling.exceptions` instead of `HTTPException`.\nNever hardcode status codes or use `starlette.status` / `fastapi.status` constants directly.**\n\nA global FastAPI exception handler converts `OnyxError` into a JSON response with the standard\n`{\"error_code\": \"...\", \"detail\": \"...\"}` shape. This eliminates boilerplate and keeps error\nhandling consistent across the entire backend.\n\n```python\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\n\n# ✅ Good\nraise OnyxError(OnyxErrorCode.NOT_FOUND, \"Session not found\")\n\n# ✅ Good — no extra message needed\nraise OnyxError(OnyxErrorCode.UNAUTHENTICATED)\n\n# ✅ Good — upstream service with dynamic status code\nraise OnyxError(OnyxErrorCode.BAD_GATEWAY, detail, status_code_override=upstream_status)\n\n# ❌ Bad — using HTTPException directly\nraise HTTPException(status_code=404, detail=\"Session not found\")\n\n# ❌ Bad — starlette constant\nraise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=\"Access denied\")\n```\n\nAvailable error codes are defined in `backend/onyx/error_handling/error_codes.py`. If a new error\ncategory is needed, add it there first — do not invent ad-hoc codes.\n\n**Upstream service errors:** When forwarding errors from an upstream service where the HTTP\nstatus code is dynamic (comes from the upstream response), use `status_code_override`:\n\n```python\nraise OnyxError(OnyxErrorCode.BAD_GATEWAY, detail, status_code_override=e.response.status_code)\n```\n\n## Best Practices\n\nIn addition to the other content in this file, best practices for contributing\nto the codebase can be found in the \"Engineering Best Practices\" section of\n`CONTRIBUTING.md`. Understand its contents and follow them.\n" }, { "path": "CONTRIBUTING.md", "content": "# Contributing to Onyx\n\nHey there! We are so excited that you're interested in Onyx.\n\n## Table of Contents\n\n- [Contribution Opportunities](#contribution-opportunities)\n- [Contribution Process](#contribution-process)\n- [Development Setup](#development-setup)\n - [Prerequisites](#prerequisites)\n - [Backend: Python Requirements](#backend-python-requirements)\n - [Frontend: Node Dependencies](#frontend-node-dependencies)\n - [Formatting and Linting](#formatting-and-linting)\n- [Running the Application](#running-the-application)\n - [VSCode Debugger (Recommended)](#vscode-debugger-recommended)\n - [Manually Running for Development](#manually-running-for-development)\n - [Running in Docker](#running-in-docker)\n- [macOS-Specific Notes](#macos-specific-notes)\n- [Engineering Best Practices](#engineering-best-practices)\n - [Principles and Collaboration](#principles-and-collaboration)\n - [Style and Maintainability](#style-and-maintainability)\n - [Performance and Correctness](#performance-and-correctness)\n - [Repository Conventions](#repository-conventions)\n- [Release Process](#release-process)\n- [Getting Help](#getting-help)\n- [Enterprise Edition Contributions](#enterprise-edition-contributions)\n\n---\n\n## Contribution Opportunities\n\nThe [GitHub Issues](https://github.com/onyx-dot-app/onyx/issues) page is a great place to look for and share contribution ideas.\n\nIf you have your own feature that you would like to build, please create an issue and community members can provide feedback and upvote if they feel a common need.\n\n---\n\n## Contribution Process\n\nTo contribute, please follow the\n[\"fork and pull request\"](https://docs.github.com/en/get-started/quickstart/contributing-to-projects) workflow.\n\n### 1. Get the feature or enhancement approved\n\nCreate a GitHub issue and see if there are upvotes. If you feel the feature is sufficiently value-additive and you would like approval to contribute it to the repo, tag [Yuhong](https://github.com/yuhongsun96) to review.\n\nIf you do not get a response within a week, feel free to email yuhong@onyx.app and include the issue in the message.\n\nNot all small features and enhancements will be accepted as there is a balance between feature richness and bloat. We strive to provide the best user experience possible so we have to be intentional about what we include in the app.\n\n### 2. Get the design approved\n\nThe Onyx team will either provide a design doc and PRD for the feature or request one from you, the contributor. The scope and detail of the design will depend on the individual feature.\n\n### 3. IP attribution for EE contributions\n\nIf you are contributing features to Onyx Enterprise Edition, you are required to sign the [IP Assignment Agreement](contributor_ip_assignment/EE_Contributor_IP_Assignment_Agreement.md).\n\n### 4. Review and testing\n\nYour features must pass all tests and all comments must be addressed prior to merging.\n\n### Implicit agreements\n\nIf we approve an issue, we are promising you the following:\n- Your work will receive timely attention and we will put aside other important items to ensure you are not blocked.\n- You will receive necessary coaching on eng quality, system design, etc. to ensure the feature is completed well.\n- The Onyx team will pull resources and bandwidth from design, PM, and engineering to ensure that you have all the resources to build the feature to the quality required for merging.\n\nBecause this is a large investment from our team, we ask that you:\n- Thoroughly read all the requirements of the design docs, engineering best practices, and try to minimize overhead for the Onyx team.\n- Complete the feature in a timely manner to reduce context switching and an ongoing resource pull from the Onyx team.\n\n---\n\n## Development Setup\n\nOnyx being a fully functional app, relies on some external software, specifically:\n\n- [Postgres](https://www.postgresql.org/) (Relational DB)\n- [OpenSearch](https://opensearch.org/) (Vector DB/Search Engine)\n- [Redis](https://redis.io/) (Cache)\n- [MinIO](https://min.io/) (File Store)\n- [Nginx](https://nginx.org/) (Not needed for development flows generally)\n\n> **Note:**\n> This guide provides instructions to build and run Onyx locally from source with Docker containers providing the above external software.\n> We believe this combination is easier for development purposes. If you prefer to use pre-built container images, see [Running in Docker](#running-in-docker) below.\n\n### Prerequisites\n\n- **Python 3.11** — If using a lower version, modifications will have to be made to the code. Higher versions may have library compatibility issues.\n- **Docker** — Required for running external services (Postgres, OpenSearch, Redis, MinIO).\n- **Node.js v22** — We recommend using [nvm](https://github.com/nvm-sh/nvm) to manage Node installations.\n\n### Backend: Python Requirements\n\nWe use [uv](https://docs.astral.sh/uv/) and recommend creating a [virtual environment](https://docs.astral.sh/uv/pip/environments/#using-a-virtual-environment).\n\n```bash\nuv venv .venv --python 3.11\nsource .venv/bin/activate\n```\n\n_For Windows, activate the virtual environment using Command Prompt:_\n\n```bash\n.venv\\Scripts\\activate\n```\n\nIf using PowerShell, the command slightly differs:\n\n```powershell\n.venv\\Scripts\\Activate.ps1\n```\n\nInstall the required Python dependencies:\n\n```bash\nuv sync --all-extras\n```\n\nInstall Playwright for Python (headless browser required by the Web Connector):\n\n```bash\nuv run playwright install\n```\n\n### Frontend: Node Dependencies\n\n```bash\nnvm install 22 && nvm use 22\nnode -v # verify your active version\n```\n\nNavigate to `onyx/web` and run:\n\n```bash\nnpm i\n```\n\n### Formatting and Linting\n\n#### Backend\n\nSet up pre-commit hooks (black / reorder-python-imports):\n\n```bash\nuv run pre-commit install\n```\n\nWe also use `mypy` for static type checking. Onyx is fully type-annotated, and we want to keep it that way! To run the mypy checks manually:\n\n```bash\nuv run mypy . # from onyx/backend\n```\n\n#### Frontend\n\nWe use `prettier` for formatting. The desired version will be installed via `npm i` from the `onyx/web` directory. To run the formatter:\n\n```bash\nnpx prettier --write . # from onyx/web\n```\n\nPre-commit will also run prettier automatically on files you've recently touched. If re-formatted, your commit will fail. Re-stage your changes and commit again.\n\n---\n\n## Running the Application\n\n### VSCode Debugger (Recommended)\n\nWe highly recommend using VSCode's debugger for development.\n\n#### Initial Setup\n\n1. Copy `.vscode/env_template.txt` to `.vscode/.env`\n2. Fill in the necessary environment variables in `.vscode/.env`\n\n#### Using the Debugger\n\nBefore starting, make sure the Docker Daemon is running.\n\n1. Open the Debug view in VSCode (Cmd+Shift+D on macOS)\n2. From the dropdown at the top, select \"Clear and Restart External Volumes and Containers\" and press the green play button\n3. From the dropdown at the top, select \"Run All Onyx Services\" and press the green play button\n4. Navigate to http://localhost:3000 in your browser to start using the app\n5. Set breakpoints by clicking to the left of line numbers to help debug while the app is running\n6. Use the debug toolbar to step through code, inspect variables, etc.\n\n> **Note:** \"Clear and Restart External Volumes and Containers\" will reset your Postgres and OpenSearch (relational-db and index). Only run this if you are okay with wiping your data.\n\n**Features:**\n- Hot reload is enabled for the web server and API servers\n- Python debugging is configured with debugpy\n- Environment variables are loaded from `.vscode/.env`\n- Console output is organized in the integrated terminal with labeled tabs\n\n### Manually Running for Development\n\n#### Docker containers for external software\n\nYou will need Docker installed to run these containers.\n\nNavigate to `onyx/deployment/docker_compose`, then start up Postgres/OpenSearch/Redis/MinIO with:\n\n```bash\ndocker compose -f docker-compose.yml -f docker-compose.dev.yml up -d index relational_db cache minio\n```\n\n(index refers to OpenSearch, relational_db refers to Postgres, and cache refers to Redis)\n\n#### Running Onyx locally\n\nTo start the frontend, navigate to `onyx/web` and run:\n\n```bash\nnpm run dev\n```\n\nNext, start the model server which runs the local NLP models. Navigate to `onyx/backend` and run:\n\n```bash\nuvicorn model_server.main:app --reload --port 9000\n```\n\n_For Windows (for compatibility with both PowerShell and Command Prompt):_\n\n```bash\npowershell -Command \"uvicorn model_server.main:app --reload --port 9000\"\n```\n\nThe first time running Onyx, you will need to run the DB migrations for Postgres. After the first time, this is no longer required unless the DB models change.\n\nNavigate to `onyx/backend` and with the venv active, run:\n\n```bash\nalembic upgrade head\n```\n\nNext, start the task queue which orchestrates the background jobs. Still in `onyx/backend`, run:\n\n```bash\npython ./scripts/dev_run_background_jobs.py\n```\n\nTo run the backend API server, navigate back to `onyx/backend` and run:\n\n```bash\nAUTH_TYPE=basic uvicorn onyx.main:app --reload --port 8080\n```\n\n_For Windows (for compatibility with both PowerShell and Command Prompt):_\n\n```bash\npowershell -Command \"\n $env:AUTH_TYPE='basic'\n uvicorn onyx.main:app --reload --port 8080\n\"\n```\n\n> **Note:** If you need finer logging, add the additional environment variable `LOG_LEVEL=DEBUG` to the relevant services.\n\n#### Wrapping up\n\nYou should now have 4 servers running:\n\n- Web server\n- Backend API\n- Model server\n- Background jobs\n\nNow, visit http://localhost:3000 in your browser. You should see the Onyx onboarding wizard where you can connect your external LLM provider to Onyx.\n\nYou've successfully set up a local Onyx instance!\n\n### Running in Docker\n\nYou can run the full Onyx application stack from pre-built images including all external software dependencies.\n\nNavigate to `onyx/deployment/docker_compose` and run:\n\n```bash\ndocker compose up -d\n```\n\nAfter Docker pulls and starts these containers, navigate to http://localhost:3000 to use Onyx.\n\nIf you want to make changes to Onyx and run those changes in Docker, you can also build a local version of the Onyx container images that incorporates your changes:\n\n```bash\ndocker compose up -d --build\n```\n\n---\n\n## macOS-Specific Notes\n\n### Setting up Python\n\nEnsure [Homebrew](https://brew.sh/) is already set up, then install Python 3.11:\n\n```bash\nbrew install python@3.11\n```\n\nAdd Python 3.11 to your path by adding the following line to `~/.zshrc`:\n\n```\nexport PATH=\"$(brew --prefix)/opt/python@3.11/libexec/bin:$PATH\"\n```\n\n> **Note:** You will need to open a new terminal for the path change above to take effect.\n\n### Setting up Docker\n\nOn macOS, you will need to install [Docker Desktop](https://www.docker.com/products/docker-desktop/) and ensure it is running before continuing with the docker commands.\n\n### Formatting and Linting\n\nmacOS will likely require you to remove some quarantine attributes on some of the hooks for them to execute properly. After installing pre-commit, run the following command:\n\n```bash\nsudo xattr -r -d com.apple.quarantine ~/.cache/pre-commit\n```\n\n---\n\n## Engineering Best Practices\n\n> These are also what we adhere to as a team internally, we love to build in the open and to uplevel our community and each other through being transparent.\n\n### Principles and Collaboration\n\n- **Use 1-way vs 2-way doors.** For 2-way doors, move faster and iterate. For 1-way doors, be more deliberate.\n- **Consistency > being \"right.\"** Prefer consistent patterns across the codebase. If something is truly bad, fix it everywhere.\n- **Fix what you touch (selectively).**\n - Don't feel obligated to fix every best-practice issue you notice.\n - Don't introduce new bad practices.\n - If your change touches code that violates best practices, fix it as part of the change.\n- **Don't tack features on.** When adding functionality, restructure logically as needed to avoid muddying interfaces and accumulating tech debt.\n\n### Style and Maintainability\n\n#### Comments and readability\nAdd clear comments:\n- At logical boundaries (e.g., interfaces) so the reader doesn't need to dig 10 layers deeper.\n- Wherever assumptions are made or something non-obvious/unexpected is done.\n- For complicated flows/functions.\n- Wherever it saves time (e.g., nontrivial regex patterns).\n\n#### Errors and exceptions\n- **Fail loudly** rather than silently skipping work.\n - Example: raise and let exceptions propagate instead of silently dropping a document.\n- **Don't overuse `try/except`.**\n - Put `try/except` at the correct logical level.\n - Do not mask exceptions unless it is clearly appropriate.\n\n#### Typing\n- Everything should be **as strictly typed as possible**.\n- Use `cast` for annoying/loose-typed interfaces (e.g., results of `run_functions_tuples_in_parallel`).\n - Only `cast` when the type checker sees `Any` or types are too loose.\n- Prefer types that are easy to read.\n - Avoid dense types like `dict[tuple[str, str], list[list[float]]]`.\n - Prefer domain models, e.g.:\n - `EmbeddingModel(provider_name, model_name)` as a Pydantic model\n - `dict[EmbeddingModel, list[EmbeddingVector]]`\n\n#### State, objects, and boundaries\n- Keep **clear logical boundaries** for state containers and objects.\n- A **config** object should never contain things like a `db_session`.\n- Avoid state containers that are overly nested, or huge + flat (use judgment).\n- Prefer **composition and functional style** over inheritance/OOP.\n- Prefer **no mutation** unless there's a strong reason.\n- State objects should be **intentional and explicit**, ideally nonmutating.\n- Use interfaces/objects to create clear separation of responsibility.\n- Prefer simplicity when there's no clear gain.\n - Avoid overcomplicated mechanisms like semaphores.\n - Prefer **hash maps (dicts)** over tree structures unless there's a strong reason.\n\n#### Naming\n- Name variables carefully and intentionally.\n- Prefer long, explicit names when undecided.\n- Avoid single-character variables except for small, self-contained utilities (or not at all).\n- Keep the same object/name consistent through the call stack and within functions when reasonable.\n - Good: `for token in tokens:`\n - Bad: `for msg in tokens:` (if iterating tokens)\n- Function names should bias toward **long + descriptive** for codebase search.\n - IntelliSense can miss call sites; search works best with unique names.\n\n#### Correctness by construction\n- Prefer self-contained correctness — don't rely on callers to \"use it right\" if you can make misuse hard.\n- Avoid redundancies: if a function takes an arg, it shouldn't also take a state object that contains that same arg.\n- No dead code (unless there's a very good reason).\n- No commented-out code in main or feature branches (unless there's a very good reason).\n- No duplicate logic:\n - Don't copy/paste into branches when shared logic can live above the conditional.\n - If you're afraid to touch the original, you don't understand it well enough.\n - LLMs often create subtle duplicate logic — review carefully and remove it.\n - Avoid \"nearly identical\" objects that confuse when to use which.\n- Avoid extremely long functions with chained logic:\n - Encapsulate steps into helpers for readability, even if not reused.\n - \"Pythonic\" multi-step expressions are OK in moderation; don't trade clarity for cleverness.\n\n### Performance and Correctness\n\n- Avoid holding resources for extended periods (DB sessions, locks/semaphores).\n- Validate objects on creation and right before use.\n- Connector code (data to Onyx documents):\n - Any in-memory structure that can grow without bound based on input must be periodically size-checked.\n - If a connector is OOMing (often shows up as \"missing celery tasks\"), this is a top thing to check retroactively.\n- Async and event loops:\n - Never introduce new async/event loop Python code, and try to make existing async code synchronous when possible if it makes sense.\n - Writing async code without 100% understanding the code and having a concrete reason to do so is likely to introduce bugs and not add any meaningful performance gains.\n\n### Repository Conventions\n\n#### Where code lives\n- Pydantic + data models: `models.py` files.\n- DB interface functions (excluding lazy loading): `db/` directory.\n- LLM prompts: `prompts/` directory, roughly mirroring the code layout that uses them.\n- API routes: `server/` directory.\n\n#### Pydantic and modeling\n- Prefer **Pydantic** over dataclasses.\n- If absolutely required, use `allow_arbitrary_types`.\n\n#### Data conventions\n- Prefer explicit `None` over sentinel empty strings (usually; depends on intent).\n- Prefer explicit identifiers: use string enums instead of integer codes.\n- Avoid magic numbers (co-location is good when necessary). **Always avoid magic strings.**\n\n#### Logging\n- Log messages where they are created.\n- Don't propagate log messages around just to log them elsewhere.\n\n#### Encapsulation\n- Don't use private attributes/methods/properties from other classes/modules.\n- \"Private\" is private — respect that boundary.\n\n#### SQLAlchemy guidance\n- Lazy loading is often bad at scale, especially across multiple list relationships.\n- Be careful when accessing SQLAlchemy object attributes:\n - It can help avoid redundant DB queries,\n - but it can also fail if accessed outside an active session,\n - and lazy loading can add hidden DB dependencies to otherwise \"simple\" functions.\n- Reference: https://www.reddit.com/r/SQLAlchemy/comments/138f248/joinedload_vs_selectinload/\n\n#### Trunk-based development and feature flags\n- **PRs should contain no more than 500 lines of real change.**\n- **Merge to main frequently.** Avoid long-lived feature branches — they create merge conflicts and integration pain.\n- **Use feature flags for incremental rollout.**\n - Large features should be merged in small, shippable increments behind a flag.\n - This allows continuous integration without exposing incomplete functionality.\n- **Keep flags short-lived.** Once a feature is fully rolled out, remove the flag and dead code paths promptly.\n- **Flag at the right level.** Prefer flagging at API/UI entry points rather than deep in business logic.\n- **Test both flag states.** Ensure the codebase works correctly with the flag on and off.\n\n#### Miscellaneous\n- Any TODOs you add in the code must be accompanied by either the name/username of the owner of that TODO, or an issue number for an issue referencing that piece of work.\n- Avoid module-level logic that runs on import, which leads to import-time side effects. Essentially every piece of meaningful logic should exist within some function that has to be explicitly invoked. Acceptable exceptions may include loading environment variables or setting up loggers.\n - If you find yourself needing something like this, you may want that logic to exist in a file dedicated for manual execution (contains `if __name__ == \"__main__\":`) which should not be imported by anything else.\n- Do not conflate Python scripts you intend to run from the command line (contains `if __name__ == \"__main__\":`) with modules you intend to import from elsewhere. If for some unlikely reason they have to be the same file, any logic specific to executing the file (including imports) should be contained in the `if __name__ == \"__main__\":` block.\n - Generally these executable files exist in `backend/scripts/`.\n\n---\n\n## Release Process\n\nOnyx loosely follows the SemVer versioning standard.\nA set of Docker containers will be pushed automatically to DockerHub with every tag.\nYou can see the containers [here](https://hub.docker.com/search?q=onyx%2F).\n\n---\n\n## Getting Help\n\nWe have support channels and generally interesting discussions on our [Discord](https://discord.gg/4NA5SbzrWb).\n\nSee you there!\n\n---\n\n## Enterprise Edition Contributions\n\nIf you are contributing features to Onyx Enterprise Edition (code under any `ee/` directory), you are required to sign the [IP Assignment Agreement](contributor_ip_assignment/EE_Contributor_IP_Assignment_Agreement.md) ([PDF version](contributor_ip_assignment/EE_Contributor_IP_Assignment_Agreement.pdf)).\n" }, { "path": "LICENSE", "content": "Copyright (c) 2023-present DanswerAI, Inc.\n\nPortions of this software are licensed as follows:\n\n- All content that resides under \"ee\" directories of this repository is licensed under the Onyx Enterprise License. Each ee directory contains an identical copy of this license at its root:\n - backend/ee/LICENSE\n - web/src/app/ee/LICENSE\n - web/src/ee/LICENSE\n- All third party components incorporated into the Onyx Software are licensed under the original license provided by the owner of the applicable component.\n- Content outside of the above mentioned directories or restrictions above is available under the \"MIT Expat\" license as defined below.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "\n\n

\n \n

\n\n

\n \n \"Discord\"\n \n \n \"Documentation\"\n \n \n \"Documentation\"\n \n \n \"License\"\n \n

\n\n

\n \n \"onyx-dot-app/onyx\n \n

\n\n# Onyx - The Open Source AI Platform\n\n**[Onyx](https://www.onyx.app/?utm_source=onyx_repo&utm_medium=github&utm_campaign=readme)** is the application layer for LLMs - bringing a feature-rich interface that can be easily hosted by anyone.\nOnyx enables LLMs through advanced capabilities like RAG, web search, code execution, file creation, deep research and more.\n\nConnect your applications with over 50+ indexing based connectors provided out of the box or via MCP.\n\n> [!TIP]\n> Deploy with a single command:\n> ```\n> curl -fsSL https://onyx.app/install_onyx.sh | bash\n> ```\n\n![Onyx Chat Silent Demo](https://github.com/onyx-dot-app/onyx/releases/download/v3.0.0/Onyx.gif)\n\n---\n\n## ⭐ Features\n\n- **🔍 Agentic RAG:** Get best in class search and answer quality based on hybrid index + AI Agents for information retrieval\n - Benchmark to release soon!\n- **🔬 Deep Research:** Get in depth reports with a multi-step research flow.\n - Top of [leaderboard](https://github.com/onyx-dot-app/onyx_deep_research_bench) as of Feb 2026.\n- **🤖 Custom Agents:** Build AI Agents with unique instructions, knowledge, and actions.\n- **🌍 Web Search:** Browse the web to get up to date information.\n - Supports Serper, Google PSE, Brave, SearXNG, and others.\n - Comes with an in house web crawler and support for Firecrawl/Exa.\n- **📄 Artifacts:** Generate documents, graphics, and other downloadable artifacts.\n- **▶️ Actions & MCP:** Let Onyx agents interact with external applications, comes with flexible Auth options.\n- **💻 Code Execution:** Execute code in a sandbox to analyze data, render graphs, or modify files.\n- **🎙️ Voice Mode:** Chat with Onyx via text-to-speech and speech-to-text.\n- **🎨 Image Generation:** Generate images based on user prompts.\n\nOnyx supports all major LLM providers, both self-hosted (like Ollama, LiteLLM, vLLM, etc.) and proprietary (like Anthropic, OpenAI, Gemini, etc.).\n\nTo learn more - check out our [docs](https://docs.onyx.app/welcome?utm_source=onyx_repo&utm_medium=github&utm_campaign=readme)!\n\n---\n\n## 🚀 Deployment Modes\n\n> Onyx supports deployments in Docker, Kubernetes, Helm/Terraform and provides guides for major cloud providers.\n> Detailed deployment guides found [here](https://docs.onyx.app/deployment/overview).\n\nOnyx supports two separate deployment options: standard and lite.\n\n#### Onyx Lite\n\nThe Lite mode can be thought of as a lightweight Chat UI. It requires less resources (under 1GB memory) and runs a less complex stack.\nIt is great for users who want to test out Onyx quickly or for teams who are only interested in the Chat UI and Agents functionalities.\n\n#### Standard Onyx\n\nThe complete feature set of Onyx which is recommended for serious users and larger teams. Additional components not included in Lite mode:\n- Vector + Keyword index for RAG.\n- Background containers to run job queues and workers for syncing knowledge from connectors.\n- AI model inference servers to run deep learning models used during indexing and inference.\n- Performance optimizations for large scale use via in memory cache (Redis) and blob store (MinIO).\n\n> [!TIP] \n> **To try Onyx for free without deploying, visit [Onyx Cloud](https://cloud.onyx.app/signup?utm_source=onyx_repo&utm_medium=github&utm_campaign=readme)**.\n\n---\n\n## 🏢 Onyx for Enterprise\n\nOnyx is built for teams of all sizes, from individual users to the largest global enterprises:\n- 👥 Collaboration: Share chats and agents with other members of your organization.\n- 🔐 Single Sign On: SSO via Google OAuth, OIDC, or SAML. Group syncing and user provisioning via SCIM.\n- 🛡️ Role Based Access Control: RBAC for sensitive resources like access to agents, actions, etc.\n- 📊 Analytics: Usage graphs broken down by teams, LLMs, or agents.\n- 🕵️ Query History: Audit usage to ensure safe adoption of AI in your organization.\n- 💻 Custom code: Run custom code to remove PII, reject sensitive queries, or to run custom analysis.\n- 🎨 Whitelabeling: Customize the look and feel of Onyx with custom naming, icons, banners, and more.\n\n## 📚 Licensing\n\nThere are two editions of Onyx:\n\n- Onyx Community Edition (CE) is available freely under the MIT license and covers all of the core features for Chat, RAG, Agents, and Actions.\n- Onyx Enterprise Edition (EE) includes extra features that are primarily useful for larger organizations.\n\nFor feature details, check out [our website](https://www.onyx.app/pricing?utm_source=onyx_repo&utm_medium=github&utm_campaign=readme).\n\n## 👪 Community\n\nJoin our open source community on **[Discord](https://discord.gg/TDJ59cGV2X)**!\n\n## 💡 Contributing\n\nLooking to contribute? Please check out the [Contribution Guide](CONTRIBUTING.md) for more details.\n" }, { "path": "backend/.dockerignore", "content": "**/__pycache__\nvenv/\nenv/\n*.egg-info\n.cache\n.git/\n.svn/\n.vscode/\n.idea/\n*.log\nlog/\n.env\nsecrets.yaml\nbuild/\ndist/\n.coverage\nhtmlcov/\nmodel_server/legacy/\n\n# Craft: demo_data directory should be unzipped at container startup, not copied\n**/demo_data/\n# Craft: templates/outputs/venv is created at container startup\n**/templates/outputs/venv\n" }, { "path": "backend/.gitignore", "content": "__pycache__/\n.mypy_cache\n.idea/\nsite_crawls/\n.ipynb_checkpoints/\napi_keys.py\n*ipynb\n.env*\nvespa-app.zip\ndynamic_config_storage/\ncelerybeat-schedule*\nonyx/connectors/salesforce/data/\n.test.env\n/generated\n" }, { "path": "backend/.trivyignore", "content": "# https://github.com/madler/zlib/issues/868\n# Pulled in with base Debian image, it's part of the contrib folder but unused\n# zlib1g is fine\n# Will be gone with Debian image upgrade\n# No impact in our settings\nCVE-2023-45853\n\n# krb5 related, worst case is denial of service by resource exhaustion\n# Accept the risk\nCVE-2024-26458\nCVE-2024-26461\nCVE-2024-26462\nCVE-2024-26458\nCVE-2024-26461\nCVE-2024-26462\nCVE-2024-26458\nCVE-2024-26461\nCVE-2024-26462\nCVE-2024-26458\nCVE-2024-26461\nCVE-2024-26462\n\n# Specific to Firefox which we do not use\n# No impact in our settings\nCVE-2024-0743\n\n# bind9 related, worst case is denial of service by CPU resource exhaustion\n# Accept the risk\nCVE-2023-50387\nCVE-2023-50868\nCVE-2023-50387\nCVE-2023-50868\n\n# libexpat1, XML parsing resource exhaustion\n# We don't parse any user provided XMLs\n# No impact in our settings\nCVE-2023-52425\nCVE-2024-28757\n\n# libharfbuzz0b, O(n^2) growth, worst case is denial of service\n# Accept the risk\nCVE-2023-25193\n" }, { "path": "backend/Dockerfile", "content": "FROM python:3.11.7-slim-bookworm\n\nLABEL com.danswer.maintainer=\"founders@onyx.app\"\nLABEL com.danswer.description=\"This image is the web/frontend container of Onyx which \\\ncontains code for both the Community and Enterprise editions of Onyx. If you do not \\\nhave a contract or agreement with DanswerAI, you are not permitted to use the Enterprise \\\nEdition features outside of personal development or testing purposes. Please reach out to \\\nfounders@onyx.app for more information. Please visit https://github.com/onyx-dot-app/onyx\"\n\n# Build argument for Craft support (disabled by default)\n# Use --build-arg ENABLE_CRAFT=true to include Node.js and opencode CLI\nARG ENABLE_CRAFT=false\n\n# DO_NOT_TRACK is used to disable telemetry for Unstructured\nENV DANSWER_RUNNING_IN_DOCKER=\"true\" \\\n DO_NOT_TRACK=\"true\" \\\n PLAYWRIGHT_BROWSERS_PATH=\"/app/.cache/ms-playwright\"\n\n# Create non-root user for security best practices\nRUN groupadd -g 1001 onyx && \\\n useradd -u 1001 -g onyx -m -s /bin/bash onyx && \\\n mkdir -p /var/log/onyx && \\\n chmod 755 /var/log/onyx && \\\n chown onyx:onyx /var/log/onyx\n\nCOPY --from=ghcr.io/astral-sh/uv:0.9.9 /uv /uvx /bin/\n\n# Install system dependencies\n# cmake needed for psycopg (postgres)\n# libpq-dev needed for psycopg (postgres)\n# curl included just for users' convenience\n# zip for Vespa step futher down\n# ca-certificates for HTTPS\nRUN apt-get update && \\\n apt-get install -y \\\n cmake \\\n curl \\\n zip \\\n ca-certificates \\\n libgnutls30 \\\n libblkid1 \\\n libmount1 \\\n libsmartcols1 \\\n libuuid1 \\\n libxmlsec1-dev \\\n pkg-config \\\n gcc \\\n nano \\\n vim \\\n # Install procps so kubernetes exec sessions can use ps aux for debugging\n procps \\\n libjemalloc2 \\\n && \\\n rm -rf /var/lib/apt/lists/* && \\\n apt-get clean\n\n# Conditionally install Node.js 20 for Craft (required for Next.js)\n# Only installed when ENABLE_CRAFT=true\nRUN if [ \"$ENABLE_CRAFT\" = \"true\" ]; then \\\n echo \"Installing Node.js 20 for Craft support...\" && \\\n curl -fsSL https://deb.nodesource.com/setup_20.x | bash - && \\\n apt-get install -y nodejs && \\\n rm -rf /var/lib/apt/lists/*; \\\n fi\n\n# Conditionally install opencode CLI for Craft agent functionality\n# Only installed when ENABLE_CRAFT=true\n# TODO: download a specific, versioned release of the opencode CLI\nRUN if [ \"$ENABLE_CRAFT\" = \"true\" ]; then \\\n echo \"Installing opencode CLI for Craft support...\" && \\\n curl -fsSL https://opencode.ai/install | bash; \\\n fi\nENV PATH=\"/root/.opencode/bin:${PATH}\"\n\n# Install Python dependencies\n# Remove py which is pulled in by retry, py is not needed and is a CVE\nCOPY ./requirements/default.txt /tmp/requirements.txt\nCOPY ./requirements/ee.txt /tmp/ee-requirements.txt\nRUN uv pip install --system --no-cache-dir --upgrade \\\n -r /tmp/requirements.txt \\\n -r /tmp/ee-requirements.txt && \\\n pip uninstall -y py && \\\n playwright install chromium && \\\n playwright install-deps chromium && \\\n chown -R onyx:onyx /app && \\\n ln -s /usr/local/bin/supervisord /usr/bin/supervisord && \\\n # Cleanup for CVEs and size reduction\n # https://github.com/tornadoweb/tornado/issues/3107\n # xserver-common and xvfb included by playwright installation but not needed after\n # perl-base is part of the base Python Debian image but not needed for Onyx functionality\n # perl-base could only be removed with --allow-remove-essential\n apt-get update && \\\n apt-get remove -y --allow-remove-essential \\\n perl-base \\\n xserver-common \\\n xvfb \\\n cmake \\\n libldap-2.5-0 \\\n libxmlsec1-dev \\\n pkg-config \\\n gcc && \\\n # Install here to avoid some packages being cleaned up above\n apt-get install -y \\\n libxmlsec1-openssl \\\n # Install postgresql-client for easy manual tests\n postgresql-client && \\\n apt-get autoremove -y && \\\n rm -rf /var/lib/apt/lists/* && \\\n rm -rf ~/.cache/uv /tmp/*.txt && \\\n rm -f /usr/local/lib/python3.11/site-packages/tornado/test/test.key\n\n# Pre-downloading models for setups with limited egress\nRUN python -c \"from tokenizers import Tokenizer; \\\nTokenizer.from_pretrained('nomic-ai/nomic-embed-text-v1')\"\n\n# Pre-downloading NLTK for setups with limited egress\nRUN python -c \"import nltk; \\\n nltk.download('stopwords', quiet=True); \\\n nltk.download('punkt_tab', quiet=True);\"\n# nltk.download('wordnet', quiet=True); introduce this back if lemmatization is needed\n\n# Pre-downloading tiktoken for setups with limited egress\nRUN python -c \"import tiktoken; \\\ntiktoken.get_encoding('cl100k_base')\"\n\n# Set up application files\nWORKDIR /app\n\n# Enterprise Version Files\nCOPY --chown=onyx:onyx ./ee /app/ee\nCOPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf\n\n# Set up application files\nCOPY --chown=onyx:onyx ./onyx /app/onyx\nCOPY --chown=onyx:onyx ./shared_configs /app/shared_configs\nCOPY --chown=onyx:onyx ./alembic /app/alembic\nCOPY --chown=onyx:onyx ./alembic_tenants /app/alembic_tenants\nCOPY --chown=onyx:onyx ./alembic.ini /app/alembic.ini\nCOPY supervisord.conf /usr/etc/supervisord.conf\nCOPY --chown=onyx:onyx ./static /app/static\nCOPY --chown=onyx:onyx ./keys /app/keys\n\n# Escape hatch scripts\nCOPY --chown=onyx:onyx ./scripts/debugging /app/scripts/debugging\nCOPY --chown=onyx:onyx ./scripts/force_delete_connector_by_id.py /app/scripts/force_delete_connector_by_id.py\nCOPY --chown=onyx:onyx ./scripts/supervisord_entrypoint.sh /app/scripts/supervisord_entrypoint.sh\nCOPY --chown=onyx:onyx ./scripts/setup_craft_templates.sh /app/scripts/setup_craft_templates.sh\nCOPY --chown=onyx:onyx ./scripts/reencrypt_secrets.py /app/scripts/reencrypt_secrets.py\nRUN chmod +x /app/scripts/supervisord_entrypoint.sh /app/scripts/setup_craft_templates.sh\n\n# Run Craft template setup at build time when ENABLE_CRAFT=true\n# This pre-bakes demo data, Python venv, and npm dependencies into the image\nRUN if [ \"$ENABLE_CRAFT\" = \"true\" ]; then \\\n echo \"Running Craft template setup at build time...\" && \\\n ENABLE_CRAFT=true /app/scripts/setup_craft_templates.sh; \\\n fi\n\n# Set Craft template paths to the in-image locations\n# These match the paths where setup_craft_templates.sh creates the templates\nENV OUTPUTS_TEMPLATE_PATH=/app/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs\nENV VENV_TEMPLATE_PATH=/app/onyx/server/features/build/sandbox/kubernetes/docker/templates/venv\n\n# Put logo in assets\nCOPY --chown=onyx:onyx ./assets /app/assets\n\nENV PYTHONPATH=/app\n\n# Default ONYX_VERSION, typically overriden during builds by GitHub Actions.\nARG ONYX_VERSION=0.0.0-dev\nENV ONYX_VERSION=${ONYX_VERSION}\n\n# Use jemalloc instead of glibc malloc to reduce memory fragmentation\n# in long-running Python processes (API server, Celery workers).\n# The soname is architecture-independent; the dynamic linker resolves\n# the correct path from standard library directories.\n# Placed after all RUN steps so build-time processes are unaffected.\nENV LD_PRELOAD=libjemalloc.so.2\n\n# Default command which does nothing\n# This container is used by api server and background which specify their own CMD\nCMD [\"tail\", \"-f\", \"/dev/null\"]\n" }, { "path": "backend/Dockerfile.model_server", "content": "# Base stage with dependencies\nFROM python:3.11.7-slim-bookworm AS base\n\nENV DANSWER_RUNNING_IN_DOCKER=\"true\" \\\n HF_HOME=/app/.cache/huggingface\n\nCOPY --from=ghcr.io/astral-sh/uv:0.9.9 /uv /uvx /bin/\n\nRUN mkdir -p /app/.cache/huggingface\n\nCOPY ./requirements/model_server.txt /tmp/requirements.txt\nRUN uv pip install --system --no-cache-dir --upgrade \\\n -r /tmp/requirements.txt && \\\n rm -rf ~/.cache/uv /tmp/*.txt\n\n# Stage for downloading embedding models\nFROM base AS embedding-models\nRUN python -c \"from huggingface_hub import snapshot_download; \\\nsnapshot_download('nomic-ai/nomic-embed-text-v1');\"\n\n# Initialize SentenceTransformer to cache the custom architecture\nRUN python -c \"from sentence_transformers import SentenceTransformer; \\\nSentenceTransformer(model_name_or_path='nomic-ai/nomic-embed-text-v1', trust_remote_code=True);\"\n\n# Final stage - combine all downloads\nFROM base AS final\n\nLABEL com.danswer.maintainer=\"founders@onyx.app\"\nLABEL com.danswer.description=\"This image is for the Onyx model server which runs all of the \\\nAI models for Onyx. This container and all the code is MIT Licensed and free for all to use. \\\nYou can find it at https://hub.docker.com/r/onyx/onyx-model-server. For more details, \\\nvisit https://github.com/onyx-dot-app/onyx.\"\n\n# Create non-root user for security best practices\nRUN groupadd -g 1001 onyx && \\\n useradd -u 1001 -g onyx -m -s /bin/bash onyx && \\\n mkdir -p /var/log/onyx && \\\n chmod 755 /var/log/onyx && \\\n chown onyx:onyx /var/log/onyx\n\n# In case the user has volumes mounted to /app/.cache/huggingface that they've downloaded while\n# running Onyx, move the current contents of the cache folder to a temporary location to ensure\n# it's preserved in order to combine with the user's cache contents\nCOPY --chown=onyx:onyx --from=embedding-models /app/.cache/huggingface /app/.cache/temp_huggingface\n\nWORKDIR /app\n\n# Utils used by model server\nCOPY ./onyx/utils/logger.py /app/onyx/utils/logger.py\nCOPY ./onyx/utils/middleware.py /app/onyx/utils/middleware.py\nCOPY ./onyx/utils/tenant.py /app/onyx/utils/tenant.py\n\n# Place to fetch version information\nCOPY ./onyx/__init__.py /app/onyx/__init__.py\n\n# Shared between Onyx Backend and Model Server\nCOPY ./shared_configs /app/shared_configs\n\n# Model Server main code\nCOPY ./model_server /app/model_server\n\nENV PYTHONPATH=/app\n\n# Default ONYX_VERSION, typically overriden during builds by GitHub Actions.\nARG ONYX_VERSION=0.0.0-dev\nENV ONYX_VERSION=${ONYX_VERSION}\n\nCMD [\"uvicorn\", \"model_server.main:app\", \"--host\", \"0.0.0.0\", \"--port\", \"9000\"]\n" }, { "path": "backend/alembic/README.md", "content": "\n\n# Alembic DB Migrations\n\nThese files are for creating/updating the tables in the Relational DB (Postgres).\nOnyx migrations use a generic single-database configuration with an async dbapi.\n\n## To generate new migrations:\n\nFrom onyx/backend, run:\n`alembic revision -m `\n\nNote: you cannot use the `--autogenerate` flag as the automatic schema parsing does not work.\n\nManually populate the upgrade and downgrade in your new migration.\n\nMore info can be found here: https://alembic.sqlalchemy.org/en/latest/autogenerate.html\n\n## Running migrations\n\nTo run all un-applied migrations:\n`alembic upgrade head`\n\nTo undo migrations:\n`alembic downgrade -X`\nwhere X is the number of migrations you want to undo from the current state\n\n### Multi-tenant migrations\n\nFor multi-tenant deployments, you can use additional options:\n\n**Upgrade all tenants:**\n```bash\nalembic -x upgrade_all_tenants=true upgrade head\n```\n\n**Upgrade specific schemas:**\n```bash\n# Single schema\nalembic -x schemas=tenant_12345678-1234-1234-1234-123456789012 upgrade head\n\n# Multiple schemas (comma-separated)\nalembic -x schemas=tenant_12345678-1234-1234-1234-123456789012,public,another_tenant upgrade head\n```\n\n**Upgrade tenants within an alphabetical range:**\n```bash\n# Upgrade tenants 100-200 when sorted alphabetically (positions 100 to 200)\nalembic -x upgrade_all_tenants=true -x tenant_range_start=100 -x tenant_range_end=200 upgrade head\n\n# Upgrade tenants starting from position 1000 alphabetically\nalembic -x upgrade_all_tenants=true -x tenant_range_start=1000 upgrade head\n\n# Upgrade first 500 tenants alphabetically\nalembic -x upgrade_all_tenants=true -x tenant_range_end=500 upgrade head\n```\n\n**Continue on error (for batch operations):**\n```bash\nalembic -x upgrade_all_tenants=true -x continue=true upgrade head\n```\n\nThe tenant range filtering works by:\n1. Sorting tenant IDs alphabetically\n2. Using 1-based position numbers (1st, 2nd, 3rd tenant, etc.)\n3. Filtering to the specified range of positions\n4. Non-tenant schemas (like 'public') are always included\n" }, { "path": "backend/alembic/env.py", "content": "from typing import Any, Literal\nfrom onyx.db.engine.iam_auth import get_iam_auth_token\nfrom onyx.configs.app_configs import USE_IAM_AUTH\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.configs.app_configs import AWS_REGION_NAME\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom sqlalchemy import event\nfrom sqlalchemy import pool\nfrom sqlalchemy import text\nfrom sqlalchemy.engine.base import Connection\nimport os\nimport ssl\nimport asyncio\nimport logging\nfrom logging.config import fileConfig\n\nfrom alembic import context\nfrom sqlalchemy.ext.asyncio import create_async_engine\nfrom sqlalchemy.sql.schema import SchemaItem\nfrom onyx.configs.constants import SSL_CERT_FILE\nfrom shared_configs.configs import (\n MULTI_TENANT,\n POSTGRES_DEFAULT_SCHEMA,\n TENANT_ID_PREFIX,\n)\nfrom onyx.db.models import Base\nfrom celery.backends.database.session import ResultModelBase # type: ignore\nfrom onyx.db.engine.sql_engine import SqlEngine\n\n# Make sure in alembic.ini [logger_root] level=INFO is set or most logging will be\n# hidden! (defaults to level=WARN)\n\n# Alembic Config object\nconfig = context.config\n\nif config.config_file_name is not None and config.attributes.get(\n \"configure_logger\", True\n):\n # disable_existing_loggers=False prevents breaking pytest's caplog fixture\n # See: https://pytest-alembic.readthedocs.io/en/latest/setup.html#caplog-issues\n fileConfig(config.config_file_name, disable_existing_loggers=False)\n\ntarget_metadata = [Base.metadata, ResultModelBase.metadata]\n\nEXCLUDE_TABLES = {\"kombu_queue\", \"kombu_message\"}\n\nlogger = logging.getLogger(__name__)\n\nssl_context: ssl.SSLContext | None = None\nif USE_IAM_AUTH:\n if not os.path.exists(SSL_CERT_FILE):\n raise FileNotFoundError(f\"Expected {SSL_CERT_FILE} when USE_IAM_AUTH is true.\")\n ssl_context = ssl.create_default_context(cafile=SSL_CERT_FILE)\n\n\ndef include_object(\n object: SchemaItem, # noqa: ARG001\n name: str | None,\n type_: Literal[\n \"schema\",\n \"table\",\n \"column\",\n \"index\",\n \"unique_constraint\",\n \"foreign_key_constraint\",\n ],\n reflected: bool, # noqa: ARG001\n compare_to: SchemaItem | None, # noqa: ARG001\n) -> bool:\n if type_ == \"table\" and name in EXCLUDE_TABLES:\n return False\n return True\n\n\ndef filter_tenants_by_range(\n tenant_ids: list[str], start_range: int | None = None, end_range: int | None = None\n) -> list[str]:\n \"\"\"\n Filter tenant IDs by alphabetical position range.\n\n Args:\n tenant_ids: List of tenant IDs to filter\n start_range: Starting position in alphabetically sorted list (1-based, inclusive)\n end_range: Ending position in alphabetically sorted list (1-based, inclusive)\n\n Returns:\n Filtered list of tenant IDs in their original order\n \"\"\"\n if start_range is None and end_range is None:\n return tenant_ids\n\n # Separate tenant IDs from non-tenant schemas\n tenant_schemas = [tid for tid in tenant_ids if tid.startswith(TENANT_ID_PREFIX)]\n non_tenant_schemas = [\n tid for tid in tenant_ids if not tid.startswith(TENANT_ID_PREFIX)\n ]\n\n # Sort tenant schemas alphabetically.\n # NOTE: can cause missed schemas if a schema is created in between workers\n # fetching of all tenant IDs. We accept this risk for now. Just re-running\n # the migration will fix the issue.\n sorted_tenant_schemas = sorted(tenant_schemas)\n\n # Apply range filtering (0-based indexing)\n start_idx = start_range if start_range is not None else 0\n end_idx = end_range if end_range is not None else len(sorted_tenant_schemas)\n\n # Ensure indices are within bounds\n start_idx = max(0, start_idx)\n end_idx = min(len(sorted_tenant_schemas), end_idx)\n\n # Get the filtered tenant schemas\n filtered_tenant_schemas = sorted_tenant_schemas[start_idx:end_idx]\n\n # Combine with non-tenant schemas and preserve original order\n filtered_tenants = []\n for tenant_id in tenant_ids:\n if tenant_id in filtered_tenant_schemas or tenant_id in non_tenant_schemas:\n filtered_tenants.append(tenant_id)\n\n return filtered_tenants\n\n\ndef get_schema_options() -> (\n tuple[bool, bool, bool, int | None, int | None, list[str] | None]\n):\n x_args_raw = context.get_x_argument()\n x_args = {}\n for arg in x_args_raw:\n if \"=\" in arg:\n key, value = arg.split(\"=\", 1)\n x_args[key.strip()] = value.strip()\n else:\n raise ValueError(f\"Invalid argument: {arg}\")\n\n create_schema = x_args.get(\"create_schema\", \"true\").lower() == \"true\"\n upgrade_all_tenants = x_args.get(\"upgrade_all_tenants\", \"false\").lower() == \"true\"\n\n # continue on error with individual tenant\n # only applies to online migrations\n continue_on_error = x_args.get(\"continue\", \"false\").lower() == \"true\"\n\n # Tenant range filtering\n tenant_range_start = None\n tenant_range_end = None\n\n if \"tenant_range_start\" in x_args:\n try:\n tenant_range_start = int(x_args[\"tenant_range_start\"])\n except ValueError:\n raise ValueError(\n f\"Invalid tenant_range_start value: {x_args['tenant_range_start']}. Must be an integer.\"\n )\n\n if \"tenant_range_end\" in x_args:\n try:\n tenant_range_end = int(x_args[\"tenant_range_end\"])\n except ValueError:\n raise ValueError(\n f\"Invalid tenant_range_end value: {x_args['tenant_range_end']}. Must be an integer.\"\n )\n\n # Validate range\n if tenant_range_start is not None and tenant_range_end is not None:\n if tenant_range_start > tenant_range_end:\n raise ValueError(\n f\"tenant_range_start ({tenant_range_start}) cannot be greater than tenant_range_end ({tenant_range_end})\"\n )\n\n # Specific schema names filtering (replaces both schema_name and the old tenant_ids approach)\n schemas = None\n if \"schemas\" in x_args:\n schema_names_str = x_args[\"schemas\"].strip()\n if schema_names_str:\n # Split by comma and strip whitespace\n schemas = [\n name.strip() for name in schema_names_str.split(\",\") if name.strip()\n ]\n if schemas:\n logger.info(f\"Specific schema names specified: {schemas}\")\n\n # Validate that only one method is used at a time\n range_filtering = tenant_range_start is not None or tenant_range_end is not None\n specific_filtering = schemas is not None and len(schemas) > 0\n\n if range_filtering and specific_filtering:\n raise ValueError(\n \"Cannot use both tenant range filtering (tenant_range_start/tenant_range_end) \"\n \"and specific schema filtering (schemas) at the same time. \"\n \"Please use only one filtering method.\"\n )\n\n if upgrade_all_tenants and specific_filtering:\n raise ValueError(\n \"Cannot use both upgrade_all_tenants=true and schemas at the same time. \"\n \"Use either upgrade_all_tenants=true for all tenants, or schemas for specific schemas.\"\n )\n\n # If any filtering parameters are specified, we're not doing the default single schema migration\n if range_filtering:\n upgrade_all_tenants = True\n\n # Validate multi-tenant requirements\n if MULTI_TENANT and not upgrade_all_tenants and not specific_filtering:\n raise ValueError(\n \"In multi-tenant mode, you must specify either upgrade_all_tenants=true \"\n \"or provide schemas. Cannot run default migration.\"\n )\n\n return (\n create_schema,\n upgrade_all_tenants,\n continue_on_error,\n tenant_range_start,\n tenant_range_end,\n schemas,\n )\n\n\ndef do_run_migrations(\n connection: Connection, schema_name: str, create_schema: bool\n) -> None:\n if create_schema:\n connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS \"{schema_name}\"'))\n\n connection.execute(text(f'SET search_path TO \"{schema_name}\"'))\n\n context.configure(\n connection=connection,\n target_metadata=target_metadata, # type: ignore\n include_object=include_object,\n version_table_schema=schema_name,\n include_schemas=True,\n compare_type=True,\n compare_server_default=True,\n script_location=config.get_main_option(\"script_location\"),\n )\n\n with context.begin_transaction():\n context.run_migrations()\n\n\ndef provide_iam_token_for_alembic(\n dialect: Any, # noqa: ARG001\n conn_rec: Any, # noqa: ARG001\n cargs: Any, # noqa: ARG001\n cparams: Any,\n) -> None:\n if USE_IAM_AUTH:\n # Database connection settings\n region = AWS_REGION_NAME\n host = POSTGRES_HOST\n port = POSTGRES_PORT\n user = POSTGRES_USER\n\n # Get IAM authentication token\n token = get_iam_auth_token(host, port, user, region)\n\n # For Alembic / SQLAlchemy in this context, set SSL and password\n cparams[\"password\"] = token\n cparams[\"ssl\"] = ssl_context\n\n\nasync def run_async_migrations() -> None:\n (\n create_schema,\n upgrade_all_tenants,\n continue_on_error,\n tenant_range_start,\n tenant_range_end,\n schemas,\n ) = get_schema_options()\n\n if not schemas and not MULTI_TENANT:\n schemas = [POSTGRES_DEFAULT_SCHEMA]\n\n # without init_engine, subsequent engine calls fail hard intentionally\n SqlEngine.init_engine(pool_size=20, max_overflow=5)\n\n engine = create_async_engine(\n build_connection_string(),\n poolclass=pool.NullPool,\n )\n\n if USE_IAM_AUTH:\n\n @event.listens_for(engine.sync_engine, \"do_connect\")\n def event_provide_iam_token_for_alembic(\n dialect: Any, conn_rec: Any, cargs: Any, cparams: Any\n ) -> None:\n provide_iam_token_for_alembic(dialect, conn_rec, cargs, cparams)\n\n if schemas:\n # Use specific schema names directly without fetching all tenants\n logger.info(f\"Migrating specific schema names: {schemas}\")\n\n i_schema = 0\n num_schemas = len(schemas)\n for schema in schemas:\n i_schema += 1\n logger.info(\n f\"Migrating schema: index={i_schema} num_schemas={num_schemas} schema={schema}\"\n )\n try:\n async with engine.connect() as connection:\n await connection.run_sync(\n do_run_migrations,\n schema_name=schema,\n create_schema=create_schema,\n )\n await connection.commit()\n except Exception as e:\n logger.error(f\"Error migrating schema {schema}: {e}\")\n if not continue_on_error:\n logger.error(\"--continue=true is not set, raising exception!\")\n raise\n\n logger.warning(\"--continue=true is set, continuing to next schema.\")\n\n elif upgrade_all_tenants:\n tenant_schemas = get_all_tenant_ids()\n\n filtered_tenant_schemas = filter_tenants_by_range(\n tenant_schemas, tenant_range_start, tenant_range_end\n )\n\n if tenant_range_start is not None or tenant_range_end is not None:\n logger.info(\n f\"Filtering tenants by range: start={tenant_range_start}, end={tenant_range_end}\"\n )\n logger.info(\n f\"Total tenants: {len(tenant_schemas)}, Filtered tenants: {len(filtered_tenant_schemas)}\"\n )\n\n i_tenant = 0\n num_tenants = len(filtered_tenant_schemas)\n for schema in filtered_tenant_schemas:\n i_tenant += 1\n logger.info(\n f\"Migrating schema: index={i_tenant} num_tenants={num_tenants} schema={schema}\"\n )\n try:\n async with engine.connect() as connection:\n await connection.run_sync(\n do_run_migrations,\n schema_name=schema,\n create_schema=create_schema,\n )\n await connection.commit()\n except Exception as e:\n logger.error(f\"Error migrating schema {schema}: {e}\")\n if not continue_on_error:\n logger.error(\"--continue=true is not set, raising exception!\")\n raise\n\n logger.warning(\"--continue=true is set, continuing to next schema.\")\n\n else:\n # This should not happen in the new design since we require either\n # upgrade_all_tenants=true or schemas in multi-tenant mode\n # and for non-multi-tenant mode, we should use schemas with the default schema\n raise ValueError(\n \"No migration target specified. Use either upgrade_all_tenants=true for all tenants or schemas for specific schemas.\"\n )\n\n await engine.dispose()\n\n\ndef run_migrations_offline() -> None:\n \"\"\"\n NOTE(rkuo): This generates a sql script that can be used to migrate the database ...\n instead of migrating the db live via an open connection\n\n Not clear on when this would be used by us or if it even works.\n\n If it is offline, then why are there calls to the db engine?\n\n This doesn't really get used when we migrate in the cloud.\"\"\"\n\n logger.info(\"run_migrations_offline starting.\")\n\n # without init_engine, subsequent engine calls fail hard intentionally\n SqlEngine.init_engine(pool_size=20, max_overflow=5)\n\n (\n create_schema,\n upgrade_all_tenants,\n continue_on_error,\n tenant_range_start,\n tenant_range_end,\n schemas,\n ) = get_schema_options()\n url = build_connection_string()\n\n if schemas:\n # Use specific schema names directly without fetching all tenants\n logger.info(f\"Migrating specific schema names: {schemas}\")\n\n for schema in schemas:\n logger.info(f\"Migrating schema: {schema}\")\n context.configure(\n url=url,\n target_metadata=target_metadata, # type: ignore\n literal_binds=True,\n include_object=include_object,\n version_table_schema=schema,\n include_schemas=True,\n script_location=config.get_main_option(\"script_location\"),\n dialect_opts={\"paramstyle\": \"named\"},\n )\n\n with context.begin_transaction():\n context.run_migrations()\n\n elif upgrade_all_tenants:\n engine = create_async_engine(url)\n\n if USE_IAM_AUTH:\n\n @event.listens_for(engine.sync_engine, \"do_connect\")\n def event_provide_iam_token_for_alembic_offline(\n dialect: Any, conn_rec: Any, cargs: Any, cparams: Any\n ) -> None:\n provide_iam_token_for_alembic(dialect, conn_rec, cargs, cparams)\n\n tenant_schemas = get_all_tenant_ids()\n engine.sync_engine.dispose()\n\n filtered_tenant_schemas = filter_tenants_by_range(\n tenant_schemas, tenant_range_start, tenant_range_end\n )\n\n if tenant_range_start is not None or tenant_range_end is not None:\n logger.info(\n f\"Filtering tenants by range: start={tenant_range_start}, end={tenant_range_end}\"\n )\n logger.info(\n f\"Total tenants: {len(tenant_schemas)}, Filtered tenants: {len(filtered_tenant_schemas)}\"\n )\n\n for schema in filtered_tenant_schemas:\n logger.info(f\"Migrating schema: {schema}\")\n context.configure(\n url=url,\n target_metadata=target_metadata, # type: ignore\n literal_binds=True,\n include_object=include_object,\n version_table_schema=schema,\n include_schemas=True,\n script_location=config.get_main_option(\"script_location\"),\n dialect_opts={\"paramstyle\": \"named\"},\n )\n\n with context.begin_transaction():\n context.run_migrations()\n else:\n # This should not happen in the new design\n raise ValueError(\n \"No migration target specified. Use either upgrade_all_tenants=true for all tenants or schemas for specific schemas.\"\n )\n\n\ndef run_migrations_online() -> None:\n \"\"\"Run migrations in 'online' mode.\n\n Supports pytest-alembic by checking for a pre-configured connection\n in context.config.attributes[\"connection\"]. If present, uses that\n connection/engine directly instead of creating a new async engine.\n \"\"\"\n # Check if pytest-alembic is providing a connection/engine\n connectable = context.config.attributes.get(\"connection\", None)\n\n if connectable is not None:\n # pytest-alembic is providing an engine - use it directly\n logger.debug(\"run_migrations_online starting (pytest-alembic mode).\")\n\n # For pytest-alembic, we use the default schema (public)\n schema_name = context.config.attributes.get(\n \"schema_name\", POSTGRES_DEFAULT_SCHEMA\n )\n\n # pytest-alembic passes an Engine, we need to get a connection from it\n with connectable.connect() as connection:\n # Set search path for the schema\n connection.execute(text(f'SET search_path TO \"{schema_name}\"'))\n\n context.configure(\n connection=connection,\n target_metadata=target_metadata, # type: ignore\n include_object=include_object,\n version_table_schema=schema_name,\n include_schemas=True,\n compare_type=True,\n compare_server_default=True,\n script_location=config.get_main_option(\"script_location\"),\n )\n\n with context.begin_transaction():\n context.run_migrations()\n\n # Commit the transaction to ensure changes are visible to next migration\n connection.commit()\n else:\n # Normal operation - use async migrations\n logger.info(\"run_migrations_online starting.\")\n asyncio.run(run_async_migrations())\n\n\nif context.is_offline_mode():\n run_migrations_offline()\nelse:\n run_migrations_online()\n" }, { "path": "backend/alembic/run_multitenant_migrations.py", "content": "#!/usr/bin/env python3\n\"\"\"Parallel Alembic Migration Runner\n\nUpgrades tenant schemas to head in batched, parallel alembic subprocesses.\nEach subprocess handles a batch of schemas (via ``-x schemas=a,b,c``),\nreducing per-process overhead compared to one-schema-per-process.\n\nUsage examples::\n\n # defaults: 6 workers, 50 schemas/batch\n python alembic/run_multitenant_migrations.py\n\n # custom settings\n python alembic/run_multitenant_migrations.py -j 8 -b 100\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport subprocess\nimport sys\nimport threading\nimport time\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom typing import NamedTuple\n\nfrom alembic.config import Config\nfrom alembic.script import ScriptDirectory\n\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.db.engine.tenant_utils import get_schemas_needing_migration\nfrom shared_configs.configs import TENANT_ID_PREFIX\n\n\n# ---------------------------------------------------------------------------\n# Data types\n# ---------------------------------------------------------------------------\n\n\nclass Args(NamedTuple):\n jobs: int\n batch_size: int\n\n\nclass BatchResult(NamedTuple):\n schemas: list[str]\n success: bool\n output: str\n elapsed_sec: float\n\n\n# ---------------------------------------------------------------------------\n# Core functions\n# ---------------------------------------------------------------------------\n\n\ndef run_alembic_for_batch(schemas: list[str]) -> BatchResult:\n \"\"\"Run ``alembic upgrade head`` for a batch of schemas in one subprocess.\n\n If the batch fails, it is automatically retried with ``-x continue=true``\n so that the remaining schemas in the batch still get migrated. The retry\n output (which contains alembic's per-schema error messages) is returned\n for diagnosis.\n \"\"\"\n csv = \",\".join(schemas)\n base_cmd = [\"alembic\", \"-x\", f\"schemas={csv}\"]\n\n start = time.monotonic()\n result = subprocess.run(\n [*base_cmd, \"upgrade\", \"head\"],\n stdout=subprocess.PIPE,\n stderr=subprocess.STDOUT,\n text=True,\n )\n\n if result.returncode == 0:\n elapsed = time.monotonic() - start\n return BatchResult(schemas, True, result.stdout or \"\", elapsed)\n\n # At least one schema failed. Print the initial error output, then\n # re-run with continue=true so the remaining schemas still get migrated.\n if result.stdout:\n print(f\"Initial error output:\\n{result.stdout}\", file=sys.stderr, flush=True)\n print(\n f\"Batch failed (exit {result.returncode}), retrying with 'continue=true'...\",\n file=sys.stderr,\n flush=True,\n )\n\n retry = subprocess.run(\n [*base_cmd, \"-x\", \"continue=true\", \"upgrade\", \"head\"],\n stdout=subprocess.PIPE,\n stderr=subprocess.STDOUT,\n text=True,\n )\n elapsed = time.monotonic() - start\n return BatchResult(schemas, False, retry.stdout or \"\", elapsed)\n\n\ndef get_head_revision() -> str | None:\n \"\"\"Get the head revision from the alembic script directory.\"\"\"\n alembic_cfg = Config(\"alembic.ini\")\n script = ScriptDirectory.from_config(alembic_cfg)\n return script.get_current_head()\n\n\ndef run_migrations_parallel(\n schemas: list[str],\n max_workers: int,\n batch_size: int,\n) -> bool:\n \"\"\"Chunk *schemas* into batches and run them in parallel.\n\n A background monitor thread prints a status line every 60 s listing\n which batches are still in-flight, making it easy to spot hung tenants.\n \"\"\"\n batches = [schemas[i : i + batch_size] for i in range(0, len(schemas), batch_size)]\n total_batches = len(batches)\n print(\n f\"{len(schemas)} schemas in {total_batches} batch(es) with {max_workers} workers (batch size: {batch_size})...\",\n flush=True,\n )\n all_success = True\n\n # Thread-safe tracking of in-flight batches for the monitor thread.\n in_flight: dict[int, list[str]] = {}\n prev_in_flight: set[int] = set()\n lock = threading.Lock()\n stop_event = threading.Event()\n\n def _monitor() -> None:\n \"\"\"Print a status line every 60 s listing batches still in-flight.\n\n Only prints batches that were also present in the previous tick,\n making it easy to spot batches that are stuck.\n \"\"\"\n nonlocal prev_in_flight\n while not stop_event.wait(60):\n with lock:\n if not in_flight:\n prev_in_flight = set()\n continue\n current = set(in_flight)\n stuck = current & prev_in_flight\n prev_in_flight = current\n\n if not stuck:\n continue\n\n schemas = [s for idx in sorted(stuck) for s in in_flight[idx]]\n print(\n f\"⏳ batch(es) still running since last check \"\n f\"({', '.join(str(i + 1) for i in sorted(stuck))}): \"\n + \", \".join(schemas),\n flush=True,\n )\n\n monitor_thread = threading.Thread(target=_monitor, daemon=True)\n monitor_thread.start()\n\n try:\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n\n def _run(batch_idx: int, batch: list[str]) -> BatchResult:\n with lock:\n in_flight[batch_idx] = batch\n print(\n f\"Batch {batch_idx + 1}/{total_batches} started ({len(batch)} schemas): {', '.join(batch)}\",\n flush=True,\n )\n result = run_alembic_for_batch(batch)\n with lock:\n in_flight.pop(batch_idx, None)\n return result\n\n future_to_idx = {\n executor.submit(_run, i, b): i for i, b in enumerate(batches)\n }\n\n for future in as_completed(future_to_idx):\n batch_idx = future_to_idx[future]\n try:\n result = future.result()\n status = \"✓\" if result.success else \"✗\"\n\n print(\n f\"Batch {batch_idx + 1}/{total_batches} \"\n f\"{status} {len(result.schemas)} schemas \"\n f\"in {result.elapsed_sec:.1f}s\",\n flush=True,\n )\n\n if not result.success:\n # Print last 20 lines of retry output for diagnosis\n tail = result.output.strip().splitlines()[-20:]\n for line in tail:\n print(f\" {line}\", flush=True)\n all_success = False\n\n except Exception as e:\n print(\n f\"Batch {batch_idx + 1}/{total_batches} ✗ exception: {e}\",\n flush=True,\n )\n all_success = False\n finally:\n stop_event.set()\n monitor_thread.join(timeout=2)\n\n return all_success\n\n\n# ---------------------------------------------------------------------------\n# CLI\n# ---------------------------------------------------------------------------\n\n\ndef parse_args() -> Args:\n parser = argparse.ArgumentParser(\n description=\"Run alembic migrations for all tenant schemas in parallel\"\n )\n parser.add_argument(\n \"-j\",\n \"--jobs\",\n type=int,\n default=6,\n metavar=\"N\",\n help=\"Number of parallel alembic processes (default: 6)\",\n )\n parser.add_argument(\n \"-b\",\n \"--batch-size\",\n type=int,\n default=50,\n metavar=\"N\",\n help=\"Schemas per alembic process (default: 50)\",\n )\n args = parser.parse_args()\n if args.jobs < 1:\n parser.error(\"--jobs must be >= 1\")\n if args.batch_size < 1:\n parser.error(\"--batch-size must be >= 1\")\n return Args(jobs=args.jobs, batch_size=args.batch_size)\n\n\ndef main() -> int:\n args = parse_args()\n\n head_rev = get_head_revision()\n if head_rev is None:\n print(\"Could not determine head revision.\", file=sys.stderr)\n return 1\n\n with SqlEngine.scoped_engine(pool_size=5, max_overflow=2):\n tenant_ids = get_all_tenant_ids()\n tenant_schemas = [tid for tid in tenant_ids if tid.startswith(TENANT_ID_PREFIX)]\n\n if not tenant_schemas:\n print(\n \"No tenant schemas found. Is MULTI_TENANT=true set?\",\n file=sys.stderr,\n )\n return 1\n\n schemas_to_migrate = get_schemas_needing_migration(tenant_schemas, head_rev)\n\n if not schemas_to_migrate:\n print(\n f\"All {len(tenant_schemas)} tenants are already at head revision ({head_rev}).\"\n )\n return 0\n\n print(\n f\"{len(schemas_to_migrate)}/{len(tenant_schemas)} tenants need migration (head: {head_rev}).\"\n )\n\n success = run_migrations_parallel(\n schemas_to_migrate,\n max_workers=args.jobs,\n batch_size=args.batch_size,\n )\n\n print(f\"\\n{'All migrations successful' if success else 'Some migrations failed'}\")\n return 0 if success else 1\n\n\nif __name__ == \"__main__\":\n raise SystemExit(main())\n" }, { "path": "backend/alembic/script.py.mako", "content": "\"\"\"${message}\n\nRevision ID: ${up_revision}\nRevises: ${down_revision | comma,n}\nCreate Date: ${create_date}\n\n\"\"\"\nfrom alembic import op\nimport sqlalchemy as sa\n${imports if imports else \"\"}\n\n# revision identifiers, used by Alembic.\nrevision = ${repr(up_revision)}\ndown_revision = ${repr(down_revision)}\nbranch_labels = ${repr(branch_labels)}\ndepends_on = ${repr(depends_on)}\n\n\ndef upgrade() -> None:\n ${upgrades if upgrades else \"pass\"}\n\n\ndef downgrade() -> None:\n ${downgrades if downgrades else \"pass\"}\n" }, { "path": "backend/alembic/versions/01f8e6d95a33_populate_flow_mapping_data.py", "content": "\"\"\"Populate flow mapping data\n\nRevision ID: 01f8e6d95a33\nRevises: d5c86e2c6dc6\nCreate Date: 2026-01-31 17:37:10.485558\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"01f8e6d95a33\"\ndown_revision = \"d5c86e2c6dc6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add each model config to the conversation flow, setting the global default if it exists\n # Exclude models that are part of ImageGenerationConfig\n op.execute(\n \"\"\"\n INSERT INTO llm_model_flow (llm_model_flow_type, is_default, model_configuration_id)\n SELECT\n 'CHAT' AS llm_model_flow_type,\n COALESCE(\n (lp.is_default_provider IS TRUE AND lp.default_model_name = mc.name),\n FALSE\n ) AS is_default,\n mc.id AS model_configuration_id\n FROM model_configuration mc\n LEFT JOIN llm_provider lp\n ON lp.id = mc.llm_provider_id\n WHERE NOT EXISTS (\n SELECT 1 FROM image_generation_config igc\n WHERE igc.model_configuration_id = mc.id\n );\n \"\"\"\n )\n\n # Add models with supports_image_input to the vision flow\n op.execute(\n \"\"\"\n INSERT INTO llm_model_flow (llm_model_flow_type, is_default, model_configuration_id)\n SELECT\n 'VISION' AS llm_model_flow_type,\n COALESCE(\n (lp.is_default_vision_provider IS TRUE AND lp.default_vision_model = mc.name),\n FALSE\n ) AS is_default,\n mc.id AS model_configuration_id\n FROM model_configuration mc\n LEFT JOIN llm_provider lp\n ON lp.id = mc.llm_provider_id\n WHERE mc.supports_image_input IS TRUE;\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n # Populate vision defaults from model_flow\n op.execute(\n \"\"\"\n UPDATE llm_provider AS lp\n SET\n is_default_vision_provider = TRUE,\n default_vision_model = mc.name\n FROM llm_model_flow mf\n JOIN model_configuration mc ON mc.id = mf.model_configuration_id\n WHERE mf.llm_model_flow_type = 'VISION'\n AND mf.is_default = TRUE\n AND mc.llm_provider_id = lp.id;\n \"\"\"\n )\n\n # Populate conversation defaults from model_flow\n op.execute(\n \"\"\"\n UPDATE llm_provider AS lp\n SET\n is_default_provider = TRUE,\n default_model_name = mc.name\n FROM llm_model_flow mf\n JOIN model_configuration mc ON mc.id = mf.model_configuration_id\n WHERE mf.llm_model_flow_type = 'CHAT'\n AND mf.is_default = TRUE\n AND mc.llm_provider_id = lp.id;\n \"\"\"\n )\n\n # For providers that have conversation flow mappings but aren't the default,\n # we still need a default_model_name (it was NOT NULL originally)\n # Pick the first visible model or any model for that provider\n op.execute(\n \"\"\"\n UPDATE llm_provider AS lp\n SET default_model_name = (\n SELECT mc.name\n FROM model_configuration mc\n JOIN llm_model_flow mf ON mf.model_configuration_id = mc.id\n WHERE mc.llm_provider_id = lp.id\n AND mf.llm_model_flow_type = 'CHAT'\n ORDER BY mc.is_visible DESC, mc.id ASC\n LIMIT 1\n )\n WHERE lp.default_model_name IS NULL;\n \"\"\"\n )\n\n # Delete all model_flow entries (reverse the inserts from upgrade)\n op.execute(\"DELETE FROM llm_model_flow;\")\n" }, { "path": "backend/alembic/versions/027381bce97c_add_shortcut_option_for_users.py", "content": "\"\"\"add shortcut option for users\n\nRevision ID: 027381bce97c\nRevises: 6fc7886d665d\nCreate Date: 2025-01-14 12:14:00.814390\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"027381bce97c\"\ndown_revision = \"6fc7886d665d\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"shortcut_enabled\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"shortcut_enabled\")\n" }, { "path": "backend/alembic/versions/03bf8be6b53a_rework_kg_config.py", "content": "\"\"\"rework-kg-config\n\nRevision ID: 03bf8be6b53a\nRevises: 65bc6e0f8500\nCreate Date: 2025-06-16 10:52:34.815335\n\n\"\"\"\n\nimport json\n\n\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy import text\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"03bf8be6b53a\"\ndown_revision = \"65bc6e0f8500\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # get current config\n current_configs = (\n op.get_bind()\n .execute(text(\"SELECT kg_variable_name, kg_variable_values FROM kg_config\"))\n .all()\n )\n current_config_dict = {\n config.kg_variable_name: (\n config.kg_variable_values[0]\n if config.kg_variable_name\n not in (\"KG_VENDOR_DOMAINS\", \"KG_IGNORE_EMAIL_DOMAINS\")\n else config.kg_variable_values\n )\n for config in current_configs\n if config.kg_variable_values\n }\n\n # not using the KGConfigSettings model here in case it changes in the future\n kg_config_settings = json.dumps(\n {\n \"KG_EXPOSED\": current_config_dict.get(\"KG_EXPOSED\", False),\n \"KG_ENABLED\": current_config_dict.get(\"KG_ENABLED\", False),\n \"KG_VENDOR\": current_config_dict.get(\"KG_VENDOR\", None),\n \"KG_VENDOR_DOMAINS\": current_config_dict.get(\"KG_VENDOR_DOMAINS\", []),\n \"KG_IGNORE_EMAIL_DOMAINS\": current_config_dict.get(\n \"KG_IGNORE_EMAIL_DOMAINS\", []\n ),\n \"KG_COVERAGE_START\": current_config_dict.get(\n \"KG_COVERAGE_START\",\n (datetime.now() - timedelta(days=90)).strftime(\"%Y-%m-%d\"),\n ),\n \"KG_MAX_COVERAGE_DAYS\": current_config_dict.get(\"KG_MAX_COVERAGE_DAYS\", 90),\n \"KG_MAX_PARENT_RECURSION_DEPTH\": current_config_dict.get(\n \"KG_MAX_PARENT_RECURSION_DEPTH\", 2\n ),\n \"KG_BETA_PERSONA_ID\": current_config_dict.get(\"KG_BETA_PERSONA_ID\", None),\n }\n )\n op.execute(\n f\"INSERT INTO key_value_store (key, value) VALUES ('kg_config', '{kg_config_settings}')\"\n )\n\n # drop kg config table\n op.drop_table(\"kg_config\")\n\n\ndef downgrade() -> None:\n # get current config\n current_config_dict = {\n \"KG_EXPOSED\": False,\n \"KG_ENABLED\": False,\n \"KG_VENDOR\": [],\n \"KG_VENDOR_DOMAINS\": [],\n \"KG_IGNORE_EMAIL_DOMAINS\": [],\n \"KG_COVERAGE_START\": (datetime.now() - timedelta(days=90)).strftime(\"%Y-%m-%d\"),\n \"KG_MAX_COVERAGE_DAYS\": 90,\n \"KG_MAX_PARENT_RECURSION_DEPTH\": 2,\n }\n current_configs = (\n op.get_bind()\n .execute(text(\"SELECT value FROM key_value_store WHERE key = 'kg_config'\"))\n .one_or_none()\n )\n if current_configs is not None:\n current_config_dict.update(current_configs[0])\n insert_values = [\n {\n \"kg_variable_name\": name,\n \"kg_variable_values\": (\n [str(val).lower() if isinstance(val, bool) else str(val)]\n if not isinstance(val, list)\n else val\n ),\n }\n for name, val in current_config_dict.items()\n ]\n\n op.create_table(\n \"kg_config\",\n sa.Column(\"id\", sa.Integer(), primary_key=True, nullable=False, index=True),\n sa.Column(\"kg_variable_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"kg_variable_values\", postgresql.ARRAY(sa.String()), nullable=False),\n sa.UniqueConstraint(\"kg_variable_name\", name=\"uq_kg_config_variable_name\"),\n )\n op.bulk_insert(\n sa.table(\n \"kg_config\",\n sa.column(\"kg_variable_name\", sa.String),\n sa.column(\"kg_variable_values\", postgresql.ARRAY(sa.String)),\n ),\n insert_values,\n )\n\n op.execute(\"DELETE FROM key_value_store WHERE key = 'kg_config'\")\n" }, { "path": "backend/alembic/versions/03d085c5c38d_backfill_account_type.py", "content": "\"\"\"backfill_account_type\n\nRevision ID: 03d085c5c38d\nRevises: 977e834c1427\nCreate Date: 2026-03-25 16:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"03d085c5c38d\"\ndown_revision = \"977e834c1427\"\nbranch_labels = None\ndepends_on = None\n\n_STANDARD = \"STANDARD\"\n_BOT = \"BOT\"\n_EXT_PERM_USER = \"EXT_PERM_USER\"\n_SERVICE_ACCOUNT = \"SERVICE_ACCOUNT\"\n_ANONYMOUS = \"ANONYMOUS\"\n\n# Well-known anonymous user UUID\nANONYMOUS_USER_ID = \"00000000-0000-0000-0000-000000000002\"\n\n# Email pattern for API key virtual users\nAPI_KEY_EMAIL_PATTERN = r\"API\\_KEY\\_\\_%\"\n\n# Reflect the table structure for use in DML\nuser_table = sa.table(\n \"user\",\n sa.column(\"id\", sa.Uuid),\n sa.column(\"email\", sa.String),\n sa.column(\"role\", sa.String),\n sa.column(\"account_type\", sa.String),\n)\n\n\ndef upgrade() -> None:\n # ------------------------------------------------------------------\n # Step 1: Backfill account_type from role.\n # Order matters — most-specific matches first so the final catch-all\n # only touches rows that haven't been classified yet.\n # ------------------------------------------------------------------\n\n # 1a. API key virtual users → SERVICE_ACCOUNT\n op.execute(\n sa.update(user_table)\n .where(\n user_table.c.email.ilike(API_KEY_EMAIL_PATTERN),\n user_table.c.account_type.is_(None),\n )\n .values(account_type=_SERVICE_ACCOUNT)\n )\n\n # 1b. Anonymous user → ANONYMOUS\n op.execute(\n sa.update(user_table)\n .where(\n user_table.c.id == ANONYMOUS_USER_ID,\n user_table.c.account_type.is_(None),\n )\n .values(account_type=_ANONYMOUS)\n )\n\n # 1c. SLACK_USER role → BOT\n op.execute(\n sa.update(user_table)\n .where(\n user_table.c.role == \"SLACK_USER\",\n user_table.c.account_type.is_(None),\n )\n .values(account_type=_BOT)\n )\n\n # 1d. EXT_PERM_USER role → EXT_PERM_USER\n op.execute(\n sa.update(user_table)\n .where(\n user_table.c.role == \"EXT_PERM_USER\",\n user_table.c.account_type.is_(None),\n )\n .values(account_type=_EXT_PERM_USER)\n )\n\n # 1e. Everything else → STANDARD\n op.execute(\n sa.update(user_table)\n .where(user_table.c.account_type.is_(None))\n .values(account_type=_STANDARD)\n )\n\n # ------------------------------------------------------------------\n # Step 2: Set account_type to NOT NULL now that every row is filled.\n # ------------------------------------------------------------------\n op.alter_column(\n \"user\",\n \"account_type\",\n nullable=False,\n server_default=\"STANDARD\",\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\"user\", \"account_type\", nullable=True, server_default=None)\n op.execute(sa.update(user_table).values(account_type=None))\n" }, { "path": "backend/alembic/versions/03d710ccf29c_add_permission_sync_attempt_tables.py", "content": "\"\"\"add permission sync attempt tables\n\nRevision ID: 03d710ccf29c\nRevises: 96a5702df6aa\nCreate Date: 2025-09-11 13:30:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"03d710ccf29c\" # Generate a new unique ID\ndown_revision = \"96a5702df6aa\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create the permission sync status enum\n permission_sync_status_enum = sa.Enum(\n \"not_started\",\n \"in_progress\",\n \"success\",\n \"canceled\",\n \"failed\",\n \"completed_with_errors\",\n name=\"permissionsyncstatus\",\n native_enum=False,\n )\n\n # Create doc_permission_sync_attempt table\n op.create_table(\n \"doc_permission_sync_attempt\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"connector_credential_pair_id\", sa.Integer(), nullable=False),\n sa.Column(\"status\", permission_sync_status_enum, nullable=False),\n sa.Column(\"total_docs_synced\", sa.Integer(), nullable=True),\n sa.Column(\"docs_with_permission_errors\", sa.Integer(), nullable=True),\n sa.Column(\"error_message\", sa.Text(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"time_started\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\"time_finished\", sa.DateTime(timezone=True), nullable=True),\n sa.ForeignKeyConstraint(\n [\"connector_credential_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Create indexes for doc_permission_sync_attempt\n op.create_index(\n \"ix_doc_permission_sync_attempt_time_created\",\n \"doc_permission_sync_attempt\",\n [\"time_created\"],\n unique=False,\n )\n op.create_index(\n \"ix_permission_sync_attempt_latest_for_cc_pair\",\n \"doc_permission_sync_attempt\",\n [\"connector_credential_pair_id\", \"time_created\"],\n unique=False,\n )\n op.create_index(\n \"ix_permission_sync_attempt_status_time\",\n \"doc_permission_sync_attempt\",\n [\"status\", sa.text(\"time_finished DESC\")],\n unique=False,\n )\n\n # Create external_group_permission_sync_attempt table\n # connector_credential_pair_id is nullable - group syncs can be global (e.g., Confluence)\n op.create_table(\n \"external_group_permission_sync_attempt\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"connector_credential_pair_id\", sa.Integer(), nullable=True),\n sa.Column(\"status\", permission_sync_status_enum, nullable=False),\n sa.Column(\"total_users_processed\", sa.Integer(), nullable=True),\n sa.Column(\"total_groups_processed\", sa.Integer(), nullable=True),\n sa.Column(\"total_group_memberships_synced\", sa.Integer(), nullable=True),\n sa.Column(\"error_message\", sa.Text(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"time_started\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\"time_finished\", sa.DateTime(timezone=True), nullable=True),\n sa.ForeignKeyConstraint(\n [\"connector_credential_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Create indexes for external_group_permission_sync_attempt\n op.create_index(\n \"ix_external_group_permission_sync_attempt_time_created\",\n \"external_group_permission_sync_attempt\",\n [\"time_created\"],\n unique=False,\n )\n op.create_index(\n \"ix_group_sync_attempt_cc_pair_time\",\n \"external_group_permission_sync_attempt\",\n [\"connector_credential_pair_id\", \"time_created\"],\n unique=False,\n )\n op.create_index(\n \"ix_group_sync_attempt_status_time\",\n \"external_group_permission_sync_attempt\",\n [\"status\", sa.text(\"time_finished DESC\")],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n # Drop indexes\n op.drop_index(\n \"ix_group_sync_attempt_status_time\",\n table_name=\"external_group_permission_sync_attempt\",\n )\n op.drop_index(\n \"ix_group_sync_attempt_cc_pair_time\",\n table_name=\"external_group_permission_sync_attempt\",\n )\n op.drop_index(\n \"ix_external_group_permission_sync_attempt_time_created\",\n table_name=\"external_group_permission_sync_attempt\",\n )\n op.drop_index(\n \"ix_permission_sync_attempt_status_time\",\n table_name=\"doc_permission_sync_attempt\",\n )\n op.drop_index(\n \"ix_permission_sync_attempt_latest_for_cc_pair\",\n table_name=\"doc_permission_sync_attempt\",\n )\n op.drop_index(\n \"ix_doc_permission_sync_attempt_time_created\",\n table_name=\"doc_permission_sync_attempt\",\n )\n\n # Drop tables\n op.drop_table(\"external_group_permission_sync_attempt\")\n op.drop_table(\"doc_permission_sync_attempt\")\n" }, { "path": "backend/alembic/versions/0568ccf46a6b_add_thread_specific_model_selection.py", "content": "\"\"\"Add thread specific model selection\n\nRevision ID: 0568ccf46a6b\nRevises: e209dc5a8156\nCreate Date: 2024-06-19 14:25:36.376046\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"0568ccf46a6b\"\ndown_revision = \"e209dc5a8156\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_session\",\n sa.Column(\"current_alternate_model\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_session\", \"current_alternate_model\")\n" }, { "path": "backend/alembic/versions/05c07bf07c00_add_search_doc_relevance_details.py", "content": "\"\"\"add search doc relevance details\n\nRevision ID: 05c07bf07c00\nRevises: b896bbd0d5a7\nCreate Date: 2024-07-10 17:48:15.886653\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"05c07bf07c00\"\ndown_revision = \"b896bbd0d5a7\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"search_doc\",\n sa.Column(\"is_relevant\", sa.Boolean(), nullable=True),\n )\n op.add_column(\n \"search_doc\",\n sa.Column(\"relevance_explanation\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"search_doc\", \"relevance_explanation\")\n op.drop_column(\"search_doc\", \"is_relevant\")\n" }, { "path": "backend/alembic/versions/07b98176f1de_code_interpreter_seed.py", "content": "\"\"\"code interpreter seed\n\nRevision ID: 07b98176f1de\nRevises: 7cb492013621\nCreate Date: 2026-02-23 15:55:07.606784\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"07b98176f1de\"\ndown_revision = \"7cb492013621\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Seed the single instance of code_interpreter_server\n # NOTE: There should only exist at most and at minimum 1 code_interpreter_server row\n op.execute(\n sa.text(\"INSERT INTO code_interpreter_server (server_enabled) VALUES (true)\")\n )\n\n\ndef downgrade() -> None:\n op.execute(sa.text(\"DELETE FROM code_interpreter_server\"))\n" }, { "path": "backend/alembic/versions/0816326d83aa_add_federated_connector_tables.py", "content": "\"\"\"add federated connector tables\n\nRevision ID: 0816326d83aa\nRevises: 12635f6655b7\nCreate Date: 2025-06-29 14:09:45.109518\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"0816326d83aa\"\ndown_revision = \"12635f6655b7\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create federated_connector table\n op.create_table(\n \"federated_connector\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"source\", sa.String(), nullable=False),\n sa.Column(\"credentials\", sa.LargeBinary(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Create federated_connector_oauth_token table\n op.create_table(\n \"federated_connector_oauth_token\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"federated_connector_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", postgresql.UUID(as_uuid=True), nullable=False),\n sa.Column(\"token\", sa.LargeBinary(), nullable=False),\n sa.Column(\"expires_at\", sa.DateTime(), nullable=True),\n sa.ForeignKeyConstraint(\n [\"federated_connector_id\"], [\"federated_connector.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Create federated_connector__document_set table\n op.create_table(\n \"federated_connector__document_set\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"federated_connector_id\", sa.Integer(), nullable=False),\n sa.Column(\"document_set_id\", sa.Integer(), nullable=False),\n sa.Column(\"entities\", postgresql.JSONB(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"federated_connector_id\"], [\"federated_connector.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint(\n [\"document_set_id\"], [\"document_set.id\"], ondelete=\"CASCADE\"\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\n \"federated_connector_id\",\n \"document_set_id\",\n name=\"uq_federated_connector_document_set\",\n ),\n )\n\n\ndef downgrade() -> None:\n # Drop tables in reverse order due to foreign key dependencies\n op.drop_table(\"federated_connector__document_set\")\n op.drop_table(\"federated_connector_oauth_token\")\n op.drop_table(\"federated_connector\")\n" }, { "path": "backend/alembic/versions/08a1eda20fe1_add_earliest_indexing_to_connector.py", "content": "\"\"\"add_indexing_start_to_connector\n\nRevision ID: 08a1eda20fe1\nRevises: 8a87bd6ec550\nCreate Date: 2024-07-23 11:12:39.462397\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"08a1eda20fe1\"\ndown_revision = \"8a87bd6ec550\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector\", sa.Column(\"indexing_start\", sa.DateTime(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector\", \"indexing_start\")\n" }, { "path": "backend/alembic/versions/09995b8811eb_add_theme_preference_to_user.py", "content": "\"\"\"add theme_preference to user\n\nRevision ID: 09995b8811eb\nRevises: 3d1cca026fe8\nCreate Date: 2025-10-24 08:58:50.246949\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom onyx.db.enums import ThemePreference\n\n\n# revision identifiers, used by Alembic.\nrevision = \"09995b8811eb\"\ndown_revision = \"3d1cca026fe8\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"theme_preference\",\n sa.Enum(ThemePreference, native_enum=False),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"theme_preference\")\n" }, { "path": "backend/alembic/versions/0a2b51deb0b8_add_starter_prompts.py", "content": "\"\"\"Add starter prompts\n\nRevision ID: 0a2b51deb0b8\nRevises: 5f4b8568a221\nCreate Date: 2024-03-02 23:23:49.960309\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"0a2b51deb0b8\"\ndown_revision = \"5f4b8568a221\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\n \"starter_messages\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"starter_messages\")\n" }, { "path": "backend/alembic/versions/0a98909f2757_enable_encrypted_fields.py", "content": "\"\"\"Enable Encrypted Fields\n\nRevision ID: 0a98909f2757\nRevises: 570282d33c49\nCreate Date: 2024-05-05 19:30:34.317972\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.sql import table\nfrom sqlalchemy.dialects import postgresql\nimport json\n\nfrom onyx.utils.encryption import encrypt_string_to_bytes\n\n# revision identifiers, used by Alembic.\nrevision = \"0a98909f2757\"\ndown_revision = \"570282d33c49\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n connection = op.get_bind()\n\n op.alter_column(\"key_value_store\", \"value\", nullable=True)\n op.add_column(\n \"key_value_store\",\n sa.Column(\n \"encrypted_value\",\n sa.LargeBinary,\n nullable=True,\n ),\n )\n\n # Need a temporary column to translate the JSONB to binary\n op.add_column(\"credential\", sa.Column(\"temp_column\", sa.LargeBinary()))\n\n creds_table = table(\n \"credential\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"credential_json\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=False,\n ),\n sa.Column(\n \"temp_column\",\n sa.LargeBinary(),\n nullable=False,\n ),\n )\n\n results = connection.execute(sa.select(creds_table))\n\n # This uses the MIT encrypt which does not actually encrypt the credentials\n # In other words, this upgrade does not apply the encryption. Porting existing sensitive data\n # and key rotation currently is not supported and will come out in the future\n for row_id, creds, _ in results:\n creds_binary = encrypt_string_to_bytes(json.dumps(creds))\n connection.execute(\n creds_table.update()\n .where(creds_table.c.id == row_id)\n .values(temp_column=creds_binary)\n )\n\n op.drop_column(\"credential\", \"credential_json\")\n op.alter_column(\"credential\", \"temp_column\", new_column_name=\"credential_json\")\n\n op.add_column(\"llm_provider\", sa.Column(\"temp_column\", sa.LargeBinary()))\n\n llm_table = table(\n \"llm_provider\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"api_key\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\n \"temp_column\",\n sa.LargeBinary(),\n nullable=False,\n ),\n )\n results = connection.execute(sa.select(llm_table))\n\n for row_id, api_key, _ in results:\n llm_key = encrypt_string_to_bytes(api_key)\n connection.execute(\n llm_table.update()\n .where(llm_table.c.id == row_id)\n .values(temp_column=llm_key)\n )\n\n op.drop_column(\"llm_provider\", \"api_key\")\n op.alter_column(\"llm_provider\", \"temp_column\", new_column_name=\"api_key\")\n\n\ndef downgrade() -> None:\n # Some information loss but this is ok. Should not allow decryption via downgrade.\n op.drop_column(\"credential\", \"credential_json\")\n op.drop_column(\"llm_provider\", \"api_key\")\n\n op.add_column(\"llm_provider\", sa.Column(\"api_key\", sa.String()))\n op.add_column(\n \"credential\",\n sa.Column(\"credential_json\", postgresql.JSONB(astext_type=sa.Text())),\n )\n\n op.execute(\"DELETE FROM key_value_store WHERE value IS NULL\")\n op.alter_column(\"key_value_store\", \"value\", nullable=False)\n op.drop_column(\"key_value_store\", \"encrypted_value\")\n" }, { "path": "backend/alembic/versions/0bb4558f35df_add_scim_username_to_scim_user_mapping.py", "content": "\"\"\"add scim_username to scim_user_mapping\n\nRevision ID: 0bb4558f35df\nRevises: 631fd2504136\nCreate Date: 2026-02-20 10:45:30.340188\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"0bb4558f35df\"\ndown_revision = \"631fd2504136\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"scim_user_mapping\",\n sa.Column(\"scim_username\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"scim_user_mapping\", \"scim_username\")\n" }, { "path": "backend/alembic/versions/0cd424f32b1d_user_file_data_preparation_and_backfill.py", "content": "\"\"\"Migration 2: User file data preparation and backfill\n\nRevision ID: 0cd424f32b1d\nRevises: 9b66d3156fc6\nCreate Date: 2025-09-22 09:44:42.727034\n\nThis migration populates the new columns added in migration 1.\nIt prepares data for the UUID transition and relationship migration.\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy import text\nimport logging\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n# revision identifiers, used by Alembic.\nrevision = \"0cd424f32b1d\"\ndown_revision = \"9b66d3156fc6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Populate new columns with data.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n # === Step 1: Populate user_file.new_id ===\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n has_new_id = \"new_id\" in user_file_columns\n\n if has_new_id:\n logger.info(\"Populating user_file.new_id with UUIDs...\")\n\n # Count rows needing UUIDs\n null_count = bind.execute(\n text(\"SELECT COUNT(*) FROM user_file WHERE new_id IS NULL\")\n ).scalar_one()\n\n if null_count > 0:\n logger.info(f\"Generating UUIDs for {null_count} user_file records...\")\n\n # Populate in batches to avoid long locks\n batch_size = 10000\n total_updated = 0\n\n while True:\n result = bind.execute(\n text(\n \"\"\"\n UPDATE user_file\n SET new_id = gen_random_uuid()\n WHERE new_id IS NULL\n AND id IN (\n SELECT id FROM user_file\n WHERE new_id IS NULL\n LIMIT :batch_size\n )\n \"\"\"\n ),\n {\"batch_size\": batch_size},\n )\n\n updated = result.rowcount\n total_updated += updated\n\n if updated < batch_size:\n break\n\n logger.info(f\" Updated {total_updated}/{null_count} records...\")\n\n logger.info(f\"Generated UUIDs for {total_updated} user_file records\")\n\n # Verify all records have UUIDs\n remaining_null = bind.execute(\n text(\"SELECT COUNT(*) FROM user_file WHERE new_id IS NULL\")\n ).scalar_one()\n\n if remaining_null > 0:\n raise Exception(\n f\"Failed to populate all user_file.new_id values ({remaining_null} NULL)\"\n )\n\n # Lock down the column\n op.alter_column(\"user_file\", \"new_id\", nullable=False)\n op.alter_column(\"user_file\", \"new_id\", server_default=None)\n logger.info(\"Locked down user_file.new_id column\")\n\n # === Step 2: Populate persona__user_file.user_file_id_uuid ===\n persona_user_file_columns = [\n col[\"name\"] for col in inspector.get_columns(\"persona__user_file\")\n ]\n\n if has_new_id and \"user_file_id_uuid\" in persona_user_file_columns:\n logger.info(\"Populating persona__user_file.user_file_id_uuid...\")\n\n # Count rows needing update\n null_count = bind.execute(\n text(\n \"\"\"\n SELECT COUNT(*) FROM persona__user_file\n WHERE user_file_id IS NOT NULL AND user_file_id_uuid IS NULL\n \"\"\"\n )\n ).scalar_one()\n\n if null_count > 0:\n logger.info(f\"Updating {null_count} persona__user_file records...\")\n\n # Update in batches\n batch_size = 10000\n total_updated = 0\n\n while True:\n result = bind.execute(\n text(\n \"\"\"\n UPDATE persona__user_file p\n SET user_file_id_uuid = uf.new_id\n FROM user_file uf\n WHERE p.user_file_id = uf.id\n AND p.user_file_id_uuid IS NULL\n AND p.persona_id IN (\n SELECT persona_id\n FROM persona__user_file\n WHERE user_file_id_uuid IS NULL\n LIMIT :batch_size\n )\n \"\"\"\n ),\n {\"batch_size\": batch_size},\n )\n\n updated = result.rowcount\n total_updated += updated\n\n if updated < batch_size:\n break\n\n logger.info(f\" Updated {total_updated}/{null_count} records...\")\n\n logger.info(f\"Updated {total_updated} persona__user_file records\")\n\n # Verify all records are populated\n remaining_null = bind.execute(\n text(\n \"\"\"\n SELECT COUNT(*) FROM persona__user_file\n WHERE user_file_id IS NOT NULL AND user_file_id_uuid IS NULL\n \"\"\"\n )\n ).scalar_one()\n\n if remaining_null > 0:\n raise Exception(\n f\"Failed to populate all persona__user_file.user_file_id_uuid values ({remaining_null} NULL)\"\n )\n\n op.alter_column(\"persona__user_file\", \"user_file_id_uuid\", nullable=False)\n logger.info(\"Locked down persona__user_file.user_file_id_uuid column\")\n\n # === Step 3: Create user_project records from chat_folder ===\n if \"chat_folder\" in inspector.get_table_names():\n logger.info(\"Creating user_project records from chat_folder...\")\n\n result = bind.execute(\n text(\n \"\"\"\n INSERT INTO user_project (user_id, name)\n SELECT cf.user_id, cf.name\n FROM chat_folder cf\n WHERE NOT EXISTS (\n SELECT 1\n FROM user_project up\n WHERE up.user_id = cf.user_id AND up.name = cf.name\n )\n \"\"\"\n )\n )\n\n logger.info(f\"Created {result.rowcount} user_project records from chat_folder\")\n\n # === Step 4: Populate chat_session.project_id ===\n chat_session_columns = [\n col[\"name\"] for col in inspector.get_columns(\"chat_session\")\n ]\n\n if \"folder_id\" in chat_session_columns and \"project_id\" in chat_session_columns:\n logger.info(\"Populating chat_session.project_id...\")\n\n # Count sessions needing update\n null_count = bind.execute(\n text(\n \"\"\"\n SELECT COUNT(*) FROM chat_session\n WHERE project_id IS NULL AND folder_id IS NOT NULL\n \"\"\"\n )\n ).scalar_one()\n\n if null_count > 0:\n logger.info(f\"Updating {null_count} chat_session records...\")\n\n result = bind.execute(\n text(\n \"\"\"\n UPDATE chat_session cs\n SET project_id = up.id\n FROM chat_folder cf\n JOIN user_project up ON up.user_id = cf.user_id AND up.name = cf.name\n WHERE cs.folder_id = cf.id AND cs.project_id IS NULL\n \"\"\"\n )\n )\n\n logger.info(f\"Updated {result.rowcount} chat_session records\")\n\n # Verify all records are populated\n remaining_null = bind.execute(\n text(\n \"\"\"\n SELECT COUNT(*) FROM chat_session\n WHERE project_id IS NULL AND folder_id IS NOT NULL\n \"\"\"\n )\n ).scalar_one()\n\n if remaining_null > 0:\n logger.warning(\n f\"Warning: {remaining_null} chat_session records could not be mapped to projects\"\n )\n\n # === Step 5: Update plaintext FileRecord IDs/display names to UUID scheme ===\n # Prior to UUID migration, plaintext cache files were stored with file_id like 'plain_text_'.\n # After migration, we use 'plaintext_' (note the name change to 'plaintext_').\n # This step remaps existing FileRecord rows to the new naming while preserving object_key/bucket.\n logger.info(\"Updating plaintext FileRecord ids and display names to UUID scheme...\")\n\n # Count legacy plaintext records that can be mapped to UUID user_file ids\n count_query = text(\n \"\"\"\n SELECT COUNT(*)\n FROM file_record fr\n JOIN user_file uf ON fr.file_id = CONCAT('plaintext_', uf.id::text)\n WHERE LOWER(fr.file_origin::text) = 'plaintext_cache'\n \"\"\"\n )\n legacy_count = bind.execute(count_query).scalar_one()\n\n if legacy_count and legacy_count > 0:\n logger.info(f\"Found {legacy_count} legacy plaintext file records to update\")\n\n # Update display_name first for readability (safe regardless of rename)\n bind.execute(\n text(\n \"\"\"\n UPDATE file_record fr\n SET display_name = CONCAT('Plaintext for user file ', uf.new_id::text)\n FROM user_file uf\n WHERE LOWER(fr.file_origin::text) = 'plaintext_cache'\n AND fr.file_id = CONCAT('plaintext_', uf.id::text)\n \"\"\"\n )\n )\n\n # Remap file_id from 'plaintext_' -> 'plaintext_' using transitional new_id\n # Use a single UPDATE ... WHERE file_id LIKE 'plain_text_%'\n # and ensure it aligns to existing user_file ids to avoid renaming unrelated rows\n result = bind.execute(\n text(\n \"\"\"\n UPDATE file_record fr\n SET file_id = CONCAT('plaintext_', uf.new_id::text)\n FROM user_file uf\n WHERE LOWER(fr.file_origin::text) = 'plaintext_cache'\n AND fr.file_id = CONCAT('plaintext_', uf.id::text)\n \"\"\"\n )\n )\n logger.info(\n f\"Updated {result.rowcount} plaintext file_record ids to UUID scheme\"\n )\n\n # === Step 6: Ensure document_id_migrated default TRUE and backfill existing FALSE ===\n # New records should default to migrated=True so the migration task won't run for them.\n # Existing rows that had a legacy document_id should be marked as not migrated to be processed.\n\n # Backfill existing records: if document_id is not null, set to FALSE\n bind.execute(\n text(\n \"\"\"\n UPDATE user_file\n SET document_id_migrated = FALSE\n WHERE document_id IS NOT NULL\n \"\"\"\n )\n )\n\n # === Step 7: Backfill user_file.status from index_attempt ===\n logger.info(\"Backfilling user_file.status from index_attempt...\")\n\n # Update user_file status based on latest index attempt\n # Using CTEs instead of temp tables for asyncpg compatibility\n result = bind.execute(\n text(\n \"\"\"\n WITH latest_attempt AS (\n SELECT DISTINCT ON (ia.connector_credential_pair_id)\n ia.connector_credential_pair_id,\n ia.status\n FROM index_attempt ia\n ORDER BY ia.connector_credential_pair_id, ia.time_updated DESC\n ),\n uf_to_ccp AS (\n SELECT DISTINCT uf.id AS uf_id, ccp.id AS cc_pair_id\n FROM user_file uf\n JOIN document_by_connector_credential_pair dcc\n ON dcc.id = REPLACE(uf.document_id, 'USER_FILE_CONNECTOR__', 'FILE_CONNECTOR__')\n JOIN connector_credential_pair ccp\n ON ccp.connector_id = dcc.connector_id\n AND ccp.credential_id = dcc.credential_id\n )\n UPDATE user_file uf\n SET status = CASE\n WHEN la.status IN ('NOT_STARTED', 'IN_PROGRESS') THEN 'PROCESSING'\n WHEN la.status = 'SUCCESS' THEN 'COMPLETED'\n ELSE 'FAILED'\n END\n FROM uf_to_ccp ufc\n LEFT JOIN latest_attempt la\n ON la.connector_credential_pair_id = ufc.cc_pair_id\n WHERE uf.id = ufc.uf_id\n AND uf.status = 'PROCESSING'\n \"\"\"\n )\n )\n\n logger.info(f\"Updated status for {result.rowcount} user_file records\")\n\n logger.info(\"Migration 2 (data preparation) completed successfully\")\n\n\ndef downgrade() -> None:\n \"\"\"Reset populated data to allow clean downgrade of schema.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n logger.info(\"Starting downgrade of data preparation...\")\n\n # Reset user_file columns to allow nulls before data removal\n if \"user_file\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n\n if \"new_id\" in columns:\n op.alter_column(\n \"user_file\",\n \"new_id\",\n nullable=True,\n server_default=sa.text(\"gen_random_uuid()\"),\n )\n # Optionally clear the data\n # bind.execute(text(\"UPDATE user_file SET new_id = NULL\"))\n logger.info(\"Reset user_file.new_id to nullable\")\n\n # Reset persona__user_file.user_file_id_uuid\n if \"persona__user_file\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"persona__user_file\")]\n\n if \"user_file_id_uuid\" in columns:\n op.alter_column(\"persona__user_file\", \"user_file_id_uuid\", nullable=True)\n # Optionally clear the data\n # bind.execute(text(\"UPDATE persona__user_file SET user_file_id_uuid = NULL\"))\n logger.info(\"Reset persona__user_file.user_file_id_uuid to nullable\")\n\n # Note: We don't delete user_project records or reset chat_session.project_id\n # as these might be in use and can be handled by the schema downgrade\n\n # Reset user_file.status to default\n if \"user_file\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n if \"status\" in columns:\n bind.execute(text(\"UPDATE user_file SET status = 'PROCESSING'\"))\n logger.info(\"Reset user_file.status to default\")\n\n logger.info(\"Downgrade completed successfully\")\n" }, { "path": "backend/alembic/versions/0ebb1d516877_add_ccpair_deletion_failure_message.py", "content": "\"\"\"add ccpair deletion failure message\n\nRevision ID: 0ebb1d516877\nRevises: 52a219fb5233\nCreate Date: 2024-09-10 15:03:48.233926\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"0ebb1d516877\"\ndown_revision = \"52a219fb5233\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"deletion_failure_message\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"deletion_failure_message\")\n" }, { "path": "backend/alembic/versions/0f7ff6d75b57_add_index_to_index_attempt_time_created.py", "content": "\"\"\"add index to index_attempt.time_created\n\nRevision ID: 0f7ff6d75b57\nRevises: 369644546676\nCreate Date: 2025-01-10 14:01:14.067144\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"0f7ff6d75b57\"\ndown_revision = \"fec3db967bf7\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_index(\n op.f(\"ix_index_attempt_status\"),\n \"index_attempt\",\n [\"status\"],\n unique=False,\n )\n\n op.create_index(\n op.f(\"ix_index_attempt_time_created\"),\n \"index_attempt\",\n [\"time_created\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(op.f(\"ix_index_attempt_time_created\"), table_name=\"index_attempt\")\n\n op.drop_index(op.f(\"ix_index_attempt_status\"), table_name=\"index_attempt\")\n" }, { "path": "backend/alembic/versions/114a638452db_add_default_app_mode_to_user.py", "content": "\"\"\"add default_app_mode to user\n\nRevision ID: 114a638452db\nRevises: feead2911109\nCreate Date: 2026-02-09 18:57:08.274640\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"114a638452db\"\ndown_revision = \"feead2911109\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"default_app_mode\",\n sa.String(),\n nullable=False,\n server_default=\"CHAT\",\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"default_app_mode\")\n" }, { "path": "backend/alembic/versions/12635f6655b7_drive_canonical_ids.py", "content": "\"\"\"drive-canonical-ids\n\nRevision ID: 12635f6655b7\nRevises: 58c50ef19f08\nCreate Date: 2025-06-20 14:44:54.241159\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom urllib.parse import urlparse, urlunparse\nfrom httpx import HTTPStatusError\nimport httpx\nfrom onyx.db.search_settings import SearchSettings\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\nfrom onyx.document_index.vespa.shared_utils.utils import (\n replace_invalid_doc_id_characters,\n)\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.utils.logger import setup_logger\nimport os\n\nlogger = setup_logger()\n\n# revision identifiers, used by Alembic.\nrevision = \"12635f6655b7\"\ndown_revision = \"58c50ef19f08\"\nbranch_labels = None\ndepends_on = None\n\nSKIP_CANON_DRIVE_IDS = os.environ.get(\"SKIP_CANON_DRIVE_IDS\", \"true\").lower() == \"true\"\n\n\ndef active_search_settings() -> tuple[SearchSettings, SearchSettings | None]:\n result = op.get_bind().execute(\n sa.text(\n \"\"\"\n SELECT * FROM search_settings WHERE status = 'PRESENT' ORDER BY id DESC LIMIT 1\n \"\"\"\n )\n )\n search_settings_fetch = result.fetchall()\n search_settings = (\n SearchSettings(**search_settings_fetch[0]._asdict())\n if search_settings_fetch\n else None\n )\n\n result2 = op.get_bind().execute(\n sa.text(\n \"\"\"\n SELECT * FROM search_settings WHERE status = 'FUTURE' ORDER BY id DESC LIMIT 1\n \"\"\"\n )\n )\n search_settings_future_fetch = result2.fetchall()\n search_settings_future = (\n SearchSettings(**search_settings_future_fetch[0]._asdict())\n if search_settings_future_fetch\n else None\n )\n\n if not isinstance(search_settings, SearchSettings):\n raise RuntimeError(\n \"current search settings is of type \" + str(type(search_settings))\n )\n if (\n not isinstance(search_settings_future, SearchSettings)\n and search_settings_future is not None\n ):\n raise RuntimeError(\n \"future search settings is of type \" + str(type(search_settings_future))\n )\n\n return search_settings, search_settings_future\n\n\ndef normalize_google_drive_url(url: str) -> str:\n \"\"\"Remove query parameters from Google Drive URLs to create canonical document IDs.\n NOTE: copied from drive doc_conversion.py\n \"\"\"\n parsed_url = urlparse(url)\n parsed_url = parsed_url._replace(query=\"\")\n spl_path = parsed_url.path.split(\"/\")\n if spl_path and (spl_path[-1] in [\"edit\", \"view\", \"preview\"]):\n spl_path.pop()\n parsed_url = parsed_url._replace(path=\"/\".join(spl_path))\n # Remove query parameters and reconstruct URL\n return urlunparse(parsed_url)\n\n\ndef get_google_drive_documents_from_database() -> list[dict]:\n \"\"\"Get all Google Drive documents from the database.\"\"\"\n bind = op.get_bind()\n result = bind.execute(\n sa.text(\n \"\"\"\n SELECT d.id\n FROM document d\n JOIN document_by_connector_credential_pair dcc ON d.id = dcc.id\n JOIN connector_credential_pair cc ON dcc.connector_id = cc.connector_id\n AND dcc.credential_id = cc.credential_id\n JOIN connector c ON cc.connector_id = c.id\n WHERE c.source = 'GOOGLE_DRIVE'\n \"\"\"\n )\n )\n\n documents = []\n for row in result:\n documents.append({\"document_id\": row.id})\n\n return documents\n\n\ndef update_document_id_in_database(\n old_doc_id: str, new_doc_id: str, index_name: str\n) -> None:\n \"\"\"Update document IDs in all relevant database tables using copy-and-swap approach.\"\"\"\n bind = op.get_bind()\n\n # print(f\"Updating database tables for document {old_doc_id} -> {new_doc_id}\")\n\n # Check if new document ID already exists\n result = bind.execute(\n sa.text(\"SELECT COUNT(*) FROM document WHERE id = :new_id\"),\n {\"new_id\": new_doc_id},\n )\n row = result.fetchone()\n if row and row[0] > 0:\n # print(f\"Document with ID {new_doc_id} already exists, deleting old one\")\n delete_document_from_db(old_doc_id, index_name)\n return\n\n # Step 1: Create a new document row with the new ID (copy all fields from old row)\n # Use a conservative approach to handle columns that might not exist in all installations\n try:\n bind.execute(\n sa.text(\n \"\"\"\n INSERT INTO document (id, from_ingestion_api, boost, hidden, semantic_id,\n link, doc_updated_at, primary_owners, secondary_owners,\n external_user_emails, external_user_group_ids, is_public,\n chunk_count, last_modified, last_synced, kg_stage, kg_processing_time)\n SELECT :new_id, from_ingestion_api, boost, hidden, semantic_id,\n link, doc_updated_at, primary_owners, secondary_owners,\n external_user_emails, external_user_group_ids, is_public,\n chunk_count, last_modified, last_synced, kg_stage, kg_processing_time\n FROM document\n WHERE id = :old_id\n \"\"\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated database tables for document {old_doc_id} -> {new_doc_id}\")\n except Exception as e:\n # If the full INSERT fails, try a more basic version with only core columns\n logger.warning(f\"Full INSERT failed, trying basic version: {e}\")\n bind.execute(\n sa.text(\n \"\"\"\n INSERT INTO document (id, from_ingestion_api, boost, hidden, semantic_id,\n link, doc_updated_at, primary_owners, secondary_owners)\n SELECT :new_id, from_ingestion_api, boost, hidden, semantic_id,\n link, doc_updated_at, primary_owners, secondary_owners\n FROM document\n WHERE id = :old_id\n \"\"\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n\n # Step 2: Update all foreign key references to point to the new ID\n\n # Update document_by_connector_credential_pair table\n bind.execute(\n sa.text(\n \"UPDATE document_by_connector_credential_pair SET id = :new_id WHERE id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated document_by_connector_credential_pair table for document {old_doc_id} -> {new_doc_id}\")\n\n # Update search_doc table (stores search results for chat replay)\n # This is critical for agent functionality\n bind.execute(\n sa.text(\n \"UPDATE search_doc SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated search_doc table for document {old_doc_id} -> {new_doc_id}\")\n # Update document_retrieval_feedback table (user feedback on documents)\n bind.execute(\n sa.text(\n \"UPDATE document_retrieval_feedback SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated document_retrieval_feedback table for document {old_doc_id} -> {new_doc_id}\")\n # Update document__tag table (document-tag relationships)\n bind.execute(\n sa.text(\n \"UPDATE document__tag SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated document__tag table for document {old_doc_id} -> {new_doc_id}\")\n # Update user_file table (user uploaded files linked to documents)\n bind.execute(\n sa.text(\n \"UPDATE user_file SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated user_file table for document {old_doc_id} -> {new_doc_id}\")\n # Update KG and chunk_stats tables (these may not exist in all installations)\n try:\n # Update kg_entity table\n bind.execute(\n sa.text(\n \"UPDATE kg_entity SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated kg_entity table for document {old_doc_id} -> {new_doc_id}\")\n # Update kg_entity_extraction_staging table\n bind.execute(\n sa.text(\n \"UPDATE kg_entity_extraction_staging SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated kg_entity_extraction_staging table for document {old_doc_id} -> {new_doc_id}\")\n # Update kg_relationship table\n bind.execute(\n sa.text(\n \"UPDATE kg_relationship SET source_document = :new_id WHERE source_document = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated kg_relationship table for document {old_doc_id} -> {new_doc_id}\")\n # Update kg_relationship_extraction_staging table\n bind.execute(\n sa.text(\n \"UPDATE kg_relationship_extraction_staging SET source_document = :new_id WHERE source_document = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated kg_relationship_extraction_staging table for document {old_doc_id} -> {new_doc_id}\")\n # Update chunk_stats table\n bind.execute(\n sa.text(\n \"UPDATE chunk_stats SET document_id = :new_id WHERE document_id = :old_id\"\n ),\n {\"new_id\": new_doc_id, \"old_id\": old_doc_id},\n )\n # print(f\"Successfully updated chunk_stats table for document {old_doc_id} -> {new_doc_id}\")\n # Update chunk_stats ID field which includes document_id\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE chunk_stats\n SET id = REPLACE(id, :old_id, :new_id)\n WHERE id LIKE :old_id_pattern\n \"\"\"\n ),\n {\n \"new_id\": new_doc_id,\n \"old_id\": old_doc_id,\n \"old_id_pattern\": f\"{old_doc_id}__%\",\n },\n )\n # print(f\"Successfully updated chunk_stats ID field for document {old_doc_id} -> {new_doc_id}\")\n except Exception as e:\n logger.warning(f\"Some KG/chunk tables may not exist or failed to update: {e}\")\n\n # Step 3: Delete the old document row (this should now be safe since all FKs point to new row)\n bind.execute(\n sa.text(\"DELETE FROM document WHERE id = :old_id\"), {\"old_id\": old_doc_id}\n )\n # print(f\"Successfully deleted document {old_doc_id} from database\")\n\n\ndef _visit_chunks(\n *,\n http_client: httpx.Client,\n index_name: str,\n selection: str,\n continuation: str | None = None,\n) -> tuple[list[dict], str | None]:\n \"\"\"Helper that calls the /document/v1 visit API once and returns (docs, next_token).\"\"\"\n\n # Use the same URL as the document API, but with visit-specific params\n base_url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)\n\n params: dict[str, str] = {\n \"selection\": selection,\n \"wantedDocumentCount\": \"1000\",\n }\n if continuation:\n params[\"continuation\"] = continuation\n\n # print(f\"Visiting chunks for selection '{selection}' with params {params}\")\n resp = http_client.get(base_url, params=params, timeout=None)\n # print(f\"Visited chunks for document {selection}\")\n resp.raise_for_status()\n\n payload = resp.json()\n return payload.get(\"documents\", []), payload.get(\"continuation\")\n\n\ndef delete_document_chunks_from_vespa(index_name: str, doc_id: str) -> None:\n \"\"\"Delete all chunks for *doc_id* from Vespa using continuation-token paging (no offset).\"\"\"\n\n total_deleted = 0\n # Use exact match instead of contains - Document Selector Language doesn't support contains\n selection = f'{index_name}.document_id==\"{doc_id}\"'\n\n with get_vespa_http_client() as http_client:\n continuation: str | None = None\n while True:\n docs, continuation = _visit_chunks(\n http_client=http_client,\n index_name=index_name,\n selection=selection,\n continuation=continuation,\n )\n\n if not docs:\n break\n\n for doc in docs:\n vespa_full_id = doc.get(\"id\")\n if not vespa_full_id:\n continue\n\n vespa_doc_uuid = vespa_full_id.split(\"::\")[-1]\n delete_url = f\"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{vespa_doc_uuid}\"\n\n try:\n resp = http_client.delete(delete_url)\n resp.raise_for_status()\n total_deleted += 1\n except Exception as e:\n print(f\"Failed to delete chunk {vespa_doc_uuid}: {e}\")\n\n if not continuation:\n break\n\n\ndef update_document_id_in_vespa(\n index_name: str, old_doc_id: str, new_doc_id: str\n) -> None:\n \"\"\"Update all chunks' document_id field from *old_doc_id* to *new_doc_id* using continuation paging.\"\"\"\n\n clean_new_doc_id = replace_invalid_doc_id_characters(new_doc_id)\n\n # Use exact match instead of contains - Document Selector Language doesn't support contains\n selection = f'{index_name}.document_id==\"{old_doc_id}\"'\n\n with get_vespa_http_client() as http_client:\n continuation: str | None = None\n while True:\n # print(f\"Visiting chunks for document {old_doc_id} -> {new_doc_id}\")\n docs, continuation = _visit_chunks(\n http_client=http_client,\n index_name=index_name,\n selection=selection,\n continuation=continuation,\n )\n\n if not docs:\n break\n\n for doc in docs:\n vespa_full_id = doc.get(\"id\")\n if not vespa_full_id:\n continue\n\n vespa_doc_uuid = vespa_full_id.split(\"::\")[-1]\n vespa_url = f\"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{vespa_doc_uuid}\"\n\n update_request = {\n \"fields\": {\"document_id\": {\"assign\": clean_new_doc_id}}\n }\n\n try:\n resp = http_client.put(vespa_url, json=update_request)\n resp.raise_for_status()\n except Exception as e:\n print(f\"Failed to update chunk {vespa_doc_uuid}: {e}\")\n raise\n\n if not continuation:\n break\n\n\ndef delete_document_from_db(current_doc_id: str, index_name: str) -> None:\n # Delete all foreign key references first, then delete the document\n try:\n bind = op.get_bind()\n\n # Delete from agent-related tables first (order matters due to foreign keys)\n # Delete from agent__sub_query__search_doc first since it references search_doc\n bind.execute(\n sa.text(\n \"\"\"\n DELETE FROM agent__sub_query__search_doc\n WHERE search_doc_id IN (\n SELECT id FROM search_doc WHERE document_id = :doc_id\n )\n \"\"\"\n ),\n {\"doc_id\": current_doc_id},\n )\n\n # Delete from chat_message__search_doc\n bind.execute(\n sa.text(\n \"\"\"\n DELETE FROM chat_message__search_doc\n WHERE search_doc_id IN (\n SELECT id FROM search_doc WHERE document_id = :doc_id\n )\n \"\"\"\n ),\n {\"doc_id\": current_doc_id},\n )\n\n # Now we can safely delete from search_doc\n bind.execute(\n sa.text(\"DELETE FROM search_doc WHERE document_id = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n # Delete from document_by_connector_credential_pair\n bind.execute(\n sa.text(\n \"DELETE FROM document_by_connector_credential_pair WHERE id = :doc_id\"\n ),\n {\"doc_id\": current_doc_id},\n )\n\n # Delete from other tables that reference this document\n bind.execute(\n sa.text(\n \"DELETE FROM document_retrieval_feedback WHERE document_id = :doc_id\"\n ),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\"DELETE FROM document__tag WHERE document_id = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\"DELETE FROM user_file WHERE document_id = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n # Delete from KG tables if they exist\n try:\n bind.execute(\n sa.text(\"DELETE FROM kg_entity WHERE document_id = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\n \"DELETE FROM kg_entity_extraction_staging WHERE document_id = :doc_id\"\n ),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\"DELETE FROM kg_relationship WHERE source_document = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\n \"DELETE FROM kg_relationship_extraction_staging WHERE source_document = :doc_id\"\n ),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\"DELETE FROM chunk_stats WHERE document_id = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n bind.execute(\n sa.text(\"DELETE FROM chunk_stats WHERE id LIKE :doc_id_pattern\"),\n {\"doc_id_pattern\": f\"{current_doc_id}__%\"},\n )\n\n except Exception as e:\n logger.warning(\n f\"Some KG/chunk tables may not exist or failed to delete from: {e}\"\n )\n\n # Finally delete the document itself\n bind.execute(\n sa.text(\"DELETE FROM document WHERE id = :doc_id\"),\n {\"doc_id\": current_doc_id},\n )\n\n # Delete chunks from vespa\n delete_document_chunks_from_vespa(index_name, current_doc_id)\n\n except Exception as e:\n print(f\"Failed to delete duplicate document {current_doc_id}: {e}\")\n # Continue with other documents instead of failing the entire migration\n\n\ndef upgrade() -> None:\n if SKIP_CANON_DRIVE_IDS:\n return\n current_search_settings, _ = active_search_settings()\n\n # Get the index name\n if hasattr(current_search_settings, \"index_name\"):\n index_name = current_search_settings.index_name\n else:\n # Default index name if we can't get it from the document_index\n index_name = \"danswer_index\"\n\n # Get all Google Drive documents from the database (this is faster and more reliable)\n gdrive_documents = get_google_drive_documents_from_database()\n\n if not gdrive_documents:\n return\n\n # Track normalized document IDs to detect duplicates\n all_normalized_doc_ids = set()\n updated_count = 0\n\n for doc_info in gdrive_documents:\n current_doc_id = doc_info[\"document_id\"]\n normalized_doc_id = normalize_google_drive_url(current_doc_id)\n\n print(f\"Processing document {current_doc_id} -> {normalized_doc_id}\")\n # Check for duplicates\n if normalized_doc_id in all_normalized_doc_ids:\n # print(f\"Deleting duplicate document {current_doc_id}\")\n delete_document_from_db(current_doc_id, index_name)\n continue\n\n all_normalized_doc_ids.add(normalized_doc_id)\n\n # If the document ID already doesn't have query parameters, skip it\n if current_doc_id == normalized_doc_id:\n # print(f\"Skipping document {current_doc_id} -> {normalized_doc_id} because it already has no query parameters\")\n continue\n\n try:\n # Update both database and Vespa in order\n # Database first to ensure consistency\n update_document_id_in_database(\n current_doc_id, normalized_doc_id, index_name\n )\n\n # For Vespa, we can now use the original document IDs since we're using contains matching\n update_document_id_in_vespa(index_name, current_doc_id, normalized_doc_id)\n updated_count += 1\n # print(f\"Finished updating document {current_doc_id} -> {normalized_doc_id}\")\n except Exception as e:\n print(f\"Failed to update document {current_doc_id}: {e}\")\n\n if isinstance(e, HTTPStatusError):\n print(f\"HTTPStatusError: {e}\")\n print(f\"Response: {e.response.text}\")\n print(f\"Status: {e.response.status_code}\")\n print(f\"Headers: {e.response.headers}\")\n print(f\"Request: {e.request.url}\")\n print(f\"Request headers: {e.request.headers}\")\n # Note: Rollback is complex with copy-and-swap approach since the old document is already deleted\n # In case of failure, manual intervention may be required\n # Continue with other documents instead of failing the entire migration\n continue\n\n logger.info(f\"Migration complete. Updated {updated_count} Google Drive documents\")\n\n\ndef downgrade() -> None:\n # this is a one way migration, so no downgrade.\n # It wouldn't make sense to store the extra query parameters\n # and duplicate documents to allow a reversal.\n pass\n" }, { "path": "backend/alembic/versions/15326fcec57e_introduce_onyx_apis.py", "content": "\"\"\"Introduce Onyx APIs\n\nRevision ID: 15326fcec57e\nRevises: 77d07dffae64\nCreate Date: 2023-11-11 20:51:24.228999\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom onyx.configs.constants import DocumentSource\n\n# revision identifiers, used by Alembic.\nrevision = \"15326fcec57e\"\ndown_revision = \"77d07dffae64\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"credential\", \"is_admin\", new_column_name=\"admin_public\")\n op.add_column(\n \"document\",\n sa.Column(\"from_ingestion_api\", sa.Boolean(), nullable=True),\n )\n op.alter_column(\n \"connector\",\n \"source\",\n type_=sa.String(length=50),\n existing_type=sa.Enum(DocumentSource, native_enum=False),\n existing_nullable=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"document\", \"from_ingestion_api\")\n op.alter_column(\"credential\", \"admin_public\", new_column_name=\"is_admin\")\n" }, { "path": "backend/alembic/versions/16c37a30adf2_user_file_relationship_migration.py", "content": "\"\"\"Migration 3: User file relationship migration\n\nRevision ID: 16c37a30adf2\nRevises: 0cd424f32b1d\nCreate Date: 2025-09-22 09:47:34.175596\n\nThis migration converts folder-based relationships to project-based relationships.\nIt migrates persona__user_folder to persona__user_file and populates project__user_file.\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy import text\nimport logging\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n# revision identifiers, used by Alembic.\nrevision = \"16c37a30adf2\"\ndown_revision = \"0cd424f32b1d\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Migrate folder-based relationships to project-based relationships.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n # === Step 1: Migrate persona__user_folder to persona__user_file ===\n table_names = inspector.get_table_names()\n\n if \"persona__user_folder\" in table_names and \"user_file\" in table_names:\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n has_new_id = \"new_id\" in user_file_columns\n\n if has_new_id and \"folder_id\" in user_file_columns:\n logger.info(\n \"Migrating persona__user_folder relationships to persona__user_file...\"\n )\n\n # Count relationships to migrate (asyncpg-compatible)\n count_query = text(\n \"\"\"\n SELECT COUNT(*)\n FROM (\n SELECT DISTINCT puf.persona_id, uf.id\n FROM persona__user_folder puf\n JOIN user_file uf ON uf.folder_id = puf.user_folder_id\n WHERE NOT EXISTS (\n SELECT 1\n FROM persona__user_file p2\n WHERE p2.persona_id = puf.persona_id\n AND p2.user_file_id = uf.id\n )\n ) AS distinct_pairs\n \"\"\"\n )\n to_migrate = bind.execute(count_query).scalar_one()\n\n if to_migrate > 0:\n logger.info(f\"Creating {to_migrate} persona-file relationships...\")\n\n # Migrate in batches to avoid memory issues\n batch_size = 10000\n total_inserted = 0\n\n while True:\n # Insert batch directly using subquery (asyncpg compatible)\n result = bind.execute(\n text(\n \"\"\"\n INSERT INTO persona__user_file (persona_id, user_file_id, user_file_id_uuid)\n SELECT DISTINCT puf.persona_id, uf.id as file_id, uf.new_id\n FROM persona__user_folder puf\n JOIN user_file uf ON uf.folder_id = puf.user_folder_id\n WHERE NOT EXISTS (\n SELECT 1\n FROM persona__user_file p2\n WHERE p2.persona_id = puf.persona_id\n AND p2.user_file_id = uf.id\n )\n LIMIT :batch_size\n \"\"\"\n ),\n {\"batch_size\": batch_size},\n )\n\n inserted = result.rowcount\n total_inserted += inserted\n\n if inserted < batch_size:\n break\n\n logger.info(\n f\" Migrated {total_inserted}/{to_migrate} relationships...\"\n )\n\n logger.info(\n f\"Created {total_inserted} persona__user_file relationships\"\n )\n\n # === Step 2: Add foreign key for chat_session.project_id ===\n chat_session_fks = inspector.get_foreign_keys(\"chat_session\")\n fk_exists = any(\n fk[\"name\"] == \"fk_chat_session_project_id\" for fk in chat_session_fks\n )\n\n if not fk_exists:\n logger.info(\"Adding foreign key constraint for chat_session.project_id...\")\n op.create_foreign_key(\n \"fk_chat_session_project_id\",\n \"chat_session\",\n \"user_project\",\n [\"project_id\"],\n [\"id\"],\n )\n logger.info(\"Added foreign key constraint\")\n\n # === Step 3: Populate project__user_file from user_file.folder_id ===\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n has_new_id = \"new_id\" in user_file_columns\n\n if has_new_id and \"folder_id\" in user_file_columns:\n logger.info(\"Populating project__user_file from folder relationships...\")\n\n # Count relationships to create\n count_query = text(\n \"\"\"\n SELECT COUNT(*)\n FROM user_file uf\n WHERE uf.folder_id IS NOT NULL\n AND NOT EXISTS (\n SELECT 1\n FROM project__user_file puf\n WHERE puf.project_id = uf.folder_id\n AND puf.user_file_id = uf.new_id\n )\n \"\"\"\n )\n to_create = bind.execute(count_query).scalar_one()\n\n if to_create > 0:\n logger.info(f\"Creating {to_create} project-file relationships...\")\n\n # Insert in batches\n batch_size = 10000\n total_inserted = 0\n\n while True:\n result = bind.execute(\n text(\n \"\"\"\n INSERT INTO project__user_file (project_id, user_file_id)\n SELECT uf.folder_id, uf.new_id\n FROM user_file uf\n WHERE uf.folder_id IS NOT NULL\n AND NOT EXISTS (\n SELECT 1\n FROM project__user_file puf\n WHERE puf.project_id = uf.folder_id\n AND puf.user_file_id = uf.new_id\n )\n LIMIT :batch_size\n ON CONFLICT (project_id, user_file_id) DO NOTHING\n \"\"\"\n ),\n {\"batch_size\": batch_size},\n )\n\n inserted = result.rowcount\n total_inserted += inserted\n\n if inserted < batch_size:\n break\n\n logger.info(f\" Created {total_inserted}/{to_create} relationships...\")\n\n logger.info(f\"Created {total_inserted} project__user_file relationships\")\n\n # === Step 4: Create index on chat_session.project_id ===\n try:\n indexes = [ix.get(\"name\") for ix in inspector.get_indexes(\"chat_session\")]\n except Exception:\n indexes = []\n\n if \"ix_chat_session_project_id\" not in indexes:\n logger.info(\"Creating index on chat_session.project_id...\")\n op.create_index(\n \"ix_chat_session_project_id\", \"chat_session\", [\"project_id\"], unique=False\n )\n logger.info(\"Created index\")\n\n logger.info(\"Migration 3 (relationship migration) completed successfully\")\n\n\ndef downgrade() -> None:\n \"\"\"Remove migrated relationships and constraints.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n logger.info(\"Starting downgrade of relationship migration...\")\n\n # Drop index on chat_session.project_id\n try:\n indexes = [ix.get(\"name\") for ix in inspector.get_indexes(\"chat_session\")]\n if \"ix_chat_session_project_id\" in indexes:\n op.drop_index(\"ix_chat_session_project_id\", \"chat_session\")\n logger.info(\"Dropped index on chat_session.project_id\")\n except Exception:\n pass\n\n # Drop foreign key constraint\n try:\n chat_session_fks = inspector.get_foreign_keys(\"chat_session\")\n fk_exists = any(\n fk[\"name\"] == \"fk_chat_session_project_id\" for fk in chat_session_fks\n )\n if fk_exists:\n op.drop_constraint(\n \"fk_chat_session_project_id\", \"chat_session\", type_=\"foreignkey\"\n )\n logger.info(\"Dropped foreign key constraint on chat_session.project_id\")\n except Exception:\n pass\n\n # Clear project__user_file relationships (but keep the table for migration 1 to handle)\n if \"project__user_file\" in inspector.get_table_names():\n result = bind.execute(text(\"DELETE FROM project__user_file\"))\n logger.info(f\"Cleared {result.rowcount} records from project__user_file\")\n\n # Remove migrated persona__user_file relationships\n # Only remove those that came from folder relationships\n if all(\n table in inspector.get_table_names()\n for table in [\"persona__user_file\", \"persona__user_folder\", \"user_file\"]\n ):\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n if \"folder_id\" in user_file_columns:\n result = bind.execute(\n text(\n \"\"\"\n DELETE FROM persona__user_file puf\n WHERE EXISTS (\n SELECT 1\n FROM user_file uf\n JOIN persona__user_folder puf2\n ON puf2.user_folder_id = uf.folder_id\n WHERE puf.persona_id = puf2.persona_id\n AND puf.user_file_id = uf.id\n )\n \"\"\"\n )\n )\n logger.info(\n f\"Removed {result.rowcount} migrated persona__user_file relationships\"\n )\n\n logger.info(\"Downgrade completed successfully\")\n" }, { "path": "backend/alembic/versions/173cae5bba26_port_config_store.py", "content": "\"\"\"Port Config Store\n\nRevision ID: 173cae5bba26\nRevises: e50154680a5c\nCreate Date: 2024-03-19 15:30:44.425436\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"173cae5bba26\"\ndown_revision = \"e50154680a5c\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"key_value_store\",\n sa.Column(\"key\", sa.String(), nullable=False),\n sa.Column(\"value\", postgresql.JSONB(astext_type=sa.Text()), nullable=False),\n sa.PrimaryKeyConstraint(\"key\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"key_value_store\")\n" }, { "path": "backend/alembic/versions/175ea04c7087_add_user_preferences.py", "content": "\"\"\"add_user_preferences\n\nRevision ID: 175ea04c7087\nRevises: d56ffa94ca32\nCreate Date: 2026-02-04 18:16:24.830873\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"175ea04c7087\"\ndown_revision = \"d56ffa94ca32\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\"user_preferences\", sa.Text(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"user_preferences\")\n" }, { "path": "backend/alembic/versions/177de57c21c9_display_custom_llm_models.py", "content": "\"\"\"display custom llm models\n\nRevision ID: 177de57c21c9\nRevises: 4ee1287bd26a\nCreate Date: 2024-11-21 11:49:04.488677\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy import and_\n\nrevision = \"177de57c21c9\"\ndown_revision = \"4ee1287bd26a\"\nbranch_labels = None\ndepends_on = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n llm_provider = sa.table(\n \"llm_provider\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"provider\", sa.String),\n sa.column(\"model_names\", postgresql.ARRAY(sa.String)),\n sa.column(\"display_model_names\", postgresql.ARRAY(sa.String)),\n )\n\n excluded_providers = [\"openai\", \"bedrock\", \"anthropic\", \"azure\"]\n\n providers_to_update = sa.select(\n llm_provider.c.id,\n llm_provider.c.model_names,\n llm_provider.c.display_model_names,\n ).where(\n and_(\n ~llm_provider.c.provider.in_(excluded_providers),\n llm_provider.c.model_names.isnot(None),\n )\n )\n\n results = conn.execute(providers_to_update).fetchall()\n\n for provider_id, model_names, display_model_names in results:\n if display_model_names is None:\n display_model_names = []\n\n combined_model_names = list(set(display_model_names + model_names))\n update_stmt = (\n llm_provider.update()\n .where(llm_provider.c.id == provider_id)\n .values(display_model_names=combined_model_names)\n )\n conn.execute(update_stmt)\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/18b5b2524446_add_is_clarification_to_chat_message.py", "content": "\"\"\"add is_clarification to chat_message\n\nRevision ID: 18b5b2524446\nRevises: 87c52ec39f84\nCreate Date: 2025-01-16\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"18b5b2524446\"\ndown_revision = \"87c52ec39f84\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"is_clarification\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"is_clarification\")\n" }, { "path": "backend/alembic/versions/19c0ccb01687_migrate_to_contextual_rag_model.py", "content": "\"\"\"Migrate to contextual rag model\n\nRevision ID: 19c0ccb01687\nRevises: 9c54986124c6\nCreate Date: 2026-02-12 11:21:41.798037\n\n\"\"\"\n\nimport sqlalchemy as sa\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"19c0ccb01687\"\ndown_revision = \"9c54986124c6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Widen the column to fit 'CONTEXTUAL_RAG' (15 chars); was varchar(10)\n # when the table was created with only CHAT/VISION values.\n op.alter_column(\n \"llm_model_flow\",\n \"llm_model_flow_type\",\n type_=sa.String(length=20),\n existing_type=sa.String(length=10),\n existing_nullable=False,\n )\n\n # For every search_settings row that has contextual rag configured,\n # create an llm_model_flow entry. is_default is TRUE if the row\n # belongs to the PRESENT search settings, FALSE otherwise.\n op.execute(\n \"\"\"\n INSERT INTO llm_model_flow (llm_model_flow_type, model_configuration_id, is_default)\n SELECT DISTINCT\n 'CONTEXTUAL_RAG',\n mc.id,\n (ss.status = 'PRESENT')\n FROM search_settings ss\n JOIN llm_provider lp\n ON lp.name = ss.contextual_rag_llm_provider\n JOIN model_configuration mc\n ON mc.llm_provider_id = lp.id\n AND mc.name = ss.contextual_rag_llm_name\n WHERE ss.enable_contextual_rag = TRUE\n AND ss.contextual_rag_llm_name IS NOT NULL\n AND ss.contextual_rag_llm_provider IS NOT NULL\n ON CONFLICT (llm_model_flow_type, model_configuration_id)\n DO UPDATE SET is_default = EXCLUDED.is_default\n WHERE EXCLUDED.is_default = TRUE\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\n \"\"\"\n DELETE FROM llm_model_flow\n WHERE llm_model_flow_type = 'CONTEXTUAL_RAG'\n \"\"\"\n )\n\n op.alter_column(\n \"llm_model_flow\",\n \"llm_model_flow_type\",\n type_=sa.String(length=10),\n existing_type=sa.String(length=20),\n existing_nullable=False,\n )\n" }, { "path": "backend/alembic/versions/1a03d2c2856b_add_indexes_to_document__tag.py", "content": "\"\"\"Add indexes to document__tag\n\nRevision ID: 1a03d2c2856b\nRevises: 9c00a2bccb83\nCreate Date: 2025-02-18 10:45:13.957807\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"1a03d2c2856b\"\ndown_revision = \"9c00a2bccb83\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_index(\n op.f(\"ix_document__tag_tag_id\"),\n \"document__tag\",\n [\"tag_id\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(op.f(\"ix_document__tag_tag_id\"), table_name=\"document__tag\")\n" }, { "path": "backend/alembic/versions/1b10e1fda030_add_additional_data_to_notifications.py", "content": "\"\"\"add additional data to notifications\n\nRevision ID: 1b10e1fda030\nRevises: 6756efa39ada\nCreate Date: 2024-10-15 19:26:44.071259\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"1b10e1fda030\"\ndown_revision = \"6756efa39ada\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"notification\", sa.Column(\"additional_data\", postgresql.JSONB(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"notification\", \"additional_data\")\n" }, { "path": "backend/alembic/versions/1b8206b29c5d_add_user_delete_cascades.py", "content": "\"\"\"add_user_delete_cascades\n\nRevision ID: 1b8206b29c5d\nRevises: 35e6853a51d5\nCreate Date: 2024-09-18 11:48:59.418726\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"1b8206b29c5d\"\ndown_revision = \"35e6853a51d5\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_constraint(\"credential_user_id_fkey\", \"credential\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"credential_user_id_fkey\",\n \"credential\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n op.drop_constraint(\"chat_session_user_id_fkey\", \"chat_session\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"chat_session_user_id_fkey\",\n \"chat_session\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n op.drop_constraint(\"chat_folder_user_id_fkey\", \"chat_folder\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"chat_folder_user_id_fkey\",\n \"chat_folder\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n op.drop_constraint(\"prompt_user_id_fkey\", \"prompt\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"prompt_user_id_fkey\", \"prompt\", \"user\", [\"user_id\"], [\"id\"], ondelete=\"CASCADE\"\n )\n\n op.drop_constraint(\"notification_user_id_fkey\", \"notification\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"notification_user_id_fkey\",\n \"notification\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n op.drop_constraint(\"inputprompt_user_id_fkey\", \"inputprompt\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"inputprompt_user_id_fkey\",\n \"inputprompt\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\"credential_user_id_fkey\", \"credential\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"credential_user_id_fkey\", \"credential\", \"user\", [\"user_id\"], [\"id\"]\n )\n\n op.drop_constraint(\"chat_session_user_id_fkey\", \"chat_session\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"chat_session_user_id_fkey\", \"chat_session\", \"user\", [\"user_id\"], [\"id\"]\n )\n\n op.drop_constraint(\"chat_folder_user_id_fkey\", \"chat_folder\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"chat_folder_user_id_fkey\", \"chat_folder\", \"user\", [\"user_id\"], [\"id\"]\n )\n\n op.drop_constraint(\"prompt_user_id_fkey\", \"prompt\", type_=\"foreignkey\")\n op.create_foreign_key(\"prompt_user_id_fkey\", \"prompt\", \"user\", [\"user_id\"], [\"id\"])\n\n op.drop_constraint(\"notification_user_id_fkey\", \"notification\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"notification_user_id_fkey\", \"notification\", \"user\", [\"user_id\"], [\"id\"]\n )\n\n op.drop_constraint(\"inputprompt_user_id_fkey\", \"inputprompt\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"inputprompt_user_id_fkey\", \"inputprompt\", \"user\", [\"user_id\"], [\"id\"]\n )\n" }, { "path": "backend/alembic/versions/1d78c0ca7853_remove_voice_provider_deleted_column.py", "content": "\"\"\"remove voice_provider deleted column\n\nRevision ID: 1d78c0ca7853\nRevises: a3f8b2c1d4e5\nCreate Date: 2026-03-26 11:30:53.883127\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"1d78c0ca7853\"\ndown_revision = \"a3f8b2c1d4e5\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Hard-delete any soft-deleted rows before dropping the column\n op.execute(\"DELETE FROM voice_provider WHERE deleted = true\")\n op.drop_column(\"voice_provider\", \"deleted\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"voice_provider\",\n sa.Column(\n \"deleted\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.text(\"false\"),\n ),\n )\n" }, { "path": "backend/alembic/versions/1f2a3b4c5d6e_add_internet_search_and_content_providers.py", "content": "\"\"\"add internet search and content provider tables\n\nRevision ID: 1f2a3b4c5d6e\nRevises: 9drpiiw74ljy\nCreate Date: 2025-11-10 19:45:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"1f2a3b4c5d6e\"\ndown_revision = \"9drpiiw74ljy\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"internet_search_provider\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"name\", sa.String(), nullable=False, unique=True),\n sa.Column(\"provider_type\", sa.String(), nullable=False),\n sa.Column(\"api_key\", sa.LargeBinary(), nullable=True),\n sa.Column(\"config\", postgresql.JSONB(astext_type=sa.Text()), nullable=True),\n sa.Column(\n \"is_active\", sa.Boolean(), nullable=False, server_default=sa.text(\"false\")\n ),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.text(\"now()\"),\n ),\n )\n op.create_index(\n \"ix_internet_search_provider_is_active\",\n \"internet_search_provider\",\n [\"is_active\"],\n )\n\n op.create_table(\n \"internet_content_provider\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"name\", sa.String(), nullable=False, unique=True),\n sa.Column(\"provider_type\", sa.String(), nullable=False),\n sa.Column(\"api_key\", sa.LargeBinary(), nullable=True),\n sa.Column(\"config\", postgresql.JSONB(astext_type=sa.Text()), nullable=True),\n sa.Column(\n \"is_active\", sa.Boolean(), nullable=False, server_default=sa.text(\"false\")\n ),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.text(\"now()\"),\n ),\n )\n op.create_index(\n \"ix_internet_content_provider_is_active\",\n \"internet_content_provider\",\n [\"is_active\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"ix_internet_content_provider_is_active\", table_name=\"internet_content_provider\"\n )\n op.drop_table(\"internet_content_provider\")\n op.drop_index(\n \"ix_internet_search_provider_is_active\", table_name=\"internet_search_provider\"\n )\n op.drop_table(\"internet_search_provider\")\n" }, { "path": "backend/alembic/versions/1f60f60c3401_embedding_model_search_settings.py", "content": "\"\"\"embedding model -> search settings\n\nRevision ID: 1f60f60c3401\nRevises: f17bf3b0d9f1\nCreate Date: 2024-08-25 12:39:51.731632\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"1f60f60c3401\"\ndown_revision = \"f17bf3b0d9f1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_constraint(\n \"index_attempt__embedding_model_fk\", \"index_attempt\", type_=\"foreignkey\"\n )\n # Rename the table\n op.rename_table(\"embedding_model\", \"search_settings\")\n\n # Add new columns\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"multipass_indexing\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n )\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"multilingual_expansion\",\n postgresql.ARRAY(sa.String()),\n nullable=False,\n server_default=\"{}\",\n ),\n )\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"disable_rerank_for_streaming\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_model_name\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_provider_type\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_api_key\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"num_rerank\",\n sa.Integer(),\n nullable=False,\n server_default=str(20),\n ),\n )\n\n # Add the new column as nullable initially\n op.add_column(\n \"index_attempt\", sa.Column(\"search_settings_id\", sa.Integer(), nullable=True)\n )\n\n # Populate the new column with data from the existing embedding_model_id\n op.execute(\"UPDATE index_attempt SET search_settings_id = embedding_model_id\")\n\n # Create the foreign key constraint\n op.create_foreign_key(\n \"fk_index_attempt_search_settings\",\n \"index_attempt\",\n \"search_settings\",\n [\"search_settings_id\"],\n [\"id\"],\n )\n\n # Make the new column non-nullable\n op.alter_column(\"index_attempt\", \"search_settings_id\", nullable=False)\n\n # Drop the old embedding_model_id column\n op.drop_column(\"index_attempt\", \"embedding_model_id\")\n\n\ndef downgrade() -> None:\n # Add back the embedding_model_id column\n op.add_column(\n \"index_attempt\", sa.Column(\"embedding_model_id\", sa.Integer(), nullable=True)\n )\n\n # Populate the old column with data from search_settings_id\n op.execute(\"UPDATE index_attempt SET embedding_model_id = search_settings_id\")\n\n # Make the old column non-nullable\n op.alter_column(\"index_attempt\", \"embedding_model_id\", nullable=False)\n\n # Drop the foreign key constraint\n op.drop_constraint(\n \"fk_index_attempt_search_settings\", \"index_attempt\", type_=\"foreignkey\"\n )\n\n # Drop the new search_settings_id column\n op.drop_column(\"index_attempt\", \"search_settings_id\")\n\n # Rename the table back\n op.rename_table(\"search_settings\", \"embedding_model\")\n\n # Remove added columns\n op.drop_column(\"embedding_model\", \"num_rerank\")\n op.drop_column(\"embedding_model\", \"rerank_api_key\")\n op.drop_column(\"embedding_model\", \"rerank_provider_type\")\n op.drop_column(\"embedding_model\", \"rerank_model_name\")\n op.drop_column(\"embedding_model\", \"disable_rerank_for_streaming\")\n op.drop_column(\"embedding_model\", \"multilingual_expansion\")\n op.drop_column(\"embedding_model\", \"multipass_indexing\")\n\n op.create_foreign_key(\n \"index_attempt__embedding_model_fk\",\n \"index_attempt\",\n \"embedding_model\",\n [\"embedding_model_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/2020d417ec84_single_onyx_craft_migration.py", "content": "\"\"\"single onyx craft migration\n\nConsolidates all buildmode/onyx craft tables into a single migration.\n\nTables created:\n- build_session: User build sessions with status tracking\n- sandbox: User-owned containerized environments (one per user)\n- artifact: Build output files (web apps, documents, images)\n- snapshot: Sandbox filesystem snapshots\n- build_message: Conversation messages for build sessions\n\nExisting table modified:\n- connector_credential_pair: Added processing_mode column\n\nRevision ID: 2020d417ec84\nRevises: 41fa44bef321\nCreate Date: 2026-01-26 14:43:54.641405\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2020d417ec84\"\ndown_revision = \"41fa44bef321\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # ==========================================================================\n # ENUMS\n # ==========================================================================\n\n # Build session status enum\n build_session_status_enum = sa.Enum(\n \"active\",\n \"idle\",\n name=\"buildsessionstatus\",\n native_enum=False,\n )\n\n # Sandbox status enum\n sandbox_status_enum = sa.Enum(\n \"provisioning\",\n \"running\",\n \"idle\",\n \"sleeping\",\n \"terminated\",\n \"failed\",\n name=\"sandboxstatus\",\n native_enum=False,\n )\n\n # Artifact type enum\n artifact_type_enum = sa.Enum(\n \"web_app\",\n \"pptx\",\n \"docx\",\n \"markdown\",\n \"excel\",\n \"image\",\n name=\"artifacttype\",\n native_enum=False,\n )\n\n # ==========================================================================\n # BUILD_SESSION TABLE\n # ==========================================================================\n\n op.create_table(\n \"build_session\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), primary_key=True),\n sa.Column(\n \"user_id\",\n postgresql.UUID(as_uuid=True),\n sa.ForeignKey(\"user.id\", ondelete=\"CASCADE\"),\n nullable=True,\n ),\n sa.Column(\"name\", sa.String(), nullable=True),\n sa.Column(\n \"status\",\n build_session_status_enum,\n nullable=False,\n server_default=\"active\",\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"last_activity_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"nextjs_port\", sa.Integer(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_index(\n \"ix_build_session_user_created\",\n \"build_session\",\n [\"user_id\", sa.text(\"created_at DESC\")],\n unique=False,\n )\n op.create_index(\n \"ix_build_session_status\",\n \"build_session\",\n [\"status\"],\n unique=False,\n )\n\n # ==========================================================================\n # SANDBOX TABLE (user-owned, one per user)\n # ==========================================================================\n\n op.create_table(\n \"sandbox\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), primary_key=True),\n sa.Column(\n \"user_id\",\n postgresql.UUID(as_uuid=True),\n sa.ForeignKey(\"user.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\"container_id\", sa.String(), nullable=True),\n sa.Column(\n \"status\",\n sandbox_status_enum,\n nullable=False,\n server_default=\"provisioning\",\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"last_heartbeat\", sa.DateTime(timezone=True), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"user_id\", name=\"sandbox_user_id_key\"),\n )\n\n op.create_index(\n \"ix_sandbox_status\",\n \"sandbox\",\n [\"status\"],\n unique=False,\n )\n op.create_index(\n \"ix_sandbox_container_id\",\n \"sandbox\",\n [\"container_id\"],\n unique=False,\n )\n\n # ==========================================================================\n # ARTIFACT TABLE\n # ==========================================================================\n\n op.create_table(\n \"artifact\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), primary_key=True),\n sa.Column(\n \"session_id\",\n postgresql.UUID(as_uuid=True),\n sa.ForeignKey(\"build_session.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\"type\", artifact_type_enum, nullable=False),\n sa.Column(\"path\", sa.String(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_index(\n \"ix_artifact_session_created\",\n \"artifact\",\n [\"session_id\", sa.text(\"created_at DESC\")],\n unique=False,\n )\n op.create_index(\n \"ix_artifact_type\",\n \"artifact\",\n [\"type\"],\n unique=False,\n )\n\n # ==========================================================================\n # SNAPSHOT TABLE\n # ==========================================================================\n\n op.create_table(\n \"snapshot\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), primary_key=True),\n sa.Column(\n \"session_id\",\n postgresql.UUID(as_uuid=True),\n sa.ForeignKey(\"build_session.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\"storage_path\", sa.String(), nullable=False),\n sa.Column(\"size_bytes\", sa.BigInteger(), nullable=False, server_default=\"0\"),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_index(\n \"ix_snapshot_session_created\",\n \"snapshot\",\n [\"session_id\", sa.text(\"created_at DESC\")],\n unique=False,\n )\n\n # ==========================================================================\n # BUILD_MESSAGE TABLE\n # ==========================================================================\n\n op.create_table(\n \"build_message\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), primary_key=True),\n sa.Column(\n \"session_id\",\n postgresql.UUID(as_uuid=True),\n sa.ForeignKey(\"build_session.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\n \"turn_index\",\n sa.Integer(),\n nullable=False,\n ),\n sa.Column(\n \"type\",\n sa.Enum(\n \"SYSTEM\",\n \"USER\",\n \"ASSISTANT\",\n \"DANSWER\",\n name=\"messagetype\",\n create_type=False,\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\n \"message_metadata\",\n postgresql.JSONB(),\n nullable=False,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_index(\n \"ix_build_message_session_turn\",\n \"build_message\",\n [\"session_id\", \"turn_index\", sa.text(\"created_at ASC\")],\n unique=False,\n )\n\n # ==========================================================================\n # CONNECTOR_CREDENTIAL_PAIR MODIFICATION\n # ==========================================================================\n\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"processing_mode\",\n sa.String(),\n nullable=False,\n server_default=\"regular\",\n ),\n )\n\n\ndef downgrade() -> None:\n # ==========================================================================\n # CONNECTOR_CREDENTIAL_PAIR MODIFICATION\n # ==========================================================================\n\n op.drop_column(\"connector_credential_pair\", \"processing_mode\")\n\n # ==========================================================================\n # BUILD_MESSAGE TABLE\n # ==========================================================================\n\n op.drop_index(\"ix_build_message_session_turn\", table_name=\"build_message\")\n op.drop_table(\"build_message\")\n\n # ==========================================================================\n # SNAPSHOT TABLE\n # ==========================================================================\n\n op.drop_index(\"ix_snapshot_session_created\", table_name=\"snapshot\")\n op.drop_table(\"snapshot\")\n\n # ==========================================================================\n # ARTIFACT TABLE\n # ==========================================================================\n\n op.drop_index(\"ix_artifact_type\", table_name=\"artifact\")\n op.drop_index(\"ix_artifact_session_created\", table_name=\"artifact\")\n op.drop_table(\"artifact\")\n sa.Enum(name=\"artifacttype\").drop(op.get_bind(), checkfirst=True)\n\n # ==========================================================================\n # SANDBOX TABLE\n # ==========================================================================\n\n op.drop_index(\"ix_sandbox_container_id\", table_name=\"sandbox\")\n op.drop_index(\"ix_sandbox_status\", table_name=\"sandbox\")\n op.drop_table(\"sandbox\")\n sa.Enum(name=\"sandboxstatus\").drop(op.get_bind(), checkfirst=True)\n\n # ==========================================================================\n # BUILD_SESSION TABLE\n # ==========================================================================\n\n op.drop_index(\"ix_build_session_status\", table_name=\"build_session\")\n op.drop_index(\"ix_build_session_user_created\", table_name=\"build_session\")\n op.drop_table(\"build_session\")\n sa.Enum(name=\"buildsessionstatus\").drop(op.get_bind(), checkfirst=True)\n" }, { "path": "backend/alembic/versions/213fd978c6d8_notifications.py", "content": "\"\"\"notifications\n\nRevision ID: 213fd978c6d8\nRevises: 5fc1f54cc252\nCreate Date: 2024-08-10 11:13:36.070790\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"213fd978c6d8\"\ndown_revision = \"5fc1f54cc252\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"notification\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"notif_type\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\n \"user_id\",\n sa.UUID(),\n nullable=True,\n ),\n sa.Column(\"dismissed\", sa.Boolean(), nullable=False),\n sa.Column(\"last_shown\", sa.DateTime(timezone=True), nullable=False),\n sa.Column(\"first_shown\", sa.DateTime(timezone=True), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"notification\")\n" }, { "path": "backend/alembic/versions/238b84885828_add_foreign_key_to_user__external_user_.py", "content": "\"\"\"Add foreign key to user__external_user_group_id\n\nRevision ID: 238b84885828\nRevises: a7688ab35c45\nCreate Date: 2025-05-19 17:15:33.424584\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"238b84885828\"\ndown_revision = \"a7688ab35c45\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # First, clean up any entries that don't have a valid cc_pair_id\n op.execute(\n \"\"\"\n DELETE FROM user__external_user_group_id\n WHERE cc_pair_id NOT IN (SELECT id FROM connector_credential_pair)\n \"\"\"\n )\n\n # Add foreign key constraint with cascade delete\n op.create_foreign_key(\n \"fk_user__external_user_group_id_cc_pair_id\",\n \"user__external_user_group_id\",\n \"connector_credential_pair\",\n [\"cc_pair_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n # Drop the foreign key constraint\n op.drop_constraint(\n \"fk_user__external_user_group_id_cc_pair_id\",\n \"user__external_user_group_id\",\n type_=\"foreignkey\",\n )\n" }, { "path": "backend/alembic/versions/23957775e5f5_remove_feedback_foreignkey_constraint.py", "content": "\"\"\"remove-feedback-foreignkey-constraint\n\nRevision ID: 23957775e5f5\nRevises: bc9771dccadf\nCreate Date: 2024-06-27 16:04:51.480437\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"23957775e5f5\"\ndown_revision = \"bc9771dccadf\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_constraint(\n \"chat_feedback__chat_message_fk\", \"chat_feedback\", type_=\"foreignkey\"\n )\n op.create_foreign_key(\n \"chat_feedback__chat_message_fk\",\n \"chat_feedback\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n op.alter_column(\n \"chat_feedback\", \"chat_message_id\", existing_type=sa.Integer(), nullable=True\n )\n op.drop_constraint(\n \"document_retrieval_feedback__chat_message_fk\",\n \"document_retrieval_feedback\",\n type_=\"foreignkey\",\n )\n op.create_foreign_key(\n \"document_retrieval_feedback__chat_message_fk\",\n \"document_retrieval_feedback\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n op.alter_column(\n \"document_retrieval_feedback\",\n \"chat_message_id\",\n existing_type=sa.Integer(),\n nullable=True,\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"chat_feedback\", \"chat_message_id\", existing_type=sa.Integer(), nullable=False\n )\n op.drop_constraint(\n \"chat_feedback__chat_message_fk\", \"chat_feedback\", type_=\"foreignkey\"\n )\n op.create_foreign_key(\n \"chat_feedback__chat_message_fk\",\n \"chat_feedback\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n )\n\n op.alter_column(\n \"document_retrieval_feedback\",\n \"chat_message_id\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n op.drop_constraint(\n \"document_retrieval_feedback__chat_message_fk\",\n \"document_retrieval_feedback\",\n type_=\"foreignkey\",\n )\n op.create_foreign_key(\n \"document_retrieval_feedback__chat_message_fk\",\n \"document_retrieval_feedback\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/25a5501dc766_group_permissions_phase1.py", "content": "\"\"\"group_permissions_phase1\n\nRevision ID: 25a5501dc766\nRevises: b728689f45b1\nCreate Date: 2026-03-23 11:41:25.557442\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\n\nfrom onyx.db.enums import AccountType\nfrom onyx.db.enums import GrantSource\nfrom onyx.db.enums import Permission\n\n\n# revision identifiers, used by Alembic.\nrevision = \"25a5501dc766\"\ndown_revision = \"b728689f45b1\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # 1. Add account_type column to user table (nullable for now).\n # TODO(subash): backfill account_type for existing rows and add NOT NULL.\n op.add_column(\n \"user\",\n sa.Column(\n \"account_type\",\n sa.Enum(AccountType, native_enum=False),\n nullable=True,\n ),\n )\n\n # 2. Add is_default column to user_group table\n op.add_column(\n \"user_group\",\n sa.Column(\n \"is_default\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n )\n\n # 3. Create permission_grant table\n op.create_table(\n \"permission_grant\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\"group_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"permission\",\n sa.Enum(Permission, native_enum=False),\n nullable=False,\n ),\n sa.Column(\n \"grant_source\",\n sa.Enum(GrantSource, native_enum=False),\n nullable=False,\n ),\n sa.Column(\n \"granted_by\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\n \"granted_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\n \"is_deleted\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.ForeignKeyConstraint(\n [\"group_id\"],\n [\"user_group.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.ForeignKeyConstraint(\n [\"granted_by\"],\n [\"user.id\"],\n ondelete=\"SET NULL\",\n ),\n sa.UniqueConstraint(\n \"group_id\", \"permission\", name=\"uq_permission_grant_group_permission\"\n ),\n )\n\n # 4. Index on user__user_group(user_id) — existing composite PK\n # has user_group_id as leading column; user-filtered queries need this\n op.create_index(\n \"ix_user__user_group_user_id\",\n \"user__user_group\",\n [\"user_id\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_user__user_group_user_id\", table_name=\"user__user_group\")\n op.drop_table(\"permission_grant\")\n op.drop_column(\"user_group\", \"is_default\")\n op.drop_column(\"user\", \"account_type\")\n" }, { "path": "backend/alembic/versions/2664261bfaab_add_cache_store_table.py", "content": "\"\"\"add cache_store table\n\nRevision ID: 2664261bfaab\nRevises: 4a1e4b1c89d2\nCreate Date: 2026-02-27 00:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"2664261bfaab\"\ndown_revision = \"4a1e4b1c89d2\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"cache_store\",\n sa.Column(\"key\", sa.String(), nullable=False),\n sa.Column(\"value\", sa.LargeBinary(), nullable=True),\n sa.Column(\"expires_at\", sa.DateTime(timezone=True), nullable=True),\n sa.PrimaryKeyConstraint(\"key\"),\n )\n op.create_index(\n \"ix_cache_store_expires\",\n \"cache_store\",\n [\"expires_at\"],\n postgresql_where=sa.text(\"expires_at IS NOT NULL\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_cache_store_expires\", table_name=\"cache_store\")\n op.drop_table(\"cache_store\")\n" }, { "path": "backend/alembic/versions/2666d766cb9b_google_oauth2.py", "content": "\"\"\"Google OAuth2\n\nRevision ID: 2666d766cb9b\nRevises: 6d387b3196c2\nCreate Date: 2023-05-05 15:49:35.716016\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2666d766cb9b\"\ndown_revision = \"6d387b3196c2\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"oauth_account\",\n sa.Column(\"id\", fastapi_users_db_sqlalchemy.generics.GUID(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\"oauth_name\", sa.String(length=100), nullable=False),\n sa.Column(\"access_token\", sa.String(length=1024), nullable=False),\n sa.Column(\"expires_at\", sa.Integer(), nullable=True),\n sa.Column(\"refresh_token\", sa.String(length=1024), nullable=True),\n sa.Column(\"account_id\", sa.String(length=320), nullable=False),\n sa.Column(\"account_email\", sa.String(length=320), nullable=False),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"cascade\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_index(\n op.f(\"ix_oauth_account_account_id\"),\n \"oauth_account\",\n [\"account_id\"],\n unique=False,\n )\n op.create_index(\n op.f(\"ix_oauth_account_oauth_name\"),\n \"oauth_account\",\n [\"oauth_name\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(op.f(\"ix_oauth_account_oauth_name\"), table_name=\"oauth_account\")\n op.drop_index(op.f(\"ix_oauth_account_account_id\"), table_name=\"oauth_account\")\n op.drop_table(\"oauth_account\")\n" }, { "path": "backend/alembic/versions/26b931506ecb_default_chosen_assistants_to_none.py", "content": "\"\"\"default chosen assistants to none\n\nRevision ID: 26b931506ecb\nRevises: 2daa494a0851\nCreate Date: 2024-11-12 13:23:29.858995\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"26b931506ecb\"\ndown_revision = \"2daa494a0851\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\", sa.Column(\"chosen_assistants_new\", postgresql.JSONB(), nullable=True)\n )\n\n op.execute(\n \"\"\"\n UPDATE \"user\"\n SET chosen_assistants_new =\n CASE\n WHEN chosen_assistants = '[-2, -1, 0]' THEN NULL\n ELSE chosen_assistants\n END\n \"\"\"\n )\n\n op.drop_column(\"user\", \"chosen_assistants\")\n\n op.alter_column(\n \"user\", \"chosen_assistants_new\", new_column_name=\"chosen_assistants\"\n )\n\n\ndef downgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"chosen_assistants_old\",\n postgresql.JSONB(),\n nullable=False,\n server_default=\"[-2, -1, 0]\",\n ),\n )\n\n op.execute(\n \"\"\"\n UPDATE \"user\"\n SET chosen_assistants_old =\n CASE\n WHEN chosen_assistants IS NULL THEN '[-2, -1, 0]'::jsonb\n ELSE chosen_assistants\n END\n \"\"\"\n )\n\n op.drop_column(\"user\", \"chosen_assistants\")\n\n op.alter_column(\n \"user\", \"chosen_assistants_old\", new_column_name=\"chosen_assistants\"\n )\n" }, { "path": "backend/alembic/versions/27c6ecc08586_permission_framework.py", "content": "\"\"\"Permission Framework\n\nRevision ID: 27c6ecc08586\nRevises: 2666d766cb9b\nCreate Date: 2023-05-24 18:45:17.244495\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\nfrom alembic import op\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"27c6ecc08586\"\ndown_revision = \"2666d766cb9b\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\"TRUNCATE TABLE index_attempt\")\n op.create_table(\n \"connector\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\n \"source\",\n sa.Enum(\n \"SLACK\",\n \"WEB\",\n \"GOOGLE_DRIVE\",\n \"GITHUB\",\n \"CONFLUENCE\",\n name=\"documentsource\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\n \"input_type\",\n sa.Enum(\n \"LOAD_STATE\",\n \"POLL\",\n \"EVENT\",\n name=\"inputtype\",\n native_enum=False,\n ),\n nullable=True,\n ),\n sa.Column(\n \"connector_specific_config\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=False,\n ),\n sa.Column(\"refresh_freq\", sa.Integer(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"disabled\", sa.Boolean(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"credential\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"credential_json\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=False,\n ),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"public_doc\", sa.Boolean(), nullable=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"connector_credential_pair\",\n sa.Column(\"connector_id\", sa.Integer(), nullable=False),\n sa.Column(\"credential_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"connector_id\"],\n [\"connector.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"credential_id\"],\n [\"credential.id\"],\n ),\n sa.PrimaryKeyConstraint(\"connector_id\", \"credential_id\"),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\"connector_id\", sa.Integer(), nullable=True),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\"credential_id\", sa.Integer(), nullable=True),\n )\n op.create_foreign_key(\n \"fk_index_attempt_credential_id\",\n \"index_attempt\",\n \"credential\",\n [\"credential_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"fk_index_attempt_connector_id\",\n \"index_attempt\",\n \"connector\",\n [\"connector_id\"],\n [\"id\"],\n )\n op.drop_column(\"index_attempt\", \"connector_specific_config\")\n op.drop_column(\"index_attempt\", \"source\")\n op.drop_column(\"index_attempt\", \"input_type\")\n\n\ndef downgrade() -> None:\n op.execute(\"TRUNCATE TABLE index_attempt\")\n conn = op.get_bind()\n inspector = sa.inspect(conn)\n existing_columns = {col[\"name\"] for col in inspector.get_columns(\"index_attempt\")}\n\n if \"input_type\" not in existing_columns:\n op.add_column(\n \"index_attempt\",\n sa.Column(\"input_type\", sa.VARCHAR(), autoincrement=False, nullable=False),\n )\n\n if \"source\" not in existing_columns:\n op.add_column(\n \"index_attempt\",\n sa.Column(\"source\", sa.VARCHAR(), autoincrement=False, nullable=False),\n )\n\n if \"connector_specific_config\" not in existing_columns:\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"connector_specific_config\",\n postgresql.JSONB(astext_type=sa.Text()),\n autoincrement=False,\n nullable=False,\n ),\n )\n\n # Check if the constraint exists before dropping\n constraints = inspector.get_foreign_keys(\"index_attempt\")\n\n if any(\n constraint[\"name\"] == \"fk_index_attempt_credential_id\"\n for constraint in constraints\n ):\n op.drop_constraint(\n \"fk_index_attempt_credential_id\", \"index_attempt\", type_=\"foreignkey\"\n )\n\n if any(\n constraint[\"name\"] == \"fk_index_attempt_connector_id\"\n for constraint in constraints\n ):\n op.drop_constraint(\n \"fk_index_attempt_connector_id\", \"index_attempt\", type_=\"foreignkey\"\n )\n\n if \"credential_id\" in existing_columns:\n op.drop_column(\"index_attempt\", \"credential_id\")\n\n if \"connector_id\" in existing_columns:\n op.drop_column(\"index_attempt\", \"connector_id\")\n\n op.execute(\"DROP TABLE IF EXISTS connector_credential_pair CASCADE\")\n op.execute(\"DROP TABLE IF EXISTS credential CASCADE\")\n op.execute(\"DROP TABLE IF EXISTS connector CASCADE\")\n" }, { "path": "backend/alembic/versions/27fb147a843f_add_timestamps_to_user_table.py", "content": "\"\"\"add timestamps to user table\n\nRevision ID: 27fb147a843f\nRevises: b5c4d7e8f9a1\nCreate Date: 2026-03-08 17:18:40.828644\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"27fb147a843f\"\ndown_revision = \"b5c4d7e8f9a1\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n )\n op.add_column(\n \"user\",\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"updated_at\")\n op.drop_column(\"user\", \"created_at\")\n" }, { "path": "backend/alembic/versions/2955778aa44c_add_chunk_count_to_document.py", "content": "\"\"\"add chunk count to document\n\nRevision ID: 2955778aa44c\nRevises: c0aab6edb6dd\nCreate Date: 2025-01-04 11:39:43.268612\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2955778aa44c\"\ndown_revision = \"c0aab6edb6dd\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\"document\", sa.Column(\"chunk_count\", sa.Integer(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"document\", \"chunk_count\")\n" }, { "path": "backend/alembic/versions/2a391f840e85_add_last_refreshed_at_mcp_server.py", "content": "\"\"\"add last refreshed at mcp server\n\nRevision ID: 2a391f840e85\nRevises: 4cebcbc9b2ae\nCreate Date: 2025-12-06 15:19:59.766066\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembi.\nrevision = \"2a391f840e85\"\ndown_revision = \"4cebcbc9b2ae\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"mcp_server\",\n sa.Column(\"last_refreshed_at\", sa.DateTime(timezone=True), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"mcp_server\", \"last_refreshed_at\")\n" }, { "path": "backend/alembic/versions/2acdef638fc2_add_switchover_type_field.py", "content": "\"\"\"add switchover_type field and remove background_reindex_enabled\n\nRevision ID: 2acdef638fc2\nRevises: a4f23d6b71c8\nCreate Date: 2025-01-XX XX:XX:XX.XXXXXX\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom onyx.db.enums import SwitchoverType\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2acdef638fc2\"\ndown_revision = \"a4f23d6b71c8\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add switchover_type column with default value of REINDEX\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"switchover_type\",\n sa.Enum(SwitchoverType, native_enum=False),\n nullable=False,\n server_default=SwitchoverType.REINDEX.value,\n ),\n )\n\n # Migrate existing data: set switchover_type based on background_reindex_enabled\n # REINDEX where background_reindex_enabled=True, INSTANT where False\n op.execute(\n \"\"\"\n UPDATE search_settings\n SET switchover_type = CASE\n WHEN background_reindex_enabled = true THEN 'REINDEX'\n ELSE 'INSTANT'\n END\n \"\"\"\n )\n\n # Remove the background_reindex_enabled column (replaced by switchover_type)\n op.drop_column(\"search_settings\", \"background_reindex_enabled\")\n\n\ndef downgrade() -> None:\n # Re-add the background_reindex_enabled column with default value of True\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"background_reindex_enabled\",\n sa.Boolean(),\n nullable=False,\n server_default=\"true\",\n ),\n )\n # Set background_reindex_enabled based on switchover_type\n op.execute(\n \"\"\"\n UPDATE search_settings\n SET background_reindex_enabled = CASE\n WHEN switchover_type = 'INSTANT' THEN false\n ELSE true\n END\n \"\"\"\n )\n # Remove the switchover_type column\n op.drop_column(\"search_settings\", \"switchover_type\")\n" }, { "path": "backend/alembic/versions/2b75d0a8ffcb_user_file_schema_cleanup.py", "content": "\"\"\"Migration 6: User file schema cleanup\n\nRevision ID: 2b75d0a8ffcb\nRevises: 3a78dba1080a\nCreate Date: 2025-09-22 10:09:26.375377\n\nThis migration removes legacy columns and tables after data migration is complete.\nIt should only be run after verifying all data has been successfully migrated.\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy import text\nimport logging\nimport fastapi_users_db_sqlalchemy\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n# revision identifiers, used by Alembic.\nrevision = \"2b75d0a8ffcb\"\ndown_revision = \"3a78dba1080a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Remove legacy columns and tables.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n logger.info(\"Starting schema cleanup...\")\n\n # === Step 1: Verify data migration is complete ===\n logger.info(\"Verifying data migration completion...\")\n\n # Check if any chat sessions still have folder_id references\n chat_session_columns = [\n col[\"name\"] for col in inspector.get_columns(\"chat_session\")\n ]\n if \"folder_id\" in chat_session_columns:\n orphaned_count = bind.execute(\n text(\n \"\"\"\n SELECT COUNT(*) FROM chat_session\n WHERE folder_id IS NOT NULL AND project_id IS NULL\n \"\"\"\n )\n ).scalar_one()\n\n if orphaned_count > 0:\n logger.warning(\n f\"WARNING: {orphaned_count} chat_session records still have folder_id without project_id. Proceeding anyway.\"\n )\n\n # === Step 2: Drop chat_session.folder_id ===\n if \"folder_id\" in chat_session_columns:\n logger.info(\"Dropping chat_session.folder_id...\")\n\n # Drop foreign key constraint first\n op.execute(\n \"ALTER TABLE chat_session DROP CONSTRAINT IF EXISTS chat_session_chat_folder_fk\"\n )\n op.execute(\n \"ALTER TABLE chat_session DROP CONSTRAINT IF EXISTS chat_session_folder_fk\"\n )\n\n # Drop the column\n op.drop_column(\"chat_session\", \"folder_id\")\n logger.info(\"Dropped chat_session.folder_id\")\n\n # === Step 3: Drop persona__user_folder table ===\n if \"persona__user_folder\" in inspector.get_table_names():\n logger.info(\"Dropping persona__user_folder table...\")\n\n # Check for any remaining data\n remaining = bind.execute(\n text(\"SELECT COUNT(*) FROM persona__user_folder\")\n ).scalar_one()\n\n if remaining > 0:\n logger.warning(\n f\"WARNING: Dropping persona__user_folder with {remaining} records\"\n )\n\n op.drop_table(\"persona__user_folder\")\n logger.info(\"Dropped persona__user_folder table\")\n\n # === Step 4: Drop chat_folder table ===\n if \"chat_folder\" in inspector.get_table_names():\n logger.info(\"Dropping chat_folder table...\")\n\n # Check for any remaining data\n remaining = bind.execute(text(\"SELECT COUNT(*) FROM chat_folder\")).scalar_one()\n\n if remaining > 0:\n logger.warning(f\"WARNING: Dropping chat_folder with {remaining} records\")\n\n op.drop_table(\"chat_folder\")\n logger.info(\"Dropped chat_folder table\")\n\n # === Step 5: Drop user_file legacy columns ===\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n\n # Drop folder_id\n if \"folder_id\" in user_file_columns:\n logger.info(\"Dropping user_file.folder_id...\")\n op.drop_column(\"user_file\", \"folder_id\")\n logger.info(\"Dropped user_file.folder_id\")\n\n # Drop cc_pair_id (already handled in migration 5, but be sure)\n if \"cc_pair_id\" in user_file_columns:\n logger.info(\"Dropping user_file.cc_pair_id...\")\n\n # Drop any remaining foreign key constraints\n bind.execute(\n text(\n \"\"\"\n DO $$\n DECLARE r RECORD;\n BEGIN\n FOR r IN (\n SELECT conname\n FROM pg_constraint c\n JOIN pg_class t ON c.conrelid = t.oid\n WHERE c.contype = 'f'\n AND t.relname = 'user_file'\n AND EXISTS (\n SELECT 1 FROM pg_attribute a\n WHERE a.attrelid = t.oid\n AND a.attname = 'cc_pair_id'\n )\n ) LOOP\n EXECUTE format('ALTER TABLE user_file DROP CONSTRAINT IF EXISTS %I', r.conname);\n END LOOP;\n END$$;\n \"\"\"\n )\n )\n\n op.drop_column(\"user_file\", \"cc_pair_id\")\n logger.info(\"Dropped user_file.cc_pair_id\")\n\n # === Step 6: Clean up any remaining constraints ===\n logger.info(\"Cleaning up remaining constraints...\")\n\n # Drop any unique constraints on removed columns\n op.execute(\n \"ALTER TABLE user_file DROP CONSTRAINT IF EXISTS user_file_cc_pair_id_key\"\n )\n\n logger.info(\"Migration 6 (schema cleanup) completed successfully\")\n logger.info(\"Legacy schema has been fully removed\")\n\n\ndef downgrade() -> None:\n \"\"\"Recreate dropped columns and tables (structure only, no data).\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n logger.warning(\"Downgrading schema cleanup - recreating structure only, no data!\")\n\n # Recreate user_file columns\n if \"user_file\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n\n if \"cc_pair_id\" not in columns:\n op.add_column(\n \"user_file\", sa.Column(\"cc_pair_id\", sa.Integer(), nullable=True)\n )\n\n if \"folder_id\" not in columns:\n op.add_column(\n \"user_file\", sa.Column(\"folder_id\", sa.Integer(), nullable=True)\n )\n\n # Recreate persona__user_folder table\n if \"persona__user_folder\" not in inspector.get_table_names():\n op.create_table(\n \"persona__user_folder\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_folder_id\", sa.Integer(), nullable=False),\n sa.PrimaryKeyConstraint(\"persona_id\", \"user_folder_id\"),\n sa.ForeignKeyConstraint([\"persona_id\"], [\"persona.id\"]),\n sa.ForeignKeyConstraint([\"user_folder_id\"], [\"user_project.id\"]),\n )\n\n # Recreate chat_folder table and related structures\n if \"chat_folder\" not in inspector.get_table_names():\n op.create_table(\n \"chat_folder\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"name\", sa.String(), nullable=True),\n sa.Column(\"display_priority\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n name=\"chat_folder_user_id_fkey\",\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Add folder_id back to chat_session\n if \"chat_session\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"chat_session\")]\n if \"folder_id\" not in columns:\n op.add_column(\n \"chat_session\", sa.Column(\"folder_id\", sa.Integer(), nullable=True)\n )\n\n # Add foreign key if chat_folder exists\n if \"chat_folder\" in inspector.get_table_names():\n op.create_foreign_key(\n \"chat_session_chat_folder_fk\",\n \"chat_session\",\n \"chat_folder\",\n [\"folder_id\"],\n [\"id\"],\n )\n\n logger.info(\"Downgrade completed - structure recreated but data is lost\")\n" }, { "path": "backend/alembic/versions/2b90f3af54b8_usage_limits.py", "content": "\"\"\"usage_limits\n\nRevision ID: 2b90f3af54b8\nRevises: 9a0296d7421e\nCreate Date: 2026-01-03 16:55:30.449692\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2b90f3af54b8\"\ndown_revision = \"9a0296d7421e\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"tenant_usage\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"window_start\", sa.DateTime(timezone=True), nullable=False, index=True\n ),\n sa.Column(\"llm_cost_cents\", sa.Float(), nullable=False, server_default=\"0.0\"),\n sa.Column(\"chunks_indexed\", sa.Integer(), nullable=False, server_default=\"0\"),\n sa.Column(\"api_calls\", sa.Integer(), nullable=False, server_default=\"0\"),\n sa.Column(\n \"non_streaming_api_calls\", sa.Integer(), nullable=False, server_default=\"0\"\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=True,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"window_start\", name=\"uq_tenant_usage_window\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_tenant_usage_window_start\", table_name=\"tenant_usage\")\n op.drop_table(\"tenant_usage\")\n" }, { "path": "backend/alembic/versions/2c2430828bdf_add_unique_constraint_to_inputprompt_.py", "content": "\"\"\"add_unique_constraint_to_inputprompt_prompt_user_id\n\nRevision ID: 2c2430828bdf\nRevises: fb80bdd256de\nCreate Date: 2026-01-20 16:01:54.314805\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2c2430828bdf\"\ndown_revision = \"fb80bdd256de\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create unique constraint on (prompt, user_id) for user-owned prompts\n # This ensures each user can only have one shortcut with a given name\n op.create_unique_constraint(\n \"uq_inputprompt_prompt_user_id\",\n \"inputprompt\",\n [\"prompt\", \"user_id\"],\n )\n\n # Create partial unique index for public prompts (where user_id IS NULL)\n # PostgreSQL unique constraints don't enforce uniqueness for NULL values,\n # so we need a partial index to ensure public prompt names are also unique\n op.execute(\n \"\"\"\n CREATE UNIQUE INDEX uq_inputprompt_prompt_public\n ON inputprompt (prompt)\n WHERE user_id IS NULL\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\"DROP INDEX IF EXISTS uq_inputprompt_prompt_public\")\n op.drop_constraint(\"uq_inputprompt_prompt_user_id\", \"inputprompt\", type_=\"unique\")\n" }, { "path": "backend/alembic/versions/2cdeff6d8c93_set_built_in_to_default.py", "content": "\"\"\"set built in to default\n\nRevision ID: 2cdeff6d8c93\nRevises: f5437cc136c5\nCreate Date: 2025-02-11 14:57:51.308775\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2cdeff6d8c93\"\ndown_revision = \"f5437cc136c5\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Prior to this migration / point in the codebase history,\n # built in personas were implicitly treated as default personas (with no option to change this)\n # This migration makes that explicit\n op.execute(\n \"\"\"\n UPDATE persona\n SET is_default_persona = TRUE\n WHERE builtin_persona = TRUE\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/2d2304e27d8c_add_above_below_to_persona.py", "content": "\"\"\"Add Above Below to Persona\n\nRevision ID: 2d2304e27d8c\nRevises: 4b08d97e175a\nCreate Date: 2024-08-21 19:15:15.762948\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"2d2304e27d8c\"\ndown_revision = \"4b08d97e175a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"persona\", sa.Column(\"chunks_above\", sa.Integer(), nullable=True))\n op.add_column(\"persona\", sa.Column(\"chunks_below\", sa.Integer(), nullable=True))\n\n op.execute(\n \"UPDATE persona SET chunks_above = 1, chunks_below = 1 WHERE chunks_above IS NULL AND chunks_below IS NULL\"\n )\n\n op.alter_column(\"persona\", \"chunks_above\", nullable=False)\n op.alter_column(\"persona\", \"chunks_below\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"chunks_below\")\n op.drop_column(\"persona\", \"chunks_above\")\n" }, { "path": "backend/alembic/versions/2daa494a0851_add_group_sync_time.py", "content": "\"\"\"add-group-sync-time\n\nRevision ID: 2daa494a0851\nRevises: c0fd6e4da83a\nCreate Date: 2024-11-11 10:57:22.991157\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"2daa494a0851\"\ndown_revision = \"c0fd6e4da83a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"last_time_external_group_sync\",\n sa.DateTime(timezone=True),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"last_time_external_group_sync\")\n" }, { "path": "backend/alembic/versions/2f80c6a2550f_add_chat_session_specific_temperature_.py", "content": "\"\"\"add chat session specific temperature override\n\nRevision ID: 2f80c6a2550f\nRevises: 33ea50e88f24\nCreate Date: 2025-01-31 10:30:27.289646\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2f80c6a2550f\"\ndown_revision = \"33ea50e88f24\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_session\", sa.Column(\"temperature_override\", sa.Float(), nullable=True)\n )\n op.add_column(\n \"user\",\n sa.Column(\n \"temperature_override_enabled\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_session\", \"temperature_override\")\n op.drop_column(\"user\", \"temperature_override_enabled\")\n" }, { "path": "backend/alembic/versions/2f95e36923e6_add_indexing_coordination.py", "content": "\"\"\"add_indexing_coordination\n\nRevision ID: 2f95e36923e6\nRevises: 0816326d83aa\nCreate Date: 2025-07-10 16:17:57.762182\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"2f95e36923e6\"\ndown_revision = \"0816326d83aa\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add database-based coordination fields (replacing Redis fencing)\n op.add_column(\n \"index_attempt\", sa.Column(\"celery_task_id\", sa.String(), nullable=True)\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"cancellation_requested\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n\n # Add batch coordination fields (replacing FileStore state)\n op.add_column(\n \"index_attempt\", sa.Column(\"total_batches\", sa.Integer(), nullable=True)\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"completed_batches\", sa.Integer(), nullable=False, server_default=\"0\"\n ),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"total_failures_batch_level\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\"total_chunks\", sa.Integer(), nullable=False, server_default=\"0\"),\n )\n\n # Progress tracking for stall detection\n op.add_column(\n \"index_attempt\",\n sa.Column(\"last_progress_time\", sa.DateTime(timezone=True), nullable=True),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"last_batches_completed_count\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n )\n\n # Heartbeat tracking for worker liveness detection\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"heartbeat_counter\", sa.Integer(), nullable=False, server_default=\"0\"\n ),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"last_heartbeat_value\", sa.Integer(), nullable=False, server_default=\"0\"\n ),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\"last_heartbeat_time\", sa.DateTime(timezone=True), nullable=True),\n )\n\n # Add index for coordination queries\n op.create_index(\n \"ix_index_attempt_active_coordination\",\n \"index_attempt\",\n [\"connector_credential_pair_id\", \"search_settings_id\", \"status\"],\n )\n\n\ndef downgrade() -> None:\n # Remove the new index\n op.drop_index(\"ix_index_attempt_active_coordination\", table_name=\"index_attempt\")\n\n # Remove the new columns\n op.drop_column(\"index_attempt\", \"last_batches_completed_count\")\n op.drop_column(\"index_attempt\", \"last_progress_time\")\n op.drop_column(\"index_attempt\", \"last_heartbeat_time\")\n op.drop_column(\"index_attempt\", \"last_heartbeat_value\")\n op.drop_column(\"index_attempt\", \"heartbeat_counter\")\n op.drop_column(\"index_attempt\", \"total_chunks\")\n op.drop_column(\"index_attempt\", \"total_failures_batch_level\")\n op.drop_column(\"index_attempt\", \"completed_batches\")\n op.drop_column(\"index_attempt\", \"total_batches\")\n op.drop_column(\"index_attempt\", \"cancellation_requested\")\n op.drop_column(\"index_attempt\", \"celery_task_id\")\n" }, { "path": "backend/alembic/versions/30c1d5744104_persona_datetime_aware.py", "content": "\"\"\"Persona Datetime Aware\n\nRevision ID: 30c1d5744104\nRevises: 7f99be1cb9f5\nCreate Date: 2023-10-16 23:21:01.283424\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"30c1d5744104\"\ndown_revision = \"7f99be1cb9f5\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"persona\", sa.Column(\"datetime_aware\", sa.Boolean(), nullable=True))\n op.execute(\"UPDATE persona SET datetime_aware = TRUE\")\n op.alter_column(\"persona\", \"datetime_aware\", nullable=False)\n op.create_index(\n \"_default_persona_name_idx\",\n \"persona\",\n [\"name\"],\n unique=True,\n postgresql_where=sa.text(\"default_persona = true\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"_default_persona_name_idx\",\n table_name=\"persona\",\n postgresql_where=sa.text(\"default_persona = true\"),\n )\n op.drop_column(\"persona\", \"datetime_aware\")\n" }, { "path": "backend/alembic/versions/325975216eb3_add_icon_color_and_icon_shape_to_persona.py", "content": "\"\"\"Add icon_color and icon_shape to Persona\n\nRevision ID: 325975216eb3\nRevises: 91ffac7e65b3\nCreate Date: 2024-07-24 21:29:31.784562\n\n\"\"\"\n\nimport random\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.sql import table, column, select\n\n# revision identifiers, used by Alembic.\nrevision = \"325975216eb3\"\ndown_revision = \"91ffac7e65b3\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ncolorOptions = [\n \"#FF6FBF\",\n \"#6FB1FF\",\n \"#B76FFF\",\n \"#FFB56F\",\n \"#6FFF8D\",\n \"#FF6F6F\",\n \"#6FFFFF\",\n]\n\n\n# Function to generate a random shape ensuring at least 3 of the middle 4 squares are filled\ndef generate_random_shape() -> int:\n center_squares = [12, 10, 6, 14, 13, 11, 7, 15]\n center_fill = random.choice(center_squares)\n remaining_squares = [i for i in range(16) if not (center_fill & (1 << i))]\n random.shuffle(remaining_squares)\n for i in range(10 - bin(center_fill).count(\"1\")):\n center_fill |= 1 << remaining_squares[i]\n return center_fill\n\n\ndef upgrade() -> None:\n op.add_column(\"persona\", sa.Column(\"icon_color\", sa.String(), nullable=True))\n op.add_column(\"persona\", sa.Column(\"icon_shape\", sa.Integer(), nullable=True))\n op.add_column(\"persona\", sa.Column(\"uploaded_image_id\", sa.String(), nullable=True))\n\n persona = table(\n \"persona\",\n column(\"id\", sa.Integer),\n column(\"icon_color\", sa.String),\n column(\"icon_shape\", sa.Integer),\n )\n\n conn = op.get_bind()\n personas = conn.execute(select(persona.c.id))\n\n for persona_id in personas:\n random_color = random.choice(colorOptions)\n random_shape = generate_random_shape()\n conn.execute(\n persona.update()\n .where(persona.c.id == persona_id[0])\n .values(icon_color=random_color, icon_shape=random_shape)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"icon_shape\")\n op.drop_column(\"persona\", \"uploaded_image_id\")\n op.drop_column(\"persona\", \"icon_color\")\n" }, { "path": "backend/alembic/versions/33cb72ea4d80_single_tool_call_per_message.py", "content": "\"\"\"single tool call per message\n\nRevision ID: 33cb72ea4d80\nRevises: 5b29123cd710\nCreate Date: 2024-11-01 12:51:01.535003\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"33cb72ea4d80\"\ndown_revision = \"5b29123cd710\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Step 1: Delete extraneous ToolCall entries\n # Keep only the ToolCall with the smallest 'id' for each 'message_id'\n op.execute(\n sa.text(\n \"\"\"\n DELETE FROM tool_call\n WHERE id NOT IN (\n SELECT MIN(id)\n FROM tool_call\n WHERE message_id IS NOT NULL\n GROUP BY message_id\n );\n \"\"\"\n )\n )\n\n # Step 2: Add a unique constraint on message_id\n op.create_unique_constraint(\n constraint_name=\"uq_tool_call_message_id\",\n table_name=\"tool_call\",\n columns=[\"message_id\"],\n )\n\n\ndef downgrade() -> None:\n # Step 1: Drop the unique constraint on message_id\n op.drop_constraint(\n constraint_name=\"uq_tool_call_message_id\",\n table_name=\"tool_call\",\n type_=\"unique\",\n )\n" }, { "path": "backend/alembic/versions/33ea50e88f24_foreign_key_input_prompts.py", "content": "\"\"\"foreign key input prompts\n\nRevision ID: 33ea50e88f24\nRevises: a6df6b88ef81\nCreate Date: 2025-01-29 10:54:22.141765\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"33ea50e88f24\"\ndown_revision = \"a6df6b88ef81\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Safely drop constraints if exists\n op.execute(\n \"\"\"\n ALTER TABLE inputprompt__user\n DROP CONSTRAINT IF EXISTS inputprompt__user_input_prompt_id_fkey\n \"\"\"\n )\n op.execute(\n \"\"\"\n ALTER TABLE inputprompt__user\n DROP CONSTRAINT IF EXISTS inputprompt__user_user_id_fkey\n \"\"\"\n )\n\n # Recreate with ON DELETE CASCADE\n op.create_foreign_key(\n \"inputprompt__user_input_prompt_id_fkey\",\n \"inputprompt__user\",\n \"inputprompt\",\n [\"input_prompt_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n op.create_foreign_key(\n \"inputprompt__user_user_id_fkey\",\n \"inputprompt__user\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n # Drop the new FKs with ondelete\n op.drop_constraint(\n \"inputprompt__user_input_prompt_id_fkey\",\n \"inputprompt__user\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"inputprompt__user_user_id_fkey\",\n \"inputprompt__user\",\n type_=\"foreignkey\",\n )\n\n # Recreate them without cascading\n op.create_foreign_key(\n \"inputprompt__user_input_prompt_id_fkey\",\n \"inputprompt__user\",\n \"inputprompt\",\n [\"input_prompt_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"inputprompt__user_user_id_fkey\",\n \"inputprompt__user\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/351faebd379d_add_curator_fields.py", "content": "\"\"\"Add curator fields\n\nRevision ID: 351faebd379d\nRevises: ee3f4b47fad5\nCreate Date: 2024-08-15 22:37:08.397052\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"351faebd379d\"\ndown_revision = \"ee3f4b47fad5\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Add is_curator column to User__UserGroup table\n op.add_column(\n \"user__user_group\",\n sa.Column(\"is_curator\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n\n # Use batch mode to modify the enum type\n with op.batch_alter_table(\"user\", schema=None) as batch_op:\n batch_op.alter_column( # type: ignore[attr-defined]\n \"role\",\n type_=sa.Enum(\n \"BASIC\",\n \"ADMIN\",\n \"CURATOR\",\n \"GLOBAL_CURATOR\",\n name=\"userrole\",\n native_enum=False,\n ),\n existing_type=sa.Enum(\"BASIC\", \"ADMIN\", name=\"userrole\", native_enum=False),\n existing_nullable=False,\n )\n # Create the association table\n op.create_table(\n \"credential__user_group\",\n sa.Column(\"credential_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_group_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"credential_id\"],\n [\"credential.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.PrimaryKeyConstraint(\"credential_id\", \"user_group_id\"),\n )\n op.add_column(\n \"credential\",\n sa.Column(\n \"curator_public\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n )\n\n\ndef downgrade() -> None:\n # Update existing records to ensure they fit within the BASIC/ADMIN roles\n op.execute(\n \"UPDATE \\\"user\\\" SET role = 'ADMIN' WHERE role IN ('CURATOR', 'GLOBAL_CURATOR')\"\n )\n\n # Remove is_curator column from User__UserGroup table\n op.drop_column(\"user__user_group\", \"is_curator\")\n\n with op.batch_alter_table(\"user\", schema=None) as batch_op:\n batch_op.alter_column( # type: ignore[attr-defined]\n \"role\",\n type_=sa.Enum(\n \"BASIC\", \"ADMIN\", name=\"userrole\", native_enum=False, length=20\n ),\n existing_type=sa.Enum(\n \"BASIC\",\n \"ADMIN\",\n \"CURATOR\",\n \"GLOBAL_CURATOR\",\n name=\"userrole\",\n native_enum=False,\n ),\n existing_nullable=False,\n )\n # Drop the association table\n op.drop_table(\"credential__user_group\")\n op.drop_column(\"credential\", \"curator_public\")\n" }, { "path": "backend/alembic/versions/35e518e0ddf4_properly_cascade.py", "content": "\"\"\"properly_cascade\n\nRevision ID: 35e518e0ddf4\nRevises: 91a0a4d62b14\nCreate Date: 2024-09-20 21:24:04.891018\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"35e518e0ddf4\"\ndown_revision = \"91a0a4d62b14\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Update chat_message foreign key constraint\n op.drop_constraint(\n \"chat_message_chat_session_id_fkey\", \"chat_message\", type_=\"foreignkey\"\n )\n op.create_foreign_key(\n \"chat_message_chat_session_id_fkey\",\n \"chat_message\",\n \"chat_session\",\n [\"chat_session_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # Update chat_message__search_doc foreign key constraints\n op.drop_constraint(\n \"chat_message__search_doc_chat_message_id_fkey\",\n \"chat_message__search_doc\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"chat_message__search_doc_search_doc_id_fkey\",\n \"chat_message__search_doc\",\n type_=\"foreignkey\",\n )\n\n op.create_foreign_key(\n \"chat_message__search_doc_chat_message_id_fkey\",\n \"chat_message__search_doc\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n op.create_foreign_key(\n \"chat_message__search_doc_search_doc_id_fkey\",\n \"chat_message__search_doc\",\n \"search_doc\",\n [\"search_doc_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # Add CASCADE delete for tool_call foreign key\n op.drop_constraint(\"tool_call_message_id_fkey\", \"tool_call\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"tool_call_message_id_fkey\",\n \"tool_call\",\n \"chat_message\",\n [\"message_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n # Revert chat_message foreign key constraint\n op.drop_constraint(\n \"chat_message_chat_session_id_fkey\", \"chat_message\", type_=\"foreignkey\"\n )\n op.create_foreign_key(\n \"chat_message_chat_session_id_fkey\",\n \"chat_message\",\n \"chat_session\",\n [\"chat_session_id\"],\n [\"id\"],\n )\n\n # Revert chat_message__search_doc foreign key constraints\n op.drop_constraint(\n \"chat_message__search_doc_chat_message_id_fkey\",\n \"chat_message__search_doc\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"chat_message__search_doc_search_doc_id_fkey\",\n \"chat_message__search_doc\",\n type_=\"foreignkey\",\n )\n\n op.create_foreign_key(\n \"chat_message__search_doc_chat_message_id_fkey\",\n \"chat_message__search_doc\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"chat_message__search_doc_search_doc_id_fkey\",\n \"chat_message__search_doc\",\n \"search_doc\",\n [\"search_doc_id\"],\n [\"id\"],\n )\n\n # Revert tool_call foreign key constraint\n op.drop_constraint(\"tool_call_message_id_fkey\", \"tool_call\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"tool_call_message_id_fkey\",\n \"tool_call\",\n \"chat_message\",\n [\"message_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/35e6853a51d5_server_default_chosen_assistants.py", "content": "\"\"\"server default chosen assistants\n\nRevision ID: 35e6853a51d5\nRevises: c99d76fcd298\nCreate Date: 2024-09-13 13:20:32.885317\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"35e6853a51d5\"\ndown_revision = \"c99d76fcd298\"\nbranch_labels = None\ndepends_on = None\n\nDEFAULT_ASSISTANTS = [-2, -1, 0]\n\n\ndef upgrade() -> None:\n # Step 1: Update any NULL values to the default value\n # This upgrades existing users without ordered assistant\n # to have default assistants set to visible assistants which are\n # accessible by them.\n op.execute(\n \"\"\"\n UPDATE \"user\" u\n SET chosen_assistants = (\n SELECT jsonb_agg(\n p.id ORDER BY\n COALESCE(p.display_priority, 2147483647) ASC,\n p.id ASC\n )\n FROM persona p\n LEFT JOIN persona__user pu ON p.id = pu.persona_id AND pu.user_id = u.id\n WHERE p.is_visible = true\n AND (p.is_public = true OR pu.user_id IS NOT NULL)\n )\n WHERE chosen_assistants IS NULL\n OR chosen_assistants = 'null'\n OR jsonb_typeof(chosen_assistants) = 'null'\n OR (jsonb_typeof(chosen_assistants) = 'string' AND chosen_assistants = '\"null\"')\n \"\"\"\n )\n\n # Step 2: Alter the column to make it non-nullable\n op.alter_column(\n \"user\",\n \"chosen_assistants\",\n type_=postgresql.JSONB(astext_type=sa.Text()),\n nullable=False,\n server_default=sa.text(f\"'{DEFAULT_ASSISTANTS}'::jsonb\"),\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"user\",\n \"chosen_assistants\",\n type_=postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n server_default=None,\n )\n" }, { "path": "backend/alembic/versions/369644546676_add_composite_index_for_index_attempt_.py", "content": "\"\"\"add composite index for index attempt time updated\n\nRevision ID: 369644546676\nRevises: 2955778aa44c\nCreate Date: 2025-01-08 15:38:17.224380\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy import text\n\n# revision identifiers, used by Alembic.\nrevision = \"369644546676\"\ndown_revision = \"2955778aa44c\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_index(\n \"ix_index_attempt_ccpair_search_settings_time_updated\",\n \"index_attempt\",\n [\n \"connector_credential_pair_id\",\n \"search_settings_id\",\n text(\"time_updated DESC\"),\n ],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"ix_index_attempt_ccpair_search_settings_time_updated\",\n table_name=\"index_attempt\",\n )\n" }, { "path": "backend/alembic/versions/36e9220ab794_update_kg_trigger_functions.py", "content": "\"\"\"update_kg_trigger_functions\n\nRevision ID: 36e9220ab794\nRevises: c9e2cd766c29\nCreate Date: 2025-06-22 17:33:25.833733\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy import text\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\n# revision identifiers, used by Alembic.\nrevision = \"36e9220ab794\"\ndown_revision = \"c9e2cd766c29\"\nbranch_labels = None\ndepends_on = None\n\n\ndef _get_tenant_contextvar(session: Session) -> str:\n \"\"\"Get the current schema for the migration\"\"\"\n current_tenant = session.execute(text(\"SELECT current_schema()\")).scalar()\n if isinstance(current_tenant, str):\n return current_tenant\n else:\n raise ValueError(\"Current tenant is not a string\")\n\n\ndef upgrade() -> None:\n\n bind = op.get_bind()\n session = Session(bind=bind)\n\n # Create kg_entity trigger to update kg_entity.name and its trigrams\n tenant_id = _get_tenant_contextvar(session)\n alphanum_pattern = r\"[^a-z0-9]+\"\n truncate_length = 1000\n function = \"update_kg_entity_name\"\n op.execute(\n text(\n f\"\"\"\n CREATE OR REPLACE FUNCTION \"{tenant_id}\".{function}()\n RETURNS TRIGGER AS $$\n DECLARE\n name text;\n cleaned_name text;\n BEGIN\n -- Set name to semantic_id if document_id is not NULL\n IF NEW.document_id IS NOT NULL THEN\n SELECT lower(semantic_id) INTO name\n FROM \"{tenant_id}\".document\n WHERE id = NEW.document_id;\n ELSE\n name = lower(NEW.name);\n END IF;\n\n -- Clean name and truncate if too long\n cleaned_name = regexp_replace(\n name,\n '{alphanum_pattern}', '', 'g'\n );\n IF length(cleaned_name) > {truncate_length} THEN\n cleaned_name = left(cleaned_name, {truncate_length});\n END IF;\n\n -- Set name and name trigrams\n NEW.name = name;\n NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name);\n RETURN NEW;\n END;\n $$ LANGUAGE plpgsql;\n \"\"\"\n )\n )\n trigger = f\"{function}_trigger\"\n op.execute(f'DROP TRIGGER IF EXISTS {trigger} ON \"{tenant_id}\".kg_entity')\n op.execute(\n f\"\"\"\n CREATE TRIGGER {trigger}\n BEFORE INSERT OR UPDATE OF name\n ON \"{tenant_id}\".kg_entity\n FOR EACH ROW\n EXECUTE FUNCTION \"{tenant_id}\".{function}();\n \"\"\"\n )\n\n # Create kg_entity trigger to update kg_entity.name and its trigrams\n function = \"update_kg_entity_name_from_doc\"\n op.execute(\n text(\n f\"\"\"\n CREATE OR REPLACE FUNCTION \"{tenant_id}\".{function}()\n RETURNS TRIGGER AS $$\n DECLARE\n doc_name text;\n cleaned_name text;\n BEGIN\n doc_name = lower(NEW.semantic_id);\n\n -- Clean name and truncate if too long\n cleaned_name = regexp_replace(\n doc_name,\n '{alphanum_pattern}', '', 'g'\n );\n IF length(cleaned_name) > {truncate_length} THEN\n cleaned_name = left(cleaned_name, {truncate_length});\n END IF;\n\n -- Set name and name trigrams for all entities referencing this document\n UPDATE \"{tenant_id}\".kg_entity\n SET\n name = doc_name,\n name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name)\n WHERE document_id = NEW.id;\n RETURN NEW;\n END;\n $$ LANGUAGE plpgsql;\n \"\"\"\n )\n )\n trigger = f\"{function}_trigger\"\n op.execute(f'DROP TRIGGER IF EXISTS {trigger} ON \"{tenant_id}\".document')\n op.execute(\n f\"\"\"\n CREATE TRIGGER {trigger}\n AFTER UPDATE OF semantic_id\n ON \"{tenant_id}\".document\n FOR EACH ROW\n EXECUTE FUNCTION \"{tenant_id}\".{function}();\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/3781a5eb12cb_add_chunk_stats_table.py", "content": "\"\"\"add chunk stats table\n\nRevision ID: 3781a5eb12cb\nRevises: df46c75b714e\nCreate Date: 2025-03-10 10:02:30.586666\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"3781a5eb12cb\"\ndown_revision = \"df46c75b714e\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"chunk_stats\",\n sa.Column(\"id\", sa.String(), primary_key=True, index=True),\n sa.Column(\n \"document_id\",\n sa.String(),\n sa.ForeignKey(\"document.id\"),\n nullable=False,\n index=True,\n ),\n sa.Column(\"chunk_in_doc_id\", sa.Integer(), nullable=False),\n sa.Column(\"information_content_boost\", sa.Float(), nullable=True),\n sa.Column(\n \"last_modified\",\n sa.DateTime(timezone=True),\n nullable=False,\n index=True,\n server_default=sa.func.now(),\n ),\n sa.Column(\"last_synced\", sa.DateTime(timezone=True), nullable=True, index=True),\n sa.UniqueConstraint(\n \"document_id\", \"chunk_in_doc_id\", name=\"uq_chunk_stats_doc_chunk\"\n ),\n )\n\n op.create_index(\n \"ix_chunk_sync_status\", \"chunk_stats\", [\"last_modified\", \"last_synced\"]\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_chunk_sync_status\", table_name=\"chunk_stats\")\n op.drop_table(\"chunk_stats\")\n" }, { "path": "backend/alembic/versions/3879338f8ba1_add_tool_table.py", "content": "\"\"\"Add tool table\n\nRevision ID: 3879338f8ba1\nRevises: f1c6478c3fd8\nCreate Date: 2024-05-11 16:11:23.718084\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"3879338f8ba1\"\ndown_revision = \"f1c6478c3fd8\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"tool\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"description\", sa.Text(), nullable=True),\n sa.Column(\"in_code_tool_id\", sa.String(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"persona__tool\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"tool_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"tool_id\"],\n [\"tool.id\"],\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"tool_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"persona__tool\")\n op.drop_table(\"tool\")\n" }, { "path": "backend/alembic/versions/38eda64af7fe_add_chat_session_sharing.py", "content": "\"\"\"Add chat session sharing\n\nRevision ID: 38eda64af7fe\nRevises: 776b3bbe9092\nCreate Date: 2024-03-27 19:41:29.073594\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"38eda64af7fe\"\ndown_revision = \"776b3bbe9092\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_session\",\n sa.Column(\n \"shared_status\",\n sa.Enum(\n \"PUBLIC\",\n \"PRIVATE\",\n name=\"chatsessionsharedstatus\",\n native_enum=False,\n ),\n nullable=True,\n ),\n )\n op.execute(\"UPDATE chat_session SET shared_status='PRIVATE'\")\n op.alter_column(\n \"chat_session\",\n \"shared_status\",\n nullable=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_session\", \"shared_status\")\n" }, { "path": "backend/alembic/versions/3934b1bc7b62_update_github_connector_repo_name_to_.py", "content": "\"\"\"Update GitHub connector repo_name to repositories\n\nRevision ID: 3934b1bc7b62\nRevises: b7c2b63c4a03\nCreate Date: 2025-03-05 10:50:30.516962\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport json\nimport logging\n\n# revision identifiers, used by Alembic.\nrevision = \"3934b1bc7b62\"\ndown_revision = \"b7c2b63c4a03\"\nbranch_labels = None\ndepends_on = None\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n\ndef upgrade() -> None:\n # Get all GitHub connectors\n conn = op.get_bind()\n\n # First get all GitHub connectors\n github_connectors = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'GITHUB'\n \"\"\"\n )\n ).fetchall()\n\n # Update each connector's config\n updated_count = 0\n for connector_id, config in github_connectors:\n try:\n if not config:\n logger.warning(f\"Connector {connector_id} has no config, skipping\")\n continue\n\n # Parse the config if it's a string\n if isinstance(config, str):\n config = json.loads(config)\n\n if \"repo_name\" not in config:\n continue\n\n # Create new config with repositories instead of repo_name\n new_config = dict(config)\n repo_name_value = new_config.pop(\"repo_name\")\n new_config[\"repositories\"] = repo_name_value\n\n # Update the connector with the new config\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = :new_config\n WHERE id = :connector_id\n \"\"\"\n ),\n {\"connector_id\": connector_id, \"new_config\": json.dumps(new_config)},\n )\n updated_count += 1\n except Exception as e:\n logger.error(f\"Error updating connector {connector_id}: {str(e)}\")\n\n\ndef downgrade() -> None:\n # Get all GitHub connectors\n conn = op.get_bind()\n\n logger.debug(\n \"Starting rollback of GitHub connectors from repositories to repo_name\"\n )\n\n github_connectors = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'GITHUB'\n \"\"\"\n )\n ).fetchall()\n\n logger.debug(f\"Found {len(github_connectors)} GitHub connectors to rollback\")\n\n # Revert each GitHub connector to use repo_name instead of repositories\n reverted_count = 0\n for connector_id, config in github_connectors:\n try:\n if not config:\n continue\n\n # Parse the config if it's a string\n if isinstance(config, str):\n config = json.loads(config)\n\n if \"repositories\" not in config:\n continue\n\n # Create new config with repo_name instead of repositories\n new_config = dict(config)\n repositories_value = new_config.pop(\"repositories\")\n new_config[\"repo_name\"] = repositories_value\n\n # Update the connector with the new config\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = :new_config\n WHERE id = :connector_id\n \"\"\"\n ),\n {\"new_config\": json.dumps(new_config), \"connector_id\": connector_id},\n )\n reverted_count += 1\n except Exception as e:\n logger.error(f\"Error reverting connector {connector_id}: {str(e)}\")\n" }, { "path": "backend/alembic/versions/3a7802814195_add_alternate_assistant_to_chat_message.py", "content": "\"\"\"add alternate assistant to chat message\n\nRevision ID: 3a7802814195\nRevises: 23957775e5f5\nCreate Date: 2024-06-05 11:18:49.966333\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"3a7802814195\"\ndown_revision = \"23957775e5f5\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\", sa.Column(\"alternate_assistant_id\", sa.Integer(), nullable=True)\n )\n op.create_foreign_key(\n \"fk_chat_message_persona\",\n \"chat_message\",\n \"persona\",\n [\"alternate_assistant_id\"],\n [\"id\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\"fk_chat_message_persona\", \"chat_message\", type_=\"foreignkey\")\n op.drop_column(\"chat_message\", \"alternate_assistant_id\")\n" }, { "path": "backend/alembic/versions/3a78dba1080a_user_file_legacy_data_cleanup.py", "content": "\"\"\"Migration 5: User file legacy data cleanup\n\nRevision ID: 3a78dba1080a\nRevises: 7cc3fcc116c1\nCreate Date: 2025-09-22 10:04:27.986294\n\nThis migration removes legacy user-file documents and connector_credential_pairs.\nIt performs bulk deletions of obsolete data after the UUID migration.\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql as psql\nfrom sqlalchemy import text\nimport logging\nfrom typing import List\nimport uuid\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n# revision identifiers, used by Alembic.\nrevision = \"3a78dba1080a\"\ndown_revision = \"7cc3fcc116c1\"\nbranch_labels = None\ndepends_on = None\n\n\ndef batch_delete(\n bind: sa.engine.Connection,\n table_name: str,\n id_column: str,\n ids: List[str | int | uuid.UUID],\n batch_size: int = 1000,\n id_type: str = \"int\",\n) -> int:\n \"\"\"Delete records in batches to avoid memory issues and timeouts.\"\"\"\n total_count = len(ids)\n if total_count == 0:\n return 0\n\n logger.info(\n f\"Starting batch deletion of {total_count} records from {table_name}...\"\n )\n\n # Determine appropriate ARRAY type\n if id_type == \"uuid\":\n array_type = psql.ARRAY(psql.UUID(as_uuid=True))\n elif id_type == \"int\":\n array_type = psql.ARRAY(sa.Integer())\n else:\n array_type = psql.ARRAY(sa.String())\n\n total_deleted = 0\n failed_batches = []\n\n for i in range(0, total_count, batch_size):\n batch_ids = ids[i : i + batch_size]\n try:\n stmt = text(\n f\"DELETE FROM {table_name} WHERE {id_column} = ANY(:ids)\"\n ).bindparams(sa.bindparam(\"ids\", value=batch_ids, type_=array_type))\n result = bind.execute(stmt)\n total_deleted += result.rowcount\n\n # Log progress every 10 batches or at completion\n batch_num = (i // batch_size) + 1\n if batch_num % 10 == 0 or i + batch_size >= total_count:\n logger.info(\n f\" Deleted {min(i + batch_size, total_count)}/{total_count} records \"\n f\"({total_deleted} actual) from {table_name}\"\n )\n except Exception as e:\n logger.error(f\"Failed to delete batch {(i // batch_size) + 1}: {e}\")\n failed_batches.append((i, min(i + batch_size, total_count)))\n\n if failed_batches:\n logger.warning(\n f\"Failed to delete {len(failed_batches)} batches from {table_name}. Total deleted: {total_deleted}/{total_count}\"\n )\n # Fail the migration to avoid silently succeeding on partial cleanup\n raise RuntimeError(\n f\"Batch deletion failed for {table_name}: \"\n f\"{len(failed_batches)} failed batches out of \"\n f\"{(total_count + batch_size - 1) // batch_size}.\"\n )\n\n return total_deleted\n\n\ndef upgrade() -> None:\n \"\"\"Remove legacy user-file documents and connector_credential_pairs.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n logger.info(\"Starting legacy data cleanup...\")\n\n # === Step 1: Identify and delete user-file documents ===\n logger.info(\"Identifying user-file documents to delete...\")\n\n # Get document IDs to delete\n doc_rows = bind.execute(\n text(\n \"\"\"\n SELECT DISTINCT dcc.id AS document_id\n FROM document_by_connector_credential_pair dcc\n JOIN connector_credential_pair u\n ON u.connector_id = dcc.connector_id\n AND u.credential_id = dcc.credential_id\n WHERE u.is_user_file IS TRUE\n \"\"\"\n )\n ).fetchall()\n\n doc_ids = [r[0] for r in doc_rows]\n\n if doc_ids:\n logger.info(f\"Found {len(doc_ids)} user-file documents to delete\")\n\n # Delete dependent rows first\n tables_to_clean = [\n (\"document_retrieval_feedback\", \"document_id\"),\n (\"document__tag\", \"document_id\"),\n (\"chunk_stats\", \"document_id\"),\n ]\n\n for table_name, column_name in tables_to_clean:\n if table_name in inspector.get_table_names():\n # document_id is a string in these tables\n deleted = batch_delete(\n bind, table_name, column_name, doc_ids, id_type=\"str\"\n )\n logger.info(f\"Deleted {deleted} records from {table_name}\")\n\n # Delete document_by_connector_credential_pair entries\n deleted = batch_delete(\n bind, \"document_by_connector_credential_pair\", \"id\", doc_ids, id_type=\"str\"\n )\n logger.info(f\"Deleted {deleted} document_by_connector_credential_pair records\")\n\n # Delete documents themselves\n deleted = batch_delete(bind, \"document\", \"id\", doc_ids, id_type=\"str\")\n logger.info(f\"Deleted {deleted} document records\")\n else:\n logger.info(\"No user-file documents found to delete\")\n\n # === Step 2: Clean up user-file connector_credential_pairs ===\n logger.info(\"Cleaning up user-file connector_credential_pairs...\")\n\n # Get cc_pair IDs\n cc_pair_rows = bind.execute(\n text(\n \"\"\"\n SELECT id AS cc_pair_id\n FROM connector_credential_pair\n WHERE is_user_file IS TRUE\n \"\"\"\n )\n ).fetchall()\n\n cc_pair_ids = [r[0] for r in cc_pair_rows]\n\n if cc_pair_ids:\n logger.info(\n f\"Found {len(cc_pair_ids)} user-file connector_credential_pairs to clean up\"\n )\n\n # Delete related records\n # Clean child tables first to satisfy foreign key constraints,\n # then the parent tables\n tables_to_clean = [\n (\"index_attempt_errors\", \"connector_credential_pair_id\"),\n (\"index_attempt\", \"connector_credential_pair_id\"),\n (\"background_error\", \"cc_pair_id\"),\n (\"document_set__connector_credential_pair\", \"connector_credential_pair_id\"),\n (\"user_group__connector_credential_pair\", \"cc_pair_id\"),\n ]\n\n for table_name, column_name in tables_to_clean:\n if table_name in inspector.get_table_names():\n deleted = batch_delete(\n bind, table_name, column_name, cc_pair_ids, id_type=\"int\"\n )\n logger.info(f\"Deleted {deleted} records from {table_name}\")\n\n # === Step 3: Identify connectors and credentials to delete ===\n logger.info(\"Identifying orphaned connectors and credentials...\")\n\n # Get connectors used only by user-file cc_pairs\n connector_rows = bind.execute(\n text(\n \"\"\"\n SELECT DISTINCT ccp.connector_id\n FROM connector_credential_pair ccp\n WHERE ccp.is_user_file IS TRUE\n AND ccp.connector_id != 0 -- Exclude system default\n AND NOT EXISTS (\n SELECT 1\n FROM connector_credential_pair c2\n WHERE c2.connector_id = ccp.connector_id\n AND c2.is_user_file IS NOT TRUE\n )\n \"\"\"\n )\n ).fetchall()\n\n userfile_only_connector_ids = [r[0] for r in connector_rows]\n\n # Get credentials used only by user-file cc_pairs\n credential_rows = bind.execute(\n text(\n \"\"\"\n SELECT DISTINCT ccp.credential_id\n FROM connector_credential_pair ccp\n WHERE ccp.is_user_file IS TRUE\n AND ccp.credential_id != 0 -- Exclude public/default\n AND NOT EXISTS (\n SELECT 1\n FROM connector_credential_pair c2\n WHERE c2.credential_id = ccp.credential_id\n AND c2.is_user_file IS NOT TRUE\n )\n \"\"\"\n )\n ).fetchall()\n\n userfile_only_credential_ids = [r[0] for r in credential_rows]\n\n # === Step 4: Delete the cc_pairs themselves ===\n if cc_pair_ids:\n # Remove FK dependency from user_file first\n bind.execute(\n text(\n \"\"\"\n DO $$\n DECLARE r RECORD;\n BEGIN\n FOR r IN (\n SELECT conname\n FROM pg_constraint c\n JOIN pg_class t ON c.conrelid = t.oid\n JOIN pg_class ft ON c.confrelid = ft.oid\n WHERE c.contype = 'f'\n AND t.relname = 'user_file'\n AND ft.relname = 'connector_credential_pair'\n ) LOOP\n EXECUTE format('ALTER TABLE user_file DROP CONSTRAINT IF EXISTS %I', r.conname);\n END LOOP;\n END$$;\n \"\"\"\n )\n )\n\n # Delete cc_pairs\n deleted = batch_delete(\n bind, \"connector_credential_pair\", \"id\", cc_pair_ids, id_type=\"int\"\n )\n logger.info(f\"Deleted {deleted} connector_credential_pair records\")\n\n # === Step 5: Delete orphaned connectors ===\n if userfile_only_connector_ids:\n deleted = batch_delete(\n bind, \"connector\", \"id\", userfile_only_connector_ids, id_type=\"int\"\n )\n logger.info(f\"Deleted {deleted} orphaned connector records\")\n\n # === Step 6: Delete orphaned credentials ===\n if userfile_only_credential_ids:\n # Clean up credential__user_group mappings first\n deleted = batch_delete(\n bind,\n \"credential__user_group\",\n \"credential_id\",\n userfile_only_credential_ids,\n id_type=\"int\",\n )\n logger.info(f\"Deleted {deleted} credential__user_group records\")\n\n # Delete credentials\n deleted = batch_delete(\n bind, \"credential\", \"id\", userfile_only_credential_ids, id_type=\"int\"\n )\n logger.info(f\"Deleted {deleted} orphaned credential records\")\n\n logger.info(\"Migration 5 (legacy data cleanup) completed successfully\")\n\n\ndef downgrade() -> None:\n \"\"\"Cannot restore deleted data - requires backup restoration.\"\"\"\n\n logger.error(\"CRITICAL: Downgrading data cleanup cannot restore deleted data!\")\n logger.error(\"Data restoration requires backup files or database backup.\")\n\n # raise NotImplementedError(\n # \"Downgrade of legacy data cleanup is not supported. \"\n # \"Deleted data must be restored from backups.\"\n # )\n" }, { "path": "backend/alembic/versions/3b25685ff73c_move_is_public_to_cc_pair.py", "content": "\"\"\"Move is_public to cc_pair\n\nRevision ID: 3b25685ff73c\nRevises: e0a68a81d434\nCreate Date: 2023-10-05 18:47:09.582849\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"3b25685ff73c\"\ndown_revision = \"e0a68a81d434\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"is_public\", sa.Boolean(), nullable=True),\n )\n # fill in is_public for existing rows\n op.execute(\n \"UPDATE connector_credential_pair SET is_public = true WHERE is_public IS NULL\"\n )\n op.alter_column(\"connector_credential_pair\", \"is_public\", nullable=False)\n\n op.add_column(\n \"credential\",\n sa.Column(\"is_admin\", sa.Boolean(), nullable=True),\n )\n op.execute(\"UPDATE credential SET is_admin = true WHERE is_admin IS NULL\")\n op.alter_column(\"credential\", \"is_admin\", nullable=False)\n\n op.drop_column(\"credential\", \"public_doc\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"credential\",\n sa.Column(\"public_doc\", sa.Boolean(), nullable=True),\n )\n # setting public_doc to false for all existing rows to be safe\n # NOTE: this is likely not the correct state of the world but it's the best we can do\n op.execute(\"UPDATE credential SET public_doc = false WHERE public_doc IS NULL\")\n op.alter_column(\"credential\", \"public_doc\", nullable=False)\n op.drop_column(\"connector_credential_pair\", \"is_public\")\n op.drop_column(\"credential\", \"is_admin\")\n" }, { "path": "backend/alembic/versions/3bd4c84fe72f_improved_index.py", "content": "\"\"\"improved index\n\nRevision ID: 3bd4c84fe72f\nRevises: 8f43500ee275\nCreate Date: 2025-02-26 13:07:56.217791\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"3bd4c84fe72f\"\ndown_revision = \"8f43500ee275\"\nbranch_labels = None\ndepends_on = None\n\n\n# NOTE:\n# This migration addresses issues with the previous migration (8f43500ee275) which caused\n# an outage by creating an index without using CONCURRENTLY. This migration:\n#\n# 1. Creates more efficient full-text search capabilities using tsvector columns and GIN indexes\n# 2. Adds indexes to both chat_message and chat_session tables for comprehensive search\n# 3. Note: CONCURRENTLY was removed due to operational issues\n\n\ndef upgrade() -> None:\n # First, drop any existing indexes to avoid conflicts\n op.execute(\"DROP INDEX IF EXISTS idx_chat_message_tsv;\")\n op.execute(\"DROP INDEX IF EXISTS idx_chat_session_desc_tsv;\")\n op.execute(\"DROP INDEX IF EXISTS idx_chat_message_message_lower;\")\n\n # Drop existing columns if they exist\n op.execute(\"ALTER TABLE chat_message DROP COLUMN IF EXISTS message_tsv;\")\n op.execute(\"ALTER TABLE chat_session DROP COLUMN IF EXISTS description_tsv;\")\n\n # Create a GIN index for full-text search on chat_message.message\n op.execute(\n \"\"\"\n ALTER TABLE chat_message\n ADD COLUMN message_tsv tsvector\n GENERATED ALWAYS AS (to_tsvector('english', message)) STORED;\n \"\"\"\n )\n\n op.execute(\n \"\"\"\n CREATE INDEX IF NOT EXISTS idx_chat_message_tsv\n ON chat_message\n USING GIN (message_tsv)\n \"\"\"\n )\n\n # Also add a stored tsvector column for chat_session.description\n op.execute(\n \"\"\"\n ALTER TABLE chat_session\n ADD COLUMN description_tsv tsvector\n GENERATED ALWAYS AS (to_tsvector('english', coalesce(description, ''))) STORED;\n \"\"\"\n )\n\n op.execute(\n \"\"\"\n CREATE INDEX IF NOT EXISTS idx_chat_session_desc_tsv\n ON chat_session\n USING GIN (description_tsv)\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n # Drop the indexes first\n op.execute(\"DROP INDEX IF EXISTS idx_chat_message_tsv;\")\n op.execute(\"DROP INDEX IF EXISTS idx_chat_session_desc_tsv;\")\n\n # Then drop the columns\n op.execute(\"ALTER TABLE chat_message DROP COLUMN IF EXISTS message_tsv;\")\n op.execute(\"ALTER TABLE chat_session DROP COLUMN IF EXISTS description_tsv;\")\n\n op.execute(\"DROP INDEX IF EXISTS idx_chat_message_message_lower;\")\n" }, { "path": "backend/alembic/versions/3c5e35aa9af0_polling_document_count.py", "content": "\"\"\"Polling Document Count\n\nRevision ID: 3c5e35aa9af0\nRevises: 27c6ecc08586\nCreate Date: 2023-06-14 23:45:51.760440\n\n\"\"\"\n\nimport sqlalchemy as sa\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"3c5e35aa9af0\"\ndown_revision = \"27c6ecc08586\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"last_successful_index_time\",\n sa.DateTime(timezone=True),\n nullable=True,\n ),\n )\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"last_attempt_status\",\n sa.Enum(\n \"NOT_STARTED\",\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n name=\"indexingstatus\",\n native_enum=False,\n ),\n nullable=False,\n ),\n )\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"total_docs_indexed\", sa.Integer(), nullable=False),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"total_docs_indexed\")\n op.drop_column(\"connector_credential_pair\", \"last_attempt_status\")\n op.drop_column(\"connector_credential_pair\", \"last_successful_index_time\")\n" }, { "path": "backend/alembic/versions/3c6531f32351_add_back_input_prompts.py", "content": "\"\"\"add back input prompts\n\nRevision ID: 3c6531f32351\nRevises: aeda5f2df4f6\nCreate Date: 2025-01-13 12:49:51.705235\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport fastapi_users_db_sqlalchemy\n\n# revision identifiers, used by Alembic.\nrevision = \"3c6531f32351\"\ndown_revision = \"aeda5f2df4f6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"inputprompt\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\"prompt\", sa.String(), nullable=False),\n sa.Column(\"content\", sa.String(), nullable=False),\n sa.Column(\"active\", sa.Boolean(), nullable=False),\n sa.Column(\"is_public\", sa.Boolean(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"inputprompt__user\",\n sa.Column(\"input_prompt_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\", fastapi_users_db_sqlalchemy.generics.GUID(), nullable=False\n ),\n sa.Column(\"disabled\", sa.Boolean(), nullable=False, default=False),\n sa.ForeignKeyConstraint(\n [\"input_prompt_id\"],\n [\"inputprompt.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"input_prompt_id\", \"user_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"inputprompt__user\")\n op.drop_table(\"inputprompt\")\n" }, { "path": "backend/alembic/versions/3c9a65f1207f_seed_exa_provider_from_env.py", "content": "\"\"\"seed_exa_provider_from_env\n\nRevision ID: 3c9a65f1207f\nRevises: 1f2a3b4c5d6e\nCreate Date: 2025-11-20 19:18:00.000000\n\n\"\"\"\n\nfrom __future__ import annotations\n\nimport os\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom dotenv import load_dotenv, find_dotenv\n\nfrom onyx.utils.encryption import encrypt_string_to_bytes\n\nrevision = \"3c9a65f1207f\"\ndown_revision = \"1f2a3b4c5d6e\"\nbranch_labels = None\ndepends_on = None\n\n\nEXA_PROVIDER_NAME = \"Exa\"\n\n\ndef _get_internet_search_table(metadata: sa.MetaData) -> sa.Table:\n return sa.Table(\n \"internet_search_provider\",\n metadata,\n sa.Column(\"id\", sa.Integer, primary_key=True),\n sa.Column(\"name\", sa.String),\n sa.Column(\"provider_type\", sa.String),\n sa.Column(\"api_key\", sa.LargeBinary),\n sa.Column(\"config\", postgresql.JSONB),\n sa.Column(\"is_active\", sa.Boolean),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.text(\"now()\"),\n ),\n )\n\n\ndef upgrade() -> None:\n load_dotenv(find_dotenv())\n\n exa_api_key = os.environ.get(\"EXA_API_KEY\")\n if not exa_api_key:\n return\n\n bind = op.get_bind()\n metadata = sa.MetaData()\n table = _get_internet_search_table(metadata)\n\n existing = bind.execute(\n sa.select(table.c.id).where(table.c.name == EXA_PROVIDER_NAME)\n ).first()\n if existing:\n return\n\n encrypted_key = encrypt_string_to_bytes(exa_api_key)\n\n has_active_provider = bind.execute(\n sa.select(table.c.id).where(table.c.is_active.is_(True))\n ).first()\n\n bind.execute(\n table.insert().values(\n name=EXA_PROVIDER_NAME,\n provider_type=\"exa\",\n api_key=encrypted_key,\n config=None,\n is_active=not bool(has_active_provider),\n )\n )\n\n\ndef downgrade() -> None:\n return\n" }, { "path": "backend/alembic/versions/3d1cca026fe8_add_oauth_config_and_user_tokens.py", "content": "\"\"\"add_oauth_config_and_user_tokens\n\nRevision ID: 3d1cca026fe8\nRevises: c8a93a2af083\nCreate Date: 2025-10-21 13:27:34.274721\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"3d1cca026fe8\"\ndown_revision = \"c8a93a2af083\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create oauth_config table\n op.create_table(\n \"oauth_config\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"authorization_url\", sa.Text(), nullable=False),\n sa.Column(\"token_url\", sa.Text(), nullable=False),\n sa.Column(\"client_id\", sa.LargeBinary(), nullable=False),\n sa.Column(\"client_secret\", sa.LargeBinary(), nullable=False),\n sa.Column(\"scopes\", postgresql.JSONB(astext_type=sa.Text()), nullable=True),\n sa.Column(\n \"additional_params\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n\n # Create oauth_user_token table\n op.create_table(\n \"oauth_user_token\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"oauth_config_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\"token_data\", sa.LargeBinary(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"oauth_config_id\"], [\"oauth_config.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"oauth_config_id\", \"user_id\", name=\"uq_oauth_user_token\"),\n )\n\n # Create index on user_id for efficient user-based token lookups\n # Note: unique constraint on (oauth_config_id, user_id) already creates\n # an index for config-based lookups\n op.create_index(\n \"ix_oauth_user_token_user_id\",\n \"oauth_user_token\",\n [\"user_id\"],\n )\n\n # Add oauth_config_id column to tool table\n op.add_column(\"tool\", sa.Column(\"oauth_config_id\", sa.Integer(), nullable=True))\n\n # Create foreign key from tool to oauth_config\n op.create_foreign_key(\n \"tool_oauth_config_fk\",\n \"tool\",\n \"oauth_config\",\n [\"oauth_config_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n\n\ndef downgrade() -> None:\n # Drop foreign key from tool to oauth_config\n op.drop_constraint(\"tool_oauth_config_fk\", \"tool\", type_=\"foreignkey\")\n\n # Drop oauth_config_id column from tool table\n op.drop_column(\"tool\", \"oauth_config_id\")\n\n # Drop index on user_id\n op.drop_index(\"ix_oauth_user_token_user_id\", table_name=\"oauth_user_token\")\n\n # Drop oauth_user_token table (will cascade delete tokens)\n op.drop_table(\"oauth_user_token\")\n\n # Drop oauth_config table\n op.drop_table(\"oauth_config\")\n" }, { "path": "backend/alembic/versions/3fc5d75723b3_add_doc_metadata_field_in_document_model.py", "content": "\"\"\"add_doc_metadata_field_in_document_model\n\nRevision ID: 3fc5d75723b3\nRevises: 2f95e36923e6\nCreate Date: 2025-07-28 18:45:37.985406\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"3fc5d75723b3\"\ndown_revision = \"2f95e36923e6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"document\",\n sa.Column(\n \"doc_metadata\", postgresql.JSONB(astext_type=sa.Text()), nullable=True\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"document\", \"doc_metadata\")\n" }, { "path": "backend/alembic/versions/401c1ac29467_add_tables_for_ui_based_llm_.py", "content": "\"\"\"Add tables for UI-based LLM configuration\n\nRevision ID: 401c1ac29467\nRevises: 703313b75876\nCreate Date: 2024-04-13 18:07:29.153817\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"401c1ac29467\"\ndown_revision = \"703313b75876\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"llm_provider\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"api_key\", sa.String(), nullable=True),\n sa.Column(\"api_base\", sa.String(), nullable=True),\n sa.Column(\"api_version\", sa.String(), nullable=True),\n sa.Column(\n \"custom_config\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n sa.Column(\"default_model_name\", sa.String(), nullable=False),\n sa.Column(\"fast_default_model_name\", sa.String(), nullable=True),\n sa.Column(\"is_default_provider\", sa.Boolean(), unique=True, nullable=True),\n sa.Column(\"model_names\", postgresql.ARRAY(sa.String()), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n\n op.add_column(\n \"persona\",\n sa.Column(\"llm_model_provider_override\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"llm_model_provider_override\")\n\n op.drop_table(\"llm_provider\")\n" }, { "path": "backend/alembic/versions/40926a4dab77_reset_userfile_document_id_migrated_.py", "content": "\"\"\"reset userfile document_id_migrated field\n\nRevision ID: 40926a4dab77\nRevises: 64bd5677aeb6\nCreate Date: 2025-10-06 16:10:32.898668\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"40926a4dab77\"\ndown_revision = \"64bd5677aeb6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Set all existing records to not migrated\n op.execute(\n \"UPDATE user_file SET document_id_migrated = FALSE WHERE document_id_migrated IS DISTINCT FROM FALSE;\"\n )\n\n\ndef downgrade() -> None:\n # No-op\n pass\n" }, { "path": "backend/alembic/versions/41fa44bef321_remove_default_prompt_shortcuts.py", "content": "\"\"\"remove default prompt shortcuts\n\nRevision ID: 41fa44bef321\nRevises: 2c2430828bdf\nCreate Date: 2025-01-21\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"41fa44bef321\"\ndown_revision = \"2c2430828bdf\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Delete any user associations for the default prompts first (foreign key constraint)\n op.execute(\n \"DELETE FROM inputprompt__user WHERE input_prompt_id IN (SELECT id FROM inputprompt WHERE id < 0)\"\n )\n # Delete the pre-seeded default prompt shortcuts (they have negative IDs)\n op.execute(\"DELETE FROM inputprompt WHERE id < 0\")\n\n\ndef downgrade() -> None:\n # We don't restore the default prompts on downgrade\n pass\n" }, { "path": "backend/alembic/versions/43cbbb3f5e6a_rename_index_origin_to_index_recursively.py", "content": "\"\"\"Rename index_origin to index_recursively\n\nRevision ID: 1d6ad76d1f37\nRevises: e1392f05e840\nCreate Date: 2024-08-01 12:38:54.466081\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"1d6ad76d1f37\"\ndown_revision = \"e1392f05e840\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = jsonb_set(\n connector_specific_config,\n '{index_recursively}',\n 'true'::jsonb\n ) - 'index_origin'\n WHERE connector_specific_config ? 'index_origin'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = jsonb_set(\n connector_specific_config,\n '{index_origin}',\n connector_specific_config->'index_recursively'\n ) - 'index_recursively'\n WHERE connector_specific_config ? 'index_recursively'\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/44f856ae2a4a_add_cloud_embedding_model.py", "content": "\"\"\"add cloud embedding model and update embedding_model\n\nRevision ID: 44f856ae2a4a\nRevises: d716b0791ddd\nCreate Date: 2024-06-28 20:01:05.927647\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"44f856ae2a4a\"\ndown_revision = \"d716b0791ddd\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Create embedding_provider table\n op.create_table(\n \"embedding_provider\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"api_key\", sa.LargeBinary(), nullable=True),\n sa.Column(\"default_model_id\", sa.Integer(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n\n # Add cloud_provider_id to embedding_model table\n op.add_column(\n \"embedding_model\", sa.Column(\"cloud_provider_id\", sa.Integer(), nullable=True)\n )\n\n # Add foreign key constraints\n op.create_foreign_key(\n \"fk_embedding_model_cloud_provider\",\n \"embedding_model\",\n \"embedding_provider\",\n [\"cloud_provider_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"fk_embedding_provider_default_model\",\n \"embedding_provider\",\n \"embedding_model\",\n [\"default_model_id\"],\n [\"id\"],\n )\n\n\ndef downgrade() -> None:\n # Remove foreign key constraints\n op.drop_constraint(\n \"fk_embedding_model_cloud_provider\", \"embedding_model\", type_=\"foreignkey\"\n )\n op.drop_constraint(\n \"fk_embedding_provider_default_model\", \"embedding_provider\", type_=\"foreignkey\"\n )\n\n # Remove cloud_provider_id column\n op.drop_column(\"embedding_model\", \"cloud_provider_id\")\n\n # Drop embedding_provider table\n op.drop_table(\"embedding_provider\")\n" }, { "path": "backend/alembic/versions/4505fd7302e1_added_is_internet_to_dbdoc.py", "content": "\"\"\"added is_internet to DBDoc\n\nRevision ID: 4505fd7302e1\nRevises: c18cdf4b497e\nCreate Date: 2024-06-18 20:46:09.095034\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"4505fd7302e1\"\ndown_revision = \"c18cdf4b497e\"\n\n\ndef upgrade() -> None:\n op.add_column(\"search_doc\", sa.Column(\"is_internet\", sa.Boolean(), nullable=True))\n op.add_column(\"tool\", sa.Column(\"display_name\", sa.String(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"tool\", \"display_name\")\n op.drop_column(\"search_doc\", \"is_internet\")\n" }, { "path": "backend/alembic/versions/465f78d9b7f9_larger_access_tokens_for_oauth.py", "content": "\"\"\"Larger Access Tokens for OAUTH\n\nRevision ID: 465f78d9b7f9\nRevises: 3c5e35aa9af0\nCreate Date: 2023-07-18 17:33:40.365034\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"465f78d9b7f9\"\ndown_revision = \"3c5e35aa9af0\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"oauth_account\", \"access_token\", type_=sa.Text())\n\n\ndef downgrade() -> None:\n op.alter_column(\"oauth_account\", \"access_token\", type_=sa.String(length=1024))\n" }, { "path": "backend/alembic/versions/46625e4745d4_remove_native_enum.py", "content": "\"\"\"Remove Native Enum\n\nRevision ID: 46625e4745d4\nRevises: 9d97fecfab7f\nCreate Date: 2023-10-27 11:38:33.803145\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy import String\n\n# revision identifiers, used by Alembic.\nrevision = \"46625e4745d4\"\ndown_revision = \"9d97fecfab7f\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # At this point, we directly changed some previous migrations,\n # https://github.com/onyx-dot-app/onyx/pull/637\n # Due to using Postgres native Enums, it caused some complications for first time users.\n # To remove those complications, all Enums are only handled application side moving forward.\n # This migration exists to ensure that existing users don't run into upgrade issues.\n op.alter_column(\"index_attempt\", \"status\", type_=String)\n op.alter_column(\"connector_credential_pair\", \"last_attempt_status\", type_=String)\n op.execute(\"DROP TYPE IF EXISTS indexingstatus\")\n\n\ndef downgrade() -> None:\n # We don't want Native Enums, do nothing\n pass\n" }, { "path": "backend/alembic/versions/46b7a812670f_fix_user__external_user_group_id_fk.py", "content": "\"\"\"fix_user__external_user_group_id_fk\n\nRevision ID: 46b7a812670f\nRevises: f32615f71aeb\nCreate Date: 2024-09-23 12:58:03.894038\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"46b7a812670f\"\ndown_revision = \"f32615f71aeb\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Drop the existing primary key\n op.drop_constraint(\n \"user__external_user_group_id_pkey\",\n \"user__external_user_group_id\",\n type_=\"primary\",\n )\n\n # Add the new composite primary key\n op.create_primary_key(\n \"user__external_user_group_id_pkey\",\n \"user__external_user_group_id\",\n [\"user_id\", \"external_user_group_id\", \"cc_pair_id\"],\n )\n\n\ndef downgrade() -> None:\n # Drop the composite primary key\n op.drop_constraint(\n \"user__external_user_group_id_pkey\",\n \"user__external_user_group_id\",\n type_=\"primary\",\n )\n # Delete all entries from the table\n op.execute(\"DELETE FROM user__external_user_group_id\")\n\n # Recreate the original primary key on user_id\n op.create_primary_key(\n \"user__external_user_group_id_pkey\", \"user__external_user_group_id\", [\"user_id\"]\n )\n" }, { "path": "backend/alembic/versions/4738e4b3bae1_pg_file_store.py", "content": "\"\"\"PG File Store\n\nRevision ID: 4738e4b3bae1\nRevises: e91df4e935ef\nCreate Date: 2024-03-20 18:53:32.461518\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"4738e4b3bae1\"\ndown_revision = \"e91df4e935ef\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"file_store\",\n sa.Column(\"file_name\", sa.String(), nullable=False),\n sa.Column(\"lobj_oid\", sa.Integer(), nullable=False),\n sa.PrimaryKeyConstraint(\"file_name\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"file_store\")\n" }, { "path": "backend/alembic/versions/473a1a7ca408_add_display_model_names_to_llm_provider.py", "content": "\"\"\"Add display_model_names to llm_provider\n\nRevision ID: 473a1a7ca408\nRevises: 325975216eb3\nCreate Date: 2024-07-25 14:31:02.002917\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"473a1a7ca408\"\ndown_revision = \"325975216eb3\"\nbranch_labels: None = None\ndepends_on: None = None\n\ndefault_models_by_provider = {\n \"openai\": [\"gpt-4\", \"gpt-4o\", \"gpt-4o-mini\"],\n \"bedrock\": [\n \"meta.llama3-1-70b-instruct-v1:0\",\n \"meta.llama3-1-8b-instruct-v1:0\",\n \"anthropic.claude-3-opus-20240229-v1:0\",\n \"mistral.mistral-large-2402-v1:0\",\n \"anthropic.claude-3-5-sonnet-20240620-v1:0\",\n ],\n \"anthropic\": [\"claude-3-opus-20240229\", \"claude-3-5-sonnet-20240620\"],\n}\n\n\ndef upgrade() -> None:\n op.add_column(\n \"llm_provider\",\n sa.Column(\"display_model_names\", postgresql.ARRAY(sa.String()), nullable=True),\n )\n\n connection = op.get_bind()\n for provider, models in default_models_by_provider.items():\n connection.execute(\n sa.text(\n \"UPDATE llm_provider SET display_model_names = :models WHERE provider = :provider\"\n ),\n {\"models\": models, \"provider\": provider},\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"llm_provider\", \"display_model_names\")\n" }, { "path": "backend/alembic/versions/47433d30de82_create_indexattempt_table.py", "content": "\"\"\"Create IndexAttempt table\n\nRevision ID: 47433d30de82\nRevises:\nCreate Date: 2023-05-04 00:55:32.971991\n\n\"\"\"\n\nimport sqlalchemy as sa\nfrom alembic import op\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"47433d30de82\"\ndown_revision: None = None\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"index_attempt\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n # String type since python enum will change often\n sa.Column(\n \"source\",\n sa.String(),\n nullable=False,\n ),\n # String type to easily accomodate new ways of pulling\n # in documents\n sa.Column(\n \"input_type\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\n \"connector_specific_config\",\n postgresql.JSONB(),\n nullable=False,\n ),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=True,\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n server_onupdate=sa.text(\"now()\"), # type: ignore\n nullable=True,\n ),\n sa.Column(\n \"status\",\n sa.Enum(\n \"NOT_STARTED\",\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n name=\"indexingstatus\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\"document_ids\", postgresql.ARRAY(sa.String()), nullable=True),\n sa.Column(\"error_msg\", sa.String(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"index_attempt\")\n" }, { "path": "backend/alembic/versions/475fcefe8826_add_name_to_api_key.py", "content": "\"\"\"Add name to api_key\n\nRevision ID: 475fcefe8826\nRevises: ecab2b3f1a3b\nCreate Date: 2024-04-11 11:05:18.414438\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"475fcefe8826\"\ndown_revision = \"ecab2b3f1a3b\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"api_key\", sa.Column(\"name\", sa.String(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"api_key\", \"name\")\n" }, { "path": "backend/alembic/versions/4794bc13e484_update_prompt_length.py", "content": "\"\"\"update prompt length\n\nRevision ID: 4794bc13e484\nRevises: f7505c5b0284\nCreate Date: 2025-04-02 11:26:36.180328\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"4794bc13e484\"\ndown_revision = \"f7505c5b0284\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"prompt\",\n \"system_prompt\",\n existing_type=sa.TEXT(),\n type_=sa.String(length=5000000),\n existing_nullable=False,\n )\n op.alter_column(\n \"prompt\",\n \"task_prompt\",\n existing_type=sa.TEXT(),\n type_=sa.String(length=5000000),\n existing_nullable=False,\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"prompt\",\n \"system_prompt\",\n existing_type=sa.String(length=5000000),\n type_=sa.TEXT(),\n existing_nullable=False,\n )\n op.alter_column(\n \"prompt\",\n \"task_prompt\",\n existing_type=sa.String(length=5000000),\n type_=sa.TEXT(),\n existing_nullable=False,\n )\n" }, { "path": "backend/alembic/versions/47a07e1a38f1_fix_invalid_model_configurations_state.py", "content": "\"\"\"Fix invalid model-configurations state\n\nRevision ID: 47a07e1a38f1\nRevises: 7a70b7664e37\nCreate Date: 2025-04-23 15:39:43.159504\n\n\"\"\"\n\nfrom alembic import op\nfrom pydantic import BaseModel, ConfigDict\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\nfrom onyx.llm.well_known_providers.llm_provider_options import (\n fetch_model_names_for_provider_as_set,\n fetch_visible_model_names_for_provider_as_set,\n)\n\n\n# revision identifiers, used by Alembic.\nrevision = \"47a07e1a38f1\"\ndown_revision = \"7a70b7664e37\"\nbranch_labels = None\ndepends_on = None\n\n\nclass _SimpleModelConfiguration(BaseModel):\n # Configure model to read from attributes\n model_config = ConfigDict(from_attributes=True)\n\n id: int\n llm_provider_id: int\n name: str\n is_visible: bool\n max_input_tokens: int | None\n\n\ndef upgrade() -> None:\n llm_provider_table = sa.sql.table(\n \"llm_provider\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"provider\", sa.String),\n sa.column(\"model_names\", postgresql.ARRAY(sa.String)),\n sa.column(\"display_model_names\", postgresql.ARRAY(sa.String)),\n sa.column(\"default_model_name\", sa.String),\n sa.column(\"fast_default_model_name\", sa.String),\n )\n model_configuration_table = sa.sql.table(\n \"model_configuration\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"llm_provider_id\", sa.Integer),\n sa.column(\"name\", sa.String),\n sa.column(\"is_visible\", sa.Boolean),\n sa.column(\"max_input_tokens\", sa.Integer),\n )\n\n connection = op.get_bind()\n\n llm_providers = connection.execute(\n sa.select(\n llm_provider_table.c.id,\n llm_provider_table.c.provider,\n )\n ).fetchall()\n\n for llm_provider in llm_providers:\n llm_provider_id, provider_name = llm_provider\n\n default_models = fetch_model_names_for_provider_as_set(provider_name)\n display_models = fetch_visible_model_names_for_provider_as_set(\n provider_name=provider_name\n )\n\n # if `fetch_model_names_for_provider_as_set` returns `None`, then\n # that means that `provider_name` is not a well-known llm provider.\n if not default_models:\n continue\n\n if not display_models:\n raise RuntimeError(\n \"If `default_models` is non-None, `display_models` must be non-None too.\"\n )\n\n model_configurations = [\n _SimpleModelConfiguration.model_validate(model_configuration)\n for model_configuration in connection.execute(\n sa.select(\n model_configuration_table.c.id,\n model_configuration_table.c.llm_provider_id,\n model_configuration_table.c.name,\n model_configuration_table.c.is_visible,\n model_configuration_table.c.max_input_tokens,\n ).where(model_configuration_table.c.llm_provider_id == llm_provider_id)\n ).fetchall()\n ]\n\n if model_configurations:\n at_least_one_is_visible = any(\n [\n model_configuration.is_visible\n for model_configuration in model_configurations\n ]\n )\n\n # If there is at least one model which is public, this is a valid state.\n # Therefore, don't touch it and move on to the next one.\n if at_least_one_is_visible:\n continue\n\n existing_visible_model_names: set[str] = set(\n [\n model_configuration.name\n for model_configuration in model_configurations\n if model_configuration.is_visible\n ]\n )\n\n difference = display_models.difference(existing_visible_model_names)\n\n for model_name in difference:\n if not model_name:\n continue\n\n insert_statement = postgresql.insert(model_configuration_table).values(\n llm_provider_id=llm_provider_id,\n name=model_name,\n is_visible=True,\n max_input_tokens=None,\n )\n\n connection.execute(\n insert_statement.on_conflict_do_update(\n index_elements=[\"llm_provider_id\", \"name\"],\n set_={\"is_visible\": insert_statement.excluded.is_visible},\n )\n )\n else:\n for model_name in default_models:\n connection.execute(\n model_configuration_table.insert().values(\n llm_provider_id=llm_provider_id,\n name=model_name,\n is_visible=model_name in display_models,\n max_input_tokens=None,\n )\n )\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/47e5bef3a1d7_add_persona_categories.py", "content": "\"\"\"add persona categories\n\nRevision ID: 47e5bef3a1d7\nRevises: dfbe9e93d3c7\nCreate Date: 2024-11-05 18:55:02.221064\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"47e5bef3a1d7\"\ndown_revision = \"dfbe9e93d3c7\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create the persona_category table\n op.create_table(\n \"persona_category\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"description\", sa.String(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n\n # Add category_id to persona table\n op.add_column(\"persona\", sa.Column(\"category_id\", sa.Integer(), nullable=True))\n op.create_foreign_key(\n \"fk_persona_category\",\n \"persona\",\n \"persona_category\",\n [\"category_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\"persona_category_id_fkey\", \"persona\", type_=\"foreignkey\")\n op.drop_column(\"persona\", \"category_id\")\n op.drop_table(\"persona_category\")\n" }, { "path": "backend/alembic/versions/48d14957fe80_add_support_for_custom_tools.py", "content": "\"\"\"Add support for custom tools\n\nRevision ID: 48d14957fe80\nRevises: b85f02ec1308\nCreate Date: 2024-06-09 14:58:19.946509\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"48d14957fe80\"\ndown_revision = \"b85f02ec1308\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"tool\",\n sa.Column(\n \"openapi_schema\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n op.add_column(\n \"tool\",\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n )\n op.create_foreign_key(\"tool_user_fk\", \"tool\", \"user\", [\"user_id\"], [\"id\"])\n\n op.create_table(\n \"tool_call\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"tool_id\", sa.Integer(), nullable=False),\n sa.Column(\"tool_name\", sa.String(), nullable=False),\n sa.Column(\n \"tool_arguments\", postgresql.JSONB(astext_type=sa.Text()), nullable=False\n ),\n sa.Column(\n \"tool_result\", postgresql.JSONB(astext_type=sa.Text()), nullable=False\n ),\n sa.Column(\n \"message_id\", sa.Integer(), sa.ForeignKey(\"chat_message.id\"), nullable=False\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"tool_call\")\n\n op.drop_constraint(\"tool_user_fk\", \"tool\", type_=\"foreignkey\")\n op.drop_column(\"tool\", \"user_id\")\n op.drop_column(\"tool\", \"openapi_schema\")\n" }, { "path": "backend/alembic/versions/495cb26ce93e_create_knowlege_graph_tables.py", "content": "\"\"\"create knowledge graph tables\n\nRevision ID: 495cb26ce93e\nRevises: ca04500b9ee8\nCreate Date: 2025-03-19 08:51:14.341989\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy import text\nfrom datetime import datetime, timedelta\n\nfrom onyx.configs.app_configs import DB_READONLY_USER\nfrom onyx.configs.app_configs import DB_READONLY_PASSWORD\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\n\n# revision identifiers, used by Alembic.\nrevision = \"495cb26ce93e\"\ndown_revision = \"ca04500b9ee8\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n\n # Create a new permission-less user to be later used for knowledge graph queries.\n # The user will later get temporary read privileges for a specific view that will be\n # ad hoc generated specific to a knowledge graph query.\n #\n # Note: in order for the migration to run, the DB_READONLY_USER and DB_READONLY_PASSWORD\n # environment variables MUST be set. Otherwise, an exception will be raised.\n\n if not MULTI_TENANT:\n # Enable pg_trgm extension if not already enabled\n op.execute(\"CREATE EXTENSION IF NOT EXISTS pg_trgm\")\n\n # Create read-only db user here only in single tenant mode. For multi-tenant mode,\n # the user is created in the alembic_tenants migration.\n if not (DB_READONLY_USER and DB_READONLY_PASSWORD):\n raise Exception(\"DB_READONLY_USER or DB_READONLY_PASSWORD is not set\")\n\n op.execute(\n text(\n f\"\"\"\n DO $$\n BEGIN\n -- Check if the read-only user already exists\n IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN\n -- Create the read-only user with the specified password\n EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');\n -- First revoke all privileges to ensure a clean slate\n EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');\n -- Grant only the CONNECT privilege to allow the user to connect to the database\n -- but not perform any operations without additional specific grants\n EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');\n END IF;\n END\n $$;\n \"\"\"\n )\n )\n\n # Grant usage on current schema to readonly user\n op.execute(\n text(\n f\"\"\"\n DO $$\n BEGIN\n IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN\n EXECUTE format('GRANT USAGE ON SCHEMA %I TO %I', current_schema(), '{DB_READONLY_USER}');\n END IF;\n END\n $$;\n \"\"\"\n )\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_config CASCADE\")\n op.create_table(\n \"kg_config\",\n sa.Column(\"id\", sa.Integer(), primary_key=True, nullable=False, index=True),\n sa.Column(\"kg_variable_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"kg_variable_values\", postgresql.ARRAY(sa.String()), nullable=False),\n sa.UniqueConstraint(\"kg_variable_name\", name=\"uq_kg_config_variable_name\"),\n )\n\n # Insert initial data into kg_config table\n op.bulk_insert(\n sa.table(\n \"kg_config\",\n sa.column(\"kg_variable_name\", sa.String),\n sa.column(\"kg_variable_values\", postgresql.ARRAY(sa.String)),\n ),\n [\n {\"kg_variable_name\": \"KG_EXPOSED\", \"kg_variable_values\": [\"false\"]},\n {\"kg_variable_name\": \"KG_ENABLED\", \"kg_variable_values\": [\"false\"]},\n {\"kg_variable_name\": \"KG_VENDOR\", \"kg_variable_values\": []},\n {\"kg_variable_name\": \"KG_VENDOR_DOMAINS\", \"kg_variable_values\": []},\n {\"kg_variable_name\": \"KG_IGNORE_EMAIL_DOMAINS\", \"kg_variable_values\": []},\n {\n \"kg_variable_name\": \"KG_EXTRACTION_IN_PROGRESS\",\n \"kg_variable_values\": [\"false\"],\n },\n {\n \"kg_variable_name\": \"KG_CLUSTERING_IN_PROGRESS\",\n \"kg_variable_values\": [\"false\"],\n },\n {\n \"kg_variable_name\": \"KG_COVERAGE_START\",\n \"kg_variable_values\": [\n (datetime.now() - timedelta(days=90)).strftime(\"%Y-%m-%d\")\n ],\n },\n {\"kg_variable_name\": \"KG_MAX_COVERAGE_DAYS\", \"kg_variable_values\": [\"90\"]},\n {\n \"kg_variable_name\": \"KG_MAX_PARENT_RECURSION_DEPTH\",\n \"kg_variable_values\": [\"2\"],\n },\n ],\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_entity_type CASCADE\")\n op.create_table(\n \"kg_entity_type\",\n sa.Column(\"id_name\", sa.String(), primary_key=True, nullable=False, index=True),\n sa.Column(\"description\", sa.String(), nullable=True),\n sa.Column(\"grounding\", sa.String(), nullable=False),\n sa.Column(\n \"attributes\",\n postgresql.JSONB,\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\"active\", sa.Boolean(), nullable=False, default=False),\n sa.Column(\"deep_extraction\", sa.Boolean(), nullable=False, default=False),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.Column(\"grounded_source_name\", sa.String(), nullable=True),\n sa.Column(\"entity_values\", postgresql.ARRAY(sa.String()), nullable=True),\n sa.Column(\n \"clustering\",\n postgresql.JSONB,\n nullable=False,\n server_default=\"{}\",\n ),\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_relationship_type CASCADE\")\n # Create KGRelationshipType table\n op.create_table(\n \"kg_relationship_type\",\n sa.Column(\"id_name\", sa.String(), primary_key=True, nullable=False, index=True),\n sa.Column(\"name\", sa.String(), nullable=False, index=True),\n sa.Column(\n \"source_entity_type_id_name\", sa.String(), nullable=False, index=True\n ),\n sa.Column(\n \"target_entity_type_id_name\", sa.String(), nullable=False, index=True\n ),\n sa.Column(\"definition\", sa.Boolean(), nullable=False, default=False),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\"type\", sa.String(), nullable=False, index=True),\n sa.Column(\"active\", sa.Boolean(), nullable=False, default=True),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.Column(\n \"clustering\",\n postgresql.JSONB,\n nullable=False,\n server_default=\"{}\",\n ),\n sa.ForeignKeyConstraint(\n [\"source_entity_type_id_name\"], [\"kg_entity_type.id_name\"]\n ),\n sa.ForeignKeyConstraint(\n [\"target_entity_type_id_name\"], [\"kg_entity_type.id_name\"]\n ),\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_relationship_type_extraction_staging CASCADE\")\n # Create KGRelationshipTypeExtractionStaging table\n op.create_table(\n \"kg_relationship_type_extraction_staging\",\n sa.Column(\"id_name\", sa.String(), primary_key=True, nullable=False, index=True),\n sa.Column(\"name\", sa.String(), nullable=False, index=True),\n sa.Column(\n \"source_entity_type_id_name\", sa.String(), nullable=False, index=True\n ),\n sa.Column(\n \"target_entity_type_id_name\", sa.String(), nullable=False, index=True\n ),\n sa.Column(\"definition\", sa.Boolean(), nullable=False, default=False),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\"type\", sa.String(), nullable=False, index=True),\n sa.Column(\"active\", sa.Boolean(), nullable=False, default=True),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.Column(\n \"clustering\",\n postgresql.JSONB,\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\"transferred\", sa.Boolean(), nullable=False, server_default=\"false\"),\n sa.ForeignKeyConstraint(\n [\"source_entity_type_id_name\"], [\"kg_entity_type.id_name\"]\n ),\n sa.ForeignKeyConstraint(\n [\"target_entity_type_id_name\"], [\"kg_entity_type.id_name\"]\n ),\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_entity CASCADE\")\n\n # Create KGEntity table\n op.create_table(\n \"kg_entity\",\n sa.Column(\"id_name\", sa.String(), primary_key=True, nullable=False, index=True),\n sa.Column(\"name\", sa.String(), nullable=False, index=True),\n sa.Column(\"entity_class\", sa.String(), nullable=True, index=True),\n sa.Column(\"entity_subtype\", sa.String(), nullable=True, index=True),\n sa.Column(\"entity_key\", sa.String(), nullable=True, index=True),\n sa.Column(\"name_trigrams\", postgresql.ARRAY(sa.String(3)), nullable=True),\n sa.Column(\"document_id\", sa.String(), nullable=True, index=True),\n sa.Column(\n \"alternative_names\",\n postgresql.ARRAY(sa.String()),\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\"entity_type_id_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"description\", sa.String(), nullable=True),\n sa.Column(\n \"keywords\",\n postgresql.ARRAY(sa.String()),\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\n \"acl\", postgresql.ARRAY(sa.String()), nullable=False, server_default=\"{}\"\n ),\n sa.Column(\"boosts\", postgresql.JSONB, nullable=False, server_default=\"{}\"),\n sa.Column(\"attributes\", postgresql.JSONB, nullable=False, server_default=\"{}\"),\n sa.Column(\"event_time\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.ForeignKeyConstraint([\"entity_type_id_name\"], [\"kg_entity_type.id_name\"]),\n sa.ForeignKeyConstraint([\"document_id\"], [\"document.id\"]),\n sa.UniqueConstraint(\n \"name\",\n \"entity_type_id_name\",\n \"document_id\",\n name=\"uq_kg_entity_name_type_doc\",\n ),\n )\n op.create_index(\"ix_entity_type_acl\", \"kg_entity\", [\"entity_type_id_name\", \"acl\"])\n op.create_index(\n \"ix_entity_name_search\", \"kg_entity\", [\"name\", \"entity_type_id_name\"]\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_entity_extraction_staging CASCADE\")\n # Create KGEntityExtractionStaging table\n op.create_table(\n \"kg_entity_extraction_staging\",\n sa.Column(\"id_name\", sa.String(), primary_key=True, nullable=False, index=True),\n sa.Column(\"name\", sa.String(), nullable=False, index=True),\n sa.Column(\"document_id\", sa.String(), nullable=True, index=True),\n sa.Column(\n \"alternative_names\",\n postgresql.ARRAY(sa.String()),\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\"entity_type_id_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"description\", sa.String(), nullable=True),\n sa.Column(\n \"keywords\",\n postgresql.ARRAY(sa.String()),\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\n \"acl\", postgresql.ARRAY(sa.String()), nullable=False, server_default=\"{}\"\n ),\n sa.Column(\"boosts\", postgresql.JSONB, nullable=False, server_default=\"{}\"),\n sa.Column(\"attributes\", postgresql.JSONB, nullable=False, server_default=\"{}\"),\n sa.Column(\"transferred_id_name\", sa.String(), nullable=True, default=None),\n sa.Column(\"entity_class\", sa.String(), nullable=True, index=True),\n sa.Column(\"entity_key\", sa.String(), nullable=True, index=True),\n sa.Column(\"entity_subtype\", sa.String(), nullable=True, index=True),\n sa.Column(\"parent_key\", sa.String(), nullable=True, index=True),\n sa.Column(\"event_time\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.ForeignKeyConstraint([\"entity_type_id_name\"], [\"kg_entity_type.id_name\"]),\n sa.ForeignKeyConstraint([\"document_id\"], [\"document.id\"]),\n )\n op.create_index(\n \"ix_entity_extraction_staging_acl\",\n \"kg_entity_extraction_staging\",\n [\"entity_type_id_name\", \"acl\"],\n )\n op.create_index(\n \"ix_entity_extraction_staging_name_search\",\n \"kg_entity_extraction_staging\",\n [\"name\", \"entity_type_id_name\"],\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_relationship CASCADE\")\n # Create KGRelationship table\n op.create_table(\n \"kg_relationship\",\n sa.Column(\"id_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"source_node\", sa.String(), nullable=False, index=True),\n sa.Column(\"target_node\", sa.String(), nullable=False, index=True),\n sa.Column(\"source_node_type\", sa.String(), nullable=False, index=True),\n sa.Column(\"target_node_type\", sa.String(), nullable=False, index=True),\n sa.Column(\"source_document\", sa.String(), nullable=True, index=True),\n sa.Column(\"type\", sa.String(), nullable=False, index=True),\n sa.Column(\"relationship_type_id_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.ForeignKeyConstraint([\"source_node\"], [\"kg_entity.id_name\"]),\n sa.ForeignKeyConstraint([\"target_node\"], [\"kg_entity.id_name\"]),\n sa.ForeignKeyConstraint([\"source_node_type\"], [\"kg_entity_type.id_name\"]),\n sa.ForeignKeyConstraint([\"target_node_type\"], [\"kg_entity_type.id_name\"]),\n sa.ForeignKeyConstraint([\"source_document\"], [\"document.id\"]),\n sa.ForeignKeyConstraint(\n [\"relationship_type_id_name\"], [\"kg_relationship_type.id_name\"]\n ),\n sa.UniqueConstraint(\n \"source_node\",\n \"target_node\",\n \"type\",\n name=\"uq_kg_relationship_source_target_type\",\n ),\n sa.PrimaryKeyConstraint(\"id_name\", \"source_document\"),\n )\n op.create_index(\n \"ix_kg_relationship_nodes\", \"kg_relationship\", [\"source_node\", \"target_node\"]\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_relationship_extraction_staging CASCADE\")\n # Create KGRelationshipExtractionStaging table\n op.create_table(\n \"kg_relationship_extraction_staging\",\n sa.Column(\"id_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"source_node\", sa.String(), nullable=False, index=True),\n sa.Column(\"target_node\", sa.String(), nullable=False, index=True),\n sa.Column(\"source_node_type\", sa.String(), nullable=False, index=True),\n sa.Column(\"target_node_type\", sa.String(), nullable=False, index=True),\n sa.Column(\"source_document\", sa.String(), nullable=True, index=True),\n sa.Column(\"type\", sa.String(), nullable=False, index=True),\n sa.Column(\"relationship_type_id_name\", sa.String(), nullable=False, index=True),\n sa.Column(\"occurrences\", sa.Integer(), server_default=\"1\", nullable=False),\n sa.Column(\"transferred\", sa.Boolean(), nullable=False, server_default=\"false\"),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n sa.ForeignKeyConstraint(\n [\"source_node\"], [\"kg_entity_extraction_staging.id_name\"]\n ),\n sa.ForeignKeyConstraint(\n [\"target_node\"], [\"kg_entity_extraction_staging.id_name\"]\n ),\n sa.ForeignKeyConstraint([\"source_node_type\"], [\"kg_entity_type.id_name\"]),\n sa.ForeignKeyConstraint([\"target_node_type\"], [\"kg_entity_type.id_name\"]),\n sa.ForeignKeyConstraint([\"source_document\"], [\"document.id\"]),\n sa.ForeignKeyConstraint(\n [\"relationship_type_id_name\"],\n [\"kg_relationship_type_extraction_staging.id_name\"],\n ),\n sa.UniqueConstraint(\n \"source_node\",\n \"target_node\",\n \"type\",\n name=\"uq_kg_relationship_extraction_staging_source_target_type\",\n ),\n sa.PrimaryKeyConstraint(\"id_name\", \"source_document\"),\n )\n op.create_index(\n \"ix_kg_relationship_extraction_staging_nodes\",\n \"kg_relationship_extraction_staging\",\n [\"source_node\", \"target_node\"],\n )\n\n op.execute(\"DROP TABLE IF EXISTS kg_term CASCADE\")\n # Create KGTerm table\n op.create_table(\n \"kg_term\",\n sa.Column(\"id_term\", sa.String(), primary_key=True, nullable=False, index=True),\n sa.Column(\n \"entity_types\",\n postgresql.ARRAY(sa.String()),\n nullable=False,\n server_default=\"{}\",\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n ),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.text(\"now()\")\n ),\n )\n op.create_index(\"ix_search_term_entities\", \"kg_term\", [\"entity_types\"])\n op.create_index(\"ix_search_term_term\", \"kg_term\", [\"id_term\"])\n\n op.add_column(\n \"document\",\n sa.Column(\"kg_stage\", sa.String(), nullable=True, index=True),\n )\n op.add_column(\n \"document\",\n sa.Column(\"kg_processing_time\", sa.DateTime(timezone=True), nullable=True),\n )\n op.add_column(\n \"connector\",\n sa.Column(\n \"kg_processing_enabled\",\n sa.Boolean(),\n nullable=True,\n server_default=\"false\",\n ),\n )\n\n op.add_column(\n \"connector\",\n sa.Column(\n \"kg_coverage_days\",\n sa.Integer(),\n nullable=True,\n server_default=None,\n ),\n )\n\n # Create GIN index for clustering and normalization\n op.execute(\n \"CREATE INDEX IF NOT EXISTS idx_kg_entity_clustering_trigrams \"\n f\"ON kg_entity USING GIN (name {POSTGRES_DEFAULT_SCHEMA}.gin_trgm_ops)\"\n )\n op.execute(\n \"CREATE INDEX IF NOT EXISTS idx_kg_entity_normalization_trigrams ON kg_entity USING GIN (name_trigrams)\"\n )\n\n # Create kg_entity trigger to update kg_entity.name and its trigrams\n alphanum_pattern = r\"[^a-z0-9]+\"\n truncate_length = 1000\n function = \"update_kg_entity_name\"\n op.execute(\n text(\n f\"\"\"\n CREATE OR REPLACE FUNCTION {function}()\n RETURNS TRIGGER AS $$\n DECLARE\n name text;\n cleaned_name text;\n BEGIN\n -- Set name to semantic_id if document_id is not NULL\n IF NEW.document_id IS NOT NULL THEN\n SELECT lower(semantic_id) INTO name\n FROM document\n WHERE id = NEW.document_id;\n ELSE\n name = lower(NEW.name);\n END IF;\n\n -- Clean name and truncate if too long\n cleaned_name = regexp_replace(\n name,\n '{alphanum_pattern}', '', 'g'\n );\n IF length(cleaned_name) > {truncate_length} THEN\n cleaned_name = left(cleaned_name, {truncate_length});\n END IF;\n\n -- Set name and name trigrams\n NEW.name = name;\n NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name);\n RETURN NEW;\n END;\n $$ LANGUAGE plpgsql;\n \"\"\"\n )\n )\n trigger = f\"{function}_trigger\"\n op.execute(f\"DROP TRIGGER IF EXISTS {trigger} ON kg_entity\")\n op.execute(\n f\"\"\"\n CREATE TRIGGER {trigger}\n BEFORE INSERT OR UPDATE OF name\n ON kg_entity\n FOR EACH ROW\n EXECUTE FUNCTION {function}();\n \"\"\"\n )\n\n # Create kg_entity trigger to update kg_entity.name and its trigrams\n function = \"update_kg_entity_name_from_doc\"\n op.execute(\n text(\n f\"\"\"\n CREATE OR REPLACE FUNCTION {function}()\n RETURNS TRIGGER AS $$\n DECLARE\n doc_name text;\n cleaned_name text;\n BEGIN\n doc_name = lower(NEW.semantic_id);\n\n -- Clean name and truncate if too long\n cleaned_name = regexp_replace(\n doc_name,\n '{alphanum_pattern}', '', 'g'\n );\n IF length(cleaned_name) > {truncate_length} THEN\n cleaned_name = left(cleaned_name, {truncate_length});\n END IF;\n\n -- Set name and name trigrams for all entities referencing this document\n UPDATE kg_entity\n SET\n name = doc_name,\n name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name)\n WHERE document_id = NEW.id;\n RETURN NEW;\n END;\n $$ LANGUAGE plpgsql;\n \"\"\"\n )\n )\n trigger = f\"{function}_trigger\"\n op.execute(f\"DROP TRIGGER IF EXISTS {trigger} ON document\")\n op.execute(\n f\"\"\"\n CREATE TRIGGER {trigger}\n AFTER UPDATE OF semantic_id\n ON document\n FOR EACH ROW\n EXECUTE FUNCTION {function}();\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n\n # Drop all views that start with 'kg_'\n op.execute(\n \"\"\"\n DO $$\n DECLARE\n view_name text;\n BEGIN\n FOR view_name IN\n SELECT c.relname\n FROM pg_catalog.pg_class c\n JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace\n WHERE c.relkind = 'v'\n AND n.nspname = current_schema()\n AND c.relname LIKE 'kg_relationships_with_access%'\n LOOP\n EXECUTE 'DROP VIEW IF EXISTS ' || quote_ident(view_name);\n END LOOP;\n END $$;\n \"\"\"\n )\n\n op.execute(\n \"\"\"\n DO $$\n DECLARE\n view_name text;\n BEGIN\n FOR view_name IN\n SELECT c.relname\n FROM pg_catalog.pg_class c\n JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace\n WHERE c.relkind = 'v'\n AND n.nspname = current_schema()\n AND c.relname LIKE 'allowed_docs%'\n LOOP\n EXECUTE 'DROP VIEW IF EXISTS ' || quote_ident(view_name);\n END LOOP;\n END $$;\n \"\"\"\n )\n\n for table, function in (\n (\"kg_entity\", \"update_kg_entity_name\"),\n (\"document\", \"update_kg_entity_name_from_doc\"),\n ):\n op.execute(f\"DROP TRIGGER IF EXISTS {function}_trigger ON {table}\")\n op.execute(f\"DROP FUNCTION IF EXISTS {function}()\")\n\n # Drop index\n op.execute(\"DROP INDEX IF EXISTS idx_kg_entity_clustering_trigrams\")\n op.execute(\"DROP INDEX IF EXISTS idx_kg_entity_normalization_trigrams\")\n\n # Drop tables in reverse order of creation to handle dependencies\n op.drop_table(\"kg_term\")\n op.drop_table(\"kg_relationship\")\n op.drop_table(\"kg_entity\")\n op.drop_table(\"kg_relationship_type\")\n op.drop_table(\"kg_relationship_extraction_staging\")\n op.drop_table(\"kg_relationship_type_extraction_staging\")\n op.drop_table(\"kg_entity_extraction_staging\")\n op.drop_table(\"kg_entity_type\")\n op.drop_column(\"connector\", \"kg_processing_enabled\")\n op.drop_column(\"connector\", \"kg_coverage_days\")\n op.drop_column(\"document\", \"kg_stage\")\n op.drop_column(\"document\", \"kg_processing_time\")\n op.drop_table(\"kg_config\")\n\n # Revoke usage on current schema for the readonly user\n op.execute(\n text(\n f\"\"\"\n DO $$\n BEGIN\n IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN\n EXECUTE format('REVOKE ALL ON SCHEMA %I FROM %I', current_schema(), '{DB_READONLY_USER}');\n END IF;\n END\n $$;\n \"\"\"\n )\n )\n\n if not MULTI_TENANT:\n # Drop read-only db user here only in single tenant mode. For multi-tenant mode,\n # the user is dropped in the alembic_tenants migration.\n\n op.execute(\n text(\n f\"\"\"\n DO $$\n BEGIN\n IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN\n -- First revoke all privileges from the database\n EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');\n -- Then drop the user\n EXECUTE format('DROP USER %I', '{DB_READONLY_USER}');\n END IF;\n END\n $$;\n \"\"\"\n )\n )\n op.execute(text(\"DROP EXTENSION IF EXISTS pg_trgm\"))\n" }, { "path": "backend/alembic/versions/4a1e4b1c89d2_add_indexing_to_userfilestatus.py", "content": "\"\"\"Add INDEXING to UserFileStatus\n\nRevision ID: 4a1e4b1c89d2\nRevises: 6b3b4083c5aa\nCreate Date: 2026-02-28 00:00:00.000000\n\n\"\"\"\n\nimport sqlalchemy as sa\nfrom alembic import op\n\nrevision = \"4a1e4b1c89d2\"\ndown_revision = \"6b3b4083c5aa\"\nbranch_labels = None\ndepends_on = None\n\nTABLE = \"user_file\"\nCOLUMN = \"status\"\nCONSTRAINT_NAME = \"ck_user_file_status\"\n\nOLD_VALUES = (\"PROCESSING\", \"COMPLETED\", \"FAILED\", \"CANCELED\", \"DELETING\")\nNEW_VALUES = (\"PROCESSING\", \"INDEXING\", \"COMPLETED\", \"FAILED\", \"CANCELED\", \"DELETING\")\n\n\ndef _drop_status_check_constraint() -> None:\n \"\"\"Drop the existing CHECK constraint on user_file.status.\n\n The constraint name is auto-generated by SQLAlchemy and unknown,\n so we look it up via the inspector.\n \"\"\"\n inspector = sa.inspect(op.get_bind())\n for constraint in inspector.get_check_constraints(TABLE):\n if COLUMN in constraint.get(\"sqltext\", \"\"):\n constraint_name = constraint[\"name\"]\n if constraint_name is not None:\n op.drop_constraint(constraint_name, TABLE, type_=\"check\")\n\n\ndef upgrade() -> None:\n _drop_status_check_constraint()\n in_clause = \", \".join(f\"'{v}'\" for v in NEW_VALUES)\n op.create_check_constraint(CONSTRAINT_NAME, TABLE, f\"{COLUMN} IN ({in_clause})\")\n\n\ndef downgrade() -> None:\n op.execute(\n f\"UPDATE {TABLE} SET {COLUMN} = 'PROCESSING' WHERE {COLUMN} = 'INDEXING'\"\n )\n op.drop_constraint(CONSTRAINT_NAME, TABLE, type_=\"check\")\n in_clause = \", \".join(f\"'{v}'\" for v in OLD_VALUES)\n op.create_check_constraint(CONSTRAINT_NAME, TABLE, f\"{COLUMN} IN ({in_clause})\")\n" }, { "path": "backend/alembic/versions/4a951134c801_moved_status_to_connector_credential_.py", "content": "\"\"\"Moved status to connector credential pair\n\nRevision ID: 4a951134c801\nRevises: 7477a5f5d728\nCreate Date: 2024-08-10 19:20:34.527559\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"4a951134c801\"\ndown_revision = \"7477a5f5d728\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"status\",\n sa.Enum(\n \"ACTIVE\",\n \"PAUSED\",\n \"DELETING\",\n name=\"connectorcredentialpairstatus\",\n native_enum=False,\n ),\n nullable=True,\n ),\n )\n\n # Update status of connector_credential_pair based on connector's disabled status\n op.execute(\n \"\"\"\n UPDATE connector_credential_pair\n SET status = CASE\n WHEN (\n SELECT disabled\n FROM connector\n WHERE connector.id = connector_credential_pair.connector_id\n ) = FALSE THEN 'ACTIVE'\n ELSE 'PAUSED'\n END\n \"\"\"\n )\n\n # Make the status column not nullable after setting values\n op.alter_column(\"connector_credential_pair\", \"status\", nullable=False)\n\n op.drop_column(\"connector\", \"disabled\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"connector\",\n sa.Column(\"disabled\", sa.BOOLEAN(), autoincrement=False, nullable=True),\n )\n\n # Update disabled status of connector based on connector_credential_pair's status\n op.execute(\n \"\"\"\n UPDATE connector\n SET disabled = CASE\n WHEN EXISTS (\n SELECT 1\n FROM connector_credential_pair\n WHERE connector_credential_pair.connector_id = connector.id\n AND connector_credential_pair.status = 'ACTIVE'\n ) THEN FALSE\n ELSE TRUE\n END\n \"\"\"\n )\n\n # Make the disabled column not nullable after setting values\n op.alter_column(\"connector\", \"disabled\", nullable=False)\n\n op.drop_column(\"connector_credential_pair\", \"status\")\n" }, { "path": "backend/alembic/versions/4b08d97e175a_change_default_prune_freq.py", "content": "\"\"\"change default prune_freq\n\nRevision ID: 4b08d97e175a\nRevises: d9ec13955951\nCreate Date: 2024-08-20 15:28:52.993827\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"4b08d97e175a\"\ndown_revision = \"d9ec13955951\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE connector\n SET prune_freq = 2592000\n WHERE prune_freq = 86400\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE connector\n SET prune_freq = 86400\n WHERE prune_freq = 2592000\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/4cebcbc9b2ae_add_tab_index_to_tool_call.py", "content": "\"\"\"add tab_index to tool_call\n\nRevision ID: 4cebcbc9b2ae\nRevises: a1b2c3d4e5f6\nCreate Date: 2025-12-16\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"4cebcbc9b2ae\"\ndown_revision = \"a1b2c3d4e5f6\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"tool_call\",\n sa.Column(\"tab_index\", sa.Integer(), nullable=False, server_default=\"0\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"tool_call\", \"tab_index\")\n" }, { "path": "backend/alembic/versions/4d58345da04a_lowercase_user_emails.py", "content": "\"\"\"lowercase_user_emails\n\nRevision ID: 4d58345da04a\nRevises: f1ca58b2f2ec\nCreate Date: 2025-01-29 07:48:46.784041\n\n\"\"\"\n\nimport logging\nfrom typing import cast\nfrom alembic import op\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.sql import text\n\n\n# revision identifiers, used by Alembic.\nrevision = \"4d58345da04a\"\ndown_revision = \"f1ca58b2f2ec\"\nbranch_labels = None\ndepends_on = None\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n\ndef upgrade() -> None:\n \"\"\"Conflicts on lowercasing will result in the uppercased email getting a\n unique integer suffix when converted to lowercase.\"\"\"\n\n connection = op.get_bind()\n\n # Fetch all user emails that are not already lowercase\n user_emails = connection.execute(\n text('SELECT id, email FROM \"user\" WHERE email != LOWER(email)')\n ).fetchall()\n\n for user_id, email in user_emails:\n email = cast(str, email)\n username, domain = email.rsplit(\"@\", 1)\n new_email = f\"{username.lower()}@{domain.lower()}\"\n attempt = 1\n\n while True:\n try:\n # Try updating the email\n connection.execute(\n text('UPDATE \"user\" SET email = :new_email WHERE id = :user_id'),\n {\"new_email\": new_email, \"user_id\": user_id},\n )\n break # Success, exit loop\n except IntegrityError:\n next_email = f\"{username.lower()}_{attempt}@{domain.lower()}\"\n # Email conflict occurred, append `_1`, `_2`, etc., to the username\n logger.warning(\n f\"Conflict while lowercasing email: old_email={email} conflicting_email={new_email} next_email={next_email}\"\n )\n new_email = next_email\n attempt += 1\n\n\ndef downgrade() -> None:\n # Cannot restore original case of emails\n pass\n" }, { "path": "backend/alembic/versions/4ea2c93919c1_add_type_to_credentials.py", "content": "\"\"\"Add type to credentials\n\nRevision ID: 4ea2c93919c1\nRevises: 473a1a7ca408\nCreate Date: 2024-07-18 13:07:13.655895\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"4ea2c93919c1\"\ndown_revision = \"473a1a7ca408\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Add the new 'source' column to the 'credential' table\n op.add_column(\n \"credential\",\n sa.Column(\n \"source\",\n sa.String(length=100), # Use String instead of Enum\n nullable=True, # Initially allow NULL values\n ),\n )\n op.add_column(\n \"credential\",\n sa.Column(\n \"name\",\n sa.String(),\n nullable=True,\n ),\n )\n\n # Create a temporary table that maps each credential to a single connector source.\n # This is needed because a credential can be associated with multiple connectors,\n # but we want to assign a single source to each credential.\n # We use DISTINCT ON to ensure we only get one row per credential_id.\n op.execute(\n \"\"\"\n CREATE TEMPORARY TABLE temp_connector_credential AS\n SELECT DISTINCT ON (cc.credential_id)\n cc.credential_id,\n c.source AS connector_source\n FROM connector_credential_pair cc\n JOIN connector c ON cc.connector_id = c.id\n \"\"\"\n )\n\n # Update the 'source' column in the 'credential' table\n op.execute(\n \"\"\"\n UPDATE credential cred\n SET source = COALESCE(\n (SELECT connector_source\n FROM temp_connector_credential temp\n WHERE cred.id = temp.credential_id),\n 'NOT_APPLICABLE'\n )\n \"\"\"\n )\n\n # Drop the temporary table to avoid conflicts if migration runs again\n # (e.g., during upgrade -> downgrade -> upgrade cycles in tests)\n op.execute(\"DROP TABLE IF EXISTS temp_connector_credential\")\n\n # If no exception was raised, alter the column\n op.alter_column(\"credential\", \"source\", nullable=True) # TODO modify\n # # ### end Alembic commands ###\n\n\ndef downgrade() -> None:\n op.drop_column(\"credential\", \"source\")\n op.drop_column(\"credential\", \"name\")\n" }, { "path": "backend/alembic/versions/4ee1287bd26a_add_multiple_slack_bot_support.py", "content": "\"\"\"add_multiple_slack_bot_support\n\nRevision ID: 4ee1287bd26a\nRevises: 47e5bef3a1d7\nCreate Date: 2024-11-06 13:15:53.302644\n\n\"\"\"\n\nfrom typing import cast\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.orm import Session\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.db.models import SlackBot\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"4ee1287bd26a\"\ndown_revision = \"47e5bef3a1d7\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Create new slack_bot table\n op.create_table(\n \"slack_bot\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"enabled\", sa.Boolean(), nullable=False, server_default=\"true\"),\n sa.Column(\"bot_token\", sa.LargeBinary(), nullable=False),\n sa.Column(\"app_token\", sa.LargeBinary(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"bot_token\"),\n sa.UniqueConstraint(\"app_token\"),\n )\n\n # # Create new slack_channel_config table\n op.create_table(\n \"slack_channel_config\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"slack_bot_id\", sa.Integer(), nullable=True),\n sa.Column(\"persona_id\", sa.Integer(), nullable=True),\n sa.Column(\"channel_config\", postgresql.JSONB(), nullable=False),\n sa.Column(\"response_type\", sa.String(), nullable=False),\n sa.Column(\n \"enable_auto_filters\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n sa.ForeignKeyConstraint(\n [\"slack_bot_id\"],\n [\"slack_bot.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Handle existing Slack bot tokens first\n bot_token = None\n app_token = None\n first_row_id = None\n\n try:\n tokens = cast(dict, get_kv_store().load(\"slack_bot_tokens_config_key\"))\n except Exception:\n tokens = {}\n\n bot_token = tokens.get(\"bot_token\")\n app_token = tokens.get(\"app_token\")\n\n if bot_token and app_token:\n session = Session(bind=op.get_bind())\n new_slack_bot = SlackBot(\n name=\"Slack Bot (Migrated)\",\n enabled=True,\n bot_token=bot_token,\n app_token=app_token,\n )\n session.add(new_slack_bot)\n session.commit()\n first_row_id = new_slack_bot.id\n\n # Create a default bot if none exists\n # This is in case there are no slack tokens but there are channels configured\n op.execute(\n sa.text(\n \"\"\"\n INSERT INTO slack_bot (name, enabled, bot_token, app_token)\n SELECT 'Default Bot', true, '', ''\n WHERE NOT EXISTS (SELECT 1 FROM slack_bot)\n RETURNING id;\n \"\"\"\n )\n )\n\n # Get the bot ID to use (either from existing migration or newly created)\n bot_id_query = sa.text(\n \"\"\"\n SELECT COALESCE(\n :first_row_id,\n (SELECT id FROM slack_bot ORDER BY id ASC LIMIT 1)\n ) as bot_id;\n \"\"\"\n )\n result = op.get_bind().execute(bot_id_query, {\"first_row_id\": first_row_id})\n bot_id = result.scalar()\n\n # CTE (Common Table Expression) that transforms the old slack_bot_config table data\n # This splits up the channel_names into their own rows\n channel_names_cte = \"\"\"\n WITH channel_names AS (\n SELECT\n sbc.id as config_id,\n sbc.persona_id,\n sbc.response_type,\n sbc.enable_auto_filters,\n jsonb_array_elements_text(sbc.channel_config->'channel_names') as channel_name,\n sbc.channel_config->>'respond_tag_only' as respond_tag_only,\n sbc.channel_config->>'respond_to_bots' as respond_to_bots,\n sbc.channel_config->'respond_member_group_list' as respond_member_group_list,\n sbc.channel_config->'answer_filters' as answer_filters,\n sbc.channel_config->'follow_up_tags' as follow_up_tags\n FROM slack_bot_config sbc\n )\n \"\"\"\n\n # Insert the channel names into the new slack_channel_config table\n insert_statement = \"\"\"\n INSERT INTO slack_channel_config (\n slack_bot_id,\n persona_id,\n channel_config,\n response_type,\n enable_auto_filters\n )\n SELECT\n :bot_id,\n channel_name.persona_id,\n jsonb_build_object(\n 'channel_name', channel_name.channel_name,\n 'respond_tag_only',\n COALESCE((channel_name.respond_tag_only)::boolean, false),\n 'respond_to_bots',\n COALESCE((channel_name.respond_to_bots)::boolean, false),\n 'respond_member_group_list',\n COALESCE(channel_name.respond_member_group_list, '[]'::jsonb),\n 'answer_filters',\n COALESCE(channel_name.answer_filters, '[]'::jsonb),\n 'follow_up_tags',\n COALESCE(channel_name.follow_up_tags, '[]'::jsonb)\n ),\n channel_name.response_type,\n channel_name.enable_auto_filters\n FROM channel_names channel_name;\n \"\"\"\n\n op.execute(sa.text(channel_names_cte + insert_statement).bindparams(bot_id=bot_id))\n\n # Clean up old tokens if they existed\n try:\n if bot_token and app_token:\n get_kv_store().delete(\"slack_bot_tokens_config_key\")\n except Exception:\n pass\n # Rename the table\n op.rename_table(\n \"slack_bot_config__standard_answer_category\",\n \"slack_channel_config__standard_answer_category\",\n )\n\n # Rename the column\n op.alter_column(\n \"slack_channel_config__standard_answer_category\",\n \"slack_bot_config_id\",\n new_column_name=\"slack_channel_config_id\",\n )\n\n # Drop the table with CASCADE to handle dependent objects\n op.execute(\"DROP TABLE slack_bot_config CASCADE\")\n\n\ndef downgrade() -> None:\n # Recreate the old slack_bot_config table\n op.create_table(\n \"slack_bot_config\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"persona_id\", sa.Integer(), nullable=True),\n sa.Column(\"channel_config\", postgresql.JSONB(), nullable=False),\n sa.Column(\"response_type\", sa.String(), nullable=False),\n sa.Column(\"enable_auto_filters\", sa.Boolean(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Migrate data back to the old format\n # Group by persona_id to combine channel names back into arrays\n op.execute(\n sa.text(\n \"\"\"\n INSERT INTO slack_bot_config (\n persona_id,\n channel_config,\n response_type,\n enable_auto_filters\n )\n SELECT DISTINCT ON (persona_id)\n persona_id,\n jsonb_build_object(\n 'channel_names', (\n SELECT jsonb_agg(c.channel_config->>'channel_name')\n FROM slack_channel_config c\n WHERE c.persona_id = scc.persona_id\n ),\n 'respond_tag_only', (channel_config->>'respond_tag_only')::boolean,\n 'respond_to_bots', (channel_config->>'respond_to_bots')::boolean,\n 'respond_member_group_list', channel_config->'respond_member_group_list',\n 'answer_filters', channel_config->'answer_filters',\n 'follow_up_tags', channel_config->'follow_up_tags'\n ),\n response_type,\n enable_auto_filters\n FROM slack_channel_config scc\n WHERE persona_id IS NOT NULL;\n \"\"\"\n )\n )\n\n # Rename the table back\n op.rename_table(\n \"slack_channel_config__standard_answer_category\",\n \"slack_bot_config__standard_answer_category\",\n )\n\n # Rename the column back\n op.alter_column(\n \"slack_bot_config__standard_answer_category\",\n \"slack_channel_config_id\",\n new_column_name=\"slack_bot_config_id\",\n )\n\n # Try to save the first bot's tokens back to KV store\n try:\n first_bot = (\n op.get_bind()\n .execute(\n sa.text(\n \"SELECT bot_token, app_token FROM slack_bot ORDER BY id LIMIT 1\"\n )\n )\n .first()\n )\n if first_bot and first_bot.bot_token and first_bot.app_token:\n tokens = {\n \"bot_token\": first_bot.bot_token,\n \"app_token\": first_bot.app_token,\n }\n get_kv_store().store(\"slack_bot_tokens_config_key\", tokens)\n except Exception:\n pass\n\n # Drop the new tables in reverse order\n op.drop_table(\"slack_channel_config\")\n op.drop_table(\"slack_bot\")\n" }, { "path": "backend/alembic/versions/4f8a2b3c1d9e_add_open_url_tool.py", "content": "\"\"\"add_open_url_tool\n\nRevision ID: 4f8a2b3c1d9e\nRevises: a852cbe15577\nCreate Date: 2025-11-24 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"4f8a2b3c1d9e\"\ndown_revision = \"a852cbe15577\"\nbranch_labels = None\ndepends_on = None\n\n\nOPEN_URL_TOOL = {\n \"name\": \"OpenURLTool\",\n \"display_name\": \"Open URL\",\n \"description\": (\n \"The Open URL Action allows the agent to fetch and read contents of web pages.\"\n ),\n \"in_code_tool_id\": \"OpenURLTool\",\n \"enabled\": True,\n}\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Check if tool already exists\n existing = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE in_code_tool_id = :in_code_tool_id\"),\n {\"in_code_tool_id\": OPEN_URL_TOOL[\"in_code_tool_id\"]},\n ).fetchone()\n\n if existing:\n tool_id = existing[0]\n # Update existing tool\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :name,\n display_name = :display_name,\n description = :description\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n OPEN_URL_TOOL,\n )\n else:\n # Insert new tool\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)\n VALUES (:name, :display_name, :description, :in_code_tool_id, :enabled)\n \"\"\"\n ),\n OPEN_URL_TOOL,\n )\n # Get the newly inserted tool's id\n result = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE in_code_tool_id = :in_code_tool_id\"),\n {\"in_code_tool_id\": OPEN_URL_TOOL[\"in_code_tool_id\"]},\n ).fetchone()\n tool_id = result[0] # type: ignore\n\n # Associate the tool with all existing personas\n # Get all persona IDs\n persona_ids = conn.execute(sa.text(\"SELECT id FROM persona\")).fetchall()\n\n for (persona_id,) in persona_ids:\n # Check if association already exists\n exists = conn.execute(\n sa.text(\n \"\"\"\n SELECT 1 FROM persona__tool\n WHERE persona_id = :persona_id AND tool_id = :tool_id\n \"\"\"\n ),\n {\"persona_id\": persona_id, \"tool_id\": tool_id},\n ).fetchone()\n\n if not exists:\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona__tool (persona_id, tool_id)\n VALUES (:persona_id, :tool_id)\n \"\"\"\n ),\n {\"persona_id\": persona_id, \"tool_id\": tool_id},\n )\n\n\ndef downgrade() -> None:\n # We don't remove the tool on downgrade since it's fine to have it around.\n # If we upgrade again, it will be a no-op.\n pass\n" }, { "path": "backend/alembic/versions/503883791c39_add_effective_permissions.py", "content": "\"\"\"add_effective_permissions\n\nAdds a JSONB column `effective_permissions` to the user table to store\ndirectly granted permissions (e.g. [\"admin\"] or [\"basic\"]). Implied\npermissions are expanded at read time, not stored.\n\nBackfill: joins user__user_group → permission_grant to collect each\nuser's granted permissions into a JSON array. Users without group\nmemberships keep the default [].\n\nRevision ID: 503883791c39\nRevises: b4b7e1028dfd\nCreate Date: 2026-03-30 14:49:22.261748\n\n\"\"\"\n\nfrom collections.abc import Sequence\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"503883791c39\"\ndown_revision = \"b4b7e1028dfd\"\nbranch_labels: str | None = None\ndepends_on: str | Sequence[str] | None = None\n\nuser_table = sa.table(\n \"user\",\n sa.column(\"id\", sa.Uuid),\n sa.column(\"effective_permissions\", postgresql.JSONB),\n)\n\nuser_user_group = sa.table(\n \"user__user_group\",\n sa.column(\"user_id\", sa.Uuid),\n sa.column(\"user_group_id\", sa.Integer),\n)\n\npermission_grant = sa.table(\n \"permission_grant\",\n sa.column(\"group_id\", sa.Integer),\n sa.column(\"permission\", sa.String),\n sa.column(\"is_deleted\", sa.Boolean),\n)\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"effective_permissions\",\n postgresql.JSONB(),\n nullable=False,\n server_default=sa.text(\"'[]'::jsonb\"),\n ),\n )\n\n conn = op.get_bind()\n\n # Deduplicated permissions per user\n deduped = (\n sa.select(\n user_user_group.c.user_id,\n permission_grant.c.permission,\n )\n .select_from(\n user_user_group.join(\n permission_grant,\n sa.and_(\n permission_grant.c.group_id == user_user_group.c.user_group_id,\n permission_grant.c.is_deleted == sa.false(),\n ),\n )\n )\n .distinct()\n .subquery(\"deduped\")\n )\n\n # Aggregate into JSONB array per user (order is not guaranteed;\n # consumers read this as a set so ordering does not matter)\n perms_per_user = (\n sa.select(\n deduped.c.user_id,\n sa.func.jsonb_agg(\n deduped.c.permission,\n type_=postgresql.JSONB,\n ).label(\"perms\"),\n )\n .group_by(deduped.c.user_id)\n .subquery(\"sub\")\n )\n\n conn.execute(\n user_table.update()\n .where(user_table.c.id == perms_per_user.c.user_id)\n .values(effective_permissions=perms_per_user.c.perms)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"effective_permissions\")\n" }, { "path": "backend/alembic/versions/505c488f6662_merge_default_assistants_into_unified.py", "content": "\"\"\"merge_default_assistants_into_unified\n\nRevision ID: 505c488f6662\nRevises: d09fc20a3c66\nCreate Date: 2025-09-09 19:00:56.816626\n\n\"\"\"\n\nimport json\nfrom typing import Any\nfrom typing import NamedTuple\nfrom uuid import UUID\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"505c488f6662\"\ndown_revision = \"d09fc20a3c66\"\nbranch_labels = None\ndepends_on = None\n\n# Constants for the unified assistant\nUNIFIED_ASSISTANT_NAME = \"Assistant\"\nUNIFIED_ASSISTANT_DESCRIPTION = (\n \"Your AI assistant with search, web browsing, and image generation capabilities.\"\n)\nUNIFIED_ASSISTANT_NUM_CHUNKS = 25\nUNIFIED_ASSISTANT_DISPLAY_PRIORITY = 0\nUNIFIED_ASSISTANT_LLM_FILTER_EXTRACTION = True\nUNIFIED_ASSISTANT_LLM_RELEVANCE_FILTER = False\nUNIFIED_ASSISTANT_RECENCY_BIAS = \"AUTO\" # NOTE: needs to be capitalized\nUNIFIED_ASSISTANT_CHUNKS_ABOVE = 0\nUNIFIED_ASSISTANT_CHUNKS_BELOW = 0\nUNIFIED_ASSISTANT_DATETIME_AWARE = True\n\n# NOTE: tool specific prompts are handled on the fly and automatically injected\n# into the prompt before passing to the LLM.\nDEFAULT_SYSTEM_PROMPT = \"\"\"\nYou are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the \\\nuser's intent, ask clarifying questions when needed, think step-by-step through complex problems, \\\nprovide clear and accurate answers, and proactively anticipate helpful follow-up information. Always \\\nprioritize being truthful, nuanced, insightful, and efficient.\nThe current date is [[CURRENT_DATETIME]]\n\nYou use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make \\\nyour responses more readable and engaging.\nYou use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, \\\nsymbols, etc.: '$$\\\\n[expression]\\\\n$$' for standalone cases and '\\\\( [expression] \\\\)' when inline.\nFor code you prefer to use Markdown and specify the language.\nYou can use Markdown horizontal rules (---) to separate sections of your responses.\nYou can use Markdown tables to format your responses for data, lists, and other structured information.\n\"\"\".strip()\n\n\nINSERT_DICT: dict[str, Any] = {\n \"name\": UNIFIED_ASSISTANT_NAME,\n \"description\": UNIFIED_ASSISTANT_DESCRIPTION,\n \"system_prompt\": DEFAULT_SYSTEM_PROMPT,\n \"num_chunks\": UNIFIED_ASSISTANT_NUM_CHUNKS,\n \"display_priority\": UNIFIED_ASSISTANT_DISPLAY_PRIORITY,\n \"llm_filter_extraction\": UNIFIED_ASSISTANT_LLM_FILTER_EXTRACTION,\n \"llm_relevance_filter\": UNIFIED_ASSISTANT_LLM_RELEVANCE_FILTER,\n \"recency_bias\": UNIFIED_ASSISTANT_RECENCY_BIAS,\n \"chunks_above\": UNIFIED_ASSISTANT_CHUNKS_ABOVE,\n \"chunks_below\": UNIFIED_ASSISTANT_CHUNKS_BELOW,\n \"datetime_aware\": UNIFIED_ASSISTANT_DATETIME_AWARE,\n}\n\nGENERAL_ASSISTANT_ID = -1\nART_ASSISTANT_ID = -3\n\n\nclass UserRow(NamedTuple):\n \"\"\"Typed representation of user row from database query.\"\"\"\n\n id: UUID\n chosen_assistants: list[int] | None\n visible_assistants: list[int] | None\n hidden_assistants: list[int] | None\n pinned_assistants: list[int] | None\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Step 1: Create or update the unified assistant (ID 0)\n search_assistant = conn.execute(\n sa.text(\"SELECT * FROM persona WHERE id = 0\")\n ).fetchone()\n\n if search_assistant:\n # Update existing Search assistant to be the unified assistant\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET name = :name,\n description = :description,\n system_prompt = :system_prompt,\n num_chunks = :num_chunks,\n is_default_persona = true,\n is_visible = true,\n deleted = false,\n display_priority = :display_priority,\n llm_filter_extraction = :llm_filter_extraction,\n llm_relevance_filter = :llm_relevance_filter,\n recency_bias = :recency_bias,\n chunks_above = :chunks_above,\n chunks_below = :chunks_below,\n datetime_aware = :datetime_aware,\n starter_messages = null\n WHERE id = 0\n \"\"\"\n ),\n INSERT_DICT,\n )\n else:\n # Create new unified assistant with ID 0\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona (\n id, name, description, system_prompt, num_chunks,\n is_default_persona, is_visible, deleted, display_priority,\n llm_filter_extraction, llm_relevance_filter, recency_bias,\n chunks_above, chunks_below, datetime_aware, starter_messages,\n builtin_persona\n ) VALUES (\n 0, :name, :description, :system_prompt, :num_chunks,\n true, true, false, :display_priority, :llm_filter_extraction,\n :llm_relevance_filter, :recency_bias, :chunks_above, :chunks_below,\n :datetime_aware, null, true\n )\n \"\"\"\n ),\n INSERT_DICT,\n )\n\n # Step 2: Mark ALL builtin assistants as deleted (except the unified assistant ID 0)\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET deleted = true, is_visible = false, is_default_persona = false\n WHERE builtin_persona = true AND id != 0\n \"\"\"\n )\n )\n\n # Step 3: Add all built-in tools to the unified assistant\n # First, get the tool IDs for SearchTool, ImageGenerationTool, and WebSearchTool\n search_tool = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE in_code_tool_id = 'SearchTool'\")\n ).fetchone()\n\n if not search_tool:\n raise ValueError(\n \"SearchTool not found in database. Ensure tools migration has run first.\"\n )\n\n image_gen_tool = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE in_code_tool_id = 'ImageGenerationTool'\")\n ).fetchone()\n\n if not image_gen_tool:\n raise ValueError(\n \"ImageGenerationTool not found in database. Ensure tools migration has run first.\"\n )\n\n # WebSearchTool is optional - may not be configured\n web_search_tool = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE in_code_tool_id = 'WebSearchTool'\")\n ).fetchone()\n\n # Clear existing tool associations for persona 0\n conn.execute(sa.text(\"DELETE FROM persona__tool WHERE persona_id = 0\"))\n\n # Add tools to the unified assistant\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona__tool (persona_id, tool_id)\n VALUES (0, :tool_id)\n ON CONFLICT DO NOTHING\n \"\"\"\n ),\n {\"tool_id\": search_tool[0]},\n )\n\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona__tool (persona_id, tool_id)\n VALUES (0, :tool_id)\n ON CONFLICT DO NOTHING\n \"\"\"\n ),\n {\"tool_id\": image_gen_tool[0]},\n )\n\n if web_search_tool:\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona__tool (persona_id, tool_id)\n VALUES (0, :tool_id)\n ON CONFLICT DO NOTHING\n \"\"\"\n ),\n {\"tool_id\": web_search_tool[0]},\n )\n\n # Step 4: Migrate existing chat sessions from all builtin assistants to unified assistant\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE chat_session\n SET persona_id = 0\n WHERE persona_id IN (\n SELECT id FROM persona WHERE builtin_persona = true AND id != 0\n )\n \"\"\"\n )\n )\n\n # Step 5: Migrate user preferences - remove references to all builtin assistants\n # First, get all builtin assistant IDs (except 0)\n builtin_assistants_result = conn.execute(\n sa.text(\n \"\"\"\n SELECT id FROM persona\n WHERE builtin_persona = true AND id != 0\n \"\"\"\n )\n ).fetchall()\n builtin_assistant_ids = [row[0] for row in builtin_assistants_result]\n\n # Get all users with preferences\n users_result = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, chosen_assistants, visible_assistants,\n hidden_assistants, pinned_assistants\n FROM \"user\"\n \"\"\"\n )\n ).fetchall()\n\n for user_row in users_result:\n user = UserRow(*user_row)\n user_id: UUID = user.id\n updates: dict[str, Any] = {}\n\n # Remove all builtin assistants from chosen_assistants\n if user.chosen_assistants:\n new_chosen: list[int] = [\n assistant_id\n for assistant_id in user.chosen_assistants\n if assistant_id not in builtin_assistant_ids\n ]\n if new_chosen != user.chosen_assistants:\n updates[\"chosen_assistants\"] = json.dumps(new_chosen)\n\n # Remove all builtin assistants from visible_assistants\n if user.visible_assistants:\n new_visible: list[int] = [\n assistant_id\n for assistant_id in user.visible_assistants\n if assistant_id not in builtin_assistant_ids\n ]\n if new_visible != user.visible_assistants:\n updates[\"visible_assistants\"] = json.dumps(new_visible)\n\n # Add all builtin assistants to hidden_assistants\n if user.hidden_assistants:\n new_hidden: list[int] = list(user.hidden_assistants)\n for old_id in builtin_assistant_ids:\n if old_id not in new_hidden:\n new_hidden.append(old_id)\n if new_hidden != user.hidden_assistants:\n updates[\"hidden_assistants\"] = json.dumps(new_hidden)\n else:\n updates[\"hidden_assistants\"] = json.dumps(builtin_assistant_ids)\n\n # Remove all builtin assistants from pinned_assistants\n if user.pinned_assistants:\n new_pinned: list[int] = [\n assistant_id\n for assistant_id in user.pinned_assistants\n if assistant_id not in builtin_assistant_ids\n ]\n if new_pinned != user.pinned_assistants:\n updates[\"pinned_assistants\"] = json.dumps(new_pinned)\n\n # Apply updates if any\n if updates:\n set_clause = \", \".join([f\"{k} = :{k}\" for k in updates.keys()])\n updates[\"user_id\"] = str(user_id) # Convert UUID to string for SQL\n conn.execute(\n sa.text(f'UPDATE \"user\" SET {set_clause} WHERE id = :user_id'),\n updates,\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n\n # Only restore General (ID -1) and Art (ID -3) assistants\n # Step 1: Keep Search assistant (ID 0) as default but restore original state\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET is_default_persona = true,\n is_visible = true,\n deleted = false\n WHERE id = 0\n \"\"\"\n )\n )\n\n # Step 2: Restore General assistant (ID -1)\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET deleted = false,\n is_visible = true,\n is_default_persona = true\n WHERE id = :general_assistant_id\n \"\"\"\n ),\n {\"general_assistant_id\": GENERAL_ASSISTANT_ID},\n )\n\n # Step 3: Restore Art assistant (ID -3)\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET deleted = false,\n is_visible = true,\n is_default_persona = true\n WHERE id = :art_assistant_id\n \"\"\"\n ),\n {\"art_assistant_id\": ART_ASSISTANT_ID},\n )\n\n # Note: We don't restore the original tool associations, names, or descriptions\n # as those would require more complex logic to determine original state.\n # We also cannot restore original chat session persona_ids as we don't\n # have the original mappings.\n # Other builtin assistants remain deleted as per the requirement.\n" }, { "path": "backend/alembic/versions/50b683a8295c_add_additional_retrieval_controls_to_.py", "content": "\"\"\"Add additional retrieval controls to Persona\n\nRevision ID: 50b683a8295c\nRevises: 7da0ae5ad583\nCreate Date: 2023-11-27 17:23:29.668422\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"50b683a8295c\"\ndown_revision = \"7da0ae5ad583\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"persona\", sa.Column(\"num_chunks\", sa.Integer(), nullable=True))\n op.add_column(\n \"persona\",\n sa.Column(\"apply_llm_relevance_filter\", sa.Boolean(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"apply_llm_relevance_filter\")\n op.drop_column(\"persona\", \"num_chunks\")\n" }, { "path": "backend/alembic/versions/52a219fb5233_add_last_synced_and_last_modified_to_document_table.py", "content": "\"\"\"Add last synced and last modified to document table\n\nRevision ID: 52a219fb5233\nRevises: f7e58d357687\nCreate Date: 2024-08-28 17:40:46.077470\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.sql import func\n\n# revision identifiers, used by Alembic.\nrevision = \"52a219fb5233\"\ndown_revision = \"f7e58d357687\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # last modified represents the last time anything needing syncing to vespa changed\n # including row metadata and the document itself. This obviously does not include\n # the last_synced column.\n op.add_column(\n \"document\",\n sa.Column(\n \"last_modified\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=func.now(),\n ),\n )\n\n # last synced represents the last time this document was synced to Vespa\n op.add_column(\n \"document\",\n sa.Column(\"last_synced\", sa.DateTime(timezone=True), nullable=True),\n )\n\n # Set last_synced to the same value as last_modified for existing rows\n op.execute(\n \"\"\"\n UPDATE document\n SET last_synced = last_modified\n \"\"\"\n )\n\n op.create_index(\n op.f(\"ix_document_last_modified\"),\n \"document\",\n [\"last_modified\"],\n unique=False,\n )\n\n op.create_index(\n op.f(\"ix_document_last_synced\"),\n \"document\",\n [\"last_synced\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(op.f(\"ix_document_last_synced\"), table_name=\"document\")\n op.drop_index(op.f(\"ix_document_last_modified\"), table_name=\"document\")\n op.drop_column(\"document\", \"last_synced\")\n op.drop_column(\"document\", \"last_modified\")\n" }, { "path": "backend/alembic/versions/54a74a0417fc_danswerbot_onyxbot.py", "content": "\"\"\"danswerbot -> onyxbot\n\nRevision ID: 54a74a0417fc\nRevises: 94dc3d0236f8\nCreate Date: 2024-12-11 18:05:05.490737\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"54a74a0417fc\"\ndown_revision = \"94dc3d0236f8\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"chat_session\", \"danswerbot_flow\", new_column_name=\"onyxbot_flow\")\n\n\ndef downgrade() -> None:\n op.alter_column(\"chat_session\", \"onyxbot_flow\", new_column_name=\"danswerbot_flow\")\n" }, { "path": "backend/alembic/versions/55546a7967ee_assistant_rework.py", "content": "\"\"\"assistant_rework\n\nRevision ID: 55546a7967ee\nRevises: 61ff3651add4\nCreate Date: 2024-09-18 17:00:23.755399\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"55546a7967ee\"\ndown_revision = \"61ff3651add4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Reworking persona and user tables for new assistant features\n # keep track of user's chosen assistants separate from their `ordering`\n op.add_column(\"persona\", sa.Column(\"builtin_persona\", sa.Boolean(), nullable=True))\n op.execute(\"UPDATE persona SET builtin_persona = default_persona\")\n op.alter_column(\"persona\", \"builtin_persona\", nullable=False)\n op.drop_index(\"_default_persona_name_idx\", table_name=\"persona\")\n op.create_index(\n \"_builtin_persona_name_idx\",\n \"persona\",\n [\"name\"],\n unique=True,\n postgresql_where=sa.text(\"builtin_persona = true\"),\n )\n\n op.add_column(\n \"user\", sa.Column(\"visible_assistants\", postgresql.JSONB(), nullable=True)\n )\n op.add_column(\n \"user\", sa.Column(\"hidden_assistants\", postgresql.JSONB(), nullable=True)\n )\n op.execute(\n \"UPDATE \\\"user\\\" SET visible_assistants = '[]'::jsonb, hidden_assistants = '[]'::jsonb\"\n )\n op.alter_column(\n \"user\",\n \"visible_assistants\",\n nullable=False,\n server_default=sa.text(\"'[]'::jsonb\"),\n )\n op.alter_column(\n \"user\",\n \"hidden_assistants\",\n nullable=False,\n server_default=sa.text(\"'[]'::jsonb\"),\n )\n op.drop_column(\"persona\", \"default_persona\")\n op.add_column(\n \"persona\", sa.Column(\"is_default_persona\", sa.Boolean(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n # Reverting changes made in upgrade\n op.drop_column(\"user\", \"hidden_assistants\")\n op.drop_column(\"user\", \"visible_assistants\")\n op.drop_index(\"_builtin_persona_name_idx\", table_name=\"persona\")\n\n op.drop_column(\"persona\", \"is_default_persona\")\n op.add_column(\"persona\", sa.Column(\"default_persona\", sa.Boolean(), nullable=True))\n op.execute(\"UPDATE persona SET default_persona = builtin_persona\")\n op.alter_column(\"persona\", \"default_persona\", nullable=False)\n op.drop_column(\"persona\", \"builtin_persona\")\n op.create_index(\n \"_default_persona_name_idx\",\n \"persona\",\n [\"name\"],\n unique=True,\n postgresql_where=sa.text(\"default_persona = true\"),\n )\n" }, { "path": "backend/alembic/versions/570282d33c49_track_onyxbot_explicitly.py", "content": "\"\"\"Track Onyxbot Explicitly\n\nRevision ID: 570282d33c49\nRevises: 7547d982db8f\nCreate Date: 2024-05-04 17:49:28.568109\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"570282d33c49\"\ndown_revision = \"7547d982db8f\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_session\", sa.Column(\"danswerbot_flow\", sa.Boolean(), nullable=True)\n )\n op.execute(\"UPDATE chat_session SET danswerbot_flow = one_shot\")\n op.alter_column(\"chat_session\", \"danswerbot_flow\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_session\", \"danswerbot_flow\")\n" }, { "path": "backend/alembic/versions/57122d037335_add_python_tool_on_default.py", "content": "\"\"\"add python tool on default\n\nRevision ID: 57122d037335\nRevises: c0c937d5c9e5\nCreate Date: 2026-02-27 10:10:40.124925\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"57122d037335\"\ndown_revision = \"c0c937d5c9e5\"\nbranch_labels = None\ndepends_on = None\n\n\nPYTHON_TOOL_NAME = \"python\"\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Look up the PythonTool id\n result = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE name = :name\"),\n {\"name\": PYTHON_TOOL_NAME},\n ).fetchone()\n\n if not result:\n return\n\n tool_id = result[0]\n\n # Attach to the default persona (id=0) if not already attached\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona__tool (persona_id, tool_id)\n VALUES (0, :tool_id)\n ON CONFLICT DO NOTHING\n \"\"\"\n ),\n {\"tool_id\": tool_id},\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n\n result = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE name = :name\"),\n {\"name\": PYTHON_TOOL_NAME},\n ).fetchone()\n\n if not result:\n return\n\n conn.execute(\n sa.text(\n \"\"\"\n DELETE FROM persona__tool\n WHERE persona_id = 0 AND tool_id = :tool_id\n \"\"\"\n ),\n {\"tool_id\": result[0]},\n )\n" }, { "path": "backend/alembic/versions/57b53544726e_add_document_set_tables.py", "content": "\"\"\"Add document set tables\n\nRevision ID: 57b53544726e\nRevises: 800f48024ae9\nCreate Date: 2023-09-20 16:59:39.097177\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"57b53544726e\"\ndown_revision = \"800f48024ae9\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"document_set\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"description\", sa.String(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"is_up_to_date\", sa.Boolean(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n op.create_table(\n \"document_set__connector_credential_pair\",\n sa.Column(\"document_set_id\", sa.Integer(), nullable=False),\n sa.Column(\"connector_credential_pair_id\", sa.Integer(), nullable=False),\n sa.Column(\"is_current\", sa.Boolean(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"connector_credential_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"document_set_id\"],\n [\"document_set.id\"],\n ),\n sa.PrimaryKeyConstraint(\n \"document_set_id\", \"connector_credential_pair_id\", \"is_current\"\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"document_set__connector_credential_pair\")\n op.drop_table(\"document_set\")\n" }, { "path": "backend/alembic/versions/5809c0787398_add_chat_sessions.py", "content": "\"\"\"Add Chat Sessions\n\nRevision ID: 5809c0787398\nRevises: d929f0c1c6af\nCreate Date: 2023-09-04 15:29:44.002164\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5809c0787398\"\ndown_revision = \"d929f0c1c6af\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"chat_session\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"description\", sa.Text(), nullable=False),\n sa.Column(\"deleted\", sa.Boolean(), nullable=False),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"chat_message\",\n sa.Column(\"chat_session_id\", sa.Integer(), nullable=False),\n sa.Column(\"message_number\", sa.Integer(), nullable=False),\n sa.Column(\"edit_number\", sa.Integer(), nullable=False),\n sa.Column(\"parent_edit_number\", sa.Integer(), nullable=True),\n sa.Column(\"latest\", sa.Boolean(), nullable=False),\n sa.Column(\"message\", sa.Text(), nullable=False),\n sa.Column(\n \"message_type\",\n sa.Enum(\n \"SYSTEM\",\n \"USER\",\n \"ASSISTANT\",\n \"DANSWER\",\n name=\"messagetype\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\n \"time_sent\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"chat_session_id\"],\n [\"chat_session.id\"],\n ),\n sa.PrimaryKeyConstraint(\"chat_session_id\", \"message_number\", \"edit_number\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"chat_message\")\n op.drop_table(\"chat_session\")\n" }, { "path": "backend/alembic/versions/58c50ef19f08_add_stale_column_to_user__external_user_.py", "content": "\"\"\"add stale column to external user group tables\n\nRevision ID: 58c50ef19f08\nRevises: 7b9b952abdf6\nCreate Date: 2025-06-25 14:08:14.162380\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"58c50ef19f08\"\ndown_revision = \"7b9b952abdf6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add the stale column with default value False to user__external_user_group_id\n op.add_column(\n \"user__external_user_group_id\",\n sa.Column(\"stale\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n\n # Create index for efficient querying of stale rows by cc_pair_id\n op.create_index(\n \"ix_user__external_user_group_id_cc_pair_id_stale\",\n \"user__external_user_group_id\",\n [\"cc_pair_id\", \"stale\"],\n unique=False,\n )\n\n # Create index for efficient querying of all stale rows\n op.create_index(\n \"ix_user__external_user_group_id_stale\",\n \"user__external_user_group_id\",\n [\"stale\"],\n unique=False,\n )\n\n # Add the stale column with default value False to public_external_user_group\n op.add_column(\n \"public_external_user_group\",\n sa.Column(\"stale\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n\n # Create index for efficient querying of stale rows by cc_pair_id\n op.create_index(\n \"ix_public_external_user_group_cc_pair_id_stale\",\n \"public_external_user_group\",\n [\"cc_pair_id\", \"stale\"],\n unique=False,\n )\n\n # Create index for efficient querying of all stale rows\n op.create_index(\n \"ix_public_external_user_group_stale\",\n \"public_external_user_group\",\n [\"stale\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n # Drop the indices for public_external_user_group first\n op.drop_index(\n \"ix_public_external_user_group_stale\", table_name=\"public_external_user_group\"\n )\n op.drop_index(\n \"ix_public_external_user_group_cc_pair_id_stale\",\n table_name=\"public_external_user_group\",\n )\n\n # Drop the stale column from public_external_user_group\n op.drop_column(\"public_external_user_group\", \"stale\")\n\n # Drop the indices for user__external_user_group_id\n op.drop_index(\n \"ix_user__external_user_group_id_stale\",\n table_name=\"user__external_user_group_id\",\n )\n op.drop_index(\n \"ix_user__external_user_group_id_cc_pair_id_stale\",\n table_name=\"user__external_user_group_id\",\n )\n\n # Drop the stale column from user__external_user_group_id\n op.drop_column(\"user__external_user_group_id\", \"stale\")\n" }, { "path": "backend/alembic/versions/5ae8240accb3_add_research_agent_database_tables_and_.py", "content": "\"\"\"add research agent database tables and chat message research fields\n\nRevision ID: 5ae8240accb3\nRevises: b558f51620b4\nCreate Date: 2025-08-06 14:29:24.691388\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5ae8240accb3\"\ndown_revision = \"b558f51620b4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add research_type and research_plan columns to chat_message table\n op.add_column(\n \"chat_message\",\n sa.Column(\"research_type\", sa.String(), nullable=True),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"research_plan\", postgresql.JSONB(), nullable=True),\n )\n\n # Create research_agent_iteration table\n op.create_table(\n \"research_agent_iteration\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\n \"primary_question_id\",\n sa.Integer(),\n sa.ForeignKey(\"chat_message.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\"iteration_nr\", sa.Integer(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\"purpose\", sa.String(), nullable=True),\n sa.Column(\"reasoning\", sa.String(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\n \"primary_question_id\",\n \"iteration_nr\",\n name=\"_research_agent_iteration_unique_constraint\",\n ),\n )\n\n # Create research_agent_iteration_sub_step table\n op.create_table(\n \"research_agent_iteration_sub_step\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\n \"primary_question_id\",\n sa.Integer(),\n sa.ForeignKey(\"chat_message.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\n \"parent_question_id\",\n sa.Integer(),\n sa.ForeignKey(\"research_agent_iteration_sub_step.id\", ondelete=\"CASCADE\"),\n nullable=True,\n ),\n sa.Column(\"iteration_nr\", sa.Integer(), nullable=False),\n sa.Column(\"iteration_sub_step_nr\", sa.Integer(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\"sub_step_instructions\", sa.String(), nullable=True),\n sa.Column(\n \"sub_step_tool_id\",\n sa.Integer(),\n sa.ForeignKey(\"tool.id\"),\n nullable=True,\n ),\n sa.Column(\"reasoning\", sa.String(), nullable=True),\n sa.Column(\"sub_answer\", sa.String(), nullable=True),\n sa.Column(\"cited_doc_results\", postgresql.JSONB(), nullable=True),\n sa.Column(\"claims\", postgresql.JSONB(), nullable=True),\n sa.Column(\"generated_images\", postgresql.JSONB(), nullable=True),\n sa.Column(\"additional_data\", postgresql.JSONB(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.ForeignKeyConstraint(\n [\"primary_question_id\", \"iteration_nr\"],\n [\n \"research_agent_iteration.primary_question_id\",\n \"research_agent_iteration.iteration_nr\",\n ],\n ondelete=\"CASCADE\",\n ),\n )\n\n\ndef downgrade() -> None:\n # Drop tables in reverse order\n op.drop_table(\"research_agent_iteration_sub_step\")\n op.drop_table(\"research_agent_iteration\")\n\n # Remove columns from chat_message table\n op.drop_column(\"chat_message\", \"research_plan\")\n op.drop_column(\"chat_message\", \"research_type\")\n" }, { "path": "backend/alembic/versions/5b29123cd710_nullable_search_settings_for_historic_.py", "content": "\"\"\"nullable search settings for historic index attempts\n\nRevision ID: 5b29123cd710\nRevises: 949b4a92a401\nCreate Date: 2024-10-30 19:37:59.630704\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5b29123cd710\"\ndown_revision = \"949b4a92a401\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Drop the existing foreign key constraint\n op.drop_constraint(\n \"fk_index_attempt_search_settings\", \"index_attempt\", type_=\"foreignkey\"\n )\n\n # Modify the column to be nullable\n op.alter_column(\n \"index_attempt\", \"search_settings_id\", existing_type=sa.INTEGER(), nullable=True\n )\n\n # Add back the foreign key with ON DELETE SET NULL\n op.create_foreign_key(\n \"fk_index_attempt_search_settings\",\n \"index_attempt\",\n \"search_settings\",\n [\"search_settings_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n\n\ndef downgrade() -> None:\n # Warning: This will delete all index attempts that don't have search settings\n op.execute(\n \"\"\"\n DELETE FROM index_attempt\n WHERE search_settings_id IS NULL\n \"\"\"\n )\n\n # Drop foreign key constraint\n op.drop_constraint(\n \"fk_index_attempt_search_settings\", \"index_attempt\", type_=\"foreignkey\"\n )\n\n # Modify the column to be not nullable\n op.alter_column(\n \"index_attempt\",\n \"search_settings_id\",\n existing_type=sa.INTEGER(),\n nullable=False,\n )\n\n # Add back the foreign key without ON DELETE SET NULL\n op.create_foreign_key(\n \"fk_index_attempt_search_settings\",\n \"index_attempt\",\n \"search_settings\",\n [\"search_settings_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/5c3dca366b35_backend_driven_notification_details.py", "content": "\"\"\"backend driven notification details\n\nRevision ID: 5c3dca366b35\nRevises: 9087b548dd69\nCreate Date: 2026-01-06 16:03:11.413724\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5c3dca366b35\"\ndown_revision = \"9087b548dd69\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"notification\",\n sa.Column(\n \"title\", sa.String(), nullable=False, server_default=\"New Notification\"\n ),\n )\n op.add_column(\n \"notification\",\n sa.Column(\"description\", sa.String(), nullable=True, server_default=\"\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"notification\", \"title\")\n op.drop_column(\"notification\", \"description\")\n" }, { "path": "backend/alembic/versions/5c448911b12f_add_content_type_to_userfile.py", "content": "\"\"\"Add content type to UserFile\n\nRevision ID: 5c448911b12f\nRevises: 47a07e1a38f1\nCreate Date: 2025-04-25 16:59:48.182672\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"5c448911b12f\"\ndown_revision = \"47a07e1a38f1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"user_file\", sa.Column(\"content_type\", sa.String(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"user_file\", \"content_type\")\n" }, { "path": "backend/alembic/versions/5c7fdadae813_match_any_keywords_flag_for_standard_.py", "content": "\"\"\"match_any_keywords flag for standard answers\n\nRevision ID: 5c7fdadae813\nRevises: efb35676026c\nCreate Date: 2024-09-13 18:52:59.256478\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"5c7fdadae813\"\ndown_revision = \"efb35676026c\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # ### commands auto generated by Alembic - please adjust! ###\n op.add_column(\n \"standard_answer\",\n sa.Column(\n \"match_any_keywords\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n )\n # ### end Alembic commands ###\n\n\ndef downgrade() -> None:\n # ### commands auto generated by Alembic - please adjust! ###\n op.drop_column(\"standard_answer\", \"match_any_keywords\")\n # ### end Alembic commands ###\n" }, { "path": "backend/alembic/versions/5d12a446f5c0_add_api_version_and_deployment_name_to_.py", "content": "\"\"\"add api_version and deployment_name to search settings\n\nRevision ID: 5d12a446f5c0\nRevises: e4334d5b33ba\nCreate Date: 2024-10-08 15:56:07.975636\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5d12a446f5c0\"\ndown_revision = \"e4334d5b33ba\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"embedding_provider\", sa.Column(\"api_version\", sa.String(), nullable=True)\n )\n op.add_column(\n \"embedding_provider\", sa.Column(\"deployment_name\", sa.String(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"embedding_provider\", \"deployment_name\")\n op.drop_column(\"embedding_provider\", \"api_version\")\n" }, { "path": "backend/alembic/versions/5e1c073d48a3_add_personal_access_token_table.py", "content": "\"\"\"add_personal_access_token_table\n\nRevision ID: 5e1c073d48a3\nRevises: 09995b8811eb\nCreate Date: 2025-10-30 17:30:24.308521\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5e1c073d48a3\"\ndown_revision = \"09995b8811eb\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create personal_access_token table\n op.create_table(\n \"personal_access_token\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"hashed_token\", sa.String(length=64), nullable=False),\n sa.Column(\"token_display\", sa.String(), nullable=False),\n sa.Column(\n \"user_id\",\n postgresql.UUID(as_uuid=True),\n nullable=False,\n ),\n sa.Column(\n \"expires_at\",\n sa.DateTime(timezone=True),\n nullable=True,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"last_used_at\",\n sa.DateTime(timezone=True),\n nullable=True,\n ),\n sa.Column(\n \"is_revoked\",\n sa.Boolean(),\n server_default=sa.text(\"false\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"hashed_token\"),\n )\n\n # Create indexes\n op.create_index(\n \"ix_personal_access_token_expires_at\",\n \"personal_access_token\",\n [\"expires_at\"],\n unique=False,\n )\n op.create_index(\n \"ix_pat_user_created\",\n \"personal_access_token\",\n [\"user_id\", sa.text(\"created_at DESC\")],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n # Drop indexes first\n op.drop_index(\"ix_pat_user_created\", table_name=\"personal_access_token\")\n op.drop_index(\n \"ix_personal_access_token_expires_at\", table_name=\"personal_access_token\"\n )\n\n # Drop table\n op.drop_table(\"personal_access_token\")\n" }, { "path": "backend/alembic/versions/5e6f7a8b9c0d_update_default_persona_prompt.py", "content": "\"\"\"update_default_persona_prompt\n\nRevision ID: 5e6f7a8b9c0d\nRevises: 4f8a2b3c1d9e\nCreate Date: 2025-11-30 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5e6f7a8b9c0d\"\ndown_revision = \"4f8a2b3c1d9e\"\nbranch_labels = None\ndepends_on = None\n\n\nDEFAULT_PERSONA_ID = 0\n\n# ruff: noqa: E501, W605 start\nDEFAULT_SYSTEM_PROMPT = \"\"\"\nYou are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent, ask clarifying questions when needed, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. Always prioritize being truthful, nuanced, insightful, and efficient.\n\nThe current date is [[CURRENT_DATETIME]].{citation_reminder_or_empty}\n\n# Response Style\nYou use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make your responses more readable and engaging.\nYou use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, symbols, etc.: '$$\\\\n[expression]\\\\n$$' for standalone cases and '\\\\( [expression] \\\\)' when inline.\nFor code you prefer to use Markdown and specify the language.\nYou can use horizontal rules (---) to separate sections of your responses.\nYou can use Markdown tables to format your responses for data, lists, and other structured information.\n\"\"\".lstrip()\n# ruff: noqa: E501, W605 end\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET system_prompt = :system_prompt\n WHERE id = :persona_id\n \"\"\"\n ),\n {\"system_prompt\": DEFAULT_SYSTEM_PROMPT, \"persona_id\": DEFAULT_PERSONA_ID},\n )\n\n\ndef downgrade() -> None:\n # We don't revert the system prompt on downgrade since we don't know\n # what the previous value was. The new prompt is a reasonable default.\n pass\n" }, { "path": "backend/alembic/versions/5e84129c8be3_add_docs_indexed_column_to_index_.py", "content": "\"\"\"Add docs_indexed_column + time_started to index_attempt table\n\nRevision ID: 5e84129c8be3\nRevises: e6a4bbc13fe4\nCreate Date: 2023-08-10 21:43:09.069523\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"5e84129c8be3\"\ndown_revision = \"e6a4bbc13fe4\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"index_attempt\",\n sa.Column(\"num_docs_indexed\", sa.Integer()),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"time_started\",\n sa.DateTime(timezone=True),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"index_attempt\", \"time_started\")\n op.drop_column(\"index_attempt\", \"num_docs_indexed\")\n" }, { "path": "backend/alembic/versions/5f4b8568a221_add_removed_documents_to_index_attempt.py", "content": "\"\"\"add removed documents to index_attempt\n\nRevision ID: 5f4b8568a221\nRevises: dbaa756c2ccf\nCreate Date: 2024-02-16 15:02:03.319907\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"5f4b8568a221\"\ndown_revision = \"8987770549c0\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"index_attempt\",\n sa.Column(\"docs_removed_from_index\", sa.Integer()),\n )\n op.execute(\"UPDATE index_attempt SET docs_removed_from_index = 0\")\n\n\ndef downgrade() -> None:\n op.drop_column(\"index_attempt\", \"docs_removed_from_index\")\n" }, { "path": "backend/alembic/versions/5fc1f54cc252_hybrid_enum.py", "content": "\"\"\"hybrid-enum\n\nRevision ID: 5fc1f54cc252\nRevises: 1d6ad76d1f37\nCreate Date: 2024-08-06 15:35:40.278485\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"5fc1f54cc252\"\ndown_revision = \"1d6ad76d1f37\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"persona\", \"search_type\")\n\n\ndef downgrade() -> None:\n op.add_column(\"persona\", sa.Column(\"search_type\", sa.String(), nullable=True))\n op.execute(\"UPDATE persona SET search_type = 'SEMANTIC'\")\n op.alter_column(\"persona\", \"search_type\", nullable=False)\n" }, { "path": "backend/alembic/versions/61ff3651add4_add_permission_syncing.py", "content": "\"\"\"Add Permission Syncing\n\nRevision ID: 61ff3651add4\nRevises: 1b8206b29c5d\nCreate Date: 2024-09-05 13:57:11.770413\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"61ff3651add4\"\ndown_revision = \"1b8206b29c5d\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Admin user who set up connectors will lose access to the docs temporarily\n # only way currently to give back access is to rerun from beginning\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"access_type\",\n sa.String(),\n nullable=True,\n ),\n )\n op.execute(\n \"UPDATE connector_credential_pair SET access_type = 'PUBLIC' WHERE is_public = true\"\n )\n op.execute(\n \"UPDATE connector_credential_pair SET access_type = 'PRIVATE' WHERE is_public = false\"\n )\n op.alter_column(\"connector_credential_pair\", \"access_type\", nullable=False)\n\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"auto_sync_options\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"last_time_perm_sync\", sa.DateTime(timezone=True), nullable=True),\n )\n op.drop_column(\"connector_credential_pair\", \"is_public\")\n\n op.add_column(\n \"document\",\n sa.Column(\"external_user_emails\", postgresql.ARRAY(sa.String()), nullable=True),\n )\n op.add_column(\n \"document\",\n sa.Column(\n \"external_user_group_ids\", postgresql.ARRAY(sa.String()), nullable=True\n ),\n )\n op.add_column(\n \"document\",\n sa.Column(\"is_public\", sa.Boolean(), nullable=True),\n )\n\n op.create_table(\n \"user__external_user_group_id\",\n sa.Column(\n \"user_id\", fastapi_users_db_sqlalchemy.generics.GUID(), nullable=False\n ),\n sa.Column(\"external_user_group_id\", sa.String(), nullable=False),\n sa.Column(\"cc_pair_id\", sa.Integer(), nullable=False),\n sa.PrimaryKeyConstraint(\"user_id\"),\n )\n\n op.drop_column(\"external_permission\", \"user_id\")\n op.drop_column(\"email_to_external_user_cache\", \"user_id\")\n op.drop_table(\"permission_sync_run\")\n op.drop_table(\"external_permission\")\n op.drop_table(\"email_to_external_user_cache\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"is_public\", sa.BOOLEAN(), nullable=True),\n )\n op.execute(\n \"UPDATE connector_credential_pair SET is_public = (access_type = 'PUBLIC')\"\n )\n op.alter_column(\"connector_credential_pair\", \"is_public\", nullable=False)\n\n op.drop_column(\"connector_credential_pair\", \"auto_sync_options\")\n op.drop_column(\"connector_credential_pair\", \"access_type\")\n op.drop_column(\"connector_credential_pair\", \"last_time_perm_sync\")\n op.drop_column(\"document\", \"external_user_emails\")\n op.drop_column(\"document\", \"external_user_group_ids\")\n op.drop_column(\"document\", \"is_public\")\n\n op.drop_table(\"user__external_user_group_id\")\n\n # Drop the enum type at the end of the downgrade\n op.create_table(\n \"permission_sync_run\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"source_type\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\"update_type\", sa.String(), nullable=False),\n sa.Column(\"cc_pair_id\", sa.Integer(), nullable=True),\n sa.Column(\n \"status\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\"error_msg\", sa.Text(), nullable=True),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"cc_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"external_permission\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", sa.UUID(), nullable=True),\n sa.Column(\"user_email\", sa.String(), nullable=False),\n sa.Column(\n \"source_type\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\"external_permission_group\", sa.String(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"email_to_external_user_cache\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"external_user_id\", sa.String(), nullable=False),\n sa.Column(\"user_id\", sa.UUID(), nullable=True),\n sa.Column(\"user_email\", sa.String(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n" }, { "path": "backend/alembic/versions/62c3a055a141_add_file_names_to_file_connector_config.py", "content": "\"\"\"add file names to file connector config\n\nRevision ID: 62c3a055a141\nRevises: 3fc5d75723b3\nCreate Date: 2025-07-30 17:01:24.417551\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport json\nimport os\nimport logging\n\n\n# revision identifiers, used by Alembic.\nrevision = \"62c3a055a141\"\ndown_revision = \"3fc5d75723b3\"\nbranch_labels = None\ndepends_on = None\n\nSKIP_FILE_NAME_MIGRATION = (\n os.environ.get(\"SKIP_FILE_NAME_MIGRATION\", \"true\").lower() == \"true\"\n)\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n\ndef upgrade() -> None:\n if SKIP_FILE_NAME_MIGRATION:\n logger.info(\n \"Skipping file name migration. Hint: set SKIP_FILE_NAME_MIGRATION=false to run this migration\"\n )\n return\n logger.info(\"Running file name migration\")\n # Get connection\n conn = op.get_bind()\n\n # Get all FILE connectors with their configs\n file_connectors = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'FILE'\n \"\"\"\n )\n ).fetchall()\n\n for connector_id, config in file_connectors:\n # Parse config if it's a string\n if isinstance(config, str):\n config = json.loads(config)\n\n # Get file_locations list\n file_locations = config.get(\"file_locations\", [])\n\n # Get display names for each file_id\n file_names = []\n for file_id in file_locations:\n result = conn.execute(\n sa.text(\n \"\"\"\n SELECT display_name\n FROM file_record\n WHERE file_id = :file_id\n \"\"\"\n ),\n {\"file_id\": file_id},\n ).fetchone()\n\n if result:\n file_names.append(result[0])\n else:\n file_names.append(file_id) # Should not happen\n\n # Add file_names to config\n new_config = dict(config)\n new_config[\"file_names\"] = file_names\n\n # Update the connector\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = :new_config\n WHERE id = :connector_id\n \"\"\"\n ),\n {\"connector_id\": connector_id, \"new_config\": json.dumps(new_config)},\n )\n\n\ndef downgrade() -> None:\n # Get connection\n conn = op.get_bind()\n\n # Remove file_names from all FILE connectors\n file_connectors = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'FILE'\n \"\"\"\n )\n ).fetchall()\n\n for connector_id, config in file_connectors:\n # Parse config if it's a string\n if isinstance(config, str):\n config = json.loads(config)\n\n # Remove file_names if it exists\n if \"file_names\" in config:\n new_config = dict(config)\n del new_config[\"file_names\"]\n\n # Update the connector\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = :new_config\n WHERE id = :connector_id\n \"\"\"\n ),\n {\n \"connector_id\": connector_id,\n \"new_config\": json.dumps(new_config),\n },\n )\n" }, { "path": "backend/alembic/versions/631fd2504136_add_approx_chunk_count_in_vespa_to_.py", "content": "\"\"\"add approx_chunk_count_in_vespa to opensearch tenant migration\n\nRevision ID: 631fd2504136\nRevises: c7f2e1b4a9d3\nCreate Date: 2026-02-18 21:07:52.831215\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"631fd2504136\"\ndown_revision = \"c7f2e1b4a9d3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"approx_chunk_count_in_vespa\",\n sa.Integer(),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"opensearch_tenant_migration_record\", \"approx_chunk_count_in_vespa\")\n" }, { "path": "backend/alembic/versions/6436661d5b65_add_created_at_in_project_userfile.py", "content": "\"\"\"add_created_at_in_project_userfile\n\nRevision ID: 6436661d5b65\nRevises: c7e9f4a3b2d1\nCreate Date: 2025-11-24 11:50:24.536052\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"6436661d5b65\"\ndown_revision = \"c7e9f4a3b2d1\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add created_at column to project__user_file table\n op.add_column(\n \"project__user_file\",\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n )\n # Add composite index on (project_id, created_at DESC)\n op.create_index(\n \"ix_project__user_file_project_id_created_at\",\n \"project__user_file\",\n [\"project_id\", sa.text(\"created_at DESC\")],\n )\n\n\ndef downgrade() -> None:\n # Remove composite index on (project_id, created_at)\n op.drop_index(\n \"ix_project__user_file_project_id_created_at\", table_name=\"project__user_file\"\n )\n # Remove created_at column from project__user_file table\n op.drop_column(\"project__user_file\", \"created_at\")\n" }, { "path": "backend/alembic/versions/643a84a42a33_add_user_configured_names_to_llmprovider.py", "content": "\"\"\"Add user-configured names to LLMProvider\n\nRevision ID: 643a84a42a33\nRevises: 0a98909f2757\nCreate Date: 2024-05-07 14:54:55.493100\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"643a84a42a33\"\ndown_revision = \"0a98909f2757\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"llm_provider\", sa.Column(\"provider\", sa.String(), nullable=True))\n # move \"name\" -> \"provider\" to match the new schema\n op.execute(\"UPDATE llm_provider SET provider = name\")\n # pretty up display name\n op.execute(\"UPDATE llm_provider SET name = 'OpenAI' WHERE name = 'openai'\")\n op.execute(\"UPDATE llm_provider SET name = 'Anthropic' WHERE name = 'anthropic'\")\n op.execute(\"UPDATE llm_provider SET name = 'Azure OpenAI' WHERE name = 'azure'\")\n op.execute(\"UPDATE llm_provider SET name = 'AWS Bedrock' WHERE name = 'bedrock'\")\n\n # update personas to use the new provider names\n op.execute(\n \"UPDATE persona SET llm_model_provider_override = 'OpenAI' WHERE llm_model_provider_override = 'openai'\"\n )\n op.execute(\n \"UPDATE persona SET llm_model_provider_override = 'Anthropic' WHERE llm_model_provider_override = 'anthropic'\"\n )\n op.execute(\n \"UPDATE persona SET llm_model_provider_override = 'Azure OpenAI' WHERE llm_model_provider_override = 'azure'\"\n )\n op.execute(\n \"UPDATE persona SET llm_model_provider_override = 'AWS Bedrock' WHERE llm_model_provider_override = 'bedrock'\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\"UPDATE llm_provider SET name = provider\")\n op.drop_column(\"llm_provider\", \"provider\")\n" }, { "path": "backend/alembic/versions/64bd5677aeb6_add_image_input_support_to_model_config.py", "content": "\"\"\"Add image input support to model config\n\nRevision ID: 64bd5677aeb6\nRevises: b30353be4eec\nCreate Date: 2025-09-28 15:48:12.003612\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"64bd5677aeb6\"\ndown_revision = \"b30353be4eec\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"model_configuration\",\n sa.Column(\"supports_image_input\", sa.Boolean(), nullable=True),\n )\n\n # Seems to be left over from when model visibility was introduced and a nullable field.\n # Set any null is_visible values to False\n connection = op.get_bind()\n connection.execute(\n sa.text(\n \"UPDATE model_configuration SET is_visible = false WHERE is_visible IS NULL\"\n )\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"model_configuration\", \"supports_image_input\")\n" }, { "path": "backend/alembic/versions/65bc6e0f8500_remove_kg_subtype_from_db.py", "content": "\"\"\"remove kg subtype from db\n\nRevision ID: 65bc6e0f8500\nRevises: cec7ec36c505\nCreate Date: 2025-06-13 10:04:27.705976\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"65bc6e0f8500\"\ndown_revision = \"cec7ec36c505\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"kg_entity\", \"entity_class\")\n op.drop_column(\"kg_entity\", \"entity_subtype\")\n op.drop_column(\"kg_entity_extraction_staging\", \"entity_class\")\n op.drop_column(\"kg_entity_extraction_staging\", \"entity_subtype\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"kg_entity_extraction_staging\",\n sa.Column(\"entity_subtype\", sa.String(), nullable=True, index=True),\n )\n op.add_column(\n \"kg_entity_extraction_staging\",\n sa.Column(\"entity_class\", sa.String(), nullable=True, index=True),\n )\n op.add_column(\n \"kg_entity\", sa.Column(\"entity_subtype\", sa.String(), nullable=True, index=True)\n )\n op.add_column(\n \"kg_entity\", sa.Column(\"entity_class\", sa.String(), nullable=True, index=True)\n )\n" }, { "path": "backend/alembic/versions/6756efa39ada_id_uuid_for_chat_session.py", "content": "\"\"\"Migrate chat_session and chat_message tables to use UUID primary keys\n\nRevision ID: 6756efa39ada\nRevises: 5d12a446f5c0\nCreate Date: 2024-10-15 17:47:44.108537\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nrevision = \"6756efa39ada\"\ndown_revision = \"5d12a446f5c0\"\nbranch_labels = None\ndepends_on = None\n\n\"\"\"\nThis script:\n1. Adds UUID columns to chat_session and chat_message\n2. Populates new columns with UUIDs\n3. Updates foreign key relationships\n4. Removes old integer ID columns\n\nNote: Downgrade will assign new integer IDs, not restore original ones.\n\"\"\"\n\n\ndef upgrade() -> None:\n op.execute(\"CREATE EXTENSION IF NOT EXISTS pgcrypto;\")\n\n op.add_column(\n \"chat_session\",\n sa.Column(\n \"new_id\",\n sa.UUID(as_uuid=True),\n server_default=sa.text(\"gen_random_uuid()\"),\n nullable=False,\n ),\n )\n\n op.execute(\"UPDATE chat_session SET new_id = gen_random_uuid();\")\n\n op.add_column(\n \"chat_message\",\n sa.Column(\"new_chat_session_id\", sa.UUID(as_uuid=True), nullable=True),\n )\n\n op.execute(\n \"\"\"\n UPDATE chat_message\n SET new_chat_session_id = cs.new_id\n FROM chat_session cs\n WHERE chat_message.chat_session_id = cs.id;\n \"\"\"\n )\n\n op.drop_constraint(\n \"chat_message_chat_session_id_fkey\", \"chat_message\", type_=\"foreignkey\"\n )\n\n op.drop_column(\"chat_message\", \"chat_session_id\")\n op.alter_column(\n \"chat_message\", \"new_chat_session_id\", new_column_name=\"chat_session_id\"\n )\n\n op.drop_constraint(\"chat_session_pkey\", \"chat_session\", type_=\"primary\")\n op.drop_column(\"chat_session\", \"id\")\n op.alter_column(\"chat_session\", \"new_id\", new_column_name=\"id\")\n\n op.create_primary_key(\"chat_session_pkey\", \"chat_session\", [\"id\"])\n\n op.create_foreign_key(\n \"chat_message_chat_session_id_fkey\",\n \"chat_message\",\n \"chat_session\",\n [\"chat_session_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\n \"chat_message_chat_session_id_fkey\", \"chat_message\", type_=\"foreignkey\"\n )\n\n op.add_column(\n \"chat_session\",\n sa.Column(\"old_id\", sa.Integer, autoincrement=True, nullable=True),\n )\n\n op.execute(\"CREATE SEQUENCE chat_session_old_id_seq OWNED BY chat_session.old_id;\")\n op.execute(\n \"ALTER TABLE chat_session ALTER COLUMN old_id SET DEFAULT nextval('chat_session_old_id_seq');\"\n )\n\n op.execute(\n \"UPDATE chat_session SET old_id = nextval('chat_session_old_id_seq') WHERE old_id IS NULL;\"\n )\n\n op.alter_column(\"chat_session\", \"old_id\", nullable=False)\n\n op.drop_constraint(\"chat_session_pkey\", \"chat_session\", type_=\"primary\")\n op.create_primary_key(\"chat_session_pkey\", \"chat_session\", [\"old_id\"])\n\n op.add_column(\n \"chat_message\",\n sa.Column(\"old_chat_session_id\", sa.Integer, nullable=True),\n )\n\n op.execute(\n \"\"\"\n UPDATE chat_message\n SET old_chat_session_id = cs.old_id\n FROM chat_session cs\n WHERE chat_message.chat_session_id = cs.id;\n \"\"\"\n )\n\n op.drop_column(\"chat_message\", \"chat_session_id\")\n op.alter_column(\n \"chat_message\", \"old_chat_session_id\", new_column_name=\"chat_session_id\"\n )\n\n op.create_foreign_key(\n \"chat_message_chat_session_id_fkey\",\n \"chat_message\",\n \"chat_session\",\n [\"chat_session_id\"],\n [\"old_id\"],\n ondelete=\"CASCADE\",\n )\n\n op.drop_column(\"chat_session\", \"id\")\n op.alter_column(\"chat_session\", \"old_id\", new_column_name=\"id\")\n\n op.alter_column(\n \"chat_session\",\n \"id\",\n type_=sa.Integer(),\n existing_type=sa.Integer(),\n existing_nullable=False,\n existing_server_default=False,\n )\n\n # Rename the sequence\n op.execute(\"ALTER SEQUENCE chat_session_old_id_seq RENAME TO chat_session_id_seq;\")\n\n # Update the default value to use the renamed sequence\n op.alter_column(\n \"chat_session\",\n \"id\",\n server_default=sa.text(\"nextval('chat_session_id_seq'::regclass)\"),\n )\n" }, { "path": "backend/alembic/versions/689433b0d8de_add_hook_and_hook_execution_log_tables.py", "content": "\"\"\"add_hook_and_hook_execution_log_tables\n\nRevision ID: 689433b0d8de\nRevises: 93a2e195e25c\nCreate Date: 2026-03-13 11:25:06.547474\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects.postgresql import UUID as PGUUID\n\n\n# revision identifiers, used by Alembic.\nrevision = \"689433b0d8de\"\ndown_revision = \"93a2e195e25c\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"hook\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\n \"hook_point\",\n sa.Enum(\"document_ingestion\", \"query_processing\", native_enum=False),\n nullable=False,\n ),\n sa.Column(\"endpoint_url\", sa.Text(), nullable=True),\n sa.Column(\"api_key\", sa.LargeBinary(), nullable=True),\n sa.Column(\"is_reachable\", sa.Boolean(), nullable=True),\n sa.Column(\n \"fail_strategy\",\n sa.Enum(\"hard\", \"soft\", native_enum=False),\n nullable=False,\n ),\n sa.Column(\"timeout_seconds\", sa.Float(), nullable=False),\n sa.Column(\n \"is_active\", sa.Boolean(), nullable=False, server_default=sa.text(\"false\")\n ),\n sa.Column(\n \"deleted\", sa.Boolean(), nullable=False, server_default=sa.text(\"false\")\n ),\n sa.Column(\"creator_id\", PGUUID(as_uuid=True), nullable=True),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint([\"creator_id\"], [\"user.id\"], ondelete=\"SET NULL\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_index(\n \"ix_hook_one_non_deleted_per_point\",\n \"hook\",\n [\"hook_point\"],\n unique=True,\n postgresql_where=sa.text(\"deleted = false\"),\n )\n\n op.create_table(\n \"hook_execution_log\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"hook_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"is_success\",\n sa.Boolean(),\n nullable=False,\n ),\n sa.Column(\"error_message\", sa.Text(), nullable=True),\n sa.Column(\"status_code\", sa.Integer(), nullable=True),\n sa.Column(\"duration_ms\", sa.Integer(), nullable=True),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint([\"hook_id\"], [\"hook.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_index(\"ix_hook_execution_log_hook_id\", \"hook_execution_log\", [\"hook_id\"])\n op.create_index(\n \"ix_hook_execution_log_created_at\", \"hook_execution_log\", [\"created_at\"]\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_hook_execution_log_created_at\", table_name=\"hook_execution_log\")\n op.drop_index(\"ix_hook_execution_log_hook_id\", table_name=\"hook_execution_log\")\n op.drop_table(\"hook_execution_log\")\n\n op.drop_index(\"ix_hook_one_non_deleted_per_point\", table_name=\"hook\")\n op.drop_table(\"hook\")\n" }, { "path": "backend/alembic/versions/699221885109_nullify_default_task_prompt.py", "content": "\"\"\"nullify_default_task_prompt\n\nRevision ID: 699221885109\nRevises: 7e490836d179\nCreate Date: 2025-12-30 10:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"699221885109\"\ndown_revision = \"7e490836d179\"\nbranch_labels = None\ndepends_on = None\n\nDEFAULT_PERSONA_ID = 0\n\n\ndef upgrade() -> None:\n # Make task_prompt column nullable\n # Note: The model had nullable=True but the DB column was NOT NULL until this point\n op.alter_column(\n \"persona\",\n \"task_prompt\",\n nullable=True,\n )\n\n # Set task_prompt to NULL for the default persona\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET task_prompt = NULL\n WHERE id = :persona_id\n \"\"\"\n ),\n {\"persona_id\": DEFAULT_PERSONA_ID},\n )\n\n\ndef downgrade() -> None:\n # Restore task_prompt to empty string for the default persona\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET task_prompt = ''\n WHERE id = :persona_id AND task_prompt IS NULL\n \"\"\"\n ),\n {\"persona_id\": DEFAULT_PERSONA_ID},\n )\n\n # Set any remaining NULL task_prompts to empty string before making non-nullable\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET task_prompt = ''\n WHERE task_prompt IS NULL\n \"\"\"\n )\n )\n\n # Revert task_prompt column to not nullable\n op.alter_column(\n \"persona\",\n \"task_prompt\",\n nullable=False,\n )\n" }, { "path": "backend/alembic/versions/6a804aeb4830_duplicated_no_harm_user_file_migration.py", "content": "\"\"\"duplicated no-harm user file migration\n\nRevision ID: 6a804aeb4830\nRevises: 8e1ac4f39a9f\nCreate Date: 2025-04-01 07:26:10.539362\n\n\"\"\"\n\n# revision identifiers, used by Alembic.\nrevision = \"6a804aeb4830\"\ndown_revision = \"8e1ac4f39a9f\"\nbranch_labels = None\ndepends_on = None\n\n\n# Leaving this around only because some people might be on this migration\n# originally was a duplicate of the user files migration\ndef upgrade() -> None:\n pass\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/6b3b4083c5aa_persona_cleanup_and_featured.py", "content": "\"\"\"persona cleanup and featured\n\nRevision ID: 6b3b4083c5aa\nRevises: 57122d037335\nCreate Date: 2026-02-26 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"6b3b4083c5aa\"\ndown_revision = \"57122d037335\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add featured column with nullable=True first\n op.add_column(\"persona\", sa.Column(\"featured\", sa.Boolean(), nullable=True))\n\n # Migrate data from is_default_persona to featured\n op.execute(\"UPDATE persona SET featured = is_default_persona\")\n\n # Make featured non-nullable with default=False\n op.alter_column(\n \"persona\",\n \"featured\",\n existing_type=sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n )\n\n # Drop is_default_persona column\n op.drop_column(\"persona\", \"is_default_persona\")\n\n # Drop unused columns\n op.drop_column(\"persona\", \"num_chunks\")\n op.drop_column(\"persona\", \"chunks_above\")\n op.drop_column(\"persona\", \"chunks_below\")\n op.drop_column(\"persona\", \"llm_relevance_filter\")\n op.drop_column(\"persona\", \"llm_filter_extraction\")\n op.drop_column(\"persona\", \"recency_bias\")\n\n\ndef downgrade() -> None:\n # Add back recency_bias column\n op.add_column(\n \"persona\",\n sa.Column(\n \"recency_bias\",\n sa.VARCHAR(),\n nullable=False,\n server_default=\"base_decay\",\n ),\n )\n\n # Add back llm_filter_extraction column\n op.add_column(\n \"persona\",\n sa.Column(\n \"llm_filter_extraction\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n )\n\n # Add back llm_relevance_filter column\n op.add_column(\n \"persona\",\n sa.Column(\n \"llm_relevance_filter\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n )\n\n # Add back chunks_below column\n op.add_column(\n \"persona\",\n sa.Column(\"chunks_below\", sa.Integer(), nullable=False, server_default=\"0\"),\n )\n\n # Add back chunks_above column\n op.add_column(\n \"persona\",\n sa.Column(\"chunks_above\", sa.Integer(), nullable=False, server_default=\"0\"),\n )\n\n # Add back num_chunks column\n op.add_column(\"persona\", sa.Column(\"num_chunks\", sa.Float(), nullable=True))\n\n # Add back is_default_persona column\n op.add_column(\n \"persona\",\n sa.Column(\n \"is_default_persona\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n ),\n )\n\n # Migrate data from featured to is_default_persona\n op.execute(\"UPDATE persona SET is_default_persona = featured\")\n\n # Drop featured column\n op.drop_column(\"persona\", \"featured\")\n" }, { "path": "backend/alembic/versions/6d387b3196c2_basic_auth.py", "content": "\"\"\"Basic Auth\n\nRevision ID: 6d387b3196c2\nRevises: 47433d30de82\nCreate Date: 2023-05-05 14:40:10.242502\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\nfrom alembic import op\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"6d387b3196c2\"\ndown_revision = \"47433d30de82\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"user\",\n sa.Column(\"id\", fastapi_users_db_sqlalchemy.generics.GUID(), nullable=False),\n sa.Column(\"email\", sa.String(length=320), nullable=False),\n sa.Column(\"hashed_password\", sa.String(length=1024), nullable=False),\n sa.Column(\"is_active\", sa.Boolean(), nullable=False),\n sa.Column(\"is_superuser\", sa.Boolean(), nullable=False),\n sa.Column(\"is_verified\", sa.Boolean(), nullable=False),\n sa.Column(\n \"role\",\n sa.Enum(\"BASIC\", \"ADMIN\", name=\"userrole\", native_enum=False),\n default=\"BASIC\",\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_index(op.f(\"ix_user_email\"), \"user\", [\"email\"], unique=True)\n op.create_table(\n \"accesstoken\",\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\"token\", sa.String(length=43), nullable=False),\n sa.Column(\n \"created_at\",\n fastapi_users_db_sqlalchemy.generics.TIMESTAMPAware(timezone=True),\n nullable=False,\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"cascade\"),\n sa.PrimaryKeyConstraint(\"token\"),\n )\n op.create_index(\n op.f(\"ix_accesstoken_created_at\"),\n \"accesstoken\",\n [\"created_at\"],\n unique=False,\n )\n op.alter_column(\n \"index_attempt\",\n \"time_created\",\n existing_type=postgresql.TIMESTAMP(timezone=True),\n nullable=False,\n existing_server_default=sa.text(\"now()\"), # type: ignore\n )\n op.alter_column(\n \"index_attempt\",\n \"time_updated\",\n existing_type=postgresql.TIMESTAMP(timezone=True),\n nullable=False,\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"index_attempt\",\n \"time_updated\",\n existing_type=postgresql.TIMESTAMP(timezone=True),\n nullable=True,\n )\n op.alter_column(\n \"index_attempt\",\n \"time_created\",\n existing_type=postgresql.TIMESTAMP(timezone=True),\n nullable=True,\n existing_server_default=sa.text(\"now()\"), # type: ignore\n )\n op.drop_index(op.f(\"ix_accesstoken_created_at\"), table_name=\"accesstoken\")\n op.drop_table(\"accesstoken\")\n op.drop_index(op.f(\"ix_user_email\"), table_name=\"user\")\n op.drop_table(\"user\")\n" }, { "path": "backend/alembic/versions/6d562f86c78b_remove_default_bot.py", "content": "\"\"\"remove default bot\n\nRevision ID: 6d562f86c78b\nRevises: 177de57c21c9\nCreate Date: 2024-11-22 11:51:29.331336\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"6d562f86c78b\"\ndown_revision = \"177de57c21c9\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.execute(\n sa.text(\n \"\"\"\n DELETE FROM slack_bot\n WHERE name = 'Default Bot'\n AND bot_token = ''\n AND app_token = ''\n AND NOT EXISTS (\n SELECT 1 FROM slack_channel_config\n WHERE slack_channel_config.slack_bot_id = slack_bot.id\n )\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n op.execute(\n sa.text(\n \"\"\"\n INSERT INTO slack_bot (name, enabled, bot_token, app_token)\n SELECT 'Default Bot', true, '', ''\n WHERE NOT EXISTS (SELECT 1 FROM slack_bot)\n RETURNING id;\n \"\"\"\n )\n )\n" }, { "path": "backend/alembic/versions/6f4f86aef280_add_queries_and_is_web_fetch_to_.py", "content": "\"\"\"add queries and is web fetch to iteration answer\n\nRevision ID: 6f4f86aef280\nRevises: 03d710ccf29c\nCreate Date: 2025-10-14 18:08:30.920123\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"6f4f86aef280\"\ndown_revision = \"03d710ccf29c\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add is_web_fetch column\n op.add_column(\n \"research_agent_iteration_sub_step\",\n sa.Column(\"is_web_fetch\", sa.Boolean(), nullable=True),\n )\n\n # Add queries column\n op.add_column(\n \"research_agent_iteration_sub_step\",\n sa.Column(\"queries\", postgresql.JSONB(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"research_agent_iteration_sub_step\", \"queries\")\n op.drop_column(\"research_agent_iteration_sub_step\", \"is_web_fetch\")\n" }, { "path": "backend/alembic/versions/6fc7886d665d_make_categories_labels_and_many_to_many.py", "content": "\"\"\"make categories labels and many to many\n\nRevision ID: 6fc7886d665d\nRevises: 3c6531f32351\nCreate Date: 2025-01-13 18:12:18.029112\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"6fc7886d665d\"\ndown_revision = \"3c6531f32351\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Rename persona_category table to persona_label\n op.rename_table(\"persona_category\", \"persona_label\")\n\n # Create the new association table\n op.create_table(\n \"persona__persona_label\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"persona_label_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"persona_label_id\"],\n [\"persona_label.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"persona_label_id\"),\n )\n\n # Copy existing relationships to the new table\n op.execute(\n \"\"\"\n INSERT INTO persona__persona_label (persona_id, persona_label_id)\n SELECT id, category_id FROM persona WHERE category_id IS NOT NULL\n \"\"\"\n )\n\n # Remove the old category_id column from persona table\n op.drop_column(\"persona\", \"category_id\")\n\n\ndef downgrade() -> None:\n # Rename persona_label table back to persona_category\n op.rename_table(\"persona_label\", \"persona_category\")\n\n # Add back the category_id column to persona table\n op.add_column(\"persona\", sa.Column(\"category_id\", sa.Integer(), nullable=True))\n op.create_foreign_key(\n \"persona_category_id_fkey\",\n \"persona\",\n \"persona_category\",\n [\"category_id\"],\n [\"id\"],\n )\n\n # Copy the first label relationship back to the persona table\n op.execute(\n \"\"\"\n UPDATE persona\n SET category_id = (\n SELECT persona_label_id\n FROM persona__persona_label\n WHERE persona__persona_label.persona_id = persona.id\n LIMIT 1\n )\n \"\"\"\n )\n\n # Drop the association table\n op.drop_table(\"persona__persona_label\")\n" }, { "path": "backend/alembic/versions/703313b75876_add_tokenratelimit_tables.py", "content": "\"\"\"Add TokenRateLimit Tables\n\nRevision ID: 703313b75876\nRevises: fad14119fb92\nCreate Date: 2024-04-15 01:36:02.952809\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"703313b75876\"\ndown_revision = \"fad14119fb92\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"token_rate_limit\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"enabled\", sa.Boolean(), nullable=False),\n sa.Column(\"token_budget\", sa.Integer(), nullable=False),\n sa.Column(\"period_hours\", sa.Integer(), nullable=False),\n sa.Column(\n \"scope\",\n sa.String(length=10),\n nullable=False,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"token_rate_limit__user_group\",\n sa.Column(\"rate_limit_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_group_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"rate_limit_id\"],\n [\"token_rate_limit.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.PrimaryKeyConstraint(\"rate_limit_id\", \"user_group_id\"),\n )\n\n # NOTE: rate limit settings used to be stored in the \"token_budget_settings\" key in the\n # KeyValueStore. This will now be lost. The KV store works differently than it used to\n # so the migration is fairly complicated and likely not worth it to support (pretty much\n # nobody will have it set)\n\n\ndef downgrade() -> None:\n op.drop_table(\"token_rate_limit__user_group\")\n op.drop_table(\"token_rate_limit\")\n" }, { "path": "backend/alembic/versions/70f00c45c0f2_more_descriptive_filestore.py", "content": "\"\"\"More Descriptive Filestore\n\nRevision ID: 70f00c45c0f2\nRevises: 3879338f8ba1\nCreate Date: 2024-05-17 17:51:41.926893\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"70f00c45c0f2\"\ndown_revision = \"3879338f8ba1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"file_store\", sa.Column(\"display_name\", sa.String(), nullable=True))\n op.add_column(\n \"file_store\",\n sa.Column(\n \"file_origin\",\n sa.String(),\n nullable=False,\n server_default=\"connector\", # Default to connector\n ),\n )\n op.add_column(\n \"file_store\",\n sa.Column(\n \"file_type\", sa.String(), nullable=False, server_default=\"text/plain\"\n ),\n )\n op.add_column(\n \"file_store\",\n sa.Column(\n \"file_metadata\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n\n op.execute(\n \"\"\"\n UPDATE file_store\n SET file_origin = CASE\n WHEN file_name LIKE 'chat__%' THEN 'chat_upload'\n ELSE 'connector'\n END,\n file_name = CASE\n WHEN file_name LIKE 'chat__%' THEN SUBSTR(file_name, 7)\n ELSE file_name\n END,\n file_type = CASE\n WHEN file_name LIKE 'chat__%' THEN 'image/png'\n ELSE 'text/plain'\n END\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"file_store\", \"file_metadata\")\n op.drop_column(\"file_store\", \"file_type\")\n op.drop_column(\"file_store\", \"file_origin\")\n op.drop_column(\"file_store\", \"display_name\")\n" }, { "path": "backend/alembic/versions/7206234e012a_add_image_generation_config_table.py", "content": "\"\"\"add image generation config table\n\nRevision ID: 7206234e012a\nRevises: 699221885109\nCreate Date: 2025-12-21 00:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"7206234e012a\"\ndown_revision = \"699221885109\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"image_generation_config\",\n sa.Column(\"image_provider_id\", sa.String(), primary_key=True),\n sa.Column(\"model_configuration_id\", sa.Integer(), nullable=False),\n sa.Column(\"is_default\", sa.Boolean(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"model_configuration_id\"],\n [\"model_configuration.id\"],\n ondelete=\"CASCADE\",\n ),\n )\n op.create_index(\n \"ix_image_generation_config_is_default\",\n \"image_generation_config\",\n [\"is_default\"],\n unique=False,\n )\n op.create_index(\n \"ix_image_generation_config_model_configuration_id\",\n \"image_generation_config\",\n [\"model_configuration_id\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"ix_image_generation_config_model_configuration_id\",\n table_name=\"image_generation_config\",\n )\n op.drop_index(\n \"ix_image_generation_config_is_default\", table_name=\"image_generation_config\"\n )\n op.drop_table(\"image_generation_config\")\n" }, { "path": "backend/alembic/versions/72aa7de2e5cf_make_processing_mode_default_all_caps.py", "content": "\"\"\"make processing mode default all caps\n\nRevision ID: 72aa7de2e5cf\nRevises: 2020d417ec84\nCreate Date: 2026-01-26 18:58:47.705253\n\nThis migration fixes the ProcessingMode enum value mismatch:\n- SQLAlchemy's Enum with native_enum=False uses enum member NAMES as valid values\n- The original migration stored lowercase VALUES ('regular', 'file_system')\n- This converts existing data to uppercase NAMES ('REGULAR', 'FILE_SYSTEM')\n- Also drops any spurious native PostgreSQL enum type that may have been auto-created\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"72aa7de2e5cf\"\ndown_revision = \"2020d417ec84\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Convert existing lowercase values to uppercase to match enum member names\n op.execute(\n \"UPDATE connector_credential_pair SET processing_mode = 'REGULAR' WHERE processing_mode = 'regular'\"\n )\n op.execute(\n \"UPDATE connector_credential_pair SET processing_mode = 'FILE_SYSTEM' WHERE processing_mode = 'file_system'\"\n )\n\n # Update the server default to use uppercase\n op.alter_column(\n \"connector_credential_pair\",\n \"processing_mode\",\n server_default=\"REGULAR\",\n )\n\n\ndef downgrade() -> None:\n # State prior to this was broken, so we don't want to revert back to it\n pass\n" }, { "path": "backend/alembic/versions/72bdc9929a46_permission_auto_sync_framework.py", "content": "\"\"\"Permission Auto Sync Framework\n\nRevision ID: 72bdc9929a46\nRevises: 475fcefe8826\nCreate Date: 2024-04-14 21:15:28.659634\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"72bdc9929a46\"\ndown_revision = \"475fcefe8826\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"email_to_external_user_cache\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"external_user_id\", sa.String(), nullable=False),\n sa.Column(\"user_id\", sa.UUID(), nullable=True),\n sa.Column(\"user_email\", sa.String(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"external_permission\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", sa.UUID(), nullable=True),\n sa.Column(\"user_email\", sa.String(), nullable=False),\n sa.Column(\n \"source_type\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\"external_permission_group\", sa.String(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"permission_sync_run\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"source_type\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\"update_type\", sa.String(), nullable=False),\n sa.Column(\"cc_pair_id\", sa.Integer(), nullable=True),\n sa.Column(\n \"status\",\n sa.String(),\n nullable=False,\n ),\n sa.Column(\"error_msg\", sa.Text(), nullable=True),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"cc_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"permission_sync_run\")\n op.drop_table(\"external_permission\")\n op.drop_table(\"email_to_external_user_cache\")\n" }, { "path": "backend/alembic/versions/73e9983e5091_add_search_query_table.py", "content": "\"\"\"add_search_query_table\n\nRevision ID: 73e9983e5091\nRevises: d1b637d7050a\nCreate Date: 2026-01-14 14:16:52.837489\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"73e9983e5091\"\ndown_revision = \"d1b637d7050a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"search_query\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), primary_key=True),\n sa.Column(\n \"user_id\",\n postgresql.UUID(as_uuid=True),\n sa.ForeignKey(\"user.id\"),\n nullable=False,\n ),\n sa.Column(\"query\", sa.String(), nullable=False),\n sa.Column(\"query_expansions\", postgresql.ARRAY(sa.String()), nullable=True),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.func.now(),\n ),\n )\n\n op.create_index(\"ix_search_query_user_id\", \"search_query\", [\"user_id\"])\n op.create_index(\"ix_search_query_created_at\", \"search_query\", [\"created_at\"])\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_search_query_created_at\", table_name=\"search_query\")\n op.drop_index(\"ix_search_query_user_id\", table_name=\"search_query\")\n op.drop_table(\"search_query\")\n" }, { "path": "backend/alembic/versions/7477a5f5d728_added_model_defaults_for_users.py", "content": "\"\"\"Added model defaults for users\n\nRevision ID: 7477a5f5d728\nRevises: 213fd978c6d8\nCreate Date: 2024-08-04 19:00:04.512634\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"7477a5f5d728\"\ndown_revision = \"213fd978c6d8\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"user\", sa.Column(\"default_model\", sa.Text(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"default_model\")\n" }, { "path": "backend/alembic/versions/7547d982db8f_chat_folders.py", "content": "\"\"\"Chat Folders\n\nRevision ID: 7547d982db8f\nRevises: ef7da92f7213\nCreate Date: 2024-05-02 15:18:56.573347\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport fastapi_users_db_sqlalchemy\n\n# revision identifiers, used by Alembic.\nrevision = \"7547d982db8f\"\ndown_revision = \"ef7da92f7213\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"chat_folder\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"name\", sa.String(), nullable=True),\n sa.Column(\"display_priority\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.add_column(\"chat_session\", sa.Column(\"folder_id\", sa.Integer(), nullable=True))\n op.create_foreign_key(\n \"chat_session_chat_folder_fk\",\n \"chat_session\",\n \"chat_folder\",\n [\"folder_id\"],\n [\"id\"],\n )\n\n\ndef downgrade() -> None:\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n if \"chat_session\" in inspector.get_table_names():\n chat_session_fks = {\n fk.get(\"name\") for fk in inspector.get_foreign_keys(\"chat_session\")\n }\n if \"chat_session_chat_folder_fk\" in chat_session_fks:\n op.drop_constraint(\n \"chat_session_chat_folder_fk\", \"chat_session\", type_=\"foreignkey\"\n )\n\n chat_session_columns = {\n col[\"name\"] for col in inspector.get_columns(\"chat_session\")\n }\n if \"folder_id\" in chat_session_columns:\n op.drop_column(\"chat_session\", \"folder_id\")\n\n if \"chat_folder\" in inspector.get_table_names():\n op.drop_table(\"chat_folder\")\n" }, { "path": "backend/alembic/versions/7616121f6e97_add_enterprise_fields_to_scim_user_mapping.py", "content": "\"\"\"add enterprise and name fields to scim_user_mapping\n\nRevision ID: 7616121f6e97\nRevises: 07b98176f1de\nCreate Date: 2026-02-23 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"7616121f6e97\"\ndown_revision = \"07b98176f1de\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"scim_user_mapping\",\n sa.Column(\"department\", sa.String(), nullable=True),\n )\n op.add_column(\n \"scim_user_mapping\",\n sa.Column(\"manager\", sa.String(), nullable=True),\n )\n op.add_column(\n \"scim_user_mapping\",\n sa.Column(\"given_name\", sa.String(), nullable=True),\n )\n op.add_column(\n \"scim_user_mapping\",\n sa.Column(\"family_name\", sa.String(), nullable=True),\n )\n op.add_column(\n \"scim_user_mapping\",\n sa.Column(\"scim_emails_json\", sa.Text(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"scim_user_mapping\", \"scim_emails_json\")\n op.drop_column(\"scim_user_mapping\", \"family_name\")\n op.drop_column(\"scim_user_mapping\", \"given_name\")\n op.drop_column(\"scim_user_mapping\", \"manager\")\n op.drop_column(\"scim_user_mapping\", \"department\")\n" }, { "path": "backend/alembic/versions/767f1c2a00eb_count_chat_tokens.py", "content": "\"\"\"Count Chat Tokens\n\nRevision ID: 767f1c2a00eb\nRevises: dba7f71618f5\nCreate Date: 2023-09-21 10:03:21.509899\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"767f1c2a00eb\"\ndown_revision = \"dba7f71618f5\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\", sa.Column(\"token_count\", sa.Integer(), nullable=False)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"token_count\")\n" }, { "path": "backend/alembic/versions/76b60d407dfb_cc_pair_name_not_unique.py", "content": "\"\"\"CC-Pair Name not Unique\n\nRevision ID: 76b60d407dfb\nRevises: b156fa702355\nCreate Date: 2023-12-22 21:42:10.018804\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"76b60d407dfb\"\ndown_revision = \"b156fa702355\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\"DELETE FROM connector_credential_pair WHERE name IS NULL\")\n op.drop_constraint(\n \"connector_credential_pair__name__key\",\n \"connector_credential_pair\",\n type_=\"unique\",\n )\n op.alter_column(\n \"connector_credential_pair\", \"name\", existing_type=sa.String(), nullable=False\n )\n\n\ndef downgrade() -> None:\n op.create_unique_constraint(\n \"connector_credential_pair__name__key\", \"connector_credential_pair\", [\"name\"]\n )\n op.alter_column(\n \"connector_credential_pair\", \"name\", existing_type=sa.String(), nullable=True\n )\n" }, { "path": "backend/alembic/versions/776b3bbe9092_remove_remaining_enums.py", "content": "\"\"\"Remove Remaining Enums\n\nRevision ID: 776b3bbe9092\nRevises: 4738e4b3bae1\nCreate Date: 2024-03-22 21:34:27.629444\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom onyx.db.models import IndexModelStatus\nfrom onyx.context.search.enums import RecencyBiasSetting, SearchType\n\n# revision identifiers, used by Alembic.\nrevision = \"776b3bbe9092\"\ndown_revision = \"4738e4b3bae1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"persona\",\n \"search_type\",\n type_=sa.String,\n existing_type=sa.Enum(SearchType, native_enum=False),\n existing_nullable=False,\n )\n op.alter_column(\n \"persona\",\n \"recency_bias\",\n type_=sa.String,\n existing_type=sa.Enum(RecencyBiasSetting, native_enum=False),\n existing_nullable=False,\n )\n\n # Because the indexmodelstatus enum does not have a mapping to a string type\n # we need this workaround instead of directly changing the type\n op.add_column(\"embedding_model\", sa.Column(\"temp_status\", sa.String))\n op.execute(\"UPDATE embedding_model SET temp_status = status::text\")\n op.drop_column(\"embedding_model\", \"status\")\n op.alter_column(\"embedding_model\", \"temp_status\", new_column_name=\"status\")\n\n op.execute(\"DROP TYPE IF EXISTS searchtype\")\n op.execute(\"DROP TYPE IF EXISTS recencybiassetting\")\n op.execute(\"DROP TYPE IF EXISTS indexmodelstatus\")\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"persona\",\n \"search_type\",\n type_=sa.Enum(SearchType, native_enum=False),\n existing_type=sa.String(length=50),\n existing_nullable=False,\n )\n op.alter_column(\n \"persona\",\n \"recency_bias\",\n type_=sa.Enum(RecencyBiasSetting, native_enum=False),\n existing_type=sa.String(length=50),\n existing_nullable=False,\n )\n op.alter_column(\n \"embedding_model\",\n \"status\",\n type_=sa.Enum(IndexModelStatus, native_enum=False),\n existing_type=sa.String(length=50),\n existing_nullable=False,\n )\n" }, { "path": "backend/alembic/versions/77d07dffae64_forcibly_remove_more_enum_types_from_.py", "content": "\"\"\"forcibly remove more enum types from postgres\n\nRevision ID: 77d07dffae64\nRevises: d61e513bef0a\nCreate Date: 2023-11-01 12:33:01.999617\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy import String\n\n\n# revision identifiers, used by Alembic.\nrevision = \"77d07dffae64\"\ndown_revision = \"d61e513bef0a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # In a PR:\n # https://github.com/onyx-dot-app/onyx/pull/397/files#diff-f05fb341f6373790b91852579631b64ca7645797a190837156a282b67e5b19c2\n # we directly changed some previous migrations. This caused some users to have native enums\n # while others wouldn't. This has caused some issues when adding new fields to these enums.\n # This migration manually changes the enum types to ensure that nobody uses native enums.\n op.alter_column(\"query_event\", \"selected_search_flow\", type_=String)\n op.alter_column(\"query_event\", \"feedback\", type_=String)\n op.alter_column(\"document_retrieval_feedback\", \"feedback\", type_=String)\n op.execute(\"DROP TYPE IF EXISTS searchtype\")\n op.execute(\"DROP TYPE IF EXISTS qafeedbacktype\")\n op.execute(\"DROP TYPE IF EXISTS searchfeedbacktype\")\n\n\ndef downgrade() -> None:\n # We don't want Native Enums, do nothing\n pass\n" }, { "path": "backend/alembic/versions/78dbe7e38469_task_tracking.py", "content": "\"\"\"Task Tracking\n\nRevision ID: 78dbe7e38469\nRevises: 7ccea01261f6\nCreate Date: 2023-10-15 23:40:50.593262\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"78dbe7e38469\"\ndown_revision = \"7ccea01261f6\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"task_queue_jobs\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"task_id\", sa.String(), nullable=False),\n sa.Column(\"task_name\", sa.String(), nullable=False),\n sa.Column(\n \"status\",\n sa.Enum(\n \"PENDING\",\n \"STARTED\",\n \"SUCCESS\",\n \"FAILURE\",\n name=\"taskstatus\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\"start_time\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"register_time\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"task_queue_jobs\")\n" }, { "path": "backend/alembic/versions/78ebc66946a0_remove_reranking_from_search_settings.py", "content": "\"\"\"remove reranking from search_settings\n\nRevision ID: 78ebc66946a0\nRevises: 849b21c732f8\nCreate Date: 2026-01-28\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"78ebc66946a0\"\ndown_revision = \"849b21c732f8\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"search_settings\", \"disable_rerank_for_streaming\")\n op.drop_column(\"search_settings\", \"rerank_model_name\")\n op.drop_column(\"search_settings\", \"rerank_provider_type\")\n op.drop_column(\"search_settings\", \"rerank_api_key\")\n op.drop_column(\"search_settings\", \"rerank_api_url\")\n op.drop_column(\"search_settings\", \"num_rerank\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"disable_rerank_for_streaming\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_model_name\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_provider_type\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_api_key\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_api_url\", sa.String(), nullable=True)\n )\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"num_rerank\",\n sa.Integer(),\n nullable=False,\n server_default=str(20),\n ),\n )\n" }, { "path": "backend/alembic/versions/795b20b85b4b_add_llm_group_permissions_control.py", "content": "\"\"\"add_llm_group_permissions_control\n\nRevision ID: 795b20b85b4b\nRevises: 05c07bf07c00\nCreate Date: 2024-07-19 11:54:35.701558\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\nrevision = \"795b20b85b4b\"\ndown_revision = \"05c07bf07c00\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"llm_provider__user_group\",\n sa.Column(\"llm_provider_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_group_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"llm_provider_id\"],\n [\"llm_provider.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.PrimaryKeyConstraint(\"llm_provider_id\", \"user_group_id\"),\n )\n op.add_column(\n \"llm_provider\",\n sa.Column(\"is_public\", sa.Boolean(), nullable=False, server_default=\"true\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"llm_provider__user_group\")\n op.drop_column(\"llm_provider\", \"is_public\")\n" }, { "path": "backend/alembic/versions/797089dfb4d2_persona_start_date.py", "content": "\"\"\"persona_start_date\n\nRevision ID: 797089dfb4d2\nRevises: 55546a7967ee\nCreate Date: 2024-09-11 14:51:49.785835\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"797089dfb4d2\"\ndown_revision = \"55546a7967ee\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"search_start_date\", sa.DateTime(timezone=True), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"search_start_date\")\n" }, { "path": "backend/alembic/versions/79acd316403a_add_api_key_table.py", "content": "\"\"\"Add api_key table\n\nRevision ID: 79acd316403a\nRevises: 904e5138fffb\nCreate Date: 2024-01-11 17:56:37.934381\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"79acd316403a\"\ndown_revision = \"904e5138fffb\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"api_key\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"hashed_api_key\", sa.String(), nullable=False),\n sa.Column(\"api_key_display\", sa.String(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\n \"owner_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"api_key_display\"),\n sa.UniqueConstraint(\"hashed_api_key\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"api_key\")\n" }, { "path": "backend/alembic/versions/7a70b7664e37_add_model_configuration_table.py", "content": "\"\"\"Add model-configuration table\n\nRevision ID: 7a70b7664e37\nRevises: d961aca62eb3\nCreate Date: 2025-04-10 15:00:35.984669\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\nfrom onyx.llm.well_known_providers.llm_provider_options import (\n fetch_model_names_for_provider_as_set,\n fetch_visible_model_names_for_provider_as_set,\n)\n\n# revision identifiers, used by Alembic.\nrevision = \"7a70b7664e37\"\ndown_revision = \"d961aca62eb3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef _resolve(\n provider_name: str,\n model_names: list[str] | None,\n display_model_names: list[str] | None,\n default_model_name: str,\n fast_default_model_name: str | None,\n) -> set[tuple[str, bool]]:\n models = set(model_names) if model_names else None\n display_models = set(display_model_names) if display_model_names else None\n\n # If both are defined, we need to make sure that `model_names` is a superset of `display_model_names`.\n if models and display_models:\n models = display_models.union(models)\n\n # If only `model_names` is defined, then:\n # - If default-model-names are available for the `provider_name`, then set `display_model_names` to it\n # and set `model_names` to the union of those default-model-names with itself.\n # - If no default-model-names are available, then set `display_models` to `models`.\n #\n # This preserves the invariant that `display_models` is a subset of `models`.\n elif models and not display_models:\n visible_default_models = fetch_visible_model_names_for_provider_as_set(\n provider_name=provider_name\n )\n if visible_default_models:\n display_models = set(visible_default_models)\n models = display_models.union(models)\n else:\n display_models = set(models)\n\n # If only the `display_model_names` are defined, then set `models` to the union of `display_model_names`\n # and the default-model-names for that provider.\n #\n # This will also preserve the invariant that `display_models` is a subset of `models`.\n elif not models and display_models:\n default_models = fetch_model_names_for_provider_as_set(\n provider_name=provider_name\n )\n if default_models:\n models = display_models.union(default_models)\n else:\n models = set(display_models)\n\n # If neither are defined, then set `models` and `display_models` to the default-model-names for the given provider.\n #\n # This will also preserve the invariant that `display_models` is a subset of `models`.\n else:\n default_models = fetch_model_names_for_provider_as_set(\n provider_name=provider_name\n )\n visible_default_models = fetch_visible_model_names_for_provider_as_set(\n provider_name=provider_name\n )\n\n if default_models:\n if not visible_default_models:\n raise RuntimeError\n raise RuntimeError(\n \"If `default_models` is non-None, `visible_default_models` must be non-None too.\"\n )\n models = default_models\n display_models = visible_default_models\n\n # This is not a well-known llm-provider; we can't provide any model suggestions.\n # Therefore, we set to the empty set and continue\n else:\n models = set()\n display_models = set()\n\n # It is possible that `default_model_name` is not in `models` and is not in `display_models`.\n # It is also possible that `fast_default_model_name` is not in `models` and is not in `display_models`.\n models.add(default_model_name)\n if fast_default_model_name:\n models.add(fast_default_model_name)\n display_models.add(default_model_name)\n if fast_default_model_name:\n display_models.add(fast_default_model_name)\n\n return set([(model, model in display_models) for model in models])\n\n\ndef upgrade() -> None:\n op.create_table(\n \"model_configuration\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"llm_provider_id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"is_visible\", sa.Boolean(), nullable=False),\n sa.Column(\"max_input_tokens\", sa.Integer(), nullable=True),\n sa.ForeignKeyConstraint(\n [\"llm_provider_id\"], [\"llm_provider.id\"], ondelete=\"CASCADE\"\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"llm_provider_id\", \"name\"),\n )\n\n # Create temporary sqlalchemy references to tables for data migration\n llm_provider_table = sa.sql.table(\n \"llm_provider\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"provider\", sa.Integer),\n sa.column(\"model_names\", postgresql.ARRAY(sa.String)),\n sa.column(\"display_model_names\", postgresql.ARRAY(sa.String)),\n sa.column(\"default_model_name\", sa.String),\n sa.column(\"fast_default_model_name\", sa.String),\n )\n model_configuration_table = sa.sql.table(\n \"model_configuration\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"llm_provider_id\", sa.Integer),\n sa.column(\"name\", sa.String),\n sa.column(\"is_visible\", sa.Boolean),\n sa.column(\"max_input_tokens\", sa.Integer),\n )\n connection = op.get_bind()\n llm_providers = connection.execute(\n sa.select(\n llm_provider_table.c.id,\n llm_provider_table.c.provider,\n llm_provider_table.c.model_names,\n llm_provider_table.c.display_model_names,\n llm_provider_table.c.default_model_name,\n llm_provider_table.c.fast_default_model_name,\n )\n ).fetchall()\n\n for llm_provider in llm_providers:\n provider_id = llm_provider[0]\n provider_name = llm_provider[1]\n model_names = llm_provider[2]\n display_model_names = llm_provider[3]\n default_model_name = llm_provider[4]\n fast_default_model_name = llm_provider[5]\n\n model_configurations = _resolve(\n provider_name=provider_name,\n model_names=model_names,\n display_model_names=display_model_names,\n default_model_name=default_model_name,\n fast_default_model_name=fast_default_model_name,\n )\n\n for model_name, is_visible in model_configurations:\n connection.execute(\n model_configuration_table.insert().values(\n llm_provider_id=provider_id,\n name=model_name,\n is_visible=is_visible,\n max_input_tokens=None,\n )\n )\n\n op.drop_column(\"llm_provider\", \"model_names\")\n op.drop_column(\"llm_provider\", \"display_model_names\")\n\n\ndef downgrade() -> None:\n llm_provider = sa.table(\n \"llm_provider\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"model_names\", postgresql.ARRAY(sa.String)),\n sa.column(\"display_model_names\", postgresql.ARRAY(sa.String)),\n )\n\n model_configuration = sa.table(\n \"model_configuration\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"llm_provider_id\", sa.Integer),\n sa.column(\"name\", sa.String),\n sa.column(\"is_visible\", sa.Boolean),\n sa.column(\"max_input_tokens\", sa.Integer),\n )\n op.add_column(\n \"llm_provider\",\n sa.Column(\n \"model_names\",\n postgresql.ARRAY(sa.VARCHAR()),\n autoincrement=False,\n nullable=True,\n ),\n )\n op.add_column(\n \"llm_provider\",\n sa.Column(\n \"display_model_names\",\n postgresql.ARRAY(sa.VARCHAR()),\n autoincrement=False,\n nullable=True,\n ),\n )\n\n connection = op.get_bind()\n provider_ids = connection.execute(sa.select(llm_provider.c.id)).fetchall()\n\n for (provider_id,) in provider_ids:\n # Get all models for this provider\n models = connection.execute(\n sa.select(\n model_configuration.c.name, model_configuration.c.is_visible\n ).where(model_configuration.c.llm_provider_id == provider_id)\n ).fetchall()\n\n all_models = [model[0] for model in models]\n visible_models = [model[0] for model in models if model[1]]\n\n # Update provider with arrays\n op.execute(\n llm_provider.update()\n .where(llm_provider.c.id == provider_id)\n .values(model_names=all_models, display_model_names=visible_models)\n )\n\n op.drop_table(\"model_configuration\")\n" }, { "path": "backend/alembic/versions/7aea705850d5_added_slack_auto_filter.py", "content": "\"\"\"added slack_auto_filter\n\nRevision ID: 7aea705850d5\nRevises: 4505fd7302e1\nCreate Date: 2024-07-10 11:01:23.581015\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nrevision = \"7aea705850d5\"\ndown_revision = \"4505fd7302e1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"slack_bot_config\",\n sa.Column(\"enable_auto_filters\", sa.Boolean(), nullable=True),\n )\n op.execute(\n \"UPDATE slack_bot_config SET enable_auto_filters = FALSE WHERE enable_auto_filters IS NULL\"\n )\n op.alter_column(\n \"slack_bot_config\",\n \"enable_auto_filters\",\n existing_type=sa.Boolean(),\n nullable=False,\n server_default=sa.false(),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"slack_bot_config\", \"enable_auto_filters\")\n" }, { "path": "backend/alembic/versions/7b9b952abdf6_update_entities.py", "content": "\"\"\"update-entities\n\nRevision ID: 7b9b952abdf6\nRevises: 36e9220ab794\nCreate Date: 2025-06-23 20:24:08.139201\n\n\"\"\"\n\nimport json\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"7b9b952abdf6\"\ndown_revision = \"36e9220ab794\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # new entity type metadata_attribute_conversion\n new_entity_type_conversion = {\n \"LINEAR\": {\n \"team\": {\"name\": \"team\", \"keep\": True, \"implication_property\": None},\n \"state\": {\"name\": \"state\", \"keep\": True, \"implication_property\": None},\n \"priority\": {\n \"name\": \"priority\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"estimate\": {\n \"name\": \"estimate\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"created_at\": {\n \"name\": \"created_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"started_at\": {\n \"name\": \"started_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"completed_at\": {\n \"name\": \"completed_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"due_date\": {\n \"name\": \"due_date\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"creator\": {\n \"name\": \"creator\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_creator_of\",\n },\n },\n \"assignee\": {\n \"name\": \"assignee\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_assignee_of\",\n },\n },\n },\n \"JIRA\": {\n \"issuetype\": {\n \"name\": \"subtype\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"status\": {\"name\": \"status\", \"keep\": True, \"implication_property\": None},\n \"priority\": {\n \"name\": \"priority\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"project_name\": {\n \"name\": \"project\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"created\": {\n \"name\": \"created_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"updated\": {\n \"name\": \"updated_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"resolution_date\": {\n \"name\": \"completed_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"duedate\": {\"name\": \"due_date\", \"keep\": True, \"implication_property\": None},\n \"reporter_email\": {\n \"name\": \"creator\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_creator_of\",\n },\n },\n \"assignee_email\": {\n \"name\": \"assignee\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_assignee_of\",\n },\n },\n \"key\": {\"name\": \"key\", \"keep\": True, \"implication_property\": None},\n \"parent\": {\"name\": \"parent\", \"keep\": True, \"implication_property\": None},\n },\n \"GITHUB_PR\": {\n \"repo\": {\"name\": \"repository\", \"keep\": True, \"implication_property\": None},\n \"state\": {\"name\": \"state\", \"keep\": True, \"implication_property\": None},\n \"num_commits\": {\n \"name\": \"num_commits\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"num_files_changed\": {\n \"name\": \"num_files_changed\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"labels\": {\"name\": \"labels\", \"keep\": True, \"implication_property\": None},\n \"merged\": {\"name\": \"merged\", \"keep\": True, \"implication_property\": None},\n \"merged_at\": {\n \"name\": \"merged_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"closed_at\": {\n \"name\": \"closed_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"created_at\": {\n \"name\": \"created_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"updated_at\": {\n \"name\": \"updated_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"user\": {\n \"name\": \"creator\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_creator_of\",\n },\n },\n \"assignees\": {\n \"name\": \"assignees\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_assignee_of\",\n },\n },\n },\n \"GITHUB_ISSUE\": {\n \"repo\": {\"name\": \"repository\", \"keep\": True, \"implication_property\": None},\n \"state\": {\"name\": \"state\", \"keep\": True, \"implication_property\": None},\n \"labels\": {\"name\": \"labels\", \"keep\": True, \"implication_property\": None},\n \"closed_at\": {\n \"name\": \"closed_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"created_at\": {\n \"name\": \"created_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"updated_at\": {\n \"name\": \"updated_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"user\": {\n \"name\": \"creator\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_creator_of\",\n },\n },\n \"assignees\": {\n \"name\": \"assignees\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"from_email\",\n \"implied_relationship_name\": \"is_assignee_of\",\n },\n },\n },\n \"FIREFLIES\": {},\n \"ACCOUNT\": {},\n \"OPPORTUNITY\": {\n \"name\": {\"name\": \"name\", \"keep\": True, \"implication_property\": None},\n \"stage_name\": {\"name\": \"stage\", \"keep\": True, \"implication_property\": None},\n \"type\": {\"name\": \"type\", \"keep\": True, \"implication_property\": None},\n \"amount\": {\"name\": \"amount\", \"keep\": True, \"implication_property\": None},\n \"fiscal_year\": {\n \"name\": \"fiscal_year\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"fiscal_quarter\": {\n \"name\": \"fiscal_quarter\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"is_closed\": {\n \"name\": \"is_closed\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"close_date\": {\n \"name\": \"close_date\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"probability\": {\n \"name\": \"close_probability\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"created_date\": {\n \"name\": \"created_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"last_modified_date\": {\n \"name\": \"updated_at\",\n \"keep\": True,\n \"implication_property\": None,\n },\n \"account\": {\n \"name\": \"account\",\n \"keep\": False,\n \"implication_property\": {\n \"implied_entity_type\": \"ACCOUNT\",\n \"implied_relationship_name\": \"is_account_of\",\n },\n },\n },\n \"VENDOR\": {},\n \"EMPLOYEE\": {},\n }\n\n current_entity_types = conn.execute(\n sa.text(\"SELECT id_name, attributes from kg_entity_type\")\n ).all()\n for entity_type, attributes in current_entity_types:\n # delete removed entity types\n if entity_type not in new_entity_type_conversion:\n op.execute(\n sa.text(f\"DELETE FROM kg_entity_type WHERE id_name = '{entity_type}'\")\n )\n continue\n\n # update entity type attributes\n if \"metadata_attributes\" in attributes:\n del attributes[\"metadata_attributes\"]\n attributes[\"metadata_attribute_conversion\"] = new_entity_type_conversion[\n entity_type\n ]\n attributes_str = json.dumps(attributes).replace(\"'\", \"''\")\n op.execute(\n sa.text(\n f\"UPDATE kg_entity_type SET attributes = '{attributes_str}'WHERE id_name = '{entity_type}'\"\n ),\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n\n current_entity_types = conn.execute(\n sa.text(\"SELECT id_name, attributes from kg_entity_type\")\n ).all()\n for entity_type, attributes in current_entity_types:\n conversion = {}\n if \"metadata_attribute_conversion\" in attributes:\n conversion = attributes.pop(\"metadata_attribute_conversion\")\n attributes[\"metadata_attributes\"] = {\n attr: prop[\"name\"] for attr, prop in conversion.items() if prop[\"keep\"]\n }\n\n attributes_str = json.dumps(attributes).replace(\"'\", \"''\")\n op.execute(\n sa.text(\n f\"UPDATE kg_entity_type SET attributes = '{attributes_str}'WHERE id_name = '{entity_type}'\"\n ),\n )\n" }, { "path": "backend/alembic/versions/7bd55f264e1b_add_display_name_to_model_configuration.py", "content": "\"\"\"Add display_name to model_configuration\n\nRevision ID: 7bd55f264e1b\nRevises: e8f0d2a38171\nCreate Date: 2025-12-04\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"7bd55f264e1b\"\ndown_revision = \"e8f0d2a38171\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"model_configuration\",\n sa.Column(\"display_name\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"model_configuration\", \"display_name\")\n" }, { "path": "backend/alembic/versions/7cb492013621_code_interpreter_server_model.py", "content": "\"\"\"code interpreter server model\n\nRevision ID: 7cb492013621\nRevises: 0bb4558f35df\nCreate Date: 2026-02-22 18:54:54.007265\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"7cb492013621\"\ndown_revision = \"0bb4558f35df\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"code_interpreter_server\",\n sa.Column(\"id\", sa.Integer, primary_key=True),\n sa.Column(\n \"server_enabled\", sa.Boolean, nullable=False, server_default=sa.true()\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"code_interpreter_server\")\n" }, { "path": "backend/alembic/versions/7cc3fcc116c1_user_file_uuid_primary_key_swap.py", "content": "\"\"\"Migration 4: User file UUID primary key swap\n\nRevision ID: 7cc3fcc116c1\nRevises: 16c37a30adf2\nCreate Date: 2025-09-22 09:54:38.292952\n\nThis migration performs the critical UUID primary key swap on user_file table.\nIt updates all foreign key references to use UUIDs instead of integers.\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql as psql\nimport logging\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n# revision identifiers, used by Alembic.\nrevision = \"7cc3fcc116c1\"\ndown_revision = \"16c37a30adf2\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Swap user_file primary key from integer to UUID.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n # Verify we're in the expected state\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n if \"new_id\" not in user_file_columns:\n logger.warning(\n \"user_file.new_id not found - migration may have already been applied\"\n )\n return\n\n logger.info(\"Starting UUID primary key swap...\")\n\n # === Step 1: Update persona__user_file foreign key to UUID ===\n logger.info(\"Updating persona__user_file foreign key...\")\n\n # Drop existing foreign key constraints\n op.execute(\n \"ALTER TABLE persona__user_file DROP CONSTRAINT IF EXISTS persona__user_file_user_file_id_uuid_fkey\"\n )\n op.execute(\n \"ALTER TABLE persona__user_file DROP CONSTRAINT IF EXISTS persona__user_file_user_file_id_fkey\"\n )\n\n # Create new foreign key to user_file.new_id\n op.create_foreign_key(\n \"persona__user_file_user_file_id_fkey\",\n \"persona__user_file\",\n \"user_file\",\n local_cols=[\"user_file_id_uuid\"],\n remote_cols=[\"new_id\"],\n )\n\n # Drop the old integer column and rename UUID column\n op.execute(\"ALTER TABLE persona__user_file DROP COLUMN IF EXISTS user_file_id\")\n op.alter_column(\n \"persona__user_file\",\n \"user_file_id_uuid\",\n new_column_name=\"user_file_id\",\n existing_type=psql.UUID(as_uuid=True),\n nullable=False,\n )\n\n # Recreate composite primary key\n op.execute(\n \"ALTER TABLE persona__user_file DROP CONSTRAINT IF EXISTS persona__user_file_pkey\"\n )\n op.execute(\n \"ALTER TABLE persona__user_file ADD PRIMARY KEY (persona_id, user_file_id)\"\n )\n\n logger.info(\"Updated persona__user_file to use UUID foreign key\")\n\n # === Step 2: Perform the primary key swap on user_file ===\n logger.info(\"Swapping user_file primary key to UUID...\")\n\n # Drop the primary key constraint\n op.execute(\"ALTER TABLE user_file DROP CONSTRAINT IF EXISTS user_file_pkey\")\n\n # Drop the old id column and rename new_id to id\n op.execute(\"ALTER TABLE user_file DROP COLUMN IF EXISTS id\")\n op.alter_column(\n \"user_file\",\n \"new_id\",\n new_column_name=\"id\",\n existing_type=psql.UUID(as_uuid=True),\n nullable=False,\n )\n\n # Set default for new inserts\n op.alter_column(\n \"user_file\",\n \"id\",\n existing_type=psql.UUID(as_uuid=True),\n server_default=sa.text(\"gen_random_uuid()\"),\n )\n\n # Create new primary key\n op.execute(\"ALTER TABLE user_file ADD PRIMARY KEY (id)\")\n\n logger.info(\"Swapped user_file primary key to UUID\")\n\n # === Step 3: Update foreign key constraints ===\n logger.info(\"Updating foreign key constraints...\")\n\n # Recreate persona__user_file foreign key to point to user_file.id\n # Drop existing FK first to break dependency on the unique constraint\n op.execute(\n \"ALTER TABLE persona__user_file DROP CONSTRAINT IF EXISTS persona__user_file_user_file_id_fkey\"\n )\n # Drop the unique constraint on (formerly) new_id BEFORE recreating the FK,\n # so the FK will bind to the primary key instead of the unique index.\n op.execute(\"ALTER TABLE user_file DROP CONSTRAINT IF EXISTS uq_user_file_new_id\")\n # Now recreate FK to the primary key column\n op.create_foreign_key(\n \"persona__user_file_user_file_id_fkey\",\n \"persona__user_file\",\n \"user_file\",\n local_cols=[\"user_file_id\"],\n remote_cols=[\"id\"],\n )\n\n # Add foreign keys for project__user_file\n existing_fks = inspector.get_foreign_keys(\"project__user_file\")\n\n has_user_file_fk = any(\n fk.get(\"referred_table\") == \"user_file\"\n and fk.get(\"constrained_columns\") == [\"user_file_id\"]\n for fk in existing_fks\n )\n\n if not has_user_file_fk:\n op.create_foreign_key(\n \"fk_project__user_file_user_file_id\",\n \"project__user_file\",\n \"user_file\",\n [\"user_file_id\"],\n [\"id\"],\n )\n logger.info(\"Added project__user_file -> user_file foreign key\")\n\n has_project_fk = any(\n fk.get(\"referred_table\") == \"user_project\"\n and fk.get(\"constrained_columns\") == [\"project_id\"]\n for fk in existing_fks\n )\n\n if not has_project_fk:\n op.create_foreign_key(\n \"fk_project__user_file_project_id\",\n \"project__user_file\",\n \"user_project\",\n [\"project_id\"],\n [\"id\"],\n )\n logger.info(\"Added project__user_file -> user_project foreign key\")\n\n # === Step 4: Mark files for document_id migration ===\n logger.info(\"Marking files for background document_id migration...\")\n\n logger.info(\"Migration 4 (UUID primary key swap) completed successfully\")\n logger.info(\n \"NOTE: Background task will update document IDs in Vespa and search_doc\"\n )\n\n\ndef downgrade() -> None:\n \"\"\"Revert UUID primary key back to integer (data destructive!).\"\"\"\n\n logger.error(\"CRITICAL: Downgrading UUID primary key swap is data destructive!\")\n logger.error(\n \"This will break all UUID-based references created after the migration.\"\n )\n logger.error(\"Only proceed if absolutely necessary and have backups.\")\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n # Capture existing primary key definitions so we can restore them after swaps\n persona_pk = inspector.get_pk_constraint(\"persona__user_file\") or {}\n persona_pk_name = persona_pk.get(\"name\")\n persona_pk_cols = persona_pk.get(\"constrained_columns\") or []\n\n project_pk = inspector.get_pk_constraint(\"project__user_file\") or {}\n project_pk_name = project_pk.get(\"name\")\n project_pk_cols = project_pk.get(\"constrained_columns\") or []\n\n # Drop foreign keys that reference the UUID primary key\n op.drop_constraint(\n \"persona__user_file_user_file_id_fkey\",\n \"persona__user_file\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"fk_project__user_file_user_file_id\",\n \"project__user_file\",\n type_=\"foreignkey\",\n )\n\n # Drop primary keys that rely on the UUID column so we can replace it\n if persona_pk_name:\n op.drop_constraint(persona_pk_name, \"persona__user_file\", type_=\"primary\")\n if project_pk_name:\n op.drop_constraint(project_pk_name, \"project__user_file\", type_=\"primary\")\n\n # Rebuild integer IDs on user_file using a sequence-backed column\n op.execute(\"CREATE SEQUENCE IF NOT EXISTS user_file_id_seq\")\n op.add_column(\n \"user_file\",\n sa.Column(\n \"id_int\",\n sa.Integer(),\n server_default=sa.text(\"nextval('user_file_id_seq')\"),\n nullable=False,\n ),\n )\n op.execute(\"ALTER SEQUENCE user_file_id_seq OWNED BY user_file.id_int\")\n\n # Prepare integer foreign key columns on referencing tables\n op.add_column(\n \"persona__user_file\",\n sa.Column(\"user_file_id_int\", sa.Integer(), nullable=True),\n )\n op.add_column(\n \"project__user_file\",\n sa.Column(\"user_file_id_int\", sa.Integer(), nullable=True),\n )\n\n # Populate the new integer foreign key columns by mapping from the UUID IDs\n op.execute(\n \"\"\"\n UPDATE persona__user_file AS p\n SET user_file_id_int = uf.id_int\n FROM user_file AS uf\n WHERE p.user_file_id = uf.id\n \"\"\"\n )\n op.execute(\n \"\"\"\n UPDATE project__user_file AS p\n SET user_file_id_int = uf.id_int\n FROM user_file AS uf\n WHERE p.user_file_id = uf.id\n \"\"\"\n )\n\n op.alter_column(\n \"persona__user_file\",\n \"user_file_id_int\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n op.alter_column(\n \"project__user_file\",\n \"user_file_id_int\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n\n # Remove the UUID foreign key columns and rename the integer replacements\n op.drop_column(\"persona__user_file\", \"user_file_id\")\n op.alter_column(\n \"persona__user_file\",\n \"user_file_id_int\",\n new_column_name=\"user_file_id\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n\n op.drop_column(\"project__user_file\", \"user_file_id\")\n op.alter_column(\n \"project__user_file\",\n \"user_file_id_int\",\n new_column_name=\"user_file_id\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n\n # Swap the user_file primary key back to the integer column\n op.drop_constraint(\"user_file_pkey\", \"user_file\", type_=\"primary\")\n op.drop_column(\"user_file\", \"id\")\n op.alter_column(\n \"user_file\",\n \"id_int\",\n new_column_name=\"id\",\n existing_type=sa.Integer(),\n )\n op.alter_column(\n \"user_file\",\n \"id\",\n existing_type=sa.Integer(),\n nullable=False,\n server_default=sa.text(\"nextval('user_file_id_seq')\"),\n )\n op.execute(\"ALTER SEQUENCE user_file_id_seq OWNED BY user_file.id\")\n op.execute(\n \"\"\"\n SELECT setval(\n 'user_file_id_seq',\n GREATEST(COALESCE(MAX(id), 1), 1),\n MAX(id) IS NOT NULL\n )\n FROM user_file\n \"\"\"\n )\n op.create_primary_key(\"user_file_pkey\", \"user_file\", [\"id\"])\n\n # Restore primary keys on referencing tables\n if persona_pk_cols:\n op.create_primary_key(\n \"persona__user_file_pkey\", \"persona__user_file\", persona_pk_cols\n )\n if project_pk_cols:\n op.create_primary_key(\n \"project__user_file_pkey\",\n \"project__user_file\",\n project_pk_cols,\n )\n\n # Recreate foreign keys pointing at the integer primary key\n op.create_foreign_key(\n \"persona__user_file_user_file_id_fkey\",\n \"persona__user_file\",\n \"user_file\",\n [\"user_file_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"fk_project__user_file_user_file_id\",\n \"project__user_file\",\n \"user_file\",\n [\"user_file_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/7ccea01261f6_store_chat_retrieval_docs.py", "content": "\"\"\"Store Chat Retrieval Docs\n\nRevision ID: 7ccea01261f6\nRevises: a570b80a5f20\nCreate Date: 2023-10-15 10:39:23.317453\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"7ccea01261f6\"\ndown_revision = \"a570b80a5f20\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"reference_docs\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"reference_docs\")\n" }, { "path": "backend/alembic/versions/7da0ae5ad583_add_description_to_persona.py", "content": "\"\"\"Add description to persona\n\nRevision ID: 7da0ae5ad583\nRevises: e86866a9c78a\nCreate Date: 2023-11-27 00:16:19.959414\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"7da0ae5ad583\"\ndown_revision = \"e86866a9c78a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"persona\", sa.Column(\"description\", sa.String(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"description\")\n" }, { "path": "backend/alembic/versions/7da543f5672f_add_slackbotconfig_table.py", "content": "\"\"\"Add SlackBotConfig table\n\nRevision ID: 7da543f5672f\nRevises: febe9eaa0644\nCreate Date: 2023-09-24 16:34:17.526128\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"7da543f5672f\"\ndown_revision = \"febe9eaa0644\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"slack_bot_config\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"persona_id\", sa.Integer(), nullable=True),\n sa.Column(\n \"channel_config\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"slack_bot_config\")\n" }, { "path": "backend/alembic/versions/7e490836d179_nullify_default_system_prompt.py", "content": "\"\"\"nullify_default_system_prompt\n\nRevision ID: 7e490836d179\nRevises: c1d2e3f4a5b6\nCreate Date: 2025-12-29 16:54:36.635574\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"7e490836d179\"\ndown_revision = \"c1d2e3f4a5b6\"\nbranch_labels = None\ndepends_on = None\n\n\n# This is the default system prompt from the previous migration (87c52ec39f84)\n# ruff: noqa: E501, W605 start\nPREVIOUS_DEFAULT_SYSTEM_PROMPT = \"\"\"\nYou are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent, ask clarifying questions when needed, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. Always prioritize being truthful, nuanced, insightful, and efficient.\n\nThe current date is [[CURRENT_DATETIME]].[[CITATION_GUIDANCE]]\n\n# Response Style\nYou use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make your responses more readable and engaging.\nYou use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, symbols, etc.: '$$\\\\n[expression]\\\\n$$' for standalone cases and '\\\\( [expression] \\\\)' when inline.\nFor code you prefer to use Markdown and specify the language.\nYou can use horizontal rules (---) to separate sections of your responses.\nYou can use Markdown tables to format your responses for data, lists, and other structured information.\n\"\"\".lstrip()\n# ruff: noqa: E501, W605 end\n\n\ndef upgrade() -> None:\n # Make system_prompt column nullable (model already has nullable=True but DB doesn't)\n op.alter_column(\n \"persona\",\n \"system_prompt\",\n nullable=True,\n )\n\n # Set system_prompt to NULL where it matches the previous default\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET system_prompt = NULL\n WHERE system_prompt = :previous_default\n \"\"\"\n ),\n {\"previous_default\": PREVIOUS_DEFAULT_SYSTEM_PROMPT},\n )\n\n\ndef downgrade() -> None:\n # Restore the default system prompt for personas that have NULL\n # Note: This may restore the prompt to personas that originally had NULL\n # before this migration, but there's no way to distinguish them\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET system_prompt = :previous_default\n WHERE system_prompt IS NULL\n \"\"\"\n ),\n {\"previous_default\": PREVIOUS_DEFAULT_SYSTEM_PROMPT},\n )\n\n # Revert system_prompt column to not nullable\n op.alter_column(\n \"persona\",\n \"system_prompt\",\n nullable=False,\n )\n" }, { "path": "backend/alembic/versions/7ed603b64d5a_add_mcp_server_and_connection_config_.py", "content": "\"\"\"add_mcp_server_and_connection_config_models\n\nRevision ID: 7ed603b64d5a\nRevises: b329d00a9ea6\nCreate Date: 2025-07-28 17:35:59.900680\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom onyx.db.enums import MCPAuthenticationType\n\n# revision identifiers, used by Alembic.\nrevision = \"7ed603b64d5a\"\ndown_revision = \"b329d00a9ea6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Create tables and columns for MCP Server support\"\"\"\n\n # 1. MCP Server main table (no FK constraints yet to avoid circular refs)\n op.create_table(\n \"mcp_server\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"owner\", sa.String(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"description\", sa.String(), nullable=True),\n sa.Column(\"server_url\", sa.String(), nullable=False),\n sa.Column(\n \"auth_type\",\n sa.Enum(\n MCPAuthenticationType,\n name=\"mcp_authentication_type\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\"admin_connection_config_id\", sa.Integer(), nullable=True),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n )\n\n # 2. MCP Connection Config table (can reference mcp_server now that it exists)\n op.create_table(\n \"mcp_connection_config\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"mcp_server_id\", sa.Integer(), nullable=True),\n sa.Column(\"user_email\", sa.String(), nullable=False, default=\"\"),\n sa.Column(\"config\", sa.LargeBinary(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"mcp_server_id\"], [\"mcp_server.id\"], ondelete=\"CASCADE\"\n ),\n )\n\n # Helpful indexes\n op.create_index(\n \"ix_mcp_connection_config_server_user\",\n \"mcp_connection_config\",\n [\"mcp_server_id\", \"user_email\"],\n )\n op.create_index(\n \"ix_mcp_connection_config_user_email\",\n \"mcp_connection_config\",\n [\"user_email\"],\n )\n\n # 3. Add the back-references from mcp_server to connection configs\n op.create_foreign_key(\n \"mcp_server_admin_config_fk\",\n \"mcp_server\",\n \"mcp_connection_config\",\n [\"admin_connection_config_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n\n # 4. Association / access-control tables\n op.create_table(\n \"mcp_server__user\",\n sa.Column(\"mcp_server_id\", sa.Integer(), primary_key=True),\n sa.Column(\"user_id\", sa.UUID(), primary_key=True),\n sa.ForeignKeyConstraint(\n [\"mcp_server_id\"], [\"mcp_server.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n )\n\n op.create_table(\n \"mcp_server__user_group\",\n sa.Column(\"mcp_server_id\", sa.Integer(), primary_key=True),\n sa.Column(\"user_group_id\", sa.Integer(), primary_key=True),\n sa.ForeignKeyConstraint(\n [\"mcp_server_id\"], [\"mcp_server.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"user_group_id\"], [\"user_group.id\"]),\n )\n\n # 5. Update existing `tool` table – allow tools to belong to an MCP server\n op.add_column(\n \"tool\",\n sa.Column(\"mcp_server_id\", sa.Integer(), nullable=True),\n )\n # Add column for MCP tool input schema\n op.add_column(\n \"tool\",\n sa.Column(\"mcp_input_schema\", postgresql.JSONB(), nullable=True),\n )\n op.create_foreign_key(\n \"tool_mcp_server_fk\",\n \"tool\",\n \"mcp_server\",\n [\"mcp_server_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # 6. Update persona__tool foreign keys to cascade delete\n # This ensures that when a tool is deleted (including via MCP server deletion),\n # the corresponding persona__tool rows are also deleted\n op.drop_constraint(\n \"persona__tool_tool_id_fkey\", \"persona__tool\", type_=\"foreignkey\"\n )\n op.drop_constraint(\n \"persona__tool_persona_id_fkey\", \"persona__tool\", type_=\"foreignkey\"\n )\n\n op.create_foreign_key(\n \"persona__tool_persona_id_fkey\",\n \"persona__tool\",\n \"persona\",\n [\"persona_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n op.create_foreign_key(\n \"persona__tool_tool_id_fkey\",\n \"persona__tool\",\n \"tool\",\n [\"tool_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # 7. Update research_agent_iteration_sub_step foreign key to SET NULL on delete\n # This ensures that when a tool is deleted, the sub_step_tool_id is set to NULL\n # instead of causing a foreign key constraint violation\n op.drop_constraint(\n \"research_agent_iteration_sub_step_sub_step_tool_id_fkey\",\n \"research_agent_iteration_sub_step\",\n type_=\"foreignkey\",\n )\n op.create_foreign_key(\n \"research_agent_iteration_sub_step_sub_step_tool_id_fkey\",\n \"research_agent_iteration_sub_step\",\n \"tool\",\n [\"sub_step_tool_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n\n\ndef downgrade() -> None:\n \"\"\"Drop all MCP-related tables / columns\"\"\"\n\n # # # 1. Drop FK & columns from tool\n # op.drop_constraint(\"tool_mcp_server_fk\", \"tool\", type_=\"foreignkey\")\n op.execute(\"DELETE FROM tool WHERE mcp_server_id IS NOT NULL\")\n\n op.drop_constraint(\n \"research_agent_iteration_sub_step_sub_step_tool_id_fkey\",\n \"research_agent_iteration_sub_step\",\n type_=\"foreignkey\",\n )\n op.create_foreign_key(\n \"research_agent_iteration_sub_step_sub_step_tool_id_fkey\",\n \"research_agent_iteration_sub_step\",\n \"tool\",\n [\"sub_step_tool_id\"],\n [\"id\"],\n )\n\n # Restore original persona__tool foreign keys (without CASCADE)\n op.drop_constraint(\n \"persona__tool_persona_id_fkey\", \"persona__tool\", type_=\"foreignkey\"\n )\n op.drop_constraint(\n \"persona__tool_tool_id_fkey\", \"persona__tool\", type_=\"foreignkey\"\n )\n\n op.create_foreign_key(\n \"persona__tool_persona_id_fkey\",\n \"persona__tool\",\n \"persona\",\n [\"persona_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"persona__tool_tool_id_fkey\",\n \"persona__tool\",\n \"tool\",\n [\"tool_id\"],\n [\"id\"],\n )\n op.drop_column(\"tool\", \"mcp_input_schema\")\n op.drop_column(\"tool\", \"mcp_server_id\")\n\n # 2. Drop association tables\n op.drop_table(\"mcp_server__user_group\")\n op.drop_table(\"mcp_server__user\")\n\n # 3. Drop FK from mcp_server to connection configs\n op.drop_constraint(\"mcp_server_admin_config_fk\", \"mcp_server\", type_=\"foreignkey\")\n\n # 4. Drop connection config indexes & table\n op.drop_index(\n \"ix_mcp_connection_config_user_email\", table_name=\"mcp_connection_config\"\n )\n op.drop_index(\n \"ix_mcp_connection_config_server_user\", table_name=\"mcp_connection_config\"\n )\n op.drop_table(\"mcp_connection_config\")\n\n # 5. Finally drop mcp_server table\n op.drop_table(\"mcp_server\")\n" }, { "path": "backend/alembic/versions/7f726bad5367_slack_followup.py", "content": "\"\"\"Slack Followup\n\nRevision ID: 7f726bad5367\nRevises: 79acd316403a\nCreate Date: 2024-01-15 00:19:55.991224\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"7f726bad5367\"\ndown_revision = \"79acd316403a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_feedback\",\n sa.Column(\"required_followup\", sa.Boolean(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_feedback\", \"required_followup\")\n" }, { "path": "backend/alembic/versions/7f99be1cb9f5_add_index_for_getting_documents_just_by_.py", "content": "\"\"\"Add index for getting documents just by connector id / credential id\n\nRevision ID: 7f99be1cb9f5\nRevises: 78dbe7e38469\nCreate Date: 2023-10-15 22:48:15.487762\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"7f99be1cb9f5\"\ndown_revision = \"78dbe7e38469\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_index(\n op.f(\n \"ix_document_by_connector_credential_pair_pkey__connector_id__credential_id\"\n ),\n \"document_by_connector_credential_pair\",\n [\"connector_id\", \"credential_id\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n op.f(\n \"ix_document_by_connector_credential_pair_pkey__connector_id__credential_id\"\n ),\n table_name=\"document_by_connector_credential_pair\",\n )\n" }, { "path": "backend/alembic/versions/800f48024ae9_add_id_to_connectorcredentialpair.py", "content": "\"\"\"Add ID to ConnectorCredentialPair\n\nRevision ID: 800f48024ae9\nRevises: 767f1c2a00eb\nCreate Date: 2023-09-19 16:13:42.299715\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.schema import Sequence, CreateSequence\n\n# revision identifiers, used by Alembic.\nrevision = \"800f48024ae9\"\ndown_revision = \"767f1c2a00eb\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n sequence = Sequence(\"connector_credential_pair_id_seq\")\n op.execute(CreateSequence(sequence)) # type: ignore\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"id\", sa.Integer(), nullable=True, server_default=sequence.next_value()\n ),\n )\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"name\", sa.String(), nullable=True),\n )\n\n # fill in IDs for existing rows\n op.execute(\n \"UPDATE connector_credential_pair SET id = nextval('connector_credential_pair_id_seq') WHERE id IS NULL\"\n )\n op.alter_column(\"connector_credential_pair\", \"id\", nullable=False)\n\n op.create_unique_constraint(\n \"connector_credential_pair__name__key\", \"connector_credential_pair\", [\"name\"]\n )\n op.create_unique_constraint(\n \"connector_credential_pair__id__key\", \"connector_credential_pair\", [\"id\"]\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\n \"connector_credential_pair__name__key\",\n \"connector_credential_pair\",\n type_=\"unique\",\n )\n op.drop_constraint(\n \"connector_credential_pair__id__key\",\n \"connector_credential_pair\",\n type_=\"unique\",\n )\n op.drop_column(\"connector_credential_pair\", \"name\")\n op.drop_column(\"connector_credential_pair\", \"id\")\n op.execute(\"DROP SEQUENCE connector_credential_pair_id_seq\")\n" }, { "path": "backend/alembic/versions/80696cf850ae_add_chat_session_to_query_event.py", "content": "\"\"\"Add chat session to query_event\n\nRevision ID: 80696cf850ae\nRevises: 15326fcec57e\nCreate Date: 2023-11-26 02:38:35.008070\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"80696cf850ae\"\ndown_revision = \"15326fcec57e\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"query_event\",\n sa.Column(\"chat_session_id\", sa.Integer(), nullable=True),\n )\n op.create_foreign_key(\n \"fk_query_event_chat_session_id\",\n \"query_event\",\n \"chat_session\",\n [\"chat_session_id\"],\n [\"id\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\n \"fk_query_event_chat_session_id\", \"query_event\", type_=\"foreignkey\"\n )\n op.drop_column(\"query_event\", \"chat_session_id\")\n" }, { "path": "backend/alembic/versions/8188861f4e92_csv_to_tabular_chat_file_type.py", "content": "\"\"\"csv to tabular chat file type\n\nRevision ID: 8188861f4e92\nRevises: d8cdfee5df80\nCreate Date: 2026-03-31 19:23:05.753184\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"8188861f4e92\"\ndown_revision = \"d8cdfee5df80\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE chat_message\n SET files = (\n SELECT jsonb_agg(\n CASE\n WHEN elem->>'type' = 'csv'\n THEN jsonb_set(elem, '{type}', '\"tabular\"')\n ELSE elem\n END\n )\n FROM jsonb_array_elements(files) AS elem\n )\n WHERE files::text LIKE '%\"type\": \"csv\"%'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE chat_message\n SET files = (\n SELECT jsonb_agg(\n CASE\n WHEN elem->>'type' = 'tabular'\n THEN jsonb_set(elem, '{type}', '\"csv\"')\n ELSE elem\n END\n )\n FROM jsonb_array_elements(files) AS elem\n )\n WHERE files::text LIKE '%\"type\": \"tabular\"%'\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/81c22b1e2e78_hierarchy_nodes_v1.py", "content": "\"\"\"hierarchy_nodes_v1\n\nRevision ID: 81c22b1e2e78\nRevises: 72aa7de2e5cf\nCreate Date: 2026-01-13 18:10:01.021451\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\nfrom onyx.configs.constants import DocumentSource\n\n\n# revision identifiers, used by Alembic.\nrevision = \"81c22b1e2e78\"\ndown_revision = \"72aa7de2e5cf\"\nbranch_labels = None\ndepends_on = None\n\n\n# Human-readable display names for each source\nSOURCE_DISPLAY_NAMES: dict[str, str] = {\n \"ingestion_api\": \"Ingestion API\",\n \"slack\": \"Slack\",\n \"web\": \"Web\",\n \"google_drive\": \"Google Drive\",\n \"gmail\": \"Gmail\",\n \"requesttracker\": \"Request Tracker\",\n \"github\": \"GitHub\",\n \"gitbook\": \"GitBook\",\n \"gitlab\": \"GitLab\",\n \"guru\": \"Guru\",\n \"bookstack\": \"BookStack\",\n \"outline\": \"Outline\",\n \"confluence\": \"Confluence\",\n \"jira\": \"Jira\",\n \"slab\": \"Slab\",\n \"productboard\": \"Productboard\",\n \"file\": \"File\",\n \"coda\": \"Coda\",\n \"notion\": \"Notion\",\n \"zulip\": \"Zulip\",\n \"linear\": \"Linear\",\n \"hubspot\": \"HubSpot\",\n \"document360\": \"Document360\",\n \"gong\": \"Gong\",\n \"google_sites\": \"Google Sites\",\n \"zendesk\": \"Zendesk\",\n \"loopio\": \"Loopio\",\n \"dropbox\": \"Dropbox\",\n \"sharepoint\": \"SharePoint\",\n \"teams\": \"Teams\",\n \"salesforce\": \"Salesforce\",\n \"discourse\": \"Discourse\",\n \"axero\": \"Axero\",\n \"clickup\": \"ClickUp\",\n \"mediawiki\": \"MediaWiki\",\n \"wikipedia\": \"Wikipedia\",\n \"asana\": \"Asana\",\n \"s3\": \"S3\",\n \"r2\": \"R2\",\n \"google_cloud_storage\": \"Google Cloud Storage\",\n \"oci_storage\": \"OCI Storage\",\n \"xenforo\": \"XenForo\",\n \"not_applicable\": \"Not Applicable\",\n \"discord\": \"Discord\",\n \"freshdesk\": \"Freshdesk\",\n \"fireflies\": \"Fireflies\",\n \"egnyte\": \"Egnyte\",\n \"airtable\": \"Airtable\",\n \"highspot\": \"Highspot\",\n \"drupal_wiki\": \"Drupal Wiki\",\n \"imap\": \"IMAP\",\n \"bitbucket\": \"Bitbucket\",\n \"testrail\": \"TestRail\",\n \"mock_connector\": \"Mock Connector\",\n \"user_file\": \"User File\",\n}\n\n\ndef upgrade() -> None:\n # 1. Create hierarchy_node table\n op.create_table(\n \"hierarchy_node\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"raw_node_id\", sa.String(), nullable=False),\n sa.Column(\"display_name\", sa.String(), nullable=False),\n sa.Column(\"link\", sa.String(), nullable=True),\n sa.Column(\"source\", sa.String(), nullable=False),\n sa.Column(\"node_type\", sa.String(), nullable=False),\n sa.Column(\"document_id\", sa.String(), nullable=True),\n sa.Column(\"parent_id\", sa.Integer(), nullable=True),\n # Permission fields - same pattern as Document table\n sa.Column(\n \"external_user_emails\",\n postgresql.ARRAY(sa.String()),\n nullable=True,\n ),\n sa.Column(\n \"external_user_group_ids\",\n postgresql.ARRAY(sa.String()),\n nullable=True,\n ),\n sa.Column(\"is_public\", sa.Boolean(), nullable=False, server_default=\"false\"),\n sa.PrimaryKeyConstraint(\"id\"),\n # When document is deleted, just unlink (node can exist without document)\n sa.ForeignKeyConstraint([\"document_id\"], [\"document.id\"], ondelete=\"SET NULL\"),\n # When parent node is deleted, orphan children (cleanup via pruning)\n sa.ForeignKeyConstraint(\n [\"parent_id\"], [\"hierarchy_node.id\"], ondelete=\"SET NULL\"\n ),\n sa.UniqueConstraint(\n \"raw_node_id\", \"source\", name=\"uq_hierarchy_node_raw_id_source\"\n ),\n )\n op.create_index(\"ix_hierarchy_node_parent_id\", \"hierarchy_node\", [\"parent_id\"])\n op.create_index(\n \"ix_hierarchy_node_source_type\", \"hierarchy_node\", [\"source\", \"node_type\"]\n )\n\n # Add partial unique index to ensure only one SOURCE-type node per source\n # This prevents duplicate source root nodes from being created\n # NOTE: node_type stores enum NAME ('SOURCE'), not value ('source')\n op.execute(\n sa.text(\n \"\"\"\n CREATE UNIQUE INDEX uq_hierarchy_node_one_source_per_type\n ON hierarchy_node (source)\n WHERE node_type = 'SOURCE'\n \"\"\"\n )\n )\n\n # 2. Create hierarchy_fetch_attempt table\n op.create_table(\n \"hierarchy_fetch_attempt\",\n sa.Column(\"id\", postgresql.UUID(as_uuid=True), nullable=False),\n sa.Column(\"connector_credential_pair_id\", sa.Integer(), nullable=False),\n sa.Column(\"status\", sa.String(), nullable=False),\n sa.Column(\"nodes_fetched\", sa.Integer(), nullable=True, server_default=\"0\"),\n sa.Column(\"nodes_updated\", sa.Integer(), nullable=True, server_default=\"0\"),\n sa.Column(\"error_msg\", sa.Text(), nullable=True),\n sa.Column(\"full_exception_trace\", sa.Text(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\"time_started\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.ForeignKeyConstraint(\n [\"connector_credential_pair_id\"],\n [\"connector_credential_pair.id\"],\n ondelete=\"CASCADE\",\n ),\n )\n op.create_index(\n \"ix_hierarchy_fetch_attempt_status\", \"hierarchy_fetch_attempt\", [\"status\"]\n )\n op.create_index(\n \"ix_hierarchy_fetch_attempt_time_created\",\n \"hierarchy_fetch_attempt\",\n [\"time_created\"],\n )\n op.create_index(\n \"ix_hierarchy_fetch_attempt_cc_pair\",\n \"hierarchy_fetch_attempt\",\n [\"connector_credential_pair_id\"],\n )\n\n # 3. Insert SOURCE-type hierarchy nodes for each DocumentSource\n # We insert these so every existing document can have a parent hierarchy node\n # NOTE: SQLAlchemy's Enum with native_enum=False stores the enum NAME (e.g., 'GOOGLE_DRIVE'),\n # not the VALUE (e.g., 'google_drive'). We must use .name for source and node_type columns.\n # SOURCE nodes are always public since they're just categorical roots.\n for source in DocumentSource:\n source_name = (\n source.name\n ) # e.g., 'GOOGLE_DRIVE' - what SQLAlchemy stores/expects\n source_value = source.value # e.g., 'google_drive' - the raw_node_id\n display_name = SOURCE_DISPLAY_NAMES.get(\n source_value, source_value.replace(\"_\", \" \").title()\n )\n op.execute(\n sa.text(\n \"\"\"\n INSERT INTO hierarchy_node (raw_node_id, display_name, source, node_type, parent_id, is_public)\n VALUES (:raw_node_id, :display_name, :source, 'SOURCE', NULL, true)\n ON CONFLICT (raw_node_id, source) DO NOTHING\n \"\"\"\n ).bindparams(\n raw_node_id=source_value, # Use .value for raw_node_id (human-readable identifier)\n display_name=display_name,\n source=source_name, # Use .name for source column (SQLAlchemy enum storage)\n )\n )\n\n # 4. Add parent_hierarchy_node_id column to document table\n op.add_column(\n \"document\",\n sa.Column(\"parent_hierarchy_node_id\", sa.Integer(), nullable=True),\n )\n # When hierarchy node is deleted, just unlink the document (SET NULL)\n op.create_foreign_key(\n \"fk_document_parent_hierarchy_node\",\n \"document\",\n \"hierarchy_node\",\n [\"parent_hierarchy_node_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n op.create_index(\n \"ix_document_parent_hierarchy_node_id\",\n \"document\",\n [\"parent_hierarchy_node_id\"],\n )\n\n # 5. Set all existing documents' parent_hierarchy_node_id to their source's SOURCE node\n # For documents with multiple connectors, we pick one source deterministically (MIN connector_id)\n # NOTE: Both connector.source and hierarchy_node.source store enum NAMEs (e.g., 'GOOGLE_DRIVE')\n # because SQLAlchemy Enum(native_enum=False) uses the enum name for storage.\n op.execute(\n sa.text(\n \"\"\"\n UPDATE document d\n SET parent_hierarchy_node_id = hn.id\n FROM (\n -- Get the source for each document (pick MIN connector_id for determinism)\n SELECT DISTINCT ON (dbcc.id)\n dbcc.id as doc_id,\n c.source as source\n FROM document_by_connector_credential_pair dbcc\n JOIN connector c ON dbcc.connector_id = c.id\n ORDER BY dbcc.id, dbcc.connector_id\n ) doc_source\n JOIN hierarchy_node hn ON hn.source = doc_source.source AND hn.node_type = 'SOURCE'\n WHERE d.id = doc_source.doc_id\n \"\"\"\n )\n )\n\n # Create the persona__hierarchy_node association table\n op.create_table(\n \"persona__hierarchy_node\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"hierarchy_node_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.ForeignKeyConstraint(\n [\"hierarchy_node_id\"],\n [\"hierarchy_node.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"hierarchy_node_id\"),\n )\n\n # Add index for efficient lookups\n op.create_index(\n \"ix_persona__hierarchy_node_hierarchy_node_id\",\n \"persona__hierarchy_node\",\n [\"hierarchy_node_id\"],\n )\n\n # Create the persona__document association table for attaching individual\n # documents directly to assistants\n op.create_table(\n \"persona__document\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"document_id\", sa.String(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.ForeignKeyConstraint(\n [\"document_id\"],\n [\"document.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"document_id\"),\n )\n\n # Add index for efficient lookups by document_id\n op.create_index(\n \"ix_persona__document_document_id\",\n \"persona__document\",\n [\"document_id\"],\n )\n\n # 6. Add last_time_hierarchy_fetch column to connector_credential_pair table\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"last_time_hierarchy_fetch\", sa.DateTime(timezone=True), nullable=True\n ),\n )\n\n\ndef downgrade() -> None:\n # Remove last_time_hierarchy_fetch from connector_credential_pair\n op.drop_column(\"connector_credential_pair\", \"last_time_hierarchy_fetch\")\n\n # Drop persona__document table\n op.drop_index(\"ix_persona__document_document_id\", table_name=\"persona__document\")\n op.drop_table(\"persona__document\")\n\n # Drop persona__hierarchy_node table\n op.drop_index(\n \"ix_persona__hierarchy_node_hierarchy_node_id\",\n table_name=\"persona__hierarchy_node\",\n )\n op.drop_table(\"persona__hierarchy_node\")\n\n # Remove parent_hierarchy_node_id from document\n op.drop_index(\"ix_document_parent_hierarchy_node_id\", table_name=\"document\")\n op.drop_constraint(\n \"fk_document_parent_hierarchy_node\", \"document\", type_=\"foreignkey\"\n )\n op.drop_column(\"document\", \"parent_hierarchy_node_id\")\n\n # Drop hierarchy_fetch_attempt table\n op.drop_index(\n \"ix_hierarchy_fetch_attempt_cc_pair\", table_name=\"hierarchy_fetch_attempt\"\n )\n op.drop_index(\n \"ix_hierarchy_fetch_attempt_time_created\", table_name=\"hierarchy_fetch_attempt\"\n )\n op.drop_index(\n \"ix_hierarchy_fetch_attempt_status\", table_name=\"hierarchy_fetch_attempt\"\n )\n op.drop_table(\"hierarchy_fetch_attempt\")\n\n # Drop hierarchy_node table\n op.drop_index(\"uq_hierarchy_node_one_source_per_type\", table_name=\"hierarchy_node\")\n op.drop_index(\"ix_hierarchy_node_source_type\", table_name=\"hierarchy_node\")\n op.drop_index(\"ix_hierarchy_node_parent_id\", table_name=\"hierarchy_node\")\n op.drop_table(\"hierarchy_node\")\n" }, { "path": "backend/alembic/versions/8405ca81cc83_notifications_constraint.py", "content": "\"\"\"notifications constraint, sort index, and cleanup old notifications\n\nRevision ID: 8405ca81cc83\nRevises: a3c1a7904cd0\nCreate Date: 2026-01-07 16:43:44.855156\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"8405ca81cc83\"\ndown_revision = \"a3c1a7904cd0\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create unique index for notification deduplication.\n # This enables atomic ON CONFLICT DO NOTHING inserts in batch_create_notifications.\n #\n # Uses COALESCE to handle NULL additional_data (NULLs are normally distinct\n # in unique constraints, but we want NULL == NULL for deduplication).\n # The '{}' represents an empty JSONB object as the NULL replacement.\n\n # Clean up legacy notifications first\n op.execute(\"DELETE FROM notification WHERE title = 'New Notification'\")\n\n op.execute(\n \"\"\"\n CREATE UNIQUE INDEX IF NOT EXISTS ix_notification_user_type_data\n ON notification (user_id, notif_type, COALESCE(additional_data, '{}'::jsonb))\n \"\"\"\n )\n\n # Create index for efficient notification sorting by user\n # Covers: WHERE user_id = ? ORDER BY dismissed, first_shown DESC\n op.execute(\n \"\"\"\n CREATE INDEX IF NOT EXISTS ix_notification_user_sort\n ON notification (user_id, dismissed, first_shown DESC)\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\"DROP INDEX IF EXISTS ix_notification_user_type_data\")\n op.execute(\"DROP INDEX IF EXISTS ix_notification_user_sort\")\n" }, { "path": "backend/alembic/versions/849b21c732f8_add_demo_data_enabled_to_build_session.py", "content": "\"\"\"add demo_data_enabled to build_session\n\nRevision ID: 849b21c732f8\nRevises: 81c22b1e2e78\nCreate Date: 2026-01-28 10:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"849b21c732f8\"\ndown_revision = \"81c22b1e2e78\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"build_session\",\n sa.Column(\n \"demo_data_enabled\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.text(\"true\"),\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"build_session\", \"demo_data_enabled\")\n" }, { "path": "backend/alembic/versions/87c52ec39f84_update_default_system_prompt.py", "content": "\"\"\"update_default_system_prompt\n\nRevision ID: 87c52ec39f84\nRevises: 7bd55f264e1b\nCreate Date: 2025-12-05 15:54:06.002452\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"87c52ec39f84\"\ndown_revision = \"7bd55f264e1b\"\nbranch_labels = None\ndepends_on = None\n\n\nDEFAULT_PERSONA_ID = 0\n\n# ruff: noqa: E501, W605 start\nDEFAULT_SYSTEM_PROMPT = \"\"\"\nYou are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent, ask clarifying questions when needed, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. Always prioritize being truthful, nuanced, insightful, and efficient.\n\nThe current date is [[CURRENT_DATETIME]].[[CITATION_GUIDANCE]]\n\n# Response Style\nYou use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make your responses more readable and engaging.\nYou use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, symbols, etc.: '$$\\\\n[expression]\\\\n$$' for standalone cases and '\\\\( [expression] \\\\)' when inline.\nFor code you prefer to use Markdown and specify the language.\nYou can use horizontal rules (---) to separate sections of your responses.\nYou can use Markdown tables to format your responses for data, lists, and other structured information.\n\"\"\".lstrip()\n# ruff: noqa: E501, W605 end\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET system_prompt = :system_prompt\n WHERE id = :persona_id\n \"\"\"\n ),\n {\"system_prompt\": DEFAULT_SYSTEM_PROMPT, \"persona_id\": DEFAULT_PERSONA_ID},\n )\n\n\ndef downgrade() -> None:\n # We don't revert the system prompt on downgrade since we don't know\n # what the previous value was. The new prompt is a reasonable default.\n pass\n" }, { "path": "backend/alembic/versions/8818cf73fa1a_drop_include_citations.py", "content": "\"\"\"drop include citations\n\nRevision ID: 8818cf73fa1a\nRevises: 7ed603b64d5a\nCreate Date: 2025-09-02 19:43:50.060680\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"8818cf73fa1a\"\ndown_revision = \"7ed603b64d5a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"prompt\", \"include_citations\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"prompt\",\n sa.Column(\n \"include_citations\",\n sa.BOOLEAN(),\n autoincrement=False,\n nullable=True,\n ),\n )\n # Set include_citations based on prompt name: FALSE for ImageGeneration, TRUE for others\n op.execute(\n sa.text(\n \"UPDATE prompt SET include_citations = CASE WHEN name = 'ImageGeneration' THEN FALSE ELSE TRUE END\"\n )\n )\n" }, { "path": "backend/alembic/versions/891cd83c87a8_add_is_visible_to_persona.py", "content": "\"\"\"Add is_visible to Persona\n\nRevision ID: 891cd83c87a8\nRevises: 76b60d407dfb\nCreate Date: 2023-12-21 11:55:54.132279\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"891cd83c87a8\"\ndown_revision = \"76b60d407dfb\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"is_visible\", sa.Boolean(), nullable=True),\n )\n op.execute(\"UPDATE persona SET is_visible = true\")\n op.alter_column(\"persona\", \"is_visible\", nullable=False)\n\n op.add_column(\n \"persona\",\n sa.Column(\"display_priority\", sa.Integer(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"is_visible\")\n op.drop_column(\"persona\", \"display_priority\")\n" }, { "path": "backend/alembic/versions/8987770549c0_add_full_exception_stack_trace.py", "content": "\"\"\"Add full exception stack trace\n\nRevision ID: 8987770549c0\nRevises: ec3ec2eabf7b\nCreate Date: 2024-02-10 19:31:28.339135\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"8987770549c0\"\ndown_revision = \"ec3ec2eabf7b\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"index_attempt\", sa.Column(\"full_exception_trace\", sa.Text(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"index_attempt\", \"full_exception_trace\")\n" }, { "path": "backend/alembic/versions/8a87bd6ec550_associate_index_attempts_with_ccpair.py", "content": "\"\"\"associate index attempts with ccpair\n\nRevision ID: 8a87bd6ec550\nRevises: 4ea2c93919c1\nCreate Date: 2024-07-22 15:15:52.558451\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"8a87bd6ec550\"\ndown_revision = \"4ea2c93919c1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Add the new connector_credential_pair_id column\n op.add_column(\n \"index_attempt\",\n sa.Column(\"connector_credential_pair_id\", sa.Integer(), nullable=True),\n )\n\n # Create a foreign key constraint to the connector_credential_pair table\n op.create_foreign_key(\n \"fk_index_attempt_connector_credential_pair_id\",\n \"index_attempt\",\n \"connector_credential_pair\",\n [\"connector_credential_pair_id\"],\n [\"id\"],\n )\n\n # Populate the new connector_credential_pair_id column using existing connector_id and credential_id\n op.execute(\n \"\"\"\n UPDATE index_attempt ia\n SET connector_credential_pair_id = (\n SELECT id FROM connector_credential_pair ccp\n WHERE\n (ia.connector_id IS NULL OR ccp.connector_id = ia.connector_id)\n AND (ia.credential_id IS NULL OR ccp.credential_id = ia.credential_id)\n LIMIT 1\n )\n WHERE ia.connector_id IS NOT NULL OR ia.credential_id IS NOT NULL\n \"\"\"\n )\n\n # For good measure\n op.execute(\n \"\"\"\n DELETE FROM index_attempt\n WHERE connector_credential_pair_id IS NULL\n \"\"\"\n )\n\n # Make the new connector_credential_pair_id column non-nullable\n op.alter_column(\"index_attempt\", \"connector_credential_pair_id\", nullable=False)\n\n # Drop the old connector_id and credential_id columns\n op.drop_column(\"index_attempt\", \"connector_id\")\n op.drop_column(\"index_attempt\", \"credential_id\")\n\n # Update the index to use connector_credential_pair_id\n op.create_index(\n \"ix_index_attempt_latest_for_connector_credential_pair\",\n \"index_attempt\",\n [\"connector_credential_pair_id\", \"time_created\"],\n )\n\n\ndef downgrade() -> None:\n # Add back the old connector_id and credential_id columns\n op.add_column(\n \"index_attempt\", sa.Column(\"connector_id\", sa.Integer(), nullable=True)\n )\n op.add_column(\n \"index_attempt\", sa.Column(\"credential_id\", sa.Integer(), nullable=True)\n )\n\n # Populate the old connector_id and credential_id columns using the connector_credential_pair_id\n op.execute(\n \"\"\"\n UPDATE index_attempt ia\n SET connector_id = ccp.connector_id, credential_id = ccp.credential_id\n FROM connector_credential_pair ccp\n WHERE ia.connector_credential_pair_id = ccp.id\n \"\"\"\n )\n\n # Make the old connector_id and credential_id columns non-nullable\n op.alter_column(\"index_attempt\", \"connector_id\", nullable=False)\n op.alter_column(\"index_attempt\", \"credential_id\", nullable=False)\n\n # Drop the new connector_credential_pair_id column\n op.drop_constraint(\n \"fk_index_attempt_connector_credential_pair_id\",\n \"index_attempt\",\n type_=\"foreignkey\",\n )\n op.drop_column(\"index_attempt\", \"connector_credential_pair_id\")\n\n op.create_index(\n \"ix_index_attempt_latest_for_connector_credential_pair\",\n \"index_attempt\",\n [\"connector_id\", \"credential_id\", \"time_created\"],\n )\n" }, { "path": "backend/alembic/versions/8aabb57f3b49_restructure_document_indices.py", "content": "\"\"\"Restructure Document Indices\n\nRevision ID: 8aabb57f3b49\nRevises: 5e84129c8be3\nCreate Date: 2023-08-18 21:15:57.629515\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"8aabb57f3b49\"\ndown_revision = \"5e84129c8be3\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_table(\"chunk\")\n op.execute(\"DROP TYPE IF EXISTS documentstoretype\")\n\n\ndef downgrade() -> None:\n op.create_table(\n \"chunk\",\n sa.Column(\"id\", sa.VARCHAR(), autoincrement=False, nullable=False),\n sa.Column(\n \"document_store_type\",\n postgresql.ENUM(\"VECTOR\", \"KEYWORD\", name=\"documentstoretype\"),\n autoincrement=False,\n nullable=False,\n ),\n sa.Column(\"document_id\", sa.VARCHAR(), autoincrement=False, nullable=False),\n sa.ForeignKeyConstraint(\n [\"document_id\"], [\"document.id\"], name=\"chunk_document_id_fkey\"\n ),\n sa.PrimaryKeyConstraint(\"id\", \"document_store_type\", name=\"chunk_pkey\"),\n )\n" }, { "path": "backend/alembic/versions/8b5ce697290e_add_discord_bot_tables.py", "content": "\"\"\"Add Discord bot tables\n\nRevision ID: 8b5ce697290e\nRevises: a1b2c3d4e5f7\nCreate Date: 2025-01-14\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"8b5ce697290e\"\ndown_revision = \"a1b2c3d4e5f7\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # DiscordBotConfig (singleton table - one per tenant)\n op.create_table(\n \"discord_bot_config\",\n sa.Column(\n \"id\",\n sa.String(),\n primary_key=True,\n server_default=sa.text(\"'SINGLETON'\"),\n ),\n sa.Column(\"bot_token\", sa.LargeBinary(), nullable=False), # EncryptedString\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.CheckConstraint(\"id = 'SINGLETON'\", name=\"ck_discord_bot_config_singleton\"),\n )\n\n # DiscordGuildConfig\n op.create_table(\n \"discord_guild_config\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"guild_id\", sa.BigInteger(), nullable=True, unique=True),\n sa.Column(\"guild_name\", sa.String(), nullable=True),\n sa.Column(\"registration_key\", sa.String(), nullable=False, unique=True),\n sa.Column(\"registered_at\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"default_persona_id\",\n sa.Integer(),\n sa.ForeignKey(\"persona.id\", ondelete=\"SET NULL\"),\n nullable=True,\n ),\n sa.Column(\n \"enabled\", sa.Boolean(), server_default=sa.text(\"true\"), nullable=False\n ),\n )\n\n # DiscordChannelConfig\n op.create_table(\n \"discord_channel_config\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\n \"guild_config_id\",\n sa.Integer(),\n sa.ForeignKey(\"discord_guild_config.id\", ondelete=\"CASCADE\"),\n nullable=False,\n ),\n sa.Column(\"channel_id\", sa.BigInteger(), nullable=False),\n sa.Column(\"channel_name\", sa.String(), nullable=False),\n sa.Column(\n \"channel_type\",\n sa.String(20),\n server_default=sa.text(\"'text'\"),\n nullable=False,\n ),\n sa.Column(\n \"is_private\",\n sa.Boolean(),\n server_default=sa.text(\"false\"),\n nullable=False,\n ),\n sa.Column(\n \"thread_only_mode\",\n sa.Boolean(),\n server_default=sa.text(\"false\"),\n nullable=False,\n ),\n sa.Column(\n \"require_bot_invocation\",\n sa.Boolean(),\n server_default=sa.text(\"true\"),\n nullable=False,\n ),\n sa.Column(\n \"persona_override_id\",\n sa.Integer(),\n sa.ForeignKey(\"persona.id\", ondelete=\"SET NULL\"),\n nullable=True,\n ),\n sa.Column(\n \"enabled\", sa.Boolean(), server_default=sa.text(\"false\"), nullable=False\n ),\n )\n\n # Unique constraint: one config per channel per guild\n op.create_unique_constraint(\n \"uq_discord_channel_guild_channel\",\n \"discord_channel_config\",\n [\"guild_config_id\", \"channel_id\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"discord_channel_config\")\n op.drop_table(\"discord_guild_config\")\n op.drop_table(\"discord_bot_config\")\n" }, { "path": "backend/alembic/versions/8e1ac4f39a9f_enable_contextual_retrieval.py", "content": "\"\"\"enable contextual retrieval\n\nRevision ID: 8e1ac4f39a9f\nRevises: 9aadf32dfeb4\nCreate Date: 2024-12-20 13:29:09.918661\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"8e1ac4f39a9f\"\ndown_revision = \"9aadf32dfeb4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"enable_contextual_rag\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"contextual_rag_llm_name\",\n sa.String(),\n nullable=True,\n ),\n )\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"contextual_rag_llm_provider\",\n sa.String(),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"search_settings\", \"enable_contextual_rag\")\n op.drop_column(\"search_settings\", \"contextual_rag_llm_name\")\n op.drop_column(\"search_settings\", \"contextual_rag_llm_provider\")\n" }, { "path": "backend/alembic/versions/8e26726b7683_chat_context_addition.py", "content": "\"\"\"Chat Context Addition\n\nRevision ID: 8e26726b7683\nRevises: 5809c0787398\nCreate Date: 2023-09-13 18:34:31.327944\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"8e26726b7683\"\ndown_revision = \"5809c0787398\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"persona\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"system_text\", sa.Text(), nullable=True),\n sa.Column(\"tools_text\", sa.Text(), nullable=True),\n sa.Column(\"hint_text\", sa.Text(), nullable=True),\n sa.Column(\"default_persona\", sa.Boolean(), nullable=False),\n sa.Column(\"deleted\", sa.Boolean(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.add_column(\"chat_message\", sa.Column(\"persona_id\", sa.Integer(), nullable=True))\n op.create_foreign_key(\n \"fk_chat_message_persona_id\", \"chat_message\", \"persona\", [\"persona_id\"], [\"id\"]\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\"fk_chat_message_persona_id\", \"chat_message\", type_=\"foreignkey\")\n op.drop_column(\"chat_message\", \"persona_id\")\n op.drop_table(\"persona\")\n" }, { "path": "backend/alembic/versions/8f43500ee275_add_index.py", "content": "\"\"\"add index\n\nRevision ID: 8f43500ee275\nRevises: da42808081e3\nCreate Date: 2025-02-24 17:35:33.072714\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"8f43500ee275\"\ndown_revision = \"da42808081e3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create a basic index on the lowercase message column for direct text matching\n # Limit to 1500 characters to stay well under the 2856 byte limit of btree version 4\n # op.execute(\n # \"\"\"\n # CREATE INDEX idx_chat_message_message_lower\n # ON chat_message (LOWER(substring(message, 1, 1500)))\n # \"\"\"\n # )\n pass\n\n\ndef downgrade() -> None:\n # Drop the index\n op.execute(\"DROP INDEX IF EXISTS idx_chat_message_message_lower;\")\n" }, { "path": "backend/alembic/versions/8ffcc2bcfc11_add_needs_persona_sync_to_user_file.py", "content": "\"\"\"add needs_persona_sync to user_file\n\nRevision ID: 8ffcc2bcfc11\nRevises: 7616121f6e97\nCreate Date: 2026-02-23 10:48:48.343826\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"8ffcc2bcfc11\"\ndown_revision = \"7616121f6e97\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user_file\",\n sa.Column(\n \"needs_persona_sync\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.text(\"false\"),\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user_file\", \"needs_persona_sync\")\n" }, { "path": "backend/alembic/versions/904451035c9b_store_tool_details.py", "content": "\"\"\"Store Tool Details\n\nRevision ID: 904451035c9b\nRevises: 3b25685ff73c\nCreate Date: 2023-10-05 12:29:26.620000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"904451035c9b\"\ndown_revision = \"3b25685ff73c\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"tools\", postgresql.JSONB(astext_type=sa.Text()), nullable=True),\n )\n op.drop_column(\"persona\", \"tools_text\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"tools_text\", sa.TEXT(), autoincrement=False, nullable=True),\n )\n op.drop_column(\"persona\", \"tools\")\n" }, { "path": "backend/alembic/versions/904e5138fffb_tags.py", "content": "\"\"\"Tags\n\nRevision ID: 904e5138fffb\nRevises: 891cd83c87a8\nCreate Date: 2024-01-01 10:44:43.733974\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"904e5138fffb\"\ndown_revision = \"891cd83c87a8\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"tag\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"tag_key\", sa.String(), nullable=False),\n sa.Column(\"tag_value\", sa.String(), nullable=False),\n sa.Column(\"source\", sa.String(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\n \"tag_key\", \"tag_value\", \"source\", name=\"_tag_key_value_source_uc\"\n ),\n )\n op.create_table(\n \"document__tag\",\n sa.Column(\"document_id\", sa.String(), nullable=False),\n sa.Column(\"tag_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"document_id\"],\n [\"document.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"tag_id\"],\n [\"tag.id\"],\n ),\n sa.PrimaryKeyConstraint(\"document_id\", \"tag_id\"),\n )\n\n op.add_column(\n \"search_doc\",\n sa.Column(\n \"doc_metadata\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n op.execute(\"UPDATE search_doc SET doc_metadata = '{}' WHERE doc_metadata IS NULL\")\n op.alter_column(\"search_doc\", \"doc_metadata\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_table(\"document__tag\")\n op.drop_table(\"tag\")\n op.drop_column(\"search_doc\", \"doc_metadata\")\n" }, { "path": "backend/alembic/versions/9087b548dd69_seed_default_image_gen_config.py", "content": "\"\"\"seed_default_image_gen_config\n\nRevision ID: 9087b548dd69\nRevises: 2b90f3af54b8\nCreate Date: 2026-01-05 00:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"9087b548dd69\"\ndown_revision = \"2b90f3af54b8\"\nbranch_labels = None\ndepends_on = None\n\n# Constants for default image generation config\n# Source: web/src/app/admin/configuration/image-generation/constants.ts\nIMAGE_PROVIDER_ID = \"openai_gpt_image_1\"\nMODEL_NAME = \"gpt-image-1\"\nPROVIDER_NAME = \"openai\"\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Check if image_generation_config table already has records\n existing_configs = (\n conn.execute(sa.text(\"SELECT COUNT(*) FROM image_generation_config\")).scalar()\n or 0\n )\n\n if existing_configs > 0:\n # Skip if configs already exist - user may have configured manually\n return\n\n # Find the first OpenAI LLM provider\n openai_provider = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, api_key\n FROM llm_provider\n WHERE provider = :provider\n ORDER BY id\n LIMIT 1\n \"\"\"\n ),\n {\"provider\": PROVIDER_NAME},\n ).fetchone()\n\n if not openai_provider:\n # No OpenAI provider found - nothing to do\n return\n\n source_provider_id, api_key = openai_provider\n\n # Create new LLM provider for image generation (clone only api_key)\n result = conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO llm_provider (\n name, provider, api_key, api_base, api_version,\n deployment_name, default_model_name, is_public,\n is_default_provider, is_default_vision_provider, is_auto_mode\n )\n VALUES (\n :name, :provider, :api_key, NULL, NULL,\n NULL, :default_model_name, :is_public,\n NULL, NULL, :is_auto_mode\n )\n RETURNING id\n \"\"\"\n ),\n {\n \"name\": f\"Image Gen - {IMAGE_PROVIDER_ID}\",\n \"provider\": PROVIDER_NAME,\n \"api_key\": api_key,\n \"default_model_name\": MODEL_NAME,\n \"is_public\": True,\n \"is_auto_mode\": False,\n },\n )\n new_provider_id = result.scalar()\n\n # Create model configuration\n result = conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO model_configuration (\n llm_provider_id, name, is_visible, max_input_tokens,\n supports_image_input, display_name\n )\n VALUES (\n :llm_provider_id, :name, :is_visible, :max_input_tokens,\n :supports_image_input, :display_name\n )\n RETURNING id\n \"\"\"\n ),\n {\n \"llm_provider_id\": new_provider_id,\n \"name\": MODEL_NAME,\n \"is_visible\": True,\n \"max_input_tokens\": None,\n \"supports_image_input\": False,\n \"display_name\": None,\n },\n )\n model_config_id = result.scalar()\n\n # Create image generation config\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO image_generation_config (\n image_provider_id, model_configuration_id, is_default\n )\n VALUES (\n :image_provider_id, :model_configuration_id, :is_default\n )\n \"\"\"\n ),\n {\n \"image_provider_id\": IMAGE_PROVIDER_ID,\n \"model_configuration_id\": model_config_id,\n \"is_default\": True,\n },\n )\n\n\ndef downgrade() -> None:\n # We don't remove the config on downgrade since it's safe to keep around\n # If we upgrade again, it will be a no-op due to the existing records check\n pass\n" }, { "path": "backend/alembic/versions/90b409d06e50_add_chat_compression_fields.py", "content": "\"\"\"add_chat_compression_fields\n\nRevision ID: 90b409d06e50\nRevises: f220515df7b4\nCreate Date: 2026-01-26 09:13:09.635427\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"90b409d06e50\"\ndown_revision = \"f220515df7b4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add last_summarized_message_id to chat_message\n # This field marks a message as a summary and indicates the last message it covers.\n # Summaries are branch-aware via their parent_message_id pointing to the branch.\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"last_summarized_message_id\",\n sa.Integer(),\n sa.ForeignKey(\"chat_message.id\", ondelete=\"SET NULL\"),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"last_summarized_message_id\")\n" }, { "path": "backend/alembic/versions/90e3b9af7da4_tag_fix.py", "content": "\"\"\"tag-fix\n\nRevision ID: 90e3b9af7da4\nRevises: 62c3a055a141\nCreate Date: 2025-08-01 20:58:14.607624\n\n\"\"\"\n\nimport json\nimport logging\nimport os\n\nfrom typing import cast\nfrom typing import Generator\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.db.search_settings import SearchSettings\nfrom onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.constants import AuthType\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n\n\n# revision identifiers, used by Alembic.\nrevision = \"90e3b9af7da4\"\ndown_revision = \"62c3a055a141\"\nbranch_labels = None\ndepends_on = None\n\nSKIP_TAG_FIX = os.environ.get(\"SKIP_TAG_FIX\", \"true\").lower() == \"true\"\n\n# override for cloud\nif AUTH_TYPE == AuthType.CLOUD:\n SKIP_TAG_FIX = True\n\n\ndef set_is_list_for_known_tags() -> None:\n \"\"\"\n Sets is_list to true for all tags that are known to be lists.\n \"\"\"\n LIST_METADATA: list[tuple[str, str]] = [\n (\"CLICKUP\", \"tags\"),\n (\"CONFLUENCE\", \"labels\"),\n (\"DISCOURSE\", \"tags\"),\n (\"FRESHDESK\", \"emails\"),\n (\"GITHUB\", \"assignees\"),\n (\"GITHUB\", \"labels\"),\n (\"GURU\", \"tags\"),\n (\"GURU\", \"folders\"),\n (\"HUBSPOT\", \"associated_contact_ids\"),\n (\"HUBSPOT\", \"associated_company_ids\"),\n (\"HUBSPOT\", \"associated_deal_ids\"),\n (\"HUBSPOT\", \"associated_ticket_ids\"),\n (\"JIRA\", \"labels\"),\n (\"MEDIAWIKI\", \"categories\"),\n (\"ZENDESK\", \"labels\"),\n (\"ZENDESK\", \"content_tags\"),\n ]\n\n bind = op.get_bind()\n for source, key in LIST_METADATA:\n bind.execute(\n sa.text(\n f\"\"\"\n UPDATE tag\n SET is_list = true\n WHERE tag_key = '{key}'\n AND source = '{source}'\n \"\"\"\n )\n )\n\n\ndef set_is_list_for_list_tags() -> None:\n \"\"\"\n Sets is_list to true for all tags which have multiple values for a given\n document, key, and source triplet. This only works if we remove old tags\n from the database.\n \"\"\"\n bind = op.get_bind()\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE tag\n SET is_list = true\n FROM (\n SELECT DISTINCT tag.tag_key, tag.source\n FROM tag\n JOIN document__tag ON tag.id = document__tag.tag_id\n GROUP BY tag.tag_key, tag.source, document__tag.document_id\n HAVING count(*) > 1\n ) AS list_tags\n WHERE tag.tag_key = list_tags.tag_key\n AND tag.source = list_tags.source\n \"\"\"\n )\n )\n\n\ndef log_list_tags() -> None:\n bind = op.get_bind()\n result = bind.execute(\n sa.text(\n \"\"\"\n SELECT DISTINCT source, tag_key\n FROM tag\n WHERE is_list\n ORDER BY source, tag_key\n \"\"\"\n )\n ).fetchall()\n logger.info(\n \"List tags:\\n\" + \"\\n\".join(f\" {source}: {key}\" for source, key in result)\n )\n\n\ndef remove_old_tags() -> None:\n \"\"\"\n Removes old tags from the database.\n Previously, there was a bug where if a document got indexed with a tag and then\n the document got reindexed, the old tag would not be removed.\n This function removes those old tags by comparing it against the tags in vespa.\n \"\"\"\n current_search_settings, _ = active_search_settings()\n\n # Get the index name\n if hasattr(current_search_settings, \"index_name\"):\n index_name = current_search_settings.index_name\n else:\n # Default index name if we can't get it from the document_index\n index_name = \"danswer_index\"\n\n for batch in _get_batch_documents_with_multiple_tags():\n n_deleted = 0\n\n for document_id in batch:\n true_metadata = _get_vespa_metadata(document_id, index_name)\n tags = _get_document_tags(document_id)\n\n # identify document__tags to delete\n to_delete: list[str] = []\n for tag_id, tag_key, tag_value in tags:\n true_val = true_metadata.get(tag_key, \"\")\n if (isinstance(true_val, list) and tag_value not in true_val) or (\n isinstance(true_val, str) and tag_value != true_val\n ):\n to_delete.append(str(tag_id))\n\n if not to_delete:\n continue\n\n # delete old document__tags\n bind = op.get_bind()\n result = bind.execute(\n sa.text(\n f\"\"\"\n DELETE FROM document__tag\n WHERE document_id = '{document_id}'\n AND tag_id IN ({\",\".join(to_delete)})\n \"\"\"\n )\n )\n n_deleted += result.rowcount\n logger.info(f\"Processed {len(batch)} documents and deleted {n_deleted} tags\")\n\n\ndef active_search_settings() -> tuple[SearchSettings, SearchSettings | None]:\n result = op.get_bind().execute(\n sa.text(\n \"\"\"\n SELECT * FROM search_settings WHERE status = 'PRESENT' ORDER BY id DESC LIMIT 1\n \"\"\"\n )\n )\n search_settings_fetch = result.fetchall()\n search_settings = (\n SearchSettings(**search_settings_fetch[0]._asdict())\n if search_settings_fetch\n else None\n )\n\n result2 = op.get_bind().execute(\n sa.text(\n \"\"\"\n SELECT * FROM search_settings WHERE status = 'FUTURE' ORDER BY id DESC LIMIT 1\n \"\"\"\n )\n )\n search_settings_future_fetch = result2.fetchall()\n search_settings_future = (\n SearchSettings(**search_settings_future_fetch[0]._asdict())\n if search_settings_future_fetch\n else None\n )\n\n if not isinstance(search_settings, SearchSettings):\n raise RuntimeError(\n \"current search settings is of type \" + str(type(search_settings))\n )\n if (\n not isinstance(search_settings_future, SearchSettings)\n and search_settings_future is not None\n ):\n raise RuntimeError(\n \"future search settings is of type \" + str(type(search_settings_future))\n )\n\n return search_settings, search_settings_future\n\n\ndef _get_batch_documents_with_multiple_tags(\n batch_size: int = 128,\n) -> Generator[list[str], None, None]:\n \"\"\"\n Returns a list of document ids which contain a one to many tag.\n The document may either contain a list metadata value, or may contain leftover\n old tags from reindexing.\n \"\"\"\n offset_clause = \"\"\n bind = op.get_bind()\n\n while True:\n batch = bind.execute(\n sa.text(\n f\"\"\"\n SELECT DISTINCT document__tag.document_id\n FROM tag\n JOIN document__tag ON tag.id = document__tag.tag_id\n GROUP BY tag.tag_key, tag.source, document__tag.document_id\n HAVING count(*) > 1 {offset_clause}\n ORDER BY document__tag.document_id\n LIMIT {batch_size}\n \"\"\"\n )\n ).fetchall()\n if not batch:\n break\n doc_ids = [document_id for (document_id,) in batch]\n yield doc_ids\n offset_clause = f\"AND document__tag.document_id > '{doc_ids[-1]}'\"\n\n\ndef _get_vespa_metadata(\n document_id: str, index_name: str\n) -> dict[str, str | list[str]]:\n url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)\n\n # Document-Selector language\n selection = (\n f\"{index_name}.document_id=='{document_id}' and {index_name}.chunk_id==0\"\n )\n\n params: dict[str, str | int] = {\n \"selection\": selection,\n \"wantedDocumentCount\": 1,\n \"fieldSet\": f\"{index_name}:metadata\",\n }\n\n with get_vespa_http_client() as client:\n resp = client.get(url, params=params)\n resp.raise_for_status()\n\n docs = resp.json().get(\"documents\", [])\n if not docs:\n raise RuntimeError(f\"No chunk-0 found for document {document_id}\")\n\n # for some reason, metadata is a string\n metadata = docs[0][\"fields\"][\"metadata\"]\n return json.loads(metadata)\n\n\ndef _get_document_tags(document_id: str) -> list[tuple[int, str, str]]:\n bind = op.get_bind()\n result = bind.execute(\n sa.text(\n f\"\"\"\n SELECT tag.id, tag.tag_key, tag.tag_value\n FROM tag\n JOIN document__tag ON tag.id = document__tag.tag_id\n WHERE document__tag.document_id = '{document_id}'\n \"\"\"\n )\n ).fetchall()\n return cast(list[tuple[int, str, str]], result)\n\n\ndef upgrade() -> None:\n op.add_column(\n \"tag\",\n sa.Column(\"is_list\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n op.drop_constraint(\n constraint_name=\"_tag_key_value_source_uc\",\n table_name=\"tag\",\n type_=\"unique\",\n )\n op.create_unique_constraint(\n constraint_name=\"_tag_key_value_source_list_uc\",\n table_name=\"tag\",\n columns=[\"tag_key\", \"tag_value\", \"source\", \"is_list\"],\n )\n set_is_list_for_known_tags()\n\n if SKIP_TAG_FIX:\n logger.warning(\n \"Skipping removal of old tags. \"\n \"This can cause issues when using the knowledge graph, or \"\n \"when filtering for documents by tags.\"\n )\n log_list_tags()\n return\n\n remove_old_tags()\n set_is_list_for_list_tags()\n\n # debug\n log_list_tags()\n\n\ndef downgrade() -> None:\n # the migration adds and populates the is_list column, and removes old bugged tags\n # there isn't a point in adding back the bugged tags, so we just drop the column\n op.drop_constraint(\n constraint_name=\"_tag_key_value_source_list_uc\",\n table_name=\"tag\",\n type_=\"unique\",\n )\n op.create_unique_constraint(\n constraint_name=\"_tag_key_value_source_uc\",\n table_name=\"tag\",\n columns=[\"tag_key\", \"tag_value\", \"source\"],\n )\n op.drop_column(\"tag\", \"is_list\")\n" }, { "path": "backend/alembic/versions/91a0a4d62b14_milestone.py", "content": "\"\"\"Milestone\n\nRevision ID: 91a0a4d62b14\nRevises: dab04867cd88\nCreate Date: 2024-12-13 19:03:30.947551\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport fastapi_users_db_sqlalchemy\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"91a0a4d62b14\"\ndown_revision = \"dab04867cd88\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"milestone\",\n sa.Column(\"id\", sa.UUID(), nullable=False),\n sa.Column(\"tenant_id\", sa.String(), nullable=True),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"event_type\", sa.String(), nullable=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"event_tracker\", postgresql.JSONB(), nullable=True),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"event_type\", name=\"uq_milestone_event_type\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"milestone\")\n" }, { "path": "backend/alembic/versions/91fd3b470d1a_remove_documentsource_from_tag.py", "content": "\"\"\"Remove DocumentSource from Tag\n\nRevision ID: 91fd3b470d1a\nRevises: 173cae5bba26\nCreate Date: 2024-03-21 12:05:23.956734\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom onyx.configs.constants import DocumentSource\n\n# revision identifiers, used by Alembic.\nrevision = \"91fd3b470d1a\"\ndown_revision = \"173cae5bba26\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"tag\",\n \"source\",\n type_=sa.String(length=50),\n existing_type=sa.Enum(DocumentSource, native_enum=False),\n existing_nullable=False,\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"tag\",\n \"source\",\n type_=sa.Enum(DocumentSource, native_enum=False),\n existing_type=sa.String(length=50),\n existing_nullable=False,\n )\n" }, { "path": "backend/alembic/versions/91ffac7e65b3_add_expiry_time.py", "content": "\"\"\"add expiry time\n\nRevision ID: 91ffac7e65b3\nRevises: bc9771dccadf\nCreate Date: 2024-06-24 09:39:56.462242\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"91ffac7e65b3\"\ndown_revision = \"795b20b85b4b\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\", sa.Column(\"oidc_expiry\", sa.DateTime(timezone=True), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"oidc_expiry\")\n" }, { "path": "backend/alembic/versions/93560ba1b118_add_web_ui_option_to_slack_config.py", "content": "\"\"\"add web ui option to slack config\n\nRevision ID: 93560ba1b118\nRevises: 6d562f86c78b\nCreate Date: 2024-11-24 06:36:17.490612\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"93560ba1b118\"\ndown_revision = \"6d562f86c78b\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add show_continue_in_web_ui with default False to all existing channel_configs\n op.execute(\n \"\"\"\n UPDATE slack_channel_config\n SET channel_config = channel_config || '{\"show_continue_in_web_ui\": false}'::jsonb\n WHERE NOT channel_config ? 'show_continue_in_web_ui'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n # Remove show_continue_in_web_ui from all channel_configs\n op.execute(\n \"\"\"\n UPDATE slack_channel_config\n SET channel_config = channel_config - 'show_continue_in_web_ui'\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/93a2e195e25c_add_voice_provider_and_user_voice_prefs.py", "content": "\"\"\"add_voice_provider_and_user_voice_prefs\n\nRevision ID: 93a2e195e25c\nRevises: 27fb147a843f\nCreate Date: 2026-02-23 15:16:39.507304\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy import column\nfrom sqlalchemy import true\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"93a2e195e25c\"\ndown_revision = \"27fb147a843f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create voice_provider table\n op.create_table(\n \"voice_provider\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"name\", sa.String(), unique=True, nullable=False),\n sa.Column(\"provider_type\", sa.String(), nullable=False),\n sa.Column(\"api_key\", sa.LargeBinary(), nullable=True),\n sa.Column(\"api_base\", sa.String(), nullable=True),\n sa.Column(\"custom_config\", postgresql.JSONB(), nullable=True),\n sa.Column(\"stt_model\", sa.String(), nullable=True),\n sa.Column(\"tts_model\", sa.String(), nullable=True),\n sa.Column(\"default_voice\", sa.String(), nullable=True),\n sa.Column(\n \"is_default_stt\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n sa.Column(\n \"is_default_tts\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n sa.Column(\"deleted\", sa.Boolean(), nullable=False, server_default=\"false\"),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n onupdate=sa.func.now(),\n nullable=False,\n ),\n )\n\n # Add partial unique indexes to enforce only one default STT/TTS provider\n op.create_index(\n \"ix_voice_provider_one_default_stt\",\n \"voice_provider\",\n [\"is_default_stt\"],\n unique=True,\n postgresql_where=column(\"is_default_stt\") == true(),\n )\n op.create_index(\n \"ix_voice_provider_one_default_tts\",\n \"voice_provider\",\n [\"is_default_tts\"],\n unique=True,\n postgresql_where=column(\"is_default_tts\") == true(),\n )\n\n # Add voice preference columns to user table\n op.add_column(\n \"user\",\n sa.Column(\n \"voice_auto_send\",\n sa.Boolean(),\n default=False,\n nullable=False,\n server_default=\"false\",\n ),\n )\n op.add_column(\n \"user\",\n sa.Column(\n \"voice_auto_playback\",\n sa.Boolean(),\n default=False,\n nullable=False,\n server_default=\"false\",\n ),\n )\n op.add_column(\n \"user\",\n sa.Column(\n \"voice_playback_speed\",\n sa.Float(),\n default=1.0,\n nullable=False,\n server_default=\"1.0\",\n ),\n )\n\n\ndef downgrade() -> None:\n # Remove user voice preference columns\n op.drop_column(\"user\", \"voice_playback_speed\")\n op.drop_column(\"user\", \"voice_auto_playback\")\n op.drop_column(\"user\", \"voice_auto_send\")\n\n op.drop_index(\"ix_voice_provider_one_default_tts\", table_name=\"voice_provider\")\n op.drop_index(\"ix_voice_provider_one_default_stt\", table_name=\"voice_provider\")\n\n # Drop voice_provider table\n op.drop_table(\"voice_provider\")\n" }, { "path": "backend/alembic/versions/93c15d6a6fbb_add_chunk_error_and_vespa_count_columns_.py", "content": "\"\"\"add chunk error and vespa count columns to opensearch tenant migration\n\nRevision ID: 93c15d6a6fbb\nRevises: d3fd499c829c\nCreate Date: 2026-02-11 23:07:34.576725\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"93c15d6a6fbb\"\ndown_revision = \"d3fd499c829c\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"total_chunks_errored\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n )\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"total_chunks_in_vespa\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"opensearch_tenant_migration_record\", \"total_chunks_in_vespa\")\n op.drop_column(\"opensearch_tenant_migration_record\", \"total_chunks_errored\")\n" }, { "path": "backend/alembic/versions/949b4a92a401_remove_rt.py", "content": "\"\"\"remove rt\n\nRevision ID: 949b4a92a401\nRevises: 1b10e1fda030\nCreate Date: 2024-10-26 13:06:06.937969\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy import text\n\n# Import your models and constants\nfrom onyx.db.models import (\n Connector,\n ConnectorCredentialPair,\n Credential,\n IndexAttempt,\n)\n\n\n# revision identifiers, used by Alembic.\nrevision = \"949b4a92a401\"\ndown_revision = \"1b10e1fda030\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Deletes all RequestTracker connectors and associated data\n bind = op.get_bind()\n session = Session(bind=bind)\n\n # Get connectors using raw SQL\n result = bind.execute(\n text(\"SELECT id FROM connector WHERE source = 'requesttracker'\")\n )\n connector_ids = [row[0] for row in result]\n\n if connector_ids:\n cc_pairs_to_delete = (\n session.query(ConnectorCredentialPair)\n .filter(ConnectorCredentialPair.connector_id.in_(connector_ids))\n .all()\n )\n\n cc_pair_ids = [cc_pair.id for cc_pair in cc_pairs_to_delete]\n\n if cc_pair_ids:\n session.query(IndexAttempt).filter(\n IndexAttempt.connector_credential_pair_id.in_(cc_pair_ids)\n ).delete(synchronize_session=False)\n\n session.query(ConnectorCredentialPair).filter(\n ConnectorCredentialPair.id.in_(cc_pair_ids)\n ).delete(synchronize_session=False)\n\n credential_ids = [cc_pair.credential_id for cc_pair in cc_pairs_to_delete]\n if credential_ids:\n session.query(Credential).filter(Credential.id.in_(credential_ids)).delete(\n synchronize_session=False\n )\n\n session.query(Connector).filter(Connector.id.in_(connector_ids)).delete(\n synchronize_session=False\n )\n\n session.commit()\n\n\ndef downgrade() -> None:\n # No-op downgrade as we cannot restore deleted data\n pass\n" }, { "path": "backend/alembic/versions/94dc3d0236f8_make_document_set_description_optional.py", "content": "\"\"\"make document set description optional\n\nRevision ID: 94dc3d0236f8\nRevises: bf7a81109301\nCreate Date: 2024-12-11 11:26:10.616722\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"94dc3d0236f8\"\ndown_revision = \"bf7a81109301\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Make document_set.description column nullable\n op.alter_column(\n \"document_set\", \"description\", existing_type=sa.String(), nullable=True\n )\n\n\ndef downgrade() -> None:\n # Revert document_set.description column to non-nullable\n op.alter_column(\n \"document_set\", \"description\", existing_type=sa.String(), nullable=False\n )\n" }, { "path": "backend/alembic/versions/96a5702df6aa_mcp_tool_enabled.py", "content": "\"\"\"mcp_tool_enabled\n\nRevision ID: 96a5702df6aa\nRevises: 40926a4dab77\nCreate Date: 2025-10-09 12:10:21.733097\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"96a5702df6aa\"\ndown_revision = \"40926a4dab77\"\nbranch_labels = None\ndepends_on = None\n\n\nDELETE_DISABLED_TOOLS_SQL = \"DELETE FROM tool WHERE enabled = false\"\n\n\ndef upgrade() -> None:\n op.add_column(\n \"tool\",\n sa.Column(\n \"enabled\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.true(),\n ),\n )\n op.create_index(\n \"ix_tool_mcp_server_enabled\",\n \"tool\",\n [\"mcp_server_id\", \"enabled\"],\n )\n # Remove the server default so application controls defaulting\n op.alter_column(\"tool\", \"enabled\", server_default=None)\n\n\ndef downgrade() -> None:\n op.execute(DELETE_DISABLED_TOOLS_SQL)\n op.drop_index(\"ix_tool_mcp_server_enabled\", table_name=\"tool\")\n op.drop_column(\"tool\", \"enabled\")\n" }, { "path": "backend/alembic/versions/977e834c1427_seed_default_groups.py", "content": "\"\"\"seed_default_groups\n\nRevision ID: 977e834c1427\nRevises: 8188861f4e92\nCreate Date: 2026-03-25 14:59:41.313091\n\n\"\"\"\n\nfrom typing import Any\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\n\n\n# revision identifiers, used by Alembic.\nrevision = \"977e834c1427\"\ndown_revision = \"8188861f4e92\"\nbranch_labels = None\ndepends_on = None\n\n# (group_name, permission_value)\nDEFAULT_GROUPS = [\n (\"Admin\", \"admin\"),\n (\"Basic\", \"basic\"),\n]\n\nCUSTOM_SUFFIX = \"(Custom)\"\n\nMAX_RENAME_ATTEMPTS = 100\n\n# Reflect table structures for use in DML\nuser_group_table = sa.table(\n \"user_group\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"name\", sa.String),\n sa.column(\"is_up_to_date\", sa.Boolean),\n sa.column(\"is_up_for_deletion\", sa.Boolean),\n sa.column(\"is_default\", sa.Boolean),\n)\n\npermission_grant_table = sa.table(\n \"permission_grant\",\n sa.column(\"group_id\", sa.Integer),\n sa.column(\"permission\", sa.String),\n sa.column(\"grant_source\", sa.String),\n)\n\nuser__user_group_table = sa.table(\n \"user__user_group\",\n sa.column(\"user_group_id\", sa.Integer),\n sa.column(\"user_id\", sa.Uuid),\n)\n\n\ndef _find_available_name(conn: sa.engine.Connection, base: str) -> str:\n \"\"\"Return a name like 'Admin (Custom)' or 'Admin (Custom 2)' that is not taken.\"\"\"\n candidate = f\"{base} {CUSTOM_SUFFIX}\"\n attempt = 1\n while attempt <= MAX_RENAME_ATTEMPTS:\n exists: Any = conn.execute(\n sa.select(sa.literal(1))\n .select_from(user_group_table)\n .where(user_group_table.c.name == candidate)\n .limit(1)\n ).fetchone()\n if exists is None:\n return candidate\n attempt += 1\n candidate = f\"{base} (Custom {attempt})\"\n raise RuntimeError(\n f\"Could not find an available name for group '{base}' \"\n f\"after {MAX_RENAME_ATTEMPTS} attempts\"\n )\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n for group_name, permission_value in DEFAULT_GROUPS:\n # Step 1: Rename ALL existing groups that clash with the canonical name.\n conflicting = conn.execute(\n sa.select(user_group_table.c.id, user_group_table.c.name).where(\n user_group_table.c.name == group_name\n )\n ).fetchall()\n\n for row_id, row_name in conflicting:\n new_name = _find_available_name(conn, row_name)\n op.execute(\n sa.update(user_group_table)\n .where(user_group_table.c.id == row_id)\n .values(name=new_name, is_up_to_date=False)\n )\n\n # Step 2: Create a fresh default group.\n result = conn.execute(\n user_group_table.insert()\n .values(\n name=group_name,\n is_up_to_date=True,\n is_up_for_deletion=False,\n is_default=True,\n )\n .returning(user_group_table.c.id)\n ).fetchone()\n assert result is not None\n group_id = result[0]\n\n # Step 3: Upsert permission grant.\n op.execute(\n pg_insert(permission_grant_table)\n .values(\n group_id=group_id,\n permission=permission_value,\n grant_source=\"SYSTEM\",\n )\n .on_conflict_do_nothing(index_elements=[\"group_id\", \"permission\"])\n )\n\n\ndef downgrade() -> None:\n # Remove the default groups created by this migration.\n # First remove user-group memberships that reference default groups\n # to avoid FK violations, then delete the groups themselves.\n default_group_ids = sa.select(user_group_table.c.id).where(\n user_group_table.c.is_default == True # noqa: E712\n )\n conn = op.get_bind()\n conn.execute(\n sa.delete(user__user_group_table).where(\n user__user_group_table.c.user_group_id.in_(default_group_ids)\n )\n )\n conn.execute(\n sa.delete(user_group_table).where(\n user_group_table.c.is_default == True # noqa: E712\n )\n )\n" }, { "path": "backend/alembic/versions/97dbb53fa8c8_add_syncrecord.py", "content": "\"\"\"Add SyncRecord\n\nRevision ID: 97dbb53fa8c8\nRevises: 369644546676\nCreate Date: 2025-01-11 19:39:50.426302\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"97dbb53fa8c8\"\ndown_revision = \"be2ab2aa50ee\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"sync_record\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"entity_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"sync_type\",\n sa.Enum(\n \"DOCUMENT_SET\",\n \"USER_GROUP\",\n \"CONNECTOR_DELETION\",\n name=\"synctype\",\n native_enum=False,\n length=40,\n ),\n nullable=False,\n ),\n sa.Column(\n \"sync_status\",\n sa.Enum(\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n \"CANCELED\",\n name=\"syncstatus\",\n native_enum=False,\n length=40,\n ),\n nullable=False,\n ),\n sa.Column(\"num_docs_synced\", sa.Integer(), nullable=False),\n sa.Column(\"sync_start_time\", sa.DateTime(timezone=True), nullable=False),\n sa.Column(\"sync_end_time\", sa.DateTime(timezone=True), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Add index for fetch_latest_sync_record query\n op.create_index(\n \"ix_sync_record_entity_id_sync_type_sync_start_time\",\n \"sync_record\",\n [\"entity_id\", \"sync_type\", \"sync_start_time\"],\n )\n\n # Add index for cleanup_sync_records query\n op.create_index(\n \"ix_sync_record_entity_id_sync_type_sync_status\",\n \"sync_record\",\n [\"entity_id\", \"sync_type\", \"sync_status\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_sync_record_entity_id_sync_type_sync_status\")\n op.drop_index(\"ix_sync_record_entity_id_sync_type_sync_start_time\")\n op.drop_table(\"sync_record\")\n" }, { "path": "backend/alembic/versions/98a5008d8711_agent_tracking.py", "content": "\"\"\"agent_tracking\n\nRevision ID: 98a5008d8711\nRevises: 2f80c6a2550f\nCreate Date: 2025-01-29 17:00:00.000001\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.dialects.postgresql import UUID\n\n# revision identifiers, used by Alembic.\nrevision = \"98a5008d8711\"\ndown_revision = \"2f80c6a2550f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"agent__search_metrics\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", postgresql.UUID(as_uuid=True), nullable=True),\n sa.Column(\"persona_id\", sa.Integer(), nullable=True),\n sa.Column(\"agent_type\", sa.String(), nullable=False),\n sa.Column(\"start_time\", sa.DateTime(timezone=True), nullable=False),\n sa.Column(\"base_duration_s\", sa.Float(), nullable=False),\n sa.Column(\"full_duration_s\", sa.Float(), nullable=False),\n sa.Column(\"base_metrics\", postgresql.JSONB(), nullable=True),\n sa.Column(\"refined_metrics\", postgresql.JSONB(), nullable=True),\n sa.Column(\"all_metrics\", postgresql.JSONB(), nullable=True),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Create sub_question table\n op.create_table(\n \"agent__sub_question\",\n sa.Column(\"id\", sa.Integer, primary_key=True),\n sa.Column(\"primary_question_id\", sa.Integer, sa.ForeignKey(\"chat_message.id\")),\n sa.Column(\n \"chat_session_id\", UUID(as_uuid=True), sa.ForeignKey(\"chat_session.id\")\n ),\n sa.Column(\"sub_question\", sa.Text),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.func.now()\n ),\n sa.Column(\"sub_answer\", sa.Text),\n sa.Column(\"sub_question_doc_results\", postgresql.JSONB(), nullable=True),\n sa.Column(\"level\", sa.Integer(), nullable=False),\n sa.Column(\"level_question_num\", sa.Integer(), nullable=False),\n )\n\n # Create sub_query table\n op.create_table(\n \"agent__sub_query\",\n sa.Column(\"id\", sa.Integer, primary_key=True),\n sa.Column(\n \"parent_question_id\", sa.Integer, sa.ForeignKey(\"agent__sub_question.id\")\n ),\n sa.Column(\n \"chat_session_id\", UUID(as_uuid=True), sa.ForeignKey(\"chat_session.id\")\n ),\n sa.Column(\"sub_query\", sa.Text),\n sa.Column(\n \"time_created\", sa.DateTime(timezone=True), server_default=sa.func.now()\n ),\n )\n\n # Create sub_query__search_doc association table\n op.create_table(\n \"agent__sub_query__search_doc\",\n sa.Column(\n \"sub_query_id\",\n sa.Integer,\n sa.ForeignKey(\"agent__sub_query.id\"),\n primary_key=True,\n ),\n sa.Column(\n \"search_doc_id\",\n sa.Integer,\n sa.ForeignKey(\"search_doc.id\"),\n primary_key=True,\n ),\n )\n\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"refined_answer_improvement\",\n sa.Boolean(),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"refined_answer_improvement\")\n op.drop_table(\"agent__sub_query__search_doc\")\n op.drop_table(\"agent__sub_query\")\n op.drop_table(\"agent__sub_question\")\n op.drop_table(\"agent__search_metrics\")\n" }, { "path": "backend/alembic/versions/9a0296d7421e_add_is_auto_mode_to_llm_provider.py", "content": "\"\"\"add_is_auto_mode_to_llm_provider\n\nRevision ID: 9a0296d7421e\nRevises: 7206234e012a\nCreate Date: 2025-12-17 18:14:29.620981\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"9a0296d7421e\"\ndown_revision = \"7206234e012a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"llm_provider\",\n sa.Column(\n \"is_auto_mode\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"llm_provider\", \"is_auto_mode\")\n" }, { "path": "backend/alembic/versions/9aadf32dfeb4_add_user_files.py", "content": "\"\"\"add user files\n\nRevision ID: 9aadf32dfeb4\nRevises: 3781a5eb12cb\nCreate Date: 2025-01-26 16:08:21.551022\n\n\"\"\"\n\nimport sqlalchemy as sa\nimport datetime\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"9aadf32dfeb4\"\ndown_revision = \"3781a5eb12cb\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create user_folder table without parent_id\n op.create_table(\n \"user_folder\",\n sa.Column(\"id\", sa.Integer(), primary_key=True, autoincrement=True),\n sa.Column(\"user_id\", sa.UUID(), sa.ForeignKey(\"user.id\"), nullable=True),\n sa.Column(\"name\", sa.String(length=255), nullable=True),\n sa.Column(\"description\", sa.String(length=255), nullable=True),\n sa.Column(\"display_priority\", sa.Integer(), nullable=True, default=0),\n sa.Column(\n \"created_at\", sa.DateTime(timezone=True), server_default=sa.func.now()\n ),\n )\n\n # Create user_file table with folder_id instead of parent_folder_id\n op.create_table(\n \"user_file\",\n sa.Column(\"id\", sa.Integer(), primary_key=True, autoincrement=True),\n sa.Column(\"user_id\", sa.UUID(), sa.ForeignKey(\"user.id\"), nullable=True),\n sa.Column(\n \"folder_id\",\n sa.Integer(),\n sa.ForeignKey(\"user_folder.id\"),\n nullable=True,\n ),\n sa.Column(\"link_url\", sa.String(), nullable=True),\n sa.Column(\"token_count\", sa.Integer(), nullable=True),\n sa.Column(\"file_type\", sa.String(), nullable=True),\n sa.Column(\"file_id\", sa.String(length=255), nullable=False),\n sa.Column(\"document_id\", sa.String(length=255), nullable=False),\n sa.Column(\"name\", sa.String(length=255), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(),\n default=datetime.datetime.utcnow,\n ),\n sa.Column(\n \"cc_pair_id\",\n sa.Integer(),\n sa.ForeignKey(\"connector_credential_pair.id\"),\n nullable=True,\n unique=True,\n ),\n )\n\n # Create persona__user_file table\n op.create_table(\n \"persona__user_file\",\n sa.Column(\n \"persona_id\", sa.Integer(), sa.ForeignKey(\"persona.id\"), primary_key=True\n ),\n sa.Column(\n \"user_file_id\",\n sa.Integer(),\n sa.ForeignKey(\"user_file.id\"),\n primary_key=True,\n ),\n )\n\n # Create persona__user_folder table\n op.create_table(\n \"persona__user_folder\",\n sa.Column(\n \"persona_id\", sa.Integer(), sa.ForeignKey(\"persona.id\"), primary_key=True\n ),\n sa.Column(\n \"user_folder_id\",\n sa.Integer(),\n sa.ForeignKey(\"user_folder.id\"),\n primary_key=True,\n ),\n )\n\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"is_user_file\", sa.Boolean(), nullable=True, default=False),\n )\n\n # Update existing records to have is_user_file=False instead of NULL\n op.execute(\n \"UPDATE connector_credential_pair SET is_user_file = FALSE WHERE is_user_file IS NULL\"\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"is_user_file\")\n # Drop the persona__user_folder table\n op.drop_table(\"persona__user_folder\")\n # Drop the persona__user_file table\n op.drop_table(\"persona__user_file\")\n # Drop the user_file table\n op.drop_table(\"user_file\")\n # Drop the user_folder table\n op.drop_table(\"user_folder\")\n" }, { "path": "backend/alembic/versions/9b66d3156fc6_user_file_schema_additions.py", "content": "\"\"\"Migration 1: User file schema additions\n\nRevision ID: 9b66d3156fc6\nRevises: b4ef3ae0bf6e\nCreate Date: 2025-09-22 09:42:06.086732\n\nThis migration adds new columns and tables without modifying existing data.\nIt is safe to run and can be easily rolled back.\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql as psql\nimport logging\n\nlogger = logging.getLogger(\"alembic.runtime.migration\")\n# revision identifiers, used by Alembic.\nrevision = \"9b66d3156fc6\"\ndown_revision = \"b4ef3ae0bf6e\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Add new columns and tables without modifying existing data.\"\"\"\n\n # Enable pgcrypto for UUID generation\n op.execute(\"CREATE EXTENSION IF NOT EXISTS pgcrypto\")\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n # === USER_FILE: Add new columns ===\n logger.info(\"Adding new columns to user_file table...\")\n\n user_file_columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n\n # Check if ID is already UUID (in case of re-run after partial migration)\n id_is_uuid = any(\n col[\"name\"] == \"id\" and \"uuid\" in str(col[\"type\"]).lower()\n for col in inspector.get_columns(\"user_file\")\n )\n\n # Add transitional UUID column only if ID is not already UUID\n if \"new_id\" not in user_file_columns and not id_is_uuid:\n op.add_column(\n \"user_file\",\n sa.Column(\n \"new_id\",\n psql.UUID(as_uuid=True),\n nullable=True,\n server_default=sa.text(\"gen_random_uuid()\"),\n ),\n )\n op.create_unique_constraint(\"uq_user_file_new_id\", \"user_file\", [\"new_id\"])\n logger.info(\"Added new_id column to user_file\")\n\n # Add status column\n if \"status\" not in user_file_columns:\n op.add_column(\n \"user_file\",\n sa.Column(\n \"status\",\n sa.Enum(\n \"PROCESSING\",\n \"COMPLETED\",\n \"FAILED\",\n \"CANCELED\",\n name=\"userfilestatus\",\n native_enum=False,\n ),\n nullable=False,\n server_default=\"PROCESSING\",\n ),\n )\n logger.info(\"Added status column to user_file\")\n\n # Add other tracking columns\n if \"chunk_count\" not in user_file_columns:\n op.add_column(\n \"user_file\", sa.Column(\"chunk_count\", sa.Integer(), nullable=True)\n )\n logger.info(\"Added chunk_count column to user_file\")\n\n if \"last_accessed_at\" not in user_file_columns:\n op.add_column(\n \"user_file\",\n sa.Column(\"last_accessed_at\", sa.DateTime(timezone=True), nullable=True),\n )\n logger.info(\"Added last_accessed_at column to user_file\")\n\n if \"needs_project_sync\" not in user_file_columns:\n op.add_column(\n \"user_file\",\n sa.Column(\n \"needs_project_sync\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.text(\"false\"),\n ),\n )\n logger.info(\"Added needs_project_sync column to user_file\")\n\n if \"last_project_sync_at\" not in user_file_columns:\n op.add_column(\n \"user_file\",\n sa.Column(\n \"last_project_sync_at\", sa.DateTime(timezone=True), nullable=True\n ),\n )\n logger.info(\"Added last_project_sync_at column to user_file\")\n\n if \"document_id_migrated\" not in user_file_columns:\n op.add_column(\n \"user_file\",\n sa.Column(\n \"document_id_migrated\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.text(\"true\"),\n ),\n )\n logger.info(\"Added document_id_migrated column to user_file\")\n\n # === USER_FOLDER -> USER_PROJECT rename ===\n table_names = set(inspector.get_table_names())\n\n if \"user_folder\" in table_names:\n logger.info(\"Updating user_folder table...\")\n # Make description nullable first\n op.alter_column(\"user_folder\", \"description\", nullable=True)\n\n # Rename table if user_project doesn't exist\n if \"user_project\" not in table_names:\n op.execute(\"ALTER TABLE user_folder RENAME TO user_project\")\n logger.info(\"Renamed user_folder to user_project\")\n elif \"user_project\" in table_names:\n # If already renamed, ensure column nullability\n project_cols = [col[\"name\"] for col in inspector.get_columns(\"user_project\")]\n if \"description\" in project_cols:\n op.alter_column(\"user_project\", \"description\", nullable=True)\n\n # Add instructions column to user_project\n inspector = sa.inspect(bind) # Refresh after rename\n if \"user_project\" in inspector.get_table_names():\n project_columns = [col[\"name\"] for col in inspector.get_columns(\"user_project\")]\n if \"instructions\" not in project_columns:\n op.add_column(\n \"user_project\",\n sa.Column(\"instructions\", sa.String(), nullable=True),\n )\n logger.info(\"Added instructions column to user_project\")\n\n # === CHAT_SESSION: Add project_id ===\n chat_session_columns = [\n col[\"name\"] for col in inspector.get_columns(\"chat_session\")\n ]\n if \"project_id\" not in chat_session_columns:\n op.add_column(\n \"chat_session\",\n sa.Column(\"project_id\", sa.Integer(), nullable=True),\n )\n logger.info(\"Added project_id column to chat_session\")\n\n # === PERSONA__USER_FILE: Add UUID column ===\n persona_user_file_columns = [\n col[\"name\"] for col in inspector.get_columns(\"persona__user_file\")\n ]\n if \"user_file_id_uuid\" not in persona_user_file_columns:\n op.add_column(\n \"persona__user_file\",\n sa.Column(\"user_file_id_uuid\", psql.UUID(as_uuid=True), nullable=True),\n )\n logger.info(\"Added user_file_id_uuid column to persona__user_file\")\n\n # === PROJECT__USER_FILE: Create new table ===\n if \"project__user_file\" not in inspector.get_table_names():\n op.create_table(\n \"project__user_file\",\n sa.Column(\"project_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_file_id\", psql.UUID(as_uuid=True), nullable=False),\n sa.PrimaryKeyConstraint(\"project_id\", \"user_file_id\"),\n )\n logger.info(\"Created project__user_file table\")\n\n # Only create the index if it doesn't exist\n existing_indexes = [\n ix[\"name\"] for ix in inspector.get_indexes(\"project__user_file\")\n ]\n if \"idx_project__user_file_user_file_id\" not in existing_indexes:\n op.create_index(\n \"idx_project__user_file_user_file_id\",\n \"project__user_file\",\n [\"user_file_id\"],\n )\n logger.info(\n \"Created index idx_project__user_file_user_file_id on project__user_file\"\n )\n\n logger.info(\"Migration 1 (schema additions) completed successfully\")\n\n\ndef downgrade() -> None:\n \"\"\"Remove added columns and tables.\"\"\"\n\n bind = op.get_bind()\n inspector = sa.inspect(bind)\n\n logger.info(\"Starting downgrade of schema additions...\")\n\n # Drop project__user_file table\n if \"project__user_file\" in inspector.get_table_names():\n # op.drop_index(\"idx_project__user_file_user_file_id\", \"project__user_file\")\n op.drop_table(\"project__user_file\")\n logger.info(\"Dropped project__user_file table\")\n\n # Remove columns from persona__user_file\n if \"persona__user_file\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"persona__user_file\")]\n if \"user_file_id_uuid\" in columns:\n op.drop_column(\"persona__user_file\", \"user_file_id_uuid\")\n logger.info(\"Dropped user_file_id_uuid from persona__user_file\")\n\n # Remove columns from chat_session\n if \"chat_session\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"chat_session\")]\n if \"project_id\" in columns:\n op.drop_column(\"chat_session\", \"project_id\")\n logger.info(\"Dropped project_id from chat_session\")\n\n # Rename user_project back to user_folder and remove instructions\n if \"user_project\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"user_project\")]\n if \"instructions\" in columns:\n op.drop_column(\"user_project\", \"instructions\")\n op.execute(\"ALTER TABLE user_project RENAME TO user_folder\")\n # Update NULL descriptions to empty string before setting NOT NULL constraint\n op.execute(\"UPDATE user_folder SET description = '' WHERE description IS NULL\")\n op.alter_column(\"user_folder\", \"description\", nullable=False)\n logger.info(\"Renamed user_project back to user_folder\")\n\n # Remove columns from user_file\n if \"user_file\" in inspector.get_table_names():\n columns = [col[\"name\"] for col in inspector.get_columns(\"user_file\")]\n\n columns_to_drop = [\n \"document_id_migrated\",\n \"last_project_sync_at\",\n \"needs_project_sync\",\n \"last_accessed_at\",\n \"chunk_count\",\n \"status\",\n ]\n\n for col in columns_to_drop:\n if col in columns:\n op.drop_column(\"user_file\", col)\n logger.info(f\"Dropped {col} from user_file\")\n\n if \"new_id\" in columns:\n op.drop_constraint(\"uq_user_file_new_id\", \"user_file\", type_=\"unique\")\n op.drop_column(\"user_file\", \"new_id\")\n logger.info(\"Dropped new_id from user_file\")\n\n # Drop enum type if no columns use it\n bind.execute(sa.text(\"DROP TYPE IF EXISTS userfilestatus\"))\n\n logger.info(\"Downgrade completed successfully\")\n" }, { "path": "backend/alembic/versions/9c00a2bccb83_chat_message_agentic.py", "content": "\"\"\"chat_message_agentic\n\nRevision ID: 9c00a2bccb83\nRevises: b7a7eee5aa15\nCreate Date: 2025-02-17 11:15:43.081150\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"9c00a2bccb83\"\ndown_revision = \"b7a7eee5aa15\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # First add the column as nullable\n op.add_column(\"chat_message\", sa.Column(\"is_agentic\", sa.Boolean(), nullable=True))\n\n # Update existing rows based on presence of SubQuestions\n op.execute(\n \"\"\"\n UPDATE chat_message\n SET is_agentic = EXISTS (\n SELECT 1\n FROM agent__sub_question\n WHERE agent__sub_question.primary_question_id = chat_message.id\n )\n WHERE is_agentic IS NULL\n \"\"\"\n )\n\n # Make the column non-nullable with a default value of False\n op.alter_column(\n \"chat_message\", \"is_agentic\", nullable=False, server_default=sa.text(\"false\")\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"is_agentic\")\n" }, { "path": "backend/alembic/versions/9c54986124c6_add_scim_tables.py", "content": "\"\"\"add_scim_tables\n\nRevision ID: 9c54986124c6\nRevises: b51c6844d1df\nCreate Date: 2026-02-12 20:29:47.448614\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"9c54986124c6\"\ndown_revision = \"b51c6844d1df\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"scim_token\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"hashed_token\", sa.String(length=64), nullable=False),\n sa.Column(\"token_display\", sa.String(), nullable=False),\n sa.Column(\n \"created_by_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\n \"is_active\",\n sa.Boolean(),\n server_default=sa.text(\"true\"),\n nullable=False,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"last_used_at\", sa.DateTime(timezone=True), nullable=True),\n sa.ForeignKeyConstraint([\"created_by_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"hashed_token\"),\n )\n op.create_table(\n \"scim_group_mapping\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"external_id\", sa.String(), nullable=False),\n sa.Column(\"user_group_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"], [\"user_group.id\"], ondelete=\"CASCADE\"\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"user_group_id\"),\n )\n op.create_index(\n op.f(\"ix_scim_group_mapping_external_id\"),\n \"scim_group_mapping\",\n [\"external_id\"],\n unique=True,\n )\n op.create_table(\n \"scim_user_mapping\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"external_id\", sa.String(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n onupdate=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"user_id\"),\n )\n op.create_index(\n op.f(\"ix_scim_user_mapping_external_id\"),\n \"scim_user_mapping\",\n [\"external_id\"],\n unique=True,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n op.f(\"ix_scim_user_mapping_external_id\"),\n table_name=\"scim_user_mapping\",\n )\n op.drop_table(\"scim_user_mapping\")\n op.drop_index(\n op.f(\"ix_scim_group_mapping_external_id\"),\n table_name=\"scim_group_mapping\",\n )\n op.drop_table(\"scim_group_mapping\")\n op.drop_table(\"scim_token\")\n" }, { "path": "backend/alembic/versions/9cf5c00f72fe_add_creator_to_cc_pair.py", "content": "\"\"\"add creator to cc pair\n\nRevision ID: 9cf5c00f72fe\nRevises: 26b931506ecb\nCreate Date: 2024-11-12 15:16:42.682902\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"9cf5c00f72fe\"\ndown_revision = \"26b931506ecb\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"creator_id\",\n sa.UUID(as_uuid=True),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"creator_id\")\n" }, { "path": "backend/alembic/versions/9d1543a37106_add_processing_duration_seconds_to_chat_.py", "content": "\"\"\"add processing_duration_seconds to chat_message\n\nRevision ID: 9d1543a37106\nRevises: cbc03e08d0f3\nCreate Date: 2026-01-21 11:42:18.546188\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"9d1543a37106\"\ndown_revision = \"cbc03e08d0f3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\",\n sa.Column(\"processing_duration_seconds\", sa.Float(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"processing_duration_seconds\")\n" }, { "path": "backend/alembic/versions/9d97fecfab7f_added_retrieved_docs_to_query_event.py", "content": "\"\"\"Added retrieved docs to query event\n\nRevision ID: 9d97fecfab7f\nRevises: ffc707a226b4\nCreate Date: 2023-10-20 12:22:31.930449\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"9d97fecfab7f\"\ndown_revision = \"ffc707a226b4\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"query_event\",\n sa.Column(\n \"retrieved_document_ids\",\n postgresql.ARRAY(sa.String()),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"query_event\", \"retrieved_document_ids\")\n" }, { "path": "backend/alembic/versions/9drpiiw74ljy_add_config_to_federated_connector.py", "content": "\"\"\"add config to federated_connector\n\nRevision ID: 9drpiiw74ljy\nRevises: 2acdef638fc2\nCreate Date: 2025-11-03 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"9drpiiw74ljy\"\ndown_revision = \"2acdef638fc2\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n connection = op.get_bind()\n\n # Check if column already exists in current schema\n result = connection.execute(\n sa.text(\n \"\"\"\n SELECT column_name\n FROM information_schema.columns\n WHERE table_schema = current_schema()\n AND table_name = 'federated_connector'\n AND column_name = 'config'\n \"\"\"\n )\n )\n column_exists = result.fetchone() is not None\n\n # Add config column with default empty object (only if it doesn't exist)\n if not column_exists:\n op.add_column(\n \"federated_connector\",\n sa.Column(\n \"config\", postgresql.JSONB(), nullable=False, server_default=\"{}\"\n ),\n )\n\n # Data migration: Single bulk update for all Slack connectors\n connection.execute(\n sa.text(\n \"\"\"\n WITH connector_configs AS (\n SELECT\n fc.id as connector_id,\n CASE\n WHEN fcds.entities->'channels' IS NOT NULL\n AND jsonb_typeof(fcds.entities->'channels') = 'array'\n AND jsonb_array_length(fcds.entities->'channels') > 0\n THEN\n jsonb_build_object(\n 'channels', fcds.entities->'channels',\n 'search_all_channels', false\n ) ||\n CASE\n WHEN fcds.entities->'include_dm' IS NOT NULL\n THEN jsonb_build_object('include_dm', fcds.entities->'include_dm')\n ELSE '{}'::jsonb\n END\n ELSE\n jsonb_build_object('search_all_channels', true) ||\n CASE\n WHEN fcds.entities->'include_dm' IS NOT NULL\n THEN jsonb_build_object('include_dm', fcds.entities->'include_dm')\n ELSE '{}'::jsonb\n END\n END as config\n FROM federated_connector fc\n LEFT JOIN LATERAL (\n SELECT entities\n FROM federated_connector__document_set\n WHERE federated_connector_id = fc.id\n AND entities IS NOT NULL\n ORDER BY id\n LIMIT 1\n ) fcds ON true\n WHERE fc.source = 'FEDERATED_SLACK'\n AND fcds.entities IS NOT NULL\n )\n UPDATE federated_connector fc\n SET config = cc.config\n FROM connector_configs cc\n WHERE fc.id = cc.connector_id\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"federated_connector\", \"config\")\n" }, { "path": "backend/alembic/versions/9f696734098f_combine_search_and_chat.py", "content": "\"\"\"Combine Search and Chat\n\nRevision ID: 9f696734098f\nRevises: a8c2065484e6\nCreate Date: 2024-11-27 15:32:19.694972\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"9f696734098f\"\ndown_revision = \"a8c2065484e6\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"chat_session\", \"description\", nullable=True)\n op.drop_column(\"chat_session\", \"one_shot\")\n op.drop_column(\"slack_channel_config\", \"response_type\")\n\n\ndef downgrade() -> None:\n op.execute(\"UPDATE chat_session SET description = '' WHERE description IS NULL\")\n op.alter_column(\"chat_session\", \"description\", nullable=False)\n op.add_column(\n \"chat_session\",\n sa.Column(\"one_shot\", sa.Boolean(), nullable=False, server_default=sa.false()),\n )\n op.add_column(\n \"slack_channel_config\",\n sa.Column(\n \"response_type\", sa.String(), nullable=False, server_default=\"citations\"\n ),\n )\n" }, { "path": "backend/alembic/versions/a01bf2971c5d_update_default_tool_descriptions.py", "content": "\"\"\"update_default_tool_descriptions\n\nRevision ID: a01bf2971c5d\nRevises: 87c52ec39f84\nCreate Date: 2025-12-16 15:21:25.656375\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"a01bf2971c5d\"\ndown_revision = \"18b5b2524446\"\nbranch_labels = None\ndepends_on = None\n\n# new tool descriptions (12/2025)\nTOOL_DESCRIPTIONS = {\n \"SearchTool\": \"The Search Action allows the agent to search through connected knowledge to help build an answer.\",\n \"ImageGenerationTool\": (\n \"The Image Generation Action allows the agent to use DALL-E 3 or GPT-IMAGE-1 to generate images. \"\n \"The action will be used when the user asks the agent to generate an image.\"\n ),\n \"WebSearchTool\": (\n \"The Web Search Action allows the agent to perform internet searches for up-to-date information.\"\n ),\n \"KnowledgeGraphTool\": (\n \"The Knowledge Graph Search Action allows the agent to search the \"\n \"Knowledge Graph for information. This tool can (for now) only be active in the KG Beta Agent, \"\n \"and it requires the Knowledge Graph to be enabled.\"\n ),\n \"OktaProfileTool\": (\n \"The Okta Profile Action allows the agent to fetch the current user's information from Okta. \"\n \"This may include the user's name, email, phone number, address, and other details such as their \"\n \"manager and direct reports.\"\n ),\n}\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n for tool_id, description in TOOL_DESCRIPTIONS.items():\n conn.execute(\n sa.text(\n \"UPDATE tool SET description = :description WHERE in_code_tool_id = :tool_id\"\n ),\n {\"description\": description, \"tool_id\": tool_id},\n )\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/a1b2c3d4e5f6_add_license_table.py", "content": "\"\"\"add license table\n\nRevision ID: a1b2c3d4e5f6\nRevises: a01bf2971c5d\nCreate Date: 2025-12-04 10:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"a1b2c3d4e5f6\"\ndown_revision = \"a01bf2971c5d\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"license\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"license_data\", sa.Text(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n )\n\n # Singleton pattern - only ever one row in this table\n op.create_index(\n \"idx_license_singleton\",\n \"license\",\n [sa.text(\"(true)\")],\n unique=True,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"idx_license_singleton\", table_name=\"license\")\n op.drop_table(\"license\")\n" }, { "path": "backend/alembic/versions/a1b2c3d4e5f7_drop_agent_search_metrics_table.py", "content": "\"\"\"drop agent_search_metrics table\n\nRevision ID: a1b2c3d4e5f7\nRevises: 73e9983e5091\nCreate Date: 2026-01-17\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"a1b2c3d4e5f7\"\ndown_revision = \"73e9983e5091\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_table(\"agent__search_metrics\")\n\n\ndef downgrade() -> None:\n op.create_table(\n \"agent__search_metrics\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", sa.UUID(), nullable=True),\n sa.Column(\"persona_id\", sa.Integer(), nullable=True),\n sa.Column(\"agent_type\", sa.String(), nullable=False),\n sa.Column(\"start_time\", sa.DateTime(timezone=True), nullable=False),\n sa.Column(\"base_duration_s\", sa.Float(), nullable=False),\n sa.Column(\"full_duration_s\", sa.Float(), nullable=False),\n sa.Column(\"base_metrics\", postgresql.JSONB(), nullable=True),\n sa.Column(\"refined_metrics\", postgresql.JSONB(), nullable=True),\n sa.Column(\"all_metrics\", postgresql.JSONB(), nullable=True),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n" }, { "path": "backend/alembic/versions/a2b3c4d5e6f7_remove_fast_default_model_name.py", "content": "\"\"\"Remove fast_default_model_name from llm_provider\n\nRevision ID: a2b3c4d5e6f7\nRevises: 2a391f840e85\nCreate Date: 2024-12-17\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"a2b3c4d5e6f7\"\ndown_revision = \"2a391f840e85\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"llm_provider\", \"fast_default_model_name\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"llm_provider\",\n sa.Column(\"fast_default_model_name\", sa.String(), nullable=True),\n )\n" }, { "path": "backend/alembic/versions/a3795dce87be_migration_confluence_to_be_explicit.py", "content": "\"\"\"migration confluence to be explicit\n\nRevision ID: a3795dce87be\nRevises: 1f60f60c3401\nCreate Date: 2024-09-01 13:52:12.006740\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.sql import table, column\n\nrevision = \"a3795dce87be\"\ndown_revision = \"1f60f60c3401\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef extract_confluence_keys_from_url(wiki_url: str) -> tuple[str, str, str, bool]:\n from urllib.parse import urlparse\n\n def _extract_confluence_keys_from_cloud_url(wiki_url: str) -> tuple[str, str, str]:\n parsed_url = urlparse(wiki_url)\n wiki_base = f\"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path.split('/spaces')[0]}\"\n path_parts = parsed_url.path.split(\"/\")\n space = path_parts[3]\n page_id = path_parts[5] if len(path_parts) > 5 else \"\"\n return wiki_base, space, page_id\n\n def _extract_confluence_keys_from_datacenter_url(\n wiki_url: str,\n ) -> tuple[str, str, str]:\n DISPLAY = \"/display/\"\n PAGE = \"/pages/\"\n parsed_url = urlparse(wiki_url)\n wiki_base = f\"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path.split(DISPLAY)[0]}\"\n space = DISPLAY.join(parsed_url.path.split(DISPLAY)[1:]).split(\"/\")[0]\n page_id = \"\"\n if (content := parsed_url.path.split(PAGE)) and len(content) > 1:\n page_id = content[1]\n return wiki_base, space, page_id\n\n is_confluence_cloud = (\n \".atlassian.net/wiki/spaces/\" in wiki_url\n or \".jira.com/wiki/spaces/\" in wiki_url\n )\n\n if is_confluence_cloud:\n wiki_base, space, page_id = _extract_confluence_keys_from_cloud_url(wiki_url)\n else:\n wiki_base, space, page_id = _extract_confluence_keys_from_datacenter_url(\n wiki_url\n )\n\n return wiki_base, space, page_id, is_confluence_cloud\n\n\ndef reconstruct_confluence_url(\n wiki_base: str, space: str, page_id: str, is_cloud: bool\n) -> str:\n if is_cloud:\n url = f\"{wiki_base}/spaces/{space}\"\n if page_id:\n url += f\"/pages/{page_id}\"\n else:\n url = f\"{wiki_base}/display/{space}\"\n if page_id:\n url += f\"/pages/{page_id}\"\n return url\n\n\ndef upgrade() -> None:\n connector = table(\n \"connector\",\n column(\"id\", sa.Integer),\n column(\"source\", sa.String()),\n column(\"input_type\", sa.String()),\n column(\"connector_specific_config\", postgresql.JSONB),\n )\n\n # Fetch all Confluence connectors\n connection = op.get_bind()\n confluence_connectors = connection.execute(\n sa.select(connector).where(\n sa.and_(\n connector.c.source == \"CONFLUENCE\", connector.c.input_type == \"POLL\"\n )\n )\n ).fetchall()\n\n for row in confluence_connectors:\n config = row.connector_specific_config\n wiki_page_url = config[\"wiki_page_url\"]\n wiki_base, space, page_id, is_cloud = extract_confluence_keys_from_url(\n wiki_page_url\n )\n\n new_config = {\n \"wiki_base\": wiki_base,\n \"space\": space,\n \"page_id\": page_id,\n \"is_cloud\": is_cloud,\n }\n\n for key, value in config.items():\n if key not in [\"wiki_page_url\"]:\n new_config[key] = value\n\n op.execute(\n connector.update()\n .where(connector.c.id == row.id)\n .values(connector_specific_config=new_config)\n )\n\n\ndef downgrade() -> None:\n connector = table(\n \"connector\",\n column(\"id\", sa.Integer),\n column(\"source\", sa.String()),\n column(\"input_type\", sa.String()),\n column(\"connector_specific_config\", postgresql.JSONB),\n )\n\n confluence_connectors = (\n op.get_bind()\n .execute(\n sa.select(connector).where(\n connector.c.source == \"CONFLUENCE\", connector.c.input_type == \"POLL\"\n )\n )\n .fetchall()\n )\n\n for row in confluence_connectors:\n config = row.connector_specific_config\n if all(key in config for key in [\"wiki_base\", \"space\", \"is_cloud\"]):\n wiki_page_url = reconstruct_confluence_url(\n config[\"wiki_base\"],\n config[\"space\"],\n config.get(\"page_id\", \"\"),\n config[\"is_cloud\"],\n )\n\n new_config = {\"wiki_page_url\": wiki_page_url}\n new_config.update(\n {\n k: v\n for k, v in config.items()\n if k not in [\"wiki_base\", \"space\", \"page_id\", \"is_cloud\"]\n }\n )\n\n op.execute(\n connector.update()\n .where(connector.c.id == row.id)\n .values(connector_specific_config=new_config)\n )\n" }, { "path": "backend/alembic/versions/a3b8d9e2f1c4_make_scim_external_id_nullable.py", "content": "\"\"\"make scim_user_mapping.external_id nullable\n\nRevision ID: a3b8d9e2f1c4\nRevises: 2664261bfaab\nCreate Date: 2026-03-02\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"a3b8d9e2f1c4\"\ndown_revision = \"2664261bfaab\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"scim_user_mapping\",\n \"external_id\",\n nullable=True,\n )\n\n\ndef downgrade() -> None:\n # Delete any rows where external_id is NULL before re-applying NOT NULL\n op.execute(\"DELETE FROM scim_user_mapping WHERE external_id IS NULL\")\n op.alter_column(\n \"scim_user_mapping\",\n \"external_id\",\n nullable=False,\n )\n" }, { "path": "backend/alembic/versions/a3bfd0d64902_add_chosen_assistants_to_user_table.py", "content": "\"\"\"Add chosen_assistants to User table\n\nRevision ID: a3bfd0d64902\nRevises: ec85f2b3c544\nCreate Date: 2024-05-26 17:22:24.834741\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"a3bfd0d64902\"\ndown_revision = \"ec85f2b3c544\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\"chosen_assistants\", postgresql.ARRAY(sa.Integer()), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"chosen_assistants\")\n" }, { "path": "backend/alembic/versions/a3c1a7904cd0_remove_userfile_related_deprecated_.py", "content": "\"\"\"remove userfile related deprecated fields\n\nRevision ID: a3c1a7904cd0\nRevises: 5c3dca366b35\nCreate Date: 2026-01-06 13:00:30.634396\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"a3c1a7904cd0\"\ndown_revision = \"5c3dca366b35\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"user_file\", \"document_id\")\n op.drop_column(\"user_file\", \"document_id_migrated\")\n op.drop_column(\"connector_credential_pair\", \"is_user_file\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"is_user_file\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n op.add_column(\n \"user_file\",\n sa.Column(\"document_id\", sa.String(), nullable=True),\n )\n op.add_column(\n \"user_file\",\n sa.Column(\n \"document_id_migrated\", sa.Boolean(), nullable=False, server_default=\"true\"\n ),\n )\n" }, { "path": "backend/alembic/versions/a3f8b2c1d4e5_add_preferred_response_id_to_chat_message.py", "content": "\"\"\"add preferred_response_id and model_display_name to chat_message\n\nRevision ID: a3f8b2c1d4e5\nCreate Date: 2026-03-22\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"a3f8b2c1d4e5\"\ndown_revision = \"25a5501dc766\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"preferred_response_id\",\n sa.Integer(),\n sa.ForeignKey(\"chat_message.id\", ondelete=\"SET NULL\"),\n nullable=True,\n ),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"model_display_name\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"model_display_name\")\n op.drop_column(\"chat_message\", \"preferred_response_id\")\n" }, { "path": "backend/alembic/versions/a4f23d6b71c8_add_llm_provider_persona_restrictions.py", "content": "\"\"\"add llm provider persona restrictions\n\nRevision ID: a4f23d6b71c8\nRevises: 5e1c073d48a3\nCreate Date: 2025-10-21 00:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"a4f23d6b71c8\"\ndown_revision = \"5e1c073d48a3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"llm_provider__persona\",\n sa.Column(\"llm_provider_id\", sa.Integer(), nullable=False),\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"llm_provider_id\"], [\"llm_provider.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"persona_id\"], [\"persona.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"llm_provider_id\", \"persona_id\"),\n )\n op.create_index(\n \"ix_llm_provider__persona_llm_provider_id\",\n \"llm_provider__persona\",\n [\"llm_provider_id\"],\n )\n op.create_index(\n \"ix_llm_provider__persona_persona_id\",\n \"llm_provider__persona\",\n [\"persona_id\"],\n )\n op.create_index(\n \"ix_llm_provider__persona_composite\",\n \"llm_provider__persona\",\n [\"persona_id\", \"llm_provider_id\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"ix_llm_provider__persona_composite\",\n table_name=\"llm_provider__persona\",\n )\n op.drop_index(\n \"ix_llm_provider__persona_persona_id\",\n table_name=\"llm_provider__persona\",\n )\n op.drop_index(\n \"ix_llm_provider__persona_llm_provider_id\",\n table_name=\"llm_provider__persona\",\n )\n op.drop_table(\"llm_provider__persona\")\n" }, { "path": "backend/alembic/versions/a570b80a5f20_usergroup_tables.py", "content": "\"\"\"UserGroup tables\n\nRevision ID: a570b80a5f20\nRevises: 904451035c9b\nCreate Date: 2023-10-02 12:27:10.265725\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"a570b80a5f20\"\ndown_revision = \"904451035c9b\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"user_group\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"is_up_to_date\", sa.Boolean(), nullable=False),\n sa.Column(\"is_up_for_deletion\", sa.Boolean(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n op.create_table(\n \"user__user_group\",\n sa.Column(\"user_group_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"user_group_id\", \"user_id\"),\n )\n op.create_table(\n \"user_group__connector_credential_pair\",\n sa.Column(\"user_group_id\", sa.Integer(), nullable=False),\n sa.Column(\"cc_pair_id\", sa.Integer(), nullable=False),\n sa.Column(\"is_current\", sa.Boolean(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"cc_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.PrimaryKeyConstraint(\"user_group_id\", \"cc_pair_id\", \"is_current\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"user_group__connector_credential_pair\")\n op.drop_table(\"user__user_group\")\n op.drop_table(\"user_group\")\n" }, { "path": "backend/alembic/versions/a6df6b88ef81_remove_recent_assistants.py", "content": "\"\"\"remove recent assistants\n\nRevision ID: a6df6b88ef81\nRevises: 4d58345da04a\nCreate Date: 2025-01-29 10:25:52.790407\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"a6df6b88ef81\"\ndown_revision = \"4d58345da04a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"user\", \"recent_assistants\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"recent_assistants\", postgresql.JSONB(), server_default=\"[]\", nullable=False\n ),\n )\n" }, { "path": "backend/alembic/versions/a7688ab35c45_add_public_external_user_group_table.py", "content": "\"\"\"Add public_external_user_group table\n\nRevision ID: a7688ab35c45\nRevises: 5c448911b12f\nCreate Date: 2025-05-06 20:55:12.747875\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"a7688ab35c45\"\ndown_revision = \"5c448911b12f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"public_external_user_group\",\n sa.Column(\"external_user_group_id\", sa.String(), nullable=False),\n sa.Column(\"cc_pair_id\", sa.Integer(), nullable=False),\n sa.PrimaryKeyConstraint(\"external_user_group_id\", \"cc_pair_id\"),\n sa.ForeignKeyConstraint(\n [\"cc_pair_id\"], [\"connector_credential_pair.id\"], ondelete=\"CASCADE\"\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"public_external_user_group\")\n" }, { "path": "backend/alembic/versions/a852cbe15577_new_chat_history.py", "content": "\"\"\"New Chat History\n\nRevision ID: a852cbe15577\nRevises: 6436661d5b65\nCreate Date: 2025-11-08 15:16:37.781308\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"a852cbe15577\"\ndown_revision = \"6436661d5b65\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # 1. Drop old research/agent tables (CASCADE handles dependencies)\n op.execute(\"DROP TABLE IF EXISTS research_agent_iteration_sub_step CASCADE\")\n op.execute(\"DROP TABLE IF EXISTS research_agent_iteration CASCADE\")\n op.execute(\"DROP TABLE IF EXISTS agent__sub_query__search_doc CASCADE\")\n op.execute(\"DROP TABLE IF EXISTS agent__sub_query CASCADE\")\n op.execute(\"DROP TABLE IF EXISTS agent__sub_question CASCADE\")\n\n # 2. ChatMessage table changes\n # Rename columns and add FKs\n op.alter_column(\n \"chat_message\", \"parent_message\", new_column_name=\"parent_message_id\"\n )\n op.create_foreign_key(\n \"fk_chat_message_parent_message_id\",\n \"chat_message\",\n \"chat_message\",\n [\"parent_message_id\"],\n [\"id\"],\n )\n op.alter_column(\n \"chat_message\",\n \"latest_child_message\",\n new_column_name=\"latest_child_message_id\",\n )\n op.create_foreign_key(\n \"fk_chat_message_latest_child_message_id\",\n \"chat_message\",\n \"chat_message\",\n [\"latest_child_message_id\"],\n [\"id\"],\n )\n\n # Add new column\n op.add_column(\n \"chat_message\", sa.Column(\"reasoning_tokens\", sa.Text(), nullable=True)\n )\n\n # Drop old columns\n op.drop_column(\"chat_message\", \"rephrased_query\")\n op.drop_column(\"chat_message\", \"alternate_assistant_id\")\n op.drop_column(\"chat_message\", \"overridden_model\")\n op.drop_column(\"chat_message\", \"is_agentic\")\n op.drop_column(\"chat_message\", \"refined_answer_improvement\")\n op.drop_column(\"chat_message\", \"research_type\")\n op.drop_column(\"chat_message\", \"research_plan\")\n op.drop_column(\"chat_message\", \"research_answer_purpose\")\n\n # 3. ToolCall table changes\n # Drop the unique constraint first\n op.drop_constraint(\"uq_tool_call_message_id\", \"tool_call\", type_=\"unique\")\n\n # Delete orphaned tool_call rows (those without valid chat_message)\n op.execute(\n \"DELETE FROM tool_call WHERE message_id NOT IN (SELECT id FROM chat_message)\"\n )\n\n # Add chat_session_id as nullable first, populate, then make NOT NULL\n op.add_column(\n \"tool_call\",\n sa.Column(\"chat_session_id\", postgresql.UUID(as_uuid=True), nullable=True),\n )\n\n # Populate chat_session_id from the related chat_message\n op.execute(\n \"\"\"\n UPDATE tool_call\n SET chat_session_id = chat_message.chat_session_id\n FROM chat_message\n WHERE tool_call.message_id = chat_message.id\n \"\"\"\n )\n\n # Now make it NOT NULL and add FK\n op.alter_column(\"tool_call\", \"chat_session_id\", nullable=False)\n op.create_foreign_key(\n \"fk_tool_call_chat_session_id\",\n \"tool_call\",\n \"chat_session\",\n [\"chat_session_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # Rename message_id and make nullable, recreate FK with CASCADE\n op.drop_constraint(\"tool_call_message_id_fkey\", \"tool_call\", type_=\"foreignkey\")\n op.alter_column(\n \"tool_call\",\n \"message_id\",\n new_column_name=\"parent_chat_message_id\",\n nullable=True,\n )\n op.create_foreign_key(\n \"fk_tool_call_parent_chat_message_id\",\n \"tool_call\",\n \"chat_message\",\n [\"parent_chat_message_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # Add parent_tool_call_id with FK\n op.add_column(\n \"tool_call\", sa.Column(\"parent_tool_call_id\", sa.Integer(), nullable=True)\n )\n op.create_foreign_key(\n \"fk_tool_call_parent_tool_call_id\",\n \"tool_call\",\n \"tool_call\",\n [\"parent_tool_call_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # Add other new columns\n op.add_column(\n \"tool_call\",\n sa.Column(\"turn_number\", sa.Integer(), nullable=False, server_default=\"0\"),\n )\n op.add_column(\n \"tool_call\",\n sa.Column(\"tool_call_id\", sa.String(), nullable=False, server_default=\"\"),\n )\n op.add_column(\"tool_call\", sa.Column(\"reasoning_tokens\", sa.Text(), nullable=True))\n op.add_column(\n \"tool_call\",\n sa.Column(\"tool_call_tokens\", sa.Integer(), nullable=False, server_default=\"0\"),\n )\n op.add_column(\n \"tool_call\",\n sa.Column(\"generated_images\", postgresql.JSONB(), nullable=True),\n )\n\n # Rename columns\n op.alter_column(\n \"tool_call\", \"tool_arguments\", new_column_name=\"tool_call_arguments\"\n )\n op.alter_column(\"tool_call\", \"tool_result\", new_column_name=\"tool_call_response\")\n\n # Change tool_call_response type from JSONB to Text\n op.execute(\n \"\"\"\n ALTER TABLE tool_call\n ALTER COLUMN tool_call_response TYPE TEXT\n USING tool_call_response::text\n \"\"\"\n )\n\n # Drop old columns\n op.drop_column(\"tool_call\", \"tool_name\")\n\n # 4. Create new association table\n op.create_table(\n \"tool_call__search_doc\",\n sa.Column(\"tool_call_id\", sa.Integer(), nullable=False),\n sa.Column(\"search_doc_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint([\"tool_call_id\"], [\"tool_call.id\"], ondelete=\"CASCADE\"),\n sa.ForeignKeyConstraint(\n [\"search_doc_id\"], [\"search_doc.id\"], ondelete=\"CASCADE\"\n ),\n sa.PrimaryKeyConstraint(\"tool_call_id\", \"search_doc_id\"),\n )\n\n # 5. Persona table change\n op.add_column(\n \"persona\",\n sa.Column(\n \"replace_base_system_prompt\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n\n\ndef downgrade() -> None:\n # Reverse persona changes\n op.drop_column(\"persona\", \"replace_base_system_prompt\")\n\n # Drop new association table\n op.drop_table(\"tool_call__search_doc\")\n\n # Reverse ToolCall changes\n op.add_column(\n \"tool_call\",\n sa.Column(\"tool_name\", sa.String(), nullable=False, server_default=\"\"),\n )\n\n # Change tool_call_response back to JSONB\n op.execute(\n \"\"\"\n ALTER TABLE tool_call\n ALTER COLUMN tool_call_response TYPE JSONB\n USING tool_call_response::jsonb\n \"\"\"\n )\n\n op.alter_column(\"tool_call\", \"tool_call_response\", new_column_name=\"tool_result\")\n op.alter_column(\n \"tool_call\", \"tool_call_arguments\", new_column_name=\"tool_arguments\"\n )\n\n op.drop_column(\"tool_call\", \"generated_images\")\n op.drop_column(\"tool_call\", \"tool_call_tokens\")\n op.drop_column(\"tool_call\", \"reasoning_tokens\")\n op.drop_column(\"tool_call\", \"tool_call_id\")\n op.drop_column(\"tool_call\", \"turn_number\")\n\n op.drop_constraint(\n \"fk_tool_call_parent_tool_call_id\", \"tool_call\", type_=\"foreignkey\"\n )\n op.drop_column(\"tool_call\", \"parent_tool_call_id\")\n\n op.drop_constraint(\n \"fk_tool_call_parent_chat_message_id\", \"tool_call\", type_=\"foreignkey\"\n )\n op.alter_column(\n \"tool_call\",\n \"parent_chat_message_id\",\n new_column_name=\"message_id\",\n nullable=False,\n )\n op.create_foreign_key(\n \"tool_call_message_id_fkey\",\n \"tool_call\",\n \"chat_message\",\n [\"message_id\"],\n [\"id\"],\n )\n\n op.drop_constraint(\"fk_tool_call_chat_session_id\", \"tool_call\", type_=\"foreignkey\")\n op.drop_column(\"tool_call\", \"chat_session_id\")\n\n op.create_unique_constraint(\"uq_tool_call_message_id\", \"tool_call\", [\"message_id\"])\n\n # Reverse ChatMessage changes\n # Note: research_answer_purpose and research_type were originally String columns,\n # not Enum types (see migrations 5ae8240accb3 and f8a9b2c3d4e5)\n op.add_column(\n \"chat_message\",\n sa.Column(\"research_answer_purpose\", sa.String(), nullable=True),\n )\n op.add_column(\n \"chat_message\", sa.Column(\"research_plan\", postgresql.JSONB(), nullable=True)\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"research_type\", sa.String(), nullable=True),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"refined_answer_improvement\", sa.Boolean(), nullable=True),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"is_agentic\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n op.add_column(\n \"chat_message\", sa.Column(\"overridden_model\", sa.String(), nullable=True)\n )\n op.add_column(\n \"chat_message\", sa.Column(\"alternate_assistant_id\", sa.Integer(), nullable=True)\n )\n # Recreate the FK constraint that was implicitly dropped when the column was dropped\n op.create_foreign_key(\n \"fk_chat_message_persona\",\n \"chat_message\",\n \"persona\",\n [\"alternate_assistant_id\"],\n [\"id\"],\n )\n op.add_column(\n \"chat_message\", sa.Column(\"rephrased_query\", sa.Text(), nullable=True)\n )\n\n op.drop_column(\"chat_message\", \"reasoning_tokens\")\n\n op.drop_constraint(\n \"fk_chat_message_latest_child_message_id\", \"chat_message\", type_=\"foreignkey\"\n )\n op.alter_column(\n \"chat_message\",\n \"latest_child_message_id\",\n new_column_name=\"latest_child_message\",\n )\n\n op.drop_constraint(\n \"fk_chat_message_parent_message_id\", \"chat_message\", type_=\"foreignkey\"\n )\n op.alter_column(\n \"chat_message\", \"parent_message_id\", new_column_name=\"parent_message\"\n )\n\n # Recreate agent sub question and sub query tables\n op.create_table(\n \"agent__sub_question\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"primary_question_id\", sa.Integer(), nullable=False),\n sa.Column(\"chat_session_id\", postgresql.UUID(as_uuid=True), nullable=False),\n sa.Column(\"sub_question\", sa.Text(), nullable=False),\n sa.Column(\"level\", sa.Integer(), nullable=False),\n sa.Column(\"level_question_num\", sa.Integer(), nullable=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"sub_answer\", sa.Text(), nullable=False),\n sa.Column(\"sub_question_doc_results\", postgresql.JSONB(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"primary_question_id\"], [\"chat_message.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"chat_session_id\"], [\"chat_session.id\"]),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_table(\n \"agent__sub_query\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"parent_question_id\", sa.Integer(), nullable=False),\n sa.Column(\"chat_session_id\", postgresql.UUID(as_uuid=True), nullable=False),\n sa.Column(\"sub_query\", sa.Text(), nullable=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"parent_question_id\"], [\"agent__sub_question.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"chat_session_id\"], [\"chat_session.id\"]),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_table(\n \"agent__sub_query__search_doc\",\n sa.Column(\"sub_query_id\", sa.Integer(), nullable=False),\n sa.Column(\"search_doc_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"sub_query_id\"], [\"agent__sub_query.id\"], ondelete=\"CASCADE\"\n ),\n sa.ForeignKeyConstraint([\"search_doc_id\"], [\"search_doc.id\"]),\n sa.PrimaryKeyConstraint(\"sub_query_id\", \"search_doc_id\"),\n )\n\n # Recreate research agent tables\n op.create_table(\n \"research_agent_iteration\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\"primary_question_id\", sa.Integer(), nullable=False),\n sa.Column(\"iteration_nr\", sa.Integer(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"purpose\", sa.String(), nullable=True),\n sa.Column(\"reasoning\", sa.String(), nullable=True),\n sa.ForeignKeyConstraint(\n [\"primary_question_id\"], [\"chat_message.id\"], ondelete=\"CASCADE\"\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\n \"primary_question_id\",\n \"iteration_nr\",\n name=\"_research_agent_iteration_unique_constraint\",\n ),\n )\n\n op.create_table(\n \"research_agent_iteration_sub_step\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\"primary_question_id\", sa.Integer(), nullable=False),\n sa.Column(\"iteration_nr\", sa.Integer(), nullable=False),\n sa.Column(\"iteration_sub_step_nr\", sa.Integer(), nullable=False),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"sub_step_instructions\", sa.String(), nullable=True),\n sa.Column(\"sub_step_tool_id\", sa.Integer(), nullable=True),\n sa.Column(\"reasoning\", sa.String(), nullable=True),\n sa.Column(\"sub_answer\", sa.String(), nullable=True),\n sa.Column(\"cited_doc_results\", postgresql.JSONB(), nullable=False),\n sa.Column(\"claims\", postgresql.JSONB(), nullable=True),\n sa.Column(\"is_web_fetch\", sa.Boolean(), nullable=True),\n sa.Column(\"queries\", postgresql.JSONB(), nullable=True),\n sa.Column(\"generated_images\", postgresql.JSONB(), nullable=True),\n sa.Column(\"additional_data\", postgresql.JSONB(), nullable=True),\n sa.Column(\"file_ids\", postgresql.JSONB(), nullable=True),\n sa.ForeignKeyConstraint(\n [\"primary_question_id\", \"iteration_nr\"],\n [\n \"research_agent_iteration.primary_question_id\",\n \"research_agent_iteration.iteration_nr\",\n ],\n ondelete=\"CASCADE\",\n ),\n sa.ForeignKeyConstraint([\"sub_step_tool_id\"], [\"tool.id\"], ondelete=\"SET NULL\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n" }, { "path": "backend/alembic/versions/a8c2065484e6_add_auto_scroll_to_user_model.py", "content": "\"\"\"add auto scroll to user model\n\nRevision ID: a8c2065484e6\nRevises: abe7378b8217\nCreate Date: 2024-11-22 17:34:09.690295\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"a8c2065484e6\"\ndown_revision = \"abe7378b8217\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\"auto_scroll\", sa.Boolean(), nullable=True, server_default=None),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"auto_scroll\")\n" }, { "path": "backend/alembic/versions/abbfec3a5ac5_merge_prompt_into_persona.py", "content": "\"\"\"merge prompt into persona\n\nRevision ID: abbfec3a5ac5\nRevises: 8818cf73fa1a\nCreate Date: 2024-12-19 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"abbfec3a5ac5\"\ndown_revision = \"8818cf73fa1a\"\nbranch_labels = None\ndepends_on = None\n\n\nMAX_PROMPT_LENGTH = 5_000_000\n\n\ndef upgrade() -> None:\n \"\"\"NOTE: Prompts without any Personas will just be lost.\"\"\"\n # Step 1: Add new columns to persona table (only if they don't exist)\n\n # Check if columns exist before adding them\n connection = op.get_bind()\n inspector = sa.inspect(connection)\n existing_columns = [col[\"name\"] for col in inspector.get_columns(\"persona\")]\n\n if \"system_prompt\" not in existing_columns:\n op.add_column(\n \"persona\",\n sa.Column(\n \"system_prompt\", sa.String(length=MAX_PROMPT_LENGTH), nullable=True\n ),\n )\n\n if \"task_prompt\" not in existing_columns:\n op.add_column(\n \"persona\",\n sa.Column(\n \"task_prompt\", sa.String(length=MAX_PROMPT_LENGTH), nullable=True\n ),\n )\n\n if \"datetime_aware\" not in existing_columns:\n op.add_column(\n \"persona\",\n sa.Column(\n \"datetime_aware\", sa.Boolean(), nullable=False, server_default=\"true\"\n ),\n )\n\n # Step 2: Migrate data from prompt table to persona table (only if tables exist)\n existing_tables = inspector.get_table_names()\n\n if \"prompt\" in existing_tables and \"persona__prompt\" in existing_tables:\n # For personas that have associated prompts, copy the prompt data\n op.execute(\n \"\"\"\n UPDATE persona\n SET\n system_prompt = p.system_prompt,\n task_prompt = p.task_prompt,\n datetime_aware = p.datetime_aware\n FROM (\n -- Get the first prompt for each persona (in case there are multiple)\n SELECT DISTINCT ON (pp.persona_id)\n pp.persona_id,\n pr.system_prompt,\n pr.task_prompt,\n pr.datetime_aware\n FROM persona__prompt pp\n JOIN prompt pr ON pp.prompt_id = pr.id\n ) p\n WHERE persona.id = p.persona_id\n \"\"\"\n )\n\n # Step 3: Update chat_message references\n # Since chat messages referenced prompt_id, we need to update them to use persona_id\n # This is complex as we need to map from prompt_id to persona_id\n\n # Check if chat_message has prompt_id column\n chat_message_columns = [\n col[\"name\"] for col in inspector.get_columns(\"chat_message\")\n ]\n if \"prompt_id\" in chat_message_columns:\n op.execute(\n \"\"\"\n ALTER TABLE chat_message\n DROP CONSTRAINT IF EXISTS chat_message__prompt_fk\n \"\"\"\n )\n op.drop_column(\"chat_message\", \"prompt_id\")\n\n # Step 4: Handle personas without prompts - set default values if needed (always run this)\n op.execute(\n \"\"\"\n UPDATE persona\n SET\n system_prompt = COALESCE(system_prompt, ''),\n task_prompt = COALESCE(task_prompt, '')\n WHERE system_prompt IS NULL OR task_prompt IS NULL\n \"\"\"\n )\n\n # Step 5: Drop the persona__prompt association table (if it exists)\n if \"persona__prompt\" in existing_tables:\n op.drop_table(\"persona__prompt\")\n\n # Step 6: Drop the prompt table (if it exists)\n if \"prompt\" in existing_tables:\n op.drop_table(\"prompt\")\n\n # Step 7: Make system_prompt and task_prompt non-nullable after migration (only if they exist)\n op.alter_column(\n \"persona\",\n \"system_prompt\",\n existing_type=sa.String(length=MAX_PROMPT_LENGTH),\n nullable=False,\n server_default=None,\n )\n\n op.alter_column(\n \"persona\",\n \"task_prompt\",\n existing_type=sa.String(length=MAX_PROMPT_LENGTH),\n nullable=False,\n server_default=None,\n )\n\n\ndef downgrade() -> None:\n # Step 1: Recreate the prompt table\n op.create_table(\n \"prompt\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", postgresql.UUID(as_uuid=True), nullable=True),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"description\", sa.String(), nullable=False),\n sa.Column(\"system_prompt\", sa.String(length=MAX_PROMPT_LENGTH), nullable=False),\n sa.Column(\"task_prompt\", sa.String(length=MAX_PROMPT_LENGTH), nullable=False),\n sa.Column(\n \"datetime_aware\", sa.Boolean(), nullable=False, server_default=\"true\"\n ),\n sa.Column(\n \"default_prompt\", sa.Boolean(), nullable=False, server_default=\"false\"\n ),\n sa.Column(\"deleted\", sa.Boolean(), nullable=False, server_default=\"false\"),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # Step 2: Recreate the persona__prompt association table\n op.create_table(\n \"persona__prompt\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"prompt_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"prompt_id\"],\n [\"prompt.id\"],\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"prompt_id\"),\n )\n\n # Step 3: Migrate data back from persona to prompt table\n op.execute(\n \"\"\"\n INSERT INTO prompt (\n name,\n description,\n system_prompt,\n task_prompt,\n datetime_aware,\n default_prompt,\n deleted,\n user_id\n )\n SELECT\n CONCAT('Prompt for ', name),\n description,\n system_prompt,\n task_prompt,\n datetime_aware,\n is_default_persona,\n deleted,\n user_id\n FROM persona\n WHERE system_prompt IS NOT NULL AND system_prompt != ''\n RETURNING id, name\n \"\"\"\n )\n\n # Step 4: Re-establish persona__prompt relationships\n op.execute(\n \"\"\"\n INSERT INTO persona__prompt (persona_id, prompt_id)\n SELECT\n p.id as persona_id,\n pr.id as prompt_id\n FROM persona p\n JOIN prompt pr ON pr.name = CONCAT('Prompt for ', p.name)\n WHERE p.system_prompt IS NOT NULL AND p.system_prompt != ''\n \"\"\"\n )\n\n # Step 5: Add prompt_id column back to chat_message\n op.add_column(\"chat_message\", sa.Column(\"prompt_id\", sa.Integer(), nullable=True))\n\n # Step 6: Re-establish foreign key constraint\n op.create_foreign_key(\n \"chat_message__prompt_fk\", \"chat_message\", \"prompt\", [\"prompt_id\"], [\"id\"]\n )\n\n # Step 7: Remove columns from persona table\n op.drop_column(\"persona\", \"datetime_aware\")\n op.drop_column(\"persona\", \"task_prompt\")\n op.drop_column(\"persona\", \"system_prompt\")\n" }, { "path": "backend/alembic/versions/abe7378b8217_add_indexing_trigger_to_cc_pair.py", "content": "\"\"\"add indexing trigger to cc_pair\n\nRevision ID: abe7378b8217\nRevises: 6d562f86c78b\nCreate Date: 2024-11-26 19:09:53.481171\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"abe7378b8217\"\ndown_revision = \"93560ba1b118\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"indexing_trigger\",\n sa.Enum(\"UPDATE\", \"REINDEX\", name=\"indexingmode\", native_enum=False),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"indexing_trigger\")\n" }, { "path": "backend/alembic/versions/ac5eaac849f9_add_last_pruned_to_connector_table.py", "content": "\"\"\"add last_pruned to the connector_credential_pair table\n\nRevision ID: ac5eaac849f9\nRevises: 52a219fb5233\nCreate Date: 2024-09-10 15:04:26.437118\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"ac5eaac849f9\"\ndown_revision = \"46b7a812670f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # last pruned represents the last time the connector was pruned\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\"last_pruned\", sa.DateTime(timezone=True), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"last_pruned\")\n" }, { "path": "backend/alembic/versions/acaab4ef4507_remove_inactive_ccpair_status_on_.py", "content": "\"\"\"remove inactive ccpair status on downgrade\n\nRevision ID: acaab4ef4507\nRevises: b388730a2899\nCreate Date: 2025-02-16 18:21:41.330212\n\n\"\"\"\n\nfrom alembic import op\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom sqlalchemy import update\n\n# revision identifiers, used by Alembic.\nrevision = \"acaab4ef4507\"\ndown_revision = \"b388730a2899\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n pass\n\n\ndef downgrade() -> None:\n op.execute(\n update(ConnectorCredentialPair)\n .where(ConnectorCredentialPair.status == ConnectorCredentialPairStatus.INVALID)\n .values(status=ConnectorCredentialPairStatus.ACTIVE)\n )\n" }, { "path": "backend/alembic/versions/ae62505e3acc_add_saml_accounts.py", "content": "\"\"\"Add SAML Accounts\n\nRevision ID: ae62505e3acc\nRevises: 7da543f5672f\nCreate Date: 2023-09-26 16:19:30.933183\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"ae62505e3acc\"\ndown_revision = \"7da543f5672f\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"saml\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\"encrypted_cookie\", sa.Text(), nullable=False),\n sa.Column(\"expires_at\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"encrypted_cookie\"),\n sa.UniqueConstraint(\"user_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"saml\")\n" }, { "path": "backend/alembic/versions/aeda5f2df4f6_add_pinned_assistants.py", "content": "\"\"\"add pinned assistants\n\nRevision ID: aeda5f2df4f6\nRevises: c5eae4a75a1b\nCreate Date: 2025-01-09 16:04:10.770636\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"aeda5f2df4f6\"\ndown_revision = \"c5eae4a75a1b\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\", sa.Column(\"pinned_assistants\", postgresql.JSONB(), nullable=True)\n )\n op.execute('UPDATE \"user\" SET pinned_assistants = chosen_assistants')\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"pinned_assistants\")\n" }, { "path": "backend/alembic/versions/b082fec533f0_make_last_attempt_status_nullable.py", "content": "\"\"\"Make 'last_attempt_status' nullable\n\nRevision ID: b082fec533f0\nRevises: df0c7ad8a076\nCreate Date: 2023-08-06 12:05:47.087325\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"b082fec533f0\"\ndown_revision = \"df0c7ad8a076\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"connector_credential_pair\",\n \"last_attempt_status\",\n existing_type=postgresql.ENUM(\n \"NOT_STARTED\",\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n name=\"indexingstatus\",\n ),\n nullable=True,\n )\n\n\ndef downgrade() -> None:\n # First, update any null values to a default value\n op.execute(\n \"UPDATE connector_credential_pair SET last_attempt_status = 'NOT_STARTED' WHERE last_attempt_status IS NULL\"\n )\n\n # Then, make the column non-nullable\n op.alter_column(\n \"connector_credential_pair\",\n \"last_attempt_status\",\n existing_type=postgresql.ENUM(\n \"NOT_STARTED\",\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n name=\"indexingstatus\",\n ),\n nullable=False,\n )\n" }, { "path": "backend/alembic/versions/b156fa702355_chat_reworked.py", "content": "\"\"\"Chat Reworked\n\nRevision ID: b156fa702355\nRevises: baf71f781b9e\nCreate Date: 2023-12-12 00:57:41.823371\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.dialects.postgresql import ENUM\nfrom onyx.configs.constants import DocumentSource\n\n# revision identifiers, used by Alembic.\nrevision = \"b156fa702355\"\ndown_revision = \"baf71f781b9e\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\nsearchtype_enum = ENUM(\n \"KEYWORD\", \"SEMANTIC\", \"HYBRID\", name=\"searchtype\", create_type=True\n)\nrecencybiassetting_enum = ENUM(\n \"FAVOR_RECENT\",\n \"BASE_DECAY\",\n \"NO_DECAY\",\n \"AUTO\",\n name=\"recencybiassetting\",\n create_type=True,\n)\n\n\ndef upgrade() -> None:\n bind = op.get_bind()\n searchtype_enum.create(bind)\n recencybiassetting_enum.create(bind)\n\n # This is irrecoverable, whatever\n op.execute(\"DELETE FROM chat_feedback\")\n op.execute(\"DELETE FROM document_retrieval_feedback\")\n\n op.create_table(\n \"search_doc\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"document_id\", sa.String(), nullable=False),\n sa.Column(\"chunk_ind\", sa.Integer(), nullable=False),\n sa.Column(\"semantic_id\", sa.String(), nullable=False),\n sa.Column(\"link\", sa.String(), nullable=True),\n sa.Column(\"blurb\", sa.String(), nullable=False),\n sa.Column(\"boost\", sa.Integer(), nullable=False),\n sa.Column(\n \"source_type\",\n sa.Enum(DocumentSource, native=False),\n nullable=False,\n ),\n sa.Column(\"hidden\", sa.Boolean(), nullable=False),\n sa.Column(\"score\", sa.Float(), nullable=False),\n sa.Column(\"match_highlights\", postgresql.ARRAY(sa.String()), nullable=False),\n sa.Column(\"updated_at\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\"primary_owners\", postgresql.ARRAY(sa.String()), nullable=True),\n sa.Column(\"secondary_owners\", postgresql.ARRAY(sa.String()), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"prompt\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.Column(\"description\", sa.String(), nullable=False),\n sa.Column(\"system_prompt\", sa.Text(), nullable=False),\n sa.Column(\"task_prompt\", sa.Text(), nullable=False),\n sa.Column(\"include_citations\", sa.Boolean(), nullable=False),\n sa.Column(\"datetime_aware\", sa.Boolean(), nullable=False),\n sa.Column(\"default_prompt\", sa.Boolean(), nullable=False),\n sa.Column(\"deleted\", sa.Boolean(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"persona__prompt\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"prompt_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"prompt_id\"],\n [\"prompt.id\"],\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"prompt_id\"),\n )\n\n # Changes to persona first so chat_sessions can have the right persona\n # The empty persona will be overwritten on server startup\n op.add_column(\n \"persona\",\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n )\n op.add_column(\n \"persona\",\n sa.Column(\n \"search_type\",\n searchtype_enum,\n nullable=True,\n ),\n )\n op.execute(\"UPDATE persona SET search_type = 'HYBRID'\")\n op.alter_column(\"persona\", \"search_type\", nullable=False)\n op.add_column(\n \"persona\",\n sa.Column(\"llm_relevance_filter\", sa.Boolean(), nullable=True),\n )\n op.execute(\"UPDATE persona SET llm_relevance_filter = TRUE\")\n op.alter_column(\"persona\", \"llm_relevance_filter\", nullable=False)\n op.add_column(\n \"persona\",\n sa.Column(\"llm_filter_extraction\", sa.Boolean(), nullable=True),\n )\n op.execute(\"UPDATE persona SET llm_filter_extraction = TRUE\")\n op.alter_column(\"persona\", \"llm_filter_extraction\", nullable=False)\n op.add_column(\n \"persona\",\n sa.Column(\n \"recency_bias\",\n recencybiassetting_enum,\n nullable=True,\n ),\n )\n op.execute(\"UPDATE persona SET recency_bias = 'BASE_DECAY'\")\n op.alter_column(\"persona\", \"recency_bias\", nullable=False)\n op.alter_column(\"persona\", \"description\", existing_type=sa.VARCHAR(), nullable=True)\n op.execute(\"UPDATE persona SET description = ''\")\n op.alter_column(\"persona\", \"description\", nullable=False)\n op.create_foreign_key(\"persona__user_fk\", \"persona\", \"user\", [\"user_id\"], [\"id\"])\n op.drop_column(\"persona\", \"datetime_aware\")\n op.drop_column(\"persona\", \"tools\")\n op.drop_column(\"persona\", \"hint_text\")\n op.drop_column(\"persona\", \"apply_llm_relevance_filter\")\n op.drop_column(\"persona\", \"retrieval_enabled\")\n op.drop_column(\"persona\", \"system_text\")\n\n # Need to create a persona row so fk can work\n result = bind.execute(sa.text(\"SELECT 1 FROM persona WHERE id = 0\"))\n exists = result.fetchone()\n if not exists:\n op.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona (\n id, user_id, name, description, search_type, num_chunks,\n llm_relevance_filter, llm_filter_extraction, recency_bias,\n llm_model_version_override, default_persona, deleted\n ) VALUES (\n 0, NULL, '', '', 'HYBRID', NULL,\n TRUE, TRUE, 'BASE_DECAY', NULL, TRUE, FALSE\n )\n \"\"\"\n )\n )\n delete_statement = sa.text(\n \"\"\"\n DELETE FROM persona\n WHERE name = 'Danswer' AND default_persona = TRUE AND id != 0\n \"\"\"\n )\n\n bind.execute(delete_statement)\n\n op.add_column(\n \"chat_feedback\",\n sa.Column(\"chat_message_id\", sa.Integer(), nullable=False),\n )\n op.drop_constraint(\n \"chat_feedback_chat_message_chat_session_id_chat_message_me_fkey\",\n \"chat_feedback\",\n type_=\"foreignkey\",\n )\n op.drop_column(\"chat_feedback\", \"chat_message_edit_number\")\n op.drop_column(\"chat_feedback\", \"chat_message_chat_session_id\")\n op.drop_column(\"chat_feedback\", \"chat_message_message_number\")\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"id\",\n sa.Integer(),\n primary_key=True,\n autoincrement=True,\n nullable=False,\n unique=True,\n ),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"parent_message\", sa.Integer(), nullable=True),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"latest_child_message\", sa.Integer(), nullable=True),\n )\n op.add_column(\n \"chat_message\", sa.Column(\"rephrased_query\", sa.Text(), nullable=True)\n )\n op.add_column(\"chat_message\", sa.Column(\"prompt_id\", sa.Integer(), nullable=True))\n op.add_column(\n \"chat_message\",\n sa.Column(\"citations\", postgresql.JSONB(astext_type=sa.Text()), nullable=True),\n )\n op.add_column(\"chat_message\", sa.Column(\"error\", sa.Text(), nullable=True))\n op.drop_constraint(\"fk_chat_message_persona_id\", \"chat_message\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"chat_message__prompt_fk\", \"chat_message\", \"prompt\", [\"prompt_id\"], [\"id\"]\n )\n op.drop_column(\"chat_message\", \"parent_edit_number\")\n op.drop_column(\"chat_message\", \"persona_id\")\n op.drop_column(\"chat_message\", \"reference_docs\")\n op.drop_column(\"chat_message\", \"edit_number\")\n op.drop_column(\"chat_message\", \"latest\")\n op.drop_column(\"chat_message\", \"message_number\")\n op.add_column(\"chat_session\", sa.Column(\"one_shot\", sa.Boolean(), nullable=True))\n op.execute(\"UPDATE chat_session SET one_shot = TRUE\")\n op.alter_column(\"chat_session\", \"one_shot\", nullable=False)\n op.alter_column(\n \"chat_session\",\n \"persona_id\",\n existing_type=sa.INTEGER(),\n nullable=True,\n )\n op.execute(\"UPDATE chat_session SET persona_id = 0\")\n op.alter_column(\"chat_session\", \"persona_id\", nullable=False)\n op.add_column(\n \"document_retrieval_feedback\",\n sa.Column(\"chat_message_id\", sa.Integer(), nullable=False),\n )\n op.drop_constraint(\n \"document_retrieval_feedback_qa_event_id_fkey\",\n \"document_retrieval_feedback\",\n type_=\"foreignkey\",\n )\n op.create_foreign_key(\n \"document_retrieval_feedback__chat_message_fk\",\n \"document_retrieval_feedback\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n )\n op.drop_column(\"document_retrieval_feedback\", \"qa_event_id\")\n\n # Relation table must be created after the other tables are correct\n op.create_table(\n \"chat_message__search_doc\",\n sa.Column(\"chat_message_id\", sa.Integer(), nullable=False),\n sa.Column(\"search_doc_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"chat_message_id\"],\n [\"chat_message.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"search_doc_id\"],\n [\"search_doc.id\"],\n ),\n sa.PrimaryKeyConstraint(\"chat_message_id\", \"search_doc_id\"),\n )\n\n # Needs to be created after chat_message id field is added\n op.create_foreign_key(\n \"chat_feedback__chat_message_fk\",\n \"chat_feedback\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n )\n\n op.drop_table(\"query_event\")\n\n\ndef downgrade() -> None:\n # NOTE: you will lose all chat history. This is to satisfy the non-nullable constraints\n # below\n op.execute(\"DELETE FROM chat_feedback\")\n op.execute(\"DELETE FROM chat_message__search_doc\")\n op.execute(\"DELETE FROM document_retrieval_feedback\")\n op.execute(\"DELETE FROM document_retrieval_feedback\")\n op.execute(\"DELETE FROM chat_message\")\n op.execute(\"DELETE FROM chat_session\")\n\n op.drop_constraint(\n \"chat_feedback__chat_message_fk\", \"chat_feedback\", type_=\"foreignkey\"\n )\n op.drop_constraint(\n \"document_retrieval_feedback__chat_message_fk\",\n \"document_retrieval_feedback\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\"persona__user_fk\", \"persona\", type_=\"foreignkey\")\n op.drop_constraint(\"chat_message__prompt_fk\", \"chat_message\", type_=\"foreignkey\")\n op.drop_constraint(\n \"chat_message__search_doc_chat_message_id_fkey\",\n \"chat_message__search_doc\",\n type_=\"foreignkey\",\n )\n op.add_column(\n \"persona\",\n sa.Column(\"system_text\", sa.TEXT(), autoincrement=False, nullable=True),\n )\n op.add_column(\n \"persona\",\n sa.Column(\n \"retrieval_enabled\",\n sa.BOOLEAN(),\n autoincrement=False,\n nullable=True,\n ),\n )\n op.execute(\"UPDATE persona SET retrieval_enabled = TRUE\")\n op.alter_column(\"persona\", \"retrieval_enabled\", nullable=False)\n op.add_column(\n \"persona\",\n sa.Column(\n \"apply_llm_relevance_filter\",\n sa.BOOLEAN(),\n autoincrement=False,\n nullable=True,\n ),\n )\n op.add_column(\n \"persona\",\n sa.Column(\"hint_text\", sa.TEXT(), autoincrement=False, nullable=True),\n )\n op.add_column(\n \"persona\",\n sa.Column(\n \"tools\",\n postgresql.JSONB(astext_type=sa.Text()),\n autoincrement=False,\n nullable=True,\n ),\n )\n op.add_column(\n \"persona\",\n sa.Column(\"datetime_aware\", sa.BOOLEAN(), autoincrement=False, nullable=True),\n )\n op.execute(\"UPDATE persona SET datetime_aware = TRUE\")\n op.alter_column(\"persona\", \"datetime_aware\", nullable=False)\n op.alter_column(\"persona\", \"description\", existing_type=sa.VARCHAR(), nullable=True)\n op.drop_column(\"persona\", \"recency_bias\")\n op.drop_column(\"persona\", \"llm_filter_extraction\")\n op.drop_column(\"persona\", \"llm_relevance_filter\")\n op.drop_column(\"persona\", \"search_type\")\n op.drop_column(\"persona\", \"user_id\")\n op.add_column(\n \"document_retrieval_feedback\",\n sa.Column(\"qa_event_id\", sa.INTEGER(), autoincrement=False, nullable=False),\n )\n op.drop_column(\"document_retrieval_feedback\", \"chat_message_id\")\n op.alter_column(\n \"chat_session\", \"persona_id\", existing_type=sa.INTEGER(), nullable=True\n )\n op.drop_column(\"chat_session\", \"one_shot\")\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"message_number\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=False,\n primary_key=True,\n ),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"latest\", sa.BOOLEAN(), autoincrement=False, nullable=False),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"edit_number\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=False,\n primary_key=True,\n ),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"reference_docs\",\n postgresql.JSONB(astext_type=sa.Text()),\n autoincrement=False,\n nullable=True,\n ),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\"persona_id\", sa.INTEGER(), autoincrement=False, nullable=True),\n )\n op.add_column(\n \"chat_message\",\n sa.Column(\n \"parent_edit_number\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=True,\n ),\n )\n op.create_foreign_key(\n \"fk_chat_message_persona_id\",\n \"chat_message\",\n \"persona\",\n [\"persona_id\"],\n [\"id\"],\n )\n op.drop_column(\"chat_message\", \"error\")\n op.drop_column(\"chat_message\", \"citations\")\n op.drop_column(\"chat_message\", \"prompt_id\")\n op.drop_column(\"chat_message\", \"rephrased_query\")\n op.drop_column(\"chat_message\", \"latest_child_message\")\n op.drop_column(\"chat_message\", \"parent_message\")\n op.drop_column(\"chat_message\", \"id\")\n op.add_column(\n \"chat_feedback\",\n sa.Column(\n \"chat_message_message_number\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=False,\n ),\n )\n op.add_column(\n \"chat_feedback\",\n sa.Column(\n \"chat_message_chat_session_id\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=False,\n primary_key=True,\n ),\n )\n op.add_column(\n \"chat_feedback\",\n sa.Column(\n \"chat_message_edit_number\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=False,\n ),\n )\n op.drop_column(\"chat_feedback\", \"chat_message_id\")\n op.create_table(\n \"query_event\",\n sa.Column(\"id\", sa.INTEGER(), autoincrement=True, nullable=False),\n sa.Column(\"query\", sa.VARCHAR(), autoincrement=False, nullable=False),\n sa.Column(\n \"selected_search_flow\",\n sa.VARCHAR(),\n autoincrement=False,\n nullable=True,\n ),\n sa.Column(\"llm_answer\", sa.VARCHAR(), autoincrement=False, nullable=True),\n sa.Column(\"feedback\", sa.VARCHAR(), autoincrement=False, nullable=True),\n sa.Column(\"user_id\", sa.UUID(), autoincrement=False, nullable=True),\n sa.Column(\n \"time_created\",\n postgresql.TIMESTAMP(timezone=True),\n server_default=sa.text(\"now()\"),\n autoincrement=False,\n nullable=False,\n ),\n sa.Column(\n \"retrieved_document_ids\",\n postgresql.ARRAY(sa.VARCHAR()),\n autoincrement=False,\n nullable=True,\n ),\n sa.Column(\"chat_session_id\", sa.INTEGER(), autoincrement=False, nullable=True),\n sa.ForeignKeyConstraint(\n [\"chat_session_id\"],\n [\"chat_session.id\"],\n name=\"fk_query_event_chat_session_id\",\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"], [\"user.id\"], name=\"query_event_user_id_fkey\"\n ),\n sa.PrimaryKeyConstraint(\"id\", name=\"query_event_pkey\"),\n )\n op.drop_table(\"chat_message__search_doc\")\n op.drop_table(\"persona__prompt\")\n op.drop_table(\"prompt\")\n op.drop_table(\"search_doc\")\n op.create_unique_constraint(\n \"uq_chat_message_combination\",\n \"chat_message\",\n [\"chat_session_id\", \"message_number\", \"edit_number\"],\n )\n op.create_foreign_key(\n \"chat_feedback_chat_message_chat_session_id_chat_message_me_fkey\",\n \"chat_feedback\",\n \"chat_message\",\n [\n \"chat_message_chat_session_id\",\n \"chat_message_message_number\",\n \"chat_message_edit_number\",\n ],\n [\"chat_session_id\", \"message_number\", \"edit_number\"],\n )\n op.create_foreign_key(\n \"document_retrieval_feedback_qa_event_id_fkey\",\n \"document_retrieval_feedback\",\n \"query_event\",\n [\"qa_event_id\"],\n [\"id\"],\n )\n\n op.execute(\"DROP TYPE IF EXISTS searchtype\")\n op.execute(\"DROP TYPE IF EXISTS recencybiassetting\")\n op.execute(\"DROP TYPE IF EXISTS documentsource\")\n" }, { "path": "backend/alembic/versions/b30353be4eec_add_mcp_auth_performer.py", "content": "\"\"\"add_mcp_auth_performer\n\nRevision ID: b30353be4eec\nRevises: 2b75d0a8ffcb\nCreate Date: 2025-09-13 14:58:08.413534\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom onyx.db.enums import MCPAuthenticationPerformer, MCPTransport\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b30353be4eec\"\ndown_revision = \"2b75d0a8ffcb\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"moving to a better way of handling auth performer and transport\"\"\"\n # Add nullable column first for backward compatibility\n op.add_column(\n \"mcp_server\",\n sa.Column(\n \"auth_performer\",\n sa.Enum(MCPAuthenticationPerformer, native_enum=False),\n nullable=True,\n ),\n )\n\n op.add_column(\n \"mcp_server\",\n sa.Column(\n \"transport\",\n sa.Enum(MCPTransport, native_enum=False),\n nullable=True,\n ),\n )\n\n # # Backfill values using existing data and inference rules\n bind = op.get_bind()\n\n # 1) OAUTH servers are always PER_USER\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE mcp_server\n SET auth_performer = 'PER_USER'\n WHERE auth_type = 'OAUTH'\n \"\"\"\n )\n )\n\n # 2) If there is no admin connection config, mark as ADMIN (and not set yet)\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE mcp_server\n SET auth_performer = 'ADMIN'\n WHERE admin_connection_config_id IS NULL\n AND auth_performer IS NULL\n \"\"\"\n )\n )\n\n # 3) If there exists any user-specific connection config (user_email != ''), mark as PER_USER\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE mcp_server AS ms\n SET auth_performer = 'PER_USER'\n FROM mcp_connection_config AS mcc\n WHERE mcc.mcp_server_id = ms.id\n AND COALESCE(mcc.user_email, '') <> ''\n AND ms.auth_performer IS NULL\n \"\"\"\n )\n )\n\n # 4) Default any remaining nulls to ADMIN (covers API_TOKEN admin-managed and NONE)\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE mcp_server\n SET auth_performer = 'ADMIN'\n WHERE auth_performer IS NULL\n \"\"\"\n )\n )\n\n # Finally, make the column non-nullable\n op.alter_column(\n \"mcp_server\",\n \"auth_performer\",\n existing_type=sa.Enum(MCPAuthenticationPerformer, native_enum=False),\n nullable=False,\n )\n\n # Backfill transport for existing rows to STREAMABLE_HTTP, then make non-nullable\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE mcp_server\n SET transport = 'STREAMABLE_HTTP'\n WHERE transport IS NULL\n \"\"\"\n )\n )\n\n op.alter_column(\n \"mcp_server\",\n \"transport\",\n existing_type=sa.Enum(MCPTransport, native_enum=False),\n nullable=False,\n )\n\n\ndef downgrade() -> None:\n \"\"\"remove cols\"\"\"\n op.drop_column(\"mcp_server\", \"transport\")\n op.drop_column(\"mcp_server\", \"auth_performer\")\n" }, { "path": "backend/alembic/versions/b329d00a9ea6_adding_assistant_specific_user_.py", "content": "\"\"\"Adding assistant-specific user preferences\n\nRevision ID: b329d00a9ea6\nRevises: f9b8c7d6e5a4\nCreate Date: 2025-08-26 23:14:44.592985\n\n\"\"\"\n\nfrom alembic import op\nimport fastapi_users_db_sqlalchemy\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"b329d00a9ea6\"\ndown_revision = \"f9b8c7d6e5a4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"assistant__user_specific_config\",\n sa.Column(\"assistant_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.Column(\"disabled_tool_ids\", postgresql.ARRAY(sa.Integer()), nullable=False),\n sa.ForeignKeyConstraint([\"assistant_id\"], [\"persona.id\"], ondelete=\"CASCADE\"),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"assistant_id\", \"user_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"assistant__user_specific_config\")\n" }, { "path": "backend/alembic/versions/b388730a2899_nullable_preferences.py", "content": "\"\"\"nullable preferences\n\nRevision ID: b388730a2899\nRevises: 1a03d2c2856b\nCreate Date: 2025-02-17 18:49:22.643902\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b388730a2899\"\ndown_revision = \"1a03d2c2856b\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"user\", \"temperature_override_enabled\", nullable=True)\n op.alter_column(\"user\", \"auto_scroll\", nullable=True)\n\n\ndef downgrade() -> None:\n # Ensure no null values before making columns non-nullable\n op.execute(\n 'UPDATE \"user\" SET temperature_override_enabled = false WHERE temperature_override_enabled IS NULL'\n )\n op.execute('UPDATE \"user\" SET auto_scroll = false WHERE auto_scroll IS NULL')\n\n op.alter_column(\"user\", \"temperature_override_enabled\", nullable=False)\n op.alter_column(\"user\", \"auto_scroll\", nullable=False)\n" }, { "path": "backend/alembic/versions/b4b7e1028dfd_grant_basic_to_existing_groups.py", "content": "\"\"\"grant_basic_to_existing_groups\n\nGrants the \"basic\" permission to all existing groups that don't already\nhave it. Every group should have at least \"basic\" so that its members\nget basic access when effective_permissions is backfilled.\n\nRevision ID: b4b7e1028dfd\nRevises: b7bcc991d722\nCreate Date: 2026-03-30 16:15:17.093498\n\n\"\"\"\n\nfrom collections.abc import Sequence\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b4b7e1028dfd\"\ndown_revision = \"b7bcc991d722\"\nbranch_labels: str | None = None\ndepends_on: str | Sequence[str] | None = None\n\nuser_group = sa.table(\n \"user_group\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"is_default\", sa.Boolean),\n)\n\npermission_grant = sa.table(\n \"permission_grant\",\n sa.column(\"group_id\", sa.Integer),\n sa.column(\"permission\", sa.String),\n sa.column(\"grant_source\", sa.String),\n sa.column(\"is_deleted\", sa.Boolean),\n)\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n already_has_basic = (\n sa.select(sa.literal(1))\n .select_from(permission_grant)\n .where(\n permission_grant.c.group_id == user_group.c.id,\n permission_grant.c.permission == \"basic\",\n )\n .exists()\n )\n\n groups_needing_basic = sa.select(\n user_group.c.id,\n sa.literal(\"basic\").label(\"permission\"),\n sa.literal(\"SYSTEM\").label(\"grant_source\"),\n sa.literal(False).label(\"is_deleted\"),\n ).where(\n user_group.c.is_default == sa.false(),\n ~already_has_basic,\n )\n\n conn.execute(\n permission_grant.insert().from_select(\n [\"group_id\", \"permission\", \"grant_source\", \"is_deleted\"],\n groups_needing_basic,\n )\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n\n non_default_group_ids = sa.select(user_group.c.id).where(\n user_group.c.is_default == sa.false()\n )\n\n conn.execute(\n permission_grant.delete().where(\n permission_grant.c.permission == \"basic\",\n permission_grant.c.grant_source == \"SYSTEM\",\n permission_grant.c.group_id.in_(non_default_group_ids),\n )\n )\n" }, { "path": "backend/alembic/versions/b4ef3ae0bf6e_add_user_oauth_token_to_slack_bot.py", "content": "\"\"\"add_user_oauth_token_to_slack_bot\n\nRevision ID: b4ef3ae0bf6e\nRevises: 505c488f6662\nCreate Date: 2025-08-26 17:47:41.788462\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b4ef3ae0bf6e\"\ndown_revision = \"505c488f6662\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add user_token column to slack_bot table\n op.add_column(\"slack_bot\", sa.Column(\"user_token\", sa.LargeBinary(), nullable=True))\n\n\ndef downgrade() -> None:\n # Remove user_token column from slack_bot table\n op.drop_column(\"slack_bot\", \"user_token\")\n" }, { "path": "backend/alembic/versions/b51c6844d1df_seed_memory_tool.py", "content": "\"\"\"seed_memory_tool and add enable_memory_tool to user\n\nRevision ID: b51c6844d1df\nRevises: 93c15d6a6fbb\nCreate Date: 2026-02-11 00:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b51c6844d1df\"\ndown_revision = \"93c15d6a6fbb\"\nbranch_labels = None\ndepends_on = None\n\n\nMEMORY_TOOL = {\n \"name\": \"MemoryTool\",\n \"display_name\": \"Add Memory\",\n \"description\": \"Save memories about the user for future conversations.\",\n \"in_code_tool_id\": \"MemoryTool\",\n \"enabled\": True,\n}\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n existing = conn.execute(\n sa.text(\n \"SELECT in_code_tool_id FROM tool WHERE in_code_tool_id = :in_code_tool_id\"\n ),\n {\"in_code_tool_id\": MEMORY_TOOL[\"in_code_tool_id\"]},\n ).fetchone()\n\n if existing:\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :name,\n display_name = :display_name,\n description = :description\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n MEMORY_TOOL,\n )\n else:\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)\n VALUES (:name, :display_name, :description, :in_code_tool_id, :enabled)\n \"\"\"\n ),\n MEMORY_TOOL,\n )\n\n op.add_column(\n \"user\",\n sa.Column(\n \"enable_memory_tool\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.true(),\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"enable_memory_tool\")\n\n conn = op.get_bind()\n conn.execute(\n sa.text(\"DELETE FROM tool WHERE in_code_tool_id = :in_code_tool_id\"),\n {\"in_code_tool_id\": MEMORY_TOOL[\"in_code_tool_id\"]},\n )\n" }, { "path": "backend/alembic/versions/b558f51620b4_pause_finished_user_file_connectors.py", "content": "\"\"\"Pause finished user file connectors\n\nRevision ID: b558f51620b4\nRevises: 90e3b9af7da4\nCreate Date: 2025-08-15 17:17:02.456704\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b558f51620b4\"\ndown_revision = \"90e3b9af7da4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Set all user file connector credential pairs with ACTIVE status to PAUSED\n # This ensures user files don't continue to run indexing tasks after processing\n op.execute(\n \"\"\"\n UPDATE connector_credential_pair\n SET status = 'PAUSED'\n WHERE is_user_file = true\n AND status = 'ACTIVE'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/b5c4d7e8f9a1_add_hierarchy_node_cc_pair_table.py", "content": "\"\"\"add hierarchy_node_by_connector_credential_pair table\n\nRevision ID: b5c4d7e8f9a1\nRevises: a3b8d9e2f1c4\nCreate Date: 2026-03-04\n\n\"\"\"\n\nimport sqlalchemy as sa\nfrom alembic import op\n\nrevision = \"b5c4d7e8f9a1\"\ndown_revision = \"a3b8d9e2f1c4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"hierarchy_node_by_connector_credential_pair\",\n sa.Column(\"hierarchy_node_id\", sa.Integer(), nullable=False),\n sa.Column(\"connector_id\", sa.Integer(), nullable=False),\n sa.Column(\"credential_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"hierarchy_node_id\"],\n [\"hierarchy_node.id\"],\n ondelete=\"CASCADE\",\n ),\n sa.ForeignKeyConstraint(\n [\"connector_id\", \"credential_id\"],\n [\n \"connector_credential_pair.connector_id\",\n \"connector_credential_pair.credential_id\",\n ],\n ondelete=\"CASCADE\",\n ),\n sa.PrimaryKeyConstraint(\"hierarchy_node_id\", \"connector_id\", \"credential_id\"),\n )\n op.create_index(\n \"ix_hierarchy_node_cc_pair_connector_credential\",\n \"hierarchy_node_by_connector_credential_pair\",\n [\"connector_id\", \"credential_id\"],\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"ix_hierarchy_node_cc_pair_connector_credential\",\n table_name=\"hierarchy_node_by_connector_credential_pair\",\n )\n op.drop_table(\"hierarchy_node_by_connector_credential_pair\")\n" }, { "path": "backend/alembic/versions/b728689f45b1_rename_persona_is_visible_to_is_listed_.py", "content": "\"\"\"rename persona is_visible to is_listed and featured to is_featured\n\nRevision ID: b728689f45b1\nRevises: 689433b0d8de\nCreate Date: 2026-03-23 12:36:26.607305\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b728689f45b1\"\ndown_revision = \"689433b0d8de\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"persona\", \"is_visible\", new_column_name=\"is_listed\")\n op.alter_column(\"persona\", \"featured\", new_column_name=\"is_featured\")\n\n\ndef downgrade() -> None:\n op.alter_column(\"persona\", \"is_listed\", new_column_name=\"is_visible\")\n op.alter_column(\"persona\", \"is_featured\", new_column_name=\"featured\")\n" }, { "path": "backend/alembic/versions/b72ed7a5db0e_remove_description_from_starter_messages.py", "content": "\"\"\"remove description from starter messages\n\nRevision ID: b72ed7a5db0e\nRevises: 33cb72ea4d80\nCreate Date: 2024-11-03 15:55:28.944408\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b72ed7a5db0e\"\ndown_revision = \"33cb72ea4d80\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET starter_messages = (\n SELECT jsonb_agg(elem - 'description')\n FROM jsonb_array_elements(starter_messages) elem\n )\n WHERE starter_messages IS NOT NULL\n AND jsonb_typeof(starter_messages) = 'array'\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n op.execute(\n sa.text(\n \"\"\"\n UPDATE persona\n SET starter_messages = (\n SELECT jsonb_agg(elem || '{\"description\": \"\"}')\n FROM jsonb_array_elements(starter_messages) elem\n )\n WHERE starter_messages IS NOT NULL\n AND jsonb_typeof(starter_messages) = 'array'\n \"\"\"\n )\n )\n" }, { "path": "backend/alembic/versions/b7a7eee5aa15_add_checkpointing_failure_handling.py", "content": "\"\"\"Add checkpointing/failure handling\n\nRevision ID: b7a7eee5aa15\nRevises: f39c5794c10a\nCreate Date: 2025-01-24 15:17:36.763172\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"b7a7eee5aa15\"\ndown_revision = \"f39c5794c10a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"index_attempt\",\n sa.Column(\"checkpoint_pointer\", sa.String(), nullable=True),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\"poll_range_start\", sa.DateTime(timezone=True), nullable=True),\n )\n op.add_column(\n \"index_attempt\",\n sa.Column(\"poll_range_end\", sa.DateTime(timezone=True), nullable=True),\n )\n\n op.create_index(\n \"ix_index_attempt_cc_pair_settings_poll\",\n \"index_attempt\",\n [\n \"connector_credential_pair_id\",\n \"search_settings_id\",\n \"status\",\n sa.text(\"time_updated DESC\"),\n ],\n )\n\n # Drop the old IndexAttemptError table\n op.drop_index(\"index_attempt_id\", table_name=\"index_attempt_errors\")\n op.drop_table(\"index_attempt_errors\")\n\n # Create the new version of the table\n op.create_table(\n \"index_attempt_errors\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"index_attempt_id\", sa.Integer(), nullable=False),\n sa.Column(\"connector_credential_pair_id\", sa.Integer(), nullable=False),\n sa.Column(\"document_id\", sa.String(), nullable=True),\n sa.Column(\"document_link\", sa.String(), nullable=True),\n sa.Column(\"entity_id\", sa.String(), nullable=True),\n sa.Column(\"failed_time_range_start\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\"failed_time_range_end\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\"failure_message\", sa.Text(), nullable=False),\n sa.Column(\"is_resolved\", sa.Boolean(), nullable=False, default=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"index_attempt_id\"],\n [\"index_attempt.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"connector_credential_pair_id\"],\n [\"connector_credential_pair.id\"],\n ),\n )\n\n\ndef downgrade() -> None:\n op.execute(\"SET lock_timeout = '5s'\")\n\n # try a few times to drop the table, this has been observed to fail due to other locks\n # blocking the drop\n NUM_TRIES = 10\n for i in range(NUM_TRIES):\n try:\n op.drop_table(\"index_attempt_errors\")\n break\n except Exception as e:\n if i == NUM_TRIES - 1:\n raise e\n print(f\"Error dropping table: {e}. Retrying...\")\n\n op.execute(\"SET lock_timeout = DEFAULT\")\n\n # Recreate the old IndexAttemptError table\n op.create_table(\n \"index_attempt_errors\",\n sa.Column(\"id\", sa.Integer(), primary_key=True),\n sa.Column(\"index_attempt_id\", sa.Integer(), nullable=True),\n sa.Column(\"batch\", sa.Integer(), nullable=True),\n sa.Column(\"doc_summaries\", postgresql.JSONB(), nullable=False),\n sa.Column(\"error_msg\", sa.Text(), nullable=True),\n sa.Column(\"traceback\", sa.Text(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n ),\n sa.ForeignKeyConstraint(\n [\"index_attempt_id\"],\n [\"index_attempt.id\"],\n ),\n )\n\n op.create_index(\n \"index_attempt_id\",\n \"index_attempt_errors\",\n [\"time_created\"],\n )\n\n op.drop_index(\"ix_index_attempt_cc_pair_settings_poll\")\n op.drop_column(\"index_attempt\", \"checkpoint_pointer\")\n op.drop_column(\"index_attempt\", \"poll_range_start\")\n op.drop_column(\"index_attempt\", \"poll_range_end\")\n" }, { "path": "backend/alembic/versions/b7bcc991d722_assign_users_to_default_groups.py", "content": "\"\"\"assign_users_to_default_groups\n\nRevision ID: b7bcc991d722\nRevises: 03d085c5c38d\nCreate Date: 2026-03-25 16:30:39.529301\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b7bcc991d722\"\ndown_revision = \"03d085c5c38d\"\nbranch_labels = None\ndepends_on = None\n\n# The no-auth placeholder user must NOT be assigned to default groups.\n# A database trigger (migrate_no_auth_data_to_user) will try to DELETE this\n# user when the first real user registers; group membership rows would cause\n# an FK violation on that DELETE.\nNO_AUTH_PLACEHOLDER_USER_UUID = \"00000000-0000-0000-0000-000000000001\"\n\n# Reflect table structures for use in DML\nuser_group_table = sa.table(\n \"user_group\",\n sa.column(\"id\", sa.Integer),\n sa.column(\"name\", sa.String),\n sa.column(\"is_default\", sa.Boolean),\n)\n\nuser_table = sa.table(\n \"user\",\n sa.column(\"id\", sa.Uuid),\n sa.column(\"role\", sa.String),\n sa.column(\"account_type\", sa.String),\n sa.column(\"is_active\", sa.Boolean),\n)\n\nuser__user_group_table = sa.table(\n \"user__user_group\",\n sa.column(\"user_group_id\", sa.Integer),\n sa.column(\"user_id\", sa.Uuid),\n)\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Look up default group IDs\n admin_row = conn.execute(\n sa.select(user_group_table.c.id).where(\n user_group_table.c.name == \"Admin\",\n user_group_table.c.is_default == True, # noqa: E712\n )\n ).fetchone()\n\n basic_row = conn.execute(\n sa.select(user_group_table.c.id).where(\n user_group_table.c.name == \"Basic\",\n user_group_table.c.is_default == True, # noqa: E712\n )\n ).fetchone()\n\n if admin_row is None:\n raise RuntimeError(\n \"Default 'Admin' group not found. \"\n \"Ensure migration 977e834c1427 (seed_default_groups) ran successfully.\"\n )\n\n if basic_row is None:\n raise RuntimeError(\n \"Default 'Basic' group not found. \"\n \"Ensure migration 977e834c1427 (seed_default_groups) ran successfully.\"\n )\n\n # Users with role=admin → Admin group\n # Include inactive users so reactivation doesn't require reconciliation.\n # Exclude non-human account types (mirrors assign_user_to_default_groups logic).\n admin_users = sa.select(\n sa.literal(admin_row[0]).label(\"user_group_id\"),\n user_table.c.id.label(\"user_id\"),\n ).where(\n user_table.c.role == \"ADMIN\",\n user_table.c.account_type.notin_([\"BOT\", \"EXT_PERM_USER\", \"ANONYMOUS\"]),\n user_table.c.id != NO_AUTH_PLACEHOLDER_USER_UUID,\n )\n op.execute(\n pg_insert(user__user_group_table)\n .from_select([\"user_group_id\", \"user_id\"], admin_users)\n .on_conflict_do_nothing(index_elements=[\"user_group_id\", \"user_id\"])\n )\n\n # STANDARD users (non-admin) and SERVICE_ACCOUNT users (role=basic) → Basic group\n # Include inactive users so reactivation doesn't require reconciliation.\n basic_users = sa.select(\n sa.literal(basic_row[0]).label(\"user_group_id\"),\n user_table.c.id.label(\"user_id\"),\n ).where(\n user_table.c.account_type.notin_([\"BOT\", \"EXT_PERM_USER\", \"ANONYMOUS\"]),\n user_table.c.id != NO_AUTH_PLACEHOLDER_USER_UUID,\n sa.or_(\n sa.and_(\n user_table.c.account_type == \"STANDARD\",\n user_table.c.role != \"ADMIN\",\n ),\n sa.and_(\n user_table.c.account_type == \"SERVICE_ACCOUNT\",\n user_table.c.role == \"BASIC\",\n ),\n ),\n )\n op.execute(\n pg_insert(user__user_group_table)\n .from_select([\"user_group_id\", \"user_id\"], basic_users)\n .on_conflict_do_nothing(index_elements=[\"user_group_id\", \"user_id\"])\n )\n\n\ndef downgrade() -> None:\n # Group memberships are left in place — removing them risks\n # deleting memberships that existed before this migration.\n pass\n" }, { "path": "backend/alembic/versions/b7c2b63c4a03_add_background_reindex_enabled_field.py", "content": "\"\"\"add background_reindex_enabled field\n\nRevision ID: b7c2b63c4a03\nRevises: f11b408e39d3\nCreate Date: 2024-03-26 12:34:56.789012\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom onyx.db.enums import EmbeddingPrecision\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b7c2b63c4a03\"\ndown_revision = \"f11b408e39d3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add background_reindex_enabled column with default value of True\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"background_reindex_enabled\",\n sa.Boolean(),\n nullable=False,\n server_default=\"true\",\n ),\n )\n\n # Add embedding_precision column with default value of FLOAT\n op.add_column(\n \"search_settings\",\n sa.Column(\n \"embedding_precision\",\n sa.Enum(EmbeddingPrecision, native_enum=False),\n nullable=False,\n server_default=EmbeddingPrecision.FLOAT.name,\n ),\n )\n\n # Add reduced_dimension column with default value of None\n op.add_column(\n \"search_settings\",\n sa.Column(\"reduced_dimension\", sa.Integer(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n # Remove the background_reindex_enabled column\n op.drop_column(\"search_settings\", \"background_reindex_enabled\")\n op.drop_column(\"search_settings\", \"embedding_precision\")\n op.drop_column(\"search_settings\", \"reduced_dimension\")\n" }, { "path": "backend/alembic/versions/b7ec9b5b505f_adjust_prompt_length.py", "content": "\"\"\"adjust prompt length\n\nRevision ID: b7ec9b5b505f\nRevises: abbfec3a5ac5\nCreate Date: 2025-09-10 18:51:15.629197\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b7ec9b5b505f\"\ndown_revision = \"abbfec3a5ac5\"\nbranch_labels = None\ndepends_on = None\n\n\nMAX_PROMPT_LENGTH = 5_000_000\n\n\ndef upgrade() -> None:\n # NOTE: need to run this since the previous migration PREVIOUSLY set the length to 8000\n op.alter_column(\n \"persona\",\n \"system_prompt\",\n existing_type=sa.String(length=8000),\n type_=sa.String(length=MAX_PROMPT_LENGTH),\n existing_nullable=False,\n )\n op.alter_column(\n \"persona\",\n \"task_prompt\",\n existing_type=sa.String(length=8000),\n type_=sa.String(length=MAX_PROMPT_LENGTH),\n existing_nullable=False,\n )\n\n\ndef downgrade() -> None:\n # Downgrade not necessary\n pass\n" }, { "path": "backend/alembic/versions/b85f02ec1308_fix_file_type_migration.py", "content": "\"\"\"fix-file-type-migration\n\nRevision ID: b85f02ec1308\nRevises: a3bfd0d64902\nCreate Date: 2024-05-31 18:09:26.658164\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"b85f02ec1308\"\ndown_revision = \"a3bfd0d64902\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE file_store\n SET file_origin = UPPER(file_origin)\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n # Let's not break anything on purpose :)\n pass\n" }, { "path": "backend/alembic/versions/b896bbd0d5a7_backfill_is_internet_data_to_false.py", "content": "\"\"\"backfill is_internet data to False\n\nRevision ID: b896bbd0d5a7\nRevises: 44f856ae2a4a\nCreate Date: 2024-07-16 15:21:05.718571\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"b896bbd0d5a7\"\ndown_revision = \"44f856ae2a4a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\"UPDATE search_doc SET is_internet = FALSE WHERE is_internet IS NULL\")\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/b8c9d0e1f2a3_drop_milestone_table.py", "content": "\"\"\"Drop milestone table\n\nRevision ID: b8c9d0e1f2a3\nRevises: a2b3c4d5e6f7\nCreate Date: 2025-12-18\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport fastapi_users_db_sqlalchemy\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"b8c9d0e1f2a3\"\ndown_revision = \"a2b3c4d5e6f7\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_table(\"milestone\")\n\n\ndef downgrade() -> None:\n op.create_table(\n \"milestone\",\n sa.Column(\"id\", sa.UUID(), nullable=False),\n sa.Column(\"tenant_id\", sa.String(), nullable=True),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\"event_type\", sa.String(), nullable=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"event_tracker\", postgresql.JSONB(), nullable=True),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"event_type\", name=\"uq_milestone_event_type\"),\n )\n" }, { "path": "backend/alembic/versions/ba98eba0f66a_add_support_for_litellm_proxy_in_.py", "content": "\"\"\"add support for litellm proxy in reranking\n\nRevision ID: ba98eba0f66a\nRevises: bceb1e139447\nCreate Date: 2024-09-06 10:36:04.507332\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"ba98eba0f66a\"\ndown_revision = \"bceb1e139447\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"search_settings\", sa.Column(\"rerank_api_url\", sa.String(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"search_settings\", \"rerank_api_url\")\n" }, { "path": "backend/alembic/versions/baf71f781b9e_add_llm_model_version_override_to_.py", "content": "\"\"\"Add llm_model_version_override to Persona\n\nRevision ID: baf71f781b9e\nRevises: 50b683a8295c\nCreate Date: 2023-12-06 21:56:50.286158\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"baf71f781b9e\"\ndown_revision = \"50b683a8295c\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"llm_model_version_override\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"llm_model_version_override\")\n" }, { "path": "backend/alembic/versions/bc9771dccadf_create_usage_reports_table.py", "content": "\"\"\"create usage reports table\n\nRevision ID: bc9771dccadf\nRevises: 0568ccf46a6b\nCreate Date: 2024-06-18 10:04:26.800282\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport fastapi_users_db_sqlalchemy\n\n# revision identifiers, used by Alembic.\nrevision = \"bc9771dccadf\"\ndown_revision = \"0568ccf46a6b\"\n\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"usage_reports\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"report_name\", sa.String(), nullable=False),\n sa.Column(\n \"requestor_user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"period_from\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\"period_to\", sa.DateTime(timezone=True), nullable=True),\n sa.ForeignKeyConstraint(\n [\"report_name\"],\n [\"file_store.file_name\"],\n ),\n sa.ForeignKeyConstraint(\n [\"requestor_user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"usage_reports\")\n" }, { "path": "backend/alembic/versions/bceb1e139447_add_base_url_to_cloudembeddingprovider.py", "content": "\"\"\"Add base_url to CloudEmbeddingProvider\n\nRevision ID: bceb1e139447\nRevises: a3795dce87be\nCreate Date: 2024-08-28 17:00:52.554580\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"bceb1e139447\"\ndown_revision = \"a3795dce87be\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"embedding_provider\", sa.Column(\"api_url\", sa.String(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"embedding_provider\", \"api_url\")\n" }, { "path": "backend/alembic/versions/bd2921608c3a_non_nullable_default_persona.py", "content": "\"\"\"non nullable default persona\n\nRevision ID: bd2921608c3a\nRevises: 797089dfb4d2\nCreate Date: 2024-09-20 10:28:37.992042\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"bd2921608c3a\"\ndown_revision = \"797089dfb4d2\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Set existing NULL values to False\n op.execute(\n \"UPDATE persona SET is_default_persona = FALSE WHERE is_default_persona IS NULL\"\n )\n\n # Alter the column to be not nullable with a default value of False\n op.alter_column(\n \"persona\",\n \"is_default_persona\",\n existing_type=sa.Boolean(),\n nullable=False,\n server_default=sa.text(\"false\"),\n )\n\n\ndef downgrade() -> None:\n # Revert the changes\n op.alter_column(\n \"persona\",\n \"is_default_persona\",\n existing_type=sa.Boolean(),\n nullable=True,\n server_default=None,\n )\n" }, { "path": "backend/alembic/versions/bd7c3bf8beba_migrate_agent_responses_to_research_.py", "content": "\"\"\"migrate_agent_sub_questions_to_research_iterations\n\nRevision ID: bd7c3bf8beba\nRevises: f8a9b2c3d4e5\nCreate Date: 2025-08-18 11:33:27.098287\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"bd7c3bf8beba\"\ndown_revision = \"f8a9b2c3d4e5\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Get connection to execute raw SQL\n connection = op.get_bind()\n\n # First, insert data into research_agent_iteration table\n # This creates one iteration record per primary_question_id using the earliest time_created\n connection.execute(\n sa.text(\n \"\"\"\n INSERT INTO research_agent_iteration (primary_question_id, created_at, iteration_nr, purpose, reasoning)\n SELECT\n primary_question_id,\n MIN(time_created) as created_at,\n 1 as iteration_nr,\n 'Generating and researching subquestions' as purpose,\n '(No previous reasoning)' as reasoning\n FROM agent__sub_question\n JOIN chat_message on agent__sub_question.primary_question_id = chat_message.id\n WHERE primary_question_id IS NOT NULL\n AND chat_message.is_agentic = true\n GROUP BY primary_question_id\n ON CONFLICT DO NOTHING;\n \"\"\"\n )\n )\n\n # Then, insert data into research_agent_iteration_sub_step table\n # This migrates each sub-question as a sub-step\n connection.execute(\n sa.text(\n \"\"\"\n INSERT INTO research_agent_iteration_sub_step (\n primary_question_id,\n iteration_nr,\n iteration_sub_step_nr,\n created_at,\n sub_step_instructions,\n sub_step_tool_id,\n sub_answer,\n cited_doc_results\n )\n SELECT\n primary_question_id,\n 1 as iteration_nr,\n level_question_num as iteration_sub_step_nr,\n time_created as created_at,\n sub_question as sub_step_instructions,\n 1 as sub_step_tool_id,\n sub_answer,\n sub_question_doc_results as cited_doc_results\n FROM agent__sub_question\n JOIN chat_message on agent__sub_question.primary_question_id = chat_message.id\n WHERE chat_message.is_agentic = true\n AND primary_question_id IS NOT NULL\n ON CONFLICT DO NOTHING;\n \"\"\"\n )\n )\n\n # Update chat_message records: set legacy agentic type and answer purpose for existing agentic messages\n connection.execute(\n sa.text(\n \"\"\"\n UPDATE chat_message\n SET research_answer_purpose = 'ANSWER'\n WHERE is_agentic = true\n AND research_type IS NULL and\n message_type = 'ASSISTANT';\n \"\"\"\n )\n )\n connection.execute(\n sa.text(\n \"\"\"\n UPDATE chat_message\n SET research_type = 'LEGACY_AGENTIC'\n WHERE is_agentic = true\n AND research_type IS NULL;\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n # Get connection to execute raw SQL\n connection = op.get_bind()\n\n # Note: This downgrade removes all research agent iteration data\n # There's no way to perfectly restore the original agent__sub_question data\n # if it was deleted after this migration\n\n # Delete all research_agent_iteration_sub_step records that were migrated\n connection.execute(\n sa.text(\n \"\"\"\n DELETE FROM research_agent_iteration_sub_step\n USING chat_message\n WHERE research_agent_iteration_sub_step.primary_question_id = chat_message.id\n AND chat_message.research_type = 'LEGACY_AGENTIC';\n \"\"\"\n )\n )\n\n # Delete all research_agent_iteration records that were migrated\n connection.execute(\n sa.text(\n \"\"\"\n DELETE FROM research_agent_iteration\n USING chat_message\n WHERE research_agent_iteration.primary_question_id = chat_message.id\n AND chat_message.research_type = 'LEGACY_AGENTIC';\n \"\"\"\n )\n )\n\n # Revert chat_message updates: clear research fields for legacy agentic messages\n connection.execute(\n sa.text(\n \"\"\"\n UPDATE chat_message\n SET research_type = NULL,\n research_answer_purpose = NULL\n WHERE is_agentic = true\n AND research_type = 'LEGACY_AGENTIC'\n AND message_type = 'ASSISTANT';\n \"\"\"\n )\n )\n" }, { "path": "backend/alembic/versions/be2ab2aa50ee_fix_capitalization.py", "content": "\"\"\"fix_capitalization\n\nRevision ID: be2ab2aa50ee\nRevises: 369644546676\nCreate Date: 2025-01-10 13:13:26.228960\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"be2ab2aa50ee\"\ndown_revision = \"369644546676\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE document\n SET\n external_user_group_ids = ARRAY(\n SELECT LOWER(unnest(external_user_group_ids))\n ),\n last_modified = NOW()\n WHERE\n external_user_group_ids IS NOT NULL\n AND external_user_group_ids::text[] <> ARRAY(\n SELECT LOWER(unnest(external_user_group_ids))\n )::text[]\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n # No way to cleanly persist the bad state through an upgrade/downgrade\n # cycle, so we just pass\n pass\n" }, { "path": "backend/alembic/versions/be87a654d5af_persona_new_default_model_configuration_.py", "content": "\"\"\"Persona new default model configuration id column\n\nRevision ID: be87a654d5af\nRevises: e7f8a9b0c1d2\nCreate Date: 2026-01-30 11:14:17.306275\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"be87a654d5af\"\ndown_revision = \"e7f8a9b0c1d2\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"default_model_configuration_id\", sa.Integer(), nullable=True),\n )\n op.create_foreign_key(\n \"fk_persona_default_model_configuration_id\",\n \"persona\",\n \"model_configuration\",\n [\"default_model_configuration_id\"],\n [\"id\"],\n ondelete=\"SET NULL\",\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\n \"fk_persona_default_model_configuration_id\", \"persona\", type_=\"foreignkey\"\n )\n\n op.drop_column(\"persona\", \"default_model_configuration_id\")\n" }, { "path": "backend/alembic/versions/bf7a81109301_delete_input_prompts.py", "content": "\"\"\"delete_input_prompts\n\nRevision ID: bf7a81109301\nRevises: f7a894b06d02\nCreate Date: 2024-12-09 12:00:49.884228\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport fastapi_users_db_sqlalchemy\n\n\n# revision identifiers, used by Alembic.\nrevision = \"bf7a81109301\"\ndown_revision = \"f7a894b06d02\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_table(\"inputprompt__user\")\n op.drop_table(\"inputprompt\")\n\n\ndef downgrade() -> None:\n op.create_table(\n \"inputprompt\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\"prompt\", sa.String(), nullable=False),\n sa.Column(\"content\", sa.String(), nullable=False),\n sa.Column(\"active\", sa.Boolean(), nullable=False),\n sa.Column(\"is_public\", sa.Boolean(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"inputprompt__user\",\n sa.Column(\"input_prompt_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"input_prompt_id\"],\n [\"inputprompt.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"inputprompt.id\"],\n ),\n sa.PrimaryKeyConstraint(\"input_prompt_id\", \"user_id\"),\n )\n" }, { "path": "backend/alembic/versions/c0aab6edb6dd_delete_workspace.py", "content": "\"\"\"delete workspace\n\nRevision ID: c0aab6edb6dd\nRevises: 35e518e0ddf4\nCreate Date: 2024-12-17 14:37:07.660631\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"c0aab6edb6dd\"\ndown_revision = \"35e518e0ddf4\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = connector_specific_config - 'workspace'\n WHERE source = 'SLACK'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n import json\n from sqlalchemy import text\n from slack_sdk import WebClient\n\n conn = op.get_bind()\n\n # Fetch all Slack credentials\n creds_result = conn.execute(\n text(\"SELECT id, credential_json FROM credential WHERE source = 'SLACK'\")\n )\n all_slack_creds = creds_result.fetchall()\n if not all_slack_creds:\n return\n\n for cred_row in all_slack_creds:\n credential_id, credential_json = cred_row\n\n credential_json = (\n credential_json.tobytes().decode(\"utf-8\")\n if isinstance(credential_json, memoryview)\n else credential_json.decode(\"utf-8\")\n )\n credential_data = json.loads(credential_json)\n slack_bot_token = credential_data.get(\"slack_bot_token\")\n if not slack_bot_token:\n print(\n f\"No slack_bot_token found for credential {credential_id}. \"\n \"Your Slack connector will not function until you upgrade and provide a valid token.\"\n )\n continue\n\n client = WebClient(token=slack_bot_token)\n try:\n auth_response = client.auth_test()\n workspace = auth_response[\"url\"].split(\"//\")[1].split(\".\")[0]\n\n # Update only the connectors linked to this credential\n # (and which are Slack connectors).\n op.execute(\n f\"\"\"\n UPDATE connector AS c\n SET connector_specific_config = jsonb_set(\n connector_specific_config,\n '{{workspace}}',\n to_jsonb('{workspace}'::text)\n )\n FROM connector_credential_pair AS ccp\n WHERE ccp.connector_id = c.id\n AND c.source = 'SLACK'\n AND ccp.credential_id = {credential_id}\n \"\"\"\n )\n except Exception:\n print(\n f\"We were unable to get the workspace url for your Slack Connector with id {credential_id}.\"\n )\n print(\"This connector will no longer work until you upgrade.\")\n continue\n" }, { "path": "backend/alembic/versions/c0c937d5c9e5_llm_provider_deprecate_fields.py", "content": "\"\"\"llm provider deprecate fields\n\nRevision ID: c0c937d5c9e5\nRevises: 8ffcc2bcfc11\nCreate Date: 2026-02-25 17:35:46.125102\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"c0c937d5c9e5\"\ndown_revision = \"8ffcc2bcfc11\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Make default_model_name nullable (was NOT NULL)\n op.alter_column(\n \"llm_provider\",\n \"default_model_name\",\n existing_type=sa.String(),\n nullable=True,\n )\n\n # Drop unique constraint on is_default_provider (defaults now tracked via LLMModelFlow)\n op.drop_constraint(\n \"llm_provider_is_default_provider_key\",\n \"llm_provider\",\n type_=\"unique\",\n )\n\n # Remove server_default from is_default_vision_provider (was server_default=false())\n op.alter_column(\n \"llm_provider\",\n \"is_default_vision_provider\",\n existing_type=sa.Boolean(),\n server_default=None,\n )\n\n\ndef downgrade() -> None:\n # Restore default_model_name to NOT NULL (set empty string for any NULLs first)\n op.execute(\n \"UPDATE llm_provider SET default_model_name = '' WHERE default_model_name IS NULL\"\n )\n op.alter_column(\n \"llm_provider\",\n \"default_model_name\",\n existing_type=sa.String(),\n nullable=False,\n )\n\n # Restore unique constraint on is_default_provider\n op.create_unique_constraint(\n \"llm_provider_is_default_provider_key\",\n \"llm_provider\",\n [\"is_default_provider\"],\n )\n\n # Restore server_default for is_default_vision_provider\n op.alter_column(\n \"llm_provider\",\n \"is_default_vision_provider\",\n existing_type=sa.Boolean(),\n server_default=sa.false(),\n )\n" }, { "path": "backend/alembic/versions/c0fd6e4da83a_add_recent_assistants.py", "content": "\"\"\"add recent assistants\n\nRevision ID: c0fd6e4da83a\nRevises: b72ed7a5db0e\nCreate Date: 2024-11-03 17:28:54.916618\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"c0fd6e4da83a\"\ndown_revision = \"b72ed7a5db0e\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"recent_assistants\", postgresql.JSONB(), server_default=\"[]\", nullable=False\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"recent_assistants\")\n" }, { "path": "backend/alembic/versions/c18cdf4b497e_add_standard_answer_tables.py", "content": "\"\"\"Add standard_answer tables\n\nRevision ID: c18cdf4b497e\nRevises: 3a7802814195\nCreate Date: 2024-06-06 15:15:02.000648\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"c18cdf4b497e\"\ndown_revision = \"3a7802814195\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"standard_answer\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"keyword\", sa.String(), nullable=False),\n sa.Column(\"answer\", sa.String(), nullable=False),\n sa.Column(\"active\", sa.Boolean(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"keyword\"),\n )\n op.create_table(\n \"standard_answer_category\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"name\", sa.String(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.UniqueConstraint(\"name\"),\n )\n op.create_table(\n \"standard_answer__standard_answer_category\",\n sa.Column(\"standard_answer_id\", sa.Integer(), nullable=False),\n sa.Column(\"standard_answer_category_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"standard_answer_category_id\"],\n [\"standard_answer_category.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"standard_answer_id\"],\n [\"standard_answer.id\"],\n ),\n sa.PrimaryKeyConstraint(\"standard_answer_id\", \"standard_answer_category_id\"),\n )\n op.create_table(\n \"slack_bot_config__standard_answer_category\",\n sa.Column(\"slack_bot_config_id\", sa.Integer(), nullable=False),\n sa.Column(\"standard_answer_category_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"slack_bot_config_id\"],\n [\"slack_bot_config.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"standard_answer_category_id\"],\n [\"standard_answer_category.id\"],\n ),\n sa.PrimaryKeyConstraint(\"slack_bot_config_id\", \"standard_answer_category_id\"),\n )\n\n op.add_column(\n \"chat_session\", sa.Column(\"slack_thread_id\", sa.String(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_session\", \"slack_thread_id\")\n\n op.drop_table(\"slack_bot_config__standard_answer_category\")\n op.drop_table(\"standard_answer__standard_answer_category\")\n op.drop_table(\"standard_answer_category\")\n op.drop_table(\"standard_answer\")\n" }, { "path": "backend/alembic/versions/c1d2e3f4a5b6_add_deep_research_tool.py", "content": "\"\"\"add_deep_research_tool\n\nRevision ID: c1d2e3f4a5b6\nRevises: b8c9d0e1f2a3\nCreate Date: 2025-12-18 16:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"c1d2e3f4a5b6\"\ndown_revision = \"b8c9d0e1f2a3\"\nbranch_labels = None\ndepends_on = None\n\n\nDEEP_RESEARCH_TOOL = {\n \"name\": \"ResearchAgent\",\n \"display_name\": \"Research Agent\",\n \"description\": \"The Research Agent is a sub-agent that conducts research on a specific topic.\",\n \"in_code_tool_id\": \"ResearchAgent\",\n}\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)\n VALUES (:name, :display_name, :description, :in_code_tool_id, false)\n \"\"\"\n ),\n DEEP_RESEARCH_TOOL,\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n conn.execute(\n sa.text(\n \"\"\"\n DELETE FROM tool\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n {\"in_code_tool_id\": DEEP_RESEARCH_TOOL[\"in_code_tool_id\"]},\n )\n" }, { "path": "backend/alembic/versions/c5b692fa265c_add_index_attempt_errors_table.py", "content": "\"\"\"Add index_attempt_errors table\n\nRevision ID: c5b692fa265c\nRevises: 4a951134c801\nCreate Date: 2024-08-08 14:06:39.581972\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"c5b692fa265c\"\ndown_revision = \"4a951134c801\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"index_attempt_errors\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"index_attempt_id\", sa.Integer(), nullable=True),\n sa.Column(\"batch\", sa.Integer(), nullable=True),\n sa.Column(\n \"doc_summaries\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=False,\n ),\n sa.Column(\"error_msg\", sa.Text(), nullable=True),\n sa.Column(\"traceback\", sa.Text(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"index_attempt_id\"],\n [\"index_attempt.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_index(\n \"index_attempt_id\",\n \"index_attempt_errors\",\n [\"time_created\"],\n unique=False,\n )\n # ### end Alembic commands ###\n\n\ndef downgrade() -> None:\n # ### commands auto generated by Alembic - please adjust! ###\n op.drop_index(\"index_attempt_id\", table_name=\"index_attempt_errors\")\n op.drop_table(\"index_attempt_errors\")\n # ### end Alembic commands ###\n" }, { "path": "backend/alembic/versions/c5eae4a75a1b_add_chat_message__standard_answer_table.py", "content": "\"\"\"Add chat_message__standard_answer table\n\nRevision ID: c5eae4a75a1b\nRevises: 0f7ff6d75b57\nCreate Date: 2025-01-15 14:08:49.688998\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"c5eae4a75a1b\"\ndown_revision = \"0f7ff6d75b57\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"chat_message__standard_answer\",\n sa.Column(\"chat_message_id\", sa.Integer(), nullable=False),\n sa.Column(\"standard_answer_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"chat_message_id\"],\n [\"chat_message.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"standard_answer_id\"],\n [\"standard_answer.id\"],\n ),\n sa.PrimaryKeyConstraint(\"chat_message_id\", \"standard_answer_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"chat_message__standard_answer\")\n" }, { "path": "backend/alembic/versions/c7bf5721733e_add_has_been_indexed_to_.py", "content": "\"\"\"Add has_been_indexed to DocumentByConnectorCredentialPair\n\nRevision ID: c7bf5721733e\nRevises: fec3db967bf7\nCreate Date: 2025-01-13 12:39:05.831693\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"c7bf5721733e\"\ndown_revision = \"027381bce97c\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # assume all existing rows have been indexed, no better approach\n op.add_column(\n \"document_by_connector_credential_pair\",\n sa.Column(\"has_been_indexed\", sa.Boolean(), nullable=True),\n )\n op.execute(\n \"UPDATE document_by_connector_credential_pair SET has_been_indexed = TRUE\"\n )\n op.alter_column(\n \"document_by_connector_credential_pair\",\n \"has_been_indexed\",\n nullable=False,\n )\n\n # Add index to optimize get_document_counts_for_cc_pairs query pattern\n op.create_index(\n \"idx_document_cc_pair_counts\",\n \"document_by_connector_credential_pair\",\n [\"connector_id\", \"credential_id\", \"has_been_indexed\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n # Remove the index first before removing the column\n op.drop_index(\n \"idx_document_cc_pair_counts\",\n table_name=\"document_by_connector_credential_pair\",\n )\n op.drop_column(\"document_by_connector_credential_pair\", \"has_been_indexed\")\n" }, { "path": "backend/alembic/versions/c7e9f4a3b2d1_add_python_tool.py", "content": "\"\"\"add_python_tool\n\nRevision ID: c7e9f4a3b2d1\nRevises: 3c9a65f1207f\nCreate Date: 2025-11-08 00:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"c7e9f4a3b2d1\"\ndown_revision = \"3c9a65f1207f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n \"\"\"Add PythonTool to built-in tools\"\"\"\n conn = op.get_bind()\n\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)\n VALUES (:name, :display_name, :description, :in_code_tool_id, :enabled)\n \"\"\"\n ),\n {\n \"name\": \"PythonTool\",\n # in the UI, call it `Code Interpreter` since this is a well known term for this tool\n \"display_name\": \"Code Interpreter\",\n \"description\": (\n \"The Code Interpreter Action allows the assistant to execute \"\n \"Python code in a secure, isolated environment for data analysis, \"\n \"computation, visualization, and file processing.\"\n ),\n \"in_code_tool_id\": \"PythonTool\",\n \"enabled\": True,\n },\n )\n\n # needed to store files generated by the python tool\n op.add_column(\n \"research_agent_iteration_sub_step\",\n sa.Column(\n \"file_ids\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n \"\"\"Remove PythonTool from built-in tools\"\"\"\n conn = op.get_bind()\n\n conn.execute(\n sa.text(\n \"\"\"\n DELETE FROM tool\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n {\n \"in_code_tool_id\": \"PythonTool\",\n },\n )\n\n op.drop_column(\"research_agent_iteration_sub_step\", \"file_ids\")\n" }, { "path": "backend/alembic/versions/c7f2e1b4a9d3_add_sharing_scope_to_build_session.py", "content": "\"\"\"add sharing_scope to build_session\n\nRevision ID: c7f2e1b4a9d3\nRevises: 19c0ccb01687\nCreate Date: 2026-02-17 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nrevision = \"c7f2e1b4a9d3\"\ndown_revision = \"19c0ccb01687\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"build_session\",\n sa.Column(\n \"sharing_scope\",\n sa.String(),\n nullable=False,\n server_default=\"private\",\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"build_session\", \"sharing_scope\")\n" }, { "path": "backend/alembic/versions/c8a93a2af083_personalization_user_info.py", "content": "\"\"\"personalization_user_info\n\nRevision ID: c8a93a2af083\nRevises: 6f4f86aef280\nCreate Date: 2025-10-14 15:59:03.577343\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n\n# revision identifiers, used by Alembic.\nrevision = \"c8a93a2af083\"\ndown_revision = \"6f4f86aef280\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\"personal_name\", sa.String(), nullable=True),\n )\n op.add_column(\n \"user\",\n sa.Column(\"personal_role\", sa.String(), nullable=True),\n )\n op.add_column(\n \"user\",\n sa.Column(\n \"use_memories\",\n sa.Boolean(),\n nullable=False,\n server_default=sa.true(),\n ),\n )\n\n op.create_table(\n \"memory\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", postgresql.UUID(as_uuid=True), nullable=False),\n sa.Column(\"memory_text\", sa.Text(), nullable=False),\n sa.Column(\"conversation_id\", postgresql.UUID(as_uuid=True), nullable=True),\n sa.Column(\"message_id\", sa.Integer(), nullable=True),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.ForeignKeyConstraint([\"user_id\"], [\"user.id\"], ondelete=\"CASCADE\"),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n op.create_index(\"ix_memory_user_id\", \"memory\", [\"user_id\"])\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_memory_user_id\", table_name=\"memory\")\n op.drop_table(\"memory\")\n\n op.drop_column(\"user\", \"use_memories\")\n op.drop_column(\"user\", \"personal_role\")\n op.drop_column(\"user\", \"personal_name\")\n" }, { "path": "backend/alembic/versions/c99d76fcd298_add_nullable_to_persona_id_in_chat_.py", "content": "\"\"\"add nullable to persona id in Chat Session\n\nRevision ID: c99d76fcd298\nRevises: 5c7fdadae813\nCreate Date: 2024-07-09 19:27:01.579697\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"c99d76fcd298\"\ndown_revision = \"5c7fdadae813\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"chat_session\", \"persona_id\", existing_type=sa.INTEGER(), nullable=True\n )\n\n\ndef downgrade() -> None:\n # Delete chat messages and feedback first since they reference chat sessions\n # Get chat messages from sessions with null persona_id\n chat_messages_query = \"\"\"\n SELECT id\n FROM chat_message\n WHERE chat_session_id IN (\n SELECT id\n FROM chat_session\n WHERE persona_id IS NULL\n )\n \"\"\"\n\n # Delete dependent records first\n op.execute(\n f\"\"\"\n DELETE FROM document_retrieval_feedback\n WHERE chat_message_id IN (\n {chat_messages_query}\n )\n \"\"\"\n )\n op.execute(\n f\"\"\"\n DELETE FROM chat_message__search_doc\n WHERE chat_message_id IN (\n {chat_messages_query}\n )\n \"\"\"\n )\n\n # Delete chat messages\n op.execute(\n \"\"\"\n DELETE FROM chat_message\n WHERE chat_session_id IN (\n SELECT id\n FROM chat_session\n WHERE persona_id IS NULL\n )\n \"\"\"\n )\n\n # Now we can safely delete the chat sessions\n op.execute(\n \"\"\"\n DELETE FROM chat_session\n WHERE persona_id IS NULL\n \"\"\"\n )\n\n op.alter_column(\n \"chat_session\",\n \"persona_id\",\n existing_type=sa.INTEGER(),\n nullable=False,\n )\n" }, { "path": "backend/alembic/versions/c9e2cd766c29_add_s3_file_store_table.py", "content": "\"\"\"modify_file_store_for_external_storage\n\nRevision ID: c9e2cd766c29\nRevises: 03bf8be6b53a\nCreate Date: 2025-06-13 14:02:09.867679\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy import text\nfrom typing import cast, Any\n\nfrom botocore.exceptions import ClientError\n\nfrom onyx.db._deprecated.pg_file_store import delete_lobj_by_id, read_lobj\nfrom onyx.file_store.file_store import get_s3_file_store\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n# revision identifiers, used by Alembic.\nrevision = \"c9e2cd766c29\"\ndown_revision = \"03bf8be6b53a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n try:\n # Modify existing file_store table to support external storage\n op.rename_table(\"file_store\", \"file_record\")\n\n # Make lobj_oid nullable (for external storage files)\n op.alter_column(\"file_record\", \"lobj_oid\", nullable=True)\n\n # Add external storage columns with generic names\n op.add_column(\n \"file_record\", sa.Column(\"bucket_name\", sa.String(), nullable=True)\n )\n op.add_column(\n \"file_record\", sa.Column(\"object_key\", sa.String(), nullable=True)\n )\n\n # Add timestamps for tracking\n op.add_column(\n \"file_record\",\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n )\n op.add_column(\n \"file_record\",\n sa.Column(\n \"updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n )\n\n op.alter_column(\"file_record\", \"file_name\", new_column_name=\"file_id\")\n except Exception as e:\n if \"does not exist\" in str(e) or 'relation \"file_store\" does not exist' in str(\n e\n ):\n print(\n f\"Ran into error - {e}. Likely means we had a partial success in the past, continuing...\"\n )\n else:\n raise\n\n print(\n \"External storage configured - migrating files from PostgreSQL to external storage...\"\n )\n # if we fail midway through this, we'll have a partial success. Running the migration\n # again should allow us to continue.\n _migrate_files_to_external_storage()\n print(\"File migration completed successfully!\")\n\n # Remove lobj_oid column\n op.drop_column(\"file_record\", \"lobj_oid\")\n\n\ndef downgrade() -> None:\n \"\"\"Revert schema changes and migrate files from external storage back to PostgreSQL large objects.\"\"\"\n\n print(\n \"Reverting to PostgreSQL-backed file store – migrating files from external storage …\"\n )\n\n # 1. Ensure `lobj_oid` exists on the current `file_record` table (nullable for now).\n op.add_column(\"file_record\", sa.Column(\"lobj_oid\", sa.Integer(), nullable=True))\n\n # 2. Move content from external storage back into PostgreSQL large objects (table is still\n # called `file_record` so application code continues to work during the copy).\n try:\n _migrate_files_to_postgres()\n except Exception:\n print(\"Error during downgrade migration, rolling back …\")\n op.drop_column(\"file_record\", \"lobj_oid\")\n raise\n\n # 3. After migration every row should now have `lobj_oid` populated – mark NOT NULL.\n op.alter_column(\"file_record\", \"lobj_oid\", nullable=False)\n\n # 4. Remove columns that are only relevant to external storage.\n op.drop_column(\"file_record\", \"updated_at\")\n op.drop_column(\"file_record\", \"created_at\")\n op.drop_column(\"file_record\", \"object_key\")\n op.drop_column(\"file_record\", \"bucket_name\")\n\n # 5. Rename `file_id` back to `file_name` (still on `file_record`).\n op.alter_column(\"file_record\", \"file_id\", new_column_name=\"file_name\")\n\n # 6. Finally, rename the table back to its original name expected by the legacy codebase.\n op.rename_table(\"file_record\", \"file_store\")\n\n print(\n \"Downgrade migration completed – files are now stored inside PostgreSQL again.\"\n )\n\n\n# -----------------------------------------------------------------------------\n# Helper: migrate from external storage (S3/MinIO) back into PostgreSQL large objects\n\n\ndef _migrate_files_to_postgres() -> None:\n \"\"\"Move any files whose content lives in external S3-compatible storage back into PostgreSQL.\n\n The logic mirrors *inverse* of `_migrate_files_to_external_storage` used on upgrade.\n \"\"\"\n\n # Obtain DB session from Alembic context\n bind = op.get_bind()\n session = Session(bind=bind)\n\n # Fetch rows that have external storage pointers (bucket/object_key not NULL)\n result = session.execute(\n text(\n \"SELECT file_id, bucket_name, object_key FROM file_record WHERE bucket_name IS NOT NULL AND object_key IS NOT NULL\"\n )\n )\n\n files_to_migrate = [row[0] for row in result.fetchall()]\n total_files = len(files_to_migrate)\n\n if total_files == 0:\n print(\"No files found in external storage to migrate back to PostgreSQL.\")\n return\n\n print(f\"Found {total_files} files to migrate back to PostgreSQL large objects.\")\n\n _set_tenant_contextvar(session)\n migrated_count = 0\n\n # only create external store if we have files to migrate. This line\n # makes it so we need to have S3/MinIO configured to run this migration.\n external_store = get_s3_file_store()\n\n for i, file_id in enumerate(files_to_migrate, 1):\n print(f\"Migrating file {i}/{total_files}: {file_id}\")\n\n # Read file content from external storage (always binary)\n try:\n file_io = external_store.read_file(\n file_id=file_id, mode=\"b\", use_tempfile=True\n )\n file_io.seek(0)\n\n # Import lazily to avoid circular deps at Alembic runtime\n from onyx.db._deprecated.pg_file_store import (\n create_populate_lobj,\n ) # noqa: E402\n\n # Create new Postgres large object and populate it\n lobj_oid = create_populate_lobj(content=file_io, db_session=session)\n\n # Update DB row: set lobj_oid, clear bucket/object_key\n session.execute(\n text(\n \"UPDATE file_record SET lobj_oid = :lobj_oid, bucket_name = NULL, object_key = NULL WHERE file_id = :file_id\"\n ),\n {\"lobj_oid\": lobj_oid, \"file_id\": file_id},\n )\n except ClientError as e:\n if \"NoSuchKey\" in str(e):\n print(\n f\"File {file_id} not found in external storage. Deleting from database.\"\n )\n session.execute(\n text(\"DELETE FROM file_record WHERE file_id = :file_id\"),\n {\"file_id\": file_id},\n )\n else:\n raise\n\n migrated_count += 1\n print(f\"✓ Successfully migrated file {i}/{total_files}: {file_id}\")\n\n # Flush the SQLAlchemy session so statements are sent to the DB, but **do not**\n # commit the transaction. The surrounding Alembic migration will commit once\n # the *entire* downgrade succeeds. This keeps the whole downgrade atomic and\n # avoids leaving the database in a partially-migrated state if a later schema\n # operation fails.\n session.flush()\n\n print(\n f\"Migration back to PostgreSQL completed: {migrated_count} files staged for commit.\"\n )\n\n\ndef _migrate_files_to_external_storage() -> None:\n \"\"\"Migrate files from PostgreSQL large objects to external storage\"\"\"\n # Get database session\n bind = op.get_bind()\n session = Session(bind=bind)\n external_store = get_s3_file_store()\n\n # Find all files currently stored in PostgreSQL (lobj_oid is not null)\n result = session.execute(\n text(\n \"SELECT file_id FROM file_record WHERE lobj_oid IS NOT NULL AND bucket_name IS NULL AND object_key IS NULL\"\n )\n )\n\n files_to_migrate = [row[0] for row in result.fetchall()]\n total_files = len(files_to_migrate)\n\n if total_files == 0:\n print(\"No files found in PostgreSQL storage to migrate.\")\n return\n\n # might need to move this above the if statement when creating a new multi-tenant\n # system. VERY extreme edge case.\n external_store.initialize()\n print(f\"Found {total_files} files to migrate from PostgreSQL to external storage.\")\n\n _set_tenant_contextvar(session)\n migrated_count = 0\n\n for i, file_id in enumerate(files_to_migrate, 1):\n print(f\"Migrating file {i}/{total_files}: {file_id}\")\n\n # Read file record to get metadata\n file_record = session.execute(\n text(\"SELECT * FROM file_record WHERE file_id = :file_id\"),\n {\"file_id\": file_id},\n ).fetchone()\n\n if file_record is None:\n print(f\"File {file_id} not found in PostgreSQL storage.\")\n continue\n\n lobj_id = cast(int, file_record.lobj_oid)\n file_metadata = cast(Any, file_record.file_metadata)\n\n # Read file content from PostgreSQL\n try:\n file_content = read_lobj(\n lobj_id, db_session=session, mode=\"b\", use_tempfile=True\n )\n except Exception as e:\n if \"large object\" in str(e) and \"does not exist\" in str(e):\n print(f\"File {file_id} not found in PostgreSQL storage.\")\n continue\n else:\n raise\n\n # Handle file_metadata type conversion\n file_metadata = None\n if file_metadata is not None:\n if isinstance(file_metadata, dict):\n file_metadata = file_metadata\n else:\n # Convert other types to dict if possible, otherwise None\n try:\n file_metadata = dict(file_record.file_metadata)\n except (TypeError, ValueError):\n file_metadata = None\n\n # Save to external storage (this will handle the database record update and cleanup)\n # NOTE: this WILL .commit() the transaction.\n external_store.save_file(\n file_id=file_id,\n content=file_content,\n display_name=file_record.display_name,\n file_origin=file_record.file_origin,\n file_type=file_record.file_type,\n file_metadata=file_metadata,\n )\n delete_lobj_by_id(lobj_id, db_session=session)\n\n migrated_count += 1\n print(f\"✓ Successfully migrated file {i}/{total_files}: {file_id}\")\n\n # See note above – flush but do **not** commit so the outer Alembic transaction\n # controls atomicity.\n session.flush()\n\n print(\n f\"Migration completed: {migrated_count} files staged for commit to external storage.\"\n )\n\n\ndef _set_tenant_contextvar(session: Session) -> None:\n \"\"\"Set the tenant contextvar to the default schema\"\"\"\n current_tenant = session.execute(text(\"SELECT current_schema()\")).scalar()\n print(f\"Migrating files for tenant: {current_tenant}\")\n CURRENT_TENANT_ID_CONTEXTVAR.set(current_tenant)\n" }, { "path": "backend/alembic/versions/ca04500b9ee8_add_cascade_deletes_to_agent_tables.py", "content": "\"\"\"add_cascade_deletes_to_agent_tables\n\nRevision ID: ca04500b9ee8\nRevises: 238b84885828\nCreate Date: 2025-05-30 16:03:51.112263\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"ca04500b9ee8\"\ndown_revision = \"238b84885828\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Drop existing foreign key constraints\n op.drop_constraint(\n \"agent__sub_question_primary_question_id_fkey\",\n \"agent__sub_question\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"agent__sub_query_parent_question_id_fkey\",\n \"agent__sub_query\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"chat_message__standard_answer_chat_message_id_fkey\",\n \"chat_message__standard_answer\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"agent__sub_query__search_doc_sub_query_id_fkey\",\n \"agent__sub_query__search_doc\",\n type_=\"foreignkey\",\n )\n\n # Recreate foreign key constraints with CASCADE delete\n op.create_foreign_key(\n \"agent__sub_question_primary_question_id_fkey\",\n \"agent__sub_question\",\n \"chat_message\",\n [\"primary_question_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n op.create_foreign_key(\n \"agent__sub_query_parent_question_id_fkey\",\n \"agent__sub_query\",\n \"agent__sub_question\",\n [\"parent_question_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n op.create_foreign_key(\n \"chat_message__standard_answer_chat_message_id_fkey\",\n \"chat_message__standard_answer\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n op.create_foreign_key(\n \"agent__sub_query__search_doc_sub_query_id_fkey\",\n \"agent__sub_query__search_doc\",\n \"agent__sub_query\",\n [\"sub_query_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n # Drop CASCADE foreign key constraints\n op.drop_constraint(\n \"agent__sub_question_primary_question_id_fkey\",\n \"agent__sub_question\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"agent__sub_query_parent_question_id_fkey\",\n \"agent__sub_query\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"chat_message__standard_answer_chat_message_id_fkey\",\n \"chat_message__standard_answer\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"agent__sub_query__search_doc_sub_query_id_fkey\",\n \"agent__sub_query__search_doc\",\n type_=\"foreignkey\",\n )\n\n # Recreate foreign key constraints without CASCADE delete\n op.create_foreign_key(\n \"agent__sub_question_primary_question_id_fkey\",\n \"agent__sub_question\",\n \"chat_message\",\n [\"primary_question_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"agent__sub_query_parent_question_id_fkey\",\n \"agent__sub_query\",\n \"agent__sub_question\",\n [\"parent_question_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"chat_message__standard_answer_chat_message_id_fkey\",\n \"chat_message__standard_answer\",\n \"chat_message\",\n [\"chat_message_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"agent__sub_query__search_doc_sub_query_id_fkey\",\n \"agent__sub_query__search_doc\",\n \"agent__sub_query\",\n [\"sub_query_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/cbc03e08d0f3_add_opensearch_migration_tables.py", "content": "\"\"\"add_opensearch_migration_tables\n\nRevision ID: cbc03e08d0f3\nRevises: be87a654d5af\nCreate Date: 2026-01-31 17:00:45.176604\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"cbc03e08d0f3\"\ndown_revision = \"be87a654d5af\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # 1. Create opensearch_document_migration_record table.\n op.create_table(\n \"opensearch_document_migration_record\",\n sa.Column(\"document_id\", sa.String(), nullable=False),\n sa.Column(\"status\", sa.String(), nullable=False, server_default=\"pending\"),\n sa.Column(\"error_message\", sa.Text(), nullable=True),\n sa.Column(\"attempts_count\", sa.Integer(), nullable=False, server_default=\"0\"),\n sa.Column(\"last_attempt_at\", sa.DateTime(timezone=True), nullable=True),\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"document_id\"),\n sa.ForeignKeyConstraint(\n [\"document_id\"],\n [\"document.id\"],\n ondelete=\"CASCADE\",\n ),\n )\n # 2. Create indices.\n op.create_index(\n \"ix_opensearch_document_migration_record_status\",\n \"opensearch_document_migration_record\",\n [\"status\"],\n )\n op.create_index(\n \"ix_opensearch_document_migration_record_attempts_count\",\n \"opensearch_document_migration_record\",\n [\"attempts_count\"],\n )\n op.create_index(\n \"ix_opensearch_document_migration_record_created_at\",\n \"opensearch_document_migration_record\",\n [\"created_at\"],\n )\n\n # 3. Create opensearch_tenant_migration_record table (singleton).\n op.create_table(\n \"opensearch_tenant_migration_record\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"document_migration_record_table_population_status\",\n sa.String(),\n nullable=False,\n server_default=\"pending\",\n ),\n sa.Column(\n \"num_times_observed_no_additional_docs_to_populate_migration_table\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n sa.Column(\n \"overall_document_migration_status\",\n sa.String(),\n nullable=False,\n server_default=\"pending\",\n ),\n sa.Column(\n \"num_times_observed_no_additional_docs_to_migrate\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n sa.Column(\n \"last_updated_at\",\n sa.DateTime(timezone=True),\n server_default=sa.func.now(),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n # 4. Create unique index on constant to enforce singleton pattern.\n op.execute(\n sa.text(\n \"\"\"\n CREATE UNIQUE INDEX idx_opensearch_tenant_migration_singleton\n ON opensearch_tenant_migration_record ((true))\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n # Drop opensearch_tenant_migration_record.\n op.drop_index(\n \"idx_opensearch_tenant_migration_singleton\",\n table_name=\"opensearch_tenant_migration_record\",\n )\n op.drop_table(\"opensearch_tenant_migration_record\")\n\n # Drop opensearch_document_migration_record.\n op.drop_index(\n \"ix_opensearch_document_migration_record_created_at\",\n table_name=\"opensearch_document_migration_record\",\n )\n op.drop_index(\n \"ix_opensearch_document_migration_record_attempts_count\",\n table_name=\"opensearch_document_migration_record\",\n )\n op.drop_index(\n \"ix_opensearch_document_migration_record_status\",\n table_name=\"opensearch_document_migration_record\",\n )\n op.drop_table(\"opensearch_document_migration_record\")\n" }, { "path": "backend/alembic/versions/cec7ec36c505_kgentity_parent.py", "content": "\"\"\"kgentity_parent\n\nRevision ID: cec7ec36c505\nRevises: 495cb26ce93e\nCreate Date: 2025-06-07 20:07:46.400770\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"cec7ec36c505\"\ndown_revision = \"495cb26ce93e\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"kg_entity\",\n sa.Column(\"parent_key\", sa.String(), nullable=True, index=True),\n )\n # NOTE: you will have to reindex the KG after this migration as the parent_key will be null\n\n\ndef downgrade() -> None:\n op.drop_column(\"kg_entity\", \"parent_key\")\n" }, { "path": "backend/alembic/versions/cf90764725d8_larger_refresh_tokens.py", "content": "\"\"\"larger refresh tokens\n\nRevision ID: cf90764725d8\nRevises: 4794bc13e484\nCreate Date: 2025-04-04 10:56:39.769294\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"cf90764725d8\"\ndown_revision = \"4794bc13e484\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.alter_column(\"oauth_account\", \"refresh_token\", type_=sa.Text())\n\n\ndef downgrade() -> None:\n op.alter_column(\"oauth_account\", \"refresh_token\", type_=sa.String(length=1024))\n" }, { "path": "backend/alembic/versions/d09fc20a3c66_seed_builtin_tools.py", "content": "\"\"\"seed_builtin_tools\n\nRevision ID: d09fc20a3c66\nRevises: b7ec9b5b505f\nCreate Date: 2025-09-09 19:32:16.824373\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d09fc20a3c66\"\ndown_revision = \"b7ec9b5b505f\"\nbranch_labels = None\ndepends_on = None\n\n\n# Tool definitions - core tools that should always be seeded\n# Names/in_code_tool_id are the same as the class names in the tool_implementations package\nBUILT_IN_TOOLS = [\n {\n \"name\": \"SearchTool\",\n \"display_name\": \"Internal Search\",\n \"description\": \"The Search Action allows the Assistant to search through connected knowledge to help build an answer.\",\n \"in_code_tool_id\": \"SearchTool\",\n },\n {\n \"name\": \"ImageGenerationTool\",\n \"display_name\": \"Image Generation\",\n \"description\": (\n \"The Image Generation Action allows the assistant to use DALL-E 3 or GPT-IMAGE-1 to generate images. \"\n \"The action will be used when the user asks the assistant to generate an image.\"\n ),\n \"in_code_tool_id\": \"ImageGenerationTool\",\n },\n {\n \"name\": \"WebSearchTool\",\n \"display_name\": \"Web Search\",\n \"description\": (\n \"The Web Search Action allows the assistant to perform internet searches for up-to-date information.\"\n ),\n \"in_code_tool_id\": \"WebSearchTool\",\n },\n {\n \"name\": \"KnowledgeGraphTool\",\n \"display_name\": \"Knowledge Graph Search\",\n \"description\": (\n \"The Knowledge Graph Search Action allows the assistant to search the \"\n \"Knowledge Graph for information. This tool can (for now) only be active in the KG Beta Assistant, \"\n \"and it requires the Knowledge Graph to be enabled.\"\n ),\n \"in_code_tool_id\": \"KnowledgeGraphTool\",\n },\n {\n \"name\": \"OktaProfileTool\",\n \"display_name\": \"Okta Profile\",\n \"description\": (\n \"The Okta Profile Action allows the assistant to fetch the current user's information from Okta. \"\n \"This may include the user's name, email, phone number, address, and other details such as their \"\n \"manager and direct reports.\"\n ),\n \"in_code_tool_id\": \"OktaProfileTool\",\n },\n]\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Get existing tools to check what already exists\n existing_tools = conn.execute(\n sa.text(\"SELECT in_code_tool_id FROM tool WHERE in_code_tool_id IS NOT NULL\")\n ).fetchall()\n existing_tool_ids = {row[0] for row in existing_tools}\n\n # Insert or update built-in tools\n for tool in BUILT_IN_TOOLS:\n in_code_id = tool[\"in_code_tool_id\"]\n\n # Handle historical rename: InternetSearchTool -> WebSearchTool\n if (\n in_code_id == \"WebSearchTool\"\n and \"WebSearchTool\" not in existing_tool_ids\n and \"InternetSearchTool\" in existing_tool_ids\n ):\n # Rename the existing InternetSearchTool row in place and update fields\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :name,\n display_name = :display_name,\n description = :description,\n in_code_tool_id = :in_code_tool_id\n WHERE in_code_tool_id = 'InternetSearchTool'\n \"\"\"\n ),\n tool,\n )\n # Keep the local view of existing ids in sync to avoid duplicate insert\n existing_tool_ids.discard(\"InternetSearchTool\")\n existing_tool_ids.add(\"WebSearchTool\")\n continue\n\n if in_code_id in existing_tool_ids:\n # Update existing tool\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :name,\n display_name = :display_name,\n description = :description\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n tool,\n )\n else:\n # Insert new tool\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO tool (name, display_name, description, in_code_tool_id)\n VALUES (:name, :display_name, :description, :in_code_tool_id)\n \"\"\"\n ),\n tool,\n )\n\n\ndef downgrade() -> None:\n # We don't remove the tools on downgrade since it's totally fine to just\n # have them around. If we upgrade again, it will be a no-op.\n pass\n" }, { "path": "backend/alembic/versions/d1b637d7050a_sync_exa_api_key_to_content_provider.py", "content": "\"\"\"sync_exa_api_key_to_content_provider\n\nRevision ID: d1b637d7050a\nRevises: d25168c2beee\nCreate Date: 2026-01-09 15:54:15.646249\n\n\"\"\"\n\nfrom alembic import op\nfrom sqlalchemy import text\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d1b637d7050a\"\ndown_revision = \"d25168c2beee\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Exa uses a shared API key between search and content providers.\n # For existing Exa search providers with API keys, create the corresponding\n # content provider if it doesn't exist yet.\n connection = op.get_bind()\n\n # Check if Exa search provider exists with an API key\n result = connection.execute(\n text(\n \"\"\"\n SELECT api_key FROM internet_search_provider\n WHERE provider_type = 'exa' AND api_key IS NOT NULL\n LIMIT 1\n \"\"\"\n )\n )\n row = result.fetchone()\n\n if row:\n api_key = row[0]\n # Create Exa content provider with the shared key\n connection.execute(\n text(\n \"\"\"\n INSERT INTO internet_content_provider\n (name, provider_type, api_key, is_active)\n VALUES ('Exa', 'exa', :api_key, false)\n ON CONFLICT (name) DO NOTHING\n \"\"\"\n ),\n {\"api_key\": api_key},\n )\n\n\ndef downgrade() -> None:\n # Remove the Exa content provider that was created by this migration\n connection = op.get_bind()\n connection.execute(\n text(\n \"\"\"\n DELETE FROM internet_content_provider\n WHERE provider_type = 'exa'\n \"\"\"\n )\n )\n" }, { "path": "backend/alembic/versions/d25168c2beee_tool_name_consistency.py", "content": "\"\"\"tool_name_consistency\n\nRevision ID: d25168c2beee\nRevises: 8405ca81cc83\nCreate Date: 2026-01-11 17:54:40.135777\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d25168c2beee\"\ndown_revision = \"8405ca81cc83\"\nbranch_labels = None\ndepends_on = None\n\n\n# Currently the seeded tools have the in_code_tool_id == name\nCURRENT_TOOL_NAME_MAPPING = [\n \"SearchTool\",\n \"WebSearchTool\",\n \"ImageGenerationTool\",\n \"PythonTool\",\n \"OpenURLTool\",\n \"KnowledgeGraphTool\",\n \"ResearchAgent\",\n]\n\n# Mapping of in_code_tool_id -> name\n# These are the expected names that we want in the database\nEXPECTED_TOOL_NAME_MAPPING = {\n \"SearchTool\": \"internal_search\",\n \"WebSearchTool\": \"web_search\",\n \"ImageGenerationTool\": \"generate_image\",\n \"PythonTool\": \"python\",\n \"OpenURLTool\": \"open_url\",\n \"KnowledgeGraphTool\": \"run_kg_search\",\n \"ResearchAgent\": \"research_agent\",\n}\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Mapping of in_code_tool_id to the NAME constant from each tool class\n # These match the .name property of each tool implementation\n tool_name_mapping = EXPECTED_TOOL_NAME_MAPPING\n\n # Update the name column for each tool based on its in_code_tool_id\n for in_code_tool_id, expected_name in tool_name_mapping.items():\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :expected_name\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n {\n \"expected_name\": expected_name,\n \"in_code_tool_id\": in_code_tool_id,\n },\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n\n # Reverse the migration by setting name back to in_code_tool_id\n # This matches the original pattern where name was the class name\n for in_code_tool_id in CURRENT_TOOL_NAME_MAPPING:\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :current_name\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n {\n \"current_name\": in_code_tool_id,\n \"in_code_tool_id\": in_code_tool_id,\n },\n )\n" }, { "path": "backend/alembic/versions/d3fd499c829c_add_file_reader_tool.py", "content": "\"\"\"add_file_reader_tool\n\nRevision ID: d3fd499c829c\nRevises: 114a638452db\nCreate Date: 2026-02-07 19:28:22.452337\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d3fd499c829c\"\ndown_revision = \"114a638452db\"\nbranch_labels = None\ndepends_on = None\n\nFILE_READER_TOOL = {\n \"name\": \"read_file\",\n \"display_name\": \"File Reader\",\n \"description\": (\n \"Read sections of user-uploaded files by character offset. \"\n \"Useful for inspecting large files that cannot fit entirely in context.\"\n ),\n \"in_code_tool_id\": \"FileReaderTool\",\n \"enabled\": True,\n}\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n\n # Check if tool already exists\n existing = conn.execute(\n sa.text(\"SELECT id FROM tool WHERE in_code_tool_id = :in_code_tool_id\"),\n {\"in_code_tool_id\": FILE_READER_TOOL[\"in_code_tool_id\"]},\n ).fetchone()\n\n if existing:\n # Update existing tool\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE tool\n SET name = :name,\n display_name = :display_name,\n description = :description\n WHERE in_code_tool_id = :in_code_tool_id\n \"\"\"\n ),\n FILE_READER_TOOL,\n )\n tool_id = existing[0]\n else:\n # Insert new tool\n result = conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO tool (name, display_name, description, in_code_tool_id, enabled)\n VALUES (:name, :display_name, :description, :in_code_tool_id, :enabled)\n RETURNING id\n \"\"\"\n ),\n FILE_READER_TOOL,\n )\n tool_id = result.scalar_one()\n\n # Attach to the default persona (id=0) if not already attached\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO persona__tool (persona_id, tool_id)\n VALUES (0, :tool_id)\n ON CONFLICT DO NOTHING\n \"\"\"\n ),\n {\"tool_id\": tool_id},\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n in_code_tool_id = FILE_READER_TOOL[\"in_code_tool_id\"]\n\n # Remove persona associations first (FK constraint)\n conn.execute(\n sa.text(\n \"\"\"\n DELETE FROM persona__tool\n WHERE tool_id IN (\n SELECT id FROM tool WHERE in_code_tool_id = :in_code_tool_id\n )\n \"\"\"\n ),\n {\"in_code_tool_id\": in_code_tool_id},\n )\n\n conn.execute(\n sa.text(\"DELETE FROM tool WHERE in_code_tool_id = :in_code_tool_id\"),\n {\"in_code_tool_id\": in_code_tool_id},\n )\n" }, { "path": "backend/alembic/versions/d5645c915d0e_remove_deletion_attempt_table.py", "content": "\"\"\"Remove deletion_attempt table\n\nRevision ID: d5645c915d0e\nRevises: 8e26726b7683\nCreate Date: 2023-09-14 15:04:14.444909\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"d5645c915d0e\"\ndown_revision = \"8e26726b7683\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_table(\"deletion_attempt\")\n\n # Remove the DeletionStatus enum\n op.execute(\"DROP TYPE IF EXISTS deletionstatus;\")\n\n\ndef downgrade() -> None:\n op.create_table(\n \"deletion_attempt\",\n sa.Column(\"id\", sa.INTEGER(), autoincrement=True, nullable=False),\n sa.Column(\"connector_id\", sa.INTEGER(), autoincrement=False, nullable=False),\n sa.Column(\"credential_id\", sa.INTEGER(), autoincrement=False, nullable=False),\n sa.Column(\n \"status\",\n postgresql.ENUM(\n \"NOT_STARTED\",\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n name=\"deletionstatus\",\n ),\n autoincrement=False,\n nullable=False,\n ),\n sa.Column(\n \"num_docs_deleted\",\n sa.INTEGER(),\n autoincrement=False,\n nullable=False,\n ),\n sa.Column(\"error_msg\", sa.VARCHAR(), autoincrement=False, nullable=True),\n sa.Column(\n \"time_created\",\n postgresql.TIMESTAMP(timezone=True),\n server_default=sa.text(\"now()\"),\n autoincrement=False,\n nullable=False,\n ),\n sa.Column(\n \"time_updated\",\n postgresql.TIMESTAMP(timezone=True),\n server_default=sa.text(\"now()\"),\n autoincrement=False,\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"connector_id\"],\n [\"connector.id\"],\n name=\"deletion_attempt_connector_id_fkey\",\n ),\n sa.ForeignKeyConstraint(\n [\"credential_id\"],\n [\"credential.id\"],\n name=\"deletion_attempt_credential_id_fkey\",\n ),\n sa.PrimaryKeyConstraint(\"id\", name=\"deletion_attempt_pkey\"),\n )\n" }, { "path": "backend/alembic/versions/d56ffa94ca32_add_file_content.py", "content": "\"\"\"add_file_content\n\nRevision ID: d56ffa94ca32\nRevises: 01f8e6d95a33\nCreate Date: 2026-02-06 15:29:34.192960\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d56ffa94ca32\"\ndown_revision = \"01f8e6d95a33\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"file_content\",\n sa.Column(\n \"file_id\",\n sa.String(),\n sa.ForeignKey(\"file_record.file_id\", ondelete=\"CASCADE\"),\n primary_key=True,\n ),\n sa.Column(\"lobj_oid\", sa.BigInteger(), nullable=False),\n sa.Column(\"file_size\", sa.BigInteger(), nullable=False, server_default=\"0\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"file_content\")\n" }, { "path": "backend/alembic/versions/d5c86e2c6dc6_add_cascade_delete_to_search_query_user_.py", "content": "\"\"\"add_cascade_delete_to_search_query_user_id\n\nRevision ID: d5c86e2c6dc6\nRevises: 90b409d06e50\nCreate Date: 2026-02-04 16:05:04.749804\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d5c86e2c6dc6\"\ndown_revision = \"90b409d06e50\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.drop_constraint(\"search_query_user_id_fkey\", \"search_query\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"search_query_user_id_fkey\",\n \"search_query\",\n \"user\",\n [\"user_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\"search_query_user_id_fkey\", \"search_query\", type_=\"foreignkey\")\n op.create_foreign_key(\n \"search_query_user_id_fkey\", \"search_query\", \"user\", [\"user_id\"], [\"id\"]\n )\n" }, { "path": "backend/alembic/versions/d61e513bef0a_add_total_docs_for_index_attempt.py", "content": "\"\"\"Add Total Docs for Index Attempt\n\nRevision ID: d61e513bef0a\nRevises: 46625e4745d4\nCreate Date: 2023-10-27 23:02:43.369964\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"d61e513bef0a\"\ndown_revision = \"46625e4745d4\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"index_attempt\",\n sa.Column(\"new_docs_indexed\", sa.Integer(), nullable=True),\n )\n op.alter_column(\n \"index_attempt\", \"num_docs_indexed\", new_column_name=\"total_docs_indexed\"\n )\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"index_attempt\", \"total_docs_indexed\", new_column_name=\"num_docs_indexed\"\n )\n op.drop_column(\"index_attempt\", \"new_docs_indexed\")\n" }, { "path": "backend/alembic/versions/d7111c1238cd_remove_document_ids.py", "content": "\"\"\"Remove Document IDs\n\nRevision ID: d7111c1238cd\nRevises: 465f78d9b7f9\nCreate Date: 2023-07-29 15:06:25.126169\n\n\"\"\"\n\nimport sqlalchemy as sa\nfrom alembic import op\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"d7111c1238cd\"\ndown_revision = \"465f78d9b7f9\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"index_attempt\", \"document_ids\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"index_attempt\",\n sa.Column(\n \"document_ids\",\n postgresql.ARRAY(sa.VARCHAR()),\n autoincrement=False,\n nullable=True,\n ),\n )\n" }, { "path": "backend/alembic/versions/d716b0791ddd_combined_slack_id_fields.py", "content": "\"\"\"combined slack id fields\n\nRevision ID: d716b0791ddd\nRevises: 7aea705850d5\nCreate Date: 2024-07-10 17:57:45.630550\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"d716b0791ddd\"\ndown_revision = \"7aea705850d5\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE slack_bot_config\n SET channel_config = jsonb_set(\n channel_config,\n '{respond_member_group_list}',\n coalesce(channel_config->'respond_team_member_list', '[]'::jsonb) ||\n coalesce(channel_config->'respond_slack_group_list', '[]'::jsonb)\n ) - 'respond_team_member_list' - 'respond_slack_group_list'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE slack_bot_config\n SET channel_config = jsonb_set(\n jsonb_set(\n channel_config - 'respond_member_group_list',\n '{respond_team_member_list}',\n '[]'::jsonb\n ),\n '{respond_slack_group_list}',\n '[]'::jsonb\n )\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/d8cdfee5df80_add_skipped_to_userfilestatus.py", "content": "\"\"\"add skipped to userfilestatus\n\nRevision ID: d8cdfee5df80\nRevises: 1d78c0ca7853\nCreate Date: 2026-04-01 10:47:12.593950\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d8cdfee5df80\"\ndown_revision = \"1d78c0ca7853\"\nbranch_labels = None\ndepends_on = None\n\n\nTABLE = \"user_file\"\nCOLUMN = \"status\"\nCONSTRAINT_NAME = \"ck_user_file_status\"\n\nOLD_VALUES = (\"PROCESSING\", \"INDEXING\", \"COMPLETED\", \"FAILED\", \"CANCELED\", \"DELETING\")\nNEW_VALUES = (\n \"PROCESSING\",\n \"INDEXING\",\n \"COMPLETED\",\n \"SKIPPED\",\n \"FAILED\",\n \"CANCELED\",\n \"DELETING\",\n)\n\n\ndef _drop_status_check_constraint() -> None:\n inspector = sa.inspect(op.get_bind())\n for constraint in inspector.get_check_constraints(TABLE):\n if COLUMN in constraint.get(\"sqltext\", \"\"):\n constraint_name = constraint[\"name\"]\n if constraint_name is not None:\n op.drop_constraint(constraint_name, TABLE, type_=\"check\")\n\n\ndef upgrade() -> None:\n _drop_status_check_constraint()\n in_clause = \", \".join(f\"'{v}'\" for v in NEW_VALUES)\n op.create_check_constraint(CONSTRAINT_NAME, TABLE, f\"{COLUMN} IN ({in_clause})\")\n\n\ndef downgrade() -> None:\n op.execute(f\"UPDATE {TABLE} SET {COLUMN} = 'COMPLETED' WHERE {COLUMN} = 'SKIPPED'\")\n _drop_status_check_constraint()\n in_clause = \", \".join(f\"'{v}'\" for v in OLD_VALUES)\n op.create_check_constraint(CONSTRAINT_NAME, TABLE, f\"{COLUMN} IN ({in_clause})\")\n" }, { "path": "backend/alembic/versions/d929f0c1c6af_feedback_feature.py", "content": "\"\"\"Feedback Feature\n\nRevision ID: d929f0c1c6af\nRevises: 8aabb57f3b49\nCreate Date: 2023-08-27 13:03:54.274987\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d929f0c1c6af\"\ndown_revision = \"8aabb57f3b49\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"query_event\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"query\", sa.String(), nullable=False),\n sa.Column(\n \"selected_search_flow\",\n sa.Enum(\"KEYWORD\", \"SEMANTIC\", name=\"searchtype\", native_enum=False),\n nullable=True,\n ),\n sa.Column(\"llm_answer\", sa.String(), nullable=True),\n sa.Column(\n \"feedback\",\n sa.Enum(\"LIKE\", \"DISLIKE\", name=\"qafeedbacktype\", native_enum=False),\n nullable=True,\n ),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"document_retrieval_feedback\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"qa_event_id\", sa.Integer(), nullable=False),\n sa.Column(\"document_id\", sa.String(), nullable=False),\n sa.Column(\"document_rank\", sa.Integer(), nullable=False),\n sa.Column(\"clicked\", sa.Boolean(), nullable=False),\n sa.Column(\n \"feedback\",\n sa.Enum(\n \"ENDORSE\",\n \"REJECT\",\n \"HIDE\",\n \"UNHIDE\",\n name=\"searchfeedbacktype\",\n native_enum=False,\n ),\n nullable=True,\n ),\n sa.ForeignKeyConstraint(\n [\"document_id\"],\n [\"document.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"qa_event_id\"],\n [\"query_event.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.add_column(\"document\", sa.Column(\"boost\", sa.Integer(), nullable=False))\n op.add_column(\"document\", sa.Column(\"hidden\", sa.Boolean(), nullable=False))\n op.add_column(\"document\", sa.Column(\"semantic_id\", sa.String(), nullable=False))\n op.add_column(\"document\", sa.Column(\"link\", sa.String(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"document\", \"link\")\n op.drop_column(\"document\", \"semantic_id\")\n op.drop_column(\"document\", \"hidden\")\n op.drop_column(\"document\", \"boost\")\n op.drop_table(\"document_retrieval_feedback\")\n op.drop_table(\"query_event\")\n" }, { "path": "backend/alembic/versions/d961aca62eb3_update_status_length.py", "content": "\"\"\"Update status length\n\nRevision ID: d961aca62eb3\nRevises: cf90764725d8\nCreate Date: 2025-03-23 16:10:05.683965\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d961aca62eb3\"\ndown_revision = \"cf90764725d8\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Drop the existing enum type constraint\n op.execute(\"ALTER TABLE connector_credential_pair ALTER COLUMN status TYPE varchar\")\n\n # Create new enum type with all values\n op.execute(\n \"ALTER TABLE connector_credential_pair ALTER COLUMN status TYPE VARCHAR(20) USING status::varchar(20)\"\n )\n\n # Update the enum type to include all possible values\n op.alter_column(\n \"connector_credential_pair\",\n \"status\",\n type_=sa.Enum(\n \"SCHEDULED\",\n \"INITIAL_INDEXING\",\n \"ACTIVE\",\n \"PAUSED\",\n \"DELETING\",\n \"INVALID\",\n name=\"connectorcredentialpairstatus\",\n native_enum=False,\n ),\n existing_type=sa.String(20),\n nullable=False,\n )\n\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"in_repeated_error_state\", sa.Boolean, default=False, server_default=\"false\"\n ),\n )\n\n\ndef downgrade() -> None:\n # no need to convert back to the old enum type, since we're not using it anymore\n op.drop_column(\"connector_credential_pair\", \"in_repeated_error_state\")\n" }, { "path": "backend/alembic/versions/d9ec13955951_remove__dim_suffix_from_model_name.py", "content": "\"\"\"Remove _alt suffix from model_name\n\nRevision ID: d9ec13955951\nRevises: da4c21c69164\nCreate Date: 2024-08-20 16:31:32.955686\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"d9ec13955951\"\ndown_revision = \"da4c21c69164\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE embedding_model\n SET model_name = regexp_replace(model_name, '__danswer_alt_index$', '')\n WHERE model_name LIKE '%__danswer_alt_index'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n # We can't reliably add the __danswer_alt_index suffix back, so we'll leave this empty\n pass\n" }, { "path": "backend/alembic/versions/da42808081e3_migrate_jira_connectors_to_new_format.py", "content": "\"\"\"migrate jira connectors to new format\n\nRevision ID: da42808081e3\nRevises: f13db29f3101\nCreate Date: 2025-02-24 11:24:54.396040\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nimport json\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.jira.utils import extract_jira_project\n\n\n# revision identifiers, used by Alembic.\nrevision = \"da42808081e3\"\ndown_revision = \"f13db29f3101\"\nbranch_labels = None\ndepends_on = None\n\n\nPRESERVED_CONFIG_KEYS = [\"comment_email_blacklist\", \"batch_size\", \"labels_to_skip\"]\n\n\ndef upgrade() -> None:\n # Get all Jira connectors\n conn = op.get_bind()\n\n # First get all Jira connectors\n jira_connectors = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = :source\n \"\"\"\n ),\n {\"source\": DocumentSource.JIRA.value.upper()},\n ).fetchall()\n\n # Update each connector's config\n for connector_id, old_config in jira_connectors:\n if not old_config:\n continue\n\n # Extract project key from URL if it exists\n new_config: dict[str, str | None] = {}\n if project_url := old_config.get(\"jira_project_url\"):\n # Parse the URL to get base and project\n try:\n jira_base, project_key = extract_jira_project(project_url)\n new_config = {\"jira_base_url\": jira_base, \"project_key\": project_key}\n except ValueError:\n # If URL parsing fails, just use the URL as the base\n new_config = {\n \"jira_base_url\": project_url.split(\"/projects/\")[0],\n \"project_key\": None,\n }\n else:\n # For connectors without a project URL, we need admin intervention\n # Mark these for review\n print(\n f\"WARNING: Jira connector {connector_id} has no project URL configured\"\n )\n continue\n for old_key in PRESERVED_CONFIG_KEYS:\n if old_key in old_config:\n new_config[old_key] = old_config[old_key]\n\n # Update the connector config\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = :new_config\n WHERE id = :id\n \"\"\"\n ),\n {\"id\": connector_id, \"new_config\": json.dumps(new_config)},\n )\n\n\ndef downgrade() -> None:\n # Get all Jira connectors\n conn = op.get_bind()\n\n # First get all Jira connectors\n jira_connectors = conn.execute(\n sa.text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = :source\n \"\"\"\n ),\n {\"source\": DocumentSource.JIRA.value.upper()},\n ).fetchall()\n\n # Update each connector's config back to the old format\n for connector_id, new_config in jira_connectors:\n if not new_config:\n continue\n\n old_config = {}\n base_url = new_config.get(\"jira_base_url\")\n project_key = new_config.get(\"project_key\")\n\n if base_url and project_key:\n old_config = {\"jira_project_url\": f\"{base_url}/projects/{project_key}\"}\n elif base_url:\n old_config = {\"jira_project_url\": base_url}\n else:\n continue\n\n for old_key in PRESERVED_CONFIG_KEYS:\n if old_key in new_config:\n old_config[old_key] = new_config[old_key]\n\n # Update the connector config\n conn.execute(\n sa.text(\n \"\"\"\n UPDATE connector\n SET connector_specific_config = :old_config\n WHERE id = :id\n \"\"\"\n ),\n {\"id\": connector_id, \"old_config\": json.dumps(old_config)},\n )\n" }, { "path": "backend/alembic/versions/da4c21c69164_chosen_assistants_changed_to_jsonb.py", "content": "\"\"\"chosen_assistants changed to jsonb\n\nRevision ID: da4c21c69164\nRevises: c5b692fa265c\nCreate Date: 2024-08-18 19:06:47.291491\n\n\"\"\"\n\nimport json\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"da4c21c69164\"\ndown_revision = \"c5b692fa265c\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n conn = op.get_bind()\n existing_ids_and_chosen_assistants = conn.execute(\n sa.text('select id, chosen_assistants from \"user\"')\n )\n op.drop_column(\n \"user\",\n \"chosen_assistants\",\n )\n op.add_column(\n \"user\",\n sa.Column(\n \"chosen_assistants\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n for id, chosen_assistants in existing_ids_and_chosen_assistants:\n conn.execute(\n sa.text(\n 'update \"user\" set chosen_assistants = :chosen_assistants where id = :id'\n ),\n {\"chosen_assistants\": json.dumps(chosen_assistants), \"id\": id},\n )\n\n\ndef downgrade() -> None:\n conn = op.get_bind()\n existing_ids_and_chosen_assistants = conn.execute(\n sa.text('select id, chosen_assistants from \"user\"')\n )\n op.drop_column(\n \"user\",\n \"chosen_assistants\",\n )\n op.add_column(\n \"user\",\n sa.Column(\"chosen_assistants\", postgresql.ARRAY(sa.Integer()), nullable=True),\n )\n for id, chosen_assistants in existing_ids_and_chosen_assistants:\n conn.execute(\n sa.text(\n 'update \"user\" set chosen_assistants = :chosen_assistants where id = :id'\n ),\n {\"chosen_assistants\": chosen_assistants, \"id\": id},\n )\n" }, { "path": "backend/alembic/versions/dab04867cd88_add_composite_index_to_document_by_.py", "content": "\"\"\"Add composite index to document_by_connector_credential_pair\n\nRevision ID: dab04867cd88\nRevises: 54a74a0417fc\nCreate Date: 2024-12-13 22:43:20.119990\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"dab04867cd88\"\ndown_revision = \"54a74a0417fc\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Composite index on (connector_id, credential_id)\n op.create_index(\n \"idx_document_cc_pair_connector_credential\",\n \"document_by_connector_credential_pair\",\n [\"connector_id\", \"credential_id\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n \"idx_document_cc_pair_connector_credential\",\n table_name=\"document_by_connector_credential_pair\",\n )\n" }, { "path": "backend/alembic/versions/dba7f71618f5_onyx_custom_tool_flow.py", "content": "\"\"\"Onyx Custom Tool Flow\n\nRevision ID: dba7f71618f5\nRevises: d5645c915d0e\nCreate Date: 2023-09-18 15:18:37.370972\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"dba7f71618f5\"\ndown_revision = \"d5645c915d0e\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"persona\",\n sa.Column(\"retrieval_enabled\", sa.Boolean(), nullable=True),\n )\n op.execute(\"UPDATE persona SET retrieval_enabled = true\")\n op.alter_column(\"persona\", \"retrieval_enabled\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"retrieval_enabled\")\n" }, { "path": "backend/alembic/versions/dbaa756c2ccf_embedding_models.py", "content": "\"\"\"Embedding Models\n\nRevision ID: dbaa756c2ccf\nRevises: 7f726bad5367\nCreate Date: 2024-01-25 17:12:31.813160\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy import table, column, String, Integer, Boolean\n\nfrom onyx.configs.model_configs import ASYM_PASSAGE_PREFIX\nfrom onyx.configs.model_configs import ASYM_QUERY_PREFIX\nfrom onyx.configs.model_configs import DOC_EMBEDDING_DIM\nfrom onyx.configs.model_configs import DOCUMENT_ENCODER_MODEL\nfrom onyx.configs.model_configs import NORMALIZE_EMBEDDINGS\nfrom onyx.configs.model_configs import OLD_DEFAULT_DOCUMENT_ENCODER_MODEL\nfrom onyx.configs.model_configs import OLD_DEFAULT_MODEL_DOC_EMBEDDING_DIM\nfrom onyx.configs.model_configs import OLD_DEFAULT_MODEL_NORMALIZE_EMBEDDINGS\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.db.models import IndexModelStatus\nfrom onyx.db.search_settings import user_has_overridden_embedding_model\nfrom onyx.indexing.models import IndexingSetting\nfrom onyx.natural_language_processing.search_nlp_models import clean_model_name\n\n# revision identifiers, used by Alembic.\nrevision = \"dbaa756c2ccf\"\ndown_revision = \"7f726bad5367\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef _get_old_default_embedding_model() -> IndexingSetting:\n is_overridden = user_has_overridden_embedding_model()\n return IndexingSetting(\n model_name=(\n DOCUMENT_ENCODER_MODEL\n if is_overridden\n else OLD_DEFAULT_DOCUMENT_ENCODER_MODEL\n ),\n model_dim=(\n DOC_EMBEDDING_DIM if is_overridden else OLD_DEFAULT_MODEL_DOC_EMBEDDING_DIM\n ),\n embedding_precision=(EmbeddingPrecision.FLOAT),\n normalize=(\n NORMALIZE_EMBEDDINGS\n if is_overridden\n else OLD_DEFAULT_MODEL_NORMALIZE_EMBEDDINGS\n ),\n query_prefix=(ASYM_QUERY_PREFIX if is_overridden else \"\"),\n passage_prefix=(ASYM_PASSAGE_PREFIX if is_overridden else \"\"),\n index_name=\"danswer_chunk\",\n multipass_indexing=False,\n enable_contextual_rag=False,\n api_url=None,\n )\n\n\ndef _get_new_default_embedding_model() -> IndexingSetting:\n return IndexingSetting(\n model_name=DOCUMENT_ENCODER_MODEL,\n model_dim=DOC_EMBEDDING_DIM,\n embedding_precision=(EmbeddingPrecision.BFLOAT16),\n normalize=NORMALIZE_EMBEDDINGS,\n query_prefix=ASYM_QUERY_PREFIX,\n passage_prefix=ASYM_PASSAGE_PREFIX,\n index_name=f\"danswer_chunk_{clean_model_name(DOCUMENT_ENCODER_MODEL)}\",\n multipass_indexing=False,\n enable_contextual_rag=False,\n api_url=None,\n )\n\n\ndef upgrade() -> None:\n op.create_table(\n \"embedding_model\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"model_name\", sa.String(), nullable=False),\n sa.Column(\"model_dim\", sa.Integer(), nullable=False),\n sa.Column(\"normalize\", sa.Boolean(), nullable=False),\n sa.Column(\"query_prefix\", sa.String(), nullable=False),\n sa.Column(\"passage_prefix\", sa.String(), nullable=False),\n sa.Column(\"index_name\", sa.String(), nullable=False),\n sa.Column(\n \"status\",\n sa.Enum(IndexModelStatus, native=False),\n nullable=False,\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n # since all index attempts must be associated with an embedding model,\n # need to put something in here to avoid nulls. On server startup,\n # this value will be overriden\n EmbeddingModel = table(\n \"embedding_model\",\n column(\"id\", Integer),\n column(\"model_name\", String),\n column(\"model_dim\", Integer),\n column(\"normalize\", Boolean),\n column(\"query_prefix\", String),\n column(\"passage_prefix\", String),\n column(\"index_name\", String),\n column(\n \"status\", sa.Enum(IndexModelStatus, name=\"indexmodelstatus\", native=False)\n ),\n )\n # insert an embedding model row that corresponds to the embedding model\n # the user selected via env variables before this change. This is needed since\n # all index_attempts must be associated with an embedding model, so without this\n # we will run into violations of non-null contraints\n old_embedding_model = _get_old_default_embedding_model()\n op.bulk_insert(\n EmbeddingModel,\n [\n {\n \"model_name\": old_embedding_model.model_name,\n \"model_dim\": old_embedding_model.model_dim,\n \"normalize\": old_embedding_model.normalize,\n \"query_prefix\": old_embedding_model.query_prefix,\n \"passage_prefix\": old_embedding_model.passage_prefix,\n \"index_name\": old_embedding_model.index_name,\n \"status\": IndexModelStatus.PRESENT,\n }\n ],\n )\n # if the user has not overridden the default embedding model via env variables,\n # insert the new default model into the database to auto-upgrade them\n if not user_has_overridden_embedding_model():\n new_embedding_model = _get_new_default_embedding_model()\n op.bulk_insert(\n EmbeddingModel,\n [\n {\n \"model_name\": new_embedding_model.model_name,\n \"model_dim\": new_embedding_model.model_dim,\n \"normalize\": new_embedding_model.normalize,\n \"query_prefix\": new_embedding_model.query_prefix,\n \"passage_prefix\": new_embedding_model.passage_prefix,\n \"index_name\": new_embedding_model.index_name,\n \"status\": IndexModelStatus.FUTURE,\n }\n ],\n )\n\n op.add_column(\n \"index_attempt\",\n sa.Column(\"embedding_model_id\", sa.Integer(), nullable=True),\n )\n op.execute(\n \"UPDATE index_attempt SET embedding_model_id=1 WHERE embedding_model_id IS NULL\"\n )\n op.alter_column(\n \"index_attempt\",\n \"embedding_model_id\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n op.create_foreign_key(\n \"index_attempt__embedding_model_fk\",\n \"index_attempt\",\n \"embedding_model\",\n [\"embedding_model_id\"],\n [\"id\"],\n )\n op.create_index(\n \"ix_embedding_model_present_unique\",\n \"embedding_model\",\n [\"status\"],\n unique=True,\n postgresql_where=sa.text(\"status = 'PRESENT'\"),\n )\n op.create_index(\n \"ix_embedding_model_future_unique\",\n \"embedding_model\",\n [\"status\"],\n unique=True,\n postgresql_where=sa.text(\"status = 'FUTURE'\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\n \"index_attempt__embedding_model_fk\", \"index_attempt\", type_=\"foreignkey\"\n )\n op.drop_column(\"index_attempt\", \"embedding_model_id\")\n op.drop_table(\"embedding_model\")\n op.execute(\"DROP TYPE IF EXISTS indexmodelstatus;\")\n" }, { "path": "backend/alembic/versions/df0c7ad8a076_added_deletion_attempt_table.py", "content": "\"\"\"Added deletion_attempt table\n\nRevision ID: df0c7ad8a076\nRevises: d7111c1238cd\nCreate Date: 2023-08-05 13:35:39.609619\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"df0c7ad8a076\"\ndown_revision = \"d7111c1238cd\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.execute(\"DROP TABLE IF EXISTS document CASCADE\")\n op.create_table(\n \"document\",\n sa.Column(\"id\", sa.String(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.execute(\"DROP TABLE IF EXISTS chunk CASCADE\")\n op.create_table(\n \"chunk\",\n sa.Column(\"id\", sa.String(), nullable=False),\n sa.Column(\n \"document_store_type\",\n sa.Enum(\n \"VECTOR\",\n \"KEYWORD\",\n name=\"documentstoretype\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\"document_id\", sa.String(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"document_id\"],\n [\"document.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\", \"document_store_type\"),\n )\n op.execute(\"DROP TABLE IF EXISTS deletion_attempt CASCADE\")\n op.create_table(\n \"deletion_attempt\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"connector_id\", sa.Integer(), nullable=False),\n sa.Column(\"credential_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"status\",\n sa.Enum(\n \"NOT_STARTED\",\n \"IN_PROGRESS\",\n \"SUCCESS\",\n \"FAILED\",\n name=\"deletionstatus\",\n native_enum=False,\n ),\n nullable=False,\n ),\n sa.Column(\"num_docs_deleted\", sa.Integer(), nullable=False),\n sa.Column(\"error_msg\", sa.String(), nullable=True),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\n \"time_updated\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"connector_id\"],\n [\"connector.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"credential_id\"],\n [\"credential.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.execute(\"DROP TABLE IF EXISTS document_by_connector_credential_pair CASCADE\")\n op.create_table(\n \"document_by_connector_credential_pair\",\n sa.Column(\"id\", sa.String(), nullable=False),\n sa.Column(\"connector_id\", sa.Integer(), nullable=False),\n sa.Column(\"credential_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"connector_id\"],\n [\"connector.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"credential_id\"],\n [\"credential.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"id\"],\n [\"document.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\", \"connector_id\", \"credential_id\"),\n )\n\n\ndef downgrade() -> None:\n # upstream tables first\n op.drop_table(\"document_by_connector_credential_pair\")\n op.drop_table(\"deletion_attempt\")\n op.drop_table(\"chunk\")\n\n # Alembic op.drop_table() has no \"cascade\" flag – issue raw SQL\n op.execute(\"DROP TABLE IF EXISTS document CASCADE\")\n" }, { "path": "backend/alembic/versions/df46c75b714e_add_default_vision_provider_to_llm_.py", "content": "\"\"\"add_default_vision_provider_to_llm_provider\n\nRevision ID: df46c75b714e\nRevises: 3934b1bc7b62\nCreate Date: 2025-03-11 16:20:19.038945\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"df46c75b714e\"\ndown_revision = \"3934b1bc7b62\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"llm_provider\",\n sa.Column(\n \"is_default_vision_provider\",\n sa.Boolean(),\n nullable=True,\n server_default=sa.false(),\n ),\n )\n op.add_column(\n \"llm_provider\", sa.Column(\"default_vision_model\", sa.String(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"llm_provider\", \"default_vision_model\")\n op.drop_column(\"llm_provider\", \"is_default_vision_provider\")\n" }, { "path": "backend/alembic/versions/dfbe9e93d3c7_extended_role_for_non_web.py", "content": "\"\"\"extended_role_for_non_web\n\nRevision ID: dfbe9e93d3c7\nRevises: 9cf5c00f72fe\nCreate Date: 2024-11-16 07:54:18.727906\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"dfbe9e93d3c7\"\ndown_revision = \"9cf5c00f72fe\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.execute(\n \"\"\"\n UPDATE \"user\"\n SET role = 'EXT_PERM_USER'\n WHERE has_web_login = false\n \"\"\"\n )\n op.drop_column(\"user\", \"has_web_login\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\"has_web_login\", sa.Boolean(), nullable=False, server_default=\"true\"),\n )\n\n op.execute(\n \"\"\"\n UPDATE \"user\"\n SET has_web_login = false,\n role = 'BASIC'\n WHERE role IN ('SLACK_USER', 'EXT_PERM_USER')\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/e0a68a81d434_add_chat_feedback.py", "content": "\"\"\"Add Chat Feedback\n\nRevision ID: e0a68a81d434\nRevises: ae62505e3acc\nCreate Date: 2023-10-04 20:22:33.380286\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"e0a68a81d434\"\ndown_revision = \"ae62505e3acc\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"chat_feedback\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"chat_message_chat_session_id\", sa.Integer(), nullable=False),\n sa.Column(\"chat_message_message_number\", sa.Integer(), nullable=False),\n sa.Column(\"chat_message_edit_number\", sa.Integer(), nullable=False),\n sa.Column(\"is_positive\", sa.Boolean(), nullable=True),\n sa.Column(\"feedback_text\", sa.Text(), nullable=True),\n sa.ForeignKeyConstraint(\n [\n \"chat_message_chat_session_id\",\n \"chat_message_message_number\",\n \"chat_message_edit_number\",\n ],\n [\n \"chat_message.chat_session_id\",\n \"chat_message.message_number\",\n \"chat_message.edit_number\",\n ],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"chat_feedback\")\n" }, { "path": "backend/alembic/versions/e1392f05e840_added_input_prompts.py", "content": "\"\"\"Added input prompts\n\nRevision ID: e1392f05e840\nRevises: 08a1eda20fe1\nCreate Date: 2024-07-13 19:09:22.556224\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"e1392f05e840\"\ndown_revision = \"08a1eda20fe1\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"inputprompt\",\n sa.Column(\"id\", sa.Integer(), autoincrement=True, nullable=False),\n sa.Column(\"prompt\", sa.String(), nullable=False),\n sa.Column(\"content\", sa.String(), nullable=False),\n sa.Column(\"active\", sa.Boolean(), nullable=False),\n sa.Column(\"is_public\", sa.Boolean(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=True,\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"id\"),\n )\n op.create_table(\n \"inputprompt__user\",\n sa.Column(\"input_prompt_id\", sa.Integer(), nullable=False),\n sa.Column(\"user_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"input_prompt_id\"],\n [\"inputprompt.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"inputprompt.id\"],\n ),\n sa.PrimaryKeyConstraint(\"input_prompt_id\", \"user_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"inputprompt__user\")\n op.drop_table(\"inputprompt\")\n" }, { "path": "backend/alembic/versions/e209dc5a8156_added_prune_frequency.py", "content": "\"\"\"added-prune-frequency\n\nRevision ID: e209dc5a8156\nRevises: 48d14957fe80\nCreate Date: 2024-06-16 16:02:35.273231\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nrevision = \"e209dc5a8156\"\ndown_revision = \"48d14957fe80\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\"connector\", sa.Column(\"prune_freq\", sa.Integer(), nullable=True))\n\n\ndef downgrade() -> None:\n op.drop_column(\"connector\", \"prune_freq\")\n" }, { "path": "backend/alembic/versions/e4334d5b33ba_add_deployment_name_to_llmprovider.py", "content": "\"\"\"add_deployment_name_to_llmprovider\n\nRevision ID: e4334d5b33ba\nRevises: ac5eaac849f9\nCreate Date: 2024-10-04 09:52:34.896867\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"e4334d5b33ba\"\ndown_revision = \"ac5eaac849f9\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"llm_provider\", sa.Column(\"deployment_name\", sa.String(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"llm_provider\", \"deployment_name\")\n" }, { "path": "backend/alembic/versions/e50154680a5c_no_source_enum.py", "content": "\"\"\"No Source Enum\n\nRevision ID: e50154680a5c\nRevises: fcd135795f21\nCreate Date: 2024-03-14 18:06:08.523106\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom onyx.configs.constants import DocumentSource\n\n# revision identifiers, used by Alembic.\nrevision = \"e50154680a5c\"\ndown_revision = \"fcd135795f21\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.alter_column(\n \"search_doc\",\n \"source_type\",\n type_=sa.String(length=50),\n existing_type=sa.Enum(DocumentSource, native_enum=False),\n existing_nullable=False,\n )\n op.execute(\"DROP TYPE IF EXISTS documentsource\")\n\n\ndef downgrade() -> None:\n op.alter_column(\n \"search_doc\",\n \"source_type\",\n type_=sa.Enum(DocumentSource, native_enum=False),\n existing_type=sa.String(length=50),\n existing_nullable=False,\n )\n" }, { "path": "backend/alembic/versions/e6a4bbc13fe4_add_index_for_retrieving_latest_index_.py", "content": "\"\"\"Add index for retrieving latest index_attempt\n\nRevision ID: e6a4bbc13fe4\nRevises: b082fec533f0\nCreate Date: 2023-08-10 12:37:23.335471\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"e6a4bbc13fe4\"\ndown_revision = \"b082fec533f0\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_index(\n op.f(\"ix_index_attempt_latest_for_connector_credential_pair\"),\n \"index_attempt\",\n [\"connector_id\", \"credential_id\", \"time_created\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\n op.f(\"ix_index_attempt_latest_for_connector_credential_pair\"),\n table_name=\"index_attempt\",\n )\n" }, { "path": "backend/alembic/versions/e7f8a9b0c1d2_create_anonymous_user.py", "content": "\"\"\"create_anonymous_user\n\nThis migration creates a permanent anonymous user in the database.\nWhen anonymous access is enabled, unauthenticated requests will use this user\ninstead of returning user_id=NULL.\n\nRevision ID: e7f8a9b0c1d2\nRevises: f7ca3e2f45d9\nCreate Date: 2026-01-15 14:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"e7f8a9b0c1d2\"\ndown_revision = \"f7ca3e2f45d9\"\nbranch_labels = None\ndepends_on = None\n\n# Must match constants in onyx/configs/constants.py file\nANONYMOUS_USER_UUID = \"00000000-0000-0000-0000-000000000002\"\nANONYMOUS_USER_EMAIL = \"anonymous@onyx.app\"\n\n# Tables with user_id foreign key that may need migration\nTABLES_WITH_USER_ID = [\n \"chat_session\",\n \"credential\",\n \"document_set\",\n \"persona\",\n \"tool\",\n \"notification\",\n \"inputprompt\",\n]\n\n\ndef _dedupe_null_notifications(connection: sa.Connection) -> None:\n # Multiple NULL-owned notifications can exist because the unique index treats\n # NULL user_id values as distinct. Before migrating them to the anonymous\n # user, collapse duplicates and remove rows that would conflict with an\n # already-existing anonymous notification.\n result = connection.execute(\n sa.text(\n \"\"\"\n WITH ranked_null_notifications AS (\n SELECT\n id,\n ROW_NUMBER() OVER (\n PARTITION BY notif_type, COALESCE(additional_data, '{}'::jsonb)\n ORDER BY first_shown DESC, last_shown DESC, id DESC\n ) AS row_num\n FROM notification\n WHERE user_id IS NULL\n )\n DELETE FROM notification\n WHERE id IN (\n SELECT id\n FROM ranked_null_notifications\n WHERE row_num > 1\n )\n \"\"\"\n )\n )\n if result.rowcount > 0:\n print(f\"Deleted {result.rowcount} duplicate NULL-owned notifications\")\n\n result = connection.execute(\n sa.text(\n \"\"\"\n DELETE FROM notification AS null_owned\n USING notification AS anonymous_owned\n WHERE null_owned.user_id IS NULL\n AND anonymous_owned.user_id = :user_id\n AND null_owned.notif_type = anonymous_owned.notif_type\n AND COALESCE(null_owned.additional_data, '{}'::jsonb) =\n COALESCE(anonymous_owned.additional_data, '{}'::jsonb)\n \"\"\"\n ),\n {\"user_id\": ANONYMOUS_USER_UUID},\n )\n if result.rowcount > 0:\n print(\n f\"Deleted {result.rowcount} NULL-owned notifications that conflict with existing anonymous-owned notifications\"\n )\n\n\ndef upgrade() -> None:\n \"\"\"\n Create the anonymous user for anonymous access feature.\n Also migrates any remaining user_id=NULL records to the anonymous user.\n \"\"\"\n connection = op.get_bind()\n\n # Create the anonymous user (using ON CONFLICT to be idempotent)\n connection.execute(\n sa.text(\n \"\"\"\n INSERT INTO \"user\" (id, email, hashed_password, is_active, is_superuser, is_verified, role)\n VALUES (:id, :email, :hashed_password, :is_active, :is_superuser, :is_verified, :role)\n ON CONFLICT (id) DO NOTHING\n \"\"\"\n ),\n {\n \"id\": ANONYMOUS_USER_UUID,\n \"email\": ANONYMOUS_USER_EMAIL,\n \"hashed_password\": \"\", # Empty password - user cannot log in directly\n \"is_active\": True, # Active so it can be used for anonymous access\n \"is_superuser\": False,\n \"is_verified\": True, # Verified since no email verification needed\n \"role\": \"LIMITED\", # Anonymous users have limited role to restrict access\n },\n )\n\n # Migrate any remaining user_id=NULL records to anonymous user\n for table in TABLES_WITH_USER_ID:\n # Dedup notifications outside the savepoint so deletions persist\n # even if the subsequent UPDATE rolls back\n if table == \"notification\":\n _dedupe_null_notifications(connection)\n\n with connection.begin_nested():\n # Exclude public credential (id=0) which must remain user_id=NULL\n # Exclude builtin tools (in_code_tool_id IS NOT NULL) which must remain user_id=NULL\n # Exclude builtin personas (builtin_persona=True) which must remain user_id=NULL\n # Exclude system input prompts (is_public=True with user_id=NULL) which must remain user_id=NULL\n if table == \"credential\":\n condition = \"user_id IS NULL AND id != 0\"\n elif table == \"tool\":\n condition = \"user_id IS NULL AND in_code_tool_id IS NULL\"\n elif table == \"persona\":\n condition = \"user_id IS NULL AND builtin_persona = false\"\n elif table == \"inputprompt\":\n condition = \"user_id IS NULL AND is_public = false\"\n else:\n condition = \"user_id IS NULL\"\n\n result = connection.execute(\n sa.text(\n f\"\"\"\n UPDATE \"{table}\"\n SET user_id = :user_id\n WHERE {condition}\n \"\"\"\n ),\n {\"user_id\": ANONYMOUS_USER_UUID},\n )\n if result.rowcount > 0:\n print(f\"Updated {result.rowcount} rows in {table} to anonymous user\")\n\n\ndef downgrade() -> None:\n \"\"\"\n Set anonymous user's records back to NULL and delete the anonymous user.\n\n Note: Duplicate NULL-owned notifications removed during upgrade are not restored.\n \"\"\"\n connection = op.get_bind()\n\n # Set records back to NULL\n for table in TABLES_WITH_USER_ID:\n with connection.begin_nested():\n connection.execute(\n sa.text(\n f\"\"\"\n UPDATE \"{table}\"\n SET user_id = NULL\n WHERE user_id = :user_id\n \"\"\"\n ),\n {\"user_id\": ANONYMOUS_USER_UUID},\n )\n\n # Delete the anonymous user\n connection.execute(\n sa.text('DELETE FROM \"user\" WHERE id = :user_id'),\n {\"user_id\": ANONYMOUS_USER_UUID},\n )\n" }, { "path": "backend/alembic/versions/e86866a9c78a_add_persona_to_chat_session.py", "content": "\"\"\"Add persona to chat_session\n\nRevision ID: e86866a9c78a\nRevises: 80696cf850ae\nCreate Date: 2023-11-26 02:51:47.657357\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"e86866a9c78a\"\ndown_revision = \"80696cf850ae\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\"chat_session\", sa.Column(\"persona_id\", sa.Integer(), nullable=True))\n op.create_foreign_key(\n \"fk_chat_session_persona_id\", \"chat_session\", \"persona\", [\"persona_id\"], [\"id\"]\n )\n\n\ndef downgrade() -> None:\n op.drop_constraint(\"fk_chat_session_persona_id\", \"chat_session\", type_=\"foreignkey\")\n op.drop_column(\"chat_session\", \"persona_id\")\n" }, { "path": "backend/alembic/versions/e8f0d2a38171_add_status_to_mcp_server_and_make_auth_.py", "content": "\"\"\"add status to mcp server and make auth fields nullable\n\nRevision ID: e8f0d2a38171\nRevises: ed9e44312505\nCreate Date: 2025-11-28 11:15:37.667340\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom onyx.db.enums import (\n MCPTransport,\n MCPAuthenticationType,\n MCPAuthenticationPerformer,\n MCPServerStatus,\n)\n\n# revision identifiers, used by Alembic.\nrevision = \"e8f0d2a38171\"\ndown_revision = \"ed9e44312505\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Make auth fields nullable\n op.alter_column(\n \"mcp_server\",\n \"transport\",\n existing_type=sa.Enum(MCPTransport, name=\"mcp_transport\", native_enum=False),\n nullable=True,\n )\n\n op.alter_column(\n \"mcp_server\",\n \"auth_type\",\n existing_type=sa.Enum(\n MCPAuthenticationType, name=\"mcp_authentication_type\", native_enum=False\n ),\n nullable=True,\n )\n\n op.alter_column(\n \"mcp_server\",\n \"auth_performer\",\n existing_type=sa.Enum(\n MCPAuthenticationPerformer,\n name=\"mcp_authentication_performer\",\n native_enum=False,\n ),\n nullable=True,\n )\n\n # Add status column with default\n op.add_column(\n \"mcp_server\",\n sa.Column(\n \"status\",\n sa.Enum(MCPServerStatus, name=\"mcp_server_status\", native_enum=False),\n nullable=False,\n server_default=\"CREATED\",\n ),\n )\n\n # For existing records, mark status as CONNECTED\n bind = op.get_bind()\n bind.execute(\n sa.text(\n \"\"\"\n UPDATE mcp_server\n SET status = 'CONNECTED'\n WHERE status != 'CONNECTED'\n and admin_connection_config_id IS NOT NULL\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n # Remove status column\n op.drop_column(\"mcp_server\", \"status\")\n\n # Make auth fields non-nullable (set defaults first)\n op.execute(\n \"UPDATE mcp_server SET transport = 'STREAMABLE_HTTP' WHERE transport IS NULL\"\n )\n op.execute(\"UPDATE mcp_server SET auth_type = 'NONE' WHERE auth_type IS NULL\")\n op.execute(\n \"UPDATE mcp_server SET auth_performer = 'ADMIN' WHERE auth_performer IS NULL\"\n )\n\n op.alter_column(\n \"mcp_server\",\n \"transport\",\n existing_type=sa.Enum(MCPTransport, name=\"mcp_transport\", native_enum=False),\n nullable=False,\n )\n op.alter_column(\n \"mcp_server\",\n \"auth_type\",\n existing_type=sa.Enum(\n MCPAuthenticationType, name=\"mcp_authentication_type\", native_enum=False\n ),\n nullable=False,\n )\n op.alter_column(\n \"mcp_server\",\n \"auth_performer\",\n existing_type=sa.Enum(\n MCPAuthenticationPerformer,\n name=\"mcp_authentication_performer\",\n native_enum=False,\n ),\n nullable=False,\n )\n" }, { "path": "backend/alembic/versions/e91df4e935ef_private_personas_documentsets.py", "content": "\"\"\"Private Personas DocumentSets\n\nRevision ID: e91df4e935ef\nRevises: 91fd3b470d1a\nCreate Date: 2024-03-17 11:47:24.675881\n\n\"\"\"\n\nimport fastapi_users_db_sqlalchemy\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"e91df4e935ef\"\ndown_revision = \"91fd3b470d1a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"document_set__user\",\n sa.Column(\"document_set_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"document_set_id\"],\n [\"document_set.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"document_set_id\", \"user_id\"),\n )\n op.create_table(\n \"persona__user\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_id\",\n fastapi_users_db_sqlalchemy.generics.GUID(),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_id\"],\n [\"user.id\"],\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"user_id\"),\n )\n op.create_table(\n \"document_set__user_group\",\n sa.Column(\"document_set_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_group_id\",\n sa.Integer(),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"document_set_id\"],\n [\"document_set.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.PrimaryKeyConstraint(\"document_set_id\", \"user_group_id\"),\n )\n op.create_table(\n \"persona__user_group\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\n \"user_group_id\",\n sa.Integer(),\n nullable=False,\n ),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"user_group_id\"],\n [\"user_group.id\"],\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"user_group_id\"),\n )\n\n op.add_column(\n \"document_set\",\n sa.Column(\"is_public\", sa.Boolean(), nullable=True),\n )\n # fill in is_public for existing rows\n op.execute(\"UPDATE document_set SET is_public = true WHERE is_public IS NULL\")\n op.alter_column(\"document_set\", \"is_public\", nullable=False)\n\n op.add_column(\n \"persona\",\n sa.Column(\"is_public\", sa.Boolean(), nullable=True),\n )\n # fill in is_public for existing rows\n op.execute(\"UPDATE persona SET is_public = true WHERE is_public IS NULL\")\n op.alter_column(\"persona\", \"is_public\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_column(\"persona\", \"is_public\")\n\n op.drop_column(\"document_set\", \"is_public\")\n\n op.drop_table(\"persona__user\")\n op.drop_table(\"document_set__user\")\n op.drop_table(\"persona__user_group\")\n op.drop_table(\"document_set__user_group\")\n" }, { "path": "backend/alembic/versions/eaa3b5593925_add_default_slack_channel_config.py", "content": "\"\"\"add default slack channel config\n\nRevision ID: eaa3b5593925\nRevises: 98a5008d8711\nCreate Date: 2025-02-03 18:07:56.552526\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"eaa3b5593925\"\ndown_revision = \"98a5008d8711\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add is_default column\n op.add_column(\n \"slack_channel_config\",\n sa.Column(\"is_default\", sa.Boolean(), nullable=False, server_default=\"false\"),\n )\n\n op.create_index(\n \"ix_slack_channel_config_slack_bot_id_default\",\n \"slack_channel_config\",\n [\"slack_bot_id\", \"is_default\"],\n unique=True,\n postgresql_where=sa.text(\"is_default IS TRUE\"),\n )\n\n # Create default channel configs for existing slack bots without one\n conn = op.get_bind()\n slack_bots = conn.execute(sa.text(\"SELECT id FROM slack_bot\")).fetchall()\n\n for slack_bot in slack_bots:\n slack_bot_id = slack_bot[0]\n existing_default = conn.execute(\n sa.text(\n \"SELECT id FROM slack_channel_config WHERE slack_bot_id = :bot_id AND is_default = TRUE\"\n ),\n {\"bot_id\": slack_bot_id},\n ).fetchone()\n\n if not existing_default:\n conn.execute(\n sa.text(\n \"\"\"\n INSERT INTO slack_channel_config (\n slack_bot_id, persona_id, channel_config, enable_auto_filters, is_default\n ) VALUES (\n :bot_id, NULL,\n '{\"channel_name\": null, '\n '\"respond_member_group_list\": [], '\n '\"answer_filters\": [], '\n '\"follow_up_tags\": [], '\n '\"respond_tag_only\": true}',\n FALSE, TRUE\n )\n \"\"\"\n ),\n {\"bot_id\": slack_bot_id},\n )\n\n\ndef downgrade() -> None:\n # Delete default slack channel configs\n conn = op.get_bind()\n conn.execute(sa.text(\"DELETE FROM slack_channel_config WHERE is_default = TRUE\"))\n\n # Remove index\n op.drop_index(\n \"ix_slack_channel_config_slack_bot_id_default\",\n table_name=\"slack_channel_config\",\n )\n\n # Remove is_default column\n op.drop_column(\"slack_channel_config\", \"is_default\")\n" }, { "path": "backend/alembic/versions/ec3ec2eabf7b_index_from_beginning.py", "content": "\"\"\"Index From Beginning\n\nRevision ID: ec3ec2eabf7b\nRevises: dbaa756c2ccf\nCreate Date: 2024-02-06 22:03:28.098158\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"ec3ec2eabf7b\"\ndown_revision = \"dbaa756c2ccf\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"index_attempt\", sa.Column(\"from_beginning\", sa.Boolean(), nullable=True)\n )\n op.execute(\"UPDATE index_attempt SET from_beginning = False\")\n op.alter_column(\"index_attempt\", \"from_beginning\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_column(\"index_attempt\", \"from_beginning\")\n" }, { "path": "backend/alembic/versions/ec85f2b3c544_remove_last_attempt_status_from_cc_pair.py", "content": "\"\"\"Remove Last Attempt Status from CC Pair\n\nRevision ID: ec85f2b3c544\nRevises: 3879338f8ba1\nCreate Date: 2024-05-23 21:39:46.126010\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"ec85f2b3c544\"\ndown_revision = \"70f00c45c0f2\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.drop_column(\"connector_credential_pair\", \"last_attempt_status\")\n\n\ndef downgrade() -> None:\n op.add_column(\n \"connector_credential_pair\",\n sa.Column(\n \"last_attempt_status\",\n sa.VARCHAR(),\n autoincrement=False,\n nullable=True,\n ),\n )\n" }, { "path": "backend/alembic/versions/ecab2b3f1a3b_add_overrides_to_the_chat_session.py", "content": "\"\"\"Add overrides to the chat session\n\nRevision ID: ecab2b3f1a3b\nRevises: 38eda64af7fe\nCreate Date: 2024-04-01 19:08:21.359102\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"ecab2b3f1a3b\"\ndown_revision = \"38eda64af7fe\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_session\",\n sa.Column(\n \"llm_override\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n op.add_column(\n \"chat_session\",\n sa.Column(\n \"prompt_override\",\n postgresql.JSONB(astext_type=sa.Text()),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_session\", \"prompt_override\")\n op.drop_column(\"chat_session\", \"llm_override\")\n" }, { "path": "backend/alembic/versions/ed9e44312505_add_icon_name_field.py", "content": "\"\"\"Add icon_name field\n\nRevision ID: ed9e44312505\nRevises: 5e6f7a8b9c0d\nCreate Date: 2025-12-03 16:35:07.828393\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"ed9e44312505\"\ndown_revision = \"5e6f7a8b9c0d\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add icon_name column\n op.add_column(\"persona\", sa.Column(\"icon_name\", sa.String(), nullable=True))\n\n # Remove old icon columns\n op.drop_column(\"persona\", \"icon_shape\")\n op.drop_column(\"persona\", \"icon_color\")\n\n\ndef downgrade() -> None:\n # Re-add old icon columns\n op.add_column(\"persona\", sa.Column(\"icon_color\", sa.String(), nullable=True))\n op.add_column(\"persona\", sa.Column(\"icon_shape\", sa.Integer(), nullable=True))\n\n # Remove icon_name column\n op.drop_column(\"persona\", \"icon_name\")\n" }, { "path": "backend/alembic/versions/ee3f4b47fad5_added_alternate_model_to_chat_message.py", "content": "\"\"\"Added alternate model to chat message\n\nRevision ID: ee3f4b47fad5\nRevises: 2d2304e27d8c\nCreate Date: 2024-08-12 00:11:50.915845\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"ee3f4b47fad5\"\ndown_revision = \"2d2304e27d8c\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\",\n sa.Column(\"overridden_model\", sa.String(length=255), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"overridden_model\")\n" }, { "path": "backend/alembic/versions/ef7da92f7213_add_files_to_chatmessage.py", "content": "\"\"\"Add files to ChatMessage\n\nRevision ID: ef7da92f7213\nRevises: 401c1ac29467\nCreate Date: 2024-04-28 16:59:33.199153\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"ef7da92f7213\"\ndown_revision = \"401c1ac29467\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_message\",\n sa.Column(\"files\", postgresql.JSONB(astext_type=sa.Text()), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_message\", \"files\")\n" }, { "path": "backend/alembic/versions/efb35676026c_standard_answer_match_regex_flag.py", "content": "\"\"\"standard answer match_regex flag\n\nRevision ID: efb35676026c\nRevises: 0ebb1d516877\nCreate Date: 2024-09-11 13:55:46.101149\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"efb35676026c\"\ndown_revision = \"0ebb1d516877\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # ### commands auto generated by Alembic - please adjust! ###\n op.add_column(\n \"standard_answer\",\n sa.Column(\n \"match_regex\", sa.Boolean(), nullable=False, server_default=sa.false()\n ),\n )\n # ### end Alembic commands ###\n\n\ndef downgrade() -> None:\n # ### commands auto generated by Alembic - please adjust! ###\n op.drop_column(\"standard_answer\", \"match_regex\")\n # ### end Alembic commands ###\n" }, { "path": "backend/alembic/versions/f11b408e39d3_force_lowercase_all_users.py", "content": "\"\"\"force lowercase all users\n\nRevision ID: f11b408e39d3\nRevises: 3bd4c84fe72f\nCreate Date: 2025-02-26 17:04:55.683500\n\n\"\"\"\n\n# revision identifiers, used by Alembic.\nrevision = \"f11b408e39d3\"\ndown_revision = \"3bd4c84fe72f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # 1) Convert all existing user emails to lowercase\n from alembic import op\n\n op.execute(\n \"\"\"\n UPDATE \"user\"\n SET email = LOWER(email)\n \"\"\"\n )\n\n # 2) Add a check constraint to ensure emails are always lowercase\n op.create_check_constraint(\"ensure_lowercase_email\", \"user\", \"email = LOWER(email)\")\n\n\ndef downgrade() -> None:\n # Drop the check constraint\n from alembic import op\n\n op.drop_constraint(\"ensure_lowercase_email\", \"user\", type_=\"check\")\n" }, { "path": "backend/alembic/versions/f13db29f3101_add_composite_index_for_last_modified_.py", "content": "\"\"\"Add composite index for last_modified and last_synced to document\n\nRevision ID: f13db29f3101\nRevises: b388730a2899\nCreate Date: 2025-02-18 22:48:11.511389\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"f13db29f3101\"\ndown_revision = \"acaab4ef4507\"\nbranch_labels: str | None = None\ndepends_on: str | None = None\n\n\ndef upgrade() -> None:\n op.create_index(\n \"ix_document_sync_status\",\n \"document\",\n [\"last_modified\", \"last_synced\"],\n unique=False,\n )\n\n\ndef downgrade() -> None:\n op.drop_index(\"ix_document_sync_status\", table_name=\"document\")\n" }, { "path": "backend/alembic/versions/f17bf3b0d9f1_embedding_provider_by_provider_type.py", "content": "\"\"\"embedding provider by provider type\n\nRevision ID: f17bf3b0d9f1\nRevises: 351faebd379d\nCreate Date: 2024-08-21 13:13:31.120460\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f17bf3b0d9f1\"\ndown_revision = \"351faebd379d\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Add provider_type column to embedding_provider\n op.add_column(\n \"embedding_provider\",\n sa.Column(\"provider_type\", sa.String(50), nullable=True),\n )\n\n # Update provider_type with existing name values\n op.execute(\"UPDATE embedding_provider SET provider_type = UPPER(name)\")\n\n # Make provider_type not nullable\n op.alter_column(\"embedding_provider\", \"provider_type\", nullable=False)\n\n # Drop the foreign key constraint in embedding_model table\n op.drop_constraint(\n \"fk_embedding_model_cloud_provider\", \"embedding_model\", type_=\"foreignkey\"\n )\n\n # Drop the existing primary key constraint\n op.drop_constraint(\"embedding_provider_pkey\", \"embedding_provider\", type_=\"primary\")\n\n # Create a new primary key constraint on provider_type\n op.create_primary_key(\n \"embedding_provider_pkey\", \"embedding_provider\", [\"provider_type\"]\n )\n\n # Add provider_type column to embedding_model\n op.add_column(\n \"embedding_model\",\n sa.Column(\"provider_type\", sa.String(50), nullable=True),\n )\n\n # Update provider_type for existing embedding models\n op.execute(\n \"\"\"\n UPDATE embedding_model\n SET provider_type = (\n SELECT provider_type\n FROM embedding_provider\n WHERE embedding_provider.id = embedding_model.cloud_provider_id\n )\n \"\"\"\n )\n\n # Drop the old id column from embedding_provider\n op.drop_column(\"embedding_provider\", \"id\")\n\n # Drop the name column from embedding_provider\n op.drop_column(\"embedding_provider\", \"name\")\n\n # Drop the default_model_id column from embedding_provider\n op.drop_column(\"embedding_provider\", \"default_model_id\")\n\n # Drop the old cloud_provider_id column from embedding_model\n op.drop_column(\"embedding_model\", \"cloud_provider_id\")\n\n # Create the new foreign key constraint\n op.create_foreign_key(\n \"fk_embedding_model_cloud_provider\",\n \"embedding_model\",\n \"embedding_provider\",\n [\"provider_type\"],\n [\"provider_type\"],\n )\n\n\ndef downgrade() -> None:\n # Drop the foreign key constraint in embedding_model table\n op.drop_constraint(\n \"fk_embedding_model_cloud_provider\", \"embedding_model\", type_=\"foreignkey\"\n )\n\n # Add back the cloud_provider_id column to embedding_model\n op.add_column(\n \"embedding_model\", sa.Column(\"cloud_provider_id\", sa.Integer(), nullable=True)\n )\n op.add_column(\"embedding_provider\", sa.Column(\"id\", sa.Integer(), nullable=True))\n\n # Assign incrementing IDs to embedding providers\n op.execute(\n \"\"\"\n CREATE SEQUENCE IF NOT EXISTS embedding_provider_id_seq;\"\"\"\n )\n op.execute(\n \"\"\"\n UPDATE embedding_provider SET id = nextval('embedding_provider_id_seq');\n \"\"\"\n )\n\n # Update cloud_provider_id based on provider_type\n op.execute(\n \"\"\"\n UPDATE embedding_model\n SET cloud_provider_id = CASE\n WHEN provider_type IS NULL THEN NULL\n ELSE (\n SELECT id\n FROM embedding_provider\n WHERE embedding_provider.provider_type = embedding_model.provider_type\n )\n END\n \"\"\"\n )\n\n # Drop the provider_type column from embedding_model\n op.drop_column(\"embedding_model\", \"provider_type\")\n\n # Add back the columns to embedding_provider\n op.add_column(\"embedding_provider\", sa.Column(\"name\", sa.String(50), nullable=True))\n op.add_column(\n \"embedding_provider\", sa.Column(\"default_model_id\", sa.Integer(), nullable=True)\n )\n\n # Drop the existing primary key constraint on provider_type\n op.drop_constraint(\"embedding_provider_pkey\", \"embedding_provider\", type_=\"primary\")\n\n # Create the original primary key constraint on id\n op.create_primary_key(\"embedding_provider_pkey\", \"embedding_provider\", [\"id\"])\n\n # Update name with existing provider_type values\n op.execute(\n \"\"\"\n UPDATE embedding_provider\n SET name = CASE\n WHEN provider_type = 'OPENAI' THEN 'OpenAI'\n WHEN provider_type = 'COHERE' THEN 'Cohere'\n WHEN provider_type = 'GOOGLE' THEN 'Google'\n WHEN provider_type = 'VOYAGE' THEN 'Voyage'\n ELSE provider_type\n END\n \"\"\"\n )\n\n # Drop the provider_type column from embedding_provider\n op.drop_column(\"embedding_provider\", \"provider_type\")\n\n # Recreate the foreign key constraint in embedding_model table\n op.create_foreign_key(\n \"fk_embedding_model_cloud_provider\",\n \"embedding_model\",\n \"embedding_provider\",\n [\"cloud_provider_id\"],\n [\"id\"],\n )\n\n # Recreate the foreign key constraint in embedding_model table\n op.create_foreign_key(\n \"fk_embedding_provider_default_model\",\n \"embedding_provider\",\n \"embedding_model\",\n [\"default_model_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/f1c6478c3fd8_add_pre_defined_feedback.py", "content": "\"\"\"Add pre-defined feedback\n\nRevision ID: f1c6478c3fd8\nRevises: 643a84a42a33\nCreate Date: 2024-05-09 18:11:49.210667\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nrevision = \"f1c6478c3fd8\"\ndown_revision = \"643a84a42a33\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"chat_feedback\",\n sa.Column(\"predefined_feedback\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"chat_feedback\", \"predefined_feedback\")\n" }, { "path": "backend/alembic/versions/f1ca58b2f2ec_add_passthrough_auth_to_tool.py", "content": "\"\"\"add passthrough auth to tool\n\nRevision ID: f1ca58b2f2ec\nRevises: c7bf5721733e\nCreate Date: 2024-03-19\n\n\"\"\"\n\nfrom typing import Sequence, Union\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision: str = \"f1ca58b2f2ec\"\ndown_revision: Union[str, None] = \"c7bf5721733e\"\nbranch_labels: Union[str, Sequence[str], None] = None\ndepends_on: Union[str, Sequence[str], None] = None\n\n\ndef upgrade() -> None:\n # Add passthrough_auth column to tool table with default value of False\n op.add_column(\n \"tool\",\n sa.Column(\n \"passthrough_auth\", sa.Boolean(), nullable=False, server_default=sa.false()\n ),\n )\n\n\ndef downgrade() -> None:\n # Remove passthrough_auth column from tool table\n op.drop_column(\"tool\", \"passthrough_auth\")\n" }, { "path": "backend/alembic/versions/f220515df7b4_add_flow_mapping_table.py", "content": "\"\"\"Add flow mapping table\n\nRevision ID: f220515df7b4\nRevises: cbc03e08d0f3\nCreate Date: 2026-01-30 12:21:24.955922\n\n\"\"\"\n\nfrom onyx.db.enums import LLMModelFlowType\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f220515df7b4\"\ndown_revision = \"9d1543a37106\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"llm_model_flow\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\n \"llm_model_flow_type\",\n sa.Enum(LLMModelFlowType, name=\"llmmodelflowtype\", native_enum=False),\n nullable=False,\n ),\n sa.Column(\n \"is_default\", sa.Boolean(), nullable=False, server_default=sa.text(\"false\")\n ),\n sa.Column(\"model_configuration_id\", sa.Integer(), nullable=False),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.ForeignKeyConstraint(\n [\"model_configuration_id\"], [\"model_configuration.id\"], ondelete=\"CASCADE\"\n ),\n sa.UniqueConstraint(\n \"llm_model_flow_type\",\n \"model_configuration_id\",\n name=\"uq_model_config_per_llm_model_flow_type\",\n ),\n )\n\n # Partial unique index so that there is at most one default for each flow type\n op.create_index(\n \"ix_one_default_per_llm_model_flow\",\n \"llm_model_flow\",\n [\"llm_model_flow_type\"],\n unique=True,\n postgresql_where=sa.text(\"is_default IS TRUE\"),\n )\n\n\ndef downgrade() -> None:\n # Drop the llm_model_flow table (index is dropped automatically with table)\n op.drop_table(\"llm_model_flow\")\n" }, { "path": "backend/alembic/versions/f32615f71aeb_add_custom_headers_to_tools.py", "content": "\"\"\"add custom headers to tools\n\nRevision ID: f32615f71aeb\nRevises: bd2921608c3a\nCreate Date: 2024-09-12 20:26:38.932377\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"f32615f71aeb\"\ndown_revision = \"bd2921608c3a\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"tool\", sa.Column(\"custom_headers\", postgresql.JSONB(), nullable=True)\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"tool\", \"custom_headers\")\n" }, { "path": "backend/alembic/versions/f39c5794c10a_add_background_errors_table.py", "content": "\"\"\"Add background errors table\n\nRevision ID: f39c5794c10a\nRevises: 2cdeff6d8c93\nCreate Date: 2025-02-12 17:11:14.527876\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"f39c5794c10a\"\ndown_revision = \"2cdeff6d8c93\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"background_error\",\n sa.Column(\"id\", sa.Integer(), nullable=False),\n sa.Column(\"message\", sa.String(), nullable=False),\n sa.Column(\n \"time_created\",\n sa.DateTime(timezone=True),\n server_default=sa.text(\"now()\"),\n nullable=False,\n ),\n sa.Column(\"cc_pair_id\", sa.Integer(), nullable=True),\n sa.PrimaryKeyConstraint(\"id\"),\n sa.ForeignKeyConstraint(\n [\"cc_pair_id\"],\n [\"connector_credential_pair.id\"],\n ondelete=\"CASCADE\",\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"background_error\")\n" }, { "path": "backend/alembic/versions/f5437cc136c5_delete_non_search_assistants.py", "content": "\"\"\"delete non-search assistants\n\nRevision ID: f5437cc136c5\nRevises: eaa3b5593925\nCreate Date: 2025-02-04 16:17:15.677256\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f5437cc136c5\"\ndown_revision = \"eaa3b5593925\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n pass\n\n\ndef downgrade() -> None:\n # Fix: split the statements into multiple op.execute() calls\n op.execute(\n \"\"\"\n WITH personas_without_search AS (\n SELECT p.id\n FROM persona p\n LEFT JOIN persona__tool pt ON p.id = pt.persona_id\n LEFT JOIN tool t ON pt.tool_id = t.id\n GROUP BY p.id\n HAVING COUNT(CASE WHEN t.in_code_tool_id = 'run_search' THEN 1 END) = 0\n )\n UPDATE slack_channel_config\n SET persona_id = NULL\n WHERE is_default = TRUE AND persona_id IN (SELECT id FROM personas_without_search)\n \"\"\"\n )\n\n op.execute(\n \"\"\"\n WITH personas_without_search AS (\n SELECT p.id\n FROM persona p\n LEFT JOIN persona__tool pt ON p.id = pt.persona_id\n LEFT JOIN tool t ON pt.tool_id = t.id\n GROUP BY p.id\n HAVING COUNT(CASE WHEN t.in_code_tool_id = 'run_search' THEN 1 END) = 0\n )\n DELETE FROM slack_channel_config\n WHERE is_default = FALSE AND persona_id IN (SELECT id FROM personas_without_search)\n \"\"\"\n )\n" }, { "path": "backend/alembic/versions/f71470ba9274_add_prompt_length_limit.py", "content": "\"\"\"add prompt length limit\n\nRevision ID: f71470ba9274\nRevises: 6a804aeb4830\nCreate Date: 2025-04-01 15:07:14.977435\n\n\"\"\"\n\n# revision identifiers, used by Alembic.\nrevision = \"f71470ba9274\"\ndown_revision = \"6a804aeb4830\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # op.alter_column(\n # \"prompt\",\n # \"system_prompt\",\n # existing_type=sa.TEXT(),\n # type_=sa.String(length=8000),\n # existing_nullable=False,\n # )\n # op.alter_column(\n # \"prompt\",\n # \"task_prompt\",\n # existing_type=sa.TEXT(),\n # type_=sa.String(length=8000),\n # existing_nullable=False,\n # )\n pass\n\n\ndef downgrade() -> None:\n # op.alter_column(\n # \"prompt\",\n # \"system_prompt\",\n # existing_type=sa.String(length=8000),\n # type_=sa.TEXT(),\n # existing_nullable=False,\n # )\n # op.alter_column(\n # \"prompt\",\n # \"task_prompt\",\n # existing_type=sa.String(length=8000),\n # type_=sa.TEXT(),\n # existing_nullable=False,\n # )\n pass\n" }, { "path": "backend/alembic/versions/f7505c5b0284_updated_constraints_for_ccpairs.py", "content": "\"\"\"updated constraints for ccpairs\n\nRevision ID: f7505c5b0284\nRevises: f71470ba9274\nCreate Date: 2025-04-01 17:50:42.504818\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f7505c5b0284\"\ndown_revision = \"f71470ba9274\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # 1) Drop the old foreign-key constraints\n op.drop_constraint(\n \"document_by_connector_credential_pair_connector_id_fkey\",\n \"document_by_connector_credential_pair\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"document_by_connector_credential_pair_credential_id_fkey\",\n \"document_by_connector_credential_pair\",\n type_=\"foreignkey\",\n )\n\n # 2) Re-add them with ondelete='CASCADE'\n op.create_foreign_key(\n \"document_by_connector_credential_pair_connector_id_fkey\",\n source_table=\"document_by_connector_credential_pair\",\n referent_table=\"connector\",\n local_cols=[\"connector_id\"],\n remote_cols=[\"id\"],\n ondelete=\"CASCADE\",\n )\n op.create_foreign_key(\n \"document_by_connector_credential_pair_credential_id_fkey\",\n source_table=\"document_by_connector_credential_pair\",\n referent_table=\"credential\",\n local_cols=[\"credential_id\"],\n remote_cols=[\"id\"],\n ondelete=\"CASCADE\",\n )\n\n\ndef downgrade() -> None:\n # Reverse the changes for rollback\n op.drop_constraint(\n \"document_by_connector_credential_pair_connector_id_fkey\",\n \"document_by_connector_credential_pair\",\n type_=\"foreignkey\",\n )\n op.drop_constraint(\n \"document_by_connector_credential_pair_credential_id_fkey\",\n \"document_by_connector_credential_pair\",\n type_=\"foreignkey\",\n )\n\n # Recreate without CASCADE\n op.create_foreign_key(\n \"document_by_connector_credential_pair_connector_id_fkey\",\n \"document_by_connector_credential_pair\",\n \"connector\",\n [\"connector_id\"],\n [\"id\"],\n )\n op.create_foreign_key(\n \"document_by_connector_credential_pair_credential_id_fkey\",\n \"document_by_connector_credential_pair\",\n \"credential\",\n [\"credential_id\"],\n [\"id\"],\n )\n" }, { "path": "backend/alembic/versions/f7a894b06d02_non_nullbale_slack_bot_id_in_channel_.py", "content": "\"\"\"non-nullbale slack bot id in channel config\n\nRevision ID: f7a894b06d02\nRevises: 9f696734098f\nCreate Date: 2024-12-06 12:55:42.845723\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f7a894b06d02\"\ndown_revision = \"9f696734098f\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Delete all rows with null slack_bot_id\n op.execute(\"DELETE FROM slack_channel_config WHERE slack_bot_id IS NULL\")\n\n # Make slack_bot_id non-nullable\n op.alter_column(\n \"slack_channel_config\",\n \"slack_bot_id\",\n existing_type=sa.Integer(),\n nullable=False,\n )\n\n\ndef downgrade() -> None:\n # Make slack_bot_id nullable again\n op.alter_column(\n \"slack_channel_config\",\n \"slack_bot_id\",\n existing_type=sa.Integer(),\n nullable=True,\n )\n" }, { "path": "backend/alembic/versions/f7ca3e2f45d9_migrate_no_auth_data_to_placeholder.py", "content": "\"\"\"migrate_no_auth_data_to_placeholder\n\nThis migration handles the transition from AUTH_TYPE=disabled to requiring\nauthentication. It creates a placeholder user and assigns all data that was\ncreated without a user (user_id=NULL) to this placeholder.\n\nA database trigger is installed that automatically transfers all data from\nthe placeholder user to the first real user who registers, then drops itself.\n\nRevision ID: f7ca3e2f45d9\nRevises: 78ebc66946a0\nCreate Date: 2026-01-15 12:49:53.802741\n\n\"\"\"\n\nimport os\n\nfrom alembic import op\nimport sqlalchemy as sa\n\nfrom shared_configs.configs import MULTI_TENANT\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f7ca3e2f45d9\"\ndown_revision = \"78ebc66946a0\"\nbranch_labels = None\ndepends_on = None\n\n# Must match constants in onyx/configs/constants.py file\nNO_AUTH_PLACEHOLDER_USER_UUID = \"00000000-0000-0000-0000-000000000001\"\nNO_AUTH_PLACEHOLDER_USER_EMAIL = \"no-auth-placeholder@onyx.app\"\n\n# Trigger and function names\nTRIGGER_NAME = \"trg_migrate_no_auth_data\"\nFUNCTION_NAME = \"migrate_no_auth_data_to_user\"\n\n# Trigger function that migrates data from placeholder to first real user\nMIGRATE_NO_AUTH_TRIGGER_FUNCTION = f\"\"\"\nCREATE OR REPLACE FUNCTION {FUNCTION_NAME}()\nRETURNS TRIGGER AS $$\nDECLARE\n placeholder_uuid UUID := '00000000-0000-0000-0000-000000000001'::uuid;\n anonymous_uuid UUID := '00000000-0000-0000-0000-000000000002'::uuid;\n placeholder_row RECORD;\n schema_name TEXT;\nBEGIN\n -- Skip if this is the placeholder user being inserted\n IF NEW.id = placeholder_uuid THEN\n RETURN NULL;\n END IF;\n\n -- Skip if this is the anonymous user being inserted (not a real user)\n IF NEW.id = anonymous_uuid THEN\n RETURN NULL;\n END IF;\n\n -- Skip if the new user is not active\n IF NEW.is_active = FALSE THEN\n RETURN NULL;\n END IF;\n\n -- Get current schema for self-cleanup\n schema_name := current_schema();\n\n -- Try to lock the placeholder user row with FOR UPDATE SKIP LOCKED\n -- This ensures only one concurrent transaction can proceed with migration\n -- SKIP LOCKED means if another transaction has the lock, we skip (don't wait)\n SELECT id INTO placeholder_row\n FROM \"user\"\n WHERE id = placeholder_uuid\n FOR UPDATE SKIP LOCKED;\n\n IF NOT FOUND THEN\n -- Either placeholder doesn't exist or another transaction has it locked\n -- Either way, drop the trigger and return without making admin\n EXECUTE format('DROP TRIGGER IF EXISTS {TRIGGER_NAME} ON %I.\"user\"', schema_name);\n EXECUTE format('DROP FUNCTION IF EXISTS %I.{FUNCTION_NAME}()', schema_name);\n RETURN NULL;\n END IF;\n\n -- We have exclusive lock on placeholder - proceed with migration\n -- The INSERT has already completed (AFTER INSERT), so NEW.id exists in the table\n\n -- Migrate chat_session\n UPDATE \"chat_session\" SET user_id = NEW.id WHERE user_id = placeholder_uuid;\n\n -- Migrate credential (exclude public credential id=0)\n UPDATE \"credential\" SET user_id = NEW.id WHERE user_id = placeholder_uuid AND id != 0;\n\n -- Migrate document_set\n UPDATE \"document_set\" SET user_id = NEW.id WHERE user_id = placeholder_uuid;\n\n -- Migrate persona (exclude builtin personas)\n UPDATE \"persona\" SET user_id = NEW.id WHERE user_id = placeholder_uuid AND builtin_persona = FALSE;\n\n -- Migrate tool (exclude builtin tools)\n UPDATE \"tool\" SET user_id = NEW.id WHERE user_id = placeholder_uuid AND in_code_tool_id IS NULL;\n\n -- Migrate notification\n UPDATE \"notification\" SET user_id = NEW.id WHERE user_id = placeholder_uuid;\n\n -- Migrate inputprompt (exclude system/public prompts)\n UPDATE \"inputprompt\" SET user_id = NEW.id WHERE user_id = placeholder_uuid AND is_public = FALSE;\n\n -- Make the new user an admin (they had admin access in no-auth mode)\n -- In AFTER INSERT trigger, we must UPDATE the row since it already exists\n UPDATE \"user\" SET role = 'ADMIN' WHERE id = NEW.id;\n\n -- Delete the placeholder user (we hold the lock so this is safe)\n DELETE FROM \"user\" WHERE id = placeholder_uuid;\n\n -- Drop the trigger and function (self-cleanup)\n EXECUTE format('DROP TRIGGER IF EXISTS {TRIGGER_NAME} ON %I.\"user\"', schema_name);\n EXECUTE format('DROP FUNCTION IF EXISTS %I.{FUNCTION_NAME}()', schema_name);\n\n RETURN NULL;\nEND;\n$$ LANGUAGE plpgsql;\n\"\"\"\n\nMIGRATE_NO_AUTH_TRIGGER = f\"\"\"\nCREATE TRIGGER {TRIGGER_NAME}\nAFTER INSERT ON \"user\"\nFOR EACH ROW\nEXECUTE FUNCTION {FUNCTION_NAME}();\n\"\"\"\n\n\ndef upgrade() -> None:\n \"\"\"\n Create a placeholder user and assign all NULL user_id records to it.\n Install a trigger that migrates data to the first real user and self-destructs.\n Only runs if AUTH_TYPE is currently disabled/none.\n\n Skipped in multi-tenant mode - each tenant starts fresh with no legacy data.\n \"\"\"\n # Skip in multi-tenant mode - this migration handles single-tenant\n # AUTH_TYPE=disabled -> auth transitions only\n if MULTI_TENANT:\n return\n\n # Only run if AUTH_TYPE is currently disabled/none\n # If they've already switched to auth-enabled, NULL data is stale anyway\n auth_type = (os.environ.get(\"AUTH_TYPE\") or \"\").lower()\n if auth_type not in (\"disabled\", \"none\", \"\"):\n print(f\"AUTH_TYPE is '{auth_type}', not disabled. Skipping migration.\")\n return\n\n connection = op.get_bind()\n\n # Check if there are any NULL user_id records that need migration\n tables_to_check = [\n \"chat_session\",\n \"credential\",\n \"document_set\",\n \"persona\",\n \"tool\",\n \"notification\",\n \"inputprompt\",\n ]\n\n has_null_records = False\n for table in tables_to_check:\n try:\n result = connection.execute(\n sa.text(f'SELECT 1 FROM \"{table}\" WHERE user_id IS NULL LIMIT 1')\n )\n if result.fetchone():\n has_null_records = True\n break\n except Exception:\n # Table might not exist\n pass\n\n if not has_null_records:\n return\n\n # Create the placeholder user\n connection.execute(\n sa.text(\n \"\"\"\n INSERT INTO \"user\" (id, email, hashed_password, is_active, is_superuser, is_verified, role)\n VALUES (:id, :email, :hashed_password, :is_active, :is_superuser, :is_verified, :role)\n \"\"\"\n ),\n {\n \"id\": NO_AUTH_PLACEHOLDER_USER_UUID,\n \"email\": NO_AUTH_PLACEHOLDER_USER_EMAIL,\n \"hashed_password\": \"\", # Empty password - user cannot log in\n \"is_active\": False, # Inactive - user cannot log in\n \"is_superuser\": False,\n \"is_verified\": False,\n \"role\": \"BASIC\",\n },\n )\n\n # Assign NULL user_id records to the placeholder user\n for table in tables_to_check:\n try:\n # Base condition for all tables\n condition = \"user_id IS NULL\"\n # Exclude public credential (id=0) which must remain user_id=NULL\n if table == \"credential\":\n condition += \" AND id != 0\"\n # Exclude builtin tools (in_code_tool_id IS NOT NULL) which must remain user_id=NULL\n elif table == \"tool\":\n condition += \" AND in_code_tool_id IS NULL\"\n # Exclude builtin personas which must remain user_id=NULL\n elif table == \"persona\":\n condition += \" AND builtin_persona = FALSE\"\n # Exclude system/public input prompts which must remain user_id=NULL\n elif table == \"inputprompt\":\n condition += \" AND is_public = FALSE\"\n result = connection.execute(\n sa.text(\n f\"\"\"\n UPDATE \"{table}\"\n SET user_id = :user_id\n WHERE {condition}\n \"\"\"\n ),\n {\"user_id\": NO_AUTH_PLACEHOLDER_USER_UUID},\n )\n if result.rowcount > 0:\n print(f\"Updated {result.rowcount} rows in {table}\")\n except Exception as e:\n print(f\"Skipping {table}: {e}\")\n\n # Install the trigger function and trigger for automatic migration on first user registration\n connection.execute(sa.text(MIGRATE_NO_AUTH_TRIGGER_FUNCTION))\n connection.execute(sa.text(MIGRATE_NO_AUTH_TRIGGER))\n print(\"Installed trigger for automatic data migration on first user registration\")\n\n\ndef downgrade() -> None:\n \"\"\"\n Drop trigger and function, set placeholder user's records back to NULL,\n and delete the placeholder user.\n \"\"\"\n # Skip in multi-tenant mode for consistency with upgrade\n if MULTI_TENANT:\n return\n\n connection = op.get_bind()\n\n # Drop trigger and function if they exist (they may have already self-destructed)\n connection.execute(sa.text(f'DROP TRIGGER IF EXISTS {TRIGGER_NAME} ON \"user\"'))\n connection.execute(sa.text(f\"DROP FUNCTION IF EXISTS {FUNCTION_NAME}()\"))\n\n tables_to_update = [\n \"chat_session\",\n \"credential\",\n \"document_set\",\n \"persona\",\n \"tool\",\n \"notification\",\n \"inputprompt\",\n ]\n\n # Set records back to NULL\n for table in tables_to_update:\n try:\n connection.execute(\n sa.text(\n f\"\"\"\n UPDATE \"{table}\"\n SET user_id = NULL\n WHERE user_id = :user_id\n \"\"\"\n ),\n {\"user_id\": NO_AUTH_PLACEHOLDER_USER_UUID},\n )\n except Exception:\n pass\n\n # Delete the placeholder user\n connection.execute(\n sa.text('DELETE FROM \"user\" WHERE id = :user_id'),\n {\"user_id\": NO_AUTH_PLACEHOLDER_USER_UUID},\n )\n" }, { "path": "backend/alembic/versions/f7e58d357687_add_has_web_column_to_user.py", "content": "\"\"\"add has_web_login column to user\n\nRevision ID: f7e58d357687\nRevises: ba98eba0f66a\nCreate Date: 2024-09-07 20:20:54.522620\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"f7e58d357687\"\ndown_revision = \"ba98eba0f66a\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\"has_web_login\", sa.Boolean(), nullable=False, server_default=\"true\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"has_web_login\")\n" }, { "path": "backend/alembic/versions/f8a9b2c3d4e5_add_research_answer_purpose_to_chat_message.py", "content": "\"\"\"add research_answer_purpose to chat_message\n\nRevision ID: f8a9b2c3d4e5\nRevises: 5ae8240accb3\nCreate Date: 2025-01-27 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f8a9b2c3d4e5\"\ndown_revision = \"5ae8240accb3\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add research_answer_purpose column to chat_message table\n op.add_column(\n \"chat_message\",\n sa.Column(\"research_answer_purpose\", sa.String(), nullable=True),\n )\n\n\ndef downgrade() -> None:\n # Remove research_answer_purpose column from chat_message table\n op.drop_column(\"chat_message\", \"research_answer_purpose\")\n" }, { "path": "backend/alembic/versions/f9b8c7d6e5a4_update_parent_question_id_foreign_key_to_research_agent_iteration.py", "content": "\"\"\"remove foreign key constraints from research_agent_iteration_sub_step\n\nRevision ID: f9b8c7d6e5a4\nRevises: bd7c3bf8beba\nCreate Date: 2025-01-27 12:00:00.000000\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"f9b8c7d6e5a4\"\ndown_revision = \"bd7c3bf8beba\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Drop the existing foreign key constraint for parent_question_id\n op.drop_constraint(\n \"research_agent_iteration_sub_step_parent_question_id_fkey\",\n \"research_agent_iteration_sub_step\",\n type_=\"foreignkey\",\n )\n\n # Drop the parent_question_id column entirely\n op.drop_column(\"research_agent_iteration_sub_step\", \"parent_question_id\")\n\n # Drop the foreign key constraint for primary_question_id to chat_message.id\n # (keep the column as it's needed for the composite foreign key)\n op.drop_constraint(\n \"research_agent_iteration_sub_step_primary_question_id_fkey\",\n \"research_agent_iteration_sub_step\",\n type_=\"foreignkey\",\n )\n\n\ndef downgrade() -> None:\n # Restore the foreign key constraint for primary_question_id to chat_message.id\n op.create_foreign_key(\n \"research_agent_iteration_sub_step_primary_question_id_fkey\",\n \"research_agent_iteration_sub_step\",\n \"chat_message\",\n [\"primary_question_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n\n # Add back the parent_question_id column\n op.add_column(\n \"research_agent_iteration_sub_step\",\n sa.Column(\n \"parent_question_id\",\n sa.Integer(),\n nullable=True,\n ),\n )\n\n # Restore the foreign key constraint pointing to research_agent_iteration_sub_step.id\n op.create_foreign_key(\n \"research_agent_iteration_sub_step_parent_question_id_fkey\",\n \"research_agent_iteration_sub_step\",\n \"research_agent_iteration_sub_step\",\n [\"parent_question_id\"],\n [\"id\"],\n ondelete=\"CASCADE\",\n )\n" }, { "path": "backend/alembic/versions/fad14119fb92_delete_tags_with_wrong_enum.py", "content": "\"\"\"Delete Tags with wrong Enum\n\nRevision ID: fad14119fb92\nRevises: 72bdc9929a46\nCreate Date: 2024-04-25 17:05:09.695703\n\n\"\"\"\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"fad14119fb92\"\ndown_revision = \"72bdc9929a46\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n # Some documents may lose their tags but this is the only way as the enum\n # mapping may have changed since tag switched to string (it will be reindexed anyway)\n op.execute(\n \"\"\"\n DELETE FROM document__tag\n WHERE tag_id IN (\n SELECT id FROM tag\n WHERE source ~ '^[0-9]+$'\n )\n \"\"\"\n )\n\n op.execute(\n \"\"\"\n DELETE FROM tag\n WHERE source ~ '^[0-9]+$'\n \"\"\"\n )\n\n\ndef downgrade() -> None:\n pass\n" }, { "path": "backend/alembic/versions/fb80bdd256de_add_chat_background_to_user.py", "content": "\"\"\"add chat_background to user\n\nRevision ID: fb80bdd256de\nRevises: 8b5ce697290e\nCreate Date: 2026-01-16 16:15:59.222617\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"fb80bdd256de\"\ndown_revision = \"8b5ce697290e\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"user\",\n sa.Column(\n \"chat_background\",\n sa.String(),\n nullable=True,\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user\", \"chat_background\")\n" }, { "path": "backend/alembic/versions/fcd135795f21_add_slack_bot_display_type.py", "content": "\"\"\"Add slack bot display type\n\nRevision ID: fcd135795f21\nRevises: 0a2b51deb0b8\nCreate Date: 2024-03-04 17:03:27.116284\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"fcd135795f21\"\ndown_revision = \"0a2b51deb0b8\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"slack_bot_config\",\n sa.Column(\n \"response_type\",\n sa.Enum(\n \"QUOTES\",\n \"CITATIONS\",\n name=\"slackbotresponsetype\",\n native_enum=False,\n ),\n nullable=True,\n ),\n )\n op.execute(\n \"UPDATE slack_bot_config SET response_type = 'QUOTES' WHERE response_type IS NULL\"\n )\n op.alter_column(\"slack_bot_config\", \"response_type\", nullable=False)\n\n\ndef downgrade() -> None:\n op.drop_column(\"slack_bot_config\", \"response_type\")\n" }, { "path": "backend/alembic/versions/febe9eaa0644_add_document_set_persona_relationship_.py", "content": "\"\"\"Add document_set / persona relationship table\n\nRevision ID: febe9eaa0644\nRevises: 57b53544726e\nCreate Date: 2023-09-24 13:06:24.018610\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"febe9eaa0644\"\ndown_revision = \"57b53544726e\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"persona__document_set\",\n sa.Column(\"persona_id\", sa.Integer(), nullable=False),\n sa.Column(\"document_set_id\", sa.Integer(), nullable=False),\n sa.ForeignKeyConstraint(\n [\"document_set_id\"],\n [\"document_set.id\"],\n ),\n sa.ForeignKeyConstraint(\n [\"persona_id\"],\n [\"persona.id\"],\n ),\n sa.PrimaryKeyConstraint(\"persona_id\", \"document_set_id\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"persona__document_set\")\n" }, { "path": "backend/alembic/versions/fec3db967bf7_add_time_updated_to_usergroup_and_.py", "content": "\"\"\"Add time_updated to UserGroup and DocumentSet\n\nRevision ID: fec3db967bf7\nRevises: 97dbb53fa8c8\nCreate Date: 2025-01-12 15:49:02.289100\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n# revision identifiers, used by Alembic.\nrevision = \"fec3db967bf7\"\ndown_revision = \"97dbb53fa8c8\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"document_set\",\n sa.Column(\n \"time_last_modified_by_user\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.func.now(),\n ),\n )\n op.add_column(\n \"user_group\",\n sa.Column(\n \"time_last_modified_by_user\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.func.now(),\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"user_group\", \"time_last_modified_by_user\")\n op.drop_column(\"document_set\", \"time_last_modified_by_user\")\n" }, { "path": "backend/alembic/versions/feead2911109_add_opensearch_tenant_migration_columns.py", "content": "\"\"\"add_opensearch_tenant_migration_columns\n\nRevision ID: feead2911109\nRevises: d56ffa94ca32\nCreate Date: 2026-02-10 17:46:34.029937\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\n\n\n# revision identifiers, used by Alembic.\nrevision = \"feead2911109\"\ndown_revision = \"175ea04c7087\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\"vespa_visit_continuation_token\", sa.Text(), nullable=True),\n )\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"total_chunks_migrated\",\n sa.Integer(),\n nullable=False,\n server_default=\"0\",\n ),\n )\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"created_at\",\n sa.DateTime(timezone=True),\n nullable=False,\n server_default=sa.func.now(),\n ),\n )\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"migration_completed_at\",\n sa.DateTime(timezone=True),\n nullable=True,\n ),\n )\n op.add_column(\n \"opensearch_tenant_migration_record\",\n sa.Column(\n \"enable_opensearch_retrieval\",\n sa.Boolean(),\n nullable=False,\n server_default=\"false\",\n ),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"opensearch_tenant_migration_record\", \"enable_opensearch_retrieval\")\n op.drop_column(\"opensearch_tenant_migration_record\", \"migration_completed_at\")\n op.drop_column(\"opensearch_tenant_migration_record\", \"created_at\")\n op.drop_column(\"opensearch_tenant_migration_record\", \"total_chunks_migrated\")\n op.drop_column(\n \"opensearch_tenant_migration_record\", \"vespa_visit_continuation_token\"\n )\n" }, { "path": "backend/alembic/versions/ffc707a226b4_basic_document_metadata.py", "content": "\"\"\"Basic Document Metadata\n\nRevision ID: ffc707a226b4\nRevises: 30c1d5744104\nCreate Date: 2023-10-18 16:52:25.967592\n\n\"\"\"\n\nfrom alembic import op\nimport sqlalchemy as sa\nfrom sqlalchemy.dialects import postgresql\n\n# revision identifiers, used by Alembic.\nrevision = \"ffc707a226b4\"\ndown_revision = \"30c1d5744104\"\nbranch_labels: None = None\ndepends_on: None = None\n\n\ndef upgrade() -> None:\n op.add_column(\n \"document\",\n sa.Column(\"doc_updated_at\", sa.DateTime(timezone=True), nullable=True),\n )\n op.add_column(\n \"document\",\n sa.Column(\"primary_owners\", postgresql.ARRAY(sa.String()), nullable=True),\n )\n op.add_column(\n \"document\",\n sa.Column(\"secondary_owners\", postgresql.ARRAY(sa.String()), nullable=True),\n )\n\n\ndef downgrade() -> None:\n op.drop_column(\"document\", \"secondary_owners\")\n op.drop_column(\"document\", \"primary_owners\")\n op.drop_column(\"document\", \"doc_updated_at\")\n" }, { "path": "backend/alembic.ini", "content": "# A generic, single database configuration.\n\n[DEFAULT]\n# path to migration scripts\nscript_location = alembic\n\n# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s\n# Uncomment the line below if you want the files to be prepended with date and time\n# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s\n\n# sys.path path, will be prepended to sys.path if present.\n# defaults to the current working directory.\nprepend_sys_path = .\n\n# timezone to use when rendering the date within the migration file\n# as well as the filename.\n# If specified, requires the python-dateutil library that can be\n# installed by adding `alembic[tz]` to the pip requirements\n# string value is passed to dateutil.tz.gettz()\n# leave blank for localtime\n# timezone =\n\n# max length of characters to apply to the\n# \"slug\" field\n# truncate_slug_length = 40\n\n# set to 'true' to run the environment during\n# the 'revision' command, regardless of autogenerate\n# revision_environment = false\n\n# set to 'true' to allow .pyc and .pyo files without\n# a source .py file to be detected as revisions in the\n# versions/ directory\n# sourceless = false\n\n# version location specification; This defaults\n# to alembic/versions. When using multiple version\n# directories, initial revisions must be specified with --version-path.\n# The path separator used here should be the separator specified by \"version_path_separator\" below.\n# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions\n\n# version path separator; As mentioned above, this is the character used to split\n# version_locations. The default within new alembic.ini files is \"os\", which uses os.pathsep.\n# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.\n# Valid values for version_path_separator are:\n#\n# version_path_separator = :\n# version_path_separator = ;\n# version_path_separator = space\nversion_path_separator = os \n# Use os.pathsep. Default configuration used for new projects.\n\n# set to 'true' to search source files recursively\n# in each \"version_locations\" directory\n# new in Alembic version 1.10\n# recursive_version_locations = false\n\n# the output encoding used when revision files\n# are written from script.py.mako\n# output_encoding = utf-8\n\n# sqlalchemy.url = driver://user:pass@localhost/dbname\n\n\n[post_write_hooks]\n# post_write_hooks defines scripts or Python functions that are run\n# on newly generated revision scripts. See the documentation for further\n# detail and examples\n\n# format using \"black\" - use the console_scripts runner, against the \"black\" entrypoint\nhooks = black\nblack.type = console_scripts\nblack.entrypoint = black\nblack.options = -l 79 REVISION_SCRIPT_FILENAME\n\n# Logging configuration\n[loggers]\nkeys = root,sqlalchemy,alembic\n\n[handlers]\nkeys = console\n\n[formatters]\nkeys = generic\n\n[logger_root]\nlevel = INFO\nhandlers = console\nqualname =\n\n[logger_sqlalchemy]\nlevel = WARN\nhandlers =\nqualname = sqlalchemy.engine\n\n[logger_alembic]\nlevel = INFO\nhandlers =\nqualname = alembic\n\n[handler_console]\nclass = StreamHandler\nargs = (sys.stderr,)\nlevel = NOTSET\nformatter = generic\n\n[formatter_generic]\nformat = %(levelname)-5.5s [%(name)s] %(message)s\ndatefmt = %H:%M:%S\n\n\n[alembic]\nscript_location = alembic\nversion_locations = %(script_location)s/versions\n\n[schema_private]\nscript_location = alembic_tenants\nversion_locations = %(script_location)s/versions\n" }, { "path": "backend/alembic_tenants/README.md", "content": "These files are for public table migrations when operating with multi tenancy.\n\nIf you are not a Onyx developer, you can ignore this directory entirely.\n" }, { "path": "backend/alembic_tenants/__init__.py", "content": "" }, { "path": "backend/alembic_tenants/env.py", "content": "import asyncio\nfrom logging.config import fileConfig\nfrom typing import Literal\n\nfrom sqlalchemy import pool\nfrom sqlalchemy.engine import Connection\nfrom sqlalchemy.ext.asyncio import create_async_engine\nfrom sqlalchemy.schema import SchemaItem\n\nfrom alembic import context\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.models import PublicBase\n\n# this is the Alembic Config object, which provides\n# access to the values within the .ini file in use.\nconfig = context.config\n\n# Interpret the config file for Python logging.\n# This line sets up loggers basically.\nif config.config_file_name is not None and config.attributes.get(\n \"configure_logger\", True\n):\n # disable_existing_loggers=False prevents breaking pytest's caplog fixture\n # See: https://pytest-alembic.readthedocs.io/en/latest/setup.html#caplog-issues\n fileConfig(config.config_file_name, disable_existing_loggers=False)\n\n# add your model's MetaData object here\n# for 'autogenerate' support\n# from myapp import mymodel\n# target_metadata = mymodel.Base.metadata\ntarget_metadata = [PublicBase.metadata]\n\n# other values from the config, defined by the needs of env.py,\n# can be acquired:\n# my_important_option = config.get_main_option(\"my_important_option\")\n# ... etc.\n\nEXCLUDE_TABLES = {\"kombu_queue\", \"kombu_message\"}\n\n\ndef include_object(\n object: SchemaItem, # noqa: ARG001\n name: str | None,\n type_: Literal[\n \"schema\",\n \"table\",\n \"column\",\n \"index\",\n \"unique_constraint\",\n \"foreign_key_constraint\",\n ],\n reflected: bool, # noqa: ARG001\n compare_to: SchemaItem | None, # noqa: ARG001\n) -> bool:\n if type_ == \"table\" and name in EXCLUDE_TABLES:\n return False\n return True\n\n\ndef run_migrations_offline() -> None:\n \"\"\"Run migrations in 'offline' mode.\n\n This configures the context with just a URL\n and not an Engine, though an Engine is acceptable\n here as well. By skipping the Engine creation\n we don't even need a DBAPI to be available.\n\n Calls to context.execute() here emit the given string to the\n script output.\n\n \"\"\"\n url = build_connection_string()\n context.configure(\n url=url,\n target_metadata=target_metadata, # type: ignore\n literal_binds=True,\n dialect_opts={\"paramstyle\": \"named\"},\n )\n\n with context.begin_transaction():\n context.run_migrations()\n\n\ndef do_run_migrations(connection: Connection) -> None:\n context.configure(\n connection=connection,\n target_metadata=target_metadata, # type: ignore[arg-type]\n include_object=include_object,\n )\n\n with context.begin_transaction():\n context.run_migrations()\n\n\nasync def run_async_migrations() -> None:\n \"\"\"In this scenario we need to create an Engine\n and associate a connection with the context.\n\n \"\"\"\n\n connectable = create_async_engine(\n build_connection_string(),\n poolclass=pool.NullPool,\n )\n\n async with connectable.connect() as connection:\n await connection.run_sync(do_run_migrations)\n\n await connectable.dispose()\n\n\ndef run_migrations_online() -> None:\n \"\"\"Run migrations in 'online' mode.\n\n Supports pytest-alembic by checking for a pre-configured connection\n in context.config.attributes[\"connection\"]. If present, uses that\n connection/engine directly instead of creating a new async engine.\n \"\"\"\n # Check if pytest-alembic is providing a connection/engine\n connectable = context.config.attributes.get(\"connection\", None)\n\n if connectable is not None:\n # pytest-alembic is providing an engine - use it directly\n with connectable.connect() as connection:\n do_run_migrations(connection)\n # Commit to ensure changes are visible to next migration\n connection.commit()\n else:\n # Normal operation - use async migrations\n asyncio.run(run_async_migrations())\n\n\nif context.is_offline_mode():\n run_migrations_offline()\nelse:\n run_migrations_online()\n" }, { "path": "backend/alembic_tenants/script.py.mako", "content": "\"\"\"${message}\n\nRevision ID: ${up_revision}\nRevises: ${down_revision | comma,n}\nCreate Date: ${create_date}\n\n\"\"\"\nfrom alembic import op\nimport sqlalchemy as sa\n${imports if imports else \"\"}\n\n# revision identifiers, used by Alembic.\nrevision = ${repr(up_revision)}\ndown_revision = ${repr(down_revision)}\nbranch_labels = ${repr(branch_labels)}\ndepends_on = ${repr(depends_on)}\n\n\ndef upgrade() -> None:\n ${upgrades if upgrades else \"pass\"}\n\n\ndef downgrade() -> None:\n ${downgrades if downgrades else \"pass\"}\n" }, { "path": "backend/alembic_tenants/versions/14a83a331951_create_usertenantmapping_table.py", "content": "import sqlalchemy as sa\n\nfrom alembic import op\n\n# revision identifiers, used by Alembic.\nrevision = \"14a83a331951\"\ndown_revision = None\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"user_tenant_mapping\",\n sa.Column(\"email\", sa.String(), nullable=False),\n sa.Column(\"tenant_id\", sa.String(), nullable=False),\n sa.UniqueConstraint(\"email\", \"tenant_id\", name=\"uq_user_tenant\"),\n sa.UniqueConstraint(\"email\", name=\"uq_email\"),\n schema=\"public\",\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"user_tenant_mapping\", schema=\"public\")\n" }, { "path": "backend/alembic_tenants/versions/34e3630c7f32_lowercase_multi_tenant_user_auth.py", "content": "\"\"\"lowercase multi-tenant user auth\n\nRevision ID: 34e3630c7f32\nRevises: a4f6ee863c47\nCreate Date: 2025-02-26 15:03:01.211894\n\n\"\"\"\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"34e3630c7f32\"\ndown_revision = \"a4f6ee863c47\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # 1) Convert all existing rows to lowercase\n op.execute(\n \"\"\"\n UPDATE user_tenant_mapping\n SET email = LOWER(email)\n \"\"\"\n )\n # 2) Add a check constraint so that emails cannot be written in uppercase\n op.create_check_constraint(\n \"ensure_lowercase_email\",\n \"user_tenant_mapping\",\n \"email = LOWER(email)\",\n schema=\"public\",\n )\n\n\ndef downgrade() -> None:\n # Drop the check constraint\n op.drop_constraint(\n \"ensure_lowercase_email\",\n \"user_tenant_mapping\",\n schema=\"public\",\n type_=\"check\",\n )\n" }, { "path": "backend/alembic_tenants/versions/3b45e0018bf1_add_new_available_tenant_table.py", "content": "\"\"\"add new available tenant table\n\nRevision ID: 3b45e0018bf1\nRevises: ac842f85f932\nCreate Date: 2025-03-06 09:55:18.229910\n\n\"\"\"\n\nimport sqlalchemy as sa\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"3b45e0018bf1\"\ndown_revision = \"ac842f85f932\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Create new_available_tenant table\n op.create_table(\n \"available_tenant\",\n sa.Column(\"tenant_id\", sa.String(), nullable=False),\n sa.Column(\"alembic_version\", sa.String(), nullable=False),\n sa.Column(\"date_created\", sa.DateTime(), nullable=False),\n sa.PrimaryKeyConstraint(\"tenant_id\"),\n )\n\n\ndef downgrade() -> None:\n # Drop new_available_tenant table\n op.drop_table(\"available_tenant\")\n" }, { "path": "backend/alembic_tenants/versions/3b9f09038764_add_read_only_kg_user.py", "content": "\"\"\"add_db_readonly_user\n\nRevision ID: 3b9f09038764\nRevises: 3b45e0018bf1\nCreate Date: 2025-05-11 11:05:11.436977\n\n\"\"\"\n\nfrom sqlalchemy import text\n\nfrom alembic import op\nfrom onyx.configs.app_configs import DB_READONLY_PASSWORD\nfrom onyx.configs.app_configs import DB_READONLY_USER\n\n\n# revision identifiers, used by Alembic.\nrevision = \"3b9f09038764\"\ndown_revision = \"3b45e0018bf1\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Enable pg_trgm extension if not already enabled\n op.execute(\"CREATE EXTENSION IF NOT EXISTS pg_trgm\")\n\n # Create the read-only db user if it does not already exist.\n if not (DB_READONLY_USER and DB_READONLY_PASSWORD):\n raise Exception(\"DB_READONLY_USER or DB_READONLY_PASSWORD is not set\")\n\n op.execute(\n text(\n f\"\"\"\n DO $$\n BEGIN\n -- Check if the read-only user already exists\n IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN\n -- Create the read-only user with the specified password\n EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}');\n -- First revoke all privileges to ensure a clean slate\n EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');\n -- Grant only the CONNECT privilege to allow the user to connect to the database\n -- but not perform any operations without additional specific grants\n EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}');\n END IF;\n END\n $$;\n \"\"\"\n )\n )\n\n\ndef downgrade() -> None:\n op.execute(\n text(\n f\"\"\"\n DO $$\n BEGIN\n IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN\n -- First revoke all privileges from the database\n EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}');\n -- Then revoke all privileges from the public schema\n EXECUTE format('REVOKE ALL ON SCHEMA public FROM %I', '{DB_READONLY_USER}');\n -- Then drop the user\n EXECUTE format('DROP USER %I', '{DB_READONLY_USER}');\n END IF;\n END\n $$;\n \"\"\"\n )\n )\n op.execute(text(\"DROP EXTENSION IF EXISTS pg_trgm\"))\n" }, { "path": "backend/alembic_tenants/versions/a4f6ee863c47_mapping_for_anonymous_user_path.py", "content": "\"\"\"mapping for anonymous user path\n\nRevision ID: a4f6ee863c47\nRevises: 14a83a331951\nCreate Date: 2025-01-04 14:16:58.697451\n\n\"\"\"\n\nimport sqlalchemy as sa\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"a4f6ee863c47\"\ndown_revision = \"14a83a331951\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n op.create_table(\n \"tenant_anonymous_user_path\",\n sa.Column(\"tenant_id\", sa.String(), primary_key=True, nullable=False),\n sa.Column(\"anonymous_user_path\", sa.String(), nullable=False),\n sa.PrimaryKeyConstraint(\"tenant_id\"),\n sa.UniqueConstraint(\"anonymous_user_path\"),\n )\n\n\ndef downgrade() -> None:\n op.drop_table(\"tenant_anonymous_user_path\")\n" }, { "path": "backend/alembic_tenants/versions/ac842f85f932_new_column_user_tenant_mapping.py", "content": "\"\"\"new column user tenant mapping\n\nRevision ID: ac842f85f932\nRevises: 34e3630c7f32\nCreate Date: 2025-03-03 13:30:14.802874\n\n\"\"\"\n\nimport sqlalchemy as sa\n\nfrom alembic import op\n\n\n# revision identifiers, used by Alembic.\nrevision = \"ac842f85f932\"\ndown_revision = \"34e3630c7f32\"\nbranch_labels = None\ndepends_on = None\n\n\ndef upgrade() -> None:\n # Add active column with default value of True\n op.add_column(\n \"user_tenant_mapping\",\n sa.Column(\n \"active\",\n sa.Boolean(),\n nullable=False,\n server_default=\"true\",\n ),\n schema=\"public\",\n )\n\n op.drop_constraint(\"uq_email\", \"user_tenant_mapping\", schema=\"public\")\n\n # Create a unique index for active=true records\n # This ensures a user can only be active in one tenant at a time\n op.execute(\n \"CREATE UNIQUE INDEX uq_user_active_email_idx ON public.user_tenant_mapping (email) WHERE active = true\"\n )\n\n\ndef downgrade() -> None:\n # Drop the unique index for active=true records\n op.execute(\"DROP INDEX IF EXISTS uq_user_active_email_idx\")\n\n op.create_unique_constraint(\n \"uq_email\", \"user_tenant_mapping\", [\"email\"], schema=\"public\"\n )\n\n # Remove the active column\n op.drop_column(\"user_tenant_mapping\", \"active\", schema=\"public\")\n" }, { "path": "backend/assets/.gitignore", "content": "*\n!.gitignore\n" }, { "path": "backend/ee/LICENSE", "content": "The Onyx Enterprise License (the \"Enterprise License\")\nCopyright (c) 2023-present DanswerAI, Inc.\n\nWith regard to the Onyx Software:\n\nThis software and associated documentation files (the \"Software\") may only be\nused in production, if you (and any entity that you represent) have agreed to,\nand are in compliance with, the Onyx Subscription Terms of Service, available\nat https://www.onyx.app/legal/self-host (the \"Enterprise Terms\"), or other\nagreement governing the use of the Software, as agreed by you and DanswerAI,\nand otherwise have a valid Onyx Enterprise License for the\ncorrect number of user seats. Subject to the foregoing sentence, you are free to\nmodify this Software and publish patches to the Software. You agree that DanswerAI\nand/or its licensors (as applicable) retain all right, title and interest in and\nto all such modifications and/or patches, and all such modifications and/or\npatches may only be used, copied, modified, displayed, distributed, or otherwise\nexploited with a valid Onyx Enterprise License for the correct\nnumber of user seats. Notwithstanding the foregoing, you may copy and modify\nthe Software for development and testing purposes, without requiring a\nsubscription. You agree that DanswerAI and/or its licensors (as applicable) retain\nall right, title and interest in and to all such modifications. You are not\ngranted any other rights beyond what is expressly stated herein. Subject to the\nforegoing, it is forbidden to copy, merge, publish, distribute, sublicense,\nand/or sell the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n\nFor all third party components incorporated into the Onyx Software, those\ncomponents are licensed under the original license provided by the owner of the\napplicable component.\n" }, { "path": "backend/ee/__init__.py", "content": "" }, { "path": "backend/ee/onyx/__init__.py", "content": "" }, { "path": "backend/ee/onyx/access/access.py", "content": "from sqlalchemy.orm import Session\n\nfrom ee.onyx.db.external_perm import fetch_external_groups_for_user\nfrom ee.onyx.db.external_perm import fetch_public_external_group_ids\nfrom ee.onyx.db.user_group import fetch_user_groups_for_documents\nfrom ee.onyx.db.user_group import fetch_user_groups_for_user\nfrom ee.onyx.external_permissions.sync_params import get_source_perm_sync_config\nfrom onyx.access.access import (\n _get_access_for_documents as get_access_for_documents_without_groups,\n)\nfrom onyx.access.access import _get_acl_for_user as get_acl_for_user_without_groups\nfrom onyx.access.access import collect_user_file_access\nfrom onyx.access.models import DocumentAccess\nfrom onyx.access.utils import prefix_external_group\nfrom onyx.access.utils import prefix_user_group\nfrom onyx.db.document import get_document_sources\nfrom onyx.db.document import get_documents_by_ids\nfrom onyx.db.models import User\nfrom onyx.db.models import UserFile\nfrom onyx.db.user_file import fetch_user_files_with_access_relationships\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _get_access_for_document(\n document_id: str,\n db_session: Session,\n) -> DocumentAccess:\n id_to_access = _get_access_for_documents([document_id], db_session)\n if len(id_to_access) == 0:\n return DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n )\n\n return next(iter(id_to_access.values()))\n\n\ndef _get_access_for_documents(\n document_ids: list[str],\n db_session: Session,\n) -> dict[str, DocumentAccess]:\n non_ee_access_dict = get_access_for_documents_without_groups(\n document_ids=document_ids,\n db_session=db_session,\n )\n user_group_info: dict[str, list[str]] = {\n document_id: group_names\n for document_id, group_names in fetch_user_groups_for_documents(\n db_session=db_session,\n document_ids=document_ids,\n )\n }\n documents = get_documents_by_ids(\n db_session=db_session,\n document_ids=document_ids,\n )\n doc_id_map = {doc.id: doc for doc in documents}\n\n # Get all sources in one batch\n doc_id_to_source_map = get_document_sources(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n all_public_ext_u_group_ids = set(fetch_public_external_group_ids(db_session))\n\n access_map = {}\n for document_id, non_ee_access in non_ee_access_dict.items():\n document = doc_id_map[document_id]\n source = doc_id_to_source_map.get(document_id)\n if source is None:\n logger.error(f\"Document {document_id} has no source\")\n continue\n\n perm_sync_config = get_source_perm_sync_config(source)\n is_only_censored = (\n perm_sync_config\n and perm_sync_config.censoring_config is not None\n and perm_sync_config.doc_sync_config is None\n )\n\n ext_u_emails = (\n set(document.external_user_emails)\n if document.external_user_emails\n else set()\n )\n\n ext_u_groups = (\n set(document.external_user_group_ids)\n if document.external_user_group_ids\n else set()\n )\n\n # If the document is determined to be \"public\" externally (through a SYNC connector)\n # then it's given the same access level as if it were marked public within Onyx\n # If its censored, then it's public anywhere during the search and then permissions are\n # applied after the search\n is_public_anywhere = (\n document.is_public\n or non_ee_access.is_public\n or is_only_censored\n or any(u_group in all_public_ext_u_group_ids for u_group in ext_u_groups)\n )\n\n # To avoid collisions of group namings between connectors, they need to be prefixed\n access_map[document_id] = DocumentAccess.build(\n user_emails=list(non_ee_access.user_emails),\n user_groups=user_group_info.get(document_id, []),\n is_public=is_public_anywhere,\n external_user_emails=list(ext_u_emails),\n external_user_group_ids=list(ext_u_groups),\n )\n return access_map\n\n\ndef _collect_user_file_group_names(user_file: UserFile) -> set[str]:\n \"\"\"Extract user-group names from the already-loaded Persona.groups\n relationships on a UserFile (skipping deleted personas).\"\"\"\n groups: set[str] = set()\n for persona in user_file.assistants:\n if persona.deleted:\n continue\n for group in persona.groups:\n groups.add(group.name)\n return groups\n\n\ndef get_access_for_user_files_impl(\n user_file_ids: list[str],\n db_session: Session,\n) -> dict[str, DocumentAccess]:\n \"\"\"EE version: extends the MIT user file ACL with user group names\n from personas shared via user groups.\n\n Uses a single DB query (via fetch_user_files_with_access_relationships)\n that eagerly loads both the MIT-needed and EE-needed relationships.\n\n NOTE: is imported in onyx.access.access by `fetch_versioned_implementation`\n DO NOT REMOVE.\"\"\"\n user_files = fetch_user_files_with_access_relationships(\n user_file_ids, db_session, eager_load_groups=True\n )\n return build_access_for_user_files_impl(user_files)\n\n\ndef build_access_for_user_files_impl(\n user_files: list[UserFile],\n) -> dict[str, DocumentAccess]:\n \"\"\"EE version: works on pre-loaded UserFile objects.\n Expects Persona.groups to be eagerly loaded.\n\n NOTE: is imported in onyx.access.access by `fetch_versioned_implementation`\n DO NOT REMOVE.\"\"\"\n result: dict[str, DocumentAccess] = {}\n for user_file in user_files:\n if user_file.user is None:\n result[str(user_file.id)] = DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n is_public=True,\n external_user_emails=[],\n external_user_group_ids=[],\n )\n continue\n\n emails, is_public = collect_user_file_access(user_file)\n group_names = _collect_user_file_group_names(user_file)\n result[str(user_file.id)] = DocumentAccess.build(\n user_emails=list(emails),\n user_groups=list(group_names),\n is_public=is_public,\n external_user_emails=[],\n external_user_group_ids=[],\n )\n return result\n\n\ndef _get_acl_for_user(user: User, db_session: Session) -> set[str]:\n \"\"\"Returns a list of ACL entries that the user has access to. This is meant to be\n used downstream to filter out documents that the user does not have access to. The\n user should have access to a document if at least one entry in the document's ACL\n matches one entry in the returned set.\n\n NOTE: is imported in onyx.access.access by `fetch_versioned_implementation`\n DO NOT REMOVE.\"\"\"\n is_anonymous = user.is_anonymous\n db_user_groups = (\n [] if is_anonymous else fetch_user_groups_for_user(db_session, user.id)\n )\n prefixed_user_groups = [\n prefix_user_group(db_user_group.name) for db_user_group in db_user_groups\n ]\n\n db_external_groups = (\n [] if is_anonymous else fetch_external_groups_for_user(db_session, user.id)\n )\n prefixed_external_groups = [\n prefix_external_group(db_external_group.external_user_group_id)\n for db_external_group in db_external_groups\n ]\n\n user_acl = set(prefixed_user_groups + prefixed_external_groups)\n user_acl.update(get_acl_for_user_without_groups(user, db_session))\n\n return user_acl\n" }, { "path": "backend/ee/onyx/access/hierarchy_access.py", "content": "from sqlalchemy.orm import Session\n\nfrom ee.onyx.db.external_perm import fetch_external_groups_for_user\nfrom onyx.db.models import User\n\n\ndef _get_user_external_group_ids(db_session: Session, user: User) -> list[str]:\n if not user:\n return []\n external_groups = fetch_external_groups_for_user(db_session, user.id)\n return [external_group.external_user_group_id for external_group in external_groups]\n" }, { "path": "backend/ee/onyx/auth/__init__.py", "content": "" }, { "path": "backend/ee/onyx/auth/users.py", "content": "import os\nfrom datetime import datetime\n\nimport jwt\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Request\nfrom fastapi import status\n\nfrom ee.onyx.configs.app_configs import SUPER_CLOUD_API_KEY\nfrom ee.onyx.configs.app_configs import SUPER_USERS\nfrom ee.onyx.server.seeding import get_seed_config\nfrom onyx.auth.users import current_admin_user\nfrom onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.app_configs import USER_AUTH_SECRET\nfrom onyx.db.models import User\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef verify_auth_setting() -> None:\n # All the Auth flows are valid for EE version, but warn about deprecated 'disabled'\n raw_auth_type = (os.environ.get(\"AUTH_TYPE\") or \"\").lower()\n if raw_auth_type == \"disabled\":\n logger.warning(\n \"AUTH_TYPE='disabled' is no longer supported. Using 'basic' instead. Please update your configuration.\"\n )\n logger.notice(f\"Using Auth Type: {AUTH_TYPE.value}\")\n\n\ndef get_default_admin_user_emails_() -> list[str]:\n seed_config = get_seed_config()\n if seed_config and seed_config.admin_user_emails:\n return seed_config.admin_user_emails\n return []\n\n\nasync def current_cloud_superuser(\n request: Request,\n user: User = Depends(current_admin_user),\n) -> User:\n api_key = request.headers.get(\"Authorization\", \"\").replace(\"Bearer \", \"\")\n if api_key != SUPER_CLOUD_API_KEY:\n raise HTTPException(status_code=401, detail=\"Invalid API key\")\n\n if user and user.email not in SUPER_USERS:\n raise HTTPException(\n status_code=status.HTTP_403_FORBIDDEN,\n detail=\"Access denied. User must be a cloud superuser to perform this action.\",\n )\n return user\n\n\ndef generate_anonymous_user_jwt_token(tenant_id: str) -> str:\n payload = {\n \"tenant_id\": tenant_id,\n # Token does not expire\n \"iat\": datetime.utcnow(), # Issued at time\n }\n\n return jwt.encode(payload, USER_AUTH_SECRET, algorithm=\"HS256\")\n\n\ndef decode_anonymous_user_jwt_token(token: str) -> dict:\n return jwt.decode(token, USER_AUTH_SECRET, algorithms=[\"HS256\"])\n" }, { "path": "backend/ee/onyx/background/celery/apps/heavy.py", "content": "from onyx.background.celery.apps import app_base\nfrom onyx.background.celery.apps.heavy import celery_app\n\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"ee.onyx.background.celery.tasks.doc_permission_syncing\",\n \"ee.onyx.background.celery.tasks.external_group_syncing\",\n \"ee.onyx.background.celery.tasks.cleanup\",\n \"ee.onyx.background.celery.tasks.query_history\",\n ]\n )\n)\n" }, { "path": "backend/ee/onyx/background/celery/apps/light.py", "content": "from onyx.background.celery.apps import app_base\nfrom onyx.background.celery.apps.light import celery_app\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"ee.onyx.background.celery.tasks.doc_permission_syncing\",\n \"ee.onyx.background.celery.tasks.external_group_syncing\",\n ]\n )\n)\n" }, { "path": "backend/ee/onyx/background/celery/apps/monitoring.py", "content": "from onyx.background.celery.apps import app_base\nfrom onyx.background.celery.apps.monitoring import celery_app\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"ee.onyx.background.celery.tasks.tenant_provisioning\",\n ]\n )\n)\n" }, { "path": "backend/ee/onyx/background/celery/apps/primary.py", "content": "from onyx.background.celery.apps import app_base\nfrom onyx.background.celery.apps.primary import celery_app\n\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"ee.onyx.background.celery.tasks.hooks\",\n \"ee.onyx.background.celery.tasks.doc_permission_syncing\",\n \"ee.onyx.background.celery.tasks.external_group_syncing\",\n \"ee.onyx.background.celery.tasks.cloud\",\n \"ee.onyx.background.celery.tasks.ttl_management\",\n \"ee.onyx.background.celery.tasks.usage_reporting\",\n ]\n )\n)\n" }, { "path": "backend/ee/onyx/background/celery/tasks/beat_schedule.py", "content": "from datetime import timedelta\nfrom typing import Any\n\nfrom ee.onyx.configs.app_configs import CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS\nfrom onyx.background.celery.tasks.beat_schedule import (\n beat_cloud_tasks as base_beat_system_tasks,\n)\nfrom onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT\nfrom onyx.background.celery.tasks.beat_schedule import (\n beat_task_templates as base_beat_task_templates,\n)\nfrom onyx.background.celery.tasks.beat_schedule import generate_cloud_tasks\nfrom onyx.background.celery.tasks.beat_schedule import (\n get_tasks_to_schedule as base_get_tasks_to_schedule,\n)\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom shared_configs.configs import MULTI_TENANT\n\nee_beat_system_tasks: list[dict] = []\n\nee_beat_task_templates: list[dict] = [\n {\n \"name\": \"autogenerate-usage-report\",\n \"task\": OnyxCeleryTask.GENERATE_USAGE_REPORT_TASK,\n \"schedule\": timedelta(days=30),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-ttl-management\",\n \"task\": OnyxCeleryTask.CHECK_TTL_MANAGEMENT_TASK,\n \"schedule\": timedelta(hours=CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"export-query-history-cleanup-task\",\n \"task\": OnyxCeleryTask.EXPORT_QUERY_HISTORY_CLEANUP_TASK,\n \"schedule\": timedelta(hours=1),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.CSV_GENERATION,\n },\n },\n]\n\nee_tasks_to_schedule: list[dict] = []\n\nif not MULTI_TENANT:\n ee_tasks_to_schedule = [\n {\n \"name\": \"hook-execution-log-cleanup\",\n \"task\": OnyxCeleryTask.HOOK_EXECUTION_LOG_CLEANUP_TASK,\n \"schedule\": timedelta(days=1),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"autogenerate-usage-report\",\n \"task\": OnyxCeleryTask.GENERATE_USAGE_REPORT_TASK,\n \"schedule\": timedelta(days=30), # TODO: change this to config flag\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-ttl-management\",\n \"task\": OnyxCeleryTask.CHECK_TTL_MANAGEMENT_TASK,\n \"schedule\": timedelta(hours=CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"export-query-history-cleanup-task\",\n \"task\": OnyxCeleryTask.EXPORT_QUERY_HISTORY_CLEANUP_TASK,\n \"schedule\": timedelta(hours=1),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.CSV_GENERATION,\n },\n },\n ]\n\n\ndef get_cloud_tasks_to_schedule(beat_multiplier: float) -> list[dict[str, Any]]:\n beat_system_tasks = ee_beat_system_tasks + base_beat_system_tasks\n beat_task_templates = ee_beat_task_templates + base_beat_task_templates\n cloud_tasks = generate_cloud_tasks(\n beat_system_tasks, beat_task_templates, beat_multiplier\n )\n return cloud_tasks\n\n\ndef get_tasks_to_schedule() -> list[dict[str, Any]]:\n return ee_tasks_to_schedule + base_get_tasks_to_schedule()\n" }, { "path": "backend/ee/onyx/background/celery/tasks/cleanup/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/cleanup/tasks.py", "content": "from datetime import datetime\nfrom datetime import timedelta\n\nfrom celery import shared_task\n\nfrom ee.onyx.db.query_history import get_all_query_history_export_tasks\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.enums import TaskStatus\nfrom onyx.db.tasks import delete_task_with_id\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n@shared_task(\n name=OnyxCeleryTask.EXPORT_QUERY_HISTORY_CLEANUP_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n)\ndef export_query_history_cleanup_task(*, tenant_id: str) -> None:\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n tasks = get_all_query_history_export_tasks(db_session=db_session)\n\n for task in tasks:\n if task.status == TaskStatus.SUCCESS:\n delete_task_with_id(db_session=db_session, task_id=task.task_id)\n elif task.status == TaskStatus.FAILURE:\n if task.start_time:\n deadline = task.start_time + timedelta(hours=24)\n now = datetime.now()\n if now < deadline:\n continue\n\n logger.error(\n f\"Task with {task.task_id=} failed; it is being deleted now\"\n )\n delete_task_with_id(db_session=db_session, task_id=task.task_id)\n" }, { "path": "backend/ee/onyx/background/celery/tasks/cloud/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/cloud/tasks.py", "content": "import time\n\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import ONYX_CLOUD_TENANT_ID\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import redis_lock_dump\nfrom shared_configs.configs import IGNORED_SYNCING_TENANT_LIST\n\n\n@shared_task(\n name=OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,\n ignore_result=True,\n trail=False,\n bind=True,\n)\ndef cloud_beat_task_generator(\n self: Task,\n task_name: str,\n queue: str = OnyxCeleryTask.DEFAULT,\n priority: int = OnyxCeleryPriority.MEDIUM,\n expires: int = BEAT_EXPIRES_DEFAULT,\n) -> bool | None:\n \"\"\"a lightweight task used to kick off individual beat tasks per tenant.\"\"\"\n time_start = time.monotonic()\n\n redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n\n lock_beat: RedisLock = redis_client.lock(\n f\"{OnyxRedisLocks.CLOUD_BEAT_TASK_GENERATOR_LOCK}:{task_name}\",\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n last_lock_time = time.monotonic()\n tenant_ids: list[str] = []\n num_processed_tenants = 0\n\n try:\n tenant_ids = get_all_tenant_ids()\n\n # NOTE: for now, we are running tasks for gated tenants, since we want to allow\n # connector deletion to run successfully. The new plan is to continously prune\n # the gated tenants set, so we won't have a build up of old, unused gated tenants.\n # Keeping this around in case we want to revert to the previous behavior.\n # gated_tenants = get_gated_tenants()\n\n for tenant_id in tenant_ids:\n # Same comment here as the above NOTE\n # if tenant_id in gated_tenants:\n # continue\n\n current_time = time.monotonic()\n if current_time - last_lock_time >= (CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4):\n lock_beat.reacquire()\n last_lock_time = current_time\n\n # needed in the cloud\n if IGNORED_SYNCING_TENANT_LIST and tenant_id in IGNORED_SYNCING_TENANT_LIST:\n continue\n\n self.app.send_task(\n task_name,\n kwargs=dict(\n tenant_id=tenant_id,\n ),\n queue=queue,\n priority=priority,\n expires=expires,\n ignore_result=True,\n )\n\n num_processed_tenants += 1\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception:\n task_logger.exception(\"Unexpected exception during cloud_beat_task_generator\")\n finally:\n if not lock_beat.owned():\n task_logger.error(\n \"cloud_beat_task_generator - Lock not owned on completion\"\n )\n redis_lock_dump(lock_beat, redis_client)\n else:\n lock_beat.release()\n\n time_elapsed = time.monotonic() - time_start\n task_logger.info(\n f\"cloud_beat_task_generator finished: \"\n f\"task={task_name} \"\n f\"num_processed_tenants={num_processed_tenants} \"\n f\"num_tenants={len(tenant_ids)} \"\n f\"elapsed={time_elapsed:.2f}\"\n )\n return True\n" }, { "path": "backend/ee/onyx/background/celery/tasks/doc_permission_syncing/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom time import sleep\nfrom typing import Any\nfrom typing import cast\nfrom uuid import uuid4\n\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom pydantic import ValidationError\nfrom redis import Redis\nfrom redis.exceptions import LockError\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\nfrom tenacity import retry\nfrom tenacity import retry_if_exception\nfrom tenacity import stop_after_delay\nfrom tenacity import wait_random_exponential\n\nfrom ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs\nfrom ee.onyx.db.document import upsert_document_external_perms\nfrom ee.onyx.external_permissions.sync_params import get_source_perm_sync_config\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_find_task\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_queue_length\nfrom onyx.background.celery.celery_redis import celery_get_queued_task_ids\nfrom onyx.background.celery.celery_redis import celery_get_unacked_task_ids\nfrom onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT\nfrom onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import OnyxRedisSignals\nfrom onyx.connectors.factory import validate_ccpair_for_user\nfrom onyx.db.connector import mark_cc_pair_as_permissions_synced\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.document import get_document_ids_for_connector_credential_pair\nfrom onyx.db.document import get_documents_for_connector_credential_pair_limited_columns\nfrom onyx.db.document import upsert_document_by_connector_credential_pair\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.hierarchy import (\n update_hierarchy_node_permissions as db_update_hierarchy_node_permissions,\n)\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.permission_sync_attempt import complete_doc_permission_sync_attempt\nfrom onyx.db.permission_sync_attempt import create_doc_permission_sync_attempt\nfrom onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_failed\nfrom onyx.db.permission_sync_attempt import (\n mark_doc_permission_sync_attempt_in_progress,\n)\nfrom onyx.db.sync_record import insert_sync_record\nfrom onyx.db.sync_record import update_sync_record_status\nfrom onyx.db.users import batch_add_ext_perm_user_if_not_exists\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import is_retryable_sqlalchemy_error\nfrom onyx.db.utils import SortOrder\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync\nfrom onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSyncPayload\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.redis.redis_pool import redis_lock_dump\nfrom onyx.server.runtime.onyx_runtime import OnyxRuntime\nfrom onyx.server.utils import make_short_id\nfrom onyx.utils.logger import doc_permission_sync_ctx\nfrom onyx.utils.logger import format_error_for_logging\nfrom onyx.utils.logger import LoggerContextVars\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\nDOCUMENT_PERMISSIONS_UPDATE_MAX_RETRIES = 3\nDOCUMENT_PERMISSIONS_UPDATE_STOP_AFTER = 10 * 60\nDOCUMENT_PERMISSIONS_UPDATE_MAX_WAIT = 60\n\n\n# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT\nLIGHT_SOFT_TIME_LIMIT = 105\nLIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15\n\n\ndef _get_fence_validation_block_expiration() -> int:\n \"\"\"\n Compute the expiration time for the fence validation block signal.\n Base expiration is 300 seconds, multiplied by the beat multiplier only in MULTI_TENANT mode.\n \"\"\"\n base_expiration = 300 # seconds\n\n if not MULTI_TENANT:\n return base_expiration\n\n try:\n beat_multiplier = OnyxRuntime.get_beat_multiplier()\n except Exception:\n beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n return int(base_expiration * beat_multiplier)\n\n\n\"\"\"Jobs / utils for kicking off doc permissions sync tasks.\"\"\"\n\n\ndef _fail_doc_permission_sync_attempt(attempt_id: int, error_msg: str) -> None:\n \"\"\"Helper to mark a doc permission sync attempt as failed with an error message.\"\"\"\n with get_session_with_current_tenant() as db_session:\n mark_doc_permission_sync_attempt_failed(\n attempt_id, db_session, error_message=error_msg\n )\n\n\ndef _is_external_doc_permissions_sync_due(cc_pair: ConnectorCredentialPair) -> bool:\n \"\"\"Returns boolean indicating if external doc permissions sync is due.\"\"\"\n\n if cc_pair.access_type != AccessType.SYNC:\n return False\n\n # skip doc permissions sync if not active\n if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:\n return False\n\n sync_config = get_source_perm_sync_config(cc_pair.connector.source)\n if sync_config is None:\n logger.error(f\"No sync config found for {cc_pair.connector.source}\")\n return False\n\n if sync_config.doc_sync_config is None:\n logger.error(f\"No doc sync config found for {cc_pair.connector.source}\")\n return False\n\n # if indexing also does perm sync, don't start running doc_sync until at\n # least one indexing is done\n if (\n sync_config.doc_sync_config.initial_index_should_sync\n and cc_pair.last_successful_index_time is None\n ):\n return False\n\n # If the last sync is None, it has never been run so we run the sync\n last_perm_sync = cc_pair.last_time_perm_sync\n if last_perm_sync is None:\n return True\n\n source_sync_period = sync_config.doc_sync_config.doc_sync_frequency\n source_sync_period *= int(OnyxRuntime.get_doc_permission_sync_multiplier())\n\n # If the last sync is greater than the full fetch period, we run the sync\n next_sync = last_perm_sync + timedelta(seconds=source_sync_period)\n if datetime.now(timezone.utc) >= next_sync:\n return True\n\n return False\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n)\ndef check_for_doc_permissions_sync(self: Task, *, tenant_id: str) -> bool | None:\n # TODO(rkuo): merge into check function after lookup table for fences is added\n\n # we need to use celery's redis client to access its redis data\n # (which lives on a different db number)\n r = get_redis_client()\n r_replica = get_redis_replica_client()\n\n lock_beat: RedisLock = r.lock(\n OnyxRedisLocks.CHECK_CONNECTOR_DOC_PERMISSIONS_SYNC_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n try:\n # get all cc pairs that need to be synced\n cc_pair_ids_to_sync: list[int] = []\n with get_session_with_current_tenant() as db_session:\n cc_pairs = get_all_auto_sync_cc_pairs(db_session)\n\n for cc_pair in cc_pairs:\n if _is_external_doc_permissions_sync_due(cc_pair):\n cc_pair_ids_to_sync.append(cc_pair.id)\n\n lock_beat.reacquire()\n for cc_pair_id in cc_pair_ids_to_sync:\n payload_id = try_creating_permissions_sync_task(\n self.app, cc_pair_id, r, tenant_id\n )\n if not payload_id:\n continue\n\n task_logger.info(\n f\"Permissions sync queued: cc_pair={cc_pair_id} id={payload_id}\"\n )\n\n # we want to run this less frequently than the overall task\n lock_beat.reacquire()\n if not r.exists(OnyxRedisSignals.BLOCK_VALIDATE_PERMISSION_SYNC_FENCES):\n # clear any permission fences that don't have associated celery tasks in progress\n # tasks can be in the queue in redis, in reserved tasks (prefetched by the worker),\n # or be currently executing\n try:\n r_celery = celery_get_broker_client(self.app)\n validate_permission_sync_fences(\n tenant_id, r, r_replica, r_celery, lock_beat\n )\n except Exception:\n task_logger.exception(\n \"Exception while validating permission sync fences\"\n )\n\n r.set(\n OnyxRedisSignals.BLOCK_VALIDATE_PERMISSION_SYNC_FENCES,\n 1,\n ex=_get_fence_validation_block_expiration(),\n )\n\n # use a lookup table to find active fences. We still have to verify the fence\n # exists since it is an optimization and not the source of truth.\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n\n if not r.exists(key_bytes):\n r.srem(OnyxRedisConstants.ACTIVE_FENCES, key_bytes)\n continue\n\n key_str = key_bytes.decode(\"utf-8\")\n if key_str.startswith(RedisConnectorPermissionSync.FENCE_PREFIX):\n with get_session_with_current_tenant() as db_session:\n monitor_ccpair_permissions_taskset(\n tenant_id, key_bytes, r, db_session\n )\n task_logger.info(f\"check_for_doc_permissions_sync finished: tenant={tenant_id}\")\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(\n f\"Unexpected check_for_doc_permissions_sync exception: tenant={tenant_id} {error_msg}\"\n )\n task_logger.exception(\n f\"Unexpected check_for_doc_permissions_sync exception: tenant={tenant_id}\"\n )\n finally:\n if lock_beat.owned():\n lock_beat.release()\n\n return True\n\n\ndef try_creating_permissions_sync_task(\n app: Celery,\n cc_pair_id: int,\n r: Redis,\n tenant_id: str,\n) -> str | None:\n \"\"\"Returns a randomized payload id on success.\n Returns None if no syncing is required.\"\"\"\n LOCK_TIMEOUT = 30\n\n payload_id: str | None = None\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n lock: RedisLock = r.lock(\n DANSWER_REDIS_FUNCTION_LOCK_PREFIX + \"try_generate_permissions_sync_tasks\",\n timeout=LOCK_TIMEOUT,\n )\n\n acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)\n if not acquired:\n return None\n\n try:\n if redis_connector.permissions.fenced:\n return None\n\n if redis_connector.delete.fenced:\n return None\n\n if redis_connector.prune.fenced:\n return None\n\n redis_connector.permissions.generator_clear()\n redis_connector.permissions.taskset_clear()\n\n custom_task_id = f\"{redis_connector.permissions.generator_task_key}_{uuid4()}\"\n\n # create before setting fence to avoid race condition where the monitoring\n # task updates the sync record before it is created\n try:\n with get_session_with_current_tenant() as db_session:\n insert_sync_record(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.EXTERNAL_PERMISSIONS,\n )\n except Exception:\n task_logger.exception(\"insert_sync_record exceptioned.\")\n\n # set a basic fence to start\n redis_connector.permissions.set_active()\n payload = RedisConnectorPermissionSyncPayload(\n id=make_short_id(),\n submitted=datetime.now(timezone.utc),\n started=None,\n celery_task_id=None,\n )\n redis_connector.permissions.set_fence(payload)\n\n result = app.send_task(\n OnyxCeleryTask.CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK,\n kwargs=dict(\n cc_pair_id=cc_pair_id,\n tenant_id=tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n )\n\n # fill in the celery task id\n payload.celery_task_id = result.id\n redis_connector.permissions.set_fence(payload)\n\n payload_id = payload.id\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(\n f\"Unexpected try_creating_permissions_sync_task exception: cc_pair={cc_pair_id} {error_msg}\"\n )\n return None\n finally:\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"try_creating_permissions_sync_task finished: cc_pair={cc_pair_id} payload_id={payload_id}\"\n )\n return payload_id\n\n\n@shared_task(\n name=OnyxCeleryTask.CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK,\n acks_late=False,\n soft_time_limit=JOB_TIMEOUT,\n track_started=True,\n trail=False,\n bind=True,\n)\ndef connector_permission_sync_generator_task(\n self: Task,\n cc_pair_id: int,\n tenant_id: str,\n) -> None:\n \"\"\"\n Permission sync task that handles document permission syncing for a given connector credential pair\n This task assumes that the task has already been properly fenced\n \"\"\"\n\n payload_id: str | None = None\n\n LoggerContextVars.reset()\n\n doc_permission_sync_ctx_dict = doc_permission_sync_ctx.get()\n doc_permission_sync_ctx_dict[\"cc_pair_id\"] = cc_pair_id\n doc_permission_sync_ctx_dict[\"request_id\"] = self.request.id\n doc_permission_sync_ctx.set(doc_permission_sync_ctx_dict)\n\n with get_session_with_current_tenant() as db_session:\n attempt_id = create_doc_permission_sync_attempt(\n connector_credential_pair_id=cc_pair_id,\n db_session=db_session,\n )\n task_logger.info(\n f\"Created doc permission sync attempt: {attempt_id} for cc_pair={cc_pair_id}\"\n )\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n r = get_redis_client()\n\n # this wait is needed to avoid a race condition where\n # the primary worker sends the task and it is immediately executed\n # before the primary worker can finalize the fence\n start = time.monotonic()\n while True:\n if time.monotonic() - start > CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT:\n error_msg = (\n f\"connector_permission_sync_generator_task - timed out waiting for fence to be ready: \"\n f\"fence={redis_connector.permissions.fence_key}\"\n )\n _fail_doc_permission_sync_attempt(attempt_id, error_msg)\n raise ValueError(error_msg)\n\n if not redis_connector.permissions.fenced: # The fence must exist\n error_msg = f\"connector_permission_sync_generator_task - fence not found: fence={redis_connector.permissions.fence_key}\"\n _fail_doc_permission_sync_attempt(attempt_id, error_msg)\n raise ValueError(error_msg)\n\n payload = redis_connector.permissions.payload # The payload must exist\n if not payload:\n error_msg = (\n \"connector_permission_sync_generator_task: payload invalid or not found\"\n )\n _fail_doc_permission_sync_attempt(attempt_id, error_msg)\n raise ValueError(error_msg)\n\n if payload.celery_task_id is None:\n logger.info(\n f\"connector_permission_sync_generator_task - Waiting for fence: fence={redis_connector.permissions.fence_key}\"\n )\n sleep(1)\n continue\n\n payload_id = payload.id\n\n logger.info(\n f\"connector_permission_sync_generator_task - Fence found, continuing...: \"\n f\"fence={redis_connector.permissions.fence_key} \"\n f\"payload_id={payload.id}\"\n )\n break\n\n lock: RedisLock = r.lock(\n OnyxRedisLocks.CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX\n + f\"_{redis_connector.cc_pair_id}\",\n timeout=CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT,\n thread_local=False,\n )\n\n acquired = lock.acquire(blocking=False)\n if not acquired:\n error_msg = (\n f\"Permission sync task already running, exiting...: cc_pair={cc_pair_id}\"\n )\n task_logger.warning(error_msg)\n _fail_doc_permission_sync_attempt(attempt_id, error_msg)\n return None\n\n try:\n with get_session_with_current_tenant() as db_session:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n eager_load_connector=True,\n eager_load_credential=True,\n )\n if cc_pair is None:\n raise ValueError(\n f\"No connector credential pair found for id: {cc_pair_id}\"\n )\n\n try:\n created = validate_ccpair_for_user(\n cc_pair.connector.id,\n cc_pair.credential.id,\n cc_pair.access_type,\n db_session,\n enforce_creation=False,\n )\n if not created:\n task_logger.warning(\n f\"Unable to create connector credential pair for id: {cc_pair_id}\"\n )\n except Exception:\n task_logger.exception(\n f\"validate_ccpair_permissions_sync exceptioned: cc_pair={cc_pair_id}\"\n )\n # TODO: add some notification to the admins here\n raise\n\n source_type = cc_pair.connector.source\n sync_config = get_source_perm_sync_config(source_type)\n if sync_config is None:\n error_msg = f\"No sync config found for {source_type}\"\n logger.error(error_msg)\n _fail_doc_permission_sync_attempt(attempt_id, error_msg)\n return None\n\n if sync_config.doc_sync_config is None:\n if sync_config.censoring_config:\n error_msg = f\"Doc sync config is None but censoring config exists for {source_type}\"\n _fail_doc_permission_sync_attempt(attempt_id, error_msg)\n return None\n\n raise ValueError(\n f\"No doc sync func found for {source_type} with cc_pair={cc_pair_id}\"\n )\n\n logger.info(f\"Syncing docs for {source_type} with cc_pair={cc_pair_id}\")\n\n mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)\n\n payload = redis_connector.permissions.payload\n if not payload:\n raise ValueError(f\"No fence payload found: cc_pair={cc_pair_id}\")\n\n new_payload = RedisConnectorPermissionSyncPayload(\n id=payload.id,\n submitted=payload.submitted,\n started=datetime.now(timezone.utc),\n celery_task_id=payload.celery_task_id,\n )\n redis_connector.permissions.set_fence(new_payload)\n\n callback = PermissionSyncCallback(\n redis_connector, lock, r, timeout_seconds=JOB_TIMEOUT\n )\n\n # pass in the capability to fetch all existing docs for the cc_pair\n # this is can be used to determine documents that are \"missing\" and thus\n # should no longer be accessible. The decision as to whether we should find\n # every document during the doc sync process is connector-specific.\n def fetch_all_existing_docs_fn(\n sort_order: SortOrder | None = None,\n ) -> list[DocumentRow]:\n result = get_documents_for_connector_credential_pair_limited_columns(\n db_session=db_session,\n connector_id=cc_pair.connector.id,\n credential_id=cc_pair.credential.id,\n sort_order=sort_order,\n )\n return list(result)\n\n def fetch_all_existing_docs_ids_fn() -> list[str]:\n result = get_document_ids_for_connector_credential_pair(\n db_session=db_session,\n connector_id=cc_pair.connector.id,\n credential_id=cc_pair.credential.id,\n )\n return result\n\n doc_sync_func = sync_config.doc_sync_config.doc_sync_func\n document_external_accesses = doc_sync_func(\n cc_pair,\n fetch_all_existing_docs_fn,\n fetch_all_existing_docs_ids_fn,\n callback,\n )\n\n task_logger.info(\n f\"RedisConnector.permissions.generate_tasks starting. cc_pair={cc_pair_id}\"\n )\n\n tasks_generated = 0\n docs_with_errors = 0\n for doc_external_access in document_external_accesses:\n if callback.should_stop():\n raise RuntimeError(\n f\"Permission sync task timed out or stop signal detected: \"\n f\"cc_pair={cc_pair_id} \"\n f\"tasks_generated={tasks_generated}\"\n )\n\n result = redis_connector.permissions.update_db(\n lock=lock,\n new_permissions=[doc_external_access],\n source_string=source_type,\n connector_id=cc_pair.connector.id,\n credential_id=cc_pair.credential.id,\n task_logger=task_logger,\n )\n tasks_generated += result.num_updated\n docs_with_errors += result.num_errors\n\n task_logger.info(\n f\"RedisConnector.permissions.generate_tasks finished. \"\n f\"cc_pair={cc_pair_id} tasks_generated={tasks_generated} docs_with_errors={docs_with_errors}\"\n )\n\n complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_docs_synced=tasks_generated,\n docs_with_permission_errors=docs_with_errors,\n )\n task_logger.info(\n f\"Completed doc permission sync attempt {attempt_id}: {tasks_generated} docs, {docs_with_errors} errors\"\n )\n\n redis_connector.permissions.generator_complete = tasks_generated\n\n except Exception as e:\n error_msg = format_error_for_logging(e)\n\n task_logger.warning(\n f\"Permission sync exceptioned: cc_pair={cc_pair_id} payload_id={payload_id} {error_msg}\"\n )\n task_logger.exception(\n f\"Permission sync exceptioned: cc_pair={cc_pair_id} payload_id={payload_id}\"\n )\n\n with get_session_with_current_tenant() as db_session:\n mark_doc_permission_sync_attempt_failed(\n attempt_id, db_session, error_message=error_msg\n )\n\n redis_connector.permissions.generator_clear()\n redis_connector.permissions.taskset_clear()\n redis_connector.permissions.set_fence(None)\n raise e\n finally:\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"Permission sync finished: cc_pair={cc_pair_id} payload_id={payload.id}\"\n )\n\n\n# NOTE(rkuo): this should probably move to the db layer\n@retry(\n retry=retry_if_exception(is_retryable_sqlalchemy_error),\n wait=wait_random_exponential(\n multiplier=1, max=DOCUMENT_PERMISSIONS_UPDATE_MAX_WAIT\n ),\n stop=stop_after_delay(DOCUMENT_PERMISSIONS_UPDATE_STOP_AFTER),\n)\ndef element_update_permissions(\n tenant_id: str,\n permissions: ElementExternalAccess,\n source_type_str: str,\n connector_id: int,\n credential_id: int,\n) -> bool:\n \"\"\"Update permissions for a document or hierarchy node.\"\"\"\n start = time.monotonic()\n external_access = permissions.external_access\n\n # Determine element type and identifier for logging\n if isinstance(permissions, DocExternalAccess):\n element_id = permissions.doc_id\n element_type = \"doc\"\n else:\n element_id = permissions.raw_node_id\n element_type = \"node\"\n\n try:\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # Add the users to the DB if they don't exist\n batch_add_ext_perm_user_if_not_exists(\n db_session=db_session,\n emails=list(external_access.external_user_emails),\n continue_on_error=True,\n )\n\n if isinstance(permissions, DocExternalAccess):\n # Document permission update\n created_new_doc = upsert_document_external_perms(\n db_session=db_session,\n doc_id=permissions.doc_id,\n external_access=external_access,\n source_type=DocumentSource(source_type_str),\n )\n\n if created_new_doc:\n # If a new document was created, we associate it with the cc_pair\n upsert_document_by_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n document_ids=[permissions.doc_id],\n )\n else:\n # Hierarchy node permission update\n db_update_hierarchy_node_permissions(\n db_session=db_session,\n raw_node_id=permissions.raw_node_id,\n source=DocumentSource(permissions.source),\n is_public=external_access.is_public,\n external_user_emails=(\n list(external_access.external_user_emails)\n if external_access.external_user_emails\n else None\n ),\n external_user_group_ids=(\n list(external_access.external_user_group_ids)\n if external_access.external_user_group_ids\n else None\n ),\n )\n\n elapsed = time.monotonic() - start\n task_logger.info(\n f\"{element_type}={element_id} action=update_permissions elapsed={elapsed:.2f}\"\n )\n except Exception as e:\n task_logger.exception(\n f\"element_update_permissions exceptioned: {element_type}={element_id}, {connector_id=} {credential_id=}\"\n )\n raise e\n finally:\n task_logger.info(\n f\"element_update_permissions completed: {element_type}={element_id}, {connector_id=} {credential_id=}\"\n )\n\n return True\n\n\ndef validate_permission_sync_fences(\n tenant_id: str,\n r: Redis,\n r_replica: Redis,\n r_celery: Redis,\n lock_beat: RedisLock,\n) -> None:\n # building lookup table can be expensive, so we won't bother\n # validating until the queue is small\n PERMISSION_SYNC_VALIDATION_MAX_QUEUE_LEN = 1024\n\n queue_len = celery_get_queue_length(\n OnyxCeleryQueues.DOC_PERMISSIONS_UPSERT, r_celery\n )\n if queue_len > PERMISSION_SYNC_VALIDATION_MAX_QUEUE_LEN:\n return\n\n queued_upsert_tasks = celery_get_queued_task_ids(\n OnyxCeleryQueues.DOC_PERMISSIONS_UPSERT, r_celery\n )\n reserved_generator_tasks = celery_get_unacked_task_ids(\n OnyxCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC, r_celery\n )\n\n # validate all existing permission sync jobs\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n key_str = key_bytes.decode(\"utf-8\")\n if not key_str.startswith(RedisConnectorPermissionSync.FENCE_PREFIX):\n continue\n\n validate_permission_sync_fence(\n tenant_id,\n key_bytes,\n queued_upsert_tasks,\n reserved_generator_tasks,\n r,\n r_celery,\n )\n\n lock_beat.reacquire()\n\n return\n\n\ndef validate_permission_sync_fence(\n tenant_id: str,\n key_bytes: bytes,\n queued_tasks: set[str],\n reserved_tasks: set[str],\n r: Redis,\n r_celery: Redis,\n) -> None:\n \"\"\"Checks for the error condition where an indexing fence is set but the associated celery tasks don't exist.\n This can happen if the indexing worker hard crashes or is terminated.\n Being in this bad state means the fence will never clear without help, so this function\n gives the help.\n\n How this works:\n 1. This function renews the active signal with a 5 minute TTL under the following conditions\n 1.2. When the task is seen in the redis queue\n 1.3. When the task is seen in the reserved / prefetched list\n\n 2. Externally, the active signal is renewed when:\n 2.1. The fence is created\n 2.2. The indexing watchdog checks the spawned task.\n\n 3. The TTL allows us to get through the transitions on fence startup\n and when the task starts executing.\n\n More TTL clarification: it is seemingly impossible to exactly query Celery for\n whether a task is in the queue or currently executing.\n 1. An unknown task id is always returned as state PENDING.\n 2. Redis can be inspected for the task id, but the task id is gone between the time a worker receives the task\n and the time it actually starts on the worker.\n\n queued_tasks: the celery queue of lightweight permission sync tasks\n reserved_tasks: prefetched tasks for sync task generator\n \"\"\"\n # if the fence doesn't exist, there's nothing to do\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n task_logger.warning(\n f\"validate_permission_sync_fence - could not parse id from {fence_key}\"\n )\n return\n\n cc_pair_id = int(cc_pair_id_str)\n # parse out metadata and initialize the helper class with it\n redis_connector = RedisConnector(tenant_id, int(cc_pair_id))\n\n # check to see if the fence/payload exists\n if not redis_connector.permissions.fenced:\n return\n\n # in the cloud, the payload format may have changed ...\n # it's a little sloppy, but just reset the fence for now if that happens\n # TODO: add intentional cleanup/abort logic\n try:\n payload = redis_connector.permissions.payload\n except ValidationError:\n task_logger.exception(\n \"validate_permission_sync_fence - \"\n \"Resetting fence because fence schema is out of date: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key}\"\n )\n\n redis_connector.permissions.reset()\n return\n\n if not payload:\n return\n\n if not payload.celery_task_id:\n return\n\n # OK, there's actually something for us to validate\n\n # either the generator task must be in flight or its subtasks must be\n found = celery_find_task(\n payload.celery_task_id,\n OnyxCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC,\n r_celery,\n )\n if found:\n # the celery task exists in the redis queue\n redis_connector.permissions.set_active()\n return\n\n if payload.celery_task_id in reserved_tasks:\n # the celery task was prefetched and is reserved within a worker\n redis_connector.permissions.set_active()\n return\n\n # look up every task in the current taskset in the celery queue\n # every entry in the taskset should have an associated entry in the celery task queue\n # because we get the celery tasks first, the entries in our own permissions taskset\n # should be roughly a subset of the tasks in celery\n\n # this check isn't very exact, but should be sufficient over a period of time\n # A single successful check over some number of attempts is sufficient.\n\n # TODO: if the number of tasks in celery is much lower than than the taskset length\n # we might be able to shortcut the lookup since by definition some of the tasks\n # must not exist in celery.\n\n tasks_scanned = 0\n tasks_not_in_celery = 0 # a non-zero number after completing our check is bad\n\n for member in r.sscan_iter(redis_connector.permissions.taskset_key):\n tasks_scanned += 1\n\n member_bytes = cast(bytes, member)\n member_str = member_bytes.decode(\"utf-8\")\n if member_str in queued_tasks:\n continue\n\n if member_str in reserved_tasks:\n continue\n\n tasks_not_in_celery += 1\n\n task_logger.info(\n f\"validate_permission_sync_fence task check: tasks_scanned={tasks_scanned} tasks_not_in_celery={tasks_not_in_celery}\"\n )\n\n # we're active if there are still tasks to run and those tasks all exist in celery\n if tasks_scanned > 0 and tasks_not_in_celery == 0:\n redis_connector.permissions.set_active()\n return\n\n # we may want to enable this check if using the active task list somehow isn't good enough\n # if redis_connector_index.generator_locked():\n # logger.info(f\"{payload.celery_task_id} is currently executing.\")\n\n # if we get here, we didn't find any direct indication that the associated celery tasks exist,\n # but they still might be there due to gaps in our ability to check states during transitions\n # Checking the active signal safeguards us against these transition periods\n # (which has a duration that allows us to bridge those gaps)\n if redis_connector.permissions.active():\n return\n\n # celery tasks don't exist and the active signal has expired, possibly due to a crash. Clean it up.\n task_logger.warning(\n \"validate_permission_sync_fence - \"\n \"Resetting fence because no associated celery tasks were found: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key} \"\n f\"payload_id={payload.id}\"\n )\n\n redis_connector.permissions.reset()\n return\n\n\nclass PermissionSyncCallback(IndexingHeartbeatInterface):\n PARENT_CHECK_INTERVAL = 60\n\n def __init__(\n self,\n redis_connector: RedisConnector,\n redis_lock: RedisLock,\n redis_client: Redis,\n timeout_seconds: int | None = None,\n ):\n super().__init__()\n self.redis_connector: RedisConnector = redis_connector\n self.redis_lock: RedisLock = redis_lock\n self.redis_client = redis_client\n\n self.started: datetime = datetime.now(timezone.utc)\n self.redis_lock.reacquire()\n\n self.last_tag: str = \"PermissionSyncCallback.__init__\"\n self.last_lock_reacquire: datetime = datetime.now(timezone.utc)\n self.last_lock_monotonic = time.monotonic()\n self.start_monotonic = time.monotonic()\n self.timeout_seconds = timeout_seconds\n\n def should_stop(self) -> bool:\n if self.redis_connector.stop.fenced:\n return True\n\n # Check if the task has exceeded its timeout\n # NOTE: Celery's soft_time_limit does not work with thread pools,\n # so we must enforce timeouts internally.\n if self.timeout_seconds is not None:\n elapsed = time.monotonic() - self.start_monotonic\n if elapsed > self.timeout_seconds:\n logger.warning(\n f\"PermissionSyncCallback - task timeout exceeded: \"\n f\"elapsed={elapsed:.0f}s timeout={self.timeout_seconds}s \"\n f\"cc_pair={self.redis_connector.cc_pair_id}\"\n )\n return True\n\n return False\n\n def progress(self, tag: str, amount: int) -> None: # noqa: ARG002\n try:\n self.redis_connector.permissions.set_active()\n\n current_time = time.monotonic()\n if current_time - self.last_lock_monotonic >= (\n CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4\n ):\n self.redis_lock.reacquire()\n self.last_lock_reacquire = datetime.now(timezone.utc)\n self.last_lock_monotonic = time.monotonic()\n\n self.last_tag = tag\n except LockError:\n logger.exception(\n f\"PermissionSyncCallback - lock.reacquire exceptioned: \"\n f\"lock_timeout={self.redis_lock.timeout} \"\n f\"start={self.started} \"\n f\"last_tag={self.last_tag} \"\n f\"last_reacquired={self.last_lock_reacquire} \"\n f\"now={datetime.now(timezone.utc)}\"\n )\n\n redis_lock_dump(self.redis_lock, self.redis_client)\n raise\n\n\n\"\"\"Monitoring CCPair permissions utils\"\"\"\n\n\ndef monitor_ccpair_permissions_taskset(\n tenant_id: str,\n key_bytes: bytes,\n r: Redis, # noqa: ARG001\n db_session: Session,\n) -> None:\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n task_logger.warning(\n f\"monitor_ccpair_permissions_taskset: could not parse cc_pair_id from {fence_key}\"\n )\n return\n\n cc_pair_id = int(cc_pair_id_str)\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n if not redis_connector.permissions.fenced:\n return\n\n initial = redis_connector.permissions.generator_complete\n if initial is None:\n return\n\n try:\n payload = redis_connector.permissions.payload\n except ValidationError:\n task_logger.exception(\n \"Permissions sync payload failed to validate. Schema may have been updated.\"\n )\n return\n\n if not payload:\n return\n\n remaining = redis_connector.permissions.get_remaining()\n task_logger.info(\n f\"Permissions sync progress: cc_pair={cc_pair_id} id={payload.id} remaining={remaining} initial={initial}\"\n )\n\n # Add telemetry for permission syncing progress\n optional_telemetry(\n record_type=RecordType.PERMISSION_SYNC_PROGRESS,\n data={\n \"cc_pair_id\": cc_pair_id,\n \"total_docs_synced\": initial if initial is not None else 0,\n \"remaining_docs_to_sync\": remaining,\n },\n tenant_id=tenant_id,\n )\n\n if remaining > 0:\n return\n\n mark_cc_pair_as_permissions_synced(db_session, int(cc_pair_id), payload.started)\n task_logger.info(\n f\"Permissions sync finished: cc_pair={cc_pair_id} id={payload.id} num_synced={initial}\"\n )\n\n # Add telemetry for permission syncing complete\n optional_telemetry(\n record_type=RecordType.PERMISSION_SYNC_COMPLETE,\n data={\"cc_pair_id\": cc_pair_id},\n tenant_id=tenant_id,\n )\n\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.EXTERNAL_PERMISSIONS,\n sync_status=SyncStatus.SUCCESS,\n num_docs_synced=initial,\n )\n\n redis_connector.permissions.reset()\n" }, { "path": "backend/ee/onyx/background/celery/tasks/external_group_syncing/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/external_group_syncing/group_sync_utils.py", "content": "from sqlalchemy.orm import Session\n\nfrom ee.onyx.external_permissions.sync_params import (\n source_group_sync_is_cc_pair_agnostic,\n)\nfrom onyx.db.connector import mark_cc_pair_as_external_group_synced\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs_for_source\nfrom onyx.db.models import ConnectorCredentialPair\n\n\ndef _get_all_cc_pair_ids_to_mark_as_group_synced(\n db_session: Session, cc_pair: ConnectorCredentialPair\n) -> list[int]:\n if not source_group_sync_is_cc_pair_agnostic(cc_pair.connector.source):\n return [cc_pair.id]\n\n cc_pairs = get_connector_credential_pairs_for_source(\n db_session, cc_pair.connector.source\n )\n return [cc_pair.id for cc_pair in cc_pairs]\n\n\ndef mark_all_relevant_cc_pairs_as_external_group_synced(\n db_session: Session, cc_pair: ConnectorCredentialPair\n) -> None:\n \"\"\"For some source types, one successful group sync run should count for all\n cc pairs of that type. This function handles that case.\"\"\"\n cc_pair_ids = _get_all_cc_pair_ids_to_mark_as_group_synced(db_session, cc_pair)\n for cc_pair_id in cc_pair_ids:\n mark_cc_pair_as_external_group_synced(db_session, cc_pair_id)\n" }, { "path": "backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom uuid import uuid4\n\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom pydantic import ValidationError\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\n\nfrom ee.onyx.background.celery.tasks.external_group_syncing.group_sync_utils import (\n mark_all_relevant_cc_pairs_as_external_group_synced,\n)\nfrom ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs\nfrom ee.onyx.db.connector_credential_pair import get_cc_pairs_by_source\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom ee.onyx.db.external_perm import mark_old_external_groups_as_stale\nfrom ee.onyx.db.external_perm import remove_stale_external_groups\nfrom ee.onyx.db.external_perm import upsert_external_groups\nfrom ee.onyx.external_permissions.sync_params import (\n get_all_cc_pair_agnostic_group_sync_sources,\n)\nfrom ee.onyx.external_permissions.sync_params import get_source_perm_sync_config\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_find_task\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_unacked_task_ids\nfrom onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT\nfrom onyx.background.error_logging import emit_background_error\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import OnyxRedisSignals\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.permission_sync_attempt import complete_external_group_sync_attempt\nfrom onyx.db.permission_sync_attempt import (\n create_external_group_sync_attempt,\n)\nfrom onyx.db.permission_sync_attempt import (\n mark_external_group_sync_attempt_failed,\n)\nfrom onyx.db.permission_sync_attempt import (\n mark_external_group_sync_attempt_in_progress,\n)\nfrom onyx.db.sync_record import insert_sync_record\nfrom onyx.db.sync_record import update_sync_record_status\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync\nfrom onyx.redis.redis_connector_ext_group_sync import (\n RedisConnectorExternalGroupSyncPayload,\n)\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.server.runtime.onyx_runtime import OnyxRuntime\nfrom onyx.server.utils import make_short_id\nfrom onyx.utils.logger import format_error_for_logging\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\n_EXTERNAL_GROUP_BATCH_SIZE = 100\n\n\ndef _fail_external_group_sync_attempt(attempt_id: int, error_msg: str) -> None:\n \"\"\"Helper to mark an external group sync attempt as failed with an error message.\"\"\"\n with get_session_with_current_tenant() as db_session:\n mark_external_group_sync_attempt_failed(\n attempt_id, db_session, error_message=error_msg\n )\n\n\ndef _get_fence_validation_block_expiration() -> int:\n \"\"\"\n Compute the expiration time for the fence validation block signal.\n Base expiration is 300 seconds, multiplied by the beat multiplier only in MULTI_TENANT mode.\n \"\"\"\n base_expiration = 300 # seconds\n\n if not MULTI_TENANT:\n return base_expiration\n\n try:\n beat_multiplier = OnyxRuntime.get_beat_multiplier()\n except Exception:\n beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n return int(base_expiration * beat_multiplier)\n\n\ndef _is_external_group_sync_due(cc_pair: ConnectorCredentialPair) -> bool:\n \"\"\"Returns boolean indicating if external group sync is due.\"\"\"\n\n if cc_pair.access_type != AccessType.SYNC:\n task_logger.error(\n f\"Received non-sync CC Pair {cc_pair.id} for external group sync. Actual access type: {cc_pair.access_type}\"\n )\n return False\n\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n task_logger.debug(\n f\"Skipping group sync for CC Pair {cc_pair.id} - CC Pair is being deleted\"\n )\n return False\n\n sync_config = get_source_perm_sync_config(cc_pair.connector.source)\n if sync_config is None:\n task_logger.debug(\n f\"Skipping group sync for CC Pair {cc_pair.id} - no sync config found for {cc_pair.connector.source}\"\n )\n return False\n\n # If there is not group sync function for the connector, we don't run the sync\n # This is fine because all sources dont necessarily have a concept of groups\n if sync_config.group_sync_config is None:\n task_logger.debug(\n f\"Skipping group sync for CC Pair {cc_pair.id} - no group sync config found for {cc_pair.connector.source}\"\n )\n return False\n\n # If the last sync is None, it has never been run so we run the sync\n last_ext_group_sync = cc_pair.last_time_external_group_sync\n if last_ext_group_sync is None:\n return True\n\n source_sync_period = sync_config.group_sync_config.group_sync_frequency\n\n # If the last sync is greater than the full fetch period, we run the sync\n next_sync = last_ext_group_sync + timedelta(seconds=source_sync_period)\n if datetime.now(timezone.utc) >= next_sync:\n return True\n\n return False\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n)\ndef check_for_external_group_sync(self: Task, *, tenant_id: str) -> bool | None:\n # we need to use celery's redis client to access its redis data\n # (which lives on a different db number)\n r = get_redis_client()\n r_replica = get_redis_replica_client()\n\n lock_beat: RedisLock = r.lock(\n OnyxRedisLocks.CHECK_CONNECTOR_EXTERNAL_GROUP_SYNC_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n task_logger.warning(\n f\"Failed to acquire beat lock for external group sync: {tenant_id}\"\n )\n return None\n\n try:\n cc_pair_ids_to_sync: list[int] = []\n with get_session_with_current_tenant() as db_session:\n cc_pairs = get_all_auto_sync_cc_pairs(db_session)\n\n # For some sources, we only want to sync one cc_pair per source type\n for source in get_all_cc_pair_agnostic_group_sync_sources():\n # These are ordered by cc_pair id so the first one is the one we want\n cc_pairs_to_dedupe = get_cc_pairs_by_source(\n db_session,\n source,\n access_type=AccessType.SYNC,\n status=ConnectorCredentialPairStatus.ACTIVE,\n )\n # dedupe cc_pairs to only keep the first one\n for cc_pair_to_remove in cc_pairs_to_dedupe[1:]:\n cc_pairs = [\n cc_pair\n for cc_pair in cc_pairs\n if cc_pair.id != cc_pair_to_remove.id\n ]\n\n for cc_pair in cc_pairs:\n if _is_external_group_sync_due(cc_pair):\n cc_pair_ids_to_sync.append(cc_pair.id)\n\n lock_beat.reacquire()\n for cc_pair_id in cc_pair_ids_to_sync:\n payload_id = try_creating_external_group_sync_task(\n self.app, cc_pair_id, r, tenant_id\n )\n if not payload_id:\n continue\n\n task_logger.info(\n f\"External group sync queued: cc_pair={cc_pair_id} id={payload_id}\"\n )\n\n # we want to run this less frequently than the overall task\n lock_beat.reacquire()\n if not r.exists(OnyxRedisSignals.BLOCK_VALIDATE_EXTERNAL_GROUP_SYNC_FENCES):\n # clear fences that don't have associated celery tasks in progress\n # tasks can be in the queue in redis, in reserved tasks (prefetched by the worker),\n # or be currently executing\n try:\n r_celery = celery_get_broker_client(self.app)\n validate_external_group_sync_fences(\n tenant_id, self.app, r, r_replica, r_celery, lock_beat\n )\n except Exception:\n task_logger.exception(\n \"Exception while validating external group sync fences\"\n )\n\n r.set(\n OnyxRedisSignals.BLOCK_VALIDATE_EXTERNAL_GROUP_SYNC_FENCES,\n 1,\n ex=_get_fence_validation_block_expiration(),\n )\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(\n f\"Unexpected check_for_external_group_sync exception: tenant={tenant_id} {error_msg}\"\n )\n task_logger.exception(f\"Unexpected exception: tenant={tenant_id}\")\n finally:\n if lock_beat.owned():\n lock_beat.release()\n\n task_logger.info(f\"check_for_external_group_sync finished: tenant={tenant_id}\")\n return True\n\n\ndef try_creating_external_group_sync_task(\n app: Celery,\n cc_pair_id: int,\n r: Redis, # noqa: ARG001\n tenant_id: str,\n) -> str | None:\n \"\"\"Returns an int if syncing is needed. The int represents the number of sync tasks generated.\n Returns None if no syncing is required.\"\"\"\n payload_id: str | None = None\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n try:\n # Dont kick off a new sync if the previous one is still running\n if redis_connector.external_group_sync.fenced:\n logger.warning(\n f\"Skipping external group sync for CC Pair {cc_pair_id} - already running.\"\n )\n return None\n\n redis_connector.external_group_sync.generator_clear()\n redis_connector.external_group_sync.taskset_clear()\n\n # create before setting fence to avoid race condition where the monitoring\n # task updates the sync record before it is created\n try:\n with get_session_with_current_tenant() as db_session:\n insert_sync_record(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.EXTERNAL_GROUP,\n )\n except Exception:\n task_logger.exception(\"insert_sync_record exceptioned.\")\n\n # Signal active before creating fence\n redis_connector.external_group_sync.set_active()\n\n payload = RedisConnectorExternalGroupSyncPayload(\n id=make_short_id(),\n submitted=datetime.now(timezone.utc),\n started=None,\n celery_task_id=None,\n )\n redis_connector.external_group_sync.set_fence(payload)\n\n custom_task_id = f\"{redis_connector.external_group_sync.taskset_key}_{uuid4()}\"\n\n result = app.send_task(\n OnyxCeleryTask.CONNECTOR_EXTERNAL_GROUP_SYNC_GENERATOR_TASK,\n kwargs=dict(\n cc_pair_id=cc_pair_id,\n tenant_id=tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_EXTERNAL_GROUP_SYNC,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n )\n\n payload.celery_task_id = result.id\n redis_connector.external_group_sync.set_fence(payload)\n\n payload_id = payload.id\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(\n f\"Unexpected try_creating_external_group_sync_task exception: cc_pair={cc_pair_id} {error_msg}\"\n )\n task_logger.exception(\n f\"Unexpected exception while trying to create external group sync task: cc_pair={cc_pair_id}\"\n )\n return None\n\n task_logger.info(\n f\"try_creating_external_group_sync_task finished: cc_pair={cc_pair_id} payload_id={payload_id}\"\n )\n return payload_id\n\n\n@shared_task(\n name=OnyxCeleryTask.CONNECTOR_EXTERNAL_GROUP_SYNC_GENERATOR_TASK,\n acks_late=False,\n soft_time_limit=JOB_TIMEOUT,\n track_started=True,\n trail=False,\n bind=True,\n)\ndef connector_external_group_sync_generator_task(\n self: Task, # noqa: ARG001\n cc_pair_id: int,\n tenant_id: str,\n) -> None:\n \"\"\"\n External group sync task for a given connector credential pair\n This task assumes that the task has already been properly fenced\n \"\"\"\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n r = get_redis_client()\n\n # this wait is needed to avoid a race condition where\n # the primary worker sends the task and it is immediately executed\n # before the primary worker can finalize the fence\n start = time.monotonic()\n while True:\n if time.monotonic() - start > CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT:\n msg = (\n f\"connector_external_group_sync_generator_task - timed out waiting for fence to be ready: \"\n f\"fence={redis_connector.external_group_sync.fence_key}\"\n )\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n raise ValueError(msg)\n\n if not redis_connector.external_group_sync.fenced: # The fence must exist\n msg = (\n f\"connector_external_group_sync_generator_task - fence not found: \"\n f\"fence={redis_connector.external_group_sync.fence_key}\"\n )\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n raise ValueError(msg)\n\n payload = redis_connector.external_group_sync.payload # The payload must exist\n if not payload:\n msg = \"connector_external_group_sync_generator_task: payload invalid or not found\"\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n raise ValueError(msg)\n\n if payload.celery_task_id is None:\n logger.info(\n f\"connector_external_group_sync_generator_task - Waiting for fence: \"\n f\"fence={redis_connector.external_group_sync.fence_key}\"\n )\n time.sleep(1)\n continue\n\n logger.info(\n f\"connector_external_group_sync_generator_task - Fence found, continuing...: \"\n f\"fence={redis_connector.external_group_sync.fence_key} \"\n f\"payload_id={payload.id}\"\n )\n break\n\n lock: RedisLock = r.lock(\n OnyxRedisLocks.CONNECTOR_EXTERNAL_GROUP_SYNC_LOCK_PREFIX\n + f\"_{redis_connector.cc_pair_id}\",\n timeout=CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT,\n )\n\n acquired = lock.acquire(blocking=False)\n if not acquired:\n msg = f\"External group sync task already running, exiting...: cc_pair={cc_pair_id}\"\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n task_logger.error(msg)\n return None\n\n try:\n payload.started = datetime.now(timezone.utc)\n redis_connector.external_group_sync.set_fence(payload)\n\n _perform_external_group_sync(\n cc_pair_id=cc_pair_id,\n tenant_id=tenant_id,\n )\n\n with get_session_with_current_tenant() as db_session:\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.EXTERNAL_GROUP,\n sync_status=SyncStatus.SUCCESS,\n )\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(\n f\"External group sync exceptioned: cc_pair={cc_pair_id} payload_id={payload.id} {error_msg}\"\n )\n task_logger.exception(\n f\"External group sync exceptioned: cc_pair={cc_pair_id} payload_id={payload.id}\"\n )\n\n msg = f\"External group sync exceptioned: cc_pair={cc_pair_id} payload_id={payload.id}\"\n task_logger.exception(msg)\n emit_background_error(msg + f\"\\n\\n{e}\", cc_pair_id=cc_pair_id)\n\n with get_session_with_current_tenant() as db_session:\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.EXTERNAL_GROUP,\n sync_status=SyncStatus.FAILED,\n )\n\n redis_connector.external_group_sync.generator_clear()\n redis_connector.external_group_sync.taskset_clear()\n raise e\n finally:\n # we always want to clear the fence after the task is done or failed so it doesn't get stuck\n redis_connector.external_group_sync.set_fence(None)\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"External group sync finished: cc_pair={cc_pair_id} payload_id={payload.id}\"\n )\n\n\ndef _perform_external_group_sync(\n cc_pair_id: int,\n tenant_id: str,\n timeout_seconds: int = JOB_TIMEOUT,\n) -> None:\n # Create attempt record at the start\n with get_session_with_current_tenant() as db_session:\n attempt_id = create_external_group_sync_attempt(\n connector_credential_pair_id=cc_pair_id,\n db_session=db_session,\n )\n logger.info(\n f\"Created external group sync attempt: {attempt_id} for cc_pair={cc_pair_id}\"\n )\n\n with get_session_with_current_tenant() as db_session:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n eager_load_credential=True,\n )\n if cc_pair is None:\n raise ValueError(f\"No connector credential pair found for id: {cc_pair_id}\")\n\n source_type = cc_pair.connector.source\n sync_config = get_source_perm_sync_config(source_type)\n if sync_config is None:\n msg = f\"No sync config found for {source_type} for cc_pair: {cc_pair_id}\"\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n _fail_external_group_sync_attempt(attempt_id, msg)\n raise ValueError(msg)\n\n if sync_config.group_sync_config is None:\n msg = f\"No group sync config found for {source_type} for cc_pair: {cc_pair_id}\"\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n _fail_external_group_sync_attempt(attempt_id, msg)\n raise ValueError(msg)\n\n ext_group_sync_func = sync_config.group_sync_config.group_sync_func\n\n logger.info(\n f\"Marking old external groups as stale for {source_type} for cc_pair: {cc_pair_id}\"\n )\n mark_old_external_groups_as_stale(db_session, cc_pair_id)\n\n # Mark attempt as in progress\n mark_external_group_sync_attempt_in_progress(attempt_id, db_session)\n logger.info(f\"Marked external group sync attempt {attempt_id} as in progress\")\n\n logger.info(\n f\"Syncing external groups for {source_type} for cc_pair: {cc_pair_id}\"\n )\n external_user_group_batch: list[ExternalUserGroup] = []\n seen_users: set[str] = set() # Track unique users across all groups\n total_groups_processed = 0\n total_group_memberships_synced = 0\n start_time = time.monotonic()\n try:\n external_user_group_generator = ext_group_sync_func(tenant_id, cc_pair)\n for external_user_group in external_user_group_generator:\n # Check if the task has exceeded its timeout\n # NOTE: Celery's soft_time_limit does not work with thread pools,\n # so we must enforce timeouts internally.\n elapsed = time.monotonic() - start_time\n if elapsed > timeout_seconds:\n raise RuntimeError(\n f\"External group sync task timed out: \"\n f\"cc_pair={cc_pair_id} \"\n f\"elapsed={elapsed:.0f}s \"\n f\"timeout={timeout_seconds}s \"\n f\"groups_processed={total_groups_processed}\"\n )\n\n external_user_group_batch.append(external_user_group)\n\n # Track progress\n total_groups_processed += 1\n total_group_memberships_synced += len(external_user_group.user_emails)\n seen_users = seen_users.union(external_user_group.user_emails)\n\n if len(external_user_group_batch) >= _EXTERNAL_GROUP_BATCH_SIZE:\n logger.debug(\n f\"New external user groups: {external_user_group_batch}\"\n )\n upsert_external_groups(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n external_groups=external_user_group_batch,\n source=cc_pair.connector.source,\n )\n external_user_group_batch = []\n\n if external_user_group_batch:\n logger.debug(f\"New external user groups: {external_user_group_batch}\")\n upsert_external_groups(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n external_groups=external_user_group_batch,\n source=cc_pair.connector.source,\n )\n except Exception as e:\n format_error_for_logging(e)\n\n # Mark as failed (this also updates progress to show partial progress)\n mark_external_group_sync_attempt_failed(\n attempt_id, db_session, error_message=str(e)\n )\n\n # TODO: add some notification to the admins here\n logger.exception(\n f\"Error syncing external groups for {source_type} for cc_pair: {cc_pair_id} {e}\"\n )\n raise e\n\n logger.info(\n f\"Removing stale external groups for {source_type} for cc_pair: {cc_pair_id}\"\n )\n remove_stale_external_groups(db_session, cc_pair_id)\n\n # Calculate total unique users processed\n total_users_processed = len(seen_users)\n\n # Complete the sync attempt with final progress\n complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_users_processed=total_users_processed,\n total_groups_processed=total_groups_processed,\n total_group_memberships_synced=total_group_memberships_synced,\n errors_encountered=0,\n )\n logger.info(\n f\"Completed external group sync attempt {attempt_id}: \"\n f\"{total_groups_processed} groups, {total_users_processed} users, \"\n f\"{total_group_memberships_synced} memberships\"\n )\n\n mark_all_relevant_cc_pairs_as_external_group_synced(db_session, cc_pair)\n\n\ndef validate_external_group_sync_fences(\n tenant_id: str,\n celery_app: Celery, # noqa: ARG001\n r: Redis, # noqa: ARG001\n r_replica: Redis,\n r_celery: Redis,\n lock_beat: RedisLock,\n) -> None:\n reserved_tasks = celery_get_unacked_task_ids(\n OnyxCeleryQueues.CONNECTOR_EXTERNAL_GROUP_SYNC, r_celery\n )\n\n # validate all existing external group sync tasks\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n key_str = key_bytes.decode(\"utf-8\")\n if not key_str.startswith(RedisConnectorExternalGroupSync.FENCE_PREFIX):\n continue\n\n validate_external_group_sync_fence(\n tenant_id,\n key_bytes,\n reserved_tasks,\n r_celery,\n )\n\n lock_beat.reacquire()\n return\n\n\ndef validate_external_group_sync_fence(\n tenant_id: str,\n key_bytes: bytes,\n reserved_tasks: set[str],\n r_celery: Redis,\n) -> None:\n \"\"\"Checks for the error condition where an indexing fence is set but the associated celery tasks don't exist.\n This can happen if the indexing worker hard crashes or is terminated.\n Being in this bad state means the fence will never clear without help, so this function\n gives the help.\n\n How this works:\n 1. This function renews the active signal with a 5 minute TTL under the following conditions\n 1.2. When the task is seen in the redis queue\n 1.3. When the task is seen in the reserved / prefetched list\n\n 2. Externally, the active signal is renewed when:\n 2.1. The fence is created\n 2.2. The indexing watchdog checks the spawned task.\n\n 3. The TTL allows us to get through the transitions on fence startup\n and when the task starts executing.\n\n More TTL clarification: it is seemingly impossible to exactly query Celery for\n whether a task is in the queue or currently executing.\n 1. An unknown task id is always returned as state PENDING.\n 2. Redis can be inspected for the task id, but the task id is gone between the time a worker receives the task\n and the time it actually starts on the worker.\n \"\"\"\n # if the fence doesn't exist, there's nothing to do\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n msg = (\n f\"validate_external_group_sync_fence - could not parse id from {fence_key}\"\n )\n emit_background_error(msg)\n task_logger.error(msg)\n return\n\n cc_pair_id = int(cc_pair_id_str)\n\n # parse out metadata and initialize the helper class with it\n redis_connector = RedisConnector(tenant_id, int(cc_pair_id))\n\n # check to see if the fence/payload exists\n if not redis_connector.external_group_sync.fenced:\n return\n\n try:\n payload = redis_connector.external_group_sync.payload\n except ValidationError:\n msg = (\n \"validate_external_group_sync_fence - \"\n \"Resetting fence because fence schema is out of date: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key}\"\n )\n task_logger.exception(msg)\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n\n redis_connector.external_group_sync.reset()\n return\n\n if not payload:\n return\n\n if not payload.celery_task_id:\n return\n\n # OK, there's actually something for us to validate\n found = celery_find_task(\n payload.celery_task_id, OnyxCeleryQueues.CONNECTOR_EXTERNAL_GROUP_SYNC, r_celery\n )\n if found:\n # the celery task exists in the redis queue\n # redis_connector_index.set_active()\n return\n\n if payload.celery_task_id in reserved_tasks:\n # the celery task was prefetched and is reserved within the indexing worker\n # redis_connector_index.set_active()\n return\n\n # we may want to enable this check if using the active task list somehow isn't good enough\n # if redis_connector_index.generator_locked():\n # logger.info(f\"{payload.celery_task_id} is currently executing.\")\n\n # if we get here, we didn't find any direct indication that the associated celery tasks exist,\n # but they still might be there due to gaps in our ability to check states during transitions\n # Checking the active signal safeguards us against these transition periods\n # (which has a duration that allows us to bridge those gaps)\n # if redis_connector_index.active():\n # return\n\n # celery tasks don't exist and the active signal has expired, possibly due to a crash. Clean it up.\n emit_background_error(\n message=(\n \"validate_external_group_sync_fence - \"\n \"Resetting fence because no associated celery tasks were found: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key} \"\n f\"payload_id={payload.id}\"\n ),\n cc_pair_id=cc_pair_id,\n )\n\n redis_connector.external_group_sync.reset()\n return\n" }, { "path": "backend/ee/onyx/background/celery/tasks/hooks/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/hooks/tasks.py", "content": "from celery import shared_task\n\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.hook import cleanup_old_execution_logs__no_commit\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_HOOK_EXECUTION_LOG_RETENTION_DAYS: int = 30\n\n\n@shared_task(\n name=OnyxCeleryTask.HOOK_EXECUTION_LOG_CLEANUP_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n trail=False,\n)\ndef hook_execution_log_cleanup_task(*, tenant_id: str) -> None: # noqa: ARG001\n try:\n with get_session_with_current_tenant() as db_session:\n deleted: int = cleanup_old_execution_logs__no_commit(\n db_session=db_session,\n max_age_days=_HOOK_EXECUTION_LOG_RETENTION_DAYS,\n )\n db_session.commit()\n if deleted:\n logger.info(\n f\"Deleted {deleted} hook execution log(s) older than \"\n f\"{_HOOK_EXECUTION_LOG_RETENTION_DAYS} days.\"\n )\n except Exception:\n logger.exception(\"Failed to clean up hook execution logs\")\n raise\n" }, { "path": "backend/ee/onyx/background/celery/tasks/query_history/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/query_history/tasks.py", "content": "import csv\nimport io\nfrom datetime import datetime\n\nfrom celery import shared_task\nfrom celery import Task\n\nfrom ee.onyx.server.query_history.api import fetch_and_process_chat_session_history\nfrom ee.onyx.server.query_history.api import ONYX_ANONYMIZED_EMAIL\nfrom ee.onyx.server.query_history.models import QuestionAnswerPairSnapshot\nfrom onyx.background.task_utils import construct_query_history_report_name\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.app_configs import ONYX_QUERY_HISTORY_TYPE\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import FileType\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import QueryHistoryType\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.tasks import delete_task_with_id\nfrom onyx.db.tasks import mark_task_as_finished_with_id\nfrom onyx.db.tasks import mark_task_as_started_with_id\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n@shared_task(\n name=OnyxCeleryTask.EXPORT_QUERY_HISTORY_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n trail=False,\n)\ndef export_query_history_task(\n self: Task,\n *,\n start: datetime,\n end: datetime,\n start_time: datetime,\n # Need to include the tenant_id since the TenantAwareTask needs this\n tenant_id: str, # noqa: ARG001\n) -> None:\n if not self.request.id:\n raise RuntimeError(\"No task id defined for this task; cannot identify it\")\n\n task_id = self.request.id\n stream = io.StringIO()\n writer = csv.DictWriter(\n stream,\n fieldnames=list(QuestionAnswerPairSnapshot.model_fields.keys()),\n )\n writer.writeheader()\n\n with get_session_with_current_tenant() as db_session:\n try:\n mark_task_as_started_with_id(\n db_session=db_session,\n task_id=task_id,\n )\n\n snapshot_generator = fetch_and_process_chat_session_history(\n db_session=db_session,\n start=start,\n end=end,\n )\n\n for snapshot in snapshot_generator:\n if ONYX_QUERY_HISTORY_TYPE == QueryHistoryType.ANONYMIZED:\n snapshot.user_email = ONYX_ANONYMIZED_EMAIL\n\n writer.writerows(\n qa_pair.to_json()\n for qa_pair in QuestionAnswerPairSnapshot.from_chat_session_snapshot(\n snapshot\n )\n )\n\n except Exception:\n logger.exception(f\"Failed to export query history with {task_id=}\")\n mark_task_as_finished_with_id(\n db_session=db_session,\n task_id=task_id,\n success=False,\n )\n raise\n\n report_name = construct_query_history_report_name(task_id)\n with get_session_with_current_tenant() as db_session:\n try:\n stream.seek(0)\n get_default_file_store().save_file(\n content=stream,\n display_name=report_name,\n file_origin=FileOrigin.QUERY_HISTORY_CSV,\n file_type=FileType.CSV,\n file_metadata={\n \"start\": start.isoformat(),\n \"end\": end.isoformat(),\n \"start_time\": start_time.isoformat(),\n },\n file_id=report_name,\n )\n\n delete_task_with_id(\n db_session=db_session,\n task_id=task_id,\n )\n except Exception:\n logger.exception(\n f\"Failed to save query history export file; {report_name=}\"\n )\n mark_task_as_finished_with_id(\n db_session=db_session,\n task_id=task_id,\n success=False,\n )\n raise\n" }, { "path": "backend/ee/onyx/background/celery/tasks/tenant_provisioning/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py", "content": "\"\"\"\nPeriodic tasks for tenant pre-provisioning.\n\"\"\"\n\nimport asyncio\nimport datetime\nimport uuid\n\nfrom celery import shared_task\nfrom celery import Task\nfrom redis.lock import Lock as RedisLock\n\nfrom ee.onyx.server.tenants.provisioning import setup_tenant\nfrom ee.onyx.server.tenants.schema_management import create_schema_if_not_exists\nfrom ee.onyx.server.tenants.schema_management import get_current_alembic_version\nfrom ee.onyx.server.tenants.schema_management import run_alembic_migrations\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.configs.app_configs import TARGET_AVAILABLE_TENANTS\nfrom onyx.configs.constants import ONYX_CLOUD_TENANT_ID\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.models import AvailableTenant\nfrom onyx.redis.redis_pool import get_redis_client\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import TENANT_ID_PREFIX\n\n# Maximum tenants to provision in a single task run.\n# Each tenant takes ~80s (alembic migrations), so 5 tenants ≈ 7 minutes.\n_MAX_TENANTS_PER_RUN = 5\n\n# Time limits sized for worst-case: provisioning up to _MAX_TENANTS_PER_RUN new tenants\n# (~90s each) plus migrating up to TARGET_AVAILABLE_TENANTS pool tenants (~90s each).\n_TENANT_PROVISIONING_SOFT_TIME_LIMIT = 60 * 20 # 20 minutes\n_TENANT_PROVISIONING_TIME_LIMIT = 60 * 25 # 25 minutes\n\n\n@shared_task(\n name=OnyxCeleryTask.CLOUD_CHECK_AVAILABLE_TENANTS,\n queue=OnyxCeleryQueues.MONITORING,\n ignore_result=True,\n soft_time_limit=_TENANT_PROVISIONING_SOFT_TIME_LIMIT,\n time_limit=_TENANT_PROVISIONING_TIME_LIMIT,\n trail=False,\n bind=True,\n)\ndef check_available_tenants(self: Task) -> None: # noqa: ARG001\n \"\"\"\n Check if we have enough pre-provisioned tenants available.\n If not, trigger the pre-provisioning of new tenants.\n \"\"\"\n task_logger.info(\"STARTING CHECK_AVAILABLE_TENANTS\")\n if not MULTI_TENANT:\n task_logger.info(\n \"Multi-tenancy is not enabled, skipping tenant pre-provisioning\"\n )\n return\n\n r = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n lock_check: RedisLock = r.lock(\n OnyxRedisLocks.CHECK_AVAILABLE_TENANTS_LOCK,\n timeout=_TENANT_PROVISIONING_TIME_LIMIT,\n )\n\n # These tasks should never overlap\n if not lock_check.acquire(blocking=False):\n task_logger.info(\n \"Skipping check_available_tenants task because it is already running\"\n )\n return\n\n try:\n # Get the current count of available tenants\n with get_session_with_shared_schema() as db_session:\n num_available_tenants = db_session.query(AvailableTenant).count()\n\n # Get the target number of available tenants\n num_minimum_available_tenants = TARGET_AVAILABLE_TENANTS\n\n # Calculate how many new tenants we need to provision\n if num_available_tenants < num_minimum_available_tenants:\n tenants_to_provision = num_minimum_available_tenants - num_available_tenants\n else:\n tenants_to_provision = 0\n\n task_logger.info(\n f\"Available tenants: {num_available_tenants}, \"\n f\"Target minimum available tenants: {num_minimum_available_tenants}, \"\n f\"To provision: {tenants_to_provision}\"\n )\n\n batch_size = min(tenants_to_provision, _MAX_TENANTS_PER_RUN)\n if batch_size < tenants_to_provision:\n task_logger.info(\n f\"Capping batch to {batch_size} (need {tenants_to_provision}, will catch up next cycle)\"\n )\n\n provisioned = 0\n for i in range(batch_size):\n task_logger.info(f\"Provisioning tenant {i + 1}/{batch_size}\")\n try:\n if pre_provision_tenant():\n provisioned += 1\n except Exception:\n task_logger.exception(\n f\"Failed to provision tenant {i + 1}/{batch_size}, continuing with remaining tenants\"\n )\n\n task_logger.info(f\"Provisioning complete: {provisioned}/{batch_size} succeeded\")\n\n # Migrate any pool tenants that were provisioned before a new migration was deployed\n _migrate_stale_pool_tenants()\n\n except Exception:\n task_logger.exception(\"Error in check_available_tenants task\")\n\n finally:\n try:\n lock_check.release()\n except Exception:\n task_logger.warning(\n \"Could not release check lock (likely expired), continuing\"\n )\n\n\ndef _migrate_stale_pool_tenants() -> None:\n \"\"\"\n Run alembic upgrade head on all pool tenants. Since alembic upgrade head is\n idempotent, tenants already at head are a fast no-op. This ensures pool\n tenants are always current so that signup doesn't hit schema mismatches\n (e.g. missing columns added after the tenant was pre-provisioned).\n \"\"\"\n with get_session_with_shared_schema() as db_session:\n pool_tenants = db_session.query(AvailableTenant).all()\n tenant_ids = [t.tenant_id for t in pool_tenants]\n\n if not tenant_ids:\n return\n\n task_logger.info(\n f\"Checking {len(tenant_ids)} pool tenant(s) for pending migrations\"\n )\n\n for tenant_id in tenant_ids:\n try:\n run_alembic_migrations(tenant_id)\n new_version = get_current_alembic_version(tenant_id)\n with get_session_with_shared_schema() as db_session:\n tenant = (\n db_session.query(AvailableTenant)\n .filter_by(tenant_id=tenant_id)\n .first()\n )\n if tenant and tenant.alembic_version != new_version:\n task_logger.info(\n f\"Migrated pool tenant {tenant_id}: {tenant.alembic_version} -> {new_version}\"\n )\n tenant.alembic_version = new_version\n db_session.commit()\n except Exception:\n task_logger.exception(\n f\"Failed to migrate pool tenant {tenant_id}, skipping\"\n )\n\n\ndef pre_provision_tenant() -> bool:\n \"\"\"\n Pre-provision a new tenant and store it in the NewAvailableTenant table.\n This function fully sets up the tenant with all necessary configurations,\n so it's ready to be assigned to a user immediately.\n\n Returns True if a tenant was successfully provisioned, False otherwise.\n \"\"\"\n # The MULTI_TENANT check is now done at the caller level (check_available_tenants)\n # rather than inside this function\n\n r = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n lock_provision: RedisLock = r.lock(\n OnyxRedisLocks.CLOUD_PRE_PROVISION_TENANT_LOCK,\n timeout=_TENANT_PROVISIONING_TIME_LIMIT,\n )\n\n # Allow multiple pre-provisioning tasks to run, but ensure they don't overlap\n if not lock_provision.acquire(blocking=False):\n task_logger.warning(\n \"Skipping pre_provision_tenant — could not acquire provision lock\"\n )\n return False\n\n tenant_id: str | None = None\n try:\n # Generate a new tenant ID\n tenant_id = TENANT_ID_PREFIX + str(uuid.uuid4())\n task_logger.info(f\"Pre-provisioning tenant: {tenant_id}\")\n\n # Create the schema for the new tenant\n schema_created = create_schema_if_not_exists(tenant_id)\n if schema_created:\n task_logger.debug(f\"Created schema for tenant: {tenant_id}\")\n else:\n task_logger.debug(f\"Schema already exists for tenant: {tenant_id}\")\n\n # Set up the tenant with all necessary configurations\n task_logger.debug(f\"Setting up tenant configuration: {tenant_id}\")\n asyncio.run(setup_tenant(tenant_id))\n task_logger.debug(f\"Tenant configuration completed: {tenant_id}\")\n\n # Get the current Alembic version\n alembic_version = get_current_alembic_version(tenant_id)\n task_logger.debug(\n f\"Tenant {tenant_id} using Alembic version: {alembic_version}\"\n )\n\n # Store the pre-provisioned tenant in the database\n task_logger.debug(f\"Storing pre-provisioned tenant in database: {tenant_id}\")\n with get_session_with_shared_schema() as db_session:\n # Use a transaction to ensure atomicity\n db_session.begin()\n try:\n new_tenant = AvailableTenant(\n tenant_id=tenant_id,\n alembic_version=alembic_version,\n date_created=datetime.datetime.now(),\n )\n db_session.add(new_tenant)\n db_session.commit()\n task_logger.info(f\"Successfully pre-provisioned tenant: {tenant_id}\")\n return True\n except Exception:\n db_session.rollback()\n task_logger.error(\n f\"Failed to store pre-provisioned tenant: {tenant_id}\",\n exc_info=True,\n )\n raise\n\n except Exception:\n task_logger.error(\"Error in pre_provision_tenant task\", exc_info=True)\n # If we have a tenant_id, attempt to rollback any partially completed provisioning\n if tenant_id:\n task_logger.info(\n f\"Rolling back failed tenant provisioning for: {tenant_id}\"\n )\n try:\n from ee.onyx.server.tenants.provisioning import (\n rollback_tenant_provisioning,\n )\n\n asyncio.run(rollback_tenant_provisioning(tenant_id))\n except Exception:\n task_logger.exception(f\"Error during rollback for tenant: {tenant_id}\")\n return False\n finally:\n try:\n lock_provision.release()\n except Exception:\n task_logger.warning(\n \"Could not release provision lock (likely expired), continuing\"\n )\n" }, { "path": "backend/ee/onyx/background/celery/tasks/ttl_management/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/ttl_management/tasks.py", "content": "from uuid import UUID\n\nfrom celery import shared_task\nfrom celery import Task\n\nfrom ee.onyx.background.celery_utils import should_perform_chat_ttl_check\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.chat import delete_chat_session\nfrom onyx.db.chat import get_chat_sessions_older_than\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.server.settings.store import load_settings\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n@shared_task(\n name=OnyxCeleryTask.PERFORM_TTL_MANAGEMENT_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n trail=False,\n)\ndef perform_ttl_management_task(\n self: Task, retention_limit_days: int, *, tenant_id: str # noqa: ARG001\n) -> None:\n task_id = self.request.id\n if not task_id:\n raise RuntimeError(\"No task id defined for this task; cannot identify it\")\n\n user_id: UUID | None = None\n session_id: UUID | None = None\n try:\n with get_session_with_current_tenant() as db_session:\n\n old_chat_sessions = get_chat_sessions_older_than(\n retention_limit_days, db_session\n )\n\n for user_id, session_id in old_chat_sessions:\n # one session per delete so that we don't blow up if a deletion fails.\n with get_session_with_current_tenant() as db_session:\n delete_chat_session(\n user_id,\n session_id,\n db_session,\n include_deleted=True,\n hard_delete=True,\n )\n\n except Exception:\n logger.exception(\n f\"delete_chat_session exceptioned. user_id={user_id} session_id={session_id}\"\n )\n raise\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_TTL_MANAGEMENT_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n)\ndef check_ttl_management_task(*, tenant_id: str) -> None:\n \"\"\"Runs periodically to check if any ttl tasks should be run and adds them\n to the queue\"\"\"\n\n settings = load_settings()\n retention_limit_days = settings.maximum_chat_retention_days\n with get_session_with_current_tenant() as db_session:\n if should_perform_chat_ttl_check(retention_limit_days, db_session):\n perform_ttl_management_task.apply_async(\n kwargs=dict(\n retention_limit_days=retention_limit_days, tenant_id=tenant_id\n ),\n )\n" }, { "path": "backend/ee/onyx/background/celery/tasks/usage_reporting/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py", "content": "from datetime import datetime\nfrom uuid import UUID\n\nfrom celery import shared_task\nfrom celery import Task\n\nfrom ee.onyx.server.reporting.usage_export_generation import create_new_usage_report\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n@shared_task(\n name=OnyxCeleryTask.GENERATE_USAGE_REPORT_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n trail=False,\n)\ndef generate_usage_report_task(\n self: Task, # noqa: ARG001\n *,\n tenant_id: str, # noqa: ARG001\n user_id: str | None = None,\n period_from: str | None = None,\n period_to: str | None = None,\n) -> None:\n \"\"\"User-initiated usage report generation task\"\"\"\n # Parse period if provided\n period = None\n if period_from and period_to:\n period = (\n datetime.fromisoformat(period_from),\n datetime.fromisoformat(period_to),\n )\n\n # Generate the report\n with get_session_with_current_tenant() as db_session:\n create_new_usage_report(\n db_session=db_session,\n user_id=UUID(user_id) if user_id else None,\n period=period,\n )\n" }, { "path": "backend/ee/onyx/background/celery/tasks/vespa/__init__.py", "content": "" }, { "path": "backend/ee/onyx/background/celery/tasks/vespa/tasks.py", "content": "from typing import cast\n\nfrom redis import Redis\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.user_group import delete_user_group\nfrom ee.onyx.db.user_group import fetch_user_group\nfrom ee.onyx.db.user_group import mark_user_group_as_synced\nfrom ee.onyx.db.user_group import prepare_user_group_for_deletion\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.sync_record import update_sync_record_status\nfrom onyx.redis.redis_usergroup import RedisUserGroup\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef monitor_usergroup_taskset(\n tenant_id: str, key_bytes: bytes, r: Redis, db_session: Session\n) -> None:\n \"\"\"This function is likely to move in the worker refactor happening next.\"\"\"\n fence_key = key_bytes.decode(\"utf-8\")\n usergroup_id_str = RedisUserGroup.get_id_from_fence_key(fence_key)\n if not usergroup_id_str:\n task_logger.warning(f\"Could not parse usergroup id from {fence_key}\")\n return\n\n try:\n usergroup_id = int(usergroup_id_str)\n except ValueError:\n task_logger.exception(f\"usergroup_id ({usergroup_id_str}) is not an integer!\")\n raise\n\n rug = RedisUserGroup(tenant_id, usergroup_id)\n if not rug.fenced:\n return\n\n initial_count = rug.payload\n if initial_count is None:\n return\n\n count = cast(int, r.scard(rug.taskset_key))\n task_logger.info(\n f\"User group sync progress: usergroup_id={usergroup_id} remaining={count} initial={initial_count}\"\n )\n if count > 0:\n update_sync_record_status(\n db_session=db_session,\n entity_id=usergroup_id,\n sync_type=SyncType.USER_GROUP,\n sync_status=SyncStatus.IN_PROGRESS,\n num_docs_synced=count,\n )\n return\n\n user_group = fetch_user_group(db_session=db_session, user_group_id=usergroup_id)\n if user_group:\n usergroup_name = user_group.name\n try:\n if user_group.is_up_for_deletion:\n # this prepare should have been run when the deletion was scheduled,\n # but run it again to be sure we're ready to go\n mark_user_group_as_synced(db_session, user_group)\n prepare_user_group_for_deletion(db_session, usergroup_id)\n delete_user_group(db_session=db_session, user_group=user_group)\n\n update_sync_record_status(\n db_session=db_session,\n entity_id=usergroup_id,\n sync_type=SyncType.USER_GROUP,\n sync_status=SyncStatus.SUCCESS,\n num_docs_synced=initial_count,\n )\n\n task_logger.info(\n f\"Deleted usergroup: name={usergroup_name} id={usergroup_id}\"\n )\n else:\n mark_user_group_as_synced(db_session=db_session, user_group=user_group)\n\n update_sync_record_status(\n db_session=db_session,\n entity_id=usergroup_id,\n sync_type=SyncType.USER_GROUP,\n sync_status=SyncStatus.SUCCESS,\n num_docs_synced=initial_count,\n )\n\n task_logger.info(\n f\"Synced usergroup. name={usergroup_name} id={usergroup_id}\"\n )\n except Exception as e:\n update_sync_record_status(\n db_session=db_session,\n entity_id=usergroup_id,\n sync_type=SyncType.USER_GROUP,\n sync_status=SyncStatus.FAILED,\n num_docs_synced=initial_count,\n )\n raise e\n\n rug.reset()\n" }, { "path": "backend/ee/onyx/background/celery_utils.py", "content": "from sqlalchemy.orm import Session\n\nfrom ee.onyx.background.task_name_builders import name_chat_ttl_task\nfrom onyx.db.tasks import check_task_is_live_and_not_timed_out\nfrom onyx.db.tasks import get_latest_task\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef should_perform_chat_ttl_check(\n retention_limit_days: float | None, db_session: Session\n) -> bool:\n # TODO: make this a check for None and add behavior for 0 day TTL\n if not retention_limit_days:\n return False\n\n task_name = name_chat_ttl_task(retention_limit_days)\n latest_task = get_latest_task(task_name, db_session)\n if not latest_task:\n return True\n\n if check_task_is_live_and_not_timed_out(latest_task, db_session):\n logger.debug(f\"{task_name} is already being performed. Skipping.\")\n return False\n return True\n" }, { "path": "backend/ee/onyx/background/task_name_builders.py", "content": "from datetime import datetime\n\nfrom onyx.configs.constants import OnyxCeleryTask\n\n\nQUERY_HISTORY_TASK_NAME_PREFIX = OnyxCeleryTask.EXPORT_QUERY_HISTORY_TASK\n\n\ndef name_chat_ttl_task(\n retention_limit_days: float,\n tenant_id: str | None = None, # noqa: ARG001\n) -> str:\n return f\"chat_ttl_{retention_limit_days}_days\"\n\n\ndef query_history_task_name(start: datetime, end: datetime) -> str:\n return f\"{QUERY_HISTORY_TASK_NAME_PREFIX}_{start}_{end}\"\n" }, { "path": "backend/ee/onyx/configs/__init__.py", "content": "" }, { "path": "backend/ee/onyx/configs/app_configs.py", "content": "import json\nimport os\n\n\n#####\n# Auto Permission Sync\n#####\n# should generally only be used for sources that support polling of permissions\n# e.g. can pull in only permission changes rather than having to go through all\n# documents every time\nDEFAULT_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY\") or 5 * 60\n)\n\n\n#####\n# Confluence\n#####\n\n# In seconds, default is 30 minutes\nCONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY = int(\n os.environ.get(\"CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY\") or 30 * 60\n)\n# In seconds, default is 30 minutes\nCONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY\") or 30 * 60\n)\n# This is a boolean that determines if anonymous access is public\n# Default behavior is to not make the page public and instead add a group\n# that contains all the users that we found in Confluence\nCONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC = (\n os.environ.get(\"CONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC\", \"\").lower() == \"true\"\n)\n\n\n#####\n# JIRA\n#####\n\n# In seconds, default is 30 minutes\nJIRA_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"JIRA_PERMISSION_DOC_SYNC_FREQUENCY\") or 30 * 60\n)\n# In seconds, default is 30 minutes\nJIRA_PERMISSION_GROUP_SYNC_FREQUENCY = int(\n os.environ.get(\"JIRA_PERMISSION_GROUP_SYNC_FREQUENCY\") or 30 * 60\n)\n\n\n#####\n# Google Drive\n#####\nGOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY = int(\n os.environ.get(\"GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY\") or 5 * 60\n)\n\n\n#####\n# GitHub\n#####\n# In seconds, default is 5 minutes\nGITHUB_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"GITHUB_PERMISSION_DOC_SYNC_FREQUENCY\") or 5 * 60\n)\n# In seconds, default is 5 minutes\nGITHUB_PERMISSION_GROUP_SYNC_FREQUENCY = int(\n os.environ.get(\"GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY\") or 5 * 60\n)\n\n\n#####\n# Slack\n#####\nSLACK_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"SLACK_PERMISSION_DOC_SYNC_FREQUENCY\") or 5 * 60\n)\n\nNUM_PERMISSION_WORKERS = int(os.environ.get(\"NUM_PERMISSION_WORKERS\") or 2)\n\n\n#####\n# Teams\n#####\n# In seconds, default is 5 minutes\nTEAMS_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"TEAMS_PERMISSION_DOC_SYNC_FREQUENCY\") or 5 * 60\n)\n\n#####\n# SharePoint\n#####\n# In seconds, default is 30 minutes\nSHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY = int(\n os.environ.get(\"SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY\") or 30 * 60\n)\n\n# In seconds, default is 5 minutes\nSHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY = int(\n os.environ.get(\"SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY\") or 5 * 60\n)\n\n\n####\n# Celery Job Frequency\n####\nCHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS = float(\n os.environ.get(\"CHECK_TTL_MANAGEMENT_TASK_FREQUENCY_IN_HOURS\") or 1\n) # float for easier testing\n\n\nSTRIPE_SECRET_KEY = os.environ.get(\"STRIPE_SECRET_KEY\")\n\n# JWT Public Key URL\nJWT_PUBLIC_KEY_URL: str | None = os.getenv(\"JWT_PUBLIC_KEY_URL\", None)\n\n\n# Super Users\nSUPER_USERS = json.loads(os.environ.get(\"SUPER_USERS\", \"[]\"))\nSUPER_CLOUD_API_KEY = os.environ.get(\"SUPER_CLOUD_API_KEY\", \"api_key\")\n\nPOSTHOG_API_KEY = os.environ.get(\"POSTHOG_API_KEY\")\nPOSTHOG_HOST = os.environ.get(\"POSTHOG_HOST\") or \"https://us.i.posthog.com\"\nPOSTHOG_DEBUG_LOGS_ENABLED = (\n os.environ.get(\"POSTHOG_DEBUG_LOGS_ENABLED\", \"\").lower() == \"true\"\n)\n\nMARKETING_POSTHOG_API_KEY = os.environ.get(\"MARKETING_POSTHOG_API_KEY\")\n\nHUBSPOT_TRACKING_URL = os.environ.get(\"HUBSPOT_TRACKING_URL\")\n\nGATED_TENANTS_KEY = \"gated_tenants\"\n\n# License enforcement - when True, blocks API access for gated/expired licenses\nLICENSE_ENFORCEMENT_ENABLED = (\n os.environ.get(\"LICENSE_ENFORCEMENT_ENABLED\", \"true\").lower() == \"true\"\n)\n\n# Cloud data plane URL - self-hosted instances call this to reach cloud proxy endpoints\n# Used when MULTI_TENANT=false (self-hosted mode)\nCLOUD_DATA_PLANE_URL = os.environ.get(\n \"CLOUD_DATA_PLANE_URL\", \"https://cloud.onyx.app/api\"\n)\n" }, { "path": "backend/ee/onyx/configs/license_enforcement_config.py", "content": "\"\"\"Constants for license enforcement.\n\nThis file is the single source of truth for:\n1. Paths that bypass license enforcement (always accessible)\n2. Paths that require an EE license (EE-only features)\n\nImport these constants in both production code and tests to ensure consistency.\n\"\"\"\n\n# Paths that are ALWAYS accessible, even when license is expired/gated.\n# These enable users to:\n# /auth - Log in/out (users can't fix billing if locked out of auth)\n# /license - Fetch, upload, or check license status\n# /health - Health checks for load balancers/orchestrators\n# /me - Basic user info needed for UI rendering\n# /settings, /enterprise-settings - View app status and branding\n# /billing - Unified billing API\n# /proxy - Self-hosted proxy endpoints (have own license-based auth)\n# /tenants/billing-* - Legacy billing endpoints (backwards compatibility)\n# /manage/users, /users - User management (needed for seat limit resolution)\n# /notifications - Needed for UI to load properly\nLICENSE_ENFORCEMENT_ALLOWED_PREFIXES: frozenset[str] = frozenset(\n {\n \"/auth\",\n \"/license\",\n \"/health\",\n \"/me\",\n \"/settings\",\n \"/enterprise-settings\",\n # Billing endpoints (unified API for both MT and self-hosted)\n \"/billing\",\n \"/admin/billing\",\n # Proxy endpoints for self-hosted billing (no tenant context)\n \"/proxy\",\n # Legacy tenant billing endpoints (kept for backwards compatibility)\n \"/tenants/billing-information\",\n \"/tenants/create-customer-portal-session\",\n \"/tenants/create-subscription-session\",\n # User management - needed to remove users when seat limit exceeded\n \"/manage/users\",\n \"/manage/admin/users\",\n \"/manage/admin/valid-domains\",\n \"/manage/admin/deactivate-user\",\n \"/manage/admin/delete-user\",\n \"/users\",\n # Notifications - needed for UI to load properly\n \"/notifications\",\n }\n)\n\n# EE-only paths that require a valid license.\n# Users without a license (community edition) cannot access these.\n# These are blocked even when user has never subscribed (no license).\nEE_ONLY_PATH_PREFIXES: frozenset[str] = frozenset(\n {\n # User groups and access control\n \"/manage/admin/user-group\",\n # Analytics and reporting\n \"/analytics\",\n # Query history (admin chat session endpoints)\n \"/admin/chat-sessions\",\n \"/admin/chat-session-history\",\n \"/admin/query-history\",\n # Usage reporting/export\n \"/admin/usage-report\",\n # Standard answers (canned responses)\n \"/manage/admin/standard-answer\",\n # Token rate limits\n \"/admin/token-rate-limits\",\n # Evals\n \"/evals\",\n # Hook extensions\n \"/admin/hooks\",\n }\n)\n" }, { "path": "backend/ee/onyx/connectors/perm_sync_valid.py", "content": "from onyx.connectors.confluence.connector import ConfluenceConnector\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.interfaces import BaseConnector\n\n\ndef validate_confluence_perm_sync(connector: ConfluenceConnector) -> None:\n \"\"\"\n Validate that the connector is configured correctly for permissions syncing.\n \"\"\"\n\n\ndef validate_drive_perm_sync(connector: GoogleDriveConnector) -> None:\n \"\"\"\n Validate that the connector is configured correctly for permissions syncing.\n \"\"\"\n\n\ndef validate_perm_sync(connector: BaseConnector) -> None:\n \"\"\"\n Override this if your connector needs to validate permissions syncing.\n Raise an exception if invalid, otherwise do nothing.\n\n Default is a no-op (always successful).\n \"\"\"\n if isinstance(connector, ConfluenceConnector):\n validate_confluence_perm_sync(connector)\n elif isinstance(connector, GoogleDriveConnector):\n validate_drive_perm_sync(connector)\n" }, { "path": "backend/ee/onyx/db/__init__.py", "content": "" }, { "path": "backend/ee/onyx/db/analytics.py", "content": "import datetime\nfrom collections.abc import Sequence\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import case\nfrom sqlalchemy import cast\nfrom sqlalchemy import Date\nfrom sqlalchemy import func\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatMessageFeedback\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import Persona\nfrom onyx.db.models import User\nfrom onyx.db.models import UserRole\n\n\ndef fetch_query_analytics(\n start: datetime.datetime,\n end: datetime.datetime,\n db_session: Session,\n) -> Sequence[tuple[int, int, int, datetime.date]]:\n stmt = (\n select(\n func.count(ChatMessage.id),\n func.sum(case((ChatMessageFeedback.is_positive, 1), else_=0)),\n func.sum(\n case(\n (ChatMessageFeedback.is_positive == False, 1), # noqa: E712\n else_=0, # noqa: E712\n )\n ),\n cast(ChatMessage.time_sent, Date),\n )\n .join(\n ChatMessageFeedback,\n ChatMessageFeedback.chat_message_id == ChatMessage.id,\n isouter=True,\n )\n .where(\n ChatMessage.time_sent >= start,\n )\n .where(\n ChatMessage.time_sent <= end,\n )\n .where(ChatMessage.message_type == MessageType.ASSISTANT)\n .group_by(cast(ChatMessage.time_sent, Date))\n .order_by(cast(ChatMessage.time_sent, Date))\n )\n\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef fetch_per_user_query_analytics(\n start: datetime.datetime,\n end: datetime.datetime,\n db_session: Session,\n) -> Sequence[tuple[int, int, int, datetime.date, UUID]]:\n stmt = (\n select(\n func.count(ChatMessage.id),\n func.sum(case((ChatMessageFeedback.is_positive, 1), else_=0)),\n func.sum(\n case(\n (ChatMessageFeedback.is_positive == False, 1), # noqa: E712\n else_=0, # noqa: E712\n )\n ),\n cast(ChatMessage.time_sent, Date),\n ChatSession.user_id,\n )\n .join(ChatSession, ChatSession.id == ChatMessage.chat_session_id)\n # Include chats that have no explicit feedback instead of dropping them\n .join(\n ChatMessageFeedback,\n ChatMessageFeedback.chat_message_id == ChatMessage.id,\n isouter=True,\n )\n .where(\n ChatMessage.time_sent >= start,\n )\n .where(\n ChatMessage.time_sent <= end,\n )\n .where(ChatMessage.message_type == MessageType.ASSISTANT)\n .group_by(cast(ChatMessage.time_sent, Date), ChatSession.user_id)\n .order_by(cast(ChatMessage.time_sent, Date), ChatSession.user_id)\n )\n\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef fetch_onyxbot_analytics(\n start: datetime.datetime,\n end: datetime.datetime,\n db_session: Session,\n) -> Sequence[tuple[int, int, datetime.date]]:\n \"\"\"Gets the:\n Date of each set of aggregated statistics\n Number of OnyxBot Queries (Chat Sessions)\n Number of instances of Negative feedback OR Needing additional help\n (only counting the last feedback)\n \"\"\"\n # Get every chat session in the time range which is a Onyxbot flow\n # along with the first Assistant message which is the response to the user question.\n # Generally there should not be more than one AI message per chat session of this type\n subquery_first_ai_response = (\n db_session.query(\n ChatMessage.chat_session_id.label(\"chat_session_id\"),\n func.min(ChatMessage.id).label(\"chat_message_id\"),\n )\n .join(ChatSession, ChatSession.id == ChatMessage.chat_session_id)\n .where(\n ChatSession.time_created >= start,\n ChatSession.time_created <= end,\n ChatSession.onyxbot_flow.is_(True),\n )\n .where(\n ChatMessage.message_type == MessageType.ASSISTANT,\n )\n .group_by(ChatMessage.chat_session_id)\n .subquery()\n )\n\n # Get the chat message ids and most recent feedback for each of those chat messages,\n # not including the messages that have no feedback\n subquery_last_feedback = (\n db_session.query(\n ChatMessageFeedback.chat_message_id.label(\"chat_message_id\"),\n func.max(ChatMessageFeedback.id).label(\"max_feedback_id\"),\n )\n .group_by(ChatMessageFeedback.chat_message_id)\n .subquery()\n )\n\n results = (\n db_session.query(\n func.count(ChatSession.id).label(\"total_sessions\"),\n # Need to explicitly specify this as False to handle the NULL case so the cases without\n # feedback aren't counted against Onyxbot\n func.sum(\n case(\n (\n or_(\n ChatMessageFeedback.is_positive.is_(False),\n ChatMessageFeedback.required_followup.is_(True),\n ),\n 1,\n ),\n else_=0,\n )\n ).label(\"negative_answer\"),\n cast(ChatSession.time_created, Date).label(\"session_date\"),\n )\n .join(\n subquery_first_ai_response,\n ChatSession.id == subquery_first_ai_response.c.chat_session_id,\n )\n # Combine the chat sessions with latest feedback to get the latest feedback for the first AI\n # message of the chat session where the chat session is Onyxbot type and within the time\n # range specified. Left/outer join used here to ensure that if no feedback, a null is used\n # for the feedback id\n .outerjoin(\n subquery_last_feedback,\n subquery_first_ai_response.c.chat_message_id\n == subquery_last_feedback.c.chat_message_id,\n )\n # Join the actual feedback table to get the feedback info for the sums\n # Outer join because the \"last feedback\" may be null\n .outerjoin(\n ChatMessageFeedback,\n ChatMessageFeedback.id == subquery_last_feedback.c.max_feedback_id,\n )\n .group_by(cast(ChatSession.time_created, Date))\n .order_by(cast(ChatSession.time_created, Date))\n .all()\n )\n\n return [tuple(row) for row in results]\n\n\ndef fetch_persona_message_analytics(\n db_session: Session,\n persona_id: int,\n start: datetime.datetime,\n end: datetime.datetime,\n) -> list[tuple[int, datetime.date]]:\n \"\"\"Gets the daily message counts for a specific persona within the given time range.\"\"\"\n query = (\n select(\n func.count(ChatMessage.id),\n cast(ChatMessage.time_sent, Date),\n )\n .join(\n ChatSession,\n ChatMessage.chat_session_id == ChatSession.id,\n )\n .where(\n ChatSession.persona_id == persona_id,\n ChatMessage.time_sent >= start,\n ChatMessage.time_sent <= end,\n ChatMessage.message_type == MessageType.ASSISTANT,\n )\n .group_by(cast(ChatMessage.time_sent, Date))\n .order_by(cast(ChatMessage.time_sent, Date))\n )\n\n return [tuple(row) for row in db_session.execute(query).all()]\n\n\ndef fetch_persona_unique_users(\n db_session: Session,\n persona_id: int,\n start: datetime.datetime,\n end: datetime.datetime,\n) -> list[tuple[int, datetime.date]]:\n \"\"\"Gets the daily unique user counts for a specific persona within the given time range.\"\"\"\n query = (\n select(\n func.count(func.distinct(ChatSession.user_id)),\n cast(ChatMessage.time_sent, Date),\n )\n .join(\n ChatSession,\n ChatMessage.chat_session_id == ChatSession.id,\n )\n .where(\n ChatSession.persona_id == persona_id,\n ChatMessage.time_sent >= start,\n ChatMessage.time_sent <= end,\n ChatMessage.message_type == MessageType.ASSISTANT,\n )\n .group_by(cast(ChatMessage.time_sent, Date))\n .order_by(cast(ChatMessage.time_sent, Date))\n )\n\n return [tuple(row) for row in db_session.execute(query).all()]\n\n\ndef fetch_assistant_message_analytics(\n db_session: Session,\n assistant_id: int,\n start: datetime.datetime,\n end: datetime.datetime,\n) -> list[tuple[int, datetime.date]]:\n \"\"\"\n Gets the daily message counts for a specific assistant in the given time range.\n \"\"\"\n query = (\n select(\n func.count(ChatMessage.id),\n cast(ChatMessage.time_sent, Date),\n )\n .join(\n ChatSession,\n ChatMessage.chat_session_id == ChatSession.id,\n )\n .where(\n ChatSession.persona_id == assistant_id,\n ChatMessage.time_sent >= start,\n ChatMessage.time_sent <= end,\n ChatMessage.message_type == MessageType.ASSISTANT,\n )\n .group_by(cast(ChatMessage.time_sent, Date))\n .order_by(cast(ChatMessage.time_sent, Date))\n )\n\n return [tuple(row) for row in db_session.execute(query).all()]\n\n\ndef fetch_assistant_unique_users(\n db_session: Session,\n assistant_id: int,\n start: datetime.datetime,\n end: datetime.datetime,\n) -> list[tuple[int, datetime.date]]:\n \"\"\"\n Gets the daily unique user counts for a specific assistant in the given time range.\n \"\"\"\n query = (\n select(\n func.count(func.distinct(ChatSession.user_id)),\n cast(ChatMessage.time_sent, Date),\n )\n .join(\n ChatSession,\n ChatMessage.chat_session_id == ChatSession.id,\n )\n .where(\n ChatSession.persona_id == assistant_id,\n ChatMessage.time_sent >= start,\n ChatMessage.time_sent <= end,\n ChatMessage.message_type == MessageType.ASSISTANT,\n )\n .group_by(cast(ChatMessage.time_sent, Date))\n .order_by(cast(ChatMessage.time_sent, Date))\n )\n\n return [tuple(row) for row in db_session.execute(query).all()]\n\n\ndef fetch_assistant_unique_users_total(\n db_session: Session,\n assistant_id: int,\n start: datetime.datetime,\n end: datetime.datetime,\n) -> int:\n \"\"\"\n Gets the total number of distinct users who have sent or received messages from\n the specified assistant in the given time range.\n \"\"\"\n query = (\n select(func.count(func.distinct(ChatSession.user_id)))\n .select_from(ChatMessage)\n .join(\n ChatSession,\n ChatMessage.chat_session_id == ChatSession.id,\n )\n .where(\n ChatSession.persona_id == assistant_id,\n ChatMessage.time_sent >= start,\n ChatMessage.time_sent <= end,\n ChatMessage.message_type == MessageType.ASSISTANT,\n )\n )\n\n result = db_session.execute(query).scalar()\n return result if result else 0\n\n\n# Users can view assistant stats if they created the persona,\n# or if they are an admin\ndef user_can_view_assistant_stats(\n db_session: Session, user: User, assistant_id: int\n) -> bool:\n if user.role == UserRole.ADMIN:\n return True\n\n # Check if the user created the persona\n stmt = select(Persona).where(\n and_(Persona.id == assistant_id, Persona.user_id == user.id)\n )\n\n persona = db_session.execute(stmt).scalar_one_or_none()\n return persona is not None\n" }, { "path": "backend/ee/onyx/db/connector.py", "content": "from sqlalchemy import distinct\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import Connector\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef fetch_sources_with_connectors(db_session: Session) -> list[DocumentSource]:\n sources = db_session.query(distinct(Connector.source)).all() # type: ignore\n\n document_sources = [source[0] for source in sources]\n\n return document_sources\n" }, { "path": "backend/ee/onyx/db/connector_credential_pair.py", "content": "from sqlalchemy import delete\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import UserGroup__ConnectorCredentialPair\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _delete_connector_credential_pair_user_groups_relationship__no_commit(\n db_session: Session, connector_id: int, credential_id: int\n) -> None:\n cc_pair = get_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n )\n if cc_pair is None:\n raise ValueError(\n f\"ConnectorCredentialPair with connector_id: {connector_id} and credential_id: {credential_id} not found\"\n )\n\n stmt = delete(UserGroup__ConnectorCredentialPair).where(\n UserGroup__ConnectorCredentialPair.cc_pair_id == cc_pair.id,\n )\n db_session.execute(stmt)\n\n\ndef get_cc_pairs_by_source(\n db_session: Session,\n source_type: DocumentSource,\n access_type: AccessType | None = None,\n status: ConnectorCredentialPairStatus | None = None,\n) -> list[ConnectorCredentialPair]:\n \"\"\"\n Get all cc_pairs for a given source type with optional filtering by access_type and status\n result is sorted by cc_pair id\n \"\"\"\n query = (\n db_session.query(ConnectorCredentialPair)\n .join(ConnectorCredentialPair.connector)\n .filter(Connector.source == source_type)\n .order_by(ConnectorCredentialPair.id)\n )\n\n if access_type is not None:\n query = query.filter(ConnectorCredentialPair.access_type == access_type)\n\n if status is not None:\n query = query.filter(ConnectorCredentialPair.status == status)\n\n cc_pairs = query.all()\n return cc_pairs\n\n\ndef get_all_auto_sync_cc_pairs(\n db_session: Session,\n) -> list[ConnectorCredentialPair]:\n return (\n db_session.query(ConnectorCredentialPair)\n .where(\n ConnectorCredentialPair.access_type == AccessType.SYNC,\n )\n .all()\n )\n" }, { "path": "backend/ee/onyx/db/document.py", "content": "from datetime import datetime\nfrom datetime import timezone\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import Document as DbDocument\n\n\ndef upsert_document_external_perms__no_commit(\n db_session: Session,\n doc_id: str,\n external_access: ExternalAccess,\n source_type: DocumentSource,\n) -> None:\n \"\"\"\n This sets the permissions for a document in postgres.\n NOTE: this will replace any existing external access, it will not do a union\n \"\"\"\n document = db_session.scalars(\n select(DbDocument).where(DbDocument.id == doc_id)\n ).first()\n\n prefixed_external_groups = [\n build_ext_group_name_for_onyx(\n ext_group_name=group_id,\n source=source_type,\n )\n for group_id in external_access.external_user_group_ids\n ]\n\n if not document:\n # If the document does not exist, still store the external access\n # So that if the document is added later, the external access is already stored\n document = DbDocument(\n id=doc_id,\n semantic_id=\"\",\n external_user_emails=external_access.external_user_emails,\n external_user_group_ids=prefixed_external_groups,\n is_public=external_access.is_public,\n )\n db_session.add(document)\n return\n\n document.external_user_emails = list(external_access.external_user_emails)\n document.external_user_group_ids = prefixed_external_groups\n document.is_public = external_access.is_public\n\n\ndef upsert_document_external_perms(\n db_session: Session,\n doc_id: str,\n external_access: ExternalAccess,\n source_type: DocumentSource,\n) -> bool:\n \"\"\"\n This sets the permissions for a document in postgres. Returns True if the\n a new document was created, False otherwise.\n NOTE: this will replace any existing external access, it will not do a union\n \"\"\"\n document = db_session.scalars(\n select(DbDocument).where(DbDocument.id == doc_id)\n ).first()\n\n prefixed_external_groups: set[str] = {\n build_ext_group_name_for_onyx(\n ext_group_name=group_id,\n source=source_type,\n )\n for group_id in external_access.external_user_group_ids\n }\n\n if not document:\n # If the document does not exist, still store the external access\n # So that if the document is added later, the external access is already stored\n # The upsert function in the indexing pipeline does not overwrite the permissions fields\n document = DbDocument(\n id=doc_id,\n semantic_id=\"\",\n external_user_emails=external_access.external_user_emails,\n external_user_group_ids=prefixed_external_groups,\n is_public=external_access.is_public,\n )\n db_session.add(document)\n db_session.commit()\n return True\n\n # If the document exists, we need to check if the external access has changed\n if (\n external_access.external_user_emails != set(document.external_user_emails or [])\n or prefixed_external_groups != set(document.external_user_group_ids or [])\n or external_access.is_public != document.is_public\n ):\n document.external_user_emails = list(external_access.external_user_emails)\n document.external_user_group_ids = list(prefixed_external_groups)\n document.is_public = external_access.is_public\n document.last_modified = datetime.now(timezone.utc)\n db_session.commit()\n\n return False\n" }, { "path": "backend/ee/onyx/db/document_set.py", "content": "from uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import DocumentSet__ConnectorCredentialPair\nfrom onyx.db.models import DocumentSet__User\nfrom onyx.db.models import DocumentSet__UserGroup\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\n\n\ndef make_doc_set_private(\n document_set_id: int,\n user_ids: list[UUID] | None,\n group_ids: list[int] | None,\n db_session: Session,\n) -> None:\n db_session.query(DocumentSet__User).filter(\n DocumentSet__User.document_set_id == document_set_id\n ).delete(synchronize_session=\"fetch\")\n db_session.query(DocumentSet__UserGroup).filter(\n DocumentSet__UserGroup.document_set_id == document_set_id\n ).delete(synchronize_session=\"fetch\")\n\n if user_ids:\n for user_uuid in user_ids:\n db_session.add(\n DocumentSet__User(document_set_id=document_set_id, user_id=user_uuid)\n )\n\n if group_ids:\n for group_id in group_ids:\n db_session.add(\n DocumentSet__UserGroup(\n document_set_id=document_set_id, user_group_id=group_id\n )\n )\n\n\ndef delete_document_set_privacy__no_commit(\n document_set_id: int, db_session: Session\n) -> None:\n db_session.query(DocumentSet__User).filter(\n DocumentSet__User.document_set_id == document_set_id\n ).delete(synchronize_session=\"fetch\")\n\n db_session.query(DocumentSet__UserGroup).filter(\n DocumentSet__UserGroup.document_set_id == document_set_id\n ).delete(synchronize_session=\"fetch\")\n\n\ndef fetch_document_sets(\n user_id: UUID | None,\n db_session: Session,\n include_outdated: bool = True, # Parameter only for versioned implementation, unused # noqa: ARG001\n) -> list[tuple[DocumentSet, list[ConnectorCredentialPair]]]:\n assert user_id is not None\n\n # Public document sets\n public_document_sets = (\n db_session.query(DocumentSet)\n .filter(DocumentSet.is_public == True) # noqa\n .all()\n )\n\n # Document sets via shared user relationships\n shared_document_sets = (\n db_session.query(DocumentSet)\n .join(DocumentSet__User, DocumentSet.id == DocumentSet__User.document_set_id)\n .filter(DocumentSet__User.user_id == user_id)\n .all()\n )\n\n # Document sets via groups\n # First, find the user groups the user belongs to\n user_groups = (\n db_session.query(UserGroup)\n .join(User__UserGroup, UserGroup.id == User__UserGroup.user_group_id)\n .filter(User__UserGroup.user_id == user_id)\n .all()\n )\n\n group_document_sets = []\n for group in user_groups:\n group_document_sets.extend(\n db_session.query(DocumentSet)\n .join(\n DocumentSet__UserGroup,\n DocumentSet.id == DocumentSet__UserGroup.document_set_id,\n )\n .filter(DocumentSet__UserGroup.user_group_id == group.id)\n .all()\n )\n\n # Combine and deduplicate document sets from all sources\n all_document_sets = list(\n set(public_document_sets + shared_document_sets + group_document_sets)\n )\n\n document_set_with_cc_pairs: list[\n tuple[DocumentSet, list[ConnectorCredentialPair]]\n ] = []\n\n for document_set in all_document_sets:\n # Fetch the associated ConnectorCredentialPairs\n cc_pairs = (\n db_session.query(ConnectorCredentialPair)\n .join(\n DocumentSet__ConnectorCredentialPair,\n ConnectorCredentialPair.id\n == DocumentSet__ConnectorCredentialPair.connector_credential_pair_id,\n )\n .filter(\n DocumentSet__ConnectorCredentialPair.document_set_id == document_set.id,\n )\n .all()\n )\n\n document_set_with_cc_pairs.append((document_set, cc_pairs))\n\n return document_set_with_cc_pairs\n" }, { "path": "backend/ee/onyx/db/external_perm.py", "content": "from collections.abc import Sequence\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import PublicExternalUserGroup\nfrom onyx.db.models import User\nfrom onyx.db.models import User__ExternalUserGroupId\nfrom onyx.db.users import batch_add_ext_perm_user_if_not_exists\nfrom onyx.db.users import get_user_by_email\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass ExternalUserGroup(BaseModel):\n id: str\n user_emails: list[str]\n # `True` for cases like a Folder in Google Drive that give domain-wide\n # or \"Anyone with link\" access to all files in the folder.\n # if this is set, `user_emails` don't really matter.\n # When this is `True`, this `ExternalUserGroup` object doesn't really represent\n # an actual \"group\" in the source.\n gives_anyone_access: bool = False\n\n\ndef delete_user__ext_group_for_user__no_commit(\n db_session: Session,\n user_id: UUID,\n) -> None:\n db_session.execute(\n delete(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.user_id == user_id\n )\n )\n\n\ndef delete_user__ext_group_for_cc_pair__no_commit(\n db_session: Session,\n cc_pair_id: int,\n) -> None:\n db_session.execute(\n delete(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.cc_pair_id == cc_pair_id\n )\n )\n\n\ndef delete_public_external_group_for_cc_pair__no_commit(\n db_session: Session,\n cc_pair_id: int,\n) -> None:\n db_session.execute(\n delete(PublicExternalUserGroup).where(\n PublicExternalUserGroup.cc_pair_id == cc_pair_id\n )\n )\n\n\ndef mark_old_external_groups_as_stale(\n db_session: Session,\n cc_pair_id: int,\n) -> None:\n db_session.execute(\n update(User__ExternalUserGroupId)\n .where(User__ExternalUserGroupId.cc_pair_id == cc_pair_id)\n .values(stale=True)\n )\n db_session.execute(\n update(PublicExternalUserGroup)\n .where(PublicExternalUserGroup.cc_pair_id == cc_pair_id)\n .values(stale=True)\n )\n\n\ndef upsert_external_groups(\n db_session: Session,\n cc_pair_id: int,\n external_groups: list[ExternalUserGroup],\n source: DocumentSource,\n) -> None:\n \"\"\"\n Performs a true upsert operation for external user groups:\n - For existing groups (same user_id, external_user_group_id, cc_pair_id), updates the stale flag to False\n - For new groups, inserts them with stale=False\n - For public groups, uses upsert logic as well\n \"\"\"\n # If there are no groups to add, return early\n if not external_groups:\n return\n\n # collect all emails from all groups to batch add all users at once for efficiency\n all_group_member_emails = set()\n for external_group in external_groups:\n for user_email in external_group.user_emails:\n all_group_member_emails.add(user_email)\n\n # batch add users if they don't exist and get their ids\n all_group_members: list[User] = batch_add_ext_perm_user_if_not_exists(\n db_session=db_session,\n # NOTE: this function handles case sensitivity for emails\n emails=list(all_group_member_emails),\n )\n\n # map emails to ids\n email_id_map = {user.email.lower(): user.id for user in all_group_members}\n\n # Process each external group\n for external_group in external_groups:\n external_group_id = build_ext_group_name_for_onyx(\n ext_group_name=external_group.id,\n source=source,\n )\n\n # Handle user-group mappings\n for user_email in external_group.user_emails:\n user_id = email_id_map.get(user_email.lower())\n if user_id is None:\n logger.warning(\n f\"User in group {external_group.id} with email {user_email} not found\"\n )\n continue\n\n # Check if the user-group mapping already exists\n existing_user_group = db_session.scalar(\n select(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.user_id == user_id,\n User__ExternalUserGroupId.external_user_group_id\n == external_group_id,\n User__ExternalUserGroupId.cc_pair_id == cc_pair_id,\n )\n )\n\n if existing_user_group:\n # Update existing record\n existing_user_group.stale = False\n else:\n # Insert new record\n new_user_group = User__ExternalUserGroupId(\n user_id=user_id,\n external_user_group_id=external_group_id,\n cc_pair_id=cc_pair_id,\n stale=False,\n )\n db_session.add(new_user_group)\n\n # Handle public group if needed\n if external_group.gives_anyone_access:\n # Check if the public group already exists\n existing_public_group = db_session.scalar(\n select(PublicExternalUserGroup).where(\n PublicExternalUserGroup.external_user_group_id == external_group_id,\n PublicExternalUserGroup.cc_pair_id == cc_pair_id,\n )\n )\n\n if existing_public_group:\n # Update existing record\n existing_public_group.stale = False\n else:\n # Insert new record\n new_public_group = PublicExternalUserGroup(\n external_user_group_id=external_group_id,\n cc_pair_id=cc_pair_id,\n stale=False,\n )\n db_session.add(new_public_group)\n\n db_session.commit()\n\n\ndef remove_stale_external_groups(\n db_session: Session,\n cc_pair_id: int,\n) -> None:\n db_session.execute(\n delete(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.cc_pair_id == cc_pair_id,\n User__ExternalUserGroupId.stale.is_(True),\n )\n )\n db_session.execute(\n delete(PublicExternalUserGroup).where(\n PublicExternalUserGroup.cc_pair_id == cc_pair_id,\n PublicExternalUserGroup.stale.is_(True),\n )\n )\n db_session.commit()\n\n\ndef fetch_external_groups_for_user(\n db_session: Session,\n user_id: UUID,\n) -> Sequence[User__ExternalUserGroupId]:\n return db_session.scalars(\n select(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.user_id == user_id\n )\n ).all()\n\n\ndef fetch_external_groups_for_user_email_and_group_ids(\n db_session: Session,\n user_email: str,\n group_ids: list[str],\n) -> list[User__ExternalUserGroupId]:\n user = get_user_by_email(db_session=db_session, email=user_email)\n if user is None:\n return []\n user_id = user.id\n user_ext_groups = db_session.scalars(\n select(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.user_id == user_id,\n User__ExternalUserGroupId.external_user_group_id.in_(group_ids),\n )\n ).all()\n return list(user_ext_groups)\n\n\ndef fetch_public_external_group_ids(\n db_session: Session,\n) -> list[str]:\n return list(\n db_session.scalars(select(PublicExternalUserGroup.external_user_group_id)).all()\n )\n" }, { "path": "backend/ee/onyx/db/hierarchy.py", "content": "\"\"\"EE version of hierarchy node access control.\n\nThis module provides permission-aware hierarchy node access for Enterprise Edition.\nIt filters hierarchy nodes based on user email and external group membership.\n\"\"\"\n\nfrom sqlalchemy import any_\nfrom sqlalchemy import cast\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy import String\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql.elements import ColumnElement\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import HierarchyNode\n\n\ndef _build_hierarchy_access_filter(\n user_email: str,\n external_group_ids: list[str],\n) -> ColumnElement[bool]:\n \"\"\"Build SQLAlchemy filter for hierarchy node access.\n\n A user can access a hierarchy node if any of the following are true:\n - The node is marked as public (is_public=True)\n - The user's email is in the node's external_user_emails list\n - Any of the user's external group IDs overlap with the node's external_user_group_ids\n \"\"\"\n access_filters: list[ColumnElement[bool]] = [HierarchyNode.is_public.is_(True)]\n if user_email:\n access_filters.append(any_(HierarchyNode.external_user_emails) == user_email)\n if external_group_ids:\n access_filters.append(\n HierarchyNode.external_user_group_ids.overlap(\n cast(postgresql.array(external_group_ids), postgresql.ARRAY(String))\n )\n )\n return or_(*access_filters)\n\n\ndef _get_accessible_hierarchy_nodes_for_source(\n db_session: Session,\n source: DocumentSource,\n user_email: str,\n external_group_ids: list[str],\n) -> list[HierarchyNode]:\n \"\"\"\n EE version: Returns hierarchy nodes filtered by user permissions.\n\n A user can access a hierarchy node if any of the following are true:\n - The node is marked as public (is_public=True)\n - The user's email is in the node's external_user_emails list\n - Any of the user's external group IDs overlap with the node's external_user_group_ids\n\n Args:\n db_session: SQLAlchemy session\n source: Document source type\n user_email: User's email for permission checking\n external_group_ids: User's external group IDs for permission checking\n\n Returns:\n List of HierarchyNode objects the user has access to\n \"\"\"\n stmt = select(HierarchyNode).where(HierarchyNode.source == source)\n stmt = stmt.where(_build_hierarchy_access_filter(user_email, external_group_ids))\n stmt = stmt.order_by(HierarchyNode.display_name)\n return list(db_session.execute(stmt).scalars().all())\n" }, { "path": "backend/ee/onyx/db/license.py", "content": "\"\"\"Database and cache operations for the license table.\"\"\"\n\nfrom datetime import datetime\nfrom typing import NamedTuple\n\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.server.license.models import LicenseMetadata\nfrom ee.onyx.server.license.models import LicensePayload\nfrom ee.onyx.server.license.models import LicenseSource\nfrom onyx.auth.schemas import UserRole\nfrom onyx.cache.factory import get_cache_backend\nfrom onyx.configs.constants import ANONYMOUS_USER_EMAIL\nfrom onyx.db.models import License\nfrom onyx.db.models import User\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nLICENSE_METADATA_KEY = \"license:metadata\"\nLICENSE_CACHE_TTL_SECONDS = 86400 # 24 hours\n\n\nclass SeatAvailabilityResult(NamedTuple):\n \"\"\"Result of a seat availability check.\"\"\"\n\n available: bool\n error_message: str | None = None\n\n\n# -----------------------------------------------------------------------------\n# Database CRUD Operations\n# -----------------------------------------------------------------------------\n\n\ndef get_license(db_session: Session) -> License | None:\n \"\"\"\n Get the current license (singleton pattern - only one row).\n\n Args:\n db_session: Database session\n\n Returns:\n License object if exists, None otherwise\n \"\"\"\n return db_session.execute(select(License)).scalars().first()\n\n\ndef upsert_license(db_session: Session, license_data: str) -> License:\n \"\"\"\n Insert or update the license (singleton pattern).\n\n Args:\n db_session: Database session\n license_data: Base64-encoded signed license blob\n\n Returns:\n The created or updated License object\n \"\"\"\n existing = get_license(db_session)\n\n if existing:\n existing.license_data = license_data\n db_session.commit()\n db_session.refresh(existing)\n logger.info(\"License updated\")\n return existing\n\n new_license = License(license_data=license_data)\n db_session.add(new_license)\n db_session.commit()\n db_session.refresh(new_license)\n logger.info(\"License created\")\n return new_license\n\n\ndef delete_license(db_session: Session) -> bool:\n \"\"\"\n Delete the current license.\n\n Args:\n db_session: Database session\n\n Returns:\n True if deleted, False if no license existed\n \"\"\"\n existing = get_license(db_session)\n if existing:\n db_session.delete(existing)\n db_session.commit()\n logger.info(\"License deleted\")\n return True\n return False\n\n\n# -----------------------------------------------------------------------------\n# Seat Counting\n# -----------------------------------------------------------------------------\n\n\ndef get_used_seats(tenant_id: str | None = None) -> int:\n \"\"\"\n Get current seat usage directly from database.\n\n For multi-tenant: counts users in UserTenantMapping for this tenant.\n For self-hosted: counts all active users (excludes EXT_PERM_USER role\n and the anonymous system user).\n\n TODO: Exclude API key dummy users from seat counting. API keys create\n users with emails like `__DANSWER_API_KEY_*` that should not count toward\n seat limits. See: https://linear.app/onyx-app/issue/ENG-3518\n \"\"\"\n if MULTI_TENANT:\n from ee.onyx.server.tenants.user_mapping import get_tenant_count\n\n return get_tenant_count(tenant_id or get_current_tenant_id())\n else:\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n with get_session_with_current_tenant() as db_session:\n result = db_session.execute(\n select(func.count())\n .select_from(User)\n .where(\n User.is_active == True, # type: ignore # noqa: E712\n User.role != UserRole.EXT_PERM_USER,\n User.email != ANONYMOUS_USER_EMAIL, # type: ignore\n )\n )\n return result.scalar() or 0\n\n\n# -----------------------------------------------------------------------------\n# Redis Cache Operations\n# -----------------------------------------------------------------------------\n\n\ndef get_cached_license_metadata(tenant_id: str | None = None) -> LicenseMetadata | None:\n \"\"\"\n Get license metadata from cache.\n\n Args:\n tenant_id: Tenant ID (for multi-tenant deployments)\n\n Returns:\n LicenseMetadata if cached, None otherwise\n \"\"\"\n cache = get_cache_backend(tenant_id=tenant_id)\n cached = cache.get(LICENSE_METADATA_KEY)\n if not cached:\n return None\n\n try:\n cached_str = (\n cached.decode(\"utf-8\") if isinstance(cached, bytes) else str(cached)\n )\n return LicenseMetadata.model_validate_json(cached_str)\n except Exception as e:\n logger.warning(f\"Failed to parse cached license metadata: {e}\")\n return None\n\n\ndef invalidate_license_cache(tenant_id: str | None = None) -> None:\n \"\"\"\n Invalidate the license metadata cache (not the license itself).\n\n Deletes the cached LicenseMetadata. The actual license in the database\n is not affected. Delete is idempotent — if the key doesn't exist, this\n is a no-op.\n\n Args:\n tenant_id: Tenant ID (for multi-tenant deployments)\n \"\"\"\n cache = get_cache_backend(tenant_id=tenant_id)\n cache.delete(LICENSE_METADATA_KEY)\n logger.info(\"License cache invalidated\")\n\n\ndef update_license_cache(\n payload: LicensePayload,\n source: LicenseSource | None = None,\n grace_period_end: datetime | None = None,\n tenant_id: str | None = None,\n) -> LicenseMetadata:\n \"\"\"\n Update the cache with license metadata.\n\n We cache all license statuses (ACTIVE, GRACE_PERIOD, GATED_ACCESS) because:\n 1. Frontend needs status to show appropriate UI/banners\n 2. Caching avoids repeated DB + crypto verification on every request\n 3. Status enforcement happens at the feature level, not here\n\n Args:\n payload: Verified license payload\n source: How the license was obtained\n grace_period_end: Optional grace period end time\n tenant_id: Tenant ID (for multi-tenant deployments)\n\n Returns:\n The cached LicenseMetadata\n \"\"\"\n from ee.onyx.utils.license import get_license_status\n\n tenant = tenant_id or get_current_tenant_id()\n cache = get_cache_backend(tenant_id=tenant_id)\n\n used_seats = get_used_seats(tenant)\n status = get_license_status(payload, grace_period_end)\n\n metadata = LicenseMetadata(\n tenant_id=payload.tenant_id,\n organization_name=payload.organization_name,\n seats=payload.seats,\n used_seats=used_seats,\n plan_type=payload.plan_type,\n issued_at=payload.issued_at,\n expires_at=payload.expires_at,\n grace_period_end=grace_period_end,\n status=status,\n source=source,\n stripe_subscription_id=payload.stripe_subscription_id,\n )\n\n cache.set(\n LICENSE_METADATA_KEY,\n metadata.model_dump_json(),\n ex=LICENSE_CACHE_TTL_SECONDS,\n )\n\n logger.info(f\"License cache updated: {metadata.seats} seats, status={status.value}\")\n return metadata\n\n\ndef refresh_license_cache(\n db_session: Session,\n tenant_id: str | None = None,\n) -> LicenseMetadata | None:\n \"\"\"\n Refresh the license cache from the database.\n\n Args:\n db_session: Database session\n tenant_id: Tenant ID (for multi-tenant deployments)\n\n Returns:\n LicenseMetadata if license exists, None otherwise\n \"\"\"\n from ee.onyx.utils.license import verify_license_signature\n\n license_record = get_license(db_session)\n if not license_record:\n invalidate_license_cache(tenant_id)\n return None\n\n try:\n payload = verify_license_signature(license_record.license_data)\n # Derive source from payload: manual licenses lack stripe_customer_id\n source: LicenseSource = (\n LicenseSource.AUTO_FETCH\n if payload.stripe_customer_id\n else LicenseSource.MANUAL_UPLOAD\n )\n return update_license_cache(\n payload,\n source=source,\n tenant_id=tenant_id,\n )\n except ValueError as e:\n logger.error(f\"Failed to verify license during cache refresh: {e}\")\n invalidate_license_cache(tenant_id)\n return None\n\n\ndef get_license_metadata(\n db_session: Session,\n tenant_id: str | None = None,\n) -> LicenseMetadata | None:\n \"\"\"\n Get license metadata, using cache if available.\n\n Args:\n db_session: Database session\n tenant_id: Tenant ID (for multi-tenant deployments)\n\n Returns:\n LicenseMetadata if license exists, None otherwise\n \"\"\"\n # Try cache first\n cached = get_cached_license_metadata(tenant_id)\n if cached:\n return cached\n\n # Refresh from database\n return refresh_license_cache(db_session, tenant_id)\n\n\ndef check_seat_availability(\n db_session: Session,\n seats_needed: int = 1,\n tenant_id: str | None = None,\n) -> SeatAvailabilityResult:\n \"\"\"\n Check if there are enough seats available to add users.\n\n Args:\n db_session: Database session\n seats_needed: Number of seats needed (default 1)\n tenant_id: Tenant ID (for multi-tenant deployments)\n\n Returns:\n SeatAvailabilityResult with available=True if seats are available,\n or available=False with error_message if limit would be exceeded.\n Returns available=True if no license exists (self-hosted = unlimited).\n \"\"\"\n metadata = get_license_metadata(db_session, tenant_id)\n\n # No license = no enforcement (self-hosted without license)\n if metadata is None:\n return SeatAvailabilityResult(available=True)\n\n # Calculate current usage directly from DB (not cache) for accuracy\n current_used = get_used_seats(tenant_id)\n total_seats = metadata.seats\n\n # Use > (not >=) to allow filling to exactly 100% capacity\n would_exceed_limit = current_used + seats_needed > total_seats\n if would_exceed_limit:\n return SeatAvailabilityResult(\n available=False,\n error_message=f\"Seat limit would be exceeded: {current_used} of {total_seats} seats used, \"\n f\"cannot add {seats_needed} more user(s).\",\n )\n\n return SeatAvailabilityResult(available=True)\n" }, { "path": "backend/ee/onyx/db/persona.py", "content": "from uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import NotificationType\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__User\nfrom onyx.db.models import Persona__UserGroup\nfrom onyx.db.notification import create_notification\nfrom onyx.db.persona import mark_persona_user_files_for_sync\nfrom onyx.server.features.persona.models import PersonaSharedNotificationData\n\n\ndef update_persona_access(\n persona_id: int,\n creator_user_id: UUID | None,\n db_session: Session,\n is_public: bool | None = None,\n user_ids: list[UUID] | None = None,\n group_ids: list[int] | None = None,\n) -> None:\n \"\"\"Updates the access settings for a persona including public status, user shares,\n and group shares.\n\n NOTE: This function batches all updates. If we don't dedupe the inputs,\n the commit will exception.\n\n NOTE: Callers are responsible for committing.\"\"\"\n\n needs_sync = False\n if is_public is not None:\n needs_sync = True\n persona = db_session.query(Persona).filter(Persona.id == persona_id).first()\n if persona:\n persona.is_public = is_public\n\n # NOTE: For user-ids and group-ids, `None` means \"leave unchanged\", `[]` means \"clear all shares\",\n # and a non-empty list means \"replace with these shares\".\n\n if user_ids is not None:\n needs_sync = True\n db_session.query(Persona__User).filter(\n Persona__User.persona_id == persona_id\n ).delete(synchronize_session=\"fetch\")\n\n user_ids_set = set(user_ids)\n for user_id in user_ids_set:\n db_session.add(Persona__User(persona_id=persona_id, user_id=user_id))\n if user_id != creator_user_id:\n create_notification(\n user_id=user_id,\n notif_type=NotificationType.PERSONA_SHARED,\n title=\"A new agent was shared with you!\",\n db_session=db_session,\n additional_data=PersonaSharedNotificationData(\n persona_id=persona_id,\n ).model_dump(),\n )\n\n if group_ids is not None:\n needs_sync = True\n db_session.query(Persona__UserGroup).filter(\n Persona__UserGroup.persona_id == persona_id\n ).delete(synchronize_session=\"fetch\")\n\n group_ids_set = set(group_ids)\n for group_id in group_ids_set:\n db_session.add(\n Persona__UserGroup(persona_id=persona_id, user_group_id=group_id)\n )\n\n # When sharing changes, user file ACLs need to be updated in the vector DB\n if needs_sync:\n mark_persona_user_files_for_sync(persona_id, db_session)\n" }, { "path": "backend/ee/onyx/db/query_history.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\n\nfrom sqlalchemy import asc\nfrom sqlalchemy import BinaryExpression\nfrom sqlalchemy import ColumnElement\nfrom sqlalchemy import desc\nfrom sqlalchemy import distinct\nfrom sqlalchemy.orm import contains_eager\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql import case\nfrom sqlalchemy.sql import func\nfrom sqlalchemy.sql import select\nfrom sqlalchemy.sql.expression import literal\nfrom sqlalchemy.sql.expression import UnaryExpression\n\nfrom ee.onyx.background.task_name_builders import QUERY_HISTORY_TASK_NAME_PREFIX\nfrom onyx.configs.constants import QAFeedbackType\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatMessageFeedback\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import TaskQueueState\nfrom onyx.db.tasks import get_all_tasks_with_prefix\n\n\ndef _build_filter_conditions(\n start_time: datetime | None,\n end_time: datetime | None,\n feedback_filter: QAFeedbackType | None,\n) -> list[ColumnElement]:\n \"\"\"\n Helper function to build all filter conditions for chat sessions.\n Filters by start and end time, feedback type, and any sessions without messages.\n start_time: Date from which to filter\n end_time: Date to which to filter\n feedback_filter: Feedback type to filter by\n Returns: List of filter conditions\n \"\"\"\n conditions = []\n\n if start_time is not None:\n conditions.append(ChatSession.time_created >= start_time)\n if end_time is not None:\n conditions.append(ChatSession.time_created <= end_time)\n\n if feedback_filter is not None:\n feedback_subq = (\n select(ChatMessage.chat_session_id)\n .join(ChatMessageFeedback)\n .group_by(ChatMessage.chat_session_id)\n .having(\n case(\n (\n case(\n {literal(feedback_filter == QAFeedbackType.LIKE): True},\n else_=False,\n ),\n func.bool_and(ChatMessageFeedback.is_positive),\n ),\n (\n case(\n {literal(feedback_filter == QAFeedbackType.DISLIKE): True},\n else_=False,\n ),\n func.bool_and(func.not_(ChatMessageFeedback.is_positive)),\n ),\n else_=func.bool_or(ChatMessageFeedback.is_positive)\n & func.bool_or(func.not_(ChatMessageFeedback.is_positive)),\n )\n )\n )\n conditions.append(ChatSession.id.in_(feedback_subq))\n\n return conditions\n\n\ndef get_total_filtered_chat_sessions_count(\n db_session: Session,\n start_time: datetime | None,\n end_time: datetime | None,\n feedback_filter: QAFeedbackType | None,\n) -> int:\n conditions = _build_filter_conditions(start_time, end_time, feedback_filter)\n stmt = (\n select(func.count(distinct(ChatSession.id)))\n .select_from(ChatSession)\n .filter(*conditions)\n )\n return db_session.scalar(stmt) or 0\n\n\ndef get_page_of_chat_sessions(\n start_time: datetime | None,\n end_time: datetime | None,\n db_session: Session,\n page_num: int,\n page_size: int,\n feedback_filter: QAFeedbackType | None = None,\n) -> Sequence[ChatSession]:\n conditions = _build_filter_conditions(start_time, end_time, feedback_filter)\n\n subquery = (\n select(ChatSession.id)\n .filter(*conditions)\n .order_by(desc(ChatSession.time_created), ChatSession.id)\n .limit(page_size)\n .offset(page_num * page_size)\n .subquery()\n )\n\n stmt = (\n select(ChatSession)\n .join(subquery, ChatSession.id == subquery.c.id)\n .outerjoin(ChatMessage, ChatSession.id == ChatMessage.chat_session_id)\n .options(\n joinedload(ChatSession.user),\n joinedload(ChatSession.persona),\n contains_eager(ChatSession.messages).joinedload(\n ChatMessage.chat_message_feedbacks\n ),\n )\n .order_by(\n desc(ChatSession.time_created),\n ChatSession.id,\n asc(ChatMessage.id), # Ensure chronological message order\n )\n )\n\n return db_session.scalars(stmt).unique().all()\n\n\ndef fetch_chat_sessions_eagerly_by_time(\n start: datetime,\n end: datetime,\n db_session: Session,\n limit: int | None = 500,\n initial_time: datetime | None = None,\n) -> list[ChatSession]:\n \"\"\"Sorted by oldest to newest, then by message id\"\"\"\n\n asc_time_order: UnaryExpression = asc(ChatSession.time_created)\n message_order: UnaryExpression = asc(ChatMessage.id)\n\n filters: list[ColumnElement | BinaryExpression] = [\n ChatSession.time_created.between(start, end)\n ]\n\n if initial_time:\n filters.append(ChatSession.time_created > initial_time)\n\n subquery = (\n db_session.query(ChatSession.id, ChatSession.time_created)\n .filter(*filters)\n .order_by(asc_time_order)\n .limit(limit)\n .subquery()\n )\n\n query = (\n db_session.query(ChatSession)\n .join(subquery, ChatSession.id == subquery.c.id)\n .outerjoin(ChatMessage, ChatSession.id == ChatMessage.chat_session_id)\n .options(\n joinedload(ChatSession.user),\n joinedload(ChatSession.persona),\n contains_eager(ChatSession.messages).joinedload(\n ChatMessage.chat_message_feedbacks\n ),\n )\n .order_by(asc_time_order, message_order)\n )\n\n chat_sessions = query.all()\n\n return chat_sessions\n\n\ndef get_all_query_history_export_tasks(\n db_session: Session,\n) -> list[TaskQueueState]:\n return get_all_tasks_with_prefix(db_session, QUERY_HISTORY_TASK_NAME_PREFIX)\n" }, { "path": "backend/ee/onyx/db/saml.py", "content": "import datetime\nfrom typing import cast\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS\nfrom onyx.db.models import SamlAccount\n\n\ndef upsert_saml_account(\n user_id: UUID,\n cookie: str,\n db_session: Session,\n expiration_offset: int = SESSION_EXPIRE_TIME_SECONDS,\n) -> datetime.datetime:\n expires_at = func.now() + datetime.timedelta(seconds=expiration_offset)\n\n existing_saml_acc = (\n db_session.query(SamlAccount)\n .filter(SamlAccount.user_id == user_id)\n .one_or_none()\n )\n\n if existing_saml_acc:\n existing_saml_acc.encrypted_cookie = cookie\n existing_saml_acc.expires_at = cast(datetime.datetime, expires_at)\n existing_saml_acc.updated_at = func.now()\n saml_acc = existing_saml_acc\n else:\n saml_acc = SamlAccount(\n user_id=user_id,\n encrypted_cookie=cookie,\n expires_at=expires_at,\n )\n db_session.add(saml_acc)\n\n db_session.commit()\n\n return saml_acc.expires_at\n\n\nasync def get_saml_account(\n cookie: str, async_db_session: AsyncSession\n) -> SamlAccount | None:\n \"\"\"NOTE: this is async, since it's used during auth\n (which is necessarily async due to FastAPI Users)\"\"\"\n stmt = (\n select(SamlAccount)\n .options(selectinload(SamlAccount.user)) # Use selectinload for collections\n .where(\n and_(\n SamlAccount.encrypted_cookie == cookie,\n SamlAccount.expires_at > func.now(),\n )\n )\n )\n\n result = await async_db_session.execute(stmt)\n return result.scalars().unique().one_or_none()\n\n\nasync def expire_saml_account(\n saml_account: SamlAccount, async_db_session: AsyncSession\n) -> None:\n saml_account.expires_at = func.now()\n await async_db_session.commit()\n" }, { "path": "backend/ee/onyx/db/scim.py", "content": "\"\"\"SCIM Data Access Layer.\n\nAll database operations for SCIM provisioning — token management, user\nmappings, and group mappings. Extends the base DAL (see ``onyx.db.dal``).\n\nUsage from FastAPI::\n\n def get_scim_dal(db_session: Session = Depends(get_session)) -> ScimDAL:\n return ScimDAL(db_session)\n\n @router.post(\"/tokens\")\n def create_token(dal: ScimDAL = Depends(get_scim_dal)) -> ...:\n token = dal.create_token(name=..., hashed_token=..., ...)\n dal.commit()\n return token\n\nUsage from background tasks::\n\n with ScimDAL.from_tenant(\"tenant_abc\") as dal:\n mapping = dal.create_user_mapping(external_id=\"idp-123\", user_id=uid)\n dal.commit()\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom uuid import UUID\n\nfrom sqlalchemy import delete as sa_delete\nfrom sqlalchemy import func\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import SQLColumnExpression\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\n\nfrom ee.onyx.server.scim.filtering import ScimFilter\nfrom ee.onyx.server.scim.filtering import ScimFilterOperator\nfrom ee.onyx.server.scim.models import ScimMappingFields\nfrom onyx.db.dal import DAL\nfrom onyx.db.enums import AccountType\nfrom onyx.db.enums import GrantSource\nfrom onyx.db.enums import Permission\nfrom onyx.db.models import PermissionGrant\nfrom onyx.db.models import ScimGroupMapping\nfrom onyx.db.models import ScimToken\nfrom onyx.db.models import ScimUserMapping\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass ScimDAL(DAL):\n \"\"\"Data Access Layer for SCIM provisioning operations.\n\n Methods mutate but do NOT commit — call ``dal.commit()`` explicitly\n when you want to persist changes. This follows the existing ``_no_commit``\n convention and lets callers batch multiple operations into one transaction.\n \"\"\"\n\n # ------------------------------------------------------------------\n # Token operations\n # ------------------------------------------------------------------\n\n def create_token(\n self,\n name: str,\n hashed_token: str,\n token_display: str,\n created_by_id: UUID,\n ) -> ScimToken:\n \"\"\"Create a new SCIM bearer token.\n\n Only one token is active at a time — this method automatically revokes\n all existing active tokens before creating the new one.\n \"\"\"\n # Revoke any currently active tokens\n active_tokens = list(\n self._session.scalars(\n select(ScimToken).where(ScimToken.is_active.is_(True))\n ).all()\n )\n for t in active_tokens:\n t.is_active = False\n\n token = ScimToken(\n name=name,\n hashed_token=hashed_token,\n token_display=token_display,\n created_by_id=created_by_id,\n )\n self._session.add(token)\n self._session.flush()\n return token\n\n def get_active_token(self) -> ScimToken | None:\n \"\"\"Return the single currently active token, or None.\"\"\"\n return self._session.scalar(\n select(ScimToken).where(ScimToken.is_active.is_(True))\n )\n\n def get_token_by_hash(self, hashed_token: str) -> ScimToken | None:\n \"\"\"Look up a token by its SHA-256 hash.\"\"\"\n return self._session.scalar(\n select(ScimToken).where(ScimToken.hashed_token == hashed_token)\n )\n\n def revoke_token(self, token_id: int) -> None:\n \"\"\"Deactivate a token by ID.\n\n Raises:\n ValueError: If the token does not exist.\n \"\"\"\n token = self._session.get(ScimToken, token_id)\n if not token:\n raise ValueError(f\"SCIM token with id {token_id} not found\")\n token.is_active = False\n\n def update_token_last_used(self, token_id: int) -> None:\n \"\"\"Update the last_used_at timestamp for a token.\"\"\"\n token = self._session.get(ScimToken, token_id)\n if token:\n token.last_used_at = func.now() # type: ignore[assignment]\n\n # ------------------------------------------------------------------\n # User mapping operations\n # ------------------------------------------------------------------\n\n def create_user_mapping(\n self,\n external_id: str | None,\n user_id: UUID,\n scim_username: str | None = None,\n fields: ScimMappingFields | None = None,\n ) -> ScimUserMapping:\n \"\"\"Create a SCIM mapping for a user.\n\n ``external_id`` may be ``None`` when the IdP omits it (RFC 7643\n allows this). The mapping still marks the user as SCIM-managed.\n \"\"\"\n f = fields or ScimMappingFields()\n mapping = ScimUserMapping(\n external_id=external_id,\n user_id=user_id,\n scim_username=scim_username,\n department=f.department,\n manager=f.manager,\n given_name=f.given_name,\n family_name=f.family_name,\n scim_emails_json=f.scim_emails_json,\n )\n self._session.add(mapping)\n self._session.flush()\n return mapping\n\n def get_user_mapping_by_external_id(\n self, external_id: str\n ) -> ScimUserMapping | None:\n \"\"\"Look up a user mapping by the IdP's external identifier.\"\"\"\n return self._session.scalar(\n select(ScimUserMapping).where(ScimUserMapping.external_id == external_id)\n )\n\n def get_user_mapping_by_user_id(self, user_id: UUID) -> ScimUserMapping | None:\n \"\"\"Look up a user mapping by the Onyx user ID.\"\"\"\n return self._session.scalar(\n select(ScimUserMapping).where(ScimUserMapping.user_id == user_id)\n )\n\n def list_user_mappings(\n self,\n start_index: int = 1,\n count: int = 100,\n ) -> tuple[list[ScimUserMapping], int]:\n \"\"\"List user mappings with SCIM-style pagination.\n\n Args:\n start_index: 1-based start index (SCIM convention).\n count: Maximum number of results to return.\n\n Returns:\n A tuple of (mappings, total_count).\n \"\"\"\n total = (\n self._session.scalar(select(func.count()).select_from(ScimUserMapping)) or 0\n )\n\n offset = max(start_index - 1, 0)\n mappings = list(\n self._session.scalars(\n select(ScimUserMapping)\n .order_by(ScimUserMapping.id)\n .offset(offset)\n .limit(count)\n ).all()\n )\n\n return mappings, total\n\n def update_user_mapping_external_id(\n self,\n mapping_id: int,\n external_id: str,\n ) -> ScimUserMapping:\n \"\"\"Update the external ID on a user mapping.\n\n Raises:\n ValueError: If the mapping does not exist.\n \"\"\"\n mapping = self._session.get(ScimUserMapping, mapping_id)\n if not mapping:\n raise ValueError(f\"SCIM user mapping with id {mapping_id} not found\")\n mapping.external_id = external_id\n return mapping\n\n def delete_user_mapping(self, mapping_id: int) -> None:\n \"\"\"Delete a user mapping by ID. No-op if already deleted.\"\"\"\n mapping = self._session.get(ScimUserMapping, mapping_id)\n if not mapping:\n logger.warning(\"SCIM user mapping %d not found during delete\", mapping_id)\n return\n self._session.delete(mapping)\n\n # ------------------------------------------------------------------\n # User query operations\n # ------------------------------------------------------------------\n\n def get_user(self, user_id: UUID) -> User | None:\n \"\"\"Fetch a user by ID.\"\"\"\n return self._session.scalar(\n select(User).where(User.id == user_id) # type: ignore[arg-type]\n )\n\n def get_user_by_email(self, email: str) -> User | None:\n \"\"\"Fetch a user by email (case-insensitive).\"\"\"\n return self._session.scalar(\n select(User).where(func.lower(User.email) == func.lower(email))\n )\n\n def add_user(self, user: User) -> None:\n \"\"\"Add a new user to the session and flush to assign an ID.\"\"\"\n self._session.add(user)\n self._session.flush()\n\n def update_user(\n self,\n user: User,\n *,\n email: str | None = None,\n is_active: bool | None = None,\n personal_name: str | None = None,\n ) -> None:\n \"\"\"Update user attributes. Only sets fields that are provided.\"\"\"\n if email is not None:\n user.email = email\n if is_active is not None:\n user.is_active = is_active\n if personal_name is not None:\n user.personal_name = personal_name\n\n def deactivate_user(self, user: User) -> None:\n \"\"\"Mark a user as inactive.\"\"\"\n user.is_active = False\n\n def list_users(\n self,\n scim_filter: ScimFilter | None,\n start_index: int = 1,\n count: int = 100,\n ) -> tuple[list[tuple[User, ScimUserMapping | None]], int]:\n \"\"\"Query users with optional SCIM filter and pagination.\n\n Returns:\n A tuple of (list of (user, mapping) pairs, total_count).\n\n Raises:\n ValueError: If the filter uses an unsupported attribute.\n \"\"\"\n # Inner-join with ScimUserMapping so only SCIM-managed users appear.\n # Pre-existing system accounts (anonymous, admin, etc.) are excluded\n # unless they were explicitly linked via SCIM provisioning.\n query = (\n select(User)\n .join(ScimUserMapping, ScimUserMapping.user_id == User.id)\n .where(\n User.account_type.notin_([AccountType.BOT, AccountType.EXT_PERM_USER])\n )\n )\n\n if scim_filter:\n attr = scim_filter.attribute.lower()\n if attr == \"username\":\n # arg-type: fastapi-users types User.email as str, not a column expression\n # assignment: union return type widens but query is still Select[tuple[User]]\n query = _apply_scim_string_op(query, User.email, scim_filter) # type: ignore[arg-type, assignment]\n elif attr == \"active\":\n query = query.where(\n User.is_active.is_(scim_filter.value.lower() == \"true\") # type: ignore[attr-defined]\n )\n elif attr == \"externalid\":\n mapping = self.get_user_mapping_by_external_id(scim_filter.value)\n if not mapping:\n return [], 0\n query = query.where(User.id == mapping.user_id) # type: ignore[arg-type]\n else:\n raise ValueError(\n f\"Unsupported filter attribute: {scim_filter.attribute}\"\n )\n\n # Count total matching rows first, then paginate. SCIM uses 1-based\n # indexing (RFC 7644 §3.4.2), so we convert to a 0-based offset.\n total = (\n self._session.scalar(select(func.count()).select_from(query.subquery()))\n or 0\n )\n\n offset = max(start_index - 1, 0)\n users = list(\n self._session.scalars(\n query.order_by(User.id).offset(offset).limit(count) # type: ignore[arg-type]\n )\n .unique()\n .all()\n )\n\n # Batch-fetch SCIM mappings to avoid N+1 queries\n mapping_map = self._get_user_mappings_batch([u.id for u in users])\n return [(u, mapping_map.get(u.id)) for u in users], total\n\n def sync_user_external_id(\n self,\n user_id: UUID,\n new_external_id: str | None,\n scim_username: str | None = None,\n fields: ScimMappingFields | None = None,\n ) -> None:\n \"\"\"Sync the SCIM mapping for a user.\n\n If a mapping already exists, its fields are updated (including\n setting ``external_id`` to ``None`` when the IdP omits it).\n If no mapping exists and ``new_external_id`` is provided, a new\n mapping is created. A mapping is never deleted here — SCIM-managed\n users must retain their mapping to remain visible in ``GET /Users``.\n\n When *fields* is provided, all mapping fields are written\n unconditionally — including ``None`` values — so that a caller can\n clear a previously-set field (e.g. removing a department).\n \"\"\"\n mapping = self.get_user_mapping_by_user_id(user_id)\n if mapping:\n if mapping.external_id != new_external_id:\n mapping.external_id = new_external_id\n if scim_username is not None:\n mapping.scim_username = scim_username\n if fields is not None:\n mapping.department = fields.department\n mapping.manager = fields.manager\n mapping.given_name = fields.given_name\n mapping.family_name = fields.family_name\n mapping.scim_emails_json = fields.scim_emails_json\n elif new_external_id:\n self.create_user_mapping(\n external_id=new_external_id,\n user_id=user_id,\n scim_username=scim_username,\n fields=fields,\n )\n\n def _get_user_mappings_batch(\n self, user_ids: list[UUID]\n ) -> dict[UUID, ScimUserMapping]:\n \"\"\"Batch-fetch SCIM user mappings keyed by user ID.\"\"\"\n if not user_ids:\n return {}\n mappings = self._session.scalars(\n select(ScimUserMapping).where(ScimUserMapping.user_id.in_(user_ids))\n ).all()\n return {m.user_id: m for m in mappings}\n\n def get_user_groups(self, user_id: UUID) -> list[tuple[int, str]]:\n \"\"\"Get groups a user belongs to as ``(group_id, group_name)`` pairs.\n\n Excludes groups marked for deletion.\n \"\"\"\n rels = self._session.scalars(\n select(User__UserGroup).where(User__UserGroup.user_id == user_id)\n ).all()\n\n group_ids = [r.user_group_id for r in rels]\n if not group_ids:\n return []\n\n groups = self._session.scalars(\n select(UserGroup).where(\n UserGroup.id.in_(group_ids),\n UserGroup.is_up_for_deletion.is_(False),\n )\n ).all()\n return [(g.id, g.name) for g in groups]\n\n def get_users_groups_batch(\n self, user_ids: list[UUID]\n ) -> dict[UUID, list[tuple[int, str]]]:\n \"\"\"Batch-fetch group memberships for multiple users.\n\n Returns a mapping of ``user_id → [(group_id, group_name), ...]``.\n Avoids N+1 queries when building user list responses.\n \"\"\"\n if not user_ids:\n return {}\n\n rels = self._session.scalars(\n select(User__UserGroup).where(User__UserGroup.user_id.in_(user_ids))\n ).all()\n\n group_ids = list({r.user_group_id for r in rels})\n if not group_ids:\n return {}\n\n groups = self._session.scalars(\n select(UserGroup).where(\n UserGroup.id.in_(group_ids),\n UserGroup.is_up_for_deletion.is_(False),\n )\n ).all()\n groups_by_id = {g.id: g.name for g in groups}\n\n result: dict[UUID, list[tuple[int, str]]] = {}\n for r in rels:\n if r.user_id and r.user_group_id in groups_by_id:\n result.setdefault(r.user_id, []).append(\n (r.user_group_id, groups_by_id[r.user_group_id])\n )\n return result\n\n # ------------------------------------------------------------------\n # Group mapping operations\n # ------------------------------------------------------------------\n\n def create_group_mapping(\n self,\n external_id: str,\n user_group_id: int,\n ) -> ScimGroupMapping:\n \"\"\"Create a mapping between a SCIM externalId and an Onyx user group.\"\"\"\n mapping = ScimGroupMapping(external_id=external_id, user_group_id=user_group_id)\n self._session.add(mapping)\n self._session.flush()\n return mapping\n\n def get_group_mapping_by_external_id(\n self, external_id: str\n ) -> ScimGroupMapping | None:\n \"\"\"Look up a group mapping by the IdP's external identifier.\"\"\"\n return self._session.scalar(\n select(ScimGroupMapping).where(ScimGroupMapping.external_id == external_id)\n )\n\n def get_group_mapping_by_group_id(\n self, user_group_id: int\n ) -> ScimGroupMapping | None:\n \"\"\"Look up a group mapping by the Onyx user group ID.\"\"\"\n return self._session.scalar(\n select(ScimGroupMapping).where(\n ScimGroupMapping.user_group_id == user_group_id\n )\n )\n\n def list_group_mappings(\n self,\n start_index: int = 1,\n count: int = 100,\n ) -> tuple[list[ScimGroupMapping], int]:\n \"\"\"List group mappings with SCIM-style pagination.\n\n Args:\n start_index: 1-based start index (SCIM convention).\n count: Maximum number of results to return.\n\n Returns:\n A tuple of (mappings, total_count).\n \"\"\"\n total = (\n self._session.scalar(select(func.count()).select_from(ScimGroupMapping))\n or 0\n )\n\n offset = max(start_index - 1, 0)\n mappings = list(\n self._session.scalars(\n select(ScimGroupMapping)\n .order_by(ScimGroupMapping.id)\n .offset(offset)\n .limit(count)\n ).all()\n )\n\n return mappings, total\n\n def delete_group_mapping(self, mapping_id: int) -> None:\n \"\"\"Delete a group mapping by ID. No-op if already deleted.\"\"\"\n mapping = self._session.get(ScimGroupMapping, mapping_id)\n if not mapping:\n logger.warning(\"SCIM group mapping %d not found during delete\", mapping_id)\n return\n self._session.delete(mapping)\n\n # ------------------------------------------------------------------\n # Group query operations\n # ------------------------------------------------------------------\n\n def get_group(self, group_id: int) -> UserGroup | None:\n \"\"\"Fetch a group by ID, returning None if deleted or missing.\"\"\"\n group = self._session.get(UserGroup, group_id)\n if group and group.is_up_for_deletion:\n return None\n return group\n\n def get_group_by_name(self, name: str) -> UserGroup | None:\n \"\"\"Fetch a group by exact name.\"\"\"\n return self._session.scalar(select(UserGroup).where(UserGroup.name == name))\n\n def add_group(self, group: UserGroup) -> None:\n \"\"\"Add a new group to the session and flush to assign an ID.\"\"\"\n self._session.add(group)\n self._session.flush()\n\n def add_permission_grant_to_group(\n self,\n group_id: int,\n permission: Permission,\n grant_source: GrantSource,\n ) -> None:\n \"\"\"Grant a permission to a group and flush.\"\"\"\n self._session.add(\n PermissionGrant(\n group_id=group_id,\n permission=permission,\n grant_source=grant_source,\n )\n )\n self._session.flush()\n\n def update_group(\n self,\n group: UserGroup,\n *,\n name: str | None = None,\n ) -> None:\n \"\"\"Update group attributes and set the modification timestamp.\"\"\"\n if name is not None:\n group.name = name\n group.time_last_modified_by_user = func.now()\n\n def delete_group(self, group: UserGroup) -> None:\n \"\"\"Delete a group from the session.\"\"\"\n self._session.delete(group)\n\n def list_groups(\n self,\n scim_filter: ScimFilter | None,\n start_index: int = 1,\n count: int = 100,\n ) -> tuple[list[tuple[UserGroup, str | None]], int]:\n \"\"\"Query groups with optional SCIM filter and pagination.\n\n Returns:\n A tuple of (list of (group, external_id) pairs, total_count).\n\n Raises:\n ValueError: If the filter uses an unsupported attribute.\n \"\"\"\n query = select(UserGroup).where(UserGroup.is_up_for_deletion.is_(False))\n\n if scim_filter:\n attr = scim_filter.attribute.lower()\n if attr == \"displayname\":\n # assignment: union return type widens but query is still Select[tuple[UserGroup]]\n query = _apply_scim_string_op(query, UserGroup.name, scim_filter) # type: ignore[assignment]\n elif attr == \"externalid\":\n mapping = self.get_group_mapping_by_external_id(scim_filter.value)\n if not mapping:\n return [], 0\n query = query.where(UserGroup.id == mapping.user_group_id)\n else:\n raise ValueError(\n f\"Unsupported filter attribute: {scim_filter.attribute}\"\n )\n\n total = (\n self._session.scalar(select(func.count()).select_from(query.subquery()))\n or 0\n )\n\n offset = max(start_index - 1, 0)\n groups = list(\n self._session.scalars(\n query.order_by(UserGroup.id).offset(offset).limit(count)\n ).all()\n )\n\n ext_id_map = self._get_group_external_ids([g.id for g in groups])\n return [(g, ext_id_map.get(g.id)) for g in groups], total\n\n def get_group_members(self, group_id: int) -> list[tuple[UUID, str | None]]:\n \"\"\"Get group members as (user_id, email) pairs.\"\"\"\n rels = self._session.scalars(\n select(User__UserGroup).where(User__UserGroup.user_group_id == group_id)\n ).all()\n\n user_ids = [r.user_id for r in rels if r.user_id]\n if not user_ids:\n return []\n\n users = (\n self._session.scalars(\n select(User).where(User.id.in_(user_ids)) # type: ignore[attr-defined]\n )\n .unique()\n .all()\n )\n users_by_id = {u.id: u for u in users}\n\n return [\n (\n r.user_id,\n users_by_id[r.user_id].email if r.user_id in users_by_id else None,\n )\n for r in rels\n if r.user_id\n ]\n\n def validate_member_ids(self, uuids: list[UUID]) -> list[UUID]:\n \"\"\"Return the subset of UUIDs that don't exist as users.\n\n Returns an empty list if all IDs are valid.\n \"\"\"\n if not uuids:\n return []\n existing_users = (\n self._session.scalars(\n select(User).where(User.id.in_(uuids)) # type: ignore[attr-defined]\n )\n .unique()\n .all()\n )\n existing_ids = {u.id for u in existing_users}\n return [uid for uid in uuids if uid not in existing_ids]\n\n def upsert_group_members(self, group_id: int, user_ids: list[UUID]) -> None:\n \"\"\"Add user-group relationships, ignoring duplicates.\"\"\"\n if not user_ids:\n return\n self._session.execute(\n pg_insert(User__UserGroup)\n .values([{\"user_id\": uid, \"user_group_id\": group_id} for uid in user_ids])\n .on_conflict_do_nothing(\n index_elements=[\n User__UserGroup.user_group_id,\n User__UserGroup.user_id,\n ]\n )\n )\n\n def replace_group_members(self, group_id: int, user_ids: list[UUID]) -> None:\n \"\"\"Replace all members of a group.\"\"\"\n self._session.execute(\n sa_delete(User__UserGroup).where(User__UserGroup.user_group_id == group_id)\n )\n self.upsert_group_members(group_id, user_ids)\n\n def remove_group_members(self, group_id: int, user_ids: list[UUID]) -> None:\n \"\"\"Remove specific members from a group.\"\"\"\n if not user_ids:\n return\n self._session.execute(\n sa_delete(User__UserGroup).where(\n User__UserGroup.user_group_id == group_id,\n User__UserGroup.user_id.in_(user_ids),\n )\n )\n\n def delete_group_with_members(self, group: UserGroup) -> None:\n \"\"\"Remove all member relationships and delete the group.\"\"\"\n self._session.execute(\n sa_delete(User__UserGroup).where(User__UserGroup.user_group_id == group.id)\n )\n self._session.delete(group)\n\n def sync_group_external_id(\n self, group_id: int, new_external_id: str | None\n ) -> None:\n \"\"\"Create, update, or delete the external ID mapping for a group.\"\"\"\n mapping = self.get_group_mapping_by_group_id(group_id)\n if new_external_id:\n if mapping:\n if mapping.external_id != new_external_id:\n mapping.external_id = new_external_id\n else:\n self.create_group_mapping(\n external_id=new_external_id, user_group_id=group_id\n )\n elif mapping:\n self.delete_group_mapping(mapping.id)\n\n def _get_group_external_ids(self, group_ids: list[int]) -> dict[int, str]:\n \"\"\"Batch-fetch external IDs for a list of group IDs.\"\"\"\n if not group_ids:\n return {}\n mappings = self._session.scalars(\n select(ScimGroupMapping).where(\n ScimGroupMapping.user_group_id.in_(group_ids)\n )\n ).all()\n return {m.user_group_id: m.external_id for m in mappings}\n\n\n# ---------------------------------------------------------------------------\n# Module-level helpers (used by DAL methods above)\n# ---------------------------------------------------------------------------\n\n\ndef _apply_scim_string_op(\n query: Select[tuple[User]] | Select[tuple[UserGroup]],\n column: SQLColumnExpression[str],\n scim_filter: ScimFilter,\n) -> Select[tuple[User]] | Select[tuple[UserGroup]]:\n \"\"\"Apply a SCIM string filter operator using SQLAlchemy column operators.\n\n Handles eq (case-insensitive exact), co (contains), and sw (starts with).\n SQLAlchemy's operators handle LIKE-pattern escaping internally.\n \"\"\"\n val = scim_filter.value\n if scim_filter.operator == ScimFilterOperator.EQUAL:\n return query.where(func.lower(column) == val.lower())\n elif scim_filter.operator == ScimFilterOperator.CONTAINS:\n return query.where(column.icontains(val, autoescape=True))\n elif scim_filter.operator == ScimFilterOperator.STARTS_WITH:\n return query.where(column.istartswith(val, autoescape=True))\n else:\n raise ValueError(f\"Unsupported string filter operator: {scim_filter.operator}\")\n" }, { "path": "backend/ee/onyx/db/search.py", "content": "import uuid\nfrom datetime import timedelta\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.models import SearchQuery\n\n\ndef create_search_query(\n db_session: Session,\n user_id: UUID,\n query: str,\n query_expansions: list[str] | None = None,\n) -> SearchQuery:\n \"\"\"Create and persist a `SearchQuery` row.\n\n Notes:\n - `SearchQuery.id` is a UUID PK without a server-side default, so we generate it.\n - `created_at` is filled by the DB (server_default=now()).\n \"\"\"\n search_query = SearchQuery(\n id=uuid.uuid4(),\n user_id=user_id,\n query=query,\n query_expansions=query_expansions,\n )\n db_session.add(search_query)\n db_session.commit()\n db_session.refresh(search_query)\n return search_query\n\n\ndef fetch_search_queries_for_user(\n db_session: Session,\n user_id: UUID,\n filter_days: int | None = None,\n limit: int | None = None,\n) -> list[SearchQuery]:\n \"\"\"Fetch `SearchQuery` rows for a user.\n\n Args:\n user_id: User UUID.\n filter_days: Optional time filter. If provided, only rows created within\n the last `filter_days` days are returned.\n limit: Optional max number of rows to return.\n \"\"\"\n if filter_days is not None and filter_days <= 0:\n raise ValueError(\"filter_days must be > 0\")\n\n stmt = select(SearchQuery).where(SearchQuery.user_id == user_id)\n\n if filter_days is not None and filter_days > 0:\n cutoff = get_db_current_time(db_session) - timedelta(days=filter_days)\n stmt = stmt.where(SearchQuery.created_at >= cutoff)\n\n stmt = stmt.order_by(SearchQuery.created_at.desc())\n\n if limit is not None:\n stmt = stmt.limit(limit)\n\n return list(db_session.scalars(stmt).all())\n" }, { "path": "backend/ee/onyx/db/standard_answer.py", "content": "import re\nimport string\nfrom collections.abc import Sequence\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import StandardAnswer\nfrom onyx.db.models import StandardAnswerCategory\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef check_category_validity(category_name: str) -> bool:\n \"\"\"If a category name is too long, it should not be used (it will cause an error in Postgres\n as the unique constraint can only apply to entries that are less than 2704 bytes).\n\n Additionally, extremely long categories are not really usable / useful.\"\"\"\n if len(category_name) > 255:\n logger.error(\n f\"Category with name '{category_name}' is too long, cannot be used\"\n )\n return False\n\n return True\n\n\ndef insert_standard_answer_category(\n category_name: str, db_session: Session\n) -> StandardAnswerCategory:\n if not check_category_validity(category_name):\n raise ValueError(f\"Invalid category name: {category_name}\")\n standard_answer_category = StandardAnswerCategory(name=category_name)\n db_session.add(standard_answer_category)\n db_session.commit()\n\n return standard_answer_category\n\n\ndef insert_standard_answer(\n keyword: str,\n answer: str,\n category_ids: list[int],\n match_regex: bool,\n match_any_keywords: bool,\n db_session: Session,\n) -> StandardAnswer:\n existing_categories = fetch_standard_answer_categories_by_ids(\n standard_answer_category_ids=category_ids,\n db_session=db_session,\n )\n if len(existing_categories) != len(category_ids):\n raise ValueError(f\"Some or all categories with ids {category_ids} do not exist\")\n\n standard_answer = StandardAnswer(\n keyword=keyword,\n answer=answer,\n categories=existing_categories,\n active=True,\n match_regex=match_regex,\n match_any_keywords=match_any_keywords,\n )\n db_session.add(standard_answer)\n db_session.commit()\n return standard_answer\n\n\ndef update_standard_answer(\n standard_answer_id: int,\n keyword: str,\n answer: str,\n category_ids: list[int],\n match_regex: bool,\n match_any_keywords: bool,\n db_session: Session,\n) -> StandardAnswer:\n standard_answer = db_session.scalar(\n select(StandardAnswer).where(StandardAnswer.id == standard_answer_id)\n )\n if standard_answer is None:\n raise ValueError(f\"No standard answer with id {standard_answer_id}\")\n\n existing_categories = fetch_standard_answer_categories_by_ids(\n standard_answer_category_ids=category_ids,\n db_session=db_session,\n )\n if len(existing_categories) != len(category_ids):\n raise ValueError(f\"Some or all categories with ids {category_ids} do not exist\")\n\n standard_answer.keyword = keyword\n standard_answer.answer = answer\n standard_answer.categories = list(existing_categories)\n standard_answer.match_regex = match_regex\n standard_answer.match_any_keywords = match_any_keywords\n\n db_session.commit()\n\n return standard_answer\n\n\ndef remove_standard_answer(\n standard_answer_id: int,\n db_session: Session,\n) -> None:\n standard_answer = db_session.scalar(\n select(StandardAnswer).where(StandardAnswer.id == standard_answer_id)\n )\n if standard_answer is None:\n raise ValueError(f\"No standard answer with id {standard_answer_id}\")\n\n standard_answer.active = False\n db_session.commit()\n\n\ndef update_standard_answer_category(\n standard_answer_category_id: int,\n category_name: str,\n db_session: Session,\n) -> StandardAnswerCategory:\n standard_answer_category = db_session.scalar(\n select(StandardAnswerCategory).where(\n StandardAnswerCategory.id == standard_answer_category_id\n )\n )\n if standard_answer_category is None:\n raise ValueError(\n f\"No standard answer category with id {standard_answer_category_id}\"\n )\n\n if not check_category_validity(category_name):\n raise ValueError(f\"Invalid category name: {category_name}\")\n\n standard_answer_category.name = category_name\n\n db_session.commit()\n\n return standard_answer_category\n\n\ndef fetch_standard_answer_category(\n standard_answer_category_id: int,\n db_session: Session,\n) -> StandardAnswerCategory | None:\n return db_session.scalar(\n select(StandardAnswerCategory).where(\n StandardAnswerCategory.id == standard_answer_category_id\n )\n )\n\n\ndef fetch_standard_answer_categories_by_ids(\n standard_answer_category_ids: list[int],\n db_session: Session,\n) -> Sequence[StandardAnswerCategory]:\n return db_session.scalars(\n select(StandardAnswerCategory).where(\n StandardAnswerCategory.id.in_(standard_answer_category_ids)\n )\n ).all()\n\n\ndef fetch_standard_answer_categories(\n db_session: Session,\n) -> Sequence[StandardAnswerCategory]:\n return db_session.scalars(select(StandardAnswerCategory)).all()\n\n\ndef fetch_standard_answer(\n standard_answer_id: int,\n db_session: Session,\n) -> StandardAnswer | None:\n return db_session.scalar(\n select(StandardAnswer).where(StandardAnswer.id == standard_answer_id)\n )\n\n\ndef fetch_standard_answers(db_session: Session) -> Sequence[StandardAnswer]:\n return db_session.scalars(\n select(StandardAnswer).where(StandardAnswer.active.is_(True))\n ).all()\n\n\ndef create_initial_default_standard_answer_category(db_session: Session) -> None:\n default_category_id = 0\n default_category_name = \"General\"\n default_category = fetch_standard_answer_category(\n standard_answer_category_id=default_category_id,\n db_session=db_session,\n )\n if default_category is not None:\n if default_category.name != default_category_name:\n raise ValueError(\n \"DB is not in a valid initial state. Default standard answer category does not have expected name.\"\n )\n return\n\n standard_answer_category = StandardAnswerCategory(\n id=default_category_id,\n name=default_category_name,\n )\n db_session.add(standard_answer_category)\n db_session.commit()\n\n\ndef fetch_standard_answer_categories_by_names(\n standard_answer_category_names: list[str],\n db_session: Session,\n) -> Sequence[StandardAnswerCategory]:\n return db_session.scalars(\n select(StandardAnswerCategory).where(\n StandardAnswerCategory.name.in_(standard_answer_category_names)\n )\n ).all()\n\n\ndef find_matching_standard_answers(\n id_in: list[int],\n query: str,\n db_session: Session,\n) -> list[tuple[StandardAnswer, str]]:\n \"\"\"\n Returns a list of tuples, where each tuple is a StandardAnswer definition matching\n the query and a string representing the match (either the regex match group or the\n set of keywords).\n\n If `answer_instance.match_regex` is true, the definition is considered \"matched\"\n if the query matches the `answer_instance.keyword` using `re.search`.\n\n Otherwise, the definition is considered \"matched\" if the space-delimited tokens\n in `keyword` exists in `query`, depending on the state of `match_any_keywords`\n \"\"\"\n stmt = (\n select(StandardAnswer)\n .where(StandardAnswer.active.is_(True))\n .where(StandardAnswer.id.in_(id_in))\n )\n possible_standard_answers: Sequence[StandardAnswer] = db_session.scalars(stmt).all()\n\n matching_standard_answers: list[tuple[StandardAnswer, str]] = []\n for standard_answer in possible_standard_answers:\n if standard_answer.match_regex:\n maybe_matches = re.search(standard_answer.keyword, query, re.IGNORECASE)\n if maybe_matches is not None:\n match_group = maybe_matches.group(0)\n matching_standard_answers.append((standard_answer, match_group))\n\n else:\n # Remove punctuation and split the keyword into individual words\n keyword_words = set(\n \"\".join(\n char\n for char in standard_answer.keyword.lower()\n if char not in string.punctuation\n ).split()\n )\n\n # Remove punctuation and split the query into individual words\n query_words = \"\".join(\n char for char in query.lower() if char not in string.punctuation\n ).split()\n\n # Check if all of the keyword words are in the query words\n if standard_answer.match_any_keywords:\n for word in query_words:\n if word in keyword_words:\n matching_standard_answers.append((standard_answer, word))\n break\n else:\n if all(word in query_words for word in keyword_words):\n matching_standard_answers.append(\n (\n standard_answer,\n re.sub(r\"\\s+?\", \", \", standard_answer.keyword),\n )\n )\n\n return matching_standard_answers\n" }, { "path": "backend/ee/onyx/db/token_limit.py", "content": "from collections.abc import Sequence\n\nfrom sqlalchemy import exists\nfrom sqlalchemy import Row\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import TokenRateLimitScope\nfrom onyx.db.models import TokenRateLimit\nfrom onyx.db.models import TokenRateLimit__UserGroup\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.db.models import UserRole\nfrom onyx.server.token_rate_limits.models import TokenRateLimitArgs\n\n\ndef _add_user_filters(stmt: Select, user: User, get_editable: bool = True) -> Select:\n if user.role == UserRole.ADMIN:\n return stmt\n\n # If anonymous user, only show global/public token_rate_limits\n if user.is_anonymous:\n where_clause = TokenRateLimit.scope == TokenRateLimitScope.GLOBAL\n return stmt.where(where_clause)\n\n stmt = stmt.distinct()\n TRLimit_UG = aliased(TokenRateLimit__UserGroup)\n User__UG = aliased(User__UserGroup)\n\n \"\"\"\n Here we select token_rate_limits by relation:\n User -> User__UserGroup -> TokenRateLimit__UserGroup ->\n TokenRateLimit\n \"\"\"\n stmt = stmt.outerjoin(TRLimit_UG).outerjoin(\n User__UG,\n User__UG.user_group_id == TRLimit_UG.user_group_id,\n )\n\n \"\"\"\n Filter token_rate_limits by:\n - if the user is in the user_group that owns the token_rate_limit\n - if the user is not a global_curator, they must also have a curator relationship\n to the user_group\n - if editing is being done, we also filter out token_rate_limits that are owned by groups\n that the user isn't a curator for\n - if we are not editing, we show all token_rate_limits in the groups the user curates\n \"\"\"\n\n where_clause = User__UG.user_id == user.id\n if user.role == UserRole.CURATOR and get_editable:\n where_clause &= User__UG.is_curator == True # noqa: E712\n if get_editable:\n user_groups = select(User__UG.user_group_id).where(User__UG.user_id == user.id)\n if user.role == UserRole.CURATOR:\n user_groups = user_groups.where(\n User__UserGroup.is_curator == True # noqa: E712\n )\n where_clause &= (\n ~exists()\n .where(TRLimit_UG.rate_limit_id == TokenRateLimit.id)\n .where(~TRLimit_UG.user_group_id.in_(user_groups))\n .correlate(TokenRateLimit)\n )\n\n return stmt.where(where_clause)\n\n\ndef fetch_all_user_group_token_rate_limits_by_group(\n db_session: Session,\n) -> Sequence[Row[tuple[TokenRateLimit, str]]]:\n query = (\n select(TokenRateLimit, UserGroup.name)\n .join(\n TokenRateLimit__UserGroup,\n TokenRateLimit.id == TokenRateLimit__UserGroup.rate_limit_id,\n )\n .join(UserGroup, UserGroup.id == TokenRateLimit__UserGroup.user_group_id)\n )\n\n return db_session.execute(query).all()\n\n\ndef insert_user_group_token_rate_limit(\n db_session: Session,\n token_rate_limit_settings: TokenRateLimitArgs,\n group_id: int,\n) -> TokenRateLimit:\n token_limit = TokenRateLimit(\n enabled=token_rate_limit_settings.enabled,\n token_budget=token_rate_limit_settings.token_budget,\n period_hours=token_rate_limit_settings.period_hours,\n scope=TokenRateLimitScope.USER_GROUP,\n )\n db_session.add(token_limit)\n db_session.flush()\n\n rate_limit = TokenRateLimit__UserGroup(\n rate_limit_id=token_limit.id, user_group_id=group_id\n )\n db_session.add(rate_limit)\n db_session.commit()\n\n return token_limit\n\n\ndef fetch_user_group_token_rate_limits_for_user(\n db_session: Session,\n group_id: int,\n user: User,\n enabled_only: bool = False,\n ordered: bool = True,\n get_editable: bool = True,\n) -> Sequence[TokenRateLimit]:\n stmt = (\n select(TokenRateLimit)\n .join(\n TokenRateLimit__UserGroup,\n TokenRateLimit.id == TokenRateLimit__UserGroup.rate_limit_id,\n )\n .where(TokenRateLimit__UserGroup.user_group_id == group_id)\n )\n stmt = _add_user_filters(stmt, user, get_editable)\n\n if enabled_only:\n stmt = stmt.where(TokenRateLimit.enabled.is_(True))\n\n if ordered:\n stmt = stmt.order_by(TokenRateLimit.created_at.desc())\n\n return db_session.scalars(stmt).all()\n" }, { "path": "backend/ee/onyx/db/usage_export.py", "content": "import uuid\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom typing import IO\nfrom typing import Optional\n\nfrom fastapi_users_db_sqlalchemy import UUID_ID\nfrom sqlalchemy import cast\nfrom sqlalchemy.dialects.postgresql import UUID\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.query_history import fetch_chat_sessions_eagerly_by_time\nfrom ee.onyx.server.reporting.usage_export_models import ChatMessageSkeleton\nfrom ee.onyx.server.reporting.usage_export_models import FlowType\nfrom ee.onyx.server.reporting.usage_export_models import UsageReportMetadata\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.models import UsageReport\nfrom onyx.db.models import User\nfrom onyx.file_store.file_store import get_default_file_store\n\n\n# Gets skeletons of all messages in the given range\ndef get_empty_chat_messages_entries__paginated(\n db_session: Session,\n period: tuple[datetime, datetime],\n limit: int | None = 500,\n initial_time: datetime | None = None,\n) -> tuple[Optional[datetime], list[ChatMessageSkeleton]]:\n \"\"\"Returns a tuple where:\n first element is the most recent timestamp out of the sessions iterated\n - this timestamp can be used to paginate forward in time\n second element is a list of messages belonging to all the sessions iterated\n\n Only messages of type USER are returned\n \"\"\"\n chat_sessions = fetch_chat_sessions_eagerly_by_time(\n start=period[0],\n end=period[1],\n db_session=db_session,\n limit=limit,\n initial_time=initial_time,\n )\n\n message_skeletons: list[ChatMessageSkeleton] = []\n for chat_session in chat_sessions:\n flow_type = FlowType.SLACK if chat_session.onyxbot_flow else FlowType.CHAT\n\n for message in chat_session.messages:\n # Only count user messages\n if message.message_type != MessageType.USER:\n continue\n\n # Get user email\n user_email = chat_session.user.email if chat_session.user else None\n\n # Get assistant name (from session persona, or alternate if specified)\n assistant_name = None\n if chat_session.persona:\n assistant_name = chat_session.persona.name\n\n message_skeletons.append(\n ChatMessageSkeleton(\n message_id=message.id,\n chat_session_id=chat_session.id,\n user_id=str(chat_session.user_id) if chat_session.user_id else None,\n flow_type=flow_type,\n time_sent=message.time_sent,\n assistant_name=assistant_name,\n user_email=user_email,\n number_of_tokens=message.token_count,\n )\n )\n if len(chat_sessions) == 0:\n return None, []\n\n return chat_sessions[-1].time_created, message_skeletons\n\n\ndef get_all_empty_chat_message_entries(\n db_session: Session,\n period: tuple[datetime, datetime],\n) -> Generator[list[ChatMessageSkeleton], None, None]:\n \"\"\"period is the range of time over which to fetch messages.\"\"\"\n initial_time: Optional[datetime] = period[0]\n while True:\n # iterate from oldest to newest\n time_created, message_skeletons = get_empty_chat_messages_entries__paginated(\n db_session,\n period,\n initial_time=initial_time,\n )\n\n if not message_skeletons:\n return\n\n yield message_skeletons\n\n # Update initial_time for the next iteration\n initial_time = time_created\n\n\ndef get_all_usage_reports(db_session: Session) -> list[UsageReportMetadata]:\n # Get the user emails\n usage_reports = db_session.query(UsageReport).all()\n user_ids = {r.requestor_user_id for r in usage_reports if r.requestor_user_id}\n user_emails = {\n user.id: user.email\n for user in db_session.query(User)\n .filter(cast(User.id, UUID).in_(user_ids))\n .all()\n }\n\n return [\n UsageReportMetadata(\n report_name=r.report_name,\n requestor=(\n user_emails.get(r.requestor_user_id) if r.requestor_user_id else None\n ),\n time_created=r.time_created,\n period_from=r.period_from,\n period_to=r.period_to,\n )\n for r in usage_reports\n ]\n\n\ndef get_usage_report_data(\n report_display_name: str,\n) -> IO:\n \"\"\"\n Get the usage report data from the file store.\n\n Args:\n db_session: The database session.\n report_display_name: The display name of the usage report. Also assumes\n that the file is stored with this as the ID in the file store.\n\n Returns:\n The usage report data.\n \"\"\"\n file_store = get_default_file_store()\n # usage report may be very large, so don't load it all into memory\n return file_store.read_file(\n file_id=report_display_name, mode=\"b\", use_tempfile=True\n )\n\n\ndef write_usage_report(\n db_session: Session,\n report_name: str,\n user_id: uuid.UUID | UUID_ID | None,\n period: tuple[datetime, datetime] | None,\n) -> UsageReport:\n new_report = UsageReport(\n report_name=report_name,\n requestor_user_id=user_id,\n period_from=period[0] if period else None,\n period_to=period[1] if period else None,\n )\n db_session.add(new_report)\n db_session.commit()\n return new_report\n" }, { "path": "backend/ee/onyx/db/user_group.py", "content": "from collections.abc import Sequence\nfrom operator import and_\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import delete\nfrom sqlalchemy import func\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.server.user_group.models import SetCuratorRequest\nfrom ee.onyx.server.user_group.models import UserGroupCreate\nfrom ee.onyx.server.user_group.models import UserGroupUpdate\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import GrantSource\nfrom onyx.db.enums import Permission\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import Credential__UserGroup\nfrom onyx.db.models import Document\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import DocumentSet__UserGroup\nfrom onyx.db.models import FederatedConnector__DocumentSet\nfrom onyx.db.models import LLMProvider__UserGroup\nfrom onyx.db.models import PermissionGrant\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__UserGroup\nfrom onyx.db.models import TokenRateLimit__UserGroup\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.db.models import UserGroup__ConnectorCredentialPair\nfrom onyx.db.models import UserRole\nfrom onyx.db.permissions import recompute_user_permissions__no_commit\nfrom onyx.db.users import fetch_user_by_id\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _cleanup_user__user_group_relationships__no_commit(\n db_session: Session,\n user_group_id: int,\n user_ids: list[UUID] | None = None,\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n where_clause = User__UserGroup.user_group_id == user_group_id\n if user_ids:\n where_clause &= User__UserGroup.user_id.in_(user_ids)\n\n user__user_group_relationships = db_session.scalars(\n select(User__UserGroup).where(where_clause)\n ).all()\n for user__user_group_relationship in user__user_group_relationships:\n db_session.delete(user__user_group_relationship)\n\n\ndef _cleanup_credential__user_group_relationships__no_commit(\n db_session: Session,\n user_group_id: int,\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n db_session.query(Credential__UserGroup).filter(\n Credential__UserGroup.user_group_id == user_group_id\n ).delete(synchronize_session=False)\n\n\ndef _cleanup_llm_provider__user_group_relationships__no_commit(\n db_session: Session, user_group_id: int\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n db_session.query(LLMProvider__UserGroup).filter(\n LLMProvider__UserGroup.user_group_id == user_group_id\n ).delete(synchronize_session=False)\n\n\ndef _cleanup_persona__user_group_relationships__no_commit(\n db_session: Session, user_group_id: int\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n db_session.query(Persona__UserGroup).filter(\n Persona__UserGroup.user_group_id == user_group_id\n ).delete(synchronize_session=False)\n\n\ndef _cleanup_token_rate_limit__user_group_relationships__no_commit(\n db_session: Session, user_group_id: int\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n token_rate_limit__user_group_relationships = db_session.scalars(\n select(TokenRateLimit__UserGroup).where(\n TokenRateLimit__UserGroup.user_group_id == user_group_id\n )\n ).all()\n for (\n token_rate_limit__user_group_relationship\n ) in token_rate_limit__user_group_relationships:\n db_session.delete(token_rate_limit__user_group_relationship)\n\n\ndef _cleanup_user_group__cc_pair_relationships__no_commit(\n db_session: Session, user_group_id: int, outdated_only: bool\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n stmt = select(UserGroup__ConnectorCredentialPair).where(\n UserGroup__ConnectorCredentialPair.user_group_id == user_group_id\n )\n if outdated_only:\n stmt = stmt.where(\n UserGroup__ConnectorCredentialPair.is_current == False # noqa: E712\n )\n user_group__cc_pair_relationships = db_session.scalars(stmt)\n for user_group__cc_pair_relationship in user_group__cc_pair_relationships:\n db_session.delete(user_group__cc_pair_relationship)\n\n\ndef _cleanup_document_set__user_group_relationships__no_commit(\n db_session: Session, user_group_id: int\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n db_session.execute(\n delete(DocumentSet__UserGroup).where(\n DocumentSet__UserGroup.user_group_id == user_group_id\n )\n )\n\n\ndef validate_object_creation_for_user(\n db_session: Session,\n user: User,\n target_group_ids: list[int] | None = None,\n object_is_public: bool | None = None,\n object_is_perm_sync: bool | None = None,\n object_is_owned_by_user: bool = False,\n object_is_new: bool = False,\n) -> None:\n \"\"\"\n All users can create/edit permission synced objects if they don't specify a group\n All admin actions are allowed.\n Curators and global curators can create public objects.\n Prevents other non-admins from creating/editing:\n - public objects\n - objects with no groups\n - objects that belong to a group they don't curate\n \"\"\"\n if object_is_perm_sync and not target_group_ids:\n return\n\n # Admins are allowed\n if user.role == UserRole.ADMIN:\n return\n\n # Allow curators and global curators to create public objects\n # w/o associated groups IF the object is new/owned by them\n if (\n object_is_public\n and user.role in [UserRole.CURATOR, UserRole.GLOBAL_CURATOR]\n and (object_is_new or object_is_owned_by_user)\n ):\n return\n\n if object_is_public and user.role == UserRole.BASIC:\n detail = \"User does not have permission to create public objects\"\n logger.error(detail)\n raise HTTPException(\n status_code=400,\n detail=detail,\n )\n\n if not target_group_ids:\n detail = \"Curators must specify 1+ groups\"\n logger.error(detail)\n raise HTTPException(\n status_code=400,\n detail=detail,\n )\n\n user_curated_groups = fetch_user_groups_for_user(\n db_session=db_session,\n user_id=user.id,\n # Global curators can curate all groups they are member of\n only_curator_groups=user.role != UserRole.GLOBAL_CURATOR,\n )\n user_curated_group_ids = set([group.id for group in user_curated_groups])\n target_group_ids_set = set(target_group_ids)\n if not target_group_ids_set.issubset(user_curated_group_ids):\n detail = \"Curators cannot control groups they don't curate\"\n logger.error(detail)\n raise HTTPException(\n status_code=400,\n detail=detail,\n )\n\n\ndef fetch_user_group(db_session: Session, user_group_id: int) -> UserGroup | None:\n stmt = select(UserGroup).where(UserGroup.id == user_group_id)\n return db_session.scalar(stmt)\n\n\ndef _add_user_group_snapshot_eager_loads(\n stmt: Select,\n) -> Select:\n \"\"\"Add eager loading options needed by UserGroup.from_model snapshot creation.\"\"\"\n return stmt.options(\n selectinload(UserGroup.users),\n selectinload(UserGroup.user_group_relationships),\n selectinload(UserGroup.cc_pair_relationships)\n .selectinload(UserGroup__ConnectorCredentialPair.cc_pair)\n .options(\n selectinload(ConnectorCredentialPair.connector),\n selectinload(ConnectorCredentialPair.credential).selectinload(\n Credential.user\n ),\n ),\n selectinload(UserGroup.document_sets).options(\n selectinload(DocumentSet.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSet.users),\n selectinload(DocumentSet.groups),\n selectinload(DocumentSet.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n ),\n selectinload(UserGroup.personas).options(\n selectinload(Persona.tools),\n selectinload(Persona.hierarchy_nodes),\n selectinload(Persona.attached_documents).selectinload(\n Document.parent_hierarchy_node\n ),\n selectinload(Persona.labels),\n selectinload(Persona.document_sets).options(\n selectinload(DocumentSet.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSet.users),\n selectinload(DocumentSet.groups),\n selectinload(DocumentSet.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n ),\n selectinload(Persona.user),\n selectinload(Persona.user_files),\n selectinload(Persona.users),\n selectinload(Persona.groups),\n ),\n )\n\n\ndef fetch_user_groups(\n db_session: Session,\n only_up_to_date: bool = True,\n eager_load_for_snapshot: bool = False,\n include_default: bool = True,\n) -> Sequence[UserGroup]:\n \"\"\"\n Fetches user groups from the database.\n\n This function retrieves a sequence of `UserGroup` objects from the database.\n If `only_up_to_date` is set to `True`, it filters the user groups to return only those\n that are marked as up-to-date (`is_up_to_date` is `True`).\n\n Args:\n db_session (Session): The SQLAlchemy session used to query the database.\n only_up_to_date (bool, optional): Flag to determine whether to filter the results\n to include only up to date user groups. Defaults to `True`.\n eager_load_for_snapshot: If True, adds eager loading for all relationships\n needed by UserGroup.from_model snapshot creation.\n include_default: If False, excludes system default groups (is_default=True).\n\n Returns:\n Sequence[UserGroup]: A sequence of `UserGroup` objects matching the query criteria.\n \"\"\"\n stmt = select(UserGroup)\n if only_up_to_date:\n stmt = stmt.where(UserGroup.is_up_to_date == True) # noqa: E712\n if not include_default:\n stmt = stmt.where(UserGroup.is_default == False) # noqa: E712\n if eager_load_for_snapshot:\n stmt = _add_user_group_snapshot_eager_loads(stmt)\n return db_session.scalars(stmt).unique().all()\n\n\ndef fetch_user_groups_for_user(\n db_session: Session,\n user_id: UUID,\n only_curator_groups: bool = False,\n eager_load_for_snapshot: bool = False,\n include_default: bool = True,\n) -> Sequence[UserGroup]:\n stmt = (\n select(UserGroup)\n .join(User__UserGroup, User__UserGroup.user_group_id == UserGroup.id)\n .join(User, User.id == User__UserGroup.user_id) # type: ignore\n .where(User.id == user_id) # type: ignore\n )\n if only_curator_groups:\n stmt = stmt.where(User__UserGroup.is_curator == True) # noqa: E712\n if not include_default:\n stmt = stmt.where(UserGroup.is_default == False) # noqa: E712\n if eager_load_for_snapshot:\n stmt = _add_user_group_snapshot_eager_loads(stmt)\n return db_session.scalars(stmt).unique().all()\n\n\ndef construct_document_id_select_by_usergroup(\n user_group_id: int,\n) -> Select:\n \"\"\"This returns a statement that should be executed using\n .yield_per() to minimize overhead. The primary consumers of this function\n are background processing task generators.\"\"\"\n stmt = (\n select(Document.id)\n .join(\n DocumentByConnectorCredentialPair,\n Document.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n UserGroup__ConnectorCredentialPair,\n UserGroup__ConnectorCredentialPair.cc_pair_id == ConnectorCredentialPair.id,\n )\n .join(\n UserGroup,\n UserGroup__ConnectorCredentialPair.user_group_id == UserGroup.id,\n )\n .where(UserGroup.id == user_group_id)\n .order_by(Document.id)\n )\n stmt = stmt.distinct()\n return stmt\n\n\ndef fetch_documents_for_user_group_paginated(\n db_session: Session,\n user_group_id: int,\n last_document_id: str | None = None,\n limit: int = 100,\n) -> tuple[Sequence[Document], str | None]:\n stmt = (\n select(Document)\n .join(\n DocumentByConnectorCredentialPair,\n Document.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n UserGroup__ConnectorCredentialPair,\n UserGroup__ConnectorCredentialPair.cc_pair_id == ConnectorCredentialPair.id,\n )\n .join(\n UserGroup,\n UserGroup__ConnectorCredentialPair.user_group_id == UserGroup.id,\n )\n .where(UserGroup.id == user_group_id)\n .order_by(Document.id)\n .limit(limit)\n )\n if last_document_id is not None:\n stmt = stmt.where(Document.id > last_document_id)\n stmt = stmt.distinct()\n\n documents = db_session.scalars(stmt).all()\n return documents, documents[-1].id if documents else None\n\n\ndef fetch_user_groups_for_documents(\n db_session: Session,\n document_ids: list[str],\n) -> Sequence[tuple[str, list[str]]]:\n \"\"\"\n Fetches all user groups that have access to the given documents.\n\n NOTE: this doesn't include groups if the cc_pair is access type SYNC\n \"\"\"\n stmt = (\n select(Document.id, func.array_agg(UserGroup.name))\n .join(\n UserGroup__ConnectorCredentialPair,\n UserGroup.id == UserGroup__ConnectorCredentialPair.user_group_id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n ConnectorCredentialPair.id\n == UserGroup__ConnectorCredentialPair.cc_pair_id,\n ConnectorCredentialPair.access_type != AccessType.SYNC,\n ),\n )\n .join(\n DocumentByConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(Document, Document.id == DocumentByConnectorCredentialPair.id)\n .where(Document.id.in_(document_ids))\n .where(UserGroup__ConnectorCredentialPair.is_current == True) # noqa: E712\n # don't include CC pairs that are being deleted\n # NOTE: CC pairs can never go from DELETING to any other state -> it's safe to ignore them\n .where(ConnectorCredentialPair.status != ConnectorCredentialPairStatus.DELETING)\n .group_by(Document.id)\n )\n\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef _check_user_group_is_modifiable(user_group: UserGroup) -> None:\n if not user_group.is_up_to_date:\n raise ValueError(\n \"Specified user group is currently syncing. Wait until the current sync has finished before editing.\"\n )\n\n\ndef _add_user__user_group_relationships__no_commit(\n db_session: Session, user_group_id: int, user_ids: list[UUID]\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\n\n This function is idempotent - it will skip users who are already in the group\n to avoid duplicate key violations during concurrent operations or re-syncs.\n Uses ON CONFLICT DO NOTHING to keep inserts atomic under concurrency.\n \"\"\"\n if not user_ids:\n return\n\n insert_stmt = (\n insert(User__UserGroup)\n .values(\n [\n {\"user_id\": user_id, \"user_group_id\": user_group_id}\n for user_id in user_ids\n ]\n )\n .on_conflict_do_nothing(\n index_elements=[User__UserGroup.user_group_id, User__UserGroup.user_id]\n )\n )\n db_session.execute(insert_stmt)\n\n\ndef _add_user_group__cc_pair_relationships__no_commit(\n db_session: Session, user_group_id: int, cc_pair_ids: list[int]\n) -> list[UserGroup__ConnectorCredentialPair]:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n relationships = [\n UserGroup__ConnectorCredentialPair(\n user_group_id=user_group_id, cc_pair_id=cc_pair_id\n )\n for cc_pair_id in cc_pair_ids\n ]\n db_session.add_all(relationships)\n return relationships\n\n\ndef insert_user_group(db_session: Session, user_group: UserGroupCreate) -> UserGroup:\n db_user_group = UserGroup(\n name=user_group.name,\n time_last_modified_by_user=func.now(),\n is_up_to_date=DISABLE_VECTOR_DB,\n )\n db_session.add(db_user_group)\n db_session.flush() # give the group an ID\n\n # Every group gets the \"basic\" permission by default\n db_session.add(\n PermissionGrant(\n group_id=db_user_group.id,\n permission=Permission.BASIC_ACCESS,\n grant_source=GrantSource.SYSTEM,\n )\n )\n db_session.flush()\n\n _add_user__user_group_relationships__no_commit(\n db_session=db_session,\n user_group_id=db_user_group.id,\n user_ids=user_group.user_ids,\n )\n _add_user_group__cc_pair_relationships__no_commit(\n db_session=db_session,\n user_group_id=db_user_group.id,\n cc_pair_ids=user_group.cc_pair_ids,\n )\n\n recompute_user_permissions__no_commit(user_group.user_ids, db_session)\n\n db_session.commit()\n return db_user_group\n\n\ndef _mark_user_group__cc_pair_relationships_outdated__no_commit(\n db_session: Session, user_group_id: int\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n user_group__cc_pair_relationships = db_session.scalars(\n select(UserGroup__ConnectorCredentialPair).where(\n UserGroup__ConnectorCredentialPair.user_group_id == user_group_id\n )\n )\n for user_group__cc_pair_relationship in user_group__cc_pair_relationships:\n user_group__cc_pair_relationship.is_current = False\n\n\ndef _validate_curator_status__no_commit(\n db_session: Session,\n users: list[User],\n) -> None:\n for user in users:\n # Check if the user is a curator in any of their groups\n curator_relationships = (\n db_session.query(User__UserGroup)\n .filter(\n User__UserGroup.user_id == user.id,\n User__UserGroup.is_curator == True, # noqa: E712\n )\n .all()\n )\n\n # if the user is a curator in any of their groups, set their role to CURATOR\n # otherwise, set their role to BASIC only if they were previously a CURATOR\n if curator_relationships:\n user.role = UserRole.CURATOR\n elif user.role == UserRole.CURATOR:\n user.role = UserRole.BASIC\n db_session.add(user)\n\n\ndef remove_curator_status__no_commit(db_session: Session, user: User) -> None:\n stmt = (\n update(User__UserGroup)\n .where(User__UserGroup.user_id == user.id)\n .values(is_curator=False)\n )\n db_session.execute(stmt)\n _validate_curator_status__no_commit(db_session, [user])\n\n\ndef _validate_curator_relationship_update_requester(\n db_session: Session,\n user_group_id: int,\n user_making_change: User,\n) -> None:\n \"\"\"\n This function validates that the user making the change has the necessary permissions\n to update the curator relationship for the target user in the given user group.\n \"\"\"\n\n # Admins can update curator relationships for any group\n if user_making_change.role == UserRole.ADMIN:\n return\n\n # check if the user making the change is a curator in the group they are changing the curator relationship for\n user_making_change_curator_groups = fetch_user_groups_for_user(\n db_session=db_session,\n user_id=user_making_change.id,\n # only check if the user making the change is a curator if they are a curator\n # otherwise, they are a global_curator and can update the curator relationship\n # for any group they are a member of\n only_curator_groups=user_making_change.role == UserRole.CURATOR,\n )\n requestor_curator_group_ids = [\n group.id for group in user_making_change_curator_groups\n ]\n if user_group_id not in requestor_curator_group_ids:\n raise ValueError(\n f\"user making change {user_making_change.email} is not a curator,\"\n f\" admin, or global_curator for group '{user_group_id}'\"\n )\n\n\ndef _validate_curator_relationship_update_request(\n db_session: Session,\n user_group_id: int,\n target_user: User,\n) -> None:\n \"\"\"\n This function validates that the curator_relationship_update request itself is valid.\n \"\"\"\n if target_user.role == UserRole.ADMIN:\n raise ValueError(\n f\"User '{target_user.email}' is an admin and therefore has all permissions \"\n \"of a curator. If you'd like this user to only have curator permissions, \"\n \"you must update their role to BASIC then assign them to be CURATOR in the \"\n \"appropriate groups.\"\n )\n elif target_user.role == UserRole.GLOBAL_CURATOR:\n raise ValueError(\n f\"User '{target_user.email}' is a global_curator and therefore has all \"\n \"permissions of a curator for all groups. If you'd like this user to only \"\n \"have curator permissions for a specific group, you must update their role \"\n \"to BASIC then assign them to be CURATOR in the appropriate groups.\"\n )\n elif target_user.role not in [UserRole.CURATOR, UserRole.BASIC]:\n raise ValueError(\n f\"This endpoint can only be used to update the curator relationship for \"\n \"users with the CURATOR or BASIC role. \\n\"\n f\"Target user: {target_user.email} \\n\"\n f\"Target user role: {target_user.role} \\n\"\n )\n\n # check if the target user is in the group they are changing the curator relationship for\n requested_user_groups = fetch_user_groups_for_user(\n db_session=db_session,\n user_id=target_user.id,\n only_curator_groups=False,\n )\n group_ids = [group.id for group in requested_user_groups]\n if user_group_id not in group_ids:\n raise ValueError(\n f\"target user {target_user.email} is not in group '{user_group_id}'\"\n )\n\n\ndef update_user_curator_relationship(\n db_session: Session,\n user_group_id: int,\n set_curator_request: SetCuratorRequest,\n user_making_change: User,\n) -> None:\n target_user = fetch_user_by_id(db_session, set_curator_request.user_id)\n if not target_user:\n raise ValueError(f\"User with id '{set_curator_request.user_id}' not found\")\n\n _validate_curator_relationship_update_request(\n db_session=db_session,\n user_group_id=user_group_id,\n target_user=target_user,\n )\n\n _validate_curator_relationship_update_requester(\n db_session=db_session,\n user_group_id=user_group_id,\n user_making_change=user_making_change,\n )\n\n logger.info(\n f\"user_making_change={user_making_change.email if user_making_change else 'None'} is \"\n f\"updating the curator relationship for user={target_user.email} \"\n f\"in group={user_group_id} to is_curator={set_curator_request.is_curator}\"\n )\n\n relationship_to_update = (\n db_session.query(User__UserGroup)\n .filter(\n User__UserGroup.user_group_id == user_group_id,\n User__UserGroup.user_id == set_curator_request.user_id,\n )\n .first()\n )\n\n if relationship_to_update:\n relationship_to_update.is_curator = set_curator_request.is_curator\n else:\n relationship_to_update = User__UserGroup(\n user_group_id=user_group_id,\n user_id=set_curator_request.user_id,\n is_curator=True,\n )\n db_session.add(relationship_to_update)\n\n _validate_curator_status__no_commit(db_session, [target_user])\n db_session.commit()\n\n\ndef add_users_to_user_group(\n db_session: Session,\n user: User,\n user_group_id: int,\n user_ids: list[UUID],\n) -> UserGroup:\n db_user_group = fetch_user_group(db_session=db_session, user_group_id=user_group_id)\n if db_user_group is None:\n raise ValueError(f\"UserGroup with id '{user_group_id}' not found\")\n\n missing_users = [\n user_id for user_id in user_ids if fetch_user_by_id(db_session, user_id) is None\n ]\n if missing_users:\n raise ValueError(\n f\"User(s) not found: {', '.join(str(user_id) for user_id in missing_users)}\"\n )\n\n _check_user_group_is_modifiable(db_user_group)\n\n current_user_ids = [user.id for user in db_user_group.users]\n current_user_ids_set = set(current_user_ids)\n new_user_ids = [\n user_id for user_id in user_ids if user_id not in current_user_ids_set\n ]\n\n if not new_user_ids:\n return db_user_group\n\n user_group_update = UserGroupUpdate(\n user_ids=current_user_ids + new_user_ids,\n cc_pair_ids=[cc_pair.id for cc_pair in db_user_group.cc_pairs],\n )\n\n return update_user_group(\n db_session=db_session,\n user=user,\n user_group_id=user_group_id,\n user_group_update=user_group_update,\n )\n\n\ndef update_user_group(\n db_session: Session,\n user: User, # noqa: ARG001\n user_group_id: int,\n user_group_update: UserGroupUpdate,\n) -> UserGroup:\n \"\"\"If successful, this can set db_user_group.is_up_to_date = False.\n That will be processed by check_for_vespa_user_groups_sync_task and trigger\n a long running background sync to Vespa.\n \"\"\"\n stmt = select(UserGroup).where(UserGroup.id == user_group_id)\n db_user_group = db_session.scalar(stmt)\n if db_user_group is None:\n raise ValueError(f\"UserGroup with id '{user_group_id}' not found\")\n\n _check_user_group_is_modifiable(db_user_group)\n\n current_user_ids = set([user.id for user in db_user_group.users])\n updated_user_ids = set(user_group_update.user_ids)\n added_user_ids = list(updated_user_ids - current_user_ids)\n removed_user_ids = list(current_user_ids - updated_user_ids)\n\n if added_user_ids:\n missing_users = [\n user_id\n for user_id in added_user_ids\n if fetch_user_by_id(db_session, user_id) is None\n ]\n if missing_users:\n raise ValueError(\n f\"User(s) not found: {', '.join(str(user_id) for user_id in missing_users)}\"\n )\n\n # LEAVING THIS HERE FOR NOW FOR GIVING DIFFERENT ROLES\n # ACCESS TO DIFFERENT PERMISSIONS\n # if (removed_user_ids or added_user_ids) and (\n # not user or user.role != UserRole.ADMIN\n # ):\n # raise ValueError(\"Only admins can add or remove users from user groups\")\n\n if removed_user_ids:\n _cleanup_user__user_group_relationships__no_commit(\n db_session=db_session,\n user_group_id=user_group_id,\n user_ids=removed_user_ids,\n )\n\n if added_user_ids:\n _add_user__user_group_relationships__no_commit(\n db_session=db_session,\n user_group_id=user_group_id,\n user_ids=added_user_ids,\n )\n\n cc_pairs_updated = set([cc_pair.id for cc_pair in db_user_group.cc_pairs]) != set(\n user_group_update.cc_pair_ids\n )\n if cc_pairs_updated:\n _mark_user_group__cc_pair_relationships_outdated__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n _add_user_group__cc_pair_relationships__no_commit(\n db_session=db_session,\n user_group_id=db_user_group.id,\n cc_pair_ids=user_group_update.cc_pair_ids,\n )\n\n if cc_pairs_updated and not DISABLE_VECTOR_DB:\n db_user_group.is_up_to_date = False\n\n removed_users = db_session.scalars(\n select(User).where(User.id.in_(removed_user_ids)) # type: ignore\n ).unique()\n\n # Filter out admin and global curator users before validating curator status\n users_to_validate = [\n user\n for user in removed_users\n if user.role not in [UserRole.ADMIN, UserRole.GLOBAL_CURATOR]\n ]\n\n if users_to_validate:\n _validate_curator_status__no_commit(db_session, users_to_validate)\n\n # update \"time_updated\" to now\n db_user_group.time_last_modified_by_user = func.now()\n\n recompute_user_permissions__no_commit(\n list(set(added_user_ids) | set(removed_user_ids)), db_session\n )\n\n db_session.commit()\n return db_user_group\n\n\ndef rename_user_group(\n db_session: Session,\n user_group_id: int,\n new_name: str,\n) -> UserGroup:\n stmt = select(UserGroup).where(UserGroup.id == user_group_id)\n db_user_group = db_session.scalar(stmt)\n if db_user_group is None:\n raise ValueError(f\"UserGroup with id '{user_group_id}' not found\")\n\n _check_user_group_is_modifiable(db_user_group)\n\n db_user_group.name = new_name\n db_user_group.time_last_modified_by_user = func.now()\n\n # CC pair documents in Vespa contain the group name, so we need to\n # trigger a sync to update them with the new name.\n _mark_user_group__cc_pair_relationships_outdated__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n if not DISABLE_VECTOR_DB:\n db_user_group.is_up_to_date = False\n\n db_session.commit()\n return db_user_group\n\n\ndef prepare_user_group_for_deletion(db_session: Session, user_group_id: int) -> None:\n stmt = select(UserGroup).where(UserGroup.id == user_group_id)\n db_user_group = db_session.scalar(stmt)\n if db_user_group is None:\n raise ValueError(f\"UserGroup with id '{user_group_id}' not found\")\n\n _check_user_group_is_modifiable(db_user_group)\n\n # Collect affected user IDs before cleanup deletes the relationships\n affected_user_ids: list[UUID] = [\n uid\n for uid in db_session.execute(\n select(User__UserGroup.user_id).where(\n User__UserGroup.user_group_id == user_group_id\n )\n )\n .scalars()\n .all()\n if uid is not None\n ]\n\n _mark_user_group__cc_pair_relationships_outdated__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n\n _cleanup_credential__user_group_relationships__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n _cleanup_user__user_group_relationships__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n _cleanup_token_rate_limit__user_group_relationships__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n _cleanup_document_set__user_group_relationships__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n _cleanup_persona__user_group_relationships__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n _cleanup_user_group__cc_pair_relationships__no_commit(\n db_session=db_session,\n user_group_id=user_group_id,\n outdated_only=False,\n )\n _cleanup_llm_provider__user_group_relationships__no_commit(\n db_session=db_session, user_group_id=user_group_id\n )\n\n # Recompute permissions for affected users now that their\n # membership in this group has been removed\n recompute_user_permissions__no_commit(affected_user_ids, db_session)\n\n db_user_group.is_up_to_date = False\n db_user_group.is_up_for_deletion = True\n db_session.commit()\n\n\ndef delete_user_group(db_session: Session, user_group: UserGroup) -> None:\n \"\"\"\n This assumes that all the fk cleanup has already been done.\n \"\"\"\n db_session.delete(user_group)\n db_session.commit()\n\n\ndef mark_user_group_as_synced(db_session: Session, user_group: UserGroup) -> None:\n # cleanup outdated relationships\n _cleanup_user_group__cc_pair_relationships__no_commit(\n db_session=db_session, user_group_id=user_group.id, outdated_only=True\n )\n user_group.is_up_to_date = True\n db_session.commit()\n\n\ndef delete_user_group_cc_pair_relationship__no_commit(\n cc_pair_id: int, db_session: Session\n) -> None:\n \"\"\"Deletes all rows from UserGroup__ConnectorCredentialPair where the\n connector_credential_pair_id matches the given cc_pair_id.\n\n Should be used very carefully (only for connectors that are being deleted).\"\"\"\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n raise ValueError(f\"Connector Credential Pair '{cc_pair_id}' does not exist\")\n\n if cc_pair.status != ConnectorCredentialPairStatus.DELETING:\n raise ValueError(\n f\"Connector Credential Pair '{cc_pair_id}' is not in the DELETING state. status={cc_pair.status}\"\n )\n\n delete_stmt = delete(UserGroup__ConnectorCredentialPair).where(\n UserGroup__ConnectorCredentialPair.cc_pair_id == cc_pair_id,\n )\n db_session.execute(delete_stmt)\n" }, { "path": "backend/ee/onyx/document_index/vespa/app_config/cloud-services.xml.jinja", "content": "\n\n \n \n \n \n \n \n \n \n \n\n\n \n \n \n \n{{ document_elements }}\n \n \n \n \n \n \n \n \n \n 2\n \n \n \n \n \n\n \n 3\n 750\n 350\n 300\n \n\n\n 2\n\n \n\n" }, { "path": "backend/ee/onyx/external_permissions/__init__.py", "content": "" }, { "path": "backend/ee/onyx/external_permissions/confluence/__init__.py", "content": "" }, { "path": "backend/ee/onyx/external_permissions/confluence/constants.py", "content": "# This is a group that we use to store all the users that we found in Confluence\n# Instead of setting a page to public, we just add this group so that the page\n# is only accessible to users who have confluence accounts.\nALL_CONF_EMAILS_GROUP_NAME = \"All_Confluence_Users_Found_By_Onyx\"\n\nVIEWSPACE_PERMISSION_TYPE = \"VIEWSPACE\"\nREQUEST_PAGINATION_LIMIT = 5000\n" }, { "path": "backend/ee/onyx/external_permissions/confluence/doc_sync.py", "content": "\"\"\"\nRules defined here:\nhttps://confluence.atlassian.com/conf85/check-who-can-view-a-page-1283360557.html\n\"\"\"\n\nfrom collections.abc import Generator\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom ee.onyx.external_permissions.utils import generic_doc_sync\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.connector import ConfluenceConnector\nfrom onyx.connectors.credentials_provider import OnyxDBCredentialsProvider\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\nCONFLUENCE_DOC_SYNC_LABEL = \"confluence_doc_sync\"\n\n\ndef confluence_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,\n callback: IndexingHeartbeatInterface | None,\n) -> Generator[ElementExternalAccess, None, None]:\n \"\"\"\n Fetches document permissions from Confluence and yields DocExternalAccess objects.\n Compares fetched documents against existing documents in the DB for the connector.\n If a document exists in the DB but not in the Confluence fetch, it's marked as restricted.\n \"\"\"\n confluence_connector = ConfluenceConnector(\n **cc_pair.connector.connector_specific_config\n )\n\n provider = OnyxDBCredentialsProvider(\n get_current_tenant_id(), \"confluence\", cc_pair.credential_id\n )\n confluence_connector.set_credentials_provider(provider)\n\n yield from generic_doc_sync(\n cc_pair=cc_pair,\n fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,\n callback=callback,\n doc_source=DocumentSource.CONFLUENCE,\n slim_connector=confluence_connector,\n label=CONFLUENCE_DOC_SYNC_LABEL,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/confluence/group_sync.py", "content": "from collections.abc import Generator\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom ee.onyx.external_permissions.confluence.constants import ALL_CONF_EMAILS_GROUP_NAME\nfrom onyx.background.error_logging import emit_background_error\nfrom onyx.configs.app_configs import CONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC\nfrom onyx.connectors.confluence.onyx_confluence import (\n get_user_email_from_username__server,\n)\nfrom onyx.connectors.confluence.onyx_confluence import OnyxConfluence\nfrom onyx.connectors.credentials_provider import OnyxDBCredentialsProvider\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.users import get_all_users\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _build_group_member_email_map(\n confluence_client: OnyxConfluence, cc_pair_id: int\n) -> dict[str, set[str]]:\n group_member_emails: dict[str, set[str]] = {}\n for user in confluence_client.paginated_cql_user_retrieval():\n logger.info(f\"Processing groups for user: {user}\")\n\n email = user.email\n if not email:\n # This field is only present in Confluence Server\n user_name = user.username\n # If it is present, try to get the email using a Server-specific method\n if user_name:\n email = get_user_email_from_username__server(\n confluence_client=confluence_client,\n user_name=user_name,\n )\n else:\n logger.error(f\"user result missing username field: {user}\")\n\n if not email:\n # If we still don't have an email, skip this user\n msg = f\"user result missing email field: {user}\"\n if user.type == \"app\":\n logger.warning(msg)\n else:\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n logger.error(msg)\n continue\n\n all_users_groups: set[str] = set()\n for group in confluence_client.paginated_groups_by_user_retrieval(user.user_id):\n # group name uniqueness is enforced by Confluence, so we can use it as a group ID\n group_id = group[\"name\"]\n group_member_emails.setdefault(group_id, set()).add(email)\n all_users_groups.add(group_id)\n\n if not all_users_groups:\n msg = f\"No groups found for user with email: {email}\"\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n logger.error(msg)\n else:\n logger.debug(f\"Found groups {all_users_groups} for user with email {email}\")\n\n if not group_member_emails:\n msg = \"No groups found for any users.\"\n emit_background_error(msg, cc_pair_id=cc_pair_id)\n logger.error(msg)\n\n return group_member_emails\n\n\ndef _build_group_member_email_map_from_onyx_users(\n confluence_client: OnyxConfluence,\n) -> dict[str, set[str]]:\n \"\"\"Hacky, but it's the only way to do this as long as the\n Confluence APIs are broken.\n\n This is fixed in Confluence Data Center 10.1.0, so first choice\n is to tell users to upgrade to 10.1.0.\n https://jira.atlassian.com/browse/CONFSERVER-95999\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n # don't include external since they are handled by the \"through confluence\"\n # user fetching mechanism\n user_emails = [\n user.email for user in get_all_users(db_session, include_external=False)\n ]\n\n def _infer_username_from_email(email: str) -> str:\n return email.split(\"@\")[0]\n\n group_member_emails: dict[str, set[str]] = {}\n for email in user_emails:\n logger.info(f\"Processing groups for user with email: {email}\")\n try:\n user_name = _infer_username_from_email(email)\n response = confluence_client.get_user_details_by_username(user_name)\n user_key = response.get(\"userKey\")\n if not user_key:\n logger.error(f\"User key not found for user with email {email}\")\n continue\n\n all_users_groups: set[str] = set()\n for group in confluence_client.paginated_groups_by_user_retrieval(user_key):\n # group name uniqueness is enforced by Confluence, so we can use it as a group ID\n group_id = group[\"name\"]\n group_member_emails.setdefault(group_id, set()).add(email)\n all_users_groups.add(group_id)\n\n if not all_users_groups:\n msg = f\"No groups found for user with email: {email}\"\n logger.error(msg)\n else:\n logger.info(\n f\"Found groups {all_users_groups} for user with email {email}\"\n )\n except Exception:\n logger.exception(f\"Error getting user details for user with email {email}\")\n\n return group_member_emails\n\n\ndef _build_final_group_to_member_email_map(\n confluence_client: OnyxConfluence,\n cc_pair_id: int,\n # if set, will infer confluence usernames from onyx users in addition to using the\n # confluence users API. This is a hacky workaround for the fact that the Confluence\n # users API is broken before Confluence Data Center 10.1.0.\n use_onyx_users: bool = CONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC,\n) -> dict[str, set[str]]:\n group_to_member_email_map = _build_group_member_email_map(\n confluence_client=confluence_client,\n cc_pair_id=cc_pair_id,\n )\n group_to_member_email_map_from_onyx_users = (\n (\n _build_group_member_email_map_from_onyx_users(\n confluence_client=confluence_client,\n )\n )\n if use_onyx_users\n else {}\n )\n\n all_group_ids = set(group_to_member_email_map.keys()) | set(\n group_to_member_email_map_from_onyx_users.keys()\n )\n final_group_to_member_email_map = {}\n for group_id in all_group_ids:\n group_member_emails = group_to_member_email_map.get(\n group_id, set()\n ) | group_to_member_email_map_from_onyx_users.get(group_id, set())\n final_group_to_member_email_map[group_id] = group_member_emails\n\n return final_group_to_member_email_map\n\n\ndef confluence_group_sync(\n tenant_id: str,\n cc_pair: ConnectorCredentialPair,\n) -> Generator[ExternalUserGroup, None, None]:\n provider = OnyxDBCredentialsProvider(tenant_id, \"confluence\", cc_pair.credential_id)\n is_cloud = cc_pair.connector.connector_specific_config.get(\"is_cloud\", False)\n wiki_base: str = cc_pair.connector.connector_specific_config[\"wiki_base\"]\n url = wiki_base.rstrip(\"/\")\n\n probe_kwargs = {\n \"max_backoff_retries\": 6,\n \"max_backoff_seconds\": 10,\n }\n\n final_kwargs = {\n \"max_backoff_retries\": 10,\n \"max_backoff_seconds\": 60,\n }\n\n confluence_client = OnyxConfluence(is_cloud, url, provider)\n confluence_client._probe_connection(**probe_kwargs)\n confluence_client._initialize_connection(**final_kwargs)\n\n group_to_member_email_map = _build_final_group_to_member_email_map(\n confluence_client, cc_pair.id\n )\n\n all_found_emails = set()\n for group_id, group_member_emails in group_to_member_email_map.items():\n yield (\n ExternalUserGroup(\n id=group_id,\n user_emails=list(group_member_emails),\n )\n )\n all_found_emails.update(group_member_emails)\n\n # This is so that when we find a public confleunce server page, we can\n # give access to all users only in if they have an email in Confluence\n if cc_pair.connector.connector_specific_config.get(\"is_cloud\", False):\n all_found_group = ExternalUserGroup(\n id=ALL_CONF_EMAILS_GROUP_NAME,\n user_emails=list(all_found_emails),\n )\n yield all_found_group\n" }, { "path": "backend/ee/onyx/external_permissions/confluence/page_access.py", "content": "from typing import Any\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.onyx_confluence import (\n get_user_email_from_username__server,\n)\nfrom onyx.connectors.confluence.onyx_confluence import OnyxConfluence\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _extract_read_access_restrictions(\n confluence_client: OnyxConfluence, restrictions: dict[str, Any]\n) -> tuple[set[str], set[str], bool]:\n \"\"\"\n Converts a page's restrictions dict into an ExternalAccess object.\n If there are no restrictions, then return None\n \"\"\"\n read_access = restrictions.get(\"read\", {})\n read_access_restrictions = read_access.get(\"restrictions\", {})\n\n # Extract the users with read access\n read_access_user = read_access_restrictions.get(\"user\", {})\n read_access_user_jsons = read_access_user.get(\"results\", [])\n # any items found means that there is a restriction\n found_any_restriction = bool(read_access_user_jsons)\n\n read_access_user_emails = []\n for user in read_access_user_jsons:\n # If the user has an email, then add it to the list\n if user.get(\"email\"):\n read_access_user_emails.append(user[\"email\"])\n # If the user has a username and not an email, then get the email from Confluence\n elif user.get(\"username\"):\n email = get_user_email_from_username__server(\n confluence_client=confluence_client, user_name=user[\"username\"]\n )\n if email:\n read_access_user_emails.append(email)\n else:\n logger.warning(\n f\"Email for user {user['username']} not found in Confluence\"\n )\n else:\n if user.get(\"email\") is not None:\n logger.warning(f\"Cant find email for user {user.get('displayName')}\")\n logger.warning(\n \"This user needs to make their email accessible in Confluence Settings\"\n )\n\n logger.warning(f\"no user email or username for {user}\")\n\n # Extract the groups with read access\n read_access_group = read_access_restrictions.get(\"group\", {})\n read_access_group_jsons = read_access_group.get(\"results\", [])\n # any items found means that there is a restriction\n found_any_restriction |= bool(read_access_group_jsons)\n read_access_group_names = [\n group[\"name\"] for group in read_access_group_jsons if group.get(\"name\")\n ]\n\n return (\n set(read_access_user_emails),\n set(read_access_group_names),\n found_any_restriction,\n )\n\n\ndef get_page_restrictions(\n confluence_client: OnyxConfluence,\n page_id: str,\n page_restrictions: dict[str, Any],\n ancestors: list[dict[str, Any]],\n add_prefix: bool = False,\n) -> ExternalAccess | None:\n \"\"\"\n This function gets the restrictions for a page. In Confluence, a child can have\n at MOST the same level accessibility as its immediate parent.\n\n If no restrictions are found anywhere, then return None, indicating that the page\n should inherit the space's restrictions.\n\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path).\n \"\"\"\n found_user_emails: set[str] = set()\n found_group_names: set[str] = set()\n\n # NOTE: need the found_any_restriction, since we can find restrictions\n # but not be able to extract any user emails or group names\n # in this case, we should just give no access\n found_user_emails, found_group_names, found_any_page_level_restriction = (\n _extract_read_access_restrictions(\n confluence_client=confluence_client,\n restrictions=page_restrictions,\n )\n )\n\n def _maybe_prefix_groups(group_names: set[str]) -> set[str]:\n if add_prefix:\n return {\n build_ext_group_name_for_onyx(g, DocumentSource.CONFLUENCE)\n for g in group_names\n }\n return group_names\n\n # if there are individual page-level restrictions, then this is the accurate\n # restriction for the page. You cannot both have page-level restrictions AND\n # inherit restrictions from the parent.\n if found_any_page_level_restriction:\n return ExternalAccess(\n external_user_emails=found_user_emails,\n external_user_group_ids=_maybe_prefix_groups(found_group_names),\n is_public=False,\n )\n\n # ancestors seem to be in order from root to immediate parent\n # https://community.atlassian.com/forums/Confluence-questions/Order-of-ancestors-in-REST-API-response-Confluence-Server-amp/qaq-p/2385981\n # we want the restrictions from the immediate parent to take precedence, so we should\n # reverse the list\n for ancestor in reversed(ancestors):\n (\n ancestor_user_emails,\n ancestor_group_names,\n found_any_restrictions_in_ancestor,\n ) = _extract_read_access_restrictions(\n confluence_client=confluence_client,\n restrictions=ancestor.get(\"restrictions\", {}),\n )\n if found_any_restrictions_in_ancestor:\n # if inheriting restrictions from the parent, then the first one we run into\n # should be applied (the reason why we'd traverse more than one ancestor is if\n # the ancestor also is in \"inherit\" mode.)\n logger.debug(\n f\"Found user restrictions {ancestor_user_emails} and group restrictions {ancestor_group_names}\"\n f\"for document {page_id} based on ancestor {ancestor}\"\n )\n return ExternalAccess(\n external_user_emails=ancestor_user_emails,\n external_user_group_ids=_maybe_prefix_groups(ancestor_group_names),\n is_public=False,\n )\n\n # we didn't find any restrictions, so the page inherits the space's restrictions\n return None\n" }, { "path": "backend/ee/onyx/external_permissions/confluence/space_access.py", "content": "from ee.onyx.configs.app_configs import CONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC\nfrom ee.onyx.external_permissions.confluence.constants import ALL_CONF_EMAILS_GROUP_NAME\nfrom ee.onyx.external_permissions.confluence.constants import REQUEST_PAGINATION_LIMIT\nfrom ee.onyx.external_permissions.confluence.constants import VIEWSPACE_PERMISSION_TYPE\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.onyx_confluence import (\n get_user_email_from_username__server,\n)\nfrom onyx.connectors.confluence.onyx_confluence import OnyxConfluence\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _get_server_space_permissions(\n confluence_client: OnyxConfluence, space_key: str\n) -> ExternalAccess:\n space_permissions = confluence_client.get_all_space_permissions_server(\n space_key=space_key\n )\n\n viewspace_permissions = []\n for permission_category in space_permissions:\n if permission_category.get(\"type\") == VIEWSPACE_PERMISSION_TYPE:\n viewspace_permissions.extend(\n permission_category.get(\"spacePermissions\", [])\n )\n\n is_public = False\n user_names = set()\n group_names = set()\n for permission in viewspace_permissions:\n if user_name := permission.get(\"userName\"):\n user_names.add(user_name)\n if group_name := permission.get(\"groupName\"):\n group_names.add(group_name)\n\n # It seems that if anonymous access is turned on for the site and space,\n # then the space is publicly accessible.\n # For confluence server, we make a group that contains all users\n # that exist in confluence and then just add that group to the space permissions\n # if anonymous access is turned on for the site and space or we set is_public = True\n # if they set the env variable CONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC to True so\n # that we can support confluence server deployments that want anonymous access\n # to be public (we cant test this because its paywalled)\n if user_name is None and group_name is None:\n # Defaults to False\n if CONFLUENCE_ANONYMOUS_ACCESS_IS_PUBLIC:\n is_public = True\n else:\n group_names.add(ALL_CONF_EMAILS_GROUP_NAME)\n\n user_emails = set()\n for user_name in user_names:\n user_email = get_user_email_from_username__server(confluence_client, user_name)\n if user_email:\n user_emails.add(user_email)\n else:\n logger.warning(f\"Email for user {user_name} not found in Confluence\")\n\n if not user_emails and not group_names:\n logger.warning(\n \"No user emails or group names found in Confluence space permissions\"\n f\"\\nSpace key: {space_key}\"\n f\"\\nSpace permissions: {space_permissions}\"\n )\n\n return ExternalAccess(\n external_user_emails=user_emails,\n external_user_group_ids=group_names,\n is_public=is_public,\n )\n\n\ndef _get_cloud_space_permissions(\n confluence_client: OnyxConfluence, space_key: str\n) -> ExternalAccess:\n space_permissions_result = confluence_client.get_space(\n space_key=space_key, expand=\"permissions\"\n )\n space_permissions = space_permissions_result.get(\"permissions\", [])\n\n user_emails = set()\n group_names = set()\n is_externally_public = False\n for permission in space_permissions:\n subs = permission.get(\"subjects\")\n if subs:\n # If there are subjects, then there are explicit users or groups with access\n if email := subs.get(\"user\", {}).get(\"results\", [{}])[0].get(\"email\"):\n user_emails.add(email)\n if group_name := subs.get(\"group\", {}).get(\"results\", [{}])[0].get(\"name\"):\n group_names.add(group_name)\n else:\n # If there are no subjects, then the permission is for everyone\n if permission.get(\"operation\", {}).get(\n \"operation\"\n ) == \"read\" and permission.get(\"anonymousAccess\", False):\n # If the permission specifies read access for anonymous users, then\n # the space is publicly accessible\n is_externally_public = True\n\n return ExternalAccess(\n external_user_emails=user_emails,\n external_user_group_ids=group_names,\n is_public=is_externally_public,\n )\n\n\ndef get_space_permission(\n confluence_client: OnyxConfluence,\n space_key: str,\n is_cloud: bool,\n add_prefix: bool = False,\n) -> ExternalAccess:\n if is_cloud:\n space_permissions = _get_cloud_space_permissions(confluence_client, space_key)\n else:\n space_permissions = _get_server_space_permissions(confluence_client, space_key)\n\n if (\n not space_permissions.is_public\n and not space_permissions.external_user_emails\n and not space_permissions.external_user_group_ids\n ):\n logger.warning(\n f\"No permissions found for space '{space_key}'. This is very unlikely \"\n \"to be correct and is more likely caused by an access token with \"\n \"insufficient permissions. Make sure that the access token has Admin \"\n f\"permissions for space '{space_key}'\"\n )\n\n # Prefix group IDs with source type if requested (for indexing path)\n if add_prefix and space_permissions.external_user_group_ids:\n prefixed_groups = {\n build_ext_group_name_for_onyx(g, DocumentSource.CONFLUENCE)\n for g in space_permissions.external_user_group_ids\n }\n return ExternalAccess(\n external_user_emails=space_permissions.external_user_emails,\n external_user_group_ids=prefixed_groups,\n is_public=space_permissions.is_public,\n )\n\n return space_permissions\n\n\ndef get_all_space_permissions(\n confluence_client: OnyxConfluence,\n is_cloud: bool,\n add_prefix: bool = False,\n) -> dict[str, ExternalAccess]:\n \"\"\"\n Get access permissions for all spaces in Confluence.\n\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path).\n \"\"\"\n logger.debug(\"Getting space permissions\")\n # Gets all the spaces in the Confluence instance\n all_space_keys = [\n key\n for space in confluence_client.retrieve_confluence_spaces(\n limit=REQUEST_PAGINATION_LIMIT,\n )\n if (key := space.get(\"key\"))\n ]\n\n # Gets the permissions for each space\n logger.debug(f\"Got {len(all_space_keys)} spaces from confluence\")\n space_permissions_by_space_key: dict[str, ExternalAccess] = {}\n for space_key in all_space_keys:\n space_permissions = get_space_permission(\n confluence_client, space_key, is_cloud, add_prefix\n )\n\n # Stores the permissions for each space\n space_permissions_by_space_key[space_key] = space_permissions\n\n return space_permissions_by_space_key\n" }, { "path": "backend/ee/onyx/external_permissions/github/doc_sync.py", "content": "import json\nfrom collections.abc import Generator\n\nfrom github import Github\nfrom github.Repository import Repository\n\nfrom ee.onyx.external_permissions.github.utils import fetch_repository_team_slugs\nfrom ee.onyx.external_permissions.github.utils import form_collaborators_group_id\nfrom ee.onyx.external_permissions.github.utils import form_organization_group_id\nfrom ee.onyx.external_permissions.github.utils import (\n form_outside_collaborators_group_id,\n)\nfrom ee.onyx.external_permissions.github.utils import get_external_access_permission\nfrom ee.onyx.external_permissions.github.utils import get_repository_visibility\nfrom ee.onyx.external_permissions.github.utils import GitHubVisibility\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.github.connector import DocMetadata\nfrom onyx.connectors.github.connector import GithubConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import SortOrder\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nGITHUB_DOC_SYNC_LABEL = \"github_doc_sync\"\n\n\ndef github_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction,\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction, # noqa: ARG001\n callback: IndexingHeartbeatInterface | None = None,\n) -> Generator[DocExternalAccess, None, None]:\n \"\"\"\n Sync GitHub documents with external access permissions.\n\n This function checks each repository for visibility/team changes and updates\n document permissions accordingly without using checkpoints.\n \"\"\"\n logger.info(f\"Starting GitHub document sync for CC pair ID: {cc_pair.id}\")\n\n # Initialize GitHub connector with credentials\n github_connector: GithubConnector = GithubConnector(\n **cc_pair.connector.connector_specific_config\n )\n\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n github_connector.load_credentials(credential_json)\n logger.info(\"GitHub connector credentials loaded successfully\")\n\n if not github_connector.github_client:\n logger.error(\"GitHub client initialization failed\")\n raise ValueError(\"github_client is required\")\n\n # Get all repositories from GitHub API\n logger.info(\"Fetching all repositories from GitHub API\")\n try:\n repos = github_connector.fetch_configured_repos()\n\n logger.info(f\"Found {len(repos)} repositories to check\")\n except Exception as e:\n logger.error(f\"Failed to fetch repositories: {e}\")\n raise\n\n repo_to_doc_list_map: dict[str, list[DocumentRow]] = {}\n # sort order is ascending because we want to get the oldest documents first\n existing_docs: list[DocumentRow] = fetch_all_existing_docs_fn(\n sort_order=SortOrder.ASC\n )\n logger.info(f\"Found {len(existing_docs)} documents to check\")\n for doc in existing_docs:\n try:\n doc_metadata = DocMetadata.model_validate_json(json.dumps(doc.doc_metadata))\n if doc_metadata.repo not in repo_to_doc_list_map:\n repo_to_doc_list_map[doc_metadata.repo] = []\n repo_to_doc_list_map[doc_metadata.repo].append(doc)\n except Exception as e:\n logger.error(f\"Failed to parse doc metadata: {e} for doc {doc.id}\")\n continue\n logger.info(f\"Found {len(repo_to_doc_list_map)} documents to check\")\n # Process each repository individually\n for repo in repos:\n try:\n logger.info(f\"Processing repository: {repo.id} (name: {repo.name})\")\n repo_doc_list: list[DocumentRow] = repo_to_doc_list_map.get(\n repo.full_name, []\n )\n if not repo_doc_list:\n logger.warning(\n f\"No documents found for repository {repo.id} ({repo.name})\"\n )\n continue\n\n current_external_group_ids = repo_doc_list[0].external_user_group_ids or []\n # Check if repository has any permission changes\n has_changes = _check_repository_for_changes(\n repo=repo,\n github_client=github_connector.github_client,\n current_external_group_ids=current_external_group_ids,\n )\n\n if has_changes:\n logger.info(\n f\"Repository {repo.id} ({repo.name}) has changes, updating documents\"\n )\n\n # Get new external access permissions for this repository\n new_external_access = get_external_access_permission(\n repo, github_connector.github_client\n )\n\n logger.info(\n f\"Found {len(repo_doc_list)} documents for repository {repo.full_name}\"\n )\n\n # Yield updated external access for each document\n for doc in repo_doc_list:\n if callback:\n callback.progress(GITHUB_DOC_SYNC_LABEL, 1)\n\n yield DocExternalAccess(\n doc_id=doc.id,\n external_access=new_external_access,\n )\n else:\n logger.info(\n f\"Repository {repo.id} ({repo.name}) has no changes, skipping\"\n )\n except Exception as e:\n logger.error(f\"Error processing repository {repo.id} ({repo.name}): {e}\")\n\n logger.info(f\"GitHub document sync completed for CC pair ID: {cc_pair.id}\")\n\n\ndef _check_repository_for_changes(\n repo: Repository,\n github_client: Github,\n current_external_group_ids: list[str],\n) -> bool:\n \"\"\"\n Check if repository has any permission changes (visibility or team updates).\n \"\"\"\n logger.info(f\"Checking repository {repo.id} ({repo.name}) for changes\")\n\n # Check for repository visibility changes using the sample document data\n if _is_repo_visibility_changed_from_groups(\n repo=repo,\n current_external_group_ids=current_external_group_ids,\n ):\n logger.info(f\"Repository {repo.id} ({repo.name}) has visibility changes\")\n return True\n\n # Check for team membership changes if repository is private\n if get_repository_visibility(\n repo\n ) == GitHubVisibility.PRIVATE and _teams_updated_from_groups(\n repo=repo,\n github_client=github_client,\n current_external_group_ids=current_external_group_ids,\n ):\n logger.info(f\"Repository {repo.id} ({repo.name}) has team changes\")\n return True\n\n logger.info(f\"Repository {repo.id} ({repo.name}) has no changes\")\n return False\n\n\ndef _is_repo_visibility_changed_from_groups(\n repo: Repository,\n current_external_group_ids: list[str],\n) -> bool:\n \"\"\"\n Check if repository visibility has changed by analyzing existing external group IDs.\n\n Args:\n repo: GitHub repository object\n current_external_group_ids: List of external group IDs from existing document\n\n Returns:\n True if visibility has changed\n \"\"\"\n current_repo_visibility = get_repository_visibility(repo)\n logger.info(f\"Current repository visibility: {current_repo_visibility.value}\")\n\n # Build expected group IDs for current visibility\n collaborators_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=form_collaborators_group_id(repo.id),\n )\n\n org_group_id = None\n if repo.organization:\n org_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=form_organization_group_id(repo.organization.id),\n )\n\n # Determine existing visibility from group IDs\n has_collaborators_group = collaborators_group_id in current_external_group_ids\n has_org_group = org_group_id and org_group_id in current_external_group_ids\n\n if has_collaborators_group:\n existing_repo_visibility = GitHubVisibility.PRIVATE\n elif has_org_group:\n existing_repo_visibility = GitHubVisibility.INTERNAL\n else:\n existing_repo_visibility = GitHubVisibility.PUBLIC\n\n logger.info(f\"Inferred existing visibility: {existing_repo_visibility.value}\")\n\n visibility_changed = existing_repo_visibility != current_repo_visibility\n if visibility_changed:\n logger.info(\n f\"Visibility changed for repo {repo.id} ({repo.name}): \"\n f\"{existing_repo_visibility.value} -> {current_repo_visibility.value}\"\n )\n\n return visibility_changed\n\n\ndef _teams_updated_from_groups(\n repo: Repository,\n github_client: Github,\n current_external_group_ids: list[str],\n) -> bool:\n \"\"\"\n Check if repository team memberships have changed using existing group IDs.\n \"\"\"\n # Fetch current team slugs for the repository\n current_teams = fetch_repository_team_slugs(repo=repo, github_client=github_client)\n logger.info(\n f\"Current teams for repository {repo.id} (name: {repo.name}): {current_teams}\"\n )\n\n # Build group IDs to exclude from team comparison (non-team groups)\n collaborators_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=form_collaborators_group_id(repo.id),\n )\n outside_collaborators_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=form_outside_collaborators_group_id(repo.id),\n )\n non_team_group_ids = {collaborators_group_id, outside_collaborators_group_id}\n\n # Extract existing team IDs from current external group IDs\n existing_team_ids = set()\n for group_id in current_external_group_ids:\n # Skip all non-team groups, keep only team groups\n if group_id not in non_team_group_ids:\n existing_team_ids.add(group_id)\n\n # Note: existing_team_ids from DB are already prefixed (e.g., \"github__team-slug\")\n # but current_teams from API are raw team slugs, so we need to add the prefix\n current_team_ids = set()\n for team_slug in current_teams:\n team_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=team_slug,\n )\n current_team_ids.add(team_group_id)\n\n logger.info(\n f\"Existing team IDs: {existing_team_ids}, Current team IDs: {current_team_ids}\"\n )\n\n # Compare actual team IDs to detect changes\n teams_changed = current_team_ids != existing_team_ids\n if teams_changed:\n logger.info(\n f\"Team changes detected for repo {repo.id} (name: {repo.name}): \"\n f\"existing={existing_team_ids}, current={current_team_ids}\"\n )\n\n return teams_changed\n" }, { "path": "backend/ee/onyx/external_permissions/github/group_sync.py", "content": "from collections.abc import Generator\n\nfrom github import Repository\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom ee.onyx.external_permissions.github.utils import get_external_user_group\nfrom onyx.connectors.github.connector import GithubConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef github_group_sync(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair,\n) -> Generator[ExternalUserGroup, None, None]:\n github_connector: GithubConnector = GithubConnector(\n **cc_pair.connector.connector_specific_config\n )\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n github_connector.load_credentials(credential_json)\n if not github_connector.github_client:\n raise ValueError(\"github_client is required\")\n\n logger.info(\"Starting GitHub group sync...\")\n repos: list[Repository.Repository] = []\n if github_connector.repositories:\n if \",\" in github_connector.repositories:\n # Multiple repositories specified\n repos = github_connector.get_github_repos(github_connector.github_client)\n else:\n # Single repository (backward compatibility)\n repos = [github_connector.get_github_repo(github_connector.github_client)]\n else:\n # All repositories\n repos = github_connector.get_all_repos(github_connector.github_client)\n\n for repo in repos:\n try:\n for external_group in get_external_user_group(\n repo, github_connector.github_client\n ):\n logger.info(f\"External group: {external_group}\")\n yield external_group\n except Exception as e:\n logger.error(f\"Error processing repository {repo.id} ({repo.name}): {e}\")\n" }, { "path": "backend/ee/onyx/external_permissions/github/utils.py", "content": "from collections.abc import Callable\nfrom enum import Enum\nfrom typing import List\nfrom typing import Optional\nfrom typing import Tuple\nfrom typing import TypeVar\n\nfrom github import Github\nfrom github import RateLimitExceededException\nfrom github.GithubException import GithubException\nfrom github.NamedUser import NamedUser\nfrom github.Organization import Organization\nfrom github.PaginatedList import PaginatedList\nfrom github.Repository import Repository\nfrom github.Team import Team\nfrom pydantic import BaseModel\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.github.rate_limit_utils import sleep_after_rate_limit_exception\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass GitHubVisibility(Enum):\n \"\"\"GitHub repository visibility options.\"\"\"\n\n PUBLIC = \"public\"\n PRIVATE = \"private\"\n INTERNAL = \"internal\"\n\n\nMAX_RETRY_COUNT = 3\n\nT = TypeVar(\"T\")\n\n# Higher-order function to wrap GitHub operations with retry and exception handling\n\n\ndef _run_with_retry(\n operation: Callable[[], T],\n description: str,\n github_client: Github,\n retry_count: int = 0,\n) -> Optional[T]:\n \"\"\"Execute a GitHub operation with retry on rate limit and exception handling.\"\"\"\n logger.debug(f\"Starting operation '{description}', attempt {retry_count + 1}\")\n try:\n result = operation()\n logger.debug(f\"Operation '{description}' completed successfully\")\n return result\n except RateLimitExceededException:\n if retry_count < MAX_RETRY_COUNT:\n sleep_after_rate_limit_exception(github_client)\n logger.warning(\n f\"Rate limit exceeded while {description}. Retrying... (attempt {retry_count + 1}/{MAX_RETRY_COUNT})\"\n )\n return _run_with_retry(\n operation, description, github_client, retry_count + 1\n )\n else:\n error_msg = f\"Max retries exceeded for {description}\"\n logger.exception(error_msg)\n raise RuntimeError(error_msg)\n except GithubException as e:\n logger.warning(f\"GitHub API error during {description}: {e}\")\n return None\n except Exception as e:\n logger.exception(f\"Unexpected error during {description}: {e}\")\n return None\n\n\nclass UserInfo(BaseModel):\n \"\"\"Represents a GitHub user with their basic information.\"\"\"\n\n login: str\n name: Optional[str] = None\n email: Optional[str] = None\n\n\nclass TeamInfo(BaseModel):\n \"\"\"Represents a GitHub team with its members.\"\"\"\n\n name: str\n slug: str\n members: List[UserInfo]\n\n\ndef _fetch_organization_members(\n github_client: Github,\n org_name: str,\n retry_count: int = 0, # noqa: ARG001\n) -> List[UserInfo]:\n \"\"\"Fetch all organization members including owners and regular members.\"\"\"\n org_members: List[UserInfo] = []\n logger.info(f\"Fetching organization members for {org_name}\")\n\n org = _run_with_retry(\n lambda: github_client.get_organization(org_name),\n f\"get organization {org_name}\",\n github_client,\n )\n if not org:\n logger.error(f\"Failed to fetch organization {org_name}\")\n raise RuntimeError(f\"Failed to fetch organization {org_name}\")\n\n member_objs: PaginatedList[NamedUser] | list[NamedUser] = (\n _run_with_retry(\n lambda: org.get_members(filter_=\"all\"),\n f\"get members for organization {org_name}\",\n github_client,\n )\n or []\n )\n\n for member in member_objs:\n user_info = UserInfo(login=member.login, name=member.name, email=member.email)\n org_members.append(user_info)\n\n logger.info(f\"Fetched {len(org_members)} members for organization {org_name}\")\n return org_members\n\n\ndef _fetch_repository_teams_detailed(\n repo: Repository,\n github_client: Github,\n retry_count: int = 0, # noqa: ARG001\n) -> List[TeamInfo]:\n \"\"\"Fetch teams with access to the repository and their members.\"\"\"\n teams_data: List[TeamInfo] = []\n logger.info(f\"Fetching teams for repository {repo.full_name}\")\n\n team_objs: PaginatedList[Team] | list[Team] = (\n _run_with_retry(\n lambda: repo.get_teams(),\n f\"get teams for repository {repo.full_name}\",\n github_client,\n )\n or []\n )\n\n for team in team_objs:\n logger.info(\n f\"Processing team {team.name} (slug: {team.slug}) for repository {repo.full_name}\"\n )\n\n members: PaginatedList[NamedUser] | list[NamedUser] = (\n _run_with_retry(\n lambda: team.get_members(),\n f\"get members for team {team.name}\",\n github_client,\n )\n or []\n )\n\n team_members = []\n for m in members:\n user_info = UserInfo(login=m.login, name=m.name, email=m.email)\n team_members.append(user_info)\n\n team_info = TeamInfo(name=team.name, slug=team.slug, members=team_members)\n teams_data.append(team_info)\n logger.info(f\"Team {team.name} has {len(team_members)} members\")\n\n logger.info(f\"Fetched {len(teams_data)} teams for repository {repo.full_name}\")\n return teams_data\n\n\ndef fetch_repository_team_slugs(\n repo: Repository,\n github_client: Github,\n retry_count: int = 0, # noqa: ARG001\n) -> List[str]:\n \"\"\"Fetch team slugs with access to the repository.\"\"\"\n logger.info(f\"Fetching team slugs for repository {repo.full_name}\")\n teams_data: List[str] = []\n\n team_objs: PaginatedList[Team] | list[Team] = (\n _run_with_retry(\n lambda: repo.get_teams(),\n f\"get teams for repository {repo.full_name}\",\n github_client,\n )\n or []\n )\n\n for team in team_objs:\n teams_data.append(team.slug)\n\n logger.info(f\"Fetched {len(teams_data)} team slugs for repository {repo.full_name}\")\n return teams_data\n\n\ndef _get_collaborators_and_outside_collaborators(\n github_client: Github,\n repo: Repository,\n) -> Tuple[List[UserInfo], List[UserInfo]]:\n \"\"\"Fetch and categorize collaborators into regular and outside collaborators.\"\"\"\n collaborators: List[UserInfo] = []\n outside_collaborators: List[UserInfo] = []\n logger.info(f\"Fetching collaborators for repository {repo.full_name}\")\n\n repo_collaborators: PaginatedList[NamedUser] | list[NamedUser] = (\n _run_with_retry(\n lambda: repo.get_collaborators(),\n f\"get collaborators for repository {repo.full_name}\",\n github_client,\n )\n or []\n )\n\n for collaborator in repo_collaborators:\n is_outside = False\n\n # Check if collaborator is outside the organization\n if repo.organization:\n org: Organization | None = _run_with_retry(\n lambda: github_client.get_organization(repo.organization.login),\n f\"get organization {repo.organization.login}\",\n github_client,\n )\n\n if org is not None:\n org_obj = org\n membership = _run_with_retry(\n lambda: org_obj.has_in_members(collaborator),\n f\"check membership for {collaborator.login} in org {org_obj.login}\",\n github_client,\n )\n is_outside = membership is not None and not membership\n\n info = UserInfo(\n login=collaborator.login, name=collaborator.name, email=collaborator.email\n )\n if repo.organization and is_outside:\n outside_collaborators.append(info)\n else:\n collaborators.append(info)\n\n logger.info(\n f\"Categorized {len(collaborators)} regular and {len(outside_collaborators)} outside collaborators for {repo.full_name}\"\n )\n return collaborators, outside_collaborators\n\n\ndef form_collaborators_group_id(repository_id: int) -> str:\n \"\"\"Generate group ID for repository collaborators.\"\"\"\n if not repository_id:\n logger.exception(\"Repository ID is required to generate collaborators group ID\")\n raise ValueError(\"Repository ID must be set to generate group ID.\")\n group_id = f\"{repository_id}_collaborators\"\n return group_id\n\n\ndef form_organization_group_id(organization_id: int) -> str:\n \"\"\"Generate group ID for organization using organization ID.\"\"\"\n if not organization_id:\n logger.exception(\n \"Organization ID is required to generate organization group ID\"\n )\n raise ValueError(\"Organization ID must be set to generate group ID.\")\n group_id = f\"{organization_id}_organization\"\n return group_id\n\n\ndef form_outside_collaborators_group_id(repository_id: int) -> str:\n \"\"\"Generate group ID for outside collaborators.\"\"\"\n if not repository_id:\n logger.exception(\n \"Repository ID is required to generate outside collaborators group ID\"\n )\n raise ValueError(\"Repository ID must be set to generate group ID.\")\n group_id = f\"{repository_id}_outside_collaborators\"\n return group_id\n\n\ndef get_repository_visibility(repo: Repository) -> GitHubVisibility:\n \"\"\"\n Get the visibility of a repository.\n Returns GitHubVisibility enum member.\n \"\"\"\n if hasattr(repo, \"visibility\"):\n visibility = repo.visibility\n logger.info(\n f\"Repository {repo.full_name} visibility from attribute: {visibility}\"\n )\n try:\n return GitHubVisibility(visibility)\n except ValueError:\n logger.warning(\n f\"Unknown visibility '{visibility}' for repo {repo.full_name}, defaulting to private\"\n )\n return GitHubVisibility.PRIVATE\n\n logger.info(f\"Repository {repo.full_name} is private\")\n return GitHubVisibility.PRIVATE\n\n\ndef get_external_access_permission(\n repo: Repository, github_client: Github, add_prefix: bool = False\n) -> ExternalAccess:\n \"\"\"\n Get the external access permission for a repository.\n Uses group-based permissions for efficiency and scalability.\n\n add_prefix: When this method is called during the initial permission sync via the connector,\n the group ID isn't prefixed with the source while inserting the document record.\n So in that case, set add_prefix to True, allowing the method itself to handle\n prefixing. However, when the same method is invoked from doc_sync, our system\n already adds the prefix to the group ID while processing the ExternalAccess object.\n \"\"\"\n # We maintain collaborators, and outside collaborators as two separate groups\n # instead of adding individual user emails to ExternalAccess.external_user_emails for two reasons:\n # 1. Changes in repo collaborators (additions/removals) would require updating all documents.\n # 2. Repo permissions can change without updating the repo's updated_at timestamp,\n # forcing full permission syncs for all documents every time, which is inefficient.\n\n repo_visibility = get_repository_visibility(repo)\n logger.info(\n f\"Generating ExternalAccess for {repo.full_name}: visibility={repo_visibility.value}\"\n )\n\n if repo_visibility == GitHubVisibility.PUBLIC:\n logger.info(\n f\"Repository {repo.full_name} is public - allowing access to all users\"\n )\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n )\n elif repo_visibility == GitHubVisibility.PRIVATE:\n logger.info(\n f\"Repository {repo.full_name} is private - setting up restricted access\"\n )\n\n collaborators_group_id = form_collaborators_group_id(repo.id)\n outside_collaborators_group_id = form_outside_collaborators_group_id(repo.id)\n if add_prefix:\n collaborators_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=collaborators_group_id,\n )\n outside_collaborators_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=outside_collaborators_group_id,\n )\n group_ids = {collaborators_group_id, outside_collaborators_group_id}\n\n team_slugs = fetch_repository_team_slugs(repo, github_client)\n if add_prefix:\n team_slugs = [\n build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=slug,\n )\n for slug in team_slugs\n ]\n group_ids.update(team_slugs)\n\n logger.info(f\"ExternalAccess groups for {repo.full_name}: {group_ids}\")\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=group_ids,\n is_public=False,\n )\n else:\n # Internal repositories - accessible to organization members\n logger.info(\n f\"Repository {repo.full_name} is internal - accessible to org members\"\n )\n org_group_id = form_organization_group_id(repo.organization.id)\n if add_prefix:\n org_group_id = build_ext_group_name_for_onyx(\n source=DocumentSource.GITHUB,\n ext_group_name=org_group_id,\n )\n group_ids = {org_group_id}\n logger.info(f\"ExternalAccess groups for {repo.full_name}: {group_ids}\")\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=group_ids,\n is_public=False,\n )\n\n\ndef get_external_user_group(\n repo: Repository, github_client: Github\n) -> list[ExternalUserGroup]:\n \"\"\"\n Get the external user group for a repository.\n Creates ExternalUserGroup objects with actual user emails for each permission group.\n \"\"\"\n repo_visibility = get_repository_visibility(repo)\n logger.info(\n f\"Generating ExternalUserGroups for {repo.full_name}: visibility={repo_visibility.value}\"\n )\n\n if repo_visibility == GitHubVisibility.PRIVATE:\n logger.info(f\"Processing private repository {repo.full_name}\")\n\n collaborators, outside_collaborators = (\n _get_collaborators_and_outside_collaborators(github_client, repo)\n )\n teams = _fetch_repository_teams_detailed(repo, github_client)\n external_user_groups = []\n\n user_emails = set()\n for collab in collaborators:\n if collab.email:\n user_emails.add(collab.email)\n else:\n logger.error(f\"Collaborator {collab.login} has no email\")\n\n if user_emails:\n collaborators_group = ExternalUserGroup(\n id=form_collaborators_group_id(repo.id),\n user_emails=list(user_emails),\n )\n external_user_groups.append(collaborators_group)\n logger.info(f\"Created collaborators group with {len(user_emails)} emails\")\n\n # Create group for outside collaborators\n user_emails = set()\n for collab in outside_collaborators:\n if collab.email:\n user_emails.add(collab.email)\n else:\n logger.error(f\"Outside collaborator {collab.login} has no email\")\n\n if user_emails:\n outside_collaborators_group = ExternalUserGroup(\n id=form_outside_collaborators_group_id(repo.id),\n user_emails=list(user_emails),\n )\n external_user_groups.append(outside_collaborators_group)\n logger.info(\n f\"Created outside collaborators group with {len(user_emails)} emails\"\n )\n\n # Create groups for teams\n for team in teams:\n user_emails = set()\n for member in team.members:\n if member.email:\n user_emails.add(member.email)\n else:\n logger.error(f\"Team member {member.login} has no email\")\n\n if user_emails:\n team_group = ExternalUserGroup(\n id=team.slug,\n user_emails=list(user_emails),\n )\n external_user_groups.append(team_group)\n logger.info(\n f\"Created team group {team.name} with {len(user_emails)} emails\"\n )\n\n logger.info(\n f\"Created {len(external_user_groups)} ExternalUserGroups for private repository {repo.full_name}\"\n )\n return external_user_groups\n\n if repo_visibility == GitHubVisibility.INTERNAL:\n logger.info(f\"Processing internal repository {repo.full_name}\")\n\n org_group_id = form_organization_group_id(repo.organization.id)\n org_members = _fetch_organization_members(\n github_client, repo.organization.login\n )\n\n user_emails = set()\n for member in org_members:\n if member.email:\n user_emails.add(member.email)\n else:\n logger.error(f\"Org member {member.login} has no email\")\n\n org_group = ExternalUserGroup(\n id=org_group_id,\n user_emails=list(user_emails),\n )\n logger.info(\n f\"Created organization group with {len(user_emails)} emails for internal repository {repo.full_name}\"\n )\n return [org_group]\n\n logger.info(f\"Repository {repo.full_name} is public - no user groups needed\")\n return []\n" }, { "path": "backend/ee/onyx/external_permissions/gmail/doc_sync.py", "content": "from collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.access.models import NodeExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.gmail.connector import GmailConnector\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _get_slim_doc_generator(\n cc_pair: ConnectorCredentialPair,\n gmail_connector: GmailConnector,\n callback: IndexingHeartbeatInterface | None = None,\n) -> GenerateSlimDocumentOutput:\n current_time = datetime.now(timezone.utc)\n start_time = (\n cc_pair.last_time_perm_sync.replace(tzinfo=timezone.utc).timestamp()\n if cc_pair.last_time_perm_sync\n else 0.0\n )\n\n return gmail_connector.retrieve_all_slim_docs_perm_sync(\n start=start_time,\n end=current_time.timestamp(),\n callback=callback,\n )\n\n\ndef gmail_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction, # noqa: ARG001\n callback: IndexingHeartbeatInterface | None,\n) -> Generator[ElementExternalAccess, None, None]:\n \"\"\"\n Adds the external permissions to the documents and hierarchy nodes in postgres.\n If the document doesn't already exist in postgres, we create\n it in postgres so that when it gets created later, the permissions are\n already populated.\n \"\"\"\n gmail_connector = GmailConnector(**cc_pair.connector.connector_specific_config)\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n gmail_connector.load_credentials(credential_json)\n\n slim_doc_generator = _get_slim_doc_generator(\n cc_pair, gmail_connector, callback=callback\n )\n\n for slim_doc_batch in slim_doc_generator:\n for slim_doc in slim_doc_batch:\n if callback:\n if callback.should_stop():\n raise RuntimeError(\"gmail_doc_sync: Stop signal detected\")\n\n callback.progress(\"gmail_doc_sync\", 1)\n\n if isinstance(slim_doc, HierarchyNode):\n # Yield hierarchy node permissions to be processed in outer layer\n if slim_doc.external_access:\n yield NodeExternalAccess(\n external_access=slim_doc.external_access,\n raw_node_id=slim_doc.raw_node_id,\n source=DocumentSource.GMAIL.value,\n )\n continue\n if slim_doc.external_access is None:\n logger.warning(f\"No permissions found for document {slim_doc.id}\")\n continue\n\n yield DocExternalAccess(\n doc_id=slim_doc.id,\n external_access=slim_doc.external_access,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/google_drive/__init__.py", "content": "" }, { "path": "backend/ee/onyx/external_permissions/google_drive/doc_sync.py", "content": "from collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission\nfrom ee.onyx.external_permissions.google_drive.models import PermissionType\nfrom ee.onyx.external_permissions.google_drive.permission_retrieval import (\n get_permissions_by_ids,\n)\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.models import NodeExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.google_drive.models import GoogleDriveFileType\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _get_slim_doc_generator(\n cc_pair: ConnectorCredentialPair,\n google_drive_connector: GoogleDriveConnector,\n callback: IndexingHeartbeatInterface | None = None,\n) -> GenerateSlimDocumentOutput:\n current_time = datetime.now(timezone.utc)\n start_time = (\n cc_pair.last_time_perm_sync.replace(tzinfo=timezone.utc).timestamp()\n if cc_pair.last_time_perm_sync\n else 0.0\n )\n\n return google_drive_connector.retrieve_all_slim_docs_perm_sync(\n start=start_time,\n end=current_time.timestamp(),\n callback=callback,\n )\n\n\ndef _merge_permissions_lists(\n permission_lists: list[list[GoogleDrivePermission]],\n) -> list[GoogleDrivePermission]:\n \"\"\"\n Merge a list of permission lists into a single list of permissions.\n \"\"\"\n seen_permission_ids: set[str] = set()\n merged_permissions: list[GoogleDrivePermission] = []\n for permission_list in permission_lists:\n for permission in permission_list:\n if permission.id not in seen_permission_ids:\n merged_permissions.append(permission)\n seen_permission_ids.add(permission.id)\n\n return merged_permissions\n\n\ndef get_external_access_for_raw_gdrive_file(\n file: GoogleDriveFileType,\n company_domain: str,\n retriever_drive_service: GoogleDriveService | None,\n admin_drive_service: GoogleDriveService,\n fallback_user_email: str,\n add_prefix: bool = False,\n) -> ExternalAccess:\n \"\"\"\n Get the external access for a raw Google Drive file.\n\n Assumes the file we retrieved has EITHER `permissions` or `permission_ids`\n\n add_prefix: When this method is called during the initial indexing via the connector,\n set add_prefix to True so group IDs are prefixed with the source type.\n When invoked from doc_sync (permission sync), use the default (False)\n since upsert_document_external_perms handles prefixing.\n fallback_user_email: When we cannot retrieve any permission info for a file\n (e.g. externally-owned files where the API returns no permissions\n and permissions.list returns 403), fall back to granting access\n to this user. This is typically the impersonated org user whose\n drive contained the file.\n \"\"\"\n doc_id = file.get(\"id\")\n if not doc_id:\n raise ValueError(\"No doc_id found in file\")\n\n permissions = file.get(\"permissions\")\n permission_ids = file.get(\"permissionIds\")\n drive_id = file.get(\"driveId\")\n\n permissions_list: list[GoogleDrivePermission] = []\n if permissions:\n permissions_list = [\n GoogleDrivePermission.from_drive_permission(p) for p in permissions\n ]\n elif permission_ids:\n\n def _get_permissions(\n drive_service: GoogleDriveService,\n ) -> list[GoogleDrivePermission]:\n return get_permissions_by_ids(\n drive_service=drive_service,\n doc_id=doc_id,\n permission_ids=permission_ids,\n )\n\n permissions_list = _get_permissions(\n retriever_drive_service or admin_drive_service\n )\n if len(permissions_list) != len(permission_ids) and retriever_drive_service:\n logger.warning(\n f\"Failed to get all permissions for file {doc_id} with retriever service, trying admin service\"\n )\n backup_permissions_list = _get_permissions(admin_drive_service)\n permissions_list = _merge_permissions_lists(\n [permissions_list, backup_permissions_list]\n )\n\n # For externally-owned files, the Drive API may return no permissions\n # and permissions.list may return 403. In this case, fall back to\n # granting access to the user who found the file in their drive.\n # Note, even if other users also have access to this file,\n # they will not be granted access in Onyx.\n # We check permissions_list (the final result after all fetch attempts)\n # rather than the raw fields, because permission_ids may be present\n # but the actual fetch can still return empty due to a 403.\n if not permissions_list:\n logger.info(\n f\"No permission info available for file {doc_id} \"\n f\"(likely owned by a user outside of your organization). \"\n f\"Falling back to granting access to retriever user: {fallback_user_email}\"\n )\n return ExternalAccess(\n external_user_emails={fallback_user_email},\n external_user_group_ids=set(),\n is_public=False,\n )\n\n folder_ids_to_inherit_permissions_from: set[str] = set()\n user_emails: set[str] = set()\n group_emails: set[str] = set()\n public = False\n\n for permission in permissions_list:\n # if the permission is inherited, do not add it directly to the file\n # instead, add the folder ID as a group that has access to the file\n # we will then handle mapping that folder to the list of Onyx users\n # in the group sync job\n # NOTE: this doesn't handle the case where a folder initially has no\n # permissioning, but then later that folder is shared with a user or group.\n # We could fetch all ancestors of the file to get the list of folders that\n # might affect the permissions of the file, but this will get replaced with\n # an audit-log based approach in the future so not doing it now.\n if permission.inherited_from:\n folder_ids_to_inherit_permissions_from.add(permission.inherited_from)\n\n if permission.type == PermissionType.USER:\n if permission.email_address:\n user_emails.add(permission.email_address)\n else:\n logger.error(\n f\"Permission is type `user` but no email address is provided for document {doc_id}\\n {permission}\"\n )\n elif permission.type == PermissionType.GROUP:\n # groups are represented as email addresses within Drive\n if permission.email_address:\n group_emails.add(permission.email_address)\n else:\n logger.error(\n f\"Permission is type `group` but no email address is provided for document {doc_id}\\n {permission}\"\n )\n elif permission.type == PermissionType.DOMAIN and company_domain:\n if permission.domain == company_domain:\n public = True\n else:\n logger.warning(\n f\"Permission is type domain but does not match company domain:\\n {permission}\"\n )\n elif permission.type == PermissionType.ANYONE:\n public = True\n\n group_ids = (\n group_emails\n | folder_ids_to_inherit_permissions_from\n | ({drive_id} if drive_id is not None else set())\n )\n\n # Prefix group IDs with source type if requested (for indexing path)\n if add_prefix:\n group_ids = {\n build_ext_group_name_for_onyx(group_id, DocumentSource.GOOGLE_DRIVE)\n for group_id in group_ids\n }\n\n return ExternalAccess(\n external_user_emails=user_emails,\n external_user_group_ids=group_ids,\n is_public=public,\n )\n\n\ndef get_external_access_for_folder(\n folder: GoogleDriveFileType,\n google_domain: str,\n drive_service: GoogleDriveService,\n add_prefix: bool = False,\n) -> ExternalAccess:\n \"\"\"\n Extract ExternalAccess from a folder's permissions.\n\n This fetches permissions using the Drive API (via permissionIds) and extracts\n user emails, group emails, and public access status.\n\n Args:\n folder: The folder metadata from Google Drive API (must include permissionIds field)\n google_domain: The company's Google Workspace domain (e.g., \"company.com\")\n drive_service: Google Drive service for fetching permission details\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path).\n\n Returns:\n ExternalAccess with extracted permission info\n \"\"\"\n folder_id = folder.get(\"id\")\n if not folder_id:\n logger.warning(\"Folder missing ID, returning empty permissions\")\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=False,\n )\n\n # Get permission IDs from folder metadata\n permission_ids = folder.get(\"permissionIds\") or []\n if not permission_ids:\n logger.debug(f\"No permissionIds found for folder {folder_id}\")\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=False,\n )\n\n # Fetch full permission objects using the permission IDs\n permissions_list = get_permissions_by_ids(\n drive_service=drive_service,\n doc_id=folder_id,\n permission_ids=permission_ids,\n )\n\n user_emails: set[str] = set()\n group_emails: set[str] = set()\n is_public = False\n\n for permission in permissions_list:\n if permission.type == PermissionType.USER:\n if permission.email_address:\n user_emails.add(permission.email_address)\n else:\n logger.warning(f\"User permission without email for folder {folder_id}\")\n elif permission.type == PermissionType.GROUP:\n # Groups are represented as email addresses in Google Drive\n if permission.email_address:\n group_emails.add(permission.email_address)\n else:\n logger.warning(f\"Group permission without email for folder {folder_id}\")\n elif permission.type == PermissionType.DOMAIN:\n # Domain permission - check if it matches company domain\n if permission.domain == google_domain:\n # Only public if discoverable (allowFileDiscovery is not False)\n # If allowFileDiscovery is False, it's \"link only\" access\n is_public = permission.allow_file_discovery is not False\n else:\n logger.debug(\n f\"Domain permission for {permission.domain} does not match \"\n f\"company domain {google_domain} for folder {folder_id}\"\n )\n elif permission.type == PermissionType.ANYONE:\n # Only public if discoverable (allowFileDiscovery is not False)\n # If allowFileDiscovery is False, it's \"link only\" access\n is_public = permission.allow_file_discovery is not False\n\n # Prefix group IDs with source type if requested (for indexing path)\n group_ids: set[str] = group_emails\n if add_prefix:\n group_ids = {\n build_ext_group_name_for_onyx(group_id, DocumentSource.GOOGLE_DRIVE)\n for group_id in group_emails\n }\n\n return ExternalAccess(\n external_user_emails=user_emails,\n external_user_group_ids=group_ids,\n is_public=is_public,\n )\n\n\ndef gdrive_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction, # noqa: ARG001\n callback: IndexingHeartbeatInterface | None,\n) -> Generator[ElementExternalAccess, None, None]:\n \"\"\"\n Adds the external permissions to the documents and hierarchy nodes in postgres.\n If the document doesn't already exist in postgres, we create\n it in postgres so that when it gets created later, the permissions are\n already populated.\n \"\"\"\n google_drive_connector = GoogleDriveConnector(\n **cc_pair.connector.connector_specific_config\n )\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n google_drive_connector.load_credentials(credential_json)\n\n slim_doc_generator = _get_slim_doc_generator(cc_pair, google_drive_connector)\n\n total_processed = 0\n for slim_doc_batch in slim_doc_generator:\n logger.info(f\"Drive perm sync: Processing {len(slim_doc_batch)} documents\")\n for slim_doc in slim_doc_batch:\n if callback:\n if callback.should_stop():\n raise RuntimeError(\"gdrive_doc_sync: Stop signal detected\")\n\n callback.progress(\"gdrive_doc_sync\", 1)\n if isinstance(slim_doc, HierarchyNode):\n # Yield hierarchy node permissions to be processed in outer layer\n if slim_doc.external_access:\n yield NodeExternalAccess(\n external_access=slim_doc.external_access,\n raw_node_id=slim_doc.raw_node_id,\n source=DocumentSource.GOOGLE_DRIVE.value,\n )\n continue\n if slim_doc.external_access is None:\n raise ValueError(\n f\"Drive perm sync: No external access for document {slim_doc.id}\"\n )\n\n yield DocExternalAccess(\n external_access=slim_doc.external_access,\n doc_id=slim_doc.id,\n )\n total_processed += len(slim_doc_batch)\n logger.info(f\"Drive perm sync: Processed {total_processed} total documents\")\n" }, { "path": "backend/ee/onyx/external_permissions/google_drive/folder_retrieval.py", "content": "from collections.abc import Iterator\n\nfrom googleapiclient.discovery import Resource # type: ignore\n\nfrom ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission\nfrom ee.onyx.external_permissions.google_drive.permission_retrieval import (\n get_permissions_by_ids,\n)\nfrom onyx.connectors.google_drive.constants import DRIVE_FOLDER_TYPE\nfrom onyx.connectors.google_drive.file_retrieval import generate_time_range_filter\nfrom onyx.connectors.google_drive.models import GoogleDriveFileType\nfrom onyx.connectors.google_utils.google_utils import execute_paginated_retrieval\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Only include fields we need - folder ID and permissions\n# IMPORTANT: must fetch permissionIds, since sometimes the drive API\n# seems to miss permissions when requesting them directly\nFOLDER_PERMISSION_FIELDS = \"nextPageToken, files(id, name, permissionIds, permissions(id, emailAddress, type, domain, permissionDetails))\"\n\n\ndef get_folder_permissions_by_ids(\n service: Resource,\n folder_id: str,\n permission_ids: list[str],\n) -> list[GoogleDrivePermission]:\n \"\"\"\n Retrieves permissions for a specific folder filtered by permission IDs.\n\n Args:\n service: The Google Drive service instance\n folder_id: The ID of the folder to fetch permissions for\n permission_ids: A list of permission IDs to filter by\n\n Returns:\n A list of permissions matching the provided permission IDs\n \"\"\"\n return get_permissions_by_ids(\n drive_service=service,\n doc_id=folder_id,\n permission_ids=permission_ids,\n )\n\n\ndef get_modified_folders(\n service: Resource,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n) -> Iterator[GoogleDriveFileType]:\n \"\"\"\n Retrieves all folders that were modified within the specified time range.\n Only includes folder ID and permission information, not any contained files.\n\n Args:\n service: The Google Drive service instance\n start: The start time as seconds since Unix epoch (inclusive)\n end: The end time as seconds since Unix epoch (inclusive)\n\n Returns:\n An iterator yielding folder information including ID and permissions\n \"\"\"\n # Build query for folders\n query = f\"mimeType = '{DRIVE_FOLDER_TYPE}'\"\n query += \" and trashed = false\"\n query += generate_time_range_filter(start, end)\n\n # Retrieve and yield folders\n for folder in execute_paginated_retrieval(\n retrieval_function=service.files().list,\n list_key=\"files\",\n continue_on_404_or_403=True,\n corpora=\"allDrives\",\n supportsAllDrives=True,\n includeItemsFromAllDrives=True,\n includePermissionsForView=\"published\",\n fields=FOLDER_PERMISSION_FIELDS,\n q=query,\n ):\n yield folder\n" }, { "path": "backend/ee/onyx/external_permissions/google_drive/group_sync.py", "content": "from collections.abc import Generator\n\nfrom googleapiclient.errors import HttpError # type: ignore\nfrom pydantic import BaseModel\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom ee.onyx.external_permissions.google_drive.folder_retrieval import (\n get_folder_permissions_by_ids,\n)\nfrom ee.onyx.external_permissions.google_drive.folder_retrieval import (\n get_modified_folders,\n)\nfrom ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission\nfrom ee.onyx.external_permissions.google_drive.models import PermissionType\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.google_utils.google_utils import execute_paginated_retrieval\nfrom onyx.connectors.google_utils.resources import AdminService\nfrom onyx.connectors.google_utils.resources import get_admin_service\nfrom onyx.connectors.google_utils.resources import get_drive_service\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n\"\"\"\nFolder Permission Sync.\n\nEach folder is treated as a group. Each file has all ancestor folders\nas groups.\n\"\"\"\n\n\nclass FolderInfo(BaseModel):\n id: str\n permissions: list[GoogleDrivePermission]\n\n\ndef _get_all_folders(\n google_drive_connector: GoogleDriveConnector, skip_folders_without_permissions: bool\n) -> list[FolderInfo]:\n \"\"\"Have to get all folders since the group syncing system assumes all groups\n are returned every time.\n\n TODO: tweak things so we can fetch deltas.\n \"\"\"\n MAX_FAILED_PERCENTAGE = 0.5\n\n all_folders: list[FolderInfo] = []\n seen_folder_ids: set[str] = set()\n\n def _get_all_folders_for_user(\n google_drive_connector: GoogleDriveConnector,\n skip_folders_without_permissions: bool,\n user_email: str,\n ) -> None:\n \"\"\"Helper to get folders for a specific user + update shared seen_folder_ids\"\"\"\n drive_service = get_drive_service(\n google_drive_connector.creds,\n user_email,\n )\n\n for folder in get_modified_folders(\n service=drive_service,\n ):\n folder_id = folder[\"id\"]\n if folder_id in seen_folder_ids:\n logger.debug(f\"Folder {folder_id} has already been seen. Skipping.\")\n continue\n\n seen_folder_ids.add(folder_id)\n\n # Check if the folder has permission IDs but no permissions\n permission_ids = folder.get(\"permissionIds\", [])\n raw_permissions = folder.get(\"permissions\", [])\n\n if not raw_permissions and permission_ids:\n # Fetch permissions using the IDs\n permissions = get_folder_permissions_by_ids(\n drive_service, folder_id, permission_ids\n )\n else:\n permissions = [\n GoogleDrivePermission.from_drive_permission(permission)\n for permission in raw_permissions\n ]\n\n # Don't include inherited permissions, those will be captured\n # by the folder/shared drive itself\n permissions = [\n permission\n for permission in permissions\n if permission.inherited_from is None\n ]\n\n if not permissions and skip_folders_without_permissions:\n logger.debug(f\"Folder {folder_id} has no permissions. Skipping.\")\n continue\n\n all_folders.append(\n FolderInfo(\n id=folder_id,\n permissions=permissions,\n )\n )\n\n failed_count = 0\n user_emails = google_drive_connector._get_all_user_emails()\n for user_email in user_emails:\n try:\n _get_all_folders_for_user(\n google_drive_connector, skip_folders_without_permissions, user_email\n )\n except Exception:\n logger.exception(f\"Error getting folders for user {user_email}\")\n failed_count += 1\n\n if failed_count > MAX_FAILED_PERCENTAGE * len(user_emails):\n raise RuntimeError(\"Too many failed folder fetches during group sync\")\n\n return all_folders\n\n\ndef _drive_folder_to_onyx_group(\n folder: FolderInfo,\n group_email_to_member_emails_map: dict[str, list[str]],\n) -> ExternalUserGroup:\n \"\"\"\n Converts a folder into an Onyx group.\n \"\"\"\n anyone_can_access = False\n folder_member_emails: set[str] = set()\n\n for permission in folder.permissions:\n if permission.type == PermissionType.USER:\n if permission.email_address is None:\n logger.warning(\n f\"User email is None for folder {folder.id} permission {permission}\"\n )\n continue\n folder_member_emails.add(permission.email_address)\n elif permission.type == PermissionType.GROUP:\n if permission.email_address not in group_email_to_member_emails_map:\n logger.warning(\n f\"Group email {permission.email_address} for folder {folder.id} not found in group_email_to_member_emails_map\"\n )\n continue\n folder_member_emails.update(\n group_email_to_member_emails_map[permission.email_address]\n )\n elif permission.type == PermissionType.ANYONE:\n anyone_can_access = True\n\n return ExternalUserGroup(\n id=folder.id,\n user_emails=list(folder_member_emails),\n gives_anyone_access=anyone_can_access,\n )\n\n\n\"\"\"Individual Shared Drive / My Drive Permission Sync\"\"\"\n\n\ndef _get_drive_members(\n google_drive_connector: GoogleDriveConnector,\n admin_service: AdminService,\n) -> dict[str, tuple[set[str], set[str]]]:\n \"\"\"\n This builds a map of drive ids to their members (group and user emails).\n E.g. {\n \"drive_id_1\": ({\"group_email_1\"}, {\"user_email_1\", \"user_email_2\"}),\n \"drive_id_2\": ({\"group_email_3\"}, {\"user_email_3\"}),\n }\n \"\"\"\n\n # fetches shared drives only\n drive_ids = google_drive_connector.get_all_drive_ids()\n\n drive_id_to_members_map: dict[str, tuple[set[str], set[str]]] = {}\n drive_service = get_drive_service(\n google_drive_connector.creds,\n google_drive_connector.primary_admin_email,\n )\n\n admin_user_info = (\n admin_service.users()\n .get(userKey=google_drive_connector.primary_admin_email)\n .execute()\n )\n is_admin = admin_user_info.get(\"isAdmin\", False) or admin_user_info.get(\n \"isDelegatedAdmin\", False\n )\n\n for drive_id in drive_ids:\n group_emails: set[str] = set()\n user_emails: set[str] = set()\n\n try:\n for permission in execute_paginated_retrieval(\n drive_service.permissions().list,\n list_key=\"permissions\",\n fileId=drive_id,\n fields=\"permissions(emailAddress, type),nextPageToken\",\n supportsAllDrives=True,\n # can only set `useDomainAdminAccess` to true if the user\n # is an admin\n useDomainAdminAccess=is_admin,\n ):\n # NOTE: don't need to check for PermissionType.ANYONE since\n # you can't share a drive with the internet\n if permission[\"type\"] == PermissionType.GROUP:\n group_emails.add(permission[\"emailAddress\"])\n elif permission[\"type\"] == PermissionType.USER:\n user_emails.add(permission[\"emailAddress\"])\n except HttpError as e:\n if e.status_code == 404:\n logger.warning(\n f\"Error getting permissions for drive id {drive_id}. \"\n f\"User '{google_drive_connector.primary_admin_email}' likely \"\n f\"does not have access to this drive. Exception: {e}\"\n )\n else:\n raise e\n\n drive_id_to_members_map[drive_id] = (group_emails, user_emails)\n return drive_id_to_members_map\n\n\ndef _drive_member_map_to_onyx_groups(\n drive_id_to_members_map: dict[str, tuple[set[str], set[str]]],\n group_email_to_member_emails_map: dict[str, list[str]],\n) -> Generator[ExternalUserGroup, None, None]:\n \"\"\"The `user_emails` for the Shared Drive should be all individuals in the\n Shared Drive + the union of all flattened group emails.\"\"\"\n for drive_id, (group_emails, user_emails) in drive_id_to_members_map.items():\n drive_member_emails: set[str] = user_emails\n for group_email in group_emails:\n if group_email not in group_email_to_member_emails_map:\n logger.warning(\n f\"Group email {group_email} for drive {drive_id} not found in group_email_to_member_emails_map\"\n )\n continue\n drive_member_emails.update(group_email_to_member_emails_map[group_email])\n yield ExternalUserGroup(\n id=drive_id,\n user_emails=list(drive_member_emails),\n )\n\n\ndef _get_all_google_groups(\n admin_service: AdminService,\n google_domain: str,\n) -> set[str]:\n \"\"\"\n This gets all the group emails.\n \"\"\"\n group_emails: set[str] = set()\n for group in execute_paginated_retrieval(\n admin_service.groups().list,\n list_key=\"groups\",\n domain=google_domain,\n fields=\"groups(email),nextPageToken\",\n ):\n group_emails.add(group[\"email\"])\n return group_emails\n\n\ndef _google_group_to_onyx_group(\n admin_service: AdminService,\n group_email: str,\n) -> ExternalUserGroup:\n \"\"\"\n This maps google group emails to their member emails.\n \"\"\"\n group_member_emails: set[str] = set()\n for member in execute_paginated_retrieval(\n admin_service.members().list,\n list_key=\"members\",\n groupKey=group_email,\n fields=\"members(email),nextPageToken\",\n ):\n group_member_emails.add(member[\"email\"])\n\n return ExternalUserGroup(\n id=group_email,\n user_emails=list(group_member_emails),\n )\n\n\ndef _map_group_email_to_member_emails(\n admin_service: AdminService,\n group_emails: set[str],\n) -> dict[str, set[str]]:\n \"\"\"\n This maps group emails to their member emails.\n \"\"\"\n group_to_member_map: dict[str, set[str]] = {}\n for group_email in group_emails:\n group_member_emails: set[str] = set()\n for member in execute_paginated_retrieval(\n admin_service.members().list,\n list_key=\"members\",\n groupKey=group_email,\n fields=\"members(email),nextPageToken\",\n ):\n group_member_emails.add(member[\"email\"])\n\n group_to_member_map[group_email] = group_member_emails\n return group_to_member_map\n\n\ndef _build_onyx_groups(\n drive_id_to_members_map: dict[str, tuple[set[str], set[str]]],\n group_email_to_member_emails_map: dict[str, set[str]],\n folder_info: list[FolderInfo],\n) -> list[ExternalUserGroup]:\n onyx_groups: list[ExternalUserGroup] = []\n\n # Convert all drive member definitions to onyx groups\n # This is because having drive level access means you have\n # irrevocable access to all the files in the drive.\n for drive_id, (group_emails, user_emails) in drive_id_to_members_map.items():\n drive_member_emails: set[str] = user_emails\n for group_email in group_emails:\n if group_email not in group_email_to_member_emails_map:\n logger.warning(\n f\"Group email {group_email} for drive {drive_id} not found in group_email_to_member_emails_map\"\n )\n continue\n drive_member_emails.update(group_email_to_member_emails_map[group_email])\n onyx_groups.append(\n ExternalUserGroup(\n id=drive_id,\n user_emails=list(drive_member_emails),\n )\n )\n\n # Convert all folder permissions to onyx groups\n for folder in folder_info:\n anyone_can_access = False\n folder_member_emails: set[str] = set()\n for permission in folder.permissions:\n if permission.type == PermissionType.USER:\n if permission.email_address is None:\n logger.warning(\n f\"User email is None for folder {folder.id} permission {permission}\"\n )\n continue\n folder_member_emails.add(permission.email_address)\n elif permission.type == PermissionType.GROUP:\n if permission.email_address not in group_email_to_member_emails_map:\n logger.warning(\n f\"Group email {permission.email_address} for folder {folder.id} \"\n \"not found in group_email_to_member_emails_map\"\n )\n continue\n folder_member_emails.update(\n group_email_to_member_emails_map[permission.email_address]\n )\n elif permission.type == PermissionType.ANYONE:\n anyone_can_access = True\n\n onyx_groups.append(\n ExternalUserGroup(\n id=folder.id,\n user_emails=list(folder_member_emails),\n gives_anyone_access=anyone_can_access,\n )\n )\n\n # Convert all group member definitions to onyx groups\n for group_email, member_emails in group_email_to_member_emails_map.items():\n onyx_groups.append(\n ExternalUserGroup(\n id=group_email,\n user_emails=list(member_emails),\n )\n )\n\n return onyx_groups\n\n\ndef gdrive_group_sync(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair,\n) -> Generator[ExternalUserGroup, None, None]:\n # Initialize connector and build credential/service objects\n google_drive_connector = GoogleDriveConnector(\n **cc_pair.connector.connector_specific_config\n )\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n google_drive_connector.load_credentials(credential_json)\n admin_service = get_admin_service(\n google_drive_connector.creds, google_drive_connector.primary_admin_email\n )\n\n # Get all drive members\n drive_id_to_members_map = _get_drive_members(google_drive_connector, admin_service)\n\n # Get all group emails\n all_group_emails = _get_all_google_groups(\n admin_service, google_drive_connector.google_domain\n )\n\n # Each google group is an Onyx group, yield those\n group_email_to_member_emails_map: dict[str, list[str]] = {}\n for group_email in all_group_emails:\n onyx_group = _google_group_to_onyx_group(admin_service, group_email)\n group_email_to_member_emails_map[group_email] = onyx_group.user_emails\n yield onyx_group\n\n # Each drive is a group, yield those\n for onyx_group in _drive_member_map_to_onyx_groups(\n drive_id_to_members_map, group_email_to_member_emails_map\n ):\n yield onyx_group\n\n # Get all folder permissions\n folder_info = _get_all_folders(\n google_drive_connector=google_drive_connector,\n skip_folders_without_permissions=True,\n )\n for folder in folder_info:\n yield _drive_folder_to_onyx_group(folder, group_email_to_member_emails_map)\n" }, { "path": "backend/ee/onyx/external_permissions/google_drive/models.py", "content": "from enum import Enum\nfrom typing import Any\n\nfrom pydantic import BaseModel\n\n\nclass PermissionType(str, Enum):\n USER = \"user\"\n GROUP = \"group\"\n DOMAIN = \"domain\"\n ANYONE = \"anyone\"\n\n\nclass GoogleDrivePermissionDetails(BaseModel):\n # this is \"file\", \"member\", etc.\n # different from the `type` field within `GoogleDrivePermission`\n # Sometimes can be not, although not sure why...\n permission_type: str | None\n # this is \"reader\", \"writer\", \"owner\", etc.\n role: str\n # this is the id of the parent permission\n inherited_from: str | None\n\n\nclass GoogleDrivePermission(BaseModel):\n id: str\n # groups are also represented as email addresses within Drive\n # will be None for domain/global permissions\n email_address: str | None\n type: PermissionType\n domain: str | None # only applies to domain permissions\n permission_details: GoogleDrivePermissionDetails | None\n # Whether this permission makes the file discoverable in search\n # False means \"anyone with the link\" (not searchable/discoverable)\n # Only applicable for domain/anyone permission types\n allow_file_discovery: bool | None\n\n @classmethod\n def from_drive_permission(\n cls, drive_permission: dict[str, Any]\n ) -> \"GoogleDrivePermission\":\n # we seem to only get details for permissions that are inherited\n # we can get multiple details if a permission is inherited from multiple\n permission_details_list = drive_permission.get(\"permissionDetails\", [])\n permission_details: dict[str, Any] | None = (\n permission_details_list[0] if permission_details_list else None\n )\n return cls(\n id=drive_permission[\"id\"],\n email_address=drive_permission.get(\"emailAddress\"),\n type=PermissionType(drive_permission[\"type\"]),\n domain=drive_permission.get(\"domain\"),\n allow_file_discovery=drive_permission.get(\"allowFileDiscovery\"),\n permission_details=(\n GoogleDrivePermissionDetails(\n permission_type=permission_details.get(\"type\"),\n role=permission_details.get(\"role\", \"\"),\n inherited_from=permission_details.get(\"inheritedFrom\"),\n )\n if permission_details\n else None\n ),\n )\n\n @property\n def inherited_from(self) -> str | None:\n if self.permission_details:\n return self.permission_details.inherited_from\n return None\n" }, { "path": "backend/ee/onyx/external_permissions/google_drive/permission_retrieval.py", "content": "from retry import retry\n\nfrom ee.onyx.external_permissions.google_drive.models import GoogleDrivePermission\nfrom onyx.connectors.google_utils.google_utils import execute_paginated_retrieval\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n@retry(tries=3, delay=2, backoff=2)\ndef get_permissions_by_ids(\n drive_service: GoogleDriveService,\n doc_id: str,\n permission_ids: list[str],\n) -> list[GoogleDrivePermission]:\n \"\"\"\n Fetches permissions for a document based on a list of permission IDs.\n\n Args:\n drive_service: The Google Drive service instance\n doc_id: The ID of the document to fetch permissions for\n permission_ids: A list of permission IDs to filter by\n\n Returns:\n A list of GoogleDrivePermission objects matching the provided permission IDs\n \"\"\"\n if not permission_ids:\n return []\n\n # Create a set for faster lookup\n permission_id_set = set(permission_ids)\n\n # Fetch all permissions for the document\n fetched_permissions = execute_paginated_retrieval(\n retrieval_function=drive_service.permissions().list,\n list_key=\"permissions\",\n fileId=doc_id,\n fields=\"permissions(id, emailAddress, type, domain, allowFileDiscovery, permissionDetails),nextPageToken\",\n supportsAllDrives=True,\n continue_on_404_or_403=True,\n )\n\n # Filter permissions by ID and convert to GoogleDrivePermission objects\n filtered_permissions = []\n for permission in fetched_permissions:\n permission_id = permission.get(\"id\")\n if permission_id in permission_id_set:\n google_drive_permission = GoogleDrivePermission.from_drive_permission(\n permission\n )\n filtered_permissions.append(google_drive_permission)\n\n # Log if we couldn't find all requested permission IDs\n if len(filtered_permissions) < len(permission_ids):\n missing_ids = permission_id_set - {p.id for p in filtered_permissions if p.id}\n logger.warning(\n f\"Could not find all requested permission IDs for document {doc_id}. Missing IDs: {missing_ids}\"\n )\n\n return filtered_permissions\n" }, { "path": "backend/ee/onyx/external_permissions/jira/__init__.py", "content": "" }, { "path": "backend/ee/onyx/external_permissions/jira/doc_sync.py", "content": "from collections.abc import Generator\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom ee.onyx.external_permissions.utils import generic_doc_sync\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.jira.connector import JiraConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nJIRA_DOC_SYNC_TAG = \"jira_doc_sync\"\n\n\ndef jira_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,\n callback: IndexingHeartbeatInterface | None = None,\n) -> Generator[ElementExternalAccess, None, None]:\n jira_connector = JiraConnector(\n **cc_pair.connector.connector_specific_config,\n )\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n jira_connector.load_credentials(credential_json)\n\n yield from generic_doc_sync(\n cc_pair=cc_pair,\n fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,\n callback=callback,\n doc_source=DocumentSource.JIRA,\n slim_connector=jira_connector,\n label=JIRA_DOC_SYNC_TAG,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/jira/group_sync.py", "content": "from collections.abc import Generator\nfrom typing import Any\n\nfrom jira import JIRA\nfrom jira.exceptions import JIRAError\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom onyx.connectors.jira.utils import build_jira_client\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_ATLASSIAN_ACCOUNT_TYPE = \"atlassian\"\n_GROUP_MEMBER_PAGE_SIZE = 50\n\n# The GET /group/member endpoint was introduced in Jira 6.0.\n# Jira versions older than 6.0 do not have group management REST APIs at all.\n_MIN_JIRA_VERSION_FOR_GROUP_MEMBER = \"6.0\"\n\n\ndef _fetch_group_member_page(\n jira_client: JIRA,\n group_name: str,\n start_at: int,\n) -> dict[str, Any]:\n \"\"\"Fetch a single page from the non-deprecated GET /group/member endpoint.\n\n The old GET /group endpoint (used by jira_client.group_members()) is deprecated\n and decommissioned in Jira Server 10.3+. This uses the replacement endpoint\n directly via the library's internal _get_json helper, following the same pattern\n as enhanced_search_ids / bulk_fetch_issues in connector.py.\n\n There is an open PR to the library to switch to this endpoint since last year:\n https://github.com/pycontribs/jira/pull/2356\n so once it is merged and released, we can switch to using the library function.\n \"\"\"\n try:\n return jira_client._get_json(\n \"group/member\",\n params={\n \"groupname\": group_name,\n \"includeInactiveUsers\": \"false\",\n \"startAt\": start_at,\n \"maxResults\": _GROUP_MEMBER_PAGE_SIZE,\n },\n )\n except JIRAError as e:\n if e.status_code == 404:\n raise RuntimeError(\n f\"GET /group/member returned 404 for group '{group_name}'. \"\n f\"This endpoint requires Jira {_MIN_JIRA_VERSION_FOR_GROUP_MEMBER}+. \"\n f\"If you are running a self-hosted Jira instance, please upgrade \"\n f\"to at least Jira {_MIN_JIRA_VERSION_FOR_GROUP_MEMBER}.\"\n ) from e\n raise\n\n\ndef _get_group_member_emails(\n jira_client: JIRA,\n group_name: str,\n) -> set[str]:\n \"\"\"Get all member emails for a single Jira group.\n\n Uses the non-deprecated GET /group/member endpoint which returns full user\n objects including accountType, so we can filter out app/customer accounts\n without making separate user() calls.\n \"\"\"\n emails: set[str] = set()\n start_at = 0\n\n while True:\n try:\n page = _fetch_group_member_page(jira_client, group_name, start_at)\n except Exception as e:\n logger.error(f\"Error fetching members for group {group_name}: {e}\")\n raise\n\n members: list[dict[str, Any]] = page.get(\"values\", [])\n for member in members:\n account_type = member.get(\"accountType\")\n # On Jira DC < 9.0, accountType is absent; include those users.\n # On Cloud / DC 9.0+, filter to real user accounts only.\n if account_type is not None and account_type != _ATLASSIAN_ACCOUNT_TYPE:\n continue\n\n email = member.get(\"emailAddress\")\n if email:\n emails.add(email)\n else:\n logger.warning(\n f\"Atlassian user {member.get('accountId', 'unknown')} in group {group_name} has no visible email address\"\n )\n\n if page.get(\"isLast\", True) or not members:\n break\n start_at += len(members)\n\n return emails\n\n\ndef jira_group_sync(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair,\n) -> Generator[ExternalUserGroup, None, None]:\n \"\"\"Sync Jira groups and their members, yielding one group at a time.\n\n Streams group-by-group rather than accumulating all groups in memory.\n \"\"\"\n jira_base_url = cc_pair.connector.connector_specific_config.get(\"jira_base_url\", \"\")\n scoped_token = cc_pair.connector.connector_specific_config.get(\n \"scoped_token\", False\n )\n\n if not jira_base_url:\n raise ValueError(\"No jira_base_url found in connector config\")\n\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n jira_client = build_jira_client(\n credentials=credential_json,\n jira_base=jira_base_url,\n scoped_token=scoped_token,\n )\n\n group_names = jira_client.groups()\n if not group_names:\n raise ValueError(f\"No groups found for cc_pair_id={cc_pair.id}\")\n\n logger.info(f\"Found {len(group_names)} groups in Jira\")\n\n for group_name in group_names:\n if not group_name:\n continue\n\n member_emails = _get_group_member_emails(\n jira_client=jira_client,\n group_name=group_name,\n )\n if not member_emails:\n logger.debug(f\"No members found for group {group_name}\")\n continue\n\n logger.debug(f\"Found {len(member_emails)} members for group {group_name}\")\n yield ExternalUserGroup(\n id=group_name,\n user_emails=list(member_emails),\n )\n" }, { "path": "backend/ee/onyx/external_permissions/jira/models.py", "content": "from typing import Any\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic.alias_generators import to_camel\n\n\nHolder = dict[str, Any]\n\n\nclass Permission(BaseModel):\n id: int\n permission: str\n holder: Holder | None\n\n\nclass User(BaseModel):\n account_id: str\n email_address: str\n display_name: str\n active: bool\n\n model_config = ConfigDict(\n alias_generator=to_camel,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/jira/page_access.py", "content": "from collections import defaultdict\n\nfrom jira import JIRA\nfrom jira.resources import PermissionScheme\nfrom pydantic import ValidationError\n\nfrom ee.onyx.external_permissions.jira.models import Holder\nfrom ee.onyx.external_permissions.jira.models import Permission\nfrom ee.onyx.external_permissions.jira.models import User\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.utils.logger import setup_logger\n\nHolderMap = dict[str, list[Holder]]\n\n\nlogger = setup_logger()\n\n\ndef _get_role_id(holder: Holder) -> str | None:\n return holder.get(\"value\") or holder.get(\"parameter\")\n\n\ndef _build_holder_map(permissions: list[dict]) -> dict[str, list[Holder]]:\n \"\"\"\n A \"Holder\" in JIRA is a person / entity who \"holds\" the corresponding permission.\n It can have different types. They can be one of (but not limited to):\n - user (an explicitly whitelisted user)\n - projectRole (for project level \"roles\")\n - reporter (the reporter of an issue)\n\n A \"Holder\" usually has following structure:\n - `{ \"type\": \"user\", \"value\": \"$USER_ID\", \"user\": { .. }, .. }`\n - `{ \"type\": \"projectRole\", \"value\": \"$PROJECT_ID\", .. }`\n\n When we fetch the PermissionSchema from JIRA, we retrieve a list of \"Holder\"s.\n The list of \"Holder\"s can have multiple \"Holder\"s of the same type in the list (e.g., you can have two `\"type\": \"user\"`s in\n there, each corresponding to a different user).\n This function constructs a map of \"Holder\" types to a list of the \"Holder\"s which contained that type.\n\n Returns:\n A dict from the \"Holder\" type to the actual \"Holder\" instance.\n\n Example:\n ```\n {\n \"user\": [\n { \"type\": \"user\", \"value\": \"10000\", \"user\": { .. }, .. },\n { \"type\": \"user\", \"value\": \"10001\", \"user\": { .. }, .. },\n ],\n \"projectRole\": [\n { \"type\": \"projectRole\", \"value\": \"10010\", .. },\n { \"type\": \"projectRole\", \"value\": \"10011\", .. },\n ],\n \"applicationRole\": [\n { \"type\": \"applicationRole\" },\n ],\n ..\n }\n ```\n \"\"\"\n\n holder_map: defaultdict[str, list[Holder]] = defaultdict(list)\n\n for raw_perm in permissions:\n if not hasattr(raw_perm, \"raw\"):\n logger.warning(f\"Expected a 'raw' field, but none was found: {raw_perm=}\")\n continue\n\n permission = Permission(**raw_perm.raw)\n\n # We only care about ability to browse through projects + issues (not other permissions such as read/write).\n if permission.permission != \"BROWSE_PROJECTS\":\n continue\n\n # In order to associate this permission to some Atlassian entity, we need the \"Holder\".\n # If this doesn't exist, then we cannot associate this permission to anyone; just skip.\n if not permission.holder:\n logger.warning(\n f\"Expected to find a permission holder, but none was found: {permission=}\"\n )\n continue\n\n type = permission.holder.get(\"type\")\n if not type:\n logger.warning(\n f\"Expected to find the type of permission holder, but none was found: {permission=}\"\n )\n continue\n\n holder_map[type].append(permission.holder)\n\n return holder_map\n\n\ndef _get_user_emails(user_holders: list[Holder]) -> list[str]:\n emails = []\n\n for user_holder in user_holders:\n if \"user\" not in user_holder:\n continue\n raw_user_dict = user_holder[\"user\"]\n\n try:\n user_model = User.model_validate(raw_user_dict)\n except ValidationError:\n logger.error(\n \"Expected to be able to serialize the raw-user-dict into an instance of `User`, but validation failed;\"\n f\"{raw_user_dict=}\"\n )\n continue\n\n emails.append(user_model.email_address)\n\n return emails\n\n\ndef _get_user_emails_and_groups_from_project_roles(\n jira_client: JIRA,\n jira_project: str,\n project_role_holders: list[Holder],\n) -> tuple[list[str], list[str]]:\n \"\"\"\n Get user emails and group names from project roles.\n Returns a tuple of (emails, group_names).\n \"\"\"\n # Get role IDs - Cloud uses \"value\", Data Center uses \"parameter\"\n role_ids = []\n for holder in project_role_holders:\n role_id = _get_role_id(holder)\n if role_id:\n role_ids.append(role_id)\n else:\n logger.warning(f\"No value or parameter in projectRole holder: {holder}\")\n\n roles = [\n jira_client.project_role(project=jira_project, id=role_id)\n for role_id in role_ids\n ]\n\n emails = []\n groups = []\n\n for role in roles:\n if not hasattr(role, \"actors\"):\n logger.warning(f\"Project role {role} has no actors attribute\")\n continue\n\n for actor in role.actors:\n # Handle group actors\n if hasattr(actor, \"actorGroup\"):\n group_name = getattr(actor.actorGroup, \"name\", None) or getattr(\n actor.actorGroup, \"displayName\", None\n )\n if group_name:\n groups.append(group_name)\n continue\n\n # Handle user actors\n if hasattr(actor, \"actorUser\"):\n account_id = getattr(actor.actorUser, \"accountId\", None)\n if not account_id:\n logger.error(f\"No accountId in actorUser: {actor.actorUser}\")\n continue\n\n user = jira_client.user(id=account_id)\n if not hasattr(user, \"accountType\") or user.accountType != \"atlassian\":\n logger.info(\n f\"Skipping user {account_id} because it is not an atlassian user\"\n )\n continue\n\n if not hasattr(user, \"emailAddress\"):\n msg = f\"User's email address was not able to be retrieved; {actor.actorUser.accountId=}\"\n if hasattr(user, \"displayName\"):\n msg += f\" {actor.displayName=}\"\n logger.warning(msg)\n continue\n\n emails.append(user.emailAddress)\n continue\n\n logger.debug(f\"Skipping actor type: {actor}\")\n\n return emails, groups\n\n\ndef _build_external_access_from_holder_map(\n jira_client: JIRA, jira_project: str, holder_map: HolderMap\n) -> ExternalAccess:\n \"\"\"\n Build ExternalAccess from the holder map.\n\n Holder types handled:\n - \"anyone\": Public project, anyone can access\n - \"applicationRole\": All users with a Jira license can access (treated as public)\n - \"user\": Specific users with access\n - \"projectRole\": Project roles containing users and/or groups\n - \"group\": Groups directly assigned in the permission scheme\n \"\"\"\n # Public access - anyone can view\n if \"anyone\" in holder_map:\n return ExternalAccess(\n external_user_emails=set(), external_user_group_ids=set(), is_public=True\n )\n\n # applicationRole means all users with a Jira license can access - treat as public\n if \"applicationRole\" in holder_map:\n return ExternalAccess(\n external_user_emails=set(), external_user_group_ids=set(), is_public=True\n )\n\n # Get emails from explicit user holders\n user_emails = (\n _get_user_emails(user_holders=holder_map[\"user\"])\n if \"user\" in holder_map\n else []\n )\n\n # Get emails and groups from project roles\n project_role_user_emails: list[str] = []\n project_role_groups: list[str] = []\n if \"projectRole\" in holder_map:\n project_role_user_emails, project_role_groups = (\n _get_user_emails_and_groups_from_project_roles(\n jira_client=jira_client,\n jira_project=jira_project,\n project_role_holders=holder_map[\"projectRole\"],\n )\n )\n\n # Get groups directly assigned in permission scheme (common in Data Center)\n # Format: {'type': 'group', 'parameter': 'group-name', 'expand': 'group'}\n direct_groups: list[str] = []\n if \"group\" in holder_map:\n for group_holder in holder_map[\"group\"]:\n group_name = _get_role_id(group_holder)\n if group_name:\n direct_groups.append(group_name)\n else:\n logger.error(f\"No parameter/value in group holder: {group_holder}\")\n\n external_user_emails = set(user_emails + project_role_user_emails)\n external_user_group_ids = set(project_role_groups + direct_groups)\n\n return ExternalAccess(\n external_user_emails=external_user_emails,\n external_user_group_ids=external_user_group_ids,\n is_public=False,\n )\n\n\ndef get_project_permissions(\n jira_client: JIRA,\n jira_project: str,\n add_prefix: bool = False,\n) -> ExternalAccess | None:\n \"\"\"\n Get project permissions from Jira.\n\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path).\n \"\"\"\n project_permissions: PermissionScheme = jira_client.project_permissionscheme(\n project=jira_project\n )\n\n if not hasattr(project_permissions, \"permissions\"):\n logger.error(f\"Project {jira_project} has no permissions attribute\")\n return None\n\n if not isinstance(project_permissions.permissions, list):\n logger.error(f\"Project {jira_project} permissions is not a list\")\n return None\n\n holder_map = _build_holder_map(permissions=project_permissions.permissions)\n\n external_access = _build_external_access_from_holder_map(\n jira_client=jira_client, jira_project=jira_project, holder_map=holder_map\n )\n\n # Prefix group IDs with source type if requested (for indexing path)\n if add_prefix and external_access and external_access.external_user_group_ids:\n prefixed_groups = {\n build_ext_group_name_for_onyx(g, DocumentSource.JIRA)\n for g in external_access.external_user_group_ids\n }\n return ExternalAccess(\n external_user_emails=external_access.external_user_emails,\n external_user_group_ids=prefixed_groups,\n is_public=external_access.is_public,\n )\n\n return external_access\n" }, { "path": "backend/ee/onyx/external_permissions/perm_sync_types.py", "content": "from collections.abc import Callable\nfrom collections.abc import Generator\nfrom typing import Optional\nfrom typing import Protocol\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup # noqa\nfrom onyx.access.models import DocExternalAccess # noqa\nfrom onyx.access.models import ElementExternalAccess # noqa\nfrom onyx.access.models import NodeExternalAccess # noqa\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.models import ConnectorCredentialPair # noqa\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import SortOrder\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface # noqa\n\n\nclass FetchAllDocumentsFunction(Protocol):\n \"\"\"Protocol for a function that fetches documents for a connector credential pair.\n\n This protocol defines the interface for functions that retrieve documents\n from the database, typically used in permission synchronization workflows.\n \"\"\"\n\n def __call__(\n self,\n sort_order: SortOrder | None,\n ) -> list[DocumentRow]:\n \"\"\"\n Fetches documents for a connector credential pair.\n \"\"\"\n ...\n\n\nclass FetchAllDocumentsIdsFunction(Protocol):\n \"\"\"Protocol for a function that fetches document IDs for a connector credential pair.\n\n This protocol defines the interface for functions that retrieve document IDs\n from the database, typically used in permission synchronization workflows.\n \"\"\"\n\n def __call__(\n self,\n ) -> list[str]:\n \"\"\"\n Fetches document IDs for a connector credential pair.\n \"\"\"\n ...\n\n\n# Defining the input/output types for the sync functions\nDocSyncFuncType = Callable[\n [\n ConnectorCredentialPair,\n FetchAllDocumentsFunction,\n FetchAllDocumentsIdsFunction,\n Optional[IndexingHeartbeatInterface],\n ],\n Generator[ElementExternalAccess, None, None],\n]\n\nGroupSyncFuncType = Callable[\n [\n str, # tenant_id\n ConnectorCredentialPair, # cc_pair\n ],\n Generator[ExternalUserGroup, None, None],\n]\n\n# list of chunks to be censored and the user email. returns censored chunks\nCensoringFuncType = Callable[[list[InferenceChunk], str], list[InferenceChunk]]\n" }, { "path": "backend/ee/onyx/external_permissions/post_query_censoring.py", "content": "from ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs\nfrom ee.onyx.external_permissions.sync_params import get_all_censoring_enabled_sources\nfrom ee.onyx.external_permissions.sync_params import get_source_perm_sync_config\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.context.search.pipeline import InferenceChunk\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import User\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _get_all_censoring_enabled_sources() -> set[DocumentSource]:\n \"\"\"\n Returns the set of sources that have censoring enabled.\n This is based on if the access_type is set to sync and the connector\n source has a censoring config.\n\n NOTE: This means if there is a source has a single cc_pair that is sync,\n all chunks for that source will be censored, even if the connector that\n indexed that chunk is not sync. This was done to avoid getting the cc_pair\n for every single chunk.\n \"\"\"\n all_censoring_enabled_sources = get_all_censoring_enabled_sources()\n with get_session_with_current_tenant() as db_session:\n enabled_sync_connectors = get_all_auto_sync_cc_pairs(db_session)\n return {\n cc_pair.connector.source\n for cc_pair in enabled_sync_connectors\n if cc_pair.connector.source in all_censoring_enabled_sources\n }\n\n\n# NOTE: This is only called if ee is enabled.\ndef _post_query_chunk_censoring(\n chunks: list[InferenceChunk],\n user: User,\n) -> list[InferenceChunk]:\n \"\"\"\n This function checks all chunks to see if they need to be sent to a censoring\n function. If they do, it sends them to the censoring function and returns the\n censored chunks. If they don't, it returns the original chunks.\n \"\"\"\n sources_to_censor = _get_all_censoring_enabled_sources()\n\n # Anonymous users can only access public (non-permission-synced) content\n if user.is_anonymous:\n return [chunk for chunk in chunks if chunk.source_type not in sources_to_censor]\n\n final_chunk_dict: dict[str, InferenceChunk] = {}\n chunks_to_process: dict[DocumentSource, list[InferenceChunk]] = {}\n for chunk in chunks:\n # Separate out chunks that require permission post-processing by source\n if chunk.source_type in sources_to_censor:\n chunks_to_process.setdefault(chunk.source_type, []).append(chunk)\n else:\n final_chunk_dict[chunk.unique_id] = chunk\n\n # For each source, filter out the chunks using the permission\n # check function for that source\n # TODO: Use a threadpool/multiprocessing to process the sources in parallel\n for source, chunks_for_source in chunks_to_process.items():\n sync_config = get_source_perm_sync_config(source)\n if sync_config is None or sync_config.censoring_config is None:\n raise ValueError(f\"No sync config found for {source}\")\n\n censor_chunks_for_source = sync_config.censoring_config.chunk_censoring_func\n try:\n censored_chunks = censor_chunks_for_source(chunks_for_source, user.email)\n except Exception as e:\n logger.exception(\n f\"Failed to censor chunks for source {source} so throwing out all chunks for this source and continuing: {e}\"\n )\n continue\n\n for censored_chunk in censored_chunks:\n final_chunk_dict[censored_chunk.unique_id] = censored_chunk\n\n # IMPORTANT: make sure to retain the same ordering as the original `chunks` passed in\n final_chunk_list: list[InferenceChunk] = []\n for chunk in chunks:\n # only if the chunk is in the final censored chunks, add it to the final list\n # if it is missing, that means it was intentionally left out\n if chunk.unique_id in final_chunk_dict:\n final_chunk_list.append(final_chunk_dict[chunk.unique_id])\n\n return final_chunk_list\n" }, { "path": "backend/ee/onyx/external_permissions/salesforce/postprocessing.py", "content": "import time\n\nfrom ee.onyx.db.external_perm import fetch_external_groups_for_user_email_and_group_ids\nfrom ee.onyx.external_permissions.salesforce.utils import (\n get_any_salesforce_client_for_doc_id,\n)\nfrom ee.onyx.external_permissions.salesforce.utils import get_objects_access_for_user_id\nfrom ee.onyx.external_permissions.salesforce.utils import (\n get_salesforce_user_id_from_email,\n)\nfrom onyx.configs.app_configs import BLURB_SIZE\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# Types\nChunkKey = tuple[str, int] # (doc_id, chunk_id)\nContentRange = tuple[int, int | None] # (start_index, end_index) None means to the end\n\n\n# NOTE: Used for testing timing\ndef _get_dummy_object_access_map(\n object_ids: set[str],\n user_email: str, # noqa: ARG001\n chunks: list[InferenceChunk], # noqa: ARG001\n) -> dict[str, bool]:\n time.sleep(0.15)\n # return {object_id: True for object_id in object_ids}\n import random\n\n return {object_id: random.choice([True, False]) for object_id in object_ids}\n\n\ndef _get_objects_access_for_user_email_from_salesforce(\n object_ids: set[str],\n user_email: str,\n chunks: list[InferenceChunk],\n) -> dict[str, bool] | None:\n \"\"\"\n This function wraps the salesforce call as we may want to change how this\n is done in the future. (E.g. replace it with the above function)\n \"\"\"\n # This is cached in the function so the first query takes an extra 0.1-0.3 seconds\n # but subsequent queries for this source are essentially instant\n first_doc_id = chunks[0].document_id\n with get_session_with_current_tenant() as db_session:\n salesforce_client = get_any_salesforce_client_for_doc_id(\n db_session, first_doc_id\n )\n\n # This is cached in the function so the first query takes an extra 0.1-0.3 seconds\n # but subsequent queries by the same user are essentially instant\n start_time = time.monotonic()\n user_id = get_salesforce_user_id_from_email(salesforce_client, user_email)\n end_time = time.monotonic()\n logger.info(\n f\"Time taken to get Salesforce user ID: {end_time - start_time} seconds\"\n )\n if user_id is None:\n logger.warning(f\"User '{user_email}' not found in Salesforce\")\n return None\n\n # This is the only query that is not cached in the function\n # so it takes 0.1-0.2 seconds total\n object_id_to_access = get_objects_access_for_user_id(\n salesforce_client, user_id, list(object_ids)\n )\n logger.debug(f\"Object ID to access: {object_id_to_access}\")\n return object_id_to_access\n\n\ndef _extract_salesforce_object_id_from_url(url: str) -> str:\n return url.split(\"/\")[-1]\n\n\ndef _get_object_ranges_for_chunk(\n chunk: InferenceChunk,\n) -> dict[str, list[ContentRange]]:\n \"\"\"\n Given a chunk, return a dictionary of salesforce object ids and the content ranges\n for that object id in the current chunk\n \"\"\"\n if chunk.source_links is None:\n return {}\n\n object_ranges: dict[str, list[ContentRange]] = {}\n end_index = None\n descending_source_links = sorted(\n chunk.source_links.items(), key=lambda x: x[0], reverse=True\n )\n for start_index, url in descending_source_links:\n object_id = _extract_salesforce_object_id_from_url(url)\n if object_id not in object_ranges:\n object_ranges[object_id] = []\n object_ranges[object_id].append((start_index, end_index))\n end_index = start_index\n return object_ranges\n\n\ndef _create_empty_censored_chunk(uncensored_chunk: InferenceChunk) -> InferenceChunk:\n \"\"\"\n Create a copy of the unfiltered chunk where potentially sensitive content is removed\n to be added later if the user has access to each of the sub-objects\n \"\"\"\n empty_censored_chunk = InferenceChunk(\n **uncensored_chunk.model_dump(),\n )\n empty_censored_chunk.content = \"\"\n empty_censored_chunk.blurb = \"\"\n empty_censored_chunk.source_links = {}\n return empty_censored_chunk\n\n\ndef _update_censored_chunk(\n censored_chunk: InferenceChunk,\n uncensored_chunk: InferenceChunk,\n content_range: ContentRange,\n) -> InferenceChunk:\n \"\"\"\n Update the filtered chunk with the content and source links from the unfiltered chunk using the content ranges\n \"\"\"\n start_index, end_index = content_range\n\n # Update the content of the filtered chunk\n permitted_content = uncensored_chunk.content[start_index:end_index]\n permitted_section_start_index = len(censored_chunk.content)\n censored_chunk.content = permitted_content + censored_chunk.content\n\n # Update the source links of the filtered chunk\n if uncensored_chunk.source_links is not None:\n if censored_chunk.source_links is None:\n censored_chunk.source_links = {}\n link_content = uncensored_chunk.source_links[start_index]\n censored_chunk.source_links[permitted_section_start_index] = link_content\n\n # Update the blurb of the filtered chunk\n censored_chunk.blurb = censored_chunk.content[:BLURB_SIZE]\n\n return censored_chunk\n\n\n# TODO: Generalize this to other sources\ndef censor_salesforce_chunks(\n chunks: list[InferenceChunk],\n user_email: str,\n # This is so we can provide a mock access map for testing\n access_map: dict[str, bool] | None = None,\n) -> list[InferenceChunk]:\n # object_id -> list[((doc_id, chunk_id), (start_index, end_index))]\n object_to_content_map: dict[str, list[tuple[ChunkKey, ContentRange]]] = {}\n\n # (doc_id, chunk_id) -> chunk\n uncensored_chunks: dict[ChunkKey, InferenceChunk] = {}\n\n # keep track of all object ids that we have seen to make it easier to get\n # the access for these object ids\n object_ids: set[str] = set()\n\n for chunk in chunks:\n chunk_key = (chunk.document_id, chunk.chunk_id)\n # create a dictionary to quickly look up the unfiltered chunk\n uncensored_chunks[chunk_key] = chunk\n\n # for each chunk, get a dictionary of object ids and the content ranges\n # for that object id in the current chunk\n object_ranges_for_chunk = _get_object_ranges_for_chunk(chunk)\n for object_id, ranges in object_ranges_for_chunk.items():\n object_ids.add(object_id)\n for start_index, end_index in ranges:\n object_to_content_map.setdefault(object_id, []).append(\n (chunk_key, (start_index, end_index))\n )\n\n # This is so we can provide a mock access map for testing\n if access_map is None:\n access_map = _get_objects_access_for_user_email_from_salesforce(\n object_ids=object_ids,\n user_email=user_email,\n chunks=chunks,\n )\n if access_map is None:\n # If the user is not found in Salesforce, access_map will be None\n # so we should just return an empty list because no chunks will be\n # censored\n return []\n\n censored_chunks: dict[ChunkKey, InferenceChunk] = {}\n for object_id, content_list in object_to_content_map.items():\n # if the user does not have access to the object, or the object is not in the\n # access_map, do not include its content in the filtered chunks\n if not access_map.get(object_id, False):\n continue\n\n # if we got this far, the user has access to the object so we can create or update\n # the filtered chunk(s) for this object\n # NOTE: we only create a censored chunk if the user has access to some\n # part of the chunk\n for chunk_key, content_range in content_list:\n if chunk_key not in censored_chunks:\n censored_chunks[chunk_key] = _create_empty_censored_chunk(\n uncensored_chunks[chunk_key]\n )\n\n uncensored_chunk = uncensored_chunks[chunk_key]\n censored_chunk = _update_censored_chunk(\n censored_chunk=censored_chunks[chunk_key],\n uncensored_chunk=uncensored_chunk,\n content_range=content_range,\n )\n censored_chunks[chunk_key] = censored_chunk\n\n return list(censored_chunks.values())\n\n\n# NOTE: This is not used anywhere.\ndef _get_objects_access_for_user_email(\n object_ids: set[str], user_email: str\n) -> dict[str, bool]:\n with get_session_with_current_tenant() as db_session:\n external_groups = fetch_external_groups_for_user_email_and_group_ids(\n db_session=db_session,\n user_email=user_email,\n # Maybe make a function that adds a salesforce prefix to the group ids\n group_ids=list(object_ids),\n )\n external_group_ids = {group.external_user_group_id for group in external_groups}\n return {group_id: group_id in external_group_ids for group_id in object_ids}\n" }, { "path": "backend/ee/onyx/external_permissions/salesforce/utils.py", "content": "from simple_salesforce import Salesforce\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.document import get_cc_pairs_for_document\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_ANY_SALESFORCE_CLIENT: Salesforce | None = None\n\n\ndef get_any_salesforce_client_for_doc_id(\n db_session: Session, doc_id: str\n) -> Salesforce:\n \"\"\"\n We create a salesforce client for the first cc_pair for the first doc_id where\n salesforce censoring is enabled. After that we just cache and reuse the same\n client for all queries.\n\n We do this to reduce the number of postgres queries we make at query time.\n\n This may be problematic if they are using multiple cc_pairs for salesforce.\n E.g. there are 2 different credential sets for 2 different salesforce cc_pairs\n but only one has the permissions to access the permissions needed for the query.\n \"\"\"\n\n # NOTE: this global seems very very bad\n global _ANY_SALESFORCE_CLIENT\n if _ANY_SALESFORCE_CLIENT is None:\n cc_pairs = get_cc_pairs_for_document(db_session, doc_id)\n first_cc_pair = cc_pairs[0]\n credential_json = (\n first_cc_pair.credential.credential_json.get_value(apply_mask=False)\n if first_cc_pair.credential.credential_json\n else {}\n )\n _ANY_SALESFORCE_CLIENT = Salesforce(\n username=credential_json[\"sf_username\"],\n password=credential_json[\"sf_password\"],\n security_token=credential_json[\"sf_security_token\"],\n )\n return _ANY_SALESFORCE_CLIENT\n\n\ndef _query_salesforce_user_id(sf_client: Salesforce, user_email: str) -> str | None:\n query = f\"SELECT Id FROM User WHERE Username = '{user_email}' AND IsActive = true\"\n result = sf_client.query(query)\n if len(result[\"records\"]) > 0:\n return result[\"records\"][0][\"Id\"]\n\n # try emails\n query = f\"SELECT Id FROM User WHERE Email = '{user_email}' AND IsActive = true\"\n result = sf_client.query(query)\n if len(result[\"records\"]) > 0:\n return result[\"records\"][0][\"Id\"]\n\n return None\n\n\n# This contains only the user_ids that we have found in Salesforce.\n# If we don't know their user_id, we don't store anything in the cache.\n_CACHED_SF_EMAIL_TO_ID_MAP: dict[str, str] = {}\n\n\ndef get_salesforce_user_id_from_email(\n sf_client: Salesforce,\n user_email: str,\n) -> str | None:\n \"\"\"\n We cache this so we don't have to query Salesforce for every query and salesforce\n user IDs never change.\n Memory usage is fine because we just store 2 small strings per user.\n\n If the email is not in the cache, we check the local salesforce database for the info.\n If the user is not found in the local salesforce database, we query Salesforce.\n Whatever we get back from Salesforce is added to the database.\n If no user_id is found, we add a NULL_ID_STRING to the database for that email so\n we don't query Salesforce again (which is slow) but we still check the local salesforce\n database every query until a user id is found. This is acceptable because the query time\n is quite fast.\n If a user_id is created in Salesforce, it will be added to the local salesforce database\n next time the connector is run. Then that value will be found in this function and cached.\n\n NOTE: First time this runs, it may be slow if it hasn't already been updated in the local\n salesforce database. (Around 0.1-0.3 seconds)\n If it's cached or stored in the local salesforce database, it's fast (<0.001 seconds).\n \"\"\"\n\n # NOTE: this global seems bad\n global _CACHED_SF_EMAIL_TO_ID_MAP\n if user_email in _CACHED_SF_EMAIL_TO_ID_MAP:\n if _CACHED_SF_EMAIL_TO_ID_MAP[user_email] is not None:\n return _CACHED_SF_EMAIL_TO_ID_MAP[user_email]\n\n # some caching via sqlite existed here before ... check history if interested\n\n # ...query Salesforce and store the result in the database\n user_id = _query_salesforce_user_id(sf_client, user_email)\n\n if user_id is None:\n return None\n\n # If the found user_id is real, cache it\n _CACHED_SF_EMAIL_TO_ID_MAP[user_email] = user_id\n return user_id\n\n\n_MAX_RECORD_IDS_PER_QUERY = 200\n\n\ndef get_objects_access_for_user_id(\n salesforce_client: Salesforce,\n user_id: str,\n record_ids: list[str],\n) -> dict[str, bool]:\n \"\"\"\n Salesforce has a limit of 200 record ids per query. So we just truncate\n the list of record ids to 200. We only ever retrieve 50 chunks at a time\n so this should be fine (unlikely that we retrieve all 50 chunks contain\n 4 unique objects).\n If we decide this isn't acceptable we can use multiple queries but they\n should be in parallel so query time doesn't get too long.\n \"\"\"\n truncated_record_ids = record_ids[:_MAX_RECORD_IDS_PER_QUERY]\n record_ids_str = \"'\" + \"','\".join(truncated_record_ids) + \"'\"\n access_query = f\"\"\"\n SELECT RecordId, HasReadAccess\n FROM UserRecordAccess\n WHERE RecordId IN ({record_ids_str})\n AND UserId = '{user_id}'\n \"\"\"\n result = salesforce_client.query_all(access_query)\n return {record[\"RecordId\"]: record[\"HasReadAccess\"] for record in result[\"records\"]}\n\n\n_CC_PAIR_ID_SALESFORCE_CLIENT_MAP: dict[int, Salesforce] = {}\n_DOC_ID_TO_CC_PAIR_ID_MAP: dict[str, int] = {}\n\n\n# NOTE: This is not used anywhere.\ndef _get_salesforce_client_for_doc_id(db_session: Session, doc_id: str) -> Salesforce:\n \"\"\"\n Uses a document id to get the cc_pair that indexed that document and uses the credentials\n for that cc_pair to create a Salesforce client.\n Problems:\n - There may be multiple cc_pairs for a document, and we don't know which one to use.\n - right now we just use the first one\n - Building a new Salesforce client for each document is slow.\n - Memory usage could be an issue as we build these dictionaries.\n \"\"\"\n if doc_id not in _DOC_ID_TO_CC_PAIR_ID_MAP:\n cc_pairs = get_cc_pairs_for_document(db_session, doc_id)\n first_cc_pair = cc_pairs[0]\n _DOC_ID_TO_CC_PAIR_ID_MAP[doc_id] = first_cc_pair.id\n\n cc_pair_id = _DOC_ID_TO_CC_PAIR_ID_MAP[doc_id]\n if cc_pair_id not in _CC_PAIR_ID_SALESFORCE_CLIENT_MAP:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if cc_pair is None:\n raise ValueError(f\"CC pair {cc_pair_id} not found\")\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n _CC_PAIR_ID_SALESFORCE_CLIENT_MAP[cc_pair_id] = Salesforce(\n username=credential_json[\"sf_username\"],\n password=credential_json[\"sf_password\"],\n security_token=credential_json[\"sf_security_token\"],\n )\n\n return _CC_PAIR_ID_SALESFORCE_CLIENT_MAP[cc_pair_id]\n" }, { "path": "backend/ee/onyx/external_permissions/sharepoint/doc_sync.py", "content": "from collections.abc import Generator\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom ee.onyx.external_permissions.utils import generic_doc_sync\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.sharepoint.connector import SharepointConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nSHAREPOINT_DOC_SYNC_TAG = \"sharepoint_doc_sync\"\n\n\ndef sharepoint_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,\n callback: IndexingHeartbeatInterface | None = None,\n) -> Generator[ElementExternalAccess, None, None]:\n sharepoint_connector = SharepointConnector(\n **cc_pair.connector.connector_specific_config,\n )\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n sharepoint_connector.load_credentials(credential_json)\n\n yield from generic_doc_sync(\n cc_pair=cc_pair,\n fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,\n callback=callback,\n doc_source=DocumentSource.SHAREPOINT,\n slim_connector=sharepoint_connector,\n label=SHAREPOINT_DOC_SYNC_TAG,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/sharepoint/group_sync.py", "content": "from collections.abc import Generator\n\nfrom office365.sharepoint.client_context import ClientContext # type: ignore[import-untyped]\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom ee.onyx.external_permissions.sharepoint.permission_utils import (\n get_sharepoint_external_groups,\n)\nfrom onyx.configs.app_configs import SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION\nfrom onyx.connectors.sharepoint.connector import acquire_token_for_rest\nfrom onyx.connectors.sharepoint.connector import SharepointConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef sharepoint_group_sync(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair,\n) -> Generator[ExternalUserGroup, None, None]:\n \"\"\"Sync SharePoint groups and their members\"\"\"\n\n # Get site URLs from connector config\n connector_config = cc_pair.connector.connector_specific_config\n\n # Create SharePoint connector instance and load credentials\n connector = SharepointConnector(**connector_config)\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n connector.load_credentials(credential_json)\n\n if not connector.msal_app:\n raise RuntimeError(\"MSAL app not initialized in connector\")\n\n if not connector.sp_tenant_domain:\n raise RuntimeError(\"Tenant domain not initialized in connector\")\n\n # Get site descriptors from connector (either configured sites or all sites)\n site_descriptors = connector.site_descriptors or connector.fetch_sites()\n\n if not site_descriptors:\n raise RuntimeError(\"No SharePoint sites found for group sync\")\n\n logger.info(f\"Processing {len(site_descriptors)} sites for group sync\")\n\n enumerate_all = connector_config.get(\n \"exhaustive_ad_enumeration\", SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION\n )\n\n msal_app = connector.msal_app\n sp_tenant_domain = connector.sp_tenant_domain\n sp_domain_suffix = connector.sharepoint_domain_suffix\n for site_descriptor in site_descriptors:\n logger.debug(f\"Processing site: {site_descriptor.url}\")\n\n ctx = ClientContext(site_descriptor.url).with_access_token(\n lambda: acquire_token_for_rest(msal_app, sp_tenant_domain, sp_domain_suffix)\n )\n\n external_groups = get_sharepoint_external_groups(\n ctx,\n connector.graph_client,\n graph_api_base=connector.graph_api_base,\n get_access_token=connector._get_graph_access_token,\n enumerate_all_ad_groups=enumerate_all,\n )\n\n # Yield each group\n for group in external_groups:\n logger.debug(\n f\"Found group: {group.id} with {len(group.user_emails)} members\"\n )\n yield group\n" }, { "path": "backend/ee/onyx/external_permissions/sharepoint/permission_utils.py", "content": "import re\nimport time\nfrom collections import deque\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom typing import Any\nfrom urllib.parse import urlparse\n\nimport requests as _requests\nfrom office365.graph_client import GraphClient # type: ignore[import-untyped]\nfrom office365.onedrive.driveitems.driveItem import DriveItem # type: ignore[import-untyped]\nfrom office365.runtime.client_request import ClientRequestException # type: ignore\nfrom office365.sharepoint.client_context import ClientContext # type: ignore[import-untyped]\nfrom office365.sharepoint.permissions.securable_object import RoleAssignmentCollection # type: ignore[import-untyped]\nfrom pydantic import BaseModel\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.sharepoint.connector import GRAPH_API_MAX_RETRIES\nfrom onyx.connectors.sharepoint.connector import GRAPH_API_RETRYABLE_STATUSES\nfrom onyx.connectors.sharepoint.connector import SHARED_DOCUMENTS_MAP_REVERSE\nfrom onyx.connectors.sharepoint.connector import sleep_and_retry\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# These values represent different types of SharePoint principals used in permission assignments\nUSER_PRINCIPAL_TYPE = 1 # Individual user accounts\nANONYMOUS_USER_PRINCIPAL_TYPE = 3 # Anonymous/unauthenticated users (public access)\nAZURE_AD_GROUP_PRINCIPAL_TYPE = 4 # Azure Active Directory security groups\nSHAREPOINT_GROUP_PRINCIPAL_TYPE = 8 # SharePoint site groups (local to the site)\nMICROSOFT_DOMAIN = \".onmicrosoft\"\n# Limited Access role type, limited access is a travel through permission not a actual permission\nLIMITED_ACCESS_ROLE_TYPES = [1, 9]\nLIMITED_ACCESS_ROLE_NAMES = [\"Limited Access\", \"Web-Only Limited Access\"]\n\n\nAD_GROUP_ENUMERATION_THRESHOLD = 100_000\n\n\ndef _graph_api_get(\n url: str,\n get_access_token: Callable[[], str],\n params: dict[str, str] | None = None,\n) -> dict[str, Any]:\n \"\"\"Authenticated Graph API GET with retry on transient errors.\"\"\"\n for attempt in range(GRAPH_API_MAX_RETRIES + 1):\n access_token = get_access_token()\n headers = {\"Authorization\": f\"Bearer {access_token}\"}\n try:\n resp = _requests.get(\n url, headers=headers, params=params, timeout=REQUEST_TIMEOUT_SECONDS\n )\n if (\n resp.status_code in GRAPH_API_RETRYABLE_STATUSES\n and attempt < GRAPH_API_MAX_RETRIES\n ):\n wait = min(int(resp.headers.get(\"Retry-After\", str(2**attempt))), 60)\n logger.warning(\n f\"Graph API {resp.status_code} on attempt {attempt + 1}, retrying in {wait}s: {url}\"\n )\n time.sleep(wait)\n continue\n resp.raise_for_status()\n return resp.json()\n except (_requests.ConnectionError, _requests.Timeout, _requests.HTTPError):\n if attempt < GRAPH_API_MAX_RETRIES:\n wait = min(2**attempt, 60)\n logger.warning(\n f\"Graph API connection error on attempt {attempt + 1}, retrying in {wait}s: {url}\"\n )\n time.sleep(wait)\n continue\n raise\n raise RuntimeError(\n f\"Graph API request failed after {GRAPH_API_MAX_RETRIES + 1} attempts: {url}\"\n )\n\n\ndef _iter_graph_collection(\n initial_url: str,\n get_access_token: Callable[[], str],\n params: dict[str, str] | None = None,\n) -> Generator[dict[str, Any], None, None]:\n \"\"\"Paginate through a Graph API collection, yielding items one at a time.\"\"\"\n url: str | None = initial_url\n while url:\n data = _graph_api_get(url, get_access_token, params)\n params = None\n yield from data.get(\"value\", [])\n url = data.get(\"@odata.nextLink\")\n\n\ndef _normalize_email(email: str) -> str:\n if MICROSOFT_DOMAIN in email:\n return email.replace(MICROSOFT_DOMAIN, \"\")\n return email\n\n\nclass SharepointGroup(BaseModel):\n model_config = {\"frozen\": True}\n\n name: str\n login_name: str\n principal_type: int\n\n\nclass GroupsResult(BaseModel):\n groups_to_emails: dict[str, set[str]]\n found_public_group: bool\n\n\ndef _get_azuread_group_guid_by_name(\n graph_client: GraphClient, group_name: str\n) -> str | None:\n try:\n # Search for groups by display name\n groups = sleep_and_retry(\n graph_client.groups.filter(f\"displayName eq '{group_name}'\").get(),\n \"get_azuread_group_guid_by_name\",\n )\n\n if groups and len(groups) > 0:\n return groups[0].id\n\n return None\n\n except Exception as e:\n logger.error(f\"Failed to get Azure AD group GUID for name {group_name}: {e}\")\n return None\n\n\ndef _extract_guid_from_claims_token(claims_token: str) -> str | None:\n\n try:\n # Pattern to match GUID in claims token\n # Claims tokens often have format: c:0o.c|provider|GUID_suffix\n guid_pattern = r\"([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\"\n\n match = re.search(guid_pattern, claims_token, re.IGNORECASE)\n if match:\n return match.group(1)\n\n return None\n\n except Exception as e:\n logger.error(f\"Failed to extract GUID from claims token {claims_token}: {e}\")\n return None\n\n\ndef _get_group_guid_from_identifier(\n graph_client: GraphClient, identifier: str\n) -> str | None:\n try:\n # Check if it's already a GUID\n guid_pattern = r\"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$\"\n if re.match(guid_pattern, identifier, re.IGNORECASE):\n return identifier\n\n # Check if it's a SharePoint claims token\n if identifier.startswith(\"c:0\") and \"|\" in identifier:\n guid = _extract_guid_from_claims_token(identifier)\n if guid:\n logger.info(f\"Extracted GUID {guid} from claims token {identifier}\")\n return guid\n\n # Try to search by display name as fallback\n return _get_azuread_group_guid_by_name(graph_client, identifier)\n\n except Exception as e:\n logger.error(f\"Failed to get group GUID from identifier {identifier}: {e}\")\n return None\n\n\ndef _get_security_group_owners(graph_client: GraphClient, group_id: str) -> list[str]:\n try:\n # Get group owners using Graph API\n group = graph_client.groups[group_id]\n owners = sleep_and_retry(\n group.owners.get_all(page_loaded=lambda _: None),\n \"get_security_group_owners\",\n )\n\n owner_emails: list[str] = []\n logger.info(f\"Owners: {owners}\")\n\n for owner in owners:\n owner_data = owner.to_json()\n\n # Extract email from the JSON data\n mail: str | None = owner_data.get(\"mail\")\n user_principal_name: str | None = owner_data.get(\"userPrincipalName\")\n\n # Check if owner is a user and has an email\n if mail:\n if MICROSOFT_DOMAIN in mail:\n mail = mail.replace(MICROSOFT_DOMAIN, \"\")\n owner_emails.append(mail)\n elif user_principal_name:\n if MICROSOFT_DOMAIN in user_principal_name:\n user_principal_name = user_principal_name.replace(\n MICROSOFT_DOMAIN, \"\"\n )\n owner_emails.append(user_principal_name)\n\n logger.info(\n f\"Retrieved {len(owner_emails)} owners from security group {group_id}\"\n )\n return owner_emails\n\n except Exception as e:\n logger.error(f\"Failed to get security group owners for group {group_id}: {e}\")\n return []\n\n\ndef _get_sharepoint_list_item_id(drive_item: DriveItem) -> str | None:\n\n try:\n # First try to get the list item directly from the drive item\n if hasattr(drive_item, \"listItem\"):\n list_item = drive_item.listItem\n if list_item:\n # Load the list item properties to get the ID\n sleep_and_retry(list_item.get(), \"get_sharepoint_list_item_id\")\n if hasattr(list_item, \"id\") and list_item.id:\n return str(list_item.id)\n\n # The SharePoint list item ID is typically available in the sharepointIds property\n sharepoint_ids = getattr(drive_item, \"sharepoint_ids\", None)\n if sharepoint_ids and hasattr(sharepoint_ids, \"listItemId\"):\n return sharepoint_ids.listItemId\n\n # Alternative: try to get it from the properties\n properties = getattr(drive_item, \"properties\", None)\n if properties:\n # Sometimes the SharePoint list item ID is in the properties\n for prop_name, prop_value in properties.items():\n if \"listitemid\" in prop_name.lower():\n return str(prop_value)\n\n return None\n except Exception as e:\n logger.error(\n f\"Error getting SharePoint list item ID for item {drive_item.id}: {e}\"\n )\n raise e\n\n\ndef _is_public_item(\n drive_item: DriveItem,\n treat_sharing_link_as_public: bool = False,\n) -> bool:\n if not treat_sharing_link_as_public:\n return False\n\n try:\n permissions = sleep_and_retry(\n drive_item.permissions.get_all(page_loaded=lambda _: None), \"is_public_item\"\n )\n for permission in permissions:\n if permission.link and permission.link.scope in (\n \"anonymous\",\n \"organization\",\n ):\n return True\n return False\n except Exception as e:\n logger.error(f\"Failed to check if item {drive_item.id} is public: {e}\")\n return False\n\n\ndef _is_public_login_name(login_name: str) -> bool:\n # Patterns that indicate public access\n # This list is derived from the below link\n # https://learn.microsoft.com/en-us/answers/questions/2085339/guid-in-the-loginname-of-site-user-everyone-except\n public_login_patterns: list[str] = [\n \"c:0-.f|rolemanager|spo-grid-all-users/\",\n \"c:0(.s|true\",\n ]\n for pattern in public_login_patterns:\n if pattern in login_name:\n logger.info(f\"Login name {login_name} is public\")\n return True\n return False\n\n\n# AD groups allows same display name for multiple groups, so we need to add the GUID to the name\ndef _get_group_name_with_suffix(\n login_name: str, group_name: str, graph_client: GraphClient\n) -> str:\n ad_group_suffix = _get_group_guid_from_identifier(graph_client, login_name)\n return f\"{group_name}_{ad_group_suffix}\"\n\n\ndef _get_sharepoint_groups(\n client_context: ClientContext, group_name: str, graph_client: GraphClient\n) -> tuple[set[SharepointGroup], set[str]]:\n\n groups: set[SharepointGroup] = set()\n user_emails: set[str] = set()\n\n def process_users(users: list[Any]) -> None:\n nonlocal groups, user_emails\n\n for user in users:\n logger.debug(f\"User: {user.to_json()}\")\n if user.principal_type == USER_PRINCIPAL_TYPE and hasattr(\n user, \"user_principal_name\"\n ):\n if user.user_principal_name:\n email = user.user_principal_name\n if MICROSOFT_DOMAIN in email:\n email = email.replace(MICROSOFT_DOMAIN, \"\")\n user_emails.add(email)\n else:\n logger.warning(\n f\"User don't have a user principal name: {user.login_name}\"\n )\n elif user.principal_type in [\n AZURE_AD_GROUP_PRINCIPAL_TYPE,\n SHAREPOINT_GROUP_PRINCIPAL_TYPE,\n ]:\n name = user.title\n if user.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:\n name = _get_group_name_with_suffix(\n user.login_name, name, graph_client\n )\n groups.add(\n SharepointGroup(\n login_name=user.login_name,\n principal_type=user.principal_type,\n name=name,\n )\n )\n\n group = client_context.web.site_groups.get_by_name(group_name)\n sleep_and_retry(\n group.users.get_all(page_loaded=process_users), \"get_sharepoint_groups\"\n )\n\n return groups, user_emails\n\n\ndef _get_azuread_groups(\n graph_client: GraphClient, group_name: str\n) -> tuple[set[SharepointGroup], set[str]]:\n\n group_id = _get_group_guid_from_identifier(graph_client, group_name)\n if not group_id:\n logger.error(f\"Failed to get Azure AD group GUID for name {group_name}\")\n return set(), set()\n group = graph_client.groups[group_id]\n groups: set[SharepointGroup] = set()\n user_emails: set[str] = set()\n\n def process_members(members: list[Any]) -> None:\n nonlocal groups, user_emails\n\n for member in members:\n member_data = member.to_json()\n logger.debug(f\"Member: {member_data}\")\n # Check for user-specific attributes\n user_principal_name = member_data.get(\"userPrincipalName\")\n mail = member_data.get(\"mail\")\n display_name = member_data.get(\"displayName\") or member_data.get(\n \"display_name\"\n )\n\n # Check object attributes directly (if available)\n is_user = False\n is_group = False\n\n # Users typically have userPrincipalName or mail\n if user_principal_name or (mail and \"@\" in str(mail)):\n is_user = True\n # Groups typically have displayName but no userPrincipalName\n elif display_name and not user_principal_name:\n # Additional check: try to access group-specific properties\n if (\n hasattr(member, \"groupTypes\")\n or member_data.get(\"groupTypes\") is not None\n ):\n is_group = True\n # Or check if it has an 'id' field typical for groups\n elif member_data.get(\"id\") and not user_principal_name:\n is_group = True\n\n # Check the object type name (fallback)\n if not is_user and not is_group:\n obj_type = type(member).__name__.lower()\n if \"user\" in obj_type:\n is_user = True\n elif \"group\" in obj_type:\n is_group = True\n\n # Process based on identification\n if is_user:\n if user_principal_name:\n email = user_principal_name\n if MICROSOFT_DOMAIN in email:\n email = email.replace(MICROSOFT_DOMAIN, \"\")\n user_emails.add(email)\n elif mail:\n email = mail\n if MICROSOFT_DOMAIN in email:\n email = email.replace(MICROSOFT_DOMAIN, \"\")\n user_emails.add(email)\n logger.info(f\"Added user: {user_principal_name or mail}\")\n elif is_group:\n if not display_name:\n logger.error(f\"No display name for group: {member_data.get('id')}\")\n continue\n name = _get_group_name_with_suffix(\n member_data.get(\"id\", \"\"), display_name, graph_client\n )\n groups.add(\n SharepointGroup(\n login_name=member_data.get(\"id\", \"\"), # Use ID for groups\n principal_type=AZURE_AD_GROUP_PRINCIPAL_TYPE,\n name=name,\n )\n )\n logger.info(f\"Added group: {name}\")\n else:\n # Log unidentified members for debugging\n logger.warning(f\"Could not identify member type for: {member_data}\")\n\n sleep_and_retry(\n group.members.get_all(page_loaded=process_members), \"get_azuread_groups\"\n )\n\n owner_emails = _get_security_group_owners(graph_client, group_id)\n user_emails.update(owner_emails)\n\n return groups, user_emails\n\n\ndef _get_groups_and_members_recursively(\n client_context: ClientContext,\n graph_client: GraphClient,\n groups: set[SharepointGroup],\n is_group_sync: bool = False,\n) -> GroupsResult:\n \"\"\"\n Get all groups and their members recursively.\n \"\"\"\n group_queue: deque[SharepointGroup] = deque(groups)\n visited_groups: set[str] = set()\n visited_group_name_to_emails: dict[str, set[str]] = {}\n found_public_group = False\n while group_queue:\n group = group_queue.popleft()\n if group.login_name in visited_groups:\n continue\n visited_groups.add(group.login_name)\n visited_group_name_to_emails[group.name] = set()\n logger.info(\n f\"Processing group: {group.name} principal type: {group.principal_type}\"\n )\n if group.principal_type == SHAREPOINT_GROUP_PRINCIPAL_TYPE:\n group_info, user_emails = _get_sharepoint_groups(\n client_context, group.login_name, graph_client\n )\n visited_group_name_to_emails[group.name].update(user_emails)\n if group_info:\n group_queue.extend(group_info)\n if group.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:\n try:\n # if the site is public, we have default groups assigned to it, so we return early\n if _is_public_login_name(group.login_name):\n found_public_group = True\n if not is_group_sync:\n return GroupsResult(\n groups_to_emails={}, found_public_group=True\n )\n else:\n # we don't want to sync public groups, so we skip them\n continue\n group_info, user_emails = _get_azuread_groups(\n graph_client, group.login_name\n )\n visited_group_name_to_emails[group.name].update(user_emails)\n if group_info:\n group_queue.extend(group_info)\n except ClientRequestException as e:\n # If the group is not found, we skip it. There is a chance that group is still referenced\n # in sharepoint but it is removed from Azure AD. There is no actual documentation on this, but based on\n # our testing we have seen this happen.\n if e.response is not None and e.response.status_code == 404:\n logger.warning(f\"Group {group.login_name} not found\")\n continue\n raise e\n\n return GroupsResult(\n groups_to_emails=visited_group_name_to_emails,\n found_public_group=found_public_group,\n )\n\n\ndef get_external_access_from_sharepoint(\n client_context: ClientContext,\n graph_client: GraphClient,\n drive_name: str | None,\n drive_item: DriveItem | None,\n site_page: dict[str, Any] | None,\n add_prefix: bool = False,\n treat_sharing_link_as_public: bool = False,\n) -> ExternalAccess:\n \"\"\"\n Get external access information from SharePoint.\n \"\"\"\n groups: set[SharepointGroup] = set()\n user_emails: set[str] = set()\n group_ids: set[str] = set()\n\n # Add all members to a processing set first\n def add_user_and_group_to_sets(\n role_assignments: RoleAssignmentCollection,\n ) -> None:\n nonlocal user_emails, groups\n for assignment in role_assignments:\n logger.debug(f\"Assignment: {assignment.to_json()}\")\n if assignment.role_definition_bindings:\n is_limited_access = True\n for role_definition_binding in assignment.role_definition_bindings:\n if (\n role_definition_binding.role_type_kind\n not in LIMITED_ACCESS_ROLE_TYPES\n or role_definition_binding.name not in LIMITED_ACCESS_ROLE_NAMES\n ):\n is_limited_access = False\n break\n\n # Skip if the role is only Limited Access, because this is not a actual permission its a travel through permission\n if is_limited_access:\n logger.info(\n \"Skipping assignment because it has only Limited Access role\"\n )\n continue\n if assignment.member:\n member = assignment.member\n if member.principal_type == USER_PRINCIPAL_TYPE and hasattr(\n member, \"user_principal_name\"\n ):\n email = member.user_principal_name\n if MICROSOFT_DOMAIN in email:\n email = email.replace(MICROSOFT_DOMAIN, \"\")\n user_emails.add(email)\n elif member.principal_type in [\n AZURE_AD_GROUP_PRINCIPAL_TYPE,\n SHAREPOINT_GROUP_PRINCIPAL_TYPE,\n ]:\n name = member.title\n if member.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:\n name = _get_group_name_with_suffix(\n member.login_name, name, graph_client\n )\n groups.add(\n SharepointGroup(\n login_name=member.login_name,\n principal_type=member.principal_type,\n name=name,\n )\n )\n\n if drive_item and drive_name:\n is_public = _is_public_item(drive_item, treat_sharing_link_as_public)\n if is_public:\n logger.info(f\"Item {drive_item.id} is public\")\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n )\n\n item_id = _get_sharepoint_list_item_id(drive_item)\n\n if not item_id:\n raise RuntimeError(\n f\"Failed to get SharePoint list item ID for item {drive_item.id}\"\n )\n\n if drive_name in SHARED_DOCUMENTS_MAP_REVERSE:\n drive_name = SHARED_DOCUMENTS_MAP_REVERSE[drive_name]\n\n item = client_context.web.lists.get_by_title(drive_name).items.get_by_id(\n item_id\n )\n\n sleep_and_retry(\n item.role_assignments.expand([\"Member\", \"RoleDefinitionBindings\"]).get_all(\n page_loaded=add_user_and_group_to_sets,\n ),\n \"get_external_access_from_sharepoint\",\n )\n elif site_page:\n site_url = site_page.get(\"webUrl\")\n # Keep percent-encoding intact so the path matches the encoding\n # used by the Office365 library's SPResPath.create_relative(),\n # which compares against urlparse(context.base_url).path.\n # Decoding (e.g. %27 → ') causes a mismatch that duplicates\n # the site prefix in the constructed URL.\n server_relative_url = urlparse(site_url).path\n file_obj = client_context.web.get_file_by_server_relative_url(\n server_relative_url\n )\n item = file_obj.listItemAllFields\n\n sleep_and_retry(\n item.role_assignments.expand([\"Member\", \"RoleDefinitionBindings\"]).get_all(\n page_loaded=add_user_and_group_to_sets,\n ),\n \"get_external_access_from_sharepoint\",\n )\n else:\n raise RuntimeError(\"No drive item or site page provided\")\n\n groups_and_members: GroupsResult = _get_groups_and_members_recursively(\n client_context, graph_client, groups\n )\n\n # If the site is public, w have default groups assigned to it, so we return early\n if groups_and_members.found_public_group:\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n )\n\n for group_name, _ in groups_and_members.groups_to_emails.items():\n if add_prefix:\n group_name = build_ext_group_name_for_onyx(\n group_name, DocumentSource.SHAREPOINT\n )\n group_ids.add(group_name.lower())\n\n logger.info(f\"User emails: {len(user_emails)}\")\n logger.info(f\"Group IDs: {len(group_ids)}\")\n\n return ExternalAccess(\n external_user_emails=user_emails,\n external_user_group_ids=group_ids,\n is_public=False,\n )\n\n\ndef _enumerate_ad_groups_paginated(\n get_access_token: Callable[[], str],\n already_resolved: set[str],\n graph_api_base: str,\n) -> Generator[ExternalUserGroup, None, None]:\n \"\"\"Paginate through all Azure AD groups and yield ExternalUserGroup for each.\n\n Skips groups whose suffixed name is already in *already_resolved*.\n Stops early if the number of groups exceeds AD_GROUP_ENUMERATION_THRESHOLD.\n \"\"\"\n groups_url = f\"{graph_api_base}/groups\"\n groups_params: dict[str, str] = {\"$select\": \"id,displayName\", \"$top\": \"999\"}\n total_groups = 0\n\n for group_json in _iter_graph_collection(\n groups_url, get_access_token, groups_params\n ):\n group_id: str = group_json.get(\"id\", \"\")\n display_name: str = group_json.get(\"displayName\", \"\")\n if not group_id or not display_name:\n continue\n\n total_groups += 1\n if total_groups > AD_GROUP_ENUMERATION_THRESHOLD:\n logger.warning(\n f\"Azure AD group enumeration exceeded {AD_GROUP_ENUMERATION_THRESHOLD} \"\n \"groups — stopping to avoid excessive memory/API usage. \"\n \"Remaining groups will be resolved from role assignments only.\"\n )\n return\n\n name = f\"{display_name}_{group_id}\"\n if name in already_resolved:\n continue\n\n member_emails: list[str] = []\n members_url = f\"{graph_api_base}/groups/{group_id}/members\"\n members_params: dict[str, str] = {\n \"$select\": \"userPrincipalName,mail\",\n \"$top\": \"999\",\n }\n for member_json in _iter_graph_collection(\n members_url, get_access_token, members_params\n ):\n email = member_json.get(\"userPrincipalName\") or member_json.get(\"mail\")\n if email:\n member_emails.append(_normalize_email(email))\n\n yield ExternalUserGroup(id=name, user_emails=member_emails)\n\n logger.info(f\"Enumerated {total_groups} Azure AD groups via paginated Graph API\")\n\n\ndef get_sharepoint_external_groups(\n client_context: ClientContext,\n graph_client: GraphClient,\n graph_api_base: str,\n get_access_token: Callable[[], str] | None = None,\n enumerate_all_ad_groups: bool = False,\n) -> list[ExternalUserGroup]:\n\n groups: set[SharepointGroup] = set()\n\n def add_group_to_sets(role_assignments: RoleAssignmentCollection) -> None:\n nonlocal groups\n for assignment in role_assignments:\n if assignment.role_definition_bindings:\n is_limited_access = True\n for role_definition_binding in assignment.role_definition_bindings:\n if (\n role_definition_binding.role_type_kind\n not in LIMITED_ACCESS_ROLE_TYPES\n or role_definition_binding.name not in LIMITED_ACCESS_ROLE_NAMES\n ):\n is_limited_access = False\n break\n\n # Skip if the role assignment is only Limited Access, because this is not a actual permission its\n # a travel through permission\n if is_limited_access:\n logger.info(\n \"Skipping assignment because it has only Limited Access role\"\n )\n continue\n if assignment.member:\n member = assignment.member\n if member.principal_type in [\n AZURE_AD_GROUP_PRINCIPAL_TYPE,\n SHAREPOINT_GROUP_PRINCIPAL_TYPE,\n ]:\n name = member.title\n if member.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:\n name = _get_group_name_with_suffix(\n member.login_name, name, graph_client\n )\n\n groups.add(\n SharepointGroup(\n login_name=member.login_name,\n principal_type=member.principal_type,\n name=name,\n )\n )\n\n sleep_and_retry(\n client_context.web.role_assignments.expand(\n [\"Member\", \"RoleDefinitionBindings\"]\n ).get_all(page_loaded=add_group_to_sets),\n \"get_sharepoint_external_groups\",\n )\n groups_and_members: GroupsResult = _get_groups_and_members_recursively(\n client_context, graph_client, groups, is_group_sync=True\n )\n\n external_user_groups: list[ExternalUserGroup] = [\n ExternalUserGroup(id=group_name, user_emails=list(emails))\n for group_name, emails in groups_and_members.groups_to_emails.items()\n ]\n\n if not enumerate_all_ad_groups or get_access_token is None:\n logger.info(\n \"Skipping exhaustive Azure AD group enumeration. Only groups found in site role assignments are included.\"\n )\n return external_user_groups\n\n already_resolved = set(groups_and_members.groups_to_emails.keys())\n for group in _enumerate_ad_groups_paginated(\n get_access_token, already_resolved, graph_api_base\n ):\n external_user_groups.append(group)\n\n return external_user_groups\n" }, { "path": "backend/ee/onyx/external_permissions/slack/channel_access.py", "content": "from slack_sdk import WebClient\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.slack.connector import ChannelType\nfrom onyx.connectors.slack.utils import expert_info_from_slack_id\nfrom onyx.connectors.slack.utils import make_paginated_slack_api_call\n\n\ndef get_channel_access(\n client: WebClient,\n channel: ChannelType,\n user_cache: dict[str, BasicExpertInfo | None],\n) -> ExternalAccess:\n \"\"\"\n Get channel access permissions for a Slack channel.\n\n Args:\n client: Slack WebClient instance\n channel: Slack channel object containing channel info\n user_cache: Cache of user IDs to BasicExpertInfo objects. May be updated in place.\n\n Returns:\n ExternalAccess object for the channel.\n \"\"\"\n channel_is_public = not channel[\"is_private\"]\n if channel_is_public:\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n )\n\n channel_id = channel[\"id\"]\n\n # Get all member IDs for the channel\n member_ids = []\n for result in make_paginated_slack_api_call(\n client.conversations_members,\n channel=channel_id,\n ):\n member_ids.extend(result.get(\"members\", []))\n\n member_emails = set()\n for member_id in member_ids:\n # Try to get user info from cache or fetch it\n user_info = expert_info_from_slack_id(\n user_id=member_id,\n client=client,\n user_cache=user_cache,\n )\n\n # If we have user info and an email, add it to the set\n if user_info and user_info.email:\n member_emails.add(user_info.email)\n\n return ExternalAccess(\n external_user_emails=member_emails,\n # NOTE: groups are not used, since adding a group to a channel just adds all\n # users that are in the group.\n external_user_group_ids=set(),\n is_public=False,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/slack/doc_sync.py", "content": "from collections.abc import Generator\n\nfrom slack_sdk import WebClient\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom ee.onyx.external_permissions.slack.utils import fetch_user_id_to_email_map\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.credentials_provider import OnyxDBCredentialsProvider\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.slack.connector import get_channels\nfrom onyx.connectors.slack.connector import make_paginated_slack_api_call\nfrom onyx.connectors.slack.connector import SlackConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\nlogger = setup_logger()\n\n\ndef _fetch_workspace_permissions(\n user_id_to_email_map: dict[str, str],\n) -> ExternalAccess:\n user_emails = set()\n for email in user_id_to_email_map.values():\n user_emails.add(email)\n return ExternalAccess(\n external_user_emails=user_emails,\n # No group<->document mapping for slack\n external_user_group_ids=set(),\n # No way to determine if slack is invite only without enterprise license\n is_public=False,\n )\n\n\ndef _fetch_channel_permissions(\n slack_client: WebClient,\n workspace_permissions: ExternalAccess,\n user_id_to_email_map: dict[str, str],\n) -> dict[str, ExternalAccess]:\n channel_permissions = {}\n public_channels = get_channels(\n client=slack_client,\n get_public=True,\n get_private=False,\n )\n public_channel_ids = [\n channel[\"id\"] for channel in public_channels if \"id\" in channel\n ]\n for channel_id in public_channel_ids:\n channel_permissions[channel_id] = workspace_permissions\n\n private_channels = get_channels(\n client=slack_client,\n get_public=False,\n get_private=True,\n )\n private_channel_ids = [\n channel[\"id\"] for channel in private_channels if \"id\" in channel\n ]\n\n for channel_id in private_channel_ids:\n # Collect all member ids for the channel pagination calls\n member_ids = []\n for result in make_paginated_slack_api_call(\n slack_client.conversations_members,\n channel=channel_id,\n ):\n member_ids.extend(result.get(\"members\", []))\n\n # Collect all member emails for the channel\n member_emails = set()\n for member_id in member_ids:\n member_email = user_id_to_email_map.get(member_id)\n\n if not member_email:\n # If the user is an external user, they wont get returned from the\n # conversations_members call so we need to make a separate call to users_info\n # and add them to the user_id_to_email_map\n member_info = slack_client.users_info(user=member_id)\n member_email = member_info[\"user\"][\"profile\"].get(\"email\")\n if not member_email:\n # If no email is found, we skip the user\n continue\n user_id_to_email_map[member_id] = member_email\n\n member_emails.add(member_email)\n\n channel_permissions[channel_id] = ExternalAccess(\n external_user_emails=member_emails,\n # No group<->document mapping for slack\n external_user_group_ids=set(),\n # No way to determine if slack is invite only without enterprise license\n is_public=False,\n )\n\n return channel_permissions\n\n\ndef _get_slack_document_access(\n slack_connector: SlackConnector,\n channel_permissions: dict[str, ExternalAccess], # noqa: ARG001\n callback: IndexingHeartbeatInterface | None,\n indexing_start: SecondsSinceUnixEpoch | None = None,\n) -> Generator[DocExternalAccess, None, None]:\n slim_doc_generator = slack_connector.retrieve_all_slim_docs_perm_sync(\n callback=callback,\n start=indexing_start,\n )\n\n for doc_metadata_batch in slim_doc_generator:\n for doc_metadata in doc_metadata_batch:\n if isinstance(doc_metadata, HierarchyNode):\n # TODO: handle hierarchynodes during sync\n continue\n if doc_metadata.external_access is None:\n raise ValueError(\n f\"No external access for document {doc_metadata.id}. \"\n \"Please check to make sure that your Slack bot token has the \"\n \"`channels:read` scope\"\n )\n\n yield DocExternalAccess(\n external_access=doc_metadata.external_access,\n doc_id=doc_metadata.id,\n )\n\n if callback:\n if callback.should_stop():\n raise RuntimeError(\"_get_slack_document_access: Stop signal detected\")\n\n callback.progress(\"_get_slack_document_access\", 1)\n\n\ndef slack_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction, # noqa: ARG001\n callback: IndexingHeartbeatInterface | None,\n) -> Generator[DocExternalAccess, None, None]:\n \"\"\"\n Adds the external permissions to the documents in postgres\n if the document doesn't already exists in postgres, we create\n it in postgres so that when it gets created later, the permissions are\n already populated\n \"\"\"\n # Use credentials provider instead of directly loading credentials\n\n tenant_id = get_current_tenant_id()\n provider = OnyxDBCredentialsProvider(tenant_id, \"slack\", cc_pair.credential.id)\n r = get_redis_client(tenant_id=tenant_id)\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n slack_client = SlackConnector.make_slack_web_client(\n provider.get_provider_key(),\n credential_json[\"slack_bot_token\"],\n SlackConnector.MAX_RETRIES,\n r,\n )\n\n user_id_to_email_map = fetch_user_id_to_email_map(slack_client)\n if not user_id_to_email_map:\n raise ValueError(\n \"No user id to email map found. Please check to make sure that your Slack bot token has the `users:read.email` scope\"\n )\n\n workspace_permissions = _fetch_workspace_permissions(\n user_id_to_email_map=user_id_to_email_map,\n )\n channel_permissions = _fetch_channel_permissions(\n slack_client=slack_client,\n workspace_permissions=workspace_permissions,\n user_id_to_email_map=user_id_to_email_map,\n )\n\n slack_connector = SlackConnector(**cc_pair.connector.connector_specific_config)\n slack_connector.set_credentials_provider(provider)\n indexing_start_ts: SecondsSinceUnixEpoch | None = (\n cc_pair.connector.indexing_start.timestamp()\n if cc_pair.connector.indexing_start is not None\n else None\n )\n\n yield from _get_slack_document_access(\n slack_connector=slack_connector,\n channel_permissions=channel_permissions,\n callback=callback,\n indexing_start=indexing_start_ts,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/slack/group_sync.py", "content": "\"\"\"\nTHIS IS NOT USEFUL OR USED FOR PERMISSION SYNCING\nWHEN USERGROUPS ARE ADDED TO A CHANNEL, IT JUST RESOLVES ALL THE USERS TO THAT CHANNEL\nSO WHEN CHECKING IF A USER CAN ACCESS A DOCUMENT, WE ONLY NEED TO CHECK THEIR EMAIL\nTHERE IS NO USERGROUP <-> DOCUMENT PERMISSION MAPPING\n\"\"\"\n\nfrom slack_sdk import WebClient\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom ee.onyx.external_permissions.slack.utils import fetch_user_id_to_email_map\nfrom onyx.connectors.credentials_provider import OnyxDBCredentialsProvider\nfrom onyx.connectors.slack.connector import SlackConnector\nfrom onyx.connectors.slack.utils import make_paginated_slack_api_call\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _get_slack_group_ids(\n slack_client: WebClient,\n) -> list[str]:\n group_ids = []\n for result in make_paginated_slack_api_call(slack_client.usergroups_list):\n for group in result.get(\"usergroups\", []):\n group_ids.append(group.get(\"id\"))\n return group_ids\n\n\ndef _get_slack_group_members_email(\n slack_client: WebClient,\n group_name: str,\n user_id_to_email_map: dict[str, str],\n) -> list[str]:\n group_member_emails = []\n for result in make_paginated_slack_api_call(\n slack_client.usergroups_users_list, usergroup=group_name\n ):\n for member_id in result.get(\"users\", []):\n member_email = user_id_to_email_map.get(member_id)\n if not member_email:\n # If the user is an external user, they wont get returned from the\n # conversations_members call so we need to make a separate call to users_info\n member_info = slack_client.users_info(user=member_id)\n member_email = member_info[\"user\"][\"profile\"].get(\"email\")\n if not member_email:\n # If no email is found, we skip the user\n continue\n user_id_to_email_map[member_id] = member_email\n group_member_emails.append(member_email)\n\n return group_member_emails\n\n\ndef slack_group_sync(\n tenant_id: str,\n cc_pair: ConnectorCredentialPair,\n) -> list[ExternalUserGroup]:\n \"\"\"NOTE: not used atm. All channel access is done at the\n individual user level. Leaving in for now in case we need it later.\"\"\"\n\n provider = OnyxDBCredentialsProvider(tenant_id, \"slack\", cc_pair.credential.id)\n r = get_redis_client(tenant_id=tenant_id)\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n slack_client = SlackConnector.make_slack_web_client(\n provider.get_provider_key(),\n credential_json[\"slack_bot_token\"],\n SlackConnector.MAX_RETRIES,\n r,\n )\n\n user_id_to_email_map = fetch_user_id_to_email_map(slack_client)\n\n onyx_groups: list[ExternalUserGroup] = []\n for group_name in _get_slack_group_ids(slack_client):\n group_member_emails = _get_slack_group_members_email(\n slack_client=slack_client,\n group_name=group_name,\n user_id_to_email_map=user_id_to_email_map,\n )\n if not group_member_emails:\n continue\n onyx_groups.append(\n ExternalUserGroup(id=group_name, user_emails=group_member_emails)\n )\n return onyx_groups\n" }, { "path": "backend/ee/onyx/external_permissions/slack/utils.py", "content": "from slack_sdk import WebClient\n\nfrom onyx.connectors.slack.utils import make_paginated_slack_api_call\n\n\ndef fetch_user_id_to_email_map(\n slack_client: WebClient,\n) -> dict[str, str]:\n user_id_to_email_map = {}\n for user_info in make_paginated_slack_api_call(\n slack_client.users_list,\n ):\n for user in user_info.get(\"members\", []):\n if user.get(\"profile\", {}).get(\"email\"):\n user_id_to_email_map[user.get(\"id\")] = user.get(\"profile\", {}).get(\n \"email\"\n )\n return user_id_to_email_map\n" }, { "path": "backend/ee/onyx/external_permissions/sync_params.py", "content": "from collections.abc import Generator\nfrom typing import Optional\nfrom typing import TYPE_CHECKING\n\nfrom pydantic import BaseModel\n\nfrom ee.onyx.configs.app_configs import CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import GITHUB_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import JIRA_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import JIRA_PERMISSION_GROUP_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import SLACK_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.configs.app_configs import TEAMS_PERMISSION_DOC_SYNC_FREQUENCY\nfrom ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync\nfrom ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync\nfrom ee.onyx.external_permissions.github.doc_sync import github_doc_sync\nfrom ee.onyx.external_permissions.github.group_sync import github_group_sync\nfrom ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync\nfrom ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync\nfrom ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync\nfrom ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync\nfrom ee.onyx.external_permissions.jira.group_sync import jira_group_sync\nfrom ee.onyx.external_permissions.perm_sync_types import CensoringFuncType\nfrom ee.onyx.external_permissions.perm_sync_types import DocSyncFuncType\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import GroupSyncFuncType\nfrom ee.onyx.external_permissions.salesforce.postprocessing import (\n censor_salesforce_chunks,\n)\nfrom ee.onyx.external_permissions.sharepoint.doc_sync import sharepoint_doc_sync\nfrom ee.onyx.external_permissions.sharepoint.group_sync import sharepoint_group_sync\nfrom ee.onyx.external_permissions.slack.doc_sync import slack_doc_sync\nfrom ee.onyx.external_permissions.teams.doc_sync import teams_doc_sync\nfrom onyx.configs.constants import DocumentSource\n\nif TYPE_CHECKING:\n from onyx.access.models import DocExternalAccess # noqa\n from onyx.db.models import ConnectorCredentialPair # noqa\n from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface # noqa\n\n\nclass DocSyncConfig(BaseModel):\n doc_sync_frequency: int\n doc_sync_func: DocSyncFuncType\n initial_index_should_sync: bool\n\n\nclass GroupSyncConfig(BaseModel):\n group_sync_frequency: int\n group_sync_func: GroupSyncFuncType\n group_sync_is_cc_pair_agnostic: bool\n\n\nclass CensoringConfig(BaseModel):\n chunk_censoring_func: CensoringFuncType\n\n\nclass SyncConfig(BaseModel):\n # None means we don't perform a doc_sync\n doc_sync_config: DocSyncConfig | None = None\n # None means we don't perform a group_sync\n group_sync_config: GroupSyncConfig | None = None\n # None means we don't perform a chunk_censoring\n censoring_config: CensoringConfig | None = None\n\n\n# Mock doc sync function for testing (no-op)\ndef mock_doc_sync(\n cc_pair: \"ConnectorCredentialPair\", # noqa: ARG001\n fetch_all_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_docs_ids_fn: FetchAllDocumentsIdsFunction, # noqa: ARG001\n callback: Optional[\"IndexingHeartbeatInterface\"], # noqa: ARG001\n) -> Generator[\"DocExternalAccess\", None, None]:\n \"\"\"Mock doc sync function for testing - returns empty list since permissions are fetched during indexing\"\"\"\n yield from []\n\n\n_SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {\n DocumentSource.GOOGLE_DRIVE: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=gdrive_doc_sync,\n initial_index_should_sync=True,\n ),\n group_sync_config=GroupSyncConfig(\n group_sync_frequency=GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY,\n group_sync_func=gdrive_group_sync,\n group_sync_is_cc_pair_agnostic=False,\n ),\n ),\n DocumentSource.CONFLUENCE: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=confluence_doc_sync,\n initial_index_should_sync=False,\n ),\n group_sync_config=GroupSyncConfig(\n group_sync_frequency=CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY,\n group_sync_func=confluence_group_sync,\n group_sync_is_cc_pair_agnostic=True,\n ),\n ),\n DocumentSource.JIRA: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=JIRA_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=jira_doc_sync,\n initial_index_should_sync=True,\n ),\n group_sync_config=GroupSyncConfig(\n group_sync_frequency=JIRA_PERMISSION_GROUP_SYNC_FREQUENCY,\n group_sync_func=jira_group_sync,\n group_sync_is_cc_pair_agnostic=True,\n ),\n ),\n # Groups are not needed for Slack.\n # All channel access is done at the individual user level.\n DocumentSource.SLACK: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=SLACK_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=slack_doc_sync,\n initial_index_should_sync=True,\n ),\n ),\n DocumentSource.GMAIL: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=gmail_doc_sync,\n initial_index_should_sync=False,\n ),\n ),\n DocumentSource.GITHUB: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=GITHUB_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=github_doc_sync,\n initial_index_should_sync=True,\n ),\n group_sync_config=GroupSyncConfig(\n group_sync_frequency=GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY,\n group_sync_func=github_group_sync,\n group_sync_is_cc_pair_agnostic=False,\n ),\n ),\n DocumentSource.SALESFORCE: SyncConfig(\n censoring_config=CensoringConfig(\n chunk_censoring_func=censor_salesforce_chunks,\n ),\n ),\n DocumentSource.MOCK_CONNECTOR: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=mock_doc_sync,\n initial_index_should_sync=True,\n ),\n ),\n # Groups are not needed for Teams.\n # All channel access is done at the individual user level.\n DocumentSource.TEAMS: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=TEAMS_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=teams_doc_sync,\n initial_index_should_sync=True,\n ),\n ),\n DocumentSource.SHAREPOINT: SyncConfig(\n doc_sync_config=DocSyncConfig(\n doc_sync_frequency=SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY,\n doc_sync_func=sharepoint_doc_sync,\n initial_index_should_sync=True,\n ),\n group_sync_config=GroupSyncConfig(\n group_sync_frequency=SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY,\n group_sync_func=sharepoint_group_sync,\n group_sync_is_cc_pair_agnostic=False,\n ),\n ),\n}\n\n\ndef source_requires_doc_sync(source: DocumentSource) -> bool:\n \"\"\"Checks if the given DocumentSource requires doc syncing.\"\"\"\n if source not in _SOURCE_TO_SYNC_CONFIG:\n return False\n return _SOURCE_TO_SYNC_CONFIG[source].doc_sync_config is not None\n\n\ndef source_requires_external_group_sync(source: DocumentSource) -> bool:\n \"\"\"Checks if the given DocumentSource requires external group syncing.\"\"\"\n if source not in _SOURCE_TO_SYNC_CONFIG:\n return False\n return _SOURCE_TO_SYNC_CONFIG[source].group_sync_config is not None\n\n\ndef get_source_perm_sync_config(source: DocumentSource) -> SyncConfig | None:\n \"\"\"Returns the frequency of the external group sync for the given DocumentSource.\"\"\"\n return _SOURCE_TO_SYNC_CONFIG.get(source)\n\n\ndef source_group_sync_is_cc_pair_agnostic(source: DocumentSource) -> bool:\n \"\"\"Checks if the given DocumentSource requires external group syncing.\"\"\"\n if source not in _SOURCE_TO_SYNC_CONFIG:\n return False\n\n group_sync_config = _SOURCE_TO_SYNC_CONFIG[source].group_sync_config\n if group_sync_config is None:\n return False\n\n return group_sync_config.group_sync_is_cc_pair_agnostic\n\n\ndef get_all_cc_pair_agnostic_group_sync_sources() -> set[DocumentSource]:\n \"\"\"Returns the set of sources that have external group syncing that is cc_pair agnostic.\"\"\"\n return {\n source\n for source, sync_config in _SOURCE_TO_SYNC_CONFIG.items()\n if sync_config.group_sync_config is not None\n and sync_config.group_sync_config.group_sync_is_cc_pair_agnostic\n }\n\n\ndef check_if_valid_sync_source(source_type: DocumentSource) -> bool:\n return source_type in _SOURCE_TO_SYNC_CONFIG\n\n\ndef get_all_censoring_enabled_sources() -> set[DocumentSource]:\n \"\"\"Returns the set of sources that have censoring enabled.\"\"\"\n return {\n source\n for source, sync_config in _SOURCE_TO_SYNC_CONFIG.items()\n if sync_config.censoring_config is not None\n }\n\n\ndef source_should_fetch_permissions_during_indexing(source: DocumentSource) -> bool:\n \"\"\"Returns True if the given DocumentSource requires permissions to be fetched during indexing.\"\"\"\n if source not in _SOURCE_TO_SYNC_CONFIG:\n return False\n\n doc_sync_config = _SOURCE_TO_SYNC_CONFIG[source].doc_sync_config\n if doc_sync_config is None:\n return False\n\n return doc_sync_config.initial_index_should_sync\n" }, { "path": "backend/ee/onyx/external_permissions/teams/doc_sync.py", "content": "from collections.abc import Generator\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom ee.onyx.external_permissions.utils import generic_doc_sync\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.teams.connector import TeamsConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nTEAMS_DOC_SYNC_LABEL = \"teams_doc_sync\"\n\n\ndef teams_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_fn: FetchAllDocumentsFunction, # noqa: ARG001\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,\n callback: IndexingHeartbeatInterface | None,\n) -> Generator[ElementExternalAccess, None, None]:\n teams_connector = TeamsConnector(\n **cc_pair.connector.connector_specific_config,\n )\n credential_json = (\n cc_pair.credential.credential_json.get_value(apply_mask=False)\n if cc_pair.credential.credential_json\n else {}\n )\n teams_connector.load_credentials(credential_json)\n\n yield from generic_doc_sync(\n cc_pair=cc_pair,\n fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,\n callback=callback,\n doc_source=DocumentSource.TEAMS,\n slim_connector=teams_connector,\n label=TEAMS_DOC_SYNC_LABEL,\n )\n" }, { "path": "backend/ee/onyx/external_permissions/utils.py", "content": "from collections.abc import Generator\n\nfrom ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.access.models import ExternalAccess\nfrom onyx.access.models import NodeExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef generic_doc_sync(\n cc_pair: ConnectorCredentialPair,\n fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,\n callback: IndexingHeartbeatInterface | None,\n doc_source: DocumentSource,\n slim_connector: SlimConnectorWithPermSync,\n label: str,\n) -> Generator[ElementExternalAccess, None, None]:\n \"\"\"\n A convenience function for performing a generic document synchronization.\n\n Notes:\n A generic doc sync includes:\n - fetching existing docs\n - fetching *all* new (slim) docs\n - yielding external-access permissions for existing docs which do not exist in the newly fetched slim-docs set (with their\n `external_access` set to \"private\")\n - yielding external-access permissions for newly fetched docs and hierarchy nodes\n\n Returns:\n A `Generator` which yields existing and newly fetched external-access permissions.\n \"\"\"\n\n logger.info(f\"Starting {doc_source} doc sync for CC Pair ID: {cc_pair.id}\")\n\n indexing_start: SecondsSinceUnixEpoch | None = (\n cc_pair.connector.indexing_start.timestamp()\n if cc_pair.connector.indexing_start is not None\n else None\n )\n\n newly_fetched_doc_ids: set[str] = set()\n\n logger.info(f\"Fetching all slim documents from {doc_source}\")\n for doc_batch in slim_connector.retrieve_all_slim_docs_perm_sync(\n start=indexing_start,\n callback=callback,\n ):\n logger.info(f\"Got {len(doc_batch)} slim documents from {doc_source}\")\n\n if callback:\n if callback.should_stop():\n raise RuntimeError(f\"{label}: Stop signal detected\")\n callback.progress(label, 1)\n\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n # Yield hierarchy node permissions to be processed in outer layer\n if doc.external_access:\n yield NodeExternalAccess(\n external_access=doc.external_access,\n raw_node_id=doc.raw_node_id,\n source=doc_source.value,\n )\n continue\n if not doc.external_access:\n raise RuntimeError(\n f\"No external access found for document ID; {cc_pair.id=} {doc_source=} {doc.id=}\"\n )\n\n newly_fetched_doc_ids.add(doc.id)\n\n yield DocExternalAccess(\n doc_id=doc.id,\n external_access=doc.external_access,\n )\n\n logger.info(f\"Querying existing document IDs for CC Pair ID: {cc_pair.id=}\")\n existing_doc_ids: list[str] = fetch_all_existing_docs_ids_fn()\n\n missing_doc_ids = set(existing_doc_ids) - newly_fetched_doc_ids\n\n if not missing_doc_ids:\n return\n\n logger.warning(\n f\"Found {len(missing_doc_ids)=} documents that are in the DB but not present in fetch. Making them inaccessible.\"\n )\n\n for missing_id in missing_doc_ids:\n logger.warning(f\"Removing access for {missing_id=}\")\n yield DocExternalAccess(\n doc_id=missing_id,\n external_access=ExternalAccess.empty(),\n )\n\n logger.info(f\"Finished {doc_source} doc sync\")\n" }, { "path": "backend/ee/onyx/feature_flags/__init__.py", "content": "" }, { "path": "backend/ee/onyx/feature_flags/factory.py", "content": "from ee.onyx.feature_flags.posthog_provider import PostHogFeatureFlagProvider\nfrom onyx.feature_flags.interface import FeatureFlagProvider\n\n\ndef get_posthog_feature_flag_provider() -> FeatureFlagProvider:\n \"\"\"\n Get the PostHog feature flag provider instance.\n\n This is the EE implementation that gets loaded by the versioned\n implementation loader.\n\n Returns:\n PostHogFeatureFlagProvider: The PostHog-based feature flag provider\n \"\"\"\n return PostHogFeatureFlagProvider()\n" }, { "path": "backend/ee/onyx/feature_flags/posthog_provider.py", "content": "from typing import Any\nfrom uuid import UUID\n\nfrom ee.onyx.utils.posthog_client import posthog\nfrom onyx.feature_flags.interface import FeatureFlagProvider\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass PostHogFeatureFlagProvider(FeatureFlagProvider):\n \"\"\"\n PostHog-based feature flag provider.\n\n Uses PostHog's feature flag API to determine if features are enabled\n for specific users. Only active in multi-tenant mode.\n \"\"\"\n\n def feature_enabled(\n self,\n flag_key: str,\n user_id: UUID,\n user_properties: dict[str, Any] | None = None,\n ) -> bool:\n \"\"\"\n Check if a feature flag is enabled for a user via PostHog.\n\n Args:\n flag_key: The identifier for the feature flag to check\n user_id: The unique identifier for the user\n user_properties: Optional dictionary of user properties/attributes\n that may influence flag evaluation\n\n Returns:\n True if the feature is enabled for the user, False otherwise.\n \"\"\"\n if not posthog:\n return False\n\n try:\n posthog.set(\n distinct_id=user_id,\n properties=user_properties,\n )\n is_enabled = posthog.feature_enabled(\n flag_key,\n str(user_id),\n person_properties=user_properties,\n )\n\n return bool(is_enabled) if is_enabled is not None else False\n\n except Exception as e:\n logger.error(\n f\"Error checking feature flag {flag_key} for user {user_id}: {e}\"\n )\n return False\n" }, { "path": "backend/ee/onyx/hooks/__init__.py", "content": "" }, { "path": "backend/ee/onyx/hooks/executor.py", "content": "\"\"\"Hook executor — calls a customer's external HTTP endpoint for a given hook point.\n\nUsage (Celery tasks and FastAPI handlers):\n result = execute_hook(\n db_session=db_session,\n hook_point=HookPoint.QUERY_PROCESSING,\n payload={\"query\": \"...\", \"user_email\": \"...\", \"chat_session_id\": \"...\"},\n response_type=QueryProcessingResponse,\n )\n\n if isinstance(result, HookSkipped):\n # no active hook configured — continue with original behavior\n ...\n elif isinstance(result, HookSoftFailed):\n # hook failed but fail strategy is SOFT — continue with original behavior\n ...\n else:\n # result is a validated Pydantic model instance (response_type)\n ...\n\nis_reachable update policy\n--------------------------\n``is_reachable`` on the Hook row is updated selectively — only when the outcome\ncarries meaningful signal about physical reachability:\n\n NetworkError (DNS, connection refused) → False (cannot reach the server)\n HTTP 401 / 403 → False (api_key revoked or invalid)\n TimeoutException → None (server may be slow, skip write)\n Other HTTP errors (4xx / 5xx) → None (server responded, skip write)\n Unknown exception → None (no signal, skip write)\n Non-JSON / non-dict response → None (server responded, skip write)\n Success (2xx, valid dict) → True (confirmed reachable)\n\nNone means \"leave the current value unchanged\" — no DB round-trip is made.\n\nDB session design\n-----------------\nThe executor uses three sessions:\n\n 1. Caller's session (db_session) — used only for the hook lookup read. All\n needed fields are extracted from the Hook object before the HTTP call, so\n the caller's session is not held open during the external HTTP request.\n\n 2. Log session — a separate short-lived session opened after the HTTP call\n completes to write the HookExecutionLog row on failure. Success runs are\n not recorded. Committed independently of everything else.\n\n 3. Reachable session — a second short-lived session to update is_reachable on\n the Hook. Kept separate from the log session so a concurrent hook deletion\n (which causes update_hook__no_commit to raise OnyxError(NOT_FOUND)) cannot\n prevent the execution log from being written. This update is best-effort.\n\"\"\"\n\nimport json\nimport time\nfrom typing import Any\nfrom typing import TypeVar\n\nimport httpx\nfrom pydantic import BaseModel\nfrom pydantic import ValidationError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import HookFailStrategy\nfrom onyx.db.enums import HookPoint\nfrom onyx.db.hook import create_hook_execution_log__no_commit\nfrom onyx.db.hook import get_non_deleted_hook_by_hook_point\nfrom onyx.db.hook import update_hook__no_commit\nfrom onyx.db.models import Hook\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.hooks.executor import HookSkipped\nfrom onyx.hooks.executor import HookSoftFailed\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\nT = TypeVar(\"T\", bound=BaseModel)\n\n\n# ---------------------------------------------------------------------------\n# Private helpers\n# ---------------------------------------------------------------------------\n\n\nclass _HttpOutcome(BaseModel):\n \"\"\"Structured result of an HTTP hook call, returned by _process_response.\"\"\"\n\n is_success: bool\n updated_is_reachable: (\n bool | None\n ) # True/False = write to DB, None = unchanged (skip write)\n status_code: int | None\n error_message: str | None\n response_payload: dict[str, Any] | None\n\n\ndef _lookup_hook(\n db_session: Session,\n hook_point: HookPoint,\n) -> Hook | HookSkipped:\n \"\"\"Return the active Hook or HookSkipped if hooks are unavailable/unconfigured.\n\n No HTTP call is made and no DB writes are performed for any HookSkipped path.\n There is nothing to log and no reachability information to update.\n \"\"\"\n if MULTI_TENANT:\n return HookSkipped()\n hook = get_non_deleted_hook_by_hook_point(\n db_session=db_session, hook_point=hook_point\n )\n if hook is None or not hook.is_active:\n return HookSkipped()\n if not hook.endpoint_url:\n return HookSkipped()\n return hook\n\n\ndef _process_response(\n *,\n response: httpx.Response | None,\n exc: Exception | None,\n timeout: float,\n) -> _HttpOutcome:\n \"\"\"Process the result of an HTTP call and return a structured outcome.\n\n Called after the client.post() try/except. If post() raised, exc is set and\n response is None. Otherwise response is set and exc is None. Handles\n raise_for_status(), JSON decoding, and the dict shape check.\n \"\"\"\n if exc is not None:\n if isinstance(exc, httpx.NetworkError):\n msg = f\"Hook network error (endpoint unreachable): {exc}\"\n logger.warning(msg, exc_info=exc)\n return _HttpOutcome(\n is_success=False,\n updated_is_reachable=False,\n status_code=None,\n error_message=msg,\n response_payload=None,\n )\n if isinstance(exc, httpx.TimeoutException):\n msg = f\"Hook timed out after {timeout}s: {exc}\"\n logger.warning(msg, exc_info=exc)\n return _HttpOutcome(\n is_success=False,\n updated_is_reachable=None, # timeout doesn't indicate unreachability\n status_code=None,\n error_message=msg,\n response_payload=None,\n )\n msg = f\"Hook call failed: {exc}\"\n logger.exception(msg, exc_info=exc)\n return _HttpOutcome(\n is_success=False,\n updated_is_reachable=None, # unknown error — don't make assumptions\n status_code=None,\n error_message=msg,\n response_payload=None,\n )\n\n if response is None:\n raise ValueError(\n \"exactly one of response or exc must be non-None; both are None\"\n )\n status_code = response.status_code\n\n try:\n response.raise_for_status()\n except httpx.HTTPStatusError as e:\n msg = f\"Hook returned HTTP {e.response.status_code}: {e.response.text}\"\n logger.warning(msg, exc_info=e)\n # 401/403 means the api_key has been revoked or is invalid — mark unreachable\n # so the operator knows to update it. All other HTTP errors keep is_reachable\n # as-is (server is up, the request just failed for application reasons).\n auth_failed = e.response.status_code in (401, 403)\n return _HttpOutcome(\n is_success=False,\n updated_is_reachable=False if auth_failed else None,\n status_code=status_code,\n error_message=msg,\n response_payload=None,\n )\n\n try:\n response_payload = response.json()\n except (json.JSONDecodeError, httpx.DecodingError) as e:\n msg = f\"Hook returned non-JSON response: {e}\"\n logger.warning(msg, exc_info=e)\n return _HttpOutcome(\n is_success=False,\n updated_is_reachable=None, # server responded — reachability unchanged\n status_code=status_code,\n error_message=msg,\n response_payload=None,\n )\n\n if not isinstance(response_payload, dict):\n msg = f\"Hook returned non-dict JSON (got {type(response_payload).__name__})\"\n logger.warning(msg)\n return _HttpOutcome(\n is_success=False,\n updated_is_reachable=None, # server responded — reachability unchanged\n status_code=status_code,\n error_message=msg,\n response_payload=None,\n )\n\n return _HttpOutcome(\n is_success=True,\n updated_is_reachable=True,\n status_code=status_code,\n error_message=None,\n response_payload=response_payload,\n )\n\n\ndef _persist_result(\n *,\n hook_id: int,\n outcome: _HttpOutcome,\n duration_ms: int,\n) -> None:\n \"\"\"Write the execution log on failure and optionally update is_reachable, each\n in its own session so a failure in one does not affect the other.\"\"\"\n # Only write the execution log on failure — success runs are not recorded.\n # Must not be skipped if the is_reachable update fails (e.g. hook concurrently\n # deleted between the initial lookup and here).\n if not outcome.is_success:\n try:\n with get_session_with_current_tenant() as log_session:\n create_hook_execution_log__no_commit(\n db_session=log_session,\n hook_id=hook_id,\n is_success=False,\n error_message=outcome.error_message,\n status_code=outcome.status_code,\n duration_ms=duration_ms,\n )\n log_session.commit()\n except Exception:\n logger.exception(\n f\"Failed to persist hook execution log for hook_id={hook_id}\"\n )\n\n # Update is_reachable separately — best-effort, non-critical.\n # None means the value is unchanged (set by the caller to skip the no-op write).\n # update_hook__no_commit can raise OnyxError(NOT_FOUND) if the hook was\n # concurrently deleted, so keep this isolated from the log write above.\n if outcome.updated_is_reachable is not None:\n try:\n with get_session_with_current_tenant() as reachable_session:\n update_hook__no_commit(\n db_session=reachable_session,\n hook_id=hook_id,\n is_reachable=outcome.updated_is_reachable,\n )\n reachable_session.commit()\n except Exception:\n logger.warning(f\"Failed to update is_reachable for hook_id={hook_id}\")\n\n\n# ---------------------------------------------------------------------------\n# Public API\n# ---------------------------------------------------------------------------\n\n\ndef _execute_hook_inner(\n hook: Hook,\n payload: dict[str, Any],\n response_type: type[T],\n) -> T | HookSoftFailed:\n \"\"\"Make the HTTP call, validate the response, and return a typed model.\n\n Raises OnyxError on HARD failure. Returns HookSoftFailed on SOFT failure.\n \"\"\"\n timeout = hook.timeout_seconds\n hook_id = hook.id\n fail_strategy = hook.fail_strategy\n endpoint_url = hook.endpoint_url\n current_is_reachable: bool | None = hook.is_reachable\n\n if not endpoint_url:\n raise ValueError(\n f\"hook_id={hook_id} is active but has no endpoint_url — \"\n \"active hooks without an endpoint_url must be rejected by _lookup_hook\"\n )\n\n start = time.monotonic()\n response: httpx.Response | None = None\n exc: Exception | None = None\n try:\n api_key: str | None = (\n hook.api_key.get_value(apply_mask=False) if hook.api_key else None\n )\n headers: dict[str, str] = {\"Content-Type\": \"application/json\"}\n if api_key:\n headers[\"Authorization\"] = f\"Bearer {api_key}\"\n with httpx.Client(\n timeout=timeout, follow_redirects=False\n ) as client: # SSRF guard: never follow redirects\n response = client.post(endpoint_url, json=payload, headers=headers)\n except Exception as e:\n exc = e\n duration_ms = int((time.monotonic() - start) * 1000)\n\n outcome = _process_response(response=response, exc=exc, timeout=timeout)\n\n # Validate the response payload against response_type.\n # A validation failure downgrades the outcome to a failure so it is logged,\n # is_reachable is left unchanged (server responded — just a bad payload),\n # and fail_strategy is respected below.\n validated_model: T | None = None\n if outcome.is_success and outcome.response_payload is not None:\n try:\n validated_model = response_type.model_validate(outcome.response_payload)\n except ValidationError as e:\n msg = (\n f\"Hook response failed validation against {response_type.__name__}: {e}\"\n )\n outcome = _HttpOutcome(\n is_success=False,\n updated_is_reachable=None, # server responded — reachability unchanged\n status_code=outcome.status_code,\n error_message=msg,\n response_payload=None,\n )\n\n # Skip the is_reachable write when the value would not change — avoids a\n # no-op DB round-trip on every call when the hook is already in the expected state.\n if outcome.updated_is_reachable == current_is_reachable:\n outcome = outcome.model_copy(update={\"updated_is_reachable\": None})\n _persist_result(hook_id=hook_id, outcome=outcome, duration_ms=duration_ms)\n\n if not outcome.is_success:\n if fail_strategy == HookFailStrategy.HARD:\n raise OnyxError(\n OnyxErrorCode.HOOK_EXECUTION_FAILED,\n outcome.error_message or \"Hook execution failed.\",\n )\n logger.warning(\n f\"Hook execution failed (soft fail) for hook_id={hook_id}: {outcome.error_message}\"\n )\n return HookSoftFailed()\n\n if validated_model is None:\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n f\"validated_model is None for successful hook call (hook_id={hook_id})\",\n )\n return validated_model\n\n\ndef _execute_hook_impl(\n *,\n db_session: Session,\n hook_point: HookPoint,\n payload: dict[str, Any],\n response_type: type[T],\n) -> T | HookSkipped | HookSoftFailed:\n \"\"\"EE implementation — loaded by CE's execute_hook via fetch_versioned_implementation.\n\n Returns HookSkipped if no active hook is configured, HookSoftFailed if the\n hook failed with SOFT fail strategy, or a validated response model on success.\n Raises OnyxError on HARD failure or if the hook is misconfigured.\n \"\"\"\n hook = _lookup_hook(db_session, hook_point)\n if isinstance(hook, HookSkipped):\n return hook\n\n fail_strategy = hook.fail_strategy\n hook_id = hook.id\n\n try:\n return _execute_hook_inner(hook, payload, response_type)\n except Exception:\n if fail_strategy == HookFailStrategy.SOFT:\n logger.exception(\n f\"Unexpected error in hook execution (soft fail) for hook_id={hook_id}\"\n )\n return HookSoftFailed()\n raise\n" }, { "path": "backend/ee/onyx/main.py", "content": "from collections.abc import AsyncGenerator\nfrom contextlib import asynccontextmanager\n\nfrom fastapi import FastAPI\nfrom httpx_oauth.clients.google import GoogleOAuth2\n\nfrom ee.onyx.server.analytics.api import router as analytics_router\nfrom ee.onyx.server.auth_check import check_ee_router_auth\nfrom ee.onyx.server.billing.api import router as billing_router\nfrom ee.onyx.server.documents.cc_pair import router as ee_document_cc_pair_router\nfrom ee.onyx.server.enterprise_settings.api import (\n admin_router as enterprise_settings_admin_router,\n)\nfrom ee.onyx.server.enterprise_settings.api import (\n basic_router as enterprise_settings_router,\n)\nfrom ee.onyx.server.evals.api import router as evals_router\nfrom ee.onyx.server.features.hooks.api import router as hook_router\nfrom ee.onyx.server.license.api import router as license_router\nfrom ee.onyx.server.manage.standard_answer import router as standard_answer_router\nfrom ee.onyx.server.middleware.license_enforcement import (\n add_license_enforcement_middleware,\n)\nfrom ee.onyx.server.middleware.tenant_tracking import (\n add_api_server_tenant_id_middleware,\n)\nfrom ee.onyx.server.oauth.api import router as ee_oauth_router\nfrom ee.onyx.server.query_and_chat.query_backend import (\n basic_router as ee_query_router,\n)\nfrom ee.onyx.server.query_and_chat.search_backend import router as search_router\nfrom ee.onyx.server.query_history.api import router as query_history_router\nfrom ee.onyx.server.reporting.usage_export_api import router as usage_export_router\nfrom ee.onyx.server.scim.api import register_scim_exception_handlers\nfrom ee.onyx.server.scim.api import scim_router\nfrom ee.onyx.server.seeding import seed_db\nfrom ee.onyx.server.tenants.api import router as tenants_router\nfrom ee.onyx.server.token_rate_limits.api import (\n router as token_rate_limit_settings_router,\n)\nfrom ee.onyx.server.user_group.api import router as user_group_router\nfrom ee.onyx.utils.encryption import test_encryption\nfrom onyx.auth.users import auth_backend\nfrom onyx.auth.users import create_onyx_oauth_router\nfrom onyx.auth.users import fastapi_users\nfrom onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.app_configs import OAUTH_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_CLIENT_SECRET\nfrom onyx.configs.app_configs import USER_AUTH_SECRET\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import AuthType\nfrom onyx.main import get_application as get_application_base\nfrom onyx.main import include_auth_router_with_prefix\nfrom onyx.main import include_router_with_global_prefix_prepended\nfrom onyx.main import lifespan as lifespan_base\nfrom onyx.main import use_route_function_names_as_operation_ids\nfrom onyx.server.query_and_chat.query_backend import (\n basic_router as query_router,\n)\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import global_version\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\n@asynccontextmanager\nasync def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:\n \"\"\"Small wrapper around the lifespan of the MIT application.\n Basically just calls the base lifespan, and then adds EE-only\n steps after.\"\"\"\n\n async with lifespan_base(app):\n # seed the Onyx environment with LLMs, Assistants, etc. based on an optional\n # environment variable. Used to automate deployment for multiple environments.\n seed_db()\n\n yield\n\n\ndef get_application() -> FastAPI:\n # Anything that happens at import time is not guaranteed to be running ee-version\n # Anything after the server startup will be running ee version\n global_version.set_ee()\n\n test_encryption()\n\n application = get_application_base(lifespan_override=lifespan)\n\n if MULTI_TENANT:\n add_api_server_tenant_id_middleware(application, logger)\n else:\n # License enforcement middleware for self-hosted deployments only\n # Checks LICENSE_ENFORCEMENT_ENABLED at runtime (can be toggled without restart)\n # MT deployments use control plane gating via is_tenant_gated() instead\n add_license_enforcement_middleware(application, logger)\n\n if AUTH_TYPE == AuthType.CLOUD:\n # For Google OAuth, refresh tokens are requested by:\n # 1. Adding the right scopes\n # 2. Properly configuring OAuth in Google Cloud Console to allow offline access\n oauth_client = GoogleOAuth2(\n OAUTH_CLIENT_ID,\n OAUTH_CLIENT_SECRET,\n # Use standard scopes that include profile and email\n scopes=[\"openid\", \"email\", \"profile\"],\n )\n include_auth_router_with_prefix(\n application,\n create_onyx_oauth_router(\n oauth_client,\n auth_backend,\n USER_AUTH_SECRET,\n associate_by_email=True,\n is_verified_by_default=True,\n # Points the user back to the login page\n redirect_url=f\"{WEB_DOMAIN}/auth/oauth/callback\",\n ),\n prefix=\"/auth/oauth\",\n )\n\n # Need basic auth router for `logout` endpoint\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_logout_router(auth_backend),\n prefix=\"/auth\",\n )\n\n # RBAC / group access control\n include_router_with_global_prefix_prepended(application, user_group_router)\n # Analytics endpoints\n include_router_with_global_prefix_prepended(application, analytics_router)\n include_router_with_global_prefix_prepended(application, query_history_router)\n # EE only backend APIs\n include_router_with_global_prefix_prepended(application, query_router)\n include_router_with_global_prefix_prepended(application, ee_query_router)\n include_router_with_global_prefix_prepended(application, search_router)\n include_router_with_global_prefix_prepended(application, standard_answer_router)\n include_router_with_global_prefix_prepended(application, ee_oauth_router)\n include_router_with_global_prefix_prepended(application, ee_document_cc_pair_router)\n include_router_with_global_prefix_prepended(application, evals_router)\n include_router_with_global_prefix_prepended(application, hook_router)\n\n # Enterprise-only global settings\n include_router_with_global_prefix_prepended(\n application, enterprise_settings_admin_router\n )\n # Token rate limit settings\n include_router_with_global_prefix_prepended(\n application, token_rate_limit_settings_router\n )\n include_router_with_global_prefix_prepended(application, enterprise_settings_router)\n include_router_with_global_prefix_prepended(application, usage_export_router)\n # License management\n include_router_with_global_prefix_prepended(application, license_router)\n\n # Unified billing API - always registered in EE.\n # Each endpoint is protected by the `current_admin_user` dependency (admin auth).\n include_router_with_global_prefix_prepended(application, billing_router)\n\n if MULTI_TENANT:\n # Tenant management\n include_router_with_global_prefix_prepended(application, tenants_router)\n\n # SCIM 2.0 — protocol endpoints (unauthenticated by Onyx session auth;\n # they use their own SCIM bearer token auth).\n # Not behind APP_API_PREFIX because IdPs expect /scim/v2/... directly.\n application.include_router(scim_router)\n register_scim_exception_handlers(application)\n\n # Ensure all routes have auth enabled or are explicitly marked as public\n check_ee_router_auth(application)\n\n # for debugging discovered routes\n # for route in application.router.routes:\n # print(f\"Path: {route.path}, Methods: {route.methods}\")\n\n use_route_function_names_as_operation_ids(application)\n\n return application\n" }, { "path": "backend/ee/onyx/onyxbot/slack/handlers/__init__.py", "content": "" }, { "path": "backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py", "content": "from slack_sdk import WebClient\nfrom slack_sdk.models.blocks import ActionsBlock\nfrom slack_sdk.models.blocks import Block\nfrom slack_sdk.models.blocks import ButtonElement\nfrom slack_sdk.models.blocks import SectionBlock\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.standard_answer import fetch_standard_answer_categories_by_names\nfrom ee.onyx.db.standard_answer import find_matching_standard_answers\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_REACT_EMOJI\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.chat import create_new_chat_message\nfrom onyx.db.chat import get_chat_messages_by_sessions\nfrom onyx.db.chat import get_chat_sessions_by_slack_thread_id\nfrom onyx.db.chat import get_or_create_root_message\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.db.models import StandardAnswer as StandardAnswerModel\nfrom onyx.onyxbot.slack.blocks import get_restate_blocks\nfrom onyx.onyxbot.slack.constants import GENERATE_ANSWER_BUTTON_ACTION_ID\nfrom onyx.onyxbot.slack.handlers.utils import send_team_member_message\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.onyxbot.slack.utils import respond_in_thread_or_channel\nfrom onyx.onyxbot.slack.utils import update_emote_react\nfrom onyx.server.manage.models import StandardAnswer as PydanticStandardAnswer\nfrom onyx.utils.logger import OnyxLoggingAdapter\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef build_standard_answer_blocks(\n answer_message: str,\n) -> list[Block]:\n generate_button_block = ButtonElement(\n action_id=GENERATE_ANSWER_BUTTON_ACTION_ID,\n text=\"Generate Full Answer\",\n )\n answer_block = SectionBlock(text=answer_message)\n return [\n answer_block,\n ActionsBlock(\n elements=[generate_button_block],\n ),\n ]\n\n\ndef oneoff_standard_answers(\n message: str,\n slack_bot_categories: list[str],\n db_session: Session,\n) -> list[PydanticStandardAnswer]:\n \"\"\"\n Respond to the user message if it matches any configured standard answers.\n\n Returns a list of matching StandardAnswers if found, otherwise None.\n \"\"\"\n configured_standard_answers = {\n standard_answer\n for category in fetch_standard_answer_categories_by_names(\n slack_bot_categories, db_session=db_session\n )\n for standard_answer in category.standard_answers\n }\n\n matching_standard_answers = find_matching_standard_answers(\n query=message,\n id_in=[answer.id for answer in configured_standard_answers],\n db_session=db_session,\n )\n\n server_standard_answers = [\n PydanticStandardAnswer.from_model(standard_answer_model)\n for (standard_answer_model, _) in matching_standard_answers\n ]\n return server_standard_answers\n\n\ndef _handle_standard_answers(\n message_info: SlackMessageInfo,\n receiver_ids: list[str] | None,\n slack_channel_config: SlackChannelConfig,\n logger: OnyxLoggingAdapter,\n client: WebClient,\n db_session: Session,\n) -> bool:\n \"\"\"\n Potentially respond to the user message depending on whether the user's message matches\n any of the configured standard answers and also whether those answers have already been\n provided in the current thread.\n\n Returns True if standard answers are found to match the user's message and therefore,\n we still need to respond to the users.\n \"\"\"\n\n slack_thread_id = message_info.thread_to_respond\n configured_standard_answer_categories = (\n slack_channel_config.standard_answer_categories\n )\n configured_standard_answers = set(\n [\n standard_answer\n for standard_answer_category in configured_standard_answer_categories\n for standard_answer in standard_answer_category.standard_answers\n ]\n )\n query_msg = message_info.thread_messages[-1]\n\n if slack_thread_id is None:\n used_standard_answer_ids = set([])\n else:\n chat_sessions = get_chat_sessions_by_slack_thread_id(\n slack_thread_id=slack_thread_id,\n user_id=None,\n db_session=db_session,\n )\n chat_messages = get_chat_messages_by_sessions(\n chat_session_ids=[chat_session.id for chat_session in chat_sessions],\n user_id=None,\n db_session=db_session,\n skip_permission_check=True,\n )\n used_standard_answer_ids = set(\n [\n standard_answer.id\n for chat_message in chat_messages\n for standard_answer in chat_message.standard_answers\n ]\n )\n\n usable_standard_answers = configured_standard_answers.difference(\n used_standard_answer_ids\n )\n\n matching_standard_answers: list[tuple[StandardAnswerModel, str]] = []\n if usable_standard_answers:\n matching_standard_answers = find_matching_standard_answers(\n query=query_msg.message,\n id_in=[standard_answer.id for standard_answer in usable_standard_answers],\n db_session=db_session,\n )\n\n if matching_standard_answers:\n chat_session = create_chat_session(\n db_session=db_session,\n description=\"\",\n user_id=None,\n persona_id=(\n slack_channel_config.persona.id if slack_channel_config.persona else 0\n ),\n onyxbot_flow=True,\n slack_thread_id=slack_thread_id,\n )\n\n root_message = get_or_create_root_message(\n chat_session_id=chat_session.id, db_session=db_session\n )\n\n new_user_message = create_new_chat_message(\n chat_session_id=chat_session.id,\n parent_message=root_message,\n message=query_msg.message,\n token_count=0,\n message_type=MessageType.USER,\n db_session=db_session,\n commit=True,\n )\n\n formatted_answers = []\n for standard_answer, match_str in matching_standard_answers:\n since_you_mentioned_pretext = (\n f'Since your question contains \"_{match_str}_\"'\n )\n block_quotified_answer = \">\" + standard_answer.answer.replace(\"\\n\", \"\\n> \")\n formatted_answer = f\"{since_you_mentioned_pretext}, I thought this might be useful: \\n\\n{block_quotified_answer}\"\n formatted_answers.append(formatted_answer)\n answer_message = \"\\n\\n\".join(formatted_answers)\n\n chat_message = create_new_chat_message(\n chat_session_id=chat_session.id,\n parent_message=new_user_message,\n message=answer_message,\n token_count=0,\n message_type=MessageType.ASSISTANT,\n error=None,\n db_session=db_session,\n commit=False,\n )\n # attach the standard answers to the chat message\n chat_message.standard_answers = [\n standard_answer for standard_answer, _ in matching_standard_answers\n ]\n db_session.commit()\n\n update_emote_react(\n emoji=ONYX_BOT_REACT_EMOJI,\n channel=message_info.channel_to_respond,\n message_ts=message_info.msg_to_respond,\n remove=True,\n client=client,\n )\n\n restate_question_blocks = get_restate_blocks(\n msg=query_msg.message,\n is_slash_command=message_info.is_slash_command,\n )\n\n answer_blocks = build_standard_answer_blocks(\n answer_message=answer_message,\n )\n\n all_blocks = restate_question_blocks + answer_blocks\n\n try:\n respond_in_thread_or_channel(\n client=client,\n channel=message_info.channel_to_respond,\n receiver_ids=receiver_ids,\n text=\"Hello! Onyx has some results for you!\",\n blocks=all_blocks,\n thread_ts=message_info.msg_to_respond,\n unfurl=False,\n )\n\n if receiver_ids and slack_thread_id:\n send_team_member_message(\n client=client,\n channel=message_info.channel_to_respond,\n thread_ts=slack_thread_id,\n receiver_ids=receiver_ids,\n )\n\n return True\n except Exception as e:\n logger.exception(f\"Unable to send standard answer message: {e}\")\n return False\n else:\n return False\n" }, { "path": "backend/ee/onyx/prompts/__init__.py", "content": "" }, { "path": "backend/ee/onyx/prompts/query_expansion.py", "content": "# Single message is likely most reliable and generally better for this task\n# No final reminders at the end since the user query is expected to be short\n# If it is not short, it should go into the chat flow so we do not need to account for this.\nKEYWORD_EXPANSION_PROMPT = \"\"\"\nGenerate a set of keyword-only queries to help find relevant documents for the provided query. \\\nThese queries will be passed to a bm25-based keyword search engine. \\\nProvide a single query per line (where each query consists of one or more keywords). \\\nThe queries must be purely keywords and not contain any filler natural language. \\\nThe each query should have as few keywords as necessary to represent the user's search intent. \\\nIf there are no useful expansions, simply return the original query with no additional keyword queries. \\\nCRITICAL: Do not include any additional formatting, comments, or anything aside from the keyword queries.\n\nThe user query is:\n{user_query}\n\"\"\".strip()\n\n\nQUERY_TYPE_PROMPT = \"\"\"\nDetermine if the provided query is better suited for a keyword search or a semantic search.\nRespond with \"keyword\" or \"semantic\" literally and nothing else.\nDo not provide any additional text or reasoning to your response.\n\nCRITICAL: It must only be 1 single word - EITHER \"keyword\" or \"semantic\".\n\nThe user query is:\n{user_query}\n\"\"\".strip()\n" }, { "path": "backend/ee/onyx/prompts/search_flow_classification.py", "content": "# ruff: noqa: E501, W605 start\nSEARCH_CLASS = \"search\"\nCHAT_CLASS = \"chat\"\n\n# Will note that with many larger LLMs the latency on running this prompt via third party APIs is as high as 2 seconds which is too slow for many\n# use cases.\nSEARCH_CHAT_PROMPT = f\"\"\"\nDetermine if the following query is better suited for a search UI or a chat UI. Respond with \"{SEARCH_CLASS}\" or \"{CHAT_CLASS}\" literally and nothing else. \\\nDo not provide any additional text or reasoning to your response. CRITICAL, IT MUST ONLY BE 1 SINGLE WORD - EITHER \"{SEARCH_CLASS}\" or \"{CHAT_CLASS}\".\n\n# Classification Guidelines:\n## {SEARCH_CLASS}\n- If the query consists entirely of keywords or query doesn't require any answer from the AI\n- If the query is a short statement that seems like a search query rather than a question\n- If the query feels nonsensical or is a short phrase that possibly describes a document or information that could be found in a internal document\n\n### Examples of {SEARCH_CLASS} queries:\n- Find me the document that goes over the onboarding process for a new hire\n- Pull requests since last week\n- Sales Runbook AMEA Region\n- Procurement process\n- Retrieve the PRD for project X\n\n## {CHAT_CLASS}\n- If the query is asking a question that requires an answer rather than a document\n- If the query is asking for a solution, suggestion, or general help\n- If the query is seeking information that is on the web and likely not in a company internal document\n- If the query should be answered without any context from additional documents or searches\n\n### Examples of {CHAT_CLASS} queries:\n- What led us to win the deal with company X? (seeking answer)\n- Google Drive not sync-ing files to my computer (seeking solution)\n- Review my email: (general help)\n- Write me a script to... (general help)\n- Cheap flights Europe to Tokyo (information likely found on the web, not internal)\n\n# User Query:\n{{user_query}}\n\nREMEMBER TO ONLY RESPOND WITH \"{SEARCH_CLASS}\" OR \"{CHAT_CLASS}\" AND NOTHING ELSE.\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/ee/onyx/search/process_search_query.py", "content": "from collections.abc import Generator\n\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.search import create_search_query\nfrom ee.onyx.secondary_llm_flows.query_expansion import expand_keywords\nfrom ee.onyx.server.query_and_chat.models import SearchDocWithContent\nfrom ee.onyx.server.query_and_chat.models import SearchFullResponse\nfrom ee.onyx.server.query_and_chat.models import SendSearchQueryRequest\nfrom ee.onyx.server.query_and_chat.streaming_models import LLMSelectedDocsPacket\nfrom ee.onyx.server.query_and_chat.streaming_models import SearchDocsPacket\nfrom ee.onyx.server.query_and_chat.streaming_models import SearchErrorPacket\nfrom ee.onyx.server.query_and_chat.streaming_models import SearchQueriesPacket\nfrom onyx.context.search.models import BaseFilters\nfrom onyx.context.search.models import ChunkSearchRequest\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.pipeline import merge_individual_chunks\nfrom onyx.context.search.pipeline import search_pipeline\nfrom onyx.db.models import User\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.document_index.factory import get_default_document_index\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.secondary_llm_flows.document_filter import select_sections_for_expansion\nfrom onyx.tools.tool_implementations.search.search_utils import (\n weighted_reciprocal_rank_fusion,\n)\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\n\nlogger = setup_logger()\n\n\n# This is just a heuristic that also happens to work well for the UI/UX\n# Users would not find it useful to see a huge list of suggested docs\n# but more than 1 is also likely good as many questions may target more than 1 doc.\nTARGET_NUM_SECTIONS_FOR_LLM_SELECTION = 3\n\n\ndef _run_single_search(\n query: str,\n filters: BaseFilters | None,\n document_index: DocumentIndex,\n user: User,\n db_session: Session,\n num_hits: int | None = None,\n hybrid_alpha: float | None = None,\n) -> list[InferenceChunk]:\n \"\"\"Execute a single search query and return chunks.\"\"\"\n chunk_search_request = ChunkSearchRequest(\n query=query,\n user_selected_filters=filters,\n limit=num_hits,\n hybrid_alpha=hybrid_alpha,\n )\n\n return search_pipeline(\n chunk_search_request=chunk_search_request,\n document_index=document_index,\n user=user,\n persona_search_info=None,\n db_session=db_session,\n )\n\n\ndef stream_search_query(\n request: SendSearchQueryRequest,\n user: User,\n db_session: Session,\n) -> Generator[\n SearchQueriesPacket | SearchDocsPacket | LLMSelectedDocsPacket | SearchErrorPacket,\n None,\n None,\n]:\n \"\"\"\n Core search function that yields streaming packets.\n Used by both streaming and non-streaming endpoints.\n \"\"\"\n # Get document index.\n search_settings = get_current_search_settings(db_session)\n # This flow is for search so we do not get all indices.\n document_index = get_default_document_index(search_settings, None, db_session)\n\n # Determine queries to execute\n original_query = request.search_query\n keyword_expansions: list[str] = []\n\n if request.run_query_expansion:\n try:\n llm = get_default_llm()\n keyword_expansions = expand_keywords(\n user_query=original_query,\n llm=llm,\n )\n if keyword_expansions:\n logger.debug(\n f\"Query expansion generated {len(keyword_expansions)} keyword queries\"\n )\n except Exception as e:\n logger.warning(f\"Query expansion failed: {e}; using original query only.\")\n keyword_expansions = []\n\n # Build list of all executed queries for tracking\n all_executed_queries = [original_query] + keyword_expansions\n\n if not user.is_anonymous:\n create_search_query(\n db_session=db_session,\n user_id=user.id,\n query=request.search_query,\n query_expansions=keyword_expansions if keyword_expansions else None,\n )\n\n # Execute search(es)\n if not keyword_expansions:\n # Single query (original only) - no threading needed\n chunks = _run_single_search(\n query=original_query,\n filters=request.filters,\n document_index=document_index,\n user=user,\n db_session=db_session,\n num_hits=request.num_hits,\n hybrid_alpha=request.hybrid_alpha,\n )\n else:\n # Multiple queries - run in parallel and merge with RRF\n # First query is the original (semantic), rest are keyword expansions\n search_functions = [\n (\n _run_single_search,\n (\n query,\n request.filters,\n document_index,\n user,\n db_session,\n request.num_hits,\n request.hybrid_alpha,\n ),\n )\n for query in all_executed_queries\n ]\n\n # Run all searches in parallel\n all_search_results: list[list[InferenceChunk]] = (\n run_functions_tuples_in_parallel(\n search_functions,\n allow_failures=True,\n )\n )\n\n # Separate original query results from keyword expansion results\n # Note that in rare cases, the original query may have failed and so we may be\n # just overweighting one set of keyword results, should be not a big deal though.\n original_result = all_search_results[0] if all_search_results else []\n keyword_results = all_search_results[1:] if len(all_search_results) > 1 else []\n\n # Build valid results and weights\n # Original query (semantic): weight 2.0\n # Keyword expansions: weight 1.0 each\n valid_results: list[list[InferenceChunk]] = []\n weights: list[float] = []\n\n if original_result:\n valid_results.append(original_result)\n weights.append(2.0)\n\n for keyword_result in keyword_results:\n if keyword_result:\n valid_results.append(keyword_result)\n weights.append(1.0)\n\n if not valid_results:\n logger.warning(\"All parallel searches returned empty results\")\n chunks = []\n else:\n chunks = weighted_reciprocal_rank_fusion(\n ranked_results=valid_results,\n weights=weights,\n id_extractor=lambda chunk: f\"{chunk.document_id}_{chunk.chunk_id}\",\n )\n\n # Merge chunks into sections\n sections = merge_individual_chunks(chunks)\n\n # Truncate to the requested number of hits\n sections = sections[: request.num_hits]\n\n # Apply LLM document selection if requested\n # num_docs_fed_to_llm_selection specifies how many sections to feed to the LLM for selection\n # The LLM will always try to select TARGET_NUM_SECTIONS_FOR_LLM_SELECTION sections from those fed to it\n # llm_selected_doc_ids will be:\n # - None if LLM selection was not requested or failed\n # - Empty list if LLM selection ran but selected nothing\n # - List of doc IDs if LLM selection succeeded\n run_llm_selection = (\n request.num_docs_fed_to_llm_selection is not None\n and request.num_docs_fed_to_llm_selection >= 1\n )\n llm_selected_doc_ids: list[str] | None = None\n llm_selection_failed = False\n if run_llm_selection and sections:\n try:\n llm = get_default_llm()\n sections_to_evaluate = sections[: request.num_docs_fed_to_llm_selection]\n selected_sections, _ = select_sections_for_expansion(\n sections=sections_to_evaluate,\n user_query=original_query,\n llm=llm,\n max_sections=TARGET_NUM_SECTIONS_FOR_LLM_SELECTION,\n try_to_fill_to_max=True,\n )\n # Extract unique document IDs from selected sections (may be empty)\n llm_selected_doc_ids = list(\n dict.fromkeys(\n section.center_chunk.document_id for section in selected_sections\n )\n )\n logger.debug(\n f\"LLM document selection evaluated {len(sections_to_evaluate)} sections, \"\n f\"selected {len(selected_sections)} sections with doc IDs: {llm_selected_doc_ids}\"\n )\n except Exception as e:\n # Allowing a blanket exception here as this step is not critical and the rest of the results are still valid\n logger.warning(f\"LLM document selection failed: {e}\")\n llm_selection_failed = True\n elif run_llm_selection and not sections:\n # LLM selection requested but no sections to evaluate\n llm_selected_doc_ids = []\n\n # Convert to SearchDocWithContent list, optionally including content\n search_docs = SearchDocWithContent.from_inference_sections(\n sections,\n include_content=request.include_content,\n is_internet=False,\n )\n\n # Yield queries packet\n yield SearchQueriesPacket(all_executed_queries=all_executed_queries)\n\n # Yield docs packet\n yield SearchDocsPacket(search_docs=search_docs)\n\n # Yield LLM selected docs packet if LLM selection was requested\n # - llm_selected_doc_ids is None if selection failed\n # - llm_selected_doc_ids is empty list if no docs were selected\n # - llm_selected_doc_ids is list of IDs if docs were selected\n if run_llm_selection:\n yield LLMSelectedDocsPacket(\n llm_selected_doc_ids=None if llm_selection_failed else llm_selected_doc_ids\n )\n\n\ndef gather_search_stream(\n packets: Generator[\n SearchQueriesPacket\n | SearchDocsPacket\n | LLMSelectedDocsPacket\n | SearchErrorPacket,\n None,\n None,\n ],\n) -> SearchFullResponse:\n \"\"\"\n Aggregate all streaming packets into SearchFullResponse.\n \"\"\"\n all_executed_queries: list[str] = []\n search_docs: list[SearchDocWithContent] = []\n llm_selected_doc_ids: list[str] | None = None\n error: str | None = None\n\n for packet in packets:\n if isinstance(packet, SearchQueriesPacket):\n all_executed_queries = packet.all_executed_queries\n elif isinstance(packet, SearchDocsPacket):\n search_docs = packet.search_docs\n elif isinstance(packet, LLMSelectedDocsPacket):\n llm_selected_doc_ids = packet.llm_selected_doc_ids\n elif isinstance(packet, SearchErrorPacket):\n error = packet.error\n\n return SearchFullResponse(\n all_executed_queries=all_executed_queries,\n search_docs=search_docs,\n doc_selection_reasoning=None,\n llm_selected_doc_ids=llm_selected_doc_ids,\n error=error,\n )\n" }, { "path": "backend/ee/onyx/secondary_llm_flows/__init__.py", "content": "" }, { "path": "backend/ee/onyx/secondary_llm_flows/query_expansion.py", "content": "import re\n\nfrom ee.onyx.prompts.query_expansion import KEYWORD_EXPANSION_PROMPT\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import LanguageModelInput\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Pattern to remove common LLM artifacts: brackets, quotes, list markers, etc.\nCLEANUP_PATTERN = re.compile(r'[\\[\\]\"\\'`]')\n\n\ndef _clean_keyword_line(line: str) -> str:\n \"\"\"Clean a keyword line by removing common LLM artifacts.\n\n Removes brackets, quotes, and other characters that LLMs may accidentally\n include in their output.\n \"\"\"\n # Remove common artifacts\n cleaned = CLEANUP_PATTERN.sub(\"\", line)\n # Remove leading list markers like \"1.\", \"2.\", \"-\", \"*\"\n cleaned = re.sub(r\"^\\s*(?:\\d+[\\.\\)]\\s*|[-*]\\s*)\", \"\", cleaned)\n return cleaned.strip()\n\n\ndef expand_keywords(\n user_query: str,\n llm: LLM,\n) -> list[str]:\n \"\"\"Expand a user query into multiple keyword-only queries for BM25 search.\n\n Uses an LLM to generate keyword-based search queries that capture different\n aspects of the user's search intent. Returns only the expanded queries,\n not the original query.\n\n Args:\n user_query: The original search query from the user\n llm: Language model to use for keyword expansion\n\n Returns:\n List of expanded keyword queries (excluding the original query).\n Returns empty list if expansion fails or produces no useful expansions.\n \"\"\"\n messages: LanguageModelInput = [\n UserMessage(content=KEYWORD_EXPANSION_PROMPT.format(user_query=user_query))\n ]\n\n try:\n response = llm.invoke(\n prompt=messages,\n reasoning_effort=ReasoningEffort.OFF,\n # Limit output - we only expect a few short keyword queries\n max_tokens=150,\n )\n\n content = llm_response_to_string(response).strip()\n\n if not content:\n logger.warning(\"Keyword expansion returned empty response.\")\n return []\n\n # Parse response - each line is a separate keyword query\n # Clean each line to remove LLM artifacts and drop empty lines\n parsed_queries = []\n for line in content.strip().split(\"\\n\"):\n cleaned = _clean_keyword_line(line)\n if cleaned:\n parsed_queries.append(cleaned)\n\n if not parsed_queries:\n logger.warning(\"Keyword expansion parsing returned no queries.\")\n return []\n\n # Filter out duplicates and queries that match the original\n expanded_queries: list[str] = []\n seen_lower: set[str] = {user_query.lower()}\n for query in parsed_queries:\n query_lower = query.lower()\n if query_lower not in seen_lower:\n seen_lower.add(query_lower)\n expanded_queries.append(query)\n\n logger.debug(f\"Keyword expansion generated {len(expanded_queries)} queries\")\n return expanded_queries\n\n except Exception as e:\n logger.warning(f\"Keyword expansion failed: {e}\")\n return []\n" }, { "path": "backend/ee/onyx/secondary_llm_flows/search_flow_classification.py", "content": "from ee.onyx.prompts.search_flow_classification import CHAT_CLASS\nfrom ee.onyx.prompts.search_flow_classification import SEARCH_CHAT_PROMPT\nfrom ee.onyx.prompts.search_flow_classification import SEARCH_CLASS\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import LanguageModelInput\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.timing import log_function_time\n\nlogger = setup_logger()\n\n\n@log_function_time(print_only=True)\ndef classify_is_search_flow(\n query: str,\n llm: LLM,\n) -> bool:\n messages: LanguageModelInput = [\n UserMessage(content=SEARCH_CHAT_PROMPT.format(user_query=query))\n ]\n response = llm.invoke(\n prompt=messages,\n reasoning_effort=ReasoningEffort.OFF,\n # Nothing can happen in the UI until this call finishes so we need to be aggressive with the timeout\n timeout_override=2,\n # Well more than necessary but just to ensure completion and in case it succeeds with classifying but\n # ends up rambling\n max_tokens=20,\n )\n\n content = llm_response_to_string(response).strip().lower()\n if not content:\n logger.warning(\n \"Search flow classification returned empty response; defaulting to chat flow.\"\n )\n return False\n\n # Prefer chat if both appear.\n if CHAT_CLASS in content:\n return False\n if SEARCH_CLASS in content:\n return True\n\n logger.warning(\n \"Search flow classification returned unexpected response; defaulting to chat flow. Response=%r\",\n content,\n )\n return False\n" }, { "path": "backend/ee/onyx/server/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/analytics/api.py", "content": "import datetime\nfrom collections import defaultdict\nfrom typing import List\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.analytics import fetch_assistant_message_analytics\nfrom ee.onyx.db.analytics import fetch_assistant_unique_users\nfrom ee.onyx.db.analytics import fetch_assistant_unique_users_total\nfrom ee.onyx.db.analytics import fetch_onyxbot_analytics\nfrom ee.onyx.db.analytics import fetch_per_user_query_analytics\nfrom ee.onyx.db.analytics import fetch_persona_message_analytics\nfrom ee.onyx.db.analytics import fetch_persona_unique_users\nfrom ee.onyx.db.analytics import fetch_query_analytics\nfrom ee.onyx.db.analytics import user_can_view_assistant_stats\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_user\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\n\nrouter = APIRouter(prefix=\"/analytics\", tags=PUBLIC_API_TAGS)\n\n\n_DEFAULT_LOOKBACK_DAYS = 30\n\n\nclass QueryAnalyticsResponse(BaseModel):\n total_queries: int\n total_likes: int\n total_dislikes: int\n date: datetime.date\n\n\n@router.get(\"/admin/query\")\ndef get_query_analytics(\n start: datetime.datetime | None = None,\n end: datetime.datetime | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[QueryAnalyticsResponse]:\n daily_query_usage_info = fetch_query_analytics(\n start=start\n or (\n datetime.datetime.utcnow() - datetime.timedelta(days=_DEFAULT_LOOKBACK_DAYS)\n ), # default is 30d lookback\n end=end or datetime.datetime.utcnow(),\n db_session=db_session,\n )\n return [\n QueryAnalyticsResponse(\n total_queries=total_queries,\n total_likes=total_likes,\n total_dislikes=total_dislikes,\n date=date,\n )\n for total_queries, total_likes, total_dislikes, date in daily_query_usage_info\n ]\n\n\nclass UserAnalyticsResponse(BaseModel):\n total_active_users: int\n date: datetime.date\n\n\n@router.get(\"/admin/user\")\ndef get_user_analytics(\n start: datetime.datetime | None = None,\n end: datetime.datetime | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[UserAnalyticsResponse]:\n daily_query_usage_info_per_user = fetch_per_user_query_analytics(\n start=start\n or (\n datetime.datetime.utcnow() - datetime.timedelta(days=_DEFAULT_LOOKBACK_DAYS)\n ), # default is 30d lookback\n end=end or datetime.datetime.utcnow(),\n db_session=db_session,\n )\n\n user_analytics: dict[datetime.date, int] = defaultdict(int)\n for __, ___, ____, date, _____ in daily_query_usage_info_per_user:\n user_analytics[date] += 1\n return [\n UserAnalyticsResponse(\n total_active_users=cnt,\n date=date,\n )\n for date, cnt in user_analytics.items()\n ]\n\n\nclass OnyxbotAnalyticsResponse(BaseModel):\n total_queries: int\n auto_resolved: int\n date: datetime.date\n\n\n@router.get(\"/admin/onyxbot\")\ndef get_onyxbot_analytics(\n start: datetime.datetime | None = None,\n end: datetime.datetime | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[OnyxbotAnalyticsResponse]:\n daily_onyxbot_info = fetch_onyxbot_analytics(\n start=start\n or (\n datetime.datetime.utcnow() - datetime.timedelta(days=_DEFAULT_LOOKBACK_DAYS)\n ), # default is 30d lookback\n end=end or datetime.datetime.utcnow(),\n db_session=db_session,\n )\n\n resolution_results = [\n OnyxbotAnalyticsResponse(\n total_queries=total_queries,\n # If it hits negatives, something has gone wrong...\n auto_resolved=max(0, total_queries - total_negatives),\n date=date,\n )\n for total_queries, total_negatives, date in daily_onyxbot_info\n ]\n\n return resolution_results\n\n\nclass PersonaMessageAnalyticsResponse(BaseModel):\n total_messages: int\n date: datetime.date\n persona_id: int\n\n\n@router.get(\"/admin/persona/messages\")\ndef get_persona_messages(\n persona_id: int,\n start: datetime.datetime | None = None,\n end: datetime.datetime | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[PersonaMessageAnalyticsResponse]:\n \"\"\"Fetch daily message counts for a single persona within the given time range.\"\"\"\n start = start or (\n datetime.datetime.utcnow() - datetime.timedelta(days=_DEFAULT_LOOKBACK_DAYS)\n )\n end = end or datetime.datetime.utcnow()\n\n persona_message_counts = []\n for count, date in fetch_persona_message_analytics(\n db_session=db_session,\n persona_id=persona_id,\n start=start,\n end=end,\n ):\n persona_message_counts.append(\n PersonaMessageAnalyticsResponse(\n total_messages=count,\n date=date,\n persona_id=persona_id,\n )\n )\n\n return persona_message_counts\n\n\nclass PersonaUniqueUsersResponse(BaseModel):\n unique_users: int\n date: datetime.date\n persona_id: int\n\n\n@router.get(\"/admin/persona/unique-users\")\ndef get_persona_unique_users(\n persona_id: int,\n start: datetime.datetime,\n end: datetime.datetime,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[PersonaUniqueUsersResponse]:\n \"\"\"Get unique users per day for a single persona.\"\"\"\n unique_user_counts = []\n daily_counts = fetch_persona_unique_users(\n db_session=db_session,\n persona_id=persona_id,\n start=start,\n end=end,\n )\n for count, date in daily_counts:\n unique_user_counts.append(\n PersonaUniqueUsersResponse(\n unique_users=count,\n date=date,\n persona_id=persona_id,\n )\n )\n return unique_user_counts\n\n\nclass AssistantDailyUsageResponse(BaseModel):\n date: datetime.date\n total_messages: int\n total_unique_users: int\n\n\nclass AssistantStatsResponse(BaseModel):\n daily_stats: List[AssistantDailyUsageResponse]\n total_messages: int\n total_unique_users: int\n\n\n@router.get(\"/assistant/{assistant_id}/stats\")\ndef get_assistant_stats(\n assistant_id: int,\n start: datetime.datetime | None = None,\n end: datetime.datetime | None = None,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> AssistantStatsResponse:\n \"\"\"\n Returns daily message and unique user counts for a user's assistant,\n along with the overall total messages and total distinct users.\n \"\"\"\n start = start or (\n datetime.datetime.utcnow() - datetime.timedelta(days=_DEFAULT_LOOKBACK_DAYS)\n )\n end = end or datetime.datetime.utcnow()\n\n if not user_can_view_assistant_stats(db_session, user, assistant_id):\n raise HTTPException(\n status_code=403, detail=\"Not allowed to access this assistant's stats.\"\n )\n\n # Pull daily usage from the DB calls\n messages_data = fetch_assistant_message_analytics(\n db_session, assistant_id, start, end\n )\n unique_users_data = fetch_assistant_unique_users(\n db_session, assistant_id, start, end\n )\n\n # Map each day => (messages, unique_users).\n daily_messages_map = {date: count for count, date in messages_data}\n daily_unique_users_map = {date: count for count, date in unique_users_data}\n all_dates = set(daily_messages_map.keys()) | set(daily_unique_users_map.keys())\n\n # Merge both sets of metrics by date\n daily_results: list[AssistantDailyUsageResponse] = []\n for date in sorted(all_dates):\n daily_results.append(\n AssistantDailyUsageResponse(\n date=date,\n total_messages=daily_messages_map.get(date, 0),\n total_unique_users=daily_unique_users_map.get(date, 0),\n )\n )\n\n # Now pull a single total distinct user count across the entire time range\n total_msgs = sum(d.total_messages for d in daily_results)\n total_users = fetch_assistant_unique_users_total(\n db_session, assistant_id, start, end\n )\n\n return AssistantStatsResponse(\n daily_stats=daily_results,\n total_messages=total_msgs,\n total_unique_users=total_users,\n )\n" }, { "path": "backend/ee/onyx/server/auth_check.py", "content": "from fastapi import FastAPI\n\nfrom onyx.server.auth_check import check_router_auth\nfrom onyx.server.auth_check import PUBLIC_ENDPOINT_SPECS\n\n\nEE_PUBLIC_ENDPOINT_SPECS = PUBLIC_ENDPOINT_SPECS + [\n # SCIM 2.0 service discovery — unauthenticated so IdPs can probe\n # before bearer token configuration is complete\n (\"/scim/v2/ServiceProviderConfig\", {\"GET\"}),\n (\"/scim/v2/ResourceTypes\", {\"GET\"}),\n (\"/scim/v2/Schemas\", {\"GET\"}),\n # needs to be accessible prior to user login\n (\"/enterprise-settings\", {\"GET\"}),\n (\"/enterprise-settings/logo\", {\"GET\"}),\n (\"/enterprise-settings/logotype\", {\"GET\"}),\n (\"/enterprise-settings/custom-analytics-script\", {\"GET\"}),\n # Stripe publishable key is safe to expose publicly\n (\"/tenants/stripe-publishable-key\", {\"GET\"}),\n (\"/admin/billing/stripe-publishable-key\", {\"GET\"}),\n # Proxy endpoints use license-based auth, not user auth\n (\"/proxy/create-checkout-session\", {\"POST\"}),\n (\"/proxy/claim-license\", {\"POST\"}),\n (\"/proxy/create-customer-portal-session\", {\"POST\"}),\n (\"/proxy/billing-information\", {\"GET\"}),\n (\"/proxy/license/{tenant_id}\", {\"GET\"}),\n (\"/proxy/seats/update\", {\"POST\"}),\n]\n\n\ndef check_ee_router_auth(\n application: FastAPI,\n public_endpoint_specs: list[tuple[str, set[str]]] = EE_PUBLIC_ENDPOINT_SPECS,\n) -> None:\n # similar to the open source version of this function, but checking for the EE-only\n # endpoints as well\n check_router_auth(application, public_endpoint_specs)\n" }, { "path": "backend/ee/onyx/server/billing/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/billing/api.py", "content": "\"\"\"Unified Billing API endpoints.\n\nThese endpoints provide Stripe billing functionality for both cloud and\nself-hosted deployments. The service layer routes requests appropriately:\n\n- Self-hosted: Routes through cloud data plane proxy\n Flow: Backend /admin/billing/* → Cloud DP /proxy/* → Control plane\n\n- Cloud (MULTI_TENANT): Routes directly to control plane\n Flow: Backend /admin/billing/* → Control plane\n\nLicense claiming is handled separately by /license/claim endpoint (self-hosted only).\n\nMigration Note (ENG-3533):\nThis /admin/billing/* API replaces the older /tenants/* billing endpoints:\n- /tenants/billing-information -> /admin/billing/billing-information\n- /tenants/create-customer-portal-session -> /admin/billing/create-customer-portal-session\n- /tenants/create-subscription-session -> /admin/billing/create-checkout-session\n- /tenants/stripe-publishable-key -> /admin/billing/stripe-publishable-key\n\nSee: https://linear.app/onyx-app/issue/ENG-3533/migrate-tenantsbilling-adminbilling\n\"\"\"\n\nimport asyncio\n\nimport httpx\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.auth.users import current_admin_user\nfrom ee.onyx.db.license import get_license\nfrom ee.onyx.db.license import get_used_seats\nfrom ee.onyx.server.billing.models import BillingInformationResponse\nfrom ee.onyx.server.billing.models import CreateCheckoutSessionRequest\nfrom ee.onyx.server.billing.models import CreateCheckoutSessionResponse\nfrom ee.onyx.server.billing.models import CreateCustomerPortalSessionRequest\nfrom ee.onyx.server.billing.models import CreateCustomerPortalSessionResponse\nfrom ee.onyx.server.billing.models import SeatUpdateRequest\nfrom ee.onyx.server.billing.models import SeatUpdateResponse\nfrom ee.onyx.server.billing.models import StripePublishableKeyResponse\nfrom ee.onyx.server.billing.models import SubscriptionStatusResponse\nfrom ee.onyx.server.billing.service import (\n create_checkout_session as create_checkout_service,\n)\nfrom ee.onyx.server.billing.service import (\n create_customer_portal_session as create_portal_service,\n)\nfrom ee.onyx.server.billing.service import (\n get_billing_information as get_billing_service,\n)\nfrom ee.onyx.server.billing.service import update_seat_count as update_seat_service\nfrom onyx.auth.users import User\nfrom onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_OVERRIDE\nfrom onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_URL\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.redis.redis_pool import get_shared_redis_client\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/admin/billing\")\n\n# Cache for Stripe publishable key to avoid hitting S3 on every request\n_stripe_publishable_key_cache: str | None = None\n_stripe_key_lock = asyncio.Lock()\n\n# Redis key for billing circuit breaker (self-hosted only)\n# When set, billing requests to Stripe are disabled until user manually retries\nBILLING_CIRCUIT_BREAKER_KEY = \"billing_circuit_open\"\n# Circuit breaker auto-expires after 1 hour (user can manually retry sooner)\nBILLING_CIRCUIT_BREAKER_TTL_SECONDS = 3600\n\n\ndef _is_billing_circuit_open() -> bool:\n \"\"\"Check if the billing circuit breaker is open (self-hosted only).\"\"\"\n if MULTI_TENANT:\n return False\n try:\n redis_client = get_shared_redis_client()\n is_open = bool(redis_client.exists(BILLING_CIRCUIT_BREAKER_KEY))\n logger.debug(\n f\"Circuit breaker check: key={BILLING_CIRCUIT_BREAKER_KEY}, is_open={is_open}\"\n )\n return is_open\n except Exception as e:\n logger.error(f\"Failed to check circuit breaker: {e}\")\n return False\n\n\ndef _open_billing_circuit() -> None:\n \"\"\"Open the billing circuit breaker after a failure (self-hosted only).\"\"\"\n if MULTI_TENANT:\n return\n try:\n redis_client = get_shared_redis_client()\n redis_client.set(\n BILLING_CIRCUIT_BREAKER_KEY,\n \"1\",\n ex=BILLING_CIRCUIT_BREAKER_TTL_SECONDS,\n )\n # Verify it was set\n exists = redis_client.exists(BILLING_CIRCUIT_BREAKER_KEY)\n logger.warning(\n f\"Billing circuit breaker opened (TTL={BILLING_CIRCUIT_BREAKER_TTL_SECONDS}s, \"\n f\"verified={exists}). Stripe billing requests are disabled until manually reset.\"\n )\n except Exception as e:\n logger.error(f\"Failed to open circuit breaker: {e}\")\n\n\ndef _close_billing_circuit() -> None:\n \"\"\"Close the billing circuit breaker (re-enable Stripe requests).\"\"\"\n if MULTI_TENANT:\n return\n try:\n redis_client = get_shared_redis_client()\n redis_client.delete(BILLING_CIRCUIT_BREAKER_KEY)\n logger.info(\n \"Billing circuit breaker closed. Stripe billing requests re-enabled.\"\n )\n except Exception as e:\n logger.error(f\"Failed to close circuit breaker: {e}\")\n\n\ndef _get_license_data(db_session: Session) -> str | None:\n \"\"\"Get license data from database if exists (self-hosted only).\"\"\"\n if MULTI_TENANT:\n return None\n license_record = get_license(db_session)\n return license_record.license_data if license_record else None\n\n\ndef _get_tenant_id() -> str | None:\n \"\"\"Get tenant ID for cloud deployments.\"\"\"\n if MULTI_TENANT:\n return get_current_tenant_id()\n return None\n\n\n@router.post(\"/create-checkout-session\")\nasync def create_checkout_session(\n request: CreateCheckoutSessionRequest | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> CreateCheckoutSessionResponse:\n \"\"\"Create a Stripe checkout session for new subscription or renewal.\n\n For new customers, no license/tenant is required.\n For renewals, existing license (self-hosted) or tenant_id (cloud) is used.\n\n After checkout completion:\n - Self-hosted: Use /license/claim to retrieve the license\n - Cloud: Subscription is automatically activated\n \"\"\"\n license_data = _get_license_data(db_session)\n tenant_id = _get_tenant_id()\n billing_period = request.billing_period if request else \"monthly\"\n seats = request.seats if request else None\n email = request.email if request else None\n\n # Validate that requested seats is not less than current used seats\n if seats is not None:\n used_seats = get_used_seats(tenant_id)\n if seats < used_seats:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n f\"Cannot subscribe with fewer seats than current usage. \"\n f\"You have {used_seats} active users/integrations but requested {seats} seats.\",\n )\n\n # Build redirect URL for after checkout completion\n redirect_url = f\"{WEB_DOMAIN}/admin/billing?checkout=success\"\n\n return await create_checkout_service(\n billing_period=billing_period,\n seats=seats,\n email=email,\n license_data=license_data,\n redirect_url=redirect_url,\n tenant_id=tenant_id,\n )\n\n\n@router.post(\"/create-customer-portal-session\")\nasync def create_customer_portal_session(\n request: CreateCustomerPortalSessionRequest | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> CreateCustomerPortalSessionResponse:\n \"\"\"Create a Stripe customer portal session for managing subscription.\n\n Requires existing license (self-hosted) or active tenant (cloud).\n \"\"\"\n license_data = _get_license_data(db_session)\n tenant_id = _get_tenant_id()\n\n # Self-hosted requires license\n if not MULTI_TENANT and not license_data:\n raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, \"No license found\")\n\n return_url = request.return_url if request else f\"{WEB_DOMAIN}/admin/billing\"\n\n return await create_portal_service(\n license_data=license_data,\n return_url=return_url,\n tenant_id=tenant_id,\n )\n\n\n@router.get(\"/billing-information\")\nasync def get_billing_information(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> BillingInformationResponse | SubscriptionStatusResponse:\n \"\"\"Get billing information for the current subscription.\n\n Returns subscription status and details from Stripe.\n For self-hosted: If the circuit breaker is open (previous failure),\n returns a 503 error without making the request.\n \"\"\"\n license_data = _get_license_data(db_session)\n tenant_id = _get_tenant_id()\n\n # Self-hosted without license = no subscription\n if not MULTI_TENANT and not license_data:\n return SubscriptionStatusResponse(subscribed=False)\n\n # Check circuit breaker (self-hosted only)\n if _is_billing_circuit_open():\n raise OnyxError(\n OnyxErrorCode.SERVICE_UNAVAILABLE,\n \"Stripe connection temporarily disabled. Click 'Connect to Stripe' to retry.\",\n )\n\n try:\n return await get_billing_service(\n license_data=license_data,\n tenant_id=tenant_id,\n )\n except OnyxError as e:\n # Open circuit breaker on connection failures (self-hosted only)\n if e.status_code in (\n OnyxErrorCode.BAD_GATEWAY.status_code,\n OnyxErrorCode.SERVICE_UNAVAILABLE.status_code,\n OnyxErrorCode.GATEWAY_TIMEOUT.status_code,\n ):\n _open_billing_circuit()\n raise\n\n\n@router.post(\"/seats/update\")\nasync def update_seats(\n request: SeatUpdateRequest,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> SeatUpdateResponse:\n \"\"\"Update the seat count for the current subscription.\n\n Handles Stripe proration and license regeneration via control plane.\n For self-hosted, the frontend should call /license/claim after a short delay\n to fetch the regenerated license.\n \"\"\"\n license_data = _get_license_data(db_session)\n tenant_id = _get_tenant_id()\n\n # Self-hosted requires license\n if not MULTI_TENANT and not license_data:\n raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, \"No license found\")\n\n # Validate that new seat count is not less than current used seats\n used_seats = get_used_seats(tenant_id)\n if request.new_seat_count < used_seats:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n f\"Cannot reduce seats below current usage. \"\n f\"You have {used_seats} active users/integrations but requested {request.new_seat_count} seats.\",\n )\n\n # Note: Don't store license here - the control plane may still be processing\n # the subscription update. The frontend should call /license/claim after a\n # short delay to get the freshly generated license.\n return await update_seat_service(\n new_seat_count=request.new_seat_count,\n license_data=license_data,\n tenant_id=tenant_id,\n )\n\n\n@router.get(\"/stripe-publishable-key\")\nasync def get_stripe_publishable_key() -> StripePublishableKeyResponse:\n \"\"\"Fetch the Stripe publishable key.\n\n Priority: env var override (for testing) > S3 bucket (production).\n This endpoint is public (no auth required) since publishable keys are safe to expose.\n The key is cached in memory to avoid hitting S3 on every request.\n \"\"\"\n global _stripe_publishable_key_cache\n\n # Fast path: return cached value without lock\n if _stripe_publishable_key_cache:\n return StripePublishableKeyResponse(\n publishable_key=_stripe_publishable_key_cache\n )\n\n # Use lock to prevent concurrent S3 requests\n async with _stripe_key_lock:\n # Double-check after acquiring lock (another request may have populated cache)\n if _stripe_publishable_key_cache:\n return StripePublishableKeyResponse(\n publishable_key=_stripe_publishable_key_cache\n )\n\n # Check for env var override first (for local testing with pk_test_* keys)\n if STRIPE_PUBLISHABLE_KEY_OVERRIDE:\n key = STRIPE_PUBLISHABLE_KEY_OVERRIDE.strip()\n if not key.startswith(\"pk_\"):\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Invalid Stripe publishable key format\",\n )\n _stripe_publishable_key_cache = key\n return StripePublishableKeyResponse(publishable_key=key)\n\n # Fall back to S3 bucket\n if not STRIPE_PUBLISHABLE_KEY_URL:\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Stripe publishable key is not configured\",\n )\n\n try:\n async with httpx.AsyncClient() as client:\n response = await client.get(STRIPE_PUBLISHABLE_KEY_URL)\n response.raise_for_status()\n key = response.text.strip()\n\n # Validate key format\n if not key.startswith(\"pk_\"):\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Invalid Stripe publishable key format\",\n )\n\n _stripe_publishable_key_cache = key\n return StripePublishableKeyResponse(publishable_key=key)\n except httpx.HTTPError:\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Failed to fetch Stripe publishable key\",\n )\n\n\nclass ResetConnectionResponse(BaseModel):\n success: bool\n message: str\n\n\n@router.post(\"/reset-connection\")\nasync def reset_stripe_connection(\n _: User = Depends(current_admin_user),\n) -> ResetConnectionResponse:\n \"\"\"Reset the Stripe connection circuit breaker.\n\n Called when user clicks \"Connect to Stripe\" to retry after a previous failure.\n This clears the circuit breaker flag, allowing billing requests to proceed again.\n Self-hosted only - cloud deployments don't use the circuit breaker.\n \"\"\"\n if MULTI_TENANT:\n return ResetConnectionResponse(\n success=True,\n message=\"Circuit breaker not applicable for cloud deployments\",\n )\n\n _close_billing_circuit()\n return ResetConnectionResponse(\n success=True,\n message=\"Stripe connection reset. Billing requests re-enabled.\",\n )\n" }, { "path": "backend/ee/onyx/server/billing/models.py", "content": "\"\"\"Pydantic models for the billing API.\"\"\"\n\nfrom datetime import datetime\nfrom typing import Literal\n\nfrom pydantic import BaseModel\n\n\nclass CreateCheckoutSessionRequest(BaseModel):\n \"\"\"Request to create a Stripe checkout session.\"\"\"\n\n billing_period: Literal[\"monthly\", \"annual\"] = \"monthly\"\n seats: int | None = None\n email: str | None = None\n\n\nclass CreateCheckoutSessionResponse(BaseModel):\n \"\"\"Response containing the Stripe checkout session URL.\"\"\"\n\n stripe_checkout_url: str\n\n\nclass CreateCustomerPortalSessionRequest(BaseModel):\n \"\"\"Request to create a Stripe customer portal session.\"\"\"\n\n return_url: str | None = None\n\n\nclass CreateCustomerPortalSessionResponse(BaseModel):\n \"\"\"Response containing the Stripe customer portal URL.\"\"\"\n\n stripe_customer_portal_url: str\n\n\nclass BillingInformationResponse(BaseModel):\n \"\"\"Billing information for the current subscription.\"\"\"\n\n tenant_id: str\n status: str | None = None\n plan_type: str | None = None\n seats: int | None = None\n billing_period: str | None = None\n current_period_start: datetime | None = None\n current_period_end: datetime | None = None\n cancel_at_period_end: bool = False\n canceled_at: datetime | None = None\n trial_start: datetime | None = None\n trial_end: datetime | None = None\n payment_method_enabled: bool = False\n\n\nclass SubscriptionStatusResponse(BaseModel):\n \"\"\"Response when no subscription exists.\"\"\"\n\n subscribed: bool = False\n\n\nclass SeatUpdateRequest(BaseModel):\n \"\"\"Request to update seat count.\"\"\"\n\n new_seat_count: int\n\n\nclass SeatUpdateResponse(BaseModel):\n \"\"\"Response from seat update operation.\"\"\"\n\n success: bool\n current_seats: int\n used_seats: int\n message: str | None = None\n license: str | None = None # Regenerated license (self-hosted stores this)\n\n\nclass StripePublishableKeyResponse(BaseModel):\n \"\"\"Response containing the Stripe publishable key.\"\"\"\n\n publishable_key: str\n" }, { "path": "backend/ee/onyx/server/billing/service.py", "content": "\"\"\"Service layer for billing operations.\n\nThis module provides functions for billing operations that route differently\nbased on deployment type:\n\n- Self-hosted (not MULTI_TENANT): Routes through cloud data plane proxy\n Flow: Self-hosted backend → Cloud DP /proxy/* → Control plane\n\n- Cloud (MULTI_TENANT): Routes directly to control plane\n Flow: Cloud backend → Control plane\n\"\"\"\n\nfrom typing import Literal\n\nimport httpx\n\nfrom ee.onyx.configs.app_configs import CLOUD_DATA_PLANE_URL\nfrom ee.onyx.server.billing.models import BillingInformationResponse\nfrom ee.onyx.server.billing.models import CreateCheckoutSessionResponse\nfrom ee.onyx.server.billing.models import CreateCustomerPortalSessionResponse\nfrom ee.onyx.server.billing.models import SeatUpdateResponse\nfrom ee.onyx.server.billing.models import SubscriptionStatusResponse\nfrom ee.onyx.server.tenants.access import generate_data_plane_token\nfrom onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n# HTTP request timeout for billing service calls\n_REQUEST_TIMEOUT = 30.0\n\n\ndef _get_proxy_headers(license_data: str | None) -> dict[str, str]:\n \"\"\"Build headers for proxy requests (self-hosted).\n\n Self-hosted instances authenticate with their license.\n \"\"\"\n headers = {\"Content-Type\": \"application/json\"}\n if license_data:\n headers[\"Authorization\"] = f\"Bearer {license_data}\"\n return headers\n\n\ndef _get_direct_headers() -> dict[str, str]:\n \"\"\"Build headers for direct control plane requests (cloud).\n\n Cloud instances authenticate with JWT.\n \"\"\"\n token = generate_data_plane_token()\n return {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {token}\",\n }\n\n\ndef _get_base_url() -> str:\n \"\"\"Get the base URL based on deployment type.\"\"\"\n if MULTI_TENANT:\n return CONTROL_PLANE_API_BASE_URL\n return f\"{CLOUD_DATA_PLANE_URL}/proxy\"\n\n\ndef _get_headers(license_data: str | None) -> dict[str, str]:\n \"\"\"Get appropriate headers based on deployment type.\"\"\"\n if MULTI_TENANT:\n return _get_direct_headers()\n return _get_proxy_headers(license_data)\n\n\nasync def _make_billing_request(\n method: Literal[\"GET\", \"POST\"],\n path: str,\n license_data: str | None = None,\n body: dict | None = None,\n params: dict | None = None,\n error_message: str = \"Billing service request failed\",\n) -> dict:\n \"\"\"Make an HTTP request to the billing service.\n\n Consolidates the common HTTP request pattern used by all billing operations.\n\n Args:\n method: HTTP method (GET or POST)\n path: URL path (appended to base URL)\n license_data: License for authentication (self-hosted)\n body: Request body for POST requests\n params: Query parameters for GET requests\n error_message: Default error message if request fails\n\n Returns:\n Response JSON as dict\n\n Raises:\n OnyxError: If request fails\n \"\"\"\n\n base_url = _get_base_url()\n url = f\"{base_url}{path}\"\n headers = _get_headers(license_data)\n\n try:\n async with httpx.AsyncClient(\n timeout=_REQUEST_TIMEOUT, follow_redirects=True\n ) as client:\n if method == \"GET\":\n response = await client.get(url, headers=headers, params=params)\n else:\n response = await client.post(url, headers=headers, json=body)\n\n response.raise_for_status()\n return response.json()\n\n except httpx.HTTPStatusError as e:\n detail = error_message\n try:\n error_data = e.response.json()\n detail = error_data.get(\"detail\", detail)\n except Exception:\n pass\n logger.error(f\"{error_message}: {e.response.status_code} - {detail}\")\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY,\n detail,\n status_code_override=e.response.status_code,\n )\n\n except httpx.RequestError:\n logger.exception(\"Failed to connect to billing service\")\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY, \"Failed to connect to billing service\"\n )\n\n\nasync def create_checkout_session(\n billing_period: str = \"monthly\",\n seats: int | None = None,\n email: str | None = None,\n license_data: str | None = None,\n redirect_url: str | None = None,\n tenant_id: str | None = None,\n) -> CreateCheckoutSessionResponse:\n \"\"\"Create a Stripe checkout session.\n\n Args:\n billing_period: \"monthly\" or \"annual\"\n seats: Number of seats to purchase (optional, uses default if not provided)\n email: Customer email for new subscriptions\n license_data: Existing license for renewals (self-hosted)\n redirect_url: URL to redirect after successful checkout\n tenant_id: Tenant ID (cloud only, for renewals)\n\n Returns:\n CreateCheckoutSessionResponse with checkout URL\n \"\"\"\n body: dict = {\"billing_period\": billing_period}\n if seats is not None:\n body[\"seats\"] = seats\n if email:\n body[\"email\"] = email\n if redirect_url:\n body[\"redirect_url\"] = redirect_url\n if tenant_id and MULTI_TENANT:\n body[\"tenant_id\"] = tenant_id\n\n data = await _make_billing_request(\n method=\"POST\",\n path=\"/create-checkout-session\",\n license_data=license_data,\n body=body,\n error_message=\"Failed to create checkout session\",\n )\n return CreateCheckoutSessionResponse(stripe_checkout_url=data[\"url\"])\n\n\nasync def create_customer_portal_session(\n license_data: str | None = None,\n return_url: str | None = None,\n tenant_id: str | None = None,\n) -> CreateCustomerPortalSessionResponse:\n \"\"\"Create a Stripe customer portal session.\n\n Args:\n license_data: License blob for authentication (self-hosted)\n return_url: URL to return to after portal session\n tenant_id: Tenant ID (cloud only)\n\n Returns:\n CreateCustomerPortalSessionResponse with portal URL\n \"\"\"\n body: dict = {}\n if return_url:\n body[\"return_url\"] = return_url\n if tenant_id and MULTI_TENANT:\n body[\"tenant_id\"] = tenant_id\n\n data = await _make_billing_request(\n method=\"POST\",\n path=\"/create-customer-portal-session\",\n license_data=license_data,\n body=body,\n error_message=\"Failed to create customer portal session\",\n )\n return CreateCustomerPortalSessionResponse(stripe_customer_portal_url=data[\"url\"])\n\n\nasync def get_billing_information(\n license_data: str | None = None,\n tenant_id: str | None = None,\n) -> BillingInformationResponse | SubscriptionStatusResponse:\n \"\"\"Fetch billing information.\n\n Args:\n license_data: License blob for authentication (self-hosted)\n tenant_id: Tenant ID (cloud only)\n\n Returns:\n BillingInformationResponse or SubscriptionStatusResponse if no subscription\n \"\"\"\n params = {}\n if tenant_id and MULTI_TENANT:\n params[\"tenant_id\"] = tenant_id\n\n data = await _make_billing_request(\n method=\"GET\",\n path=\"/billing-information\",\n license_data=license_data,\n params=params or None,\n error_message=\"Failed to fetch billing information\",\n )\n\n # Check if no subscription\n if isinstance(data, dict) and data.get(\"subscribed\") is False:\n return SubscriptionStatusResponse(subscribed=False)\n\n return BillingInformationResponse(**data)\n\n\nasync def update_seat_count(\n new_seat_count: int,\n license_data: str | None = None,\n tenant_id: str | None = None,\n) -> SeatUpdateResponse:\n \"\"\"Update the seat count for the current subscription.\n\n Args:\n new_seat_count: New number of seats\n license_data: License blob for authentication (self-hosted)\n tenant_id: Tenant ID (cloud only)\n\n Returns:\n SeatUpdateResponse with updated seat information\n \"\"\"\n body: dict = {\"new_seat_count\": new_seat_count}\n if tenant_id and MULTI_TENANT:\n body[\"tenant_id\"] = tenant_id\n\n data = await _make_billing_request(\n method=\"POST\",\n path=\"/seats/update\",\n license_data=license_data,\n body=body,\n error_message=\"Failed to update seat count\",\n )\n\n return SeatUpdateResponse(\n success=data.get(\"success\", False),\n current_seats=data.get(\"current_seats\", 0),\n used_seats=data.get(\"used_seats\", 0),\n message=data.get(\"message\"),\n license=data.get(\"license\"),\n )\n" }, { "path": "backend/ee/onyx/server/documents/cc_pair.py", "content": "from datetime import datetime\nfrom http import HTTPStatus\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.background.celery.tasks.doc_permission_syncing.tasks import (\n try_creating_permissions_sync_task,\n)\nfrom ee.onyx.background.celery.tasks.external_group_syncing.tasks import (\n try_creating_external_group_sync_task,\n)\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.background.celery.versioned_apps.client import app as client_app\nfrom onyx.db.connector_credential_pair import (\n get_connector_credential_pair_from_id_for_user,\n)\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.models import StatusResponse\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\nrouter = APIRouter(prefix=\"/manage\")\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/sync-permissions\")\ndef get_cc_pair_latest_sync(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> datetime | None:\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"cc_pair not found for current user's permissions\",\n )\n\n return cc_pair.last_time_perm_sync\n\n\n@router.post(\"/admin/cc-pair/{cc_pair_id}/sync-permissions\")\ndef sync_cc_pair(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[None]:\n \"\"\"Triggers permissions sync on a particular cc_pair immediately\"\"\"\n tenant_id = get_current_tenant_id()\n\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"Connection not found for current user's permissions\",\n )\n\n r = get_redis_client()\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n if redis_connector.permissions.fenced:\n raise HTTPException(\n status_code=HTTPStatus.CONFLICT,\n detail=\"Permissions sync task already in progress.\",\n )\n\n logger.info(\n f\"Permissions sync cc_pair={cc_pair_id} \"\n f\"connector_id={cc_pair.connector_id} \"\n f\"credential_id={cc_pair.credential_id} \"\n f\"{cc_pair.connector.name} connector.\"\n )\n payload_id = try_creating_permissions_sync_task(\n client_app, cc_pair_id, r, tenant_id\n )\n if not payload_id:\n raise HTTPException(\n status_code=HTTPStatus.INTERNAL_SERVER_ERROR,\n detail=\"Permissions sync task creation failed.\",\n )\n\n logger.info(f\"Permissions sync queued: cc_pair={cc_pair_id} id={payload_id}\")\n\n return StatusResponse(\n success=True,\n message=\"Successfully created the permissions sync task.\",\n )\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/sync-groups\")\ndef get_cc_pair_latest_group_sync(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> datetime | None:\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"cc_pair not found for current user's permissions\",\n )\n\n return cc_pair.last_time_external_group_sync\n\n\n@router.post(\"/admin/cc-pair/{cc_pair_id}/sync-groups\")\ndef sync_cc_pair_groups(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[None]:\n \"\"\"Triggers group sync on a particular cc_pair immediately\"\"\"\n tenant_id = get_current_tenant_id()\n\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"Connection not found for current user's permissions\",\n )\n\n r = get_redis_client()\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n if redis_connector.external_group_sync.fenced:\n raise HTTPException(\n status_code=HTTPStatus.CONFLICT,\n detail=\"External group sync task already in progress.\",\n )\n\n logger.info(\n f\"External group sync cc_pair={cc_pair_id} \"\n f\"connector_id={cc_pair.connector_id} \"\n f\"credential_id={cc_pair.credential_id} \"\n f\"{cc_pair.connector.name} connector.\"\n )\n payload_id = try_creating_external_group_sync_task(\n client_app, cc_pair_id, r, tenant_id\n )\n if not payload_id:\n raise HTTPException(\n status_code=HTTPStatus.INTERNAL_SERVER_ERROR,\n detail=\"External group sync task creation failed.\",\n )\n\n logger.info(f\"External group sync queued: cc_pair={cc_pair_id} id={payload_id}\")\n\n return StatusResponse(\n success=True,\n message=\"Successfully created the external group sync task.\",\n )\n" }, { "path": "backend/ee/onyx/server/enterprise_settings/api.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nimport httpx\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Response\nfrom fastapi import status\nfrom fastapi import UploadFile\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.scim import ScimDAL\nfrom ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload\nfrom ee.onyx.server.enterprise_settings.models import EnterpriseSettings\nfrom ee.onyx.server.enterprise_settings.store import get_logo_filename\nfrom ee.onyx.server.enterprise_settings.store import get_logotype_filename\nfrom ee.onyx.server.enterprise_settings.store import load_analytics_script\nfrom ee.onyx.server.enterprise_settings.store import load_settings\nfrom ee.onyx.server.enterprise_settings.store import store_analytics_script\nfrom ee.onyx.server.enterprise_settings.store import store_settings\nfrom ee.onyx.server.enterprise_settings.store import upload_logo\nfrom ee.onyx.server.scim.auth import generate_scim_token\nfrom ee.onyx.server.scim.models import ScimTokenCreate\nfrom ee.onyx.server.scim.models import ScimTokenCreatedResponse\nfrom ee.onyx.server.scim.models import ScimTokenResponse\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_user_with_expired_token\nfrom onyx.auth.users import get_user_manager\nfrom onyx.auth.users import UserManager\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.server.utils import BasicAuthenticationError\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.contextvars import get_current_tenant_id\n\nadmin_router = APIRouter(prefix=\"/admin/enterprise-settings\")\nbasic_router = APIRouter(prefix=\"/enterprise-settings\")\n\nlogger = setup_logger()\n\n\nclass RefreshTokenData(BaseModel):\n access_token: str\n refresh_token: str\n session: dict = Field(..., description=\"Contains session information\")\n userinfo: dict = Field(..., description=\"Contains user information\")\n\n def __init__(self, **data: Any) -> None:\n super().__init__(**data)\n if \"exp\" not in self.session:\n raise ValueError(\"'exp' must be set in the session dictionary\")\n if \"userId\" not in self.userinfo or \"email\" not in self.userinfo:\n raise ValueError(\n \"'userId' and 'email' must be set in the userinfo dictionary\"\n )\n\n\n@basic_router.post(\"/refresh-token\")\nasync def refresh_access_token(\n refresh_token: RefreshTokenData,\n user: User = Depends(current_user_with_expired_token),\n user_manager: UserManager = Depends(get_user_manager),\n) -> None:\n try:\n logger.debug(f\"Received response from Meechum auth URL for user {user.id}\")\n\n # Extract new tokens\n new_access_token = refresh_token.access_token\n new_refresh_token = refresh_token.refresh_token\n\n new_expiry = datetime.fromtimestamp(\n refresh_token.session[\"exp\"] / 1000, tz=timezone.utc\n )\n expires_at_timestamp = int(new_expiry.timestamp())\n\n logger.debug(f\"Access token has been refreshed for user {user.id}\")\n\n await user_manager.oauth_callback(\n oauth_name=\"custom\",\n access_token=new_access_token,\n account_id=refresh_token.userinfo[\"userId\"],\n account_email=refresh_token.userinfo[\"email\"],\n expires_at=expires_at_timestamp,\n refresh_token=new_refresh_token,\n associate_by_email=True,\n )\n\n logger.info(f\"Successfully refreshed tokens for user {user.id}\")\n\n except httpx.HTTPStatusError as e:\n if e.response.status_code == 401:\n logger.warning(f\"Full authentication required for user {user.id}\")\n raise HTTPException(\n status_code=status.HTTP_401_UNAUTHORIZED,\n detail=\"Full authentication required\",\n )\n logger.error(\n f\"HTTP error occurred while refreshing token for user {user.id}: {str(e)}\"\n )\n raise HTTPException(\n status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,\n detail=\"Failed to refresh token\",\n )\n except Exception as e:\n logger.error(\n f\"Unexpected error occurred while refreshing token for user {user.id}: {str(e)}\"\n )\n raise HTTPException(\n status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,\n detail=\"An unexpected error occurred\",\n )\n\n\n@admin_router.put(\"\")\ndef admin_ee_put_settings(\n settings: EnterpriseSettings, _: User = Depends(current_admin_user)\n) -> None:\n store_settings(settings)\n\n\n@basic_router.get(\"\")\ndef ee_fetch_settings() -> EnterpriseSettings:\n if MULTI_TENANT:\n tenant_id = get_current_tenant_id()\n if not tenant_id or tenant_id == POSTGRES_DEFAULT_SCHEMA:\n raise BasicAuthenticationError(detail=\"User must authenticate\")\n\n return load_settings()\n\n\n@admin_router.put(\"/logo\")\ndef put_logo(\n file: UploadFile,\n is_logotype: bool = False,\n _: User = Depends(current_admin_user),\n) -> None:\n upload_logo(file=file, is_logotype=is_logotype)\n\n\ndef fetch_logo_helper(db_session: Session) -> Response: # noqa: ARG001\n try:\n file_store = get_default_file_store()\n onyx_file = file_store.get_file_with_mime_type(get_logo_filename())\n if not onyx_file:\n raise ValueError(\"get_onyx_file returned None!\")\n except Exception:\n logger.exception(\"Faield to fetch logo file\")\n raise HTTPException(\n status_code=404,\n detail=\"No logo file found\",\n )\n else:\n return Response(\n content=onyx_file.data,\n media_type=onyx_file.mime_type,\n headers={\"Cache-Control\": \"no-cache\"},\n )\n\n\ndef fetch_logotype_helper(db_session: Session) -> Response: # noqa: ARG001\n try:\n file_store = get_default_file_store()\n onyx_file = file_store.get_file_with_mime_type(get_logotype_filename())\n if not onyx_file:\n raise ValueError(\"get_onyx_file returned None!\")\n except Exception:\n raise HTTPException(\n status_code=404,\n detail=\"No logotype file found\",\n )\n else:\n return Response(content=onyx_file.data, media_type=onyx_file.mime_type)\n\n\n@basic_router.get(\"/logotype\")\ndef fetch_logotype(db_session: Session = Depends(get_session)) -> Response:\n return fetch_logotype_helper(db_session)\n\n\n@basic_router.get(\"/logo\")\ndef fetch_logo(\n is_logotype: bool = False, db_session: Session = Depends(get_session)\n) -> Response:\n if is_logotype:\n return fetch_logotype_helper(db_session)\n\n return fetch_logo_helper(db_session)\n\n\n@admin_router.put(\"/custom-analytics-script\")\ndef upload_custom_analytics_script(\n script_upload: AnalyticsScriptUpload, _: User = Depends(current_admin_user)\n) -> None:\n try:\n store_analytics_script(script_upload)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n\n@basic_router.get(\"/custom-analytics-script\")\ndef fetch_custom_analytics_script() -> str | None:\n return load_analytics_script()\n\n\n# ---------------------------------------------------------------------------\n# SCIM token management\n# ---------------------------------------------------------------------------\n\n\ndef _get_scim_dal(db_session: Session = Depends(get_session)) -> ScimDAL:\n return ScimDAL(db_session)\n\n\n@admin_router.get(\"/scim/token\")\ndef get_active_scim_token(\n _: User = Depends(current_admin_user),\n dal: ScimDAL = Depends(_get_scim_dal),\n) -> ScimTokenResponse:\n \"\"\"Return the currently active SCIM token's metadata, or 404 if none.\"\"\"\n token = dal.get_active_token()\n if not token:\n raise HTTPException(status_code=404, detail=\"No active SCIM token\")\n\n # Derive the IdP domain from the first synced user as a heuristic.\n idp_domain: str | None = None\n mappings, _total = dal.list_user_mappings(start_index=1, count=1)\n if mappings:\n user = dal.get_user(mappings[0].user_id)\n if user and \"@\" in user.email:\n idp_domain = user.email.rsplit(\"@\", 1)[1]\n\n return ScimTokenResponse(\n id=token.id,\n name=token.name,\n token_display=token.token_display,\n is_active=token.is_active,\n created_at=token.created_at,\n last_used_at=token.last_used_at,\n idp_domain=idp_domain,\n )\n\n\n@admin_router.post(\"/scim/token\", status_code=201)\ndef create_scim_token(\n body: ScimTokenCreate,\n user: User = Depends(current_admin_user),\n dal: ScimDAL = Depends(_get_scim_dal),\n) -> ScimTokenCreatedResponse:\n \"\"\"Create a new SCIM bearer token.\n\n Only one token is active at a time — creating a new token automatically\n revokes all previous tokens. The raw token value is returned exactly once\n in the response; it cannot be retrieved again.\n \"\"\"\n raw_token, hashed_token, token_display = generate_scim_token()\n token = dal.create_token(\n name=body.name,\n hashed_token=hashed_token,\n token_display=token_display,\n created_by_id=user.id,\n )\n dal.commit()\n\n return ScimTokenCreatedResponse(\n id=token.id,\n name=token.name,\n token_display=token.token_display,\n is_active=token.is_active,\n created_at=token.created_at,\n last_used_at=token.last_used_at,\n raw_token=raw_token,\n )\n" }, { "path": "backend/ee/onyx/server/enterprise_settings/models.py", "content": "from enum import Enum\nfrom typing import Any\nfrom typing import List\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\n\nclass NavigationItem(BaseModel):\n link: str\n title: str\n # Right now must be one of the FA icons\n icon: str | None = None\n # NOTE: SVG must not have a width / height specified\n # This is the actual SVG as a string. Done this way to reduce\n # complexity / having to store additional \"logos\" in Postgres\n svg_logo: str | None = None\n\n @classmethod\n def model_validate(cls, *args: Any, **kwargs: Any) -> \"NavigationItem\":\n instance = super().model_validate(*args, **kwargs)\n if bool(instance.icon) == bool(instance.svg_logo):\n raise ValueError(\"Exactly one of fa_icon or svg_logo must be specified\")\n return instance\n\n\nclass LogoDisplayStyle(str, Enum):\n LOGO_AND_NAME = \"logo_and_name\"\n LOGO_ONLY = \"logo_only\"\n NAME_ONLY = \"name_only\"\n\n\nclass EnterpriseSettings(BaseModel):\n \"\"\"General settings that only apply to the Enterprise Edition of Onyx\n\n NOTE: don't put anything sensitive in here, as this is accessible without auth.\"\"\"\n\n application_name: str | None = None\n use_custom_logo: bool = False\n use_custom_logotype: bool = False\n logo_display_style: LogoDisplayStyle | None = None\n\n # custom navigation\n custom_nav_items: List[NavigationItem] = Field(default_factory=list)\n\n # custom Chat components\n two_lines_for_chat_header: bool | None = None\n custom_lower_disclaimer_content: str | None = None\n custom_header_content: str | None = None\n custom_popup_header: str | None = None\n custom_popup_content: str | None = None\n enable_consent_screen: bool | None = None\n consent_screen_prompt: str | None = None\n show_first_visit_notice: bool | None = None\n custom_greeting_message: str | None = None\n\n def check_validity(self) -> None:\n return\n\n\nclass AnalyticsScriptUpload(BaseModel):\n script: str\n secret_key: str\n" }, { "path": "backend/ee/onyx/server/enterprise_settings/store.py", "content": "import os\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import cast\nfrom typing import IO\n\nfrom fastapi import HTTPException\nfrom fastapi import UploadFile\n\nfrom ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload\nfrom ee.onyx.server.enterprise_settings.models import EnterpriseSettings\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import KV_CUSTOM_ANALYTICS_SCRIPT_KEY\nfrom onyx.configs.constants import KV_ENTERPRISE_SETTINGS_KEY\nfrom onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n_LOGO_FILENAME = \"__logo__\"\n_LOGOTYPE_FILENAME = \"__logotype__\"\n\n\ndef load_settings() -> EnterpriseSettings:\n \"\"\"Loads settings data directly from DB. This should be used primarily\n for checking what is actually in the DB, aka for editing and saving back settings.\n\n Runtime settings actually used by the application should be checked with\n load_runtime_settings as defaults may be applied at runtime.\n \"\"\"\n\n dynamic_config_store = get_kv_store()\n try:\n settings = EnterpriseSettings(\n **cast(dict, dynamic_config_store.load(KV_ENTERPRISE_SETTINGS_KEY))\n )\n except KvKeyNotFoundError:\n settings = EnterpriseSettings()\n dynamic_config_store.store(KV_ENTERPRISE_SETTINGS_KEY, settings.model_dump())\n\n return settings\n\n\ndef store_settings(settings: EnterpriseSettings) -> None:\n \"\"\"Stores settings directly to the kv store / db.\"\"\"\n\n get_kv_store().store(KV_ENTERPRISE_SETTINGS_KEY, settings.model_dump())\n\n\ndef load_runtime_settings() -> EnterpriseSettings:\n \"\"\"Loads settings from DB and applies any defaults or transformations for use\n at runtime.\n\n Should not be stored back to the DB.\n \"\"\"\n enterprise_settings = load_settings()\n if not enterprise_settings.application_name:\n enterprise_settings.application_name = ONYX_DEFAULT_APPLICATION_NAME\n\n return enterprise_settings\n\n\n_CUSTOM_ANALYTICS_SECRET_KEY = os.environ.get(\"CUSTOM_ANALYTICS_SECRET_KEY\")\n\n\ndef load_analytics_script() -> str | None:\n dynamic_config_store = get_kv_store()\n try:\n return cast(str, dynamic_config_store.load(KV_CUSTOM_ANALYTICS_SCRIPT_KEY))\n except KvKeyNotFoundError:\n return None\n\n\ndef store_analytics_script(analytics_script_upload: AnalyticsScriptUpload) -> None:\n if (\n not _CUSTOM_ANALYTICS_SECRET_KEY\n or analytics_script_upload.secret_key != _CUSTOM_ANALYTICS_SECRET_KEY\n ):\n raise ValueError(\"Invalid secret key\")\n\n get_kv_store().store(KV_CUSTOM_ANALYTICS_SCRIPT_KEY, analytics_script_upload.script)\n\n\ndef is_valid_file_type(filename: str) -> bool:\n valid_extensions = (\".png\", \".jpg\", \".jpeg\")\n return filename.endswith(valid_extensions)\n\n\ndef guess_file_type(filename: str) -> str:\n if filename.lower().endswith(\".png\"):\n return \"image/png\"\n elif filename.lower().endswith(\".jpg\") or filename.lower().endswith(\".jpeg\"):\n return \"image/jpeg\"\n return \"application/octet-stream\"\n\n\ndef upload_logo(file: UploadFile | str, is_logotype: bool = False) -> bool:\n content: IO[Any]\n\n if isinstance(file, str):\n logger.notice(f\"Uploading logo from local path {file}\")\n if not os.path.isfile(file) or not is_valid_file_type(file):\n logger.error(\n \"Invalid file type- only .png, .jpg, and .jpeg files are allowed\"\n )\n return False\n\n with open(file, \"rb\") as file_handle:\n file_content = file_handle.read()\n content = BytesIO(file_content)\n display_name = file\n file_type = guess_file_type(file)\n\n else:\n logger.notice(\"Uploading logo from uploaded file\")\n if not file.filename or not is_valid_file_type(file.filename):\n raise HTTPException(\n status_code=400,\n detail=\"Invalid file type- only .png, .jpg, and .jpeg files are allowed\",\n )\n content = file.file\n display_name = file.filename\n file_type = file.content_type or \"image/jpeg\"\n\n file_store = get_default_file_store()\n file_store.save_file(\n content=content,\n display_name=display_name,\n file_origin=FileOrigin.OTHER,\n file_type=file_type,\n file_id=_LOGOTYPE_FILENAME if is_logotype else _LOGO_FILENAME,\n )\n return True\n\n\ndef get_logo_filename() -> str:\n return _LOGO_FILENAME\n\n\ndef get_logotype_filename() -> str:\n return _LOGOTYPE_FILENAME\n" }, { "path": "backend/ee/onyx/server/evals/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/evals/api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\n\nfrom ee.onyx.auth.users import current_cloud_superuser\nfrom onyx.background.celery.apps.client import celery_app as client_app\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.models import User\nfrom onyx.evals.models import EvalConfigurationOptions\nfrom onyx.server.evals.models import EvalRunAck\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/evals\")\n\n\n@router.post(\"/eval_run\", response_model=EvalRunAck)\ndef eval_run(\n request: EvalConfigurationOptions,\n user: User = Depends(current_cloud_superuser), # noqa: ARG001\n) -> EvalRunAck:\n \"\"\"\n Run an evaluation with the given message and optional dataset.\n This endpoint requires a valid API key for authentication.\n \"\"\"\n client_app.send_task(\n OnyxCeleryTask.EVAL_RUN_TASK,\n kwargs={\n \"configuration_dict\": request.model_dump(),\n },\n )\n return EvalRunAck(success=True)\n" }, { "path": "backend/ee/onyx/server/features/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/features/hooks/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/features/hooks/api.py", "content": "import httpx\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import Query\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import User\nfrom onyx.db.constants import UNSET\nfrom onyx.db.constants import UnsetType\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.hook import create_hook__no_commit\nfrom onyx.db.hook import delete_hook__no_commit\nfrom onyx.db.hook import get_hook_by_id\nfrom onyx.db.hook import get_hook_execution_logs\nfrom onyx.db.hook import get_hooks\nfrom onyx.db.hook import update_hook__no_commit\nfrom onyx.db.models import Hook\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.hooks.api_dependencies import require_hook_enabled\nfrom onyx.hooks.models import HookCreateRequest\nfrom onyx.hooks.models import HookExecutionRecord\nfrom onyx.hooks.models import HookPointMetaResponse\nfrom onyx.hooks.models import HookResponse\nfrom onyx.hooks.models import HookUpdateRequest\nfrom onyx.hooks.models import HookValidateResponse\nfrom onyx.hooks.models import HookValidateStatus\nfrom onyx.hooks.registry import get_all_specs\nfrom onyx.hooks.registry import get_hook_point_spec\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.url import SSRFException\nfrom onyx.utils.url import validate_outbound_http_url\n\nlogger = setup_logger()\n\n# ---------------------------------------------------------------------------\n# SSRF protection\n# ---------------------------------------------------------------------------\n\n\ndef _check_ssrf_safety(endpoint_url: str) -> None:\n \"\"\"Raise OnyxError if endpoint_url could be used for SSRF.\n\n Delegates to validate_outbound_http_url with https_only=True.\n Uses BAD_GATEWAY so the frontend maps the error to the Endpoint URL field.\n \"\"\"\n try:\n validate_outbound_http_url(endpoint_url, https_only=True)\n except (SSRFException, ValueError) as e:\n raise OnyxError(OnyxErrorCode.BAD_GATEWAY, str(e))\n\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _hook_to_response(hook: Hook, creator_email: str | None = None) -> HookResponse:\n return HookResponse(\n id=hook.id,\n name=hook.name,\n hook_point=hook.hook_point,\n endpoint_url=hook.endpoint_url,\n api_key_masked=(\n hook.api_key.get_value(apply_mask=True) if hook.api_key else None\n ),\n fail_strategy=hook.fail_strategy,\n timeout_seconds=hook.timeout_seconds,\n is_active=hook.is_active,\n is_reachable=hook.is_reachable,\n creator_email=(\n creator_email\n if creator_email is not None\n else (hook.creator.email if hook.creator else None)\n ),\n created_at=hook.created_at,\n updated_at=hook.updated_at,\n )\n\n\ndef _get_hook_or_404(\n db_session: Session,\n hook_id: int,\n include_creator: bool = False,\n) -> Hook:\n hook = get_hook_by_id(\n db_session=db_session,\n hook_id=hook_id,\n include_creator=include_creator,\n )\n if hook is None:\n raise OnyxError(OnyxErrorCode.NOT_FOUND, f\"Hook {hook_id} not found.\")\n return hook\n\n\ndef _raise_for_validation_failure(validation: HookValidateResponse) -> None:\n \"\"\"Raise an appropriate OnyxError for a non-passed validation result.\"\"\"\n if validation.status == HookValidateStatus.auth_failed:\n raise OnyxError(OnyxErrorCode.CREDENTIAL_INVALID, validation.error_message)\n if validation.status == HookValidateStatus.timeout:\n raise OnyxError(\n OnyxErrorCode.GATEWAY_TIMEOUT,\n f\"Endpoint validation failed: {validation.error_message}\",\n )\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY,\n f\"Endpoint validation failed: {validation.error_message}\",\n )\n\n\ndef _validate_endpoint(\n endpoint_url: str,\n api_key: str | None,\n timeout_seconds: float,\n) -> HookValidateResponse:\n \"\"\"Check whether endpoint_url is reachable by sending an empty POST request.\n\n We use POST since hook endpoints expect POST requests. The server will typically\n respond with 4xx (missing/invalid body) — that is fine. Any HTTP response means\n the server is up and routable. A 401/403 response returns auth_failed\n (not reachable — indicates the api_key is invalid).\n\n Timeout handling:\n - Any httpx.TimeoutException (ConnectTimeout, ReadTimeout, WriteTimeout, PoolTimeout) →\n timeout (operator should consider increasing timeout_seconds).\n - All other exceptions → cannot_connect.\n \"\"\"\n _check_ssrf_safety(endpoint_url)\n headers: dict[str, str] = {}\n if api_key:\n headers[\"Authorization\"] = f\"Bearer {api_key}\"\n try:\n with httpx.Client(timeout=timeout_seconds, follow_redirects=False) as client:\n response = client.post(endpoint_url, headers=headers)\n if response.status_code in (401, 403):\n return HookValidateResponse(\n status=HookValidateStatus.auth_failed,\n error_message=f\"Authentication failed (HTTP {response.status_code})\",\n )\n return HookValidateResponse(status=HookValidateStatus.passed)\n except httpx.TimeoutException as exc:\n # Any timeout (connect, read, or write) means the configured timeout_seconds\n # is too low for this endpoint. Report as timeout so the UI directs the user\n # to increase the timeout setting.\n logger.warning(\n \"Hook endpoint validation: timeout for %s\",\n endpoint_url,\n exc_info=exc,\n )\n return HookValidateResponse(\n status=HookValidateStatus.timeout,\n error_message=\"Endpoint timed out — consider increasing timeout_seconds.\",\n )\n except Exception as exc:\n logger.warning(\n \"Hook endpoint validation: connection error for %s\",\n endpoint_url,\n exc_info=exc,\n )\n return HookValidateResponse(\n status=HookValidateStatus.cannot_connect, error_message=str(exc)\n )\n\n\n# ---------------------------------------------------------------------------\n# Routers\n# ---------------------------------------------------------------------------\n\nrouter = APIRouter(prefix=\"/admin/hooks\")\n\n\n# ---------------------------------------------------------------------------\n# Hook endpoints\n# ---------------------------------------------------------------------------\n\n\n@router.get(\"/specs\")\ndef get_hook_point_specs(\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n) -> list[HookPointMetaResponse]:\n return [\n HookPointMetaResponse(\n hook_point=spec.hook_point,\n display_name=spec.display_name,\n description=spec.description,\n docs_url=spec.docs_url,\n input_schema=spec.input_schema,\n output_schema=spec.output_schema,\n default_timeout_seconds=spec.default_timeout_seconds,\n default_fail_strategy=spec.default_fail_strategy,\n fail_hard_description=spec.fail_hard_description,\n )\n for spec in get_all_specs()\n ]\n\n\n@router.get(\"\")\ndef list_hooks(\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> list[HookResponse]:\n hooks = get_hooks(db_session=db_session, include_creator=True)\n return [_hook_to_response(h) for h in hooks]\n\n\n@router.post(\"\")\ndef create_hook(\n req: HookCreateRequest,\n user: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> HookResponse:\n \"\"\"Create a new hook. The endpoint is validated before persisting — creation fails if\n the endpoint cannot be reached or the api_key is invalid. Hooks are created active.\n \"\"\"\n spec = get_hook_point_spec(req.hook_point)\n api_key = req.api_key.get_secret_value() if req.api_key else None\n validation = _validate_endpoint(\n endpoint_url=req.endpoint_url,\n api_key=api_key,\n timeout_seconds=req.timeout_seconds or spec.default_timeout_seconds,\n )\n if validation.status != HookValidateStatus.passed:\n _raise_for_validation_failure(validation)\n\n hook = create_hook__no_commit(\n db_session=db_session,\n name=req.name,\n hook_point=req.hook_point,\n endpoint_url=req.endpoint_url,\n api_key=api_key,\n fail_strategy=req.fail_strategy or spec.default_fail_strategy,\n timeout_seconds=req.timeout_seconds or spec.default_timeout_seconds,\n is_active=True,\n is_reachable=True,\n creator_id=user.id,\n )\n db_session.commit()\n return _hook_to_response(hook, creator_email=user.email)\n\n\n@router.get(\"/{hook_id}\")\ndef get_hook(\n hook_id: int,\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> HookResponse:\n hook = _get_hook_or_404(db_session, hook_id, include_creator=True)\n return _hook_to_response(hook)\n\n\n@router.patch(\"/{hook_id}\")\ndef update_hook(\n hook_id: int,\n req: HookUpdateRequest,\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> HookResponse:\n \"\"\"Update hook fields. If endpoint_url, api_key, or timeout_seconds changes, the\n endpoint is re-validated using the effective values. For active hooks the update is\n rejected on validation failure, keeping live traffic unaffected. For inactive hooks\n the update goes through regardless and is_reachable is updated to reflect the result.\n\n Note: if an active hook's endpoint is currently down, even a timeout_seconds-only\n increase will be rejected. The recovery flow is: deactivate → update → reactivate.\n \"\"\"\n # api_key: UNSET = no change, None = clear, value = update\n api_key: str | None | UnsetType\n if \"api_key\" not in req.model_fields_set:\n api_key = UNSET\n elif req.api_key is None:\n api_key = None\n else:\n api_key = req.api_key.get_secret_value()\n\n endpoint_url_changing = \"endpoint_url\" in req.model_fields_set\n api_key_changing = not isinstance(api_key, UnsetType)\n timeout_changing = \"timeout_seconds\" in req.model_fields_set\n\n validated_is_reachable: bool | None = None\n if endpoint_url_changing or api_key_changing or timeout_changing:\n existing = _get_hook_or_404(db_session, hook_id)\n effective_url: str = (\n req.endpoint_url if endpoint_url_changing else existing.endpoint_url # type: ignore[assignment] # endpoint_url is required on create and cannot be cleared on update\n )\n effective_api_key: str | None = (\n (api_key if not isinstance(api_key, UnsetType) else None)\n if api_key_changing\n else (\n existing.api_key.get_value(apply_mask=False)\n if existing.api_key\n else None\n )\n )\n effective_timeout: float = (\n req.timeout_seconds if timeout_changing else existing.timeout_seconds # type: ignore[assignment] # req.timeout_seconds is non-None when timeout_changing (validated by HookUpdateRequest)\n )\n validation = _validate_endpoint(\n endpoint_url=effective_url,\n api_key=effective_api_key,\n timeout_seconds=effective_timeout,\n )\n if existing.is_active and validation.status != HookValidateStatus.passed:\n _raise_for_validation_failure(validation)\n validated_is_reachable = validation.status == HookValidateStatus.passed\n\n hook = update_hook__no_commit(\n db_session=db_session,\n hook_id=hook_id,\n name=req.name,\n endpoint_url=(req.endpoint_url if endpoint_url_changing else UNSET),\n api_key=api_key,\n fail_strategy=req.fail_strategy,\n timeout_seconds=req.timeout_seconds,\n is_reachable=validated_is_reachable,\n include_creator=True,\n )\n db_session.commit()\n return _hook_to_response(hook)\n\n\n@router.delete(\"/{hook_id}\")\ndef delete_hook(\n hook_id: int,\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> None:\n delete_hook__no_commit(db_session=db_session, hook_id=hook_id)\n db_session.commit()\n\n\n@router.post(\"/{hook_id}/activate\")\ndef activate_hook(\n hook_id: int,\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> HookResponse:\n hook = _get_hook_or_404(db_session, hook_id)\n if not hook.endpoint_url:\n raise OnyxError(\n OnyxErrorCode.INVALID_INPUT, \"Hook has no endpoint URL configured.\"\n )\n\n api_key = hook.api_key.get_value(apply_mask=False) if hook.api_key else None\n validation = _validate_endpoint(\n endpoint_url=hook.endpoint_url,\n api_key=api_key,\n timeout_seconds=hook.timeout_seconds,\n )\n if validation.status != HookValidateStatus.passed:\n # Persist is_reachable=False in a separate session so the request\n # session has no commits on the failure path and the transaction\n # boundary stays clean.\n if hook.is_reachable is not False:\n with get_session_with_current_tenant() as side_session:\n update_hook__no_commit(\n db_session=side_session, hook_id=hook_id, is_reachable=False\n )\n side_session.commit()\n _raise_for_validation_failure(validation)\n\n hook = update_hook__no_commit(\n db_session=db_session,\n hook_id=hook_id,\n is_active=True,\n is_reachable=True,\n include_creator=True,\n )\n db_session.commit()\n return _hook_to_response(hook)\n\n\n@router.post(\"/{hook_id}/validate\")\ndef validate_hook(\n hook_id: int,\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> HookValidateResponse:\n hook = _get_hook_or_404(db_session, hook_id)\n if not hook.endpoint_url:\n raise OnyxError(\n OnyxErrorCode.INVALID_INPUT, \"Hook has no endpoint URL configured.\"\n )\n\n api_key = hook.api_key.get_value(apply_mask=False) if hook.api_key else None\n validation = _validate_endpoint(\n endpoint_url=hook.endpoint_url,\n api_key=api_key,\n timeout_seconds=hook.timeout_seconds,\n )\n validation_passed = validation.status == HookValidateStatus.passed\n if hook.is_reachable != validation_passed:\n update_hook__no_commit(\n db_session=db_session, hook_id=hook_id, is_reachable=validation_passed\n )\n db_session.commit()\n return validation\n\n\n@router.post(\"/{hook_id}/deactivate\")\ndef deactivate_hook(\n hook_id: int,\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> HookResponse:\n hook = update_hook__no_commit(\n db_session=db_session,\n hook_id=hook_id,\n is_active=False,\n include_creator=True,\n )\n db_session.commit()\n return _hook_to_response(hook)\n\n\n# ---------------------------------------------------------------------------\n# Execution log endpoints\n# ---------------------------------------------------------------------------\n\n\n@router.get(\"/{hook_id}/execution-logs\")\ndef list_hook_execution_logs(\n hook_id: int,\n limit: int = Query(default=10, ge=1, le=100),\n _: User = Depends(current_admin_user),\n _hook_enabled: None = Depends(require_hook_enabled),\n db_session: Session = Depends(get_session),\n) -> list[HookExecutionRecord]:\n _get_hook_or_404(db_session, hook_id)\n logs = get_hook_execution_logs(db_session=db_session, hook_id=hook_id, limit=limit)\n return [\n HookExecutionRecord(\n error_message=log.error_message,\n status_code=log.status_code,\n duration_ms=log.duration_ms,\n created_at=log.created_at,\n )\n for log in logs\n ]\n" }, { "path": "backend/ee/onyx/server/license/api.py", "content": "\"\"\"License API endpoints for self-hosted deployments.\n\nThese endpoints allow self-hosted Onyx instances to:\n1. Claim a license after Stripe checkout (via cloud data plane proxy)\n2. Upload a license file manually (for air-gapped deployments)\n3. View license status and seat usage\n4. Refresh/delete the local license\n\nNOTE: Cloud (MULTI_TENANT) deployments do NOT use these endpoints.\nCloud licensing is managed via the control plane and gated_tenants Redis key.\n\"\"\"\n\nimport requests\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import File\nfrom fastapi import UploadFile\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.auth.users import current_admin_user\nfrom ee.onyx.configs.app_configs import CLOUD_DATA_PLANE_URL\nfrom ee.onyx.db.license import delete_license as db_delete_license\nfrom ee.onyx.db.license import get_license\nfrom ee.onyx.db.license import get_license_metadata\nfrom ee.onyx.db.license import invalidate_license_cache\nfrom ee.onyx.db.license import refresh_license_cache\nfrom ee.onyx.db.license import update_license_cache\nfrom ee.onyx.db.license import upsert_license\nfrom ee.onyx.server.license.models import LicenseResponse\nfrom ee.onyx.server.license.models import LicenseSource\nfrom ee.onyx.server.license.models import LicenseStatusResponse\nfrom ee.onyx.server.license.models import LicenseUploadResponse\nfrom ee.onyx.server.license.models import SeatUsageResponse\nfrom ee.onyx.utils.license import verify_license_signature\nfrom onyx.auth.users import User\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/license\")\n\n# PEM-style delimiters used in license file format\n_PEM_BEGIN = \"-----BEGIN ONYX LICENSE-----\"\n_PEM_END = \"-----END ONYX LICENSE-----\"\n\n\ndef _strip_pem_delimiters(content: str) -> str:\n \"\"\"Strip PEM-style delimiters from license content if present.\"\"\"\n content = content.strip()\n if content.startswith(_PEM_BEGIN) and content.endswith(_PEM_END):\n # Remove first and last lines (the delimiters)\n lines = content.split(\"\\n\")\n return \"\\n\".join(lines[1:-1]).strip()\n return content\n\n\n@router.get(\"\")\nasync def get_license_status(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> LicenseStatusResponse:\n \"\"\"Get current license status and seat usage.\"\"\"\n metadata = get_license_metadata(db_session)\n\n if not metadata:\n return LicenseStatusResponse(has_license=False)\n\n return LicenseStatusResponse(\n has_license=True,\n seats=metadata.seats,\n used_seats=metadata.used_seats,\n plan_type=metadata.plan_type,\n issued_at=metadata.issued_at,\n expires_at=metadata.expires_at,\n grace_period_end=metadata.grace_period_end,\n status=metadata.status,\n source=metadata.source,\n )\n\n\n@router.get(\"/seats\")\nasync def get_seat_usage(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> SeatUsageResponse:\n \"\"\"Get detailed seat usage information.\"\"\"\n metadata = get_license_metadata(db_session)\n\n if not metadata:\n return SeatUsageResponse(\n total_seats=0,\n used_seats=0,\n available_seats=0,\n )\n\n return SeatUsageResponse(\n total_seats=metadata.seats,\n used_seats=metadata.used_seats,\n available_seats=max(0, metadata.seats - metadata.used_seats),\n )\n\n\n@router.post(\"/claim\")\nasync def claim_license(\n session_id: str | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> LicenseResponse:\n \"\"\"\n Claim a license from the control plane (self-hosted only).\n\n Two modes:\n 1. With session_id: After Stripe checkout, exchange session_id for license\n 2. Without session_id: Re-claim using existing license for auth\n\n Use without session_id after:\n - Updating seats via the billing API\n - Returning from the Stripe customer portal\n - Any operation that regenerates the license on control plane\n Claim a license from the control plane (self-hosted only).\n\n Two modes:\n 1. With session_id: After Stripe checkout, exchange session_id for license\n 2. Without session_id: Re-claim using existing license for auth\n \"\"\"\n if MULTI_TENANT:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"License claiming is only available for self-hosted deployments\",\n )\n\n try:\n if session_id:\n # Claim license after checkout using session_id\n url = f\"{CLOUD_DATA_PLANE_URL}/proxy/claim-license\"\n response = requests.post(\n url,\n json={\"session_id\": session_id},\n headers={\"Content-Type\": \"application/json\"},\n timeout=30,\n )\n else:\n # Re-claim using existing license for auth\n metadata = get_license_metadata(db_session)\n if not metadata or not metadata.tenant_id:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"No license found. Provide session_id after checkout.\",\n )\n\n license_row = get_license(db_session)\n if not license_row or not license_row.license_data:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"No license found in database\",\n )\n\n url = f\"{CLOUD_DATA_PLANE_URL}/proxy/license/{metadata.tenant_id}\"\n response = requests.get(\n url,\n headers={\n \"Authorization\": f\"Bearer {license_row.license_data}\",\n \"Content-Type\": \"application/json\",\n },\n timeout=30,\n )\n\n response.raise_for_status()\n\n data = response.json()\n license_data = data.get(\"license\")\n\n if not license_data:\n raise OnyxError(OnyxErrorCode.NOT_FOUND, \"No license in response\")\n\n # Verify signature before persisting\n payload = verify_license_signature(license_data)\n\n # Store in DB\n upsert_license(db_session, license_data)\n\n try:\n update_license_cache(payload, source=LicenseSource.AUTO_FETCH)\n except Exception as cache_error:\n logger.warning(f\"Failed to update license cache: {cache_error}\")\n\n logger.info(\n f\"License claimed: seats={payload.seats}, expires={payload.expires_at.date()}\"\n )\n return LicenseResponse(success=True, license=payload)\n\n except requests.HTTPError as e:\n status_code = e.response.status_code if e.response is not None else 502\n detail = \"Failed to claim license\"\n try:\n error_data = e.response.json() if e.response is not None else {}\n detail = error_data.get(\"detail\", detail)\n except Exception:\n pass\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY, detail, status_code_override=status_code\n )\n except ValueError as e:\n raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, str(e))\n except requests.RequestException:\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY, \"Failed to connect to license server\"\n )\n\n\n@router.post(\"/upload\")\nasync def upload_license(\n license_file: UploadFile = File(...),\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> LicenseUploadResponse:\n \"\"\"\n Upload a license file manually (self-hosted only).\n\n Used for air-gapped deployments where the cloud data plane is not accessible.\n The license file must be cryptographically signed by Onyx.\n \"\"\"\n if MULTI_TENANT:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"License upload is only available for self-hosted deployments\",\n )\n\n try:\n content = await license_file.read()\n license_data = content.decode(\"utf-8\").strip()\n # Strip PEM-style delimiters if present (used in .lic file format)\n license_data = _strip_pem_delimiters(license_data)\n # Remove any stray whitespace/newlines from user input\n license_data = license_data.strip()\n except UnicodeDecodeError:\n raise OnyxError(OnyxErrorCode.INVALID_INPUT, \"Invalid license file format\")\n\n # Verify cryptographic signature - this is the only validation needed\n # The license's tenant_id identifies the customer in control plane, not locally\n try:\n payload = verify_license_signature(license_data)\n except ValueError as e:\n raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, str(e))\n\n # Persist to DB and update cache\n upsert_license(db_session, license_data)\n\n try:\n update_license_cache(payload, source=LicenseSource.MANUAL_UPLOAD)\n except Exception as cache_error:\n logger.warning(f\"Failed to update license cache: {cache_error}\")\n\n return LicenseUploadResponse(\n success=True,\n message=f\"License uploaded successfully. {payload.seats} seats, expires {payload.expires_at.date()}\",\n )\n\n\n@router.post(\"/refresh\")\nasync def refresh_license_cache_endpoint(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> LicenseStatusResponse:\n \"\"\"\n Force refresh the license cache from the local database.\n\n Useful after manual database changes or to verify license validity.\n Does NOT fetch from control plane - use /claim for that.\n \"\"\"\n metadata = refresh_license_cache(db_session)\n\n if not metadata:\n return LicenseStatusResponse(has_license=False)\n\n return LicenseStatusResponse(\n has_license=True,\n seats=metadata.seats,\n used_seats=metadata.used_seats,\n plan_type=metadata.plan_type,\n issued_at=metadata.issued_at,\n expires_at=metadata.expires_at,\n grace_period_end=metadata.grace_period_end,\n status=metadata.status,\n source=metadata.source,\n )\n\n\n@router.delete(\"\")\nasync def delete_license(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> dict[str, bool]:\n \"\"\"\n Delete the current license.\n\n Admin only - removes license from database and invalidates cache.\n \"\"\"\n if MULTI_TENANT:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"License deletion is only available for self-hosted deployments\",\n )\n\n try:\n invalidate_license_cache()\n except Exception as cache_error:\n logger.warning(f\"Failed to invalidate license cache: {cache_error}\")\n\n deleted = db_delete_license(db_session)\n\n return {\"deleted\": deleted}\n" }, { "path": "backend/ee/onyx/server/license/models.py", "content": "from datetime import datetime\nfrom enum import Enum\n\nfrom pydantic import BaseModel\n\nfrom onyx.server.settings.models import ApplicationStatus\n\n\nclass PlanType(str, Enum):\n MONTHLY = \"monthly\"\n ANNUAL = \"annual\"\n\n\nclass LicenseSource(str, Enum):\n AUTO_FETCH = \"auto_fetch\"\n MANUAL_UPLOAD = \"manual_upload\"\n\n\nclass LicensePayload(BaseModel):\n \"\"\"The payload portion of a signed license.\"\"\"\n\n version: str\n tenant_id: str\n organization_name: str | None = None\n issued_at: datetime\n expires_at: datetime\n seats: int\n plan_type: PlanType\n billing_cycle: str | None = None\n grace_period_days: int = 30\n stripe_subscription_id: str | None = None\n stripe_customer_id: str | None = None\n\n\nclass LicenseData(BaseModel):\n \"\"\"Full signed license structure.\"\"\"\n\n payload: LicensePayload\n signature: str\n\n\nclass LicenseMetadata(BaseModel):\n \"\"\"Cached license metadata stored in Redis.\"\"\"\n\n tenant_id: str\n organization_name: str | None = None\n seats: int\n used_seats: int\n plan_type: PlanType\n issued_at: datetime\n expires_at: datetime\n grace_period_end: datetime | None = None\n status: ApplicationStatus\n source: LicenseSource | None = None\n stripe_subscription_id: str | None = None\n\n\nclass LicenseStatusResponse(BaseModel):\n \"\"\"Response for license status API.\"\"\"\n\n has_license: bool\n seats: int = 0\n used_seats: int = 0\n plan_type: PlanType | None = None\n issued_at: datetime | None = None\n expires_at: datetime | None = None\n grace_period_end: datetime | None = None\n status: ApplicationStatus | None = None\n source: LicenseSource | None = None\n\n\nclass LicenseResponse(BaseModel):\n \"\"\"Response after license fetch/upload.\"\"\"\n\n success: bool\n message: str | None = None\n license: LicensePayload | None = None\n\n\nclass LicenseUploadResponse(BaseModel):\n \"\"\"Response after license upload.\"\"\"\n\n success: bool\n message: str | None = None\n\n\nclass SeatUsageResponse(BaseModel):\n \"\"\"Response for seat usage API.\"\"\"\n\n total_seats: int\n used_seats: int\n available_seats: int\n" }, { "path": "backend/ee/onyx/server/manage/standard_answer.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.standard_answer import fetch_standard_answer\nfrom ee.onyx.db.standard_answer import fetch_standard_answer_categories\nfrom ee.onyx.db.standard_answer import fetch_standard_answer_category\nfrom ee.onyx.db.standard_answer import fetch_standard_answers\nfrom ee.onyx.db.standard_answer import insert_standard_answer\nfrom ee.onyx.db.standard_answer import insert_standard_answer_category\nfrom ee.onyx.db.standard_answer import remove_standard_answer\nfrom ee.onyx.db.standard_answer import update_standard_answer\nfrom ee.onyx.db.standard_answer import update_standard_answer_category\nfrom onyx.auth.users import current_admin_user\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.server.manage.models import StandardAnswer\nfrom onyx.server.manage.models import StandardAnswerCategory\nfrom onyx.server.manage.models import StandardAnswerCategoryCreationRequest\nfrom onyx.server.manage.models import StandardAnswerCreationRequest\n\nrouter = APIRouter(prefix=\"/manage\")\n\n\n@router.post(\"/admin/standard-answer\")\ndef create_standard_answer(\n standard_answer_creation_request: StandardAnswerCreationRequest,\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> StandardAnswer:\n standard_answer_model = insert_standard_answer(\n keyword=standard_answer_creation_request.keyword,\n answer=standard_answer_creation_request.answer,\n category_ids=standard_answer_creation_request.categories,\n match_regex=standard_answer_creation_request.match_regex,\n match_any_keywords=standard_answer_creation_request.match_any_keywords,\n db_session=db_session,\n )\n return StandardAnswer.from_model(standard_answer_model)\n\n\n@router.get(\"/admin/standard-answer\")\ndef list_standard_answers(\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> list[StandardAnswer]:\n standard_answer_models = fetch_standard_answers(db_session=db_session)\n return [\n StandardAnswer.from_model(standard_answer_model)\n for standard_answer_model in standard_answer_models\n ]\n\n\n@router.patch(\"/admin/standard-answer/{standard_answer_id}\")\ndef patch_standard_answer(\n standard_answer_id: int,\n standard_answer_creation_request: StandardAnswerCreationRequest,\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> StandardAnswer:\n existing_standard_answer = fetch_standard_answer(\n standard_answer_id=standard_answer_id,\n db_session=db_session,\n )\n\n if existing_standard_answer is None:\n raise HTTPException(status_code=404, detail=\"Standard answer not found\")\n\n standard_answer_model = update_standard_answer(\n standard_answer_id=standard_answer_id,\n keyword=standard_answer_creation_request.keyword,\n answer=standard_answer_creation_request.answer,\n category_ids=standard_answer_creation_request.categories,\n match_regex=standard_answer_creation_request.match_regex,\n match_any_keywords=standard_answer_creation_request.match_any_keywords,\n db_session=db_session,\n )\n return StandardAnswer.from_model(standard_answer_model)\n\n\n@router.delete(\"/admin/standard-answer/{standard_answer_id}\")\ndef delete_standard_answer(\n standard_answer_id: int,\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> None:\n return remove_standard_answer(\n standard_answer_id=standard_answer_id,\n db_session=db_session,\n )\n\n\n@router.post(\"/admin/standard-answer/category\")\ndef create_standard_answer_category(\n standard_answer_category_creation_request: StandardAnswerCategoryCreationRequest,\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> StandardAnswerCategory:\n standard_answer_category_model = insert_standard_answer_category(\n category_name=standard_answer_category_creation_request.name,\n db_session=db_session,\n )\n return StandardAnswerCategory.from_model(standard_answer_category_model)\n\n\n@router.get(\"/admin/standard-answer/category\")\ndef list_standard_answer_categories(\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> list[StandardAnswerCategory]:\n standard_answer_category_models = fetch_standard_answer_categories(\n db_session=db_session\n )\n return [\n StandardAnswerCategory.from_model(standard_answer_category_model)\n for standard_answer_category_model in standard_answer_category_models\n ]\n\n\n@router.patch(\"/admin/standard-answer/category/{standard_answer_category_id}\")\ndef patch_standard_answer_category(\n standard_answer_category_id: int,\n standard_answer_category_creation_request: StandardAnswerCategoryCreationRequest,\n db_session: Session = Depends(get_session),\n _: User = Depends(current_admin_user),\n) -> StandardAnswerCategory:\n existing_standard_answer_category = fetch_standard_answer_category(\n standard_answer_category_id=standard_answer_category_id,\n db_session=db_session,\n )\n\n if existing_standard_answer_category is None:\n raise HTTPException(\n status_code=404, detail=\"Standard answer category not found\"\n )\n\n standard_answer_category_model = update_standard_answer_category(\n standard_answer_category_id=standard_answer_category_id,\n category_name=standard_answer_category_creation_request.name,\n db_session=db_session,\n )\n return StandardAnswerCategory.from_model(standard_answer_category_model)\n" }, { "path": "backend/ee/onyx/server/middleware/license_enforcement.py", "content": "\"\"\"Middleware to enforce license status for SELF-HOSTED deployments only.\n\nNOTE: This middleware is NOT used for multi-tenant (cloud) deployments.\nMulti-tenant gating is handled separately by the control plane via the\n/tenants/product-gating endpoint and is_tenant_gated() checks.\n\nIMPORTANT: Mutual Exclusivity with ENTERPRISE_EDITION_ENABLED\n============================================================\nThis middleware is controlled by LICENSE_ENFORCEMENT_ENABLED env var.\nIt works alongside the legacy ENTERPRISE_EDITION_ENABLED system:\n\n- LICENSE_ENFORCEMENT_ENABLED=false (default):\n Middleware is disabled. EE features are controlled solely by\n ENTERPRISE_EDITION_ENABLED. This preserves legacy behavior.\n\n- LICENSE_ENFORCEMENT_ENABLED=true:\n Middleware actively enforces license status. EE features require\n a valid license, regardless of ENTERPRISE_EDITION_ENABLED.\n\nEventually, ENTERPRISE_EDITION_ENABLED will be removed and license\nenforcement will be the only mechanism for gating EE features.\n\nLicense Enforcement States (when enabled)\n=========================================\nFor self-hosted deployments:\n\n1. No license (never subscribed):\n - Allow community features (basic connectors, search, chat)\n - Block EE-only features (analytics, user groups, etc.)\n\n2. GATED_ACCESS (fully expired):\n - Block all routes except billing/auth/license\n - User must renew subscription to continue\n\n3. Valid license (ACTIVE, GRACE_PERIOD, PAYMENT_REMINDER):\n - Full access to all EE features\n - Seat limits enforced\n - GRACE_PERIOD/PAYMENT_REMINDER are for notifications only, not blocking\n\"\"\"\n\nimport logging\nfrom collections.abc import Awaitable\nfrom collections.abc import Callable\n\nfrom fastapi import FastAPI\nfrom fastapi import Request\nfrom fastapi import Response\nfrom fastapi.responses import JSONResponse\nfrom sqlalchemy.exc import SQLAlchemyError\n\nfrom ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED\nfrom ee.onyx.configs.license_enforcement_config import EE_ONLY_PATH_PREFIXES\nfrom ee.onyx.configs.license_enforcement_config import (\n LICENSE_ENFORCEMENT_ALLOWED_PREFIXES,\n)\nfrom ee.onyx.db.license import get_cached_license_metadata\nfrom ee.onyx.db.license import refresh_license_cache\nfrom onyx.cache.interface import CACHE_TRANSIENT_ERRORS\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.server.settings.models import ApplicationStatus\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\ndef _is_path_allowed(path: str) -> bool:\n \"\"\"Check if path is in allowlist (prefix match).\"\"\"\n return any(\n path.startswith(prefix) for prefix in LICENSE_ENFORCEMENT_ALLOWED_PREFIXES\n )\n\n\ndef _is_ee_only_path(path: str) -> bool:\n \"\"\"Check if path requires EE license (prefix match).\"\"\"\n return any(path.startswith(prefix) for prefix in EE_ONLY_PATH_PREFIXES)\n\n\ndef add_license_enforcement_middleware(\n app: FastAPI, logger: logging.LoggerAdapter\n) -> None:\n logger.info(\"License enforcement middleware registered\")\n\n @app.middleware(\"http\")\n async def enforce_license(\n request: Request, call_next: Callable[[Request], Awaitable[Response]]\n ) -> Response:\n \"\"\"Block requests when license is expired/gated.\"\"\"\n if not LICENSE_ENFORCEMENT_ENABLED:\n return await call_next(request)\n\n path = request.url.path\n if path.startswith(\"/api\"):\n path = path[4:]\n\n if _is_path_allowed(path):\n return await call_next(request)\n\n is_gated = False\n tenant_id = get_current_tenant_id()\n\n try:\n metadata = get_cached_license_metadata(tenant_id)\n\n # If no cached metadata, check database (cache may have been cleared)\n if not metadata:\n logger.debug(\n \"[license_enforcement] No cached license, checking database...\"\n )\n try:\n with get_session_with_current_tenant() as db_session:\n metadata = refresh_license_cache(db_session, tenant_id)\n if metadata:\n logger.info(\n \"[license_enforcement] Loaded license from database\"\n )\n except SQLAlchemyError as db_error:\n logger.warning(\n f\"[license_enforcement] Failed to check database for license: {db_error}\"\n )\n\n if metadata:\n # User HAS a license (current or expired)\n if metadata.status == ApplicationStatus.GATED_ACCESS:\n # License fully expired - gate the user\n # Note: GRACE_PERIOD and PAYMENT_REMINDER are for notifications only,\n # they don't block access\n is_gated = True\n else:\n # License is active - check seat limit\n # used_seats in cache is kept accurate via invalidation\n # when users are added/removed\n if metadata.used_seats > metadata.seats:\n logger.info(\n f\"[license_enforcement] Blocking request: \"\n f\"seat limit exceeded ({metadata.used_seats}/{metadata.seats})\"\n )\n return JSONResponse(\n status_code=402,\n content={\n \"detail\": {\n \"error\": \"seat_limit_exceeded\",\n \"message\": f\"Seat limit exceeded: {metadata.used_seats} of {metadata.seats} seats used.\",\n \"used_seats\": metadata.used_seats,\n \"seats\": metadata.seats,\n }\n },\n )\n else:\n # No license in cache OR database = never subscribed\n # Allow community features, but block EE-only features\n if _is_ee_only_path(path):\n logger.info(\n f\"[license_enforcement] Blocking EE-only path (no license): {path}\"\n )\n return JSONResponse(\n status_code=402,\n content={\n \"detail\": {\n \"error\": \"enterprise_license_required\",\n \"message\": \"This feature requires an Enterprise license. \"\n \"Please upgrade to access this functionality.\",\n }\n },\n )\n logger.debug(\n \"[license_enforcement] No license, allowing community features\"\n )\n is_gated = False\n except CACHE_TRANSIENT_ERRORS as e:\n logger.warning(f\"Failed to check license metadata: {e}\")\n # Fail open - don't block users due to cache connectivity issues\n is_gated = False\n\n if is_gated:\n logger.info(\n f\"[license_enforcement] Blocking request (license expired): {path}\"\n )\n\n return JSONResponse(\n status_code=402,\n content={\n \"detail\": {\n \"error\": \"license_expired\",\n \"message\": \"Your subscription has expired. Please update your billing.\",\n }\n },\n )\n\n return await call_next(request)\n" }, { "path": "backend/ee/onyx/server/middleware/tenant_tracking.py", "content": "import logging\nfrom collections.abc import Awaitable\nfrom collections.abc import Callable\n\nfrom fastapi import FastAPI\nfrom fastapi import HTTPException\nfrom fastapi import Request\nfrom fastapi import Response\n\nfrom ee.onyx.auth.users import decode_anonymous_user_jwt_token\nfrom onyx.auth.utils import extract_tenant_from_auth_header\nfrom onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME\nfrom onyx.configs.constants import TENANT_ID_COOKIE_NAME\nfrom onyx.db.engine.sql_engine import is_valid_schema_name\nfrom onyx.redis.redis_pool import retrieve_auth_token_data_from_redis\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n\ndef add_api_server_tenant_id_middleware(\n app: FastAPI, logger: logging.LoggerAdapter\n) -> None:\n @app.middleware(\"http\")\n async def set_tenant_id(\n request: Request, call_next: Callable[[Request], Awaitable[Response]]\n ) -> Response:\n \"\"\"Extracts the tenant id from multiple locations and sets the context var.\n\n This is very specific to the api server and probably not something you'd want\n to use elsewhere.\n \"\"\"\n try:\n if MULTI_TENANT:\n tenant_id = await _get_tenant_id_from_request(request, logger)\n else:\n tenant_id = POSTGRES_DEFAULT_SCHEMA\n\n CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n return await call_next(request)\n\n except Exception as e:\n logger.exception(f\"Error in tenant ID middleware: {str(e)}\")\n raise\n\n\nasync def _get_tenant_id_from_request(\n request: Request, logger: logging.LoggerAdapter\n) -> str:\n \"\"\"\n Attempt to extract tenant_id from:\n 1) The API key or PAT (Personal Access Token) header\n 2) The Redis-based token (stored in Cookie: fastapiusersauth)\n 3) The anonymous user cookie\n Fallback: POSTGRES_DEFAULT_SCHEMA\n \"\"\"\n # Check for API key or PAT in Authorization header\n tenant_id = extract_tenant_from_auth_header(request)\n if tenant_id is not None:\n return tenant_id\n\n try:\n # Look up token data in Redis\n\n token_data = await retrieve_auth_token_data_from_redis(request)\n\n if token_data:\n tenant_id_from_payload = token_data.get(\n \"tenant_id\", POSTGRES_DEFAULT_SCHEMA\n )\n\n tenant_id = (\n str(tenant_id_from_payload)\n if tenant_id_from_payload is not None\n else None\n )\n\n if tenant_id and not is_valid_schema_name(tenant_id):\n raise HTTPException(status_code=400, detail=\"Invalid tenant ID format\")\n\n # Check for anonymous user cookie\n anonymous_user_cookie = request.cookies.get(ANONYMOUS_USER_COOKIE_NAME)\n if anonymous_user_cookie:\n try:\n anonymous_user_data = decode_anonymous_user_jwt_token(\n anonymous_user_cookie\n )\n tenant_id = anonymous_user_data.get(\n \"tenant_id\", POSTGRES_DEFAULT_SCHEMA\n )\n\n if not tenant_id or not is_valid_schema_name(tenant_id):\n raise HTTPException(\n status_code=400, detail=\"Invalid tenant ID format\"\n )\n\n return tenant_id\n\n except Exception as e:\n logger.error(f\"Error decoding anonymous user cookie: {str(e)}\")\n # Continue and attempt to authenticate\n\n logger.debug(\n \"Token data not found or expired in Redis, defaulting to POSTGRES_DEFAULT_SCHEMA\"\n )\n\n # Return POSTGRES_DEFAULT_SCHEMA, so non-authenticated requests are sent to the default schema\n # The CURRENT_TENANT_ID_CONTEXTVAR is initialized with POSTGRES_DEFAULT_SCHEMA,\n # so we maintain consistency by returning it here when no valid tenant is found.\n return POSTGRES_DEFAULT_SCHEMA\n\n except Exception as e:\n logger.error(f\"Unexpected error in _get_tenant_id_from_request: {str(e)}\")\n raise HTTPException(status_code=500, detail=\"Internal server error\")\n\n finally:\n if tenant_id:\n return tenant_id\n\n # As a final step, check for explicit tenant_id cookie\n tenant_id_cookie = request.cookies.get(TENANT_ID_COOKIE_NAME)\n if tenant_id_cookie and is_valid_schema_name(tenant_id_cookie):\n return tenant_id_cookie\n\n # If we've reached this point, return the default schema\n return POSTGRES_DEFAULT_SCHEMA\n" }, { "path": "backend/ee/onyx/server/oauth/api.py", "content": "import base64\nimport uuid\n\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi.responses import JSONResponse\n\nfrom ee.onyx.server.oauth.api_router import router\nfrom ee.onyx.server.oauth.confluence_cloud import ConfluenceCloudOAuth\nfrom ee.onyx.server.oauth.google_drive import GoogleDriveOAuth\nfrom ee.onyx.server.oauth.slack import SlackOAuth\nfrom onyx.auth.users import current_admin_user\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import User\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\n@router.post(\"/prepare-authorization-request\")\ndef prepare_authorization_request(\n connector: DocumentSource,\n redirect_on_success: str | None,\n user: User = Depends(current_admin_user),\n tenant_id: str | None = Depends(get_current_tenant_id),\n) -> JSONResponse:\n \"\"\"Used by the frontend to generate the url for the user's browser during auth request.\n\n Example: https://www.oauth.com/oauth2-servers/authorization/the-authorization-request/\n \"\"\"\n\n # create random oauth state param for security and to retrieve user data later\n oauth_uuid = uuid.uuid4()\n oauth_uuid_str = str(oauth_uuid)\n\n # urlsafe b64 encode the uuid for the oauth url\n oauth_state = (\n base64.urlsafe_b64encode(oauth_uuid.bytes).rstrip(b\"=\").decode(\"utf-8\")\n )\n\n session: str | None = None\n if connector == DocumentSource.SLACK:\n if not DEV_MODE:\n oauth_url = SlackOAuth.generate_oauth_url(oauth_state)\n else:\n oauth_url = SlackOAuth.generate_dev_oauth_url(oauth_state)\n\n session = SlackOAuth.session_dump_json(\n email=user.email, redirect_on_success=redirect_on_success\n )\n elif connector == DocumentSource.CONFLUENCE:\n if not DEV_MODE:\n oauth_url = ConfluenceCloudOAuth.generate_oauth_url(oauth_state)\n else:\n oauth_url = ConfluenceCloudOAuth.generate_dev_oauth_url(oauth_state)\n session = ConfluenceCloudOAuth.session_dump_json(\n email=user.email, redirect_on_success=redirect_on_success\n )\n elif connector == DocumentSource.GOOGLE_DRIVE:\n if not DEV_MODE:\n oauth_url = GoogleDriveOAuth.generate_oauth_url(oauth_state)\n else:\n oauth_url = GoogleDriveOAuth.generate_dev_oauth_url(oauth_state)\n session = GoogleDriveOAuth.session_dump_json(\n email=user.email, redirect_on_success=redirect_on_success\n )\n else:\n oauth_url = None\n\n if not oauth_url:\n raise HTTPException(\n status_code=404,\n detail=f\"The document source type {connector} does not have OAuth implemented\",\n )\n\n if not session:\n raise HTTPException(\n status_code=500,\n detail=f\"The document source type {connector} failed to generate an OAuth session.\",\n )\n\n r = get_redis_client(tenant_id=tenant_id)\n\n # store important session state to retrieve when the user is redirected back\n # 10 min is the max we want an oauth flow to be valid\n r.set(f\"da_oauth:{oauth_uuid_str}\", session, ex=600)\n\n return JSONResponse(content={\"url\": oauth_url})\n" }, { "path": "backend/ee/onyx/server/oauth/api_router.py", "content": "from fastapi import APIRouter\n\nrouter: APIRouter = APIRouter(prefix=\"/oauth\")\n" }, { "path": "backend/ee/onyx/server/oauth/confluence_cloud.py", "content": "import base64\nimport uuid\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi.responses import JSONResponse\nfrom pydantic import BaseModel\nfrom pydantic import ValidationError\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.server.oauth.api_router import router\nfrom onyx.auth.users import current_admin_user\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.utils import CONFLUENCE_OAUTH_TOKEN_URL\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.credentials import fetch_credential_by_id_for_user\nfrom onyx.db.credentials import update_credential_json\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\nclass ConfluenceCloudOAuth:\n # https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/\n\n class OAuthSession(BaseModel):\n \"\"\"Stored in redis to be looked up on callback\"\"\"\n\n email: str\n redirect_on_success: str | None # Where to send the user if OAuth flow succeeds\n\n class TokenResponse(BaseModel):\n access_token: str\n expires_in: int\n token_type: str\n refresh_token: str\n scope: str\n\n class AccessibleResources(BaseModel):\n id: str\n name: str\n url: str\n scopes: list[str]\n avatarUrl: str\n\n CLIENT_ID = OAUTH_CONFLUENCE_CLOUD_CLIENT_ID\n CLIENT_SECRET = OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET\n TOKEN_URL = CONFLUENCE_OAUTH_TOKEN_URL\n\n ACCESSIBLE_RESOURCE_URL = (\n \"https://api.atlassian.com/oauth/token/accessible-resources\"\n )\n\n # All read scopes per https://developer.atlassian.com/cloud/confluence/scopes-for-oauth-2-3LO-and-forge-apps/\n CONFLUENCE_OAUTH_SCOPE = (\n # classic scope\n \"read:confluence-space.summary%20\"\n \"read:confluence-props%20\"\n \"read:confluence-content.all%20\"\n \"read:confluence-content.summary%20\"\n \"read:confluence-content.permission%20\"\n \"read:confluence-user%20\"\n \"read:confluence-groups%20\"\n \"read:space:confluence%20\"\n \"readonly:content.attachment:confluence%20\"\n \"search:confluence%20\"\n # granular scope\n \"read:attachment:confluence%20\" # possibly unneeded unless calling v2 attachments api\n \"read:content-details:confluence%20\" # for permission sync\n \"offline_access\"\n )\n\n REDIRECT_URI = f\"{WEB_DOMAIN}/admin/connectors/confluence/oauth/callback\"\n DEV_REDIRECT_URI = f\"https://redirectmeto.com/{REDIRECT_URI}\"\n\n # eventually for Confluence Data Center\n # oauth_url = (\n # f\"http://localhost:8090/rest/oauth/v2/authorize?client_id={CONFLUENCE_OAUTH_CLIENT_ID}\"\n # f\"&scope={CONFLUENCE_OAUTH_SCOPE_2}\"\n # f\"&redirect_uri={redirectme_uri}\"\n # )\n\n @classmethod\n def generate_oauth_url(cls, state: str) -> str:\n return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)\n\n @classmethod\n def generate_dev_oauth_url(cls, state: str) -> str:\n \"\"\"dev mode workaround for localhost testing\n - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https\n \"\"\"\n return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)\n\n @classmethod\n def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:\n # https://developer.atlassian.com/cloud/jira/platform/oauth-2-3lo-apps/#1--direct-the-user-to-the-authorization-url-to-get-an-authorization-code\n\n url = (\n \"https://auth.atlassian.com/authorize\"\n f\"?audience=api.atlassian.com\"\n f\"&client_id={cls.CLIENT_ID}\"\n f\"&scope={cls.CONFLUENCE_OAUTH_SCOPE}\"\n f\"&redirect_uri={redirect_uri}\"\n f\"&state={state}\"\n \"&response_type=code\"\n \"&prompt=consent\"\n )\n return url\n\n @classmethod\n def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:\n \"\"\"Temporary state to store in redis. to be looked up on auth response.\n Returns a json string.\n \"\"\"\n session = ConfluenceCloudOAuth.OAuthSession(\n email=email, redirect_on_success=redirect_on_success\n )\n return session.model_dump_json()\n\n @classmethod\n def parse_session(cls, session_json: str) -> OAuthSession:\n session = ConfluenceCloudOAuth.OAuthSession.model_validate_json(session_json)\n return session\n\n @classmethod\n def generate_finalize_url(cls, credential_id: int) -> str:\n return f\"{WEB_DOMAIN}/admin/connectors/confluence/oauth/finalize?credential={credential_id}\"\n\n\n@router.post(\"/connector/confluence/callback\")\ndef confluence_oauth_callback(\n code: str,\n state: str,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n tenant_id: str | None = Depends(get_current_tenant_id),\n) -> JSONResponse:\n \"\"\"Handles the backend logic for the frontend page that the user is redirected to\n after visiting the oauth authorization url.\"\"\"\n\n if not ConfluenceCloudOAuth.CLIENT_ID or not ConfluenceCloudOAuth.CLIENT_SECRET:\n raise HTTPException(\n status_code=500,\n detail=\"Confluence Cloud client ID or client secret is not configured.\",\n )\n\n r = get_redis_client(tenant_id=tenant_id)\n\n # recover the state\n padded_state = state + \"=\" * (\n -len(state) % 4\n ) # Add padding back (Base64 decoding requires padding)\n uuid_bytes = base64.urlsafe_b64decode(\n padded_state\n ) # Decode the Base64 string back to bytes\n\n # Convert bytes back to a UUID\n oauth_uuid = uuid.UUID(bytes=uuid_bytes)\n oauth_uuid_str = str(oauth_uuid)\n\n r_key = f\"da_oauth:{oauth_uuid_str}\"\n\n session_json_bytes = cast(bytes, r.get(r_key))\n if not session_json_bytes:\n raise HTTPException(\n status_code=400,\n detail=f\"Confluence Cloud OAuth failed - OAuth state key not found: key={r_key}\",\n )\n\n session_json = session_json_bytes.decode(\"utf-8\")\n try:\n session = ConfluenceCloudOAuth.parse_session(session_json)\n\n if not DEV_MODE:\n redirect_uri = ConfluenceCloudOAuth.REDIRECT_URI\n else:\n redirect_uri = ConfluenceCloudOAuth.DEV_REDIRECT_URI\n\n # Exchange the authorization code for an access token\n response = requests.post(\n ConfluenceCloudOAuth.TOKEN_URL,\n headers={\"Content-Type\": \"application/x-www-form-urlencoded\"},\n data={\n \"client_id\": ConfluenceCloudOAuth.CLIENT_ID,\n \"client_secret\": ConfluenceCloudOAuth.CLIENT_SECRET,\n \"code\": code,\n \"redirect_uri\": redirect_uri,\n \"grant_type\": \"authorization_code\",\n },\n )\n\n token_response: ConfluenceCloudOAuth.TokenResponse | None = None\n\n try:\n token_response = ConfluenceCloudOAuth.TokenResponse.model_validate_json(\n response.text\n )\n except Exception:\n raise RuntimeError(\n \"Confluence Cloud OAuth failed during code/token exchange.\"\n )\n\n now = datetime.now(timezone.utc)\n expires_at = now + timedelta(seconds=token_response.expires_in)\n\n credential_info = CredentialBase(\n credential_json={\n \"confluence_access_token\": token_response.access_token,\n \"confluence_refresh_token\": token_response.refresh_token,\n \"created_at\": now.isoformat(),\n \"expires_at\": expires_at.isoformat(),\n \"expires_in\": token_response.expires_in,\n \"scope\": token_response.scope,\n },\n admin_public=True,\n source=DocumentSource.CONFLUENCE,\n name=\"Confluence Cloud OAuth\",\n )\n\n credential = create_credential(credential_info, user, db_session)\n except Exception as e:\n return JSONResponse(\n status_code=500,\n content={\n \"success\": False,\n \"message\": f\"An error occurred during Confluence Cloud OAuth: {str(e)}\",\n },\n )\n finally:\n r.delete(r_key)\n\n # return the result\n return JSONResponse(\n content={\n \"success\": True,\n \"message\": \"Confluence Cloud OAuth completed successfully.\",\n \"finalize_url\": ConfluenceCloudOAuth.generate_finalize_url(credential.id),\n \"redirect_on_success\": session.redirect_on_success,\n }\n )\n\n\n@router.get(\"/connector/confluence/accessible-resources\")\ndef confluence_oauth_accessible_resources(\n credential_id: int,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n tenant_id: str | None = Depends(get_current_tenant_id), # noqa: ARG001\n) -> JSONResponse:\n \"\"\"Atlassian's API is weird and does not supply us with enough info to be in a\n usable state after authorizing. All API's require a cloud id. We have to list\n the accessible resources/sites and let the user choose which site to use.\"\"\"\n\n credential = fetch_credential_by_id_for_user(credential_id, user, db_session)\n if not credential:\n raise HTTPException(400, f\"Credential {credential_id} not found.\")\n\n credential_dict = (\n credential.credential_json.get_value(apply_mask=False)\n if credential.credential_json\n else {}\n )\n access_token = credential_dict[\"confluence_access_token\"]\n\n try:\n # Exchange the authorization code for an access token\n response = requests.get(\n ConfluenceCloudOAuth.ACCESSIBLE_RESOURCE_URL,\n headers={\n \"Authorization\": f\"Bearer {access_token}\",\n \"Accept\": \"application/json\",\n },\n )\n\n response.raise_for_status()\n accessible_resources_data = response.json()\n\n # Validate the list of AccessibleResources\n try:\n accessible_resources = [\n ConfluenceCloudOAuth.AccessibleResources(**resource)\n for resource in accessible_resources_data\n ]\n except ValidationError as e:\n raise RuntimeError(f\"Failed to parse accessible resources: {e}\")\n except Exception as e:\n return JSONResponse(\n status_code=500,\n content={\n \"success\": False,\n \"message\": f\"An error occurred retrieving Confluence Cloud accessible resources: {str(e)}\",\n },\n )\n\n # return the result\n return JSONResponse(\n content={\n \"success\": True,\n \"message\": \"Confluence Cloud get accessible resources completed successfully.\",\n \"accessible_resources\": [\n resource.model_dump() for resource in accessible_resources\n ],\n }\n )\n\n\n@router.post(\"/connector/confluence/finalize\")\ndef confluence_oauth_finalize(\n credential_id: int,\n cloud_id: str,\n cloud_name: str,\n cloud_url: str,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n tenant_id: str | None = Depends(get_current_tenant_id), # noqa: ARG001\n) -> JSONResponse:\n \"\"\"Saves the info for the selected cloud site to the credential.\n This is the final step in the confluence oauth flow where after the traditional\n OAuth process, the user has to select a site to associate with the credentials.\n After this, the credential is usable.\"\"\"\n\n credential = fetch_credential_by_id_for_user(credential_id, user, db_session)\n if not credential:\n raise HTTPException(\n status_code=400,\n detail=f\"Confluence Cloud OAuth failed - credential {credential_id} not found.\",\n )\n\n existing_credential_json = (\n credential.credential_json.get_value(apply_mask=False)\n if credential.credential_json\n else {}\n )\n new_credential_json: dict[str, Any] = dict(existing_credential_json)\n new_credential_json[\"cloud_id\"] = cloud_id\n new_credential_json[\"cloud_name\"] = cloud_name\n new_credential_json[\"wiki_base\"] = cloud_url\n\n try:\n update_credential_json(credential_id, new_credential_json, user, db_session)\n except Exception as e:\n return JSONResponse(\n status_code=500,\n content={\n \"success\": False,\n \"message\": f\"An error occurred during Confluence Cloud OAuth: {str(e)}\",\n },\n )\n\n # return the result\n return JSONResponse(\n content={\n \"success\": True,\n \"message\": \"Confluence Cloud OAuth finalized successfully.\",\n \"redirect_url\": f\"{WEB_DOMAIN}/admin/connectors/confluence\",\n }\n )\n" }, { "path": "backend/ee/onyx/server/oauth/google_drive.py", "content": "import base64\nimport json\nimport uuid\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi.responses import JSONResponse\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.server.oauth.api_router import router\nfrom onyx.auth.users import current_admin_user\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_SECRET\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.google_utils.google_auth import get_google_oauth_creds\nfrom onyx.connectors.google_utils.google_auth import sanitize_oauth_credentials\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_AUTHENTICATION_METHOD,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_TOKEN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GoogleOAuthAuthenticationMethod,\n)\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import CredentialBase\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\nclass GoogleDriveOAuth:\n # https://developers.google.com/identity/protocols/oauth2\n # https://developers.google.com/identity/protocols/oauth2/web-server\n\n class OAuthSession(BaseModel):\n \"\"\"Stored in redis to be looked up on callback\"\"\"\n\n email: str\n redirect_on_success: str | None # Where to send the user if OAuth flow succeeds\n\n CLIENT_ID = OAUTH_GOOGLE_DRIVE_CLIENT_ID\n CLIENT_SECRET = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET\n\n TOKEN_URL = \"https://oauth2.googleapis.com/token\"\n\n # SCOPE is per https://docs.danswer.dev/connectors/google-drive\n # TODO: Merge with or use google_utils.GOOGLE_SCOPES\n SCOPE = (\n \"https://www.googleapis.com/auth/drive.readonly%20\"\n \"https://www.googleapis.com/auth/drive.metadata.readonly%20\"\n \"https://www.googleapis.com/auth/admin.directory.user.readonly%20\"\n \"https://www.googleapis.com/auth/admin.directory.group.readonly\"\n )\n\n REDIRECT_URI = f\"{WEB_DOMAIN}/admin/connectors/google-drive/oauth/callback\"\n DEV_REDIRECT_URI = f\"https://redirectmeto.com/{REDIRECT_URI}\"\n\n @classmethod\n def generate_oauth_url(cls, state: str) -> str:\n return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)\n\n @classmethod\n def generate_dev_oauth_url(cls, state: str) -> str:\n \"\"\"dev mode workaround for localhost testing\n - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https\n \"\"\"\n\n return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)\n\n @classmethod\n def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:\n # without prompt=consent, a refresh token is only issued the first time the user approves\n url = (\n f\"https://accounts.google.com/o/oauth2/v2/auth\"\n f\"?client_id={cls.CLIENT_ID}\"\n f\"&redirect_uri={redirect_uri}\"\n \"&response_type=code\"\n f\"&scope={cls.SCOPE}\"\n \"&access_type=offline\"\n f\"&state={state}\"\n \"&prompt=consent\"\n )\n return url\n\n @classmethod\n def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:\n \"\"\"Temporary state to store in redis. to be looked up on auth response.\n Returns a json string.\n \"\"\"\n session = GoogleDriveOAuth.OAuthSession(\n email=email, redirect_on_success=redirect_on_success\n )\n return session.model_dump_json()\n\n @classmethod\n def parse_session(cls, session_json: str) -> OAuthSession:\n session = GoogleDriveOAuth.OAuthSession.model_validate_json(session_json)\n return session\n\n\n@router.post(\"/connector/google-drive/callback\")\ndef handle_google_drive_oauth_callback(\n code: str,\n state: str,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n tenant_id: str | None = Depends(get_current_tenant_id),\n) -> JSONResponse:\n if not GoogleDriveOAuth.CLIENT_ID or not GoogleDriveOAuth.CLIENT_SECRET:\n raise HTTPException(\n status_code=500,\n detail=\"Google Drive client ID or client secret is not configured.\",\n )\n\n r = get_redis_client(tenant_id=tenant_id)\n\n # recover the state\n padded_state = state + \"=\" * (\n -len(state) % 4\n ) # Add padding back (Base64 decoding requires padding)\n uuid_bytes = base64.urlsafe_b64decode(\n padded_state\n ) # Decode the Base64 string back to bytes\n\n # Convert bytes back to a UUID\n oauth_uuid = uuid.UUID(bytes=uuid_bytes)\n oauth_uuid_str = str(oauth_uuid)\n\n r_key = f\"da_oauth:{oauth_uuid_str}\"\n\n session_json_bytes = cast(bytes, r.get(r_key))\n if not session_json_bytes:\n raise HTTPException(\n status_code=400,\n detail=f\"Google Drive OAuth failed - OAuth state key not found: key={r_key}\",\n )\n\n session_json = session_json_bytes.decode(\"utf-8\")\n try:\n session = GoogleDriveOAuth.parse_session(session_json)\n\n if not DEV_MODE:\n redirect_uri = GoogleDriveOAuth.REDIRECT_URI\n else:\n redirect_uri = GoogleDriveOAuth.DEV_REDIRECT_URI\n\n # Exchange the authorization code for an access token\n response = requests.post(\n GoogleDriveOAuth.TOKEN_URL,\n headers={\"Content-Type\": \"application/x-www-form-urlencoded\"},\n data={\n \"client_id\": GoogleDriveOAuth.CLIENT_ID,\n \"client_secret\": GoogleDriveOAuth.CLIENT_SECRET,\n \"code\": code,\n \"redirect_uri\": redirect_uri,\n \"grant_type\": \"authorization_code\",\n },\n )\n\n response.raise_for_status()\n\n authorization_response: dict[str, Any] = response.json()\n\n # the connector wants us to store the json in its authorized_user_info format\n # returned from OAuthCredentials.get_authorized_user_info().\n # So refresh immediately via get_google_oauth_creds with the params filled in\n # from fields in authorization_response to get the json we need\n authorized_user_info = {}\n authorized_user_info[\"client_id\"] = OAUTH_GOOGLE_DRIVE_CLIENT_ID\n authorized_user_info[\"client_secret\"] = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET\n authorized_user_info[\"refresh_token\"] = authorization_response[\"refresh_token\"]\n\n token_json_str = json.dumps(authorized_user_info)\n oauth_creds = get_google_oauth_creds(\n token_json_str=token_json_str, source=DocumentSource.GOOGLE_DRIVE\n )\n if not oauth_creds:\n raise RuntimeError(\"get_google_oauth_creds returned None.\")\n\n # save off the credentials\n oauth_creds_sanitized_json_str = sanitize_oauth_credentials(oauth_creds)\n\n credential_dict: dict[str, str] = {}\n credential_dict[DB_CREDENTIALS_DICT_TOKEN_KEY] = oauth_creds_sanitized_json_str\n credential_dict[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] = session.email\n credential_dict[DB_CREDENTIALS_AUTHENTICATION_METHOD] = (\n GoogleOAuthAuthenticationMethod.OAUTH_INTERACTIVE.value\n )\n\n credential_info = CredentialBase(\n credential_json=credential_dict,\n admin_public=True,\n source=DocumentSource.GOOGLE_DRIVE,\n name=\"OAuth (interactive)\",\n )\n\n create_credential(credential_info, user, db_session)\n except Exception as e:\n return JSONResponse(\n status_code=500,\n content={\n \"success\": False,\n \"message\": f\"An error occurred during Google Drive OAuth: {str(e)}\",\n },\n )\n finally:\n r.delete(r_key)\n\n # return the result\n return JSONResponse(\n content={\n \"success\": True,\n \"message\": \"Google Drive OAuth completed successfully.\",\n \"finalize_url\": None,\n \"redirect_on_success\": session.redirect_on_success,\n }\n )\n" }, { "path": "backend/ee/onyx/server/oauth/slack.py", "content": "import base64\nimport uuid\nfrom typing import cast\n\nimport requests\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi.responses import JSONResponse\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.server.oauth.api_router import router\nfrom onyx.auth.users import current_admin_user\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.configs.app_configs import OAUTH_SLACK_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_SLACK_CLIENT_SECRET\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import CredentialBase\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\nclass SlackOAuth:\n # https://knock.app/blog/how-to-authenticate-users-in-slack-using-oauth\n # Example: https://api.slack.com/authentication/oauth-v2#exchanging\n\n class OAuthSession(BaseModel):\n \"\"\"Stored in redis to be looked up on callback\"\"\"\n\n email: str\n redirect_on_success: str | None # Where to send the user if OAuth flow succeeds\n\n CLIENT_ID = OAUTH_SLACK_CLIENT_ID\n CLIENT_SECRET = OAUTH_SLACK_CLIENT_SECRET\n\n TOKEN_URL = \"https://slack.com/api/oauth.v2.access\"\n\n # SCOPE is per https://docs.danswer.dev/connectors/slack\n BOT_SCOPE = (\n \"channels:history,\"\n \"channels:read,\"\n \"groups:history,\"\n \"groups:read,\"\n \"channels:join,\"\n \"im:history,\"\n \"users:read,\"\n \"users:read.email,\"\n \"usergroups:read\"\n )\n\n REDIRECT_URI = f\"{WEB_DOMAIN}/admin/connectors/slack/oauth/callback\"\n DEV_REDIRECT_URI = f\"https://redirectmeto.com/{REDIRECT_URI}\"\n\n @classmethod\n def generate_oauth_url(cls, state: str) -> str:\n return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)\n\n @classmethod\n def generate_dev_oauth_url(cls, state: str) -> str:\n \"\"\"dev mode workaround for localhost testing\n - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https\n \"\"\"\n\n return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)\n\n @classmethod\n def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:\n url = (\n f\"https://slack.com/oauth/v2/authorize\"\n f\"?client_id={cls.CLIENT_ID}\"\n f\"&redirect_uri={redirect_uri}\"\n f\"&scope={cls.BOT_SCOPE}\"\n f\"&state={state}\"\n )\n return url\n\n @classmethod\n def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:\n \"\"\"Temporary state to store in redis. to be looked up on auth response.\n Returns a json string.\n \"\"\"\n session = SlackOAuth.OAuthSession(\n email=email, redirect_on_success=redirect_on_success\n )\n return session.model_dump_json()\n\n @classmethod\n def parse_session(cls, session_json: str) -> OAuthSession:\n session = SlackOAuth.OAuthSession.model_validate_json(session_json)\n return session\n\n\n@router.post(\"/connector/slack/callback\")\ndef handle_slack_oauth_callback(\n code: str,\n state: str,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n tenant_id: str | None = Depends(get_current_tenant_id),\n) -> JSONResponse:\n if not SlackOAuth.CLIENT_ID or not SlackOAuth.CLIENT_SECRET:\n raise HTTPException(\n status_code=500,\n detail=\"Slack client ID or client secret is not configured.\",\n )\n\n r = get_redis_client(tenant_id=tenant_id)\n\n # recover the state\n padded_state = state + \"=\" * (\n -len(state) % 4\n ) # Add padding back (Base64 decoding requires padding)\n uuid_bytes = base64.urlsafe_b64decode(\n padded_state\n ) # Decode the Base64 string back to bytes\n\n # Convert bytes back to a UUID\n oauth_uuid = uuid.UUID(bytes=uuid_bytes)\n oauth_uuid_str = str(oauth_uuid)\n\n r_key = f\"da_oauth:{oauth_uuid_str}\"\n\n session_json_bytes = cast(bytes, r.get(r_key))\n if not session_json_bytes:\n raise HTTPException(\n status_code=400,\n detail=f\"Slack OAuth failed - OAuth state key not found: key={r_key}\",\n )\n\n session_json = session_json_bytes.decode(\"utf-8\")\n try:\n session = SlackOAuth.parse_session(session_json)\n\n if not DEV_MODE:\n redirect_uri = SlackOAuth.REDIRECT_URI\n else:\n redirect_uri = SlackOAuth.DEV_REDIRECT_URI\n\n # Exchange the authorization code for an access token\n response = requests.post(\n SlackOAuth.TOKEN_URL,\n headers={\"Content-Type\": \"application/x-www-form-urlencoded\"},\n data={\n \"client_id\": SlackOAuth.CLIENT_ID,\n \"client_secret\": SlackOAuth.CLIENT_SECRET,\n \"code\": code,\n \"redirect_uri\": redirect_uri,\n },\n )\n\n response_data = response.json()\n\n if not response_data.get(\"ok\"):\n raise HTTPException(\n status_code=400,\n detail=f\"Slack OAuth failed: {response_data.get('error')}\",\n )\n\n # Extract token and team information\n access_token: str = response_data.get(\"access_token\")\n team_id: str = response_data.get(\"team\", {}).get(\"id\")\n authed_user_id: str = response_data.get(\"authed_user\", {}).get(\"id\")\n\n credential_info = CredentialBase(\n credential_json={\"slack_bot_token\": access_token},\n admin_public=True,\n source=DocumentSource.SLACK,\n name=\"Slack OAuth\",\n )\n\n create_credential(credential_info, user, db_session)\n except Exception as e:\n return JSONResponse(\n status_code=500,\n content={\n \"success\": False,\n \"message\": f\"An error occurred during Slack OAuth: {str(e)}\",\n },\n )\n finally:\n r.delete(r_key)\n\n # return the result\n return JSONResponse(\n content={\n \"success\": True,\n \"message\": \"Slack OAuth completed successfully.\",\n \"finalize_url\": None,\n \"redirect_on_success\": session.redirect_on_success,\n \"team_id\": team_id,\n \"authed_user_id\": authed_user_id,\n }\n )\n" }, { "path": "backend/ee/onyx/server/query_and_chat/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/query_and_chat/models.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.context.search.models import BaseFilters\nfrom onyx.context.search.models import InferenceSection\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.server.manage.models import StandardAnswer\n\n\nclass StandardAnswerRequest(BaseModel):\n message: str\n slack_bot_categories: list[str]\n\n\nclass StandardAnswerResponse(BaseModel):\n standard_answers: list[StandardAnswer] = Field(default_factory=list)\n\n\nclass SearchFlowClassificationRequest(BaseModel):\n user_query: str\n\n\nclass SearchFlowClassificationResponse(BaseModel):\n is_search_flow: bool\n\n\n# NOTE: This model is used for the core flow of the Onyx application, any\n# changes to it should be reviewed and approved by an experienced team member.\n# It is very important to 1. avoid bloat and 2. that this remains backwards\n# compatible across versions.\nclass SendSearchQueryRequest(BaseModel):\n search_query: str\n filters: BaseFilters | None = None\n num_docs_fed_to_llm_selection: int | None = None\n run_query_expansion: bool = False\n num_hits: int = 30\n hybrid_alpha: float | None = None\n include_content: bool = False\n stream: bool = False\n\n\nclass SearchDocWithContent(SearchDoc):\n # Allows None because this is determined by a flag but the object used in code\n # of the search path uses this type\n content: str | None\n\n @classmethod\n def from_inference_sections(\n cls,\n sections: Sequence[InferenceSection],\n include_content: bool = False,\n is_internet: bool = False,\n ) -> list[\"SearchDocWithContent\"]:\n \"\"\"Convert InferenceSections to SearchDocWithContent objects.\n\n Args:\n sections: Sequence of InferenceSection objects\n include_content: If True, populate content field with combined_content\n is_internet: Whether these are internet search results\n\n Returns:\n List of SearchDocWithContent with optional content\n \"\"\"\n if not sections:\n return []\n\n return [\n cls(\n document_id=(chunk := section.center_chunk).document_id,\n chunk_ind=chunk.chunk_id,\n semantic_identifier=chunk.semantic_identifier or \"Unknown\",\n link=chunk.source_links[0] if chunk.source_links else None,\n blurb=chunk.blurb,\n source_type=chunk.source_type,\n boost=chunk.boost,\n hidden=chunk.hidden,\n metadata=chunk.metadata,\n score=chunk.score,\n match_highlights=chunk.match_highlights,\n updated_at=chunk.updated_at,\n primary_owners=chunk.primary_owners,\n secondary_owners=chunk.secondary_owners,\n is_internet=is_internet,\n content=section.combined_content if include_content else None,\n )\n for section in sections\n ]\n\n\nclass SearchFullResponse(BaseModel):\n all_executed_queries: list[str]\n search_docs: list[SearchDocWithContent]\n # Reasoning tokens output by the LLM for the document selection\n doc_selection_reasoning: str | None = None\n # This a list of document ids that are in the search_docs list\n llm_selected_doc_ids: list[str] | None = None\n # Error message if the search failed partway through\n error: str | None = None\n\n\nclass SearchQueryResponse(BaseModel):\n query: str\n query_expansions: list[str] | None\n created_at: datetime\n\n\nclass SearchHistoryResponse(BaseModel):\n search_queries: list[SearchQueryResponse]\n" }, { "path": "backend/ee/onyx/server/query_and_chat/query_backend.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.onyxbot.slack.handlers.handle_standard_answers import (\n oneoff_standard_answers,\n)\nfrom ee.onyx.server.query_and_chat.models import StandardAnswerRequest\nfrom ee.onyx.server.query_and_chat.models import StandardAnswerResponse\nfrom onyx.auth.users import current_user\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nbasic_router = APIRouter(prefix=\"/query\")\n\n\n@basic_router.get(\"/standard-answer\")\ndef get_standard_answer(\n request: StandardAnswerRequest,\n db_session: Session = Depends(get_session),\n _: User = Depends(current_user),\n) -> StandardAnswerResponse:\n try:\n standard_answers = oneoff_standard_answers(\n message=request.message,\n slack_bot_categories=request.slack_bot_categories,\n db_session=db_session,\n )\n return StandardAnswerResponse(standard_answers=standard_answers)\n except Exception as e:\n logger.error(f\"Error in get_standard_answer: {str(e)}\", exc_info=True)\n raise HTTPException(status_code=500, detail=\"An internal server error occurred\")\n" }, { "path": "backend/ee/onyx/server/query_and_chat/search_backend.py", "content": "from collections.abc import Generator\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi.responses import StreamingResponse\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.search import fetch_search_queries_for_user\nfrom ee.onyx.search.process_search_query import gather_search_stream\nfrom ee.onyx.search.process_search_query import stream_search_query\nfrom ee.onyx.secondary_llm_flows.search_flow_classification import (\n classify_is_search_flow,\n)\nfrom ee.onyx.server.query_and_chat.models import SearchFlowClassificationRequest\nfrom ee.onyx.server.query_and_chat.models import SearchFlowClassificationResponse\nfrom ee.onyx.server.query_and_chat.models import SearchFullResponse\nfrom ee.onyx.server.query_and_chat.models import SearchHistoryResponse\nfrom ee.onyx.server.query_and_chat.models import SearchQueryResponse\nfrom ee.onyx.server.query_and_chat.models import SendSearchQueryRequest\nfrom ee.onyx.server.query_and_chat.streaming_models import SearchErrorPacket\nfrom onyx.auth.users import current_user\nfrom onyx.configs.app_configs import ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import User\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.server.usage_limits import check_llm_cost_limit_for_provider\nfrom onyx.server.utils import get_json_line\nfrom onyx.server.utils_vector_db import require_vector_db\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/search\")\n\n\n@router.post(\"/search-flow-classification\")\ndef search_flow_classification(\n request: SearchFlowClassificationRequest,\n _: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> SearchFlowClassificationResponse:\n query = request.user_query\n # This is a heuristic that if the user is typing a lot of text, it's unlikely they're looking for some specific document\n # Most likely something needs to be done with the text included so we'll just classify it as a chat flow\n if len(query) > 200:\n return SearchFlowClassificationResponse(is_search_flow=False)\n\n llm = get_default_llm()\n\n check_llm_cost_limit_for_provider(\n db_session=db_session,\n tenant_id=get_current_tenant_id(),\n llm_provider_api_key=llm.config.api_key,\n )\n\n try:\n is_search_flow = classify_is_search_flow(query=query, llm=llm)\n except Exception as e:\n logger.exception(\n \"Search flow classification failed; defaulting to chat flow\",\n exc_info=e,\n )\n is_search_flow = False\n\n return SearchFlowClassificationResponse(is_search_flow=is_search_flow)\n\n\n# NOTE: This endpoint is used for the core flow of the Onyx application, any\n# changes to it should be reviewed and approved by an experienced team member.\n# It is very important to 1. avoid bloat and 2. that this remains backwards\n# compatible across versions.\n@router.post(\n \"/send-search-message\",\n response_model=None,\n dependencies=[Depends(require_vector_db)],\n)\ndef handle_send_search_message(\n request: SendSearchQueryRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StreamingResponse | SearchFullResponse:\n \"\"\"\n Executes a search query with optional streaming.\n\n If hybrid_alpha is unset and ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH\n is True, executes pure keyword search.\n\n Returns:\n StreamingResponse with SSE if stream=True, otherwise SearchFullResponse.\n \"\"\"\n logger.debug(f\"Received search query: {request.search_query}\")\n\n if request.hybrid_alpha is None and ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH:\n request.hybrid_alpha = 0.0\n\n # Non-streaming path\n if not request.stream:\n try:\n packets = stream_search_query(request, user, db_session)\n return gather_search_stream(packets)\n except NotImplementedError as e:\n return SearchFullResponse(\n all_executed_queries=[],\n search_docs=[],\n error=str(e),\n )\n\n # Streaming path\n def stream_generator() -> Generator[str, None, None]:\n try:\n with get_session_with_current_tenant() as streaming_db_session:\n for packet in stream_search_query(request, user, streaming_db_session):\n yield get_json_line(packet.model_dump())\n except NotImplementedError as e:\n yield get_json_line(SearchErrorPacket(error=str(e)).model_dump())\n except HTTPException:\n raise\n except Exception as e:\n logger.exception(\"Error in search streaming\")\n yield get_json_line(SearchErrorPacket(error=str(e)).model_dump())\n\n return StreamingResponse(stream_generator(), media_type=\"text/event-stream\")\n\n\n@router.get(\"/search-history\")\ndef get_search_history(\n limit: int = 100,\n filter_days: int | None = None,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> SearchHistoryResponse:\n \"\"\"\n Fetch past search queries for the authenticated user.\n\n Args:\n limit: Maximum number of queries to return (default 100)\n filter_days: Only return queries from the last N days (optional)\n\n Returns:\n SearchHistoryResponse with list of search queries, ordered by most recent first.\n \"\"\"\n # Validate limit\n if limit <= 0:\n raise HTTPException(\n status_code=400,\n detail=\"limit must be greater than 0\",\n )\n if limit > 1000:\n raise HTTPException(\n status_code=400,\n detail=\"limit must be at most 1000\",\n )\n\n # Validate filter_days\n if filter_days is not None and filter_days <= 0:\n raise HTTPException(\n status_code=400,\n detail=\"filter_days must be greater than 0\",\n )\n\n search_queries = fetch_search_queries_for_user(\n db_session=db_session,\n user_id=user.id,\n filter_days=filter_days,\n limit=limit,\n )\n\n return SearchHistoryResponse(\n search_queries=[\n SearchQueryResponse(\n query=sq.query,\n query_expansions=sq.query_expansions,\n created_at=sq.created_at,\n )\n for sq in search_queries\n ]\n )\n" }, { "path": "backend/ee/onyx/server/query_and_chat/streaming_models.py", "content": "from typing import Literal\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\n\nfrom ee.onyx.server.query_and_chat.models import SearchDocWithContent\n\n\nclass SearchQueriesPacket(BaseModel):\n model_config = ConfigDict(frozen=True)\n\n type: Literal[\"search_queries\"] = \"search_queries\"\n all_executed_queries: list[str]\n\n\nclass SearchDocsPacket(BaseModel):\n model_config = ConfigDict(frozen=True)\n\n type: Literal[\"search_docs\"] = \"search_docs\"\n search_docs: list[SearchDocWithContent]\n\n\nclass SearchErrorPacket(BaseModel):\n model_config = ConfigDict(frozen=True)\n\n type: Literal[\"search_error\"] = \"search_error\"\n error: str\n\n\nclass LLMSelectedDocsPacket(BaseModel):\n model_config = ConfigDict(frozen=True)\n\n type: Literal[\"llm_selected_docs\"] = \"llm_selected_docs\"\n # None if LLM selection failed, empty list if no docs selected, list of IDs otherwise\n llm_selected_doc_ids: list[str] | None\n" }, { "path": "backend/ee/onyx/server/query_and_chat/token_limit.py", "content": "from collections import defaultdict\nfrom collections.abc import Sequence\nfrom datetime import datetime\nfrom itertools import groupby\nfrom typing import Dict\nfrom typing import List\nfrom typing import Tuple\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.api_key import is_api_key_email_address\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import TokenRateLimit\nfrom onyx.db.models import TokenRateLimit__UserGroup\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.db.token_limit import fetch_all_user_token_rate_limits\nfrom onyx.server.query_and_chat.token_limit import _get_cutoff_time\nfrom onyx.server.query_and_chat.token_limit import _is_rate_limited\nfrom onyx.server.query_and_chat.token_limit import _user_is_rate_limited_by_global\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\n\n\ndef _check_token_rate_limits(user: User) -> None:\n # Anonymous users are only rate limited by global settings\n if user.is_anonymous:\n _user_is_rate_limited_by_global()\n\n elif is_api_key_email_address(user.email):\n # API keys are only rate limited by global settings\n _user_is_rate_limited_by_global()\n\n else:\n run_functions_tuples_in_parallel(\n [\n (_user_is_rate_limited, (user.id,)),\n (_user_is_rate_limited_by_group, (user.id,)),\n (_user_is_rate_limited_by_global, ()),\n ]\n )\n\n\n\"\"\"\nUser rate limits\n\"\"\"\n\n\ndef _user_is_rate_limited(user_id: UUID) -> None:\n with get_session_with_current_tenant() as db_session:\n user_rate_limits = fetch_all_user_token_rate_limits(\n db_session=db_session, enabled_only=True, ordered=False\n )\n\n if user_rate_limits:\n user_cutoff_time = _get_cutoff_time(user_rate_limits)\n user_usage = _fetch_user_usage(user_id, user_cutoff_time, db_session)\n\n if _is_rate_limited(user_rate_limits, user_usage):\n raise HTTPException(\n status_code=429,\n detail=\"Token budget exceeded for user. Try again later.\",\n )\n\n\ndef _fetch_user_usage(\n user_id: UUID, cutoff_time: datetime, db_session: Session\n) -> Sequence[tuple[datetime, int]]:\n \"\"\"\n Fetch user usage within the cutoff time, grouped by minute\n \"\"\"\n result = db_session.execute(\n select(\n func.date_trunc(\"minute\", ChatMessage.time_sent),\n func.sum(ChatMessage.token_count),\n )\n .join(ChatSession, ChatMessage.chat_session_id == ChatSession.id)\n .where(ChatSession.user_id == user_id, ChatMessage.time_sent >= cutoff_time)\n .group_by(func.date_trunc(\"minute\", ChatMessage.time_sent))\n ).all()\n\n return [(row[0], row[1]) for row in result]\n\n\n\"\"\"\nUser Group rate limits\n\"\"\"\n\n\ndef _user_is_rate_limited_by_group(user_id: UUID) -> None:\n with get_session_with_current_tenant() as db_session:\n group_rate_limits = _fetch_all_user_group_rate_limits(user_id, db_session)\n\n if group_rate_limits:\n # Group cutoff time is the same for all groups.\n # This could be optimized to only fetch the maximum cutoff time for\n # a specific group, but seems unnecessary for now.\n group_cutoff_time = _get_cutoff_time(\n [e for sublist in group_rate_limits.values() for e in sublist]\n )\n\n user_group_ids = list(group_rate_limits.keys())\n group_usage = _fetch_user_group_usage(\n user_group_ids, group_cutoff_time, db_session\n )\n\n has_at_least_one_untriggered_limit = False\n for user_group_id, rate_limits in group_rate_limits.items():\n usage = group_usage.get(user_group_id, [])\n\n if not _is_rate_limited(rate_limits, usage):\n has_at_least_one_untriggered_limit = True\n break\n\n if not has_at_least_one_untriggered_limit:\n raise HTTPException(\n status_code=429,\n detail=\"Token budget exceeded for user's groups. Try again later.\",\n )\n\n\ndef _fetch_all_user_group_rate_limits(\n user_id: UUID, db_session: Session\n) -> Dict[int, List[TokenRateLimit]]:\n group_limits = (\n select(TokenRateLimit, User__UserGroup.user_group_id)\n .join(\n TokenRateLimit__UserGroup,\n TokenRateLimit.id == TokenRateLimit__UserGroup.rate_limit_id,\n )\n .join(\n UserGroup,\n UserGroup.id == TokenRateLimit__UserGroup.user_group_id,\n )\n .join(\n User__UserGroup,\n User__UserGroup.user_group_id == UserGroup.id,\n )\n .where(\n User__UserGroup.user_id == user_id,\n TokenRateLimit.enabled.is_(True),\n )\n )\n\n raw_rate_limits = db_session.execute(group_limits).all()\n\n group_rate_limits = defaultdict(list)\n for rate_limit, user_group_id in raw_rate_limits:\n group_rate_limits[user_group_id].append(rate_limit)\n\n return group_rate_limits\n\n\ndef _fetch_user_group_usage(\n user_group_ids: list[int], cutoff_time: datetime, db_session: Session\n) -> dict[int, list[Tuple[datetime, int]]]:\n \"\"\"\n Fetch user group usage within the cutoff time, grouped by minute\n \"\"\"\n user_group_usage = db_session.execute(\n select(\n func.sum(ChatMessage.token_count),\n func.date_trunc(\"minute\", ChatMessage.time_sent),\n UserGroup.id,\n )\n .join(ChatSession, ChatMessage.chat_session_id == ChatSession.id)\n .join(User__UserGroup, User__UserGroup.user_id == ChatSession.user_id)\n .join(UserGroup, UserGroup.id == User__UserGroup.user_group_id)\n .filter(UserGroup.id.in_(user_group_ids), ChatMessage.time_sent >= cutoff_time)\n .group_by(func.date_trunc(\"minute\", ChatMessage.time_sent), UserGroup.id)\n ).all()\n\n return {\n user_group_id: [(usage, time_sent) for time_sent, usage, _ in group_usage]\n for user_group_id, group_usage in groupby(\n user_group_usage, key=lambda row: row[2]\n )\n }\n" }, { "path": "backend/ee/onyx/server/query_history/api.py", "content": "import uuid\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom http import HTTPStatus\nfrom uuid import UUID\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi.responses import StreamingResponse\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.background.task_name_builders import query_history_task_name\nfrom ee.onyx.db.query_history import get_all_query_history_export_tasks\nfrom ee.onyx.db.query_history import get_page_of_chat_sessions\nfrom ee.onyx.db.query_history import get_total_filtered_chat_sessions_count\nfrom ee.onyx.server.query_history.models import ChatSessionMinimal\nfrom ee.onyx.server.query_history.models import ChatSessionSnapshot\nfrom ee.onyx.server.query_history.models import MessageSnapshot\nfrom ee.onyx.server.query_history.models import QueryHistoryExport\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import get_display_email\nfrom onyx.background.celery.versioned_apps.client import app as client_app\nfrom onyx.background.task_utils import construct_query_history_report_name\nfrom onyx.chat.chat_utils import create_chat_history_chain\nfrom onyx.configs.app_configs import ONYX_QUERY_HISTORY_TYPE\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import FileType\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.configs.constants import QAFeedbackType\nfrom onyx.configs.constants import QueryHistoryType\nfrom onyx.configs.constants import SessionType\nfrom onyx.db.chat import get_chat_session_by_id\nfrom onyx.db.chat import get_chat_sessions_by_user\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import TaskStatus\nfrom onyx.db.file_record import get_query_history_export_files\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import User\nfrom onyx.db.tasks import get_task_with_id\nfrom onyx.db.tasks import register_task\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.server.documents.models import PaginatedReturn\nfrom onyx.server.query_and_chat.models import ChatSessionDetails\nfrom onyx.server.query_and_chat.models import ChatSessionsResponse\nfrom onyx.utils.threadpool_concurrency import parallel_yield\nfrom shared_configs.contextvars import get_current_tenant_id\n\nrouter = APIRouter()\n\nONYX_ANONYMIZED_EMAIL = \"anonymous@anonymous.invalid\"\n\n\ndef ensure_query_history_is_enabled(\n disallowed: list[QueryHistoryType],\n) -> None:\n if ONYX_QUERY_HISTORY_TYPE in disallowed:\n raise HTTPException(\n status_code=HTTPStatus.FORBIDDEN,\n detail=\"Query history has been disabled by the administrator.\",\n )\n\n\ndef yield_snapshot_from_chat_session(\n chat_session: ChatSession,\n db_session: Session,\n) -> Generator[ChatSessionSnapshot | None]:\n yield snapshot_from_chat_session(chat_session=chat_session, db_session=db_session)\n\n\ndef fetch_and_process_chat_session_history(\n db_session: Session,\n start: datetime,\n end: datetime,\n limit: int | None = 500, # noqa: ARG001\n) -> Generator[ChatSessionSnapshot]:\n PAGE_SIZE = 100\n\n page = 0\n while True:\n paged_chat_sessions = get_page_of_chat_sessions(\n start_time=start,\n end_time=end,\n db_session=db_session,\n page_num=page,\n page_size=PAGE_SIZE,\n )\n\n if not paged_chat_sessions:\n break\n\n paged_snapshots = parallel_yield(\n [\n yield_snapshot_from_chat_session(\n db_session=db_session,\n chat_session=chat_session,\n )\n for chat_session in paged_chat_sessions\n ]\n )\n\n for snapshot in paged_snapshots:\n if snapshot:\n yield snapshot\n\n # If we've fetched *less* than a `PAGE_SIZE` worth\n # of data, we have reached the end of the\n # pagination sequence; break.\n if len(paged_chat_sessions) < PAGE_SIZE:\n break\n\n page += 1\n\n\ndef snapshot_from_chat_session(\n chat_session: ChatSession,\n db_session: Session,\n) -> ChatSessionSnapshot | None:\n try:\n # Older chats may not have the right structure\n messages = create_chat_history_chain(\n chat_session_id=chat_session.id, db_session=db_session\n )\n except RuntimeError:\n return None\n\n flow_type = SessionType.SLACK if chat_session.onyxbot_flow else SessionType.CHAT\n\n return ChatSessionSnapshot(\n id=chat_session.id,\n user_email=get_display_email(\n chat_session.user.email if chat_session.user else None\n ),\n name=chat_session.description,\n messages=[\n MessageSnapshot.build(message)\n for message in messages\n if message.message_type != MessageType.SYSTEM\n ],\n assistant_id=chat_session.persona_id,\n assistant_name=chat_session.persona.name if chat_session.persona else None,\n time_created=chat_session.time_created,\n flow_type=flow_type,\n )\n\n\n@router.get(\"/admin/chat-sessions\")\ndef admin_get_chat_sessions(\n user_id: UUID,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> ChatSessionsResponse:\n # we specifically don't allow this endpoint if \"anonymized\" since\n # this is a direct query on the user id\n ensure_query_history_is_enabled(\n [\n QueryHistoryType.DISABLED,\n QueryHistoryType.ANONYMIZED,\n ]\n )\n\n try:\n chat_sessions = get_chat_sessions_by_user(\n user_id=user_id, deleted=False, db_session=db_session, limit=0\n )\n\n except ValueError:\n raise ValueError(\"Chat session does not exist or has been deleted\")\n\n return ChatSessionsResponse(\n sessions=[\n ChatSessionDetails(\n id=chat.id,\n name=chat.description,\n persona_id=chat.persona_id,\n time_created=chat.time_created.isoformat(),\n time_updated=chat.time_updated.isoformat(),\n shared_status=chat.shared_status,\n current_alternate_model=chat.current_alternate_model,\n )\n for chat in chat_sessions\n ]\n )\n\n\n@router.get(\"/admin/chat-session-history\")\ndef get_chat_session_history(\n page_num: int = Query(0, ge=0),\n page_size: int = Query(10, ge=1),\n feedback_type: QAFeedbackType | None = None,\n start_time: datetime | None = None,\n end_time: datetime | None = None,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> PaginatedReturn[ChatSessionMinimal]:\n ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])\n\n page_of_chat_sessions = get_page_of_chat_sessions(\n page_num=page_num,\n page_size=page_size,\n db_session=db_session,\n start_time=start_time,\n end_time=end_time,\n feedback_filter=feedback_type,\n )\n\n total_filtered_chat_sessions_count = get_total_filtered_chat_sessions_count(\n db_session=db_session,\n start_time=start_time,\n end_time=end_time,\n feedback_filter=feedback_type,\n )\n\n minimal_chat_sessions: list[ChatSessionMinimal] = []\n\n for chat_session in page_of_chat_sessions:\n minimal_chat_session = ChatSessionMinimal.from_chat_session(chat_session)\n if ONYX_QUERY_HISTORY_TYPE == QueryHistoryType.ANONYMIZED:\n minimal_chat_session.user_email = ONYX_ANONYMIZED_EMAIL\n minimal_chat_sessions.append(minimal_chat_session)\n\n return PaginatedReturn(\n items=minimal_chat_sessions,\n total_items=total_filtered_chat_sessions_count,\n )\n\n\n@router.get(\"/admin/chat-session-history/{chat_session_id}\")\ndef get_chat_session_admin(\n chat_session_id: UUID,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> ChatSessionSnapshot:\n ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])\n\n try:\n chat_session = get_chat_session_by_id(\n chat_session_id=chat_session_id,\n user_id=None, # view chat regardless of user\n db_session=db_session,\n include_deleted=True,\n )\n except ValueError:\n raise HTTPException(\n HTTPStatus.BAD_REQUEST,\n f\"Chat session with id '{chat_session_id}' does not exist.\",\n )\n snapshot = snapshot_from_chat_session(\n chat_session=chat_session, db_session=db_session\n )\n\n if snapshot is None:\n raise HTTPException(\n HTTPStatus.BAD_REQUEST,\n f\"Could not create snapshot for chat session with id '{chat_session_id}'\",\n )\n\n if ONYX_QUERY_HISTORY_TYPE == QueryHistoryType.ANONYMIZED:\n snapshot.user_email = ONYX_ANONYMIZED_EMAIL\n\n return snapshot\n\n\n@router.get(\"/admin/query-history/list\")\ndef list_all_query_history_exports(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[QueryHistoryExport]:\n ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])\n try:\n pending_tasks = [\n QueryHistoryExport.from_task(task)\n for task in get_all_query_history_export_tasks(db_session=db_session)\n ]\n generated_files = [\n QueryHistoryExport.from_file(file)\n for file in get_query_history_export_files(db_session=db_session)\n ]\n merged = pending_tasks + generated_files\n\n # We sort based off of the start-time of the task.\n # We also return it in reverse order since viewing generated reports in most-recent to least-recent is most common.\n merged.sort(key=lambda task: task.start_time, reverse=True)\n\n return merged\n except Exception as e:\n raise HTTPException(\n HTTPStatus.INTERNAL_SERVER_ERROR, f\"Failed to get all tasks: {e}\"\n )\n\n\n@router.post(\"/admin/query-history/start-export\", tags=PUBLIC_API_TAGS)\ndef start_query_history_export(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n start: datetime | None = None,\n end: datetime | None = None,\n) -> dict[str, str]:\n ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])\n\n start = start or datetime.fromtimestamp(0, tz=timezone.utc)\n end = end or datetime.now(tz=timezone.utc)\n\n if start >= end:\n raise HTTPException(\n HTTPStatus.BAD_REQUEST,\n f\"Start time must come before end time, but instead got the start time coming after; {start=} {end=}\",\n )\n\n task_id_uuid = uuid.uuid4()\n task_id = str(task_id_uuid)\n start_time = datetime.now(tz=timezone.utc)\n\n register_task(\n db_session=db_session,\n task_name=query_history_task_name(start=start, end=end),\n task_id=task_id,\n status=TaskStatus.PENDING,\n start_time=start_time,\n )\n\n client_app.send_task(\n OnyxCeleryTask.EXPORT_QUERY_HISTORY_TASK,\n task_id=task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n queue=OnyxCeleryQueues.CSV_GENERATION,\n kwargs={\n \"start\": start,\n \"end\": end,\n \"start_time\": start_time,\n \"tenant_id\": get_current_tenant_id(),\n },\n )\n\n return {\"request_id\": task_id}\n\n\n@router.get(\"/admin/query-history/export-status\", tags=PUBLIC_API_TAGS)\ndef get_query_history_export_status(\n request_id: str,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> dict[str, str]:\n ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])\n\n task = get_task_with_id(db_session=db_session, task_id=request_id)\n\n if task:\n return {\"status\": task.status}\n\n # If task is None, then it's possible that the task has already finished processing.\n # Therefore, we should then check if the export file has already been stored inside of the file-store.\n # If that *also* doesn't exist, then we can return a 404.\n file_store = get_default_file_store()\n\n report_name = construct_query_history_report_name(request_id)\n has_file = file_store.has_file(\n file_id=report_name,\n file_origin=FileOrigin.QUERY_HISTORY_CSV,\n file_type=FileType.CSV,\n )\n\n if not has_file:\n raise HTTPException(\n HTTPStatus.NOT_FOUND,\n f\"No task with {request_id=} was found\",\n )\n\n return {\"status\": TaskStatus.SUCCESS}\n\n\n@router.get(\"/admin/query-history/download\", tags=PUBLIC_API_TAGS)\ndef download_query_history_csv(\n request_id: str,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> StreamingResponse:\n ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])\n\n report_name = construct_query_history_report_name(request_id)\n file_store = get_default_file_store()\n has_file = file_store.has_file(\n file_id=report_name,\n file_origin=FileOrigin.QUERY_HISTORY_CSV,\n file_type=FileType.CSV,\n )\n\n if has_file:\n try:\n csv_stream = file_store.read_file(report_name)\n except Exception as e:\n raise HTTPException(\n HTTPStatus.INTERNAL_SERVER_ERROR,\n f\"Failed to read query history file: {str(e)}\",\n )\n csv_stream.seek(0)\n return StreamingResponse(\n iter(csv_stream),\n media_type=FileType.CSV,\n headers={\"Content-Disposition\": f\"attachment;filename={report_name}\"},\n )\n\n # If the file doesn't exist yet, it may still be processing.\n # Therefore, we check the task queue to determine its status, if there is any.\n task = get_task_with_id(db_session=db_session, task_id=request_id)\n if not task:\n raise HTTPException(\n HTTPStatus.NOT_FOUND,\n f\"No task with {request_id=} was found\",\n )\n\n if task.status in [TaskStatus.STARTED, TaskStatus.PENDING]:\n raise HTTPException(\n HTTPStatus.ACCEPTED, f\"Task with {request_id=} is still being worked on\"\n )\n\n elif task.status == TaskStatus.FAILURE:\n raise HTTPException(\n HTTPStatus.INTERNAL_SERVER_ERROR,\n f\"Task with {request_id=} failed to be processed\",\n )\n else:\n # This is the final case in which `task.status == SUCCESS`\n raise RuntimeError(\n \"The task was marked as success, the file was not found in the file store; this is an internal error...\"\n )\n" }, { "path": "backend/ee/onyx/server/query_history/models.py", "content": "from datetime import datetime\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\n\nfrom ee.onyx.background.task_name_builders import QUERY_HISTORY_TASK_NAME_PREFIX\nfrom onyx.auth.users import get_display_email\nfrom onyx.background.task_utils import extract_task_id_from_query_history_report_name\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import QAFeedbackType\nfrom onyx.configs.constants import SessionType\nfrom onyx.db.enums import TaskStatus\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import FileRecord\nfrom onyx.db.models import TaskQueueState\n\n\nclass AbridgedSearchDoc(BaseModel):\n \"\"\"A subset of the info present in `SearchDoc`\"\"\"\n\n document_id: str\n semantic_identifier: str\n link: str | None\n\n\nclass MessageSnapshot(BaseModel):\n id: int\n message: str\n message_type: MessageType\n documents: list[AbridgedSearchDoc]\n feedback_type: QAFeedbackType | None\n feedback_text: str | None\n time_created: datetime\n\n @classmethod\n def build(cls, message: ChatMessage) -> \"MessageSnapshot\":\n latest_messages_feedback_obj = (\n message.chat_message_feedbacks[-1]\n if len(message.chat_message_feedbacks) > 0\n else None\n )\n feedback_type = (\n (\n QAFeedbackType.LIKE\n if latest_messages_feedback_obj.is_positive\n else QAFeedbackType.DISLIKE\n )\n if latest_messages_feedback_obj\n else None\n )\n feedback_text = (\n latest_messages_feedback_obj.feedback_text\n if latest_messages_feedback_obj\n else None\n )\n return cls(\n id=message.id,\n message=message.message,\n message_type=message.message_type,\n documents=[\n AbridgedSearchDoc(\n document_id=document.document_id,\n semantic_identifier=document.semantic_id,\n link=document.link,\n )\n for document in message.search_docs\n ],\n feedback_type=feedback_type,\n feedback_text=feedback_text,\n time_created=message.time_sent,\n )\n\n\nclass ChatSessionMinimal(BaseModel):\n id: UUID\n user_email: str\n name: str | None\n first_user_message: str\n first_ai_message: str\n assistant_id: int | None\n assistant_name: str | None\n time_created: datetime\n feedback_type: QAFeedbackType | None\n flow_type: SessionType\n conversation_length: int\n\n @classmethod\n def from_chat_session(cls, chat_session: ChatSession) -> \"ChatSessionMinimal\":\n first_user_message = next(\n (\n message.message\n for message in chat_session.messages\n if message.message_type == MessageType.USER\n ),\n \"\",\n )\n first_ai_message = next(\n (\n message.message\n for message in chat_session.messages\n if message.message_type == MessageType.ASSISTANT\n ),\n \"\",\n )\n\n list_of_message_feedbacks = [\n feedback.is_positive\n for message in chat_session.messages\n for feedback in message.chat_message_feedbacks\n ]\n session_feedback_type = None\n if list_of_message_feedbacks:\n if all(list_of_message_feedbacks):\n session_feedback_type = QAFeedbackType.LIKE\n elif not any(list_of_message_feedbacks):\n session_feedback_type = QAFeedbackType.DISLIKE\n else:\n session_feedback_type = QAFeedbackType.MIXED\n\n return cls(\n id=chat_session.id,\n user_email=get_display_email(\n chat_session.user.email if chat_session.user else None\n ),\n name=chat_session.description,\n first_user_message=first_user_message,\n first_ai_message=first_ai_message,\n assistant_id=chat_session.persona_id,\n assistant_name=(\n chat_session.persona.name if chat_session.persona else None\n ),\n time_created=chat_session.time_created,\n feedback_type=session_feedback_type,\n flow_type=(\n SessionType.SLACK if chat_session.onyxbot_flow else SessionType.CHAT\n ),\n conversation_length=len(\n [\n message\n for message in chat_session.messages\n if message.message_type != MessageType.SYSTEM\n ]\n ),\n )\n\n\nclass ChatSessionSnapshot(BaseModel):\n id: UUID\n user_email: str\n name: str | None\n messages: list[MessageSnapshot]\n assistant_id: int | None\n assistant_name: str | None\n time_created: datetime\n flow_type: SessionType\n\n\nclass QuestionAnswerPairSnapshot(BaseModel):\n chat_session_id: UUID\n # 1-indexed message number in the chat_session\n # e.g. the first message pair in the chat_session is 1, the second is 2, etc.\n message_pair_num: int\n user_message: str\n ai_response: str\n retrieved_documents: list[AbridgedSearchDoc]\n feedback_type: QAFeedbackType | None\n feedback_text: str | None\n persona_name: str | None\n user_email: str\n time_created: datetime\n flow_type: SessionType\n\n @classmethod\n def from_chat_session_snapshot(\n cls,\n chat_session_snapshot: ChatSessionSnapshot,\n ) -> list[\"QuestionAnswerPairSnapshot\"]:\n message_pairs: list[tuple[MessageSnapshot, MessageSnapshot]] = []\n for ind in range(1, len(chat_session_snapshot.messages), 2):\n message_pairs.append(\n (\n chat_session_snapshot.messages[ind - 1],\n chat_session_snapshot.messages[ind],\n )\n )\n\n return [\n cls(\n chat_session_id=chat_session_snapshot.id,\n message_pair_num=ind + 1,\n user_message=user_message.message,\n ai_response=ai_message.message,\n retrieved_documents=ai_message.documents,\n feedback_type=ai_message.feedback_type,\n feedback_text=ai_message.feedback_text,\n persona_name=chat_session_snapshot.assistant_name,\n user_email=get_display_email(chat_session_snapshot.user_email),\n time_created=user_message.time_created,\n flow_type=chat_session_snapshot.flow_type,\n )\n for ind, (user_message, ai_message) in enumerate(message_pairs)\n ]\n\n def to_json(self) -> dict[str, str | None]:\n return {\n \"chat_session_id\": str(self.chat_session_id),\n \"message_pair_num\": str(self.message_pair_num),\n \"user_message\": self.user_message,\n \"ai_response\": self.ai_response,\n \"retrieved_documents\": \"|\".join(\n [\n doc.link or doc.semantic_identifier\n for doc in self.retrieved_documents\n ]\n ),\n \"feedback_type\": self.feedback_type.value if self.feedback_type else \"\",\n \"feedback_text\": self.feedback_text or \"\",\n \"persona_name\": self.persona_name,\n \"user_email\": self.user_email,\n \"time_created\": str(self.time_created),\n \"flow_type\": self.flow_type,\n }\n\n\nclass QueryHistoryExport(BaseModel):\n task_id: str\n status: TaskStatus\n start: datetime\n end: datetime\n start_time: datetime\n\n @classmethod\n def from_task(\n cls,\n task_queue_state: TaskQueueState,\n ) -> \"QueryHistoryExport\":\n start_end = task_queue_state.task_name.removeprefix(\n f\"{QUERY_HISTORY_TASK_NAME_PREFIX}_\"\n )\n start, end = start_end.split(\"_\")\n\n if not task_queue_state.start_time:\n raise RuntimeError(\"The start time of the task must always be present\")\n\n return cls(\n task_id=task_queue_state.task_id,\n status=task_queue_state.status,\n start=datetime.fromisoformat(start),\n end=datetime.fromisoformat(end),\n start_time=task_queue_state.start_time,\n )\n\n @classmethod\n def from_file(\n cls,\n file: FileRecord,\n ) -> \"QueryHistoryExport\":\n if not file.file_metadata or not isinstance(file.file_metadata, dict):\n raise RuntimeError(\n \"The file metadata must be non-null, and must be of type `dict[str, str]`\"\n )\n\n metadata = QueryHistoryFileMetadata.model_validate(dict(file.file_metadata))\n task_id = extract_task_id_from_query_history_report_name(file.file_id)\n\n return cls(\n task_id=task_id,\n status=TaskStatus.SUCCESS,\n start=metadata.start,\n end=metadata.end,\n start_time=metadata.start_time,\n )\n\n\nclass QueryHistoryFileMetadata(BaseModel):\n start: datetime\n end: datetime\n start_time: datetime\n" }, { "path": "backend/ee/onyx/server/reporting/usage_export_api.py", "content": "from collections.abc import Generator\nfrom datetime import datetime\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Response\nfrom fastapi.responses import StreamingResponse\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.usage_export import get_all_usage_reports\nfrom ee.onyx.db.usage_export import get_usage_report_data\nfrom ee.onyx.db.usage_export import UsageReportMetadata\nfrom onyx.auth.users import current_admin_user\nfrom onyx.background.celery.versioned_apps.client import app as client_app\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.file_store.constants import STANDARD_CHUNK_SIZE\nfrom shared_configs.contextvars import get_current_tenant_id\n\nrouter = APIRouter()\n\n\nclass GenerateUsageReportParams(BaseModel):\n period_from: str | None = None\n period_to: str | None = None\n\n\n@router.post(\"/admin/usage-report\", status_code=204)\ndef generate_report(\n params: GenerateUsageReportParams,\n user: User = Depends(current_admin_user),\n) -> None:\n # Validate period parameters\n if params.period_from and params.period_to:\n try:\n datetime.fromisoformat(params.period_from)\n datetime.fromisoformat(params.period_to)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n tenant_id = get_current_tenant_id()\n client_app.send_task(\n OnyxCeleryTask.GENERATE_USAGE_REPORT_TASK,\n kwargs={\n \"tenant_id\": tenant_id,\n \"user_id\": str(user.id) if user else None,\n \"period_from\": params.period_from,\n \"period_to\": params.period_to,\n },\n )\n\n return None\n\n\n@router.get(\"/admin/usage-report/{report_name}\")\ndef read_usage_report(\n report_name: str,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session), # noqa: ARG001\n) -> Response:\n try:\n file = get_usage_report_data(report_name)\n except (ValueError, RuntimeError) as e:\n raise HTTPException(status_code=404, detail=str(e))\n\n def iterfile() -> Generator[bytes, None, None]:\n while True:\n chunk = file.read(STANDARD_CHUNK_SIZE)\n if not chunk:\n break\n yield chunk\n\n return StreamingResponse(\n content=iterfile(),\n media_type=\"application/zip\",\n headers={\"Content-Disposition\": f\"attachment; filename={report_name}\"},\n )\n\n\n@router.get(\"/admin/usage-report\")\ndef fetch_usage_reports(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[UsageReportMetadata]:\n try:\n return get_all_usage_reports(db_session)\n except ValueError as e:\n raise HTTPException(status_code=404, detail=str(e))\n" }, { "path": "backend/ee/onyx/server/reporting/usage_export_generation.py", "content": "import csv\nimport tempfile\nimport uuid\nimport zipfile\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom fastapi_users_db_sqlalchemy import UUID_ID\nfrom sqlalchemy import cast\nfrom sqlalchemy.dialects.postgresql import UUID\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.usage_export import get_all_empty_chat_message_entries\nfrom ee.onyx.db.usage_export import write_usage_report\nfrom ee.onyx.server.reporting.usage_export_models import UsageReportMetadata\nfrom ee.onyx.server.reporting.usage_export_models import UserSkeleton\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.db.models import User\nfrom onyx.db.users import get_all_users\nfrom onyx.file_store.constants import MAX_IN_MEMORY_SIZE\nfrom onyx.file_store.file_store import FileStore\nfrom onyx.file_store.file_store import get_default_file_store\n\n\ndef generate_chat_messages_report(\n db_session: Session,\n file_store: FileStore,\n report_id: str,\n period: tuple[datetime, datetime] | None,\n) -> str:\n file_name = f\"{report_id}_chat_sessions\"\n\n if period is None:\n period = (\n datetime.fromtimestamp(0, tz=timezone.utc),\n datetime.now(tz=timezone.utc),\n )\n else:\n # time-picker sends a time which is at the beginning of the day\n # so we need to add one day to the end time to make it inclusive\n period = (\n period[0],\n period[1] + timedelta(days=1),\n )\n\n with tempfile.SpooledTemporaryFile(\n max_size=MAX_IN_MEMORY_SIZE, mode=\"w+\"\n ) as temp_file:\n csvwriter = csv.writer(temp_file, delimiter=\",\")\n csvwriter.writerow(\n [\n \"session_id\",\n \"user_id\",\n \"flow_type\",\n \"time_sent\",\n \"assistant_name\",\n \"user_email\",\n \"number_of_tokens\",\n ]\n )\n for chat_message_skeleton_batch in get_all_empty_chat_message_entries(\n db_session, period\n ):\n for chat_message_skeleton in chat_message_skeleton_batch:\n csvwriter.writerow(\n [\n chat_message_skeleton.chat_session_id,\n chat_message_skeleton.user_id,\n chat_message_skeleton.flow_type,\n chat_message_skeleton.time_sent.isoformat(),\n chat_message_skeleton.assistant_name,\n chat_message_skeleton.user_email,\n chat_message_skeleton.number_of_tokens,\n ]\n )\n\n # after writing seek to beginning of buffer\n temp_file.seek(0)\n file_id = file_store.save_file(\n content=temp_file,\n display_name=file_name,\n file_origin=FileOrigin.GENERATED_REPORT,\n file_type=\"text/csv\",\n )\n\n return file_id\n\n\ndef generate_user_report(\n db_session: Session,\n file_store: FileStore,\n report_id: str,\n) -> str:\n file_name = f\"{report_id}_users\"\n\n with tempfile.SpooledTemporaryFile(\n max_size=MAX_IN_MEMORY_SIZE, mode=\"w+\"\n ) as temp_file:\n csvwriter = csv.writer(temp_file, delimiter=\",\")\n csvwriter.writerow([\"user_id\", \"is_active\"])\n\n users = get_all_users(db_session)\n for user in users:\n user_skeleton = UserSkeleton(\n user_id=str(user.id),\n is_active=user.is_active,\n )\n csvwriter.writerow([user_skeleton.user_id, user_skeleton.is_active])\n\n temp_file.seek(0)\n file_id = file_store.save_file(\n content=temp_file,\n display_name=file_name,\n file_origin=FileOrigin.GENERATED_REPORT,\n file_type=\"text/csv\",\n )\n\n return file_id\n\n\ndef create_new_usage_report(\n db_session: Session,\n user_id: UUID_ID | None, # None = auto-generated\n period: tuple[datetime, datetime] | None,\n) -> UsageReportMetadata:\n report_id = str(uuid.uuid4())\n file_store = get_default_file_store()\n\n messages_file_id = generate_chat_messages_report(\n db_session, file_store, report_id, period\n )\n users_file_id = generate_user_report(db_session, file_store, report_id)\n\n with tempfile.SpooledTemporaryFile(max_size=MAX_IN_MEMORY_SIZE) as zip_buffer:\n with zipfile.ZipFile(zip_buffer, \"a\", zipfile.ZIP_DEFLATED) as zip_file:\n # write messages\n chat_messages_tmpfile = file_store.read_file(\n messages_file_id, mode=\"b\", use_tempfile=True\n )\n zip_file.writestr(\n \"chat_messages.csv\",\n chat_messages_tmpfile.read(),\n )\n\n # write users\n users_tmpfile = file_store.read_file(\n users_file_id, mode=\"b\", use_tempfile=True\n )\n zip_file.writestr(\"users.csv\", users_tmpfile.read())\n\n zip_buffer.seek(0)\n\n # store zip blob to file_store\n report_name = f\"{datetime.now(tz=timezone.utc).strftime('%Y-%m-%d')}_{report_id}_usage_report.zip\"\n file_store.save_file(\n content=zip_buffer,\n display_name=report_name,\n file_origin=FileOrigin.GENERATED_REPORT,\n file_type=\"application/zip\",\n file_id=report_name,\n )\n\n # add report after zip file is written\n new_report = write_usage_report(db_session, report_name, user_id, period)\n\n # get user email\n requestor_user = (\n db_session.query(User)\n .filter(cast(User.id, UUID) == new_report.requestor_user_id)\n .one_or_none()\n if new_report.requestor_user_id\n else None\n )\n requestor_email = requestor_user.email if requestor_user else None\n\n return UsageReportMetadata(\n report_name=new_report.report_name,\n requestor=requestor_email,\n time_created=new_report.time_created,\n period_from=new_report.period_from,\n period_to=new_report.period_to,\n )\n" }, { "path": "backend/ee/onyx/server/reporting/usage_export_models.py", "content": "from datetime import datetime\nfrom enum import Enum\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\n\n\nclass FlowType(str, Enum):\n CHAT = \"chat\"\n SLACK = \"slack\"\n\n\nclass ChatMessageSkeleton(BaseModel):\n message_id: int\n chat_session_id: UUID\n user_id: str | None\n flow_type: FlowType\n time_sent: datetime\n assistant_name: str | None\n user_email: str | None\n number_of_tokens: int\n\n\nclass UserSkeleton(BaseModel):\n user_id: str\n is_active: bool\n\n\nclass UsageReportMetadata(BaseModel):\n report_name: str\n requestor: str | None\n time_created: datetime\n period_from: datetime | None # None = All time\n period_to: datetime | None\n" }, { "path": "backend/ee/onyx/server/scim/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/scim/api.py", "content": "\"\"\"SCIM 2.0 API endpoints (RFC 7644).\n\nThis module provides the FastAPI router for SCIM service discovery,\nUser CRUD, and Group CRUD. Identity providers (Okta, Azure AD) call\nthese endpoints to provision and manage users and groups.\n\nService discovery endpoints are unauthenticated — IdPs may probe them\nbefore bearer token configuration is complete. All other endpoints\nrequire a valid SCIM bearer token.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom uuid import UUID\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import FastAPI\nfrom fastapi import Query\nfrom fastapi import Request\nfrom fastapi import Response\nfrom fastapi.responses import JSONResponse\nfrom fastapi_users.password import PasswordHelper\nfrom sqlalchemy import func\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.scim import ScimDAL\nfrom ee.onyx.server.scim.auth import ScimAuthError\nfrom ee.onyx.server.scim.auth import verify_scim_token\nfrom ee.onyx.server.scim.filtering import parse_scim_filter\nfrom ee.onyx.server.scim.models import SCIM_LIST_RESPONSE_SCHEMA\nfrom ee.onyx.server.scim.models import ScimError\nfrom ee.onyx.server.scim.models import ScimGroupMember\nfrom ee.onyx.server.scim.models import ScimGroupResource\nfrom ee.onyx.server.scim.models import ScimListResponse\nfrom ee.onyx.server.scim.models import ScimMappingFields\nfrom ee.onyx.server.scim.models import ScimName\nfrom ee.onyx.server.scim.models import ScimPatchRequest\nfrom ee.onyx.server.scim.models import ScimServiceProviderConfig\nfrom ee.onyx.server.scim.models import ScimUserResource\nfrom ee.onyx.server.scim.patch import apply_group_patch\nfrom ee.onyx.server.scim.patch import apply_user_patch\nfrom ee.onyx.server.scim.patch import ScimPatchError\nfrom ee.onyx.server.scim.providers.base import get_default_provider\nfrom ee.onyx.server.scim.providers.base import ScimProvider\nfrom ee.onyx.server.scim.providers.base import serialize_emails\nfrom ee.onyx.server.scim.schema_definitions import ENTERPRISE_USER_SCHEMA_DEF\nfrom ee.onyx.server.scim.schema_definitions import GROUP_RESOURCE_TYPE\nfrom ee.onyx.server.scim.schema_definitions import GROUP_SCHEMA_DEF\nfrom ee.onyx.server.scim.schema_definitions import SERVICE_PROVIDER_CONFIG\nfrom ee.onyx.server.scim.schema_definitions import USER_RESOURCE_TYPE\nfrom ee.onyx.server.scim.schema_definitions import USER_SCHEMA_DEF\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import AccountType\nfrom onyx.db.enums import GrantSource\nfrom onyx.db.enums import Permission\nfrom onyx.db.models import ScimToken\nfrom onyx.db.models import ScimUserMapping\nfrom onyx.db.models import User\nfrom onyx.db.models import UserGroup\nfrom onyx.db.models import UserRole\nfrom onyx.db.permissions import recompute_permissions_for_group__no_commit\nfrom onyx.db.permissions import recompute_user_permissions__no_commit\nfrom onyx.db.users import assign_user_to_default_groups__no_commit\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\nlogger = setup_logger()\n\n# Group names reserved for system default groups (seeded by migration).\n_RESERVED_GROUP_NAMES = frozenset({\"Admin\", \"Basic\"})\n\n\nclass ScimJSONResponse(JSONResponse):\n \"\"\"JSONResponse with Content-Type: application/scim+json (RFC 7644 §3.1).\"\"\"\n\n media_type = \"application/scim+json\"\n\n\n# NOTE: All URL paths in this router (/ServiceProviderConfig, /ResourceTypes,\n# /Schemas, /Users, /Groups) are mandated by the SCIM spec (RFC 7643/7644).\n# IdPs like Okta and Azure AD hardcode these exact paths, so they cannot be\n# changed to kebab-case.\n\n\nscim_router = APIRouter(prefix=\"/scim/v2\", tags=[\"SCIM\"])\n\n_pw_helper = PasswordHelper()\n\n\ndef register_scim_exception_handlers(app: FastAPI) -> None:\n \"\"\"Register SCIM-specific exception handlers on the FastAPI app.\n\n Call this after ``app.include_router(scim_router)`` so that auth\n failures from ``verify_scim_token`` return RFC 7644 §3.12 error\n envelopes (with ``schemas`` and ``status`` fields) instead of\n FastAPI's default ``{\"detail\": \"...\"}`` format.\n \"\"\"\n\n @app.exception_handler(ScimAuthError)\n async def _handle_scim_auth_error(\n _request: Request, exc: ScimAuthError\n ) -> ScimJSONResponse:\n return _scim_error_response(exc.status_code, exc.detail)\n\n\ndef _get_provider(\n _token: ScimToken = Depends(verify_scim_token),\n) -> ScimProvider:\n \"\"\"Resolve the SCIM provider for the current request.\n\n Currently returns OktaProvider for all requests. When multi-provider\n support is added (ENG-3652), this will resolve based on token metadata\n or tenant configuration — no endpoint changes required.\n \"\"\"\n return get_default_provider()\n\n\n# ---------------------------------------------------------------------------\n# Service Discovery Endpoints (unauthenticated)\n# ---------------------------------------------------------------------------\n\n\n@scim_router.get(\"/ServiceProviderConfig\")\ndef get_service_provider_config() -> ScimServiceProviderConfig:\n \"\"\"Advertise supported SCIM features (RFC 7643 §5).\"\"\"\n return SERVICE_PROVIDER_CONFIG\n\n\n@scim_router.get(\"/ResourceTypes\")\ndef get_resource_types() -> ScimJSONResponse:\n \"\"\"List available SCIM resource types (RFC 7643 §6).\n\n Wrapped in a ListResponse envelope (RFC 7644 §3.4.2) because IdPs\n like Entra ID expect a JSON object, not a bare array.\n \"\"\"\n resources = [USER_RESOURCE_TYPE, GROUP_RESOURCE_TYPE]\n return ScimJSONResponse(\n content={\n \"schemas\": [SCIM_LIST_RESPONSE_SCHEMA],\n \"totalResults\": len(resources),\n \"Resources\": [\n r.model_dump(exclude_none=True, by_alias=True) for r in resources\n ],\n }\n )\n\n\n@scim_router.get(\"/Schemas\")\ndef get_schemas() -> ScimJSONResponse:\n \"\"\"Return SCIM schema definitions (RFC 7643 §7).\n\n Wrapped in a ListResponse envelope (RFC 7644 §3.4.2) because IdPs\n like Entra ID expect a JSON object, not a bare array.\n \"\"\"\n schemas = [USER_SCHEMA_DEF, GROUP_SCHEMA_DEF, ENTERPRISE_USER_SCHEMA_DEF]\n return ScimJSONResponse(\n content={\n \"schemas\": [SCIM_LIST_RESPONSE_SCHEMA],\n \"totalResults\": len(schemas),\n \"Resources\": [s.model_dump(exclude_none=True) for s in schemas],\n }\n )\n\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _scim_error_response(status: int, detail: str) -> ScimJSONResponse:\n \"\"\"Build a SCIM-compliant error response (RFC 7644 §3.12).\"\"\"\n logger.warning(\"SCIM error response: status=%s detail=%s\", status, detail)\n body = ScimError(status=str(status), detail=detail)\n return ScimJSONResponse(\n status_code=status,\n content=body.model_dump(exclude_none=True),\n )\n\n\ndef _parse_excluded_attributes(raw: str | None) -> set[str]:\n \"\"\"Parse the ``excludedAttributes`` query parameter (RFC 7644 §3.4.2.5).\n\n Returns a set of lowercased attribute names to omit from responses.\n \"\"\"\n if not raw:\n return set()\n return {attr.strip().lower() for attr in raw.split(\",\") if attr.strip()}\n\n\ndef _apply_exclusions(\n resource: ScimUserResource | ScimGroupResource,\n excluded: set[str],\n) -> dict:\n \"\"\"Serialize a SCIM resource, omitting attributes the IdP excluded.\n\n RFC 7644 §3.4.2.5 lets the IdP pass ``?excludedAttributes=groups,emails``\n to reduce response payload size. We strip those fields after serialization\n so the rest of the pipeline doesn't need to know about them.\n \"\"\"\n data = resource.model_dump(exclude_none=True, by_alias=True)\n for attr in excluded:\n # Match case-insensitively against the camelCase field names\n keys_to_remove = [k for k in data if k.lower() == attr]\n for k in keys_to_remove:\n del data[k]\n return data\n\n\ndef _check_seat_availability(dal: ScimDAL) -> str | None:\n \"\"\"Return an error message if seat limit is reached, else None.\"\"\"\n check_fn = fetch_ee_implementation_or_noop(\n \"onyx.db.license\", \"check_seat_availability\", None\n )\n if check_fn is None:\n return None\n result = check_fn(dal.session, seats_needed=1)\n if not result.available:\n return result.error_message or \"Seat limit reached\"\n return None\n\n\ndef _fetch_user_or_404(user_id: str, dal: ScimDAL) -> User | ScimJSONResponse:\n \"\"\"Parse *user_id* as UUID, look up the user, or return a 404 error.\"\"\"\n try:\n uid = UUID(user_id)\n except ValueError:\n return _scim_error_response(404, f\"User {user_id} not found\")\n user = dal.get_user(uid)\n if not user:\n return _scim_error_response(404, f\"User {user_id} not found\")\n return user\n\n\ndef _scim_name_to_str(name: ScimName | None) -> str | None:\n \"\"\"Extract a display name string from a SCIM name object.\n\n Returns None if no name is provided, so the caller can decide\n whether to update the user's personal_name.\n \"\"\"\n if not name:\n return None\n # If the client explicitly provides ``formatted``, prefer it — the client\n # knows what display string it wants. Otherwise build from components.\n if name.formatted:\n return name.formatted\n parts = \" \".join(part for part in [name.givenName, name.familyName] if part)\n return parts or None\n\n\ndef _scim_resource_response(\n resource: ScimUserResource | ScimGroupResource | ScimListResponse,\n status_code: int = 200,\n) -> ScimJSONResponse:\n \"\"\"Serialize a SCIM resource as ``application/scim+json``.\"\"\"\n content = resource.model_dump(exclude_none=True, by_alias=True)\n return ScimJSONResponse(\n status_code=status_code,\n content=content,\n )\n\n\ndef _build_list_response(\n resources: list[ScimUserResource | ScimGroupResource],\n total: int,\n start_index: int,\n count: int,\n excluded: set[str] | None = None,\n) -> ScimListResponse | ScimJSONResponse:\n \"\"\"Build a SCIM list response, optionally applying attribute exclusions.\n\n RFC 7644 §3.4.2.5 — IdPs may request certain attributes be omitted via\n the ``excludedAttributes`` query parameter.\n \"\"\"\n if excluded:\n envelope = ScimListResponse(\n totalResults=total,\n startIndex=start_index,\n itemsPerPage=count,\n )\n data = envelope.model_dump(exclude_none=True)\n data[\"Resources\"] = [_apply_exclusions(r, excluded) for r in resources]\n return ScimJSONResponse(content=data)\n\n return _scim_resource_response(\n ScimListResponse(\n totalResults=total,\n startIndex=start_index,\n itemsPerPage=count,\n Resources=resources,\n )\n )\n\n\ndef _extract_enterprise_fields(\n resource: ScimUserResource,\n) -> tuple[str | None, str | None]:\n \"\"\"Extract department and manager from enterprise extension.\"\"\"\n ext = resource.enterprise_extension\n if not ext:\n return None, None\n department = ext.department\n manager = ext.manager.value if ext.manager else None\n return department, manager\n\n\ndef _mapping_to_fields(\n mapping: ScimUserMapping | None,\n) -> ScimMappingFields | None:\n \"\"\"Extract round-trip fields from a SCIM user mapping.\"\"\"\n if not mapping:\n return None\n return ScimMappingFields(\n department=mapping.department,\n manager=mapping.manager,\n given_name=mapping.given_name,\n family_name=mapping.family_name,\n scim_emails_json=mapping.scim_emails_json,\n )\n\n\ndef _fields_from_resource(resource: ScimUserResource) -> ScimMappingFields:\n \"\"\"Build mapping fields from an incoming SCIM user resource.\"\"\"\n department, manager = _extract_enterprise_fields(resource)\n return ScimMappingFields(\n department=department,\n manager=manager,\n given_name=resource.name.givenName if resource.name else None,\n family_name=resource.name.familyName if resource.name else None,\n scim_emails_json=serialize_emails(resource.emails),\n )\n\n\n# ---------------------------------------------------------------------------\n# User CRUD (RFC 7644 §3)\n# ---------------------------------------------------------------------------\n\n\n@scim_router.get(\"/Users\", response_model=None)\ndef list_users(\n filter: str | None = Query(None),\n excludedAttributes: str | None = None,\n startIndex: int = Query(1, ge=1),\n count: int = Query(100, ge=0, le=500),\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimListResponse | ScimJSONResponse:\n \"\"\"List users with optional SCIM filter and pagination.\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n dal.commit()\n\n try:\n scim_filter = parse_scim_filter(filter)\n except ValueError as e:\n return _scim_error_response(400, str(e))\n\n try:\n users_with_mappings, total = dal.list_users(scim_filter, startIndex, count)\n except ValueError as e:\n return _scim_error_response(400, str(e))\n\n user_groups_map = dal.get_users_groups_batch([u.id for u, _ in users_with_mappings])\n resources: list[ScimUserResource | ScimGroupResource] = [\n provider.build_user_resource(\n user,\n mapping.external_id if mapping else None,\n groups=user_groups_map.get(user.id, []),\n scim_username=mapping.scim_username if mapping else None,\n fields=_mapping_to_fields(mapping),\n )\n for user, mapping in users_with_mappings\n ]\n\n return _build_list_response(\n resources,\n total,\n startIndex,\n count,\n excluded=_parse_excluded_attributes(excludedAttributes),\n )\n\n\n@scim_router.get(\"/Users/{user_id}\", response_model=None)\ndef get_user(\n user_id: str,\n excludedAttributes: str | None = None,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimUserResource | ScimJSONResponse:\n \"\"\"Get a single user by ID.\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n dal.commit()\n\n result = _fetch_user_or_404(user_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n user = result\n\n mapping = dal.get_user_mapping_by_user_id(user.id)\n\n resource = provider.build_user_resource(\n user,\n mapping.external_id if mapping else None,\n groups=dal.get_user_groups(user.id),\n scim_username=mapping.scim_username if mapping else None,\n fields=_mapping_to_fields(mapping),\n )\n\n # RFC 7644 §3.4.2.5 — IdP may request certain attributes be omitted\n excluded = _parse_excluded_attributes(excludedAttributes)\n if excluded:\n return ScimJSONResponse(content=_apply_exclusions(resource, excluded))\n\n return _scim_resource_response(resource)\n\n\n@scim_router.post(\"/Users\", status_code=201, response_model=None)\ndef create_user(\n user_resource: ScimUserResource,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimUserResource | ScimJSONResponse:\n \"\"\"Create a new user from a SCIM provisioning request.\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n email = user_resource.userName.strip()\n\n # Check for existing user — if they exist but aren't SCIM-managed yet,\n # link them to the IdP rather than rejecting with 409.\n external_id: str | None = user_resource.externalId\n scim_username: str = user_resource.userName.strip()\n fields: ScimMappingFields = _fields_from_resource(user_resource)\n\n existing_user = dal.get_user_by_email(email)\n if existing_user:\n existing_mapping = dal.get_user_mapping_by_user_id(existing_user.id)\n if existing_mapping:\n return _scim_error_response(409, f\"User with email {email} already exists\")\n\n # Adopt pre-existing user into SCIM management.\n # Reactivating a deactivated user consumes a seat, so enforce the\n # seat limit the same way replace_user does.\n if user_resource.active and not existing_user.is_active:\n seat_error = _check_seat_availability(dal)\n if seat_error:\n return _scim_error_response(403, seat_error)\n\n personal_name = _scim_name_to_str(user_resource.name)\n dal.update_user(\n existing_user,\n is_active=user_resource.active,\n **({\"personal_name\": personal_name} if personal_name else {}),\n )\n\n try:\n dal.create_user_mapping(\n external_id=external_id,\n user_id=existing_user.id,\n scim_username=scim_username,\n fields=fields,\n )\n dal.commit()\n except IntegrityError:\n dal.rollback()\n return _scim_error_response(\n 409, f\"User with email {email} already has a SCIM mapping\"\n )\n\n return _scim_resource_response(\n provider.build_user_resource(\n existing_user,\n external_id,\n scim_username=scim_username,\n fields=fields,\n ),\n status_code=201,\n )\n\n # Only enforce seat limit for net-new users — adopting a pre-existing\n # user doesn't consume a new seat.\n seat_error = _check_seat_availability(dal)\n if seat_error:\n return _scim_error_response(403, seat_error)\n\n # Create user with a random password (SCIM users authenticate via IdP)\n personal_name = _scim_name_to_str(user_resource.name)\n user = User(\n email=email,\n hashed_password=_pw_helper.hash(_pw_helper.generate()),\n role=UserRole.BASIC,\n account_type=AccountType.STANDARD,\n is_active=user_resource.active,\n is_verified=True,\n personal_name=personal_name,\n )\n\n try:\n dal.add_user(user)\n except IntegrityError:\n dal.rollback()\n return _scim_error_response(409, f\"User with email {email} already exists\")\n\n # Always create a SCIM mapping so that the user is marked as\n # SCIM-managed. externalId may be None (RFC 7643 says it's optional).\n try:\n dal.create_user_mapping(\n external_id=external_id,\n user_id=user.id,\n scim_username=scim_username,\n fields=fields,\n )\n except IntegrityError:\n dal.rollback()\n return _scim_error_response(\n 409, f\"User with email {email} already has a SCIM mapping\"\n )\n\n # Assign user to default group BEFORE commit so everything is atomic.\n # If this fails, the entire user creation rolls back and IdP can retry.\n try:\n assign_user_to_default_groups__no_commit(db_session, user)\n except Exception:\n dal.rollback()\n logger.exception(f\"Failed to assign SCIM user {email} to default groups\")\n return _scim_error_response(\n 500, f\"Failed to assign user {email} to default group\"\n )\n\n dal.commit()\n\n return _scim_resource_response(\n provider.build_user_resource(\n user,\n external_id,\n scim_username=scim_username,\n fields=fields,\n ),\n status_code=201,\n )\n\n\n@scim_router.put(\"/Users/{user_id}\", response_model=None)\ndef replace_user(\n user_id: str,\n user_resource: ScimUserResource,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimUserResource | ScimJSONResponse:\n \"\"\"Replace a user entirely (RFC 7644 §3.5.1).\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n result = _fetch_user_or_404(user_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n user = result\n\n # Handle activation (need seat check) / deactivation\n is_reactivation = user_resource.active and not user.is_active\n if is_reactivation:\n seat_error = _check_seat_availability(dal)\n if seat_error:\n return _scim_error_response(403, seat_error)\n\n personal_name = _scim_name_to_str(user_resource.name)\n\n dal.update_user(\n user,\n email=user_resource.userName.strip(),\n is_active=user_resource.active,\n personal_name=personal_name,\n )\n\n # Reconcile default-group membership on reactivation\n if is_reactivation:\n assign_user_to_default_groups__no_commit(\n db_session, user, is_admin=(user.role == UserRole.ADMIN)\n )\n\n new_external_id = user_resource.externalId\n scim_username = user_resource.userName.strip()\n fields = _fields_from_resource(user_resource)\n dal.sync_user_external_id(\n user.id,\n new_external_id,\n scim_username=scim_username,\n fields=fields,\n )\n\n dal.commit()\n\n return _scim_resource_response(\n provider.build_user_resource(\n user,\n new_external_id,\n groups=dal.get_user_groups(user.id),\n scim_username=scim_username,\n fields=fields,\n )\n )\n\n\n@scim_router.patch(\"/Users/{user_id}\", response_model=None)\ndef patch_user(\n user_id: str,\n patch_request: ScimPatchRequest,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimUserResource | ScimJSONResponse:\n \"\"\"Partially update a user (RFC 7644 §3.5.2).\n\n This is the primary endpoint for user deprovisioning — Okta sends\n ``PATCH {\"active\": false}`` rather than DELETE.\n \"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n result = _fetch_user_or_404(user_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n user = result\n\n mapping = dal.get_user_mapping_by_user_id(user.id)\n external_id = mapping.external_id if mapping else None\n current_scim_username = mapping.scim_username if mapping else None\n current_fields = _mapping_to_fields(mapping)\n\n current = provider.build_user_resource(\n user,\n external_id,\n groups=dal.get_user_groups(user.id),\n scim_username=current_scim_username,\n fields=current_fields,\n )\n\n try:\n patched, ent_data = apply_user_patch(\n patch_request.Operations, current, provider.ignored_patch_paths\n )\n except ScimPatchError as e:\n return _scim_error_response(e.status, e.detail)\n\n # Apply changes back to the DB model\n is_reactivation = patched.active and not user.is_active\n if patched.active != user.is_active:\n if patched.active:\n seat_error = _check_seat_availability(dal)\n if seat_error:\n return _scim_error_response(403, seat_error)\n\n # Track the scim_username — if userName was patched, update it\n new_scim_username = patched.userName.strip() if patched.userName else None\n\n # If displayName was explicitly patched (different from the original), use\n # it as personal_name directly. Otherwise, derive from name components.\n personal_name: str | None\n if patched.displayName and patched.displayName != current.displayName:\n personal_name = patched.displayName\n else:\n personal_name = _scim_name_to_str(patched.name)\n\n dal.update_user(\n user,\n email=(\n patched.userName.strip()\n if patched.userName.strip().lower() != user.email.lower()\n else None\n ),\n is_active=patched.active if patched.active != user.is_active else None,\n personal_name=personal_name,\n )\n\n # Reconcile default-group membership on reactivation\n if is_reactivation:\n assign_user_to_default_groups__no_commit(\n db_session, user, is_admin=(user.role == UserRole.ADMIN)\n )\n\n # Build updated fields by merging PATCH enterprise data with current values\n cf = current_fields or ScimMappingFields()\n fields = ScimMappingFields(\n department=ent_data.get(\"department\", cf.department),\n manager=ent_data.get(\"manager\", cf.manager),\n given_name=patched.name.givenName if patched.name else cf.given_name,\n family_name=patched.name.familyName if patched.name else cf.family_name,\n scim_emails_json=(\n serialize_emails(patched.emails)\n if patched.emails is not None\n else cf.scim_emails_json\n ),\n )\n\n dal.sync_user_external_id(\n user.id,\n patched.externalId,\n scim_username=new_scim_username,\n fields=fields,\n )\n\n dal.commit()\n\n return _scim_resource_response(\n provider.build_user_resource(\n user,\n patched.externalId,\n groups=dal.get_user_groups(user.id),\n scim_username=new_scim_username,\n fields=fields,\n )\n )\n\n\n@scim_router.delete(\"/Users/{user_id}\", status_code=204, response_model=None)\ndef delete_user(\n user_id: str,\n _token: ScimToken = Depends(verify_scim_token),\n db_session: Session = Depends(get_session),\n) -> Response | ScimJSONResponse:\n \"\"\"Delete a user (RFC 7644 §3.6).\n\n Deactivates the user and removes the SCIM mapping. Note that Okta\n typically uses PATCH active=false instead of DELETE.\n A second DELETE returns 404 per RFC 7644 §3.6.\n \"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n result = _fetch_user_or_404(user_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n user = result\n\n # If no SCIM mapping exists, the user was already deleted from\n # SCIM's perspective — return 404 per RFC 7644 §3.6.\n mapping = dal.get_user_mapping_by_user_id(user.id)\n if not mapping:\n return _scim_error_response(404, f\"User {user_id} not found\")\n\n dal.deactivate_user(user)\n dal.delete_user_mapping(mapping.id)\n\n dal.commit()\n\n return Response(status_code=204)\n\n\n# ---------------------------------------------------------------------------\n# Group helpers\n# ---------------------------------------------------------------------------\n\n\ndef _fetch_group_or_404(group_id: str, dal: ScimDAL) -> UserGroup | ScimJSONResponse:\n \"\"\"Parse *group_id* as int, look up the group, or return a 404 error.\"\"\"\n try:\n gid = int(group_id)\n except ValueError:\n return _scim_error_response(404, f\"Group {group_id} not found\")\n group = dal.get_group(gid)\n if not group:\n return _scim_error_response(404, f\"Group {group_id} not found\")\n return group\n\n\ndef _parse_member_uuids(\n members: list[ScimGroupMember],\n) -> tuple[list[UUID], str | None]:\n \"\"\"Parse member value strings to UUIDs.\n\n Returns (uuid_list, error_message). error_message is None on success.\n \"\"\"\n uuids: list[UUID] = []\n for m in members:\n try:\n uuids.append(UUID(m.value))\n except ValueError:\n return [], f\"Invalid member ID: {m.value}\"\n return uuids, None\n\n\ndef _validate_and_parse_members(\n members: list[ScimGroupMember], dal: ScimDAL\n) -> tuple[list[UUID], str | None]:\n \"\"\"Parse and validate member UUIDs exist in the database.\n\n Returns (uuid_list, error_message). error_message is None on success.\n \"\"\"\n uuids, err = _parse_member_uuids(members)\n if err:\n return [], err\n\n if uuids:\n missing = dal.validate_member_ids(uuids)\n if missing:\n return [], f\"Member(s) not found: {', '.join(str(u) for u in missing)}\"\n\n return uuids, None\n\n\n# ---------------------------------------------------------------------------\n# Group CRUD (RFC 7644 §3)\n# ---------------------------------------------------------------------------\n\n\n@scim_router.get(\"/Groups\", response_model=None)\ndef list_groups(\n filter: str | None = Query(None),\n excludedAttributes: str | None = None,\n startIndex: int = Query(1, ge=1),\n count: int = Query(100, ge=0, le=500),\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimListResponse | ScimJSONResponse:\n \"\"\"List groups with optional SCIM filter and pagination.\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n dal.commit()\n\n try:\n scim_filter = parse_scim_filter(filter)\n except ValueError as e:\n return _scim_error_response(400, str(e))\n\n try:\n groups_with_ext_ids, total = dal.list_groups(scim_filter, startIndex, count)\n except ValueError as e:\n return _scim_error_response(400, str(e))\n\n resources: list[ScimUserResource | ScimGroupResource] = [\n provider.build_group_resource(group, dal.get_group_members(group.id), ext_id)\n for group, ext_id in groups_with_ext_ids\n ]\n\n return _build_list_response(\n resources,\n total,\n startIndex,\n count,\n excluded=_parse_excluded_attributes(excludedAttributes),\n )\n\n\n@scim_router.get(\"/Groups/{group_id}\", response_model=None)\ndef get_group(\n group_id: str,\n excludedAttributes: str | None = None,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimGroupResource | ScimJSONResponse:\n \"\"\"Get a single group by ID.\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n dal.commit()\n\n result = _fetch_group_or_404(group_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n group = result\n\n mapping = dal.get_group_mapping_by_group_id(group.id)\n members = dal.get_group_members(group.id)\n\n resource = provider.build_group_resource(\n group, members, mapping.external_id if mapping else None\n )\n\n # RFC 7644 §3.4.2.5 — IdP may request certain attributes be omitted\n excluded = _parse_excluded_attributes(excludedAttributes)\n if excluded:\n return ScimJSONResponse(content=_apply_exclusions(resource, excluded))\n\n return _scim_resource_response(resource)\n\n\n@scim_router.post(\"/Groups\", status_code=201, response_model=None)\ndef create_group(\n group_resource: ScimGroupResource,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimGroupResource | ScimJSONResponse:\n \"\"\"Create a new group from a SCIM provisioning request.\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n if group_resource.displayName in _RESERVED_GROUP_NAMES:\n return _scim_error_response(\n 409, f\"'{group_resource.displayName}' is a reserved group name.\"\n )\n\n if dal.get_group_by_name(group_resource.displayName):\n return _scim_error_response(\n 409, f\"Group with name '{group_resource.displayName}' already exists\"\n )\n\n member_uuids, err = _validate_and_parse_members(group_resource.members, dal)\n if err:\n return _scim_error_response(400, err)\n\n db_group = UserGroup(\n name=group_resource.displayName,\n is_up_to_date=True,\n time_last_modified_by_user=func.now(),\n )\n try:\n dal.add_group(db_group)\n except IntegrityError:\n dal.rollback()\n return _scim_error_response(\n 409, f\"Group with name '{group_resource.displayName}' already exists\"\n )\n\n # Every group gets the \"basic\" permission by default.\n dal.add_permission_grant_to_group(\n group_id=db_group.id,\n permission=Permission.BASIC_ACCESS,\n grant_source=GrantSource.SYSTEM,\n )\n\n dal.upsert_group_members(db_group.id, member_uuids)\n\n # Recompute permissions for initial members.\n recompute_user_permissions__no_commit(member_uuids, db_session)\n\n external_id = group_resource.externalId\n if external_id:\n dal.create_group_mapping(external_id=external_id, user_group_id=db_group.id)\n\n dal.commit()\n\n members = dal.get_group_members(db_group.id)\n return _scim_resource_response(\n provider.build_group_resource(db_group, members, external_id),\n status_code=201,\n )\n\n\n@scim_router.put(\"/Groups/{group_id}\", response_model=None)\ndef replace_group(\n group_id: str,\n group_resource: ScimGroupResource,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimGroupResource | ScimJSONResponse:\n \"\"\"Replace a group entirely (RFC 7644 §3.5.1).\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n result = _fetch_group_or_404(group_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n group = result\n\n if group.name in _RESERVED_GROUP_NAMES and group_resource.displayName != group.name:\n return _scim_error_response(\n 409, f\"'{group.name}' is a reserved group name and cannot be renamed.\"\n )\n\n if (\n group_resource.displayName in _RESERVED_GROUP_NAMES\n and group_resource.displayName != group.name\n ):\n return _scim_error_response(\n 409, f\"'{group_resource.displayName}' is a reserved group name.\"\n )\n\n member_uuids, err = _validate_and_parse_members(group_resource.members, dal)\n if err:\n return _scim_error_response(400, err)\n\n # Capture old member IDs before replacing so we can recompute their\n # permissions after they are removed from the group.\n old_member_ids = {uid for uid, _ in dal.get_group_members(group.id)}\n\n dal.update_group(group, name=group_resource.displayName)\n dal.replace_group_members(group.id, member_uuids)\n dal.sync_group_external_id(group.id, group_resource.externalId)\n\n # Recompute permissions for current members (batch) and removed members.\n recompute_permissions_for_group__no_commit(group.id, db_session)\n removed_ids = list(old_member_ids - set(member_uuids))\n recompute_user_permissions__no_commit(removed_ids, db_session)\n\n dal.commit()\n\n members = dal.get_group_members(group.id)\n return _scim_resource_response(\n provider.build_group_resource(group, members, group_resource.externalId)\n )\n\n\n@scim_router.patch(\"/Groups/{group_id}\", response_model=None)\ndef patch_group(\n group_id: str,\n patch_request: ScimPatchRequest,\n _token: ScimToken = Depends(verify_scim_token),\n provider: ScimProvider = Depends(_get_provider),\n db_session: Session = Depends(get_session),\n) -> ScimGroupResource | ScimJSONResponse:\n \"\"\"Partially update a group (RFC 7644 §3.5.2).\n\n Handles member add/remove operations from Okta and Azure AD.\n \"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n result = _fetch_group_or_404(group_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n group = result\n\n mapping = dal.get_group_mapping_by_group_id(group.id)\n external_id = mapping.external_id if mapping else None\n\n current_members = dal.get_group_members(group.id)\n current = provider.build_group_resource(group, current_members, external_id)\n\n try:\n patched, added_ids, removed_ids = apply_group_patch(\n patch_request.Operations, current, provider.ignored_patch_paths\n )\n except ScimPatchError as e:\n return _scim_error_response(e.status, e.detail)\n\n new_name = patched.displayName if patched.displayName != group.name else None\n\n if group.name in _RESERVED_GROUP_NAMES and new_name:\n return _scim_error_response(\n 409, f\"'{group.name}' is a reserved group name and cannot be renamed.\"\n )\n\n if new_name and new_name in _RESERVED_GROUP_NAMES:\n return _scim_error_response(409, f\"'{new_name}' is a reserved group name.\")\n\n dal.update_group(group, name=new_name)\n\n affected_uuids: list[UUID] = []\n\n if added_ids:\n add_uuids = [UUID(mid) for mid in added_ids if _is_valid_uuid(mid)]\n if add_uuids:\n missing = dal.validate_member_ids(add_uuids)\n if missing:\n return _scim_error_response(\n 400,\n f\"Member(s) not found: {', '.join(str(u) for u in missing)}\",\n )\n dal.upsert_group_members(group.id, add_uuids)\n affected_uuids.extend(add_uuids)\n\n if removed_ids:\n remove_uuids = [UUID(mid) for mid in removed_ids if _is_valid_uuid(mid)]\n dal.remove_group_members(group.id, remove_uuids)\n affected_uuids.extend(remove_uuids)\n\n # Recompute permissions for all users whose group membership changed.\n recompute_user_permissions__no_commit(affected_uuids, db_session)\n\n dal.sync_group_external_id(group.id, patched.externalId)\n dal.commit()\n\n members = dal.get_group_members(group.id)\n return _scim_resource_response(\n provider.build_group_resource(group, members, patched.externalId)\n )\n\n\n@scim_router.delete(\"/Groups/{group_id}\", status_code=204, response_model=None)\ndef delete_group(\n group_id: str,\n _token: ScimToken = Depends(verify_scim_token),\n db_session: Session = Depends(get_session),\n) -> Response | ScimJSONResponse:\n \"\"\"Delete a group (RFC 7644 §3.6).\"\"\"\n dal = ScimDAL(db_session)\n dal.update_token_last_used(_token.id)\n\n result = _fetch_group_or_404(group_id, dal)\n if isinstance(result, ScimJSONResponse):\n return result\n group = result\n\n if group.name in _RESERVED_GROUP_NAMES:\n return _scim_error_response(409, f\"'{group.name}' is a reserved group name.\")\n\n # Capture member IDs before deletion so we can recompute their permissions.\n affected_user_ids = [uid for uid, _ in dal.get_group_members(group.id)]\n\n mapping = dal.get_group_mapping_by_group_id(group.id)\n if mapping:\n dal.delete_group_mapping(mapping.id)\n\n dal.delete_group_with_members(group)\n\n # Recompute permissions for users who lost this group membership.\n recompute_user_permissions__no_commit(affected_user_ids, db_session)\n\n dal.commit()\n\n return Response(status_code=204)\n\n\ndef _is_valid_uuid(value: str) -> bool:\n \"\"\"Check if a string is a valid UUID.\"\"\"\n try:\n UUID(value)\n return True\n except ValueError:\n return False\n" }, { "path": "backend/ee/onyx/server/scim/auth.py", "content": "\"\"\"SCIM bearer token authentication.\n\nSCIM endpoints are authenticated via bearer tokens that admins create in the\nOnyx UI. This module provides:\n\n - ``verify_scim_token``: FastAPI dependency that extracts, hashes, and\n validates the token from the Authorization header.\n - ``generate_scim_token``: Creates a new cryptographically random token\n and returns the raw value, its SHA-256 hash, and a display suffix.\n\nToken format: ``onyx_scim_`` where ```` is 48 bytes of\nURL-safe base64 from ``secrets.token_urlsafe``.\n\nThe hash is stored in the ``scim_token`` table; the raw value is shown to\nthe admin exactly once at creation time.\n\"\"\"\n\nimport hashlib\nimport secrets\n\nfrom fastapi import Depends\nfrom fastapi import Request\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.scim import ScimDAL\nfrom onyx.auth.utils import get_hashed_bearer_token_from_request\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import ScimToken\n\n\nclass ScimAuthError(Exception):\n \"\"\"Raised when SCIM bearer token authentication fails.\n\n Unlike HTTPException, this carries the status and detail so the SCIM\n exception handler can wrap them in an RFC 7644 §3.12 error envelope\n with ``schemas`` and ``status`` fields.\n \"\"\"\n\n def __init__(self, status_code: int, detail: str) -> None:\n self.status_code = status_code\n self.detail = detail\n super().__init__(detail)\n\n\nSCIM_TOKEN_PREFIX = \"onyx_scim_\"\nSCIM_TOKEN_LENGTH = 48\n\n\ndef _hash_scim_token(token: str) -> str:\n \"\"\"SHA-256 hash a SCIM token. No salt needed — tokens are random.\"\"\"\n return hashlib.sha256(token.encode(\"utf-8\")).hexdigest()\n\n\ndef generate_scim_token() -> tuple[str, str, str]:\n \"\"\"Generate a new SCIM bearer token.\n\n Returns:\n A tuple of ``(raw_token, hashed_token, token_display)`` where\n ``token_display`` is a masked version showing only the last 4 chars.\n \"\"\"\n raw_token = SCIM_TOKEN_PREFIX + secrets.token_urlsafe(SCIM_TOKEN_LENGTH)\n hashed_token = _hash_scim_token(raw_token)\n token_display = SCIM_TOKEN_PREFIX + \"****\" + raw_token[-4:]\n return raw_token, hashed_token, token_display\n\n\ndef _get_hashed_scim_token_from_request(request: Request) -> str | None:\n \"\"\"Extract and hash a SCIM token from the request Authorization header.\"\"\"\n return get_hashed_bearer_token_from_request(\n request,\n valid_prefixes=[SCIM_TOKEN_PREFIX],\n hash_fn=_hash_scim_token,\n )\n\n\ndef _get_scim_dal(db_session: Session = Depends(get_session)) -> ScimDAL:\n return ScimDAL(db_session)\n\n\ndef verify_scim_token(\n request: Request,\n dal: ScimDAL = Depends(_get_scim_dal),\n) -> ScimToken:\n \"\"\"FastAPI dependency that authenticates SCIM requests.\n\n Extracts the bearer token from the Authorization header, hashes it,\n looks it up in the database, and verifies it is active.\n\n Note:\n This dependency does NOT update ``last_used_at`` — the endpoint\n should do that via ``ScimDAL.update_token_last_used()`` so the\n timestamp write is part of the endpoint's transaction.\n\n Raises:\n HTTPException(401): If the token is missing, invalid, or inactive.\n \"\"\"\n hashed = _get_hashed_scim_token_from_request(request)\n if not hashed:\n raise ScimAuthError(401, \"Missing or invalid SCIM bearer token\")\n\n token = dal.get_token_by_hash(hashed)\n\n if not token:\n raise ScimAuthError(401, \"Invalid SCIM bearer token\")\n\n if not token.is_active:\n raise ScimAuthError(401, \"SCIM token has been revoked\")\n\n return token\n" }, { "path": "backend/ee/onyx/server/scim/filtering.py", "content": "\"\"\"SCIM filter expression parser (RFC 7644 §3.4.2.2).\n\nIdentity providers (Okta, Azure AD, OneLogin, etc.) use filters to look up\nresources before deciding whether to create or update them. For example, when\nan admin assigns a user to the Onyx app, the IdP first checks whether that\nuser already exists::\n\n GET /scim/v2/Users?filter=userName eq \"john@example.com\"\n\nIf zero results come back the IdP creates the user (``POST``); if a match is\nfound it links to the existing record and uses ``PUT``/``PATCH`` going forward.\nThe same pattern applies to groups (``displayName eq \"Engineering\"``).\n\nThis module parses the subset of the SCIM filter grammar that identity\nproviders actually send in practice:\n\n attribute SP operator SP value\n\nSupported operators: ``eq``, ``co`` (contains), ``sw`` (starts with).\nCompound filters (``and`` / ``or``) are not supported; if an IdP sends one\nthe parser returns ``None`` and the caller falls back to an unfiltered list.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport re\nfrom dataclasses import dataclass\nfrom enum import Enum\n\n\nclass ScimFilterOperator(str, Enum):\n \"\"\"Supported SCIM filter operators.\"\"\"\n\n EQUAL = \"eq\"\n CONTAINS = \"co\"\n STARTS_WITH = \"sw\"\n\n\n@dataclass(frozen=True, slots=True)\nclass ScimFilter:\n \"\"\"Parsed SCIM filter expression.\"\"\"\n\n attribute: str\n operator: ScimFilterOperator\n value: str\n\n\n# Matches: attribute operator \"value\" (with or without quotes around value)\n# Groups: (attribute) (operator) (\"quoted value\" | unquoted_value)\n_FILTER_RE = re.compile(\n r\"^(\\S+)\\s+(eq|co|sw)\\s+\" # attribute + operator\n r'(?:\"([^\"]*)\"' # quoted value\n r\"|'([^']*)')\" # or single-quoted value\n r\"$\",\n re.IGNORECASE,\n)\n\n\ndef parse_scim_filter(filter_string: str | None) -> ScimFilter | None:\n \"\"\"Parse a simple SCIM filter expression.\n\n Args:\n filter_string: Raw filter query parameter value, e.g.\n ``'userName eq \"john@example.com\"'``\n\n Returns:\n A ``ScimFilter`` if the expression is valid and uses a supported\n operator, or ``None`` if the input is empty / missing.\n\n Raises:\n ValueError: If the filter string is present but malformed or uses\n an unsupported operator.\n \"\"\"\n if not filter_string or not filter_string.strip():\n return None\n\n match = _FILTER_RE.match(filter_string.strip())\n if not match:\n raise ValueError(f\"Unsupported or malformed SCIM filter: {filter_string}\")\n\n return _build_filter(match, filter_string)\n\n\ndef _build_filter(match: re.Match[str], raw: str) -> ScimFilter:\n \"\"\"Extract fields from a regex match and construct a ScimFilter.\"\"\"\n attribute = match.group(1)\n op_str = match.group(2).lower()\n # Value is in group 3 (double-quoted) or group 4 (single-quoted)\n value = match.group(3) if match.group(3) is not None else match.group(4)\n\n if value is None:\n raise ValueError(f\"Unsupported or malformed SCIM filter: {raw}\")\n\n operator = ScimFilterOperator(op_str)\n\n return ScimFilter(attribute=attribute, operator=operator, value=value)\n" }, { "path": "backend/ee/onyx/server/scim/models.py", "content": "\"\"\"Pydantic schemas for SCIM 2.0 provisioning (RFC 7643 / RFC 7644).\n\nSCIM protocol schemas follow the wire format defined in:\n - Core Schema: https://datatracker.ietf.org/doc/html/rfc7643\n - Protocol: https://datatracker.ietf.org/doc/html/rfc7644\n\nAdmin API schemas are internal to Onyx and used for SCIM token management.\n\"\"\"\n\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom enum import Enum\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import Field\nfrom pydantic import field_validator\n\n\n# ---------------------------------------------------------------------------\n# SCIM Schema URIs (RFC 7643 §8)\n# Every SCIM JSON payload includes a \"schemas\" array identifying its type.\n# IdPs like Okta/Azure AD use these URIs to determine how to parse responses.\n# ---------------------------------------------------------------------------\n\nSCIM_USER_SCHEMA = \"urn:ietf:params:scim:schemas:core:2.0:User\"\nSCIM_GROUP_SCHEMA = \"urn:ietf:params:scim:schemas:core:2.0:Group\"\nSCIM_LIST_RESPONSE_SCHEMA = \"urn:ietf:params:scim:api:messages:2.0:ListResponse\"\nSCIM_PATCH_OP_SCHEMA = \"urn:ietf:params:scim:api:messages:2.0:PatchOp\"\nSCIM_ERROR_SCHEMA = \"urn:ietf:params:scim:api:messages:2.0:Error\"\nSCIM_SERVICE_PROVIDER_CONFIG_SCHEMA = (\n \"urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig\"\n)\nSCIM_RESOURCE_TYPE_SCHEMA = \"urn:ietf:params:scim:schemas:core:2.0:ResourceType\"\nSCIM_SCHEMA_SCHEMA = \"urn:ietf:params:scim:schemas:core:2.0:Schema\"\nSCIM_ENTERPRISE_USER_SCHEMA = (\n \"urn:ietf:params:scim:schemas:extension:enterprise:2.0:User\"\n)\n\n\n# ---------------------------------------------------------------------------\n# SCIM Protocol Schemas\n# ---------------------------------------------------------------------------\n\n\nclass ScimName(BaseModel):\n \"\"\"User name components (RFC 7643 §4.1.1).\"\"\"\n\n givenName: str | None = None\n familyName: str | None = None\n formatted: str | None = None\n\n\nclass ScimEmail(BaseModel):\n \"\"\"Email sub-attribute (RFC 7643 §4.1.2).\"\"\"\n\n value: str\n type: str | None = None\n primary: bool = False\n\n\nclass ScimMeta(BaseModel):\n \"\"\"Resource metadata (RFC 7643 §3.1).\"\"\"\n\n resourceType: str | None = None\n created: datetime | None = None\n lastModified: datetime | None = None\n location: str | None = None\n\n\nclass ScimUserGroupRef(BaseModel):\n \"\"\"Group reference within a User resource (RFC 7643 §4.1.2, read-only).\"\"\"\n\n value: str\n display: str | None = None\n\n\nclass ScimManagerRef(BaseModel):\n \"\"\"Manager sub-attribute for the enterprise extension (RFC 7643 §4.3).\"\"\"\n\n value: str | None = None\n\n\nclass ScimEnterpriseExtension(BaseModel):\n \"\"\"Enterprise User extension attributes (RFC 7643 §4.3).\"\"\"\n\n department: str | None = None\n manager: ScimManagerRef | None = None\n\n\n@dataclass\nclass ScimMappingFields:\n \"\"\"Stored SCIM mapping fields that need to round-trip through the IdP.\n\n Entra ID sends structured name components, email metadata, and enterprise\n extension attributes that must be returned verbatim in subsequent GET\n responses. These fields are persisted on ScimUserMapping and threaded\n through the DAL, provider, and endpoint layers.\n \"\"\"\n\n department: str | None = None\n manager: str | None = None\n given_name: str | None = None\n family_name: str | None = None\n scim_emails_json: str | None = None\n\n\nclass ScimUserResource(BaseModel):\n \"\"\"SCIM User resource representation (RFC 7643 §4.1).\n\n This is the JSON shape that IdPs send when creating/updating a user via\n SCIM, and the shape we return in GET responses. Field names use camelCase\n to match the SCIM wire format (not Python convention).\n \"\"\"\n\n model_config = ConfigDict(populate_by_name=True)\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_USER_SCHEMA])\n id: str | None = None # Onyx's internal user ID, set on responses\n externalId: str | None = None # IdP's identifier for this user\n userName: str # Typically the user's email address\n name: ScimName | None = None\n displayName: str | None = None\n emails: list[ScimEmail] = Field(default_factory=list)\n active: bool = True\n groups: list[ScimUserGroupRef] = Field(default_factory=list)\n meta: ScimMeta | None = None\n enterprise_extension: ScimEnterpriseExtension | None = Field(\n default=None,\n alias=\"urn:ietf:params:scim:schemas:extension:enterprise:2.0:User\",\n )\n\n\nclass ScimGroupMember(BaseModel):\n \"\"\"Group member reference (RFC 7643 §4.2).\n\n Represents a user within a SCIM group. The IdP sends these when adding\n or removing users from groups. ``value`` is the Onyx user ID.\n \"\"\"\n\n value: str # User ID of the group member\n display: str | None = None\n\n\nclass ScimGroupResource(BaseModel):\n \"\"\"SCIM Group resource representation (RFC 7643 §4.2).\"\"\"\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_GROUP_SCHEMA])\n id: str | None = None\n externalId: str | None = None\n displayName: str\n members: list[ScimGroupMember] = Field(default_factory=list)\n meta: ScimMeta | None = None\n\n\nclass ScimListResponse(BaseModel):\n \"\"\"Paginated list response (RFC 7644 §3.4.2).\"\"\"\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_LIST_RESPONSE_SCHEMA])\n totalResults: int\n startIndex: int = 1\n itemsPerPage: int = 100\n Resources: list[ScimUserResource | ScimGroupResource] = Field(default_factory=list)\n\n\nclass ScimPatchOperationType(str, Enum):\n \"\"\"Supported PATCH operations (RFC 7644 §3.5.2).\"\"\"\n\n ADD = \"add\"\n REPLACE = \"replace\"\n REMOVE = \"remove\"\n\n\nclass ScimPatchResourceValue(BaseModel):\n \"\"\"Partial resource dict for path-less PATCH replace operations.\n\n When an IdP sends a PATCH without a ``path``, the ``value`` is a dict\n of resource attributes to set. IdPs may include read-only fields\n (``id``, ``schemas``, ``meta``) alongside actual changes — these are\n stripped by the provider's ``ignored_patch_paths`` before processing.\n\n ``extra=\"allow\"`` lets unknown attributes pass through so the patch\n handler can decide what to do with them (ignore or reject).\n \"\"\"\n\n model_config = ConfigDict(extra=\"allow\")\n\n active: bool | None = None\n userName: str | None = None\n displayName: str | None = None\n externalId: str | None = None\n name: ScimName | None = None\n members: list[ScimGroupMember] | None = None\n id: str | None = None\n schemas: list[str] | None = None\n meta: ScimMeta | None = None\n\n\nScimPatchValue = str | bool | list[ScimGroupMember] | ScimPatchResourceValue | None\n\n\nclass ScimPatchOperation(BaseModel):\n \"\"\"Single PATCH operation (RFC 7644 §3.5.2).\"\"\"\n\n op: ScimPatchOperationType\n path: str | None = None\n value: ScimPatchValue = None\n\n @field_validator(\"op\", mode=\"before\")\n @classmethod\n def normalize_operation(cls, v: object) -> object:\n \"\"\"Normalize op to lowercase for case-insensitive matching.\n\n Some IdPs (e.g. Entra ID) send capitalized ops like ``\"Replace\"``\n instead of ``\"replace\"``. This is safe for all providers since the\n enum values are lowercase. If a future provider requires other\n pre-processing quirks, move patch deserialization into the provider\n subclass instead of adding more special cases here.\n \"\"\"\n return v.lower() if isinstance(v, str) else v\n\n\nclass ScimPatchRequest(BaseModel):\n \"\"\"PATCH request body (RFC 7644 §3.5.2).\n\n IdPs use PATCH to make incremental changes — e.g. deactivating a user\n (replace active=false) or adding/removing group members — instead of\n replacing the entire resource with PUT.\n \"\"\"\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_PATCH_OP_SCHEMA])\n Operations: list[ScimPatchOperation]\n\n\nclass ScimError(BaseModel):\n \"\"\"SCIM error response (RFC 7644 §3.12).\"\"\"\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_ERROR_SCHEMA])\n status: str\n detail: str | None = None\n scimType: str | None = None\n\n\n# ---------------------------------------------------------------------------\n# Service Provider Configuration (RFC 7643 §5)\n# ---------------------------------------------------------------------------\n\n\nclass ScimSupported(BaseModel):\n \"\"\"Generic supported/not-supported flag used in ServiceProviderConfig.\"\"\"\n\n supported: bool\n\n\nclass ScimFilterConfig(BaseModel):\n \"\"\"Filter configuration within ServiceProviderConfig (RFC 7643 §5).\"\"\"\n\n supported: bool\n maxResults: int = 100\n\n\nclass ScimServiceProviderConfig(BaseModel):\n \"\"\"SCIM ServiceProviderConfig resource (RFC 7643 §5).\n\n Served at GET /scim/v2/ServiceProviderConfig. IdPs fetch this during\n initial setup to discover which SCIM features our server supports\n (e.g. PATCH yes, bulk no, filtering yes).\n \"\"\"\n\n schemas: list[str] = Field(\n default_factory=lambda: [SCIM_SERVICE_PROVIDER_CONFIG_SCHEMA]\n )\n patch: ScimSupported = ScimSupported(supported=True)\n bulk: ScimSupported = ScimSupported(supported=False)\n filter: ScimFilterConfig = ScimFilterConfig(supported=True)\n changePassword: ScimSupported = ScimSupported(supported=False)\n sort: ScimSupported = ScimSupported(supported=False)\n etag: ScimSupported = ScimSupported(supported=False)\n authenticationSchemes: list[dict[str, str]] = Field(\n default_factory=lambda: [\n {\n \"type\": \"oauthbearertoken\",\n \"name\": \"OAuth Bearer Token\",\n \"description\": \"Authentication scheme using a SCIM bearer token\",\n }\n ]\n )\n\n\nclass ScimSchemaAttribute(BaseModel):\n \"\"\"Attribute definition within a SCIM Schema (RFC 7643 §7).\"\"\"\n\n name: str\n type: str\n multiValued: bool = False\n required: bool = False\n description: str = \"\"\n caseExact: bool = False\n mutability: str = \"readWrite\"\n returned: str = \"default\"\n uniqueness: str = \"none\"\n subAttributes: list[\"ScimSchemaAttribute\"] = Field(default_factory=list)\n\n\nclass ScimSchemaDefinition(BaseModel):\n \"\"\"SCIM Schema definition (RFC 7643 §7).\n\n Served at GET /scim/v2/Schemas. Describes the attributes available\n on each resource type so IdPs know which fields they can provision.\n \"\"\"\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_SCHEMA_SCHEMA])\n id: str\n name: str\n description: str\n attributes: list[ScimSchemaAttribute] = Field(default_factory=list)\n\n\nclass ScimSchemaExtension(BaseModel):\n \"\"\"Schema extension reference within ResourceType (RFC 7643 §6).\"\"\"\n\n model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)\n\n schema_: str = Field(alias=\"schema\")\n required: bool\n\n\nclass ScimResourceType(BaseModel):\n \"\"\"SCIM ResourceType resource (RFC 7643 §6).\n\n Served at GET /scim/v2/ResourceTypes. Tells the IdP which resource\n types are available (Users, Groups) and their respective endpoints.\n \"\"\"\n\n model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)\n\n schemas: list[str] = Field(default_factory=lambda: [SCIM_RESOURCE_TYPE_SCHEMA])\n id: str\n name: str\n endpoint: str\n description: str | None = None\n schema_: str = Field(alias=\"schema\")\n schemaExtensions: list[ScimSchemaExtension] = Field(default_factory=list)\n\n\n# ---------------------------------------------------------------------------\n# Admin API Schemas (Onyx-internal, for SCIM token management)\n# These are NOT part of the SCIM protocol. They power the Onyx admin UI\n# where admins create/revoke the bearer tokens that IdPs use to authenticate.\n# ---------------------------------------------------------------------------\n\n\nclass ScimTokenCreate(BaseModel):\n \"\"\"Request to create a new SCIM bearer token.\"\"\"\n\n name: str\n\n\nclass ScimTokenResponse(BaseModel):\n \"\"\"SCIM token metadata returned in list/get responses.\"\"\"\n\n id: int\n name: str\n token_display: str\n is_active: bool\n created_at: datetime\n last_used_at: datetime | None = None\n idp_domain: str | None = None\n\n\nclass ScimTokenCreatedResponse(ScimTokenResponse):\n \"\"\"Response returned when a new SCIM token is created.\n\n Includes the raw token value which is only available at creation time.\n \"\"\"\n\n raw_token: str\n" }, { "path": "backend/ee/onyx/server/scim/patch.py", "content": "\"\"\"SCIM PATCH operation handler (RFC 7644 §3.5.2).\n\nIdentity providers use PATCH to make incremental changes to SCIM resources\ninstead of replacing the entire resource with PUT. Common operations include:\n\n - Deactivating a user: ``replace`` ``active`` with ``false``\n - Adding group members: ``add`` to ``members``\n - Removing group members: ``remove`` from ``members[value eq \"...\"]``\n\nThis module applies PATCH operations to Pydantic SCIM resource objects and\nreturns the modified result. It does NOT touch the database — the caller is\nresponsible for persisting changes.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport logging\nimport re\nfrom dataclasses import dataclass\nfrom dataclasses import field\nfrom typing import Any\n\nfrom ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA\nfrom ee.onyx.server.scim.models import ScimGroupMember\nfrom ee.onyx.server.scim.models import ScimGroupResource\nfrom ee.onyx.server.scim.models import ScimPatchOperation\nfrom ee.onyx.server.scim.models import ScimPatchOperationType\nfrom ee.onyx.server.scim.models import ScimPatchResourceValue\nfrom ee.onyx.server.scim.models import ScimPatchValue\nfrom ee.onyx.server.scim.models import ScimUserResource\n\nlogger = logging.getLogger(__name__)\n\n# Lowercased enterprise extension URN for case-insensitive matching\n_ENTERPRISE_URN_LOWER = SCIM_ENTERPRISE_USER_SCHEMA.lower()\n\n# Pattern for email filter paths, e.g.:\n# emails[primary eq true].value (Okta)\n# emails[type eq \"work\"].value (Azure AD / Entra ID)\n_EMAIL_FILTER_RE = re.compile(\n r\"^emails\\[.+\\]\\.value$\",\n re.IGNORECASE,\n)\n\n# Pattern for member removal path: members[value eq \"user-id\"]\n_MEMBER_FILTER_RE = re.compile(\n r'^members\\[value\\s+eq\\s+\"([^\"]+)\"\\]$',\n re.IGNORECASE,\n)\n\n# ---------------------------------------------------------------------------\n# Dispatch tables for user PATCH paths\n#\n# Maps lowercased SCIM path → (camelCase key, target dict name).\n# \"data\" writes to the top-level resource dict, \"name\" writes to the\n# name sub-object dict. This replaces the elif chains for simple fields.\n# ---------------------------------------------------------------------------\n\n_USER_REPLACE_PATHS: dict[str, tuple[str, str]] = {\n \"active\": (\"active\", \"data\"),\n \"username\": (\"userName\", \"data\"),\n \"externalid\": (\"externalId\", \"data\"),\n \"name.givenname\": (\"givenName\", \"name\"),\n \"name.familyname\": (\"familyName\", \"name\"),\n \"name.formatted\": (\"formatted\", \"name\"),\n}\n\n_USER_REMOVE_PATHS: dict[str, tuple[str, str]] = {\n \"externalid\": (\"externalId\", \"data\"),\n \"name.givenname\": (\"givenName\", \"name\"),\n \"name.familyname\": (\"familyName\", \"name\"),\n \"name.formatted\": (\"formatted\", \"name\"),\n \"displayname\": (\"displayName\", \"data\"),\n}\n\n_GROUP_REPLACE_PATHS: dict[str, tuple[str, str]] = {\n \"displayname\": (\"displayName\", \"data\"),\n \"externalid\": (\"externalId\", \"data\"),\n}\n\n\nclass ScimPatchError(Exception):\n \"\"\"Raised when a PATCH operation cannot be applied.\"\"\"\n\n def __init__(self, detail: str, status: int = 400) -> None:\n self.detail = detail\n self.status = status\n super().__init__(detail)\n\n\n@dataclass\nclass _UserPatchCtx:\n \"\"\"Bundles the mutable state for user PATCH operations.\"\"\"\n\n data: dict[str, Any]\n name_data: dict[str, Any]\n ent_data: dict[str, str | None] = field(default_factory=dict)\n\n\n# ---------------------------------------------------------------------------\n# User PATCH\n# ---------------------------------------------------------------------------\n\n\ndef apply_user_patch(\n operations: list[ScimPatchOperation],\n current: ScimUserResource,\n ignored_paths: frozenset[str] = frozenset(),\n) -> tuple[ScimUserResource, dict[str, str | None]]:\n \"\"\"Apply SCIM PATCH operations to a user resource.\n\n Args:\n operations: The PATCH operations to apply.\n current: The current user resource state.\n ignored_paths: SCIM attribute paths to silently skip (from provider).\n\n Returns:\n A tuple of (modified user resource, enterprise extension data dict).\n The enterprise dict has keys ``\"department\"`` and ``\"manager\"``\n with values set only when a PATCH operation touched them.\n\n Raises:\n ScimPatchError: If an operation targets an unsupported path.\n \"\"\"\n data = current.model_dump()\n ctx = _UserPatchCtx(data=data, name_data=data.get(\"name\") or {})\n\n for op in operations:\n if op.op in (ScimPatchOperationType.REPLACE, ScimPatchOperationType.ADD):\n _apply_user_replace(op, ctx, ignored_paths)\n elif op.op == ScimPatchOperationType.REMOVE:\n _apply_user_remove(op, ctx, ignored_paths)\n else:\n raise ScimPatchError(\n f\"Unsupported operation '{op.op.value}' on User resource\"\n )\n\n ctx.data[\"name\"] = ctx.name_data\n return ScimUserResource.model_validate(ctx.data), ctx.ent_data\n\n\ndef _apply_user_replace(\n op: ScimPatchOperation,\n ctx: _UserPatchCtx,\n ignored_paths: frozenset[str],\n) -> None:\n \"\"\"Apply a replace/add operation to user data.\"\"\"\n path = (op.path or \"\").lower()\n\n if not path:\n # No path — value is a resource dict of top-level attributes to set.\n if isinstance(op.value, ScimPatchResourceValue):\n for key, val in op.value.model_dump(exclude_unset=True).items():\n _set_user_field(key.lower(), val, ctx, ignored_paths, strict=False)\n else:\n raise ScimPatchError(\"Replace without path requires a dict value\")\n return\n\n _set_user_field(path, op.value, ctx, ignored_paths)\n\n\ndef _apply_user_remove(\n op: ScimPatchOperation,\n ctx: _UserPatchCtx,\n ignored_paths: frozenset[str],\n) -> None:\n \"\"\"Apply a remove operation to user data — clears the target field.\"\"\"\n path = (op.path or \"\").lower()\n if not path:\n raise ScimPatchError(\"Remove operation requires a path\")\n\n if path in ignored_paths:\n return\n\n entry = _USER_REMOVE_PATHS.get(path)\n if entry:\n key, target = entry\n target_dict = ctx.data if target == \"data\" else ctx.name_data\n target_dict[key] = None\n return\n\n raise ScimPatchError(f\"Unsupported remove path '{path}' for User PATCH\")\n\n\ndef _set_user_field(\n path: str,\n value: ScimPatchValue,\n ctx: _UserPatchCtx,\n ignored_paths: frozenset[str],\n *,\n strict: bool = True,\n) -> None:\n \"\"\"Set a single field on user data by SCIM path.\n\n Args:\n strict: When ``False`` (path-less replace), unknown attributes are\n silently skipped. When ``True`` (explicit path), they raise.\n \"\"\"\n if path in ignored_paths:\n return\n\n # Simple field writes handled by the dispatch table\n entry = _USER_REPLACE_PATHS.get(path)\n if entry:\n key, target = entry\n target_dict = ctx.data if target == \"data\" else ctx.name_data\n target_dict[key] = value\n return\n\n # displayName sets both the top-level field and the name.formatted sub-field\n if path == \"displayname\":\n ctx.data[\"displayName\"] = value\n ctx.name_data[\"formatted\"] = value\n elif path == \"name\":\n if isinstance(value, dict):\n for k, v in value.items():\n ctx.name_data[k] = v\n elif path == \"emails\":\n if isinstance(value, list):\n ctx.data[\"emails\"] = value\n elif _EMAIL_FILTER_RE.match(path):\n _update_primary_email(ctx.data, value)\n elif path.startswith(_ENTERPRISE_URN_LOWER):\n _set_enterprise_field(path, value, ctx.ent_data)\n elif not strict:\n return\n else:\n raise ScimPatchError(f\"Unsupported path '{path}' for User PATCH\")\n\n\ndef _update_primary_email(data: dict[str, Any], value: ScimPatchValue) -> None:\n \"\"\"Update the primary email entry via an email filter path.\"\"\"\n emails: list[dict] = data.get(\"emails\") or []\n for email_entry in emails:\n if email_entry.get(\"primary\"):\n email_entry[\"value\"] = value\n break\n else:\n emails.append({\"value\": value, \"type\": \"work\", \"primary\": True})\n data[\"emails\"] = emails\n\n\ndef _to_dict(value: ScimPatchValue) -> dict | None:\n \"\"\"Coerce a SCIM patch value to a plain dict if possible.\n\n Pydantic may parse raw dicts as ``ScimPatchResourceValue`` (which uses\n ``extra=\"allow\"``), so we also dump those back to a dict.\n \"\"\"\n if isinstance(value, dict):\n return value\n if isinstance(value, ScimPatchResourceValue):\n return value.model_dump(exclude_unset=True)\n return None\n\n\ndef _set_enterprise_field(\n path: str,\n value: ScimPatchValue,\n ent_data: dict[str, str | None],\n) -> None:\n \"\"\"Handle enterprise extension URN paths or value dicts.\"\"\"\n # Full URN as key with dict value (path-less PATCH)\n # e.g. key=\"urn:...:user\", value={\"department\": \"Eng\", \"manager\": {...}}\n if path == _ENTERPRISE_URN_LOWER:\n d = _to_dict(value)\n if d is not None:\n if \"department\" in d:\n ent_data[\"department\"] = d[\"department\"]\n if \"manager\" in d:\n mgr = d[\"manager\"]\n if isinstance(mgr, dict):\n ent_data[\"manager\"] = mgr.get(\"value\")\n return\n\n # Dotted URN path, e.g. \"urn:...:user:department\"\n suffix = path[len(_ENTERPRISE_URN_LOWER) :].lstrip(\":\").lower()\n if suffix == \"department\":\n ent_data[\"department\"] = str(value) if value is not None else None\n elif suffix == \"manager\":\n d = _to_dict(value)\n if d is not None:\n ent_data[\"manager\"] = d.get(\"value\")\n elif isinstance(value, str):\n ent_data[\"manager\"] = value\n else:\n # Unknown enterprise attributes are silently ignored rather than\n # rejected — IdPs may send attributes we don't model yet.\n logger.warning(\"Ignoring unknown enterprise extension attribute '%s'\", suffix)\n\n\n# ---------------------------------------------------------------------------\n# Group PATCH\n# ---------------------------------------------------------------------------\n\n\ndef apply_group_patch(\n operations: list[ScimPatchOperation],\n current: ScimGroupResource,\n ignored_paths: frozenset[str] = frozenset(),\n) -> tuple[ScimGroupResource, list[str], list[str]]:\n \"\"\"Apply SCIM PATCH operations to a group resource.\n\n Args:\n operations: The PATCH operations to apply.\n current: The current group resource state.\n ignored_paths: SCIM attribute paths to silently skip (from provider).\n\n Returns:\n A tuple of (modified group, added member IDs, removed member IDs).\n The caller uses the member ID lists to update the database.\n\n Raises:\n ScimPatchError: If an operation targets an unsupported path.\n \"\"\"\n data = current.model_dump()\n current_members: list[dict] = list(data.get(\"members\") or [])\n added_ids: list[str] = []\n removed_ids: list[str] = []\n\n for op in operations:\n if op.op == ScimPatchOperationType.REPLACE:\n _apply_group_replace(\n op, data, current_members, added_ids, removed_ids, ignored_paths\n )\n elif op.op == ScimPatchOperationType.ADD:\n _apply_group_add(op, current_members, added_ids)\n elif op.op == ScimPatchOperationType.REMOVE:\n _apply_group_remove(op, current_members, removed_ids)\n else:\n raise ScimPatchError(\n f\"Unsupported operation '{op.op.value}' on Group resource\"\n )\n\n data[\"members\"] = current_members\n group = ScimGroupResource.model_validate(data)\n return group, added_ids, removed_ids\n\n\ndef _apply_group_replace(\n op: ScimPatchOperation,\n data: dict,\n current_members: list[dict],\n added_ids: list[str],\n removed_ids: list[str],\n ignored_paths: frozenset[str],\n) -> None:\n \"\"\"Apply a replace operation to group data.\"\"\"\n path = (op.path or \"\").lower()\n\n if not path:\n if isinstance(op.value, ScimPatchResourceValue):\n dumped = op.value.model_dump(exclude_unset=True)\n for key, val in dumped.items():\n if key.lower() == \"members\":\n _replace_members(val, current_members, added_ids, removed_ids)\n else:\n _set_group_field(key.lower(), val, data, ignored_paths)\n else:\n raise ScimPatchError(\"Replace without path requires a dict value\")\n return\n\n if path == \"members\":\n _replace_members(\n _members_to_dicts(op.value), current_members, added_ids, removed_ids\n )\n return\n\n _set_group_field(path, op.value, data, ignored_paths)\n\n\ndef _members_to_dicts(\n value: str | bool | list[ScimGroupMember] | ScimPatchResourceValue | None,\n) -> list[dict]:\n \"\"\"Convert a member list value to a list of dicts for internal processing.\"\"\"\n if not isinstance(value, list):\n raise ScimPatchError(\"Replace members requires a list value\")\n return [m.model_dump(exclude_none=True) for m in value]\n\n\ndef _replace_members(\n value: list[dict],\n current_members: list[dict],\n added_ids: list[str],\n removed_ids: list[str],\n) -> None:\n \"\"\"Replace the entire group member list.\"\"\"\n old_ids = {m[\"value\"] for m in current_members}\n new_ids = {m.get(\"value\", \"\") for m in value}\n\n removed_ids.extend(old_ids - new_ids)\n added_ids.extend(new_ids - old_ids)\n\n current_members[:] = value\n\n\ndef _set_group_field(\n path: str,\n value: ScimPatchValue,\n data: dict,\n ignored_paths: frozenset[str],\n) -> None:\n \"\"\"Set a single field on group data by SCIM path.\"\"\"\n if path in ignored_paths:\n return\n\n entry = _GROUP_REPLACE_PATHS.get(path)\n if entry:\n key, _ = entry\n data[key] = value\n return\n\n raise ScimPatchError(f\"Unsupported path '{path}' for Group PATCH\")\n\n\ndef _apply_group_add(\n op: ScimPatchOperation,\n members: list[dict],\n added_ids: list[str],\n) -> None:\n \"\"\"Add members to a group.\"\"\"\n path = (op.path or \"\").lower()\n\n if path and path != \"members\":\n raise ScimPatchError(f\"Unsupported add path '{op.path}' for Group\")\n\n if not isinstance(op.value, list):\n raise ScimPatchError(\"Add members requires a list value\")\n\n member_dicts = [m.model_dump(exclude_none=True) for m in op.value]\n\n existing_ids = {m[\"value\"] for m in members}\n for member_data in member_dicts:\n member_id = member_data.get(\"value\", \"\")\n if member_id and member_id not in existing_ids:\n members.append(member_data)\n added_ids.append(member_id)\n existing_ids.add(member_id)\n\n\ndef _apply_group_remove(\n op: ScimPatchOperation,\n members: list[dict],\n removed_ids: list[str],\n) -> None:\n \"\"\"Remove members from a group.\"\"\"\n if not op.path:\n raise ScimPatchError(\"Remove operation requires a path\")\n\n match = _MEMBER_FILTER_RE.match(op.path)\n if not match:\n raise ScimPatchError(\n f\"Unsupported remove path '{op.path}'. Expected: members[value eq \\\"user-id\\\"]\"\n )\n\n target_id = match.group(1)\n original_len = len(members)\n members[:] = [m for m in members if m.get(\"value\") != target_id]\n\n if len(members) < original_len:\n removed_ids.append(target_id)\n" }, { "path": "backend/ee/onyx/server/scim/providers/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/scim/providers/base.py", "content": "\"\"\"Base SCIM provider abstraction.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport logging\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom uuid import UUID\n\nfrom pydantic import ValidationError\n\nfrom ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA\nfrom ee.onyx.server.scim.models import SCIM_USER_SCHEMA\nfrom ee.onyx.server.scim.models import ScimEmail\nfrom ee.onyx.server.scim.models import ScimEnterpriseExtension\nfrom ee.onyx.server.scim.models import ScimGroupMember\nfrom ee.onyx.server.scim.models import ScimGroupResource\nfrom ee.onyx.server.scim.models import ScimManagerRef\nfrom ee.onyx.server.scim.models import ScimMappingFields\nfrom ee.onyx.server.scim.models import ScimMeta\nfrom ee.onyx.server.scim.models import ScimName\nfrom ee.onyx.server.scim.models import ScimUserGroupRef\nfrom ee.onyx.server.scim.models import ScimUserResource\nfrom onyx.db.models import User\nfrom onyx.db.models import UserGroup\n\n\nlogger = logging.getLogger(__name__)\n\nCOMMON_IGNORED_PATCH_PATHS: frozenset[str] = frozenset(\n {\n \"id\",\n \"schemas\",\n \"meta\",\n }\n)\n\n\nclass ScimProvider(ABC):\n \"\"\"Base class for provider-specific SCIM behavior.\n\n Subclass this to handle IdP-specific quirks. The base class provides\n RFC 7643-compliant response builders that populate all standard fields.\n \"\"\"\n\n @property\n @abstractmethod\n def name(self) -> str:\n \"\"\"Short identifier for this provider (e.g. ``\"okta\"``).\"\"\"\n ...\n\n @property\n @abstractmethod\n def ignored_patch_paths(self) -> frozenset[str]:\n \"\"\"SCIM attribute paths to silently skip in PATCH value-object dicts.\n\n IdPs may include read-only or meta fields alongside actual changes\n (e.g. Okta sends ``{\"id\": \"...\", \"active\": false}``). Paths listed\n here are silently dropped instead of raising an error.\n \"\"\"\n ...\n\n @property\n def user_schemas(self) -> list[str]:\n \"\"\"Schema URIs to include in User resource responses.\n\n Override in subclasses to advertise additional schemas (e.g. the\n enterprise extension for Entra ID).\n \"\"\"\n return [SCIM_USER_SCHEMA]\n\n def build_user_resource(\n self,\n user: User,\n external_id: str | None = None,\n groups: list[tuple[int, str]] | None = None,\n scim_username: str | None = None,\n fields: ScimMappingFields | None = None,\n ) -> ScimUserResource:\n \"\"\"Build a SCIM User response from an Onyx User.\n\n Args:\n user: The Onyx user model.\n external_id: The IdP's external identifier for this user.\n groups: List of ``(group_id, group_name)`` tuples for the\n ``groups`` read-only attribute. Pass ``None`` or ``[]``\n for newly-created users.\n scim_username: The original-case userName from the IdP. Falls\n back to ``user.email`` (lowercase) when not available.\n fields: Stored mapping fields that the IdP expects round-tripped.\n \"\"\"\n f = fields or ScimMappingFields()\n group_refs = [\n ScimUserGroupRef(value=str(gid), display=gname)\n for gid, gname in (groups or [])\n ]\n\n username = scim_username or user.email\n\n # Build enterprise extension when at least one value is present.\n # Dynamically add the enterprise URN to schemas per RFC 7643 §3.0.\n enterprise_ext: ScimEnterpriseExtension | None = None\n schemas = list(self.user_schemas)\n if f.department is not None or f.manager is not None:\n manager_ref = (\n ScimManagerRef(value=f.manager) if f.manager is not None else None\n )\n enterprise_ext = ScimEnterpriseExtension(\n department=f.department,\n manager=manager_ref,\n )\n if SCIM_ENTERPRISE_USER_SCHEMA not in schemas:\n schemas.append(SCIM_ENTERPRISE_USER_SCHEMA)\n\n name = self.build_scim_name(user, f)\n emails = _deserialize_emails(f.scim_emails_json, username)\n\n resource = ScimUserResource(\n schemas=schemas,\n id=str(user.id),\n externalId=external_id,\n userName=username,\n name=name,\n displayName=user.personal_name,\n emails=emails,\n active=user.is_active,\n groups=group_refs,\n meta=ScimMeta(resourceType=\"User\"),\n )\n resource.enterprise_extension = enterprise_ext\n return resource\n\n def build_group_resource(\n self,\n group: UserGroup,\n members: list[tuple[UUID, str | None]],\n external_id: str | None = None,\n ) -> ScimGroupResource:\n \"\"\"Build a SCIM Group response from an Onyx UserGroup.\"\"\"\n scim_members = [\n ScimGroupMember(value=str(uid), display=email) for uid, email in members\n ]\n return ScimGroupResource(\n id=str(group.id),\n externalId=external_id,\n displayName=group.name,\n members=scim_members,\n meta=ScimMeta(resourceType=\"Group\"),\n )\n\n def build_scim_name(\n self,\n user: User,\n fields: ScimMappingFields,\n ) -> ScimName:\n \"\"\"Build SCIM name components for the response.\n\n Round-trips stored ``given_name``/``family_name`` when available (so\n the IdP gets back what it sent). Falls back to splitting\n ``personal_name`` for users provisioned before we stored components.\n Always returns a ScimName — Okta's spec tests expect ``name``\n (with ``givenName``/``familyName``) on every user resource.\n Providers may override for custom behavior.\n \"\"\"\n if fields.given_name is not None or fields.family_name is not None:\n return ScimName(\n givenName=fields.given_name or \"\",\n familyName=fields.family_name or \"\",\n formatted=user.personal_name or \"\",\n )\n if not user.personal_name:\n # Derive a reasonable name from the email so that SCIM spec tests\n # see non-empty givenName / familyName for every user resource.\n local = user.email.split(\"@\")[0] if user.email else \"\"\n return ScimName(givenName=local, familyName=\"\", formatted=local)\n parts = user.personal_name.split(\" \", 1)\n return ScimName(\n givenName=parts[0],\n familyName=parts[1] if len(parts) > 1 else \"\",\n formatted=user.personal_name,\n )\n\n\ndef _deserialize_emails(stored_json: str | None, username: str) -> list[ScimEmail]:\n \"\"\"Deserialize stored email entries or build a default work email.\"\"\"\n if stored_json:\n try:\n entries = json.loads(stored_json)\n if isinstance(entries, list) and entries:\n return [ScimEmail(**e) for e in entries]\n except (json.JSONDecodeError, TypeError, ValidationError):\n logger.warning(\n \"Corrupt scim_emails_json, falling back to default: %s\", stored_json\n )\n return [ScimEmail(value=username, type=\"work\", primary=True)]\n\n\ndef serialize_emails(emails: list[ScimEmail]) -> str | None:\n \"\"\"Serialize SCIM email entries to JSON for storage.\"\"\"\n if not emails:\n return None\n return json.dumps([e.model_dump(exclude_none=True) for e in emails])\n\n\ndef get_default_provider() -> ScimProvider:\n \"\"\"Return the default SCIM provider.\n\n Currently returns ``OktaProvider`` since Okta is the primary supported\n IdP. When provider detection is added (via token metadata or tenant\n config), this can be replaced with dynamic resolution.\n \"\"\"\n from ee.onyx.server.scim.providers.okta import OktaProvider\n\n return OktaProvider()\n" }, { "path": "backend/ee/onyx/server/scim/providers/entra.py", "content": "\"\"\"Entra ID (Azure AD) SCIM provider.\"\"\"\n\nfrom __future__ import annotations\n\nfrom ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA\nfrom ee.onyx.server.scim.models import SCIM_USER_SCHEMA\nfrom ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS\nfrom ee.onyx.server.scim.providers.base import ScimProvider\n\n_ENTRA_IGNORED_PATCH_PATHS = COMMON_IGNORED_PATCH_PATHS\n\n\nclass EntraProvider(ScimProvider):\n \"\"\"Entra ID (Azure AD) SCIM provider.\n\n Entra behavioral notes:\n - Sends capitalized PATCH ops (``\"Add\"``, ``\"Replace\"``, ``\"Remove\"``)\n — handled by ``ScimPatchOperation.normalize_op`` validator.\n - Sends the enterprise extension URN as a key in path-less PATCH value\n dicts — handled by ``_set_enterprise_field`` in ``patch.py`` to\n store department/manager values.\n - Expects the enterprise extension schema in ``schemas`` arrays and\n ``/Schemas`` + ``/ResourceTypes`` discovery endpoints.\n \"\"\"\n\n @property\n def name(self) -> str:\n return \"entra\"\n\n @property\n def ignored_patch_paths(self) -> frozenset[str]:\n return _ENTRA_IGNORED_PATCH_PATHS\n\n @property\n def user_schemas(self) -> list[str]:\n return [SCIM_USER_SCHEMA, SCIM_ENTERPRISE_USER_SCHEMA]\n" }, { "path": "backend/ee/onyx/server/scim/providers/okta.py", "content": "\"\"\"Okta SCIM provider.\"\"\"\n\nfrom __future__ import annotations\n\nfrom ee.onyx.server.scim.providers.base import COMMON_IGNORED_PATCH_PATHS\nfrom ee.onyx.server.scim.providers.base import ScimProvider\n\n\nclass OktaProvider(ScimProvider):\n \"\"\"Okta SCIM provider.\n\n Okta behavioral notes:\n - Uses ``PATCH {\"active\": false}`` for deprovisioning (not DELETE)\n - Sends path-less PATCH with value dicts containing extra fields\n (``id``, ``schemas``)\n - Expects ``displayName`` and ``groups`` in user responses\n - Only uses ``eq`` operator for ``userName`` filter\n \"\"\"\n\n @property\n def name(self) -> str:\n return \"okta\"\n\n @property\n def ignored_patch_paths(self) -> frozenset[str]:\n return COMMON_IGNORED_PATCH_PATHS\n" }, { "path": "backend/ee/onyx/server/scim/schema_definitions.py", "content": "\"\"\"Static SCIM service discovery responses (RFC 7643 §5, §6, §7).\n\nPre-built at import time — these never change at runtime. Separated from\napi.py to keep the endpoint module focused on request handling.\n\"\"\"\n\nfrom ee.onyx.server.scim.models import SCIM_ENTERPRISE_USER_SCHEMA\nfrom ee.onyx.server.scim.models import SCIM_GROUP_SCHEMA\nfrom ee.onyx.server.scim.models import SCIM_USER_SCHEMA\nfrom ee.onyx.server.scim.models import ScimResourceType\nfrom ee.onyx.server.scim.models import ScimSchemaAttribute\nfrom ee.onyx.server.scim.models import ScimSchemaDefinition\nfrom ee.onyx.server.scim.models import ScimServiceProviderConfig\n\nSERVICE_PROVIDER_CONFIG = ScimServiceProviderConfig()\n\nUSER_RESOURCE_TYPE = ScimResourceType.model_validate(\n {\n \"id\": \"User\",\n \"name\": \"User\",\n \"endpoint\": \"/scim/v2/Users\",\n \"description\": \"SCIM User resource\",\n \"schema\": SCIM_USER_SCHEMA,\n \"schemaExtensions\": [\n {\"schema\": SCIM_ENTERPRISE_USER_SCHEMA, \"required\": False}\n ],\n }\n)\n\nGROUP_RESOURCE_TYPE = ScimResourceType.model_validate(\n {\n \"id\": \"Group\",\n \"name\": \"Group\",\n \"endpoint\": \"/scim/v2/Groups\",\n \"description\": \"SCIM Group resource\",\n \"schema\": SCIM_GROUP_SCHEMA,\n }\n)\n\nUSER_SCHEMA_DEF = ScimSchemaDefinition(\n id=SCIM_USER_SCHEMA,\n name=\"User\",\n description=\"SCIM core User schema\",\n attributes=[\n ScimSchemaAttribute(\n name=\"userName\",\n type=\"string\",\n required=True,\n uniqueness=\"server\",\n description=\"Unique identifier for the user, typically an email address.\",\n ),\n ScimSchemaAttribute(\n name=\"name\",\n type=\"complex\",\n description=\"The components of the user's name.\",\n subAttributes=[\n ScimSchemaAttribute(\n name=\"givenName\",\n type=\"string\",\n description=\"The user's first name.\",\n ),\n ScimSchemaAttribute(\n name=\"familyName\",\n type=\"string\",\n description=\"The user's last name.\",\n ),\n ScimSchemaAttribute(\n name=\"formatted\",\n type=\"string\",\n description=\"The full name, including all middle names and titles.\",\n ),\n ],\n ),\n ScimSchemaAttribute(\n name=\"emails\",\n type=\"complex\",\n multiValued=True,\n description=\"Email addresses for the user.\",\n subAttributes=[\n ScimSchemaAttribute(\n name=\"value\",\n type=\"string\",\n description=\"Email address value.\",\n ),\n ScimSchemaAttribute(\n name=\"type\",\n type=\"string\",\n description=\"Label for this email (e.g. 'work').\",\n ),\n ScimSchemaAttribute(\n name=\"primary\",\n type=\"boolean\",\n description=\"Whether this is the primary email.\",\n ),\n ],\n ),\n ScimSchemaAttribute(\n name=\"active\",\n type=\"boolean\",\n description=\"Whether the user account is active.\",\n ),\n ScimSchemaAttribute(\n name=\"externalId\",\n type=\"string\",\n description=\"Identifier from the provisioning client (IdP).\",\n caseExact=True,\n ),\n ],\n)\n\nENTERPRISE_USER_SCHEMA_DEF = ScimSchemaDefinition(\n id=SCIM_ENTERPRISE_USER_SCHEMA,\n name=\"EnterpriseUser\",\n description=\"Enterprise User extension (RFC 7643 §4.3)\",\n attributes=[\n ScimSchemaAttribute(\n name=\"department\",\n type=\"string\",\n description=\"Department.\",\n ),\n ScimSchemaAttribute(\n name=\"manager\",\n type=\"complex\",\n description=\"The user's manager.\",\n subAttributes=[\n ScimSchemaAttribute(\n name=\"value\",\n type=\"string\",\n description=\"Manager user ID.\",\n ),\n ],\n ),\n ],\n)\n\nGROUP_SCHEMA_DEF = ScimSchemaDefinition(\n id=SCIM_GROUP_SCHEMA,\n name=\"Group\",\n description=\"SCIM core Group schema\",\n attributes=[\n ScimSchemaAttribute(\n name=\"displayName\",\n type=\"string\",\n required=True,\n description=\"Human-readable name for the group.\",\n ),\n ScimSchemaAttribute(\n name=\"members\",\n type=\"complex\",\n multiValued=True,\n description=\"Members of the group.\",\n subAttributes=[\n ScimSchemaAttribute(\n name=\"value\",\n type=\"string\",\n description=\"User ID of the group member.\",\n ),\n ScimSchemaAttribute(\n name=\"display\",\n type=\"string\",\n mutability=\"readOnly\",\n description=\"Display name of the group member.\",\n ),\n ],\n ),\n ScimSchemaAttribute(\n name=\"externalId\",\n type=\"string\",\n description=\"Identifier from the provisioning client (IdP).\",\n caseExact=True,\n ),\n ],\n)\n" }, { "path": "backend/ee/onyx/server/seeding.py", "content": "import json\nimport os\nfrom copy import deepcopy\nfrom typing import List\nfrom typing import Optional\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.standard_answer import (\n create_initial_default_standard_answer_category,\n)\nfrom ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload\nfrom ee.onyx.server.enterprise_settings.models import EnterpriseSettings\nfrom ee.onyx.server.enterprise_settings.models import NavigationItem\nfrom ee.onyx.server.enterprise_settings.store import store_analytics_script\nfrom ee.onyx.server.enterprise_settings.store import (\n store_settings as store_ee_settings,\n)\nfrom ee.onyx.server.enterprise_settings.store import upload_logo\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import Tool\nfrom onyx.db.persona import upsert_persona\nfrom onyx.server.features.persona.models import PersonaUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.server.settings.models import Settings\nfrom onyx.server.settings.store import store_settings as store_base_settings\nfrom onyx.utils.logger import setup_logger\n\n\nclass CustomToolSeed(BaseModel):\n name: str\n description: str\n definition_path: str\n custom_headers: Optional[List[dict]] = None\n display_name: Optional[str] = None\n in_code_tool_id: Optional[str] = None\n user_id: Optional[str] = None\n\n\nlogger = setup_logger()\n\n_SEED_CONFIG_ENV_VAR_NAME = \"ENV_SEED_CONFIGURATION\"\n\n\nclass NavigationItemSeed(BaseModel):\n link: str\n title: str\n # NOTE: SVG at this path must not have a width / height specified\n svg_path: str\n\n\nclass SeedConfiguration(BaseModel):\n llms: list[LLMProviderUpsertRequest] | None = None\n admin_user_emails: list[str] | None = None\n seeded_logo_path: str | None = None\n personas: list[PersonaUpsertRequest] | None = None\n settings: Settings | None = None\n enterprise_settings: EnterpriseSettings | None = None\n\n # allows for specifying custom navigation items that have your own custom SVG logos\n nav_item_overrides: list[NavigationItemSeed] | None = None\n\n # Use existing `CUSTOM_ANALYTICS_SECRET_KEY` for reference\n analytics_script_path: str | None = None\n custom_tools: List[CustomToolSeed] | None = None\n\n\ndef _parse_env() -> SeedConfiguration | None:\n seed_config_str = os.getenv(_SEED_CONFIG_ENV_VAR_NAME)\n if not seed_config_str:\n return None\n seed_config = SeedConfiguration.model_validate_json(seed_config_str)\n return seed_config\n\n\ndef _seed_custom_tools(db_session: Session, tools: List[CustomToolSeed]) -> None:\n if tools:\n logger.notice(\"Seeding Custom Tools\")\n for tool in tools:\n try:\n logger.debug(f\"Attempting to seed tool: {tool.name}\")\n logger.debug(f\"Reading definition from: {tool.definition_path}\")\n with open(tool.definition_path, \"r\") as file:\n file_content = file.read()\n if not file_content.strip():\n raise ValueError(\"File is empty\")\n openapi_schema = json.loads(file_content)\n db_tool = Tool(\n name=tool.name,\n description=tool.description,\n openapi_schema=openapi_schema,\n custom_headers=tool.custom_headers,\n display_name=tool.display_name,\n in_code_tool_id=tool.in_code_tool_id,\n user_id=tool.user_id,\n )\n db_session.add(db_tool)\n logger.debug(f\"Successfully added tool: {tool.name}\")\n except FileNotFoundError:\n logger.error(\n f\"Definition file not found for tool {tool.name}: {tool.definition_path}\"\n )\n except json.JSONDecodeError as e:\n logger.error(\n f\"Invalid JSON in definition file for tool {tool.name}: {str(e)}\"\n )\n except Exception as e:\n logger.error(f\"Failed to seed tool {tool.name}: {str(e)}\")\n db_session.commit()\n logger.notice(f\"Successfully seeded {len(tools)} Custom Tools\")\n\n\ndef _seed_llms(\n db_session: Session, llm_upsert_requests: list[LLMProviderUpsertRequest]\n) -> None:\n if not llm_upsert_requests:\n return\n\n logger.notice(\"Seeding LLMs\")\n for request in llm_upsert_requests:\n existing = fetch_existing_llm_provider(name=request.name, db_session=db_session)\n if existing:\n request.id = existing.id\n seeded_providers: list[LLMProviderView] = []\n for llm_upsert_request in llm_upsert_requests:\n try:\n seeded_providers.append(upsert_llm_provider(llm_upsert_request, db_session))\n except ValueError as e:\n logger.warning(\n \"Failed to upsert LLM provider '%s' during seeding: %s\",\n llm_upsert_request.name,\n e,\n )\n\n default_provider = next(\n (p for p in seeded_providers if p.model_configurations), None\n )\n if not default_provider:\n return\n\n visible_configs = [\n mc for mc in default_provider.model_configurations if mc.is_visible\n ]\n default_config = (\n visible_configs[0]\n if visible_configs\n else default_provider.model_configurations[0]\n )\n update_default_provider(\n provider_id=default_provider.id,\n model_name=default_config.name,\n db_session=db_session,\n )\n\n\ndef _seed_personas(db_session: Session, personas: list[PersonaUpsertRequest]) -> None:\n if personas:\n logger.notice(\"Seeding Personas\")\n try:\n for persona in personas:\n upsert_persona(\n user=None, # Seeding is done as admin\n name=persona.name,\n description=persona.description,\n document_set_ids=persona.document_set_ids,\n llm_model_provider_override=persona.llm_model_provider_override,\n llm_model_version_override=persona.llm_model_version_override,\n starter_messages=persona.starter_messages,\n is_public=persona.is_public,\n db_session=db_session,\n tool_ids=persona.tool_ids,\n display_priority=persona.display_priority,\n system_prompt=persona.system_prompt,\n task_prompt=persona.task_prompt,\n datetime_aware=persona.datetime_aware,\n is_featured=persona.is_featured,\n commit=False,\n )\n db_session.commit()\n except Exception:\n logger.exception(\"Failed to seed personas.\")\n raise\n\n\ndef _seed_settings(settings: Settings) -> None:\n logger.notice(\"Seeding Settings\")\n try:\n store_base_settings(settings)\n logger.notice(\"Successfully seeded Settings\")\n except ValueError as e:\n logger.error(f\"Failed to seed Settings: {str(e)}\")\n\n\ndef _seed_enterprise_settings(seed_config: SeedConfiguration) -> None:\n if (\n seed_config.enterprise_settings is not None\n or seed_config.nav_item_overrides is not None\n ):\n final_enterprise_settings = (\n deepcopy(seed_config.enterprise_settings)\n if seed_config.enterprise_settings\n else EnterpriseSettings()\n )\n\n final_nav_items = final_enterprise_settings.custom_nav_items\n if seed_config.nav_item_overrides is not None:\n final_nav_items = []\n for item in seed_config.nav_item_overrides:\n with open(item.svg_path, \"r\") as file:\n svg_content = file.read().strip()\n\n final_nav_items.append(\n NavigationItem(\n link=item.link,\n title=item.title,\n svg_logo=svg_content,\n )\n )\n\n final_enterprise_settings.custom_nav_items = final_nav_items\n\n logger.notice(\"Seeding enterprise settings\")\n store_ee_settings(final_enterprise_settings)\n\n\ndef _seed_logo(logo_path: str | None) -> None:\n if logo_path:\n logger.notice(\"Uploading logo\")\n upload_logo(file=logo_path)\n\n\ndef _seed_analytics_script(seed_config: SeedConfiguration) -> None:\n custom_analytics_secret_key = os.environ.get(\"CUSTOM_ANALYTICS_SECRET_KEY\")\n if seed_config.analytics_script_path and custom_analytics_secret_key:\n logger.notice(\"Seeding analytics script\")\n try:\n with open(seed_config.analytics_script_path, \"r\") as file:\n script_content = file.read()\n analytics_script = AnalyticsScriptUpload(\n script=script_content, secret_key=custom_analytics_secret_key\n )\n store_analytics_script(analytics_script)\n except FileNotFoundError:\n logger.error(\n f\"Analytics script file not found: {seed_config.analytics_script_path}\"\n )\n except ValueError as e:\n logger.error(f\"Failed to seed analytics script: {str(e)}\")\n\n\ndef get_seed_config() -> SeedConfiguration | None:\n return _parse_env()\n\n\ndef seed_db() -> None:\n seed_config = _parse_env()\n if seed_config is None:\n logger.debug(\"No seeding configuration file passed\")\n return\n\n with get_session_with_current_tenant() as db_session:\n if seed_config.llms is not None:\n _seed_llms(db_session, seed_config.llms)\n if seed_config.personas is not None:\n _seed_personas(db_session, seed_config.personas)\n if seed_config.settings is not None:\n _seed_settings(seed_config.settings)\n if seed_config.custom_tools is not None:\n _seed_custom_tools(db_session, seed_config.custom_tools)\n\n _seed_logo(seed_config.seeded_logo_path)\n _seed_enterprise_settings(seed_config)\n _seed_analytics_script(seed_config)\n\n logger.notice(\"Verifying default standard answer category exists.\")\n create_initial_default_standard_answer_category(db_session)\n" }, { "path": "backend/ee/onyx/server/settings/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/settings/api.py", "content": "\"\"\"EE Settings API - provides license-aware settings override.\"\"\"\n\nfrom redis.exceptions import RedisError\nfrom sqlalchemy.exc import SQLAlchemyError\n\nfrom ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED\nfrom ee.onyx.db.license import get_cached_license_metadata\nfrom ee.onyx.db.license import refresh_license_cache\nfrom onyx.cache.interface import CACHE_TRANSIENT_ERRORS\nfrom onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.server.settings.models import ApplicationStatus\nfrom onyx.server.settings.models import Settings\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n# Only GATED_ACCESS actually blocks access - other statuses are for notifications\n_BLOCKING_STATUS = ApplicationStatus.GATED_ACCESS\n\n\ndef check_ee_features_enabled() -> bool:\n \"\"\"EE version: checks if EE features should be available.\n\n Returns True if:\n - LICENSE_ENFORCEMENT_ENABLED is False (legacy/rollout mode)\n - Cloud mode (MULTI_TENANT) - cloud handles its own gating\n - Self-hosted with a valid (non-expired) license\n\n Returns False if:\n - Self-hosted with no license (never subscribed)\n - Self-hosted with expired license\n \"\"\"\n if not LICENSE_ENFORCEMENT_ENABLED:\n # License enforcement disabled - allow EE features (legacy behavior)\n return True\n\n if MULTI_TENANT:\n # Cloud mode - EE features always available (gating handled by is_tenant_gated)\n return True\n\n # Self-hosted with enforcement - check for valid license\n tenant_id = get_current_tenant_id()\n try:\n metadata = get_cached_license_metadata(tenant_id)\n if not metadata:\n # Cache miss — warm from DB so cold-start doesn't block EE features\n try:\n with get_session_with_current_tenant() as db_session:\n metadata = refresh_license_cache(db_session, tenant_id)\n except SQLAlchemyError as db_error:\n logger.warning(f\"Failed to load license from DB: {db_error}\")\n\n if metadata and metadata.status != _BLOCKING_STATUS:\n # Has a valid license (GRACE_PERIOD/PAYMENT_REMINDER still allow EE features)\n return True\n except RedisError as e:\n logger.warning(f\"Failed to check license for EE features: {e}\")\n # Fail closed - if Redis is down, other things will break anyway\n return False\n\n # No license or GATED_ACCESS - no EE features\n return False\n\n\ndef apply_license_status_to_settings(settings: Settings) -> Settings:\n \"\"\"EE version: checks license status for self-hosted deployments.\n\n For self-hosted, looks up license metadata and overrides application_status\n if the license indicates GATED_ACCESS (fully expired).\n\n Also sets ee_features_enabled based on license status to control\n visibility of EE features in the UI.\n\n For multi-tenant (cloud), the settings already have the correct status\n from the control plane, so no override is needed.\n\n If LICENSE_ENFORCEMENT_ENABLED is false, ee_features_enabled is set to True\n (since EE code was loaded via ENABLE_PAID_ENTERPRISE_EDITION_FEATURES).\n \"\"\"\n if not LICENSE_ENFORCEMENT_ENABLED:\n # License enforcement disabled - EE code is loaded via\n # ENABLE_PAID_ENTERPRISE_EDITION_FEATURES, so EE features are on\n settings.ee_features_enabled = True\n return settings\n\n if MULTI_TENANT:\n # Cloud mode - EE features always available (gating handled by is_tenant_gated)\n settings.ee_features_enabled = True\n return settings\n\n tenant_id = get_current_tenant_id()\n try:\n metadata = get_cached_license_metadata(tenant_id)\n if not metadata:\n # Cache miss (e.g. after TTL expiry). Fall back to DB so\n # the /settings request doesn't falsely return GATED_ACCESS\n # while the cache is cold.\n try:\n with get_session_with_current_tenant() as db_session:\n metadata = refresh_license_cache(db_session, tenant_id)\n except SQLAlchemyError as db_error:\n logger.warning(\n f\"Failed to load license from DB for settings: {db_error}\"\n )\n\n if metadata:\n if metadata.status == _BLOCKING_STATUS:\n settings.application_status = metadata.status\n settings.ee_features_enabled = False\n elif metadata.used_seats > metadata.seats:\n # License is valid but seat limit exceeded\n settings.application_status = ApplicationStatus.SEAT_LIMIT_EXCEEDED\n settings.seat_count = metadata.seats\n settings.used_seats = metadata.used_seats\n settings.ee_features_enabled = True\n else:\n # Has a valid license (GRACE_PERIOD/PAYMENT_REMINDER still allow EE features)\n settings.ee_features_enabled = True\n else:\n # No license found in cache or DB.\n if ENTERPRISE_EDITION_ENABLED:\n # Legacy EE flag is set → prior EE usage (e.g. permission\n # syncing) means indexed data may need protection.\n settings.application_status = _BLOCKING_STATUS\n settings.ee_features_enabled = False\n except CACHE_TRANSIENT_ERRORS as e:\n logger.warning(f\"Failed to check license metadata for settings: {e}\")\n # Fail closed - disable EE features if we can't verify license\n settings.ee_features_enabled = False\n\n return settings\n" }, { "path": "backend/ee/onyx/server/tenant_usage_limits.py", "content": "\"\"\"Tenant-specific usage limit overrides from the control plane (EE version).\"\"\"\n\nimport time\n\nimport requests\n\nfrom ee.onyx.server.tenants.access import generate_data_plane_token\nfrom onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.server.tenant_usage_limits import TenantUsageLimitOverrides\nfrom onyx.server.usage_limits import NO_LIMIT\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# In-memory storage for tenant overrides (populated at startup)\n_tenant_usage_limit_overrides: dict[str, TenantUsageLimitOverrides] | None = None\n_last_fetch_time: float = 0.0\n_FETCH_INTERVAL = 60 * 60 * 24 # 24 hours\n_ERROR_FETCH_INTERVAL = 30 * 60 # 30 minutes (if the last fetch failed)\n\n\ndef fetch_usage_limit_overrides() -> dict[str, TenantUsageLimitOverrides] | None:\n \"\"\"\n Fetch tenant-specific usage limit overrides from the control plane.\n\n Returns:\n Dictionary mapping tenant_id to their specific limit overrides.\n Returns empty dict on any error (falls back to defaults).\n \"\"\"\n try:\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n url = f\"{CONTROL_PLANE_API_BASE_URL}/usage-limit-overrides\"\n response = requests.get(url, headers=headers, timeout=30)\n response.raise_for_status()\n\n tenant_overrides = response.json()\n\n # Parse each tenant's overrides\n result: dict[str, TenantUsageLimitOverrides] = {}\n for override_data in tenant_overrides:\n tenant_id = override_data[\"tenant_id\"]\n try:\n result[tenant_id] = TenantUsageLimitOverrides(**override_data)\n except Exception as e:\n logger.warning(\n f\"Failed to parse usage limit overrides for tenant {tenant_id}: {e}\"\n )\n\n return (\n result or None\n ) # if empty dictionary, something went wrong and we shouldn't enforce limits\n\n except requests.exceptions.RequestException as e:\n logger.warning(f\"Failed to fetch usage limit overrides from control plane: {e}\")\n return None\n except Exception as e:\n logger.error(f\"Error parsing usage limit overrides: {e}\")\n return None\n\n\ndef load_usage_limit_overrides() -> None:\n \"\"\"\n Load tenant usage limit overrides from the control plane.\n \"\"\"\n global _tenant_usage_limit_overrides\n global _last_fetch_time\n\n logger.info(\"Loading tenant usage limit overrides from control plane...\")\n overrides = fetch_usage_limit_overrides()\n\n _last_fetch_time = time.time()\n\n # use the new result if it exists, otherwise use the old result\n # (prevents us from updating to a failed fetch result)\n _tenant_usage_limit_overrides = overrides or _tenant_usage_limit_overrides\n\n if overrides:\n logger.info(f\"Loaded usage limit overrides for {len(overrides)} tenants\")\n else:\n logger.info(\"No tenant-specific usage limit overrides found\")\n\n\ndef unlimited(tenant_id: str) -> TenantUsageLimitOverrides:\n return TenantUsageLimitOverrides(\n tenant_id=tenant_id,\n llm_cost_cents_trial=NO_LIMIT,\n llm_cost_cents_paid=NO_LIMIT,\n chunks_indexed_trial=NO_LIMIT,\n chunks_indexed_paid=NO_LIMIT,\n api_calls_trial=NO_LIMIT,\n api_calls_paid=NO_LIMIT,\n non_streaming_calls_trial=NO_LIMIT,\n non_streaming_calls_paid=NO_LIMIT,\n )\n\n\ndef get_tenant_usage_limit_overrides(\n tenant_id: str,\n) -> TenantUsageLimitOverrides | None:\n \"\"\"\n Get the usage limit overrides for a specific tenant.\n\n Args:\n tenant_id: The tenant ID to look up\n\n Returns:\n TenantUsageLimitOverrides if the tenant has overrides, None otherwise.\n \"\"\"\n\n if DEV_MODE: # in dev mode, we return unlimited limits for all tenants\n return unlimited(tenant_id)\n\n global _tenant_usage_limit_overrides\n time_since = time.time() - _last_fetch_time\n if (\n _tenant_usage_limit_overrides is None and time_since > _ERROR_FETCH_INTERVAL\n ) or (time_since > _FETCH_INTERVAL):\n logger.debug(\n f\"Last fetch time: {_last_fetch_time}, time since last fetch: {time_since}\"\n )\n\n load_usage_limit_overrides()\n\n # If we have failed to fetch from the control plane or we're in dev mode, don't usage limit anyone.\n if _tenant_usage_limit_overrides is None or DEV_MODE:\n return unlimited(tenant_id)\n return _tenant_usage_limit_overrides.get(tenant_id)\n" }, { "path": "backend/ee/onyx/server/tenants/__init__.py", "content": "" }, { "path": "backend/ee/onyx/server/tenants/access.py", "content": "from datetime import datetime\nfrom datetime import timedelta\n\nimport jwt\nfrom fastapi import HTTPException\nfrom fastapi import Request\n\nfrom onyx.configs.app_configs import DATA_PLANE_SECRET\nfrom onyx.configs.app_configs import EXPECTED_API_KEY\nfrom onyx.configs.app_configs import JWT_ALGORITHM\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef generate_data_plane_token() -> str:\n if DATA_PLANE_SECRET is None:\n raise ValueError(\"DATA_PLANE_SECRET is not set\")\n\n payload = {\n \"iss\": \"data_plane\",\n \"exp\": datetime.utcnow() + timedelta(minutes=5),\n \"iat\": datetime.utcnow(),\n \"scope\": \"api_access\",\n }\n\n token = jwt.encode(payload, DATA_PLANE_SECRET, algorithm=JWT_ALGORITHM)\n return token\n\n\nasync def control_plane_dep(request: Request) -> None:\n api_key = request.headers.get(\"X-API-KEY\")\n if api_key != EXPECTED_API_KEY:\n logger.warning(\"Invalid API key\")\n raise HTTPException(status_code=401, detail=\"Invalid API key\")\n\n auth_header = request.headers.get(\"Authorization\")\n if not auth_header or not auth_header.startswith(\"Bearer \"):\n logger.warning(\"Invalid authorization header\")\n raise HTTPException(status_code=401, detail=\"Invalid authorization header\")\n\n token = auth_header.split(\" \")[1]\n try:\n payload = jwt.decode(token, DATA_PLANE_SECRET, algorithms=[JWT_ALGORITHM])\n if payload.get(\"scope\") != \"tenant:create\":\n logger.warning(\"Insufficient permissions\")\n raise HTTPException(status_code=403, detail=\"Insufficient permissions\")\n except jwt.ExpiredSignatureError:\n logger.warning(\"Token has expired\")\n raise HTTPException(status_code=401, detail=\"Token has expired\")\n except jwt.InvalidTokenError:\n logger.warning(\"Invalid token\")\n raise HTTPException(status_code=401, detail=\"Invalid token\")\n" }, { "path": "backend/ee/onyx/server/tenants/admin_api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Response\nfrom fastapi_users import exceptions\n\nfrom ee.onyx.auth.users import current_cloud_superuser\nfrom ee.onyx.server.tenants.models import ImpersonateRequest\nfrom ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email\nfrom onyx.auth.users import auth_backend\nfrom onyx.auth.users import get_redis_strategy\nfrom onyx.auth.users import User\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.users import get_user_by_email\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/tenants\")\n\n\n@router.post(\"/impersonate\")\nasync def impersonate_user(\n impersonate_request: ImpersonateRequest,\n _: User = Depends(current_cloud_superuser),\n) -> Response:\n \"\"\"Allows a cloud superuser to impersonate another user by generating an impersonation JWT token\"\"\"\n try:\n tenant_id = get_tenant_id_for_email(impersonate_request.email)\n except exceptions.UserNotExists:\n detail = f\"User has no tenant mapping: {impersonate_request.email=}\"\n logger.warning(detail)\n raise HTTPException(status_code=422, detail=detail)\n\n with get_session_with_tenant(tenant_id=tenant_id) as tenant_session:\n user_to_impersonate = get_user_by_email(\n impersonate_request.email, tenant_session\n )\n if user_to_impersonate is None:\n detail = (\n f\"User not found in tenant: {impersonate_request.email=} {tenant_id=}\"\n )\n logger.warning(detail)\n raise HTTPException(status_code=422, detail=detail)\n\n token = await get_redis_strategy().write_token(user_to_impersonate)\n\n response = await auth_backend.transport.get_login_response(token)\n response.set_cookie(\n key=\"fastapiusersauth\",\n value=token,\n httponly=True,\n secure=True,\n samesite=\"lax\",\n )\n return response\n" }, { "path": "backend/ee/onyx/server/tenants/anonymous_user_path.py", "content": "from sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import TenantAnonymousUserPath\n\n\ndef get_anonymous_user_path(tenant_id: str, db_session: Session) -> str | None:\n result = db_session.execute(\n select(TenantAnonymousUserPath).where(\n TenantAnonymousUserPath.tenant_id == tenant_id\n )\n )\n result_scalar = result.scalar_one_or_none()\n if result_scalar:\n return result_scalar.anonymous_user_path\n else:\n return None\n\n\ndef modify_anonymous_user_path(\n tenant_id: str, anonymous_user_path: str, db_session: Session\n) -> None:\n # Enforce lowercase path at DB operation level\n anonymous_user_path = anonymous_user_path.lower()\n\n existing_entry = (\n db_session.query(TenantAnonymousUserPath).filter_by(tenant_id=tenant_id).first()\n )\n\n if existing_entry:\n existing_entry.anonymous_user_path = anonymous_user_path\n\n else:\n new_entry = TenantAnonymousUserPath(\n tenant_id=tenant_id, anonymous_user_path=anonymous_user_path\n )\n db_session.add(new_entry)\n\n db_session.commit()\n\n\ndef get_tenant_id_for_anonymous_user_path(\n anonymous_user_path: str, db_session: Session\n) -> str | None:\n result = db_session.execute(\n select(TenantAnonymousUserPath).where(\n TenantAnonymousUserPath.anonymous_user_path == anonymous_user_path\n )\n )\n result_scalar = result.scalar_one_or_none()\n if result_scalar:\n return result_scalar.tenant_id\n else:\n return None\n\n\ndef validate_anonymous_user_path(path: str) -> None:\n if not path or \"/\" in path or not path.replace(\"-\", \"\").isalnum():\n raise ValueError(\"Invalid path. Use only letters, numbers, and hyphens.\")\n" }, { "path": "backend/ee/onyx/server/tenants/anonymous_users_api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Response\nfrom sqlalchemy.exc import IntegrityError\n\nfrom ee.onyx.auth.users import generate_anonymous_user_jwt_token\nfrom ee.onyx.server.tenants.anonymous_user_path import get_anonymous_user_path\nfrom ee.onyx.server.tenants.anonymous_user_path import (\n get_tenant_id_for_anonymous_user_path,\n)\nfrom ee.onyx.server.tenants.anonymous_user_path import modify_anonymous_user_path\nfrom ee.onyx.server.tenants.anonymous_user_path import validate_anonymous_user_path\nfrom ee.onyx.server.tenants.models import AnonymousUserPath\nfrom onyx.auth.users import anonymous_user_enabled\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import User\nfrom onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME\nfrom onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/tenants\")\n\n\n@router.get(\"/anonymous-user-path\")\nasync def get_anonymous_user_path_api(\n _: User = Depends(current_admin_user),\n) -> AnonymousUserPath:\n tenant_id = get_current_tenant_id()\n\n if tenant_id is None:\n raise HTTPException(status_code=404, detail=\"Tenant not found\")\n\n with get_session_with_shared_schema() as db_session:\n current_path = get_anonymous_user_path(tenant_id, db_session)\n\n return AnonymousUserPath(anonymous_user_path=current_path)\n\n\n@router.post(\"/anonymous-user-path\")\nasync def set_anonymous_user_path_api(\n anonymous_user_path: str,\n _: User = Depends(current_admin_user),\n) -> None:\n tenant_id = get_current_tenant_id()\n try:\n validate_anonymous_user_path(anonymous_user_path)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n with get_session_with_shared_schema() as db_session:\n try:\n modify_anonymous_user_path(tenant_id, anonymous_user_path, db_session)\n except IntegrityError:\n raise HTTPException(\n status_code=409,\n detail=\"The anonymous user path is already in use. Please choose a different path.\",\n )\n except Exception as e:\n logger.exception(f\"Failed to modify anonymous user path: {str(e)}\")\n raise HTTPException(\n status_code=500,\n detail=\"An unexpected error occurred while modifying the anonymous user path\",\n )\n\n\n@router.post(\"/anonymous-user\")\nasync def login_as_anonymous_user(\n anonymous_user_path: str,\n) -> Response:\n with get_session_with_shared_schema() as db_session:\n tenant_id = get_tenant_id_for_anonymous_user_path(\n anonymous_user_path, db_session\n )\n if not tenant_id:\n raise HTTPException(status_code=404, detail=\"Tenant not found\")\n\n if not anonymous_user_enabled(tenant_id=tenant_id):\n raise HTTPException(status_code=403, detail=\"Anonymous user is not enabled\")\n\n token = generate_anonymous_user_jwt_token(tenant_id)\n\n response = Response()\n response.delete_cookie(FASTAPI_USERS_AUTH_COOKIE_NAME)\n response.set_cookie(\n key=ANONYMOUS_USER_COOKIE_NAME,\n value=token,\n httponly=True,\n secure=True,\n samesite=\"strict\",\n )\n return response\n" }, { "path": "backend/ee/onyx/server/tenants/api.py", "content": "from fastapi import APIRouter\n\nfrom ee.onyx.server.tenants.admin_api import router as admin_router\nfrom ee.onyx.server.tenants.anonymous_users_api import router as anonymous_users_router\nfrom ee.onyx.server.tenants.billing_api import router as billing_router\nfrom ee.onyx.server.tenants.proxy import router as proxy_router\nfrom ee.onyx.server.tenants.team_membership_api import router as team_membership_router\nfrom ee.onyx.server.tenants.tenant_management_api import (\n router as tenant_management_router,\n)\nfrom ee.onyx.server.tenants.user_invitations_api import (\n router as user_invitations_router,\n)\n\n# Create a main router to include all sub-routers\n# Note: We don't add a prefix here as each router already has the /tenants prefix\nrouter = APIRouter()\n\n# Include all the individual routers\nrouter.include_router(admin_router)\nrouter.include_router(anonymous_users_router)\nrouter.include_router(billing_router)\nrouter.include_router(team_membership_router)\nrouter.include_router(tenant_management_router)\nrouter.include_router(user_invitations_router)\nrouter.include_router(proxy_router)\n" }, { "path": "backend/ee/onyx/server/tenants/billing.py", "content": "from typing import cast\nfrom typing import Literal\n\nimport requests\nimport stripe\n\nfrom ee.onyx.configs.app_configs import STRIPE_SECRET_KEY\nfrom ee.onyx.server.tenants.access import generate_data_plane_token\nfrom ee.onyx.server.tenants.models import BillingInformation\nfrom ee.onyx.server.tenants.models import SubscriptionStatusResponse\nfrom onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL\nfrom onyx.utils.logger import setup_logger\n\nstripe.api_key = STRIPE_SECRET_KEY\n\nlogger = setup_logger()\n\n\ndef fetch_stripe_checkout_session(\n tenant_id: str,\n billing_period: Literal[\"monthly\", \"annual\"] = \"monthly\",\n seats: int | None = None,\n) -> str:\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n url = f\"{CONTROL_PLANE_API_BASE_URL}/create-checkout-session\"\n payload = {\n \"tenant_id\": tenant_id,\n \"billing_period\": billing_period,\n \"seats\": seats,\n }\n response = requests.post(url, headers=headers, json=payload)\n if not response.ok:\n try:\n data = response.json()\n error_msg = (\n data.get(\"error\")\n or f\"Request failed with status {response.status_code}\"\n )\n except (ValueError, requests.exceptions.JSONDecodeError):\n error_msg = f\"Request failed with status {response.status_code}: {response.text[:200]}\"\n raise Exception(error_msg)\n data = response.json()\n if data.get(\"error\"):\n raise Exception(data[\"error\"])\n return data[\"sessionId\"]\n\n\ndef fetch_tenant_stripe_information(tenant_id: str) -> dict:\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n url = f\"{CONTROL_PLANE_API_BASE_URL}/tenant-stripe-information\"\n params = {\"tenant_id\": tenant_id}\n response = requests.get(url, headers=headers, params=params)\n response.raise_for_status()\n return response.json()\n\n\ndef fetch_billing_information(\n tenant_id: str,\n) -> BillingInformation | SubscriptionStatusResponse:\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n url = f\"{CONTROL_PLANE_API_BASE_URL}/billing-information\"\n params = {\"tenant_id\": tenant_id}\n response = requests.get(url, headers=headers, params=params)\n response.raise_for_status()\n\n response_data = response.json()\n\n # Check if the response indicates no subscription\n if (\n isinstance(response_data, dict)\n and \"subscribed\" in response_data\n and not response_data[\"subscribed\"]\n ):\n return SubscriptionStatusResponse(**response_data)\n\n # Otherwise, parse as BillingInformation\n return BillingInformation(**response_data)\n\n\ndef fetch_customer_portal_session(tenant_id: str, return_url: str | None = None) -> str:\n \"\"\"\n Fetch a Stripe customer portal session URL from the control plane.\n NOTE: This is currently only used for multi-tenant (cloud) deployments.\n Self-hosted proxy endpoints will be added in a future phase.\n \"\"\"\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n url = f\"{CONTROL_PLANE_API_BASE_URL}/create-customer-portal-session\"\n payload = {\"tenant_id\": tenant_id}\n if return_url:\n payload[\"return_url\"] = return_url\n response = requests.post(url, headers=headers, json=payload)\n response.raise_for_status()\n return response.json()[\"url\"]\n\n\ndef register_tenant_users(tenant_id: str, number_of_users: int) -> stripe.Subscription:\n \"\"\"\n Update the number of seats for a tenant's subscription.\n Preserves the existing price (monthly, annual, or grandfathered).\n \"\"\"\n response = fetch_tenant_stripe_information(tenant_id)\n stripe_subscription_id = cast(str, response.get(\"stripe_subscription_id\"))\n\n subscription = stripe.Subscription.retrieve(stripe_subscription_id)\n subscription_item = subscription[\"items\"][\"data\"][0]\n\n # Use existing price to preserve the customer's current plan\n current_price_id = subscription_item.price.id\n\n updated_subscription = stripe.Subscription.modify(\n stripe_subscription_id,\n items=[\n {\n \"id\": subscription_item.id,\n \"price\": current_price_id,\n \"quantity\": number_of_users,\n }\n ],\n metadata={\"tenant_id\": str(tenant_id)},\n )\n return updated_subscription\n" }, { "path": "backend/ee/onyx/server/tenants/billing_api.py", "content": "\"\"\"Billing API endpoints for cloud multi-tenant deployments.\n\nDEPRECATED: These /tenants/* billing endpoints are being replaced by /admin/billing/*\nwhich provides a unified API for both self-hosted and cloud deployments.\n\nTODO(ENG-3533): Migrate frontend to use /admin/billing/* endpoints and remove this file.\nhttps://linear.app/onyx-app/issue/ENG-3533/migrate-tenantsbilling-adminbilling\n\nCurrent endpoints to migrate:\n- GET /tenants/billing-information -> GET /admin/billing/information\n- POST /tenants/create-customer-portal-session -> POST /admin/billing/portal-session\n- POST /tenants/create-subscription-session -> POST /admin/billing/checkout-session\n- GET /tenants/stripe-publishable-key -> (keep as-is, shared endpoint)\n\nNote: /tenants/product-gating/* endpoints are control-plane-to-data-plane calls\nand are NOT part of this migration - they stay here.\n\"\"\"\n\nimport asyncio\n\nimport httpx\nfrom fastapi import APIRouter\nfrom fastapi import Depends\n\nfrom ee.onyx.auth.users import current_admin_user\nfrom ee.onyx.server.tenants.access import control_plane_dep\nfrom ee.onyx.server.tenants.billing import fetch_billing_information\nfrom ee.onyx.server.tenants.billing import fetch_customer_portal_session\nfrom ee.onyx.server.tenants.billing import fetch_stripe_checkout_session\nfrom ee.onyx.server.tenants.models import BillingInformation\nfrom ee.onyx.server.tenants.models import CreateCheckoutSessionRequest\nfrom ee.onyx.server.tenants.models import CreateSubscriptionSessionRequest\nfrom ee.onyx.server.tenants.models import ProductGatingFullSyncRequest\nfrom ee.onyx.server.tenants.models import ProductGatingRequest\nfrom ee.onyx.server.tenants.models import ProductGatingResponse\nfrom ee.onyx.server.tenants.models import StripePublishableKeyResponse\nfrom ee.onyx.server.tenants.models import SubscriptionSessionResponse\nfrom ee.onyx.server.tenants.models import SubscriptionStatusResponse\nfrom ee.onyx.server.tenants.product_gating import overwrite_full_gated_set\nfrom ee.onyx.server.tenants.product_gating import store_product_gating\nfrom onyx.auth.users import User\nfrom onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_OVERRIDE\nfrom onyx.configs.app_configs import STRIPE_PUBLISHABLE_KEY_URL\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/tenants\")\n\n# Cache for Stripe publishable key to avoid hitting S3 on every request\n_stripe_publishable_key_cache: str | None = None\n_stripe_key_lock = asyncio.Lock()\n\n\n@router.post(\"/product-gating\")\ndef gate_product(\n product_gating_request: ProductGatingRequest, _: None = Depends(control_plane_dep)\n) -> ProductGatingResponse:\n \"\"\"\n Gating the product means that the product is not available to the tenant.\n They will be directed to the billing page.\n We gate the product when their subscription has ended.\n \"\"\"\n try:\n store_product_gating(\n product_gating_request.tenant_id, product_gating_request.application_status\n )\n return ProductGatingResponse(updated=True, error=None)\n\n except Exception as e:\n logger.exception(\"Failed to gate product\")\n return ProductGatingResponse(updated=False, error=str(e))\n\n\n@router.post(\"/product-gating/full-sync\")\ndef gate_product_full_sync(\n product_gating_request: ProductGatingFullSyncRequest,\n _: None = Depends(control_plane_dep),\n) -> ProductGatingResponse:\n \"\"\"\n Bulk operation to overwrite the entire gated tenant set.\n This replaces all currently gated tenants with the provided list.\n Gated tenants are not available to access the product and will be\n directed to the billing page when their subscription has ended.\n \"\"\"\n try:\n overwrite_full_gated_set(product_gating_request.gated_tenant_ids)\n return ProductGatingResponse(updated=True, error=None)\n\n except Exception as e:\n logger.exception(\"Failed to gate products during full sync\")\n return ProductGatingResponse(updated=False, error=str(e))\n\n\n@router.get(\"/billing-information\")\nasync def billing_information(\n _: User = Depends(current_admin_user),\n) -> BillingInformation | SubscriptionStatusResponse:\n logger.info(\"Fetching billing information\")\n tenant_id = get_current_tenant_id()\n return fetch_billing_information(tenant_id)\n\n\n@router.post(\"/create-customer-portal-session\")\nasync def create_customer_portal_session(\n _: User = Depends(current_admin_user),\n) -> dict:\n \"\"\"Create a Stripe customer portal session via the control plane.\"\"\"\n tenant_id = get_current_tenant_id()\n return_url = f\"{WEB_DOMAIN}/admin/billing\"\n\n try:\n portal_url = fetch_customer_portal_session(tenant_id, return_url)\n return {\"stripe_customer_portal_url\": portal_url}\n except OnyxError:\n raise\n except Exception:\n logger.exception(\"Failed to create customer portal session\")\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Failed to create customer portal session\",\n )\n\n\n@router.post(\"/create-checkout-session\")\nasync def create_checkout_session(\n request: CreateCheckoutSessionRequest | None = None,\n _: User = Depends(current_admin_user),\n) -> dict:\n \"\"\"Create a Stripe checkout session via the control plane.\"\"\"\n tenant_id = get_current_tenant_id()\n billing_period = request.billing_period if request else \"monthly\"\n seats = request.seats if request else None\n\n try:\n checkout_url = fetch_stripe_checkout_session(tenant_id, billing_period, seats)\n return {\"stripe_checkout_url\": checkout_url}\n except OnyxError:\n raise\n except Exception:\n logger.exception(\"Failed to create checkout session\")\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Failed to create checkout session\",\n )\n\n\n@router.post(\"/create-subscription-session\")\nasync def create_subscription_session(\n request: CreateSubscriptionSessionRequest | None = None,\n _: User = Depends(current_admin_user),\n) -> SubscriptionSessionResponse:\n try:\n tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()\n if not tenant_id:\n raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, \"Tenant ID not found\")\n\n billing_period = request.billing_period if request else \"monthly\"\n session_id = fetch_stripe_checkout_session(tenant_id, billing_period)\n return SubscriptionSessionResponse(sessionId=session_id)\n\n except OnyxError:\n raise\n except Exception:\n logger.exception(\"Failed to create subscription session\")\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Failed to create subscription session\",\n )\n\n\n@router.get(\"/stripe-publishable-key\")\nasync def get_stripe_publishable_key() -> StripePublishableKeyResponse:\n \"\"\"\n Fetch the Stripe publishable key.\n Priority: env var override (for testing) > S3 bucket (production).\n This endpoint is public (no auth required) since publishable keys are safe to expose.\n The key is cached in memory to avoid hitting S3 on every request.\n \"\"\"\n global _stripe_publishable_key_cache\n\n # Fast path: return cached value without lock\n if _stripe_publishable_key_cache:\n return StripePublishableKeyResponse(\n publishable_key=_stripe_publishable_key_cache\n )\n\n # Use lock to prevent concurrent S3 requests\n async with _stripe_key_lock:\n # Double-check after acquiring lock (another request may have populated cache)\n if _stripe_publishable_key_cache:\n return StripePublishableKeyResponse(\n publishable_key=_stripe_publishable_key_cache\n )\n\n # Check for env var override first (for local testing with pk_test_* keys)\n if STRIPE_PUBLISHABLE_KEY_OVERRIDE:\n key = STRIPE_PUBLISHABLE_KEY_OVERRIDE.strip()\n if not key.startswith(\"pk_\"):\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Invalid Stripe publishable key format\",\n )\n _stripe_publishable_key_cache = key\n return StripePublishableKeyResponse(publishable_key=key)\n\n # Fall back to S3 bucket\n if not STRIPE_PUBLISHABLE_KEY_URL:\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Stripe publishable key is not configured\",\n )\n\n try:\n async with httpx.AsyncClient() as client:\n response = await client.get(STRIPE_PUBLISHABLE_KEY_URL)\n response.raise_for_status()\n key = response.text.strip()\n\n # Validate key format\n if not key.startswith(\"pk_\"):\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Invalid Stripe publishable key format\",\n )\n\n _stripe_publishable_key_cache = key\n return StripePublishableKeyResponse(publishable_key=key)\n except httpx.HTTPError:\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR,\n \"Failed to fetch Stripe publishable key\",\n )\n" }, { "path": "backend/ee/onyx/server/tenants/models.py", "content": "from datetime import datetime\nfrom typing import Literal\n\nfrom pydantic import BaseModel\n\nfrom onyx.server.settings.models import ApplicationStatus\n\n\nclass CheckoutSessionCreationRequest(BaseModel):\n quantity: int\n\n\nclass CreateTenantRequest(BaseModel):\n tenant_id: str\n initial_admin_email: str\n\n\nclass ProductGatingRequest(BaseModel):\n tenant_id: str\n application_status: ApplicationStatus\n\n\nclass ProductGatingFullSyncRequest(BaseModel):\n gated_tenant_ids: list[str]\n\n\nclass SubscriptionStatusResponse(BaseModel):\n subscribed: bool\n\n\nclass BillingInformation(BaseModel):\n stripe_subscription_id: str\n status: str\n current_period_start: datetime\n current_period_end: datetime\n number_of_seats: int\n cancel_at_period_end: bool\n canceled_at: datetime | None\n trial_start: datetime | None\n trial_end: datetime | None\n seats: int\n payment_method_enabled: bool\n\n\nclass CreateCheckoutSessionRequest(BaseModel):\n billing_period: Literal[\"monthly\", \"annual\"] = \"monthly\"\n seats: int | None = None\n email: str | None = None\n\n\nclass CheckoutSessionCreationResponse(BaseModel):\n id: str\n\n\nclass ImpersonateRequest(BaseModel):\n email: str\n\n\nclass TenantCreationPayload(BaseModel):\n tenant_id: str\n email: str\n referral_source: str | None = None\n\n\nclass TenantDeletionPayload(BaseModel):\n tenant_id: str\n email: str\n\n\nclass AnonymousUserPath(BaseModel):\n anonymous_user_path: str | None\n\n\nclass ProductGatingResponse(BaseModel):\n updated: bool\n error: str | None\n\n\nclass SubscriptionSessionResponse(BaseModel):\n sessionId: str\n\n\nclass CreateSubscriptionSessionRequest(BaseModel):\n \"\"\"Request to create a subscription checkout session.\"\"\"\n\n billing_period: Literal[\"monthly\", \"annual\"] = \"monthly\"\n\n\nclass TenantByDomainResponse(BaseModel):\n tenant_id: str\n number_of_users: int\n creator_email: str\n\n\nclass TenantByDomainRequest(BaseModel):\n email: str\n\n\nclass RequestInviteRequest(BaseModel):\n tenant_id: str\n\n\nclass RequestInviteResponse(BaseModel):\n success: bool\n message: str\n\n\nclass PendingUserSnapshot(BaseModel):\n email: str\n\n\nclass ApproveUserRequest(BaseModel):\n email: str\n\n\nclass StripePublishableKeyResponse(BaseModel):\n publishable_key: str\n" }, { "path": "backend/ee/onyx/server/tenants/product_gating.py", "content": "from typing import cast\n\nfrom ee.onyx.configs.app_configs import GATED_TENANTS_KEY\nfrom onyx.configs.constants import ONYX_CLOUD_TENANT_ID\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.server.settings.models import ApplicationStatus\nfrom onyx.server.settings.store import load_settings\nfrom onyx.server.settings.store import store_settings\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\n\ndef update_tenant_gating(tenant_id: str, status: ApplicationStatus) -> None:\n redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n\n # Maintain the GATED_ACCESS set\n if status == ApplicationStatus.GATED_ACCESS:\n redis_client.sadd(GATED_TENANTS_KEY, tenant_id)\n else:\n redis_client.srem(GATED_TENANTS_KEY, tenant_id)\n\n\ndef store_product_gating(tenant_id: str, application_status: ApplicationStatus) -> None:\n try:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n settings = load_settings()\n settings.application_status = application_status\n store_settings(settings)\n\n # Store gated tenant information in Redis\n update_tenant_gating(tenant_id, application_status)\n\n if token is not None:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n except Exception:\n logger.exception(\"Failed to gate product\")\n raise\n\n\ndef overwrite_full_gated_set(tenant_ids: list[str]) -> None:\n redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n\n pipeline = redis_client.pipeline()\n\n # using pipeline doesn't automatically add the tenant_id prefix\n full_gated_set_key = f\"{ONYX_CLOUD_TENANT_ID}:{GATED_TENANTS_KEY}\"\n\n # Clear the existing set\n pipeline.delete(full_gated_set_key)\n\n # Add all tenant IDs to the set and set their status\n for tenant_id in tenant_ids:\n pipeline.sadd(full_gated_set_key, tenant_id)\n\n # Execute all commands at once\n pipeline.execute()\n\n\ndef get_gated_tenants() -> set[str]:\n redis_client = get_redis_replica_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n gated_tenants_bytes = cast(set[bytes], redis_client.smembers(GATED_TENANTS_KEY))\n return {tenant_id.decode(\"utf-8\") for tenant_id in gated_tenants_bytes}\n\n\ndef is_tenant_gated(tenant_id: str) -> bool:\n \"\"\"Fast O(1) check if tenant is in gated set (multi-tenant only).\"\"\"\n redis_client = get_redis_replica_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n return bool(redis_client.sismember(GATED_TENANTS_KEY, tenant_id))\n" }, { "path": "backend/ee/onyx/server/tenants/provisioning.py", "content": "import asyncio\nimport uuid\n\nimport aiohttp # Async HTTP client\nimport httpx\nimport requests\nfrom fastapi import HTTPException\nfrom fastapi import Request\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.configs.app_configs import HUBSPOT_TRACKING_URL\nfrom ee.onyx.server.tenants.access import generate_data_plane_token\nfrom ee.onyx.server.tenants.models import TenantByDomainResponse\nfrom ee.onyx.server.tenants.models import TenantCreationPayload\nfrom ee.onyx.server.tenants.models import TenantDeletionPayload\nfrom ee.onyx.server.tenants.schema_management import create_schema_if_not_exists\nfrom ee.onyx.server.tenants.schema_management import drop_schema\nfrom ee.onyx.server.tenants.schema_management import run_alembic_migrations\nfrom ee.onyx.server.tenants.user_mapping import add_users_to_tenant\nfrom ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email\nfrom ee.onyx.server.tenants.user_mapping import user_owns_a_tenant\nfrom onyx.auth.users import exceptions\nfrom onyx.configs.app_configs import ANTHROPIC_DEFAULT_API_KEY\nfrom onyx.configs.app_configs import COHERE_DEFAULT_API_KEY\nfrom onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.configs.app_configs import OPENAI_DEFAULT_API_KEY\nfrom onyx.configs.app_configs import OPENROUTER_DEFAULT_API_KEY\nfrom onyx.configs.app_configs import VERTEXAI_DEFAULT_CREDENTIALS\nfrom onyx.configs.app_configs import VERTEXAI_DEFAULT_LOCATION\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.image_generation import create_default_image_gen_config_from_api_key\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_cloud_embedding_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import AvailableTenant\nfrom onyx.db.models import IndexModelStatus\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.models import UserTenantMapping\nfrom onyx.llm.well_known_providers.auto_update_models import LLMRecommendations\nfrom onyx.llm.well_known_providers.constants import ANTHROPIC_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import OPENAI_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import OPENROUTER_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import VERTEX_CREDENTIALS_FILE_KWARG\nfrom onyx.llm.well_known_providers.constants import VERTEX_LOCATION_KWARG\nfrom onyx.llm.well_known_providers.constants import VERTEXAI_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.llm_provider_options import (\n get_recommendations,\n)\nfrom onyx.llm.well_known_providers.llm_provider_options import (\n model_configurations_for_provider,\n)\nfrom onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom onyx.setup import setup_onyx\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import TENANT_ID_PREFIX\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.enums import EmbeddingProvider\n\n\nlogger = setup_logger()\n\n\nasync def get_or_provision_tenant(\n email: str,\n referral_source: str | None = None,\n request: Request | None = None,\n) -> str:\n \"\"\"\n Get existing tenant ID for an email or create a new tenant if none exists.\n This function should only be called after we have verified we want this user's tenant to exist.\n It returns the tenant ID associated with the email, creating a new tenant if necessary.\n \"\"\"\n # Early return for non-multi-tenant mode\n if not MULTI_TENANT:\n return POSTGRES_DEFAULT_SCHEMA\n\n if referral_source and request:\n await submit_to_hubspot(email, referral_source, request)\n\n # First, check if the user already has a tenant\n tenant_id: str | None = None\n try:\n tenant_id = get_tenant_id_for_email(email)\n return tenant_id\n except exceptions.UserNotExists:\n # User doesn't exist, so we need to create a new tenant or assign an existing one\n pass\n\n try:\n # Try to get a pre-provisioned tenant\n tenant_id = await get_available_tenant()\n\n if tenant_id:\n # Run migrations to ensure the pre-provisioned tenant schema is current.\n # Pool tenants may have been created before a new migration was deployed.\n # Capture as a non-optional local so mypy can type the lambda correctly.\n _tenant_id: str = tenant_id\n loop = asyncio.get_running_loop()\n try:\n await loop.run_in_executor(\n None, lambda: run_alembic_migrations(_tenant_id)\n )\n except Exception:\n # The tenant was already dequeued from the pool — roll it back so\n # it doesn't end up orphaned (schema exists, but not assigned to anyone).\n logger.exception(\n f\"Migration failed for pre-provisioned tenant {_tenant_id}; rolling back\"\n )\n try:\n await rollback_tenant_provisioning(_tenant_id)\n except Exception:\n logger.exception(f\"Failed to rollback orphaned tenant {_tenant_id}\")\n raise\n # If we have a pre-provisioned tenant, assign it to the user\n await assign_tenant_to_user(tenant_id, email, referral_source)\n logger.info(f\"Assigned pre-provisioned tenant {tenant_id} to user {email}\")\n else:\n # If no pre-provisioned tenant is available, create a new one on-demand\n tenant_id = await create_tenant(email, referral_source)\n\n # Notify control plane if we have created / assigned a new tenant\n if not DEV_MODE:\n await notify_control_plane(tenant_id, email, referral_source)\n\n return tenant_id\n\n except Exception as e:\n # If we've encountered an error, log and raise an exception\n error_msg = \"Failed to provision tenant\"\n logger.error(error_msg, exc_info=e)\n raise HTTPException(\n status_code=500,\n detail=\"Failed to provision tenant. Please try again later.\",\n )\n\n\nasync def create_tenant(\n email: str,\n referral_source: str | None = None, # noqa: ARG001\n) -> str:\n \"\"\"\n Create a new tenant on-demand when no pre-provisioned tenants are available.\n This is the fallback method when we can't use a pre-provisioned tenant.\n\n \"\"\"\n tenant_id = TENANT_ID_PREFIX + str(uuid.uuid4())\n logger.info(f\"Creating new tenant {tenant_id} for user {email}\")\n\n try:\n # Provision tenant on data plane\n await provision_tenant(tenant_id, email)\n\n except Exception as e:\n logger.exception(f\"Tenant provisioning failed: {str(e)}\")\n # Attempt to rollback the tenant provisioning\n try:\n await rollback_tenant_provisioning(tenant_id)\n except Exception:\n logger.exception(f\"Failed to rollback tenant provisioning for {tenant_id}\")\n raise HTTPException(status_code=500, detail=\"Failed to provision tenant.\")\n\n return tenant_id\n\n\nasync def provision_tenant(tenant_id: str, email: str) -> None:\n if not MULTI_TENANT:\n raise HTTPException(status_code=403, detail=\"Multi-tenancy is not enabled\")\n\n if user_owns_a_tenant(email):\n raise HTTPException(\n status_code=409, detail=\"User already belongs to an organization\"\n )\n\n logger.debug(f\"Provisioning tenant {tenant_id} for user {email}\")\n\n try:\n # Create the schema for the tenant\n if not create_schema_if_not_exists(tenant_id):\n logger.debug(f\"Created schema for tenant {tenant_id}\")\n else:\n logger.debug(f\"Schema already exists for tenant {tenant_id}\")\n\n # Set up the tenant with all necessary configurations\n await setup_tenant(tenant_id)\n\n # Assign the tenant to the user\n await assign_tenant_to_user(tenant_id, email)\n\n except Exception as e:\n logger.exception(f\"Failed to create tenant {tenant_id}\")\n raise HTTPException(\n status_code=500, detail=f\"Failed to create tenant: {str(e)}\"\n )\n\n\nasync def notify_control_plane(\n tenant_id: str, email: str, referral_source: str | None = None\n) -> None:\n logger.info(\"Fetching billing information\")\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n payload = TenantCreationPayload(\n tenant_id=tenant_id, email=email, referral_source=referral_source\n )\n\n async with aiohttp.ClientSession() as session:\n async with session.post(\n f\"{CONTROL_PLANE_API_BASE_URL}/tenants/create\",\n headers=headers,\n json=payload.model_dump(),\n ) as response:\n if response.status != 200:\n error_text = await response.text()\n logger.error(f\"Control plane tenant creation failed: {error_text}\")\n raise Exception(\n f\"Failed to create tenant on control plane: {error_text}\"\n )\n\n\nasync def rollback_tenant_provisioning(tenant_id: str) -> None:\n \"\"\"\n Logic to rollback tenant provisioning on data plane.\n Handles each step independently to ensure maximum cleanup even if some steps fail.\n \"\"\"\n logger.info(f\"Rolling back tenant provisioning for tenant_id: {tenant_id}\")\n\n # Track if any part of the rollback fails\n rollback_errors = []\n\n # 1. Try to drop the tenant's schema\n try:\n drop_schema(tenant_id)\n logger.info(f\"Successfully dropped schema for tenant {tenant_id}\")\n except Exception as e:\n error_msg = f\"Failed to drop schema for tenant {tenant_id}: {str(e)}\"\n logger.error(error_msg)\n rollback_errors.append(error_msg)\n\n # 2. Try to remove tenant mapping\n try:\n with get_session_with_shared_schema() as db_session:\n db_session.begin()\n try:\n db_session.query(UserTenantMapping).filter(\n UserTenantMapping.tenant_id == tenant_id\n ).delete()\n db_session.commit()\n logger.info(\n f\"Successfully removed user mappings for tenant {tenant_id}\"\n )\n except Exception as e:\n db_session.rollback()\n raise e\n except Exception as e:\n error_msg = f\"Failed to remove user mappings for tenant {tenant_id}: {str(e)}\"\n logger.error(error_msg)\n rollback_errors.append(error_msg)\n\n # 3. If this tenant was in the available tenants table, remove it\n try:\n with get_session_with_shared_schema() as db_session:\n db_session.begin()\n try:\n available_tenant = (\n db_session.query(AvailableTenant)\n .filter(AvailableTenant.tenant_id == tenant_id)\n .first()\n )\n\n if available_tenant:\n db_session.delete(available_tenant)\n db_session.commit()\n logger.info(\n f\"Removed tenant {tenant_id} from available tenants table\"\n )\n except Exception as e:\n db_session.rollback()\n raise e\n except Exception as e:\n error_msg = f\"Failed to remove tenant {tenant_id} from available tenants table: {str(e)}\"\n logger.error(error_msg)\n rollback_errors.append(error_msg)\n\n # Log summary of rollback operation\n if rollback_errors:\n logger.error(f\"Tenant rollback completed with {len(rollback_errors)} errors\")\n else:\n logger.info(f\"Tenant rollback completed successfully for tenant {tenant_id}\")\n\n\ndef _build_model_configuration_upsert_requests(\n provider_name: str,\n recommendations: LLMRecommendations,\n) -> list[ModelConfigurationUpsertRequest]:\n model_configurations = model_configurations_for_provider(\n provider_name, recommendations\n )\n return [\n ModelConfigurationUpsertRequest(\n name=model_configuration.name,\n is_visible=model_configuration.is_visible,\n max_input_tokens=model_configuration.max_input_tokens,\n supports_image_input=model_configuration.supports_image_input,\n )\n for model_configuration in model_configurations\n ]\n\n\ndef configure_default_api_keys(db_session: Session) -> None:\n \"\"\"Configure default LLM providers using recommended-models.json for model selection.\"\"\"\n # Load recommendations from JSON config\n recommendations = get_recommendations()\n\n has_set_default_provider = False\n\n def _upsert(request: LLMProviderUpsertRequest, default_model: str) -> None:\n nonlocal has_set_default_provider\n try:\n existing = fetch_existing_llm_provider(\n name=request.name, db_session=db_session\n )\n if existing:\n request.id = existing.id\n provider = upsert_llm_provider(request, db_session)\n if not has_set_default_provider:\n update_default_provider(provider.id, default_model, db_session)\n has_set_default_provider = True\n except Exception as e:\n logger.error(f\"Failed to configure {request.provider} provider: {e}\")\n\n # Configure OpenAI provider\n if OPENAI_DEFAULT_API_KEY:\n default_model = recommendations.get_default_model(OPENAI_PROVIDER_NAME)\n if default_model is None:\n logger.error(\n f\"No default model found for {OPENAI_PROVIDER_NAME} in recommendations\"\n )\n default_model_name = default_model.name if default_model else \"gpt-5.2\"\n\n openai_provider = LLMProviderUpsertRequest(\n name=\"OpenAI\",\n provider=OPENAI_PROVIDER_NAME,\n api_key=OPENAI_DEFAULT_API_KEY,\n model_configurations=_build_model_configuration_upsert_requests(\n OPENAI_PROVIDER_NAME, recommendations\n ),\n api_key_changed=True,\n is_auto_mode=True,\n )\n _upsert(openai_provider, default_model_name)\n\n # Create default image generation config using the OpenAI API key\n try:\n create_default_image_gen_config_from_api_key(\n db_session, OPENAI_DEFAULT_API_KEY\n )\n except Exception as e:\n logger.error(f\"Failed to create default image gen config: {e}\")\n else:\n logger.info(\n \"OPENAI_DEFAULT_API_KEY not set, skipping OpenAI provider configuration\"\n )\n\n # Configure Anthropic provider\n if ANTHROPIC_DEFAULT_API_KEY:\n default_model = recommendations.get_default_model(ANTHROPIC_PROVIDER_NAME)\n if default_model is None:\n logger.error(\n f\"No default model found for {ANTHROPIC_PROVIDER_NAME} in recommendations\"\n )\n default_model_name = (\n default_model.name if default_model else \"claude-sonnet-4-5\"\n )\n\n anthropic_provider = LLMProviderUpsertRequest(\n name=\"Anthropic\",\n provider=ANTHROPIC_PROVIDER_NAME,\n api_key=ANTHROPIC_DEFAULT_API_KEY,\n model_configurations=_build_model_configuration_upsert_requests(\n ANTHROPIC_PROVIDER_NAME, recommendations\n ),\n api_key_changed=True,\n is_auto_mode=True,\n )\n _upsert(anthropic_provider, default_model_name)\n else:\n logger.info(\n \"ANTHROPIC_DEFAULT_API_KEY not set, skipping Anthropic provider configuration\"\n )\n\n # Configure Vertex AI provider\n if VERTEXAI_DEFAULT_CREDENTIALS:\n default_model = recommendations.get_default_model(VERTEXAI_PROVIDER_NAME)\n if default_model is None:\n logger.error(\n f\"No default model found for {VERTEXAI_PROVIDER_NAME} in recommendations\"\n )\n default_model_name = default_model.name if default_model else \"gemini-2.5-pro\"\n\n # Vertex AI uses custom_config for credentials and location\n custom_config = {\n VERTEX_CREDENTIALS_FILE_KWARG: VERTEXAI_DEFAULT_CREDENTIALS,\n VERTEX_LOCATION_KWARG: VERTEXAI_DEFAULT_LOCATION,\n }\n\n vertexai_provider = LLMProviderUpsertRequest(\n name=\"Google Vertex AI\",\n provider=VERTEXAI_PROVIDER_NAME,\n custom_config=custom_config,\n model_configurations=_build_model_configuration_upsert_requests(\n VERTEXAI_PROVIDER_NAME, recommendations\n ),\n api_key_changed=True,\n is_auto_mode=True,\n )\n _upsert(vertexai_provider, default_model_name)\n else:\n logger.info(\n \"VERTEXAI_DEFAULT_CREDENTIALS not set, skipping Vertex AI provider configuration\"\n )\n\n # Configure OpenRouter provider\n if OPENROUTER_DEFAULT_API_KEY:\n default_model = recommendations.get_default_model(OPENROUTER_PROVIDER_NAME)\n if default_model is None:\n logger.error(\n f\"No default model found for {OPENROUTER_PROVIDER_NAME} in recommendations\"\n )\n default_model_name = default_model.name if default_model else \"z-ai/glm-4.7\"\n\n # For OpenRouter, we use the visible models from recommendations as model_configurations\n # since OpenRouter models are dynamic (fetched from their API)\n visible_models = recommendations.get_visible_models(OPENROUTER_PROVIDER_NAME)\n model_configurations = [\n ModelConfigurationUpsertRequest(\n name=model.name,\n is_visible=True,\n max_input_tokens=None,\n display_name=model.display_name,\n )\n for model in visible_models\n ]\n\n openrouter_provider = LLMProviderUpsertRequest(\n name=\"OpenRouter\",\n provider=OPENROUTER_PROVIDER_NAME,\n api_key=OPENROUTER_DEFAULT_API_KEY,\n model_configurations=model_configurations,\n api_key_changed=True,\n is_auto_mode=True,\n )\n _upsert(openrouter_provider, default_model_name)\n else:\n logger.info(\n \"OPENROUTER_DEFAULT_API_KEY not set, skipping OpenRouter provider configuration\"\n )\n\n # Configure Cohere embedding provider\n if COHERE_DEFAULT_API_KEY:\n cloud_embedding_provider = CloudEmbeddingProviderCreationRequest(\n provider_type=EmbeddingProvider.COHERE,\n api_key=COHERE_DEFAULT_API_KEY,\n )\n\n try:\n logger.info(\"Attempting to upsert Cohere cloud embedding provider\")\n upsert_cloud_embedding_provider(db_session, cloud_embedding_provider)\n logger.info(\"Successfully upserted Cohere cloud embedding provider\")\n\n logger.info(\"Updating search settings with Cohere embedding model details\")\n query = (\n select(SearchSettings)\n .where(SearchSettings.status == IndexModelStatus.FUTURE)\n .order_by(SearchSettings.id.desc())\n )\n result = db_session.execute(query)\n current_search_settings = result.scalars().first()\n\n if current_search_settings:\n current_search_settings.model_name = (\n \"embed-english-v3.0\" # Cohere's latest model as of now\n )\n current_search_settings.model_dim = (\n 1024 # Cohere's embed-english-v3.0 dimension\n )\n current_search_settings.provider_type = EmbeddingProvider.COHERE\n current_search_settings.index_name = (\n \"danswer_chunk_cohere_embed_english_v3_0\"\n )\n current_search_settings.query_prefix = \"\"\n current_search_settings.passage_prefix = \"\"\n db_session.commit()\n else:\n raise RuntimeError(\n \"No search settings specified, DB is not in a valid state\"\n )\n logger.info(\"Fetching updated search settings to verify changes\")\n updated_query = (\n select(SearchSettings)\n .where(SearchSettings.status == IndexModelStatus.PRESENT)\n .order_by(SearchSettings.id.desc())\n )\n updated_result = db_session.execute(updated_query)\n updated_result.scalars().first()\n\n except Exception:\n logger.exception(\"Failed to configure Cohere embedding provider\")\n else:\n logger.info(\n \"COHERE_DEFAULT_API_KEY not set, skipping Cohere embedding provider configuration\"\n )\n\n\nasync def submit_to_hubspot(\n email: str, referral_source: str | None, request: Request\n) -> None:\n if not HUBSPOT_TRACKING_URL:\n logger.info(\"HUBSPOT_TRACKING_URL not set, skipping HubSpot submission\")\n return\n\n # HubSpot tracking cookie\n hubspot_cookie = request.cookies.get(\"hubspotutk\")\n\n # IP address\n ip_address = request.client.host if request.client else None\n\n data = {\n \"fields\": [\n {\"name\": \"email\", \"value\": email},\n {\"name\": \"referral_source\", \"value\": referral_source or \"\"},\n ],\n \"context\": {\n \"hutk\": hubspot_cookie,\n \"ipAddress\": ip_address,\n \"pageUri\": str(request.url),\n \"pageName\": \"User Registration\",\n },\n }\n\n async with httpx.AsyncClient() as client:\n response = await client.post(HUBSPOT_TRACKING_URL, json=data)\n\n if response.status_code != 200:\n logger.error(f\"Failed to submit to HubSpot: {response.text}\")\n\n\nasync def delete_user_from_control_plane(tenant_id: str, email: str) -> None:\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n payload = TenantDeletionPayload(tenant_id=tenant_id, email=email)\n\n async with aiohttp.ClientSession() as session:\n async with session.delete(\n f\"{CONTROL_PLANE_API_BASE_URL}/tenants/delete\",\n headers=headers,\n json=payload.model_dump(),\n ) as response:\n if response.status != 200:\n error_text = await response.text()\n logger.error(f\"Control plane tenant creation failed: {error_text}\")\n raise Exception(\n f\"Failed to delete tenant on control plane: {error_text}\"\n )\n\n\ndef get_tenant_by_domain_from_control_plane(\n domain: str,\n tenant_id: str,\n) -> TenantByDomainResponse | None:\n \"\"\"\n Fetches tenant information from the control plane based on the email domain.\n\n Args:\n domain: The email domain to search for (e.g., \"example.com\")\n\n Returns:\n A dictionary containing tenant information if found, None otherwise\n \"\"\"\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n\n try:\n response = requests.get(\n f\"{CONTROL_PLANE_API_BASE_URL}/tenant-by-domain\",\n headers=headers,\n json={\"domain\": domain, \"tenant_id\": tenant_id},\n )\n\n if response.status_code != 200:\n logger.error(f\"Control plane tenant lookup failed: {response.text}\")\n return None\n\n response_data = response.json()\n if not response_data:\n return None\n\n return TenantByDomainResponse(\n tenant_id=response_data.get(\"tenant_id\"),\n number_of_users=response_data.get(\"number_of_users\"),\n creator_email=response_data.get(\"creator_email\"),\n )\n except Exception as e:\n logger.error(f\"Error fetching tenant by domain: {str(e)}\")\n return None\n\n\nasync def get_available_tenant() -> str | None:\n \"\"\"\n Get an available pre-provisioned tenant from the NewAvailableTenant table.\n Returns the tenant_id if one is available, None otherwise.\n Uses row-level locking to prevent race conditions when multiple processes\n try to get an available tenant simultaneously.\n \"\"\"\n if not MULTI_TENANT:\n return None\n\n with get_session_with_shared_schema() as db_session:\n try:\n db_session.begin()\n\n # Get the oldest available tenant with FOR UPDATE lock to prevent race conditions\n available_tenant = (\n db_session.query(AvailableTenant)\n .order_by(AvailableTenant.date_created)\n .with_for_update(skip_locked=True) # Skip locked rows to avoid blocking\n .first()\n )\n\n if available_tenant:\n tenant_id = available_tenant.tenant_id\n # Remove the tenant from the available tenants table\n db_session.delete(available_tenant)\n db_session.commit()\n logger.info(f\"Using pre-provisioned tenant {tenant_id}\")\n return tenant_id\n else:\n db_session.rollback()\n return None\n except Exception:\n logger.exception(\"Error getting available tenant\")\n db_session.rollback()\n return None\n\n\nasync def setup_tenant(tenant_id: str) -> None:\n \"\"\"\n Set up a tenant with all necessary configurations.\n This is a centralized function that handles all tenant setup logic.\n \"\"\"\n token = None\n try:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n # Run Alembic migrations in a way that isolates it from the current event loop\n # Create a new event loop for this synchronous operation\n loop = asyncio.get_event_loop()\n # Use run_in_executor which properly isolates the thread execution\n await loop.run_in_executor(None, lambda: run_alembic_migrations(tenant_id))\n\n # Configure the tenant with default settings\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # Configure default API keys\n configure_default_api_keys(db_session)\n\n # Set up Onyx with appropriate settings\n current_search_settings = (\n db_session.query(SearchSettings)\n .filter_by(status=IndexModelStatus.FUTURE)\n .first()\n )\n cohere_enabled = (\n current_search_settings is not None\n and current_search_settings.provider_type == EmbeddingProvider.COHERE\n )\n setup_onyx(db_session, tenant_id, cohere_enabled=cohere_enabled)\n\n except Exception as e:\n logger.exception(f\"Failed to set up tenant {tenant_id}\")\n raise e\n finally:\n if token is not None:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\nasync def assign_tenant_to_user(\n tenant_id: str,\n email: str,\n referral_source: str | None = None, # noqa: ARG001\n) -> None:\n \"\"\"\n Assign a tenant to a user and perform necessary operations.\n Uses transaction handling to ensure atomicity and includes retry logic\n for control plane notifications.\n \"\"\"\n # First, add the user to the tenant in a transaction\n\n try:\n add_users_to_tenant([email], tenant_id)\n except Exception:\n logger.exception(f\"Failed to assign tenant {tenant_id} to user {email}\")\n raise Exception(\"Failed to assign tenant to user\")\n" }, { "path": "backend/ee/onyx/server/tenants/proxy.py", "content": "\"\"\"Proxy endpoints for billing operations.\n\nThese endpoints run on the CLOUD DATA PLANE (cloud.onyx.app) and serve as a proxy\nfor self-hosted instances to reach the control plane.\n\nFlow:\n Self-hosted backend → Cloud DP /proxy/* (license auth) → Control plane (JWT auth)\n\nSelf-hosted instances call these endpoints with their license in the Authorization\nheader. The cloud data plane validates the license signature and forwards the\nrequest to the control plane using JWT authentication.\n\nAuth levels by endpoint:\n- /create-checkout-session: No auth (new customer) or expired license OK (renewal)\n- /claim-license: Session ID based (one-time after Stripe payment)\n- /create-customer-portal-session: Expired license OK (need portal to fix payment)\n- /billing-information: Valid license required\n- /license/{tenant_id}: Valid license required\n- /seats/update: Valid license required\n\"\"\"\n\nfrom typing import Literal\n\nimport httpx\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import Header\nfrom fastapi import HTTPException\nfrom pydantic import BaseModel\n\nfrom ee.onyx.configs.app_configs import LICENSE_ENFORCEMENT_ENABLED\nfrom ee.onyx.server.billing.models import SeatUpdateRequest\nfrom ee.onyx.server.billing.models import SeatUpdateResponse\nfrom ee.onyx.server.license.models import LicensePayload\nfrom ee.onyx.server.tenants.access import generate_data_plane_token\nfrom ee.onyx.utils.license import is_license_valid\nfrom ee.onyx.utils.license import verify_license_signature\nfrom onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/proxy\")\n\n\ndef _check_license_enforcement_enabled() -> None:\n \"\"\"Ensure LICENSE_ENFORCEMENT_ENABLED is true (proxy endpoints only work on cloud DP).\"\"\"\n if not LICENSE_ENFORCEMENT_ENABLED:\n raise HTTPException(\n status_code=501,\n detail=\"Proxy endpoints are only available on cloud data plane\",\n )\n\n\ndef _extract_license_from_header(\n authorization: str | None,\n required: bool = True,\n) -> str | None:\n \"\"\"Extract license data from Authorization header.\n\n Self-hosted instances authenticate to these proxy endpoints by sending their\n license as a Bearer token: `Authorization: Bearer `.\n\n We use the Bearer scheme (RFC 6750) because:\n 1. It's the standard HTTP auth scheme for token-based authentication\n 2. The license blob is cryptographically signed (RSA), so it's self-validating\n 3. No other auth schemes (Basic, Digest, etc.) are supported for license auth\n\n The license data is the base64-encoded signed blob that contains tenant_id,\n seats, expiration, etc. We verify the signature to authenticate the caller.\n\n Args:\n authorization: The Authorization header value (e.g., \"Bearer \")\n required: If True, raise 401 when header is missing/invalid\n\n Returns:\n License data string (base64-encoded), or None if not required and missing\n\n Raises:\n HTTPException: 401 if required and header is missing/invalid\n \"\"\"\n if not authorization or not authorization.startswith(\"Bearer \"):\n if required:\n raise HTTPException(\n status_code=401, detail=\"Missing or invalid authorization header\"\n )\n return None\n\n return authorization.split(\" \", 1)[1]\n\n\ndef verify_license_auth(\n license_data: str,\n allow_expired: bool = False,\n) -> LicensePayload:\n \"\"\"Verify license signature and optionally check expiry.\n\n Args:\n license_data: Base64-encoded signed license blob\n allow_expired: If True, accept expired licenses (for renewal flows)\n\n Returns:\n LicensePayload if valid\n\n Raises:\n HTTPException: If license is invalid or expired (when not allowed)\n \"\"\"\n _check_license_enforcement_enabled()\n\n try:\n payload = verify_license_signature(license_data)\n except ValueError as e:\n raise HTTPException(status_code=401, detail=f\"Invalid license: {e}\")\n\n if not allow_expired and not is_license_valid(payload):\n raise HTTPException(status_code=401, detail=\"License has expired\")\n\n return payload\n\n\nasync def get_license_payload(\n authorization: str | None = Header(None, alias=\"Authorization\"),\n) -> LicensePayload:\n \"\"\"Dependency: Require valid (non-expired) license.\n\n Used for endpoints that require an active subscription.\n \"\"\"\n license_data = _extract_license_from_header(authorization, required=True)\n # license_data is guaranteed non-None when required=True\n assert license_data is not None\n return verify_license_auth(license_data, allow_expired=False)\n\n\nasync def get_license_payload_allow_expired(\n authorization: str | None = Header(None, alias=\"Authorization\"),\n) -> LicensePayload:\n \"\"\"Dependency: Require license with valid signature, expired OK.\n\n Used for endpoints needed to fix payment issues (portal, renewal checkout).\n \"\"\"\n license_data = _extract_license_from_header(authorization, required=True)\n # license_data is guaranteed non-None when required=True\n assert license_data is not None\n return verify_license_auth(license_data, allow_expired=True)\n\n\nasync def get_optional_license_payload(\n authorization: str | None = Header(None, alias=\"Authorization\"),\n) -> LicensePayload | None:\n \"\"\"Dependency: Optional license auth (for checkout - new customers have none).\n\n Returns None if no license provided, otherwise validates and returns payload.\n Expired licenses are allowed for renewal flows.\n \"\"\"\n _check_license_enforcement_enabled()\n\n license_data = _extract_license_from_header(authorization, required=False)\n if license_data is None:\n return None\n\n return verify_license_auth(license_data, allow_expired=True)\n\n\nasync def forward_to_control_plane(\n method: str,\n path: str,\n body: dict | None = None,\n params: dict | None = None,\n) -> dict:\n \"\"\"Forward a request to the control plane with proper authentication.\"\"\"\n token = generate_data_plane_token()\n headers = {\n \"Authorization\": f\"Bearer {token}\",\n \"Content-Type\": \"application/json\",\n }\n\n url = f\"{CONTROL_PLANE_API_BASE_URL}{path}\"\n\n try:\n async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:\n if method == \"GET\":\n response = await client.get(url, headers=headers, params=params)\n elif method == \"POST\":\n response = await client.post(url, headers=headers, json=body)\n else:\n raise ValueError(f\"Unsupported HTTP method: {method}\")\n\n response.raise_for_status()\n return response.json()\n\n except httpx.HTTPStatusError as e:\n status_code = e.response.status_code\n detail = \"Control plane request failed\"\n try:\n error_data = e.response.json()\n detail = error_data.get(\"detail\", detail)\n except Exception:\n pass\n logger.error(f\"Control plane returned {status_code}: {detail}\")\n raise HTTPException(status_code=status_code, detail=detail)\n except httpx.RequestError:\n logger.exception(\"Failed to connect to control plane\")\n raise HTTPException(\n status_code=502, detail=\"Failed to connect to control plane\"\n )\n\n\n# -----------------------------------------------------------------------------\n# Endpoints\n# -----------------------------------------------------------------------------\n\n\nclass CreateCheckoutSessionRequest(BaseModel):\n billing_period: Literal[\"monthly\", \"annual\"] = \"monthly\"\n seats: int | None = None\n email: str | None = None\n # Redirect URL after successful checkout - self-hosted passes their instance URL\n redirect_url: str | None = None\n # Cancel URL when user exits checkout - returns to upgrade page\n cancel_url: str | None = None\n\n\nclass CreateCheckoutSessionResponse(BaseModel):\n url: str\n\n\n@router.post(\"/create-checkout-session\")\nasync def proxy_create_checkout_session(\n request_body: CreateCheckoutSessionRequest,\n license_payload: LicensePayload | None = Depends(get_optional_license_payload),\n) -> CreateCheckoutSessionResponse:\n \"\"\"Proxy checkout session creation to control plane.\n\n Auth: Optional license (new customers don't have one yet).\n If license provided, expired is OK (for renewals).\n \"\"\"\n # license_payload is None for new customers who don't have a license yet.\n # In that case, tenant_id is omitted from the request body and the control\n # plane will create a new tenant during checkout completion.\n tenant_id = license_payload.tenant_id if license_payload else None\n\n body: dict = {\n \"billing_period\": request_body.billing_period,\n }\n if tenant_id:\n body[\"tenant_id\"] = tenant_id\n if request_body.seats is not None:\n body[\"seats\"] = request_body.seats\n if request_body.email:\n body[\"email\"] = request_body.email\n if request_body.redirect_url:\n body[\"redirect_url\"] = request_body.redirect_url\n if request_body.cancel_url:\n body[\"cancel_url\"] = request_body.cancel_url\n\n result = await forward_to_control_plane(\n \"POST\", \"/create-checkout-session\", body=body\n )\n return CreateCheckoutSessionResponse(url=result[\"url\"])\n\n\nclass ClaimLicenseRequest(BaseModel):\n session_id: str\n\n\nclass ClaimLicenseResponse(BaseModel):\n tenant_id: str\n license: str\n message: str | None = None\n\n\n@router.post(\"/claim-license\")\nasync def proxy_claim_license(\n request_body: ClaimLicenseRequest,\n) -> ClaimLicenseResponse:\n \"\"\"Claim a license after successful Stripe checkout.\n\n Auth: Session ID based (one-time use after payment).\n The control plane verifies the session_id is valid and unclaimed.\n\n Returns the license to the caller. For self-hosted instances, they will\n store the license locally. The cloud DP doesn't need to store it.\n \"\"\"\n _check_license_enforcement_enabled()\n\n result = await forward_to_control_plane(\n \"POST\",\n \"/claim-license\",\n body={\"session_id\": request_body.session_id},\n )\n\n tenant_id = result.get(\"tenant_id\")\n license_data = result.get(\"license\")\n\n if not tenant_id or not license_data:\n logger.error(f\"Control plane returned incomplete claim response: {result}\")\n raise HTTPException(\n status_code=502,\n detail=\"Control plane returned incomplete license data\",\n )\n\n return ClaimLicenseResponse(\n tenant_id=tenant_id,\n license=license_data,\n message=\"License claimed successfully\",\n )\n\n\nclass CreateCustomerPortalSessionRequest(BaseModel):\n return_url: str | None = None\n\n\nclass CreateCustomerPortalSessionResponse(BaseModel):\n url: str\n\n\n@router.post(\"/create-customer-portal-session\")\nasync def proxy_create_customer_portal_session(\n request_body: CreateCustomerPortalSessionRequest | None = None,\n license_payload: LicensePayload = Depends(get_license_payload_allow_expired),\n) -> CreateCustomerPortalSessionResponse:\n \"\"\"Proxy customer portal session creation to control plane.\n\n Auth: License required, expired OK (need portal to fix payment issues).\n \"\"\"\n # tenant_id is a required field in LicensePayload (Pydantic validates this),\n # but we check explicitly for defense in depth\n if not license_payload.tenant_id:\n raise HTTPException(status_code=401, detail=\"License missing tenant_id\")\n\n tenant_id = license_payload.tenant_id\n\n body: dict = {\"tenant_id\": tenant_id}\n if request_body and request_body.return_url:\n body[\"return_url\"] = request_body.return_url\n\n result = await forward_to_control_plane(\n \"POST\", \"/create-customer-portal-session\", body=body\n )\n return CreateCustomerPortalSessionResponse(url=result[\"url\"])\n\n\nclass BillingInformationResponse(BaseModel):\n tenant_id: str\n status: str | None = None\n plan_type: str | None = None\n seats: int | None = None\n billing_period: str | None = None\n current_period_start: str | None = None\n current_period_end: str | None = None\n cancel_at_period_end: bool = False\n canceled_at: str | None = None\n trial_start: str | None = None\n trial_end: str | None = None\n payment_method_enabled: bool = False\n stripe_subscription_id: str | None = None\n\n\n@router.get(\"/billing-information\")\nasync def proxy_billing_information(\n license_payload: LicensePayload = Depends(get_license_payload),\n) -> BillingInformationResponse:\n \"\"\"Proxy billing information request to control plane.\n\n Auth: Valid (non-expired) license required.\n \"\"\"\n # tenant_id is a required field in LicensePayload (Pydantic validates this),\n # but we check explicitly for defense in depth\n if not license_payload.tenant_id:\n raise HTTPException(status_code=401, detail=\"License missing tenant_id\")\n\n tenant_id = license_payload.tenant_id\n\n result = await forward_to_control_plane(\n \"GET\", \"/billing-information\", params={\"tenant_id\": tenant_id}\n )\n # Add tenant_id from license if not in response (control plane may not include it)\n if \"tenant_id\" not in result:\n result[\"tenant_id\"] = tenant_id\n return BillingInformationResponse(**result)\n\n\nclass LicenseFetchResponse(BaseModel):\n license: str\n tenant_id: str\n\n\n@router.get(\"/license/{tenant_id}\")\nasync def proxy_license_fetch(\n tenant_id: str,\n license_payload: LicensePayload = Depends(get_license_payload),\n) -> LicenseFetchResponse:\n \"\"\"Proxy license fetch to control plane.\n\n Auth: Valid license required.\n The tenant_id in path must match the authenticated tenant.\n \"\"\"\n # tenant_id is a required field in LicensePayload (Pydantic validates this),\n # but we check explicitly for defense in depth\n if not license_payload.tenant_id:\n raise HTTPException(status_code=401, detail=\"License missing tenant_id\")\n\n if tenant_id != license_payload.tenant_id:\n raise HTTPException(\n status_code=403,\n detail=\"Cannot fetch license for a different tenant\",\n )\n\n result = await forward_to_control_plane(\"GET\", f\"/license/{tenant_id}\")\n\n license_data = result.get(\"license\")\n if not license_data:\n logger.error(f\"Control plane returned incomplete license response: {result}\")\n raise HTTPException(\n status_code=502,\n detail=\"Control plane returned incomplete license data\",\n )\n\n # Return license to caller - self-hosted instance stores it via /api/license/claim\n return LicenseFetchResponse(license=license_data, tenant_id=tenant_id)\n\n\n@router.post(\"/seats/update\")\nasync def proxy_seat_update(\n request_body: SeatUpdateRequest,\n license_payload: LicensePayload = Depends(get_license_payload),\n) -> SeatUpdateResponse:\n \"\"\"Proxy seat update to control plane.\n\n Auth: Valid (non-expired) license required.\n Handles Stripe proration and license regeneration.\n Returns the regenerated license in the response for the caller to store.\n \"\"\"\n if not license_payload.tenant_id:\n raise HTTPException(status_code=401, detail=\"License missing tenant_id\")\n\n tenant_id = license_payload.tenant_id\n\n result = await forward_to_control_plane(\n \"POST\",\n \"/seats/update\",\n body={\n \"tenant_id\": tenant_id,\n \"new_seat_count\": request_body.new_seat_count,\n },\n )\n\n # Return license in response - self-hosted instance stores it via /api/license/claim\n return SeatUpdateResponse(\n success=result.get(\"success\", False),\n current_seats=result.get(\"current_seats\", 0),\n used_seats=result.get(\"used_seats\", 0),\n message=result.get(\"message\"),\n license=result.get(\"license\"),\n )\n" }, { "path": "backend/ee/onyx/server/tenants/schema_management.py", "content": "import logging\nimport os\nimport re\nfrom types import SimpleNamespace\n\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.schema import CreateSchema\n\nfrom alembic import command\nfrom alembic.config import Config\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.engine.sql_engine import get_sqlalchemy_engine\nfrom shared_configs.configs import TENANT_ID_PREFIX\n\nlogger = logging.getLogger(__name__)\n\n# Regex pattern for valid tenant IDs:\n# - UUID format: tenant_xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx\n# - AWS instance ID format: tenant_i-xxxxxxxxxxxxxxxxx\n# Also useful for not accidentally dropping `public` schema\nTENANT_ID_PATTERN = re.compile(\n rf\"^{re.escape(TENANT_ID_PREFIX)}(\"\n r\"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}\" # UUID\n r\"|i-[a-f0-9]+\" # AWS instance ID\n r\")$\"\n)\n\n\ndef validate_tenant_id(tenant_id: str) -> bool:\n \"\"\"Validate that tenant_id matches expected format.\n\n This is important for SQL injection prevention since schema names\n cannot be parameterized in SQL and must be formatted directly.\n \"\"\"\n return bool(TENANT_ID_PATTERN.match(tenant_id))\n\n\ndef run_alembic_migrations(schema_name: str) -> None:\n logger.info(f\"Starting Alembic migrations for schema: {schema_name}\")\n\n try:\n current_dir = os.path.dirname(os.path.abspath(__file__))\n root_dir = os.path.abspath(os.path.join(current_dir, \"..\", \"..\", \"..\", \"..\"))\n alembic_ini_path = os.path.join(root_dir, \"alembic.ini\")\n\n # Configure Alembic\n alembic_cfg = Config(alembic_ini_path)\n alembic_cfg.set_main_option(\"sqlalchemy.url\", build_connection_string())\n alembic_cfg.set_main_option(\n \"script_location\", os.path.join(root_dir, \"alembic\")\n )\n\n # Ensure that logging isn't broken\n alembic_cfg.attributes[\"configure_logger\"] = False\n\n # Mimic command-line options by adding 'cmd_opts' to the config\n alembic_cfg.cmd_opts = SimpleNamespace() # type: ignore\n alembic_cfg.cmd_opts.x = [f\"schemas={schema_name}\"] # type: ignore\n\n # Run migrations programmatically\n command.upgrade(alembic_cfg, \"head\")\n\n # Run migrations programmatically\n logger.info(\n f\"Alembic migrations completed successfully for schema: {schema_name}\"\n )\n\n except Exception as e:\n logger.exception(f\"Alembic migration failed for schema {schema_name}: {str(e)}\")\n raise\n\n\ndef create_schema_if_not_exists(tenant_id: str) -> bool:\n with Session(get_sqlalchemy_engine()) as db_session:\n with db_session.begin():\n result = db_session.execute(\n text(\n \"SELECT schema_name FROM information_schema.schemata WHERE schema_name = :schema_name\"\n ),\n {\"schema_name\": tenant_id},\n )\n schema_exists = result.scalar() is not None\n if not schema_exists:\n stmt = CreateSchema(tenant_id)\n db_session.execute(stmt)\n return True\n return False\n\n\ndef drop_schema(tenant_id: str) -> None:\n \"\"\"Drop a tenant's schema.\n\n Uses strict regex validation to reject unexpected formats early,\n preventing SQL injection since schema names cannot be parameterized.\n \"\"\"\n if not validate_tenant_id(tenant_id):\n raise ValueError(f\"Invalid tenant_id format: {tenant_id}\")\n\n with get_sqlalchemy_engine().connect() as connection:\n with connection.begin():\n # Use string formatting with validated tenant_id (safe after validation)\n connection.execute(text(f'DROP SCHEMA IF EXISTS \"{tenant_id}\" CASCADE'))\n\n\ndef get_current_alembic_version(tenant_id: str) -> str:\n \"\"\"Get the current Alembic version for a tenant.\"\"\"\n from alembic.runtime.migration import MigrationContext\n from sqlalchemy import text\n\n engine = get_sqlalchemy_engine()\n\n # Set the search path to the tenant's schema\n with engine.connect() as connection:\n connection.execute(text(f'SET search_path TO \"{tenant_id}\"'))\n\n # Get the current version from the alembic_version table\n context = MigrationContext.configure(connection)\n current_rev = context.get_current_revision()\n\n return current_rev or \"head\"\n" }, { "path": "backend/ee/onyx/server/tenants/team_membership_api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.server.tenants.provisioning import delete_user_from_control_plane\nfrom ee.onyx.server.tenants.user_mapping import remove_all_users_from_tenant\nfrom ee.onyx.server.tenants.user_mapping import remove_users_from_tenant\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import User\nfrom onyx.db.auth import get_user_count\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.users import delete_user_from_db\nfrom onyx.db.users import get_user_by_email\nfrom onyx.server.manage.models import UserByEmail\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/tenants\")\n\n\n@router.post(\"/leave-team\")\nasync def leave_organization(\n user_email: UserByEmail,\n current_user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> None:\n tenant_id = get_current_tenant_id()\n\n if current_user.email != user_email.user_email:\n raise HTTPException(\n status_code=403, detail=\"You can only leave the organization as yourself\"\n )\n\n user_to_delete = get_user_by_email(user_email.user_email, db_session)\n if user_to_delete is None:\n raise HTTPException(status_code=404, detail=\"User not found\")\n\n num_admin_users = await get_user_count(only_admin_users=True)\n\n should_delete_tenant = num_admin_users == 1\n\n if should_delete_tenant:\n logger.info(\n \"Last admin user is leaving the organization. Deleting tenant from control plane.\"\n )\n try:\n await delete_user_from_control_plane(tenant_id, user_to_delete.email)\n logger.debug(\"User deleted from control plane\")\n except Exception as e:\n logger.exception(\n f\"Failed to delete user from control plane for tenant {tenant_id}: {e}\"\n )\n raise HTTPException(\n status_code=500,\n detail=f\"Failed to remove user from control plane: {str(e)}\",\n )\n\n db_session.expunge(user_to_delete)\n delete_user_from_db(user_to_delete, db_session)\n\n if should_delete_tenant:\n remove_all_users_from_tenant(tenant_id)\n else:\n remove_users_from_tenant([user_to_delete.email], tenant_id)\n" }, { "path": "backend/ee/onyx/server/tenants/tenant_management_api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\n\nfrom ee.onyx.server.tenants.models import TenantByDomainResponse\nfrom ee.onyx.server.tenants.provisioning import get_tenant_by_domain_from_control_plane\nfrom onyx.auth.users import current_user\nfrom onyx.auth.users import User\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/tenants\")\n\nFORBIDDEN_COMMON_EMAIL_SUBSTRINGS = [\n \"gmail\",\n \"outlook\",\n \"yahoo\",\n \"hotmail\",\n \"icloud\",\n \"msn\",\n \"hotmail\",\n \"hotmail.co.uk\",\n]\n\n\n@router.get(\"/existing-team-by-domain\")\ndef get_existing_tenant_by_domain(\n user: User = Depends(current_user),\n) -> TenantByDomainResponse | None:\n domain = user.email.split(\"@\")[1]\n if any(substring in domain for substring in FORBIDDEN_COMMON_EMAIL_SUBSTRINGS):\n return None\n\n tenant_id = get_current_tenant_id()\n\n return get_tenant_by_domain_from_control_plane(domain, tenant_id)\n" }, { "path": "backend/ee/onyx/server/tenants/user_invitations_api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\n\nfrom ee.onyx.server.tenants.models import ApproveUserRequest\nfrom ee.onyx.server.tenants.models import PendingUserSnapshot\nfrom ee.onyx.server.tenants.models import RequestInviteRequest\nfrom ee.onyx.server.tenants.user_mapping import accept_user_invite\nfrom ee.onyx.server.tenants.user_mapping import approve_user_invite\nfrom ee.onyx.server.tenants.user_mapping import deny_user_invite\nfrom ee.onyx.server.tenants.user_mapping import invite_self_to_tenant\nfrom onyx.auth.invited_users import get_pending_users\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_user\nfrom onyx.auth.users import User\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/tenants\")\n\n\n@router.post(\"/users/invite/request\")\nasync def request_invite(\n invite_request: RequestInviteRequest,\n user: User = Depends(current_admin_user),\n) -> None:\n try:\n invite_self_to_tenant(user.email, invite_request.tenant_id)\n except Exception as e:\n logger.exception(\n f\"Failed to invite self to tenant {invite_request.tenant_id}: {e}\"\n )\n raise HTTPException(status_code=500, detail=str(e))\n\n\n@router.get(\"/users/pending\")\ndef list_pending_users(\n _: User = Depends(current_admin_user),\n) -> list[PendingUserSnapshot]:\n pending_emails = get_pending_users()\n return [PendingUserSnapshot(email=email) for email in pending_emails]\n\n\n@router.post(\"/users/invite/approve\")\nasync def approve_user(\n approve_user_request: ApproveUserRequest,\n _: User = Depends(current_admin_user),\n) -> None:\n tenant_id = get_current_tenant_id()\n approve_user_invite(approve_user_request.email, tenant_id)\n\n\n@router.post(\"/users/invite/accept\")\nasync def accept_invite(\n invite_request: RequestInviteRequest,\n user: User = Depends(current_user),\n) -> None:\n \"\"\"\n Accept an invitation to join a tenant.\n \"\"\"\n try:\n accept_user_invite(user.email, invite_request.tenant_id)\n except Exception as e:\n logger.exception(f\"Failed to accept invite: {str(e)}\")\n raise HTTPException(status_code=500, detail=\"Failed to accept invitation\")\n\n\n@router.post(\"/users/invite/deny\")\nasync def deny_invite(\n invite_request: RequestInviteRequest,\n user: User = Depends(current_user),\n) -> None:\n \"\"\"\n Deny an invitation to join a tenant.\n \"\"\"\n try:\n deny_user_invite(user.email, invite_request.tenant_id)\n except Exception as e:\n logger.exception(f\"Failed to deny invite: {str(e)}\")\n raise HTTPException(status_code=500, detail=\"Failed to deny invitation\")\n" }, { "path": "backend/ee/onyx/server/tenants/user_mapping.py", "content": "from fastapi_users import exceptions\nfrom sqlalchemy import select\n\nfrom onyx.auth.invited_users import get_invited_users\nfrom onyx.auth.invited_users import get_pending_users\nfrom onyx.auth.invited_users import write_invited_users\nfrom onyx.auth.invited_users import write_pending_users\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.models import UserTenantMapping\nfrom onyx.server.manage.models import TenantSnapshot\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\n\ndef get_tenant_id_for_email(email: str) -> str:\n if not MULTI_TENANT:\n return POSTGRES_DEFAULT_SCHEMA\n # Implement logic to get tenant_id from the mapping table\n try:\n with get_session_with_shared_schema() as db_session:\n # First try to get an active tenant\n result = db_session.execute(\n select(UserTenantMapping).where(\n UserTenantMapping.email == email,\n UserTenantMapping.active == True, # noqa: E712\n )\n )\n mapping = result.scalar_one_or_none()\n tenant_id = mapping.tenant_id if mapping else None\n\n # If no active tenant found, try to get the first inactive one\n if tenant_id is None:\n result = db_session.execute(\n select(UserTenantMapping).where(\n UserTenantMapping.email == email,\n UserTenantMapping.active == False, # noqa: E712\n )\n )\n mapping = result.scalar_one_or_none()\n if mapping:\n # Mark this mapping as active\n mapping.active = True\n db_session.commit()\n tenant_id = mapping.tenant_id\n except Exception as e:\n logger.exception(f\"Error getting tenant id for email {email}: {e}\")\n raise exceptions.UserNotExists()\n\n if tenant_id is None:\n raise exceptions.UserNotExists()\n return tenant_id\n\n\ndef user_owns_a_tenant(email: str) -> bool:\n with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session:\n result = (\n db_session.query(UserTenantMapping)\n .filter(UserTenantMapping.email == email)\n .first()\n )\n return result is not None\n\n\ndef add_users_to_tenant(emails: list[str], tenant_id: str) -> None:\n \"\"\"\n Add users to a tenant with proper transaction handling.\n Checks if users already have a tenant mapping to avoid duplicates.\n\n If a user already has an active mapping to a different tenant, they receive\n an inactive mapping (invitation) to this tenant. They can accept the\n invitation later to switch tenants.\n\n \"\"\"\n unique_emails = set(emails)\n if not unique_emails:\n return\n\n with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session:\n try:\n # Start a transaction\n db_session.begin()\n\n # Batch query 1: Get all existing mappings for these emails to this tenant\n # Lock rows to prevent concurrent modifications\n existing_mappings = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email.in_(unique_emails),\n UserTenantMapping.tenant_id == tenant_id,\n )\n .with_for_update()\n .all()\n )\n emails_with_mapping = {m.email for m in existing_mappings}\n\n # Batch query 2: Get all active mappings for these emails (any tenant)\n active_mappings = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email.in_(unique_emails),\n UserTenantMapping.active == True, # noqa: E712\n )\n .all()\n )\n emails_with_active_mapping = {m.email for m in active_mappings}\n\n # Add mappings for emails that don't already have one to this tenant\n for email in unique_emails:\n if email in emails_with_mapping:\n continue\n\n # Create mapping: inactive if user belongs to another tenant (invitation),\n # active otherwise\n db_session.add(\n UserTenantMapping(\n email=email,\n tenant_id=tenant_id,\n active=email not in emails_with_active_mapping,\n )\n )\n\n # Commit the transaction\n db_session.commit()\n logger.info(f\"Successfully added users {emails} to tenant {tenant_id}\")\n\n except Exception:\n logger.exception(f\"Failed to add users to tenant {tenant_id}\")\n db_session.rollback()\n raise\n\n\ndef remove_users_from_tenant(emails: list[str], tenant_id: str) -> None:\n with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session:\n try:\n mappings_to_delete = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email.in_(emails),\n UserTenantMapping.tenant_id == tenant_id,\n )\n .all()\n )\n\n for mapping in mappings_to_delete:\n db_session.delete(mapping)\n\n db_session.commit()\n except Exception as e:\n logger.exception(\n f\"Failed to remove users from tenant {tenant_id}: {str(e)}\"\n )\n db_session.rollback()\n\n\ndef remove_all_users_from_tenant(tenant_id: str) -> None:\n with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session:\n db_session.query(UserTenantMapping).filter(\n UserTenantMapping.tenant_id == tenant_id\n ).delete()\n db_session.commit()\n\n\ndef invite_self_to_tenant(email: str, tenant_id: str) -> None:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n pending_users = get_pending_users()\n if email in pending_users:\n return\n write_pending_users(pending_users + [email])\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\ndef approve_user_invite(email: str, tenant_id: str) -> None:\n \"\"\"\n Approve a user invite to a tenant.\n This will delete all existing records for this email and create a new mapping entry for the user in this tenant.\n \"\"\"\n with get_session_with_shared_schema() as db_session:\n # Delete all existing records for this email\n db_session.query(UserTenantMapping).filter(\n UserTenantMapping.email == email\n ).delete()\n\n # Create a new mapping entry for the user in this tenant\n new_mapping = UserTenantMapping(email=email, tenant_id=tenant_id, active=True)\n db_session.add(new_mapping)\n db_session.commit()\n\n # Also remove the user from pending users list\n # Remove from pending users\n pending_users = get_pending_users()\n if email in pending_users:\n pending_users.remove(email)\n write_pending_users(pending_users)\n\n # Add to invited users\n invited_users = get_invited_users()\n if email not in invited_users:\n invited_users.append(email)\n write_invited_users(invited_users)\n\n\ndef accept_user_invite(email: str, tenant_id: str) -> None:\n \"\"\"\n Accept an invitation to join a tenant.\n This activates the user's mapping to the tenant.\n \"\"\"\n with get_session_with_shared_schema() as db_session:\n try:\n # Lock the user's mappings first to prevent race conditions.\n # This ensures no concurrent request can modify this user's mappings.\n active_mapping = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email == email,\n UserTenantMapping.active == True, # noqa: E712\n )\n .with_for_update()\n .first()\n )\n\n # If an active mapping exists, delete it\n if active_mapping:\n db_session.delete(active_mapping)\n logger.info(\n f\"Deleted existing active mapping for user {email} in tenant {tenant_id}\"\n )\n\n # Find the inactive mapping for this user and tenant\n mapping = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email == email,\n UserTenantMapping.tenant_id == tenant_id,\n UserTenantMapping.active == False, # noqa: E712\n )\n .first()\n )\n\n if mapping:\n # Set all other mappings for this user to inactive\n db_session.query(UserTenantMapping).filter(\n UserTenantMapping.email == email,\n UserTenantMapping.active == True, # noqa: E712\n ).update({\"active\": False})\n\n # Activate this mapping\n mapping.active = True\n db_session.commit()\n logger.info(f\"User {email} accepted invitation to tenant {tenant_id}\")\n else:\n logger.warning(\n f\"No invitation found for user {email} in tenant {tenant_id}\"\n )\n\n except Exception as e:\n db_session.rollback()\n logger.exception(\n f\"Failed to accept invitation for user {email} to tenant {tenant_id}: {str(e)}\"\n )\n raise\n\n # Remove from invited users list since they've accepted\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n invited_users = get_invited_users()\n if email in invited_users:\n invited_users.remove(email)\n write_invited_users(invited_users)\n logger.info(f\"Removed {email} from invited users list after acceptance\")\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\ndef deny_user_invite(email: str, tenant_id: str) -> None:\n \"\"\"\n Deny an invitation to join a tenant.\n This removes the user's mapping to the tenant.\n \"\"\"\n with get_session_with_shared_schema() as db_session:\n # Delete the mapping for this user and tenant\n result = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email == email,\n UserTenantMapping.tenant_id == tenant_id,\n UserTenantMapping.active == False, # noqa: E712\n )\n .delete()\n )\n\n db_session.commit()\n if result:\n logger.info(f\"User {email} denied invitation to tenant {tenant_id}\")\n else:\n logger.warning(\n f\"No invitation found for user {email} in tenant {tenant_id}\"\n )\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n pending_users = get_invited_users()\n if email in pending_users:\n pending_users.remove(email)\n write_invited_users(pending_users)\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\ndef get_tenant_count(tenant_id: str) -> int:\n \"\"\"\n Get the number of active users for this tenant.\n\n A user counts toward the seat count if:\n 1. They have an active mapping to this tenant (UserTenantMapping.active == True)\n 2. AND the User is active (User.is_active == True)\n 3. AND the User is not the anonymous system user\n\n TODO: Exclude API key dummy users from seat counting. API keys create\n users with emails like `__DANSWER_API_KEY_*` that should not count toward\n seat limits. See: https://linear.app/onyx-app/issue/ENG-3518\n \"\"\"\n from onyx.configs.constants import ANONYMOUS_USER_EMAIL\n from onyx.db.models import User\n\n # First get all emails with active mappings to this tenant\n with get_session_with_shared_schema() as db_session:\n active_mapping_emails = (\n db_session.query(UserTenantMapping.email)\n .filter(\n UserTenantMapping.tenant_id == tenant_id,\n UserTenantMapping.active == True, # noqa: E712\n UserTenantMapping.email != ANONYMOUS_USER_EMAIL,\n )\n .all()\n )\n emails = [email for (email,) in active_mapping_emails]\n\n if not emails:\n return 0\n\n # Now count how many of those users are actually active in the tenant's User table\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n user_count = (\n db_session.query(User)\n .filter(\n User.email.in_(emails), # type: ignore\n User.is_active == True, # type: ignore # noqa: E712\n )\n .count()\n )\n\n return user_count\n\n\ndef get_tenant_invitation(email: str) -> TenantSnapshot | None:\n \"\"\"\n Get the first tenant invitation for this user\n \"\"\"\n with get_session_with_shared_schema() as db_session:\n # Get the first tenant invitation for this user\n invitation = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.email == email,\n UserTenantMapping.active == False, # noqa: E712\n )\n .first()\n )\n\n if invitation:\n # Get the user count for this tenant\n user_count = (\n db_session.query(UserTenantMapping)\n .filter(\n UserTenantMapping.tenant_id == invitation.tenant_id,\n UserTenantMapping.active == True, # noqa: E712\n )\n .count()\n )\n return TenantSnapshot(\n tenant_id=invitation.tenant_id, number_of_users=user_count\n )\n\n return None\n" }, { "path": "backend/ee/onyx/server/token_rate_limits/api.py", "content": "from collections import defaultdict\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.token_limit import fetch_all_user_group_token_rate_limits_by_group\nfrom ee.onyx.db.token_limit import fetch_user_group_token_rate_limits_for_user\nfrom ee.onyx.db.token_limit import insert_user_group_token_rate_limit\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.db.token_limit import fetch_all_user_token_rate_limits\nfrom onyx.db.token_limit import insert_user_token_rate_limit\nfrom onyx.server.query_and_chat.token_limit import any_rate_limit_exists\nfrom onyx.server.token_rate_limits.models import TokenRateLimitArgs\nfrom onyx.server.token_rate_limits.models import TokenRateLimitDisplay\n\nrouter = APIRouter(prefix=\"/admin/token-rate-limits\", tags=PUBLIC_API_TAGS)\n\n\n\"\"\"\nGroup Token Limit Settings\n\"\"\"\n\n\n@router.get(\"/user-groups\")\ndef get_all_group_token_limit_settings(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> dict[str, list[TokenRateLimitDisplay]]:\n user_groups_to_token_rate_limits = fetch_all_user_group_token_rate_limits_by_group(\n db_session\n )\n\n token_rate_limits_by_group = defaultdict(list)\n for token_rate_limit, group_name in user_groups_to_token_rate_limits:\n token_rate_limits_by_group[group_name].append(\n TokenRateLimitDisplay.from_db(token_rate_limit)\n )\n\n return dict(token_rate_limits_by_group)\n\n\n@router.get(\"/user-group/{group_id}\")\ndef get_group_token_limit_settings(\n group_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[TokenRateLimitDisplay]:\n return [\n TokenRateLimitDisplay.from_db(token_rate_limit)\n for token_rate_limit in fetch_user_group_token_rate_limits_for_user(\n db_session=db_session,\n group_id=group_id,\n user=user,\n )\n ]\n\n\n@router.post(\"/user-group/{group_id}\")\ndef create_group_token_limit_settings(\n group_id: int,\n token_limit_settings: TokenRateLimitArgs,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> TokenRateLimitDisplay:\n rate_limit_display = TokenRateLimitDisplay.from_db(\n insert_user_group_token_rate_limit(\n db_session=db_session,\n token_rate_limit_settings=token_limit_settings,\n group_id=group_id,\n )\n )\n # clear cache in case this was the first rate limit created\n any_rate_limit_exists.cache_clear()\n return rate_limit_display\n\n\n\"\"\"\nUser Token Limit Settings\n\"\"\"\n\n\n@router.get(\"/users\")\ndef get_user_token_limit_settings(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[TokenRateLimitDisplay]:\n return [\n TokenRateLimitDisplay.from_db(token_rate_limit)\n for token_rate_limit in fetch_all_user_token_rate_limits(db_session)\n ]\n\n\n@router.post(\"/users\")\ndef create_user_token_limit_settings(\n token_limit_settings: TokenRateLimitArgs,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> TokenRateLimitDisplay:\n rate_limit_display = TokenRateLimitDisplay.from_db(\n insert_user_token_rate_limit(db_session, token_limit_settings)\n )\n # clear cache in case this was the first rate limit created\n any_rate_limit_exists.cache_clear()\n return rate_limit_display\n" }, { "path": "backend/ee/onyx/server/usage_limits.py", "content": "\"\"\"EE Usage limits - trial detection via billing information.\"\"\"\n\nfrom ee.onyx.server.tenants.billing import fetch_billing_information\nfrom ee.onyx.server.tenants.models import BillingInformation\nfrom ee.onyx.server.tenants.models import SubscriptionStatusResponse\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef is_tenant_on_trial(tenant_id: str) -> bool:\n \"\"\"\n Determine if a tenant is currently on a trial subscription.\n\n In multi-tenant mode, we fetch billing information from the control plane\n to determine if the tenant has an active trial.\n \"\"\"\n if not MULTI_TENANT:\n return False\n\n try:\n billing_info = fetch_billing_information(tenant_id)\n\n # If not subscribed at all, check if we have trial information\n if isinstance(billing_info, SubscriptionStatusResponse):\n # No subscription means they're likely on trial (new tenant)\n return True\n\n if isinstance(billing_info, BillingInformation):\n return billing_info.status == \"trialing\"\n\n return False\n\n except Exception as e:\n logger.warning(f\"Failed to fetch billing info for trial check: {e}\")\n # Default to trial limits on error (more restrictive = safer)\n return True\n" }, { "path": "backend/ee/onyx/server/user_group/api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.persona import update_persona_access\nfrom ee.onyx.db.user_group import add_users_to_user_group\nfrom ee.onyx.db.user_group import delete_user_group as db_delete_user_group\nfrom ee.onyx.db.user_group import fetch_user_group\nfrom ee.onyx.db.user_group import fetch_user_groups\nfrom ee.onyx.db.user_group import fetch_user_groups_for_user\nfrom ee.onyx.db.user_group import insert_user_group\nfrom ee.onyx.db.user_group import prepare_user_group_for_deletion\nfrom ee.onyx.db.user_group import rename_user_group\nfrom ee.onyx.db.user_group import update_user_curator_relationship\nfrom ee.onyx.db.user_group import update_user_group\nfrom ee.onyx.server.user_group.models import AddUsersToUserGroupRequest\nfrom ee.onyx.server.user_group.models import MinimalUserGroupSnapshot\nfrom ee.onyx.server.user_group.models import SetCuratorRequest\nfrom ee.onyx.server.user_group.models import UpdateGroupAgentsRequest\nfrom ee.onyx.server.user_group.models import UserGroup\nfrom ee.onyx.server.user_group.models import UserGroupCreate\nfrom ee.onyx.server.user_group.models import UserGroupRename\nfrom ee.onyx.server.user_group.models import UserGroupUpdate\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.auth.users import current_user\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.db.models import UserRole\nfrom onyx.db.persona import get_persona_by_id\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/manage\", tags=PUBLIC_API_TAGS)\n\n\n@router.get(\"/admin/user-group\")\ndef list_user_groups(\n include_default: bool = False,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[UserGroup]:\n if user.role == UserRole.ADMIN:\n user_groups = fetch_user_groups(\n db_session,\n only_up_to_date=False,\n eager_load_for_snapshot=True,\n include_default=include_default,\n )\n else:\n user_groups = fetch_user_groups_for_user(\n db_session=db_session,\n user_id=user.id,\n only_curator_groups=user.role == UserRole.CURATOR,\n eager_load_for_snapshot=True,\n include_default=include_default,\n )\n return [UserGroup.from_model(user_group) for user_group in user_groups]\n\n\n@router.get(\"/user-groups/minimal\")\ndef list_minimal_user_groups(\n include_default: bool = False,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> list[MinimalUserGroupSnapshot]:\n if user.role == UserRole.ADMIN:\n user_groups = fetch_user_groups(\n db_session,\n only_up_to_date=False,\n include_default=include_default,\n )\n else:\n user_groups = fetch_user_groups_for_user(\n db_session=db_session,\n user_id=user.id,\n include_default=include_default,\n )\n return [\n MinimalUserGroupSnapshot.from_model(user_group) for user_group in user_groups\n ]\n\n\n@router.get(\"/admin/user-group/{user_group_id}/permissions\")\ndef get_user_group_permissions(\n user_group_id: int,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[str]:\n group = fetch_user_group(db_session, user_group_id)\n if group is None:\n raise OnyxError(OnyxErrorCode.NOT_FOUND, \"User group not found\")\n return [\n grant.permission.value\n for grant in group.permission_grants\n if not grant.is_deleted\n ]\n\n\n@router.post(\"/admin/user-group\")\ndef create_user_group(\n user_group: UserGroupCreate,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> UserGroup:\n try:\n db_user_group = insert_user_group(db_session, user_group)\n except IntegrityError:\n raise HTTPException(\n 400,\n f\"User group with name '{user_group.name}' already exists. Please \"\n + \"choose a different name.\",\n )\n return UserGroup.from_model(db_user_group)\n\n\n@router.patch(\"/admin/user-group/rename\")\ndef rename_user_group_endpoint(\n rename_request: UserGroupRename,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> UserGroup:\n group = fetch_user_group(db_session, rename_request.id)\n if group and group.is_default:\n raise OnyxError(OnyxErrorCode.CONFLICT, \"Cannot rename a default system group.\")\n try:\n return UserGroup.from_model(\n rename_user_group(\n db_session=db_session,\n user_group_id=rename_request.id,\n new_name=rename_request.name,\n )\n )\n except IntegrityError:\n raise OnyxError(\n OnyxErrorCode.DUPLICATE_RESOURCE,\n f\"User group with name '{rename_request.name}' already exists.\",\n )\n except ValueError as e:\n msg = str(e)\n if \"not found\" in msg.lower():\n raise OnyxError(OnyxErrorCode.NOT_FOUND, msg)\n raise OnyxError(OnyxErrorCode.CONFLICT, msg)\n\n\n@router.patch(\"/admin/user-group/{user_group_id}\")\ndef patch_user_group(\n user_group_id: int,\n user_group_update: UserGroupUpdate,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> UserGroup:\n try:\n return UserGroup.from_model(\n update_user_group(\n db_session=db_session,\n user=user,\n user_group_id=user_group_id,\n user_group_update=user_group_update,\n )\n )\n except ValueError as e:\n raise HTTPException(status_code=404, detail=str(e))\n\n\n@router.post(\"/admin/user-group/{user_group_id}/add-users\")\ndef add_users(\n user_group_id: int,\n add_users_request: AddUsersToUserGroupRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> UserGroup:\n try:\n return UserGroup.from_model(\n add_users_to_user_group(\n db_session=db_session,\n user=user,\n user_group_id=user_group_id,\n user_ids=add_users_request.user_ids,\n )\n )\n except ValueError as e:\n raise HTTPException(status_code=404, detail=str(e))\n\n\n@router.post(\"/admin/user-group/{user_group_id}/set-curator\")\ndef set_user_curator(\n user_group_id: int,\n set_curator_request: SetCuratorRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> None:\n try:\n update_user_curator_relationship(\n db_session=db_session,\n user_group_id=user_group_id,\n set_curator_request=set_curator_request,\n user_making_change=user,\n )\n except ValueError as e:\n logger.error(f\"Error setting user curator: {e}\")\n raise HTTPException(status_code=404, detail=str(e))\n\n\n@router.delete(\"/admin/user-group/{user_group_id}\")\ndef delete_user_group(\n user_group_id: int,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> None:\n group = fetch_user_group(db_session, user_group_id)\n if group and group.is_default:\n raise OnyxError(OnyxErrorCode.CONFLICT, \"Cannot delete a default system group.\")\n try:\n prepare_user_group_for_deletion(db_session, user_group_id)\n except ValueError as e:\n raise HTTPException(status_code=404, detail=str(e))\n\n if DISABLE_VECTOR_DB:\n user_group = fetch_user_group(db_session, user_group_id)\n if user_group:\n db_delete_user_group(db_session, user_group)\n\n\n@router.patch(\"/admin/user-group/{user_group_id}/agents\")\ndef update_group_agents(\n user_group_id: int,\n request: UpdateGroupAgentsRequest,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> None:\n for agent_id in request.added_agent_ids:\n persona = get_persona_by_id(\n persona_id=agent_id, user=user, db_session=db_session\n )\n current_group_ids = [g.id for g in persona.groups]\n if user_group_id not in current_group_ids:\n update_persona_access(\n persona_id=agent_id,\n creator_user_id=user.id,\n db_session=db_session,\n group_ids=current_group_ids + [user_group_id],\n )\n\n for agent_id in request.removed_agent_ids:\n persona = get_persona_by_id(\n persona_id=agent_id, user=user, db_session=db_session\n )\n current_group_ids = [g.id for g in persona.groups]\n update_persona_access(\n persona_id=agent_id,\n creator_user_id=user.id,\n db_session=db_session,\n group_ids=[gid for gid in current_group_ids if gid != user_group_id],\n )\n\n db_session.commit()\n" }, { "path": "backend/ee/onyx/server/user_group/models.py", "content": "from uuid import UUID\n\nfrom pydantic import BaseModel\n\nfrom onyx.db.models import UserGroup as UserGroupModel\nfrom onyx.server.documents.models import ConnectorCredentialPairDescriptor\nfrom onyx.server.documents.models import ConnectorSnapshot\nfrom onyx.server.documents.models import CredentialSnapshot\nfrom onyx.server.features.document_set.models import DocumentSet\nfrom onyx.server.features.persona.models import PersonaSnapshot\nfrom onyx.server.manage.models import UserInfo\nfrom onyx.server.manage.models import UserPreferences\n\n\nclass UserGroup(BaseModel):\n id: int\n name: str\n users: list[UserInfo]\n curator_ids: list[UUID]\n cc_pairs: list[ConnectorCredentialPairDescriptor]\n document_sets: list[DocumentSet]\n personas: list[PersonaSnapshot]\n is_up_to_date: bool\n is_up_for_deletion: bool\n is_default: bool\n\n @classmethod\n def from_model(cls, user_group_model: UserGroupModel) -> \"UserGroup\":\n return cls(\n id=user_group_model.id,\n name=user_group_model.name,\n users=[\n UserInfo(\n id=str(user.id),\n email=user.email,\n is_active=user.is_active,\n is_superuser=user.is_superuser,\n is_verified=user.is_verified,\n role=user.role,\n preferences=UserPreferences(\n default_model=user.default_model,\n chosen_assistants=user.chosen_assistants,\n ),\n )\n for user in user_group_model.users\n ],\n curator_ids=[\n user.user_id\n for user in user_group_model.user_group_relationships\n if user.is_curator and user.user_id is not None\n ],\n cc_pairs=[\n ConnectorCredentialPairDescriptor(\n id=cc_pair_relationship.cc_pair.id,\n name=cc_pair_relationship.cc_pair.name,\n connector=ConnectorSnapshot.from_connector_db_model(\n cc_pair_relationship.cc_pair.connector,\n credential_ids=[cc_pair_relationship.cc_pair.credential_id],\n ),\n credential=CredentialSnapshot.from_credential_db_model(\n cc_pair_relationship.cc_pair.credential\n ),\n access_type=cc_pair_relationship.cc_pair.access_type,\n )\n for cc_pair_relationship in user_group_model.cc_pair_relationships\n if cc_pair_relationship.is_current\n ],\n document_sets=[\n DocumentSet.from_model(ds) for ds in user_group_model.document_sets\n ],\n personas=[\n PersonaSnapshot.from_model(persona)\n for persona in user_group_model.personas\n if not persona.deleted\n ],\n is_up_to_date=user_group_model.is_up_to_date,\n is_up_for_deletion=user_group_model.is_up_for_deletion,\n is_default=user_group_model.is_default,\n )\n\n\nclass MinimalUserGroupSnapshot(BaseModel):\n id: int\n name: str\n is_default: bool\n\n @classmethod\n def from_model(cls, user_group_model: UserGroupModel) -> \"MinimalUserGroupSnapshot\":\n return cls(\n id=user_group_model.id,\n name=user_group_model.name,\n is_default=user_group_model.is_default,\n )\n\n\nclass UserGroupCreate(BaseModel):\n name: str\n user_ids: list[UUID]\n cc_pair_ids: list[int]\n\n\nclass UserGroupUpdate(BaseModel):\n user_ids: list[UUID]\n cc_pair_ids: list[int]\n\n\nclass AddUsersToUserGroupRequest(BaseModel):\n user_ids: list[UUID]\n\n\nclass UserGroupRename(BaseModel):\n id: int\n name: str\n\n\nclass SetCuratorRequest(BaseModel):\n user_id: UUID\n is_curator: bool\n\n\nclass UpdateGroupAgentsRequest(BaseModel):\n added_agent_ids: list[int]\n removed_agent_ids: list[int]\n" }, { "path": "backend/ee/onyx/utils/__init__.py", "content": "" }, { "path": "backend/ee/onyx/utils/encryption.py", "content": "from functools import lru_cache\nfrom os import urandom\n\nfrom cryptography.hazmat.backends import default_backend\nfrom cryptography.hazmat.primitives import padding\nfrom cryptography.hazmat.primitives.ciphers import algorithms\nfrom cryptography.hazmat.primitives.ciphers import Cipher\nfrom cryptography.hazmat.primitives.ciphers import modes\n\nfrom onyx.configs.app_configs import ENCRYPTION_KEY_SECRET\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\nlogger = setup_logger()\n\n\n@lru_cache(maxsize=2)\ndef _get_trimmed_key(key: str) -> bytes:\n encoded_key = key.encode()\n key_length = len(encoded_key)\n if key_length < 16:\n raise RuntimeError(\"Invalid ENCRYPTION_KEY_SECRET - too short\")\n\n # Trim to the largest valid AES key size that fits\n valid_lengths = [32, 24, 16]\n for size in valid_lengths:\n if key_length >= size:\n return encoded_key[:size]\n\n raise AssertionError(\"unreachable\")\n\n\ndef _encrypt_string(input_str: str, key: str | None = None) -> bytes:\n effective_key = key if key is not None else ENCRYPTION_KEY_SECRET\n if not effective_key:\n return input_str.encode()\n\n trimmed = _get_trimmed_key(effective_key)\n iv = urandom(16)\n padder = padding.PKCS7(algorithms.AES.block_size).padder()\n padded_data = padder.update(input_str.encode()) + padder.finalize()\n\n cipher = Cipher(algorithms.AES(trimmed), modes.CBC(iv), backend=default_backend())\n encryptor = cipher.encryptor()\n encrypted_data = encryptor.update(padded_data) + encryptor.finalize()\n\n return iv + encrypted_data\n\n\ndef _decrypt_bytes(input_bytes: bytes, key: str | None = None) -> str:\n effective_key = key if key is not None else ENCRYPTION_KEY_SECRET\n if not effective_key:\n return input_bytes.decode()\n\n trimmed = _get_trimmed_key(effective_key)\n try:\n iv = input_bytes[:16]\n encrypted_data = input_bytes[16:]\n\n cipher = Cipher(\n algorithms.AES(trimmed), modes.CBC(iv), backend=default_backend()\n )\n decryptor = cipher.decryptor()\n decrypted_padded_data = decryptor.update(encrypted_data) + decryptor.finalize()\n\n unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()\n decrypted_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()\n\n return decrypted_data.decode()\n except (ValueError, UnicodeDecodeError):\n if key is not None:\n # Explicit key was provided — don't fall back silently\n raise\n # Read path: attempt raw UTF-8 decode as a fallback for legacy data.\n # Does NOT handle data encrypted with a different key — that\n # ciphertext is not valid UTF-8 and will raise below.\n logger.warning(\n \"AES decryption failed — falling back to raw decode. Run the re-encrypt secrets script to rotate to the current key.\"\n )\n try:\n return input_bytes.decode()\n except UnicodeDecodeError:\n raise ValueError(\n \"Data is not valid UTF-8 — likely encrypted with a different key. \"\n \"Run the re-encrypt secrets script to rotate to the current key.\"\n ) from None\n\n\ndef encrypt_string_to_bytes(input_str: str, key: str | None = None) -> bytes:\n versioned_encryption_fn = fetch_versioned_implementation(\n \"onyx.utils.encryption\", \"_encrypt_string\"\n )\n return versioned_encryption_fn(input_str, key=key)\n\n\ndef decrypt_bytes_to_string(input_bytes: bytes, key: str | None = None) -> str:\n versioned_decryption_fn = fetch_versioned_implementation(\n \"onyx.utils.encryption\", \"_decrypt_bytes\"\n )\n return versioned_decryption_fn(input_bytes, key=key)\n\n\ndef test_encryption() -> None:\n test_string = \"Onyx is the BEST!\"\n encrypted_bytes = encrypt_string_to_bytes(test_string)\n decrypted_string = decrypt_bytes_to_string(encrypted_bytes)\n if test_string != decrypted_string:\n raise RuntimeError(\"Encryption decryption test failed\")\n" }, { "path": "backend/ee/onyx/utils/license.py", "content": "\"\"\"RSA-4096 license signature verification utilities.\"\"\"\n\nimport base64\nimport json\nimport os\nfrom datetime import datetime\nfrom datetime import timezone\nfrom pathlib import Path\n\nfrom cryptography.exceptions import InvalidSignature\nfrom cryptography.hazmat.primitives import hashes\nfrom cryptography.hazmat.primitives import serialization\nfrom cryptography.hazmat.primitives.asymmetric import padding\nfrom cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey\n\nfrom ee.onyx.server.license.models import LicenseData\nfrom ee.onyx.server.license.models import LicensePayload\nfrom onyx.server.settings.models import ApplicationStatus\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Path to the license public key file\n_LICENSE_PUBLIC_KEY_PATH = (\n Path(__file__).parent.parent.parent.parent / \"keys\" / \"license_public_key.pem\"\n)\n\n\ndef _get_public_key() -> RSAPublicKey:\n \"\"\"Load the public key from file, with env var override.\"\"\"\n # Allow env var override for flexibility\n key_pem = os.environ.get(\"LICENSE_PUBLIC_KEY_PEM\")\n\n if not key_pem:\n # Read from file\n if not _LICENSE_PUBLIC_KEY_PATH.exists():\n raise ValueError(\n f\"License public key not found at {_LICENSE_PUBLIC_KEY_PATH}. \"\n \"License verification requires the control plane public key.\"\n )\n key_pem = _LICENSE_PUBLIC_KEY_PATH.read_text()\n\n key = serialization.load_pem_public_key(key_pem.encode())\n if not isinstance(key, RSAPublicKey):\n raise ValueError(\"Expected RSA public key\")\n return key\n\n\ndef verify_license_signature(license_data: str) -> LicensePayload:\n \"\"\"\n Verify RSA-4096 signature and return payload if valid.\n\n Args:\n license_data: Base64-encoded JSON containing payload and signature\n\n Returns:\n LicensePayload if signature is valid\n\n Raises:\n ValueError: If license data is invalid or signature verification fails\n \"\"\"\n try:\n decoded = json.loads(base64.b64decode(license_data))\n\n # Parse into LicenseData to validate structure\n license_obj = LicenseData(**decoded)\n\n # IMPORTANT: Use the ORIGINAL payload JSON for signature verification,\n # not re-serialized through Pydantic. Pydantic may format fields differently\n # (e.g., datetime \"+00:00\" vs \"Z\") which would break signature verification.\n original_payload = decoded.get(\"payload\", {})\n payload_json = json.dumps(original_payload, sort_keys=True)\n signature_bytes = base64.b64decode(license_obj.signature)\n\n # Verify signature using PSS padding (modern standard)\n public_key = _get_public_key()\n\n public_key.verify(\n signature_bytes,\n payload_json.encode(),\n padding.PSS(\n mgf=padding.MGF1(hashes.SHA256()),\n salt_length=padding.PSS.MAX_LENGTH,\n ),\n hashes.SHA256(),\n )\n\n return license_obj.payload\n\n except InvalidSignature:\n logger.error(\"[verify_license] FAILED: Signature verification failed\")\n raise ValueError(\"Invalid license signature\")\n except json.JSONDecodeError as e:\n logger.error(f\"[verify_license] FAILED: JSON decode error: {e}\")\n raise ValueError(\"Invalid license format: not valid JSON\")\n except (ValueError, KeyError, TypeError) as e:\n logger.error(\n f\"[verify_license] FAILED: Validation error: {type(e).__name__}: {e}\"\n )\n raise ValueError(f\"Invalid license format: {type(e).__name__}: {e}\")\n except Exception:\n logger.exception(\"[verify_license] FAILED: Unexpected error\")\n raise ValueError(\"License verification failed: unexpected error\")\n\n\ndef get_license_status(\n payload: LicensePayload,\n grace_period_end: datetime | None = None,\n) -> ApplicationStatus:\n \"\"\"\n Determine current license status based on expiry.\n\n Args:\n payload: The verified license payload\n grace_period_end: Optional grace period end datetime\n\n Returns:\n ApplicationStatus indicating current license state\n \"\"\"\n now = datetime.now(timezone.utc)\n\n # Check if grace period has expired\n if grace_period_end and now > grace_period_end:\n return ApplicationStatus.GATED_ACCESS\n\n # Check if license has expired\n if now > payload.expires_at:\n if grace_period_end and now <= grace_period_end:\n return ApplicationStatus.GRACE_PERIOD\n return ApplicationStatus.GATED_ACCESS\n\n # License is valid\n return ApplicationStatus.ACTIVE\n\n\ndef is_license_valid(payload: LicensePayload) -> bool:\n \"\"\"Check if a license is currently valid (not expired).\"\"\"\n now = datetime.now(timezone.utc)\n return now <= payload.expires_at\n" }, { "path": "backend/ee/onyx/utils/posthog_client.py", "content": "import json\nfrom typing import Any\nfrom urllib.parse import unquote\n\nfrom posthog import Posthog\n\nfrom ee.onyx.configs.app_configs import MARKETING_POSTHOG_API_KEY\nfrom ee.onyx.configs.app_configs import POSTHOG_API_KEY\nfrom ee.onyx.configs.app_configs import POSTHOG_DEBUG_LOGS_ENABLED\nfrom ee.onyx.configs.app_configs import POSTHOG_HOST\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef posthog_on_error(error: Any, items: Any) -> None:\n \"\"\"Log any PostHog delivery errors.\"\"\"\n logger.error(f\"PostHog error: {error}, items: {items}\")\n\n\nposthog: Posthog | None = None\nif POSTHOG_API_KEY:\n posthog = Posthog(\n project_api_key=POSTHOG_API_KEY,\n host=POSTHOG_HOST,\n debug=POSTHOG_DEBUG_LOGS_ENABLED,\n on_error=posthog_on_error,\n )\nelif MULTI_TENANT:\n logger.warning(\n \"POSTHOG_API_KEY is not set but MULTI_TENANT is enabled — \"\n \"PostHog telemetry and feature flags will be disabled\"\n )\n\n# For cross referencing between cloud and www Onyx sites\n# NOTE: These clients are separate because they are separate posthog projects.\n# We should eventually unify them into a single posthog project,\n# which would no longer require this workaround\nmarketing_posthog = None\nif MARKETING_POSTHOG_API_KEY:\n marketing_posthog = Posthog(\n project_api_key=MARKETING_POSTHOG_API_KEY,\n host=POSTHOG_HOST,\n debug=POSTHOG_DEBUG_LOGS_ENABLED,\n on_error=posthog_on_error,\n )\n\n\ndef capture_and_sync_with_alternate_posthog(\n alternate_distinct_id: str, event: str, properties: dict[str, Any]\n) -> None:\n \"\"\"\n Identify in both PostHog projects and capture the event in marketing.\n - Marketing keeps the marketing distinct_id (for feature flags).\n - Cloud identify uses the cloud distinct_id\n \"\"\"\n if not marketing_posthog:\n return\n\n props = properties.copy()\n\n try:\n marketing_posthog.identify(distinct_id=alternate_distinct_id, properties=props)\n marketing_posthog.capture(alternate_distinct_id, event, props)\n marketing_posthog.flush()\n except Exception as e:\n logger.error(f\"Error capturing marketing posthog event: {e}\")\n\n try:\n if posthog and (cloud_user_id := props.get(\"onyx_cloud_user_id\")):\n cloud_props = props.copy()\n cloud_props.pop(\"onyx_cloud_user_id\", None)\n\n posthog.identify(\n distinct_id=cloud_user_id,\n properties=cloud_props,\n )\n except Exception as e:\n logger.error(f\"Error identifying cloud posthog user: {e}\")\n\n\ndef alias_user(distinct_id: str, anonymous_id: str) -> None:\n \"\"\"Link an anonymous distinct_id to an identified user, merging person profiles.\n\n No-ops when the IDs match (e.g. returning users whose PostHog cookie\n already contains their identified user ID).\n \"\"\"\n if not posthog or anonymous_id == distinct_id:\n return\n\n try:\n posthog.alias(previous_id=anonymous_id, distinct_id=distinct_id)\n posthog.flush()\n except Exception as e:\n logger.error(f\"Error aliasing PostHog user: {e}\")\n\n\ndef get_anon_id_from_request(request: Any) -> str | None:\n \"\"\"Extract the anonymous distinct_id from the app PostHog cookie on a request.\"\"\"\n if not POSTHOG_API_KEY:\n return None\n\n cookie_name = f\"ph_{POSTHOG_API_KEY}_posthog\"\n if (cookie_value := request.cookies.get(cookie_name)) and (\n parsed := parse_posthog_cookie(cookie_value)\n ):\n return parsed.get(\"distinct_id\")\n\n return None\n\n\ndef get_marketing_posthog_cookie_name() -> str | None:\n if not MARKETING_POSTHOG_API_KEY:\n return None\n return f\"onyx_custom_ph_{MARKETING_POSTHOG_API_KEY}_posthog\"\n\n\ndef parse_posthog_cookie(cookie_value: str) -> dict[str, Any] | None:\n \"\"\"\n Parse a URL-encoded JSON PostHog cookie\n\n Expected format (URL-encoded):\n {\"distinct_id\":\"...\", \"featureFlags\":{\"landing_page_variant\":\"...\"}, ...}\n\n Returns:\n Dict with 'distinct_id' explicitly required and all other cookie values\n passed through as-is, or None if parsing fails or distinct_id is missing.\n \"\"\"\n try:\n decoded_cookie = unquote(cookie_value)\n cookie_data = json.loads(decoded_cookie)\n\n distinct_id = cookie_data.get(\"distinct_id\")\n if not distinct_id or not isinstance(distinct_id, str):\n return None\n\n return cookie_data\n except (json.JSONDecodeError, KeyError, TypeError, AttributeError) as e:\n logger.warning(f\"Failed to parse cookie: {e}\")\n return None\n" }, { "path": "backend/ee/onyx/utils/telemetry.py", "content": "from typing import Any\n\nfrom ee.onyx.utils.posthog_client import posthog\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef event_telemetry(\n distinct_id: str, event: str, properties: dict[str, Any] | None = None\n) -> None:\n \"\"\"Capture and send an event to PostHog, flushing immediately.\"\"\"\n if not posthog:\n return\n\n logger.info(f\"Capturing PostHog event: {distinct_id} {event} {properties}\")\n try:\n posthog.capture(distinct_id, event, properties)\n posthog.flush()\n except Exception as e:\n logger.error(f\"Error capturing PostHog event: {e}\")\n\n\ndef identify_user(distinct_id: str, properties: dict[str, Any] | None = None) -> None:\n \"\"\"Create/update a PostHog person profile, flushing immediately.\"\"\"\n if not posthog:\n return\n\n try:\n posthog.identify(distinct_id, properties)\n posthog.flush()\n except Exception as e:\n logger.error(f\"Error identifying PostHog user: {e}\")\n" }, { "path": "backend/generated/README.md", "content": "- Generated Files\n* Generated files live here. This directory should be git ignored." }, { "path": "backend/keys/license_public_key.pem", "content": "-----BEGIN PUBLIC KEY-----\nMIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEA5DpchQujdxjCwpc4/RQP\nHej6rc3SS/5ENCXL0I8NAfMogel0fqG6PKRhonyEh/Bt3P4q18y8vYzAShwf4b6Q\naS0WwshbvnkjyWlsK0BY4HLBKPkTpes7kaz8MwmPZDeelvGJ7SNv3FvyJR4QsoSQ\nGSoB5iTH7hi63TjzdxtckkXoNG+GdVd/koxVDUv2uWcAoWIFTTcbKWyuq2SS/5Sf\nxdVaIArqfAhLpnNbnM9OS7lZ1xP+29ZXpHxDoeluz35tJLMNBYn9u0y+puo1kW1E\nTOGizlAq5kmEMsTJ55e9ZuyIV3gZAUaUKe8CxYJPkOGt0Gj6e1jHoHZCBJmaq97Y\nstKj//84HNBzajaryEZuEfRecJ94ANEjkD8u9cGmW+9VxRe5544zWguP5WMT/nv1\n0Q+jkOBW2hkY5SS0Rug4cblxiB7bDymWkaX6+sC0VWd5g6WXp36EuP2T0v3mYuHU\nGDEiWbD44ToREPVwE/M07ny8qhLo/HYk2l8DKFt83hXe7ePBnyQdcsrVbQWOO1na\nj43OkoU5gOFyOkrk2RmmtCjA8jSnw+tGCTpRaRcshqoWC1MjZyU+8/kDteXNkmv9\n/B5VxzYSyX+abl7yAu5wLiUPW8l+mOazzWu0nPkmiA160ArxnRyxbGnmp4dUIrt5\nazYku4tQYLSsSabfhcpeiCsCAwEAAQ==\n-----END PUBLIC KEY-----\n" }, { "path": "backend/model_server/__init__.py", "content": "" }, { "path": "backend/model_server/constants.py", "content": "MODEL_WARM_UP_STRING = \"hi \" * 512\n\n\nclass GPUStatus:\n CUDA = \"cuda\"\n MAC_MPS = \"mps\"\n NONE = \"none\"\n" }, { "path": "backend/model_server/encoders.py", "content": "import asyncio\nimport time\nfrom typing import Any\nfrom typing import TYPE_CHECKING\n\nfrom fastapi import APIRouter\nfrom fastapi import HTTPException\nfrom fastapi import Request\n\nfrom model_server.utils import simple_log_function_time\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.enums import EmbedTextType\nfrom shared_configs.model_server_models import Embedding\nfrom shared_configs.model_server_models import EmbedRequest\nfrom shared_configs.model_server_models import EmbedResponse\n\nif TYPE_CHECKING:\n from sentence_transformers import SentenceTransformer\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/encoder\")\n\n\n_GLOBAL_MODELS_DICT: dict[str, \"SentenceTransformer\"] = {}\n\n\ndef get_embedding_model(\n model_name: str,\n max_context_length: int,\n) -> \"SentenceTransformer\":\n \"\"\"\n Loads or returns a cached SentenceTransformer, sets max_seq_length, pins device,\n pre-warms rotary caches once, and wraps encode() with a lock to avoid cache races.\n \"\"\"\n from sentence_transformers import SentenceTransformer\n\n def _prewarm_rope(st_model: \"SentenceTransformer\", target_len: int) -> None:\n \"\"\"\n Build RoPE cos/sin caches once on the final device/dtype so later forwards only read.\n Works by calling the underlying HF model directly with dummy IDs/attention.\n \"\"\"\n try:\n # ensure > max seq after tokenization\n # Ideally we would use the saved tokenizer, but whatever it's ok\n # we'll make an assumption about tokenization here\n long_text = \"x \" * (target_len * 2)\n _ = st_model.encode(\n [long_text],\n batch_size=1,\n convert_to_tensor=True,\n show_progress_bar=False,\n normalize_embeddings=False,\n )\n logger.info(\"RoPE pre-warm successful\")\n except Exception as e:\n logger.warning(f\"RoPE pre-warm skipped/failed: {e}\")\n\n global _GLOBAL_MODELS_DICT\n\n if model_name not in _GLOBAL_MODELS_DICT:\n logger.notice(f\"Loading {model_name}\")\n model = SentenceTransformer(\n model_name_or_path=model_name,\n trust_remote_code=True,\n )\n model.max_seq_length = max_context_length\n _prewarm_rope(model, max_context_length)\n _GLOBAL_MODELS_DICT[model_name] = model\n else:\n model = _GLOBAL_MODELS_DICT[model_name]\n if max_context_length != model.max_seq_length:\n model.max_seq_length = max_context_length\n prev = getattr(model, \"_rope_prewarmed_to\", 0)\n if max_context_length > int(prev or 0):\n _prewarm_rope(model, max_context_length)\n\n return _GLOBAL_MODELS_DICT[model_name]\n\n\nENCODING_RETRIES = 3\nENCODING_RETRY_DELAY = 0.1\n\n\ndef _concurrent_embedding(\n texts: list[str], model: \"SentenceTransformer\", normalize_embeddings: bool\n) -> Any:\n \"\"\"Synchronous wrapper for concurrent_embedding to use with run_in_executor.\"\"\"\n for _ in range(ENCODING_RETRIES):\n try:\n return model.encode(texts, normalize_embeddings=normalize_embeddings)\n except RuntimeError as e:\n # There is a concurrency bug in the SentenceTransformer library that causes\n # the model to fail to encode texts. It's pretty rare and we want to allow\n # concurrent embedding, hence we retry (the specific error is\n # \"RuntimeError: Already borrowed\" and occurs in the transformers library)\n logger.warning(f\"Error encoding texts, retrying: {e}\")\n time.sleep(ENCODING_RETRY_DELAY)\n return model.encode(texts, normalize_embeddings=normalize_embeddings)\n\n\n@simple_log_function_time()\nasync def embed_text(\n texts: list[str],\n model_name: str | None,\n max_context_length: int,\n normalize_embeddings: bool,\n prefix: str | None,\n gpu_type: str = \"UNKNOWN\",\n) -> list[Embedding]:\n if not all(texts):\n logger.error(\"Empty strings provided for embedding\")\n raise ValueError(\"Empty strings are not allowed for embedding.\")\n\n if not texts:\n logger.error(\"No texts provided for embedding\")\n raise ValueError(\"No texts provided for embedding.\")\n\n start = time.monotonic()\n\n total_chars = 0\n for text in texts:\n total_chars += len(text)\n\n # Only local models should call this function now\n # API providers should go directly to API server\n\n if model_name is not None:\n logger.info(\n f\"Embedding {len(texts)} texts with {total_chars} total characters with local model: {model_name}\"\n )\n\n prefixed_texts = [f\"{prefix}{text}\" for text in texts] if prefix else texts\n\n local_model = get_embedding_model(\n model_name=model_name, max_context_length=max_context_length\n )\n # Run CPU-bound embedding in a thread pool\n embeddings_vectors = await asyncio.get_event_loop().run_in_executor(\n None,\n lambda: _concurrent_embedding(\n prefixed_texts, local_model, normalize_embeddings\n ),\n )\n embeddings = [\n embedding if isinstance(embedding, list) else embedding.tolist()\n for embedding in embeddings_vectors\n ]\n\n elapsed = time.monotonic() - start\n logger.info(\n f\"Successfully embedded {len(texts)} texts with {total_chars} total characters \"\n f\"with local model {model_name} in {elapsed:.2f}\"\n )\n logger.info(\n f\"event=embedding_model \"\n f\"texts={len(texts)} \"\n f\"chars={total_chars} \"\n f\"model={model_name} \"\n f\"gpu={gpu_type} \"\n f\"elapsed={elapsed:.2f}\"\n )\n else:\n logger.error(\"Model name not specified for embedding\")\n raise ValueError(\"Model name must be provided to run embeddings.\")\n\n return embeddings\n\n\n@router.post(\"/bi-encoder-embed\")\nasync def route_bi_encoder_embed(\n request: Request,\n embed_request: EmbedRequest,\n) -> EmbedResponse:\n return await process_embed_request(embed_request, request.app.state.gpu_type)\n\n\nasync def process_embed_request(\n embed_request: EmbedRequest, gpu_type: str = \"UNKNOWN\"\n) -> EmbedResponse:\n from litellm.exceptions import RateLimitError\n\n # Only local models should use this endpoint - API providers should make direct API calls\n if embed_request.provider_type is not None:\n raise ValueError(\n f\"Model server embedding endpoint should only be used for local models. \"\n f\"API provider '{embed_request.provider_type}' should make direct API calls instead.\"\n )\n\n if not embed_request.texts:\n raise HTTPException(status_code=400, detail=\"No texts to be embedded\")\n\n if not all(embed_request.texts):\n raise ValueError(\"Empty strings are not allowed for embedding.\")\n\n try:\n if embed_request.text_type == EmbedTextType.QUERY:\n prefix = embed_request.manual_query_prefix\n elif embed_request.text_type == EmbedTextType.PASSAGE:\n prefix = embed_request.manual_passage_prefix\n else:\n prefix = None\n\n embeddings = await embed_text(\n texts=embed_request.texts,\n model_name=embed_request.model_name,\n max_context_length=embed_request.max_context_length,\n normalize_embeddings=embed_request.normalize_embeddings,\n prefix=prefix,\n gpu_type=gpu_type,\n )\n return EmbedResponse(embeddings=embeddings)\n except RateLimitError as e:\n raise HTTPException(\n status_code=429,\n detail=str(e),\n )\n except Exception as e:\n logger.exception(\n f\"Error during embedding process: provider={embed_request.provider_type} model={embed_request.model_name}\"\n )\n raise HTTPException(\n status_code=500, detail=f\"Error during embedding process: {e}\"\n )\n" }, { "path": "backend/model_server/legacy/README.md", "content": "This directory contains code that was useful and may become useful again in the future.\n\nWe stopped using rerankers because the state of the art rerankers are not significantly better than the biencoders and much worse than LLMs which are also capable of acting on a small set of documents for filtering, reranking, etc.\n\nWe stopped using the internal query classifier as that's now offloaded to the LLM which does query expansion so we know ahead of time if it's a keyword or semantic query.\n" }, { "path": "backend/model_server/legacy/__init__.py", "content": "" }, { "path": "backend/model_server/legacy/custom_models.py", "content": "# from typing import cast\n# from typing import Optional\n# from typing import TYPE_CHECKING\n\n# import numpy as np\n# import torch\n# import torch.nn.functional as F\n# from fastapi import APIRouter\n# from huggingface_hub import snapshot_download\n# from pydantic import BaseModel\n\n# from model_server.constants import MODEL_WARM_UP_STRING\n# from model_server.legacy.onyx_torch_model import ConnectorClassifier\n# from model_server.legacy.onyx_torch_model import HybridClassifier\n# from model_server.utils import simple_log_function_time\n# from onyx.utils.logger import setup_logger\n# from shared_configs.configs import CONNECTOR_CLASSIFIER_MODEL_REPO\n# from shared_configs.configs import CONNECTOR_CLASSIFIER_MODEL_TAG\n# from shared_configs.configs import INDEXING_ONLY\n# from shared_configs.configs import INTENT_MODEL_TAG\n# from shared_configs.configs import INTENT_MODEL_VERSION\n# from shared_configs.model_server_models import IntentRequest\n# from shared_configs.model_server_models import IntentResponse\n\n# if TYPE_CHECKING:\n# from setfit import SetFitModel # type: ignore[import-untyped]\n# from transformers import PreTrainedTokenizer, BatchEncoding\n\n\n# INFORMATION_CONTENT_MODEL_WARM_UP_STRING = \"hi\" * 50\n\n# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MAX = 1.0\n# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN = 0.7\n# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_TEMPERATURE = 4.0\n# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_CUTOFF_LENGTH = 10\n# INFORMATION_CONTENT_MODEL_VERSION = \"onyx-dot-app/information-content-model\"\n# INFORMATION_CONTENT_MODEL_TAG: str | None = None\n\n\n# class ConnectorClassificationRequest(BaseModel):\n# available_connectors: list[str]\n# query: str\n\n\n# class ConnectorClassificationResponse(BaseModel):\n# connectors: list[str]\n\n\n# class ContentClassificationPrediction(BaseModel):\n# predicted_label: int\n# content_boost_factor: float\n\n\n# logger = setup_logger()\n\n# router = APIRouter(prefix=\"/custom\")\n\n# _CONNECTOR_CLASSIFIER_TOKENIZER: Optional[\"PreTrainedTokenizer\"] = None\n# _CONNECTOR_CLASSIFIER_MODEL: ConnectorClassifier | None = None\n\n# _INTENT_TOKENIZER: Optional[\"PreTrainedTokenizer\"] = None\n# _INTENT_MODEL: HybridClassifier | None = None\n\n# _INFORMATION_CONTENT_MODEL: Optional[\"SetFitModel\"] = None\n\n# _INFORMATION_CONTENT_MODEL_PROMPT_PREFIX: str = \"\" # spec to model version!\n\n\n# def get_connector_classifier_tokenizer() -> \"PreTrainedTokenizer\":\n# global _CONNECTOR_CLASSIFIER_TOKENIZER\n# from transformers import AutoTokenizer, PreTrainedTokenizer\n\n# if _CONNECTOR_CLASSIFIER_TOKENIZER is None:\n# # The tokenizer details are not uploaded to the HF hub since it's just the\n# # unmodified distilbert tokenizer.\n# _CONNECTOR_CLASSIFIER_TOKENIZER = cast(\n# PreTrainedTokenizer,\n# AutoTokenizer.from_pretrained(\"distilbert-base-uncased\"),\n# )\n# return _CONNECTOR_CLASSIFIER_TOKENIZER\n\n\n# def get_local_connector_classifier(\n# model_name_or_path: str = CONNECTOR_CLASSIFIER_MODEL_REPO,\n# tag: str = CONNECTOR_CLASSIFIER_MODEL_TAG,\n# ) -> ConnectorClassifier:\n# global _CONNECTOR_CLASSIFIER_MODEL\n# if _CONNECTOR_CLASSIFIER_MODEL is None:\n# try:\n# # Calculate where the cache should be, then load from local if available\n# local_path = snapshot_download(\n# repo_id=model_name_or_path, revision=tag, local_files_only=True\n# )\n# _CONNECTOR_CLASSIFIER_MODEL = ConnectorClassifier.from_pretrained(\n# local_path\n# )\n# except Exception as e:\n# logger.warning(f\"Failed to load model directly: {e}\")\n# try:\n# # Attempt to download the model snapshot\n# logger.info(f\"Downloading model snapshot for {model_name_or_path}\")\n# local_path = snapshot_download(repo_id=model_name_or_path, revision=tag)\n# _CONNECTOR_CLASSIFIER_MODEL = ConnectorClassifier.from_pretrained(\n# local_path\n# )\n# except Exception as e:\n# logger.error(\n# f\"Failed to load model even after attempted snapshot download: {e}\"\n# )\n# raise\n# return _CONNECTOR_CLASSIFIER_MODEL\n\n\n# def get_intent_model_tokenizer() -> \"PreTrainedTokenizer\":\n# from transformers import AutoTokenizer, PreTrainedTokenizer\n\n# global _INTENT_TOKENIZER\n# if _INTENT_TOKENIZER is None:\n# # The tokenizer details are not uploaded to the HF hub since it's just the\n# # unmodified distilbert tokenizer.\n# _INTENT_TOKENIZER = cast(\n# PreTrainedTokenizer,\n# AutoTokenizer.from_pretrained(\"distilbert-base-uncased\"),\n# )\n# return _INTENT_TOKENIZER\n\n\n# def get_local_intent_model(\n# model_name_or_path: str = INTENT_MODEL_VERSION,\n# tag: str | None = INTENT_MODEL_TAG,\n# ) -> HybridClassifier:\n# global _INTENT_MODEL\n# if _INTENT_MODEL is None:\n# try:\n# # Calculate where the cache should be, then load from local if available\n# logger.notice(f\"Loading model from local cache: {model_name_or_path}\")\n# local_path = snapshot_download(\n# repo_id=model_name_or_path, revision=tag, local_files_only=True\n# )\n# _INTENT_MODEL = HybridClassifier.from_pretrained(local_path)\n# logger.notice(f\"Loaded model from local cache: {local_path}\")\n# except Exception as e:\n# logger.warning(f\"Failed to load model directly: {e}\")\n# try:\n# # Attempt to download the model snapshot\n# logger.notice(f\"Downloading model snapshot for {model_name_or_path}\")\n# local_path = snapshot_download(\n# repo_id=model_name_or_path, revision=tag, local_files_only=False\n# )\n# _INTENT_MODEL = HybridClassifier.from_pretrained(local_path)\n# except Exception as e:\n# logger.error(\n# f\"Failed to load model even after attempted snapshot download: {e}\"\n# )\n# raise\n# return _INTENT_MODEL\n\n\n# def get_local_information_content_model(\n# model_name_or_path: str = INFORMATION_CONTENT_MODEL_VERSION,\n# tag: str | None = INFORMATION_CONTENT_MODEL_TAG,\n# ) -> \"SetFitModel\":\n# from setfit import SetFitModel\n\n# global _INFORMATION_CONTENT_MODEL\n# if _INFORMATION_CONTENT_MODEL is None:\n# try:\n# # Calculate where the cache should be, then load from local if available\n# logger.notice(\n# f\"Loading content information model from local cache: {model_name_or_path}\"\n# )\n# local_path = snapshot_download(\n# repo_id=model_name_or_path, revision=tag, local_files_only=True\n# )\n# _INFORMATION_CONTENT_MODEL = SetFitModel.from_pretrained(local_path)\n# logger.notice(\n# f\"Loaded content information model from local cache: {local_path}\"\n# )\n# except Exception as e:\n# logger.warning(f\"Failed to load content information model directly: {e}\")\n# try:\n# # Attempt to download the model snapshot\n# logger.notice(\n# f\"Downloading content information model snapshot for {model_name_or_path}\"\n# )\n# local_path = snapshot_download(\n# repo_id=model_name_or_path, revision=tag, local_files_only=False\n# )\n# _INFORMATION_CONTENT_MODEL = SetFitModel.from_pretrained(local_path)\n# except Exception as e:\n# logger.error(\n# f\"Failed to load content information model even after attempted snapshot download: {e}\"\n# )\n# raise\n\n# return _INFORMATION_CONTENT_MODEL\n\n\n# def tokenize_connector_classification_query(\n# connectors: list[str],\n# query: str,\n# tokenizer: \"PreTrainedTokenizer\",\n# connector_token_end_id: int,\n# ) -> tuple[torch.Tensor, torch.Tensor]:\n# \"\"\"\n# Tokenize the connectors & user query into one prompt for the forward pass of ConnectorClassifier models\n\n# The attention mask is just all 1s. The prompt is CLS + each connector name suffixed with the connector end\n# token and then the user query.\n# \"\"\"\n\n# input_ids = torch.tensor([tokenizer.cls_token_id], dtype=torch.long)\n\n# for connector in connectors:\n# connector_token_ids = tokenizer(\n# connector,\n# add_special_tokens=False,\n# return_tensors=\"pt\",\n# )\n\n# input_ids = torch.cat(\n# (\n# input_ids,\n# connector_token_ids[\"input_ids\"].squeeze(dim=0),\n# torch.tensor([connector_token_end_id], dtype=torch.long),\n# ),\n# dim=-1,\n# )\n# query_token_ids = tokenizer(\n# query,\n# add_special_tokens=False,\n# return_tensors=\"pt\",\n# )\n\n# input_ids = torch.cat(\n# (\n# input_ids,\n# query_token_ids[\"input_ids\"].squeeze(dim=0),\n# torch.tensor([tokenizer.sep_token_id], dtype=torch.long),\n# ),\n# dim=-1,\n# )\n# attention_mask = torch.ones(input_ids.numel(), dtype=torch.long)\n\n# return input_ids.unsqueeze(0), attention_mask.unsqueeze(0)\n\n\n# def warm_up_connector_classifier_model() -> None:\n# logger.info(\n# f\"Warming up connector_classifier model {CONNECTOR_CLASSIFIER_MODEL_TAG}\"\n# )\n# connector_classifier_tokenizer = get_connector_classifier_tokenizer()\n# connector_classifier = get_local_connector_classifier()\n\n# input_ids, attention_mask = tokenize_connector_classification_query(\n# [\"GitHub\"],\n# \"onyx classifier query google doc\",\n# connector_classifier_tokenizer,\n# connector_classifier.connector_end_token_id,\n# )\n# input_ids = input_ids.to(connector_classifier.device)\n# attention_mask = attention_mask.to(connector_classifier.device)\n\n# connector_classifier(input_ids, attention_mask)\n\n\n# def warm_up_intent_model() -> None:\n# logger.notice(f\"Warming up Intent Model: {INTENT_MODEL_VERSION}\")\n# intent_tokenizer = get_intent_model_tokenizer()\n# tokens = intent_tokenizer(\n# MODEL_WARM_UP_STRING, return_tensors=\"pt\", truncation=True, padding=True\n# )\n\n# intent_model = get_local_intent_model()\n# device = intent_model.device\n# intent_model(\n# query_ids=tokens[\"input_ids\"].to(device),\n# query_mask=tokens[\"attention_mask\"].to(device),\n# )\n\n\n# def warm_up_information_content_model() -> None:\n# logger.notice(\"Warming up Content Model\") # TODO: add version if needed\n\n# information_content_model = get_local_information_content_model()\n# information_content_model(INFORMATION_CONTENT_MODEL_WARM_UP_STRING)\n\n\n# @simple_log_function_time()\n# def run_inference(tokens: \"BatchEncoding\") -> tuple[list[float], list[float]]:\n# intent_model = get_local_intent_model()\n# device = intent_model.device\n\n# outputs = intent_model(\n# query_ids=tokens[\"input_ids\"].to(device),\n# query_mask=tokens[\"attention_mask\"].to(device),\n# )\n\n# token_logits = outputs[\"token_logits\"]\n# intent_logits = outputs[\"intent_logits\"]\n\n# # Move tensors to CPU before applying softmax and converting to numpy\n# intent_probabilities = F.softmax(intent_logits.cpu(), dim=-1).numpy()[0]\n# token_probabilities = F.softmax(token_logits.cpu(), dim=-1).numpy()[0]\n\n# # Extract the probabilities for the positive class (index 1) for each token\n# token_positive_probs = token_probabilities[:, 1].tolist()\n\n# return intent_probabilities.tolist(), token_positive_probs\n\n\n# @simple_log_function_time()\n# def run_content_classification_inference(\n# text_inputs: list[str],\n# ) -> list[ContentClassificationPrediction]:\n# \"\"\"\n# Assign a score to the segments in question. The model stored in get_local_information_content_model()\n# creates the 'model score' based on its training, and the scores are then converted to a 0.0-1.0 scale.\n# In the code outside of the model/inference model servers that score will be converted into the actual\n# boost factor.\n# \"\"\"\n\n# def _prob_to_score(prob: float) -> float:\n# \"\"\"\n# Conversion of base score to 0.0 - 1.0 score. Note that the min/max values depend on the model!\n# \"\"\"\n# _MIN_BASE_SCORE = 0.25\n# _MAX_BASE_SCORE = 0.75\n# if prob < _MIN_BASE_SCORE:\n# raw_score = 0.0\n# elif prob < _MAX_BASE_SCORE:\n# raw_score = (prob - _MIN_BASE_SCORE) / (_MAX_BASE_SCORE - _MIN_BASE_SCORE)\n# else:\n# raw_score = 1.0\n# return (\n# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN\n# + (\n# INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MAX\n# - INDEXING_INFORMATION_CONTENT_CLASSIFICATION_MIN\n# )\n# * raw_score\n# )\n\n# _BATCH_SIZE = 32\n# content_model = get_local_information_content_model()\n\n# # Process inputs in batches\n# all_output_classes: list[int] = []\n# all_base_output_probabilities: list[float] = []\n\n# for i in range(0, len(text_inputs), _BATCH_SIZE):\n# batch = text_inputs[i : i + _BATCH_SIZE]\n# batch_with_prefix = []\n# batch_indices = []\n\n# # Pre-allocate results for this batch\n# batch_output_classes: list[np.ndarray] = [np.array(1)] * len(batch)\n# batch_probabilities: list[np.ndarray] = [np.array(1.0)] * len(batch)\n\n# # Pre-process batch to handle long input exceptions\n# for j, text in enumerate(batch):\n# if len(text) == 0:\n# # if no input, treat as non-informative from the model's perspective\n# batch_output_classes[j] = np.array(0)\n# batch_probabilities[j] = np.array(0.0)\n# logger.warning(\"Input for Content Information Model is empty\")\n\n# elif (\n# len(text.split())\n# <= INDEXING_INFORMATION_CONTENT_CLASSIFICATION_CUTOFF_LENGTH\n# ):\n# # if input is short, use the model\n# batch_with_prefix.append(\n# _INFORMATION_CONTENT_MODEL_PROMPT_PREFIX + text\n# )\n# batch_indices.append(j)\n# else:\n# # if longer than cutoff, treat as informative (stay with default), but issue warning\n# logger.warning(\"Input for Content Information Model too long\")\n\n# if batch_with_prefix: # Only run model if we have valid inputs\n# # Get predictions for the batch\n# model_output_classes = content_model(batch_with_prefix)\n# model_output_probabilities = content_model.predict_proba(batch_with_prefix)\n\n# # Place results in the correct positions\n# for idx, batch_idx in enumerate(batch_indices):\n# batch_output_classes[batch_idx] = model_output_classes[idx].numpy()\n# batch_probabilities[batch_idx] = model_output_probabilities[idx][\n# 1\n# ].numpy() # x[1] is prob of the positive class\n\n# all_output_classes.extend([int(x) for x in batch_output_classes])\n# all_base_output_probabilities.extend([float(x) for x in batch_probabilities])\n\n# logits = [\n# np.log(p / (1 - p)) if p != 0.0 and p != 1.0 else (100 if p == 1.0 else -100)\n# for p in all_base_output_probabilities\n# ]\n# scaled_logits = [\n# logit / INDEXING_INFORMATION_CONTENT_CLASSIFICATION_TEMPERATURE\n# for logit in logits\n# ]\n# output_probabilities_with_temp = [\n# np.exp(scaled_logit) / (1 + np.exp(scaled_logit))\n# for scaled_logit in scaled_logits\n# ]\n\n# prediction_scores = [\n# _prob_to_score(p_temp) for p_temp in output_probabilities_with_temp\n# ]\n\n# content_classification_predictions = [\n# ContentClassificationPrediction(\n# predicted_label=predicted_label, content_boost_factor=output_score\n# )\n# for predicted_label, output_score in zip(all_output_classes, prediction_scores)\n# ]\n\n# return content_classification_predictions\n\n\n# def map_keywords(\n# input_ids: torch.Tensor, tokenizer: \"PreTrainedTokenizer\", is_keyword: list[bool]\n# ) -> list[str]:\n# tokens = tokenizer.convert_ids_to_tokens(input_ids) # type: ignore\n\n# if not len(tokens) == len(is_keyword):\n# raise ValueError(\"Length of tokens and keyword predictions must match\")\n\n# if input_ids[0] == tokenizer.cls_token_id:\n# tokens = tokens[1:]\n# is_keyword = is_keyword[1:]\n\n# if input_ids[-1] == tokenizer.sep_token_id:\n# tokens = tokens[:-1]\n# is_keyword = is_keyword[:-1]\n\n# unk_token = tokenizer.unk_token\n# if unk_token in tokens:\n# raise ValueError(\"Unknown token detected in the input\")\n\n# keywords = []\n# current_keyword = \"\"\n\n# for ind, token in enumerate(tokens):\n# if is_keyword[ind]:\n# if token.startswith(\"##\"):\n# current_keyword += token[2:]\n# else:\n# if current_keyword:\n# keywords.append(current_keyword)\n# current_keyword = token\n# else:\n# # If mispredicted a later token of a keyword, add it to the current keyword\n# # to complete it\n# if current_keyword:\n# if len(current_keyword) > 2 and current_keyword.startswith(\"##\"):\n# current_keyword = current_keyword[2:]\n\n# else:\n# keywords.append(current_keyword)\n# current_keyword = \"\"\n\n# if current_keyword:\n# keywords.append(current_keyword)\n\n# return keywords\n\n\n# def clean_keywords(keywords: list[str]) -> list[str]:\n# cleaned_words = []\n# for word in keywords:\n# word = word[:-2] if word.endswith(\"'s\") else word\n# word = word.replace(\"/\", \" \")\n# word = word.replace(\"'\", \"\").replace('\"', \"\")\n# cleaned_words.extend([w for w in word.strip().split() if w and not w.isspace()])\n# return cleaned_words\n\n\n# def run_connector_classification(req: ConnectorClassificationRequest) -> list[str]:\n# tokenizer = get_connector_classifier_tokenizer()\n# model = get_local_connector_classifier()\n\n# connector_names = req.available_connectors\n\n# input_ids, attention_mask = tokenize_connector_classification_query(\n# connector_names,\n# req.query,\n# tokenizer,\n# model.connector_end_token_id,\n# )\n# input_ids = input_ids.to(model.device)\n# attention_mask = attention_mask.to(model.device)\n\n# global_confidence, classifier_confidence = model(input_ids, attention_mask)\n\n# if global_confidence.item() < 0.5:\n# return []\n\n# passed_connectors = []\n\n# for i, connector_name in enumerate(connector_names):\n# if classifier_confidence.view(-1)[i].item() > 0.5:\n# passed_connectors.append(connector_name)\n\n# return passed_connectors\n\n\n# def run_analysis(intent_req: IntentRequest) -> tuple[bool, list[str]]:\n# tokenizer = get_intent_model_tokenizer()\n# model_input = tokenizer(\n# intent_req.query, return_tensors=\"pt\", truncation=False, padding=False\n# )\n\n# if len(model_input.input_ids[0]) > 512:\n# # If the user text is too long, assume it is semantic and keep all words\n# return True, intent_req.query.split()\n\n# intent_probs, token_probs = run_inference(model_input)\n\n# is_keyword_sequence = intent_probs[0] >= intent_req.keyword_percent_threshold\n\n# keyword_preds = [\n# token_prob >= intent_req.keyword_percent_threshold for token_prob in token_probs\n# ]\n\n# try:\n# keywords = map_keywords(model_input.input_ids[0], tokenizer, keyword_preds)\n# except Exception as e:\n# logger.warning(\n# f\"Failed to extract keywords for query: {intent_req.query} due to {e}\"\n# )\n# # Fallback to keeping all words\n# keywords = intent_req.query.split()\n\n# cleaned_keywords = clean_keywords(keywords)\n\n# return is_keyword_sequence, cleaned_keywords\n\n\n# @router.post(\"/connector-classification\")\n# async def process_connector_classification_request(\n# classification_request: ConnectorClassificationRequest,\n# ) -> ConnectorClassificationResponse:\n# if INDEXING_ONLY:\n# raise RuntimeError(\n# \"Indexing model server should not call connector classification endpoint\"\n# )\n\n# if len(classification_request.available_connectors) == 0:\n# return ConnectorClassificationResponse(connectors=[])\n\n# connectors = run_connector_classification(classification_request)\n# return ConnectorClassificationResponse(connectors=connectors)\n\n\n# @router.post(\"/query-analysis\")\n# async def process_analysis_request(\n# intent_request: IntentRequest,\n# ) -> IntentResponse:\n# if INDEXING_ONLY:\n# raise RuntimeError(\"Indexing model server should not call intent endpoint\")\n\n# is_keyword, keywords = run_analysis(intent_request)\n# return IntentResponse(is_keyword=is_keyword, keywords=keywords)\n\n\n# @router.post(\"/content-classification\")\n# async def process_content_classification_request(\n# content_classification_requests: list[str],\n# ) -> list[ContentClassificationPrediction]:\n# return run_content_classification_inference(content_classification_requests)\n" }, { "path": "backend/model_server/legacy/onyx_torch_model.py", "content": "# import json\n# import os\n# from typing import cast\n# from typing import TYPE_CHECKING\n\n# import torch\n# import torch.nn as nn\n\n\n# if TYPE_CHECKING:\n# from transformers import DistilBertConfig\n\n\n# class HybridClassifier(nn.Module):\n# def __init__(self) -> None:\n# from transformers import DistilBertConfig, DistilBertModel\n\n# super().__init__()\n# config = DistilBertConfig()\n# self.distilbert = DistilBertModel(config)\n# config = self.distilbert.config # type: ignore\n\n# # Keyword tokenwise binary classification layer\n# self.keyword_classifier = nn.Linear(config.dim, 2)\n\n# # Intent Classifier layers\n# self.pre_classifier = nn.Linear(config.dim, config.dim)\n# self.intent_classifier = nn.Linear(config.dim, 2)\n\n# self.device = torch.device(\"cpu\")\n\n# def forward(\n# self,\n# query_ids: torch.Tensor,\n# query_mask: torch.Tensor,\n# ) -> dict[str, torch.Tensor]:\n# outputs = self.distilbert(input_ids=query_ids, attention_mask=query_mask)\n# sequence_output = outputs.last_hidden_state\n\n# # Intent classification on the CLS token\n# cls_token_state = sequence_output[:, 0, :]\n# pre_classifier_out = self.pre_classifier(cls_token_state)\n# intent_logits = self.intent_classifier(pre_classifier_out)\n\n# # Keyword classification on all tokens\n# token_logits = self.keyword_classifier(sequence_output)\n\n# return {\"intent_logits\": intent_logits, \"token_logits\": token_logits}\n\n# @classmethod\n# def from_pretrained(cls, load_directory: str) -> \"HybridClassifier\":\n# model_path = os.path.join(load_directory, \"pytorch_model.bin\")\n# config_path = os.path.join(load_directory, \"config.json\")\n\n# with open(config_path, \"r\") as f:\n# config = json.load(f)\n# model = cls(**config)\n\n# if torch.backends.mps.is_available():\n# # Apple silicon GPU\n# device = torch.device(\"mps\")\n# elif torch.cuda.is_available():\n# device = torch.device(\"cuda\")\n# else:\n# device = torch.device(\"cpu\")\n\n# model.load_state_dict(torch.load(model_path, map_location=device))\n# model = model.to(device)\n\n# model.device = device\n\n# model.eval()\n# # Eval doesn't set requires_grad to False, do it manually to save memory and have faster inference\n# for param in model.parameters():\n# param.requires_grad = False\n\n# return model\n\n\n# class ConnectorClassifier(nn.Module):\n# def __init__(self, config: \"DistilBertConfig\") -> None:\n# from transformers import DistilBertTokenizer, DistilBertModel\n\n# super().__init__()\n\n# self.config = config\n# self.distilbert = DistilBertModel(config)\n# config = self.distilbert.config # type: ignore\n# self.connector_global_classifier = nn.Linear(config.dim, 1)\n# self.connector_match_classifier = nn.Linear(config.dim, 1)\n# self.tokenizer = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")\n\n# # Token indicating end of connector name, and on which classifier is used\n# self.connector_end_token_id = self.tokenizer.get_vocab()[\n# self.config.connector_end_token\n# ]\n\n# self.device = torch.device(\"cpu\")\n\n# def forward(\n# self,\n# input_ids: torch.Tensor,\n# attention_mask: torch.Tensor,\n# ) -> tuple[torch.Tensor, torch.Tensor]:\n# hidden_states = self.distilbert(\n# input_ids=input_ids, attention_mask=attention_mask\n# ).last_hidden_state\n\n# cls_hidden_states = hidden_states[\n# :, 0, :\n# ] # Take leap of faith that first token is always [CLS]\n# global_logits = self.connector_global_classifier(cls_hidden_states).view(-1)\n# global_confidence = torch.sigmoid(global_logits).view(-1)\n\n# connector_end_position_ids = input_ids == self.connector_end_token_id\n# connector_end_hidden_states = hidden_states[connector_end_position_ids]\n# classifier_output = self.connector_match_classifier(connector_end_hidden_states)\n# classifier_confidence = torch.nn.functional.sigmoid(classifier_output).view(-1)\n\n# return global_confidence, classifier_confidence\n\n# @classmethod\n# def from_pretrained(cls, repo_dir: str) -> \"ConnectorClassifier\":\n# from transformers import DistilBertConfig\n\n# config = cast(\n# DistilBertConfig,\n# DistilBertConfig.from_pretrained(os.path.join(repo_dir, \"config.json\")),\n# )\n# device = (\n# torch.device(\"cuda\")\n# if torch.cuda.is_available()\n# else (\n# torch.device(\"mps\")\n# if torch.backends.mps.is_available()\n# else torch.device(\"cpu\")\n# )\n# )\n# state_dict = torch.load(\n# os.path.join(repo_dir, \"pytorch_model.pt\"),\n# map_location=device,\n# weights_only=True,\n# )\n\n# model = cls(config)\n# model.load_state_dict(state_dict)\n# model.to(device)\n# model.device = device\n# model.eval()\n\n# for param in model.parameters():\n# param.requires_grad = False\n\n# return model\n" }, { "path": "backend/model_server/legacy/reranker.py", "content": "# import asyncio\n# from typing import Optional\n# from typing import TYPE_CHECKING\n\n# from fastapi import APIRouter\n# from fastapi import HTTPException\n\n# from model_server.utils import simple_log_function_time\n# from onyx.utils.logger import setup_logger\n# from shared_configs.configs import INDEXING_ONLY\n# from shared_configs.model_server_models import RerankRequest\n# from shared_configs.model_server_models import RerankResponse\n\n# if TYPE_CHECKING:\n# from sentence_transformers import CrossEncoder\n\n# logger = setup_logger()\n\n# router = APIRouter(prefix=\"/encoder\")\n\n# _RERANK_MODEL: Optional[\"CrossEncoder\"] = None\n\n\n# def get_local_reranking_model(\n# model_name: str,\n# ) -> \"CrossEncoder\":\n# global _RERANK_MODEL\n# from sentence_transformers import CrossEncoder\n\n# if _RERANK_MODEL is None:\n# logger.notice(f\"Loading {model_name}\")\n# model = CrossEncoder(model_name)\n# _RERANK_MODEL = model\n# return _RERANK_MODEL\n\n\n# @simple_log_function_time()\n# async def local_rerank(query: str, docs: list[str], model_name: str) -> list[float]:\n# cross_encoder = get_local_reranking_model(model_name)\n# # Run CPU-bound reranking in a thread pool\n# return await asyncio.get_event_loop().run_in_executor(\n# None,\n# lambda: cross_encoder.predict([(query, doc) for doc in docs]).tolist(),\n# )\n\n\n# @router.post(\"/cross-encoder-scores\")\n# async def process_rerank_request(rerank_request: RerankRequest) -> RerankResponse:\n# \"\"\"Cross encoders can be purely black box from the app perspective\"\"\"\n# # Only local models should use this endpoint - API providers should make direct API calls\n# if rerank_request.provider_type is not None:\n# raise ValueError(\n# f\"Model server reranking endpoint should only be used for local models. \"\n# f\"API provider '{rerank_request.provider_type}' should make direct API calls instead.\"\n# )\n\n# if INDEXING_ONLY:\n# raise RuntimeError(\"Indexing model server should not call reranking endpoint\")\n\n# if not rerank_request.documents or not rerank_request.query:\n# raise HTTPException(\n# status_code=400, detail=\"Missing documents or query for reranking\"\n# )\n# if not all(rerank_request.documents):\n# raise ValueError(\"Empty documents cannot be reranked.\")\n\n# try:\n# # At this point, provider_type is None, so handle local reranking\n# sim_scores = await local_rerank(\n# query=rerank_request.query,\n# docs=rerank_request.documents,\n# model_name=rerank_request.model_name,\n# )\n# return RerankResponse(scores=sim_scores)\n\n# except Exception as e:\n# logger.exception(f\"Error during reranking process:\\n{str(e)}\")\n# raise HTTPException(\n# status_code=500, detail=\"Failed to run Cross-Encoder reranking\"\n# )\n" }, { "path": "backend/model_server/main.py", "content": "import logging\nimport os\nimport shutil\nfrom collections.abc import AsyncGenerator\nfrom contextlib import asynccontextmanager\nfrom pathlib import Path\n\nimport sentry_sdk\nimport torch\nimport uvicorn\nfrom fastapi import FastAPI\nfrom prometheus_fastapi_instrumentator import Instrumentator\nfrom sentry_sdk.integrations.fastapi import FastApiIntegration\nfrom sentry_sdk.integrations.starlette import StarletteIntegration\nfrom transformers import logging as transformer_logging\n\nfrom model_server.encoders import router as encoders_router\nfrom model_server.management_endpoints import router as management_router\nfrom model_server.utils import get_gpu_type\nfrom onyx import __version__\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.logger import setup_uvicorn_logger\nfrom onyx.utils.middleware import add_onyx_request_id_middleware\nfrom onyx.utils.middleware import add_onyx_tenant_id_middleware\nfrom shared_configs.configs import INDEXING_ONLY\nfrom shared_configs.configs import MIN_THREADS_ML_MODELS\nfrom shared_configs.configs import MODEL_SERVER_ALLOWED_HOST\nfrom shared_configs.configs import MODEL_SERVER_PORT\nfrom shared_configs.configs import SENTRY_DSN\n\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\nos.environ[\"HF_HUB_DISABLE_TELEMETRY\"] = \"1\"\n\nHF_CACHE_PATH = Path(\".cache/huggingface\")\nTEMP_HF_CACHE_PATH = Path(\".cache/temp_huggingface\")\n\ntransformer_logging.set_verbosity_error()\n\nlogger = setup_logger()\n\nfile_handlers = [\n h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)\n]\n\nsetup_uvicorn_logger(shared_file_handlers=file_handlers)\n\n\ndef _move_files_recursively(source: Path, dest: Path, overwrite: bool = False) -> None:\n \"\"\"\n This moves the files from the temp huggingface cache to the huggingface cache\n\n We have to move each file individually because the directories might\n have the same name but not the same contents and we dont want to remove\n the files in the existing huggingface cache that don't exist in the temp\n huggingface cache.\n \"\"\"\n\n for item in source.iterdir():\n target_path = dest / item.relative_to(source)\n if item.is_dir():\n _move_files_recursively(item, target_path, overwrite)\n else:\n target_path.parent.mkdir(parents=True, exist_ok=True)\n if target_path.exists() and not overwrite:\n continue\n shutil.move(str(item), str(target_path))\n\n\n@asynccontextmanager\nasync def lifespan(app: FastAPI) -> AsyncGenerator:\n gpu_type = get_gpu_type()\n logger.notice(f\"Torch GPU Detection: gpu_type={gpu_type}\")\n\n app.state.gpu_type = gpu_type\n\n try:\n if TEMP_HF_CACHE_PATH.is_dir():\n logger.notice(\"Moving contents of temp_huggingface to huggingface cache.\")\n _move_files_recursively(TEMP_HF_CACHE_PATH, HF_CACHE_PATH)\n shutil.rmtree(TEMP_HF_CACHE_PATH, ignore_errors=True)\n logger.notice(\"Moved contents of temp_huggingface to huggingface cache.\")\n except Exception as e:\n logger.warning(\n f\"Error moving contents of temp_huggingface to huggingface cache: {e}. \"\n \"This is not a critical error and the model server will continue to run.\"\n )\n\n torch.set_num_threads(max(MIN_THREADS_ML_MODELS, torch.get_num_threads()))\n logger.notice(f\"Torch Threads: {torch.get_num_threads()}\")\n\n yield\n\n\ndef get_model_app() -> FastAPI:\n application = FastAPI(\n title=\"Onyx Model Server\", version=__version__, lifespan=lifespan\n )\n if SENTRY_DSN:\n sentry_sdk.init(\n dsn=SENTRY_DSN,\n integrations=[StarletteIntegration(), FastApiIntegration()],\n traces_sample_rate=0.1,\n release=__version__,\n )\n logger.info(\"Sentry initialized\")\n else:\n logger.debug(\"Sentry DSN not provided, skipping Sentry initialization\")\n\n application.include_router(management_router)\n application.include_router(encoders_router)\n\n request_id_prefix = \"INF\"\n if INDEXING_ONLY:\n request_id_prefix = \"IDX\"\n\n add_onyx_tenant_id_middleware(application, logger)\n add_onyx_request_id_middleware(application, request_id_prefix, logger)\n\n # Initialize and instrument the app\n Instrumentator().instrument(application).expose(application)\n\n return application\n\n\napp = get_model_app()\n\n\nif __name__ == \"__main__\":\n logger.notice(\n f\"Starting Onyx Model Server on http://{MODEL_SERVER_ALLOWED_HOST}:{str(MODEL_SERVER_PORT)}/\"\n )\n logger.notice(f\"Model Server Version: {__version__}\")\n uvicorn.run(app, host=MODEL_SERVER_ALLOWED_HOST, port=MODEL_SERVER_PORT)\n" }, { "path": "backend/model_server/management_endpoints.py", "content": "from fastapi import APIRouter\nfrom fastapi import Response\n\nfrom model_server.constants import GPUStatus\nfrom model_server.utils import get_gpu_type\n\nrouter = APIRouter(prefix=\"/api\")\n\n\n@router.get(\"/health\")\nasync def healthcheck() -> Response:\n return Response(status_code=200)\n\n\n@router.get(\"/gpu-status\")\nasync def route_gpu_status() -> dict[str, bool | str]:\n gpu_type = get_gpu_type()\n gpu_available = gpu_type != GPUStatus.NONE\n return {\"gpu_available\": gpu_available, \"type\": gpu_type}\n" }, { "path": "backend/model_server/utils.py", "content": "import asyncio\nimport time\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\nfrom functools import wraps\nfrom typing import Any\nfrom typing import cast\nfrom typing import TypeVar\n\nimport torch\n\nfrom model_server.constants import GPUStatus\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nF = TypeVar(\"F\", bound=Callable)\nFG = TypeVar(\"FG\", bound=Callable[..., Generator | Iterator])\n\n\ndef simple_log_function_time(\n func_name: str | None = None,\n debug_only: bool = False,\n include_args: bool = False,\n) -> Callable[[F], F]:\n def decorator(func: F) -> F:\n if asyncio.iscoroutinefunction(func):\n\n @wraps(func)\n async def wrapped_async_func(*args: Any, **kwargs: Any) -> Any:\n start_time = time.time()\n result = await func(*args, **kwargs)\n elapsed_time_str = str(time.time() - start_time)\n log_name = func_name or func.__name__\n args_str = f\" args={args} kwargs={kwargs}\" if include_args else \"\"\n final_log = f\"{log_name}{args_str} took {elapsed_time_str} seconds\"\n if debug_only:\n logger.debug(final_log)\n else:\n logger.notice(final_log)\n return result\n\n return cast(F, wrapped_async_func)\n else:\n\n @wraps(func)\n def wrapped_sync_func(*args: Any, **kwargs: Any) -> Any:\n start_time = time.time()\n result = func(*args, **kwargs)\n elapsed_time_str = str(time.time() - start_time)\n log_name = func_name or func.__name__\n args_str = f\" args={args} kwargs={kwargs}\" if include_args else \"\"\n final_log = f\"{log_name}{args_str} took {elapsed_time_str} seconds\"\n if debug_only:\n logger.debug(final_log)\n else:\n logger.notice(final_log)\n return result\n\n return cast(F, wrapped_sync_func)\n\n return decorator\n\n\ndef get_gpu_type() -> str:\n if torch.cuda.is_available():\n return GPUStatus.CUDA\n if torch.backends.mps.is_available():\n return GPUStatus.MAC_MPS\n\n return GPUStatus.NONE\n" }, { "path": "backend/onyx/__init__.py", "content": "import os\n\n__version__ = os.environ.get(\"ONYX_VERSION\", \"\") or \"Development\"\n" }, { "path": "backend/onyx/access/__init__.py", "content": "" }, { "path": "backend/onyx/access/access.py", "content": "from collections.abc import Callable\nfrom typing import cast\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.access.utils import prefix_user_email\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\nfrom onyx.db.document import get_access_info_for_document\nfrom onyx.db.document import get_access_info_for_documents\nfrom onyx.db.models import User\nfrom onyx.db.models import UserFile\nfrom onyx.db.user_file import fetch_user_files_with_access_relationships\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\n\ndef _get_access_for_document(\n document_id: str,\n db_session: Session,\n) -> DocumentAccess:\n info = get_access_info_for_document(\n db_session=db_session,\n document_id=document_id,\n )\n\n doc_access = DocumentAccess.build(\n user_emails=info[1] if info and info[1] else [],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=info[2] if info else False,\n )\n\n return doc_access\n\n\ndef get_access_for_document(\n document_id: str,\n db_session: Session,\n) -> DocumentAccess:\n versioned_get_access_for_document_fn = fetch_versioned_implementation(\n \"onyx.access.access\", \"_get_access_for_document\"\n )\n return versioned_get_access_for_document_fn(document_id, db_session)\n\n\ndef get_null_document_access() -> DocumentAccess:\n return DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n is_public=False,\n external_user_emails=[],\n external_user_group_ids=[],\n )\n\n\ndef _get_access_for_documents(\n document_ids: list[str],\n db_session: Session,\n) -> dict[str, DocumentAccess]:\n document_access_info = get_access_info_for_documents(\n db_session=db_session,\n document_ids=document_ids,\n )\n doc_access = {}\n for document_id, user_emails, is_public in document_access_info:\n doc_access[document_id] = DocumentAccess.build(\n user_emails=[email for email in user_emails if email],\n # MIT version will wipe all groups and external groups on update\n user_groups=[],\n is_public=is_public,\n external_user_emails=[],\n external_user_group_ids=[],\n )\n\n # Sometimes the document has not been indexed by the indexing job yet, in those cases\n # the document does not exist and so we use least permissive. Specifically the EE version\n # checks the MIT version permissions and creates a superset. This ensures that this flow\n # does not fail even if the Document has not yet been indexed.\n for doc_id in document_ids:\n if doc_id not in doc_access:\n doc_access[doc_id] = get_null_document_access()\n return doc_access\n\n\ndef get_access_for_documents(\n document_ids: list[str],\n db_session: Session,\n) -> dict[str, DocumentAccess]:\n \"\"\"Fetches all access information for the given documents.\"\"\"\n versioned_get_access_for_documents_fn = fetch_versioned_implementation(\n \"onyx.access.access\", \"_get_access_for_documents\"\n )\n return versioned_get_access_for_documents_fn(document_ids, db_session)\n\n\ndef _get_acl_for_user(\n user: User, db_session: Session # noqa: ARG001\n) -> set[str]: # noqa: ARG001\n \"\"\"Returns a list of ACL entries that the user has access to. This is meant to be\n used downstream to filter out documents that the user does not have access to. The\n user should have access to a document if at least one entry in the document's ACL\n matches one entry in the returned set.\n\n Anonymous users only have access to public documents.\n \"\"\"\n if user.is_anonymous:\n return {PUBLIC_DOC_PAT}\n return {prefix_user_email(user.email), PUBLIC_DOC_PAT}\n\n\ndef get_acl_for_user(user: User, db_session: Session | None = None) -> set[str]:\n versioned_acl_for_user_fn = fetch_versioned_implementation(\n \"onyx.access.access\", \"_get_acl_for_user\"\n )\n return versioned_acl_for_user_fn(user, db_session)\n\n\ndef source_should_fetch_permissions_during_indexing(source: DocumentSource) -> bool:\n _source_should_fetch_permissions_during_indexing_func = cast(\n Callable[[DocumentSource], bool],\n fetch_ee_implementation_or_noop(\n \"onyx.external_permissions.sync_params\",\n \"source_should_fetch_permissions_during_indexing\",\n False,\n ),\n )\n return _source_should_fetch_permissions_during_indexing_func(source)\n\n\ndef get_access_for_user_files(\n user_file_ids: list[str],\n db_session: Session,\n) -> dict[str, DocumentAccess]:\n versioned_fn = fetch_versioned_implementation(\n \"onyx.access.access\", \"get_access_for_user_files_impl\"\n )\n return versioned_fn(user_file_ids, db_session)\n\n\ndef get_access_for_user_files_impl(\n user_file_ids: list[str],\n db_session: Session,\n) -> dict[str, DocumentAccess]:\n user_files = fetch_user_files_with_access_relationships(user_file_ids, db_session)\n return build_access_for_user_files_impl(user_files)\n\n\ndef build_access_for_user_files(\n user_files: list[UserFile],\n) -> dict[str, DocumentAccess]:\n \"\"\"Compute access from pre-loaded UserFile objects (with relationships).\n Callers must ensure UserFile.user, Persona.users, and Persona.user are\n eagerly loaded (and Persona.groups for the EE path).\"\"\"\n versioned_fn = fetch_versioned_implementation(\n \"onyx.access.access\", \"build_access_for_user_files_impl\"\n )\n return versioned_fn(user_files)\n\n\ndef build_access_for_user_files_impl(\n user_files: list[UserFile],\n) -> dict[str, DocumentAccess]:\n result: dict[str, DocumentAccess] = {}\n for user_file in user_files:\n emails, is_public = collect_user_file_access(user_file)\n result[str(user_file.id)] = DocumentAccess.build(\n user_emails=list(emails),\n user_groups=[],\n is_public=is_public,\n external_user_emails=[],\n external_user_group_ids=[],\n )\n return result\n\n\ndef collect_user_file_access(user_file: UserFile) -> tuple[set[str], bool]:\n \"\"\"Collect all user emails that should have access to this user file.\n Includes the owner plus any users who have access via shared personas.\n Returns (emails, is_public).\"\"\"\n emails: set[str] = {user_file.user.email}\n is_public = False\n for persona in user_file.assistants:\n if persona.deleted:\n continue\n if persona.is_public:\n is_public = True\n if persona.user_id is not None and persona.user:\n emails.add(persona.user.email)\n for shared_user in persona.users:\n emails.add(shared_user.email)\n return emails, is_public\n" }, { "path": "backend/onyx/access/hierarchy_access.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.db.models import User\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\n\ndef _get_user_external_group_ids(\n db_session: Session, # noqa: ARG001\n user: User, # noqa: ARG001\n) -> list[str]:\n return []\n\n\ndef get_user_external_group_ids(db_session: Session, user: User) -> list[str]:\n versioned_get_user_external_group_ids = fetch_versioned_implementation(\n \"onyx.access.hierarchy_access\", \"_get_user_external_group_ids\"\n )\n return versioned_get_user_external_group_ids(db_session, user)\n" }, { "path": "backend/onyx/access/models.py", "content": "from dataclasses import dataclass\n\nfrom onyx.access.utils import prefix_external_group\nfrom onyx.access.utils import prefix_user_email\nfrom onyx.access.utils import prefix_user_group\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\n\n\n@dataclass(frozen=True)\nclass ExternalAccess:\n # arbitrary limit to prevent excessively large permissions sets\n # not internally enforced ... the caller can check this before using the instance\n MAX_NUM_ENTRIES = 5000\n\n # Emails of external users with access to the doc externally\n external_user_emails: set[str]\n # Names or external IDs of groups with access to the doc\n external_user_group_ids: set[str]\n # Whether the document is public in the external system or Onyx\n is_public: bool\n\n def __str__(self) -> str:\n \"\"\"Prevent extremely long logs\"\"\"\n\n def truncate_set(s: set[str], max_len: int = 100) -> str:\n s_str = str(s)\n if len(s_str) > max_len:\n return f\"{s_str[:max_len]}... ({len(s)} items)\"\n return s_str\n\n return (\n f\"ExternalAccess(\"\n f\"external_user_emails={truncate_set(self.external_user_emails)}, \"\n f\"external_user_group_ids={truncate_set(self.external_user_group_ids)}, \"\n f\"is_public={self.is_public})\"\n )\n\n @property\n def num_entries(self) -> int:\n return len(self.external_user_emails) + len(self.external_user_group_ids)\n\n @classmethod\n def public(cls) -> \"ExternalAccess\":\n return cls(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n )\n\n @classmethod\n def empty(cls) -> \"ExternalAccess\":\n \"\"\"\n A helper function that returns an *empty* set of external user-emails and group-ids, and sets `is_public` to `False`.\n This effectively makes the document in question \"private\" or inaccessible to anyone else.\n\n This is especially helpful to use when you are performing permission-syncing, and some document's permissions aren't able\n to be determined (for whatever reason). Setting its `ExternalAccess` to \"private\" is a feasible fallback.\n \"\"\"\n\n return cls(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=False,\n )\n\n\n@dataclass(frozen=True)\nclass DocExternalAccess:\n \"\"\"\n This is just a class to wrap the external access and the document ID\n together. It's used for syncing document permissions to Vespa.\n \"\"\"\n\n external_access: ExternalAccess\n # The document ID\n doc_id: str\n\n def to_dict(self) -> dict:\n return {\n \"external_access\": {\n \"external_user_emails\": list(self.external_access.external_user_emails),\n \"external_user_group_ids\": list(\n self.external_access.external_user_group_ids\n ),\n \"is_public\": self.external_access.is_public,\n },\n \"doc_id\": self.doc_id,\n }\n\n @classmethod\n def from_dict(cls, data: dict) -> \"DocExternalAccess\":\n external_access = ExternalAccess(\n external_user_emails=set(\n data[\"external_access\"].get(\"external_user_emails\", [])\n ),\n external_user_group_ids=set(\n data[\"external_access\"].get(\"external_user_group_ids\", [])\n ),\n is_public=data[\"external_access\"][\"is_public\"],\n )\n return cls(\n external_access=external_access,\n doc_id=data[\"doc_id\"],\n )\n\n\n@dataclass(frozen=True)\nclass NodeExternalAccess:\n \"\"\"\n Wraps external access with a hierarchy node's raw ID.\n Used for syncing hierarchy node permissions (e.g., folder permissions).\n \"\"\"\n\n external_access: ExternalAccess\n # The raw node ID from the source system (e.g., Google Drive folder ID)\n raw_node_id: str\n # The source type (e.g., \"google_drive\")\n source: str\n\n def to_dict(self) -> dict:\n return {\n \"external_access\": {\n \"external_user_emails\": list(self.external_access.external_user_emails),\n \"external_user_group_ids\": list(\n self.external_access.external_user_group_ids\n ),\n \"is_public\": self.external_access.is_public,\n },\n \"raw_node_id\": self.raw_node_id,\n \"source\": self.source,\n }\n\n @classmethod\n def from_dict(cls, data: dict) -> \"NodeExternalAccess\":\n external_access = ExternalAccess(\n external_user_emails=set(\n data[\"external_access\"].get(\"external_user_emails\", [])\n ),\n external_user_group_ids=set(\n data[\"external_access\"].get(\"external_user_group_ids\", [])\n ),\n is_public=data[\"external_access\"][\"is_public\"],\n )\n return cls(\n external_access=external_access,\n raw_node_id=data[\"raw_node_id\"],\n source=data[\"source\"],\n )\n\n\n# Union type for elements that can have permissions synced\nElementExternalAccess = DocExternalAccess | NodeExternalAccess\n\n\n# TODO(andrei): First refactor this into a pydantic model, then get rid of\n# duplicate fields.\n@dataclass(frozen=True, init=False)\nclass DocumentAccess(ExternalAccess):\n # User emails for Onyx users, None indicates admin\n user_emails: set[str | None]\n\n # Names of user groups associated with this document\n user_groups: set[str]\n\n external_user_emails: set[str]\n external_user_group_ids: set[str]\n is_public: bool\n\n def __init__(self) -> None:\n raise TypeError(\n \"Use `DocumentAccess.build(...)` instead of creating an instance directly.\"\n )\n\n def to_acl(self) -> set[str]:\n \"\"\"Converts the access state to a set of formatted ACL strings.\n\n NOTE: When querying for documents, the supplied ACL filter strings must\n be formatted in the same way as this function.\n \"\"\"\n acl_set: set[str] = set()\n for user_email in self.user_emails:\n if user_email:\n acl_set.add(prefix_user_email(user_email))\n\n for group_name in self.user_groups:\n acl_set.add(prefix_user_group(group_name))\n\n for external_user_email in self.external_user_emails:\n acl_set.add(prefix_user_email(external_user_email))\n\n for external_group_id in self.external_user_group_ids:\n acl_set.add(prefix_external_group(external_group_id))\n\n if self.is_public:\n acl_set.add(PUBLIC_DOC_PAT)\n\n return acl_set\n\n @classmethod\n def build(\n cls,\n user_emails: list[str | None],\n user_groups: list[str],\n external_user_emails: list[str],\n external_user_group_ids: list[str],\n is_public: bool,\n ) -> \"DocumentAccess\":\n \"\"\"Don't prefix incoming data wth acl type, prefix on read from to_acl!\"\"\"\n\n obj = object.__new__(cls)\n object.__setattr__(\n obj, \"user_emails\", {user_email for user_email in user_emails if user_email}\n )\n object.__setattr__(obj, \"user_groups\", set(user_groups))\n object.__setattr__(\n obj,\n \"external_user_emails\",\n {external_email for external_email in external_user_emails},\n )\n object.__setattr__(\n obj,\n \"external_user_group_ids\",\n {external_group_id for external_group_id in external_user_group_ids},\n )\n object.__setattr__(obj, \"is_public\", is_public)\n\n return obj\n\n\ndefault_public_access = DocumentAccess.build(\n external_user_emails=[],\n external_user_group_ids=[],\n user_emails=[],\n user_groups=[],\n is_public=True,\n)\n" }, { "path": "backend/onyx/access/utils.py", "content": "from onyx.configs.constants import DocumentSource\n\n\ndef prefix_user_email(user_email: str) -> str:\n \"\"\"Prefixes a user email to eliminate collision with group names.\n This applies to both a Onyx user and an External user, this is to make the query time\n more efficient\"\"\"\n return f\"user_email:{user_email}\"\n\n\ndef prefix_user_group(user_group_name: str) -> str:\n \"\"\"Prefixes a user group name to eliminate collision with user emails.\n This assumes that user ids are prefixed with a different prefix.\"\"\"\n return f\"group:{user_group_name}\"\n\n\ndef prefix_external_group(ext_group_name: str) -> str:\n \"\"\"Prefixes an external group name to eliminate collision with user emails / Onyx groups.\"\"\"\n return f\"external_group:{ext_group_name}\"\n\n\ndef build_ext_group_name_for_onyx(ext_group_name: str, source: DocumentSource) -> str:\n \"\"\"\n External groups may collide across sources, every source needs its own prefix.\n NOTE: the name is lowercased to handle case sensitivity for group names\n \"\"\"\n return f\"{source.value}_{ext_group_name}\".lower()\n" }, { "path": "backend/onyx/auth/__init__.py", "content": "" }, { "path": "backend/onyx/auth/anonymous_user.py", "content": "from collections.abc import Mapping\nfrom typing import Any\nfrom typing import cast\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import ANONYMOUS_USER_EMAIL\nfrom onyx.configs.constants import ANONYMOUS_USER_INFO_ID\nfrom onyx.configs.constants import KV_ANONYMOUS_USER_PERSONALIZATION_KEY\nfrom onyx.configs.constants import KV_ANONYMOUS_USER_PREFERENCES_KEY\nfrom onyx.key_value_store.store import KeyValueStore\nfrom onyx.key_value_store.store import KvKeyNotFoundError\nfrom onyx.server.manage.models import UserInfo\nfrom onyx.server.manage.models import UserPersonalization\nfrom onyx.server.manage.models import UserPreferences\n\n\ndef set_anonymous_user_preferences(\n store: KeyValueStore, preferences: UserPreferences\n) -> None:\n store.store(KV_ANONYMOUS_USER_PREFERENCES_KEY, preferences.model_dump())\n\n\ndef set_anonymous_user_personalization(\n store: KeyValueStore, personalization: UserPersonalization\n) -> None:\n store.store(KV_ANONYMOUS_USER_PERSONALIZATION_KEY, personalization.model_dump())\n\n\ndef load_anonymous_user_preferences(store: KeyValueStore) -> UserPreferences:\n try:\n preferences_data = cast(\n Mapping[str, Any], store.load(KV_ANONYMOUS_USER_PREFERENCES_KEY)\n )\n return UserPreferences(**preferences_data)\n except KvKeyNotFoundError:\n return UserPreferences(\n chosen_assistants=None, default_model=None, auto_scroll=True\n )\n\n\ndef fetch_anonymous_user_info(store: KeyValueStore) -> UserInfo:\n \"\"\"Fetch a UserInfo object for anonymous users (used for API responses).\"\"\"\n personalization = UserPersonalization()\n try:\n personalization_data = cast(\n Mapping[str, Any], store.load(KV_ANONYMOUS_USER_PERSONALIZATION_KEY)\n )\n personalization = UserPersonalization(**personalization_data)\n except KvKeyNotFoundError:\n pass\n\n return UserInfo(\n id=ANONYMOUS_USER_INFO_ID,\n email=ANONYMOUS_USER_EMAIL,\n is_active=True,\n is_superuser=False,\n is_verified=True,\n role=UserRole.LIMITED,\n preferences=load_anonymous_user_preferences(store),\n personalization=personalization,\n is_anonymous_user=True,\n password_configured=False,\n )\n" }, { "path": "backend/onyx/auth/api_key.py", "content": "import hashlib\nimport secrets\nimport uuid\nfrom urllib.parse import quote\n\nfrom fastapi import Request\nfrom passlib.hash import sha256_crypt\nfrom pydantic import BaseModel\n\nfrom onyx.auth.constants import API_KEY_LENGTH\nfrom onyx.auth.constants import API_KEY_PREFIX\nfrom onyx.auth.constants import DEPRECATED_API_KEY_PREFIX\nfrom onyx.auth.schemas import UserRole\nfrom onyx.auth.utils import get_hashed_bearer_token_from_request\nfrom onyx.configs.app_configs import API_KEY_HASH_ROUNDS\nfrom shared_configs.configs import MULTI_TENANT\n\n\nclass ApiKeyDescriptor(BaseModel):\n api_key_id: int\n api_key_display: str\n api_key: str | None = None # only present on initial creation\n api_key_name: str | None = None\n api_key_role: UserRole\n\n user_id: uuid.UUID\n\n\ndef generate_api_key(tenant_id: str | None = None) -> str:\n if not MULTI_TENANT or not tenant_id:\n return API_KEY_PREFIX + secrets.token_urlsafe(API_KEY_LENGTH)\n\n encoded_tenant = quote(tenant_id) # URL encode the tenant ID\n return f\"{API_KEY_PREFIX}{encoded_tenant}.{secrets.token_urlsafe(API_KEY_LENGTH)}\"\n\n\ndef _deprecated_hash_api_key(api_key: str) -> str:\n return sha256_crypt.hash(api_key, salt=\"\", rounds=API_KEY_HASH_ROUNDS)\n\n\ndef hash_api_key(api_key: str) -> str:\n # NOTE: no salt is needed, as the API key is randomly generated\n # and overlaps are impossible\n if api_key.startswith(API_KEY_PREFIX):\n return hashlib.sha256(api_key.encode(\"utf-8\")).hexdigest()\n\n if api_key.startswith(DEPRECATED_API_KEY_PREFIX):\n return _deprecated_hash_api_key(api_key)\n\n raise ValueError(f\"Invalid API key prefix: {api_key[:3]}\")\n\n\ndef build_displayable_api_key(api_key: str) -> str:\n if api_key.startswith(API_KEY_PREFIX):\n api_key = api_key[len(API_KEY_PREFIX) :]\n\n return API_KEY_PREFIX + api_key[:4] + \"********\" + api_key[-4:]\n\n\ndef get_hashed_api_key_from_request(request: Request) -> str | None:\n \"\"\"Extract and hash API key from Authorization header.\n\n Accepts both \"Bearer \" and raw key formats.\n \"\"\"\n return get_hashed_bearer_token_from_request(\n request,\n valid_prefixes=[API_KEY_PREFIX, DEPRECATED_API_KEY_PREFIX],\n hash_fn=hash_api_key,\n allow_non_bearer=True, # API keys historically support both formats\n )\n" }, { "path": "backend/onyx/auth/captcha.py", "content": "\"\"\"Captcha verification for user registration.\"\"\"\n\nimport httpx\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.configs.app_configs import CAPTCHA_ENABLED\nfrom onyx.configs.app_configs import RECAPTCHA_SCORE_THRESHOLD\nfrom onyx.configs.app_configs import RECAPTCHA_SECRET_KEY\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nRECAPTCHA_VERIFY_URL = \"https://www.google.com/recaptcha/api/siteverify\"\n\n\nclass CaptchaVerificationError(Exception):\n \"\"\"Raised when captcha verification fails.\"\"\"\n\n\nclass RecaptchaResponse(BaseModel):\n \"\"\"Response from Google reCAPTCHA verification API.\"\"\"\n\n success: bool\n score: float | None = None # Only present for reCAPTCHA v3\n action: str | None = None\n challenge_ts: str | None = None\n hostname: str | None = None\n error_codes: list[str] | None = Field(default=None, alias=\"error-codes\")\n\n\ndef is_captcha_enabled() -> bool:\n \"\"\"Check if captcha verification is enabled.\"\"\"\n return CAPTCHA_ENABLED and bool(RECAPTCHA_SECRET_KEY)\n\n\nasync def verify_captcha_token(\n token: str,\n expected_action: str = \"signup\",\n) -> None:\n \"\"\"\n Verify a reCAPTCHA token with Google's API.\n\n Args:\n token: The reCAPTCHA response token from the client\n expected_action: Expected action name for v3 verification\n\n Raises:\n CaptchaVerificationError: If verification fails\n \"\"\"\n if not is_captcha_enabled():\n return\n\n if not token:\n raise CaptchaVerificationError(\"Captcha token is required\")\n\n try:\n async with httpx.AsyncClient() as client:\n response = await client.post(\n RECAPTCHA_VERIFY_URL,\n data={\n \"secret\": RECAPTCHA_SECRET_KEY,\n \"response\": token,\n },\n timeout=10.0,\n )\n response.raise_for_status()\n\n data = response.json()\n result = RecaptchaResponse(**data)\n\n if not result.success:\n error_codes = result.error_codes or [\"unknown-error\"]\n logger.warning(f\"Captcha verification failed: {error_codes}\")\n raise CaptchaVerificationError(\n f\"Captcha verification failed: {', '.join(error_codes)}\"\n )\n\n # For reCAPTCHA v3, also check the score\n if result.score is not None:\n if result.score < RECAPTCHA_SCORE_THRESHOLD:\n logger.warning(\n f\"Captcha score too low: {result.score} < {RECAPTCHA_SCORE_THRESHOLD}\"\n )\n raise CaptchaVerificationError(\n \"Captcha verification failed: suspicious activity detected\"\n )\n\n # Optionally verify the action matches\n if result.action and result.action != expected_action:\n logger.warning(\n f\"Captcha action mismatch: {result.action} != {expected_action}\"\n )\n raise CaptchaVerificationError(\n \"Captcha verification failed: action mismatch\"\n )\n\n logger.debug(\n f\"Captcha verification passed: score={result.score}, action={result.action}\"\n )\n\n except httpx.HTTPError as e:\n logger.error(f\"Captcha API request failed: {e}\")\n # In case of API errors, we might want to allow registration\n # to prevent blocking legitimate users. This is a policy decision.\n raise CaptchaVerificationError(\"Captcha verification service unavailable\")\n" }, { "path": "backend/onyx/auth/constants.py", "content": "\"\"\"Authentication constants shared across auth modules.\"\"\"\n\n# API Key constants\nAPI_KEY_PREFIX = \"on_\"\nDEPRECATED_API_KEY_PREFIX = \"dn_\"\nAPI_KEY_LENGTH = 192\n\n# PAT constants\nPAT_PREFIX = \"onyx_pat_\"\nPAT_LENGTH = 192\n\n# Shared header constants\nAPI_KEY_HEADER_NAME = \"Authorization\"\nAPI_KEY_HEADER_ALTERNATIVE_NAME = \"X-Onyx-Authorization\"\nBEARER_PREFIX = \"Bearer \"\n" }, { "path": "backend/onyx/auth/disposable_email_validator.py", "content": "\"\"\"\nUtility to validate and block disposable/temporary email addresses.\n\nThis module fetches a list of known disposable email domains from a remote source\nand caches them for performance. It's used during user registration to prevent\nabuse from temporary email services.\n\"\"\"\n\nimport threading\nimport time\nfrom typing import Set\n\nimport httpx\n\nfrom onyx.configs.app_configs import DISPOSABLE_EMAIL_DOMAINS_URL\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass DisposableEmailValidator:\n \"\"\"\n Thread-safe singleton validator for disposable email domains.\n\n Fetches and caches the list of disposable domains, with periodic refresh.\n \"\"\"\n\n _instance: \"DisposableEmailValidator | None\" = None\n _lock = threading.Lock()\n\n def __new__(cls) -> \"DisposableEmailValidator\":\n if cls._instance is None:\n with cls._lock:\n if cls._instance is None:\n cls._instance = super().__new__(cls)\n return cls._instance\n\n def __init__(self) -> None:\n # Check if already initialized using a try/except to avoid type issues\n try:\n if self._initialized:\n return\n except AttributeError:\n pass\n\n self._domains: Set[str] = set()\n self._last_fetch_time: float = 0\n self._fetch_lock = threading.Lock()\n # Cache for 1 hour\n self._cache_duration = 3600\n # Hardcoded fallback list of common disposable domains\n # This ensures we block at least these even if the remote fetch fails\n self._fallback_domains = {\n \"trashlify.com\",\n \"10minutemail.com\",\n \"guerrillamail.com\",\n \"mailinator.com\",\n \"tempmail.com\",\n \"chat-tempmail.com\",\n \"throwaway.email\",\n \"yopmail.com\",\n \"temp-mail.org\",\n \"getnada.com\",\n \"maildrop.cc\",\n }\n # Set initialized flag last to prevent race conditions\n self._initialized: bool = True\n\n def _should_refresh(self) -> bool:\n \"\"\"Check if the cached domains should be refreshed.\"\"\"\n return (time.time() - self._last_fetch_time) > self._cache_duration\n\n def _fetch_domains(self) -> Set[str]:\n \"\"\"\n Fetch disposable email domains from the configured URL.\n\n Returns:\n Set of domain strings (lowercased)\n \"\"\"\n if not DISPOSABLE_EMAIL_DOMAINS_URL:\n logger.debug(\"DISPOSABLE_EMAIL_DOMAINS_URL not configured\")\n return self._fallback_domains.copy()\n\n try:\n logger.info(\n f\"Fetching disposable email domains from {DISPOSABLE_EMAIL_DOMAINS_URL}\"\n )\n with httpx.Client(timeout=10.0) as client:\n response = client.get(DISPOSABLE_EMAIL_DOMAINS_URL)\n response.raise_for_status()\n\n domains_list = response.json()\n\n if not isinstance(domains_list, list):\n logger.error(\n f\"Expected list from disposable domains URL, got {type(domains_list)}\"\n )\n return self._fallback_domains.copy()\n\n # Convert all to lowercase and create set\n domains = {domain.lower().strip() for domain in domains_list if domain}\n\n # Always include fallback domains\n domains.update(self._fallback_domains)\n\n logger.info(\n f\"Successfully fetched {len(domains)} disposable email domains\"\n )\n return domains\n\n except httpx.HTTPError as e:\n logger.warning(f\"Failed to fetch disposable domains (HTTP error): {e}\")\n except Exception as e:\n logger.warning(f\"Failed to fetch disposable domains: {e}\")\n\n # On error, return fallback domains\n return self._fallback_domains.copy()\n\n def get_domains(self) -> Set[str]:\n \"\"\"\n Get the cached set of disposable email domains.\n Refreshes the cache if needed.\n\n Returns:\n Set of disposable domain strings (lowercased)\n \"\"\"\n # Fast path: return cached domains if still fresh\n if self._domains and not self._should_refresh():\n return self._domains.copy()\n\n # Slow path: need to refresh\n with self._fetch_lock:\n # Double-check after acquiring lock\n if self._domains and not self._should_refresh():\n return self._domains.copy()\n\n self._domains = self._fetch_domains()\n self._last_fetch_time = time.time()\n return self._domains.copy()\n\n def is_disposable(self, email: str) -> bool:\n \"\"\"\n Check if an email address uses a disposable domain.\n\n Args:\n email: The email address to check\n\n Returns:\n True if the email domain is disposable, False otherwise\n \"\"\"\n if not email or \"@\" not in email:\n return False\n\n parts = email.split(\"@\")\n if len(parts) != 2 or not parts[0]: # Must have user@domain with non-empty user\n return False\n\n domain = parts[1].lower().strip()\n if not domain: # Domain part must not be empty\n return False\n\n disposable_domains = self.get_domains()\n return domain in disposable_domains\n\n\n# Global singleton instance\n_validator = DisposableEmailValidator()\n\n\ndef is_disposable_email(email: str) -> bool:\n \"\"\"\n Check if an email address uses a disposable/temporary domain.\n\n This is a convenience function that uses the global validator instance.\n\n Args:\n email: The email address to check\n\n Returns:\n True if the email uses a disposable domain, False otherwise\n \"\"\"\n return _validator.is_disposable(email)\n\n\ndef refresh_disposable_domains() -> None:\n \"\"\"\n Force a refresh of the disposable domains list.\n\n This can be called manually if you want to update the list\n without waiting for the cache to expire.\n \"\"\"\n _validator._last_fetch_time = 0\n _validator.get_domains()\n" }, { "path": "backend/onyx/auth/email_utils.py", "content": "import base64\nimport smtplib\nfrom datetime import datetime\nfrom email.mime.image import MIMEImage\nfrom email.mime.multipart import MIMEMultipart\nfrom email.mime.text import MIMEText\nfrom email.utils import formatdate\nfrom email.utils import make_msgid\n\nimport sendgrid # type: ignore\nfrom sendgrid.helpers.mail import Attachment # type: ignore\nfrom sendgrid.helpers.mail import Content\nfrom sendgrid.helpers.mail import ContentId\nfrom sendgrid.helpers.mail import Disposition\nfrom sendgrid.helpers.mail import Email\nfrom sendgrid.helpers.mail import FileContent\nfrom sendgrid.helpers.mail import FileName\nfrom sendgrid.helpers.mail import FileType\nfrom sendgrid.helpers.mail import Mail\nfrom sendgrid.helpers.mail import To\n\nfrom onyx.configs.app_configs import EMAIL_CONFIGURED\nfrom onyx.configs.app_configs import EMAIL_FROM\nfrom onyx.configs.app_configs import SENDGRID_API_KEY\nfrom onyx.configs.app_configs import SMTP_PASS\nfrom onyx.configs.app_configs import SMTP_PORT\nfrom onyx.configs.app_configs import SMTP_SERVER\nfrom onyx.configs.app_configs import SMTP_USER\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import AuthType\nfrom onyx.configs.constants import ONYX_DEFAULT_APPLICATION_NAME\nfrom onyx.configs.constants import ONYX_DISCORD_URL\nfrom onyx.db.models import User\nfrom onyx.server.runtime.onyx_runtime import OnyxRuntime\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.url import add_url_params\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\nHTML_EMAIL_TEMPLATE = \"\"\"\\\n\n\n\n \n \n {title}\n \n\n\n \n \n \n \n \n \n \n \n \n \n
\n \n
\n

{heading}

\n
\n {message}\n
\n {cta_block}\n
\n © {year} {application_name}. All rights reserved.\n {community_link_fragment}\n
\n\n\n\"\"\"\n\n\ndef build_html_email(\n application_name: str | None,\n heading: str,\n message: str,\n cta_text: str | None = None,\n cta_link: str | None = None,\n) -> str:\n community_link_fragment = \"\"\n if application_name == ONYX_DEFAULT_APPLICATION_NAME:\n community_link_fragment = f'
Have questions? Join our Discord community here.'\n\n if cta_text and cta_link:\n cta_block = f'{cta_text}'\n else:\n cta_block = \"\"\n return HTML_EMAIL_TEMPLATE.format(\n application_name=application_name,\n title=heading,\n heading=heading,\n message=message,\n cta_block=cta_block,\n community_link_fragment=community_link_fragment,\n year=datetime.now().year,\n )\n\n\ndef send_email(\n user_email: str,\n subject: str,\n html_body: str,\n text_body: str,\n mail_from: str = EMAIL_FROM,\n inline_png: tuple[str, bytes] | None = None,\n) -> None:\n if not EMAIL_CONFIGURED:\n raise ValueError(\"Email is not configured.\")\n\n if SENDGRID_API_KEY:\n send_email_with_sendgrid(\n user_email, subject, html_body, text_body, mail_from, inline_png\n )\n return\n\n send_email_with_smtplib(\n user_email, subject, html_body, text_body, mail_from, inline_png\n )\n\n\ndef send_email_with_sendgrid(\n user_email: str,\n subject: str,\n html_body: str,\n text_body: str,\n mail_from: str = EMAIL_FROM,\n inline_png: tuple[str, bytes] | None = None,\n) -> None:\n from_email = Email(mail_from) if mail_from else Email(\"noreply@onyx.app\")\n to_email = To(user_email)\n\n mail = Mail(\n from_email=from_email,\n to_emails=to_email,\n subject=subject,\n plain_text_content=Content(\"text/plain\", text_body),\n )\n\n # Add HTML content\n mail.add_content(Content(\"text/html\", html_body))\n\n if inline_png:\n image_name, image_data = inline_png\n\n # Create attachment\n encoded_image = base64.b64encode(image_data).decode()\n attachment = Attachment()\n attachment.file_content = FileContent(encoded_image)\n attachment.file_name = FileName(image_name)\n attachment.file_type = FileType(\"image/png\")\n attachment.disposition = Disposition(\"inline\")\n attachment.content_id = ContentId(image_name)\n\n mail.add_attachment(attachment)\n\n # Get a JSON-ready representation of the Mail object\n mail_json = mail.get()\n\n sg = sendgrid.SendGridAPIClient(api_key=SENDGRID_API_KEY)\n response = sg.client.mail.send.post(request_body=mail_json) # can raise\n if response.status_code != 202:\n logger.warning(f\"Unexpected status code {response.status_code}\")\n\n\ndef send_email_with_smtplib(\n user_email: str,\n subject: str,\n html_body: str,\n text_body: str,\n mail_from: str = EMAIL_FROM,\n inline_png: tuple[str, bytes] | None = None,\n) -> None:\n\n # Create a multipart/alternative message - this indicates these are alternative versions of the same content\n msg = MIMEMultipart(\"alternative\")\n msg[\"Subject\"] = subject\n msg[\"To\"] = user_email\n if mail_from:\n msg[\"From\"] = mail_from\n msg[\"Date\"] = formatdate(localtime=True)\n msg[\"Message-ID\"] = make_msgid(domain=\"onyx.app\")\n\n # Add text part first (lowest priority)\n text_part = MIMEText(text_body, \"plain\")\n msg.attach(text_part)\n\n if inline_png:\n # For HTML with images, create a multipart/related container\n related = MIMEMultipart(\"related\")\n\n # Add the HTML part to the related container\n html_part = MIMEText(html_body, \"html\")\n related.attach(html_part)\n\n # Add image with proper Content-ID to the related container\n img = MIMEImage(inline_png[1], _subtype=\"png\")\n img.add_header(\"Content-ID\", f\"<{inline_png[0]}>\")\n img.add_header(\"Content-Disposition\", \"inline\", filename=inline_png[0])\n related.attach(img)\n\n # Add the related part to the message (higher priority than text)\n msg.attach(related)\n else:\n # No images, just add HTML directly (higher priority than text)\n html_part = MIMEText(html_body, \"html\")\n msg.attach(html_part)\n\n with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as s:\n s.starttls()\n s.login(SMTP_USER, SMTP_PASS)\n s.send_message(msg)\n\n\ndef send_subscription_cancellation_email(user_email: str) -> None:\n \"\"\"This is templated but isn't meaningful for whitelabeling.\"\"\"\n\n # Example usage of the reusable HTML\n try:\n load_runtime_settings_fn = fetch_versioned_implementation(\n \"onyx.server.enterprise_settings.store\", \"load_runtime_settings\"\n )\n settings = load_runtime_settings_fn()\n application_name = settings.application_name\n except ModuleNotFoundError:\n application_name = ONYX_DEFAULT_APPLICATION_NAME\n\n onyx_file = OnyxRuntime.get_emailable_logo()\n\n subject = f\"Your {application_name} Subscription Has Been Canceled\"\n heading = \"Subscription Canceled\"\n message = (\n \"

We're sorry to see you go.

\"\n \"

Your subscription has been canceled and will end on your next billing date.

\"\n \"

If you change your mind, you can always come back!

\"\n )\n cta_text = \"Renew Subscription\"\n cta_link = \"https://www.onyx.app/pricing\"\n html_content = build_html_email(\n application_name,\n heading,\n message,\n cta_text,\n cta_link,\n )\n text_content = (\n \"We're sorry to see you go.\\n\"\n \"Your subscription has been canceled and will end on your next billing date.\\n\"\n \"If you change your mind, visit https://www.onyx.app/pricing\"\n )\n send_email(\n user_email,\n subject,\n html_content,\n text_content,\n inline_png=(\"logo.png\", onyx_file.data),\n )\n\n\ndef build_user_email_invite(\n from_email: str, to_email: str, application_name: str, auth_type: AuthType\n) -> tuple[str, str]:\n heading = \"You've Been Invited!\"\n\n # the exact action taken by the user, and thus the message, depends on the auth type\n message = f\"

You have been invited by {from_email} to join an organization on {application_name}.

\"\n if auth_type == AuthType.CLOUD:\n message += (\n \"

To join the organization, please click the button below to set a password \"\n \"or login with Google and complete your registration.

\"\n )\n elif auth_type == AuthType.BASIC:\n message += \"

To join the organization, please click the button below to set a password and complete your registration.

\"\n elif auth_type == AuthType.GOOGLE_OAUTH:\n message += \"

To join the organization, please click the button below to login with Google and complete your registration.

\"\n elif auth_type == AuthType.OIDC or auth_type == AuthType.SAML:\n message += \"

To join the organization, please click the button below to complete your registration.

\"\n else:\n raise ValueError(f\"Invalid auth type: {auth_type}\")\n\n cta_text = \"Join Organization\"\n cta_link = f\"{WEB_DOMAIN}/auth/signup?email={to_email}\"\n\n html_content = build_html_email(\n application_name,\n heading,\n message,\n cta_text,\n cta_link,\n )\n\n # text content is the fallback for clients that don't support HTML\n # not as critical, so not having special cases for each auth type\n text_content = (\n f\"You have been invited by {from_email} to join an organization on {application_name}.\\n\"\n \"To join the organization, please visit the following link:\\n\"\n f\"{WEB_DOMAIN}/auth/signup?email={to_email}\\n\"\n )\n if auth_type == AuthType.CLOUD:\n text_content += \"You'll be asked to set a password or login with Google to complete your registration.\"\n\n return text_content, html_content\n\n\ndef send_user_email_invite(\n user_email: str, current_user: User, auth_type: AuthType\n) -> None:\n try:\n load_runtime_settings_fn = fetch_versioned_implementation(\n \"onyx.server.enterprise_settings.store\", \"load_runtime_settings\"\n )\n settings = load_runtime_settings_fn()\n application_name = settings.application_name\n except ModuleNotFoundError:\n application_name = ONYX_DEFAULT_APPLICATION_NAME\n\n onyx_file = OnyxRuntime.get_emailable_logo()\n\n subject = f\"Invitation to Join {application_name} Organization\"\n\n text_content, html_content = build_user_email_invite(\n current_user.email, user_email, application_name, auth_type\n )\n\n send_email(\n user_email,\n subject,\n html_content,\n text_content,\n inline_png=(\"logo.png\", onyx_file.data),\n )\n\n\ndef send_forgot_password_email(\n user_email: str,\n token: str,\n tenant_id: str,\n mail_from: str = EMAIL_FROM,\n) -> None:\n # Builds a forgot password email with or without fancy HTML\n try:\n load_runtime_settings_fn = fetch_versioned_implementation(\n \"onyx.server.enterprise_settings.store\", \"load_runtime_settings\"\n )\n settings = load_runtime_settings_fn()\n application_name = settings.application_name\n except ModuleNotFoundError:\n application_name = ONYX_DEFAULT_APPLICATION_NAME\n\n onyx_file = OnyxRuntime.get_emailable_logo()\n\n subject = f\"Reset Your {application_name} Password\"\n heading = \"Reset Your Password\"\n tenant_param = f\"&tenant={tenant_id}\" if tenant_id and MULTI_TENANT else \"\"\n message = \"

Please click the button below to reset your password. This link will expire in 24 hours.

\"\n cta_text = \"Reset Password\"\n cta_link = f\"{WEB_DOMAIN}/auth/reset-password?token={token}{tenant_param}\"\n html_content = build_html_email(\n application_name,\n heading,\n message,\n cta_text,\n cta_link,\n )\n text_content = (\n f\"Please click the following link to reset your password. This link will expire in 24 hours.\\n\"\n f\"{WEB_DOMAIN}/auth/reset-password?token={token}{tenant_param}\"\n )\n send_email(\n user_email,\n subject,\n html_content,\n text_content,\n mail_from,\n inline_png=(\"logo.png\", onyx_file.data),\n )\n\n\ndef send_user_verification_email(\n user_email: str,\n token: str,\n new_organization: bool = False,\n mail_from: str = EMAIL_FROM,\n) -> None:\n # Builds a verification email\n try:\n load_runtime_settings_fn = fetch_versioned_implementation(\n \"onyx.server.enterprise_settings.store\", \"load_runtime_settings\"\n )\n settings = load_runtime_settings_fn()\n application_name = settings.application_name\n except ModuleNotFoundError:\n application_name = ONYX_DEFAULT_APPLICATION_NAME\n\n onyx_file = OnyxRuntime.get_emailable_logo()\n\n subject = f\"{application_name} Email Verification\"\n link = f\"{WEB_DOMAIN}/auth/verify-email?token={token}\"\n if new_organization:\n link = add_url_params(link, {\"first_user\": \"true\"})\n message = (\n f\"

Click the following link to verify your email address:

{link}

\"\n )\n html_content = build_html_email(\n application_name,\n \"Verify Your Email\",\n message,\n )\n text_content = f\"Click the following link to verify your email address: {link}\"\n send_email(\n user_email,\n subject,\n html_content,\n text_content,\n mail_from,\n inline_png=(\"logo.png\", onyx_file.data),\n )\n" }, { "path": "backend/onyx/auth/invited_users.py", "content": "from typing import cast\n\nfrom onyx.configs.constants import KV_PENDING_USERS_KEY\nfrom onyx.configs.constants import KV_USER_STORE_KEY\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.utils.special_types import JSON_ro\n\n\ndef remove_user_from_invited_users(email: str) -> int:\n try:\n store = get_kv_store()\n user_emails = cast(list, store.load(KV_USER_STORE_KEY))\n remaining_users = [user for user in user_emails if user != email]\n store.store(KV_USER_STORE_KEY, cast(JSON_ro, remaining_users))\n return len(remaining_users)\n except KvKeyNotFoundError:\n return 0\n\n\ndef get_invited_users() -> list[str]:\n try:\n store = get_kv_store()\n return cast(list, store.load(KV_USER_STORE_KEY))\n except KvKeyNotFoundError:\n return list()\n\n\ndef write_invited_users(emails: list[str]) -> int:\n store = get_kv_store()\n store.store(KV_USER_STORE_KEY, cast(JSON_ro, emails))\n return len(emails)\n\n\ndef get_pending_users() -> list[str]:\n try:\n store = get_kv_store()\n return cast(list, store.load(KV_PENDING_USERS_KEY))\n except KvKeyNotFoundError:\n return list()\n\n\ndef write_pending_users(emails: list[str]) -> int:\n store = get_kv_store()\n store.store(KV_PENDING_USERS_KEY, cast(JSON_ro, emails))\n return len(emails)\n" }, { "path": "backend/onyx/auth/jwt.py", "content": "import json\nfrom enum import Enum\nfrom functools import lru_cache\nfrom typing import Any\nfrom typing import cast\n\nimport jwt\nimport requests\nfrom cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey\nfrom jwt import decode as jwt_decode\nfrom jwt import InvalidTokenError\nfrom jwt import PyJWTError\nfrom jwt.algorithms import RSAAlgorithm\n\nfrom onyx.configs.app_configs import JWT_PUBLIC_KEY_URL\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n_PUBLIC_KEY_FETCH_ATTEMPTS = 2\n\n\nclass PublicKeyFormat(Enum):\n JWKS = \"jwks\"\n PEM = \"pem\"\n\n\n@lru_cache()\ndef _fetch_public_key_payload() -> tuple[str | dict[str, Any], PublicKeyFormat] | None:\n \"\"\"Fetch and cache the raw JWT verification material.\"\"\"\n if JWT_PUBLIC_KEY_URL is None:\n logger.error(\"JWT_PUBLIC_KEY_URL is not set\")\n return None\n\n try:\n response = requests.get(JWT_PUBLIC_KEY_URL)\n response.raise_for_status()\n except requests.RequestException as exc:\n logger.error(f\"Failed to fetch JWT public key: {str(exc)}\")\n return None\n content_type = response.headers.get(\"Content-Type\", \"\").lower()\n raw_body = response.text\n body_lstripped = raw_body.lstrip()\n\n if \"application/json\" in content_type or body_lstripped.startswith(\"{\"):\n try:\n data = response.json()\n except ValueError:\n logger.error(\"JWT public key URL returned invalid JSON\")\n return None\n\n if isinstance(data, dict) and \"keys\" in data:\n return data, PublicKeyFormat.JWKS\n\n logger.error(\n \"JWT public key URL returned JSON but no JWKS 'keys' field was found\"\n )\n return None\n\n body = raw_body.strip()\n if not body:\n logger.error(\"JWT public key URL returned an empty response\")\n return None\n\n return body, PublicKeyFormat.PEM\n\n\ndef get_public_key(token: str) -> RSAPublicKey | str | None:\n \"\"\"Return the concrete public key used to verify the provided JWT token.\"\"\"\n payload = _fetch_public_key_payload()\n if payload is None:\n logger.error(\"Failed to retrieve public key payload\")\n return None\n\n key_material, key_format = payload\n\n if key_format is PublicKeyFormat.JWKS:\n jwks_data = cast(dict[str, Any], key_material)\n return _resolve_public_key_from_jwks(token, jwks_data)\n\n return cast(str, key_material)\n\n\ndef _resolve_public_key_from_jwks(\n token: str, jwks_payload: dict[str, Any]\n) -> RSAPublicKey | None:\n try:\n header = jwt.get_unverified_header(token)\n except PyJWTError as e:\n logger.error(f\"Unable to parse JWT header: {str(e)}\")\n return None\n\n keys = jwks_payload.get(\"keys\", []) if isinstance(jwks_payload, dict) else []\n if not keys:\n logger.error(\"JWKS payload did not contain any keys\")\n return None\n\n kid = header.get(\"kid\")\n thumbprint = header.get(\"x5t\")\n\n candidates = []\n if kid:\n candidates = [k for k in keys if k.get(\"kid\") == kid]\n if not candidates and thumbprint:\n candidates = [k for k in keys if k.get(\"x5t\") == thumbprint]\n if not candidates and len(keys) == 1:\n candidates = keys\n\n if not candidates:\n logger.warning(\n \"No matching JWK found for token header (kid=%s, x5t=%s)\", kid, thumbprint\n )\n return None\n\n if len(candidates) > 1:\n logger.warning(\n \"Multiple JWKs matched token header kid=%s; selecting the first occurrence\",\n kid,\n )\n\n jwk = candidates[0]\n try:\n return cast(RSAPublicKey, RSAAlgorithm.from_jwk(json.dumps(jwk)))\n except ValueError as e:\n logger.error(f\"Failed to construct RSA key from JWK: {str(e)}\")\n return None\n\n\nasync def verify_jwt_token(token: str) -> dict[str, Any] | None:\n for attempt in range(_PUBLIC_KEY_FETCH_ATTEMPTS):\n public_key = get_public_key(token)\n if public_key is None:\n logger.error(\"Unable to resolve a public key for JWT verification\")\n if attempt < _PUBLIC_KEY_FETCH_ATTEMPTS - 1:\n _fetch_public_key_payload.cache_clear()\n continue\n return None\n\n try:\n payload = jwt_decode(\n token,\n public_key,\n algorithms=[\"RS256\"],\n options={\"verify_aud\": False},\n )\n except InvalidTokenError as e:\n logger.error(f\"Invalid JWT token: {str(e)}\")\n if attempt < _PUBLIC_KEY_FETCH_ATTEMPTS - 1:\n _fetch_public_key_payload.cache_clear()\n continue\n return None\n except PyJWTError as e:\n logger.error(f\"JWT decoding error: {str(e)}\")\n if attempt < _PUBLIC_KEY_FETCH_ATTEMPTS - 1:\n _fetch_public_key_payload.cache_clear()\n continue\n return None\n\n return payload\n\n return None\n" }, { "path": "backend/onyx/auth/oauth_refresher.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\nfrom typing import List\nfrom typing import Optional\n\nimport httpx\nfrom fastapi_users.manager import BaseUserManager\nfrom sqlalchemy.ext.asyncio import AsyncSession\n\nfrom onyx.configs.app_configs import OAUTH_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_CLIENT_SECRET\nfrom onyx.configs.app_configs import TRACK_EXTERNAL_IDP_EXPIRY\nfrom onyx.db.models import OAuthAccount\nfrom onyx.db.models import User\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Standard OAuth refresh token endpoints\nREFRESH_ENDPOINTS = {\n \"google\": \"https://oauth2.googleapis.com/token\",\n}\n\n\n# NOTE: Keeping this as a utility function for potential future debugging,\n# but not using it in production code\nasync def _test_expire_oauth_token(\n user: User,\n oauth_account: OAuthAccount,\n db_session: AsyncSession, # noqa: ARG001\n user_manager: BaseUserManager[User, Any],\n expire_in_seconds: int = 10,\n) -> bool:\n \"\"\"\n Utility function for testing - Sets an OAuth token to expire in a short time\n to facilitate testing of the refresh flow.\n Not used in production code.\n \"\"\"\n try:\n new_expires_at = int(\n (datetime.now(timezone.utc).timestamp() + expire_in_seconds)\n )\n\n updated_data: Dict[str, Any] = {\"expires_at\": new_expires_at}\n\n await user_manager.user_db.update_oauth_account(\n user, cast(Any, oauth_account), updated_data\n )\n\n return True\n except Exception as e:\n logger.exception(f\"Error setting artificial expiration: {str(e)}\")\n return False\n\n\nasync def refresh_oauth_token(\n user: User,\n oauth_account: OAuthAccount,\n db_session: AsyncSession, # noqa: ARG001\n user_manager: BaseUserManager[User, Any],\n) -> bool:\n \"\"\"\n Attempt to refresh an OAuth token that's about to expire or has expired.\n Returns True if successful, False otherwise.\n \"\"\"\n if not oauth_account.refresh_token:\n logger.warning(\n f\"No refresh token available for {user.email}'s {oauth_account.oauth_name} account\"\n )\n return False\n\n provider = oauth_account.oauth_name\n if provider not in REFRESH_ENDPOINTS:\n logger.warning(f\"Refresh endpoint not configured for provider: {provider}\")\n return False\n\n try:\n logger.info(f\"Refreshing OAuth token for {user.email}'s {provider} account\")\n\n async with httpx.AsyncClient() as client:\n response = await client.post(\n REFRESH_ENDPOINTS[provider],\n data={\n \"client_id\": OAUTH_CLIENT_ID,\n \"client_secret\": OAUTH_CLIENT_SECRET,\n \"refresh_token\": oauth_account.refresh_token,\n \"grant_type\": \"refresh_token\",\n },\n headers={\"Content-Type\": \"application/x-www-form-urlencoded\"},\n )\n\n if response.status_code != 200:\n logger.error(\n f\"Failed to refresh OAuth token: Status {response.status_code}\"\n )\n return False\n\n token_data = response.json()\n\n new_access_token = token_data.get(\"access_token\")\n new_refresh_token = token_data.get(\n \"refresh_token\", oauth_account.refresh_token\n )\n expires_in = token_data.get(\"expires_in\")\n\n # Calculate new expiry time if provided\n new_expires_at: Optional[int] = None\n if expires_in:\n new_expires_at = int(\n (datetime.now(timezone.utc).timestamp() + expires_in)\n )\n\n # Update the OAuth account\n updated_data: Dict[str, Any] = {\n \"access_token\": new_access_token,\n \"refresh_token\": new_refresh_token,\n }\n\n if new_expires_at:\n updated_data[\"expires_at\"] = new_expires_at\n\n # Update oidc_expiry in user model if we're tracking it\n if TRACK_EXTERNAL_IDP_EXPIRY:\n oidc_expiry = datetime.fromtimestamp(\n new_expires_at, tz=timezone.utc\n )\n await user_manager.user_db.update(\n user, {\"oidc_expiry\": oidc_expiry}\n )\n\n # Update the OAuth account\n await user_manager.user_db.update_oauth_account(\n user, cast(Any, oauth_account), updated_data\n )\n\n logger.info(f\"Successfully refreshed OAuth token for {user.email}\")\n return True\n\n except Exception as e:\n logger.exception(f\"Error refreshing OAuth token: {str(e)}\")\n return False\n\n\nasync def check_and_refresh_oauth_tokens(\n user: User,\n db_session: AsyncSession,\n user_manager: BaseUserManager[User, Any],\n) -> None:\n \"\"\"\n Check if any OAuth tokens are expired or about to expire and refresh them.\n \"\"\"\n if not hasattr(user, \"oauth_accounts\") or not user.oauth_accounts:\n return\n\n now_timestamp = datetime.now(timezone.utc).timestamp()\n\n # Buffer time to refresh tokens before they expire (in seconds)\n buffer_seconds = 300 # 5 minutes\n\n for oauth_account in user.oauth_accounts:\n # Skip accounts without refresh tokens\n if not oauth_account.refresh_token:\n continue\n\n # If token is about to expire, refresh it\n if (\n oauth_account.expires_at\n and oauth_account.expires_at - now_timestamp < buffer_seconds\n ):\n logger.info(f\"OAuth token for {user.email} is about to expire - refreshing\")\n success = await refresh_oauth_token(\n user, oauth_account, db_session, user_manager\n )\n\n if not success:\n logger.warning(\n \"Failed to refresh OAuth token. User may need to re-authenticate.\"\n )\n\n\nasync def check_oauth_account_has_refresh_token(\n user: User, # noqa: ARG001\n oauth_account: OAuthAccount,\n) -> bool:\n \"\"\"\n Check if an OAuth account has a refresh token.\n Returns True if a refresh token exists, False otherwise.\n \"\"\"\n return bool(oauth_account.refresh_token)\n\n\nasync def get_oauth_accounts_requiring_refresh_token(user: User) -> List[OAuthAccount]:\n \"\"\"\n Returns a list of OAuth accounts for a user that are missing refresh tokens.\n These accounts will need re-authentication to get refresh tokens.\n \"\"\"\n if not hasattr(user, \"oauth_accounts\") or not user.oauth_accounts:\n return []\n\n accounts_needing_refresh = []\n for oauth_account in user.oauth_accounts:\n has_refresh_token = await check_oauth_account_has_refresh_token(\n user, oauth_account\n )\n if not has_refresh_token:\n accounts_needing_refresh.append(oauth_account)\n\n return accounts_needing_refresh\n" }, { "path": "backend/onyx/auth/oauth_token_manager.py", "content": "import time\nfrom typing import Any\nfrom urllib.parse import urlencode\nfrom uuid import UUID\n\nimport requests\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import OAuthConfig\nfrom onyx.db.models import OAuthUserToken\nfrom onyx.db.oauth_config import get_user_oauth_token\nfrom onyx.db.oauth_config import upsert_user_oauth_token\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.sensitive import SensitiveValue\n\n\nlogger = setup_logger()\n\n\nclass OAuthTokenManager:\n \"\"\"Manages OAuth token retrieval, refresh, and validation\"\"\"\n\n def __init__(self, oauth_config: OAuthConfig, user_id: UUID, db_session: Session):\n self.oauth_config = oauth_config\n self.user_id = user_id\n self.db_session = db_session\n\n def get_valid_access_token(self) -> str | None:\n \"\"\"Get valid access token, refreshing if necessary\"\"\"\n user_token = get_user_oauth_token(\n self.oauth_config.id, self.user_id, self.db_session\n )\n\n if not user_token:\n return None\n\n if not user_token.token_data:\n return None\n\n token_data = self._unwrap_token_data(user_token.token_data)\n\n # Check if token is expired\n if OAuthTokenManager.is_token_expired(token_data):\n # Try to refresh if we have a refresh token\n if \"refresh_token\" in token_data:\n try:\n return self.refresh_token(user_token)\n except Exception as e:\n logger.warning(f\"Failed to refresh token: {e}\")\n return None\n else:\n return None\n\n return token_data.get(\"access_token\")\n\n def refresh_token(self, user_token: OAuthUserToken) -> str:\n \"\"\"Refresh access token using refresh token\"\"\"\n if not user_token.token_data:\n raise ValueError(\"No token data available for refresh\")\n\n if (\n self.oauth_config.client_id is None\n or self.oauth_config.client_secret is None\n ):\n raise ValueError(\n \"OAuth client_id and client_secret are required for token refresh\"\n )\n\n token_data = self._unwrap_token_data(user_token.token_data)\n\n data: dict[str, str] = {\n \"grant_type\": \"refresh_token\",\n \"refresh_token\": token_data[\"refresh_token\"],\n \"client_id\": self._unwrap_sensitive_str(self.oauth_config.client_id),\n \"client_secret\": self._unwrap_sensitive_str(\n self.oauth_config.client_secret\n ),\n }\n response = requests.post(\n self.oauth_config.token_url,\n data=data,\n headers={\"Accept\": \"application/json\"},\n )\n response.raise_for_status()\n\n new_token_data = response.json()\n\n # Calculate expires_at if expires_in is present\n if \"expires_in\" in new_token_data:\n new_token_data[\"expires_at\"] = (\n int(time.time()) + new_token_data[\"expires_in\"]\n )\n\n # Preserve refresh_token if not returned (some providers don't return it)\n if \"refresh_token\" not in new_token_data and \"refresh_token\" in token_data:\n new_token_data[\"refresh_token\"] = token_data[\"refresh_token\"]\n\n # Update token in DB\n upsert_user_oauth_token(\n self.oauth_config.id,\n self.user_id,\n new_token_data,\n self.db_session,\n )\n\n return new_token_data[\"access_token\"]\n\n @classmethod\n def token_expiration_time(cls, token_data: dict[str, Any]) -> int | None:\n \"\"\"Get the token expiration time\"\"\"\n expires_at = token_data.get(\"expires_at\")\n if not expires_at:\n return None\n\n return expires_at\n\n @classmethod\n def is_token_expired(cls, token_data: dict[str, Any]) -> bool:\n \"\"\"Check if token is expired (with 60 second buffer)\"\"\"\n expires_at = cls.token_expiration_time(token_data)\n if not expires_at:\n return False # No expiration data, assume valid\n\n # Add 60 second buffer to avoid race conditions\n return int(time.time()) + 60 >= expires_at\n\n def exchange_code_for_token(self, code: str, redirect_uri: str) -> dict[str, Any]:\n \"\"\"Exchange authorization code for access token\"\"\"\n if (\n self.oauth_config.client_id is None\n or self.oauth_config.client_secret is None\n ):\n raise ValueError(\n \"OAuth client_id and client_secret are required for code exchange\"\n )\n\n data: dict[str, str] = {\n \"grant_type\": \"authorization_code\",\n \"code\": code,\n \"client_id\": self._unwrap_sensitive_str(self.oauth_config.client_id),\n \"client_secret\": self._unwrap_sensitive_str(\n self.oauth_config.client_secret\n ),\n \"redirect_uri\": redirect_uri,\n }\n response = requests.post(\n self.oauth_config.token_url,\n data=data,\n headers={\"Accept\": \"application/json\"},\n )\n response.raise_for_status()\n\n token_data = response.json()\n\n # Calculate expires_at if expires_in is present\n if \"expires_in\" in token_data:\n token_data[\"expires_at\"] = int(time.time()) + token_data[\"expires_in\"]\n\n return token_data\n\n @staticmethod\n def build_authorization_url(\n oauth_config: OAuthConfig, redirect_uri: str, state: str\n ) -> str:\n \"\"\"Build OAuth authorization URL\"\"\"\n if oauth_config.client_id is None:\n raise ValueError(\"OAuth client_id is required to build authorization URL\")\n\n params: dict[str, Any] = {\n \"client_id\": OAuthTokenManager._unwrap_sensitive_str(\n oauth_config.client_id\n ),\n \"redirect_uri\": redirect_uri,\n \"response_type\": \"code\",\n \"state\": state,\n }\n\n # Add scopes if configured\n if oauth_config.scopes:\n params[\"scope\"] = \" \".join(oauth_config.scopes)\n\n # Add any additional provider-specific parameters\n if oauth_config.additional_params:\n params.update(oauth_config.additional_params)\n\n # Check if URL already has query parameters\n separator = \"&\" if \"?\" in oauth_config.authorization_url else \"?\"\n\n return f\"{oauth_config.authorization_url}{separator}{urlencode(params)}\"\n\n @staticmethod\n def _unwrap_sensitive_str(value: SensitiveValue[str] | str) -> str:\n if isinstance(value, SensitiveValue):\n return value.get_value(apply_mask=False)\n return value\n\n @staticmethod\n def _unwrap_token_data(\n token_data: SensitiveValue[dict[str, Any]] | dict[str, Any],\n ) -> dict[str, Any]:\n if isinstance(token_data, SensitiveValue):\n return token_data.get_value(apply_mask=False)\n return token_data\n" }, { "path": "backend/onyx/auth/pat.py", "content": "\"\"\"Personal Access Token generation and validation.\"\"\"\n\nimport hashlib\nimport secrets\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom urllib.parse import quote\n\nfrom fastapi import Request\n\nfrom onyx.auth.constants import PAT_LENGTH\nfrom onyx.auth.constants import PAT_PREFIX\nfrom onyx.auth.utils import get_hashed_bearer_token_from_request\nfrom shared_configs.configs import MULTI_TENANT\n\n\ndef generate_pat(tenant_id: str | None = None) -> str:\n \"\"\"Generate cryptographically secure PAT.\"\"\"\n if MULTI_TENANT and tenant_id:\n encoded_tenant = quote(tenant_id)\n return f\"{PAT_PREFIX}{encoded_tenant}.{secrets.token_urlsafe(PAT_LENGTH)}\"\n return PAT_PREFIX + secrets.token_urlsafe(PAT_LENGTH)\n\n\ndef hash_pat(token: str) -> str:\n \"\"\"Hash PAT using SHA256 (no salt needed due to cryptographic randomness).\"\"\"\n return hashlib.sha256(token.encode(\"utf-8\")).hexdigest()\n\n\ndef build_displayable_pat(token: str) -> str:\n \"\"\"Create masked display version: show prefix + first 4 random chars, mask middle, show last 4.\n\n Example: onyx_pat_abc1****xyz9\n \"\"\"\n # Show first 12 chars (onyx_pat_ + 4 random chars) and last 4 chars\n return f\"{token[:12]}****{token[-4:]}\"\n\n\ndef get_hashed_pat_from_request(request: Request) -> str | None:\n \"\"\"Extract and hash PAT from Authorization header.\n\n Only accepts \"Bearer \" format (unlike API keys which support raw format).\n \"\"\"\n return get_hashed_bearer_token_from_request(\n request,\n valid_prefixes=[PAT_PREFIX],\n hash_fn=hash_pat,\n allow_non_bearer=False, # PATs require Bearer prefix\n )\n\n\ndef calculate_expiration(days: int | None) -> datetime | None:\n \"\"\"Calculate expiration at 23:59:59.999999 UTC on the target date. None = no expiration.\"\"\"\n if days is None:\n return None\n expiry_date = datetime.now(timezone.utc).date() + timedelta(days=days)\n return datetime.combine(expiry_date, datetime.max.time()).replace(\n tzinfo=timezone.utc\n )\n" }, { "path": "backend/onyx/auth/permissions.py", "content": "\"\"\"\nPermission resolution for group-based authorization.\n\nGranted permissions are stored as a JSONB column on the User table and\nloaded for free with every auth query. Implied permissions are expanded\nat read time — only directly granted permissions are persisted.\n\"\"\"\n\nfrom collections.abc import Callable\nfrom collections.abc import Coroutine\nfrom typing import Any\n\nfrom fastapi import Depends\n\nfrom onyx.auth.users import current_user\nfrom onyx.db.enums import Permission\nfrom onyx.db.models import User\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nALL_PERMISSIONS: frozenset[str] = frozenset(p.value for p in Permission)\n\n# Implication map: granted permission -> set of permissions it implies.\nIMPLIED_PERMISSIONS: dict[str, set[str]] = {\n Permission.ADD_AGENTS.value: {Permission.READ_AGENTS.value},\n Permission.MANAGE_AGENTS.value: {\n Permission.ADD_AGENTS.value,\n Permission.READ_AGENTS.value,\n },\n Permission.MANAGE_DOCUMENT_SETS.value: {\n Permission.READ_DOCUMENT_SETS.value,\n Permission.READ_CONNECTORS.value,\n },\n Permission.ADD_CONNECTORS.value: {Permission.READ_CONNECTORS.value},\n Permission.MANAGE_CONNECTORS.value: {\n Permission.ADD_CONNECTORS.value,\n Permission.READ_CONNECTORS.value,\n },\n Permission.MANAGE_USER_GROUPS.value: {\n Permission.READ_CONNECTORS.value,\n Permission.READ_DOCUMENT_SETS.value,\n Permission.READ_AGENTS.value,\n Permission.READ_USERS.value,\n },\n}\n\n\ndef resolve_effective_permissions(granted: set[str]) -> set[str]:\n \"\"\"Expand granted permissions with their implied permissions.\n\n If \"admin\" is present, returns all 19 permissions.\n \"\"\"\n if Permission.FULL_ADMIN_PANEL_ACCESS.value in granted:\n return set(ALL_PERMISSIONS)\n\n effective = set(granted)\n changed = True\n while changed:\n changed = False\n for perm in list(effective):\n implied = IMPLIED_PERMISSIONS.get(perm)\n if implied and not implied.issubset(effective):\n effective |= implied\n changed = True\n return effective\n\n\ndef get_effective_permissions(user: User) -> set[Permission]:\n \"\"\"Read granted permissions from the column and expand implied permissions.\"\"\"\n granted: set[Permission] = set()\n for p in user.effective_permissions:\n try:\n granted.add(Permission(p))\n except ValueError:\n logger.warning(f\"Skipping unknown permission '{p}' for user {user.id}\")\n if Permission.FULL_ADMIN_PANEL_ACCESS in granted:\n return set(Permission)\n expanded = resolve_effective_permissions({p.value for p in granted})\n return {Permission(p) for p in expanded}\n\n\ndef require_permission(\n required: Permission,\n) -> Callable[..., Coroutine[Any, Any, User]]:\n \"\"\"FastAPI dependency factory for permission-based access control.\n\n Usage:\n @router.get(\"/endpoint\")\n def endpoint(user: User = Depends(require_permission(Permission.MANAGE_CONNECTORS))):\n ...\n \"\"\"\n\n async def dependency(user: User = Depends(current_user)) -> User:\n effective = get_effective_permissions(user)\n\n if Permission.FULL_ADMIN_PANEL_ACCESS in effective:\n return user\n\n if required not in effective:\n raise OnyxError(\n OnyxErrorCode.INSUFFICIENT_PERMISSIONS,\n \"You do not have the required permissions for this action.\",\n )\n\n return user\n\n return dependency\n" }, { "path": "backend/onyx/auth/schemas.py", "content": "import uuid\nfrom enum import Enum\nfrom typing import Any\n\nfrom fastapi_users import schemas\nfrom typing_extensions import override\n\nfrom onyx.db.enums import AccountType\n\n\nclass UserRole(str, Enum):\n \"\"\"\n User roles\n - Basic can't perform any admin actions\n - Admin can perform all admin actions\n - Curator can perform admin actions for\n groups they are curators of\n - Global Curator can perform admin actions\n for all groups they are a member of\n - Limited can access a limited set of basic api endpoints\n - Slack are users that have used onyx via slack but dont have a web login\n - External permissioned users that have been picked up during the external permissions sync process but don't have a web login\n \"\"\"\n\n LIMITED = \"limited\"\n BASIC = \"basic\"\n ADMIN = \"admin\"\n CURATOR = \"curator\"\n GLOBAL_CURATOR = \"global_curator\"\n SLACK_USER = \"slack_user\"\n EXT_PERM_USER = \"ext_perm_user\"\n\n def is_web_login(self) -> bool:\n return self not in [\n UserRole.SLACK_USER,\n UserRole.EXT_PERM_USER,\n ]\n\n\nclass UserRead(schemas.BaseUser[uuid.UUID]):\n role: UserRole\n\n\nclass UserCreate(schemas.BaseUserCreate):\n role: UserRole = UserRole.BASIC\n account_type: AccountType = AccountType.STANDARD\n tenant_id: str | None = None\n # Captcha token for cloud signup protection (optional, only used when captcha is enabled)\n # Excluded from create_update_dict so it never reaches the DB layer\n captcha_token: str | None = None\n\n @override\n def create_update_dict(self) -> dict[str, Any]:\n d = super().create_update_dict()\n d.pop(\"captcha_token\", None)\n # Force STANDARD for self-registration; only trusted paths\n # (SCIM, API key creation) supply a different account_type directly.\n d[\"account_type\"] = AccountType.STANDARD\n return d\n\n @override\n def create_update_dict_superuser(self) -> dict[str, Any]:\n d = super().create_update_dict_superuser()\n d.pop(\"captcha_token\", None)\n d.setdefault(\"account_type\", self.account_type)\n return d\n\n\nclass UserUpdate(schemas.BaseUserUpdate):\n \"\"\"\n Role updates are not allowed through the user update endpoint for security reasons\n Role changes should be handled through a separate, admin-only process\n \"\"\"\n\n\nclass AuthBackend(str, Enum):\n REDIS = \"redis\"\n POSTGRES = \"postgres\"\n JWT = \"jwt\"\n" }, { "path": "backend/onyx/auth/users.py", "content": "import base64\nimport hashlib\nimport json\nimport os\nimport random\nimport secrets\nimport string\nimport uuid\nfrom collections.abc import AsyncGenerator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\nfrom typing import List\nfrom typing import Literal\nfrom typing import Optional\nfrom typing import Protocol\nfrom typing import Tuple\nfrom typing import TypeVar\nfrom urllib.parse import urlparse\n\nimport jwt\nfrom email_validator import EmailNotValidError\nfrom email_validator import EmailUndeliverableError\nfrom email_validator import validate_email\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi import Request\nfrom fastapi import Response\nfrom fastapi import status\nfrom fastapi import WebSocket\nfrom fastapi.responses import JSONResponse\nfrom fastapi.responses import RedirectResponse\nfrom fastapi.security import OAuth2PasswordRequestForm\nfrom fastapi_users import BaseUserManager\nfrom fastapi_users import exceptions\nfrom fastapi_users import FastAPIUsers\nfrom fastapi_users import models\nfrom fastapi_users import schemas\nfrom fastapi_users import UUIDIDMixin\nfrom fastapi_users.authentication import AuthenticationBackend\nfrom fastapi_users.authentication import CookieTransport\nfrom fastapi_users.authentication import JWTStrategy\nfrom fastapi_users.authentication import RedisStrategy\nfrom fastapi_users.authentication import Strategy\nfrom fastapi_users.authentication.strategy.db import AccessTokenDatabase\nfrom fastapi_users.authentication.strategy.db import DatabaseStrategy\nfrom fastapi_users.exceptions import UserAlreadyExists\nfrom fastapi_users.jwt import decode_jwt\nfrom fastapi_users.jwt import generate_jwt\nfrom fastapi_users.jwt import SecretType\nfrom fastapi_users.manager import UserManagerDependency\nfrom fastapi_users.openapi import OpenAPIResponseType\nfrom fastapi_users.router.common import ErrorCode\nfrom fastapi_users.router.common import ErrorModel\nfrom fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase\nfrom httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback\nfrom httpx_oauth.oauth2 import BaseOAuth2\nfrom httpx_oauth.oauth2 import GetAccessTokenError\nfrom httpx_oauth.oauth2 import OAuth2Token\nfrom pydantic import BaseModel\nfrom sqlalchemy import nulls_last\nfrom sqlalchemy import select\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.api_key import get_hashed_api_key_from_request\nfrom onyx.auth.disposable_email_validator import is_disposable_email\nfrom onyx.auth.email_utils import send_forgot_password_email\nfrom onyx.auth.email_utils import send_user_verification_email\nfrom onyx.auth.invited_users import get_invited_users\nfrom onyx.auth.invited_users import remove_user_from_invited_users\nfrom onyx.auth.jwt import verify_jwt_token\nfrom onyx.auth.pat import get_hashed_pat_from_request\nfrom onyx.auth.schemas import AuthBackend\nfrom onyx.auth.schemas import UserCreate\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.app_configs import AUTH_BACKEND\nfrom onyx.configs.app_configs import AUTH_COOKIE_EXPIRE_TIME_SECONDS\nfrom onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.app_configs import EMAIL_CONFIGURED\nfrom onyx.configs.app_configs import JWT_PUBLIC_KEY_URL\nfrom onyx.configs.app_configs import PASSWORD_MAX_LENGTH\nfrom onyx.configs.app_configs import PASSWORD_MIN_LENGTH\nfrom onyx.configs.app_configs import PASSWORD_REQUIRE_DIGIT\nfrom onyx.configs.app_configs import PASSWORD_REQUIRE_LOWERCASE\nfrom onyx.configs.app_configs import PASSWORD_REQUIRE_SPECIAL_CHAR\nfrom onyx.configs.app_configs import PASSWORD_REQUIRE_UPPERCASE\nfrom onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX\nfrom onyx.configs.app_configs import REQUIRE_EMAIL_VERIFICATION\nfrom onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS\nfrom onyx.configs.app_configs import TRACK_EXTERNAL_IDP_EXPIRY\nfrom onyx.configs.app_configs import USER_AUTH_SECRET\nfrom onyx.configs.app_configs import VALID_EMAIL_DOMAINS\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME\nfrom onyx.configs.constants import ANONYMOUS_USER_EMAIL\nfrom onyx.configs.constants import ANONYMOUS_USER_UUID\nfrom onyx.configs.constants import AuthType\nfrom onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN\nfrom onyx.configs.constants import DANSWER_API_KEY_PREFIX\nfrom onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME\nfrom onyx.configs.constants import MilestoneRecordType\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import PASSWORD_SPECIAL_CHARS\nfrom onyx.configs.constants import UNNAMED_KEY_PLACEHOLDER\nfrom onyx.db.api_key import fetch_user_for_api_key\nfrom onyx.db.auth import get_access_token_db\nfrom onyx.db.auth import get_default_admin_user_emails\nfrom onyx.db.auth import get_user_count\nfrom onyx.db.auth import get_user_db\nfrom onyx.db.auth import SQLAlchemyUserAdminDB\nfrom onyx.db.engine.async_sql_engine import get_async_session\nfrom onyx.db.engine.async_sql_engine import get_async_session_context_manager\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import AccessToken\nfrom onyx.db.models import OAuthAccount\nfrom onyx.db.models import Persona\nfrom onyx.db.models import User\nfrom onyx.db.pat import fetch_user_for_pat\nfrom onyx.db.users import assign_user_to_default_groups__no_commit\nfrom onyx.db.users import get_user_by_email\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import log_onyx_error\nfrom onyx.error_handling.exceptions import onyx_error_to_json_response\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.redis.redis_pool import get_async_redis_connection\nfrom onyx.redis.redis_pool import retrieve_ws_token_data\nfrom onyx.server.settings.store import load_settings\nfrom onyx.server.utils import BasicAuthenticationError\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import mt_cloud_alias\nfrom onyx.utils.telemetry import mt_cloud_get_anon_id\nfrom onyx.utils.telemetry import mt_cloud_identify\nfrom onyx.utils.telemetry import mt_cloud_telemetry\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\nfrom onyx.utils.timing import log_function_time\nfrom onyx.utils.url import add_url_params\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom shared_configs.configs import async_return_default_schema\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nREGISTER_INVITE_ONLY_CODE = \"REGISTER_INVITE_ONLY\"\n\n\ndef is_user_admin(user: User) -> bool:\n return user.role == UserRole.ADMIN\n\n\ndef verify_auth_setting() -> None:\n \"\"\"Log warnings for AUTH_TYPE issues.\n\n This only runs on app startup not during migrations/scripts.\n \"\"\"\n raw_auth_type = (os.environ.get(\"AUTH_TYPE\") or \"\").lower()\n\n if raw_auth_type == \"cloud\":\n raise ValueError(\n \"'cloud' is not a valid auth type for self-hosted deployments.\"\n )\n if raw_auth_type == \"disabled\":\n logger.warning(\n \"AUTH_TYPE='disabled' is no longer supported. Using 'basic' instead. Please update your configuration.\"\n )\n\n logger.notice(f\"Using Auth Type: {AUTH_TYPE.value}\")\n\n\ndef get_display_email(email: str | None, space_less: bool = False) -> str:\n if email and email.endswith(DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN):\n name = email.split(\"@\")[0]\n if name == DANSWER_API_KEY_PREFIX + UNNAMED_KEY_PLACEHOLDER:\n return \"Unnamed API Key\"\n\n if space_less:\n return name\n\n return name.replace(\"API_KEY__\", \"API Key: \")\n\n return email or \"\"\n\n\ndef generate_password() -> str:\n lowercase_letters = string.ascii_lowercase\n uppercase_letters = string.ascii_uppercase\n digits = string.digits\n special_characters = string.punctuation\n\n # Ensure at least one of each required character type\n password = [\n secrets.choice(uppercase_letters),\n secrets.choice(digits),\n secrets.choice(special_characters),\n ]\n\n # Fill the rest with a mix of characters\n remaining_length = 12 - len(password)\n all_characters = lowercase_letters + uppercase_letters + digits + special_characters\n password.extend(secrets.choice(all_characters) for _ in range(remaining_length))\n\n # Shuffle the password to randomize the position of the required characters\n random.shuffle(password)\n\n return \"\".join(password)\n\n\ndef user_needs_to_be_verified() -> bool:\n if AUTH_TYPE == AuthType.BASIC or AUTH_TYPE == AuthType.CLOUD:\n return REQUIRE_EMAIL_VERIFICATION\n\n # For other auth types, if the user is authenticated it's assumed that\n # the user is already verified via the external IDP\n return False\n\n\ndef anonymous_user_enabled(*, tenant_id: str | None = None) -> bool:\n from onyx.cache.factory import get_cache_backend\n\n cache = get_cache_backend(tenant_id=tenant_id)\n value = cache.get(OnyxRedisLocks.ANONYMOUS_USER_ENABLED)\n\n if value is None:\n return False\n\n return int(value.decode(\"utf-8\")) == 1\n\n\ndef workspace_invite_only_enabled() -> bool:\n settings = load_settings()\n return settings.invite_only_enabled\n\n\ndef verify_email_is_invited(email: str) -> None:\n if AUTH_TYPE in {AuthType.SAML, AuthType.OIDC}:\n # SSO providers manage membership; allow JIT provisioning regardless of invites\n return\n\n if not workspace_invite_only_enabled():\n return\n\n whitelist = get_invited_users()\n\n if not email:\n raise OnyxError(OnyxErrorCode.INVALID_INPUT, \"Email must be specified\")\n\n try:\n email_info = validate_email(email, check_deliverability=False)\n except EmailUndeliverableError:\n raise OnyxError(OnyxErrorCode.INVALID_INPUT, \"Email is not valid\")\n\n for email_whitelist in whitelist:\n try:\n # normalized emails are now being inserted into the db\n # we can remove this normalization on read after some time has passed\n email_info_whitelist = validate_email(\n email_whitelist, check_deliverability=False\n )\n except EmailNotValidError:\n continue\n\n # oddly, normalization does not include lowercasing the user part of the\n # email address ... which we want to allow\n if email_info.normalized.lower() == email_info_whitelist.normalized.lower():\n return\n\n raise OnyxError(\n OnyxErrorCode.UNAUTHORIZED,\n \"This workspace is invite-only. Please ask your admin to invite you.\",\n )\n\n\ndef verify_email_in_whitelist(email: str, tenant_id: str) -> None:\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n if not get_user_by_email(email, db_session):\n verify_email_is_invited(email)\n\n\ndef verify_email_domain(email: str, *, is_registration: bool = False) -> None:\n if email.count(\"@\") != 1:\n raise OnyxError(OnyxErrorCode.INVALID_INPUT, \"Email is not valid\")\n\n local_part, domain = email.split(\"@\")\n domain = domain.lower()\n local_part = local_part.lower()\n\n if AUTH_TYPE == AuthType.CLOUD:\n # Normalize googlemail.com to gmail.com (they deliver to the same inbox)\n if domain == \"googlemail.com\":\n raise OnyxError(\n OnyxErrorCode.INVALID_INPUT,\n \"Please use @gmail.com instead of @googlemail.com.\",\n )\n\n # Only block dotted Gmail on new signups — existing users must still be\n # able to sign in with the address they originally registered with.\n if is_registration and domain == \"gmail.com\" and \".\" in local_part:\n raise OnyxError(\n OnyxErrorCode.INVALID_INPUT,\n \"Gmail addresses with '.' are not allowed. Please use your base email address.\",\n )\n\n if \"+\" in local_part and domain != \"onyx.app\":\n raise OnyxError(\n OnyxErrorCode.INVALID_INPUT,\n \"Email addresses with '+' are not allowed. Please use your base email address.\",\n )\n\n # Check if email uses a disposable/temporary domain\n if is_disposable_email(email):\n raise OnyxError(\n OnyxErrorCode.INVALID_INPUT,\n \"Disposable email addresses are not allowed. Please use a permanent email address.\",\n )\n\n # Check domain whitelist if configured\n if VALID_EMAIL_DOMAINS:\n if domain not in VALID_EMAIL_DOMAINS:\n raise OnyxError(OnyxErrorCode.INVALID_INPUT, \"Email domain is not valid\")\n\n\ndef enforce_seat_limit(db_session: Session, seats_needed: int = 1) -> None:\n \"\"\"Raise HTTPException(402) if adding users would exceed the seat limit.\n\n No-op for multi-tenant or CE deployments.\n \"\"\"\n if MULTI_TENANT:\n return\n\n result = fetch_ee_implementation_or_noop(\n \"onyx.db.license\", \"check_seat_availability\", None\n )(db_session, seats_needed=seats_needed)\n\n if result is not None and not result.available:\n raise OnyxError(OnyxErrorCode.SEAT_LIMIT_EXCEEDED, result.error_message)\n\n\nclass UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):\n reset_password_token_secret = USER_AUTH_SECRET\n verification_token_secret = USER_AUTH_SECRET\n verification_token_lifetime_seconds = AUTH_COOKIE_EXPIRE_TIME_SECONDS\n user_db: SQLAlchemyUserDatabase[User, uuid.UUID]\n\n async def get_by_email(self, user_email: str) -> User:\n tenant_id = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.user_mapping\", \"get_tenant_id_for_email\", None\n )(user_email)\n async with get_async_session_context_manager(tenant_id) as db_session:\n if MULTI_TENANT:\n tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](\n db_session, User, OAuthAccount\n )\n user = await tenant_user_db.get_by_email(user_email)\n else:\n user = await self.user_db.get_by_email(user_email)\n\n if not user:\n raise exceptions.UserNotExists()\n\n return user\n\n async def create(\n self,\n user_create: schemas.UC | UserCreate,\n safe: bool = False,\n request: Optional[Request] = None,\n ) -> User:\n # Verify captcha if enabled (for cloud signup protection)\n from onyx.auth.captcha import CaptchaVerificationError\n from onyx.auth.captcha import is_captcha_enabled\n from onyx.auth.captcha import verify_captcha_token\n\n if is_captcha_enabled() and request is not None:\n # Get captcha token from request body or headers\n captcha_token = None\n if hasattr(user_create, \"captcha_token\"):\n captcha_token = getattr(user_create, \"captcha_token\", None)\n\n # Also check headers as a fallback\n if not captcha_token:\n captcha_token = request.headers.get(\"X-Captcha-Token\")\n\n try:\n await verify_captcha_token(\n captcha_token or \"\", expected_action=\"signup\"\n )\n except CaptchaVerificationError as e:\n raise OnyxError(OnyxErrorCode.INVALID_INPUT, str(e))\n\n # We verify the password here to make sure it's valid before we proceed\n await self.validate_password(\n user_create.password, cast(schemas.UC, user_create)\n )\n\n # Check for disposable emails BEFORE provisioning tenant\n # This prevents creating tenants for throwaway email addresses\n try:\n verify_email_domain(user_create.email, is_registration=True)\n except OnyxError as e:\n # Log blocked disposable email attempts\n if \"Disposable email\" in e.detail:\n domain = (\n user_create.email.split(\"@\")[-1]\n if \"@\" in user_create.email\n else \"unknown\"\n )\n logger.warning(\n f\"Blocked disposable email registration attempt: {domain}\",\n extra={\"email_domain\": domain},\n )\n raise\n\n user_count: int | None = None\n referral_source = (\n request.cookies.get(\"referral_source\", None)\n if request is not None\n else None\n )\n\n tenant_id = await fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.provisioning\",\n \"get_or_provision_tenant\",\n async_return_default_schema,\n )(\n email=user_create.email,\n referral_source=referral_source,\n request=request,\n )\n user: User\n\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n async with get_async_session_context_manager(tenant_id) as db_session:\n # Check invite list based on deployment mode\n if MULTI_TENANT:\n # Multi-tenant: Only require invite for existing tenants\n # New tenant creation (first user) doesn't require an invite\n user_count = await get_user_count()\n if user_count > 0:\n # Tenant already has users - require invite for new users\n verify_email_is_invited(user_create.email)\n else:\n # Single-tenant: Check invite list (skips if SAML/OIDC or no list configured)\n verify_email_is_invited(user_create.email)\n if MULTI_TENANT:\n tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](\n db_session, User, OAuthAccount\n )\n self.user_db = tenant_user_db\n\n if hasattr(user_create, \"role\"):\n user_create.role = UserRole.BASIC\n\n user_count = await get_user_count()\n if (\n user_count == 0\n or user_create.email in get_default_admin_user_emails()\n ):\n user_create.role = UserRole.ADMIN\n\n # Check seat availability for new users (single-tenant only)\n with get_session_with_current_tenant() as sync_db:\n existing = get_user_by_email(user_create.email, sync_db)\n if existing is None:\n enforce_seat_limit(sync_db)\n\n user_created = False\n try:\n user = await super().create(user_create, safe=safe, request=request)\n user_created = True\n except IntegrityError as error:\n # Race condition: another request created the same user after the\n # pre-insert existence check but before our commit.\n await self.user_db.session.rollback()\n logger.warning(\n \"IntegrityError while creating user %s, assuming duplicate: %s\",\n user_create.email,\n str(error),\n )\n try:\n user = await self.get_by_email(user_create.email)\n except exceptions.UserNotExists:\n # Unexpected integrity error, surface it for handling upstream.\n raise error\n\n if MULTI_TENANT:\n user_by_session = await db_session.get(User, user.id)\n if user_by_session:\n user = user_by_session\n\n if (\n user.account_type.is_web_login()\n or not isinstance(user_create, UserCreate)\n or not user_create.account_type.is_web_login()\n ):\n raise exceptions.UserAlreadyExists()\n\n # Cache id before expire — accessing attrs on an expired\n # object triggers a sync lazy-load which raises MissingGreenlet\n # in this async context.\n user_id = user.id\n self._upgrade_user_to_standard__sync(user_id, user_create)\n # Expire so the async session re-fetches the row updated by\n # the sync session above.\n self.user_db.session.expire(user)\n user = await self.user_db.get(user_id) # type: ignore[assignment]\n except exceptions.UserAlreadyExists:\n user = await self.get_by_email(user_create.email)\n\n # we must use the existing user in the session if it matches\n # the user we just got by email. Note that this only applies\n # to multi-tenant, due to the overwriting of the user_db\n if MULTI_TENANT:\n user_by_session = await db_session.get(User, user.id)\n if user_by_session:\n user = user_by_session\n\n # Handle case where user has used product outside of web and is now creating an account through web\n if (\n user.account_type.is_web_login()\n or not isinstance(user_create, UserCreate)\n or not user_create.account_type.is_web_login()\n ):\n raise exceptions.UserAlreadyExists()\n\n # Cache id before expire — accessing attrs on an expired\n # object triggers a sync lazy-load which raises MissingGreenlet\n # in this async context.\n user_id = user.id\n self._upgrade_user_to_standard__sync(user_id, user_create)\n # Expire so the async session re-fetches the row updated by\n # the sync session above.\n self.user_db.session.expire(user)\n user = await self.user_db.get(user_id) # type: ignore[assignment]\n if user_created:\n await self._assign_default_pinned_assistants(user, db_session)\n remove_user_from_invited_users(user_create.email)\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n return user\n\n async def _assign_default_pinned_assistants(\n self, user: User, db_session: AsyncSession\n ) -> None:\n if user.pinned_assistants is not None:\n return\n\n result = await db_session.execute(\n select(Persona.id)\n .where(\n Persona.is_featured.is_(True),\n Persona.is_public.is_(True),\n Persona.is_listed.is_(True),\n Persona.deleted.is_(False),\n )\n .order_by(\n nulls_last(Persona.display_priority.asc()),\n Persona.id.asc(),\n )\n )\n default_persona_ids = list(result.scalars().all())\n if not default_persona_ids:\n return\n\n await self.user_db.update(\n user,\n {\"pinned_assistants\": default_persona_ids},\n )\n user.pinned_assistants = default_persona_ids\n\n def _upgrade_user_to_standard__sync(\n self,\n user_id: uuid.UUID,\n user_create: UserCreate,\n ) -> None:\n \"\"\"Upgrade a non-web user to STANDARD and assign default groups atomically.\n\n All writes happen in a single sync transaction so neither the field\n update nor the group assignment is visible without the other.\n \"\"\"\n with get_session_with_current_tenant() as sync_db:\n sync_user = sync_db.query(User).filter(User.id == user_id).first() # type: ignore[arg-type]\n if sync_user:\n sync_user.hashed_password = self.password_helper.hash(\n user_create.password\n )\n sync_user.is_verified = user_create.is_verified or False\n sync_user.role = user_create.role\n sync_user.account_type = AccountType.STANDARD\n assign_user_to_default_groups__no_commit(\n sync_db,\n sync_user,\n is_admin=(user_create.role == UserRole.ADMIN),\n )\n sync_db.commit()\n else:\n logger.warning(\n \"User %s not found in sync session during upgrade to standard; \"\n \"skipping upgrade\",\n user_id,\n )\n\n async def validate_password(self, password: str, _: schemas.UC | models.UP) -> None:\n # Validate password according to configurable security policy (defined via environment variables)\n if len(password) < PASSWORD_MIN_LENGTH:\n raise exceptions.InvalidPasswordException(\n reason=f\"Password must be at least {PASSWORD_MIN_LENGTH} characters long.\"\n )\n if len(password) > PASSWORD_MAX_LENGTH:\n raise exceptions.InvalidPasswordException(\n reason=f\"Password must not exceed {PASSWORD_MAX_LENGTH} characters.\"\n )\n if PASSWORD_REQUIRE_UPPERCASE and not any(char.isupper() for char in password):\n raise exceptions.InvalidPasswordException(\n reason=\"Password must contain at least one uppercase letter.\"\n )\n if PASSWORD_REQUIRE_LOWERCASE and not any(char.islower() for char in password):\n raise exceptions.InvalidPasswordException(\n reason=\"Password must contain at least one lowercase letter.\"\n )\n if PASSWORD_REQUIRE_DIGIT and not any(char.isdigit() for char in password):\n raise exceptions.InvalidPasswordException(\n reason=\"Password must contain at least one number.\"\n )\n if PASSWORD_REQUIRE_SPECIAL_CHAR and not any(\n char in PASSWORD_SPECIAL_CHARS for char in password\n ):\n raise exceptions.InvalidPasswordException(\n reason=f\"Password must contain at least one special character from the following set: {PASSWORD_SPECIAL_CHARS}.\"\n )\n return\n\n @log_function_time(print_only=True)\n async def oauth_callback(\n self,\n oauth_name: str,\n access_token: str,\n account_id: str,\n account_email: str,\n expires_at: Optional[int] = None,\n refresh_token: Optional[str] = None,\n request: Optional[Request] = None,\n *,\n associate_by_email: bool = False,\n is_verified_by_default: bool = False,\n ) -> User:\n referral_source = (\n getattr(request.state, \"referral_source\", None) if request else None\n )\n\n tenant_id = await fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.provisioning\",\n \"get_or_provision_tenant\",\n async_return_default_schema,\n )(\n email=account_email,\n referral_source=referral_source,\n request=request,\n )\n\n if not tenant_id:\n raise HTTPException(status_code=401, detail=\"User not found\")\n\n # Proceed with the tenant context\n token = None\n async with get_async_session_context_manager(tenant_id) as db_session:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n verify_email_in_whitelist(account_email, tenant_id)\n verify_email_domain(account_email)\n\n # NOTE(rkuo): If this UserManager is instantiated per connection\n # should we even be doing this here?\n if MULTI_TENANT:\n tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](\n db_session, User, OAuthAccount\n )\n self.user_db = tenant_user_db\n\n oauth_account_dict = {\n \"oauth_name\": oauth_name,\n \"access_token\": access_token,\n \"account_id\": account_id,\n \"account_email\": account_email,\n \"expires_at\": expires_at,\n \"refresh_token\": refresh_token,\n }\n\n user: User | None = None\n\n try:\n # Attempt to get user by OAuth account\n user = await self.get_by_oauth_account(oauth_name, account_id)\n\n except exceptions.UserNotExists:\n try:\n # Attempt to get user by email\n user = await self.user_db.get_by_email(account_email)\n if not associate_by_email:\n raise exceptions.UserAlreadyExists()\n\n # Make sure user is not None before adding OAuth account\n if user is not None:\n user = await self.user_db.add_oauth_account(\n user, oauth_account_dict\n )\n else:\n # This shouldn't happen since get_by_email would raise UserNotExists\n # but adding as a safeguard\n raise exceptions.UserNotExists()\n\n except exceptions.UserNotExists:\n verify_email_domain(account_email, is_registration=True)\n\n # Check seat availability before creating (single-tenant only)\n with get_session_with_current_tenant() as sync_db:\n enforce_seat_limit(sync_db)\n\n password = self.password_helper.generate()\n user_dict = {\n \"email\": account_email,\n \"hashed_password\": self.password_helper.hash(password),\n \"is_verified\": is_verified_by_default,\n \"account_type\": AccountType.STANDARD,\n }\n\n user = await self.user_db.create(user_dict)\n await self.user_db.add_oauth_account(user, oauth_account_dict)\n await self._assign_default_pinned_assistants(user, db_session)\n await self.on_after_register(user, request)\n\n else:\n # User exists, update OAuth account if needed\n if user is not None: # Add explicit check\n for existing_oauth_account in user.oauth_accounts:\n if (\n existing_oauth_account.account_id == account_id\n and existing_oauth_account.oauth_name == oauth_name\n ):\n user = await self.user_db.update_oauth_account(\n user,\n # NOTE: OAuthAccount DOES implement the OAuthAccountProtocol\n # but the type checker doesn't know that :(\n existing_oauth_account, # type: ignore\n oauth_account_dict,\n )\n\n # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to\n # re-authenticate that frequently, so by default this is disabled\n if expires_at and TRACK_EXTERNAL_IDP_EXPIRY:\n oidc_expiry = datetime.fromtimestamp(expires_at, tz=timezone.utc)\n await self.user_db.update(\n user, update_dict={\"oidc_expiry\": oidc_expiry}\n )\n\n # Handle case where user has used product outside of web and is now creating an account through web\n if not user.account_type.is_web_login():\n # We must use the existing user in the session if it matches\n # the user we just got by email/oauth. Note that this only applies\n # to multi-tenant, due to the overwriting of the user_db\n if MULTI_TENANT:\n if user.id:\n user_by_session = await db_session.get(User, user.id)\n if user_by_session:\n user = user_by_session\n\n # If the user is inactive, check seat availability before\n # upgrading role — otherwise they'd become an inactive BASIC\n # user who still can't log in.\n if not user.is_active:\n with get_session_with_current_tenant() as sync_db:\n enforce_seat_limit(sync_db)\n\n # Upgrade the user and assign default groups in a single\n # transaction so neither change is visible without the other.\n was_inactive = not user.is_active\n with get_session_with_current_tenant() as sync_db:\n sync_user = sync_db.query(User).filter(User.id == user.id).first() # type: ignore[arg-type]\n if sync_user:\n sync_user.is_verified = is_verified_by_default\n sync_user.role = UserRole.BASIC\n sync_user.account_type = AccountType.STANDARD\n if was_inactive:\n sync_user.is_active = True\n assign_user_to_default_groups__no_commit(sync_db, sync_user)\n sync_db.commit()\n\n # Refresh the async user object so downstream code\n # (e.g. oidc_expiry check) sees the updated fields.\n self.user_db.session.expire(user)\n user = await self.user_db.get(user.id)\n assert user is not None\n\n # this is needed if an organization goes from `TRACK_EXTERNAL_IDP_EXPIRY=true` to `false`\n # otherwise, the oidc expiry will always be old, and the user will never be able to login\n if user.oidc_expiry is not None and not TRACK_EXTERNAL_IDP_EXPIRY:\n await self.user_db.update(user, {\"oidc_expiry\": None})\n user.oidc_expiry = None # type: ignore\n remove_user_from_invited_users(user.email)\n if token:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n return user\n\n async def on_after_login(\n self,\n user: User,\n request: Optional[Request] = None,\n response: Optional[Response] = None,\n ) -> None:\n try:\n if response and request and ANONYMOUS_USER_COOKIE_NAME in request.cookies:\n response.delete_cookie(\n ANONYMOUS_USER_COOKIE_NAME,\n # Ensure cookie deletion doesn't override other cookies by setting the same path/domain\n path=\"/\",\n domain=None,\n secure=WEB_DOMAIN.startswith(\"https\"),\n )\n logger.debug(f\"Deleted anonymous user cookie for user {user.email}\")\n except Exception:\n logger.exception(\"Error deleting anonymous user cookie\")\n\n tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()\n\n # Link the anonymous PostHog session to the identified user so that\n # pre-login session recordings and events merge into one person profile.\n if anon_id := mt_cloud_get_anon_id(request):\n mt_cloud_alias(distinct_id=str(user.id), anonymous_id=anon_id)\n\n mt_cloud_identify(\n distinct_id=str(user.id),\n properties={\"email\": user.email, \"tenant_id\": tenant_id},\n )\n\n async def on_after_register(\n self, user: User, request: Optional[Request] = None\n ) -> None:\n tenant_id = await fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.provisioning\",\n \"get_or_provision_tenant\",\n async_return_default_schema,\n )(\n email=user.email,\n request=request,\n )\n\n user_count = None\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n user_count = await get_user_count()\n logger.debug(f\"Current tenant user count: {user_count}\")\n\n # Link the anonymous PostHog session to the identified user so\n # that pre-signup session recordings merge into one person profile.\n if anon_id := mt_cloud_get_anon_id(request):\n mt_cloud_alias(distinct_id=str(user.id), anonymous_id=anon_id)\n\n # Ensure a PostHog person profile exists for this user.\n mt_cloud_identify(\n distinct_id=str(user.id),\n properties={\"email\": user.email, \"tenant_id\": tenant_id},\n )\n\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id),\n event=MilestoneRecordType.USER_SIGNED_UP,\n )\n\n if user_count == 1:\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id),\n event=MilestoneRecordType.TENANT_CREATED,\n )\n\n # Assign user to the appropriate default group (Admin or Basic).\n # Must happen inside the try block while tenant context is active,\n # otherwise get_session_with_current_tenant() targets the wrong schema.\n is_admin = user_count == 1 or user.email in get_default_admin_user_emails()\n with get_session_with_current_tenant() as db_session:\n assign_user_to_default_groups__no_commit(\n db_session, user, is_admin=is_admin\n )\n db_session.commit()\n\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n # Fetch EE PostHog functions if available\n get_marketing_posthog_cookie_name = fetch_ee_implementation_or_noop(\n module=\"onyx.utils.posthog_client\",\n attribute=\"get_marketing_posthog_cookie_name\",\n noop_return_value=None,\n )\n parse_posthog_cookie = fetch_ee_implementation_or_noop(\n module=\"onyx.utils.posthog_client\",\n attribute=\"parse_posthog_cookie\",\n noop_return_value=None,\n )\n capture_and_sync_with_alternate_posthog = fetch_ee_implementation_or_noop(\n module=\"onyx.utils.posthog_client\",\n attribute=\"capture_and_sync_with_alternate_posthog\",\n noop_return_value=None,\n )\n\n if (\n request\n and user_count is not None\n and (marketing_cookie_name := get_marketing_posthog_cookie_name())\n and (marketing_cookie_value := request.cookies.get(marketing_cookie_name))\n and (parsed_cookie := parse_posthog_cookie(marketing_cookie_value))\n ):\n marketing_anonymous_id = parsed_cookie[\"distinct_id\"]\n\n # Technically, USER_SIGNED_UP is only fired from the cloud site when\n # it is the first user in a tenant. However, it is semantically correct\n # for the marketing site and should probably be refactored for the cloud site\n # to also be semantically correct.\n properties = {\n \"email\": user.email,\n \"onyx_cloud_user_id\": str(user.id),\n \"tenant_id\": str(tenant_id) if tenant_id else None,\n \"role\": user.role.value,\n \"is_first_user\": user_count == 1,\n \"source\": \"marketing_site_signup\",\n \"conversion_timestamp\": datetime.now(timezone.utc).isoformat(),\n }\n\n # Add all other values from the marketing cookie (featureFlags, etc.)\n for key, value in parsed_cookie.items():\n if key != \"distinct_id\":\n properties.setdefault(key, value)\n\n capture_and_sync_with_alternate_posthog(\n alternate_distinct_id=marketing_anonymous_id,\n event=MilestoneRecordType.USER_SIGNED_UP,\n properties=properties,\n )\n\n logger.debug(f\"User {user.id} has registered.\")\n optional_telemetry(\n record_type=RecordType.SIGN_UP,\n data={\"action\": \"create\"},\n user_id=str(user.id),\n )\n\n async def on_after_forgot_password(\n self,\n user: User,\n token: str,\n request: Optional[Request] = None, # noqa: ARG002\n ) -> None:\n if not EMAIL_CONFIGURED:\n logger.error(\n \"Email is not configured. Please configure email in the admin panel\"\n )\n raise HTTPException(\n status.HTTP_500_INTERNAL_SERVER_ERROR,\n \"Your admin has not enabled this feature.\",\n )\n tenant_id = await fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.provisioning\",\n \"get_or_provision_tenant\",\n async_return_default_schema,\n )(email=user.email)\n\n send_forgot_password_email(user.email, tenant_id=tenant_id, token=token)\n\n async def on_after_request_verify(\n self,\n user: User,\n token: str,\n request: Optional[Request] = None, # noqa: ARG002\n ) -> None:\n verify_email_domain(user.email)\n\n logger.notice(\n f\"Verification requested for user {user.id}. Verification token: {token}\"\n )\n user_count = await get_user_count()\n send_user_verification_email(\n user.email, token, new_organization=user_count == 1\n )\n\n @log_function_time(print_only=True)\n async def authenticate(\n self, credentials: OAuth2PasswordRequestForm\n ) -> Optional[User]:\n email = credentials.username\n\n tenant_id: str | None = None\n try:\n tenant_id = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.provisioning\",\n \"get_tenant_id_for_email\",\n POSTGRES_DEFAULT_SCHEMA,\n )(\n email=email,\n )\n except Exception as e:\n logger.warning(\n f\"User attempted to login with invalid credentials: {str(e)}\"\n )\n\n if not tenant_id:\n # User not found in mapping\n self.password_helper.hash(credentials.password)\n return None\n\n # Create a tenant-specific session\n async with get_async_session_context_manager(tenant_id) as tenant_session:\n tenant_user_db: SQLAlchemyUserDatabase = SQLAlchemyUserDatabase(\n tenant_session, User\n )\n self.user_db = tenant_user_db\n\n # Proceed with authentication\n try:\n user = await self.get_by_email(email)\n\n except exceptions.UserNotExists:\n self.password_helper.hash(credentials.password)\n return None\n\n if not user.account_type.is_web_login():\n raise BasicAuthenticationError(\n detail=\"NO_WEB_LOGIN_AND_HAS_NO_PASSWORD\",\n )\n\n verified, updated_password_hash = self.password_helper.verify_and_update(\n credentials.password, user.hashed_password\n )\n if not verified:\n return None\n\n if updated_password_hash is not None:\n await self.user_db.update(\n user, {\"hashed_password\": updated_password_hash}\n )\n\n return user\n\n async def reset_password_as_admin(self, user_id: uuid.UUID) -> str:\n \"\"\"Admin-only. Generate a random password for a user and return it.\"\"\"\n user = await self.get(user_id)\n new_password = generate_password()\n await self._update(user, {\"password\": new_password})\n return new_password\n\n async def change_password_if_old_matches(\n self, user: User, old_password: str, new_password: str\n ) -> None:\n \"\"\"\n For normal users to change password if they know the old one.\n Raises 400 if old password doesn't match.\n \"\"\"\n verified, updated_password_hash = self.password_helper.verify_and_update(\n old_password, user.hashed_password\n )\n if not verified:\n # Raise some HTTPException (or your custom exception) if old password is invalid:\n from fastapi import HTTPException, status\n\n raise HTTPException(\n status_code=status.HTTP_400_BAD_REQUEST,\n detail=\"Invalid current password\",\n )\n\n # If the hash was upgraded behind the scenes, we can keep it before setting the new password:\n if updated_password_hash:\n user.hashed_password = updated_password_hash\n\n # Now apply and validate the new password\n await self._update(user, {\"password\": new_password})\n\n\nasync def get_user_manager(\n user_db: SQLAlchemyUserDatabase = Depends(get_user_db),\n) -> AsyncGenerator[UserManager, None]:\n yield UserManager(user_db)\n\n\ncookie_transport = CookieTransport(\n cookie_max_age=SESSION_EXPIRE_TIME_SECONDS,\n cookie_secure=WEB_DOMAIN.startswith(\"https\"),\n cookie_name=FASTAPI_USERS_AUTH_COOKIE_NAME,\n)\n\n\nT = TypeVar(\"T\", covariant=True)\nID = TypeVar(\"ID\", contravariant=True)\n\n\n# Protocol for strategies that support token refreshing without inheritance.\nclass RefreshableStrategy(Protocol):\n \"\"\"Protocol for authentication strategies that support token refreshing.\"\"\"\n\n async def refresh_token(self, token: Optional[str], user: Any) -> str:\n \"\"\"\n Refresh an existing token by extending its lifetime.\n Returns either the same token with extended expiration or a new token.\n \"\"\"\n ...\n\n\nclass TenantAwareRedisStrategy(RedisStrategy[User, uuid.UUID]):\n \"\"\"\n A custom strategy that fetches the actual async Redis connection inside each method.\n We do NOT pass a synchronous or \"coroutine\" redis object to the constructor.\n \"\"\"\n\n def __init__(\n self,\n lifetime_seconds: Optional[int] = SESSION_EXPIRE_TIME_SECONDS,\n key_prefix: str = REDIS_AUTH_KEY_PREFIX,\n ):\n self.lifetime_seconds = lifetime_seconds\n self.key_prefix = key_prefix\n\n async def write_token(self, user: User) -> str:\n redis = await get_async_redis_connection()\n\n tenant_id = await fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.provisioning\",\n \"get_or_provision_tenant\",\n async_return_default_schema,\n )(email=user.email)\n\n token_data = {\n \"sub\": str(user.id),\n \"tenant_id\": tenant_id,\n }\n token = secrets.token_urlsafe()\n await redis.set(\n f\"{self.key_prefix}{token}\",\n json.dumps(token_data),\n ex=self.lifetime_seconds,\n )\n return token\n\n async def read_token(\n self, token: Optional[str], user_manager: BaseUserManager[User, uuid.UUID]\n ) -> Optional[User]:\n redis = await get_async_redis_connection()\n token_data_str = await redis.get(f\"{self.key_prefix}{token}\")\n if not token_data_str:\n return None\n\n try:\n token_data = json.loads(token_data_str)\n user_id = token_data[\"sub\"]\n parsed_id = user_manager.parse_id(user_id)\n return await user_manager.get(parsed_id)\n except (exceptions.UserNotExists, exceptions.InvalidID, KeyError):\n return None\n\n async def destroy_token(self, token: str, user: User) -> None: # noqa: ARG002\n \"\"\"Properly delete the token from async redis.\"\"\"\n redis = await get_async_redis_connection()\n await redis.delete(f\"{self.key_prefix}{token}\")\n\n async def refresh_token(self, token: Optional[str], user: User) -> str:\n \"\"\"Refresh a token by extending its expiration time in Redis.\"\"\"\n if token is None:\n # If no token provided, create a new one\n return await self.write_token(user)\n\n redis = await get_async_redis_connection()\n token_key = f\"{self.key_prefix}{token}\"\n\n # Check if token exists\n token_data_str = await redis.get(token_key)\n if not token_data_str:\n # Token not found, create new one\n return await self.write_token(user)\n\n # Token exists, extend its lifetime\n token_data = json.loads(token_data_str)\n await redis.set(\n token_key,\n json.dumps(token_data),\n ex=self.lifetime_seconds,\n )\n\n return token\n\n\nclass RefreshableDatabaseStrategy(DatabaseStrategy[User, uuid.UUID, AccessToken]):\n \"\"\"Database strategy with token refreshing capabilities.\"\"\"\n\n def __init__(\n self,\n access_token_db: AccessTokenDatabase[AccessToken],\n lifetime_seconds: Optional[int] = None,\n ):\n super().__init__(access_token_db, lifetime_seconds)\n self._access_token_db = access_token_db\n\n async def refresh_token(self, token: Optional[str], user: User) -> str:\n \"\"\"Refresh a token by updating its expiration time in the database.\"\"\"\n if token is None:\n return await self.write_token(user)\n\n # Find the token in database\n access_token = await self._access_token_db.get_by_token(token)\n\n if access_token is None:\n # Token not found, create new one\n return await self.write_token(user)\n\n # Update expiration time\n new_expires = datetime.now(timezone.utc) + timedelta(\n seconds=float(self.lifetime_seconds or SESSION_EXPIRE_TIME_SECONDS)\n )\n await self._access_token_db.update(access_token, {\"expires\": new_expires})\n\n return token\n\n\nclass SingleTenantJWTStrategy(JWTStrategy[User, uuid.UUID]):\n \"\"\"Stateless JWT strategy for single-tenant deployments.\n\n Tokens are self-contained and verified via signature — no Redis or DB\n lookup required per request. An ``iat`` claim is embedded so that\n downstream code can determine when the token was created without\n querying an external store.\n\n Refresh is implemented by issuing a brand-new JWT (the old one remains\n valid until its natural expiry). ``destroy_token`` is a no-op because\n JWTs cannot be server-side invalidated.\n \"\"\"\n\n def __init__(\n self,\n secret: SecretType,\n lifetime_seconds: int | None = SESSION_EXPIRE_TIME_SECONDS,\n token_audience: list[str] | None = None,\n algorithm: str = \"HS256\",\n public_key: SecretType | None = None,\n ):\n super().__init__(\n secret=secret,\n lifetime_seconds=lifetime_seconds,\n token_audience=token_audience or [\"fastapi-users:auth\"],\n algorithm=algorithm,\n public_key=public_key,\n )\n\n async def write_token(self, user: User) -> str:\n data = {\n \"sub\": str(user.id),\n \"aud\": self.token_audience,\n \"iat\": int(datetime.now(timezone.utc).timestamp()),\n }\n return generate_jwt(\n data, self.encode_key, self.lifetime_seconds, algorithm=self.algorithm\n )\n\n async def destroy_token(self, token: str, user: User) -> None: # noqa: ARG002\n # JWTs are stateless — nothing to invalidate server-side.\n # NOTE: a compromise that makes JWT auth stateful but revocable\n # is to include a token_version claim in the JWT payload. The token_version\n # is incremented whenever the user logs out (or gets login revoked). Whenever\n # the JWT is used, it is only valid if the token_version claim is the same as the one\n # in the db. If not, the JWT is invalid and the user needs to login again.\n return\n\n async def refresh_token(\n self,\n token: Optional[str], # noqa: ARG002\n user: User, # noqa: ARG002\n ) -> str:\n \"\"\"Issue a fresh JWT with a new expiry.\"\"\"\n return await self.write_token(user)\n\n\ndef get_redis_strategy() -> TenantAwareRedisStrategy:\n return TenantAwareRedisStrategy()\n\n\ndef get_database_strategy(\n access_token_db: AccessTokenDatabase[AccessToken] = Depends(get_access_token_db),\n) -> RefreshableDatabaseStrategy:\n return RefreshableDatabaseStrategy(\n access_token_db, lifetime_seconds=SESSION_EXPIRE_TIME_SECONDS\n )\n\n\ndef get_jwt_strategy() -> SingleTenantJWTStrategy:\n return SingleTenantJWTStrategy(\n secret=USER_AUTH_SECRET,\n lifetime_seconds=SESSION_EXPIRE_TIME_SECONDS,\n )\n\n\nif AUTH_BACKEND == AuthBackend.JWT:\n if MULTI_TENANT or AUTH_TYPE == AuthType.CLOUD:\n raise ValueError(\n \"JWT auth backend is only supported for single-tenant, self-hosted deployments. Use 'redis' or 'postgres' instead.\"\n )\n if not USER_AUTH_SECRET:\n raise ValueError(\"USER_AUTH_SECRET is required for JWT auth backend.\")\n\nif AUTH_BACKEND == AuthBackend.REDIS:\n auth_backend = AuthenticationBackend(\n name=\"redis\", transport=cookie_transport, get_strategy=get_redis_strategy\n )\nelif AUTH_BACKEND == AuthBackend.POSTGRES:\n auth_backend = AuthenticationBackend(\n name=\"postgres\", transport=cookie_transport, get_strategy=get_database_strategy\n )\nelif AUTH_BACKEND == AuthBackend.JWT:\n auth_backend = AuthenticationBackend(\n name=\"jwt\", transport=cookie_transport, get_strategy=get_jwt_strategy\n )\nelse:\n raise ValueError(f\"Invalid auth backend: {AUTH_BACKEND}\")\n\n\nclass FastAPIUserWithLogoutRouter(FastAPIUsers[models.UP, models.ID]):\n def get_logout_router(\n self,\n backend: AuthenticationBackend,\n requires_verification: bool = REQUIRE_EMAIL_VERIFICATION,\n ) -> APIRouter:\n \"\"\"\n Provide a router for logout only for OAuth/OIDC Flows.\n This way the login router does not need to be included\n \"\"\"\n router = APIRouter()\n\n get_current_user_token = self.authenticator.current_user_token(\n active=True, verified=requires_verification\n )\n\n logout_responses: OpenAPIResponseType = {\n **{\n status.HTTP_401_UNAUTHORIZED: {\n \"description\": \"Missing token or inactive user.\"\n }\n },\n **backend.transport.get_openapi_logout_responses_success(),\n }\n\n @router.post(\n \"/logout\", name=f\"auth:{backend.name}.logout\", responses=logout_responses\n )\n async def logout(\n user_token: Tuple[models.UP, str] = Depends(get_current_user_token),\n strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),\n ) -> Response:\n user, token = user_token\n return await backend.logout(strategy, user, token)\n\n return router\n\n def get_refresh_router(\n self,\n backend: AuthenticationBackend,\n requires_verification: bool = REQUIRE_EMAIL_VERIFICATION,\n ) -> APIRouter:\n \"\"\"\n Provide a router for session token refreshing.\n \"\"\"\n # Import the oauth_refresher here to avoid circular imports\n from onyx.auth.oauth_refresher import check_and_refresh_oauth_tokens\n\n router = APIRouter()\n\n get_current_user_token = self.authenticator.current_user_token(\n active=True, verified=requires_verification\n )\n\n refresh_responses: OpenAPIResponseType = {\n **{\n status.HTTP_401_UNAUTHORIZED: {\n \"description\": \"Missing token or inactive user.\"\n }\n },\n **backend.transport.get_openapi_login_responses_success(),\n }\n\n @router.post(\n \"/refresh\", name=f\"auth:{backend.name}.refresh\", responses=refresh_responses\n )\n async def refresh(\n user_token: Tuple[models.UP, str] = Depends(get_current_user_token),\n strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),\n user_manager: BaseUserManager[models.UP, models.ID] = Depends(\n get_user_manager\n ),\n db_session: AsyncSession = Depends(get_async_session),\n ) -> Response:\n try:\n user, token = user_token\n logger.info(f\"Processing token refresh request for user {user.email}\")\n\n # Check if user has OAuth accounts that need refreshing\n await check_and_refresh_oauth_tokens(\n user=cast(User, user),\n db_session=db_session,\n user_manager=cast(Any, user_manager),\n )\n\n # Check if strategy supports refreshing\n supports_refresh = hasattr(strategy, \"refresh_token\") and callable(\n getattr(strategy, \"refresh_token\")\n )\n\n if supports_refresh:\n try:\n refresh_method = getattr(strategy, \"refresh_token\")\n new_token = await refresh_method(token, user)\n logger.info(\n f\"Successfully refreshed session token for user {user.email}\"\n )\n return await backend.transport.get_login_response(new_token)\n except Exception as e:\n logger.error(f\"Error refreshing session token: {str(e)}\")\n # Fallback to logout and login if refresh fails\n await backend.logout(strategy, user, token)\n return await backend.login(strategy, user)\n\n # Fallback: logout and login again\n logger.info(\n \"Strategy doesn't support refresh - using logout/login flow\"\n )\n await backend.logout(strategy, user, token)\n return await backend.login(strategy, user)\n except Exception as e:\n logger.error(f\"Unexpected error in refresh endpoint: {str(e)}\")\n raise HTTPException(\n status_code=status.HTTP_400_BAD_REQUEST,\n detail=f\"Token refresh failed: {str(e)}\",\n )\n\n return router\n\n\nfastapi_users = FastAPIUserWithLogoutRouter[User, uuid.UUID](\n get_user_manager, [auth_backend]\n)\n\n\n# NOTE: verified=REQUIRE_EMAIL_VERIFICATION is not used here since we\n# take care of that in `double_check_user` ourself. This is needed, since\n# we want the /me endpoint to still return a user even if they are not\n# yet verified, so that the frontend knows they exist\noptional_fastapi_current_user = fastapi_users.current_user(active=True, optional=True)\n\n\n_JWT_EMAIL_CLAIM_KEYS = (\"email\", \"preferred_username\", \"upn\")\n\n\ndef _extract_email_from_jwt(payload: dict[str, Any]) -> str | None:\n \"\"\"Return the best-effort email/username from a decoded JWT payload.\"\"\"\n for key in _JWT_EMAIL_CLAIM_KEYS:\n value = payload.get(key)\n if isinstance(value, str) and value:\n try:\n email_info = validate_email(value, check_deliverability=False)\n except EmailNotValidError:\n continue\n normalized_email = email_info.normalized or email_info.email\n return normalized_email.lower()\n return None\n\n\nasync def _sync_jwt_oidc_expiry(\n user_manager: UserManager, user: User, payload: dict[str, Any]\n) -> None:\n if TRACK_EXTERNAL_IDP_EXPIRY:\n expires_at = payload.get(\"exp\")\n if expires_at is None:\n return\n try:\n expiry_timestamp = int(expires_at)\n except (TypeError, ValueError):\n logger.warning(\"Invalid exp claim on JWT for user %s\", user.email)\n return\n\n oidc_expiry = datetime.fromtimestamp(expiry_timestamp, tz=timezone.utc)\n if user.oidc_expiry == oidc_expiry:\n return\n\n await user_manager.user_db.update(user, {\"oidc_expiry\": oidc_expiry})\n user.oidc_expiry = oidc_expiry\n return\n\n if user.oidc_expiry is not None:\n await user_manager.user_db.update(user, {\"oidc_expiry\": None})\n user.oidc_expiry = None # type: ignore\n\n\nasync def _get_or_create_user_from_jwt(\n payload: dict[str, Any],\n request: Request,\n async_db_session: AsyncSession,\n) -> User | None:\n email = _extract_email_from_jwt(payload)\n if email is None:\n logger.warning(\n \"JWT token decoded successfully but no email claim found; skipping auth\"\n )\n return None\n\n # Enforce the same allowlist/domain policies as other auth flows\n verify_email_is_invited(email)\n verify_email_domain(email)\n\n user_db: SQLAlchemyUserAdminDB[User, uuid.UUID] = SQLAlchemyUserAdminDB(\n async_db_session, User, OAuthAccount\n )\n user_manager = UserManager(user_db)\n\n try:\n user = await user_manager.get_by_email(email)\n if not user.is_active:\n logger.warning(\"Inactive user %s attempted JWT login; skipping\", email)\n return None\n if not user.account_type.is_web_login():\n raise exceptions.UserNotExists()\n except exceptions.UserNotExists:\n logger.info(\"Provisioning user %s from JWT login\", email)\n try:\n user = await user_manager.create(\n UserCreate(\n email=email,\n password=generate_password(),\n is_verified=True,\n ),\n request=request,\n )\n except exceptions.UserAlreadyExists:\n user = await user_manager.get_by_email(email)\n if not user.is_active:\n logger.warning(\n \"Inactive user %s attempted JWT login during provisioning race; skipping\",\n email,\n )\n return None\n if not user.account_type.is_web_login():\n logger.warning(\n \"Non-web-login user %s attempted JWT login during provisioning race; skipping\",\n email,\n )\n return None\n\n await _sync_jwt_oidc_expiry(user_manager, user, payload)\n return user\n\n\nasync def _check_for_saml_and_jwt(\n request: Request,\n user: User | None,\n async_db_session: AsyncSession,\n) -> User | None:\n # If user is None, check for JWT in Authorization header\n if user is None and JWT_PUBLIC_KEY_URL is not None:\n auth_header = request.headers.get(\"Authorization\")\n if auth_header and auth_header.startswith(\"Bearer \"):\n token = auth_header[len(\"Bearer \") :].strip()\n payload = await verify_jwt_token(token)\n if payload is not None:\n user = await _get_or_create_user_from_jwt(\n payload, request, async_db_session\n )\n\n return user\n\n\nasync def optional_user(\n request: Request,\n async_db_session: AsyncSession = Depends(get_async_session),\n user: User | None = Depends(optional_fastapi_current_user),\n) -> User | None:\n\n if user := await _check_for_saml_and_jwt(request, user, async_db_session):\n # If user is already set, _check_for_saml_and_jwt returns the same user object\n return user\n\n try:\n if hashed_pat := get_hashed_pat_from_request(request):\n user = await fetch_user_for_pat(hashed_pat, async_db_session)\n elif hashed_api_key := get_hashed_api_key_from_request(request):\n user = await fetch_user_for_api_key(hashed_api_key, async_db_session)\n except ValueError:\n logger.warning(\"Issue with validating authentication token\")\n return None\n\n return user\n\n\ndef get_anonymous_user() -> User:\n \"\"\"Create anonymous user object.\"\"\"\n user = User(\n id=uuid.UUID(ANONYMOUS_USER_UUID),\n email=ANONYMOUS_USER_EMAIL,\n hashed_password=\"\",\n is_active=True,\n is_verified=True,\n is_superuser=False,\n role=UserRole.LIMITED,\n account_type=AccountType.ANONYMOUS,\n use_memories=False,\n enable_memory_tool=False,\n )\n return user\n\n\nasync def double_check_user(\n user: User | None,\n include_expired: bool = False,\n allow_anonymous_access: bool = False,\n) -> User:\n if user is not None:\n # If user attempted to authenticate, verify them, do not default\n # to anonymous access if it fails.\n if user_needs_to_be_verified() and not user.is_verified:\n raise BasicAuthenticationError(\n detail=\"Access denied. User is not verified.\",\n )\n\n if (\n user.oidc_expiry\n and user.oidc_expiry < datetime.now(timezone.utc)\n and not include_expired\n ):\n raise BasicAuthenticationError(\n detail=\"Access denied. User's OIDC token has expired.\",\n )\n\n return user\n\n if allow_anonymous_access:\n return get_anonymous_user()\n\n raise BasicAuthenticationError(\n detail=\"Access denied. User is not authenticated.\",\n )\n\n\nasync def current_user_with_expired_token(\n user: User | None = Depends(optional_user),\n) -> User:\n return await double_check_user(user, include_expired=True)\n\n\nasync def current_limited_user(\n user: User | None = Depends(optional_user),\n) -> User:\n return await double_check_user(user)\n\n\nasync def current_chat_accessible_user(\n user: User | None = Depends(optional_user),\n) -> User:\n tenant_id = get_current_tenant_id()\n\n return await double_check_user(\n user, allow_anonymous_access=anonymous_user_enabled(tenant_id=tenant_id)\n )\n\n\nasync def current_user(\n user: User | None = Depends(optional_user),\n) -> User:\n user = await double_check_user(user)\n\n if user.role == UserRole.LIMITED:\n raise BasicAuthenticationError(\n detail=\"Access denied. User role is LIMITED. BASIC or higher permissions are required.\",\n )\n return user\n\n\nasync def current_curator_or_admin_user(\n user: User = Depends(current_user),\n) -> User:\n allowed_roles = {UserRole.GLOBAL_CURATOR, UserRole.CURATOR, UserRole.ADMIN}\n if user.role not in allowed_roles:\n raise BasicAuthenticationError(\n detail=\"Access denied. User is not a curator or admin.\",\n )\n\n return user\n\n\nasync def current_admin_user(user: User = Depends(current_user)) -> User:\n if user.role != UserRole.ADMIN:\n raise BasicAuthenticationError(\n detail=\"Access denied. User must be an admin to perform this action.\",\n )\n\n return user\n\n\nasync def _get_user_from_token_data(token_data: dict) -> User | None:\n \"\"\"Shared logic: token data dict → User object.\n\n Args:\n token_data: Decoded token data containing 'sub' (user ID).\n\n Returns:\n User object if found and active, None otherwise.\n \"\"\"\n user_id = token_data.get(\"sub\")\n if not user_id:\n return None\n\n try:\n user_uuid = uuid.UUID(user_id)\n except ValueError:\n return None\n\n async with get_async_session_context_manager() as async_db_session:\n user = await async_db_session.get(User, user_uuid)\n if user is None or not user.is_active:\n return None\n return user\n\n\n_LOOPBACK_HOSTNAMES = frozenset({\"localhost\", \"127.0.0.1\", \"::1\"})\n\n\ndef _is_same_origin(actual: str, expected: str) -> bool:\n \"\"\"Compare two origins for the WebSocket CSWSH check.\n\n Scheme and hostname must match exactly. Port must also match, except\n when the hostname is a loopback address (localhost / 127.0.0.1 / ::1),\n where port is ignored. On loopback, all ports belong to the same\n operator, so port differences carry no security significance — the\n CSWSH threat is remote origins, not local ones.\n \"\"\"\n a = urlparse(actual.rstrip(\"/\"))\n e = urlparse(expected.rstrip(\"/\"))\n\n if a.scheme != e.scheme or a.hostname != e.hostname:\n return False\n\n if a.hostname in _LOOPBACK_HOSTNAMES:\n return True\n\n actual_port = a.port or (443 if a.scheme == \"https\" else 80)\n expected_port = e.port or (443 if e.scheme == \"https\" else 80)\n\n return actual_port == expected_port\n\n\nasync def current_user_from_websocket(\n websocket: WebSocket,\n token: str = Query(..., description=\"WebSocket authentication token\"),\n) -> User:\n \"\"\"\n WebSocket authentication dependency using query parameter.\n\n Validates the WS token from query param and returns the User.\n Raises BasicAuthenticationError if authentication fails.\n\n The token must be obtained from POST /voice/ws-token before connecting.\n Tokens are single-use and expire after 60 seconds.\n\n Usage:\n 1. POST /voice/ws-token -> {\"token\": \"xxx\"}\n 2. Connect to ws://host/path?token=xxx\n\n This applies the same auth checks as current_user() for HTTP endpoints.\n \"\"\"\n # Check Origin header to prevent Cross-Site WebSocket Hijacking (CSWSH).\n # Browsers always send Origin on WebSocket connections.\n origin = websocket.headers.get(\"origin\")\n if not origin:\n logger.warning(\"WS auth: missing Origin header\")\n raise BasicAuthenticationError(detail=\"Access denied. Missing origin.\")\n\n if not _is_same_origin(origin, WEB_DOMAIN):\n logger.warning(f\"WS auth: origin mismatch. Expected {WEB_DOMAIN}, got {origin}\")\n raise BasicAuthenticationError(detail=\"Access denied. Invalid origin.\")\n\n # Validate WS token in Redis (single-use, deleted after retrieval)\n try:\n token_data = await retrieve_ws_token_data(token)\n if token_data is None:\n raise BasicAuthenticationError(\n detail=\"Access denied. Invalid or expired authentication token.\"\n )\n except BasicAuthenticationError:\n raise\n except Exception as e:\n logger.error(f\"WS auth: error during token validation: {e}\")\n raise BasicAuthenticationError(\n detail=\"Authentication verification failed.\"\n ) from e\n\n # Get user from token data\n user = await _get_user_from_token_data(token_data)\n if user is None:\n logger.warning(f\"WS auth: user not found for id={token_data.get('sub')}\")\n raise BasicAuthenticationError(\n detail=\"Access denied. User not found or inactive.\"\n )\n\n # Apply same checks as HTTP auth (verification, OIDC expiry, role)\n user = await double_check_user(user)\n\n # Block LIMITED users (same as current_user)\n if user.role == UserRole.LIMITED:\n logger.warning(f\"WS auth: user {user.email} has LIMITED role\")\n raise BasicAuthenticationError(\n detail=\"Access denied. User role is LIMITED. BASIC or higher permissions are required.\",\n )\n\n logger.debug(f\"WS auth: authenticated {user.email}\")\n return user\n\n\ndef get_default_admin_user_emails_() -> list[str]:\n # No default seeding available for Onyx MIT\n return []\n\n\nSTATE_TOKEN_AUDIENCE = \"fastapi-users:oauth-state\"\nSTATE_TOKEN_LIFETIME_SECONDS = 3600\nCSRF_TOKEN_KEY = \"csrftoken\"\nCSRF_TOKEN_COOKIE_NAME = \"fastapiusersoauthcsrf\"\nPKCE_COOKIE_NAME_PREFIX = \"fastapiusersoauthpkce\"\n\n\nclass OAuth2AuthorizeResponse(BaseModel):\n authorization_url: str\n\n\ndef generate_state_token(\n data: Dict[str, str],\n secret: SecretType,\n lifetime_seconds: int = STATE_TOKEN_LIFETIME_SECONDS,\n) -> str:\n data[\"aud\"] = STATE_TOKEN_AUDIENCE\n\n return generate_jwt(data, secret, lifetime_seconds)\n\n\ndef generate_csrf_token() -> str:\n return secrets.token_urlsafe(32)\n\n\ndef _base64url_encode(data: bytes) -> str:\n return base64.urlsafe_b64encode(data).rstrip(b\"=\").decode(\"ascii\")\n\n\ndef generate_pkce_pair() -> tuple[str, str]:\n verifier = secrets.token_urlsafe(64)\n challenge = _base64url_encode(hashlib.sha256(verifier.encode(\"ascii\")).digest())\n return verifier, challenge\n\n\ndef get_pkce_cookie_name(state: str) -> str:\n state_hash = hashlib.sha256(state.encode(\"utf-8\")).hexdigest()\n return f\"{PKCE_COOKIE_NAME_PREFIX}_{state_hash}\"\n\n\n# refer to https://github.com/fastapi-users/fastapi-users/blob/42ddc241b965475390e2bce887b084152ae1a2cd/fastapi_users/fastapi_users.py#L91\ndef create_onyx_oauth_router(\n oauth_client: BaseOAuth2,\n backend: AuthenticationBackend,\n state_secret: SecretType,\n redirect_url: Optional[str] = None,\n associate_by_email: bool = False,\n is_verified_by_default: bool = False,\n enable_pkce: bool = False,\n) -> APIRouter:\n return get_oauth_router(\n oauth_client,\n backend,\n get_user_manager,\n state_secret,\n redirect_url,\n associate_by_email,\n is_verified_by_default,\n enable_pkce=enable_pkce,\n )\n\n\ndef get_oauth_router(\n oauth_client: BaseOAuth2,\n backend: AuthenticationBackend,\n get_user_manager: UserManagerDependency[models.UP, models.ID],\n state_secret: SecretType,\n redirect_url: Optional[str] = None,\n associate_by_email: bool = False,\n is_verified_by_default: bool = False,\n *,\n csrf_token_cookie_name: str = CSRF_TOKEN_COOKIE_NAME,\n csrf_token_cookie_path: str = \"/\",\n csrf_token_cookie_domain: Optional[str] = None,\n csrf_token_cookie_secure: Optional[bool] = None,\n csrf_token_cookie_httponly: bool = True,\n csrf_token_cookie_samesite: Optional[Literal[\"lax\", \"strict\", \"none\"]] = \"lax\",\n enable_pkce: bool = False,\n) -> APIRouter:\n \"\"\"Generate a router with the OAuth routes.\"\"\"\n router = APIRouter()\n callback_route_name = f\"oauth:{oauth_client.name}.{backend.name}.callback\"\n\n if redirect_url is not None:\n oauth2_authorize_callback = OAuth2AuthorizeCallback(\n oauth_client,\n redirect_url=redirect_url,\n )\n else:\n oauth2_authorize_callback = OAuth2AuthorizeCallback(\n oauth_client,\n route_name=callback_route_name,\n )\n\n async def null_access_token_state() -> tuple[OAuth2Token, Optional[str]] | None:\n return None\n\n access_token_state_dependency = (\n oauth2_authorize_callback if not enable_pkce else null_access_token_state\n )\n\n if csrf_token_cookie_secure is None:\n csrf_token_cookie_secure = WEB_DOMAIN.startswith(\"https\")\n\n @router.get(\n \"/authorize\",\n name=f\"oauth:{oauth_client.name}.{backend.name}.authorize\",\n response_model=OAuth2AuthorizeResponse,\n )\n async def authorize(\n request: Request,\n response: Response,\n redirect: bool = Query(False),\n scopes: List[str] = Query(None),\n ) -> Response | OAuth2AuthorizeResponse:\n referral_source = request.cookies.get(\"referral_source\", None)\n\n if redirect_url is not None:\n authorize_redirect_url = redirect_url\n else:\n # Use WEB_DOMAIN instead of request.url_for() to prevent host\n # header poisoning — request.url_for() trusts the Host header.\n callback_path = request.app.url_path_for(callback_route_name)\n authorize_redirect_url = f\"{WEB_DOMAIN}{callback_path}\"\n\n next_url = request.query_params.get(\"next\", \"/\")\n\n csrf_token = generate_csrf_token()\n state_data: Dict[str, str] = {\n \"next_url\": next_url,\n \"referral_source\": referral_source or \"default_referral\",\n CSRF_TOKEN_KEY: csrf_token,\n }\n state = generate_state_token(state_data, state_secret)\n pkce_cookie: tuple[str, str] | None = None\n\n if enable_pkce:\n code_verifier, code_challenge = generate_pkce_pair()\n pkce_cookie_name = get_pkce_cookie_name(state)\n pkce_cookie = (pkce_cookie_name, code_verifier)\n authorization_url = await oauth_client.get_authorization_url(\n authorize_redirect_url,\n state,\n scopes,\n code_challenge=code_challenge,\n code_challenge_method=\"S256\",\n )\n else:\n # Get the basic authorization URL\n authorization_url = await oauth_client.get_authorization_url(\n authorize_redirect_url,\n state,\n scopes,\n )\n\n # For Google OAuth, add parameters to request refresh tokens\n if oauth_client.name == \"google\":\n authorization_url = add_url_params(\n authorization_url, {\"access_type\": \"offline\", \"prompt\": \"consent\"}\n )\n\n def set_oauth_cookie(\n target_response: Response,\n *,\n key: str,\n value: str,\n ) -> None:\n target_response.set_cookie(\n key=key,\n value=value,\n max_age=STATE_TOKEN_LIFETIME_SECONDS,\n path=csrf_token_cookie_path,\n domain=csrf_token_cookie_domain,\n secure=csrf_token_cookie_secure,\n httponly=csrf_token_cookie_httponly,\n samesite=csrf_token_cookie_samesite,\n )\n\n response_with_cookies: Response\n if redirect:\n response_with_cookies = RedirectResponse(authorization_url, status_code=302)\n else:\n response_with_cookies = response\n\n set_oauth_cookie(\n response_with_cookies,\n key=csrf_token_cookie_name,\n value=csrf_token,\n )\n if pkce_cookie is not None:\n pkce_cookie_name, code_verifier = pkce_cookie\n set_oauth_cookie(\n response_with_cookies,\n key=pkce_cookie_name,\n value=code_verifier,\n )\n\n if redirect:\n return response_with_cookies\n\n return OAuth2AuthorizeResponse(authorization_url=authorization_url)\n\n @log_function_time(print_only=True)\n @router.get(\n \"/callback\",\n name=callback_route_name,\n description=\"The response varies based on the authentication backend used.\",\n responses={\n status.HTTP_400_BAD_REQUEST: {\n \"model\": ErrorModel,\n \"content\": {\n \"application/json\": {\n \"examples\": {\n \"INVALID_STATE_TOKEN\": {\n \"summary\": \"Invalid state token.\",\n \"value\": None,\n },\n ErrorCode.LOGIN_BAD_CREDENTIALS: {\n \"summary\": \"User is inactive.\",\n \"value\": {\"detail\": ErrorCode.LOGIN_BAD_CREDENTIALS},\n },\n }\n }\n },\n },\n },\n )\n async def callback(\n request: Request,\n access_token_state: Tuple[OAuth2Token, Optional[str]] | None = Depends(\n access_token_state_dependency\n ),\n code: Optional[str] = None,\n state: Optional[str] = None,\n error: Optional[str] = None,\n user_manager: BaseUserManager[models.UP, models.ID] = Depends(get_user_manager),\n strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),\n ) -> Response:\n pkce_cookie_name: str | None = None\n\n def delete_pkce_cookie(response: Response) -> None:\n if enable_pkce and pkce_cookie_name:\n response.delete_cookie(\n key=pkce_cookie_name,\n path=csrf_token_cookie_path,\n domain=csrf_token_cookie_domain,\n secure=csrf_token_cookie_secure,\n httponly=csrf_token_cookie_httponly,\n samesite=csrf_token_cookie_samesite,\n )\n\n def build_error_response(exc: OnyxError) -> JSONResponse:\n log_onyx_error(exc)\n error_response = onyx_error_to_json_response(exc)\n delete_pkce_cookie(error_response)\n return error_response\n\n def decode_and_validate_state(state_value: str) -> Dict[str, str]:\n try:\n state_data = decode_jwt(\n state_value, state_secret, [STATE_TOKEN_AUDIENCE]\n )\n except jwt.DecodeError:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n getattr(\n ErrorCode,\n \"ACCESS_TOKEN_DECODE_ERROR\",\n \"ACCESS_TOKEN_DECODE_ERROR\",\n ),\n )\n except jwt.ExpiredSignatureError:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n getattr(\n ErrorCode,\n \"ACCESS_TOKEN_ALREADY_EXPIRED\",\n \"ACCESS_TOKEN_ALREADY_EXPIRED\",\n ),\n )\n except jwt.PyJWTError:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n getattr(\n ErrorCode,\n \"ACCESS_TOKEN_DECODE_ERROR\",\n \"ACCESS_TOKEN_DECODE_ERROR\",\n ),\n )\n\n cookie_csrf_token = request.cookies.get(csrf_token_cookie_name)\n state_csrf_token = state_data.get(CSRF_TOKEN_KEY)\n if (\n not cookie_csrf_token\n or not state_csrf_token\n or not secrets.compare_digest(cookie_csrf_token, state_csrf_token)\n ):\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n getattr(ErrorCode, \"OAUTH_INVALID_STATE\", \"OAUTH_INVALID_STATE\"),\n )\n\n return state_data\n\n token: OAuth2Token\n state_data: Dict[str, str]\n\n # `code`, `state`, and `error` are read directly only in the PKCE path.\n # In the non-PKCE path, `oauth2_authorize_callback` consumes them.\n if enable_pkce:\n if state is not None:\n pkce_cookie_name = get_pkce_cookie_name(state)\n\n if error is not None:\n return build_error_response(\n OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"Authorization request failed or was denied\",\n )\n )\n if code is None:\n return build_error_response(\n OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"Missing authorization code in OAuth callback\",\n )\n )\n if state is None:\n return build_error_response(\n OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"Missing state parameter in OAuth callback\",\n )\n )\n\n state_value = state\n\n if redirect_url is not None:\n callback_redirect_url = redirect_url\n else:\n callback_path = request.app.url_path_for(callback_route_name)\n callback_redirect_url = f\"{WEB_DOMAIN}{callback_path}\"\n\n code_verifier = request.cookies.get(cast(str, pkce_cookie_name))\n if not code_verifier:\n return build_error_response(\n OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"Missing PKCE verifier cookie in OAuth callback\",\n )\n )\n\n try:\n state_data = decode_and_validate_state(state_value)\n except OnyxError as e:\n return build_error_response(e)\n\n try:\n token = await oauth_client.get_access_token(\n code, callback_redirect_url, code_verifier\n )\n except GetAccessTokenError:\n return build_error_response(\n OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"Authorization code exchange failed\",\n )\n )\n else:\n if access_token_state is None:\n raise OnyxError(\n OnyxErrorCode.INTERNAL_ERROR, \"Missing OAuth callback state\"\n )\n token, callback_state = access_token_state\n if callback_state is None:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n \"Missing state parameter in OAuth callback\",\n )\n state_data = decode_and_validate_state(callback_state)\n\n async def complete_login_flow(\n token: OAuth2Token, state_data: Dict[str, str]\n ) -> RedirectResponse:\n account_id, account_email = await oauth_client.get_id_email(\n token[\"access_token\"]\n )\n\n if account_email is None:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n ErrorCode.OAUTH_NOT_AVAILABLE_EMAIL,\n )\n\n next_url = state_data.get(\"next_url\", \"/\")\n referral_source = state_data.get(\"referral_source\", None)\n try:\n tenant_id = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.user_mapping\", \"get_tenant_id_for_email\", None\n )(account_email)\n except exceptions.UserNotExists:\n tenant_id = None\n\n request.state.referral_source = referral_source\n\n # Proceed to authenticate or create the user\n try:\n user = await user_manager.oauth_callback(\n oauth_client.name,\n token[\"access_token\"],\n account_id,\n account_email,\n token.get(\"expires_at\"),\n token.get(\"refresh_token\"),\n request,\n associate_by_email=associate_by_email,\n is_verified_by_default=is_verified_by_default,\n )\n except UserAlreadyExists:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n ErrorCode.OAUTH_USER_ALREADY_EXISTS,\n )\n\n if not user.is_active:\n raise OnyxError(\n OnyxErrorCode.VALIDATION_ERROR,\n ErrorCode.LOGIN_BAD_CREDENTIALS,\n )\n\n # Login user\n response = await backend.login(strategy, user)\n await user_manager.on_after_login(user, request, response)\n\n # Prepare redirect response\n if tenant_id is None:\n # Use URL utility to add parameters\n redirect_destination = add_url_params(next_url, {\"new_team\": \"true\"})\n redirect_response = RedirectResponse(\n redirect_destination, status_code=302\n )\n else:\n # No parameters to add\n redirect_response = RedirectResponse(next_url, status_code=302)\n\n # Copy headers from auth response to redirect response, with special handling for Set-Cookie\n for header_name, header_value in response.headers.items():\n header_name_lower = header_name.lower()\n if header_name_lower == \"set-cookie\":\n redirect_response.headers.append(header_name, header_value)\n continue\n if header_name_lower in {\"location\", \"content-length\"}:\n continue\n redirect_response.headers[header_name] = header_value\n\n return redirect_response\n\n if enable_pkce:\n try:\n redirect_response = await complete_login_flow(token, state_data)\n except OnyxError as e:\n return build_error_response(e)\n delete_pkce_cookie(redirect_response)\n return redirect_response\n\n return await complete_login_flow(token, state_data)\n\n return router\n" }, { "path": "backend/onyx/auth/utils.py", "content": "\"\"\"Shared authentication utilities for bearer token extraction and validation.\"\"\"\n\nfrom collections.abc import Callable\nfrom urllib.parse import unquote\n\nfrom fastapi import Request\n\nfrom onyx.auth.constants import API_KEY_HEADER_ALTERNATIVE_NAME\nfrom onyx.auth.constants import API_KEY_HEADER_NAME\nfrom onyx.auth.constants import API_KEY_PREFIX\nfrom onyx.auth.constants import BEARER_PREFIX\nfrom onyx.auth.constants import DEPRECATED_API_KEY_PREFIX\nfrom onyx.auth.constants import PAT_PREFIX\n\n\ndef get_hashed_bearer_token_from_request(\n request: Request,\n valid_prefixes: list[str],\n hash_fn: Callable[[str], str],\n allow_non_bearer: bool = False,\n) -> str | None:\n \"\"\"Generic extraction and hashing of bearer tokens from request headers.\n\n Args:\n request: The FastAPI request\n valid_prefixes: List of valid token prefixes (e.g., [\"on_\", \"onyx_pat_\"])\n hash_fn: Function to hash the token (e.g., hash_api_key or hash_pat)\n allow_non_bearer: If True, accept raw tokens without \"Bearer \" prefix\n\n Returns:\n Hashed token if valid format, else None\n \"\"\"\n auth_header = request.headers.get(\n API_KEY_HEADER_ALTERNATIVE_NAME\n ) or request.headers.get(API_KEY_HEADER_NAME)\n\n if not auth_header:\n return None\n\n # Handle bearer format\n if auth_header.startswith(BEARER_PREFIX):\n token = auth_header[len(BEARER_PREFIX) :].strip()\n elif allow_non_bearer:\n token = auth_header\n else:\n return None\n\n # Check if token starts with any valid prefix\n if valid_prefixes:\n valid = any(token.startswith(prefix) for prefix in valid_prefixes)\n if not valid:\n return None\n\n return hash_fn(token)\n\n\ndef _extract_tenant_from_bearer_token(\n request: Request, valid_prefixes: list[str]\n) -> str | None:\n \"\"\"Generic tenant extraction from bearer token. Returns None if invalid format.\n\n Args:\n request: The FastAPI request\n valid_prefixes: List of valid token prefixes (e.g., [\"on_\", \"dn_\"])\n\n Returns:\n Tenant ID if found in format ., else None\n \"\"\"\n auth_header = request.headers.get(\n API_KEY_HEADER_ALTERNATIVE_NAME\n ) or request.headers.get(API_KEY_HEADER_NAME)\n\n if not auth_header or not auth_header.startswith(BEARER_PREFIX):\n return None\n\n token = auth_header[len(BEARER_PREFIX) :].strip()\n\n # Check if token starts with any valid prefix\n matched_prefix = None\n for prefix in valid_prefixes:\n if token.startswith(prefix):\n matched_prefix = prefix\n break\n\n if not matched_prefix:\n return None\n\n # Parse tenant from token format: .\n parts = token[len(matched_prefix) :].split(\".\", 1)\n if len(parts) != 2:\n return None\n\n tenant_id = parts[0]\n return unquote(tenant_id) if tenant_id else None\n\n\ndef extract_tenant_from_auth_header(request: Request) -> str | None:\n \"\"\"Extract tenant ID from API key or PAT header.\n\n Unified function for extracting tenant from any bearer token (API key or PAT).\n Checks all known token prefixes in order.\n\n Returns:\n Tenant ID if found, else None\n \"\"\"\n return _extract_tenant_from_bearer_token(\n request, [API_KEY_PREFIX, DEPRECATED_API_KEY_PREFIX, PAT_PREFIX]\n )\n" }, { "path": "backend/onyx/background/README.md", "content": "# Overview of Onyx Background Jobs\n\nThe background jobs take care of:\n1. Pulling/Indexing documents (from connectors)\n2. Updating document metadata (from connectors)\n3. Cleaning up checkpoints and logic around indexing work (indexing indexing checkpoints and index attempt metadata)\n4. Handling user uploaded files and deletions (from the Projects feature and uploads via the Chat)\n5. Reporting metrics on things like queue length for monitoring purposes\n\n## Worker → Queue Mapping\n\n| Worker | File | Queues |\n|--------|------|--------|\n| Primary | `apps/primary.py` | `celery` |\n| Light | `apps/light.py` | `vespa_metadata_sync`, `connector_deletion`, `doc_permissions_upsert`, `checkpoint_cleanup`, `index_attempt_cleanup` |\n| Heavy | `apps/heavy.py` | `connector_pruning`, `connector_doc_permissions_sync`, `connector_external_group_sync`, `csv_generation`, `sandbox` |\n| Docprocessing | `apps/docprocessing.py` | `docprocessing` |\n| Docfetching | `apps/docfetching.py` | `connector_doc_fetching` |\n| User File Processing | `apps/user_file_processing.py` | `user_file_processing`, `user_file_project_sync`, `user_file_delete` |\n| Monitoring | `apps/monitoring.py` | `monitoring` |\n| Background (consolidated) | `apps/background.py` | All queues above except `celery` |\n\n## Non-Worker Apps\n| App | File | Purpose |\n|-----|------|---------|\n| **Beat** | `beat.py` | Celery beat scheduler with `DynamicTenantScheduler` that generates per-tenant periodic task schedules |\n| **Client** | `client.py` | Minimal app for task submission from non-worker processes (e.g., API server) |\n\n### Shared Module\n`app_base.py` provides:\n- `TenantAwareTask` - Base task class that sets tenant context\n- Signal handlers for logging, cleanup, and lifecycle events\n- Readiness probes and health checks\n\n\n## Worker Details\n\n### Primary (Coordinator and task dispatcher)\nIt is the single worker which handles tasks from the default celery queue. It is a singleton worker ensured by the `PRIMARY_WORKER` Redis lock\nwhich it touches every `CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 8` seconds (using Celery Bootsteps)\n\nOn startup:\n- waits for redis, postgres, document index to all be healthy\n- acquires the singleton lock\n- cleans all the redis states associated with background jobs\n- mark orphaned index attempts failed\n\nThen it cycles through its tasks as scheduled by Celery Beat:\n\n| Task | Frequency | Description |\n|------|-----------|-------------|\n| `check_for_indexing` | 15s | Scans for connectors needing indexing → dispatches to `DOCFETCHING` queue |\n| `check_for_vespa_sync_task` | 20s | Finds stale documents/document sets → dispatches sync tasks to `VESPA_METADATA_SYNC` queue |\n| `check_for_pruning` | 20s | Finds connectors due for pruning → dispatches to `CONNECTOR_PRUNING` queue |\n| `check_for_connector_deletion` | 20s | Processes deletion requests → dispatches to `CONNECTOR_DELETION` queue |\n| `check_for_user_file_processing` | 20s | Checks for user uploads → dispatches to `USER_FILE_PROCESSING` queue |\n| `check_for_checkpoint_cleanup` | 1h | Cleans up old indexing checkpoints |\n| `check_for_index_attempt_cleanup` | 30m | Cleans up old index attempts |\n| `kombu_message_cleanup_task` | periodic | Cleans orphaned Kombu messages from DB (Kombu being the messaging framework used by Celery) |\n| `celery_beat_heartbeat` | 1m | Heartbeat for Beat watchdog |\n\nWatchdog is a separate Python process managed by supervisord which runs alongside celery workers. It checks the ONYX_CELERY_BEAT_HEARTBEAT_KEY in\nRedis to ensure Celery Beat is not dead. Beat schedules the celery_beat_heartbeat for Primary to touch the key and share that it's still alive.\nSee supervisord.conf for watchdog config.\n\n\n### Light\nFast and short living tasks that are not resource intensive. High concurrency:\nCan have 24 concurrent workers, each with a prefetch of 8 for a total of 192 tasks in flight at once.\n\nTasks it handles:\n- Syncs access/permissions, document sets, boosts, hidden state\n- Deletes documents that are marked for deletion in Postgres\n- Cleanup of checkpoints and index attempts\n\n\n### Heavy\nLong running, resource intensive tasks, handles pruning and sandbox operations. Low concurrency - max concurrency of 4 with 1 prefetch.\n\nDoes not interact with the Document Index, it handles the syncs with external systems. Large volume API calls to handle pruning and fetching permissions, etc.\n\nGenerates CSV exports which may take a long time with significant data in Postgres.\n\nSandbox (new feature) for running Next.js, Python virtual env, OpenCode AI Agent, and access to knowledge files\n\n\n### Docprocessing, Docfetching, User File Processing\nDocprocessing and Docfetching are for indexing documents:\n- Docfetching runs connectors to pull documents from external APIs (Google Drive, Confluence, etc.), stores batches to file storage, and dispatches docprocessing tasks\n- Docprocessing retrieves batches, runs the indexing pipeline (chunking, embedding), and indexes into the Document Index \nUser Files come from uploads directly via the input bar\n\n\n### Monitoring\nObservability and metrics collections:\n- Queue lengths, connector success/failure, lconnector latencies\n- Memory of supervisor managed processes (workers, beat, slack)\n- Cloud and multitenant specific monitorings\n" }, { "path": "backend/onyx/background/celery/apps/app_base.py", "content": "import logging\nimport multiprocessing\nimport os\nimport time\nfrom typing import Any\nfrom typing import cast\n\nimport sentry_sdk\nfrom celery import bootsteps # type: ignore\nfrom celery import Task\nfrom celery.app import trace\nfrom celery.exceptions import WorkerShutdown\nfrom celery.signals import task_postrun\nfrom celery.signals import task_prerun\nfrom celery.states import READY_STATES\nfrom celery.utils.log import get_task_logger\nfrom celery.worker import strategy # type: ignore\nfrom redis.lock import Lock as RedisLock\nfrom sentry_sdk.integrations.celery import CeleryIntegration\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\n\nfrom onyx import __version__\nfrom onyx.background.celery.apps.task_formatters import CeleryTaskColoredFormatter\nfrom onyx.background.celery.apps.task_formatters import CeleryTaskPlainFormatter\nfrom onyx.background.celery.celery_utils import celery_is_worker_primary\nfrom onyx.background.celery.celery_utils import make_probe_path\nfrom onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_PREFIX\nfrom onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_TASKSET_KEY\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\nfrom onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.db.engine.sql_engine import get_sqlalchemy_engine\nfrom onyx.document_index.opensearch.client import (\n wait_for_opensearch_with_timeout,\n)\nfrom onyx.document_index.vespa.shared_utils.utils import wait_for_vespa_with_timeout\nfrom onyx.httpx.httpx_pool import HttpxPool\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_connector_delete import RedisConnectorDelete\nfrom onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync\nfrom onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync\nfrom onyx.redis.redis_connector_prune import RedisConnectorPrune\nfrom onyx.redis.redis_document_set import RedisDocumentSet\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_usergroup import RedisUserGroup\nfrom onyx.tracing.setup import setup_tracing\nfrom onyx.utils.logger import ColoredFormatter\nfrom onyx.utils.logger import LoggerContextVars\nfrom onyx.utils.logger import PlainFormatter\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import DEV_LOGGING_ENABLED\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import SENTRY_DSN\nfrom shared_configs.configs import TENANT_ID_PREFIX\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\ntask_logger = get_task_logger(__name__)\n\nif SENTRY_DSN:\n sentry_sdk.init(\n dsn=SENTRY_DSN,\n integrations=[CeleryIntegration()],\n traces_sample_rate=0.1,\n release=__version__,\n )\n logger.info(\"Sentry initialized\")\nelse:\n logger.debug(\"Sentry DSN not provided, skipping Sentry initialization\")\n\n\nclass TenantAwareTask(Task):\n \"\"\"A custom base Task that sets tenant_id in a contextvar before running.\"\"\"\n\n abstract = True # So Celery knows not to register this as a real task.\n\n def __call__(self, *args: Any, **kwargs: Any) -> Any:\n # Grab tenant_id from the kwargs, or fallback to default if missing.\n tenant_id = kwargs.get(\"tenant_id\", None) or POSTGRES_DEFAULT_SCHEMA\n\n # Set the context var\n CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n # Actually run the task now\n try:\n return super().__call__(*args, **kwargs)\n finally:\n # Clear or reset after the task runs\n # so it does not leak into any subsequent tasks on the same worker process\n CURRENT_TENANT_ID_CONTEXTVAR.set(None)\n\n\n@task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None, # noqa: ARG001\n task_id: str | None = None, # noqa: ARG001\n task: Task | None = None, # noqa: ARG001\n args: tuple[Any, ...] | None = None, # noqa: ARG001\n kwargs: dict[str, Any] | None = None, # noqa: ARG001\n **other_kwargs: Any, # noqa: ARG001\n) -> None:\n # Reset any per-task logging context so that prefixes (e.g. pruning_ctx)\n # from a previous task executed in the same worker process do not leak\n # into the next task's log messages. This fixes incorrect [CC Pair:/Index Attempt]\n # prefixes observed when a pruning task finishes and an indexing task\n # runs in the same process.\n\n LoggerContextVars.reset()\n\n\ndef on_task_postrun(\n sender: Any | None = None, # noqa: ARG001\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None, # noqa: ARG001\n kwargs: dict[str, Any] | None = None,\n retval: Any | None = None, # noqa: ARG001\n state: str | None = None,\n **kwds: Any, # noqa: ARG001\n) -> None:\n \"\"\"We handle this signal in order to remove completed tasks\n from their respective tasksets. This allows us to track the progress of document set\n and user group syncs.\n\n This function runs after any task completes (both success and failure)\n Note that this signal does not fire on a task that failed to complete and is going\n to be retried.\n\n This also does not fire if a worker with acks_late=False crashes (which all of our\n long running workers are)\n \"\"\"\n if not task:\n return\n\n task_logger.debug(f\"Task {task.name} (ID: {task_id}) completed with state: {state}\")\n\n if state not in READY_STATES:\n return\n\n if not task_id:\n return\n\n if task.name.startswith(ONYX_CLOUD_CELERY_TASK_PREFIX):\n # this is a cloud / all tenant task ... no postrun is needed\n return\n\n # Get tenant_id directly from kwargs- each celery task has a tenant_id kwarg\n if not kwargs:\n logger.error(f\"Task {task.name} (ID: {task_id}) is missing kwargs\")\n tenant_id = POSTGRES_DEFAULT_SCHEMA\n else:\n tenant_id = cast(str, kwargs.get(\"tenant_id\", POSTGRES_DEFAULT_SCHEMA))\n\n task_logger.debug(\n f\"Task {task.name} (ID: {task_id}) completed with state: {state} {f'for tenant_id={tenant_id}' if tenant_id else ''}\"\n )\n\n r = get_redis_client(tenant_id=tenant_id)\n\n # NOTE: we want to remove the `Redis*` classes, prefer to just have functions to\n # do these things going forward. In short, things should generally be like the doc\n # sync task rather than the others below\n if task_id.startswith(DOCUMENT_SYNC_PREFIX):\n r.srem(DOCUMENT_SYNC_TASKSET_KEY, task_id)\n return\n\n if task_id.startswith(RedisDocumentSet.PREFIX):\n document_set_id = RedisDocumentSet.get_id_from_task_id(task_id)\n if document_set_id is not None:\n rds = RedisDocumentSet(tenant_id, int(document_set_id))\n r.srem(rds.taskset_key, task_id)\n return\n\n if task_id.startswith(RedisUserGroup.PREFIX):\n usergroup_id = RedisUserGroup.get_id_from_task_id(task_id)\n if usergroup_id is not None:\n rug = RedisUserGroup(tenant_id, int(usergroup_id))\n r.srem(rug.taskset_key, task_id)\n return\n\n if task_id.startswith(RedisConnectorDelete.PREFIX):\n cc_pair_id = RedisConnector.get_id_from_task_id(task_id)\n if cc_pair_id is not None:\n RedisConnectorDelete.remove_from_taskset(int(cc_pair_id), task_id, r)\n return\n\n if task_id.startswith(RedisConnectorPrune.SUBTASK_PREFIX):\n cc_pair_id = RedisConnector.get_id_from_task_id(task_id)\n if cc_pair_id is not None:\n RedisConnectorPrune.remove_from_taskset(int(cc_pair_id), task_id, r)\n return\n\n if task_id.startswith(RedisConnectorPermissionSync.SUBTASK_PREFIX):\n cc_pair_id = RedisConnector.get_id_from_task_id(task_id)\n if cc_pair_id is not None:\n RedisConnectorPermissionSync.remove_from_taskset(\n int(cc_pair_id), task_id, r\n )\n return\n\n if task_id.startswith(RedisConnectorExternalGroupSync.SUBTASK_PREFIX):\n cc_pair_id = RedisConnector.get_id_from_task_id(task_id)\n if cc_pair_id is not None:\n RedisConnectorExternalGroupSync.remove_from_taskset(\n int(cc_pair_id), task_id, r\n )\n return\n\n\ndef on_celeryd_init(\n sender: str, # noqa: ARG001\n conf: Any = None, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n) -> None:\n \"\"\"The first signal sent on celery worker startup\"\"\"\n\n # NOTE(rkuo): start method \"fork\" is unsafe and we really need it to be \"spawn\"\n # But something is blocking set_start_method from working in the cloud unless\n # force=True. so we use force=True as a fallback.\n\n all_start_methods: list[str] = multiprocessing.get_all_start_methods()\n logger.info(f\"Multiprocessing all start methods: {all_start_methods}\")\n\n try:\n multiprocessing.set_start_method(\"spawn\") # fork is unsafe, set to spawn\n except Exception:\n logger.info(\n \"Multiprocessing set_start_method exceptioned. Trying force=True...\"\n )\n try:\n multiprocessing.set_start_method(\n \"spawn\", force=True\n ) # fork is unsafe, set to spawn\n except Exception:\n logger.info(\n \"Multiprocessing set_start_method force=True exceptioned even with force=True.\"\n )\n\n logger.info(\n f\"Multiprocessing selected start method: {multiprocessing.get_start_method()}\"\n )\n\n # Initialize tracing in workers if credentials are available.\n setup_tracing()\n\n\ndef wait_for_redis(sender: Any, **kwargs: Any) -> None: # noqa: ARG001\n \"\"\"Waits for redis to become ready subject to a hardcoded timeout.\n Will raise WorkerShutdown to kill the celery worker if the timeout\n is reached.\"\"\"\n\n r = get_redis_client(tenant_id=POSTGRES_DEFAULT_SCHEMA)\n\n WAIT_INTERVAL = 5\n WAIT_LIMIT = 60\n\n ready = False\n time_start = time.monotonic()\n logger.info(\"Redis: Readiness probe starting.\")\n while True:\n try:\n if r.ping():\n ready = True\n break\n except Exception:\n pass\n\n time_elapsed = time.monotonic() - time_start\n if time_elapsed > WAIT_LIMIT:\n break\n\n logger.info(\n f\"Redis: Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}\"\n )\n\n time.sleep(WAIT_INTERVAL)\n\n if not ready:\n msg = f\"Redis: Readiness probe did not succeed within the timeout ({WAIT_LIMIT} seconds). Exiting...\"\n logger.error(msg)\n raise WorkerShutdown(msg)\n\n logger.info(\"Redis: Readiness probe succeeded. Continuing...\")\n return\n\n\ndef wait_for_db(sender: Any, **kwargs: Any) -> None: # noqa: ARG001\n \"\"\"Waits for the db to become ready subject to a hardcoded timeout.\n Will raise WorkerShutdown to kill the celery worker if the timeout is reached.\"\"\"\n\n WAIT_INTERVAL = 5\n WAIT_LIMIT = 60\n\n ready = False\n time_start = time.monotonic()\n logger.info(\"Database: Readiness probe starting.\")\n while True:\n try:\n with Session(get_sqlalchemy_engine()) as db_session:\n result = db_session.execute(text(\"SELECT NOW()\")).scalar()\n if result:\n ready = True\n break\n except Exception:\n pass\n\n time_elapsed = time.monotonic() - time_start\n if time_elapsed > WAIT_LIMIT:\n break\n\n logger.info(\n f\"Database: Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}\"\n )\n\n time.sleep(WAIT_INTERVAL)\n\n if not ready:\n msg = f\"Database: Readiness probe did not succeed within the timeout ({WAIT_LIMIT} seconds). Exiting...\"\n logger.error(msg)\n raise WorkerShutdown(msg)\n\n logger.info(\"Database: Readiness probe succeeded. Continuing...\")\n return\n\n\ndef on_secondary_worker_init(sender: Any, **kwargs: Any) -> None: # noqa: ARG001\n logger.info(f\"Running as a secondary celery worker: pid={os.getpid()}\")\n\n # Set up variables for waiting on primary worker\n WAIT_INTERVAL = 5\n WAIT_LIMIT = 60\n r = get_redis_client(tenant_id=POSTGRES_DEFAULT_SCHEMA)\n time_start = time.monotonic()\n\n logger.info(\"Waiting for primary worker to be ready...\")\n while True:\n if r.exists(OnyxRedisLocks.PRIMARY_WORKER):\n break\n\n time_elapsed = time.monotonic() - time_start\n logger.info(\n f\"Primary worker is not ready yet. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}\"\n )\n if time_elapsed > WAIT_LIMIT:\n msg = f\"Primary worker was not ready within the timeout. ({WAIT_LIMIT} seconds). Exiting...\"\n logger.error(msg)\n raise WorkerShutdown(msg)\n\n time.sleep(WAIT_INTERVAL)\n\n logger.info(\"Wait for primary worker completed successfully. Continuing...\")\n return\n\n\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None: # noqa: ARG001\n task_logger.info(\"worker_ready signal received.\")\n\n # file based way to do readiness/liveness probes\n # https://medium.com/ambient-innovation/health-checks-for-celery-in-kubernetes-cf3274a3e106\n # https://github.com/celery/celery/issues/4079#issuecomment-1270085680\n\n hostname: str = cast(str, sender.hostname)\n path = make_probe_path(\"readiness\", hostname)\n path.touch()\n logger.info(f\"Readiness signal touched at {path}.\")\n\n\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None: # noqa: ARG001\n HttpxPool.close_all()\n\n hostname: str = cast(str, sender.hostname)\n path = make_probe_path(\"readiness\", hostname)\n path.unlink(missing_ok=True)\n\n if not celery_is_worker_primary(sender):\n return\n\n if not hasattr(sender, \"primary_worker_lock\"):\n # primary_worker_lock will not exist when MULTI_TENANT is True\n return\n\n if not sender.primary_worker_lock:\n return\n\n logger.info(\"Releasing primary worker lock.\")\n lock: RedisLock = sender.primary_worker_lock\n try:\n if lock.owned():\n try:\n lock.release()\n sender.primary_worker_lock = None\n except Exception:\n logger.exception(\"Failed to release primary worker lock\")\n except Exception:\n logger.exception(\"Failed to check if primary worker lock is owned\")\n\n\ndef on_setup_logging(\n loglevel: int,\n logfile: str | None,\n format: str, # noqa: ARG001\n colorize: bool, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n) -> None:\n # TODO: could unhardcode format and colorize and accept these as options from\n # celery's config\n\n root_logger = logging.getLogger()\n root_logger.handlers = []\n\n # Define the log format\n log_format = (\n \"%(levelname)-8s %(asctime)s %(filename)15s:%(lineno)-4d: %(name)s %(message)s\"\n )\n\n # Set up the root handler\n root_handler = logging.StreamHandler()\n root_formatter = ColoredFormatter(\n log_format,\n datefmt=\"%m/%d/%Y %I:%M:%S %p\",\n )\n root_handler.setFormatter(root_formatter)\n root_logger.addHandler(root_handler)\n\n if logfile:\n # Truncate log file if DEV_LOGGING_ENABLED (for clean dev experience)\n if DEV_LOGGING_ENABLED and os.path.exists(logfile):\n try:\n open(logfile, \"w\").close() # Truncate the file\n except Exception:\n pass # Ignore errors, just proceed with normal logging\n\n root_file_handler = logging.FileHandler(logfile)\n root_file_formatter = PlainFormatter(\n log_format,\n datefmt=\"%m/%d/%Y %I:%M:%S %p\",\n )\n root_file_handler.setFormatter(root_file_formatter)\n root_logger.addHandler(root_file_handler)\n\n root_logger.setLevel(loglevel)\n\n # Configure the task logger\n task_logger.handlers = []\n\n task_handler = logging.StreamHandler()\n task_handler.addFilter(TenantContextFilter())\n task_formatter = CeleryTaskColoredFormatter(\n log_format,\n datefmt=\"%m/%d/%Y %I:%M:%S %p\",\n )\n task_handler.setFormatter(task_formatter)\n task_logger.addHandler(task_handler)\n\n if logfile:\n # No need to truncate again, already done above for root logger\n task_file_handler = logging.FileHandler(logfile)\n task_file_handler.addFilter(TenantContextFilter())\n task_file_formatter = CeleryTaskPlainFormatter(\n log_format,\n datefmt=\"%m/%d/%Y %I:%M:%S %p\",\n )\n task_file_handler.setFormatter(task_file_formatter)\n task_logger.addHandler(task_file_handler)\n\n task_logger.setLevel(loglevel)\n task_logger.propagate = False\n\n # hide celery task received spam\n # e.g. \"Task check_for_pruning[a1e96171-0ba8-4e00-887b-9fbf7442eab3] received\"\n strategy.logger.setLevel(logging.WARNING)\n\n # uncomment this to hide celery task succeeded/failed spam\n # e.g. \"Task check_for_pruning[a1e96171-0ba8-4e00-887b-9fbf7442eab3] succeeded in 0.03137450001668185s: None\"\n trace.logger.setLevel(logging.WARNING)\n\n\ndef set_task_finished_log_level(logLevel: int) -> None:\n \"\"\"call this to override the setLevel in on_setup_logging. We are interested\n in the task timings in the cloud but it can be spammy for self hosted.\"\"\"\n trace.logger.setLevel(logLevel)\n\n\nclass TenantContextFilter(logging.Filter):\n \"\"\"Logging filter to inject tenant ID into the logger's name.\"\"\"\n\n def filter(self, record: logging.LogRecord) -> bool:\n if not MULTI_TENANT:\n record.name = \"\"\n return True\n\n tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()\n if tenant_id:\n # Match the 8 character tenant abbreviation used in OnyxLoggingAdapter\n tenant_id = tenant_id.split(TENANT_ID_PREFIX)[-1][:8]\n record.name = f\"[t:{tenant_id}]\"\n else:\n record.name = \"\"\n return True\n\n\n@task_postrun.connect\ndef reset_tenant_id(\n sender: Any | None = None, # noqa: ARG001\n task_id: str | None = None, # noqa: ARG001\n task: Task | None = None, # noqa: ARG001\n args: tuple[Any, ...] | None = None, # noqa: ARG001\n kwargs: dict[str, Any] | None = None, # noqa: ARG001\n **other_kwargs: Any, # noqa: ARG001\n) -> None:\n \"\"\"Signal handler to reset tenant ID in context var after task ends.\"\"\"\n CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA)\n\n\ndef wait_for_vespa_or_shutdown(\n sender: Any, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Waits for Vespa to become ready subject to a timeout.\n Raises WorkerShutdown if the timeout is reached.\"\"\"\n\n if DISABLE_VECTOR_DB:\n logger.info(\n \"DISABLE_VECTOR_DB is set — skipping Vespa/OpenSearch readiness check.\"\n )\n return\n\n if not wait_for_vespa_with_timeout():\n msg = \"[Vespa] Readiness probe did not succeed within the timeout. Exiting...\"\n logger.error(msg)\n raise WorkerShutdown(msg)\n\n if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX:\n if not wait_for_opensearch_with_timeout():\n msg = \"[OpenSearch] Readiness probe did not succeed within the timeout. Exiting...\"\n logger.error(msg)\n raise WorkerShutdown(msg)\n\n\n# File for validating worker liveness\nclass LivenessProbe(bootsteps.StartStopStep):\n requires = {\"celery.worker.components:Timer\"}\n\n def __init__(self, worker: Any, **kwargs: Any) -> None:\n super().__init__(worker, **kwargs)\n self.requests: list[Any] = []\n self.task_tref = None\n self.path = make_probe_path(\"liveness\", worker.hostname)\n\n def start(self, worker: Any) -> None:\n self.task_tref = worker.timer.call_repeatedly(\n 15.0,\n self.update_liveness_file,\n (worker,),\n priority=10,\n )\n\n def stop(self, worker: Any) -> None: # noqa: ARG002\n self.path.unlink(missing_ok=True)\n if self.task_tref:\n self.task_tref.cancel()\n\n def update_liveness_file(self, worker: Any) -> None: # noqa: ARG002\n self.path.touch()\n\n\ndef get_bootsteps() -> list[type]:\n return [LivenessProbe]\n\n\n# Task modules that require a vector DB (Vespa/OpenSearch).\n# When DISABLE_VECTOR_DB is True these are excluded from autodiscover lists.\n_VECTOR_DB_TASK_MODULES: set[str] = {\n \"onyx.background.celery.tasks.connector_deletion\",\n \"onyx.background.celery.tasks.docprocessing\",\n \"onyx.background.celery.tasks.docfetching\",\n \"onyx.background.celery.tasks.pruning\",\n \"onyx.background.celery.tasks.vespa\",\n \"onyx.background.celery.tasks.opensearch_migration\",\n \"onyx.background.celery.tasks.doc_permission_syncing\",\n \"onyx.background.celery.tasks.hierarchyfetching\",\n # EE modules that are vector-DB-dependent\n \"ee.onyx.background.celery.tasks.doc_permission_syncing\",\n \"ee.onyx.background.celery.tasks.external_group_syncing\",\n}\n# NOTE: \"onyx.background.celery.tasks.shared\" is intentionally NOT in the set\n# above. It contains celery_beat_heartbeat (which only writes to Redis) alongside\n# document cleanup tasks. The cleanup tasks won't be invoked in minimal mode\n# because the periodic tasks that trigger them are in other filtered modules.\n\n\ndef filter_task_modules(modules: list[str]) -> list[str]:\n \"\"\"Remove vector-DB-dependent task modules when DISABLE_VECTOR_DB is True.\"\"\"\n if not DISABLE_VECTOR_DB:\n return modules\n return [m for m in modules if m not in _VECTOR_DB_TASK_MODULES]\n" }, { "path": "backend/onyx/background/celery/apps/beat.py", "content": "from datetime import timedelta\nfrom typing import Any\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery.beat import PersistentScheduler # type: ignore\nfrom celery.signals import beat_init\nfrom celery.utils.log import get_task_logger\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.background.celery.celery_utils import make_probe_path\nfrom onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT\nfrom onyx.configs.constants import POSTGRES_CELERY_BEAT_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.server.runtime.onyx_runtime import OnyxRuntime\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom shared_configs.configs import IGNORED_SYNCING_TENANT_LIST\nfrom shared_configs.configs import MULTI_TENANT\n\ntask_logger = get_task_logger(__name__)\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.beat\")\n\n\nclass DynamicTenantScheduler(PersistentScheduler):\n \"\"\"This scheduler is useful because we can dynamically adjust task generation rates\n through it.\"\"\"\n\n RELOAD_INTERVAL = 60\n\n def __init__(self, *args: Any, **kwargs: Any) -> None:\n super().__init__(*args, **kwargs)\n\n self.last_beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n self._reload_interval = timedelta(\n seconds=DynamicTenantScheduler.RELOAD_INTERVAL\n )\n self._last_reload = self.app.now() - self._reload_interval\n\n # Let the parent class handle store initialization\n self.setup_schedule()\n task_logger.info(\n f\"DynamicTenantScheduler initialized: reload_interval={self._reload_interval}\"\n )\n\n self._liveness_probe_path = make_probe_path(\"liveness\", \"beat@hostname\")\n\n # do not set the initial schedule here because we don't have db access yet.\n # do it in beat_init after the db engine is initialized\n\n # An initial schedule is required ... otherwise, the scheduler will delay\n # for 5 minutes before calling tick()\n\n def setup_schedule(self) -> None:\n super().setup_schedule()\n\n def tick(self) -> float:\n retval = super().tick()\n now = self.app.now()\n if (\n self._last_reload is None\n or (now - self._last_reload) > self._reload_interval\n ):\n task_logger.debug(\"Reload interval reached, initiating task update\")\n self._liveness_probe_path.touch()\n\n try:\n self._try_updating_schedule()\n except (AttributeError, KeyError):\n task_logger.exception(\"Failed to process task configuration\")\n except Exception:\n task_logger.exception(\"Unexpected error updating tasks\")\n\n self._last_reload = now\n\n return retval\n\n def _generate_schedule(\n self, tenant_ids: list[str] | list[None], beat_multiplier: float\n ) -> dict[str, dict[str, Any]]:\n \"\"\"Given a list of tenant id's, generates a new beat schedule for celery.\"\"\"\n new_schedule: dict[str, dict[str, Any]] = {}\n\n if MULTI_TENANT:\n # cloud tasks are system wide and thus only need to be on the beat schedule\n # once for all tenants\n get_cloud_tasks_to_schedule = fetch_versioned_implementation(\n \"onyx.background.celery.tasks.beat_schedule\",\n \"get_cloud_tasks_to_schedule\",\n )\n\n cloud_tasks_to_schedule: list[dict[str, Any]] = get_cloud_tasks_to_schedule(\n beat_multiplier\n )\n for task in cloud_tasks_to_schedule:\n task_name = task[\"name\"]\n cloud_task = {\n \"task\": task[\"task\"],\n \"schedule\": task[\"schedule\"],\n \"kwargs\": task.get(\"kwargs\", {}),\n }\n if options := task.get(\"options\"):\n task_logger.debug(f\"Adding options to task {task_name}: {options}\")\n cloud_task[\"options\"] = options\n new_schedule[task_name] = cloud_task\n\n # regular task beats are multiplied across all tenants\n # note that currently this just schedules for a single tenant in self hosted\n # and doesn't do anything in the cloud because it's much more scalable\n # to schedule a single cloud beat task to dispatch per tenant tasks.\n get_tasks_to_schedule = fetch_versioned_implementation(\n \"onyx.background.celery.tasks.beat_schedule\", \"get_tasks_to_schedule\"\n )\n\n tasks_to_schedule: list[dict[str, Any]] = get_tasks_to_schedule()\n\n for tenant_id in tenant_ids:\n if IGNORED_SYNCING_TENANT_LIST and tenant_id in IGNORED_SYNCING_TENANT_LIST:\n task_logger.debug(\n f\"Skipping tenant {tenant_id} as it is in the ignored syncing list\"\n )\n continue\n\n for task in tasks_to_schedule:\n task_name = task[\"name\"]\n tenant_task_name = f\"{task['name']}-{tenant_id}\"\n\n task_logger.debug(f\"Creating task configuration for {tenant_task_name}\")\n tenant_task = {\n \"task\": task[\"task\"],\n \"schedule\": task[\"schedule\"],\n \"kwargs\": {\"tenant_id\": tenant_id},\n }\n if options := task.get(\"options\"):\n task_logger.debug(\n f\"Adding options to task {tenant_task_name}: {options}\"\n )\n tenant_task[\"options\"] = options\n\n new_schedule[tenant_task_name] = tenant_task\n\n return new_schedule\n\n def _try_updating_schedule(self) -> None:\n \"\"\"Only updates the actual beat schedule on the celery app when it changes\"\"\"\n do_update = False\n\n task_logger.debug(\"_try_updating_schedule starting\")\n\n tenant_ids = get_all_tenant_ids()\n task_logger.debug(f\"Found {len(tenant_ids)} IDs\")\n\n # get current schedule and extract current tenants\n current_schedule = self.schedule.items()\n\n # get potential new state\n try:\n beat_multiplier = OnyxRuntime.get_beat_multiplier()\n except Exception:\n beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n new_schedule = self._generate_schedule(tenant_ids, beat_multiplier)\n\n # if the schedule or beat multiplier has changed, update\n while True:\n if beat_multiplier != self.last_beat_multiplier:\n do_update = True\n break\n\n if not DynamicTenantScheduler._compare_schedules(\n current_schedule, new_schedule\n ):\n do_update = True\n break\n\n break\n\n if not do_update:\n # exit early if nothing changed\n task_logger.info(\n f\"_try_updating_schedule - Schedule unchanged: tasks={len(new_schedule)} beat_multiplier={beat_multiplier}\"\n )\n return\n\n # schedule needs updating\n task_logger.debug(\n \"Schedule update required\",\n extra={\n \"new_tasks\": len(new_schedule),\n \"current_tasks\": len(current_schedule),\n },\n )\n\n # Create schedule entries\n entries = {}\n for name, entry in new_schedule.items():\n entries[name] = self.Entry(\n name=name,\n app=self.app,\n task=entry[\"task\"],\n schedule=entry[\"schedule\"],\n options=entry.get(\"options\", {}),\n kwargs=entry.get(\"kwargs\", {}),\n )\n\n # Update the schedule using the scheduler's methods\n self.schedule.clear()\n self.schedule.update(entries)\n\n # Ensure changes are persisted\n self.sync()\n\n task_logger.info(\n f\"_try_updating_schedule - Schedule updated: \"\n f\"prev_num_tasks={len(current_schedule)} \"\n f\"prev_beat_multiplier={self.last_beat_multiplier} \"\n f\"tasks={len(new_schedule)} \"\n f\"beat_multiplier={beat_multiplier}\"\n )\n\n self.last_beat_multiplier = beat_multiplier\n\n @staticmethod\n def _compare_schedules(schedule1: dict, schedule2: dict) -> bool:\n \"\"\"Compare schedules by task name only to determine if an update is needed.\n True if equivalent, False if not.\"\"\"\n current_tasks = set(name for name, _ in schedule1)\n new_tasks = set(schedule2.keys())\n return current_tasks == new_tasks\n\n\n@beat_init.connect\ndef on_beat_init(sender: Any, **kwargs: Any) -> None:\n task_logger.info(\"beat_init signal received.\")\n\n # Celery beat shouldn't touch the db at all. But just setting a low minimum here.\n SqlEngine.set_app_name(POSTGRES_CELERY_BEAT_APP_NAME)\n SqlEngine.init_engine(pool_size=2, max_overflow=0)\n\n app_base.wait_for_redis(sender, **kwargs)\n path = make_probe_path(\"readiness\", \"beat@hostname\")\n path.touch()\n task_logger.info(f\"Readiness signal touched at {path}.\")\n\n # first time init of the scheduler after db has been init'ed\n scheduler: DynamicTenantScheduler = sender.scheduler\n scheduler._try_updating_schedule()\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\ncelery_app.conf.beat_scheduler = DynamicTenantScheduler\ncelery_app.conf.task_default_base = app_base.TenantAwareTask\n" }, { "path": "backend/onyx/background/celery/apps/client.py", "content": "from celery import Celery\n\nimport onyx.background.celery.apps.app_base as app_base\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.client\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n" }, { "path": "backend/onyx/background/celery/apps/docfetching.py", "content": "from typing import Any\nfrom typing import cast\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.apps.worker import Worker\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_DOCFETCHING_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_postrun\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_prerun\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_rejected\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_retry\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_revoked\nfrom onyx.server.metrics.indexing_task_metrics import on_indexing_task_postrun\nfrom onyx.server.metrics.indexing_task_metrics import on_indexing_task_prerun\nfrom onyx.server.metrics.metrics_server import start_metrics_server\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.docfetching\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n on_celery_task_prerun(task_id, task)\n on_indexing_task_prerun(task_id, task, kwargs)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n on_celery_task_postrun(task_id, task, state)\n on_indexing_task_postrun(task_id, task, kwargs, state)\n\n\n@signals.task_retry.connect\ndef on_task_retry(sender: Any | None = None, **kwargs: Any) -> None: # noqa: ARG001\n # task_retry signal doesn't pass task_id in kwargs; get it from\n # the sender (the task instance) via sender.request.id.\n task_id = getattr(getattr(sender, \"request\", None), \"id\", None)\n on_celery_task_retry(task_id, sender)\n\n\n@signals.task_revoked.connect\ndef on_task_revoked(sender: Any | None = None, **kwargs: Any) -> None:\n task_name = getattr(sender, \"name\", None) or str(sender)\n on_celery_task_revoked(kwargs.get(\"task_id\"), task_name)\n\n\n@signals.task_rejected.connect\ndef on_task_rejected(sender: Any | None = None, **kwargs: Any) -> None: # noqa: ARG001\n # task_rejected sends the Consumer as sender, not the task instance.\n # The task name must be extracted from the Celery message headers.\n message = kwargs.get(\"message\")\n task_name: str | None = None\n if message is not None:\n headers = getattr(message, \"headers\", None) or {}\n task_name = headers.get(\"task\")\n if task_name is None:\n task_name = \"unknown\"\n on_celery_task_rejected(None, task_name)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n@worker_init.connect\ndef on_worker_init(sender: Worker, **kwargs: Any) -> None:\n logger.info(\"worker_init signal received.\")\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_DOCFETCHING_APP_NAME)\n pool_size = cast(int, sender.concurrency) # type: ignore\n SqlEngine.init_engine(pool_size=pool_size, max_overflow=8)\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n app_base.wait_for_vespa_or_shutdown(sender, **kwargs)\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n app_base.on_secondary_worker_init(sender, **kwargs)\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n start_metrics_server(\"docfetching\")\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.docfetching\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/apps/docprocessing.py", "content": "from typing import Any\nfrom typing import cast\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.apps.worker import Worker\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_process_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_DOCPROCESSING_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_postrun\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_prerun\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_rejected\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_retry\nfrom onyx.server.metrics.celery_task_metrics import on_celery_task_revoked\nfrom onyx.server.metrics.indexing_task_metrics import on_indexing_task_postrun\nfrom onyx.server.metrics.indexing_task_metrics import on_indexing_task_prerun\nfrom onyx.server.metrics.metrics_server import start_metrics_server\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.docprocessing\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n on_celery_task_prerun(task_id, task)\n on_indexing_task_prerun(task_id, task, kwargs)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n on_celery_task_postrun(task_id, task, state)\n on_indexing_task_postrun(task_id, task, kwargs, state)\n\n\n@signals.task_retry.connect\ndef on_task_retry(sender: Any | None = None, **kwargs: Any) -> None: # noqa: ARG001\n # task_retry signal doesn't pass task_id in kwargs; get it from\n # the sender (the task instance) via sender.request.id.\n task_id = getattr(getattr(sender, \"request\", None), \"id\", None)\n on_celery_task_retry(task_id, sender)\n\n\n@signals.task_revoked.connect\ndef on_task_revoked(sender: Any | None = None, **kwargs: Any) -> None:\n task_name = getattr(sender, \"name\", None) or str(sender)\n on_celery_task_revoked(kwargs.get(\"task_id\"), task_name)\n\n\n@signals.task_rejected.connect\ndef on_task_rejected(sender: Any | None = None, **kwargs: Any) -> None: # noqa: ARG001\n # task_rejected sends the Consumer as sender, not the task instance.\n # The task name must be extracted from the Celery message headers.\n message = kwargs.get(\"message\")\n task_name: str | None = None\n if message is not None:\n headers = getattr(message, \"headers\", None) or {}\n task_name = headers.get(\"task\")\n if task_name is None:\n task_name = \"unknown\"\n on_celery_task_rejected(None, task_name)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n@worker_init.connect\ndef on_worker_init(sender: Worker, **kwargs: Any) -> None:\n logger.info(\"worker_init signal received.\")\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_DOCPROCESSING_APP_NAME)\n\n # rkuo: Transient errors keep happening in the indexing watchdog threads.\n # \"SSL connection has been closed unexpectedly\"\n # actually setting the spawn method in the cloud fixes 95% of these.\n # setting pre ping might help even more, but not worrying about that yet\n pool_size = cast(int, sender.concurrency) # type: ignore\n SqlEngine.init_engine(pool_size=pool_size, max_overflow=8)\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n app_base.wait_for_vespa_or_shutdown(sender, **kwargs)\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n app_base.on_secondary_worker_init(sender, **kwargs)\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n start_metrics_server(\"docprocessing\")\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n# Note: worker_process_init only fires in prefork pool mode. Docprocessing uses\n# worker_pool=\"threads\" (see configs/docprocessing.py), so this handler is\n# effectively a no-op in normal operation. It remains as a safety net in case\n# the pool type is ever changed to prefork. Prometheus metrics are safe in\n# thread-pool mode since all threads share the same process memory and can\n# update the same Counter/Gauge/Histogram objects directly.\n@worker_process_init.connect\ndef init_worker(**kwargs: Any) -> None: # noqa: ARG001\n SqlEngine.reset_engine()\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.docprocessing\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/apps/heavy.py", "content": "from typing import Any\nfrom typing import cast\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.apps.worker import Worker\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_HEAVY_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.heavy\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n@worker_init.connect\ndef on_worker_init(sender: Worker, **kwargs: Any) -> None:\n logger.info(\"worker_init signal received.\")\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_HEAVY_APP_NAME)\n pool_size = cast(int, sender.concurrency) # type: ignore\n SqlEngine.init_engine(pool_size=pool_size, max_overflow=8)\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n app_base.wait_for_vespa_or_shutdown(sender, **kwargs)\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n app_base.on_secondary_worker_init(sender, **kwargs)\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.pruning\",\n # Sandbox tasks (file sync, cleanup)\n \"onyx.server.features.build.sandbox.tasks\",\n \"onyx.background.celery.tasks.hierarchyfetching\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/apps/light.py", "content": "from typing import Any\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.apps.worker import Worker\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.background.celery.celery_utils import httpx_init_vespa_pool\nfrom onyx.configs.app_configs import MANAGED_VESPA\nfrom onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH\nfrom onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_LIGHT_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.light\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n@worker_init.connect\ndef on_worker_init(sender: Worker, **kwargs: Any) -> None:\n EXTRA_CONCURRENCY = 8 # small extra fudge factor for connection limits\n\n logger.info(\"worker_init signal received.\")\n\n logger.info(f\"Concurrency: {sender.concurrency}\") # type: ignore\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_LIGHT_APP_NAME)\n SqlEngine.init_engine(pool_size=sender.concurrency, max_overflow=EXTRA_CONCURRENCY) # type: ignore\n\n if MANAGED_VESPA:\n httpx_init_vespa_pool(\n sender.concurrency + EXTRA_CONCURRENCY, # type: ignore\n ssl_cert=VESPA_CLOUD_CERT_PATH,\n ssl_key=VESPA_CLOUD_KEY_PATH,\n )\n else:\n httpx_init_vespa_pool(sender.concurrency + EXTRA_CONCURRENCY) # type: ignore\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n app_base.wait_for_vespa_or_shutdown(sender, **kwargs)\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n app_base.on_secondary_worker_init(sender, **kwargs)\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.shared\",\n \"onyx.background.celery.tasks.vespa\",\n \"onyx.background.celery.tasks.connector_deletion\",\n \"onyx.background.celery.tasks.doc_permission_syncing\",\n \"onyx.background.celery.tasks.docprocessing\",\n \"onyx.background.celery.tasks.opensearch_migration\",\n # Sandbox cleanup tasks (isolated in build feature)\n \"onyx.server.features.build.sandbox.tasks\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/apps/monitoring.py", "content": "import multiprocessing\nfrom typing import Any\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_MONITORING_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.monitoring\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: Any = None, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n# Set by on_worker_init so on_worker_ready knows whether to start the server.\n_prometheus_collectors_ok: bool = False\n\n\n@worker_init.connect\ndef on_worker_init(sender: Any, **kwargs: Any) -> None:\n global _prometheus_collectors_ok\n\n logger.info(\"worker_init signal received.\")\n logger.info(f\"Multiprocessing start method: {multiprocessing.get_start_method()}\")\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_MONITORING_APP_NAME)\n SqlEngine.init_engine(pool_size=sender.concurrency, max_overflow=3)\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n\n _prometheus_collectors_ok = _setup_prometheus_collectors(sender)\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n app_base.on_secondary_worker_init(sender, **kwargs)\n\n\ndef _setup_prometheus_collectors(sender: Any) -> bool:\n \"\"\"Register Prometheus collectors that need Redis/DB access.\n\n Passes the Celery app so the queue depth collector can obtain a fresh\n broker Redis client on each scrape (rather than holding a stale reference).\n\n Returns True if registration succeeded, False otherwise.\n \"\"\"\n try:\n from onyx.server.metrics.indexing_pipeline_setup import (\n setup_indexing_pipeline_metrics,\n )\n\n setup_indexing_pipeline_metrics(sender.app)\n logger.info(\"Prometheus indexing pipeline collectors registered\")\n return True\n except Exception:\n logger.exception(\"Failed to register Prometheus indexing pipeline collectors\")\n return False\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n if _prometheus_collectors_ok:\n from onyx.server.metrics.metrics_server import start_metrics_server\n\n start_metrics_server(\"monitoring\")\n else:\n logger.warning(\n \"Skipping Prometheus metrics server — collector registration failed\"\n )\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.monitoring\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/apps/primary.py", "content": "import logging\nimport os\nfrom typing import Any\nfrom typing import cast\n\nfrom celery import bootsteps # type: ignore\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.apps.worker import Worker\nfrom celery.exceptions import WorkerShutdown\nfrom celery.result import AsyncResult\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\nfrom redis.lock import Lock as RedisLock\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_utils import celery_is_worker_primary\nfrom onyx.background.celery.tasks.vespa.document_sync import reset_document_sync\nfrom onyx.configs.app_configs import CELERY_WORKER_PRIMARY_POOL_OVERFLOW\nfrom onyx.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.index_attempt import get_index_attempt\nfrom onyx.db.index_attempt import mark_attempt_canceled\nfrom onyx.db.indexing_coordination import IndexingCoordination\nfrom onyx.redis.redis_connector_delete import RedisConnectorDelete\nfrom onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync\nfrom onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync\nfrom onyx.redis.redis_connector_prune import RedisConnectorPrune\nfrom onyx.redis.redis_connector_stop import RedisConnectorStop\nfrom onyx.redis.redis_document_set import RedisDocumentSet\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_usergroup import RedisUserGroup\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.primary\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n@worker_init.connect\ndef on_worker_init(sender: Worker, **kwargs: Any) -> None:\n logger.info(\"worker_init signal received.\")\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME)\n pool_size = cast(int, sender.concurrency) # type: ignore\n SqlEngine.init_engine(\n pool_size=pool_size, max_overflow=CELERY_WORKER_PRIMARY_POOL_OVERFLOW\n )\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n app_base.wait_for_vespa_or_shutdown(sender, **kwargs)\n\n logger.info(f\"Running as the primary celery worker: pid={os.getpid()}\")\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n # This is singleton work that should be done on startup exactly once\n # by the primary worker. This is unnecessary in the multi tenant scenario\n r = get_redis_client(tenant_id=POSTGRES_DEFAULT_SCHEMA)\n\n # Log the role and slave count - being connected to a slave or slave count > 0 could be problematic\n replication_info: dict[str, Any] = cast(dict, r.info(\"replication\"))\n role: str = cast(str, replication_info.get(\"role\", \"\"))\n connected_slaves: int = replication_info.get(\"connected_slaves\", 0)\n\n logger.info(\n f\"Redis INFO REPLICATION: role={role} connected_slaves={connected_slaves}\"\n )\n\n memory_info: dict[str, Any] = cast(dict, r.info(\"memory\"))\n maxmemory_policy: str = cast(str, memory_info.get(\"maxmemory_policy\", \"\"))\n\n logger.info(f\"Redis INFO MEMORY: maxmemory_policy={maxmemory_policy}\")\n\n # For the moment, we're assuming that we are the only primary worker\n # that should be running.\n # TODO: maybe check for or clean up another zombie primary worker if we detect it\n r.delete(OnyxRedisLocks.PRIMARY_WORKER)\n\n # this process wide lock is taken to help other workers start up in order.\n # it is planned to use this lock to enforce singleton behavior on the primary\n # worker, since the primary worker does redis cleanup on startup, but this isn't\n # implemented yet.\n\n # set thread_local=False since we don't control what thread the periodic task might\n # reacquire the lock with\n lock: RedisLock = r.lock(\n OnyxRedisLocks.PRIMARY_WORKER,\n timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,\n thread_local=False,\n )\n\n logger.info(\"Primary worker lock: Acquire starting.\")\n acquired = lock.acquire(blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2)\n if acquired:\n logger.info(\"Primary worker lock: Acquire succeeded.\")\n else:\n logger.error(\"Primary worker lock: Acquire failed!\")\n raise WorkerShutdown(\"Primary worker lock could not be acquired!\")\n\n # tacking on our own user data to the sender\n sender.primary_worker_lock = lock # type: ignore\n\n # As currently designed, when this worker starts as \"primary\", we reinitialize redis\n # to a clean state (for our purposes, anyway)\n r.delete(OnyxRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK)\n\n r.delete(OnyxRedisConstants.ACTIVE_FENCES)\n\n # NOTE: we want to remove the `Redis*` classes, prefer to just have functions\n # This is the preferred way to do this going forward\n reset_document_sync(r)\n\n RedisDocumentSet.reset_all(r)\n RedisUserGroup.reset_all(r)\n RedisConnectorDelete.reset_all(r)\n RedisConnectorPrune.reset_all(r)\n RedisConnectorStop.reset_all(r)\n RedisConnectorPermissionSync.reset_all(r)\n RedisConnectorExternalGroupSync.reset_all(r)\n\n # mark orphaned index attempts as failed\n # This uses database coordination instead of Redis fencing\n with get_session_with_current_tenant() as db_session:\n # Get potentially orphaned attempts (those with active status and task IDs)\n potentially_orphaned_ids = IndexingCoordination.get_orphaned_index_attempt_ids(\n db_session\n )\n\n for attempt_id in potentially_orphaned_ids:\n attempt = get_index_attempt(db_session, attempt_id)\n\n # handle case where not started or docfetching is done but indexing is not\n if (\n not attempt\n or not attempt.celery_task_id\n or attempt.total_batches is not None\n ):\n continue\n\n # Check if the Celery task actually exists\n try:\n result: AsyncResult = AsyncResult(attempt.celery_task_id)\n\n # If the task is not in PENDING state, it exists in Celery\n if result.state != \"PENDING\":\n continue\n\n # Task is orphaned - mark as failed\n failure_reason = (\n f\"Orphaned index attempt found on startup - Celery task not found: \"\n f\"index_attempt={attempt.id} \"\n f\"cc_pair={attempt.connector_credential_pair_id} \"\n f\"search_settings={attempt.search_settings_id} \"\n f\"celery_task_id={attempt.celery_task_id}\"\n )\n logger.warning(failure_reason)\n mark_attempt_canceled(attempt.id, db_session, failure_reason)\n\n except Exception:\n # If we can't check the task status, be conservative and continue\n logger.warning(\n f\"Could not verify Celery task status on startup for attempt {attempt.id}, task_id={attempt.celery_task_id}\"\n )\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n # this can be spammy, so just enable it in the cloud for now\n if MULTI_TENANT:\n app_base.set_task_finished_log_level(logging.INFO)\n\n\nclass HubPeriodicTask(bootsteps.StartStopStep):\n \"\"\"Regularly reacquires the primary worker lock outside of the task queue.\n Use the task_logger in this class to avoid double logging.\n\n This cannot be done inside a regular beat task because it must run on schedule and\n a queue of existing work would starve the task from running.\n \"\"\"\n\n # it's unclear to me whether using the hub's timer or the bootstep timer is better\n requires = {\"celery.worker.components:Hub\"}\n\n def __init__(self, worker: Any, **kwargs: Any) -> None: # noqa: ARG002\n self.interval = CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 8 # Interval in seconds\n self.task_tref = None\n\n def start(self, worker: Any) -> None:\n if not celery_is_worker_primary(worker):\n return\n\n # Access the worker's event loop (hub)\n hub = worker.consumer.controller.hub\n\n # Schedule the periodic task\n self.task_tref = hub.call_repeatedly(\n self.interval, self.run_periodic_task, worker\n )\n task_logger.info(\"Scheduled periodic task with hub.\")\n\n def run_periodic_task(self, worker: Any) -> None:\n try:\n if not celery_is_worker_primary(worker):\n return\n\n if not hasattr(worker, \"primary_worker_lock\"):\n return\n\n lock: RedisLock = worker.primary_worker_lock\n\n r = get_redis_client(tenant_id=POSTGRES_DEFAULT_SCHEMA)\n\n if lock.owned():\n task_logger.debug(\"Reacquiring primary worker lock.\")\n lock.reacquire()\n else:\n task_logger.warning(\n \"Full acquisition of primary worker lock. Reasons could be worker restart or lock expiration.\"\n )\n lock = r.lock(\n OnyxRedisLocks.PRIMARY_WORKER,\n timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,\n )\n\n task_logger.info(\"Primary worker lock: Acquire starting.\")\n acquired = lock.acquire(\n blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2\n )\n if acquired:\n task_logger.info(\"Primary worker lock: Acquire succeeded.\")\n worker.primary_worker_lock = lock\n else:\n task_logger.error(\"Primary worker lock: Acquire failed!\")\n raise TimeoutError(\"Primary worker lock could not be acquired!\")\n\n except Exception:\n task_logger.exception(\"Periodic task failed.\")\n\n def stop(self, worker: Any) -> None: # noqa: ARG002\n # Cancel the scheduled task when the worker stops\n if self.task_tref:\n self.task_tref.cancel()\n task_logger.info(\"Canceled periodic task with hub.\")\n\n\ncelery_app.steps[\"worker\"].add(HubPeriodicTask)\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.connector_deletion\",\n \"onyx.background.celery.tasks.docprocessing\",\n \"onyx.background.celery.tasks.evals\",\n \"onyx.background.celery.tasks.hierarchyfetching\",\n \"onyx.background.celery.tasks.periodic\",\n \"onyx.background.celery.tasks.pruning\",\n \"onyx.background.celery.tasks.shared\",\n \"onyx.background.celery.tasks.vespa\",\n \"onyx.background.celery.tasks.llm_model_update\",\n \"onyx.background.celery.tasks.user_file_processing\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/apps/task_formatters.py", "content": "import logging\n\nfrom celery import current_task\n\nfrom onyx.utils.logger import ColoredFormatter\nfrom onyx.utils.logger import PlainFormatter\n\n\nclass CeleryTaskPlainFormatter(PlainFormatter):\n def format(self, record: logging.LogRecord) -> str:\n task = current_task\n if task and task.request:\n record.__dict__.update(task_id=task.request.id, task_name=task.name)\n record.msg = f\"[{task.name}({task.request.id})] {record.msg}\"\n\n return super().format(record)\n\n\nclass CeleryTaskColoredFormatter(ColoredFormatter):\n def format(self, record: logging.LogRecord) -> str:\n task = current_task\n if task and task.request:\n record.__dict__.update(task_id=task.request.id, task_name=task.name)\n record.msg = f\"[{task.name}({task.request.id})] {record.msg}\"\n\n return super().format(record)\n" }, { "path": "backend/onyx/background/celery/apps/user_file_processing.py", "content": "from typing import Any\nfrom typing import cast\n\nfrom celery import Celery\nfrom celery import signals\nfrom celery import Task\nfrom celery.apps.worker import Worker\nfrom celery.signals import celeryd_init\nfrom celery.signals import worker_init\nfrom celery.signals import worker_process_init\nfrom celery.signals import worker_ready\nfrom celery.signals import worker_shutdown\n\nimport onyx.background.celery.apps.app_base as app_base\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_USER_FILE_PROCESSING_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\n\nlogger = setup_logger()\n\ncelery_app = Celery(__name__)\ncelery_app.config_from_object(\"onyx.background.celery.configs.user_file_processing\")\ncelery_app.Task = app_base.TenantAwareTask # type: ignore [misc]\n\n\n@signals.task_prerun.connect\ndef on_task_prerun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)\n\n\n@signals.task_postrun.connect\ndef on_task_postrun(\n sender: Any | None = None,\n task_id: str | None = None,\n task: Task | None = None,\n args: tuple | None = None,\n kwargs: dict | None = None,\n retval: Any | None = None,\n state: str | None = None,\n **kwds: Any,\n) -> None:\n app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)\n\n\n@celeryd_init.connect\ndef on_celeryd_init(sender: str, conf: Any = None, **kwargs: Any) -> None:\n app_base.on_celeryd_init(sender, conf, **kwargs)\n\n\n@worker_init.connect\ndef on_worker_init(sender: Worker, **kwargs: Any) -> None:\n logger.info(\"worker_init signal received.\")\n\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_USER_FILE_PROCESSING_APP_NAME)\n\n # rkuo: Transient errors keep happening in the indexing watchdog threads.\n # \"SSL connection has been closed unexpectedly\"\n # actually setting the spawn method in the cloud fixes 95% of these.\n # setting pre ping might help even more, but not worrying about that yet\n pool_size = cast(int, sender.concurrency) # type: ignore\n SqlEngine.init_engine(pool_size=pool_size, max_overflow=8)\n\n app_base.wait_for_redis(sender, **kwargs)\n app_base.wait_for_db(sender, **kwargs)\n app_base.wait_for_vespa_or_shutdown(sender, **kwargs)\n\n # Less startup checks in multi-tenant case\n if MULTI_TENANT:\n return\n\n app_base.on_secondary_worker_init(sender, **kwargs)\n\n\n@worker_ready.connect\ndef on_worker_ready(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_ready(sender, **kwargs)\n\n\n@worker_shutdown.connect\ndef on_worker_shutdown(sender: Any, **kwargs: Any) -> None:\n app_base.on_worker_shutdown(sender, **kwargs)\n\n\n@worker_process_init.connect\ndef init_worker(**kwargs: Any) -> None: # noqa: ARG001\n SqlEngine.reset_engine()\n\n\n@signals.setup_logging.connect\ndef on_setup_logging(\n loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any\n) -> None:\n app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)\n\n\nbase_bootsteps = app_base.get_bootsteps()\nfor bootstep in base_bootsteps:\n celery_app.steps[\"worker\"].add(bootstep)\n\ncelery_app.autodiscover_tasks(\n app_base.filter_task_modules(\n [\n \"onyx.background.celery.tasks.user_file_processing\",\n ]\n )\n)\n" }, { "path": "backend/onyx/background/celery/celery_k8s_probe.py", "content": "# script to use as a kubernetes readiness / liveness probe\n\nimport argparse\nimport sys\nimport time\nfrom pathlib import Path\n\n\ndef main_readiness(filename: str) -> int:\n \"\"\"Checks if the file exists.\"\"\"\n path = Path(filename)\n if not path.is_file():\n return 1\n\n return 0\n\n\ndef main_liveness(filename: str) -> int:\n \"\"\"Checks if the file exists AND was recently modified.\"\"\"\n path = Path(filename)\n if not path.is_file():\n return 1\n\n stats = path.stat()\n liveness_timestamp = stats.st_mtime\n current_timestamp = time.time()\n time_diff = current_timestamp - liveness_timestamp\n if time_diff > 60:\n return 1\n\n return 0\n\n\nif __name__ == \"__main__\":\n exit_code: int\n\n parser = argparse.ArgumentParser(description=\"k8s readiness/liveness probe\")\n parser.add_argument(\n \"--probe\",\n type=str,\n choices=[\"readiness\", \"liveness\"],\n help=\"The type of probe\",\n required=True,\n )\n parser.add_argument(\"--filename\", help=\"The filename to watch\", required=True)\n args = parser.parse_args()\n\n if args.probe == \"readiness\":\n exit_code = main_readiness(args.filename)\n elif args.probe == \"liveness\":\n exit_code = main_liveness(args.filename)\n else:\n raise ValueError(f\"Unknown probe type: {args.probe}\")\n\n sys.exit(exit_code)\n" }, { "path": "backend/onyx/background/celery/celery_redis.py", "content": "# These are helper objects for tracking the keys we need to write in redis\nimport json\nimport threading\nfrom typing import Any\nfrom typing import cast\n\nfrom celery import Celery\nfrom redis import Redis\n\nfrom onyx.background.celery.configs.base import CELERY_SEPARATOR\nfrom onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS\n\n\n_broker_client: Redis | None = None\n_broker_url: str | None = None\n_broker_client_lock = threading.Lock()\n\n\ndef celery_get_broker_client(app: Celery) -> Redis:\n \"\"\"Return a shared Redis client connected to the Celery broker DB.\n\n Uses a module-level singleton so all tasks on a worker share one\n connection instead of creating a new one per call. The client\n connects directly to the broker Redis DB (parsed from the broker URL).\n\n Thread-safe via lock — safe for use in Celery thread-pool workers.\n\n Usage:\n r_celery = celery_get_broker_client(self.app)\n length = celery_get_queue_length(queue, r_celery)\n \"\"\"\n global _broker_client, _broker_url\n with _broker_client_lock:\n url = app.conf.broker_url\n if _broker_client is not None and _broker_url == url:\n try:\n _broker_client.ping()\n return _broker_client\n except Exception:\n try:\n _broker_client.close()\n except Exception:\n pass\n _broker_client = None\n elif _broker_client is not None:\n try:\n _broker_client.close()\n except Exception:\n pass\n _broker_client = None\n\n _broker_url = url\n _broker_client = Redis.from_url(\n url,\n decode_responses=False,\n health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,\n socket_keepalive=True,\n socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,\n retry_on_timeout=True,\n )\n return _broker_client\n\n\ndef celery_get_unacked_length(r: Redis) -> int:\n \"\"\"Checking the unacked queue is useful because a non-zero length tells us there\n may be prefetched tasks.\n\n There can be other tasks in here besides indexing tasks, so this is mostly useful\n just to see if the task count is non zero.\n\n ref: https://blog.hikaru.run/2022/08/29/get-waiting-tasks-count-in-celery.html\n \"\"\"\n length = cast(int, r.hlen(\"unacked\"))\n return length\n\n\ndef celery_get_unacked_task_ids(queue: str, r: Redis) -> set[str]:\n \"\"\"Gets the set of task id's matching the given queue in the unacked hash.\n\n Unacked entries belonging to the indexing queues are \"prefetched\", so this gives\n us crucial visibility as to what tasks are in that state.\n \"\"\"\n tasks: set[str] = set()\n\n for _, v in r.hscan_iter(\"unacked\"):\n v_bytes = cast(bytes, v)\n v_str = v_bytes.decode(\"utf-8\")\n task = json.loads(v_str)\n\n task_description = task[0]\n task_queue = task[2]\n\n if task_queue != queue:\n continue\n\n task_id = task_description.get(\"headers\", {}).get(\"id\")\n if not task_id:\n continue\n\n # if the queue matches and we see the task_id, add it\n tasks.add(task_id)\n return tasks\n\n\ndef celery_get_queue_length(queue: str, r: Redis) -> int:\n \"\"\"This is a redis specific way to get the length of a celery queue.\n It is priority aware and knows how to count across the multiple redis lists\n used to implement task prioritization.\n This operation is not atomic.\"\"\"\n total_length = 0\n for i in range(len(OnyxCeleryPriority)):\n queue_name = queue\n if i > 0:\n queue_name += CELERY_SEPARATOR\n queue_name += str(i)\n\n length = r.llen(queue_name)\n total_length += cast(int, length)\n\n return total_length\n\n\ndef celery_find_task(task_id: str, queue: str, r: Redis) -> int:\n \"\"\"This is a redis specific way to find a task for a particular queue in redis.\n It is priority aware and knows how to look through the multiple redis lists\n used to implement task prioritization.\n This operation is not atomic.\n\n This is a linear search O(n) ... so be careful using it when the task queues can be larger.\n\n Returns true if the id is in the queue, False if not.\n \"\"\"\n for priority in range(len(OnyxCeleryPriority)):\n queue_name = f\"{queue}{CELERY_SEPARATOR}{priority}\" if priority > 0 else queue\n\n tasks = cast(list[bytes], r.lrange(queue_name, 0, -1))\n for task in tasks:\n task_dict: dict[str, Any] = json.loads(task.decode(\"utf-8\"))\n if task_dict.get(\"headers\", {}).get(\"id\") == task_id:\n return True\n\n return False\n\n\ndef celery_get_queued_task_ids(queue: str, r: Redis) -> set[str]:\n \"\"\"This is a redis specific way to build a list of tasks in a queue and return them\n as a set.\n\n This helps us read the queue once and then efficiently look for missing tasks\n in the queue.\n \"\"\"\n\n task_set: set[str] = set()\n\n for priority in range(len(OnyxCeleryPriority)):\n queue_name = f\"{queue}{CELERY_SEPARATOR}{priority}\" if priority > 0 else queue\n\n tasks = cast(list[bytes], r.lrange(queue_name, 0, -1))\n for task in tasks:\n task_dict: dict[str, Any] = json.loads(task.decode(\"utf-8\"))\n task_id = task_dict.get(\"headers\", {}).get(\"id\")\n if task_id:\n task_set.add(task_id)\n\n return task_set\n\n\ndef celery_inspect_get_workers(name_filter: str | None, app: Celery) -> list[str]:\n \"\"\"Returns a list of current workers containing name_filter, or all workers if\n name_filter is None.\n\n We've empirically discovered that the celery inspect API is potentially unstable\n and may hang or return empty results when celery is under load. Suggest using this\n more to debug and troubleshoot than in production code.\n \"\"\"\n worker_names: list[str] = []\n\n # filter for and create an indexing specific inspect object\n inspect = app.control.inspect()\n workers: dict[str, Any] = inspect.ping() # type: ignore\n if workers:\n for worker_name in list(workers.keys()):\n # if the name filter not set, return all worker names\n if not name_filter:\n worker_names.append(worker_name)\n continue\n\n # if the name filter is set, return only worker names that contain the name filter\n if name_filter not in worker_name:\n continue\n\n worker_names.append(worker_name)\n\n return worker_names\n\n\ndef celery_inspect_get_reserved(worker_names: list[str], app: Celery) -> set[str]:\n \"\"\"Returns a list of reserved tasks on the specified workers.\n\n We've empirically discovered that the celery inspect API is potentially unstable\n and may hang or return empty results when celery is under load. Suggest using this\n more to debug and troubleshoot than in production code.\n \"\"\"\n reserved_task_ids: set[str] = set()\n\n inspect = app.control.inspect(destination=worker_names)\n\n # get the list of reserved tasks\n reserved_tasks: dict[str, list] | None = inspect.reserved() # type: ignore\n if reserved_tasks:\n for _, task_list in reserved_tasks.items():\n for task in task_list:\n reserved_task_ids.add(task[\"id\"])\n\n return reserved_task_ids\n\n\ndef celery_inspect_get_active(worker_names: list[str], app: Celery) -> set[str]:\n \"\"\"Returns a list of active tasks on the specified workers.\n\n We've empirically discovered that the celery inspect API is potentially unstable\n and may hang or return empty results when celery is under load. Suggest using this\n more to debug and troubleshoot than in production code.\n \"\"\"\n active_task_ids: set[str] = set()\n\n inspect = app.control.inspect(destination=worker_names)\n\n # get the list of reserved tasks\n active_tasks: dict[str, list] | None = inspect.active() # type: ignore\n if active_tasks:\n for _, task_list in active_tasks.items():\n for task in task_list:\n active_task_ids.add(task[\"id\"])\n\n return active_task_ids\n" }, { "path": "backend/onyx/background/celery/celery_utils.py", "content": "from collections.abc import Generator\nfrom collections.abc import Iterator\nfrom collections.abc import Sequence\nfrom datetime import datetime\nfrom datetime import timezone\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import cast\nfrom typing import TypeVar\n\nimport httpx\nfrom pydantic import BaseModel\n\nfrom onyx.configs.app_configs import MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE\nfrom onyx.configs.app_configs import VESPA_REQUEST_TIMEOUT\nfrom onyx.connectors.connector_runner import CheckpointOutputWrapper\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.interfaces import BaseConnector\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import ConnectorCheckpoint\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SlimConnector\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.httpx.httpx_pool import HttpxPool\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\nCT = TypeVar(\"CT\", bound=ConnectorCheckpoint)\n\n\nclass SlimConnectorExtractionResult(BaseModel):\n \"\"\"Result of extracting document IDs and hierarchy nodes from a connector.\n\n raw_id_to_parent maps document ID → parent_hierarchy_raw_node_id (or None).\n Use raw_id_to_parent.keys() wherever the old set of IDs was needed.\n \"\"\"\n\n raw_id_to_parent: dict[str, str | None]\n hierarchy_nodes: list[HierarchyNode]\n\n\ndef _checkpointed_batched_items(\n connector: CheckpointedConnector[CT],\n start: float,\n end: float,\n) -> Generator[list[Document | HierarchyNode | ConnectorFailure], None, None]:\n \"\"\"Loop through all checkpoint steps and yield batched items.\n\n Some checkpointed connectors (e.g. IMAP) are multi-step: the first\n checkpoint call may only initialize internal state without yielding\n any documents. This function loops until checkpoint.has_more is False\n to ensure all items are collected across every step.\n \"\"\"\n checkpoint = connector.build_dummy_checkpoint()\n while True:\n checkpoint_output = connector.load_from_checkpoint(\n start=start, end=end, checkpoint=checkpoint\n )\n wrapper: CheckpointOutputWrapper[CT] = CheckpointOutputWrapper()\n batch: list[Document | HierarchyNode | ConnectorFailure] = []\n for document, hierarchy_node, failure, next_checkpoint in wrapper(\n checkpoint_output\n ):\n if document is not None:\n batch.append(document)\n elif hierarchy_node is not None:\n batch.append(hierarchy_node)\n elif failure is not None:\n batch.append(failure)\n\n if next_checkpoint is not None:\n checkpoint = next_checkpoint\n\n if batch:\n yield batch\n\n if not checkpoint.has_more:\n break\n\n\ndef _get_failure_id(failure: ConnectorFailure) -> str | None:\n \"\"\"Extract the document/entity ID from a ConnectorFailure.\"\"\"\n if failure.failed_document:\n return failure.failed_document.document_id\n if failure.failed_entity:\n return failure.failed_entity.entity_id\n return None\n\n\nclass BatchResult(BaseModel):\n raw_id_to_parent: dict[str, str | None]\n hierarchy_nodes: list[HierarchyNode]\n\n\ndef _extract_from_batch(\n doc_list: Sequence[Document | SlimDocument | HierarchyNode | ConnectorFailure],\n) -> BatchResult:\n \"\"\"Separate a batch into document IDs (with parent mapping) and hierarchy nodes.\n\n ConnectorFailure items have their failed document/entity IDs added to the\n ID dict so that failed-to-retrieve documents are not accidentally pruned.\n \"\"\"\n ids: dict[str, str | None] = {}\n hierarchy_nodes: list[HierarchyNode] = []\n for item in doc_list:\n if isinstance(item, HierarchyNode):\n hierarchy_nodes.append(item)\n elif isinstance(item, ConnectorFailure):\n failed_id = _get_failure_id(item)\n if failed_id:\n ids[failed_id] = None\n logger.warning(\n f\"Failed to retrieve document {failed_id}: {item.failure_message}\"\n )\n else:\n ids[item.id] = item.parent_hierarchy_raw_node_id\n return BatchResult(raw_id_to_parent=ids, hierarchy_nodes=hierarchy_nodes)\n\n\ndef extract_ids_from_runnable_connector(\n runnable_connector: BaseConnector,\n callback: IndexingHeartbeatInterface | None = None,\n) -> SlimConnectorExtractionResult:\n \"\"\"\n Extract document IDs and hierarchy nodes from a runnable connector.\n\n Hierarchy nodes yielded alongside documents/slim docs are collected and\n returned in the result. ConnectorFailure items have their IDs preserved\n so that failed-to-retrieve documents are not accidentally pruned.\n\n Optionally, a callback can be passed to handle the length of each document batch.\n \"\"\"\n all_raw_id_to_parent: dict[str, str | None] = {}\n all_hierarchy_nodes: list[HierarchyNode] = []\n\n # Sequence (covariant) lets all the specific list[...] iterator types unify here\n raw_batch_generator: (\n Iterator[Sequence[Document | SlimDocument | HierarchyNode | ConnectorFailure]]\n | None\n ) = None\n\n if isinstance(runnable_connector, SlimConnector):\n raw_batch_generator = runnable_connector.retrieve_all_slim_docs()\n elif isinstance(runnable_connector, SlimConnectorWithPermSync):\n raw_batch_generator = runnable_connector.retrieve_all_slim_docs_perm_sync()\n # If the connector isn't slim, fall back to running it normally to get ids\n elif isinstance(runnable_connector, LoadConnector):\n raw_batch_generator = runnable_connector.load_from_state()\n elif isinstance(runnable_connector, PollConnector):\n start = datetime(1970, 1, 1, tzinfo=timezone.utc).timestamp()\n end = datetime.now(timezone.utc).timestamp()\n raw_batch_generator = runnable_connector.poll_source(start=start, end=end)\n elif isinstance(runnable_connector, CheckpointedConnector):\n start = datetime(1970, 1, 1, tzinfo=timezone.utc).timestamp()\n end = datetime.now(timezone.utc).timestamp()\n raw_batch_generator = _checkpointed_batched_items(\n runnable_connector, start, end\n )\n else:\n raise RuntimeError(\"Pruning job could not find a valid runnable_connector.\")\n\n # this function is called per batch for rate limiting\n doc_batch_processing_func = (\n rate_limit_builder(\n max_calls=MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE, period=60\n )(lambda x: x)\n if MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE\n else lambda x: x\n )\n\n # process raw batches to extract both IDs and hierarchy nodes\n for doc_list in raw_batch_generator:\n if callback and callback.should_stop():\n raise RuntimeError(\n \"extract_ids_from_runnable_connector: Stop signal detected\"\n )\n\n batch_result = _extract_from_batch(doc_list)\n batch_ids = batch_result.raw_id_to_parent\n batch_nodes = batch_result.hierarchy_nodes\n doc_batch_processing_func(batch_ids)\n all_raw_id_to_parent.update(batch_ids)\n all_hierarchy_nodes.extend(batch_nodes)\n\n if callback:\n callback.progress(\"extract_ids_from_runnable_connector\", len(batch_ids))\n\n return SlimConnectorExtractionResult(\n raw_id_to_parent=all_raw_id_to_parent,\n hierarchy_nodes=all_hierarchy_nodes,\n )\n\n\ndef celery_is_listening_to_queue(worker: Any, name: str) -> bool:\n \"\"\"Checks to see if we're listening to the named queue\"\"\"\n\n # how to get a list of queues this worker is listening to\n # https://stackoverflow.com/questions/29790523/how-to-determine-which-queues-a-celery-worker-is-consuming-at-runtime\n queue_names = list(worker.app.amqp.queues.consume_from.keys())\n for queue_name in queue_names:\n if queue_name == name:\n return True\n\n return False\n\n\ndef celery_is_worker_primary(worker: Any) -> bool:\n \"\"\"There are multiple approaches that could be taken to determine if a celery worker\n is 'primary', as defined by us. But the way we do it is to check the hostname set\n for the celery worker, which can be done on the\n command line with '--hostname'.\"\"\"\n hostname = worker.hostname\n if hostname.startswith(\"primary\"):\n return True\n\n return False\n\n\ndef httpx_init_vespa_pool(\n max_keepalive_connections: int,\n timeout: int = VESPA_REQUEST_TIMEOUT,\n ssl_cert: str | None = None,\n ssl_key: str | None = None,\n) -> None:\n httpx_cert = None\n httpx_verify = False\n if ssl_cert and ssl_key:\n httpx_cert = cast(tuple[str, str], (ssl_cert, ssl_key))\n httpx_verify = True\n\n HttpxPool.init_client(\n name=\"vespa\",\n cert=httpx_cert,\n verify=httpx_verify,\n timeout=timeout,\n http2=False,\n limits=httpx.Limits(max_keepalive_connections=max_keepalive_connections),\n )\n\n\ndef make_probe_path(probe: str, hostname: str) -> Path:\n \"\"\"templates the path for a k8s probe file.\n\n e.g. /tmp/onyx_k8s_indexing_readiness.txt\n \"\"\"\n hostname_parts = hostname.split(\"@\")\n if len(hostname_parts) != 2:\n raise ValueError(f\"hostname could not be split! {hostname=}\")\n\n name = hostname_parts[0]\n if not name:\n raise ValueError(f\"name cannot be empty! {name=}\")\n\n safe_name = \"\".join(c for c in name if c.isalnum()).rstrip()\n return Path(f\"/tmp/onyx_k8s_{safe_name}_{probe}.txt\")\n" }, { "path": "backend/onyx/background/celery/configs/base.py", "content": "# docs: https://docs.celeryq.dev/en/stable/userguide/configuration.html\nimport urllib.parse\n\nfrom onyx.configs.app_configs import CELERY_BROKER_POOL_LIMIT\nfrom onyx.configs.app_configs import CELERY_RESULT_EXPIRES\nfrom onyx.configs.app_configs import REDIS_DB_NUMBER_CELERY\nfrom onyx.configs.app_configs import REDIS_DB_NUMBER_CELERY_RESULT_BACKEND\nfrom onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL\nfrom onyx.configs.app_configs import REDIS_HOST\nfrom onyx.configs.app_configs import REDIS_PASSWORD\nfrom onyx.configs.app_configs import REDIS_PORT\nfrom onyx.configs.app_configs import REDIS_SSL\nfrom onyx.configs.app_configs import REDIS_SSL_CA_CERTS\nfrom onyx.configs.app_configs import REDIS_SSL_CERT_REQS\nfrom onyx.configs.app_configs import USE_REDIS_IAM_AUTH\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS\n\nCELERY_SEPARATOR = \":\"\n\nCELERY_PASSWORD_PART = \"\"\nif REDIS_PASSWORD:\n CELERY_PASSWORD_PART = \":\" + urllib.parse.quote(REDIS_PASSWORD, safe=\"\") + \"@\"\n\nREDIS_SCHEME = \"redis\"\n\n# SSL-specific query parameters for Redis URL\nSSL_QUERY_PARAMS = \"\"\nif REDIS_SSL and not USE_REDIS_IAM_AUTH:\n REDIS_SCHEME = \"rediss\"\n SSL_QUERY_PARAMS = f\"?ssl_cert_reqs={REDIS_SSL_CERT_REQS}\"\n if REDIS_SSL_CA_CERTS:\n SSL_QUERY_PARAMS += f\"&ssl_ca_certs={REDIS_SSL_CA_CERTS}\"\n\n# region Broker settings\n# example celery_broker_url: \"redis://:password@localhost:6379/15\"\nbroker_url = f\"{REDIS_SCHEME}://{CELERY_PASSWORD_PART}{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB_NUMBER_CELERY}{SSL_QUERY_PARAMS}\"\n\nbroker_connection_retry_on_startup = True\nbroker_pool_limit = CELERY_BROKER_POOL_LIMIT\n\n# redis broker settings\n# https://docs.celeryq.dev/projects/kombu/en/stable/reference/kombu.transport.redis.html\nbroker_transport_options = {\n \"priority_steps\": list(range(len(OnyxCeleryPriority))),\n \"sep\": CELERY_SEPARATOR,\n \"queue_order_strategy\": \"priority\",\n \"retry_on_timeout\": True,\n \"health_check_interval\": REDIS_HEALTH_CHECK_INTERVAL,\n \"socket_keepalive\": True,\n \"socket_keepalive_options\": REDIS_SOCKET_KEEPALIVE_OPTIONS,\n}\n# endregion\n\n# redis backend settings\n# https://docs.celeryq.dev/en/stable/userguide/configuration.html#redis-backend-settings\n\n# there doesn't appear to be a way to set socket_keepalive_options on the redis result backend\nredis_socket_keepalive = True\nredis_retry_on_timeout = True\nredis_backend_health_check_interval = REDIS_HEALTH_CHECK_INTERVAL\n\n\ntask_default_priority = OnyxCeleryPriority.MEDIUM\ntask_acks_late = True\n\n# region Task result backend settings\n# It's possible we don't even need celery's result backend, in which case all of the optimization below\n# might be irrelevant\nresult_backend = f\"{REDIS_SCHEME}://{CELERY_PASSWORD_PART}{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB_NUMBER_CELERY_RESULT_BACKEND}{SSL_QUERY_PARAMS}\"\nresult_expires = CELERY_RESULT_EXPIRES # 86400 seconds is the default\n# endregion\n\n# Leaving this to the default of True may cause double logging since both our own app\n# and celery think they are controlling the logger.\n# TODO: Configure celery's logger entirely manually and set this to False\n# worker_hijack_root_logger = False\n\n# region Notes on serialization performance\n# Option 0: Defaults (json serializer, no compression)\n# about 1.5 KB per queued task. 1KB in queue, 400B for result, 100 as a child entry in generator result\n\n# Option 1: Reduces generator task result sizes by roughly 20%\n# task_compression = \"bzip2\"\n# task_serializer = \"pickle\"\n# result_compression = \"bzip2\"\n# result_serializer = \"pickle\"\n# accept_content=[\"pickle\"]\n\n# Option 2: this significantly reduces the size of the result for generator tasks since the list of children\n# can be large. small tasks change very little\n# def pickle_bz2_encoder(data):\n# return bz2.compress(pickle.dumps(data))\n\n# def pickle_bz2_decoder(data):\n# return pickle.loads(bz2.decompress(data))\n\n# from kombu import serialization # To register custom serialization with Celery/Kombu\n\n# serialization.register('pickle-bzip2', pickle_bz2_encoder, pickle_bz2_decoder, 'application/x-pickle-bz2', 'binary')\n\n# task_serializer = \"pickle-bzip2\"\n# result_serializer = \"pickle-bzip2\"\n# accept_content=[\"pickle\", \"pickle-bzip2\"]\n# endregion\n" }, { "path": "backend/onyx/background/celery/configs/beat.py", "content": "# docs: https://docs.celeryq.dev/en/stable/userguide/configuration.html\nimport onyx.background.celery.configs.base as shared_config\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n" }, { "path": "backend/onyx/background/celery/configs/client.py", "content": "import onyx.background.celery.configs.base as shared_config\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n" }, { "path": "backend/onyx/background/celery/configs/docfetching.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_DOCFETCHING_CONCURRENCY\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\n# Docfetching worker configuration\nworker_concurrency = CELERY_WORKER_DOCFETCHING_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = 1\n" }, { "path": "backend/onyx/background/celery/configs/docprocessing.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_DOCPROCESSING_CONCURRENCY\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\n# Indexing worker specific ... this lets us track the transition to STARTED in redis\n# We don't currently rely on this but it has the potential to be useful and\n# indexing tasks are not high volume\n\n# we don't turn this on yet because celery occasionally runs tasks more than once\n# which means a duplicate run might change the task state unexpectedly\n# task_track_started = True\n\nworker_concurrency = CELERY_WORKER_DOCPROCESSING_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = 1\n" }, { "path": "backend/onyx/background/celery/configs/heavy.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_HEAVY_CONCURRENCY\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\nworker_concurrency = CELERY_WORKER_HEAVY_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = 1\n" }, { "path": "backend/onyx/background/celery/configs/light.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_LIGHT_CONCURRENCY\nfrom onyx.configs.app_configs import CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\nworker_concurrency = CELERY_WORKER_LIGHT_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER\n" }, { "path": "backend/onyx/background/celery/configs/monitoring.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_MONITORING_CONCURRENCY\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\n# Monitoring worker specific settings\nworker_concurrency = CELERY_WORKER_MONITORING_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = 1\n" }, { "path": "backend/onyx/background/celery/configs/primary.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_PRIMARY_CONCURRENCY\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\nworker_concurrency = CELERY_WORKER_PRIMARY_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = 1\n" }, { "path": "backend/onyx/background/celery/configs/user_file_processing.py", "content": "import onyx.background.celery.configs.base as shared_config\nfrom onyx.configs.app_configs import CELERY_WORKER_USER_FILE_PROCESSING_CONCURRENCY\n\nbroker_url = shared_config.broker_url\nbroker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup\nbroker_pool_limit = shared_config.broker_pool_limit\nbroker_transport_options = shared_config.broker_transport_options\n\nredis_socket_keepalive = shared_config.redis_socket_keepalive\nredis_retry_on_timeout = shared_config.redis_retry_on_timeout\nredis_backend_health_check_interval = shared_config.redis_backend_health_check_interval\n\nresult_backend = shared_config.result_backend\nresult_expires = shared_config.result_expires # 86400 seconds is the default\n\ntask_default_priority = shared_config.task_default_priority\ntask_acks_late = shared_config.task_acks_late\n\n# User file processing worker configuration\nworker_concurrency = CELERY_WORKER_USER_FILE_PROCESSING_CONCURRENCY\nworker_pool = \"threads\"\nworker_prefetch_multiplier = 1\n" }, { "path": "backend/onyx/background/celery/memory_monitoring.py", "content": "# backend/onyx/background/celery/memory_monitoring.py\nimport logging\nimport os\nfrom logging.handlers import RotatingFileHandler\n\nimport psutil\n\nfrom onyx.utils.logger import is_running_in_container\nfrom onyx.utils.logger import setup_logger\n\n# Regular application logger\nlogger = setup_logger()\n\n# Only set up memory monitoring in container environment\nif is_running_in_container():\n # Set up a dedicated memory monitoring logger\n MEMORY_LOG_DIR = \"/var/log/onyx/memory\"\n MEMORY_LOG_FILE = os.path.join(MEMORY_LOG_DIR, \"memory_usage.log\")\n MEMORY_LOG_MAX_BYTES = 10 * 1024 * 1024 # 10MB\n MEMORY_LOG_BACKUP_COUNT = 5 # Keep 5 backup files\n\n # Ensure log directory exists\n os.makedirs(MEMORY_LOG_DIR, exist_ok=True)\n\n # Create a dedicated logger for memory monitoring\n memory_logger = logging.getLogger(\"memory_monitoring\")\n memory_logger.setLevel(logging.INFO)\n\n # Create a rotating file handler\n memory_handler = RotatingFileHandler(\n MEMORY_LOG_FILE,\n maxBytes=MEMORY_LOG_MAX_BYTES,\n backupCount=MEMORY_LOG_BACKUP_COUNT,\n )\n\n # Create a formatter that includes all relevant information\n memory_formatter = logging.Formatter(\n \"%(asctime)s [%(levelname)s] %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\"\n )\n memory_handler.setFormatter(memory_formatter)\n memory_logger.addHandler(memory_handler)\nelse:\n # Create a null logger when not in container\n memory_logger = logging.getLogger(\"memory_monitoring\")\n memory_logger.addHandler(logging.NullHandler())\n\n\ndef emit_process_memory(\n pid: int, process_name: str, additional_metadata: dict[str, str | int]\n) -> None:\n # Skip memory monitoring if not in container\n if not is_running_in_container():\n return\n\n try:\n process = psutil.Process(pid)\n memory_info = process.memory_info()\n cpu_percent = process.cpu_percent(interval=0.1)\n\n # Build metadata string from additional_metadata dictionary\n metadata_str = \" \".join(\n [f\"{key}={value}\" for key, value in additional_metadata.items()]\n )\n metadata_str = f\" {metadata_str}\" if metadata_str else \"\"\n\n memory_logger.info(\n f\"PROCESS_MEMORY process_name={process_name} pid={pid} \"\n f\"rss_mb={memory_info.rss / (1024 * 1024):.2f} \"\n f\"vms_mb={memory_info.vms / (1024 * 1024):.2f} \"\n f\"cpu={cpu_percent:.2f}{metadata_str}\"\n )\n except Exception:\n logger.exception(\"Error monitoring process memory.\")\n" }, { "path": "backend/onyx/background/celery/tasks/beat_schedule.py", "content": "import copy\nfrom datetime import timedelta\nfrom typing import Any\n\nfrom celery.schedules import crontab\n\nfrom onyx.configs.app_configs import AUTO_LLM_CONFIG_URL\nfrom onyx.configs.app_configs import AUTO_LLM_UPDATE_INTERVAL_SECONDS\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\nfrom onyx.configs.app_configs import ENTERPRISE_EDITION_ENABLED\nfrom onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES\nfrom onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom shared_configs.configs import MULTI_TENANT\n\n# choosing 15 minutes because it roughly gives us enough time to process many tasks\n# we might be able to reduce this greatly if we can run a unified\n# loop across all tenants rather than tasks per tenant\n\n# we set expires because it isn't necessary to queue up these tasks\n# it's only important that they run relatively regularly\nBEAT_EXPIRES_DEFAULT = 15 * 60 # 15 minutes (in seconds)\n\n# hack to slow down task dispatch in the cloud until\n# we have a better implementation (backpressure, etc)\n# Note that DynamicTenantScheduler can adjust the runtime value for this via Redis\nCLOUD_BEAT_MULTIPLIER_DEFAULT = 8.0\nCLOUD_DOC_PERMISSION_SYNC_MULTIPLIER_DEFAULT = 1.0\n\n# tasks that run in either self-hosted on cloud\nbeat_task_templates: list[dict] = [\n {\n \"name\": \"check-for-user-file-processing\",\n \"task\": OnyxCeleryTask.CHECK_FOR_USER_FILE_PROCESSING,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-user-file-project-sync\",\n \"task\": OnyxCeleryTask.CHECK_FOR_USER_FILE_PROJECT_SYNC,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-user-file-delete\",\n \"task\": OnyxCeleryTask.CHECK_FOR_USER_FILE_DELETE,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-indexing\",\n \"task\": OnyxCeleryTask.CHECK_FOR_INDEXING,\n \"schedule\": timedelta(seconds=15),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-checkpoint-cleanup\",\n \"task\": OnyxCeleryTask.CHECK_FOR_CHECKPOINT_CLEANUP,\n \"schedule\": timedelta(hours=1),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-index-attempt-cleanup\",\n \"task\": OnyxCeleryTask.CHECK_FOR_INDEX_ATTEMPT_CLEANUP,\n \"schedule\": timedelta(minutes=30),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-connector-deletion\",\n \"task\": OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-vespa-sync\",\n \"task\": OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-pruning\",\n \"task\": OnyxCeleryTask.CHECK_FOR_PRUNING,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-hierarchy-fetching\",\n \"task\": OnyxCeleryTask.CHECK_FOR_HIERARCHY_FETCHING,\n \"schedule\": timedelta(hours=1), # Check hourly, but only fetch once per day\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"monitor-background-processes\",\n \"task\": OnyxCeleryTask.MONITOR_BACKGROUND_PROCESSES,\n \"schedule\": timedelta(minutes=5),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.MONITORING,\n },\n },\n # Sandbox cleanup tasks\n {\n \"name\": \"cleanup-idle-sandboxes\",\n \"task\": OnyxCeleryTask.CLEANUP_IDLE_SANDBOXES,\n \"schedule\": timedelta(minutes=1),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.SANDBOX,\n },\n },\n {\n \"name\": \"cleanup-old-snapshots\",\n \"task\": OnyxCeleryTask.CLEANUP_OLD_SNAPSHOTS,\n \"schedule\": timedelta(hours=24),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.SANDBOX,\n },\n },\n]\n\nif ENTERPRISE_EDITION_ENABLED:\n beat_task_templates.extend(\n [\n {\n \"name\": \"check-for-doc-permissions-sync\",\n \"task\": OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,\n \"schedule\": timedelta(seconds=30),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": \"check-for-external-group-sync\",\n \"task\": OnyxCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,\n \"schedule\": timedelta(seconds=20),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n ]\n )\n\n# Add the Auto LLM update task if the config URL is set (has a default)\nif AUTO_LLM_CONFIG_URL:\n beat_task_templates.append(\n {\n \"name\": \"check-for-auto-llm-update\",\n \"task\": OnyxCeleryTask.CHECK_FOR_AUTO_LLM_UPDATE,\n \"schedule\": timedelta(seconds=AUTO_LLM_UPDATE_INTERVAL_SECONDS),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n }\n )\n\n# Add scheduled eval task if datasets are configured\nif SCHEDULED_EVAL_DATASET_NAMES:\n beat_task_templates.append(\n {\n \"name\": \"scheduled-eval-pipeline\",\n \"task\": OnyxCeleryTask.SCHEDULED_EVAL_TASK,\n # run every Sunday at midnight UTC\n \"schedule\": crontab(\n hour=0,\n minute=0,\n day_of_week=0,\n ),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n }\n )\n\n# Add OpenSearch migration task if enabled.\nif ENABLE_OPENSEARCH_INDEXING_FOR_ONYX:\n beat_task_templates.append(\n {\n \"name\": \"migrate-chunks-from-vespa-to-opensearch\",\n \"task\": OnyxCeleryTask.MIGRATE_CHUNKS_FROM_VESPA_TO_OPENSEARCH_TASK,\n # Try to enqueue an invocation of this task with this frequency.\n \"schedule\": timedelta(seconds=120), # 2 minutes\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n # If the task was not dequeued in this time, revoke it.\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.OPENSEARCH_MIGRATION,\n },\n }\n )\n\n\n# Beat task names that require a vector DB. Filtered out when DISABLE_VECTOR_DB.\n_VECTOR_DB_BEAT_TASK_NAMES: set[str] = {\n \"check-for-indexing\",\n \"check-for-connector-deletion\",\n \"check-for-vespa-sync\",\n \"check-for-pruning\",\n \"check-for-hierarchy-fetching\",\n \"check-for-checkpoint-cleanup\",\n \"check-for-index-attempt-cleanup\",\n \"check-for-doc-permissions-sync\",\n \"check-for-external-group-sync\",\n \"migrate-chunks-from-vespa-to-opensearch\",\n}\n\nif DISABLE_VECTOR_DB:\n beat_task_templates = [\n t for t in beat_task_templates if t[\"name\"] not in _VECTOR_DB_BEAT_TASK_NAMES\n ]\n\n\ndef make_cloud_generator_task(task: dict[str, Any]) -> dict[str, Any]:\n cloud_task: dict[str, Any] = {}\n\n # constant options for cloud beat task generators\n task_schedule: timedelta = task[\"schedule\"]\n cloud_task[\"schedule\"] = task_schedule\n cloud_task[\"options\"] = {}\n cloud_task[\"options\"][\"priority\"] = OnyxCeleryPriority.HIGHEST\n cloud_task[\"options\"][\"expires\"] = BEAT_EXPIRES_DEFAULT\n\n # settings dependent on the original task\n cloud_task[\"name\"] = f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_{task['name']}\"\n cloud_task[\"task\"] = OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR\n cloud_task[\"kwargs\"] = {}\n cloud_task[\"kwargs\"][\"task_name\"] = task[\"task\"]\n\n optional_fields = [\"queue\", \"priority\", \"expires\"]\n for field in optional_fields:\n if field in task[\"options\"]:\n cloud_task[\"kwargs\"][field] = task[\"options\"][field]\n\n return cloud_task\n\n\n# tasks that only run in the cloud and are system wide\n# the name attribute must start with ONYX_CLOUD_CELERY_TASK_PREFIX = \"cloud\" to be seen\n# by the DynamicTenantScheduler as system wide task and not a per tenant task\nbeat_cloud_tasks: list[dict] = [\n # cloud specific tasks\n {\n \"name\": f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-alembic\",\n \"task\": OnyxCeleryTask.CLOUD_MONITOR_ALEMBIC,\n \"schedule\": timedelta(hours=1),\n \"options\": {\n \"queue\": OnyxCeleryQueues.MONITORING,\n \"priority\": OnyxCeleryPriority.HIGH,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-celery-queues\",\n \"task\": OnyxCeleryTask.CLOUD_MONITOR_CELERY_QUEUES,\n \"schedule\": timedelta(seconds=30),\n \"options\": {\n \"queue\": OnyxCeleryQueues.MONITORING,\n \"priority\": OnyxCeleryPriority.HIGH,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-available-tenants\",\n \"task\": OnyxCeleryTask.CLOUD_CHECK_AVAILABLE_TENANTS,\n \"schedule\": timedelta(minutes=10),\n \"options\": {\n \"queue\": OnyxCeleryQueues.MONITORING,\n \"priority\": OnyxCeleryPriority.HIGH,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n {\n \"name\": f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-celery-pidbox\",\n \"task\": OnyxCeleryTask.CLOUD_MONITOR_CELERY_PIDBOX,\n \"schedule\": timedelta(hours=4),\n \"options\": {\n \"queue\": OnyxCeleryQueues.MONITORING,\n \"priority\": OnyxCeleryPriority.HIGH,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n },\n },\n]\n\n# tasks that only run self hosted\ntasks_to_schedule: list[dict] = []\nif not MULTI_TENANT:\n tasks_to_schedule.extend(\n [\n {\n \"name\": \"monitor-celery-queues\",\n \"task\": OnyxCeleryTask.MONITOR_CELERY_QUEUES,\n \"schedule\": timedelta(seconds=10),\n \"options\": {\n \"priority\": OnyxCeleryPriority.MEDIUM,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.MONITORING,\n },\n },\n {\n \"name\": \"monitor-process-memory\",\n \"task\": OnyxCeleryTask.MONITOR_PROCESS_MEMORY,\n \"schedule\": timedelta(minutes=5),\n \"options\": {\n \"priority\": OnyxCeleryPriority.LOW,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.MONITORING,\n },\n },\n {\n \"name\": \"celery-beat-heartbeat\",\n \"task\": OnyxCeleryTask.CELERY_BEAT_HEARTBEAT,\n \"schedule\": timedelta(minutes=1),\n \"options\": {\n \"priority\": OnyxCeleryPriority.HIGHEST,\n \"expires\": BEAT_EXPIRES_DEFAULT,\n \"queue\": OnyxCeleryQueues.PRIMARY,\n },\n },\n ]\n )\n\n tasks_to_schedule.extend(beat_task_templates)\n\n\ndef generate_cloud_tasks(\n beat_tasks: list[dict], beat_templates: list[dict], beat_multiplier: float\n) -> list[dict[str, Any]]:\n \"\"\"\n beat_tasks: system wide tasks that can be sent as is\n beat_templates: task templates that will be transformed into per tenant tasks via\n the cloud_beat_task_generator\n beat_multiplier: a multiplier that can be applied on top of the task schedule\n to speed up or slow down the task generation rate. useful in production.\n\n Returns a list of cloud tasks, which consists of incoming tasks + tasks generated\n from incoming templates.\n \"\"\"\n\n if beat_multiplier <= 0:\n raise ValueError(\"beat_multiplier must be positive!\")\n\n cloud_tasks: list[dict] = []\n\n # generate our tenant aware cloud tasks from the templates\n for beat_template in beat_templates:\n cloud_task = make_cloud_generator_task(beat_template)\n cloud_tasks.append(cloud_task)\n\n # factor in the cloud multiplier for the above\n for cloud_task in cloud_tasks:\n cloud_task[\"schedule\"] = cloud_task[\"schedule\"] * beat_multiplier\n\n # add the fixed cloud/system beat tasks. No multiplier for these.\n cloud_tasks.extend(copy.deepcopy(beat_tasks))\n return cloud_tasks\n\n\ndef get_cloud_tasks_to_schedule(beat_multiplier: float) -> list[dict[str, Any]]:\n return generate_cloud_tasks(beat_cloud_tasks, beat_task_templates, beat_multiplier)\n\n\ndef get_tasks_to_schedule() -> list[dict[str, Any]]:\n return tasks_to_schedule\n" }, { "path": "backend/onyx/background/celery/tasks/connector_deletion/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/connector_deletion/tasks.py", "content": "import traceback\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom pydantic import ValidationError\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_queue_length\nfrom onyx.background.celery.celery_redis import celery_get_queued_task_ids\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import OnyxRedisSignals\nfrom onyx.db.connector import fetch_connector_by_id\nfrom onyx.db.connector_credential_pair import add_deletion_failure_message\nfrom onyx.db.connector_credential_pair import (\n delete_connector_credential_pair__no_commit,\n)\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs\nfrom onyx.db.document import (\n delete_all_documents_by_connector_credential_pair__no_commit,\n)\nfrom onyx.db.document import get_document_ids_for_connector_credential_pair\nfrom onyx.db.document_set import delete_document_set_cc_pair_relationship__no_commit\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.index_attempt import delete_index_attempts\nfrom onyx.db.index_attempt import get_recent_attempts_for_cc_pair\nfrom onyx.db.permission_sync_attempt import (\n delete_doc_permission_sync_attempts__no_commit,\n)\nfrom onyx.db.permission_sync_attempt import (\n delete_external_group_permission_sync_attempts__no_commit,\n)\nfrom onyx.db.search_settings import get_all_search_settings\nfrom onyx.db.sync_record import cleanup_sync_records\nfrom onyx.db.sync_record import insert_sync_record\nfrom onyx.db.sync_record import update_sync_record_status\nfrom onyx.db.tag import delete_orphan_tags__no_commit\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_connector_delete import RedisConnectorDelete\nfrom onyx.redis.redis_connector_delete import RedisConnectorDeletePayload\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\nfrom onyx.utils.variable_functionality import noop_fallback\n\n\nclass TaskDependencyError(RuntimeError):\n \"\"\"Raised to the caller to indicate dependent tasks are running that would interfere\n with connector deletion.\"\"\"\n\n\ndef revoke_tasks_blocking_deletion(\n redis_connector: RedisConnector, db_session: Session, app: Celery\n) -> None:\n search_settings_list = get_all_search_settings(db_session)\n for search_settings in search_settings_list:\n try:\n recent_index_attempts = get_recent_attempts_for_cc_pair(\n cc_pair_id=redis_connector.cc_pair_id,\n search_settings_id=search_settings.id,\n limit=1,\n db_session=db_session,\n )\n if (\n recent_index_attempts\n and recent_index_attempts[0].status == IndexingStatus.IN_PROGRESS\n and recent_index_attempts[0].celery_task_id\n ):\n app.control.revoke(recent_index_attempts[0].celery_task_id)\n task_logger.info(\n f\"Revoked indexing task {recent_index_attempts[0].celery_task_id}.\"\n )\n except Exception:\n task_logger.exception(\"Exception while revoking indexing task\")\n\n try:\n permissions_sync_payload = redis_connector.permissions.payload\n if permissions_sync_payload and permissions_sync_payload.celery_task_id:\n app.control.revoke(permissions_sync_payload.celery_task_id)\n task_logger.info(\n f\"Revoked permissions sync task {permissions_sync_payload.celery_task_id}.\"\n )\n except Exception:\n task_logger.exception(\"Exception while revoking pruning task\")\n\n try:\n prune_payload = redis_connector.prune.payload\n if prune_payload and prune_payload.celery_task_id:\n app.control.revoke(prune_payload.celery_task_id)\n task_logger.info(f\"Revoked pruning task {prune_payload.celery_task_id}.\")\n except Exception:\n task_logger.exception(\"Exception while revoking permissions sync task\")\n\n try:\n external_group_sync_payload = redis_connector.external_group_sync.payload\n if external_group_sync_payload and external_group_sync_payload.celery_task_id:\n app.control.revoke(external_group_sync_payload.celery_task_id)\n task_logger.info(\n f\"Revoked external group sync task {external_group_sync_payload.celery_task_id}.\"\n )\n except Exception:\n task_logger.exception(\"Exception while revoking external group sync task\")\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n trail=False,\n bind=True,\n)\ndef check_for_connector_deletion_task(self: Task, *, tenant_id: str) -> bool | None:\n r = get_redis_client()\n r_replica = get_redis_replica_client()\n\n lock_beat: RedisLock = r.lock(\n OnyxRedisLocks.CHECK_CONNECTOR_DELETION_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # Prevent this task from overlapping with itself\n if not lock_beat.acquire(blocking=False):\n return None\n\n try:\n # we want to run this less frequently than the overall task\n lock_beat.reacquire()\n if not r.exists(OnyxRedisSignals.BLOCK_VALIDATE_CONNECTOR_DELETION_FENCES):\n # clear fences that don't have associated celery tasks in progress\n try:\n r_celery = celery_get_broker_client(self.app)\n validate_connector_deletion_fences(\n tenant_id, r, r_replica, r_celery, lock_beat\n )\n except Exception:\n task_logger.exception(\n \"Exception while validating connector deletion fences\"\n )\n\n r.set(OnyxRedisSignals.BLOCK_VALIDATE_CONNECTOR_DELETION_FENCES, 1, ex=300)\n\n # collect cc_pair_ids\n cc_pair_ids: list[int] = []\n with get_session_with_current_tenant() as db_session:\n cc_pairs = get_connector_credential_pairs(db_session)\n for cc_pair in cc_pairs:\n cc_pair_ids.append(cc_pair.id)\n\n # try running cleanup on the cc_pair_ids\n for cc_pair_id in cc_pair_ids:\n with get_session_with_current_tenant() as db_session:\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n try:\n try_generate_document_cc_pair_cleanup_tasks(\n self.app, cc_pair_id, db_session, lock_beat, tenant_id\n )\n except TaskDependencyError as e:\n # this means we wanted to start deleting but dependent tasks were running\n # on the first error, we set a stop signal and revoke the dependent tasks\n # on subsequent errors, we hard reset blocking fences after our specified timeout\n # is exceeded\n task_logger.info(str(e))\n\n if not redis_connector.stop.fenced:\n # one time revoke of celery tasks\n task_logger.info(\"Revoking any tasks blocking deletion.\")\n revoke_tasks_blocking_deletion(\n redis_connector, db_session, self.app\n )\n redis_connector.stop.set_fence(True)\n redis_connector.stop.set_timeout()\n else:\n # stop signal already set\n if redis_connector.stop.timed_out:\n # waiting too long, just reset blocking fences\n task_logger.info(\n \"Timed out waiting for tasks blocking deletion. Resetting blocking fences.\"\n )\n\n redis_connector.prune.reset()\n redis_connector.permissions.reset()\n redis_connector.external_group_sync.reset()\n else:\n # just wait\n pass\n else:\n # clear the stop signal if it exists ... no longer needed\n redis_connector.stop.set_fence(False)\n\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n\n if not r.exists(key_bytes):\n r.srem(OnyxRedisConstants.ACTIVE_FENCES, key_bytes)\n continue\n\n key_str = key_bytes.decode(\"utf-8\")\n if key_str.startswith(RedisConnectorDelete.FENCE_PREFIX):\n monitor_connector_deletion_taskset(tenant_id, key_bytes, r)\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception:\n task_logger.exception(\"Unexpected exception during connector deletion check\")\n finally:\n if lock_beat.owned():\n lock_beat.release()\n\n return True\n\n\ndef try_generate_document_cc_pair_cleanup_tasks(\n app: Celery,\n cc_pair_id: int,\n db_session: Session,\n lock_beat: RedisLock,\n tenant_id: str,\n) -> int | None:\n \"\"\"Returns an int if syncing is needed. The int represents the number of sync tasks generated.\n Note that syncing can still be required even if the number of sync tasks generated is zero.\n Returns None if no syncing is required.\n\n Will raise TaskDependencyError if dependent tasks such as indexing and pruning are\n still running. In our case, the caller reacts by setting a stop signal in Redis to\n exit those tasks as quickly as possible.\n \"\"\"\n\n lock_beat.reacquire()\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n # don't generate sync tasks if tasks are still pending\n if redis_connector.delete.fenced:\n return None\n\n # we need to load the state of the object inside the fence\n # to avoid a race condition with db.commit/fence deletion\n # at the end of this taskset\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n return None\n\n if cc_pair.status != ConnectorCredentialPairStatus.DELETING:\n # there should be no in-progress sync records if this is up to date\n # clean it up just in case things got into a bad state\n cleanup_sync_records(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.CONNECTOR_DELETION,\n )\n return None\n\n # set a basic fence to start\n redis_connector.delete.set_active()\n fence_payload = RedisConnectorDeletePayload(\n num_tasks=None,\n submitted=datetime.now(timezone.utc),\n )\n\n redis_connector.delete.set_fence(fence_payload)\n\n try:\n # do not proceed if connector indexing or connector pruning are running\n search_settings_list = get_all_search_settings(db_session)\n for search_settings in search_settings_list:\n recent_index_attempts = get_recent_attempts_for_cc_pair(\n cc_pair_id=cc_pair_id,\n search_settings_id=search_settings.id,\n limit=1,\n db_session=db_session,\n )\n if (\n recent_index_attempts\n and recent_index_attempts[0].status == IndexingStatus.IN_PROGRESS\n ):\n raise TaskDependencyError(\n \"Connector deletion - Delayed (indexing in progress): \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings.id}\"\n )\n\n if redis_connector.prune.fenced:\n raise TaskDependencyError(\n f\"Connector deletion - Delayed (pruning in progress): cc_pair={cc_pair_id}\"\n )\n\n if redis_connector.permissions.fenced:\n raise TaskDependencyError(\n f\"Connector deletion - Delayed (permissions in progress): cc_pair={cc_pair_id}\"\n )\n\n # add tasks to celery and build up the task set to monitor in redis\n redis_connector.delete.taskset_clear()\n\n # Add all documents that need to be updated into the queue\n task_logger.info(\n f\"RedisConnectorDeletion.generate_tasks starting. cc_pair={cc_pair_id}\"\n )\n tasks_generated = redis_connector.delete.generate_tasks(\n app, db_session, lock_beat\n )\n if tasks_generated is None:\n raise ValueError(\"RedisConnectorDeletion.generate_tasks returned None\")\n\n try:\n insert_sync_record(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.CONNECTOR_DELETION,\n )\n except Exception:\n task_logger.exception(\"insert_sync_record exceptioned.\")\n\n except TaskDependencyError:\n redis_connector.delete.set_fence(None)\n raise\n except Exception:\n task_logger.exception(\"Unexpected exception\")\n redis_connector.delete.set_fence(None)\n return None\n else:\n # Currently we are allowing the sync to proceed with 0 tasks.\n # It's possible for sets/groups to be generated initially with no entries\n # and they still need to be marked as up to date.\n # if tasks_generated == 0:\n # return 0\n\n task_logger.info(\n f\"RedisConnectorDeletion.generate_tasks finished. cc_pair={cc_pair_id} tasks_generated={tasks_generated}\"\n )\n\n # set this only after all tasks have been added\n fence_payload.num_tasks = tasks_generated\n redis_connector.delete.set_fence(fence_payload)\n\n return tasks_generated\n\n\ndef monitor_connector_deletion_taskset(\n tenant_id: str,\n key_bytes: bytes,\n r: Redis, # noqa: ARG001\n) -> None:\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n task_logger.warning(f\"could not parse cc_pair_id from {fence_key}\")\n return\n\n cc_pair_id = int(cc_pair_id_str)\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n fence_data = redis_connector.delete.payload\n if not fence_data:\n task_logger.warning(\n f\"Connector deletion - fence payload invalid: cc_pair={cc_pair_id}\"\n )\n return\n\n if fence_data.num_tasks is None:\n # the fence is setting up but isn't ready yet\n return\n\n remaining = redis_connector.delete.get_remaining()\n task_logger.info(\n f\"Connector deletion progress: cc_pair={cc_pair_id} remaining={remaining} initial={fence_data.num_tasks}\"\n )\n if remaining > 0:\n with get_session_with_current_tenant() as db_session:\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.CONNECTOR_DELETION,\n sync_status=SyncStatus.IN_PROGRESS,\n num_docs_synced=remaining,\n )\n return\n\n with get_session_with_current_tenant() as db_session:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n credential_id_to_delete: int | None = None\n connector_id_to_delete: int | None = None\n if not cc_pair:\n task_logger.warning(\n f\"Connector deletion - cc_pair not found: cc_pair={cc_pair_id}\"\n )\n return\n\n try:\n doc_ids = get_document_ids_for_connector_credential_pair(\n db_session, cc_pair.connector_id, cc_pair.credential_id\n )\n if len(doc_ids) > 0:\n # NOTE(rkuo): if this happens, documents somehow got added while\n # deletion was in progress. Likely a bug gating off pruning and indexing\n # work before deletion starts.\n task_logger.warning(\n \"Connector deletion - documents still found after taskset completion. \"\n \"Clearing the current deletion attempt and allowing deletion to restart: \"\n f\"cc_pair={cc_pair_id} \"\n f\"docs_deleted={fence_data.num_tasks} \"\n f\"docs_remaining={len(doc_ids)}\"\n )\n\n # We don't want to waive off why we get into this state, but resetting\n # our attempt and letting the deletion restart is a good way to recover\n redis_connector.delete.reset()\n raise RuntimeError(\n \"Connector deletion - documents still found after taskset completion\"\n )\n\n # clean up the rest of the related Postgres entities\n # index attempts\n delete_index_attempts(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n\n # permission sync attempts\n delete_doc_permission_sync_attempts__no_commit(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n delete_external_group_permission_sync_attempts__no_commit(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n\n # document sets\n delete_document_set_cc_pair_relationship__no_commit(\n db_session=db_session,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n )\n\n # user groups\n cleanup_user_groups = fetch_versioned_implementation_with_fallback(\n \"onyx.db.user_group\",\n \"delete_user_group_cc_pair_relationship__no_commit\",\n noop_fallback,\n )\n cleanup_user_groups(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n )\n\n # delete orphan tags\n delete_orphan_tags__no_commit(db_session)\n\n # Store IDs before potentially expiring cc_pair\n connector_id_to_delete = cc_pair.connector_id\n credential_id_to_delete = cc_pair.credential_id\n\n # Explicitly delete document by connector credential pair records before deleting the connector\n # This is needed because connector_id is a primary key in that table and cascading deletes won't work\n delete_all_documents_by_connector_credential_pair__no_commit(\n db_session=db_session,\n connector_id=connector_id_to_delete,\n credential_id=credential_id_to_delete,\n )\n\n # Flush to ensure document deletion happens before connector deletion\n db_session.flush()\n\n # Expire the cc_pair to ensure SQLAlchemy doesn't try to manage its state\n # related to the deleted DocumentByConnectorCredentialPair during commit\n db_session.expire(cc_pair)\n\n # finally, delete the cc-pair\n delete_connector_credential_pair__no_commit(\n db_session=db_session,\n connector_id=connector_id_to_delete,\n credential_id=credential_id_to_delete,\n )\n # if there are no credentials left, delete the connector\n connector = fetch_connector_by_id(\n db_session=db_session,\n connector_id=connector_id_to_delete,\n )\n if not connector or not len(connector.credentials):\n task_logger.info(\n \"Connector deletion - Found no credentials left for connector, deleting connector\"\n )\n db_session.delete(connector)\n db_session.commit()\n\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.CONNECTOR_DELETION,\n sync_status=SyncStatus.SUCCESS,\n num_docs_synced=fence_data.num_tasks,\n )\n\n except Exception as e:\n db_session.rollback()\n stack_trace = traceback.format_exc()\n error_message = f\"Error: {str(e)}\\n\\nStack Trace:\\n{stack_trace}\"\n add_deletion_failure_message(db_session, cc_pair_id, error_message)\n\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.CONNECTOR_DELETION,\n sync_status=SyncStatus.FAILED,\n num_docs_synced=fence_data.num_tasks,\n )\n\n task_logger.exception(\n f\"Connector deletion exceptioned: \"\n f\"cc_pair={cc_pair_id} connector={connector_id_to_delete} credential={credential_id_to_delete}\"\n )\n raise e\n\n task_logger.info(\n f\"Connector deletion succeeded: \"\n f\"cc_pair={cc_pair_id} \"\n f\"connector={connector_id_to_delete} \"\n f\"credential={credential_id_to_delete} \"\n f\"docs_deleted={fence_data.num_tasks}\"\n )\n\n redis_connector.delete.reset()\n\n\ndef validate_connector_deletion_fences(\n tenant_id: str,\n r: Redis,\n r_replica: Redis,\n r_celery: Redis,\n lock_beat: RedisLock,\n) -> None:\n # building lookup table can be expensive, so we won't bother\n # validating until the queue is small\n CONNECTION_DELETION_VALIDATION_MAX_QUEUE_LEN = 1024\n\n queue_len = celery_get_queue_length(OnyxCeleryQueues.CONNECTOR_DELETION, r_celery)\n if queue_len > CONNECTION_DELETION_VALIDATION_MAX_QUEUE_LEN:\n return\n\n queued_upsert_tasks = celery_get_queued_task_ids(\n OnyxCeleryQueues.CONNECTOR_DELETION, r_celery\n )\n\n # validate all existing connector deletion jobs\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n key_str = key_bytes.decode(\"utf-8\")\n if not key_str.startswith(RedisConnectorDelete.FENCE_PREFIX):\n continue\n\n validate_connector_deletion_fence(\n tenant_id,\n key_bytes,\n queued_upsert_tasks,\n r,\n )\n\n lock_beat.reacquire()\n\n return\n\n\ndef validate_connector_deletion_fence(\n tenant_id: str,\n key_bytes: bytes,\n queued_upsert_tasks: set[str],\n r: Redis,\n) -> None:\n \"\"\"Checks for the error condition where an indexing fence is set but the associated celery tasks don't exist.\n This can happen if the indexing worker hard crashes or is terminated.\n Being in this bad state means the fence will never clear without help, so this function\n gives the help.\n\n How this works:\n 1. This function renews the active signal with a 5 minute TTL under the following conditions\n 1.2. When the task is seen in the redis queue\n 1.3. When the task is seen in the reserved / prefetched list\n\n 2. Externally, the active signal is renewed when:\n 2.1. The fence is created\n 2.2. The indexing watchdog checks the spawned task.\n\n 3. The TTL allows us to get through the transitions on fence startup\n and when the task starts executing.\n\n More TTL clarification: it is seemingly impossible to exactly query Celery for\n whether a task is in the queue or currently executing.\n 1. An unknown task id is always returned as state PENDING.\n 2. Redis can be inspected for the task id, but the task id is gone between the time a worker receives the task\n and the time it actually starts on the worker.\n\n queued_tasks: the celery queue of lightweight permission sync tasks\n reserved_tasks: prefetched tasks for sync task generator\n \"\"\"\n # if the fence doesn't exist, there's nothing to do\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n task_logger.warning(\n f\"validate_connector_deletion_fence - could not parse id from {fence_key}\"\n )\n return\n\n cc_pair_id = int(cc_pair_id_str)\n # parse out metadata and initialize the helper class with it\n redis_connector = RedisConnector(tenant_id, int(cc_pair_id))\n\n # check to see if the fence/payload exists\n if not redis_connector.delete.fenced:\n return\n\n # in the cloud, the payload format may have changed ...\n # it's a little sloppy, but just reset the fence for now if that happens\n # TODO: add intentional cleanup/abort logic\n try:\n payload = redis_connector.delete.payload\n except ValidationError:\n task_logger.exception(\n \"validate_connector_deletion_fence - \"\n \"Resetting fence because fence schema is out of date: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key}\"\n )\n\n redis_connector.delete.reset()\n return\n\n if not payload:\n return\n\n # OK, there's actually something for us to validate\n\n # look up every task in the current taskset in the celery queue\n # every entry in the taskset should have an associated entry in the celery task queue\n # because we get the celery tasks first, the entries in our own permissions taskset\n # should be roughly a subset of the tasks in celery\n\n # this check isn't very exact, but should be sufficient over a period of time\n # A single successful check over some number of attempts is sufficient.\n\n # TODO: if the number of tasks in celery is much lower than than the taskset length\n # we might be able to shortcut the lookup since by definition some of the tasks\n # must not exist in celery.\n\n tasks_scanned = 0\n tasks_not_in_celery = 0 # a non-zero number after completing our check is bad\n\n for member in r.sscan_iter(redis_connector.delete.taskset_key):\n tasks_scanned += 1\n\n member_bytes = cast(bytes, member)\n member_str = member_bytes.decode(\"utf-8\")\n if member_str in queued_upsert_tasks:\n continue\n\n tasks_not_in_celery += 1\n\n task_logger.info(\n f\"validate_connector_deletion_fence task check: tasks_scanned={tasks_scanned} tasks_not_in_celery={tasks_not_in_celery}\"\n )\n\n # we're active if there are still tasks to run and those tasks all exist in celery\n if tasks_scanned > 0 and tasks_not_in_celery == 0:\n redis_connector.delete.set_active()\n return\n\n # we may want to enable this check if using the active task list somehow isn't good enough\n # if redis_connector_index.generator_locked():\n # logger.info(f\"{payload.celery_task_id} is currently executing.\")\n\n # if we get here, we didn't find any direct indication that the associated celery tasks exist,\n # but they still might be there due to gaps in our ability to check states during transitions\n # Checking the active signal safeguards us against these transition periods\n # (which has a duration that allows us to bridge those gaps)\n if redis_connector.delete.active():\n return\n\n # celery tasks don't exist and the active signal has expired, possibly due to a crash. Clean it up.\n task_logger.warning(\n \"validate_connector_deletion_fence - \"\n \"Resetting fence because no associated celery tasks were found: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key}\"\n )\n\n redis_connector.delete.reset()\n return\n" }, { "path": "backend/onyx/background/celery/tasks/docfetching/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/docfetching/task_creation_utils.py", "content": "from uuid import uuid4\n\nfrom celery import Celery\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.index_attempt import mark_attempt_failed\nfrom onyx.db.indexing_coordination import IndexingCoordination\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import SearchSettings\n\n\ndef try_creating_docfetching_task(\n celery_app: Celery,\n cc_pair: ConnectorCredentialPair,\n search_settings: SearchSettings,\n reindex: bool,\n db_session: Session,\n r: Redis,\n tenant_id: str,\n) -> int | None:\n \"\"\"Checks for any conditions that should block the indexing task from being\n created, then creates the task.\n\n Does not check for scheduling related conditions as this function\n is used to trigger indexing immediately.\n\n Now uses database-based coordination instead of Redis fencing.\n \"\"\"\n\n LOCK_TIMEOUT = 30\n\n # we need to serialize any attempt to trigger indexing since it can be triggered\n # either via celery beat or manually (API call)\n lock: RedisLock = r.lock(\n DANSWER_REDIS_FUNCTION_LOCK_PREFIX + \"try_creating_indexing_task\",\n timeout=LOCK_TIMEOUT,\n )\n\n acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)\n if not acquired:\n return None\n\n index_attempt_id = None\n try:\n # Basic status checks\n db_session.refresh(cc_pair)\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n return None\n\n # Generate custom task ID for tracking\n custom_task_id = f\"docfetching_{cc_pair.id}_{search_settings.id}_{uuid4()}\"\n\n # Try to create a new index attempt using database coordination\n # This replaces the Redis fencing mechanism\n index_attempt_id = IndexingCoordination.try_create_index_attempt(\n db_session=db_session,\n cc_pair_id=cc_pair.id,\n search_settings_id=search_settings.id,\n celery_task_id=custom_task_id,\n from_beginning=reindex,\n )\n\n if index_attempt_id is None:\n # Another indexing attempt is already running\n return None\n\n # Use higher priority for first-time indexing to ensure new connectors\n # get processed before re-indexing of existing connectors\n has_successful_attempt = cc_pair.last_successful_index_time is not None\n priority = (\n OnyxCeleryPriority.MEDIUM\n if has_successful_attempt\n else OnyxCeleryPriority.HIGH\n )\n\n # Send the task to Celery\n result = celery_app.send_task(\n OnyxCeleryTask.CONNECTOR_DOC_FETCHING_TASK,\n kwargs=dict(\n index_attempt_id=index_attempt_id,\n cc_pair_id=cc_pair.id,\n search_settings_id=search_settings.id,\n tenant_id=tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_DOC_FETCHING,\n task_id=custom_task_id,\n priority=priority,\n )\n if not result:\n raise RuntimeError(\"send_task for connector_doc_fetching_task failed.\")\n\n task_logger.info(\n f\"Created docfetching task: \"\n f\"cc_pair={cc_pair.id} \"\n f\"search_settings={search_settings.id} \"\n f\"attempt_id={index_attempt_id} \"\n f\"celery_task_id={custom_task_id}\"\n )\n\n return index_attempt_id\n\n except Exception:\n task_logger.exception(\n f\"try_creating_indexing_task - Unexpected exception: cc_pair={cc_pair.id} search_settings={search_settings.id}\"\n )\n\n # Clean up on failure\n if index_attempt_id is not None:\n mark_attempt_failed(index_attempt_id, db_session)\n\n return None\n finally:\n if lock.owned():\n lock.release()\n\n return index_attempt_id\n" }, { "path": "backend/onyx/background/celery/tasks/docfetching/tasks.py", "content": "import multiprocessing\nimport os\nimport time\nimport traceback\nfrom time import sleep\n\nimport sentry_sdk\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\n\nfrom onyx import __version__\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.memory_monitoring import emit_process_memory\nfrom onyx.background.celery.tasks.docprocessing.heartbeat import start_heartbeat\nfrom onyx.background.celery.tasks.docprocessing.heartbeat import stop_heartbeat\nfrom onyx.background.celery.tasks.docprocessing.tasks import ConnectorIndexingLogBuilder\nfrom onyx.background.celery.tasks.docprocessing.utils import IndexingCallback\nfrom onyx.background.celery.tasks.models import DocProcessingContext\nfrom onyx.background.celery.tasks.models import IndexingWatchdogTerminalStatus\nfrom onyx.background.celery.tasks.models import SimpleJobResult\nfrom onyx.background.indexing.job_client import SimpleJob\nfrom onyx.background.indexing.job_client import SimpleJobClient\nfrom onyx.background.indexing.job_client import SimpleJobException\nfrom onyx.background.indexing.run_docfetching import run_docfetching_entrypoint\nfrom onyx.configs.constants import CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.index_attempt import get_index_attempt\nfrom onyx.db.index_attempt import mark_attempt_canceled\nfrom onyx.db.index_attempt import mark_attempt_failed\nfrom onyx.db.indexing_coordination import IndexingCoordination\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import global_version\nfrom shared_configs.configs import SENTRY_DSN\n\nlogger = setup_logger()\n\n\ndef _verify_indexing_attempt(\n index_attempt_id: int,\n cc_pair_id: int,\n search_settings_id: int,\n) -> None:\n \"\"\"\n Verify that the indexing attempt exists and is in the correct state.\n \"\"\"\n\n with get_session_with_current_tenant() as db_session:\n attempt = get_index_attempt(db_session, index_attempt_id)\n\n if not attempt:\n raise SimpleJobException(\n f\"docfetching_task - IndexAttempt not found: attempt_id={index_attempt_id}\",\n code=IndexingWatchdogTerminalStatus.FENCE_NOT_FOUND.code,\n )\n\n if attempt.connector_credential_pair_id != cc_pair_id:\n raise SimpleJobException(\n f\"docfetching_task - CC pair mismatch: expected={cc_pair_id} actual={attempt.connector_credential_pair_id}\",\n code=IndexingWatchdogTerminalStatus.FENCE_MISMATCH.code,\n )\n\n if attempt.search_settings_id != search_settings_id:\n raise SimpleJobException(\n f\"docfetching_task - Search settings mismatch: expected={search_settings_id} actual={attempt.search_settings_id}\",\n code=IndexingWatchdogTerminalStatus.FENCE_MISMATCH.code,\n )\n\n if attempt.status not in [\n IndexingStatus.NOT_STARTED,\n IndexingStatus.IN_PROGRESS,\n ]:\n raise SimpleJobException(\n f\"docfetching_task - Invalid attempt status: attempt_id={index_attempt_id} status={attempt.status}\",\n code=IndexingWatchdogTerminalStatus.FENCE_MISMATCH.code,\n )\n\n # Check for cancellation\n if IndexingCoordination.check_cancellation_requested(\n db_session, index_attempt_id\n ):\n raise SimpleJobException(\n f\"docfetching_task - Cancellation requested: attempt_id={index_attempt_id}\",\n code=IndexingWatchdogTerminalStatus.BLOCKED_BY_STOP_SIGNAL.code,\n )\n\n logger.info(\n f\"docfetching_task - IndexAttempt verified: \"\n f\"attempt_id={index_attempt_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id}\"\n )\n\n\ndef docfetching_task(\n app: Celery,\n index_attempt_id: int,\n cc_pair_id: int,\n search_settings_id: int,\n is_ee: bool,\n tenant_id: str,\n) -> None:\n \"\"\"\n This function is run in a SimpleJob as a new process. It is responsible for validating\n some stuff, but basically it just calls run_indexing_entrypoint.\n\n NOTE: if an exception is raised out of this task, the primary worker will detect\n that the task transitioned to a \"READY\" state but the generator_complete_key doesn't exist.\n This will cause the primary worker to abort the indexing attempt and clean up.\n \"\"\"\n\n # Start heartbeat for this indexing attempt\n heartbeat_thread, stop_event = start_heartbeat(index_attempt_id)\n try:\n _docfetching_task(\n app, index_attempt_id, cc_pair_id, search_settings_id, is_ee, tenant_id\n )\n finally:\n stop_heartbeat(heartbeat_thread, stop_event) # Stop heartbeat before exiting\n\n\ndef _docfetching_task(\n app: Celery,\n index_attempt_id: int,\n cc_pair_id: int,\n search_settings_id: int,\n is_ee: bool,\n tenant_id: str,\n) -> None:\n # Since connector_indexing_proxy_task spawns a new process using this function as\n # the entrypoint, we init Sentry here.\n if SENTRY_DSN:\n sentry_sdk.init(\n dsn=SENTRY_DSN,\n traces_sample_rate=0.1,\n release=__version__,\n )\n logger.info(\"Sentry initialized\")\n else:\n logger.debug(\"Sentry DSN not provided, skipping Sentry initialization\")\n\n logger.info(\n f\"Indexing spawned task starting: \"\n f\"attempt={index_attempt_id} \"\n f\"tenant={tenant_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id}\"\n )\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n # TODO: remove all fences, cause all signals to be set in postgres\n if redis_connector.delete.fenced:\n raise SimpleJobException(\n f\"Indexing will not start because connector deletion is in progress: \"\n f\"attempt={index_attempt_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={redis_connector.delete.fence_key}\",\n code=IndexingWatchdogTerminalStatus.BLOCKED_BY_DELETION.code,\n )\n\n if redis_connector.stop.fenced:\n raise SimpleJobException(\n f\"Indexing will not start because a connector stop signal was detected: \"\n f\"attempt={index_attempt_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={redis_connector.stop.fence_key}\",\n code=IndexingWatchdogTerminalStatus.BLOCKED_BY_STOP_SIGNAL.code,\n )\n\n # Verify the indexing attempt exists and is valid\n # This replaces the Redis fence payload waiting\n _verify_indexing_attempt(index_attempt_id, cc_pair_id, search_settings_id)\n\n try:\n with get_session_with_current_tenant() as db_session:\n attempt = get_index_attempt(db_session, index_attempt_id)\n if not attempt:\n raise SimpleJobException(\n f\"Index attempt not found: index_attempt={index_attempt_id}\",\n code=IndexingWatchdogTerminalStatus.INDEX_ATTEMPT_MISMATCH.code,\n )\n\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n\n if not cc_pair:\n raise SimpleJobException(\n f\"cc_pair not found: cc_pair={cc_pair_id}\",\n code=IndexingWatchdogTerminalStatus.INDEX_ATTEMPT_MISMATCH.code,\n )\n\n # define a callback class\n callback = IndexingCallback(\n redis_connector,\n )\n\n logger.info(\n f\"Indexing spawned task running entrypoint: attempt={index_attempt_id} \"\n f\"tenant={tenant_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id}\"\n )\n\n # This is where the heavy/real work happens\n run_docfetching_entrypoint(\n app,\n index_attempt_id,\n tenant_id,\n cc_pair_id,\n is_ee,\n callback=callback,\n )\n\n except ConnectorValidationError:\n raise SimpleJobException(\n f\"Indexing task failed: attempt={index_attempt_id} \"\n f\"tenant={tenant_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id}\",\n code=IndexingWatchdogTerminalStatus.CONNECTOR_VALIDATION_ERROR.code,\n )\n\n except Exception as e:\n logger.exception(\n f\"Indexing spawned task failed: attempt={index_attempt_id} \"\n f\"tenant={tenant_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id}\"\n )\n\n # special bulletproofing ... truncate long exception messages\n # for exception types that require more args, this will fail\n # thus the try/except\n try:\n sanitized_e = type(e)(str(e)[:1024])\n sanitized_e.__traceback__ = e.__traceback__\n raise sanitized_e\n except Exception:\n raise e\n\n logger.info(\n f\"Indexing spawned task finished: attempt={index_attempt_id} cc_pair={cc_pair_id} search_settings={search_settings_id}\"\n )\n os._exit(0) # ensure process exits cleanly\n\n\ndef process_job_result(\n job: SimpleJob,\n connector_source: str | None,\n index_attempt_id: int,\n log_builder: ConnectorIndexingLogBuilder,\n) -> SimpleJobResult:\n result = SimpleJobResult()\n result.connector_source = connector_source\n\n if job.process:\n result.exit_code = job.process.exitcode\n\n if job.status != \"error\":\n result.status = IndexingWatchdogTerminalStatus.SUCCEEDED\n return result\n\n ignore_exitcode = False\n\n # In EKS, there is an edge case where successful tasks return exit\n # code 1 in the cloud due to the set_spawn_method not sticking.\n # Workaround: check that the total number of batches is set, since this only\n # happens when docfetching completed successfully\n with get_session_with_current_tenant() as db_session:\n index_attempt = get_index_attempt(db_session, index_attempt_id)\n if index_attempt and index_attempt.total_batches is not None:\n ignore_exitcode = True\n\n if ignore_exitcode:\n result.status = IndexingWatchdogTerminalStatus.SUCCEEDED\n task_logger.warning(\n log_builder.build(\n \"Indexing watchdog - spawned task has non-zero exit code but completion signal is OK. Continuing...\",\n exit_code=str(result.exit_code),\n )\n )\n else:\n if result.exit_code is not None:\n result.status = IndexingWatchdogTerminalStatus.from_code(result.exit_code)\n\n job_level_exception = job.exception()\n result.exception_str = f\"Docfetching returned exit code {result.exit_code} with exception: {job_level_exception}\"\n\n return result\n\n\n@shared_task(\n name=OnyxCeleryTask.CONNECTOR_DOC_FETCHING_TASK,\n bind=True,\n acks_late=False,\n track_started=True,\n)\ndef docfetching_proxy_task(\n self: Task,\n index_attempt_id: int,\n cc_pair_id: int,\n search_settings_id: int,\n tenant_id: str,\n) -> None:\n \"\"\"\n This task is the entrypoint for the full indexing pipeline, which is composed of two tasks:\n docfetching and docprocessing.\n This task is spawned by \"try_creating_indexing_task\" which is called in the \"check_for_indexing\" task.\n\n This task spawns a new process for a new scheduled index attempt. That\n new process (which runs the docfetching_task function) does the following:\n\n 1) determines parameters of the indexing attempt (which connector indexing function to run,\n start and end time, from prev checkpoint or not), then run that connector. Specifically,\n connectors are responsible for reading data from an outside source and converting it to Onyx documents.\n At the moment these two steps (reading external data and converting to an Onyx document)\n are not parallelized in most connectors; that's a subject for future work.\n\n Each document batch produced by step 1 is stored in the file store, and a docprocessing task is spawned\n to process it. docprocessing involves the steps listed below.\n\n 2) upserts documents to postgres (index_doc_batch_prepare)\n 3) chunks each document (optionally adds context for contextual rag)\n 4) embeds chunks (embed_chunks_with_failure_handling) via a call to the model server\n 5) write chunks to vespa (write_chunks_to_vector_db_with_backoff)\n 6) update document and indexing metadata in postgres\n 7) pulls all document IDs from the source and compares those IDs to locally stored documents and deletes\n all locally stored IDs missing from the most recently pulled document ID list\n\n Some important notes:\n Invariants:\n - docfetching proxy tasks are spawned by check_for_indexing. The proxy then runs the docfetching_task wrapped in a watchdog.\n The watchdog is responsible for monitoring the docfetching_task and marking the index attempt as failed\n if it is not making progress.\n - All docprocessing tasks are spawned by a docfetching task.\n - all docfetching tasks, docprocessing tasks, and document batches in the file store are\n associated with a specific index attempt.\n - the index attempt status is the source of truth for what is currently happening with the index attempt.\n It is coupled with the creation/running of docfetching and docprocessing tasks as much as possible.\n\n How we deal with failures/ partial indexing:\n - non-checkpointed connectors/ new runs in general => delete the old document batches from the file store and do the new run\n - checkpointed connectors + resuming from checkpoint => reissue the old document batches and do a new run\n\n Misc:\n - most inter-process communication is handled in postgres, some is still in redis and we're trying to remove it\n - Heartbeat spawned in docfetching and docprocessing is how check_for_indexing monitors liveliness\n - progress based liveliness check: if nothing is done in 3-6 hours, mark the attempt as failed\n - TODO: task level timeouts (i.e. a connector stuck in an infinite loop)\n\n\n Comments below are from the old version and some may no longer be valid.\n TODO(rkuo): refactor this so that there is a single return path where we canonically\n log the result of running this function.\n\n Some more Richard notes:\n celery out of process task execution strategy is pool=prefork, but it uses fork,\n and forking is inherently unstable.\n\n To work around this, we use pool=threads and proxy our work to a spawned task.\n\n acks_late must be set to False. Otherwise, celery's visibility timeout will\n cause any task that runs longer than the timeout to be redispatched by the broker.\n There appears to be no good workaround for this, so we need to handle redispatching\n manually.\n NOTE: we try/except all db access in this function because as a watchdog, this function\n needs to be extremely stable.\n \"\"\"\n # TODO: remove dependence on Redis\n start = time.monotonic()\n\n result = SimpleJobResult()\n\n ctx = DocProcessingContext(\n tenant_id=tenant_id,\n cc_pair_id=cc_pair_id,\n search_settings_id=search_settings_id,\n index_attempt_id=index_attempt_id,\n )\n\n log_builder = ConnectorIndexingLogBuilder(ctx)\n\n task_logger.info(\n log_builder.build(\n \"Indexing watchdog - starting\",\n mp_start_method=str(multiprocessing.get_start_method()),\n )\n )\n\n if not self.request.id:\n task_logger.error(\"self.request.id is None!\")\n\n client = SimpleJobClient()\n task_logger.info(f\"submitting docfetching_task with tenant_id={tenant_id}\")\n\n job = client.submit(\n docfetching_task,\n self.app,\n index_attempt_id,\n cc_pair_id,\n search_settings_id,\n global_version.is_ee_version(),\n tenant_id,\n )\n\n if not job or not job.process:\n result.status = IndexingWatchdogTerminalStatus.SPAWN_FAILED\n task_logger.info(\n log_builder.build(\n \"Indexing watchdog - finished\",\n status=str(result.status.value),\n exit_code=str(result.exit_code),\n )\n )\n return\n\n # Ensure the process has moved out of the starting state\n num_waits = 0\n while True:\n if num_waits > 15:\n result.status = IndexingWatchdogTerminalStatus.SPAWN_NOT_ALIVE\n task_logger.info(\n log_builder.build(\n \"Indexing watchdog - finished\",\n status=str(result.status.value),\n exit_code=str(result.exit_code),\n )\n )\n job.release()\n return\n\n if job.process.is_alive() or job.process.exitcode is not None:\n break\n\n sleep(1)\n num_waits += 1\n\n task_logger.info(\n log_builder.build(\n \"Indexing watchdog - spawn succeeded\",\n pid=str(job.process.pid),\n )\n )\n\n # Track the last time memory info was emitted\n last_memory_emit_time = 0.0\n\n try:\n with get_session_with_current_tenant() as db_session:\n index_attempt = get_index_attempt(\n db_session=db_session,\n index_attempt_id=index_attempt_id,\n eager_load_cc_pair=True,\n )\n if not index_attempt:\n raise RuntimeError(\"Index attempt not found\")\n\n result.connector_source = (\n index_attempt.connector_credential_pair.connector.source.value\n )\n\n while True:\n sleep(5)\n\n time.monotonic()\n\n # if the job is done, clean up and break\n if job.done():\n try:\n result = process_job_result(\n job, result.connector_source, index_attempt_id, log_builder\n )\n except Exception:\n task_logger.exception(\n log_builder.build(\n \"Indexing watchdog - spawned task exceptioned\"\n )\n )\n finally:\n job.release()\n break\n\n # log the memory usage for tracking down memory leaks / connector-specific memory issues\n pid = job.process.pid\n if pid is not None:\n # Only emit memory info once per minute (60 seconds)\n current_time = time.monotonic()\n if current_time - last_memory_emit_time >= 60.0:\n emit_process_memory(\n pid,\n \"indexing_worker\",\n {\n \"cc_pair_id\": cc_pair_id,\n \"search_settings_id\": search_settings_id,\n \"index_attempt_id\": index_attempt_id,\n },\n )\n last_memory_emit_time = current_time\n\n # if the spawned task is still running, restart the check once again\n # if the index attempt is not in a finished status\n try:\n with get_session_with_current_tenant() as db_session:\n index_attempt = get_index_attempt(\n db_session=db_session, index_attempt_id=index_attempt_id\n )\n\n if not index_attempt:\n continue\n\n if not index_attempt.is_finished():\n continue\n\n except Exception:\n task_logger.exception(\n log_builder.build(\n \"Indexing watchdog - transient exception looking up index attempt\"\n )\n )\n continue\n\n except Exception as e:\n result.status = IndexingWatchdogTerminalStatus.WATCHDOG_EXCEPTIONED\n if isinstance(e, ConnectorValidationError):\n # No need to expose full stack trace for validation errors\n result.exception_str = str(e)\n else:\n result.exception_str = traceback.format_exc()\n\n # handle exit and reporting\n elapsed = time.monotonic() - start\n if result.exception_str is not None:\n # print with exception\n try:\n with get_session_with_current_tenant() as db_session:\n attempt = get_index_attempt(db_session, ctx.index_attempt_id)\n\n # only mark failures if not already terminal,\n # otherwise we're overwriting potential real stack traces\n if attempt and not attempt.status.is_terminal():\n failure_reason = (\n f\"Spawned task exceptioned: exit_code={result.exit_code}\"\n )\n mark_attempt_failed(\n ctx.index_attempt_id,\n db_session,\n failure_reason=failure_reason,\n full_exception_trace=result.exception_str,\n )\n except Exception:\n task_logger.exception(\n log_builder.build(\n \"Indexing watchdog - transient exception marking index attempt as failed\"\n )\n )\n\n normalized_exception_str = \"None\"\n if result.exception_str:\n normalized_exception_str = result.exception_str.replace(\n \"\\n\", \"\\\\n\"\n ).replace('\"', '\\\\\"')\n\n task_logger.warning(\n log_builder.build(\n \"Indexing watchdog - finished\",\n source=result.connector_source,\n status=result.status.value,\n exit_code=str(result.exit_code),\n exception=f'\"{normalized_exception_str}\"',\n elapsed=f\"{elapsed:.2f}s\",\n )\n )\n raise RuntimeError(f\"Exception encountered: traceback={result.exception_str}\")\n\n # print without exception\n if result.status == IndexingWatchdogTerminalStatus.TERMINATED_BY_SIGNAL:\n try:\n with get_session_with_current_tenant() as db_session:\n logger.exception(\n f\"Marking attempt {index_attempt_id} as canceled due to termination signal\"\n )\n mark_attempt_canceled(\n index_attempt_id,\n db_session,\n \"Connector termination signal detected\",\n )\n except Exception:\n task_logger.exception(\n log_builder.build(\n \"Indexing watchdog - transient exception marking index attempt as canceled\"\n )\n )\n\n job.cancel()\n elif result.status == IndexingWatchdogTerminalStatus.TERMINATED_BY_ACTIVITY_TIMEOUT:\n try:\n with get_session_with_current_tenant() as db_session:\n mark_attempt_failed(\n index_attempt_id,\n db_session,\n \"Indexing watchdog - activity timeout exceeded: \"\n f\"attempt={index_attempt_id} \"\n f\"timeout={CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT}s\",\n )\n except Exception:\n logger.exception(\n log_builder.build(\n \"Indexing watchdog - transient exception marking index attempt as failed\"\n )\n )\n job.cancel()\n else:\n pass\n\n task_logger.info(\n log_builder.build(\n \"Indexing watchdog - finished\",\n source=result.connector_source,\n status=str(result.status.value),\n exit_code=str(result.exit_code),\n elapsed=f\"{elapsed:.2f}s\",\n )\n )\n" }, { "path": "backend/onyx/background/celery/tasks/docprocessing/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/docprocessing/heartbeat.py", "content": "import contextvars\nimport threading\n\nfrom sqlalchemy import update\n\nfrom onyx.configs.constants import INDEXING_WORKER_HEARTBEAT_INTERVAL\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import IndexAttempt\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef start_heartbeat(index_attempt_id: int) -> tuple[threading.Thread, threading.Event]:\n \"\"\"Start a heartbeat thread for the given index attempt\"\"\"\n stop_event = threading.Event()\n\n def heartbeat_loop() -> None:\n while not stop_event.wait(INDEXING_WORKER_HEARTBEAT_INTERVAL):\n try:\n with get_session_with_current_tenant() as db_session:\n db_session.execute(\n update(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .values(heartbeat_counter=IndexAttempt.heartbeat_counter + 1)\n )\n db_session.commit()\n except Exception:\n logger.exception(\n \"Failed to update heartbeat counter for index attempt %s\",\n index_attempt_id,\n )\n\n # Ensure contextvars from the outer context are available in the thread\n context = contextvars.copy_context()\n thread = threading.Thread(target=context.run, args=(heartbeat_loop,), daemon=True)\n thread.start()\n return thread, stop_event\n\n\ndef stop_heartbeat(thread: threading.Thread, stop_event: threading.Event) -> None:\n \"\"\"Stop the heartbeat thread\"\"\"\n stop_event.set()\n thread.join(timeout=5) # Wait up to 5 seconds for clean shutdown\n" }, { "path": "backend/onyx/background/celery/tasks/docprocessing/tasks.py", "content": "import gc\nimport os\nimport time\nimport traceback\nfrom collections import defaultdict\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\n\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom fastapi import HTTPException\nfrom pydantic import BaseModel\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy import exists\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_find_task\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_unacked_task_ids\nfrom onyx.background.celery.celery_utils import httpx_init_vespa_pool\nfrom onyx.background.celery.memory_monitoring import emit_process_memory\nfrom onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT\nfrom onyx.background.celery.tasks.docfetching.task_creation_utils import (\n try_creating_docfetching_task,\n)\nfrom onyx.background.celery.tasks.docprocessing.heartbeat import start_heartbeat\nfrom onyx.background.celery.tasks.docprocessing.heartbeat import stop_heartbeat\nfrom onyx.background.celery.tasks.docprocessing.utils import IndexingCallback\nfrom onyx.background.celery.tasks.docprocessing.utils import is_in_repeated_error_state\nfrom onyx.background.celery.tasks.docprocessing.utils import should_index\nfrom onyx.background.celery.tasks.models import DocProcessingContext\nfrom onyx.background.indexing.checkpointing_utils import cleanup_checkpoint\nfrom onyx.background.indexing.checkpointing_utils import (\n get_index_attempts_with_old_checkpoints,\n)\nfrom onyx.background.indexing.index_attempt_utils import cleanup_index_attempts\nfrom onyx.background.indexing.index_attempt_utils import get_old_index_attempts\nfrom onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.app_configs import MANAGED_VESPA\nfrom onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH\nfrom onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH\nfrom onyx.configs.constants import AuthType\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_INDEXING_LOCK_TIMEOUT\nfrom onyx.configs.constants import MilestoneRecordType\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import OnyxRedisSignals\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import IndexAttemptMetadata\nfrom onyx.db.connector import mark_ccpair_with_indexing_trigger\nfrom onyx.db.connector_credential_pair import (\n fetch_indexable_standard_connector_credential_pair_ids,\n)\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.connector_credential_pair import set_cc_pair_repeated_error_state\nfrom onyx.db.connector_credential_pair import update_connector_credential_pair_from_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingMode\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import SwitchoverType\nfrom onyx.db.index_attempt import create_index_attempt_error\nfrom onyx.db.index_attempt import get_index_attempt\nfrom onyx.db.index_attempt import get_index_attempt_errors_for_cc_pair\nfrom onyx.db.index_attempt import IndexAttemptError\nfrom onyx.db.index_attempt import mark_attempt_canceled\nfrom onyx.db.index_attempt import mark_attempt_failed\nfrom onyx.db.index_attempt import mark_attempt_partially_succeeded\nfrom onyx.db.index_attempt import mark_attempt_succeeded\nfrom onyx.db.indexing_coordination import CoordinationStatus\nfrom onyx.db.indexing_coordination import INDEXING_PROGRESS_TIMEOUT_HOURS\nfrom onyx.db.indexing_coordination import IndexingCoordination\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.db.search_settings import get_secondary_search_settings\nfrom onyx.db.swap_index import check_and_perform_index_swap\nfrom onyx.document_index.factory import get_all_document_indices\nfrom onyx.file_store.document_batch_storage import DocumentBatchStorage\nfrom onyx.file_store.document_batch_storage import get_document_batch_storage\nfrom onyx.httpx.httpx_pool import HttpxPool\nfrom onyx.indexing.adapters.document_indexing_adapter import (\n DocumentIndexingBatchAdapter,\n)\nfrom onyx.indexing.embedder import DefaultIndexingEmbedder\nfrom onyx.indexing.indexing_pipeline import run_indexing_pipeline\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.redis.redis_pool import redis_lock_dump\nfrom onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT\nfrom onyx.redis.redis_utils import is_fence\nfrom onyx.server.runtime.onyx_runtime import OnyxRuntime\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.middleware import make_randomized_onyx_request_id\nfrom onyx.utils.telemetry import mt_cloud_telemetry\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\nfrom shared_configs.configs import INDEXING_MODEL_SERVER_HOST\nfrom shared_configs.configs import INDEXING_MODEL_SERVER_PORT\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import USAGE_LIMITS_ENABLED\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR\n\nlogger = setup_logger()\n\nDOCPROCESSING_STALL_TIMEOUT_MULTIPLIER = 4\nDOCPROCESSING_HEARTBEAT_TIMEOUT_MULTIPLIER = 24\n# Heartbeat timeout: if no heartbeat received for 30 minutes, consider it dead\n# This should be much longer than INDEXING_WORKER_HEARTBEAT_INTERVAL (30s)\nHEARTBEAT_TIMEOUT_SECONDS = 30 * 60 # 30 minutes\nINDEX_ATTEMPT_BATCH_SIZE = 500\n\n\ndef _get_fence_validation_block_expiration() -> int:\n \"\"\"\n Compute the expiration time for the fence validation block signal.\n Base expiration is 60 seconds, multiplied by the beat multiplier only in MULTI_TENANT mode.\n \"\"\"\n base_expiration = 60 # seconds\n\n if not MULTI_TENANT:\n return base_expiration\n\n try:\n beat_multiplier = OnyxRuntime.get_beat_multiplier()\n except Exception:\n beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n return int(base_expiration * beat_multiplier)\n\n\ndef validate_active_indexing_attempts(\n lock_beat: RedisLock,\n) -> None:\n \"\"\"\n Validates that active indexing attempts are still alive by checking heartbeat.\n If no heartbeat has been received for a certain amount of time, mark the attempt as failed.\n\n This uses the heartbeat_counter field which is incremented by active worker threads\n every INDEXING_WORKER_HEARTBEAT_INTERVAL seconds.\n \"\"\"\n logger.info(\"Validating active indexing attempts\")\n\n with get_session_with_current_tenant() as db_session:\n # Find all active indexing attempts\n active_attempts = (\n db_session.execute(\n select(IndexAttempt).where(\n IndexAttempt.status.in_([IndexingStatus.IN_PROGRESS]),\n IndexAttempt.celery_task_id.isnot(None),\n )\n )\n .scalars()\n .all()\n )\n\n for attempt in active_attempts:\n lock_beat.reacquire()\n\n # Initialize timeout for each attempt to prevent state pollution\n heartbeat_timeout_seconds = HEARTBEAT_TIMEOUT_SECONDS\n\n # Double-check the attempt still exists and has the same status\n fresh_attempt = get_index_attempt(db_session, attempt.id)\n if not fresh_attempt or fresh_attempt.status.is_terminal():\n continue\n\n # Check if this attempt has been updated with heartbeat tracking\n if fresh_attempt.last_heartbeat_time is None:\n # First time seeing this attempt - initialize heartbeat tracking\n fresh_attempt.last_heartbeat_value = fresh_attempt.heartbeat_counter\n fresh_attempt.last_heartbeat_time = datetime.now(timezone.utc)\n db_session.commit()\n\n task_logger.info(\n f\"Initialized heartbeat tracking for attempt {fresh_attempt.id}: counter={fresh_attempt.heartbeat_counter}\"\n )\n continue\n\n # Check if the heartbeat counter has advanced since last check\n current_counter = fresh_attempt.heartbeat_counter\n last_known_counter = fresh_attempt.last_heartbeat_value\n last_check_time = fresh_attempt.last_heartbeat_time\n\n task_logger.debug(\n f\"Checking heartbeat for attempt {fresh_attempt.id}: \"\n f\"current_counter={current_counter} \"\n f\"last_known_counter={last_known_counter} \"\n f\"last_check_time={last_check_time}\"\n )\n\n if current_counter > last_known_counter:\n # Heartbeat has advanced - worker is alive\n fresh_attempt.last_heartbeat_value = current_counter\n fresh_attempt.last_heartbeat_time = datetime.now(timezone.utc)\n db_session.commit()\n\n task_logger.debug(\n f\"Heartbeat advanced for attempt {fresh_attempt.id}: new_counter={current_counter}\"\n )\n continue\n\n if fresh_attempt.total_batches and fresh_attempt.completed_batches == 0:\n heartbeat_timeout_seconds = (\n HEARTBEAT_TIMEOUT_SECONDS\n * DOCPROCESSING_HEARTBEAT_TIMEOUT_MULTIPLIER\n )\n cutoff_time = datetime.now(timezone.utc) - timedelta(\n seconds=heartbeat_timeout_seconds\n )\n\n # Heartbeat hasn't advanced - check if it's been too long\n if last_check_time >= cutoff_time:\n task_logger.debug(\n f\"Heartbeat hasn't advanced for attempt {fresh_attempt.id} but still within timeout window\"\n )\n continue\n\n # No heartbeat for too long - mark as failed\n failure_reason = (\n f\"No heartbeat received for {heartbeat_timeout_seconds} seconds\"\n )\n\n task_logger.warning(\n f\"Heartbeat timeout for attempt {fresh_attempt.id}: \"\n f\"last_heartbeat_time={last_check_time} \"\n f\"cutoff_time={cutoff_time} \"\n f\"counter={current_counter}\"\n )\n\n try:\n mark_attempt_failed(\n fresh_attempt.id,\n db_session,\n failure_reason=failure_reason,\n )\n\n task_logger.error(\n f\"Marked attempt {fresh_attempt.id} as failed due to heartbeat timeout\"\n )\n\n except Exception:\n task_logger.exception(\n f\"Failed to mark attempt {fresh_attempt.id} as failed due to heartbeat timeout\"\n )\n\n\nclass ConnectorIndexingLogBuilder:\n def __init__(self, ctx: DocProcessingContext):\n self.ctx = ctx\n\n def build(self, msg: str, **kwargs: Any) -> str:\n msg_final = (\n f\"{msg}: \"\n f\"tenant_id={self.ctx.tenant_id} \"\n f\"attempt={self.ctx.index_attempt_id} \"\n f\"cc_pair={self.ctx.cc_pair_id} \"\n f\"search_settings={self.ctx.search_settings_id}\"\n )\n\n # Append extra keyword arguments in logfmt style\n if kwargs:\n extra_logfmt = \" \".join(f\"{key}={value}\" for key, value in kwargs.items())\n msg_final = f\"{msg_final} {extra_logfmt}\"\n\n return msg_final\n\n\ndef monitor_indexing_attempt_progress(\n attempt: IndexAttempt, tenant_id: str, db_session: Session, task: Task\n) -> None:\n \"\"\"\n TODO: rewrite this docstring\n Monitor the progress of an indexing attempt using database coordination.\n This replaces the Redis fence-based monitoring.\n\n Race condition handling:\n - Uses database coordination status to track progress\n - Only updates CC pair status based on confirmed database state\n - Handles concurrent completion gracefully\n \"\"\"\n if not attempt.celery_task_id:\n # Attempt hasn't been assigned a task yet\n return\n\n cc_pair = get_connector_credential_pair_from_id(\n db_session, attempt.connector_credential_pair_id\n )\n if not cc_pair:\n task_logger.warning(f\"CC pair not found for attempt {attempt.id}\")\n return\n\n # Check if the CC Pair should be moved to INITIAL_INDEXING\n if cc_pair.status == ConnectorCredentialPairStatus.SCHEDULED:\n cc_pair.status = ConnectorCredentialPairStatus.INITIAL_INDEXING\n db_session.commit()\n\n # Get coordination status to track progress\n\n coordination_status = IndexingCoordination.get_coordination_status(\n db_session, attempt.id\n )\n\n current_db_time = get_db_current_time(db_session)\n total_batches: int | str = (\n coordination_status.total_batches\n if coordination_status.total_batches is not None\n else \"?\"\n )\n if coordination_status.found:\n task_logger.info(\n f\"Indexing attempt progress: \"\n f\"attempt={attempt.id} \"\n f\"cc_pair={attempt.connector_credential_pair_id} \"\n f\"search_settings={attempt.search_settings_id} \"\n f\"completed_batches={coordination_status.completed_batches} \"\n f\"total_batches={total_batches} \"\n f\"total_docs={coordination_status.total_docs} \"\n f\"total_failures={coordination_status.total_failures}\"\n f\"elapsed={(current_db_time - attempt.time_created).seconds}\"\n )\n\n if coordination_status.cancellation_requested:\n task_logger.info(f\"Indexing attempt {attempt.id} has been cancelled\")\n mark_attempt_canceled(attempt.id, db_session)\n return\n\n storage = get_document_batch_storage(\n attempt.connector_credential_pair_id, attempt.id\n )\n\n # Check task completion using Celery\n try:\n check_indexing_completion(\n attempt.id, coordination_status, storage, tenant_id, task\n )\n except Exception as e:\n logger.exception(\n f\"Failed to monitor document processing completion: attempt={attempt.id} error={str(e)}\"\n )\n\n # Mark the attempt as failed if monitoring fails\n try:\n with get_session_with_current_tenant() as db_session:\n mark_attempt_failed(\n attempt.id,\n db_session,\n failure_reason=f\"Processing monitoring failed: {str(e)}\",\n full_exception_trace=traceback.format_exc(),\n )\n\n except Exception:\n logger.exception(\"Failed to mark attempt as failed\")\n\n # Try to clean up storage\n try:\n logger.info(f\"Cleaning up storage after monitoring failure: {storage}\")\n storage.cleanup_all_batches()\n except Exception:\n logger.exception(\"Failed to cleanup storage after monitoring failure\")\n\n\ndef _resolve_indexing_entity_errors(\n cc_pair_id: int,\n db_session: Session,\n) -> None:\n unresolved_errors = get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair_id,\n unresolved_only=True,\n db_session=db_session,\n )\n for error in unresolved_errors:\n if error.entity_id:\n error.is_resolved = True\n db_session.add(error)\n db_session.commit()\n\n\ndef check_indexing_completion(\n index_attempt_id: int,\n coordination_status: CoordinationStatus,\n storage: DocumentBatchStorage,\n tenant_id: str,\n task: Task,\n) -> None:\n\n logger.info(\n f\"Checking for indexing completion: attempt={index_attempt_id} tenant={tenant_id}\"\n )\n\n # Check if indexing is complete and all batches are processed\n batches_total = coordination_status.total_batches\n batches_processed = coordination_status.completed_batches\n indexing_completed = (\n batches_total is not None and batches_processed >= batches_total\n )\n\n logger.info(\n f\"Indexing status: \"\n f\"indexing_completed={indexing_completed} \"\n f\"batches_processed={batches_processed}/{batches_total if batches_total is not None else '?'} \"\n f\"total_docs={coordination_status.total_docs} \"\n f\"total_chunks={coordination_status.total_chunks} \"\n f\"total_failures={coordination_status.total_failures}\"\n )\n\n # Update progress tracking and check for stalls\n with get_session_with_current_tenant() as db_session:\n stalled_timeout_hours = INDEXING_PROGRESS_TIMEOUT_HOURS\n # Index attempts that are waiting between docfetching and\n # docprocessing get a generous stalling timeout\n if batches_total is not None and batches_processed == 0:\n stalled_timeout_hours = (\n stalled_timeout_hours * DOCPROCESSING_STALL_TIMEOUT_MULTIPLIER\n )\n\n timed_out = not IndexingCoordination.update_progress_tracking(\n db_session,\n index_attempt_id,\n batches_processed,\n timeout_hours=stalled_timeout_hours,\n )\n\n # Check for stalls (3-6 hour timeout). Only applies to in-progress attempts.\n attempt = get_index_attempt(db_session, index_attempt_id)\n if attempt and timed_out:\n if attempt.status == IndexingStatus.IN_PROGRESS:\n logger.error(\n f\"Indexing attempt {index_attempt_id} has been indexing for \"\n f\"{stalled_timeout_hours // 2}-{stalled_timeout_hours} hours without progress. \"\n f\"Marking it as failed.\"\n )\n mark_attempt_failed(\n index_attempt_id, db_session, failure_reason=\"Stalled indexing\"\n )\n elif (\n attempt.status == IndexingStatus.NOT_STARTED and attempt.celery_task_id\n ):\n # Check if the task exists in the celery queue\n # This handles the case where Redis dies after task creation but before task execution\n redis_celery = celery_get_broker_client(task.app)\n task_exists = celery_find_task(\n attempt.celery_task_id,\n OnyxCeleryQueues.CONNECTOR_DOC_FETCHING,\n redis_celery,\n )\n unacked_task_ids = celery_get_unacked_task_ids(\n OnyxCeleryQueues.CONNECTOR_DOC_FETCHING, redis_celery\n )\n\n if not task_exists and attempt.celery_task_id not in unacked_task_ids:\n # there is a race condition where the docfetching task has been taken off\n # the queues (i.e. started) but the indexing attempt still has a status of\n # Not Started because the switch to in progress takes like 0.1 seconds.\n # sleep a bit and confirm that the attempt is still not in progress.\n time.sleep(1)\n attempt = get_index_attempt(db_session, index_attempt_id)\n if attempt and attempt.status == IndexingStatus.NOT_STARTED:\n logger.error(\n f\"Task {attempt.celery_task_id} attached to indexing attempt \"\n f\"{index_attempt_id} does not exist in the queue. \"\n f\"Marking indexing attempt as failed.\"\n )\n mark_attempt_failed(\n index_attempt_id,\n db_session,\n failure_reason=\"Task not in queue\",\n )\n else:\n logger.info(\n f\"Indexing attempt {index_attempt_id} is {attempt.status}. 3-6 hours without heartbeat \"\n \"but task is in the queue. Likely underprovisioned docfetching worker.\"\n )\n # Update last progress time so we won't time out again for another 3 hours\n IndexingCoordination.update_progress_tracking(\n db_session,\n index_attempt_id,\n batches_processed,\n force_update_progress=True,\n )\n\n # check again on the next check_for_indexing task\n # TODO: on the cloud this is currently 25 minutes at most, which\n # is honestly too slow. We should either increase the frequency of\n # this task or change where we check for completion.\n if not indexing_completed:\n return\n\n # If processing is complete, handle completion\n logger.info(f\"Connector indexing finished for index attempt {index_attempt_id}.\")\n\n # All processing is complete\n total_failures = coordination_status.total_failures\n\n with get_session_with_current_tenant() as db_session:\n if total_failures == 0:\n attempt = mark_attempt_succeeded(index_attempt_id, db_session)\n logger.info(f\"Index attempt {index_attempt_id} completed successfully\")\n else:\n attempt = mark_attempt_partially_succeeded(index_attempt_id, db_session)\n logger.info(\n f\"Index attempt {index_attempt_id} completed with {total_failures} failures\"\n )\n\n # Update CC pair status if successful\n cc_pair = get_connector_credential_pair_from_id(\n db_session, attempt.connector_credential_pair_id\n )\n if cc_pair is None:\n raise RuntimeError(\n f\"CC pair {attempt.connector_credential_pair_id} not found in database\"\n )\n\n if attempt.status.is_successful():\n # NOTE: we define the last successful index time as the time the last successful\n # attempt finished. This is distinct from the poll_range_end of the last successful\n # attempt, which is the time up to which documents have been fetched.\n cc_pair.last_successful_index_time = attempt.time_updated\n if cc_pair.status in [\n ConnectorCredentialPairStatus.SCHEDULED,\n ConnectorCredentialPairStatus.INITIAL_INDEXING,\n ]:\n # User file connectors must be paused on success\n # NOTE: _run_indexing doesn't update connectors if the index attempt is the future embedding model\n cc_pair.status = ConnectorCredentialPairStatus.ACTIVE\n db_session.commit()\n\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=tenant_id,\n event=MilestoneRecordType.CONNECTOR_SUCCEEDED,\n )\n\n # Clear repeated error state on success\n if cc_pair.in_repeated_error_state:\n cc_pair.in_repeated_error_state = False\n db_session.commit()\n\n if attempt.status == IndexingStatus.SUCCESS:\n logger.info(\n f\"Resolving indexing entity errors for attempt {index_attempt_id}\"\n )\n _resolve_indexing_entity_errors(\n cc_pair_id=attempt.connector_credential_pair_id,\n db_session=db_session,\n )\n\n # Clean up FileStore storage (still needed for document batches during transition)\n try:\n logger.info(f\"Cleaning up storage after indexing completion: {storage}\")\n storage.cleanup_all_batches()\n except Exception:\n logger.exception(\"Failed to clean up document batches - continuing\")\n\n logger.info(f\"Database coordination completed for attempt {index_attempt_id}\")\n\n\ndef active_indexing_attempt(\n cc_pair_id: int,\n search_settings_id: int,\n db_session: Session,\n) -> bool:\n \"\"\"\n Check if there's already an active indexing attempt for this CC pair + search settings.\n This prevents race conditions where multiple indexing attempts could be created.\n We check for any non-terminal status (NOT_STARTED, IN_PROGRESS).\n\n Returns True if there's an active indexing attempt, False otherwise.\n \"\"\"\n active_indexing_attempt = db_session.execute(\n select(\n exists().where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n IndexAttempt.search_settings_id == search_settings_id,\n IndexAttempt.status.in_(\n [\n IndexingStatus.NOT_STARTED,\n IndexingStatus.IN_PROGRESS,\n ]\n ),\n )\n )\n ).scalar()\n\n if active_indexing_attempt:\n task_logger.debug(\n f\"active_indexing_attempt - Skipping due to active indexing attempt: \"\n f\"cc_pair={cc_pair_id} search_settings={search_settings_id}\"\n )\n\n return bool(active_indexing_attempt)\n\n\ndef _kickoff_indexing_tasks(\n celery_app: Celery,\n db_session: Session,\n search_settings: SearchSettings,\n cc_pair_ids: list[int],\n secondary_index_building: bool,\n redis_client: Redis,\n lock_beat: RedisLock,\n tenant_id: str,\n) -> int:\n \"\"\"Kick off indexing tasks for the given cc_pair_ids and search_settings.\n\n Returns the number of tasks successfully created.\n \"\"\"\n tasks_created = 0\n\n for cc_pair_id in cc_pair_ids:\n lock_beat.reacquire()\n\n # Lightweight check prior to fetching cc pair\n if active_indexing_attempt(\n cc_pair_id=cc_pair_id,\n search_settings_id=search_settings.id,\n db_session=db_session,\n ):\n continue\n\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n task_logger.warning(\n f\"_kickoff_indexing_tasks - CC pair not found: cc_pair={cc_pair_id}\"\n )\n continue\n\n # Heavyweight check after fetching cc pair\n if not should_index(\n cc_pair=cc_pair,\n search_settings_instance=search_settings,\n secondary_index_building=secondary_index_building,\n db_session=db_session,\n ):\n task_logger.debug(\n f\"_kickoff_indexing_tasks - Not indexing cc_pair_id: {cc_pair_id} \"\n f\"search_settings={search_settings.id}, \"\n f\"secondary_index_building={secondary_index_building}\"\n )\n continue\n\n task_logger.debug(\n f\"_kickoff_indexing_tasks - Will index cc_pair_id: {cc_pair_id} \"\n f\"search_settings={search_settings.id}, \"\n f\"secondary_index_building={secondary_index_building}\"\n )\n\n reindex = False\n # the indexing trigger is only checked and cleared with the current search settings\n if search_settings.status.is_current() and cc_pair.indexing_trigger is not None:\n if cc_pair.indexing_trigger == IndexingMode.REINDEX:\n reindex = True\n\n task_logger.info(\n f\"_kickoff_indexing_tasks - Connector indexing manual trigger detected: \"\n f\"cc_pair={cc_pair.id} \"\n f\"search_settings={search_settings.id} \"\n f\"indexing_mode={cc_pair.indexing_trigger}\"\n )\n\n mark_ccpair_with_indexing_trigger(cc_pair.id, None, db_session)\n\n # using a task queue and only allowing one task per cc_pair/search_setting\n # prevents us from starving out certain attempts\n attempt_id = try_creating_docfetching_task(\n celery_app,\n cc_pair,\n search_settings,\n reindex,\n db_session,\n redis_client,\n tenant_id,\n )\n\n if attempt_id is not None:\n task_logger.info(\n f\"Connector indexing queued: index_attempt={attempt_id} cc_pair={cc_pair.id} search_settings={search_settings.id}\"\n )\n tasks_created += 1\n else:\n task_logger.error(\n f\"Failed to create indexing task: cc_pair={cc_pair.id} search_settings={search_settings.id}\"\n )\n\n return tasks_created\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_INDEXING,\n soft_time_limit=300,\n bind=True,\n)\ndef check_for_indexing(self: Task, *, tenant_id: str) -> int | None:\n \"\"\"a lightweight task used to kick off the pipeline of indexing tasks.\n Occcasionally does some validation of existing state to clear up error conditions.\n\n This task is the entrypoint for the full \"indexing pipeline\", which is composed\n of two tasks: \"docfetching\" and \"docprocessing\". More details in\n the docfetching task (OnyxCeleryTask.CONNECTOR_DOC_FETCHING_TASK).\n\n For cc pairs that should be indexed (see should_index()), this task\n calls try_creating_docfetching_task, which creates a docfetching task.\n All the logic for determining what state the indexing pipeline is in\n w.r.t previous failed attempt, checkpointing, etc is handled in the docfetching task.\n \"\"\"\n\n time_start = time.monotonic()\n task_logger.warning(\"check_for_indexing - Starting\")\n\n tasks_created = 0\n locked = False\n redis_client = get_redis_client()\n redis_client_replica = get_redis_replica_client()\n\n # we need to use celery's redis client to access its redis data\n # (which lives on a different db number)\n # redis_client_celery: Redis = self.app.broker_connection().channel().client # type: ignore\n\n lock_beat: RedisLock = redis_client.lock(\n OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n try:\n locked = True\n\n # SPECIAL 0/3: sync lookup table for active fences\n # we want to run this less frequently than the overall task\n if not redis_client.exists(OnyxRedisSignals.BLOCK_BUILD_FENCE_LOOKUP_TABLE):\n # build a lookup table of existing fences\n # this is just a migration concern and should be unnecessary once\n # lookup tables are rolled out\n for key_bytes in redis_client_replica.scan_iter(\n count=SCAN_ITER_COUNT_DEFAULT\n ):\n if is_fence(key_bytes) and not redis_client.sismember(\n OnyxRedisConstants.ACTIVE_FENCES, key_bytes\n ):\n logger.warning(f\"Adding {key_bytes} to the lookup table.\")\n redis_client.sadd(OnyxRedisConstants.ACTIVE_FENCES, key_bytes)\n\n redis_client.set(\n OnyxRedisSignals.BLOCK_BUILD_FENCE_LOOKUP_TABLE,\n 1,\n ex=OnyxRuntime.get_build_fence_lookup_table_interval(),\n )\n\n # 1/3: KICKOFF\n\n # check for search settings swap\n with get_session_with_current_tenant() as db_session:\n old_search_settings = check_and_perform_index_swap(db_session=db_session)\n current_search_settings = get_current_search_settings(db_session)\n # So that the first time users aren't surprised by really slow speed of first\n # batch of documents indexed\n if current_search_settings.provider_type is None and not MULTI_TENANT:\n if old_search_settings:\n embedding_model = EmbeddingModel.from_db_model(\n search_settings=current_search_settings,\n server_host=INDEXING_MODEL_SERVER_HOST,\n server_port=INDEXING_MODEL_SERVER_PORT,\n )\n\n # only warm up if search settings were changed\n warm_up_bi_encoder(\n embedding_model=embedding_model,\n )\n\n # gather search settings and indexable cc_pair_ids\n # indexable CC pairs include everything for future model and only active cc pairs for current model\n lock_beat.reacquire()\n with get_session_with_current_tenant() as db_session:\n # Get CC pairs for primary search settings\n standard_cc_pair_ids = (\n fetch_indexable_standard_connector_credential_pair_ids(\n db_session, active_cc_pairs_only=True\n )\n )\n\n primary_cc_pair_ids = standard_cc_pair_ids\n\n # Get CC pairs for secondary search settings\n secondary_cc_pair_ids: list[int] = []\n secondary_search_settings = get_secondary_search_settings(db_session)\n if secondary_search_settings:\n # For ACTIVE_ONLY, we skip paused connectors\n include_paused = (\n secondary_search_settings.switchover_type\n != SwitchoverType.ACTIVE_ONLY\n )\n standard_cc_pair_ids = (\n fetch_indexable_standard_connector_credential_pair_ids(\n db_session, active_cc_pairs_only=not include_paused\n )\n )\n\n secondary_cc_pair_ids = standard_cc_pair_ids\n\n # Flag CC pairs in repeated error state for primary/current search settings\n with get_session_with_current_tenant() as db_session:\n for cc_pair_id in primary_cc_pair_ids:\n lock_beat.reacquire()\n\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n\n # if already in repeated error state, don't do anything\n # this is important so that we don't keep pausing the connector\n # immediately upon a user un-pausing it to manually re-trigger and\n # recover.\n if (\n cc_pair\n and not cc_pair.in_repeated_error_state\n and is_in_repeated_error_state(\n cc_pair=cc_pair,\n search_settings_id=current_search_settings.id,\n db_session=db_session,\n )\n ):\n set_cc_pair_repeated_error_state(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n in_repeated_error_state=True,\n )\n # When entering repeated error state, also pause the connector\n # to prevent continued indexing retry attempts burning through embedding credits.\n # NOTE: only for Cloud, since most self-hosted users use self-hosted embedding\n # models. Also, they are more prone to repeated failures -> eventual success.\n if AUTH_TYPE == AuthType.CLOUD:\n update_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair.id,\n status=ConnectorCredentialPairStatus.PAUSED,\n )\n\n # NOTE: At this point, we haven't done heavy checks on whether or not the CC pairs should actually be indexed\n # Heavy check, should_index(), is called in _kickoff_indexing_tasks\n with get_session_with_current_tenant() as db_session:\n # Primary first\n tasks_created += _kickoff_indexing_tasks(\n celery_app=self.app,\n db_session=db_session,\n search_settings=current_search_settings,\n cc_pair_ids=primary_cc_pair_ids,\n secondary_index_building=secondary_search_settings is not None,\n redis_client=redis_client,\n lock_beat=lock_beat,\n tenant_id=tenant_id,\n )\n\n # Secondary indexing (only if secondary search settings exist and switchover_type is not INSTANT)\n if (\n secondary_search_settings\n and secondary_search_settings.switchover_type != SwitchoverType.INSTANT\n and secondary_cc_pair_ids\n ):\n tasks_created += _kickoff_indexing_tasks(\n celery_app=self.app,\n db_session=db_session,\n search_settings=secondary_search_settings,\n cc_pair_ids=secondary_cc_pair_ids,\n secondary_index_building=True,\n redis_client=redis_client,\n lock_beat=lock_beat,\n tenant_id=tenant_id,\n )\n elif (\n secondary_search_settings\n and secondary_search_settings.switchover_type == SwitchoverType.INSTANT\n ):\n task_logger.info(\n f\"Skipping secondary indexing: switchover_type=INSTANT for search_settings={secondary_search_settings.id}\"\n )\n\n # 2/3: VALIDATE\n # Check for inconsistent index attempts - active attempts without task IDs\n # This can happen if attempt creation fails partway through\n lock_beat.reacquire()\n with get_session_with_current_tenant() as db_session:\n inconsistent_attempts = (\n db_session.execute(\n select(IndexAttempt).where(\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n ),\n IndexAttempt.celery_task_id.is_(None),\n )\n )\n .scalars()\n .all()\n )\n\n for attempt in inconsistent_attempts:\n lock_beat.reacquire()\n\n # Double-check the attempt still has the inconsistent state\n fresh_attempt = get_index_attempt(db_session, attempt.id)\n if (\n not fresh_attempt\n or fresh_attempt.celery_task_id\n or fresh_attempt.status.is_terminal()\n ):\n continue\n\n failure_reason = (\n f\"Inconsistent index attempt found - active status without Celery task: \"\n f\"index_attempt={attempt.id} \"\n f\"cc_pair={attempt.connector_credential_pair_id} \"\n f\"search_settings={attempt.search_settings_id}\"\n )\n task_logger.error(failure_reason)\n mark_attempt_failed(\n attempt.id, db_session, failure_reason=failure_reason\n )\n\n lock_beat.reacquire()\n # we want to run this less frequently than the overall task\n if not redis_client.exists(OnyxRedisSignals.BLOCK_VALIDATE_INDEXING_FENCES):\n # Check for orphaned index attempts that have Celery task IDs but no actual running tasks\n # This can happen if workers crash or tasks are terminated unexpectedly\n # We reuse the same Redis signal name for backwards compatibility\n try:\n validate_active_indexing_attempts(lock_beat)\n except Exception:\n task_logger.exception(\n \"Exception while validating active indexing attempts\"\n )\n\n redis_client.set(\n OnyxRedisSignals.BLOCK_VALIDATE_INDEXING_FENCES,\n 1,\n ex=_get_fence_validation_block_expiration(),\n )\n\n # 3/3: FINALIZE - Monitor active indexing attempts using database\n lock_beat.reacquire()\n with get_session_with_current_tenant() as db_session:\n # Monitor all active indexing attempts directly from the database\n # This replaces the Redis fence-based monitoring\n active_attempts = (\n db_session.execute(\n select(IndexAttempt).where(\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n )\n )\n )\n .scalars()\n .all()\n )\n\n for attempt in active_attempts:\n try:\n monitor_indexing_attempt_progress(\n attempt, tenant_id, db_session, self\n )\n except Exception:\n task_logger.exception(f\"Error monitoring attempt {attempt.id}\")\n\n lock_beat.reacquire()\n\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception:\n task_logger.exception(\"Unexpected exception during indexing check\")\n finally:\n if locked:\n if lock_beat.owned():\n lock_beat.release()\n else:\n task_logger.error(\n f\"check_for_indexing - Lock not owned on completion: tenant={tenant_id}\"\n )\n redis_lock_dump(lock_beat, redis_client)\n\n time_elapsed = time.monotonic() - time_start\n task_logger.info(f\"check_for_indexing finished: elapsed={time_elapsed:.2f}\")\n return tasks_created\n\n\n# primary\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_CHECKPOINT_CLEANUP,\n soft_time_limit=300,\n bind=True,\n)\ndef check_for_checkpoint_cleanup(self: Task, *, tenant_id: str) -> None:\n \"\"\"Clean up old checkpoints that are older than 7 days.\"\"\"\n locked = False\n redis_client = get_redis_client(tenant_id=tenant_id)\n lock: RedisLock = redis_client.lock(\n OnyxRedisLocks.CHECK_CHECKPOINT_CLEANUP_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock.acquire(blocking=False):\n return None\n\n try:\n locked = True\n with get_session_with_current_tenant() as db_session:\n old_attempts = get_index_attempts_with_old_checkpoints(db_session)\n for attempt in old_attempts:\n task_logger.info(\n f\"Cleaning up checkpoint for index attempt {attempt.id}\"\n )\n self.app.send_task(\n OnyxCeleryTask.CLEANUP_CHECKPOINT,\n kwargs={\n \"index_attempt_id\": attempt.id,\n \"tenant_id\": tenant_id,\n },\n queue=OnyxCeleryQueues.CHECKPOINT_CLEANUP,\n priority=OnyxCeleryPriority.MEDIUM,\n )\n except Exception:\n task_logger.exception(\"Unexpected exception during checkpoint cleanup\")\n return None\n finally:\n if locked:\n if lock.owned():\n lock.release()\n else:\n task_logger.error(\n f\"check_for_checkpoint_cleanup - Lock not owned on completion: tenant={tenant_id}\"\n )\n\n\n# light worker\n@shared_task(\n name=OnyxCeleryTask.CLEANUP_CHECKPOINT,\n bind=True,\n)\ndef cleanup_checkpoint_task(\n self: Task, # noqa: ARG001\n *,\n index_attempt_id: int,\n tenant_id: str | None,\n) -> None:\n \"\"\"Clean up a checkpoint for a given index attempt\"\"\"\n\n start = time.monotonic()\n\n try:\n with get_session_with_current_tenant() as db_session:\n cleanup_checkpoint(db_session, index_attempt_id)\n finally:\n elapsed = time.monotonic() - start\n\n task_logger.info(\n f\"cleanup_checkpoint_task completed: tenant_id={tenant_id} index_attempt_id={index_attempt_id} elapsed={elapsed:.2f}\"\n )\n\n\n# primary\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_INDEX_ATTEMPT_CLEANUP,\n soft_time_limit=300,\n bind=True,\n)\ndef check_for_index_attempt_cleanup(self: Task, *, tenant_id: str) -> None:\n \"\"\"Clean up old index attempts that are older than 7 days.\"\"\"\n locked = False\n redis_client = get_redis_client(tenant_id=tenant_id)\n lock: RedisLock = redis_client.lock(\n OnyxRedisLocks.CHECK_INDEX_ATTEMPT_CLEANUP_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock.acquire(blocking=False):\n task_logger.info(\n f\"check_for_index_attempt_cleanup - Lock not acquired: tenant={tenant_id}\"\n )\n return None\n\n try:\n locked = True\n batch_size = INDEX_ATTEMPT_BATCH_SIZE\n with get_session_with_current_tenant() as db_session:\n old_attempts = get_old_index_attempts(db_session)\n # We need to batch this because during the initial run, the system might have a large number\n # of index attempts since they were never deleted. After that, the number will be\n # significantly lower.\n if len(old_attempts) == 0:\n task_logger.info(\n \"check_for_index_attempt_cleanup - No index attempts to cleanup\"\n )\n return\n\n for i in range(0, len(old_attempts), batch_size):\n batch = old_attempts[i : i + batch_size]\n task_logger.info(\n f\"check_for_index_attempt_cleanup - Cleaning up index attempts {len(batch)}\"\n )\n self.app.send_task(\n OnyxCeleryTask.CLEANUP_INDEX_ATTEMPT,\n kwargs={\n \"index_attempt_ids\": [attempt.id for attempt in batch],\n \"tenant_id\": tenant_id,\n },\n queue=OnyxCeleryQueues.INDEX_ATTEMPT_CLEANUP,\n priority=OnyxCeleryPriority.MEDIUM,\n )\n except Exception:\n task_logger.exception(\"Unexpected exception during index attempt cleanup check\")\n return None\n finally:\n if locked:\n if lock.owned():\n lock.release()\n else:\n task_logger.error(\n f\"check_for_index_attempt_cleanup - Lock not owned on completion: tenant={tenant_id}\"\n )\n\n\n# light worker\n@shared_task(\n name=OnyxCeleryTask.CLEANUP_INDEX_ATTEMPT,\n bind=True,\n)\ndef cleanup_index_attempt_task(\n self: Task, # noqa: ARG001\n *,\n index_attempt_ids: list[int],\n tenant_id: str,\n) -> None:\n \"\"\"Clean up an index attempt\"\"\"\n start = time.monotonic()\n\n try:\n with get_session_with_current_tenant() as db_session:\n cleanup_index_attempts(db_session, index_attempt_ids)\n\n finally:\n elapsed = time.monotonic() - start\n\n task_logger.info(\n f\"cleanup_index_attempt_task completed: tenant_id={tenant_id} \"\n f\"index_attempt_ids={index_attempt_ids} \"\n f\"elapsed={elapsed:.2f}\"\n )\n\n\nclass DocumentProcessingBatch(BaseModel):\n \"\"\"Data structure for a document processing batch.\"\"\"\n\n batch_id: str\n index_attempt_id: int\n cc_pair_id: int\n tenant_id: str\n batch_num: int\n\n\ndef _check_failure_threshold(\n total_failures: int,\n document_count: int,\n batch_num: int,\n last_failure: ConnectorFailure | None,\n) -> None:\n \"\"\"Check if we've hit the failure threshold and raise an appropriate exception if so.\n\n We consider the threshold hit if:\n 1. We have more than 3 failures AND\n 2. Failures account for more than 10% of processed documents\n \"\"\"\n failure_ratio = total_failures / (document_count or 1)\n\n FAILURE_THRESHOLD = 3\n FAILURE_RATIO_THRESHOLD = 0.1\n if total_failures > FAILURE_THRESHOLD and failure_ratio > FAILURE_RATIO_THRESHOLD:\n logger.error(\n f\"Connector run failed with '{total_failures}' errors after '{batch_num}' batches.\"\n )\n if last_failure and last_failure.exception:\n raise last_failure.exception from last_failure.exception\n\n raise RuntimeError(\n f\"Connector run encountered too many errors, aborting. Last error: {last_failure}\"\n )\n\n\ndef _resolve_indexing_document_errors(\n cc_pair_id: int,\n failures: list[ConnectorFailure],\n document_batch: list[Document],\n) -> None:\n with get_session_with_current_tenant() as db_session_temp:\n # get previously unresolved errors\n unresolved_errors = get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair_id,\n unresolved_only=True,\n db_session=db_session_temp,\n )\n doc_id_to_unresolved_errors: dict[str, list[IndexAttemptError]] = defaultdict(\n list\n )\n for error in unresolved_errors:\n if error.document_id:\n doc_id_to_unresolved_errors[error.document_id].append(error)\n\n # resolve errors for documents that were successfully indexed\n failed_document_ids = [\n failure.failed_document.document_id\n for failure in failures\n if failure.failed_document\n ]\n successful_document_ids = [\n document.id\n for document in document_batch\n if document.id not in failed_document_ids\n ]\n for document_id in successful_document_ids:\n if document_id not in doc_id_to_unresolved_errors:\n continue\n\n logger.info(f\"Resolving IndexAttemptError for document '{document_id}'\")\n for error in doc_id_to_unresolved_errors[document_id]:\n error.is_resolved = True\n db_session_temp.add(error)\n\n db_session_temp.commit()\n\n\n@shared_task(\n name=OnyxCeleryTask.DOCPROCESSING_TASK,\n bind=True,\n)\ndef docprocessing_task(\n self: Task, # noqa: ARG001\n index_attempt_id: int,\n cc_pair_id: int,\n tenant_id: str,\n batch_num: int,\n) -> None:\n \"\"\"Process a batch of documents through the indexing pipeline.\n\n This task retrieves documents from storage and processes them through\n the indexing pipeline (embedding + vector store indexing).\n \"\"\"\n # Start heartbeat for this indexing attempt\n heartbeat_thread, stop_event = start_heartbeat(index_attempt_id)\n try:\n # Cannot use the TaskSingleton approach here because the worker is multithreaded\n token = INDEX_ATTEMPT_INFO_CONTEXTVAR.set((cc_pair_id, index_attempt_id))\n _docprocessing_task(index_attempt_id, cc_pair_id, tenant_id, batch_num)\n finally:\n stop_heartbeat(heartbeat_thread, stop_event) # Stop heartbeat before exiting\n INDEX_ATTEMPT_INFO_CONTEXTVAR.reset(token)\n\n\ndef _check_chunk_usage_limit(tenant_id: str) -> None:\n \"\"\"Check if chunk indexing usage limit has been exceeded.\n\n Raises UsageLimitExceededError if the limit is exceeded.\n \"\"\"\n if not USAGE_LIMITS_ENABLED:\n return\n\n from onyx.db.usage import UsageType\n from onyx.server.usage_limits import check_usage_and_raise\n\n with get_session_with_current_tenant() as db_session:\n check_usage_and_raise(\n db_session=db_session,\n usage_type=UsageType.CHUNKS_INDEXED,\n tenant_id=tenant_id,\n pending_amount=0, # Just check current usage\n )\n\n\ndef _docprocessing_task(\n index_attempt_id: int,\n cc_pair_id: int,\n tenant_id: str,\n batch_num: int,\n) -> None:\n start_time = time.monotonic()\n\n if tenant_id:\n CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n # Check if chunk indexing usage limit has been exceeded before processing\n if USAGE_LIMITS_ENABLED:\n try:\n _check_chunk_usage_limit(tenant_id)\n except HTTPException as e:\n # Log the error and fail the indexing attempt\n task_logger.error(\n f\"Chunk indexing usage limit exceeded for tenant {tenant_id}: {e}\"\n )\n with get_session_with_current_tenant() as db_session:\n from onyx.db.index_attempt import mark_attempt_failed\n\n mark_attempt_failed(\n index_attempt_id=index_attempt_id,\n db_session=db_session,\n failure_reason=str(e),\n )\n raise\n\n task_logger.info(\n f\"Processing document batch: attempt={index_attempt_id} batch_num={batch_num} \"\n )\n\n # Get the document batch storage\n storage = get_document_batch_storage(cc_pair_id, index_attempt_id)\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n r = get_redis_client(tenant_id=tenant_id)\n\n # 20 is the documented default for httpx max_keepalive_connections\n if MANAGED_VESPA:\n httpx_init_vespa_pool(\n 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH\n )\n else:\n httpx_init_vespa_pool(20)\n\n # dummy lock to satisfy linter\n per_batch_lock: RedisLock | None = None\n try:\n # FIX: Monitor memory before loading documents to track problematic batches\n emit_process_memory(\n os.getpid(),\n \"docprocessing\",\n {\n \"phase\": \"before_load\",\n \"tenant_id\": tenant_id,\n \"cc_pair_id\": cc_pair_id,\n \"index_attempt_id\": index_attempt_id,\n \"batch_num\": batch_num,\n },\n )\n\n # Retrieve documents from storage\n documents = storage.get_batch(batch_num)\n if not documents:\n task_logger.error(f\"No documents found for batch {batch_num}\")\n return\n\n # FIX: Monitor memory after loading documents\n emit_process_memory(\n os.getpid(),\n \"docprocessing\",\n {\n \"phase\": \"after_load\",\n \"tenant_id\": tenant_id,\n \"cc_pair_id\": cc_pair_id,\n \"index_attempt_id\": index_attempt_id,\n \"batch_num\": batch_num,\n \"doc_count\": len(documents),\n },\n )\n\n with get_session_with_current_tenant() as db_session:\n # matches parts of _run_indexing\n index_attempt = get_index_attempt(\n db_session,\n index_attempt_id,\n eager_load_cc_pair=True,\n eager_load_search_settings=True,\n )\n if not index_attempt:\n raise RuntimeError(f\"Index attempt {index_attempt_id} not found\")\n\n if index_attempt.search_settings is None:\n raise ValueError(\"Search settings must be set for indexing\")\n\n if (\n index_attempt.celery_task_id is None\n or index_attempt.status.is_terminal()\n ):\n raise RuntimeError(\n f\"Index attempt {index_attempt_id} is not running, status {index_attempt.status}\"\n )\n\n cross_batch_db_lock: RedisLock = r.lock(\n redis_connector.db_lock_key(index_attempt.search_settings.id),\n timeout=CELERY_INDEXING_LOCK_TIMEOUT,\n thread_local=False,\n )\n\n callback = IndexingCallback(\n redis_connector,\n )\n # TODO: right now this is the only thing the callback is used for,\n # probably there is a simpler way to handle pausing\n if callback.should_stop():\n raise RuntimeError(\"Docprocessing cancelled by connector pausing\")\n\n # Set up indexing pipeline components\n embedding_model = DefaultIndexingEmbedder.from_db_search_settings(\n search_settings=index_attempt.search_settings,\n callback=callback,\n )\n\n document_indices = get_all_document_indices(\n index_attempt.search_settings,\n None,\n httpx_client=HttpxPool.get(\"vespa\"),\n )\n\n # Set up metadata for this batch\n index_attempt_metadata = IndexAttemptMetadata(\n attempt_id=index_attempt_id,\n connector_id=index_attempt.connector_credential_pair.connector.id,\n credential_id=index_attempt.connector_credential_pair.credential.id,\n request_id=make_randomized_onyx_request_id(\"DIP\"),\n structured_id=f\"{tenant_id}:{cc_pair_id}:{index_attempt_id}:{batch_num}\",\n batch_num=batch_num,\n )\n\n # Process documents through indexing pipeline\n connector_source = (\n index_attempt.connector_credential_pair.connector.source.value\n )\n task_logger.info(\n f\"Processing {len(documents)} documents through indexing pipeline: \"\n f\"cc_pair_id={cc_pair_id}, source={connector_source}, \"\n f\"batch_num={batch_num}\"\n )\n\n adapter = DocumentIndexingBatchAdapter(\n db_session=db_session,\n connector_id=index_attempt.connector_credential_pair.connector.id,\n credential_id=index_attempt.connector_credential_pair.credential.id,\n tenant_id=tenant_id,\n index_attempt_metadata=index_attempt_metadata,\n )\n\n # real work happens here!\n index_pipeline_result = run_indexing_pipeline(\n embedder=embedding_model,\n document_indices=document_indices,\n ignore_time_skip=True, # Documents are already filtered during extraction\n db_session=db_session,\n tenant_id=tenant_id,\n document_batch=documents,\n request_id=index_attempt_metadata.request_id,\n adapter=adapter,\n )\n\n # Track chunk indexing usage for cloud usage limits\n if USAGE_LIMITS_ENABLED and index_pipeline_result.total_chunks > 0:\n try:\n from onyx.db.usage import increment_usage\n from onyx.db.usage import UsageType\n\n with get_session_with_current_tenant() as usage_db_session:\n increment_usage(\n db_session=usage_db_session,\n usage_type=UsageType.CHUNKS_INDEXED,\n amount=index_pipeline_result.total_chunks,\n )\n usage_db_session.commit()\n except Exception as e:\n # Log but don't fail indexing if usage tracking fails\n task_logger.warning(f\"Failed to track chunk indexing usage: {e}\")\n\n # Update batch completion and document counts atomically using database coordination\n\n with get_session_with_current_tenant() as db_session, cross_batch_db_lock:\n IndexingCoordination.update_batch_completion_and_docs(\n db_session=db_session,\n index_attempt_id=index_attempt_id,\n total_docs_indexed=index_pipeline_result.total_docs,\n new_docs_indexed=index_pipeline_result.new_docs,\n total_chunks=index_pipeline_result.total_chunks,\n )\n\n _resolve_indexing_document_errors(\n cc_pair_id,\n index_pipeline_result.failures,\n documents,\n )\n\n coordination_status = None\n # Record failures in the database\n if index_pipeline_result.failures:\n with get_session_with_current_tenant() as db_session:\n for failure in index_pipeline_result.failures:\n create_index_attempt_error(\n index_attempt_id,\n cc_pair_id,\n failure,\n db_session,\n )\n # Use database state instead of FileStore for failure checking\n with get_session_with_current_tenant() as db_session:\n coordination_status = IndexingCoordination.get_coordination_status(\n db_session, index_attempt_id\n )\n _check_failure_threshold(\n coordination_status.total_failures,\n coordination_status.total_docs,\n batch_num,\n index_pipeline_result.failures[-1],\n )\n\n # Add telemetry for indexing progress using database coordination status\n # only re-fetch coordination status if necessary\n if coordination_status is None:\n with get_session_with_current_tenant() as db_session:\n coordination_status = IndexingCoordination.get_coordination_status(\n db_session, index_attempt_id\n )\n\n optional_telemetry(\n record_type=RecordType.INDEXING_PROGRESS,\n data={\n \"index_attempt_id\": index_attempt_id,\n \"cc_pair_id\": cc_pair_id,\n \"current_docs_indexed\": coordination_status.total_docs,\n \"current_chunks_indexed\": coordination_status.total_chunks,\n \"source\": index_attempt.connector_credential_pair.connector.source.value,\n \"completed_batches\": coordination_status.completed_batches,\n \"total_batches\": coordination_status.total_batches,\n },\n tenant_id=tenant_id,\n )\n # Clean up this batch after successful processing\n storage.delete_batch_by_num(batch_num)\n\n # FIX: Explicitly clear document batch from memory and force garbage collection\n # This helps prevent memory accumulation across multiple batches\n # NOTE: Thread-local event loops in embedding threads are cleaned up automatically\n # via the _cleanup_thread_local decorator in search_nlp_models.py\n del documents\n gc.collect()\n\n # FIX: Log final memory usage to track problematic tenants/CC pairs\n emit_process_memory(\n os.getpid(),\n \"docprocessing\",\n {\n \"phase\": \"after_processing\",\n \"tenant_id\": tenant_id,\n \"cc_pair_id\": cc_pair_id,\n \"index_attempt_id\": index_attempt_id,\n \"batch_num\": batch_num,\n \"chunks_processed\": index_pipeline_result.total_chunks,\n },\n )\n\n elapsed_time = time.monotonic() - start_time\n task_logger.info(\n f\"Completed document batch processing: \"\n f\"index_attempt={index_attempt_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={index_attempt.search_settings.id} \"\n f\"batch_num={batch_num} \"\n f\"docs={len(index_pipeline_result.failures) + index_pipeline_result.total_docs} \"\n f\"chunks={index_pipeline_result.total_chunks} \"\n f\"failures={len(index_pipeline_result.failures)} \"\n f\"elapsed={elapsed_time:.2f}s\"\n )\n\n except Exception:\n task_logger.exception(\n f\"Document batch processing failed: batch_num={batch_num} attempt={index_attempt_id} \"\n )\n\n raise\n finally:\n if per_batch_lock and per_batch_lock.owned():\n per_batch_lock.release()\n" }, { "path": "backend/onyx/background/celery/tasks/docprocessing/utils.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom redis import Redis\nfrom redis.exceptions import LockError\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.index_attempt import get_last_attempt_for_cc_pair\nfrom onyx.db.index_attempt import get_recent_attempts_for_cc_pair\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import SearchSettings\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_pool import redis_lock_dump\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nNUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE = 5\n\n\nclass IndexingCallbackBase(IndexingHeartbeatInterface):\n PARENT_CHECK_INTERVAL = 60\n\n def __init__(\n self,\n parent_pid: int,\n redis_connector: RedisConnector,\n redis_lock: RedisLock,\n redis_client: Redis,\n timeout_seconds: int | None = None,\n ):\n super().__init__()\n self.parent_pid = parent_pid\n self.redis_connector: RedisConnector = redis_connector\n self.redis_lock: RedisLock = redis_lock\n self.redis_client = redis_client\n self.started: datetime = datetime.now(timezone.utc)\n self.redis_lock.reacquire()\n\n self.last_tag: str = f\"{self.__class__.__name__}.__init__\"\n self.last_lock_reacquire: datetime = datetime.now(timezone.utc)\n self.last_lock_monotonic = time.monotonic()\n\n self.last_parent_check = time.monotonic()\n self.start_monotonic = time.monotonic()\n self.timeout_seconds = timeout_seconds\n\n def should_stop(self) -> bool:\n # Check if the associated indexing attempt has been cancelled\n # TODO: Pass index_attempt_id to the callback and check cancellation using the db\n if bool(self.redis_connector.stop.fenced):\n return True\n\n # Check if the task has exceeded its timeout\n # NOTE: Celery's soft_time_limit does not work with thread pools,\n # so we must enforce timeouts internally.\n if self.timeout_seconds is not None:\n elapsed = time.monotonic() - self.start_monotonic\n if elapsed > self.timeout_seconds:\n logger.warning(\n f\"IndexingCallback Docprocessing - task timeout exceeded: \"\n f\"elapsed={elapsed:.0f}s timeout={self.timeout_seconds}s \"\n f\"cc_pair={self.redis_connector.cc_pair_id}\"\n )\n return True\n\n return False\n\n def progress(self, tag: str, amount: int) -> None: # noqa: ARG002\n \"\"\"Amount isn't used yet.\"\"\"\n\n # rkuo: this shouldn't be necessary yet because we spawn the process this runs inside\n # with daemon=True. It seems likely some indexing tasks will need to spawn other processes\n # eventually, which daemon=True prevents, so leave this code in until we're ready to test it.\n\n # if self.parent_pid:\n # # check if the parent pid is alive so we aren't running as a zombie\n # now = time.monotonic()\n # if now - self.last_parent_check > IndexingCallback.PARENT_CHECK_INTERVAL:\n # try:\n # # this is unintuitive, but it checks if the parent pid is still running\n # os.kill(self.parent_pid, 0)\n # except Exception:\n # logger.exception(\"IndexingCallback - parent pid check exceptioned\")\n # raise\n # self.last_parent_check = now\n\n try:\n current_time = time.monotonic()\n if current_time - self.last_lock_monotonic >= (\n CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4\n ):\n self.redis_lock.reacquire()\n self.last_lock_reacquire = datetime.now(timezone.utc)\n self.last_lock_monotonic = time.monotonic()\n\n self.last_tag = tag\n except LockError:\n logger.exception(\n f\"{self.__class__.__name__} - lock.reacquire exceptioned: \"\n f\"lock_timeout={self.redis_lock.timeout} \"\n f\"start={self.started} \"\n f\"last_tag={self.last_tag} \"\n f\"last_reacquired={self.last_lock_reacquire} \"\n f\"now={datetime.now(timezone.utc)}\"\n )\n\n redis_lock_dump(self.redis_lock, self.redis_client)\n raise\n\n\n# NOTE: we're in the process of removing all fences from indexing; this will\n# eventually no longer be used. For now, it is used only for connector pausing.\nclass IndexingCallback(IndexingHeartbeatInterface):\n def __init__(\n self,\n redis_connector: RedisConnector,\n ):\n self.redis_connector = redis_connector\n\n def should_stop(self) -> bool:\n # Check if the associated indexing attempt has been cancelled\n # TODO: Pass index_attempt_id to the callback and check cancellation using the db\n return bool(self.redis_connector.stop.fenced)\n\n # included to satisfy old interface\n def progress(self, tag: str, amount: int) -> None:\n pass\n\n\n# NOTE: The validate_indexing_fence and validate_indexing_fences functions have been removed\n# as they are no longer needed with database-based coordination. The new validation is\n# handled by validate_active_indexing_attempts in the main indexing tasks module.\n\n\ndef is_in_repeated_error_state(\n cc_pair: ConnectorCredentialPair, search_settings_id: int, db_session: Session\n) -> bool:\n \"\"\"Checks if the cc pair / search setting combination is in a repeated error state.\"\"\"\n # if the connector doesn't have a refresh_freq, a single failed attempt is enough\n number_of_failed_attempts_in_a_row_needed = (\n NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE\n if cc_pair.connector.refresh_freq is not None\n else 1\n )\n\n most_recent_index_attempts = get_recent_attempts_for_cc_pair(\n cc_pair_id=cc_pair.id,\n search_settings_id=search_settings_id,\n limit=number_of_failed_attempts_in_a_row_needed,\n db_session=db_session,\n )\n return len(\n most_recent_index_attempts\n ) >= number_of_failed_attempts_in_a_row_needed and all(\n attempt.status == IndexingStatus.FAILED\n for attempt in most_recent_index_attempts\n )\n\n\ndef should_index(\n cc_pair: ConnectorCredentialPair,\n search_settings_instance: SearchSettings,\n secondary_index_building: bool,\n db_session: Session,\n) -> bool:\n \"\"\"Checks various global settings and past indexing attempts to determine if\n we should try to start indexing the cc pair / search setting combination.\n\n Note that tactical checks such as preventing overlap with a currently running task\n are not handled here.\n\n Return True if we should try to index, False if not.\n \"\"\"\n connector = cc_pair.connector\n last_index_attempt = get_last_attempt_for_cc_pair(\n cc_pair_id=cc_pair.id,\n search_settings_id=search_settings_instance.id,\n db_session=db_session,\n )\n all_recent_errored = is_in_repeated_error_state(\n cc_pair=cc_pair,\n search_settings_id=search_settings_instance.id,\n db_session=db_session,\n )\n\n # uncomment for debugging\n # task_logger.debug(\n # f\"_should_index: \"\n # f\"cc_pair={cc_pair.id} \"\n # f\"connector={cc_pair.connector_id} \"\n # f\"refresh_freq={connector.refresh_freq}\"\n # )\n\n # don't kick off indexing for `NOT_APPLICABLE` sources\n if connector.source == DocumentSource.NOT_APPLICABLE:\n # print(f\"Not indexing cc_pair={cc_pair.id}: NOT_APPLICABLE source\")\n return False\n\n # User can still manually create single indexing attempts via the UI for the\n # currently in use index\n if DISABLE_INDEX_UPDATE_ON_SWAP:\n if (\n search_settings_instance.status == IndexModelStatus.PRESENT\n and secondary_index_building\n ):\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: DISABLE_INDEX_UPDATE_ON_SWAP is True and secondary index building\"\n # )\n return False\n\n # When switching over models, always index at least once\n if search_settings_instance.status == IndexModelStatus.FUTURE:\n if last_index_attempt:\n # No new index if the last index attempt succeeded\n # Once is enough. The model will never be able to swap otherwise.\n if last_index_attempt.status == IndexingStatus.SUCCESS:\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: FUTURE model with successful last index attempt={last_index.id}\"\n # )\n return False\n\n # No new index if the last index attempt is waiting to start\n if last_index_attempt.status == IndexingStatus.NOT_STARTED:\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: FUTURE model with NOT_STARTED last index attempt={last_index.id}\"\n # )\n return False\n\n # No new index if the last index attempt is running\n if last_index_attempt.status == IndexingStatus.IN_PROGRESS:\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: FUTURE model with IN_PROGRESS last index attempt={last_index.id}\"\n # )\n return False\n else:\n if (\n connector.id == 0 or connector.source == DocumentSource.INGESTION_API\n ): # Ingestion API\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: FUTURE model with Ingestion API source\"\n # )\n return False\n return True\n\n # If the connector is paused or is the ingestion API, don't index\n # NOTE: during an embedding model switch over, the following logic\n # is bypassed by the above check for a future model\n if (\n not cc_pair.status.is_active()\n or connector.id == 0\n or connector.source == DocumentSource.INGESTION_API\n ):\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: Connector is paused or is Ingestion API\"\n # )\n return False\n\n if search_settings_instance.status.is_current():\n if cc_pair.indexing_trigger is not None:\n # if a manual indexing trigger is on the cc pair, honor it for live search settings\n return True\n\n # if no attempt has ever occurred, we should index regardless of refresh_freq\n if not last_index_attempt:\n return True\n\n if connector.refresh_freq is None:\n # print(f\"Not indexing cc_pair={cc_pair.id}: refresh_freq is None\")\n return False\n\n # if in the \"initial\" phase, we should always try and kick-off indexing\n # as soon as possible if there is no ongoing attempt. In other words,\n # no delay UNLESS we're repeatedly failing to index.\n if (\n cc_pair.status == ConnectorCredentialPairStatus.INITIAL_INDEXING\n and not all_recent_errored\n ):\n return True\n\n current_db_time = get_db_current_time(db_session)\n time_since_index = current_db_time - last_index_attempt.time_updated\n if time_since_index.total_seconds() < connector.refresh_freq:\n # print(\n # f\"Not indexing cc_pair={cc_pair.id}: Last index attempt={last_index_attempt.id} \"\n # f\"too recent ({time_since_index.total_seconds()}s < {connector.refresh_freq}s)\"\n # )\n return False\n\n return True\n" }, { "path": "backend/onyx/background/celery/tasks/evals/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/evals/tasks.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nfrom celery import shared_task\nfrom celery import Task\n\nfrom onyx.configs.app_configs import BRAINTRUST_API_KEY\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES\nfrom onyx.configs.app_configs import SCHEDULED_EVAL_PERMISSIONS_EMAIL\nfrom onyx.configs.app_configs import SCHEDULED_EVAL_PROJECT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.evals.eval import run_eval\nfrom onyx.evals.models import EvalConfigurationOptions\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n@shared_task(\n name=OnyxCeleryTask.EVAL_RUN_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n trail=False,\n)\ndef eval_run_task(\n self: Task, # noqa: ARG001\n *,\n configuration_dict: dict[str, Any],\n) -> None:\n \"\"\"Background task to run an evaluation with the given configuration\"\"\"\n try:\n configuration = EvalConfigurationOptions.model_validate(configuration_dict)\n run_eval(configuration, remote_dataset_name=configuration.dataset_name)\n logger.info(\"Successfully completed eval run task\")\n\n except Exception:\n logger.error(\"Failed to run eval task\")\n raise\n\n\n@shared_task(\n name=OnyxCeleryTask.SCHEDULED_EVAL_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT * 5, # Allow more time for multiple datasets\n bind=True,\n trail=False,\n)\ndef scheduled_eval_task(self: Task, **kwargs: Any) -> None: # noqa: ARG001\n \"\"\"\n Scheduled task to run evaluations on configured datasets.\n Runs weekly on Sunday at midnight UTC.\n\n Configure via environment variables (with defaults):\n - SCHEDULED_EVAL_DATASET_NAMES: Comma-separated list of Braintrust dataset names\n - SCHEDULED_EVAL_PERMISSIONS_EMAIL: Email for search permissions (default: roshan@onyx.app)\n - SCHEDULED_EVAL_PROJECT: Braintrust project name\n \"\"\"\n if not BRAINTRUST_API_KEY:\n logger.error(\"BRAINTRUST_API_KEY is not configured, cannot run scheduled evals\")\n return\n\n if not SCHEDULED_EVAL_PROJECT:\n logger.error(\n \"SCHEDULED_EVAL_PROJECT is not configured, cannot run scheduled evals\"\n )\n return\n\n if not SCHEDULED_EVAL_DATASET_NAMES:\n logger.info(\"No scheduled eval datasets configured, skipping\")\n return\n\n if not SCHEDULED_EVAL_PERMISSIONS_EMAIL:\n logger.error(\"SCHEDULED_EVAL_PERMISSIONS_EMAIL not configured\")\n return\n\n project_name = SCHEDULED_EVAL_PROJECT\n dataset_names = SCHEDULED_EVAL_DATASET_NAMES\n permissions_email = SCHEDULED_EVAL_PERMISSIONS_EMAIL\n\n # Create a timestamp for the scheduled run\n run_timestamp = datetime.now(timezone.utc).strftime(\"%Y-%m-%d\")\n\n logger.info(\n f\"Starting scheduled eval pipeline for project '{project_name}' with {len(dataset_names)} dataset(s): {dataset_names}\"\n )\n\n pipeline_start = datetime.now(timezone.utc)\n results: list[dict[str, Any]] = []\n\n for dataset_name in dataset_names:\n start_time = datetime.now(timezone.utc)\n error_message: str | None = None\n success = False\n\n # Create informative experiment name for scheduled runs\n experiment_name = f\"{dataset_name} - {run_timestamp}\"\n\n try:\n logger.info(\n f\"Running scheduled eval for dataset: {dataset_name} (project: {project_name})\"\n )\n\n configuration = EvalConfigurationOptions(\n search_permissions_email=permissions_email,\n dataset_name=dataset_name,\n no_send_logs=False,\n braintrust_project=project_name,\n experiment_name=experiment_name,\n )\n\n result = run_eval(\n configuration=configuration,\n remote_dataset_name=dataset_name,\n )\n success = result.success\n logger.info(f\"Completed eval for {dataset_name}: success={success}\")\n\n except Exception as e:\n logger.exception(f\"Failed to run scheduled eval for {dataset_name}\")\n error_message = str(e)\n success = False\n\n end_time = datetime.now(timezone.utc)\n\n results.append(\n {\n \"dataset_name\": dataset_name,\n \"success\": success,\n \"start_time\": start_time,\n \"end_time\": end_time,\n \"error_message\": error_message,\n }\n )\n\n pipeline_end = datetime.now(timezone.utc)\n total_duration = (pipeline_end - pipeline_start).total_seconds()\n\n passed_count = sum(1 for r in results if r[\"success\"])\n logger.info(\n f\"Scheduled eval pipeline completed: {passed_count}/{len(results)} passed in {total_duration:.1f}s\"\n )\n" }, { "path": "backend/onyx/background/celery/tasks/hierarchyfetching/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/hierarchyfetching/tasks.py", "content": "\"\"\"Celery tasks for hierarchy fetching.\n\nThis module provides tasks for fetching hierarchy node information from connectors.\nHierarchy nodes represent structural elements like folders, spaces, and pages that\ncan be used to filter search results.\n\nThe hierarchy fetching pipeline runs once per day per connector and fetches\nstructural information from the connector source.\n\"\"\"\n\nimport time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom uuid import uuid4\n\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.connectors.factory import ConnectorMissingException\nfrom onyx.connectors.factory import identify_connector_class\nfrom onyx.connectors.factory import instantiate_connector\nfrom onyx.connectors.interfaces import HierarchyConnector\nfrom onyx.connectors.models import HierarchyNode as PydanticHierarchyNode\nfrom onyx.db.connector import mark_cc_pair_as_hierarchy_fetched\nfrom onyx.db.connector_credential_pair import (\n fetch_indexable_standard_connector_credential_pair_ids,\n)\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries\nfrom onyx.db.hierarchy import upsert_hierarchy_nodes_batch\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.redis.redis_hierarchy import cache_hierarchy_nodes_batch\nfrom onyx.redis.redis_hierarchy import ensure_source_node_exists\nfrom onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Hierarchy fetching runs once per day (24 hours in seconds)\nHIERARCHY_FETCH_INTERVAL_SECONDS = 24 * 60 * 60\n\n\ndef _connector_supports_hierarchy_fetching(\n cc_pair: ConnectorCredentialPair,\n) -> bool:\n \"\"\"Return True only for connectors whose class implements HierarchyConnector.\"\"\"\n try:\n connector_class = identify_connector_class(\n cc_pair.connector.source,\n )\n except ConnectorMissingException as e:\n task_logger.warning(\n \"Skipping hierarchy fetching enqueue for source=%s input_type=%s: %s\",\n cc_pair.connector.source,\n cc_pair.connector.input_type,\n str(e),\n )\n return False\n\n return issubclass(connector_class, HierarchyConnector)\n\n\ndef _is_hierarchy_fetching_due(cc_pair: ConnectorCredentialPair) -> bool:\n \"\"\"Returns boolean indicating if hierarchy fetching is due for this connector.\n\n Hierarchy fetching should run once per day for active connectors.\n \"\"\"\n # Skip if not active\n if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:\n return False\n\n # Skip if connector has never successfully indexed\n if not cc_pair.last_successful_index_time:\n return False\n\n # Check if we've fetched hierarchy recently\n last_fetch = cc_pair.last_time_hierarchy_fetch\n if last_fetch is None:\n # Never fetched before - fetch now\n return True\n\n # Check if enough time has passed since last fetch\n next_fetch_time = last_fetch + timedelta(seconds=HIERARCHY_FETCH_INTERVAL_SECONDS)\n return datetime.now(timezone.utc) >= next_fetch_time\n\n\ndef _try_creating_hierarchy_fetching_task(\n celery_app: Celery,\n cc_pair: ConnectorCredentialPair,\n db_session: Session,\n r: Redis,\n tenant_id: str,\n) -> str | None:\n \"\"\"Try to create a hierarchy fetching task for a connector.\n\n Returns the task ID if created, None otherwise.\n \"\"\"\n LOCK_TIMEOUT = 30\n\n # Serialize task creation attempts\n lock: RedisLock = r.lock(\n DANSWER_REDIS_FUNCTION_LOCK_PREFIX + f\"hierarchy_fetching_{cc_pair.id}\",\n timeout=LOCK_TIMEOUT,\n )\n\n acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)\n if not acquired:\n return None\n\n try:\n # Refresh to get latest state\n db_session.refresh(cc_pair)\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n return None\n\n # Generate task ID\n custom_task_id = f\"hierarchy_fetching_{cc_pair.id}_{uuid4()}\"\n\n # Send the task\n result = celery_app.send_task(\n OnyxCeleryTask.CONNECTOR_HIERARCHY_FETCHING_TASK,\n kwargs=dict(\n cc_pair_id=cc_pair.id,\n tenant_id=tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_HIERARCHY_FETCHING,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.LOW,\n )\n\n if not result:\n raise RuntimeError(\"send_task for hierarchy_fetching_task failed.\")\n\n task_logger.info(\n f\"Created hierarchy fetching task: cc_pair={cc_pair.id} celery_task_id={custom_task_id}\"\n )\n\n return custom_task_id\n\n except Exception:\n task_logger.exception(\n f\"Failed to create hierarchy fetching task: cc_pair={cc_pair.id}\"\n )\n return None\n finally:\n if lock.owned():\n lock.release()\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_HIERARCHY_FETCHING,\n soft_time_limit=300,\n bind=True,\n)\ndef check_for_hierarchy_fetching(self: Task, *, tenant_id: str) -> int | None:\n \"\"\"Check for connectors that need hierarchy fetching and spawn tasks.\n\n This task runs periodically (once per day) and checks all active connectors\n to see if they need hierarchy information fetched.\n \"\"\"\n time_start = time.monotonic()\n task_logger.info(\"check_for_hierarchy_fetching - Starting\")\n\n tasks_created = 0\n locked = False\n redis_client = get_redis_client()\n\n lock_beat: RedisLock = redis_client.lock(\n OnyxRedisLocks.CHECK_HIERARCHY_FETCHING_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # These tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n try:\n locked = True\n\n with get_session_with_current_tenant() as db_session:\n # Get all active connector credential pairs\n cc_pair_ids = fetch_indexable_standard_connector_credential_pair_ids(\n db_session=db_session,\n active_cc_pairs_only=True,\n )\n\n for cc_pair_id in cc_pair_ids:\n lock_beat.reacquire()\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n\n if not cc_pair or not _connector_supports_hierarchy_fetching(cc_pair):\n continue\n\n if not _is_hierarchy_fetching_due(cc_pair):\n continue\n\n task_id = _try_creating_hierarchy_fetching_task(\n celery_app=self.app,\n cc_pair=cc_pair,\n db_session=db_session,\n r=redis_client,\n tenant_id=tenant_id,\n )\n\n if task_id:\n tasks_created += 1\n\n except Exception:\n task_logger.exception(\"check_for_hierarchy_fetching - Unexpected error\")\n finally:\n if locked:\n if lock_beat.owned():\n lock_beat.release()\n else:\n task_logger.error(\n \"check_for_hierarchy_fetching - Lock not owned on completion\"\n )\n\n time_elapsed = time.monotonic() - time_start\n task_logger.info(\n f\"check_for_hierarchy_fetching finished: tasks_created={tasks_created} elapsed={time_elapsed:.2f}s\"\n )\n return tasks_created\n\n\n# Batch size for hierarchy node processing\nHIERARCHY_NODE_BATCH_SIZE = 100\n\n\ndef _run_hierarchy_extraction(\n db_session: Session,\n cc_pair: ConnectorCredentialPair,\n source: DocumentSource,\n tenant_id: str,\n) -> int:\n \"\"\"\n Run the hierarchy extraction for a connector.\n\n Instantiates the connector and calls load_hierarchy() if the connector\n implements HierarchyConnector.\n\n Returns the total number of hierarchy nodes extracted.\n \"\"\"\n connector = cc_pair.connector\n credential = cc_pair.credential\n\n # Instantiate the connector using its configured input type\n runnable_connector = instantiate_connector(\n db_session=db_session,\n source=source,\n input_type=connector.input_type,\n connector_specific_config=connector.connector_specific_config,\n credential=credential,\n )\n\n # Check if the connector supports hierarchy fetching\n if not isinstance(runnable_connector, HierarchyConnector):\n task_logger.debug(\n f\"Connector {source} does not implement HierarchyConnector, skipping\"\n )\n return 0\n\n redis_client = get_redis_client(tenant_id=tenant_id)\n\n # Ensure the SOURCE-type root node exists before processing hierarchy nodes.\n # This is the root of the hierarchy tree - all other nodes for this source\n # should ultimately have this as an ancestor.\n ensure_source_node_exists(redis_client, db_session, source)\n\n # Determine time range: start from last hierarchy fetch, end at now\n last_fetch = cc_pair.last_time_hierarchy_fetch\n start_time = last_fetch.timestamp() if last_fetch else 0\n end_time = datetime.now(timezone.utc).timestamp()\n\n # Check if connector is public - all hierarchy nodes from public connectors\n # should be accessible to all users\n is_connector_public = cc_pair.access_type == AccessType.PUBLIC\n\n total_nodes = 0\n node_batch: list[PydanticHierarchyNode] = []\n\n def _process_batch() -> int:\n \"\"\"Process accumulated hierarchy nodes batch.\"\"\"\n if not node_batch:\n return 0\n\n upserted_nodes = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=node_batch,\n source=source,\n commit=True,\n is_connector_public=is_connector_public,\n )\n\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=[n.id for n in upserted_nodes],\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n commit=True,\n )\n\n # Cache in Redis for fast ancestor resolution\n cache_entries = [\n HierarchyNodeCacheEntry.from_db_model(node) for node in upserted_nodes\n ]\n cache_hierarchy_nodes_batch(\n redis_client=redis_client,\n source=source,\n entries=cache_entries,\n )\n\n count = len(node_batch)\n node_batch.clear()\n return count\n\n # Fetch hierarchy nodes from the connector\n for node in runnable_connector.load_hierarchy(start=start_time, end=end_time):\n node_batch.append(node)\n if len(node_batch) >= HIERARCHY_NODE_BATCH_SIZE:\n total_nodes += _process_batch()\n\n # Process any remaining nodes\n total_nodes += _process_batch()\n\n return total_nodes\n\n\n@shared_task(\n name=OnyxCeleryTask.CONNECTOR_HIERARCHY_FETCHING_TASK,\n soft_time_limit=3600, # 1 hour soft limit\n time_limit=3900, # 1 hour 5 min hard limit\n bind=True,\n)\ndef connector_hierarchy_fetching_task(\n self: Task, # noqa: ARG001\n *,\n cc_pair_id: int,\n tenant_id: str,\n) -> None:\n \"\"\"Fetch hierarchy information from a connector.\n\n This task fetches structural information (folders, spaces, pages, etc.)\n from the connector source and stores it in the database.\n \"\"\"\n task_logger.info(\n f\"connector_hierarchy_fetching_task starting: cc_pair={cc_pair_id} tenant={tenant_id}\"\n )\n\n try:\n with get_session_with_current_tenant() as db_session:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n\n if not cc_pair:\n task_logger.warning(\n f\"CC pair not found for hierarchy fetching: cc_pair={cc_pair_id}\"\n )\n return\n\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n task_logger.info(\n f\"Skipping hierarchy fetching for deleting connector: cc_pair={cc_pair_id}\"\n )\n return\n\n source = cc_pair.connector.source\n total_nodes = _run_hierarchy_extraction(\n db_session=db_session,\n cc_pair=cc_pair,\n source=source,\n tenant_id=tenant_id,\n )\n\n task_logger.info(\n f\"connector_hierarchy_fetching_task: Extracted {total_nodes} hierarchy nodes for cc_pair={cc_pair_id}\"\n )\n\n # Update the last fetch time to prevent re-running until next interval\n mark_cc_pair_as_hierarchy_fetched(db_session, cc_pair_id)\n\n except Exception:\n task_logger.exception(\n f\"connector_hierarchy_fetching_task failed: cc_pair={cc_pair_id}\"\n )\n raise\n\n task_logger.info(\n f\"connector_hierarchy_fetching_task completed: cc_pair={cc_pair_id}\"\n )\n" }, { "path": "backend/onyx/background/celery/tasks/llm_model_update/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/llm_model_update/tasks.py", "content": "from celery import shared_task\nfrom celery import Task\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.configs.app_configs import AUTO_LLM_CONFIG_URL\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.llm.well_known_providers.auto_update_service import (\n sync_llm_models_from_github,\n)\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_AUTO_LLM_UPDATE,\n ignore_result=True,\n soft_time_limit=300, # 5 minute timeout\n trail=False,\n bind=True,\n)\ndef check_for_auto_llm_updates(\n self: Task, # noqa: ARG001\n *,\n tenant_id: str, # noqa: ARG001\n) -> bool | None:\n \"\"\"Periodic task to fetch LLM model updates from GitHub\n and sync them to providers in Auto mode.\n\n This task checks the GitHub-hosted config file and updates all\n providers that have is_auto_mode=True.\n \"\"\"\n if not AUTO_LLM_CONFIG_URL:\n task_logger.debug(\"AUTO_LLM_CONFIG_URL not configured, skipping\")\n return None\n\n try:\n # Sync to database\n with get_session_with_current_tenant() as db_session:\n results = sync_llm_models_from_github(db_session)\n\n if results:\n task_logger.info(f\"Auto mode sync results: {results}\")\n else:\n task_logger.debug(\"No model updates applied\")\n\n except Exception:\n task_logger.exception(\"Error in auto LLM update task\")\n raise\n\n return True\n" }, { "path": "backend/onyx/background/celery/tasks/models.py", "content": "from enum import Enum\n\nfrom pydantic import BaseModel\n\n\nclass DocProcessingContext(BaseModel):\n tenant_id: str\n cc_pair_id: int\n search_settings_id: int\n index_attempt_id: int\n\n\nclass IndexingWatchdogTerminalStatus(str, Enum):\n \"\"\"The different statuses the watchdog can finish with.\n\n TODO: create broader success/failure/abort categories\n \"\"\"\n\n UNDEFINED = \"undefined\"\n\n SUCCEEDED = \"succeeded\"\n\n SPAWN_FAILED = \"spawn_failed\" # connector spawn failed\n SPAWN_NOT_ALIVE = (\n \"spawn_not_alive\" # spawn succeeded but process did not come alive\n )\n\n BLOCKED_BY_DELETION = \"blocked_by_deletion\"\n BLOCKED_BY_STOP_SIGNAL = \"blocked_by_stop_signal\"\n FENCE_NOT_FOUND = \"fence_not_found\" # fence does not exist\n FENCE_READINESS_TIMEOUT = (\n \"fence_readiness_timeout\" # fence exists but wasn't ready within the timeout\n )\n FENCE_MISMATCH = \"fence_mismatch\" # task and fence metadata mismatch\n TASK_ALREADY_RUNNING = \"task_already_running\" # task appears to be running already\n INDEX_ATTEMPT_MISMATCH = (\n \"index_attempt_mismatch\" # expected index attempt metadata not found in db\n )\n\n CONNECTOR_VALIDATION_ERROR = (\n \"connector_validation_error\" # the connector validation failed\n )\n CONNECTOR_EXCEPTIONED = \"connector_exceptioned\" # the connector itself exceptioned\n WATCHDOG_EXCEPTIONED = \"watchdog_exceptioned\" # the watchdog exceptioned\n\n # the watchdog received a termination signal\n TERMINATED_BY_SIGNAL = \"terminated_by_signal\"\n\n # the watchdog terminated the task due to no activity\n TERMINATED_BY_ACTIVITY_TIMEOUT = \"terminated_by_activity_timeout\"\n\n # NOTE: this may actually be the same as SIGKILL, but parsed differently by python\n # consolidate once we know more\n OUT_OF_MEMORY = \"out_of_memory\"\n\n PROCESS_SIGNAL_SIGKILL = \"process_signal_sigkill\"\n\n @property\n def code(self) -> int:\n _ENUM_TO_CODE: dict[IndexingWatchdogTerminalStatus, int] = {\n IndexingWatchdogTerminalStatus.PROCESS_SIGNAL_SIGKILL: -9,\n IndexingWatchdogTerminalStatus.OUT_OF_MEMORY: 137,\n IndexingWatchdogTerminalStatus.CONNECTOR_VALIDATION_ERROR: 247,\n IndexingWatchdogTerminalStatus.BLOCKED_BY_DELETION: 248,\n IndexingWatchdogTerminalStatus.BLOCKED_BY_STOP_SIGNAL: 249,\n IndexingWatchdogTerminalStatus.FENCE_NOT_FOUND: 250,\n IndexingWatchdogTerminalStatus.FENCE_READINESS_TIMEOUT: 251,\n IndexingWatchdogTerminalStatus.FENCE_MISMATCH: 252,\n IndexingWatchdogTerminalStatus.TASK_ALREADY_RUNNING: 253,\n IndexingWatchdogTerminalStatus.INDEX_ATTEMPT_MISMATCH: 254,\n IndexingWatchdogTerminalStatus.CONNECTOR_EXCEPTIONED: 255,\n }\n\n return _ENUM_TO_CODE[self]\n\n @classmethod\n def from_code(cls, code: int) -> \"IndexingWatchdogTerminalStatus\":\n _CODE_TO_ENUM: dict[int, IndexingWatchdogTerminalStatus] = {\n -9: IndexingWatchdogTerminalStatus.PROCESS_SIGNAL_SIGKILL,\n 137: IndexingWatchdogTerminalStatus.OUT_OF_MEMORY,\n 247: IndexingWatchdogTerminalStatus.CONNECTOR_VALIDATION_ERROR,\n 248: IndexingWatchdogTerminalStatus.BLOCKED_BY_DELETION,\n 249: IndexingWatchdogTerminalStatus.BLOCKED_BY_STOP_SIGNAL,\n 250: IndexingWatchdogTerminalStatus.FENCE_NOT_FOUND,\n 251: IndexingWatchdogTerminalStatus.FENCE_READINESS_TIMEOUT,\n 252: IndexingWatchdogTerminalStatus.FENCE_MISMATCH,\n 253: IndexingWatchdogTerminalStatus.TASK_ALREADY_RUNNING,\n 254: IndexingWatchdogTerminalStatus.INDEX_ATTEMPT_MISMATCH,\n 255: IndexingWatchdogTerminalStatus.CONNECTOR_EXCEPTIONED,\n }\n\n if code in _CODE_TO_ENUM:\n return _CODE_TO_ENUM[code]\n\n return IndexingWatchdogTerminalStatus.UNDEFINED\n\n\nclass SimpleJobResult:\n \"\"\"The data we want to have when the watchdog finishes\"\"\"\n\n def __init__(self) -> None:\n self.status = IndexingWatchdogTerminalStatus.UNDEFINED\n self.connector_source = None\n self.exit_code = None\n self.exception_str = None\n\n status: IndexingWatchdogTerminalStatus\n connector_source: str | None\n exit_code: int | None\n exception_str: str | None\n" }, { "path": "backend/onyx/background/celery/tasks/monitoring/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/monitoring/tasks.py", "content": "import json\nimport time\nfrom datetime import timedelta\nfrom itertools import islice\nfrom typing import Any\nfrom typing import cast\nfrom typing import Literal\n\nimport psutil\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom pydantic import BaseModel\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy import select\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_queue_length\nfrom onyx.background.celery.celery_redis import celery_get_unacked_task_ids\nfrom onyx.background.celery.memory_monitoring import emit_process_memory\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import ONYX_CLOUD_TENANT_ID\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import SyncRecord\nfrom onyx.db.models import UserGroup\nfrom onyx.db.search_settings import get_active_search_settings_list\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import redis_lock_dump\nfrom onyx.utils.logger import is_running_in_container\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n_MONITORING_SOFT_TIME_LIMIT = 60 * 5 # 5 minutes\n_MONITORING_TIME_LIMIT = _MONITORING_SOFT_TIME_LIMIT + 60 # 6 minutes\n\n_CONNECTOR_INDEX_ATTEMPT_START_LATENCY_KEY_FMT = (\n \"monitoring_connector_index_attempt_start_latency:{cc_pair_id}:{index_attempt_id}\"\n)\n\n_CONNECTOR_INDEX_ATTEMPT_RUN_SUCCESS_KEY_FMT = (\n \"monitoring_connector_index_attempt_run_success:{cc_pair_id}:{index_attempt_id}\"\n)\n\n_FINAL_METRIC_KEY_FMT = \"sync_final_metrics:{sync_type}:{entity_id}:{sync_record_id}\"\n\n_SYNC_START_LATENCY_KEY_FMT = (\n \"sync_start_latency:{sync_type}:{entity_id}:{sync_record_id}\"\n)\n\n_CONNECTOR_START_TIME_KEY_FMT = \"connector_start_time:{cc_pair_id}:{index_attempt_id}\"\n_CONNECTOR_END_TIME_KEY_FMT = \"connector_end_time:{cc_pair_id}:{index_attempt_id}\"\n_SYNC_START_TIME_KEY_FMT = \"sync_start_time:{sync_type}:{entity_id}:{sync_record_id}\"\n_SYNC_END_TIME_KEY_FMT = \"sync_end_time:{sync_type}:{entity_id}:{sync_record_id}\"\n\n\ndef _mark_metric_as_emitted(redis_std: Redis, key: str) -> None:\n \"\"\"Mark a metric as having been emitted by setting a Redis key with expiration\"\"\"\n redis_std.set(key, \"1\", ex=24 * 60 * 60) # Expire after 1 day\n\n\ndef _has_metric_been_emitted(redis_std: Redis, key: str) -> bool:\n \"\"\"Check if a metric has been emitted by checking for existence of Redis key\"\"\"\n return bool(redis_std.exists(key))\n\n\nclass Metric(BaseModel):\n key: (\n str | None\n ) # only required if we need to store that we have emitted this metric\n name: str\n value: Any\n tags: dict[str, str]\n\n def log(self) -> None:\n \"\"\"Log the metric in a standardized format\"\"\"\n data = {\n \"metric\": self.name,\n \"value\": self.value,\n \"tags\": self.tags,\n }\n task_logger.info(json.dumps(data))\n\n def emit(self, tenant_id: str) -> None:\n # Convert value to appropriate type based on the input value\n bool_value = None\n float_value = None\n int_value = None\n string_value = None\n # NOTE: have to do bool first, since `isinstance(True, int)` is true\n # e.g. bool is a subclass of int\n if isinstance(self.value, bool):\n bool_value = self.value\n elif isinstance(self.value, int):\n int_value = self.value\n elif isinstance(self.value, float):\n float_value = self.value\n elif isinstance(self.value, str):\n string_value = self.value\n else:\n task_logger.error(\n f\"Invalid metric value type: {type(self.value)} ({self.value}) for metric {self.name}.\"\n )\n return\n\n # don't send None values over the wire\n data = {\n k: v\n for k, v in {\n \"metric_name\": self.name,\n \"float_value\": float_value,\n \"int_value\": int_value,\n \"string_value\": string_value,\n \"bool_value\": bool_value,\n \"tags\": self.tags,\n }.items()\n if v is not None\n }\n task_logger.info(f\"Emitting metric: {data}\")\n optional_telemetry(\n record_type=RecordType.METRIC,\n data=data,\n tenant_id=tenant_id,\n )\n\n\ndef _collect_queue_metrics(redis_celery: Redis) -> list[Metric]:\n \"\"\"Collect metrics about queue lengths for different Celery queues\"\"\"\n metrics = []\n queue_mappings = {\n \"celery_queue_length\": OnyxCeleryQueues.PRIMARY,\n \"docprocessing_queue_length\": OnyxCeleryQueues.DOCPROCESSING,\n \"docfetching_queue_length\": OnyxCeleryQueues.CONNECTOR_DOC_FETCHING,\n \"sync_queue_length\": OnyxCeleryQueues.VESPA_METADATA_SYNC,\n \"deletion_queue_length\": OnyxCeleryQueues.CONNECTOR_DELETION,\n \"pruning_queue_length\": OnyxCeleryQueues.CONNECTOR_PRUNING,\n \"permissions_sync_queue_length\": OnyxCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC,\n \"external_group_sync_queue_length\": OnyxCeleryQueues.CONNECTOR_EXTERNAL_GROUP_SYNC,\n \"permissions_upsert_queue_length\": OnyxCeleryQueues.DOC_PERMISSIONS_UPSERT,\n \"hierarchy_fetching_queue_length\": OnyxCeleryQueues.CONNECTOR_HIERARCHY_FETCHING,\n \"llm_model_update_queue_length\": OnyxCeleryQueues.LLM_MODEL_UPDATE,\n \"checkpoint_cleanup_queue_length\": OnyxCeleryQueues.CHECKPOINT_CLEANUP,\n \"index_attempt_cleanup_queue_length\": OnyxCeleryQueues.INDEX_ATTEMPT_CLEANUP,\n \"csv_generation_queue_length\": OnyxCeleryQueues.CSV_GENERATION,\n \"user_file_processing_queue_length\": OnyxCeleryQueues.USER_FILE_PROCESSING,\n \"user_file_project_sync_queue_length\": OnyxCeleryQueues.USER_FILE_PROJECT_SYNC,\n \"user_file_delete_queue_length\": OnyxCeleryQueues.USER_FILE_DELETE,\n \"monitoring_queue_length\": OnyxCeleryQueues.MONITORING,\n \"sandbox_queue_length\": OnyxCeleryQueues.SANDBOX,\n \"opensearch_migration_queue_length\": OnyxCeleryQueues.OPENSEARCH_MIGRATION,\n }\n\n for name, queue in queue_mappings.items():\n metrics.append(\n Metric(\n key=None,\n name=name,\n value=celery_get_queue_length(queue, redis_celery),\n tags={\"queue\": name},\n )\n )\n\n return metrics\n\n\ndef _build_connector_start_latency_metric(\n cc_pair: ConnectorCredentialPair,\n recent_attempt: IndexAttempt,\n second_most_recent_attempt: IndexAttempt | None,\n redis_std: Redis,\n) -> Metric | None:\n if not recent_attempt.time_started:\n return None\n\n # check if we already emitted a metric for this index attempt\n metric_key = _CONNECTOR_INDEX_ATTEMPT_START_LATENCY_KEY_FMT.format(\n cc_pair_id=cc_pair.id,\n index_attempt_id=recent_attempt.id,\n )\n if _has_metric_been_emitted(redis_std, metric_key):\n task_logger.info(\n f\"Skipping metric for connector {cc_pair.connector.id} \"\n f\"index attempt {recent_attempt.id} because it has already been \"\n \"emitted\"\n )\n return None\n\n # Connector start latency\n # first run case - we should start as soon as it's created\n if not second_most_recent_attempt:\n desired_start_time = cc_pair.connector.time_created\n else:\n if not cc_pair.connector.refresh_freq:\n task_logger.debug(\n \"Connector has no refresh_freq and this is a non-initial index attempt. \"\n \"Assuming user manually triggered indexing, so we'll skip start latency metric.\"\n )\n return None\n\n desired_start_time = second_most_recent_attempt.time_updated + timedelta(\n seconds=cc_pair.connector.refresh_freq\n )\n\n start_latency = (recent_attempt.time_started - desired_start_time).total_seconds()\n\n task_logger.info(\n f\"Start latency for index attempt {recent_attempt.id}: {start_latency:.2f}s \"\n f\"(desired: {desired_start_time}, actual: {recent_attempt.time_started})\"\n )\n\n job_id = build_job_id(\"connector\", str(cc_pair.id), str(recent_attempt.id))\n\n return Metric(\n key=metric_key,\n name=\"connector_start_latency\",\n value=start_latency,\n tags={\n \"job_id\": job_id,\n \"connector_id\": str(cc_pair.connector.id),\n \"source\": str(cc_pair.connector.source),\n },\n )\n\n\ndef _build_connector_final_metrics(\n cc_pair: ConnectorCredentialPair,\n recent_attempts: list[IndexAttempt],\n redis_std: Redis,\n) -> list[Metric]:\n \"\"\"\n Final metrics for connector index attempts:\n - Boolean success/fail metric\n - If success, emit:\n * duration (seconds)\n * doc_count\n \"\"\"\n metrics = []\n for attempt in recent_attempts:\n metric_key = _CONNECTOR_INDEX_ATTEMPT_RUN_SUCCESS_KEY_FMT.format(\n cc_pair_id=cc_pair.id,\n index_attempt_id=attempt.id,\n )\n if _has_metric_been_emitted(redis_std, metric_key):\n task_logger.info(\n f\"Skipping final metrics for connector {cc_pair.connector.id} index attempt {attempt.id}, already emitted.\"\n )\n continue\n\n # We only emit final metrics if the attempt is in a terminal state\n if attempt.status not in [\n IndexingStatus.SUCCESS,\n IndexingStatus.FAILED,\n IndexingStatus.CANCELED,\n ]:\n # Not finished; skip\n continue\n\n job_id = build_job_id(\"connector\", str(cc_pair.id), str(attempt.id))\n success = attempt.status == IndexingStatus.SUCCESS\n metrics.append(\n Metric(\n key=metric_key, # We'll mark the same key for any final metrics\n name=\"connector_run_succeeded\",\n value=success,\n tags={\n \"job_id\": job_id,\n \"connector_id\": str(cc_pair.connector.id),\n \"source\": str(cc_pair.connector.source),\n \"status\": attempt.status.value,\n },\n )\n )\n\n if success:\n # Make sure we have valid time_started\n if attempt.time_started and attempt.time_updated:\n duration_seconds = (\n attempt.time_updated - attempt.time_started\n ).total_seconds()\n metrics.append(\n Metric(\n key=None, # No need for a new key, or you can reuse the same if you prefer\n name=\"connector_index_duration_seconds\",\n value=duration_seconds,\n tags={\n \"job_id\": job_id,\n \"connector_id\": str(cc_pair.connector.id),\n \"source\": str(cc_pair.connector.source),\n },\n )\n )\n else:\n task_logger.error(\n f\"Index attempt {attempt.id} succeeded but has missing time \"\n f\"(time_started={attempt.time_started}, time_updated={attempt.time_updated}).\"\n )\n\n # For doc counts, choose whichever field is more relevant\n doc_count = attempt.total_docs_indexed or 0\n metrics.append(\n Metric(\n key=None,\n name=\"connector_index_doc_count\",\n value=doc_count,\n tags={\n \"job_id\": job_id,\n \"connector_id\": str(cc_pair.connector.id),\n \"source\": str(cc_pair.connector.source),\n },\n )\n )\n\n return metrics\n\n\ndef _collect_connector_metrics(db_session: Session, redis_std: Redis) -> list[Metric]:\n \"\"\"Collect metrics about connector runs from the past hour\"\"\"\n one_hour_ago = get_db_current_time(db_session) - timedelta(hours=1)\n\n # Get all connector credential pairs\n cc_pairs = db_session.scalars(select(ConnectorCredentialPair)).all()\n # Might be more than one search setting, or just one\n active_search_settings_list = get_active_search_settings_list(db_session)\n\n metrics = []\n\n # If you want to process each cc_pair against each search setting:\n for cc_pair in cc_pairs:\n for search_settings in active_search_settings_list:\n recent_attempts = (\n db_session.query(IndexAttempt)\n .filter(\n IndexAttempt.connector_credential_pair_id == cc_pair.id,\n IndexAttempt.search_settings_id == search_settings.id,\n )\n .order_by(IndexAttempt.time_created.desc())\n .limit(2)\n .all()\n )\n\n if not recent_attempts:\n continue\n\n most_recent_attempt = recent_attempts[0]\n second_most_recent_attempt = (\n recent_attempts[1] if len(recent_attempts) > 1 else None\n )\n\n if one_hour_ago > most_recent_attempt.time_created:\n continue\n\n # Build a job_id for correlation\n job_id = build_job_id(\n \"connector\", str(cc_pair.id), str(most_recent_attempt.id)\n )\n\n # Add raw start time metric if available\n if most_recent_attempt.time_started:\n start_time_key = _CONNECTOR_START_TIME_KEY_FMT.format(\n cc_pair_id=cc_pair.id,\n index_attempt_id=most_recent_attempt.id,\n )\n metrics.append(\n Metric(\n key=start_time_key,\n name=\"connector_start_time\",\n value=most_recent_attempt.time_started.timestamp(),\n tags={\n \"job_id\": job_id,\n \"connector_id\": str(cc_pair.connector.id),\n \"source\": str(cc_pair.connector.source),\n },\n )\n )\n\n # Add raw end time metric if available and in terminal state\n if (\n most_recent_attempt.status.is_terminal()\n and most_recent_attempt.time_updated\n ):\n end_time_key = _CONNECTOR_END_TIME_KEY_FMT.format(\n cc_pair_id=cc_pair.id,\n index_attempt_id=most_recent_attempt.id,\n )\n metrics.append(\n Metric(\n key=end_time_key,\n name=\"connector_end_time\",\n value=most_recent_attempt.time_updated.timestamp(),\n tags={\n \"job_id\": job_id,\n \"connector_id\": str(cc_pair.connector.id),\n \"source\": str(cc_pair.connector.source),\n },\n )\n )\n\n # Connector start latency\n start_latency_metric = _build_connector_start_latency_metric(\n cc_pair, most_recent_attempt, second_most_recent_attempt, redis_std\n )\n\n if start_latency_metric:\n metrics.append(start_latency_metric)\n\n # Connector run success/failure\n final_metrics = _build_connector_final_metrics(\n cc_pair, recent_attempts, redis_std\n )\n metrics.extend(final_metrics)\n\n return metrics\n\n\ndef _collect_sync_metrics(db_session: Session, redis_std: Redis) -> list[Metric]:\n \"\"\"\n Collect metrics for document set and group syncing:\n - Success/failure status\n - Start latency (for doc sets / user groups)\n - Duration & doc count (only if success)\n - Throughput (docs/min) (only if success)\n - Raw start/end times for each sync\n \"\"\"\n\n one_hour_ago = get_db_current_time(db_session) - timedelta(hours=1)\n\n # Get all sync records that ended in the last hour\n recent_sync_records = db_session.scalars(\n select(SyncRecord)\n .where(SyncRecord.sync_end_time.isnot(None))\n .where(SyncRecord.sync_end_time >= one_hour_ago)\n .order_by(SyncRecord.sync_end_time.desc())\n ).all()\n\n task_logger.info(\n f\"Collecting sync metrics for {len(recent_sync_records)} sync records\"\n )\n\n metrics = []\n\n for sync_record in recent_sync_records:\n # Build a job_id for correlation\n job_id = build_job_id(\"sync_record\", str(sync_record.id))\n\n # Add raw start time metric\n start_time_key = _SYNC_START_TIME_KEY_FMT.format(\n sync_type=sync_record.sync_type,\n entity_id=sync_record.entity_id,\n sync_record_id=sync_record.id,\n )\n metrics.append(\n Metric(\n key=start_time_key,\n name=\"sync_start_time\",\n value=sync_record.sync_start_time.timestamp(),\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n },\n )\n )\n\n # Add raw end time metric if available\n if sync_record.sync_end_time:\n end_time_key = _SYNC_END_TIME_KEY_FMT.format(\n sync_type=sync_record.sync_type,\n entity_id=sync_record.entity_id,\n sync_record_id=sync_record.id,\n )\n metrics.append(\n Metric(\n key=end_time_key,\n name=\"sync_end_time\",\n value=sync_record.sync_end_time.timestamp(),\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n },\n )\n )\n\n # Emit a SUCCESS/FAIL boolean metric\n # Use a single Redis key to avoid re-emitting final metrics\n final_metric_key = _FINAL_METRIC_KEY_FMT.format(\n sync_type=sync_record.sync_type,\n entity_id=sync_record.entity_id,\n sync_record_id=sync_record.id,\n )\n if not _has_metric_been_emitted(redis_std, final_metric_key):\n # Evaluate success\n sync_succeeded = sync_record.sync_status == SyncStatus.SUCCESS\n\n metrics.append(\n Metric(\n key=final_metric_key,\n name=\"sync_run_succeeded\",\n value=sync_succeeded,\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n \"status\": str(sync_record.sync_status),\n },\n )\n )\n\n # If successful, emit additional metrics\n if sync_succeeded:\n if sync_record.sync_end_time and sync_record.sync_start_time:\n duration_seconds = (\n sync_record.sync_end_time - sync_record.sync_start_time\n ).total_seconds()\n else:\n task_logger.error(\n f\"Invalid times for sync record {sync_record.id}: \"\n f\"start={sync_record.sync_start_time}, end={sync_record.sync_end_time}\"\n )\n duration_seconds = None\n\n doc_count = sync_record.num_docs_synced or 0\n\n sync_speed = None\n if duration_seconds and duration_seconds > 0:\n duration_mins = duration_seconds / 60.0\n sync_speed = (\n doc_count / duration_mins if duration_mins > 0 else None\n )\n\n # Emit duration, doc count, speed\n if duration_seconds is not None:\n metrics.append(\n Metric(\n key=final_metric_key,\n name=\"sync_duration_seconds\",\n value=duration_seconds,\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n },\n )\n )\n else:\n task_logger.error(\n f\"Invalid sync record {sync_record.id} with no duration\"\n )\n\n metrics.append(\n Metric(\n key=final_metric_key,\n name=\"sync_doc_count\",\n value=doc_count,\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n },\n )\n )\n\n if sync_speed is not None:\n metrics.append(\n Metric(\n key=final_metric_key,\n name=\"sync_speed_docs_per_min\",\n value=sync_speed,\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n },\n )\n )\n else:\n task_logger.error(\n f\"Invalid sync record {sync_record.id} with no duration\"\n )\n\n # Emit start latency\n start_latency_key = _SYNC_START_LATENCY_KEY_FMT.format(\n sync_type=sync_record.sync_type,\n entity_id=sync_record.entity_id,\n sync_record_id=sync_record.id,\n )\n if not _has_metric_been_emitted(redis_std, start_latency_key):\n # Get the entity's last update time based on sync type\n entity: DocumentSet | UserGroup | None = None\n if sync_record.sync_type == SyncType.DOCUMENT_SET:\n entity = db_session.scalar(\n select(DocumentSet).where(DocumentSet.id == sync_record.entity_id)\n )\n elif sync_record.sync_type == SyncType.USER_GROUP:\n entity = db_session.scalar(\n select(UserGroup).where(UserGroup.id == sync_record.entity_id)\n )\n else:\n # Only user groups and document set sync records have\n # an associated entity we can use for latency metrics\n continue\n\n if entity is None:\n task_logger.error(\n f\"Sync record of type {sync_record.sync_type} doesn't have an entity \"\n f\"associated with it (id={sync_record.entity_id}). Skipping start latency metric.\"\n )\n\n # Calculate start latency in seconds:\n # (actual sync start) - (last modified time)\n if (\n entity is not None\n and entity.time_last_modified_by_user\n and sync_record.sync_start_time\n ):\n start_latency = (\n sync_record.sync_start_time - entity.time_last_modified_by_user\n ).total_seconds()\n\n if start_latency < 0:\n task_logger.error(\n f\"Negative start latency for sync record {sync_record.id} \"\n f\"(start={sync_record.sync_start_time}, entity_modified={entity.time_last_modified_by_user})\"\n )\n continue\n\n metrics.append(\n Metric(\n key=start_latency_key,\n name=\"sync_start_latency_seconds\",\n value=start_latency,\n tags={\n \"job_id\": job_id,\n \"sync_type\": str(sync_record.sync_type),\n },\n )\n )\n\n return metrics\n\n\ndef build_job_id(\n job_type: Literal[\"connector\", \"sync_record\"],\n primary_id: str,\n secondary_id: str | None = None,\n) -> str:\n if job_type == \"connector\":\n if secondary_id is None:\n raise ValueError(\n \"secondary_id (attempt_id) is required for connector job_type\"\n )\n return f\"connector:{primary_id}:attempt:{secondary_id}\"\n elif job_type == \"sync_record\":\n return f\"sync_record:{primary_id}\"\n\n\n@shared_task(\n name=OnyxCeleryTask.MONITOR_BACKGROUND_PROCESSES,\n ignore_result=True,\n soft_time_limit=_MONITORING_SOFT_TIME_LIMIT,\n time_limit=_MONITORING_TIME_LIMIT,\n queue=OnyxCeleryQueues.MONITORING,\n bind=True,\n)\ndef monitor_background_processes(self: Task, *, tenant_id: str) -> None:\n \"\"\"Collect and emit metrics about background processes.\n This task runs periodically to gather metrics about:\n - Queue lengths for different Celery queues\n - Connector run metrics (start latency, success rate)\n - Syncing speed metrics\n - Worker status and task counts\n \"\"\"\n if tenant_id is not None:\n CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n task_logger.info(\"Starting background monitoring\")\n r = get_redis_client()\n\n lock_monitoring: RedisLock = r.lock(\n OnyxRedisLocks.MONITOR_BACKGROUND_PROCESSES_LOCK,\n timeout=_MONITORING_SOFT_TIME_LIMIT,\n )\n\n # these tasks should never overlap\n if not lock_monitoring.acquire(blocking=False):\n task_logger.info(\"Skipping monitoring task because it is already running\")\n return None\n\n try:\n redis_std = get_redis_client()\n\n # Collect queue metrics with broker connection\n r_celery = celery_get_broker_client(self.app)\n queue_metrics = _collect_queue_metrics(r_celery)\n\n # Collect remaining metrics (no broker connection needed)\n with get_session_with_current_tenant() as db_session:\n all_metrics: list[Metric] = queue_metrics\n all_metrics.extend(_collect_connector_metrics(db_session, redis_std))\n all_metrics.extend(_collect_sync_metrics(db_session, redis_std))\n\n for metric in all_metrics:\n if metric.key is None or not _has_metric_been_emitted(\n redis_std, metric.key\n ):\n metric.log()\n metric.emit(tenant_id)\n\n if metric.key is not None:\n _mark_metric_as_emitted(redis_std, metric.key)\n\n task_logger.info(\"Successfully collected background metrics\")\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception as e:\n task_logger.exception(\"Error collecting background process metrics\")\n raise e\n finally:\n if lock_monitoring.owned():\n lock_monitoring.release()\n\n task_logger.info(\"Background monitoring task finished\")\n\n\n@shared_task(\n name=OnyxCeleryTask.CLOUD_MONITOR_ALEMBIC,\n)\ndef cloud_check_alembic() -> bool | None:\n \"\"\"A task to verify that all tenants are on the same alembic revision.\n\n This check is expected to fail if a cloud alembic migration is currently running\n across all tenants.\n\n TODO: have the cloud migration script set an activity signal that this check\n uses to know it doesn't make sense to run a check at the present time.\n \"\"\"\n\n # Used as a placeholder if the alembic revision cannot be retrieved\n ALEMBIC_NULL_REVISION = \"000000000000\"\n\n time_start = time.monotonic()\n\n redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)\n\n lock_beat: RedisLock = redis_client.lock(\n OnyxRedisLocks.CLOUD_CHECK_ALEMBIC_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n last_lock_time = time.monotonic()\n\n tenant_to_revision: dict[str, str] = {}\n revision_counts: dict[str, int] = {}\n out_of_date_tenants: dict[str, str] = {}\n top_revision: str = \"\"\n tenant_ids: list[str] | list[None] = []\n\n try:\n # map tenant_id to revision (or ALEMBIC_NULL_REVISION if the query fails)\n tenant_ids = get_all_tenant_ids()\n for tenant_id in tenant_ids:\n current_time = time.monotonic()\n if current_time - last_lock_time >= (CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4):\n lock_beat.reacquire()\n last_lock_time = current_time\n\n if tenant_id is None:\n continue\n\n with get_session_with_shared_schema() as session:\n try:\n result = session.execute(\n text(f'SELECT * FROM \"{tenant_id}\".alembic_version LIMIT 1')\n )\n result_scalar: str | None = result.scalar_one_or_none()\n if result_scalar is None:\n raise ValueError(\"Alembic version should not be None.\")\n\n tenant_to_revision[tenant_id] = result_scalar\n except Exception:\n task_logger.error(f\"Tenant {tenant_id} has no revision!\")\n tenant_to_revision[tenant_id] = ALEMBIC_NULL_REVISION\n\n # get the total count of each revision\n for k, v in tenant_to_revision.items():\n revision_counts[v] = revision_counts.get(v, 0) + 1\n\n # error if any null revision tenants are found\n if ALEMBIC_NULL_REVISION in revision_counts:\n num_null_revisions = revision_counts[ALEMBIC_NULL_REVISION]\n raise ValueError(f\"No revision was found for {num_null_revisions} tenants!\")\n\n # get the revision with the most counts\n sorted_revision_counts = sorted(\n revision_counts.items(), key=lambda item: item[1], reverse=True\n )\n\n if len(sorted_revision_counts) == 0:\n raise ValueError(\n f\"cloud_check_alembic - No revisions found for {len(tenant_ids)} tenant ids!\"\n )\n\n top_revision, _ = sorted_revision_counts[0]\n\n # build a list of out of date tenants\n for k, v in tenant_to_revision.items():\n if v == top_revision:\n continue\n\n out_of_date_tenants[k] = v\n\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n raise\n except Exception:\n task_logger.exception(\"Unexpected exception during cloud alembic check\")\n raise\n finally:\n if lock_beat.owned():\n lock_beat.release()\n else:\n task_logger.error(\"cloud_check_alembic - Lock not owned on completion\")\n redis_lock_dump(lock_beat, redis_client)\n\n if len(out_of_date_tenants) > 0:\n task_logger.error(\n f\"Found out of date tenants: \"\n f\"num_out_of_date_tenants={len(out_of_date_tenants)} \"\n f\"num_tenants={len(tenant_ids)} \"\n f\"revision={top_revision}\"\n )\n\n num_to_log = min(5, len(out_of_date_tenants))\n task_logger.info(\n f\"Logging {num_to_log}/{len(out_of_date_tenants)} out of date tenants.\"\n )\n for k, v in islice(out_of_date_tenants.items(), 5):\n task_logger.info(f\"Out of date tenant: tenant={k} revision={v}\")\n else:\n task_logger.info(\n f\"All tenants are up to date: num_tenants={len(tenant_ids)} revision={top_revision}\"\n )\n\n time_elapsed = time.monotonic() - time_start\n task_logger.info(\n f\"cloud_check_alembic finished: num_tenants={len(tenant_ids)} elapsed={time_elapsed:.2f}\"\n )\n return True\n\n\n@shared_task(\n name=OnyxCeleryTask.CLOUD_MONITOR_CELERY_QUEUES, ignore_result=True, bind=True\n)\ndef cloud_monitor_celery_queues(\n self: Task,\n) -> None:\n return monitor_celery_queues_helper(self)\n\n\n@shared_task(name=OnyxCeleryTask.MONITOR_CELERY_QUEUES, ignore_result=True, bind=True)\ndef monitor_celery_queues(self: Task, *, tenant_id: str) -> None: # noqa: ARG001\n return monitor_celery_queues_helper(self)\n\n\ndef monitor_celery_queues_helper(\n task: Task,\n) -> None:\n \"\"\"A task to monitor all celery queue lengths.\"\"\"\n\n r_celery = celery_get_broker_client(task.app)\n n_celery = celery_get_queue_length(OnyxCeleryQueues.PRIMARY, r_celery)\n n_docfetching = celery_get_queue_length(\n OnyxCeleryQueues.CONNECTOR_DOC_FETCHING, r_celery\n )\n n_docprocessing = celery_get_queue_length(OnyxCeleryQueues.DOCPROCESSING, r_celery)\n\n n_user_file_processing = celery_get_queue_length(\n OnyxCeleryQueues.USER_FILE_PROCESSING, r_celery\n )\n n_user_file_project_sync = celery_get_queue_length(\n OnyxCeleryQueues.USER_FILE_PROJECT_SYNC, r_celery\n )\n n_user_file_delete = celery_get_queue_length(\n OnyxCeleryQueues.USER_FILE_DELETE, r_celery\n )\n n_sync = celery_get_queue_length(OnyxCeleryQueues.VESPA_METADATA_SYNC, r_celery)\n n_deletion = celery_get_queue_length(OnyxCeleryQueues.CONNECTOR_DELETION, r_celery)\n n_pruning = celery_get_queue_length(OnyxCeleryQueues.CONNECTOR_PRUNING, r_celery)\n n_permissions_sync = celery_get_queue_length(\n OnyxCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC, r_celery\n )\n n_external_group_sync = celery_get_queue_length(\n OnyxCeleryQueues.CONNECTOR_EXTERNAL_GROUP_SYNC, r_celery\n )\n n_permissions_upsert = celery_get_queue_length(\n OnyxCeleryQueues.DOC_PERMISSIONS_UPSERT, r_celery\n )\n n_hierarchy_fetching = celery_get_queue_length(\n OnyxCeleryQueues.CONNECTOR_HIERARCHY_FETCHING, r_celery\n )\n n_llm_model_update = celery_get_queue_length(\n OnyxCeleryQueues.LLM_MODEL_UPDATE, r_celery\n )\n n_checkpoint_cleanup = celery_get_queue_length(\n OnyxCeleryQueues.CHECKPOINT_CLEANUP, r_celery\n )\n n_index_attempt_cleanup = celery_get_queue_length(\n OnyxCeleryQueues.INDEX_ATTEMPT_CLEANUP, r_celery\n )\n n_csv_generation = celery_get_queue_length(\n OnyxCeleryQueues.CSV_GENERATION, r_celery\n )\n n_monitoring = celery_get_queue_length(OnyxCeleryQueues.MONITORING, r_celery)\n n_sandbox = celery_get_queue_length(OnyxCeleryQueues.SANDBOX, r_celery)\n n_opensearch_migration = celery_get_queue_length(\n OnyxCeleryQueues.OPENSEARCH_MIGRATION, r_celery\n )\n\n n_docfetching_prefetched = celery_get_unacked_task_ids(\n OnyxCeleryQueues.CONNECTOR_DOC_FETCHING, r_celery\n )\n n_docprocessing_prefetched = celery_get_unacked_task_ids(\n OnyxCeleryQueues.DOCPROCESSING, r_celery\n )\n\n task_logger.info(\n f\"Queue lengths: celery={n_celery} \"\n f\"docfetching={n_docfetching} \"\n f\"docfetching_prefetched={len(n_docfetching_prefetched)} \"\n f\"docprocessing={n_docprocessing} \"\n f\"docprocessing_prefetched={len(n_docprocessing_prefetched)} \"\n f\"user_file_processing={n_user_file_processing} \"\n f\"user_file_project_sync={n_user_file_project_sync} \"\n f\"user_file_delete={n_user_file_delete} \"\n f\"sync={n_sync} \"\n f\"deletion={n_deletion} \"\n f\"pruning={n_pruning} \"\n f\"permissions_sync={n_permissions_sync} \"\n f\"external_group_sync={n_external_group_sync} \"\n f\"permissions_upsert={n_permissions_upsert} \"\n f\"hierarchy_fetching={n_hierarchy_fetching} \"\n f\"llm_model_update={n_llm_model_update} \"\n f\"checkpoint_cleanup={n_checkpoint_cleanup} \"\n f\"index_attempt_cleanup={n_index_attempt_cleanup} \"\n f\"csv_generation={n_csv_generation} \"\n f\"monitoring={n_monitoring} \"\n f\"sandbox={n_sandbox} \"\n f\"opensearch_migration={n_opensearch_migration} \"\n )\n\n\n\"\"\"Memory monitoring\"\"\"\n\n\ndef _get_cmdline_for_process(process: psutil.Process) -> str | None:\n try:\n return \" \".join(process.cmdline())\n except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):\n return None\n\n\n@shared_task(\n name=OnyxCeleryTask.MONITOR_PROCESS_MEMORY,\n ignore_result=True,\n soft_time_limit=_MONITORING_SOFT_TIME_LIMIT,\n time_limit=_MONITORING_TIME_LIMIT,\n queue=OnyxCeleryQueues.MONITORING,\n bind=True,\n)\ndef monitor_process_memory(self: Task, *, tenant_id: str) -> None: # noqa: ARG001\n \"\"\"\n Task to monitor memory usage of supervisor-managed processes.\n This periodically checks the memory usage of processes and logs information\n in a standardized format.\n\n The task looks for processes managed by supervisor and logs their\n memory usage statistics. This is useful for monitoring memory consumption\n over time and identifying potential memory leaks.\n \"\"\"\n # don't run this task in multi-tenant mode, have other, better means of monitoring\n if MULTI_TENANT:\n return\n\n # Skip memory monitoring if not in container\n if not is_running_in_container():\n return\n\n try:\n # Get all supervisor-managed processes\n supervisor_processes: dict[int, str] = {}\n\n # Map cmd line elements to more readable process names\n process_type_mapping = {\n \"--hostname=primary\": \"primary\",\n \"--hostname=light\": \"light\",\n \"--hostname=heavy\": \"heavy\",\n \"--hostname=indexing\": \"indexing\",\n \"--hostname=monitoring\": \"monitoring\",\n \"beat\": \"beat\",\n \"slack/listener.py\": \"slack\",\n }\n\n # Find all python processes that are likely celery workers\n for proc in psutil.process_iter():\n cmdline = _get_cmdline_for_process(proc)\n if not cmdline:\n continue\n\n # Match supervisor-managed processes\n for process_name, process_type in process_type_mapping.items():\n if process_name in cmdline:\n if process_type in supervisor_processes.values():\n task_logger.error(\n f\"Duplicate process type for type {process_type} with cmd {cmdline} with pid={proc.pid}.\"\n )\n continue\n\n supervisor_processes[proc.pid] = process_type\n break\n\n if len(supervisor_processes) != len(process_type_mapping):\n task_logger.error(\n f\"Missing processes: {set(process_type_mapping.keys()).symmetric_difference(supervisor_processes.values())}\"\n )\n\n # Log memory usage for each process\n for pid, process_type in supervisor_processes.items():\n try:\n emit_process_memory(pid, process_type, {})\n except psutil.NoSuchProcess:\n # Process may have terminated since we obtained the list\n continue\n except Exception as e:\n task_logger.exception(f\"Error monitoring process {pid}: {str(e)}\")\n\n except Exception:\n task_logger.exception(\"Error in monitor_process_memory task\")\n\n\n@shared_task(\n name=OnyxCeleryTask.CLOUD_MONITOR_CELERY_PIDBOX, ignore_result=True, bind=True\n)\ndef cloud_monitor_celery_pidbox(\n self: Task,\n) -> None:\n \"\"\"\n Celery can leave behind orphaned pidboxes from old workers that are idle and never cleaned up.\n This task removes them based on idle time to avoid Redis clutter and overflowing the instance.\n This is a real issue we've observed in production.\n\n Note:\n - Setting CELERY_ENABLE_REMOTE_CONTROL = False would prevent pidbox keys entirely,\n but might also disable features like inspect, broadcast, and worker remote control.\n Use with caution.\n \"\"\"\n\n num_deleted = 0\n\n MAX_PIDBOX_IDLE = 24 * 3600 # 1 day in seconds\n r_celery = celery_get_broker_client(self.app)\n for key in r_celery.scan_iter(\"*.reply.celery.pidbox\"):\n key_bytes = cast(bytes, key)\n key_str = key_bytes.decode(\"utf-8\")\n if key_str.startswith(\"_kombu\"):\n continue\n\n idletime_raw = r_celery.object(\"idletime\", key)\n if idletime_raw is None:\n continue\n\n idletime = cast(int, idletime_raw)\n if idletime < MAX_PIDBOX_IDLE:\n continue\n\n r_celery.delete(key)\n task_logger.info(\n f\"Deleted idle pidbox: pidbox={key_str} idletime={idletime} max_idletime={MAX_PIDBOX_IDLE}\"\n )\n num_deleted += 1\n\n # Enable later in case we want some aggregate metrics\n # task_logger.info(f\"Deleted idle pidbox: pidbox={key_str}\")\n" }, { "path": "backend/onyx/background/celery/tasks/opensearch_migration/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/opensearch_migration/constants.py", "content": "# Tasks are expected to cease execution and do cleanup after the soft time\n# limit. In principle they are also forceably terminated after the hard time\n# limit, in practice this does not happen since we use threadpools for Celery\n# task execution, and we simple hope that the total task time plus cleanup does\n# not exceed this. Therefore tasks should regularly check their timeout and lock\n# status. The lock timeout is the maximum time the lock manager (Redis in this\n# case) will enforce the lock, independent of what is happening in the task. To\n# reduce the chances that a task is still doing work while a lock has expired,\n# make the lock timeout well above the task timeouts. In practice we should\n# never see locks be held for this long anyway because a task should release the\n# lock after its cleanup which happens at most after its soft timeout.\n\n# Constants corresponding to migrate_documents_from_vespa_to_opensearch_task.\nfrom onyx.configs.app_configs import OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE\n\n\nMIGRATION_TASK_SOFT_TIME_LIMIT_S = 60 * 5 # 5 minutes.\nMIGRATION_TASK_TIME_LIMIT_S = 60 * 6 # 6 minutes.\n# The maximum time the lock can be held for. Will automatically be released\n# after this time.\nMIGRATION_TASK_LOCK_TIMEOUT_S = 60 * 7 # 7 minutes.\nassert (\n MIGRATION_TASK_SOFT_TIME_LIMIT_S < MIGRATION_TASK_TIME_LIMIT_S\n), \"The soft time limit must be less than the time limit.\"\nassert (\n MIGRATION_TASK_TIME_LIMIT_S < MIGRATION_TASK_LOCK_TIMEOUT_S\n), \"The time limit must be less than the lock timeout.\"\n# Time to wait to acquire the lock.\nMIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S = 60 * 2 # 2 minutes.\n\n# Constants corresponding to check_for_documents_for_opensearch_migration_task.\nCHECK_FOR_DOCUMENTS_TASK_SOFT_TIME_LIMIT_S = 60 # 60 seconds / 1 minute.\nCHECK_FOR_DOCUMENTS_TASK_TIME_LIMIT_S = 90 # 90 seconds.\n# The maximum time the lock can be held for. Will automatically be released\n# after this time.\nCHECK_FOR_DOCUMENTS_TASK_LOCK_TIMEOUT_S = 120 # 120 seconds / 2 minutes.\nassert (\n CHECK_FOR_DOCUMENTS_TASK_SOFT_TIME_LIMIT_S < CHECK_FOR_DOCUMENTS_TASK_TIME_LIMIT_S\n), \"The soft time limit must be less than the time limit.\"\nassert (\n CHECK_FOR_DOCUMENTS_TASK_TIME_LIMIT_S < CHECK_FOR_DOCUMENTS_TASK_LOCK_TIMEOUT_S\n), \"The time limit must be less than the lock timeout.\"\n# Time to wait to acquire the lock.\nCHECK_FOR_DOCUMENTS_TASK_LOCK_BLOCKING_TIMEOUT_S = 30 # 30 seconds.\n\nTOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE = 15\n\n# WARNING: Do not change these values without knowing what changes also need to\n# be made to OpenSearchTenantMigrationRecord.\nGET_VESPA_CHUNKS_PAGE_SIZE = OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE\nGET_VESPA_CHUNKS_SLICE_COUNT = 4\n\n# String used to indicate in the vespa_visit_continuation_token mapping that the\n# slice has finished and there is nothing left to visit.\nFINISHED_VISITING_SLICE_CONTINUATION_TOKEN = (\n \"FINISHED_VISITING_SLICE_CONTINUATION_TOKEN\"\n)\n" }, { "path": "backend/onyx/background/celery/tasks/opensearch_migration/tasks.py", "content": "\"\"\"Celery tasks for migrating documents from Vespa to OpenSearch.\"\"\"\n\nimport time\nimport traceback\n\nfrom celery import shared_task\nfrom celery import Task\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n FINISHED_VISITING_SLICE_CONTINUATION_TOKEN,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n GET_VESPA_CHUNKS_PAGE_SIZE,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n MIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n MIGRATION_TASK_LOCK_TIMEOUT_S,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n MIGRATION_TASK_SOFT_TIME_LIMIT_S,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n MIGRATION_TASK_TIME_LIMIT_S,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.transformer import (\n transform_vespa_chunks_to_opensearch_chunks,\n)\nfrom onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\nfrom onyx.configs.app_configs import VESPA_MIGRATION_REQUEST_TIMEOUT_S\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.opensearch_migration import build_sanitized_to_original_doc_id_mapping\nfrom onyx.db.opensearch_migration import get_vespa_visit_state\nfrom onyx.db.opensearch_migration import is_migration_completed\nfrom onyx.db.opensearch_migration import (\n mark_migration_completed_time_if_not_set_with_commit,\n)\nfrom onyx.db.opensearch_migration import (\n try_insert_opensearch_tenant_migration_record_with_commit,\n)\nfrom onyx.db.opensearch_migration import update_vespa_visit_progress_with_commit\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchDocumentIndex,\n)\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\nfrom onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex\nfrom onyx.indexing.models import IndexingSetting\nfrom onyx.redis.redis_pool import get_redis_client\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\ndef is_continuation_token_done_for_all_slices(\n continuation_token_map: dict[int, str | None],\n) -> bool:\n return all(\n continuation_token == FINISHED_VISITING_SLICE_CONTINUATION_TOKEN\n for continuation_token in continuation_token_map.values()\n )\n\n\n# shared_task allows this task to be shared across celery app instances.\n@shared_task(\n name=OnyxCeleryTask.MIGRATE_CHUNKS_FROM_VESPA_TO_OPENSEARCH_TASK,\n # Does not store the task's return value in the result backend.\n ignore_result=True,\n # WARNING: This is here just for rigor but since we use threads for Celery\n # this config is not respected and timeout logic must be implemented in the\n # task.\n soft_time_limit=MIGRATION_TASK_SOFT_TIME_LIMIT_S,\n # WARNING: This is here just for rigor but since we use threads for Celery\n # this config is not respected and timeout logic must be implemented in the\n # task.\n time_limit=MIGRATION_TASK_TIME_LIMIT_S,\n # Passed in self to the task to get task metadata.\n bind=True,\n)\ndef migrate_chunks_from_vespa_to_opensearch_task(\n self: Task, # noqa: ARG001\n *,\n tenant_id: str,\n) -> bool | None:\n \"\"\"\n Periodic task to migrate chunks from Vespa to OpenSearch via the Visit API.\n\n Uses Vespa's Visit API to iterate through ALL chunks in bulk (not\n per-document), transform them, and index them into OpenSearch. Progress is\n tracked via a continuation token map stored in the\n OpenSearchTenantMigrationRecord.\n\n The first time we see no continuation token map and non-zero chunks\n migrated, we consider the migration complete and all subsequent invocations\n are no-ops.\n\n We divide the index into GET_VESPA_CHUNKS_SLICE_COUNT independent slices\n where progress is tracked for each slice.\n\n Returns:\n None if OpenSearch migration is not enabled, or if the lock could not be\n acquired; effectively a no-op. True if the task completed\n successfully. False if the task errored.\n \"\"\"\n # 1. Check if we should run the task.\n # 1.a. If OpenSearch indexing is disabled, we don't run the task.\n if not ENABLE_OPENSEARCH_INDEXING_FOR_ONYX:\n task_logger.warning(\n \"OpenSearch migration is not enabled, skipping chunk migration task.\"\n )\n return None\n task_logger.info(\"Starting chunk-level migration from Vespa to OpenSearch.\")\n task_start_time = time.monotonic()\n\n # 1.b. Only one instance per tenant of this task may run concurrently at\n # once. If we fail to acquire a lock, we assume it is because another task\n # has one and we exit.\n r = get_redis_client()\n lock: RedisLock = r.lock(\n name=OnyxRedisLocks.OPENSEARCH_MIGRATION_BEAT_LOCK,\n # The maximum time the lock can be held for. Will automatically be\n # released after this time.\n timeout=MIGRATION_TASK_LOCK_TIMEOUT_S,\n # .acquire will block until the lock is acquired.\n blocking=True,\n # Time to wait to acquire the lock.\n blocking_timeout=MIGRATION_TASK_LOCK_BLOCKING_TIMEOUT_S,\n )\n if not lock.acquire():\n task_logger.warning(\n \"The OpenSearch migration task timed out waiting for the lock.\"\n )\n return None\n else:\n task_logger.info(\n f\"Acquired the OpenSearch migration lock. Took {time.monotonic() - task_start_time:.3f} seconds. \"\n f\"Token: {lock.local.token}\"\n )\n\n # 2. Prepare to migrate.\n total_chunks_migrated_this_task = 0\n total_chunks_errored_this_task = 0\n try:\n # 2.a. Double-check that tenant info is correct.\n if tenant_id != get_current_tenant_id():\n err_str = (\n f\"Tenant ID mismatch in the OpenSearch migration task: \"\n f\"{tenant_id} != {get_current_tenant_id()}. This should never happen.\"\n )\n task_logger.error(err_str)\n return False\n\n # Do as much as we can with a DB session in one spot to not hold a\n # session during a migration batch.\n with get_session_with_current_tenant() as db_session:\n # 2.b. Immediately check to see if this tenant is done, to save\n # having to do any other work. This function does not require a\n # migration record to necessarily exist.\n if is_migration_completed(db_session):\n return True\n\n # 2.c. Try to insert the OpenSearchTenantMigrationRecord table if it\n # does not exist.\n try_insert_opensearch_tenant_migration_record_with_commit(db_session)\n\n # 2.d. Get search settings.\n search_settings = get_current_search_settings(db_session)\n indexing_setting = IndexingSetting.from_db_model(search_settings)\n\n # 2.e. Build sanitized to original doc ID mapping to check for\n # conflicts in the event we sanitize a doc ID to an\n # already-existing doc ID.\n # We reconstruct this mapping for every task invocation because\n # a document may have been added in the time between two tasks.\n sanitized_doc_start_time = time.monotonic()\n sanitized_to_original_doc_id_mapping = (\n build_sanitized_to_original_doc_id_mapping(db_session)\n )\n task_logger.debug(\n f\"Built sanitized_to_original_doc_id_mapping with {len(sanitized_to_original_doc_id_mapping)} entries \"\n f\"in {time.monotonic() - sanitized_doc_start_time:.3f} seconds.\"\n )\n\n # 2.f. Get the current migration state.\n continuation_token_map, total_chunks_migrated = get_vespa_visit_state(\n db_session\n )\n # 2.f.1. Double-check that the migration state does not imply\n # completion. Really we should never have to enter this block as we\n # would expect is_migration_completed to return True, but in the\n # strange event that the migration is complete but the migration\n # completed time was never stamped, we do so here.\n if is_continuation_token_done_for_all_slices(continuation_token_map):\n task_logger.info(\n f\"OpenSearch migration COMPLETED for tenant {tenant_id}. Total chunks migrated: {total_chunks_migrated}.\"\n )\n mark_migration_completed_time_if_not_set_with_commit(db_session)\n return True\n task_logger.debug(\n f\"Read the tenant migration record. Total chunks migrated: {total_chunks_migrated}. \"\n f\"Continuation token map: {continuation_token_map}\"\n )\n\n with get_vespa_http_client(\n timeout=VESPA_MIGRATION_REQUEST_TIMEOUT_S\n ) as vespa_client:\n # 2.g. Create the OpenSearch and Vespa document indexes.\n tenant_state = TenantState(tenant_id=tenant_id, multitenant=MULTI_TENANT)\n opensearch_document_index = OpenSearchDocumentIndex(\n tenant_state=tenant_state,\n index_name=search_settings.index_name,\n embedding_dim=indexing_setting.final_embedding_dim,\n embedding_precision=indexing_setting.embedding_precision,\n )\n vespa_document_index = VespaDocumentIndex(\n index_name=search_settings.index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=False,\n httpx_client=vespa_client,\n )\n\n # 2.h. Get the approximate chunk count in Vespa as of this time to\n # update the migration record.\n approx_chunk_count_in_vespa: int | None = None\n get_chunk_count_start_time = time.monotonic()\n try:\n approx_chunk_count_in_vespa = vespa_document_index.get_chunk_count()\n except Exception:\n # This failure should not be blocking.\n task_logger.exception(\n \"Error getting approximate chunk count in Vespa. Moving on...\"\n )\n task_logger.debug(\n f\"Took {time.monotonic() - get_chunk_count_start_time:.3f} seconds to attempt to get \"\n f\"approximate chunk count in Vespa. Got {approx_chunk_count_in_vespa}.\"\n )\n\n # 3. Do the actual migration in batches until we run out of time.\n while (\n time.monotonic() - task_start_time < MIGRATION_TASK_SOFT_TIME_LIMIT_S\n and lock.owned()\n ):\n # 3.a. Get the next batch of raw chunks from Vespa.\n get_vespa_chunks_start_time = time.monotonic()\n raw_vespa_chunks, next_continuation_token_map = (\n vespa_document_index.get_all_raw_document_chunks_paginated(\n continuation_token_map=continuation_token_map,\n page_size=GET_VESPA_CHUNKS_PAGE_SIZE,\n )\n )\n task_logger.debug(\n f\"Read {len(raw_vespa_chunks)} chunks from Vespa in {time.monotonic() - get_vespa_chunks_start_time:.3f} \"\n f\"seconds. Next continuation token map: {next_continuation_token_map}\"\n )\n\n # 3.b. Transform the raw chunks to OpenSearch chunks in memory.\n opensearch_document_chunks, errored_chunks = (\n transform_vespa_chunks_to_opensearch_chunks(\n raw_vespa_chunks,\n tenant_state,\n sanitized_to_original_doc_id_mapping,\n )\n )\n if len(opensearch_document_chunks) != len(raw_vespa_chunks):\n task_logger.error(\n f\"Migration task error: Number of candidate chunks to migrate ({len(opensearch_document_chunks)}) does \"\n f\"not match number of chunks in Vespa ({len(raw_vespa_chunks)}). {len(errored_chunks)} chunks \"\n \"errored.\"\n )\n\n # 3.c. Index the OpenSearch chunks into OpenSearch.\n index_opensearch_chunks_start_time = time.monotonic()\n opensearch_document_index.index_raw_chunks(\n chunks=opensearch_document_chunks\n )\n task_logger.debug(\n f\"Indexed {len(opensearch_document_chunks)} chunks into OpenSearch in \"\n f\"{time.monotonic() - index_opensearch_chunks_start_time:.3f} seconds.\"\n )\n\n total_chunks_migrated_this_task += len(opensearch_document_chunks)\n total_chunks_errored_this_task += len(errored_chunks)\n\n # Do as much as we can with a DB session in one spot to not hold a\n # session during a migration batch.\n with get_session_with_current_tenant() as db_session:\n # 3.d. Update the migration state.\n update_vespa_visit_progress_with_commit(\n db_session,\n continuation_token_map=next_continuation_token_map,\n chunks_processed=len(opensearch_document_chunks),\n chunks_errored=len(errored_chunks),\n approx_chunk_count_in_vespa=approx_chunk_count_in_vespa,\n )\n\n # 3.e. Get the current migration state. Even thought we\n # technically have it in-memory since we just wrote it, we\n # want to reference the DB as the source of truth at all\n # times.\n continuation_token_map, total_chunks_migrated = (\n get_vespa_visit_state(db_session)\n )\n # 3.e.1. Check if the migration is done.\n if is_continuation_token_done_for_all_slices(\n continuation_token_map\n ):\n task_logger.info(\n f\"OpenSearch migration COMPLETED for tenant {tenant_id}. Total chunks migrated: {total_chunks_migrated}.\"\n )\n mark_migration_completed_time_if_not_set_with_commit(db_session)\n return True\n task_logger.debug(\n f\"Read the tenant migration record. Total chunks migrated: {total_chunks_migrated}. \"\n f\"Continuation token map: {continuation_token_map}\"\n )\n except Exception:\n traceback.print_exc()\n task_logger.exception(\"Error in the OpenSearch migration task.\")\n return False\n finally:\n if lock.owned():\n lock.release()\n else:\n task_logger.warning(\n \"The OpenSearch migration lock was not owned on completion of the migration task.\"\n )\n\n task_logger.info(\n f\"OpenSearch chunk migration task pausing (time limit reached). \"\n f\"Total chunks migrated this task: {total_chunks_migrated_this_task}. \"\n f\"Total chunks errored this task: {total_chunks_errored_this_task}. \"\n f\"Elapsed: {time.monotonic() - task_start_time:.3f}s. \"\n \"Will resume from continuation token on next invocation.\"\n )\n\n return True\n" }, { "path": "backend/onyx/background/celery/tasks/opensearch_migration/transformer.py", "content": "import traceback\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.schema import DocumentChunk\nfrom onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST\nfrom onyx.document_index.vespa_constants import BLURB\nfrom onyx.document_index.vespa_constants import BOOST\nfrom onyx.document_index.vespa_constants import CHUNK_CONTEXT\nfrom onyx.document_index.vespa_constants import CHUNK_ID\nfrom onyx.document_index.vespa_constants import CONTENT\nfrom onyx.document_index.vespa_constants import DOC_SUMMARY\nfrom onyx.document_index.vespa_constants import DOC_UPDATED_AT\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID\nfrom onyx.document_index.vespa_constants import DOCUMENT_SETS\nfrom onyx.document_index.vespa_constants import EMBEDDINGS\nfrom onyx.document_index.vespa_constants import FULL_CHUNK_EMBEDDING_KEY\nfrom onyx.document_index.vespa_constants import HIDDEN\nfrom onyx.document_index.vespa_constants import IMAGE_FILE_NAME\nfrom onyx.document_index.vespa_constants import METADATA_LIST\nfrom onyx.document_index.vespa_constants import METADATA_SUFFIX\nfrom onyx.document_index.vespa_constants import PERSONAS\nfrom onyx.document_index.vespa_constants import PRIMARY_OWNERS\nfrom onyx.document_index.vespa_constants import SECONDARY_OWNERS\nfrom onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER\nfrom onyx.document_index.vespa_constants import SOURCE_LINKS\nfrom onyx.document_index.vespa_constants import SOURCE_TYPE\nfrom onyx.document_index.vespa_constants import TENANT_ID\nfrom onyx.document_index.vespa_constants import TITLE\nfrom onyx.document_index.vespa_constants import TITLE_EMBEDDING\nfrom onyx.document_index.vespa_constants import USER_PROJECT\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger(__name__)\n\n\nFIELDS_NEEDED_FOR_TRANSFORMATION: list[str] = [\n DOCUMENT_ID,\n CHUNK_ID,\n TITLE,\n TITLE_EMBEDDING,\n CONTENT,\n EMBEDDINGS,\n SOURCE_TYPE,\n METADATA_LIST,\n DOC_UPDATED_AT,\n HIDDEN,\n BOOST,\n SEMANTIC_IDENTIFIER,\n IMAGE_FILE_NAME,\n SOURCE_LINKS,\n BLURB,\n DOC_SUMMARY,\n CHUNK_CONTEXT,\n METADATA_SUFFIX,\n DOCUMENT_SETS,\n USER_PROJECT,\n PERSONAS,\n PRIMARY_OWNERS,\n SECONDARY_OWNERS,\n ACCESS_CONTROL_LIST,\n]\nif MULTI_TENANT:\n FIELDS_NEEDED_FOR_TRANSFORMATION.append(TENANT_ID)\n\n\ndef _extract_content_vector(embeddings: Any) -> list[float]:\n \"\"\"Extracts the full chunk embedding vector from Vespa's embeddings tensor.\n\n Vespa stores embeddings as a tensor(t{},x[dim]) where 't' maps\n embedding names (like \"full_chunk\") to vectors. The API can return this in\n different formats:\n 1. Direct list: {\"full_chunk\": [...]}\n 2. Blocks format: {\"blocks\": {\"full_chunk\": [0.1, 0.2, ...]}}\n 3. Possibly other formats.\n\n We only support formats 1 and 2. Any other supplied format will raise an\n error.\n\n Raises:\n ValueError: If the embeddings format is not supported.\n\n Returns:\n The full chunk content embedding vector as a list of floats.\n \"\"\"\n if isinstance(embeddings, dict):\n # Handle format 1.\n full_chunk_embedding = embeddings.get(FULL_CHUNK_EMBEDDING_KEY)\n if isinstance(full_chunk_embedding, list):\n # Double check that within the list we have floats and not another\n # list or dict.\n if not full_chunk_embedding:\n raise ValueError(\"Full chunk embedding is empty.\")\n if isinstance(full_chunk_embedding[0], float):\n return full_chunk_embedding\n\n # Handle format 2.\n blocks = embeddings.get(\"blocks\")\n if isinstance(blocks, dict):\n full_chunk_embedding = blocks.get(FULL_CHUNK_EMBEDDING_KEY)\n if isinstance(full_chunk_embedding, list):\n # Double check that within the list we have floats and not another\n # list or dict.\n if not full_chunk_embedding:\n raise ValueError(\"Full chunk embedding is empty.\")\n if isinstance(full_chunk_embedding[0], float):\n return full_chunk_embedding\n\n raise ValueError(f\"Unknown embedding format: {type(embeddings)}\")\n\n\ndef _extract_title_vector(title_embedding: Any | None) -> list[float] | None:\n \"\"\"Extract the title embedding vector.\n\n Returns None if no title embedding exists.\n\n Vespa returns title_embedding as tensor(x[dim]) which can be in\n formats:\n 1. Direct list: [0.1, 0.2, ...]\n 2. Values format: {\"values\": [0.1, 0.2, ...]}\n 3. Possibly other formats.\n\n Only formats 1 and 2 are supported. Any other supplied format will raise an\n error.\n\n Raises:\n ValueError: If the title embedding format is not supported.\n\n Returns:\n The title embedding vector as a list of floats.\n \"\"\"\n if title_embedding is None:\n return None\n\n # Handle format 1.\n if isinstance(title_embedding, list):\n # Double check that within the list we have floats and not another\n # list or dict.\n if not title_embedding:\n return None\n if isinstance(title_embedding[0], float):\n return title_embedding\n\n # Handle format 2.\n if isinstance(title_embedding, dict):\n # Try values format.\n values = title_embedding.get(\"values\")\n if values is not None and isinstance(values, list):\n # Double check that within the list we have floats and not another\n # list or dict.\n if not values:\n return None\n if isinstance(values[0], float):\n return values\n\n raise ValueError(f\"Unknown title embedding format: {type(title_embedding)}\")\n\n\ndef _transform_vespa_document_sets_to_opensearch_document_sets(\n vespa_document_sets: dict[str, int] | None,\n) -> list[str] | None:\n if not vespa_document_sets:\n return None\n return list(vespa_document_sets.keys())\n\n\ndef _transform_vespa_acl_to_opensearch_acl(\n vespa_acl: dict[str, int] | None,\n) -> tuple[bool, list[str]]:\n if not vespa_acl:\n return False, []\n acl_list = list(vespa_acl.keys())\n is_public = PUBLIC_DOC_PAT in acl_list\n if is_public:\n acl_list.remove(PUBLIC_DOC_PAT)\n return is_public, acl_list\n\n\ndef transform_vespa_chunks_to_opensearch_chunks(\n vespa_chunks: list[dict[str, Any]],\n tenant_state: TenantState,\n sanitized_to_original_doc_id_mapping: dict[str, str],\n) -> tuple[list[DocumentChunk], list[dict[str, Any]]]:\n result: list[DocumentChunk] = []\n errored_chunks: list[dict[str, Any]] = []\n for vespa_chunk in vespa_chunks:\n try:\n # This should exist; fail loudly if it does not.\n vespa_document_id: str = vespa_chunk[DOCUMENT_ID]\n if not vespa_document_id:\n raise ValueError(\"Missing document_id in Vespa chunk.\")\n # Vespa doc IDs were sanitized using\n # replace_invalid_doc_id_characters. This was a poor design choice\n # and we don't want this in OpenSearch; whatever restrictions there\n # may be on indexed chunk ID should have no bearing on the chunk's\n # document ID field, even if document ID is an argument to the chunk\n # ID. Deliberately choose to use the real doc ID supplied to this\n # function.\n if vespa_document_id in sanitized_to_original_doc_id_mapping:\n logger.warning(\n f\"Migration warning: Vespa document ID {vespa_document_id} does not match the document ID supplied \"\n f\"{sanitized_to_original_doc_id_mapping[vespa_document_id]}. \"\n \"The Vespa ID will be discarded.\"\n )\n document_id = sanitized_to_original_doc_id_mapping.get(\n vespa_document_id, vespa_document_id\n )\n\n # This should exist; fail loudly if it does not.\n chunk_index: int = vespa_chunk[CHUNK_ID]\n\n title: str | None = vespa_chunk.get(TITLE)\n # WARNING: Should supply format.tensors=short-value to the Vespa\n # client in order to get a supported format for the tensors.\n title_vector: list[float] | None = _extract_title_vector(\n vespa_chunk.get(TITLE_EMBEDDING)\n )\n\n # This should exist; fail loudly if it does not.\n content: str = vespa_chunk[CONTENT]\n if not content:\n raise ValueError(\n f\"Missing content in Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index}.\"\n )\n # This should exist; fail loudly if it does not.\n # WARNING: Should supply format.tensors=short-value to the Vespa\n # client in order to get a supported format for the tensors.\n content_vector: list[float] = _extract_content_vector(\n vespa_chunk[EMBEDDINGS]\n )\n if not content_vector:\n raise ValueError(\n f\"Missing content_vector in Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index}.\"\n )\n\n # This should exist; fail loudly if it does not.\n source_type: str = vespa_chunk[SOURCE_TYPE]\n if not source_type:\n raise ValueError(\n f\"Missing source_type in Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index}.\"\n )\n\n metadata_list: list[str] | None = vespa_chunk.get(METADATA_LIST)\n\n _raw_doc_updated_at: int | None = vespa_chunk.get(DOC_UPDATED_AT)\n last_updated: datetime | None = (\n datetime.fromtimestamp(_raw_doc_updated_at, tz=timezone.utc)\n if _raw_doc_updated_at is not None\n else None\n )\n\n hidden: bool = vespa_chunk.get(HIDDEN, False)\n\n # This should exist; fail loudly if it does not.\n global_boost: int = vespa_chunk[BOOST]\n\n # This should exist; fail loudly if it does not.\n semantic_identifier: str = vespa_chunk[SEMANTIC_IDENTIFIER]\n if not semantic_identifier:\n raise ValueError(\n f\"Missing semantic_identifier in Vespa chunk with document ID {vespa_document_id} and chunk \"\n f\"index {chunk_index}.\"\n )\n\n image_file_id: str | None = vespa_chunk.get(IMAGE_FILE_NAME)\n source_links: str | None = vespa_chunk.get(SOURCE_LINKS)\n blurb: str = vespa_chunk.get(BLURB, \"\")\n doc_summary: str = vespa_chunk.get(DOC_SUMMARY, \"\")\n chunk_context: str = vespa_chunk.get(CHUNK_CONTEXT, \"\")\n metadata_suffix: str | None = vespa_chunk.get(METADATA_SUFFIX)\n document_sets: list[str] | None = (\n _transform_vespa_document_sets_to_opensearch_document_sets(\n vespa_chunk.get(DOCUMENT_SETS)\n )\n )\n user_projects: list[int] | None = vespa_chunk.get(USER_PROJECT)\n personas: list[int] | None = vespa_chunk.get(PERSONAS)\n primary_owners: list[str] | None = vespa_chunk.get(PRIMARY_OWNERS)\n secondary_owners: list[str] | None = vespa_chunk.get(SECONDARY_OWNERS)\n\n is_public, acl_list = _transform_vespa_acl_to_opensearch_acl(\n vespa_chunk.get(ACCESS_CONTROL_LIST)\n )\n if not is_public and not acl_list:\n logger.warning(\n f\"Migration warning: Vespa chunk with document ID {vespa_document_id} and chunk index {chunk_index} has no \"\n \"public ACL and no access control list. This does not make sense as it implies the document is never \"\n \"searchable. Continuing with the migration...\"\n )\n\n chunk_tenant_id: str | None = vespa_chunk.get(TENANT_ID)\n if MULTI_TENANT:\n if not chunk_tenant_id:\n raise ValueError(\n \"Missing tenant_id in Vespa chunk in a multi-tenant environment.\"\n )\n if chunk_tenant_id != tenant_state.tenant_id:\n raise ValueError(\n f\"Chunk tenant_id {chunk_tenant_id} does not match expected tenant_id {tenant_state.tenant_id}\"\n )\n\n opensearch_chunk = DocumentChunk(\n # We deliberately choose to use the doc ID supplied to this function\n # over the Vespa doc ID.\n document_id=document_id,\n chunk_index=chunk_index,\n title=title,\n title_vector=title_vector,\n content=content,\n content_vector=content_vector,\n source_type=source_type,\n metadata_list=metadata_list,\n last_updated=last_updated,\n public=is_public,\n access_control_list=acl_list,\n hidden=hidden,\n global_boost=global_boost,\n semantic_identifier=semantic_identifier,\n image_file_id=image_file_id,\n source_links=source_links,\n blurb=blurb,\n doc_summary=doc_summary,\n chunk_context=chunk_context,\n metadata_suffix=metadata_suffix,\n document_sets=document_sets,\n user_projects=user_projects,\n personas=personas,\n primary_owners=primary_owners,\n secondary_owners=secondary_owners,\n tenant_id=tenant_state,\n )\n\n result.append(opensearch_chunk)\n except Exception:\n traceback.print_exc()\n logger.exception(\n f\"Migration error: Error transforming Vespa chunk with document ID {vespa_chunk.get(DOCUMENT_ID)} \"\n f\"and chunk index {vespa_chunk.get(CHUNK_ID)} into an OpenSearch chunk. Continuing with \"\n \"the migration...\"\n )\n errored_chunks.append(vespa_chunk)\n\n return result, errored_chunks\n" }, { "path": "backend/onyx/background/celery/tasks/periodic/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/periodic/tasks.py", "content": "#####\n# Periodic Tasks\n#####\nimport json\nfrom typing import Any\n\nfrom celery import shared_task\nfrom celery.contrib.abortable import AbortableTask # type: ignore\nfrom celery.exceptions import TaskRevokedError\nfrom sqlalchemy import inspect\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import PostgresAdvisoryLocks\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n\n@shared_task(\n name=OnyxCeleryTask.KOMBU_MESSAGE_CLEANUP_TASK,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n base=AbortableTask,\n)\ndef kombu_message_cleanup_task(self: Any, tenant_id: str) -> int: # noqa: ARG001\n \"\"\"Runs periodically to clean up the kombu_message table\"\"\"\n\n # we will select messages older than this amount to clean up\n KOMBU_MESSAGE_CLEANUP_AGE = 7 # days\n KOMBU_MESSAGE_CLEANUP_PAGE_LIMIT = 1000\n\n ctx = {}\n ctx[\"last_processed_id\"] = 0\n ctx[\"deleted\"] = 0\n ctx[\"cleanup_age\"] = KOMBU_MESSAGE_CLEANUP_AGE\n ctx[\"page_limit\"] = KOMBU_MESSAGE_CLEANUP_PAGE_LIMIT\n with get_session_with_current_tenant() as db_session:\n # Exit the task if we can't take the advisory lock\n result = db_session.execute(\n text(\"SELECT pg_try_advisory_lock(:id)\"),\n {\"id\": PostgresAdvisoryLocks.KOMBU_MESSAGE_CLEANUP_LOCK_ID.value},\n ).scalar()\n if not result:\n return 0\n\n while True:\n if self.is_aborted():\n raise TaskRevokedError(\"kombu_message_cleanup_task was aborted.\")\n\n b = kombu_message_cleanup_task_helper(ctx, db_session)\n if not b:\n break\n\n db_session.commit()\n\n if ctx[\"deleted\"] > 0:\n task_logger.info(\n f\"Deleted {ctx['deleted']} orphaned messages from kombu_message.\"\n )\n\n return ctx[\"deleted\"]\n\n\ndef kombu_message_cleanup_task_helper(ctx: dict, db_session: Session) -> bool:\n \"\"\"\n Helper function to clean up old messages from the `kombu_message` table that are no longer relevant.\n\n This function retrieves messages from the `kombu_message` table that are no longer visible and\n older than a specified interval. It checks if the corresponding task_id exists in the\n `celery_taskmeta` table. If the task_id does not exist, the message is deleted.\n\n Args:\n ctx (dict): A context dictionary containing configuration parameters such as:\n - 'cleanup_age' (int): The age in days after which messages are considered old.\n - 'page_limit' (int): The maximum number of messages to process in one batch.\n - 'last_processed_id' (int): The ID of the last processed message to handle pagination.\n - 'deleted' (int): A counter to track the number of deleted messages.\n db_session (Session): The SQLAlchemy database session for executing queries.\n\n Returns:\n bool: Returns True if there are more rows to process, False if not.\n \"\"\"\n\n inspector = inspect(db_session.bind)\n if not inspector:\n return False\n\n # With the move to redis as celery's broker and backend, kombu tables may not even exist.\n # We can fail silently.\n if not inspector.has_table(\"kombu_message\"):\n return False\n\n query = text(\n \"\"\"\n SELECT id, timestamp, payload\n FROM kombu_message WHERE visible = 'false'\n AND timestamp < CURRENT_TIMESTAMP - INTERVAL :interval_days\n AND id > :last_processed_id\n ORDER BY id\n LIMIT :page_limit\n\"\"\"\n )\n kombu_messages = db_session.execute(\n query,\n {\n \"interval_days\": f\"{ctx['cleanup_age']} days\",\n \"page_limit\": ctx[\"page_limit\"],\n \"last_processed_id\": ctx[\"last_processed_id\"],\n },\n ).fetchall()\n\n if len(kombu_messages) == 0:\n return False\n\n for msg in kombu_messages:\n payload = json.loads(msg[2])\n task_id = payload[\"headers\"][\"id\"]\n\n # Check if task_id exists in celery_taskmeta\n task_exists = db_session.execute(\n text(\"SELECT 1 FROM celery_taskmeta WHERE task_id = :task_id\"),\n {\"task_id\": task_id},\n ).fetchone()\n\n # If task_id does not exist, delete the message\n if not task_exists:\n result = db_session.execute(\n text(\"DELETE FROM kombu_message WHERE id = :message_id\"),\n {\"message_id\": msg[0]},\n )\n if result.rowcount > 0: # type: ignore\n ctx[\"deleted\"] += 1\n\n ctx[\"last_processed_id\"] = msg[0]\n\n return True\n" }, { "path": "backend/onyx/background/celery/tasks/pruning/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/pruning/tasks.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom uuid import uuid4\n\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom pydantic import ValidationError\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_find_task\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_queue_length\nfrom onyx.background.celery.celery_redis import celery_get_queued_task_ids\nfrom onyx.background.celery.celery_redis import celery_get_unacked_task_ids\nfrom onyx.background.celery.celery_utils import extract_ids_from_runnable_connector\nfrom onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEFAULT\nfrom onyx.background.celery.tasks.docprocessing.utils import IndexingCallbackBase\nfrom onyx.configs.app_configs import ALLOW_SIMULTANEOUS_PRUNING\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_PRUNING_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT\nfrom onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import OnyxRedisSignals\nfrom onyx.connectors.factory import instantiate_connector\nfrom onyx.connectors.models import InputType\nfrom onyx.db.connector import mark_ccpair_as_pruned\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs\nfrom onyx.db.document import get_documents_for_connector_credential_pair\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.hierarchy import delete_orphaned_hierarchy_nodes\nfrom onyx.db.hierarchy import link_hierarchy_nodes_to_documents\nfrom onyx.db.hierarchy import remove_stale_hierarchy_node_cc_pair_entries\nfrom onyx.db.hierarchy import reparent_orphaned_hierarchy_nodes\nfrom onyx.db.hierarchy import update_document_parent_hierarchy_nodes\nfrom onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries\nfrom onyx.db.hierarchy import upsert_hierarchy_nodes_batch\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import HierarchyNode as DBHierarchyNode\nfrom onyx.db.sync_record import insert_sync_record\nfrom onyx.db.sync_record import update_sync_record_status\nfrom onyx.db.tag import delete_orphan_tags__no_commit\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_connector_prune import RedisConnectorPrune\nfrom onyx.redis.redis_connector_prune import RedisConnectorPrunePayload\nfrom onyx.redis.redis_hierarchy import cache_hierarchy_nodes_batch\nfrom onyx.redis.redis_hierarchy import ensure_source_node_exists\nfrom onyx.redis.redis_hierarchy import evict_hierarchy_nodes_from_cache\nfrom onyx.redis.redis_hierarchy import get_node_id_from_raw_id\nfrom onyx.redis.redis_hierarchy import get_source_node_id_from_cache\nfrom onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.server.runtime.onyx_runtime import OnyxRuntime\nfrom onyx.server.utils import make_short_id\nfrom onyx.utils.logger import format_error_for_logging\nfrom onyx.utils.logger import LoggerContextVars\nfrom onyx.utils.logger import pruning_ctx\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef _get_pruning_block_expiration() -> int:\n \"\"\"\n Compute the expiration time for the pruning block signal.\n Base expiration is 60 seconds (1 minute), multiplied by the beat multiplier only in MULTI_TENANT mode.\n \"\"\"\n base_expiration = 60 # seconds\n\n if not MULTI_TENANT:\n return base_expiration\n\n try:\n beat_multiplier = OnyxRuntime.get_beat_multiplier()\n except Exception:\n beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n return int(base_expiration * beat_multiplier)\n\n\ndef _get_fence_validation_block_expiration() -> int:\n \"\"\"\n Compute the expiration time for the fence validation block signal.\n Base expiration is 300 seconds, multiplied by the beat multiplier only in MULTI_TENANT mode.\n \"\"\"\n base_expiration = 300 # seconds\n\n if not MULTI_TENANT:\n return base_expiration\n\n try:\n beat_multiplier = OnyxRuntime.get_beat_multiplier()\n except Exception:\n beat_multiplier = CLOUD_BEAT_MULTIPLIER_DEFAULT\n\n return int(base_expiration * beat_multiplier)\n\n\nclass PruneCallback(IndexingCallbackBase):\n def progress(self, tag: str, amount: int) -> None:\n self.redis_connector.prune.set_active()\n super().progress(tag, amount)\n\n\ndef _resolve_and_update_document_parents(\n db_session: Session,\n redis_client: Redis,\n source: DocumentSource,\n raw_id_to_parent: dict[str, str | None],\n) -> None:\n \"\"\"Resolve parent_hierarchy_raw_node_id → parent_hierarchy_node_id for\n each document and bulk-update the DB. Mirrors the resolution logic in\n run_docfetching.py.\"\"\"\n source_node_id = get_source_node_id_from_cache(redis_client, db_session, source)\n\n resolved: dict[str, int | None] = {}\n for doc_id, raw_parent_id in raw_id_to_parent.items():\n if raw_parent_id is None:\n continue\n node_id, found = get_node_id_from_raw_id(redis_client, source, raw_parent_id)\n resolved[doc_id] = node_id if found else source_node_id\n\n if not resolved:\n return\n\n update_document_parent_hierarchy_nodes(\n db_session=db_session,\n doc_parent_map=resolved,\n commit=True,\n )\n task_logger.info(\n f\"Pruning: resolved and updated parent hierarchy for {len(resolved)} documents (source={source.value})\"\n )\n\n\n\"\"\"Jobs / utils for kicking off pruning tasks.\"\"\"\n\n\ndef _is_pruning_due(cc_pair: ConnectorCredentialPair) -> bool:\n \"\"\"Returns boolean indicating if pruning is due.\n\n Next pruning time is calculated as a delta from the last successful prune, or the\n last successful indexing if pruning has never succeeded.\n\n TODO(rkuo): consider whether we should allow pruning to be immediately rescheduled\n if pruning fails (which is what it does now). A backoff could be reasonable.\n \"\"\"\n\n # skip pruning if no prune frequency is set\n # pruning can still be forced via the API which will run a pruning task directly\n if not cc_pair.connector.prune_freq:\n return False\n\n # skip pruning if not active\n if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:\n return False\n\n # skip pruning if the next scheduled prune time hasn't been reached yet\n last_pruned = cc_pair.last_pruned\n if not last_pruned:\n if not cc_pair.last_successful_index_time:\n # if we've never indexed, we can't prune\n return False\n\n # if never pruned, use the connector creation time. We could also\n # compute the completion time of the first successful index attempt, but\n # that is a reasonably heavy operation. This is a reasonable approximation —\n # in the worst case, we'll prune a little bit earlier than we should.\n last_pruned = cc_pair.connector.time_created\n\n next_prune = last_pruned + timedelta(seconds=cc_pair.connector.prune_freq)\n return datetime.now(timezone.utc) >= next_prune\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_PRUNING,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n bind=True,\n)\ndef check_for_pruning(self: Task, *, tenant_id: str) -> bool | None:\n r = get_redis_client()\n r_replica = get_redis_replica_client()\n\n lock_beat: RedisLock = r.lock(\n OnyxRedisLocks.CHECK_PRUNE_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n try:\n # the entire task needs to run frequently in order to finalize pruning\n\n # but pruning only kicks off once per hour\n if not r.exists(OnyxRedisSignals.BLOCK_PRUNING):\n task_logger.info(\"Checking for pruning due\")\n\n cc_pair_ids: list[int] = []\n with get_session_with_current_tenant() as db_session:\n cc_pairs = get_connector_credential_pairs(db_session)\n for cc_pair_entry in cc_pairs:\n cc_pair_ids.append(cc_pair_entry.id)\n\n for cc_pair_id in cc_pair_ids:\n lock_beat.reacquire()\n with get_session_with_current_tenant() as db_session:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n logger.error(f\"CC pair not found: {cc_pair_id}\")\n continue\n\n if not _is_pruning_due(cc_pair):\n logger.info(f\"CC pair not due for pruning: {cc_pair_id}\")\n continue\n\n payload_id = try_creating_prune_generator_task(\n self.app, cc_pair, db_session, r, tenant_id\n )\n if not payload_id:\n logger.info(f\"Pruning not created: {cc_pair_id}\")\n continue\n\n task_logger.info(\n f\"Pruning queued: cc_pair={cc_pair.id} id={payload_id}\"\n )\n r.set(OnyxRedisSignals.BLOCK_PRUNING, 1, ex=_get_pruning_block_expiration())\n\n # we want to run this less frequently than the overall task\n lock_beat.reacquire()\n if not r.exists(OnyxRedisSignals.BLOCK_VALIDATE_PRUNING_FENCES):\n # clear any permission fences that don't have associated celery tasks in progress\n # tasks can be in the queue in redis, in reserved tasks (prefetched by the worker),\n # or be currently executing\n try:\n r_celery = celery_get_broker_client(self.app)\n validate_pruning_fences(tenant_id, r, r_replica, r_celery, lock_beat)\n except Exception:\n task_logger.exception(\"Exception while validating pruning fences\")\n\n r.set(\n OnyxRedisSignals.BLOCK_VALIDATE_PRUNING_FENCES,\n 1,\n ex=_get_fence_validation_block_expiration(),\n )\n\n # use a lookup table to find active fences. We still have to verify the fence\n # exists since it is an optimization and not the source of truth.\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n\n if not r.exists(key_bytes):\n r.srem(OnyxRedisConstants.ACTIVE_FENCES, key_bytes)\n continue\n\n key_str = key_bytes.decode(\"utf-8\")\n if key_str.startswith(RedisConnectorPrune.FENCE_PREFIX):\n with get_session_with_current_tenant() as db_session:\n monitor_ccpair_pruning_taskset(tenant_id, key_bytes, r, db_session)\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(f\"Unexpected pruning check exception: {error_msg}\")\n task_logger.exception(\"Unexpected exception during pruning check\")\n finally:\n if lock_beat.owned():\n lock_beat.release()\n task_logger.info(f\"check_for_pruning finished: tenant={tenant_id}\")\n return True\n\n\ndef try_creating_prune_generator_task(\n celery_app: Celery,\n cc_pair: ConnectorCredentialPair,\n db_session: Session,\n r: Redis,\n tenant_id: str,\n) -> str | None:\n \"\"\"Checks for any conditions that should block the pruning generator task from being\n created, then creates the task.\n\n Does not check for scheduling related conditions as this function\n is used to trigger prunes immediately, e.g. via the web ui.\n \"\"\"\n\n logger.info(f\"try_creating_prune_generator_task: cc_pair={cc_pair.id}\")\n\n redis_connector = RedisConnector(tenant_id, cc_pair.id)\n\n if not ALLOW_SIMULTANEOUS_PRUNING:\n count = redis_connector.prune.get_active_task_count()\n if count > 0:\n logger.info(\n f\"try_creating_prune_generator_task: cc_pair={cc_pair.id} no simultaneous pruning allowed\"\n )\n return None\n\n LOCK_TIMEOUT = 30\n\n # we need to serialize starting pruning since it can be triggered either via\n # celery beat or manually (API call)\n lock: RedisLock = r.lock(\n DANSWER_REDIS_FUNCTION_LOCK_PREFIX + \"try_creating_prune_generator_task\",\n timeout=LOCK_TIMEOUT,\n )\n\n acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)\n if not acquired:\n logger.info(\n f\"try_creating_prune_generator_task: cc_pair={cc_pair.id} lock not acquired\"\n )\n return None\n\n try:\n # skip pruning if already pruning\n if redis_connector.prune.fenced:\n logger.info(\n f\"try_creating_prune_generator_task: cc_pair={cc_pair.id} already pruning\"\n )\n return None\n\n # skip pruning if the cc_pair is deleting\n if redis_connector.delete.fenced:\n logger.info(\n f\"try_creating_prune_generator_task: cc_pair={cc_pair.id} deleting\"\n )\n return None\n\n # skip pruning if doc permissions sync is running\n if redis_connector.permissions.fenced:\n logger.info(\n f\"try_creating_prune_generator_task: cc_pair={cc_pair.id} permissions sync running\"\n )\n return None\n\n db_session.refresh(cc_pair)\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n logger.info(\n f\"try_creating_prune_generator_task: cc_pair={cc_pair.id} deleting\"\n )\n return None\n\n # add a long running generator task to the queue\n redis_connector.prune.generator_clear()\n redis_connector.prune.taskset_clear()\n\n custom_task_id = f\"{redis_connector.prune.generator_task_key}_{uuid4()}\"\n\n # create before setting fence to avoid race condition where the monitoring\n # task updates the sync record before it is created\n try:\n insert_sync_record(\n db_session=db_session,\n entity_id=cc_pair.id,\n sync_type=SyncType.PRUNING,\n )\n except Exception:\n task_logger.exception(\"insert_sync_record exceptioned.\")\n\n # signal active before the fence is set\n redis_connector.prune.set_active()\n\n # set a basic fence to start\n payload = RedisConnectorPrunePayload(\n id=make_short_id(),\n submitted=datetime.now(timezone.utc),\n started=None,\n celery_task_id=None,\n )\n redis_connector.prune.set_fence(payload)\n\n result = celery_app.send_task(\n OnyxCeleryTask.CONNECTOR_PRUNING_GENERATOR_TASK,\n kwargs=dict(\n cc_pair_id=cc_pair.id,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n tenant_id=tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_PRUNING,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.LOW,\n )\n\n # fill in the celery task id\n payload.celery_task_id = result.id\n redis_connector.prune.set_fence(payload)\n\n payload_id = payload.id\n except Exception as e:\n error_msg = format_error_for_logging(e)\n task_logger.warning(\n f\"Unexpected try_creating_prune_generator_task exception: cc_pair={cc_pair.id} {error_msg}\"\n )\n task_logger.exception(f\"Unexpected exception: cc_pair={cc_pair.id}\")\n return None\n finally:\n if lock.owned():\n lock.release()\n task_logger.info(\n f\"try_creating_prune_generator_task finished: cc_pair={cc_pair.id} payload_id={payload_id}\"\n )\n return payload_id\n\n\n@shared_task(\n name=OnyxCeleryTask.CONNECTOR_PRUNING_GENERATOR_TASK,\n acks_late=False,\n soft_time_limit=JOB_TIMEOUT,\n track_started=True,\n trail=False,\n bind=True,\n)\ndef connector_pruning_generator_task(\n self: Task,\n cc_pair_id: int,\n connector_id: int,\n credential_id: int,\n tenant_id: str,\n) -> None:\n \"\"\"connector pruning task. For a cc pair, this task pulls all document IDs from the source\n and compares those IDs to locally stored documents and deletes all locally stored IDs missing\n from the most recently pulled document ID list\"\"\"\n\n payload_id: str | None = None\n\n LoggerContextVars.reset()\n\n pruning_ctx_dict = pruning_ctx.get()\n pruning_ctx_dict[\"cc_pair_id\"] = cc_pair_id\n pruning_ctx_dict[\"request_id\"] = self.request.id\n pruning_ctx.set(pruning_ctx_dict)\n\n task_logger.info(f\"Pruning generator starting: cc_pair={cc_pair_id}\")\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n\n r = get_redis_client()\n\n # this wait is needed to avoid a race condition where\n # the primary worker sends the task and it is immediately executed\n # before the primary worker can finalize the fence\n start = time.monotonic()\n while True:\n if time.monotonic() - start > CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT:\n raise ValueError(\n f\"connector_prune_generator_task - timed out waiting for fence to be ready: \"\n f\"fence={redis_connector.prune.fence_key}\"\n )\n\n if not redis_connector.prune.fenced: # The fence must exist\n raise ValueError(\n f\"connector_prune_generator_task - fence not found: fence={redis_connector.prune.fence_key}\"\n )\n\n payload = redis_connector.prune.payload # The payload must exist\n if not payload:\n raise ValueError(\n \"connector_prune_generator_task: payload invalid or not found\"\n )\n\n if payload.celery_task_id is None:\n logger.info(\n f\"connector_prune_generator_task - Waiting for fence: fence={redis_connector.prune.fence_key}\"\n )\n time.sleep(1)\n continue\n\n payload_id = payload.id\n\n logger.info(\n f\"connector_prune_generator_task - Fence found, continuing...: \"\n f\"fence={redis_connector.prune.fence_key} \"\n f\"payload_id={payload.id}\"\n )\n break\n\n # set thread_local=False since we don't control what thread the indexing/pruning\n # might run our callback with\n lock: RedisLock = r.lock(\n OnyxRedisLocks.PRUNING_LOCK_PREFIX + f\"_{redis_connector.cc_pair_id}\",\n timeout=CELERY_PRUNING_LOCK_TIMEOUT,\n thread_local=False,\n )\n\n acquired = lock.acquire(blocking=False)\n if not acquired:\n task_logger.warning(\n f\"Pruning task already running, exiting...: cc_pair={cc_pair_id}\"\n )\n return None\n\n try:\n with get_session_with_current_tenant() as db_session:\n cc_pair = get_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n )\n\n if not cc_pair:\n task_logger.warning(\n f\"cc_pair not found for {connector_id} {credential_id}\"\n )\n return\n\n payload = redis_connector.prune.payload\n if not payload:\n raise ValueError(f\"No fence payload found: cc_pair={cc_pair_id}\")\n\n new_payload = RedisConnectorPrunePayload(\n id=payload.id,\n submitted=payload.submitted,\n started=datetime.now(timezone.utc),\n celery_task_id=payload.celery_task_id,\n )\n redis_connector.prune.set_fence(new_payload)\n\n task_logger.info(\n f\"Pruning generator running connector: cc_pair={cc_pair_id} connector_source={cc_pair.connector.source}\"\n )\n\n runnable_connector = instantiate_connector(\n db_session,\n cc_pair.connector.source,\n InputType.SLIM_RETRIEVAL,\n cc_pair.connector.connector_specific_config,\n cc_pair.credential,\n )\n\n callback = PruneCallback(\n 0,\n redis_connector,\n lock,\n r,\n timeout_seconds=JOB_TIMEOUT,\n )\n\n # Extract docs and hierarchy nodes from the source\n extraction_result = extract_ids_from_runnable_connector(\n runnable_connector, callback\n )\n all_connector_doc_ids = extraction_result.raw_id_to_parent\n\n # Process hierarchy nodes (same as docfetching):\n # upsert to Postgres and cache in Redis\n source = cc_pair.connector.source\n redis_client = get_redis_client(tenant_id=tenant_id)\n\n ensure_source_node_exists(redis_client, db_session, source)\n\n upserted_nodes: list[DBHierarchyNode] = []\n if extraction_result.hierarchy_nodes:\n is_connector_public = cc_pair.access_type == AccessType.PUBLIC\n\n upserted_nodes = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=extraction_result.hierarchy_nodes,\n source=source,\n commit=True,\n is_connector_public=is_connector_public,\n )\n\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=[n.id for n in upserted_nodes],\n connector_id=connector_id,\n credential_id=credential_id,\n commit=True,\n )\n\n cache_entries = [\n HierarchyNodeCacheEntry.from_db_model(node)\n for node in upserted_nodes\n ]\n cache_hierarchy_nodes_batch(\n redis_client=redis_client,\n source=source,\n entries=cache_entries,\n )\n\n task_logger.info(\n f\"Pruning: persisted and cached {len(extraction_result.hierarchy_nodes)} \"\n f\"hierarchy nodes for cc_pair={cc_pair_id}\"\n )\n\n # Resolve parent_hierarchy_raw_node_id → parent_hierarchy_node_id\n # and bulk-update documents, mirroring the docfetching resolution\n _resolve_and_update_document_parents(\n db_session=db_session,\n redis_client=redis_client,\n source=source,\n raw_id_to_parent=all_connector_doc_ids,\n )\n\n # Link hierarchy nodes to documents for sources where pages can be\n # both hierarchy nodes AND documents (e.g. Notion, Confluence)\n all_doc_id_list = list(all_connector_doc_ids.keys())\n link_hierarchy_nodes_to_documents(\n db_session=db_session,\n document_ids=all_doc_id_list,\n source=source,\n commit=True,\n )\n\n # a list of docs in our local index\n all_indexed_document_ids = {\n doc.id\n for doc in get_documents_for_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n )\n }\n\n # generate list of docs to remove (no longer in the source)\n doc_ids_to_remove = list(\n all_indexed_document_ids - all_connector_doc_ids.keys()\n )\n\n task_logger.info(\n \"Pruning set collected: \"\n f\"cc_pair={cc_pair_id} \"\n f\"connector_source={cc_pair.connector.source} \"\n f\"docs_to_remove={len(doc_ids_to_remove)}\"\n )\n\n task_logger.info(\n f\"RedisConnector.prune.generate_tasks starting. cc_pair={cc_pair_id}\"\n )\n tasks_generated = redis_connector.prune.generate_tasks(\n set(doc_ids_to_remove), self.app, db_session, None\n )\n if tasks_generated is None:\n return None\n\n task_logger.info(\n f\"RedisConnector.prune.generate_tasks finished. cc_pair={cc_pair_id} tasks_generated={tasks_generated}\"\n )\n\n redis_connector.prune.generator_complete = tasks_generated\n\n # --- Hierarchy node pruning ---\n live_node_ids = {n.id for n in upserted_nodes}\n stale_removed = remove_stale_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n live_hierarchy_node_ids=live_node_ids,\n commit=True,\n )\n deleted_raw_ids = delete_orphaned_hierarchy_nodes(\n db_session=db_session,\n source=source,\n commit=True,\n )\n reparented_nodes = reparent_orphaned_hierarchy_nodes(\n db_session=db_session,\n source=source,\n commit=True,\n )\n if deleted_raw_ids:\n evict_hierarchy_nodes_from_cache(redis_client, source, deleted_raw_ids)\n if reparented_nodes:\n reparented_cache_entries = [\n HierarchyNodeCacheEntry.from_db_model(node)\n for node in reparented_nodes\n ]\n cache_hierarchy_nodes_batch(\n redis_client, source, reparented_cache_entries\n )\n if stale_removed or deleted_raw_ids or reparented_nodes:\n task_logger.info(\n f\"Hierarchy node pruning: cc_pair={cc_pair_id} \"\n f\"stale_entries_removed={stale_removed} \"\n f\"nodes_deleted={len(deleted_raw_ids)} \"\n f\"nodes_reparented={len(reparented_nodes)}\"\n )\n except Exception as e:\n task_logger.exception(\n f\"Pruning exceptioned: cc_pair={cc_pair_id} connector={connector_id} payload_id={payload_id}\"\n )\n\n redis_connector.prune.reset()\n raise e\n finally:\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"Pruning generator finished: cc_pair={cc_pair_id} payload_id={payload_id}\"\n )\n\n\n\"\"\"Monitoring pruning utils\"\"\"\n\n\ndef monitor_ccpair_pruning_taskset(\n tenant_id: str,\n key_bytes: bytes,\n r: Redis, # noqa: ARG001\n db_session: Session,\n) -> None:\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n task_logger.warning(\n f\"monitor_ccpair_pruning_taskset: could not parse cc_pair_id from {fence_key}\"\n )\n return\n\n cc_pair_id = int(cc_pair_id_str)\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n if not redis_connector.prune.fenced:\n return\n\n initial = redis_connector.prune.generator_complete\n if initial is None:\n return\n\n remaining = redis_connector.prune.get_remaining()\n task_logger.info(\n f\"Connector pruning progress: cc_pair={cc_pair_id} remaining={remaining} initial={initial}\"\n )\n if remaining > 0:\n return\n\n mark_ccpair_as_pruned(int(cc_pair_id), db_session)\n task_logger.info(\n f\"Connector pruning finished: cc_pair={cc_pair_id} num_pruned={initial}\"\n )\n\n update_sync_record_status(\n db_session=db_session,\n entity_id=cc_pair_id,\n sync_type=SyncType.PRUNING,\n sync_status=SyncStatus.SUCCESS,\n num_docs_synced=initial,\n )\n\n delete_orphan_tags__no_commit(db_session)\n\n redis_connector.prune.taskset_clear()\n redis_connector.prune.generator_clear()\n redis_connector.prune.set_fence(None)\n\n\ndef validate_pruning_fences(\n tenant_id: str,\n r: Redis,\n r_replica: Redis,\n r_celery: Redis,\n lock_beat: RedisLock,\n) -> None:\n # building lookup table can be expensive, so we won't bother\n # validating until the queue is small\n PERMISSION_SYNC_VALIDATION_MAX_QUEUE_LEN = 1024\n\n queue_len = celery_get_queue_length(OnyxCeleryQueues.CONNECTOR_DELETION, r_celery)\n if queue_len > PERMISSION_SYNC_VALIDATION_MAX_QUEUE_LEN:\n return\n\n # the queue for a single pruning generator task\n reserved_generator_tasks = celery_get_unacked_task_ids(\n OnyxCeleryQueues.CONNECTOR_PRUNING, r_celery\n )\n\n # the queue for a reasonably large set of lightweight deletion tasks\n queued_upsert_tasks = celery_get_queued_task_ids(\n OnyxCeleryQueues.CONNECTOR_DELETION, r_celery\n )\n\n # Use replica for this because the worst thing that happens\n # is that we don't run the validation on this pass\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n key_str = key_bytes.decode(\"utf-8\")\n if not key_str.startswith(RedisConnectorPrune.FENCE_PREFIX):\n continue\n\n validate_pruning_fence(\n tenant_id,\n key_bytes,\n reserved_generator_tasks,\n queued_upsert_tasks,\n r,\n r_celery,\n )\n\n lock_beat.reacquire()\n\n return\n\n\ndef validate_pruning_fence(\n tenant_id: str,\n key_bytes: bytes,\n reserved_tasks: set[str],\n queued_tasks: set[str],\n r: Redis,\n r_celery: Redis,\n) -> None:\n \"\"\"See validate_indexing_fence for an overall idea of validation flows.\n\n queued_tasks: the celery queue of lightweight permission sync tasks\n reserved_tasks: prefetched tasks for sync task generator\n \"\"\"\n # if the fence doesn't exist, there's nothing to do\n fence_key = key_bytes.decode(\"utf-8\")\n cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)\n if cc_pair_id_str is None:\n task_logger.warning(\n f\"validate_pruning_fence - could not parse id from {fence_key}\"\n )\n return\n\n cc_pair_id = int(cc_pair_id_str)\n # parse out metadata and initialize the helper class with it\n redis_connector = RedisConnector(tenant_id, int(cc_pair_id))\n\n # check to see if the fence/payload exists\n if not redis_connector.prune.fenced:\n return\n\n # in the cloud, the payload format may have changed ...\n # it's a little sloppy, but just reset the fence for now if that happens\n # TODO: add intentional cleanup/abort logic\n try:\n payload = redis_connector.prune.payload\n except ValidationError:\n task_logger.exception(\n \"validate_pruning_fence - \"\n \"Resetting fence because fence schema is out of date: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key}\"\n )\n\n redis_connector.prune.reset()\n return\n\n if not payload:\n return\n\n if not payload.celery_task_id:\n return\n\n # OK, there's actually something for us to validate\n\n # either the generator task must be in flight or its subtasks must be\n found = celery_find_task(\n payload.celery_task_id,\n OnyxCeleryQueues.CONNECTOR_PRUNING,\n r_celery,\n )\n if found:\n # the celery task exists in the redis queue\n redis_connector.prune.set_active()\n return\n\n if payload.celery_task_id in reserved_tasks:\n # the celery task was prefetched and is reserved within a worker\n redis_connector.prune.set_active()\n return\n\n # look up every task in the current taskset in the celery queue\n # every entry in the taskset should have an associated entry in the celery task queue\n # because we get the celery tasks first, the entries in our own pruning taskset\n # should be roughly a subset of the tasks in celery\n\n # this check isn't very exact, but should be sufficient over a period of time\n # A single successful check over some number of attempts is sufficient.\n\n # TODO: if the number of tasks in celery is much lower than than the taskset length\n # we might be able to shortcut the lookup since by definition some of the tasks\n # must not exist in celery.\n\n tasks_scanned = 0\n tasks_not_in_celery = 0 # a non-zero number after completing our check is bad\n\n for member in r.sscan_iter(redis_connector.prune.taskset_key):\n tasks_scanned += 1\n\n member_bytes = cast(bytes, member)\n member_str = member_bytes.decode(\"utf-8\")\n if member_str in queued_tasks:\n continue\n\n if member_str in reserved_tasks:\n continue\n\n tasks_not_in_celery += 1\n\n task_logger.info(\n f\"validate_pruning_fence task check: tasks_scanned={tasks_scanned} tasks_not_in_celery={tasks_not_in_celery}\"\n )\n\n # we're active if there are still tasks to run and those tasks all exist in celery\n if tasks_scanned > 0 and tasks_not_in_celery == 0:\n redis_connector.prune.set_active()\n return\n\n # we may want to enable this check if using the active task list somehow isn't good enough\n # if redis_connector_index.generator_locked():\n # logger.info(f\"{payload.celery_task_id} is currently executing.\")\n\n # if we get here, we didn't find any direct indication that the associated celery tasks exist,\n # but they still might be there due to gaps in our ability to check states during transitions\n # Checking the active signal safeguards us against these transition periods\n # (which has a duration that allows us to bridge those gaps)\n if redis_connector.prune.active():\n return\n\n # celery tasks don't exist and the active signal has expired, possibly due to a crash. Clean it up.\n task_logger.warning(\n \"validate_pruning_fence - \"\n \"Resetting fence because no associated celery tasks were found: \"\n f\"cc_pair={cc_pair_id} \"\n f\"fence={fence_key} \"\n f\"payload_id={payload.id}\"\n )\n\n redis_connector.prune.reset()\n return\n" }, { "path": "backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py", "content": "import httpx\nfrom tenacity import retry\nfrom tenacity import retry_if_exception_type\nfrom tenacity import stop_after_delay\nfrom tenacity import wait_random_exponential\n\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\n\n\nclass RetryDocumentIndex:\n \"\"\"A wrapper class to help with specific retries against Vespa involving\n read timeouts.\n\n wait_random_exponential implements full jitter as per this article:\n https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/\"\"\"\n\n MAX_WAIT = 30\n\n # STOP_AFTER + MAX_WAIT should be slightly less (5?) than the celery soft_time_limit\n STOP_AFTER = 70\n\n def __init__(self, index: DocumentIndex):\n self.index: DocumentIndex = index\n\n @retry(\n retry=retry_if_exception_type(httpx.ReadTimeout),\n wait=wait_random_exponential(multiplier=1, max=MAX_WAIT),\n stop=stop_after_delay(STOP_AFTER),\n )\n def delete_single(\n self,\n doc_id: str,\n *,\n tenant_id: str,\n chunk_count: int | None,\n ) -> int:\n return self.index.delete_single(\n doc_id,\n tenant_id=tenant_id,\n chunk_count=chunk_count,\n )\n\n @retry(\n retry=retry_if_exception_type(httpx.ReadTimeout),\n wait=wait_random_exponential(multiplier=1, max=MAX_WAIT),\n stop=stop_after_delay(STOP_AFTER),\n )\n def update_single(\n self,\n doc_id: str,\n *,\n tenant_id: str,\n chunk_count: int | None,\n fields: VespaDocumentFields | None,\n user_fields: VespaDocumentUserFields | None,\n ) -> None:\n self.index.update_single(\n doc_id,\n tenant_id=tenant_id,\n chunk_count=chunk_count,\n fields=fields,\n user_fields=user_fields,\n )\n" }, { "path": "backend/onyx/background/celery/tasks/shared/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/shared/tasks.py", "content": "import time\nfrom enum import Enum\nfrom http import HTTPStatus\n\nimport httpx\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom redis import Redis\nfrom tenacity import RetryError\n\nfrom onyx.access.access import get_access_for_document\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex\nfrom onyx.configs.constants import ONYX_CELERY_BEAT_HEARTBEAT_KEY\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.document import delete_document_by_connector_credential_pair__no_commit\nfrom onyx.db.document import delete_documents_complete__no_commit\nfrom onyx.db.document import fetch_chunk_count_for_document\nfrom onyx.db.document import get_document\nfrom onyx.db.document import get_document_connector_count\nfrom onyx.db.document import mark_document_as_modified\nfrom onyx.db.document import mark_document_as_synced\nfrom onyx.db.document_set import fetch_document_sets_for_document\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.relationships import delete_document_references_from_kg\nfrom onyx.db.search_settings import get_active_search_settings\nfrom onyx.document_index.factory import get_all_document_indices\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.httpx.httpx_pool import HttpxPool\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import ConnectorCredentialPairIdentifier\n\nDOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES = 3\n\n\n# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT\nLIGHT_SOFT_TIME_LIMIT = 105\nLIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15\n\n\nclass OnyxCeleryTaskCompletionStatus(str, Enum):\n \"\"\"The different statuses the watchdog can finish with.\n\n TODO: create broader success/failure/abort categories\n \"\"\"\n\n UNDEFINED = \"undefined\"\n\n SUCCEEDED = \"succeeded\"\n\n SKIPPED = \"skipped\"\n\n SOFT_TIME_LIMIT = \"soft_time_limit\"\n\n NON_RETRYABLE_EXCEPTION = \"non_retryable_exception\"\n RETRYABLE_EXCEPTION = \"retryable_exception\"\n\n\n@shared_task(\n name=OnyxCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_TASK,\n soft_time_limit=LIGHT_SOFT_TIME_LIMIT,\n time_limit=LIGHT_TIME_LIMIT,\n max_retries=DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES,\n bind=True,\n)\ndef document_by_cc_pair_cleanup_task(\n self: Task,\n document_id: str,\n connector_id: int,\n credential_id: int,\n tenant_id: str,\n) -> bool:\n \"\"\"A lightweight subtask used to clean up document to cc pair relationships.\n Created by connection deletion and connector pruning parent tasks.\"\"\"\n\n \"\"\"\n To delete a connector / credential pair:\n (1) find all documents associated with connector / credential pair where there\n this the is only connector / credential pair that has indexed it\n (2) delete all documents from document stores\n (3) delete all entries from postgres\n (4) find all documents associated with connector / credential pair where there\n are multiple connector / credential pairs that have indexed it\n (5) update document store entries to remove access associated with the\n connector / credential pair from the access list\n (6) delete all relevant entries from postgres\n \"\"\"\n task_logger.debug(f\"Task start: doc={document_id}\")\n\n start = time.monotonic()\n\n completion_status = OnyxCeleryTaskCompletionStatus.UNDEFINED\n\n try:\n with get_session_with_current_tenant() as db_session:\n action = \"skip\"\n\n active_search_settings = get_active_search_settings(db_session)\n # This flow is for updates and deletion so we get all indices.\n document_indices = get_all_document_indices(\n active_search_settings.primary,\n active_search_settings.secondary,\n httpx_client=HttpxPool.get(\"vespa\"),\n )\n\n retry_document_indices: list[RetryDocumentIndex] = [\n RetryDocumentIndex(document_index)\n for document_index in document_indices\n ]\n\n count = get_document_connector_count(db_session, document_id)\n if count == 1:\n # count == 1 means this is the only remaining cc_pair reference to the doc\n # delete it from vespa and the db\n action = \"delete\"\n\n chunk_count = fetch_chunk_count_for_document(document_id, db_session)\n\n for retry_document_index in retry_document_indices:\n _ = retry_document_index.delete_single(\n document_id,\n tenant_id=tenant_id,\n chunk_count=chunk_count,\n )\n\n delete_document_references_from_kg(\n db_session=db_session,\n document_id=document_id,\n )\n\n delete_documents_complete__no_commit(\n db_session=db_session,\n document_ids=[document_id],\n )\n db_session.commit()\n\n completion_status = OnyxCeleryTaskCompletionStatus.SUCCEEDED\n elif count > 1:\n action = \"update\"\n\n # count > 1 means the document still has cc_pair references\n doc = get_document(document_id, db_session)\n if not doc:\n return False\n\n # the below functions do not include cc_pairs being deleted.\n # i.e. they will correctly omit access for the current cc_pair\n doc_access = get_access_for_document(\n document_id=document_id, db_session=db_session\n )\n\n doc_sets = fetch_document_sets_for_document(document_id, db_session)\n update_doc_sets: set[str] = set(doc_sets)\n\n fields = VespaDocumentFields(\n document_sets=update_doc_sets,\n access=doc_access,\n boost=doc.boost,\n hidden=doc.hidden,\n )\n\n for retry_document_index in retry_document_indices:\n # TODO(andrei): Previously there was a comment here saying\n # it was ok if a doc did not exist in the document index. I\n # don't agree with that claim, so keep an eye on this task\n # to see if this raises.\n retry_document_index.update_single(\n document_id,\n tenant_id=tenant_id,\n chunk_count=doc.chunk_count,\n fields=fields,\n user_fields=None,\n )\n\n # there are still other cc_pair references to the doc, so just resync to Vespa\n delete_document_by_connector_credential_pair__no_commit(\n db_session=db_session,\n document_id=document_id,\n connector_credential_pair_identifier=ConnectorCredentialPairIdentifier(\n connector_id=connector_id,\n credential_id=credential_id,\n ),\n )\n\n mark_document_as_synced(document_id, db_session)\n db_session.commit()\n\n completion_status = OnyxCeleryTaskCompletionStatus.SUCCEEDED\n else:\n completion_status = OnyxCeleryTaskCompletionStatus.SKIPPED\n\n elapsed = time.monotonic() - start\n task_logger.info(\n f\"doc={document_id} action={action} refcount={count} elapsed={elapsed:.2f}\"\n )\n except SoftTimeLimitExceeded:\n task_logger.info(f\"SoftTimeLimitExceeded exception. doc={document_id}\")\n completion_status = OnyxCeleryTaskCompletionStatus.SOFT_TIME_LIMIT\n except Exception as ex:\n e: Exception | None = None\n while True:\n if isinstance(ex, RetryError):\n task_logger.warning(\n f\"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}\"\n )\n\n # only set the inner exception if it is of type Exception\n e_temp = ex.last_attempt.exception()\n if isinstance(e_temp, Exception):\n e = e_temp\n else:\n e = ex\n\n if isinstance(e, httpx.HTTPStatusError):\n if e.response.status_code == HTTPStatus.BAD_REQUEST:\n task_logger.exception(\n f\"Non-retryable HTTPStatusError: doc={document_id} status={e.response.status_code}\"\n )\n completion_status = (\n OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION\n )\n break\n\n task_logger.exception(\n f\"document_by_cc_pair_cleanup_task exceptioned: doc={document_id}\"\n )\n\n completion_status = OnyxCeleryTaskCompletionStatus.RETRYABLE_EXCEPTION\n if (\n self.max_retries is not None\n and self.request.retries >= self.max_retries\n ):\n # This is the last attempt! mark the document as dirty in the db so that it\n # eventually gets fixed out of band via stale document reconciliation\n task_logger.warning(\n f\"Max celery task retries reached. Marking doc as dirty for reconciliation: doc={document_id}\"\n )\n with get_session_with_current_tenant() as db_session:\n # delete the cc pair relationship now and let reconciliation clean it up\n # in vespa\n delete_document_by_connector_credential_pair__no_commit(\n db_session=db_session,\n document_id=document_id,\n connector_credential_pair_identifier=ConnectorCredentialPairIdentifier(\n connector_id=connector_id,\n credential_id=credential_id,\n ),\n )\n mark_document_as_modified(document_id, db_session)\n completion_status = (\n OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION\n )\n break\n\n # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64\n countdown = 2 ** (self.request.retries + 4)\n self.retry(exc=e, countdown=countdown) # this will raise a celery exception\n break # we won't hit this, but it looks weird not to have it\n finally:\n task_logger.info(\n f\"document_by_cc_pair_cleanup_task completed: status={completion_status.value} doc={document_id}\"\n )\n\n if completion_status != OnyxCeleryTaskCompletionStatus.SUCCEEDED:\n return False\n\n task_logger.info(f\"document_by_cc_pair_cleanup_task finished: doc={document_id}\")\n return True\n\n\n@shared_task(name=OnyxCeleryTask.CELERY_BEAT_HEARTBEAT, ignore_result=True, bind=True)\ndef celery_beat_heartbeat(self: Task, *, tenant_id: str) -> None: # noqa: ARG001\n \"\"\"When this task runs, it writes a key to Redis with a TTL.\n\n An external observer can check this key to figure out if the celery beat is still running.\n \"\"\"\n time_start = time.monotonic()\n r: Redis = get_redis_client()\n r.set(ONYX_CELERY_BEAT_HEARTBEAT_KEY, 1, ex=600)\n time_elapsed = time.monotonic() - time_start\n task_logger.info(f\"celery_beat_heartbeat finished: elapsed={time_elapsed:.2f}\")\n" }, { "path": "backend/onyx/background/celery/tasks/user_file_processing/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/user_file_processing/tasks.py", "content": "import datetime\nimport time\nfrom typing import Any\nfrom uuid import UUID\n\nimport httpx\nimport sqlalchemy as sa\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom retry import retry\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.access import build_access_for_user_files\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.celery_redis import celery_get_broker_client\nfrom onyx.background.celery.celery_redis import celery_get_queue_length\nfrom onyx.background.celery.celery_utils import httpx_init_vespa_pool\nfrom onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import MANAGED_VESPA\nfrom onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH\nfrom onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_USER_FILE_DELETE_TASK_EXPIRES\nfrom onyx.configs.constants import CELERY_USER_FILE_PROCESSING_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES\nfrom onyx.configs.constants import CELERY_USER_FILE_PROJECT_SYNC_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.constants import USER_FILE_DELETE_MAX_QUEUE_DEPTH\nfrom onyx.configs.constants import USER_FILE_PROCESSING_MAX_QUEUE_DEPTH\nfrom onyx.configs.constants import USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH\nfrom onyx.connectors.file.connector import LocalFileConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import UserFile\nfrom onyx.db.search_settings import get_active_search_settings\nfrom onyx.db.search_settings import get_active_search_settings_list\nfrom onyx.db.user_file import fetch_user_files_with_access_relationships\nfrom onyx.document_index.factory import get_all_document_indices\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.file_store.utils import store_user_file_plaintext\nfrom onyx.file_store.utils import user_file_id_to_plaintext_file_name\nfrom onyx.httpx.httpx_pool import HttpxPool\nfrom onyx.indexing.adapters.user_file_indexing_adapter import UserFileIndexingAdapter\nfrom onyx.indexing.embedder import DefaultIndexingEmbedder\nfrom onyx.indexing.indexing_pipeline import run_indexing_pipeline\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.variable_functionality import global_version\n\n\ndef _as_uuid(value: str | UUID) -> UUID:\n \"\"\"Return a UUID, accepting either a UUID or a string-like value.\"\"\"\n return value if isinstance(value, UUID) else UUID(str(value))\n\n\ndef _user_file_lock_key(user_file_id: str | UUID) -> str:\n return f\"{OnyxRedisLocks.USER_FILE_PROCESSING_LOCK_PREFIX}:{user_file_id}\"\n\n\ndef _user_file_queued_key(user_file_id: str | UUID) -> str:\n \"\"\"Key that exists while a process_single_user_file task is sitting in the queue.\n\n The beat generator sets this with a TTL equal to CELERY_USER_FILE_PROCESSING_TASK_EXPIRES\n before enqueuing and the worker deletes it as its first action. This prevents\n the beat from adding duplicate tasks for files that already have a live task\n in flight.\n \"\"\"\n return f\"{OnyxRedisLocks.USER_FILE_QUEUED_PREFIX}:{user_file_id}\"\n\n\ndef user_file_project_sync_lock_key(user_file_id: str | UUID) -> str:\n return f\"{OnyxRedisLocks.USER_FILE_PROJECT_SYNC_LOCK_PREFIX}:{user_file_id}\"\n\n\ndef _user_file_project_sync_queued_key(user_file_id: str | UUID) -> str:\n return f\"{OnyxRedisLocks.USER_FILE_PROJECT_SYNC_QUEUED_PREFIX}:{user_file_id}\"\n\n\ndef _user_file_delete_lock_key(user_file_id: str | UUID) -> str:\n return f\"{OnyxRedisLocks.USER_FILE_DELETE_LOCK_PREFIX}:{user_file_id}\"\n\n\ndef _user_file_delete_queued_key(user_file_id: str | UUID) -> str:\n \"\"\"Key that exists while a delete_single_user_file task is sitting in the queue.\n\n The beat generator sets this with a TTL equal to CELERY_USER_FILE_DELETE_TASK_EXPIRES\n before enqueuing and the worker deletes it as its first action. This prevents\n the beat from adding duplicate tasks for files that already have a live task\n in flight.\n \"\"\"\n return f\"{OnyxRedisLocks.USER_FILE_DELETE_QUEUED_PREFIX}:{user_file_id}\"\n\n\ndef get_user_file_project_sync_queue_depth(celery_app: Celery) -> int:\n redis_celery = celery_get_broker_client(celery_app)\n return celery_get_queue_length(\n OnyxCeleryQueues.USER_FILE_PROJECT_SYNC, redis_celery\n )\n\n\ndef enqueue_user_file_project_sync_task(\n *,\n celery_app: Celery,\n redis_client: Redis,\n user_file_id: str | UUID,\n tenant_id: str,\n priority: OnyxCeleryPriority = OnyxCeleryPriority.HIGH,\n) -> bool:\n \"\"\"Enqueue a project-sync task if no matching queued task already exists.\"\"\"\n queued_key = _user_file_project_sync_queued_key(user_file_id)\n\n # NX+EX gives us atomic dedupe and a self-healing TTL.\n queued_guard_set = redis_client.set(\n queued_key,\n 1,\n nx=True,\n ex=CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES,\n )\n if not queued_guard_set:\n return False\n\n try:\n celery_app.send_task(\n OnyxCeleryTask.PROCESS_SINGLE_USER_FILE_PROJECT_SYNC,\n kwargs={\"user_file_id\": str(user_file_id), \"tenant_id\": tenant_id},\n queue=OnyxCeleryQueues.USER_FILE_PROJECT_SYNC,\n priority=priority,\n expires=CELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES,\n )\n except Exception:\n # Roll back the queued guard if task publish fails.\n redis_client.delete(queued_key)\n raise\n\n return True\n\n\n@retry(tries=3, delay=1, backoff=2, jitter=(0.0, 1.0))\ndef _visit_chunks(\n *,\n http_client: httpx.Client,\n index_name: str,\n selection: str,\n continuation: str | None = None,\n) -> tuple[list[dict[str, Any]], str | None]:\n task_logger.info(\n f\"Visiting chunks for index={index_name} with selection={selection}\"\n )\n base_url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)\n params: dict[str, str] = {\n \"selection\": selection,\n \"wantedDocumentCount\": \"100\", # Use smaller batch size to avoid timeouts\n }\n if continuation:\n params[\"continuation\"] = continuation\n resp = http_client.get(base_url, params=params, timeout=None)\n resp.raise_for_status()\n payload = resp.json()\n return payload.get(\"documents\", []), payload.get(\"continuation\")\n\n\ndef _get_document_chunk_count(\n *,\n index_name: str,\n selection: str,\n) -> int:\n chunk_count = 0\n continuation = None\n while True:\n docs, continuation = _visit_chunks(\n http_client=HttpxPool.get(\"vespa\"),\n index_name=index_name,\n selection=selection,\n continuation=continuation,\n )\n if not docs:\n break\n chunk_count += len(docs)\n if not continuation:\n break\n return chunk_count\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_USER_FILE_PROCESSING,\n soft_time_limit=300,\n bind=True,\n ignore_result=True,\n)\ndef check_user_file_processing(self: Task, *, tenant_id: str) -> None:\n \"\"\"Scan for user files with PROCESSING status and enqueue per-file tasks.\n\n Three mechanisms prevent queue runaway:\n\n 1. **Queue depth backpressure** – if the broker queue already has more than\n USER_FILE_PROCESSING_MAX_QUEUE_DEPTH items we skip this beat cycle\n entirely. Workers are clearly behind; adding more tasks would only make\n the backlog worse.\n\n 2. **Per-file queued guard** – before enqueuing a task we set a short-lived\n Redis key (TTL = CELERY_USER_FILE_PROCESSING_TASK_EXPIRES). If that key\n already exists the file already has a live task in the queue, so we skip\n it. The worker deletes the key the moment it picks up the task so the\n next beat cycle can re-enqueue if the file is still PROCESSING.\n\n 3. **Task expiry** – every enqueued task carries an `expires` value equal to\n CELERY_USER_FILE_PROCESSING_TASK_EXPIRES. If a task is still sitting in\n the queue after that deadline, Celery discards it without touching the DB.\n This is a belt-and-suspenders defence: even if the guard key is lost (e.g.\n Redis restart), stale tasks evict themselves rather than piling up forever.\n \"\"\"\n task_logger.info(\"check_user_file_processing - Starting\")\n\n redis_client = get_redis_client(tenant_id=tenant_id)\n lock: RedisLock = redis_client.lock(\n OnyxRedisLocks.USER_FILE_PROCESSING_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n # Do not overlap generator runs\n if not lock.acquire(blocking=False):\n return None\n\n enqueued = 0\n skipped_guard = 0\n try:\n # --- Protection 1: queue depth backpressure ---\n r_celery = celery_get_broker_client(self.app)\n queue_len = celery_get_queue_length(\n OnyxCeleryQueues.USER_FILE_PROCESSING, r_celery\n )\n if queue_len > USER_FILE_PROCESSING_MAX_QUEUE_DEPTH:\n task_logger.warning(\n f\"check_user_file_processing - Queue depth {queue_len} exceeds \"\n f\"{USER_FILE_PROCESSING_MAX_QUEUE_DEPTH}, skipping enqueue for \"\n f\"tenant={tenant_id}\"\n )\n return None\n\n with get_session_with_current_tenant() as db_session:\n user_file_ids = (\n db_session.execute(\n select(UserFile.id).where(\n UserFile.status == UserFileStatus.PROCESSING\n )\n )\n .scalars()\n .all()\n )\n\n for user_file_id in user_file_ids:\n # --- Protection 2: per-file queued guard ---\n queued_key = _user_file_queued_key(user_file_id)\n guard_set = redis_client.set(\n queued_key,\n 1,\n ex=CELERY_USER_FILE_PROCESSING_TASK_EXPIRES,\n nx=True,\n )\n if not guard_set:\n skipped_guard += 1\n continue\n\n # --- Protection 3: task expiry ---\n # If task submission fails, clear the guard immediately so the\n # next beat cycle can retry enqueuing this file.\n try:\n self.app.send_task(\n OnyxCeleryTask.PROCESS_SINGLE_USER_FILE,\n kwargs={\n \"user_file_id\": str(user_file_id),\n \"tenant_id\": tenant_id,\n },\n queue=OnyxCeleryQueues.USER_FILE_PROCESSING,\n priority=OnyxCeleryPriority.HIGH,\n expires=CELERY_USER_FILE_PROCESSING_TASK_EXPIRES,\n )\n except Exception:\n redis_client.delete(queued_key)\n raise\n enqueued += 1\n\n finally:\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"check_user_file_processing - Enqueued {enqueued} skipped_guard={skipped_guard} tasks for tenant={tenant_id}\"\n )\n return None\n\n\ndef _process_user_file_without_vector_db(\n uf: UserFile,\n documents: list[Document],\n db_session: Session,\n) -> None:\n \"\"\"Process a user file when the vector DB is disabled.\n\n Extracts raw text and computes a token count, stores the plaintext in\n the file store, and marks the file as COMPLETED. Skips embedding and\n the indexing pipeline entirely.\n \"\"\"\n from onyx.llm.factory import get_default_llm\n from onyx.llm.factory import get_llm_tokenizer_encode_func\n\n # Combine section text from all document sections\n combined_text = \" \".join(\n section.text for doc in documents for section in doc.sections if section.text\n )\n\n # Compute token count using the user's default LLM tokenizer\n try:\n llm = get_default_llm()\n encode = get_llm_tokenizer_encode_func(llm)\n token_count: int | None = len(encode(combined_text))\n except Exception:\n task_logger.warning(\n f\"_process_user_file_without_vector_db - Failed to compute token count for {uf.id}, falling back to None\"\n )\n token_count = None\n\n # Persist plaintext for fast FileReaderTool loads\n store_user_file_plaintext(\n user_file_id=uf.id,\n plaintext_content=combined_text,\n )\n\n # Update the DB record\n if uf.status != UserFileStatus.DELETING:\n uf.status = UserFileStatus.COMPLETED\n uf.token_count = token_count\n uf.chunk_count = 0 # no chunks without vector DB\n uf.last_project_sync_at = datetime.datetime.now(datetime.timezone.utc)\n db_session.add(uf)\n db_session.commit()\n\n task_logger.info(\n f\"_process_user_file_without_vector_db - Completed id={uf.id} tokens={token_count}\"\n )\n\n\ndef _process_user_file_with_indexing(\n uf: UserFile,\n user_file_id: str,\n documents: list[Document],\n tenant_id: str,\n db_session: Session,\n) -> None:\n \"\"\"Process a user file through the full indexing pipeline (vector DB path).\"\"\"\n # 20 is the documented default for httpx max_keepalive_connections\n if MANAGED_VESPA:\n httpx_init_vespa_pool(\n 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH\n )\n else:\n httpx_init_vespa_pool(20)\n\n search_settings_list = get_active_search_settings_list(db_session)\n current_search_settings = next(\n (ss for ss in search_settings_list if ss.status.is_current()),\n None,\n )\n if current_search_settings is None:\n raise RuntimeError(\n f\"_process_user_file_with_indexing - No current search settings found for tenant={tenant_id}\"\n )\n\n adapter = UserFileIndexingAdapter(\n tenant_id=tenant_id,\n db_session=db_session,\n )\n\n embedding_model = DefaultIndexingEmbedder.from_db_search_settings(\n search_settings=current_search_settings,\n )\n\n document_indices = get_all_document_indices(\n current_search_settings,\n None,\n httpx_client=HttpxPool.get(\"vespa\"),\n )\n\n index_pipeline_result = run_indexing_pipeline(\n embedder=embedding_model,\n document_indices=document_indices,\n ignore_time_skip=True,\n db_session=db_session,\n tenant_id=tenant_id,\n document_batch=documents,\n request_id=None,\n adapter=adapter,\n )\n\n task_logger.info(\n f\"_process_user_file_with_indexing - Indexing pipeline completed ={index_pipeline_result}\"\n )\n\n if (\n index_pipeline_result.failures\n or index_pipeline_result.total_docs != len(documents)\n or index_pipeline_result.total_chunks == 0\n ):\n task_logger.error(\n f\"_process_user_file_with_indexing - Indexing pipeline failed id={user_file_id}\"\n )\n if uf.status != UserFileStatus.DELETING:\n uf.status = UserFileStatus.FAILED\n db_session.add(uf)\n db_session.commit()\n raise RuntimeError(f\"Indexing pipeline failed for user file {user_file_id}\")\n\n\ndef process_user_file_impl(\n *, user_file_id: str, tenant_id: str, redis_locking: bool\n) -> None:\n \"\"\"Core implementation for processing a single user file.\n\n When redis_locking=True, acquires a per-file Redis lock and clears the\n queued-key guard (Celery path). When redis_locking=False, skips all Redis\n operations (BackgroundTask path).\n \"\"\"\n task_logger.info(f\"process_user_file_impl - Starting id={user_file_id}\")\n start = time.monotonic()\n\n file_lock: RedisLock | None = None\n if redis_locking:\n redis_client = get_redis_client(tenant_id=tenant_id)\n redis_client.delete(_user_file_queued_key(user_file_id))\n file_lock = redis_client.lock(\n _user_file_lock_key(user_file_id),\n timeout=CELERY_USER_FILE_PROCESSING_LOCK_TIMEOUT,\n )\n if file_lock is not None and not file_lock.acquire(blocking=False):\n task_logger.info(\n f\"process_user_file_impl - Lock held, skipping user_file_id={user_file_id}\"\n )\n return\n\n documents: list[Document] = []\n try:\n with get_session_with_current_tenant() as db_session:\n uf = db_session.get(UserFile, _as_uuid(user_file_id))\n if not uf:\n task_logger.warning(\n f\"process_user_file_impl - UserFile not found id={user_file_id}\"\n )\n return\n\n if uf.status not in (\n UserFileStatus.PROCESSING,\n UserFileStatus.INDEXING,\n ):\n task_logger.info(\n f\"process_user_file_impl - Skipping id={user_file_id} status={uf.status}\"\n )\n return\n\n connector = LocalFileConnector(\n file_locations=[uf.file_id],\n file_names=[uf.name] if uf.name else None,\n )\n connector.load_credentials({})\n\n try:\n for batch in connector.load_from_state():\n documents.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n\n for document in documents:\n document.id = str(user_file_id)\n document.source = DocumentSource.USER_FILE\n\n if DISABLE_VECTOR_DB:\n _process_user_file_without_vector_db(\n uf=uf,\n documents=documents,\n db_session=db_session,\n )\n else:\n _process_user_file_with_indexing(\n uf=uf,\n user_file_id=user_file_id,\n documents=documents,\n tenant_id=tenant_id,\n db_session=db_session,\n )\n\n except Exception as e:\n task_logger.exception(\n f\"process_user_file_impl - Error processing file id={user_file_id} - {e.__class__.__name__}\"\n )\n current_user_file = db_session.get(UserFile, _as_uuid(user_file_id))\n if (\n current_user_file\n and current_user_file.status != UserFileStatus.DELETING\n ):\n uf.status = UserFileStatus.FAILED\n db_session.add(uf)\n db_session.commit()\n return\n\n elapsed = time.monotonic() - start\n task_logger.info(\n f\"process_user_file_impl - Finished id={user_file_id} docs={len(documents)} elapsed={elapsed:.2f}s\"\n )\n except Exception as e:\n with get_session_with_current_tenant() as db_session:\n uf = db_session.get(UserFile, _as_uuid(user_file_id))\n if uf:\n if uf.status != UserFileStatus.DELETING:\n uf.status = UserFileStatus.FAILED\n db_session.add(uf)\n db_session.commit()\n\n task_logger.exception(\n f\"process_user_file_impl - Error processing file id={user_file_id} - {e.__class__.__name__}\"\n )\n raise\n finally:\n if file_lock is not None and file_lock.owned():\n file_lock.release()\n\n\n@shared_task(\n name=OnyxCeleryTask.PROCESS_SINGLE_USER_FILE,\n bind=True,\n ignore_result=True,\n)\ndef process_single_user_file(\n self: Task, # noqa: ARG001\n *,\n user_file_id: str,\n tenant_id: str,\n) -> None:\n process_user_file_impl(\n user_file_id=user_file_id, tenant_id=tenant_id, redis_locking=True\n )\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_USER_FILE_DELETE,\n soft_time_limit=300,\n bind=True,\n ignore_result=True,\n)\ndef check_for_user_file_delete(self: Task, *, tenant_id: str) -> None:\n \"\"\"Scan for user files with DELETING status and enqueue per-file tasks.\n\n Three mechanisms prevent queue runaway (mirrors check_user_file_processing):\n\n 1. **Queue depth backpressure** – if the broker queue already has more than\n USER_FILE_DELETE_MAX_QUEUE_DEPTH items we skip this beat cycle entirely.\n\n 2. **Per-file queued guard** – before enqueuing a task we set a short-lived\n Redis key (TTL = CELERY_USER_FILE_DELETE_TASK_EXPIRES). If that key\n already exists the file already has a live task in the queue, so we skip\n it. The worker deletes the key the moment it picks up the task so the\n next beat cycle can re-enqueue if the file is still DELETING.\n\n 3. **Task expiry** – every enqueued task carries an `expires` value equal to\n CELERY_USER_FILE_DELETE_TASK_EXPIRES. If a task is still sitting in\n the queue after that deadline, Celery discards it without touching the DB.\n \"\"\"\n task_logger.info(\"check_for_user_file_delete - Starting\")\n redis_client = get_redis_client(tenant_id=tenant_id)\n lock: RedisLock = redis_client.lock(\n OnyxRedisLocks.USER_FILE_DELETE_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n if not lock.acquire(blocking=False):\n return None\n\n enqueued = 0\n skipped_guard = 0\n try:\n # --- Protection 1: queue depth backpressure ---\n # NOTE: must use the broker's Redis client (not redis_client) because\n # Celery queues live on a separate Redis DB with CELERY_SEPARATOR keys.\n r_celery = celery_get_broker_client(self.app)\n queue_len = celery_get_queue_length(OnyxCeleryQueues.USER_FILE_DELETE, r_celery)\n if queue_len > USER_FILE_DELETE_MAX_QUEUE_DEPTH:\n task_logger.warning(\n f\"check_for_user_file_delete - Queue depth {queue_len} exceeds \"\n f\"{USER_FILE_DELETE_MAX_QUEUE_DEPTH}, skipping enqueue for \"\n f\"tenant={tenant_id}\"\n )\n return None\n\n with get_session_with_current_tenant() as db_session:\n user_file_ids = (\n db_session.execute(\n select(UserFile.id).where(\n UserFile.status == UserFileStatus.DELETING\n )\n )\n .scalars()\n .all()\n )\n for user_file_id in user_file_ids:\n # --- Protection 2: per-file queued guard ---\n queued_key = _user_file_delete_queued_key(user_file_id)\n guard_set = redis_client.set(\n queued_key,\n 1,\n ex=CELERY_USER_FILE_DELETE_TASK_EXPIRES,\n nx=True,\n )\n if not guard_set:\n skipped_guard += 1\n continue\n\n # --- Protection 3: task expiry ---\n try:\n self.app.send_task(\n OnyxCeleryTask.DELETE_SINGLE_USER_FILE,\n kwargs={\n \"user_file_id\": str(user_file_id),\n \"tenant_id\": tenant_id,\n },\n queue=OnyxCeleryQueues.USER_FILE_DELETE,\n priority=OnyxCeleryPriority.HIGH,\n expires=CELERY_USER_FILE_DELETE_TASK_EXPIRES,\n )\n except Exception:\n redis_client.delete(queued_key)\n raise\n enqueued += 1\n finally:\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"check_for_user_file_delete - Enqueued {enqueued} tasks, skipped_guard={skipped_guard} for tenant={tenant_id}\"\n )\n return None\n\n\ndef delete_user_file_impl(\n *, user_file_id: str, tenant_id: str, redis_locking: bool\n) -> None:\n \"\"\"Core implementation for deleting a single user file.\n\n When redis_locking=True, acquires a per-file Redis lock (Celery path).\n When redis_locking=False, skips Redis operations (BackgroundTask path).\n \"\"\"\n task_logger.info(f\"delete_user_file_impl - Starting id={user_file_id}\")\n\n file_lock: RedisLock | None = None\n if redis_locking:\n redis_client = get_redis_client(tenant_id=tenant_id)\n # Clear the queued guard so the beat can re-enqueue if deletion fails\n # and the file remains in DELETING status.\n redis_client.delete(_user_file_delete_queued_key(user_file_id))\n file_lock = redis_client.lock(\n _user_file_delete_lock_key(user_file_id),\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n if file_lock is not None and not file_lock.acquire(blocking=False):\n task_logger.info(\n f\"delete_user_file_impl - Lock held, skipping user_file_id={user_file_id}\"\n )\n return\n\n try:\n with get_session_with_current_tenant() as db_session:\n user_file = db_session.get(UserFile, _as_uuid(user_file_id))\n if not user_file:\n task_logger.info(\n f\"delete_user_file_impl - User file not found id={user_file_id}\"\n )\n return\n\n if not DISABLE_VECTOR_DB:\n if MANAGED_VESPA:\n httpx_init_vespa_pool(\n 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH\n )\n else:\n httpx_init_vespa_pool(20)\n\n active_search_settings = get_active_search_settings(db_session)\n document_indices = get_all_document_indices(\n search_settings=active_search_settings.primary,\n secondary_search_settings=active_search_settings.secondary,\n httpx_client=HttpxPool.get(\"vespa\"),\n )\n retry_document_indices: list[RetryDocumentIndex] = [\n RetryDocumentIndex(document_index)\n for document_index in document_indices\n ]\n index_name = active_search_settings.primary.index_name\n selection = f\"{index_name}.document_id=='{user_file_id}'\"\n\n chunk_count = 0\n if user_file.chunk_count is None or user_file.chunk_count == 0:\n chunk_count = _get_document_chunk_count(\n index_name=index_name,\n selection=selection,\n )\n else:\n chunk_count = user_file.chunk_count\n\n for retry_document_index in retry_document_indices:\n retry_document_index.delete_single(\n doc_id=user_file_id,\n tenant_id=tenant_id,\n chunk_count=chunk_count,\n )\n\n file_store = get_default_file_store()\n try:\n file_store.delete_file(user_file.file_id)\n file_store.delete_file(\n user_file_id_to_plaintext_file_name(user_file.id)\n )\n except Exception as e:\n task_logger.exception(\n f\"delete_user_file_impl - Error deleting file id={user_file.id} - {e.__class__.__name__}\"\n )\n\n db_session.delete(user_file)\n db_session.commit()\n task_logger.info(f\"delete_user_file_impl - Completed id={user_file_id}\")\n except Exception as e:\n task_logger.exception(\n f\"delete_user_file_impl - Error processing file id={user_file_id} - {e.__class__.__name__}\"\n )\n raise\n finally:\n if file_lock is not None and file_lock.owned():\n file_lock.release()\n\n\n@shared_task(\n name=OnyxCeleryTask.DELETE_SINGLE_USER_FILE,\n bind=True,\n ignore_result=True,\n)\ndef process_single_user_file_delete(\n self: Task, # noqa: ARG001\n *,\n user_file_id: str,\n tenant_id: str,\n) -> None:\n delete_user_file_impl(\n user_file_id=user_file_id, tenant_id=tenant_id, redis_locking=True\n )\n\n\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_USER_FILE_PROJECT_SYNC,\n soft_time_limit=300,\n bind=True,\n ignore_result=True,\n)\ndef check_for_user_file_project_sync(self: Task, *, tenant_id: str) -> None:\n \"\"\"Scan for user files needing project sync and enqueue per-file tasks.\"\"\"\n task_logger.info(\"Starting\")\n\n redis_client = get_redis_client(tenant_id=tenant_id)\n lock: RedisLock = redis_client.lock(\n OnyxRedisLocks.USER_FILE_PROJECT_SYNC_BEAT_LOCK,\n timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,\n )\n\n if not lock.acquire(blocking=False):\n return None\n\n enqueued = 0\n skipped_guard = 0\n try:\n queue_depth = get_user_file_project_sync_queue_depth(self.app)\n if queue_depth > USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH:\n task_logger.warning(\n f\"Queue depth {queue_depth} exceeds \"\n f\"{USER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH}, skipping enqueue for tenant={tenant_id}\"\n )\n return None\n\n with get_session_with_current_tenant() as db_session:\n user_file_ids = (\n db_session.execute(\n select(UserFile.id).where(\n sa.and_(\n sa.or_(\n UserFile.needs_project_sync.is_(True),\n UserFile.needs_persona_sync.is_(True),\n ),\n UserFile.status == UserFileStatus.COMPLETED,\n )\n )\n )\n .scalars()\n .all()\n )\n\n for user_file_id in user_file_ids:\n if not enqueue_user_file_project_sync_task(\n celery_app=self.app,\n redis_client=redis_client,\n user_file_id=user_file_id,\n tenant_id=tenant_id,\n priority=OnyxCeleryPriority.HIGH,\n ):\n skipped_guard += 1\n continue\n enqueued += 1\n finally:\n if lock.owned():\n lock.release()\n\n task_logger.info(\n f\"Enqueued {enqueued} Skipped guard {skipped_guard} tasks for tenant={tenant_id}\"\n )\n return None\n\n\ndef project_sync_user_file_impl(\n *, user_file_id: str, tenant_id: str, redis_locking: bool\n) -> None:\n \"\"\"Core implementation for syncing a user file's project/persona metadata.\n\n When redis_locking=True, acquires a per-file Redis lock and clears the\n queued-key guard (Celery path). When redis_locking=False, skips Redis\n operations (BackgroundTask path).\n \"\"\"\n task_logger.info(f\"project_sync_user_file_impl - Starting id={user_file_id}\")\n\n file_lock: RedisLock | None = None\n if redis_locking:\n redis_client = get_redis_client(tenant_id=tenant_id)\n redis_client.delete(_user_file_project_sync_queued_key(user_file_id))\n file_lock = redis_client.lock(\n user_file_project_sync_lock_key(user_file_id),\n timeout=CELERY_USER_FILE_PROJECT_SYNC_LOCK_TIMEOUT,\n )\n if file_lock is not None and not file_lock.acquire(blocking=False):\n task_logger.info(\n f\"project_sync_user_file_impl - Lock held, skipping user_file_id={user_file_id}\"\n )\n return\n\n try:\n with get_session_with_current_tenant() as db_session:\n user_files = fetch_user_files_with_access_relationships(\n [user_file_id],\n db_session,\n eager_load_groups=global_version.is_ee_version(),\n )\n user_file = user_files[0] if user_files else None\n if not user_file:\n task_logger.info(\n f\"project_sync_user_file_impl - User file not found id={user_file_id}\"\n )\n return\n\n if not DISABLE_VECTOR_DB:\n if MANAGED_VESPA:\n httpx_init_vespa_pool(\n 20, ssl_cert=VESPA_CLOUD_CERT_PATH, ssl_key=VESPA_CLOUD_KEY_PATH\n )\n else:\n httpx_init_vespa_pool(20)\n\n active_search_settings = get_active_search_settings(db_session)\n document_indices = get_all_document_indices(\n search_settings=active_search_settings.primary,\n secondary_search_settings=active_search_settings.secondary,\n httpx_client=HttpxPool.get(\"vespa\"),\n )\n retry_document_indices: list[RetryDocumentIndex] = [\n RetryDocumentIndex(document_index)\n for document_index in document_indices\n ]\n\n project_ids = [project.id for project in user_file.projects]\n persona_ids = [p.id for p in user_file.assistants if not p.deleted]\n\n file_id_str = str(user_file.id)\n access_map = build_access_for_user_files([user_file])\n access = access_map.get(file_id_str)\n\n for retry_document_index in retry_document_indices:\n retry_document_index.update_single(\n doc_id=file_id_str,\n tenant_id=tenant_id,\n chunk_count=user_file.chunk_count,\n fields=(\n VespaDocumentFields(access=access)\n if access is not None\n else None\n ),\n user_fields=VespaDocumentUserFields(\n user_projects=project_ids,\n personas=persona_ids,\n ),\n )\n\n task_logger.info(\n f\"project_sync_user_file_impl - User file id={user_file_id}\"\n )\n\n user_file.needs_project_sync = False\n user_file.needs_persona_sync = False\n user_file.last_project_sync_at = datetime.datetime.now(\n datetime.timezone.utc\n )\n db_session.add(user_file)\n db_session.commit()\n\n except Exception as e:\n task_logger.exception(\n f\"project_sync_user_file_impl - Error syncing project for file id={user_file_id} - {e.__class__.__name__}\"\n )\n raise\n finally:\n if file_lock is not None and file_lock.owned():\n file_lock.release()\n\n\n@shared_task(\n name=OnyxCeleryTask.PROCESS_SINGLE_USER_FILE_PROJECT_SYNC,\n bind=True,\n ignore_result=True,\n)\ndef process_single_user_file_project_sync(\n self: Task, # noqa: ARG001\n *,\n user_file_id: str,\n tenant_id: str,\n) -> None:\n project_sync_user_file_impl(\n user_file_id=user_file_id, tenant_id=tenant_id, redis_locking=True\n )\n" }, { "path": "backend/onyx/background/celery/tasks/vespa/__init__.py", "content": "" }, { "path": "backend/onyx/background/celery/tasks/vespa/document_sync.py", "content": "import time\nfrom typing import cast\nfrom uuid import uuid4\n\nfrom celery import Celery\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DB_YIELD_PER_DEFAULT\nfrom onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.db.document import construct_document_id_select_by_needs_sync\nfrom onyx.db.document import count_documents_by_needs_sync\nfrom onyx.utils.logger import setup_logger\n\n# Redis keys for document sync tracking\nDOCUMENT_SYNC_PREFIX = \"documentsync\"\nDOCUMENT_SYNC_FENCE_KEY = f\"{DOCUMENT_SYNC_PREFIX}_fence\"\nDOCUMENT_SYNC_TASKSET_KEY = f\"{DOCUMENT_SYNC_PREFIX}_taskset\"\nFENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\nTASKSET_TTL = FENCE_TTL\n\nlogger = setup_logger()\n\n\ndef is_document_sync_fenced(r: Redis) -> bool:\n \"\"\"Check if document sync tasks are currently in progress.\"\"\"\n return bool(r.exists(DOCUMENT_SYNC_FENCE_KEY))\n\n\ndef get_document_sync_payload(r: Redis) -> int | None:\n \"\"\"Get the initial number of tasks that were created.\"\"\"\n bytes_result = r.get(DOCUMENT_SYNC_FENCE_KEY)\n if bytes_result is None:\n return None\n return int(cast(int, bytes_result))\n\n\ndef get_document_sync_remaining(r: Redis) -> int:\n \"\"\"Get the number of tasks still pending completion.\"\"\"\n return cast(int, r.scard(DOCUMENT_SYNC_TASKSET_KEY))\n\n\ndef set_document_sync_fence(r: Redis, payload: int | None) -> None:\n \"\"\"Set up the fence and register with active fences.\"\"\"\n if payload is None:\n r.srem(OnyxRedisConstants.ACTIVE_FENCES, DOCUMENT_SYNC_FENCE_KEY)\n r.delete(DOCUMENT_SYNC_FENCE_KEY)\n return\n\n r.set(DOCUMENT_SYNC_FENCE_KEY, payload, ex=FENCE_TTL)\n r.sadd(OnyxRedisConstants.ACTIVE_FENCES, DOCUMENT_SYNC_FENCE_KEY)\n\n\ndef delete_document_sync_taskset(r: Redis) -> None:\n \"\"\"Clear the document sync taskset.\"\"\"\n r.delete(DOCUMENT_SYNC_TASKSET_KEY)\n\n\ndef reset_document_sync(r: Redis) -> None:\n \"\"\"Reset all document sync tracking data.\"\"\"\n r.srem(OnyxRedisConstants.ACTIVE_FENCES, DOCUMENT_SYNC_FENCE_KEY)\n r.delete(DOCUMENT_SYNC_TASKSET_KEY)\n r.delete(DOCUMENT_SYNC_FENCE_KEY)\n\n\ndef generate_document_sync_tasks(\n r: Redis,\n max_tasks: int,\n celery_app: Celery,\n db_session: Session,\n lock: RedisLock,\n tenant_id: str,\n) -> tuple[int, int]:\n \"\"\"Generate sync tasks for all documents that need syncing.\n\n Args:\n r: Redis client\n max_tasks: Maximum number of tasks to generate\n celery_app: Celery application instance\n db_session: Database session\n lock: Redis lock for coordination\n tenant_id: Tenant identifier\n\n Returns:\n tuple[int, int]: (tasks_generated, total_docs_found)\n \"\"\"\n last_lock_time = time.monotonic()\n num_tasks_sent = 0\n num_docs = 0\n\n # Get all documents that need syncing\n stmt = construct_document_id_select_by_needs_sync()\n\n for doc_id in db_session.scalars(stmt).yield_per(DB_YIELD_PER_DEFAULT):\n doc_id = cast(str, doc_id)\n current_time = time.monotonic()\n\n # Reacquire lock periodically to prevent timeout\n if current_time - last_lock_time >= (CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4):\n lock.reacquire()\n last_lock_time = current_time\n\n num_docs += 1\n\n # Create a unique task ID\n custom_task_id = f\"{DOCUMENT_SYNC_PREFIX}_{uuid4()}\"\n\n # Add to the tracking taskset in Redis BEFORE creating the celery task\n r.sadd(DOCUMENT_SYNC_TASKSET_KEY, custom_task_id)\n r.expire(DOCUMENT_SYNC_TASKSET_KEY, TASKSET_TTL)\n\n # Create the Celery task\n celery_app.send_task(\n OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,\n kwargs=dict(document_id=doc_id, tenant_id=tenant_id),\n queue=OnyxCeleryQueues.VESPA_METADATA_SYNC,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n ignore_result=True,\n )\n\n num_tasks_sent += 1\n\n if num_tasks_sent >= max_tasks:\n break\n\n return num_tasks_sent, num_docs\n\n\ndef try_generate_stale_document_sync_tasks(\n celery_app: Celery,\n max_tasks: int,\n db_session: Session,\n r: Redis,\n lock_beat: RedisLock,\n tenant_id: str,\n) -> int | None:\n # the fence is up, do nothing\n if is_document_sync_fenced(r):\n return None\n\n # add tasks to celery and build up the task set to monitor in redis\n stale_doc_count = count_documents_by_needs_sync(db_session)\n if stale_doc_count == 0:\n logger.info(\"No stale documents found. Skipping sync tasks generation.\")\n return None\n\n logger.info(\n f\"Stale documents found (at least {stale_doc_count}). Generating sync tasks in one batch.\"\n )\n\n logger.info(\"generate_document_sync_tasks starting for all documents.\")\n\n # Generate all tasks in one pass\n result = generate_document_sync_tasks(\n r, max_tasks, celery_app, db_session, lock_beat, tenant_id\n )\n\n if result is None:\n return None\n\n tasks_generated, total_docs = result\n\n if tasks_generated >= max_tasks:\n logger.info(\n f\"generate_document_sync_tasks reached the task generation limit: \"\n f\"tasks_generated={tasks_generated} max_tasks={max_tasks}\"\n )\n else:\n logger.info(\n f\"generate_document_sync_tasks finished for all documents. \"\n f\"tasks_generated={tasks_generated} total_docs_found={total_docs}\"\n )\n\n set_document_sync_fence(r, tasks_generated)\n return tasks_generated\n" }, { "path": "backend/onyx/background/celery/tasks/vespa/tasks.py", "content": "import time\nfrom collections.abc import Callable\nfrom http import HTTPStatus\nfrom typing import Any\nfrom typing import cast\n\nimport httpx\nfrom celery import Celery\nfrom celery import shared_task\nfrom celery import Task\nfrom celery.exceptions import SoftTimeLimitExceeded\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\nfrom tenacity import RetryError\n\nfrom onyx.access.access import get_access_for_document\nfrom onyx.background.celery.apps.app_base import task_logger\nfrom onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex\nfrom onyx.background.celery.tasks.shared.tasks import LIGHT_SOFT_TIME_LIMIT\nfrom onyx.background.celery.tasks.shared.tasks import LIGHT_TIME_LIMIT\nfrom onyx.background.celery.tasks.shared.tasks import OnyxCeleryTaskCompletionStatus\nfrom onyx.background.celery.tasks.vespa.document_sync import DOCUMENT_SYNC_FENCE_KEY\nfrom onyx.background.celery.tasks.vespa.document_sync import get_document_sync_payload\nfrom onyx.background.celery.tasks.vespa.document_sync import get_document_sync_remaining\nfrom onyx.background.celery.tasks.vespa.document_sync import reset_document_sync\nfrom onyx.background.celery.tasks.vespa.document_sync import (\n try_generate_stale_document_sync_tasks,\n)\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.configs.app_configs import VESPA_SYNC_MAX_TASKS\nfrom onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.db.document import get_document\nfrom onyx.db.document import mark_document_as_synced\nfrom onyx.db.document_set import delete_document_set\nfrom onyx.db.document_set import fetch_document_sets\nfrom onyx.db.document_set import fetch_document_sets_for_document\nfrom onyx.db.document_set import get_document_set_by_id\nfrom onyx.db.document_set import mark_document_set_as_synced\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import UserGroup\nfrom onyx.db.search_settings import get_active_search_settings\nfrom onyx.db.sync_record import cleanup_sync_records\nfrom onyx.db.sync_record import insert_sync_record\nfrom onyx.db.sync_record import update_sync_record_status\nfrom onyx.document_index.factory import get_all_document_indices\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.httpx.httpx_pool import HttpxPool\nfrom onyx.redis.redis_document_set import RedisDocumentSet\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.redis.redis_pool import get_redis_replica_client\nfrom onyx.redis.redis_pool import redis_lock_dump\nfrom onyx.redis.redis_usergroup import RedisUserGroup\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\nfrom onyx.utils.variable_functionality import global_version\nfrom onyx.utils.variable_functionality import noop_fallback\n\nlogger = setup_logger()\n\n\n# celery auto associates tasks created inside another task,\n# which bloats the result metadata considerably. trail=False prevents this.\n# TODO(andrei): Rename all these kinds of functions from *vespa* to a more\n# generic *document_index*.\n@shared_task(\n name=OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,\n ignore_result=True,\n soft_time_limit=JOB_TIMEOUT,\n trail=False,\n bind=True,\n)\ndef check_for_vespa_sync_task(self: Task, *, tenant_id: str) -> bool | None:\n \"\"\"Runs periodically to check if any document needs syncing.\n Generates sets of tasks for Celery if syncing is needed.\"\"\"\n\n # Useful for debugging timing issues with reacquisitions.\n # TODO: remove once more generalized logging is in place\n task_logger.info(\"check_for_vespa_sync_task started\")\n\n time_start = time.monotonic()\n\n r = get_redis_client()\n r_replica = get_redis_replica_client()\n\n lock_beat: RedisLock = r.lock(\n OnyxRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK,\n timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,\n )\n\n # these tasks should never overlap\n if not lock_beat.acquire(blocking=False):\n return None\n\n try:\n # 1/3: KICKOFF\n with get_session_with_current_tenant() as db_session:\n try_generate_stale_document_sync_tasks(\n self.app, VESPA_SYNC_MAX_TASKS, db_session, r, lock_beat, tenant_id\n )\n\n # region document set scan\n lock_beat.reacquire()\n document_set_ids: list[int] = []\n with get_session_with_current_tenant() as db_session:\n # check if any document sets are not synced\n document_set_info = fetch_document_sets(\n user_id=None, db_session=db_session, include_outdated=True\n )\n\n for document_set, _ in document_set_info:\n document_set_ids.append(document_set.id)\n\n for document_set_id in document_set_ids:\n lock_beat.reacquire()\n with get_session_with_current_tenant() as db_session:\n try_generate_document_set_sync_tasks(\n self.app, document_set_id, db_session, r, lock_beat, tenant_id\n )\n # endregion\n\n # check if any user groups are not synced\n lock_beat.reacquire()\n if global_version.is_ee_version():\n try:\n fetch_user_groups = fetch_versioned_implementation(\n \"onyx.db.user_group\", \"fetch_user_groups\"\n )\n except ModuleNotFoundError:\n # Always exceptions on the MIT version, which is expected\n # We shouldn't actually get here if the ee version check works\n pass\n else:\n usergroup_ids: list[int] = []\n with get_session_with_current_tenant() as db_session:\n user_groups = fetch_user_groups(\n db_session=db_session, only_up_to_date=False\n )\n\n for usergroup in user_groups:\n usergroup_ids.append(usergroup.id)\n\n for usergroup_id in usergroup_ids:\n lock_beat.reacquire()\n with get_session_with_current_tenant() as db_session:\n try_generate_user_group_sync_tasks(\n self.app, usergroup_id, db_session, r, lock_beat, tenant_id\n )\n\n # 2/3: VALIDATE: TODO\n\n # 3/3: FINALIZE\n lock_beat.reacquire()\n keys = cast(set[Any], r_replica.smembers(OnyxRedisConstants.ACTIVE_FENCES))\n for key in keys:\n key_bytes = cast(bytes, key)\n\n if not r.exists(key_bytes):\n r.srem(OnyxRedisConstants.ACTIVE_FENCES, key_bytes)\n continue\n\n key_str = key_bytes.decode(\"utf-8\")\n # NOTE: removing the \"Redis*\" classes, prefer to just have functions to\n # do these things going forward. In short, things should generally be like the doc\n # sync task rather than the others\n if key_str == DOCUMENT_SYNC_FENCE_KEY:\n monitor_document_sync_taskset(r)\n elif key_str.startswith(RedisDocumentSet.FENCE_PREFIX):\n with get_session_with_current_tenant() as db_session:\n monitor_document_set_taskset(tenant_id, key_bytes, r, db_session)\n elif key_str.startswith(RedisUserGroup.FENCE_PREFIX):\n monitor_usergroup_taskset = (\n fetch_versioned_implementation_with_fallback(\n \"onyx.background.celery.tasks.vespa.tasks\",\n \"monitor_usergroup_taskset\",\n noop_fallback,\n )\n )\n with get_session_with_current_tenant() as db_session:\n monitor_usergroup_taskset(tenant_id, key_bytes, r, db_session)\n\n except SoftTimeLimitExceeded:\n task_logger.info(\n \"Soft time limit exceeded, task is being terminated gracefully.\"\n )\n except Exception:\n task_logger.exception(\"Unexpected exception during vespa metadata sync\")\n finally:\n if lock_beat.owned():\n lock_beat.release()\n else:\n task_logger.error(\n f\"check_for_vespa_sync_task - Lock not owned on completion: tenant={tenant_id}\"\n )\n redis_lock_dump(lock_beat, r)\n\n time_elapsed = time.monotonic() - time_start\n task_logger.debug(f\"check_for_vespa_sync_task finished: elapsed={time_elapsed:.2f}\")\n return True\n\n\ndef try_generate_document_set_sync_tasks(\n celery_app: Celery,\n document_set_id: int,\n db_session: Session,\n r: Redis,\n lock_beat: RedisLock,\n tenant_id: str,\n) -> int | None:\n lock_beat.reacquire()\n\n rds = RedisDocumentSet(tenant_id, document_set_id)\n\n # don't generate document set sync tasks if tasks are still pending\n if rds.fenced:\n return None\n\n # don't generate sync tasks if we're up to date\n # race condition with the monitor/cleanup function if we use a cached result!\n document_set = get_document_set_by_id(\n db_session=db_session,\n document_set_id=document_set_id,\n )\n if not document_set:\n return None\n\n if document_set.is_up_to_date:\n # there should be no in-progress sync records if this is up to date\n # clean it up just in case things got into a bad state\n cleanup_sync_records(\n db_session=db_session,\n entity_id=document_set_id,\n sync_type=SyncType.DOCUMENT_SET,\n )\n return None\n\n # add tasks to celery and build up the task set to monitor in redis\n r.delete(rds.taskset_key)\n\n task_logger.info(\n f\"RedisDocumentSet.generate_tasks starting. document_set_id={document_set.id}\"\n )\n\n # Add all documents that need to be updated into the queue\n result = rds.generate_tasks(\n VESPA_SYNC_MAX_TASKS, celery_app, db_session, r, lock_beat, tenant_id\n )\n if result is None:\n return None\n\n tasks_generated = result[0]\n # Currently we are allowing the sync to proceed with 0 tasks.\n # It's possible for sets/groups to be generated initially with no entries\n # and they still need to be marked as up to date.\n # if tasks_generated == 0:\n # return 0\n\n task_logger.info(\n f\"RedisDocumentSet.generate_tasks finished. document_set={document_set.id} tasks_generated={tasks_generated}\"\n )\n\n # create before setting fence to avoid race condition where the monitoring\n # task updates the sync record before it is created\n try:\n insert_sync_record(\n db_session=db_session,\n entity_id=document_set_id,\n sync_type=SyncType.DOCUMENT_SET,\n )\n except Exception:\n task_logger.exception(\"insert_sync_record exceptioned.\")\n\n # set this only after all tasks have been added\n rds.set_fence(tasks_generated)\n return tasks_generated\n\n\ndef try_generate_user_group_sync_tasks(\n celery_app: Celery,\n usergroup_id: int,\n db_session: Session,\n r: Redis,\n lock_beat: RedisLock,\n tenant_id: str,\n) -> int | None:\n lock_beat.reacquire()\n\n rug = RedisUserGroup(tenant_id, usergroup_id)\n if rug.fenced:\n # don't generate sync tasks if tasks are still pending\n return None\n\n # race condition with the monitor/cleanup function if we use a cached result!\n fetch_user_group = cast(\n Callable[[Session, int], UserGroup | None],\n fetch_versioned_implementation(\"onyx.db.user_group\", \"fetch_user_group\"),\n )\n\n usergroup = fetch_user_group(db_session, usergroup_id)\n if not usergroup:\n return None\n\n if usergroup.is_up_to_date:\n # there should be no in-progress sync records if this is up to date\n # clean it up just in case things got into a bad state\n cleanup_sync_records(\n db_session=db_session,\n entity_id=usergroup_id,\n sync_type=SyncType.USER_GROUP,\n )\n return None\n\n # add tasks to celery and build up the task set to monitor in redis\n r.delete(rug.taskset_key)\n\n # Add all documents that need to be updated into the queue\n task_logger.info(\n f\"RedisUserGroup.generate_tasks starting. usergroup_id={usergroup.id}\"\n )\n result = rug.generate_tasks(\n VESPA_SYNC_MAX_TASKS, celery_app, db_session, r, lock_beat, tenant_id\n )\n if result is None:\n return None\n\n tasks_generated = result[0]\n # Currently we are allowing the sync to proceed with 0 tasks.\n # It's possible for sets/groups to be generated initially with no entries\n # and they still need to be marked as up to date.\n # if tasks_generated == 0:\n # return 0\n\n task_logger.info(\n f\"RedisUserGroup.generate_tasks finished. usergroup={usergroup.id} tasks_generated={tasks_generated}\"\n )\n\n # create before setting fence to avoid race condition where the monitoring\n # task updates the sync record before it is created\n try:\n insert_sync_record(\n db_session=db_session,\n entity_id=usergroup_id,\n sync_type=SyncType.USER_GROUP,\n )\n except Exception:\n task_logger.exception(\"insert_sync_record exceptioned.\")\n\n # set this only after all tasks have been added\n rug.set_fence(tasks_generated)\n\n return tasks_generated\n\n\ndef monitor_document_sync_taskset(r: Redis) -> None:\n initial_count = get_document_sync_payload(r)\n if initial_count is None:\n return\n\n remaining = get_document_sync_remaining(r)\n task_logger.info(\n f\"Document sync progress: remaining={remaining} initial={initial_count}\"\n )\n if remaining == 0:\n reset_document_sync(r)\n task_logger.info(f\"Successfully synced all documents. count={initial_count}\")\n\n\ndef monitor_document_set_taskset(\n tenant_id: str, key_bytes: bytes, r: Redis, db_session: Session\n) -> None:\n fence_key = key_bytes.decode(\"utf-8\")\n document_set_id_str = RedisDocumentSet.get_id_from_fence_key(fence_key)\n if document_set_id_str is None:\n task_logger.warning(f\"could not parse document set id from {fence_key}\")\n return\n\n document_set_id = int(document_set_id_str)\n\n rds = RedisDocumentSet(tenant_id, document_set_id)\n if not rds.fenced:\n return\n\n initial_count = rds.payload\n if initial_count is None:\n return\n\n count = cast(int, r.scard(rds.taskset_key))\n task_logger.info(\n f\"Document set sync progress: document_set={document_set_id} remaining={count} initial={initial_count}\"\n )\n if count > 0:\n update_sync_record_status(\n db_session=db_session,\n entity_id=document_set_id,\n sync_type=SyncType.DOCUMENT_SET,\n sync_status=SyncStatus.IN_PROGRESS,\n num_docs_synced=count,\n )\n return\n\n document_set = cast(\n DocumentSet,\n get_document_set_by_id(db_session=db_session, document_set_id=document_set_id),\n ) # casting since we \"know\" a document set with this ID exists\n if document_set:\n has_connector_pairs = bool(document_set.connector_credential_pairs)\n # Federated connectors should keep a document set alive even without cc pairs.\n has_federated_connectors = bool(\n getattr(document_set, \"federated_connectors\", [])\n )\n\n if not has_connector_pairs and not has_federated_connectors:\n # If there are no connectors of any kind, delete the document set.\n delete_document_set(document_set_row=document_set, db_session=db_session)\n task_logger.info(\n f\"Successfully deleted document set: document_set={document_set_id}\"\n )\n else:\n mark_document_set_as_synced(document_set_id, db_session)\n task_logger.info(\n f\"Successfully synced document set: document_set={document_set_id}\"\n )\n\n try:\n update_sync_record_status(\n db_session=db_session,\n entity_id=document_set_id,\n sync_type=SyncType.DOCUMENT_SET,\n sync_status=SyncStatus.SUCCESS,\n num_docs_synced=initial_count,\n )\n except Exception:\n task_logger.exception(\n f\"update_sync_record_status exceptioned. document_set_id={document_set_id} Resetting document set regardless.\"\n )\n\n rds.reset()\n\n\n@shared_task(\n name=OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,\n bind=True,\n soft_time_limit=LIGHT_SOFT_TIME_LIMIT,\n time_limit=LIGHT_TIME_LIMIT,\n max_retries=3,\n)\ndef vespa_metadata_sync_task(self: Task, document_id: str, *, tenant_id: str) -> bool:\n start = time.monotonic()\n\n completion_status = OnyxCeleryTaskCompletionStatus.UNDEFINED\n\n try:\n with get_session_with_current_tenant() as db_session:\n active_search_settings = get_active_search_settings(db_session)\n # This flow is for updates so we get all indices.\n document_indices = get_all_document_indices(\n search_settings=active_search_settings.primary,\n secondary_search_settings=active_search_settings.secondary,\n httpx_client=HttpxPool.get(\"vespa\"),\n )\n\n retry_document_indices: list[RetryDocumentIndex] = [\n RetryDocumentIndex(document_index)\n for document_index in document_indices\n ]\n\n doc = get_document(document_id, db_session)\n if not doc:\n elapsed = time.monotonic() - start\n task_logger.info(\n f\"doc={document_id} action=no_operation elapsed={elapsed:.2f}\"\n )\n completion_status = OnyxCeleryTaskCompletionStatus.SKIPPED\n else:\n # document set sync\n doc_sets = fetch_document_sets_for_document(document_id, db_session)\n update_doc_sets: set[str] = set(doc_sets)\n\n # User group sync\n doc_access = get_access_for_document(\n document_id=document_id, db_session=db_session\n )\n\n fields = VespaDocumentFields(\n document_sets=update_doc_sets,\n access=doc_access,\n boost=doc.boost,\n hidden=doc.hidden,\n # aggregated_boost_factor=doc.aggregated_boost_factor,\n )\n\n for retry_document_index in retry_document_indices:\n # TODO(andrei): Previously there was a comment here saying\n # it was ok if a doc did not exist in the document index. I\n # don't agree with that claim, so keep an eye on this task\n # to see if this raises.\n retry_document_index.update_single(\n document_id,\n tenant_id=tenant_id,\n chunk_count=doc.chunk_count,\n fields=fields,\n user_fields=None,\n )\n\n # update db last. Worst case = we crash right before this and\n # the sync might repeat again later\n mark_document_as_synced(document_id, db_session)\n\n elapsed = time.monotonic() - start\n task_logger.info(f\"doc={document_id} action=sync elapsed={elapsed:.2f}\")\n completion_status = OnyxCeleryTaskCompletionStatus.SUCCEEDED\n except SoftTimeLimitExceeded:\n task_logger.info(f\"SoftTimeLimitExceeded exception. doc={document_id}\")\n completion_status = OnyxCeleryTaskCompletionStatus.SOFT_TIME_LIMIT\n except Exception as ex:\n e: Exception | None = None\n while True:\n if isinstance(ex, RetryError):\n task_logger.warning(\n f\"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}\"\n )\n\n # only set the inner exception if it is of type Exception\n e_temp = ex.last_attempt.exception()\n if isinstance(e_temp, Exception):\n e = e_temp\n else:\n e = ex\n\n if isinstance(e, httpx.HTTPStatusError):\n if e.response.status_code == HTTPStatus.BAD_REQUEST:\n task_logger.exception(\n f\"Non-retryable HTTPStatusError: doc={document_id} status={e.response.status_code}\"\n )\n completion_status = (\n OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION\n )\n break\n\n task_logger.exception(\n f\"vespa_metadata_sync_task exceptioned: doc={document_id}\"\n )\n\n completion_status = OnyxCeleryTaskCompletionStatus.RETRYABLE_EXCEPTION\n if (\n self.max_retries is not None\n and self.request.retries >= self.max_retries\n ):\n completion_status = (\n OnyxCeleryTaskCompletionStatus.NON_RETRYABLE_EXCEPTION\n )\n\n # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64\n countdown = 2 ** (self.request.retries + 4)\n self.retry(exc=e, countdown=countdown) # this will raise a celery exception\n break # we won't hit this, but it looks weird not to have it\n finally:\n task_logger.info(\n f\"vespa_metadata_sync_task completed: status={completion_status.value} doc={document_id}\"\n )\n\n return completion_status == OnyxCeleryTaskCompletionStatus.SUCCEEDED\n" }, { "path": "backend/onyx/background/celery/versioned_apps/beat.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.background.celery.apps.beat import celery_app\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\napp: Celery = celery_app\n" }, { "path": "backend/onyx/background/celery/versioned_apps/client.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\nThis code is different from the primary/beat stubs because there is no EE version to\nfetch. Port over the code in those files if we add an EE version of this worker.\n\nThis is an app stub purely for sending tasks as a client.\n\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\n\n\ndef get_app() -> Celery:\n from onyx.background.celery.apps.client import celery_app\n\n return celery_app\n\n\napp = get_app()\n" }, { "path": "backend/onyx/background/celery/versioned_apps/docfetching.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\nThis code is different from the primary/beat stubs because there is no EE version to\nfetch. Port over the code in those files if we add an EE version of this worker.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\n\n\ndef get_app() -> Celery:\n from onyx.background.celery.apps.docfetching import celery_app\n\n return celery_app\n\n\napp = get_app()\n" }, { "path": "backend/onyx/background/celery/versioned_apps/docprocessing.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\nThis code is different from the primary/beat stubs because there is no EE version to\nfetch. Port over the code in those files if we add an EE version of this worker.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\n\n\ndef get_app() -> Celery:\n from onyx.background.celery.apps.docprocessing import celery_app\n\n return celery_app\n\n\napp = get_app()\n" }, { "path": "backend/onyx/background/celery/versioned_apps/heavy.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\nThis code is different from the primary/beat stubs because there is no EE version to\nfetch. Port over the code in those files if we add an EE version of this worker.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\napp: Celery = fetch_versioned_implementation(\n \"onyx.background.celery.apps.heavy\",\n \"celery_app\",\n)\n" }, { "path": "backend/onyx/background/celery/versioned_apps/light.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\nThis code is different from the primary/beat stubs because there is no EE version to\nfetch. Port over the code in those files if we add an EE version of this worker.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\napp: Celery = fetch_versioned_implementation(\n \"onyx.background.celery.apps.light\",\n \"celery_app\",\n)\n" }, { "path": "backend/onyx/background/celery/versioned_apps/monitoring.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\napp: Celery = fetch_versioned_implementation(\n \"onyx.background.celery.apps.monitoring\",\n \"celery_app\",\n)\n" }, { "path": "backend/onyx/background/celery/versioned_apps/primary.py", "content": "\"\"\"Factory stub for running celery worker / celery beat.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\napp: Celery = fetch_versioned_implementation(\n \"onyx.background.celery.apps.primary\",\n \"celery_app\",\n)\n" }, { "path": "backend/onyx/background/celery/versioned_apps/user_file_processing.py", "content": "\"\"\"Factory stub for running the user file processing Celery worker.\"\"\"\n\nfrom celery import Celery\n\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\nset_is_ee_based_on_env_variable()\n\n\ndef get_app() -> Celery:\n from onyx.background.celery.apps.user_file_processing import celery_app\n\n return celery_app\n\n\napp = get_app()\n" }, { "path": "backend/onyx/background/error_logging.py", "content": "from sqlalchemy.exc import IntegrityError\n\nfrom onyx.db.background_error import create_background_error\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n\ndef emit_background_error(\n message: str,\n cc_pair_id: int | None = None,\n) -> None:\n \"\"\"Currently just saves a row in the background_errors table.\n\n In the future, could create notifications based on the severity.\"\"\"\n error_message = \"\"\n\n # try to write to the db, but handle IntegrityError specifically\n try:\n with get_session_with_current_tenant() as db_session:\n create_background_error(db_session, message, cc_pair_id)\n except IntegrityError as e:\n # Log an error if the cc_pair_id was deleted or any other exception occurs\n error_message = (\n f\"Failed to create background error: {str(e)}. Original message: {message}\"\n )\n except Exception:\n pass\n\n if not error_message:\n return\n\n # if we get here from an IntegrityError, try to write the error message to the db\n # we need a new session because the first session is now invalid\n try:\n with get_session_with_current_tenant() as db_session:\n create_background_error(db_session, error_message, None)\n except Exception:\n pass\n" }, { "path": "backend/onyx/background/indexing/checkpointing_utils.py", "content": "from datetime import datetime\nfrom datetime import timedelta\nfrom io import BytesIO\n\nfrom sqlalchemy import and_\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import NUM_DAYS_TO_KEEP_CHECKPOINTS\nfrom onyx.connectors.interfaces import BaseConnector\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.index_attempt import get_index_attempt\nfrom onyx.db.index_attempt import get_recent_completed_attempts_for_cc_pair\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexingStatus\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.object_size_check import deep_getsizeof\n\nlogger = setup_logger()\n\n_NUM_RECENT_ATTEMPTS_TO_CONSIDER = 50\n\n\ndef _build_checkpoint_pointer(index_attempt_id: int) -> str:\n return f\"checkpoint_{index_attempt_id}.json\"\n\n\ndef save_checkpoint(\n db_session: Session, index_attempt_id: int, checkpoint: ConnectorCheckpoint\n) -> str:\n \"\"\"Save a checkpoint for a given index attempt to the file store\"\"\"\n checkpoint_pointer = _build_checkpoint_pointer(index_attempt_id)\n\n file_store = get_default_file_store()\n file_store.save_file(\n content=BytesIO(checkpoint.model_dump_json().encode()),\n display_name=checkpoint_pointer,\n file_origin=FileOrigin.INDEXING_CHECKPOINT,\n file_type=\"application/json\",\n file_id=checkpoint_pointer,\n )\n\n index_attempt = get_index_attempt(db_session, index_attempt_id)\n if not index_attempt:\n raise RuntimeError(f\"Index attempt {index_attempt_id} not found in DB.\")\n index_attempt.checkpoint_pointer = checkpoint_pointer\n db_session.add(index_attempt)\n db_session.commit()\n return checkpoint_pointer\n\n\ndef load_checkpoint(\n index_attempt_id: int, connector: BaseConnector\n) -> ConnectorCheckpoint:\n \"\"\"Load a checkpoint for a given index attempt from the file store\"\"\"\n checkpoint_pointer = _build_checkpoint_pointer(index_attempt_id)\n file_store = get_default_file_store()\n checkpoint_io = file_store.read_file(checkpoint_pointer, mode=\"rb\")\n checkpoint_data = checkpoint_io.read().decode(\"utf-8\")\n if isinstance(connector, CheckpointedConnector):\n return connector.validate_checkpoint_json(checkpoint_data)\n return ConnectorCheckpoint.model_validate_json(checkpoint_data)\n\n\ndef get_latest_valid_checkpoint(\n db_session: Session,\n cc_pair_id: int,\n search_settings_id: int,\n window_start: datetime,\n window_end: datetime,\n connector: BaseConnector,\n) -> tuple[ConnectorCheckpoint, bool]:\n \"\"\"Get the latest valid checkpoint for a given connector credential pair\"\"\"\n checkpoint_candidates = get_recent_completed_attempts_for_cc_pair(\n cc_pair_id=cc_pair_id,\n search_settings_id=search_settings_id,\n db_session=db_session,\n limit=_NUM_RECENT_ATTEMPTS_TO_CONSIDER,\n )\n\n # don't keep using checkpoints if we've had a bunch of failed attempts in a row\n # where we make no progress. Only do this if we have had at least\n # _NUM_RECENT_ATTEMPTS_TO_CONSIDER completed attempts.\n if len(checkpoint_candidates) >= _NUM_RECENT_ATTEMPTS_TO_CONSIDER:\n had_any_progress = False\n for candidate in checkpoint_candidates:\n if (\n candidate.total_docs_indexed is not None\n and candidate.total_docs_indexed > 0\n ) or candidate.status.is_successful():\n had_any_progress = True\n break\n\n if not had_any_progress:\n logger.warning(\n f\"{_NUM_RECENT_ATTEMPTS_TO_CONSIDER} consecutive failed attempts without progress \"\n f\"found for cc_pair={cc_pair_id}. Ignoring checkpoint to let the run start \"\n \"from scratch.\"\n )\n return connector.build_dummy_checkpoint(), False\n\n # filter out any candidates that don't meet the criteria\n checkpoint_candidates = [\n candidate\n for candidate in checkpoint_candidates\n if (\n candidate.poll_range_start == window_start\n and candidate.poll_range_end == window_end\n and (\n candidate.status == IndexingStatus.FAILED\n # if the background job was killed (and thus the attempt was canceled)\n # we still want to use the checkpoint so that we can pick up where we left off\n or candidate.status == IndexingStatus.CANCELED\n )\n and candidate.checkpoint_pointer is not None\n # NOTE: There are a couple connectors that may make progress but not have\n # any \"total_docs_indexed\". E.g. they are going through\n # Slack channels, and tons of them don't have any updates.\n # Leaving the below in as historical context / in-case we want to use it again.\n # we want to make sure that the checkpoint is actually useful\n # if it's only gone through a few docs, it's probably not worth\n # using. This also avoids weird cases where a connector is basically\n # non-functional but still \"makes progress\" by slowly moving the\n # checkpoint forward run after run\n # and candidate.total_docs_indexed\n # and candidate.total_docs_indexed > 100\n )\n ]\n\n # assumes latest checkpoint is the furthest along. This only isn't true\n # if something else has gone wrong.\n latest_valid_checkpoint_candidate = (\n checkpoint_candidates[0] if checkpoint_candidates else None\n )\n\n checkpoint = connector.build_dummy_checkpoint()\n if latest_valid_checkpoint_candidate is None:\n logger.info(\n f\"No valid checkpoint found for cc_pair={cc_pair_id}. Starting from scratch.\"\n )\n return checkpoint, False\n\n try:\n previous_checkpoint = load_checkpoint(\n index_attempt_id=latest_valid_checkpoint_candidate.id,\n connector=connector,\n )\n except Exception:\n logger.exception(\n f\"Failed to load checkpoint from previous failed attempt with ID \"\n f\"{latest_valid_checkpoint_candidate.id}. Falling back to default checkpoint.\"\n )\n return checkpoint, False\n\n logger.info(\n f\"Using checkpoint from previous failed attempt with ID \"\n f\"{latest_valid_checkpoint_candidate.id}. Previous checkpoint: \"\n f\"{previous_checkpoint}\"\n )\n return previous_checkpoint, True\n\n\ndef get_index_attempts_with_old_checkpoints(\n db_session: Session, days_to_keep: int = NUM_DAYS_TO_KEEP_CHECKPOINTS\n) -> list[IndexAttempt]:\n \"\"\"Get all index attempts with checkpoints older than the specified number of days.\n\n Args:\n db_session: The database session\n days_to_keep: Number of days to keep checkpoints for (default: NUM_DAYS_TO_KEEP_CHECKPOINTS)\n\n Returns:\n List of IndexAttempt objects with old checkpoints\n \"\"\"\n cutoff_date = get_db_current_time(db_session) - timedelta(days=days_to_keep)\n\n # Find all index attempts with checkpoints older than cutoff_date\n old_attempts = (\n db_session.query(IndexAttempt)\n .filter(\n and_(\n IndexAttempt.checkpoint_pointer.isnot(None),\n IndexAttempt.time_created < cutoff_date,\n )\n )\n .all()\n )\n\n return old_attempts\n\n\ndef cleanup_checkpoint(db_session: Session, index_attempt_id: int) -> None:\n \"\"\"Clean up a checkpoint for a given index attempt\"\"\"\n index_attempt = get_index_attempt(db_session, index_attempt_id)\n if not index_attempt:\n raise RuntimeError(f\"Index attempt {index_attempt_id} not found in DB.\")\n\n if not index_attempt.checkpoint_pointer:\n return None\n\n file_store = get_default_file_store()\n file_store.delete_file(index_attempt.checkpoint_pointer)\n\n index_attempt.checkpoint_pointer = None\n db_session.add(index_attempt)\n db_session.commit()\n\n return None\n\n\ndef check_checkpoint_size(checkpoint: ConnectorCheckpoint) -> None:\n \"\"\"Check if the checkpoint content size exceeds the limit (200MB)\"\"\"\n content_size = deep_getsizeof(checkpoint.model_dump())\n if content_size > 200_000_000: # 200MB in bytes\n raise ValueError(\n f\"Checkpoint content size ({content_size} bytes) exceeds 200MB limit\"\n )\n" }, { "path": "backend/onyx/background/indexing/dask_utils.py", "content": "import asyncio\n\nimport psutil\nfrom dask.distributed import WorkerPlugin\nfrom distributed import Worker\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass ResourceLogger(WorkerPlugin):\n def __init__(self, log_interval: int = 60 * 5):\n self.log_interval = log_interval\n\n def setup(self, worker: Worker) -> None:\n \"\"\"This method will be called when the plugin is attached to a worker.\"\"\"\n self.worker = worker\n worker.loop.add_callback(self.log_resources)\n\n async def log_resources(self) -> None:\n \"\"\"Periodically log CPU and memory usage.\n\n NOTE: must be async or else will clog up the worker indefinitely due to the fact that\n Dask uses Tornado under the hood (which is async)\"\"\"\n while True:\n cpu_percent = psutil.cpu_percent(interval=None)\n memory_available_gb = psutil.virtual_memory().available / (1024.0**3)\n # You can now log these values or send them to a monitoring service\n logger.debug(\n f\"Worker {self.worker.address}: CPU usage {cpu_percent}%, Memory available {memory_available_gb}GB\"\n )\n await asyncio.sleep(self.log_interval)\n" }, { "path": "backend/onyx/background/indexing/index_attempt_utils.py", "content": "from datetime import timedelta\n\nfrom sqlalchemy import func\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexAttemptError\n\n\n# Always retain at least this many attempts per connector/search settings pair\nNUM_RECENT_INDEX_ATTEMPTS_TO_KEEP = 10\n\n\ndef get_old_index_attempts(\n db_session: Session, days_to_keep: int = NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS\n) -> list[IndexAttempt]:\n \"\"\"\n Get index attempts older than the specified number of days while retaining\n the latest NUM_RECENT_INDEX_ATTEMPTS_TO_KEEP per connector/search settings pair.\n \"\"\"\n cutoff_date = get_db_current_time(db_session) - timedelta(days=days_to_keep)\n ranked_attempts = (\n db_session.query(\n IndexAttempt.id.label(\"attempt_id\"),\n IndexAttempt.time_created.label(\"time_created\"),\n func.row_number()\n .over(\n partition_by=(\n IndexAttempt.connector_credential_pair_id,\n IndexAttempt.search_settings_id,\n ),\n order_by=IndexAttempt.time_created.desc(),\n )\n .label(\"attempt_rank\"),\n )\n ).subquery()\n\n return (\n db_session.query(IndexAttempt)\n .join(\n ranked_attempts,\n IndexAttempt.id == ranked_attempts.c.attempt_id,\n )\n .filter(\n ranked_attempts.c.time_created < cutoff_date,\n ranked_attempts.c.attempt_rank > NUM_RECENT_INDEX_ATTEMPTS_TO_KEEP,\n )\n .all()\n )\n\n\ndef cleanup_index_attempts(db_session: Session, index_attempt_ids: list[int]) -> None:\n \"\"\"Clean up multiple index attempts\"\"\"\n db_session.query(IndexAttemptError).filter(\n IndexAttemptError.index_attempt_id.in_(index_attempt_ids)\n ).delete(synchronize_session=False)\n\n db_session.query(IndexAttempt).filter(\n IndexAttempt.id.in_(index_attempt_ids)\n ).delete(synchronize_session=False)\n db_session.commit()\n" }, { "path": "backend/onyx/background/indexing/job_client.py", "content": "\"\"\"Custom client that works similarly to Dask, but simpler and more lightweight.\nDask jobs behaved very strangely - they would die all the time, retries would\nnot follow the expected behavior, etc.\n\nNOTE: cannot use Celery directly due to\nhttps://github.com/celery/celery/issues/7007#issuecomment-1740139367\"\"\"\n\nimport multiprocessing as mp\nimport sys\nimport traceback\nfrom collections.abc import Callable\nfrom dataclasses import dataclass\nfrom multiprocessing.context import SpawnProcess\nfrom typing import Any\nfrom typing import Literal\nfrom typing import Optional\n\nfrom onyx.configs.constants import POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import TENANT_ID_PREFIX\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\n\nclass SimpleJobException(Exception):\n \"\"\"lets us raise an exception that will return a specific error code\"\"\"\n\n def __init__(self, *args: Any, **kwargs: Any) -> None:\n code: int | None = kwargs.pop(\"code\", None)\n self.code = code\n super().__init__(*args, **kwargs)\n\n\nJobStatusType = (\n Literal[\"error\"]\n | Literal[\"finished\"]\n | Literal[\"pending\"]\n | Literal[\"running\"]\n | Literal[\"cancelled\"]\n)\n\n\ndef _initializer(\n func: Callable,\n queue: mp.Queue,\n args: list | tuple,\n kwargs: dict[str, Any] | None = None,\n) -> Any:\n \"\"\"Initialize the child process with a fresh SQLAlchemy Engine.\n\n Based on SQLAlchemy's recommendations to handle multiprocessing:\n https://docs.sqlalchemy.org/en/20/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork\n \"\"\"\n if kwargs is None:\n kwargs = {}\n\n logger.info(\"Initializing spawned worker child process.\")\n # 1. Get tenant_id from args or fallback to default\n tenant_id = POSTGRES_DEFAULT_SCHEMA\n for arg in reversed(args):\n if isinstance(arg, str) and arg.startswith(TENANT_ID_PREFIX):\n tenant_id = arg\n break\n\n # 2. Set the tenant context before running anything\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n # Reset the engine in the child process\n SqlEngine.reset_engine()\n\n # Optionally set a custom app name for database logging purposes\n SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME)\n\n # Initialize a new engine with desired parameters\n SqlEngine.init_engine(\n pool_size=4, max_overflow=12, pool_recycle=60, pool_pre_ping=True\n )\n\n # Proceed with executing the target function\n try:\n return func(*args, **kwargs)\n except SimpleJobException as e:\n logger.exception(\"SimpleJob raised a SimpleJobException\")\n error_msg = traceback.format_exc()\n queue.put(error_msg) # Send the exception to the parent process\n\n sys.exit(e.code) # use the given exit code\n except Exception:\n logger.exception(\"SimpleJob raised an exception\")\n error_msg = traceback.format_exc()\n queue.put(error_msg) # Send the exception to the parent process\n\n sys.exit(255) # use 255 to indicate a generic exception\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\ndef _run_in_process(\n func: Callable,\n queue: mp.Queue,\n args: list | tuple,\n kwargs: dict[str, Any] | None = None,\n) -> None:\n _initializer(func, queue, args, kwargs)\n\n\n@dataclass\nclass SimpleJob:\n \"\"\"Drop in replacement for `dask.distributed.Future`\"\"\"\n\n id: int\n process: Optional[\"SpawnProcess\"] = None\n queue: Optional[mp.Queue] = None\n _exception: Optional[str] = None\n\n def cancel(self) -> bool:\n return self.release()\n\n def release(self) -> bool:\n if self.process is not None and self.process.is_alive():\n self.process.terminate()\n return True\n return False\n\n @property\n def status(self) -> JobStatusType:\n if not self.process:\n return \"pending\"\n elif self.process.is_alive():\n return \"running\"\n elif self.process.exitcode is None:\n return \"cancelled\"\n elif self.process.exitcode != 0:\n return \"error\"\n else:\n return \"finished\"\n\n def done(self) -> bool:\n return (\n self.status == \"finished\"\n or self.status == \"cancelled\"\n or self.status == \"error\"\n )\n\n def exception(self) -> str:\n \"\"\"Needed to match the Dask API, but not implemented since we don't currently\n have a way to get back the exception information from the child process.\"\"\"\n\n \"\"\"Retrieve exception from the multiprocessing queue if available.\"\"\"\n if self._exception is None and self.queue and not self.queue.empty():\n self._exception = self.queue.get() # Get exception from queue\n\n return (\n self._exception or f\"Job with ID '{self.id}' did not report an exception.\"\n )\n\n\nclass SimpleJobClient:\n \"\"\"Drop in replacement for `dask.distributed.Client`\"\"\"\n\n def __init__(self, n_workers: int = 1) -> None:\n self.n_workers = n_workers\n self.job_id_counter = 0\n self.jobs: dict[int, SimpleJob] = {}\n\n def _cleanup_completed_jobs(self) -> None:\n current_job_ids = list(self.jobs.keys())\n for job_id in current_job_ids:\n job = self.jobs.get(job_id)\n if job and job.done():\n logger.debug(f\"Cleaning up job with id: '{job.id}'\")\n del self.jobs[job.id]\n\n def submit(\n self,\n func: Callable,\n *args: Any,\n pure: bool = True, # noqa: ARG002\n ) -> SimpleJob | None:\n \"\"\"NOTE: `pure` arg is needed so this can be a drop in replacement for Dask\"\"\"\n self._cleanup_completed_jobs()\n if len(self.jobs) >= self.n_workers:\n logger.debug(\n f\"No available workers to run job. Currently running '{len(self.jobs)}' jobs, with a limit of '{self.n_workers}'.\"\n )\n return None\n\n job_id = self.job_id_counter\n self.job_id_counter += 1\n\n # this approach allows us to always \"spawn\" a new process regardless of\n # get_start_method's current setting\n ctx = mp.get_context(\"spawn\")\n queue = ctx.Queue()\n process = ctx.Process(\n target=_run_in_process, args=(func, queue, args), daemon=True\n )\n job = SimpleJob(id=job_id, process=process, queue=queue)\n process.start()\n\n self.jobs[job_id] = job\n\n return job\n" }, { "path": "backend/onyx/background/indexing/memory_tracer.py", "content": "import tracemalloc\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nDANSWER_TRACEMALLOC_FRAMES = 10\n\n\nclass MemoryTracer:\n def __init__(self, interval: int = 0, num_print_entries: int = 5):\n self.interval = interval\n self.num_print_entries = num_print_entries\n self.snapshot_first: tracemalloc.Snapshot | None = None\n self.snapshot_prev: tracemalloc.Snapshot | None = None\n self.snapshot: tracemalloc.Snapshot | None = None\n self.counter = 0\n\n def start(self) -> None:\n \"\"\"Start the memory tracer if interval is greater than 0.\"\"\"\n if self.interval > 0:\n logger.debug(f\"Memory tracer starting: interval={self.interval}\")\n tracemalloc.start(DANSWER_TRACEMALLOC_FRAMES)\n self._take_snapshot()\n\n def stop(self) -> None:\n \"\"\"Stop the memory tracer if it's running.\"\"\"\n if self.interval > 0:\n self.log_final_diff()\n tracemalloc.stop()\n logger.debug(\"Memory tracer stopped.\")\n\n def _take_snapshot(self) -> None:\n \"\"\"Take a snapshot and update internal snapshot states.\"\"\"\n snapshot = tracemalloc.take_snapshot()\n # Filter out irrelevant frames\n snapshot = snapshot.filter_traces(\n (\n tracemalloc.Filter(False, tracemalloc.__file__),\n tracemalloc.Filter(False, \"\"),\n tracemalloc.Filter(False, \"\"),\n )\n )\n\n if not self.snapshot_first:\n self.snapshot_first = snapshot\n\n if self.snapshot:\n self.snapshot_prev = self.snapshot\n\n self.snapshot = snapshot\n\n def _log_diff(\n self, current: tracemalloc.Snapshot, previous: tracemalloc.Snapshot\n ) -> None:\n \"\"\"Log the memory difference between two snapshots.\"\"\"\n stats = current.compare_to(previous, \"traceback\")\n for s in stats[: self.num_print_entries]:\n logger.debug(f\"Tracer diff: {s}\")\n for line in s.traceback.format():\n logger.debug(f\"* {line}\")\n\n def increment_and_maybe_trace(self) -> None:\n \"\"\"Increment counter and perform trace if interval is hit.\"\"\"\n if self.interval <= 0:\n return\n\n self.counter += 1\n if self.counter % self.interval == 0:\n logger.debug(\n f\"Running trace comparison for batch {self.counter}. interval={self.interval}\"\n )\n self._take_snapshot()\n if self.snapshot and self.snapshot_prev:\n self._log_diff(self.snapshot, self.snapshot_prev)\n\n def log_final_diff(self) -> None:\n \"\"\"Log the final memory diff between start and end of indexing.\"\"\"\n if self.interval <= 0:\n return\n\n logger.debug(\n f\"Running trace comparison between start and end of indexing. {self.counter} batches processed.\"\n )\n self._take_snapshot()\n if self.snapshot and self.snapshot_first:\n self._log_diff(self.snapshot, self.snapshot_first)\n" }, { "path": "backend/onyx/background/indexing/models.py", "content": "from datetime import datetime\n\nfrom pydantic import BaseModel\n\nfrom onyx.db.models import IndexAttemptError\n\n\nclass IndexAttemptErrorPydantic(BaseModel):\n id: int\n connector_credential_pair_id: int\n\n document_id: str | None\n document_link: str | None\n\n entity_id: str | None\n failed_time_range_start: datetime | None\n failed_time_range_end: datetime | None\n\n failure_message: str\n is_resolved: bool = False\n\n time_created: datetime\n\n index_attempt_id: int\n\n @classmethod\n def from_model(cls, model: IndexAttemptError) -> \"IndexAttemptErrorPydantic\":\n return cls(\n id=model.id,\n connector_credential_pair_id=model.connector_credential_pair_id,\n document_id=model.document_id,\n document_link=model.document_link,\n entity_id=model.entity_id,\n failed_time_range_start=model.failed_time_range_start,\n failed_time_range_end=model.failed_time_range_end,\n failure_message=model.failure_message,\n is_resolved=model.is_resolved,\n time_created=model.time_created,\n index_attempt_id=model.index_attempt_id,\n )\n" }, { "path": "backend/onyx/background/indexing/run_docfetching.py", "content": "import sys\nimport time\nimport traceback\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom celery import Celery\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.access import source_should_fetch_permissions_during_indexing\nfrom onyx.background.indexing.checkpointing_utils import check_checkpoint_size\nfrom onyx.background.indexing.checkpointing_utils import get_latest_valid_checkpoint\nfrom onyx.background.indexing.checkpointing_utils import save_checkpoint\nfrom onyx.background.indexing.memory_tracer import MemoryTracer\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import INDEXING_SIZE_WARNING_THRESHOLD\nfrom onyx.configs.app_configs import INDEXING_TRACER_INTERVAL\nfrom onyx.configs.app_configs import INTEGRATION_TESTS_MODE\nfrom onyx.configs.app_configs import LEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE\nfrom onyx.configs.app_configs import MAX_FILE_SIZE_BYTES\nfrom onyx.configs.app_configs import POLL_CONNECTOR_OFFSET\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.connectors.connector_runner import ConnectorRunner\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.factory import instantiate_connector\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorStopSignal\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import IndexAttemptMetadata\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.connector import mark_ccpair_with_indexing_trigger\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.connector_credential_pair import get_last_successful_attempt_poll_range_end\nfrom onyx.db.connector_credential_pair import update_connector_credential_pair\nfrom onyx.db.constants import CONNECTOR_VALIDATION_ERROR_MESSAGE_PREFIX\nfrom onyx.db.document import mark_document_as_indexed_for_cc_pair__no_commit\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.enums import ProcessingMode\nfrom onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries\nfrom onyx.db.hierarchy import upsert_hierarchy_nodes_batch\nfrom onyx.db.index_attempt import create_index_attempt_error\nfrom onyx.db.index_attempt import get_index_attempt\nfrom onyx.db.index_attempt import get_recent_completed_attempts_for_cc_pair\nfrom onyx.db.index_attempt import mark_attempt_canceled\nfrom onyx.db.index_attempt import mark_attempt_failed\nfrom onyx.db.index_attempt import transition_attempt_to_in_progress\nfrom onyx.db.indexing_coordination import IndexingCoordination\nfrom onyx.db.models import IndexAttempt\nfrom onyx.file_store.document_batch_storage import DocumentBatchStorage\nfrom onyx.file_store.document_batch_storage import get_document_batch_storage\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.indexing.indexing_pipeline import index_doc_batch_prepare\nfrom onyx.redis.redis_hierarchy import cache_hierarchy_nodes_batch\nfrom onyx.redis.redis_hierarchy import ensure_source_node_exists\nfrom onyx.redis.redis_hierarchy import get_node_id_from_raw_id\nfrom onyx.redis.redis_hierarchy import get_source_node_id_from_cache\nfrom onyx.redis.redis_hierarchy import HierarchyNodeCacheEntry\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.features.build.indexing.persistent_document_writer import (\n get_persistent_document_writer,\n)\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.middleware import make_randomized_onyx_request_id\nfrom onyx.utils.postgres_sanitization import sanitize_document_for_postgres\nfrom onyx.utils.postgres_sanitization import sanitize_hierarchy_nodes_for_postgres\nfrom onyx.utils.variable_functionality import global_version\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR\n\nlogger = setup_logger(propagate=False)\n\nINDEXING_TRACER_NUM_PRINT_ENTRIES = 5\n\n\ndef _get_connector_runner(\n db_session: Session,\n attempt: IndexAttempt,\n batch_size: int,\n start_time: datetime,\n end_time: datetime,\n include_permissions: bool,\n leave_connector_active: bool = LEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE,\n) -> ConnectorRunner:\n \"\"\"\n NOTE: `start_time` and `end_time` are only used for poll connectors\n\n Returns an iterator of document batches and whether the returned documents\n are the complete list of existing documents of the connector. If the task\n of type LOAD_STATE, the list will be considered complete and otherwise incomplete.\n \"\"\"\n\n task = attempt.connector_credential_pair.connector.input_type\n\n try:\n runnable_connector = instantiate_connector(\n db_session=db_session,\n source=attempt.connector_credential_pair.connector.source,\n input_type=task,\n connector_specific_config=attempt.connector_credential_pair.connector.connector_specific_config,\n credential=attempt.connector_credential_pair.credential,\n )\n\n # validate the connector settings\n if not INTEGRATION_TESTS_MODE:\n runnable_connector.validate_connector_settings()\n if attempt.connector_credential_pair.access_type == AccessType.SYNC:\n runnable_connector.validate_perm_sync()\n\n except UnexpectedValidationError as e:\n logger.exception(\n \"Unable to instantiate connector due to an unexpected temporary issue.\"\n )\n raise e\n except Exception as e:\n logger.exception(\"Unable to instantiate connector. Pausing until fixed.\")\n # since we failed to even instantiate the connector, we pause the CCPair since\n # it will never succeed\n\n # Sometimes there are cases where the connector will\n # intermittently fail to initialize in which case we should pass in\n # leave_connector_active=True to allow it to continue.\n # For example, if there is nightly maintenance on a Confluence Server instance,\n # the connector will fail to initialize every night.\n if not leave_connector_active:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=attempt.connector_credential_pair.id,\n )\n if cc_pair and cc_pair.status == ConnectorCredentialPairStatus.ACTIVE:\n update_connector_credential_pair(\n db_session=db_session,\n connector_id=attempt.connector_credential_pair.connector.id,\n credential_id=attempt.connector_credential_pair.credential.id,\n status=ConnectorCredentialPairStatus.PAUSED,\n )\n raise e\n\n return ConnectorRunner(\n connector=runnable_connector,\n batch_size=batch_size,\n include_permissions=include_permissions,\n time_range=(start_time, end_time),\n )\n\n\ndef strip_null_characters(doc_batch: list[Document]) -> list[Document]:\n cleaned_batch = []\n for doc in doc_batch:\n if sys.getsizeof(doc) > MAX_FILE_SIZE_BYTES:\n logger.warning(\n f\"doc {doc.id} too large, Document size: {sys.getsizeof(doc)}\"\n )\n cleaned_batch.append(sanitize_document_for_postgres(doc))\n\n return cleaned_batch\n\n\ndef _check_connector_and_attempt_status(\n db_session_temp: Session,\n cc_pair_id: int,\n search_settings_status: IndexModelStatus,\n index_attempt_id: int,\n) -> None:\n \"\"\"\n Checks the status of the connector credential pair and index attempt.\n Raises a RuntimeError if any conditions are not met.\n \"\"\"\n cc_pair_loop = get_connector_credential_pair_from_id(\n db_session_temp,\n cc_pair_id,\n )\n if not cc_pair_loop:\n raise RuntimeError(f\"CC pair {cc_pair_id} not found in DB.\")\n\n if (\n cc_pair_loop.status == ConnectorCredentialPairStatus.PAUSED\n and search_settings_status != IndexModelStatus.FUTURE\n ) or cc_pair_loop.status == ConnectorCredentialPairStatus.DELETING:\n raise ConnectorStopSignal(f\"Connector {cc_pair_loop.status.value.lower()}\")\n\n index_attempt_loop = get_index_attempt(db_session_temp, index_attempt_id)\n if not index_attempt_loop:\n raise RuntimeError(f\"Index attempt {index_attempt_id} not found in DB.\")\n\n if index_attempt_loop.status == IndexingStatus.CANCELED:\n raise ConnectorStopSignal(f\"Index attempt {index_attempt_id} was canceled\")\n\n if index_attempt_loop.status != IndexingStatus.IN_PROGRESS:\n error_str = \"\"\n if index_attempt_loop.error_msg:\n error_str = f\" Original error: {index_attempt_loop.error_msg}\"\n\n raise RuntimeError(\n f\"Index Attempt is not running, status is {index_attempt_loop.status}.{error_str}\"\n )\n\n if index_attempt_loop.celery_task_id is None:\n raise RuntimeError(f\"Index attempt {index_attempt_id} has no celery task id\")\n\n\n# TODO: delete from here if ends up unused\ndef _check_failure_threshold(\n total_failures: int,\n document_count: int,\n batch_num: int,\n last_failure: ConnectorFailure | None,\n) -> None:\n \"\"\"Check if we've hit the failure threshold and raise an appropriate exception if so.\n\n We consider the threshold hit if:\n 1. We have more than 3 failures AND\n 2. Failures account for more than 10% of processed documents\n \"\"\"\n failure_ratio = total_failures / (document_count or 1)\n\n FAILURE_THRESHOLD = 3\n FAILURE_RATIO_THRESHOLD = 0.1\n if total_failures > FAILURE_THRESHOLD and failure_ratio > FAILURE_RATIO_THRESHOLD:\n logger.error(\n f\"Connector run failed with '{total_failures}' errors after '{batch_num}' batches.\"\n )\n if last_failure and last_failure.exception:\n raise last_failure.exception from last_failure.exception\n\n raise RuntimeError(\n f\"Connector run encountered too many errors, aborting. Last error: {last_failure}\"\n )\n\n\ndef run_docfetching_entrypoint(\n app: Celery,\n index_attempt_id: int,\n tenant_id: str,\n connector_credential_pair_id: int,\n is_ee: bool = False,\n callback: IndexingHeartbeatInterface | None = None,\n) -> None:\n \"\"\"Don't swallow exceptions here ... propagate them up.\"\"\"\n\n if is_ee:\n global_version.set_ee()\n\n # set the indexing attempt ID so that all log messages from this process\n # will have it added as a prefix\n token = INDEX_ATTEMPT_INFO_CONTEXTVAR.set(\n (connector_credential_pair_id, index_attempt_id)\n )\n with get_session_with_current_tenant() as db_session:\n attempt = transition_attempt_to_in_progress(index_attempt_id, db_session)\n\n tenant_str = \"\"\n if MULTI_TENANT:\n tenant_str = f\" for tenant {tenant_id}\"\n\n connector_name = attempt.connector_credential_pair.connector.name\n connector_config = (\n attempt.connector_credential_pair.connector.connector_specific_config\n )\n credential_id = attempt.connector_credential_pair.credential_id\n\n logger.info(\n f\"Docfetching starting{tenant_str}: \"\n f\"connector='{connector_name}' \"\n f\"config='{connector_config}' \"\n f\"credentials='{credential_id}'\"\n )\n\n connector_document_extraction(\n app,\n index_attempt_id,\n attempt.connector_credential_pair_id,\n attempt.search_settings_id,\n tenant_id,\n callback,\n )\n\n logger.info(\n f\"Docfetching finished{tenant_str}: \"\n f\"connector='{connector_name}' \"\n f\"config='{connector_config}' \"\n f\"credentials='{credential_id}'\"\n )\n\n INDEX_ATTEMPT_INFO_CONTEXTVAR.reset(token)\n\n\ndef connector_document_extraction(\n app: Celery,\n index_attempt_id: int,\n cc_pair_id: int,\n search_settings_id: int,\n tenant_id: str,\n callback: IndexingHeartbeatInterface | None = None,\n) -> None:\n \"\"\"Extract documents from connector and queue them for indexing pipeline processing.\n\n This is the first part of the split indexing process that runs the connector\n and extracts documents, storing them in the filestore for later processing.\n \"\"\"\n\n start_time = time.monotonic()\n\n logger.info(\n f\"Document extraction starting: \"\n f\"attempt={index_attempt_id} \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id} \"\n f\"tenant={tenant_id}\"\n )\n\n # Get batch storage (transition to IN_PROGRESS is handled by run_indexing_entrypoint)\n batch_storage = get_document_batch_storage(cc_pair_id, index_attempt_id)\n\n # Initialize memory tracer. NOTE: won't actually do anything if\n # `INDEXING_TRACER_INTERVAL` is 0.\n memory_tracer = MemoryTracer(interval=INDEXING_TRACER_INTERVAL)\n memory_tracer.start()\n\n index_attempt = None\n last_batch_num = 0 # used to continue from checkpointing\n # comes from _run_indexing\n with get_session_with_current_tenant() as db_session:\n index_attempt = get_index_attempt(\n db_session,\n index_attempt_id,\n eager_load_cc_pair=True,\n eager_load_search_settings=True,\n )\n if not index_attempt:\n raise RuntimeError(f\"Index attempt {index_attempt_id} not found\")\n\n if index_attempt.search_settings is None:\n raise ValueError(\"Search settings must be set for indexing\")\n\n # Clear the indexing trigger if it was set, to prevent duplicate indexing attempts\n if index_attempt.connector_credential_pair.indexing_trigger is not None:\n logger.info(\n \"Clearing indexing trigger: \"\n f\"cc_pair={index_attempt.connector_credential_pair.id} \"\n f\"trigger={index_attempt.connector_credential_pair.indexing_trigger}\"\n )\n mark_ccpair_with_indexing_trigger(\n index_attempt.connector_credential_pair.id, None, db_session\n )\n\n db_connector = index_attempt.connector_credential_pair.connector\n db_credential = index_attempt.connector_credential_pair.credential\n processing_mode = index_attempt.connector_credential_pair.processing_mode\n is_primary = index_attempt.search_settings.status == IndexModelStatus.PRESENT\n is_connector_public = (\n index_attempt.connector_credential_pair.access_type == AccessType.PUBLIC\n )\n\n from_beginning = index_attempt.from_beginning\n has_successful_attempt = (\n index_attempt.connector_credential_pair.last_successful_index_time\n is not None\n )\n # Use higher priority for first-time indexing to ensure new connectors\n # get processed before re-indexing of existing connectors\n docprocessing_priority = (\n OnyxCeleryPriority.MEDIUM\n if has_successful_attempt\n else OnyxCeleryPriority.HIGH\n )\n\n earliest_index_time = (\n db_connector.indexing_start.timestamp()\n if db_connector.indexing_start\n else 0\n )\n should_fetch_permissions_during_indexing = (\n index_attempt.connector_credential_pair.access_type == AccessType.SYNC\n and source_should_fetch_permissions_during_indexing(db_connector.source)\n and is_primary\n # if we've already successfully indexed, let the doc_sync job\n # take care of doc-level permissions\n and (from_beginning or not has_successful_attempt)\n )\n\n # Set up time windows for polling\n last_successful_index_poll_range_end = (\n earliest_index_time\n if from_beginning\n else get_last_successful_attempt_poll_range_end(\n cc_pair_id=cc_pair_id,\n earliest_index=earliest_index_time,\n search_settings=index_attempt.search_settings,\n db_session=db_session,\n )\n )\n\n if last_successful_index_poll_range_end > POLL_CONNECTOR_OFFSET:\n window_start = datetime.fromtimestamp(\n last_successful_index_poll_range_end, tz=timezone.utc\n ) - timedelta(minutes=POLL_CONNECTOR_OFFSET)\n else:\n # don't go into \"negative\" time if we've never indexed before\n window_start = datetime.fromtimestamp(0, tz=timezone.utc)\n\n most_recent_attempt = next(\n iter(\n get_recent_completed_attempts_for_cc_pair(\n cc_pair_id=cc_pair_id,\n search_settings_id=index_attempt.search_settings_id,\n db_session=db_session,\n limit=1,\n )\n ),\n None,\n )\n\n # if the last attempt failed, try and use the same window. This is necessary\n # to ensure correctness with checkpointing. If we don't do this, things like\n # new slack channels could be missed (since existing slack channels are\n # cached as part of the checkpoint).\n if (\n most_recent_attempt\n and most_recent_attempt.poll_range_end\n and (\n most_recent_attempt.status == IndexingStatus.FAILED\n or most_recent_attempt.status == IndexingStatus.CANCELED\n )\n ):\n window_end = most_recent_attempt.poll_range_end\n else:\n window_end = datetime.now(tz=timezone.utc)\n\n # set time range in db\n index_attempt.poll_range_start = window_start\n index_attempt.poll_range_end = window_end\n db_session.commit()\n\n # TODO: maybe memory tracer here\n\n # Set up connector runner\n connector_runner = _get_connector_runner(\n db_session=db_session,\n attempt=index_attempt,\n batch_size=INDEX_BATCH_SIZE,\n start_time=window_start,\n end_time=window_end,\n include_permissions=should_fetch_permissions_during_indexing,\n )\n\n # don't use a checkpoint if we're explicitly indexing from\n # the beginning in order to avoid weird interactions between\n # checkpointing / failure handling\n # OR\n # if the last attempt was successful\n if index_attempt.from_beginning or (\n most_recent_attempt and most_recent_attempt.status.is_successful()\n ):\n logger.info(\n f\"Cleaning up all old batches for index attempt {index_attempt_id} before starting new run\"\n )\n batch_storage.cleanup_all_batches()\n checkpoint = connector_runner.connector.build_dummy_checkpoint()\n else:\n logger.info(\n f\"Getting latest valid checkpoint for index attempt {index_attempt_id}\"\n )\n checkpoint, resuming_from_checkpoint = get_latest_valid_checkpoint(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n search_settings_id=index_attempt.search_settings_id,\n window_start=window_start,\n window_end=window_end,\n connector=connector_runner.connector,\n )\n\n # checkpoint resumption OR the connector already finished.\n if (\n isinstance(connector_runner.connector, CheckpointedConnector)\n and resuming_from_checkpoint\n ) or (\n most_recent_attempt\n and most_recent_attempt.total_batches is not None\n and not checkpoint.has_more\n ):\n reissued_batch_count, completed_batches = reissue_old_batches(\n batch_storage,\n index_attempt_id,\n cc_pair_id,\n tenant_id,\n app,\n most_recent_attempt,\n docprocessing_priority,\n )\n last_batch_num = reissued_batch_count + completed_batches\n index_attempt.completed_batches = completed_batches\n db_session.commit()\n else:\n logger.info(\n f\"Cleaning up all batches for index attempt {index_attempt_id} before starting new run\"\n )\n # for non-checkpointed connectors, throw out batches from previous unsuccessful attempts\n # because we'll be getting those documents again anyways.\n batch_storage.cleanup_all_batches()\n\n # Save initial checkpoint\n save_checkpoint(\n db_session=db_session,\n index_attempt_id=index_attempt_id,\n checkpoint=checkpoint,\n )\n\n try:\n batch_num = last_batch_num # starts at 0 if no last batch\n total_doc_batches_queued = 0\n total_failures = 0\n document_count = 0\n\n # Ensure the SOURCE-type root hierarchy node exists before processing.\n # This is the root of the hierarchy tree for this source - all other\n # hierarchy nodes should ultimately have this as an ancestor.\n redis_client = get_redis_client(tenant_id=tenant_id)\n with get_session_with_current_tenant() as db_session:\n ensure_source_node_exists(redis_client, db_session, db_connector.source)\n\n # Main extraction loop\n while checkpoint.has_more:\n logger.info(\n f\"Running '{db_connector.source.value}' connector with checkpoint: {checkpoint}\"\n )\n for (\n document_batch,\n hierarchy_node_batch,\n failure,\n next_checkpoint,\n ) in connector_runner.run(checkpoint):\n # Check if connector is disabled mid run and stop if so unless it's the secondary\n # index being built. We want to populate it even for paused connectors\n # Often paused connectors are sources that aren't updated frequently but the\n # contents still need to be initially pulled.\n if callback and callback.should_stop():\n raise ConnectorStopSignal(\"Connector stop signal detected\")\n\n # will exception if the connector/index attempt is marked as paused/failed\n with get_session_with_current_tenant() as db_session_tmp:\n _check_connector_and_attempt_status(\n db_session_tmp,\n cc_pair_id,\n index_attempt.search_settings.status,\n index_attempt_id,\n )\n\n # save record of any failures at the connector level\n if failure is not None:\n total_failures += 1\n with get_session_with_current_tenant() as db_session:\n create_index_attempt_error(\n index_attempt_id,\n cc_pair_id,\n failure,\n db_session,\n )\n _check_failure_threshold(\n total_failures, document_count, batch_num, failure\n )\n\n # Save checkpoint if provided\n if next_checkpoint:\n checkpoint = next_checkpoint\n\n # Process hierarchy nodes batch - upsert to Postgres and cache in Redis\n if hierarchy_node_batch:\n hierarchy_node_batch_cleaned = (\n sanitize_hierarchy_nodes_for_postgres(hierarchy_node_batch)\n )\n with get_session_with_current_tenant() as db_session:\n upserted_nodes = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=hierarchy_node_batch_cleaned,\n source=db_connector.source,\n commit=True,\n is_connector_public=is_connector_public,\n )\n\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=[n.id for n in upserted_nodes],\n connector_id=db_connector.id,\n credential_id=db_credential.id,\n commit=True,\n )\n\n # Cache in Redis for fast ancestor resolution during doc processing\n redis_client = get_redis_client(tenant_id=tenant_id)\n cache_entries = [\n HierarchyNodeCacheEntry.from_db_model(node)\n for node in upserted_nodes\n ]\n cache_hierarchy_nodes_batch(\n redis_client=redis_client,\n source=db_connector.source,\n entries=cache_entries,\n )\n\n logger.debug(\n f\"Persisted and cached {len(hierarchy_node_batch_cleaned)} hierarchy nodes for attempt={index_attempt_id}\"\n )\n\n # below is all document processing task, so if no batch we can just continue\n if not document_batch:\n continue\n\n # Clean documents and create batch\n doc_batch_cleaned = strip_null_characters(document_batch)\n\n # Resolve parent_hierarchy_raw_node_id to parent_hierarchy_node_id\n # using the Redis cache (just populated from hierarchy nodes batch)\n with get_session_with_current_tenant() as db_session_tmp:\n source_node_id = get_source_node_id_from_cache(\n redis_client, db_session_tmp, db_connector.source\n )\n for doc in doc_batch_cleaned:\n if doc.parent_hierarchy_raw_node_id is not None:\n node_id, found = get_node_id_from_raw_id(\n redis_client,\n db_connector.source,\n doc.parent_hierarchy_raw_node_id,\n )\n doc.parent_hierarchy_node_id = (\n node_id if found else source_node_id\n )\n else:\n doc.parent_hierarchy_node_id = source_node_id\n\n batch_description = []\n\n for doc in doc_batch_cleaned:\n batch_description.append(doc.to_short_descriptor())\n\n doc_size = 0\n for section in doc.sections:\n if (\n isinstance(section, TextSection)\n and section.text is not None\n ):\n doc_size += len(section.text)\n\n if doc_size > INDEXING_SIZE_WARNING_THRESHOLD:\n logger.warning(\n f\"Document size: doc='{doc.to_short_descriptor()}' \"\n f\"size={doc_size} \"\n f\"threshold={INDEXING_SIZE_WARNING_THRESHOLD}\"\n )\n\n logger.debug(f\"Indexing batch of documents: {batch_description}\")\n memory_tracer.increment_and_maybe_trace()\n\n if processing_mode == ProcessingMode.FILE_SYSTEM:\n # File system only - write directly to persistent storage,\n # skip chunking/embedding/Vespa but still track documents in DB\n\n # IMPORTANT: Write to S3 FIRST, before marking as indexed in DB.\n\n # Write documents to persistent file system\n # Use creator_id for user-segregated storage paths (sandbox isolation)\n creator_id = index_attempt.connector_credential_pair.creator_id\n if creator_id is None:\n raise ValueError(\n f\"ConnectorCredentialPair {index_attempt.connector_credential_pair.id} \"\n \"must have a creator_id for persistent document storage\"\n )\n user_id_str: str = str(creator_id)\n writer = get_persistent_document_writer(\n user_id=user_id_str,\n tenant_id=tenant_id,\n )\n written_paths = writer.write_documents(doc_batch_cleaned)\n\n # Only after successful S3 write, mark documents as indexed in DB\n with get_session_with_current_tenant() as db_session:\n # Create metadata for the batch\n index_attempt_metadata = IndexAttemptMetadata(\n attempt_id=index_attempt_id,\n connector_id=db_connector.id,\n credential_id=db_credential.id,\n request_id=make_randomized_onyx_request_id(\"FSI\"),\n structured_id=f\"{tenant_id}:{cc_pair_id}:{index_attempt_id}:{batch_num}\",\n batch_num=batch_num,\n )\n\n # Upsert documents to PostgreSQL (document table + cc_pair relationship)\n # This is a subset of what docprocessing does - just DB tracking, no chunking/embedding\n index_doc_batch_prepare(\n documents=doc_batch_cleaned,\n index_attempt_metadata=index_attempt_metadata,\n db_session=db_session,\n ignore_time_skip=True, # Documents already filtered during extraction\n )\n\n # Mark documents as indexed for the CC pair\n mark_document_as_indexed_for_cc_pair__no_commit(\n connector_id=db_connector.id,\n credential_id=db_credential.id,\n document_ids=[doc.id for doc in doc_batch_cleaned],\n db_session=db_session,\n )\n db_session.commit()\n\n # Update coordination directly (no docprocessing task)\n with get_session_with_current_tenant() as db_session:\n IndexingCoordination.update_batch_completion_and_docs(\n db_session=db_session,\n index_attempt_id=index_attempt_id,\n total_docs_indexed=len(doc_batch_cleaned),\n new_docs_indexed=len(doc_batch_cleaned),\n total_chunks=0, # No chunks for file system mode\n )\n\n batch_num += 1\n total_doc_batches_queued += 1\n\n logger.info(\n f\"Wrote documents to file system: \"\n f\"batch_num={batch_num} \"\n f\"docs={len(written_paths)} \"\n f\"attempt={index_attempt_id}\"\n )\n else:\n # REGULAR mode (default): Full pipeline - store and queue docprocessing\n batch_storage.store_batch(batch_num, doc_batch_cleaned)\n\n # Create processing task data\n processing_batch_data = {\n \"index_attempt_id\": index_attempt_id,\n \"cc_pair_id\": cc_pair_id,\n \"tenant_id\": tenant_id,\n \"batch_num\": batch_num, # 0-indexed\n }\n\n # Queue document processing task\n app.send_task(\n OnyxCeleryTask.DOCPROCESSING_TASK,\n kwargs=processing_batch_data,\n queue=OnyxCeleryQueues.DOCPROCESSING,\n priority=docprocessing_priority,\n )\n\n batch_num += 1\n total_doc_batches_queued += 1\n\n logger.info(\n f\"Queued document processing batch: \"\n f\"batch_num={batch_num} \"\n f\"docs={len(doc_batch_cleaned)} \"\n f\"attempt={index_attempt_id}\"\n )\n\n # Check checkpoint size periodically\n CHECKPOINT_SIZE_CHECK_INTERVAL = 100\n if batch_num % CHECKPOINT_SIZE_CHECK_INTERVAL == 0:\n check_checkpoint_size(checkpoint)\n\n # Save latest checkpoint\n # NOTE: checkpointing is used to track which batches have\n # been sent to the filestore, NOT which batches have been fully indexed\n # as it used to be.\n with get_session_with_current_tenant() as db_session:\n save_checkpoint(\n db_session=db_session,\n index_attempt_id=index_attempt_id,\n checkpoint=checkpoint,\n )\n\n elapsed_time = time.monotonic() - start_time\n\n logger.info(\n f\"Document extraction completed: \"\n f\"attempt={index_attempt_id} \"\n f\"batches_queued={total_doc_batches_queued} \"\n f\"elapsed={elapsed_time:.2f}s\"\n )\n\n # Set total batches in database to signal extraction completion.\n # Used by check_for_indexing to determine if the index attempt is complete.\n with get_session_with_current_tenant() as db_session:\n IndexingCoordination.set_total_batches(\n db_session=db_session,\n index_attempt_id=index_attempt_id,\n total_batches=batch_num,\n )\n\n # Trigger file sync to user's sandbox (if running) - only for FILE_SYSTEM mode\n # This syncs the newly written documents from S3 to any running sandbox pod\n if processing_mode == ProcessingMode.FILE_SYSTEM:\n creator_id = index_attempt.connector_credential_pair.creator_id\n if creator_id:\n source_value = db_connector.source.value\n app.send_task(\n OnyxCeleryTask.SANDBOX_FILE_SYNC,\n kwargs={\n \"user_id\": str(creator_id),\n \"tenant_id\": tenant_id,\n \"source\": source_value,\n },\n queue=OnyxCeleryQueues.SANDBOX,\n )\n logger.info(\n f\"Triggered sandbox file sync for user {creator_id} source={source_value} after indexing complete\"\n )\n\n except Exception as e:\n logger.exception(\n f\"Document extraction failed: attempt={index_attempt_id} error={str(e)}\"\n )\n\n # Do NOT clean up batches on failure; future runs will use those batches\n # while docfetching will continue from the saved checkpoint if one exists\n\n if isinstance(e, ConnectorValidationError):\n # On validation errors during indexing, we want to cancel the indexing attempt\n # and mark the CCPair as invalid. This prevents the connector from being\n # used in the future until the credentials are updated.\n with get_session_with_current_tenant() as db_session_temp:\n logger.exception(\n f\"Marking attempt {index_attempt_id} as canceled due to validation error.\"\n )\n mark_attempt_canceled(\n index_attempt_id,\n db_session_temp,\n reason=f\"{CONNECTOR_VALIDATION_ERROR_MESSAGE_PREFIX}{str(e)}\",\n )\n\n if is_primary:\n if not index_attempt:\n # should always be set by now\n raise RuntimeError(\"Should never happen.\")\n\n VALIDATION_ERROR_THRESHOLD = 5\n\n recent_index_attempts = get_recent_completed_attempts_for_cc_pair(\n cc_pair_id=cc_pair_id,\n search_settings_id=index_attempt.search_settings_id,\n limit=VALIDATION_ERROR_THRESHOLD,\n db_session=db_session_temp,\n )\n num_validation_errors = len(\n [\n index_attempt\n for index_attempt in recent_index_attempts\n if index_attempt.error_msg\n and index_attempt.error_msg.startswith(\n CONNECTOR_VALIDATION_ERROR_MESSAGE_PREFIX\n )\n ]\n )\n\n if num_validation_errors >= VALIDATION_ERROR_THRESHOLD:\n logger.warning(\n f\"Connector {db_connector.id} has {num_validation_errors} consecutive validation\"\n f\" errors. Marking the CC Pair as invalid.\"\n )\n update_connector_credential_pair(\n db_session=db_session_temp,\n connector_id=db_connector.id,\n credential_id=db_credential.id,\n status=ConnectorCredentialPairStatus.INVALID,\n )\n raise e\n elif isinstance(e, ConnectorStopSignal):\n with get_session_with_current_tenant() as db_session_temp:\n logger.exception(\n f\"Marking attempt {index_attempt_id} as canceled due to stop signal.\"\n )\n mark_attempt_canceled(\n index_attempt_id,\n db_session_temp,\n reason=str(e),\n )\n\n else:\n with get_session_with_current_tenant() as db_session_temp:\n # don't overwrite attempts that are already failed/canceled for another reason\n index_attempt = get_index_attempt(db_session_temp, index_attempt_id)\n if index_attempt and index_attempt.status in [\n IndexingStatus.CANCELED,\n IndexingStatus.FAILED,\n ]:\n logger.info(\n f\"Attempt {index_attempt_id} is already failed/canceled, skipping marking as failed.\"\n )\n raise e\n\n mark_attempt_failed(\n index_attempt_id,\n db_session_temp,\n failure_reason=str(e),\n full_exception_trace=traceback.format_exc(),\n )\n\n raise e\n\n finally:\n memory_tracer.stop()\n\n\ndef reissue_old_batches(\n batch_storage: DocumentBatchStorage,\n index_attempt_id: int,\n cc_pair_id: int,\n tenant_id: str,\n app: Celery,\n most_recent_attempt: IndexAttempt | None,\n priority: OnyxCeleryPriority,\n) -> tuple[int, int]:\n # When loading from a checkpoint, we need to start new docprocessing tasks\n # tied to the new index attempt for any batches left over in the file store\n old_batches = batch_storage.get_all_batches_for_cc_pair()\n batch_storage.update_old_batches_to_new_index_attempt(old_batches)\n for batch_id in old_batches:\n logger.info(\n f\"Re-issuing docprocessing task for batch {batch_id} for index attempt {index_attempt_id}\"\n )\n path_info = batch_storage.extract_path_info(batch_id)\n if path_info is None:\n logger.warning(\n f\"Could not extract path info from batch {batch_id}, skipping\"\n )\n continue\n if path_info.cc_pair_id != cc_pair_id:\n raise RuntimeError(f\"Batch {batch_id} is not for cc pair {cc_pair_id}\")\n\n app.send_task(\n OnyxCeleryTask.DOCPROCESSING_TASK,\n kwargs={\n \"index_attempt_id\": index_attempt_id,\n \"cc_pair_id\": cc_pair_id,\n \"tenant_id\": tenant_id,\n \"batch_num\": path_info.batch_num, # use same batch num as previously\n },\n queue=OnyxCeleryQueues.DOCPROCESSING,\n priority=priority,\n )\n recent_batches = most_recent_attempt.completed_batches if most_recent_attempt else 0\n # resume from the batch num of the last attempt. This should be one more\n # than the last batch created by docfetching regardless of whether the batch\n # is still in the filestore waiting for processing or not.\n last_batch_num = len(old_batches) + recent_batches\n logger.info(\n f\"Starting from batch {last_batch_num} due to re-issued batches: {old_batches}, completed batches: {recent_batches}\"\n )\n return len(old_batches), recent_batches\n" }, { "path": "backend/onyx/background/periodic_poller.py", "content": "\"\"\"Periodic poller for NO_VECTOR_DB deployments.\n\nReplaces Celery Beat and background workers with a lightweight daemon thread\nthat runs from the API server process. Two responsibilities:\n\n1. Recovery polling (every 30 s): re-processes user files stuck in\n PROCESSING / DELETING / needs_sync states via the drain loops defined\n in ``task_utils.py``.\n\n2. Periodic task execution (configurable intervals): runs LLM model updates\n and scheduled evals at their configured cadences, with Postgres advisory\n lock deduplication across multiple API server instances.\n\"\"\"\n\nimport threading\nimport time\nfrom collections.abc import Callable\nfrom dataclasses import dataclass\nfrom dataclasses import field\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nRECOVERY_INTERVAL_SECONDS = 30\nPERIODIC_TASK_LOCK_BASE = 20_000\nPERIODIC_TASK_KV_PREFIX = \"periodic_poller:last_claimed:\"\n\n\n# ------------------------------------------------------------------\n# Periodic task definitions\n# ------------------------------------------------------------------\n\n\n_NEVER_RAN: float = -1e18\n\n\n@dataclass\nclass _PeriodicTaskDef:\n name: str\n interval_seconds: float\n lock_id: int\n run_fn: Callable[[], None]\n last_run_at: float = field(default=_NEVER_RAN)\n\n\ndef _run_auto_llm_update() -> None:\n from onyx.configs.app_configs import AUTO_LLM_CONFIG_URL\n\n if not AUTO_LLM_CONFIG_URL:\n return\n\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n from onyx.llm.well_known_providers.auto_update_service import (\n sync_llm_models_from_github,\n )\n\n with get_session_with_current_tenant() as db_session:\n sync_llm_models_from_github(db_session)\n\n\ndef _run_cache_cleanup() -> None:\n from onyx.cache.postgres_backend import cleanup_expired_cache_entries\n\n cleanup_expired_cache_entries()\n\n\ndef _run_scheduled_eval() -> None:\n from onyx.configs.app_configs import BRAINTRUST_API_KEY\n from onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES\n from onyx.configs.app_configs import SCHEDULED_EVAL_PERMISSIONS_EMAIL\n from onyx.configs.app_configs import SCHEDULED_EVAL_PROJECT\n\n if not all(\n [\n BRAINTRUST_API_KEY,\n SCHEDULED_EVAL_PROJECT,\n SCHEDULED_EVAL_DATASET_NAMES,\n SCHEDULED_EVAL_PERMISSIONS_EMAIL,\n ]\n ):\n return\n\n from datetime import datetime\n from datetime import timezone\n\n from onyx.evals.eval import run_eval\n from onyx.evals.models import EvalConfigurationOptions\n\n run_timestamp = datetime.now(timezone.utc).strftime(\"%Y-%m-%d\")\n for dataset_name in SCHEDULED_EVAL_DATASET_NAMES:\n try:\n run_eval(\n configuration=EvalConfigurationOptions(\n search_permissions_email=SCHEDULED_EVAL_PERMISSIONS_EMAIL,\n dataset_name=dataset_name,\n no_send_logs=False,\n braintrust_project=SCHEDULED_EVAL_PROJECT,\n experiment_name=f\"{dataset_name} - {run_timestamp}\",\n ),\n remote_dataset_name=dataset_name,\n )\n except Exception:\n logger.exception(\n f\"Periodic poller - Failed scheduled eval for dataset {dataset_name}\"\n )\n\n\n_CACHE_CLEANUP_INTERVAL_SECONDS = 300\n\n\ndef _build_periodic_tasks() -> list[_PeriodicTaskDef]:\n from onyx.cache.interface import CacheBackendType\n from onyx.configs.app_configs import AUTO_LLM_CONFIG_URL\n from onyx.configs.app_configs import AUTO_LLM_UPDATE_INTERVAL_SECONDS\n from onyx.configs.app_configs import CACHE_BACKEND\n from onyx.configs.app_configs import SCHEDULED_EVAL_DATASET_NAMES\n\n tasks: list[_PeriodicTaskDef] = []\n if CACHE_BACKEND == CacheBackendType.POSTGRES:\n tasks.append(\n _PeriodicTaskDef(\n name=\"cache-cleanup\",\n interval_seconds=_CACHE_CLEANUP_INTERVAL_SECONDS,\n lock_id=PERIODIC_TASK_LOCK_BASE + 2,\n run_fn=_run_cache_cleanup,\n )\n )\n if AUTO_LLM_CONFIG_URL:\n tasks.append(\n _PeriodicTaskDef(\n name=\"auto-llm-update\",\n interval_seconds=AUTO_LLM_UPDATE_INTERVAL_SECONDS,\n lock_id=PERIODIC_TASK_LOCK_BASE,\n run_fn=_run_auto_llm_update,\n )\n )\n if SCHEDULED_EVAL_DATASET_NAMES:\n tasks.append(\n _PeriodicTaskDef(\n name=\"scheduled-eval\",\n interval_seconds=7 * 24 * 3600,\n lock_id=PERIODIC_TASK_LOCK_BASE + 1,\n run_fn=_run_scheduled_eval,\n )\n )\n return tasks\n\n\n# ------------------------------------------------------------------\n# Periodic task runner with advisory-lock-guarded claim\n# ------------------------------------------------------------------\n\n\ndef _try_claim_task(task_def: _PeriodicTaskDef) -> bool:\n \"\"\"Atomically check whether *task_def* should run and record a claim.\n\n Uses a transaction-scoped advisory lock for atomicity combined with a\n ``KVStore`` timestamp for cross-instance dedup. The DB session is held\n only for this brief claim transaction, not during task execution.\n \"\"\"\n from datetime import datetime\n from datetime import timezone\n\n from sqlalchemy import text\n\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n from onyx.db.models import KVStore\n\n kv_key = PERIODIC_TASK_KV_PREFIX + task_def.name\n\n with get_session_with_current_tenant() as db_session:\n acquired = db_session.execute(\n text(\"SELECT pg_try_advisory_xact_lock(:id)\"),\n {\"id\": task_def.lock_id},\n ).scalar()\n if not acquired:\n return False\n\n row = db_session.query(KVStore).filter_by(key=kv_key).first()\n if row and row.value is not None:\n last_claimed = datetime.fromisoformat(str(row.value))\n elapsed = (datetime.now(timezone.utc) - last_claimed).total_seconds()\n if elapsed < task_def.interval_seconds:\n return False\n\n now_ts = datetime.now(timezone.utc).isoformat()\n if row:\n row.value = now_ts\n else:\n db_session.add(KVStore(key=kv_key, value=now_ts))\n db_session.commit()\n\n return True\n\n\ndef _try_run_periodic_task(task_def: _PeriodicTaskDef) -> None:\n \"\"\"Run *task_def* if its interval has elapsed and no peer holds the lock.\"\"\"\n now = time.monotonic()\n if now - task_def.last_run_at < task_def.interval_seconds:\n return\n\n if not _try_claim_task(task_def):\n return\n\n try:\n task_def.run_fn()\n task_def.last_run_at = now\n except Exception:\n logger.exception(\n f\"Periodic poller - Error running periodic task {task_def.name}\"\n )\n\n\n# ------------------------------------------------------------------\n# Recovery / drain loop runner\n# ------------------------------------------------------------------\n\n\ndef _run_drain_loops(tenant_id: str) -> None:\n from onyx.background.task_utils import drain_delete_loop\n from onyx.background.task_utils import drain_processing_loop\n from onyx.background.task_utils import drain_project_sync_loop\n\n drain_processing_loop(tenant_id)\n drain_delete_loop(tenant_id)\n drain_project_sync_loop(tenant_id)\n\n\n# ------------------------------------------------------------------\n# Startup recovery (10g)\n# ------------------------------------------------------------------\n\n\ndef recover_stuck_user_files(tenant_id: str) -> None:\n \"\"\"Run all drain loops once to re-process files left in intermediate states.\n\n Called from ``lifespan()`` on startup when ``DISABLE_VECTOR_DB`` is set.\n \"\"\"\n logger.info(\"recover_stuck_user_files - Checking for stuck user files\")\n try:\n _run_drain_loops(tenant_id)\n except Exception:\n logger.exception(\"recover_stuck_user_files - Error during recovery\")\n\n\n# ------------------------------------------------------------------\n# Daemon thread (10f)\n# ------------------------------------------------------------------\n\n_shutdown_event = threading.Event()\n_poller_thread: threading.Thread | None = None\n\n\ndef _poller_loop(tenant_id: str) -> None:\n from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n\n periodic_tasks = _build_periodic_tasks()\n logger.info(\n f\"Periodic poller started with {len(periodic_tasks)} periodic task(s): {[t.name for t in periodic_tasks]}\"\n )\n\n while not _shutdown_event.is_set():\n try:\n _run_drain_loops(tenant_id)\n except Exception:\n logger.exception(\"Periodic poller - Error in recovery polling\")\n\n for task_def in periodic_tasks:\n try:\n _try_run_periodic_task(task_def)\n except Exception:\n logger.exception(\n f\"Periodic poller - Unhandled error checking task {task_def.name}\"\n )\n\n _shutdown_event.wait(RECOVERY_INTERVAL_SECONDS)\n\n\ndef start_periodic_poller(tenant_id: str) -> None:\n \"\"\"Start the periodic poller daemon thread.\"\"\"\n global _poller_thread # noqa: PLW0603\n _shutdown_event.clear()\n _poller_thread = threading.Thread(\n target=_poller_loop,\n args=(tenant_id,),\n daemon=True,\n name=\"no-vectordb-periodic-poller\",\n )\n _poller_thread.start()\n logger.info(\"Periodic poller thread started\")\n\n\ndef stop_periodic_poller() -> None:\n \"\"\"Signal the periodic poller to stop and wait for it to exit.\"\"\"\n global _poller_thread # noqa: PLW0603\n if _poller_thread is None:\n return\n _shutdown_event.set()\n _poller_thread.join(timeout=10)\n if _poller_thread.is_alive():\n logger.warning(\"Periodic poller thread did not stop within timeout\")\n _poller_thread = None\n logger.info(\"Periodic poller thread stopped\")\n" }, { "path": "backend/onyx/background/task_utils.py", "content": "\"\"\"Background task utilities.\n\nContains query-history report helpers (used by all deployment modes) and\nin-process background task execution helpers for NO_VECTOR_DB mode:\n\n- Atomic claim-and-mark helpers that prevent duplicate processing\n- Drain loops that process all pending user file work\n\nEach claim function runs a short-lived transaction: SELECT ... FOR UPDATE\nSKIP LOCKED, UPDATE the row to remove it from future queries, COMMIT.\nAfter the commit the row lock is released, but the row is no longer\neligible for re-claiming. No long-lived sessions or advisory locks.\n\"\"\"\n\nfrom uuid import UUID\n\nimport sqlalchemy as sa\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import UserFile\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# ------------------------------------------------------------------\n# Query-history report helpers (pre-existing, used by all modes)\n# ------------------------------------------------------------------\n\nQUERY_REPORT_NAME_PREFIX = \"query-history\"\n\n\ndef construct_query_history_report_name(\n task_id: str,\n) -> str:\n return f\"{QUERY_REPORT_NAME_PREFIX}-{task_id}.csv\"\n\n\ndef extract_task_id_from_query_history_report_name(name: str) -> str:\n return name.removeprefix(f\"{QUERY_REPORT_NAME_PREFIX}-\").removesuffix(\".csv\")\n\n\n# ------------------------------------------------------------------\n# Atomic claim-and-mark helpers\n# ------------------------------------------------------------------\n# Each function runs inside a single short-lived session/transaction:\n# 1. SELECT ... FOR UPDATE SKIP LOCKED (locks one eligible row)\n# 2. UPDATE the row so it is no longer eligible\n# 3. COMMIT (releases the row lock)\n# After the commit, no other drain loop can claim the same row.\n\n\ndef _claim_next_processing_file(db_session: Session) -> UUID | None:\n \"\"\"Claim the next PROCESSING file by transitioning it to INDEXING.\n\n Returns the file id, or None when no eligible files remain.\n \"\"\"\n file_id = db_session.execute(\n select(UserFile.id)\n .where(UserFile.status == UserFileStatus.PROCESSING)\n .order_by(UserFile.created_at)\n .limit(1)\n .with_for_update(skip_locked=True)\n ).scalar_one_or_none()\n if file_id is None:\n return None\n\n db_session.execute(\n sa.update(UserFile)\n .where(UserFile.id == file_id)\n .values(status=UserFileStatus.INDEXING)\n )\n db_session.commit()\n return file_id\n\n\ndef _claim_next_deleting_file(\n db_session: Session,\n exclude_ids: set[UUID] | None = None,\n) -> UUID | None:\n \"\"\"Claim the next DELETING file.\n\n No status transition needed — the impl deletes the row on success.\n The short-lived FOR UPDATE lock prevents concurrent claims.\n *exclude_ids* prevents re-processing the same file if the impl fails.\n \"\"\"\n stmt = (\n select(UserFile.id)\n .where(UserFile.status == UserFileStatus.DELETING)\n .order_by(UserFile.created_at)\n .limit(1)\n .with_for_update(skip_locked=True)\n )\n if exclude_ids:\n stmt = stmt.where(UserFile.id.notin_(exclude_ids))\n file_id = db_session.execute(stmt).scalar_one_or_none()\n db_session.commit()\n return file_id\n\n\ndef _claim_next_sync_file(\n db_session: Session,\n exclude_ids: set[UUID] | None = None,\n) -> UUID | None:\n \"\"\"Claim the next file needing project/persona sync.\n\n No status transition needed — the impl clears the sync flags on\n success. The short-lived FOR UPDATE lock prevents concurrent claims.\n *exclude_ids* prevents re-processing the same file if the impl fails.\n \"\"\"\n stmt = (\n select(UserFile.id)\n .where(\n sa.and_(\n sa.or_(\n UserFile.needs_project_sync.is_(True),\n UserFile.needs_persona_sync.is_(True),\n ),\n UserFile.status == UserFileStatus.COMPLETED,\n )\n )\n .order_by(UserFile.created_at)\n .limit(1)\n .with_for_update(skip_locked=True)\n )\n if exclude_ids:\n stmt = stmt.where(UserFile.id.notin_(exclude_ids))\n file_id = db_session.execute(stmt).scalar_one_or_none()\n db_session.commit()\n return file_id\n\n\n# ------------------------------------------------------------------\n# Drain loops — process *all* pending work of each type\n# ------------------------------------------------------------------\n\n\ndef drain_processing_loop(tenant_id: str) -> None:\n \"\"\"Process all pending PROCESSING user files.\"\"\"\n from onyx.background.celery.tasks.user_file_processing.tasks import (\n process_user_file_impl,\n )\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n while True:\n with get_session_with_current_tenant() as session:\n file_id = _claim_next_processing_file(session)\n if file_id is None:\n break\n try:\n process_user_file_impl(\n user_file_id=str(file_id),\n tenant_id=tenant_id,\n redis_locking=False,\n )\n except Exception:\n logger.exception(f\"Failed to process user file {file_id}\")\n\n\ndef drain_delete_loop(tenant_id: str) -> None:\n \"\"\"Delete all pending DELETING user files.\"\"\"\n from onyx.background.celery.tasks.user_file_processing.tasks import (\n delete_user_file_impl,\n )\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n failed: set[UUID] = set()\n while True:\n with get_session_with_current_tenant() as session:\n file_id = _claim_next_deleting_file(session, exclude_ids=failed)\n if file_id is None:\n break\n try:\n delete_user_file_impl(\n user_file_id=str(file_id),\n tenant_id=tenant_id,\n redis_locking=False,\n )\n except Exception:\n logger.exception(f\"Failed to delete user file {file_id}\")\n failed.add(file_id)\n\n\ndef drain_project_sync_loop(tenant_id: str) -> None:\n \"\"\"Sync all pending project/persona metadata for user files.\"\"\"\n from onyx.background.celery.tasks.user_file_processing.tasks import (\n project_sync_user_file_impl,\n )\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n failed: set[UUID] = set()\n while True:\n with get_session_with_current_tenant() as session:\n file_id = _claim_next_sync_file(session, exclude_ids=failed)\n if file_id is None:\n break\n try:\n project_sync_user_file_impl(\n user_file_id=str(file_id),\n tenant_id=tenant_id,\n redis_locking=False,\n )\n except Exception:\n logger.exception(f\"Failed to sync user file {file_id}\")\n failed.add(file_id)\n" }, { "path": "backend/onyx/cache/factory.py", "content": "from collections.abc import Callable\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.cache.interface import CacheBackendType\nfrom onyx.configs.app_configs import CACHE_BACKEND\n\n\ndef _build_redis_backend(tenant_id: str) -> CacheBackend:\n from onyx.cache.redis_backend import RedisCacheBackend\n from onyx.redis.redis_pool import redis_pool\n\n return RedisCacheBackend(redis_pool.get_client(tenant_id))\n\n\ndef _build_postgres_backend(tenant_id: str) -> CacheBackend:\n from onyx.cache.postgres_backend import PostgresCacheBackend\n\n return PostgresCacheBackend(tenant_id)\n\n\n_BACKEND_BUILDERS: dict[CacheBackendType, Callable[[str], CacheBackend]] = {\n CacheBackendType.REDIS: _build_redis_backend,\n CacheBackendType.POSTGRES: _build_postgres_backend,\n}\n\n\ndef get_cache_backend(*, tenant_id: str | None = None) -> CacheBackend:\n \"\"\"Return a tenant-aware ``CacheBackend``.\n\n If *tenant_id* is ``None``, the current tenant is read from the\n thread-local context variable (same behaviour as ``get_redis_client``).\n \"\"\"\n if tenant_id is None:\n from shared_configs.contextvars import get_current_tenant_id\n\n tenant_id = get_current_tenant_id()\n\n builder = _BACKEND_BUILDERS.get(CACHE_BACKEND)\n if builder is None:\n raise ValueError(\n f\"Unsupported CACHE_BACKEND={CACHE_BACKEND!r}. Supported values: {[t.value for t in CacheBackendType]}\"\n )\n return builder(tenant_id)\n\n\ndef get_shared_cache_backend() -> CacheBackend:\n \"\"\"Return a ``CacheBackend`` in the shared (cross-tenant) namespace.\"\"\"\n from shared_configs.configs import DEFAULT_REDIS_PREFIX\n\n return get_cache_backend(tenant_id=DEFAULT_REDIS_PREFIX)\n" }, { "path": "backend/onyx/cache/interface.py", "content": "import abc\nfrom enum import Enum\n\nfrom redis.exceptions import RedisError\nfrom sqlalchemy.exc import SQLAlchemyError\n\nTTL_KEY_NOT_FOUND = -2\nTTL_NO_EXPIRY = -1\n\nCACHE_TRANSIENT_ERRORS: tuple[type[Exception], ...] = (RedisError, SQLAlchemyError)\n\"\"\"Exception types that represent transient cache connectivity / operational\nfailures. Callers that want to fail-open (or fail-closed) on cache errors\nshould catch this tuple instead of bare ``Exception``.\n\nWhen adding a new ``CacheBackend`` implementation, add its transient error\nbase class(es) here so all call-sites pick it up automatically.\"\"\"\n\n\nclass CacheBackendType(str, Enum):\n REDIS = \"redis\"\n POSTGRES = \"postgres\"\n\n\nclass CacheLock(abc.ABC):\n \"\"\"Abstract distributed lock returned by CacheBackend.lock().\"\"\"\n\n @abc.abstractmethod\n def acquire(\n self,\n blocking: bool = True,\n blocking_timeout: float | None = None,\n ) -> bool:\n raise NotImplementedError\n\n @abc.abstractmethod\n def release(self) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def owned(self) -> bool:\n raise NotImplementedError\n\n def __enter__(self) -> \"CacheLock\":\n if not self.acquire():\n raise RuntimeError(\"Failed to acquire lock\")\n return self\n\n def __exit__(self, *args: object) -> None:\n self.release()\n\n\nclass CacheBackend(abc.ABC):\n \"\"\"Thin abstraction over a key-value cache with TTL, locks, and blocking lists.\n\n Covers the subset of Redis operations used outside of Celery. When\n CACHE_BACKEND=postgres, a PostgreSQL-backed implementation is used instead.\n \"\"\"\n\n # -- basic key/value ---------------------------------------------------\n\n @abc.abstractmethod\n def get(self, key: str) -> bytes | None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def set(\n self,\n key: str,\n value: str | bytes | int | float,\n ex: int | None = None,\n ) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def delete(self, key: str) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def exists(self, key: str) -> bool:\n raise NotImplementedError\n\n # -- TTL ---------------------------------------------------------------\n\n @abc.abstractmethod\n def expire(self, key: str, seconds: int) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def ttl(self, key: str) -> int:\n \"\"\"Return remaining TTL in seconds.\n\n Returns ``TTL_NO_EXPIRY`` (-1) if key exists without expiry,\n ``TTL_KEY_NOT_FOUND`` (-2) if key is missing or expired.\n \"\"\"\n raise NotImplementedError\n\n # -- distributed lock --------------------------------------------------\n\n @abc.abstractmethod\n def lock(self, name: str, timeout: float | None = None) -> CacheLock:\n raise NotImplementedError\n\n # -- blocking list (used by MCP OAuth BLPOP pattern) -------------------\n\n @abc.abstractmethod\n def rpush(self, key: str, value: str | bytes) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def blpop(self, keys: list[str], timeout: int = 0) -> tuple[bytes, bytes] | None:\n \"\"\"Block until a value is available on one of *keys*, or *timeout* expires.\n\n Returns ``(key, value)`` or ``None`` on timeout.\n \"\"\"\n raise NotImplementedError\n" }, { "path": "backend/onyx/cache/postgres_backend.py", "content": "\"\"\"PostgreSQL-backed ``CacheBackend`` for NO_VECTOR_DB deployments.\n\nUses the ``cache_store`` table for key-value storage, PostgreSQL advisory locks\nfor distributed locking, and a polling loop for the BLPOP pattern.\n\"\"\"\n\nimport hashlib\nimport struct\nimport time\nimport uuid\nfrom contextlib import AbstractContextManager\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom sqlalchemy import delete\nfrom sqlalchemy import func\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.orm import Session\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.cache.interface import CacheLock\nfrom onyx.cache.interface import TTL_KEY_NOT_FOUND\nfrom onyx.cache.interface import TTL_NO_EXPIRY\nfrom onyx.db.models import CacheStore\n\n_LIST_KEY_PREFIX = \"_q:\"\n# ASCII: ':' (0x3A) < ';' (0x3B). Upper bound for range queries so [prefix+, prefix;)\n# captures all list-item keys (e.g. _q:mylist:123:uuid) without including other\n# lists whose names share a prefix (e.g. _q:mylist2:...).\n_LIST_KEY_RANGE_TERMINATOR = \";\"\n_LIST_ITEM_TTL_SECONDS = 3600\n_LOCK_POLL_INTERVAL = 0.1\n_BLPOP_POLL_INTERVAL = 0.25\n\n\ndef _list_item_key(key: str) -> str:\n \"\"\"Unique key for a list item. Timestamp for FIFO ordering; UUID prevents\n collision when concurrent rpush calls occur within the same nanosecond.\n \"\"\"\n return f\"{_LIST_KEY_PREFIX}{key}:{time.time_ns()}:{uuid.uuid4().hex}\"\n\n\ndef _to_bytes(value: str | bytes | int | float) -> bytes:\n if isinstance(value, bytes):\n return value\n return str(value).encode()\n\n\n# ------------------------------------------------------------------\n# Lock\n# ------------------------------------------------------------------\n\n\nclass PostgresCacheLock(CacheLock):\n \"\"\"Advisory-lock-based distributed lock.\n\n Uses ``get_session_with_tenant`` for connection lifecycle. The lock is tied\n to the session's connection; releasing or closing the session frees it.\n\n NOTE: Unlike Redis locks, advisory locks do not auto-expire after\n ``timeout`` seconds. They are released when ``release()`` is\n called or when the session is closed.\n \"\"\"\n\n def __init__(self, lock_id: int, timeout: float | None, tenant_id: str) -> None:\n self._lock_id = lock_id\n self._timeout = timeout\n self._tenant_id = tenant_id\n self._session_cm: AbstractContextManager[Session] | None = None\n self._session: Session | None = None\n self._acquired = False\n\n def acquire(\n self,\n blocking: bool = True,\n blocking_timeout: float | None = None,\n ) -> bool:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n self._session_cm = get_session_with_tenant(tenant_id=self._tenant_id)\n self._session = self._session_cm.__enter__()\n try:\n if not blocking:\n return self._try_lock()\n\n effective_timeout = blocking_timeout or self._timeout\n deadline = (\n (time.monotonic() + effective_timeout) if effective_timeout else None\n )\n while True:\n if self._try_lock():\n return True\n if deadline is not None and time.monotonic() >= deadline:\n return False\n time.sleep(_LOCK_POLL_INTERVAL)\n finally:\n if not self._acquired:\n self._close_session()\n\n def release(self) -> None:\n if not self._acquired or self._session is None:\n return\n try:\n self._session.execute(select(func.pg_advisory_unlock(self._lock_id)))\n finally:\n self._acquired = False\n self._close_session()\n\n def owned(self) -> bool:\n return self._acquired\n\n def _close_session(self) -> None:\n if self._session_cm is not None:\n try:\n self._session_cm.__exit__(None, None, None)\n finally:\n self._session_cm = None\n self._session = None\n\n def _try_lock(self) -> bool:\n assert self._session is not None\n result = self._session.execute(\n select(func.pg_try_advisory_lock(self._lock_id))\n ).scalar()\n if result:\n self._acquired = True\n return True\n return False\n\n\n# ------------------------------------------------------------------\n# Backend\n# ------------------------------------------------------------------\n\n\nclass PostgresCacheBackend(CacheBackend):\n \"\"\"``CacheBackend`` backed by the ``cache_store`` table in PostgreSQL.\n\n Each operation opens and closes its own database session so the backend\n is safe to share across threads. Tenant isolation is handled by\n SQLAlchemy's ``schema_translate_map`` (set by ``get_session_with_tenant``).\n \"\"\"\n\n def __init__(self, tenant_id: str) -> None:\n self._tenant_id = tenant_id\n\n # -- basic key/value ---------------------------------------------------\n\n def get(self, key: str) -> bytes | None:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n stmt = select(CacheStore.value).where(\n CacheStore.key == key,\n or_(CacheStore.expires_at.is_(None), CacheStore.expires_at > func.now()),\n )\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n value = session.execute(stmt).scalar_one_or_none()\n if value is None:\n return None\n return bytes(value)\n\n def set(\n self,\n key: str,\n value: str | bytes | int | float,\n ex: int | None = None,\n ) -> None:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n value_bytes = _to_bytes(value)\n expires_at = (\n datetime.now(timezone.utc) + timedelta(seconds=ex)\n if ex is not None\n else None\n )\n stmt = (\n pg_insert(CacheStore)\n .values(key=key, value=value_bytes, expires_at=expires_at)\n .on_conflict_do_update(\n index_elements=[CacheStore.key],\n set_={\"value\": value_bytes, \"expires_at\": expires_at},\n )\n )\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n session.execute(stmt)\n session.commit()\n\n def delete(self, key: str) -> None:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n session.execute(delete(CacheStore).where(CacheStore.key == key))\n session.commit()\n\n def exists(self, key: str) -> bool:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n stmt = (\n select(CacheStore.key)\n .where(\n CacheStore.key == key,\n or_(\n CacheStore.expires_at.is_(None),\n CacheStore.expires_at > func.now(),\n ),\n )\n .limit(1)\n )\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n return session.execute(stmt).first() is not None\n\n # -- TTL ---------------------------------------------------------------\n\n def expire(self, key: str, seconds: int) -> None:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n new_exp = datetime.now(timezone.utc) + timedelta(seconds=seconds)\n stmt = (\n update(CacheStore).where(CacheStore.key == key).values(expires_at=new_exp)\n )\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n session.execute(stmt)\n session.commit()\n\n def ttl(self, key: str) -> int:\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n stmt = select(CacheStore.expires_at).where(CacheStore.key == key)\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n result = session.execute(stmt).first()\n if result is None:\n return TTL_KEY_NOT_FOUND\n expires_at: datetime | None = result[0]\n if expires_at is None:\n return TTL_NO_EXPIRY\n remaining = (expires_at - datetime.now(timezone.utc)).total_seconds()\n if remaining <= 0:\n return TTL_KEY_NOT_FOUND\n return int(remaining)\n\n # -- distributed lock --------------------------------------------------\n\n def lock(self, name: str, timeout: float | None = None) -> CacheLock:\n return PostgresCacheLock(\n self._lock_id_for(name), timeout, tenant_id=self._tenant_id\n )\n\n # -- blocking list (MCP OAuth BLPOP pattern) ---------------------------\n\n def rpush(self, key: str, value: str | bytes) -> None:\n self.set(_list_item_key(key), value, ex=_LIST_ITEM_TTL_SECONDS)\n\n def blpop(self, keys: list[str], timeout: int = 0) -> tuple[bytes, bytes] | None:\n if timeout <= 0:\n raise ValueError(\n \"PostgresCacheBackend.blpop requires timeout > 0. \"\n \"timeout=0 would block the calling thread indefinitely \"\n \"with no way to interrupt short of process termination.\"\n )\n from onyx.db.engine.sql_engine import get_session_with_tenant\n\n deadline = time.monotonic() + timeout\n while True:\n for key in keys:\n lower = f\"{_LIST_KEY_PREFIX}{key}:\"\n upper = f\"{_LIST_KEY_PREFIX}{key}{_LIST_KEY_RANGE_TERMINATOR}\"\n stmt = (\n select(CacheStore)\n .where(\n CacheStore.key >= lower,\n CacheStore.key < upper,\n or_(\n CacheStore.expires_at.is_(None),\n CacheStore.expires_at > func.now(),\n ),\n )\n .order_by(CacheStore.key)\n .limit(1)\n .with_for_update(skip_locked=True)\n )\n with get_session_with_tenant(tenant_id=self._tenant_id) as session:\n row = session.execute(stmt).scalars().first()\n if row is not None:\n value = bytes(row.value) if row.value else b\"\"\n session.delete(row)\n session.commit()\n return (key.encode(), value)\n if time.monotonic() >= deadline:\n return None\n time.sleep(_BLPOP_POLL_INTERVAL)\n\n # -- helpers -----------------------------------------------------------\n\n def _lock_id_for(self, name: str) -> int:\n \"\"\"Map *name* to a 64-bit signed int for ``pg_advisory_lock``.\"\"\"\n h = hashlib.md5(\n f\"{self._tenant_id}:{name}\".encode(), usedforsecurity=False\n ).digest()\n return struct.unpack(\"q\", h[:8])[0]\n\n\n# ------------------------------------------------------------------\n# Periodic cleanup\n# ------------------------------------------------------------------\n\n\ndef cleanup_expired_cache_entries() -> None:\n \"\"\"Delete rows whose ``expires_at`` is in the past.\n\n Called by the periodic poller every 5 minutes.\n \"\"\"\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n with get_session_with_current_tenant() as session:\n session.execute(\n delete(CacheStore).where(\n CacheStore.expires_at.is_not(None),\n CacheStore.expires_at < func.now(),\n )\n )\n session.commit()\n" }, { "path": "backend/onyx/cache/redis_backend.py", "content": "from typing import cast\n\nfrom redis.client import Redis\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.cache.interface import CacheLock\n\n\nclass RedisCacheLock(CacheLock):\n \"\"\"Wraps ``redis.lock.Lock`` behind the ``CacheLock`` interface.\"\"\"\n\n def __init__(self, lock: RedisLock) -> None:\n self._lock = lock\n\n def acquire(\n self,\n blocking: bool = True,\n blocking_timeout: float | None = None,\n ) -> bool:\n return bool(\n self._lock.acquire(\n blocking=blocking,\n blocking_timeout=blocking_timeout,\n )\n )\n\n def release(self) -> None:\n self._lock.release()\n\n def owned(self) -> bool:\n return bool(self._lock.owned())\n\n\nclass RedisCacheBackend(CacheBackend):\n \"\"\"``CacheBackend`` implementation that delegates to a ``redis.Redis`` client.\n\n This is a thin pass-through — every method maps 1-to-1 to the underlying\n Redis command. ``TenantRedis`` key-prefixing is handled by the client\n itself (provided by ``get_redis_client``).\n \"\"\"\n\n def __init__(self, redis_client: Redis) -> None:\n self._r = redis_client\n\n # -- basic key/value ---------------------------------------------------\n\n def get(self, key: str) -> bytes | None:\n val = self._r.get(key)\n if val is None:\n return None\n if isinstance(val, bytes):\n return val\n return str(val).encode()\n\n def set(\n self,\n key: str,\n value: str | bytes | int | float,\n ex: int | None = None,\n ) -> None:\n self._r.set(key, value, ex=ex)\n\n def delete(self, key: str) -> None:\n self._r.delete(key)\n\n def exists(self, key: str) -> bool:\n return bool(self._r.exists(key))\n\n # -- TTL ---------------------------------------------------------------\n\n def expire(self, key: str, seconds: int) -> None:\n self._r.expire(key, seconds)\n\n def ttl(self, key: str) -> int:\n return cast(int, self._r.ttl(key))\n\n # -- distributed lock --------------------------------------------------\n\n def lock(self, name: str, timeout: float | None = None) -> CacheLock:\n return RedisCacheLock(self._r.lock(name, timeout=timeout))\n\n # -- blocking list (MCP OAuth BLPOP pattern) ---------------------------\n\n def rpush(self, key: str, value: str | bytes) -> None:\n self._r.rpush(key, value)\n\n def blpop(self, keys: list[str], timeout: int = 0) -> tuple[bytes, bytes] | None:\n result = cast(list[bytes] | None, self._r.blpop(keys, timeout=timeout))\n if result is None:\n return None\n return (result[0], result[1])\n" }, { "path": "backend/onyx/chat/COMPRESSION.md", "content": "# Chat History Compression\n\nCompresses long chat histories by summarizing older messages while keeping recent ones verbatim.\n\n## Architecture Decisions\n\n### Branch-Aware via Tree Structure\nSummaries are stored as `ChatMessage` records with two key fields:\n- `parent_message_id` → last message when compression triggered (places summary in the tree)\n- `last_summarized_message_id` → pointer to an older message up the chain (the cutoff). Messages after this are kept verbatim.\n\n**Why store summary as a separate message?** If we embedded the summary in the `last_summarized_message_id` message itself, that message would contain context from messages that came after it—context that doesn't exist in other branches. By creating the summary as a new message attached to the branch tip, it only applies to the specific branch where compression occurred. It's only back-pointed to by the\nbranch which it applies to. All of this is necessary because we keep the last few messages verbatim and also to support branching logic.\n\n### Progressive Summarization\nSubsequent compressions incorporate the existing summary text + new messages, preventing information loss in very long conversations.\n\n### Cutoff Marker Prompt Strategy\nThe LLM receives older messages, a cutoff marker, then recent messages. It summarizes only content before the marker while using recent context to inform what's important.\n\n## Token Budget\n\nContext window breakdown:\n- `max_context_tokens` — LLM's total context window\n- `reserved_tokens` — space for system prompt, tools, files, etc.\n- Available for chat history = `max_context_tokens - reserved_tokens`\nNote: If there is a lot of reserved tokens, chat compression may happen fairly frequently which is costly, slow, and leads to a bad user experience. Possible area of future improvement.\n\nConfigurable ratios:\n- `COMPRESSION_TRIGGER_RATIO` (default 0.75) — compress when chat history exceeds this ratio of available space\n- `RECENT_MESSAGES_RATIO` (default 0.2) — portion of chat history to keep verbatim when compressing\n\n## Flow\n\n1. Trigger when `history_tokens > available * 0.75`\n2. Find existing summary for branch (if any)\n3. Split messages: older (summarize) / recent (keep 25%)\n4. Generate summary via LLM\n5. Save as `ChatMessage` with `parent_message_id` + `last_summarized_message_id`\n\n## Key Functions\n\n| Function | Purpose |\n|----------|---------|\n| `get_compression_params` | Check if compression needed based on token counts |\n| `find_summary_for_branch` | Find applicable summary by checking `parent_message_id` membership |\n| `get_messages_to_summarize` | Split messages at token budget boundary |\n| `compress_chat_history` | Orchestrate flow, save summary message |\n" }, { "path": "backend/onyx/chat/README.md", "content": "# Overview of Context Management\n\nThis document reviews some design decisions around the main agent-loop powering Onyx's chat flow.\nIt is highly recommended for all engineers contributing to this flow to be familiar with the concepts here.\n\n> Note: it is assumed the reader is familiar with the Onyx product and features such as Projects, User files, Citations, etc. \n\n## System Prompt\n\nThe system prompt is a default prompt that comes packaged with the system. Users can edit the default prompt and it will be persisted in the database.\n\nSome parts of the system prompt are dynamically updated / inserted:\n\n- Datetime of the message sent\n- Tools description of when to use certain tools depending on if the tool is available in that cycle\n- If the user has just called a search related tool, then a section about citations is included\n\n## Custom Agent Prompt\n\nThe custom agent is inserted as a user message above the most recent user message, it is dynamically moved in the history as the user sends more messages.\nIf the user has opted to completely replace the System Prompt, then this Custom Agent prompt replaces the system prompt and does not move along the history.\n\n## How Files are handled\n\nOn upload, Files are processed for tokens, if too many tokens to fit in the context, it’s considered a failed inclusion. This is done using the LLM tokenizer.\n\n- In many cases, there is not a known tokenizer for each LLM so there is a default tokenizer used as a catchall.\n- File upload happens in 2 parts - the actual upload + token counting.\n- Files are added into chat context as a “point in time” inclusion and move up the context window as the conversation progresses.\n Every file knows how many tokens it is (model agnostic), image files have some assumed number of tokens.\n\nImage files are attached to User Messages also as point in time inclusions.\n\n**Future Extension**:\nFiles selected from the search results are also counted as “point in time” inclusions. Files that are too large cannot be selected.\nFor these files, the \"entire file\" does not exist for most connectors, it's pieced back together from the search engine.\n\n## Projects\n\nIf a Project contains few enough files that it all fits in the model context, we keep it close enough in the history to ensure it is easy for the LLM to\naccess. Note that the project documents are assumed to be quite useful and that they should 1. never be dropped from context, 2. is not just a needle in\na haystack type search with a strong keyword to make the LLM attend to it.\n\nProject files are vectorized and stored in the Search Engine so that if the user chooses a model with less context than the number of tokens in the project,\nthe system can RAG over the project files.\n\n## How documents are represented\n\nDocuments from search or uploaded Project files are represented as a json so that the LLM can easily understand it. It is represented with a prefix string to\nmake the context clearer to the LLM. Note that for search results (whether web or internal, it will just be the json) and it will be a Tool Call type of\nmessage rather than a user message.\n\n```\nHere are some documents provided for context, they may not all be relevant:\n{\n \"documents\": [\n {\"document\": 1, \"title\": \"Hello\", \"metadata\": \"status closed\", \"contents\": \"Foo\"},\n {\"document\": 2, \"title\": \"World\", \"contents\": \"Bar\"}\n ]\n}\n```\n\nDocuments are represented with the `document` key so that the LLM can easily cite them with a single number. The tool returns have to be richer to be able to\ntranslate this into links and other UI elements. What the LLM sees is far simpler to reduce noise/hallucinations.\n\nNote that documents included in a single turn should be collapsed into a single user message.\n\nSearch tools also give URLs to the LLM so that open_url (a separate tool) can be called on them.\n\n## Reminders\n\nTo ensure the LLM follows certain specific instructions, instructions are added at the very end of the chat context as a user message. If a search related\ntool is used, a citation reminder is always added. Otherwise, by default there is no reminder. If the user configures reminders, those are added to the\nfinal message. If a search related tool just ran and the user has reminders, both appear in a single message.\n\nIf a search related tool is called at any point during the turn, the reminder will remain at the end until the turn is over and the agent has responded.\n\n## Tool Calls\n\nAs tool call responses can get very long (like an internal search can be many thousands of tokens), tool responses are current replaced with a hardcoded\nstring saying it is no longer available. Tool Call details like the search query and other arguments are kept in the history as this is information\nrich and generally very few tokens.\n\n> Note: in the Internal Search flow with query expansion, the Tool Call which was actually run differs from what the LLM provided as arguments.\n> What the LLM sees in the history (to be most informative for future calls) is the full set of expanded queries.\n\n**Possible Future Extension**:\nInstead of dropping the Tool Call response, we might summarize it using an LLM so that it is just 1-2 sentences and captures the main points. That said,\nthis is questionable value add because anything relevant and useful should be already captured in the Agent response.\n\n## Examples\n\n```\nS -> System Message\nCA -> Custom Agent as a User Message\nA -> Agent Message response to user\nU -> User Message\nTC -> Agent Message for a tool call\nTR -> Tool response\nR -> Reminder\nF -> Point in time File\nP -> Project Files (not overflowed case)\n1,2,3 etc. to represent turn number. A turn consists of a user input and a final response from the Agent\n\nFlow with Custom Agent\nS, U1, TC, TR, A1, CA, U2, A2 -- user sends another message, triggers tool call -> S, U1, TC, TR, A1, U2, A2, CA, U3, TC, TR, R, A3\n- Custom agent response moves\n- Reminder inserted after TR\n\nFlow with Project and File Upload\nS, CA, P, F, U1, A1 -- user sends another message -> S, F, U1, A1, CA, P, U2, A2\n- File stays in place, above the user message\n- Project files move along the chain as new messages are sent\n- Custom Agent prompt comes before project files which come before user uploaded files in each turn\n\nReminders during a single Turn\nS, U1, TC, TR, R -- agent calls another tool -> S, U1, TC, TR, TC, TR, R, A1\n- Reminder moved to the end\n```\n\n## Product considerations\n\nProject files are important to the entire duration of the chat session. If the user has uploaded project files, they are likely very intent on working with\nthose files. The LLM is much better at referencing documents close to the end of the context window so keeping it there for ease of access.\n\nUser uploaded files are considered relevant for that point in time, it is ok if the Agent forgets about it as the chat gets long. If every uploaded file is\nconstantly moved towards the end of the chat, it would degrade quality as these stack up. Even with a single file, there is some cost of making the previous\nUser Message further away. This tradeoff is accepted for Projects because of the intent of the feature.\n\nReminder are absolutely necessary to ensure 1-2 specific instructions get followed with a very high probability. It is less detailed than the system prompt\nand should be very targetted for it to work reliably and also not interfere with the last user message.\n\n## Reasons / Experiments\n\nCustom Agent instructions being placed in the system prompt is poorly followed. It also degrades performance of the system especially when the instructions\nare orthogonal (or even possibly contradictory) to the system prompt. For weaker models, it causes strange artifacts in tool calls and final responses\nthat completely ruins the user experience. Empirically, this way works better across a range of models especially when the history gets longer.\nHaving the Custom Agent instructions not move means it fades more as the chat gets long which is also not ok from a UX perspective.\n\nDifferent LLMs vary in this but some now have a section that cannot be set via the API layer called the \"System Prompt\" (OpenAI terminology) which contains\ninformation like the model cutoff date, identity, and some other basic non-changing information. The System prompt described above is in that convention called\nthe \"Developer Prompt\". It seems the distribution of the System Prompt, by which I mean the style of wording and terms used can also affect the behavior. This\nis different between different models and not necessarily scientific so the system prompt is built from an exploration across different models. It currently\nstarts with: \"You are a highly capable, thoughtful, and precise assistant. Your goal is to deeply understand the user's intent...\"\n\nLLMs are able to handle changes in topic best at message boundaries. There are special tokens under the hood for this. We also use this property to slice up\nthe history in the way presented above.\n\nReminder messages are placed at the end of the prompt because all model fine tuning approaches cause the LLMs to attend very strongly to the tokens at the very\nback of the context closest to generation. This is the only way to get the LLMs to not miss critical information and for the product to be reliable. Specifically\nthe built-in reminders are around citations and what tools it should call in certain situations.\n\nThe document json includes a field for the LLM to cite (it's a single number) to make citations reliable and avoid weird artifacts. It's called \"document\" so\nthat the LLM does not create weird artifacts in reasoning like \"I should reference citation_id: 5 for...\". It is also strategically placed so that it is easy to\nreference. It is followed by a couple short sections like the metadata and title before the long content section. It seems LLMs are still better at local\nattention despite having global access.\n\nIn a similar concept, LLM instructions in the system prompt are structured specifically so that there are coherent sections for the LLM to attend to. This is\nfairly surprising actually but if there is a line of instructions effectively saying \"If you try to use some tools and find that you need more information or\nneed to call additional tools, you are encouraged to do this\", having this in the Tool section of the System prompt makes all the LLMs follow it well but if it's\neven just a paragraph away like near the beginning of the prompt, it is often ignored. The difference is as drastic as a 30% follow rate to a 90% follow\nrate by even just moving the same statement a few sentences.\n\n## Other related pointers\n\n- How messages, files, images are stored can be found in backend/onyx/db/models.py, there is also a README.md under that directory that may be helpful.\n\n---\n\n# Overview of LLM flow architecture\n\n**Concepts:**\nTurn: User sends a message and AI does some set of things and responds\nStep/Cycle: 1 single LLM inference given some context and some tools\n\n## 1. Top Level (process_message function):\n\nThis function can be thought of as the set-up and validation layer. It ensures that the database is in a valid state, reads the\nmessages in the session and sets up all the necessary items to run the chat loop and state containers. The major things it does\nare:\n\n- Validates the request\n- Builds the chat history for the session\n- Fetches any additional context such as files and images\n- Prepares all of the tools for the LLM\n- Creates the state container objects for use in the loop\n\n### Execution (`_run_models` function):\n\nEach model runs in its own worker thread inside a `ThreadPoolExecutor`. Workers write packets to a shared\n`merged_queue` via an `Emitter`; the main thread drains the queue and yields packets in arrival order. This\nmeans the top level is isolated from the LLM flow and can yield packets as soon as they are produced. If a\nworker fails, the main thread yields a `StreamingError` for that model and keeps the other models running.\nAll saving and database operations are handled by the main thread after the workers complete (or by the\nworkers themselves via self-completion if the drain loop exits early).\n\n### Emitter\n\nThe emitter is an object that lower levels use to send packets without needing to yield them all the way back\nup the call stack. Each `Emitter` tags every packet with a `model_index` and places it on the shared\n`merged_queue` as a `(model_idx, packet)` tuple. The drain loop in `_run_models` consumes these tuples and\nyields the packets to the caller. Both the emitter and the state container are mutating state objects used\nonly to accumulate state. There should be no logic dependent on the states of these objects, especially in\nthe lower levels. The emitter should only take packets and should not be used for other things.\n\n### State Container\n\nThe state container is used to accumulate state during the LLM flow. Similar to the emitter, it should not be used for logic,\nonly for accumulating state. It is used to gather all of the necessary information for saving the chat turn into the database.\nSo it will accumulate answer tokens, reasoning tokens, tool calls, citation info, etc. This is used at the end of the flow once\nthe lower level is completed whether on its own or stopped by the user. At that point, all of the state is read and stored into\nthe database. The state container can be added to by any of the underlying layers, this is fine.\n\n### Stopping Generation\n\nThe drain loop in `_run_models` checks `check_is_connected()` every 50 ms (on queue timeout). The signal itself\nis stored in Redis and is set by the user calling the stop endpoint. On disconnect, the drain loop saves\npartial state for every model, yields an `OverallStop(stop_reason=\"user_cancelled\")` packet, and returns.\nA `drain_done` event signals emitters to stop blocking so worker threads can exit quickly. Workers that\nalready completed successfully will self-complete (persist their response) if the drain loop exited before\nreaching the normal completion path.\n\n## 2. LLM Loop (run_llm_loop function)\n\nThis function handles the logic of the Turn. It's essentially a while loop where context is added and modified (according what\nis outlined in the first half of this doc). Its main functionality is:\n\n- Translate and truncate the context for the LLM inference\n- Add context modifiers like reminders, updates to the system prompts, etc.\n- Run tool calls and gather results\n- Build some of the objects stored in the state container.\n\n## 3. LLM Step (run_llm_step function)\n\nThis function is a single inference of the LLM. It's a wrapper around the LLM stream function which handles packet translations\nso that the Emitter can emit individual tokens as soon as they arrive. It also keeps track of the different sections since they\ndo not all come at once (reasoning, answers, tool calls are all built up token by token). This layer also tracks the different\ntool calls and returns that to the LLM Loop to execute.\n\n## Things to know\n\n- Packets are labeled with a \"turn_index\" field as part of the Placement of the packet. This is not the same as the backend\n concept of a turn. The turn_index for the frontend is which block does this packet belong to. So while a reasoning + tool call\n comes from the same LLM inference (same backend LLM step), they are 2 turns to the frontend because that's how it's rendered.\n\n- There are 3 representations of a message, each scoped to a different layer:\n 1. **ChatMessage** — The database model. Should be converted into ChatMessageSimple early and never passed deep into the flow.\n 2. **ChatMessageSimple** — The canonical data model used throughout the codebase. This is the rich, full-featured representation\n of a message. Any modifications or additions to message structure should be made here.\n 3. **LanguageModelInput** — The LLM-facing representation. Intentionally minimal so the LLM interface layer stays clean and\n easy to maintain/extend.\n" }, { "path": "backend/onyx/chat/__init__.py", "content": "" }, { "path": "backend/onyx/chat/chat_processing_checker.py", "content": "from uuid import UUID\n\nfrom onyx.cache.interface import CacheBackend\n\nPREFIX = \"chatprocessing\"\nFENCE_PREFIX = f\"{PREFIX}_fence\"\nFENCE_TTL = 30 * 60 # 30 minutes\n\n\ndef _get_fence_key(chat_session_id: UUID) -> str:\n \"\"\"Generate the cache key for a chat session processing fence.\n\n Args:\n chat_session_id: The UUID of the chat session\n\n Returns:\n The fence key string. Tenant isolation is handled automatically\n by the cache backend (Redis key-prefixing or Postgres schema routing).\n \"\"\"\n return f\"{FENCE_PREFIX}_{chat_session_id}\"\n\n\ndef set_processing_status(\n chat_session_id: UUID, cache: CacheBackend, value: bool\n) -> None:\n \"\"\"Set or clear the fence for a chat session processing a message.\n\n If the key exists, a message is being processed.\n\n Args:\n chat_session_id: The UUID of the chat session\n cache: Tenant-aware cache backend\n value: True to set the fence, False to clear it\n \"\"\"\n fence_key = _get_fence_key(chat_session_id)\n if value:\n cache.set(fence_key, 0, ex=FENCE_TTL)\n else:\n cache.delete(fence_key)\n\n\ndef is_chat_session_processing(chat_session_id: UUID, cache: CacheBackend) -> bool:\n \"\"\"Check if the chat session is processing a message.\n\n Args:\n chat_session_id: The UUID of the chat session\n cache: Tenant-aware cache backend\n\n Returns:\n True if the chat session is processing a message, False otherwise\n \"\"\"\n return cache.exists(_get_fence_key(chat_session_id))\n" }, { "path": "backend/onyx/chat/chat_state.py", "content": "import threading\nfrom collections.abc import Callable\nfrom dataclasses import dataclass\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.chat.citation_processor import CitationMapping\nfrom onyx.chat.models import ChatLoadedFile\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.chat.models import ExtractedContextFiles\nfrom onyx.chat.models import FileToolMetadata\nfrom onyx.chat.models import SearchParams\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.memory import UserMemoryContext\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import Persona\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.onyxbot.slack.models import SlackContext\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.tools.models import ChatFile\nfrom onyx.tools.models import ToolCallInfo\n\n# Type alias for search doc deduplication key\n# Simple key: just document_id (str)\n# Full key: (document_id, chunk_ind, match_highlights)\nSearchDocKey = str | tuple[str, int, tuple[str, ...]]\n\n\nclass ChatStateContainer:\n \"\"\"Container for accumulating state during LLM loop execution.\n\n This container holds the partial state that can be saved to the database\n if the generation is stopped by the user or completes normally.\n\n Thread-safe: All write operations are protected by a lock to ensure safe\n concurrent access from multiple threads. For thread-safe reads, use the\n getter methods. Direct attribute access is not thread-safe.\n \"\"\"\n\n def __init__(self) -> None:\n self._lock = threading.Lock()\n # These are collected at the end after the entire tool call is completed\n self.tool_calls: list[ToolCallInfo] = []\n # This is accumulated during the streaming\n self.reasoning_tokens: str | None = None\n # This is accumulated during the streaming of the answer\n self.answer_tokens: str | None = None\n # Store citation mapping for building citation_docs_info during partial saves\n self.citation_to_doc: CitationMapping = {}\n # True if this turn is a clarification question (deep research flow)\n self.is_clarification: bool = False\n # Pre-answer processing time (time before answer starts) in seconds\n self.pre_answer_processing_time: float | None = None\n # Note: LLM cost tracking is now handled in multi_llm.py\n # Search doc collection - maps dedup key to SearchDoc for all docs from tool calls\n self._all_search_docs: dict[SearchDocKey, SearchDoc] = {}\n # Track which citation numbers were actually emitted during streaming\n self._emitted_citations: set[int] = set()\n\n def add_tool_call(self, tool_call: ToolCallInfo) -> None:\n \"\"\"Add a tool call to the accumulated state.\"\"\"\n with self._lock:\n self.tool_calls.append(tool_call)\n\n def set_reasoning_tokens(self, reasoning: str | None) -> None:\n \"\"\"Set the reasoning tokens from the final answer generation.\"\"\"\n with self._lock:\n self.reasoning_tokens = reasoning\n\n def set_answer_tokens(self, answer: str | None) -> None:\n \"\"\"Set the answer tokens from the final answer generation.\"\"\"\n with self._lock:\n self.answer_tokens = answer\n\n def set_citation_mapping(self, citation_to_doc: CitationMapping) -> None:\n \"\"\"Set the citation mapping from citation processor.\"\"\"\n with self._lock:\n self.citation_to_doc = citation_to_doc\n\n def set_is_clarification(self, is_clarification: bool) -> None:\n \"\"\"Set whether this turn is a clarification question.\"\"\"\n with self._lock:\n self.is_clarification = is_clarification\n\n def get_answer_tokens(self) -> str | None:\n \"\"\"Thread-safe getter for answer_tokens.\"\"\"\n with self._lock:\n return self.answer_tokens\n\n def get_reasoning_tokens(self) -> str | None:\n \"\"\"Thread-safe getter for reasoning_tokens.\"\"\"\n with self._lock:\n return self.reasoning_tokens\n\n def get_tool_calls(self) -> list[ToolCallInfo]:\n \"\"\"Thread-safe getter for tool_calls (returns a copy).\"\"\"\n with self._lock:\n return self.tool_calls.copy()\n\n def get_citation_to_doc(self) -> CitationMapping:\n \"\"\"Thread-safe getter for citation_to_doc (returns a copy).\"\"\"\n with self._lock:\n return self.citation_to_doc.copy()\n\n def get_is_clarification(self) -> bool:\n \"\"\"Thread-safe getter for is_clarification.\"\"\"\n with self._lock:\n return self.is_clarification\n\n def set_pre_answer_processing_time(self, duration: float | None) -> None:\n \"\"\"Set the pre-answer processing time (time before answer starts).\"\"\"\n with self._lock:\n self.pre_answer_processing_time = duration\n\n def get_pre_answer_processing_time(self) -> float | None:\n \"\"\"Thread-safe getter for pre_answer_processing_time.\"\"\"\n with self._lock:\n return self.pre_answer_processing_time\n\n @staticmethod\n def create_search_doc_key(\n search_doc: SearchDoc, use_simple_key: bool = True\n ) -> SearchDocKey:\n \"\"\"Create a unique key for a SearchDoc for deduplication.\n\n Args:\n search_doc: The SearchDoc to create a key for\n use_simple_key: If True (default), use only document_id for deduplication.\n If False, include chunk_ind and match_highlights so that the same\n document/chunk with different highlights are stored separately.\n \"\"\"\n if use_simple_key:\n return search_doc.document_id\n match_highlights_tuple = tuple(sorted(search_doc.match_highlights or []))\n return (search_doc.document_id, search_doc.chunk_ind, match_highlights_tuple)\n\n def add_search_docs(\n self, search_docs: list[SearchDoc], use_simple_key: bool = True\n ) -> None:\n \"\"\"Add search docs to the accumulated collection with deduplication.\n\n Args:\n search_docs: List of SearchDoc objects to add\n use_simple_key: If True (default), deduplicate by document_id only.\n If False, deduplicate by document_id + chunk_ind + match_highlights.\n \"\"\"\n with self._lock:\n for doc in search_docs:\n key = self.create_search_doc_key(doc, use_simple_key)\n if key not in self._all_search_docs:\n self._all_search_docs[key] = doc\n\n def get_all_search_docs(self) -> dict[SearchDocKey, SearchDoc]:\n \"\"\"Thread-safe getter for all accumulated search docs (returns a copy).\"\"\"\n with self._lock:\n return self._all_search_docs.copy()\n\n def add_emitted_citation(self, citation_num: int) -> None:\n \"\"\"Add a citation number that was actually emitted during streaming.\"\"\"\n with self._lock:\n self._emitted_citations.add(citation_num)\n\n def get_emitted_citations(self) -> set[int]:\n \"\"\"Thread-safe getter for emitted citations (returns a copy).\"\"\"\n with self._lock:\n return self._emitted_citations.copy()\n\n\nclass AvailableFiles(BaseModel):\n \"\"\"Separated file IDs for the FileReaderTool so it knows which loader to use.\"\"\"\n\n # IDs from the ``user_file`` table (project / persona-attached files).\n user_file_ids: list[UUID] = []\n # IDs from the ``file_record`` table (chat-attached files).\n chat_file_ids: list[UUID] = []\n\n\n@dataclass(frozen=True)\nclass ChatTurnSetup:\n \"\"\"Immutable context produced by ``build_chat_turn`` and consumed by ``_run_models``.\"\"\"\n\n new_msg_req: SendMessageRequest\n chat_session: ChatSession\n persona: Persona\n user_message: ChatMessage\n user_identity: LLMUserIdentity\n llms: list[LLM] # length 1 for single-model, N for multi-model\n model_display_names: list[str] # parallel to llms\n simple_chat_history: list[ChatMessageSimple]\n extracted_context_files: ExtractedContextFiles\n reserved_messages: list[ChatMessage] # length 1 for single, N for multi\n reserved_token_count: int\n search_params: SearchParams\n all_injected_file_metadata: dict[str, FileToolMetadata]\n available_files: AvailableFiles\n tool_id_to_name_map: dict[int, str]\n forced_tool_id: int | None\n files: list[ChatLoadedFile]\n chat_files_for_tools: list[ChatFile]\n custom_agent_prompt: str | None\n user_memory_context: UserMemoryContext\n # For deep research: was the last assistant message a clarification request?\n skip_clarification: bool\n check_is_connected: Callable[[], bool]\n cache: CacheBackend\n # Execution params forwarded to per-model tool construction\n bypass_acl: bool\n slack_context: SlackContext | None\n custom_tool_additional_headers: dict[str, str] | None\n mcp_headers: dict[str, str] | None\n" }, { "path": "backend/onyx/chat/chat_utils.py", "content": "import json\nimport re\nfrom collections.abc import Callable\nfrom typing import cast\nfrom uuid import UUID\n\nfrom fastapi.datastructures import Headers\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.models import ChatHistoryResult\nfrom onyx.chat.models import ChatLoadedFile\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.chat.models import FileToolMetadata\nfrom onyx.chat.models import ToolCallSimple\nfrom onyx.configs.constants import DEFAULT_PERSONA_ID\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import TMP_DRALPHA_PERSONA_NAME\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.chat import get_chat_messages_by_session\nfrom onyx.db.chat import get_or_create_root_message\nfrom onyx.db.kg_config import get_kg_config_settings\nfrom onyx.db.kg_config import is_kg_config_settings_enabled_valid\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import Persona\nfrom onyx.db.models import SearchDoc as DbSearchDoc\nfrom onyx.db.models import UserFile\nfrom onyx.db.projects import check_project_ownership\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.file_store.models import ChatFileType\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.file_store.utils import plaintext_file_name_for_id\nfrom onyx.file_store.utils import store_plaintext\nfrom onyx.kg.models import KGException\nfrom onyx.kg.setup.kg_default_entity_definitions import (\n populate_missing_default_entity_types__commit,\n)\nfrom onyx.prompts.chat_prompts import ADDITIONAL_CONTEXT_PROMPT\nfrom onyx.prompts.chat_prompts import TOOL_CALL_RESPONSE_CROSS_MESSAGE\nfrom onyx.prompts.tool_prompts import TOOL_CALL_FAILURE_PROMPT\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.tools.models import ToolCallKickoff\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom onyx.utils.timing import log_function_time\n\n\nlogger = setup_logger()\nIMAGE_GENERATION_TOOL_NAME = \"generate_image\"\n\n\nclass FileContextResult(BaseModel):\n \"\"\"Result of building a file's LLM context representation.\"\"\"\n\n message: ChatMessageSimple\n tool_metadata: FileToolMetadata\n\n\ndef build_file_context(\n tool_file_id: str,\n filename: str,\n file_type: ChatFileType,\n content_text: str | None = None,\n token_count: int = 0,\n approx_char_count: int | None = None,\n) -> FileContextResult:\n \"\"\"Build the LLM context representation for a single file.\n\n Centralises how files should appear in the LLM prompt\n — the ID that FileReaderTool accepts (``UserFile.id`` for user files).\n \"\"\"\n if file_type.use_metadata_only():\n message_text = (\n f\"File: {filename} (id={tool_file_id})\\n\"\n \"Use the file_reader or python tools to access \"\n \"this file's contents.\"\n )\n message = ChatMessageSimple(\n message=message_text,\n token_count=max(1, len(message_text) // 4),\n message_type=MessageType.USER,\n file_id=tool_file_id,\n )\n else:\n message_text = f\"File: {filename}\\n{content_text or ''}\\nEnd of File\"\n message = ChatMessageSimple(\n message=message_text,\n token_count=token_count,\n message_type=MessageType.USER,\n file_id=tool_file_id,\n )\n\n metadata = FileToolMetadata(\n file_id=tool_file_id,\n filename=filename,\n approx_char_count=(\n approx_char_count\n if approx_char_count is not None\n else len(content_text or \"\")\n ),\n )\n\n return FileContextResult(message=message, tool_metadata=metadata)\n\n\ndef create_chat_session_from_request(\n chat_session_request: ChatSessionCreationRequest,\n user_id: UUID | None,\n db_session: Session,\n) -> ChatSession:\n \"\"\"Create a chat session from a ChatSessionCreationRequest.\n\n Includes project ownership validation when project_id is provided.\n\n Args:\n chat_session_request: The request containing persona_id, description, and project_id\n user_id: The ID of the user creating the session (can be None for anonymous)\n db_session: The database session\n\n Returns:\n The newly created ChatSession\n\n Raises:\n ValueError: If user lacks access to the specified project\n Exception: If the persona is invalid\n \"\"\"\n project_id = chat_session_request.project_id\n if project_id:\n if not check_project_ownership(project_id, user_id, db_session):\n raise ValueError(\"User does not have access to project\")\n\n return create_chat_session(\n db_session=db_session,\n description=chat_session_request.description or \"\",\n user_id=user_id,\n persona_id=chat_session_request.persona_id,\n project_id=chat_session_request.project_id,\n )\n\n\ndef create_chat_history_chain(\n chat_session_id: UUID,\n db_session: Session,\n prefetch_top_two_level_tool_calls: bool = True,\n # Optional id at which we finish processing\n stop_at_message_id: int | None = None,\n) -> list[ChatMessage]:\n \"\"\"Build the linear chain of messages without including the root message\"\"\"\n mainline_messages: list[ChatMessage] = []\n\n all_chat_messages = get_chat_messages_by_session(\n chat_session_id=chat_session_id,\n user_id=None,\n db_session=db_session,\n skip_permission_check=True,\n prefetch_top_two_level_tool_calls=prefetch_top_two_level_tool_calls,\n )\n\n if not all_chat_messages:\n root_message = get_or_create_root_message(\n chat_session_id=chat_session_id, db_session=db_session\n )\n else:\n root_message = all_chat_messages[0]\n if root_message.parent_message is not None:\n raise RuntimeError(\n \"Invalid root message, unable to fetch valid chat message sequence\"\n )\n\n current_message: ChatMessage | None = root_message\n previous_message: ChatMessage | None = None\n while current_message is not None:\n child_msg = current_message.latest_child_message\n\n # Break if at the end of the chain\n # or have reached the `final_id` of the submitted message\n if not child_msg or (\n stop_at_message_id and current_message.id == stop_at_message_id\n ):\n break\n current_message = child_msg\n\n if (\n current_message.message_type == MessageType.ASSISTANT\n and previous_message is not None\n and previous_message.message_type == MessageType.ASSISTANT\n and mainline_messages\n ):\n # Note that 2 user messages in a row is fine since this is often used for\n # adding custom prompts and reminders\n raise RuntimeError(\n \"Invalid message chain, cannot have two assistant messages in a row\"\n )\n else:\n mainline_messages.append(current_message)\n\n previous_message = current_message\n\n return mainline_messages\n\n\ndef reorganize_citations(\n answer: str, citations: list[CitationInfo]\n) -> tuple[str, list[CitationInfo]]:\n \"\"\"For a complete, citation-aware response, we want to reorganize the citations so that\n they are in the order of the documents that were used in the response. This just looks nicer / avoids\n confusion (\"Why is there [7] when only 2 documents are cited?\").\"\"\"\n\n # Regular expression to find all instances of [[x]](LINK)\n pattern = r\"\\[\\[(.*?)\\]\\]\\((.*?)\\)\"\n\n all_citation_matches = re.findall(pattern, answer)\n\n new_citation_info: dict[int, CitationInfo] = {}\n for citation_match in all_citation_matches:\n try:\n citation_num = int(citation_match[0])\n if citation_num in new_citation_info:\n continue\n\n matching_citation = next(\n iter([c for c in citations if c.citation_number == int(citation_num)]),\n None,\n )\n if matching_citation is None:\n continue\n\n new_citation_info[citation_num] = CitationInfo(\n citation_number=len(new_citation_info) + 1,\n document_id=matching_citation.document_id,\n )\n except Exception:\n pass\n\n # Function to replace citations with their new number\n def slack_link_format(match: re.Match) -> str:\n link_text = match.group(1)\n try:\n citation_num = int(link_text)\n if citation_num in new_citation_info:\n link_text = new_citation_info[citation_num].citation_number\n except Exception:\n pass\n\n link_url = match.group(2)\n return f\"[[{link_text}]]({link_url})\"\n\n # Substitute all matches in the input text\n new_answer = re.sub(pattern, slack_link_format, answer)\n\n # if any citations weren't parsable, just add them back to be safe\n for citation in citations:\n if citation.citation_number not in new_citation_info:\n new_citation_info[citation.citation_number] = citation\n\n return new_answer, list(new_citation_info.values())\n\n\ndef build_citation_map_from_infos(\n citations_list: list[CitationInfo], db_docs: list[DbSearchDoc]\n) -> dict[int, int]:\n \"\"\"Translate a list of streaming CitationInfo objects into a mapping of\n citation number -> saved search doc DB id.\n\n Always cites the first instance of a document_id and assumes db_docs are\n ordered as shown to the user (display order).\n \"\"\"\n doc_id_to_saved_doc_id_map: dict[str, int] = {}\n for db_doc in db_docs:\n if db_doc.document_id not in doc_id_to_saved_doc_id_map:\n doc_id_to_saved_doc_id_map[db_doc.document_id] = db_doc.id\n\n citation_to_saved_doc_id_map: dict[int, int] = {}\n for citation in citations_list:\n if citation.citation_number not in citation_to_saved_doc_id_map:\n saved_id = doc_id_to_saved_doc_id_map.get(citation.document_id)\n if saved_id is not None:\n citation_to_saved_doc_id_map[citation.citation_number] = saved_id\n\n return citation_to_saved_doc_id_map\n\n\ndef build_citation_map_from_numbers(\n cited_numbers: list[int] | set[int], db_docs: list[DbSearchDoc]\n) -> dict[int, int]:\n \"\"\"Translate parsed citation numbers (e.g., from [[n]]) into a mapping of\n citation number -> saved search doc DB id by positional index.\n \"\"\"\n citation_to_saved_doc_id_map: dict[int, int] = {}\n for num in sorted(set(cited_numbers)):\n idx = num - 1\n if 0 <= idx < len(db_docs):\n citation_to_saved_doc_id_map[num] = db_docs[idx].id\n\n return citation_to_saved_doc_id_map\n\n\ndef extract_headers(\n headers: dict[str, str] | Headers, pass_through_headers: list[str] | None\n) -> dict[str, str]:\n \"\"\"\n Extract headers specified in pass_through_headers from input headers.\n Handles both dict and FastAPI Headers objects, accounting for lowercase keys.\n\n Args:\n headers: Input headers as dict or Headers object.\n\n Returns:\n dict: Filtered headers based on pass_through_headers.\n \"\"\"\n if not pass_through_headers:\n return {}\n\n extracted_headers: dict[str, str] = {}\n for key in pass_through_headers:\n if key in headers:\n extracted_headers[key] = headers[key]\n else:\n # fastapi makes all header keys lowercase, handling that here\n lowercase_key = key.lower()\n if lowercase_key in headers:\n extracted_headers[lowercase_key] = headers[lowercase_key]\n return extracted_headers\n\n\ndef process_kg_commands(\n message: str,\n persona_name: str,\n tenant_id: str, # noqa: ARG001\n db_session: Session,\n) -> None:\n # Temporarily, until we have a draft UI for the KG Operations/Management\n # TODO: move to api endpoint once we get frontend\n if not persona_name.startswith(TMP_DRALPHA_PERSONA_NAME):\n return\n\n kg_config_settings = get_kg_config_settings()\n if not is_kg_config_settings_enabled_valid(kg_config_settings):\n return\n\n if message == \"kg_setup\":\n populate_missing_default_entity_types__commit(db_session=db_session)\n raise KGException(\"KG setup done\")\n\n\ndef _get_or_extract_plaintext(\n file_id: str,\n extract_fn: Callable[[], str],\n) -> str:\n \"\"\"Load cached plaintext for a file, or extract and store it.\n\n Tries to read pre-stored plaintext from the file store. On a miss,\n calls extract_fn to produce the text, then stores the result so\n future calls skip the expensive extraction.\n \"\"\"\n file_store = get_default_file_store()\n plaintext_key = plaintext_file_name_for_id(file_id)\n\n # Try cached plaintext first.\n try:\n plaintext_io = file_store.read_file(plaintext_key, mode=\"b\")\n return plaintext_io.read().decode(\"utf-8\")\n except Exception:\n logger.exception(f\"Error when reading file, id={file_id}\")\n\n # Cache miss — extract and store.\n content_text = extract_fn()\n if content_text:\n store_plaintext(file_id, content_text)\n return content_text\n\n\n@log_function_time(print_only=True)\ndef load_chat_file(\n file_descriptor: FileDescriptor, db_session: Session\n) -> ChatLoadedFile:\n file_io = get_default_file_store().read_file(file_descriptor[\"id\"], mode=\"b\")\n content = file_io.read()\n\n # Extract text content if it's a text file type (not an image)\n content_text = None\n # `FileDescriptor` is often JSON-roundtripped (e.g. JSONB / API), so `type`\n # may arrive as a raw string value instead of a `ChatFileType`.\n file_type = ChatFileType(file_descriptor[\"type\"])\n\n if file_type.is_text_file():\n file_id = file_descriptor[\"id\"]\n\n def _extract() -> str:\n return extract_file_text(\n file=file_io,\n file_name=file_descriptor.get(\"name\") or \"\",\n break_on_unprocessable=False,\n )\n\n # Use the user_file_id as cache key when available (matches what\n # the celery indexing worker stores), otherwise fall back to the\n # file store id (covers code-interpreter-generated files, etc.).\n user_file_id_str = file_descriptor.get(\"user_file_id\")\n cache_key = user_file_id_str or file_id\n\n try:\n content_text = _get_or_extract_plaintext(cache_key, _extract)\n except Exception as e:\n logger.warning(\n f\"Failed to retrieve content for file {file_descriptor['id']}: {str(e)}\"\n )\n\n # Get token count from UserFile if available\n token_count = 0\n user_file_id_str = file_descriptor.get(\"user_file_id\")\n if user_file_id_str:\n try:\n user_file_id = UUID(user_file_id_str)\n user_file = (\n db_session.query(UserFile).filter(UserFile.id == user_file_id).first()\n )\n if user_file and user_file.token_count:\n token_count = user_file.token_count\n except (ValueError, TypeError) as e:\n logger.warning(\n f\"Failed to get token count for file {file_descriptor['id']}: {e}\"\n )\n\n return ChatLoadedFile(\n file_id=file_descriptor[\"id\"],\n content=content,\n file_type=file_type,\n filename=file_descriptor.get(\"name\"),\n content_text=content_text,\n token_count=token_count,\n )\n\n\ndef load_all_chat_files(\n chat_messages: list[ChatMessage],\n db_session: Session,\n) -> list[ChatLoadedFile]:\n # TODO There is likely a more efficient/standard way to load the files here.\n file_descriptors_for_history: list[FileDescriptor] = []\n for chat_message in chat_messages:\n if chat_message.files:\n file_descriptors_for_history.extend(chat_message.files)\n\n files = cast(\n list[ChatLoadedFile],\n run_functions_tuples_in_parallel(\n [\n (load_chat_file, (file, db_session))\n for file in file_descriptors_for_history\n ]\n ),\n )\n return files\n\n\ndef convert_chat_history_basic(\n chat_history: list[ChatMessage],\n token_counter: Callable[[str], int],\n max_individual_message_tokens: int | None = None,\n max_total_tokens: int | None = None,\n) -> list[ChatMessageSimple]:\n \"\"\"Convert ChatMessage history to ChatMessageSimple format with no tool calls or files included.\n\n Args:\n chat_history: List of ChatMessage objects to convert\n token_counter: Function to count tokens in a message string\n max_individual_message_tokens: If set, messages exceeding this number of tokens are dropped.\n If None, no messages are dropped based on individual token count.\n max_total_tokens: If set, maximum number of tokens allowed for the entire history.\n If None, the history is not trimmed based on total token count.\n\n Returns:\n List of ChatMessageSimple objects\n \"\"\"\n # Defensive: treat a non-positive total budget as \"no history\".\n if max_total_tokens is not None and max_total_tokens <= 0:\n return []\n\n # Convert only the core USER/ASSISTANT messages; omit files and tool calls.\n converted: list[ChatMessageSimple] = []\n for chat_message in chat_history:\n if chat_message.message_type not in (MessageType.USER, MessageType.ASSISTANT):\n continue\n\n message = chat_message.message or \"\"\n token_count = getattr(chat_message, \"token_count\", None)\n if token_count is None:\n token_count = token_counter(message)\n\n # Drop any single message that would dominate the context window.\n if (\n max_individual_message_tokens is not None\n and token_count > max_individual_message_tokens\n ):\n continue\n\n converted.append(\n ChatMessageSimple(\n message=message,\n token_count=token_count,\n message_type=chat_message.message_type,\n image_files=None,\n )\n )\n\n if max_total_tokens is None:\n return converted\n\n # Enforce a max total budget by keeping a contiguous suffix of the conversation.\n trimmed_reversed: list[ChatMessageSimple] = []\n total_tokens = 0\n for msg in reversed(converted):\n if total_tokens + msg.token_count > max_total_tokens:\n break\n trimmed_reversed.append(msg)\n total_tokens += msg.token_count\n\n return list(reversed(trimmed_reversed))\n\n\ndef _build_tool_call_response_history_message(\n tool_name: str,\n generated_images: list[dict] | None,\n tool_call_response: str | None,\n) -> str:\n if tool_name != IMAGE_GENERATION_TOOL_NAME:\n return TOOL_CALL_RESPONSE_CROSS_MESSAGE\n\n if generated_images:\n llm_image_context: list[dict[str, str]] = []\n for image in generated_images:\n file_id = image.get(\"file_id\")\n revised_prompt = image.get(\"revised_prompt\")\n if not isinstance(file_id, str):\n continue\n\n llm_image_context.append(\n {\n \"file_id\": file_id,\n \"revised_prompt\": (\n revised_prompt if isinstance(revised_prompt, str) else \"\"\n ),\n }\n )\n\n if llm_image_context:\n return json.dumps(llm_image_context)\n\n if tool_call_response:\n return tool_call_response\n\n return TOOL_CALL_RESPONSE_CROSS_MESSAGE\n\n\ndef convert_chat_history(\n chat_history: list[ChatMessage],\n files: list[ChatLoadedFile],\n context_image_files: list[ChatLoadedFile],\n additional_context: str | None,\n token_counter: Callable[[str], int],\n tool_id_to_name_map: dict[int, str],\n) -> ChatHistoryResult:\n \"\"\"Convert ChatMessage history to ChatMessageSimple format.\n\n For user messages: includes attached files (images attached to message, text files as separate messages)\n For assistant messages with tool calls: creates ONE ASSISTANT message with tool_calls array,\n followed by N TOOL_CALL_RESPONSE messages (OpenAI parallel tool calling format)\n For assistant messages without tool calls: creates a simple ASSISTANT message\n\n Every injected text-file message is tagged with ``file_id`` and its\n metadata is collected in ``ChatHistoryResult.all_injected_file_metadata``.\n After context-window truncation, callers compare surviving ``file_id`` tags\n against this map to discover \"forgotten\" files and provide their metadata\n to the FileReaderTool.\n \"\"\"\n simple_messages: list[ChatMessageSimple] = []\n all_injected_file_metadata: dict[str, FileToolMetadata] = {}\n\n # Create a mapping of file IDs to loaded files for quick lookup\n file_map = {str(f.file_id): f for f in files}\n\n # Find the index of the last USER message\n last_user_message_idx = None\n for i in range(len(chat_history) - 1, -1, -1):\n if chat_history[i].message_type == MessageType.USER:\n last_user_message_idx = i\n break\n\n for idx, chat_message in enumerate(chat_history):\n if chat_message.message_type == MessageType.USER:\n # Process files attached to this message\n text_files: list[tuple[ChatLoadedFile, FileDescriptor]] = []\n image_files: list[ChatLoadedFile] = []\n\n if chat_message.files:\n for file_descriptor in chat_message.files:\n file_id = file_descriptor[\"id\"]\n loaded_file = file_map.get(file_id)\n if loaded_file:\n if loaded_file.file_type == ChatFileType.IMAGE:\n image_files.append(loaded_file)\n else:\n # Text files (DOC, PLAIN_TEXT, TABULAR) are added as separate messages\n text_files.append((loaded_file, file_descriptor))\n\n # Add text files as separate messages before the user message.\n # Each message is tagged with ``file_id`` so that forgotten files\n # can be detected after context-window truncation.\n for text_file, fd in text_files:\n # Use user_file_id as the FileReaderTool accepts that.\n # Fall back to the file-store path id.\n tool_id = fd.get(\"user_file_id\") or text_file.file_id\n filename = text_file.filename or \"unknown\"\n ctx = build_file_context(\n tool_file_id=tool_id,\n filename=filename,\n file_type=text_file.file_type,\n content_text=text_file.content_text,\n token_count=text_file.token_count,\n )\n simple_messages.append(ctx.message)\n all_injected_file_metadata[tool_id] = ctx.tool_metadata\n\n # Sum token counts from image files (excluding project image files)\n image_token_count = (\n sum(img.token_count for img in image_files) if image_files else 0\n )\n\n # Add the user message with image files attached\n # If this is the last USER message, also include context_image_files\n # Note: context image file tokens are NOT counted in the token count\n if idx == last_user_message_idx:\n if context_image_files:\n image_files.extend(context_image_files)\n\n if additional_context:\n simple_messages.append(\n ChatMessageSimple(\n message=ADDITIONAL_CONTEXT_PROMPT.format(\n additional_context=additional_context\n ),\n token_count=token_counter(additional_context),\n message_type=MessageType.USER,\n image_files=None,\n )\n )\n\n simple_messages.append(\n ChatMessageSimple(\n message=chat_message.message,\n token_count=chat_message.token_count + image_token_count,\n message_type=MessageType.USER,\n image_files=image_files if image_files else None,\n )\n )\n\n elif chat_message.message_type == MessageType.ASSISTANT:\n # Handle tool calls if present using OpenAI parallel tool calling format:\n # 1. Group tool calls by turn_number\n # 2. For each turn: ONE ASSISTANT message with tool_calls array\n # 3. Followed by N TOOL_CALL_RESPONSE messages (one per tool call)\n if chat_message.tool_calls:\n # Group tool calls by turn number\n tool_calls_by_turn: dict[int, list] = {}\n for tool_call in chat_message.tool_calls:\n if tool_call.turn_number not in tool_calls_by_turn:\n tool_calls_by_turn[tool_call.turn_number] = []\n tool_calls_by_turn[tool_call.turn_number].append(tool_call)\n\n # Sort turns and process each turn\n for turn_number in sorted(tool_calls_by_turn.keys()):\n turn_tool_calls = tool_calls_by_turn[turn_number]\n # Sort by tool_id within the turn for consistent ordering\n turn_tool_calls.sort(key=lambda tc: tc.tool_id)\n\n # Build ToolCallSimple list for this turn\n tool_calls_simple: list[ToolCallSimple] = []\n for tool_call in turn_tool_calls:\n tool_name = tool_id_to_name_map.get(\n tool_call.tool_id, \"unknown\"\n )\n tool_calls_simple.append(\n ToolCallSimple(\n tool_call_id=tool_call.tool_call_id,\n tool_name=tool_name,\n tool_arguments=tool_call.tool_call_arguments or {},\n token_count=tool_call.tool_call_tokens,\n )\n )\n\n # Create ONE ASSISTANT message with all tool calls for this turn\n total_tool_call_tokens = sum(\n tc.token_count for tc in tool_calls_simple\n )\n simple_messages.append(\n ChatMessageSimple(\n message=\"\", # No text content when making tool calls\n token_count=total_tool_call_tokens,\n message_type=MessageType.ASSISTANT,\n tool_calls=tool_calls_simple,\n image_files=None,\n )\n )\n\n # Add TOOL_CALL_RESPONSE messages for each tool call in this turn\n for tool_call in turn_tool_calls:\n tool_name = tool_id_to_name_map.get(\n tool_call.tool_id, \"unknown\"\n )\n tool_response_message = (\n _build_tool_call_response_history_message(\n tool_name=tool_name,\n generated_images=tool_call.generated_images,\n tool_call_response=tool_call.tool_call_response,\n )\n )\n simple_messages.append(\n ChatMessageSimple(\n message=tool_response_message,\n token_count=(\n token_counter(tool_response_message)\n if tool_name == IMAGE_GENERATION_TOOL_NAME\n else 20\n ),\n message_type=MessageType.TOOL_CALL_RESPONSE,\n tool_call_id=tool_call.tool_call_id,\n image_files=None,\n )\n )\n\n # Add the assistant message itself (the final answer)\n simple_messages.append(\n ChatMessageSimple(\n message=chat_message.message,\n token_count=chat_message.token_count,\n message_type=MessageType.ASSISTANT,\n image_files=None,\n )\n )\n else:\n raise ValueError(\n f\"Invalid message type when constructing simple history: {chat_message.message_type}\"\n )\n\n return ChatHistoryResult(\n simple_messages=simple_messages,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n\n\ndef get_custom_agent_prompt(persona: Persona, chat_session: ChatSession) -> str | None:\n \"\"\"Get the custom agent prompt from persona or project instructions. If it's replacing the base system prompt,\n it does not count as a custom agent prompt (logic exists later also to drop it in this case).\n\n Chat Sessions in Projects that are using a custom agent will retain the custom agent prompt.\n Priority: persona.system_prompt (if not default Agent) > chat_session.project.instructions\n\n # NOTE: Logic elsewhere allows saving empty strings for potentially other purposes but for constructing the prompts\n # we never want to return an empty string for a prompt so it's translated into an explicit None.\n\n Args:\n persona: The Persona object\n chat_session: The ChatSession object\n\n Returns:\n The prompt to use for the custom Agent part of the prompt.\n \"\"\"\n # If using a custom Agent, always respect its prompt, even if in a Project, and even if it's an empty custom prompt.\n if persona.id != DEFAULT_PERSONA_ID:\n # Logic exists later also to drop it in this case but this is strictly correct anyhow.\n if persona.replace_base_system_prompt:\n return None\n return persona.system_prompt or None\n\n # If in a project and using the default Agent, respect the project instructions.\n if chat_session.project and chat_session.project.instructions:\n return chat_session.project.instructions\n\n return None\n\n\ndef is_last_assistant_message_clarification(chat_history: list[ChatMessage]) -> bool:\n \"\"\"Check if the last assistant message in chat history was a clarification question.\n\n This is used in the deep research flow to determine whether to skip the\n clarification step when the user has already responded to a clarification.\n\n Args:\n chat_history: List of ChatMessage objects in chronological order\n\n Returns:\n True if the last assistant message has is_clarification=True, False otherwise\n \"\"\"\n for message in reversed(chat_history):\n if message.message_type == MessageType.ASSISTANT:\n return message.is_clarification\n return False\n\n\ndef create_tool_call_failure_messages(\n tool_calls: list[ToolCallKickoff], token_counter: Callable[[str], int]\n) -> list[ChatMessageSimple]:\n \"\"\"Create ChatMessageSimple objects for failed tool calls.\n\n Creates messages using OpenAI parallel tool calling format:\n 1. An ASSISTANT message with tool_calls field containing all failed tool calls\n 2. A TOOL_CALL_RESPONSE failure message for each tool call\n\n Args:\n tool_calls: List of ToolCallKickoff objects representing the failed tool calls\n token_counter: Function to count tokens in a message string\n\n Returns:\n List containing ChatMessageSimple objects: one assistant message with all tool calls\n followed by a failure response for each tool call\n \"\"\"\n if not tool_calls:\n return []\n\n # Create ToolCallSimple for each failed tool call\n tool_calls_simple: list[ToolCallSimple] = []\n for tool_call in tool_calls:\n tool_call_token_count = token_counter(tool_call.to_msg_str())\n tool_calls_simple.append(\n ToolCallSimple(\n tool_call_id=tool_call.tool_call_id,\n tool_name=tool_call.tool_name,\n tool_arguments=tool_call.tool_args,\n token_count=tool_call_token_count,\n )\n )\n\n total_token_count = sum(tc.token_count for tc in tool_calls_simple)\n\n # Create ONE ASSISTANT message with all tool_calls (OpenAI format)\n assistant_msg = ChatMessageSimple(\n message=\"\", # No text content when making tool calls\n token_count=total_token_count,\n message_type=MessageType.ASSISTANT,\n tool_calls=tool_calls_simple,\n image_files=None,\n )\n\n messages: list[ChatMessageSimple] = [assistant_msg]\n\n # Create a TOOL_CALL_RESPONSE failure message for each tool call\n for tool_call in tool_calls:\n failure_response_msg = ChatMessageSimple(\n message=TOOL_CALL_FAILURE_PROMPT,\n token_count=50, # Tiny overestimate\n message_type=MessageType.TOOL_CALL_RESPONSE,\n tool_call_id=tool_call.tool_call_id,\n image_files=None,\n )\n messages.append(failure_response_msg)\n\n return messages\n" }, { "path": "backend/onyx/chat/citation_processor.py", "content": "\"\"\"\nDynamic Citation Processor for LLM Responses\n\nThis module provides a citation processor that can:\n- Accept citation number to SearchDoc mappings dynamically\n- Process token streams from LLMs to extract citations\n- Handle citations in three modes: REMOVE, KEEP_MARKERS, or HYPERLINK\n- Emit CitationInfo objects for detected citations (in HYPERLINK mode)\n- Track all seen citations regardless of mode\n- Maintain a list of cited documents in order of first citation\n\"\"\"\n\nimport re\nfrom collections.abc import Generator\nfrom enum import Enum\nfrom typing import TypeAlias\n\nfrom onyx.configs.chat_configs import STOP_STREAM_PAT\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.prompts.constants import TRIPLE_BACKTICK\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass CitationMode(Enum):\n \"\"\"Defines how citations should be handled in the output.\n\n REMOVE: Citations are completely removed from output text.\n No CitationInfo objects are emitted.\n Use case: When you need to remove citations from the output if they are not shared with the user\n (e.g. in discord bot, public slack bot).\n\n KEEP_MARKERS: Original citation markers like [1], [2] are preserved unchanged.\n No CitationInfo objects are emitted.\n Use case: When you need to track citations in research agent and later process\n them with collapse_citations() to renumber.\n\n HYPERLINK: Citations are replaced with markdown links like [[1]](url).\n CitationInfo objects are emitted for UI tracking.\n Use case: Final reports shown to users with clickable links.\n \"\"\"\n\n REMOVE = \"remove\"\n KEEP_MARKERS = \"keep_markers\"\n HYPERLINK = \"hyperlink\"\n\n\nCitationMapping: TypeAlias = dict[int, SearchDoc]\n\n\n# ============================================================================\n# Utility functions\n# ============================================================================\n\n\ndef in_code_block(llm_text: str) -> bool:\n \"\"\"Check if we're currently inside a code block by counting triple backticks.\"\"\"\n count = llm_text.count(TRIPLE_BACKTICK)\n return count % 2 != 0\n\n\n# ============================================================================\n# Main Citation Processor with Dynamic Mapping\n# ============================================================================\n\n\nclass DynamicCitationProcessor:\n \"\"\"\n A citation processor that accepts dynamic citation mappings.\n\n This processor is designed for multi-turn conversations where the citation\n number to document mapping is provided externally. It processes streaming\n tokens from an LLM, detects citations (e.g., [1], [2,3], [[4]]), and handles\n them according to the configured CitationMode:\n\n CitationMode.HYPERLINK (default):\n 1. Replaces citation markers with formatted markdown links (e.g., [[1]](url))\n 2. Emits CitationInfo objects for tracking\n 3. Maintains the order in which documents were first cited\n Use case: Final reports shown to users with clickable links.\n\n CitationMode.KEEP_MARKERS:\n 1. Preserves original citation markers like [1], [2] unchanged\n 2. Does NOT emit CitationInfo objects\n 3. Still tracks all seen citations via get_seen_citations()\n Use case: When citations need later processing (e.g., renumbering).\n\n CitationMode.REMOVE:\n 1. Removes citation markers entirely from the output text\n 2. Does NOT emit CitationInfo objects\n 3. Still tracks all seen citations via get_seen_citations()\n Use case: Research agent intermediate reports.\n\n Features:\n - Accepts citation number → SearchDoc mapping via update_citation_mapping()\n - Configurable citation mode at initialization\n - Always tracks seen citations regardless of mode\n - Holds back tokens that might be partial citations\n - Maintains list of cited SearchDocs in order of first citation\n - Handles unicode bracket variants (【】, ［］)\n - Skips citation processing inside code blocks\n\n Example (HYPERLINK mode - default):\n processor = DynamicCitationProcessor()\n\n # Set up citation mapping\n processor.update_citation_mapping({1: search_doc1, 2: search_doc2})\n\n # Process tokens from LLM\n for token in llm_stream:\n for result in processor.process_token(token):\n if isinstance(result, str):\n print(result) # Display text with [[1]](url) format\n elif isinstance(result, CitationInfo):\n handle_citation(result) # Track citation\n\n # Get cited documents at the end\n cited_docs = processor.get_cited_documents()\n\n Example (KEEP_MARKERS mode):\n processor = DynamicCitationProcessor(citation_mode=CitationMode.KEEP_MARKERS)\n processor.update_citation_mapping({1: search_doc1, 2: search_doc2})\n\n # Process tokens from LLM\n for token in llm_stream:\n for result in processor.process_token(token):\n # Only strings are yielded, no CitationInfo objects\n print(result) # Display text with original [1] format preserved\n\n # Get all seen citations after processing\n seen_citations = processor.get_seen_citations() # {1: search_doc1, ...}\n\n Example (REMOVE mode):\n processor = DynamicCitationProcessor(citation_mode=CitationMode.REMOVE)\n processor.update_citation_mapping({1: search_doc1, 2: search_doc2})\n\n # Process tokens - citations are removed but tracked\n for token in llm_stream:\n for result in processor.process_token(token):\n print(result) # Text without any citation markers\n\n # Citations are still tracked\n seen_citations = processor.get_seen_citations()\n \"\"\"\n\n def __init__(\n self,\n citation_mode: CitationMode = CitationMode.HYPERLINK,\n stop_stream: str | None = STOP_STREAM_PAT,\n ):\n \"\"\"\n Initialize the citation processor.\n\n Args:\n citation_mode: How to handle citations in the output. One of:\n - CitationMode.HYPERLINK (default): Replace [1] with [[1]](url)\n and emit CitationInfo objects.\n - CitationMode.KEEP_MARKERS: Keep original [1] markers unchanged,\n no CitationInfo objects emitted.\n - CitationMode.REMOVE: Remove citations entirely from output,\n no CitationInfo objects emitted.\n All modes track seen citations via get_seen_citations().\n stop_stream: Optional stop token pattern to halt processing early.\n When this pattern is detected in the token stream, processing stops.\n Defaults to STOP_STREAM_PAT from chat configs.\n \"\"\"\n\n # Citation mapping from citation number to SearchDoc\n self.citation_to_doc: CitationMapping = {}\n self.seen_citations: CitationMapping = {} # citation num -> SearchDoc\n\n # Token processing state\n self.llm_out = \"\" # entire output so far\n self.curr_segment = \"\" # tokens held for citation processing\n self.hold = \"\" # tokens held for stop token processing\n self.stop_stream = stop_stream\n self.citation_mode = citation_mode\n\n # Citation tracking\n self.cited_documents_in_order: list[SearchDoc] = (\n []\n ) # SearchDocs in citation order\n self.cited_document_ids: set[str] = set() # all cited document_ids\n self.recent_cited_documents: set[str] = (\n set()\n ) # recently cited (for deduplication)\n self.non_citation_count = 0\n\n # Citation patterns\n # Matches potential incomplete citations: '[', '[[', '[1', '[[1', '[1,', '[1, ', etc.\n # Also matches unicode bracket variants: 【, ［\n self.possible_citation_pattern = re.compile(r\"([\\[【［]+(?:\\d+,? ?)*$)\")\n\n # Matches complete citations:\n # group 1: '[[1]]', [[2]], etc. (also matches 【【1】】, ［［1］］, 【1】, ［1］)\n # group 2: '[1]', '[1, 2]', '[1,2,16]', etc. (also matches unicode variants)\n self.citation_pattern = re.compile(\n r\"([\\[【［]{2}\\d+[\\]】］]{2})|([\\[【［]\\d+(?:, ?\\d+)*[\\]】］])\"\n )\n\n def update_citation_mapping(\n self,\n citation_mapping: CitationMapping,\n update_duplicate_keys: bool = False,\n ) -> None:\n \"\"\"\n Update the citation number to SearchDoc mapping.\n\n This can be called multiple times to add or update mappings. New mappings\n will be merged with existing ones.\n\n Args:\n citation_mapping: Dictionary mapping citation numbers (1, 2, 3, ...) to SearchDoc objects\n update_duplicate_keys: If True, update existing mappings with new values when keys overlap.\n If False (default), filter out duplicate keys and only add non-duplicates.\n The default behavior is useful when OpenURL may have the same citation number as a\n Web Search result - in those cases, we keep the web search citation and snippet etc.\n \"\"\"\n if update_duplicate_keys:\n # Update all mappings, including duplicates\n self.citation_to_doc.update(citation_mapping)\n else:\n # Filter out duplicate keys and only add non-duplicates\n # Reason for this is that OpenURL may have the same citation number as a Web Search result\n # For those, we should just keep the web search citation and snippet etc.\n duplicate_keys = set(citation_mapping.keys()) & set(\n self.citation_to_doc.keys()\n )\n non_duplicate_mapping = {\n k: v for k, v in citation_mapping.items() if k not in duplicate_keys\n }\n self.citation_to_doc.update(non_duplicate_mapping)\n\n def process_token(\n self, token: str | None\n ) -> Generator[str | CitationInfo, None, None]:\n \"\"\"\n Process a token from the LLM stream.\n\n This method:\n 1. Accumulates tokens until a complete citation or non-citation is found\n 2. Holds back potential partial citations (e.g., \"[\", \"[1\")\n 3. Yields text chunks when they're safe to display\n 4. Handles code blocks (avoids processing citations inside code)\n 5. Handles stop tokens\n 6. Always tracks seen citations in self.seen_citations\n\n Behavior depends on the `citation_mode` setting from __init__:\n - HYPERLINK: Citations are replaced with [[n]](url) format and CitationInfo\n objects are yielded before each formatted citation\n - KEEP_MARKERS: Original citation markers like [1] are preserved unchanged,\n no CitationInfo objects are yielded\n - REMOVE: Citations are removed entirely from output,\n no CitationInfo objects are yielded\n\n Args:\n token: The next token from the LLM stream, or None to signal end of stream.\n Pass None to flush any remaining buffered text at end of stream.\n\n Yields:\n str: Text chunks to display. Citation format depends on citation_mode.\n CitationInfo: Citation metadata (only when citation_mode=HYPERLINK)\n \"\"\"\n # None -> end of stream, flush remaining segment\n if token is None:\n if self.curr_segment:\n yield self.curr_segment\n return\n\n # Handle stop stream token\n if self.stop_stream:\n next_hold = self.hold + token\n if self.stop_stream in next_hold:\n # Extract text before the stop pattern\n stop_pos = next_hold.find(self.stop_stream)\n text_before_stop = next_hold[:stop_pos]\n # Process the text before stop pattern if any exists\n if text_before_stop:\n # Process text_before_stop through normal flow\n self.hold = \"\"\n token = text_before_stop\n # Continue to normal processing below\n else:\n # Stop pattern at the beginning, nothing to yield\n return\n elif next_hold == self.stop_stream[: len(next_hold)]:\n self.hold = next_hold\n return\n else:\n token = next_hold\n self.hold = \"\"\n\n self.curr_segment += token\n self.llm_out += token\n\n # Handle code blocks without language tags\n # If we see ``` followed by \\n, add \"plaintext\" language specifier\n if \"`\" in self.curr_segment:\n if self.curr_segment.endswith(\"`\"):\n pass\n elif \"```\" in self.curr_segment:\n parts = self.curr_segment.split(\"```\")\n if len(parts) > 1 and len(parts[1]) > 0:\n piece_that_comes_after = parts[1][0]\n if piece_that_comes_after == \"\\n\" and in_code_block(self.llm_out):\n self.curr_segment = self.curr_segment.replace(\n \"```\", \"```plaintext\"\n )\n\n # Look for citations in current segment\n citation_matches = list(self.citation_pattern.finditer(self.curr_segment))\n possible_citation_found = bool(\n re.search(self.possible_citation_pattern, self.curr_segment)\n )\n\n result = \"\"\n if citation_matches and not in_code_block(self.llm_out):\n match_idx = 0\n for match in citation_matches:\n match_span = match.span()\n\n # Get text before/between citations\n intermatch_str = self.curr_segment[match_idx : match_span[0]]\n self.non_citation_count += len(intermatch_str)\n match_idx = match_span[1]\n\n # Check if there is already a space before this citation\n if intermatch_str:\n has_leading_space = intermatch_str[-1].isspace()\n else:\n # No text between citations (consecutive citations)\n # If match_idx > 0, we've already processed a citation, so don't add space\n if match_idx > 0:\n # Consecutive citations - don't add space between them\n has_leading_space = True\n else:\n # Citation at start of segment - check if previous output has space\n segment_start_idx = len(self.llm_out) - len(self.curr_segment)\n if segment_start_idx > 0:\n has_leading_space = self.llm_out[\n segment_start_idx - 1\n ].isspace()\n else:\n has_leading_space = False\n\n # Reset recent citations if no citations found for a while\n if self.non_citation_count > 5:\n self.recent_cited_documents.clear()\n\n # Process the citation (returns formatted citation text and CitationInfo objects)\n # Always tracks seen citations regardless of citation_mode\n citation_text, citation_info_list = self._process_citation(\n match, has_leading_space\n )\n\n if self.citation_mode == CitationMode.HYPERLINK:\n # HYPERLINK mode: Replace citations with markdown links [[n]](url)\n # Yield text before citation FIRST (preserve order)\n if intermatch_str:\n yield intermatch_str\n # Yield CitationInfo objects BEFORE the citation text\n # This allows the frontend to receive citation metadata before the token\n # that contains [[n]](link), enabling immediate rendering\n for citation in citation_info_list:\n yield citation\n # Then yield the formatted citation text\n if citation_text:\n yield citation_text\n\n elif self.citation_mode == CitationMode.KEEP_MARKERS:\n # KEEP_MARKERS mode: Preserve original citation markers unchanged\n # Yield text before citation\n if intermatch_str:\n yield intermatch_str\n # Yield the original citation marker as-is\n yield match.group()\n\n else: # CitationMode.REMOVE\n # REMOVE mode: Remove citations entirely from output\n # This strips citation markers like [1], [2], 【1】 from the output text\n # When removing citations, we need to handle spacing to avoid issues like:\n # - \"text [1] more\" -> \"text more\" (double space)\n # - \"text [1].\" -> \"text .\" (space before punctuation)\n if intermatch_str:\n remaining_text = self.curr_segment[match_span[1] :]\n # Strip trailing space from intermatch if:\n # 1. Remaining text starts with space (avoids double space)\n # 2. Remaining text starts with punctuation (avoids space before punctuation)\n if intermatch_str[-1].isspace() and remaining_text:\n first_char = remaining_text[0]\n # Check if next char is space or common punctuation\n if first_char.isspace() or first_char in \".,;:!?)]}\":\n intermatch_str = intermatch_str.rstrip()\n if intermatch_str:\n yield intermatch_str\n\n self.non_citation_count = 0\n\n # Leftover text could be part of next citation\n self.curr_segment = self.curr_segment[match_idx:]\n self.non_citation_count = len(self.curr_segment)\n\n # Hold onto the current segment if potential citations found, otherwise stream it\n if not possible_citation_found:\n result += self.curr_segment\n self.non_citation_count += len(self.curr_segment)\n self.curr_segment = \"\"\n\n if result:\n yield result\n\n def _process_citation(\n self, match: re.Match, has_leading_space: bool\n ) -> tuple[str, list[CitationInfo]]:\n \"\"\"\n Process a single citation match and return formatted citation text and citation info objects.\n\n This is an internal method called by process_token(). The match string can be\n in various formats: '[1]', '[1, 13, 6]', '[[4]]', '【1】', '［1］', etc.\n\n This method always:\n 1. Extracts citation numbers from the match\n 2. Looks up the corresponding SearchDoc from the mapping\n 3. Tracks seen citations in self.seen_citations (regardless of citation_mode)\n\n When citation_mode is HYPERLINK:\n 4. Creates formatted citation text as [[n]](url)\n 5. Creates CitationInfo objects for new citations\n 6. Handles deduplication of recently cited documents\n\n When citation_mode is REMOVE or KEEP_MARKERS:\n 4. Returns empty string and empty list (caller handles output based on mode)\n\n Args:\n match: Regex match object containing the citation pattern\n has_leading_space: Whether the text immediately before this citation\n ends with whitespace. Used to determine if a leading space should\n be added to the formatted output.\n\n Returns:\n Tuple of (formatted_citation_text, citation_info_list):\n - formatted_citation_text: Markdown-formatted citation text like\n \"[[1]](https://example.com)\" or empty string if not in HYPERLINK mode\n - citation_info_list: List of CitationInfo objects for newly cited\n documents, or empty list if not in HYPERLINK mode\n \"\"\"\n citation_str: str = match.group() # e.g., '[1]', '[1, 2, 3]', '[[1]]', '【1】'\n formatted = (\n match.lastindex == 1\n ) # True means already in form '[[1]]' or '【【1】】'\n\n citation_info_list: list[CitationInfo] = []\n formatted_citation_parts: list[str] = []\n\n # Extract citation numbers - regex ensures matched brackets, so we can simply slice\n citation_content = citation_str[2:-2] if formatted else citation_str[1:-1]\n\n for num_str in citation_content.split(\",\"):\n num_str = num_str.strip()\n if not num_str:\n continue\n\n try:\n num = int(num_str)\n except ValueError:\n # Invalid citation, skip it\n logger.warning(f\"Invalid citation number format: {num_str}\")\n continue\n\n # Check if we have a mapping for this citation number\n if num not in self.citation_to_doc:\n logger.warning(\n f\"Citation number {num} not found in mapping. Available: {list(self.citation_to_doc.keys())}\"\n )\n continue\n\n # Get the SearchDoc\n search_doc = self.citation_to_doc[num]\n doc_id = search_doc.document_id\n link = search_doc.link or \"\"\n\n # Always track seen citations regardless of citation_mode setting\n self.seen_citations[num] = search_doc\n\n # Only generate formatted citations and CitationInfo in HYPERLINK mode\n if self.citation_mode != CitationMode.HYPERLINK:\n continue\n\n # Format the citation text as [[n]](link)\n formatted_citation_parts.append(f\"[[{num}]]({link})\")\n\n # Skip creating CitationInfo for citations of the same work if cited recently (deduplication)\n if doc_id in self.recent_cited_documents:\n continue\n self.recent_cited_documents.add(doc_id)\n\n # Track cited documents and create CitationInfo only for new citations\n if doc_id not in self.cited_document_ids:\n self.cited_document_ids.add(doc_id)\n self.cited_documents_in_order.append(search_doc)\n citation_info_list.append(\n CitationInfo(\n citation_number=num,\n document_id=doc_id,\n )\n )\n\n # Join all citation parts with spaces\n formatted_citation_text = \" \".join(formatted_citation_parts)\n\n # Apply leading space only if the text didn't already have one\n if formatted_citation_text and not has_leading_space:\n formatted_citation_text = \" \" + formatted_citation_text\n\n return formatted_citation_text, citation_info_list\n\n def get_cited_documents(self) -> list[SearchDoc]:\n \"\"\"\n Get the list of cited SearchDoc objects in the order they were first cited.\n\n Note: This list is only populated when `citation_mode=HYPERLINK`.\n When using REMOVE or KEEP_MARKERS mode, this will return an empty list.\n Use get_seen_citations() instead if you need to track citations without\n emitting CitationInfo objects.\n\n Returns:\n List of SearchDoc objects in the order they were first cited.\n Empty list if citation_mode is not HYPERLINK.\n \"\"\"\n return self.cited_documents_in_order\n\n def get_cited_document_ids(self) -> list[str]:\n \"\"\"\n Get the list of cited document IDs in the order they were first cited.\n\n Note: This list is only populated when `citation_mode=HYPERLINK`.\n When using REMOVE or KEEP_MARKERS mode, this will return an empty list.\n Use get_seen_citations() instead if you need to track citations without\n emitting CitationInfo objects.\n\n Returns:\n List of document IDs (strings) in the order they were first cited.\n Empty list if citation_mode is not HYPERLINK.\n \"\"\"\n return [doc.document_id for doc in self.cited_documents_in_order]\n\n def get_seen_citations(self) -> CitationMapping:\n \"\"\"\n Get all seen citations as a mapping from citation number to SearchDoc.\n\n This returns all citations that have been encountered during processing,\n regardless of the `citation_mode` setting. Citations are tracked\n whenever they are parsed, making this useful for cases where you need to\n know which citations appeared in the text without emitting CitationInfo objects.\n\n This is particularly useful when using REMOVE or KEEP_MARKERS mode, as\n get_cited_documents() will be empty in those cases, but get_seen_citations()\n will still contain all the citations that were found.\n\n Returns:\n Dictionary mapping citation numbers (int) to SearchDoc objects.\n The dictionary is keyed by the citation number as it appeared in\n the text (e.g., {1: SearchDoc(...), 3: SearchDoc(...)}).\n \"\"\"\n return self.seen_citations\n\n @property\n def num_cited_documents(self) -> int:\n \"\"\"\n Get the number of unique documents that have been cited.\n\n Note: This count is only updated when `citation_mode=HYPERLINK`.\n When using REMOVE or KEEP_MARKERS mode, this will always return 0.\n Use len(get_seen_citations()) instead if you need to count citations\n without emitting CitationInfo objects.\n\n Returns:\n Number of unique documents cited. 0 if citation_mode is not HYPERLINK.\n \"\"\"\n return len(self.cited_document_ids)\n\n def reset_recent_citations(self) -> None:\n \"\"\"\n Reset the recent citations tracker.\n\n The processor tracks \"recently cited\" documents to avoid emitting duplicate\n CitationInfo objects for the same document when it's cited multiple times\n in close succession. This method clears that tracker.\n\n This is primarily useful when `citation_mode=HYPERLINK` to allow\n previously cited documents to emit CitationInfo objects again. Has no\n effect when using REMOVE or KEEP_MARKERS mode.\n\n The recent citation tracker is also automatically cleared when more than\n 5 non-citation characters are processed between citations.\n \"\"\"\n self.recent_cited_documents.clear()\n\n def get_next_citation_number(self) -> int:\n \"\"\"\n Get the next available citation number for adding new documents to the mapping.\n\n This method returns the next citation number that should be used when adding\n new documents via update_citation_mapping(). Useful when dynamically adding\n citations during processing (e.g., from tool results like web search).\n\n If no citations exist yet in the mapping, returns 1.\n Otherwise, returns max(existing_citation_numbers) + 1.\n\n Returns:\n The next available citation number (1-indexed integer).\n\n Example:\n # After adding citations 1, 2, 3\n processor.get_next_citation_number() # Returns 4\n\n # With non-sequential citations 1, 5, 10\n processor.get_next_citation_number() # Returns 11\n \"\"\"\n if not self.citation_to_doc:\n return 1\n return max(self.citation_to_doc.keys()) + 1\n" }, { "path": "backend/onyx/chat/citation_utils.py", "content": "import re\n\nfrom onyx.chat.citation_processor import CitationMapping\nfrom onyx.chat.citation_processor import DynamicCitationProcessor\nfrom onyx.context.search.models import SearchDocsResponse\nfrom onyx.tools.built_in_tools import CITEABLE_TOOLS_NAMES\nfrom onyx.tools.models import ToolResponse\n\n\ndef update_citation_processor_from_tool_response(\n tool_response: ToolResponse,\n citation_processor: DynamicCitationProcessor,\n) -> None:\n \"\"\"Update citation processor if this was a citeable tool with a SearchDocsResponse.\n\n Checks if the tool call is citeable and if the response contains a SearchDocsResponse,\n then creates a mapping from citation numbers to SearchDoc objects and updates the\n citation processor.\n\n Args:\n tool_response: The response from the tool execution (must have tool_call set)\n citation_processor: The DynamicCitationProcessor to update\n \"\"\"\n # Early return if tool_call is not set\n if tool_response.tool_call is None:\n return\n\n # Update citation processor if this was a search tool\n if tool_response.tool_call.tool_name in CITEABLE_TOOLS_NAMES:\n # Check if the rich_response is a SearchDocsResponse\n if isinstance(tool_response.rich_response, SearchDocsResponse):\n search_response = tool_response.rich_response\n\n # Create mapping from citation number to SearchDoc\n citation_to_doc: CitationMapping = {}\n for (\n citation_num,\n doc_id,\n ) in search_response.citation_mapping.items():\n # Find the SearchDoc with this doc_id\n matching_doc = next(\n (\n doc\n for doc in search_response.search_docs\n if doc.document_id == doc_id\n ),\n None,\n )\n if matching_doc:\n citation_to_doc[citation_num] = matching_doc\n\n # Update the citation processor\n citation_processor.update_citation_mapping(citation_to_doc)\n\n\ndef extract_citation_order_from_text(text: str) -> list[int]:\n \"\"\"Extract citation numbers from text in order of first appearance.\n\n Parses citation patterns like [1], [1, 2], [[1]], 【1】 etc. and returns\n the citation numbers in the order they first appear in the text.\n\n Args:\n text: The text containing citations\n\n Returns:\n List of citation numbers in order of first appearance (no duplicates)\n \"\"\"\n # Same pattern used in collapse_citations and DynamicCitationProcessor\n # Group 2 captures the number in double bracket format: [[1]], 【【1】】\n # Group 4 captures the numbers in single bracket format: [1], [1, 2]\n citation_pattern = re.compile(\n r\"([\\[【［]{2}(\\d+)[\\]】］]{2})|([\\[【［]([\\d]+(?: *, *\\d+)*)[\\]】］])\"\n )\n seen: set[int] = set()\n order: list[int] = []\n\n for match in citation_pattern.finditer(text):\n # Group 2 is for double bracket single number, group 4 is for single bracket\n if match.group(2):\n nums_str = match.group(2)\n elif match.group(4):\n nums_str = match.group(4)\n else:\n continue\n\n for num_str in nums_str.split(\",\"):\n num_str = num_str.strip()\n if num_str:\n try:\n num = int(num_str)\n if num not in seen:\n seen.add(num)\n order.append(num)\n except ValueError:\n continue\n\n return order\n\n\ndef collapse_citations(\n answer_text: str,\n existing_citation_mapping: CitationMapping,\n new_citation_mapping: CitationMapping,\n) -> tuple[str, CitationMapping]:\n \"\"\"Collapse the citations in the text to use the smallest possible numbers.\n\n This function takes citations in the text (like [25], [30], etc.) and replaces them\n with the smallest possible numbers. It starts numbering from the next available\n integer after the existing citation mapping. If a citation refers to a document\n that already exists in the existing citation mapping (matched by document_id),\n it uses the existing citation number instead of assigning a new one.\n\n Args:\n answer_text: The text containing citations to collapse (e.g., \"See [25] and [30]\")\n existing_citation_mapping: Citations already processed/displayed. These mappings\n are preserved unchanged in the output.\n new_citation_mapping: Citations from the current text that need to be collapsed.\n The keys are the citation numbers as they appear in answer_text.\n\n Returns:\n A tuple of (updated_text, combined_mapping) where:\n - updated_text: The text with citations replaced with collapsed numbers\n - combined_mapping: All values from existing_citation_mapping plus the new\n mappings with their (possibly renumbered) keys\n \"\"\"\n # Build a reverse lookup: document_id -> existing citation number\n doc_id_to_existing_citation: dict[str, int] = {\n doc.document_id: citation_num\n for citation_num, doc in existing_citation_mapping.items()\n }\n\n # Determine the next available citation number\n if existing_citation_mapping:\n next_citation_num = max(existing_citation_mapping.keys()) + 1\n else:\n next_citation_num = 1\n\n # Build the mapping from old citation numbers (in new_citation_mapping) to new numbers\n old_to_new: dict[int, int] = {}\n additional_mappings: CitationMapping = {}\n\n for old_num, search_doc in new_citation_mapping.items():\n doc_id = search_doc.document_id\n\n # Check if this document already exists in existing citations\n if doc_id in doc_id_to_existing_citation:\n # Use the existing citation number\n old_to_new[old_num] = doc_id_to_existing_citation[doc_id]\n else:\n # Check if we've already assigned a new number to this document\n # (handles case where same doc appears with different old numbers)\n existing_new_num = None\n for mapped_old, mapped_new in old_to_new.items():\n if (\n mapped_old in new_citation_mapping\n and new_citation_mapping[mapped_old].document_id == doc_id\n ):\n existing_new_num = mapped_new\n break\n\n if existing_new_num is not None:\n old_to_new[old_num] = existing_new_num\n else:\n # Assign the next available number\n old_to_new[old_num] = next_citation_num\n additional_mappings[next_citation_num] = search_doc\n next_citation_num += 1\n\n # Pattern to match citations like [25], [1, 2, 3], [[25]], etc.\n # Also matches unicode bracket variants: 【】, ［］\n citation_pattern = re.compile(\n r\"([\\[【［]{2}\\d+[\\]】］]{2})|([\\[【［]\\d+(?:, ?\\d+)*[\\]】］])\"\n )\n\n def replace_citation(match: re.Match) -> str:\n \"\"\"Replace citation numbers in a match with their new collapsed values.\"\"\"\n citation_str = match.group()\n\n # Determine bracket style\n if (\n citation_str.startswith(\"[[\")\n or citation_str.startswith(\"【【\")\n or citation_str.startswith(\"［［\")\n ):\n open_bracket = citation_str[:2]\n close_bracket = citation_str[-2:]\n content = citation_str[2:-2]\n else:\n open_bracket = citation_str[0]\n close_bracket = citation_str[-1]\n content = citation_str[1:-1]\n\n # Parse and replace citation numbers\n new_nums = []\n for num_str in content.split(\",\"):\n num_str = num_str.strip()\n if not num_str:\n continue\n try:\n num = int(num_str)\n # Only replace if we have a mapping for this number\n if num in old_to_new:\n new_nums.append(str(old_to_new[num]))\n else:\n # Keep original if not in our mapping\n new_nums.append(num_str)\n except ValueError:\n new_nums.append(num_str)\n\n # Reconstruct the citation with original bracket style\n new_content = \", \".join(new_nums)\n return f\"{open_bracket}{new_content}{close_bracket}\"\n\n # Replace all citations in the text\n updated_text = citation_pattern.sub(replace_citation, answer_text)\n\n # Build the combined mapping\n combined_mapping: CitationMapping = dict(existing_citation_mapping)\n combined_mapping.update(additional_mappings)\n\n return updated_text, combined_mapping\n" }, { "path": "backend/onyx/chat/compression.py", "content": "\"\"\"\nChat history compression via summarization.\n\nThis module handles compressing long chat histories by summarizing older messages\nwhile keeping recent messages verbatim.\n\nSummaries are branch-aware: each summary's parent_message_id points to the last\nmessage when compression triggered, making it part of the tree structure.\n\"\"\"\n\nfrom typing import NamedTuple\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.chat_configs import COMPRESSION_TRIGGER_RATIO\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.models import ChatMessage\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import AssistantMessage\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import SystemMessage\nfrom onyx.llm.models import UserMessage\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.prompts.compression_prompts import PROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK\nfrom onyx.prompts.compression_prompts import PROGRESSIVE_USER_REMINDER\nfrom onyx.prompts.compression_prompts import SUMMARIZATION_CUTOFF_MARKER\nfrom onyx.prompts.compression_prompts import SUMMARIZATION_PROMPT\nfrom onyx.prompts.compression_prompts import USER_REMINDER\nfrom onyx.tracing.framework.create import ensure_trace\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n# Ratio of available context to allocate for recent messages after compression\nRECENT_MESSAGES_RATIO = 0.2\n\n\nclass CompressionResult(BaseModel):\n \"\"\"Result of a compression operation.\"\"\"\n\n summary_created: bool\n messages_summarized: int\n error: str | None = None\n\n\nclass CompressionParams(BaseModel):\n \"\"\"Parameters for compression operation.\"\"\"\n\n should_compress: bool\n tokens_for_recent: int = 0\n\n\nclass SummaryContent(NamedTuple):\n \"\"\"Messages split for summarization.\"\"\"\n\n older_messages: list[ChatMessage]\n recent_messages: list[ChatMessage]\n\n\ndef calculate_total_history_tokens(chat_history: list[ChatMessage]) -> int:\n \"\"\"\n Calculate the total token count for the given chat history.\n\n Args:\n chat_history: Branch-aware list of messages\n\n Returns:\n Total token count for the history\n \"\"\"\n return sum(m.token_count or 0 for m in chat_history)\n\n\ndef get_compression_params(\n max_input_tokens: int,\n current_history_tokens: int,\n reserved_tokens: int,\n) -> CompressionParams:\n \"\"\"\n Calculate compression parameters based on model's context window.\n\n Args:\n max_input_tokens: The maximum input tokens for the LLM\n current_history_tokens: Current total tokens in chat history\n reserved_tokens: Tokens reserved for system prompt, tools, files, etc.\n\n Returns:\n CompressionParams indicating whether to compress and token budgets\n \"\"\"\n available = max_input_tokens - reserved_tokens\n\n # Check trigger threshold\n trigger_threshold = int(available * COMPRESSION_TRIGGER_RATIO)\n\n if current_history_tokens <= trigger_threshold:\n return CompressionParams(should_compress=False)\n\n # Calculate token budget for recent messages as a percentage of current history\n # This ensures we always have messages to summarize when compression triggers\n tokens_for_recent = int(current_history_tokens * RECENT_MESSAGES_RATIO)\n\n return CompressionParams(\n should_compress=True,\n tokens_for_recent=tokens_for_recent,\n )\n\n\ndef find_summary_for_branch(\n db_session: Session,\n chat_history: list[ChatMessage],\n) -> ChatMessage | None:\n \"\"\"\n Find the most recent summary that applies to the current branch.\n\n A summary applies if its parent_message_id is in the current chat history,\n meaning it was created on this branch.\n\n Args:\n db_session: Database session\n chat_history: Branch-aware list of messages\n\n Returns:\n The applicable summary message, or None if no summary exists for this branch\n \"\"\"\n if not chat_history:\n return None\n\n history_ids = {m.id for m in chat_history}\n chat_session_id = chat_history[0].chat_session_id\n\n # Query all summaries for this session (typically few), then filter in Python.\n # Order by time_sent descending to get the most recent summary first.\n summaries = (\n db_session.query(ChatMessage)\n .filter(\n ChatMessage.chat_session_id == chat_session_id,\n ChatMessage.last_summarized_message_id.isnot(None),\n )\n .order_by(ChatMessage.time_sent.desc())\n .all()\n )\n # Optimization to avoid using IN clause for large histories\n for summary in summaries:\n if summary.parent_message_id in history_ids:\n return summary\n\n return None\n\n\ndef get_messages_to_summarize(\n chat_history: list[ChatMessage],\n existing_summary: ChatMessage | None,\n tokens_for_recent: int,\n) -> SummaryContent:\n \"\"\"\n Split messages into those to summarize and those to keep verbatim.\n\n Args:\n chat_history: Branch-aware list of messages\n existing_summary: Existing summary for this branch (if any)\n tokens_for_recent: Token budget for recent messages to keep\n\n Returns:\n SummaryContent with older_messages to summarize and recent_messages to keep\n \"\"\"\n # Filter to messages after the existing summary's cutoff using timestamp\n if existing_summary and existing_summary.last_summarized_message_id:\n cutoff_id = existing_summary.last_summarized_message_id\n last_summarized_msg = next(m for m in chat_history if m.id == cutoff_id)\n messages = [\n m for m in chat_history if m.time_sent > last_summarized_msg.time_sent\n ]\n else:\n messages = list(chat_history)\n\n # Filter out empty messages\n messages = [m for m in messages if m.message]\n\n if not messages:\n return SummaryContent(older_messages=[], recent_messages=[])\n\n # Work backwards from most recent, keeping messages until we exceed budget\n recent_messages: list[ChatMessage] = []\n tokens_used = 0\n\n for msg in reversed(messages):\n msg_tokens = msg.token_count or 0\n if tokens_used + msg_tokens > tokens_for_recent and recent_messages:\n break\n recent_messages.insert(0, msg)\n tokens_used += msg_tokens\n\n # Ensure cutoff is right before a user message by moving any leading\n # non-user messages from recent_messages to older_messages\n while recent_messages and recent_messages[0].message_type != MessageType.USER:\n recent_messages.pop(0)\n\n # Everything else gets summarized\n recent_ids = {m.id for m in recent_messages}\n older_messages = [m for m in messages if m.id not in recent_ids]\n\n return SummaryContent(\n older_messages=older_messages, recent_messages=recent_messages\n )\n\n\ndef _build_llm_messages_for_summarization(\n messages: list[ChatMessage],\n tool_id_to_name: dict[int, str],\n) -> list[UserMessage | AssistantMessage]:\n \"\"\"Convert ChatMessage objects to LLM message format for summarization.\n\n This is intentionally different from translate_history_to_llm_format in llm_step.py:\n - Compacts tool calls to \"[Used tools: tool1, tool2]\" to save tokens in summaries\n - Skips TOOL_CALL_RESPONSE messages entirely (tool usage captured in assistant message)\n - No image/multimodal handling (summaries are text-only)\n - No caching or LLMConfig-specific behavior needed\n \"\"\"\n result: list[UserMessage | AssistantMessage] = []\n\n for msg in messages:\n # Skip empty messages\n if not msg.message:\n continue\n\n # Handle assistant messages with tool calls compactly\n if msg.message_type == MessageType.ASSISTANT:\n if msg.tool_calls:\n tool_names = [\n tool_id_to_name.get(tc.tool_id, \"unknown\") for tc in msg.tool_calls\n ]\n result.append(\n AssistantMessage(content=f\"[Used tools: {', '.join(tool_names)}]\")\n )\n else:\n result.append(AssistantMessage(content=msg.message))\n continue\n\n # Skip tool call response messages - tool calls are captured above via assistant messages\n if msg.message_type == MessageType.TOOL_CALL_RESPONSE:\n continue\n\n # Handle user messages\n if msg.message_type == MessageType.USER:\n result.append(UserMessage(content=msg.message))\n\n return result\n\n\ndef generate_summary(\n older_messages: list[ChatMessage],\n recent_messages: list[ChatMessage],\n llm: LLM,\n tool_id_to_name: dict[int, str],\n existing_summary: str | None = None,\n) -> str:\n \"\"\"\n Generate a summary using cutoff marker approach.\n\n The cutoff marker tells the LLM to summarize only older messages,\n while using recent messages as context to inform what's important.\n\n Messages are sent as separate UserMessage/AssistantMessage objects rather\n than being concatenated into a single message.\n\n Args:\n older_messages: Messages to compress into summary (before cutoff)\n recent_messages: Messages kept verbatim (after cutoff, for context only)\n llm: LLM to use for summarization\n tool_id_to_name: Mapping of tool IDs to display names\n existing_summary: Previous summary text to incorporate (progressive)\n\n Returns:\n Summary text\n \"\"\"\n # Build system prompt\n system_content = SUMMARIZATION_PROMPT\n if existing_summary:\n # Progressive summarization: append existing summary to system prompt\n system_content += PROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK.format(\n previous_summary=existing_summary\n )\n final_reminder = PROGRESSIVE_USER_REMINDER\n else:\n final_reminder = USER_REMINDER\n\n # Convert messages to LLM format (using compression-specific conversion)\n older_llm_messages = _build_llm_messages_for_summarization(\n older_messages, tool_id_to_name\n )\n recent_llm_messages = _build_llm_messages_for_summarization(\n recent_messages, tool_id_to_name\n )\n\n # Build message list with separate messages\n input_messages: list[ChatCompletionMessage] = [\n SystemMessage(content=system_content),\n ]\n\n # Add older messages (to be summarized)\n input_messages.extend(older_llm_messages)\n\n # Add cutoff marker as a user message\n input_messages.append(UserMessage(content=SUMMARIZATION_CUTOFF_MARKER))\n\n # Add recent messages (for context only)\n input_messages.extend(recent_llm_messages)\n\n # Add final reminder\n input_messages.append(UserMessage(content=final_reminder))\n\n with llm_generation_span(\n llm=llm,\n flow=\"chat_history_summarization\",\n input_messages=input_messages,\n ) as span_generation:\n response = llm.invoke(input_messages)\n record_llm_response(span_generation, response)\n\n content = response.choice.message.content\n if not (content and content.strip()):\n raise ValueError(\"LLM returned empty summary\")\n return content.strip()\n\n\ndef compress_chat_history(\n db_session: Session,\n chat_history: list[ChatMessage],\n llm: LLM,\n compression_params: CompressionParams,\n tool_id_to_name: dict[int, str],\n) -> CompressionResult:\n \"\"\"\n Main compression function. Creates a summary ChatMessage.\n\n The summary message's parent_message_id points to the last message in\n chat_history, making it branch-aware via the tree structure.\n\n Note: This takes the entire chat history as input, splits it into older\n messages (to summarize) and recent messages (kept verbatim within the\n token budget), generates a summary of the older part, and persists the\n new summary message with its parent set to the last message in history.\n\n Past summary is taken into context (progressive summarization): we find\n at most one existing summary for this branch. If present, only messages\n after that summary's last_summarized_message_id are considered; the\n existing summary text is passed into the LLM so the new summary\n incorporates it instead of summarizing from scratch.\n\n For more details, see the COMPRESSION.md file.\n\n Args:\n db_session: Database session\n chat_history: Branch-aware list of messages\n llm: LLM to use for summarization\n compression_params: Parameters from get_compression_params\n tool_id_to_name: Mapping of tool IDs to display names\n\n Returns:\n CompressionResult indicating success/failure\n \"\"\"\n if not chat_history:\n return CompressionResult(summary_created=False, messages_summarized=0)\n\n chat_session_id = chat_history[0].chat_session_id\n\n logger.info(\n f\"Starting compression for session {chat_session_id}, \"\n f\"history_len={len(chat_history)}, tokens_for_recent={compression_params.tokens_for_recent}\"\n )\n\n with ensure_trace(\n \"chat_history_compression\",\n group_id=str(chat_session_id),\n metadata={\n \"tenant_id\": get_current_tenant_id(),\n \"chat_session_id\": str(chat_session_id),\n },\n ):\n try:\n # Find existing summary for this branch\n existing_summary = find_summary_for_branch(db_session, chat_history)\n\n # Get messages to summarize\n summary_content = get_messages_to_summarize(\n chat_history,\n existing_summary,\n tokens_for_recent=compression_params.tokens_for_recent,\n )\n\n if not summary_content.older_messages:\n logger.debug(\"No messages to summarize, skipping compression\")\n return CompressionResult(summary_created=False, messages_summarized=0)\n\n # Generate summary (incorporate existing summary if present)\n existing_summary_text = (\n existing_summary.message if existing_summary else None\n )\n summary_text = generate_summary(\n older_messages=summary_content.older_messages,\n recent_messages=summary_content.recent_messages,\n llm=llm,\n tool_id_to_name=tool_id_to_name,\n existing_summary=existing_summary_text,\n )\n\n # Calculate token count for the summary\n tokenizer = get_tokenizer(None, None)\n summary_token_count = len(tokenizer.encode(summary_text))\n logger.debug(\n f\"Generated summary ({summary_token_count} tokens): {summary_text[:200]}...\"\n )\n\n # Create new summary as a ChatMessage\n # Parent is the last message in history - this makes the summary branch-aware\n summary_message = ChatMessage(\n chat_session_id=chat_session_id,\n message_type=MessageType.ASSISTANT,\n message=summary_text,\n token_count=summary_token_count,\n parent_message_id=chat_history[-1].id,\n last_summarized_message_id=summary_content.older_messages[-1].id,\n )\n db_session.add(summary_message)\n db_session.commit()\n\n logger.info(\n f\"Compressed {len(summary_content.older_messages)} messages into summary \"\n f\"(session_id={chat_session_id}, \"\n f\"summary_tokens={summary_token_count})\"\n )\n\n return CompressionResult(\n summary_created=True,\n messages_summarized=len(summary_content.older_messages),\n )\n\n except Exception as e:\n logger.exception(f\"Compression failed for session {chat_session_id}: {e}\")\n db_session.rollback()\n return CompressionResult(\n summary_created=False,\n messages_summarized=0,\n error=str(e),\n )\n" }, { "path": "backend/onyx/chat/emitter.py", "content": "import threading\nfrom queue import Queue\n\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import Packet\n\n\nclass Emitter:\n \"\"\"Routes packets from LLM/tool execution to the ``_run_models`` drain loop.\n\n Tags every packet with ``model_index`` and places it on ``merged_queue``\n as a ``(model_idx, packet)`` tuple for ordered consumption downstream.\n\n Args:\n merged_queue: Shared queue owned by ``_run_models``.\n model_idx: Index embedded in packet placements (``0`` for N=1 runs).\n drain_done: Optional event set by ``_run_models`` when the drain loop\n exits early (e.g. HTTP disconnect). When set, ``emit`` returns\n immediately so worker threads can exit fast.\n \"\"\"\n\n def __init__(\n self,\n merged_queue: Queue[tuple[int, Packet | Exception | object]],\n model_idx: int = 0,\n drain_done: threading.Event | None = None,\n ) -> None:\n self._model_idx = model_idx\n self._merged_queue = merged_queue\n self._drain_done = drain_done\n\n def emit(self, packet: Packet) -> None:\n if self._drain_done is not None and self._drain_done.is_set():\n return\n base = packet.placement or Placement(turn_index=0)\n tagged = Packet(\n placement=base.model_copy(update={\"model_index\": self._model_idx}),\n obj=packet.obj,\n )\n self._merged_queue.put((self._model_idx, tagged))\n" }, { "path": "backend/onyx/chat/llm_loop.py", "content": "import json\nimport time\nfrom collections.abc import Callable\nfrom typing import Any\nfrom typing import Literal\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.chat_state import ChatStateContainer\nfrom onyx.chat.chat_utils import create_tool_call_failure_messages\nfrom onyx.chat.citation_processor import CitationMapping\nfrom onyx.chat.citation_processor import CitationMode\nfrom onyx.chat.citation_processor import DynamicCitationProcessor\nfrom onyx.chat.citation_utils import update_citation_processor_from_tool_response\nfrom onyx.chat.emitter import Emitter\nfrom onyx.chat.llm_step import extract_tool_calls_from_response_text\nfrom onyx.chat.llm_step import run_llm_step\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.chat.models import ContextFileMetadata\nfrom onyx.chat.models import ExtractedContextFiles\nfrom onyx.chat.models import FileToolMetadata\nfrom onyx.chat.models import LlmStepResult\nfrom onyx.chat.models import ToolCallSimple\nfrom onyx.chat.prompt_utils import build_reminder_message\nfrom onyx.chat.prompt_utils import build_system_prompt\nfrom onyx.chat.prompt_utils import (\n get_default_base_system_prompt,\n)\nfrom onyx.configs.app_configs import INTEGRATION_TESTS_MODE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import MessageType\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.context.search.models import SearchDocsResponse\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.memory import add_memory\nfrom onyx.db.memory import update_memory_at_index\nfrom onyx.db.memory import UserMemoryContext\nfrom onyx.db.models import Persona\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.interfaces import ToolChoiceOptions\nfrom onyx.llm.utils import is_true_openai_model\nfrom onyx.prompts.chat_prompts import IMAGE_GEN_REMINDER\nfrom onyx.prompts.chat_prompts import OPEN_URL_REMINDER\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import OverallStop\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import ToolCallDebug\nfrom onyx.server.query_and_chat.streaming_models import TopLevelBranching\nfrom onyx.tools.built_in_tools import CITEABLE_TOOLS_NAMES\nfrom onyx.tools.built_in_tools import STOPPING_TOOLS_NAMES\nfrom onyx.tools.interface import Tool\nfrom onyx.tools.models import ChatFile\nfrom onyx.tools.models import CustomToolCallSummary\nfrom onyx.tools.models import MemoryToolResponseSnapshot\nfrom onyx.tools.models import PythonToolRichResponse\nfrom onyx.tools.models import ToolCallInfo\nfrom onyx.tools.models import ToolCallKickoff\nfrom onyx.tools.models import ToolResponse\nfrom onyx.tools.tool_implementations.images.models import (\n FinalImageGenerationResponse,\n)\nfrom onyx.tools.tool_implementations.memory.models import MemoryToolResponse\nfrom onyx.tools.tool_implementations.python.python_tool import PythonTool\nfrom onyx.tools.tool_implementations.search.search_tool import SearchTool\nfrom onyx.tools.tool_implementations.web_search.utils import extract_url_snippet_map\nfrom onyx.tools.tool_implementations.web_search.web_search_tool import WebSearchTool\nfrom onyx.tools.tool_runner import run_tool_calls\nfrom onyx.tracing.framework.create import trace\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\nclass EmptyLLMResponseError(RuntimeError):\n \"\"\"Raised when the streamed LLM response completes without a usable answer.\"\"\"\n\n def __init__(\n self,\n *,\n provider: str,\n model: str,\n tool_choice: ToolChoiceOptions,\n client_error_msg: str,\n error_code: str = \"EMPTY_LLM_RESPONSE\",\n is_retryable: bool = True,\n ) -> None:\n super().__init__(client_error_msg)\n self.provider = provider\n self.model = model\n self.tool_choice = tool_choice\n self.client_error_msg = client_error_msg\n self.error_code = error_code\n self.is_retryable = is_retryable\n\n\ndef _build_empty_llm_response_error(\n llm: LLM,\n llm_step_result: LlmStepResult,\n tool_choice: ToolChoiceOptions,\n) -> EmptyLLMResponseError:\n provider = llm.config.model_provider\n model = llm.config.model_name\n\n # OpenAI quota exhaustion has reached us as a streamed \"stop\" with zero content.\n # When the stream is completely empty and there is no reasoning/tool output, surface\n # the likely account-level cause instead of a generic tool-calling error.\n if (\n not llm_step_result.reasoning\n and provider == LlmProviderNames.OPENAI\n and is_true_openai_model(provider, model)\n ):\n return EmptyLLMResponseError(\n provider=provider,\n model=model,\n tool_choice=tool_choice,\n client_error_msg=(\n \"The selected OpenAI model returned an empty streamed response \"\n \"before producing any tokens. This commonly happens when the API \"\n \"key or project has no remaining quota or billing is not enabled. \"\n \"Verify quota and billing for this key and try again.\"\n ),\n error_code=\"BUDGET_EXCEEDED\",\n is_retryable=False,\n )\n\n return EmptyLLMResponseError(\n provider=provider,\n model=model,\n tool_choice=tool_choice,\n client_error_msg=(\n \"The selected model returned no final answer before the stream \"\n \"completed. No text or tool calls were received from the upstream \"\n \"provider.\"\n ),\n )\n\n\ndef _looks_like_xml_tool_call_payload(text: str | None) -> bool:\n \"\"\"Detect XML-style marshaled tool calls emitted as plain text.\"\"\"\n if not text:\n return False\n lowered = text.lower()\n return (\n \" tuple[LlmStepResult, bool]:\n \"\"\"Attempt to extract tool calls from response text as a fallback.\n\n This is a last resort fallback for low quality LLMs or those that don't have\n tool calling from the serving layer. Also triggers if there's reasoning but\n no answer and no tool calls.\n\n Args:\n llm_step_result: The result from the LLM step\n tool_choice: The tool choice option used for this step\n fallback_extraction_attempted: Whether fallback extraction was already attempted\n tool_defs: List of tool definitions\n turn_index: The current turn index for placement\n\n Returns:\n Tuple of (possibly updated LlmStepResult, whether fallback was attempted this call)\n \"\"\"\n if fallback_extraction_attempted:\n return llm_step_result, False\n\n no_tool_calls = (\n not llm_step_result.tool_calls or len(llm_step_result.tool_calls) == 0\n )\n reasoning_but_no_answer_or_tools = (\n llm_step_result.reasoning and not llm_step_result.answer and no_tool_calls\n )\n xml_tool_call_text_detected = no_tool_calls and (\n _looks_like_xml_tool_call_payload(llm_step_result.answer)\n or _looks_like_xml_tool_call_payload(llm_step_result.raw_answer)\n or _looks_like_xml_tool_call_payload(llm_step_result.reasoning)\n )\n should_try_fallback = (\n (tool_choice == ToolChoiceOptions.REQUIRED and no_tool_calls)\n or reasoning_but_no_answer_or_tools\n or xml_tool_call_text_detected\n )\n\n if not should_try_fallback:\n return llm_step_result, False\n\n # Try to extract from answer first, then fall back to reasoning\n extracted_tool_calls: list[ToolCallKickoff] = []\n\n if llm_step_result.answer:\n extracted_tool_calls = extract_tool_calls_from_response_text(\n response_text=llm_step_result.answer,\n tool_definitions=tool_defs,\n placement=Placement(turn_index=turn_index),\n )\n if (\n not extracted_tool_calls\n and llm_step_result.raw_answer\n and llm_step_result.raw_answer != llm_step_result.answer\n ):\n extracted_tool_calls = extract_tool_calls_from_response_text(\n response_text=llm_step_result.raw_answer,\n tool_definitions=tool_defs,\n placement=Placement(turn_index=turn_index),\n )\n if not extracted_tool_calls and llm_step_result.reasoning:\n extracted_tool_calls = extract_tool_calls_from_response_text(\n response_text=llm_step_result.reasoning,\n tool_definitions=tool_defs,\n placement=Placement(turn_index=turn_index),\n )\n if extracted_tool_calls:\n logger.info(\n f\"Extracted {len(extracted_tool_calls)} tool call(s) from response text as fallback\"\n )\n return (\n LlmStepResult(\n reasoning=llm_step_result.reasoning,\n answer=llm_step_result.answer,\n tool_calls=extracted_tool_calls,\n raw_answer=llm_step_result.raw_answer,\n ),\n True,\n )\n\n return llm_step_result, True\n\n\n# Hardcoded oppinionated value, might breaks down to something like:\n# Cycle 1: Calls web_search for something\n# Cycle 2: Calls open_url for some results\n# Cycle 3: Calls web_search for some other aspect of the question\n# Cycle 4: Calls open_url for some results\n# Cycle 5: Maybe call open_url for some additional results or because last set failed\n# Cycle 6: No more tools available, forced to answer\nMAX_LLM_CYCLES = 6\n\n\ndef _build_context_file_citation_mapping(\n file_metadata: list[ContextFileMetadata],\n starting_citation_num: int = 1,\n) -> CitationMapping:\n \"\"\"Build citation mapping for context files.\n\n Converts context file metadata into SearchDoc objects that can be cited.\n Citation numbers start from the provided starting number.\n\n Args:\n file_metadata: List of context file metadata\n starting_citation_num: Starting citation number (default: 1)\n\n Returns:\n Dictionary mapping citation numbers to SearchDoc objects\n \"\"\"\n citation_mapping: CitationMapping = {}\n\n for idx, file_meta in enumerate(file_metadata, start=starting_citation_num):\n search_doc = SearchDoc(\n document_id=file_meta.file_id,\n chunk_ind=0,\n semantic_identifier=file_meta.filename,\n link=None,\n blurb=file_meta.file_content,\n source_type=DocumentSource.FILE,\n boost=1,\n hidden=False,\n metadata={},\n score=0.0,\n match_highlights=[file_meta.file_content],\n )\n citation_mapping[idx] = search_doc\n\n return citation_mapping\n\n\ndef _build_project_message(\n context_files: ExtractedContextFiles | None,\n token_counter: Callable[[str], int] | None,\n) -> list[ChatMessageSimple]:\n \"\"\"Build messages for context-injected / tool-backed files.\n\n Returns up to two messages:\n 1. The full-text files message (if file_texts is populated).\n 2. A lightweight metadata message for files the LLM should access via the\n FileReaderTool (e.g. oversized files that don't fit in context).\n \"\"\"\n if not context_files:\n return []\n\n messages: list[ChatMessageSimple] = []\n if context_files.file_texts:\n messages.append(\n _create_context_files_message(context_files, token_counter=None)\n )\n if context_files.file_metadata_for_tool and token_counter:\n messages.append(\n _create_file_tool_metadata_message(\n context_files.file_metadata_for_tool, token_counter\n )\n )\n return messages\n\n\ndef construct_message_history(\n system_prompt: ChatMessageSimple | None,\n custom_agent_prompt: ChatMessageSimple | None,\n simple_chat_history: list[ChatMessageSimple],\n reminder_message: ChatMessageSimple | None,\n context_files: ExtractedContextFiles | None,\n available_tokens: int,\n last_n_user_messages: int | None = None,\n token_counter: Callable[[str], int] | None = None,\n all_injected_file_metadata: dict[str, FileToolMetadata] | None = None,\n) -> list[ChatMessageSimple]:\n if last_n_user_messages is not None:\n if last_n_user_messages <= 0:\n raise ValueError(\n \"filtering chat history by last N user messages must be a value greater than 0\"\n )\n\n # Build the project / file-metadata messages up front so we can use their\n # actual token counts for the budget.\n project_messages = _build_project_message(context_files, token_counter)\n project_messages_tokens = sum(m.token_count for m in project_messages)\n\n history_token_budget = available_tokens\n history_token_budget -= system_prompt.token_count if system_prompt else 0\n history_token_budget -= (\n custom_agent_prompt.token_count if custom_agent_prompt else 0\n )\n history_token_budget -= project_messages_tokens\n history_token_budget -= reminder_message.token_count if reminder_message else 0\n\n if history_token_budget < 0:\n raise ValueError(\"Not enough tokens available to construct message history\")\n\n if system_prompt:\n system_prompt.should_cache = True\n\n # If no history, build minimal context\n if not simple_chat_history:\n result = [system_prompt] if system_prompt else []\n if custom_agent_prompt:\n result.append(custom_agent_prompt)\n result.extend(project_messages)\n if reminder_message:\n result.append(reminder_message)\n return result\n\n # If last_n_user_messages is set, filter history to only include the last n user messages\n if last_n_user_messages is not None:\n # Find all user message indices\n user_msg_indices = [\n i\n for i, msg in enumerate(simple_chat_history)\n if msg.message_type == MessageType.USER\n ]\n\n if not user_msg_indices:\n raise ValueError(\"No user message found in simple_chat_history\")\n\n # If we have more than n user messages, keep only the last n\n if len(user_msg_indices) > last_n_user_messages:\n # Find the index of the n-th user message from the end\n # For example, if last_n_user_messages=2, we want the 2nd-to-last user message\n nth_user_msg_index = user_msg_indices[-(last_n_user_messages)]\n # Keep everything from that user message onwards\n simple_chat_history = simple_chat_history[nth_user_msg_index:]\n\n # Find the last USER message in the history\n # The history may contain tool calls and responses after the last user message\n last_user_msg_index = None\n for i in range(len(simple_chat_history) - 1, -1, -1):\n if simple_chat_history[i].message_type == MessageType.USER:\n last_user_msg_index = i\n break\n\n if last_user_msg_index is None:\n raise ValueError(\"No user message found in simple_chat_history\")\n\n # Split history into three parts:\n # 1. History before the last user message\n # 2. The last user message\n # 3. Messages after the last user message (tool calls, responses, etc.)\n history_before_last_user = simple_chat_history[:last_user_msg_index]\n last_user_message = simple_chat_history[last_user_msg_index]\n messages_after_last_user = simple_chat_history[last_user_msg_index + 1 :]\n\n # Calculate tokens needed for the last user message and everything after it\n last_user_tokens = last_user_message.token_count\n after_user_tokens = sum(msg.token_count for msg in messages_after_last_user)\n\n # Check if we can fit at least the last user message and messages after it\n required_tokens = last_user_tokens + after_user_tokens\n if required_tokens > history_token_budget:\n raise ValueError(\n f\"Not enough tokens to include the last user message and subsequent messages. \"\n f\"Required: {required_tokens}, Available: {history_token_budget}\"\n )\n\n # Calculate remaining budget for history before the last user message\n remaining_budget = history_token_budget - required_tokens\n\n # Truncate history_before_last_user from the top to fit in remaining budget.\n # Track dropped file messages so we can provide their metadata to the\n # FileReaderTool instead.\n truncated_history_before: list[ChatMessageSimple] = []\n dropped_file_ids: list[str] = []\n current_token_count = 0\n\n for msg in reversed(history_before_last_user):\n if current_token_count + msg.token_count <= remaining_budget:\n msg.should_cache = True\n truncated_history_before.insert(0, msg)\n current_token_count += msg.token_count\n else:\n # Can't fit this message, stop truncating.\n # This message and everything older is dropped.\n break\n\n # Collect file_ids from ALL dropped messages (those not in\n # truncated_history_before). The truncation loop above keeps the most\n # recent messages, so the dropped ones are at the start of the original\n # list up to (len(history) - len(kept)).\n num_kept = len(truncated_history_before)\n for msg in history_before_last_user[: len(history_before_last_user) - num_kept]:\n if msg.file_id is not None:\n dropped_file_ids.append(msg.file_id)\n\n # Also treat \"orphaned\" metadata entries as dropped -- these are files\n # from messages removed by summary truncation (before convert_chat_history\n # ran), so no ChatMessageSimple was ever tagged with their file_id.\n if all_injected_file_metadata:\n surviving_file_ids = {\n msg.file_id for msg in simple_chat_history if msg.file_id is not None\n }\n for fid in all_injected_file_metadata:\n if fid not in surviving_file_ids and fid not in dropped_file_ids:\n dropped_file_ids.append(fid)\n\n # Build a forgotten-files metadata message if any file messages were\n # dropped AND we have metadata for them (meaning the FileReaderTool is\n # available). Reserve tokens for this message in the budget.\n forgotten_files_message: ChatMessageSimple | None = None\n if dropped_file_ids and all_injected_file_metadata and token_counter:\n forgotten_meta = [\n all_injected_file_metadata[fid]\n for fid in dropped_file_ids\n if fid in all_injected_file_metadata\n ]\n if forgotten_meta:\n logger.debug(\n f\"FileReader: building forgotten-files message for {[(m.file_id, m.filename) for m in forgotten_meta]}\"\n )\n forgotten_files_message = _create_file_tool_metadata_message(\n forgotten_meta, token_counter\n )\n # Shrink the remaining budget. If the metadata message doesn't\n # fit we may need to drop more history messages.\n remaining_budget -= forgotten_files_message.token_count\n while truncated_history_before and current_token_count > remaining_budget:\n evicted = truncated_history_before.pop(0)\n current_token_count -= evicted.token_count\n # If the evicted message is itself a file, add it to the\n # forgotten metadata (it's now dropped too).\n if (\n evicted.file_id is not None\n and evicted.file_id in all_injected_file_metadata\n and evicted.file_id not in {m.file_id for m in forgotten_meta}\n ):\n forgotten_meta.append(all_injected_file_metadata[evicted.file_id])\n # Rebuild the message with the new entry\n forgotten_files_message = _create_file_tool_metadata_message(\n forgotten_meta, token_counter\n )\n\n # Attach project images to the last user message\n if context_files and context_files.image_files:\n existing_images = last_user_message.image_files or []\n last_user_message = ChatMessageSimple(\n message=last_user_message.message,\n token_count=last_user_message.token_count,\n message_type=last_user_message.message_type,\n image_files=existing_images + context_files.image_files,\n )\n\n # Build the final message list according to README ordering:\n # [system], [history_before_last_user], [custom_agent], [context_files],\n # [forgotten_files], [last_user_message], [messages_after_last_user], [reminder]\n result = [system_prompt] if system_prompt else []\n\n # 1. Add truncated history before last user message\n result.extend(truncated_history_before)\n\n # 2. Add custom agent prompt (inserted before last user message)\n if custom_agent_prompt:\n result.append(custom_agent_prompt)\n\n # 3. Add context files / file-metadata messages (inserted before last user message)\n result.extend(project_messages)\n\n # 4. Add forgotten-files metadata (right before the user's question)\n if forgotten_files_message:\n result.append(forgotten_files_message)\n\n # 5. Add last user message (with context images attached)\n result.append(last_user_message)\n\n # 6. Add messages after last user message (tool calls, responses, etc.)\n result.extend(messages_after_last_user)\n\n # 7. Add reminder message at the very end\n if reminder_message:\n result.append(reminder_message)\n\n return _drop_orphaned_tool_call_responses(result)\n\n\ndef _drop_orphaned_tool_call_responses(\n messages: list[ChatMessageSimple],\n) -> list[ChatMessageSimple]:\n \"\"\"Drop tool response messages whose tool_call_id is not in prior assistant tool calls.\n\n This can happen when history truncation drops an ASSISTANT tool-call message but\n leaves a later TOOL_CALL_RESPONSE message in context. Some providers (e.g. Ollama)\n reject such history with an \"unexpected tool call id\" error.\n \"\"\"\n known_tool_call_ids: set[str] = set()\n sanitized: list[ChatMessageSimple] = []\n\n for msg in messages:\n if msg.message_type == MessageType.ASSISTANT and msg.tool_calls:\n for tool_call in msg.tool_calls:\n known_tool_call_ids.add(tool_call.tool_call_id)\n sanitized.append(msg)\n continue\n\n if msg.message_type == MessageType.TOOL_CALL_RESPONSE:\n if msg.tool_call_id and msg.tool_call_id in known_tool_call_ids:\n sanitized.append(msg)\n else:\n logger.debug(\n \"Dropping orphaned tool response with tool_call_id=%s while constructing message history\",\n msg.tool_call_id,\n )\n continue\n\n sanitized.append(msg)\n\n return sanitized\n\n\ndef _create_file_tool_metadata_message(\n file_metadata: list[FileToolMetadata],\n token_counter: Callable[[str], int],\n) -> ChatMessageSimple:\n \"\"\"Build a lightweight metadata-only message listing files available via FileReaderTool.\n\n Used when files are too large to fit in context and the vector DB is\n disabled, so the LLM must use ``read_file`` to inspect them.\n \"\"\"\n lines = [\n \"You have access to the following files. Use the read_file tool to \"\n \"read sections of any file. You MUST pass the file_id UUID (not the \"\n \"filename) to read_file:\"\n ]\n for meta in file_metadata:\n lines.append(\n f'- file_id=\"{meta.file_id}\" filename=\"{meta.filename}\" (~{meta.approx_char_count:,} chars)'\n )\n\n message_content = \"\\n\".join(lines)\n return ChatMessageSimple(\n message=message_content,\n token_count=token_counter(message_content),\n message_type=MessageType.USER,\n )\n\n\ndef _create_context_files_message(\n context_files: ExtractedContextFiles,\n token_counter: Callable[[str], int] | None, # noqa: ARG001\n) -> ChatMessageSimple:\n \"\"\"Convert context files to a ChatMessageSimple message.\n\n Format follows the README specification for document representation.\n \"\"\"\n import json\n\n # Format as documents JSON as described in README\n documents_list = []\n for idx, file_text in enumerate(context_files.file_texts, start=1):\n title = (\n context_files.file_metadata[idx - 1].filename\n if idx - 1 < len(context_files.file_metadata)\n else None\n )\n entry: dict[str, Any] = {\"document\": idx}\n if title:\n entry[\"title\"] = title\n entry[\"contents\"] = file_text\n documents_list.append(entry)\n\n documents_json = json.dumps({\"documents\": documents_list}, indent=2)\n message_content = f\"Here are some documents provided for context, they may not all be relevant:\\n{documents_json}\"\n\n # Use pre-calculated token count from context_files\n return ChatMessageSimple(\n message=message_content,\n token_count=context_files.total_token_count,\n message_type=MessageType.USER,\n )\n\n\ndef run_llm_loop(\n emitter: Emitter,\n state_container: ChatStateContainer,\n simple_chat_history: list[ChatMessageSimple],\n tools: list[Tool],\n custom_agent_prompt: str | None,\n context_files: ExtractedContextFiles,\n persona: Persona | None,\n user_memory_context: UserMemoryContext | None,\n llm: LLM,\n token_counter: Callable[[str], int],\n db_session: Session,\n forced_tool_id: int | None = None,\n user_identity: LLMUserIdentity | None = None,\n chat_session_id: str | None = None,\n chat_files: list[ChatFile] | None = None,\n include_citations: bool = True,\n all_injected_file_metadata: dict[str, FileToolMetadata] | None = None,\n inject_memories_in_prompt: bool = True,\n) -> None:\n with trace(\n \"run_llm_loop\",\n group_id=chat_session_id,\n metadata={\n \"tenant_id\": get_current_tenant_id(),\n \"chat_session_id\": chat_session_id,\n },\n ):\n # Fix some LiteLLM issues,\n from onyx.llm.litellm_singleton.config import (\n initialize_litellm,\n ) # Here for lazy load LiteLLM\n\n initialize_litellm()\n\n # Track when the loop starts for calculating time-to-answer\n loop_start_time = time.monotonic()\n\n # Initialize citation processor for handling citations dynamically\n # When include_citations is True, use HYPERLINK mode to format citations as [[1]](url)\n # When include_citations is False, use REMOVE mode to strip citations from output\n citation_processor = DynamicCitationProcessor(\n citation_mode=(\n CitationMode.HYPERLINK if include_citations else CitationMode.REMOVE\n )\n )\n\n # Add project file citation mappings if project files are present\n project_citation_mapping: CitationMapping = {}\n if context_files.file_metadata:\n project_citation_mapping = _build_context_file_citation_mapping(\n context_files.file_metadata\n )\n citation_processor.update_citation_mapping(project_citation_mapping)\n\n llm_step_result = LlmStepResult(\n reasoning=None,\n answer=None,\n tool_calls=None,\n raw_answer=None,\n )\n\n # Pass the total budget to construct_message_history, which will handle token allocation\n available_tokens = llm.config.max_input_tokens\n tool_choice: ToolChoiceOptions = ToolChoiceOptions.AUTO\n # Initialize gathered_documents with project files if present\n gathered_documents: list[SearchDoc] | None = (\n list(project_citation_mapping.values())\n if project_citation_mapping\n else None\n )\n # TODO allow citing of images in Projects. Since attached to the last user message, it has no text associated with it.\n # One future workaround is to include the images as separate user messages with citation information and process those.\n always_cite_documents: bool = bool(\n context_files.use_as_search_filter or context_files.file_texts\n )\n should_cite_documents: bool = False\n ran_image_gen: bool = False\n just_ran_web_search: bool = False\n has_called_search_tool: bool = False\n code_interpreter_file_generated: bool = False\n fallback_extraction_attempted: bool = False\n citation_mapping: dict[int, str] = {} # Maps citation_num -> document_id/URL\n\n # Fetch this in a short-lived session so the long-running stream loop does\n # not pin a connection just to keep read state alive.\n with get_session_with_current_tenant() as prompt_db_session:\n default_base_system_prompt: str = get_default_base_system_prompt(\n prompt_db_session\n )\n system_prompt = None\n custom_agent_prompt_msg = None\n\n reasoning_cycles = 0\n for llm_cycle_count in range(MAX_LLM_CYCLES):\n # Handling tool calls based on cycle count and past cycle conditions\n out_of_cycles = llm_cycle_count == MAX_LLM_CYCLES - 1\n if forced_tool_id:\n # Needs to be just the single one because the \"required\" currently doesn't have a specified tool, just a binary\n final_tools = [tool for tool in tools if tool.id == forced_tool_id]\n if not final_tools:\n raise ValueError(f\"Tool {forced_tool_id} not found in tools\")\n tool_choice = ToolChoiceOptions.REQUIRED\n forced_tool_id = None\n elif out_of_cycles or ran_image_gen:\n # Last cycle, no tools allowed, just answer!\n tool_choice = ToolChoiceOptions.NONE\n final_tools = []\n else:\n tool_choice = ToolChoiceOptions.AUTO\n final_tools = tools\n\n # Handling the system prompt and custom agent prompt\n # The section below calculates the available tokens for history a bit more accurately\n # now that project files are loaded in.\n if persona and persona.replace_base_system_prompt:\n # Handles the case where user has checked off the \"Replace base system prompt\" checkbox\n system_prompt = (\n ChatMessageSimple(\n message=persona.system_prompt,\n token_count=token_counter(persona.system_prompt),\n message_type=MessageType.SYSTEM,\n )\n if persona.system_prompt\n else None\n )\n custom_agent_prompt_msg = None\n else:\n # If it's an empty string, we assume the user does not want to include it as an empty System message\n if default_base_system_prompt:\n prompt_memory_context = (\n user_memory_context\n if inject_memories_in_prompt\n else (\n user_memory_context.without_memories()\n if user_memory_context\n else None\n )\n )\n system_prompt_str = build_system_prompt(\n base_system_prompt=default_base_system_prompt,\n datetime_aware=persona.datetime_aware if persona else True,\n user_memory_context=prompt_memory_context,\n tools=tools,\n should_cite_documents=should_cite_documents\n or always_cite_documents,\n )\n system_prompt = ChatMessageSimple(\n message=system_prompt_str,\n token_count=token_counter(system_prompt_str),\n message_type=MessageType.SYSTEM,\n )\n custom_agent_prompt_msg = (\n ChatMessageSimple(\n message=custom_agent_prompt,\n token_count=token_counter(custom_agent_prompt),\n message_type=MessageType.USER,\n )\n if custom_agent_prompt\n else None\n )\n else:\n # If there is a custom agent prompt, it replaces the system prompt when the default system prompt is empty\n system_prompt = (\n ChatMessageSimple(\n message=custom_agent_prompt,\n token_count=token_counter(custom_agent_prompt),\n message_type=MessageType.SYSTEM,\n )\n if custom_agent_prompt\n else None\n )\n custom_agent_prompt_msg = None\n\n reminder_message_text: str | None\n if ran_image_gen:\n # Some models are trained to give back images to the user for some similar tool\n # This is to prevent it generating things like:\n # [Cute Cat](attachment://a_cute_cat_sitting_playfully.png)\n reminder_message_text = IMAGE_GEN_REMINDER\n elif just_ran_web_search and not out_of_cycles:\n reminder_message_text = OPEN_URL_REMINDER\n else:\n # This is the default case, the LLM at this point may answer so it is important\n # to include the reminder. Potentially this should also mention citation\n reminder_message_text = build_reminder_message(\n reminder_text=(\n persona.task_prompt if persona and persona.task_prompt else None\n ),\n include_citation_reminder=should_cite_documents\n or always_cite_documents,\n include_file_reminder=code_interpreter_file_generated,\n is_last_cycle=out_of_cycles,\n )\n\n reminder_msg = (\n ChatMessageSimple(\n message=reminder_message_text,\n token_count=token_counter(reminder_message_text),\n message_type=MessageType.USER_REMINDER,\n )\n if reminder_message_text\n else None\n )\n\n truncated_message_history = construct_message_history(\n system_prompt=system_prompt,\n custom_agent_prompt=custom_agent_prompt_msg,\n simple_chat_history=simple_chat_history,\n reminder_message=reminder_msg,\n context_files=context_files,\n available_tokens=available_tokens,\n token_counter=token_counter,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n\n # This calls the LLM, yields packets (reasoning, answers, etc.) and returns the result\n # It also pre-processes the tool calls in preparation for running them\n tool_defs = [tool.tool_definition() for tool in final_tools]\n\n # Calculate total processing time from loop start until now\n # This measures how long the user waits before the answer starts streaming\n pre_answer_processing_time = time.monotonic() - loop_start_time\n\n llm_step_result, has_reasoned = run_llm_step(\n emitter=emitter,\n history=truncated_message_history,\n tool_definitions=tool_defs,\n tool_choice=tool_choice,\n llm=llm,\n placement=Placement(turn_index=llm_cycle_count + reasoning_cycles),\n citation_processor=citation_processor,\n state_container=state_container,\n # The rich docs representation is passed in so that when yielding the answer, it can also\n # immediately yield the full set of found documents. This gives us the option to show the\n # final set of documents immediately if desired.\n final_documents=gathered_documents,\n user_identity=user_identity,\n pre_answer_processing_time=pre_answer_processing_time,\n )\n if has_reasoned:\n reasoning_cycles += 1\n\n # Fallback extraction for LLMs that don't support tool calling natively or are lower quality\n # and might incorrectly output tool calls in other channels\n llm_step_result, attempted = _try_fallback_tool_extraction(\n llm_step_result=llm_step_result,\n tool_choice=tool_choice,\n fallback_extraction_attempted=fallback_extraction_attempted,\n tool_defs=tool_defs,\n turn_index=llm_cycle_count + reasoning_cycles,\n )\n if attempted:\n # To prevent the case of excessive looping with bad models, we only allow one fallback attempt\n fallback_extraction_attempted = True\n\n # Save citation mapping after each LLM step for incremental state updates\n state_container.set_citation_mapping(citation_processor.citation_to_doc)\n\n # Run the LLM selected tools, there is some more logic here than a simple execution\n # each tool might have custom logic here\n tool_responses: list[ToolResponse] = []\n tool_calls = llm_step_result.tool_calls or []\n\n if INTEGRATION_TESTS_MODE and tool_calls:\n for tool_call in tool_calls:\n emitter.emit(\n Packet(\n placement=tool_call.placement,\n obj=ToolCallDebug(\n tool_call_id=tool_call.tool_call_id,\n tool_name=tool_call.tool_name,\n tool_args=tool_call.tool_args,\n ),\n )\n )\n\n if len(tool_calls) > 1:\n emitter.emit(\n Packet(\n placement=Placement(\n turn_index=tool_calls[0].placement.turn_index\n ),\n obj=TopLevelBranching(num_parallel_branches=len(tool_calls)),\n )\n )\n\n # Quick note for why citation_mapping and citation_processors are both needed:\n # 1. Tools return lightweight string mappings, not SearchDoc objects\n # 2. The SearchDoc resolution is deliberately deferred to llm_loop.py\n # 3. The citation_processor operates on SearchDoc objects and can't provide a complete reverse URL lookup for\n # in-flight citations\n # It can be cleaned up but not super trivial or worthwhile right now\n just_ran_web_search = False\n parallel_tool_call_results = run_tool_calls(\n tool_calls=tool_calls,\n tools=final_tools,\n message_history=truncated_message_history,\n user_memory_context=user_memory_context,\n user_info=None, # TODO, this is part of memories right now, might want to separate it out\n citation_mapping=citation_mapping,\n next_citation_num=citation_processor.get_next_citation_number(),\n max_concurrent_tools=None,\n skip_search_query_expansion=has_called_search_tool,\n chat_files=chat_files,\n url_snippet_map=extract_url_snippet_map(gathered_documents or []),\n inject_memories_in_prompt=inject_memories_in_prompt,\n )\n tool_responses = parallel_tool_call_results.tool_responses\n citation_mapping = parallel_tool_call_results.updated_citation_mapping\n\n # Failure case, give something reasonable to the LLM to try again\n if tool_calls and not tool_responses:\n failure_messages = create_tool_call_failure_messages(\n tool_calls, token_counter\n )\n simple_chat_history.extend(failure_messages)\n continue\n\n for tool_response in tool_responses:\n # Extract tool_call from the response (set by run_tool_calls)\n if tool_response.tool_call is None:\n raise ValueError(\"Tool response missing tool_call reference\")\n\n tool_call = tool_response.tool_call\n tab_index = tool_call.placement.tab_index\n\n # Track if search tool was called (for skipping query expansion on subsequent calls)\n if tool_call.tool_name == SearchTool.NAME:\n has_called_search_tool = True\n\n # Track if code interpreter generated files with download links\n if (\n tool_call.tool_name == PythonTool.NAME\n and not code_interpreter_file_generated\n ):\n try:\n parsed = json.loads(tool_response.llm_facing_response)\n if parsed.get(\"generated_files\"):\n code_interpreter_file_generated = True\n except (json.JSONDecodeError, AttributeError):\n pass\n\n # Build a mapping of tool names to tool objects for getting tool_id\n tools_by_name = {tool.name: tool for tool in final_tools}\n\n # Add the results to the chat history. Even though tools may run in parallel,\n # LLM APIs require linear history, so results are added sequentially.\n # Get the tool object to retrieve tool_id\n tool = tools_by_name.get(tool_call.tool_name)\n if not tool:\n raise ValueError(\n f\"Tool '{tool_call.tool_name}' not found in tools list\"\n )\n\n # Extract search_docs if this is a search tool response\n search_docs = None\n displayed_docs = None\n if isinstance(tool_response.rich_response, SearchDocsResponse):\n search_docs = tool_response.rich_response.search_docs\n displayed_docs = tool_response.rich_response.displayed_docs\n\n # Add ALL search docs to state container for DB persistence\n if search_docs:\n state_container.add_search_docs(search_docs)\n\n if gathered_documents:\n gathered_documents.extend(search_docs)\n else:\n gathered_documents = search_docs\n\n # This is used for the Open URL reminder in the next cycle\n # only do this if the web search tool yielded results\n if search_docs and tool_call.tool_name == WebSearchTool.NAME:\n just_ran_web_search = True\n\n # Extract generated_images if this is an image generation tool response\n generated_images = None\n if isinstance(\n tool_response.rich_response, FinalImageGenerationResponse\n ):\n generated_images = tool_response.rich_response.generated_images\n\n # Extract generated_files if this is a code interpreter response\n generated_files = None\n if isinstance(tool_response.rich_response, PythonToolRichResponse):\n generated_files = (\n tool_response.rich_response.generated_files or None\n )\n\n # Persist memory if this is a memory tool response\n memory_snapshot: MemoryToolResponseSnapshot | None = None\n if isinstance(tool_response.rich_response, MemoryToolResponse):\n persisted_memory_id: int | None = None\n if user_memory_context and user_memory_context.user_id:\n if tool_response.rich_response.index_to_replace is not None:\n memory = update_memory_at_index(\n user_id=user_memory_context.user_id,\n index=tool_response.rich_response.index_to_replace,\n new_text=tool_response.rich_response.memory_text,\n db_session=db_session,\n )\n persisted_memory_id = memory.id if memory else None\n else:\n memory = add_memory(\n user_id=user_memory_context.user_id,\n memory_text=tool_response.rich_response.memory_text,\n db_session=db_session,\n )\n persisted_memory_id = memory.id\n operation: Literal[\"add\", \"update\"] = (\n \"update\"\n if tool_response.rich_response.index_to_replace is not None\n else \"add\"\n )\n memory_snapshot = MemoryToolResponseSnapshot(\n memory_text=tool_response.rich_response.memory_text,\n operation=operation,\n memory_id=persisted_memory_id,\n index=tool_response.rich_response.index_to_replace,\n )\n\n if memory_snapshot:\n saved_response = json.dumps(memory_snapshot.model_dump())\n elif isinstance(tool_response.rich_response, CustomToolCallSummary):\n saved_response = json.dumps(\n tool_response.rich_response.model_dump()\n )\n elif isinstance(tool_response.rich_response, str):\n saved_response = tool_response.rich_response\n else:\n saved_response = tool_response.llm_facing_response\n\n tool_call_info = ToolCallInfo(\n parent_tool_call_id=None, # Top-level tool calls are attached to the chat message\n turn_index=llm_cycle_count + reasoning_cycles,\n tab_index=tab_index,\n tool_name=tool_call.tool_name,\n tool_call_id=tool_call.tool_call_id,\n tool_id=tool.id,\n reasoning_tokens=llm_step_result.reasoning, # All tool calls from this loop share the same reasoning\n tool_call_arguments=tool_call.tool_args,\n tool_call_response=saved_response,\n search_docs=displayed_docs or search_docs,\n generated_images=generated_images,\n generated_files=generated_files,\n )\n # Add to state container for partial save support\n state_container.add_tool_call(tool_call_info)\n\n # Update citation processor if this was a search tool\n update_citation_processor_from_tool_response(\n tool_response, citation_processor\n )\n\n # After processing all tool responses for this turn, add messages to history\n # using OpenAI parallel tool calling format:\n # 1. ONE ASSISTANT message with tool_calls array\n # 2. N TOOL_CALL_RESPONSE messages (one per tool call)\n if tool_responses:\n # Filter to only responses with valid tool_call references\n valid_tool_responses = [\n tr for tr in tool_responses if tr.tool_call is not None\n ]\n\n # Build ToolCallSimple list for all tool calls in this turn\n tool_calls_simple: list[ToolCallSimple] = []\n for tool_response in valid_tool_responses:\n tc = tool_response.tool_call\n assert (\n tc is not None\n ) # Already filtered above, this is just for typing purposes\n\n tool_call_message = tc.to_msg_str()\n tool_call_token_count = token_counter(tool_call_message)\n\n tool_calls_simple.append(\n ToolCallSimple(\n tool_call_id=tc.tool_call_id,\n tool_name=tc.tool_name,\n tool_arguments=tc.tool_args,\n token_count=tool_call_token_count,\n )\n )\n\n # Create ONE ASSISTANT message with all tool calls for this turn\n total_tool_call_tokens = sum(tc.token_count for tc in tool_calls_simple)\n assistant_with_tools = ChatMessageSimple(\n message=\"\", # No text content when making tool calls\n token_count=total_tool_call_tokens,\n message_type=MessageType.ASSISTANT,\n tool_calls=tool_calls_simple,\n image_files=None,\n )\n simple_chat_history.append(assistant_with_tools)\n\n # Add TOOL_CALL_RESPONSE messages for each tool call\n for tool_response in valid_tool_responses:\n tc = tool_response.tool_call\n assert tc is not None # Already filtered above\n\n tool_response_message = tool_response.llm_facing_response\n tool_response_token_count = token_counter(tool_response_message)\n\n tool_response_msg = ChatMessageSimple(\n message=tool_response_message,\n token_count=tool_response_token_count,\n message_type=MessageType.TOOL_CALL_RESPONSE,\n tool_call_id=tc.tool_call_id,\n image_files=None,\n )\n simple_chat_history.append(tool_response_msg)\n\n # If no tool calls, then it must have answered, wrap up\n if not llm_step_result.tool_calls or len(llm_step_result.tool_calls) == 0:\n break\n\n # Certain tools do not allow further actions, force the LLM wrap up on the next cycle\n if any(\n tool.tool_name in STOPPING_TOOLS_NAMES\n for tool in llm_step_result.tool_calls\n ):\n ran_image_gen = True\n\n if llm_step_result.tool_calls and any(\n tool.tool_name in CITEABLE_TOOLS_NAMES\n for tool in llm_step_result.tool_calls\n ):\n # As long as 1 tool with citeable documents is called at any point, we ask the LLM to try to cite\n should_cite_documents = True\n\n if not llm_step_result.answer and not llm_step_result.tool_calls:\n raise _build_empty_llm_response_error(\n llm=llm,\n llm_step_result=llm_step_result,\n tool_choice=tool_choice,\n )\n\n if not llm_step_result.answer:\n raise RuntimeError(\n \"The LLM did not return a final answer after tool execution. \"\n \"Typically this indicates invalid tool-call output, a model/provider mismatch, \"\n \"or serving API misconfiguration.\"\n )\n\n emitter.emit(\n Packet(\n placement=Placement(turn_index=llm_cycle_count + reasoning_cycles),\n obj=OverallStop(type=\"stop\"),\n )\n )\n" }, { "path": "backend/onyx/chat/llm_step.py", "content": "import json\nimport re\nimport time\nimport uuid\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom collections.abc import Mapping\nfrom collections.abc import Sequence\nfrom html import unescape\nfrom typing import Any\nfrom typing import cast\n\nfrom onyx.chat.chat_state import ChatStateContainer\nfrom onyx.chat.citation_processor import DynamicCitationProcessor\nfrom onyx.chat.emitter import Emitter\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.chat.models import LlmStepResult\nfrom onyx.chat.tool_call_args_streaming import maybe_emit_argument_delta\nfrom onyx.configs.app_configs import LOG_ONYX_MODEL_INTERACTIONS\nfrom onyx.configs.app_configs import PROMPT_CACHE_CHAT_HISTORY\nfrom onyx.configs.constants import MessageType\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.file_store.models import ChatFileType\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMConfig\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.interfaces import ToolChoiceOptions\nfrom onyx.llm.model_response import Delta\nfrom onyx.llm.models import AssistantMessage\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import FunctionCall\nfrom onyx.llm.models import ImageContentPart\nfrom onyx.llm.models import ImageUrlDetail\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import SystemMessage\nfrom onyx.llm.models import TextContentPart\nfrom onyx.llm.models import ToolCall\nfrom onyx.llm.models import ToolMessage\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.prompt_cache.processor import process_with_prompt_cache\nfrom onyx.llm.utils import model_needs_formatting_reenabled\nfrom onyx.prompts.chat_prompts import CODE_BLOCK_MARKDOWN\nfrom onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE\nfrom onyx.prompts.constants import SYSTEM_REMINDER_TAG_OPEN\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import ReasoningDelta\nfrom onyx.server.query_and_chat.streaming_models import ReasoningDone\nfrom onyx.server.query_and_chat.streaming_models import ReasoningStart\nfrom onyx.tools.models import ToolCallKickoff\nfrom onyx.tracing.framework.create import generation_span\nfrom onyx.utils.b64 import get_image_type_from_bytes\nfrom onyx.utils.jsonriver import Parser\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.postgres_sanitization import sanitize_string\nfrom onyx.utils.text_processing import find_all_json_objects\n\nlogger = setup_logger()\n\n_XML_INVOKE_BLOCK_RE = re.compile(\n r\"[^>]*)>(?P.*?)\",\n re.IGNORECASE | re.DOTALL,\n)\n_XML_PARAMETER_RE = re.compile(\n r\"[^>]*)>(?P.*?)\",\n re.IGNORECASE | re.DOTALL,\n)\n_FUNCTION_CALLS_OPEN_MARKER = \"\"\n\n\nclass _XmlToolCallContentFilter:\n \"\"\"Streaming filter that strips XML-style tool call payload blocks from text.\"\"\"\n\n def __init__(self) -> None:\n self._pending = \"\"\n self._inside_function_calls_block = False\n\n def process(self, content: str) -> str:\n if not content:\n return \"\"\n\n self._pending += content\n output_parts: list[str] = []\n\n while self._pending:\n pending_lower = self._pending.lower()\n\n if self._inside_function_calls_block:\n end_idx = pending_lower.find(_FUNCTION_CALLS_CLOSE_MARKER)\n if end_idx == -1:\n # Keep buffering until we see the close marker.\n return \"\".join(output_parts)\n\n # Drop the whole function_calls block.\n self._pending = self._pending[\n end_idx + len(_FUNCTION_CALLS_CLOSE_MARKER) :\n ]\n self._inside_function_calls_block = False\n continue\n\n start_idx = _find_function_calls_open_marker(pending_lower)\n if start_idx == -1:\n # Keep only a possible prefix of \" 0:\n output_parts.append(self._pending[:emit_upto])\n self._pending = self._pending[emit_upto:]\n return \"\".join(output_parts)\n\n if start_idx > 0:\n output_parts.append(self._pending[:start_idx])\n\n # Enter block-stripping mode and keep scanning for close marker.\n self._pending = self._pending[start_idx:]\n self._inside_function_calls_block = True\n\n return \"\".join(output_parts)\n\n def flush(self) -> str:\n if self._inside_function_calls_block:\n # Drop any incomplete block at stream end.\n self._pending = \"\"\n self._inside_function_calls_block = False\n return \"\"\n\n remaining = self._pending\n self._pending = \"\"\n return remaining\n\n\ndef _matching_open_marker_prefix_len(text: str) -> int:\n \"\"\"Return longest suffix of text that matches prefix of \" bool:\n return char is None or char in {\">\", \" \", \"\\t\", \"\\n\", \"\\r\"}\n\n\ndef _find_function_calls_open_marker(text_lower: str) -> int:\n \"\"\"Find ' Any:\n \"\"\"Attempt to parse a JSON string value into its Python equivalent.\n\n If value is a string that looks like a JSON array or object, parse it.\n Otherwise return the value unchanged.\n\n This handles the case where the LLM returns arguments like:\n - queries: '[\"query1\", \"query2\"]' instead of [\"query1\", \"query2\"]\n \"\"\"\n if not isinstance(value, str):\n return value\n\n stripped = value.strip()\n # Only attempt to parse if it looks like a JSON array or object\n if not (\n (stripped.startswith(\"[\") and stripped.endswith(\"]\"))\n or (stripped.startswith(\"{\") and stripped.endswith(\"}\"))\n ):\n return value\n\n try:\n return json.loads(stripped)\n except json.JSONDecodeError:\n return value\n\n\ndef _parse_tool_args_to_dict(raw_args: Any) -> dict[str, Any]:\n \"\"\"Parse tool arguments into a dict.\n\n Normal case:\n - raw_args == '{\"queries\":[...]}' -> dict via json.loads\n\n Defensive case (JSON string literal of an object):\n - raw_args == '\"{\\\\\"queries\\\\\":[...]}\"' -> json.loads -> str -> json.loads -> dict\n\n Also handles the case where argument values are JSON strings that need parsing:\n - {\"queries\": '[\"q1\", \"q2\"]'} -> {\"queries\": [\"q1\", \"q2\"]}\n\n Anything else returns {}.\n \"\"\"\n\n if raw_args is None:\n return {}\n\n if isinstance(raw_args, dict):\n # Parse any string values that look like JSON arrays/objects\n return {\n k: _try_parse_json_string(sanitize_string(v) if isinstance(v, str) else v)\n for k, v in raw_args.items()\n }\n\n if not isinstance(raw_args, str):\n return {}\n\n # Sanitize before parsing to remove NULL bytes and surrogates\n raw_args = sanitize_string(raw_args)\n\n try:\n parsed1: Any = json.loads(raw_args)\n except json.JSONDecodeError:\n return {}\n\n if isinstance(parsed1, dict):\n # Parse any string values that look like JSON arrays/objects\n return {k: _try_parse_json_string(v) for k, v in parsed1.items()}\n\n if isinstance(parsed1, str):\n try:\n parsed2: Any = json.loads(parsed1)\n except json.JSONDecodeError:\n return {}\n if isinstance(parsed2, dict):\n # Parse any string values that look like JSON arrays/objects\n return {k: _try_parse_json_string(v) for k, v in parsed2.items()}\n return {}\n\n return {}\n\n\ndef _format_message_history_for_logging(\n message_history: LanguageModelInput,\n) -> str:\n \"\"\"Format message history for logging, with special handling for tool calls.\n\n Tool calls are formatted as JSON with 4-space indentation for readability.\n \"\"\"\n formatted_lines = []\n\n separator = \"================================================\"\n\n # Handle single ChatCompletionMessage - wrap in list for uniform processing\n if isinstance(\n message_history, (SystemMessage, UserMessage, AssistantMessage, ToolMessage)\n ):\n message_history = [message_history]\n\n # Handle sequence of messages\n for i, msg in enumerate(message_history):\n if isinstance(msg, SystemMessage):\n formatted_lines.append(f\"Message {i + 1} [system]:\")\n formatted_lines.append(separator)\n formatted_lines.append(f\"{msg.content}\")\n\n elif isinstance(msg, UserMessage):\n formatted_lines.append(f\"Message {i + 1} [user]:\")\n formatted_lines.append(separator)\n if isinstance(msg.content, str):\n formatted_lines.append(f\"{msg.content}\")\n elif isinstance(msg.content, list):\n # Handle multimodal content (text + images)\n for part in msg.content:\n if isinstance(part, TextContentPart):\n formatted_lines.append(f\"{part.text}\")\n elif isinstance(part, ImageContentPart):\n url = part.image_url.url\n formatted_lines.append(f\"[Image: {url[:50]}...]\")\n\n elif isinstance(msg, AssistantMessage):\n formatted_lines.append(f\"Message {i + 1} [assistant]:\")\n formatted_lines.append(separator)\n if msg.content:\n formatted_lines.append(f\"{msg.content}\")\n\n if msg.tool_calls:\n formatted_lines.append(\"Tool calls:\")\n for tool_call in msg.tool_calls:\n tool_call_dict: dict[str, Any] = {\n \"id\": tool_call.id,\n \"type\": tool_call.type,\n \"function\": {\n \"name\": tool_call.function.name,\n \"arguments\": tool_call.function.arguments,\n },\n }\n tool_call_json = json.dumps(tool_call_dict, indent=4)\n formatted_lines.append(tool_call_json)\n\n elif isinstance(msg, ToolMessage):\n formatted_lines.append(f\"Message {i + 1} [tool]:\")\n formatted_lines.append(separator)\n formatted_lines.append(f\"Tool call ID: {msg.tool_call_id}\")\n formatted_lines.append(f\"Response: {msg.content}\")\n\n else:\n # Fallback for unknown message types\n formatted_lines.append(f\"Message {i + 1} [unknown]:\")\n formatted_lines.append(separator)\n formatted_lines.append(f\"{msg}\")\n\n # Add separator before next message (or at end)\n if i < len(message_history) - 1:\n formatted_lines.append(separator)\n\n return \"\\n\".join(formatted_lines)\n\n\ndef _update_tool_call_with_delta(\n tool_calls_in_progress: dict[int, dict[str, Any]],\n tool_call_delta: Any,\n) -> None:\n index = tool_call_delta.index\n\n if index not in tool_calls_in_progress:\n tool_calls_in_progress[index] = {\n # Fallback ID in case the provider never sends one via deltas.\n \"id\": f\"fallback_{uuid.uuid4().hex}\",\n \"name\": None,\n \"arguments\": \"\",\n }\n\n if tool_call_delta.id:\n tool_calls_in_progress[index][\"id\"] = tool_call_delta.id\n\n if tool_call_delta.function:\n if tool_call_delta.function.name:\n tool_calls_in_progress[index][\"name\"] = tool_call_delta.function.name\n\n if tool_call_delta.function.arguments:\n tool_calls_in_progress[index][\n \"arguments\"\n ] += tool_call_delta.function.arguments\n\n\ndef _extract_tool_call_kickoffs(\n id_to_tool_call_map: dict[int, dict[str, Any]],\n turn_index: int,\n tab_index: int | None = None,\n sub_turn_index: int | None = None,\n) -> list[ToolCallKickoff]:\n \"\"\"Extract ToolCallKickoff objects from the tool call map.\n\n Returns a list of ToolCallKickoff objects for valid tool calls (those with both id and name).\n Each tool call is assigned the given turn_index and a tab_index based on its order.\n\n Args:\n id_to_tool_call_map: Map of tool call index to tool call data\n turn_index: The turn index for this set of tool calls\n tab_index: If provided, use this tab_index for all tool calls (otherwise auto-increment)\n sub_turn_index: The sub-turn index for nested tool calls\n \"\"\"\n tool_calls: list[ToolCallKickoff] = []\n tab_index_calculated = 0\n for tool_call_data in id_to_tool_call_map.values():\n if tool_call_data.get(\"id\") and tool_call_data.get(\"name\"):\n tool_args = _parse_tool_args_to_dict(tool_call_data.get(\"arguments\"))\n\n tool_calls.append(\n ToolCallKickoff(\n tool_call_id=tool_call_data[\"id\"],\n tool_name=tool_call_data[\"name\"],\n tool_args=tool_args,\n placement=Placement(\n turn_index=turn_index,\n tab_index=(\n tab_index_calculated if tab_index is None else tab_index\n ),\n sub_turn_index=sub_turn_index,\n ),\n )\n )\n tab_index_calculated += 1\n return tool_calls\n\n\ndef extract_tool_calls_from_response_text(\n response_text: str | None,\n tool_definitions: list[dict],\n placement: Placement,\n) -> list[ToolCallKickoff]:\n \"\"\"Extract tool calls from LLM response text by matching JSON against tool definitions.\n\n This is a fallback mechanism for when the LLM was expected to return tool calls\n but didn't use the proper tool call format. It searches for tool calls embedded\n in response text (JSON first, then XML-like invoke blocks) that match available\n tool definitions.\n\n Args:\n response_text: The LLM's text response to search for tool calls\n tool_definitions: List of tool definitions to match against\n placement: Placement information for the tool calls\n\n Returns:\n List of ToolCallKickoff objects for any matched tool calls\n \"\"\"\n if not response_text or not tool_definitions:\n return []\n\n # Build a map of tool names to their definitions\n tool_name_to_def: dict[str, dict] = {}\n for tool_def in tool_definitions:\n if tool_def.get(\"type\") == \"function\" and \"function\" in tool_def:\n func_def = tool_def[\"function\"]\n tool_name = func_def.get(\"name\")\n if tool_name:\n tool_name_to_def[tool_name] = func_def\n\n if not tool_name_to_def:\n return []\n\n matched_tool_calls: list[tuple[str, dict[str, Any]]] = []\n # Find all JSON objects in the response text\n json_objects = find_all_json_objects(response_text)\n prev_json_obj: dict[str, Any] | None = None\n prev_tool_call: tuple[str, dict[str, Any]] | None = None\n\n for json_obj in json_objects:\n matched_tool_call = _try_match_json_to_tool(json_obj, tool_name_to_def)\n if not matched_tool_call:\n continue\n\n # `find_all_json_objects` can return both an outer tool-call object and\n # its nested arguments object. If both resolve to the same tool call,\n # drop only this nested duplicate artifact.\n if (\n prev_json_obj is not None\n and prev_tool_call is not None\n and matched_tool_call == prev_tool_call\n and _is_nested_arguments_duplicate(\n previous_json_obj=prev_json_obj,\n current_json_obj=json_obj,\n tool_name_to_def=tool_name_to_def,\n )\n ):\n continue\n\n matched_tool_calls.append(matched_tool_call)\n prev_json_obj = json_obj\n prev_tool_call = matched_tool_call\n\n # Some providers/models emit XML-style function calls instead of JSON objects.\n # Keep this as a fallback behind JSON extraction to preserve current behavior.\n if not matched_tool_calls:\n matched_tool_calls = _extract_xml_tool_calls_from_response_text(\n response_text=response_text,\n tool_name_to_def=tool_name_to_def,\n )\n\n tool_calls: list[ToolCallKickoff] = []\n for tab_index, (tool_name, tool_args) in enumerate(matched_tool_calls):\n tool_calls.append(\n ToolCallKickoff(\n tool_call_id=f\"extracted_{uuid.uuid4().hex[:8]}\",\n tool_name=tool_name,\n tool_args=tool_args,\n placement=Placement(\n turn_index=placement.turn_index,\n tab_index=tab_index,\n sub_turn_index=placement.sub_turn_index,\n ),\n )\n )\n\n logger.info(\n f\"Extracted {len(tool_calls)} tool call(s) from response text as fallback\"\n )\n\n return tool_calls\n\n\ndef _extract_xml_tool_calls_from_response_text(\n response_text: str,\n tool_name_to_def: dict[str, dict],\n) -> list[tuple[str, dict[str, Any]]]:\n \"\"\"Extract XML-style tool calls from response text.\n\n Supports formats such as:\n \n \n [\"foo\"]\n \n \n \"\"\"\n matched_tool_calls: list[tuple[str, dict[str, Any]]] = []\n\n for invoke_match in _XML_INVOKE_BLOCK_RE.finditer(response_text):\n invoke_attrs = invoke_match.group(\"attrs\")\n tool_name = _extract_xml_attribute(invoke_attrs, \"name\")\n if not tool_name or tool_name not in tool_name_to_def:\n continue\n\n tool_args: dict[str, Any] = {}\n invoke_body = invoke_match.group(\"body\")\n for parameter_match in _XML_PARAMETER_RE.finditer(invoke_body):\n parameter_attrs = parameter_match.group(\"attrs\")\n parameter_name = _extract_xml_attribute(parameter_attrs, \"name\")\n if not parameter_name:\n continue\n\n string_attr = _extract_xml_attribute(parameter_attrs, \"string\")\n tool_args[parameter_name] = _parse_xml_parameter_value(\n raw_value=parameter_match.group(\"value\"),\n string_attr=string_attr,\n )\n\n matched_tool_calls.append((tool_name, tool_args))\n\n return matched_tool_calls\n\n\ndef _extract_xml_attribute(attrs: str, attr_name: str) -> str | None:\n \"\"\"Extract a single XML-style attribute value from a tag attribute string.\"\"\"\n attr_match = re.search(\n rf\"\"\"\\b{re.escape(attr_name)}\\s*=\\s*(['\"])(.*?)\\1\"\"\",\n attrs,\n flags=re.IGNORECASE | re.DOTALL,\n )\n if not attr_match:\n return None\n return sanitize_string(unescape(attr_match.group(2).strip()))\n\n\ndef _parse_xml_parameter_value(raw_value: str, string_attr: str | None) -> Any:\n \"\"\"Parse a parameter value from XML-style tool call payloads.\"\"\"\n value = sanitize_string(unescape(raw_value).strip())\n\n if string_attr and string_attr.lower() == \"true\":\n return value\n\n try:\n return json.loads(value)\n except json.JSONDecodeError:\n return value\n\n\ndef _resolve_tool_arguments(obj: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"Extract and parse an arguments/parameters value from a tool-call-like object.\n\n Looks for \"arguments\" or \"parameters\" keys, handles JSON-string values,\n and returns a dict if successful, or None otherwise.\n \"\"\"\n arguments = obj.get(\"arguments\", obj.get(\"parameters\", {}))\n if isinstance(arguments, str):\n arguments = sanitize_string(arguments)\n try:\n arguments = json.loads(arguments)\n except json.JSONDecodeError:\n arguments = {}\n if isinstance(arguments, dict):\n return arguments\n return None\n\n\ndef _try_match_json_to_tool(\n json_obj: dict[str, Any],\n tool_name_to_def: dict[str, dict],\n) -> tuple[str, dict[str, Any]] | None:\n \"\"\"Try to match a JSON object to a tool definition.\n\n Supports several formats:\n 1. Direct tool call format: {\"name\": \"tool_name\", \"arguments\": {...}}\n 2. Function call format: {\"function\": {\"name\": \"tool_name\", \"arguments\": {...}}}\n 3. Tool name as key: {\"tool_name\": {...arguments...}}\n 4. Arguments matching a tool's parameter schema\n\n Args:\n json_obj: The JSON object to match\n tool_name_to_def: Map of tool names to their function definitions\n\n Returns:\n Tuple of (tool_name, tool_args) if matched, None otherwise\n \"\"\"\n # Format 1: Direct tool call format {\"name\": \"...\", \"arguments\": {...}}\n if \"name\" in json_obj and json_obj[\"name\"] in tool_name_to_def:\n tool_name = json_obj[\"name\"]\n arguments = _resolve_tool_arguments(json_obj)\n if arguments is not None:\n return (tool_name, arguments)\n\n # Format 2: Function call format {\"function\": {\"name\": \"...\", \"arguments\": {...}}}\n if \"function\" in json_obj and isinstance(json_obj[\"function\"], dict):\n func_obj = json_obj[\"function\"]\n if \"name\" in func_obj and func_obj[\"name\"] in tool_name_to_def:\n tool_name = func_obj[\"name\"]\n arguments = _resolve_tool_arguments(func_obj)\n if arguments is not None:\n return (tool_name, arguments)\n\n # Format 3: Tool name as key {\"tool_name\": {...arguments...}}\n for tool_name in tool_name_to_def:\n if tool_name in json_obj:\n arguments = json_obj[tool_name]\n if isinstance(arguments, dict):\n return (tool_name, arguments)\n\n # Format 4: Check if the JSON object matches a tool's parameter schema\n for tool_name, func_def in tool_name_to_def.items():\n params = func_def.get(\"parameters\", {})\n properties = params.get(\"properties\", {})\n required = params.get(\"required\", [])\n\n if not properties:\n continue\n\n # Check if all required parameters are present (empty required = all optional)\n if all(req in json_obj for req in required):\n # Check if any of the tool's properties are in the JSON object\n matching_props = [prop for prop in properties if prop in json_obj]\n if matching_props:\n # Filter to only include known properties\n filtered_args = {k: v for k, v in json_obj.items() if k in properties}\n return (tool_name, filtered_args)\n\n return None\n\n\ndef _is_nested_arguments_duplicate(\n previous_json_obj: dict[str, Any],\n current_json_obj: dict[str, Any],\n tool_name_to_def: dict[str, dict],\n) -> bool:\n \"\"\"Detect when current object is the nested args object from previous tool call.\"\"\"\n extracted_args = _extract_nested_arguments_obj(previous_json_obj, tool_name_to_def)\n return extracted_args is not None and current_json_obj == extracted_args\n\n\ndef _extract_nested_arguments_obj(\n json_obj: dict[str, Any],\n tool_name_to_def: dict[str, dict],\n) -> dict[str, Any] | None:\n # Format 1: {\"name\": \"...\", \"arguments\": {...}} or {\"name\": \"...\", \"parameters\": {...}}\n if \"name\" in json_obj and json_obj[\"name\"] in tool_name_to_def:\n args_obj = json_obj.get(\"arguments\", json_obj.get(\"parameters\"))\n if isinstance(args_obj, dict):\n return args_obj\n\n # Format 2: {\"function\": {\"name\": \"...\", \"arguments\": {...}}}\n if \"function\" in json_obj and isinstance(json_obj[\"function\"], dict):\n function_obj = json_obj[\"function\"]\n if \"name\" in function_obj and function_obj[\"name\"] in tool_name_to_def:\n args_obj = function_obj.get(\"arguments\", function_obj.get(\"parameters\"))\n if isinstance(args_obj, dict):\n return args_obj\n\n # Format 3: {\"tool_name\": {...arguments...}}\n for tool_name in tool_name_to_def:\n if tool_name in json_obj and isinstance(json_obj[tool_name], dict):\n return json_obj[tool_name]\n\n return None\n\n\ndef _build_structured_assistant_message(msg: ChatMessageSimple) -> AssistantMessage:\n tool_calls_list: list[ToolCall] | None = None\n if msg.tool_calls:\n tool_calls_list = [\n ToolCall(\n id=tc.tool_call_id,\n type=\"function\",\n function=FunctionCall(\n name=tc.tool_name,\n arguments=json.dumps(tc.tool_arguments),\n ),\n )\n for tc in msg.tool_calls\n ]\n\n return AssistantMessage(\n role=\"assistant\",\n content=msg.message or None,\n tool_calls=tool_calls_list,\n )\n\n\ndef _build_structured_tool_response_message(msg: ChatMessageSimple) -> ToolMessage:\n if not msg.tool_call_id:\n raise ValueError(\n f\"Tool call response message encountered but tool_call_id is not available. Message: {msg}\"\n )\n\n return ToolMessage(\n role=\"tool\",\n content=msg.message,\n tool_call_id=msg.tool_call_id,\n )\n\n\nclass _HistoryMessageFormatter:\n def format_assistant_message(self, msg: ChatMessageSimple) -> AssistantMessage:\n raise NotImplementedError\n\n def format_tool_response_message(\n self, msg: ChatMessageSimple\n ) -> ToolMessage | UserMessage:\n raise NotImplementedError\n\n\nclass _DefaultHistoryMessageFormatter(_HistoryMessageFormatter):\n def format_assistant_message(self, msg: ChatMessageSimple) -> AssistantMessage:\n return _build_structured_assistant_message(msg)\n\n def format_tool_response_message(self, msg: ChatMessageSimple) -> ToolMessage:\n return _build_structured_tool_response_message(msg)\n\n\nclass _OllamaHistoryMessageFormatter(_HistoryMessageFormatter):\n def format_assistant_message(self, msg: ChatMessageSimple) -> AssistantMessage:\n if not msg.tool_calls:\n return _build_structured_assistant_message(msg)\n\n tool_call_lines = [\n (\n f\"[Tool Call] name={tc.tool_name} id={tc.tool_call_id} args={json.dumps(tc.tool_arguments)}\"\n )\n for tc in msg.tool_calls\n ]\n assistant_content = (\n \"\\n\".join([msg.message, *tool_call_lines])\n if msg.message\n else \"\\n\".join(tool_call_lines)\n )\n return AssistantMessage(\n role=\"assistant\",\n content=assistant_content,\n tool_calls=None,\n )\n\n def format_tool_response_message(self, msg: ChatMessageSimple) -> UserMessage:\n if not msg.tool_call_id:\n raise ValueError(\n f\"Tool call response message encountered but tool_call_id is not available. Message: {msg}\"\n )\n\n return UserMessage(\n role=\"user\",\n content=f\"[Tool Result] id={msg.tool_call_id}\\n{msg.message}\",\n )\n\n\n_DEFAULT_HISTORY_MESSAGE_FORMATTER = _DefaultHistoryMessageFormatter()\n_OLLAMA_HISTORY_MESSAGE_FORMATTER = _OllamaHistoryMessageFormatter()\n\n\ndef _get_history_message_formatter(llm_config: LLMConfig) -> _HistoryMessageFormatter:\n if llm_config.model_provider == LlmProviderNames.OLLAMA_CHAT:\n return _OLLAMA_HISTORY_MESSAGE_FORMATTER\n\n return _DEFAULT_HISTORY_MESSAGE_FORMATTER\n\n\ndef translate_history_to_llm_format(\n history: list[ChatMessageSimple],\n llm_config: LLMConfig,\n) -> LanguageModelInput:\n \"\"\"Convert a list of ChatMessageSimple to LanguageModelInput format.\n\n Converts ChatMessageSimple messages to ChatCompletionMessage format,\n handling different message types and image files for multimodal support.\n \"\"\"\n messages: list[ChatCompletionMessage] = []\n history_message_formatter = _get_history_message_formatter(llm_config)\n # Note: cacheability is computed from pre-translation ChatMessageSimple types.\n # Some providers flatten tool history into plain assistant/user text, so this split\n # may be less semantically meaningful, but it remains safe and order-preserving.\n last_cacheable_msg_idx = -1\n all_previous_msgs_cacheable = True\n\n for idx, msg in enumerate(history):\n # if the message is being added to the history\n if PROMPT_CACHE_CHAT_HISTORY and msg.message_type in [\n MessageType.SYSTEM,\n MessageType.USER,\n MessageType.USER_REMINDER,\n MessageType.ASSISTANT,\n MessageType.TOOL_CALL_RESPONSE,\n ]:\n all_previous_msgs_cacheable = (\n all_previous_msgs_cacheable and msg.should_cache\n )\n if all_previous_msgs_cacheable:\n last_cacheable_msg_idx = idx\n\n if msg.message_type == MessageType.SYSTEM:\n system_msg = SystemMessage(\n role=\"system\",\n content=msg.message,\n )\n messages.append(system_msg)\n\n elif msg.message_type == MessageType.USER:\n # Handle user messages with potential images\n if msg.image_files:\n # Build content parts: text + images\n content_parts: list[TextContentPart | ImageContentPart] = [\n TextContentPart(\n type=\"text\",\n text=msg.message,\n )\n ]\n\n # Add image parts\n for img_file in msg.image_files:\n if img_file.file_type == ChatFileType.IMAGE:\n try:\n image_type = get_image_type_from_bytes(img_file.content)\n base64_data = img_file.to_base64()\n image_url = f\"data:{image_type};base64,{base64_data}\"\n\n image_part = ImageContentPart(\n type=\"image_url\",\n image_url=ImageUrlDetail(\n url=image_url,\n detail=None,\n ),\n )\n content_parts.append(image_part)\n except Exception as e:\n logger.warning(\n f\"Failed to process image file {img_file.file_id}: {e}. Skipping image.\"\n )\n user_msg = UserMessage(\n role=\"user\",\n content=content_parts,\n )\n messages.append(user_msg)\n else:\n # Simple text-only user message\n user_msg_text = UserMessage(\n role=\"user\",\n content=msg.message,\n )\n messages.append(user_msg_text)\n\n elif msg.message_type == MessageType.USER_REMINDER:\n # User reminder messages are wrapped with system-reminder tags\n # and converted to UserMessage (LLM APIs don't have a native reminder type)\n wrapped_content = f\"{SYSTEM_REMINDER_TAG_OPEN}\\n{msg.message}\\n{SYSTEM_REMINDER_TAG_CLOSE}\"\n reminder_msg = UserMessage(\n role=\"user\",\n content=wrapped_content,\n )\n messages.append(reminder_msg)\n\n elif msg.message_type == MessageType.ASSISTANT:\n messages.append(history_message_formatter.format_assistant_message(msg))\n\n elif msg.message_type == MessageType.TOOL_CALL_RESPONSE:\n messages.append(history_message_formatter.format_tool_response_message(msg))\n\n else:\n logger.warning(\n f\"Unknown message type {msg.message_type} in history. Skipping message.\"\n )\n\n # Apply model-specific formatting when translating to LLM format (e.g. OpenAI\n # reasoning models need CODE_BLOCK_MARKDOWN prefix for correct markdown generation)\n if model_needs_formatting_reenabled(llm_config.model_name):\n for i, m in enumerate(messages):\n if isinstance(m, SystemMessage):\n messages[i] = SystemMessage(\n role=\"system\",\n content=CODE_BLOCK_MARKDOWN + m.content,\n )\n break\n\n # prompt caching: rely on should_cache in ChatMessageSimple to\n # pick the split point for the cacheable prefix and suffix\n if last_cacheable_msg_idx != -1:\n processed_messages, _ = process_with_prompt_cache(\n llm_config=llm_config,\n cacheable_prefix=messages[: last_cacheable_msg_idx + 1],\n suffix=messages[last_cacheable_msg_idx + 1 :],\n continuation=False,\n )\n assert isinstance(processed_messages, list) # for mypy\n messages = processed_messages\n\n return messages\n\n\ndef _increment_turns(\n turn_index: int, sub_turn_index: int | None\n) -> tuple[int, int | None]:\n if sub_turn_index is None:\n return turn_index + 1, None\n else:\n return turn_index, sub_turn_index + 1\n\n\ndef _delta_has_action(delta: Delta) -> bool:\n return bool(delta.content or delta.reasoning_content or delta.tool_calls)\n\n\ndef run_llm_step_pkt_generator(\n history: list[ChatMessageSimple],\n tool_definitions: list[dict],\n tool_choice: ToolChoiceOptions,\n llm: LLM,\n placement: Placement,\n state_container: ChatStateContainer | None,\n citation_processor: DynamicCitationProcessor | None,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n final_documents: list[SearchDoc] | None = None,\n user_identity: LLMUserIdentity | None = None,\n custom_token_processor: (\n Callable[[Delta | None, Any], tuple[Delta | None, Any]] | None\n ) = None,\n max_tokens: int | None = None,\n # TODO: Temporary handling of nested tool calls with agents, figure out a better way to handle this\n use_existing_tab_index: bool = False,\n is_deep_research: bool = False,\n pre_answer_processing_time: float | None = None,\n timeout_override: int | None = None,\n) -> Generator[Packet, None, tuple[LlmStepResult, bool]]:\n \"\"\"Run an LLM step and stream the response as packets.\n NOTE: DO NOT TOUCH THIS FUNCTION BEFORE ASKING YUHONG, this is very finicky and\n delicate logic that is core to the app's main functionality.\n\n This generator function streams LLM responses, processing reasoning content,\n answer content, tool calls, and citations. It yields Packet objects for\n real-time streaming to clients and accumulates the final result.\n\n Args:\n history: List of chat messages in the conversation history.\n tool_definitions: List of tool definitions available to the LLM.\n tool_choice: Tool choice configuration (e.g., \"auto\", \"required\", \"none\").\n llm: Language model interface to use for generation.\n placement: Placement info (turn_index, tab_index, sub_turn_index) for\n positioning packets in the conversation UI.\n state_container: Container for storing chat state (reasoning, answers).\n citation_processor: Optional processor for extracting and formatting citations\n from the response. If provided, processes tokens to identify citations.\n reasoning_effort: Optional reasoning effort configuration for models that\n support reasoning (e.g., o1 models).\n final_documents: Optional list of search documents to include in the response\n start packet.\n user_identity: Optional user identity information for the LLM.\n custom_token_processor: Optional callable that processes each token delta\n before yielding. Receives (delta, processor_state) and returns\n (modified_delta, new_processor_state). Can return None for delta to skip.\n max_tokens: Optional maximum number of tokens for the LLM response.\n use_existing_tab_index: If True, use the tab_index from placement for all\n tool calls instead of auto-incrementing.\n is_deep_research: If True, treat content before tool calls as reasoning\n when tool_choice is REQUIRED.\n pre_answer_processing_time: Optional time spent processing before the\n answer started, recorded in state_container for analytics.\n timeout_override: Optional timeout override for the LLM call.\n\n Yields:\n Packet: Streaming packets containing:\n - ReasoningStart/ReasoningDelta/ReasoningDone for reasoning content\n - AgentResponseStart/AgentResponseDelta for answer content\n - CitationInfo for extracted citations\n - ToolCallKickoff for tool calls (extracted at the end)\n\n Returns:\n tuple[LlmStepResult, bool]: A tuple containing:\n - LlmStepResult: The final result with accumulated reasoning, answer,\n and tool calls (if any).\n - bool: Whether reasoning occurred during this step. This should be used to\n increment the turn index or sub_turn index for the rest of the LLM loop.\n\n Note:\n The function handles incremental state updates, saving reasoning and answer\n tokens to the state container as they are generated. Tool calls are extracted\n and yielded only after the stream completes.\n \"\"\"\n\n turn_index = placement.turn_index\n tab_index = placement.tab_index\n sub_turn_index = placement.sub_turn_index\n\n def _current_placement() -> Placement:\n return Placement(\n turn_index=turn_index,\n tab_index=tab_index,\n sub_turn_index=sub_turn_index,\n )\n\n llm_msg_history = translate_history_to_llm_format(history, llm.config)\n has_reasoned = False\n\n if LOG_ONYX_MODEL_INTERACTIONS:\n logger.debug(\n f\"Message history:\\n{_format_message_history_for_logging(llm_msg_history)}\"\n )\n\n id_to_tool_call_map: dict[int, dict[str, Any]] = {}\n arg_parsers: dict[int, Parser] = {}\n reasoning_start = False\n answer_start = False\n accumulated_reasoning = \"\"\n accumulated_answer = \"\"\n accumulated_raw_answer = \"\"\n stream_chunk_count = 0\n actionable_chunk_count = 0\n empty_chunk_count = 0\n finish_reasons: set[str] = set()\n xml_tool_call_content_filter = _XmlToolCallContentFilter()\n\n processor_state: Any = None\n\n with generation_span(\n model=llm.config.model_name,\n model_config={\n \"base_url\": str(llm.config.api_base or \"\"),\n \"model_impl\": \"litellm\",\n },\n ) as span_generation:\n span_generation.span_data.input = cast(\n Sequence[Mapping[str, Any]], llm_msg_history\n )\n stream_start_time = time.monotonic()\n first_action_recorded = False\n\n def _emit_citation_results(\n results: Generator[str | CitationInfo, None, None],\n ) -> Generator[Packet, None, None]:\n \"\"\"Yield packets for citation processor results (str or CitationInfo).\"\"\"\n nonlocal accumulated_answer\n\n for result in results:\n if isinstance(result, str):\n accumulated_answer += result\n if state_container:\n state_container.set_answer_tokens(accumulated_answer)\n yield Packet(\n placement=_current_placement(),\n obj=AgentResponseDelta(content=result),\n )\n elif isinstance(result, CitationInfo):\n yield Packet(\n placement=_current_placement(),\n obj=result,\n )\n if state_container:\n state_container.add_emitted_citation(result.citation_number)\n\n def _close_reasoning_if_active() -> Generator[Packet, None, None]:\n \"\"\"Emit ReasoningDone and increment turns if reasoning is in progress.\"\"\"\n nonlocal reasoning_start\n nonlocal has_reasoned\n nonlocal turn_index\n nonlocal sub_turn_index\n\n if reasoning_start:\n yield Packet(\n placement=Placement(\n turn_index=turn_index,\n tab_index=tab_index,\n sub_turn_index=sub_turn_index,\n ),\n obj=ReasoningDone(),\n )\n has_reasoned = True\n turn_index, sub_turn_index = _increment_turns(\n turn_index, sub_turn_index\n )\n reasoning_start = False\n\n def _emit_content_chunk(content_chunk: str) -> Generator[Packet, None, None]:\n nonlocal accumulated_answer\n nonlocal accumulated_reasoning\n nonlocal answer_start\n nonlocal reasoning_start\n nonlocal turn_index\n nonlocal sub_turn_index\n\n # When tool_choice is REQUIRED, content before tool calls is reasoning/thinking\n # about which tool to call, not an actual answer to the user.\n # Treat this content as reasoning instead of answer.\n if is_deep_research and tool_choice == ToolChoiceOptions.REQUIRED:\n accumulated_reasoning += content_chunk\n if state_container:\n state_container.set_reasoning_tokens(accumulated_reasoning)\n if not reasoning_start:\n yield Packet(\n placement=_current_placement(),\n obj=ReasoningStart(),\n )\n yield Packet(\n placement=_current_placement(),\n obj=ReasoningDelta(reasoning=content_chunk),\n )\n reasoning_start = True\n return\n\n # Normal flow for AUTO or NONE tool choice\n yield from _close_reasoning_if_active()\n\n if not answer_start:\n # Store pre-answer processing time in state container for save_chat\n if state_container and pre_answer_processing_time is not None:\n state_container.set_pre_answer_processing_time(\n pre_answer_processing_time\n )\n\n yield Packet(\n placement=_current_placement(),\n obj=AgentResponseStart(\n final_documents=final_documents,\n pre_answer_processing_seconds=pre_answer_processing_time,\n ),\n )\n answer_start = True\n\n if citation_processor:\n yield from _emit_citation_results(\n citation_processor.process_token(content_chunk)\n )\n else:\n accumulated_answer += content_chunk\n # Save answer incrementally to state container\n if state_container:\n state_container.set_answer_tokens(accumulated_answer)\n yield Packet(\n placement=_current_placement(),\n obj=AgentResponseDelta(content=content_chunk),\n )\n\n for packet in llm.stream(\n prompt=llm_msg_history,\n tools=tool_definitions,\n tool_choice=tool_choice,\n structured_response_format=None, # TODO\n max_tokens=max_tokens,\n reasoning_effort=reasoning_effort,\n user_identity=user_identity,\n timeout_override=timeout_override,\n ):\n stream_chunk_count += 1\n if packet.usage:\n usage = packet.usage\n span_generation.span_data.usage = {\n \"input_tokens\": usage.prompt_tokens,\n \"output_tokens\": usage.completion_tokens,\n \"cache_read_input_tokens\": usage.cache_read_input_tokens,\n \"cache_creation_input_tokens\": usage.cache_creation_input_tokens,\n }\n # Note: LLM cost tracking is now handled in multi_llm.py\n finish_reason = packet.choice.finish_reason\n if finish_reason:\n finish_reasons.add(str(finish_reason))\n delta = packet.choice.delta\n\n # Weird behavior from some model providers, just log and ignore for now\n if (\n not delta.content\n and delta.reasoning_content is None\n and not delta.tool_calls\n ):\n empty_chunk_count += 1\n logger.warning(\n \"LLM packet is empty (no content, reasoning, or tool calls). \"\n f\"finish_reason={finish_reason}. Skipping: {packet}\"\n )\n continue\n\n if not first_action_recorded and _delta_has_action(delta):\n span_generation.span_data.time_to_first_action_seconds = (\n time.monotonic() - stream_start_time\n )\n first_action_recorded = True\n if _delta_has_action(delta):\n actionable_chunk_count += 1\n\n if custom_token_processor:\n # The custom token processor can modify the deltas for specific custom logic\n # It can also return a state so that it can handle aggregated delta logic etc.\n # Loosely typed so the function can be flexible\n modified_delta, processor_state = custom_token_processor(\n delta, processor_state\n )\n if modified_delta is None:\n continue\n delta = modified_delta\n\n # Should only happen once, frontend does not expect multiple\n # ReasoningStart or ReasoningDone packets.\n if delta.reasoning_content:\n accumulated_reasoning += delta.reasoning_content\n # Save reasoning incrementally to state container\n if state_container:\n state_container.set_reasoning_tokens(accumulated_reasoning)\n if not reasoning_start:\n yield Packet(\n placement=_current_placement(),\n obj=ReasoningStart(),\n )\n yield Packet(\n placement=_current_placement(),\n obj=ReasoningDelta(reasoning=delta.reasoning_content),\n )\n reasoning_start = True\n\n if delta.content:\n # Keep raw content for fallback extraction. Display content can be\n # filtered and, in deep-research REQUIRED mode, routed as reasoning.\n accumulated_raw_answer += delta.content\n filtered_content = xml_tool_call_content_filter.process(delta.content)\n if filtered_content:\n yield from _emit_content_chunk(filtered_content)\n\n if delta.tool_calls:\n yield from _close_reasoning_if_active()\n\n for tool_call_delta in delta.tool_calls:\n # maybe_emit depends and update being called first and attaching the delta\n _update_tool_call_with_delta(id_to_tool_call_map, tool_call_delta)\n yield from maybe_emit_argument_delta(\n tool_calls_in_progress=id_to_tool_call_map,\n tool_call_delta=tool_call_delta,\n placement=_current_placement(),\n parsers=arg_parsers,\n )\n\n # Flush any tail text buffered while checking for split \" tuple[LlmStepResult, bool]:\n \"\"\"Wrapper around run_llm_step_pkt_generator that consumes packets and emits them.\n\n Returns:\n tuple[LlmStepResult, bool]: The LLM step result and whether reasoning occurred.\n \"\"\"\n step_generator = run_llm_step_pkt_generator(\n history=history,\n tool_definitions=tool_definitions,\n tool_choice=tool_choice,\n llm=llm,\n placement=placement,\n state_container=state_container,\n citation_processor=citation_processor,\n reasoning_effort=reasoning_effort,\n final_documents=final_documents,\n user_identity=user_identity,\n custom_token_processor=custom_token_processor,\n max_tokens=max_tokens,\n use_existing_tab_index=use_existing_tab_index,\n is_deep_research=is_deep_research,\n pre_answer_processing_time=pre_answer_processing_time,\n timeout_override=timeout_override,\n )\n\n while True:\n try:\n packet = next(step_generator)\n emitter.emit(packet)\n except StopIteration as e:\n llm_step_result, has_reasoned = e.value\n return llm_step_result, has_reasoned\n" }, { "path": "backend/onyx/chat/models.py", "content": "from collections.abc import Iterator\nfrom typing import Any\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.file_store.models import InMemoryChatFile\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.models import MultiModelMessageResponseIDInfo\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.server.query_and_chat.streaming_models import GeneratedImage\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.tools.models import SearchToolUsage\nfrom onyx.tools.models import ToolCallKickoff\nfrom onyx.tools.tool_implementations.custom.base_tool_types import ToolResultType\n\n\nclass StreamingError(BaseModel):\n error: str\n stack_trace: str | None = None\n error_code: str | None = (\n None # e.g., \"RATE_LIMIT\", \"AUTH_ERROR\", \"TOOL_CALL_FAILED\"\n )\n is_retryable: bool = True # Hint to frontend if retry might help\n details: dict | None = None # Additional context (tool name, model name, etc.)\n\n\nclass CustomToolResponse(BaseModel):\n response: ToolResultType\n tool_name: str\n\n\nclass CreateChatSessionID(BaseModel):\n chat_session_id: UUID\n\n\nAnswerStreamPart = (\n Packet\n | MessageResponseIDInfo\n | MultiModelMessageResponseIDInfo\n | StreamingError\n | CreateChatSessionID\n)\n\nAnswerStream = Iterator[AnswerStreamPart]\n\n\nclass ToolCallResponse(BaseModel):\n \"\"\"Tool call with full details for non-streaming response.\"\"\"\n\n tool_name: str\n tool_arguments: dict[str, Any]\n tool_result: str\n search_docs: list[SearchDoc] | None = None\n generated_images: list[GeneratedImage] | None = None\n # Reasoning that led to the tool call\n pre_reasoning: str | None = None\n\n\nclass ChatBasicResponse(BaseModel):\n # This is built piece by piece, any of these can be None as the flow could break\n answer: str\n answer_citationless: str\n\n top_documents: list[SearchDoc]\n\n error_msg: str | None\n message_id: int\n citation_info: list[CitationInfo]\n\n\nclass ChatFullResponse(BaseModel):\n \"\"\"Complete non-streaming response with all available data.\n NOTE: This model is used for the core flow of the Onyx application, any changes to it should be reviewed and approved by an\n experienced team member. It is very important to 1. avoid bloat and 2. that this remains backwards compatible across versions.\n \"\"\"\n\n # Core response fields\n answer: str\n answer_citationless: str\n pre_answer_reasoning: str | None = None\n tool_calls: list[ToolCallResponse] = []\n\n # Documents & citations\n top_documents: list[SearchDoc]\n citation_info: list[CitationInfo]\n\n # Metadata\n message_id: int\n chat_session_id: UUID | None = None\n error_msg: str | None = None\n\n\nclass ChatLoadedFile(InMemoryChatFile):\n content_text: str | None\n token_count: int\n\n\nclass ToolCallSimple(BaseModel):\n \"\"\"Tool call for ChatMessageSimple representation (mirrors OpenAI format).\n\n Used when an ASSISTANT message contains one or more tool calls.\n Each tool call has an ID, name, arguments, and token count for tracking.\n \"\"\"\n\n tool_call_id: str\n tool_name: str\n tool_arguments: dict[str, Any]\n token_count: int = 0\n\n\nclass ChatMessageSimple(BaseModel):\n message: str\n token_count: int\n message_type: MessageType\n # Only for USER type messages\n image_files: list[ChatLoadedFile] | None = None\n # Only for TOOL_CALL_RESPONSE type messages\n tool_call_id: str | None = None\n # For ASSISTANT messages with tool calls (OpenAI parallel tool calling format)\n tool_calls: list[ToolCallSimple] | None = None\n # The last message for which this is true\n # AND is true for all previous messages\n # (counting from the start of the history)\n # represents the end of the cacheable prefix\n # used for prompt caching\n should_cache: bool = False\n # When this message represents an injected text file, this is the file's ID.\n # Used to detect which file messages survive context-window truncation.\n file_id: str | None = None\n\n\nclass ContextFileMetadata(BaseModel):\n \"\"\"Metadata for a context-injected file to enable citation support.\"\"\"\n\n file_id: str\n filename: str\n file_content: str\n\n\nclass FileToolMetadata(BaseModel):\n \"\"\"Lightweight metadata for exposing files to the FileReaderTool.\n\n Used when files cannot be loaded directly into context (project too large\n or persona-attached user_files without direct-load path). The LLM receives\n a listing of these so it knows which files it can read via ``read_file``.\n \"\"\"\n\n file_id: str\n filename: str\n approx_char_count: int\n\n\nclass ChatHistoryResult(BaseModel):\n \"\"\"Result of converting chat history to simple format.\n\n Bundles the simple messages with metadata for every text file that was\n injected into the history. After context-window truncation drops older\n messages, callers compare surviving ``file_id`` tags against this map\n to discover \"forgotten\" files whose metadata should be provided to the\n FileReaderTool.\n \"\"\"\n\n simple_messages: list[ChatMessageSimple]\n all_injected_file_metadata: dict[str, FileToolMetadata]\n\n\nclass ExtractedContextFiles(BaseModel):\n \"\"\"Result of attempting to load user files (from a project or persona) into context.\"\"\"\n\n file_texts: list[str]\n image_files: list[ChatLoadedFile]\n use_as_search_filter: bool\n total_token_count: int\n # Lightweight metadata for files exposed via FileReaderTool\n # (populated when files don't fit in context and vector DB is disabled).\n file_metadata: list[ContextFileMetadata]\n uncapped_token_count: int | None\n file_metadata_for_tool: list[FileToolMetadata] = []\n\n\nclass SearchParams(BaseModel):\n \"\"\"Resolved search filter IDs and search-tool usage for a chat turn.\"\"\"\n\n project_id_filter: int | None\n persona_id_filter: int | None\n search_usage: SearchToolUsage\n\n\nclass LlmStepResult(BaseModel):\n reasoning: str | None\n answer: str | None\n tool_calls: list[ToolCallKickoff] | None\n # Raw LLM text before any display-oriented filtering/sanitization.\n # Used for fallback tool-call extraction when providers emit calls as text.\n raw_answer: str | None = None\n" }, { "path": "backend/onyx/chat/process_message.py", "content": "\"\"\"\nIMPORTANT: familiarize yourself with the design concepts prior to contributing to this file.\nAn overview can be found in the README.md file in this directory.\n\"\"\"\n\nimport contextvars\nimport io\nimport queue\nimport re\nimport threading\nimport traceback\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom concurrent.futures import ThreadPoolExecutor\nfrom contextvars import Token\nfrom typing import Final\nfrom uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.cache.factory import get_cache_backend\nfrom onyx.chat.chat_processing_checker import set_processing_status\nfrom onyx.chat.chat_state import AvailableFiles\nfrom onyx.chat.chat_state import ChatStateContainer\nfrom onyx.chat.chat_state import ChatTurnSetup\nfrom onyx.chat.chat_utils import build_file_context\nfrom onyx.chat.chat_utils import convert_chat_history\nfrom onyx.chat.chat_utils import create_chat_history_chain\nfrom onyx.chat.chat_utils import create_chat_session_from_request\nfrom onyx.chat.chat_utils import get_custom_agent_prompt\nfrom onyx.chat.chat_utils import is_last_assistant_message_clarification\nfrom onyx.chat.chat_utils import load_all_chat_files\nfrom onyx.chat.compression import calculate_total_history_tokens\nfrom onyx.chat.compression import compress_chat_history\nfrom onyx.chat.compression import find_summary_for_branch\nfrom onyx.chat.compression import get_compression_params\nfrom onyx.chat.emitter import Emitter\nfrom onyx.chat.llm_loop import EmptyLLMResponseError\nfrom onyx.chat.llm_loop import run_llm_loop\nfrom onyx.chat.models import AnswerStream\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.chat.models import ChatBasicResponse\nfrom onyx.chat.models import ChatFullResponse\nfrom onyx.chat.models import ChatLoadedFile\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.chat.models import ContextFileMetadata\nfrom onyx.chat.models import CreateChatSessionID\nfrom onyx.chat.models import ExtractedContextFiles\nfrom onyx.chat.models import FileToolMetadata\nfrom onyx.chat.models import SearchParams\nfrom onyx.chat.models import StreamingError\nfrom onyx.chat.models import ToolCallResponse\nfrom onyx.chat.prompt_utils import calculate_reserved_tokens\nfrom onyx.chat.save_chat import save_chat_turn\nfrom onyx.chat.stop_signal_checker import is_connected as check_stop_signal\nfrom onyx.chat.stop_signal_checker import reset_cancel_status\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import INTEGRATION_TESTS_MODE\nfrom onyx.configs.constants import DEFAULT_PERSONA_ID\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import MilestoneRecordType\nfrom onyx.context.search.models import BaseFilters\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.chat import create_new_chat_message\nfrom onyx.db.chat import get_chat_session_by_id\nfrom onyx.db.chat import get_or_create_root_message\nfrom onyx.db.chat import reserve_message_id\nfrom onyx.db.chat import reserve_multi_model_message_ids\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import HookPoint\nfrom onyx.db.memory import get_memories\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import Persona\nfrom onyx.db.models import User\nfrom onyx.db.models import UserFile\nfrom onyx.db.projects import get_user_files_from_project\nfrom onyx.db.tools import get_tools\nfrom onyx.deep_research.dr_loop import run_deep_research_llm_loop\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import log_onyx_error\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_store.models import ChatFileType\nfrom onyx.file_store.models import InMemoryChatFile\nfrom onyx.file_store.utils import load_in_memory_chat_files\nfrom onyx.file_store.utils import verify_user_files\nfrom onyx.hooks.executor import execute_hook\nfrom onyx.hooks.executor import HookSkipped\nfrom onyx.hooks.executor import HookSoftFailed\nfrom onyx.hooks.points.query_processing import QueryProcessingPayload\nfrom onyx.hooks.points.query_processing import QueryProcessingResponse\nfrom onyx.llm.factory import get_llm_for_persona\nfrom onyx.llm.factory import get_llm_token_counter\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.llm.request_context import reset_llm_mock_response\nfrom onyx.llm.request_context import set_llm_mock_response\nfrom onyx.llm.utils import litellm_exception_to_error_msg\nfrom onyx.onyxbot.slack.models import SlackContext\nfrom onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type\nfrom onyx.server.query_and_chat.models import AUTO_PLACE_AFTER_LATEST_MESSAGE\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.models import ModelResponseSlot\nfrom onyx.server.query_and_chat.models import MultiModelMessageResponseIDInfo\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.server.query_and_chat.streaming_models import OverallStop\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.usage_limits import check_llm_cost_limit_for_provider\nfrom onyx.tools.constants import FILE_READER_TOOL_ID\nfrom onyx.tools.constants import SEARCH_TOOL_ID\nfrom onyx.tools.models import ChatFile\nfrom onyx.tools.models import SearchToolUsage\nfrom onyx.tools.tool_constructor import construct_tools\nfrom onyx.tools.tool_constructor import CustomToolConfig\nfrom onyx.tools.tool_constructor import FileReaderToolConfig\nfrom onyx.tools.tool_constructor import SearchToolConfig\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import mt_cloud_telemetry\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\nERROR_TYPE_CANCELLED = \"cancelled\"\nAPPROX_CHARS_PER_TOKEN = 4\n\n\ndef _collect_available_file_ids(\n chat_history: list[ChatMessage],\n project_id: int | None,\n user_id: UUID | None,\n db_session: Session,\n) -> AvailableFiles:\n \"\"\"Collect all file IDs the FileReaderTool should be allowed to access.\n\n Returns *separate* lists for chat-attached files (``file_record`` IDs) and\n project/user files (``user_file`` IDs) so the tool can pick the right\n loader without a try/except fallback.\"\"\"\n chat_file_ids: set[UUID] = set()\n user_file_ids: set[UUID] = set()\n\n for msg in chat_history:\n if not msg.files:\n continue\n for fd in msg.files:\n try:\n chat_file_ids.add(UUID(fd[\"id\"]))\n except (ValueError, KeyError):\n pass\n\n if project_id:\n user_files = get_user_files_from_project(\n project_id=project_id,\n user_id=user_id,\n db_session=db_session,\n )\n for uf in user_files:\n user_file_ids.add(uf.id)\n\n return AvailableFiles(\n user_file_ids=list(user_file_ids),\n chat_file_ids=list(chat_file_ids),\n )\n\n\ndef _should_enable_slack_search(\n persona: Persona,\n filters: BaseFilters | None,\n) -> bool:\n \"\"\"Determine if Slack search should be enabled.\n\n Returns True if:\n - Source type filter exists and includes Slack, OR\n - Default persona with no source type filter\n \"\"\"\n source_types = filters.source_type if filters else None\n return (source_types is not None and DocumentSource.SLACK in source_types) or (\n persona.id == DEFAULT_PERSONA_ID and source_types is None\n )\n\n\ndef _convert_loaded_files_to_chat_files(\n loaded_files: list[ChatLoadedFile],\n) -> list[ChatFile]:\n \"\"\"Convert ChatLoadedFile objects to ChatFile for tool usage (e.g., PythonTool).\n\n Args:\n loaded_files: List of ChatLoadedFile objects from the chat history\n\n Returns:\n List of ChatFile objects that can be passed to tools\n \"\"\"\n chat_files = []\n for loaded_file in loaded_files:\n if len(loaded_file.content) > 0:\n chat_files.append(\n ChatFile(\n filename=loaded_file.filename or f\"file_{loaded_file.file_id}\",\n content=loaded_file.content,\n )\n )\n return chat_files\n\n\ndef resolve_context_user_files(\n persona: Persona,\n project_id: int | None,\n user_id: UUID | None,\n db_session: Session,\n) -> list[UserFile]:\n \"\"\"Apply the precedence rule to decide which user files to load.\n\n A custom persona fully supersedes the project. When a chat uses a\n custom persona, the project is purely organisational — its files are\n never loaded and never made searchable.\n\n Custom persona → persona's own user_files (may be empty).\n Default persona inside a project → project files.\n Otherwise → empty list.\n \"\"\"\n if persona.id != DEFAULT_PERSONA_ID:\n return list(persona.user_files) if persona.user_files else []\n if project_id:\n return get_user_files_from_project(\n project_id=project_id,\n user_id=user_id,\n db_session=db_session,\n )\n return []\n\n\ndef _empty_extracted_context_files() -> ExtractedContextFiles:\n return ExtractedContextFiles(\n file_texts=[],\n image_files=[],\n use_as_search_filter=False,\n total_token_count=0,\n file_metadata=[],\n uncapped_token_count=None,\n )\n\n\ndef _extract_text_from_in_memory_file(f: InMemoryChatFile) -> str | None:\n \"\"\"Extract text content from an InMemoryChatFile.\n\n PLAIN_TEXT: the content is pre-extracted UTF-8 plaintext stored during\n ingestion — decode directly.\n DOC / CSV / other text types: the content is the original file bytes —\n use extract_file_text which handles encoding detection and format parsing.\n \"\"\"\n try:\n if f.file_type == ChatFileType.PLAIN_TEXT:\n return f.content.decode(\"utf-8\", errors=\"ignore\").replace(\"\\x00\", \"\")\n return extract_file_text(\n file=io.BytesIO(f.content),\n file_name=f.filename or \"\",\n break_on_unprocessable=False,\n )\n except Exception:\n logger.warning(f\"Failed to extract text from file {f.file_id}\", exc_info=True)\n return None\n\n\ndef extract_context_files(\n user_files: list[UserFile],\n llm_max_context_window: int,\n reserved_token_count: int,\n db_session: Session,\n # Because the tokenizer is a generic tokenizer, the token count may be incorrect.\n # to account for this, the maximum context that is allowed for this function is\n # 60% of the LLM's max context window. The other benefit is that for projects with\n # more files, this makes it so that we don't throw away the history too quickly every time.\n max_llm_context_percentage: float = 0.6,\n) -> ExtractedContextFiles:\n \"\"\"Load user files into context if they fit; otherwise flag for search.\n\n The caller is responsible for deciding *which* user files to pass in\n (project files, persona files, etc.). This function only cares about\n the all-or-nothing fit check and the actual content loading.\n\n Args:\n project_id: The project ID to load files from\n user_id: The user ID for authorization\n llm_max_context_window: Maximum tokens allowed in the LLM context window\n reserved_token_count: Number of tokens to reserve for other content\n db_session: Database session\n max_llm_context_percentage: Maximum percentage of the LLM context window to use.\n Returns:\n ExtractedContextFiles containing:\n - List of text content strings from context files (text files only)\n - List of image files from context (ChatLoadedFile objects)\n - Total token count of all extracted files\n - File metadata for context files\n - Uncapped token count of all extracted files\n - File metadata for files that don't fit in context and vector DB is disabled\n \"\"\"\n # TODO(yuhong): I believe this is not handling all file types correctly.\n\n if not user_files:\n return _empty_extracted_context_files()\n\n # Aggregate tokens for the file content that will be added\n # Skip tokens for those with metadata only\n aggregate_tokens = sum(\n uf.token_count or 0\n for uf in user_files\n if not mime_type_to_chat_file_type(uf.file_type).use_metadata_only()\n )\n max_actual_tokens = (\n llm_max_context_window - reserved_token_count\n ) * max_llm_context_percentage\n\n if aggregate_tokens >= max_actual_tokens:\n use_as_search_filter = not DISABLE_VECTOR_DB\n if DISABLE_VECTOR_DB:\n overflow_tool_metadata = [_build_tool_metadata(uf) for uf in user_files]\n else:\n overflow_tool_metadata = [\n _build_tool_metadata(uf)\n for uf in user_files\n if mime_type_to_chat_file_type(uf.file_type).use_metadata_only()\n ]\n return ExtractedContextFiles(\n file_texts=[],\n image_files=[],\n use_as_search_filter=use_as_search_filter,\n total_token_count=0,\n file_metadata=[],\n uncapped_token_count=aggregate_tokens,\n file_metadata_for_tool=overflow_tool_metadata,\n )\n\n # Files fit — load them into context\n user_file_map = {uf.file_id: uf for uf in user_files}\n in_memory_files = load_in_memory_chat_files(\n user_file_ids=[uf.id for uf in user_files],\n db_session=db_session,\n )\n\n file_texts: list[str] = []\n image_files: list[ChatLoadedFile] = []\n file_metadata: list[ContextFileMetadata] = []\n tool_metadata: list[FileToolMetadata] = []\n total_token_count = 0\n\n for f in in_memory_files:\n uf = user_file_map.get(str(f.file_id))\n filename = f.filename or f\"file_{f.file_id}\"\n\n if f.file_type.use_metadata_only():\n # Metadata-only files are not injected as full text.\n # Only the metadata is provided, with LLM using tools\n if not uf:\n logger.error(\n f\"File with id={f.file_id} in metadata-only path with no associated user file\"\n )\n continue\n tool_metadata.append(_build_tool_metadata(uf))\n elif f.file_type.is_text_file():\n text_content = _extract_text_from_in_memory_file(f)\n if not text_content:\n continue\n if not uf:\n logger.warning(f\"No user file for file_id={f.file_id}\")\n continue\n file_texts.append(text_content)\n file_metadata.append(\n ContextFileMetadata(\n file_id=str(uf.id),\n filename=filename,\n file_content=text_content,\n )\n )\n if uf.token_count:\n total_token_count += uf.token_count\n elif f.file_type == ChatFileType.IMAGE:\n token_count = uf.token_count if uf and uf.token_count else 0\n total_token_count += token_count\n image_files.append(\n ChatLoadedFile(\n file_id=f.file_id,\n content=f.content,\n file_type=f.file_type,\n filename=f.filename,\n content_text=None,\n token_count=token_count,\n )\n )\n\n return ExtractedContextFiles(\n file_texts=file_texts,\n image_files=image_files,\n use_as_search_filter=False,\n total_token_count=total_token_count,\n file_metadata=file_metadata,\n uncapped_token_count=aggregate_tokens,\n file_metadata_for_tool=tool_metadata,\n )\n\n\ndef _build_tool_metadata(user_file: UserFile) -> FileToolMetadata:\n \"\"\"Build lightweight FileToolMetadata from a UserFile record.\n\n Delegates to ``build_file_context`` so that the file ID exposed to the\n LLM is always consistent with what FileReaderTool expects.\n \"\"\"\n return build_file_context(\n tool_file_id=str(user_file.id),\n filename=user_file.name,\n file_type=mime_type_to_chat_file_type(user_file.file_type),\n approx_char_count=(user_file.token_count or 0) * APPROX_CHARS_PER_TOKEN,\n ).tool_metadata\n\n\ndef determine_search_params(\n persona_id: int,\n project_id: int | None,\n extracted_context_files: ExtractedContextFiles,\n) -> SearchParams:\n \"\"\"Decide which search filter IDs and search-tool usage apply for a chat turn.\n\n A custom persona fully supersedes the project — project files are never\n searchable and the search tool config is entirely controlled by the\n persona. The project_id filter is only set for the default persona.\n\n For the default persona inside a project:\n - Files overflow → ENABLED (vector DB scopes to these files)\n - Files fit → DISABLED (content already in prompt)\n - No files at all → DISABLED (nothing to search)\n \"\"\"\n is_custom_persona = persona_id != DEFAULT_PERSONA_ID\n\n project_id_filter: int | None = None\n persona_id_filter: int | None = None\n if extracted_context_files.use_as_search_filter:\n if is_custom_persona:\n persona_id_filter = persona_id\n else:\n project_id_filter = project_id\n\n search_usage = SearchToolUsage.AUTO\n if not is_custom_persona and project_id:\n has_context_files = bool(extracted_context_files.uncapped_token_count)\n files_loaded_in_context = bool(extracted_context_files.file_texts)\n\n if extracted_context_files.use_as_search_filter:\n search_usage = SearchToolUsage.ENABLED\n elif files_loaded_in_context or not has_context_files:\n search_usage = SearchToolUsage.DISABLED\n\n return SearchParams(\n project_id_filter=project_id_filter,\n persona_id_filter=persona_id_filter,\n search_usage=search_usage,\n )\n\n\ndef _resolve_query_processing_hook_result(\n hook_result: QueryProcessingResponse | HookSkipped | HookSoftFailed,\n message_text: str,\n) -> str:\n \"\"\"Apply the Query Processing hook result to the message text.\n\n Returns the (possibly rewritten) message text, or raises OnyxError with\n QUERY_REJECTED if the hook signals rejection (query is null or empty).\n HookSkipped and HookSoftFailed are pass-throughs — the original text is\n returned unchanged.\n \"\"\"\n if isinstance(hook_result, (HookSkipped, HookSoftFailed)):\n return message_text\n if not (hook_result.query and hook_result.query.strip()):\n raise OnyxError(\n OnyxErrorCode.QUERY_REJECTED,\n hook_result.rejection_message\n or \"The hook extension for query processing did not return a valid query. No rejection reason was provided.\",\n )\n return hook_result.query.strip()\n\n\ndef build_chat_turn(\n new_msg_req: SendMessageRequest,\n user: User,\n db_session: Session,\n # None → single-model (persona default LLM); non-empty list → multi-model (one LLM per override)\n llm_overrides: list[LLMOverride] | None,\n *,\n litellm_additional_headers: dict[str, str] | None = None,\n custom_tool_additional_headers: dict[str, str] | None = None,\n mcp_headers: dict[str, str] | None = None,\n bypass_acl: bool = False,\n # Slack context for federated Slack search\n slack_context: SlackContext | None = None,\n # Additional context to include in the chat history, e.g. Slack threads where the\n # conversation cannot be represented by a chain of User/Assistant messages.\n # NOTE: not stored in the database, only passed in to the LLM as context\n additional_context: str | None = None,\n) -> Generator[AnswerStreamPart, None, ChatTurnSetup]:\n \"\"\"Shared setup generator for both single-model and multi-model chat turns.\n\n Yields the packet(s) the frontend needs for request tracking, then returns an\n immutable ``ChatTurnSetup`` containing everything the execution strategy needs.\n\n Callers use::\n\n setup = yield from build_chat_turn(new_msg_req, ..., llm_overrides=...)\n\n to forward yielded packets upstream while receiving the return value locally.\n\n Args:\n llm_overrides: ``None`` → single-model (persona default LLM).\n Non-empty list → multi-model (one LLM per override).\n \"\"\"\n tenant_id = get_current_tenant_id()\n is_multi = bool(llm_overrides)\n\n user_id = user.id\n llm_user_identifier = (\n \"anonymous_user\" if user.is_anonymous else (user.email or str(user_id))\n )\n\n # ── Session resolution ───────────────────────────────────────────────────\n if not new_msg_req.chat_session_id:\n if not new_msg_req.chat_session_info:\n raise RuntimeError(\"Must specify a chat session id or chat session info\")\n chat_session = create_chat_session_from_request(\n chat_session_request=new_msg_req.chat_session_info,\n user_id=user_id,\n db_session=db_session,\n )\n yield CreateChatSessionID(chat_session_id=chat_session.id)\n chat_session = get_chat_session_by_id(\n chat_session_id=chat_session.id,\n user_id=user_id,\n db_session=db_session,\n eager_load_persona=True,\n )\n else:\n chat_session = get_chat_session_by_id(\n chat_session_id=new_msg_req.chat_session_id,\n user_id=user_id,\n db_session=db_session,\n eager_load_persona=True,\n )\n\n persona = chat_session.persona\n message_text = new_msg_req.message\n\n user_identity = LLMUserIdentity(\n user_id=llm_user_identifier, session_id=str(chat_session.id)\n )\n\n # Milestone tracking, most devs using the API don't need to understand this\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id) if not user.is_anonymous else tenant_id,\n event=MilestoneRecordType.MULTIPLE_ASSISTANTS,\n )\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id) if not user.is_anonymous else tenant_id,\n event=MilestoneRecordType.USER_MESSAGE_SENT,\n properties={\n \"origin\": new_msg_req.origin.value,\n \"has_files\": len(new_msg_req.file_descriptors) > 0,\n \"has_project\": chat_session.project_id is not None,\n \"has_persona\": persona is not None and persona.id != DEFAULT_PERSONA_ID,\n \"deep_research\": new_msg_req.deep_research,\n },\n )\n\n # Check LLM cost limits before using the LLM (only for Onyx-managed keys),\n # then build the LLM instance(s).\n llms: list[LLM] = []\n model_display_names: list[str] = []\n selected_overrides: list[LLMOverride | None] = (\n list(llm_overrides or [])\n if is_multi\n else [new_msg_req.llm_override or chat_session.llm_override]\n )\n for override in selected_overrides:\n llm = get_llm_for_persona(\n persona=persona,\n user=user,\n llm_override=override,\n additional_headers=litellm_additional_headers,\n )\n check_llm_cost_limit_for_provider(\n db_session=db_session,\n tenant_id=tenant_id,\n llm_provider_api_key=llm.config.api_key,\n )\n llms.append(llm)\n model_display_names.append(_build_model_display_name(override))\n token_counter = get_llm_token_counter(llms[0])\n\n # not sure why we do this, but to maintain parity with previous code:\n if not is_multi:\n model_display_names = [\"\"]\n\n # Verify that the user-specified files actually belong to the user\n verify_user_files(\n user_files=new_msg_req.file_descriptors,\n user_id=user_id,\n db_session=db_session,\n project_id=chat_session.project_id,\n )\n\n # Re-create linear history of messages\n chat_history = create_chat_history_chain(\n chat_session_id=chat_session.id, db_session=db_session\n )\n\n # Determine the parent message based on the request:\n # - AUTO_PLACE_AFTER_LATEST_MESSAGE (-1): auto-place after latest message in chain\n # - None or root ID: regeneration from root (first message)\n # - positive int: place after that specific parent message\n root_message = get_or_create_root_message(\n chat_session_id=chat_session.id, db_session=db_session\n )\n\n if new_msg_req.parent_message_id == AUTO_PLACE_AFTER_LATEST_MESSAGE:\n parent_message = chat_history[-1] if chat_history else root_message\n elif (\n new_msg_req.parent_message_id is None\n or new_msg_req.parent_message_id == root_message.id\n ):\n # Regeneration from root — clear history so we start fresh\n parent_message = root_message\n chat_history = []\n else:\n parent_message = None\n for i in range(len(chat_history) - 1, -1, -1):\n if chat_history[i].id == new_msg_req.parent_message_id:\n parent_message = chat_history[i]\n # Truncate to only messages up to and including the parent\n chat_history = chat_history[: i + 1]\n break\n\n if parent_message is None:\n raise ValueError(\n \"The new message sent is not on the latest mainline of messages\"\n )\n\n # ── Query Processing hook + user message ─────────────────────────────────\n # Skipped on regeneration (parent is USER type): message already exists/was accepted.\n if parent_message.message_type == MessageType.USER:\n user_message = parent_message\n else:\n # New message — run the Query Processing hook before saving to DB.\n # Skipped on regeneration: the message already exists and was accepted previously.\n # Skip for empty/whitespace-only messages — no meaningful query to process,\n # and SendMessageRequest.message has no min_length guard.\n if message_text.strip():\n hook_result = execute_hook(\n db_session=db_session,\n hook_point=HookPoint.QUERY_PROCESSING,\n payload=QueryProcessingPayload(\n query=message_text,\n # Pass None for anonymous users or authenticated users without an email\n # (e.g. some SSO flows). QueryProcessingPayload.user_email is str | None,\n # so None is accepted and serialised as null in both cases.\n user_email=None if user.is_anonymous else user.email,\n chat_session_id=str(chat_session.id),\n ).model_dump(),\n response_type=QueryProcessingResponse,\n )\n message_text = _resolve_query_processing_hook_result(\n hook_result, message_text\n )\n\n user_message = create_new_chat_message(\n chat_session_id=chat_session.id,\n parent_message=parent_message,\n message=message_text,\n token_count=token_counter(message_text),\n message_type=MessageType.USER,\n files=new_msg_req.file_descriptors,\n db_session=db_session,\n commit=True,\n )\n chat_history.append(user_message)\n\n # Collect file IDs for the file reader tool *before* summary truncation so\n # that files attached to older (summarized-away) messages are still accessible\n # via the FileReaderTool.\n available_files = _collect_available_file_ids(\n chat_history=chat_history,\n project_id=chat_session.project_id,\n user_id=user_id,\n db_session=db_session,\n )\n\n # Find applicable summary for the current branch\n summary_message = find_summary_for_branch(db_session, chat_history)\n # Collect file metadata from messages that will be dropped by summary truncation.\n # These become \"pre-summarized\" file metadata so the forgotten-file mechanism can\n # still tell the LLM about them.\n summarized_file_metadata: dict[str, FileToolMetadata] = {}\n if summary_message and summary_message.last_summarized_message_id:\n cutoff_id = summary_message.last_summarized_message_id\n for msg in chat_history:\n if msg.id > cutoff_id or not msg.files:\n continue\n for fd in msg.files:\n file_id = fd.get(\"id\")\n if not file_id:\n continue\n summarized_file_metadata[file_id] = FileToolMetadata(\n file_id=file_id,\n filename=fd.get(\"name\") or \"unknown\",\n # We don't know the exact size without loading the file,\n # but 0 signals \"unknown\" to the LLM.\n approx_char_count=0,\n )\n # Filter chat_history to only messages after the cutoff\n chat_history = [m for m in chat_history if m.id > cutoff_id]\n\n # Compute skip-clarification flag for deep research path (cheap, always available)\n skip_clarification = is_last_assistant_message_clarification(chat_history)\n\n user_memory_context = get_memories(user, db_session)\n\n # This prompt may come from the Agent or Project. Fetched here (before run_llm_loop)\n # because the inner loop shouldn't need to access the DB-form chat history, but we\n # need it early for token reservation.\n custom_agent_prompt = get_custom_agent_prompt(persona, chat_session)\n\n # When use_memories is disabled, strip memories from the prompt context but keep\n # user info/preferences. The full context is still passed to the LLM loop for\n # memory tool persistence.\n prompt_memory_context = (\n user_memory_context\n if user.use_memories\n else user_memory_context.without_memories()\n )\n\n # ── Token reservation ────────────────────────────────────────────────────\n max_reserved_system_prompt_tokens_str = (persona.system_prompt or \"\") + (\n custom_agent_prompt or \"\"\n )\n reserved_token_count = calculate_reserved_tokens(\n db_session=db_session,\n persona_system_prompt=max_reserved_system_prompt_tokens_str,\n token_counter=token_counter,\n files=new_msg_req.file_descriptors,\n user_memory_context=prompt_memory_context,\n )\n\n # Determine which user files to use. A custom persona fully supersedes the project —\n # project files are never loaded or searchable when a custom persona is in play.\n # Only the default persona inside a project uses the project's files.\n context_user_files = resolve_context_user_files(\n persona=persona,\n project_id=chat_session.project_id,\n user_id=user_id,\n db_session=db_session,\n )\n\n # Use the smallest context window across models for safety (harmless for N=1).\n llm_max_context_window = min(llm.config.max_input_tokens for llm in llms)\n\n extracted_context_files = extract_context_files(\n user_files=context_user_files,\n llm_max_context_window=llm_max_context_window,\n reserved_token_count=reserved_token_count,\n db_session=db_session,\n )\n\n search_params = determine_search_params(\n persona_id=persona.id,\n project_id=chat_session.project_id,\n extracted_context_files=extracted_context_files,\n )\n\n # Also grant access to persona-attached user files for FileReaderTool\n if persona.user_files:\n existing = set(available_files.user_file_ids)\n for uf in persona.user_files:\n if uf.id not in existing:\n available_files.user_file_ids.append(uf.id)\n\n all_tools = get_tools(db_session)\n tool_id_to_name_map = {tool.id: tool.name for tool in all_tools}\n\n search_tool_id = next(\n (tool.id for tool in all_tools if tool.in_code_tool_id == SEARCH_TOOL_ID), None\n )\n\n forced_tool_id = new_msg_req.forced_tool_id\n if (\n search_params.search_usage == SearchToolUsage.DISABLED\n and forced_tool_id is not None\n and search_tool_id is not None\n and forced_tool_id == search_tool_id\n ):\n forced_tool_id = None\n\n # TODO(nmgarza5): Once summarization is done, we don't need to load all files from the beginning.\n # Load all files needed for this chat chain into memory.\n files = load_all_chat_files(chat_history, db_session)\n # Convert loaded files to ChatFile format for tools like PythonTool\n chat_files_for_tools = _convert_loaded_files_to_chat_files(files)\n\n # ── Reserve assistant message ID(s) → yield to frontend ──────────────────\n if is_multi:\n assert llm_overrides is not None\n reserved_messages = reserve_multi_model_message_ids(\n db_session=db_session,\n chat_session_id=chat_session.id,\n parent_message_id=user_message.id,\n model_display_names=model_display_names,\n )\n yield MultiModelMessageResponseIDInfo(\n user_message_id=user_message.id,\n responses=[\n ModelResponseSlot(message_id=m.id, model_name=name)\n for m, name in zip(reserved_messages, model_display_names)\n ],\n )\n else:\n assistant_response = reserve_message_id(\n db_session=db_session,\n chat_session_id=chat_session.id,\n parent_message=user_message.id,\n message_type=MessageType.ASSISTANT,\n )\n reserved_messages = [assistant_response]\n yield MessageResponseIDInfo(\n user_message_id=user_message.id,\n reserved_assistant_message_id=assistant_response.id,\n )\n\n # Convert the chat history into a simple format that is free of any DB objects\n # and is easy to parse for the agent loop.\n has_file_reader_tool = any(\n tool.in_code_tool_id == FILE_READER_TOOL_ID for tool in persona.tools\n )\n\n chat_history_result = convert_chat_history(\n chat_history=chat_history,\n files=files,\n context_image_files=extracted_context_files.image_files,\n additional_context=additional_context or new_msg_req.additional_context,\n token_counter=token_counter,\n tool_id_to_name_map=tool_id_to_name_map,\n )\n simple_chat_history = chat_history_result.simple_messages\n\n # Metadata for every text file injected into the history. After context-window\n # truncation drops older messages, the LLM loop compares surviving file_id tags\n # against this map to discover \"forgotten\" files and provide their metadata to\n # FileReaderTool.\n all_injected_file_metadata: dict[str, FileToolMetadata] = (\n chat_history_result.all_injected_file_metadata if has_file_reader_tool else {}\n )\n\n # Merge in file metadata from messages dropped by summary truncation. These files\n # are no longer in simple_chat_history so they'd be invisible to the forgotten-file\n # mechanism — they'll always appear as \"forgotten\" since no surviving message carries\n # their file_id tag.\n if summarized_file_metadata:\n for fid, meta in summarized_file_metadata.items():\n all_injected_file_metadata.setdefault(fid, meta)\n\n if all_injected_file_metadata:\n logger.debug(\n f\"FileReader: file metadata for LLM: {[(fid, m.filename) for fid, m in all_injected_file_metadata.items()]}\"\n )\n\n if summary_message is not None:\n summary_simple = ChatMessageSimple(\n message=summary_message.message,\n token_count=summary_message.token_count,\n message_type=MessageType.ASSISTANT,\n )\n simple_chat_history.insert(0, summary_simple)\n\n # ── Stop signal and processing status ────────────────────────────────────\n cache = get_cache_backend()\n reset_cancel_status(chat_session.id, cache)\n\n def check_is_connected() -> bool:\n return check_stop_signal(chat_session.id, cache)\n\n set_processing_status(\n chat_session_id=chat_session.id,\n cache=cache,\n value=True,\n )\n\n # Release any read transaction before the long-running LLM stream.\n # If commit fails here, reset the processing status before propagating —\n # otherwise the chat session appears stuck at \"processing\" permanently.\n try:\n db_session.commit()\n except Exception:\n set_processing_status(chat_session_id=chat_session.id, cache=cache, value=False)\n raise\n\n return ChatTurnSetup(\n new_msg_req=new_msg_req,\n chat_session=chat_session,\n persona=persona,\n user_message=user_message,\n user_identity=user_identity,\n llms=llms,\n model_display_names=model_display_names,\n simple_chat_history=simple_chat_history,\n extracted_context_files=extracted_context_files,\n reserved_messages=reserved_messages,\n reserved_token_count=reserved_token_count,\n search_params=search_params,\n all_injected_file_metadata=all_injected_file_metadata,\n available_files=available_files,\n tool_id_to_name_map=tool_id_to_name_map,\n forced_tool_id=forced_tool_id,\n files=files,\n chat_files_for_tools=chat_files_for_tools,\n custom_agent_prompt=custom_agent_prompt,\n user_memory_context=user_memory_context,\n skip_clarification=skip_clarification,\n check_is_connected=check_is_connected,\n cache=cache,\n bypass_acl=bypass_acl,\n slack_context=slack_context,\n custom_tool_additional_headers=custom_tool_additional_headers,\n mcp_headers=mcp_headers,\n )\n\n\n# Sentinel placed on the merged queue when a model thread finishes.\n_MODEL_DONE = object()\n\n# How often the drain loop polls for user-initiated cancellation (stop button).\n_CANCEL_POLL_INTERVAL_S: Final[float] = 0.05\n\n\ndef _run_models(\n setup: ChatTurnSetup,\n user: User,\n db_session: Session,\n external_state_container: ChatStateContainer | None = None,\n) -> AnswerStream:\n \"\"\"Stream packets from one or more LLM loops running in parallel worker threads.\n\n Each model gets its own worker thread, DB session, and ``Emitter``. Threads write\n packets to a shared unbounded queue as they are produced; the drain loop yields them\n in arrival order so the caller receives a single interleaved stream regardless of\n how many models are running.\n\n Single-model (N=1) and multi-model (N>1) use the same execution path. Every\n packet is tagged with ``model_index`` by the model's Emitter — ``0`` for N=1,\n ``0``/``1``/``2`` for multi-model.\n\n Args:\n setup: Fully constructed turn context — LLMs, persona, history, tool config.\n user: Authenticated user making the request.\n db_session: Caller's DB session (used for setup reads; each worker opens its own\n session because SQLAlchemy sessions are not thread-safe).\n external_state_container: Pre-constructed state container for the first model.\n Used by evals and the non-streaming API path so the caller can inspect\n accumulated state (tool calls, answer tokens, citations) after the stream\n is consumed. When ``None`` a fresh container is created automatically.\n\n Returns:\n Generator yielding ``Packet`` objects as they arrive from worker threads —\n answer tokens, tool output, citations — followed by a terminal ``Packet``\n containing ``OverallStop`` once all models complete (or one containing\n ``OverallStop(stop_reason=\"user_cancelled\")`` if the connection drops).\n \"\"\"\n n_models = len(setup.llms)\n\n merged_queue: queue.Queue[tuple[int, Packet | Exception | object]] = queue.Queue()\n\n state_containers: list[ChatStateContainer] = [\n (\n external_state_container\n if (external_state_container is not None and i == 0)\n else ChatStateContainer()\n )\n for i in range(n_models)\n ]\n model_succeeded: list[bool] = [False] * n_models\n # Set to True when a model raises an exception (distinct from \"still running\").\n # Used in the stop-button path to avoid calling completion for errored models.\n model_errored: list[bool] = [False] * n_models\n\n # Set when the drain loop exits early (HTTP disconnect / GeneratorExit).\n # Signals emitters to skip future puts so workers exit promptly.\n drain_done = threading.Event()\n\n def _run_model(model_idx: int) -> None:\n \"\"\"Run one LLM loop inside a worker thread, writing packets to ``merged_queue``.\"\"\"\n model_emitter = Emitter(\n model_idx=model_idx,\n merged_queue=merged_queue,\n drain_done=drain_done,\n )\n sc = state_containers[model_idx]\n model_llm = setup.llms[model_idx]\n\n try:\n # Each worker opens its own session — SQLAlchemy sessions are not thread-safe.\n # Do NOT write to the outer db_session (or any shared DB state) from here;\n # all DB writes in this thread must go through thread_db_session.\n with get_session_with_current_tenant() as thread_db_session:\n thread_tool_dict = construct_tools(\n persona=setup.persona,\n db_session=thread_db_session,\n emitter=model_emitter,\n user=user,\n llm=model_llm,\n search_tool_config=SearchToolConfig(\n user_selected_filters=setup.new_msg_req.internal_search_filters,\n project_id_filter=setup.search_params.project_id_filter,\n persona_id_filter=setup.search_params.persona_id_filter,\n bypass_acl=setup.bypass_acl,\n slack_context=setup.slack_context,\n enable_slack_search=_should_enable_slack_search(\n setup.persona, setup.new_msg_req.internal_search_filters\n ),\n ),\n custom_tool_config=CustomToolConfig(\n chat_session_id=setup.chat_session.id,\n message_id=setup.user_message.id,\n additional_headers=setup.custom_tool_additional_headers,\n mcp_headers=setup.mcp_headers,\n ),\n file_reader_tool_config=FileReaderToolConfig(\n user_file_ids=setup.available_files.user_file_ids,\n chat_file_ids=setup.available_files.chat_file_ids,\n ),\n allowed_tool_ids=setup.new_msg_req.allowed_tool_ids,\n search_usage_forcing_setting=setup.search_params.search_usage,\n )\n model_tools = [\n tool\n for tool_list in thread_tool_dict.values()\n for tool in tool_list\n ]\n\n if setup.forced_tool_id and setup.forced_tool_id not in {\n tool.id for tool in model_tools\n }:\n raise ValueError(\n f\"Forced tool {setup.forced_tool_id} not found in tools\"\n )\n\n # Per-thread copy: run_llm_loop mutates simple_chat_history in-place.\n if n_models == 1 and setup.new_msg_req.deep_research:\n if setup.chat_session.project_id:\n raise RuntimeError(\n \"Deep research is not supported for projects\"\n )\n run_deep_research_llm_loop(\n emitter=model_emitter,\n state_container=sc,\n simple_chat_history=list(setup.simple_chat_history),\n tools=model_tools,\n custom_agent_prompt=setup.custom_agent_prompt,\n llm=model_llm,\n token_counter=get_llm_token_counter(model_llm),\n db_session=thread_db_session,\n skip_clarification=setup.skip_clarification,\n user_identity=setup.user_identity,\n chat_session_id=str(setup.chat_session.id),\n all_injected_file_metadata=setup.all_injected_file_metadata,\n )\n else:\n run_llm_loop(\n emitter=model_emitter,\n state_container=sc,\n simple_chat_history=list(setup.simple_chat_history),\n tools=model_tools,\n custom_agent_prompt=setup.custom_agent_prompt,\n context_files=setup.extracted_context_files,\n persona=setup.persona,\n user_memory_context=setup.user_memory_context,\n llm=model_llm,\n token_counter=get_llm_token_counter(model_llm),\n db_session=thread_db_session,\n forced_tool_id=setup.forced_tool_id,\n user_identity=setup.user_identity,\n chat_session_id=str(setup.chat_session.id),\n chat_files=setup.chat_files_for_tools,\n include_citations=setup.new_msg_req.include_citations,\n all_injected_file_metadata=setup.all_injected_file_metadata,\n inject_memories_in_prompt=user.use_memories,\n )\n\n model_succeeded[model_idx] = True\n\n except Exception as e:\n model_errored[model_idx] = True\n merged_queue.put((model_idx, e))\n\n finally:\n merged_queue.put((model_idx, _MODEL_DONE))\n\n def _delete_orphaned_message(model_idx: int, context: str) -> None:\n \"\"\"Delete a reserved ChatMessage that was never populated due to a model error.\"\"\"\n try:\n orphaned = db_session.get(\n ChatMessage, setup.reserved_messages[model_idx].id\n )\n if orphaned is not None:\n db_session.delete(orphaned)\n db_session.commit()\n except Exception:\n logger.exception(\n \"%s orphan cleanup failed for model %d (%s)\",\n context,\n model_idx,\n setup.model_display_names[model_idx],\n )\n\n # Copy contextvars before submitting futures — ThreadPoolExecutor does NOT\n # auto-propagate contextvars in Python 3.11; threads would inherit a blank context.\n worker_context = contextvars.copy_context()\n executor = ThreadPoolExecutor(\n max_workers=n_models, thread_name_prefix=\"multi-model\"\n )\n completion_persisted: bool = False\n try:\n for i in range(n_models):\n executor.submit(worker_context.run, _run_model, i)\n\n # ── Main thread: merge and yield packets ────────────────────────────\n models_remaining = n_models\n while models_remaining > 0:\n try:\n model_idx, item = merged_queue.get(timeout=_CANCEL_POLL_INTERVAL_S)\n except queue.Empty:\n # Check for user-initiated cancellation every 50 ms.\n if not setup.check_is_connected():\n # Save state for every model before exiting.\n # - Succeeded models: full answer (is_connected=True).\n # - Still-in-flight models: partial answer + \"stopped by user\".\n # - Errored models: delete the orphaned reserved message; do NOT\n # save \"stopped by user\" for a model that actually threw an exception.\n for i in range(n_models):\n if model_errored[i]:\n _delete_orphaned_message(i, \"stop-button\")\n continue\n try:\n succeeded = model_succeeded[i]\n llm_loop_completion_handle(\n state_container=state_containers[i],\n is_connected=lambda: succeeded,\n db_session=db_session,\n assistant_message=setup.reserved_messages[i],\n llm=setup.llms[i],\n reserved_tokens=setup.reserved_token_count,\n )\n except Exception:\n logger.exception(\n \"stop-button completion failed for model %d (%s)\",\n i,\n setup.model_display_names[i],\n )\n yield Packet(\n placement=Placement(turn_index=0),\n obj=OverallStop(type=\"stop\", stop_reason=\"user_cancelled\"),\n )\n completion_persisted = True\n return\n continue\n else:\n if item is _MODEL_DONE:\n models_remaining -= 1\n elif isinstance(item, Exception):\n # Yield a tagged error for this model but keep the other models running.\n # Do NOT decrement models_remaining — _run_model's finally always posts\n # _MODEL_DONE, which is the sole completion signal.\n error_msg = str(item)\n stack_trace = \"\".join(\n traceback.format_exception(type(item), item, item.__traceback__)\n )\n model_llm = setup.llms[model_idx]\n if model_llm.config.api_key and len(model_llm.config.api_key) > 2:\n error_msg = error_msg.replace(\n model_llm.config.api_key, \"[REDACTED_API_KEY]\"\n )\n stack_trace = stack_trace.replace(\n model_llm.config.api_key, \"[REDACTED_API_KEY]\"\n )\n yield StreamingError(\n error=error_msg,\n stack_trace=stack_trace,\n error_code=\"MODEL_ERROR\",\n is_retryable=True,\n details={\n \"model\": model_llm.config.model_name,\n \"provider\": model_llm.config.model_provider,\n \"model_index\": model_idx,\n },\n )\n elif isinstance(item, Packet):\n # model_index already embedded by the model's Emitter in _run_model\n yield item\n\n # ── Completion: save each successful model's response ───────────────\n # All model loops have completed (run_llm_loop returned) — no more writes\n # to state_containers. Worker threads may still be closing their own DB\n # sessions, but the main-thread db_session is unshared and safe to use.\n for i in range(n_models):\n if not model_succeeded[i]:\n # Model errored — delete its orphaned reserved message.\n _delete_orphaned_message(i, \"normal\")\n continue\n try:\n llm_loop_completion_handle(\n state_container=state_containers[i],\n is_connected=setup.check_is_connected,\n db_session=db_session,\n assistant_message=setup.reserved_messages[i],\n llm=setup.llms[i],\n reserved_tokens=setup.reserved_token_count,\n )\n except Exception:\n logger.exception(\n \"normal completion failed for model %d (%s)\",\n i,\n setup.model_display_names[i],\n )\n completion_persisted = True\n\n finally:\n if completion_persisted:\n # Normal exit or stop-button exit: completion already persisted.\n # Threads are done (normal path) or can finish in the background (stop-button).\n executor.shutdown(wait=False)\n else:\n # Early exit (GeneratorExit from raw HTTP disconnect, or unhandled\n # exception in the drain loop).\n # 1. Signal emitters to stop — future emit() calls return immediately,\n # so workers exit their LLM loops promptly.\n drain_done.set()\n # 2. Wait for all workers to finish. Once drain_done is set the Emitter\n # short-circuits, so workers should exit quickly.\n executor.shutdown(wait=True)\n # 3. All workers are done — complete from the main thread only.\n for i in range(n_models):\n if model_succeeded[i]:\n try:\n llm_loop_completion_handle(\n state_container=state_containers[i],\n # Model already finished — persist full response.\n is_connected=lambda: True,\n db_session=db_session,\n assistant_message=setup.reserved_messages[i],\n llm=setup.llms[i],\n reserved_tokens=setup.reserved_token_count,\n )\n except Exception:\n logger.exception(\n \"disconnect completion failed for model %d (%s)\",\n i,\n setup.model_display_names[i],\n )\n elif model_errored[i]:\n _delete_orphaned_message(i, \"disconnect\")\n # 4. Drain buffered packets from memory — no consumer is running.\n while not merged_queue.empty():\n try:\n merged_queue.get_nowait()\n except queue.Empty:\n break\n\n\ndef _stream_chat_turn(\n new_msg_req: SendMessageRequest,\n user: User,\n db_session: Session,\n llm_overrides: list[LLMOverride] | None = None,\n litellm_additional_headers: dict[str, str] | None = None,\n custom_tool_additional_headers: dict[str, str] | None = None,\n mcp_headers: dict[str, str] | None = None,\n bypass_acl: bool = False,\n additional_context: str | None = None,\n slack_context: SlackContext | None = None,\n external_state_container: ChatStateContainer | None = None,\n) -> AnswerStream:\n \"\"\"Private implementation for single-model and multi-model chat turn streaming.\n\n Builds the turn context via ``build_chat_turn``, then streams packets from\n ``_run_models`` back to the caller. Handles setup errors, LLM errors, and\n cancellation uniformly, saving whatever partial state has been accumulated\n before re-raising or yielding a terminal error packet.\n\n Not called directly — use the public wrappers:\n - ``handle_stream_message_objects`` for single-model (N=1) requests.\n - ``handle_multi_model_stream`` for side-by-side multi-model comparison (N>1).\n\n Args:\n new_msg_req: The incoming chat request from the user.\n user: Authenticated user; may be anonymous for public personas.\n db_session: Database session for this request.\n llm_overrides: ``None`` → single-model (persona default LLM).\n Non-empty list → multi-model (one LLM per override, 2–3 items).\n litellm_additional_headers: Extra headers forwarded to the LLM provider.\n custom_tool_additional_headers: Extra headers for custom tool HTTP calls.\n mcp_headers: Extra headers for MCP tool calls.\n bypass_acl: If ``True``, document ACL checks are skipped (used by Slack bot).\n additional_context: Extra context prepended to the LLM's chat history, not\n stored in the DB (used for Slack thread hydration).\n slack_context: Federated Slack search context passed through to the search tool.\n external_state_container: Optional pre-constructed state container. When\n provided, accumulated state (tool calls, citations, answer tokens) is\n written into it so the caller can inspect the result after streaming.\n\n Returns:\n Generator yielding ``Packet`` objects — answer tokens, tool output, citations —\n followed by a terminal ``Packet`` containing ``OverallStop``.\n \"\"\"\n if new_msg_req.mock_llm_response is not None and not INTEGRATION_TESTS_MODE:\n raise ValueError(\n \"mock_llm_response can only be used when INTEGRATION_TESTS_MODE=true\"\n )\n\n mock_response_token: Token[str | None] | None = None\n setup: ChatTurnSetup | None = None\n\n try:\n setup = yield from build_chat_turn(\n new_msg_req=new_msg_req,\n user=user,\n db_session=db_session,\n llm_overrides=llm_overrides,\n litellm_additional_headers=litellm_additional_headers,\n custom_tool_additional_headers=custom_tool_additional_headers,\n mcp_headers=mcp_headers,\n bypass_acl=bypass_acl,\n slack_context=slack_context,\n additional_context=additional_context,\n )\n\n # Set mock response token right before the LLM stream begins so that\n # run_in_background threads inherit the correct context.\n if new_msg_req.mock_llm_response is not None:\n mock_response_token = set_llm_mock_response(new_msg_req.mock_llm_response)\n\n yield from _run_models(\n setup=setup,\n user=user,\n db_session=db_session,\n external_state_container=external_state_container,\n )\n\n except OnyxError as e:\n if e.error_code is not OnyxErrorCode.QUERY_REJECTED:\n log_onyx_error(e)\n yield StreamingError(\n error=e.detail,\n error_code=e.error_code.code,\n is_retryable=e.status_code >= 500,\n )\n db_session.rollback()\n return\n\n except ValueError as e:\n logger.exception(\"Failed to process chat message.\")\n yield StreamingError(\n error=str(e),\n error_code=\"VALIDATION_ERROR\",\n is_retryable=True,\n )\n db_session.rollback()\n return\n\n except EmptyLLMResponseError as e:\n stack_trace = traceback.format_exc()\n logger.warning(\n f\"LLM returned an empty response (provider={e.provider}, model={e.model}, tool_choice={e.tool_choice})\"\n )\n yield StreamingError(\n error=e.client_error_msg,\n stack_trace=stack_trace,\n error_code=e.error_code,\n is_retryable=e.is_retryable,\n details={\n \"model\": e.model,\n \"provider\": e.provider,\n \"tool_choice\": e.tool_choice.value,\n },\n )\n db_session.rollback()\n\n except Exception as e:\n logger.exception(f\"Failed to process chat message due to {e}\")\n stack_trace = traceback.format_exc()\n\n llm = setup.llms[0] if setup else None\n if llm:\n client_error_msg, error_code, is_retryable = litellm_exception_to_error_msg(\n e, llm\n )\n if llm.config.api_key and len(llm.config.api_key) > 2:\n client_error_msg = client_error_msg.replace(\n llm.config.api_key, \"[REDACTED_API_KEY]\"\n )\n stack_trace = stack_trace.replace(\n llm.config.api_key, \"[REDACTED_API_KEY]\"\n )\n yield StreamingError(\n error=client_error_msg,\n stack_trace=stack_trace,\n error_code=error_code,\n is_retryable=is_retryable,\n details={\n \"model\": llm.config.model_name,\n \"provider\": llm.config.model_provider,\n },\n )\n else:\n yield StreamingError(\n error=\"Failed to initialize the chat. Please check your configuration and try again.\",\n stack_trace=stack_trace,\n error_code=\"INIT_FAILED\",\n is_retryable=True,\n )\n db_session.rollback()\n\n finally:\n if mock_response_token is not None:\n reset_llm_mock_response(mock_response_token)\n try:\n if setup is not None:\n set_processing_status(\n chat_session_id=setup.chat_session.id,\n cache=setup.cache,\n value=False,\n )\n except Exception:\n logger.exception(\"Error in setting processing status\")\n\n\ndef handle_stream_message_objects(\n new_msg_req: SendMessageRequest,\n user: User,\n db_session: Session,\n litellm_additional_headers: dict[str, str] | None = None,\n custom_tool_additional_headers: dict[str, str] | None = None,\n mcp_headers: dict[str, str] | None = None,\n bypass_acl: bool = False,\n additional_context: str | None = None,\n slack_context: SlackContext | None = None,\n external_state_container: ChatStateContainer | None = None,\n) -> AnswerStream:\n \"\"\"Single-model streaming entrypoint. For multi-model comparison, use ``handle_multi_model_stream``.\"\"\"\n yield from _stream_chat_turn(\n new_msg_req=new_msg_req,\n user=user,\n db_session=db_session,\n llm_overrides=None,\n litellm_additional_headers=litellm_additional_headers,\n custom_tool_additional_headers=custom_tool_additional_headers,\n mcp_headers=mcp_headers,\n bypass_acl=bypass_acl,\n additional_context=additional_context,\n slack_context=slack_context,\n external_state_container=external_state_container,\n )\n\n\ndef _build_model_display_name(override: LLMOverride | None) -> str:\n \"\"\"Build a human-readable display name from an LLM override.\"\"\"\n if override is None:\n return \"unknown\"\n return override.display_name or override.model_version or \"unknown\"\n\n\ndef handle_multi_model_stream(\n new_msg_req: SendMessageRequest,\n user: User,\n db_session: Session,\n llm_overrides: list[LLMOverride],\n litellm_additional_headers: dict[str, str] | None = None,\n custom_tool_additional_headers: dict[str, str] | None = None,\n mcp_headers: dict[str, str] | None = None,\n) -> AnswerStream:\n \"\"\"Thin wrapper for side-by-side multi-model comparison (2–3 models).\n\n Validates the override list and delegates to ``_stream_chat_turn``,\n which handles both single-model and multi-model execution via the same path.\n\n Args:\n new_msg_req: The incoming chat request. ``deep_research`` must be ``False``.\n user: Authenticated user making the request.\n db_session: Database session for this request.\n llm_overrides: Exactly 2 or 3 ``LLMOverride`` objects — one per model to run.\n litellm_additional_headers: Extra headers forwarded to each LLM provider.\n custom_tool_additional_headers: Extra headers for custom tool HTTP calls.\n mcp_headers: Extra headers for MCP tool calls.\n\n Returns:\n Generator yielding interleaved ``Packet`` objects from all models, each tagged\n with ``model_index`` in its placement.\n \"\"\"\n n_models = len(llm_overrides)\n if n_models < 2 or n_models > 3:\n yield StreamingError(\n error=f\"Multi-model requires 2-3 overrides, got {n_models}\",\n error_code=\"VALIDATION_ERROR\",\n is_retryable=False,\n )\n return\n if new_msg_req.deep_research:\n yield StreamingError(\n error=\"Multi-model is not supported with deep research\",\n error_code=\"VALIDATION_ERROR\",\n is_retryable=False,\n )\n return\n yield from _stream_chat_turn(\n new_msg_req=new_msg_req,\n user=user,\n db_session=db_session,\n llm_overrides=llm_overrides,\n litellm_additional_headers=litellm_additional_headers,\n custom_tool_additional_headers=custom_tool_additional_headers,\n mcp_headers=mcp_headers,\n )\n\n\ndef llm_loop_completion_handle(\n state_container: ChatStateContainer,\n is_connected: Callable[[], bool],\n db_session: Session,\n assistant_message: ChatMessage,\n llm: LLM,\n reserved_tokens: int,\n) -> None:\n chat_session_id = assistant_message.chat_session_id\n\n # Snapshot all state under the container's lock before any DB write.\n # Worker threads may still be running (e.g. user-cancellation path), so\n # direct attribute access is not thread-safe — use the provided getters.\n answer_tokens = state_container.get_answer_tokens()\n reasoning_tokens = state_container.get_reasoning_tokens()\n citation_to_doc = state_container.get_citation_to_doc()\n tool_calls = state_container.get_tool_calls()\n is_clarification = state_container.get_is_clarification()\n all_search_docs = state_container.get_all_search_docs()\n emitted_citations = state_container.get_emitted_citations()\n pre_answer_processing_time = state_container.get_pre_answer_processing_time()\n\n completed_normally = is_connected()\n if completed_normally:\n if answer_tokens is None:\n raise RuntimeError(\n \"LLM run completed normally but did not return an answer.\"\n )\n final_answer = answer_tokens\n else:\n # Stopped by user - append stop message\n logger.debug(f\"Chat session {chat_session_id} stopped by user\")\n if answer_tokens:\n final_answer = (\n answer_tokens + \" ... \\n\\nGeneration was stopped by the user.\"\n )\n else:\n final_answer = \"The generation was stopped by the user.\"\n\n save_chat_turn(\n message_text=final_answer,\n reasoning_tokens=reasoning_tokens,\n citation_to_doc=citation_to_doc,\n tool_calls=tool_calls,\n all_search_docs=all_search_docs,\n db_session=db_session,\n assistant_message=assistant_message,\n is_clarification=is_clarification,\n emitted_citations=emitted_citations,\n pre_answer_processing_time=pre_answer_processing_time,\n )\n\n # Check if compression is needed after saving the message\n updated_chat_history = create_chat_history_chain(\n chat_session_id=chat_session_id,\n db_session=db_session,\n )\n total_tokens = calculate_total_history_tokens(updated_chat_history)\n\n compression_params = get_compression_params(\n max_input_tokens=llm.config.max_input_tokens,\n current_history_tokens=total_tokens,\n reserved_tokens=reserved_tokens,\n )\n if compression_params.should_compress:\n # Build tool mapping for formatting messages\n all_tools = get_tools(db_session)\n tool_id_to_name = {tool.id: tool.name for tool in all_tools}\n\n compress_chat_history(\n db_session=db_session,\n chat_history=updated_chat_history,\n llm=llm,\n compression_params=compression_params,\n tool_id_to_name=tool_id_to_name,\n )\n\n\n_CITATION_LINK_START_PATTERN = re.compile(r\"\\s*\\[\\[\\d+\\]\\]\\(\")\n\n\ndef _find_markdown_link_end(text: str, destination_start: int) -> int | None:\n depth = 0\n i = destination_start\n\n while i < len(text):\n curr = text[i]\n if curr == \"\\\\\":\n i += 2\n continue\n\n if curr == \"(\":\n depth += 1\n elif curr == \")\":\n if depth == 0:\n return i\n depth -= 1\n\n i += 1\n\n return None\n\n\ndef remove_answer_citations(answer: str) -> str:\n stripped_parts: list[str] = []\n cursor = 0\n\n while match := _CITATION_LINK_START_PATTERN.search(answer, cursor):\n stripped_parts.append(answer[cursor : match.start()])\n link_end = _find_markdown_link_end(answer, match.end())\n if link_end is None:\n stripped_parts.append(answer[match.start() :])\n return \"\".join(stripped_parts)\n\n cursor = link_end + 1\n\n stripped_parts.append(answer[cursor:])\n return \"\".join(stripped_parts)\n\n\n@log_function_time()\ndef gather_stream(\n packets: AnswerStream,\n) -> ChatBasicResponse:\n answer: str | None = None\n citations: list[CitationInfo] = []\n error_msg: str | None = None\n message_id: int | None = None\n top_documents: list[SearchDoc] = []\n\n for packet in packets:\n if isinstance(packet, Packet):\n # Handle the different packet object types\n if isinstance(packet.obj, AgentResponseStart):\n # AgentResponseStart contains the final documents\n if packet.obj.final_documents:\n top_documents = packet.obj.final_documents\n elif isinstance(packet.obj, AgentResponseDelta):\n # AgentResponseDelta contains incremental content updates\n if answer is None:\n answer = \"\"\n if packet.obj.content:\n answer += packet.obj.content\n elif isinstance(packet.obj, CitationInfo):\n # CitationInfo contains citation information\n citations.append(packet.obj)\n elif isinstance(packet, StreamingError):\n error_msg = packet.error\n elif isinstance(packet, MessageResponseIDInfo):\n message_id = packet.reserved_assistant_message_id\n\n if message_id is None:\n raise ValueError(\"Message ID is required\")\n\n if answer is None:\n if error_msg is not None:\n answer = \"\"\n else:\n # This should never be the case as these non-streamed flows do not have a stop-generation signal\n raise RuntimeError(\"Answer was not generated\")\n\n return ChatBasicResponse(\n answer=answer,\n answer_citationless=remove_answer_citations(answer),\n citation_info=citations,\n message_id=message_id,\n error_msg=error_msg,\n top_documents=top_documents,\n )\n\n\n@log_function_time()\ndef gather_stream_full(\n packets: AnswerStream,\n state_container: ChatStateContainer,\n) -> ChatFullResponse:\n \"\"\"\n Aggregate streaming packets and state container into a complete ChatFullResponse.\n\n This function consumes all packets from the stream and combines them with\n the accumulated state from the ChatStateContainer to build a complete response\n including answer, reasoning, citations, and tool calls.\n\n Args:\n packets: The stream of packets from handle_stream_message_objects\n state_container: The state container that accumulates tool calls, reasoning, etc.\n\n Returns:\n ChatFullResponse with all available data\n \"\"\"\n answer: str | None = None\n citations: list[CitationInfo] = []\n error_msg: str | None = None\n message_id: int | None = None\n top_documents: list[SearchDoc] = []\n chat_session_id: UUID | None = None\n\n for packet in packets:\n if isinstance(packet, Packet):\n if isinstance(packet.obj, AgentResponseStart):\n if packet.obj.final_documents:\n top_documents = packet.obj.final_documents\n elif isinstance(packet.obj, AgentResponseDelta):\n if answer is None:\n answer = \"\"\n if packet.obj.content:\n answer += packet.obj.content\n elif isinstance(packet.obj, CitationInfo):\n citations.append(packet.obj)\n elif isinstance(packet, StreamingError):\n error_msg = packet.error\n elif isinstance(packet, MessageResponseIDInfo):\n message_id = packet.reserved_assistant_message_id\n elif isinstance(packet, CreateChatSessionID):\n chat_session_id = packet.chat_session_id\n\n if message_id is None:\n raise ValueError(\"Message ID is required\")\n\n # Use state_container for complete answer (handles edge cases gracefully)\n final_answer = state_container.get_answer_tokens() or answer or \"\"\n\n # Get reasoning from state container (None when model doesn't produce reasoning)\n reasoning = state_container.get_reasoning_tokens()\n\n # Convert ToolCallInfo list to ToolCallResponse list\n tool_call_responses = [\n ToolCallResponse(\n tool_name=tc.tool_name,\n tool_arguments=tc.tool_call_arguments,\n tool_result=tc.tool_call_response,\n search_docs=tc.search_docs,\n generated_images=tc.generated_images,\n pre_reasoning=tc.reasoning_tokens,\n )\n for tc in state_container.get_tool_calls()\n ]\n\n return ChatFullResponse(\n answer=final_answer,\n answer_citationless=remove_answer_citations(final_answer),\n pre_answer_reasoning=reasoning,\n tool_calls=tool_call_responses,\n top_documents=top_documents,\n citation_info=citations,\n message_id=message_id,\n chat_session_id=chat_session_id,\n error_msg=error_msg,\n )\n" }, { "path": "backend/onyx/chat/prompt_utils.py", "content": "from collections.abc import Callable\nfrom collections.abc import Sequence\nfrom uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.memory import UserMemoryContext\nfrom onyx.db.persona import get_default_behavior_persona\nfrom onyx.db.user_file import calculate_user_files_token_count\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.prompts.chat_prompts import CITATION_REMINDER\nfrom onyx.prompts.chat_prompts import DEFAULT_SYSTEM_PROMPT\nfrom onyx.prompts.chat_prompts import FILE_REMINDER\nfrom onyx.prompts.chat_prompts import LAST_CYCLE_CITATION_REMINDER\nfrom onyx.prompts.chat_prompts import REQUIRE_CITATION_GUIDANCE\nfrom onyx.prompts.prompt_utils import get_company_context\nfrom onyx.prompts.prompt_utils import handle_onyx_date_awareness\nfrom onyx.prompts.prompt_utils import replace_citation_guidance_tag\nfrom onyx.prompts.prompt_utils import replace_reminder_tag\nfrom onyx.prompts.tool_prompts import GENERATE_IMAGE_GUIDANCE\nfrom onyx.prompts.tool_prompts import INTERNAL_SEARCH_GUIDANCE\nfrom onyx.prompts.tool_prompts import MEMORY_GUIDANCE\nfrom onyx.prompts.tool_prompts import OPEN_URLS_GUIDANCE\nfrom onyx.prompts.tool_prompts import PYTHON_TOOL_GUIDANCE\nfrom onyx.prompts.tool_prompts import TOOL_DESCRIPTION_SEARCH_GUIDANCE\nfrom onyx.prompts.tool_prompts import TOOL_SECTION_HEADER\nfrom onyx.prompts.tool_prompts import WEB_SEARCH_GUIDANCE\nfrom onyx.prompts.tool_prompts import WEB_SEARCH_SITE_DISABLED_GUIDANCE\nfrom onyx.prompts.user_info import BASIC_INFORMATION_PROMPT\nfrom onyx.prompts.user_info import TEAM_INFORMATION_PROMPT\nfrom onyx.prompts.user_info import USER_INFORMATION_HEADER\nfrom onyx.prompts.user_info import USER_MEMORIES_PROMPT\nfrom onyx.prompts.user_info import USER_PREFERENCES_PROMPT\nfrom onyx.prompts.user_info import USER_ROLE_PROMPT\nfrom onyx.tools.interface import Tool\nfrom onyx.tools.tool_implementations.images.image_generation_tool import (\n ImageGenerationTool,\n)\nfrom onyx.tools.tool_implementations.memory.memory_tool import MemoryTool\nfrom onyx.tools.tool_implementations.open_url.open_url_tool import OpenURLTool\nfrom onyx.tools.tool_implementations.python.python_tool import PythonTool\nfrom onyx.tools.tool_implementations.search.search_tool import SearchTool\nfrom onyx.tools.tool_implementations.web_search.web_search_tool import WebSearchTool\nfrom onyx.utils.timing import log_function_time\n\n\ndef get_default_base_system_prompt(db_session: Session) -> str:\n default_persona = get_default_behavior_persona(db_session)\n return (\n default_persona.system_prompt\n if default_persona and default_persona.system_prompt is not None\n else DEFAULT_SYSTEM_PROMPT\n )\n\n\n@log_function_time(print_only=True)\ndef calculate_reserved_tokens(\n db_session: Session,\n persona_system_prompt: str,\n token_counter: Callable[[str], int],\n files: list[FileDescriptor] | None = None,\n user_memory_context: UserMemoryContext | None = None,\n) -> int:\n \"\"\"\n Calculate reserved token count for system prompt and user files.\n\n This is used for token estimation purposes to reserve space for:\n - The system prompt (base + custom agent prompt + all guidance)\n - User files attached to the message\n\n Args:\n db_session: Database session\n persona_system_prompt: Custom agent system prompt (can be empty string)\n token_counter: Function that counts tokens in text\n files: List of file descriptors from the chat message (optional)\n user_memory_context: User memory context (optional)\n\n Returns:\n Total reserved token count\n \"\"\"\n base_system_prompt = get_default_base_system_prompt(db_session)\n\n # This is for token estimation purposes\n fake_system_prompt = build_system_prompt(\n base_system_prompt=base_system_prompt,\n datetime_aware=True,\n user_memory_context=user_memory_context,\n tools=None,\n should_cite_documents=True,\n include_all_guidance=True,\n )\n\n custom_agent_prompt = persona_system_prompt if persona_system_prompt else \"\"\n\n reserved_token_count = token_counter(\n # Annoying that the dict has no attributes now\n custom_agent_prompt\n + \" \"\n + fake_system_prompt\n )\n\n # Calculate total token count for files in the last message\n file_token_count = 0\n if files:\n # Extract user_file_id from each file descriptor\n user_file_ids: list[UUID] = []\n for file in files:\n uid = file.get(\"user_file_id\")\n if not uid:\n continue\n try:\n user_file_ids.append(UUID(uid))\n except (TypeError, ValueError, AttributeError):\n # Skip invalid user_file_id values\n continue\n if user_file_ids:\n file_token_count = calculate_user_files_token_count(\n user_file_ids, db_session\n )\n\n reserved_token_count += file_token_count\n\n return reserved_token_count\n\n\ndef build_reminder_message(\n reminder_text: str | None,\n include_citation_reminder: bool,\n include_file_reminder: bool,\n is_last_cycle: bool,\n) -> str | None:\n reminder = reminder_text.strip() if reminder_text else \"\"\n if is_last_cycle:\n reminder += \"\\n\\n\" + LAST_CYCLE_CITATION_REMINDER\n if include_citation_reminder:\n reminder += \"\\n\\n\" + CITATION_REMINDER\n if include_file_reminder:\n reminder += \"\\n\\n\" + FILE_REMINDER\n reminder = reminder.strip()\n return reminder if reminder else None\n\n\ndef _build_user_information_section(\n user_memory_context: UserMemoryContext | None,\n company_context: str | None,\n) -> str:\n \"\"\"Build the complete '# User Information' section with all sub-sections\n in the correct order: Basic Info → Team Info → Preferences → Memories.\"\"\"\n sections: list[str] = []\n\n if user_memory_context:\n ctx = user_memory_context\n has_basic_info = ctx.user_info.name or ctx.user_info.email or ctx.user_info.role\n\n if has_basic_info:\n role_line = (\n USER_ROLE_PROMPT.format(user_role=ctx.user_info.role).strip()\n if ctx.user_info.role\n else \"\"\n )\n if role_line:\n role_line = \"\\n\" + role_line\n sections.append(\n BASIC_INFORMATION_PROMPT.format(\n user_name=ctx.user_info.name or \"\",\n user_email=ctx.user_info.email or \"\",\n user_role=role_line,\n )\n )\n\n if company_context:\n sections.append(\n TEAM_INFORMATION_PROMPT.format(team_information=company_context.strip())\n )\n\n if user_memory_context:\n ctx = user_memory_context\n\n if ctx.user_preferences:\n sections.append(\n USER_PREFERENCES_PROMPT.format(user_preferences=ctx.user_preferences)\n )\n\n if ctx.memories:\n formatted_memories = \"\\n\".join(f\"- {memory}\" for memory in ctx.memories)\n sections.append(\n USER_MEMORIES_PROMPT.format(user_memories=formatted_memories)\n )\n\n if not sections:\n return \"\"\n\n return USER_INFORMATION_HEADER + \"\\n\".join(sections)\n\n\ndef build_system_prompt(\n base_system_prompt: str,\n datetime_aware: bool = False,\n user_memory_context: UserMemoryContext | None = None,\n tools: Sequence[Tool] | None = None,\n should_cite_documents: bool = False,\n include_all_guidance: bool = False,\n) -> str:\n \"\"\"Should only be called with the default behavior system prompt.\n If the user has replaced the default behavior prompt with their custom agent prompt, do not call this function.\n \"\"\"\n system_prompt = handle_onyx_date_awareness(base_system_prompt, datetime_aware)\n\n # Replace citation guidance placeholder if present\n system_prompt, should_append_citation_guidance = replace_citation_guidance_tag(\n system_prompt,\n should_cite_documents=should_cite_documents,\n include_all_guidance=include_all_guidance,\n )\n\n # Replace reminder tag placeholder if present\n system_prompt = replace_reminder_tag(system_prompt)\n\n company_context = get_company_context()\n user_info_section = _build_user_information_section(\n user_memory_context, company_context\n )\n system_prompt += user_info_section\n\n # Append citation guidance after company context if placeholder was not present\n # This maintains backward compatibility and ensures citations are always enforced when needed\n if should_append_citation_guidance:\n system_prompt += REQUIRE_CITATION_GUIDANCE\n\n if include_all_guidance:\n tool_sections = [\n TOOL_DESCRIPTION_SEARCH_GUIDANCE,\n INTERNAL_SEARCH_GUIDANCE,\n WEB_SEARCH_GUIDANCE.format(\n site_colon_disabled=WEB_SEARCH_SITE_DISABLED_GUIDANCE\n ),\n OPEN_URLS_GUIDANCE,\n PYTHON_TOOL_GUIDANCE,\n GENERATE_IMAGE_GUIDANCE,\n MEMORY_GUIDANCE,\n ]\n system_prompt += TOOL_SECTION_HEADER + \"\\n\".join(tool_sections)\n return system_prompt\n\n if tools:\n has_web_search = any(isinstance(tool, WebSearchTool) for tool in tools)\n has_internal_search = any(isinstance(tool, SearchTool) for tool in tools)\n has_open_urls = any(isinstance(tool, OpenURLTool) for tool in tools)\n has_python = any(isinstance(tool, PythonTool) for tool in tools)\n has_generate_image = any(\n isinstance(tool, ImageGenerationTool) for tool in tools\n )\n has_memory = any(isinstance(tool, MemoryTool) for tool in tools)\n\n tool_guidance_sections: list[str] = []\n\n if has_web_search or has_internal_search or include_all_guidance:\n tool_guidance_sections.append(TOOL_DESCRIPTION_SEARCH_GUIDANCE)\n\n # These are not included at the Tool level because the ordering may matter.\n if has_internal_search or include_all_guidance:\n tool_guidance_sections.append(INTERNAL_SEARCH_GUIDANCE)\n\n if has_web_search or include_all_guidance:\n site_disabled_guidance = \"\"\n if has_web_search:\n web_search_tool = next(\n (t for t in tools if isinstance(t, WebSearchTool)), None\n )\n if web_search_tool and not web_search_tool.supports_site_filter:\n site_disabled_guidance = WEB_SEARCH_SITE_DISABLED_GUIDANCE\n tool_guidance_sections.append(\n WEB_SEARCH_GUIDANCE.format(site_colon_disabled=site_disabled_guidance)\n )\n\n if has_open_urls or include_all_guidance:\n tool_guidance_sections.append(OPEN_URLS_GUIDANCE)\n\n if has_python or include_all_guidance:\n tool_guidance_sections.append(PYTHON_TOOL_GUIDANCE)\n\n if has_generate_image or include_all_guidance:\n tool_guidance_sections.append(GENERATE_IMAGE_GUIDANCE)\n\n if has_memory or include_all_guidance:\n tool_guidance_sections.append(MEMORY_GUIDANCE)\n\n if tool_guidance_sections:\n system_prompt += TOOL_SECTION_HEADER + \"\\n\".join(tool_guidance_sections)\n\n return system_prompt\n" }, { "path": "backend/onyx/chat/save_chat.py", "content": "import json\nimport mimetypes\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.chat_state import ChatStateContainer\nfrom onyx.chat.chat_state import SearchDocKey\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.chat import add_search_docs_to_chat_message\nfrom onyx.db.chat import add_search_docs_to_tool_call\nfrom onyx.db.chat import create_db_search_doc\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ToolCall\nfrom onyx.db.tools import create_tool_call_no_commit\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.natural_language_processing.utils import BaseTokenizer\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type\nfrom onyx.tools.models import ToolCallInfo\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.postgres_sanitization import sanitize_string\n\nlogger = setup_logger()\n\n\ndef _extract_referenced_file_descriptors(\n tool_calls: list[ToolCallInfo],\n message_text: str,\n) -> list[FileDescriptor]:\n \"\"\"Extract FileDescriptors for code interpreter files referenced in the message text.\"\"\"\n descriptors: list[FileDescriptor] = []\n for tool_call_info in tool_calls:\n if not tool_call_info.generated_files:\n continue\n for gen_file in tool_call_info.generated_files:\n file_id = (\n gen_file.file_link.rsplit(\"/\", 1)[-1] if gen_file.file_link else \"\"\n )\n if file_id and file_id in message_text:\n mime_type, _ = mimetypes.guess_type(gen_file.filename)\n descriptors.append(\n FileDescriptor(\n id=file_id,\n type=mime_type_to_chat_file_type(mime_type),\n name=gen_file.filename,\n )\n )\n return descriptors\n\n\ndef _create_and_link_tool_calls(\n tool_calls: list[ToolCallInfo],\n assistant_message: ChatMessage,\n db_session: Session,\n default_tokenizer: BaseTokenizer,\n tool_call_to_search_doc_ids: dict[str, list[int]],\n) -> None:\n \"\"\"\n Create ToolCall entries and link parent references and SearchDocs.\n\n This function handles the logic of:\n 1. Creating all ToolCall objects (with temporary parent references)\n 2. Flushing to get DB IDs\n 3. Building mappings and updating parent references\n 4. Linking SearchDocs to ToolCalls\n\n\n Args:\n tool_calls: List of tool call information to create\n assistant_message: The ChatMessage these tool calls belong to\n db_session: Database session\n default_tokenizer: Tokenizer for calculating token counts\n tool_call_to_search_doc_ids: Mapping from tool_call_id to list of search_doc IDs\n \"\"\"\n # Create all ToolCall objects first (without parent_tool_call_id set)\n # We'll update parent references after flushing to get IDs\n tool_call_objects: list[ToolCall] = []\n tool_call_info_map: dict[str, ToolCallInfo] = {}\n\n for tool_call_info in tool_calls:\n tool_call_info_map[tool_call_info.tool_call_id] = tool_call_info\n\n # Calculate tool_call_tokens from arguments\n try:\n arguments_json_str = json.dumps(tool_call_info.tool_call_arguments)\n tool_call_tokens = len(default_tokenizer.encode(arguments_json_str))\n except Exception as e:\n logger.warning(\n f\"Failed to tokenize tool call arguments for {tool_call_info.tool_call_id}: {e}. Using length as (over) estimate.\"\n )\n arguments_json_str = json.dumps(tool_call_info.tool_call_arguments)\n tool_call_tokens = len(arguments_json_str)\n\n parent_message_id = (\n assistant_message.id if tool_call_info.parent_tool_call_id is None else None\n )\n\n # Create ToolCall DB entry (parent_tool_call_id will be set after flush)\n # This is needed to get the IDs for the parent pointers\n tool_call = create_tool_call_no_commit(\n chat_session_id=assistant_message.chat_session_id,\n parent_chat_message_id=parent_message_id,\n turn_number=tool_call_info.turn_index,\n tool_id=tool_call_info.tool_id,\n tool_call_id=tool_call_info.tool_call_id,\n tool_call_arguments=tool_call_info.tool_call_arguments,\n tool_call_response=tool_call_info.tool_call_response,\n tool_call_tokens=tool_call_tokens,\n db_session=db_session,\n parent_tool_call_id=None, # Will be updated after flush\n reasoning_tokens=tool_call_info.reasoning_tokens,\n generated_images=(\n [img.model_dump() for img in tool_call_info.generated_images]\n if tool_call_info.generated_images\n else None\n ),\n tab_index=tool_call_info.tab_index,\n add_only=True,\n )\n\n # Flush to get all of the IDs\n db_session.flush()\n\n tool_call_objects.append(tool_call)\n\n # Build mapping of tool calls (tool_call_id string -> DB id int)\n tool_call_map: dict[str, int] = {}\n for tool_call_obj in tool_call_objects:\n tool_call_map[tool_call_obj.tool_call_id] = tool_call_obj.id\n\n # Update parent_tool_call_id for all tool calls\n # Filter out orphaned children (whose parents don't exist) - this can happen\n # when generation is stopped mid-execution and parent tool calls were cancelled\n valid_tool_calls: list[ToolCall] = []\n for tool_call_obj in tool_call_objects:\n tool_call_info = tool_call_info_map[tool_call_obj.tool_call_id]\n if tool_call_info.parent_tool_call_id is not None:\n parent_id = tool_call_map.get(tool_call_info.parent_tool_call_id)\n if parent_id is not None:\n tool_call_obj.parent_tool_call_id = parent_id\n valid_tool_calls.append(tool_call_obj)\n else:\n # Parent doesn't exist (likely cancelled) - skip this orphaned child\n logger.warning(\n f\"Skipping tool call '{tool_call_obj.tool_call_id}' with missing parent \"\n f\"'{tool_call_info.parent_tool_call_id}' (likely cancelled during execution)\"\n )\n # Remove from DB session to prevent saving\n db_session.delete(tool_call_obj)\n else:\n # Top-level tool call (no parent)\n valid_tool_calls.append(tool_call_obj)\n\n # Link SearchDocs only to valid ToolCalls\n for tool_call_obj in valid_tool_calls:\n search_doc_ids = tool_call_to_search_doc_ids.get(tool_call_obj.tool_call_id, [])\n if search_doc_ids:\n add_search_docs_to_tool_call(\n tool_call_id=tool_call_obj.id,\n search_doc_ids=search_doc_ids,\n db_session=db_session,\n )\n\n\ndef save_chat_turn(\n message_text: str,\n reasoning_tokens: str | None,\n tool_calls: list[ToolCallInfo],\n citation_to_doc: dict[int, SearchDoc],\n all_search_docs: dict[SearchDocKey, SearchDoc],\n db_session: Session,\n assistant_message: ChatMessage,\n is_clarification: bool = False,\n emitted_citations: set[int] | None = None,\n pre_answer_processing_time: float | None = None,\n) -> None:\n \"\"\"\n Save a chat turn by populating the assistant_message and creating related entities.\n\n This function:\n 1. Updates the ChatMessage with text, reasoning tokens, and token count\n 2. Creates DB SearchDoc entries from pre-deduplicated all_search_docs\n 3. Builds tool_call -> search_doc mapping for displayed docs\n 4. Builds citation mapping from citation_to_doc\n 5. Links all unique SearchDocs to the ChatMessage\n 6. Creates ToolCall entries and links SearchDocs to them\n 7. Builds the citations mapping for the ChatMessage\n\n Args:\n message_text: The message content to save\n reasoning_tokens: Optional reasoning tokens for the message\n tool_calls: List of tool call information to create ToolCall entries (may include search_docs)\n citation_to_doc: Mapping from citation number to SearchDoc for building citations\n all_search_docs: Pre-deduplicated search docs from ChatStateContainer\n db_session: Database session for persistence\n assistant_message: The ChatMessage object to populate (should already exist in DB)\n is_clarification: Whether this assistant message is a clarification question (deep research flow)\n emitted_citations: Set of citation numbers that were actually emitted during streaming.\n If provided, only citations in this set will be saved; others are filtered out.\n pre_answer_processing_time: Duration of processing before answer starts (in seconds)\n \"\"\"\n # 1. Update ChatMessage with message content, reasoning tokens, and token count\n sanitized_message_text = (\n sanitize_string(message_text) if message_text else message_text\n )\n assistant_message.message = sanitized_message_text\n assistant_message.reasoning_tokens = (\n sanitize_string(reasoning_tokens) if reasoning_tokens else reasoning_tokens\n )\n assistant_message.is_clarification = is_clarification\n\n # Use pre-answer processing time (captured when MESSAGE_START was emitted)\n if pre_answer_processing_time is not None:\n assistant_message.processing_duration_seconds = pre_answer_processing_time\n\n # Calculate token count using default tokenizer, when storing, this should not use the LLM\n # specific one so we use a system default tokenizer here.\n default_tokenizer = get_tokenizer(None, None)\n if sanitized_message_text:\n assistant_message.token_count = len(\n default_tokenizer.encode(sanitized_message_text)\n )\n else:\n assistant_message.token_count = 0\n\n # 2. Create DB SearchDoc entries from pre-deduplicated all_search_docs\n search_doc_key_to_id: dict[SearchDocKey, int] = {}\n for key, search_doc_py in all_search_docs.items():\n db_search_doc = create_db_search_doc(\n server_search_doc=search_doc_py,\n db_session=db_session,\n commit=False,\n )\n search_doc_key_to_id[key] = db_search_doc.id\n\n # 3. Build tool_call -> search_doc mapping (for displayed docs in each tool call)\n tool_call_to_search_doc_ids: dict[str, list[int]] = {}\n for tool_call_info in tool_calls:\n if tool_call_info.search_docs:\n search_doc_ids_for_tool: list[int] = []\n for search_doc_py in tool_call_info.search_docs:\n key = ChatStateContainer.create_search_doc_key(search_doc_py)\n if key in search_doc_key_to_id:\n search_doc_ids_for_tool.append(search_doc_key_to_id[key])\n else:\n # Displayed doc not in all_search_docs - create it\n # This can happen if displayed_docs contains docs not in search_docs\n db_search_doc = create_db_search_doc(\n server_search_doc=search_doc_py,\n db_session=db_session,\n commit=False,\n )\n search_doc_key_to_id[key] = db_search_doc.id\n search_doc_ids_for_tool.append(db_search_doc.id)\n tool_call_to_search_doc_ids[tool_call_info.tool_call_id] = list(\n set(search_doc_ids_for_tool)\n )\n\n # Collect all search doc IDs for ChatMessage linking\n all_search_doc_ids_set: set[int] = set(search_doc_key_to_id.values())\n\n # 4. Build a citation mapping from the citation number to the saved DB SearchDoc ID\n # Only include citations that were actually emitted during streaming\n citation_number_to_search_doc_id: dict[int, int] = {}\n\n for citation_num, search_doc_py in citation_to_doc.items():\n # Skip citations that weren't actually emitted (if emitted_citations is provided)\n if emitted_citations is not None and citation_num not in emitted_citations:\n continue\n\n # Create the unique key for this SearchDoc version\n search_doc_key = ChatStateContainer.create_search_doc_key(search_doc_py)\n\n # Get the search doc ID (should already exist from processing tool_calls)\n if search_doc_key in search_doc_key_to_id:\n db_search_doc_id = search_doc_key_to_id[search_doc_key]\n else:\n # Citation doc not found in tool call search_docs\n # Expected case: Project files (source_type=FILE) are cited but don't come from tool calls\n # Unexpected case: Other citation-only docs (indicates a potential issue upstream)\n is_project_file = search_doc_py.source_type == DocumentSource.FILE\n\n if is_project_file:\n logger.info(\n f\"Project file citation {search_doc_py.document_id} not in tool calls, creating it\"\n )\n else:\n logger.warning(\n f\"Citation doc {search_doc_py.document_id} not found in tool call search_docs, creating it\"\n )\n\n # Create the SearchDoc in the database\n # NOTE: It's important that this maps to the saved DB Document ID, because\n # the match-highlights are specific to this saved version, not any document that has\n # the same document_id.\n db_search_doc = create_db_search_doc(\n server_search_doc=search_doc_py,\n db_session=db_session,\n commit=False,\n )\n db_search_doc_id = db_search_doc.id\n search_doc_key_to_id[search_doc_key] = db_search_doc_id\n\n # Link project files to ChatMessage to enable frontend preview\n if is_project_file:\n all_search_doc_ids_set.add(db_search_doc_id)\n\n # Build mapping from citation number to search doc ID\n citation_number_to_search_doc_id[citation_num] = db_search_doc_id\n\n # 5. Link all unique SearchDocs (from both tool calls and citations) to ChatMessage\n final_search_doc_ids: list[int] = list(all_search_doc_ids_set)\n if final_search_doc_ids:\n add_search_docs_to_chat_message(\n chat_message_id=assistant_message.id,\n search_doc_ids=final_search_doc_ids,\n db_session=db_session,\n )\n\n # 6. Create ToolCall entries and link SearchDocs to them\n _create_and_link_tool_calls(\n tool_calls=tool_calls,\n assistant_message=assistant_message,\n db_session=db_session,\n default_tokenizer=default_tokenizer,\n tool_call_to_search_doc_ids=tool_call_to_search_doc_ids,\n )\n\n # 7. Build citations mapping - use the mapping we already built in step 4\n assistant_message.citations = (\n citation_number_to_search_doc_id if citation_number_to_search_doc_id else None\n )\n\n # 8. Attach code interpreter generated files that the assistant actually\n # referenced in its response, so they are available via load_all_chat_files\n # on subsequent turns. Files not mentioned are intermediate artifacts.\n if sanitized_message_text:\n referenced = _extract_referenced_file_descriptors(\n tool_calls, sanitized_message_text\n )\n if referenced:\n existing_files = assistant_message.files or []\n assistant_message.files = existing_files + referenced\n\n # Finally save the messages, tool calls, and docs\n db_session.commit()\n" }, { "path": "backend/onyx/chat/stop_signal_checker.py", "content": "from uuid import UUID\n\nfrom onyx.cache.interface import CacheBackend\n\nPREFIX = \"chatsessionstop\"\nFENCE_PREFIX = f\"{PREFIX}_fence\"\nFENCE_TTL = 10 * 60 # 10 minutes\n\n\ndef _get_fence_key(chat_session_id: UUID) -> str:\n \"\"\"Generate the cache key for a chat session stop signal fence.\n\n Args:\n chat_session_id: The UUID of the chat session\n\n Returns:\n The fence key string. Tenant isolation is handled automatically\n by the cache backend (Redis key-prefixing or Postgres schema routing).\n \"\"\"\n return f\"{FENCE_PREFIX}_{chat_session_id}\"\n\n\ndef set_fence(chat_session_id: UUID, cache: CacheBackend, value: bool) -> None:\n \"\"\"Set or clear the stop signal fence for a chat session.\n\n Args:\n chat_session_id: The UUID of the chat session\n cache: Tenant-aware cache backend\n value: True to set the fence (stop signal), False to clear it\n \"\"\"\n fence_key = _get_fence_key(chat_session_id)\n if not value:\n cache.delete(fence_key)\n return\n cache.set(fence_key, 0, ex=FENCE_TTL)\n\n\ndef is_connected(chat_session_id: UUID, cache: CacheBackend) -> bool:\n \"\"\"Check if the chat session should continue (not stopped).\n\n Args:\n chat_session_id: The UUID of the chat session to check\n cache: Tenant-aware cache backend\n\n Returns:\n True if the session should continue, False if it should stop\n \"\"\"\n return not cache.exists(_get_fence_key(chat_session_id))\n\n\ndef reset_cancel_status(chat_session_id: UUID, cache: CacheBackend) -> None:\n \"\"\"Clear the stop signal for a chat session.\n\n Args:\n chat_session_id: The UUID of the chat session\n cache: Tenant-aware cache backend\n \"\"\"\n cache.delete(_get_fence_key(chat_session_id))\n" }, { "path": "backend/onyx/chat/tool_call_args_streaming.py", "content": "from collections.abc import Generator\nfrom collections.abc import Mapping\nfrom typing import Any\nfrom typing import Type\n\nfrom onyx.llm.model_response import ChatCompletionDeltaToolCall\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta\nfrom onyx.tools.built_in_tools import TOOL_NAME_TO_CLASS\nfrom onyx.tools.interface import Tool\nfrom onyx.utils.jsonriver import Parser\n\n\ndef _get_tool_class(\n tool_calls_in_progress: Mapping[int, Mapping[str, Any]],\n tool_call_delta: ChatCompletionDeltaToolCall,\n) -> Type[Tool] | None:\n \"\"\"Look up the Tool subclass for a streaming tool call delta.\"\"\"\n tool_name = tool_calls_in_progress.get(tool_call_delta.index, {}).get(\"name\")\n if not tool_name:\n return None\n return TOOL_NAME_TO_CLASS.get(tool_name)\n\n\ndef maybe_emit_argument_delta(\n tool_calls_in_progress: Mapping[int, Mapping[str, Any]],\n tool_call_delta: ChatCompletionDeltaToolCall,\n placement: Placement,\n parsers: dict[int, Parser],\n) -> Generator[Packet, None, None]:\n \"\"\"Emit decoded tool-call argument deltas to the frontend.\n\n Uses a ``jsonriver.Parser`` per tool-call index to incrementally parse\n the JSON argument string and extract only the newly-appended content\n for each string-valued argument.\n\n NOTE: Non-string arguments (numbers, booleans, null, arrays, objects)\n are skipped — they are available in the final tool-call kickoff packet.\n\n ``parsers`` is a mutable dict keyed by tool-call index. A new\n ``Parser`` is created automatically for each new index.\n \"\"\"\n tool_cls = _get_tool_class(tool_calls_in_progress, tool_call_delta)\n if not tool_cls or not tool_cls.should_emit_argument_deltas():\n return\n\n fn = tool_call_delta.function\n delta_fragment = fn.arguments if fn else None\n if not delta_fragment:\n return\n\n idx = tool_call_delta.index\n if idx not in parsers:\n parsers[idx] = Parser()\n parser = parsers[idx]\n\n deltas = parser.feed(delta_fragment)\n\n argument_deltas: dict[str, str] = {}\n for delta in deltas:\n if isinstance(delta, dict):\n for key, value in delta.items():\n if isinstance(value, str):\n argument_deltas[key] = argument_deltas.get(key, \"\") + value\n\n if not argument_deltas:\n return\n\n tc_data = tool_calls_in_progress[tool_call_delta.index]\n yield Packet(\n placement=placement,\n obj=ToolCallArgumentDelta(\n tool_type=tc_data.get(\"name\", \"\"),\n argument_deltas=argument_deltas,\n ),\n )\n" }, { "path": "backend/onyx/configs/__init__.py", "content": "" }, { "path": "backend/onyx/configs/agent_configs.py", "content": "import os\n\n\nAGENT_DEFAULT_RETRIEVAL_HITS = 15\nAGENT_DEFAULT_RERANKING_HITS = 10\nAGENT_DEFAULT_SUB_QUESTION_MAX_CONTEXT_HITS = 8\nAGENT_DEFAULT_NUM_DOCS_FOR_INITIAL_DECOMPOSITION = 3\nAGENT_DEFAULT_NUM_DOCS_FOR_REFINED_DECOMPOSITION = 5\n\nAGENT_DEFAULT_MAX_STREAMED_DOCS_FOR_INITIAL_ANSWER = 25\nAGENT_DEFAULT_MAX_STREAMED_DOCS_FOR_REFINED_ANSWER = 35\n\n\nAGENT_DEFAULT_EXPLORATORY_SEARCH_RESULTS = 5\nAGENT_DEFAULT_MIN_ORIG_QUESTION_DOCS = 3\nAGENT_DEFAULT_MAX_ANSWER_CONTEXT_DOCS = 10\nAGENT_DEFAULT_MAX_STATIC_HISTORY_WORD_LENGTH = 2000\n\nINITIAL_SEARCH_DECOMPOSITION_ENABLED = True\n\nAGENT_DEFAULT_RETRIEVAL_HITS = 15\nAGENT_DEFAULT_RERANKING_HITS = 10\nAGENT_DEFAULT_MAX_VERIFIVATION_HITS = 30\nAGENT_DEFAULT_SUB_QUESTION_MAX_CONTEXT_HITS = 8\nAGENT_DEFAULT_NUM_DOCS_FOR_INITIAL_DECOMPOSITION = 3\nAGENT_DEFAULT_NUM_DOCS_FOR_REFINED_DECOMPOSITION = 5\nAGENT_DEFAULT_EXPLORATORY_SEARCH_RESULTS = 5\nAGENT_DEFAULT_MIN_ORIG_QUESTION_DOCS = 3\nAGENT_DEFAULT_MAX_ANSWER_CONTEXT_DOCS = 10\nAGENT_DEFAULT_MAX_STATIC_HISTORY_WORD_LENGTH = 2000\n\n\nAGENT_ALLOW_REFINEMENT = os.environ.get(\"AGENT_ALLOW_REFINEMENT\", \"\").lower() == \"true\"\n\nAGENT_ANSWER_GENERATION_BY_FAST_LLM = (\n os.environ.get(\"AGENT_ANSWER_GENERATION_BY_FAST_LLM\", \"\").lower() == \"true\"\n)\n\nAGENT_RETRIEVAL_STATS = (\n not os.environ.get(\"AGENT_RETRIEVAL_STATS\") == \"False\"\n) or True # default True\n\n\nAGENT_MAX_VERIFICATION_HITS = int(\n os.environ.get(\"AGENT_MAX_VERIFICATION_HITS\") or AGENT_DEFAULT_MAX_VERIFIVATION_HITS\n) # 30\n\nAGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_MAX_QUERY_RETRIEVAL_RESULTS\") or AGENT_DEFAULT_RETRIEVAL_HITS\n) # 15\n\nAGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_MAX_QUERY_RETRIEVAL_RESULTS\") or AGENT_DEFAULT_RETRIEVAL_HITS\n) # 15\n\n# Reranking agent configs\n# Reranking stats - no influence on flow outside of stats collection\nAGENT_RERANKING_STATS = (\n not os.environ.get(\"AGENT_RERANKING_STATS\") == \"True\"\n) or False # default False\n\nAGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_MAX_QUERY_RETRIEVAL_RESULTS\") or AGENT_DEFAULT_RETRIEVAL_HITS\n) # 15\n\nAGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS\")\n or AGENT_DEFAULT_RERANKING_HITS\n) # 10\n\nAGENT_NUM_DOCS_FOR_DECOMPOSITION = int(\n os.environ.get(\"AGENT_NUM_DOCS_FOR_DECOMPOSITION\")\n or AGENT_DEFAULT_NUM_DOCS_FOR_INITIAL_DECOMPOSITION\n) # 3\n\nAGENT_NUM_DOCS_FOR_REFINED_DECOMPOSITION = int(\n os.environ.get(\"AGENT_NUM_DOCS_FOR_REFINED_DECOMPOSITION\")\n or AGENT_DEFAULT_NUM_DOCS_FOR_REFINED_DECOMPOSITION\n) # 5\n\nAGENT_EXPLORATORY_SEARCH_RESULTS = int(\n os.environ.get(\"AGENT_EXPLORATORY_SEARCH_RESULTS\")\n or AGENT_DEFAULT_EXPLORATORY_SEARCH_RESULTS\n) # 5\n\nAGENT_MIN_ORIG_QUESTION_DOCS = int(\n os.environ.get(\"AGENT_MIN_ORIG_QUESTION_DOCS\")\n or AGENT_DEFAULT_MIN_ORIG_QUESTION_DOCS\n) # 3\n\nAGENT_MAX_ANSWER_CONTEXT_DOCS = int(\n os.environ.get(\"AGENT_MAX_ANSWER_CONTEXT_DOCS\")\n or AGENT_DEFAULT_SUB_QUESTION_MAX_CONTEXT_HITS\n) # 8\n\n\nAGENT_MAX_STATIC_HISTORY_WORD_LENGTH = int(\n os.environ.get(\"AGENT_MAX_STATIC_HISTORY_WORD_LENGTH\")\n or AGENT_DEFAULT_MAX_STATIC_HISTORY_WORD_LENGTH\n) # 2000\n\nAGENT_MAX_STREAMED_DOCS_FOR_INITIAL_ANSWER = int(\n os.environ.get(\"AGENT_MAX_STREAMED_DOCS_FOR_INITIAL_ANSWER\")\n or AGENT_DEFAULT_MAX_STREAMED_DOCS_FOR_INITIAL_ANSWER\n) # 25\n\nAGENT_MAX_STREAMED_DOCS_FOR_REFINED_ANSWER = int(\n os.environ.get(\"AGENT_MAX_STREAMED_DOCS_FOR_REFINED_ANSWER\")\n or AGENT_DEFAULT_MAX_STREAMED_DOCS_FOR_REFINED_ANSWER\n) # 35\n\n\nAGENT_RETRIEVAL_STATS = (\n not os.environ.get(\"AGENT_RETRIEVAL_STATS\") == \"False\"\n) or True # default True\n\n\nAGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_MAX_QUERY_RETRIEVAL_RESULTS\") or AGENT_DEFAULT_RETRIEVAL_HITS\n) # 15\n\nAGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_MAX_QUERY_RETRIEVAL_RESULTS\") or AGENT_DEFAULT_RETRIEVAL_HITS\n) # 15\n\n# Reranking agent configs\n# Reranking stats - no influence on flow outside of stats collection\nAGENT_RERANKING_STATS = (\n not os.environ.get(\"AGENT_RERANKING_STATS\") == \"True\"\n) or False # default False\n\nAGENT_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_MAX_QUERY_RETRIEVAL_RESULTS\") or AGENT_DEFAULT_RETRIEVAL_HITS\n) # 15\n\nAGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS = int(\n os.environ.get(\"AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS\")\n or AGENT_DEFAULT_RERANKING_HITS\n) # 10\n\nAGENT_NUM_DOCS_FOR_DECOMPOSITION = int(\n os.environ.get(\"AGENT_NUM_DOCS_FOR_DECOMPOSITION\")\n or AGENT_DEFAULT_NUM_DOCS_FOR_INITIAL_DECOMPOSITION\n) # 3\n\nAGENT_NUM_DOCS_FOR_REFINED_DECOMPOSITION = int(\n os.environ.get(\"AGENT_NUM_DOCS_FOR_REFINED_DECOMPOSITION\")\n or AGENT_DEFAULT_NUM_DOCS_FOR_REFINED_DECOMPOSITION\n) # 5\n\nAGENT_EXPLORATORY_SEARCH_RESULTS = int(\n os.environ.get(\"AGENT_EXPLORATORY_SEARCH_RESULTS\")\n or AGENT_DEFAULT_EXPLORATORY_SEARCH_RESULTS\n) # 5\n\nAGENT_MIN_ORIG_QUESTION_DOCS = int(\n os.environ.get(\"AGENT_MIN_ORIG_QUESTION_DOCS\")\n or AGENT_DEFAULT_MIN_ORIG_QUESTION_DOCS\n) # 3\n\nAGENT_MAX_ANSWER_CONTEXT_DOCS = int(\n os.environ.get(\"AGENT_MAX_ANSWER_CONTEXT_DOCS\")\n or AGENT_DEFAULT_SUB_QUESTION_MAX_CONTEXT_HITS\n) # 8\n\n\nAGENT_MAX_STATIC_HISTORY_WORD_LENGTH = int(\n os.environ.get(\"AGENT_MAX_STATIC_HISTORY_WORD_LENGTH\")\n or AGENT_DEFAULT_MAX_STATIC_HISTORY_WORD_LENGTH\n) # 2000\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_ENTITY_TERM_EXTRACTION = 15 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_ENTITY_TERM_EXTRACTION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_ENTITY_TERM_EXTRACTION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_ENTITY_TERM_EXTRACTION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_ENTITY_TERM_EXTRACTION = 45 # in seconds\nAGENT_TIMEOUT_LLM_ENTITY_TERM_EXTRACTION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_ENTITY_TERM_EXTRACTION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_ENTITY_TERM_EXTRACTION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_DOCUMENT_VERIFICATION = 5 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_DOCUMENT_VERIFICATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_DOCUMENT_VERIFICATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_DOCUMENT_VERIFICATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_DOCUMENT_VERIFICATION = 8 # in seconds\nAGENT_TIMEOUT_LLM_DOCUMENT_VERIFICATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_DOCUMENT_VERIFICATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_DOCUMENT_VERIFICATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_GENERAL_GENERATION = 8 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_GENERAL_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_GENERAL_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_GENERAL_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_GENERAL_GENERATION = 45 # in seconds\nAGENT_TIMEOUT_LLM_GENERAL_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_GENERAL_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_GENERAL_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_SUBQUESTION_GENERATION = 8 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_SUBQUESTION_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_SUBQUESTION_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_SUBQUESTION_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_SUBQUESTION_GENERATION = 10 # in seconds\nAGENT_TIMEOUT_LLM_SUBQUESTION_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_SUBQUESTION_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_SUBQUESTION_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_SUBANSWER_GENERATION = 9 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_SUBANSWER_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_SUBANSWER_GENERATION = 45 # in seconds\nAGENT_TIMEOUT_LLM_SUBANSWER_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_SUBANSWER_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_SUBANSWER_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION = 15 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION = 40 # in seconds\nAGENT_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_GENERATION = 20 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_REFINED_ANSWER_GENERATION = 60 # in seconds\nAGENT_TIMEOUT_LLM_REFINED_ANSWER_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_REFINED_ANSWER_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_REFINED_ANSWER_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_SUBANSWER_CHECK = 6 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_CHECK = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_CHECK\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_SUBANSWER_CHECK\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_SUBANSWER_CHECK = 12 # in seconds\nAGENT_TIMEOUT_LLM_SUBANSWER_CHECK = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_SUBANSWER_CHECK\")\n or AGENT_DEFAULT_TIMEOUT_LLM_SUBANSWER_CHECK\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_REFINED_SUBQUESTION_GENERATION = 6 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_REFINED_SUBQUESTION_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_REFINED_SUBQUESTION_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_REFINED_SUBQUESTION_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_REFINED_SUBQUESTION_GENERATION = 12 # in seconds\nAGENT_TIMEOUT_LLM_REFINED_SUBQUESTION_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_REFINED_SUBQUESTION_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_REFINED_SUBQUESTION_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_QUERY_REWRITING_GENERATION = 4 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_QUERY_REWRITING_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_QUERY_REWRITING_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_QUERY_REWRITING_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_QUERY_REWRITING_GENERATION = 6 # in seconds\nAGENT_TIMEOUT_LLM_QUERY_REWRITING_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_QUERY_REWRITING_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_QUERY_REWRITING_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_HISTORY_SUMMARY_GENERATION = 6 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_HISTORY_SUMMARY_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_HISTORY_SUMMARY_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_HISTORY_SUMMARY_GENERATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_HISTORY_SUMMARY_GENERATION = 8 # in seconds\nAGENT_TIMEOUT_LLM_HISTORY_SUMMARY_GENERATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_HISTORY_SUMMARY_GENERATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_HISTORY_SUMMARY_GENERATION\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_COMPARE_ANSWERS = 6 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_COMPARE_ANSWERS = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_COMPARE_ANSWERS\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_COMPARE_ANSWERS\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_COMPARE_ANSWERS = 12 # in seconds\nAGENT_TIMEOUT_LLM_COMPARE_ANSWERS = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_COMPARE_ANSWERS\")\n or AGENT_DEFAULT_TIMEOUT_LLM_COMPARE_ANSWERS\n)\n\n\nAGENT_DEFAULT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_VALIDATION = 6 # in seconds\nAGENT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_VALIDATION = int(\n os.environ.get(\"AGENT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_VALIDATION\")\n or AGENT_DEFAULT_TIMEOUT_CONNECT_LLM_REFINED_ANSWER_VALIDATION\n)\n\nAGENT_DEFAULT_TIMEOUT_LLM_REFINED_ANSWER_VALIDATION = 12 # in seconds\nAGENT_TIMEOUT_LLM_REFINED_ANSWER_VALIDATION = int(\n os.environ.get(\"AGENT_TIMEOUT_LLM_REFINED_ANSWER_VALIDATION\")\n or AGENT_DEFAULT_TIMEOUT_LLM_REFINED_ANSWER_VALIDATION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_VALIDATION = 4\nAGENT_MAX_TOKENS_VALIDATION = int(\n os.environ.get(\"AGENT_MAX_TOKENS_VALIDATION\") or AGENT_DEFAULT_MAX_TOKENS_VALIDATION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_SUBANSWER_GENERATION = 256\nAGENT_MAX_TOKENS_SUBANSWER_GENERATION = int(\n os.environ.get(\"AGENT_MAX_TOKENS_SUBANSWER_GENERATION\")\n or AGENT_DEFAULT_MAX_TOKENS_SUBANSWER_GENERATION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_ANSWER_GENERATION = 1024\nAGENT_MAX_TOKENS_ANSWER_GENERATION = int(\n os.environ.get(\"AGENT_MAX_TOKENS_ANSWER_GENERATION\")\n or AGENT_DEFAULT_MAX_TOKENS_ANSWER_GENERATION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_SUBQUESTION_GENERATION = 256\nAGENT_MAX_TOKENS_SUBQUESTION_GENERATION = int(\n os.environ.get(\"AGENT_MAX_TOKENS_SUBQUESTION_GENERATION\")\n or AGENT_DEFAULT_MAX_TOKENS_SUBQUESTION_GENERATION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_ENTITY_TERM_EXTRACTION = 1024\nAGENT_MAX_TOKENS_ENTITY_TERM_EXTRACTION = int(\n os.environ.get(\"AGENT_MAX_TOKENS_ENTITY_TERM_EXTRACTION\")\n or AGENT_DEFAULT_MAX_TOKENS_ENTITY_TERM_EXTRACTION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_SUBQUERY_GENERATION = 64\nAGENT_MAX_TOKENS_SUBQUERY_GENERATION = int(\n os.environ.get(\"AGENT_MAX_TOKENS_SUBQUERY_GENERATION\")\n or AGENT_DEFAULT_MAX_TOKENS_SUBQUERY_GENERATION\n)\n\nAGENT_DEFAULT_MAX_TOKENS_HISTORY_SUMMARY = 128\nAGENT_MAX_TOKENS_HISTORY_SUMMARY = int(\n os.environ.get(\"AGENT_MAX_TOKENS_HISTORY_SUMMARY\")\n or AGENT_DEFAULT_MAX_TOKENS_HISTORY_SUMMARY\n)\n\n# Parameters for the Thoughtful/Deep Research flows\nTF_DR_TIMEOUT_LONG = int(os.environ.get(\"TF_DR_TIMEOUT_LONG\") or 120)\nTF_DR_TIMEOUT_SHORT = int(os.environ.get(\"TF_DR_TIMEOUT_SHORT\") or 60)\n\n\nTF_DR_DEFAULT_FAST = (os.environ.get(\"TF_DR_DEFAULT_FAST\") or \"False\").lower() == \"true\"\n\nGRAPH_VERSION_NAME: str = \"a\"\n" }, { "path": "backend/onyx/configs/app_configs.py", "content": "import json\nimport os\nimport urllib.parse\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import cast\n\nfrom onyx.auth.schemas import AuthBackend\nfrom onyx.cache.interface import CacheBackendType\nfrom onyx.configs.constants import AuthType\nfrom onyx.configs.constants import QueryHistoryType\nfrom onyx.file_processing.enums import HtmlBasedConnectorTransformLinksStrategy\nfrom onyx.prompts.image_analysis import DEFAULT_IMAGE_SUMMARIZATION_SYSTEM_PROMPT\nfrom onyx.prompts.image_analysis import DEFAULT_IMAGE_SUMMARIZATION_USER_PROMPT\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n#####\n# App Configs\n#####\nAPP_HOST = \"0.0.0.0\"\nAPP_PORT = 8080\n# API_PREFIX is used to prepend a base path for all API routes\n# generally used if using a reverse proxy which doesn't support stripping the `/api`\n# prefix from requests directed towards the API server. In these cases, set this to `/api`\nAPP_API_PREFIX = os.environ.get(\"API_PREFIX\", \"\")\n\n# Certain services need to make HTTP requests to the API server, such as the MCP server and Discord bot\nAPI_SERVER_PROTOCOL = os.environ.get(\"API_SERVER_PROTOCOL\", \"http\")\nAPI_SERVER_HOST = os.environ.get(\"API_SERVER_HOST\", \"127.0.0.1\")\n# This override allows self-hosting the MCP server with Onyx Cloud backend.\nAPI_SERVER_URL_OVERRIDE_FOR_HTTP_REQUESTS = os.environ.get(\n \"API_SERVER_URL_OVERRIDE_FOR_HTTP_REQUESTS\"\n)\n\n# Whether to send user metadata (user_id/email and session_id) to the LLM provider.\n# Disabled by default.\nSEND_USER_METADATA_TO_LLM_PROVIDER = (\n os.environ.get(\"SEND_USER_METADATA_TO_LLM_PROVIDER\", \"\")\n).lower() == \"true\"\n\n#####\n# User Facing Features Configs\n#####\nBLURB_SIZE = 128 # Number Encoder Tokens included in the chunk blurb\n\n# Hard ceiling for the admin-configurable file upload size (in MB).\n# Self-hosted customers can raise or lower this via the environment variable.\n_raw_max_upload_size_mb = int(os.environ.get(\"MAX_ALLOWED_UPLOAD_SIZE_MB\", \"250\"))\nif _raw_max_upload_size_mb < 0:\n logger.warning(\n \"MAX_ALLOWED_UPLOAD_SIZE_MB=%d is negative; falling back to 250\",\n _raw_max_upload_size_mb,\n )\n _raw_max_upload_size_mb = 250\nMAX_ALLOWED_UPLOAD_SIZE_MB = _raw_max_upload_size_mb\n\n# Default fallback for the per-user file upload size limit (in MB) when no\n# admin-configured value exists. Clamped to MAX_ALLOWED_UPLOAD_SIZE_MB at\n# runtime so this never silently exceeds the hard ceiling.\n_raw_default_upload_size_mb = int(\n os.environ.get(\"DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB\", \"100\")\n)\nif _raw_default_upload_size_mb < 0:\n logger.warning(\n \"DEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB=%d is negative; falling back to 100\",\n _raw_default_upload_size_mb,\n )\n _raw_default_upload_size_mb = 100\nDEFAULT_USER_FILE_MAX_UPLOAD_SIZE_MB = _raw_default_upload_size_mb\nGENERATIVE_MODEL_ACCESS_CHECK_FREQ = int(\n os.environ.get(\"GENERATIVE_MODEL_ACCESS_CHECK_FREQ\") or 86400\n) # 1 day\n\n# Controls whether users can use User Knowledge (personal documents) in assistants\nDISABLE_USER_KNOWLEDGE = os.environ.get(\"DISABLE_USER_KNOWLEDGE\", \"\").lower() == \"true\"\n\n# Disables vector DB (Vespa/OpenSearch) entirely. When True, connectors and RAG search\n# are disabled but core chat, tools, user file uploads, and Projects still work.\nDISABLE_VECTOR_DB = os.environ.get(\"DISABLE_VECTOR_DB\", \"\").lower() == \"true\"\n\n# Which backend to use for caching, locks, and ephemeral state.\n# \"redis\" (default) or \"postgres\" (only valid when DISABLE_VECTOR_DB=true).\nCACHE_BACKEND = CacheBackendType(\n os.environ.get(\"CACHE_BACKEND\", CacheBackendType.REDIS)\n)\n\n# If set to true, will show extra/uncommon connectors in the \"Other\" category\nSHOW_EXTRA_CONNECTORS = os.environ.get(\"SHOW_EXTRA_CONNECTORS\", \"\").lower() == \"true\"\n\n# Controls whether to allow admin query history reports with:\n# 1. associated user emails\n# 2. anonymized user emails\n# 3. no queries\nONYX_QUERY_HISTORY_TYPE = QueryHistoryType(\n (os.environ.get(\"ONYX_QUERY_HISTORY_TYPE\") or QueryHistoryType.NORMAL.value).lower()\n)\n\n#####\n# Web Configs\n#####\n# WEB_DOMAIN is used to set the redirect_uri after login flows\n# NOTE: if you are having problems accessing the Onyx web UI locally (especially\n# on Windows, try setting this to `http://127.0.0.1:3000` instead and see if that\n# fixes it)\nWEB_DOMAIN = os.environ.get(\"WEB_DOMAIN\") or \"http://localhost:3000\"\n\n\n#####\n# Auth Configs\n#####\n# Silently default to basic - warnings/errors logged in verify_auth_setting()\n# which only runs on app startup, not during migrations/scripts\n_auth_type_str = (os.environ.get(\"AUTH_TYPE\") or \"\").lower()\nif _auth_type_str in [auth_type.value for auth_type in AuthType]:\n AUTH_TYPE = AuthType(_auth_type_str)\nelse:\n AUTH_TYPE = AuthType.BASIC\n\nPASSWORD_MIN_LENGTH = int(os.getenv(\"PASSWORD_MIN_LENGTH\", 8))\nPASSWORD_MAX_LENGTH = int(os.getenv(\"PASSWORD_MAX_LENGTH\", 64))\nPASSWORD_REQUIRE_UPPERCASE = (\n os.environ.get(\"PASSWORD_REQUIRE_UPPERCASE\", \"false\").lower() == \"true\"\n)\nPASSWORD_REQUIRE_LOWERCASE = (\n os.environ.get(\"PASSWORD_REQUIRE_LOWERCASE\", \"false\").lower() == \"true\"\n)\nPASSWORD_REQUIRE_DIGIT = (\n os.environ.get(\"PASSWORD_REQUIRE_DIGIT\", \"false\").lower() == \"true\"\n)\nPASSWORD_REQUIRE_SPECIAL_CHAR = (\n os.environ.get(\"PASSWORD_REQUIRE_SPECIAL_CHAR\", \"false\").lower() == \"true\"\n)\n\n# Encryption key secret is used to encrypt connector credentials, api keys, and other sensitive\n# information. This provides an extra layer of security on top of Postgres access controls\n# and is available in Onyx EE\nENCRYPTION_KEY_SECRET = os.environ.get(\"ENCRYPTION_KEY_SECRET\") or \"\"\n\n# Turn off mask if admin users should see full credentials for data connectors.\nMASK_CREDENTIAL_PREFIX = (\n os.environ.get(\"MASK_CREDENTIAL_PREFIX\", \"True\").lower() != \"false\"\n)\n\nAUTH_BACKEND = AuthBackend(os.environ.get(\"AUTH_BACKEND\") or AuthBackend.REDIS.value)\n\nSESSION_EXPIRE_TIME_SECONDS = int(\n os.environ.get(\"SESSION_EXPIRE_TIME_SECONDS\")\n or os.environ.get(\"REDIS_AUTH_EXPIRE_TIME_SECONDS\")\n or 86400 * 7\n) # 7 days\n\n# Default request timeout, mostly used by connectors\nREQUEST_TIMEOUT_SECONDS = int(os.environ.get(\"REQUEST_TIMEOUT_SECONDS\") or 60)\n\n# set `VALID_EMAIL_DOMAINS` to a comma seperated list of domains in order to\n# restrict access to Onyx to only users with emails from those domains.\n# E.g. `VALID_EMAIL_DOMAINS=example.com,example.org` will restrict Onyx\n# signups to users with either an @example.com or an @example.org email.\n# NOTE: maintaining `VALID_EMAIL_DOMAIN` to keep backwards compatibility\n_VALID_EMAIL_DOMAIN = os.environ.get(\"VALID_EMAIL_DOMAIN\", \"\")\n_VALID_EMAIL_DOMAINS_STR = (\n os.environ.get(\"VALID_EMAIL_DOMAINS\", \"\") or _VALID_EMAIL_DOMAIN\n)\nVALID_EMAIL_DOMAINS = (\n [\n domain.strip().lower()\n for domain in _VALID_EMAIL_DOMAINS_STR.split(\",\")\n if domain.strip()\n ]\n if _VALID_EMAIL_DOMAINS_STR\n else []\n)\n\n# Disposable email blocking - blocks temporary/throwaway email addresses\n# Set to empty string to disable disposable email blocking\nDISPOSABLE_EMAIL_DOMAINS_URL = os.environ.get(\n \"DISPOSABLE_EMAIL_DOMAINS_URL\",\n \"https://disposable.github.io/disposable-email-domains/domains.json\",\n)\n\n# OAuth Login Flow\n# Used for both Google OAuth2 and OIDC flows\nOAUTH_CLIENT_ID = (\n os.environ.get(\"OAUTH_CLIENT_ID\", os.environ.get(\"GOOGLE_OAUTH_CLIENT_ID\")) or \"\"\n)\nOAUTH_CLIENT_SECRET = (\n os.environ.get(\"OAUTH_CLIENT_SECRET\", os.environ.get(\"GOOGLE_OAUTH_CLIENT_SECRET\"))\n or \"\"\n)\n\n# Whether Google OAuth is enabled (requires both client ID and secret)\nOAUTH_ENABLED = bool(OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET)\n\n# OpenID Connect configuration URL for OIDC integrations\nOPENID_CONFIG_URL = os.environ.get(\"OPENID_CONFIG_URL\") or \"\"\n\n# Applicable for OIDC Auth, allows you to override the scopes that\n# are requested from the OIDC provider. Currently used when passing\n# over access tokens to tool calls and the tool needs more scopes\nOIDC_SCOPE_OVERRIDE: list[str] | None = None\n_OIDC_SCOPE_OVERRIDE = os.environ.get(\"OIDC_SCOPE_OVERRIDE\")\n\nif _OIDC_SCOPE_OVERRIDE:\n try:\n OIDC_SCOPE_OVERRIDE = [\n scope.strip() for scope in _OIDC_SCOPE_OVERRIDE.split(\",\")\n ]\n except Exception:\n pass\n\n# Enables PKCE for OIDC login flow. Disabled by default to preserve\n# backwards compatibility for existing OIDC deployments.\nOIDC_PKCE_ENABLED = os.environ.get(\"OIDC_PKCE_ENABLED\", \"\").lower() == \"true\"\n\n# Applicable for SAML Auth\nSAML_CONF_DIR = os.environ.get(\"SAML_CONF_DIR\") or \"/app/onyx/configs/saml_config\"\n\n# JWT Public Key URL for JWT token verification\nJWT_PUBLIC_KEY_URL: str | None = os.getenv(\"JWT_PUBLIC_KEY_URL\", None)\n\nUSER_AUTH_SECRET = os.environ.get(\"USER_AUTH_SECRET\", \"\")\n\nif AUTH_TYPE == AuthType.BASIC and not USER_AUTH_SECRET:\n logger.warning(\n \"USER_AUTH_SECRET is not set. This is required for secure password reset \"\n \"and email verification tokens. Please set USER_AUTH_SECRET in production.\"\n )\n\n# Duration (in seconds) for which the FastAPI Users JWT token remains valid in the user's browser.\n# By default, this is set to match the Redis expiry time for consistency.\nAUTH_COOKIE_EXPIRE_TIME_SECONDS = int(\n os.environ.get(\"AUTH_COOKIE_EXPIRE_TIME_SECONDS\") or 86400 * 7\n) # 7 days\n\n# for basic auth\nREQUIRE_EMAIL_VERIFICATION = (\n os.environ.get(\"REQUIRE_EMAIL_VERIFICATION\", \"\").lower() == \"true\"\n)\nSMTP_SERVER = os.environ.get(\"SMTP_SERVER\") or \"\"\nSMTP_PORT = int(os.environ.get(\"SMTP_PORT\") or \"587\")\nSMTP_USER = os.environ.get(\"SMTP_USER\") or \"\"\nSMTP_PASS = os.environ.get(\"SMTP_PASS\") or \"\"\nEMAIL_FROM = os.environ.get(\"EMAIL_FROM\") or SMTP_USER\n\nSENDGRID_API_KEY = os.environ.get(\"SENDGRID_API_KEY\") or \"\"\nEMAIL_CONFIGURED = all([SMTP_SERVER, SMTP_USER, SMTP_PASS]) or SENDGRID_API_KEY\n\n# If set, Onyx will listen to the `expires_at` returned by the identity\n# provider (e.g. Okta, Google, etc.) and force the user to re-authenticate\n# after this time has elapsed. Disabled since by default many auth providers\n# have very short expiry times (e.g. 1 hour) which provide a poor user experience\nTRACK_EXTERNAL_IDP_EXPIRY = (\n os.environ.get(\"TRACK_EXTERNAL_IDP_EXPIRY\", \"\").lower() == \"true\"\n)\n\n\n#####\n# DB Configs\n#####\nDOCUMENT_INDEX_NAME = \"danswer_index\"\n\n# OpenSearch Configs\nOPENSEARCH_HOST = os.environ.get(\"OPENSEARCH_HOST\") or \"localhost\"\nOPENSEARCH_REST_API_PORT = int(os.environ.get(\"OPENSEARCH_REST_API_PORT\") or 9200)\n# TODO(andrei): 60 seconds is too much, we're just setting a high default\n# timeout for now to examine why queries are slow.\n# NOTE: This timeout applies to all requests the client makes, including bulk\n# indexing.\nDEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S = int(\n os.environ.get(\"DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S\") or 60\n)\n# TODO(andrei): 50 seconds is too much, we're just setting a high default\n# timeout for now to examine why queries are slow.\n# NOTE: To get useful partial results, this value should be less than the client\n# timeout above.\nDEFAULT_OPENSEARCH_QUERY_TIMEOUT_S = int(\n os.environ.get(\"DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S\") or 50\n)\nOPENSEARCH_ADMIN_USERNAME = os.environ.get(\"OPENSEARCH_ADMIN_USERNAME\", \"admin\")\nOPENSEARCH_ADMIN_PASSWORD = os.environ.get(\n \"OPENSEARCH_ADMIN_PASSWORD\", \"StrongPassword123!\"\n)\nUSING_AWS_MANAGED_OPENSEARCH = (\n os.environ.get(\"USING_AWS_MANAGED_OPENSEARCH\", \"\").lower() == \"true\"\n)\n# Profiling adds some overhead to OpenSearch operations. This overhead is\n# unknown right now. Defaults to True.\nOPENSEARCH_PROFILING_DISABLED = (\n os.environ.get(\"OPENSEARCH_PROFILING_DISABLED\", \"true\").lower() == \"true\"\n)\n# Whether to disable match highlights for OpenSearch. Defaults to True for now\n# as we investigate query performance.\nOPENSEARCH_MATCH_HIGHLIGHTS_DISABLED = (\n os.environ.get(\"OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED\", \"true\").lower() == \"true\"\n)\n# When enabled, OpenSearch returns detailed score breakdowns for each hit.\n# Useful for debugging and tuning search relevance. Has ~10-30% performance overhead according to documentation.\n# Seems for Hybrid Search in practice, the impact is actually more like 1000x slower.\nOPENSEARCH_EXPLAIN_ENABLED = (\n os.environ.get(\"OPENSEARCH_EXPLAIN_ENABLED\", \"\").lower() == \"true\"\n)\n# Analyzer used for full-text fields (title, content). Use OpenSearch built-in analyzer\n# names (e.g. \"english\", \"standard\", \"german\"). Affects stemming and tokenization;\n# existing indices need reindexing after a change.\nOPENSEARCH_TEXT_ANALYZER = os.environ.get(\"OPENSEARCH_TEXT_ANALYZER\") or \"english\"\n\n# This is the \"base\" config for now, the idea is that at least for our dev\n# environments we always want to be dual indexing into both OpenSearch and Vespa\n# to stress test the new codepaths. Only enable this if there is some instance\n# of OpenSearch running for the relevant Onyx instance.\n# NOTE: Now enabled on by default, unless the env indicates otherwise.\nENABLE_OPENSEARCH_INDEXING_FOR_ONYX = (\n os.environ.get(\"ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\", \"true\").lower() == \"true\"\n)\n# NOTE: This effectively does nothing anymore, admins can now toggle whether\n# retrieval is through OpenSearch. This value is only used as a final fallback\n# in case that doesn't work for whatever reason.\n# Given that the \"base\" config above is true, this enables whether we want to\n# retrieve from OpenSearch or Vespa. We want to be able to quickly toggle this\n# in the event we see issues with OpenSearch retrieval in our dev environments.\nENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX = (\n ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\n and os.environ.get(\"ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX\", \"\").lower() == \"true\"\n)\n# Whether we should check for and create an index if necessary every time we\n# instantiate an OpenSearchDocumentIndex on multitenant cloud. Defaults to True.\nVERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT = (\n os.environ.get(\"VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT\", \"true\").lower()\n == \"true\"\n)\nOPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE = int(\n os.environ.get(\"OPENSEARCH_MIGRATION_GET_VESPA_CHUNKS_PAGE_SIZE\") or 500\n)\n# If set, will override the default number of shards and replicas for the index.\nOPENSEARCH_INDEX_NUM_SHARDS: int | None = (\n int(os.environ[\"OPENSEARCH_INDEX_NUM_SHARDS\"])\n if os.environ.get(\"OPENSEARCH_INDEX_NUM_SHARDS\", None) is not None\n else None\n)\nOPENSEARCH_INDEX_NUM_REPLICAS: int | None = (\n int(os.environ[\"OPENSEARCH_INDEX_NUM_REPLICAS\"])\n if os.environ.get(\"OPENSEARCH_INDEX_NUM_REPLICAS\", None) is not None\n else None\n)\nONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH = (\n os.environ.get(\"ONYX_SEARCH_UI_USES_OPENSEARCH_KEYWORD_SEARCH\", \"\").lower()\n == \"true\"\n)\n\nVESPA_HOST = os.environ.get(\"VESPA_HOST\") or \"localhost\"\n# NOTE: this is used if and only if the vespa config server is accessible via a\n# different host than the main vespa application\nVESPA_CONFIG_SERVER_HOST = os.environ.get(\"VESPA_CONFIG_SERVER_HOST\") or VESPA_HOST\nVESPA_PORT = os.environ.get(\"VESPA_PORT\") or \"8081\"\nVESPA_TENANT_PORT = os.environ.get(\"VESPA_TENANT_PORT\") or \"19071\"\n# the number of times to try and connect to vespa on startup before giving up\nVESPA_NUM_ATTEMPTS_ON_STARTUP = int(os.environ.get(\"NUM_RETRIES_ON_STARTUP\") or 10)\n\nVESPA_CLOUD_URL = os.environ.get(\"VESPA_CLOUD_URL\", \"\")\n\nVESPA_CLOUD_CERT_PATH = os.environ.get(\"VESPA_CLOUD_CERT_PATH\")\nVESPA_CLOUD_KEY_PATH = os.environ.get(\"VESPA_CLOUD_KEY_PATH\")\n\n# Number of documents in a batch during indexing (further batching done by chunks before passing to bi-encoder)\nINDEX_BATCH_SIZE = int(os.environ.get(\"INDEX_BATCH_SIZE\") or 16)\n\nMAX_DRIVE_WORKERS = int(os.environ.get(\"MAX_DRIVE_WORKERS\", 4))\n\n# Below are intended to match the env variables names used by the official postgres docker image\n# https://hub.docker.com/_/postgres\nPOSTGRES_USER = os.environ.get(\"POSTGRES_USER\") or \"postgres\"\n# URL-encode the password for asyncpg to avoid issues with special characters on some machines.\nPOSTGRES_PASSWORD = urllib.parse.quote_plus(\n os.environ.get(\"POSTGRES_PASSWORD\") or \"password\"\n)\nPOSTGRES_HOST = os.environ.get(\"POSTGRES_HOST\") or \"127.0.0.1\"\nPOSTGRES_PORT = os.environ.get(\"POSTGRES_PORT\") or \"5432\"\nPOSTGRES_DB = os.environ.get(\"POSTGRES_DB\") or \"postgres\"\nAWS_REGION_NAME = os.environ.get(\"AWS_REGION_NAME\") or \"us-east-2\"\n\nPOSTGRES_API_SERVER_POOL_SIZE = int(\n os.environ.get(\"POSTGRES_API_SERVER_POOL_SIZE\") or 40\n)\nPOSTGRES_API_SERVER_POOL_OVERFLOW = int(\n os.environ.get(\"POSTGRES_API_SERVER_POOL_OVERFLOW\") or 10\n)\n\nPOSTGRES_API_SERVER_READ_ONLY_POOL_SIZE = int(\n os.environ.get(\"POSTGRES_API_SERVER_READ_ONLY_POOL_SIZE\") or 10\n)\nPOSTGRES_API_SERVER_READ_ONLY_POOL_OVERFLOW = int(\n os.environ.get(\"POSTGRES_API_SERVER_READ_ONLY_POOL_OVERFLOW\") or 5\n)\n\n# defaults to False\n# generally should only be used for\nPOSTGRES_USE_NULL_POOL = os.environ.get(\"POSTGRES_USE_NULL_POOL\", \"\").lower() == \"true\"\n\n# defaults to False\nPOSTGRES_POOL_PRE_PING = os.environ.get(\"POSTGRES_POOL_PRE_PING\", \"\").lower() == \"true\"\n\n# recycle timeout in seconds\nPOSTGRES_POOL_RECYCLE_DEFAULT = 60 * 20 # 20 minutes\ntry:\n POSTGRES_POOL_RECYCLE = int(\n os.environ.get(\"POSTGRES_POOL_RECYCLE\", POSTGRES_POOL_RECYCLE_DEFAULT)\n )\nexcept ValueError:\n POSTGRES_POOL_RECYCLE = POSTGRES_POOL_RECYCLE_DEFAULT\n\n# RDS IAM authentication - enables IAM-based authentication for PostgreSQL\nUSE_IAM_AUTH = os.getenv(\"USE_IAM_AUTH\", \"False\").lower() == \"true\"\n\n# Redis IAM authentication - enables IAM-based authentication for Redis ElastiCache\n# Note: This is separate from RDS IAM auth as they use different authentication mechanisms\nUSE_REDIS_IAM_AUTH = os.getenv(\"USE_REDIS_IAM_AUTH\", \"False\").lower() == \"true\"\nREDIS_SSL = os.getenv(\"REDIS_SSL\", \"\").lower() == \"true\"\nREDIS_HOST = os.environ.get(\"REDIS_HOST\") or \"localhost\"\nREDIS_PORT = int(os.environ.get(\"REDIS_PORT\", 6379))\nREDIS_PASSWORD = os.environ.get(\"REDIS_PASSWORD\") or \"\"\n\n# this assumes that other redis settings remain the same as the primary\nREDIS_REPLICA_HOST = os.environ.get(\"REDIS_REPLICA_HOST\") or REDIS_HOST\n\nREDIS_AUTH_KEY_PREFIX = \"fastapi_users_token:\"\n\n# Rate limiting for auth endpoints\nRATE_LIMIT_WINDOW_SECONDS: int | None = None\n_rate_limit_window_seconds_str = os.environ.get(\"RATE_LIMIT_WINDOW_SECONDS\")\nif _rate_limit_window_seconds_str is not None:\n try:\n RATE_LIMIT_WINDOW_SECONDS = int(_rate_limit_window_seconds_str)\n except ValueError:\n pass\n\nRATE_LIMIT_MAX_REQUESTS: int | None = None\n_rate_limit_max_requests_str = os.environ.get(\"RATE_LIMIT_MAX_REQUESTS\")\nif _rate_limit_max_requests_str is not None:\n try:\n RATE_LIMIT_MAX_REQUESTS = int(_rate_limit_max_requests_str)\n except ValueError:\n pass\n\nAUTH_RATE_LIMITING_ENABLED = RATE_LIMIT_MAX_REQUESTS and RATE_LIMIT_WINDOW_SECONDS\n# Used for general redis things\nREDIS_DB_NUMBER = int(os.environ.get(\"REDIS_DB_NUMBER\", 0))\n\n# Used by celery as broker and backend\nREDIS_DB_NUMBER_CELERY_RESULT_BACKEND = int(\n os.environ.get(\"REDIS_DB_NUMBER_CELERY_RESULT_BACKEND\", 14)\n)\nREDIS_DB_NUMBER_CELERY = int(os.environ.get(\"REDIS_DB_NUMBER_CELERY\", 15)) # broker\n\n# will propagate to both our redis client as well as celery's redis client\nREDIS_HEALTH_CHECK_INTERVAL = int(os.environ.get(\"REDIS_HEALTH_CHECK_INTERVAL\", 60))\n\n# our redis client only, not celery's\nREDIS_POOL_MAX_CONNECTIONS = int(os.environ.get(\"REDIS_POOL_MAX_CONNECTIONS\", 128))\n\n# https://docs.celeryq.dev/en/stable/userguide/configuration.html#redis-backend-settings\n# should be one of \"required\", \"optional\", or \"none\"\nREDIS_SSL_CERT_REQS = os.getenv(\"REDIS_SSL_CERT_REQS\", \"none\")\nREDIS_SSL_CA_CERTS = os.getenv(\"REDIS_SSL_CA_CERTS\", None)\n\nCELERY_RESULT_EXPIRES = int(os.environ.get(\"CELERY_RESULT_EXPIRES\", 86400)) # seconds\n\n# https://docs.celeryq.dev/en/stable/userguide/configuration.html#broker-pool-limit\n# Setting to None may help when there is a proxy in the way closing idle connections\n_CELERY_BROKER_POOL_LIMIT_DEFAULT = 10\ntry:\n CELERY_BROKER_POOL_LIMIT = int(\n os.environ.get(\"CELERY_BROKER_POOL_LIMIT\", _CELERY_BROKER_POOL_LIMIT_DEFAULT)\n )\nexcept ValueError:\n CELERY_BROKER_POOL_LIMIT = _CELERY_BROKER_POOL_LIMIT_DEFAULT\n\n_CELERY_WORKER_LIGHT_CONCURRENCY_DEFAULT = 24\ntry:\n CELERY_WORKER_LIGHT_CONCURRENCY = int(\n os.environ.get(\n \"CELERY_WORKER_LIGHT_CONCURRENCY\",\n _CELERY_WORKER_LIGHT_CONCURRENCY_DEFAULT,\n )\n )\nexcept ValueError:\n CELERY_WORKER_LIGHT_CONCURRENCY = _CELERY_WORKER_LIGHT_CONCURRENCY_DEFAULT\n\n_CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER_DEFAULT = 8\ntry:\n CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER = int(\n os.environ.get(\n \"CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER\",\n _CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER_DEFAULT,\n )\n )\nexcept ValueError:\n CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER = (\n _CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER_DEFAULT\n )\n\n_CELERY_WORKER_DOCPROCESSING_CONCURRENCY_DEFAULT = 6\ntry:\n env_value = os.environ.get(\"CELERY_WORKER_DOCPROCESSING_CONCURRENCY\")\n if not env_value:\n env_value = os.environ.get(\"NUM_INDEXING_WORKERS\")\n\n if not env_value:\n env_value = str(_CELERY_WORKER_DOCPROCESSING_CONCURRENCY_DEFAULT)\n CELERY_WORKER_DOCPROCESSING_CONCURRENCY = int(env_value)\nexcept ValueError:\n CELERY_WORKER_DOCPROCESSING_CONCURRENCY = (\n _CELERY_WORKER_DOCPROCESSING_CONCURRENCY_DEFAULT\n )\n\n_CELERY_WORKER_DOCFETCHING_CONCURRENCY_DEFAULT = 1\ntry:\n env_value = os.environ.get(\"CELERY_WORKER_DOCFETCHING_CONCURRENCY\")\n if not env_value:\n env_value = os.environ.get(\"NUM_DOCFETCHING_WORKERS\")\n\n if not env_value:\n env_value = str(_CELERY_WORKER_DOCFETCHING_CONCURRENCY_DEFAULT)\n CELERY_WORKER_DOCFETCHING_CONCURRENCY = int(env_value)\nexcept ValueError:\n CELERY_WORKER_DOCFETCHING_CONCURRENCY = (\n _CELERY_WORKER_DOCFETCHING_CONCURRENCY_DEFAULT\n )\n\nCELERY_WORKER_PRIMARY_CONCURRENCY = int(\n os.environ.get(\"CELERY_WORKER_PRIMARY_CONCURRENCY\") or 4\n)\n\nCELERY_WORKER_PRIMARY_POOL_OVERFLOW = int(\n os.environ.get(\"CELERY_WORKER_PRIMARY_POOL_OVERFLOW\") or 4\n)\n\n# Individual worker concurrency settings\nCELERY_WORKER_HEAVY_CONCURRENCY = int(\n os.environ.get(\"CELERY_WORKER_HEAVY_CONCURRENCY\") or 4\n)\n\nCELERY_WORKER_MONITORING_CONCURRENCY = int(\n os.environ.get(\"CELERY_WORKER_MONITORING_CONCURRENCY\") or 1\n)\n\nCELERY_WORKER_USER_FILE_PROCESSING_CONCURRENCY = int(\n os.environ.get(\"CELERY_WORKER_USER_FILE_PROCESSING_CONCURRENCY\") or 2\n)\n\n# The maximum number of tasks that can be queued up to sync to Vespa in a single pass\nVESPA_SYNC_MAX_TASKS = 8192\n\nDB_YIELD_PER_DEFAULT = 64\n\n#####\n# Connector Configs\n#####\nPOLL_CONNECTOR_OFFSET = 30 # Minutes overlap between poll windows\n\n# View the list here:\n# https://github.com/onyx-dot-app/onyx/blob/main/backend/onyx/connectors/factory.py\n# If this is empty, all connectors are enabled, this is an option for security heavy orgs where\n# only very select connectors are enabled and admins cannot add other connector types\nENABLED_CONNECTOR_TYPES = os.environ.get(\"ENABLED_CONNECTOR_TYPES\") or \"\"\n\n# If set to true, curators can only access and edit assistants that they created\nCURATORS_CANNOT_VIEW_OR_EDIT_NON_OWNED_ASSISTANTS = (\n os.environ.get(\"CURATORS_CANNOT_VIEW_OR_EDIT_NON_OWNED_ASSISTANTS\", \"\").lower()\n == \"true\"\n)\n\n# Some calls to get information on expert users are quite costly especially with rate limiting\n# Since experts are not used in the actual user experience, currently it is turned off\n# for some connectors\nENABLE_EXPENSIVE_EXPERT_CALLS = False\n\n\n# TODO these should be available for frontend configuration, via advanced options expandable\nWEB_CONNECTOR_IGNORED_CLASSES = os.environ.get(\n \"WEB_CONNECTOR_IGNORED_CLASSES\", \"sidebar,footer\"\n).split(\",\")\nWEB_CONNECTOR_IGNORED_ELEMENTS = os.environ.get(\n \"WEB_CONNECTOR_IGNORED_ELEMENTS\", \"nav,footer,meta,script,style,symbol,aside\"\n).split(\",\")\nWEB_CONNECTOR_OAUTH_CLIENT_ID = os.environ.get(\"WEB_CONNECTOR_OAUTH_CLIENT_ID\")\nWEB_CONNECTOR_OAUTH_CLIENT_SECRET = os.environ.get(\"WEB_CONNECTOR_OAUTH_CLIENT_SECRET\")\nWEB_CONNECTOR_OAUTH_TOKEN_URL = os.environ.get(\"WEB_CONNECTOR_OAUTH_TOKEN_URL\")\nWEB_CONNECTOR_VALIDATE_URLS = os.environ.get(\"WEB_CONNECTOR_VALIDATE_URLS\")\n\nHTML_BASED_CONNECTOR_TRANSFORM_LINKS_STRATEGY = os.environ.get(\n \"HTML_BASED_CONNECTOR_TRANSFORM_LINKS_STRATEGY\",\n HtmlBasedConnectorTransformLinksStrategy.STRIP,\n)\n\nNOTION_CONNECTOR_DISABLE_RECURSIVE_PAGE_LOOKUP = (\n os.environ.get(\"NOTION_CONNECTOR_DISABLE_RECURSIVE_PAGE_LOOKUP\", \"\").lower()\n == \"true\"\n)\n\n\n#####\n# Confluence Connector Configs\n#####\n\nCONFLUENCE_CONNECTOR_LABELS_TO_SKIP = [\n ignored_tag\n for ignored_tag in os.environ.get(\"CONFLUENCE_CONNECTOR_LABELS_TO_SKIP\", \"\").split(\n \",\"\n )\n if ignored_tag\n]\n\n# Attachments exceeding this size will not be retrieved (in bytes)\nCONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD = int(\n os.environ.get(\"CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD\", 10 * 1024 * 1024)\n)\n# Attachments with more chars than this will not be indexed. This is to prevent extremely\n# large files from freezing indexing. 200,000 is ~100 google doc pages.\nCONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD = int(\n os.environ.get(\"CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD\", 200_000)\n)\n\n# A JSON-formatted array. Each item in the array should have the following structure:\n# {\n# \"user_id\": \"1234567890\",\n# \"username\": \"bob\",\n# \"display_name\": \"Bob Fitzgerald\",\n# \"email\": \"bob@example.com\",\n# \"type\": \"known\"\n# }\n_RAW_CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE = os.environ.get(\n \"CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE\", \"\"\n)\nCONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE = cast(\n list[dict[str, str]] | None,\n (\n json.loads(_RAW_CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE)\n if _RAW_CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE\n else None\n ),\n)\n\n# Due to breakages in the confluence API, the timezone offset must be specified client side\n# to match the user's specified timezone.\n\n# The current state of affairs:\n# CQL queries are parsed in the user's timezone and cannot be specified in UTC\n# no API retrieves the user's timezone\n# All data is returned in UTC, so we can't derive the user's timezone from that\n\n# https://community.developer.atlassian.com/t/confluence-cloud-time-zone-get-via-rest-api/35954/16\n# https://jira.atlassian.com/browse/CONFCLOUD-69670\n\n\ndef get_current_tz_offset() -> int:\n # datetime now() gets local time, datetime.now(timezone.utc) gets UTC time.\n # remove tzinfo to compare non-timezone-aware objects.\n time_diff = datetime.now() - datetime.now(timezone.utc).replace(tzinfo=None)\n return round(time_diff.total_seconds() / 3600)\n\n\n# enter as a floating point offset from UTC in hours (-24 < val < 24)\n# this will be applied globally, so it probably makes sense to transition this to per\n# connector as some point.\n# For the default value, we assume that the user's local timezone is more likely to be\n# correct (i.e. the configured user's timezone or the default server one) than UTC.\n# https://developer.atlassian.com/cloud/confluence/cql-fields/#created\nCONFLUENCE_TIMEZONE_OFFSET = float(\n os.environ.get(\"CONFLUENCE_TIMEZONE_OFFSET\", get_current_tz_offset())\n)\n\nCONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC = (\n os.environ.get(\"CONFLUENCE_USE_ONYX_USERS_FOR_GROUP_SYNC\", \"\").lower() == \"true\"\n)\n\nGOOGLE_DRIVE_CONNECTOR_SIZE_THRESHOLD = int(\n os.environ.get(\"GOOGLE_DRIVE_CONNECTOR_SIZE_THRESHOLD\", 10 * 1024 * 1024)\n)\n\n# Default size threshold for Drupal Wiki attachments (10MB)\nDRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD = int(\n os.environ.get(\"DRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD\", 10 * 1024 * 1024)\n)\n\n# Default size threshold for SharePoint files (20MB)\nSHAREPOINT_CONNECTOR_SIZE_THRESHOLD = int(\n os.environ.get(\"SHAREPOINT_CONNECTOR_SIZE_THRESHOLD\", 20 * 1024 * 1024)\n)\n\n# When True, group sync enumerates every Azure AD group in the tenant (expensive).\n# When False (default), only groups found in site role assignments are synced.\n# Can be overridden per-connector via the \"exhaustive_ad_enumeration\" key in\n# connector_specific_config.\nSHAREPOINT_EXHAUSTIVE_AD_ENUMERATION = (\n os.environ.get(\"SHAREPOINT_EXHAUSTIVE_AD_ENUMERATION\", \"\").lower() == \"true\"\n)\n\nBLOB_STORAGE_SIZE_THRESHOLD = int(\n os.environ.get(\"BLOB_STORAGE_SIZE_THRESHOLD\", 20 * 1024 * 1024)\n)\n\nJIRA_CONNECTOR_LABELS_TO_SKIP = [\n ignored_tag\n for ignored_tag in os.environ.get(\"JIRA_CONNECTOR_LABELS_TO_SKIP\", \"\").split(\",\")\n if ignored_tag\n]\n# Maximum size for Jira tickets in bytes (default: 100KB)\nJIRA_CONNECTOR_MAX_TICKET_SIZE = int(\n os.environ.get(\"JIRA_CONNECTOR_MAX_TICKET_SIZE\", 100 * 1024)\n)\nJIRA_SLIM_PAGE_SIZE = int(os.environ.get(\"JIRA_SLIM_PAGE_SIZE\", 500))\n\nGONG_CONNECTOR_START_TIME = os.environ.get(\"GONG_CONNECTOR_START_TIME\")\n\nGITHUB_CONNECTOR_BASE_URL = os.environ.get(\"GITHUB_CONNECTOR_BASE_URL\") or None\n\nGITLAB_CONNECTOR_INCLUDE_CODE_FILES = (\n os.environ.get(\"GITLAB_CONNECTOR_INCLUDE_CODE_FILES\", \"\").lower() == \"true\"\n)\n\n# Typically set to http://localhost:3000 for OAuth connector development\nCONNECTOR_LOCALHOST_OVERRIDE = os.getenv(\"CONNECTOR_LOCALHOST_OVERRIDE\")\n\n# Egnyte specific configs\nEGNYTE_CLIENT_ID = os.getenv(\"EGNYTE_CLIENT_ID\")\nEGNYTE_CLIENT_SECRET = os.getenv(\"EGNYTE_CLIENT_SECRET\")\n\n# Linear specific configs\nLINEAR_CLIENT_ID = os.getenv(\"LINEAR_CLIENT_ID\")\nLINEAR_CLIENT_SECRET = os.getenv(\"LINEAR_CLIENT_SECRET\")\n\n# Slack specific configs\nSLACK_NUM_THREADS = int(os.getenv(\"SLACK_NUM_THREADS\") or 8)\nMAX_SLACK_QUERY_EXPANSIONS = int(os.environ.get(\"MAX_SLACK_QUERY_EXPANSIONS\", \"5\"))\n\n# Slack federated search thread context settings\n# Batch size for fetching thread context (controls concurrent API calls per batch)\nSLACK_THREAD_CONTEXT_BATCH_SIZE = int(\n os.environ.get(\"SLACK_THREAD_CONTEXT_BATCH_SIZE\", \"5\")\n)\n# Maximum messages to fetch thread context for (top N by relevance get full context)\nMAX_SLACK_THREAD_CONTEXT_MESSAGES = int(\n os.environ.get(\"MAX_SLACK_THREAD_CONTEXT_MESSAGES\", \"5\")\n)\n\n# TestRail specific configs\nTESTRAIL_BASE_URL = os.environ.get(\"TESTRAIL_BASE_URL\", \"\")\nTESTRAIL_USERNAME = os.environ.get(\"TESTRAIL_USERNAME\", \"\")\nTESTRAIL_API_KEY = os.environ.get(\"TESTRAIL_API_KEY\", \"\")\n\nLEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE = (\n os.environ.get(\"LEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE\", \"\").lower()\n == \"true\"\n)\n\nDEFAULT_PRUNING_FREQ = 60 * 60 * 24 # Once a day\n\nALLOW_SIMULTANEOUS_PRUNING = (\n os.environ.get(\"ALLOW_SIMULTANEOUS_PRUNING\", \"\").lower() == \"true\"\n)\n\n# This is the maximum rate at which documents are queried for a pruning job. 0 disables the limitation.\nMAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE = int(\n os.environ.get(\"MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE\", 0)\n)\n\n# comma delimited list of zendesk article labels to skip indexing for\nZENDESK_CONNECTOR_SKIP_ARTICLE_LABELS = os.environ.get(\n \"ZENDESK_CONNECTOR_SKIP_ARTICLE_LABELS\", \"\"\n).split(\",\")\n\n\n#####\n# Indexing Configs\n#####\n# NOTE: Currently only supported in the Confluence and Google Drive connectors +\n# only handles some failures (Confluence = handles API call failures, Google\n# Drive = handles failures pulling files / parsing them)\nCONTINUE_ON_CONNECTOR_FAILURE = os.environ.get(\n \"CONTINUE_ON_CONNECTOR_FAILURE\", \"\"\n).lower() not in [\"false\", \"\"]\n# When swapping to a new embedding model, a secondary index is created in the background, to conserve\n# resources, we pause updates on the primary index by default while the secondary index is created\nDISABLE_INDEX_UPDATE_ON_SWAP = (\n os.environ.get(\"DISABLE_INDEX_UPDATE_ON_SWAP\", \"\").lower() == \"true\"\n)\n# More accurate results at the expense of indexing speed and index size (stores additional 4 MINI_CHUNK vectors)\nENABLE_MULTIPASS_INDEXING = (\n os.environ.get(\"ENABLE_MULTIPASS_INDEXING\", \"\").lower() == \"true\"\n)\n# Enable contextual retrieval\nENABLE_CONTEXTUAL_RAG = os.environ.get(\"ENABLE_CONTEXTUAL_RAG\", \"\").lower() == \"true\"\n\nDEFAULT_CONTEXTUAL_RAG_LLM_NAME = \"gpt-4o-mini\"\nDEFAULT_CONTEXTUAL_RAG_LLM_PROVIDER = \"DevEnvPresetOpenAI\"\n# Finer grained chunking for more detail retention\n# Slightly larger since the sentence aware split is a max cutoff so most minichunks will be under MINI_CHUNK_SIZE\n# tokens. But we need it to be at least as big as 1/4th chunk size to avoid having a tiny mini-chunk at the end\nMINI_CHUNK_SIZE = 150\n\n# This is the number of regular chunks per large chunk\nLARGE_CHUNK_RATIO = 4\n\n# The maximum number of chunks that can be held for 1 document processing batch\n# The purpose of this is to set an upper bound on memory usage\nMAX_CHUNKS_PER_DOC_BATCH = int(os.environ.get(\"MAX_CHUNKS_PER_DOC_BATCH\") or 1000)\n\n# Include the document level metadata in each chunk. If the metadata is too long, then it is thrown out\n# We don't want the metadata to overwhelm the actual contents of the chunk\nSKIP_METADATA_IN_CHUNK = os.environ.get(\"SKIP_METADATA_IN_CHUNK\", \"\").lower() == \"true\"\n\n# The indexer will warn in the logs whenver a document exceeds this threshold (in bytes)\nINDEXING_SIZE_WARNING_THRESHOLD = int(\n os.environ.get(\"INDEXING_SIZE_WARNING_THRESHOLD\") or 100 * 1024 * 1024\n)\n\n# during indexing, will log verbose memory diff stats every x batches and at the end.\n# 0 disables this behavior and is the default.\nINDEXING_TRACER_INTERVAL = int(os.environ.get(\"INDEXING_TRACER_INTERVAL\") or 0)\n\n# Enable multi-threaded embedding model calls for parallel processing\n# Note: only applies for API-based embedding models\nINDEXING_EMBEDDING_MODEL_NUM_THREADS = int(\n os.environ.get(\"INDEXING_EMBEDDING_MODEL_NUM_THREADS\") or 8\n)\n\n# Maximum file size in a document to be indexed\nMAX_DOCUMENT_CHARS = int(os.environ.get(\"MAX_DOCUMENT_CHARS\") or 5_000_000)\nMAX_FILE_SIZE_BYTES = int(\n os.environ.get(\"MAX_FILE_SIZE_BYTES\") or 2 * 1024 * 1024 * 1024\n) # 2GB in bytes\n\n# Use document summary for contextual rag\nUSE_DOCUMENT_SUMMARY = os.environ.get(\"USE_DOCUMENT_SUMMARY\", \"true\").lower() == \"true\"\n# Use chunk summary for contextual rag\nUSE_CHUNK_SUMMARY = os.environ.get(\"USE_CHUNK_SUMMARY\", \"true\").lower() == \"true\"\n# Average summary embeddings for contextual rag (not yet implemented)\nAVERAGE_SUMMARY_EMBEDDINGS = (\n os.environ.get(\"AVERAGE_SUMMARY_EMBEDDINGS\", \"false\").lower() == \"true\"\n)\n\nMAX_TOKENS_FOR_FULL_INCLUSION = 4096\n\n# The intent was to have this be configurable per query, but I don't think any\n# codepath was actually configuring this, so for the migrated Vespa interface\n# we'll just use the default value, but also have it be configurable by env var.\nRECENCY_BIAS_MULTIPLIER = float(os.environ.get(\"RECENCY_BIAS_MULTIPLIER\") or 1.0)\n\n# Should match the rerank-count value set in\n# backend/onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja.\nRERANK_COUNT = int(os.environ.get(\"RERANK_COUNT\") or 1000)\n\n\n#####\n# Tool Configs\n#####\n# Code Interpreter Service Configuration\nCODE_INTERPRETER_BASE_URL = os.environ.get(\n \"CODE_INTERPRETER_BASE_URL\", \"http://localhost:8000\"\n)\n\nCODE_INTERPRETER_DEFAULT_TIMEOUT_MS = int(\n os.environ.get(\"CODE_INTERPRETER_DEFAULT_TIMEOUT_MS\") or 60_000\n)\n\nCODE_INTERPRETER_MAX_OUTPUT_LENGTH = int(\n os.environ.get(\"CODE_INTERPRETER_MAX_OUTPUT_LENGTH\") or 50_000\n)\n\n\n#####\n# Miscellaneous\n#####\nJOB_TIMEOUT = 60 * 60 * 6 # 6 hours default\n# Logs Onyx only model interactions like prompts, responses, messages etc.\nLOG_ONYX_MODEL_INTERACTIONS = (\n os.environ.get(\"LOG_ONYX_MODEL_INTERACTIONS\", \"\").lower() == \"true\"\n)\n\nPROMPT_CACHE_CHAT_HISTORY = (\n os.environ.get(\"PROMPT_CACHE_CHAT_HISTORY\", \"\").lower() == \"true\"\n)\n# If set to `true` will enable additional logs about Vespa query performance\n# (time spent on finding the right docs + time spent fetching summaries from disk)\nLOG_VESPA_TIMING_INFORMATION = (\n os.environ.get(\"LOG_VESPA_TIMING_INFORMATION\", \"\").lower() == \"true\"\n)\nLOG_ENDPOINT_LATENCY = os.environ.get(\"LOG_ENDPOINT_LATENCY\", \"\").lower() == \"true\"\nLOG_POSTGRES_LATENCY = os.environ.get(\"LOG_POSTGRES_LATENCY\", \"\").lower() == \"true\"\nLOG_POSTGRES_CONN_COUNTS = (\n os.environ.get(\"LOG_POSTGRES_CONN_COUNTS\", \"\").lower() == \"true\"\n)\n# Anonymous usage telemetry\nDISABLE_TELEMETRY = os.environ.get(\"DISABLE_TELEMETRY\", \"\").lower() == \"true\"\n\n#####\n# Braintrust Configuration\n#####\n# Braintrust project name\nBRAINTRUST_PROJECT = os.environ.get(\"BRAINTRUST_PROJECT\", \"Onyx\")\n# Braintrust API key - if provided, Braintrust tracing will be enabled\nBRAINTRUST_API_KEY = os.environ.get(\"BRAINTRUST_API_KEY\") or \"\"\n# Maximum concurrency for Braintrust evaluations\n# None means unlimited concurrency, otherwise specify a number\n_braintrust_concurrency = os.environ.get(\"BRAINTRUST_MAX_CONCURRENCY\")\nBRAINTRUST_MAX_CONCURRENCY = (\n int(_braintrust_concurrency) if _braintrust_concurrency else None\n)\n\n#####\n# Scheduled Evals Configuration\n#####\n# Comma-separated list of Braintrust dataset names to run on schedule\nSCHEDULED_EVAL_DATASET_NAMES = [\n name.strip()\n for name in os.environ.get(\"SCHEDULED_EVAL_DATASET_NAMES\", \"\").split(\",\")\n if name.strip()\n]\n# Email address to use for search permissions during scheduled evals\nSCHEDULED_EVAL_PERMISSIONS_EMAIL = os.environ.get(\n \"SCHEDULED_EVAL_PERMISSIONS_EMAIL\", \"roshan@onyx.app\"\n)\n# Braintrust project name to use for scheduled evals\nSCHEDULED_EVAL_PROJECT = os.environ.get(\"SCHEDULED_EVAL_PROJECT\", \"st-dev\")\n\n#####\n# Langfuse Configuration\n#####\n# Langfuse API credentials - if provided, Langfuse tracing will be enabled\nLANGFUSE_SECRET_KEY = os.environ.get(\"LANGFUSE_SECRET_KEY\") or \"\"\nLANGFUSE_PUBLIC_KEY = os.environ.get(\"LANGFUSE_PUBLIC_KEY\") or \"\"\nLANGFUSE_HOST = os.environ.get(\"LANGFUSE_HOST\") or \"\" # For self-hosted Langfuse\n\n# Defined custom query/answer conditions to validate the query and the LLM answer.\n# Format: list of strings\nCUSTOM_ANSWER_VALIDITY_CONDITIONS = json.loads(\n os.environ.get(\"CUSTOM_ANSWER_VALIDITY_CONDITIONS\", \"[]\")\n)\n\nVESPA_REQUEST_TIMEOUT = int(os.environ.get(\"VESPA_REQUEST_TIMEOUT\") or \"15\")\n# This is the timeout for the client side of the Vespa migration task. When\n# exceeded, an exception is raised in our code. This value should be higher than\n# VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT.\nVESPA_MIGRATION_REQUEST_TIMEOUT_S = int(\n os.environ.get(\"VESPA_MIGRATION_REQUEST_TIMEOUT_S\") or \"120\"\n)\n# This is the timeout Vespa uses on the server side to know when to wrap up its\n# traversal and try to report partial results. This differs from the client\n# timeout above which raises an exception in our code when exceeded. This\n# timeout allows Vespa to return gracefully. This value should be lower than\n# VESPA_MIGRATION_REQUEST_TIMEOUT_S. Formatted as s.\nVESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT = os.environ.get(\n \"VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT\", \"110s\"\n)\n\nSYSTEM_RECURSION_LIMIT = int(os.environ.get(\"SYSTEM_RECURSION_LIMIT\") or \"1000\")\n\nPARSE_WITH_TRAFILATURA = os.environ.get(\"PARSE_WITH_TRAFILATURA\", \"\").lower() == \"true\"\n\n# allow for custom error messages for different errors returned by litellm\n# for example, can specify: {\"Violated content safety policy\": \"EVIL REQUEST!!!\"}\n# to make it so that if an LLM call returns an error containing \"Violated content safety policy\"\n# the end user will see \"EVIL REQUEST!!!\" instead of the default error message.\n_LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS = os.environ.get(\n \"LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS\", \"\"\n)\nLITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS: dict[str, str] | None = None\ntry:\n LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS = cast(\n dict[str, str], json.loads(_LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS)\n )\nexcept json.JSONDecodeError:\n pass\n\n# Auto LLM Configuration - fetches model configs from GitHub for providers in Auto mode\nAUTO_LLM_CONFIG_URL = os.environ.get(\n \"AUTO_LLM_CONFIG_URL\",\n \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/onyx/llm/well_known_providers/recommended-models.json\",\n)\n\n# How often to check for auto LLM model updates (in seconds)\nAUTO_LLM_UPDATE_INTERVAL_SECONDS = int(\n os.environ.get(\"AUTO_LLM_UPDATE_INTERVAL_SECONDS\", 1800) # 30 minutes\n)\n\n#####\n# Enterprise Edition Configs\n#####\n# NOTE: this should only be enabled if you have purchased an enterprise license.\n# if you're interested in an enterprise license, please reach out to us at\n# founders@onyx.app OR message Chris Weaver or Yuhong Sun in the Onyx\n# Discord community https://discord.gg/4NA5SbzrWb\nENTERPRISE_EDITION_ENABLED = (\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() == \"true\"\n)\n\n#####\n# Image Generation Configuration (DEPRECATED)\n# These environment variables will be deprecated soon.\n# To configure image generation, please visit the Image Generation page in the Admin Panel.\n#####\n# Azure Image Configurations\nAZURE_IMAGE_API_VERSION = os.environ.get(\"AZURE_IMAGE_API_VERSION\") or os.environ.get(\n \"AZURE_DALLE_API_VERSION\"\n)\nAZURE_IMAGE_API_KEY = os.environ.get(\"AZURE_IMAGE_API_KEY\") or os.environ.get(\n \"AZURE_DALLE_API_KEY\"\n)\nAZURE_IMAGE_API_BASE = os.environ.get(\"AZURE_IMAGE_API_BASE\") or os.environ.get(\n \"AZURE_DALLE_API_BASE\"\n)\nAZURE_IMAGE_DEPLOYMENT_NAME = os.environ.get(\n \"AZURE_IMAGE_DEPLOYMENT_NAME\"\n) or os.environ.get(\"AZURE_DALLE_DEPLOYMENT_NAME\")\n\n# configurable image model\nIMAGE_MODEL_NAME = os.environ.get(\"IMAGE_MODEL_NAME\", \"gpt-image-1\")\nIMAGE_MODEL_PROVIDER = os.environ.get(\"IMAGE_MODEL_PROVIDER\", \"openai\")\n\n# Use managed Vespa (Vespa Cloud). If set, must also set VESPA_CLOUD_URL, VESPA_CLOUD_CERT_PATH and VESPA_CLOUD_KEY_PATH\nMANAGED_VESPA = os.environ.get(\"MANAGED_VESPA\", \"\").lower() == \"true\"\n\nENABLE_EMAIL_INVITES = os.environ.get(\"ENABLE_EMAIL_INVITES\", \"\").lower() == \"true\"\n\n# Limit on number of users a free trial tenant can invite (cloud only)\nNUM_FREE_TRIAL_USER_INVITES = int(os.environ.get(\"NUM_FREE_TRIAL_USER_INVITES\", \"10\"))\n\n# Security and authentication\nDATA_PLANE_SECRET = os.environ.get(\n \"DATA_PLANE_SECRET\", \"\"\n) # Used for secure communication between the control and data plane\nEXPECTED_API_KEY = os.environ.get(\n \"EXPECTED_API_KEY\", \"\"\n) # Additional security check for the control plane API\n\n# API configuration\nCONTROL_PLANE_API_BASE_URL = os.environ.get(\n \"CONTROL_PLANE_API_BASE_URL\", \"http://localhost:8082\"\n)\n\nOAUTH_SLACK_CLIENT_ID = os.environ.get(\"OAUTH_SLACK_CLIENT_ID\", \"\")\nOAUTH_SLACK_CLIENT_SECRET = os.environ.get(\"OAUTH_SLACK_CLIENT_SECRET\", \"\")\nOAUTH_CONFLUENCE_CLOUD_CLIENT_ID = os.environ.get(\n \"OAUTH_CONFLUENCE_CLOUD_CLIENT_ID\", \"\"\n)\nOAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET = os.environ.get(\n \"OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET\", \"\"\n)\nOAUTH_GOOGLE_DRIVE_CLIENT_ID = os.environ.get(\"OAUTH_GOOGLE_DRIVE_CLIENT_ID\", \"\")\nOAUTH_GOOGLE_DRIVE_CLIENT_SECRET = os.environ.get(\n \"OAUTH_GOOGLE_DRIVE_CLIENT_SECRET\", \"\"\n)\n\n# JWT configuration\nJWT_ALGORITHM = \"HS256\"\n\n#####\n# API Key Configs\n#####\n# refers to the rounds described here: https://passlib.readthedocs.io/en/stable/lib/passlib.hash.sha256_crypt.html\n_API_KEY_HASH_ROUNDS_RAW = os.environ.get(\"API_KEY_HASH_ROUNDS\")\nAPI_KEY_HASH_ROUNDS = (\n int(_API_KEY_HASH_ROUNDS_RAW) if _API_KEY_HASH_ROUNDS_RAW else None\n)\n\n#####\n# MCP Server Configs\n#####\nMCP_SERVER_ENABLED = os.environ.get(\"MCP_SERVER_ENABLED\", \"\").lower() == \"true\"\nMCP_SERVER_HOST = os.environ.get(\"MCP_SERVER_HOST\", \"0.0.0.0\")\nMCP_SERVER_PORT = int(os.environ.get(\"MCP_SERVER_PORT\") or 8090)\n\n# CORS origins for MCP clients (comma-separated)\n# Local dev: \"http://localhost:*\"\n# Production: \"https://trusted-client.com,https://another-client.com\"\nMCP_SERVER_CORS_ORIGINS = [\n origin.strip()\n for origin in os.environ.get(\"MCP_SERVER_CORS_ORIGINS\", \"\").split(\",\")\n if origin.strip()\n]\n\n\nPOD_NAME = os.environ.get(\"POD_NAME\")\nPOD_NAMESPACE = os.environ.get(\"POD_NAMESPACE\")\n\n\nDEV_MODE = os.environ.get(\"DEV_MODE\", \"\").lower() == \"true\"\n\n\nINTEGRATION_TESTS_MODE = os.environ.get(\"INTEGRATION_TESTS_MODE\", \"\").lower() == \"true\"\n\n#####\n# Captcha Configuration (for cloud signup protection)\n#####\n# Enable captcha verification for new user registration\nCAPTCHA_ENABLED = os.environ.get(\"CAPTCHA_ENABLED\", \"\").lower() == \"true\"\n\n# Google reCAPTCHA secret key (server-side validation)\nRECAPTCHA_SECRET_KEY = os.environ.get(\"RECAPTCHA_SECRET_KEY\", \"\")\n\n# Minimum score threshold for reCAPTCHA v3 (0.0-1.0, higher = more likely human)\n# 0.5 is the recommended default\nRECAPTCHA_SCORE_THRESHOLD = float(os.environ.get(\"RECAPTCHA_SCORE_THRESHOLD\", \"0.5\"))\n\nMOCK_CONNECTOR_FILE_PATH = os.environ.get(\"MOCK_CONNECTOR_FILE_PATH\")\n\n# Set to true to mock LLM responses for testing purposes\nMOCK_LLM_RESPONSE = (\n os.environ.get(\"MOCK_LLM_RESPONSE\") if os.environ.get(\"MOCK_LLM_RESPONSE\") else None\n)\n\n\nDEFAULT_IMAGE_ANALYSIS_MAX_SIZE_MB = 20\n\n# Number of pre-provisioned tenants to maintain\nTARGET_AVAILABLE_TENANTS = int(os.environ.get(\"TARGET_AVAILABLE_TENANTS\", \"5\"))\n\n\n# Image summarization configuration\nIMAGE_SUMMARIZATION_SYSTEM_PROMPT = os.environ.get(\n \"IMAGE_SUMMARIZATION_SYSTEM_PROMPT\",\n DEFAULT_IMAGE_SUMMARIZATION_SYSTEM_PROMPT,\n)\n\n# The user prompt for image summarization - the image filename will be automatically prepended\nIMAGE_SUMMARIZATION_USER_PROMPT = os.environ.get(\n \"IMAGE_SUMMARIZATION_USER_PROMPT\",\n DEFAULT_IMAGE_SUMMARIZATION_USER_PROMPT,\n)\n\n# Knowledge Graph Read Only User Configuration\nDB_READONLY_USER: str = os.environ.get(\"DB_READONLY_USER\", \"db_readonly_user\")\nDB_READONLY_PASSWORD: str = urllib.parse.quote_plus(\n os.environ.get(\"DB_READONLY_PASSWORD\") or \"password\"\n)\n\n# File Store Configuration\n# Which backend to use for file storage: \"s3\" (S3/MinIO) or \"postgres\" (PostgreSQL Large Objects)\nFILE_STORE_BACKEND = os.environ.get(\"FILE_STORE_BACKEND\", \"s3\")\n\nS3_FILE_STORE_BUCKET_NAME = (\n os.environ.get(\"S3_FILE_STORE_BUCKET_NAME\") or \"onyx-file-store-bucket\"\n)\nS3_FILE_STORE_PREFIX = os.environ.get(\"S3_FILE_STORE_PREFIX\") or \"onyx-files\"\n# S3_ENDPOINT_URL is for MinIO and other S3-compatible storage. Leave blank for AWS S3.\nS3_ENDPOINT_URL = os.environ.get(\"S3_ENDPOINT_URL\")\nS3_VERIFY_SSL = os.environ.get(\"S3_VERIFY_SSL\", \"\").lower() == \"true\"\n\n# S3/MinIO Access Keys\nS3_AWS_ACCESS_KEY_ID = os.environ.get(\"S3_AWS_ACCESS_KEY_ID\")\nS3_AWS_SECRET_ACCESS_KEY = os.environ.get(\"S3_AWS_SECRET_ACCESS_KEY\")\n\n# Should we force S3 local checksumming\nS3_GENERATE_LOCAL_CHECKSUM = (\n os.environ.get(\"S3_GENERATE_LOCAL_CHECKSUM\", \"\").lower() == \"true\"\n)\n\n# Forcing Vespa Language\n# English: en, German:de, etc. See: https://docs.vespa.ai/en/linguistics.html\nVESPA_LANGUAGE_OVERRIDE = os.environ.get(\"VESPA_LANGUAGE_OVERRIDE\")\n\n\n#####\n# Default LLM API Keys (for cloud deployments)\n# These are Onyx-managed API keys provided to tenants by default\n#####\nOPENAI_DEFAULT_API_KEY = os.environ.get(\"OPENAI_DEFAULT_API_KEY\")\nANTHROPIC_DEFAULT_API_KEY = os.environ.get(\"ANTHROPIC_DEFAULT_API_KEY\")\nCOHERE_DEFAULT_API_KEY = os.environ.get(\"COHERE_DEFAULT_API_KEY\")\nVERTEXAI_DEFAULT_CREDENTIALS = os.environ.get(\"VERTEXAI_DEFAULT_CREDENTIALS\")\nVERTEXAI_DEFAULT_LOCATION = os.environ.get(\"VERTEXAI_DEFAULT_LOCATION\", \"global\")\nOPENROUTER_DEFAULT_API_KEY = os.environ.get(\"OPENROUTER_DEFAULT_API_KEY\")\n\nINSTANCE_TYPE = (\n \"managed\"\n if os.environ.get(\"IS_MANAGED_INSTANCE\", \"\").lower() == \"true\"\n else \"cloud\" if AUTH_TYPE == AuthType.CLOUD else \"self_hosted\"\n)\n\n\n## Discord Bot Configuration\nDISCORD_BOT_TOKEN = os.environ.get(\"DISCORD_BOT_TOKEN\")\nDISCORD_BOT_INVOKE_CHAR = os.environ.get(\"DISCORD_BOT_INVOKE_CHAR\", \"!\")\n\n\n## Stripe Configuration\n# URL to fetch the Stripe publishable key from a public S3 bucket.\n# Publishable keys are safe to expose publicly - they can only initialize\n# Stripe.js and tokenize payment info, not make charges or access data.\nSTRIPE_PUBLISHABLE_KEY_URL = (\n \"https://onyx-stripe-public.s3.amazonaws.com/publishable-key.txt\"\n)\n# Override for local testing with Stripe test keys (pk_test_*)\nSTRIPE_PUBLISHABLE_KEY_OVERRIDE = os.environ.get(\"STRIPE_PUBLISHABLE_KEY\")\n" }, { "path": "backend/onyx/configs/chat_configs.py", "content": "import os\n\nPROMPTS_YAML = \"./onyx/seeding/prompts.yaml\"\nPERSONAS_YAML = \"./onyx/seeding/personas.yaml\"\nNUM_RETURNED_HITS = 50\n\n# May be less depending on model\nMAX_CHUNKS_FED_TO_CHAT = int(os.environ.get(\"MAX_CHUNKS_FED_TO_CHAT\") or 25)\n\n# 1 / (1 + DOC_TIME_DECAY * doc-age-in-years), set to 0 to have no decay\n# Capped in Vespa at 0.5\nDOC_TIME_DECAY = float(\n os.environ.get(\"DOC_TIME_DECAY\") or 0.5 # Hits limit at 2 years by default\n)\nBASE_RECENCY_DECAY = 0.5\nFAVOR_RECENT_DECAY_MULTIPLIER = 2.0\n# For the highest matching base size chunk, how many chunks above and below do we pull in by default\n# Note this is not in any of the deployment configs yet\n# Currently only applies to search flow not chat\nCONTEXT_CHUNKS_ABOVE = int(os.environ.get(\"CONTEXT_CHUNKS_ABOVE\") or 1)\nCONTEXT_CHUNKS_BELOW = int(os.environ.get(\"CONTEXT_CHUNKS_BELOW\") or 1)\n# Fairly long but this is to account for edge cases where the LLM pauses for much longer than usual\n# The alternative is to fail the request completely so this is intended to be fairly lenient.\nLLM_SOCKET_READ_TIMEOUT = int(\n os.environ.get(\"LLM_SOCKET_READ_TIMEOUT\") or \"60\"\n) # 60 seconds\n# Weighting factor between vector and keyword Search; 1 for completely vector\n# search, 0 for keyword. Enforces a valid range of [0, 1]. A supplied value from\n# the env outside of this range will be clipped to the respective end of the\n# range. Defaults to 0.5.\nHYBRID_ALPHA = max(0, min(1, float(os.environ.get(\"HYBRID_ALPHA\") or 0.5)))\n# Weighting factor between Title and Content of documents during search, 1 for completely\n# Title based. Default heavily favors Content because Title is also included at the top of\n# Content. This is to avoid cases where the Content is very relevant but it may not be clear\n# if the title is separated out. Title is most of a \"boost\" than a separate field.\nTITLE_CONTENT_RATIO = max(\n 0, min(1, float(os.environ.get(\"TITLE_CONTENT_RATIO\") or 0.10))\n)\n\n# Stops streaming answers back to the UI if this pattern is seen:\nSTOP_STREAM_PAT = os.environ.get(\"STOP_STREAM_PAT\") or None\n\n# Set this to \"true\" to hard delete chats\n# This will make chats unviewable by admins after a user deletes them\n# As opposed to soft deleting them, which just hides them from non-admin users\nHARD_DELETE_CHATS = os.environ.get(\"HARD_DELETE_CHATS\", \"\").lower() == \"true\"\n\n# Internet Search\nNUM_INTERNET_SEARCH_RESULTS = int(os.environ.get(\"NUM_INTERNET_SEARCH_RESULTS\") or 10)\nNUM_INTERNET_SEARCH_CHUNKS = int(os.environ.get(\"NUM_INTERNET_SEARCH_CHUNKS\") or 50)\n\nVESPA_SEARCHER_THREADS = int(os.environ.get(\"VESPA_SEARCHER_THREADS\") or 2)\n\n# Whether or not to use the semantic & keyword search expansions for Basic Search\nUSE_SEMANTIC_KEYWORD_EXPANSIONS_BASIC_SEARCH = (\n os.environ.get(\"USE_SEMANTIC_KEYWORD_EXPANSIONS_BASIC_SEARCH\", \"false\").lower()\n == \"true\"\n)\n\n# Chat History Compression\n# Trigger compression when history exceeds this ratio of available context window\nCOMPRESSION_TRIGGER_RATIO = float(os.environ.get(\"COMPRESSION_TRIGGER_RATIO\", \"0.75\"))\n\nSKIP_DEEP_RESEARCH_CLARIFICATION = (\n os.environ.get(\"SKIP_DEEP_RESEARCH_CLARIFICATION\", \"false\").lower() == \"true\"\n)\n" }, { "path": "backend/onyx/configs/constants.py", "content": "import platform\nimport re\nimport socket\nfrom enum import auto\nfrom enum import Enum\n\n\nONYX_DEFAULT_APPLICATION_NAME = \"Onyx\"\nONYX_DISCORD_URL = \"https://discord.gg/4NA5SbzrWb\"\nONYX_UTM_SOURCE = \"onyx_app\"\nSLACK_USER_TOKEN_PREFIX = \"xoxp-\"\nSLACK_BOT_TOKEN_PREFIX = \"xoxb-\"\nONYX_EMAILABLE_LOGO_MAX_DIM = 512\n\nSOURCE_TYPE = \"source_type\"\n# stored in the `metadata` of a chunk. Used to signify that this chunk should\n# not be used for QA. For example, Google Drive file types which can't be parsed\n# are still useful as a search result but not for QA.\nIGNORE_FOR_QA = \"ignore_for_qa\"\n# NOTE: deprecated, only used for porting key from old system\nGEN_AI_API_KEY_STORAGE_KEY = \"genai_api_key\"\nPUBLIC_DOC_PAT = \"PUBLIC\"\nID_SEPARATOR = \":;:\"\nDEFAULT_BOOST = 0\n\n# Tag for endpoints that should be included in the public API documentation\nPUBLIC_API_TAGS: list[str | Enum] = [\"public\"]\n\n# Cookies\nFASTAPI_USERS_AUTH_COOKIE_NAME = (\n \"fastapiusersauth\" # Currently a constant, but logic allows for configuration\n)\nTENANT_ID_COOKIE_NAME = \"onyx_tid\" # tenant id - for workaround cases\nANONYMOUS_USER_COOKIE_NAME = \"onyx_anonymous_user\"\n\n# ID used in UserInfo API responses for anonymous users (not a UUID, just a string identifier)\nANONYMOUS_USER_INFO_ID = \"__anonymous_user__\"\n# Placeholder user for migrating no-auth data to first registered user\nNO_AUTH_PLACEHOLDER_USER_UUID = \"00000000-0000-0000-0000-000000000001\"\nNO_AUTH_PLACEHOLDER_USER_EMAIL = \"no-auth-placeholder@onyx.app\"\n# Real anonymous user in DB for anonymous access feature\nANONYMOUS_USER_UUID = \"00000000-0000-0000-0000-000000000002\"\nANONYMOUS_USER_EMAIL = \"anonymous@onyx.app\"\n\n# For chunking/processing chunks\nRETURN_SEPARATOR = \"\\n\\r\\n\"\nSECTION_SEPARATOR = \"\\n\\n\"\n# For combining attributes, doesn't have to be unique/perfect to work\nINDEX_SEPARATOR = \"===\"\n\n# For File Connector Metadata override file\nONYX_METADATA_FILENAME = \".onyx_metadata.json\"\n\n# Messages\nDISABLED_GEN_AI_MSG = (\n \"Your System Admin has disabled the Generative AI functionalities of Onyx.\\n\"\n \"Please contact them if you wish to have this enabled.\\n\"\n \"You can still use Onyx as a search engine.\"\n)\n\n#####\n# Version Pattern Configs\n#####\n# Version patterns for Docker image tags\nSTABLE_VERSION_PATTERN = re.compile(r\"^v(\\d+)\\.(\\d+)\\.(\\d+)$\")\nDEV_VERSION_PATTERN = re.compile(r\"^v(\\d+)\\.(\\d+)\\.(\\d+)-beta\\.(\\d+)$\")\n\nDEFAULT_PERSONA_ID = 0\n\nDEFAULT_CC_PAIR_ID = 1\n\n\nCANCEL_CHECK_INTERVAL = 20\nDISPATCH_SEP_CHAR = \"\\n\"\nFORMAT_DOCS_SEPARATOR = \"\\n\\n\"\nNUM_EXPLORATORY_DOCS = 15\n# Postgres connection constants for application_name\nPOSTGRES_WEB_APP_NAME = \"web\"\nPOSTGRES_INDEXER_APP_NAME = \"indexer\"\nPOSTGRES_CELERY_APP_NAME = \"celery\"\nPOSTGRES_CELERY_BEAT_APP_NAME = \"celery_beat\"\nPOSTGRES_CELERY_WORKER_PRIMARY_APP_NAME = \"celery_worker_primary\"\nPOSTGRES_CELERY_WORKER_LIGHT_APP_NAME = \"celery_worker_light\"\nPOSTGRES_CELERY_WORKER_DOCPROCESSING_APP_NAME = \"celery_worker_docprocessing\"\nPOSTGRES_CELERY_WORKER_DOCFETCHING_APP_NAME = \"celery_worker_docfetching\"\nPOSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME = \"celery_worker_indexing_child\"\nPOSTGRES_CELERY_WORKER_HEAVY_APP_NAME = \"celery_worker_heavy\"\nPOSTGRES_CELERY_WORKER_MONITORING_APP_NAME = \"celery_worker_monitoring\"\nPOSTGRES_CELERY_WORKER_USER_FILE_PROCESSING_APP_NAME = (\n \"celery_worker_user_file_processing\"\n)\nPOSTGRES_PERMISSIONS_APP_NAME = \"permissions\"\nPOSTGRES_UNKNOWN_APP_NAME = \"unknown\"\n\nSSL_CERT_FILE = \"bundle.pem\"\n# API Keys\nDANSWER_API_KEY_PREFIX = \"API_KEY__\"\nDANSWER_API_KEY_DUMMY_EMAIL_DOMAIN = \"onyxapikey.ai\"\nUNNAMED_KEY_PLACEHOLDER = \"Unnamed\"\nDISCORD_SERVICE_API_KEY_NAME = \"discord-bot-service\"\n\n# Key-Value store keys\nKV_REINDEX_KEY = \"needs_reindexing\"\nKV_UNSTRUCTURED_API_KEY = \"unstructured_api_key\"\nKV_USER_STORE_KEY = \"INVITED_USERS\"\nKV_PENDING_USERS_KEY = \"PENDING_USERS\"\nKV_ANONYMOUS_USER_PREFERENCES_KEY = \"anonymous_user_preferences\"\nKV_ANONYMOUS_USER_PERSONALIZATION_KEY = \"anonymous_user_personalization\"\nKV_CRED_KEY = \"credential_id_{}\"\nKV_GMAIL_CRED_KEY = \"gmail_app_credential\"\nKV_GMAIL_SERVICE_ACCOUNT_KEY = \"gmail_service_account_key\"\nKV_GOOGLE_DRIVE_CRED_KEY = \"google_drive_app_credential\"\nKV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY = \"google_drive_service_account_key\"\nKV_GEN_AI_KEY_CHECK_TIME = \"genai_api_key_last_check_time\"\nKV_SETTINGS_KEY = \"onyx_settings\"\nKV_CUSTOMER_UUID_KEY = \"customer_uuid\"\nKV_INSTANCE_DOMAIN_KEY = \"instance_domain\"\nKV_ENTERPRISE_SETTINGS_KEY = \"onyx_enterprise_settings\"\nKV_CUSTOM_ANALYTICS_SCRIPT_KEY = \"__custom_analytics_script__\"\nKV_KG_CONFIG_KEY = \"kg_config\"\n\n# NOTE: we use this timeout / 4 in various places to refresh a lock\n# might be worth separating this timeout into separate timeouts for each situation\nCELERY_GENERIC_BEAT_LOCK_TIMEOUT = 120\n\nCELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT = 120\n\n\nCELERY_PRIMARY_WORKER_LOCK_TIMEOUT = 120\n\n\n# hard timeout applied by the watchdog to the indexing connector run\n# to handle hung connectors\nCELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT = 3 * 60 * 60 # 3 hours (in seconds)\n\n# soft timeout for the lock taken by the indexing connector run\n# allows the lock to eventually expire if the managing code around it dies\n# if we can get callbacks as object bytes download, we could lower this a lot.\n# CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT + 15 minutes\n# hard termination should always fire first if the connector is hung\nCELERY_INDEXING_LOCK_TIMEOUT = CELERY_INDEXING_WATCHDOG_CONNECTOR_TIMEOUT + 900\n\n# Heartbeat interval for indexing worker liveness detection\nINDEXING_WORKER_HEARTBEAT_INTERVAL = 30 # seconds\n\n# how long a task should wait for associated fence to be ready\nCELERY_TASK_WAIT_FOR_FENCE_TIMEOUT = 5 * 60 # 5 min\n\n# needs to be long enough to cover the maximum time it takes to download an object\n# if we can get callbacks as object bytes download, we could lower this a lot.\nCELERY_PRUNING_LOCK_TIMEOUT = 3600 # 1 hour (in seconds)\n\nCELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT = 3600 # 1 hour (in seconds)\n\nCELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT = 300 # 5 min\n\nCELERY_USER_FILE_PROCESSING_LOCK_TIMEOUT = 30 * 60 # 30 minutes (in seconds)\n\n# How long a queued user-file task is valid before workers discard it.\n# Should be longer than the beat interval (20 s) but short enough to prevent\n# indefinite queue growth. Workers drop tasks older than this without touching\n# the DB, so a shorter value = faster drain of stale duplicates.\nCELERY_USER_FILE_PROCESSING_TASK_EXPIRES = 60 # 1 minute (in seconds)\n\n# Maximum number of tasks allowed in the user-file-processing queue before the\n# beat generator stops adding more. Prevents unbounded queue growth when workers\n# fall behind.\nUSER_FILE_PROCESSING_MAX_QUEUE_DEPTH = 500\n# How long a queued user-file-project-sync task remains valid.\n# Should be short enough to discard stale queue entries under load while still\n# allowing workers enough time to pick up new tasks.\nCELERY_USER_FILE_PROJECT_SYNC_TASK_EXPIRES = 60 # 1 minute (in seconds)\n\n# Max queue depth before user-file-project-sync producers stop enqueuing.\n# This applies backpressure when workers are falling behind.\nUSER_FILE_PROJECT_SYNC_MAX_QUEUE_DEPTH = 500\n\nCELERY_USER_FILE_PROJECT_SYNC_LOCK_TIMEOUT = 5 * 60 # 5 minutes (in seconds)\n\n# How long a queued user-file-delete task is valid before workers discard it.\n# Mirrors the processing task expiry to prevent indefinite queue growth when\n# files are stuck in DELETING status and the beat keeps re-enqueuing them.\nCELERY_USER_FILE_DELETE_TASK_EXPIRES = 60 # 1 minute (in seconds)\n\n# Max queue depth before the delete beat stops enqueuing more delete tasks.\nUSER_FILE_DELETE_MAX_QUEUE_DEPTH = 500\n\nCELERY_SANDBOX_FILE_SYNC_LOCK_TIMEOUT = 5 * 60 # 5 minutes (in seconds)\n\nDANSWER_REDIS_FUNCTION_LOCK_PREFIX = \"da_function_lock:\"\n\nTMP_DRALPHA_PERSONA_NAME = \"KG Beta\"\n\n\nclass DocumentSource(str, Enum):\n # Special case, document passed in via Onyx APIs without specifying a source type\n INGESTION_API = \"ingestion_api\"\n SLACK = \"slack\"\n WEB = \"web\"\n GOOGLE_DRIVE = \"google_drive\"\n GMAIL = \"gmail\"\n REQUESTTRACKER = \"requesttracker\"\n GITHUB = \"github\"\n GITBOOK = \"gitbook\"\n GITLAB = \"gitlab\"\n GURU = \"guru\"\n BOOKSTACK = \"bookstack\"\n OUTLINE = \"outline\"\n CONFLUENCE = \"confluence\"\n JIRA = \"jira\"\n SLAB = \"slab\"\n PRODUCTBOARD = \"productboard\"\n FILE = \"file\"\n CODA = \"coda\"\n CANVAS = \"canvas\"\n NOTION = \"notion\"\n ZULIP = \"zulip\"\n LINEAR = \"linear\"\n HUBSPOT = \"hubspot\"\n DOCUMENT360 = \"document360\"\n GONG = \"gong\"\n GOOGLE_SITES = \"google_sites\"\n ZENDESK = \"zendesk\"\n LOOPIO = \"loopio\"\n DROPBOX = \"dropbox\"\n SHAREPOINT = \"sharepoint\"\n TEAMS = \"teams\"\n SALESFORCE = \"salesforce\"\n DISCOURSE = \"discourse\"\n AXERO = \"axero\"\n CLICKUP = \"clickup\"\n MEDIAWIKI = \"mediawiki\"\n WIKIPEDIA = \"wikipedia\"\n ASANA = \"asana\"\n S3 = \"s3\"\n R2 = \"r2\"\n GOOGLE_CLOUD_STORAGE = \"google_cloud_storage\"\n OCI_STORAGE = \"oci_storage\"\n XENFORO = \"xenforo\"\n NOT_APPLICABLE = \"not_applicable\"\n DISCORD = \"discord\"\n FRESHDESK = \"freshdesk\"\n FIREFLIES = \"fireflies\"\n EGNYTE = \"egnyte\"\n AIRTABLE = \"airtable\"\n HIGHSPOT = \"highspot\"\n DRUPAL_WIKI = \"drupal_wiki\"\n\n IMAP = \"imap\"\n BITBUCKET = \"bitbucket\"\n TESTRAIL = \"testrail\"\n\n # Special case just for integration tests\n MOCK_CONNECTOR = \"mock_connector\"\n # Special case for user files\n USER_FILE = \"user_file\"\n # Raw files for Craft sandbox access (xlsx, pptx, docx, etc.)\n # Uses RAW_BINARY processing mode - no text extraction\n CRAFT_FILE = \"craft_file\"\n\n\nclass FederatedConnectorSource(str, Enum):\n FEDERATED_SLACK = \"federated_slack\"\n\n def to_non_federated_source(self) -> DocumentSource | None:\n if self == FederatedConnectorSource.FEDERATED_SLACK:\n return DocumentSource.SLACK\n return None\n\n\nDocumentSourceRequiringTenantContext: list[DocumentSource] = [DocumentSource.FILE]\n\n\nclass NotificationType(str, Enum):\n REINDEX = \"reindex\"\n PERSONA_SHARED = \"persona_shared\"\n TRIAL_ENDS_TWO_DAYS = \"two_day_trial_ending\" # 2 days left in trial\n RELEASE_NOTES = \"release_notes\"\n ASSISTANT_FILES_READY = \"assistant_files_ready\"\n FEATURE_ANNOUNCEMENT = \"feature_announcement\"\n\n\nclass BlobType(str, Enum):\n R2 = \"r2\"\n S3 = \"s3\"\n GOOGLE_CLOUD_STORAGE = \"google_cloud_storage\"\n OCI_STORAGE = \"oci_storage\"\n\n\nclass DocumentIndexType(str, Enum):\n COMBINED = \"combined\" # Vespa\n SPLIT = \"split\" # Typesense + Qdrant\n\n\nclass AuthType(str, Enum):\n BASIC = \"basic\"\n GOOGLE_OAUTH = \"google_oauth\"\n OIDC = \"oidc\"\n SAML = \"saml\"\n\n # google auth and basic\n CLOUD = \"cloud\"\n\n\nclass QueryHistoryType(str, Enum):\n DISABLED = \"disabled\"\n ANONYMIZED = \"anonymized\"\n NORMAL = \"normal\"\n\n\n# Special characters for password validation\nPASSWORD_SPECIAL_CHARS = \"!@#$%^&*()_+-=[]{}|;:,.<>?\"\n\n\nclass SessionType(str, Enum):\n CHAT = \"Chat\"\n SEARCH = \"Search\"\n SLACK = \"Slack\"\n\n\nclass QAFeedbackType(str, Enum):\n LIKE = \"like\" # User likes the answer, used for metrics\n DISLIKE = \"dislike\" # User dislikes the answer, used for metrics\n MIXED = \"mixed\" # User likes some answers and dislikes other, used for chat session metrics\n\n\nclass SearchFeedbackType(str, Enum):\n ENDORSE = \"endorse\" # boost this document for all future queries\n REJECT = \"reject\" # down-boost this document for all future queries\n HIDE = \"hide\" # mark this document as untrusted, hide from LLM\n UNHIDE = \"unhide\"\n\n\nclass MessageType(str, Enum):\n # Using OpenAI standards, Langchain equivalent shown in comment\n # System message is always constructed on the fly, not saved\n SYSTEM = \"system\" # SystemMessage\n USER = \"user\" # HumanMessage\n ASSISTANT = \"assistant\" # AIMessage - Can include tool_calls field for parallel tool calling\n TOOL_CALL_RESPONSE = \"tool_call_response\"\n USER_REMINDER = \"user_reminder\" # Custom Onyx message type which is translated into a USER message when passed to the LLM\n\n\nclass ChatMessageSimpleType(str, Enum):\n USER = \"user\"\n ASSISTANT = \"assistant\"\n TOOL_CALL = \"tool_call\"\n FILE_TEXT = \"file_text\"\n\n\nclass TokenRateLimitScope(str, Enum):\n USER = \"user\"\n USER_GROUP = \"user_group\"\n GLOBAL = \"global\"\n\n\nclass FileStoreType(str, Enum):\n S3 = \"s3\"\n POSTGRES = \"postgres\"\n\n\nclass FileOrigin(str, Enum):\n CHAT_UPLOAD = \"chat_upload\"\n CHAT_IMAGE_GEN = \"chat_image_gen\"\n CONNECTOR = \"connector\"\n CONNECTOR_METADATA = \"connector_metadata\"\n GENERATED_REPORT = \"generated_report\"\n INDEXING_CHECKPOINT = \"indexing_checkpoint\"\n PLAINTEXT_CACHE = \"plaintext_cache\"\n OTHER = \"other\"\n QUERY_HISTORY_CSV = \"query_history_csv\"\n SANDBOX_SNAPSHOT = \"sandbox_snapshot\"\n USER_FILE = \"user_file\"\n\n\nclass FileType(str, Enum):\n CSV = \"text/csv\"\n\n\nclass MilestoneRecordType(str, Enum):\n TENANT_CREATED = \"tenant_created\"\n USER_SIGNED_UP = \"user_signed_up\"\n VISITED_ADMIN_PAGE = \"visited_admin_page\"\n CREATED_CONNECTOR = \"created_connector\"\n CONNECTOR_SUCCEEDED = \"connector_succeeded\"\n RAN_QUERY = \"ran_query\"\n USER_MESSAGE_SENT = \"user_message_sent\"\n MULTIPLE_ASSISTANTS = \"multiple_assistants\"\n CREATED_ASSISTANT = \"created_assistant\"\n CREATED_ONYX_BOT = \"created_onyx_bot\"\n REQUESTED_CONNECTOR = \"requested_connector\"\n\n\nclass PostgresAdvisoryLocks(Enum):\n KOMBU_MESSAGE_CLEANUP_LOCK_ID = auto()\n\n\nclass OnyxCeleryQueues:\n # \"celery\" is the default queue defined by celery and also the queue\n # we are running in the primary worker to run system tasks\n # Tasks running in this queue should be designed specifically to run quickly\n PRIMARY = \"celery\"\n\n # Light queue\n VESPA_METADATA_SYNC = \"vespa_metadata_sync\"\n DOC_PERMISSIONS_UPSERT = \"doc_permissions_upsert\"\n CONNECTOR_DELETION = \"connector_deletion\"\n LLM_MODEL_UPDATE = \"llm_model_update\"\n CHECKPOINT_CLEANUP = \"checkpoint_cleanup\"\n INDEX_ATTEMPT_CLEANUP = \"index_attempt_cleanup\"\n # Heavy queue\n CONNECTOR_PRUNING = \"connector_pruning\"\n CONNECTOR_DOC_PERMISSIONS_SYNC = \"connector_doc_permissions_sync\"\n CONNECTOR_EXTERNAL_GROUP_SYNC = \"connector_external_group_sync\"\n CONNECTOR_HIERARCHY_FETCHING = \"connector_hierarchy_fetching\"\n CSV_GENERATION = \"csv_generation\"\n\n # User file processing queue\n USER_FILE_PROCESSING = \"user_file_processing\"\n USER_FILE_PROJECT_SYNC = \"user_file_project_sync\"\n USER_FILE_DELETE = \"user_file_delete\"\n # Document processing pipeline queue\n DOCPROCESSING = \"docprocessing\"\n CONNECTOR_DOC_FETCHING = \"connector_doc_fetching\"\n\n # Monitoring queue\n MONITORING = \"monitoring\"\n\n # Sandbox processing queue\n SANDBOX = \"sandbox\"\n\n OPENSEARCH_MIGRATION = \"opensearch_migration\"\n\n\nclass OnyxRedisLocks:\n PRIMARY_WORKER = \"da_lock:primary_worker\"\n CHECK_VESPA_SYNC_BEAT_LOCK = \"da_lock:check_vespa_sync_beat\"\n CHECK_CONNECTOR_DELETION_BEAT_LOCK = \"da_lock:check_connector_deletion_beat\"\n CHECK_PRUNE_BEAT_LOCK = \"da_lock:check_prune_beat\"\n CHECK_HIERARCHY_FETCHING_BEAT_LOCK = \"da_lock:check_hierarchy_fetching_beat\"\n CHECK_INDEXING_BEAT_LOCK = \"da_lock:check_indexing_beat\"\n CHECK_CHECKPOINT_CLEANUP_BEAT_LOCK = \"da_lock:check_checkpoint_cleanup_beat\"\n CHECK_INDEX_ATTEMPT_CLEANUP_BEAT_LOCK = \"da_lock:check_index_attempt_cleanup_beat\"\n CHECK_CONNECTOR_DOC_PERMISSIONS_SYNC_BEAT_LOCK = (\n \"da_lock:check_connector_doc_permissions_sync_beat\"\n )\n CHECK_CONNECTOR_EXTERNAL_GROUP_SYNC_BEAT_LOCK = (\n \"da_lock:check_connector_external_group_sync_beat\"\n )\n OPENSEARCH_MIGRATION_BEAT_LOCK = \"da_lock:opensearch_migration_beat\"\n\n MONITOR_BACKGROUND_PROCESSES_LOCK = \"da_lock:monitor_background_processes\"\n CHECK_AVAILABLE_TENANTS_LOCK = \"da_lock:check_available_tenants\"\n CLOUD_PRE_PROVISION_TENANT_LOCK = \"da_lock:pre_provision_tenant\"\n\n CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX = (\n \"da_lock:connector_doc_permissions_sync\"\n )\n CONNECTOR_EXTERNAL_GROUP_SYNC_LOCK_PREFIX = \"da_lock:connector_external_group_sync\"\n PRUNING_LOCK_PREFIX = \"da_lock:pruning\"\n INDEXING_METADATA_PREFIX = \"da_metadata:indexing\"\n\n SLACK_BOT_LOCK = \"da_lock:slack_bot\"\n SLACK_BOT_HEARTBEAT_PREFIX = \"da_heartbeat:slack_bot\"\n ANONYMOUS_USER_ENABLED = \"anonymous_user_enabled\"\n\n CLOUD_BEAT_TASK_GENERATOR_LOCK = \"da_lock:cloud_beat_task_generator\"\n CLOUD_CHECK_ALEMBIC_BEAT_LOCK = \"da_lock:cloud_check_alembic\"\n\n # User file processing\n USER_FILE_PROCESSING_BEAT_LOCK = \"da_lock:check_user_file_processing_beat\"\n USER_FILE_PROCESSING_LOCK_PREFIX = \"da_lock:user_file_processing\"\n # Short-lived key set when a task is enqueued; cleared when the worker picks it up.\n # Prevents the beat from re-enqueuing the same file while a task is already queued.\n USER_FILE_QUEUED_PREFIX = \"da_lock:user_file_queued\"\n USER_FILE_PROJECT_SYNC_BEAT_LOCK = \"da_lock:check_user_file_project_sync_beat\"\n USER_FILE_PROJECT_SYNC_LOCK_PREFIX = \"da_lock:user_file_project_sync\"\n USER_FILE_PROJECT_SYNC_QUEUED_PREFIX = \"da_lock:user_file_project_sync_queued\"\n USER_FILE_DELETE_BEAT_LOCK = \"da_lock:check_user_file_delete_beat\"\n USER_FILE_DELETE_LOCK_PREFIX = \"da_lock:user_file_delete\"\n # Short-lived key set when a delete task is enqueued; cleared when the worker picks it up.\n # Prevents the beat from re-enqueuing the same file while a delete task is already queued.\n USER_FILE_DELETE_QUEUED_PREFIX = \"da_lock:user_file_delete_queued\"\n\n # Release notes\n RELEASE_NOTES_FETCH_LOCK = \"da_lock:release_notes_fetch\"\n\n # Sandbox cleanup\n CLEANUP_IDLE_SANDBOXES_BEAT_LOCK = \"da_lock:cleanup_idle_sandboxes_beat\"\n CLEANUP_OLD_SNAPSHOTS_BEAT_LOCK = \"da_lock:cleanup_old_snapshots_beat\"\n\n # Sandbox file sync\n SANDBOX_FILE_SYNC_LOCK_PREFIX = \"da_lock:sandbox_file_sync\"\n\n\nclass OnyxRedisSignals:\n BLOCK_VALIDATE_INDEXING_FENCES = \"signal:block_validate_indexing_fences\"\n BLOCK_VALIDATE_EXTERNAL_GROUP_SYNC_FENCES = (\n \"signal:block_validate_external_group_sync_fences\"\n )\n BLOCK_VALIDATE_PERMISSION_SYNC_FENCES = (\n \"signal:block_validate_permission_sync_fences\"\n )\n BLOCK_PRUNING = \"signal:block_pruning\"\n BLOCK_VALIDATE_PRUNING_FENCES = \"signal:block_validate_pruning_fences\"\n BLOCK_BUILD_FENCE_LOOKUP_TABLE = \"signal:block_build_fence_lookup_table\"\n BLOCK_VALIDATE_CONNECTOR_DELETION_FENCES = (\n \"signal:block_validate_connector_deletion_fences\"\n )\n\n\nclass OnyxRedisConstants:\n ACTIVE_FENCES = \"active_fences\"\n\n\nclass OnyxCeleryPriority(int, Enum):\n HIGHEST = 0\n HIGH = auto()\n MEDIUM = auto()\n LOW = auto()\n LOWEST = auto()\n\n\n# a prefix used to distinguish system wide tasks in the cloud\nONYX_CLOUD_CELERY_TASK_PREFIX = \"cloud\"\n\n# the tenant id we use for system level redis operations\nONYX_CLOUD_TENANT_ID = \"cloud\"\n\n# the redis namespace for runtime variables\nONYX_CLOUD_REDIS_RUNTIME = \"runtime\"\nCLOUD_BUILD_FENCE_LOOKUP_TABLE_INTERVAL_DEFAULT = 600\n\n\nclass OnyxCeleryTask:\n DEFAULT = \"celery\"\n\n CLOUD_BEAT_TASK_GENERATOR = f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_generate_beat_tasks\"\n CLOUD_MONITOR_ALEMBIC = f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor_alembic\"\n CLOUD_MONITOR_CELERY_QUEUES = (\n f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor_celery_queues\"\n )\n CLOUD_CHECK_AVAILABLE_TENANTS = (\n f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check_available_tenants\"\n )\n CLOUD_MONITOR_CELERY_PIDBOX = (\n f\"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor_celery_pidbox\"\n )\n\n CHECK_FOR_CONNECTOR_DELETION = \"check_for_connector_deletion_task\"\n CHECK_FOR_VESPA_SYNC_TASK = \"check_for_vespa_sync_task\"\n CHECK_FOR_INDEXING = \"check_for_indexing\"\n CHECK_FOR_PRUNING = \"check_for_pruning\"\n CHECK_FOR_HIERARCHY_FETCHING = \"check_for_hierarchy_fetching\"\n CHECK_FOR_DOC_PERMISSIONS_SYNC = \"check_for_doc_permissions_sync\"\n CHECK_FOR_EXTERNAL_GROUP_SYNC = \"check_for_external_group_sync\"\n CHECK_FOR_AUTO_LLM_UPDATE = \"check_for_auto_llm_update\"\n\n # User file processing\n CHECK_FOR_USER_FILE_PROCESSING = \"check_for_user_file_processing\"\n PROCESS_SINGLE_USER_FILE = \"process_single_user_file\"\n CHECK_FOR_USER_FILE_PROJECT_SYNC = \"check_for_user_file_project_sync\"\n PROCESS_SINGLE_USER_FILE_PROJECT_SYNC = \"process_single_user_file_project_sync\"\n CHECK_FOR_USER_FILE_DELETE = \"check_for_user_file_delete\"\n DELETE_SINGLE_USER_FILE = \"delete_single_user_file\"\n\n # Connector checkpoint cleanup\n CHECK_FOR_CHECKPOINT_CLEANUP = \"check_for_checkpoint_cleanup\"\n CLEANUP_CHECKPOINT = \"cleanup_checkpoint\"\n\n # Connector index attempt cleanup\n CHECK_FOR_INDEX_ATTEMPT_CLEANUP = \"check_for_index_attempt_cleanup\"\n CLEANUP_INDEX_ATTEMPT = \"cleanup_index_attempt\"\n\n MONITOR_BACKGROUND_PROCESSES = \"monitor_background_processes\"\n MONITOR_CELERY_QUEUES = \"monitor_celery_queues\"\n MONITOR_PROCESS_MEMORY = \"monitor_process_memory\"\n CELERY_BEAT_HEARTBEAT = \"celery_beat_heartbeat\"\n\n KOMBU_MESSAGE_CLEANUP_TASK = \"kombu_message_cleanup_task\"\n CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK = (\n \"connector_permission_sync_generator_task\"\n )\n UPDATE_EXTERNAL_DOCUMENT_PERMISSIONS_TASK = (\n \"update_external_document_permissions_task\"\n )\n CONNECTOR_EXTERNAL_GROUP_SYNC_GENERATOR_TASK = (\n \"connector_external_group_sync_generator_task\"\n )\n\n # New split indexing tasks\n CONNECTOR_DOC_FETCHING_TASK = \"connector_doc_fetching_task\"\n DOCPROCESSING_TASK = \"docprocessing_task\"\n\n CONNECTOR_PRUNING_GENERATOR_TASK = \"connector_pruning_generator_task\"\n CONNECTOR_HIERARCHY_FETCHING_TASK = \"connector_hierarchy_fetching_task\"\n DOCUMENT_BY_CC_PAIR_CLEANUP_TASK = \"document_by_cc_pair_cleanup_task\"\n VESPA_METADATA_SYNC_TASK = \"vespa_metadata_sync_task\"\n\n # chat retention\n CHECK_TTL_MANAGEMENT_TASK = \"check_ttl_management_task\"\n PERFORM_TTL_MANAGEMENT_TASK = \"perform_ttl_management_task\"\n\n GENERATE_USAGE_REPORT_TASK = \"generate_usage_report_task\"\n\n EVAL_RUN_TASK = \"eval_run_task\"\n SCHEDULED_EVAL_TASK = \"scheduled_eval_task\"\n\n EXPORT_QUERY_HISTORY_TASK = \"export_query_history_task\"\n EXPORT_QUERY_HISTORY_CLEANUP_TASK = \"export_query_history_cleanup_task\"\n\n # Hook execution log retention\n HOOK_EXECUTION_LOG_CLEANUP_TASK = \"hook_execution_log_cleanup_task\"\n\n # Sandbox cleanup\n CLEANUP_IDLE_SANDBOXES = \"cleanup_idle_sandboxes\"\n CLEANUP_OLD_SNAPSHOTS = \"cleanup_old_snapshots\"\n\n # Sandbox file sync\n SANDBOX_FILE_SYNC = \"sandbox_file_sync\"\n\n CHECK_FOR_DOCUMENTS_FOR_OPENSEARCH_MIGRATION_TASK = (\n \"check_for_documents_for_opensearch_migration_task\"\n )\n MIGRATE_DOCUMENTS_FROM_VESPA_TO_OPENSEARCH_TASK = (\n \"migrate_documents_from_vespa_to_opensearch_task\"\n )\n MIGRATE_CHUNKS_FROM_VESPA_TO_OPENSEARCH_TASK = (\n \"migrate_chunks_from_vespa_to_opensearch_task\"\n )\n\n\n# this needs to correspond to the matching entry in supervisord\nONYX_CELERY_BEAT_HEARTBEAT_KEY = \"onyx:celery:beat:heartbeat\"\n\nREDIS_SOCKET_KEEPALIVE_OPTIONS = {}\nREDIS_SOCKET_KEEPALIVE_OPTIONS[socket.TCP_KEEPINTVL] = 15\nREDIS_SOCKET_KEEPALIVE_OPTIONS[socket.TCP_KEEPCNT] = 3\n\nif platform.system() == \"Darwin\":\n REDIS_SOCKET_KEEPALIVE_OPTIONS[socket.TCP_KEEPALIVE] = 60 # type: ignore[attr-defined,unused-ignore]\nelse:\n REDIS_SOCKET_KEEPALIVE_OPTIONS[socket.TCP_KEEPIDLE] = 60 # type: ignore[attr-defined,unused-ignore]\n\n\nclass OnyxCallTypes(str, Enum):\n FIREFLIES = \"FIREFLIES\"\n GONG = \"GONG\"\n\n\nNUM_DAYS_TO_KEEP_CHECKPOINTS = 7\n# checkpoints are queried based on index attempts, so we need to keep index attempts for one more day\nNUM_DAYS_TO_KEEP_INDEX_ATTEMPTS = NUM_DAYS_TO_KEEP_CHECKPOINTS + 1\n\n# TODO: this should be stored likely in database\nDocumentSourceDescription: dict[DocumentSource, str] = {\n # Special case, document passed in via Onyx APIs without specifying a source type\n DocumentSource.INGESTION_API: \"ingestion_api\",\n DocumentSource.SLACK: \"slack channels for discussions and collaboration\",\n DocumentSource.WEB: \"indexed web pages\",\n DocumentSource.GOOGLE_DRIVE: \"google drive documents (docs, sheets, etc.)\",\n DocumentSource.GMAIL: \"email messages\",\n DocumentSource.REQUESTTRACKER: \"requesttracker\",\n DocumentSource.GITHUB: \"github data (issues, PRs)\",\n DocumentSource.GITBOOK: \"gitbook data\",\n DocumentSource.GITLAB: \"gitlab data\",\n DocumentSource.BITBUCKET: \"bitbucket data\",\n DocumentSource.GURU: \"guru data\",\n DocumentSource.BOOKSTACK: \"bookstack data\",\n DocumentSource.OUTLINE: \"outline data\",\n DocumentSource.CONFLUENCE: \"confluence data (pages, spaces, etc.)\",\n DocumentSource.JIRA: \"jira data (issues, tickets, projects, etc.)\",\n DocumentSource.SLAB: \"slab data\",\n DocumentSource.PRODUCTBOARD: \"productboard data (boards, etc.)\",\n DocumentSource.FILE: \"files\",\n DocumentSource.CANVAS: \"canvas lms - courses, pages, assignments, and announcements\",\n DocumentSource.CODA: \"coda - team workspace with docs, tables, and pages\",\n DocumentSource.NOTION: \"notion data - a workspace that combines note-taking, \\\nproject management, and collaboration tools into a single, customizable platform\",\n DocumentSource.ZULIP: \"zulip data\",\n DocumentSource.LINEAR: \"linear data - project management tool, including tickets etc.\",\n DocumentSource.HUBSPOT: \"hubspot data - CRM and marketing automation data\",\n DocumentSource.DOCUMENT360: \"document360 data\",\n DocumentSource.GONG: \"gong - call transcripts\",\n DocumentSource.GOOGLE_SITES: \"google_sites - websites\",\n DocumentSource.ZENDESK: \"zendesk - customer support data\",\n DocumentSource.LOOPIO: \"loopio - rfp data\",\n DocumentSource.DROPBOX: \"dropbox - files\",\n DocumentSource.SHAREPOINT: \"sharepoint - files\",\n DocumentSource.TEAMS: \"teams - chat and collaboration\",\n DocumentSource.SALESFORCE: \"salesforce - CRM data\",\n DocumentSource.DISCOURSE: \"discourse - discussion forums\",\n DocumentSource.AXERO: \"axero - employee engagement data\",\n DocumentSource.CLICKUP: \"clickup - project management tool\",\n DocumentSource.MEDIAWIKI: \"mediawiki - wiki data\",\n DocumentSource.WIKIPEDIA: \"wikipedia - encyclopedia data\",\n DocumentSource.ASANA: \"asana\",\n DocumentSource.S3: \"s3\",\n DocumentSource.R2: \"r2\",\n DocumentSource.GOOGLE_CLOUD_STORAGE: \"google_cloud_storage - cloud storage\",\n DocumentSource.OCI_STORAGE: \"oci_storage - cloud storage\",\n DocumentSource.XENFORO: \"xenforo - forum data\",\n DocumentSource.DISCORD: \"discord - chat and collaboration\",\n DocumentSource.FRESHDESK: \"freshdesk - customer support data\",\n DocumentSource.FIREFLIES: \"fireflies - call transcripts\",\n DocumentSource.EGNYTE: \"egnyte - files\",\n DocumentSource.AIRTABLE: \"airtable - database\",\n DocumentSource.HIGHSPOT: \"highspot - CRM data\",\n DocumentSource.DRUPAL_WIKI: \"drupal wiki - knowledge base content (pages, spaces, attachments)\",\n DocumentSource.IMAP: \"imap - email data\",\n DocumentSource.TESTRAIL: \"testrail - test case management tool for QA processes\",\n}\n" }, { "path": "backend/onyx/configs/embedding_configs.py", "content": "from pydantic import BaseModel\n\nfrom onyx.db.enums import EmbeddingPrecision\n\n\nclass _BaseEmbeddingModel(BaseModel):\n \"\"\"Private model for defining base embedding model configurations.\"\"\"\n\n name: str\n dim: int\n index_name: str\n\n\nclass SupportedEmbeddingModel(BaseModel):\n name: str\n dim: int\n index_name: str\n embedding_precision: EmbeddingPrecision\n\n\n# Base embedding model configurations (without precision)\n_BASE_EMBEDDING_MODELS = [\n # Cloud-based models\n _BaseEmbeddingModel(\n name=\"cohere/embed-english-v3.0\",\n dim=1024,\n index_name=\"danswer_chunk_cohere_embed_english_v3_0\",\n ),\n _BaseEmbeddingModel(\n name=\"cohere/embed-english-v3.0\",\n dim=1024,\n index_name=\"danswer_chunk_embed_english_v3_0\",\n ),\n _BaseEmbeddingModel(\n name=\"cohere/embed-english-light-v3.0\",\n dim=384,\n index_name=\"danswer_chunk_cohere_embed_english_light_v3_0\",\n ),\n _BaseEmbeddingModel(\n name=\"cohere/embed-english-light-v3.0\",\n dim=384,\n index_name=\"danswer_chunk_embed_english_light_v3_0\",\n ),\n _BaseEmbeddingModel(\n name=\"openai/text-embedding-3-large\",\n dim=3072,\n index_name=\"danswer_chunk_openai_text_embedding_3_large\",\n ),\n _BaseEmbeddingModel(\n name=\"openai/text-embedding-3-large\",\n dim=3072,\n index_name=\"danswer_chunk_text_embedding_3_large\",\n ),\n _BaseEmbeddingModel(\n name=\"openai/text-embedding-3-small\",\n dim=1536,\n index_name=\"danswer_chunk_openai_text_embedding_3_small\",\n ),\n _BaseEmbeddingModel(\n name=\"openai/text-embedding-3-small\",\n dim=1536,\n index_name=\"danswer_chunk_text_embedding_3_small\",\n ),\n _BaseEmbeddingModel(\n name=\"google/gemini-embedding-001\",\n dim=3072,\n index_name=\"danswer_chunk_gemini_embedding_001\",\n ),\n _BaseEmbeddingModel(\n name=\"google/text-embedding-005\",\n dim=768,\n index_name=\"danswer_chunk_text_embedding_005\",\n ),\n _BaseEmbeddingModel(\n name=\"voyage/voyage-large-2-instruct\",\n dim=1024,\n index_name=\"danswer_chunk_voyage_large_2_instruct\",\n ),\n _BaseEmbeddingModel(\n name=\"voyage/voyage-large-2-instruct\",\n dim=1024,\n index_name=\"danswer_chunk_large_2_instruct\",\n ),\n _BaseEmbeddingModel(\n name=\"voyage/voyage-light-2-instruct\",\n dim=384,\n index_name=\"danswer_chunk_voyage_light_2_instruct\",\n ),\n _BaseEmbeddingModel(\n name=\"voyage/voyage-light-2-instruct\",\n dim=384,\n index_name=\"danswer_chunk_light_2_instruct\",\n ),\n # Self-hosted models\n _BaseEmbeddingModel(\n name=\"nomic-ai/nomic-embed-text-v1\",\n dim=768,\n index_name=\"danswer_chunk_nomic_ai_nomic_embed_text_v1\",\n ),\n _BaseEmbeddingModel(\n name=\"nomic-ai/nomic-embed-text-v1\",\n dim=768,\n index_name=\"danswer_chunk_nomic_embed_text_v1\",\n ),\n _BaseEmbeddingModel(\n name=\"intfloat/e5-base-v2\",\n dim=768,\n index_name=\"danswer_chunk_intfloat_e5_base_v2\",\n ),\n _BaseEmbeddingModel(\n name=\"intfloat/e5-small-v2\",\n dim=384,\n index_name=\"danswer_chunk_intfloat_e5_small_v2\",\n ),\n _BaseEmbeddingModel(\n name=\"intfloat/multilingual-e5-base\",\n dim=768,\n index_name=\"danswer_chunk_intfloat_multilingual_e5_base\",\n ),\n _BaseEmbeddingModel(\n name=\"intfloat/multilingual-e5-small\",\n dim=384,\n index_name=\"danswer_chunk_intfloat_multilingual_e5_small\",\n ),\n]\n\n# Automatically generate both FLOAT and BFLOAT16 versions of all models\nSUPPORTED_EMBEDDING_MODELS = [\n # BFLOAT16 precision versions\n *[\n SupportedEmbeddingModel(\n name=model.name,\n dim=model.dim,\n index_name=f\"{model.index_name}_bfloat16\",\n embedding_precision=EmbeddingPrecision.BFLOAT16,\n )\n for model in _BASE_EMBEDDING_MODELS\n ],\n # FLOAT precision versions\n # NOTE: need to keep this one for backwards compatibility. We now default to\n # BFLOAT16.\n *[\n SupportedEmbeddingModel(\n name=model.name,\n dim=model.dim,\n index_name=model.index_name,\n embedding_precision=EmbeddingPrecision.FLOAT,\n )\n for model in _BASE_EMBEDDING_MODELS\n ],\n]\n" }, { "path": "backend/onyx/configs/kg_configs.py", "content": "import os\n\nKG_RESEARCH_NUM_RETRIEVED_DOCS: int = int(\n os.environ.get(\"KG_RESEARCH_NUM_RETRIEVED_DOCS\", \"25\")\n)\n\n\nKG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES: int = int(\n os.environ.get(\"KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES\", \"10\")\n)\n\n\nKG_ENTITY_EXTRACTION_TIMEOUT: int = int(\n os.environ.get(\"KG_ENTITY_EXTRACTION_TIMEOUT\", \"15\")\n)\n\nKG_RELATIONSHIP_EXTRACTION_TIMEOUT: int = int(\n os.environ.get(\"KG_RELATIONSHIP_EXTRACTION_TIMEOUT\", \"15\")\n)\n\nKG_STRATEGY_GENERATION_TIMEOUT: int = int(\n os.environ.get(\"KG_STRATEGY_GENERATION_TIMEOUT\", \"20\")\n)\n\nKG_SQL_GENERATION_TIMEOUT: int = int(os.environ.get(\"KG_SQL_GENERATION_TIMEOUT\", \"40\"))\n\nKG_SQL_GENERATION_TIMEOUT_OVERRIDE: int = int(\n os.environ.get(\"KG_SQL_GENERATION_TIMEOUT_OVERRIDE\", \"40\")\n)\n\nKG_SQL_GENERATION_MAX_TOKENS: int = int(\n os.environ.get(\"KG_SQL_GENERATION_MAX_TOKENS\", \"1500\")\n)\n\nKG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX: str = os.environ.get(\n \"KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX\", \"allowed_docs\"\n)\n\nKG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX: str = os.environ.get(\n \"KG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX\", \"kg_relationships_with_access\"\n)\n\nKG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX: str = os.environ.get(\n \"KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX\", \"kg_entities_with_access\"\n)\n\n\nKG_FILTER_CONSTRUCTION_TIMEOUT: int = int(\n os.environ.get(\"KG_FILTER_CONSTRUCTION_TIMEOUT\", \"15\")\n)\n\n\nKG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT: int = int(\n os.environ.get(\"KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT\", \"100\")\n)\n\nKG_FILTERED_SEARCH_TIMEOUT: int = int(\n os.environ.get(\"KG_FILTERED_SEARCH_TIMEOUT\", \"30\")\n)\n\n\nKG_OBJECT_SOURCE_RESEARCH_TIMEOUT: int = int(\n os.environ.get(\"KG_OBJECT_SOURCE_RESEARCH_TIMEOUT\", \"30\")\n)\n\nKG_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION: int = int(\n os.environ.get(\"KG_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION\", \"45\")\n)\n\nKG_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION: int = int(\n os.environ.get(\"KG_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION\", \"15\")\n)\n\nKG_MAX_TOKENS_ANSWER_GENERATION: int = int(\n os.environ.get(\"KG_MAX_TOKENS_ANSWER_GENERATION\", \"1024\")\n)\n\nKG_MAX_DEEP_SEARCH_RESULTS: int = int(\n os.environ.get(\"KG_MAX_DEEP_SEARCH_RESULTS\", \"30\")\n)\n\n\nKG_METADATA_TRACKING_THRESHOLD: int = int(\n os.environ.get(\"KG_METADATA_TRACKING_THRESHOLD\", \"10\")\n)\n\n\nKG_DEFAULT_MAX_PARENT_RECURSION_DEPTH: int = int(\n os.environ.get(\"KG_DEFAULT_MAX_PARENT_RECURSION_DEPTH\", \"2\")\n)\n\n\n_KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT: float = max(\n 1e-3,\n min(1, float(os.environ.get(\"KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT\", \"0.25\"))),\n)\n_KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT: float = max(\n 1e-3,\n min(1, float(os.environ.get(\"KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT\", \"0.25\"))),\n)\n_KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT: float = max(\n 1e-3,\n min(1, float(os.environ.get(\"KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT\", \"0.5\"))),\n)\n_KG_NORMALIZATION_RERANK_NGRAM_SUMS: float = (\n _KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT\n + _KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT\n + _KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT\n)\n\nKG_NORMALIZATION_RERANK_NGRAM_WEIGHTS: tuple[float, float, float] = (\n _KG_NORMALIZATION_RERANK_UNIGRAM_WEIGHT / _KG_NORMALIZATION_RERANK_NGRAM_SUMS,\n _KG_NORMALIZATION_RERANK_BIGRAM_WEIGHT / _KG_NORMALIZATION_RERANK_NGRAM_SUMS,\n _KG_NORMALIZATION_RERANK_TRIGRAM_WEIGHT / _KG_NORMALIZATION_RERANK_NGRAM_SUMS,\n)\n\n\nKG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT: float = max(\n 0,\n min(1, float(os.environ.get(\"KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT\", \"0.25\"))),\n)\n\n\nKG_NORMALIZATION_RERANK_THRESHOLD: float = float(\n os.environ.get(\"KG_NORMALIZATION_RERANK_THRESHOLD\", \"0.3\")\n)\n\n\nKG_CLUSTERING_RETRIEVE_THRESHOLD: float = float(\n os.environ.get(\"KG_CLUSTERING_RETRIEVE_THRESHOLD\", \"0.6\")\n)\n\n\nKG_CLUSTERING_THRESHOLD: float = float(\n os.environ.get(\"KG_CLUSTERING_THRESHOLD\", \"0.96\")\n)\n\nKG_MAX_SEARCH_DOCUMENTS: int = int(os.environ.get(\"KG_MAX_SEARCH_DOCUMENTS\", \"15\"))\n\nKG_MAX_DECOMPOSITION_SEGMENTS: int = int(\n os.environ.get(\"KG_MAX_DECOMPOSITION_SEGMENTS\", \"10\")\n)\nKG_BETA_ASSISTANT_DESCRIPTION = \"The KG Beta assistant uses the Onyx Knowledge Graph (beta) structure \\\nto answer questions\"\n" }, { "path": "backend/onyx/configs/llm_configs.py", "content": "from onyx.configs.app_configs import DEFAULT_IMAGE_ANALYSIS_MAX_SIZE_MB\nfrom onyx.server.settings.store import load_settings\n\n\ndef get_image_extraction_and_analysis_enabled() -> bool:\n \"\"\"Get image extraction and analysis enabled setting from workspace settings or fallback to False\"\"\"\n try:\n settings = load_settings()\n if settings.image_extraction_and_analysis_enabled is not None:\n return settings.image_extraction_and_analysis_enabled\n except Exception:\n pass\n\n return False\n\n\ndef get_search_time_image_analysis_enabled() -> bool:\n \"\"\"Get search time image analysis enabled setting from workspace settings or fallback to False\"\"\"\n try:\n settings = load_settings()\n if settings.search_time_image_analysis_enabled is not None:\n return settings.search_time_image_analysis_enabled\n except Exception:\n pass\n\n return False\n\n\ndef get_image_analysis_max_size_mb() -> int:\n \"\"\"Get image analysis max size MB setting from workspace settings or fallback to environment variable\"\"\"\n try:\n settings = load_settings()\n if settings.image_analysis_max_size_mb is not None:\n return settings.image_analysis_max_size_mb\n except Exception:\n pass\n\n return DEFAULT_IMAGE_ANALYSIS_MAX_SIZE_MB\n" }, { "path": "backend/onyx/configs/model_configs.py", "content": "import json\nimport os\n\n#####\n# Embedding/Reranking Model Configs\n#####\n# Important considerations when choosing models\n# Max tokens count needs to be high considering use case (at least 512)\n# Models used must be MIT or Apache license\n# Inference/Indexing speed\n# https://huggingface.co/DOCUMENT_ENCODER_MODEL\n# The useable models configured as below must be SentenceTransformer compatible\n# NOTE: DO NOT CHANGE SET THESE UNLESS YOU KNOW WHAT YOU ARE DOING\n# IDEALLY, YOU SHOULD CHANGE EMBEDDING MODELS VIA THE UI\nDEFAULT_DOCUMENT_ENCODER_MODEL = \"nomic-ai/nomic-embed-text-v1\"\nDOCUMENT_ENCODER_MODEL = (\n os.environ.get(\"DOCUMENT_ENCODER_MODEL\") or DEFAULT_DOCUMENT_ENCODER_MODEL\n)\n# If the below is changed, Vespa deployment must also be changed\nDOC_EMBEDDING_DIM = int(os.environ.get(\"DOC_EMBEDDING_DIM\") or 768)\nNORMALIZE_EMBEDDINGS = (\n os.environ.get(\"NORMALIZE_EMBEDDINGS\") or \"true\"\n).lower() == \"true\"\n\n# Old default model settings, which are needed for an automatic easy upgrade\nOLD_DEFAULT_DOCUMENT_ENCODER_MODEL = \"thenlper/gte-small\"\nOLD_DEFAULT_MODEL_DOC_EMBEDDING_DIM = 384\nOLD_DEFAULT_MODEL_NORMALIZE_EMBEDDINGS = False\n\n# These are only used if reranking is turned off, to normalize the direct retrieval scores for display\n# Currently unused\nSIM_SCORE_RANGE_LOW = float(os.environ.get(\"SIM_SCORE_RANGE_LOW\") or 0.0)\nSIM_SCORE_RANGE_HIGH = float(os.environ.get(\"SIM_SCORE_RANGE_HIGH\") or 1.0)\n# Certain models like e5, BGE, etc use a prefix for asymmetric retrievals (query generally shorter than docs)\nASYM_QUERY_PREFIX = os.environ.get(\"ASYM_QUERY_PREFIX\", \"search_query: \")\nASYM_PASSAGE_PREFIX = os.environ.get(\"ASYM_PASSAGE_PREFIX\", \"search_document: \")\n# Purely an optimization, memory limitation consideration\n\n# User's set embedding batch size overrides the default encoding batch sizes\nEMBEDDING_BATCH_SIZE = int(os.environ.get(\"EMBEDDING_BATCH_SIZE\") or 0) or None\n\nBATCH_SIZE_ENCODE_CHUNKS = EMBEDDING_BATCH_SIZE or 8\n# don't send over too many chunks at once, as sending too many could cause timeouts\nBATCH_SIZE_ENCODE_CHUNKS_FOR_API_EMBEDDING_SERVICES = EMBEDDING_BATCH_SIZE or 512\n# For score display purposes, only way is to know the expected ranges\nCROSS_ENCODER_RANGE_MAX = 1\nCROSS_ENCODER_RANGE_MIN = 0\n\n\n#####\n# Generative AI Model Configs\n#####\n\n# NOTE: the 2 below should only be used for dev.\nGEN_AI_API_KEY = os.environ.get(\"GEN_AI_API_KEY\")\nGEN_AI_MODEL_VERSION = os.environ.get(\"GEN_AI_MODEL_VERSION\")\n\n# Override the auto-detection of LLM max context length\nGEN_AI_MAX_TOKENS = int(os.environ.get(\"GEN_AI_MAX_TOKENS\") or 0) or None\n\n# Set this to be enough for an answer + quotes. Also used for Chat\n# This is the minimum token context we will leave for the LLM to generate an answer\nGEN_AI_NUM_RESERVED_OUTPUT_TOKENS = int(\n os.environ.get(\"GEN_AI_NUM_RESERVED_OUTPUT_TOKENS\") or 1024\n)\n\n# Fallback token limit for models where the max context is unknown\n# Set conservatively at 32K to handle most modern models\nGEN_AI_MODEL_FALLBACK_MAX_TOKENS = int(\n os.environ.get(\"GEN_AI_MODEL_FALLBACK_MAX_TOKENS\") or 32000\n)\n\n# This is used when computing how much context space is available for documents\n# ahead of time in order to let the user know if they can \"select\" more documents\n# It represents a maximum \"expected\" number of input tokens from the latest user\n# message. At query time, we don't actually enforce this - we will only throw an\n# error if the total # of tokens exceeds the max input tokens.\nGEN_AI_SINGLE_USER_MESSAGE_EXPECTED_MAX_TOKENS = 512\nGEN_AI_TEMPERATURE = float(os.environ.get(\"GEN_AI_TEMPERATURE\") or 0)\n\n# should be used if you are using a custom LLM inference provider that doesn't support\n# streaming format AND you are still using the langchain/litellm LLM class\nDISABLE_LITELLM_STREAMING = (\n os.environ.get(\"DISABLE_LITELLM_STREAMING\") or \"false\"\n).lower() == \"true\"\n\n# extra headers to pass to LiteLLM\nLITELLM_EXTRA_HEADERS: dict[str, str] | None = None\n_LITELLM_EXTRA_HEADERS_RAW = os.environ.get(\"LITELLM_EXTRA_HEADERS\")\nif _LITELLM_EXTRA_HEADERS_RAW:\n try:\n LITELLM_EXTRA_HEADERS = json.loads(_LITELLM_EXTRA_HEADERS_RAW)\n except Exception:\n # need to import here to avoid circular imports\n from onyx.utils.logger import setup_logger\n\n logger = setup_logger()\n logger.error(\n \"Failed to parse LITELLM_EXTRA_HEADERS, must be a valid JSON object\"\n )\n\n# if specified, will pass through request headers to the call to the LLM\nLITELLM_PASS_THROUGH_HEADERS: list[str] | None = None\n_LITELLM_PASS_THROUGH_HEADERS_RAW = os.environ.get(\"LITELLM_PASS_THROUGH_HEADERS\")\nif _LITELLM_PASS_THROUGH_HEADERS_RAW:\n try:\n LITELLM_PASS_THROUGH_HEADERS = json.loads(_LITELLM_PASS_THROUGH_HEADERS_RAW)\n except Exception:\n # need to import here to avoid circular imports\n from onyx.utils.logger import setup_logger\n\n logger = setup_logger()\n logger.error(\n \"Failed to parse LITELLM_PASS_THROUGH_HEADERS, must be a valid JSON object\"\n )\n\n\n# if specified, will merge the specified JSON with the existing body of the\n# request before sending it to the LLM\nLITELLM_EXTRA_BODY: dict | None = None\n_LITELLM_EXTRA_BODY_RAW = os.environ.get(\"LITELLM_EXTRA_BODY\")\nif _LITELLM_EXTRA_BODY_RAW:\n try:\n LITELLM_EXTRA_BODY = json.loads(_LITELLM_EXTRA_BODY_RAW)\n except Exception:\n pass\n\n#####\n# Prompt Caching Configs\n#####\n# Enable prompt caching framework\nENABLE_PROMPT_CACHING = (\n os.environ.get(\"ENABLE_PROMPT_CACHING\", \"true\").lower() != \"false\"\n)\n\n# Cache TTL multiplier - store caches slightly longer than provider TTL\n# This allows for some clock skew and ensures we don't lose cache metadata prematurely\nPROMPT_CACHE_REDIS_TTL_MULTIPLIER = float(\n os.environ.get(\"PROMPT_CACHE_REDIS_TTL_MULTIPLIER\") or 1.2\n)\n" }, { "path": "backend/onyx/configs/onyxbot_configs.py", "content": "import os\n\n#####\n# Onyx Slack Bot Configs\n#####\nONYX_BOT_NUM_RETRIES = int(os.environ.get(\"ONYX_BOT_NUM_RETRIES\", \"5\"))\n# Number of docs to display in \"Reference Documents\"\nONYX_BOT_NUM_DOCS_TO_DISPLAY = int(os.environ.get(\"ONYX_BOT_NUM_DOCS_TO_DISPLAY\", \"5\"))\n# If the LLM fails to answer, Onyx can still show the \"Reference Documents\"\nONYX_BOT_DISABLE_DOCS_ONLY_ANSWER = os.environ.get(\n \"ONYX_BOT_DISABLE_DOCS_ONLY_ANSWER\", \"\"\n).lower() not in [\"false\", \"\"]\n# When Onyx is considering a message, what emoji does it react with\nONYX_BOT_REACT_EMOJI = os.environ.get(\"ONYX_BOT_REACT_EMOJI\") or \"eyes\"\n# When User needs more help, what should the emoji be\nONYX_BOT_FOLLOWUP_EMOJI = os.environ.get(\"ONYX_BOT_FOLLOWUP_EMOJI\") or \"sos\"\n# What kind of message should be shown when someone gives an AI answer feedback to OnyxBot\n# Defaults to Private if not provided or invalid\n# Private: Only visible to user clicking the feedback\n# Anonymous: Public but anonymous\n# Public: Visible with the user name who submitted the feedback\nONYX_BOT_FEEDBACK_VISIBILITY = (\n os.environ.get(\"ONYX_BOT_FEEDBACK_VISIBILITY\") or \"private\"\n)\n# Should OnyxBot send an apology message if it's not able to find an answer\n# That way the user isn't confused as to why OnyxBot reacted but then said nothing\n# Off by default to be less intrusive (don't want to give a notif that just says we couldnt help)\nNOTIFY_SLACKBOT_NO_ANSWER = (\n os.environ.get(\"NOTIFY_SLACKBOT_NO_ANSWER\", \"\").lower() == \"true\"\n)\n# Mostly for debugging purposes but it's for explaining what went wrong\n# if OnyxBot couldn't find an answer\nONYX_BOT_DISPLAY_ERROR_MSGS = os.environ.get(\n \"ONYX_BOT_DISPLAY_ERROR_MSGS\", \"\"\n).lower() not in [\n \"false\",\n \"\",\n]\n\n# Maximum Questions Per Minute, Default Uncapped\nONYX_BOT_MAX_QPM = int(os.environ.get(\"ONYX_BOT_MAX_QPM\") or 0) or None\n# Maximum time to wait when a question is queued\nONYX_BOT_MAX_WAIT_TIME = int(os.environ.get(\"ONYX_BOT_MAX_WAIT_TIME\") or 180)\n\n# Time (in minutes) after which a Slack message is sent to the user to remind him to give feedback.\n# Set to 0 to disable it (default)\nONYX_BOT_FEEDBACK_REMINDER = int(os.environ.get(\"ONYX_BOT_FEEDBACK_REMINDER\") or 0)\n\n# ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD is the number of\n# responses OnyxBot can send in a given time period.\n# Set to 0 to disable the limit.\nONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD = int(\n os.environ.get(\"ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD\", \"5000\")\n)\n# ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS is the number\n# of seconds until the response limit is reset.\nONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS = int(\n os.environ.get(\"ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS\", \"86400\")\n)\n" }, { "path": "backend/onyx/configs/research_configs.py", "content": "" }, { "path": "backend/onyx/configs/saml_config/template.settings.json", "content": "{\n \"strict\": true,\n \"debug\": false,\n \"idp\": {\n \"entityId\": \"\",\n \"singleSignOnService\": {\n \"url\": \" https://trial-1234567.okta.com/home/trial-1234567_onyx/somevalues/somevalues\",\n \"binding\": \"urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect\"\n },\n \"x509cert\": \"\"\n },\n \"sp\": {\n \"entityId\": \"\",\n \"assertionConsumerService\": {\n \"url\": \"http://127.0.0.1:3000/auth/saml/callback\",\n \"binding\": \"urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST\"\n },\n \"x509cert\": \"\"\n }\n}\n" }, { "path": "backend/onyx/configs/tool_configs.py", "content": "import json\nimport os\n\n\nIMAGE_GENERATION_OUTPUT_FORMAT = os.environ.get(\n \"IMAGE_GENERATION_OUTPUT_FORMAT\", \"b64_json\"\n)\n\n# if specified, will pass through request headers to the call to API calls made by custom tools\nCUSTOM_TOOL_PASS_THROUGH_HEADERS: list[str] | None = None\n_CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW = os.environ.get(\n \"CUSTOM_TOOL_PASS_THROUGH_HEADERS\"\n)\nif _CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW:\n try:\n CUSTOM_TOOL_PASS_THROUGH_HEADERS = json.loads(\n _CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW\n )\n except Exception:\n # need to import here to avoid circular imports\n from onyx.utils.logger import setup_logger\n\n logger = setup_logger()\n logger.error(\n \"Failed to parse CUSTOM_TOOL_PASS_THROUGH_HEADERS, must be a valid JSON object\"\n )\n" }, { "path": "backend/onyx/connectors/README.md", "content": "\n\n# Writing a new Onyx Connector\n\nThis README covers how to contribute a new Connector for Onyx. It includes an overview of the design, interfaces,\nand required changes.\n\nThank you for your contribution!\n\n### Connector Overview\n\nConnectors come in 3 different flows:\n\n- Load Connector:\n - Bulk indexes documents to reflect a point in time. This type of connector generally works by either pulling all\n documents via a connector's API or loads the documents from some sort of a dump file.\n- Poll Connector:\n - Incrementally updates documents based on a provided time range. It is used by the background job to pull the latest\n changes and additions since the last round of polling. This connector helps keep the document index up to date\n without needing to fetch/embed/index every document which would be too slow to do frequently on large sets of\n documents.\n- Slim Connector:\n - This connector should be a lighter weight method of checking all documents in the source to see if they still exist.\n - This connector should be identical to the Poll or Load Connector except that it only fetches the IDs of the documents, not the documents themselves.\n - This is used by our pruning job which removes old documents from the index.\n - The optional start and end datetimes can be ignored.\n- Event Based connectors:\n - Connectors that listen to events and update documents accordingly.\n - Currently not used by the background job, this exists for future design purposes.\n\n### Connector Implementation\n\nRefer to [interfaces.py](https://github.com/onyx-dot-app/onyx/blob/main/backend/onyx/connectors/interfaces.py)\nand this first contributor created Pull Request for a new connector (Shoutout to Dan Brown):\n[Reference Pull Request](https://github.com/onyx-dot-app/onyx/pull/139)\n\nFor implementing a Slim Connector, refer to the comments in this PR:\n[Slim Connector PR](https://github.com/onyx-dot-app/onyx/pull/3303/files)\n\nAll new connectors should have tests added to the `backend/tests/daily/connectors` directory. Refer to the above PR for an example of adding tests for a new connector.\n\n#### Implementing the new Connector\n\nThe connector must subclass one or more of LoadConnector, PollConnector, CheckpointedConnector, or CheckpointedConnectorWithPermSync\n\nThe `__init__` should take arguments for configuring what documents the connector will and where it finds those\ndocuments. For example, if you have a wiki site, it may include the configuration for the team, topic, folder, etc. of\nthe documents to fetch. It may also include the base domain of the wiki. Alternatively, if all the access information\nof the connector is stored in the credential/token, then there may be no required arguments.\n\n`load_credentials` should take a dictionary which provides all the access information that the connector might need.\nFor example this could be the user's username and access token.\n\nRefer to the existing connectors for `load_from_state` and `poll_source` examples. There is not yet a process to listen\nfor EventConnector events, this will come down the line.\n\n#### Development Tip\n\nIt may be handy to test your new connector separate from the rest of the stack while developing.\nFollow the below template:\n\n```commandline\nif __name__ == \"__main__\":\n import time\n test_connector = NewConnector(space=\"engineering\")\n test_connector.load_credentials({\n \"user_id\": \"foobar\",\n \"access_token\": \"fake_token\"\n })\n all_docs = test_connector.load_from_state()\n\n current = time.time()\n one_day_ago = current - 24 * 60 * 60 # 1 day\n latest_docs = test_connector.poll_source(one_day_ago, current)\n```\n\n> Note: Be sure to set PYTHONPATH to onyx/backend before running the above main.\n\n### Additional Required Changes:\n\n#### Backend Changes\n\n- Add a new type to\n [DocumentSource](https://github.com/onyx-dot-app/onyx/blob/main/backend/onyx/configs/constants.py)\n- Add a mapping from DocumentSource (and optionally connector type) to the right connector class\n [here](https://github.com/onyx-dot-app/onyx/blob/main/backend/onyx/connectors/factory.py#L33)\n\n#### Frontend Changes\n\n- Add the new Connector definition to the `SOURCE_METADATA_MAP` [here](https://github.com/onyx-dot-app/onyx/blob/main/web/src/lib/sources.ts#L59).\n- Add the definition for the new Form to the `connectorConfigs` object [here](https://github.com/onyx-dot-app/onyx/blob/main/web/src/lib/connectors/connectors.ts#L79).\n\n#### Docs Changes\n\nCreate the new connector page (with guiding images!) with how to get the connector credentials and how to set up the\nconnector in Onyx. Then create a Pull Request in [https://github.com/onyx-dot-app/documentation](https://github.com/onyx-dot-app/documentation).\n\n### Before opening PR\n\n1. Be sure to fully test changes end to end with setting up the connector and updating the index with new docs from the\n new connector. To make it easier to review, please attach a video showing the successful creation of the connector via the UI (starting from the `Add Connector` page).\n2. Add a folder + tests under `backend/tests/daily/connectors` director. For an example, checkout the [test for Confluence](https://github.com/onyx-dot-app/onyx/blob/main/backend/tests/daily/connectors/confluence/test_confluence_basic.py). In the PR description, include a guide on how to setup the new source to pass the test. Before merging, we will re-create the environment and make sure the test(s) pass.\n3. Be sure to run the linting/formatting, refer to the formatting and linting section in\n [CONTRIBUTING.md](https://github.com/onyx-dot-app/onyx/blob/main/CONTRIBUTING.md#formatting-and-linting)\n" }, { "path": "backend/onyx/connectors/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/airtable/airtable_connector.py", "content": "import contextvars\nimport re\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import Future\nfrom concurrent.futures import ThreadPoolExecutor\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom pyairtable import Api as AirtableApi\nfrom pyairtable.api.types import RecordDict\nfrom pyairtable.models.schema import TableSchema\nfrom retry import retry\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# NOTE: all are made lowercase to avoid case sensitivity issues\n# These field types are considered metadata by default when\n# treat_all_non_attachment_fields_as_metadata is False\nDEFAULT_METADATA_FIELD_TYPES = {\n \"singlecollaborator\",\n \"collaborator\",\n \"createdby\",\n \"singleselect\",\n \"multipleselects\",\n \"checkbox\",\n \"date\",\n \"datetime\",\n \"email\",\n \"phone\",\n \"url\",\n \"number\",\n \"currency\",\n \"duration\",\n \"percent\",\n \"rating\",\n \"createdtime\",\n \"lastmodifiedtime\",\n \"autonumber\",\n \"rollup\",\n \"lookup\",\n \"count\",\n \"formula\",\n \"date\",\n}\n\n\nclass AirtableClientNotSetUpError(PermissionError):\n def __init__(self) -> None:\n super().__init__(\"Airtable Client is not set up, was load_credentials called?\")\n\n\n# Matches URLs like https://airtable.com/appXXX/tblYYY/viwZZZ?blocks=hide\n# Captures: base_id (appXXX), table_id (tblYYY), and optionally view_id (viwZZZ)\n_AIRTABLE_URL_PATTERN = re.compile(\n r\"https?://airtable\\.com/(app[A-Za-z0-9]+)/(tbl[A-Za-z0-9]+)(?:/(viw[A-Za-z0-9]+))?\",\n)\n\n\ndef parse_airtable_url(\n url: str,\n) -> tuple[str, str, str | None]:\n \"\"\"Parse an Airtable URL into (base_id, table_id, view_id).\n\n Accepts URLs like:\n https://airtable.com/appXXX/tblYYY\n https://airtable.com/appXXX/tblYYY/viwZZZ\n https://airtable.com/appXXX/tblYYY/viwZZZ?blocks=hide\n\n Returns:\n (base_id, table_id, view_id or None)\n\n Raises:\n ValueError if the URL doesn't match the expected format.\n \"\"\"\n match = _AIRTABLE_URL_PATTERN.search(url.strip())\n if not match:\n raise ValueError(\n f\"Could not parse Airtable URL: '{url}'. Expected format: https://airtable.com/appXXX/tblYYY[/viwZZZ]\"\n )\n return match.group(1), match.group(2), match.group(3)\n\n\nclass AirtableConnector(LoadConnector):\n def __init__(\n self,\n base_id: str = \"\",\n table_name_or_id: str = \"\",\n airtable_url: str = \"\",\n treat_all_non_attachment_fields_as_metadata: bool = False,\n view_id: str | None = None,\n share_id: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n \"\"\"Initialize an AirtableConnector.\n\n Args:\n base_id: The ID of the Airtable base (not required when airtable_url is set)\n table_name_or_id: The name or ID of the table (not required when airtable_url is set)\n airtable_url: An Airtable URL to parse base_id, table_id, and view_id from.\n Overrides base_id, table_name_or_id, and view_id if provided.\n treat_all_non_attachment_fields_as_metadata: If True, all fields except attachments will be treated as metadata.\n If False, only fields with types in DEFAULT_METADATA_FIELD_TYPES will be treated as metadata.\n view_id: Optional ID of a specific view to use\n share_id: Optional ID of a \"share\" to use for generating record URLs\n batch_size: Number of records to process in each batch\n\n Mode is auto-detected: if a specific table is identified (via URL or\n base_id + table_name_or_id), the connector indexes that single table.\n Otherwise, it discovers and indexes all accessible bases and tables.\n \"\"\"\n # If a URL is provided, parse it to extract base_id, table_id, and view_id\n if airtable_url:\n parsed_base_id, parsed_table_id, parsed_view_id = parse_airtable_url(\n airtable_url\n )\n base_id = parsed_base_id\n table_name_or_id = parsed_table_id\n if parsed_view_id:\n view_id = parsed_view_id\n\n self.base_id = base_id\n self.table_name_or_id = table_name_or_id\n self.index_all = not (base_id and table_name_or_id)\n self.view_id = view_id\n self.share_id = share_id\n self.batch_size = batch_size\n self._airtable_client: AirtableApi | None = None\n self.treat_all_non_attachment_fields_as_metadata = (\n treat_all_non_attachment_fields_as_metadata\n )\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self._airtable_client = AirtableApi(credentials[\"airtable_access_token\"])\n return None\n\n @property\n def airtable_client(self) -> AirtableApi:\n if not self._airtable_client:\n raise AirtableClientNotSetUpError()\n return self._airtable_client\n\n def validate_connector_settings(self) -> None:\n if self.index_all:\n try:\n bases = self.airtable_client.bases()\n if not bases:\n raise ConnectorValidationError(\n \"No bases found. Ensure your API token has access to at least one base.\"\n )\n except ConnectorValidationError:\n raise\n except Exception as e:\n raise ConnectorValidationError(f\"Failed to list Airtable bases: {e}\")\n else:\n if not self.base_id or not self.table_name_or_id:\n raise ConnectorValidationError(\n \"A valid Airtable URL or base_id and table_name_or_id are required when not using index_all mode.\"\n )\n try:\n table = self.airtable_client.table(self.base_id, self.table_name_or_id)\n table.schema()\n except Exception as e:\n raise ConnectorValidationError(\n f\"Failed to access table '{self.table_name_or_id}' in base '{self.base_id}': {e}\"\n )\n\n @classmethod\n def _get_record_url(\n cls,\n base_id: str,\n table_id: str,\n record_id: str,\n share_id: str | None,\n view_id: str | None,\n field_id: str | None = None,\n attachment_id: str | None = None,\n ) -> str:\n \"\"\"Constructs the URL for a record, optionally including field and attachment IDs\n\n Full possible structure is:\n\n https://airtable.com/BASE_ID/SHARE_ID/TABLE_ID/VIEW_ID/RECORD_ID/FIELD_ID/ATTACHMENT_ID\n \"\"\"\n # If we have a shared link, use that view for better UX\n if share_id:\n base_url = f\"https://airtable.com/{base_id}/{share_id}/{table_id}\"\n else:\n base_url = f\"https://airtable.com/{base_id}/{table_id}\"\n\n if view_id:\n base_url = f\"{base_url}/{view_id}\"\n\n base_url = f\"{base_url}/{record_id}\"\n\n if field_id and attachment_id:\n return f\"{base_url}/{field_id}/{attachment_id}?blocks=hide\"\n\n return base_url\n\n def _extract_field_values(\n self,\n field_id: str,\n field_name: str,\n field_info: Any,\n field_type: str,\n base_id: str,\n table_id: str,\n view_id: str | None,\n record_id: str,\n ) -> list[tuple[str, str]]:\n \"\"\"\n Extract value(s) + links from a field regardless of its type.\n Attachments are represented as multiple sections, and therefore\n returned as a list of tuples (value, link).\n \"\"\"\n if field_info is None:\n return []\n\n # skip references to other records for now (would need to do another\n # request to get the actual record name/type)\n # TODO: support this\n if field_type == \"multipleRecordLinks\":\n return []\n\n # Get the base URL for this record\n default_link = self._get_record_url(\n base_id, table_id, record_id, self.share_id, self.view_id or view_id\n )\n\n if field_type == \"multipleAttachments\":\n attachment_texts: list[tuple[str, str]] = []\n for attachment in field_info:\n url = attachment.get(\"url\")\n filename = attachment.get(\"filename\", \"\")\n if not url:\n continue\n\n @retry(\n tries=5,\n delay=1,\n backoff=2,\n max_delay=10,\n )\n def get_attachment_with_retry(url: str, record_id: str) -> bytes | None:\n try:\n attachment_response = requests.get(url)\n attachment_response.raise_for_status()\n return attachment_response.content\n except requests.exceptions.HTTPError as e:\n if e.response.status_code == 410:\n logger.info(f\"Refreshing attachment for {filename}\")\n # Re-fetch the record to get a fresh URL\n refreshed_record = self.airtable_client.table(\n base_id, table_id\n ).get(record_id)\n for refreshed_attachment in refreshed_record[\"fields\"][\n field_name\n ]:\n if refreshed_attachment.get(\"filename\") == filename:\n new_url = refreshed_attachment.get(\"url\")\n if new_url:\n attachment_response = requests.get(new_url)\n attachment_response.raise_for_status()\n return attachment_response.content\n\n logger.error(f\"Failed to refresh attachment for {filename}\")\n raise\n\n attachment_content = get_attachment_with_retry(url, record_id)\n if attachment_content:\n try:\n file_ext = get_file_ext(filename)\n attachment_id = attachment[\"id\"]\n attachment_text = extract_file_text(\n BytesIO(attachment_content),\n filename,\n break_on_unprocessable=False,\n extension=file_ext,\n )\n if attachment_text:\n # Use the helper method to construct attachment URLs\n attachment_link = self._get_record_url(\n base_id,\n table_id,\n record_id,\n self.share_id,\n self.view_id or view_id,\n field_id,\n attachment_id,\n )\n attachment_texts.append(\n (f\"{filename}:\\n{attachment_text}\", attachment_link)\n )\n except Exception as e:\n logger.warning(\n f\"Failed to process attachment {filename}: {str(e)}\"\n )\n return attachment_texts\n\n if field_type in [\"singleCollaborator\", \"collaborator\", \"createdBy\"]:\n combined = []\n collab_name = field_info.get(\"name\")\n collab_email = field_info.get(\"email\")\n if collab_name:\n combined.append(collab_name)\n if collab_email:\n combined.append(f\"({collab_email})\")\n return [(\" \".join(combined) if combined else str(field_info), default_link)]\n\n if isinstance(field_info, list):\n return [(str(item), default_link) for item in field_info]\n\n return [(str(field_info), default_link)]\n\n def _should_be_metadata(self, field_type: str) -> bool:\n \"\"\"Determine if a field type should be treated as metadata.\n\n When treat_all_non_attachment_fields_as_metadata is True, all fields except\n attachments are treated as metadata. Otherwise, only fields with types listed\n in DEFAULT_METADATA_FIELD_TYPES are treated as metadata.\"\"\"\n if self.treat_all_non_attachment_fields_as_metadata:\n return field_type.lower() != \"multipleattachments\"\n return field_type.lower() in DEFAULT_METADATA_FIELD_TYPES\n\n def _process_field(\n self,\n field_id: str,\n field_name: str,\n field_info: Any,\n field_type: str,\n base_id: str,\n table_id: str,\n view_id: str | None,\n record_id: str,\n ) -> tuple[list[TextSection], dict[str, str | list[str]]]:\n \"\"\"\n Process a single Airtable field and return sections or metadata.\n\n Args:\n field_name: Name of the field\n field_info: Raw field information from Airtable\n field_type: Airtable field type\n\n Returns:\n (list of Sections, dict of metadata)\n \"\"\"\n if field_info is None:\n return [], {}\n\n # Get the value(s) for the field\n field_value_and_links = self._extract_field_values(\n field_id=field_id,\n field_name=field_name,\n field_info=field_info,\n field_type=field_type,\n base_id=base_id,\n table_id=table_id,\n view_id=view_id,\n record_id=record_id,\n )\n if len(field_value_and_links) == 0:\n return [], {}\n\n # Determine if it should be metadata or a section\n if self._should_be_metadata(field_type):\n field_values = [value for value, _ in field_value_and_links]\n if len(field_values) > 1:\n return [], {field_name: field_values}\n return [], {field_name: field_values[0]}\n\n # Otherwise, create relevant sections\n sections = [\n TextSection(\n link=link,\n text=(\n f\"{field_name}:\\n------------------------\\n{text}\\n------------------------\"\n ),\n )\n for text, link in field_value_and_links\n ]\n return sections, {}\n\n def _process_record(\n self,\n record: RecordDict,\n table_schema: TableSchema,\n primary_field_name: str | None,\n base_id: str,\n base_name: str | None = None,\n ) -> Document | None:\n \"\"\"Process a single Airtable record into a Document.\n\n Args:\n record: The Airtable record to process\n table_schema: Schema information for the table\n primary_field_name: Name of the primary field, if any\n base_id: The ID of the base this record belongs to\n base_name: The name of the base (used in semantic ID for index_all mode)\n\n Returns:\n Document object representing the record\n \"\"\"\n table_id = table_schema.id\n table_name = table_schema.name\n record_id = record[\"id\"]\n fields = record[\"fields\"]\n sections: list[TextSection] = []\n metadata: dict[str, str | list[str]] = {}\n\n # Get primary field value if it exists\n primary_field_value = (\n fields.get(primary_field_name) if primary_field_name else None\n )\n view_id = table_schema.views[0].id if table_schema.views else None\n\n for field_schema in table_schema.fields:\n field_name = field_schema.name\n field_val = fields.get(field_name)\n field_type = field_schema.type\n\n logger.debug(\n f\"Processing field '{field_name}' of type '{field_type}' for record '{record_id}'.\"\n )\n\n field_sections, field_metadata = self._process_field(\n field_id=field_schema.id,\n field_name=field_name,\n field_info=field_val,\n field_type=field_type,\n base_id=base_id,\n table_id=table_id,\n view_id=view_id,\n record_id=record_id,\n )\n\n sections.extend(field_sections)\n metadata.update(field_metadata)\n\n if not sections:\n logger.warning(f\"No sections found for record {record_id}\")\n return None\n\n # Include base name in semantic ID only in index_all mode\n if self.index_all and base_name:\n semantic_id = (\n f\"{base_name} > {table_name}: {primary_field_value}\"\n if primary_field_value\n else f\"{base_name} > {table_name}\"\n )\n else:\n semantic_id = (\n f\"{table_name}: {primary_field_value}\"\n if primary_field_value\n else table_name\n )\n\n # Build hierarchy source_path for Craft file system subdirectory structure.\n # This creates: airtable/{base_name}/{table_name}/record.json\n source_path: list[str] = []\n if base_name:\n source_path.append(base_name)\n source_path.append(table_name)\n\n return Document(\n id=f\"airtable__{record_id}\",\n sections=(cast(list[TextSection | ImageSection], sections)),\n source=DocumentSource.AIRTABLE,\n semantic_identifier=semantic_id,\n metadata=metadata,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": source_path,\n \"base_id\": base_id,\n \"table_id\": table_id,\n \"table_name\": table_name,\n **({\"base_name\": base_name} if base_name else {}),\n }\n },\n )\n\n def _resolve_base_name(self, base_id: str) -> str | None:\n \"\"\"Try to resolve a human-readable base name from the API.\"\"\"\n try:\n for base_info in self.airtable_client.bases():\n if base_info.id == base_id:\n return base_info.name\n except Exception:\n logger.debug(f\"Could not resolve base name for {base_id}\")\n return None\n\n def _index_table(\n self,\n base_id: str,\n table_name_or_id: str,\n base_name: str | None = None,\n ) -> GenerateDocumentsOutput:\n \"\"\"Index all records from a single table. Yields batches of Documents.\"\"\"\n # Resolve base name for hierarchy if not provided\n if base_name is None:\n base_name = self._resolve_base_name(base_id)\n\n table = self.airtable_client.table(base_id, table_name_or_id)\n records = table.all()\n\n table_schema = table.schema()\n primary_field_name = None\n\n # Find a primary field from the schema\n for field in table_schema.fields:\n if field.id == table_schema.primary_field_id:\n primary_field_name = field.name\n break\n\n logger.info(\n f\"Processing {len(records)} records from table '{table_schema.name}' in base '{base_name or base_id}'.\"\n )\n\n if not records:\n return\n\n # Process records in parallel batches using ThreadPoolExecutor\n PARALLEL_BATCH_SIZE = 8\n max_workers = min(PARALLEL_BATCH_SIZE, len(records))\n\n for i in range(0, len(records), PARALLEL_BATCH_SIZE):\n batch_records = records[i : i + PARALLEL_BATCH_SIZE]\n record_documents: list[Document | HierarchyNode] = []\n\n with ThreadPoolExecutor(max_workers=max_workers) as executor:\n # Submit batch tasks\n future_to_record: dict[Future[Document | None], RecordDict] = {}\n for record in batch_records:\n # Capture the current context so that the thread gets the current tenant ID\n current_context = contextvars.copy_context()\n future_to_record[\n executor.submit(\n current_context.run,\n self._process_record,\n record=record,\n table_schema=table_schema,\n primary_field_name=primary_field_name,\n base_id=base_id,\n base_name=base_name,\n )\n ] = record\n\n # Wait for all tasks in this batch to complete\n for future in as_completed(future_to_record):\n record = future_to_record[future]\n try:\n document = future.result()\n if document:\n record_documents.append(document)\n except Exception as e:\n logger.exception(f\"Failed to process record {record['id']}\")\n raise e\n\n if record_documents:\n yield record_documents\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"\n Fetch all records from one or all tables.\n\n NOTE: Airtable does not support filtering by time updated, so\n we have to fetch all records every time.\n \"\"\"\n if not self.airtable_client:\n raise AirtableClientNotSetUpError()\n\n if self.index_all:\n yield from self._load_all()\n else:\n yield from self._index_table(\n base_id=self.base_id,\n table_name_or_id=self.table_name_or_id,\n )\n\n def _load_all(self) -> GenerateDocumentsOutput:\n \"\"\"Discover all bases and tables, then index everything.\"\"\"\n bases = self.airtable_client.bases()\n logger.info(f\"Discovered {len(bases)} Airtable base(s).\")\n\n for base_info in bases:\n base_id = base_info.id\n base_name = base_info.name\n logger.info(f\"Listing tables for base '{base_name}' ({base_id}).\")\n\n try:\n base = self.airtable_client.base(base_id)\n tables = base.tables()\n except Exception:\n logger.exception(\n f\"Failed to list tables for base '{base_name}' ({base_id}), skipping.\"\n )\n continue\n\n logger.info(f\"Found {len(tables)} table(s) in base '{base_name}'.\")\n\n for table in tables:\n try:\n yield from self._index_table(\n base_id=base_id,\n table_name_or_id=table.id,\n base_name=base_name,\n )\n except Exception:\n logger.exception(\n f\"Failed to index table '{table.name}' ({table.id}) in base '{base_name}' ({base_id}), skipping.\"\n )\n continue\n" }, { "path": "backend/onyx/connectors/asana/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/asana/asana_api.py", "content": "import time\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom typing import Dict\n\nimport asana # type: ignore\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# https://github.com/Asana/python-asana/tree/master?tab=readme-ov-file#documentation-for-api-endpoints\nclass AsanaTask:\n def __init__(\n self,\n id: str,\n title: str,\n text: str,\n link: str,\n last_modified: datetime,\n project_gid: str,\n project_name: str,\n ) -> None:\n self.id = id\n self.title = title\n self.text = text\n self.link = link\n self.last_modified = last_modified\n self.project_gid = project_gid\n self.project_name = project_name\n\n def __str__(self) -> str:\n return f\"ID: {self.id}\\nTitle: {self.title}\\nLast modified: {self.last_modified}\\nText: {self.text}\"\n\n\nclass AsanaAPI:\n def __init__(\n self, api_token: str, workspace_gid: str, team_gid: str | None\n ) -> None:\n self._user = None\n self.workspace_gid = workspace_gid\n self.team_gid = team_gid\n\n self.configuration = asana.Configuration()\n self.api_client = asana.ApiClient(self.configuration)\n self.tasks_api = asana.TasksApi(self.api_client)\n self.stories_api = asana.StoriesApi(self.api_client)\n self.users_api = asana.UsersApi(self.api_client)\n self.project_api = asana.ProjectsApi(self.api_client)\n self.workspaces_api = asana.WorkspacesApi(self.api_client)\n\n self.api_error_count = 0\n self.configuration.access_token = api_token\n self.task_count = 0\n\n def get_tasks(\n self, project_gids: list[str] | None, start_date: str\n ) -> Iterator[AsanaTask]:\n \"\"\"Get all tasks from the projects with the given gids that were modified since the given date.\n If project_gids is None, get all tasks from all projects in the workspace.\"\"\"\n logger.info(\"Starting to fetch Asana projects\")\n projects = self.project_api.get_projects(\n opts={\n \"workspace\": self.workspace_gid,\n \"opt_fields\": \"gid,name,archived,modified_at\",\n }\n )\n start_seconds = int(time.mktime(datetime.now().timetuple()))\n projects_list = []\n project_count = 0\n for project_info in projects:\n project_gid = project_info[\"gid\"]\n if project_gids is None or project_gid in project_gids:\n projects_list.append(project_gid)\n else:\n logger.debug(\n f\"Skipping project: {project_gid} - not in accepted project_gids\"\n )\n project_count += 1\n if project_count % 100 == 0:\n logger.info(f\"Processed {project_count} projects\")\n\n logger.info(f\"Found {len(projects_list)} projects to process\")\n for project_gid in projects_list:\n for task in self._get_tasks_for_project(\n project_gid, start_date, start_seconds\n ):\n yield task\n logger.info(f\"Completed fetching {self.task_count} tasks from Asana\")\n if self.api_error_count > 0:\n logger.warning(\n f\"Encountered {self.api_error_count} API errors during task fetching\"\n )\n\n def _get_tasks_for_project(\n self, project_gid: str, start_date: str, start_seconds: int\n ) -> Iterator[AsanaTask]:\n project = self.project_api.get_project(project_gid, opts={})\n project_name = project.get(\"name\", project_gid)\n team = project.get(\"team\") or {}\n team_gid = team.get(\"gid\")\n\n if project.get(\"archived\"):\n logger.info(f\"Skipping archived project: {project_name} ({project_gid})\")\n return\n if not team_gid:\n logger.info(\n f\"Skipping project without a team: {project_name} ({project_gid})\"\n )\n return\n if project.get(\"privacy_setting\") == \"private\":\n if self.team_gid and team_gid != self.team_gid:\n logger.info(\n f\"Skipping private project not in configured team: {project_name} ({project_gid})\"\n )\n return\n logger.info(\n f\"Processing private project in configured team: {project_name} ({project_gid})\"\n )\n\n simple_start_date = start_date.split(\".\")[0].split(\"+\")[0]\n logger.info(\n f\"Fetching tasks modified since {simple_start_date} for project: {project_name} ({project_gid})\"\n )\n\n opts = {\n \"opt_fields\": \"name,memberships,memberships.project,completed_at,completed_by,created_at,\"\n \"created_by,custom_fields,dependencies,due_at,due_on,external,html_notes,liked,likes,\"\n \"modified_at,notes,num_hearts,parent,projects,resource_subtype,resource_type,start_on,\"\n \"workspace,permalink_url\",\n \"modified_since\": start_date,\n }\n tasks_from_api = self.tasks_api.get_tasks_for_project(project_gid, opts)\n for data in tasks_from_api:\n self.task_count += 1\n if self.task_count % 10 == 0:\n end_seconds = time.mktime(datetime.now().timetuple())\n runtime_seconds = end_seconds - start_seconds\n if runtime_seconds > 0:\n logger.info(\n f\"Processed {self.task_count} tasks in {runtime_seconds:.0f} seconds \"\n f\"({self.task_count / runtime_seconds:.2f} tasks/second)\"\n )\n\n logger.debug(f\"Processing Asana task: {data['name']}\")\n\n text = self._construct_task_text(data)\n\n try:\n text += self._fetch_and_add_comments(data[\"gid\"])\n\n last_modified_date = self.format_date(data[\"modified_at\"])\n text += f\"Last modified: {last_modified_date}\\n\"\n\n task = AsanaTask(\n id=data[\"gid\"],\n title=data[\"name\"],\n text=text,\n link=data[\"permalink_url\"],\n last_modified=datetime.fromisoformat(data[\"modified_at\"]),\n project_gid=project_gid,\n project_name=project_name,\n )\n yield task\n except Exception:\n logger.error(\n f\"Error processing task {data['gid']} in project {project_gid}\",\n exc_info=True,\n )\n self.api_error_count += 1\n\n def _construct_task_text(self, data: Dict) -> str:\n text = f\"{data['name']}\\n\\n\"\n\n if data[\"notes\"]:\n text += f\"{data['notes']}\\n\\n\"\n\n if data[\"created_by\"] and data[\"created_by\"][\"gid\"]:\n creator = self.get_user(data[\"created_by\"][\"gid\"])[\"name\"]\n created_date = self.format_date(data[\"created_at\"])\n text += f\"Created by: {creator} on {created_date}\\n\"\n\n if data[\"due_on\"]:\n due_date = self.format_date(data[\"due_on\"])\n text += f\"Due date: {due_date}\\n\"\n\n if data[\"completed_at\"]:\n completed_date = self.format_date(data[\"completed_at\"])\n text += f\"Completed on: {completed_date}\\n\"\n\n text += \"\\n\"\n return text\n\n def _fetch_and_add_comments(self, task_gid: str) -> str:\n text = \"\"\n stories_opts: Dict[str, str] = {}\n story_start = time.time()\n stories = self.stories_api.get_stories_for_task(task_gid, stories_opts)\n\n story_count = 0\n comment_count = 0\n\n for story in stories:\n story_count += 1\n if story[\"resource_subtype\"] == \"comment_added\":\n comment = self.stories_api.get_story(\n story[\"gid\"], opts={\"opt_fields\": \"text,created_by,created_at\"}\n )\n commenter = self.get_user(comment[\"created_by\"][\"gid\"])[\"name\"]\n text += f\"Comment by {commenter}: {comment['text']}\\n\\n\"\n comment_count += 1\n\n story_duration = time.time() - story_start\n logger.debug(\n f\"Processed {story_count} stories (including {comment_count} comments) in {story_duration:.2f} seconds\"\n )\n\n return text\n\n def get_user(self, user_gid: str) -> Dict:\n if self._user is not None:\n return self._user\n self._user = self.users_api.get_user(user_gid, {\"opt_fields\": \"name,email\"})\n\n if not self._user:\n logger.warning(f\"Unable to fetch user information for user_gid: {user_gid}\")\n return {\"name\": \"Unknown\"}\n return self._user\n\n def format_date(self, date_str: str) -> str:\n date = datetime.fromisoformat(date_str)\n return time.strftime(\"%Y-%m-%d\", date.timetuple())\n\n def get_time(self) -> str:\n return time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime())\n" }, { "path": "backend/onyx/connectors/asana/connector.py", "content": "import datetime\nfrom typing import Any\n\nfrom onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.asana import asana_api\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass AsanaConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n asana_workspace_id: str,\n asana_project_ids: str | None = None,\n asana_team_id: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n continue_on_failure: bool = CONTINUE_ON_CONNECTOR_FAILURE,\n ) -> None:\n self.workspace_id = asana_workspace_id.strip()\n if asana_project_ids:\n project_ids = [\n project_id.strip()\n for project_id in asana_project_ids.split(\",\")\n if project_id.strip()\n ]\n self.project_ids_to_index = project_ids or None\n else:\n self.project_ids_to_index = None\n self.asana_team_id = (asana_team_id.strip() or None) if asana_team_id else None\n self.batch_size = batch_size\n self.continue_on_failure = continue_on_failure\n logger.info(\n f\"AsanaConnector initialized with workspace_id: {asana_workspace_id}\"\n )\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.api_token = credentials[\"asana_api_token_secret\"]\n self.asana_client = asana_api.AsanaAPI(\n api_token=self.api_token,\n workspace_gid=self.workspace_id,\n team_gid=self.asana_team_id,\n )\n logger.info(\"Asana credentials loaded and API client initialized\")\n return None\n\n def poll_source(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch | None, # noqa: ARG002\n ) -> GenerateDocumentsOutput:\n start_time = datetime.datetime.fromtimestamp(start).isoformat()\n logger.info(f\"Starting Asana poll from {start_time}\")\n asana = asana_api.AsanaAPI(\n api_token=self.api_token,\n workspace_gid=self.workspace_id,\n team_gid=self.asana_team_id,\n )\n docs_batch: list[Document | HierarchyNode] = []\n tasks = asana.get_tasks(self.project_ids_to_index, start_time)\n\n for task in tasks:\n doc = self._message_to_doc(task)\n docs_batch.append(doc)\n\n if len(docs_batch) >= self.batch_size:\n logger.info(f\"Yielding batch of {len(docs_batch)} documents\")\n yield docs_batch\n docs_batch = []\n\n if docs_batch:\n logger.info(f\"Yielding final batch of {len(docs_batch)} documents\")\n yield docs_batch\n\n logger.info(\"Asana poll completed\")\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n logger.notice(\"Starting full index of all Asana tasks\")\n return self.poll_source(start=0, end=None)\n\n def _message_to_doc(self, task: asana_api.AsanaTask) -> Document:\n logger.debug(f\"Converting Asana task {task.id} to Document\")\n return Document(\n id=task.id,\n sections=[TextSection(link=task.link, text=task.text)],\n doc_updated_at=task.last_modified,\n source=DocumentSource.ASANA,\n semantic_identifier=task.title,\n metadata={\n \"group\": task.project_gid,\n \"project\": task.project_name,\n },\n )\n\n\nif __name__ == \"__main__\":\n import time\n import os\n\n logger.notice(\"Starting Asana connector test\")\n connector = AsanaConnector(\n os.environ[\"WORKSPACE_ID\"],\n os.environ[\"PROJECT_IDS\"],\n os.environ[\"TEAM_ID\"],\n )\n connector.load_credentials(\n {\n \"asana_api_token_secret\": os.environ[\"API_TOKEN\"],\n }\n )\n logger.info(\"Loading all documents from Asana\")\n all_docs = connector.load_from_state()\n current = time.time()\n one_day_ago = current - 24 * 60 * 60 # 1 day\n logger.info(\"Polling for documents updated in the last 24 hours\")\n latest_docs = connector.poll_source(one_day_ago, current)\n for docs in latest_docs:\n for doc in docs:\n if isinstance(doc, HierarchyNode):\n print(\"hierarchynode:\", doc.display_name)\n else:\n print(doc.id)\n logger.notice(\"Asana connector test completed\")\n" }, { "path": "backend/onyx/connectors/axero/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/axero/connector.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nimport requests\nfrom pydantic import BaseModel\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n process_in_batches,\n)\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\n\nlogger = setup_logger()\n\n\nENTITY_NAME_MAP = {1: \"Forum\", 3: \"Article\", 4: \"Blog\", 9: \"Wiki\"}\n\n\ndef _get_auth_header(api_key: str) -> dict[str, str]:\n return {\"Rest-Api-Key\": api_key}\n\n\n@retry_builder()\n@rate_limit_builder(max_calls=5, period=1)\ndef _rate_limited_request(\n endpoint: str, headers: dict, params: dict | None = None\n) -> Any:\n # https://my.axerosolutions.com/spaces/5/communifire-documentation/wiki/view/370/rest-api\n return requests.get(endpoint, headers=headers, params=params)\n\n\n# https://my.axerosolutions.com/spaces/5/communifire-documentation/wiki/view/595/rest-api-get-content-list\ndef _get_entities(\n entity_type: int,\n api_key: str,\n axero_base_url: str,\n start: datetime,\n end: datetime,\n space_id: str | None = None,\n) -> list[dict]:\n endpoint = axero_base_url + \"api/content/list\"\n page_num = 1\n pages_fetched = 0\n pages_to_return = []\n break_out = False\n while True:\n params = {\n \"EntityType\": str(entity_type),\n \"SortColumn\": \"DateUpdated\",\n \"SortOrder\": \"1\", # descending\n \"StartPage\": str(page_num),\n }\n\n if space_id is not None:\n params[\"SpaceID\"] = space_id\n\n res = _rate_limited_request(\n endpoint, headers=_get_auth_header(api_key), params=params\n )\n res.raise_for_status()\n\n # Axero limitations:\n # No next page token, can paginate but things may have changed\n # for example, a doc that hasn't been read in by Onyx is updated and is now front of the list\n # due to this limitation and the fact that Axero has no rate limiting but API calls can cause\n # increased latency for the team, we have to just fetch all the pages quickly to reduce the\n # chance of missing a document due to an update (it will still get updated next pass)\n # Assumes the volume of data isn't too big to store in memory (probably fine)\n data = res.json()\n total_records = data[\"TotalRecords\"]\n contents = data[\"ResponseData\"]\n pages_fetched += len(contents)\n logger.debug(f\"Fetched {pages_fetched} {ENTITY_NAME_MAP[entity_type]}\")\n\n for page in contents:\n update_time = time_str_to_utc(page[\"DateUpdated\"])\n\n if update_time > end:\n continue\n\n if update_time < start:\n break_out = True\n break\n\n pages_to_return.append(page)\n\n if pages_fetched >= total_records:\n break\n\n page_num += 1\n\n if break_out:\n break\n\n return pages_to_return\n\n\ndef _get_obj_by_id(obj_id: int, api_key: str, axero_base_url: str) -> dict:\n endpoint = axero_base_url + f\"api/content/{obj_id}\"\n res = _rate_limited_request(endpoint, headers=_get_auth_header(api_key))\n res.raise_for_status()\n\n return res.json()\n\n\nclass AxeroForum(BaseModel):\n doc_id: str\n title: str\n link: str\n initial_content: str\n responses: list[str]\n last_update: datetime\n\n\ndef _map_post_to_parent(\n posts: dict,\n api_key: str,\n axero_base_url: str,\n) -> list[AxeroForum]:\n \"\"\"Cannot handle in batches since the posts aren't ordered or structured in any way\n may need to map any number of them to the initial post\"\"\"\n epoch_str = \"1970-01-01T00:00:00.000\"\n post_map: dict[int, AxeroForum] = {}\n\n for ind, post in enumerate(posts):\n if (ind + 1) % 25 == 0:\n logger.debug(f\"Processed {ind + 1} posts or responses\")\n\n post_time = time_str_to_utc(\n post.get(\"DateUpdated\") or post.get(\"DateCreated\") or epoch_str\n )\n p_id = post.get(\"ParentContentID\")\n if p_id in post_map:\n axero_forum = post_map[p_id]\n axero_forum.responses.insert(0, post.get(\"ContentSummary\"))\n axero_forum.last_update = max(axero_forum.last_update, post_time)\n else:\n initial_post_d = _get_obj_by_id(p_id, api_key, axero_base_url)[\n \"ResponseData\"\n ]\n initial_post_time = time_str_to_utc(\n initial_post_d.get(\"DateUpdated\")\n or initial_post_d.get(\"DateCreated\")\n or epoch_str\n )\n post_map[p_id] = AxeroForum(\n doc_id=\"AXERO_\" + str(initial_post_d.get(\"ContentID\")),\n title=initial_post_d.get(\"ContentTitle\"),\n link=initial_post_d.get(\"ContentURL\"),\n initial_content=initial_post_d.get(\"ContentSummary\"),\n responses=[post.get(\"ContentSummary\")],\n last_update=max(post_time, initial_post_time),\n )\n\n return list(post_map.values())\n\n\ndef _get_forums(\n api_key: str,\n axero_base_url: str,\n space_id: str | None = None,\n) -> list[dict]:\n endpoint = axero_base_url + \"api/content/list\"\n page_num = 1\n pages_fetched = 0\n pages_to_return = []\n break_out = False\n\n while True:\n params = {\n \"EntityType\": \"54\",\n \"SortColumn\": \"DateUpdated\",\n \"SortOrder\": \"1\", # descending\n \"StartPage\": str(page_num),\n }\n\n if space_id is not None:\n params[\"SpaceID\"] = space_id\n\n res = _rate_limited_request(\n endpoint, headers=_get_auth_header(api_key), params=params\n )\n res.raise_for_status()\n\n data = res.json()\n total_records = data[\"TotalRecords\"]\n contents = data[\"ResponseData\"]\n pages_fetched += len(contents)\n logger.debug(f\"Fetched {pages_fetched} forums\")\n\n for page in contents:\n pages_to_return.append(page)\n\n if pages_fetched >= total_records:\n break\n\n page_num += 1\n\n if break_out:\n break\n\n return pages_to_return\n\n\ndef _translate_forum_to_doc(af: AxeroForum) -> Document:\n doc = Document(\n id=af.doc_id,\n sections=[TextSection(link=af.link, text=reply) for reply in af.responses],\n source=DocumentSource.AXERO,\n semantic_identifier=af.title,\n doc_updated_at=af.last_update,\n metadata={},\n )\n\n return doc\n\n\ndef _translate_content_to_doc(content: dict) -> Document:\n page_text = \"\"\n summary = content.get(\"ContentSummary\")\n body = content.get(\"ContentBody\")\n if summary:\n page_text += f\"{summary}\\n\"\n\n if body:\n content_parsed = parse_html_page_basic(body)\n page_text += content_parsed\n\n doc = Document(\n id=\"AXERO_\" + str(content[\"ContentID\"]),\n sections=[TextSection(link=content[\"ContentURL\"], text=page_text)],\n source=DocumentSource.AXERO,\n semantic_identifier=content[\"ContentTitle\"],\n doc_updated_at=time_str_to_utc(content[\"DateUpdated\"]),\n metadata={\"space\": content[\"SpaceName\"]},\n )\n\n return doc\n\n\nclass AxeroConnector(PollConnector):\n def __init__(\n self,\n # Strings of the integer ids of the spaces\n spaces: list[str] | None = None,\n include_article: bool = True,\n include_blog: bool = True,\n include_wiki: bool = True,\n include_forum: bool = True,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.include_article = include_article\n self.include_blog = include_blog\n self.include_wiki = include_wiki\n self.include_forum = include_forum\n self.batch_size = batch_size\n self.space_ids = spaces\n self.axero_key = None\n self.base_url = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.axero_key = credentials[\"axero_api_token\"]\n # As the API key specifically applies to a particular deployment, this is\n # included as part of the credential\n base_url = credentials[\"base_url\"]\n if not base_url.endswith(\"/\"):\n base_url += \"/\"\n self.base_url = base_url\n return None\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n if not self.axero_key or not self.base_url:\n raise ConnectorMissingCredentialError(\"Axero\")\n\n start_datetime = datetime.utcfromtimestamp(start).replace(tzinfo=timezone.utc)\n end_datetime = datetime.utcfromtimestamp(end).replace(tzinfo=timezone.utc)\n\n entity_types = []\n if self.include_article:\n entity_types.append(3)\n if self.include_blog:\n entity_types.append(4)\n if self.include_wiki:\n entity_types.append(9)\n\n iterable_space_ids = self.space_ids if self.space_ids else [None]\n\n for space_id in iterable_space_ids:\n for entity in entity_types:\n axero_obj = _get_entities(\n entity_type=entity,\n api_key=self.axero_key,\n axero_base_url=self.base_url,\n start=start_datetime,\n end=end_datetime,\n space_id=space_id,\n )\n yield from process_in_batches(\n objects=axero_obj,\n process_function=_translate_content_to_doc,\n batch_size=self.batch_size,\n )\n\n if self.include_forum:\n forums_posts = _get_forums(\n api_key=self.axero_key,\n axero_base_url=self.base_url,\n space_id=space_id,\n )\n\n all_axero_forums = _map_post_to_parent(\n posts=forums_posts,\n api_key=self.axero_key,\n axero_base_url=self.base_url,\n )\n\n filtered_forums = [\n f\n for f in all_axero_forums\n if f.last_update >= start_datetime and f.last_update <= end_datetime\n ]\n\n yield from process_in_batches(\n objects=filtered_forums,\n process_function=_translate_forum_to_doc,\n batch_size=self.batch_size,\n )\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = AxeroConnector()\n connector.load_credentials(\n {\n \"axero_api_token\": os.environ[\"AXERO_API_TOKEN\"],\n \"base_url\": os.environ[\"AXERO_BASE_URL\"],\n }\n )\n current = time.time()\n\n one_year_ago = current - 24 * 60 * 60 * 360\n latest_docs = connector.poll_source(one_year_ago, current)\n\n print(next(latest_docs))\n" }, { "path": "backend/onyx/connectors/bitbucket/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/bitbucket/connector.py", "content": "from __future__ import annotations\n\nimport copy\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import TYPE_CHECKING\n\nfrom typing_extensions import override\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.bitbucket.utils import build_auth_client\nfrom onyx.connectors.bitbucket.utils import list_repositories\nfrom onyx.connectors.bitbucket.utils import map_pr_to_document\nfrom onyx.connectors.bitbucket.utils import paginate\nfrom onyx.connectors.bitbucket.utils import PR_LIST_RESPONSE_FIELDS\nfrom onyx.connectors.bitbucket.utils import SLIM_PR_LIST_RESPONSE_FIELDS\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nif TYPE_CHECKING:\n import httpx\n\nlogger = setup_logger()\n\n\nclass BitbucketConnectorCheckpoint(ConnectorCheckpoint):\n \"\"\"Checkpoint state for resumable Bitbucket PR indexing.\n\n Fields:\n repos_queue: Materialized list of repository slugs to process.\n current_repo_index: Index of the repository currently being processed.\n next_url: Bitbucket \"next\" URL for continuing pagination within the current repo.\n \"\"\"\n\n repos_queue: list[str] = []\n current_repo_index: int = 0\n next_url: str | None = None\n\n\nclass BitbucketConnector(\n CheckpointedConnector[BitbucketConnectorCheckpoint],\n SlimConnectorWithPermSync,\n):\n \"\"\"Connector for indexing Bitbucket Cloud pull requests.\n\n Args:\n workspace: Bitbucket workspace ID.\n repositories: Comma-separated list of repository slugs to index.\n projects: Comma-separated list of project keys to index all repositories within.\n batch_size: Max number of documents to yield per batch.\n \"\"\"\n\n def __init__(\n self,\n workspace: str,\n repositories: str | None = None,\n projects: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.workspace = workspace\n self._repositories = (\n [s.strip() for s in repositories.split(\",\") if s.strip()]\n if repositories\n else None\n )\n self._projects: list[str] | None = (\n [s.strip() for s in projects.split(\",\") if s.strip()] if projects else None\n )\n self.batch_size = batch_size\n self.email: str | None = None\n self.api_token: str | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"Load API token-based credentials.\n\n Expects a dict with keys: `bitbucket_email`, `bitbucket_api_token`.\n \"\"\"\n self.email = credentials.get(\"bitbucket_email\")\n self.api_token = credentials.get(\"bitbucket_api_token\")\n if not self.email or not self.api_token:\n raise ConnectorMissingCredentialError(\"Bitbucket\")\n return None\n\n def _client(self) -> httpx.Client:\n \"\"\"Build an authenticated HTTP client or raise if credentials missing.\"\"\"\n if not self.email or not self.api_token:\n raise ConnectorMissingCredentialError(\"Bitbucket\")\n return build_auth_client(self.email, self.api_token)\n\n def _iter_pull_requests_for_repo(\n self,\n client: httpx.Client,\n repo_slug: str,\n params: dict[str, Any] | None = None,\n start_url: str | None = None,\n on_page: Callable[[str | None], None] | None = None,\n ) -> Iterator[dict[str, Any]]:\n base = f\"https://api.bitbucket.org/2.0/repositories/{self.workspace}/{repo_slug}/pullrequests\"\n yield from paginate(\n client,\n base,\n params,\n start_url=start_url,\n on_page=on_page,\n )\n\n def _build_params(\n self,\n fields: str = PR_LIST_RESPONSE_FIELDS,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> dict[str, Any]:\n \"\"\"Build Bitbucket fetch params.\n\n Always include OPEN, MERGED, and DECLINED PRs. If both ``start`` and\n ``end`` are provided, apply a single updated_on time window.\n \"\"\"\n\n def _iso(ts: SecondsSinceUnixEpoch) -> str:\n return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat()\n\n def _tc_epoch(\n lower_epoch: SecondsSinceUnixEpoch | None,\n upper_epoch: SecondsSinceUnixEpoch | None,\n ) -> str | None:\n if lower_epoch is not None and upper_epoch is not None:\n lower_iso = _iso(lower_epoch)\n upper_iso = _iso(upper_epoch)\n return f'(updated_on >= \"{lower_iso}\" AND updated_on <= \"{upper_iso}\")'\n return None\n\n params: dict[str, Any] = {\"fields\": fields, \"pagelen\": 50}\n time_clause = _tc_epoch(start, end)\n q = '(state = \"OPEN\" OR state = \"MERGED\" OR state = \"DECLINED\")'\n if time_clause:\n q = f\"{q} AND {time_clause}\"\n params[\"q\"] = q\n return params\n\n def _iter_target_repositories(self, client: httpx.Client) -> Iterator[str]:\n \"\"\"Yield repository slugs based on configuration.\n\n Priority:\n - repositories list\n - projects list (list repos by project key)\n - workspace (all repos)\n \"\"\"\n if self._repositories:\n for slug in self._repositories:\n yield slug\n return\n if self._projects:\n for project_key in self._projects:\n for repo in list_repositories(client, self.workspace, project_key):\n slug_val = repo.get(\"slug\")\n if isinstance(slug_val, str) and slug_val:\n yield slug_val\n return\n for repo in list_repositories(client, self.workspace, None):\n slug_val = repo.get(\"slug\")\n if isinstance(slug_val, str) and slug_val:\n yield slug_val\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: BitbucketConnectorCheckpoint,\n ) -> CheckpointOutput[BitbucketConnectorCheckpoint]:\n \"\"\"Resumable PR ingestion across repos and pages within a time window.\n\n Yields Documents (or ConnectorFailure for per-PR mapping failures) and returns\n an updated checkpoint that records repo position and next page URL.\n \"\"\"\n new_checkpoint = copy.deepcopy(checkpoint)\n\n with self._client() as client:\n # Materialize target repositories once\n if not new_checkpoint.repos_queue:\n # Preserve explicit order; otherwise ensure deterministic ordering\n repos_list = list(self._iter_target_repositories(client))\n new_checkpoint.repos_queue = sorted(set(repos_list))\n new_checkpoint.current_repo_index = 0\n new_checkpoint.next_url = None\n\n repos = new_checkpoint.repos_queue\n if not repos or new_checkpoint.current_repo_index >= len(repos):\n new_checkpoint.has_more = False\n return new_checkpoint\n\n repo_slug = repos[new_checkpoint.current_repo_index]\n\n first_page_params = self._build_params(\n fields=PR_LIST_RESPONSE_FIELDS,\n start=start,\n end=end,\n )\n\n def _on_page(next_url: str | None) -> None:\n new_checkpoint.next_url = next_url\n\n for pr in self._iter_pull_requests_for_repo(\n client,\n repo_slug,\n params=first_page_params,\n start_url=new_checkpoint.next_url,\n on_page=_on_page,\n ):\n try:\n document = map_pr_to_document(pr, self.workspace, repo_slug)\n yield document\n except Exception as e:\n pr_id = pr.get(\"id\")\n pr_link = (\n f\"https://bitbucket.org/{self.workspace}/{repo_slug}/pull-requests/{pr_id}\"\n if pr_id is not None\n else None\n )\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=(\n f\"{DocumentSource.BITBUCKET.value}:{self.workspace}:{repo_slug}:pr:{pr_id}\"\n if pr_id is not None\n else f\"{DocumentSource.BITBUCKET.value}:{self.workspace}:{repo_slug}:pr:unknown\"\n ),\n document_link=pr_link,\n ),\n failure_message=f\"Failed to process Bitbucket PR: {e}\",\n exception=e,\n )\n\n # Advance to next repository (if any) and set has_more accordingly\n new_checkpoint.current_repo_index += 1\n new_checkpoint.next_url = None\n new_checkpoint.has_more = new_checkpoint.current_repo_index < len(repos)\n\n return new_checkpoint\n\n @override\n def build_dummy_checkpoint(self) -> BitbucketConnectorCheckpoint:\n \"\"\"Create an initial checkpoint with work remaining.\"\"\"\n return BitbucketConnectorCheckpoint(has_more=True)\n\n @override\n def validate_checkpoint_json(\n self, checkpoint_json: str\n ) -> BitbucketConnectorCheckpoint:\n \"\"\"Validate and deserialize a checkpoint instance from JSON.\"\"\"\n return BitbucketConnectorCheckpoint.model_validate_json(checkpoint_json)\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> Iterator[list[SlimDocument | HierarchyNode]]:\n \"\"\"Return only document IDs for all existing pull requests.\"\"\"\n batch: list[SlimDocument | HierarchyNode] = []\n params = self._build_params(\n fields=SLIM_PR_LIST_RESPONSE_FIELDS,\n start=start,\n end=end,\n )\n with self._client() as client:\n for slug in self._iter_target_repositories(client):\n for pr in self._iter_pull_requests_for_repo(\n client, slug, params=params\n ):\n pr_id = pr[\"id\"]\n doc_id = f\"{DocumentSource.BITBUCKET.value}:{self.workspace}:{slug}:pr:{pr_id}\"\n batch.append(SlimDocument(id=doc_id))\n if len(batch) >= self.batch_size:\n yield batch\n batch = []\n if callback:\n if callback.should_stop():\n # Note: this is not actually used for permission sync yet, just pruning\n raise RuntimeError(\n \"bitbucket_pr_sync: Stop signal detected\"\n )\n callback.progress(\"bitbucket_pr_sync\", len(batch))\n if batch:\n yield batch\n\n def validate_connector_settings(self) -> None:\n \"\"\"Validate Bitbucket credentials and workspace access by probing a lightweight endpoint.\n\n Raises:\n CredentialExpiredError: on HTTP 401\n InsufficientPermissionsError: on HTTP 403\n UnexpectedValidationError: on any other failure\n \"\"\"\n try:\n with self._client() as client:\n url = f\"https://api.bitbucket.org/2.0/repositories/{self.workspace}\"\n resp = client.get(\n url,\n params={\"pagelen\": 1, \"fields\": \"pagelen\"},\n timeout=REQUEST_TIMEOUT_SECONDS,\n )\n if resp.status_code == 401:\n raise CredentialExpiredError(\n \"Invalid or expired Bitbucket credentials (HTTP 401).\"\n )\n if resp.status_code == 403:\n raise InsufficientPermissionsError(\n \"Insufficient permissions to access Bitbucket workspace (HTTP 403).\"\n )\n if resp.status_code < 200 or resp.status_code >= 300:\n raise UnexpectedValidationError(\n f\"Unexpected Bitbucket error (status={resp.status_code}).\"\n )\n except Exception as e:\n # Network or other unexpected errors\n if isinstance(\n e,\n (\n CredentialExpiredError,\n InsufficientPermissionsError,\n UnexpectedValidationError,\n ConnectorMissingCredentialError,\n ),\n ):\n raise\n raise UnexpectedValidationError(\n f\"Unexpected error while validating Bitbucket settings: {e}\"\n )\n" }, { "path": "backend/onyx/connectors/bitbucket/utils.py", "content": "from __future__ import annotations\n\nimport time\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nimport httpx\n\nfrom onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\nlogger = setup_logger()\n\n# Fields requested from Bitbucket PR list endpoint to ensure rich PR data\nPR_LIST_RESPONSE_FIELDS: str = \",\".join(\n [\n \"next\",\n \"page\",\n \"pagelen\",\n \"values.author\",\n \"values.close_source_branch\",\n \"values.closed_by\",\n \"values.comment_count\",\n \"values.created_on\",\n \"values.description\",\n \"values.destination\",\n \"values.draft\",\n \"values.id\",\n \"values.links\",\n \"values.merge_commit\",\n \"values.participants\",\n \"values.reason\",\n \"values.rendered\",\n \"values.reviewers\",\n \"values.source\",\n \"values.state\",\n \"values.summary\",\n \"values.task_count\",\n \"values.title\",\n \"values.type\",\n \"values.updated_on\",\n ]\n)\n\n# Minimal fields for slim retrieval (IDs only)\nSLIM_PR_LIST_RESPONSE_FIELDS: str = \",\".join(\n [\n \"next\",\n \"page\",\n \"pagelen\",\n \"values.id\",\n ]\n)\n\n\n# Minimal fields for repository list calls\nREPO_LIST_RESPONSE_FIELDS: str = \",\".join(\n [\n \"next\",\n \"page\",\n \"pagelen\",\n \"values.slug\",\n \"values.full_name\",\n \"values.project.key\",\n ]\n)\n\n\nclass BitbucketRetriableError(Exception):\n \"\"\"Raised for retriable Bitbucket conditions (429, 5xx).\"\"\"\n\n\nclass BitbucketNonRetriableError(Exception):\n \"\"\"Raised for non-retriable Bitbucket client errors (4xx except 429).\"\"\"\n\n\n@retry_builder(\n tries=6,\n delay=1,\n backoff=2,\n max_delay=30,\n exceptions=(BitbucketRetriableError, httpx.RequestError),\n)\n@rate_limit_builder(max_calls=60, period=60)\ndef bitbucket_get(\n client: httpx.Client, url: str, params: dict[str, Any] | None = None\n) -> httpx.Response:\n \"\"\"Perform a GET against Bitbucket with retry and rate limiting.\n\n Retries on 429 and 5xx responses, and on transport errors. Honors\n `Retry-After` header for 429 when present by sleeping before retrying.\n \"\"\"\n try:\n response = client.get(url, params=params, timeout=REQUEST_TIMEOUT_SECONDS)\n except httpx.RequestError:\n # Allow retry_builder to handle retries of transport errors\n raise\n\n try:\n response.raise_for_status()\n except httpx.HTTPStatusError as e:\n status = e.response.status_code if e.response is not None else None\n if status == 429:\n retry_after = e.response.headers.get(\"Retry-After\") if e.response else None\n if retry_after is not None:\n try:\n time.sleep(int(retry_after))\n except (TypeError, ValueError):\n pass\n raise BitbucketRetriableError(\"Bitbucket rate limit exceeded (429)\") from e\n if status is not None and 500 <= status < 600:\n raise BitbucketRetriableError(f\"Bitbucket server error: {status}\") from e\n if status is not None and 400 <= status < 500:\n raise BitbucketNonRetriableError(f\"Bitbucket client error: {status}\") from e\n # Unknown status, propagate\n raise\n\n return response\n\n\ndef build_auth_client(email: str, api_token: str) -> httpx.Client:\n \"\"\"Create an authenticated httpx client for Bitbucket Cloud API.\"\"\"\n return httpx.Client(auth=(email, api_token), http2=True)\n\n\ndef paginate(\n client: httpx.Client,\n url: str,\n params: dict[str, Any] | None = None,\n start_url: str | None = None,\n on_page: Callable[[str | None], None] | None = None,\n) -> Iterator[dict[str, Any]]:\n \"\"\"Iterate over paginated Bitbucket API responses yielding individual values.\n\n Args:\n client: Authenticated HTTP client.\n url: Base collection URL (first page when start_url is None).\n params: Query params for the first page.\n start_url: If provided, start from this absolute URL (ignores params).\n on_page: Optional callback invoked after each page with the next page URL.\n \"\"\"\n next_url = start_url or url\n # If resuming from a next URL, do not pass params again\n query = params.copy() if params else None\n query = None if start_url else query\n while next_url:\n resp = bitbucket_get(client, next_url, params=query)\n data = resp.json()\n values = data.get(\"values\", [])\n for item in values:\n yield item\n next_url = data.get(\"next\")\n if on_page is not None:\n on_page(next_url)\n # only include params on first call, next_url will contain all necessary params\n query = None\n\n\ndef list_repositories(\n client: httpx.Client, workspace: str, project_key: str | None = None\n) -> Iterator[dict[str, Any]]:\n \"\"\"List repositories in a workspace, optionally filtered by project key.\"\"\"\n base_url = f\"https://api.bitbucket.org/2.0/repositories/{workspace}\"\n params: dict[str, Any] = {\n \"fields\": REPO_LIST_RESPONSE_FIELDS,\n \"pagelen\": 100,\n # Ensure deterministic ordering\n \"sort\": \"full_name\",\n }\n if project_key:\n params[\"q\"] = f'project.key=\"{project_key}\"'\n yield from paginate(client, base_url, params)\n\n\ndef map_pr_to_document(pr: dict[str, Any], workspace: str, repo_slug: str) -> Document:\n \"\"\"Map a Bitbucket pull request JSON to Onyx Document.\"\"\"\n pr_id = pr[\"id\"]\n title = pr.get(\"title\") or f\"PR {pr_id}\"\n description = pr.get(\"description\") or \"\"\n state = pr.get(\"state\")\n draft = pr.get(\"draft\", False)\n author = pr.get(\"author\", {})\n reviewers = pr.get(\"reviewers\", [])\n participants = pr.get(\"participants\", [])\n\n link = pr.get(\"links\", {}).get(\"html\", {}).get(\"href\") or (\n f\"https://bitbucket.org/{workspace}/{repo_slug}/pull-requests/{pr_id}\"\n )\n\n created_on = pr.get(\"created_on\")\n updated_on = pr.get(\"updated_on\")\n updated_dt = (\n datetime.fromisoformat(updated_on.replace(\"Z\", \"+00:00\")).astimezone(\n timezone.utc\n )\n if isinstance(updated_on, str)\n else None\n )\n\n source_branch = pr.get(\"source\", {}).get(\"branch\", {}).get(\"name\", \"\")\n destination_branch = pr.get(\"destination\", {}).get(\"branch\", {}).get(\"name\", \"\")\n\n approved_by = [\n _get_user_name(p.get(\"user\", {})) for p in participants if p.get(\"approved\")\n ]\n\n primary_owner = None\n if author:\n primary_owner = BasicExpertInfo(\n display_name=_get_user_name(author),\n )\n\n secondary_owners = [\n BasicExpertInfo(display_name=_get_user_name(r)) for r in reviewers\n ] or None\n\n reviewer_names = [_get_user_name(r) for r in reviewers]\n\n # Create a concise summary of key PR info\n created_date = created_on.split(\"T\")[0] if created_on else \"N/A\"\n updated_date = updated_on.split(\"T\")[0] if updated_on else \"N/A\"\n content_text = (\n \"Pull Request Information:\\n\"\n f\"- Pull Request ID: {pr_id}\\n\"\n f\"- Title: {title}\\n\"\n f\"- State: {state or 'N/A'} {'(Draft)' if draft else ''}\\n\"\n )\n if state == \"DECLINED\":\n content_text += f\"- Reason: {pr.get('reason', 'N/A')}\\n\"\n content_text += (\n f\"- Author: {_get_user_name(author) if author else 'N/A'}\\n\"\n f\"- Reviewers: {', '.join(reviewer_names) if reviewer_names else 'N/A'}\\n\"\n f\"- Branch: {source_branch} -> {destination_branch}\\n\"\n f\"- Created: {created_date}\\n\"\n f\"- Updated: {updated_date}\"\n )\n if description:\n content_text += f\"\\n\\nDescription:\\n{description}\"\n sections: list[TextSection | ImageSection] = [\n TextSection(link=link, text=content_text)\n ]\n\n metadata: dict[str, str | list[str]] = {\n \"object_type\": \"PullRequest\",\n \"workspace\": workspace,\n \"repository\": repo_slug,\n \"pr_key\": f\"{workspace}/{repo_slug}#{pr_id}\",\n \"id\": str(pr_id),\n \"title\": title,\n \"state\": state or \"\",\n \"draft\": str(bool(draft)),\n \"link\": link,\n \"author\": _get_user_name(author) if author else \"\",\n \"reviewers\": reviewer_names,\n \"approved_by\": approved_by,\n \"comment_count\": str(pr.get(\"comment_count\", \"\")),\n \"task_count\": str(pr.get(\"task_count\", \"\")),\n \"created_on\": created_on or \"\",\n \"updated_on\": updated_on or \"\",\n \"source_branch\": source_branch,\n \"destination_branch\": destination_branch,\n \"closed_by\": (\n _get_user_name(pr.get(\"closed_by\", {})) if pr.get(\"closed_by\") else \"\"\n ),\n \"close_source_branch\": str(bool(pr.get(\"close_source_branch\", False))),\n }\n\n return Document(\n id=f\"{DocumentSource.BITBUCKET.value}:{workspace}:{repo_slug}:pr:{pr_id}\",\n sections=sections,\n source=DocumentSource.BITBUCKET,\n semantic_identifier=f\"#{pr_id}: {title}\",\n title=title,\n doc_updated_at=updated_dt,\n primary_owners=[primary_owner] if primary_owner else None,\n secondary_owners=secondary_owners,\n metadata=metadata,\n )\n\n\ndef _get_user_name(user: dict[str, Any]) -> str:\n return user.get(\"display_name\") or user.get(\"nickname\") or \"unknown\"\n" }, { "path": "backend/onyx/connectors/blob/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/blob/connector.py", "content": "import os\nimport time\nfrom collections.abc import Mapping\nfrom datetime import datetime\nfrom datetime import timezone\nfrom io import BytesIO\nfrom numbers import Integral\nfrom typing import Any\nfrom typing import Optional\nfrom urllib.parse import quote\n\nimport boto3\nfrom botocore.client import Config\nfrom botocore.credentials import RefreshableCredentials\nfrom botocore.exceptions import ClientError\nfrom botocore.exceptions import NoCredentialsError\nfrom botocore.exceptions import PartialCredentialsError\nfrom botocore.session import get_session\nfrom mypy_boto3_s3 import S3Client\n\nfrom onyx.configs.app_configs import BLOB_STORAGE_SIZE_THRESHOLD\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import BlobType\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n process_onyx_metadata,\n)\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_text_and_images\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.image_utils import store_image_and_create_section\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nDOWNLOAD_CHUNK_SIZE = 1024 * 1024\nSIZE_THRESHOLD_BUFFER = 64\n\n\nclass BlobStorageConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n bucket_type: str,\n bucket_name: str,\n prefix: str = \"\",\n batch_size: int = INDEX_BATCH_SIZE,\n european_residency: bool = False,\n ) -> None:\n self.bucket_type: BlobType = BlobType(bucket_type)\n self.bucket_name = bucket_name.strip()\n self.prefix = prefix if not prefix or prefix.endswith(\"/\") else prefix + \"/\"\n self.batch_size = batch_size\n self.s3_client: Optional[S3Client] = None\n self._allow_images: bool | None = None\n self.size_threshold: int | None = BLOB_STORAGE_SIZE_THRESHOLD\n self.bucket_region: Optional[str] = None\n self.european_residency: bool = european_residency\n\n def set_allow_images(self, allow_images: bool) -> None:\n \"\"\"Set whether to process images in this connector.\"\"\"\n logger.info(f\"Setting allow_images to {allow_images}.\")\n self._allow_images = allow_images\n\n def _detect_bucket_region(self) -> None:\n \"\"\"Detect and cache the actual region of the S3 bucket using head_bucket.\"\"\"\n if self.s3_client is None:\n logger.warning(\n \"S3 client not initialized. Skipping bucket region detection.\"\n )\n return\n\n try:\n response = self.s3_client.head_bucket(Bucket=self.bucket_name)\n # The region is in the response headers as 'x-amz-bucket-region'\n self.bucket_region = response.get(\"BucketRegion\") or response.get(\n \"ResponseMetadata\", {}\n ).get(\"HTTPHeaders\", {}).get(\"x-amz-bucket-region\")\n\n if self.bucket_region:\n logger.debug(f\"Detected bucket region: {self.bucket_region}\")\n else:\n logger.warning(\"Bucket region not found in head_bucket response\")\n except Exception as e:\n logger.warning(f\"Failed to detect bucket region via head_bucket: {e}\")\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"Checks for boto3 credentials based on the bucket type.\n (1) R2: Access Key ID, Secret Access Key, Account ID\n (2) S3: AWS Access Key ID, AWS Secret Access Key or IAM role or Assume Role\n (3) GOOGLE_CLOUD_STORAGE: Access Key ID, Secret Access Key, Project ID\n (4) OCI_STORAGE: Namespace, Region, Access Key ID, Secret Access Key\n\n For each bucket type, the method initializes the appropriate S3 client:\n - R2: Uses Cloudflare R2 endpoint with S3v4 signature\n - S3: Creates a standard boto3 S3 client\n - GOOGLE_CLOUD_STORAGE: Uses Google Cloud Storage endpoint\n - OCI_STORAGE: Uses Oracle Cloud Infrastructure Object Storage endpoint\n\n Raises ConnectorMissingCredentialError if required credentials are missing.\n Raises ValueError for unsupported bucket types.\n \"\"\"\n\n logger.debug(\n f\"Loading credentials for {self.bucket_name} or type {self.bucket_type}\"\n )\n\n if self.bucket_type == BlobType.R2:\n if not all(\n credentials.get(key)\n for key in [\"r2_access_key_id\", \"r2_secret_access_key\", \"account_id\"]\n ):\n raise ConnectorMissingCredentialError(\"Cloudflare R2\")\n\n # Use EU endpoint if european_residency is enabled\n subdomain = \"eu.\" if self.european_residency else \"\"\n endpoint_url = f\"https://{credentials['account_id']}.{subdomain}r2.cloudflarestorage.com\"\n\n self.s3_client = boto3.client(\n \"s3\",\n endpoint_url=endpoint_url,\n aws_access_key_id=credentials[\"r2_access_key_id\"],\n aws_secret_access_key=credentials[\"r2_secret_access_key\"],\n region_name=\"auto\",\n config=Config(signature_version=\"s3v4\"),\n )\n\n elif self.bucket_type == BlobType.S3:\n # For S3, we can use either access keys or IAM roles.\n authentication_method = credentials.get(\n \"authentication_method\", \"access_key\"\n )\n logger.debug(\n f\"Using authentication method: {authentication_method} for S3 bucket.\"\n )\n if authentication_method == \"access_key\":\n logger.debug(\"Using access key authentication for S3 bucket.\")\n if not all(\n credentials.get(key)\n for key in [\"aws_access_key_id\", \"aws_secret_access_key\"]\n ):\n raise ConnectorMissingCredentialError(\"Amazon S3\")\n\n session = boto3.Session(\n aws_access_key_id=credentials[\"aws_access_key_id\"],\n aws_secret_access_key=credentials[\"aws_secret_access_key\"],\n )\n self.s3_client = session.client(\"s3\")\n elif authentication_method == \"iam_role\":\n # If using IAM roles, we assume the role and let boto3 handle the credentials.\n role_arn = credentials.get(\"aws_role_arn\")\n # create session name using timestamp\n if not role_arn:\n raise ConnectorMissingCredentialError(\n \"Amazon S3 IAM role ARN is required for assuming role.\"\n )\n\n def _refresh_credentials() -> dict[str, str]:\n \"\"\"Refreshes the credentials for the assumed role.\"\"\"\n sts_client = boto3.client(\"sts\")\n assumed_role_object = sts_client.assume_role(\n RoleArn=role_arn,\n RoleSessionName=f\"onyx_blob_storage_{int(time.time())}\",\n )\n creds = assumed_role_object[\"Credentials\"]\n return {\n \"access_key\": creds[\"AccessKeyId\"],\n \"secret_key\": creds[\"SecretAccessKey\"],\n \"token\": creds[\"SessionToken\"],\n \"expiry_time\": creds[\"Expiration\"].isoformat(),\n }\n\n refreshable = RefreshableCredentials.create_from_metadata(\n metadata=_refresh_credentials(),\n refresh_using=_refresh_credentials,\n method=\"sts-assume-role\",\n )\n botocore_session = get_session()\n botocore_session._credentials = refreshable # type: ignore[attr-defined]\n session = boto3.Session(botocore_session=botocore_session)\n self.s3_client = session.client(\"s3\")\n elif authentication_method == \"assume_role\":\n # We will assume the instance role to access S3.\n logger.debug(\"Using instance role authentication for S3 bucket.\")\n self.s3_client = boto3.client(\"s3\")\n else:\n raise ConnectorValidationError(\"Invalid authentication method for S3. \")\n\n # This is important for correct citation links\n # NOTE: the client region actually doesn't matter for accessing the bucket\n self._detect_bucket_region()\n\n elif self.bucket_type == BlobType.GOOGLE_CLOUD_STORAGE:\n if not all(\n credentials.get(key) for key in [\"access_key_id\", \"secret_access_key\"]\n ):\n raise ConnectorMissingCredentialError(\"Google Cloud Storage\")\n\n self.s3_client = boto3.client(\n \"s3\",\n endpoint_url=\"https://storage.googleapis.com\",\n aws_access_key_id=credentials[\"access_key_id\"],\n aws_secret_access_key=credentials[\"secret_access_key\"],\n region_name=\"auto\",\n )\n\n elif self.bucket_type == BlobType.OCI_STORAGE:\n if not all(\n credentials.get(key)\n for key in [\"namespace\", \"region\", \"access_key_id\", \"secret_access_key\"]\n ):\n raise ConnectorMissingCredentialError(\"Oracle Cloud Infrastructure\")\n\n self.s3_client = boto3.client(\n \"s3\",\n endpoint_url=f\"https://{credentials['namespace']}.compat.objectstorage.{credentials['region']}.oraclecloud.com\",\n aws_access_key_id=credentials[\"access_key_id\"],\n aws_secret_access_key=credentials[\"secret_access_key\"],\n region_name=credentials[\"region\"],\n )\n\n else:\n raise ValueError(f\"Unsupported bucket type: {self.bucket_type}\")\n\n return None\n\n def _download_object(self, key: str) -> bytes | None:\n if self.s3_client is None:\n raise ConnectorMissingCredentialError(\"Blob storage\")\n response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)\n body = response[\"Body\"]\n\n try:\n if self.size_threshold is None:\n return body.read()\n\n return self._read_stream_with_limit(body, key)\n finally:\n body.close()\n\n def _read_stream_with_limit(self, body: Any, key: str) -> bytes | None:\n if self.size_threshold is None:\n return body.read()\n\n bytes_read = 0\n chunks: list[bytes] = []\n chunk_size = min(\n DOWNLOAD_CHUNK_SIZE, self.size_threshold + SIZE_THRESHOLD_BUFFER\n )\n\n for chunk in body.iter_chunks(chunk_size=chunk_size):\n if not chunk:\n continue\n chunks.append(chunk)\n bytes_read += len(chunk)\n\n if bytes_read > self.size_threshold + SIZE_THRESHOLD_BUFFER:\n logger.warning(\n f\"{key} exceeds size threshold of {self.size_threshold}. Skipping.\"\n )\n return None\n\n return b\"\".join(chunks)\n\n # NOTE: Left in as may be useful for one-off access to documents and sharing across orgs.\n # def _get_presigned_url(self, key: str) -> str:\n # if self.s3_client is None:\n # raise ConnectorMissingCredentialError(\"Blog storage\")\n\n # url = self.s3_client.generate_presigned_url(\n # \"get_object\",\n # Params={\"Bucket\": self.bucket_name, \"Key\": key},\n # ExpiresIn=self.presign_length,\n # )\n # return url\n\n def _get_blob_link(self, key: str) -> str:\n # NOTE: We store the object dashboard URL instead of the actual object URL\n # This is because the actual object URL requires S3 client authentication\n # Accessing through the browser will always return an unauthorized error\n\n if self.s3_client is None:\n raise ConnectorMissingCredentialError(\"Blob storage\")\n\n # URL encode the key to handle special characters, spaces, etc.\n # safe='/' keeps forward slashes unencoded for proper path structure\n encoded_key = quote(key, safe=\"/\")\n\n if self.bucket_type == BlobType.R2:\n account_id = self.s3_client.meta.endpoint_url.split(\"//\")[1].split(\".\")[0]\n subdomain = \"eu/\" if self.european_residency else \"default/\"\n\n return f\"https://dash.cloudflare.com/{account_id}/r2/{subdomain}buckets/{self.bucket_name}/objects/{encoded_key}/details\"\n\n elif self.bucket_type == BlobType.S3:\n region = self.bucket_region or self.s3_client.meta.region_name\n return f\"https://s3.console.aws.amazon.com/s3/object/{self.bucket_name}?region={region}&prefix={encoded_key}\"\n\n elif self.bucket_type == BlobType.GOOGLE_CLOUD_STORAGE:\n return f\"https://console.cloud.google.com/storage/browser/_details/{self.bucket_name}/{encoded_key}\"\n\n elif self.bucket_type == BlobType.OCI_STORAGE:\n namespace = self.s3_client.meta.endpoint_url.split(\"//\")[1].split(\".\")[0]\n region = self.s3_client.meta.region_name\n return f\"https://objectstorage.{region}.oraclecloud.com/n/{namespace}/b/{self.bucket_name}/o/{encoded_key}\"\n\n else:\n # This should never happen!\n raise ValueError(f\"Unsupported bucket type: {self.bucket_type}\")\n\n @staticmethod\n def _extract_size_bytes(obj: Mapping[str, Any]) -> int | None:\n \"\"\"Return the first numeric size field found on the object metadata.\"\"\"\n\n candidate_keys = (\n \"Size\",\n \"size\",\n \"ContentLength\",\n \"content_length\",\n \"Content-Length\",\n \"contentLength\",\n \"bytes\",\n \"Bytes\",\n )\n\n def _normalize(value: Any) -> int | None:\n if value is None or isinstance(value, bool):\n return None\n if isinstance(value, Integral):\n return int(value)\n try:\n numeric = float(value)\n except (TypeError, ValueError):\n return None\n if numeric >= 0 and numeric.is_integer():\n return int(numeric)\n return None\n\n for key in candidate_keys:\n if key in obj:\n normalized = _normalize(obj.get(key))\n if normalized is not None:\n return normalized\n\n for key, value in obj.items():\n if not isinstance(key, str):\n continue\n lowered_key = key.lower()\n if \"size\" in lowered_key or \"length\" in lowered_key:\n normalized = _normalize(value)\n if normalized is not None:\n return normalized\n\n return None\n\n def _yield_blob_objects(\n self,\n start: datetime,\n end: datetime,\n ) -> GenerateDocumentsOutput:\n if self.s3_client is None:\n raise ConnectorMissingCredentialError(\"Blob storage\")\n\n paginator = self.s3_client.get_paginator(\"list_objects_v2\")\n pages = paginator.paginate(Bucket=self.bucket_name, Prefix=self.prefix)\n\n batch: list[Document | HierarchyNode] = []\n for page in pages:\n if \"Contents\" not in page:\n continue\n\n for obj in page[\"Contents\"]:\n if obj[\"Key\"].endswith(\"/\"):\n continue\n\n last_modified = obj[\"LastModified\"].replace(tzinfo=timezone.utc)\n\n if not start <= last_modified <= end:\n continue\n\n file_name = os.path.basename(obj[\"Key\"])\n file_ext = get_file_ext(file_name)\n key = obj[\"Key\"]\n link = self._get_blob_link(key)\n\n size_bytes = self._extract_size_bytes(obj)\n if (\n self.size_threshold is not None\n and isinstance(size_bytes, int)\n and self.size_threshold is not None\n and size_bytes > self.size_threshold\n ):\n logger.warning(\n f\"{file_name} exceeds size threshold of {self.size_threshold}. Skipping.\"\n )\n continue\n\n # Handle image files\n if file_ext in OnyxFileExtensions.IMAGE_EXTENSIONS:\n if not self._allow_images:\n logger.debug(\n f\"Skipping image file: {key} (image processing not enabled)\"\n )\n continue\n\n # Process the image file\n try:\n downloaded_file = self._download_object(key)\n if downloaded_file is None:\n continue\n\n # TODO: Refactor to avoid direct DB access in connector\n # This will require broader refactoring across the codebase\n image_section, _ = store_image_and_create_section(\n image_data=downloaded_file,\n file_id=f\"{self.bucket_type}_{self.bucket_name}_{key.replace('/', '_')}\",\n display_name=file_name,\n link=link,\n file_origin=FileOrigin.CONNECTOR,\n )\n\n batch.append(\n Document(\n id=f\"{self.bucket_type}:{self.bucket_name}:{key}\",\n sections=[image_section],\n source=DocumentSource(self.bucket_type.value),\n semantic_identifier=file_name,\n doc_updated_at=last_modified,\n metadata={},\n )\n )\n\n if len(batch) == self.batch_size:\n yield batch\n batch = []\n except Exception:\n logger.exception(f\"Error processing image {key}\")\n continue\n\n # Handle text and document files\n try:\n downloaded_file = self._download_object(key)\n if downloaded_file is None:\n continue\n extraction_result = extract_text_and_images(\n BytesIO(downloaded_file), file_name=file_name\n )\n\n onyx_metadata, custom_tags = process_onyx_metadata(\n extraction_result.metadata\n )\n file_display_name = onyx_metadata.file_display_name or file_name\n time_updated = onyx_metadata.doc_updated_at or last_modified\n link = onyx_metadata.link or link\n primary_owners = onyx_metadata.primary_owners\n secondary_owners = onyx_metadata.secondary_owners\n source_type = onyx_metadata.source_type or DocumentSource(\n self.bucket_type.value\n )\n\n sections: list[TextSection | ImageSection] = []\n if extraction_result.text_content.strip():\n logger.debug(\n f\"Creating TextSection for {file_name} with link: {link}\"\n )\n sections.append(\n TextSection(\n link=link,\n text=extraction_result.text_content.strip(),\n )\n )\n\n batch.append(\n Document(\n id=f\"{self.bucket_type}:{self.bucket_name}:{key}\",\n sections=(\n sections\n if sections\n else [TextSection(link=link, text=\"\")]\n ),\n source=source_type,\n semantic_identifier=file_display_name,\n doc_updated_at=time_updated,\n metadata=custom_tags,\n primary_owners=primary_owners,\n secondary_owners=secondary_owners,\n )\n )\n if len(batch) == self.batch_size:\n yield batch\n batch = []\n\n except Exception:\n logger.exception(f\"Error decoding object {key} as UTF-8\")\n if batch:\n yield batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n logger.debug(\"Loading blob objects\")\n return self._yield_blob_objects(\n start=datetime(1970, 1, 1, tzinfo=timezone.utc),\n end=datetime.now(timezone.utc),\n )\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n if self.s3_client is None:\n raise ConnectorMissingCredentialError(\"Blob storage\")\n\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n\n for batch in self._yield_blob_objects(start_datetime, end_datetime):\n yield batch\n\n return None\n\n def validate_connector_settings(self) -> None:\n if self.s3_client is None:\n raise ConnectorMissingCredentialError(\n \"Blob storage credentials not loaded.\"\n )\n\n if not self.bucket_name:\n raise ConnectorValidationError(\n \"No bucket name was provided in connector settings.\"\n )\n\n try:\n # We only fetch one object/page as a light-weight validation step.\n # This ensures we trigger typical S3 permission checks (ListObjectsV2, etc.).\n self.s3_client.list_objects_v2(\n Bucket=self.bucket_name, Prefix=self.prefix, MaxKeys=1\n )\n\n except NoCredentialsError:\n raise ConnectorMissingCredentialError(\n \"No valid blob storage credentials found or provided to boto3.\"\n )\n except PartialCredentialsError:\n raise ConnectorMissingCredentialError(\n \"Partial or incomplete blob storage credentials provided to boto3.\"\n )\n except ClientError as e:\n error_code = e.response[\"Error\"].get(\"Code\", \"\")\n status_code = e.response[\"ResponseMetadata\"].get(\"HTTPStatusCode\")\n\n # Most common S3 error cases\n if error_code in [\n \"AccessDenied\",\n \"InvalidAccessKeyId\",\n \"SignatureDoesNotMatch\",\n ]:\n if status_code == 403 or error_code == \"AccessDenied\":\n raise InsufficientPermissionsError(\n f\"Insufficient permissions to list objects in bucket '{self.bucket_name}'. \"\n \"Please check your bucket policy and/or IAM policy.\"\n )\n if status_code == 401 or error_code == \"SignatureDoesNotMatch\":\n raise CredentialExpiredError(\n \"Provided blob storage credentials appear invalid or expired.\"\n )\n\n raise CredentialExpiredError(\n f\"Credential issue encountered ({error_code}).\"\n )\n\n if error_code == \"NoSuchBucket\" or status_code == 404:\n raise ConnectorValidationError(\n f\"Bucket '{self.bucket_name}' does not exist or cannot be found.\"\n )\n\n raise ConnectorValidationError(\n f\"Unexpected S3 client error (code={error_code}, status={status_code}): {e}\"\n )\n\n except Exception as e:\n # Catch-all for anything not captured by the above\n # Since we are unsure of the error and it may not disable the connector,\n # raise an unexpected error (does not disable connector)\n raise UnexpectedValidationError(\n f\"Unexpected error during blob storage settings validation: {e}\"\n )\n\n\nif __name__ == \"__main__\":\n credentials_dict = {\n \"aws_access_key_id\": os.environ.get(\"AWS_ACCESS_KEY_ID\"),\n \"aws_secret_access_key\": os.environ.get(\"AWS_SECRET_ACCESS_KEY\"),\n }\n\n # Initialize the connector\n connector = BlobStorageConnector(\n bucket_type=os.environ.get(\"BUCKET_TYPE\") or \"s3\",\n bucket_name=os.environ.get(\"BUCKET_NAME\") or \"test\",\n prefix=\"\",\n )\n\n try:\n connector.load_credentials(credentials_dict)\n document_batch_generator = connector.load_from_state()\n for document_batch in document_batch_generator:\n print(\"First batch of documents:\")\n for doc in document_batch:\n if isinstance(doc, HierarchyNode):\n print(\"hierarchynode:\", doc.display_name)\n continue\n\n print(f\"Document ID: {doc.id}\")\n print(f\"Semantic Identifier: {doc.semantic_identifier}\")\n print(f\"Source: {doc.source}\")\n print(f\"Updated At: {doc.doc_updated_at}\")\n print(\"Sections:\")\n for section in doc.sections:\n print(f\" - Link: {section.link}\")\n if isinstance(section, TextSection) and section.text is not None:\n print(f\" - Text: {section.text[:100]}...\")\n elif hasattr(section, \"image_file_id\") and section.image_file_id:\n print(f\" - Image: {section.image_file_id}\")\n else:\n print(\"Error: Unknown section type\")\n print(\"---\")\n break\n\n except ConnectorMissingCredentialError as e:\n print(f\"Error: {e}\")\n except Exception as e:\n print(f\"An unexpected error occurred: {e}\")\n" }, { "path": "backend/onyx/connectors/bookstack/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/bookstack/client.py", "content": "from typing import Any\n\nimport requests\n\n\nclass BookStackClientRequestFailedError(ConnectionError):\n def __init__(self, status: int, error: str) -> None:\n self.status_code = status\n self.error = error\n super().__init__(\n \"BookStack Client request failed with status {status}: {error}\".format(\n status=status, error=error\n )\n )\n\n\nclass BookStackApiClient:\n def __init__(\n self,\n base_url: str,\n token_id: str,\n token_secret: str,\n ) -> None:\n self.base_url = base_url\n self.token_id = token_id\n self.token_secret = token_secret\n\n def get(self, endpoint: str, params: dict[str, str]) -> dict[str, Any]:\n url: str = self._build_url(endpoint)\n headers = self._build_headers()\n response = requests.get(url, headers=headers, params=params)\n\n try:\n json = response.json()\n except Exception:\n json = {}\n\n if response.status_code >= 300:\n error = response.reason\n response_error = json.get(\"error\", {}).get(\"message\", \"\")\n if response_error:\n error = response_error\n raise BookStackClientRequestFailedError(response.status_code, error)\n\n return json\n\n def _build_headers(self) -> dict[str, str]:\n auth = \"Token \" + self.token_id + \":\" + self.token_secret\n return {\n \"Authorization\": auth,\n \"Accept\": \"application/json\",\n }\n\n def _build_url(self, endpoint: str) -> str:\n return self.base_url.rstrip(\"/\") + \"/api/\" + endpoint.lstrip(\"/\")\n\n def build_app_url(self, endpoint: str) -> str:\n return self.base_url.rstrip(\"/\") + \"/\" + endpoint.lstrip(\"/\")\n" }, { "path": "backend/onyx/connectors/bookstack/connector.py", "content": "import html\nimport time\nfrom collections.abc import Callable\nfrom datetime import datetime\nfrom typing import Any\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.bookstack.client import BookStackApiClient\nfrom onyx.connectors.bookstack.client import BookStackClientRequestFailedError\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\n\n\nclass BookstackConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.batch_size = batch_size\n self.bookstack_client: BookStackApiClient | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.bookstack_client = BookStackApiClient(\n base_url=credentials[\"bookstack_base_url\"],\n token_id=credentials[\"bookstack_api_token_id\"],\n token_secret=credentials[\"bookstack_api_token_secret\"],\n )\n return None\n\n @staticmethod\n def _get_doc_batch(\n batch_size: int,\n bookstack_client: BookStackApiClient,\n endpoint: str,\n transformer: Callable[[BookStackApiClient, dict], Document],\n start_ind: int,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> tuple[list[Document | HierarchyNode], int]:\n params = {\n \"count\": str(batch_size),\n \"offset\": str(start_ind),\n \"sort\": \"+id\",\n }\n\n if start:\n params[\"filter[updated_at:gte]\"] = datetime.utcfromtimestamp(\n start\n ).strftime(\"%Y-%m-%d\")\n\n if end:\n params[\"filter[updated_at:lte]\"] = datetime.utcfromtimestamp(end).strftime(\n \"%Y-%m-%d\"\n )\n\n batch = bookstack_client.get(endpoint, params=params).get(\"data\", [])\n doc_batch: list[Document | HierarchyNode] = [\n transformer(bookstack_client, item) for item in batch\n ]\n\n return doc_batch, len(batch)\n\n @staticmethod\n def _book_to_document(\n bookstack_client: BookStackApiClient, book: dict[str, Any]\n ) -> Document:\n url = bookstack_client.build_app_url(\"/books/\" + str(book.get(\"slug\")))\n title = str(book.get(\"name\", \"\"))\n text = book.get(\"name\", \"\") + \"\\n\" + book.get(\"description\", \"\")\n updated_at_str = (\n str(book.get(\"updated_at\")) if book.get(\"updated_at\") is not None else None\n )\n return Document(\n id=\"book__\" + str(book.get(\"id\")),\n sections=[TextSection(link=url, text=text)],\n source=DocumentSource.BOOKSTACK,\n semantic_identifier=\"Book: \" + title,\n title=title,\n doc_updated_at=(\n time_str_to_utc(updated_at_str) if updated_at_str is not None else None\n ),\n metadata={\"type\": \"book\"},\n )\n\n @staticmethod\n def _chapter_to_document(\n bookstack_client: BookStackApiClient, chapter: dict[str, Any]\n ) -> Document:\n url = bookstack_client.build_app_url(\n \"/books/\"\n + str(chapter.get(\"book_slug\"))\n + \"/chapter/\"\n + str(chapter.get(\"slug\"))\n )\n title = str(chapter.get(\"name\", \"\"))\n text = chapter.get(\"name\", \"\") + \"\\n\" + chapter.get(\"description\", \"\")\n updated_at_str = (\n str(chapter.get(\"updated_at\"))\n if chapter.get(\"updated_at\") is not None\n else None\n )\n return Document(\n id=\"chapter__\" + str(chapter.get(\"id\")),\n sections=[TextSection(link=url, text=text)],\n source=DocumentSource.BOOKSTACK,\n semantic_identifier=\"Chapter: \" + title,\n title=title,\n doc_updated_at=(\n time_str_to_utc(updated_at_str) if updated_at_str is not None else None\n ),\n metadata={\"type\": \"chapter\"},\n )\n\n @staticmethod\n def _shelf_to_document(\n bookstack_client: BookStackApiClient, shelf: dict[str, Any]\n ) -> Document:\n url = bookstack_client.build_app_url(\"/shelves/\" + str(shelf.get(\"slug\")))\n title = str(shelf.get(\"name\", \"\"))\n text = shelf.get(\"name\", \"\") + \"\\n\" + shelf.get(\"description\", \"\")\n updated_at_str = (\n str(shelf.get(\"updated_at\"))\n if shelf.get(\"updated_at\") is not None\n else None\n )\n return Document(\n id=\"shelf:\" + str(shelf.get(\"id\")),\n sections=[TextSection(link=url, text=text)],\n source=DocumentSource.BOOKSTACK,\n semantic_identifier=\"Shelf: \" + title,\n title=title,\n doc_updated_at=(\n time_str_to_utc(updated_at_str) if updated_at_str is not None else None\n ),\n metadata={\"type\": \"shelf\"},\n )\n\n @staticmethod\n def _page_to_document(\n bookstack_client: BookStackApiClient, page: dict[str, Any]\n ) -> Document:\n page_id = str(page.get(\"id\"))\n title = str(page.get(\"name\", \"\"))\n page_data = bookstack_client.get(\"/pages/\" + page_id, {})\n url = bookstack_client.build_app_url(\n \"/books/\"\n + str(page.get(\"book_slug\"))\n + \"/page/\"\n + str(page_data.get(\"slug\"))\n )\n page_html = \"

\" + html.escape(title) + \"

\" + str(page_data.get(\"html\"))\n text = parse_html_page_basic(page_html)\n updated_at_str = (\n str(page_data.get(\"updated_at\"))\n if page_data.get(\"updated_at\") is not None\n else None\n )\n time.sleep(0.1)\n return Document(\n id=\"page:\" + page_id,\n sections=[TextSection(link=url, text=text)],\n source=DocumentSource.BOOKSTACK,\n semantic_identifier=\"Page: \" + str(title),\n title=str(title),\n doc_updated_at=(\n time_str_to_utc(updated_at_str) if updated_at_str is not None else None\n ),\n metadata={\"type\": \"page\"},\n )\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n if self.bookstack_client is None:\n raise ConnectorMissingCredentialError(\"Bookstack\")\n\n return self.poll_source(None, None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch | None, end: SecondsSinceUnixEpoch | None\n ) -> GenerateDocumentsOutput:\n if self.bookstack_client is None:\n raise ConnectorMissingCredentialError(\"Bookstack\")\n\n transform_by_endpoint: dict[\n str, Callable[[BookStackApiClient, dict], Document]\n ] = {\n \"/books\": self._book_to_document,\n \"/chapters\": self._chapter_to_document,\n \"/shelves\": self._shelf_to_document,\n \"/pages\": self._page_to_document,\n }\n\n for endpoint, transform in transform_by_endpoint.items():\n start_ind = 0\n while True:\n doc_batch, num_results = self._get_doc_batch(\n batch_size=self.batch_size,\n bookstack_client=self.bookstack_client,\n endpoint=endpoint,\n transformer=transform,\n start_ind=start_ind,\n start=start,\n end=end,\n )\n start_ind += num_results\n if doc_batch:\n yield doc_batch\n\n if num_results < self.batch_size:\n break\n else:\n time.sleep(0.2)\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n Validate that the BookStack credentials and connector settings are correct.\n Specifically checks that we can make an authenticated request to BookStack.\n \"\"\"\n if not self.bookstack_client:\n raise ConnectorMissingCredentialError(\n \"BookStack credentials have not been loaded.\"\n )\n\n try:\n # Attempt to fetch a small batch of books (arbitrary endpoint) to verify credentials\n _ = self.bookstack_client.get(\n \"/books\", params={\"count\": \"1\", \"offset\": \"0\"}\n )\n\n except BookStackClientRequestFailedError as e:\n # Check for HTTP status codes\n if e.status_code == 401:\n raise CredentialExpiredError(\n \"Your BookStack credentials appear to be invalid or expired (HTTP 401).\"\n ) from e\n elif e.status_code == 403:\n raise InsufficientPermissionsError(\n \"The configured BookStack token does not have sufficient permissions (HTTP 403).\"\n ) from e\n else:\n raise ConnectorValidationError(\n f\"Unexpected BookStack error (status={e.status_code}): {e}\"\n ) from e\n\n except Exception as exc:\n raise ConnectorValidationError(\n f\"Unexpected error while validating BookStack connector settings: {exc}\"\n ) from exc\n" }, { "path": "backend/onyx/connectors/canvas/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/canvas/access.py", "content": "\"\"\"\nPermissioning / AccessControl logic for Canvas courses.\n\nCE stub — returns None (no permissions). The EE implementation is loaded\nat runtime via ``fetch_versioned_implementation``.\n\"\"\"\n\nfrom collections.abc import Callable\nfrom typing import cast\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.canvas.client import CanvasApiClient\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\n\ndef get_course_permissions(\n canvas_client: CanvasApiClient,\n course_id: int,\n) -> ExternalAccess | None:\n if not global_version.is_ee_version():\n return None\n\n ee_get_course_permissions = cast(\n Callable[[CanvasApiClient, int], ExternalAccess | None],\n fetch_versioned_implementation(\n \"onyx.external_permissions.canvas.access\",\n \"get_course_permissions\",\n ),\n )\n\n return ee_get_course_permissions(canvas_client, course_id)\n" }, { "path": "backend/onyx/connectors/canvas/client.py", "content": "from __future__ import annotations\n\nimport logging\nimport re\nfrom collections.abc import Iterator\nfrom typing import Any\nfrom urllib.parse import urlparse\n\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rl_requests,\n)\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\n\nlogger = logging.getLogger(__name__)\n\n# Requests timeout in seconds.\n_CANVAS_CALL_TIMEOUT: int = 30\n_CANVAS_API_VERSION: str = \"/api/v1\"\n# Matches the \"next\" URL in a Canvas Link header, e.g.:\n# ; rel=\"next\"\n# Captures the URL inside the angle brackets.\n_NEXT_LINK_PATTERN: re.Pattern[str] = re.compile(r'<([^>]+)>;\\s*rel=\"next\"')\n\n\n_STATUS_TO_ERROR_CODE: dict[int, OnyxErrorCode] = {\n 401: OnyxErrorCode.CREDENTIAL_EXPIRED,\n 403: OnyxErrorCode.INSUFFICIENT_PERMISSIONS,\n 404: OnyxErrorCode.BAD_GATEWAY,\n 429: OnyxErrorCode.RATE_LIMITED,\n}\n\n\ndef _error_code_for_status(status_code: int) -> OnyxErrorCode:\n \"\"\"Map an HTTP status code to the appropriate OnyxErrorCode.\n\n Expects a >= 400 status code. Known codes (401, 403, 404, 429) are\n mapped to specific error codes; all other codes (unrecognised 4xx\n and 5xx) map to BAD_GATEWAY as unexpected upstream errors.\n \"\"\"\n if status_code in _STATUS_TO_ERROR_CODE:\n return _STATUS_TO_ERROR_CODE[status_code]\n return OnyxErrorCode.BAD_GATEWAY\n\n\nclass CanvasApiClient:\n def __init__(\n self,\n bearer_token: str,\n canvas_base_url: str,\n ) -> None:\n parsed_base = urlparse(canvas_base_url)\n if not parsed_base.hostname:\n raise ValueError(\"canvas_base_url must include a valid host\")\n if parsed_base.scheme != \"https\":\n raise ValueError(\"canvas_base_url must use https\")\n\n self._bearer_token = bearer_token\n self.base_url = (\n canvas_base_url.rstrip(\"/\").removesuffix(_CANVAS_API_VERSION)\n + _CANVAS_API_VERSION\n )\n # Hostname is already validated above; reuse parsed_base instead\n # of re-parsing. Used by _parse_next_link to validate pagination URLs.\n self._expected_host: str = parsed_base.hostname\n\n def get(\n self,\n endpoint: str = \"\",\n params: dict[str, Any] | None = None,\n full_url: str | None = None,\n ) -> tuple[Any, str | None]:\n \"\"\"Make a GET request to the Canvas API.\n\n Returns a tuple of (json_body, next_url).\n next_url is parsed from the Link header and is None if there are no more pages.\n If full_url is provided, it is used directly (for following pagination links).\n\n Security note: full_url must only be set to values returned by\n ``_parse_next_link``, which validates the host against the configured\n Canvas base URL. Passing an arbitrary URL would leak the bearer token.\n \"\"\"\n # full_url is used when following pagination (Canvas returns the\n # next-page URL in the Link header). For the first request we build\n # the URL from the endpoint name instead.\n url = full_url if full_url else self._build_url(endpoint)\n headers = self._build_headers()\n\n response = rl_requests.get(\n url,\n headers=headers,\n params=params if not full_url else None,\n timeout=_CANVAS_CALL_TIMEOUT,\n )\n\n try:\n response_json = response.json()\n except ValueError as e:\n if response.status_code < 300:\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY,\n detail=f\"Invalid JSON in Canvas response: {e}\",\n )\n logger.warning(\n \"Failed to parse JSON from Canvas error response (status=%d): %s\",\n response.status_code,\n e,\n )\n response_json = {}\n\n if response.status_code >= 400:\n # Try to extract the most specific error message from the\n # Canvas response body. Canvas uses three different shapes\n # depending on the endpoint and error type:\n default_error: str = response.reason or f\"HTTP {response.status_code}\"\n error = default_error\n if isinstance(response_json, dict):\n # Shape 1: {\"error\": {\"message\": \"Not authorized\"}}\n error_field = response_json.get(\"error\")\n if isinstance(error_field, dict):\n response_error = error_field.get(\"message\", \"\")\n if response_error:\n error = response_error\n # Shape 2: {\"error\": \"Invalid access token\"}\n elif isinstance(error_field, str):\n error = error_field\n # Shape 3: {\"errors\": [{\"message\": \"...\"}]}\n # Used for validation errors. Only use as fallback if\n # we didn't already find a more specific message above.\n if error == default_error:\n errors_list = response_json.get(\"errors\")\n if isinstance(errors_list, list) and errors_list:\n first_error = errors_list[0]\n if isinstance(first_error, dict):\n msg = first_error.get(\"message\", \"\")\n if msg:\n error = msg\n raise OnyxError(\n _error_code_for_status(response.status_code),\n detail=error,\n status_code_override=response.status_code,\n )\n\n next_url = self._parse_next_link(response.headers.get(\"Link\", \"\"))\n return response_json, next_url\n\n def _parse_next_link(self, link_header: str) -> str | None:\n \"\"\"Extract the 'next' URL from a Canvas Link header.\n\n Only returns URLs whose host matches the configured Canvas base URL\n to prevent leaking the bearer token to arbitrary hosts.\n \"\"\"\n expected_host = self._expected_host\n for match in _NEXT_LINK_PATTERN.finditer(link_header):\n url = match.group(1)\n parsed_url = urlparse(url)\n if parsed_url.hostname != expected_host:\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY,\n detail=(\n \"Canvas pagination returned an unexpected host \"\n f\"({parsed_url.hostname}); expected {expected_host}\"\n ),\n )\n if parsed_url.scheme != \"https\":\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY,\n detail=(\n \"Canvas pagination link must use https, \"\n f\"got {parsed_url.scheme!r}\"\n ),\n )\n return url\n return None\n\n def _build_headers(self) -> dict[str, str]:\n \"\"\"Return the Authorization header with the bearer token.\"\"\"\n return {\"Authorization\": f\"Bearer {self._bearer_token}\"}\n\n def _build_url(self, endpoint: str) -> str:\n \"\"\"Build a full Canvas API URL from an endpoint path.\n\n Assumes endpoint is non-empty (e.g. ``\"courses\"``, ``\"announcements\"``).\n Only called on a first request, endpoint must be set for first request.\n Verify endpoint exists in case of future changes where endpoint might be optional.\n Leading slashes are stripped to avoid double-slash in the result.\n self.base_url is already normalized with no trailing slash.\n \"\"\"\n final_url = self.base_url\n clean_endpoint = endpoint.lstrip(\"/\")\n if clean_endpoint:\n final_url += \"/\" + clean_endpoint\n return final_url\n\n def paginate(\n self,\n endpoint: str,\n params: dict[str, Any] | None = None,\n ) -> Iterator[list[Any]]:\n \"\"\"Yield each page of results, following Link-header pagination.\n\n Makes the first request with endpoint + params, then follows\n next_url from Link headers for subsequent pages.\n \"\"\"\n response, next_url = self.get(endpoint, params=params)\n while True:\n if not response:\n break\n yield response\n if not next_url:\n break\n response, next_url = self.get(full_url=next_url)\n" }, { "path": "backend/onyx/connectors/canvas/connector.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import Literal\nfrom typing import NoReturn\nfrom typing import TypeAlias\n\nfrom pydantic import BaseModel\nfrom retry import retry\nfrom typing_extensions import override\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.canvas.access import get_course_permissions\nfrom onyx.connectors.canvas.client import CanvasApiClient\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _handle_canvas_api_error(e: OnyxError) -> NoReturn:\n \"\"\"Map Canvas API errors to connector framework exceptions.\"\"\"\n if e.status_code == 401:\n raise CredentialExpiredError(\n \"Canvas API token is invalid or expired (HTTP 401).\"\n )\n elif e.status_code == 403:\n raise InsufficientPermissionsError(\n \"Canvas API token does not have sufficient permissions (HTTP 403).\"\n )\n elif e.status_code == 429:\n raise ConnectorValidationError(\n \"Canvas rate-limit exceeded (HTTP 429). Please try again later.\"\n )\n elif e.status_code >= 500:\n raise UnexpectedValidationError(\n f\"Unexpected Canvas HTTP error (status={e.status_code}): {e}\"\n )\n else:\n raise ConnectorValidationError(\n f\"Canvas API error (status={e.status_code}): {e}\"\n )\n\n\nclass CanvasCourse(BaseModel):\n id: int\n name: str | None = None\n course_code: str | None = None\n created_at: str | None = None\n workflow_state: str | None = None\n\n @classmethod\n def from_api(cls, payload: dict[str, Any]) -> \"CanvasCourse\":\n return cls(\n id=payload[\"id\"],\n name=payload.get(\"name\"),\n course_code=payload.get(\"course_code\"),\n created_at=payload.get(\"created_at\"),\n workflow_state=payload.get(\"workflow_state\"),\n )\n\n\nclass CanvasPage(BaseModel):\n page_id: int\n url: str\n title: str\n body: str | None = None\n created_at: str | None = None\n updated_at: str | None = None\n course_id: int\n\n @classmethod\n def from_api(cls, payload: dict[str, Any], course_id: int) -> \"CanvasPage\":\n return cls(\n page_id=payload[\"page_id\"],\n url=payload[\"url\"],\n title=payload[\"title\"],\n body=payload.get(\"body\"),\n created_at=payload.get(\"created_at\"),\n updated_at=payload.get(\"updated_at\"),\n course_id=course_id,\n )\n\n\nclass CanvasAssignment(BaseModel):\n id: int\n name: str\n description: str | None = None\n html_url: str\n course_id: int\n created_at: str | None = None\n updated_at: str | None = None\n due_at: str | None = None\n\n @classmethod\n def from_api(cls, payload: dict[str, Any], course_id: int) -> \"CanvasAssignment\":\n return cls(\n id=payload[\"id\"],\n name=payload[\"name\"],\n description=payload.get(\"description\"),\n html_url=payload[\"html_url\"],\n course_id=course_id,\n created_at=payload.get(\"created_at\"),\n updated_at=payload.get(\"updated_at\"),\n due_at=payload.get(\"due_at\"),\n )\n\n\nclass CanvasAnnouncement(BaseModel):\n id: int\n title: str\n message: str | None = None\n html_url: str\n posted_at: str | None = None\n course_id: int\n\n @classmethod\n def from_api(cls, payload: dict[str, Any], course_id: int) -> \"CanvasAnnouncement\":\n return cls(\n id=payload[\"id\"],\n title=payload[\"title\"],\n message=payload.get(\"message\"),\n html_url=payload[\"html_url\"],\n posted_at=payload.get(\"posted_at\"),\n course_id=course_id,\n )\n\n\nCanvasStage: TypeAlias = Literal[\"pages\", \"assignments\", \"announcements\"]\n\n\nclass CanvasConnectorCheckpoint(ConnectorCheckpoint):\n \"\"\"Checkpoint state for resumable Canvas indexing.\n\n Fields:\n course_ids: Materialized list of course IDs to process.\n current_course_index: Index into course_ids for current course.\n stage: Which item type we're processing for the current course.\n next_url: Pagination cursor within the current stage. None means\n start from the first page; a URL means resume from that page.\n\n Invariant:\n If current_course_index is incremented, stage must be reset to\n \"pages\" and next_url must be reset to None.\n \"\"\"\n\n course_ids: list[int] = []\n current_course_index: int = 0\n stage: CanvasStage = \"pages\"\n next_url: str | None = None\n\n def advance_course(self) -> None:\n \"\"\"Move to the next course and reset within-course state.\"\"\"\n self.current_course_index += 1\n self.stage = \"pages\"\n self.next_url = None\n\n\nclass CanvasConnector(\n CheckpointedConnectorWithPermSync[CanvasConnectorCheckpoint],\n SlimConnectorWithPermSync,\n):\n def __init__(\n self,\n canvas_base_url: str,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.canvas_base_url = canvas_base_url.rstrip(\"/\").removesuffix(\"/api/v1\")\n self.batch_size = batch_size\n self._canvas_client: CanvasApiClient | None = None\n self._course_permissions_cache: dict[int, ExternalAccess | None] = {}\n\n @property\n def canvas_client(self) -> CanvasApiClient:\n if self._canvas_client is None:\n raise ConnectorMissingCredentialError(\"Canvas\")\n return self._canvas_client\n\n def _get_course_permissions(self, course_id: int) -> ExternalAccess | None:\n \"\"\"Get course permissions with caching.\"\"\"\n if course_id not in self._course_permissions_cache:\n self._course_permissions_cache[course_id] = get_course_permissions(\n canvas_client=self.canvas_client,\n course_id=course_id,\n )\n return self._course_permissions_cache[course_id]\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_courses(self) -> list[CanvasCourse]:\n \"\"\"Fetch all courses accessible to the authenticated user.\"\"\"\n logger.debug(\"Fetching Canvas courses\")\n\n courses: list[CanvasCourse] = []\n for page in self.canvas_client.paginate(\n \"courses\", params={\"per_page\": \"100\", \"state[]\": \"available\"}\n ):\n courses.extend(CanvasCourse.from_api(c) for c in page)\n return courses\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_pages(self, course_id: int) -> list[CanvasPage]:\n \"\"\"Fetch all pages for a given course.\"\"\"\n logger.debug(f\"Fetching pages for course {course_id}\")\n\n pages: list[CanvasPage] = []\n for page in self.canvas_client.paginate(\n f\"courses/{course_id}/pages\",\n params={\"per_page\": \"100\", \"include[]\": \"body\", \"published\": \"true\"},\n ):\n pages.extend(CanvasPage.from_api(p, course_id=course_id) for p in page)\n return pages\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_assignments(self, course_id: int) -> list[CanvasAssignment]:\n \"\"\"Fetch all assignments for a given course.\"\"\"\n logger.debug(f\"Fetching assignments for course {course_id}\")\n\n assignments: list[CanvasAssignment] = []\n for page in self.canvas_client.paginate(\n f\"courses/{course_id}/assignments\",\n params={\"per_page\": \"100\", \"published\": \"true\"},\n ):\n assignments.extend(\n CanvasAssignment.from_api(a, course_id=course_id) for a in page\n )\n return assignments\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_announcements(self, course_id: int) -> list[CanvasAnnouncement]:\n \"\"\"Fetch all announcements for a given course.\"\"\"\n logger.debug(f\"Fetching announcements for course {course_id}\")\n\n announcements: list[CanvasAnnouncement] = []\n for page in self.canvas_client.paginate(\n \"announcements\",\n params={\n \"per_page\": \"100\",\n \"context_codes[]\": f\"course_{course_id}\",\n \"active_only\": \"true\",\n },\n ):\n announcements.extend(\n CanvasAnnouncement.from_api(a, course_id=course_id) for a in page\n )\n return announcements\n\n def _build_document(\n self,\n doc_id: str,\n link: str,\n text: str,\n semantic_identifier: str,\n doc_updated_at: datetime | None,\n course_id: int,\n doc_type: str,\n ) -> Document:\n \"\"\"Build a Document with standard Canvas fields.\"\"\"\n return Document(\n id=doc_id,\n sections=cast(\n list[TextSection | ImageSection],\n [TextSection(link=link, text=text)],\n ),\n source=DocumentSource.CANVAS,\n semantic_identifier=semantic_identifier,\n doc_updated_at=doc_updated_at,\n metadata={\"course_id\": str(course_id), \"type\": doc_type},\n )\n\n def _convert_page_to_document(self, page: CanvasPage) -> Document:\n \"\"\"Convert a Canvas page to a Document.\"\"\"\n link = f\"{self.canvas_base_url}/courses/{page.course_id}/pages/{page.url}\"\n\n text_parts = [page.title]\n body_text = parse_html_page_basic(page.body) if page.body else \"\"\n if body_text:\n text_parts.append(body_text)\n\n doc_updated_at = (\n datetime.fromisoformat(page.updated_at.replace(\"Z\", \"+00:00\")).astimezone(\n timezone.utc\n )\n if page.updated_at\n else None\n )\n\n document = self._build_document(\n doc_id=f\"canvas-page-{page.course_id}-{page.page_id}\",\n link=link,\n text=\"\\n\\n\".join(text_parts),\n semantic_identifier=page.title or f\"Page {page.page_id}\",\n doc_updated_at=doc_updated_at,\n course_id=page.course_id,\n doc_type=\"page\",\n )\n return document\n\n def _convert_assignment_to_document(self, assignment: CanvasAssignment) -> Document:\n \"\"\"Convert a Canvas assignment to a Document.\"\"\"\n text_parts = [assignment.name]\n desc_text = (\n parse_html_page_basic(assignment.description)\n if assignment.description\n else \"\"\n )\n if desc_text:\n text_parts.append(desc_text)\n if assignment.due_at:\n due_dt = datetime.fromisoformat(\n assignment.due_at.replace(\"Z\", \"+00:00\")\n ).astimezone(timezone.utc)\n text_parts.append(f\"Due: {due_dt.strftime('%B %d, %Y %H:%M UTC')}\")\n\n doc_updated_at = (\n datetime.fromisoformat(\n assignment.updated_at.replace(\"Z\", \"+00:00\")\n ).astimezone(timezone.utc)\n if assignment.updated_at\n else None\n )\n\n document = self._build_document(\n doc_id=f\"canvas-assignment-{assignment.course_id}-{assignment.id}\",\n link=assignment.html_url,\n text=\"\\n\\n\".join(text_parts),\n semantic_identifier=assignment.name or f\"Assignment {assignment.id}\",\n doc_updated_at=doc_updated_at,\n course_id=assignment.course_id,\n doc_type=\"assignment\",\n )\n return document\n\n def _convert_announcement_to_document(\n self, announcement: CanvasAnnouncement\n ) -> Document:\n \"\"\"Convert a Canvas announcement to a Document.\"\"\"\n text_parts = [announcement.title]\n msg_text = (\n parse_html_page_basic(announcement.message) if announcement.message else \"\"\n )\n if msg_text:\n text_parts.append(msg_text)\n\n doc_updated_at = (\n datetime.fromisoformat(\n announcement.posted_at.replace(\"Z\", \"+00:00\")\n ).astimezone(timezone.utc)\n if announcement.posted_at\n else None\n )\n\n document = self._build_document(\n doc_id=f\"canvas-announcement-{announcement.course_id}-{announcement.id}\",\n link=announcement.html_url,\n text=\"\\n\\n\".join(text_parts),\n semantic_identifier=announcement.title or f\"Announcement {announcement.id}\",\n doc_updated_at=doc_updated_at,\n course_id=announcement.course_id,\n doc_type=\"announcement\",\n )\n return document\n\n @override\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"Load and validate Canvas credentials.\"\"\"\n access_token = credentials.get(\"canvas_access_token\")\n if not access_token:\n raise ConnectorMissingCredentialError(\"Canvas\")\n\n try:\n client = CanvasApiClient(\n bearer_token=access_token,\n canvas_base_url=self.canvas_base_url,\n )\n client.get(\"courses\", params={\"per_page\": \"1\"})\n except ValueError as e:\n raise ConnectorValidationError(f\"Invalid Canvas base URL: {e}\")\n except OnyxError as e:\n _handle_canvas_api_error(e)\n\n self._canvas_client = client\n return None\n\n @override\n def validate_connector_settings(self) -> None:\n \"\"\"Validate Canvas connector settings by testing API access.\"\"\"\n try:\n self.canvas_client.get(\"courses\", params={\"per_page\": \"1\"})\n logger.info(\"Canvas connector settings validated successfully\")\n except OnyxError as e:\n _handle_canvas_api_error(e)\n except ConnectorMissingCredentialError:\n raise\n except Exception as exc:\n raise UnexpectedValidationError(\n f\"Unexpected error during Canvas settings validation: {exc}\"\n )\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: CanvasConnectorCheckpoint,\n ) -> CheckpointOutput[CanvasConnectorCheckpoint]:\n # TODO(benwu408): implemented in PR3 (checkpoint)\n raise NotImplementedError\n\n @override\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: CanvasConnectorCheckpoint,\n ) -> CheckpointOutput[CanvasConnectorCheckpoint]:\n # TODO(benwu408): implemented in PR3 (checkpoint)\n raise NotImplementedError\n\n @override\n def build_dummy_checkpoint(self) -> CanvasConnectorCheckpoint:\n # TODO(benwu408): implemented in PR3 (checkpoint)\n raise NotImplementedError\n\n @override\n def validate_checkpoint_json(\n self, checkpoint_json: str\n ) -> CanvasConnectorCheckpoint:\n # TODO(benwu408): implemented in PR3 (checkpoint)\n raise NotImplementedError\n\n @override\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n # TODO(benwu408): implemented in PR4 (perm sync)\n raise NotImplementedError\n" }, { "path": "backend/onyx/connectors/clickup/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/clickup/connector.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import Optional\n\nimport requests\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.retry_wrapper import retry_builder\n\n\nCLICKUP_API_BASE_URL = \"https://api.clickup.com/api/v2\"\n\n\nclass ClickupConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n api_token: str | None = None,\n team_id: str | None = None,\n connector_type: str | None = None,\n connector_ids: list[str] | None = None,\n retrieve_task_comments: bool = True,\n ) -> None:\n self.batch_size = batch_size\n self.api_token = api_token\n self.team_id = team_id\n self.connector_type = connector_type if connector_type else \"workspace\"\n self.connector_ids = connector_ids\n self.retrieve_task_comments = retrieve_task_comments\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.api_token = credentials[\"clickup_api_token\"]\n self.team_id = credentials[\"clickup_team_id\"]\n return None\n\n @retry_builder()\n @rate_limit_builder(max_calls=100, period=60)\n def _make_request(self, endpoint: str, params: Optional[dict] = None) -> Any:\n if not self.api_token:\n raise ConnectorMissingCredentialError(\"Clickup\")\n\n headers = {\"Authorization\": self.api_token}\n\n response = requests.get(\n f\"{CLICKUP_API_BASE_URL}/{endpoint}\", headers=headers, params=params\n )\n\n response.raise_for_status()\n\n return response.json()\n\n def _get_task_comments(self, task_id: str) -> list[TextSection]:\n url_endpoint = f\"/task/{task_id}/comment\"\n response = self._make_request(url_endpoint)\n comments = [\n TextSection(\n link=f\"https://app.clickup.com/t/{task_id}?comment={comment_dict['id']}\",\n text=comment_dict[\"comment_text\"],\n )\n for comment_dict in response[\"comments\"]\n ]\n\n return comments\n\n def _get_all_tasks_filtered(\n self,\n start: int | None = None,\n end: int | None = None,\n ) -> GenerateDocumentsOutput:\n doc_batch: list[Document | HierarchyNode] = []\n page: int = 0\n params = {\n \"include_markdown_description\": \"true\",\n \"include_closed\": \"true\",\n \"page\": page,\n }\n\n if start is not None:\n params[\"date_updated_gt\"] = start\n if end is not None:\n params[\"date_updated_lt\"] = end\n\n if self.connector_type == \"list\":\n params[\"list_ids[]\"] = self.connector_ids\n elif self.connector_type == \"folder\":\n params[\"project_ids[]\"] = self.connector_ids\n elif self.connector_type == \"space\":\n params[\"space_ids[]\"] = self.connector_ids\n\n url_endpoint = f\"/team/{self.team_id}/task\"\n\n while True:\n response = self._make_request(url_endpoint, params)\n\n page += 1\n params[\"page\"] = page\n\n for task in response[\"tasks\"]:\n document = Document(\n id=task[\"id\"],\n source=DocumentSource.CLICKUP,\n semantic_identifier=task[\"name\"],\n doc_updated_at=(\n datetime.fromtimestamp(\n round(float(task[\"date_updated\"]) / 1000, 3)\n ).replace(tzinfo=timezone.utc)\n ),\n primary_owners=[\n BasicExpertInfo(\n display_name=task[\"creator\"][\"username\"],\n email=task[\"creator\"][\"email\"],\n )\n ],\n secondary_owners=[\n BasicExpertInfo(\n display_name=assignee[\"username\"],\n email=assignee[\"email\"],\n )\n for assignee in task[\"assignees\"]\n ],\n title=task[\"name\"],\n sections=[\n TextSection(\n link=task[\"url\"],\n text=(\n task[\"markdown_description\"]\n if \"markdown_description\" in task\n else task[\"description\"]\n ),\n )\n ],\n metadata={\n \"id\": task[\"id\"],\n \"status\": task[\"status\"][\"status\"],\n \"list\": task[\"list\"][\"name\"],\n \"project\": task[\"project\"][\"name\"],\n \"folder\": task[\"folder\"][\"name\"],\n \"space_id\": task[\"space\"][\"id\"],\n \"tags\": [tag[\"name\"] for tag in task[\"tags\"]],\n \"priority\": (\n task[\"priority\"][\"priority\"]\n if \"priority\" in task and task[\"priority\"] is not None\n else \"\"\n ),\n },\n )\n\n extra_fields = [\n \"date_created\",\n \"date_updated\",\n \"date_closed\",\n \"date_done\",\n \"due_date\",\n ]\n for extra_field in extra_fields:\n if extra_field in task and task[extra_field] is not None:\n document.metadata[extra_field] = task[extra_field]\n\n if self.retrieve_task_comments:\n document.sections.extend(self._get_task_comments(task[\"id\"]))\n\n doc_batch.append(document)\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if response.get(\"last_page\") is True or len(response[\"tasks\"]) < 100:\n break\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n if self.api_token is None:\n raise ConnectorMissingCredentialError(\"Clickup\")\n\n return self._get_all_tasks_filtered(None, None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n if self.api_token is None:\n raise ConnectorMissingCredentialError(\"Clickup\")\n\n return self._get_all_tasks_filtered(int(start * 1000), int(end * 1000))\n\n\nif __name__ == \"__main__\":\n import os\n\n clickup_connector = ClickupConnector()\n\n clickup_connector.load_credentials(\n {\n \"clickup_api_token\": os.environ[\"clickup_api_token\"],\n \"clickup_team_id\": os.environ[\"clickup_team_id\"],\n }\n )\n\n latest_docs = clickup_connector.load_from_state()\n\n for doc in latest_docs:\n print(doc)\n" }, { "path": "backend/onyx/connectors/coda/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/coda/connector.py", "content": "import os\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\nfrom typing import List\nfrom typing import Optional\n\nfrom pydantic import BaseModel\nfrom retry import retry\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rl_requests,\n)\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.batching import batch_generator\nfrom onyx.utils.logger import setup_logger\n\n_CODA_CALL_TIMEOUT = 30\n_CODA_BASE_URL = \"https://coda.io/apis/v1\"\n\nlogger = setup_logger()\n\n\nclass CodaClientRequestFailedError(ConnectionError):\n def __init__(self, message: str, status_code: int):\n super().__init__(\n f\"Coda API request failed with status {status_code}: {message}\"\n )\n self.status_code = status_code\n\n\nclass CodaDoc(BaseModel):\n id: str\n browser_link: str\n name: str\n created_at: str\n updated_at: str\n workspace_id: str\n workspace_name: str\n folder_id: str | None\n folder_name: str | None\n\n\nclass CodaPage(BaseModel):\n id: str\n browser_link: str\n name: str\n content_type: str\n created_at: str\n updated_at: str\n doc_id: str\n\n\nclass CodaTable(BaseModel):\n id: str\n name: str\n browser_link: str\n created_at: str\n updated_at: str\n doc_id: str\n\n\nclass CodaRow(BaseModel):\n id: str\n name: Optional[str] = None\n index: Optional[int] = None\n browser_link: str\n created_at: str\n updated_at: str\n values: Dict[str, Any]\n table_id: str\n doc_id: str\n\n\nclass CodaApiClient:\n def __init__(\n self,\n bearer_token: str,\n ) -> None:\n self.bearer_token = bearer_token\n self.base_url = os.environ.get(\"CODA_BASE_URL\", _CODA_BASE_URL)\n\n def get(\n self, endpoint: str, params: Optional[dict[str, str]] = None\n ) -> dict[str, Any]:\n url = self._build_url(endpoint)\n headers = self._build_headers()\n\n response = rl_requests.get(\n url, headers=headers, params=params, timeout=_CODA_CALL_TIMEOUT\n )\n\n try:\n json = response.json()\n except Exception:\n json = {}\n\n if response.status_code >= 300:\n error = response.reason\n response_error = json.get(\"error\", {}).get(\"message\", \"\")\n if response_error:\n error = response_error\n raise CodaClientRequestFailedError(error, response.status_code)\n\n return json\n\n def _build_headers(self) -> Dict[str, str]:\n return {\"Authorization\": f\"Bearer {self.bearer_token}\"}\n\n def _build_url(self, endpoint: str) -> str:\n return self.base_url.rstrip(\"/\") + \"/\" + endpoint.lstrip(\"/\")\n\n\nclass CodaConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n index_page_content: bool = True,\n workspace_id: str | None = None,\n ) -> None:\n self.batch_size = batch_size\n self.index_page_content = index_page_content\n self.workspace_id = workspace_id\n self._coda_client: CodaApiClient | None = None\n\n @property\n def coda_client(self) -> CodaApiClient:\n if self._coda_client is None:\n raise ConnectorMissingCredentialError(\"Coda\")\n return self._coda_client\n\n @retry(tries=3, delay=1, backoff=2)\n def _get_doc(self, doc_id: str) -> CodaDoc:\n \"\"\"Fetch a specific Coda document by its ID.\"\"\"\n logger.debug(f\"Fetching Coda doc with ID: {doc_id}\")\n try:\n response = self.coda_client.get(f\"docs/{doc_id}\")\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(f\"Failed to fetch doc: {doc_id}\") from e\n else:\n raise\n\n return CodaDoc(\n id=response[\"id\"],\n browser_link=response[\"browserLink\"],\n name=response[\"name\"],\n created_at=response[\"createdAt\"],\n updated_at=response[\"updatedAt\"],\n workspace_id=response[\"workspace\"][\"id\"],\n workspace_name=response[\"workspace\"][\"name\"],\n folder_id=response[\"folder\"][\"id\"] if response.get(\"folder\") else None,\n folder_name=response[\"folder\"][\"name\"] if response.get(\"folder\") else None,\n )\n\n @retry(tries=3, delay=1, backoff=2)\n def _get_page(self, doc_id: str, page_id: str) -> CodaPage:\n \"\"\"Fetch a specific page from a Coda document.\"\"\"\n logger.debug(f\"Fetching Coda page with ID: {page_id}\")\n try:\n response = self.coda_client.get(f\"docs/{doc_id}/pages/{page_id}\")\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\n f\"Failed to fetch page: {page_id} from doc: {doc_id}\"\n ) from e\n else:\n raise\n\n return CodaPage(\n id=response[\"id\"],\n doc_id=doc_id,\n browser_link=response[\"browserLink\"],\n name=response[\"name\"],\n content_type=response[\"contentType\"],\n created_at=response[\"createdAt\"],\n updated_at=response[\"updatedAt\"],\n )\n\n @retry(tries=3, delay=1, backoff=2)\n def _get_table(self, doc_id: str, table_id: str) -> CodaTable:\n \"\"\"Fetch a specific table from a Coda document.\"\"\"\n logger.debug(f\"Fetching Coda table with ID: {table_id}\")\n try:\n response = self.coda_client.get(f\"docs/{doc_id}/tables/{table_id}\")\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\n f\"Failed to fetch table: {table_id} from doc: {doc_id}\"\n ) from e\n else:\n raise\n\n return CodaTable(\n id=response[\"id\"],\n name=response[\"name\"],\n browser_link=response[\"browserLink\"],\n created_at=response[\"createdAt\"],\n updated_at=response[\"updatedAt\"],\n doc_id=doc_id,\n )\n\n @retry(tries=3, delay=1, backoff=2)\n def _get_row(self, doc_id: str, table_id: str, row_id: str) -> CodaRow:\n \"\"\"Fetch a specific row from a Coda table.\"\"\"\n logger.debug(f\"Fetching Coda row with ID: {row_id}\")\n try:\n response = self.coda_client.get(\n f\"docs/{doc_id}/tables/{table_id}/rows/{row_id}\"\n )\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\n f\"Failed to fetch row: {row_id} from table: {table_id} in doc: {doc_id}\"\n ) from e\n else:\n raise\n\n values = {}\n for col_name, col_value in response.get(\"values\", {}).items():\n values[col_name] = col_value\n\n return CodaRow(\n id=response[\"id\"],\n name=response.get(\"name\"),\n index=response.get(\"index\"),\n browser_link=response[\"browserLink\"],\n created_at=response[\"createdAt\"],\n updated_at=response[\"updatedAt\"],\n values=values,\n table_id=table_id,\n doc_id=doc_id,\n )\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_all_docs(\n self, endpoint: str = \"docs\", params: Optional[Dict[str, str]] = None\n ) -> List[CodaDoc]:\n \"\"\"List all Coda documents in the workspace.\"\"\"\n logger.debug(\"Listing documents in Coda\")\n\n all_docs: List[CodaDoc] = []\n next_page_token: str | None = None\n params = params or {}\n\n if self.workspace_id:\n params[\"workspaceId\"] = self.workspace_id\n\n while True:\n if next_page_token:\n params[\"pageToken\"] = next_page_token\n\n try:\n response = self.coda_client.get(endpoint, params=params)\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\"Failed to list docs\") from e\n else:\n raise\n\n items = response.get(\"items\", [])\n\n for item in items:\n doc = CodaDoc(\n id=item[\"id\"],\n browser_link=item[\"browserLink\"],\n name=item[\"name\"],\n created_at=item[\"createdAt\"],\n updated_at=item[\"updatedAt\"],\n workspace_id=item[\"workspace\"][\"id\"],\n workspace_name=item[\"workspace\"][\"name\"],\n folder_id=item[\"folder\"][\"id\"] if item.get(\"folder\") else None,\n folder_name=item[\"folder\"][\"name\"] if item.get(\"folder\") else None,\n )\n all_docs.append(doc)\n\n next_page_token = response.get(\"nextPageToken\")\n if not next_page_token:\n break\n\n logger.debug(f\"Found {len(all_docs)} docs\")\n return all_docs\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_pages_in_doc(self, doc_id: str) -> List[CodaPage]:\n \"\"\"List all pages in a Coda document.\"\"\"\n logger.debug(f\"Listing pages in Coda doc with ID: {doc_id}\")\n\n pages: List[CodaPage] = []\n endpoint = f\"docs/{doc_id}/pages\"\n params: Dict[str, str] = {}\n next_page_token: str | None = None\n\n while True:\n if next_page_token:\n params[\"pageToken\"] = next_page_token\n\n try:\n response = self.coda_client.get(endpoint, params=params)\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\n f\"Failed to list pages for doc: {doc_id}\"\n ) from e\n else:\n raise\n\n items = response.get(\"items\", [])\n for item in items:\n # can be removed if we don't care to skip hidden pages\n if item.get(\"isHidden\", False):\n continue\n\n pages.append(\n CodaPage(\n id=item[\"id\"],\n browser_link=item[\"browserLink\"],\n name=item[\"name\"],\n content_type=item[\"contentType\"],\n created_at=item[\"createdAt\"],\n updated_at=item[\"updatedAt\"],\n doc_id=doc_id,\n )\n )\n\n next_page_token = response.get(\"nextPageToken\")\n if not next_page_token:\n break\n\n logger.debug(f\"Found {len(pages)} pages in doc {doc_id}\")\n return pages\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_page_content(self, doc_id: str, page_id: str) -> str:\n \"\"\"Fetch the content of a Coda page.\"\"\"\n logger.debug(f\"Fetching content for page {page_id} in doc {doc_id}\")\n\n content_parts = []\n next_page_token: str | None = None\n params: Dict[str, str] = {}\n\n while True:\n if next_page_token:\n params[\"pageToken\"] = next_page_token\n\n try:\n response = self.coda_client.get(\n f\"docs/{doc_id}/pages/{page_id}/content\", params=params\n )\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n logger.debug(f\"No content available for page {page_id}\")\n return \"\"\n raise\n\n items = response.get(\"items\", [])\n\n for item in items:\n item_content = item.get(\"itemContent\", {})\n\n content_text = item_content.get(\"content\", \"\")\n if content_text:\n content_parts.append(content_text)\n\n next_page_token = response.get(\"nextPageToken\")\n if not next_page_token:\n break\n\n return \"\\n\\n\".join(content_parts)\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_tables(self, doc_id: str) -> List[CodaTable]:\n \"\"\"List all tables in a Coda document.\"\"\"\n logger.debug(f\"Listing tables in Coda doc with ID: {doc_id}\")\n\n tables: List[CodaTable] = []\n endpoint = f\"docs/{doc_id}/tables\"\n params: Dict[str, str] = {}\n next_page_token: str | None = None\n\n while True:\n if next_page_token:\n params[\"pageToken\"] = next_page_token\n\n try:\n response = self.coda_client.get(endpoint, params=params)\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\n f\"Failed to list tables for doc: {doc_id}\"\n ) from e\n else:\n raise\n\n items = response.get(\"items\", [])\n for item in items:\n tables.append(\n CodaTable(\n id=item[\"id\"],\n browser_link=item[\"browserLink\"],\n name=item[\"name\"],\n created_at=item[\"createdAt\"],\n updated_at=item[\"updatedAt\"],\n doc_id=doc_id,\n )\n )\n\n next_page_token = response.get(\"nextPageToken\")\n if not next_page_token:\n break\n\n logger.debug(f\"Found {len(tables)} tables in doc {doc_id}\")\n return tables\n\n @retry(tries=3, delay=1, backoff=2)\n def _list_rows_and_values(self, doc_id: str, table_id: str) -> List[CodaRow]:\n \"\"\"List all rows and their values in a table.\"\"\"\n logger.debug(f\"Listing rows in Coda table: {table_id} in Coda doc: {doc_id}\")\n\n rows: List[CodaRow] = []\n endpoint = f\"docs/{doc_id}/tables/{table_id}/rows\"\n params: Dict[str, str] = {\"valueFormat\": \"rich\"}\n next_page_token: str | None = None\n\n while True:\n if next_page_token:\n params[\"pageToken\"] = next_page_token\n\n try:\n response = self.coda_client.get(endpoint, params=params)\n except CodaClientRequestFailedError as e:\n if e.status_code == 404:\n raise ConnectorValidationError(\n f\"Failed to list rows for table: {table_id} in doc: {doc_id}\"\n ) from e\n else:\n raise\n\n items = response.get(\"items\", [])\n for item in items:\n values = {}\n for col_name, col_value in item.get(\"values\", {}).items():\n values[col_name] = col_value\n\n rows.append(\n CodaRow(\n id=item[\"id\"],\n name=item[\"name\"],\n index=item[\"index\"],\n browser_link=item[\"browserLink\"],\n created_at=item[\"createdAt\"],\n updated_at=item[\"updatedAt\"],\n values=values,\n table_id=table_id,\n doc_id=doc_id,\n )\n )\n\n next_page_token = response.get(\"nextPageToken\")\n if not next_page_token:\n break\n\n logger.debug(f\"Found {len(rows)} rows in table {table_id}\")\n return rows\n\n def _convert_page_to_document(self, page: CodaPage, content: str = \"\") -> Document:\n \"\"\"Convert a page into a Document.\"\"\"\n page_updated = datetime.fromisoformat(page.updated_at).astimezone(timezone.utc)\n\n text_parts = [page.name, page.browser_link]\n if content:\n text_parts.append(content)\n\n sections = [TextSection(link=page.browser_link, text=\"\\n\\n\".join(text_parts))]\n\n return Document(\n id=f\"coda-page-{page.doc_id}-{page.id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.CODA,\n semantic_identifier=page.name or f\"Page {page.id}\",\n doc_updated_at=page_updated,\n metadata={\n \"browser_link\": page.browser_link,\n \"doc_id\": page.doc_id,\n \"content_type\": page.content_type,\n },\n )\n\n def _convert_table_with_rows_to_document(\n self, table: CodaTable, rows: List[CodaRow]\n ) -> Document:\n \"\"\"Convert a table and its rows into a single Document with multiple sections (one per row).\"\"\"\n table_updated = datetime.fromisoformat(table.updated_at).astimezone(\n timezone.utc\n )\n\n sections: List[TextSection] = []\n for row in rows:\n content_text = \" \".join(\n str(v) if not isinstance(v, list) else \" \".join(map(str, v))\n for v in row.values.values()\n )\n\n row_name = row.name or f\"Row {row.index or row.id}\"\n text = f\"{row_name}: {content_text}\" if content_text else row_name\n\n sections.append(TextSection(link=row.browser_link, text=text))\n\n # If no rows, create a single section for the table itself\n if not sections:\n sections = [\n TextSection(link=table.browser_link, text=f\"Table: {table.name}\")\n ]\n\n return Document(\n id=f\"coda-table-{table.doc_id}-{table.id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.CODA,\n semantic_identifier=table.name or f\"Table {table.id}\",\n doc_updated_at=table_updated,\n metadata={\n \"browser_link\": table.browser_link,\n \"doc_id\": table.doc_id,\n \"row_count\": str(len(rows)),\n },\n )\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"Load and validate Coda credentials.\"\"\"\n self._coda_client = CodaApiClient(bearer_token=credentials[\"coda_bearer_token\"])\n\n try:\n self._coda_client.get(\"docs\", params={\"limit\": \"1\"})\n except CodaClientRequestFailedError as e:\n if e.status_code == 401:\n raise ConnectorMissingCredentialError(\"Invalid Coda API token\")\n raise\n\n return None\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"Load all documents from Coda workspace.\"\"\"\n\n def _iter_documents() -> Generator[Document, None, None]:\n docs = self._list_all_docs()\n logger.info(f\"Found {len(docs)} Coda docs to process\")\n\n for doc in docs:\n logger.debug(f\"Processing doc: {doc.name} ({doc.id})\")\n\n try:\n pages = self._list_pages_in_doc(doc.id)\n for page in pages:\n content = \"\"\n if self.index_page_content:\n try:\n content = self._fetch_page_content(doc.id, page.id)\n except Exception as e:\n logger.warning(\n f\"Failed to fetch content for page {page.id}: {e}\"\n )\n yield self._convert_page_to_document(page, content)\n except ConnectorValidationError as e:\n logger.warning(f\"Failed to list pages for doc {doc.id}: {e}\")\n\n try:\n tables = self._list_tables(doc.id)\n for table in tables:\n try:\n rows = self._list_rows_and_values(doc.id, table.id)\n yield self._convert_table_with_rows_to_document(table, rows)\n except ConnectorValidationError as e:\n logger.warning(\n f\"Failed to list rows for table {table.id}: {e}\"\n )\n yield self._convert_table_with_rows_to_document(table, [])\n except ConnectorValidationError as e:\n logger.warning(f\"Failed to list tables for doc {doc.id}: {e}\")\n\n return batch_generator(_iter_documents(), self.batch_size)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n \"\"\"\n Polls the Coda API for documents updated between start and end timestamps.\n We refer to page and table update times to determine if they need to be re-indexed.\n \"\"\"\n\n def _iter_documents() -> Generator[Document, None, None]:\n docs = self._list_all_docs()\n logger.info(\n f\"Polling {len(docs)} Coda docs for updates between {start} and {end}\"\n )\n\n for doc in docs:\n try:\n pages = self._list_pages_in_doc(doc.id)\n for page in pages:\n page_timestamp = (\n datetime.fromisoformat(page.updated_at)\n .astimezone(timezone.utc)\n .timestamp()\n )\n if start < page_timestamp <= end:\n content = \"\"\n if self.index_page_content:\n try:\n content = self._fetch_page_content(doc.id, page.id)\n except Exception as e:\n logger.warning(\n f\"Failed to fetch content for page {page.id}: {e}\"\n )\n yield self._convert_page_to_document(page, content)\n except ConnectorValidationError as e:\n logger.warning(f\"Failed to list pages for doc {doc.id}: {e}\")\n\n try:\n tables = self._list_tables(doc.id)\n for table in tables:\n table_timestamp = (\n datetime.fromisoformat(table.updated_at)\n .astimezone(timezone.utc)\n .timestamp()\n )\n\n try:\n rows = self._list_rows_and_values(doc.id, table.id)\n\n table_or_rows_updated = start < table_timestamp <= end\n if not table_or_rows_updated:\n for row in rows:\n row_timestamp = (\n datetime.fromisoformat(row.updated_at)\n .astimezone(timezone.utc)\n .timestamp()\n )\n if start < row_timestamp <= end:\n table_or_rows_updated = True\n break\n\n if table_or_rows_updated:\n yield self._convert_table_with_rows_to_document(\n table, rows\n )\n\n except ConnectorValidationError as e:\n logger.warning(\n f\"Failed to list rows for table {table.id}: {e}\"\n )\n if table_timestamp > start and table_timestamp <= end:\n yield self._convert_table_with_rows_to_document(\n table, []\n )\n\n except ConnectorValidationError as e:\n logger.warning(f\"Failed to list tables for doc {doc.id}: {e}\")\n\n return batch_generator(_iter_documents(), self.batch_size)\n\n def validate_connector_settings(self) -> None:\n \"\"\"Validates the Coda connector settings calling the 'whoami' endpoint.\"\"\"\n try:\n response = self.coda_client.get(\"whoami\")\n logger.info(\n f\"Coda connector validated for user: {response.get('name', 'Unknown')}\"\n )\n\n if self.workspace_id:\n params = {\"workspaceId\": self.workspace_id, \"limit\": \"1\"}\n self.coda_client.get(\"docs\", params=params)\n logger.info(f\"Validated access to workspace: {self.workspace_id}\")\n\n except CodaClientRequestFailedError as e:\n if e.status_code == 401:\n raise CredentialExpiredError(\n \"Coda credential appears to be invalid or expired (HTTP 401).\"\n )\n elif e.status_code == 404:\n raise ConnectorValidationError(\n \"Coda workspace not found or not accessible (HTTP 404). \"\n \"Please verify the workspace_id is correct and shared with the integration.\"\n )\n elif e.status_code == 429:\n raise ConnectorValidationError(\n \"Validation failed due to Coda rate-limits being exceeded (HTTP 429). Please try again later.\"\n )\n else:\n raise UnexpectedValidationError(\n f\"Unexpected Coda HTTP error (status={e.status_code}): {e}\"\n )\n except Exception as exc:\n raise UnexpectedValidationError(\n f\"Unexpected error during Coda settings validation: {exc}\"\n )\n" }, { "path": "backend/onyx/connectors/confluence/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/confluence/access.py", "content": "from collections.abc import Callable\nfrom typing import Any\nfrom typing import cast\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.confluence.onyx_confluence import OnyxConfluence\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\n\ndef get_page_restrictions(\n confluence_client: OnyxConfluence,\n page_id: str,\n page_restrictions: dict[str, Any],\n ancestors: list[dict[str, Any]],\n) -> ExternalAccess | None:\n \"\"\"\n Get page access restrictions for a Confluence page.\n This functionality requires Enterprise Edition.\n\n Note: This wrapper is only called from permission sync path. Group IDs are\n left unprefixed here because upsert_document_external_perms handles prefixing.\n\n Args:\n confluence_client: OnyxConfluence client instance\n page_id: The ID of the page\n page_restrictions: Dictionary containing page restriction data\n ancestors: List of ancestor pages with their restriction data\n\n Returns:\n ExternalAccess object for the page. None if EE is not enabled or no restrictions found.\n \"\"\"\n # Check if EE is enabled\n if not global_version.is_ee_version():\n return None\n\n # Fetch the EE implementation\n ee_get_all_page_restrictions = cast(\n Callable[\n [OnyxConfluence, str, dict[str, Any], list[dict[str, Any]], bool],\n ExternalAccess | None,\n ],\n fetch_versioned_implementation(\n \"onyx.external_permissions.confluence.page_access\", \"get_page_restrictions\"\n ),\n )\n\n # add_prefix=False: permission sync path - upsert_document_external_perms handles prefixing\n return ee_get_all_page_restrictions(\n confluence_client, page_id, page_restrictions, ancestors, False\n )\n\n\ndef get_all_space_permissions(\n confluence_client: OnyxConfluence,\n is_cloud: bool,\n) -> dict[str, ExternalAccess]:\n \"\"\"\n Get access permissions for all spaces in Confluence.\n This functionality requires Enterprise Edition.\n\n Note: This wrapper is only called from permission sync path. Group IDs are\n left unprefixed here because upsert_document_external_perms handles prefixing.\n\n Args:\n confluence_client: OnyxConfluence client instance\n is_cloud: Whether this is a Confluence Cloud instance\n\n Returns:\n Dictionary mapping space keys to ExternalAccess objects. Empty dict if EE is not enabled.\n \"\"\"\n # Check if EE is enabled\n if not global_version.is_ee_version():\n return {}\n\n # Fetch the EE implementation\n ee_get_all_space_permissions = cast(\n Callable[\n [OnyxConfluence, bool, bool],\n dict[str, ExternalAccess],\n ],\n fetch_versioned_implementation(\n \"onyx.external_permissions.confluence.space_access\",\n \"get_all_space_permissions\",\n ),\n )\n\n # add_prefix=False: permission sync path - upsert_document_external_perms handles prefixing\n return ee_get_all_space_permissions(confluence_client, is_cloud, False)\n" }, { "path": "backend/onyx/connectors/confluence/connector.py", "content": "import copy\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom urllib.parse import quote\n\nfrom atlassian.errors import ApiError # type: ignore\nfrom requests.exceptions import HTTPError\nfrom typing_extensions import override\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.app_configs import CONFLUENCE_CONNECTOR_LABELS_TO_SKIP\nfrom onyx.configs.app_configs import CONFLUENCE_TIMEZONE_OFFSET\nfrom onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.access import get_all_space_permissions\nfrom onyx.connectors.confluence.access import get_page_restrictions\nfrom onyx.connectors.confluence.onyx_confluence import extract_text_from_confluence_html\nfrom onyx.connectors.confluence.onyx_confluence import OnyxConfluence\nfrom onyx.connectors.confluence.utils import build_confluence_document_id\nfrom onyx.connectors.confluence.utils import convert_attachment_to_content\nfrom onyx.connectors.confluence.utils import datetime_from_string\nfrom onyx.connectors.confluence.utils import update_param_in_path\nfrom onyx.connectors.confluence.utils import validate_attachment_filetype\nfrom onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n is_atlassian_date_error,\n)\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import ConnectorCheckpoint\nfrom onyx.connectors.interfaces import ConnectorFailure\nfrom onyx.connectors.interfaces import CredentialsConnector\nfrom onyx.connectors.interfaces import CredentialsProviderInterface\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnector\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n# Potential Improvements\n# 1. Segment into Sections for more accurate linking, can split by headers but make sure no text/ordering is lost\n_COMMENT_EXPANSION_FIELDS = [\"body.storage.value\"]\n_PAGE_EXPANSION_FIELDS = [\n \"body.storage.value\",\n \"version\",\n \"space\",\n \"metadata.labels\",\n \"history.lastUpdated\",\n \"ancestors\", # For hierarchy node tracking\n]\n_ATTACHMENT_EXPANSION_FIELDS = [\n \"version\",\n \"space\",\n \"metadata.labels\",\n]\n_RESTRICTIONS_EXPANSION_FIELDS = [\n \"space\",\n \"restrictions.read.restrictions.user\",\n \"restrictions.read.restrictions.group\",\n \"ancestors.restrictions.read.restrictions.user\",\n \"ancestors.restrictions.read.restrictions.group\",\n]\n\n_SLIM_DOC_BATCH_SIZE = 5000\n\nONE_HOUR = 3600\nONE_DAY = ONE_HOUR * 24\n\nMAX_CACHED_IDS = 100\n\n\ndef _get_page_id(page: dict[str, Any], allow_missing: bool = False) -> str:\n if allow_missing and \"id\" not in page:\n return \"unknown\"\n return str(page[\"id\"])\n\n\nclass ConfluenceCheckpoint(ConnectorCheckpoint):\n next_page_url: str | None\n\n\nclass ConfluenceConnector(\n CheckpointedConnector[ConfluenceCheckpoint],\n SlimConnector,\n SlimConnectorWithPermSync,\n CredentialsConnector,\n):\n def __init__(\n self,\n wiki_base: str,\n is_cloud: bool,\n space: str = \"\",\n page_id: str = \"\",\n index_recursively: bool = False,\n cql_query: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n continue_on_failure: bool = CONTINUE_ON_CONNECTOR_FAILURE,\n # if a page has one of the labels specified in this list, we will just\n # skip it. This is generally used to avoid indexing extra sensitive\n # pages.\n labels_to_skip: list[str] = CONFLUENCE_CONNECTOR_LABELS_TO_SKIP,\n timezone_offset: float = CONFLUENCE_TIMEZONE_OFFSET,\n scoped_token: bool = False,\n ) -> None:\n self.wiki_base = wiki_base\n self.is_cloud = is_cloud\n self.space = space\n self.page_id = page_id\n self.index_recursively = index_recursively\n self.cql_query = cql_query\n self.batch_size = batch_size\n self.labels_to_skip = labels_to_skip\n self.timezone_offset = timezone_offset\n self.scoped_token = scoped_token\n self._confluence_client: OnyxConfluence | None = None\n self._low_timeout_confluence_client: OnyxConfluence | None = None\n self._fetched_titles: set[str] = set()\n self.allow_images = False\n\n # Track hierarchy nodes we've already yielded to avoid duplicates\n self.seen_hierarchy_node_raw_ids: set[str] = set()\n\n # Remove trailing slash from wiki_base if present\n self.wiki_base = wiki_base.rstrip(\"/\")\n \"\"\"\n If nothing is provided, we default to fetching all pages\n Only one or none of the following options should be specified so\n the order shouldn't matter\n However, we use elif to ensure that only of the following is enforced\n \"\"\"\n base_cql_page_query = \"type=page\"\n if cql_query:\n base_cql_page_query = cql_query\n elif page_id:\n if index_recursively:\n base_cql_page_query += f\" and (ancestor='{page_id}' or id='{page_id}')\"\n else:\n base_cql_page_query += f\" and id='{page_id}'\"\n elif space:\n uri_safe_space = quote(space)\n base_cql_page_query += f\" and space='{uri_safe_space}'\"\n\n self.base_cql_page_query = base_cql_page_query\n\n self.cql_label_filter = \"\"\n if labels_to_skip:\n labels_to_skip = list(set(labels_to_skip))\n comma_separated_labels = \",\".join(\n f\"'{quote(label)}'\" for label in labels_to_skip\n )\n self.cql_label_filter = f\" and label not in ({comma_separated_labels})\"\n\n self.timezone: timezone = timezone(offset=timedelta(hours=timezone_offset))\n self.credentials_provider: CredentialsProviderInterface | None = None\n\n self.probe_kwargs = {\n \"max_backoff_retries\": 6,\n \"max_backoff_seconds\": 10,\n }\n\n self.final_kwargs = {\n \"max_backoff_retries\": 10,\n \"max_backoff_seconds\": 60,\n }\n\n # deprecated\n self.continue_on_failure = continue_on_failure\n\n def set_allow_images(self, value: bool) -> None:\n logger.info(f\"Setting allow_images to {value}.\")\n self.allow_images = value\n\n def _yield_space_hierarchy_nodes(\n self,\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield hierarchy nodes for all spaces we're indexing.\"\"\"\n space_keys = [self.space] if self.space else None\n\n for space in self.confluence_client.retrieve_confluence_spaces(\n space_keys=space_keys,\n limit=50,\n ):\n space_key = space.get(\"key\")\n if not space_key or space_key in self.seen_hierarchy_node_raw_ids:\n continue\n\n self.seen_hierarchy_node_raw_ids.add(space_key)\n\n # Build space link\n space_link = f\"{self.wiki_base}/spaces/{space_key}\"\n\n yield HierarchyNode(\n raw_node_id=space_key,\n raw_parent_id=None, # Parent is SOURCE\n display_name=space.get(\"name\", space_key),\n link=space_link,\n node_type=HierarchyNodeType.SPACE,\n )\n\n def _yield_ancestor_hierarchy_nodes(\n self,\n page: dict[str, Any],\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield hierarchy nodes for all unseen ancestors of this page.\n\n Any page that appears as an ancestor of another page IS a hierarchy node\n (it has at least one child - the page we're currently processing).\n\n This ensures parent nodes are always yielded before child documents.\n\n Note: raw_node_id for page hierarchy nodes uses the page URL (same as document.id)\n to enable document<->hierarchy node linking in the indexing pipeline.\n Space hierarchy nodes use the space key since they don't have documents.\n \"\"\"\n ancestors = page.get(\"ancestors\", [])\n space_key = page.get(\"space\", {}).get(\"key\")\n\n # Ensure space is yielded first (if not already)\n if space_key and space_key not in self.seen_hierarchy_node_raw_ids:\n self.seen_hierarchy_node_raw_ids.add(space_key)\n space = page.get(\"space\", {})\n yield HierarchyNode(\n raw_node_id=space_key,\n raw_parent_id=None, # Parent is SOURCE\n display_name=space.get(\"name\", space_key),\n link=f\"{self.wiki_base}/spaces/{space_key}\",\n node_type=HierarchyNodeType.SPACE,\n )\n\n # Walk through ancestors (root to immediate parent)\n # Build a list of (ancestor_url, ancestor_data) pairs first\n ancestor_urls: list[str | None] = []\n for ancestor in ancestors:\n if \"_links\" in ancestor and \"webui\" in ancestor[\"_links\"]:\n ancestor_urls.append(\n build_confluence_document_id(\n self.wiki_base, ancestor[\"_links\"][\"webui\"], self.is_cloud\n )\n )\n else:\n ancestor_urls.append(None)\n\n for i, ancestor in enumerate(ancestors):\n ancestor_url = ancestor_urls[i]\n if not ancestor_url:\n # Can't build URL for this ancestor, skip it\n continue\n\n if ancestor_url in self.seen_hierarchy_node_raw_ids:\n continue\n\n self.seen_hierarchy_node_raw_ids.add(ancestor_url)\n\n # Determine parent of this ancestor\n if i == 0:\n # First ancestor - parent is the space\n parent_raw_id = space_key\n else:\n # Parent is the previous ancestor (use URL)\n parent_raw_id = ancestor_urls[i - 1]\n\n yield HierarchyNode(\n raw_node_id=ancestor_url, # Use URL to match document.id\n raw_parent_id=parent_raw_id,\n display_name=ancestor.get(\"title\", f\"Page {ancestor.get('id')}\"),\n link=ancestor_url,\n node_type=HierarchyNodeType.PAGE,\n )\n\n def _get_parent_hierarchy_raw_id(self, page: dict[str, Any]) -> str | None:\n \"\"\"Get the raw hierarchy node ID of this page's parent.\n\n Returns:\n - Parent page URL if page has a parent page (last item in ancestors)\n - Space key if page is at top level of space\n - None if we can't determine\n\n Note: For pages, we return URLs (to match document.id and hierarchy node raw_node_id).\n For spaces, we return the space key (spaces don't have documents).\n \"\"\"\n ancestors = page.get(\"ancestors\", [])\n if ancestors:\n # Last ancestor is the immediate parent page - use URL\n parent = ancestors[-1]\n if \"_links\" in parent and \"webui\" in parent[\"_links\"]:\n return build_confluence_document_id(\n self.wiki_base, parent[\"_links\"][\"webui\"], self.is_cloud\n )\n # Fallback to page ID if URL not available (shouldn't happen normally)\n return str(parent.get(\"id\"))\n\n # Top-level page - parent is the space (use space key)\n return page.get(\"space\", {}).get(\"key\")\n\n def _maybe_yield_page_hierarchy_node(\n self, page: dict[str, Any]\n ) -> HierarchyNode | None:\n \"\"\"Yield a hierarchy node for this page if not already yielded.\n\n Used when a page has attachments - attachments are children of the page\n in the hierarchy, so the page must be a hierarchy node.\n\n Note: raw_node_id uses the page URL (same as document.id) to enable\n document<->hierarchy node linking in the indexing pipeline.\n \"\"\"\n # Build page URL - we use this as raw_node_id to match document.id\n if \"_links\" not in page or \"webui\" not in page[\"_links\"]:\n return None # Can't build URL, skip\n\n page_url = build_confluence_document_id(\n self.wiki_base, page[\"_links\"][\"webui\"], self.is_cloud\n )\n\n if page_url in self.seen_hierarchy_node_raw_ids:\n return None\n\n self.seen_hierarchy_node_raw_ids.add(page_url)\n\n # Get parent hierarchy ID\n parent_raw_id = self._get_parent_hierarchy_raw_id(page)\n\n return HierarchyNode(\n raw_node_id=page_url, # Use URL to match document.id\n raw_parent_id=parent_raw_id,\n display_name=page.get(\"title\", f\"Page {_get_page_id(page)}\"),\n link=page_url,\n node_type=HierarchyNodeType.PAGE,\n )\n\n @property\n def confluence_client(self) -> OnyxConfluence:\n if self._confluence_client is None:\n raise ConnectorMissingCredentialError(\"Confluence\")\n return self._confluence_client\n\n @property\n def low_timeout_confluence_client(self) -> OnyxConfluence:\n if self._low_timeout_confluence_client is None:\n raise ConnectorMissingCredentialError(\"Confluence\")\n return self._low_timeout_confluence_client\n\n def set_credentials_provider(\n self, credentials_provider: CredentialsProviderInterface\n ) -> None:\n self.credentials_provider = credentials_provider\n\n # raises exception if there's a problem\n confluence_client = OnyxConfluence(\n is_cloud=self.is_cloud,\n url=self.wiki_base,\n credentials_provider=credentials_provider,\n scoped_token=self.scoped_token,\n )\n confluence_client._probe_connection(**self.probe_kwargs)\n confluence_client._initialize_connection(**self.final_kwargs)\n\n self._confluence_client = confluence_client\n\n # create a low timeout confluence client for sync flows\n low_timeout_confluence_client = OnyxConfluence(\n is_cloud=self.is_cloud,\n url=self.wiki_base,\n credentials_provider=credentials_provider,\n timeout=3,\n scoped_token=self.scoped_token,\n )\n low_timeout_confluence_client._probe_connection(**self.probe_kwargs)\n low_timeout_confluence_client._initialize_connection(**self.final_kwargs)\n\n self._low_timeout_confluence_client = low_timeout_confluence_client\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n raise NotImplementedError(\"Use set_credentials_provider with this connector.\")\n\n def _construct_page_cql_query(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> str:\n \"\"\"\n Constructs a CQL query for use in the confluence API. See\n https://developer.atlassian.com/server/confluence/advanced-searching-using-cql/\n for more information. This is JUST the CQL, not the full URL used to hit the API.\n Use _build_page_retrieval_url to get the full URL.\n \"\"\"\n page_query = self.base_cql_page_query + self.cql_label_filter\n # Add time filters\n if start:\n formatted_start_time = datetime.fromtimestamp(\n start, tz=self.timezone\n ).strftime(\"%Y-%m-%d %H:%M\")\n page_query += f\" and lastmodified >= '{formatted_start_time}'\"\n if end:\n formatted_end_time = datetime.fromtimestamp(end, tz=self.timezone).strftime(\n \"%Y-%m-%d %H:%M\"\n )\n page_query += f\" and lastmodified <= '{formatted_end_time}'\"\n\n page_query += \" order by lastmodified asc\"\n return page_query\n\n def _construct_attachment_query(\n self,\n confluence_page_id: str,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> str:\n attachment_query = f\"type=attachment and container='{confluence_page_id}'\"\n attachment_query += self.cql_label_filter\n # Add time filters to avoid reprocessing unchanged attachments during refresh\n if start:\n formatted_start_time = datetime.fromtimestamp(\n start, tz=self.timezone\n ).strftime(\"%Y-%m-%d %H:%M\")\n attachment_query += f\" and lastmodified >= '{formatted_start_time}'\"\n if end:\n formatted_end_time = datetime.fromtimestamp(end, tz=self.timezone).strftime(\n \"%Y-%m-%d %H:%M\"\n )\n attachment_query += f\" and lastmodified <= '{formatted_end_time}'\"\n attachment_query += \" order by lastmodified asc\"\n return attachment_query\n\n def _get_comment_string_for_page_id(self, page_id: str) -> str:\n comment_string = \"\"\n comment_cql = f\"type=comment and container='{page_id}'\"\n comment_cql += self.cql_label_filter\n expand = \",\".join(_COMMENT_EXPANSION_FIELDS)\n\n for comment in self.confluence_client.paginated_cql_retrieval(\n cql=comment_cql,\n expand=expand,\n ):\n comment_string += \"\\nComment:\\n\"\n comment_string += extract_text_from_confluence_html(\n confluence_client=self.confluence_client,\n confluence_object=comment,\n fetched_titles=set(),\n )\n return comment_string\n\n def _convert_page_to_document(\n self, page: dict[str, Any]\n ) -> Document | ConnectorFailure:\n \"\"\"\n Converts a Confluence page to a Document object.\n Includes the page content, comments, and attachments.\n \"\"\"\n page_id = page_url = \"\"\n try:\n # Extract basic page information\n page_id = _get_page_id(page)\n page_title = page[\"title\"]\n logger.info(f\"Converting page {page_title} to document\")\n page_url = build_confluence_document_id(\n self.wiki_base, page[\"_links\"][\"webui\"], self.is_cloud\n )\n\n # Get the page content\n page_content = extract_text_from_confluence_html(\n self.confluence_client, page, self._fetched_titles\n )\n\n # Create the main section for the page content\n sections: list[TextSection | ImageSection] = [\n TextSection(text=page_content, link=page_url)\n ]\n\n # Process comments if available\n comment_text = self._get_comment_string_for_page_id(page_id)\n if comment_text:\n sections.append(\n TextSection(text=comment_text, link=f\"{page_url}#comments\")\n )\n # Note: attachments are no longer merged into the page document.\n # They are indexed as separate documents downstream.\n\n # Extract metadata\n metadata = {}\n if \"space\" in page:\n metadata[\"space\"] = page[\"space\"].get(\"name\", \"\")\n\n # Extract labels\n labels = []\n if \"metadata\" in page and \"labels\" in page[\"metadata\"]:\n for label in page[\"metadata\"][\"labels\"].get(\"results\", []):\n labels.append(label.get(\"name\", \"\"))\n if labels:\n metadata[\"labels\"] = labels\n\n # Extract owners\n primary_owners = []\n if \"version\" in page and \"by\" in page[\"version\"]:\n author = page[\"version\"][\"by\"]\n display_name = author.get(\"displayName\", \"Unknown\")\n email = author.get(\"email\", \"unknown@domain.invalid\")\n primary_owners.append(\n BasicExpertInfo(display_name=display_name, email=email)\n )\n\n # Determine parent hierarchy node\n parent_hierarchy_raw_node_id = self._get_parent_hierarchy_raw_id(page)\n\n # Create the document\n return Document(\n id=page_url,\n sections=sections,\n source=DocumentSource.CONFLUENCE,\n semantic_identifier=page_title,\n metadata=metadata,\n doc_updated_at=datetime_from_string(page[\"version\"][\"when\"]),\n primary_owners=primary_owners if primary_owners else None,\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n )\n except Exception as e:\n logger.error(f\"Error converting page {page.get('id', 'unknown')}: {e}\")\n if is_atlassian_date_error(e): # propagate error to be caught and retried\n raise\n return ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=page_id,\n document_link=page_url,\n ),\n failure_message=f\"Error converting page {page.get('id', 'unknown')}: {e}\",\n exception=e,\n )\n\n def _fetch_page_attachments(\n self,\n page: dict[str, Any],\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> tuple[list[Document | HierarchyNode], list[ConnectorFailure]]:\n \"\"\"\n Inline attachments are added directly to the document as text or image sections by\n this function. The returned documents/connectorfailures are for non-inline attachments\n and those at the end of the page.\n\n If there are valid attachments, the page itself is yielded as a hierarchy node\n (since attachments are children of the page in the hierarchy).\n \"\"\"\n attachment_query = self._construct_attachment_query(\n _get_page_id(page), start, end\n )\n attachment_failures: list[ConnectorFailure] = []\n attachment_docs: list[Document | HierarchyNode] = []\n page_url = \"\"\n page_hierarchy_node_yielded = False\n\n try:\n for attachment in self.confluence_client.paginated_cql_retrieval(\n cql=attachment_query,\n expand=\",\".join(_ATTACHMENT_EXPANSION_FIELDS),\n ):\n media_type: str = attachment.get(\"metadata\", {}).get(\"mediaType\", \"\")\n\n # TODO(rkuo): this check is partially redundant with validate_attachment_filetype\n # and checks in convert_attachment_to_content/process_attachment\n # but doing the check here avoids an unnecessary download. Due for refactoring.\n if not self.allow_images:\n if media_type.startswith(\"image/\"):\n logger.info(\n f\"Skipping attachment because allow images is False: {attachment['title']}\"\n )\n continue\n\n if not validate_attachment_filetype(\n attachment,\n ):\n logger.info(\n f\"Skipping attachment because it is not an accepted file type: {attachment['title']}\"\n )\n continue\n\n logger.info(\n f\"Processing attachment: {attachment['title']} attached to page {page['title']}\"\n )\n # Attachment document id: use the download URL for stable identity\n try:\n object_url = build_confluence_document_id(\n self.wiki_base, attachment[\"_links\"][\"download\"], self.is_cloud\n )\n except Exception as e:\n logger.warning(\n f\"Invalid attachment url for id {attachment['id']}, skipping\"\n )\n logger.debug(f\"Error building attachment url: {e}\")\n continue\n try:\n response = convert_attachment_to_content(\n confluence_client=self.confluence_client,\n attachment=attachment,\n page_id=_get_page_id(page),\n allow_images=self.allow_images,\n )\n if response is None:\n continue\n\n content_text, file_storage_name = response\n\n sections: list[TextSection | ImageSection] = []\n if content_text:\n sections.append(TextSection(text=content_text, link=object_url))\n elif file_storage_name:\n sections.append(\n ImageSection(\n link=object_url, image_file_id=file_storage_name\n )\n )\n\n # Build attachment-specific metadata\n attachment_metadata: dict[str, str | list[str]] = {}\n if \"space\" in attachment:\n attachment_metadata[\"space\"] = attachment[\"space\"].get(\n \"name\", \"\"\n )\n labels: list[str] = []\n if \"metadata\" in attachment and \"labels\" in attachment[\"metadata\"]:\n for label in attachment[\"metadata\"][\"labels\"].get(\n \"results\", []\n ):\n labels.append(label.get(\"name\", \"\"))\n if labels:\n attachment_metadata[\"labels\"] = labels\n page_url = page_url or build_confluence_document_id(\n self.wiki_base, page[\"_links\"][\"webui\"], self.is_cloud\n )\n attachment_metadata[\"parent_page_id\"] = page_url\n attachment_id = build_confluence_document_id(\n self.wiki_base, attachment[\"_links\"][\"webui\"], self.is_cloud\n )\n\n primary_owners: list[BasicExpertInfo] | None = None\n if \"version\" in attachment and \"by\" in attachment[\"version\"]:\n author = attachment[\"version\"][\"by\"]\n display_name = author.get(\"displayName\", \"Unknown\")\n email = author.get(\"email\", \"unknown@domain.invalid\")\n primary_owners = [\n BasicExpertInfo(display_name=display_name, email=email)\n ]\n\n # Attachments have their parent page as the hierarchy parent\n # Use page URL to match the hierarchy node's raw_node_id\n attachment_parent_hierarchy_raw_id = page_url\n\n attachment_doc = Document(\n id=attachment_id,\n sections=sections,\n source=DocumentSource.CONFLUENCE,\n semantic_identifier=attachment.get(\"title\", object_url),\n metadata=attachment_metadata,\n doc_updated_at=(\n datetime_from_string(attachment[\"version\"][\"when\"])\n if attachment.get(\"version\")\n and attachment[\"version\"].get(\"when\")\n else None\n ),\n primary_owners=primary_owners,\n parent_hierarchy_raw_node_id=attachment_parent_hierarchy_raw_id,\n )\n\n # If this is the first valid attachment, yield the page as a\n # hierarchy node (attachments are children of the page)\n if not page_hierarchy_node_yielded:\n page_hierarchy_node = self._maybe_yield_page_hierarchy_node(\n page\n )\n if page_hierarchy_node:\n attachment_docs.append(page_hierarchy_node)\n page_hierarchy_node_yielded = True\n\n attachment_docs.append(attachment_doc)\n except Exception as e:\n logger.error(\n f\"Failed to extract/summarize attachment {attachment['title']}\",\n exc_info=e,\n )\n if is_atlassian_date_error(e):\n # propagate error to be caught and retried\n raise\n attachment_failures.append(\n ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=object_url,\n document_link=object_url,\n ),\n failure_message=f\"Failed to extract/summarize attachment {attachment['title']} for doc {object_url}\",\n exception=e,\n )\n )\n except HTTPError as e:\n # If we get a 403 after all retries, the user likely doesn't have permission\n # to access attachments on this page. Log and skip rather than failing the whole job.\n page_id = _get_page_id(page, allow_missing=True)\n page_title = page.get(\"title\", \"unknown\")\n if e.response and e.response.status_code in [401, 403]:\n failure_message_prefix = (\n \"Invalid credentials (401)\"\n if e.response.status_code == 401\n else \"Permission denied (403)\"\n )\n failure_message = (\n f\"{failure_message_prefix} when fetching attachments for page '{page_title}' \"\n f\"(ID: {page_id}). The user may not have permission to query attachments on this page. \"\n \"Skipping attachments for this page.\"\n )\n logger.warning(failure_message)\n\n # Build the page URL for the failure record\n try:\n page_url = build_confluence_document_id(\n self.wiki_base, page[\"_links\"][\"webui\"], self.is_cloud\n )\n except Exception:\n page_url = f\"page_id:{page_id}\"\n\n return [], [\n ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=page_id,\n document_link=page_url,\n ),\n failure_message=failure_message,\n exception=e,\n )\n ]\n else:\n raise\n\n return attachment_docs, attachment_failures\n\n def _fetch_document_batches(\n self,\n checkpoint: ConfluenceCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> CheckpointOutput[ConfluenceCheckpoint]:\n \"\"\"\n Yields batches of Documents and HierarchyNodes. For each page:\n - Yield hierarchy nodes for spaces and ancestor pages (parent-before-child ordering)\n - Create a Document with 1 Section for the page text/comments\n - Then fetch attachments. For each attachment:\n - Attempt to convert it with convert_attachment_to_content(...)\n - If successful, create a new Section with the extracted text or summary.\n \"\"\"\n checkpoint = copy.deepcopy(checkpoint)\n\n # Yield space hierarchy nodes FIRST (only once per connector run)\n if not checkpoint.next_page_url:\n yield from self._yield_space_hierarchy_nodes()\n\n # use \"start\" when last_updated is 0 or for confluence server\n start_ts = start\n page_query_url = checkpoint.next_page_url or self._build_page_retrieval_url(\n start_ts, end, self.batch_size\n )\n logger.debug(f\"page_query_url: {page_query_url}\")\n\n # store the next page start for confluence server, cursor for confluence cloud\n def store_next_page_url(next_page_url: str) -> None:\n checkpoint.next_page_url = next_page_url\n\n for page in self.confluence_client.paginated_page_retrieval(\n cql_url=page_query_url,\n limit=self.batch_size,\n next_page_callback=store_next_page_url,\n ):\n # Yield hierarchy nodes for all ancestors (parent-before-child ordering)\n yield from self._yield_ancestor_hierarchy_nodes(page)\n\n # Build doc from page\n doc_or_failure = self._convert_page_to_document(page)\n\n if isinstance(doc_or_failure, ConnectorFailure):\n yield doc_or_failure\n continue\n\n # yield completed document (or failure)\n yield doc_or_failure\n\n # Now get attachments for that page:\n attachment_docs, attachment_failures = self._fetch_page_attachments(\n page, start, end\n )\n # yield attached docs and failures\n yield from attachment_docs\n yield from attachment_failures\n\n # Create checkpoint once a full page of results is returned\n if checkpoint.next_page_url and checkpoint.next_page_url != page_query_url:\n return checkpoint\n\n checkpoint.has_more = False\n return checkpoint\n\n def _build_page_retrieval_url(\n self,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n limit: int,\n ) -> str:\n \"\"\"\n Builds the full URL used to retrieve pages from the confluence API.\n This can be used as input to the confluence client's _paginate_url\n or paginated_page_retrieval methods.\n \"\"\"\n page_query = self._construct_page_cql_query(start, end)\n cql_url = self.confluence_client.build_cql_url(\n page_query, expand=\",\".join(_PAGE_EXPANSION_FIELDS)\n )\n return update_param_in_path(cql_url, \"limit\", str(limit))\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: ConfluenceCheckpoint,\n ) -> CheckpointOutput[ConfluenceCheckpoint]:\n end += ONE_DAY # handle time zone weirdness\n try:\n return self._fetch_document_batches(checkpoint, start, end)\n except Exception as e:\n if is_atlassian_date_error(e) and start is not None:\n logger.warning(\n \"Confluence says we provided an invalid 'updated' field. This may indicate\"\n \"a real issue, but can also appear during edge cases like daylight\"\n f\"savings time changes. Retrying with a 1 hour offset. Error: {e}\"\n )\n return self._fetch_document_batches(checkpoint, start - ONE_HOUR, end)\n raise\n\n @override\n def build_dummy_checkpoint(self) -> ConfluenceCheckpoint:\n return ConfluenceCheckpoint(has_more=True, next_page_url=None)\n\n @override\n def validate_checkpoint_json(self, checkpoint_json: str) -> ConfluenceCheckpoint:\n return ConfluenceCheckpoint.model_validate_json(checkpoint_json)\n\n @override\n def retrieve_all_slim_docs(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n return self._retrieve_all_slim_docs(\n start=start,\n end=end,\n callback=callback,\n include_permissions=False,\n )\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n \"\"\"\n Return 'slim' docs (IDs + minimal permission data).\n Does not fetch actual text. Used primarily for incremental permission sync.\n \"\"\"\n return self._retrieve_all_slim_docs(\n start=start,\n end=end,\n callback=callback,\n include_permissions=True,\n )\n\n def _retrieve_all_slim_docs(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n include_permissions: bool = True,\n ) -> GenerateSlimDocumentOutput:\n doc_metadata_list: list[SlimDocument | HierarchyNode] = []\n restrictions_expand = \",\".join(_RESTRICTIONS_EXPANSION_FIELDS)\n\n space_level_access_info: dict[str, ExternalAccess] = {}\n if include_permissions:\n space_level_access_info = get_all_space_permissions(\n self.confluence_client, self.is_cloud\n )\n\n # Yield space hierarchy nodes first\n for node in self._yield_space_hierarchy_nodes():\n doc_metadata_list.append(node)\n\n def get_external_access(\n doc_id: str, restrictions: dict[str, Any], ancestors: list[dict[str, Any]]\n ) -> ExternalAccess | None:\n return get_page_restrictions(\n self.confluence_client, doc_id, restrictions, ancestors\n ) or space_level_access_info.get(page_space_key)\n\n # Query pages (with optional time filtering for indexing_start)\n page_query = self._construct_page_cql_query(start, end)\n for page in self.confluence_client.cql_paginate_all_expansions(\n cql=page_query,\n expand=restrictions_expand,\n limit=_SLIM_DOC_BATCH_SIZE,\n ):\n # Yield ancestor hierarchy nodes for this page\n for node in self._yield_ancestor_hierarchy_nodes(page):\n doc_metadata_list.append(node)\n\n page_id = _get_page_id(page)\n page_restrictions = page.get(\"restrictions\") or {}\n page_space_key = page.get(\"space\", {}).get(\"key\")\n page_ancestors = page.get(\"ancestors\", [])\n\n page_id = build_confluence_document_id(\n self.wiki_base, page[\"_links\"][\"webui\"], self.is_cloud\n )\n doc_metadata_list.append(\n SlimDocument(\n id=page_id,\n external_access=(\n get_external_access(page_id, page_restrictions, page_ancestors)\n if include_permissions\n else None\n ),\n parent_hierarchy_raw_node_id=self._get_parent_hierarchy_raw_id(\n page\n ),\n )\n )\n\n # Query attachments for each page\n page_hierarchy_node_yielded = False\n attachment_query = self._construct_attachment_query(\n _get_page_id(page), start, end\n )\n for attachment in self.confluence_client.cql_paginate_all_expansions(\n cql=attachment_query,\n expand=restrictions_expand,\n limit=_SLIM_DOC_BATCH_SIZE,\n ):\n # If you skip images, you'll skip them in the permission sync\n attachment[\"metadata\"].get(\"mediaType\", \"\")\n if not validate_attachment_filetype(\n attachment,\n ):\n continue\n\n # If this page has valid attachments and we haven't yielded it as a\n # hierarchy node yet, do so now (attachments are children of the page)\n if not page_hierarchy_node_yielded:\n page_node = self._maybe_yield_page_hierarchy_node(page)\n if page_node:\n doc_metadata_list.append(page_node)\n page_hierarchy_node_yielded = True\n\n attachment_restrictions = attachment.get(\"restrictions\", {})\n if not attachment_restrictions:\n attachment_restrictions = page_restrictions or {}\n\n attachment_space_key = attachment.get(\"space\", {}).get(\"key\")\n if not attachment_space_key:\n attachment_space_key = page_space_key\n\n attachment_id = build_confluence_document_id(\n self.wiki_base,\n attachment[\"_links\"][\"webui\"],\n self.is_cloud,\n )\n doc_metadata_list.append(\n SlimDocument(\n id=attachment_id,\n external_access=(\n get_external_access(\n attachment_id, attachment_restrictions, []\n )\n if include_permissions\n else None\n ),\n parent_hierarchy_raw_node_id=page_id,\n )\n )\n\n if len(doc_metadata_list) > _SLIM_DOC_BATCH_SIZE:\n yield doc_metadata_list[:_SLIM_DOC_BATCH_SIZE]\n doc_metadata_list = doc_metadata_list[_SLIM_DOC_BATCH_SIZE:]\n\n if callback and callback.should_stop():\n raise RuntimeError(\n \"retrieve_all_slim_docs_perm_sync: Stop signal detected\"\n )\n if callback:\n callback.progress(\"retrieve_all_slim_docs_perm_sync\", 1)\n\n yield doc_metadata_list\n\n def validate_connector_settings(self) -> None:\n try:\n spaces_iter = self.low_timeout_confluence_client.retrieve_confluence_spaces(\n limit=1,\n )\n first_space = next(spaces_iter, None)\n except HTTPError as e:\n status_code = e.response.status_code if e.response else None\n if status_code == 401:\n raise CredentialExpiredError(\n \"Invalid or expired Confluence credentials (HTTP 401).\"\n )\n elif status_code == 403:\n raise InsufficientPermissionsError(\n \"Insufficient permissions to access Confluence resources (HTTP 403).\"\n )\n raise UnexpectedValidationError(\n f\"Unexpected Confluence error (status={status_code}): {e}\"\n )\n except Exception as e:\n raise UnexpectedValidationError(\n f\"Unexpected error while validating Confluence settings: {e}\"\n )\n\n if not first_space:\n raise ConnectorValidationError(\n \"No Confluence spaces found. Either your credentials lack permissions, or \"\n \"there truly are no spaces in this Confluence instance.\"\n )\n\n if self.space:\n try:\n self.low_timeout_confluence_client.get_space(self.space)\n except ApiError as e:\n raise ConnectorValidationError(\n \"Invalid Confluence space key provided\"\n ) from e\n\n\nif __name__ == \"__main__\":\n import os\n from onyx.utils.variable_functionality import global_version\n from tests.daily.connectors.utils import load_all_from_connector\n\n # For connector permission testing, set EE to true.\n global_version.set_ee()\n\n # base url\n wiki_base = os.environ[\"CONFLUENCE_URL\"]\n\n # auth stuff\n username = os.environ[\"CONFLUENCE_USERNAME\"]\n access_token = os.environ[\"CONFLUENCE_ACCESS_TOKEN\"]\n is_cloud = os.environ[\"CONFLUENCE_IS_CLOUD\"].lower() == \"true\"\n\n # space + page\n space = os.environ[\"CONFLUENCE_SPACE_KEY\"]\n # page_id = os.environ[\"CONFLUENCE_PAGE_ID\"]\n\n confluence_connector = ConfluenceConnector(\n wiki_base=wiki_base,\n space=space,\n is_cloud=is_cloud,\n # page_id=page_id,\n )\n\n credentials_provider = OnyxStaticCredentialsProvider(\n None,\n DocumentSource.CONFLUENCE,\n {\n \"confluence_username\": username,\n \"confluence_access_token\": access_token,\n },\n )\n confluence_connector.set_credentials_provider(credentials_provider)\n\n start = 0.0\n end = datetime.now().timestamp()\n\n # Fetch all `SlimDocuments`.\n for slim_doc in confluence_connector.retrieve_all_slim_docs_perm_sync():\n print(slim_doc)\n\n # Fetch all `Documents`.\n for doc in load_all_from_connector(\n connector=confluence_connector,\n start=start,\n end=end,\n ).documents:\n print(doc)\n" }, { "path": "backend/onyx/connectors/confluence/models.py", "content": "from pydantic import BaseModel\n\n\nclass ConfluenceUser(BaseModel):\n user_id: str # accountId in Cloud, userKey in Server\n username: str | None # Confluence Cloud doesn't give usernames\n display_name: str\n # Confluence Data Center doesn't give email back by default,\n # have to fetch it with a different endpoint\n email: str | None\n type: str\n" }, { "path": "backend/onyx/connectors/confluence/onyx_confluence.py", "content": "\"\"\"\n# README (notes on Confluence pagination):\n\nWe've noticed that the `search/users` and `users/memberof` endpoints for Confluence Cloud use offset-based pagination as\nopposed to cursor-based. We also know that page-retrieval uses cursor-based pagination.\n\nOur default pagination strategy right now for cloud is to assume cursor-based.\nHowever, if you notice that a cloud API is not being properly paginated (i.e., if the `_links.next` is not appearing in the\nreturned payload), then you can force offset-based pagination.\n\n# TODO (@raunakab)\nWe haven't explored all of the cloud APIs' pagination strategies. @raunakab take time to go through this and figure them out.\n\"\"\"\n\nimport json\nimport time\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import TypeVar\nfrom urllib.parse import quote\n\nimport bs4\nfrom atlassian import Confluence # type:ignore\nfrom redis import Redis\nfrom requests import HTTPError\n\nfrom onyx.configs.app_configs import CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE\nfrom onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET\nfrom onyx.connectors.confluence.models import ConfluenceUser\nfrom onyx.connectors.confluence.user_profile_override import (\n process_confluence_user_profiles_override,\n)\nfrom onyx.connectors.confluence.utils import _handle_http_error\nfrom onyx.connectors.confluence.utils import confluence_refresh_tokens\nfrom onyx.connectors.confluence.utils import get_start_param_from_url\nfrom onyx.connectors.confluence.utils import update_param_in_path\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import scoped_url\nfrom onyx.connectors.interfaces import CredentialsProviderInterface\nfrom onyx.file_processing.html_utils import format_document_soup\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nF = TypeVar(\"F\", bound=Callable[..., Any])\n\n\n# https://jira.atlassian.com/browse/CONFCLOUD-76433\n_PROBLEMATIC_EXPANSIONS = \"body.storage.value\"\n_REPLACEMENT_EXPANSIONS = \"body.view.value\"\n\n_USER_NOT_FOUND = \"Unknown Confluence User\"\n_USER_ID_TO_DISPLAY_NAME_CACHE: dict[str, str | None] = {}\n_USER_EMAIL_CACHE: dict[str, str | None] = {}\n_DEFAULT_PAGINATION_LIMIT = 1000\n\n_CONFLUENCE_SPACES_API_V1 = \"rest/api/space\"\n_CONFLUENCE_SPACES_API_V2 = \"wiki/api/v2/spaces\"\n\n\nclass ConfluenceRateLimitError(Exception):\n pass\n\n\nclass OnyxConfluence:\n \"\"\"\n This is a custom Confluence class that:\n\n A. overrides the default Confluence class to add a custom CQL method.\n B.\n This is necessary because the default Confluence class does not properly support cql expansions.\n All methods are automatically wrapped with handle_confluence_rate_limit.\n \"\"\"\n\n CREDENTIAL_PREFIX = \"connector:confluence:credential\"\n CREDENTIAL_TTL = 300 # 5 min\n PROBE_TIMEOUT = 5 # 5 seconds\n\n def __init__(\n self,\n is_cloud: bool,\n url: str,\n credentials_provider: CredentialsProviderInterface,\n timeout: int | None = None,\n scoped_token: bool = False,\n # should generally not be passed in, but making it overridable for\n # easier testing\n confluence_user_profiles_override: list[dict[str, str]] | None = (\n CONFLUENCE_CONNECTOR_USER_PROFILES_OVERRIDE\n ),\n ) -> None:\n self.base_url = url #'/'.join(url.rstrip(\"/\").split(\"/\")[:-1])\n url = scoped_url(url, \"confluence\") if scoped_token else url\n\n self._is_cloud = is_cloud\n self._url = url.rstrip(\"/\")\n self._credentials_provider = credentials_provider\n self.scoped_token = scoped_token\n self.redis_client: Redis | None = None\n self.static_credentials: dict[str, Any] | None = None\n if self._credentials_provider.is_dynamic():\n self.redis_client = get_redis_client(\n tenant_id=credentials_provider.get_tenant_id()\n )\n else:\n self.static_credentials = self._credentials_provider.get_credentials()\n\n self._confluence = Confluence(url)\n self.credential_key: str = (\n self.CREDENTIAL_PREFIX\n + f\":credential_{self._credentials_provider.get_provider_key()}\"\n )\n\n self._kwargs: Any = None\n\n self.shared_base_kwargs: dict[str, str | int | bool] = {\n \"api_version\": \"cloud\" if is_cloud else \"latest\",\n \"backoff_and_retry\": False,\n \"cloud\": is_cloud,\n }\n if timeout:\n self.shared_base_kwargs[\"timeout\"] = timeout\n\n self._confluence_user_profiles_override = (\n process_confluence_user_profiles_override(confluence_user_profiles_override)\n if confluence_user_profiles_override\n else None\n )\n\n def _renew_credentials(self) -> tuple[dict[str, Any], bool]:\n \"\"\"credential_json - the current json credentials\n Returns a tuple\n 1. The up to date credentials\n 2. True if the credentials were updated\n\n This method is intended to be used within a distributed lock.\n Lock, call this, update credentials if the tokens were refreshed, then release\n \"\"\"\n # static credentials are preloaded, so no locking/redis required\n if self.static_credentials:\n return self.static_credentials, False\n\n if not self.redis_client:\n raise RuntimeError(\"self.redis_client is None\")\n\n # dynamic credentials need locking\n # check redis first, then fallback to the DB\n credential_raw = self.redis_client.get(self.credential_key)\n if credential_raw is not None:\n credential_bytes = cast(bytes, credential_raw)\n credential_str = credential_bytes.decode(\"utf-8\")\n credential_json: dict[str, Any] = json.loads(credential_str)\n else:\n credential_json = self._credentials_provider.get_credentials()\n\n if \"confluence_refresh_token\" not in credential_json:\n # static credentials ... cache them permanently and return\n self.static_credentials = credential_json\n return credential_json, False\n\n if not OAUTH_CONFLUENCE_CLOUD_CLIENT_ID:\n raise RuntimeError(\"OAUTH_CONFLUENCE_CLOUD_CLIENT_ID must be set!\")\n\n if not OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET:\n raise RuntimeError(\"OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET must be set!\")\n\n # check if we should refresh tokens. we're deciding to refresh halfway\n # to expiration\n now = datetime.now(timezone.utc)\n created_at = datetime.fromisoformat(credential_json[\"created_at\"])\n expires_in: int = credential_json[\"expires_in\"]\n renew_at = created_at + timedelta(seconds=expires_in // 2)\n if now <= renew_at:\n # cached/current credentials are reasonably up to date\n return credential_json, False\n\n # we need to refresh\n logger.info(\"Renewing Confluence Cloud credentials...\")\n new_credentials = confluence_refresh_tokens(\n OAUTH_CONFLUENCE_CLOUD_CLIENT_ID,\n OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET,\n credential_json[\"cloud_id\"],\n credential_json[\"confluence_refresh_token\"],\n )\n\n # store the new credentials to redis and to the db thru the provider\n # redis: we use a 5 min TTL because we are given a 10 minute grace period\n # when keys are rotated. it's easier to expire the cached credentials\n # reasonably frequently rather than trying to handle strong synchronization\n # between the db and redis everywhere the credentials might be updated\n new_credential_str = json.dumps(new_credentials)\n self.redis_client.set(\n self.credential_key, new_credential_str, nx=True, ex=self.CREDENTIAL_TTL\n )\n self._credentials_provider.set_credentials(new_credentials)\n\n return new_credentials, True\n\n @staticmethod\n def _make_oauth2_dict(credentials: dict[str, Any]) -> dict[str, Any]:\n oauth2_dict: dict[str, Any] = {}\n if \"confluence_refresh_token\" in credentials:\n oauth2_dict[\"client_id\"] = OAUTH_CONFLUENCE_CLOUD_CLIENT_ID\n oauth2_dict[\"token\"] = {}\n oauth2_dict[\"token\"][\"access_token\"] = credentials[\n \"confluence_access_token\"\n ]\n return oauth2_dict\n\n def _build_spaces_url(\n self,\n is_v2: bool,\n base_url: str,\n limit: int,\n space_keys: list[str] | None,\n start: int | None = None,\n ) -> str:\n \"\"\"Build URL for Confluence spaces API with query parameters.\"\"\"\n key_param = \"keys\" if is_v2 else \"spaceKey\"\n\n params = [f\"limit={limit}\"]\n if space_keys:\n params.append(f\"{key_param}={','.join(space_keys)}\")\n if start is not None and not is_v2:\n params.append(f\"start={start}\")\n\n return f\"{base_url}?{'&'.join(params)}\"\n\n def _paginate_spaces_for_endpoint(\n self,\n is_v2: bool,\n base_url: str,\n limit: int,\n space_keys: list[str] | None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"Internal helper to paginate through spaces for a specific API endpoint.\"\"\"\n start = 0\n url = self._build_spaces_url(\n is_v2, base_url, limit, space_keys, start if not is_v2 else None\n )\n\n while url:\n response = self.get(url, advanced_mode=True)\n response.raise_for_status()\n data = response.json()\n\n results = data.get(\"results\", [])\n if not results:\n return\n\n yield from results\n\n if is_v2:\n url = data.get(\"_links\", {}).get(\"next\", \"\")\n else:\n if len(results) < limit:\n return\n start += len(results)\n url = self._build_spaces_url(is_v2, base_url, limit, space_keys, start)\n\n def retrieve_confluence_spaces(\n self,\n space_keys: list[str] | None = None,\n limit: int = 50,\n ) -> Iterator[dict[str, str]]:\n \"\"\"\n Retrieve spaces from Confluence using v2 API (Cloud) or v1 API (Server/fallback).\n\n Args:\n space_keys: Optional list of space keys to filter by\n limit: Results per page (default 50)\n\n Yields:\n Space dictionaries with keys: id, key, name, type, status, etc.\n\n Note:\n For Cloud instances, attempts v2 API first. If v2 returns 404,\n automatically falls back to v1 API for compatibility with older instances.\n \"\"\"\n # Determine API version once\n use_v2 = self._is_cloud and not self.scoped_token\n base_url = _CONFLUENCE_SPACES_API_V2 if use_v2 else _CONFLUENCE_SPACES_API_V1\n\n try:\n yield from self._paginate_spaces_for_endpoint(\n use_v2, base_url, limit, space_keys\n )\n except HTTPError as e:\n if e.response.status_code == 404 and use_v2:\n logger.warning(\n \"v2 spaces API returned 404, falling back to v1 API. This may indicate an older Confluence Cloud instance.\"\n )\n # Fallback to v1\n yield from self._paginate_spaces_for_endpoint(\n False, _CONFLUENCE_SPACES_API_V1, limit, space_keys\n )\n else:\n raise\n\n def _probe_connection(\n self,\n **kwargs: Any,\n ) -> None:\n merged_kwargs = {**self.shared_base_kwargs, **kwargs}\n # add special timeout to make sure that we don't hang indefinitely\n merged_kwargs[\"timeout\"] = self.PROBE_TIMEOUT\n\n with self._credentials_provider:\n credentials, _ = self._renew_credentials()\n if self.scoped_token:\n # v2 endpoint doesn't always work with scoped tokens, use v1\n token = credentials[\"confluence_access_token\"]\n probe_url = f\"{self.base_url}/{_CONFLUENCE_SPACES_API_V1}?limit=1\"\n import requests\n\n try:\n r = requests.get(\n probe_url,\n headers={\"Authorization\": f\"Bearer {token}\"},\n timeout=10,\n )\n r.raise_for_status()\n except HTTPError as e:\n if e.response.status_code == 403:\n logger.warning(\n \"scoped token authenticated but not valid for probe endpoint (spaces)\"\n )\n else:\n if \"WWW-Authenticate\" in e.response.headers:\n logger.warning(\n f\"WWW-Authenticate: {e.response.headers['WWW-Authenticate']}\"\n )\n logger.warning(f\"Full error: {e.response.text}\")\n raise e\n return\n\n # Initialize connection with probe timeout settings\n self._confluence = self._initialize_connection_helper(\n credentials, **merged_kwargs\n )\n\n # Retrieve first space to validate connection\n spaces_iter = self.retrieve_confluence_spaces(limit=1)\n first_space = next(spaces_iter, None)\n\n if not first_space:\n raise RuntimeError(\n f\"No spaces found at {self._url}! Check your credentials and wiki_base and make sure is_cloud is set correctly.\"\n )\n\n logger.info(\"Confluence probe succeeded.\")\n\n def _initialize_connection(\n self,\n **kwargs: Any,\n ) -> None:\n \"\"\"Called externally to init the connection in a thread safe manner.\"\"\"\n merged_kwargs = {**self.shared_base_kwargs, **kwargs}\n with self._credentials_provider:\n credentials, _ = self._renew_credentials()\n self._confluence = self._initialize_connection_helper(\n credentials, **merged_kwargs\n )\n self._kwargs = merged_kwargs\n\n def _initialize_connection_helper(\n self,\n credentials: dict[str, Any],\n **kwargs: Any,\n ) -> Confluence:\n \"\"\"Called internally to init the connection. Distributed locking\n to prevent multiple threads from modifying the credentials\n must be handled around this function.\"\"\"\n\n confluence = None\n\n # probe connection with direct client, no retries\n if \"confluence_refresh_token\" in credentials:\n logger.info(\"Connecting to Confluence Cloud with OAuth Access Token.\")\n\n oauth2_dict: dict[str, Any] = OnyxConfluence._make_oauth2_dict(credentials)\n url = f\"https://api.atlassian.com/ex/confluence/{credentials['cloud_id']}\"\n confluence = Confluence(url=url, oauth2=oauth2_dict, **kwargs)\n else:\n logger.info(\n f\"Connecting to Confluence with Personal Access Token as user: {credentials['confluence_username']}\"\n )\n if self._is_cloud:\n confluence = Confluence(\n url=self._url,\n username=credentials[\"confluence_username\"],\n password=credentials[\"confluence_access_token\"],\n **kwargs,\n )\n else:\n confluence = Confluence(\n url=self._url,\n token=credentials[\"confluence_access_token\"],\n **kwargs,\n )\n\n return confluence\n\n # https://developer.atlassian.com/cloud/confluence/rate-limiting/\n # This uses the native rate limiting option provided by the\n # confluence client and otherwise applies a simpler set of error handling.\n def _make_rate_limited_confluence_method(\n self, name: str, credential_provider: CredentialsProviderInterface | None\n ) -> Callable[..., Any]:\n def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:\n MAX_RETRIES = 5\n\n TIMEOUT = 600\n timeout_at = time.monotonic() + TIMEOUT\n\n for attempt in range(MAX_RETRIES):\n if time.monotonic() > timeout_at:\n raise TimeoutError(\n f\"Confluence call attempts took longer than {TIMEOUT} seconds.\"\n )\n\n # we're relying more on the client to rate limit itself\n # and applying our own retries in a more specific set of circumstances\n try:\n if credential_provider:\n with credential_provider:\n credentials, renewed = self._renew_credentials()\n if renewed:\n self._confluence = self._initialize_connection_helper(\n credentials, **self._kwargs\n )\n attr = getattr(self._confluence, name, None)\n if attr is None:\n # The underlying Confluence client doesn't have this attribute\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n\n return attr(*args, **kwargs)\n else:\n attr = getattr(self._confluence, name, None)\n if attr is None:\n # The underlying Confluence client doesn't have this attribute\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n\n return attr(*args, **kwargs)\n\n except HTTPError as e:\n delay_until = _handle_http_error(e, attempt, MAX_RETRIES)\n logger.warning(\n f\"HTTPError in confluence call. Retrying in {delay_until} seconds...\"\n )\n while time.monotonic() < delay_until:\n # in the future, check a signal here to exit\n time.sleep(1)\n except AttributeError as e:\n # Some error within the Confluence library, unclear why it fails.\n # Users reported it to be intermittent, so just retry\n if attempt == MAX_RETRIES - 1:\n raise e\n\n logger.exception(\n \"Confluence Client raised an AttributeError. Retrying...\"\n )\n time.sleep(5)\n\n return wrapped_call\n\n def __getattr__(self, name: str) -> Any:\n \"\"\"Dynamically intercept attribute/method access.\"\"\"\n attr = getattr(self._confluence, name, None)\n if attr is None:\n # The underlying Confluence client doesn't have this attribute\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n\n # If it's not a method, just return it after ensuring token validity\n if not callable(attr):\n return attr\n\n # skip methods that start with \"_\"\n if name.startswith(\"_\"):\n return attr\n\n # wrap the method with our retry handler\n rate_limited_method: Callable[..., Any] = (\n self._make_rate_limited_confluence_method(name, self._credentials_provider)\n )\n\n return rate_limited_method\n\n def _try_one_by_one_for_paginated_url(\n self,\n url_suffix: str,\n initial_start: int,\n limit: int,\n ) -> Generator[dict[str, Any], None, str | None]:\n \"\"\"\n Go through `limit` items, starting at `initial_start` one by one (e.g. using\n `limit=1` for each call).\n\n If we encounter an error, we skip the item and try the next one. We will return\n the items we were able to retrieve successfully.\n\n Returns the expected next url_suffix. Returns None if it thinks we've hit the end.\n\n TODO (chris): make this yield failures as well as successes.\n TODO (chris): make this work for confluence cloud somehow.\n \"\"\"\n if self._is_cloud:\n raise RuntimeError(\"This method is not implemented for Confluence Cloud.\")\n\n found_empty_page = False\n temp_url_suffix = url_suffix\n\n for ind in range(limit):\n try:\n temp_url_suffix = update_param_in_path(\n url_suffix, \"start\", str(initial_start + ind)\n )\n temp_url_suffix = update_param_in_path(temp_url_suffix, \"limit\", \"1\")\n logger.info(f\"Making recovery confluence call to {temp_url_suffix}\")\n raw_response = self.get(path=temp_url_suffix, advanced_mode=True)\n raw_response.raise_for_status()\n\n latest_results = raw_response.json().get(\"results\", [])\n yield from latest_results\n\n if not latest_results:\n # no more results, break out of the loop\n logger.info(\n f\"No results found for call '{temp_url_suffix}'Stopping pagination.\"\n )\n found_empty_page = True\n break\n except Exception:\n logger.exception(\n f\"Error in confluence call to {temp_url_suffix}. Continuing.\"\n )\n\n if found_empty_page:\n return None\n\n # if we got here, we successfully tried `limit` items\n return update_param_in_path(url_suffix, \"start\", str(initial_start + limit))\n\n def _paginate_url(\n self,\n url_suffix: str,\n limit: int | None = None,\n # Called with the next url to use to get the next page\n next_page_callback: Callable[[str], None] | None = None,\n force_offset_pagination: bool = False,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n This will paginate through the top level query.\n \"\"\"\n if not limit:\n limit = _DEFAULT_PAGINATION_LIMIT\n\n url_suffix = update_param_in_path(url_suffix, \"limit\", str(limit))\n\n while url_suffix:\n logger.debug(f\"Making confluence call to {url_suffix}\")\n try:\n # Only pass params if they're not already in the URL to avoid duplicate\n # params accumulating. Confluence's _links.next already includes these.\n params = {}\n if \"body-format=\" not in url_suffix:\n params[\"body-format\"] = \"atlas_doc_format\"\n if \"expand=\" not in url_suffix:\n params[\"expand\"] = \"body.atlas_doc_format\"\n\n raw_response = self.get(\n path=url_suffix,\n advanced_mode=True,\n params=params,\n )\n except Exception as e:\n logger.exception(f\"Error in confluence call to {url_suffix}\")\n raise e\n\n try:\n raw_response.raise_for_status()\n except Exception as e:\n logger.warning(f\"Error in confluence call to {url_suffix}\")\n\n # If the problematic expansion is in the url, replace it\n # with the replacement expansion and try again\n # If that fails, raise the error\n if _PROBLEMATIC_EXPANSIONS in url_suffix:\n logger.warning(\n f\"Replacing {_PROBLEMATIC_EXPANSIONS} with {_REPLACEMENT_EXPANSIONS} and trying again.\"\n )\n url_suffix = url_suffix.replace(\n _PROBLEMATIC_EXPANSIONS,\n _REPLACEMENT_EXPANSIONS,\n )\n continue\n\n # If we fail due to a 500, try one by one.\n # NOTE: this iterative approach only works for server, since cloud uses cursor-based\n # pagination\n if raw_response.status_code == 500 and not self._is_cloud:\n initial_start = get_start_param_from_url(url_suffix)\n if initial_start is None:\n # can't handle this if we don't have offset-based pagination\n raise\n\n # this will just yield the successful items from the batch\n new_url_suffix = yield from self._try_one_by_one_for_paginated_url(\n url_suffix,\n initial_start=initial_start,\n limit=limit,\n )\n\n # this means we ran into an empty page\n if new_url_suffix is None:\n if next_page_callback:\n next_page_callback(\"\")\n break\n\n url_suffix = new_url_suffix\n continue\n\n else:\n logger.exception(\n f\"Error in confluence call to {url_suffix} \\n\"\n f\"Raw Response Text: {raw_response.text} \\n\"\n f\"Full Response: {raw_response.__dict__} \\n\"\n f\"Error: {e} \\n\"\n )\n raise\n\n try:\n next_response = raw_response.json()\n except Exception as e:\n logger.exception(\n f\"Failed to parse response as JSON. Response: {raw_response.__dict__}\"\n )\n raise e\n\n # Yield the results individually.\n results = cast(list[dict[str, Any]], next_response.get(\"results\", []))\n\n # Note 1:\n # Make sure we don't update the start by more than the amount\n # of results we were able to retrieve. The Confluence API has a\n # weird behavior where if you pass in a limit that is too large for\n # the configured server, it will artificially limit the amount of\n # results returned BUT will not apply this to the start parameter.\n # This will cause us to miss results.\n #\n # Note 2:\n # We specifically perform manual yielding (i.e., `for x in xs: yield x`) as opposed to using a `yield from xs`\n # because we *have to call the `next_page_callback`* prior to yielding the last element!\n #\n # If we did:\n #\n # ```py\n # yield from results\n # if next_page_callback:\n # next_page_callback(url_suffix)\n # ```\n #\n # then the logic would fail since the iterator would finish (and the calling scope would exit out of its driving\n # loop) prior to the callback being called.\n\n old_url_suffix = url_suffix\n updated_start = get_start_param_from_url(old_url_suffix)\n url_suffix = cast(str, next_response.get(\"_links\", {}).get(\"next\", \"\"))\n for i, result in enumerate(results):\n updated_start += 1\n if url_suffix and next_page_callback and i == len(results) - 1:\n # update the url if we're on the last result in the page\n if not self._is_cloud:\n # If confluence claims there are more results, we update the start param\n # based on how many results were returned and try again.\n url_suffix = update_param_in_path(\n url_suffix, \"start\", str(updated_start)\n )\n # notify the caller of the new url\n next_page_callback(url_suffix)\n\n elif force_offset_pagination and i == len(results) - 1:\n url_suffix = update_param_in_path(\n old_url_suffix, \"start\", str(updated_start)\n )\n\n yield result\n\n # we've observed that Confluence sometimes returns a next link despite giving\n # 0 results. This is a bug with Confluence, so we need to check for it and\n # stop paginating.\n if url_suffix and not results:\n logger.info(\n f\"No results found for call '{old_url_suffix}' despite next link being present. Stopping pagination.\"\n )\n break\n\n def build_cql_url(self, cql: str, expand: str | None = None) -> str:\n expand_string = f\"&expand={expand}\" if expand else \"\"\n return f\"rest/api/content/search?cql={cql}{expand_string}\"\n\n def paginated_cql_retrieval(\n self,\n cql: str,\n expand: str | None = None,\n limit: int | None = None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n The content/search endpoint can be used to fetch pages, attachments, and comments.\n \"\"\"\n cql_url = self.build_cql_url(cql, expand)\n yield from self._paginate_url(cql_url, limit)\n\n def paginated_page_retrieval(\n self,\n cql_url: str,\n limit: int,\n # Called with the next url to use to get the next page\n next_page_callback: Callable[[str], None] | None = None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n Error handling (and testing) wrapper for _paginate_url,\n because the current approach to page retrieval involves handling the\n next page links manually.\n \"\"\"\n try:\n yield from self._paginate_url(\n cql_url, limit=limit, next_page_callback=next_page_callback\n )\n except Exception as e:\n logger.exception(f\"Error in paginated_page_retrieval: {e}\")\n raise e\n\n def cql_paginate_all_expansions(\n self,\n cql: str,\n expand: str | None = None,\n limit: int | None = None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n This function will paginate through the top level query first, then\n paginate through all of the expansions.\n \"\"\"\n\n def _traverse_and_update(data: dict | list) -> None:\n if isinstance(data, dict):\n next_url = data.get(\"_links\", {}).get(\"next\")\n if next_url and \"results\" in data:\n data[\"results\"].extend(self._paginate_url(next_url, limit=limit))\n\n for value in data.values():\n _traverse_and_update(value)\n elif isinstance(data, list):\n for item in data:\n _traverse_and_update(item)\n\n for confluence_object in self.paginated_cql_retrieval(cql, expand, limit):\n _traverse_and_update(confluence_object)\n yield confluence_object\n\n def paginated_cql_user_retrieval(\n self,\n expand: str | None = None,\n limit: int | None = None,\n ) -> Iterator[ConfluenceUser]:\n \"\"\"\n The search/user endpoint can be used to fetch users.\n It's a separate endpoint from the content/search endpoint used only for users.\n Otherwise it's very similar to the content/search endpoint.\n \"\"\"\n\n # this is needed since there is a live bug with Confluence Server/Data Center\n # where not all users are returned by the APIs. This is a workaround needed until\n # that is patched.\n if self._confluence_user_profiles_override:\n yield from self._confluence_user_profiles_override\n\n elif self._is_cloud:\n cql = \"type=user\"\n url = \"rest/api/search/user\"\n expand_string = f\"&expand={expand}\" if expand else \"\"\n url += f\"?cql={cql}{expand_string}\"\n for user_result in self._paginate_url(\n url, limit, force_offset_pagination=True\n ):\n # Example response:\n # {\n # 'user': {\n # 'type': 'known',\n # 'accountId': '712020:35e60fbb-d0f3-4c91-b8c1-f2dd1d69462d',\n # 'accountType': 'atlassian',\n # 'email': 'chris@danswer.ai',\n # 'publicName': 'Chris Weaver',\n # 'profilePicture': {\n # 'path': '/wiki/aa-avatar/712020:35e60fbb-d0f3-4c91-b8c1-f2dd1d69462d',\n # 'width': 48,\n # 'height': 48,\n # 'isDefault': False\n # },\n # 'displayName': 'Chris Weaver',\n # 'isExternalCollaborator': False,\n # '_expandable': {\n # 'operations': '',\n # 'personalSpace': ''\n # },\n # '_links': {\n # 'self': 'https://danswerai.atlassian.net/wiki/rest/api/user?accountId=712020:35e60fbb-d0f3-4c91-b8c1-f2dd1d69462d'\n # }\n # },\n # 'title': 'Chris Weaver',\n # 'excerpt': '',\n # 'url': '/people/712020:35e60fbb-d0f3-4c91-b8c1-f2dd1d69462d',\n # 'breadcrumbs': [],\n # 'entityType': 'user',\n # 'iconCssClass': 'aui-icon content-type-profile',\n # 'lastModified': '2025-02-18T04:08:03.579Z',\n # 'score': 0.0\n # }\n user = user_result[\"user\"]\n yield ConfluenceUser(\n user_id=user[\"accountId\"],\n username=None,\n display_name=user[\"displayName\"],\n email=user.get(\"email\"),\n type=user[\"accountType\"],\n )\n else:\n for user in self._paginate_url(\"rest/api/user/list\", limit):\n yield ConfluenceUser(\n user_id=user[\"userKey\"],\n username=user[\"username\"],\n display_name=user[\"displayName\"],\n email=None,\n type=user.get(\"type\", \"user\"),\n )\n\n def paginated_groups_by_user_retrieval(\n self,\n user_id: str, # accountId in Cloud, userKey in Server\n limit: int | None = None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n This is not an SQL like query.\n It's a confluence specific endpoint that can be used to fetch groups.\n \"\"\"\n user_field = \"accountId\" if self._is_cloud else \"key\"\n user_value = user_id\n # Server uses userKey (but calls it key during the API call), Cloud uses accountId\n user_query = f\"{user_field}={quote(user_value)}\"\n\n url = f\"rest/api/user/memberof?{user_query}\"\n yield from self._paginate_url(url, limit, force_offset_pagination=True)\n\n def paginated_groups_retrieval(\n self,\n limit: int | None = None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n This is not an SQL like query.\n It's a confluence specific endpoint that can be used to fetch groups.\n \"\"\"\n yield from self._paginate_url(\"rest/api/group\", limit)\n\n def paginated_group_members_retrieval(\n self,\n group_name: str,\n limit: int | None = None,\n ) -> Iterator[dict[str, Any]]:\n \"\"\"\n This is not an SQL like query.\n It's a confluence specific endpoint that can be used to fetch the members of a group.\n THIS DOESN'T WORK FOR SERVER because it breaks when there is a slash in the group name.\n E.g. neither \"test/group\" nor \"test%2Fgroup\" works for confluence.\n \"\"\"\n group_name = quote(group_name)\n yield from self._paginate_url(f\"rest/api/group/{group_name}/member\", limit)\n\n def get_all_space_permissions_server(\n self,\n space_key: str,\n ) -> list[dict[str, Any]]:\n \"\"\"\n This is a confluence server/data center specific method that can be used to\n fetch the permissions of a space.\n\n NOTE: This uses the JSON-RPC API which is the ONLY way to get space permissions\n on Confluence Server/Data Center. The REST API equivalent (expand=permissions)\n is Cloud-only and not available on Data Center as of version 8.9.x.\n\n If this fails with 401 Unauthorized, the customer needs to enable JSON-RPC:\n Confluence Admin -> General Configuration -> Further Configuration\n -> Enable \"Remote API (XML-RPC & SOAP)\"\n \"\"\"\n url = \"rpc/json-rpc/confluenceservice-v2\"\n data = {\n \"jsonrpc\": \"2.0\",\n \"method\": \"getSpacePermissionSets\",\n \"id\": 7,\n \"params\": [space_key],\n }\n try:\n response = self.post(url, data=data)\n except HTTPError as e:\n if e.response is not None and e.response.status_code == 401:\n raise HTTPError(\n \"Unauthorized (401) when calling JSON-RPC API for space permissions. \"\n \"This is likely because the Remote API is disabled. \"\n \"To fix: Confluence Admin -> General Configuration -> Further Configuration \"\n \"-> Enable 'Remote API (XML-RPC & SOAP)'\",\n response=e.response,\n ) from e\n raise\n logger.debug(f\"jsonrpc response: {response}\")\n if not response.get(\"result\"):\n logger.warning(\n f\"No jsonrpc response for space permissions for space {space_key}\\nResponse: {response}\"\n )\n\n return response.get(\"result\", [])\n\n def get_current_user(self, expand: str | None = None) -> Any:\n \"\"\"\n Implements a method that isn't in the third party client.\n\n Get information about the current user\n :param expand: OPTIONAL expand for get status of user.\n Possible param is \"status\". Results are \"Active, Deactivated\"\n :return: Returns the user details\n \"\"\"\n\n from atlassian.errors import ApiPermissionError # type:ignore\n\n url = \"rest/api/user/current\"\n params = {}\n if expand:\n params[\"expand\"] = expand\n try:\n response = self.get(url, params=params)\n except HTTPError as e:\n if e.response.status_code == 403:\n raise ApiPermissionError(\n \"The calling user does not have permission\", reason=e\n )\n raise\n return response\n\n\ndef get_user_email_from_username__server(\n confluence_client: OnyxConfluence, user_name: str\n) -> str | None:\n global _USER_EMAIL_CACHE\n if _USER_EMAIL_CACHE.get(user_name) is None:\n try:\n response = confluence_client.get_mobile_parameters(user_name)\n email = response.get(\"email\")\n except HTTPError as e:\n status_code = e.response.status_code if e.response is not None else \"N/A\"\n logger.warning(\n f\"Failed to get confluence email for {user_name}: HTTP {status_code} - {e}\"\n )\n # For now, we'll just return None and log a warning. This means\n # we will keep retrying to get the email every group sync.\n email = None\n except Exception as e:\n logger.warning(\n f\"Failed to get confluence email for {user_name}: {type(e).__name__} - {e}\"\n )\n email = None\n _USER_EMAIL_CACHE[user_name] = email\n return _USER_EMAIL_CACHE[user_name]\n\n\ndef _get_user(confluence_client: OnyxConfluence, user_id: str) -> str:\n \"\"\"Get Confluence Display Name based on the account-id or userkey value\n\n Args:\n user_id (str): The user id (i.e: the account-id or userkey)\n confluence_client (Confluence): The Confluence Client\n\n Returns:\n str: The User Display Name. 'Unknown User' if the user is deactivated or not found\n \"\"\"\n global _USER_ID_TO_DISPLAY_NAME_CACHE\n if _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) is None:\n try:\n result = confluence_client.get_user_details_by_userkey(user_id)\n found_display_name = result.get(\"displayName\")\n except Exception:\n found_display_name = None\n\n if not found_display_name:\n try:\n result = confluence_client.get_user_details_by_accountid(user_id)\n found_display_name = result.get(\"displayName\")\n except Exception:\n found_display_name = None\n\n _USER_ID_TO_DISPLAY_NAME_CACHE[user_id] = found_display_name\n\n return _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) or _USER_NOT_FOUND\n\n\ndef sanitize_attachment_title(title: str) -> str:\n \"\"\"\n Sanitize the attachment title to be a valid HTML attribute.\n \"\"\"\n return title.replace(\"<\", \"_\").replace(\">\", \"_\").replace(\" \", \"_\").replace(\":\", \"_\")\n\n\ndef extract_text_from_confluence_html(\n confluence_client: OnyxConfluence,\n confluence_object: dict[str, Any],\n fetched_titles: set[str],\n) -> str:\n \"\"\"Parse a Confluence html page and replace the 'user Id' by the real\n User Display Name\n\n Args:\n confluence_object (dict): The confluence object as a dict\n confluence_client (Confluence): Confluence client\n fetched_titles (set[str]): The titles of the pages that have already been fetched\n Returns:\n str: loaded and formated Confluence page\n \"\"\"\n body = confluence_object[\"body\"]\n object_html = body.get(\"storage\", body.get(\"view\", {})).get(\"value\")\n\n soup = bs4.BeautifulSoup(object_html, \"html.parser\")\n\n _remove_macro_stylings(soup=soup)\n\n for user in soup.findAll(\"ri:user\"):\n user_id = (\n user.attrs[\"ri:account-id\"]\n if \"ri:account-id\" in user.attrs\n else user.get(\"ri:userkey\")\n )\n if not user_id:\n logger.warning(\n f\"ri:userkey not found in ri:user element. Found attrs: {user.attrs}\"\n )\n continue\n # Include @ sign for tagging, more clear for LLM\n user.replaceWith(\"@\" + _get_user(confluence_client, user_id))\n\n for html_page_reference in soup.findAll(\"ac:structured-macro\"):\n # Here, we only want to process page within page macros\n if html_page_reference.attrs.get(\"ac:name\") != \"include\":\n continue\n\n page_data = html_page_reference.find(\"ri:page\")\n if not page_data:\n logger.warning(\n f\"Skipping retrieval of {html_page_reference} because because page data is missing\"\n )\n continue\n\n page_title = page_data.attrs.get(\"ri:content-title\")\n if not page_title:\n # only fetch pages that have a title\n logger.warning(\n f\"Skipping retrieval of {html_page_reference} because it has no title\"\n )\n continue\n\n if page_title in fetched_titles:\n # prevent recursive fetching of pages\n logger.debug(f\"Skipping {page_title} because it has already been fetched\")\n continue\n\n fetched_titles.add(page_title)\n\n # Wrap this in a try-except because there are some pages that might not exist\n try:\n page_query = f\"type=page and title='{quote(page_title)}'\"\n\n page_contents: dict[str, Any] | None = None\n # Confluence enforces title uniqueness, so we should only get one result here\n for page in confluence_client.paginated_cql_retrieval(\n cql=page_query,\n expand=\"body.storage.value\",\n limit=1,\n ):\n page_contents = page\n break\n except Exception as e:\n logger.warning(\n f\"Error getting page contents for object {confluence_object}: {e}\"\n )\n continue\n\n if not page_contents:\n continue\n\n text_from_page = extract_text_from_confluence_html(\n confluence_client=confluence_client,\n confluence_object=page_contents,\n fetched_titles=fetched_titles,\n )\n\n html_page_reference.replaceWith(text_from_page)\n\n for html_link_body in soup.findAll(\"ac:link-body\"):\n # This extracts the text from inline links in the page so they can be\n # represented in the document text as plain text\n try:\n text_from_link = html_link_body.text\n html_link_body.replaceWith(f\"(LINK TEXT: {text_from_link})\")\n except Exception as e:\n logger.warning(f\"Error processing ac:link-body: {e}\")\n\n for html_attachment in soup.findAll(\"ri:attachment\"):\n # This extracts the text from inline attachments in the page so they can be\n # represented in the document text as plain text\n try:\n html_attachment.replaceWith(\n f\"{sanitize_attachment_title(html_attachment.attrs['ri:filename'])}\"\n ) # to be replaced later\n except Exception as e:\n logger.warning(f\"Error processing ac:attachment: {e}\")\n\n return format_document_soup(soup)\n\n\ndef _remove_macro_stylings(soup: bs4.BeautifulSoup) -> None:\n for macro_root in soup.findAll(\"ac:structured-macro\"):\n if not isinstance(macro_root, bs4.Tag):\n continue\n\n macro_styling = macro_root.find(name=\"ac:parameter\", attrs={\"ac:name\": \"page\"})\n if not macro_styling or not isinstance(macro_styling, bs4.Tag):\n continue\n\n macro_styling.extract()\n" }, { "path": "backend/onyx/connectors/confluence/user_profile_override.py", "content": "from onyx.connectors.confluence.models import ConfluenceUser\n\n\ndef process_confluence_user_profiles_override(\n confluence_user_email_override: list[dict[str, str]],\n) -> list[ConfluenceUser]:\n return [\n ConfluenceUser(\n user_id=override[\"user_id\"],\n # username is not returned by the Confluence Server API anyways\n username=override[\"username\"],\n display_name=override[\"display_name\"],\n email=override[\"email\"],\n type=override[\"type\"],\n )\n for override in confluence_user_email_override\n if override is not None\n ]\n" }, { "path": "backend/onyx/connectors/confluence/utils.py", "content": "import math\nimport time\nfrom collections.abc import Callable\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom io import BytesIO\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import cast\nfrom typing import TYPE_CHECKING\nfrom typing import TypeVar\nfrom urllib.parse import parse_qs\nfrom urllib.parse import quote\nfrom urllib.parse import urljoin\nfrom urllib.parse import urlparse\n\nimport requests\nfrom pydantic import BaseModel\n\nfrom onyx.configs.app_configs import (\n CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD,\n)\nfrom onyx.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.file_types import OnyxMimeTypes\nfrom onyx.file_processing.image_utils import store_image_and_create_section\nfrom onyx.utils.logger import setup_logger\n\nif TYPE_CHECKING:\n from onyx.connectors.confluence.onyx_confluence import OnyxConfluence\n\n\nlogger = setup_logger()\n\nCONFLUENCE_OAUTH_TOKEN_URL = \"https://auth.atlassian.com/oauth/token\"\nRATE_LIMIT_MESSAGE_LOWERCASE = \"Rate limit exceeded\".lower()\n\n\nclass TokenResponse(BaseModel):\n access_token: str\n expires_in: int\n token_type: str\n refresh_token: str\n scope: str\n\n\ndef validate_attachment_filetype(\n attachment: dict[str, Any],\n) -> bool:\n \"\"\"\n Validates if the attachment is a supported file type.\n \"\"\"\n media_type = attachment.get(\"metadata\", {}).get(\"mediaType\", \"\")\n if media_type.startswith(\"image/\"):\n return media_type in OnyxMimeTypes.IMAGE_MIME_TYPES\n\n # For non-image files, check if we support the extension\n title = attachment.get(\"title\", \"\")\n extension = get_file_ext(title)\n\n return extension in OnyxFileExtensions.ALL_ALLOWED_EXTENSIONS\n\n\nclass AttachmentProcessingResult(BaseModel):\n \"\"\"\n A container for results after processing a Confluence attachment.\n 'text' is the textual content of the attachment.\n 'file_name' is the final file name used in FileStore to store the content.\n 'error' holds an exception or string if something failed.\n \"\"\"\n\n text: str | None\n file_name: str | None\n error: str | None = None\n\n\ndef _make_attachment_link(\n confluence_client: \"OnyxConfluence\",\n attachment: dict[str, Any],\n parent_content_id: str | None = None,\n) -> str | None:\n download_link = \"\"\n\n if \"api.atlassian.com\" in confluence_client.url:\n # https://developer.atlassian.com/cloud/confluence/rest/v1/api-group-content---attachments/#api-wiki-rest-api-content-id-child-attachment-attachmentid-download-get\n if not parent_content_id:\n logger.warning(\n \"parent_content_id is required to download attachments from Confluence Cloud!\"\n )\n return None\n\n download_link = (\n confluence_client.url\n + f\"/rest/api/content/{parent_content_id}/child/attachment/{attachment['id']}/download\"\n )\n else:\n download_link = confluence_client.url + attachment[\"_links\"][\"download\"]\n\n return download_link\n\n\ndef process_attachment(\n confluence_client: \"OnyxConfluence\",\n attachment: dict[str, Any],\n parent_content_id: str | None,\n allow_images: bool,\n) -> AttachmentProcessingResult:\n \"\"\"\n Processes a Confluence attachment. If it's a document, extracts text,\n or if it's an image, stores it for later analysis. Returns a structured result.\n \"\"\"\n try:\n # Get the media type from the attachment metadata\n media_type: str = attachment.get(\"metadata\", {}).get(\"mediaType\", \"\")\n # Validate the attachment type\n if not validate_attachment_filetype(attachment):\n return AttachmentProcessingResult(\n text=None,\n file_name=None,\n error=f\"Unsupported file type: {media_type}\",\n )\n\n attachment_link = _make_attachment_link(\n confluence_client, attachment, parent_content_id\n )\n if not attachment_link:\n return AttachmentProcessingResult(\n text=None, file_name=None, error=\"Failed to make attachment link\"\n )\n\n attachment_size = attachment[\"extensions\"][\"fileSize\"]\n\n if media_type.startswith(\"image/\"):\n if not allow_images:\n return AttachmentProcessingResult(\n text=None,\n file_name=None,\n error=\"Image downloading is not enabled\",\n )\n else:\n if attachment_size > CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD:\n logger.warning(\n f\"Skipping {attachment_link} due to size. \"\n f\"size={attachment_size} \"\n f\"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD}\"\n )\n return AttachmentProcessingResult(\n text=None,\n file_name=None,\n error=f\"Attachment text too long: {attachment_size} chars\",\n )\n\n logger.info(\n f\"Downloading attachment: title={attachment['title']} length={attachment_size} link={attachment_link}\"\n )\n\n # Download the attachment\n resp: requests.Response = confluence_client._session.get(attachment_link)\n if resp.status_code != 200:\n logger.warning(\n f\"Failed to fetch {attachment_link} with status code {resp.status_code}\"\n )\n return AttachmentProcessingResult(\n text=None,\n file_name=None,\n error=f\"Attachment download status code is {resp.status_code}\",\n )\n\n raw_bytes = resp.content\n if not raw_bytes:\n return AttachmentProcessingResult(\n text=None, file_name=None, error=\"attachment.content is None\"\n )\n\n # Process image attachments\n if media_type.startswith(\"image/\"):\n return _process_image_attachment(\n confluence_client, attachment, raw_bytes, media_type\n )\n\n # Process document attachments\n try:\n text = extract_file_text(\n file=BytesIO(raw_bytes),\n file_name=attachment[\"title\"],\n )\n\n # Skip if the text is too long\n if len(text) > CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD:\n return AttachmentProcessingResult(\n text=None,\n file_name=None,\n error=f\"Attachment text too long: {len(text)} chars\",\n )\n\n return AttachmentProcessingResult(text=text, file_name=None, error=None)\n except Exception as e:\n return AttachmentProcessingResult(\n text=None, file_name=None, error=f\"Failed to extract text: {e}\"\n )\n\n except Exception as e:\n return AttachmentProcessingResult(\n text=None, file_name=None, error=f\"Failed to process attachment: {e}\"\n )\n\n\ndef _process_image_attachment(\n confluence_client: \"OnyxConfluence\", # noqa: ARG001\n attachment: dict[str, Any],\n raw_bytes: bytes,\n media_type: str,\n) -> AttachmentProcessingResult:\n \"\"\"Process an image attachment by saving it without generating a summary.\"\"\"\n try:\n # Use the standardized image storage and section creation\n section, file_name = store_image_and_create_section(\n image_data=raw_bytes,\n file_id=Path(attachment[\"id\"]).name,\n display_name=attachment[\"title\"],\n media_type=media_type,\n file_origin=FileOrigin.CONNECTOR,\n )\n logger.info(f\"Stored image attachment with file name: {file_name}\")\n\n # Return empty text but include the file_name for later processing\n return AttachmentProcessingResult(text=\"\", file_name=file_name, error=None)\n except Exception as e:\n msg = f\"Image storage failed for {attachment['title']}: {e}\"\n logger.error(msg, exc_info=e)\n return AttachmentProcessingResult(text=None, file_name=None, error=msg)\n\n\ndef convert_attachment_to_content(\n confluence_client: \"OnyxConfluence\",\n attachment: dict[str, Any],\n page_id: str,\n allow_images: bool,\n) -> tuple[str | None, str | None] | None:\n \"\"\"\n Facade function which:\n 1. Validates attachment type\n 2. Extracts content or stores image for later processing\n 3. Returns (content_text, stored_file_name) or None if we should skip it\n \"\"\"\n media_type = attachment.get(\"metadata\", {}).get(\"mediaType\", \"\")\n # Quick check for unsupported types:\n if media_type.startswith(\"video/\") or media_type == \"application/gliffy+json\":\n logger.warning(\n f\"Skipping unsupported attachment type: '{media_type}' for {attachment['title']}\"\n )\n return None\n\n result = process_attachment(confluence_client, attachment, page_id, allow_images)\n if result.error is not None:\n logger.warning(\n f\"Attachment {attachment['title']} encountered error: {result.error}\"\n )\n return None\n\n # Return the text and the file name\n return result.text, result.file_name\n\n\ndef build_confluence_document_id(\n base_url: str, content_url: str, is_cloud: bool\n) -> str:\n \"\"\"For confluence, the document id is the page url for a page based document\n or the attachment download url for an attachment based document\n\n Args:\n base_url (str): The base url of the Confluence instance\n content_url (str): The url of the page or attachment download url\n\n Returns:\n str: The document id\n \"\"\"\n\n # NOTE: urljoin is tricky and will drop the last segment of the base if it doesn't\n # end with \"/\" because it believes that makes it a file.\n final_url = base_url.rstrip(\"/\") + \"/\"\n if is_cloud and not final_url.endswith(\"/wiki/\"):\n final_url = urljoin(final_url, \"wiki\") + \"/\"\n final_url = urljoin(final_url, content_url.lstrip(\"/\"))\n return final_url\n\n\ndef datetime_from_string(datetime_string: str) -> datetime:\n datetime_object = datetime.fromisoformat(datetime_string)\n\n if datetime_object.tzinfo is None:\n # If no timezone info, assume it is UTC\n datetime_object = datetime_object.replace(tzinfo=timezone.utc)\n else:\n # If not in UTC, translate it\n datetime_object = datetime_object.astimezone(timezone.utc)\n\n return datetime_object\n\n\ndef confluence_refresh_tokens(\n client_id: str, client_secret: str, cloud_id: str, refresh_token: str\n) -> dict[str, Any]:\n # rotate the refresh and access token\n # Note that access tokens are only good for an hour in confluence cloud,\n # so we're going to have problems if the connector runs for longer\n # https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/#use-a-refresh-token-to-get-another-access-token-and-refresh-token-pair\n response = requests.post(\n CONFLUENCE_OAUTH_TOKEN_URL,\n headers={\"Content-Type\": \"application/x-www-form-urlencoded\"},\n data={\n \"grant_type\": \"refresh_token\",\n \"client_id\": client_id,\n \"client_secret\": client_secret,\n \"refresh_token\": refresh_token,\n },\n )\n\n try:\n token_response = TokenResponse.model_validate_json(response.text)\n except Exception:\n raise RuntimeError(\"Confluence Cloud token refresh failed.\")\n\n now = datetime.now(timezone.utc)\n expires_at = now + timedelta(seconds=token_response.expires_in)\n\n new_credentials: dict[str, Any] = {}\n new_credentials[\"confluence_access_token\"] = token_response.access_token\n new_credentials[\"confluence_refresh_token\"] = token_response.refresh_token\n new_credentials[\"created_at\"] = now.isoformat()\n new_credentials[\"expires_at\"] = expires_at.isoformat()\n new_credentials[\"expires_in\"] = token_response.expires_in\n new_credentials[\"scope\"] = token_response.scope\n new_credentials[\"cloud_id\"] = cloud_id\n return new_credentials\n\n\nF = TypeVar(\"F\", bound=Callable[..., Any])\n\n\n# https://developer.atlassian.com/cloud/confluence/rate-limiting/\n# this uses the native rate limiting option provided by the\n# confluence client and otherwise applies a simpler set of error handling\ndef handle_confluence_rate_limit(confluence_call: F) -> F:\n def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:\n MAX_RETRIES = 5\n\n TIMEOUT = 600\n timeout_at = time.monotonic() + TIMEOUT\n\n for attempt in range(MAX_RETRIES):\n if time.monotonic() > timeout_at:\n raise TimeoutError(\n f\"Confluence call attempts took longer than {TIMEOUT} seconds.\"\n )\n\n try:\n # we're relying more on the client to rate limit itself\n # and applying our own retries in a more specific set of circumstances\n return confluence_call(*args, **kwargs)\n except requests.HTTPError as e:\n delay_until = _handle_http_error(e, attempt, MAX_RETRIES)\n logger.warning(\n f\"HTTPError in confluence call. Retrying in {delay_until} seconds...\"\n )\n while time.monotonic() < delay_until:\n # in the future, check a signal here to exit\n time.sleep(1)\n except AttributeError as e:\n # Some error within the Confluence library, unclear why it fails.\n # Users reported it to be intermittent, so just retry\n if attempt == MAX_RETRIES - 1:\n raise e\n\n logger.exception(\n \"Confluence Client raised an AttributeError. Retrying...\"\n )\n time.sleep(5)\n\n return cast(F, wrapped_call)\n\n\ndef _handle_http_error(e: requests.HTTPError, attempt: int, max_retries: int) -> int:\n MIN_DELAY = 2\n MAX_DELAY = 60\n STARTING_DELAY = 5\n BACKOFF = 2\n\n # Check if the response or headers are None to avoid potential AttributeError\n if e.response is None or e.response.headers is None:\n logger.warning(\"HTTPError with `None` as response or as headers\")\n raise e\n\n # Confluence Server returns 403 when rate limited\n if e.response.status_code == 403:\n FORBIDDEN_MAX_RETRY_ATTEMPTS = 7\n FORBIDDEN_RETRY_DELAY = 10\n if attempt < FORBIDDEN_MAX_RETRY_ATTEMPTS:\n logger.warning(\n \"403 error. This sometimes happens when we hit \"\n f\"Confluence rate limits. Retrying in {FORBIDDEN_RETRY_DELAY} seconds...\"\n )\n return FORBIDDEN_RETRY_DELAY\n\n raise e\n\n if e.response.status_code >= 500:\n if attempt >= max_retries - 1:\n raise e\n\n delay = min(STARTING_DELAY * (BACKOFF**attempt), MAX_DELAY)\n logger.warning(\n f\"Server error {e.response.status_code}. \"\n f\"Retrying in {delay} seconds (attempt {attempt + 1})...\"\n )\n return math.ceil(time.monotonic() + delay)\n\n if (\n e.response.status_code != 429\n and RATE_LIMIT_MESSAGE_LOWERCASE not in e.response.text.lower()\n ):\n raise e\n\n retry_after = None\n\n retry_after_header = e.response.headers.get(\"Retry-After\")\n if retry_after_header is not None:\n try:\n retry_after = int(retry_after_header)\n if retry_after > MAX_DELAY:\n logger.warning(\n f\"Clamping retry_after from {retry_after} to {MAX_DELAY} seconds...\"\n )\n retry_after = MAX_DELAY\n if retry_after < MIN_DELAY:\n retry_after = MIN_DELAY\n except ValueError:\n pass\n\n if retry_after is not None:\n logger.warning(\n f\"Rate limiting with retry header. Retrying after {retry_after} seconds...\"\n )\n delay = retry_after\n else:\n logger.warning(\n \"Rate limiting without retry header. Retrying with exponential backoff...\"\n )\n delay = min(STARTING_DELAY * (BACKOFF**attempt), MAX_DELAY)\n\n delay_until = math.ceil(time.monotonic() + delay)\n return delay_until\n\n\ndef get_single_param_from_url(url: str, param: str) -> str | None:\n \"\"\"Get a parameter from a url\"\"\"\n parsed_url = urlparse(url)\n return parse_qs(parsed_url.query).get(param, [None])[0]\n\n\ndef get_start_param_from_url(url: str) -> int:\n \"\"\"Get the start parameter from a url\"\"\"\n start_str = get_single_param_from_url(url, \"start\")\n return int(start_str) if start_str else 0\n\n\ndef update_param_in_path(path: str, param: str, value: str) -> str:\n \"\"\"Update a parameter in a path. Path should look something like:\n\n /api/rest/users?start=0&limit=10\n \"\"\"\n parsed_url = urlparse(path)\n query_params = parse_qs(parsed_url.query)\n query_params[param] = [value]\n return (\n path.split(\"?\")[0]\n + \"?\"\n + \"&\".join(f\"{k}={quote(v[0])}\" for k, v in query_params.items())\n )\n" }, { "path": "backend/onyx/connectors/connector_runner.py", "content": "import sys\nimport time\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom typing import Generic\nfrom typing import TypeVar\n\nfrom onyx.connectors.interfaces import BaseConnector\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\nTimeRange = tuple[datetime, datetime]\n\nCT = TypeVar(\"CT\", bound=ConnectorCheckpoint)\n\n\ndef batched_doc_ids(\n checkpoint_connector_generator: CheckpointOutput[CT],\n batch_size: int,\n) -> Generator[set[str], None, None]:\n batch: set[str] = set()\n for document, hierarchy_node, failure, next_checkpoint in CheckpointOutputWrapper[\n CT\n ]()(checkpoint_connector_generator):\n if document is not None:\n batch.add(document.id)\n elif (\n failure and failure.failed_document and failure.failed_document.document_id\n ):\n batch.add(failure.failed_document.document_id)\n # HierarchyNodes don't have IDs that need to be batched for doc processing\n\n if len(batch) >= batch_size:\n yield batch\n batch = set()\n if len(batch) > 0:\n yield batch\n\n\nclass CheckpointOutputWrapper(Generic[CT]):\n \"\"\"\n Wraps a CheckpointOutput generator to give things back in a more digestible format,\n specifically for Document outputs.\n The connector format is easier for the connector implementor (e.g. it enforces exactly\n one new checkpoint is returned AND that the checkpoint is at the end), thus the different\n formats.\n \"\"\"\n\n def __init__(self) -> None:\n self.next_checkpoint: CT | None = None\n\n def __call__(\n self,\n checkpoint_connector_generator: CheckpointOutput[CT],\n ) -> Generator[\n tuple[\n Document | None, HierarchyNode | None, ConnectorFailure | None, CT | None\n ],\n None,\n None,\n ]:\n # grabs the final return value and stores it in the `next_checkpoint` variable\n def _inner_wrapper(\n checkpoint_connector_generator: CheckpointOutput[CT],\n ) -> CheckpointOutput[CT]:\n self.next_checkpoint = yield from checkpoint_connector_generator\n return self.next_checkpoint # not used\n\n for item in _inner_wrapper(checkpoint_connector_generator):\n if isinstance(item, Document):\n yield item, None, None, None\n elif isinstance(item, HierarchyNode):\n yield None, item, None, None\n elif isinstance(item, ConnectorFailure):\n yield None, None, item, None\n else:\n raise ValueError(f\"Invalid connector output type: {type(item)}\")\n\n if self.next_checkpoint is None:\n raise RuntimeError(\n \"Checkpoint is None. This should never happen - the connector should always return a checkpoint.\"\n )\n\n yield None, None, None, self.next_checkpoint\n\n\nclass ConnectorRunner(Generic[CT]):\n \"\"\"\n Handles:\n - Batching\n - Additional exception logging\n - Combining different connector types to a single interface\n \"\"\"\n\n def __init__(\n self,\n connector: BaseConnector,\n batch_size: int,\n # cannot be True for non-checkpointed connectors\n include_permissions: bool,\n time_range: TimeRange | None = None,\n ):\n if not isinstance(connector, CheckpointedConnector) and include_permissions:\n raise ValueError(\n \"include_permissions cannot be True for non-checkpointed connectors\"\n )\n\n self.connector = connector\n self.time_range = time_range\n self.batch_size = batch_size\n self.include_permissions = include_permissions\n\n self.doc_batch: list[Document] = []\n self.hierarchy_node_batch: list[HierarchyNode] = []\n\n def run(self, checkpoint: CT) -> Generator[\n tuple[\n list[Document] | None,\n list[HierarchyNode] | None,\n ConnectorFailure | None,\n CT | None,\n ],\n None,\n None,\n ]:\n \"\"\"\n Yields batches of Documents, HierarchyNodes, failures, and checkpoints.\n\n Returns tuples of:\n - (doc_batch, None, None, None) - batch of documents\n - (None, hierarchy_batch, None, None) - batch of hierarchy nodes\n - (None, None, failure, None) - a connector failure\n - (None, None, None, checkpoint) - new checkpoint\n \"\"\"\n try:\n if isinstance(self.connector, CheckpointedConnector):\n if self.time_range is None:\n raise ValueError(\"time_range is required for CheckpointedConnector\")\n\n start = time.monotonic()\n if self.include_permissions:\n if not isinstance(\n self.connector, CheckpointedConnectorWithPermSync\n ):\n raise ValueError(\n \"Connector does not support permission syncing\"\n )\n load_from_checkpoint = (\n self.connector.load_from_checkpoint_with_perm_sync\n )\n else:\n load_from_checkpoint = self.connector.load_from_checkpoint\n checkpoint_connector_generator = load_from_checkpoint(\n start=self.time_range[0].timestamp(),\n end=self.time_range[1].timestamp(),\n checkpoint=checkpoint,\n )\n next_checkpoint: CT | None = None\n # this is guaranteed to always run at least once with next_checkpoint being non-None\n for (\n document,\n hierarchy_node,\n failure,\n next_checkpoint,\n ) in CheckpointOutputWrapper[CT]()(checkpoint_connector_generator):\n if document is not None:\n self.doc_batch.append(document)\n\n if hierarchy_node is not None:\n self.hierarchy_node_batch.append(hierarchy_node)\n\n if failure is not None:\n yield None, None, failure, None\n\n # Yield hierarchy nodes batch if it reaches batch_size\n # (yield nodes before docs to maintain parent-before-child invariant)\n if len(self.hierarchy_node_batch) >= self.batch_size:\n yield None, self.hierarchy_node_batch, None, None\n self.hierarchy_node_batch = []\n\n # Yield document batch if it reaches batch_size\n # First flush any pending hierarchy nodes to ensure parents exist\n if len(self.doc_batch) >= self.batch_size:\n if len(self.hierarchy_node_batch) > 0:\n yield None, self.hierarchy_node_batch, None, None\n self.hierarchy_node_batch = []\n yield self.doc_batch, None, None, None\n self.doc_batch = []\n\n # yield remaining hierarchy nodes first (parents before children)\n if len(self.hierarchy_node_batch) > 0:\n yield None, self.hierarchy_node_batch, None, None\n self.hierarchy_node_batch = []\n\n # yield remaining documents\n if len(self.doc_batch) > 0:\n yield self.doc_batch, None, None, None\n self.doc_batch = []\n\n yield None, None, None, next_checkpoint\n\n logger.debug(\n f\"Connector took {time.monotonic() - start} seconds to get to the next checkpoint.\"\n )\n\n else:\n finished_checkpoint = self.connector.build_dummy_checkpoint()\n finished_checkpoint.has_more = False\n\n if isinstance(self.connector, PollConnector):\n if self.time_range is None:\n raise ValueError(\"time_range is required for PollConnector\")\n\n for batch in self.connector.poll_source(\n start=self.time_range[0].timestamp(),\n end=self.time_range[1].timestamp(),\n ):\n docs, nodes = self._separate_batch(batch)\n if nodes:\n yield None, nodes, None, None\n if docs:\n yield docs, None, None, None\n\n yield None, None, None, finished_checkpoint\n elif isinstance(self.connector, LoadConnector):\n for batch in self.connector.load_from_state():\n docs, nodes = self._separate_batch(batch)\n if nodes:\n yield None, nodes, None, None\n if docs:\n yield docs, None, None, None\n\n yield None, None, None, finished_checkpoint\n else:\n raise ValueError(f\"Invalid connector. type: {type(self.connector)}\")\n except Exception:\n exc_type, _, exc_traceback = sys.exc_info()\n\n # Traverse the traceback to find the last frame where the exception was raised\n tb = exc_traceback\n if tb is None:\n logger.error(\"No traceback found for exception\")\n raise\n\n while tb.tb_next:\n tb = tb.tb_next # Move to the next frame in the traceback\n\n # Get the local variables from the frame where the exception occurred\n local_vars = tb.tb_frame.f_locals\n local_vars_str = \"\\n\".join(\n f\"{key}: {value}\" for key, value in local_vars.items()\n )\n logger.error(\n f\"Error in connector. type: {exc_type};\\nlocal_vars below -> \\n{local_vars_str[:1024]}\"\n )\n raise\n\n def _separate_batch(\n self, batch: list[Document | HierarchyNode]\n ) -> tuple[list[Document], list[HierarchyNode]]:\n \"\"\"Separate a mixed batch into Documents and HierarchyNodes.\"\"\"\n docs: list[Document] = []\n nodes: list[HierarchyNode] = []\n for item in batch:\n if isinstance(item, Document):\n docs.append(item)\n elif isinstance(item, HierarchyNode):\n nodes.append(item)\n return docs, nodes\n" }, { "path": "backend/onyx/connectors/credentials_provider.py", "content": "import uuid\nfrom types import TracebackType\nfrom typing import Any\n\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy import select\n\nfrom onyx.connectors.interfaces import CredentialsProviderInterface\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.models import Credential\nfrom onyx.redis.redis_pool import get_redis_client\n\n\nclass OnyxDBCredentialsProvider(\n CredentialsProviderInterface[\"OnyxDBCredentialsProvider\"]\n):\n \"\"\"Implementation to allow the connector to callback and update credentials in the db.\n Required in cases where credentials can rotate while the connector is running.\n \"\"\"\n\n LOCK_TTL = 900 # TTL of the lock\n\n def __init__(self, tenant_id: str, connector_name: str, credential_id: int):\n self._tenant_id = tenant_id\n self._connector_name = connector_name\n self._credential_id = credential_id\n\n self.redis_client = get_redis_client(tenant_id=tenant_id)\n\n # lock used to prevent overlapping renewal of credentials\n self.lock_key = f\"da_lock:connector:{connector_name}:credential_{credential_id}\"\n self._lock: RedisLock = self.redis_client.lock(self.lock_key, self.LOCK_TTL)\n\n def __enter__(self) -> \"OnyxDBCredentialsProvider\":\n acquired = self._lock.acquire(blocking_timeout=self.LOCK_TTL)\n if not acquired:\n raise RuntimeError(f\"Could not acquire lock for key: {self.lock_key}\")\n\n return self\n\n def __exit__(\n self,\n exc_type: type[BaseException] | None,\n exc_value: BaseException | None,\n traceback: TracebackType | None,\n ) -> None:\n \"\"\"Release the lock when exiting the context.\"\"\"\n if self._lock and self._lock.owned():\n self._lock.release()\n\n def get_tenant_id(self) -> str | None:\n return self._tenant_id\n\n def get_provider_key(self) -> str:\n return str(self._credential_id)\n\n def get_credentials(self) -> dict[str, Any]:\n with get_session_with_tenant(tenant_id=self._tenant_id) as db_session:\n credential = db_session.execute(\n select(Credential).where(Credential.id == self._credential_id)\n ).scalar_one()\n\n if credential is None:\n raise ValueError(\n f\"No credential found: credential={self._credential_id}\"\n )\n\n if credential.credential_json is None:\n return {}\n return credential.credential_json.get_value(apply_mask=False)\n\n def set_credentials(self, credential_json: dict[str, Any]) -> None:\n with get_session_with_tenant(tenant_id=self._tenant_id) as db_session:\n try:\n credential = db_session.execute(\n select(Credential)\n .where(Credential.id == self._credential_id)\n .with_for_update()\n ).scalar_one()\n\n if credential is None:\n raise ValueError(\n f\"No credential found: credential={self._credential_id}\"\n )\n\n credential.credential_json = credential_json # type: ignore[assignment]\n db_session.commit()\n except Exception:\n db_session.rollback()\n raise\n\n def is_dynamic(self) -> bool:\n return True\n\n\nclass OnyxStaticCredentialsProvider(\n CredentialsProviderInterface[\"OnyxStaticCredentialsProvider\"]\n):\n \"\"\"Implementation (a very simple one!) to handle static credentials.\"\"\"\n\n def __init__(\n self,\n tenant_id: str | None,\n connector_name: str,\n credential_json: dict[str, Any],\n ):\n self._tenant_id = tenant_id\n self._connector_name = connector_name\n self._credential_json = credential_json\n\n self._provider_key = str(uuid.uuid4())\n\n def __enter__(self) -> \"OnyxStaticCredentialsProvider\":\n return self\n\n def __exit__(\n self,\n exc_type: type[BaseException] | None,\n exc_value: BaseException | None,\n traceback: TracebackType | None,\n ) -> None:\n pass\n\n def get_tenant_id(self) -> str | None:\n return self._tenant_id\n\n def get_provider_key(self) -> str:\n return self._provider_key\n\n def get_credentials(self) -> dict[str, Any]:\n return self._credential_json\n\n def set_credentials(self, credential_json: dict[str, Any]) -> None:\n self._credential_json = credential_json\n\n def is_dynamic(self) -> bool:\n return False\n" }, { "path": "backend/onyx/connectors/cross_connector_utils/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/cross_connector_utils/miscellaneous_utils.py", "content": "import re\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import TypeVar\nfrom urllib.parse import urljoin\nfrom urllib.parse import urlparse\n\nimport requests\nfrom dateutil.parser import parse\nfrom dateutil.parser import ParserError\n\nfrom onyx.configs.app_configs import CONNECTOR_LOCALHOST_OVERRIDE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import IGNORE_FOR_QA\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import OnyxMetadata\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.text_processing import is_valid_email\n\n\nT = TypeVar(\"T\")\nU = TypeVar(\"U\")\nlogger = setup_logger()\n\n\ndef datetime_to_utc(dt: datetime) -> datetime:\n if dt.tzinfo is None or dt.tzinfo.utcoffset(dt) is None:\n dt = dt.replace(tzinfo=timezone.utc)\n\n return dt.astimezone(timezone.utc)\n\n\ndef time_str_to_utc(datetime_str: str) -> datetime:\n # Remove all timezone abbreviations in parentheses\n normalized = re.sub(r\"\\([A-Z]+\\)\", \"\", datetime_str).strip()\n\n # Remove any remaining parentheses and their contents\n normalized = re.sub(r\"\\(.*?\\)\", \"\", normalized).strip()\n\n candidates: list[str] = [normalized]\n\n # Some sources (e.g. Gmail) may prefix the value with labels like \"Date:\"\n label_stripped = re.sub(\n r\"^\\s*[A-Za-z][A-Za-z\\s_-]*:\\s*\", \"\", normalized, count=1\n ).strip()\n if label_stripped and label_stripped != normalized:\n candidates.append(label_stripped)\n\n # Fix common format issues (e.g. \"0000\" => \"+0000\")\n for candidate in list(candidates):\n if \" 0000\" in candidate:\n fixed = candidate.replace(\" 0000\", \" +0000\")\n if fixed not in candidates:\n candidates.append(fixed)\n\n last_exception: Exception | None = None\n for candidate in candidates:\n try:\n dt = parse(candidate)\n return datetime_to_utc(dt)\n except (ValueError, ParserError) as exc:\n last_exception = exc\n\n if last_exception is not None:\n raise last_exception\n\n # Fallback in case parsing failed without raising (should not happen)\n raise ValueError(f\"Unable to parse datetime string: {datetime_str}\")\n\n\n# TODO: use this function in other connectors\ndef datetime_from_utc_timestamp(timestamp: int) -> datetime:\n \"\"\"Convert a Unix timestamp to a datetime object in UTC\"\"\"\n\n return datetime.fromtimestamp(timestamp, tz=timezone.utc)\n\n\ndef basic_expert_info_representation(info: BasicExpertInfo) -> str | None:\n if info.first_name and info.last_name:\n return f\"{info.first_name} {info.middle_initial} {info.last_name}\"\n\n if info.display_name:\n return info.display_name\n\n if info.email and is_valid_email(info.email):\n return info.email\n\n if info.first_name:\n return info.first_name\n\n return None\n\n\ndef get_experts_stores_representations(\n experts: list[BasicExpertInfo] | None,\n) -> list[str] | None:\n \"\"\"Gets string representations of experts supplied.\n\n If an expert cannot be represented as a string, it is omitted from the\n result.\n \"\"\"\n if not experts:\n return None\n\n reps: list[str | None] = [\n basic_expert_info_representation(owner) for owner in experts\n ]\n return [owner for owner in reps if owner is not None]\n\n\ndef process_in_batches(\n objects: list[T], process_function: Callable[[T], U], batch_size: int\n) -> Iterator[list[U]]:\n for i in range(0, len(objects), batch_size):\n yield [process_function(obj) for obj in objects[i : i + batch_size]]\n\n\ndef get_metadata_keys_to_ignore() -> list[str]:\n return [IGNORE_FOR_QA]\n\n\ndef _parse_document_source(connector_type: Any) -> DocumentSource | None:\n if connector_type is None:\n return None\n\n if isinstance(connector_type, DocumentSource):\n return connector_type\n\n if not isinstance(connector_type, str):\n logger.warning(f\"Invalid connector_type type: {type(connector_type).__name__}\")\n return None\n\n normalized = re.sub(r\"[\\s\\-]+\", \"_\", connector_type.strip().lower())\n try:\n return DocumentSource(normalized)\n except ValueError:\n logger.warning(\n f\"Invalid connector_type value: '{connector_type}' (normalized: '{normalized}')\"\n )\n return None\n\n\ndef process_onyx_metadata(\n metadata: dict[str, Any],\n) -> tuple[OnyxMetadata, dict[str, Any]]:\n \"\"\"\n Users may set Onyx metadata and custom tags in text files. https://docs.onyx.app/admins/connectors/official/file\n Any unrecognized fields are treated as custom tags.\n \"\"\"\n p_owner_names = metadata.get(\"primary_owners\")\n p_owners = (\n [BasicExpertInfo(display_name=name) for name in p_owner_names]\n if p_owner_names\n else None\n )\n\n s_owner_names = metadata.get(\"secondary_owners\")\n s_owners = (\n [BasicExpertInfo(display_name=name) for name in s_owner_names]\n if s_owner_names\n else None\n )\n source_type = _parse_document_source(metadata.get(\"connector_type\"))\n\n dt_str = metadata.get(\"doc_updated_at\")\n doc_updated_at = time_str_to_utc(dt_str) if dt_str else None\n\n return (\n OnyxMetadata(\n document_id=metadata.get(\"id\"),\n source_type=source_type,\n link=metadata.get(\"link\"),\n file_display_name=metadata.get(\"file_display_name\"),\n title=metadata.get(\"title\"),\n primary_owners=p_owners,\n secondary_owners=s_owners,\n doc_updated_at=doc_updated_at,\n ),\n {\n k: v\n for k, v in metadata.items()\n if k\n not in [\n \"document_id\",\n \"time_updated\",\n \"doc_updated_at\",\n \"link\",\n \"primary_owners\",\n \"secondary_owners\",\n \"filename\",\n \"file_display_name\",\n \"title\",\n \"connector_type\",\n \"pdf_password\",\n \"mime_type\",\n ]\n },\n )\n\n\ndef get_oauth_callback_uri(base_domain: str, connector_id: str) -> str:\n if CONNECTOR_LOCALHOST_OVERRIDE:\n # Used for development\n base_domain = CONNECTOR_LOCALHOST_OVERRIDE\n return f\"{base_domain.strip('/')}/connector/oauth/callback/{connector_id}\"\n\n\ndef is_atlassian_date_error(e: Exception) -> bool:\n return \"field 'updated' is invalid\" in str(e)\n\n\ndef get_cloudId(base_url: str) -> str:\n tenant_info_url = urljoin(base_url, \"/_edge/tenant_info\")\n response = requests.get(tenant_info_url, timeout=10)\n response.raise_for_status()\n return response.json()[\"cloudId\"]\n\n\ndef scoped_url(url: str, product: str) -> str:\n parsed = urlparse(url)\n base_url = parsed.scheme + \"://\" + parsed.netloc\n cloud_id = get_cloudId(base_url)\n return f\"https://api.atlassian.com/ex/{product}/{cloud_id}{parsed.path}\"\n" }, { "path": "backend/onyx/connectors/cross_connector_utils/rate_limit_wrapper.py", "content": "import time\nfrom collections.abc import Callable\nfrom functools import wraps\nfrom typing import Any\nfrom typing import cast\nfrom typing import TypeVar\n\nimport requests\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nF = TypeVar(\"F\", bound=Callable[..., Any])\n\n\nclass RateLimitTriedTooManyTimesError(Exception):\n pass\n\n\nclass _RateLimitDecorator:\n \"\"\"Builds a generic wrapper/decorator for calls to external APIs that\n prevents making more than `max_calls` requests per `period`\n\n Implementation inspired by the `ratelimit` library:\n https://github.com/tomasbasham/ratelimit.\n\n NOTE: is not thread safe.\n \"\"\"\n\n def __init__(\n self,\n max_calls: int,\n period: float, # in seconds\n sleep_time: float = 2, # in seconds\n sleep_backoff: float = 2, # applies exponential backoff\n max_num_sleep: int = 0,\n ):\n self.max_calls = max_calls\n self.period = period\n self.sleep_time = sleep_time\n self.sleep_backoff = sleep_backoff\n self.max_num_sleep = max_num_sleep\n\n self.call_history: list[float] = []\n self.curr_calls = 0\n\n def __call__(self, func: F) -> F:\n @wraps(func)\n def wrapped_func(*args: list, **kwargs: dict[str, Any]) -> Any:\n # cleanup calls which are no longer relevant\n self._cleanup()\n\n # check if we've exceeded the rate limit\n sleep_cnt = 0\n while len(self.call_history) == self.max_calls:\n sleep_time = self.sleep_time * (self.sleep_backoff**sleep_cnt)\n logger.notice(\n f\"Rate limit exceeded for function {func.__name__}. Waiting {sleep_time} seconds before retrying.\"\n )\n time.sleep(sleep_time)\n sleep_cnt += 1\n if self.max_num_sleep != 0 and sleep_cnt >= self.max_num_sleep:\n raise RateLimitTriedTooManyTimesError(\n f\"Exceeded '{self.max_num_sleep}' retries for function '{func.__name__}'\"\n )\n\n self._cleanup()\n\n # add the current call to the call history\n self.call_history.append(time.monotonic())\n return func(*args, **kwargs)\n\n return cast(F, wrapped_func)\n\n def _cleanup(self) -> None:\n curr_time = time.monotonic()\n time_to_expire_before = curr_time - self.period\n self.call_history = [\n call_time\n for call_time in self.call_history\n if call_time > time_to_expire_before\n ]\n\n\nrate_limit_builder = _RateLimitDecorator\n\n\n\"\"\"If you want to allow the external service to tell you when you've hit the rate limit,\nuse the following instead\"\"\"\n\nR = TypeVar(\"R\", bound=Callable[..., requests.Response])\n\n\ndef wrap_request_to_handle_ratelimiting(\n request_fn: R, default_wait_time_sec: int = 30, max_waits: int = 30\n) -> R:\n def wrapped_request(*args: list, **kwargs: dict[str, Any]) -> requests.Response:\n for _ in range(max_waits):\n response = request_fn(*args, **kwargs)\n if response.status_code == 429:\n try:\n wait_time = int(\n response.headers.get(\"Retry-After\", default_wait_time_sec)\n )\n except ValueError:\n wait_time = default_wait_time_sec\n\n time.sleep(wait_time)\n continue\n\n return response\n\n raise RateLimitTriedTooManyTimesError(f\"Exceeded '{max_waits}' retries\")\n\n return cast(R, wrapped_request)\n\n\n_rate_limited_get = wrap_request_to_handle_ratelimiting(requests.get)\n_rate_limited_post = wrap_request_to_handle_ratelimiting(requests.post)\n\n\nclass _RateLimitedRequest:\n get = _rate_limited_get\n post = _rate_limited_post\n\n\nrl_requests = _RateLimitedRequest\n" }, { "path": "backend/onyx/connectors/discord/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/discord/connector.py", "content": "import asyncio\nfrom collections.abc import AsyncGenerator\nfrom collections.abc import AsyncIterable\nfrom collections.abc import Iterable\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nfrom discord import Client\nfrom discord.channel import TextChannel\nfrom discord.channel import Thread\nfrom discord.enums import MessageType\nfrom discord.errors import LoginFailure\nfrom discord.flags import Intents\nfrom discord.message import Message as DiscordMessage\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import CredentialInvalidError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n_DISCORD_DOC_ID_PREFIX = \"DISCORD_\"\n_SNIPPET_LENGTH = 30\n\n\ndef _convert_message_to_document(\n message: DiscordMessage,\n sections: list[TextSection],\n) -> Document:\n \"\"\"\n Convert a discord message to a document\n Sections are collected before calling this function because it relies on async\n calls to fetch the thread history if there is one\n \"\"\"\n\n metadata: dict[str, str | list[str]] = {}\n semantic_substring = \"\"\n\n # Only messages from TextChannels will make it here but we have to check for it anyways\n if isinstance(message.channel, TextChannel) and (\n channel_name := message.channel.name\n ):\n metadata[\"Channel\"] = channel_name\n semantic_substring += f\" in Channel: #{channel_name}\"\n\n # Single messages dont have a title\n title = \"\"\n\n # If there is a thread, add more detail to the metadata, title, and semantic identifier\n if isinstance(message.channel, Thread):\n # Threads do have a title\n title = message.channel.name\n\n # If its a thread, update the metadata, title, and semantic_substring\n metadata[\"Thread\"] = title\n\n # Add more detail to the semantic identifier if available\n semantic_substring += f\" in Thread: {title}\"\n\n snippet: str = (\n message.content[:_SNIPPET_LENGTH].rstrip() + \"...\"\n if len(message.content) > _SNIPPET_LENGTH\n else message.content\n )\n\n semantic_identifier = f\"{message.author.name} said{semantic_substring}: {snippet}\"\n\n return Document(\n id=f\"{_DISCORD_DOC_ID_PREFIX}{message.id}\",\n source=DocumentSource.DISCORD,\n semantic_identifier=semantic_identifier,\n doc_updated_at=message.edited_at,\n title=title,\n sections=(cast(list[TextSection | ImageSection], sections)),\n metadata=metadata,\n )\n\n\nasync def _fetch_filtered_channels(\n discord_client: Client,\n server_ids: list[int] | None,\n channel_names: list[str] | None,\n) -> list[TextChannel]:\n filtered_channels: list[TextChannel] = []\n\n for channel in discord_client.get_all_channels():\n if not channel.permissions_for(channel.guild.me).read_message_history:\n continue\n if not isinstance(channel, TextChannel):\n continue\n if server_ids and len(server_ids) > 0 and channel.guild.id not in server_ids:\n continue\n if channel_names and channel.name not in channel_names:\n continue\n filtered_channels.append(channel)\n\n logger.info(f\"Found {len(filtered_channels)} channels for the authenticated user\")\n return filtered_channels\n\n\nasync def _fetch_documents_from_channel(\n channel: TextChannel,\n start_time: datetime | None,\n end_time: datetime | None,\n) -> AsyncIterable[Document]:\n # Discord's epoch starts at 2015-01-01\n discord_epoch = datetime(2015, 1, 1, tzinfo=timezone.utc)\n if start_time and start_time < discord_epoch:\n start_time = discord_epoch\n\n # NOTE: limit=None is the correct way to fetch all messages and threads with pagination\n # The discord package erroneously uses limit for both pagination AND number of results\n # This causes the history and archived_threads methods to return 100 results even if there are more results within the filters\n # Pagination is handled automatically (100 results at a time) when limit=None\n\n async for channel_message in channel.history(\n limit=None,\n after=start_time,\n before=end_time,\n ):\n # Skip messages that are not the default type\n if channel_message.type != MessageType.default:\n continue\n\n sections: list[TextSection] = [\n TextSection(\n text=channel_message.content,\n link=channel_message.jump_url,\n )\n ]\n\n yield _convert_message_to_document(channel_message, sections)\n\n for active_thread in channel.threads:\n async for thread_message in active_thread.history(\n limit=None,\n after=start_time,\n before=end_time,\n ):\n # Skip messages that are not the default type\n if thread_message.type != MessageType.default:\n continue\n\n sections = [\n TextSection(\n text=thread_message.content,\n link=thread_message.jump_url,\n )\n ]\n\n yield _convert_message_to_document(thread_message, sections)\n\n async for archived_thread in channel.archived_threads(\n limit=None,\n ):\n async for thread_message in archived_thread.history(\n limit=None,\n after=start_time,\n before=end_time,\n ):\n # Skip messages that are not the default type\n if thread_message.type != MessageType.default:\n continue\n\n sections = [\n TextSection(\n text=thread_message.content,\n link=thread_message.jump_url,\n )\n ]\n\n yield _convert_message_to_document(thread_message, sections)\n\n\ndef _manage_async_retrieval(\n token: str,\n requested_start_date_string: str,\n channel_names: list[str],\n server_ids: list[int],\n start: datetime | None = None,\n end: datetime | None = None,\n) -> Iterable[Document]:\n # parse requested_start_date_string to datetime\n pull_date: datetime | None = (\n datetime.strptime(requested_start_date_string, \"%Y-%m-%d\").replace(\n tzinfo=timezone.utc\n )\n if requested_start_date_string\n else None\n )\n\n # Set start_time to the later of start and pull_date, or whichever is provided\n start_time = max(filter(None, [start, pull_date])) if start or pull_date else None\n\n end_time: datetime | None = end\n\n async def _async_fetch() -> AsyncGenerator[Document, None]:\n intents = Intents.default()\n intents.message_content = True\n async with Client(intents=intents) as discord_client:\n start_task = asyncio.create_task(discord_client.start(token))\n ready_task = asyncio.create_task(discord_client.wait_until_ready())\n\n done, _ = await asyncio.wait(\n {start_task, ready_task},\n return_when=asyncio.FIRST_COMPLETED,\n )\n\n # start() runs indefinitely once connected, so it only lands\n # in `done` when login/connection failed — propagate the error.\n if start_task in done:\n ready_task.cancel()\n start_task.result()\n\n filtered_channels: list[TextChannel] = await _fetch_filtered_channels(\n discord_client=discord_client,\n server_ids=server_ids,\n channel_names=channel_names,\n )\n\n for channel in filtered_channels:\n async for doc in _fetch_documents_from_channel(\n channel=channel,\n start_time=start_time,\n end_time=end_time,\n ):\n yield doc\n\n def run_and_yield() -> Iterable[Document]:\n loop = asyncio.new_event_loop()\n async_gen = _async_fetch()\n try:\n while True:\n try:\n doc = loop.run_until_complete(anext(async_gen))\n yield doc\n except StopAsyncIteration:\n break\n finally:\n # Must close the async generator before the loop so the Discord\n # client's `async with` block can await its shutdown coroutine.\n # The nested try/finally ensures the loop always closes even if\n # aclose() raises (same pattern as cursor.close() before conn.close()).\n try:\n loop.run_until_complete(async_gen.aclose())\n finally:\n loop.close()\n\n return run_and_yield()\n\n\nclass DiscordConnector(PollConnector, LoadConnector):\n def __init__(\n self,\n server_ids: list[str] = [],\n channel_names: list[str] = [],\n # YYYY-MM-DD\n start_date: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ):\n self.batch_size = batch_size\n self.channel_names: list[str] = channel_names if channel_names else []\n self.server_ids: list[int] = (\n [int(server_id) for server_id in server_ids] if server_ids else []\n )\n self._discord_bot_token: str | None = None\n self.requested_start_date_string: str = start_date or \"\"\n\n @property\n def discord_bot_token(self) -> str:\n if self._discord_bot_token is None:\n raise ConnectorMissingCredentialError(\"Discord\")\n return self._discord_bot_token\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self._discord_bot_token = credentials[\"discord_bot_token\"]\n return None\n\n def validate_connector_settings(self) -> None:\n loop = asyncio.new_event_loop()\n try:\n client = Client(intents=Intents.default())\n try:\n loop.run_until_complete(client.login(self.discord_bot_token))\n except LoginFailure as e:\n raise CredentialInvalidError(f\"Invalid Discord bot token: {e}\")\n finally:\n loop.run_until_complete(client.close())\n finally:\n loop.close()\n\n def _manage_doc_batching(\n self,\n start: datetime | None = None,\n end: datetime | None = None,\n ) -> GenerateDocumentsOutput:\n doc_batch: list[Document | HierarchyNode] = []\n for doc in _manage_async_retrieval(\n token=self.discord_bot_token,\n requested_start_date_string=self.requested_start_date_string,\n channel_names=self.channel_names,\n server_ids=self.server_ids,\n start=start,\n end=end,\n ):\n doc_batch.append(doc)\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n return self._manage_doc_batching(\n datetime.fromtimestamp(start, tz=timezone.utc),\n datetime.fromtimestamp(end, tz=timezone.utc),\n )\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._manage_doc_batching(None, None)\n\n\nif __name__ == \"__main__\":\n import os\n import time\n\n end = time.time()\n # 1 day\n start = end - 24 * 60 * 60 * 1\n # \"1,2,3\"\n server_ids: str | None = os.environ.get(\"server_ids\", None)\n # \"channel1,channel2\"\n channel_names: str | None = os.environ.get(\"channel_names\", None)\n\n connector = DiscordConnector(\n server_ids=server_ids.split(\",\") if server_ids else [],\n channel_names=channel_names.split(\",\") if channel_names else [],\n start_date=os.environ.get(\"start_date\", None),\n )\n connector.load_credentials(\n {\"discord_bot_token\": os.environ.get(\"discord_bot_token\")}\n )\n\n for doc_batch in connector.poll_source(start, end):\n for doc in doc_batch:\n print(doc)\n" }, { "path": "backend/onyx/connectors/discourse/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/discourse/connector.py", "content": "import time\nimport urllib.parse\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom pydantic import BaseModel\nfrom requests import Response\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\nlogger = setup_logger()\n\n\nclass DiscoursePerms(BaseModel):\n api_key: str\n api_username: str\n\n\n@retry_builder()\ndef discourse_request(\n endpoint: str, perms: DiscoursePerms, params: dict | None = None\n) -> Response:\n headers = {\"Api-Key\": perms.api_key, \"Api-Username\": perms.api_username}\n\n response = requests.get(endpoint, headers=headers, params=params)\n response.raise_for_status()\n\n return response\n\n\nclass DiscourseConnector(PollConnector):\n def __init__(\n self,\n base_url: str,\n categories: list[str] | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n parsed_url = urllib.parse.urlparse(base_url)\n if not parsed_url.scheme:\n base_url = \"https://\" + base_url\n self.base_url = base_url\n\n self.categories = [c.lower() for c in categories] if categories else []\n self.category_id_map: dict[int, dict] = {}\n\n self.batch_size = batch_size\n self.permissions: DiscoursePerms | None = None\n self.active_categories: set | None = None\n\n @rate_limit_builder(max_calls=50, period=60)\n def _make_request(self, endpoint: str, params: dict | None = None) -> Response:\n if not self.permissions:\n raise ConnectorMissingCredentialError(\"Discourse\")\n return discourse_request(endpoint, self.permissions, params)\n\n def _get_categories_map(\n self,\n ) -> None:\n assert self.permissions is not None\n categories_endpoint = urllib.parse.urljoin(self.base_url, \"categories.json\")\n response = self._make_request(\n endpoint=categories_endpoint,\n params={\"include_subcategories\": True},\n )\n categories = response.json()[\"category_list\"][\"categories\"]\n self.category_id_map = {\n cat[\"id\"]: {\"name\": cat[\"name\"], \"slug\": cat[\"slug\"]}\n for cat in categories\n if not self.categories or cat[\"name\"].lower() in self.categories\n }\n self.active_categories = set(self.category_id_map)\n\n def _get_doc_from_topic(self, topic_id: int) -> Document:\n assert self.permissions is not None\n topic_endpoint = urllib.parse.urljoin(self.base_url, f\"t/{topic_id}.json\")\n response = self._make_request(endpoint=topic_endpoint)\n topic = response.json()\n\n topic_url = urllib.parse.urljoin(self.base_url, f\"t/{topic['slug']}\")\n\n sections = []\n poster = None\n responders = []\n seen_names = set()\n for ind, post in enumerate(topic[\"post_stream\"][\"posts\"]):\n if ind == 0:\n poster_name = post.get(\"name\")\n if poster_name:\n seen_names.add(poster_name)\n poster = BasicExpertInfo(display_name=poster_name)\n else:\n responder_name = post.get(\"name\")\n if responder_name and responder_name not in seen_names:\n seen_names.add(responder_name)\n responders.append(BasicExpertInfo(display_name=responder_name))\n\n sections.append(\n TextSection(link=topic_url, text=parse_html_page_basic(post[\"cooked\"]))\n )\n category_name = self.category_id_map.get(topic[\"category_id\"], {}).get(\"name\")\n\n metadata: dict[str, str | list[str]] = (\n {\n \"category\": category_name,\n }\n if category_name\n else {}\n )\n\n if topic.get(\"tags\"):\n metadata[\"tags\"] = topic[\"tags\"]\n\n doc = Document(\n id=\"_\".join([DocumentSource.DISCOURSE.value, str(topic[\"id\"])]),\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.DISCOURSE,\n semantic_identifier=topic[\"title\"],\n doc_updated_at=time_str_to_utc(topic[\"last_posted_at\"]),\n primary_owners=[poster] if poster else None,\n secondary_owners=responders or None,\n metadata=metadata,\n )\n return doc\n\n def _get_latest_topics(\n self, start: datetime | None, end: datetime | None, page: int\n ) -> list[int]:\n assert self.permissions is not None\n topic_ids = []\n\n if not self.categories:\n latest_endpoint = urllib.parse.urljoin(\n self.base_url, f\"latest.json?page={page}\"\n )\n response = self._make_request(endpoint=latest_endpoint)\n topics = response.json()[\"topic_list\"][\"topics\"]\n\n else:\n topics = []\n empty_categories = []\n\n for category_id, category_dict in self.category_id_map.items():\n category_endpoint = urllib.parse.urljoin(\n self.base_url,\n f\"c/{category_dict['slug']}/{category_id}.json?page={page}&sys=latest\",\n )\n response = self._make_request(endpoint=category_endpoint)\n new_topics = response.json()[\"topic_list\"][\"topics\"]\n\n if len(new_topics) == 0:\n empty_categories.append(category_id)\n topics.extend(new_topics)\n\n for empty_category in empty_categories:\n self.category_id_map.pop(empty_category)\n\n for topic in topics:\n last_time = topic.get(\"last_posted_at\")\n if not last_time:\n continue\n\n last_time_dt = time_str_to_utc(last_time)\n if (start and start > last_time_dt) or (end and end < last_time_dt):\n continue\n\n topic_ids.append(topic[\"id\"])\n\n return topic_ids\n\n def _yield_discourse_documents(\n self,\n start: datetime,\n end: datetime,\n ) -> GenerateDocumentsOutput:\n page = 0\n while topic_ids := self._get_latest_topics(start, end, page):\n doc_batch: list[Document | HierarchyNode] = []\n for topic_id in topic_ids:\n doc_batch.append(self._get_doc_from_topic(topic_id))\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n page += 1\n\n def load_credentials(\n self,\n credentials: dict[str, Any],\n ) -> dict[str, Any] | None:\n self.permissions = DiscoursePerms(\n api_key=credentials[\"discourse_api_key\"],\n api_username=credentials[\"discourse_api_username\"],\n )\n return None\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n if self.permissions is None:\n raise ConnectorMissingCredentialError(\"Discourse\")\n\n start_datetime = datetime.utcfromtimestamp(start).replace(tzinfo=timezone.utc)\n end_datetime = datetime.utcfromtimestamp(end).replace(tzinfo=timezone.utc)\n\n self._get_categories_map()\n\n yield from self._yield_discourse_documents(start_datetime, end_datetime)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = DiscourseConnector(base_url=os.environ[\"DISCOURSE_BASE_URL\"])\n connector.load_credentials(\n {\n \"discourse_api_key\": os.environ[\"DISCOURSE_API_KEY\"],\n \"discourse_api_username\": os.environ[\"DISCOURSE_API_USERNAME\"],\n }\n )\n\n current = time.time()\n one_year_ago = current - 24 * 60 * 60 * 360\n latest_docs = connector.poll_source(one_year_ago, current)\n print(next(latest_docs))\n" }, { "path": "backend/onyx/connectors/document360/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/document360/connector.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import List\nfrom typing import Optional\n\nimport requests\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.document360.utils import flatten_child_categories\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.utils.retry_wrapper import retry_builder\n\n# Limitations and Potential Improvements\n# 1. The \"Categories themselves contain potentially relevant information\" but they're not pulled in\n# 2. Only the HTML Articles are supported, Document360 also has a Markdown and \"Block\" format\n# 3. The contents are not as cleaned up as other HTML connectors\n\nDOCUMENT360_BASE_URL = \"https://portal.document360.io\"\nDOCUMENT360_API_BASE_URL = \"https://apihub.document360.io/v2\"\n\n\nclass Document360Connector(LoadConnector, PollConnector):\n def __init__(\n self,\n workspace: str,\n categories: List[str] | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n portal_id: Optional[str] = None,\n api_token: Optional[str] = None,\n ) -> None:\n self.portal_id = portal_id\n self.workspace = workspace\n self.categories = categories\n self.batch_size = batch_size\n self.api_token = api_token\n\n def load_credentials(self, credentials: dict[str, Any]) -> Optional[dict[str, Any]]:\n self.api_token = credentials.get(\"document360_api_token\")\n self.portal_id = credentials.get(\"portal_id\")\n return None\n\n # rate limiting set based on the enterprise plan: https://apidocs.document360.com/apidocs/rate-limiting\n # NOTE: retry will handle cases where user is not on enterprise plan - we will just hit the rate limit\n # and then retry after a period\n @retry_builder()\n @rate_limit_builder(max_calls=100, period=60)\n def _make_request(self, endpoint: str, params: Optional[dict] = None) -> Any:\n if not self.api_token:\n raise ConnectorMissingCredentialError(\"Document360\")\n\n headers = {\"accept\": \"application/json\", \"api_token\": self.api_token}\n\n response = requests.get(\n f\"{DOCUMENT360_API_BASE_URL}/{endpoint}\", headers=headers, params=params\n )\n response.raise_for_status()\n\n return response.json()[\"data\"]\n\n def _get_workspace_id_by_name(self) -> str:\n projects = self._make_request(\"ProjectVersions\")\n workspace_id = next(\n (\n project[\"id\"]\n for project in projects\n if project[\"version_code_name\"] == self.workspace\n ),\n None,\n )\n if workspace_id is None:\n raise ValueError(\"Not able to find Workspace ID by the user provided name\")\n\n return workspace_id\n\n def _get_articles_with_category(self, workspace_id: str) -> Any:\n all_categories = self._make_request(\n f\"ProjectVersions/{workspace_id}/categories\"\n )\n articles_with_category = []\n\n for category in all_categories:\n if not self.categories or category[\"name\"] in self.categories:\n for article in category[\"articles\"]:\n articles_with_category.append(\n {\"id\": article[\"id\"], \"category_name\": category[\"name\"]}\n )\n for child_category in category[\"child_categories\"]:\n all_nested_categories = flatten_child_categories(child_category)\n for nested_category in all_nested_categories:\n for article in nested_category[\"articles\"]:\n articles_with_category.append(\n {\n \"id\": article[\"id\"],\n \"category_name\": nested_category[\"name\"],\n }\n )\n\n return articles_with_category\n\n def _process_articles(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n if self.api_token is None:\n raise ConnectorMissingCredentialError(\"Document360\")\n\n workspace_id = self._get_workspace_id_by_name()\n articles = self._get_articles_with_category(workspace_id)\n\n doc_batch: List[Document | HierarchyNode] = []\n\n for article in articles:\n article_details = self._make_request(\n f\"Articles/{article['id']}\", {\"langCode\": \"en\"}\n )\n\n updated_at = datetime.strptime(\n article_details[\"modified_at\"], \"%Y-%m-%dT%H:%M:%S.%fZ\"\n ).replace(tzinfo=timezone.utc)\n if start is not None and updated_at < start:\n continue\n if end is not None and updated_at > end:\n continue\n\n authors = [\n BasicExpertInfo(\n display_name=author.get(\"name\"), email=author[\"email_id\"]\n )\n for author in article_details.get(\"authors\", [])\n if author[\"email_id\"]\n ]\n\n doc_link = (\n article_details[\"url\"]\n if article_details.get(\"url\")\n else f\"{DOCUMENT360_BASE_URL}/{self.portal_id}/document/v1/view/{article['id']}\"\n )\n\n html_content = article_details[\"html_content\"]\n article_content = (\n parse_html_page_basic(html_content) if html_content is not None else \"\"\n )\n doc_text = (\n f\"{article_details.get('description', '')}\\n{article_content}\".strip()\n )\n\n document = Document(\n id=article_details[\"id\"],\n sections=[TextSection(link=doc_link, text=doc_text)],\n source=DocumentSource.DOCUMENT360,\n semantic_identifier=article_details[\"title\"],\n doc_updated_at=updated_at,\n primary_owners=authors,\n metadata={\n \"workspace\": self.workspace,\n \"category\": article[\"category_name\"],\n },\n )\n\n doc_batch.append(document)\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._process_articles()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n return self._process_articles(start_datetime, end_datetime)\n\n\nif __name__ == \"__main__\":\n import time\n import os\n\n document360_connector = Document360Connector(os.environ[\"DOCUMENT360_WORKSPACE\"])\n document360_connector.load_credentials(\n {\n \"portal_id\": os.environ[\"DOCUMENT360_PORTAL_ID\"],\n \"document360_api_token\": os.environ[\"DOCUMENT360_API_TOKEN\"],\n }\n )\n\n current = time.time()\n one_year_ago = current - 24 * 60 * 60 * 360\n latest_docs = document360_connector.poll_source(one_year_ago, current)\n\n for doc in latest_docs:\n print(doc)\n" }, { "path": "backend/onyx/connectors/document360/utils.py", "content": "def flatten_child_categories(category: dict) -> list[dict]:\n if not category[\"child_categories\"]:\n return [category]\n else:\n flattened_categories = [category]\n for child_category in category[\"child_categories\"]:\n flattened_categories.extend(flatten_child_categories(child_category))\n return flattened_categories\n" }, { "path": "backend/onyx/connectors/dropbox/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/dropbox/connector.py", "content": "from datetime import timezone\nfrom io import BytesIO\nfrom typing import Any\n\nfrom dropbox import Dropbox # type: ignore[import-untyped]\nfrom dropbox.exceptions import ApiError # type: ignore[import-untyped]\nfrom dropbox.exceptions import AuthError\nfrom dropbox.files import FileMetadata # type: ignore[import-untyped]\nfrom dropbox.files import FolderMetadata\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialInvalidError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\nclass DropboxConnector(LoadConnector, PollConnector):\n def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:\n self.batch_size = batch_size\n self.dropbox_client: Dropbox | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.dropbox_client = Dropbox(credentials[\"dropbox_access_token\"])\n return None\n\n def _download_file(self, path: str) -> bytes:\n \"\"\"Download a single file from Dropbox.\"\"\"\n if self.dropbox_client is None:\n raise ConnectorMissingCredentialError(\"Dropbox\")\n _, resp = self.dropbox_client.files_download(path)\n return resp.content\n\n def _get_shared_link(self, path: str) -> str:\n \"\"\"Create a shared link for a file in Dropbox.\"\"\"\n if self.dropbox_client is None:\n raise ConnectorMissingCredentialError(\"Dropbox\")\n\n try:\n # Check if a shared link already exists\n shared_links = self.dropbox_client.sharing_list_shared_links(path=path)\n if shared_links.links:\n return shared_links.links[0].url\n\n link_metadata = (\n self.dropbox_client.sharing_create_shared_link_with_settings(path)\n )\n return link_metadata.url\n except ApiError as err:\n logger.exception(f\"Failed to create a shared link for {path}: {err}\")\n return \"\"\n\n def _yield_files_recursive(\n self,\n path: str,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n ) -> GenerateDocumentsOutput:\n \"\"\"Yield files in batches from a specified Dropbox folder, including subfolders.\"\"\"\n if self.dropbox_client is None:\n raise ConnectorMissingCredentialError(\"Dropbox\")\n\n result = self.dropbox_client.files_list_folder(\n path,\n limit=self.batch_size,\n recursive=False,\n include_non_downloadable_files=False,\n )\n\n while True:\n batch: list[Document | HierarchyNode] = []\n for entry in result.entries:\n if isinstance(entry, FileMetadata):\n modified_time = entry.client_modified\n if modified_time.tzinfo is None:\n # If no timezone info, assume it is UTC\n modified_time = modified_time.replace(tzinfo=timezone.utc)\n else:\n # If not in UTC, translate it\n modified_time = modified_time.astimezone(timezone.utc)\n\n time_as_seconds = int(modified_time.timestamp())\n if start and time_as_seconds < start:\n continue\n if end and time_as_seconds > end:\n continue\n\n downloaded_file = self._download_file(entry.path_display)\n link = self._get_shared_link(entry.path_display)\n try:\n text = extract_file_text(\n BytesIO(downloaded_file),\n file_name=entry.name,\n break_on_unprocessable=False,\n )\n batch.append(\n Document(\n id=f\"doc:{entry.id}\",\n sections=[TextSection(link=link, text=text)],\n source=DocumentSource.DROPBOX,\n semantic_identifier=entry.name,\n doc_updated_at=modified_time,\n metadata={\"type\": \"article\"},\n )\n )\n except Exception as e:\n logger.exception(\n f\"Error decoding file {entry.path_display} as utf-8 error occurred: {e}\"\n )\n\n elif isinstance(entry, FolderMetadata):\n yield from self._yield_files_recursive(entry.path_lower, start, end)\n\n if batch:\n yield batch\n\n if not result.has_more:\n break\n\n result = self.dropbox_client.files_list_folder_continue(result.cursor)\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self.poll_source(None, None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch | None, end: SecondsSinceUnixEpoch | None\n ) -> GenerateDocumentsOutput:\n if self.dropbox_client is None:\n raise ConnectorMissingCredentialError(\"Dropbox\")\n\n for batch in self._yield_files_recursive(\"\", start, end):\n yield batch\n\n return None\n\n def validate_connector_settings(self) -> None:\n if self.dropbox_client is None:\n raise ConnectorMissingCredentialError(\"Dropbox credentials not loaded.\")\n\n try:\n self.dropbox_client.files_list_folder(path=\"\", limit=1)\n except AuthError as e:\n logger.exception(\"Failed to validate Dropbox credentials\")\n raise CredentialInvalidError(f\"Dropbox credential is invalid: {e.error}\")\n except ApiError as e:\n if (\n e.error is not None\n and \"insufficient_permissions\" in str(e.error).lower()\n ):\n raise InsufficientPermissionsError(\n \"Your Dropbox token does not have sufficient permissions.\"\n )\n raise ConnectorValidationError(\n f\"Unexpected Dropbox error during validation: {e.user_message_text or e}\"\n )\n except Exception as e:\n raise Exception(f\"Unexpected error during Dropbox settings validation: {e}\")\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = DropboxConnector()\n connector.load_credentials(\n {\n \"dropbox_access_token\": os.environ[\"DROPBOX_ACCESS_TOKEN\"],\n }\n )\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/drupal_wiki/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/drupal_wiki/connector.py", "content": "import mimetypes\nfrom io import BytesIO\nfrom typing import Any\n\nimport requests\nfrom typing_extensions import override\n\nfrom onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE\nfrom onyx.configs.app_configs import DRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n datetime_from_utc_timestamp,\n)\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import rate_limit_builder\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import rl_requests\nfrom onyx.connectors.drupal_wiki.models import DrupalWikiCheckpoint\nfrom onyx.connectors.drupal_wiki.models import DrupalWikiPage\nfrom onyx.connectors.drupal_wiki.models import DrupalWikiPageResponse\nfrom onyx.connectors.drupal_wiki.models import DrupalWikiSpaceResponse\nfrom onyx.connectors.drupal_wiki.utils import build_drupal_wiki_document_id\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import ConnectorFailure\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnector\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_text_and_images\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.file_processing.image_utils import store_image_and_create_section\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.b64 import get_image_type_from_bytes\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\nlogger = setup_logger()\n\nMAX_API_PAGE_SIZE = 2000 # max allowed by API\nDRUPAL_WIKI_SPACE_KEY = \"space\"\n\n\nrate_limited_get = retry_builder()(\n rate_limit_builder(max_calls=10, period=1)(rl_requests.get)\n)\n\n\nclass DrupalWikiConnector(\n CheckpointedConnector[DrupalWikiCheckpoint],\n SlimConnector,\n):\n # Deprecated parameters that may exist in old connector configurations\n _DEPRECATED_PARAMS = {\"drupal_wiki_scope\", \"include_all_spaces\"}\n\n def __init__(\n self,\n base_url: str,\n spaces: list[str] | None = None,\n pages: list[str] | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n continue_on_failure: bool = CONTINUE_ON_CONNECTOR_FAILURE,\n include_attachments: bool = False,\n allow_images: bool = False,\n **kwargs: Any,\n ) -> None:\n \"\"\"\n Initialize the Drupal Wiki connector.\n\n Args:\n base_url: The base URL of the Drupal Wiki instance (e.g., https://help.drupal-wiki.com)\n spaces: List of space IDs to index. If None and pages is also None, all spaces will be indexed.\n pages: List of page IDs to index. If provided, these specific pages will be indexed.\n batch_size: Number of documents to process in a batch.\n continue_on_failure: If True, continue indexing even if some documents fail.\n include_attachments: If True, enable processing of page attachments including images and documents.\n allow_images: If True, enable processing of image attachments.\n \"\"\"\n\n #########################################################\n # TODO: Remove this after 02/01/2026 and remove **kwargs from the function signature\n # Check for deprecated parameters from old connector configurations\n # If attempting to update without deleting the connector:\n # Remove the deprecated parameters from the custom_connector_config in the relevant connector table rows\n deprecated_found = set(kwargs.keys()) & self._DEPRECATED_PARAMS\n if deprecated_found:\n raise ConnectorValidationError(\n f\"Outdated Drupal Wiki connector configuration detected \"\n f\"(found deprecated parameters: {', '.join(deprecated_found)}). \"\n f\"Please delete and recreate this connector, or contact Onyx support \"\n f\"for assistance with updating the configuration without deleting the connector.\"\n )\n # Reject any other unexpected parameters\n if kwargs:\n raise ConnectorValidationError(\n f\"Unexpected parameters for Drupal Wiki connector: {', '.join(kwargs.keys())}\"\n )\n #########################################################\n\n self.base_url = base_url.rstrip(\"/\")\n self.spaces = spaces or []\n self.pages = pages or []\n\n # If no specific spaces or pages are provided, index all spaces\n self.include_all_spaces = not self.spaces and not self.pages\n\n self.batch_size = batch_size\n self.continue_on_failure = continue_on_failure\n\n # Attachment processing configuration\n self.include_attachments = include_attachments\n self.allow_images = allow_images\n\n self.headers: dict[str, str] = {\"Accept\": \"application/json\"}\n self._api_token: str | None = None # set by load_credentials\n\n def set_allow_images(self, value: bool) -> None:\n logger.info(f\"Setting allow_images to {value}.\")\n self.allow_images = value\n\n def _get_page_attachments(self, page_id: int) -> list[dict[str, Any]]:\n \"\"\"\n Get all attachments for a specific page.\n\n Args:\n page_id: ID of the page.\n\n Returns:\n List of attachment dictionaries.\n \"\"\"\n url = f\"{self.base_url}/api/rest/scope/api/attachment\"\n params = {\"pageId\": str(page_id)}\n logger.debug(f\"Fetching attachments for page {page_id} from {url}\")\n\n try:\n response = rate_limited_get(url, headers=self.headers, params=params)\n response.raise_for_status()\n attachments = response.json()\n logger.info(f\"Found {len(attachments)} attachments for page {page_id}\")\n return attachments\n except Exception as e:\n logger.warning(f\"Failed to fetch attachments for page {page_id}: {e}\")\n return []\n\n def _download_attachment(self, attachment_id: int) -> bytes:\n \"\"\"\n Download attachment content.\n\n Args:\n attachment_id: ID of the attachment to download.\n\n Returns:\n Raw bytes of the attachment.\n \"\"\"\n url = f\"{self.base_url}/api/rest/scope/api/attachment/{attachment_id}/download\"\n logger.info(f\"Downloading attachment {attachment_id} from {url}\")\n\n # Use headers without Accept for binary downloads\n download_headers = {\"Authorization\": f\"Bearer {self._api_token}\"}\n\n response = rate_limited_get(url, headers=download_headers)\n response.raise_for_status()\n\n return response.content\n\n def _validate_attachment_filetype(self, attachment: dict[str, Any]) -> bool:\n \"\"\"\n Validate if the attachment file type is supported.\n\n Args:\n attachment: Attachment dictionary from Drupal Wiki API.\n\n Returns:\n True if the file type is supported, False otherwise.\n \"\"\"\n file_name = attachment.get(\"fileName\", \"\")\n if not file_name:\n return False\n\n # Get file extension\n file_extension = get_file_ext(file_name)\n\n if file_extension in OnyxFileExtensions.ALL_ALLOWED_EXTENSIONS:\n return True\n\n logger.warning(f\"Unsupported file type: {file_extension} for {file_name}\")\n return False\n\n def _get_media_type_from_filename(self, filename: str) -> str:\n \"\"\"\n Get media type from filename using the standard mimetypes library.\n\n Args:\n filename: The filename.\n\n Returns:\n Media type string.\n \"\"\"\n mime_type, _encoding = mimetypes.guess_type(filename)\n return mime_type or \"application/octet-stream\"\n\n def _process_attachment(\n self,\n attachment: dict[str, Any],\n page_id: int,\n download_url: str,\n ) -> tuple[list[TextSection | ImageSection], str | None]:\n \"\"\"\n Process a single attachment and return generated sections.\n\n Args:\n attachment: Attachment dictionary from Drupal Wiki API.\n page_id: ID of the parent page.\n download_url: Direct download URL for the attachment.\n\n Returns:\n Tuple of (sections, error_message). If error_message is not None, the\n sections list should be treated as invalid.\n \"\"\"\n sections: list[TextSection | ImageSection] = []\n\n try:\n if not self._validate_attachment_filetype(attachment):\n return (\n [],\n f\"Unsupported file type: {attachment.get('fileName', 'unknown')}\",\n )\n\n attachment_id = attachment[\"id\"]\n file_name = attachment.get(\"fileName\", f\"attachment_{attachment_id}\")\n file_size = attachment.get(\"fileSize\", 0)\n media_type = self._get_media_type_from_filename(file_name)\n\n if file_size > DRUPAL_WIKI_ATTACHMENT_SIZE_THRESHOLD:\n return [], f\"Attachment too large: {file_size} bytes\"\n\n try:\n raw_bytes = self._download_attachment(attachment_id)\n except Exception as e:\n return [], f\"Failed to download attachment: {e}\"\n\n if media_type.startswith(\"image/\"):\n if not self.allow_images:\n logger.info(\n f\"Skipping image attachment {file_name} because allow_images is False\",\n )\n return [], None\n\n try:\n image_section, _ = store_image_and_create_section(\n image_data=raw_bytes,\n file_id=str(attachment_id),\n display_name=attachment.get(\n \"name\", attachment.get(\"fileName\", \"Unknown\")\n ),\n link=download_url,\n media_type=media_type,\n file_origin=FileOrigin.CONNECTOR,\n )\n sections.append(image_section)\n logger.debug(f\"Stored image attachment with file name: {file_name}\")\n except Exception as e:\n return [], f\"Image storage failed: {e}\"\n\n return sections, None\n\n image_counter = 0\n\n def _store_embedded_image(image_data: bytes, image_name: str) -> None:\n nonlocal image_counter\n\n if not self.allow_images:\n return\n\n media_for_image = self._get_media_type_from_filename(image_name)\n if media_for_image == \"application/octet-stream\":\n try:\n media_for_image = get_image_type_from_bytes(image_data)\n except ValueError:\n logger.warning(\n f\"Unable to determine media type for embedded image {image_name} on attachment {file_name}\"\n )\n\n image_counter += 1\n display_name = (\n image_name\n or f\"{attachment.get('name', file_name)} - embedded image {image_counter}\"\n )\n\n try:\n image_section, _ = store_image_and_create_section(\n image_data=image_data,\n file_id=f\"{attachment_id}_embedded_{image_counter}\",\n display_name=display_name,\n link=download_url,\n media_type=media_for_image,\n file_origin=FileOrigin.CONNECTOR,\n )\n sections.append(image_section)\n except Exception as err:\n logger.warning(\n f\"Failed to store embedded image {image_name or image_counter} for attachment {file_name}: {err}\"\n )\n\n extraction_result = extract_text_and_images(\n file=BytesIO(raw_bytes),\n file_name=file_name,\n content_type=media_type,\n image_callback=_store_embedded_image if self.allow_images else None,\n )\n\n text_content = extraction_result.text_content.strip()\n if text_content:\n sections.insert(0, TextSection(text=text_content, link=download_url))\n logger.info(\n f\"Extracted {len(text_content)} characters from {file_name}\"\n )\n elif not sections:\n return [], f\"No text extracted for {file_name}\"\n\n return sections, None\n\n except Exception as e:\n logger.error(\n f\"Failed to process attachment {attachment.get('name', 'unknown')} on page {page_id}: {e}\"\n )\n return [], f\"Failed to process attachment: {e}\"\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"\n Load credentials for the Drupal Wiki connector.\n\n Args:\n credentials: Dictionary containing the API token.\n\n Returns:\n None\n \"\"\"\n\n api_token = credentials.get(\"drupal_wiki_api_token\", \"\").strip()\n\n if not api_token:\n raise ConnectorValidationError(\n \"API token is required for Drupal Wiki connector\"\n )\n\n self._api_token = api_token\n self.headers.update(\n {\n \"Authorization\": f\"Bearer {api_token}\",\n }\n )\n\n return None\n\n def _get_space_ids(self) -> list[int]:\n \"\"\"\n Get all space IDs from the Drupal Wiki instance.\n\n Returns:\n List of space IDs (deduplicated). The list is sorted to be deterministic.\n \"\"\"\n url = f\"{self.base_url}/api/rest/scope/api/space\"\n size = MAX_API_PAGE_SIZE\n page = 0\n all_space_ids: set[int] = set()\n has_more = True\n last_num_ids = -1\n\n while has_more and len(all_space_ids) > last_num_ids:\n last_num_ids = len(all_space_ids)\n params = {\"size\": size, \"page\": page}\n logger.debug(f\"Fetching spaces from {url} (page={page}, size={size})\")\n response = rate_limited_get(url, headers=self.headers, params=params)\n response.raise_for_status()\n resp_json = response.json()\n space_response = DrupalWikiSpaceResponse.model_validate(resp_json)\n\n logger.info(f\"Fetched {len(space_response.content)} spaces from {page}\")\n # Collect ids into the set to deduplicate\n for space in space_response.content:\n all_space_ids.add(space.id)\n\n # Continue if we got a full page, indicating there might be more\n has_more = len(space_response.content) >= size\n\n page += 1\n\n # Return a deterministic, sorted list of ids\n space_id_list = list(sorted(all_space_ids))\n logger.debug(f\"Total spaces fetched: {len(space_id_list)}\")\n return space_id_list\n\n def _get_pages_for_space(\n self, space_id: int, modified_after: SecondsSinceUnixEpoch | None = None\n ) -> list[DrupalWikiPage]:\n \"\"\"\n Get all pages for a specific space, optionally filtered by modification time.\n\n Args:\n space_id: ID of the space.\n modified_after: Only return pages modified after this timestamp (seconds since Unix epoch).\n\n Returns:\n List of DrupalWikiPage objects.\n \"\"\"\n url = f\"{self.base_url}/api/rest/scope/api/page\"\n size = MAX_API_PAGE_SIZE\n page = 0\n all_pages = []\n has_more = True\n\n while has_more:\n params: dict[str, str | int] = {\n DRUPAL_WIKI_SPACE_KEY: str(space_id),\n \"size\": size,\n \"page\": page,\n }\n\n # Add modifiedAfter parameter if provided\n if modified_after is not None:\n params[\"modifiedAfter\"] = int(modified_after)\n\n logger.debug(\n f\"Fetching pages for space {space_id} from {url} ({page=}, {size=}, {modified_after=})\"\n )\n response = rate_limited_get(url, headers=self.headers, params=params)\n response.raise_for_status()\n resp_json = response.json()\n\n try:\n page_response = DrupalWikiPageResponse.model_validate(resp_json)\n except Exception as e:\n logger.error(f\"Failed to validate Drupal Wiki page response: {e}\")\n raise ConnectorValidationError(f\"Invalid API response format: {e}\")\n\n logger.info(\n f\"Fetched {len(page_response.content)} pages in space {space_id} (page={page})\"\n )\n\n # Pydantic should automatically parse content items as DrupalWikiPage objects\n # If validation fails, it will raise an exception which we should catch\n all_pages.extend(page_response.content)\n\n # Continue if we got a full page, indicating there might be more\n has_more = len(page_response.content) >= size\n\n page += 1\n\n logger.debug(f\"Total pages fetched for space {space_id}: {len(all_pages)}\")\n return all_pages\n\n def _get_page_content(self, page_id: int) -> DrupalWikiPage:\n \"\"\"\n Get the content of a specific page.\n\n Args:\n page_id: ID of the page.\n\n Returns:\n DrupalWikiPage object.\n \"\"\"\n url = f\"{self.base_url}/api/rest/scope/api/page/{page_id}\"\n response = rate_limited_get(url, headers=self.headers)\n response.raise_for_status()\n\n return DrupalWikiPage.model_validate(response.json())\n\n def _process_page(self, page: DrupalWikiPage) -> Document | ConnectorFailure:\n \"\"\"\n Process a page and convert it to a Document.\n\n Args:\n page: DrupalWikiPage object.\n\n Returns:\n Document object or ConnectorFailure.\n \"\"\"\n try:\n # Extract text from HTML, handle None body\n text_content = parse_html_page_basic(page.body or \"\")\n\n # Ensure text_content is a string, not None\n if text_content is None:\n text_content = \"\"\n\n # Create document URL\n page_url = build_drupal_wiki_document_id(self.base_url, page.id)\n\n # Create sections with just the page content\n sections: list[TextSection | ImageSection] = [\n TextSection(text=text_content, link=page_url)\n ]\n\n # Only process attachments if self.include_attachments is True\n if self.include_attachments:\n attachments = self._get_page_attachments(page.id)\n for attachment in attachments:\n logger.info(\n f\"Processing attachment: {attachment.get('name', 'Unknown')} (ID: {attachment['id']})\"\n )\n # Use downloadUrl from API; fallback to page URL\n raw_download = attachment.get(\"downloadUrl\")\n if raw_download:\n download_url = (\n raw_download\n if raw_download.startswith(\"http\")\n else f\"{self.base_url.rstrip('/')}\" + raw_download\n )\n else:\n download_url = page_url\n # Process the attachment\n attachment_sections, error = self._process_attachment(\n attachment, page.id, download_url\n )\n if error:\n logger.warning(\n f\"Error processing attachment {attachment.get('name', 'Unknown')}: {error}\"\n )\n continue\n\n if attachment_sections:\n sections.extend(attachment_sections)\n logger.debug(\n f\"Added {len(attachment_sections)} section(s) for attachment {attachment.get('name', 'Unknown')}\"\n )\n\n # Create metadata\n metadata: dict[str, str | list[str]] = {\n \"space_id\": str(page.homeSpace),\n \"page_id\": str(page.id),\n \"type\": page.type,\n }\n\n # Create document\n return Document(\n id=page_url,\n sections=sections,\n source=DocumentSource.DRUPAL_WIKI,\n semantic_identifier=page.title,\n metadata=metadata,\n doc_updated_at=datetime_from_utc_timestamp(page.lastModified),\n )\n except Exception as e:\n logger.error(f\"Error processing page {page.id}: {e}\")\n return ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=str(page.id),\n document_link=build_drupal_wiki_document_id(self.base_url, page.id),\n ),\n failure_message=f\"Error processing page {page.id}: {e}\",\n exception=e,\n )\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: DrupalWikiCheckpoint,\n ) -> CheckpointOutput[DrupalWikiCheckpoint]:\n \"\"\"\n Load documents from a checkpoint.\n\n Args:\n start: Start time as seconds since Unix epoch.\n end: End time as seconds since Unix epoch.\n checkpoint: Checkpoint to resume from.\n\n Returns:\n Generator yielding documents and the updated checkpoint.\n \"\"\"\n # Ensure page_ids is not None\n if checkpoint.page_ids is None:\n checkpoint.page_ids = []\n\n # Initialize page_ids from self.pages if not already set\n if not checkpoint.page_ids and self.pages:\n logger.info(f\"Initializing page_ids from self.pages: {self.pages}\")\n checkpoint.page_ids = [int(page_id.strip()) for page_id in self.pages]\n\n # Ensure spaces is not None\n if checkpoint.spaces is None:\n checkpoint.spaces = []\n\n while checkpoint.current_page_id_index < len(checkpoint.page_ids):\n page_id = checkpoint.page_ids[checkpoint.current_page_id_index]\n logger.debug(f\"Processing page ID: {page_id}\")\n\n try:\n # Get the page content directly\n page = self._get_page_content(page_id)\n\n # Skip pages outside the time range\n if not self._is_page_in_time_range(page.lastModified, start, end):\n logger.info(f\"Skipping page {page_id} - outside time range\")\n checkpoint.current_page_id_index += 1\n continue\n\n # Process the page\n doc_or_failure = self._process_page(page)\n yield doc_or_failure\n\n except Exception as e:\n logger.error(f\"Error processing page ID {page_id}: {e}\")\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=str(page_id),\n document_link=build_drupal_wiki_document_id(\n self.base_url, page_id\n ),\n ),\n failure_message=f\"Error processing page ID {page_id}: {e}\",\n exception=e,\n )\n\n # Move to the next page ID\n checkpoint.current_page_id_index += 1\n\n # TODO: The main benefit of CheckpointedConnectors is that they can \"save their work\"\n # by storing a checkpoint so transient errors are easy to recover from: simply resume\n # from the last checkpoint. The way to get checkpoints saved is to return them somewhere\n # in the middle of this function. The guarantee our checkpointing system gives to you,\n # the connector implementer, is that when you return a checkpoint, this connector will\n # at a later time (generally within a few seconds) call the load_from_checkpoint function\n # again with the checkpoint you last returned as long as has_more=True.\n\n # Process spaces if include_all_spaces is True or spaces are provided\n if self.include_all_spaces or self.spaces:\n # If include_all_spaces is True, always fetch all spaces\n if self.include_all_spaces:\n logger.info(\"Fetching all spaces\")\n # Fetch all spaces\n all_space_ids = self._get_space_ids()\n # checkpoint.spaces expects a list of ints; assign returned list\n checkpoint.spaces = all_space_ids\n logger.info(f\"Found {len(checkpoint.spaces)} spaces to process\")\n # Otherwise, use provided spaces if checkpoint is empty\n elif not checkpoint.spaces:\n logger.info(f\"Using provided spaces: {self.spaces}\")\n # Use provided spaces\n checkpoint.spaces = [int(space_id.strip()) for space_id in self.spaces]\n\n # Process spaces from the checkpoint\n while checkpoint.current_space_index < len(checkpoint.spaces):\n space_id = checkpoint.spaces[checkpoint.current_space_index]\n logger.debug(f\"Processing space ID: {space_id}\")\n\n # Get pages for the current space, filtered by start time if provided\n pages = self._get_pages_for_space(space_id, modified_after=start)\n\n # Process pages from the checkpoint\n while checkpoint.current_page_index < len(pages):\n page = pages[checkpoint.current_page_index]\n logger.debug(f\"Processing page: {page.title} (ID: {page.id})\")\n\n # For space-based pages, we already filtered by modifiedAfter in the API call\n # Only need to check the end time boundary\n if end and page.lastModified >= end:\n logger.info(\n f\"Skipping page {page.id} - outside time range (after end)\"\n )\n checkpoint.current_page_index += 1\n continue\n\n # Process the page\n doc_or_failure = self._process_page(page)\n yield doc_or_failure\n\n # Move to the next page\n checkpoint.current_page_index += 1\n\n # Move to the next space\n checkpoint.current_space_index += 1\n checkpoint.current_page_index = 0\n\n # All spaces and pages processed\n logger.info(\"Finished processing all spaces and pages\")\n checkpoint.has_more = False\n return checkpoint\n\n @override\n def build_dummy_checkpoint(self) -> DrupalWikiCheckpoint:\n \"\"\"\n Build a dummy checkpoint.\n\n Returns:\n DrupalWikiCheckpoint with default values.\n \"\"\"\n return DrupalWikiCheckpoint(\n has_more=True,\n current_space_index=0,\n current_page_index=0,\n current_page_id_index=0,\n spaces=[],\n page_ids=[],\n is_processing_specific_pages=False,\n )\n\n @override\n def validate_checkpoint_json(self, checkpoint_json: str) -> DrupalWikiCheckpoint:\n \"\"\"\n Validate a checkpoint JSON string.\n\n Args:\n checkpoint_json: JSON string representing a checkpoint.\n\n Returns:\n Validated DrupalWikiCheckpoint.\n \"\"\"\n return DrupalWikiCheckpoint.model_validate_json(checkpoint_json)\n\n # TODO: unify approach with load_from_checkpoint.\n # Ideally slim retrieval shares a lot of the same code with non-slim\n # and we pass in a param is_slim to the main helper function\n # that does the retrieval.\n @override\n def retrieve_all_slim_docs(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n \"\"\"\n Retrieve all slim documents.\n\n Args:\n start: Start time as seconds since Unix epoch.\n end: End time as seconds since Unix epoch.\n callback: Callback for indexing heartbeat.\n\n Returns:\n Generator yielding batches of SlimDocument objects.\n \"\"\"\n slim_docs: list[SlimDocument | HierarchyNode] = []\n logger.info(\n f\"Starting retrieve_all_slim_docs with include_all_spaces={self.include_all_spaces}, spaces={self.spaces}\"\n )\n\n # Process specific page IDs if provided\n if self.pages:\n logger.info(f\"Processing specific pages: {self.pages}\")\n for page_id in self.pages:\n try:\n # Get the page content directly\n page_content = self._get_page_content(int(page_id.strip()))\n\n # Skip pages outside the time range\n if not self._is_page_in_time_range(\n page_content.lastModified, start, end\n ):\n logger.info(f\"Skipping page {page_id} - outside time range\")\n continue\n\n # Create slim document for the page\n page_url = build_drupal_wiki_document_id(\n self.base_url, page_content.id\n )\n slim_docs.append(\n SlimDocument(\n id=page_url,\n )\n )\n logger.debug(f\"Added slim document for page {page_content.id}\")\n\n # Process attachments for this page\n attachments = self._get_page_attachments(page_content.id)\n for attachment in attachments:\n if self._validate_attachment_filetype(attachment):\n attachment_url = f\"{page_url}#attachment-{attachment['id']}\"\n slim_docs.append(\n SlimDocument(\n id=attachment_url,\n )\n )\n logger.debug(\n f\"Added slim document for attachment {attachment['id']}\"\n )\n\n # Yield batch if it reaches the batch size\n if len(slim_docs) >= self.batch_size:\n logger.debug(\n f\"Yielding batch of {len(slim_docs)} slim documents\"\n )\n yield slim_docs\n slim_docs = []\n\n if callback and callback.should_stop():\n return\n if callback:\n callback.progress(\"retrieve_all_slim_docs\", 1)\n\n except Exception as e:\n logger.error(\n f\"Error processing page ID {page_id} for slim documents: {e}\"\n )\n\n # Process spaces if include_all_spaces is True or spaces are provided\n if self.include_all_spaces or self.spaces:\n logger.info(\"Processing spaces for slim documents\")\n # Get spaces to process\n spaces_to_process = []\n if self.include_all_spaces:\n logger.info(\"Fetching all spaces for slim documents\")\n # Fetch all spaces\n all_space_ids = self._get_space_ids()\n spaces_to_process = all_space_ids\n logger.info(f\"Found {len(spaces_to_process)} spaces to process\")\n else:\n logger.info(f\"Using provided spaces: {self.spaces}\")\n # Use provided spaces\n spaces_to_process = [int(space_id.strip()) for space_id in self.spaces]\n\n # Process each space\n for space_id in spaces_to_process:\n logger.info(f\"Processing space ID: {space_id}\")\n # Get pages for the current space, filtered by start time if provided\n pages = self._get_pages_for_space(space_id, modified_after=start)\n\n # Process each page\n for page in pages:\n logger.debug(f\"Processing page: {page.title} (ID: {page.id})\")\n # Skip pages outside the time range\n if end and page.lastModified >= end:\n logger.info(\n f\"Skipping page {page.id} - outside time range (after end)\"\n )\n continue\n\n # Create slim document for the page\n page_url = build_drupal_wiki_document_id(self.base_url, page.id)\n slim_docs.append(\n SlimDocument(\n id=page_url,\n )\n )\n logger.info(f\"Added slim document for page {page.id}\")\n\n # Process attachments for this page\n attachments = self._get_page_attachments(page.id)\n for attachment in attachments:\n if self._validate_attachment_filetype(attachment):\n attachment_url = f\"{page_url}#attachment-{attachment['id']}\"\n slim_docs.append(\n SlimDocument(\n id=attachment_url,\n )\n )\n logger.info(\n f\"Added slim document for attachment {attachment['id']}\"\n )\n\n # Yield batch if it reaches the batch size\n if len(slim_docs) >= self.batch_size:\n logger.info(\n f\"Yielding batch of {len(slim_docs)} slim documents\"\n )\n yield slim_docs\n slim_docs = []\n\n if callback and callback.should_stop():\n return\n if callback:\n callback.progress(\"retrieve_all_slim_docs\", 1)\n\n # Yield remaining documents\n if slim_docs:\n logger.debug(f\"Yielding final batch of {len(slim_docs)} slim documents\")\n yield slim_docs\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n Validate the connector settings.\n\n Raises:\n ConnectorValidationError: If the settings are invalid.\n \"\"\"\n if not self.headers:\n raise ConnectorMissingCredentialError(\"Drupal Wiki\")\n\n try:\n # Try to fetch spaces to validate the connection\n # Call the new helper which returns the list of space ids\n self._get_space_ids()\n except requests.exceptions.RequestException as e:\n raise ConnectorValidationError(f\"Failed to connect to Drupal Wiki: {e}\")\n\n def _is_page_in_time_range(\n self,\n last_modified: int,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n ) -> bool:\n \"\"\"\n Check if a page's last modified timestamp falls within the specified time range.\n\n Args:\n last_modified: The page's last modified timestamp.\n start: Start time as seconds since Unix epoch (inclusive).\n end: End time as seconds since Unix epoch (exclusive).\n\n Returns:\n True if the page is within the time range, False otherwise.\n \"\"\"\n return (not start or last_modified >= start) and (\n not end or last_modified < end\n )\n" }, { "path": "backend/onyx/connectors/drupal_wiki/models.py", "content": "from enum import Enum\nfrom typing import Generic\nfrom typing import List\nfrom typing import Optional\nfrom typing import TypeVar\n\nfrom pydantic import BaseModel\n\nfrom onyx.connectors.interfaces import ConnectorCheckpoint\n\n\nclass SpaceAccessStatus(str, Enum):\n \"\"\"Enum for Drupal Wiki space access status\"\"\"\n\n PRIVATE = \"PRIVATE\"\n ANONYMOUS = \"ANONYMOUS\"\n AUTHENTICATED = \"AUTHENTICATED\"\n\n\nclass DrupalWikiSpace(BaseModel):\n \"\"\"Model for a Drupal Wiki space\"\"\"\n\n id: int\n name: str\n type: str\n description: Optional[str] = None\n accessStatus: Optional[SpaceAccessStatus] = None\n color: Optional[str] = None\n\n\nclass DrupalWikiPage(BaseModel):\n \"\"\"Model for a Drupal Wiki page\"\"\"\n\n id: int\n title: str\n homeSpace: int\n lastModified: int\n type: str\n body: Optional[str] = None\n\n\nT = TypeVar(\"T\")\n\n\nclass DrupalWikiBaseResponse(BaseModel, Generic[T]):\n \"\"\"Base model for Drupal Wiki API responses\"\"\"\n\n totalPages: int\n totalElements: int\n size: int\n content: List[T]\n number: int\n first: bool\n last: bool\n numberOfElements: int\n empty: bool\n\n\nclass DrupalWikiSpaceResponse(DrupalWikiBaseResponse[DrupalWikiSpace]):\n \"\"\"Model for the response from the Drupal Wiki spaces API\"\"\"\n\n\nclass DrupalWikiPageResponse(DrupalWikiBaseResponse[DrupalWikiPage]):\n \"\"\"Model for the response from the Drupal Wiki pages API\"\"\"\n\n\nclass DrupalWikiCheckpoint(ConnectorCheckpoint):\n \"\"\"Checkpoint for the Drupal Wiki connector\"\"\"\n\n current_space_index: int = 0\n current_page_index: int = 0\n current_page_id_index: int = 0\n spaces: List[int] = []\n page_ids: List[int] = []\n is_processing_specific_pages: bool = False\n" }, { "path": "backend/onyx/connectors/drupal_wiki/utils.py", "content": "from onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef build_drupal_wiki_document_id(base_url: str, page_id: int) -> str:\n \"\"\"Build a document ID for a Drupal Wiki page using the real URL format\"\"\"\n # Ensure base_url ends with a slash\n base_url = base_url.rstrip(\"/\") + \"/\"\n return f\"{base_url}node/{page_id}\"\n" }, { "path": "backend/onyx/connectors/egnyte/connector.py", "content": "import io\nimport os\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import IO\nfrom urllib.parse import quote\n\nfrom pydantic import Field\n\nfrom onyx.configs.app_configs import EGNYTE_CLIENT_ID\nfrom onyx.configs.app_configs import EGNYTE_CLIENT_SECRET\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n get_oauth_callback_uri,\n)\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import OAuthConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import detect_encoding\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.extract_file_text import read_text_file\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import request_with_retries\n\n\nlogger = setup_logger()\n\n_EGNYTE_API_BASE = \"https://{domain}.egnyte.com/pubapi/v1\"\n_EGNYTE_APP_BASE = \"https://{domain}.egnyte.com\"\n\n\ndef _parse_last_modified(last_modified: str) -> datetime:\n return datetime.strptime(last_modified, \"%a, %d %b %Y %H:%M:%S %Z\").replace(\n tzinfo=timezone.utc\n )\n\n\ndef _process_egnyte_file(\n file_metadata: dict[str, Any],\n file_content: IO,\n base_url: str,\n folder_path: str | None = None,\n) -> Document | None:\n \"\"\"Process an Egnyte file into a Document object\n\n Args:\n file_data: The file data from Egnyte API\n file_content: The raw content of the file in bytes\n base_url: The base URL for the Egnyte instance\n folder_path: Optional folder path to filter results\n \"\"\"\n # Skip if file path doesn't match folder path filter\n if folder_path and not file_metadata[\"path\"].startswith(folder_path):\n raise ValueError(\n f\"File path {file_metadata['path']} does not match folder path {folder_path}\"\n )\n\n file_name = file_metadata[\"name\"]\n extension = get_file_ext(file_name)\n\n # Explicitly excluding image extensions here. TODO: consider allowing images\n if extension not in OnyxFileExtensions.TEXT_AND_DOCUMENT_EXTENSIONS:\n logger.warning(f\"Skipping file '{file_name}' with extension '{extension}'\")\n return None\n\n # Extract text content based on file type\n # TODO @wenxi-onyx: convert to extract_text_and_images\n if extension in OnyxFileExtensions.PLAIN_TEXT_EXTENSIONS:\n encoding = detect_encoding(file_content)\n file_content_raw, file_metadata = read_text_file(\n file_content, encoding=encoding, ignore_onyx_metadata=False\n )\n else:\n file_content_raw = extract_file_text(\n file=file_content,\n file_name=file_name,\n break_on_unprocessable=True,\n )\n\n # Build the web URL for the file\n web_url = f\"{base_url}/navigate/file/{file_metadata['group_id']}\"\n\n # Create document metadata\n metadata: dict[str, str | list[str]] = {\n \"file_path\": file_metadata[\"path\"],\n \"last_modified\": file_metadata.get(\"last_modified\", \"\"),\n }\n\n # Add lock info if present\n if lock_info := file_metadata.get(\"lock_info\"):\n metadata[\"lock_owner\"] = (\n f\"{lock_info.get('first_name', '')} {lock_info.get('last_name', '')}\"\n )\n\n # Create the document owners\n primary_owner = None\n if uploaded_by := file_metadata.get(\"uploaded_by\"):\n primary_owner = BasicExpertInfo(\n email=uploaded_by, # Using username as email since that's what we have\n )\n\n # Create the document\n return Document(\n id=f\"egnyte-{file_metadata['entry_id']}\",\n sections=[TextSection(text=file_content_raw.strip(), link=web_url)],\n source=DocumentSource.EGNYTE,\n semantic_identifier=file_name,\n metadata=metadata,\n doc_updated_at=(\n _parse_last_modified(file_metadata[\"last_modified\"])\n if \"last_modified\" in file_metadata\n else None\n ),\n primary_owners=[primary_owner] if primary_owner else None,\n )\n\n\nclass EgnyteConnector(LoadConnector, PollConnector, OAuthConnector):\n class AdditionalOauthKwargs(OAuthConnector.AdditionalOauthKwargs):\n egnyte_domain: str = Field(\n title=\"Egnyte Domain\",\n description=(\n \"The domain for the Egnyte instance (e.g. 'company' for company.egnyte.com)\"\n ),\n )\n\n def __init__(\n self,\n folder_path: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.domain = \"\" # will always be set in `load_credentials`\n self.folder_path = folder_path or \"\" # Root folder if not specified\n self.batch_size = batch_size\n self.access_token: str | None = None\n\n @classmethod\n def oauth_id(cls) -> DocumentSource:\n return DocumentSource.EGNYTE\n\n @classmethod\n def oauth_authorization_url(\n cls,\n base_domain: str,\n state: str,\n additional_kwargs: dict[str, str],\n ) -> str:\n if not EGNYTE_CLIENT_ID:\n raise ValueError(\"EGNYTE_CLIENT_ID environment variable must be set\")\n\n oauth_kwargs = cls.AdditionalOauthKwargs(**additional_kwargs)\n\n callback_uri = get_oauth_callback_uri(base_domain, \"egnyte\")\n return (\n f\"https://{oauth_kwargs.egnyte_domain}.egnyte.com/puboauth/token\"\n f\"?client_id={EGNYTE_CLIENT_ID}\"\n f\"&redirect_uri={callback_uri}\"\n f\"&scope=Egnyte.filesystem\"\n f\"&state={state}\"\n f\"&response_type=code\"\n )\n\n @classmethod\n def oauth_code_to_token(\n cls,\n base_domain: str,\n code: str,\n additional_kwargs: dict[str, str],\n ) -> dict[str, Any]:\n if not EGNYTE_CLIENT_ID:\n raise ValueError(\"EGNYTE_CLIENT_ID environment variable must be set\")\n if not EGNYTE_CLIENT_SECRET:\n raise ValueError(\"EGNYTE_CLIENT_SECRET environment variable must be set\")\n\n oauth_kwargs = cls.AdditionalOauthKwargs(**additional_kwargs)\n\n # Exchange code for token\n url = f\"https://{oauth_kwargs.egnyte_domain}.egnyte.com/puboauth/token\"\n redirect_uri = get_oauth_callback_uri(base_domain, \"egnyte\")\n\n data = {\n \"client_id\": EGNYTE_CLIENT_ID,\n \"client_secret\": EGNYTE_CLIENT_SECRET,\n \"code\": code,\n \"grant_type\": \"authorization_code\",\n \"redirect_uri\": redirect_uri,\n \"scope\": \"Egnyte.filesystem\",\n }\n headers = {\"Content-Type\": \"application/x-www-form-urlencoded\"}\n\n response = request_with_retries(\n method=\"POST\",\n url=url,\n data=data,\n headers=headers,\n # try a lot faster since this is a realtime flow\n backoff=0,\n delay=0.1,\n )\n if not response.ok:\n raise RuntimeError(f\"Failed to exchange code for token: {response.text}\")\n\n token_data = response.json()\n return {\n \"domain\": oauth_kwargs.egnyte_domain,\n \"access_token\": token_data[\"access_token\"],\n }\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.domain = credentials[\"domain\"]\n self.access_token = credentials[\"access_token\"]\n return None\n\n def _get_files_list(\n self,\n path: str,\n ) -> Generator[dict[str, Any], None, None]:\n if not self.access_token or not self.domain:\n raise ConnectorMissingCredentialError(\"Egnyte\")\n\n headers = {\n \"Authorization\": f\"Bearer {self.access_token}\",\n }\n\n params: dict[str, Any] = {\n \"list_content\": True,\n }\n\n url_encoded_path = quote(path or \"\")\n url = f\"{_EGNYTE_API_BASE.format(domain=self.domain)}/fs/{url_encoded_path}\"\n response = request_with_retries(\n method=\"GET\", url=url, headers=headers, params=params\n )\n if not response.ok:\n raise RuntimeError(f\"Failed to fetch files from Egnyte: {response.text}\")\n\n data = response.json()\n\n # Yield files from current directory\n for file in data.get(\"files\", []):\n yield file\n\n # Recursively traverse folders\n for folder in data.get(\"folders\", []):\n yield from self._get_files_list(folder[\"path\"])\n\n def _should_index_file(\n self,\n file: dict[str, Any],\n start_time: datetime | None = None,\n end_time: datetime | None = None,\n ) -> bool:\n \"\"\"Return True if file should be included based on filters.\"\"\"\n if file[\"is_folder\"]:\n return False\n\n file_modified = _parse_last_modified(file[\"last_modified\"])\n if start_time and file_modified < start_time:\n return False\n if end_time and file_modified > end_time:\n return False\n\n return True\n\n def _process_files(\n self,\n start_time: datetime | None = None,\n end_time: datetime | None = None,\n ) -> Generator[list[Document | HierarchyNode], None, None]:\n current_batch: list[Document | HierarchyNode] = []\n\n # Iterate through yielded files and filter them\n for file in self._get_files_list(self.folder_path):\n if not self._should_index_file(file, start_time, end_time):\n logger.debug(f\"Skipping file '{file['path']}'.\")\n continue\n\n try:\n # Set up request with streaming enabled\n headers = {\n \"Authorization\": f\"Bearer {self.access_token}\",\n }\n url_encoded_path = quote(file[\"path\"])\n url = f\"{_EGNYTE_API_BASE.format(domain=self.domain)}/fs-content/{url_encoded_path}\"\n response = request_with_retries(\n method=\"GET\",\n url=url,\n headers=headers,\n stream=True,\n )\n\n if not response.ok:\n logger.error(\n f\"Failed to fetch file content: {file['path']} (status code: {response.status_code})\"\n )\n continue\n\n # Stream the response content into a BytesIO buffer\n buffer = io.BytesIO()\n for chunk in response.iter_content(chunk_size=8192):\n if chunk:\n buffer.write(chunk)\n\n # Reset buffer's position to the start\n buffer.seek(0)\n\n # Process the streamed file content\n doc = _process_egnyte_file(\n file_metadata=file,\n file_content=buffer,\n base_url=_EGNYTE_APP_BASE.format(domain=self.domain),\n folder_path=self.folder_path,\n )\n\n if doc is not None:\n current_batch.append(doc)\n\n if len(current_batch) >= self.batch_size:\n yield current_batch\n current_batch = []\n\n except Exception:\n logger.exception(f\"Failed to process file {file['path']}\")\n continue\n\n if current_batch:\n yield current_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n yield from self._process_files()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_time = datetime.fromtimestamp(start, tz=timezone.utc)\n end_time = datetime.fromtimestamp(end, tz=timezone.utc)\n\n yield from self._process_files(start_time=start_time, end_time=end_time)\n\n\nif __name__ == \"__main__\":\n connector = EgnyteConnector()\n connector.load_credentials(\n {\n \"domain\": os.environ[\"EGNYTE_DOMAIN\"],\n \"access_token\": os.environ[\"EGNYTE_ACCESS_TOKEN\"],\n }\n )\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/exceptions.py", "content": "class ValidationError(Exception):\n \"\"\"General exception for validation errors.\"\"\"\n\n def __init__(self, message: str):\n self.message = message\n super().__init__(self.message)\n\n\nclass ConnectorValidationError(ValidationError):\n \"\"\"General exception for connector validation errors.\"\"\"\n\n def __init__(self, message: str):\n self.message = message\n super().__init__(self.message)\n\n\nclass UnexpectedValidationError(ValidationError):\n \"\"\"Raised when an unexpected error occurs during connector validation.\n\n Unexpected errors don't necessarily mean the credential is invalid,\n but rather that there was an error during the validation process\n or we encountered a currently unhandled error case.\n\n Currently, unexpected validation errors are defined as transient and should not be\n used to disable the connector.\n \"\"\"\n\n def __init__(self, message: str = \"Unexpected error during connector validation\"):\n super().__init__(message)\n\n\nclass CredentialInvalidError(ConnectorValidationError):\n \"\"\"Raised when a connector's credential is invalid.\"\"\"\n\n def __init__(self, message: str = \"Credential is invalid\"):\n super().__init__(message)\n\n\nclass CredentialExpiredError(ConnectorValidationError):\n \"\"\"Raised when a connector's credential is expired.\"\"\"\n\n def __init__(self, message: str = \"Credential has expired\"):\n super().__init__(message)\n\n\nclass InsufficientPermissionsError(ConnectorValidationError):\n \"\"\"Raised when the credential does not have sufficient API permissions.\"\"\"\n\n def __init__(\n self, message: str = \"Insufficient permissions for the requested operation\"\n ):\n super().__init__(message)\n" }, { "path": "backend/onyx/connectors/factory.py", "content": "import importlib\nfrom typing import Any\nfrom typing import Type\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import INTEGRATION_TESTS_MODE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled\nfrom onyx.connectors.credentials_provider import OnyxDBCredentialsProvider\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.interfaces import BaseConnector\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CredentialsConnector\nfrom onyx.connectors.interfaces import EventConnector\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.models import InputType\nfrom onyx.connectors.registry import CONNECTOR_CLASS_MAP\nfrom onyx.db.connector import fetch_connector_by_id\nfrom onyx.db.credentials import backend_update_credential_json\nfrom onyx.db.credentials import fetch_credential_by_id\nfrom onyx.db.enums import AccessType\nfrom onyx.db.models import Credential\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\nclass ConnectorMissingException(Exception):\n pass\n\n\n# Cache for already imported connector classes\n_connector_cache: dict[DocumentSource, Type[BaseConnector]] = {}\n\n\ndef _load_connector_class(source: DocumentSource) -> Type[BaseConnector]:\n \"\"\"Dynamically load and cache a connector class.\"\"\"\n if source in _connector_cache:\n return _connector_cache[source]\n\n if source not in CONNECTOR_CLASS_MAP:\n raise ConnectorMissingException(f\"Connector not found for source={source}\")\n\n mapping = CONNECTOR_CLASS_MAP[source]\n\n try:\n module = importlib.import_module(mapping.module_path)\n connector_class = getattr(module, mapping.class_name)\n _connector_cache[source] = connector_class\n return connector_class\n except (ImportError, AttributeError) as e:\n raise ConnectorMissingException(\n f\"Failed to import {mapping.class_name} from {mapping.module_path}: {e}\"\n )\n\n\ndef _validate_connector_supports_input_type(\n connector: Type[BaseConnector],\n input_type: InputType | None,\n source: DocumentSource,\n) -> None:\n \"\"\"Validate that a connector supports the requested input type.\"\"\"\n if input_type is None:\n return\n\n # Check each input type requirement separately for clarity\n load_state_unsupported = input_type == InputType.LOAD_STATE and not issubclass(\n connector, LoadConnector\n )\n\n poll_unsupported = (\n input_type == InputType.POLL\n # Either poll or checkpoint works for this, in the future\n # all connectors should be checkpoint connectors\n and (\n not issubclass(connector, PollConnector)\n and not issubclass(connector, CheckpointedConnector)\n )\n )\n\n event_unsupported = input_type == InputType.EVENT and not issubclass(\n connector, EventConnector\n )\n\n if any([load_state_unsupported, poll_unsupported, event_unsupported]):\n raise ConnectorMissingException(\n f\"Connector for source={source} does not accept input_type={input_type}\"\n )\n\n\ndef identify_connector_class(\n source: DocumentSource,\n input_type: InputType | None = None,\n) -> Type[BaseConnector]:\n # Load the connector class using lazy loading\n connector = _load_connector_class(source)\n\n # Validate connector supports the requested input_type\n _validate_connector_supports_input_type(connector, input_type, source)\n\n return connector\n\n\ndef instantiate_connector(\n db_session: Session,\n source: DocumentSource,\n input_type: InputType,\n connector_specific_config: dict[str, Any],\n credential: Credential,\n) -> BaseConnector:\n connector_class = identify_connector_class(source, input_type)\n\n connector = connector_class(**connector_specific_config)\n\n if isinstance(connector, CredentialsConnector):\n provider = OnyxDBCredentialsProvider(\n get_current_tenant_id(), str(source), credential.id\n )\n connector.set_credentials_provider(provider)\n else:\n credential_json = (\n credential.credential_json.get_value(apply_mask=False)\n if credential.credential_json\n else {}\n )\n new_credentials = connector.load_credentials(credential_json)\n\n if new_credentials is not None:\n backend_update_credential_json(credential, new_credentials, db_session)\n\n connector.set_allow_images(get_image_extraction_and_analysis_enabled())\n\n return connector\n\n\ndef validate_ccpair_for_user(\n connector_id: int,\n credential_id: int,\n access_type: AccessType,\n db_session: Session,\n enforce_creation: bool = True,\n) -> bool:\n if INTEGRATION_TESTS_MODE:\n return True\n\n # Validate the connector settings\n connector = fetch_connector_by_id(connector_id, db_session)\n credential = fetch_credential_by_id(\n credential_id,\n db_session,\n )\n\n if not connector:\n raise ValueError(\"Connector not found\")\n\n if (\n connector.source == DocumentSource.INGESTION_API\n or connector.source == DocumentSource.MOCK_CONNECTOR\n ):\n return True\n\n if not credential:\n raise ValueError(\"Credential not found\")\n\n try:\n runnable_connector = instantiate_connector(\n db_session=db_session,\n source=connector.source,\n input_type=connector.input_type,\n connector_specific_config=connector.connector_specific_config,\n credential=credential,\n )\n except ConnectorValidationError as e:\n raise e\n except Exception as e:\n if enforce_creation:\n raise ConnectorValidationError(str(e))\n else:\n return False\n\n runnable_connector.validate_connector_settings()\n if access_type == AccessType.SYNC:\n runnable_connector.validate_perm_sync()\n return True\n" }, { "path": "backend/onyx/connectors/file/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/file/connector.py", "content": "import json\nimport os\nfrom datetime import datetime\nfrom datetime import timezone\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import IO\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n process_onyx_metadata,\n)\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_text_and_images\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.image_utils import store_image_and_create_section\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _create_image_section(\n image_data: bytes,\n parent_file_name: str,\n display_name: str,\n media_type: str | None = None,\n link: str | None = None,\n idx: int = 0,\n) -> tuple[ImageSection, str | None]:\n \"\"\"\n Creates an ImageSection for an image file or embedded image.\n Stores the image in FileStore but does not generate a summary.\n\n Args:\n image_data: Raw image bytes\n db_session: Database session\n parent_file_name: Name of the parent file (for embedded images)\n display_name: Display name for the image\n idx: Index for embedded images\n\n Returns:\n Tuple of (ImageSection, stored_file_name or None)\n \"\"\"\n # Create a unique identifier for the image\n file_id = f\"{parent_file_name}_embedded_{idx}\" if idx > 0 else parent_file_name\n\n # Store the image and create a section\n try:\n section, stored_file_name = store_image_and_create_section(\n image_data=image_data,\n file_id=file_id,\n display_name=display_name,\n media_type=(\n media_type if media_type is not None else \"application/octet-stream\"\n ),\n link=link,\n file_origin=FileOrigin.CONNECTOR,\n )\n return section, stored_file_name\n except Exception as e:\n logger.error(f\"Failed to store image {display_name}: {e}\")\n raise e\n\n\ndef _process_file(\n file_id: str,\n file_name: str,\n file: IO[Any],\n metadata: dict[str, Any] | None,\n pdf_pass: str | None,\n file_type: str | None,\n) -> list[Document]:\n \"\"\"\n Process a file and return a list of Documents.\n For images, creates ImageSection objects without summarization.\n For documents with embedded images, extracts and stores the images.\n \"\"\"\n if metadata is None:\n metadata = {}\n\n # Get file extension and determine file type\n extension = get_file_ext(file_name)\n\n if extension not in OnyxFileExtensions.ALL_ALLOWED_EXTENSIONS:\n logger.warning(\n f\"Skipping file '{file_name}' with unrecognized extension '{extension}'\"\n )\n return []\n\n # If a zip is uploaded with a metadata file, we can process it here\n onyx_metadata, custom_tags = process_onyx_metadata(metadata)\n file_display_name = onyx_metadata.file_display_name or os.path.basename(file_name)\n time_updated = onyx_metadata.doc_updated_at or datetime.now(timezone.utc)\n primary_owners = onyx_metadata.primary_owners\n secondary_owners = onyx_metadata.secondary_owners\n link = onyx_metadata.link\n\n # These metadata items are not settable by the user\n source_type = onyx_metadata.source_type or DocumentSource.FILE\n\n doc_id = onyx_metadata.document_id or f\"FILE_CONNECTOR__{file_id}\"\n title = metadata.get(\"title\") or file_display_name\n\n # 1) If the file itself is an image, handle that scenario quickly\n if extension in OnyxFileExtensions.IMAGE_EXTENSIONS:\n # Read the image data\n image_data = file.read()\n if not image_data:\n logger.warning(f\"Empty image file: {file_name}\")\n return []\n\n # Create an ImageSection for the image\n try:\n section, _ = _create_image_section(\n image_data=image_data,\n parent_file_name=file_id,\n display_name=title,\n media_type=file_type,\n )\n\n return [\n Document(\n id=doc_id,\n sections=[section],\n source=source_type,\n semantic_identifier=file_display_name,\n title=title,\n doc_updated_at=time_updated,\n primary_owners=primary_owners,\n secondary_owners=secondary_owners,\n metadata=custom_tags,\n )\n ]\n except Exception as e:\n logger.error(f\"Failed to process image file {file_name}: {e}\")\n return []\n\n # 2) Otherwise: text-based approach. Possibly with embedded images.\n file.seek(0)\n\n # Extract text and images from the file\n extraction_result = extract_text_and_images(\n file=file,\n file_name=file_name,\n pdf_pass=pdf_pass,\n content_type=file_type,\n )\n\n # Each file may have file-specific ONYX_METADATA https://docs.onyx.app/admins/connectors/official/file\n # If so, we should add it to any metadata processed so far\n if extraction_result.metadata:\n logger.debug(\n f\"Found file-specific metadata for {file_name}: {extraction_result.metadata}\"\n )\n onyx_metadata, more_custom_tags = process_onyx_metadata(\n extraction_result.metadata\n )\n\n # Add file-specific tags\n custom_tags.update(more_custom_tags)\n\n # File-specific metadata overrides metadata processed so far\n source_type = onyx_metadata.source_type or source_type\n primary_owners = onyx_metadata.primary_owners or primary_owners\n secondary_owners = onyx_metadata.secondary_owners or secondary_owners\n time_updated = onyx_metadata.doc_updated_at or time_updated\n file_display_name = onyx_metadata.file_display_name or file_display_name\n title = onyx_metadata.title or onyx_metadata.file_display_name or title\n link = onyx_metadata.link or link\n\n # Build sections: first the text as a single Section\n sections: list[TextSection | ImageSection] = []\n if extraction_result.text_content.strip():\n logger.debug(f\"Creating TextSection for {file_name} with link: {link}\")\n sections.append(\n TextSection(link=link, text=extraction_result.text_content.strip())\n )\n\n # Then any extracted images from docx, PDFs, etc.\n for idx, (img_data, img_name) in enumerate(\n extraction_result.embedded_images, start=1\n ):\n # Store each embedded image as a separate file in FileStore\n # and create a section with the image reference\n try:\n image_section, stored_file_name = _create_image_section(\n image_data=img_data,\n parent_file_name=file_id,\n display_name=f\"{title} - image {idx}\",\n media_type=\"application/octet-stream\", # Default media type for embedded images\n idx=idx,\n )\n sections.append(image_section)\n logger.debug(\n f\"Created ImageSection for embedded image {idx} in {file_name}, stored as: {stored_file_name}\"\n )\n except Exception as e:\n logger.warning(\n f\"Failed to process embedded image {idx} in {file_name}: {e}\"\n )\n\n return [\n Document(\n id=doc_id,\n sections=sections,\n source=source_type,\n semantic_identifier=file_display_name,\n title=title,\n doc_updated_at=time_updated,\n primary_owners=primary_owners,\n secondary_owners=secondary_owners,\n metadata=custom_tags,\n )\n ]\n\n\nclass LocalFileConnector(LoadConnector):\n \"\"\"\n Connector that reads files from Postgres and yields Documents, including\n embedded image extraction without summarization.\n\n file_locations are S3/Filestore UUIDs\n file_names are the names of the files\n \"\"\"\n\n # Note: file_names is a required parameter, but should not break backwards compatibility.\n # If add_file_names migration is not run, old file connector configs will not have file_names.\n # file_names is only used for display purposes in the UI and file_locations is used as a fallback.\n def __init__(\n self,\n file_locations: list[Path | str],\n file_names: list[str] | None = None, # noqa: ARG002\n zip_metadata_file_id: str | None = None,\n zip_metadata: dict[str, Any] | None = None, # Deprecated, for backwards compat\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.file_locations = [str(loc) for loc in file_locations]\n self.batch_size = batch_size\n self.pdf_pass: str | None = None\n self._zip_metadata_file_id = zip_metadata_file_id\n self._zip_metadata_deprecated = zip_metadata\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.pdf_pass = credentials.get(\"pdf_password\")\n\n return None\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"\n Iterates over each file path, fetches from Postgres, tries to parse text\n or images, and yields Document batches.\n \"\"\"\n # Load metadata dict at start (from file store or deprecated inline format)\n zip_metadata: dict[str, Any] = {}\n if self._zip_metadata_file_id:\n try:\n file_store = get_default_file_store()\n metadata_io = file_store.read_file(\n file_id=self._zip_metadata_file_id, mode=\"b\"\n )\n metadata_bytes = metadata_io.read()\n loaded_metadata = json.loads(metadata_bytes)\n if isinstance(loaded_metadata, list):\n zip_metadata = {d[\"filename\"]: d for d in loaded_metadata}\n else:\n zip_metadata = loaded_metadata\n except Exception as e:\n logger.warning(f\"Failed to load metadata from file store: {e}\")\n elif self._zip_metadata_deprecated:\n logger.warning(\n \"Using deprecated inline zip_metadata dict. Re-upload files to use the new file store format.\"\n )\n zip_metadata = self._zip_metadata_deprecated\n\n documents: list[Document | HierarchyNode] = []\n\n for file_id in self.file_locations:\n file_store = get_default_file_store()\n file_record = file_store.read_file_record(file_id=file_id)\n if not file_record:\n # typically an unsupported extension\n logger.warning(f\"No file record found for '{file_id}' in PG; skipping.\")\n continue\n\n metadata = zip_metadata.get(\n file_record.display_name, {}\n ) or zip_metadata.get(os.path.basename(file_record.display_name), {})\n file_io = file_store.read_file(file_id=file_id, mode=\"b\")\n new_docs = _process_file(\n file_id=file_id,\n file_name=file_record.display_name,\n file=file_io,\n metadata=metadata,\n pdf_pass=self.pdf_pass,\n file_type=file_record.file_type,\n )\n documents.extend(new_docs)\n\n if len(documents) >= self.batch_size:\n yield documents\n\n documents = []\n\n if documents:\n yield documents\n\n\nif __name__ == \"__main__\":\n connector = LocalFileConnector(\n file_locations=[os.environ[\"TEST_FILE\"]],\n file_names=[os.environ[\"TEST_FILE\"]],\n )\n connector.load_credentials({\"pdf_password\": os.environ.get(\"PDF_PASSWORD\")})\n doc_batches = connector.load_from_state()\n for batch in doc_batches:\n print(\"BATCH:\", batch)\n" }, { "path": "backend/onyx/connectors/fireflies/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/fireflies/connector.py", "content": "from collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import cast\nfrom typing import List\n\nimport requests\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_FIREFLIES_ID_PREFIX = \"FIREFLIES_\"\n\n_FIREFLIES_API_URL = \"https://api.fireflies.ai/graphql\"\n\n_FIREFLIES_TRANSCRIPT_QUERY_SIZE = 50 # Max page size is 50\n\n_FIREFLIES_API_QUERY = \"\"\"\n query Transcripts($fromDate: DateTime, $toDate: DateTime, $limit: Int!, $skip: Int!) {\n transcripts(fromDate: $fromDate, toDate: $toDate, limit: $limit, skip: $skip) {\n id\n title\n organizer_email\n participants\n date\n duration\n transcript_url\n sentences {\n text\n speaker_name\n start_time\n }\n }\n }\n\"\"\"\n\nONE_MINUTE = 60\n\n\ndef _create_doc_from_transcript(transcript: dict) -> Document | None:\n sections: List[TextSection] = []\n current_speaker_name = None\n current_link = \"\"\n current_text = \"\"\n\n if transcript[\"sentences\"] is None:\n return None\n\n for sentence in transcript[\"sentences\"]:\n if sentence[\"speaker_name\"] != current_speaker_name:\n if current_speaker_name is not None:\n sections.append(\n TextSection(\n link=current_link,\n text=current_text.strip(),\n )\n )\n current_speaker_name = sentence.get(\"speaker_name\") or \"Unknown Speaker\"\n current_link = f\"{transcript['transcript_url']}?t={sentence['start_time']}\"\n current_text = f\"{current_speaker_name}: \"\n\n cleaned_text = sentence[\"text\"].replace(\"\\xa0\", \" \")\n current_text += f\"{cleaned_text} \"\n\n # Sometimes these links (links with a timestamp) do not work, it is a bug with Fireflies.\n sections.append(\n TextSection(\n link=current_link,\n text=current_text.strip(),\n )\n )\n\n fireflies_id = _FIREFLIES_ID_PREFIX + transcript[\"id\"]\n\n meeting_title = transcript[\"title\"] or \"No Title\"\n\n meeting_date_unix = transcript[\"date\"]\n meeting_date = datetime.fromtimestamp(meeting_date_unix / 1000, tz=timezone.utc)\n\n # Build hierarchy based on meeting date (year-month)\n year_month = meeting_date.strftime(\"%Y-%m\")\n\n meeting_organizer_email = transcript[\"organizer_email\"]\n organizer_email_user_info = [BasicExpertInfo(email=meeting_organizer_email)]\n\n meeting_participants_email_list = []\n for participant in transcript.get(\"participants\", []):\n if participant != meeting_organizer_email and participant:\n meeting_participants_email_list.append(BasicExpertInfo(email=participant))\n\n return Document(\n id=fireflies_id,\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.FIREFLIES,\n semantic_identifier=meeting_title,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [year_month],\n \"year_month\": year_month,\n \"meeting_title\": meeting_title,\n \"organizer_email\": meeting_organizer_email,\n }\n },\n metadata={\n k: str(v)\n for k, v in {\n \"meeting_date\": meeting_date,\n \"duration_min\": transcript.get(\"duration\"),\n }.items()\n if v is not None\n },\n doc_updated_at=meeting_date,\n primary_owners=organizer_email_user_info,\n secondary_owners=meeting_participants_email_list,\n )\n\n\n# If not all transcripts are being indexed, try using a more-recently-generated\n# API key.\nclass FirefliesConnector(PollConnector, LoadConnector):\n def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:\n self.batch_size = batch_size\n\n def load_credentials(self, credentials: dict[str, str]) -> None:\n api_key = credentials.get(\"fireflies_api_key\")\n\n if not isinstance(api_key, str):\n raise ConnectorMissingCredentialError(\n \"The Fireflies API key must be a string\"\n )\n\n self.api_key = api_key\n\n return None\n\n def _fetch_transcripts(\n self, start_datetime: str | None = None, end_datetime: str | None = None\n ) -> Iterator[List[dict]]:\n if self.api_key is None:\n raise ConnectorMissingCredentialError(\"Missing API key\")\n\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": \"Bearer \" + self.api_key,\n }\n\n skip = 0\n variables: dict[str, int | str] = {\n \"limit\": _FIREFLIES_TRANSCRIPT_QUERY_SIZE,\n }\n\n if start_datetime:\n variables[\"fromDate\"] = start_datetime\n if end_datetime:\n variables[\"toDate\"] = end_datetime\n\n while True:\n variables[\"skip\"] = skip\n response = requests.post(\n _FIREFLIES_API_URL,\n headers=headers,\n json={\"query\": _FIREFLIES_API_QUERY, \"variables\": variables},\n )\n\n response.raise_for_status()\n\n if response.status_code == 204:\n break\n\n received_transcripts = response.json()\n parsed_transcripts = received_transcripts.get(\"data\", {}).get(\n \"transcripts\", []\n )\n\n yield parsed_transcripts\n\n if len(parsed_transcripts) < _FIREFLIES_TRANSCRIPT_QUERY_SIZE:\n break\n\n skip += _FIREFLIES_TRANSCRIPT_QUERY_SIZE\n\n def _process_transcripts(\n self, start: str | None = None, end: str | None = None\n ) -> GenerateDocumentsOutput:\n doc_batch: List[Document | HierarchyNode] = []\n\n for transcript_batch in self._fetch_transcripts(start, end):\n for transcript in transcript_batch:\n if doc := _create_doc_from_transcript(transcript):\n doc_batch.append(doc)\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._process_transcripts()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n # add some leeway to account for any timezone funkiness and/or bad handling\n # of start time on the Fireflies side\n start = max(0, start - ONE_MINUTE)\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc).strftime(\n \"%Y-%m-%dT%H:%M:%S.000Z\"\n )\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc).strftime(\n \"%Y-%m-%dT%H:%M:%S.000Z\"\n )\n\n yield from self._process_transcripts(start_datetime, end_datetime)\n" }, { "path": "backend/onyx/connectors/freshdesk/__init__,py", "content": "" }, { "path": "backend/onyx/connectors/freshdesk/connector.py", "content": "import json\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import List\n\nimport requests\nfrom retry import retry\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rl_requests,\n)\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_FRESHDESK_ID_PREFIX = \"FRESHDESK_\"\n\n\n_TICKET_FIELDS_TO_INCLUDE = {\n \"fr_escalated\",\n \"spam\",\n \"priority\",\n \"source\",\n \"status\",\n \"type\",\n \"is_escalated\",\n \"tags\",\n \"nr_due_by\",\n \"nr_escalated\",\n \"cc_emails\",\n \"fwd_emails\",\n \"reply_cc_emails\",\n \"ticket_cc_emails\",\n \"support_email\",\n \"to_emails\",\n}\n\n_SOURCE_NUMBER_TYPE_MAP: dict[int, str] = {\n 1: \"Email\",\n 2: \"Portal\",\n 3: \"Phone\",\n 7: \"Chat\",\n 9: \"Feedback Widget\",\n 10: \"Outbound Email\",\n}\n\n_PRIORITY_NUMBER_TYPE_MAP: dict[int, str] = {\n 1: \"low\",\n 2: \"medium\",\n 3: \"high\",\n 4: \"urgent\",\n}\n\n_STATUS_NUMBER_TYPE_MAP: dict[int, str] = {\n 2: \"open\",\n 3: \"pending\",\n 4: \"resolved\",\n 5: \"closed\",\n}\n\n\n# TODO: unify this with other generic rate limited requests with retries (e.g. Axero, Notion?)\n@retry(tries=3, delay=1, backoff=2)\ndef _rate_limited_freshdesk_get(\n url: str, auth: tuple, params: dict\n) -> requests.Response:\n return rl_requests.get(url, auth=auth, params=params)\n\n\ndef _create_metadata_from_ticket(ticket: dict) -> dict:\n metadata: dict[str, str | list[str]] = {}\n # Combine all emails into a list so there are no repeated emails\n email_data: set[str] = set()\n\n for key, value in ticket.items():\n # Skip fields that aren't useful for embedding\n if key not in _TICKET_FIELDS_TO_INCLUDE:\n continue\n\n # Skip empty fields\n if not value or value == \"[]\":\n continue\n\n # Convert strings or lists to strings\n stringified_value: str | list[str]\n if isinstance(value, list):\n stringified_value = [str(item) for item in value]\n else:\n stringified_value = str(value)\n\n if \"email\" in key:\n if isinstance(stringified_value, list):\n email_data.update(stringified_value)\n else:\n email_data.add(stringified_value)\n else:\n metadata[key] = stringified_value\n\n if email_data:\n metadata[\"emails\"] = list(email_data)\n\n # Convert source numbers to human-parsable string\n if source_number := ticket.get(\"source\"):\n metadata[\"source\"] = _SOURCE_NUMBER_TYPE_MAP.get(\n source_number, \"Unknown Source Type\"\n )\n\n # Convert priority numbers to human-parsable string\n if priority_number := ticket.get(\"priority\"):\n metadata[\"priority\"] = _PRIORITY_NUMBER_TYPE_MAP.get(\n priority_number, \"Unknown Priority\"\n )\n\n # Convert status to human-parsable string\n if status_number := ticket.get(\"status\"):\n metadata[\"status\"] = _STATUS_NUMBER_TYPE_MAP.get(\n status_number, \"Unknown Status\"\n )\n\n due_by = datetime.fromisoformat(ticket[\"due_by\"].replace(\"Z\", \"+00:00\"))\n metadata[\"overdue\"] = str(datetime.now(timezone.utc) > due_by)\n\n return metadata\n\n\ndef _create_doc_from_ticket(ticket: dict, domain: str) -> Document:\n # Use the ticket description as the text\n text = f\"Ticket description: {parse_html_page_basic(ticket.get('description_text', ''))}\"\n metadata = _create_metadata_from_ticket(ticket)\n\n # This is also used in the ID because it is more unique than the just the ticket ID\n link = f\"https://{domain}.freshdesk.com/helpdesk/tickets/{ticket['id']}\"\n\n return Document(\n id=_FRESHDESK_ID_PREFIX + link,\n sections=[\n TextSection(\n link=link,\n text=text,\n )\n ],\n source=DocumentSource.FRESHDESK,\n semantic_identifier=ticket[\"subject\"],\n metadata=metadata,\n doc_updated_at=datetime.fromisoformat(\n ticket[\"updated_at\"].replace(\"Z\", \"+00:00\")\n ),\n )\n\n\nclass FreshdeskConnector(PollConnector, LoadConnector):\n def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:\n self.batch_size = batch_size\n\n def load_credentials(self, credentials: dict[str, str | int]) -> None:\n api_key = credentials.get(\"freshdesk_api_key\")\n domain = credentials.get(\"freshdesk_domain\")\n if not all(isinstance(cred, str) for cred in [domain, api_key]):\n raise ConnectorMissingCredentialError(\n \"All Freshdesk credentials must be strings\"\n )\n\n # TODO: Move the domain to the connector-specific configuration instead of part of the credential\n # Then apply normalization and validation against the config\n # Clean and normalize the domain URL\n domain = str(domain).strip().lower()\n\n # Remove any trailing slashes\n domain = domain.rstrip(\"/\")\n\n # Remove protocol if present\n if domain.startswith((\"http://\", \"https://\")):\n domain = domain.replace(\"http://\", \"\").replace(\"https://\", \"\")\n\n # Remove .freshdesk.com suffix and any API paths if present\n if \".freshdesk.com\" in domain:\n domain = domain.split(\".freshdesk.com\")[0]\n\n if not domain:\n raise ConnectorMissingCredentialError(\"Freshdesk domain cannot be empty\")\n\n self.api_key = str(api_key)\n self.domain = domain\n\n def _fetch_tickets(\n self,\n start: datetime | None = None,\n end: datetime | None = None, # noqa: ARG002\n ) -> Iterator[List[dict]]:\n \"\"\"\n 'end' is not currently used, so we may double fetch tickets created after the indexing\n starts but before the actual call is made.\n\n To use 'end' would require us to use the search endpoint but it has limitations,\n namely having to fetch all IDs and then individually fetch each ticket because there is no\n 'include' field available for this endpoint:\n https://developers.freshdesk.com/api/#filter_tickets\n \"\"\"\n if self.api_key is None or self.domain is None:\n raise ConnectorMissingCredentialError(\"freshdesk\")\n\n base_url = f\"https://{self.domain}.freshdesk.com/api/v2/tickets\"\n params: dict[str, int | str] = {\n \"include\": \"description\",\n \"per_page\": 50,\n \"page\": 1,\n }\n\n if start:\n params[\"updated_since\"] = start.isoformat()\n\n while True:\n # Freshdesk API uses API key as the username and any value as the password.\n response = _rate_limited_freshdesk_get(\n base_url,\n auth=(self.api_key, \"CanYouBelieveFreshdeskDoesThis\"),\n params=params,\n )\n response.raise_for_status()\n\n if response.status_code == 204:\n break\n\n tickets = json.loads(response.content)\n logger.info(\n f\"Fetched {len(tickets)} tickets from Freshdesk API (Page {params['page']})\"\n )\n\n yield tickets\n\n if len(tickets) < int(params[\"per_page\"]):\n break\n\n params[\"page\"] = int(params[\"page\"]) + 1\n\n def _process_tickets(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n doc_batch: List[Document | HierarchyNode] = []\n\n for ticket_batch in self._fetch_tickets(start, end):\n for ticket in ticket_batch:\n doc_batch.append(_create_doc_from_ticket(ticket, self.domain))\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._process_tickets()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n\n yield from self._process_tickets(start_datetime, end_datetime)\n" }, { "path": "backend/onyx/connectors/gitbook/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/gitbook/connector.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom urllib.parse import urljoin\n\nimport requests\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\nGITBOOK_API_BASE = \"https://api.gitbook.com/v1/\"\n\n\nclass GitbookApiClient:\n def __init__(self, access_token: str) -> None:\n self.access_token = access_token\n\n def get(self, endpoint: str, params: dict[str, Any] | None = None) -> Any:\n headers = {\n \"Authorization\": f\"Bearer {self.access_token}\",\n \"Content-Type\": \"application/json\",\n }\n\n url = urljoin(GITBOOK_API_BASE, endpoint.lstrip(\"/\"))\n response = requests.get(url, headers=headers, params=params)\n response.raise_for_status()\n return response.json()\n\n def get_page_content(self, space_id: str, page_id: str) -> dict[str, Any]:\n return self.get(f\"/spaces/{space_id}/content/page/{page_id}\")\n\n\ndef _extract_text_from_document(document: dict[str, Any]) -> str:\n \"\"\"Extract text content from GitBook document structure by parsing the document nodes\n into markdown format.\"\"\"\n\n def parse_leaf(leaf: dict[str, Any]) -> str:\n text = leaf.get(\"text\", \"\")\n leaf.get(\"marks\", [])\n return text\n\n def parse_text_node(node: dict[str, Any]) -> str:\n text = \"\"\n for leaf in node.get(\"leaves\", []):\n text += parse_leaf(leaf)\n return text\n\n def parse_block_node(node: dict[str, Any]) -> str:\n block_type = node.get(\"type\", \"\")\n result = \"\"\n\n if block_type == \"heading-1\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"# {text}\\n\\n\"\n\n elif block_type == \"heading-2\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"## {text}\\n\\n\"\n\n elif block_type == \"heading-3\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"### {text}\\n\\n\"\n\n elif block_type == \"heading-4\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"#### {text}\\n\\n\"\n\n elif block_type == \"heading-5\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"##### {text}\\n\\n\"\n\n elif block_type == \"heading-6\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"###### {text}\\n\\n\"\n\n elif block_type == \"list-unordered\":\n for list_item in node.get(\"nodes\", []):\n paragraph = list_item.get(\"nodes\", [])[0]\n text = \"\".join(parse_text_node(n) for n in paragraph.get(\"nodes\", []))\n result += f\"* {text}\\n\"\n result += \"\\n\"\n\n elif block_type == \"paragraph\":\n text = \"\".join(parse_text_node(n) for n in node.get(\"nodes\", []))\n result = f\"{text}\\n\\n\"\n\n elif block_type == \"list-tasks\":\n for task_item in node.get(\"nodes\", []):\n checked = task_item.get(\"data\", {}).get(\"checked\", False)\n paragraph = task_item.get(\"nodes\", [])[0]\n text = \"\".join(parse_text_node(n) for n in paragraph.get(\"nodes\", []))\n checkbox = \"[x]\" if checked else \"[ ]\"\n result += f\"- {checkbox} {text}\\n\"\n result += \"\\n\"\n\n elif block_type == \"code\":\n for code_line in node.get(\"nodes\", []):\n if code_line.get(\"type\") == \"code-line\":\n text = \"\".join(\n parse_text_node(n) for n in code_line.get(\"nodes\", [])\n )\n result += f\"{text}\\n\"\n result += \"\\n\"\n\n elif block_type == \"blockquote\":\n for quote_node in node.get(\"nodes\", []):\n if quote_node.get(\"type\") == \"paragraph\":\n text = \"\".join(\n parse_text_node(n) for n in quote_node.get(\"nodes\", [])\n )\n result += f\"> {text}\\n\"\n result += \"\\n\"\n\n elif block_type == \"table\":\n records = node.get(\"data\", {}).get(\"records\", {})\n definition = node.get(\"data\", {}).get(\"definition\", {})\n view = node.get(\"data\", {}).get(\"view\", {})\n\n columns = view.get(\"columns\", [])\n\n header_cells = []\n for col_id in columns:\n col_def = definition.get(col_id, {})\n header_cells.append(col_def.get(\"title\", \"\"))\n\n result = \"| \" + \" | \".join(header_cells) + \" |\\n\"\n result += \"|\" + \"---|\" * len(header_cells) + \"\\n\"\n\n sorted_records = sorted(\n records.items(), key=lambda x: x[1].get(\"orderIndex\", \"\")\n )\n\n for record_id, record_data in sorted_records:\n values = record_data.get(\"values\", {})\n row_cells = []\n for col_id in columns:\n fragment_id = values.get(col_id, \"\")\n fragment_text = \"\"\n for fragment in node.get(\"fragments\", []):\n if fragment.get(\"fragment\") == fragment_id:\n for frag_node in fragment.get(\"nodes\", []):\n if frag_node.get(\"type\") == \"paragraph\":\n fragment_text = \"\".join(\n parse_text_node(n)\n for n in frag_node.get(\"nodes\", [])\n )\n break\n row_cells.append(fragment_text)\n result += \"| \" + \" | \".join(row_cells) + \" |\\n\"\n\n result += \"\\n\"\n return result\n\n if not document or \"document\" not in document:\n return \"\"\n\n markdown = \"\"\n nodes = document[\"document\"].get(\"nodes\", [])\n\n for node in nodes:\n markdown += parse_block_node(node)\n\n return markdown\n\n\ndef _convert_page_to_document(\n client: GitbookApiClient, space_id: str, page: dict[str, Any]\n) -> Document:\n page_id = page[\"id\"]\n page_content = client.get_page_content(space_id, page_id)\n\n return Document(\n id=f\"gitbook-{space_id}-{page_id}\",\n sections=[\n TextSection(\n link=page.get(\"urls\", {}).get(\"app\", \"\"),\n text=_extract_text_from_document(page_content),\n )\n ],\n source=DocumentSource.GITBOOK,\n semantic_identifier=page.get(\"title\", \"\"),\n doc_updated_at=datetime.fromisoformat(page[\"updatedAt\"]).replace(\n tzinfo=timezone.utc\n ),\n metadata={\n \"path\": page.get(\"path\", \"\"),\n \"type\": page.get(\"type\", \"\"),\n \"kind\": page.get(\"kind\", \"\"),\n },\n )\n\n\nclass GitbookConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n space_id: str,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.space_id = space_id\n self.batch_size = batch_size\n self.access_token: str | None = None\n self.client: GitbookApiClient | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> None:\n access_token = credentials.get(\"gitbook_api_key\")\n if not access_token:\n raise ConnectorMissingCredentialError(\"GitBook access token\")\n self.access_token = access_token\n self.client = GitbookApiClient(access_token)\n\n def _fetch_all_pages(\n self,\n start: datetime | None = None,\n end: datetime | None = None,\n ) -> GenerateDocumentsOutput:\n if not self.client:\n raise ConnectorMissingCredentialError(\"GitBook\")\n\n try:\n content = self.client.get(f\"/spaces/{self.space_id}/content/pages\")\n pages: list[dict[str, Any]] = content.get(\"pages\", [])\n current_batch: list[Document | HierarchyNode] = []\n\n logger.info(f\"Found {len(pages)} root pages.\")\n logger.info(\n f\"First 20 Page Ids: {[page.get('id', 'Unknown') for page in pages[:20]]}\"\n )\n\n while pages:\n page = pages.pop(0)\n\n updated_at_raw = page.get(\"updatedAt\")\n if updated_at_raw is None:\n # if updatedAt is not present, that means the page has never been edited\n continue\n\n updated_at = datetime.fromisoformat(updated_at_raw)\n if start and updated_at < start:\n continue\n if end and updated_at > end:\n continue\n\n current_batch.append(\n _convert_page_to_document(self.client, self.space_id, page)\n )\n\n if len(current_batch) >= self.batch_size:\n yield current_batch\n current_batch = []\n\n pages.extend(page.get(\"pages\", []))\n\n if current_batch:\n yield current_batch\n\n except requests.RequestException as e:\n logger.error(f\"Error fetching GitBook content: {str(e)}\")\n raise\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._fetch_all_pages()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n return self._fetch_all_pages(start_datetime, end_datetime)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = GitbookConnector(\n space_id=os.environ[\"GITBOOK_SPACE_ID\"],\n )\n connector.load_credentials({\"gitbook_api_key\": os.environ[\"GITBOOK_API_KEY\"]})\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/github/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/github/connector.py", "content": "import copy\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\n\nfrom github import Github\nfrom github import RateLimitExceededException\nfrom github import Repository\nfrom github.GithubException import GithubException\nfrom github.Issue import Issue\nfrom github.NamedUser import NamedUser\nfrom github.PaginatedList import PaginatedList\nfrom github.PullRequest import PullRequest\nfrom pydantic import BaseModel\nfrom typing_extensions import override\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.app_configs import GITHUB_CONNECTOR_BASE_URL\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.connector_runner import ConnectorRunner\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.github.models import SerializedRepository\nfrom onyx.connectors.github.rate_limit_utils import sleep_after_rate_limit_exception\nfrom onyx.connectors.github.utils import deserialize_repository\nfrom onyx.connectors.github.utils import get_external_access_permission\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import ConnectorCheckpoint\nfrom onyx.connectors.interfaces import ConnectorFailure\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nITEMS_PER_PAGE = 100\nCURSOR_LOG_FREQUENCY = 50\n\n_MAX_NUM_RATE_LIMIT_RETRIES = 5\n\nONE_DAY = timedelta(days=1)\nSLIM_BATCH_SIZE = 100\n# Cases\n# X (from start) standard run, no fallback to cursor-based pagination\n# X (from start) standard run errors, fallback to cursor-based pagination\n# X error in the middle of a page\n# X no errors: run to completion\n# X (from checkpoint) standard run, no fallback to cursor-based pagination\n# X (from checkpoint) continue from cursor-based pagination\n# - retrying\n# - no retrying\n\n# things to check:\n# checkpoint state on return\n# checkpoint progress (no infinite loop)\n\n\nclass DocMetadata(BaseModel):\n repo: str\n\n\ndef get_nextUrl_key(pag_list: PaginatedList[PullRequest | Issue]) -> str:\n if \"_PaginatedList__nextUrl\" in pag_list.__dict__:\n return \"_PaginatedList__nextUrl\"\n for key in pag_list.__dict__:\n if \"__nextUrl\" in key:\n return key\n for key in pag_list.__dict__:\n if \"nextUrl\" in key:\n return key\n return \"\"\n\n\ndef get_nextUrl(\n pag_list: PaginatedList[PullRequest | Issue], nextUrl_key: str\n) -> str | None:\n return getattr(pag_list, nextUrl_key) if nextUrl_key else None\n\n\ndef set_nextUrl(\n pag_list: PaginatedList[PullRequest | Issue], nextUrl_key: str, nextUrl: str\n) -> None:\n if nextUrl_key:\n setattr(pag_list, nextUrl_key, nextUrl)\n elif nextUrl:\n raise ValueError(\"Next URL key not found: \" + str(pag_list.__dict__))\n\n\ndef _paginate_until_error(\n git_objs: Callable[[], PaginatedList[PullRequest | Issue]],\n cursor_url: str | None,\n prev_num_objs: int,\n cursor_url_callback: Callable[[str | None, int], None],\n retrying: bool = False,\n) -> Generator[PullRequest | Issue, None, None]:\n num_objs = prev_num_objs\n pag_list = git_objs()\n nextUrl_key = get_nextUrl_key(pag_list)\n if cursor_url:\n set_nextUrl(pag_list, nextUrl_key, cursor_url)\n elif retrying:\n # if we are retrying, we want to skip the objects retrieved\n # over previous calls. Unfortunately, this WILL retrieve all\n # pages before the one we are resuming from, so we really\n # don't want this case to be hit often\n logger.warning(\n \"Retrying from a previous cursor-based pagination call. \"\n \"This will retrieve all pages before the one we are resuming from, \"\n \"which may take a while and consume many API calls.\"\n )\n pag_list = cast(PaginatedList[PullRequest | Issue], pag_list[prev_num_objs:])\n num_objs = 0\n\n try:\n # this for loop handles cursor-based pagination\n for issue_or_pr in pag_list:\n num_objs += 1\n yield issue_or_pr\n # used to store the current cursor url in the checkpoint. This value\n # is updated during iteration over pag_list.\n cursor_url_callback(get_nextUrl(pag_list, nextUrl_key), num_objs)\n\n if num_objs % CURSOR_LOG_FREQUENCY == 0:\n logger.info(\n f\"Retrieved {num_objs} objects with current cursor url: {get_nextUrl(pag_list, nextUrl_key)}\"\n )\n\n except Exception as e:\n logger.exception(f\"Error during cursor-based pagination: {e}\")\n if num_objs - prev_num_objs > 0:\n raise\n\n if get_nextUrl(pag_list, nextUrl_key) is not None and not retrying:\n logger.info(\n \"Assuming that this error is due to cursor \"\n \"expiration because no objects were retrieved. \"\n \"Retrying from the first page.\"\n )\n yield from _paginate_until_error(\n git_objs, None, prev_num_objs, cursor_url_callback, retrying=True\n )\n return\n\n # for no cursor url or if we reach this point after a retry, raise the error\n raise\n\n\ndef _get_batch_rate_limited(\n # We pass in a callable because we want git_objs to produce a fresh\n # PaginatedList each time it's called to avoid using the same object for cursor-based pagination\n # from a partial offset-based pagination call.\n git_objs: Callable[[], PaginatedList],\n page_num: int,\n cursor_url: str | None,\n prev_num_objs: int,\n cursor_url_callback: Callable[[str | None, int], None],\n github_client: Github,\n attempt_num: int = 0,\n) -> Generator[PullRequest | Issue, None, None]:\n if attempt_num > _MAX_NUM_RATE_LIMIT_RETRIES:\n raise RuntimeError(\n \"Re-tried fetching batch too many times. Something is going wrong with fetching objects from Github\"\n )\n try:\n if cursor_url:\n # when this is set, we are resuming from an earlier\n # cursor-based pagination call.\n yield from _paginate_until_error(\n git_objs, cursor_url, prev_num_objs, cursor_url_callback\n )\n return\n objs = list(git_objs().get_page(page_num))\n # fetch all data here to disable lazy loading later\n # this is needed to capture the rate limit exception here (if one occurs)\n for obj in objs:\n if hasattr(obj, \"raw_data\"):\n getattr(obj, \"raw_data\")\n yield from objs\n except RateLimitExceededException:\n sleep_after_rate_limit_exception(github_client)\n yield from _get_batch_rate_limited(\n git_objs,\n page_num,\n cursor_url,\n prev_num_objs,\n cursor_url_callback,\n github_client,\n attempt_num + 1,\n )\n except GithubException as e:\n if not (\n e.status == 422\n and (\n \"cursor\" in (e.message or \"\")\n or \"cursor\" in (e.data or {}).get(\"message\", \"\")\n )\n ):\n raise\n # Fallback to a cursor-based pagination strategy\n # This can happen for \"large datasets,\" but there's no documentation\n # On the error on the web as far as we can tell.\n # Error message:\n # \"Pagination with the page parameter is not supported for large datasets,\n # please use cursor based pagination (after/before)\"\n yield from _paginate_until_error(\n git_objs, cursor_url, prev_num_objs, cursor_url_callback\n )\n\n\ndef _get_userinfo(user: NamedUser) -> dict[str, str]:\n def _safe_get(attr_name: str) -> str | None:\n try:\n return cast(str | None, getattr(user, attr_name))\n except GithubException:\n logger.debug(f\"Error getting {attr_name} for user\")\n return None\n\n return {\n k: v\n for k, v in {\n \"login\": _safe_get(\"login\"),\n \"name\": _safe_get(\"name\"),\n \"email\": _safe_get(\"email\"),\n }.items()\n if v is not None\n }\n\n\ndef _convert_pr_to_document(\n pull_request: PullRequest, repo_external_access: ExternalAccess | None\n) -> Document:\n repo_full_name = pull_request.base.repo.full_name if pull_request.base else \"\"\n # Split full_name (e.g., \"owner/repo\") into owner and repo\n parts = repo_full_name.split(\"/\", 1)\n owner_name = parts[0] if parts else \"\"\n repo_name = parts[1] if len(parts) > 1 else repo_full_name\n\n doc_metadata = {\n \"repo\": repo_full_name,\n \"hierarchy\": {\n \"source_path\": [owner_name, repo_name, \"pull_requests\"],\n \"owner\": owner_name,\n \"repo\": repo_name,\n \"object_type\": \"pull_request\",\n },\n }\n return Document(\n id=pull_request.html_url,\n sections=[\n TextSection(link=pull_request.html_url, text=pull_request.body or \"\")\n ],\n external_access=repo_external_access,\n source=DocumentSource.GITHUB,\n semantic_identifier=f\"{pull_request.number}: {pull_request.title}\",\n # updated_at is UTC time but is timezone unaware, explicitly add UTC\n # as there is logic in indexing to prevent wrong timestamped docs\n # due to local time discrepancies with UTC\n doc_updated_at=(\n pull_request.updated_at.replace(tzinfo=timezone.utc)\n if pull_request.updated_at\n else None\n ),\n # this metadata is used in perm sync\n doc_metadata=doc_metadata,\n metadata={\n k: [str(vi) for vi in v] if isinstance(v, list) else str(v)\n for k, v in {\n \"object_type\": \"PullRequest\",\n \"id\": pull_request.number,\n \"merged\": pull_request.merged,\n \"state\": pull_request.state,\n \"user\": _get_userinfo(pull_request.user) if pull_request.user else None,\n \"assignees\": [\n _get_userinfo(assignee) for assignee in pull_request.assignees\n ],\n \"repo\": (\n pull_request.base.repo.full_name if pull_request.base else None\n ),\n \"num_commits\": str(pull_request.commits),\n \"num_files_changed\": str(pull_request.changed_files),\n \"labels\": [label.name for label in pull_request.labels],\n \"created_at\": (\n pull_request.created_at.replace(tzinfo=timezone.utc)\n if pull_request.created_at\n else None\n ),\n \"updated_at\": (\n pull_request.updated_at.replace(tzinfo=timezone.utc)\n if pull_request.updated_at\n else None\n ),\n \"closed_at\": (\n pull_request.closed_at.replace(tzinfo=timezone.utc)\n if pull_request.closed_at\n else None\n ),\n \"merged_at\": (\n pull_request.merged_at.replace(tzinfo=timezone.utc)\n if pull_request.merged_at\n else None\n ),\n \"merged_by\": (\n _get_userinfo(pull_request.merged_by)\n if pull_request.merged_by\n else None\n ),\n }.items()\n if v is not None\n },\n )\n\n\ndef _fetch_issue_comments(issue: Issue) -> str:\n comments = issue.get_comments()\n return \"\\nComment: \".join(comment.body for comment in comments)\n\n\ndef _convert_issue_to_document(\n issue: Issue, repo_external_access: ExternalAccess | None\n) -> Document:\n repo_full_name = issue.repository.full_name if issue.repository else \"\"\n # Split full_name (e.g., \"owner/repo\") into owner and repo\n parts = repo_full_name.split(\"/\", 1)\n owner_name = parts[0] if parts else \"\"\n repo_name = parts[1] if len(parts) > 1 else repo_full_name\n\n doc_metadata = {\n \"repo\": repo_full_name,\n \"hierarchy\": {\n \"source_path\": [owner_name, repo_name, \"issues\"],\n \"owner\": owner_name,\n \"repo\": repo_name,\n \"object_type\": \"issue\",\n },\n }\n return Document(\n id=issue.html_url,\n sections=[TextSection(link=issue.html_url, text=issue.body or \"\")],\n source=DocumentSource.GITHUB,\n external_access=repo_external_access,\n semantic_identifier=f\"{issue.number}: {issue.title}\",\n # updated_at is UTC time but is timezone unaware\n doc_updated_at=issue.updated_at.replace(tzinfo=timezone.utc),\n # this metadata is used in perm sync\n doc_metadata=doc_metadata,\n metadata={\n k: [str(vi) for vi in v] if isinstance(v, list) else str(v)\n for k, v in {\n \"object_type\": \"Issue\",\n \"id\": issue.number,\n \"state\": issue.state,\n \"user\": _get_userinfo(issue.user) if issue.user else None,\n \"assignees\": [_get_userinfo(assignee) for assignee in issue.assignees],\n \"repo\": issue.repository.full_name if issue.repository else None,\n \"labels\": [label.name for label in issue.labels],\n \"created_at\": (\n issue.created_at.replace(tzinfo=timezone.utc)\n if issue.created_at\n else None\n ),\n \"updated_at\": (\n issue.updated_at.replace(tzinfo=timezone.utc)\n if issue.updated_at\n else None\n ),\n \"closed_at\": (\n issue.closed_at.replace(tzinfo=timezone.utc)\n if issue.closed_at\n else None\n ),\n \"closed_by\": (\n _get_userinfo(issue.closed_by) if issue.closed_by else None\n ),\n }.items()\n if v is not None\n },\n )\n\n\nclass GithubConnectorStage(Enum):\n START = \"start\"\n PRS = \"prs\"\n ISSUES = \"issues\"\n\n\nclass GithubConnectorCheckpoint(ConnectorCheckpoint):\n stage: GithubConnectorStage\n curr_page: int\n\n cached_repo_ids: list[int] | None = None\n cached_repo: SerializedRepository | None = None\n\n # Used for the fallback cursor-based pagination strategy\n num_retrieved: int\n cursor_url: str | None = None\n\n def reset(self) -> None:\n \"\"\"\n Resets curr_page, num_retrieved, and cursor_url to their initial values (0, 0, None)\n \"\"\"\n self.curr_page = 0\n self.num_retrieved = 0\n self.cursor_url = None\n\n\ndef make_cursor_url_callback(\n checkpoint: GithubConnectorCheckpoint,\n) -> Callable[[str | None, int], None]:\n def cursor_url_callback(cursor_url: str | None, num_objs: int) -> None:\n # we want to maintain the old cursor url so code after retrieval\n # can determine that we are using the fallback cursor-based pagination strategy\n if cursor_url:\n checkpoint.cursor_url = cursor_url\n checkpoint.num_retrieved = num_objs\n\n return cursor_url_callback\n\n\nclass GithubConnector(CheckpointedConnectorWithPermSync[GithubConnectorCheckpoint]):\n def __init__(\n self,\n repo_owner: str,\n repositories: str | None = None,\n state_filter: str = \"all\",\n include_prs: bool = True,\n include_issues: bool = False,\n ) -> None:\n self.repo_owner = repo_owner\n self.repositories = repositories\n self.state_filter = state_filter\n self.include_prs = include_prs\n self.include_issues = include_issues\n self.github_client: Github | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n # defaults to 30 items per page, can be set to as high as 100\n self.github_client = (\n Github(\n credentials[\"github_access_token\"],\n base_url=GITHUB_CONNECTOR_BASE_URL,\n per_page=ITEMS_PER_PAGE,\n )\n if GITHUB_CONNECTOR_BASE_URL\n else Github(credentials[\"github_access_token\"], per_page=ITEMS_PER_PAGE)\n )\n return None\n\n def get_github_repo(\n self, github_client: Github, attempt_num: int = 0\n ) -> Repository.Repository:\n if attempt_num > _MAX_NUM_RATE_LIMIT_RETRIES:\n raise RuntimeError(\n \"Re-tried fetching repo too many times. Something is going wrong with fetching objects from Github\"\n )\n\n try:\n return github_client.get_repo(f\"{self.repo_owner}/{self.repositories}\")\n except RateLimitExceededException:\n sleep_after_rate_limit_exception(github_client)\n return self.get_github_repo(github_client, attempt_num + 1)\n\n def get_github_repos(\n self, github_client: Github, attempt_num: int = 0\n ) -> list[Repository.Repository]:\n \"\"\"Get specific repositories based on comma-separated repo_name string.\"\"\"\n if attempt_num > _MAX_NUM_RATE_LIMIT_RETRIES:\n raise RuntimeError(\n \"Re-tried fetching repos too many times. Something is going wrong with fetching objects from Github\"\n )\n\n try:\n repos = []\n # Split repo_name by comma and strip whitespace\n repo_names = [\n name.strip() for name in (cast(str, self.repositories)).split(\",\")\n ]\n\n for repo_name in repo_names:\n if repo_name: # Skip empty strings\n try:\n repo = github_client.get_repo(f\"{self.repo_owner}/{repo_name}\")\n repos.append(repo)\n except GithubException as e:\n logger.warning(\n f\"Could not fetch repo {self.repo_owner}/{repo_name}: {e}\"\n )\n\n return repos\n except RateLimitExceededException:\n sleep_after_rate_limit_exception(github_client)\n return self.get_github_repos(github_client, attempt_num + 1)\n\n def get_all_repos(\n self, github_client: Github, attempt_num: int = 0\n ) -> list[Repository.Repository]:\n if attempt_num > _MAX_NUM_RATE_LIMIT_RETRIES:\n raise RuntimeError(\n \"Re-tried fetching repos too many times. Something is going wrong with fetching objects from Github\"\n )\n\n try:\n # Try to get organization first\n try:\n org = github_client.get_organization(self.repo_owner)\n return list(org.get_repos())\n\n except GithubException:\n # If not an org, try as a user\n user = github_client.get_user(self.repo_owner)\n return list(user.get_repos())\n except RateLimitExceededException:\n sleep_after_rate_limit_exception(github_client)\n return self.get_all_repos(github_client, attempt_num + 1)\n\n def fetch_configured_repos(self) -> list[Repository.Repository]:\n \"\"\"\n Fetch the configured repositories based on the connector settings.\n\n Returns:\n list[Repository.Repository]: The configured repositories.\n \"\"\"\n assert self.github_client is not None # mypy\n if self.repositories:\n if \",\" in self.repositories:\n return self.get_github_repos(self.github_client)\n else:\n return [self.get_github_repo(self.github_client)]\n else:\n return self.get_all_repos(self.github_client)\n\n def _pull_requests_func(\n self, repo: Repository.Repository\n ) -> Callable[[], PaginatedList[PullRequest]]:\n return lambda: repo.get_pulls(\n state=self.state_filter, sort=\"updated\", direction=\"desc\"\n )\n\n def _issues_func(\n self, repo: Repository.Repository\n ) -> Callable[[], PaginatedList[Issue]]:\n return lambda: repo.get_issues(\n state=self.state_filter, sort=\"updated\", direction=\"desc\"\n )\n\n def _fetch_from_github(\n self,\n checkpoint: GithubConnectorCheckpoint,\n start: datetime | None = None,\n end: datetime | None = None,\n include_permissions: bool = False,\n ) -> Generator[Document | ConnectorFailure, None, GithubConnectorCheckpoint]:\n if self.github_client is None:\n raise ConnectorMissingCredentialError(\"GitHub\")\n\n checkpoint = copy.deepcopy(checkpoint)\n\n # First run of the connector, fetch all repos and store in checkpoint\n if checkpoint.cached_repo_ids is None:\n repos = self.fetch_configured_repos()\n if not repos:\n checkpoint.has_more = False\n return checkpoint\n\n curr_repo = repos.pop()\n checkpoint.cached_repo_ids = [repo.id for repo in repos]\n checkpoint.cached_repo = SerializedRepository(\n id=curr_repo.id,\n headers=curr_repo.raw_headers,\n raw_data=curr_repo.raw_data,\n )\n checkpoint.stage = GithubConnectorStage.PRS\n checkpoint.curr_page = 0\n # save checkpoint with repo ids retrieved\n return checkpoint\n\n if checkpoint.cached_repo is None:\n raise ValueError(\"No repo saved in checkpoint\")\n\n # Deserialize the repository from the checkpoint\n repo = deserialize_repository(checkpoint.cached_repo, self.github_client)\n\n cursor_url_callback = make_cursor_url_callback(checkpoint)\n repo_external_access: ExternalAccess | None = None\n if include_permissions:\n repo_external_access = get_external_access_permission(\n repo, self.github_client\n )\n if self.include_prs and checkpoint.stage == GithubConnectorStage.PRS:\n logger.info(f\"Fetching PRs for repo: {repo.name}\")\n\n pr_batch = _get_batch_rate_limited(\n self._pull_requests_func(repo),\n checkpoint.curr_page,\n checkpoint.cursor_url,\n checkpoint.num_retrieved,\n cursor_url_callback,\n self.github_client,\n )\n checkpoint.curr_page += 1 # NOTE: not used for cursor-based fallback\n done_with_prs = False\n num_prs = 0\n pr = None\n for pr in pr_batch:\n num_prs += 1\n\n # we iterate backwards in time, so at this point we stop processing prs\n if (\n start is not None\n and pr.updated_at\n and pr.updated_at.replace(tzinfo=timezone.utc) < start\n ):\n done_with_prs = True\n break\n # Skip PRs updated after the end date\n if (\n end is not None\n and pr.updated_at\n and pr.updated_at.replace(tzinfo=timezone.utc) > end\n ):\n continue\n try:\n yield _convert_pr_to_document(\n cast(PullRequest, pr), repo_external_access\n )\n except Exception as e:\n error_msg = f\"Error converting PR to document: {e}\"\n logger.exception(error_msg)\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=str(pr.id), document_link=pr.html_url\n ),\n failure_message=error_msg,\n exception=e,\n )\n continue\n\n # If we reach this point with a cursor url in the checkpoint, we were using\n # the fallback cursor-based pagination strategy. That strategy tries to get all\n # PRs, so having curosr_url set means we are done with prs. However, we need to\n # return AFTER the checkpoint reset to avoid infinite loops.\n\n # if we found any PRs on the page and there are more PRs to get, return the checkpoint.\n # In offset mode, while indexing without time constraints, the pr batch\n # will be empty when we're done.\n used_cursor = checkpoint.cursor_url is not None\n logger.info(f\"Fetched {num_prs} PRs for repo: {repo.name}\")\n if num_prs > 0 and not done_with_prs and not used_cursor:\n return checkpoint\n\n # if we went past the start date during the loop or there are no more\n # prs to get, we move on to issues\n checkpoint.stage = GithubConnectorStage.ISSUES\n checkpoint.reset()\n\n if used_cursor:\n # save the checkpoint after changing stage; next run will continue from issues\n return checkpoint\n\n checkpoint.stage = GithubConnectorStage.ISSUES\n\n if self.include_issues and checkpoint.stage == GithubConnectorStage.ISSUES:\n logger.info(f\"Fetching issues for repo: {repo.name}\")\n\n issue_batch = list(\n _get_batch_rate_limited(\n self._issues_func(repo),\n checkpoint.curr_page,\n checkpoint.cursor_url,\n checkpoint.num_retrieved,\n cursor_url_callback,\n self.github_client,\n )\n )\n logger.info(f\"Fetched {len(issue_batch)} issues for repo: {repo.name}\")\n checkpoint.curr_page += 1\n done_with_issues = False\n num_issues = 0\n for issue in issue_batch:\n num_issues += 1\n issue = cast(Issue, issue)\n # we iterate backwards in time, so at this point we stop processing prs\n if (\n start is not None\n and issue.updated_at.replace(tzinfo=timezone.utc) < start\n ):\n done_with_issues = True\n break\n # Skip PRs updated after the end date\n if (\n end is not None\n and issue.updated_at.replace(tzinfo=timezone.utc) > end\n ):\n continue\n\n if issue.pull_request is not None:\n # PRs are handled separately\n continue\n\n try:\n yield _convert_issue_to_document(issue, repo_external_access)\n except Exception as e:\n error_msg = f\"Error converting issue to document: {e}\"\n logger.exception(error_msg)\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=str(issue.id),\n document_link=issue.html_url,\n ),\n failure_message=error_msg,\n exception=e,\n )\n continue\n\n logger.info(f\"Fetched {num_issues} issues for repo: {repo.name}\")\n # if we found any issues on the page, and we're not done, return the checkpoint.\n # don't return if we're using cursor-based pagination to avoid infinite loops\n if num_issues > 0 and not done_with_issues and not checkpoint.cursor_url:\n return checkpoint\n\n # if we went past the start date during the loop or there are no more\n # issues to get, we move on to the next repo\n checkpoint.stage = GithubConnectorStage.PRS\n checkpoint.reset()\n\n checkpoint.has_more = len(checkpoint.cached_repo_ids) > 0\n if checkpoint.cached_repo_ids:\n next_id = checkpoint.cached_repo_ids.pop()\n next_repo = self.github_client.get_repo(next_id)\n checkpoint.cached_repo = SerializedRepository(\n id=next_id,\n headers=next_repo.raw_headers,\n raw_data=next_repo.raw_data,\n )\n checkpoint.stage = GithubConnectorStage.PRS\n checkpoint.reset()\n\n if checkpoint.cached_repo_ids:\n logger.info(\n f\"{len(checkpoint.cached_repo_ids)} repos remaining (IDs: {checkpoint.cached_repo_ids})\"\n )\n else:\n logger.info(\"No more repos remaining\")\n\n return checkpoint\n\n def _load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GithubConnectorCheckpoint,\n include_permissions: bool = False,\n ) -> CheckpointOutput[GithubConnectorCheckpoint]:\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n # add a day for timezone safety\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc) + ONE_DAY\n\n # Move start time back by 3 hours, since some Issues/PRs are getting dropped\n # Could be due to delayed processing on GitHub side\n # The non-updated issues since last poll will be shortcut-ed and not embedded\n adjusted_start_datetime = start_datetime - timedelta(hours=3)\n\n epoch = datetime.fromtimestamp(0, tz=timezone.utc)\n if adjusted_start_datetime < epoch:\n adjusted_start_datetime = epoch\n\n return self._fetch_from_github(\n checkpoint,\n start=adjusted_start_datetime,\n end=end_datetime,\n include_permissions=include_permissions,\n )\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GithubConnectorCheckpoint,\n ) -> CheckpointOutput[GithubConnectorCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=False\n )\n\n @override\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GithubConnectorCheckpoint,\n ) -> CheckpointOutput[GithubConnectorCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=True\n )\n\n def validate_connector_settings(self) -> None:\n if self.github_client is None:\n raise ConnectorMissingCredentialError(\"GitHub credentials not loaded.\")\n\n if not self.repo_owner:\n raise ConnectorValidationError(\n \"Invalid connector settings: 'repo_owner' must be provided.\"\n )\n\n try:\n if self.repositories:\n if \",\" in self.repositories:\n # Multiple repositories specified\n repo_names = [name.strip() for name in self.repositories.split(\",\")]\n if not repo_names:\n raise ConnectorValidationError(\n \"Invalid connector settings: No valid repository names provided.\"\n )\n\n # Validate at least one repository exists and is accessible\n valid_repos = False\n validation_errors = []\n\n for repo_name in repo_names:\n if not repo_name:\n continue\n\n try:\n test_repo = self.github_client.get_repo(\n f\"{self.repo_owner}/{repo_name}\"\n )\n logger.info(\n f\"Successfully accessed repository: {self.repo_owner}/{repo_name}\"\n )\n test_repo.get_contents(\"\")\n valid_repos = True\n # If at least one repo is valid, we can proceed\n break\n except GithubException as e:\n validation_errors.append(\n f\"Repository '{repo_name}': {e.data.get('message', str(e))}\"\n )\n\n if not valid_repos:\n error_msg = (\n \"None of the specified repositories could be accessed: \"\n )\n error_msg += \", \".join(validation_errors)\n raise ConnectorValidationError(error_msg)\n else:\n # Single repository (backward compatibility)\n test_repo = self.github_client.get_repo(\n f\"{self.repo_owner}/{self.repositories}\"\n )\n test_repo.get_contents(\"\")\n else:\n # Try to get organization first\n try:\n org = self.github_client.get_organization(self.repo_owner)\n total_count = org.get_repos().totalCount\n if total_count == 0:\n raise ConnectorValidationError(\n f\"Found no repos for organization: {self.repo_owner}. Does the credential have the right scopes?\"\n )\n except GithubException as e:\n # Check for missing SSO\n MISSING_SSO_ERROR_MESSAGE = \"You must grant your Personal Access token access to this organization\".lower()\n if MISSING_SSO_ERROR_MESSAGE in str(e).lower():\n SSO_GUIDE_LINK = (\n \"https://docs.github.com/en/enterprise-cloud@latest/authentication/\"\n \"authenticating-with-saml-single-sign-on/\"\n \"authorizing-a-personal-access-token-for-use-with-saml-single-sign-on\"\n )\n raise ConnectorValidationError(\n f\"Your GitHub token is missing authorization to access the \"\n f\"`{self.repo_owner}` organization. Please follow the guide to \"\n f\"authorize your token: {SSO_GUIDE_LINK}\"\n )\n # If not an org, try as a user\n user = self.github_client.get_user(self.repo_owner)\n\n # Check if we can access any repos\n total_count = user.get_repos().totalCount\n if total_count == 0:\n raise ConnectorValidationError(\n f\"Found no repos for user: {self.repo_owner}. Does the credential have the right scopes?\"\n )\n\n except RateLimitExceededException:\n raise UnexpectedValidationError(\n \"Validation failed due to GitHub rate-limits being exceeded. Please try again later.\"\n )\n\n except GithubException as e:\n if e.status == 401:\n raise CredentialExpiredError(\n \"GitHub credential appears to be invalid or expired (HTTP 401).\"\n )\n elif e.status == 403:\n raise InsufficientPermissionsError(\n \"Your GitHub token does not have sufficient permissions for this repository (HTTP 403).\"\n )\n elif e.status == 404:\n if self.repositories:\n if \",\" in self.repositories:\n raise ConnectorValidationError(\n f\"None of the specified GitHub repositories could be found for owner: {self.repo_owner}\"\n )\n else:\n raise ConnectorValidationError(\n f\"GitHub repository not found with name: {self.repo_owner}/{self.repositories}\"\n )\n else:\n raise ConnectorValidationError(\n f\"GitHub user or organization not found: {self.repo_owner}\"\n )\n else:\n raise ConnectorValidationError(\n f\"Unexpected GitHub error (status={e.status}): {e.data}\"\n )\n\n except Exception as exc:\n raise Exception(\n f\"Unexpected error during GitHub settings validation: {exc}\"\n )\n\n def validate_checkpoint_json(\n self, checkpoint_json: str\n ) -> GithubConnectorCheckpoint:\n return GithubConnectorCheckpoint.model_validate_json(checkpoint_json)\n\n def build_dummy_checkpoint(self) -> GithubConnectorCheckpoint:\n return GithubConnectorCheckpoint(\n stage=GithubConnectorStage.PRS, curr_page=0, has_more=True, num_retrieved=0\n )\n\n\nif __name__ == \"__main__\":\n import os\n\n # Initialize the connector\n connector = GithubConnector(\n repo_owner=os.environ[\"REPO_OWNER\"],\n repositories=os.environ.get(\"REPOSITORIES\"),\n )\n connector.load_credentials(\n {\"github_access_token\": os.environ[\"ACCESS_TOKEN_GITHUB\"]}\n )\n\n if connector.github_client:\n get_external_access_permission(\n connector.get_github_repos(connector.github_client).pop(),\n connector.github_client,\n )\n\n # Create a time range from epoch to now\n end_time = datetime.now(timezone.utc)\n start_time = datetime.fromtimestamp(0, tz=timezone.utc)\n time_range = (start_time, end_time)\n\n # Initialize the runner with a batch size of 10\n runner: ConnectorRunner[GithubConnectorCheckpoint] = ConnectorRunner(\n connector, batch_size=10, include_permissions=False, time_range=time_range\n )\n\n # Get initial checkpoint\n checkpoint = connector.build_dummy_checkpoint()\n\n # Run the connector\n while checkpoint.has_more:\n for doc_batch, hierarchy_node_batch, failure, next_checkpoint in runner.run(\n checkpoint\n ):\n if doc_batch:\n print(f\"Retrieved batch of {len(doc_batch)} documents\")\n for doc in doc_batch:\n print(f\"Document: {doc.semantic_identifier}\")\n if failure:\n print(f\"Failure: {failure.failure_message}\")\n if next_checkpoint:\n checkpoint = next_checkpoint\n" }, { "path": "backend/onyx/connectors/github/models.py", "content": "from typing import Any\n\nfrom github import Repository\nfrom github.Requester import Requester\nfrom pydantic import BaseModel\n\n\nclass SerializedRepository(BaseModel):\n # id is part of the raw_data as well, just pulled out for convenience\n id: int\n headers: dict[str, str | int]\n raw_data: dict[str, Any]\n\n def to_Repository(self, requester: Requester) -> Repository.Repository:\n return Repository.Repository(\n requester, self.headers, self.raw_data, completed=True\n )\n" }, { "path": "backend/onyx/connectors/github/rate_limit_utils.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom github import Github\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef sleep_after_rate_limit_exception(github_client: Github) -> None:\n \"\"\"\n Sleep until the GitHub rate limit resets.\n\n Args:\n github_client: The GitHub client that hit the rate limit\n \"\"\"\n sleep_time = github_client.get_rate_limit().core.reset.replace(\n tzinfo=timezone.utc\n ) - datetime.now(tz=timezone.utc)\n sleep_time += timedelta(minutes=1) # add an extra minute just to be safe\n logger.notice(f\"Ran into Github rate-limit. Sleeping {sleep_time.seconds} seconds.\")\n time.sleep(sleep_time.total_seconds())\n" }, { "path": "backend/onyx/connectors/github/utils.py", "content": "from collections.abc import Callable\nfrom typing import cast\n\nfrom github import Github\nfrom github.Repository import Repository\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.github.models import SerializedRepository\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\nlogger = setup_logger()\n\n\ndef get_external_access_permission(\n repo: Repository, github_client: Github\n) -> ExternalAccess:\n \"\"\"\n Get the external access permission for a repository.\n This functionality requires Enterprise Edition.\n \"\"\"\n # Check if EE is enabled\n if not global_version.is_ee_version():\n # For the MIT version, return an empty ExternalAccess (private document)\n return ExternalAccess.empty()\n\n # Fetch the EE implementation\n ee_get_external_access_permission = cast(\n Callable[[Repository, Github, bool], ExternalAccess],\n fetch_versioned_implementation(\n \"onyx.external_permissions.github.utils\",\n \"get_external_access_permission\",\n ),\n )\n\n return ee_get_external_access_permission(repo, github_client, True)\n\n\ndef deserialize_repository(\n cached_repo: SerializedRepository, github_client: Github\n) -> Repository:\n \"\"\"\n Deserialize a SerializedRepository back into a Repository object.\n \"\"\"\n # Try to access the requester - different PyGithub versions may use different attribute names\n try:\n # Try to get the requester using getattr to avoid linter errors\n requester = getattr(github_client, \"_requester\", None)\n if requester is None:\n requester = getattr(github_client, \"_Github__requester\", None)\n if requester is None:\n # If we can't find the requester attribute, we need to fall back to recreating the repo\n raise AttributeError(\"Could not find requester attribute\")\n\n return cached_repo.to_Repository(requester)\n except Exception as e:\n # If all else fails, re-fetch the repo directly\n logger.warning(\n f\"Failed to deserialize repository: {e}. Attempting to re-fetch.\"\n )\n repo_id = cached_repo.id\n return github_client.get_repo(repo_id)\n" }, { "path": "backend/onyx/connectors/gitlab/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/gitlab/connector.py", "content": "import fnmatch\nimport itertools\nfrom collections import deque\nfrom collections.abc import Iterable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import TypeVar\n\nimport gitlab\nimport pytz\nfrom gitlab.v4.objects import Project\n\nfrom onyx.configs.app_configs import GITLAB_CONNECTOR_INCLUDE_CODE_FILES\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nT = TypeVar(\"T\")\n\n\nlogger = setup_logger()\n\n# List of directories/Files to exclude\nexclude_patterns = [\n \"logs\",\n \".github/\",\n \".gitlab/\",\n \".pre-commit-config.yaml\",\n]\n\n\ndef _batch_gitlab_objects(git_objs: Iterable[T], batch_size: int) -> Iterator[list[T]]:\n it = iter(git_objs)\n while True:\n batch = list(itertools.islice(it, batch_size))\n if not batch:\n break\n yield batch\n\n\ndef get_author(author: Any) -> BasicExpertInfo:\n return BasicExpertInfo(\n display_name=author.get(\"name\"),\n )\n\n\ndef _convert_merge_request_to_document(mr: Any) -> Document:\n doc = Document(\n id=mr.web_url,\n sections=[TextSection(link=mr.web_url, text=mr.description or \"\")],\n source=DocumentSource.GITLAB,\n semantic_identifier=mr.title,\n # updated_at is UTC time but is timezone unaware, explicitly add UTC\n # as there is logic in indexing to prevent wrong timestamped docs\n # due to local time discrepancies with UTC\n doc_updated_at=mr.updated_at.replace(tzinfo=timezone.utc),\n primary_owners=[get_author(mr.author)],\n metadata={\"state\": mr.state, \"type\": \"MergeRequest\"},\n )\n return doc\n\n\ndef _convert_issue_to_document(issue: Any) -> Document:\n doc = Document(\n id=issue.web_url,\n sections=[TextSection(link=issue.web_url, text=issue.description or \"\")],\n source=DocumentSource.GITLAB,\n semantic_identifier=issue.title,\n # updated_at is UTC time but is timezone unaware, explicitly add UTC\n # as there is logic in indexing to prevent wrong timestamped docs\n # due to local time discrepancies with UTC\n doc_updated_at=issue.updated_at.replace(tzinfo=timezone.utc),\n primary_owners=[get_author(issue.author)],\n metadata={\"state\": issue.state, \"type\": issue.type if issue.type else \"Issue\"},\n )\n return doc\n\n\ndef _convert_code_to_document(\n project: Project, file: Any, url: str, projectName: str, projectOwner: str\n) -> Document:\n # Dynamically get the default branch from the project object\n default_branch = project.default_branch\n\n # Fetch the file content using the correct branch\n file_content_obj = project.files.get(\n file_path=file[\"path\"],\n ref=default_branch, # Use the default branch\n )\n try:\n file_content = file_content_obj.decode().decode(\"utf-8\")\n except UnicodeDecodeError:\n file_content = file_content_obj.decode().decode(\"latin-1\")\n\n # Construct the file URL dynamically using the default branch\n file_url = (\n f\"{url}/{projectOwner}/{projectName}/-/blob/{default_branch}/{file['path']}\"\n )\n\n # Create and return a Document object\n doc = Document(\n id=file[\"id\"],\n sections=[TextSection(link=file_url, text=file_content)],\n source=DocumentSource.GITLAB,\n semantic_identifier=file[\"name\"],\n doc_updated_at=datetime.now().replace(tzinfo=timezone.utc),\n primary_owners=[], # Add owners if needed\n metadata={\"type\": \"CodeFile\"},\n )\n return doc\n\n\ndef _should_exclude(path: str) -> bool:\n \"\"\"Check if a path matches any of the exclude patterns.\"\"\"\n return any(fnmatch.fnmatch(path, pattern) for pattern in exclude_patterns)\n\n\nclass GitlabConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n project_owner: str,\n project_name: str,\n batch_size: int = INDEX_BATCH_SIZE,\n state_filter: str = \"all\",\n include_mrs: bool = True,\n include_issues: bool = True,\n include_code_files: bool = GITLAB_CONNECTOR_INCLUDE_CODE_FILES,\n ) -> None:\n self.project_owner = project_owner\n self.project_name = project_name\n self.batch_size = batch_size\n self.state_filter = state_filter\n self.include_mrs = include_mrs\n self.include_issues = include_issues\n self.include_code_files = include_code_files\n self.gitlab_client: gitlab.Gitlab | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.gitlab_client = gitlab.Gitlab(\n credentials[\"gitlab_url\"], private_token=credentials[\"gitlab_access_token\"]\n )\n return None\n\n def _fetch_from_gitlab(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n if self.gitlab_client is None:\n raise ConnectorMissingCredentialError(\"Gitlab\")\n project: Project = self.gitlab_client.projects.get(\n f\"{self.project_owner}/{self.project_name}\"\n )\n\n # Fetch code files\n if self.include_code_files:\n # Fetching using BFS as project.report_tree with recursion causing slow load\n queue = deque([\"\"]) # Start with the root directory\n while queue:\n current_path = queue.popleft()\n files = project.repository_tree(path=current_path, all=True)\n for file_batch in _batch_gitlab_objects(files, self.batch_size):\n code_doc_batch: list[Document | HierarchyNode] = []\n for file in file_batch:\n if _should_exclude(file[\"path\"]):\n continue\n\n if file[\"type\"] == \"blob\":\n code_doc_batch.append(\n _convert_code_to_document(\n project,\n file,\n self.gitlab_client.url,\n self.project_name,\n self.project_owner,\n )\n )\n elif file[\"type\"] == \"tree\":\n queue.append(file[\"path\"])\n\n if code_doc_batch:\n yield code_doc_batch\n\n if self.include_mrs:\n merge_requests = project.mergerequests.list(\n state=self.state_filter,\n order_by=\"updated_at\",\n sort=\"desc\",\n iterator=True,\n )\n\n for mr_batch in _batch_gitlab_objects(merge_requests, self.batch_size):\n mr_doc_batch: list[Document | HierarchyNode] = []\n for mr in mr_batch:\n mr.updated_at = datetime.strptime(\n mr.updated_at, \"%Y-%m-%dT%H:%M:%S.%f%z\"\n )\n if start is not None and mr.updated_at < start.replace(\n tzinfo=pytz.UTC\n ):\n yield mr_doc_batch\n return\n if end is not None and mr.updated_at > end.replace(tzinfo=pytz.UTC):\n continue\n mr_doc_batch.append(_convert_merge_request_to_document(mr))\n yield mr_doc_batch\n\n if self.include_issues:\n issues = project.issues.list(state=self.state_filter, iterator=True)\n\n for issue_batch in _batch_gitlab_objects(issues, self.batch_size):\n issue_doc_batch: list[Document | HierarchyNode] = []\n for issue in issue_batch:\n issue.updated_at = datetime.strptime(\n issue.updated_at, \"%Y-%m-%dT%H:%M:%S.%f%z\"\n )\n if start is not None:\n start = start.replace(tzinfo=pytz.UTC)\n if issue.updated_at < start:\n yield issue_doc_batch\n return\n if end is not None:\n end = end.replace(tzinfo=pytz.UTC)\n if issue.updated_at > end:\n continue\n issue_doc_batch.append(_convert_issue_to_document(issue))\n yield issue_doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._fetch_from_gitlab()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n return self._fetch_from_gitlab(start_datetime, end_datetime)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = GitlabConnector(\n # gitlab_url=\"https://gitlab.com/api/v4\",\n project_owner=os.environ[\"PROJECT_OWNER\"],\n project_name=os.environ[\"PROJECT_NAME\"],\n batch_size=10,\n state_filter=\"all\",\n include_mrs=True,\n include_issues=True,\n include_code_files=GITLAB_CONNECTOR_INCLUDE_CODE_FILES,\n )\n\n connector.load_credentials(\n {\n \"gitlab_access_token\": os.environ[\"GITLAB_ACCESS_TOKEN\"],\n \"gitlab_url\": os.environ[\"GITLAB_URL\"],\n }\n )\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/gmail/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/gmail/connector.py", "content": "from base64 import urlsafe_b64decode\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\n\nfrom google.oauth2.credentials import Credentials as OAuthCredentials\nfrom google.oauth2.service_account import Credentials as ServiceAccountCredentials\nfrom googleapiclient.errors import HttpError # type: ignore\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.google_utils.google_auth import get_google_creds\nfrom onyx.connectors.google_utils.google_utils import execute_paginated_retrieval\nfrom onyx.connectors.google_utils.google_utils import (\n execute_paginated_retrieval_with_max_pages,\n)\nfrom onyx.connectors.google_utils.google_utils import execute_single_retrieval\nfrom onyx.connectors.google_utils.google_utils import PAGE_TOKEN_KEY\nfrom onyx.connectors.google_utils.resources import get_admin_service\nfrom onyx.connectors.google_utils.resources import get_gmail_service\nfrom onyx.connectors.google_utils.resources import GmailService\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR\nfrom onyx.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS\nfrom onyx.connectors.google_utils.shared_constants import SLIM_BATCH_SIZE\nfrom onyx.connectors.google_utils.shared_constants import USER_FIELDS\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import ConnectorFailure\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\n\nlogger = setup_logger()\n\n# This is for the initial list call to get the thread ids\nTHREAD_LIST_FIELDS = \"nextPageToken, threads(id)\"\n\n# These are the fields to retrieve using the ID from the initial list call\nPARTS_FIELDS = \"parts(body(data), mimeType)\"\nPAYLOAD_FIELDS = f\"payload(headers, {PARTS_FIELDS})\"\nMESSAGES_FIELDS = f\"messages(id, {PAYLOAD_FIELDS})\"\nTHREADS_FIELDS = f\"threads(id, {MESSAGES_FIELDS})\"\nTHREAD_FIELDS = f\"id, {MESSAGES_FIELDS}\"\n\nEMAIL_FIELDS = [\n \"cc\",\n \"bcc\",\n \"from\",\n \"to\",\n]\n\nMAX_MESSAGE_BODY_BYTES = 10 * 1024 * 1024 # 10MB cap to keep large threads safe\n\nPAGES_PER_CHECKPOINT = 1\n\nadd_retries = retry_builder(tries=50, max_delay=30)\n\n\ndef _is_mail_service_disabled_error(error: HttpError) -> bool:\n \"\"\"Detect if the Gmail API is telling us the mailbox is not provisioned.\"\"\"\n\n if error.resp.status != 400:\n return False\n\n error_message = str(error)\n return (\n \"Mail service not enabled\" in error_message\n or \"failedPrecondition\" in error_message\n )\n\n\ndef _build_time_range_query(\n time_range_start: SecondsSinceUnixEpoch | None = None,\n time_range_end: SecondsSinceUnixEpoch | None = None,\n) -> str | None:\n query = \"\"\n if time_range_start is not None and time_range_start != 0:\n query += f\"after:{int(time_range_start)}\"\n if time_range_end is not None and time_range_end != 0:\n query += f\" before:{int(time_range_end)}\"\n query = query.strip()\n\n if len(query) == 0:\n return None\n\n return query\n\n\ndef _clean_email_and_extract_name(email: str) -> tuple[str, str | None]:\n email = email.strip()\n if \"<\" in email and \">\" in email:\n # Handle format: \"Display Name \"\n display_name = email[: email.find(\"<\")].strip()\n email_address = email[email.find(\"<\") + 1 : email.find(\">\")].strip()\n return email_address, display_name if display_name else None\n else:\n # Handle plain email address\n return email.strip(), None\n\n\ndef _get_owners_from_emails(emails: dict[str, str | None]) -> list[BasicExpertInfo]:\n owners = []\n for email, names in emails.items():\n if names:\n name_parts = names.split(\" \")\n first_name = \" \".join(name_parts[:-1])\n last_name = name_parts[-1]\n else:\n first_name = None\n last_name = None\n owners.append(\n BasicExpertInfo(email=email, first_name=first_name, last_name=last_name)\n )\n return owners\n\n\ndef _get_message_body(payload: dict[str, Any]) -> str:\n \"\"\"\n Gmail threads can contain large inline parts (including attachments\n transmitted as base64). Only decode text/plain parts and skip anything\n that breaches the safety threshold to protect against OOMs.\n \"\"\"\n\n message_body_chunks: list[str] = []\n stack = [payload]\n\n while stack:\n part = stack.pop()\n if not part:\n continue\n\n children = part.get(\"parts\", [])\n stack.extend(reversed(children))\n\n mime_type = part.get(\"mimeType\")\n if mime_type != \"text/plain\":\n continue\n\n body = part.get(\"body\", {})\n data = body.get(\"data\", \"\")\n\n if not data:\n continue\n\n # base64 inflates storage by ~4/3; work with decoded size estimate\n approx_decoded_size = (len(data) * 3) // 4\n if approx_decoded_size > MAX_MESSAGE_BODY_BYTES:\n logger.warning(\n \"Skipping oversized Gmail message part (%s bytes > %s limit)\",\n approx_decoded_size,\n MAX_MESSAGE_BODY_BYTES,\n )\n continue\n\n try:\n text = urlsafe_b64decode(data).decode()\n except (ValueError, UnicodeDecodeError) as error:\n logger.warning(\"Failed to decode Gmail message part: %s\", error)\n continue\n\n message_body_chunks.append(text)\n\n return \"\".join(message_body_chunks)\n\n\ndef _build_document_link(thread_id: str) -> str:\n return f\"https://mail.google.com/mail/u/0/#inbox/{thread_id}\"\n\n\ndef message_to_section(message: Dict[str, Any]) -> tuple[TextSection, dict[str, str]]:\n link = _build_document_link(message[\"id\"])\n\n payload = message.get(\"payload\", {})\n headers = payload.get(\"headers\", [])\n metadata: dict[str, Any] = {}\n for header in headers:\n name = header.get(\"name\").lower()\n value = header.get(\"value\")\n if name in EMAIL_FIELDS:\n metadata[name] = value\n if name == \"subject\":\n metadata[\"subject\"] = value\n if name == \"date\":\n metadata[\"updated_at\"] = value\n\n if labels := message.get(\"labelIds\"):\n metadata[\"labels\"] = labels\n\n message_data = \"\"\n for name, value in metadata.items():\n # updated at isnt super useful for the llm\n if name != \"updated_at\":\n message_data += f\"{name}: {value}\\n\"\n\n message_body_text: str = _get_message_body(payload)\n\n return TextSection(link=link, text=message_body_text + message_data), metadata\n\n\ndef thread_to_document(\n full_thread: Dict[str, Any], email_used_to_fetch_thread: str\n) -> Document | None:\n all_messages = full_thread.get(\"messages\", [])\n if not all_messages:\n return None\n\n sections = []\n semantic_identifier = \"\"\n updated_at = None\n from_emails: dict[str, str | None] = {}\n other_emails: dict[str, str | None] = {}\n for message in all_messages:\n section, message_metadata = message_to_section(message)\n sections.append(section)\n\n for name, value in message_metadata.items():\n if name in EMAIL_FIELDS:\n email, display_name = _clean_email_and_extract_name(value)\n if name == \"from\":\n from_emails[email] = (\n display_name if not from_emails.get(email) else None\n )\n else:\n other_emails[email] = (\n display_name if not other_emails.get(email) else None\n )\n\n # If we haven't set the semantic identifier yet, set it to the subject of the first message\n if not semantic_identifier:\n semantic_identifier = message_metadata.get(\"subject\", \"\")\n\n if message_metadata.get(\"updated_at\"):\n updated_at = message_metadata.get(\"updated_at\")\n\n updated_at_datetime = None\n if updated_at:\n updated_at_datetime = time_str_to_utc(updated_at)\n\n id = full_thread.get(\"id\")\n if not id:\n raise ValueError(\"Thread ID is required\")\n\n primary_owners = _get_owners_from_emails(from_emails)\n secondary_owners = _get_owners_from_emails(other_emails)\n\n # If emails have no subject, match Gmail's default \"no subject\"\n # Search will break without a semantic identifier\n if not semantic_identifier:\n semantic_identifier = \"(no subject)\"\n\n # NOTE: we're choosing to unconditionally include perm sync info\n # (external_access) as it doesn't cost much space\n return Document(\n id=id,\n semantic_identifier=semantic_identifier,\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.GMAIL,\n # This is used to perform permission sync\n primary_owners=primary_owners,\n secondary_owners=secondary_owners,\n doc_updated_at=updated_at_datetime,\n # Not adding emails to metadata because it's already in the sections\n metadata={},\n external_access=ExternalAccess(\n external_user_emails={email_used_to_fetch_thread},\n external_user_group_ids=set(),\n is_public=False,\n ),\n )\n\n\ndef _full_thread_from_id(\n thread_id: str,\n user_email: str,\n gmail_service: GmailService,\n) -> Document | ConnectorFailure | None:\n try:\n thread = next(\n execute_single_retrieval(\n retrieval_function=gmail_service.users().threads().get,\n list_key=None,\n userId=user_email,\n fields=THREAD_FIELDS,\n id=thread_id,\n continue_on_404_or_403=True,\n ),\n None,\n )\n if thread is None:\n raise ValueError(f\"Thread {thread_id} not found\")\n return thread_to_document(thread, user_email)\n except Exception as e:\n return ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=thread_id, document_link=_build_document_link(thread_id)\n ),\n failure_message=f\"Failed to retrieve thread {thread_id}\",\n exception=e,\n )\n\n\ndef _slim_thread_from_id(\n thread_id: str,\n user_email: str,\n gmail_service: GmailService, # noqa: ARG001\n) -> SlimDocument:\n return SlimDocument(\n id=thread_id,\n external_access=ExternalAccess(\n external_user_emails={user_email},\n external_user_group_ids=set(),\n is_public=False,\n ),\n )\n\n\nclass GmailCheckpoint(ConnectorCheckpoint):\n user_emails: list[str] = [] # stack of user emails to process\n page_token: str | None = None\n\n\nclass GmailConnector(\n SlimConnectorWithPermSync, CheckpointedConnectorWithPermSync[GmailCheckpoint]\n):\n def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:\n self.batch_size = batch_size\n\n self._creds: OAuthCredentials | ServiceAccountCredentials | None = None\n self._primary_admin_email: str | None = None\n\n @property\n def primary_admin_email(self) -> str:\n if self._primary_admin_email is None:\n raise RuntimeError(\n \"Primary admin email missing, should not call this property before calling load_credentials\"\n )\n return self._primary_admin_email\n\n @property\n def google_domain(self) -> str:\n if self._primary_admin_email is None:\n raise RuntimeError(\n \"Primary admin email missing, should not call this property before calling load_credentials\"\n )\n return self._primary_admin_email.split(\"@\")[-1]\n\n @property\n def creds(self) -> OAuthCredentials | ServiceAccountCredentials:\n if self._creds is None:\n raise RuntimeError(\n \"Creds missing, should not call this property before calling load_credentials\"\n )\n return self._creds\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, str] | None:\n primary_admin_email = credentials[DB_CREDENTIALS_PRIMARY_ADMIN_KEY]\n self._primary_admin_email = primary_admin_email\n\n self._creds, new_creds_dict = get_google_creds(\n credentials=credentials,\n source=DocumentSource.GMAIL,\n )\n return new_creds_dict\n\n def _get_all_user_emails(self) -> list[str]:\n \"\"\"\n List all user emails if we are on a Google Workspace domain.\n If the domain is gmail.com, or if we attempt to call the Admin SDK and\n get a 404 or 403, fall back to using the single user.\n A 404 indicates a personal Gmail account with no Workspace domain.\n A 403 indicates insufficient permissions (e.g., OAuth user without admin privileges).\n \"\"\"\n\n try:\n admin_service = get_admin_service(self.creds, self.primary_admin_email)\n emails = []\n for user in execute_paginated_retrieval(\n retrieval_function=admin_service.users().list,\n list_key=\"users\",\n fields=USER_FIELDS,\n domain=self.google_domain,\n ):\n if email := user.get(\"primaryEmail\"):\n emails.append(email)\n return emails\n\n except HttpError as e:\n if e.resp.status == 404:\n logger.warning(\n \"Received 404 from Admin SDK; this may indicate a personal Gmail account \"\n \"with no Workspace domain. Falling back to single user.\"\n )\n return [self.primary_admin_email]\n elif e.resp.status == 403:\n logger.warning(\n \"Received 403 from Admin SDK; this may indicate insufficient permissions \"\n \"(e.g., OAuth user without admin privileges or service account without \"\n \"domain-wide delegation). Falling back to single user.\"\n )\n return [self.primary_admin_email]\n raise\n\n def _fetch_threads_impl(\n self,\n user_email: str,\n time_range_start: SecondsSinceUnixEpoch | None = None,\n time_range_end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n page_token: str | None = None,\n set_page_token: Callable[[str | None], None] = lambda x: None, # noqa: ARG005\n is_slim: bool = False,\n ) -> Iterator[Document | ConnectorFailure] | GenerateSlimDocumentOutput:\n query = _build_time_range_query(time_range_start, time_range_end)\n slim_doc_batch: list[SlimDocument | HierarchyNode] = []\n logger.info(\n f\"Fetching {'slim' if is_slim else 'full'} threads for user: {user_email}\"\n )\n gmail_service = get_gmail_service(self.creds, user_email)\n try:\n for thread in execute_paginated_retrieval_with_max_pages(\n max_num_pages=PAGES_PER_CHECKPOINT,\n retrieval_function=gmail_service.users().threads().list,\n list_key=\"threads\",\n userId=user_email,\n fields=THREAD_LIST_FIELDS,\n q=query,\n continue_on_404_or_403=True,\n **({PAGE_TOKEN_KEY: page_token} if page_token else {}),\n ):\n # if a page token is returned, set it and leave the function\n if isinstance(thread, str):\n set_page_token(thread)\n return\n if is_slim:\n slim_doc_batch.append(\n SlimDocument(\n id=thread[\"id\"],\n external_access=ExternalAccess(\n external_user_emails={user_email},\n external_user_group_ids=set(),\n is_public=False,\n ),\n )\n )\n if len(slim_doc_batch) >= SLIM_BATCH_SIZE:\n yield slim_doc_batch\n slim_doc_batch = []\n else:\n result = _full_thread_from_id(\n thread[\"id\"], user_email, gmail_service\n )\n if result is not None:\n yield result\n if callback:\n tag = (\n \"retrieve_all_slim_docs_perm_sync\"\n if is_slim\n else \"gmail_retrieve_all_docs\"\n )\n if callback.should_stop():\n raise RuntimeError(f\"{tag}: Stop signal detected\")\n\n callback.progress(tag, 1)\n if slim_doc_batch:\n yield slim_doc_batch\n\n # done with user\n set_page_token(None)\n except HttpError as e:\n if _is_mail_service_disabled_error(e):\n logger.warning(\n \"Skipping Gmail sync for %s because the mailbox is disabled.\",\n user_email,\n )\n return\n raise\n\n def _fetch_threads(\n self,\n user_email: str,\n page_token: str | None = None,\n set_page_token: Callable[[str | None], None] = lambda x: None, # noqa: ARG005\n time_range_start: SecondsSinceUnixEpoch | None = None,\n time_range_end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> Iterator[Document | ConnectorFailure]:\n yield from cast(\n Iterator[Document | ConnectorFailure],\n self._fetch_threads_impl(\n user_email,\n time_range_start,\n time_range_end,\n callback,\n page_token,\n set_page_token,\n False,\n ),\n )\n\n def _fetch_slim_threads(\n self,\n user_email: str,\n page_token: str | None = None,\n set_page_token: Callable[[str | None], None] = lambda x: None, # noqa: ARG005\n time_range_start: SecondsSinceUnixEpoch | None = None,\n time_range_end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n yield from cast(\n GenerateSlimDocumentOutput,\n self._fetch_threads_impl(\n user_email,\n time_range_start,\n time_range_end,\n callback,\n page_token,\n set_page_token,\n True,\n ),\n )\n\n def _load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GmailCheckpoint,\n ) -> CheckpointOutput[GmailCheckpoint]:\n if not checkpoint.user_emails:\n checkpoint.user_emails = self._get_all_user_emails()\n try:\n\n def set_page_token(page_token: str | None) -> None:\n checkpoint.page_token = page_token\n\n yield from self._fetch_threads(\n checkpoint.user_emails[-1],\n checkpoint.page_token,\n set_page_token,\n start,\n end,\n callback=None,\n )\n if checkpoint.page_token is None:\n # we're done with this user\n checkpoint.user_emails.pop()\n\n if len(checkpoint.user_emails) == 0:\n checkpoint.has_more = False\n return checkpoint\n except Exception as e:\n if MISSING_SCOPES_ERROR_STR in str(e):\n raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e\n raise e\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GmailCheckpoint,\n ) -> CheckpointOutput[GmailCheckpoint]:\n return self._load_from_checkpoint(\n start=start,\n end=end,\n checkpoint=checkpoint,\n )\n\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GmailCheckpoint,\n ) -> CheckpointOutput[GmailCheckpoint]:\n # NOTE: we're choosing to unconditionally include perm sync info\n # (external_access) as it doesn't cost much space\n return self._load_from_checkpoint(\n start=start,\n end=end,\n checkpoint=checkpoint,\n )\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n try:\n pt_dict: dict[str, str | None] = {PAGE_TOKEN_KEY: None}\n\n def set_page_token(page_token: str | None) -> None:\n pt_dict[PAGE_TOKEN_KEY] = page_token\n\n for user_email in self._get_all_user_emails():\n yield from self._fetch_slim_threads(\n user_email,\n pt_dict[PAGE_TOKEN_KEY],\n set_page_token,\n start,\n end,\n callback=callback,\n )\n except Exception as e:\n if MISSING_SCOPES_ERROR_STR in str(e):\n raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e\n raise e\n\n def build_dummy_checkpoint(self) -> GmailCheckpoint:\n return GmailCheckpoint(has_more=True)\n\n def validate_checkpoint_json(self, checkpoint_json: str) -> GmailCheckpoint:\n return GmailCheckpoint.model_validate_json(checkpoint_json)\n\n\nif __name__ == \"__main__\":\n pass\n" }, { "path": "backend/onyx/connectors/gong/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/gong/connector.py", "content": "import base64\nimport time\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom requests.adapters import HTTPAdapter\nfrom urllib3.util import Retry\n\nfrom onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE\nfrom onyx.configs.app_configs import GONG_CONNECTOR_START_TIME\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass GongConnector(LoadConnector, PollConnector):\n BASE_URL = \"https://api.gong.io\"\n MAX_CALL_DETAILS_ATTEMPTS = 6\n CALL_DETAILS_DELAY = 30 # in seconds\n # Gong API limit is 3 calls/sec — stay safely under it\n MIN_REQUEST_INTERVAL = 0.5 # seconds between requests\n\n def __init__(\n self,\n workspaces: list[str] | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n continue_on_fail: bool = CONTINUE_ON_CONNECTOR_FAILURE,\n hide_user_info: bool = False,\n ) -> None:\n self.workspaces = workspaces\n self.batch_size: int = batch_size\n self.continue_on_fail = continue_on_fail\n self.auth_token_basic: str | None = None\n self.hide_user_info = hide_user_info\n self._last_request_time: float = 0.0\n\n # urllib3 Retry already respects the Retry-After header by default\n # (respect_retry_after_header=True), so on 429 it will sleep for the\n # duration Gong specifies before retrying.\n retry_strategy = Retry(\n total=10,\n backoff_factor=2,\n status_forcelist=[429, 500, 502, 503, 504],\n )\n\n session = requests.Session()\n session.mount(GongConnector.BASE_URL, HTTPAdapter(max_retries=retry_strategy))\n self._session = session\n\n @staticmethod\n def make_url(endpoint: str) -> str:\n url = f\"{GongConnector.BASE_URL}{endpoint}\"\n return url\n\n def _throttled_request(\n self, method: str, url: str, **kwargs: Any\n ) -> requests.Response:\n \"\"\"Rate-limited request wrapper. Enforces MIN_REQUEST_INTERVAL between\n calls to stay under Gong's 3 calls/sec limit and avoid triggering 429s.\"\"\"\n now = time.monotonic()\n elapsed = now - self._last_request_time\n if elapsed < self.MIN_REQUEST_INTERVAL:\n time.sleep(self.MIN_REQUEST_INTERVAL - elapsed)\n\n response = self._session.request(method, url, **kwargs)\n self._last_request_time = time.monotonic()\n return response\n\n def _get_workspace_id_map(self) -> dict[str, str]:\n response = self._throttled_request(\n \"GET\", GongConnector.make_url(\"/v2/workspaces\")\n )\n response.raise_for_status()\n\n workspaces_details = response.json().get(\"workspaces\")\n name_id_map = {\n workspace[\"name\"]: workspace[\"id\"] for workspace in workspaces_details\n }\n id_id_map = {\n workspace[\"id\"]: workspace[\"id\"] for workspace in workspaces_details\n }\n # In very rare case, if a workspace is given a name which is the id of another workspace,\n # Then the user input is treated as the name\n return {**id_id_map, **name_id_map}\n\n def _get_transcript_batches(\n self, start_datetime: str | None = None, end_datetime: str | None = None\n ) -> Generator[list[dict[str, Any]], None, None]:\n body: dict[str, dict] = {\"filter\": {}}\n if start_datetime:\n body[\"filter\"][\"fromDateTime\"] = start_datetime\n if end_datetime:\n body[\"filter\"][\"toDateTime\"] = end_datetime\n\n # The batch_ids in the previous method appears to be batches of call_ids to process\n # In this method, we will retrieve transcripts for them in batches.\n transcripts: list[dict[str, Any]] = []\n workspace_list = self.workspaces or [None] # type: ignore\n workspace_map = self._get_workspace_id_map() if self.workspaces else {}\n\n for workspace in workspace_list:\n if workspace:\n logger.info(f\"Updating Gong workspace: {workspace}\")\n workspace_id = workspace_map.get(workspace)\n if not workspace_id:\n logger.error(f\"Invalid Gong workspace: {workspace}\")\n if not self.continue_on_fail:\n raise ValueError(f\"Invalid workspace: {workspace}\")\n continue\n body[\"filter\"][\"workspaceId\"] = workspace_id\n else:\n if \"workspaceId\" in body[\"filter\"]:\n del body[\"filter\"][\"workspaceId\"]\n\n while True:\n response = self._throttled_request(\n \"POST\", GongConnector.make_url(\"/v2/calls/transcript\"), json=body\n )\n # If no calls in the range, just break out\n if response.status_code == 404:\n break\n\n try:\n response.raise_for_status()\n except Exception:\n logger.error(f\"Error fetching transcripts: {response.text}\")\n raise\n\n data = response.json()\n call_transcripts = data.get(\"callTranscripts\", [])\n transcripts.extend(call_transcripts)\n\n while len(transcripts) >= self.batch_size:\n yield transcripts[: self.batch_size]\n transcripts = transcripts[self.batch_size :]\n\n cursor = data.get(\"records\", {}).get(\"cursor\")\n if cursor:\n body[\"cursor\"] = cursor\n else:\n break\n\n if transcripts:\n yield transcripts\n\n def _get_call_details_by_ids(self, call_ids: list[str]) -> dict:\n body = {\n \"filter\": {\"callIds\": call_ids},\n \"contentSelector\": {\"exposedFields\": {\"parties\": True}},\n }\n\n response = self._throttled_request(\n \"POST\", GongConnector.make_url(\"/v2/calls/extensive\"), json=body\n )\n response.raise_for_status()\n\n calls = response.json().get(\"calls\")\n call_to_metadata = {}\n for call in calls:\n call_to_metadata[call[\"metaData\"][\"id\"]] = call\n\n return call_to_metadata\n\n @staticmethod\n def _parse_parties(parties: list[dict]) -> dict[str, str]:\n id_mapping = {}\n for party in parties:\n name = party.get(\"name\")\n email = party.get(\"emailAddress\")\n\n if name and email:\n full_identifier = f\"{name} ({email})\"\n elif name:\n full_identifier = name\n elif email:\n full_identifier = email\n else:\n full_identifier = \"Unknown\"\n\n id_mapping[party[\"speakerId\"]] = full_identifier\n\n return id_mapping\n\n def _fetch_calls(\n self, start_datetime: str | None = None, end_datetime: str | None = None\n ) -> GenerateDocumentsOutput:\n num_calls = 0\n\n for transcript_batch in self._get_transcript_batches(\n start_datetime, end_datetime\n ):\n doc_batch: list[Document | HierarchyNode] = []\n\n transcript_call_ids = cast(\n list[str],\n [t.get(\"callId\") for t in transcript_batch if t.get(\"callId\")],\n )\n\n call_details_map: dict[str, Any] = {}\n\n # There's a likely race condition in the API where a transcript will have a\n # call id but the call to v2/calls/extensive will not return all of the id's\n # retry with exponential backoff has been observed to mitigate this\n # in ~2 minutes. After max attempts, proceed with whatever we have —\n # the per-call loop below will skip missing IDs gracefully.\n current_attempt = 0\n while True:\n current_attempt += 1\n call_details_map = self._get_call_details_by_ids(transcript_call_ids)\n if set(transcript_call_ids) == set(call_details_map.keys()):\n # we got all the id's we were expecting ... break and continue\n break\n\n # we are missing some id's. Log and retry with exponential backoff\n missing_call_ids = set(transcript_call_ids) - set(\n call_details_map.keys()\n )\n logger.warning(\n f\"_get_call_details_by_ids is missing call id's: \"\n f\"current_attempt={current_attempt} \"\n f\"missing_call_ids={missing_call_ids}\"\n )\n if current_attempt >= self.MAX_CALL_DETAILS_ATTEMPTS:\n logger.error(\n f\"Giving up on missing call id's after \"\n f\"{self.MAX_CALL_DETAILS_ATTEMPTS} attempts: \"\n f\"missing_call_ids={missing_call_ids} — \"\n f\"proceeding with {len(call_details_map)} of \"\n f\"{len(transcript_call_ids)} calls\"\n )\n break\n\n wait_seconds = self.CALL_DETAILS_DELAY * pow(2, current_attempt - 1)\n logger.warning(\n f\"_get_call_details_by_ids waiting to retry: \"\n f\"wait={wait_seconds}s \"\n f\"current_attempt={current_attempt} \"\n f\"next_attempt={current_attempt + 1} \"\n f\"max_attempts={self.MAX_CALL_DETAILS_ATTEMPTS}\"\n )\n time.sleep(wait_seconds)\n\n # now we can iterate per call/transcript\n for transcript in transcript_batch:\n call_id = transcript.get(\"callId\")\n\n if not call_id or call_id not in call_details_map:\n # NOTE(rkuo): seeing odd behavior where call_ids from the transcript\n # don't have call details. adding error debugging logs to trace.\n logger.error(\n f\"Couldn't get call information for Call ID: {call_id}\"\n )\n if call_id:\n logger.error(\n f\"Call debug info: call_id={call_id} \"\n f\"call_ids={transcript_call_ids} \"\n f\"call_details_map={call_details_map.keys()}\"\n )\n if not self.continue_on_fail:\n raise RuntimeError(\n f\"Couldn't get call information for Call ID: {call_id}\"\n )\n continue\n\n call_details = call_details_map[call_id]\n call_metadata = call_details[\"metaData\"]\n\n call_time_str = call_metadata[\"started\"]\n call_title = call_metadata[\"title\"]\n logger.info(\n f\"{num_calls + 1}: Indexing Gong call id {call_id} from {call_time_str.split('T', 1)[0]}: {call_title}\"\n )\n\n call_parties = cast(list[dict] | None, call_details.get(\"parties\"))\n if call_parties is None:\n logger.error(f\"Couldn't get parties for Call ID: {call_id}\")\n call_parties = []\n\n id_to_name_map = self._parse_parties(call_parties)\n\n # Keeping a separate dict here in case the parties info is incomplete\n speaker_to_name: dict[str, str] = {}\n\n transcript_text = \"\"\n call_purpose = call_metadata[\"purpose\"]\n if call_purpose:\n transcript_text += f\"Call Description: {call_purpose}\\n\\n\"\n\n contents = transcript[\"transcript\"]\n for segment in contents:\n speaker_id = segment.get(\"speakerId\", \"\")\n if speaker_id not in speaker_to_name:\n if self.hide_user_info:\n speaker_to_name[speaker_id] = (\n f\"User {len(speaker_to_name) + 1}\"\n )\n else:\n speaker_to_name[speaker_id] = id_to_name_map.get(\n speaker_id, \"Unknown\"\n )\n\n speaker_name = speaker_to_name[speaker_id]\n\n sentences = segment.get(\"sentences\", {})\n monolog = \" \".join(\n [sentence.get(\"text\", \"\") for sentence in sentences]\n )\n transcript_text += f\"{speaker_name}: {monolog}\\n\\n\"\n\n metadata = {}\n if call_metadata.get(\"system\"):\n metadata[\"client\"] = call_metadata.get(\"system\")\n # TODO calls have a clientUniqueId field, can pull that in later\n\n doc_batch.append(\n Document(\n id=call_id,\n sections=[\n TextSection(link=call_metadata[\"url\"], text=transcript_text)\n ],\n source=DocumentSource.GONG,\n # Should not ever be Untitled as a call cannot be made without a Title\n semantic_identifier=call_title or \"Untitled\",\n doc_updated_at=datetime.fromisoformat(call_time_str).astimezone(\n timezone.utc\n ),\n metadata={\"client\": call_metadata.get(\"system\")},\n )\n )\n\n num_calls += 1\n\n yield doc_batch\n\n logger.info(f\"_fetch_calls finished: num_calls={num_calls}\")\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n combined = (\n f\"{credentials['gong_access_key']}:{credentials['gong_access_key_secret']}\"\n )\n self.auth_token_basic = base64.b64encode(combined.encode(\"utf-8\")).decode(\n \"utf-8\"\n )\n\n if self.auth_token_basic is None:\n raise ConnectorMissingCredentialError(\"Gong\")\n\n self._session.headers.update(\n {\"Authorization\": f\"Basic {self.auth_token_basic}\"}\n )\n return None\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._fetch_calls()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n\n # if this env variable is set, don't start from a timestamp before the specified\n # start time\n # TODO: remove this once this is globally available\n if GONG_CONNECTOR_START_TIME:\n special_start_datetime = datetime.fromisoformat(GONG_CONNECTOR_START_TIME)\n special_start_datetime = special_start_datetime.replace(tzinfo=timezone.utc)\n else:\n special_start_datetime = datetime.fromtimestamp(0, tz=timezone.utc)\n\n # don't let the special start dt be past the end time, this causes issues when\n # the Gong API (`filter.fromDateTime: must be before toDateTime`)\n special_start_datetime = min(special_start_datetime, end_datetime)\n\n start_datetime = max(\n datetime.fromtimestamp(start, tz=timezone.utc), special_start_datetime\n )\n\n # Because these are meeting start times, the meeting needs to end and be processed\n # so adding a 1 day buffer and fetching by default till current time\n start_one_day_offset = start_datetime - timedelta(days=1)\n start_time = start_one_day_offset.isoformat()\n\n end_time = datetime.fromtimestamp(end, tz=timezone.utc).isoformat()\n\n logger.info(f\"Fetching Gong calls between {start_time} and {end_time}\")\n return self._fetch_calls(start_time, end_time)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = GongConnector()\n connector.load_credentials(\n {\n \"gong_access_key\": os.environ[\"GONG_ACCESS_KEY\"],\n \"gong_access_key_secret\": os.environ[\"GONG_ACCESS_KEY_SECRET\"],\n }\n )\n\n latest_docs = connector.load_from_state()\n print(next(latest_docs))\n" }, { "path": "backend/onyx/connectors/google_drive/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/google_drive/connector.py", "content": "import copy\nimport json\nimport os\nimport sys\nimport threading\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\nfrom typing import Protocol\nfrom urllib.parse import parse_qs\nfrom urllib.parse import urlparse\nfrom urllib.parse import urlunparse\n\nfrom google.auth.exceptions import RefreshError\nfrom google.oauth2.credentials import Credentials as OAuthCredentials\nfrom google.oauth2.service_account import Credentials as ServiceAccountCredentials\nfrom googleapiclient.errors import HttpError # type: ignore\nfrom typing_extensions import override\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.app_configs import GOOGLE_DRIVE_CONNECTOR_SIZE_THRESHOLD\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import MAX_DRIVE_WORKERS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.google_drive.doc_conversion import build_slim_document\nfrom onyx.connectors.google_drive.doc_conversion import (\n convert_drive_item_to_document,\n)\nfrom onyx.connectors.google_drive.doc_conversion import onyx_document_id_from_drive_file\nfrom onyx.connectors.google_drive.doc_conversion import PermissionSyncContext\nfrom onyx.connectors.google_drive.file_retrieval import crawl_folders_for_files\nfrom onyx.connectors.google_drive.file_retrieval import DriveFileFieldType\nfrom onyx.connectors.google_drive.file_retrieval import get_all_files_for_oauth\nfrom onyx.connectors.google_drive.file_retrieval import (\n get_all_files_in_my_drive_and_shared,\n)\nfrom onyx.connectors.google_drive.file_retrieval import get_external_access_for_folder\nfrom onyx.connectors.google_drive.file_retrieval import get_files_in_shared_drive\nfrom onyx.connectors.google_drive.file_retrieval import get_folder_metadata\nfrom onyx.connectors.google_drive.file_retrieval import get_root_folder_id\nfrom onyx.connectors.google_drive.file_retrieval import get_shared_drive_name\nfrom onyx.connectors.google_drive.file_retrieval import has_link_only_permission\nfrom onyx.connectors.google_drive.models import DriveRetrievalStage\nfrom onyx.connectors.google_drive.models import GoogleDriveCheckpoint\nfrom onyx.connectors.google_drive.models import GoogleDriveFileType\nfrom onyx.connectors.google_drive.models import RetrievedDriveFile\nfrom onyx.connectors.google_drive.models import StageCompletion\nfrom onyx.connectors.google_utils.google_auth import get_google_creds\nfrom onyx.connectors.google_utils.google_utils import execute_paginated_retrieval\nfrom onyx.connectors.google_utils.google_utils import get_file_owners\nfrom onyx.connectors.google_utils.google_utils import GoogleFields\nfrom onyx.connectors.google_utils.resources import get_admin_service\nfrom onyx.connectors.google_utils.resources import get_drive_service\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR\nfrom onyx.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS\nfrom onyx.connectors.google_utils.shared_constants import SLIM_BATCH_SIZE\nfrom onyx.connectors.google_utils.shared_constants import USER_FIELDS\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import NormalizationResult\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import EntityFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\nfrom onyx.utils.threadpool_concurrency import parallel_yield\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom onyx.utils.threadpool_concurrency import ThreadSafeDict\nfrom onyx.utils.threadpool_concurrency import ThreadSafeSet\n\nlogger = setup_logger()\n# TODO: Improve this by using the batch utility: https://googleapis.github.io/google-api-python-client/docs/batch.html\n# All file retrievals could be batched and made at once\n\nBATCHES_PER_CHECKPOINT = 1\n\nDRIVE_BATCH_SIZE = 80\n\nSHARED_DRIVE_PAGES_PER_CHECKPOINT = 2\nMY_DRIVE_PAGES_PER_CHECKPOINT = 2\nOAUTH_PAGES_PER_CHECKPOINT = 2\nFOLDERS_PER_CHECKPOINT = 1\n\n\ndef _extract_str_list_from_comma_str(string: str | None) -> list[str]:\n if not string:\n return []\n return [s.strip() for s in string.split(\",\") if s.strip()]\n\n\ndef _extract_ids_from_urls(urls: list[str]) -> list[str]:\n return [urlparse(url).path.strip(\"/\").split(\"/\")[-1] for url in urls]\n\n\ndef _clean_requested_drive_ids(\n requested_drive_ids: set[str],\n requested_folder_ids: set[str],\n all_drive_ids_available: set[str],\n) -> tuple[list[str], list[str]]:\n invalid_requested_drive_ids = requested_drive_ids - all_drive_ids_available\n filtered_folder_ids = requested_folder_ids - all_drive_ids_available\n if invalid_requested_drive_ids:\n logger.warning(\n f\"Some shared drive IDs were not found. IDs: {invalid_requested_drive_ids}\"\n )\n logger.warning(\"Checking for folder access instead...\")\n filtered_folder_ids.update(invalid_requested_drive_ids)\n\n valid_requested_drive_ids = requested_drive_ids - invalid_requested_drive_ids\n return sorted(valid_requested_drive_ids), sorted(filtered_folder_ids)\n\n\ndef _get_parent_id_from_file(drive_file: GoogleDriveFileType) -> str | None:\n \"\"\"Extract the first parent ID from a drive file.\"\"\"\n parents = drive_file.get(\"parents\")\n if parents and len(parents) > 0:\n return parents[0] # files have a unique parent\n return None\n\n\ndef _is_shared_drive_root(folder: GoogleDriveFileType) -> bool:\n \"\"\"\n Check if a folder is a verified shared drive root.\n\n For shared drives, we can verify using driveId:\n - If driveId is set and folder_id == driveId AND no parents, it's the shared drive root\n - If driveId is set but folder_id != driveId with empty parents, it's a permission issue\n\n Returns True only for verified shared drive roots.\n \"\"\"\n folder_id = folder.get(\"id\")\n drive_id = folder.get(\"driveId\")\n parents = folder.get(\"parents\", [])\n\n # Must have no parents to be a root\n if parents:\n return False\n\n # For shared drive content, the root has id == driveId\n return bool(drive_id and folder_id == drive_id)\n\n\ndef _public_access() -> ExternalAccess:\n return ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n )\n\n\nclass CredentialedRetrievalMethod(Protocol):\n def __call__(\n self,\n field_type: DriveFileFieldType,\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]: ...\n\n\ndef add_retrieval_info(\n drive_files: Iterator[GoogleDriveFileType | str],\n user_email: str,\n completion_stage: DriveRetrievalStage,\n parent_id: str | None = None,\n) -> Iterator[RetrievedDriveFile | str]:\n for file in drive_files:\n if isinstance(file, str):\n yield file\n continue\n yield RetrievedDriveFile(\n drive_file=file,\n user_email=user_email,\n parent_id=parent_id,\n completion_stage=completion_stage,\n )\n\n\nclass DriveIdStatus(Enum):\n AVAILABLE = \"available\"\n IN_PROGRESS = \"in_progress\"\n FINISHED = \"finished\"\n\n\nclass GoogleDriveConnector(\n SlimConnectorWithPermSync, CheckpointedConnectorWithPermSync[GoogleDriveCheckpoint]\n):\n def __init__(\n self,\n include_shared_drives: bool = False,\n include_my_drives: bool = False,\n include_files_shared_with_me: bool = False,\n shared_drive_urls: str | None = None,\n my_drive_emails: str | None = None,\n shared_folder_urls: str | None = None,\n specific_user_emails: str | None = None,\n exclude_domain_link_only: bool = False,\n batch_size: int = INDEX_BATCH_SIZE, # noqa: ARG002\n # OLD PARAMETERS\n folder_paths: list[str] | None = None,\n include_shared: bool | None = None,\n follow_shortcuts: bool | None = None,\n only_org_public: bool | None = None,\n continue_on_failure: bool | None = None,\n ) -> None:\n # Check for old input parameters\n if folder_paths is not None:\n logger.warning(\n \"The 'folder_paths' parameter is deprecated. Use 'shared_folder_urls' instead.\"\n )\n if include_shared is not None:\n logger.warning(\n \"The 'include_shared' parameter is deprecated. Use 'include_files_shared_with_me' instead.\"\n )\n if follow_shortcuts is not None:\n logger.warning(\"The 'follow_shortcuts' parameter is deprecated.\")\n if only_org_public is not None:\n logger.warning(\"The 'only_org_public' parameter is deprecated.\")\n if continue_on_failure is not None:\n logger.warning(\"The 'continue_on_failure' parameter is deprecated.\")\n\n if not any(\n (\n include_shared_drives,\n include_my_drives,\n include_files_shared_with_me,\n shared_folder_urls,\n my_drive_emails,\n shared_drive_urls,\n )\n ):\n raise ConnectorValidationError(\n \"Nothing to index. Please specify at least one of the following: \"\n \"include_shared_drives, include_my_drives, include_files_shared_with_me, \"\n \"shared_folder_urls, or my_drive_emails\"\n )\n\n specific_requests_made = False\n if bool(shared_drive_urls) or bool(my_drive_emails) or bool(shared_folder_urls):\n specific_requests_made = True\n self.specific_requests_made = specific_requests_made\n\n # NOTE: potentially modified in load_credentials if using service account\n self.include_files_shared_with_me = (\n False if specific_requests_made else include_files_shared_with_me\n )\n self.include_my_drives = False if specific_requests_made else include_my_drives\n self.include_shared_drives = (\n False if specific_requests_made else include_shared_drives\n )\n\n shared_drive_url_list = _extract_str_list_from_comma_str(shared_drive_urls)\n self._requested_shared_drive_ids = set(\n _extract_ids_from_urls(shared_drive_url_list)\n )\n\n self._requested_my_drive_emails = set(\n _extract_str_list_from_comma_str(my_drive_emails)\n )\n\n shared_folder_url_list = _extract_str_list_from_comma_str(shared_folder_urls)\n self._requested_folder_ids = set(_extract_ids_from_urls(shared_folder_url_list))\n self._specific_user_emails = _extract_str_list_from_comma_str(\n specific_user_emails\n )\n self.exclude_domain_link_only = exclude_domain_link_only\n\n self._primary_admin_email: str | None = None\n\n self._creds: OAuthCredentials | ServiceAccountCredentials | None = None\n self._creds_dict: dict[str, Any] | None = None\n\n # ids of folders and shared drives that have been traversed\n self._retrieved_folder_and_drive_ids: set[str] = set()\n\n # Cache of known My Drive root IDs (user_email -> root_id)\n # Used to verify if a folder with no parents is actually a My Drive root\n # Thread-safe because multiple impersonation threads access this concurrently\n self._my_drive_root_id_cache: ThreadSafeDict[str, str] = ThreadSafeDict()\n\n self.allow_images = False\n\n self.size_threshold = GOOGLE_DRIVE_CONNECTOR_SIZE_THRESHOLD\n\n def set_allow_images(self, value: bool) -> None:\n self.allow_images = value\n\n @property\n def primary_admin_email(self) -> str:\n if self._primary_admin_email is None:\n raise RuntimeError(\n \"Primary admin email missing, should not call this property before calling load_credentials\"\n )\n return self._primary_admin_email\n\n @property\n def google_domain(self) -> str:\n if self._primary_admin_email is None:\n raise RuntimeError(\n \"Primary admin email missing, should not call this property before calling load_credentials\"\n )\n return self._primary_admin_email.split(\"@\")[-1]\n\n @property\n def creds(self) -> OAuthCredentials | ServiceAccountCredentials:\n if self._creds is None:\n raise RuntimeError(\n \"Creds missing, should not call this property before calling load_credentials\"\n )\n return self._creds\n\n @classmethod\n @override\n def normalize_url(cls, url: str) -> NormalizationResult:\n \"\"\"Normalize a Google Drive URL to match the canonical Document.id format.\n\n Reuses the connector's existing document ID creation logic from\n onyx_document_id_from_drive_file.\n \"\"\"\n parsed = urlparse(url)\n netloc = parsed.netloc.lower()\n\n if not (\n netloc.startswith(\"docs.google.com\")\n or netloc.startswith(\"drive.google.com\")\n ):\n return NormalizationResult(normalized_url=None, use_default=False)\n\n # Handle ?id= query parameter case\n query_params = parse_qs(parsed.query)\n doc_id = query_params.get(\"id\", [None])[0]\n if doc_id:\n scheme = parsed.scheme or \"https\"\n netloc = \"drive.google.com\"\n path = f\"/file/d/{doc_id}\"\n params = \"\"\n query = \"\"\n fragment = \"\"\n normalized = urlunparse(\n (scheme, netloc, path, params, query, fragment)\n ).rstrip(\"/\")\n return NormalizationResult(normalized_url=normalized, use_default=False)\n\n # Extract file ID and use connector's function\n path_parts = parsed.path.split(\"/\")\n file_id = None\n for i, part in enumerate(path_parts):\n if part == \"d\" and i + 1 < len(path_parts):\n file_id = path_parts[i + 1]\n break\n\n if not file_id:\n return NormalizationResult(normalized_url=None, use_default=False)\n\n # Create minimal file object for connector function\n file_obj = {\"webViewLink\": url, \"id\": file_id}\n normalized = onyx_document_id_from_drive_file(file_obj).rstrip(\"/\")\n return NormalizationResult(normalized_url=normalized, use_default=False)\n\n # TODO: ensure returned new_creds_dict is actually persisted when this is called?\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, str] | None:\n try:\n self._primary_admin_email = credentials[DB_CREDENTIALS_PRIMARY_ADMIN_KEY]\n except KeyError:\n raise ValueError(\"Credentials json missing primary admin key\")\n\n self._creds, new_creds_dict = get_google_creds(\n credentials=credentials,\n source=DocumentSource.GOOGLE_DRIVE,\n )\n\n # Service account connectors don't have a specific setting determining whether\n # to include \"shared with me\" for each user, so we default to true unless the connector\n # is in specific folders/drives mode. Note that shared files are only picked up during\n # the My Drive stage, so this does nothing if the connector is set to only index shared drives.\n if (\n isinstance(self._creds, ServiceAccountCredentials)\n and not self.specific_requests_made\n ):\n self.include_files_shared_with_me = True\n\n self._creds_dict = new_creds_dict\n\n return new_creds_dict\n\n def _update_traversed_parent_ids(self, folder_id: str) -> None:\n self._retrieved_folder_and_drive_ids.add(folder_id)\n\n def _get_all_user_emails(self) -> list[str]:\n if self._specific_user_emails:\n return self._specific_user_emails\n\n # Start with primary admin email\n user_emails = [self.primary_admin_email]\n\n # Only fetch additional users if using service account\n if isinstance(self.creds, OAuthCredentials):\n return user_emails\n\n admin_service = get_admin_service(\n creds=self.creds,\n user_email=self.primary_admin_email,\n )\n\n # Get admins first since they're more likely to have access to most files\n for is_admin in [True, False]:\n query = \"isAdmin=true\" if is_admin else \"isAdmin=false\"\n for user in execute_paginated_retrieval(\n retrieval_function=admin_service.users().list,\n list_key=\"users\",\n fields=USER_FIELDS,\n domain=self.google_domain,\n query=query,\n ):\n if email := user.get(\"primaryEmail\"):\n if email not in user_emails:\n user_emails.append(email)\n return user_emails\n\n def _get_my_drive_root_id(self, user_email: str) -> str | None:\n \"\"\"\n Get the My Drive root folder ID for a user.\n\n Uses a cache to avoid repeated API calls. Returns None if the user\n doesn't have access to Drive APIs or the call fails.\n \"\"\"\n if user_email in self._my_drive_root_id_cache:\n return self._my_drive_root_id_cache[user_email]\n\n try:\n drive_service = get_drive_service(self.creds, user_email)\n root_id = get_root_folder_id(drive_service)\n self._my_drive_root_id_cache[user_email] = root_id\n return root_id\n except Exception:\n # User might not have access to Drive APIs\n return None\n\n def _is_my_drive_root(\n self, folder: GoogleDriveFileType, retriever_email: str\n ) -> bool:\n \"\"\"\n Check if a folder is a My Drive root.\n\n For My Drive folders (no driveId), we verify by comparing the folder ID\n to the actual My Drive root ID obtained via files().get(fileId='root').\n \"\"\"\n folder_id = folder.get(\"id\")\n drive_id = folder.get(\"driveId\")\n parents = folder.get(\"parents\", [])\n\n # If there are parents, this is not a root\n if parents:\n return False\n\n # If driveId is set, this is shared drive content, not My Drive\n if drive_id:\n return False\n\n # Get the My Drive root ID for this user and compare\n root_id = self._get_my_drive_root_id(retriever_email)\n if root_id and folder_id == root_id:\n return True\n\n # Also check with admin in case the retriever doesn't have access\n admin_root_id = self._get_my_drive_root_id(self.primary_admin_email)\n if admin_root_id and folder_id == admin_root_id:\n return True\n\n return False\n\n def _get_new_ancestors_for_files(\n self,\n files: list[RetrievedDriveFile],\n seen_hierarchy_node_raw_ids: ThreadSafeSet[str],\n fully_walked_hierarchy_node_raw_ids: ThreadSafeSet[str],\n permission_sync_context: PermissionSyncContext | None = None,\n add_prefix: bool = False,\n ) -> list[HierarchyNode]:\n \"\"\"\n Get all NEW ancestor hierarchy nodes for a batch of files.\n\n For each file, walks up the parent chain until reaching a root/drive\n (terminal node with no parent). Returns HierarchyNode objects for all\n new ancestors.\n\n The function tracks two separate sets:\n - seen_hierarchy_node_raw_ids: Nodes we've already yielded (to avoid duplicates)\n - fully_walked_hierarchy_node_raw_ids: Nodes where we've successfully walked\n to a terminal root. Only skip walking from a node if it's in this set.\n\n This separation ensures that if User A can access folder C but not its parent B,\n a later User B who has access to both can still complete the walk to the root.\n\n Args:\n files: List of retrieved drive files to get ancestors for\n seen_hierarchy_node_raw_ids: Set of already-yielded node IDs (modified in place)\n fully_walked_hierarchy_node_raw_ids: Set of node IDs where the walk to root\n succeeded (modified in place)\n permission_sync_context: If provided, permissions will be fetched for hierarchy nodes.\n Contains google_domain and primary_admin_email needed for permission syncing.\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path).\n\n Returns:\n List of HierarchyNode objects for new ancestors (ordered parent-first)\n \"\"\"\n service = get_drive_service(self.creds, self.primary_admin_email)\n field_type = (\n DriveFileFieldType.WITH_PERMISSIONS\n if permission_sync_context\n else DriveFileFieldType.STANDARD\n )\n new_nodes: list[HierarchyNode] = []\n\n for file in files:\n parent_id = _get_parent_id_from_file(file.drive_file)\n if not parent_id:\n continue\n\n # Only skip if we've already successfully walked from this node to a root.\n # Don't skip just because it's \"seen\" - a previous user may have failed\n # to walk to the root, and this user might have better access.\n if parent_id in fully_walked_hierarchy_node_raw_ids:\n continue\n\n # Walk up the parent chain\n ancestors_to_add: list[HierarchyNode] = []\n node_ids_in_walk: list[str] = []\n current_id: str | None = parent_id\n reached_terminal = False\n\n while current_id:\n node_ids_in_walk.append(current_id)\n\n # If we hit a node that's already been fully walked, we know\n # the path from here to root is complete\n if current_id in fully_walked_hierarchy_node_raw_ids:\n reached_terminal = True\n break\n\n # Fetch folder metadata\n folder = self._get_folder_metadata(\n current_id, file.user_email, field_type\n )\n if not folder:\n # Can't access this folder - stop climbing\n # Don't mark as fully walked since we didn't reach root\n break\n\n folder_parent_id = _get_parent_id_from_file(folder)\n\n # Create the node BEFORE marking as seen to avoid a race condition where:\n # 1. Thread A marks node as \"seen\"\n # 2. Thread A fails to create node (e.g., API error in get_external_access)\n # 3. Thread B sees node as \"already seen\" and skips it\n # 4. Result: node is never yielded\n #\n # By creating first and then atomically checking/marking, we ensure that\n # if creation fails, another thread can still try. If both succeed,\n # only one will add to ancestors_to_add (the one that wins check_and_add).\n if permission_sync_context:\n external_access = get_external_access_for_folder(\n folder,\n permission_sync_context.google_domain,\n service,\n add_prefix,\n )\n else:\n external_access = _public_access()\n\n node = HierarchyNode(\n raw_node_id=current_id,\n raw_parent_id=folder_parent_id,\n display_name=folder.get(\"name\", \"Unknown Folder\"),\n link=folder.get(\"webViewLink\"),\n node_type=HierarchyNodeType.FOLDER,\n external_access=external_access,\n )\n\n # Now atomically check and add - only append if we're the first thread\n # to successfully create this node\n already_seen = seen_hierarchy_node_raw_ids.check_and_add(current_id)\n if not already_seen:\n ancestors_to_add.append(node)\n\n # Check if this is a verified terminal node (actual root, not just\n # empty parents due to permission limitations)\n # Check shared drive root first (simple ID comparison)\n if _is_shared_drive_root(folder):\n # files().get() returns 'Drive' for shared drive roots;\n # fetch the real name via drives().get().\n # Try both the retriever and admin since the admin may\n # not have access to private shared drives.\n drive_name = self._get_shared_drive_name(\n current_id, file.user_email\n )\n if drive_name:\n node.display_name = drive_name\n node.node_type = HierarchyNodeType.SHARED_DRIVE\n reached_terminal = True\n break\n\n # Check if this is a My Drive root (requires API call, but cached)\n if self._is_my_drive_root(folder, file.user_email):\n reached_terminal = True\n break\n\n # If parents is empty but we couldn't verify it's a true root,\n # stop walking but don't mark as fully walked (another user\n # with better access might be able to continue)\n if folder_parent_id is None:\n break\n\n # Move to parent\n current_id = folder_parent_id\n\n # If we successfully reached a terminal node (or a fully-walked node),\n # mark all nodes in this walk as fully walked\n if reached_terminal:\n fully_walked_hierarchy_node_raw_ids.update(set(node_ids_in_walk))\n\n new_nodes += ancestors_to_add\n\n return new_nodes\n\n def _get_folder_metadata(\n self, folder_id: str, retriever_email: str, field_type: DriveFileFieldType\n ) -> GoogleDriveFileType | None:\n \"\"\"\n Fetch metadata for a folder by ID.\n\n Important: When a user has access to a shared folder but NOT its parent,\n the Google Drive API returns the folder metadata WITHOUT the parent info.\n To handle this, if the retriever gets a folder without parents, we also\n try with admin who may have better access and can see the parent chain.\n \"\"\"\n best_folder: GoogleDriveFileType | None = None\n\n # Use a set to deduplicate if retriever_email == primary_admin_email\n for email in {retriever_email, self.primary_admin_email}:\n service = get_drive_service(self.creds, email)\n folder = get_folder_metadata(service, folder_id, field_type)\n\n if not folder:\n logger.debug(f\"Failed to fetch folder {folder_id} using {email}\")\n continue\n\n logger.debug(f\"Successfully fetched folder {folder_id} using {email}\")\n\n # If this folder has parents, use it\n if folder.get(\"parents\"):\n return folder\n\n # Folder has no parents - could be a root OR user lacks access to parent\n # Keep this as a fallback but try admin to see if they can see parents\n if best_folder is None:\n best_folder = folder\n logger.debug(\n f\"Folder {folder_id} has no parents when fetched by {email}, will try admin to check for parent access\"\n )\n\n if best_folder:\n logger.debug(\n f\"Successfully fetched folder {folder_id} but no parents found\"\n )\n return best_folder\n\n logger.debug(\n f\"All attempts failed to fetch folder {folder_id} (tried {retriever_email} and {self.primary_admin_email})\"\n )\n return None\n\n def _get_shared_drive_name(self, drive_id: str, retriever_email: str) -> str | None:\n \"\"\"Fetch the name of a shared drive, trying both the retriever and admin.\"\"\"\n for email in {retriever_email, self.primary_admin_email}:\n svc = get_drive_service(self.creds, email)\n name = get_shared_drive_name(svc, drive_id)\n if name:\n return name\n return None\n\n def get_all_drive_ids(self) -> set[str]:\n return self._get_all_drives_for_user(self.primary_admin_email)\n\n def _get_all_drives_for_user(self, user_email: str) -> set[str]:\n drive_service = get_drive_service(self.creds, user_email)\n is_service_account = isinstance(self.creds, ServiceAccountCredentials)\n logger.info(\n f\"Getting all drives for user {user_email} with service account: {is_service_account}\"\n )\n all_drive_ids: set[str] = set()\n for drive in execute_paginated_retrieval(\n retrieval_function=drive_service.drives().list,\n list_key=\"drives\",\n useDomainAdminAccess=is_service_account,\n fields=\"drives(id),nextPageToken\",\n ):\n all_drive_ids.add(drive[\"id\"])\n\n if not all_drive_ids:\n logger.warning(\n \"No drives found even though indexing shared drives was requested.\"\n )\n\n return all_drive_ids\n\n def make_drive_id_getter(\n self, drive_ids: list[str], checkpoint: GoogleDriveCheckpoint\n ) -> Callable[[str], str | None]:\n status_lock = threading.Lock()\n\n in_progress_drive_ids = {\n completion.current_folder_or_drive_id: user_email\n for user_email, completion in checkpoint.completion_map.items()\n if completion.stage == DriveRetrievalStage.SHARED_DRIVE_FILES\n and completion.current_folder_or_drive_id is not None\n }\n drive_id_status: dict[str, DriveIdStatus] = {}\n for drive_id in drive_ids:\n if drive_id in self._retrieved_folder_and_drive_ids:\n drive_id_status[drive_id] = DriveIdStatus.FINISHED\n elif drive_id in in_progress_drive_ids:\n drive_id_status[drive_id] = DriveIdStatus.IN_PROGRESS\n else:\n drive_id_status[drive_id] = DriveIdStatus.AVAILABLE\n\n def get_available_drive_id(thread_id: str) -> str | None:\n completion = checkpoint.completion_map[thread_id]\n with status_lock:\n future_work = None\n for drive_id, status in drive_id_status.items():\n if drive_id in self._retrieved_folder_and_drive_ids:\n drive_id_status[drive_id] = DriveIdStatus.FINISHED\n continue\n if drive_id in completion.processed_drive_ids:\n continue\n\n if status == DriveIdStatus.AVAILABLE:\n # add to processed drive ids so if this user fails to retrieve once\n # they won't try again on the next checkpoint run\n completion.processed_drive_ids.add(drive_id)\n return drive_id\n elif status == DriveIdStatus.IN_PROGRESS:\n logger.debug(f\"Drive id in progress: {drive_id}\")\n future_work = drive_id\n\n if future_work:\n # in this case, all drive ids are either finished or in progress.\n # This thread will pick up one of the in progress ones in case it fails.\n # This is a much simpler approach than waiting for a failure picking it up,\n # at the cost of some repeated work until all shared drives are retrieved.\n # we avoid apocalyptic cases like all threads focusing on one huge drive\n # because the drive id is added to _retrieved_folder_and_drive_ids after any thread\n # manages to retrieve any file from it (unfortunately, this is also the reason we currently\n # sometimes fail to retrieve restricted access folders/files)\n completion.processed_drive_ids.add(future_work)\n return future_work\n return None # no work available, return None\n\n return get_available_drive_id\n\n def _impersonate_user_for_retrieval(\n self,\n user_email: str,\n field_type: DriveFileFieldType,\n checkpoint: GoogleDriveCheckpoint,\n get_new_drive_id: Callable[[str], str | None],\n sorted_filtered_folder_ids: list[str],\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]:\n logger.info(f\"Impersonating user {user_email}\")\n curr_stage = checkpoint.completion_map[user_email]\n resuming = True\n if curr_stage.stage == DriveRetrievalStage.START:\n logger.info(f\"Setting stage to {DriveRetrievalStage.MY_DRIVE_FILES.value}\")\n curr_stage.stage = DriveRetrievalStage.MY_DRIVE_FILES\n resuming = False\n drive_service = get_drive_service(self.creds, user_email)\n\n # validate that the user has access to the drive APIs by performing a simple\n # request and checking for a 401\n try:\n logger.debug(f\"Getting root folder id for user {user_email}\")\n # default is ~17mins of retries, don't do that here for cases so we don't\n # waste 17mins everytime we run into a user without access to drive APIs\n retry_builder(tries=3, delay=1)(get_root_folder_id)(drive_service)\n except HttpError as e:\n if e.status_code == 401:\n # fail gracefully, let the other impersonations continue\n # one user without access shouldn't block the entire connector\n logger.warning(\n f\"User '{user_email}' does not have access to the drive APIs.\"\n )\n # mark this user as done so we don't try to retrieve anything for them\n # again\n curr_stage.stage = DriveRetrievalStage.DONE\n return\n raise\n except RefreshError as e:\n logger.warning(\n f\"User '{user_email}' could not refresh their token. Error: {e}\"\n )\n # mark this user as done so we don't try to retrieve anything for them\n # again\n yield RetrievedDriveFile(\n completion_stage=DriveRetrievalStage.DONE,\n drive_file={},\n user_email=user_email,\n error=e,\n )\n curr_stage.stage = DriveRetrievalStage.DONE\n return\n # if we are including my drives, try to get the current user's my\n # drive if any of the following are true:\n # - include_my_drives is true\n # - the current user's email is in the requested emails\n if curr_stage.stage == DriveRetrievalStage.MY_DRIVE_FILES:\n if self.include_my_drives or user_email in self._requested_my_drive_emails:\n logger.info(\n f\"Getting all files in my drive as '{user_email}. Resuming: {resuming}. \"\n f\"Stage completed until: {curr_stage.completed_until}. \"\n f\"Next page token: {curr_stage.next_page_token}\"\n )\n\n for file_or_token in add_retrieval_info(\n get_all_files_in_my_drive_and_shared(\n service=drive_service,\n update_traversed_ids_func=self._update_traversed_parent_ids,\n field_type=field_type,\n include_shared_with_me=self.include_files_shared_with_me,\n max_num_pages=MY_DRIVE_PAGES_PER_CHECKPOINT,\n start=curr_stage.completed_until if resuming else start,\n end=end,\n cache_folders=not bool(curr_stage.completed_until),\n page_token=curr_stage.next_page_token,\n ),\n user_email,\n DriveRetrievalStage.MY_DRIVE_FILES,\n ):\n if isinstance(file_or_token, str):\n logger.debug(f\"Done with max num pages for user {user_email}\")\n checkpoint.completion_map[user_email].next_page_token = (\n file_or_token\n )\n return # done with the max num pages, return checkpoint\n yield file_or_token\n\n checkpoint.completion_map[user_email].next_page_token = None\n curr_stage.stage = DriveRetrievalStage.SHARED_DRIVE_FILES\n curr_stage.current_folder_or_drive_id = None\n return # resume from next stage on the next run\n\n if curr_stage.stage == DriveRetrievalStage.SHARED_DRIVE_FILES:\n\n def _yield_from_drive(\n drive_id: str, drive_start: SecondsSinceUnixEpoch | None\n ) -> Iterator[RetrievedDriveFile | str]:\n yield from add_retrieval_info(\n get_files_in_shared_drive(\n service=drive_service,\n drive_id=drive_id,\n field_type=field_type,\n max_num_pages=SHARED_DRIVE_PAGES_PER_CHECKPOINT,\n update_traversed_ids_func=self._update_traversed_parent_ids,\n cache_folders=not bool(\n drive_start\n ), # only cache folders for 0 or None\n start=drive_start,\n end=end,\n page_token=curr_stage.next_page_token,\n ),\n user_email,\n DriveRetrievalStage.SHARED_DRIVE_FILES,\n parent_id=drive_id,\n )\n\n # resume from a checkpoint\n if resuming and (drive_id := curr_stage.current_folder_or_drive_id):\n resume_start = curr_stage.completed_until\n for file_or_token in _yield_from_drive(drive_id, resume_start):\n if isinstance(file_or_token, str):\n checkpoint.completion_map[user_email].next_page_token = (\n file_or_token\n )\n return # done with the max num pages, return checkpoint\n yield file_or_token\n\n drive_id = get_new_drive_id(user_email)\n if drive_id:\n logger.info(\n f\"Getting files in shared drive '{drive_id}' as '{user_email}. Resuming: {resuming}\"\n )\n curr_stage.completed_until = 0\n curr_stage.current_folder_or_drive_id = drive_id\n for file_or_token in _yield_from_drive(drive_id, start):\n if isinstance(file_or_token, str):\n checkpoint.completion_map[user_email].next_page_token = (\n file_or_token\n )\n return # done with the max num pages, return checkpoint\n yield file_or_token\n curr_stage.current_folder_or_drive_id = None\n return # get a new drive id on the next run\n\n checkpoint.completion_map[user_email].next_page_token = None\n curr_stage.stage = DriveRetrievalStage.FOLDER_FILES\n curr_stage.current_folder_or_drive_id = None\n return # resume from next stage on the next run\n\n # In the folder files section of service account retrieval we take extra care\n # to not retrieve duplicate docs. In particular, we only add a folder to\n # retrieved_folder_and_drive_ids when all users are finished retrieving files\n # from that folder, and maintain a set of all file ids that have been retrieved\n # for each folder. This might get rather large; in practice we assume that the\n # specific folders users choose to index don't have too many files.\n if curr_stage.stage == DriveRetrievalStage.FOLDER_FILES:\n\n def _yield_from_folder_crawl(\n folder_id: str, folder_start: SecondsSinceUnixEpoch | None\n ) -> Iterator[RetrievedDriveFile]:\n for retrieved_file in crawl_folders_for_files(\n service=drive_service,\n parent_id=folder_id,\n field_type=field_type,\n user_email=user_email,\n traversed_parent_ids=self._retrieved_folder_and_drive_ids,\n update_traversed_ids_func=self._update_traversed_parent_ids,\n start=folder_start,\n end=end,\n ):\n yield retrieved_file\n\n # resume from a checkpoint\n last_processed_folder = None\n if resuming:\n folder_id = curr_stage.current_folder_or_drive_id\n if folder_id is None:\n logger.warning(\n f\"folder id not set in checkpoint for user {user_email}. \"\n \"This happens occasionally when the connector is interrupted \"\n \"and resumed.\"\n )\n else:\n resume_start = curr_stage.completed_until\n yield from _yield_from_folder_crawl(folder_id, resume_start)\n last_processed_folder = folder_id\n\n skipping_seen_folders = last_processed_folder is not None\n # NOTE: this assumes a small number of folders to crawl. If someone\n # really wants to specify a large number of folders, we should use\n # binary search to find the first unseen folder.\n num_completed_folders = 0\n for folder_id in sorted_filtered_folder_ids:\n if skipping_seen_folders:\n skipping_seen_folders = folder_id != last_processed_folder\n continue\n\n if folder_id in self._retrieved_folder_and_drive_ids:\n continue\n\n curr_stage.completed_until = 0\n curr_stage.current_folder_or_drive_id = folder_id\n\n if num_completed_folders >= FOLDERS_PER_CHECKPOINT:\n return # resume from this folder on the next run\n\n logger.info(f\"Getting files in folder '{folder_id}' as '{user_email}'\")\n yield from _yield_from_folder_crawl(folder_id, start)\n num_completed_folders += 1\n\n curr_stage.stage = DriveRetrievalStage.DONE\n\n def _manage_service_account_retrieval(\n self,\n field_type: DriveFileFieldType,\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]:\n \"\"\"\n The current implementation of the service account retrieval does some\n initial setup work using the primary admin email, then runs MAX_DRIVE_WORKERS\n concurrent threads, each of which impersonates a different user and retrieves\n files for that user. Technically, the actual work each thread does is \"yield the\n next file retrieved by the user\", at which point it returns to the thread pool;\n see parallel_yield for more details.\n \"\"\"\n if checkpoint.completion_stage == DriveRetrievalStage.START:\n checkpoint.completion_stage = DriveRetrievalStage.USER_EMAILS\n\n if checkpoint.completion_stage == DriveRetrievalStage.USER_EMAILS:\n all_org_emails: list[str] = self._get_all_user_emails()\n checkpoint.user_emails = all_org_emails\n checkpoint.completion_stage = DriveRetrievalStage.DRIVE_IDS\n else:\n if checkpoint.user_emails is None:\n raise ValueError(\"user emails not set\")\n all_org_emails = checkpoint.user_emails\n\n sorted_drive_ids, sorted_folder_ids = self._determine_retrieval_ids(\n checkpoint, DriveRetrievalStage.MY_DRIVE_FILES\n )\n\n # Setup initial completion map on first connector run\n for email in all_org_emails:\n # don't overwrite existing completion map on resuming runs\n if email in checkpoint.completion_map:\n continue\n checkpoint.completion_map[email] = StageCompletion(\n stage=DriveRetrievalStage.START,\n completed_until=0,\n processed_drive_ids=set(),\n )\n\n # we've found all users and drives, now time to actually start\n # fetching stuff\n logger.info(f\"Found {len(all_org_emails)} users to impersonate\")\n logger.debug(f\"Users: {all_org_emails}\")\n logger.info(f\"Found {len(sorted_drive_ids)} drives to retrieve\")\n logger.debug(f\"Drives: {sorted_drive_ids}\")\n logger.info(f\"Found {len(sorted_folder_ids)} folders to retrieve\")\n logger.debug(f\"Folders: {sorted_folder_ids}\")\n\n drive_id_getter = self.make_drive_id_getter(sorted_drive_ids, checkpoint)\n\n # only process emails that we haven't already completed retrieval for\n non_completed_org_emails = [\n user_email\n for user_email, stage_completion in checkpoint.completion_map.items()\n if stage_completion.stage != DriveRetrievalStage.DONE\n ]\n\n logger.debug(f\"Non-completed users remaining: {len(non_completed_org_emails)}\")\n\n # don't process too many emails before returning a checkpoint. This is\n # to resolve the case where there are a ton of emails that don't have access\n # to the drive APIs. Without this, we could loop through these emails for\n # more than 3 hours, causing a timeout and stalling progress.\n email_batch_takes_us_to_completion = True\n MAX_EMAILS_TO_PROCESS_BEFORE_CHECKPOINTING = MAX_DRIVE_WORKERS\n if len(non_completed_org_emails) > MAX_EMAILS_TO_PROCESS_BEFORE_CHECKPOINTING:\n non_completed_org_emails = non_completed_org_emails[\n :MAX_EMAILS_TO_PROCESS_BEFORE_CHECKPOINTING\n ]\n email_batch_takes_us_to_completion = False\n\n user_retrieval_gens = [\n self._impersonate_user_for_retrieval(\n email,\n field_type,\n checkpoint,\n drive_id_getter,\n sorted_folder_ids,\n start,\n end,\n )\n for email in non_completed_org_emails\n ]\n yield from parallel_yield(user_retrieval_gens, max_workers=MAX_DRIVE_WORKERS)\n\n # if there are more emails to process, don't mark as complete\n if not email_batch_takes_us_to_completion:\n return\n\n remaining_folders = (\n set(sorted_drive_ids) | set(sorted_folder_ids)\n ) - self._retrieved_folder_and_drive_ids\n if remaining_folders:\n logger.warning(\n f\"Some folders/drives were not retrieved. IDs: {remaining_folders}\"\n )\n if any(\n checkpoint.completion_map[user_email].stage != DriveRetrievalStage.DONE\n for user_email in all_org_emails\n ):\n logger.info(\n \"some users did not complete retrieval, returning checkpoint for another run\"\n )\n return\n checkpoint.completion_stage = DriveRetrievalStage.DONE\n\n def _determine_retrieval_ids(\n self,\n checkpoint: GoogleDriveCheckpoint,\n next_stage: DriveRetrievalStage,\n ) -> tuple[list[str], list[str]]:\n all_drive_ids = self.get_all_drive_ids()\n sorted_drive_ids: list[str] = []\n sorted_folder_ids: list[str] = []\n if checkpoint.completion_stage == DriveRetrievalStage.DRIVE_IDS:\n if self._requested_shared_drive_ids or self._requested_folder_ids:\n (\n sorted_drive_ids,\n sorted_folder_ids,\n ) = _clean_requested_drive_ids(\n requested_drive_ids=self._requested_shared_drive_ids,\n requested_folder_ids=self._requested_folder_ids,\n all_drive_ids_available=all_drive_ids,\n )\n elif self.include_shared_drives:\n sorted_drive_ids = sorted(all_drive_ids)\n\n checkpoint.drive_ids_to_retrieve = sorted_drive_ids\n checkpoint.folder_ids_to_retrieve = sorted_folder_ids\n checkpoint.completion_stage = next_stage\n else:\n if checkpoint.drive_ids_to_retrieve is None:\n raise ValueError(\"drive ids to retrieve not set in checkpoint\")\n if checkpoint.folder_ids_to_retrieve is None:\n raise ValueError(\"folder ids to retrieve not set in checkpoint\")\n # When loading from a checkpoint, load the previously cached drive and folder ids\n sorted_drive_ids = checkpoint.drive_ids_to_retrieve\n sorted_folder_ids = checkpoint.folder_ids_to_retrieve\n\n return sorted_drive_ids, sorted_folder_ids\n\n def _oauth_retrieval_all_files(\n self,\n field_type: DriveFileFieldType,\n drive_service: GoogleDriveService,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n page_token: str | None = None,\n ) -> Iterator[RetrievedDriveFile | str]:\n if not self.include_files_shared_with_me and not self.include_my_drives:\n return\n\n logger.info(\n f\"Getting shared files/my drive files for OAuth \"\n f\"with include_files_shared_with_me={self.include_files_shared_with_me}, \"\n f\"include_my_drives={self.include_my_drives}, \"\n f\"include_shared_drives={self.include_shared_drives}.\"\n f\"Using '{self.primary_admin_email}' as the account.\"\n )\n yield from add_retrieval_info(\n get_all_files_for_oauth(\n service=drive_service,\n include_files_shared_with_me=self.include_files_shared_with_me,\n include_my_drives=self.include_my_drives,\n include_shared_drives=self.include_shared_drives,\n field_type=field_type,\n max_num_pages=OAUTH_PAGES_PER_CHECKPOINT,\n start=start,\n end=end,\n page_token=page_token,\n ),\n self.primary_admin_email,\n DriveRetrievalStage.OAUTH_FILES,\n )\n\n def _oauth_retrieval_drives(\n self,\n field_type: DriveFileFieldType,\n drive_service: GoogleDriveService,\n drive_ids_to_retrieve: list[str],\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile | str]:\n def _yield_from_drive(\n drive_id: str, drive_start: SecondsSinceUnixEpoch | None\n ) -> Iterator[RetrievedDriveFile | str]:\n yield from add_retrieval_info(\n get_files_in_shared_drive(\n service=drive_service,\n drive_id=drive_id,\n field_type=field_type,\n max_num_pages=SHARED_DRIVE_PAGES_PER_CHECKPOINT,\n cache_folders=not bool(\n drive_start\n ), # only cache folders for 0 or None\n update_traversed_ids_func=self._update_traversed_parent_ids,\n start=drive_start,\n end=end,\n page_token=checkpoint.completion_map[\n self.primary_admin_email\n ].next_page_token,\n ),\n self.primary_admin_email,\n DriveRetrievalStage.SHARED_DRIVE_FILES,\n parent_id=drive_id,\n )\n\n # If we are resuming from a checkpoint, we need to finish retrieving the files from the last drive we retrieved\n if (\n checkpoint.completion_map[self.primary_admin_email].stage\n == DriveRetrievalStage.SHARED_DRIVE_FILES\n ):\n drive_id = checkpoint.completion_map[\n self.primary_admin_email\n ].current_folder_or_drive_id\n if drive_id is None:\n raise ValueError(\"drive id not set in checkpoint\")\n resume_start = checkpoint.completion_map[\n self.primary_admin_email\n ].completed_until\n for file_or_token in _yield_from_drive(drive_id, resume_start):\n if isinstance(file_or_token, str):\n checkpoint.completion_map[\n self.primary_admin_email\n ].next_page_token = file_or_token\n return # done with the max num pages, return checkpoint\n yield file_or_token\n checkpoint.completion_map[self.primary_admin_email].next_page_token = None\n\n for drive_id in drive_ids_to_retrieve:\n if drive_id in self._retrieved_folder_and_drive_ids:\n logger.info(\n f\"Skipping drive '{drive_id}' as it has already been retrieved\"\n )\n continue\n logger.info(\n f\"Getting files in shared drive '{drive_id}' as '{self.primary_admin_email}'\"\n )\n for file_or_token in _yield_from_drive(drive_id, start):\n if isinstance(file_or_token, str):\n checkpoint.completion_map[\n self.primary_admin_email\n ].next_page_token = file_or_token\n return # done with the max num pages, return checkpoint\n yield file_or_token\n checkpoint.completion_map[self.primary_admin_email].next_page_token = None\n\n def _oauth_retrieval_folders(\n self,\n field_type: DriveFileFieldType,\n drive_service: GoogleDriveService,\n drive_ids_to_retrieve: set[str],\n folder_ids_to_retrieve: set[str],\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]:\n \"\"\"\n If there are any remaining folder ids to retrieve found earlier in the\n retrieval process, we recursively descend the file tree and retrieve all\n files in the folder(s).\n \"\"\"\n # Even if no folders were requested, we still check if any drives were requested\n # that could be folders.\n remaining_folders = (\n folder_ids_to_retrieve - self._retrieved_folder_and_drive_ids\n )\n\n def _yield_from_folder_crawl(\n folder_id: str, folder_start: SecondsSinceUnixEpoch | None\n ) -> Iterator[RetrievedDriveFile]:\n yield from crawl_folders_for_files(\n service=drive_service,\n parent_id=folder_id,\n field_type=field_type,\n user_email=self.primary_admin_email,\n traversed_parent_ids=self._retrieved_folder_and_drive_ids,\n update_traversed_ids_func=self._update_traversed_parent_ids,\n start=folder_start,\n end=end,\n )\n\n # resume from a checkpoint\n # TODO: actually checkpoint folder retrieval. Since we moved towards returning from\n # generator functions to indicate when a checkpoint should be returned, this code\n # shouldn't be used currently. Unfortunately folder crawling is quite difficult to checkpoint\n # effectively (likely need separate folder crawling and file retrieval stages),\n # so we'll revisit this later.\n if checkpoint.completion_map[\n self.primary_admin_email\n ].stage == DriveRetrievalStage.FOLDER_FILES and (\n folder_id := checkpoint.completion_map[\n self.primary_admin_email\n ].current_folder_or_drive_id\n ):\n resume_start = checkpoint.completion_map[\n self.primary_admin_email\n ].completed_until\n yield from _yield_from_folder_crawl(folder_id, resume_start)\n\n # the times stored in the completion_map aren't used due to the crawling behavior\n # instead, the traversed_parent_ids are used to determine what we have left to retrieve\n for folder_id in remaining_folders:\n logger.info(\n f\"Getting files in folder '{folder_id}' as '{self.primary_admin_email}'\"\n )\n yield from _yield_from_folder_crawl(folder_id, start)\n\n remaining_folders = (\n drive_ids_to_retrieve | folder_ids_to_retrieve\n ) - self._retrieved_folder_and_drive_ids\n if remaining_folders:\n logger.warning(\n f\"Some folders/drives were not retrieved. IDs: {remaining_folders}\"\n )\n\n def _checkpointed_retrieval(\n self,\n retrieval_method: CredentialedRetrievalMethod,\n field_type: DriveFileFieldType,\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]:\n drive_files = retrieval_method(\n field_type=field_type,\n checkpoint=checkpoint,\n start=start,\n end=end,\n )\n\n for file in drive_files:\n drive_file = file.drive_file or {}\n completion = checkpoint.completion_map[file.user_email]\n\n completed_until = completion.completed_until\n modified_time = drive_file.get(GoogleFields.MODIFIED_TIME.value)\n if isinstance(modified_time, str):\n try:\n completed_until = datetime.fromisoformat(modified_time).timestamp()\n except ValueError:\n logger.warning(\n \"Invalid modifiedTime for file '%s' (stage=%s, user=%s).\",\n drive_file.get(\"id\"),\n file.completion_stage,\n file.user_email,\n )\n\n completion.update(\n stage=file.completion_stage,\n completed_until=completed_until,\n current_folder_or_drive_id=file.parent_id,\n )\n\n if file.error is not None or not drive_file:\n yield file\n continue\n\n try:\n document_id = onyx_document_id_from_drive_file(drive_file)\n except KeyError as exc:\n logger.warning(\n \"Drive file missing id/webViewLink (stage=%s user=%s). Skipping.\",\n file.completion_stage,\n file.user_email,\n )\n if file.error is None:\n file.error = exc\n yield file\n continue\n\n logger.debug(\n f\"Updating checkpoint for file: {drive_file.get('name')}. \"\n f\"Seen: {document_id in checkpoint.all_retrieved_file_ids}\"\n )\n if document_id in checkpoint.all_retrieved_file_ids:\n continue\n\n checkpoint.all_retrieved_file_ids.add(document_id)\n yield file\n\n def _manage_oauth_retrieval(\n self,\n field_type: DriveFileFieldType,\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]:\n if checkpoint.completion_stage == DriveRetrievalStage.START:\n checkpoint.completion_stage = DriveRetrievalStage.OAUTH_FILES\n checkpoint.completion_map[self.primary_admin_email] = StageCompletion(\n stage=DriveRetrievalStage.START,\n completed_until=0,\n current_folder_or_drive_id=None,\n )\n\n drive_service = get_drive_service(self.creds, self.primary_admin_email)\n\n if checkpoint.completion_stage == DriveRetrievalStage.OAUTH_FILES:\n completion = checkpoint.completion_map[self.primary_admin_email]\n all_files_start = start\n # if resuming from a checkpoint\n if completion.stage == DriveRetrievalStage.OAUTH_FILES:\n all_files_start = completion.completed_until\n\n for file_or_token in self._oauth_retrieval_all_files(\n field_type=field_type,\n drive_service=drive_service,\n start=all_files_start,\n end=end,\n page_token=checkpoint.completion_map[\n self.primary_admin_email\n ].next_page_token,\n ):\n if isinstance(file_or_token, str):\n checkpoint.completion_map[\n self.primary_admin_email\n ].next_page_token = file_or_token\n return # done with the max num pages, return checkpoint\n yield file_or_token\n checkpoint.completion_stage = DriveRetrievalStage.DRIVE_IDS\n checkpoint.completion_map[self.primary_admin_email].next_page_token = None\n return # create a new checkpoint\n\n all_requested = (\n self.include_files_shared_with_me\n and self.include_my_drives\n and self.include_shared_drives\n )\n if all_requested:\n # If all 3 are true, we already yielded from get_all_files_for_oauth\n checkpoint.completion_stage = DriveRetrievalStage.DONE\n return\n\n sorted_drive_ids, sorted_folder_ids = self._determine_retrieval_ids(\n checkpoint, DriveRetrievalStage.SHARED_DRIVE_FILES\n )\n\n if checkpoint.completion_stage == DriveRetrievalStage.SHARED_DRIVE_FILES:\n for file_or_token in self._oauth_retrieval_drives(\n field_type=field_type,\n drive_service=drive_service,\n drive_ids_to_retrieve=sorted_drive_ids,\n checkpoint=checkpoint,\n start=start,\n end=end,\n ):\n if isinstance(file_or_token, str):\n checkpoint.completion_map[\n self.primary_admin_email\n ].next_page_token = file_or_token\n return # done with the max num pages, return checkpoint\n yield file_or_token\n checkpoint.completion_stage = DriveRetrievalStage.FOLDER_FILES\n checkpoint.completion_map[self.primary_admin_email].next_page_token = None\n return # create a new checkpoint\n\n if checkpoint.completion_stage == DriveRetrievalStage.FOLDER_FILES:\n yield from self._oauth_retrieval_folders(\n field_type=field_type,\n drive_service=drive_service,\n drive_ids_to_retrieve=set(sorted_drive_ids),\n folder_ids_to_retrieve=set(sorted_folder_ids),\n checkpoint=checkpoint,\n start=start,\n end=end,\n )\n\n checkpoint.completion_stage = DriveRetrievalStage.DONE\n\n def _fetch_drive_items(\n self,\n field_type: DriveFileFieldType,\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> Iterator[RetrievedDriveFile]:\n retrieval_method = (\n self._manage_service_account_retrieval\n if isinstance(self.creds, ServiceAccountCredentials)\n else self._manage_oauth_retrieval\n )\n\n return self._checkpointed_retrieval(\n retrieval_method=retrieval_method,\n field_type=field_type,\n checkpoint=checkpoint,\n start=start,\n end=end,\n )\n\n def _convert_retrieved_files_to_documents(\n self,\n drive_files_iter: Iterator[RetrievedDriveFile],\n checkpoint: GoogleDriveCheckpoint,\n include_permissions: bool,\n ) -> Iterator[Document | ConnectorFailure | HierarchyNode]:\n \"\"\"\n Converts retrieved files to documents, yielding HierarchyNode\n objects for ancestor folders before the converted documents.\n \"\"\"\n permission_sync_context = (\n PermissionSyncContext(\n primary_admin_email=self.primary_admin_email,\n google_domain=self.google_domain,\n )\n if include_permissions\n else None\n )\n\n files_batch: list[RetrievedDriveFile] = []\n for retrieved_file in drive_files_iter:\n if self.exclude_domain_link_only and has_link_only_permission(\n retrieved_file.drive_file\n ):\n continue\n if retrieved_file.error is None:\n files_batch.append(retrieved_file)\n continue\n\n failure_stage = retrieved_file.completion_stage.value\n failure_message = f\"retrieval failure during stage: {failure_stage},\"\n failure_message += f\"user: {retrieved_file.user_email},\"\n failure_message += f\"parent drive/folder: {retrieved_file.parent_id},\"\n failure_message += f\"error: {retrieved_file.error}\"\n logger.error(failure_message)\n yield ConnectorFailure(\n failed_entity=EntityFailure(\n entity_id=retrieved_file.drive_file.get(\"id\", failure_stage),\n ),\n failure_message=failure_message,\n exception=retrieved_file.error,\n )\n\n new_ancestors = self._get_new_ancestors_for_files(\n files=files_batch,\n seen_hierarchy_node_raw_ids=checkpoint.seen_hierarchy_node_raw_ids,\n fully_walked_hierarchy_node_raw_ids=checkpoint.fully_walked_hierarchy_node_raw_ids,\n permission_sync_context=permission_sync_context,\n add_prefix=True,\n )\n if new_ancestors:\n logger.debug(f\"Yielding {len(new_ancestors)} new hierarchy nodes\")\n yield from new_ancestors\n\n func_with_args = [\n (\n self._convert_retrieved_file_to_document,\n (retrieved_file, permission_sync_context),\n )\n for retrieved_file in files_batch\n ]\n raw_results = cast(\n list[Document | ConnectorFailure | None],\n run_functions_tuples_in_parallel(func_with_args, max_workers=8),\n )\n\n results: list[Document | ConnectorFailure] = [\n r for r in raw_results if r is not None\n ]\n logger.debug(f\"batch has {len(results)} docs or failures\")\n yield from results\n\n checkpoint.retrieved_folder_and_drive_ids = self._retrieved_folder_and_drive_ids\n\n def _convert_retrieved_file_to_document(\n self,\n retrieved_file: RetrievedDriveFile,\n permission_sync_context: PermissionSyncContext | None,\n ) -> Document | ConnectorFailure | None:\n \"\"\"\n Converts a single retrieved file to a document.\n \"\"\"\n try:\n return convert_drive_item_to_document(\n self.creds,\n self.allow_images,\n self.size_threshold,\n permission_sync_context,\n [retrieved_file.user_email, self.primary_admin_email]\n + get_file_owners(retrieved_file.drive_file, self.primary_admin_email),\n retrieved_file.drive_file,\n )\n except Exception as e:\n logger.exception(\n f\"Error extracting document: \"\n f\"{retrieved_file.drive_file.get('name')} from Google Drive\"\n )\n return ConnectorFailure(\n failed_entity=EntityFailure(\n entity_id=retrieved_file.drive_file.get(\"id\", \"unknown\"),\n ),\n failure_message=(\n f\"Error extracting document: \"\n f\"{retrieved_file.drive_file.get('name')}\"\n ),\n exception=e,\n )\n\n def _load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GoogleDriveCheckpoint,\n include_permissions: bool,\n ) -> CheckpointOutput[GoogleDriveCheckpoint]:\n \"\"\"\n Entrypoint for the connector; first run is with an empty checkpoint.\n \"\"\"\n if self._creds is None or self._primary_admin_email is None:\n raise RuntimeError(\n \"Credentials missing, should not call this method before calling load_credentials\"\n )\n\n logger.info(\n f\"Loading from checkpoint with completion stage: {checkpoint.completion_stage},\"\n f\"num retrieved ids: {len(checkpoint.all_retrieved_file_ids)}\"\n )\n checkpoint = copy.deepcopy(checkpoint)\n self._retrieved_folder_and_drive_ids = checkpoint.retrieved_folder_and_drive_ids\n try:\n field_type = (\n DriveFileFieldType.WITH_PERMISSIONS\n if include_permissions or self.exclude_domain_link_only\n else DriveFileFieldType.STANDARD\n )\n drive_files_iter = self._fetch_drive_items(\n field_type=field_type,\n checkpoint=checkpoint,\n start=start,\n end=end,\n )\n yield from self._convert_retrieved_files_to_documents(\n drive_files_iter, checkpoint, include_permissions\n )\n except Exception as e:\n if MISSING_SCOPES_ERROR_STR in str(e):\n raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e\n raise e\n checkpoint.retrieved_folder_and_drive_ids = self._retrieved_folder_and_drive_ids\n\n logger.info(\n f\"num drive files retrieved: {len(checkpoint.all_retrieved_file_ids)}\"\n )\n if checkpoint.completion_stage == DriveRetrievalStage.DONE:\n checkpoint.has_more = False\n return checkpoint\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GoogleDriveCheckpoint,\n ) -> CheckpointOutput[GoogleDriveCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=False\n )\n\n @override\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: GoogleDriveCheckpoint,\n ) -> CheckpointOutput[GoogleDriveCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=True\n )\n\n def _extract_slim_docs_from_google_drive(\n self,\n checkpoint: GoogleDriveCheckpoint,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n files_batch: list[RetrievedDriveFile] = []\n slim_batch: list[SlimDocument | HierarchyNode] = []\n\n def _yield_slim_batch() -> list[SlimDocument | HierarchyNode]:\n \"\"\"Process files batch and return items to yield (hierarchy nodes + slim docs).\"\"\"\n nonlocal files_batch, slim_batch\n\n # Get new ancestor hierarchy nodes first\n permission_sync_context = PermissionSyncContext(\n primary_admin_email=self.primary_admin_email,\n google_domain=self.google_domain,\n )\n new_ancestors = self._get_new_ancestors_for_files(\n files=files_batch,\n seen_hierarchy_node_raw_ids=checkpoint.seen_hierarchy_node_raw_ids,\n fully_walked_hierarchy_node_raw_ids=checkpoint.fully_walked_hierarchy_node_raw_ids,\n permission_sync_context=permission_sync_context,\n )\n\n # Build slim documents\n for file in files_batch:\n if doc := build_slim_document(\n self.creds,\n file.drive_file,\n PermissionSyncContext(\n primary_admin_email=self.primary_admin_email,\n google_domain=self.google_domain,\n ),\n retriever_email=file.user_email,\n ):\n slim_batch.append(doc)\n\n # Combine: hierarchy nodes first, then slim docs\n result: list[SlimDocument | HierarchyNode] = []\n result.extend(new_ancestors)\n result.extend(slim_batch)\n files_batch = []\n slim_batch = []\n return result\n\n for file in self._fetch_drive_items(\n field_type=DriveFileFieldType.SLIM,\n checkpoint=checkpoint,\n start=start,\n end=end,\n ):\n if file.error is not None:\n raise file.error\n if self.exclude_domain_link_only and has_link_only_permission(\n file.drive_file\n ):\n continue\n files_batch.append(file)\n\n if len(files_batch) >= SLIM_BATCH_SIZE:\n yield _yield_slim_batch()\n if callback:\n if callback.should_stop():\n raise RuntimeError(\n \"_extract_slim_docs_from_google_drive: Stop signal detected\"\n )\n callback.progress(\"_extract_slim_docs_from_google_drive\", 1)\n\n # Yield remaining files\n if files_batch:\n yield _yield_slim_batch()\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n try:\n checkpoint = self.build_dummy_checkpoint()\n while checkpoint.completion_stage != DriveRetrievalStage.DONE:\n yield from self._extract_slim_docs_from_google_drive(\n checkpoint=checkpoint,\n start=start,\n end=end,\n callback=callback,\n )\n logger.info(\"Drive perm sync: Slim doc retrieval complete\")\n\n except Exception as e:\n if MISSING_SCOPES_ERROR_STR in str(e):\n raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e\n raise e\n\n def validate_connector_settings(self) -> None:\n if self._creds is None:\n raise ConnectorMissingCredentialError(\n \"Google Drive credentials not loaded.\"\n )\n\n if self._primary_admin_email is None:\n raise ConnectorValidationError(\n \"Primary admin email not found in credentials. Ensure DB_CREDENTIALS_PRIMARY_ADMIN_KEY is set.\"\n )\n\n try:\n drive_service = get_drive_service(self._creds, self._primary_admin_email)\n drive_service.files().list(pageSize=1, fields=\"files(id)\").execute()\n\n if isinstance(self._creds, ServiceAccountCredentials):\n # default is ~17mins of retries, don't do that here since this is called from\n # the UI\n retry_builder(tries=3, delay=0.1)(get_root_folder_id)(drive_service)\n\n except HttpError as e:\n status_code = e.resp.status if e.resp else None\n if status_code == 401:\n raise CredentialExpiredError(\n \"Invalid or expired Google Drive credentials (401).\"\n )\n elif status_code == 403:\n raise InsufficientPermissionsError(\n \"Google Drive app lacks required permissions (403). \"\n \"Please ensure the necessary scopes are granted and Drive \"\n \"apps are enabled.\"\n )\n else:\n raise ConnectorValidationError(\n f\"Unexpected Google Drive error (status={status_code}): {e}\"\n )\n\n except Exception as e:\n # Check for scope-related hints from the error message\n if MISSING_SCOPES_ERROR_STR in str(e):\n raise InsufficientPermissionsError(\n f\"Google Drive credentials are missing required scopes. {ONYX_SCOPE_INSTRUCTIONS}\"\n )\n raise ConnectorValidationError(\n f\"Unexpected error during Google Drive validation: {e}\"\n )\n\n @override\n def build_dummy_checkpoint(self) -> GoogleDriveCheckpoint:\n return GoogleDriveCheckpoint(\n retrieved_folder_and_drive_ids=set(),\n completion_stage=DriveRetrievalStage.START,\n completion_map=ThreadSafeDict(),\n all_retrieved_file_ids=set(),\n has_more=True,\n )\n\n @override\n def validate_checkpoint_json(self, checkpoint_json: str) -> GoogleDriveCheckpoint:\n return GoogleDriveCheckpoint.model_validate_json(checkpoint_json)\n\n\ndef get_credentials_from_env(email: str, oauth: bool) -> dict:\n if oauth:\n raw_credential_string = os.environ[\"GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR\"]\n else:\n raw_credential_string = os.environ[\"GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR\"]\n\n refried_credential_string = json.dumps(json.loads(raw_credential_string))\n\n # This is the Oauth token\n DB_CREDENTIALS_DICT_TOKEN_KEY = \"google_tokens\"\n # This is the service account key\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY = \"google_service_account_key\"\n # The email saved for both auth types\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY = \"google_primary_admin\"\n DB_CREDENTIALS_AUTHENTICATION_METHOD = \"authentication_method\"\n cred_key = (\n DB_CREDENTIALS_DICT_TOKEN_KEY\n if oauth\n else DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY\n )\n return {\n cred_key: refried_credential_string,\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: email,\n DB_CREDENTIALS_AUTHENTICATION_METHOD: \"uploaded\",\n }\n\n\nclass CheckpointOutputWrapper:\n \"\"\"\n Wraps a CheckpointOutput generator to give things back in a more digestible format.\n The connector format is easier for the connector implementor (e.g. it enforces exactly\n one new checkpoint is returned AND that the checkpoint is at the end), thus the different\n formats.\n \"\"\"\n\n def __init__(self) -> None:\n self.next_checkpoint: GoogleDriveCheckpoint | None = None\n\n def __call__(\n self,\n checkpoint_connector_generator: CheckpointOutput[GoogleDriveCheckpoint],\n ) -> Generator[\n tuple[Document | None, ConnectorFailure | None, GoogleDriveCheckpoint | None],\n None,\n None,\n ]:\n # grabs the final return value and stores it in the `next_checkpoint` variable\n def _inner_wrapper(\n checkpoint_connector_generator: CheckpointOutput[GoogleDriveCheckpoint],\n ) -> CheckpointOutput[GoogleDriveCheckpoint]:\n self.next_checkpoint = yield from checkpoint_connector_generator\n return self.next_checkpoint # not used\n\n for document_or_failure in _inner_wrapper(checkpoint_connector_generator):\n if isinstance(document_or_failure, Document):\n yield document_or_failure, None, None\n elif isinstance(document_or_failure, ConnectorFailure):\n yield None, document_or_failure, None\n else:\n raise ValueError(\n f\"Invalid document_or_failure type: {type(document_or_failure)}\"\n )\n\n if self.next_checkpoint is None:\n raise RuntimeError(\n \"Checkpoint is None. This should never happen - the connector should always return a checkpoint.\"\n )\n\n yield None, None, self.next_checkpoint\n\n\ndef yield_all_docs_from_checkpoint_connector(\n connector: GoogleDriveConnector,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n) -> Iterator[Document | ConnectorFailure]:\n num_iterations = 0\n\n checkpoint = connector.build_dummy_checkpoint()\n while checkpoint.has_more:\n doc_batch_generator = CheckpointOutputWrapper()(\n connector.load_from_checkpoint(start, end, checkpoint)\n )\n for document, failure, next_checkpoint in doc_batch_generator:\n if failure is not None:\n yield failure\n if document is not None:\n yield document\n if next_checkpoint is not None:\n checkpoint = next_checkpoint\n\n num_iterations += 1\n if num_iterations > 100_000:\n raise RuntimeError(\"Too many iterations. Infinite loop?\")\n\n\nif __name__ == \"__main__\":\n import time\n\n creds = get_credentials_from_env(\n os.environ[\"GOOGLE_DRIVE_PRIMARY_ADMIN_EMAIL\"], False\n )\n connector = GoogleDriveConnector(\n include_shared_drives=True,\n shared_drive_urls=None,\n include_my_drives=True,\n my_drive_emails=None,\n shared_folder_urls=None,\n include_files_shared_with_me=True,\n specific_user_emails=None,\n )\n connector.load_credentials(creds)\n max_fsize = 0\n biggest_fsize = 0\n num_errors = 0\n start_time = time.time()\n with open(\"stats.txt\", \"w\") as f:\n for num, doc_or_failure in enumerate(\n yield_all_docs_from_checkpoint_connector(connector, 0, time.time())\n ):\n if num % 200 == 0:\n f.write(f\"Processed {num} files\\n\")\n f.write(f\"Max file size: {max_fsize / 1000_000:.2f} MB\\n\")\n f.write(f\"Time so far: {time.time() - start_time:.2f} seconds\\n\")\n f.write(\n f\"Docs per minute: {num / (time.time() - start_time) * 60:.2f}\\n\"\n )\n biggest_fsize = max(biggest_fsize, max_fsize)\n max_fsize = 0\n if isinstance(doc_or_failure, Document):\n max_fsize = max(max_fsize, sys.getsizeof(doc_or_failure))\n elif isinstance(doc_or_failure, ConnectorFailure):\n num_errors += 1\n print(f\"Num errors: {num_errors}\")\n print(f\"Biggest file size: {biggest_fsize / 1000_000:.2f} MB\")\n print(f\"Time taken: {time.time() - start_time:.2f} seconds\")\n" }, { "path": "backend/onyx/connectors/google_drive/constants.py", "content": "UNSUPPORTED_FILE_TYPE_CONTENT = \"\" # keep empty for now\nDRIVE_FOLDER_TYPE = \"application/vnd.google-apps.folder\"\nDRIVE_SHORTCUT_TYPE = \"application/vnd.google-apps.shortcut\"\nDRIVE_FILE_TYPE = \"application/vnd.google-apps.file\"\n" }, { "path": "backend/onyx/connectors/google_drive/doc_conversion.py", "content": "import io\nfrom collections.abc import Callable\nfrom datetime import datetime\nfrom typing import Any\nfrom typing import cast\nfrom urllib.parse import urlparse\nfrom urllib.parse import urlunparse\n\nfrom googleapiclient.errors import HttpError # type: ignore\nfrom googleapiclient.http import MediaIoBaseDownload # type: ignore\nfrom pydantic import BaseModel\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.google_drive.constants import DRIVE_FOLDER_TYPE\nfrom onyx.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE\nfrom onyx.connectors.google_drive.models import GDriveMimeType\nfrom onyx.connectors.google_drive.models import GoogleDriveFileType\nfrom onyx.connectors.google_drive.section_extraction import get_document_sections\nfrom onyx.connectors.google_drive.section_extraction import HEADING_DELIMITER\nfrom onyx.connectors.google_utils.resources import get_drive_service\nfrom onyx.connectors.google_utils.resources import get_google_docs_service\nfrom onyx.connectors.google_utils.resources import GoogleDocsService\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.extract_file_text import pptx_to_text\nfrom onyx.file_processing.extract_file_text import read_docx_file\nfrom onyx.file_processing.extract_file_text import read_pdf_file\nfrom onyx.file_processing.extract_file_text import xlsx_to_text\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.file_types import OnyxMimeTypes\nfrom onyx.file_processing.image_utils import store_image_and_create_section\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\nfrom onyx.utils.variable_functionality import noop_fallback\n\nlogger = setup_logger()\n\n# Cache for folder path lookups to avoid redundant API calls\n# Maps folder_id -> (folder_name, parent_id)\n_folder_cache: dict[str, tuple[str, str | None]] = {}\n\n\ndef _get_folder_info(\n service: GoogleDriveService, folder_id: str\n) -> tuple[str, str | None]:\n \"\"\"Fetch folder name and parent ID, with caching.\"\"\"\n if folder_id in _folder_cache:\n return _folder_cache[folder_id]\n\n try:\n folder = (\n service.files()\n .get(\n fileId=folder_id,\n fields=\"name, parents\",\n supportsAllDrives=True,\n )\n .execute()\n )\n folder_name = folder.get(\"name\", \"Unknown\")\n parents = folder.get(\"parents\", [])\n parent_id = parents[0] if parents else None\n _folder_cache[folder_id] = (folder_name, parent_id)\n return folder_name, parent_id\n except HttpError as e:\n logger.warning(f\"Failed to get folder info for {folder_id}: {e}\")\n _folder_cache[folder_id] = (\"Unknown\", None)\n return \"Unknown\", None\n\n\ndef _get_drive_name(service: GoogleDriveService, drive_id: str) -> str:\n \"\"\"Fetch shared drive name.\"\"\"\n cache_key = f\"drive_{drive_id}\"\n if cache_key in _folder_cache:\n return _folder_cache[cache_key][0]\n\n try:\n drive = service.drives().get(driveId=drive_id).execute()\n drive_name = drive.get(\"name\", f\"Shared Drive {drive_id}\")\n _folder_cache[cache_key] = (drive_name, None)\n return drive_name\n except HttpError as e:\n logger.warning(f\"Failed to get drive name for {drive_id}: {e}\")\n _folder_cache[cache_key] = (f\"Shared Drive {drive_id}\", None)\n return f\"Shared Drive {drive_id}\"\n\n\ndef build_folder_path(\n file: GoogleDriveFileType,\n service: GoogleDriveService,\n drive_id: str | None = None,\n user_email: str | None = None,\n) -> list[str]:\n \"\"\"\n Build the full folder path for a file by walking up the parent chain.\n Returns a list of folder names from root to immediate parent.\n\n Args:\n file: The Google Drive file object\n service: Google Drive service instance\n drive_id: Optional drive ID (will be extracted from file if not provided)\n user_email: Optional user email to check ownership for \"My Drive\" vs \"Shared with me\"\n \"\"\"\n path_parts: list[str] = []\n\n # Get drive_id from file if not provided\n if drive_id is None:\n drive_id = file.get(\"driveId\")\n\n # Check if file is owned by the user (for distinguishing \"My Drive\" vs \"Shared with me\")\n is_owned_by_user = False\n if user_email:\n owners = file.get(\"owners\", [])\n is_owned_by_user = any(\n owner.get(\"emailAddress\", \"\").lower() == user_email.lower()\n for owner in owners\n )\n\n # Get the file's parent folder ID\n parents = file.get(\"parents\", [])\n if not parents:\n # File is at root level\n if drive_id:\n return [_get_drive_name(service, drive_id)]\n # If not in a shared drive, check if it's owned by the user\n if is_owned_by_user:\n return [\"My Drive\"]\n else:\n return [\"Shared with me\"]\n\n parent_id: str | None = parents[0]\n\n # Walk up the folder hierarchy (limit to 50 levels to prevent infinite loops)\n visited: set[str] = set()\n for _ in range(50):\n if not parent_id or parent_id in visited:\n break\n visited.add(parent_id)\n\n folder_name, next_parent = _get_folder_info(service, parent_id)\n\n # Check if we've reached the root (parent is the drive itself or no parent)\n if next_parent is None:\n # This folder's name is either the drive root, My Drive, or Shared with me\n if drive_id:\n path_parts.insert(0, _get_drive_name(service, drive_id))\n else:\n # Not in a shared drive - determine if it's \"My Drive\" or \"Shared with me\"\n if is_owned_by_user:\n path_parts.insert(0, \"My Drive\")\n else:\n path_parts.insert(0, \"Shared with me\")\n break\n else:\n path_parts.insert(0, folder_name)\n parent_id = next_parent\n\n # If we didn't find a root, determine the root based on ownership and drive\n if not path_parts:\n if drive_id:\n return [_get_drive_name(service, drive_id)]\n elif is_owned_by_user:\n return [\"My Drive\"]\n else:\n return [\"Shared with me\"]\n\n return path_parts\n\n\n# This is not a standard valid unicode char, it is used by the docs advanced API to\n# represent smart chips (elements like dates and doc links).\nSMART_CHIP_CHAR = \"\\ue907\"\nWEB_VIEW_LINK_KEY = \"webViewLink\"\n# Fallback templates for generating web links when Drive omits webViewLink.\n_FALLBACK_WEB_VIEW_LINK_TEMPLATES = {\n GDriveMimeType.DOC.value: \"https://docs.google.com/document/d/{}/view\",\n GDriveMimeType.SPREADSHEET.value: \"https://docs.google.com/spreadsheets/d/{}/view\",\n GDriveMimeType.PPT.value: \"https://docs.google.com/presentation/d/{}/view\",\n}\n\nMAX_RETRIEVER_EMAILS = 20\nCHUNK_SIZE_BUFFER = 64 # extra bytes past the limit to read\n\n# Mapping of Google Drive mime types to export formats\nGOOGLE_MIME_TYPES_TO_EXPORT = {\n GDriveMimeType.DOC.value: \"text/plain\",\n GDriveMimeType.SPREADSHEET.value: \"text/csv\",\n GDriveMimeType.PPT.value: \"text/plain\",\n}\n\n# Define Google MIME types mapping\nGOOGLE_MIME_TYPES = {\n GDriveMimeType.DOC.value: \"text/plain\",\n GDriveMimeType.SPREADSHEET.value: \"text/csv\",\n GDriveMimeType.PPT.value: \"text/plain\",\n}\n\n\nclass PermissionSyncContext(BaseModel):\n \"\"\"\n This is the information that is needed to sync permissions for a document.\n \"\"\"\n\n primary_admin_email: str\n google_domain: str\n\n\ndef onyx_document_id_from_drive_file(file: GoogleDriveFileType) -> str:\n link = file.get(WEB_VIEW_LINK_KEY)\n if not link:\n file_id = file.get(\"id\")\n if not file_id:\n raise KeyError(\n f\"Google Drive file missing both '{WEB_VIEW_LINK_KEY}' and 'id' fields.\"\n )\n mime_type = file.get(\"mimeType\", \"\")\n template = _FALLBACK_WEB_VIEW_LINK_TEMPLATES.get(mime_type)\n if template is None:\n link = f\"https://drive.google.com/file/d/{file_id}/view\"\n else:\n link = template.format(file_id)\n logger.debug(\n \"Missing webViewLink for Google Drive file with id %s. Falling back to constructed link %s\",\n file_id,\n link,\n )\n parsed_url = urlparse(link)\n parsed_url = parsed_url._replace(query=\"\") # remove query parameters\n spl_path = parsed_url.path.split(\"/\")\n if spl_path and (spl_path[-1] in [\"edit\", \"view\", \"preview\"]):\n spl_path.pop()\n parsed_url = parsed_url._replace(path=\"/\".join(spl_path))\n # Remove query parameters and reconstruct URL\n return urlunparse(parsed_url)\n\n\ndef download_request(\n service: GoogleDriveService, file_id: str, size_threshold: int\n) -> bytes:\n \"\"\"\n Download the file from Google Drive.\n \"\"\"\n # For other file types, download the file\n # Use the correct API call for downloading files\n request = service.files().get_media(fileId=file_id)\n return _download_request(request, file_id, size_threshold)\n\n\n_DOWNLOAD_NUM_RETRIES = 3\n\n\ndef _download_request(request: Any, file_id: str, size_threshold: int) -> bytes:\n response_bytes = io.BytesIO()\n downloader = MediaIoBaseDownload(\n response_bytes, request, chunksize=size_threshold + CHUNK_SIZE_BUFFER\n )\n done = False\n while not done:\n # num_retries enables automatic retry with exponential backoff for transient errors\n download_progress, done = downloader.next_chunk(\n num_retries=_DOWNLOAD_NUM_RETRIES\n )\n if download_progress.resumable_progress > size_threshold:\n logger.warning(\n f\"File {file_id} exceeds size threshold of {size_threshold}. Skipping2.\"\n )\n return bytes()\n\n response = response_bytes.getvalue()\n if not response:\n logger.warning(f\"Failed to download {file_id}\")\n return bytes()\n return response\n\n\ndef _download_and_extract_sections_basic(\n file: dict[str, str],\n service: GoogleDriveService,\n allow_images: bool,\n size_threshold: int,\n) -> list[TextSection | ImageSection]:\n \"\"\"Extract text and images from a Google Drive file.\"\"\"\n file_id = file[\"id\"]\n file_name = file[\"name\"]\n mime_type = file[\"mimeType\"]\n link = file.get(WEB_VIEW_LINK_KEY, \"\")\n\n # For non-Google files, download the file\n # Use the correct API call for downloading files\n # lazy evaluation to only download the file if necessary\n def response_call() -> bytes:\n return download_request(service, file_id, size_threshold)\n\n if mime_type in OnyxMimeTypes.IMAGE_MIME_TYPES:\n # Skip images if not explicitly enabled\n if not allow_images:\n return []\n\n # Store images for later processing\n sections: list[TextSection | ImageSection] = []\n try:\n section, embedded_id = store_image_and_create_section(\n image_data=response_call(),\n file_id=file_id,\n display_name=file_name,\n media_type=mime_type,\n file_origin=FileOrigin.CONNECTOR,\n link=link,\n )\n sections.append(section)\n except Exception as e:\n logger.error(f\"Failed to process image {file_name}: {e}\")\n return sections\n\n # For Google Docs, Sheets, and Slides, export as plain text\n if mime_type in GOOGLE_MIME_TYPES_TO_EXPORT:\n export_mime_type = GOOGLE_MIME_TYPES_TO_EXPORT[mime_type]\n # Use the correct API call for exporting files\n request = service.files().export_media(\n fileId=file_id, mimeType=export_mime_type\n )\n response = _download_request(request, file_id, size_threshold)\n if not response:\n logger.warning(f\"Failed to export {file_name} as {export_mime_type}\")\n return []\n\n text = response.decode(\"utf-8\")\n return [TextSection(link=link, text=text)]\n\n # Process based on mime type\n if mime_type == \"text/plain\":\n try:\n text = response_call().decode(\"utf-8\")\n return [TextSection(link=link, text=text)]\n except UnicodeDecodeError as e:\n logger.warning(f\"Failed to extract text from {file_name}: {e}\")\n return []\n\n elif (\n mime_type\n == \"application/vnd.openxmlformats-officedocument.wordprocessingml.document\"\n ):\n text, _ = read_docx_file(io.BytesIO(response_call()))\n return [TextSection(link=link, text=text)]\n\n elif (\n mime_type == \"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet\"\n ):\n text = xlsx_to_text(io.BytesIO(response_call()), file_name=file_name)\n return [TextSection(link=link, text=text)] if text else []\n\n elif (\n mime_type\n == \"application/vnd.openxmlformats-officedocument.presentationml.presentation\"\n ):\n text = pptx_to_text(io.BytesIO(response_call()), file_name=file_name)\n return [TextSection(link=link, text=text)] if text else []\n\n elif mime_type == \"application/pdf\":\n text, _pdf_meta, images = read_pdf_file(io.BytesIO(response_call()))\n pdf_sections: list[TextSection | ImageSection] = [\n TextSection(link=link, text=text)\n ]\n\n # Process embedded images in the PDF\n try:\n for idx, (img_data, img_name) in enumerate(images):\n section, embedded_id = store_image_and_create_section(\n image_data=img_data,\n file_id=f\"{file_id}_img_{idx}\",\n display_name=img_name or f\"{file_name} - image {idx}\",\n file_origin=FileOrigin.CONNECTOR,\n )\n pdf_sections.append(section)\n except Exception as e:\n logger.error(f\"Failed to process PDF images in {file_name}: {e}\")\n return pdf_sections\n\n # Final attempt at extracting text\n file_ext = get_file_ext(file.get(\"name\", \"\"))\n if file_ext not in OnyxFileExtensions.ALL_ALLOWED_EXTENSIONS:\n logger.warning(f\"Skipping file {file.get('name')} due to extension.\")\n return []\n\n try:\n text = extract_file_text(io.BytesIO(response_call()), file_name)\n return [TextSection(link=link, text=text)]\n except Exception as e:\n logger.warning(f\"Failed to extract text from {file_name}: {e}\")\n return []\n\n\ndef _find_nth(haystack: str, needle: str, n: int, start: int = 0) -> int:\n start = haystack.find(needle, start)\n while start >= 0 and n > 1:\n start = haystack.find(needle, start + len(needle))\n n -= 1\n return start\n\n\ndef align_basic_advanced(\n basic_sections: list[TextSection | ImageSection], adv_sections: list[TextSection]\n) -> list[TextSection | ImageSection]:\n \"\"\"Align the basic sections with the advanced sections.\n In particular, the basic sections contain all content of the file,\n including smart chips like dates and doc links. The advanced sections\n are separated by section headers and contain header-based links that\n improve user experience when they click on the source in the UI.\n\n There are edge cases in text matching (i.e. the heading is a smart chip or\n there is a smart chip in the doc with text containing the actual heading text)\n that make the matching imperfect; this is hence done on a best-effort basis.\n \"\"\"\n if len(adv_sections) <= 1:\n return basic_sections # no benefit from aligning\n\n basic_full_text = \"\".join(\n [section.text for section in basic_sections if isinstance(section, TextSection)]\n )\n new_sections: list[TextSection | ImageSection] = []\n heading_start = 0\n for adv_ind in range(1, len(adv_sections)):\n heading = adv_sections[adv_ind].text.split(HEADING_DELIMITER)[0]\n # retrieve the longest part of the heading that is not a smart chip\n heading_key = max(heading.split(SMART_CHIP_CHAR), key=len).strip()\n if heading_key == \"\":\n logger.warning(\n f\"Cannot match heading: {heading}, its link will come from the following section\"\n )\n continue\n heading_offset = heading.find(heading_key)\n\n # count occurrences of heading str in previous section\n heading_count = adv_sections[adv_ind - 1].text.count(heading_key)\n\n prev_start = heading_start\n heading_start = (\n _find_nth(basic_full_text, heading_key, heading_count, start=prev_start)\n - heading_offset\n )\n if heading_start < 0:\n logger.warning(\n f\"Heading key {heading_key} from heading {heading} not found in basic text\"\n )\n heading_start = prev_start\n continue\n\n new_sections.append(\n TextSection(\n link=adv_sections[adv_ind - 1].link,\n text=basic_full_text[prev_start:heading_start],\n )\n )\n\n # handle last section\n new_sections.append(\n TextSection(link=adv_sections[-1].link, text=basic_full_text[heading_start:])\n )\n return new_sections\n\n\ndef _get_external_access_for_raw_gdrive_file(\n file: GoogleDriveFileType,\n company_domain: str,\n retriever_drive_service: GoogleDriveService | None,\n admin_drive_service: GoogleDriveService,\n fallback_user_email: str,\n add_prefix: bool = False,\n) -> ExternalAccess:\n \"\"\"\n Get the external access for a raw Google Drive file.\n\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path\n where upsert_document_external_perms handles prefixing).\n fallback_user_email: When permission info can't be retrieved (e.g. externally-owned\n files), fall back to granting access to this user.\n \"\"\"\n external_access_fn = cast(\n Callable[\n [\n GoogleDriveFileType,\n str,\n GoogleDriveService | None,\n GoogleDriveService,\n str,\n bool,\n ],\n ExternalAccess,\n ],\n fetch_versioned_implementation_with_fallback(\n \"onyx.external_permissions.google_drive.doc_sync\",\n \"get_external_access_for_raw_gdrive_file\",\n fallback=noop_fallback,\n ),\n )\n return external_access_fn(\n file,\n company_domain,\n retriever_drive_service,\n admin_drive_service,\n fallback_user_email,\n add_prefix,\n )\n\n\ndef convert_drive_item_to_document(\n creds: Any,\n allow_images: bool,\n size_threshold: int,\n # if not specified, we will not sync permissions\n # will also be a no-op if EE is not enabled\n permission_sync_context: PermissionSyncContext | None,\n retriever_emails: list[str],\n file: GoogleDriveFileType,\n) -> Document | ConnectorFailure | None:\n \"\"\"\n Attempt to convert a drive item to a document with each retriever email\n in order. returns upon a successful retrieval or a non-403 error.\n\n We used to always get the user email from the file owners when available,\n but this was causing issues with shared folders where the owner was not included in the service account\n now we use the email of the account that successfully listed the file. There are cases where a\n user that can list a file cannot download it, so we retry with file owners and admin email.\n \"\"\"\n first_error = None\n doc_or_failure = None\n retriever_emails = retriever_emails[:MAX_RETRIEVER_EMAILS]\n # use seen instead of list(set()) to avoid re-ordering the retriever emails\n seen = set()\n for retriever_email in retriever_emails:\n if retriever_email in seen:\n continue\n seen.add(retriever_email)\n doc_or_failure = _convert_drive_item_to_document(\n creds,\n allow_images,\n size_threshold,\n retriever_email,\n file,\n permission_sync_context,\n )\n\n # There are a variety of permissions-based errors that occasionally occur\n # when retrieving files. Often when these occur, there is another user\n # that can successfully retrieve the file, so we try the next user.\n if (\n doc_or_failure is None\n or isinstance(doc_or_failure, Document)\n or not (\n isinstance(doc_or_failure.exception, HttpError)\n and doc_or_failure.exception.status_code in [401, 403, 404]\n )\n ):\n return doc_or_failure\n\n if first_error is None:\n first_error = doc_or_failure\n else:\n first_error.failure_message += f\"\\n\\n{doc_or_failure.failure_message}\"\n\n if (\n first_error\n and isinstance(first_error.exception, HttpError)\n and first_error.exception.status_code == 403\n ):\n # This SHOULD happen very rarely, and we don't want to break the indexing process when\n # a high volume of 403s occurs early. We leave a verbose log to help investigate.\n logger.error(\n f\"Skipping file id: {file.get('id')} name: {file.get('name')} due to 403 error.\"\n f\"Attempted to retrieve with {retriever_emails},\"\n f\"got the following errors: {first_error.failure_message}\"\n )\n return None\n return first_error\n\n\ndef _convert_drive_item_to_document(\n creds: Any,\n allow_images: bool,\n size_threshold: int,\n retriever_email: str,\n file: GoogleDriveFileType,\n # if not specified, we will not sync permissions\n # will also be a no-op if EE is not enabled\n permission_sync_context: PermissionSyncContext | None,\n) -> Document | ConnectorFailure | None:\n \"\"\"\n Main entry point for converting a Google Drive file => Document object.\n \"\"\"\n sections: list[TextSection | ImageSection] = []\n\n # Only construct these services when needed\n def _get_drive_service() -> GoogleDriveService:\n return get_drive_service(creds, user_email=retriever_email)\n\n def _get_docs_service() -> GoogleDocsService:\n return get_google_docs_service(creds, user_email=retriever_email)\n\n doc_id = \"unknown\"\n\n try:\n # skip shortcuts or folders\n if file.get(\"mimeType\") in [DRIVE_SHORTCUT_TYPE, DRIVE_FOLDER_TYPE]:\n logger.info(\"Skipping shortcut/folder.\")\n return None\n\n size_str = file.get(\"size\")\n if size_str:\n try:\n size_int = int(size_str)\n except ValueError:\n logger.warning(f\"Parsing string to int failed: size_str={size_str}\")\n else:\n if size_int > size_threshold:\n logger.warning(\n f\"{file.get('name')} exceeds size threshold of {size_threshold}. Skipping.\"\n )\n return None\n\n # If it's a Google Doc, we might do advanced parsing\n if file.get(\"mimeType\") == GDriveMimeType.DOC.value:\n try:\n logger.debug(f\"starting advanced parsing for {file.get('name')}\")\n # get_document_sections is the advanced approach for Google Docs\n doc_sections = get_document_sections(\n docs_service=_get_docs_service(),\n doc_id=file.get(\"id\", \"\"),\n )\n if doc_sections:\n sections = cast(list[TextSection | ImageSection], doc_sections)\n if any(SMART_CHIP_CHAR in section.text for section in doc_sections):\n logger.debug(\n f\"found smart chips in {file.get('name')}, aligning with basic sections\"\n )\n basic_sections = _download_and_extract_sections_basic(\n file, _get_drive_service(), allow_images, size_threshold\n )\n sections = align_basic_advanced(basic_sections, doc_sections)\n\n except Exception as e:\n logger.warning(\n f\"Error in advanced parsing: {e}. Falling back to basic extraction.\"\n )\n # Not Google Doc, attempt basic extraction\n else:\n sections = _download_and_extract_sections_basic(\n file, _get_drive_service(), allow_images, size_threshold\n )\n\n # If we still don't have any sections, skip this file\n if not sections:\n logger.warning(f\"No content extracted from {file.get('name')}. Skipping.\")\n return None\n\n doc_id = onyx_document_id_from_drive_file(file)\n external_access = (\n _get_external_access_for_raw_gdrive_file(\n file=file,\n company_domain=permission_sync_context.google_domain,\n # try both retriever_email and primary_admin_email if necessary\n retriever_drive_service=_get_drive_service(),\n admin_drive_service=get_drive_service(\n creds, user_email=permission_sync_context.primary_admin_email\n ),\n add_prefix=True, # Indexing path - prefix here\n fallback_user_email=retriever_email,\n )\n if permission_sync_context\n else None\n )\n\n # Build doc_metadata with hierarchy information\n file_name = file.get(\"name\", \"\")\n mime_type = file.get(\"mimeType\", \"\")\n drive_id = file.get(\"driveId\")\n\n # Build full folder path by walking up the parent chain\n # Pass retriever_email to determine if file is in \"My Drive\" vs \"Shared with me\"\n source_path = build_folder_path(\n file, _get_drive_service(), drive_id, retriever_email\n )\n\n doc_metadata = {\n \"hierarchy\": {\n \"source_path\": source_path,\n \"drive_id\": drive_id,\n \"file_name\": file_name,\n \"mime_type\": mime_type,\n }\n }\n\n # Create the document\n return Document(\n id=doc_id,\n sections=sections,\n source=DocumentSource.GOOGLE_DRIVE,\n semantic_identifier=file_name,\n doc_metadata=doc_metadata,\n metadata={\n \"owner_names\": \", \".join(\n owner.get(\"displayName\", \"\") for owner in file.get(\"owners\", [])\n ),\n },\n doc_updated_at=datetime.fromisoformat(\n file.get(\"modifiedTime\", \"\").replace(\"Z\", \"+00:00\")\n ),\n external_access=external_access,\n parent_hierarchy_raw_node_id=(file.get(\"parents\") or [None])[0],\n )\n except Exception as e:\n doc_id = \"unknown\"\n try:\n doc_id = onyx_document_id_from_drive_file(file)\n except Exception as e2:\n logger.warning(f\"Error getting document id from file: {e2}\")\n\n file_name = file.get(\"name\")\n error_str = (\n f\"Error converting file '{file_name}' to Document as {retriever_email}: {e}\"\n )\n if isinstance(e, HttpError) and e.status_code == 403:\n logger.warning(\n f\"Uncommon permissions error while downloading file. User \"\n f\"{retriever_email} was able to see file {file_name} \"\n \"but cannot download it.\"\n )\n logger.warning(error_str)\n\n return ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=doc_id,\n document_link=(\n sections[0].link if sections else None\n ), # TODO: see if this is the best way to get a link\n ),\n failed_entity=None,\n failure_message=error_str,\n exception=e,\n )\n\n\ndef build_slim_document(\n creds: Any,\n file: GoogleDriveFileType,\n # if not specified, we will not sync permissions\n # will also be a no-op if EE is not enabled\n permission_sync_context: PermissionSyncContext | None,\n retriever_email: str,\n) -> SlimDocument | None:\n if file.get(\"mimeType\") in [DRIVE_FOLDER_TYPE, DRIVE_SHORTCUT_TYPE]:\n return None\n\n owner_email = cast(str | None, file.get(\"owners\", [{}])[0].get(\"emailAddress\"))\n external_access = (\n _get_external_access_for_raw_gdrive_file(\n file=file,\n company_domain=permission_sync_context.google_domain,\n retriever_drive_service=(\n get_drive_service(\n creds,\n user_email=owner_email,\n )\n if owner_email\n else None\n ),\n admin_drive_service=get_drive_service(\n creds,\n user_email=permission_sync_context.primary_admin_email,\n ),\n fallback_user_email=retriever_email,\n )\n if permission_sync_context\n else None\n )\n return SlimDocument(\n id=onyx_document_id_from_drive_file(file),\n external_access=external_access,\n parent_hierarchy_raw_node_id=(file.get(\"parents\") or [None])[0],\n )\n" }, { "path": "backend/onyx/connectors/google_drive/file_retrieval.py", "content": "from collections.abc import Callable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom enum import Enum\nfrom typing import cast\nfrom urllib.parse import parse_qs\nfrom urllib.parse import urlparse\n\nfrom googleapiclient.discovery import Resource # type: ignore\nfrom googleapiclient.errors import HttpError # type: ignore\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.google_drive.constants import DRIVE_FOLDER_TYPE\nfrom onyx.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE\nfrom onyx.connectors.google_drive.models import DriveRetrievalStage\nfrom onyx.connectors.google_drive.models import GoogleDriveFileType\nfrom onyx.connectors.google_drive.models import RetrievedDriveFile\nfrom onyx.connectors.google_utils.google_utils import execute_paginated_retrieval\nfrom onyx.connectors.google_utils.google_utils import (\n execute_paginated_retrieval_with_max_pages,\n)\nfrom onyx.connectors.google_utils.google_utils import GoogleFields\nfrom onyx.connectors.google_utils.google_utils import ORDER_BY_KEY\nfrom onyx.connectors.google_utils.google_utils import PAGE_TOKEN_KEY\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\nfrom onyx.utils.variable_functionality import noop_fallback\n\n\nlogger = setup_logger()\n\n\nclass DriveFileFieldType(Enum):\n \"\"\"Enum to specify which fields to retrieve from Google Drive files\"\"\"\n\n SLIM = \"slim\" # Minimal fields for basic file info\n STANDARD = \"standard\" # Standard fields including content metadata\n WITH_PERMISSIONS = \"with_permissions\" # Full fields including permissions\n\n\nPERMISSION_FULL_DESCRIPTION = (\n \"permissions(id, emailAddress, type, domain, allowFileDiscovery, permissionDetails)\"\n)\nFILE_FIELDS = (\n \"nextPageToken, files(mimeType, id, name, driveId, parents, \"\n \"modifiedTime, webViewLink, shortcutDetails, owners(emailAddress), size)\"\n)\nFILE_FIELDS_WITH_PERMISSIONS = (\n f\"nextPageToken, files(mimeType, id, name, driveId, parents, {PERMISSION_FULL_DESCRIPTION}, permissionIds, \"\n \"modifiedTime, webViewLink, shortcutDetails, owners(emailAddress), size)\"\n)\nSLIM_FILE_FIELDS = (\n f\"nextPageToken, files(mimeType, driveId, id, name, parents, {PERMISSION_FULL_DESCRIPTION}, \"\n \"permissionIds, webViewLink, owners(emailAddress), modifiedTime)\"\n)\nFOLDER_FIELDS = \"nextPageToken, files(id, name, permissions, modifiedTime, webViewLink, shortcutDetails)\"\n\nHIERARCHY_FIELDS = \"id, name, parents, webViewLink, mimeType, driveId\"\n\nHIERARCHY_FIELDS_WITH_PERMISSIONS = (\n \"id, name, parents, webViewLink, mimeType, permissionIds, driveId\"\n)\n\n\ndef generate_time_range_filter(\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n) -> str:\n time_range_filter = \"\"\n if start is not None:\n time_start = datetime.fromtimestamp(start, tz=timezone.utc).isoformat()\n time_range_filter += (\n f\" and {GoogleFields.MODIFIED_TIME.value} >= '{time_start}'\"\n )\n if end is not None:\n time_stop = datetime.fromtimestamp(end, tz=timezone.utc).isoformat()\n time_range_filter += f\" and {GoogleFields.MODIFIED_TIME.value} <= '{time_stop}'\"\n return time_range_filter\n\n\nLINK_ONLY_PERMISSION_TYPES = {\"domain\", \"anyone\"}\n\n\ndef has_link_only_permission(file: GoogleDriveFileType) -> bool:\n \"\"\"\n Return True if any permission requires a direct link to access\n (allowFileDiscovery is explicitly false for supported types).\n \"\"\"\n permissions = file.get(\"permissions\") or []\n for permission in permissions:\n if permission.get(\"type\") not in LINK_ONLY_PERMISSION_TYPES:\n continue\n if permission.get(\"allowFileDiscovery\") is False:\n return True\n return False\n\n\ndef _get_folders_in_parent(\n service: Resource,\n parent_id: str | None = None,\n) -> Iterator[GoogleDriveFileType]:\n # Follow shortcuts to folders\n query = f\"(mimeType = '{DRIVE_FOLDER_TYPE}' or mimeType = '{DRIVE_SHORTCUT_TYPE}')\"\n query += \" and trashed = false\"\n\n if parent_id:\n query += f\" and '{parent_id}' in parents\"\n\n for file in execute_paginated_retrieval(\n retrieval_function=service.files().list,\n list_key=\"files\",\n continue_on_404_or_403=True,\n corpora=\"allDrives\",\n supportsAllDrives=True,\n includeItemsFromAllDrives=True,\n fields=FOLDER_FIELDS,\n q=query,\n ):\n yield file\n\n\ndef get_folder_metadata(\n service: Resource,\n folder_id: str,\n field_type: DriveFileFieldType,\n) -> GoogleDriveFileType | None:\n \"\"\"Fetch metadata for a folder by ID.\"\"\"\n fields = _get_hierarchy_fields_for_file_type(field_type)\n try:\n return (\n service.files()\n .get(\n fileId=folder_id,\n fields=fields,\n supportsAllDrives=True,\n )\n .execute()\n )\n except HttpError as e:\n if e.resp.status in (403, 404):\n logger.debug(f\"Cannot access folder {folder_id}: {e}\")\n else:\n raise e\n return None\n\n\ndef _get_hierarchy_fields_for_file_type(field_type: DriveFileFieldType) -> str:\n if field_type == DriveFileFieldType.WITH_PERMISSIONS:\n return HIERARCHY_FIELDS_WITH_PERMISSIONS\n else:\n return HIERARCHY_FIELDS\n\n\ndef get_shared_drive_name(\n service: Resource,\n drive_id: str,\n) -> str | None:\n \"\"\"Fetch the actual name of a shared drive via the drives().get() API.\n\n The files().get() API returns 'Drive' as the name for shared drive root\n folders. Only drives().get() returns the real user-assigned name.\n \"\"\"\n try:\n drive = service.drives().get(driveId=drive_id, fields=\"name\").execute()\n return drive.get(\"name\")\n except HttpError as e:\n if e.resp.status in (403, 404):\n logger.debug(f\"Cannot access drive {drive_id}: {e}\")\n else:\n raise\n return None\n\n\ndef get_external_access_for_folder(\n folder: GoogleDriveFileType,\n google_domain: str,\n drive_service: GoogleDriveService,\n add_prefix: bool = False,\n) -> ExternalAccess:\n \"\"\"\n Extract ExternalAccess from a folder's permissions.\n\n This fetches permissions using the Drive API (via permissionIds) and extracts\n user emails, group emails, and public access status.\n\n Uses the EE implementation if available, otherwise returns public access\n (fallback for non-EE deployments).\n\n Args:\n folder: The folder metadata from Google Drive API (must include permissionIds field)\n google_domain: The company's Google Workspace domain (e.g., \"company.com\")\n drive_service: Google Drive service for fetching permission details\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path\n where upsert_document_external_perms handles prefixing).\n\n Returns:\n ExternalAccess with extracted permission info\n \"\"\"\n # Try to get the EE implementation\n get_folder_access_fn = cast(\n Callable[[GoogleDriveFileType, str, GoogleDriveService, bool], ExternalAccess],\n fetch_versioned_implementation_with_fallback(\n \"onyx.external_permissions.google_drive.doc_sync\",\n \"get_external_access_for_folder\",\n noop_fallback,\n ),\n )\n\n return get_folder_access_fn(folder, google_domain, drive_service, add_prefix)\n\n\ndef _get_fields_for_file_type(field_type: DriveFileFieldType) -> str:\n \"\"\"Get the appropriate fields string based on the field type enum\"\"\"\n if field_type == DriveFileFieldType.SLIM:\n return SLIM_FILE_FIELDS\n elif field_type == DriveFileFieldType.WITH_PERMISSIONS:\n return FILE_FIELDS_WITH_PERMISSIONS\n else: # DriveFileFieldType.STANDARD\n return FILE_FIELDS\n\n\ndef _get_files_in_parent(\n service: Resource,\n parent_id: str,\n field_type: DriveFileFieldType,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n) -> Iterator[GoogleDriveFileType]:\n query = f\"mimeType != '{DRIVE_FOLDER_TYPE}' and '{parent_id}' in parents\"\n query += \" and trashed = false\"\n query += generate_time_range_filter(start, end)\n\n kwargs = {ORDER_BY_KEY: GoogleFields.MODIFIED_TIME.value}\n\n for file in execute_paginated_retrieval(\n retrieval_function=service.files().list,\n list_key=\"files\",\n continue_on_404_or_403=True,\n corpora=\"allDrives\",\n supportsAllDrives=True,\n includeItemsFromAllDrives=True,\n fields=_get_fields_for_file_type(field_type),\n q=query,\n **kwargs,\n ):\n yield file\n\n\ndef crawl_folders_for_files(\n service: Resource,\n parent_id: str,\n field_type: DriveFileFieldType,\n user_email: str,\n traversed_parent_ids: set[str],\n update_traversed_ids_func: Callable[[str], None],\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n) -> Iterator[RetrievedDriveFile]:\n \"\"\"\n This function starts crawling from any folder. It is slower though.\n \"\"\"\n logger.info(\"Entered crawl_folders_for_files with parent_id: \" + parent_id)\n if parent_id not in traversed_parent_ids:\n logger.info(\"Parent id not in traversed parent ids, getting files\")\n found_files = False\n file = {}\n try:\n for file in _get_files_in_parent(\n service=service,\n parent_id=parent_id,\n field_type=field_type,\n start=start,\n end=end,\n ):\n logger.info(f\"Found file: {file['name']}, user email: {user_email}\")\n found_files = True\n yield RetrievedDriveFile(\n drive_file=file,\n user_email=user_email,\n parent_id=parent_id,\n completion_stage=DriveRetrievalStage.FOLDER_FILES,\n )\n # Only mark a folder as done if it was fully traversed without errors\n # This usually indicates that the owner of the folder was impersonated.\n # In cases where this never happens, most likely the folder owner is\n # not part of the google workspace in question (or for oauth, the authenticated\n # user doesn't own the folder)\n if found_files:\n update_traversed_ids_func(parent_id)\n except Exception as e:\n if isinstance(e, HttpError) and e.status_code == 403:\n # don't yield an error here because this is expected behavior\n # when a user doesn't have access to a folder\n logger.debug(f\"Error getting files in parent {parent_id}: {e}\")\n else:\n logger.error(f\"Error getting files in parent {parent_id}: {e}\")\n yield RetrievedDriveFile(\n drive_file=file,\n user_email=user_email,\n parent_id=parent_id,\n completion_stage=DriveRetrievalStage.FOLDER_FILES,\n error=e,\n )\n else:\n logger.info(f\"Skipping subfolder files since already traversed: {parent_id}\")\n\n for subfolder in _get_folders_in_parent(\n service=service,\n parent_id=parent_id,\n ):\n logger.info(\"Fetching all files in subfolder: \" + subfolder[\"name\"])\n yield from crawl_folders_for_files(\n service=service,\n parent_id=subfolder[\"id\"],\n field_type=field_type,\n user_email=user_email,\n traversed_parent_ids=traversed_parent_ids,\n update_traversed_ids_func=update_traversed_ids_func,\n start=start,\n end=end,\n )\n\n\ndef get_files_in_shared_drive(\n service: Resource,\n drive_id: str,\n field_type: DriveFileFieldType,\n max_num_pages: int,\n update_traversed_ids_func: Callable[[str], None] = lambda _: None,\n cache_folders: bool = True,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n page_token: str | None = None,\n) -> Iterator[GoogleDriveFileType | str]:\n kwargs = {ORDER_BY_KEY: GoogleFields.MODIFIED_TIME.value}\n if page_token:\n logger.info(f\"Using page token: {page_token}\")\n kwargs[PAGE_TOKEN_KEY] = page_token\n\n if cache_folders:\n # If we know we are going to folder crawl later, we can cache the folders here\n # Get all folders being queried and add them to the traversed set\n folder_query = f\"mimeType = '{DRIVE_FOLDER_TYPE}'\"\n folder_query += \" and trashed = false\"\n for folder in execute_paginated_retrieval(\n retrieval_function=service.files().list,\n list_key=\"files\",\n continue_on_404_or_403=True,\n corpora=\"drive\",\n driveId=drive_id,\n supportsAllDrives=True,\n includeItemsFromAllDrives=True,\n fields=\"nextPageToken, files(id)\",\n q=folder_query,\n ):\n update_traversed_ids_func(folder[\"id\"])\n\n # Get all files in the shared drive\n file_query = f\"mimeType != '{DRIVE_FOLDER_TYPE}'\"\n file_query += \" and trashed = false\"\n file_query += generate_time_range_filter(start, end)\n\n for file in execute_paginated_retrieval_with_max_pages(\n retrieval_function=service.files().list,\n max_num_pages=max_num_pages,\n list_key=\"files\",\n continue_on_404_or_403=True,\n corpora=\"drive\",\n driveId=drive_id,\n supportsAllDrives=True,\n includeItemsFromAllDrives=True,\n fields=_get_fields_for_file_type(field_type),\n q=file_query,\n **kwargs,\n ):\n # If we found any files, mark this drive as traversed. When a user has access to a drive,\n # they have access to all the files in the drive. Also not a huge deal if we re-traverse\n # empty drives.\n # NOTE: ^^ the above is not actually true due to folder restrictions:\n # https://support.google.com/a/users/answer/12380484?hl=en\n # So we may have to change this logic for people who use folder restrictions.\n update_traversed_ids_func(drive_id)\n yield file\n\n\ndef get_all_files_in_my_drive_and_shared(\n service: GoogleDriveService,\n update_traversed_ids_func: Callable,\n field_type: DriveFileFieldType,\n include_shared_with_me: bool,\n max_num_pages: int,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n cache_folders: bool = True,\n page_token: str | None = None,\n) -> Iterator[GoogleDriveFileType | str]:\n kwargs = {ORDER_BY_KEY: GoogleFields.MODIFIED_TIME.value}\n if page_token:\n logger.info(f\"Using page token: {page_token}\")\n kwargs[PAGE_TOKEN_KEY] = page_token\n\n if cache_folders:\n # If we know we are going to folder crawl later, we can cache the folders here\n # Get all folders being queried and add them to the traversed set\n folder_query = f\"mimeType = '{DRIVE_FOLDER_TYPE}'\"\n folder_query += \" and trashed = false\"\n if not include_shared_with_me:\n folder_query += \" and 'me' in owners\"\n found_folders = False\n for folder in execute_paginated_retrieval(\n retrieval_function=service.files().list,\n list_key=\"files\",\n corpora=\"user\",\n fields=_get_fields_for_file_type(field_type),\n q=folder_query,\n ):\n update_traversed_ids_func(folder[GoogleFields.ID])\n found_folders = True\n if found_folders:\n update_traversed_ids_func(get_root_folder_id(service))\n\n # Then get the files\n file_query = f\"mimeType != '{DRIVE_FOLDER_TYPE}'\"\n file_query += \" and trashed = false\"\n if not include_shared_with_me:\n file_query += \" and 'me' in owners\"\n file_query += generate_time_range_filter(start, end)\n yield from execute_paginated_retrieval_with_max_pages(\n retrieval_function=service.files().list,\n max_num_pages=max_num_pages,\n list_key=\"files\",\n continue_on_404_or_403=False,\n corpora=\"user\",\n fields=_get_fields_for_file_type(field_type),\n q=file_query,\n **kwargs,\n )\n\n\ndef get_all_files_for_oauth(\n service: GoogleDriveService,\n include_files_shared_with_me: bool,\n include_my_drives: bool,\n # One of the above 2 should be true\n include_shared_drives: bool,\n field_type: DriveFileFieldType,\n max_num_pages: int,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n page_token: str | None = None,\n) -> Iterator[GoogleDriveFileType | str]:\n kwargs = {ORDER_BY_KEY: GoogleFields.MODIFIED_TIME.value}\n if page_token:\n logger.info(f\"Using page token: {page_token}\")\n kwargs[PAGE_TOKEN_KEY] = page_token\n\n should_get_all = (\n include_shared_drives and include_my_drives and include_files_shared_with_me\n )\n corpora = \"allDrives\" if should_get_all else \"user\"\n\n file_query = f\"mimeType != '{DRIVE_FOLDER_TYPE}'\"\n file_query += \" and trashed = false\"\n file_query += generate_time_range_filter(start, end)\n\n if not should_get_all:\n if include_files_shared_with_me and not include_my_drives:\n file_query += \" and not 'me' in owners\"\n if not include_files_shared_with_me and include_my_drives:\n file_query += \" and 'me' in owners\"\n\n yield from execute_paginated_retrieval_with_max_pages(\n max_num_pages=max_num_pages,\n retrieval_function=service.files().list,\n list_key=\"files\",\n continue_on_404_or_403=False,\n corpora=corpora,\n includeItemsFromAllDrives=should_get_all,\n supportsAllDrives=should_get_all,\n fields=_get_fields_for_file_type(field_type),\n q=file_query,\n **kwargs,\n )\n\n\n# Just in case we need to get the root folder id\ndef get_root_folder_id(service: Resource) -> str:\n # we dont paginate here because there is only one root folder per user\n # https://developers.google.com/drive/api/guides/v2-to-v3-reference\n return (\n service.files()\n .get(fileId=\"root\", fields=GoogleFields.ID.value)\n .execute()[GoogleFields.ID.value]\n )\n\n\ndef _extract_file_id_from_web_view_link(web_view_link: str) -> str:\n parsed = urlparse(web_view_link)\n path_parts = [part for part in parsed.path.split(\"/\") if part]\n\n if \"d\" in path_parts:\n idx = path_parts.index(\"d\")\n if idx + 1 < len(path_parts):\n return path_parts[idx + 1]\n\n query_params = parse_qs(parsed.query)\n for key in (\"id\", \"fileId\"):\n value = query_params.get(key)\n if value and value[0]:\n return value[0]\n\n raise ValueError(\n f\"Unable to extract Drive file id from webViewLink: {web_view_link}\"\n )\n\n\ndef get_file_by_web_view_link(\n service: GoogleDriveService,\n web_view_link: str,\n fields: str,\n) -> GoogleDriveFileType:\n \"\"\"Retrieve a Google Drive file using its webViewLink.\"\"\"\n file_id = _extract_file_id_from_web_view_link(web_view_link)\n return (\n service.files()\n .get(\n fileId=file_id,\n supportsAllDrives=True,\n fields=fields,\n )\n .execute()\n )\n" }, { "path": "backend/onyx/connectors/google_drive/models.py", "content": "from enum import Enum\nfrom typing import Any\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import Field\nfrom pydantic import field_serializer\nfrom pydantic import field_validator\n\nfrom onyx.connectors.interfaces import ConnectorCheckpoint\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.utils.threadpool_concurrency import ThreadSafeDict\nfrom onyx.utils.threadpool_concurrency import ThreadSafeSet\n\n\nclass GDriveMimeType(str, Enum):\n DOC = \"application/vnd.google-apps.document\"\n SPREADSHEET = \"application/vnd.google-apps.spreadsheet\"\n SPREADSHEET_OPEN_FORMAT = (\n \"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet\"\n )\n SPREADSHEET_MS_EXCEL = \"application/vnd.ms-excel\"\n PDF = \"application/pdf\"\n WORD_DOC = \"application/vnd.openxmlformats-officedocument.wordprocessingml.document\"\n PPT = \"application/vnd.google-apps.presentation\"\n POWERPOINT = (\n \"application/vnd.openxmlformats-officedocument.presentationml.presentation\"\n )\n PLAIN_TEXT = \"text/plain\"\n MARKDOWN = \"text/markdown\"\n\n\nGoogleDriveFileType = dict[str, Any]\n\n\nTOKEN_EXPIRATION_TIME = 3600 # 1 hour\n\n\n# These correspond to The major stages of retrieval for google drive.\n# The stages for the oauth flow are:\n# get_all_files_for_oauth(),\n# get_all_drive_ids(),\n# get_files_in_shared_drive(),\n# crawl_folders_for_files()\n#\n# The stages for the service account flow are roughly:\n# get_all_user_emails(),\n# get_all_drive_ids(),\n# get_files_in_shared_drive(),\n# Then for each user:\n# get_files_in_my_drive()\n# get_files_in_shared_drive()\n# crawl_folders_for_files()\nclass DriveRetrievalStage(str, Enum):\n START = \"start\"\n DONE = \"done\"\n # OAuth specific stages\n OAUTH_FILES = \"oauth_files\"\n\n # Service account specific stages\n USER_EMAILS = \"user_emails\"\n MY_DRIVE_FILES = \"my_drive_files\"\n\n # Used for both oauth and service account flows\n DRIVE_IDS = \"drive_ids\"\n SHARED_DRIVE_FILES = \"shared_drive_files\"\n FOLDER_FILES = \"folder_files\"\n\n\nclass StageCompletion(BaseModel):\n \"\"\"\n Describes the point in the retrieval+indexing process that the\n connector is at. completed_until is the timestamp of the latest\n file that has been retrieved or error that has been yielded.\n Optional fields are used for retrieval stages that need more information\n for resuming than just the timestamp of the latest file.\n \"\"\"\n\n stage: DriveRetrievalStage\n completed_until: SecondsSinceUnixEpoch\n current_folder_or_drive_id: str | None = None\n next_page_token: str | None = None\n\n # only used for shared drives\n processed_drive_ids: set[str] = set()\n\n def update(\n self,\n stage: DriveRetrievalStage,\n completed_until: SecondsSinceUnixEpoch,\n current_folder_or_drive_id: str | None = None,\n ) -> None:\n self.stage = stage\n self.completed_until = completed_until\n self.current_folder_or_drive_id = current_folder_or_drive_id\n\n\nclass RetrievedDriveFile(BaseModel):\n \"\"\"\n Describes a file that has been retrieved from google drive.\n user_email is the email of the user that the file was retrieved\n by impersonating. If an error worthy of being reported is encountered,\n error should be set and later propagated as a ConnectorFailure.\n \"\"\"\n\n # The stage at which this file was retrieved\n completion_stage: DriveRetrievalStage\n\n # The file that was retrieved\n drive_file: GoogleDriveFileType\n\n # The email of the user that the file was retrieved by impersonating\n user_email: str\n\n # The id of the parent folder or drive of the file\n parent_id: str | None = None\n\n # Any unexpected error that occurred while retrieving the file.\n # In particular, this is not used for 403/404 errors, which are expected\n # in the context of impersonating all the users to try to retrieve all\n # files from all their Drives and Folders.\n error: Exception | None = None\n\n model_config = ConfigDict(arbitrary_types_allowed=True)\n\n\nclass GoogleDriveCheckpoint(ConnectorCheckpoint):\n # Checkpoint version of _retrieved_ids\n retrieved_folder_and_drive_ids: set[str]\n\n # Describes the point in the retrieval+indexing process that the\n # checkpoint is at. when this is set to a given stage, the connector\n # has finished yielding all values from the previous stage.\n completion_stage: DriveRetrievalStage\n\n # The latest timestamp of a file that has been retrieved per user email.\n # StageCompletion is used to track the completion of each stage, but the\n # timestamp part is not used for folder crawling.\n completion_map: ThreadSafeDict[str, StageCompletion]\n\n # all file ids that have been retrieved\n all_retrieved_file_ids: set[str] = set()\n\n # cached version of the drive and folder ids to retrieve\n drive_ids_to_retrieve: list[str] | None = None\n folder_ids_to_retrieve: list[str] | None = None\n\n # cached user emails\n user_emails: list[str] | None = None\n\n # Hierarchy node raw IDs that have already been yielded.\n # Used to avoid yielding duplicate hierarchy nodes across checkpoints.\n # Thread-safe because multiple impersonation threads access this concurrently.\n # Uses default_factory to ensure each checkpoint instance gets a fresh set.\n seen_hierarchy_node_raw_ids: ThreadSafeSet[str] = Field(\n default_factory=ThreadSafeSet\n )\n\n # Hierarchy node raw IDs where we have successfully walked up to a terminal\n # node (a drive root with no parent). This is separate from seen_hierarchy_node_raw_ids\n # because a node might be yielded before we've walked its full ancestry chain.\n # We only skip walking from a node if it's in this set, ensuring that if one user\n # fails to walk to the root, another user with better access can still complete the walk.\n # Thread-safe because multiple impersonation threads access this concurrently.\n # Uses default_factory to ensure each checkpoint instance gets a fresh set.\n fully_walked_hierarchy_node_raw_ids: ThreadSafeSet[str] = Field(\n default_factory=ThreadSafeSet\n )\n\n @field_serializer(\"completion_map\")\n def serialize_completion_map(\n self, completion_map: ThreadSafeDict[str, StageCompletion], _info: Any\n ) -> dict[str, StageCompletion]:\n return completion_map._dict\n\n @field_serializer(\"seen_hierarchy_node_raw_ids\")\n def serialize_seen_hierarchy(\n self, seen_hierarchy_node_raw_ids: ThreadSafeSet[str], _info: Any\n ) -> set[str]:\n return seen_hierarchy_node_raw_ids.copy()\n\n @field_serializer(\"fully_walked_hierarchy_node_raw_ids\")\n def serialize_fully_walked_hierarchy(\n self, fully_walked_hierarchy_node_raw_ids: ThreadSafeSet[str], _info: Any\n ) -> set[str]:\n return fully_walked_hierarchy_node_raw_ids.copy()\n\n @field_validator(\"completion_map\", mode=\"before\")\n def validate_completion_map(cls, v: Any) -> ThreadSafeDict[str, StageCompletion]:\n assert isinstance(v, dict) or isinstance(v, ThreadSafeDict)\n return ThreadSafeDict(\n {k: StageCompletion.model_validate(val) for k, val in v.items()}\n )\n\n @field_validator(\"seen_hierarchy_node_raw_ids\", mode=\"before\")\n def validate_seen_hierarchy(cls, v: Any) -> ThreadSafeSet[str]:\n if isinstance(v, ThreadSafeSet):\n return v\n if isinstance(v, set):\n return ThreadSafeSet(v)\n if isinstance(v, list):\n return ThreadSafeSet(set(v))\n return ThreadSafeSet()\n\n @field_validator(\"fully_walked_hierarchy_node_raw_ids\", mode=\"before\")\n def validate_fully_walked_hierarchy(cls, v: Any) -> ThreadSafeSet[str]:\n if isinstance(v, ThreadSafeSet):\n return v\n if isinstance(v, set):\n return ThreadSafeSet(v)\n if isinstance(v, list):\n return ThreadSafeSet(set(v))\n return ThreadSafeSet()\n" }, { "path": "backend/onyx/connectors/google_drive/section_extraction.py", "content": "from typing import Any\n\nfrom pydantic import BaseModel\n\nfrom onyx.connectors.google_utils.resources import GoogleDocsService\nfrom onyx.connectors.models import TextSection\n\nHEADING_DELIMITER = \"\\n\"\n\n\nclass CurrentHeading(BaseModel):\n id: str | None\n text: str\n\n\ndef _build_gdoc_section_link(doc_id: str, tab_id: str, heading_id: str | None) -> str:\n \"\"\"Builds a Google Doc link that jumps to a specific heading\"\"\"\n # NOTE: doesn't support docs with multiple tabs atm, if we need that ask\n # @Chris\n heading_str = f\"#heading={heading_id}\" if heading_id else \"\"\n return f\"https://docs.google.com/document/d/{doc_id}/edit?tab={tab_id}{heading_str}\"\n\n\ndef _extract_id_from_heading(paragraph: dict[str, Any]) -> str:\n \"\"\"Extracts the id from a heading paragraph element\"\"\"\n return paragraph[\"paragraphStyle\"][\"headingId\"]\n\n\ndef _extract_text_from_paragraph(paragraph: dict[str, Any]) -> str:\n \"\"\"Extracts the text content from a paragraph element\"\"\"\n text_elements = []\n for element in paragraph.get(\"elements\", []):\n if \"textRun\" in element:\n text_elements.append(element[\"textRun\"].get(\"content\", \"\"))\n\n # Handle links\n if \"textStyle\" in element and \"link\" in element[\"textStyle\"]:\n text_elements.append(f\"({element['textStyle']['link'].get('url', '')})\")\n\n if \"person\" in element:\n name = element[\"person\"].get(\"personProperties\", {}).get(\"name\", \"\")\n email = element[\"person\"].get(\"personProperties\", {}).get(\"email\", \"\")\n person_str = \"\"\n text_elements.append(person_str)\n\n if \"richLink\" in element:\n props = element[\"richLink\"].get(\"richLinkProperties\", {})\n title = props.get(\"title\", \"\")\n uri = props.get(\"uri\", \"\")\n link_str = f\"[{title}]({uri})\"\n text_elements.append(link_str)\n\n return \"\".join(text_elements)\n\n\ndef _extract_text_from_table(table: dict[str, Any]) -> str:\n \"\"\"\n Extracts the text content from a table element.\n \"\"\"\n row_strs = []\n\n for row in table.get(\"tableRows\", []):\n cells = row.get(\"tableCells\", [])\n cell_strs = []\n for cell in cells:\n child_elements = cell.get(\"content\", {})\n cell_str = []\n for child_elem in child_elements:\n if \"paragraph\" not in child_elem:\n continue\n cell_str.append(_extract_text_from_paragraph(child_elem[\"paragraph\"]))\n cell_strs.append(\"\".join(cell_str))\n row_strs.append(\", \".join(cell_strs))\n return \"\\n\".join(row_strs)\n\n\ndef get_document_sections(\n docs_service: GoogleDocsService,\n doc_id: str,\n) -> list[TextSection]:\n \"\"\"Extracts sections from a Google Doc, including their headings and content\"\"\"\n # Fetch the document structure\n http_request = docs_service.documents().get(documentId=doc_id)\n\n # Google has poor support for tabs in the docs api, see\n # https://cloud.google.com/python/docs/reference/cloudtasks/\n # latest/google.cloud.tasks_v2.types.HttpRequest\n # https://developers.google.com/workspace/docs/api/how-tos/tabs\n # https://developers.google.com/workspace/docs/api/reference/rest/v1/documents/get\n # this is a hack to use the param mentioned in the rest api docs\n # TODO: check if it can be specified i.e. in documents()\n http_request.uri += \"&includeTabsContent=true\"\n doc = http_request.execute()\n\n # Get the content\n tabs = doc.get(\"tabs\", {})\n sections: list[TextSection] = []\n for tab in tabs:\n sections.extend(get_tab_sections(tab, doc_id))\n return sections\n\n\ndef _is_heading(paragraph: dict[str, Any]) -> bool:\n \"\"\"Checks if a paragraph (a block of text in a drive document) is a heading\"\"\"\n if not (\n \"paragraphStyle\" in paragraph\n and \"namedStyleType\" in paragraph[\"paragraphStyle\"]\n ):\n return False\n\n style = paragraph[\"paragraphStyle\"][\"namedStyleType\"]\n is_heading = style.startswith(\"HEADING_\")\n is_title = style.startswith(\"TITLE\")\n return is_heading or is_title\n\n\ndef _add_finished_section(\n sections: list[TextSection],\n doc_id: str,\n tab_id: str,\n current_heading: CurrentHeading,\n current_section: list[str],\n) -> None:\n \"\"\"Adds a finished section to the list of sections if the section has content.\n Returns the list of sections to use going forward, which may be the old list\n if a new section was not added.\n \"\"\"\n if not (current_section or current_heading.text):\n return\n # If we were building a previous section, add it to sections list\n\n # this is unlikely to ever matter, but helps if the doc contains weird headings\n header_text = current_heading.text.replace(HEADING_DELIMITER, \"\")\n section_text = f\"{header_text}{HEADING_DELIMITER}\" + \"\\n\".join(current_section)\n sections.append(\n TextSection(\n text=section_text.strip(),\n link=_build_gdoc_section_link(doc_id, tab_id, current_heading.id),\n )\n )\n\n\ndef get_tab_sections(tab: dict[str, Any], doc_id: str) -> list[TextSection]:\n tab_id = tab[\"tabProperties\"][\"tabId\"]\n content = tab.get(\"documentTab\", {}).get(\"body\", {}).get(\"content\", [])\n\n sections: list[TextSection] = []\n current_section: list[str] = []\n current_heading = CurrentHeading(id=None, text=\"\")\n\n for element in content:\n if \"paragraph\" in element:\n paragraph = element[\"paragraph\"]\n\n # If this is not a heading, add content to current section\n if not _is_heading(paragraph):\n text = _extract_text_from_paragraph(paragraph)\n if text.strip():\n current_section.append(text)\n continue\n\n _add_finished_section(\n sections, doc_id, tab_id, current_heading, current_section\n )\n\n current_section = []\n\n # Start new heading\n heading_id = _extract_id_from_heading(paragraph)\n heading_text = _extract_text_from_paragraph(paragraph)\n current_heading = CurrentHeading(\n id=heading_id,\n text=heading_text,\n )\n elif \"table\" in element:\n text = _extract_text_from_table(element[\"table\"])\n if text.strip():\n current_section.append(text)\n\n # Don't forget to add the last section\n _add_finished_section(sections, doc_id, tab_id, current_heading, current_section)\n\n return sections\n" }, { "path": "backend/onyx/connectors/google_site/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/google_site/connector.py", "content": "import os\nimport re\nfrom typing import Any\nfrom typing import cast\n\nfrom bs4 import BeautifulSoup\nfrom bs4 import Tag\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import load_files_from_zip\nfrom onyx.file_processing.extract_file_text import read_text_file\nfrom onyx.file_processing.html_utils import web_html_cleanup\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef a_tag_text_to_path(atag: Tag) -> str:\n page_path = atag.text.strip().lower()\n page_path = re.sub(r\"[^a-zA-Z0-9\\s]\", \"\", page_path)\n page_path = \"-\".join(page_path.split())\n\n return page_path\n\n\ndef find_google_sites_page_path_from_navbar(\n element: BeautifulSoup | Tag, path: str, depth: int\n) -> str | None:\n lis = cast(\n list[Tag],\n element.find_all(\"li\", attrs={\"data-nav-level\": f\"{depth}\"}),\n )\n for li in lis:\n a = cast(Tag, li.find(\"a\"))\n if a.get(\"aria-selected\") == \"true\":\n return f\"{path}/{a_tag_text_to_path(a)}\"\n elif a.get(\"aria-expanded\") == \"true\":\n sub_path = find_google_sites_page_path_from_navbar(\n element, f\"{path}/{a_tag_text_to_path(a)}\", depth + 1\n )\n if sub_path:\n return sub_path\n\n return None\n\n\nclass GoogleSitesConnector(LoadConnector):\n def __init__(\n self,\n zip_path: str,\n base_url: str,\n batch_size: int = INDEX_BATCH_SIZE,\n ):\n self.zip_path = zip_path\n self.base_url = base_url\n self.batch_size = batch_size\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n pass\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n documents: list[Document | HierarchyNode] = []\n\n file_content_io = get_default_file_store().read_file(self.zip_path, mode=\"b\")\n\n # load the HTML files\n files = load_files_from_zip(file_content_io)\n count = 0\n for file_info, file_io in files:\n # skip non-published files\n if \"/PUBLISHED/\" not in file_info.filename:\n continue\n\n file_path, extension = os.path.splitext(file_info.filename)\n if extension != \".html\":\n continue\n\n file_content, _ = read_text_file(file_io)\n soup = BeautifulSoup(file_content, \"html.parser\")\n\n # get the link out of the navbar\n header = cast(Tag, soup.find(\"header\"))\n nav = cast(Tag, header.find(\"nav\"))\n path = find_google_sites_page_path_from_navbar(nav, \"\", 1)\n if not path:\n count += 1\n logger.error(\n f\"Could not find path for '{file_info.filename}'. \"\n + \"This page will not have a working link.\\n\\n\"\n + f\"# of broken links so far - {count}\"\n )\n logger.info(f\"Path to page: {path}\")\n # cleanup the hidden `Skip to main content` and `Skip to navigation` that\n # appears at the top of every page\n for div in soup.find_all(\"div\", attrs={\"data-is-touch-wrapper\": \"true\"}):\n div.extract()\n\n # get the body of the page\n parsed_html = web_html_cleanup(\n soup, additional_element_types_to_discard=[\"header\", \"nav\"]\n )\n\n title = parsed_html.title or file_path.split(\"/\")[-1]\n documents.append(\n Document(\n id=f\"{DocumentSource.GOOGLE_SITES.value}:{path}\",\n source=DocumentSource.GOOGLE_SITES,\n semantic_identifier=title,\n sections=[\n TextSection(\n link=(\n (self.base_url.rstrip(\"/\") + \"/\" + path.lstrip(\"/\"))\n if path\n else \"\"\n ),\n text=parsed_html.cleaned_text,\n )\n ],\n metadata={},\n )\n )\n\n if len(documents) >= self.batch_size:\n yield documents\n documents = []\n\n if documents:\n yield documents\n\n\nif __name__ == \"__main__\":\n connector = GoogleSitesConnector(\n os.environ[\"GOOGLE_SITES_ZIP_PATH\"],\n os.environ.get(\"GOOGLE_SITES_BASE_URL\", \"\"),\n )\n for doc_batch in connector.load_from_state():\n for doc in doc_batch:\n print(doc)\n" }, { "path": "backend/onyx/connectors/google_utils/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/google_utils/google_auth.py", "content": "import json\nfrom typing import Any\n\nfrom google.auth.transport.requests import Request\nfrom google.oauth2.credentials import Credentials as OAuthCredentials\nfrom google.oauth2.service_account import Credentials as ServiceAccountCredentials\n\nfrom onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_SECRET\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_AUTHENTICATION_METHOD,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_TOKEN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GOOGLE_SCOPES,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GoogleOAuthAuthenticationMethod,\n)\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef sanitize_oauth_credentials(oauth_creds: OAuthCredentials) -> str:\n \"\"\"we really don't want to be persisting the client id and secret anywhere but the\n environment.\n\n Returns a string of serialized json.\n \"\"\"\n\n # strip the client id and secret\n oauth_creds_json_str = oauth_creds.to_json()\n oauth_creds_sanitized_json: dict[str, Any] = json.loads(oauth_creds_json_str)\n oauth_creds_sanitized_json.pop(\"client_id\", None)\n oauth_creds_sanitized_json.pop(\"client_secret\", None)\n oauth_creds_sanitized_json_str = json.dumps(oauth_creds_sanitized_json)\n return oauth_creds_sanitized_json_str\n\n\ndef get_google_oauth_creds(\n token_json_str: str, source: DocumentSource\n) -> OAuthCredentials | None:\n \"\"\"creds_json only needs to contain client_id, client_secret and refresh_token to\n refresh the creds.\n\n expiry and token are optional ... however, if passing in expiry, token\n should also be passed in or else we may not return any creds.\n (probably a sign we should refactor the function)\n \"\"\"\n creds_json = json.loads(token_json_str)\n creds = OAuthCredentials.from_authorized_user_info(\n info=creds_json,\n scopes=GOOGLE_SCOPES[source],\n )\n if creds.valid:\n return creds\n\n if creds.expired and creds.refresh_token:\n try:\n creds.refresh(Request())\n if creds.valid:\n logger.notice(\"Refreshed Google Drive tokens.\")\n return creds\n except Exception:\n logger.exception(\"Failed to refresh google drive access token\")\n return None\n\n return None\n\n\ndef get_google_creds(\n credentials: dict[str, str],\n source: DocumentSource,\n) -> tuple[ServiceAccountCredentials | OAuthCredentials, dict[str, str] | None]:\n \"\"\"Checks for two different types of credentials.\n (1) A credential which holds a token acquired via a user going through\n the Google OAuth flow.\n (2) A credential which holds a service account key JSON file, which\n can then be used to impersonate any user in the workspace.\n\n Return a tuple where:\n The first element is the requested credentials\n The second element is a new credentials dict that the caller should write back\n to the db. This happens if token rotation occurs while loading credentials.\n \"\"\"\n oauth_creds = None\n service_creds = None\n new_creds_dict = None\n if DB_CREDENTIALS_DICT_TOKEN_KEY in credentials:\n # OAUTH\n authentication_method: str = credentials.get(\n DB_CREDENTIALS_AUTHENTICATION_METHOD,\n GoogleOAuthAuthenticationMethod.UPLOADED.value,\n )\n\n credentials_dict_str = credentials[DB_CREDENTIALS_DICT_TOKEN_KEY]\n credentials_dict = json.loads(credentials_dict_str)\n\n # only send what get_google_oauth_creds needs\n authorized_user_info = {}\n\n # oauth_interactive is sanitized and needs credentials from the environment\n if (\n authentication_method\n == GoogleOAuthAuthenticationMethod.OAUTH_INTERACTIVE.value\n ):\n authorized_user_info[\"client_id\"] = OAUTH_GOOGLE_DRIVE_CLIENT_ID\n authorized_user_info[\"client_secret\"] = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET\n else:\n authorized_user_info[\"client_id\"] = credentials_dict[\"client_id\"]\n authorized_user_info[\"client_secret\"] = credentials_dict[\"client_secret\"]\n\n authorized_user_info[\"refresh_token\"] = credentials_dict[\"refresh_token\"]\n\n authorized_user_info[\"token\"] = credentials_dict[\"token\"]\n authorized_user_info[\"expiry\"] = credentials_dict[\"expiry\"]\n\n token_json_str = json.dumps(authorized_user_info)\n oauth_creds = get_google_oauth_creds(\n token_json_str=token_json_str, source=source\n )\n\n # tell caller to update token stored in DB if the refresh token changed\n if oauth_creds:\n if oauth_creds.refresh_token != authorized_user_info[\"refresh_token\"]:\n # if oauth_interactive, sanitize the credentials so they don't get stored in the db\n if (\n authentication_method\n == GoogleOAuthAuthenticationMethod.OAUTH_INTERACTIVE.value\n ):\n oauth_creds_json_str = sanitize_oauth_credentials(oauth_creds)\n else:\n oauth_creds_json_str = oauth_creds.to_json()\n\n new_creds_dict = {\n DB_CREDENTIALS_DICT_TOKEN_KEY: oauth_creds_json_str,\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: credentials[\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY\n ],\n DB_CREDENTIALS_AUTHENTICATION_METHOD: authentication_method,\n }\n elif DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY in credentials:\n # SERVICE ACCOUNT\n service_account_key_json_str = credentials[\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY\n ]\n service_account_key = json.loads(service_account_key_json_str)\n\n service_creds = ServiceAccountCredentials.from_service_account_info(\n service_account_key, scopes=GOOGLE_SCOPES[source]\n )\n\n if not service_creds.valid or not service_creds.expired:\n service_creds.refresh(Request())\n\n if not service_creds.valid:\n raise PermissionError(\n f\"Unable to access {source} - service account credentials are invalid.\"\n )\n\n creds: ServiceAccountCredentials | OAuthCredentials | None = (\n oauth_creds or service_creds\n )\n if creds is None:\n raise PermissionError(\n f\"Unable to access {source} - unknown credential structure.\"\n )\n\n return creds, new_creds_dict\n" }, { "path": "backend/onyx/connectors/google_utils/google_kv.py", "content": "import json\nfrom typing import cast\nfrom urllib.parse import parse_qs\nfrom urllib.parse import ParseResult\nfrom urllib.parse import urlparse\n\nfrom google.oauth2.credentials import Credentials as OAuthCredentials\nfrom google_auth_oauthlib.flow import InstalledAppFlow # type: ignore\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import KV_CRED_KEY\nfrom onyx.configs.constants import KV_GMAIL_CRED_KEY\nfrom onyx.configs.constants import KV_GMAIL_SERVICE_ACCOUNT_KEY\nfrom onyx.configs.constants import KV_GOOGLE_DRIVE_CRED_KEY\nfrom onyx.configs.constants import KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY\nfrom onyx.connectors.google_utils.resources import get_drive_service\nfrom onyx.connectors.google_utils.resources import get_gmail_service\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_AUTHENTICATION_METHOD,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_TOKEN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GOOGLE_SCOPES,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GoogleOAuthAuthenticationMethod,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n MISSING_SCOPES_ERROR_STR,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n ONYX_SCOPE_INSTRUCTIONS,\n)\nfrom onyx.db.credentials import update_credential_json\nfrom onyx.db.models import User\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.key_value_store.interface import unwrap_str\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.server.documents.models import GoogleAppCredentials\nfrom onyx.server.documents.models import GoogleServiceAccountKey\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _build_frontend_google_drive_redirect(source: DocumentSource) -> str:\n if source == DocumentSource.GOOGLE_DRIVE:\n return f\"{WEB_DOMAIN}/admin/connectors/google-drive/auth/callback\"\n elif source == DocumentSource.GMAIL:\n return f\"{WEB_DOMAIN}/admin/connectors/gmail/auth/callback\"\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n\n\ndef _get_current_oauth_user(creds: OAuthCredentials, source: DocumentSource) -> str:\n if source == DocumentSource.GOOGLE_DRIVE:\n drive_service = get_drive_service(creds)\n user_info = (\n drive_service.about()\n .get(\n fields=\"user(emailAddress)\",\n )\n .execute()\n )\n email = user_info.get(\"user\", {}).get(\"emailAddress\")\n elif source == DocumentSource.GMAIL:\n gmail_service = get_gmail_service(creds)\n user_info = (\n gmail_service.users()\n .getProfile(\n userId=\"me\",\n fields=\"emailAddress\",\n )\n .execute()\n )\n email = user_info.get(\"emailAddress\")\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n return email\n\n\ndef verify_csrf(credential_id: int, state: str) -> None:\n csrf = unwrap_str(get_kv_store().load(KV_CRED_KEY.format(str(credential_id))))\n if csrf != state:\n raise PermissionError(\n \"State from Google Drive Connector callback does not match expected\"\n )\n\n\ndef update_credential_access_tokens(\n auth_code: str,\n credential_id: int,\n user: User,\n db_session: Session,\n source: DocumentSource,\n auth_method: GoogleOAuthAuthenticationMethod,\n) -> OAuthCredentials | None:\n app_credentials = get_google_app_cred(source)\n flow = InstalledAppFlow.from_client_config(\n app_credentials.model_dump(),\n scopes=GOOGLE_SCOPES[source],\n redirect_uri=_build_frontend_google_drive_redirect(source),\n )\n flow.fetch_token(code=auth_code)\n creds = flow.credentials\n token_json_str = creds.to_json()\n\n # Get user email from Google API so we know who\n # the primary admin is for this connector\n try:\n email = _get_current_oauth_user(creds, source)\n except Exception as e:\n if MISSING_SCOPES_ERROR_STR in str(e):\n raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e\n raise e\n\n new_creds_dict = {\n DB_CREDENTIALS_DICT_TOKEN_KEY: token_json_str,\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: email,\n DB_CREDENTIALS_AUTHENTICATION_METHOD: auth_method.value,\n }\n\n if not update_credential_json(credential_id, new_creds_dict, user, db_session):\n return None\n return creds\n\n\ndef build_service_account_creds(\n source: DocumentSource,\n primary_admin_email: str | None = None,\n name: str | None = None,\n) -> CredentialBase:\n service_account_key = get_service_account_key(source=source)\n\n credential_dict = {\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: service_account_key.json(),\n }\n if primary_admin_email:\n credential_dict[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] = primary_admin_email\n\n credential_dict[DB_CREDENTIALS_AUTHENTICATION_METHOD] = (\n GoogleOAuthAuthenticationMethod.UPLOADED.value\n )\n\n return CredentialBase(\n credential_json=credential_dict,\n admin_public=True,\n source=source,\n name=name,\n )\n\n\ndef get_auth_url(credential_id: int, source: DocumentSource) -> str:\n if source == DocumentSource.GOOGLE_DRIVE:\n creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))\n elif source == DocumentSource.GMAIL:\n creds_str = str(get_kv_store().load(KV_GMAIL_CRED_KEY))\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n credential_json = json.loads(creds_str)\n flow = InstalledAppFlow.from_client_config(\n credential_json,\n scopes=GOOGLE_SCOPES[source],\n redirect_uri=_build_frontend_google_drive_redirect(source),\n )\n auth_url, _ = flow.authorization_url(prompt=\"consent\")\n\n parsed_url = cast(ParseResult, urlparse(auth_url))\n params = parse_qs(parsed_url.query)\n\n get_kv_store().store(\n KV_CRED_KEY.format(credential_id),\n {\"value\": params.get(\"state\", [None])[0]},\n encrypt=True,\n )\n return str(auth_url)\n\n\ndef get_google_app_cred(source: DocumentSource) -> GoogleAppCredentials:\n if source == DocumentSource.GOOGLE_DRIVE:\n creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))\n elif source == DocumentSource.GMAIL:\n creds_str = str(get_kv_store().load(KV_GMAIL_CRED_KEY))\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n return GoogleAppCredentials(**json.loads(creds_str))\n\n\ndef upsert_google_app_cred(\n app_credentials: GoogleAppCredentials, source: DocumentSource\n) -> None:\n if source == DocumentSource.GOOGLE_DRIVE:\n get_kv_store().store(\n KV_GOOGLE_DRIVE_CRED_KEY, app_credentials.json(), encrypt=True\n )\n elif source == DocumentSource.GMAIL:\n get_kv_store().store(KV_GMAIL_CRED_KEY, app_credentials.json(), encrypt=True)\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n\n\ndef delete_google_app_cred(source: DocumentSource) -> None:\n if source == DocumentSource.GOOGLE_DRIVE:\n get_kv_store().delete(KV_GOOGLE_DRIVE_CRED_KEY)\n elif source == DocumentSource.GMAIL:\n get_kv_store().delete(KV_GMAIL_CRED_KEY)\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n\n\ndef get_service_account_key(source: DocumentSource) -> GoogleServiceAccountKey:\n if source == DocumentSource.GOOGLE_DRIVE:\n creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY))\n elif source == DocumentSource.GMAIL:\n creds_str = str(get_kv_store().load(KV_GMAIL_SERVICE_ACCOUNT_KEY))\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n return GoogleServiceAccountKey(**json.loads(creds_str))\n\n\ndef upsert_service_account_key(\n service_account_key: GoogleServiceAccountKey, source: DocumentSource\n) -> None:\n if source == DocumentSource.GOOGLE_DRIVE:\n get_kv_store().store(\n KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY,\n service_account_key.json(),\n encrypt=True,\n )\n elif source == DocumentSource.GMAIL:\n get_kv_store().store(\n KV_GMAIL_SERVICE_ACCOUNT_KEY, service_account_key.json(), encrypt=True\n )\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n\n\ndef delete_service_account_key(source: DocumentSource) -> None:\n if source == DocumentSource.GOOGLE_DRIVE:\n get_kv_store().delete(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY)\n elif source == DocumentSource.GMAIL:\n get_kv_store().delete(KV_GMAIL_SERVICE_ACCOUNT_KEY)\n else:\n raise ValueError(f\"Unsupported source: {source}\")\n" }, { "path": "backend/onyx/connectors/google_utils/google_utils.py", "content": "import re\nimport socket\nimport time\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom enum import Enum\nfrom typing import Any\n\nfrom googleapiclient.errors import HttpError # type: ignore\n\nfrom onyx.connectors.google_drive.models import GoogleDriveFileType\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\nlogger = setup_logger()\n\n_RATE_LIMIT_REASONS = {\"userRateLimitExceeded\", \"rateLimitExceeded\"}\n\n\ndef _is_rate_limit_error(error: HttpError) -> bool:\n \"\"\"Google sometimes returns rate-limit errors as 403 with reason\n 'userRateLimitExceeded' instead of 429. This helper detects both.\"\"\"\n if error.resp.status == 429:\n return True\n if error.resp.status != 403:\n return False\n error_details = getattr(error, \"error_details\", None) or []\n for detail in error_details:\n if isinstance(detail, dict) and detail.get(\"reason\") in _RATE_LIMIT_REASONS:\n return True\n return \"userRateLimitExceeded\" in str(error) or \"rateLimitExceeded\" in str(error)\n\n\n# Google Drive APIs are quite flakey and may 500 for an\n# extended period of time. This is now addressed by checkpointing.\n#\n# NOTE: We previously tried to combat this here by adding a very\n# long retry period (~20 minutes of trying, one request a minute.)\n# This is no longer necessary due to checkpointing.\nadd_retries = retry_builder(tries=5, max_delay=10)\n\nNEXT_PAGE_TOKEN_KEY = \"nextPageToken\"\nPAGE_TOKEN_KEY = \"pageToken\"\nORDER_BY_KEY = \"orderBy\"\n\n\n# See https://developers.google.com/drive/api/reference/rest/v3/files/list for more\nclass GoogleFields(str, Enum):\n ID = \"id\"\n CREATED_TIME = \"createdTime\"\n MODIFIED_TIME = \"modifiedTime\"\n NAME = \"name\"\n SIZE = \"size\"\n PARENTS = \"parents\"\n\n\ndef _execute_with_retry(request: Any) -> Any:\n max_attempts = 6\n attempt = 1\n\n while attempt < max_attempts:\n # Note for reasons unknown, the Google API will sometimes return a 429\n # and even after waiting the retry period, it will return another 429.\n # It could be due to a few possibilities:\n # 1. Other things are also requesting from the Drive/Gmail API with the same key\n # 2. It's a rolling rate limit so the moment we get some amount of requests cleared, we hit it again very quickly\n # 3. The retry-after has a maximum and we've already hit the limit for the day\n # or it's something else...\n try:\n return request.execute()\n except HttpError as error:\n attempt += 1\n\n if _is_rate_limit_error(error):\n # Attempt to get 'Retry-After' from headers\n retry_after = error.resp.get(\"Retry-After\")\n if retry_after:\n sleep_time = int(retry_after)\n else:\n # Extract 'Retry after' timestamp from error message\n match = re.search(\n r\"Retry after (\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}\\.\\d+Z)\",\n str(error),\n )\n if match:\n retry_after_timestamp = match.group(1)\n retry_after_dt = datetime.strptime(\n retry_after_timestamp, \"%Y-%m-%dT%H:%M:%S.%fZ\"\n ).replace(tzinfo=timezone.utc)\n current_time = datetime.now(timezone.utc)\n sleep_time = max(\n int((retry_after_dt - current_time).total_seconds()),\n 0,\n )\n else:\n logger.error(\n f\"No Retry-After header or timestamp found in error message: {error}\"\n )\n sleep_time = 60\n\n sleep_time += 3 # Add a buffer to be safe\n\n logger.info(\n f\"Rate limit exceeded. Attempt {attempt}/{max_attempts}. Sleeping for {sleep_time} seconds.\"\n )\n time.sleep(sleep_time)\n\n else:\n raise\n\n # If we've exhausted all attempts\n raise Exception(f\"Failed to execute request after {max_attempts} attempts\")\n\n\ndef get_file_owners(file: GoogleDriveFileType, primary_admin_email: str) -> list[str]:\n \"\"\"\n Get the owners of a file if the attribute is present.\n \"\"\"\n return [\n email\n for owner in file.get(\"owners\", [])\n if (email := owner.get(\"emailAddress\"))\n and email.split(\"@\")[-1] == primary_admin_email.split(\"@\")[-1]\n ]\n\n\ndef _execute_single_retrieval(\n retrieval_function: Callable,\n continue_on_404_or_403: bool = False,\n **request_kwargs: Any,\n) -> GoogleDriveFileType:\n \"\"\"Execute a single retrieval from Google Drive API\"\"\"\n try:\n results = retrieval_function(**request_kwargs).execute()\n except HttpError as e:\n if e.resp.status >= 500:\n results = add_retries(\n lambda: retrieval_function(**request_kwargs).execute()\n )()\n elif e.resp.status == 400:\n if (\n \"pageToken\" in request_kwargs\n and \"Invalid Value\" in str(e)\n and \"pageToken\" in str(e)\n ):\n logger.warning(\n f\"Invalid page token: {request_kwargs['pageToken']}, retrying from start of request\"\n )\n request_kwargs.pop(\"pageToken\")\n return _execute_single_retrieval(\n retrieval_function,\n continue_on_404_or_403,\n **request_kwargs,\n )\n logger.error(f\"Error executing request: {e}\")\n raise e\n elif _is_rate_limit_error(e):\n results = _execute_with_retry(retrieval_function(**request_kwargs))\n elif e.resp.status == 404 or e.resp.status == 403:\n if continue_on_404_or_403:\n logger.debug(f\"Error executing request: {e}\")\n results = {}\n else:\n raise e\n else:\n logger.exception(\"Error executing request:\")\n raise e\n except (TimeoutError, socket.timeout) as error:\n logger.warning(\n \"Timed out executing Google API request; retrying with backoff. Details: %s\",\n error,\n )\n results = add_retries(lambda: retrieval_function(**request_kwargs).execute())()\n\n return results\n\n\ndef execute_single_retrieval(\n retrieval_function: Callable,\n list_key: str | None = None,\n continue_on_404_or_403: bool = False,\n **request_kwargs: Any,\n) -> Iterator[GoogleDriveFileType]:\n results = _execute_single_retrieval(\n retrieval_function,\n continue_on_404_or_403,\n **request_kwargs,\n )\n if list_key:\n for item in results.get(list_key, []):\n yield item\n else:\n yield results\n\n\n# included for type purposes; caller should not need to address\n# Nones unless max_num_pages is specified. Use\n# execute_paginated_retrieval_with_max_pages instead if you want\n# the early stop + yield None after max_num_pages behavior.\ndef execute_paginated_retrieval(\n retrieval_function: Callable,\n list_key: str | None = None,\n continue_on_404_or_403: bool = False,\n **kwargs: Any,\n) -> Iterator[GoogleDriveFileType]:\n for item in _execute_paginated_retrieval(\n retrieval_function,\n list_key,\n continue_on_404_or_403,\n **kwargs,\n ):\n if not isinstance(item, str):\n yield item\n\n\ndef execute_paginated_retrieval_with_max_pages(\n retrieval_function: Callable,\n max_num_pages: int,\n list_key: str | None = None,\n continue_on_404_or_403: bool = False,\n **kwargs: Any,\n) -> Iterator[GoogleDriveFileType | str]:\n yield from _execute_paginated_retrieval(\n retrieval_function,\n list_key,\n continue_on_404_or_403,\n max_num_pages=max_num_pages,\n **kwargs,\n )\n\n\ndef _execute_paginated_retrieval(\n retrieval_function: Callable,\n list_key: str | None = None,\n continue_on_404_or_403: bool = False,\n max_num_pages: int | None = None,\n **kwargs: Any,\n) -> Iterator[GoogleDriveFileType | str]:\n \"\"\"Execute a paginated retrieval from Google Drive API\n Args:\n retrieval_function: The specific list function to call (e.g., service.files().list)\n list_key: If specified, each object returned by the retrieval function\n will be accessed at the specified key and yielded from.\n continue_on_404_or_403: If True, the retrieval will continue even if the request returns a 404 or 403 error.\n max_num_pages: If specified, the retrieval will stop after the specified number of pages and yield None.\n **kwargs: Arguments to pass to the list function\n \"\"\"\n if \"fields\" not in kwargs or \"nextPageToken\" not in kwargs[\"fields\"]:\n raise ValueError(\n \"fields must contain nextPageToken for execute_paginated_retrieval\"\n )\n next_page_token = kwargs.get(PAGE_TOKEN_KEY, \"\")\n num_pages = 0\n while next_page_token is not None:\n if max_num_pages is not None and num_pages >= max_num_pages:\n yield next_page_token\n return\n num_pages += 1\n request_kwargs = kwargs.copy()\n if next_page_token:\n request_kwargs[PAGE_TOKEN_KEY] = next_page_token\n results = _execute_single_retrieval(\n retrieval_function,\n continue_on_404_or_403,\n **request_kwargs,\n )\n\n next_page_token = results.get(NEXT_PAGE_TOKEN_KEY)\n if list_key:\n for item in results.get(list_key, []):\n yield item\n else:\n yield results\n" }, { "path": "backend/onyx/connectors/google_utils/resources.py", "content": "from collections.abc import Callable\nfrom typing import Any\n\nfrom google.auth.exceptions import RefreshError\nfrom google.oauth2.credentials import Credentials as OAuthCredentials\nfrom google.oauth2.service_account import Credentials as ServiceAccountCredentials\nfrom googleapiclient.discovery import build # type: ignore[import-untyped]\nfrom googleapiclient.discovery import Resource\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass GoogleDriveService(Resource):\n pass\n\n\nclass GoogleDocsService(Resource):\n pass\n\n\nclass AdminService(Resource):\n pass\n\n\nclass GmailService(Resource):\n pass\n\n\nclass RefreshableDriveObject:\n \"\"\"\n Running Google drive service retrieval functions\n involves accessing methods of the service object (ie. files().list())\n which can raise a RefreshError if the access token is expired.\n This class is a wrapper that propagates the ability to refresh the access token\n and retry the final retrieval function until execute() is called.\n \"\"\"\n\n def __init__(\n self,\n call_stack: Callable[[ServiceAccountCredentials | OAuthCredentials], Any],\n creds: ServiceAccountCredentials | OAuthCredentials,\n creds_getter: Callable[..., ServiceAccountCredentials | OAuthCredentials],\n ):\n self.call_stack = call_stack\n self.creds = creds\n self.creds_getter = creds_getter\n\n def __getattr__(self, name: str) -> Any:\n if name == \"execute\":\n return self.make_refreshable_execute()\n return RefreshableDriveObject(\n lambda creds: getattr(self.call_stack(creds), name),\n self.creds,\n self.creds_getter,\n )\n\n def __call__(self, *args: Any, **kwargs: Any) -> Any:\n return RefreshableDriveObject(\n lambda creds: self.call_stack(creds)(*args, **kwargs),\n self.creds,\n self.creds_getter,\n )\n\n def make_refreshable_execute(self) -> Callable:\n def execute(*args: Any, **kwargs: Any) -> Any:\n try:\n return self.call_stack(self.creds).execute(*args, **kwargs)\n except RefreshError as e:\n logger.warning(\n f\"RefreshError, going to attempt a creds refresh and retry: {e}\"\n )\n # Refresh the access token\n self.creds = self.creds_getter()\n return self.call_stack(self.creds).execute(*args, **kwargs)\n\n return execute\n\n\ndef _get_google_service(\n service_name: str,\n service_version: str,\n creds: ServiceAccountCredentials | OAuthCredentials,\n user_email: str | None = None,\n) -> GoogleDriveService | GoogleDocsService | AdminService | GmailService:\n service: Resource\n if isinstance(creds, ServiceAccountCredentials):\n # NOTE: https://developers.google.com/identity/protocols/oauth2/service-account#error-codes\n creds = creds.with_subject(user_email)\n service = build(service_name, service_version, credentials=creds)\n elif isinstance(creds, OAuthCredentials):\n service = build(service_name, service_version, credentials=creds)\n\n return service\n\n\ndef get_google_docs_service(\n creds: ServiceAccountCredentials | OAuthCredentials,\n user_email: str | None = None,\n) -> GoogleDocsService:\n return _get_google_service(\"docs\", \"v1\", creds, user_email)\n\n\ndef get_drive_service(\n creds: ServiceAccountCredentials | OAuthCredentials,\n user_email: str | None = None,\n) -> GoogleDriveService:\n return _get_google_service(\"drive\", \"v3\", creds, user_email)\n\n\ndef get_admin_service(\n creds: ServiceAccountCredentials | OAuthCredentials,\n user_email: str | None = None,\n) -> AdminService:\n return _get_google_service(\"admin\", \"directory_v1\", creds, user_email)\n\n\ndef get_gmail_service(\n creds: ServiceAccountCredentials | OAuthCredentials,\n user_email: str | None = None,\n) -> GmailService:\n return _get_google_service(\"gmail\", \"v1\", creds, user_email)\n" }, { "path": "backend/onyx/connectors/google_utils/shared_constants.py", "content": "from enum import Enum as PyEnum\n\nfrom onyx.configs.constants import DocumentSource\n\n# NOTE: do not need https://www.googleapis.com/auth/documents.readonly\n# this is counted under `/auth/drive.readonly`\nGOOGLE_SCOPES = {\n DocumentSource.GOOGLE_DRIVE: [\n \"https://www.googleapis.com/auth/drive.readonly\",\n \"https://www.googleapis.com/auth/drive.metadata.readonly\",\n \"https://www.googleapis.com/auth/admin.directory.group.readonly\",\n \"https://www.googleapis.com/auth/admin.directory.user.readonly\",\n ],\n DocumentSource.GMAIL: [\n \"https://www.googleapis.com/auth/gmail.readonly\",\n \"https://www.googleapis.com/auth/admin.directory.user.readonly\",\n \"https://www.googleapis.com/auth/admin.directory.group.readonly\",\n ],\n}\n\n# This is the Oauth token\nDB_CREDENTIALS_DICT_TOKEN_KEY = \"google_tokens\"\n# This is the service account key\nDB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY = \"google_service_account_key\"\n# The email saved for both auth types\nDB_CREDENTIALS_PRIMARY_ADMIN_KEY = \"google_primary_admin\"\n\n# https://developers.google.com/workspace/guides/create-credentials\n# Internally defined authentication method type.\n# The value must be one of \"oauth_interactive\" or \"uploaded\"\n# Used to disambiguate whether credentials have already been created via\n# certain methods and what actions we allow users to take\nDB_CREDENTIALS_AUTHENTICATION_METHOD = \"authentication_method\"\n\n\nclass GoogleOAuthAuthenticationMethod(str, PyEnum):\n OAUTH_INTERACTIVE = \"oauth_interactive\"\n UPLOADED = \"uploaded\"\n\n\nUSER_FIELDS = \"nextPageToken, users(primaryEmail)\"\n\n# Error message substrings\nMISSING_SCOPES_ERROR_STR = \"client not authorized for any of the scopes requested\"\n\n# Documentation and error messages\nSCOPE_DOC_URL = \"https://docs.onyx.app/admins/connectors/official/google_drive/overview\"\nONYX_SCOPE_INSTRUCTIONS = (\n \"You have upgraded Onyx without updating the Google Auth scopes. \"\n f\"Please refer to the documentation to learn how to update the scopes: {SCOPE_DOC_URL}\"\n)\n\n\n# This is the maximum number of threads that can be retrieved at once\nSLIM_BATCH_SIZE = 500\n" }, { "path": "backend/onyx/connectors/guru/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/guru/connector.py", "content": "import json\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nimport requests\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n# Potential Improvements\n# 1. Support fetching per collection via collection token (configured at connector creation)\nGURU_API_BASE = \"https://api.getguru.com/api/v1/\"\nGURU_QUERY_ENDPOINT = GURU_API_BASE + \"search/query\"\nGURU_CARDS_URL = \"https://app.getguru.com/card/\"\n\n\ndef unixtime_to_guru_time_str(unix_time: SecondsSinceUnixEpoch) -> str:\n date_obj = datetime.fromtimestamp(unix_time, tz=timezone.utc)\n date_str = date_obj.strftime(\"%Y-%m-%dT%H:%M:%S.%f\")[:-3]\n tz_str = date_obj.strftime(\"%z\")\n return date_str + tz_str\n\n\nclass GuruConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n guru_user: str | None = None,\n guru_user_token: str | None = None,\n ) -> None:\n self.batch_size = batch_size\n self.guru_user = guru_user\n self.guru_user_token = guru_user_token\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.guru_user = credentials[\"guru_user\"]\n self.guru_user_token = credentials[\"guru_user_token\"]\n return None\n\n def _process_cards(\n self, start_str: str | None = None, end_str: str | None = None\n ) -> GenerateDocumentsOutput:\n if self.guru_user is None or self.guru_user_token is None:\n raise ConnectorMissingCredentialError(\"Guru\")\n\n doc_batch: list[Document | HierarchyNode] = []\n\n session = requests.Session()\n session.auth = (self.guru_user, self.guru_user_token)\n\n params: dict[str, str | int] = {\"maxResults\": self.batch_size}\n\n if start_str is not None and end_str is not None:\n params[\"q\"] = f\"lastModified >= {start_str} AND lastModified < {end_str}\"\n\n current_url = GURU_QUERY_ENDPOINT # This is how they handle pagination, a different url will be provided\n while True:\n response = session.get(current_url, params=params)\n response.raise_for_status()\n\n if response.status_code == 204:\n break\n\n cards = json.loads(response.text)\n for card in cards:\n title = card[\"preferredPhrase\"]\n link = GURU_CARDS_URL + card[\"slug\"]\n content_text = parse_html_page_basic(card[\"content\"])\n last_updated = time_str_to_utc(card[\"lastModified\"])\n last_verified = (\n time_str_to_utc(card.get(\"lastVerified\"))\n if card.get(\"lastVerified\")\n else None\n )\n\n # For Onyx, we decay document score overtime, either last_updated or\n # last_verified is a good enough signal for the document's recency\n latest_time = (\n max(last_verified, last_updated) if last_verified else last_updated\n )\n\n metadata_dict: dict[str, str | list[str]] = {}\n tags = [tag.get(\"value\") for tag in card.get(\"tags\", [])]\n if tags:\n metadata_dict[\"tags\"] = tags\n\n boards = [board.get(\"title\") for board in card.get(\"boards\", [])]\n if boards:\n # In UI it's called Folders\n metadata_dict[\"folders\"] = boards\n\n collection = card.get(\"collection\", {})\n if collection:\n metadata_dict[\"collection_name\"] = collection.get(\"name\", \"\")\n\n owner = card.get(\"owner\", {})\n author = None\n if owner:\n author = BasicExpertInfo(\n email=owner.get(\"email\"),\n first_name=owner.get(\"firstName\"),\n last_name=owner.get(\"lastName\"),\n )\n\n doc_batch.append(\n Document(\n id=card[\"id\"],\n sections=[TextSection(link=link, text=content_text)],\n source=DocumentSource.GURU,\n semantic_identifier=title,\n doc_updated_at=latest_time,\n primary_owners=[author] if author is not None else None,\n # Can add verifies and commenters later\n metadata=metadata_dict,\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if not hasattr(response, \"links\") or not response.links:\n break\n current_url = response.links[\"next-page\"][\"url\"]\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._process_cards()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_time = unixtime_to_guru_time_str(start)\n end_time = unixtime_to_guru_time_str(end)\n\n return self._process_cards(start_time, end_time)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = GuruConnector()\n connector.load_credentials(\n {\n \"guru_user\": os.environ[\"GURU_USER\"],\n \"guru_user_token\": os.environ[\"GURU_USER_TOKEN\"],\n }\n )\n\n latest_docs = connector.load_from_state()\n print(next(latest_docs))\n" }, { "path": "backend/onyx/connectors/highspot/__init__.py", "content": "\"\"\"\nHighspot connector package for Onyx.\nEnables integration with Highspot's knowledge base.\n\"\"\"\n" }, { "path": "backend/onyx/connectors/highspot/client.py", "content": "import base64\nfrom typing import Any\nfrom typing import Dict\nfrom typing import List\nfrom typing import Optional\nfrom urllib.parse import urljoin\n\nimport requests\nfrom requests.adapters import HTTPAdapter\nfrom requests.exceptions import HTTPError\nfrom requests.exceptions import RequestException\nfrom requests.exceptions import Timeout\nfrom urllib3.util.retry import Retry\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\nPAGE_SIZE = 100\n\n\nclass HighspotClientError(Exception):\n \"\"\"Base exception for Highspot API client errors.\"\"\"\n\n def __init__(self, message: str, status_code: Optional[int] = None):\n self.message = message\n self.status_code = status_code\n super().__init__(self.message)\n\n\nclass HighspotAuthenticationError(HighspotClientError):\n \"\"\"Exception raised for authentication errors.\"\"\"\n\n\nclass HighspotRateLimitError(HighspotClientError):\n \"\"\"Exception raised when rate limit is exceeded.\"\"\"\n\n def __init__(self, message: str, retry_after: Optional[str] = None):\n self.retry_after = retry_after\n super().__init__(message)\n\n\nclass HighspotClient:\n \"\"\"\n Client for interacting with the Highspot API.\n\n Uses basic authentication with provided key (username) and secret (password).\n Implements retry logic, error handling, and connection pooling.\n \"\"\"\n\n BASE_URL = \"https://api-su2.highspot.com/v1.0/\"\n\n def __init__(\n self,\n key: str,\n secret: str,\n base_url: str = BASE_URL,\n timeout: int = 30,\n max_retries: int = 3,\n backoff_factor: float = 0.5,\n status_forcelist: Optional[List[int]] = None,\n ):\n \"\"\"\n Initialize the Highspot API client.\n\n Args:\n key: API key (used as username)\n secret: API secret (used as password)\n base_url: Base URL for the Highspot API\n timeout: Request timeout in seconds\n max_retries: Maximum number of retries for failed requests\n backoff_factor: Backoff factor for retries\n status_forcelist: HTTP status codes to retry on\n \"\"\"\n if not key or not secret:\n raise ValueError(\"API key and secret are required\")\n\n self.key = key\n self.secret = secret\n self.base_url = base_url.rstrip(\"/\") + \"/\"\n self.timeout = timeout\n\n # Set up session with retry logic\n self.session = requests.Session()\n retry_strategy = Retry(\n total=max_retries,\n backoff_factor=backoff_factor,\n status_forcelist=status_forcelist or [429, 500, 502, 503, 504],\n allowed_methods=[\"GET\", \"POST\", \"PUT\", \"DELETE\"],\n )\n adapter = HTTPAdapter(max_retries=retry_strategy)\n self.session.mount(\"http://\", adapter)\n self.session.mount(\"https://\", adapter)\n\n # Set up authentication\n self._setup_auth()\n\n def _setup_auth(self) -> None:\n \"\"\"Set up basic authentication for the session.\"\"\"\n auth = f\"{self.key}:{self.secret}\"\n encoded_auth = base64.b64encode(auth.encode()).decode()\n self.session.headers.update(\n {\n \"Authorization\": f\"Basic {encoded_auth}\",\n \"Content-Type\": \"application/json\",\n \"Accept\": \"application/json\",\n }\n )\n\n def _make_request(\n self,\n method: str,\n endpoint: str,\n params: Optional[Dict[str, Any]] = None,\n data: Optional[Dict[str, Any]] = None,\n json_data: Optional[Dict[str, Any]] = None,\n headers: Optional[Dict[str, str]] = None,\n ) -> Dict[str, Any]:\n \"\"\"\n Make a request to the Highspot API.\n\n Args:\n method: HTTP method (GET, POST, etc.)\n endpoint: API endpoint\n params: URL parameters\n data: Form data\n json_data: JSON data\n headers: Additional headers\n\n Returns:\n API response as a dictionary\n\n Raises:\n HighspotClientError: On API errors\n HighspotAuthenticationError: On authentication errors\n HighspotRateLimitError: On rate limiting\n requests.exceptions.RequestException: On request failures\n \"\"\"\n url = urljoin(self.base_url, endpoint)\n request_headers = {}\n if headers:\n request_headers.update(headers)\n\n try:\n logger.debug(f\"Making {method} request to {url}\")\n response = self.session.request(\n method=method,\n url=url,\n params=params,\n data=data,\n json=json_data,\n headers=request_headers,\n timeout=self.timeout,\n )\n response.raise_for_status()\n\n if response.content and response.content.strip():\n return response.json()\n return {}\n\n except HTTPError as e:\n status_code = e.response.status_code\n error_msg = str(e)\n\n try:\n error_data = e.response.json()\n if isinstance(error_data, dict):\n error_msg = error_data.get(\"message\", str(e))\n except (ValueError, KeyError):\n pass\n\n if status_code == 401:\n raise HighspotAuthenticationError(f\"Authentication failed: {error_msg}\")\n elif status_code == 429:\n retry_after = e.response.headers.get(\"Retry-After\")\n raise HighspotRateLimitError(\n f\"Rate limit exceeded: {error_msg}\", retry_after=retry_after\n )\n else:\n raise HighspotClientError(\n f\"API error {status_code}: {error_msg}\", status_code=status_code\n )\n\n except Timeout:\n raise HighspotClientError(\"Request timed out\")\n except RequestException as e:\n raise HighspotClientError(f\"Request failed: {str(e)}\")\n\n def get_spots(self) -> List[Dict[str, Any]]:\n \"\"\"\n Get all available spots, paginated.\n\n Returns:\n List of spots with their names and IDs\n \"\"\"\n all_spots = []\n has_more = True\n current_offset = 0\n\n while has_more:\n params = {\"right\": \"view\", \"start\": current_offset, \"limit\": PAGE_SIZE}\n response = self._make_request(\"GET\", \"spots\", params=params)\n found_spots = response.get(\"collection\", [])\n logger.info(f\"Received {len(found_spots)} spots at offset {current_offset}\")\n all_spots.extend(found_spots)\n if len(found_spots) < PAGE_SIZE:\n has_more = False\n else:\n current_offset += PAGE_SIZE\n logger.info(f\"Total spots retrieved: {len(all_spots)}\")\n return all_spots\n\n def get_spot(self, spot_id: str) -> Dict[str, Any]:\n \"\"\"\n Get details for a specific spot.\n\n Args:\n spot_id: ID of the spot\n\n Returns:\n Spot details\n \"\"\"\n if not spot_id:\n raise ValueError(\"spot_id is required\")\n return self._make_request(\"GET\", f\"spots/{spot_id}\")\n\n def get_spot_items(\n self, spot_id: str, offset: int = 0, page_size: int = PAGE_SIZE\n ) -> Dict[str, Any]:\n \"\"\"\n Get items in a specific spot.\n\n Args:\n spot_id: ID of the spot\n offset: offset number\n page_size: Number of items per page\n\n Returns:\n Items in the spot\n \"\"\"\n if not spot_id:\n raise ValueError(\"spot_id is required\")\n\n params = {\"spot\": spot_id, \"start\": offset, \"limit\": page_size}\n return self._make_request(\"GET\", \"items\", params=params)\n\n def get_item(self, item_id: str) -> Dict[str, Any]:\n \"\"\"\n Get details for a specific item.\n\n Args:\n item_id: ID of the item\n\n Returns:\n Item details\n \"\"\"\n if not item_id:\n raise ValueError(\"item_id is required\")\n return self._make_request(\"GET\", f\"items/{item_id}\")\n\n def get_item_content(self, item_id: str) -> bytes:\n \"\"\"\n Get the raw content of an item.\n\n Args:\n item_id: ID of the item\n\n Returns:\n Raw content bytes\n \"\"\"\n if not item_id:\n raise ValueError(\"item_id is required\")\n\n url = urljoin(self.base_url, f\"items/{item_id}/content\")\n response = self.session.get(url, timeout=self.timeout)\n response.raise_for_status()\n return response.content\n\n def health_check(self) -> bool:\n \"\"\"\n Check if the API is accessible and credentials are valid.\n\n Returns:\n True if API is accessible, False otherwise\n \"\"\"\n try:\n self._make_request(\"GET\", \"spots\", params={\"limit\": 1})\n return True\n except (HighspotClientError, HighspotAuthenticationError):\n return False\n" }, { "path": "backend/onyx/connectors/highspot/connector.py", "content": "import os\nfrom datetime import datetime\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import Dict\nfrom typing import List\nfrom typing import Optional\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.highspot.client import HighspotClient\nfrom onyx.connectors.highspot.client import HighspotClientError\nfrom onyx.connectors.highspot.utils import scrape_url_content\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import extract_file_text\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n_SLIM_BATCH_SIZE = 1000\n\n\nclass HighspotSpot(BaseModel):\n id: str\n name: str\n\n\nclass HighspotConnector(LoadConnector, PollConnector, SlimConnectorWithPermSync):\n \"\"\"\n Connector for loading data from Highspot.\n\n Retrieves content from specified spots using the Highspot API.\n If no spots are specified, retrieves content from all available spots.\n \"\"\"\n\n def __init__(\n self,\n spot_names: list[str] | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ):\n \"\"\"\n Initialize the Highspot connector.\n\n Args:\n spot_names: List of spot names to retrieve content from (if empty, gets all spots)\n batch_size: Number of items to retrieve in each batch\n \"\"\"\n self.spot_names = spot_names or []\n self.batch_size = batch_size\n\n self._client: Optional[HighspotClient] = None\n self.highspot_url: Optional[str] = None\n self.key: Optional[str] = None\n self.secret: Optional[str] = None\n\n @property\n def client(self) -> HighspotClient:\n if self._client is None:\n if not self.key or not self.secret:\n raise ConnectorMissingCredentialError(\"Highspot\")\n # Ensure highspot_url is a string, use default if None\n base_url = (\n self.highspot_url\n if self.highspot_url is not None\n else HighspotClient.BASE_URL\n )\n self._client = HighspotClient(self.key, self.secret, base_url=base_url)\n return self._client\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n logger.info(\"Loading Highspot credentials\")\n self.highspot_url = credentials.get(\"highspot_url\")\n self.key = credentials.get(\"highspot_key\")\n self.secret = credentials.get(\"highspot_secret\")\n return None\n\n def _fetch_spots(self) -> list[HighspotSpot]:\n \"\"\"\n Populate the spot ID map with all available spots.\n Keys are stored as lowercase for case-insensitive lookups.\n \"\"\"\n return [\n HighspotSpot(id=spot[\"id\"], name=spot[\"title\"])\n for spot in self.client.get_spots()\n ]\n\n def _fetch_spots_to_process(self) -> list[HighspotSpot]:\n \"\"\"\n Fetch spots to process based on the configured spot names.\n \"\"\"\n spots = self._fetch_spots()\n if not spots:\n raise ValueError(\"No spots found in Highspot.\")\n\n if self.spot_names:\n lower_spot_names = [name.lower() for name in self.spot_names]\n spots_to_process = [\n spot for spot in spots if spot.name.lower() in lower_spot_names\n ]\n if not spots_to_process:\n raise ValueError(\n f\"No valid spots found in Highspot. Found {spots} but {self.spot_names} were requested.\"\n )\n return spots_to_process\n\n return spots\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"\n Load content from configured spots in Highspot.\n If no spots are configured, loads from all spots.\n\n Yields:\n Batches of Document objects\n \"\"\"\n return self.poll_source(None, None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch | None, end: SecondsSinceUnixEpoch | None\n ) -> GenerateDocumentsOutput:\n \"\"\"\n Poll Highspot for content updated since the start time.\n\n Args:\n start: Start time as seconds since Unix epoch\n end: End time as seconds since Unix epoch\n\n Yields:\n Batches of Document objects\n \"\"\"\n spots_to_process = self._fetch_spots_to_process()\n\n doc_batch: list[Document | HierarchyNode] = []\n try:\n for spot in spots_to_process:\n try:\n offset = 0\n has_more = True\n\n while has_more:\n logger.info(\n f\"Retrieving items from spot {spot.name}, offset {offset}\"\n )\n response = self.client.get_spot_items(\n spot_id=spot.id, offset=offset, page_size=self.batch_size\n )\n items = response.get(\"collection\", [])\n logger.info(\n f\"Received {len(items)} items from spot {spot.name}\"\n )\n if not items:\n has_more = False\n continue\n\n for item in items:\n try:\n item_id = item.get(\"id\")\n if not item_id:\n logger.warning(\"Item without ID found, skipping\")\n continue\n\n item_details = self.client.get_item(item_id)\n if not item_details:\n logger.warning(\n f\"Item {item_id} details not found, skipping\"\n )\n continue\n # Apply time filter if specified\n if start or end:\n updated_at = item_details.get(\"date_updated\")\n if updated_at:\n # Convert to datetime for comparison\n try:\n updated_time = datetime.fromisoformat(\n updated_at.replace(\"Z\", \"+00:00\")\n )\n if (\n start\n and updated_time.timestamp() < start\n ) or (\n end and updated_time.timestamp() > end\n ):\n continue\n except (ValueError, TypeError):\n # Skip if date cannot be parsed\n logger.warning(\n f\"Invalid date format for item {item_id}: {updated_at}\"\n )\n continue\n\n content = self._get_item_content(item_details)\n\n title = item_details.get(\"title\", \"\")\n\n doc_batch.append(\n Document(\n id=f\"HIGHSPOT_{item_id}\",\n sections=[\n TextSection(\n link=item_details.get(\n \"url\",\n f\"https://www.highspot.com/items/{item_id}\",\n ),\n text=content,\n )\n ],\n source=DocumentSource.HIGHSPOT,\n semantic_identifier=title,\n metadata={\n \"spot_name\": spot.name,\n \"type\": item_details.get(\n \"content_type\", \"\"\n ),\n \"created_at\": item_details.get(\n \"date_added\", \"\"\n ),\n \"author\": item_details.get(\"author\", \"\"),\n \"language\": item_details.get(\n \"language\", \"\"\n ),\n \"can_download\": str(\n item_details.get(\"can_download\", False)\n ),\n },\n doc_updated_at=item_details.get(\"date_updated\"),\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n except HighspotClientError as e:\n item_id = \"ID\" if not item_id else item_id\n logger.error(\n f\"Error retrieving item {item_id}: {str(e)}\"\n )\n except Exception as e:\n item_id = \"ID\" if not item_id else item_id\n logger.error(\n f\"Unexpected error for item {item_id}: {str(e)}\"\n )\n\n has_more = len(items) >= self.batch_size\n offset += self.batch_size\n\n except (HighspotClientError, ValueError) as e:\n logger.error(f\"Error processing spot {spot.name}: {str(e)}\")\n raise\n except Exception as e:\n logger.error(\n f\"Unexpected error processing spot {spot.name}: {str(e)}\"\n )\n raise\n\n except Exception as e:\n logger.error(f\"Error in Highspot connector: {str(e)}\")\n raise\n\n if doc_batch:\n yield doc_batch\n\n def _get_item_content(self, item_details: Dict[str, Any]) -> str:\n \"\"\"\n Get the text content of an item.\n\n Args:\n item_details: Item details from the API\n\n Returns:\n Text content of the item\n \"\"\"\n item_id = item_details.get(\"id\", \"\")\n content_name = item_details.get(\"content_name\", \"\")\n is_valid_format = content_name and \".\" in content_name\n file_extension = content_name.split(\".\")[-1].lower() if is_valid_format else \"\"\n file_extension = \".\" + file_extension if file_extension else \"\"\n can_download = item_details.get(\"can_download\", False)\n content_type = item_details.get(\"content_type\", \"\")\n\n # Extract title and description once at the beginning\n title, description = self._extract_title_and_description(item_details)\n default_content = f\"{title}\\n{description}\"\n logger.info(\n f\"Processing item {item_id} with extension {file_extension} and file name {content_name}\"\n )\n\n try:\n if content_type == \"WebLink\":\n url = item_details.get(\"url\")\n if not url:\n return default_content\n content = scrape_url_content(url, True)\n return content if content else default_content\n\n elif (\n is_valid_format\n and file_extension in OnyxFileExtensions.TEXT_AND_DOCUMENT_EXTENSIONS\n and can_download\n ):\n content_response = self.client.get_item_content(item_id)\n # Process and extract text from binary content based on type\n if content_response:\n text_content = extract_file_text(\n BytesIO(content_response), content_name, False\n )\n return text_content if text_content else default_content\n return default_content\n\n else:\n logger.warning(\n f\"Item {item_id} has unsupported format: {file_extension}\"\n )\n return default_content\n\n except HighspotClientError as e:\n error_context = f\"item {item_id}\" if item_id else \"(item id not found)\"\n logger.warning(f\"Could not retrieve content for {error_context}: {str(e)}\")\n return default_content\n except ValueError as e:\n error_context = f\"item {item_id}\" if item_id else \"(item id not found)\"\n logger.error(f\"Value error for {error_context}: {str(e)}\")\n return default_content\n\n except Exception as e:\n error_context = f\"item {item_id}\" if item_id else \"(item id not found)\"\n logger.error(\n f\"Unexpected error retrieving content for {error_context}: {str(e)}\"\n )\n return default_content\n\n def _extract_title_and_description(\n self, item_details: Dict[str, Any]\n ) -> tuple[str, str]:\n \"\"\"\n Extract the title and description from item details.\n\n Args:\n item_details: Item details from the API\n\n Returns:\n Tuple of title and description\n \"\"\"\n title = item_details.get(\"title\", \"\")\n description = item_details.get(\"description\", \"\")\n return title, description\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n end: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n \"\"\"\n Retrieve all document IDs from the configured spots.\n If no spots are configured, retrieves from all spots.\n\n Args:\n start: Optional start time filter\n end: Optional end time filter\n callback: Optional indexing heartbeat callback\n\n Yields:\n Batches of SlimDocument objects\n \"\"\"\n spots_to_process = self._fetch_spots_to_process()\n\n slim_doc_batch: list[SlimDocument | HierarchyNode] = []\n try:\n for spot in spots_to_process:\n try:\n offset = 0\n has_more = True\n\n while has_more:\n logger.info(\n f\"Retrieving slim documents from spot {spot.name}, offset {offset}\"\n )\n response = self.client.get_spot_items(\n spot_id=spot.id, offset=offset, page_size=self.batch_size\n )\n\n items = response.get(\"collection\", [])\n if not items:\n has_more = False\n continue\n\n for item in items:\n item_id = item.get(\"id\")\n if not item_id:\n logger.warning(\"Item without ID found, skipping\")\n continue\n\n slim_doc_batch.append(\n SlimDocument(id=f\"HIGHSPOT_{item_id}\")\n )\n\n if len(slim_doc_batch) >= _SLIM_BATCH_SIZE:\n yield slim_doc_batch\n slim_doc_batch = []\n\n has_more = len(items) >= self.batch_size\n offset += self.batch_size\n\n except (HighspotClientError, ValueError):\n logger.exception(\n f\"Error retrieving slim documents from spot {spot.name}\"\n )\n raise\n\n if slim_doc_batch:\n yield slim_doc_batch\n except Exception:\n logger.exception(\"Error in Highspot Slim Connector\")\n raise\n\n def validate_credentials(self) -> bool:\n \"\"\"\n Validate that the provided credentials can access the Highspot API.\n\n Returns:\n True if credentials are valid, False otherwise\n \"\"\"\n try:\n return self.client.health_check()\n except Exception as e:\n logger.error(f\"Failed to validate credentials: {str(e)}\")\n return False\n\n\nif __name__ == \"__main__\":\n spot_names: List[str] = []\n connector = HighspotConnector(spot_names)\n credentials = {\n \"highspot_key\": os.environ.get(\"HIGHSPOT_KEY\"),\n \"highspot_secret\": os.environ.get(\"HIGHSPOT_SECRET\"),\n }\n connector.load_credentials(credentials=credentials)\n for doc in connector.load_from_state():\n print(doc)\n" }, { "path": "backend/onyx/connectors/highspot/utils.py", "content": "from typing import Optional\nfrom urllib.parse import urlparse\n\nfrom bs4 import BeautifulSoup\nfrom playwright.sync_api import sync_playwright\n\nfrom onyx.file_processing.html_utils import web_html_cleanup\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Constants\nWEB_CONNECTOR_MAX_SCROLL_ATTEMPTS = 20\nJAVASCRIPT_DISABLED_MESSAGE = \"You have JavaScript disabled in your browser\"\nDEFAULT_TIMEOUT = 60000 # 60 seconds\n\n\ndef scrape_url_content(\n url: str, scroll_before_scraping: bool = False, timeout_ms: int = DEFAULT_TIMEOUT\n) -> Optional[str]:\n \"\"\"\n Scrapes content from a given URL and returns the cleaned text.\n\n Args:\n url: The URL to scrape\n scroll_before_scraping: Whether to scroll through the page to load lazy content\n timeout_ms: Timeout in milliseconds for page navigation and loading\n\n Returns:\n The cleaned text content of the page or None if scraping fails\n \"\"\"\n playwright = None\n browser = None\n try:\n validate_url(url)\n playwright = sync_playwright().start()\n browser = playwright.chromium.launch(headless=True)\n context = browser.new_context()\n page = context.new_page()\n\n logger.info(f\"Navigating to URL: {url}\")\n try:\n page.goto(url, timeout=timeout_ms)\n except Exception as e:\n logger.error(f\"Failed to navigate to {url}: {str(e)}\")\n return None\n\n if scroll_before_scraping:\n logger.debug(\"Scrolling page to load lazy content\")\n scroll_attempts = 0\n previous_height = page.evaluate(\"document.body.scrollHeight\")\n while scroll_attempts < WEB_CONNECTOR_MAX_SCROLL_ATTEMPTS:\n page.evaluate(\"window.scrollTo(0, document.body.scrollHeight)\")\n try:\n page.wait_for_load_state(\"networkidle\", timeout=timeout_ms)\n except Exception as e:\n logger.warning(f\"Network idle wait timed out: {str(e)}\")\n break\n\n new_height = page.evaluate(\"document.body.scrollHeight\")\n if new_height == previous_height:\n break\n previous_height = new_height\n scroll_attempts += 1\n\n content = page.content()\n soup = BeautifulSoup(content, \"html.parser\")\n\n parsed_html = web_html_cleanup(soup)\n\n if JAVASCRIPT_DISABLED_MESSAGE in parsed_html.cleaned_text:\n logger.debug(\"JavaScript disabled message detected, checking iframes\")\n try:\n iframe_count = page.frame_locator(\"iframe\").locator(\"html\").count()\n if iframe_count > 0:\n iframe_texts = (\n page.frame_locator(\"iframe\").locator(\"html\").all_inner_texts()\n )\n iframe_content = \"\\n\".join(iframe_texts)\n\n if len(parsed_html.cleaned_text) < 700:\n parsed_html.cleaned_text = iframe_content\n else:\n parsed_html.cleaned_text += \"\\n\" + iframe_content\n except Exception as e:\n logger.warning(f\"Error processing iframes: {str(e)}\")\n\n return parsed_html.cleaned_text\n\n except Exception as e:\n logger.error(f\"Error scraping URL {url}: {str(e)}\")\n return None\n\n finally:\n if browser:\n try:\n browser.close()\n except Exception as e:\n logger.debug(f\"Error closing browser: {str(e)}\")\n if playwright:\n try:\n playwright.stop()\n except Exception as e:\n logger.debug(f\"Error stopping playwright: {str(e)}\")\n\n\ndef validate_url(url: str) -> None:\n \"\"\"\n Validates that a URL is properly formatted.\n\n Args:\n url: The URL to validate\n\n Raises:\n ValueError: If URL is not valid\n \"\"\"\n parse = urlparse(url)\n if parse.scheme != \"http\" and parse.scheme != \"https\":\n raise ValueError(\"URL must be of scheme https?://\")\n\n if not parse.hostname:\n raise ValueError(\"URL must include a hostname\")\n" }, { "path": "backend/onyx/connectors/hubspot/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/hubspot/connector.py", "content": "import re\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import TypeVar\n\nimport requests\nfrom hubspot import HubSpot # type: ignore\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.hubspot.rate_limit import HubSpotRateLimiter\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nHUBSPOT_BASE_URL = \"https://app.hubspot.com\"\nHUBSPOT_API_URL = \"https://api.hubapi.com/integrations/v1/me\"\n\nAVAILABLE_OBJECT_TYPES = {\"tickets\", \"companies\", \"deals\", \"contacts\"}\n\nHUBSPOT_PAGE_SIZE = 100\n\nT = TypeVar(\"T\")\n\nlogger = setup_logger()\n\n\nclass HubSpotConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n access_token: str | None = None,\n object_types: list[str] | None = None,\n ) -> None:\n self.batch_size = batch_size\n self._access_token = access_token\n self._portal_id: str | None = None\n self._rate_limiter = HubSpotRateLimiter()\n\n # Set object types to fetch, default to all available types\n if object_types is None:\n self.object_types = AVAILABLE_OBJECT_TYPES.copy()\n else:\n object_types_set = set(object_types)\n\n # Validate provided object types\n invalid_types = object_types_set - AVAILABLE_OBJECT_TYPES\n if invalid_types:\n raise ValueError(\n f\"Invalid object types: {invalid_types}. Available types: {AVAILABLE_OBJECT_TYPES}\"\n )\n self.object_types = object_types_set.copy()\n\n @property\n def access_token(self) -> str:\n \"\"\"Get the access token, raising an exception if not set.\"\"\"\n if self._access_token is None:\n raise ConnectorMissingCredentialError(\"HubSpot access token not set\")\n return self._access_token\n\n @access_token.setter\n def access_token(self, value: str | None) -> None:\n \"\"\"Set the access token.\"\"\"\n self._access_token = value\n\n @property\n def portal_id(self) -> str:\n \"\"\"Get the portal ID, raising an exception if not set.\"\"\"\n if self._portal_id is None:\n raise ConnectorMissingCredentialError(\"HubSpot portal ID not set\")\n return self._portal_id\n\n @portal_id.setter\n def portal_id(self, value: str | None) -> None:\n \"\"\"Set the portal ID.\"\"\"\n self._portal_id = value\n\n def _call_hubspot(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> T:\n return self._rate_limiter.call(func, *args, **kwargs)\n\n def _paginated_results(\n self,\n fetch_page: Callable[..., Any],\n **kwargs: Any,\n ) -> Generator[Any, None, None]:\n base_kwargs = dict(kwargs)\n base_kwargs.setdefault(\"limit\", HUBSPOT_PAGE_SIZE)\n\n after: str | None = None\n while True:\n page_kwargs = base_kwargs.copy()\n if after is not None:\n page_kwargs[\"after\"] = after\n\n page = self._call_hubspot(fetch_page, **page_kwargs)\n results = getattr(page, \"results\", [])\n for result in results:\n yield result\n\n paging = getattr(page, \"paging\", None)\n next_page = getattr(paging, \"next\", None) if paging else None\n if next_page is None:\n break\n\n after = getattr(next_page, \"after\", None)\n if after is None:\n break\n\n def _clean_html_content(self, html_content: str) -> str:\n \"\"\"Clean HTML content and extract raw text\"\"\"\n if not html_content:\n return \"\"\n\n # Remove HTML tags using regex\n clean_text = re.sub(r\"<[^>]+>\", \"\", html_content)\n\n # Decode common HTML entities\n clean_text = clean_text.replace(\" \", \" \")\n clean_text = clean_text.replace(\"&\", \"&\")\n clean_text = clean_text.replace(\"<\", \"<\")\n clean_text = clean_text.replace(\">\", \">\")\n clean_text = clean_text.replace(\""\", '\"')\n clean_text = clean_text.replace(\"'\", \"'\")\n\n # Clean up whitespace\n clean_text = \" \".join(clean_text.split())\n\n return clean_text.strip()\n\n def get_portal_id(self) -> str:\n headers = {\n \"Authorization\": f\"Bearer {self.access_token}\",\n \"Content-Type\": \"application/json\",\n }\n\n response = requests.get(HUBSPOT_API_URL, headers=headers)\n if response.status_code != 200:\n raise Exception(\"Error fetching portal ID\")\n\n data = response.json()\n return str(data[\"portalId\"])\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.access_token = cast(str, credentials[\"hubspot_access_token\"])\n self.portal_id = self.get_portal_id()\n return None\n\n def _get_object_url(self, object_type: str, object_id: str) -> str:\n \"\"\"Generate HubSpot URL for different object types\"\"\"\n if object_type == \"tickets\":\n return (\n f\"{HUBSPOT_BASE_URL}/contacts/{self.portal_id}/record/0-5/{object_id}\"\n )\n elif object_type == \"companies\":\n return (\n f\"{HUBSPOT_BASE_URL}/contacts/{self.portal_id}/record/0-2/{object_id}\"\n )\n elif object_type == \"deals\":\n return (\n f\"{HUBSPOT_BASE_URL}/contacts/{self.portal_id}/record/0-3/{object_id}\"\n )\n elif object_type == \"contacts\":\n return (\n f\"{HUBSPOT_BASE_URL}/contacts/{self.portal_id}/record/0-1/{object_id}\"\n )\n elif object_type == \"notes\":\n return (\n f\"{HUBSPOT_BASE_URL}/contacts/{self.portal_id}/objects/0-4/{object_id}\"\n )\n else:\n return f\"{HUBSPOT_BASE_URL}/contacts/{self.portal_id}/{object_type}/{object_id}\"\n\n def _get_associated_objects(\n self,\n api_client: HubSpot,\n object_id: str,\n from_object_type: str,\n to_object_type: str,\n ) -> list[dict[str, Any]]:\n \"\"\"Get associated objects for a given object\"\"\"\n try:\n associations_iter = self._paginated_results(\n api_client.crm.associations.v4.basic_api.get_page,\n object_type=from_object_type,\n object_id=object_id,\n to_object_type=to_object_type,\n )\n\n object_ids = [assoc.to_object_id for assoc in associations_iter]\n\n associated_objects: list[dict[str, Any]] = []\n\n if to_object_type == \"contacts\":\n for obj_id in object_ids:\n try:\n obj = self._call_hubspot(\n api_client.crm.contacts.basic_api.get_by_id,\n contact_id=obj_id,\n properties=[\n \"firstname\",\n \"lastname\",\n \"email\",\n \"company\",\n \"jobtitle\",\n ],\n )\n associated_objects.append(obj.to_dict())\n except Exception as e:\n logger.warning(f\"Failed to fetch contact {obj_id}: {e}\")\n\n elif to_object_type == \"companies\":\n for obj_id in object_ids:\n try:\n obj = self._call_hubspot(\n api_client.crm.companies.basic_api.get_by_id,\n company_id=obj_id,\n properties=[\n \"name\",\n \"domain\",\n \"industry\",\n \"city\",\n \"state\",\n ],\n )\n associated_objects.append(obj.to_dict())\n except Exception as e:\n logger.warning(f\"Failed to fetch company {obj_id}: {e}\")\n\n elif to_object_type == \"deals\":\n for obj_id in object_ids:\n try:\n obj = self._call_hubspot(\n api_client.crm.deals.basic_api.get_by_id,\n deal_id=obj_id,\n properties=[\n \"dealname\",\n \"amount\",\n \"dealstage\",\n \"closedate\",\n \"pipeline\",\n ],\n )\n associated_objects.append(obj.to_dict())\n except Exception as e:\n logger.warning(f\"Failed to fetch deal {obj_id}: {e}\")\n\n elif to_object_type == \"tickets\":\n for obj_id in object_ids:\n try:\n obj = self._call_hubspot(\n api_client.crm.tickets.basic_api.get_by_id,\n ticket_id=obj_id,\n properties=[\"subject\", \"content\", \"hs_ticket_priority\"],\n )\n associated_objects.append(obj.to_dict())\n except Exception as e:\n logger.warning(f\"Failed to fetch ticket {obj_id}: {e}\")\n\n return associated_objects\n\n except Exception as e:\n logger.warning(\n f\"Failed to get associations from {from_object_type} to {to_object_type}: {e}\"\n )\n return []\n\n def _get_associated_notes(\n self,\n api_client: HubSpot,\n object_id: str,\n object_type: str,\n ) -> list[dict[str, Any]]:\n \"\"\"Get notes associated with a given object\"\"\"\n try:\n associations_iter = self._paginated_results(\n api_client.crm.associations.v4.basic_api.get_page,\n object_type=object_type,\n object_id=object_id,\n to_object_type=\"notes\",\n )\n\n note_ids = [assoc.to_object_id for assoc in associations_iter]\n\n associated_notes = []\n\n for note_id in note_ids:\n try:\n # Notes are engagements in HubSpot, use the engagements API\n note = self._call_hubspot(\n api_client.crm.objects.notes.basic_api.get_by_id,\n note_id=note_id,\n properties=[\n \"hs_note_body\",\n \"hs_timestamp\",\n \"hs_created_by\",\n \"hubspot_owner_id\",\n ],\n )\n associated_notes.append(note.to_dict())\n except Exception as e:\n logger.warning(f\"Failed to fetch note {note_id}: {e}\")\n\n return associated_notes\n\n except Exception as e:\n logger.warning(f\"Failed to get notes for {object_type} {object_id}: {e}\")\n return []\n\n def _create_object_section(\n self, obj: dict[str, Any], object_type: str\n ) -> TextSection:\n \"\"\"Create a TextSection for an associated object\"\"\"\n obj_id = obj.get(\"id\", \"\")\n properties = obj.get(\"properties\", {})\n\n if object_type == \"contacts\":\n name_parts = []\n if properties.get(\"firstname\"):\n name_parts.append(properties[\"firstname\"])\n if properties.get(\"lastname\"):\n name_parts.append(properties[\"lastname\"])\n\n if name_parts:\n name = \" \".join(name_parts)\n elif properties.get(\"email\"):\n # Use email as fallback if no first/last name\n name = properties[\"email\"]\n else:\n name = \"Unknown Contact\"\n\n content_parts = [f\"Contact: {name}\"]\n if properties.get(\"email\"):\n content_parts.append(f\"Email: {properties['email']}\")\n if properties.get(\"company\"):\n content_parts.append(f\"Company: {properties['company']}\")\n if properties.get(\"jobtitle\"):\n content_parts.append(f\"Job Title: {properties['jobtitle']}\")\n\n elif object_type == \"companies\":\n name = properties.get(\"name\", \"Unknown Company\")\n content_parts = [f\"Company: {name}\"]\n if properties.get(\"domain\"):\n content_parts.append(f\"Domain: {properties['domain']}\")\n if properties.get(\"industry\"):\n content_parts.append(f\"Industry: {properties['industry']}\")\n if properties.get(\"city\") and properties.get(\"state\"):\n content_parts.append(\n f\"Location: {properties['city']}, {properties['state']}\"\n )\n\n elif object_type == \"deals\":\n name = properties.get(\"dealname\", \"Unknown Deal\")\n content_parts = [f\"Deal: {name}\"]\n if properties.get(\"amount\"):\n content_parts.append(f\"Amount: ${properties['amount']}\")\n if properties.get(\"dealstage\"):\n content_parts.append(f\"Stage: {properties['dealstage']}\")\n if properties.get(\"closedate\"):\n content_parts.append(f\"Close Date: {properties['closedate']}\")\n if properties.get(\"pipeline\"):\n content_parts.append(f\"Pipeline: {properties['pipeline']}\")\n\n elif object_type == \"tickets\":\n name = properties.get(\"subject\", \"Unknown Ticket\")\n content_parts = [f\"Ticket: {name}\"]\n if properties.get(\"content\"):\n content_parts.append(f\"Content: {properties['content']}\")\n if properties.get(\"hs_ticket_priority\"):\n content_parts.append(f\"Priority: {properties['hs_ticket_priority']}\")\n elif object_type == \"notes\":\n # Notes have a body property that contains the note content\n body = properties.get(\"hs_note_body\", \"\")\n timestamp = properties.get(\"hs_timestamp\", \"\")\n\n # Clean HTML content to get raw text\n clean_body = self._clean_html_content(body)\n\n # Use full content, not truncated\n content_parts = [f\"Note: {clean_body}\"]\n if timestamp:\n content_parts.append(f\"Created: {timestamp}\")\n else:\n content_parts = [f\"{object_type.capitalize()}: {obj_id}\"]\n\n content = \"\\n\".join(content_parts)\n link = self._get_object_url(object_type, obj_id)\n\n return TextSection(link=link, text=content)\n\n def _process_tickets(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n api_client = HubSpot(access_token=self.access_token)\n\n tickets_iter = self._paginated_results(\n api_client.crm.tickets.basic_api.get_page,\n properties=[\n \"subject\",\n \"content\",\n \"hs_ticket_priority\",\n \"createdate\",\n \"hs_lastmodifieddate\",\n ],\n associations=[\"contacts\", \"companies\", \"deals\"],\n )\n\n doc_batch: list[Document | HierarchyNode] = []\n\n for ticket in tickets_iter:\n updated_at = ticket.updated_at.replace(tzinfo=None)\n if start is not None and updated_at < start.replace(tzinfo=None):\n continue\n if end is not None and updated_at > end.replace(tzinfo=None):\n continue\n\n title = ticket.properties.get(\"subject\") or f\"Ticket {ticket.id}\"\n link = self._get_object_url(\"tickets\", ticket.id)\n content_text = ticket.properties.get(\"content\") or \"\"\n\n # Main ticket section\n sections = [TextSection(link=link, text=content_text)]\n\n # Metadata with parent object IDs\n metadata: dict[str, str | list[str]] = {\n \"object_type\": \"ticket\",\n }\n\n if ticket.properties.get(\"hs_ticket_priority\"):\n metadata[\"priority\"] = ticket.properties[\"hs_ticket_priority\"]\n\n # Add associated objects as sections\n associated_contact_ids = []\n associated_company_ids = []\n associated_deal_ids = []\n\n # Get associated contacts\n associated_contacts = self._get_associated_objects(\n api_client, ticket.id, \"tickets\", \"contacts\"\n )\n for contact in associated_contacts:\n sections.append(self._create_object_section(contact, \"contacts\"))\n associated_contact_ids.append(contact[\"id\"])\n\n # Get associated companies\n associated_companies = self._get_associated_objects(\n api_client, ticket.id, \"tickets\", \"companies\"\n )\n for company in associated_companies:\n sections.append(self._create_object_section(company, \"companies\"))\n associated_company_ids.append(company[\"id\"])\n\n # Get associated deals\n associated_deals = self._get_associated_objects(\n api_client, ticket.id, \"tickets\", \"deals\"\n )\n for deal in associated_deals:\n sections.append(self._create_object_section(deal, \"deals\"))\n associated_deal_ids.append(deal[\"id\"])\n\n # Get associated notes\n associated_notes = self._get_associated_notes(\n api_client, ticket.id, \"tickets\"\n )\n for note in associated_notes:\n sections.append(self._create_object_section(note, \"notes\"))\n\n # Add association IDs to metadata\n if associated_contact_ids:\n metadata[\"associated_contact_ids\"] = associated_contact_ids\n if associated_company_ids:\n metadata[\"associated_company_ids\"] = associated_company_ids\n if associated_deal_ids:\n metadata[\"associated_deal_ids\"] = associated_deal_ids\n\n doc_batch.append(\n Document(\n id=f\"hubspot_ticket_{ticket.id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.HUBSPOT,\n semantic_identifier=title,\n doc_updated_at=ticket.updated_at.replace(tzinfo=timezone.utc),\n metadata=metadata,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [\"Tickets\"],\n \"object_type\": \"ticket\",\n \"object_id\": ticket.id,\n }\n },\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def _process_companies(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n api_client = HubSpot(access_token=self.access_token)\n\n companies_iter = self._paginated_results(\n api_client.crm.companies.basic_api.get_page,\n properties=[\n \"name\",\n \"domain\",\n \"industry\",\n \"city\",\n \"state\",\n \"description\",\n \"createdate\",\n \"hs_lastmodifieddate\",\n ],\n associations=[\"contacts\", \"deals\", \"tickets\"],\n )\n\n doc_batch: list[Document | HierarchyNode] = []\n\n for company in companies_iter:\n updated_at = company.updated_at.replace(tzinfo=None)\n if start is not None and updated_at < start.replace(tzinfo=None):\n continue\n if end is not None and updated_at > end.replace(tzinfo=None):\n continue\n\n title = company.properties.get(\"name\") or f\"Company {company.id}\"\n link = self._get_object_url(\"companies\", company.id)\n\n # Build main content\n content_parts = [f\"Company: {title}\"]\n if company.properties.get(\"domain\"):\n content_parts.append(f\"Domain: {company.properties['domain']}\")\n if company.properties.get(\"industry\"):\n content_parts.append(f\"Industry: {company.properties['industry']}\")\n if company.properties.get(\"city\") and company.properties.get(\"state\"):\n content_parts.append(\n f\"Location: {company.properties['city']}, {company.properties['state']}\"\n )\n if company.properties.get(\"description\"):\n content_parts.append(\n f\"Description: {company.properties['description']}\"\n )\n\n content_text = \"\\n\".join(content_parts)\n\n # Main company section\n sections = [TextSection(link=link, text=content_text)]\n\n # Metadata with parent object IDs\n metadata: dict[str, str | list[str]] = {\n \"company_id\": company.id,\n \"object_type\": \"company\",\n }\n\n if company.properties.get(\"industry\"):\n metadata[\"industry\"] = company.properties[\"industry\"]\n if company.properties.get(\"domain\"):\n metadata[\"domain\"] = company.properties[\"domain\"]\n\n # Add associated objects as sections\n associated_contact_ids = []\n associated_deal_ids = []\n associated_ticket_ids = []\n\n # Get associated contacts\n associated_contacts = self._get_associated_objects(\n api_client, company.id, \"companies\", \"contacts\"\n )\n for contact in associated_contacts:\n sections.append(self._create_object_section(contact, \"contacts\"))\n associated_contact_ids.append(contact[\"id\"])\n\n # Get associated deals\n associated_deals = self._get_associated_objects(\n api_client, company.id, \"companies\", \"deals\"\n )\n for deal in associated_deals:\n sections.append(self._create_object_section(deal, \"deals\"))\n associated_deal_ids.append(deal[\"id\"])\n\n # Get associated tickets\n associated_tickets = self._get_associated_objects(\n api_client, company.id, \"companies\", \"tickets\"\n )\n for ticket in associated_tickets:\n sections.append(self._create_object_section(ticket, \"tickets\"))\n associated_ticket_ids.append(ticket[\"id\"])\n\n # Get associated notes\n associated_notes = self._get_associated_notes(\n api_client, company.id, \"companies\"\n )\n for note in associated_notes:\n sections.append(self._create_object_section(note, \"notes\"))\n\n # Add association IDs to metadata\n if associated_contact_ids:\n metadata[\"associated_contact_ids\"] = associated_contact_ids\n if associated_deal_ids:\n metadata[\"associated_deal_ids\"] = associated_deal_ids\n if associated_ticket_ids:\n metadata[\"associated_ticket_ids\"] = associated_ticket_ids\n\n doc_batch.append(\n Document(\n id=f\"hubspot_company_{company.id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.HUBSPOT,\n semantic_identifier=title,\n doc_updated_at=company.updated_at.replace(tzinfo=timezone.utc),\n metadata=metadata,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [\"Companies\"],\n \"object_type\": \"company\",\n \"object_id\": company.id,\n }\n },\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def _process_deals(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n api_client = HubSpot(access_token=self.access_token)\n\n deals_iter = self._paginated_results(\n api_client.crm.deals.basic_api.get_page,\n properties=[\n \"dealname\",\n \"amount\",\n \"dealstage\",\n \"closedate\",\n \"pipeline\",\n \"description\",\n \"createdate\",\n \"hs_lastmodifieddate\",\n ],\n associations=[\"contacts\", \"companies\", \"tickets\"],\n )\n\n doc_batch: list[Document | HierarchyNode] = []\n\n for deal in deals_iter:\n updated_at = deal.updated_at.replace(tzinfo=None)\n if start is not None and updated_at < start.replace(tzinfo=None):\n continue\n if end is not None and updated_at > end.replace(tzinfo=None):\n continue\n\n title = deal.properties.get(\"dealname\") or f\"Deal {deal.id}\"\n link = self._get_object_url(\"deals\", deal.id)\n\n # Build main content\n content_parts = [f\"Deal: {title}\"]\n if deal.properties.get(\"amount\"):\n content_parts.append(f\"Amount: ${deal.properties['amount']}\")\n if deal.properties.get(\"dealstage\"):\n content_parts.append(f\"Stage: {deal.properties['dealstage']}\")\n if deal.properties.get(\"closedate\"):\n content_parts.append(f\"Close Date: {deal.properties['closedate']}\")\n if deal.properties.get(\"pipeline\"):\n content_parts.append(f\"Pipeline: {deal.properties['pipeline']}\")\n if deal.properties.get(\"description\"):\n content_parts.append(f\"Description: {deal.properties['description']}\")\n\n content_text = \"\\n\".join(content_parts)\n\n # Main deal section\n sections = [TextSection(link=link, text=content_text)]\n\n # Metadata with parent object IDs\n metadata: dict[str, str | list[str]] = {\n \"deal_id\": deal.id,\n \"object_type\": \"deal\",\n }\n\n if deal.properties.get(\"dealstage\"):\n metadata[\"deal_stage\"] = deal.properties[\"dealstage\"]\n if deal.properties.get(\"pipeline\"):\n metadata[\"pipeline\"] = deal.properties[\"pipeline\"]\n if deal.properties.get(\"amount\"):\n metadata[\"amount\"] = deal.properties[\"amount\"]\n\n # Add associated objects as sections\n associated_contact_ids = []\n associated_company_ids = []\n associated_ticket_ids = []\n\n # Get associated contacts\n associated_contacts = self._get_associated_objects(\n api_client, deal.id, \"deals\", \"contacts\"\n )\n for contact in associated_contacts:\n sections.append(self._create_object_section(contact, \"contacts\"))\n associated_contact_ids.append(contact[\"id\"])\n\n # Get associated companies\n associated_companies = self._get_associated_objects(\n api_client, deal.id, \"deals\", \"companies\"\n )\n for company in associated_companies:\n sections.append(self._create_object_section(company, \"companies\"))\n associated_company_ids.append(company[\"id\"])\n\n # Get associated tickets\n associated_tickets = self._get_associated_objects(\n api_client, deal.id, \"deals\", \"tickets\"\n )\n for ticket in associated_tickets:\n sections.append(self._create_object_section(ticket, \"tickets\"))\n associated_ticket_ids.append(ticket[\"id\"])\n\n # Get associated notes\n associated_notes = self._get_associated_notes(api_client, deal.id, \"deals\")\n for note in associated_notes:\n sections.append(self._create_object_section(note, \"notes\"))\n\n # Add association IDs to metadata\n if associated_contact_ids:\n metadata[\"associated_contact_ids\"] = associated_contact_ids\n if associated_company_ids:\n metadata[\"associated_company_ids\"] = associated_company_ids\n if associated_ticket_ids:\n metadata[\"associated_ticket_ids\"] = associated_ticket_ids\n\n doc_batch.append(\n Document(\n id=f\"hubspot_deal_{deal.id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.HUBSPOT,\n semantic_identifier=title,\n doc_updated_at=deal.updated_at.replace(tzinfo=timezone.utc),\n metadata=metadata,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [\"Deals\"],\n \"object_type\": \"deal\",\n \"object_id\": deal.id,\n }\n },\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def _process_contacts(\n self, start: datetime | None = None, end: datetime | None = None\n ) -> GenerateDocumentsOutput:\n api_client = HubSpot(access_token=self.access_token)\n\n contacts_iter = self._paginated_results(\n api_client.crm.contacts.basic_api.get_page,\n properties=[\n \"firstname\",\n \"lastname\",\n \"email\",\n \"company\",\n \"jobtitle\",\n \"phone\",\n \"city\",\n \"state\",\n \"createdate\",\n \"lastmodifieddate\",\n ],\n associations=[\"companies\", \"deals\", \"tickets\"],\n )\n\n doc_batch: list[Document | HierarchyNode] = []\n\n for contact in contacts_iter:\n updated_at = contact.updated_at.replace(tzinfo=None)\n if start is not None and updated_at < start.replace(tzinfo=None):\n continue\n if end is not None and updated_at > end.replace(tzinfo=None):\n continue\n\n # Build contact name\n name_parts = []\n if contact.properties.get(\"firstname\"):\n name_parts.append(contact.properties[\"firstname\"])\n if contact.properties.get(\"lastname\"):\n name_parts.append(contact.properties[\"lastname\"])\n\n if name_parts:\n title = \" \".join(name_parts)\n elif contact.properties.get(\"email\"):\n # Use email as fallback if no first/last name\n title = contact.properties[\"email\"]\n else:\n title = f\"Contact {contact.id}\"\n\n link = self._get_object_url(\"contacts\", contact.id)\n\n # Build main content\n content_parts = [f\"Contact: {title}\"]\n if contact.properties.get(\"email\"):\n content_parts.append(f\"Email: {contact.properties['email']}\")\n if contact.properties.get(\"company\"):\n content_parts.append(f\"Company: {contact.properties['company']}\")\n if contact.properties.get(\"jobtitle\"):\n content_parts.append(f\"Job Title: {contact.properties['jobtitle']}\")\n if contact.properties.get(\"phone\"):\n content_parts.append(f\"Phone: {contact.properties['phone']}\")\n if contact.properties.get(\"city\") and contact.properties.get(\"state\"):\n content_parts.append(\n f\"Location: {contact.properties['city']}, {contact.properties['state']}\"\n )\n\n content_text = \"\\n\".join(content_parts)\n\n # Main contact section\n sections = [TextSection(link=link, text=content_text)]\n\n # Metadata with parent object IDs\n metadata: dict[str, str | list[str]] = {\n \"contact_id\": contact.id,\n \"object_type\": \"contact\",\n }\n\n if contact.properties.get(\"email\"):\n metadata[\"email\"] = contact.properties[\"email\"]\n if contact.properties.get(\"company\"):\n metadata[\"company\"] = contact.properties[\"company\"]\n if contact.properties.get(\"jobtitle\"):\n metadata[\"job_title\"] = contact.properties[\"jobtitle\"]\n\n # Add associated objects as sections\n associated_company_ids = []\n associated_deal_ids = []\n associated_ticket_ids = []\n\n # Get associated companies\n associated_companies = self._get_associated_objects(\n api_client, contact.id, \"contacts\", \"companies\"\n )\n for company in associated_companies:\n sections.append(self._create_object_section(company, \"companies\"))\n associated_company_ids.append(company[\"id\"])\n\n # Get associated deals\n associated_deals = self._get_associated_objects(\n api_client, contact.id, \"contacts\", \"deals\"\n )\n for deal in associated_deals:\n sections.append(self._create_object_section(deal, \"deals\"))\n associated_deal_ids.append(deal[\"id\"])\n\n # Get associated tickets\n associated_tickets = self._get_associated_objects(\n api_client, contact.id, \"contacts\", \"tickets\"\n )\n for ticket in associated_tickets:\n sections.append(self._create_object_section(ticket, \"tickets\"))\n associated_ticket_ids.append(ticket[\"id\"])\n\n # Get associated notes\n associated_notes = self._get_associated_notes(\n api_client, contact.id, \"contacts\"\n )\n for note in associated_notes:\n sections.append(self._create_object_section(note, \"notes\"))\n\n # Add association IDs to metadata\n if associated_company_ids:\n metadata[\"associated_company_ids\"] = associated_company_ids\n if associated_deal_ids:\n metadata[\"associated_deal_ids\"] = associated_deal_ids\n if associated_ticket_ids:\n metadata[\"associated_ticket_ids\"] = associated_ticket_ids\n\n doc_batch.append(\n Document(\n id=f\"hubspot_contact_{contact.id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.HUBSPOT,\n semantic_identifier=title,\n doc_updated_at=contact.updated_at.replace(tzinfo=timezone.utc),\n metadata=metadata,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [\"Contacts\"],\n \"object_type\": \"contact\",\n \"object_id\": contact.id,\n }\n },\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"Load all HubSpot objects (tickets, companies, deals, contacts)\"\"\"\n # Process each object type based on configuration\n if \"tickets\" in self.object_types:\n yield from self._process_tickets()\n if \"companies\" in self.object_types:\n yield from self._process_companies()\n if \"deals\" in self.object_types:\n yield from self._process_deals()\n if \"contacts\" in self.object_types:\n yield from self._process_contacts()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)\n end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n\n # Process each object type with time filtering based on configuration\n if \"tickets\" in self.object_types:\n yield from self._process_tickets(start_datetime, end_datetime)\n if \"companies\" in self.object_types:\n yield from self._process_companies(start_datetime, end_datetime)\n if \"deals\" in self.object_types:\n yield from self._process_deals(start_datetime, end_datetime)\n if \"contacts\" in self.object_types:\n yield from self._process_contacts(start_datetime, end_datetime)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = HubSpotConnector()\n connector.load_credentials(\n {\"hubspot_access_token\": os.environ[\"HUBSPOT_ACCESS_TOKEN\"]}\n )\n # Run the first example\n document_batches = connector.load_from_state()\n first_batch = next(document_batches)\n for doc in first_batch:\n print(doc.model_dump_json(indent=2))\n" }, { "path": "backend/onyx/connectors/hubspot/rate_limit.py", "content": "from __future__ import annotations\n\nimport time\nfrom collections.abc import Callable\nfrom typing import Any\nfrom typing import TypeVar\n\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n RateLimitTriedTooManyTimesError,\n)\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nT = TypeVar(\"T\")\n\n# HubSpot exposes a ten second rolling window (x-hubspot-ratelimit-interval-milliseconds)\n# with a maximum of 190 requests, and a per-second limit of 19 requests.\n_HUBSPOT_TEN_SECOND_LIMIT = 190\n_HUBSPOT_TEN_SECOND_PERIOD = 10 # seconds\n_HUBSPOT_SECONDLY_LIMIT = 19\n_HUBSPOT_SECONDLY_PERIOD = 1 # second\n_DEFAULT_SLEEP_SECONDS = 10\n_SLEEP_PADDING_SECONDS = 1.0\n_MAX_RATE_LIMIT_RETRIES = 5\n\n\ndef _extract_header(headers: Any, key: str) -> str | None:\n if headers is None:\n return None\n\n getter = getattr(headers, \"get\", None)\n if callable(getter):\n value = getter(key)\n if value is not None:\n return value\n\n if isinstance(headers, dict):\n value = headers.get(key)\n if value is not None:\n return value\n\n return None\n\n\ndef is_rate_limit_error(exception: Exception) -> bool:\n status = getattr(exception, \"status\", None)\n if status == 429:\n return True\n\n headers = getattr(exception, \"headers\", None)\n if headers is not None:\n remaining = _extract_header(headers, \"x-hubspot-ratelimit-remaining\")\n if remaining == \"0\":\n return True\n secondly_remaining = _extract_header(\n headers, \"x-hubspot-ratelimit-secondly-remaining\"\n )\n if secondly_remaining == \"0\":\n return True\n\n message = str(exception)\n return \"RATE_LIMIT\" in message or \"Too Many Requests\" in message\n\n\ndef get_rate_limit_retry_delay_seconds(exception: Exception) -> float:\n headers = getattr(exception, \"headers\", None)\n\n retry_after = _extract_header(headers, \"Retry-After\")\n if retry_after:\n try:\n return float(retry_after) + _SLEEP_PADDING_SECONDS\n except ValueError:\n logger.debug(\n \"Failed to parse Retry-After header '%s' as float\", retry_after\n )\n\n interval_ms = _extract_header(headers, \"x-hubspot-ratelimit-interval-milliseconds\")\n if interval_ms:\n try:\n return float(interval_ms) / 1000.0 + _SLEEP_PADDING_SECONDS\n except ValueError:\n logger.debug(\n \"Failed to parse x-hubspot-ratelimit-interval-milliseconds '%s' as float\",\n interval_ms,\n )\n\n secondly_limit = _extract_header(headers, \"x-hubspot-ratelimit-secondly\")\n if secondly_limit:\n try:\n per_second = max(float(secondly_limit), 1.0)\n return (1.0 / per_second) + _SLEEP_PADDING_SECONDS\n except ValueError:\n logger.debug(\n \"Failed to parse x-hubspot-ratelimit-secondly '%s' as float\",\n secondly_limit,\n )\n\n return _DEFAULT_SLEEP_SECONDS + _SLEEP_PADDING_SECONDS\n\n\nclass HubSpotRateLimiter:\n def __init__(\n self,\n *,\n ten_second_limit: int = _HUBSPOT_TEN_SECOND_LIMIT,\n ten_second_period: int = _HUBSPOT_TEN_SECOND_PERIOD,\n secondly_limit: int = _HUBSPOT_SECONDLY_LIMIT,\n secondly_period: int = _HUBSPOT_SECONDLY_PERIOD,\n max_retries: int = _MAX_RATE_LIMIT_RETRIES,\n ) -> None:\n self._max_retries = max_retries\n\n @rate_limit_builder(max_calls=secondly_limit, period=secondly_period)\n @rate_limit_builder(max_calls=ten_second_limit, period=ten_second_period)\n def _execute(callable_: Callable[[], T]) -> T:\n return callable_()\n\n self._execute = _execute\n\n def call(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> T:\n attempts = 0\n\n while True:\n try:\n return self._execute(lambda: func(*args, **kwargs))\n except Exception as exc: # pylint: disable=broad-except\n if not is_rate_limit_error(exc):\n raise\n\n attempts += 1\n if attempts > self._max_retries:\n raise RateLimitTriedTooManyTimesError(\n \"Exceeded configured HubSpot rate limit retries\"\n ) from exc\n\n wait_time = get_rate_limit_retry_delay_seconds(exc)\n logger.notice(\n \"HubSpot rate limit reached. Sleeping %.2f seconds before retrying.\",\n wait_time,\n )\n time.sleep(wait_time)\n" }, { "path": "backend/onyx/connectors/imap/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/imap/connector.py", "content": "import copy\nimport email\nimport imaplib\nimport os\nimport re\nfrom datetime import datetime\nfrom datetime import timezone\nfrom email.message import Message\nfrom email.utils import parseaddr\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\n\nimport bs4\nfrom pydantic import BaseModel\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.imap.models import EmailHeaders\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import CredentialsConnector\nfrom onyx.connectors.interfaces import CredentialsProviderInterface\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n_DEFAULT_IMAP_PORT_NUMBER = int(os.environ.get(\"IMAP_PORT\", 993))\n_IMAP_OKAY_STATUS = \"OK\"\n_PAGE_SIZE = 100\n_USERNAME_KEY = \"imap_username\"\n_PASSWORD_KEY = \"imap_password\"\n\n\nclass CurrentMailbox(BaseModel):\n mailbox: str\n todo_email_ids: list[str]\n\n\n# An email has a list of mailboxes.\n# Each mailbox has a list of email-ids inside of it.\n#\n# Usage:\n# To use this checkpointer, first fetch all the mailboxes.\n# Then, pop a mailbox and fetch all of its email-ids.\n# Then, pop each email-id and fetch its content (and parse it, etc..).\n# When you have popped all email-ids for this mailbox, pop the next mailbox and repeat the above process until you're done.\n#\n# For initial checkpointing, set both fields to `None`.\nclass ImapCheckpoint(ConnectorCheckpoint):\n todo_mailboxes: list[str] | None = None\n current_mailbox: CurrentMailbox | None = None\n\n\nclass LoginState(str, Enum):\n LoggedIn = \"logged_in\"\n LoggedOut = \"logged_out\"\n\n\nclass ImapConnector(\n CredentialsConnector,\n CheckpointedConnectorWithPermSync[ImapCheckpoint],\n):\n def __init__(\n self,\n host: str,\n port: int = _DEFAULT_IMAP_PORT_NUMBER,\n mailboxes: list[str] | None = None,\n ) -> None:\n self._host = host\n self._port = port\n self._mailboxes = mailboxes\n self._credentials: dict[str, Any] | None = None\n\n @property\n def credentials(self) -> dict[str, Any]:\n if not self._credentials:\n raise RuntimeError(\n \"Credentials have not been initialized; call `set_credentials_provider` first\"\n )\n return self._credentials\n\n def _get_mail_client(self) -> imaplib.IMAP4_SSL:\n \"\"\"\n Returns a new `imaplib.IMAP4_SSL` instance.\n\n The `imaplib.IMAP4_SSL` object is supposed to be an \"ephemeral\" object; it's not something that you can login,\n logout, then log back into again. I.e., the following will fail:\n\n ```py\n mail_client.login(..)\n mail_client.logout();\n mail_client.login(..)\n ```\n\n Therefore, you need a fresh, new instance in order to operate with IMAP. This function gives one to you.\n\n # Notes\n This function will throw an error if the credentials have not yet been set.\n \"\"\"\n\n def get_or_raise(name: str) -> str:\n value = self.credentials.get(name)\n if not value:\n raise RuntimeError(f\"Credential item {name=} was not found\")\n if not isinstance(value, str):\n raise RuntimeError(\n f\"Credential item {name=} must be of type str, instead received {type(name)=}\"\n )\n return value\n\n username = get_or_raise(_USERNAME_KEY)\n password = get_or_raise(_PASSWORD_KEY)\n\n mail_client = imaplib.IMAP4_SSL(host=self._host, port=self._port)\n status, _data = mail_client.login(user=username, password=password)\n\n if status != _IMAP_OKAY_STATUS:\n raise RuntimeError(f\"Failed to log into imap server; {status=}\")\n\n return mail_client\n\n def _load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: ImapCheckpoint,\n include_perm_sync: bool,\n ) -> CheckpointOutput[ImapCheckpoint]:\n checkpoint = cast(ImapCheckpoint, copy.deepcopy(checkpoint))\n checkpoint.has_more = True\n\n mail_client = self._get_mail_client()\n\n if checkpoint.todo_mailboxes is None:\n # This is the dummy checkpoint.\n # Fill it with mailboxes first.\n if self._mailboxes:\n checkpoint.todo_mailboxes = _sanitize_mailbox_names(self._mailboxes)\n else:\n fetched_mailboxes = _fetch_all_mailboxes_for_email_account(\n mail_client=mail_client\n )\n if not fetched_mailboxes:\n raise RuntimeError(\n \"Failed to find any mailboxes for this email account\"\n )\n checkpoint.todo_mailboxes = _sanitize_mailbox_names(fetched_mailboxes)\n\n return checkpoint\n\n if (\n not checkpoint.current_mailbox\n or not checkpoint.current_mailbox.todo_email_ids\n ):\n if not checkpoint.todo_mailboxes:\n checkpoint.has_more = False\n return checkpoint\n\n mailbox = checkpoint.todo_mailboxes.pop()\n email_ids = _fetch_email_ids_in_mailbox(\n mail_client=mail_client,\n mailbox=mailbox,\n start=start,\n end=end,\n )\n checkpoint.current_mailbox = CurrentMailbox(\n mailbox=mailbox,\n todo_email_ids=email_ids,\n )\n\n _select_mailbox(\n mail_client=mail_client, mailbox=checkpoint.current_mailbox.mailbox\n )\n current_todos = cast(\n list, copy.deepcopy(checkpoint.current_mailbox.todo_email_ids[:_PAGE_SIZE])\n )\n checkpoint.current_mailbox.todo_email_ids = (\n checkpoint.current_mailbox.todo_email_ids[_PAGE_SIZE:]\n )\n\n for email_id in current_todos:\n email_msg = _fetch_email(mail_client=mail_client, email_id=email_id)\n if not email_msg:\n logger.warn(f\"Failed to fetch message {email_id=}; skipping\")\n continue\n\n email_headers = EmailHeaders.from_email_msg(email_msg=email_msg)\n\n yield _convert_email_headers_and_body_into_document(\n email_msg=email_msg,\n email_headers=email_headers,\n include_perm_sync=include_perm_sync,\n )\n\n return checkpoint\n\n # impls for BaseConnector\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n raise NotImplementedError(\"Use `set_credentials_provider` instead\")\n\n def validate_connector_settings(self) -> None:\n self._get_mail_client()\n\n # impls for CredentialsConnector\n\n def set_credentials_provider(\n self, credentials_provider: CredentialsProviderInterface\n ) -> None:\n self._credentials = credentials_provider.get_credentials()\n\n # impls for CheckpointedConnector\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: ImapCheckpoint,\n ) -> CheckpointOutput[ImapCheckpoint]:\n return self._load_from_checkpoint(\n start=start, end=end, checkpoint=checkpoint, include_perm_sync=False\n )\n\n def build_dummy_checkpoint(self) -> ImapCheckpoint:\n return ImapCheckpoint(has_more=True)\n\n def validate_checkpoint_json(self, checkpoint_json: str) -> ImapCheckpoint:\n return ImapCheckpoint.model_validate_json(json_data=checkpoint_json)\n\n # impls for CheckpointedConnectorWithPermSync\n\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: ImapCheckpoint,\n ) -> CheckpointOutput[ImapCheckpoint]:\n return self._load_from_checkpoint(\n start=start, end=end, checkpoint=checkpoint, include_perm_sync=True\n )\n\n\ndef _fetch_all_mailboxes_for_email_account(mail_client: imaplib.IMAP4_SSL) -> list[str]:\n status, mailboxes_data = mail_client.list(directory=\"*\", pattern=\"*\")\n if status != _IMAP_OKAY_STATUS:\n raise RuntimeError(f\"Failed to fetch mailboxes; {status=}\")\n\n mailboxes = []\n\n for mailboxes_raw in mailboxes_data:\n if isinstance(mailboxes_raw, bytes):\n mailboxes_str = mailboxes_raw.decode()\n elif isinstance(mailboxes_raw, str):\n mailboxes_str = mailboxes_raw\n else:\n logger.warn(\n f\"Expected the mailbox data to be of type str, instead got {type(mailboxes_raw)=} {mailboxes_raw}; skipping\"\n )\n continue\n\n # The mailbox LIST response output can be found here:\n # https://www.rfc-editor.org/rfc/rfc3501.html#section-7.2.2\n #\n # The general format is:\n # `() `\n #\n # The below regex matches on that pattern; from there, we select the 3rd match (index 2), which is the mailbox-name.\n match = re.match(r'\\([^)]*\\)\\s+\"([^\"]+)\"\\s+\"?(.+?)\"?$', mailboxes_str)\n if not match:\n logger.warn(\n f\"Invalid mailbox-data formatting structure: {mailboxes_str=}; skipping\"\n )\n continue\n\n mailbox = match.group(2)\n mailboxes.append(mailbox)\n\n return mailboxes\n\n\ndef _select_mailbox(mail_client: imaplib.IMAP4_SSL, mailbox: str) -> None:\n status, _ids = mail_client.select(mailbox=mailbox, readonly=True)\n if status != _IMAP_OKAY_STATUS:\n raise RuntimeError(f\"Failed to select {mailbox=}\")\n\n\ndef _fetch_email_ids_in_mailbox(\n mail_client: imaplib.IMAP4_SSL,\n mailbox: str,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n) -> list[str]:\n _select_mailbox(mail_client=mail_client, mailbox=mailbox)\n\n start_str = datetime.fromtimestamp(start, tz=timezone.utc).strftime(\"%d-%b-%Y\")\n end_str = datetime.fromtimestamp(end, tz=timezone.utc).strftime(\"%d-%b-%Y\")\n search_criteria = f'(SINCE \"{start_str}\" BEFORE \"{end_str}\")'\n\n status, email_ids_byte_array = mail_client.search(None, search_criteria)\n\n if status != _IMAP_OKAY_STATUS or not email_ids_byte_array:\n raise RuntimeError(f\"Failed to fetch email ids; {status=}\")\n\n email_ids: bytes = email_ids_byte_array[0]\n\n return [email_id.decode() for email_id in email_ids.split()]\n\n\ndef _fetch_email(mail_client: imaplib.IMAP4_SSL, email_id: str) -> Message | None:\n status, msg_data = mail_client.fetch(message_set=email_id, message_parts=\"(RFC822)\")\n if status != _IMAP_OKAY_STATUS or not msg_data:\n return None\n\n data = msg_data[0]\n if not isinstance(data, tuple):\n raise RuntimeError(\n f\"Message data should be a tuple; instead got a {type(data)=} {data=}\"\n )\n\n _metadata, raw_email = data\n return email.message_from_bytes(raw_email)\n\n\ndef _convert_email_headers_and_body_into_document(\n email_msg: Message,\n email_headers: EmailHeaders,\n include_perm_sync: bool,\n) -> Document:\n sender_name, sender_addr = _parse_singular_addr(raw_header=email_headers.sender)\n parsed_recipients = (\n _parse_addrs(raw_header=email_headers.recipients)\n if email_headers.recipients\n else []\n )\n\n expert_info_map = {\n recipient_addr: BasicExpertInfo(\n display_name=recipient_name, email=recipient_addr\n )\n for recipient_name, recipient_addr in parsed_recipients\n }\n if sender_addr not in expert_info_map:\n expert_info_map[sender_addr] = BasicExpertInfo(\n display_name=sender_name, email=sender_addr\n )\n\n email_body = _parse_email_body(email_msg=email_msg, email_headers=email_headers)\n primary_owners = list(expert_info_map.values())\n external_access = (\n ExternalAccess(\n external_user_emails=set(expert_info_map.keys()),\n external_user_group_ids=set(),\n is_public=False,\n )\n if include_perm_sync\n else None\n )\n\n return Document(\n id=email_headers.id,\n title=email_headers.subject,\n semantic_identifier=email_headers.subject,\n metadata={},\n source=DocumentSource.IMAP,\n sections=[TextSection(text=email_body)],\n primary_owners=primary_owners,\n external_access=external_access,\n )\n\n\ndef _parse_email_body(\n email_msg: Message,\n email_headers: EmailHeaders,\n) -> str:\n body = None\n for part in email_msg.walk():\n if part.is_multipart():\n # Multipart parts are *containers* for other parts, not the actual content itself.\n # Therefore, we skip until we find the individual parts instead.\n continue\n\n charset = part.get_content_charset() or \"utf-8\"\n\n try:\n raw_payload = part.get_payload(decode=True)\n if not isinstance(raw_payload, bytes):\n logger.warn(\n \"Payload section from email was expected to be an array of bytes, instead got \"\n f\"{type(raw_payload)=}, {raw_payload=}\"\n )\n continue\n body = raw_payload.decode(charset)\n break\n except (UnicodeDecodeError, LookupError) as e:\n print(f\"Warning: Could not decode part with charset {charset}. Error: {e}\")\n continue\n\n if not body:\n logger.warn(\n f\"Email with {email_headers.id=} has an empty body; returning an empty string\"\n )\n return \"\"\n\n soup = bs4.BeautifulSoup(markup=body, features=\"html.parser\")\n\n return \" \".join(str_section for str_section in soup.stripped_strings)\n\n\ndef _sanitize_mailbox_names(mailboxes: list[str]) -> list[str]:\n \"\"\"\n Mailboxes with special characters in them must be enclosed by double-quotes, as per the IMAP protocol.\n Just to be safe, we wrap *all* mailboxes with double-quotes.\n \"\"\"\n return [f'\"{mailbox}\"' for mailbox in mailboxes if mailbox]\n\n\ndef _parse_addrs(raw_header: str) -> list[tuple[str, str]]:\n addrs = raw_header.split(\",\")\n name_addr_pairs = [parseaddr(addr=addr) for addr in addrs if addr]\n return [(name, addr) for name, addr in name_addr_pairs if addr]\n\n\ndef _parse_singular_addr(raw_header: str) -> tuple[str, str]:\n addrs = _parse_addrs(raw_header=raw_header)\n if not addrs:\n raise RuntimeError(\n f\"Parsing email header resulted in no addresses being found; {raw_header=}\"\n )\n elif len(addrs) >= 2:\n raise RuntimeError(\n f\"Expected a singular address, but instead got multiple; {raw_header=} {addrs=}\"\n )\n\n return addrs[0]\n\n\nif __name__ == \"__main__\":\n import time\n from tests.daily.connectors.utils import load_all_from_connector\n from onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\n\n host = os.environ.get(\"IMAP_HOST\")\n mailboxes_str = os.environ.get(\"IMAP_MAILBOXES\")\n username = os.environ.get(\"IMAP_USERNAME\")\n password = os.environ.get(\"IMAP_PASSWORD\")\n\n mailboxes = (\n [mailbox.strip() for mailbox in mailboxes_str.split(\",\")]\n if mailboxes_str\n else []\n )\n\n if not host:\n raise RuntimeError(\"`IMAP_HOST` must be set\")\n\n imap_connector = ImapConnector(\n host=host,\n mailboxes=mailboxes,\n )\n\n imap_connector.set_credentials_provider(\n OnyxStaticCredentialsProvider(\n tenant_id=None,\n connector_name=DocumentSource.IMAP,\n credential_json={\n _USERNAME_KEY: username,\n _PASSWORD_KEY: password,\n },\n )\n )\n\n for doc in load_all_from_connector(\n connector=imap_connector,\n start=0,\n end=time.time(),\n ).documents:\n print(doc)\n" }, { "path": "backend/onyx/connectors/imap/models.py", "content": "import email\nfrom datetime import datetime\nfrom email.message import Message\nfrom enum import Enum\n\nfrom pydantic import BaseModel\n\n\nclass Header(str, Enum):\n SUBJECT_HEADER = \"subject\"\n FROM_HEADER = \"from\"\n TO_HEADER = \"to\"\n DELIVERED_TO_HEADER = (\n \"Delivered-To\" # Used in mailing lists instead of the \"to\" header.\n )\n DATE_HEADER = \"date\"\n MESSAGE_ID_HEADER = \"Message-ID\"\n\n\nclass EmailHeaders(BaseModel):\n \"\"\"\n Model for email headers extracted from IMAP messages.\n \"\"\"\n\n id: str\n subject: str\n sender: str\n recipients: str | None\n date: datetime\n\n @classmethod\n def from_email_msg(cls, email_msg: Message) -> \"EmailHeaders\":\n def _decode(header: str, default: str | None = None) -> str | None:\n value = email_msg.get(header, default)\n if not value:\n return None\n\n decoded_value, encoding = email.header.decode_header(value)[0]\n if isinstance(decoded_value, bytes):\n encoding = encoding or \"utf-8\"\n return decoded_value.decode(encoding, errors=\"replace\")\n elif isinstance(decoded_value, str):\n return decoded_value\n else:\n return None\n\n def _parse_date(date_str: str | None) -> datetime | None:\n if not date_str:\n return None\n try:\n return email.utils.parsedate_to_datetime(date_str)\n except (TypeError, ValueError):\n return None\n\n message_id = _decode(header=Header.MESSAGE_ID_HEADER)\n # It's possible for the subject line to not exist or be an empty string.\n subject = _decode(header=Header.SUBJECT_HEADER) or \"Unknown Subject\"\n from_ = _decode(header=Header.FROM_HEADER)\n to = _decode(header=Header.TO_HEADER)\n if not to:\n to = _decode(header=Header.DELIVERED_TO_HEADER)\n date_str = _decode(header=Header.DATE_HEADER)\n date = _parse_date(date_str=date_str)\n\n # If any of the above are `None`, model validation will fail.\n # Therefore, no guards (i.e.: `if
is None: raise RuntimeError(..)`) were written.\n return cls.model_validate(\n {\n \"id\": message_id,\n \"subject\": subject,\n \"sender\": from_,\n \"recipients\": to,\n \"date\": date,\n }\n )\n" }, { "path": "backend/onyx/connectors/interfaces.py", "content": "import abc\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\nfrom types import TracebackType\nfrom typing import Any\nfrom typing import Generic\nfrom typing import TypeAlias\nfrom typing import TypeVar\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\nSecondsSinceUnixEpoch = float\n\n# Output types that can include HierarchyNode alongside Documents/SlimDocuments\nGenerateDocumentsOutput = Iterator[list[Document | HierarchyNode]]\nGenerateSlimDocumentOutput = Iterator[list[SlimDocument | HierarchyNode]]\n\nCT = TypeVar(\"CT\", bound=ConnectorCheckpoint)\n\n\nclass NormalizationResult(BaseModel):\n \"\"\"Result of URL normalization attempt.\n\n Attributes:\n normalized_url: The normalized URL string, or None if normalization failed\n use_default: If True, fall back to default normalizer. If False, return None.\n \"\"\"\n\n normalized_url: str | None\n use_default: bool = False\n\n\nclass BaseConnector(abc.ABC, Generic[CT]):\n REDIS_KEY_PREFIX = \"da_connector_data:\"\n\n @abc.abstractmethod\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n raise NotImplementedError\n\n @staticmethod\n def parse_metadata(metadata: dict[str, Any]) -> list[str]:\n \"\"\"Parse the metadata for a document/chunk into a string to pass to Generative AI as additional context\"\"\"\n custom_parser_req_msg = (\n \"Specific metadata parsing required, connector has not implemented it.\"\n )\n metadata_lines = []\n for metadata_key, metadata_value in metadata.items():\n if isinstance(metadata_value, str):\n metadata_lines.append(f\"{metadata_key}: {metadata_value}\")\n elif isinstance(metadata_value, list):\n if not all([isinstance(val, str) for val in metadata_value]):\n raise RuntimeError(custom_parser_req_msg)\n metadata_lines.append(f\"{metadata_key}: {', '.join(metadata_value)}\")\n else:\n raise RuntimeError(custom_parser_req_msg)\n return metadata_lines\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n Override this if your connector needs to validate credentials or settings.\n Raise an exception if invalid, otherwise do nothing.\n\n Default is a no-op (always successful).\n \"\"\"\n\n def validate_perm_sync(self) -> None:\n \"\"\"\n Don't override this; add a function to perm_sync_valid.py in the ee package\n to do permission sync validation\n \"\"\"\n validate_connector_settings_fn = fetch_ee_implementation_or_noop(\n \"onyx.connectors.perm_sync_valid\",\n \"validate_perm_sync\",\n noop_return_value=None,\n )\n validate_connector_settings_fn(self)\n\n def set_allow_images(self, value: bool) -> None:\n \"\"\"Implement if the underlying connector wants to skip/allow image downloading\n based on the application level image analysis setting.\"\"\"\n\n @classmethod\n def normalize_url(cls, url: str) -> \"NormalizationResult\": # noqa: ARG003\n \"\"\"Normalize a URL to match the canonical Document.id format used during ingestion.\n\n Connectors that use URLs as document IDs should override this method.\n Returns NormalizationResult with use_default=True if not implemented.\n \"\"\"\n return NormalizationResult(normalized_url=None, use_default=True)\n\n def build_dummy_checkpoint(self) -> CT:\n # TODO: find a way to make this work without type: ignore\n return ConnectorCheckpoint(has_more=True) # type: ignore\n\n\n# Large set update or reindex, generally pulling a complete state or from a savestate file\nclass LoadConnector(BaseConnector):\n @abc.abstractmethod\n def load_from_state(self) -> GenerateDocumentsOutput:\n raise NotImplementedError\n\n\n# Small set updates by time\nclass PollConnector(BaseConnector):\n @abc.abstractmethod\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n raise NotImplementedError\n\n\n# Slim connectors retrieve just the ids of documents\nclass SlimConnector(BaseConnector):\n @abc.abstractmethod\n def retrieve_all_slim_docs(\n self,\n ) -> GenerateSlimDocumentOutput:\n raise NotImplementedError\n\n\n# Slim connectors retrieve both the ids AND\n# permission syncing information for connected documents\nclass SlimConnectorWithPermSync(BaseConnector):\n @abc.abstractmethod\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n raise NotImplementedError\n\n\nclass OAuthConnector(BaseConnector):\n class AdditionalOauthKwargs(BaseModel):\n # if overridden, all fields should be str type\n pass\n\n @classmethod\n @abc.abstractmethod\n def oauth_id(cls) -> DocumentSource:\n raise NotImplementedError\n\n @classmethod\n @abc.abstractmethod\n def oauth_authorization_url(\n cls,\n base_domain: str,\n state: str,\n additional_kwargs: dict[str, str],\n ) -> str:\n raise NotImplementedError\n\n @classmethod\n @abc.abstractmethod\n def oauth_code_to_token(\n cls,\n base_domain: str,\n code: str,\n additional_kwargs: dict[str, str],\n ) -> dict[str, Any]:\n raise NotImplementedError\n\n\nT = TypeVar(\"T\", bound=\"CredentialsProviderInterface\")\n\n\nclass CredentialsProviderInterface(abc.ABC, Generic[T]):\n @abc.abstractmethod\n def __enter__(self) -> T:\n raise NotImplementedError\n\n @abc.abstractmethod\n def __exit__(\n self,\n exc_type: type[BaseException] | None,\n exc_value: BaseException | None,\n traceback: TracebackType | None,\n ) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def get_tenant_id(self) -> str | None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def get_provider_key(self) -> str:\n \"\"\"a unique key that the connector can use to lock around a credential\n that might be used simultaneously.\n\n Will typically be the credential id, but can also just be something random\n in cases when there is nothing to lock (aka static credentials)\n \"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def get_credentials(self) -> dict[str, Any]:\n raise NotImplementedError\n\n @abc.abstractmethod\n def set_credentials(self, credential_json: dict[str, Any]) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def is_dynamic(self) -> bool:\n \"\"\"If dynamic, the credentials may change during usage ... meaning the client\n needs to use the locking features of the credentials provider to operate\n correctly.\n\n If static, the client can simply reference the credentials once and use them\n through the entire indexing run.\n \"\"\"\n raise NotImplementedError\n\n\nclass CredentialsConnector(BaseConnector):\n \"\"\"Implement this if the connector needs to be able to read and write credentials\n on the fly. Typically used with shared credentials/tokens that might be renewed\n at any time.\"\"\"\n\n @abc.abstractmethod\n def set_credentials_provider(\n self, credentials_provider: CredentialsProviderInterface\n ) -> None:\n raise NotImplementedError\n\n\n# Event driven\nclass EventConnector(BaseConnector):\n @abc.abstractmethod\n def handle_event(self, event: Any) -> GenerateDocumentsOutput:\n raise NotImplementedError\n\n\nCheckpointOutput: TypeAlias = Generator[\n Document | HierarchyNode | ConnectorFailure, None, CT\n]\n\nHierarchyOutput: TypeAlias = Generator[HierarchyNode, None, None]\n\n\nclass CheckpointedConnector(BaseConnector[CT]):\n @abc.abstractmethod\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: CT,\n ) -> CheckpointOutput[CT]:\n \"\"\"Yields back documents or failures. Final return is the new checkpoint.\n\n Final return can be access via either:\n\n ```\n try:\n for document_or_failure in connector.load_from_checkpoint(start, end, checkpoint):\n print(document_or_failure)\n except StopIteration as e:\n checkpoint = e.value # Extracting the return value\n print(checkpoint)\n ```\n\n OR\n\n ```\n checkpoint = yield from connector.load_from_checkpoint(start, end, checkpoint)\n ```\n \"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def build_dummy_checkpoint(self) -> CT:\n raise NotImplementedError\n\n @abc.abstractmethod\n def validate_checkpoint_json(self, checkpoint_json: str) -> CT:\n \"\"\"Validate the checkpoint json and return the checkpoint object\"\"\"\n raise NotImplementedError\n\n\nclass CheckpointedConnectorWithPermSync(CheckpointedConnector[CT]):\n @abc.abstractmethod\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: CT,\n ) -> CheckpointOutput[CT]:\n raise NotImplementedError\n\n\nclass HierarchyConnector(BaseConnector):\n @abc.abstractmethod\n def load_hierarchy(\n self,\n start: SecondsSinceUnixEpoch, # may be unused if the connector must load the full hierarchy each time\n end: SecondsSinceUnixEpoch,\n ) -> HierarchyOutput:\n raise NotImplementedError\n" }, { "path": "backend/onyx/connectors/jira/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/jira/access.py", "content": "\"\"\"\nPermissioning / AccessControl logic for JIRA Projects + Issues.\n\"\"\"\n\nfrom collections.abc import Callable\nfrom typing import cast\n\nfrom jira import JIRA\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\n\ndef get_project_permissions(\n jira_client: JIRA,\n jira_project: str,\n add_prefix: bool = False,\n) -> ExternalAccess | None:\n \"\"\"\n Fetch the project + issue level permissions / access-control.\n This functionality requires Enterprise Edition.\n\n Args:\n jira_client: The JIRA client instance.\n jira_project: The JIRA project string.\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path\n where upsert_document_external_perms handles prefixing).\n\n Returns:\n ExternalAccess object for the page. None if EE is not enabled or no restrictions found.\n \"\"\"\n\n # Check if EE is enabled\n if not global_version.is_ee_version():\n return None\n\n ee_get_project_permissions = cast(\n Callable[\n [JIRA, str, bool],\n ExternalAccess | None,\n ],\n fetch_versioned_implementation(\n \"onyx.external_permissions.jira.page_access\", \"get_project_permissions\"\n ),\n )\n\n return ee_get_project_permissions(\n jira_client,\n jira_project,\n add_prefix,\n )\n" }, { "path": "backend/onyx/connectors/jira/connector.py", "content": "import copy\nimport json\nimport os\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom collections.abc import Iterable\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\n\nimport requests\nfrom jira import JIRA\nfrom jira.exceptions import JIRAError\nfrom jira.resources import Issue\nfrom more_itertools import chunked\nfrom typing_extensions import override\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import JIRA_CONNECTOR_LABELS_TO_SKIP\nfrom onyx.configs.app_configs import JIRA_CONNECTOR_MAX_TICKET_SIZE\nfrom onyx.configs.app_configs import JIRA_SLIM_PAGE_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n is_atlassian_date_error,\n)\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.jira.access import get_project_permissions\nfrom onyx.connectors.jira.utils import best_effort_basic_expert_info\nfrom onyx.connectors.jira.utils import best_effort_get_field_from_issue\nfrom onyx.connectors.jira.utils import build_jira_client\nfrom onyx.connectors.jira.utils import build_jira_url\nfrom onyx.connectors.jira.utils import extract_text_from_adf\nfrom onyx.connectors.jira.utils import get_comment_strs\nfrom onyx.connectors.jira.utils import JIRA_CLOUD_API_VERSION\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\nONE_HOUR = 3600\n\n_MAX_RESULTS_FETCH_IDS = 5000 # 5000\n_JIRA_FULL_PAGE_SIZE = 50\n\n# Constants for Jira field names\n_FIELD_REPORTER = \"reporter\"\n_FIELD_ASSIGNEE = \"assignee\"\n_FIELD_PRIORITY = \"priority\"\n_FIELD_STATUS = \"status\"\n_FIELD_RESOLUTION = \"resolution\"\n_FIELD_LABELS = \"labels\"\n_FIELD_KEY = \"key\"\n_FIELD_CREATED = \"created\"\n_FIELD_DUEDATE = \"duedate\"\n_FIELD_ISSUETYPE = \"issuetype\"\n_FIELD_PARENT = \"parent\"\n_FIELD_ASSIGNEE_EMAIL = \"assignee_email\"\n_FIELD_REPORTER_EMAIL = \"reporter_email\"\n_FIELD_PROJECT = \"project\"\n_FIELD_PROJECT_NAME = \"project_name\"\n_FIELD_UPDATED = \"updated\"\n_FIELD_RESOLUTION_DATE = \"resolutiondate\"\n_FIELD_RESOLUTION_DATE_KEY = \"resolution_date\"\n\n\ndef _is_cloud_client(jira_client: JIRA) -> bool:\n return jira_client._options[\"rest_api_version\"] == JIRA_CLOUD_API_VERSION\n\n\ndef _perform_jql_search(\n jira_client: JIRA,\n jql: str,\n start: int,\n max_results: int,\n fields: str | None = None,\n all_issue_ids: list[list[str]] | None = None,\n checkpoint_callback: (\n Callable[[Iterator[list[str]], str | None], None] | None\n ) = None,\n nextPageToken: str | None = None,\n ids_done: bool = False,\n) -> Iterable[Issue]:\n \"\"\"\n The caller should expect\n a) this function returns an iterable of issues of length 0 < len(issues) <= max_results.\n - caveat; if all_issue_ids is provided, the iterable will be the size of some sub-list.\n - this will only not match the above bound if a recent deployment changed max_results.\n\n IF the v3 API is used (i.e. the jira instance is a cloud instance), then the caller should expect:\n\n b) this function will call checkpoint_callback ONCE after at least one of the following has happened:\n - a new batch of ids has been fetched via enhanced search\n - a batch of issues has been bulk-fetched\n c) checkpoint_callback is called with the new all_issue_ids and the pageToken of the enhanced\n search request. We pass in a pageToken of None once we've fetched all the issue ids.\n\n Note: nextPageToken is valid for 7 days according to a post from a year ago, so for now\n we won't add any handling for restarting (just re-index, since there's no easy\n way to recover from this).\n \"\"\"\n # it would be preferable to use one approach for both versions, but\n # v2 doesnt have the bulk fetch api and v3 has fully deprecated the search\n # api that v2 uses\n if _is_cloud_client(jira_client):\n if all_issue_ids is None:\n raise ValueError(\"all_issue_ids is required for v3\")\n return _perform_jql_search_v3(\n jira_client,\n jql,\n max_results,\n all_issue_ids,\n fields=fields,\n checkpoint_callback=checkpoint_callback,\n nextPageToken=nextPageToken,\n ids_done=ids_done,\n )\n else:\n return _perform_jql_search_v2(jira_client, jql, start, max_results, fields)\n\n\ndef _handle_jira_search_error(e: Exception, jql: str) -> None:\n \"\"\"Handle common Jira search errors and raise appropriate exceptions.\n\n Args:\n e: The exception raised by the Jira API\n jql: The JQL query that caused the error\n\n Raises:\n ConnectorValidationError: For HTTP 400 errors (invalid JQL or project)\n CredentialExpiredError: For HTTP 401 errors\n InsufficientPermissionsError: For HTTP 403 errors\n Exception: Re-raises the original exception for other error types\n \"\"\"\n # Extract error information from the exception\n error_text = \"\"\n status_code = None\n\n def _format_error_text(error_payload: Any) -> str:\n error_messages = (\n error_payload.get(\"errorMessages\", [])\n if isinstance(error_payload, dict)\n else []\n )\n if error_messages:\n return (\n \"; \".join(error_messages)\n if isinstance(error_messages, list)\n else str(error_messages)\n )\n return str(error_payload)\n\n # Try to get status code and error text from JIRAError or requests response\n if hasattr(e, \"status_code\"):\n status_code = e.status_code\n raw_text = getattr(e, \"text\", \"\")\n if isinstance(raw_text, str):\n try:\n error_text = _format_error_text(json.loads(raw_text))\n except Exception:\n error_text = raw_text\n else:\n error_text = str(raw_text)\n elif hasattr(e, \"response\") and e.response is not None:\n status_code = e.response.status_code\n # Try JSON first, fall back to text\n try:\n error_json = e.response.json()\n error_text = _format_error_text(error_json)\n except Exception:\n error_text = e.response.text\n\n # Handle specific status codes\n if status_code == 400:\n if \"does not exist for the field 'project'\" in error_text:\n raise ConnectorValidationError(\n f\"The specified Jira project does not exist or you don't have access to it. JQL query: {jql}. Error: {error_text}\"\n )\n raise ConnectorValidationError(\n f\"Invalid JQL query. JQL: {jql}. Error: {error_text}\"\n )\n elif status_code == 401:\n raise CredentialExpiredError(\n \"Jira credentials are expired or invalid (HTTP 401).\"\n )\n elif status_code == 403:\n raise InsufficientPermissionsError(\n f\"Insufficient permissions to execute JQL query. JQL: {jql}\"\n )\n\n # Re-raise for other error types\n raise e\n\n\ndef enhanced_search_ids(\n jira_client: JIRA, jql: str, nextPageToken: str | None = None\n) -> tuple[list[str], str | None]:\n # https://community.atlassian.com/forums/Jira-articles/\n # Avoiding-Pitfalls-A-Guide-to-Smooth-Migration-to-Enhanced-JQL/ba-p/2985433\n # For cloud, it's recommended that we fetch all ids first then use the bulk fetch API.\n # The enhanced search isn't currently supported by our python library, so we have to\n # do this janky thing where we use the session directly.\n enhanced_search_path = jira_client._get_url(\"search/jql\")\n params: dict[str, str | int | None] = {\n \"jql\": jql,\n \"maxResults\": _MAX_RESULTS_FETCH_IDS,\n \"nextPageToken\": nextPageToken,\n \"fields\": \"id\",\n }\n try:\n response = jira_client._session.get(enhanced_search_path, params=params)\n response.raise_for_status()\n response_json = response.json()\n except Exception as e:\n _handle_jira_search_error(e, jql)\n raise # Explicitly re-raise for type checker, should never reach here\n\n return [str(issue[\"id\"]) for issue in response_json[\"issues\"]], response_json.get(\n \"nextPageToken\"\n )\n\n\ndef _bulk_fetch_request(\n jira_client: JIRA, issue_ids: list[str], fields: str | None\n) -> list[dict[str, Any]]:\n \"\"\"Raw POST to the bulkfetch endpoint. Returns the list of raw issue dicts.\"\"\"\n bulk_fetch_path = jira_client._get_url(\"issue/bulkfetch\")\n # Prepare the payload according to Jira API v3 specification\n payload: dict[str, Any] = {\"issueIdsOrKeys\": issue_ids}\n # Only restrict fields if specified, might want to explicitly do this in the future\n # to avoid reading unnecessary data\n payload[\"fields\"] = fields.split(\",\") if fields else [\"*all\"]\n\n resp = jira_client._session.post(bulk_fetch_path, json=payload)\n return resp.json()[\"issues\"]\n\n\ndef bulk_fetch_issues(\n jira_client: JIRA, issue_ids: list[str], fields: str | None = None\n) -> list[Issue]:\n # TODO(evan): move away from this jira library if they continue to not support\n # the endpoints we need. Using private fields is not ideal, but\n # is likely fine for now since we pin the library version\n\n try:\n raw_issues = _bulk_fetch_request(jira_client, issue_ids, fields)\n except requests.exceptions.JSONDecodeError:\n if len(issue_ids) <= 1:\n logger.exception(\n f\"Jira bulk-fetch response for issue(s) {issue_ids} could not \"\n f\"be decoded as JSON (response too large or truncated).\"\n )\n raise\n\n mid = len(issue_ids) // 2\n logger.warning(\n f\"Jira bulk-fetch JSON decode failed for batch of {len(issue_ids)} issues. \"\n f\"Splitting into sub-batches of {mid} and {len(issue_ids) - mid}.\"\n )\n left = bulk_fetch_issues(jira_client, issue_ids[:mid], fields)\n right = bulk_fetch_issues(jira_client, issue_ids[mid:], fields)\n return left + right\n except Exception as e:\n logger.error(f\"Error fetching issues: {e}\")\n raise\n\n return [\n Issue(jira_client._options, jira_client._session, raw=issue)\n for issue in raw_issues\n ]\n\n\ndef _perform_jql_search_v3(\n jira_client: JIRA,\n jql: str,\n max_results: int,\n all_issue_ids: list[list[str]],\n fields: str | None = None,\n checkpoint_callback: (\n Callable[[Iterator[list[str]], str | None], None] | None\n ) = None,\n nextPageToken: str | None = None,\n ids_done: bool = False,\n) -> Iterable[Issue]:\n \"\"\"\n The way this works is we get all the issue ids and bulk fetch them in batches.\n However, for really large deployments we can't do these operations sequentially,\n as it might take several hours to fetch all the issue ids.\n\n So, each run of this function does at least one of:\n - fetch a batch of issue ids\n - bulk fetch a batch of issues\n\n If all_issue_ids is not None, we use it to bulk fetch issues.\n \"\"\"\n\n # with some careful synchronization these steps can be done in parallel,\n # leaving that out for now to avoid rate limit issues\n if not ids_done:\n new_ids, pageToken = enhanced_search_ids(jira_client, jql, nextPageToken)\n if checkpoint_callback is not None:\n checkpoint_callback(chunked(new_ids, max_results), pageToken)\n\n # bulk fetch issues from ids. Note that the above callback MAY mutate all_issue_ids,\n # but this fetch always just takes the last id batch.\n if all_issue_ids:\n yield from bulk_fetch_issues(jira_client, all_issue_ids.pop(), fields)\n\n\ndef _perform_jql_search_v2(\n jira_client: JIRA,\n jql: str,\n start: int,\n max_results: int,\n fields: str | None = None,\n) -> Iterable[Issue]:\n \"\"\"\n Unfortunately, jira server/data center will forever use the v2 APIs that are now deprecated.\n \"\"\"\n logger.debug(\n f\"Fetching Jira issues with JQL: {jql}, starting at {start}, max results: {max_results}\"\n )\n try:\n issues = jira_client.search_issues(\n jql_str=jql,\n startAt=start,\n maxResults=max_results,\n fields=fields,\n )\n except JIRAError as e:\n _handle_jira_search_error(e, jql)\n raise # Explicitly re-raise for type checker, should never reach here\n\n for issue in issues:\n if isinstance(issue, Issue):\n yield issue\n else:\n raise RuntimeError(f\"Found Jira object not of type Issue: {issue}\")\n\n\ndef process_jira_issue(\n jira_base_url: str,\n issue: Issue,\n comment_email_blacklist: tuple[str, ...] = (),\n labels_to_skip: set[str] | None = None,\n parent_hierarchy_raw_node_id: str | None = None,\n) -> Document | None:\n if labels_to_skip:\n if any(label in issue.fields.labels for label in labels_to_skip):\n logger.info(\n f\"Skipping {issue.key} because it has a label to skip. Found \"\n f\"labels: {issue.fields.labels}. Labels to skip: {labels_to_skip}.\"\n )\n return None\n\n if isinstance(issue.fields.description, str):\n description = issue.fields.description\n else:\n description = extract_text_from_adf(issue.raw[\"fields\"][\"description\"])\n\n comments = get_comment_strs(\n issue=issue,\n comment_email_blacklist=comment_email_blacklist,\n )\n ticket_content = f\"{description}\\n\" + \"\\n\".join(\n [f\"Comment: {comment}\" for comment in comments if comment]\n )\n\n # Check ticket size\n if len(ticket_content.encode(\"utf-8\")) > JIRA_CONNECTOR_MAX_TICKET_SIZE:\n logger.info(\n f\"Skipping {issue.key} because it exceeds the maximum size of {JIRA_CONNECTOR_MAX_TICKET_SIZE} bytes.\"\n )\n return None\n\n page_url = build_jira_url(jira_base_url, issue.key)\n\n metadata_dict: dict[str, str | list[str]] = {}\n people = set()\n\n creator = best_effort_get_field_from_issue(issue, _FIELD_REPORTER)\n if creator is not None and (\n basic_expert_info := best_effort_basic_expert_info(creator)\n ):\n people.add(basic_expert_info)\n metadata_dict[_FIELD_REPORTER] = basic_expert_info.get_semantic_name()\n if email := basic_expert_info.get_email():\n metadata_dict[_FIELD_REPORTER_EMAIL] = email\n\n assignee = best_effort_get_field_from_issue(issue, _FIELD_ASSIGNEE)\n if assignee is not None and (\n basic_expert_info := best_effort_basic_expert_info(assignee)\n ):\n people.add(basic_expert_info)\n metadata_dict[_FIELD_ASSIGNEE] = basic_expert_info.get_semantic_name()\n if email := basic_expert_info.get_email():\n metadata_dict[_FIELD_ASSIGNEE_EMAIL] = email\n\n metadata_dict[_FIELD_KEY] = issue.key\n if priority := best_effort_get_field_from_issue(issue, _FIELD_PRIORITY):\n metadata_dict[_FIELD_PRIORITY] = priority.name\n if status := best_effort_get_field_from_issue(issue, _FIELD_STATUS):\n metadata_dict[_FIELD_STATUS] = status.name\n if resolution := best_effort_get_field_from_issue(issue, _FIELD_RESOLUTION):\n metadata_dict[_FIELD_RESOLUTION] = resolution.name\n if labels := best_effort_get_field_from_issue(issue, _FIELD_LABELS):\n metadata_dict[_FIELD_LABELS] = labels\n if created := best_effort_get_field_from_issue(issue, _FIELD_CREATED):\n metadata_dict[_FIELD_CREATED] = created\n if updated := best_effort_get_field_from_issue(issue, _FIELD_UPDATED):\n metadata_dict[_FIELD_UPDATED] = updated\n if duedate := best_effort_get_field_from_issue(issue, _FIELD_DUEDATE):\n metadata_dict[_FIELD_DUEDATE] = duedate\n if issuetype := best_effort_get_field_from_issue(issue, _FIELD_ISSUETYPE):\n metadata_dict[_FIELD_ISSUETYPE] = issuetype.name\n if resolutiondate := best_effort_get_field_from_issue(\n issue, _FIELD_RESOLUTION_DATE\n ):\n metadata_dict[_FIELD_RESOLUTION_DATE_KEY] = resolutiondate\n\n parent = best_effort_get_field_from_issue(issue, _FIELD_PARENT)\n if parent is not None:\n metadata_dict[_FIELD_PARENT] = parent.key\n\n project = best_effort_get_field_from_issue(issue, _FIELD_PROJECT)\n if project is not None:\n metadata_dict[_FIELD_PROJECT_NAME] = project.name\n metadata_dict[_FIELD_PROJECT] = project.key\n else:\n logger.error(f\"Project should exist but does not for {issue.key}\")\n\n return Document(\n id=page_url,\n sections=[TextSection(link=page_url, text=ticket_content)],\n source=DocumentSource.JIRA,\n semantic_identifier=f\"{issue.key}: {issue.fields.summary}\",\n title=f\"{issue.key} {issue.fields.summary}\",\n doc_updated_at=time_str_to_utc(issue.fields.updated),\n primary_owners=list(people) or None,\n metadata=metadata_dict,\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n )\n\n\nclass JiraConnectorCheckpoint(ConnectorCheckpoint):\n # used for v3 (cloud) endpoint\n all_issue_ids: list[list[str]] = []\n ids_done: bool = False\n cursor: str | None = None\n # deprecated\n # Used for v2 endpoint (server/data center)\n offset: int | None = None\n # Track hierarchy nodes we've already yielded to avoid duplicates across restarts\n seen_hierarchy_node_ids: list[str] = []\n\n\nclass JiraConnector(\n CheckpointedConnectorWithPermSync[JiraConnectorCheckpoint],\n SlimConnectorWithPermSync,\n):\n def __init__(\n self,\n jira_base_url: str,\n project_key: str | None = None,\n comment_email_blacklist: list[str] | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n # if a ticket has one of the labels specified in this list, we will just\n # skip it. This is generally used to avoid indexing extra sensitive\n # tickets.\n labels_to_skip: list[str] = JIRA_CONNECTOR_LABELS_TO_SKIP,\n # Custom JQL query to filter Jira issues\n jql_query: str | None = None,\n scoped_token: bool = False,\n ) -> None:\n self.batch_size = batch_size\n\n # dealing with scoped tokens is a bit tricky becasue we need to hit api.atlassian.net\n # when making jira requests but still want correct links to issues in the UI.\n # So, the user's base url is stored here, but converted to a scoped url when passed\n # to the jira client.\n self.jira_base = jira_base_url.rstrip(\"/\") # Remove trailing slash if present\n self.jira_project = project_key\n self._comment_email_blacklist = comment_email_blacklist or []\n self.labels_to_skip = set(labels_to_skip)\n self.jql_query = jql_query\n self.scoped_token = scoped_token\n self._jira_client: JIRA | None = None\n # Cache project permissions to avoid fetching them repeatedly across runs\n self._project_permissions_cache: dict[str, Any] = {}\n\n @property\n def comment_email_blacklist(self) -> tuple:\n return tuple(email.strip() for email in self._comment_email_blacklist)\n\n @property\n def jira_client(self) -> JIRA:\n if self._jira_client is None:\n raise ConnectorMissingCredentialError(\"Jira\")\n return self._jira_client\n\n @property\n def quoted_jira_project(self) -> str:\n # Quote the project name to handle reserved words\n if not self.jira_project:\n return \"\"\n return f'\"{self.jira_project}\"'\n\n def _get_project_permissions(\n self, project_key: str, add_prefix: bool = False\n ) -> Any:\n \"\"\"Get project permissions with caching.\n\n Args:\n project_key: The Jira project key\n add_prefix: When True, prefix group IDs with source type (for indexing path).\n When False (default), leave unprefixed (for permission sync path).\n\n Returns:\n The external access permissions for the project\n \"\"\"\n # Use different cache keys for prefixed vs unprefixed to avoid mixing\n cache_key = f\"{project_key}:{'prefixed' if add_prefix else 'unprefixed'}\"\n if cache_key not in self._project_permissions_cache:\n self._project_permissions_cache[cache_key] = get_project_permissions(\n jira_client=self.jira_client,\n jira_project=project_key,\n add_prefix=add_prefix,\n )\n return self._project_permissions_cache[cache_key]\n\n def _is_epic(self, issue: Issue) -> bool:\n \"\"\"Check if issue is an Epic.\"\"\"\n issuetype = best_effort_get_field_from_issue(issue, _FIELD_ISSUETYPE)\n if issuetype is None:\n return False\n return issuetype.name.lower() == \"epic\"\n\n def _is_parent_epic(self, parent: Any) -> bool:\n \"\"\"Check if a parent reference is an Epic.\n\n The parent object from issue.fields.parent has a different structure\n than a full Issue, so we handle it separately.\n \"\"\"\n parent_issuetype = (\n getattr(parent.fields, \"issuetype\", None)\n if hasattr(parent, \"fields\")\n else None\n )\n if parent_issuetype is None:\n return False\n return parent_issuetype.name.lower() == \"epic\"\n\n def _yield_project_hierarchy_node(\n self,\n project_key: str,\n project_name: str | None,\n seen_hierarchy_node_ids: set[str],\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield a hierarchy node for a project if not already yielded.\"\"\"\n if project_key in seen_hierarchy_node_ids:\n return\n\n seen_hierarchy_node_ids.add(project_key)\n\n yield HierarchyNode(\n raw_node_id=project_key,\n raw_parent_id=None, # Parent is SOURCE\n display_name=project_name or project_key,\n link=f\"{self.jira_base}/projects/{project_key}\",\n node_type=HierarchyNodeType.PROJECT,\n )\n\n def _yield_epic_hierarchy_node(\n self,\n issue: Issue,\n project_key: str,\n seen_hierarchy_node_ids: set[str],\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield a hierarchy node for an Epic issue.\"\"\"\n issue_key = issue.key\n if issue_key in seen_hierarchy_node_ids:\n return\n\n seen_hierarchy_node_ids.add(issue_key)\n\n yield HierarchyNode(\n raw_node_id=issue_key,\n raw_parent_id=project_key,\n display_name=f\"{issue_key}: {issue.fields.summary}\",\n link=build_jira_url(self.jira_base, issue_key),\n node_type=HierarchyNodeType.FOLDER, # don't have a separate epic node type\n )\n\n def _yield_parent_hierarchy_node_if_epic(\n self,\n parent: Any,\n project_key: str,\n seen_hierarchy_node_ids: set[str],\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield hierarchy node for parent issue if it's an Epic we haven't seen.\"\"\"\n parent_key = parent.key\n if parent_key in seen_hierarchy_node_ids:\n return\n\n if not self._is_parent_epic(parent):\n # Not an epic, don't create hierarchy node for it\n return\n\n seen_hierarchy_node_ids.add(parent_key)\n\n # Get summary if available\n parent_summary = (\n getattr(parent.fields, \"summary\", None)\n if hasattr(parent, \"fields\")\n else None\n )\n display_name = (\n f\"{parent_key}: {parent_summary}\" if parent_summary else parent_key\n )\n\n yield HierarchyNode(\n raw_node_id=parent_key,\n raw_parent_id=project_key,\n display_name=display_name,\n link=build_jira_url(self.jira_base, parent_key),\n node_type=HierarchyNodeType.FOLDER, # don't have a separate epic node type\n )\n\n def _get_parent_hierarchy_raw_node_id(self, issue: Issue, project_key: str) -> str:\n \"\"\"Determine the parent hierarchy node ID for an issue.\n\n Returns:\n - Epic key if issue's parent is an Epic\n - Project key otherwise (for top-level issues or non-epic parents)\n \"\"\"\n parent = best_effort_get_field_from_issue(issue, _FIELD_PARENT)\n if parent is None:\n # No parent, directly under project\n return project_key\n\n if self._is_parent_epic(parent):\n return parent.key\n\n # For non-epic parents (e.g., story with subtasks),\n # the document belongs directly under the project in the hierarchy\n return project_key\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self._jira_client = build_jira_client(\n credentials=credentials,\n jira_base=self.jira_base,\n scoped_token=self.scoped_token,\n )\n return None\n\n def _get_jql_query(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> str:\n \"\"\"Get the JQL query based on configuration and time range\n\n If a custom JQL query is provided, it will be used and combined with time constraints.\n Otherwise, the query will be constructed based on project key (if provided).\n \"\"\"\n start_date_str = datetime.fromtimestamp(start, tz=timezone.utc).strftime(\n \"%Y-%m-%d %H:%M\"\n )\n end_date_str = datetime.fromtimestamp(end, tz=timezone.utc).strftime(\n \"%Y-%m-%d %H:%M\"\n )\n\n time_jql = f\"updated >= '{start_date_str}' AND updated <= '{end_date_str}'\"\n\n # If custom JQL query is provided, use it and combine with time constraints\n if self.jql_query:\n return f\"({self.jql_query}) AND {time_jql}\"\n\n # Otherwise, use project key if provided\n if self.jira_project:\n base_jql = f\"project = {self.quoted_jira_project}\"\n return f\"{base_jql} AND {time_jql}\"\n\n return time_jql\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: JiraConnectorCheckpoint,\n ) -> CheckpointOutput[JiraConnectorCheckpoint]:\n jql = self._get_jql_query(start, end)\n try:\n return self._load_from_checkpoint(\n jql, checkpoint, include_permissions=False\n )\n except Exception as e:\n if is_atlassian_date_error(e):\n jql = self._get_jql_query(start - ONE_HOUR, end)\n return self._load_from_checkpoint(\n jql, checkpoint, include_permissions=False\n )\n raise e\n\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: JiraConnectorCheckpoint,\n ) -> CheckpointOutput[JiraConnectorCheckpoint]:\n \"\"\"Load documents from checkpoint with permission information included.\"\"\"\n jql = self._get_jql_query(start, end)\n try:\n return self._load_from_checkpoint(jql, checkpoint, include_permissions=True)\n except Exception as e:\n if is_atlassian_date_error(e):\n jql = self._get_jql_query(start - ONE_HOUR, end)\n return self._load_from_checkpoint(\n jql, checkpoint, include_permissions=True\n )\n raise e\n\n def _load_from_checkpoint(\n self, jql: str, checkpoint: JiraConnectorCheckpoint, include_permissions: bool\n ) -> CheckpointOutput[JiraConnectorCheckpoint]:\n # Get the current offset from checkpoint or start at 0\n starting_offset = checkpoint.offset or 0\n current_offset = starting_offset\n new_checkpoint = copy.deepcopy(checkpoint)\n\n # Convert checkpoint list to set for efficient lookups\n seen_hierarchy_node_ids = set(new_checkpoint.seen_hierarchy_node_ids)\n\n checkpoint_callback = make_checkpoint_callback(new_checkpoint)\n\n for issue in _perform_jql_search(\n jira_client=self.jira_client,\n jql=jql,\n start=current_offset,\n max_results=_JIRA_FULL_PAGE_SIZE,\n all_issue_ids=new_checkpoint.all_issue_ids,\n checkpoint_callback=checkpoint_callback,\n nextPageToken=new_checkpoint.cursor,\n ids_done=new_checkpoint.ids_done,\n ):\n issue_key = issue.key\n try:\n # Get project info for hierarchy\n project = best_effort_get_field_from_issue(issue, _FIELD_PROJECT)\n project_key = project.key if project else None\n project_name = project.name if project else None\n\n # Yield hierarchy nodes BEFORE the document (parent-before-child)\n if project_key:\n # 1. Yield project hierarchy node (if not already yielded)\n yield from self._yield_project_hierarchy_node(\n project_key, project_name, seen_hierarchy_node_ids\n )\n\n # 2. If parent is an Epic, yield hierarchy node for it\n parent = best_effort_get_field_from_issue(issue, _FIELD_PARENT)\n if parent:\n yield from self._yield_parent_hierarchy_node_if_epic(\n parent, project_key, seen_hierarchy_node_ids\n )\n\n # 3. If this issue IS an Epic, yield it as hierarchy node\n if self._is_epic(issue):\n yield from self._yield_epic_hierarchy_node(\n issue, project_key, seen_hierarchy_node_ids\n )\n\n # Determine parent hierarchy node ID for the document\n parent_hierarchy_raw_node_id = (\n self._get_parent_hierarchy_raw_node_id(issue, project_key)\n if project_key\n else None\n )\n\n if document := process_jira_issue(\n jira_base_url=self.jira_base,\n issue=issue,\n comment_email_blacklist=self.comment_email_blacklist,\n labels_to_skip=self.labels_to_skip,\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n ):\n # Add permission information to the document if requested\n if include_permissions:\n document.external_access = self._get_project_permissions(\n project_key,\n add_prefix=True, # Indexing path - prefix here\n )\n yield document\n\n except Exception as e:\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=issue_key,\n document_link=build_jira_url(self.jira_base, issue_key),\n ),\n failure_message=f\"Failed to process Jira issue: {str(e)}\",\n exception=e,\n )\n\n current_offset += 1\n\n # Update checkpoint with seen hierarchy nodes\n new_checkpoint.seen_hierarchy_node_ids = list(seen_hierarchy_node_ids)\n\n # Update checkpoint\n self.update_checkpoint_for_next_run(\n new_checkpoint, current_offset, starting_offset, _JIRA_FULL_PAGE_SIZE\n )\n\n return new_checkpoint\n\n def update_checkpoint_for_next_run(\n self,\n checkpoint: JiraConnectorCheckpoint,\n current_offset: int,\n starting_offset: int,\n page_size: int,\n ) -> None:\n if _is_cloud_client(self.jira_client):\n # other updates done in the checkpoint callback\n checkpoint.has_more = (\n len(checkpoint.all_issue_ids) > 0 or not checkpoint.ids_done\n )\n else:\n checkpoint.offset = current_offset\n # if we didn't retrieve a full batch, we're done\n checkpoint.has_more = current_offset - starting_offset == page_size\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n one_day = timedelta(hours=24).total_seconds()\n\n start = start or 0\n end = (\n end or datetime.now().timestamp() + one_day\n ) # we add one day to account for any potential timezone issues\n\n jql = self._get_jql_query(start, end)\n checkpoint = self.build_dummy_checkpoint()\n checkpoint_callback = make_checkpoint_callback(checkpoint)\n prev_offset = 0\n current_offset = 0\n slim_doc_batch: list[SlimDocument | HierarchyNode] = []\n\n # Track seen hierarchy nodes within this sync run\n seen_hierarchy_node_ids: set[str] = set()\n\n while checkpoint.has_more:\n for issue in _perform_jql_search(\n jira_client=self.jira_client,\n jql=jql,\n start=current_offset,\n max_results=JIRA_SLIM_PAGE_SIZE,\n all_issue_ids=checkpoint.all_issue_ids,\n checkpoint_callback=checkpoint_callback,\n nextPageToken=checkpoint.cursor,\n ids_done=checkpoint.ids_done,\n ):\n # Get project info\n project = best_effort_get_field_from_issue(issue, _FIELD_PROJECT)\n project_key = project.key if project else None\n project_name = project.name if project else None\n\n if not project_key:\n continue\n\n # Yield hierarchy nodes BEFORE the slim document (parent-before-child)\n # 1. Yield project hierarchy node (if not already yielded)\n for node in self._yield_project_hierarchy_node(\n project_key, project_name, seen_hierarchy_node_ids\n ):\n slim_doc_batch.append(node)\n\n # 2. If parent is an Epic, yield hierarchy node for it\n parent = best_effort_get_field_from_issue(issue, _FIELD_PARENT)\n if parent:\n for node in self._yield_parent_hierarchy_node_if_epic(\n parent, project_key, seen_hierarchy_node_ids\n ):\n slim_doc_batch.append(node)\n\n # 3. If this issue IS an Epic, yield it as hierarchy node\n if self._is_epic(issue):\n for node in self._yield_epic_hierarchy_node(\n issue, project_key, seen_hierarchy_node_ids\n ):\n slim_doc_batch.append(node)\n\n # Now add the slim document\n issue_key = best_effort_get_field_from_issue(issue, _FIELD_KEY)\n doc_id = build_jira_url(self.jira_base, issue_key)\n\n slim_doc_batch.append(\n SlimDocument(\n id=doc_id,\n # Permission sync path - don't prefix, upsert_document_external_perms handles it\n external_access=self._get_project_permissions(\n project_key, add_prefix=False\n ),\n parent_hierarchy_raw_node_id=(\n self._get_parent_hierarchy_raw_node_id(issue, project_key)\n if project_key\n else None\n ),\n )\n )\n current_offset += 1\n if len(slim_doc_batch) >= JIRA_SLIM_PAGE_SIZE:\n yield slim_doc_batch\n slim_doc_batch = []\n self.update_checkpoint_for_next_run(\n checkpoint, current_offset, prev_offset, JIRA_SLIM_PAGE_SIZE\n )\n prev_offset = current_offset\n\n if slim_doc_batch:\n yield slim_doc_batch\n\n def validate_connector_settings(self) -> None:\n if self._jira_client is None:\n raise ConnectorMissingCredentialError(\"Jira\")\n\n # If a custom JQL query is set, validate it's valid\n if self.jql_query:\n try:\n # Try to execute the JQL query with a small limit to validate its syntax\n # Use next(iter(...), None) to get just the first result without\n # forcing evaluation of all results\n next(\n iter(\n _perform_jql_search(\n jira_client=self.jira_client,\n jql=self.jql_query,\n start=0,\n max_results=1,\n all_issue_ids=[],\n )\n ),\n None,\n )\n except Exception as e:\n self._handle_jira_connector_settings_error(e)\n\n # If a specific project is set, validate it exists\n elif self.jira_project:\n try:\n self.jira_client.project(self.jira_project)\n except Exception as e:\n self._handle_jira_connector_settings_error(e)\n else:\n # If neither JQL nor project specified, validate we can access the Jira API\n try:\n # Try to list projects to validate access\n self.jira_client.projects()\n except Exception as e:\n self._handle_jira_connector_settings_error(e)\n\n def _handle_jira_connector_settings_error(self, e: Exception) -> None:\n \"\"\"Helper method to handle Jira API errors consistently.\n\n Extracts error messages from the Jira API response for all status codes when possible,\n providing more user-friendly error messages.\n\n Args:\n e: The exception raised by the Jira API\n\n Raises:\n CredentialExpiredError: If the status code is 401\n InsufficientPermissionsError: If the status code is 403\n ConnectorValidationError: For other HTTP errors with extracted error messages\n \"\"\"\n status_code = getattr(e, \"status_code\", None)\n logger.error(f\"Jira API error during validation: {e}\")\n\n # Handle specific status codes with appropriate exceptions\n if status_code == 401:\n raise CredentialExpiredError(\n \"Jira credential appears to be expired or invalid (HTTP 401).\"\n )\n elif status_code == 403:\n raise InsufficientPermissionsError(\n \"Your Jira token does not have sufficient permissions for this configuration (HTTP 403).\"\n )\n elif status_code == 429:\n raise ConnectorValidationError(\n \"Validation failed due to Jira rate-limits being exceeded. Please try again later.\"\n )\n\n # Try to extract original error message from the response\n error_message = getattr(e, \"text\", None)\n if error_message is None:\n raise UnexpectedValidationError(\n f\"Unexpected Jira error during validation: {e}\"\n )\n\n raise ConnectorValidationError(\n f\"Validation failed due to Jira error: {error_message}\"\n )\n\n @override\n def validate_checkpoint_json(self, checkpoint_json: str) -> JiraConnectorCheckpoint:\n return JiraConnectorCheckpoint.model_validate_json(checkpoint_json)\n\n @override\n def build_dummy_checkpoint(self) -> JiraConnectorCheckpoint:\n return JiraConnectorCheckpoint(\n has_more=True,\n )\n\n\ndef make_checkpoint_callback(\n checkpoint: JiraConnectorCheckpoint,\n) -> Callable[[Iterator[list[str]], str | None], None]:\n def checkpoint_callback(\n issue_ids: Iterator[list[str]], pageToken: str | None\n ) -> None:\n for id_batch in issue_ids:\n checkpoint.all_issue_ids.append(id_batch)\n checkpoint.cursor = pageToken\n # pageToken starts out as None and is only None once we've fetched all the issue ids\n checkpoint.ids_done = pageToken is None\n\n return checkpoint_callback\n\n\nif __name__ == \"__main__\":\n import os\n from onyx.utils.variable_functionality import global_version\n from tests.daily.connectors.utils import load_all_from_connector\n\n # For connector permission testing, set EE to true.\n global_version.set_ee()\n\n connector = JiraConnector(\n jira_base_url=os.environ[\"JIRA_BASE_URL\"],\n project_key=os.environ.get(\"JIRA_PROJECT_KEY\"),\n comment_email_blacklist=[],\n )\n\n connector.load_credentials(\n {\n \"jira_user_email\": os.environ[\"JIRA_USER_EMAIL\"],\n \"jira_api_token\": os.environ[\"JIRA_API_TOKEN\"],\n }\n )\n\n start = 0\n end = datetime.now().timestamp()\n\n for slim_doc in connector.retrieve_all_slim_docs_perm_sync(\n start=start,\n end=end,\n ):\n print(slim_doc)\n\n for doc in load_all_from_connector(\n connector=connector,\n start=start,\n end=end,\n ).documents:\n print(doc)\n" }, { "path": "backend/onyx/connectors/jira/utils.py", "content": "\"\"\"Module with custom fields processing functions\"\"\"\n\nimport os\nfrom typing import Any\nfrom typing import List\nfrom urllib.parse import urlparse\n\nfrom jira import JIRA\nfrom jira.resources import CustomFieldOption\nfrom jira.resources import Issue\nfrom jira.resources import User\n\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import scoped_url\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nPROJECT_URL_PAT = \"projects\"\nJIRA_SERVER_API_VERSION = os.environ.get(\"JIRA_SERVER_API_VERSION\") or \"2\"\nJIRA_CLOUD_API_VERSION = os.environ.get(\"JIRA_CLOUD_API_VERSION\") or \"3\"\n\n\ndef best_effort_basic_expert_info(obj: Any) -> BasicExpertInfo | None:\n display_name = None\n email = None\n\n try:\n if hasattr(obj, \"displayName\"):\n display_name = obj.displayName\n else:\n display_name = obj.get(\"displayName\")\n\n if hasattr(obj, \"emailAddress\"):\n email = obj.emailAddress\n else:\n email = obj.get(\"emailAddress\")\n\n except Exception:\n return None\n\n if not email and not display_name:\n return None\n\n return BasicExpertInfo(display_name=display_name, email=email)\n\n\ndef best_effort_get_field_from_issue(jira_issue: Issue, field: str) -> Any:\n if hasattr(jira_issue, field):\n return getattr(jira_issue, field)\n\n if hasattr(jira_issue, \"fields\") and hasattr(jira_issue.fields, field):\n return getattr(jira_issue.fields, field)\n\n try:\n return jira_issue.raw[\"fields\"][field]\n except Exception:\n return None\n\n\ndef extract_text_from_adf(adf: dict | None) -> str:\n \"\"\"Extracts plain text from Atlassian Document Format:\n https://developer.atlassian.com/cloud/jira/platform/apis/document/structure/\n\n WARNING: This function is incomplete and will e.g. skip lists!\n \"\"\"\n # TODO: complete this function\n texts = []\n if adf is not None and \"content\" in adf:\n for block in adf[\"content\"]:\n if \"content\" in block:\n for item in block[\"content\"]:\n if item[\"type\"] == \"text\":\n texts.append(item[\"text\"])\n return \" \".join(texts)\n\n\ndef build_jira_url(jira_base_url: str, issue_key: str) -> str:\n \"\"\"\n Get the url used to access an issue in the UI.\n \"\"\"\n return f\"{jira_base_url}/browse/{issue_key}\"\n\n\ndef build_jira_client(\n credentials: dict[str, Any], jira_base: str, scoped_token: bool = False\n) -> JIRA:\n\n jira_base = scoped_url(jira_base, \"jira\") if scoped_token else jira_base\n api_token = credentials[\"jira_api_token\"]\n # if user provide an email we assume it's cloud\n if \"jira_user_email\" in credentials:\n email = credentials[\"jira_user_email\"]\n return JIRA(\n basic_auth=(email, api_token),\n server=jira_base,\n options={\"rest_api_version\": JIRA_CLOUD_API_VERSION},\n )\n else:\n return JIRA(\n token_auth=api_token,\n server=jira_base,\n options={\"rest_api_version\": JIRA_SERVER_API_VERSION},\n )\n\n\ndef extract_jira_project(url: str) -> tuple[str, str]:\n parsed_url = urlparse(url)\n jira_base = parsed_url.scheme + \"://\" + parsed_url.netloc\n\n # Split the path by '/' and find the position of 'projects' to get the project name\n split_path = parsed_url.path.split(\"/\")\n if PROJECT_URL_PAT in split_path:\n project_pos = split_path.index(PROJECT_URL_PAT)\n if len(split_path) > project_pos + 1:\n jira_project = split_path[project_pos + 1]\n else:\n raise ValueError(\"No project name found in the URL\")\n else:\n raise ValueError(\"'projects' not found in the URL\")\n\n return jira_base, jira_project\n\n\ndef get_comment_strs(\n issue: Issue, comment_email_blacklist: tuple[str, ...] = ()\n) -> list[str]:\n comment_strs = []\n for comment in issue.fields.comment.comments:\n try:\n if isinstance(comment.body, str):\n body_text = comment.body\n else:\n body_text = extract_text_from_adf(comment.raw[\"body\"])\n\n if (\n hasattr(comment, \"author\")\n and hasattr(comment.author, \"emailAddress\")\n and comment.author.emailAddress in comment_email_blacklist\n ):\n continue # Skip adding comment if author's email is in blacklist\n\n comment_strs.append(body_text)\n except Exception as e:\n logger.error(f\"Failed to process comment due to an error: {e}\")\n continue\n\n return comment_strs\n\n\ndef get_jira_project_key_from_issue(issue: Issue) -> str | None:\n if not hasattr(issue, \"fields\"):\n return None\n if not hasattr(issue.fields, \"project\"):\n return None\n if not hasattr(issue.fields.project, \"key\"):\n return None\n\n return issue.fields.project.key\n\n\nclass CustomFieldExtractor:\n @staticmethod\n def _process_custom_field_value(value: Any) -> str:\n \"\"\"\n Process a custom field value to a string\n \"\"\"\n try:\n if isinstance(value, str):\n return value\n elif isinstance(value, CustomFieldOption):\n return value.value\n elif isinstance(value, User):\n return value.displayName\n elif isinstance(value, List):\n return \" \".join(\n [CustomFieldExtractor._process_custom_field_value(v) for v in value]\n )\n else:\n return str(value)\n except Exception as e:\n logger.error(f\"Error processing custom field value {value}: {e}\")\n return \"\"\n\n @staticmethod\n def get_issue_custom_fields(\n jira: Issue, custom_fields: dict, max_value_length: int = 250\n ) -> dict:\n \"\"\"\n Process all custom fields of an issue to a dictionary of strings\n :param jira: jira_issue, bug or similar\n :param custom_fields: custom fields dictionary\n :param max_value_length: maximum length of the value to be processed, if exceeded, it will be truncated\n \"\"\"\n\n issue_custom_fields = {\n custom_fields[key]: value\n for key, value in jira.fields.__dict__.items()\n if value and key in custom_fields.keys()\n }\n\n processed_fields = {}\n\n if issue_custom_fields:\n for key, value in issue_custom_fields.items():\n processed = CustomFieldExtractor._process_custom_field_value(value)\n # We need max length parameter, because there are some plugins that often has very long description\n # and there is just a technical information so we just avoid long values\n if len(processed) < max_value_length:\n processed_fields[key] = processed\n\n return processed_fields\n\n @staticmethod\n def get_all_custom_fields(jira_client: JIRA) -> dict:\n \"\"\"Get all custom fields from Jira\"\"\"\n fields = jira_client.fields()\n fields_dct = {\n field[\"id\"]: field[\"name\"] for field in fields if field[\"custom\"] is True\n }\n return fields_dct\n\n\nclass CommonFieldExtractor:\n @staticmethod\n def get_issue_common_fields(jira: Issue) -> dict:\n return {\n \"Priority\": jira.fields.priority.name if jira.fields.priority else None,\n \"Reporter\": (\n jira.fields.reporter.displayName if jira.fields.reporter else None\n ),\n \"Assignee\": (\n jira.fields.assignee.displayName if jira.fields.assignee else None\n ),\n \"Status\": jira.fields.status.name if jira.fields.status else None,\n \"Resolution\": (\n jira.fields.resolution.name if jira.fields.resolution else None\n ),\n }\n" }, { "path": "backend/onyx/connectors/linear/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/linear/connector.py", "content": "import os\nimport re\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom urllib.parse import urlparse\n\nimport requests\nfrom typing_extensions import override\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import LINEAR_CLIENT_ID\nfrom onyx.configs.app_configs import LINEAR_CLIENT_SECRET\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n get_oauth_callback_uri,\n)\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import NormalizationResult\nfrom onyx.connectors.interfaces import OAuthConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import request_with_retries\n\n\nlogger = setup_logger()\n\n_NUM_RETRIES = 5\n_TIMEOUT = 60\n_LINEAR_GRAPHQL_URL = \"https://api.linear.app/graphql\"\n\n\ndef _make_query(request_body: dict[str, Any], api_key: str) -> requests.Response:\n headers = {\n \"Authorization\": api_key,\n \"Content-Type\": \"application/json\",\n }\n\n for i in range(_NUM_RETRIES):\n try:\n response = requests.post(\n _LINEAR_GRAPHQL_URL,\n headers=headers,\n json=request_body,\n timeout=_TIMEOUT,\n )\n if not response.ok:\n raise RuntimeError(\n f\"Error fetching issues from Linear: {response.text}\"\n )\n\n return response\n except Exception as e:\n if i == _NUM_RETRIES - 1:\n raise e\n\n logger.warning(f\"A Linear GraphQL error occurred: {e}. Retrying...\")\n\n raise RuntimeError(\n \"Unexpected execution when querying Linear. This should never happen.\"\n )\n\n\nclass LinearConnector(LoadConnector, PollConnector, OAuthConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.batch_size = batch_size\n self.linear_api_key: str | None = None\n\n @classmethod\n def oauth_id(cls) -> DocumentSource:\n return DocumentSource.LINEAR\n\n @classmethod\n def oauth_authorization_url(\n cls,\n base_domain: str,\n state: str,\n additional_kwargs: dict[str, str], # noqa: ARG003\n ) -> str:\n if not LINEAR_CLIENT_ID:\n raise ValueError(\"LINEAR_CLIENT_ID environment variable must be set\")\n\n callback_uri = get_oauth_callback_uri(base_domain, DocumentSource.LINEAR.value)\n return (\n f\"https://linear.app/oauth/authorize\"\n f\"?client_id={LINEAR_CLIENT_ID}\"\n f\"&redirect_uri={callback_uri}\"\n f\"&response_type=code\"\n f\"&scope=read\"\n f\"&state={state}\"\n f\"&prompt=consent\" # prompts user for access; allows choosing workspace\n )\n\n @classmethod\n def oauth_code_to_token(\n cls,\n base_domain: str,\n code: str,\n additional_kwargs: dict[str, str], # noqa: ARG003\n ) -> dict[str, Any]:\n data = {\n \"code\": code,\n \"redirect_uri\": get_oauth_callback_uri(\n base_domain, DocumentSource.LINEAR.value\n ),\n \"client_id\": LINEAR_CLIENT_ID,\n \"client_secret\": LINEAR_CLIENT_SECRET,\n \"grant_type\": \"authorization_code\",\n }\n headers = {\"Content-Type\": \"application/x-www-form-urlencoded\"}\n\n response = request_with_retries(\n method=\"POST\",\n url=\"https://api.linear.app/oauth/token\",\n data=data,\n headers=headers,\n backoff=0,\n delay=0.1,\n )\n if not response.ok:\n raise RuntimeError(f\"Failed to exchange code for token: {response.text}\")\n\n token_data = response.json()\n\n return {\n \"access_token\": token_data[\"access_token\"],\n }\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n if \"linear_api_key\" in credentials:\n self.linear_api_key = cast(str, credentials[\"linear_api_key\"])\n elif \"access_token\" in credentials:\n self.linear_api_key = \"Bearer \" + cast(str, credentials[\"access_token\"])\n else:\n # May need to handle case in the future if the OAuth flow expires\n raise ConnectorMissingCredentialError(\"Linear\")\n\n return None\n\n def _process_issues(\n self, start_str: datetime | None = None, end_str: datetime | None = None\n ) -> GenerateDocumentsOutput:\n if self.linear_api_key is None:\n raise ConnectorMissingCredentialError(\"Linear\")\n\n lte_filter = f'lte: \"{end_str}\"' if end_str else \"\"\n gte_filter = f'gte: \"{start_str}\"' if start_str else \"\"\n updatedAtFilter = f\"\"\"\n {lte_filter}\n {gte_filter}\n \"\"\"\n\n query = (\n \"\"\"\n query IterateIssueBatches($first: Int, $after: String) {\n issues(\n orderBy: updatedAt,\n first: $first,\n after: $after,\n filter: {\n updatedAt: {\n \"\"\"\n + updatedAtFilter\n + \"\"\"\n },\n\n }\n ) {\n edges {\n node {\n id\n createdAt\n updatedAt\n archivedAt\n number\n title\n priority\n estimate\n sortOrder\n startedAt\n completedAt\n startedTriageAt\n triagedAt\n canceledAt\n autoClosedAt\n autoArchivedAt\n dueDate\n slaStartedAt\n slaBreachesAt\n trashed\n snoozedUntilAt\n team {\n name\n }\n creator {\n name\n email\n }\n assignee {\n name\n email\n }\n previousIdentifiers\n subIssueSortOrder\n priorityLabel\n identifier\n url\n branchName\n state {\n id\n name\n }\n customerTicketCount\n description\n comments {\n nodes {\n url\n body\n }\n }\n }\n }\n pageInfo {\n hasNextPage\n endCursor\n }\n }\n }\n \"\"\"\n )\n\n has_more = True\n endCursor = None\n while has_more:\n graphql_query = {\n \"query\": query,\n \"variables\": {\n \"first\": self.batch_size,\n \"after\": endCursor,\n },\n }\n logger.debug(f\"Requesting issues from Linear with query: {graphql_query}\")\n\n response = _make_query(graphql_query, self.linear_api_key)\n response_json = response.json()\n logger.debug(f\"Raw response from Linear: {response_json}\")\n edges = response_json[\"data\"][\"issues\"][\"edges\"]\n\n documents: list[Document | HierarchyNode] = []\n for edge in edges:\n node = edge[\"node\"]\n # Create sections for description and comments\n sections = [\n TextSection(\n link=node[\"url\"],\n text=node[\"description\"] or \"\",\n )\n ]\n\n # Add comment sections\n for comment in node[\"comments\"][\"nodes\"]:\n sections.append(\n TextSection(\n link=node[\"url\"],\n text=comment[\"body\"] or \"\",\n )\n )\n\n # Cast the sections list to the expected type\n typed_sections = cast(list[TextSection | ImageSection], sections)\n\n # Extract team name for hierarchy\n team_name = (node.get(\"team\") or {}).get(\"name\") or \"Unknown Team\"\n identifier = node.get(\"identifier\", node[\"id\"])\n\n documents.append(\n Document(\n id=node[\"id\"],\n sections=typed_sections,\n source=DocumentSource.LINEAR,\n semantic_identifier=f\"[{node['identifier']}] {node['title']}\",\n title=node[\"title\"],\n doc_updated_at=time_str_to_utc(node[\"updatedAt\"]),\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [team_name],\n \"team_name\": team_name,\n \"identifier\": identifier,\n }\n },\n metadata={\n k: str(v)\n for k, v in {\n \"team\": (node.get(\"team\") or {}).get(\"name\"),\n \"creator\": node.get(\"creator\"),\n \"assignee\": node.get(\"assignee\"),\n \"state\": (node.get(\"state\") or {}).get(\"name\"),\n \"priority\": node.get(\"priority\"),\n \"estimate\": node.get(\"estimate\"),\n \"started_at\": node.get(\"startedAt\"),\n \"completed_at\": node.get(\"completedAt\"),\n \"created_at\": node.get(\"createdAt\"),\n \"due_date\": node.get(\"dueDate\"),\n }.items()\n if v is not None\n },\n )\n )\n yield documents\n\n endCursor = response_json[\"data\"][\"issues\"][\"pageInfo\"][\"endCursor\"]\n has_more = response_json[\"data\"][\"issues\"][\"pageInfo\"][\"hasNextPage\"]\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n yield from self._process_issues()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_time = datetime.fromtimestamp(start, tz=timezone.utc)\n end_time = datetime.fromtimestamp(end, tz=timezone.utc)\n\n yield from self._process_issues(start_str=start_time, end_str=end_time)\n\n @classmethod\n @override\n def normalize_url(cls, url: str) -> NormalizationResult:\n \"\"\"Extract Linear issue identifier from URL.\n\n Linear URLs are like: https://linear.app/team/issue/IDENTIFIER/...\n Returns the identifier (e.g., \"DAN-2327\") which can be used to match Document.link.\n \"\"\"\n parsed = urlparse(url)\n netloc = parsed.netloc.lower()\n\n if \"linear.app\" not in netloc:\n return NormalizationResult(normalized_url=None, use_default=False)\n\n # Extract identifier from path: /team/issue/IDENTIFIER/...\n # Pattern: /{team}/issue/{identifier}/...\n path_parts = [p for p in parsed.path.split(\"/\") if p]\n if len(path_parts) >= 3 and path_parts[1] == \"issue\":\n identifier = path_parts[2]\n # Validate identifier format (e.g., \"DAN-2327\")\n if re.match(r\"^[A-Z]+-\\d+$\", identifier):\n return NormalizationResult(normalized_url=identifier, use_default=False)\n\n return NormalizationResult(normalized_url=None, use_default=False)\n\n\nif __name__ == \"__main__\":\n connector = LinearConnector()\n connector.load_credentials({\"linear_api_key\": os.environ[\"LINEAR_API_KEY\"]})\n\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/loopio/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/loopio/connector.py", "content": "import json\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nfrom oauthlib.oauth2 import BackendApplicationClient\nfrom requests_oauthlib import OAuth2Session # type: ignore\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.file_processing.html_utils import strip_excessive_newlines_and_spaces\nfrom onyx.utils.logger import setup_logger\n\nLOOPIO_API_BASE = \"https://api.loopio.com/\"\nLOOPIO_AUTH_URL = LOOPIO_API_BASE + \"oauth2/access_token\"\nLOOPIO_DATA_URL = LOOPIO_API_BASE + \"data/\"\n\nlogger = setup_logger()\n\n\nclass LoopioConnector(LoadConnector, PollConnector):\n def __init__(\n self,\n loopio_stack_name: str | None = None,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.batch_size = batch_size\n self.loopio_client_id: str | None = None\n self.loopio_client_token: str | None = None\n self.loopio_stack_name = loopio_stack_name\n\n def _fetch_data(\n self, resource: str, params: dict[str, str | int]\n ) -> Generator[dict[str, Any], None, None]:\n client = BackendApplicationClient(\n client_id=self.loopio_client_id, scope=[\"library:read\"]\n )\n session = OAuth2Session(client=client)\n session.fetch_token(\n token_url=LOOPIO_AUTH_URL,\n client_id=self.loopio_client_id,\n client_secret=self.loopio_client_token,\n )\n page = 0\n stop_at_page = 1\n while (page := page + 1) <= stop_at_page:\n params[\"page\"] = page\n response = session.request(\n \"GET\",\n LOOPIO_DATA_URL + resource,\n headers={\"Accept\": \"application/json\"},\n params=params,\n )\n if response.status_code == 400:\n logger.error(\n f\"Loopio API returned 400 for {resource} with params {params}\",\n )\n logger.error(response.text)\n response.raise_for_status()\n response_data = json.loads(response.text)\n stop_at_page = response_data.get(\"totalPages\", 1)\n yield response_data\n\n def _build_search_filter(\n self, stack_name: str | None, start: str | None, end: str | None\n ) -> dict[str, Any]:\n filter: dict[str, Any] = {}\n if start is not None and end is not None:\n filter[\"lastUpdatedDate\"] = {\"gte\": start, \"lt\": end}\n\n if stack_name is not None:\n # Right now this is fetching the stacks every time, which is not ideal.\n # We should update this later to store the ID when we create the Connector\n for stack in self._fetch_data(resource=\"v2/stacks\", params={}):\n for item in stack[\"items\"]:\n if item[\"name\"] == stack_name:\n filter[\"locations\"] = [{\"stackID\": item[\"id\"]}]\n break\n if \"locations\" not in filter:\n raise ValueError(f\"Stack {stack_name} not found in Loopio\")\n return filter\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.loopio_subdomain = credentials[\"loopio_subdomain\"]\n self.loopio_client_id = credentials[\"loopio_client_id\"]\n self.loopio_client_token = credentials[\"loopio_client_token\"]\n return None\n\n def _process_entries(\n self, start: str | None = None, end: str | None = None\n ) -> GenerateDocumentsOutput:\n if self.loopio_client_id is None or self.loopio_client_token is None:\n raise ConnectorMissingCredentialError(\"Loopio\")\n\n filter = self._build_search_filter(\n stack_name=self.loopio_stack_name, start=start, end=end\n )\n params: dict[str, str | int] = {\"pageSize\": self.batch_size}\n params[\"filter\"] = json.dumps(filter)\n\n doc_batch: list[Document | HierarchyNode] = []\n for library_entries in self._fetch_data(\n resource=\"v2/libraryEntries\", params=params\n ):\n for entry in library_entries.get(\"items\", []):\n link = f\"https://{self.loopio_subdomain}.loopio.com/library?entry={entry['id']}\"\n topic = \"/\".join(\n part[\"name\"] for part in entry[\"location\"].values() if part\n )\n\n answer_text = entry.get(\"answer\", {}).get(\"text\", \"\")\n if not answer_text:\n logger.warning(\n f\"The Library entry {entry['id']} has no answer text. Skipping.\"\n )\n continue\n\n try:\n answer = parse_html_page_basic(answer_text)\n except Exception as e:\n logger.error(f\"Error parsing HTML for entry {entry['id']}: {e}\")\n continue\n\n questions = [\n question.get(\"text\").replace(\"\\xa0\", \" \")\n for question in entry[\"questions\"]\n if question.get(\"text\")\n ]\n questions_string = strip_excessive_newlines_and_spaces(\n \"\\n\".join(questions)\n )\n content_text = f\"{answer}\\n\\nRelated Questions: {questions_string}\"\n content_text = strip_excessive_newlines_and_spaces(\n content_text.replace(\"\\xa0\", \" \")\n )\n\n last_updated = time_str_to_utc(entry[\"lastUpdatedDate\"])\n last_reviewed = (\n time_str_to_utc(entry[\"lastReviewedDate\"])\n if entry.get(\"lastReviewedDate\")\n else None\n )\n\n # For Onyx, we decay document score overtime, either last_updated or\n # last_reviewed is a good enough signal for the document's recency\n latest_time = (\n max(last_reviewed, last_updated) if last_reviewed else last_updated\n )\n creator = entry.get(\"creator\")\n last_updated_by = entry.get(\"lastUpdatedBy\")\n last_reviewed_by = entry.get(\"lastReviewedBy\")\n\n primary_owners: list[BasicExpertInfo] = [\n BasicExpertInfo(display_name=owner.get(\"name\"))\n for owner in [creator, last_updated_by]\n if owner is not None\n ]\n secondary_owners: list[BasicExpertInfo] = [\n BasicExpertInfo(display_name=owner.get(\"name\"))\n for owner in [last_reviewed_by]\n if owner is not None\n ]\n doc_batch.append(\n Document(\n id=str(entry[\"id\"]),\n sections=[TextSection(link=link, text=content_text)],\n source=DocumentSource.LOOPIO,\n semantic_identifier=questions[0],\n doc_updated_at=latest_time,\n primary_owners=primary_owners,\n secondary_owners=secondary_owners,\n metadata={\n \"topic\": topic,\n \"questions\": \"\\n\".join(questions),\n \"creator\": creator.get(\"name\") if creator else \"\",\n },\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n if len(doc_batch) > 0:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._process_entries()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_time = datetime.fromtimestamp(start, tz=timezone.utc).isoformat(\n timespec=\"seconds\"\n )\n end_time = datetime.fromtimestamp(end, tz=timezone.utc).isoformat(\n timespec=\"seconds\"\n )\n\n return self._process_entries(start_time, end_time)\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = LoopioConnector(\n loopio_stack_name=os.environ.get(\"LOOPIO_STACK_NAME\", None)\n )\n connector.load_credentials(\n {\n \"loopio_client_id\": os.environ[\"LOOPIO_CLIENT_ID\"],\n \"loopio_client_token\": os.environ[\"LOOPIO_CLIENT_TOKEN\"],\n \"loopio_subdomain\": os.environ[\"LOOPIO_SUBDOMAIN\"],\n }\n )\n\n latest_docs = connector.load_from_state()\n print(next(latest_docs))\n" }, { "path": "backend/onyx/connectors/mediawiki/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/mediawiki/family.py", "content": "from __future__ import annotations\n\nimport builtins\nimport functools\nimport itertools\nimport tempfile\nfrom typing import Any\nfrom unittest import mock\nfrom urllib.parse import urlparse\nfrom urllib.parse import urlunparse\n\nfrom pywikibot import family # type: ignore[import-untyped]\nfrom pywikibot import pagegenerators\nfrom pywikibot.scripts import generate_family_file # type: ignore[import-untyped]\nfrom pywikibot.scripts.generate_user_files import pywikibot # type: ignore[import-untyped]\n\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\npywikibot.config.base_dir = tempfile.TemporaryDirectory().name\n\n\n@mock.patch.object(\n builtins, \"print\", lambda *args: logger.info(\"\\t\".join(map(str, args)))\n)\nclass FamilyFileGeneratorInMemory(generate_family_file.FamilyFileGenerator):\n \"\"\"A subclass of FamilyFileGenerator that writes the family file to memory instead of to disk.\"\"\"\n\n def __init__(\n self,\n url: str,\n name: str,\n dointerwiki: str | bool = True,\n verify: str | bool = True,\n ):\n \"\"\"Initialize the FamilyFileGeneratorInMemory.\"\"\"\n\n url_parse = urlparse(url, \"https\")\n if not url_parse.netloc and url_parse.path:\n url = urlunparse(\n (url_parse.scheme, url_parse.path, url_parse.netloc, *url_parse[3:])\n )\n else:\n url = urlunparse(url_parse)\n assert isinstance(url, str)\n\n if any(x not in generate_family_file.NAME_CHARACTERS for x in name):\n raise ValueError(\n f'ERROR: Name of family \"{name}\" must be ASCII letters and digits [a-zA-Z0-9]',\n )\n\n if isinstance(dointerwiki, bool):\n dointerwiki = \"Y\" if dointerwiki else \"N\"\n assert isinstance(dointerwiki, str)\n\n if isinstance(verify, bool):\n verify = \"Y\" if verify else \"N\"\n assert isinstance(verify, str)\n\n super().__init__(url, name, dointerwiki, verify)\n self.family_definition: type[family.Family] | None = None\n\n def get_params(self) -> bool:\n \"\"\"Get the parameters for the family class definition.\n\n This override prevents the method from prompting the user for input (which would be impossible in this context).\n We do all the input validation in the constructor.\n \"\"\"\n return True\n\n def writefile(self, verify: Any) -> None: # noqa: ARG002\n \"\"\"Write the family file.\n\n This overrides the method in the parent class to write the family definition to memory instead of to disk.\n\n Args:\n verify: unused argument necessary to match the signature of the method in the parent class.\n \"\"\"\n code_hostname_pairs = {\n f\"{k}\": f\"{urlparse(w.server).netloc}\" for k, w in self.wikis.items()\n }\n\n code_path_pairs = {f\"{k}\": f\"{w.scriptpath}\" for k, w in self.wikis.items()}\n\n code_protocol_pairs = {\n f\"{k}\": f\"{urlparse(w.server).scheme}\" for k, w in self.wikis.items()\n }\n\n class Family(family.Family): # noqa: D101\n \"\"\"The family definition for the wiki.\"\"\"\n\n name = \"%(name)s\"\n langs = code_hostname_pairs\n\n def scriptpath(self, code: str) -> str:\n return code_path_pairs[code]\n\n def protocol(self, code: str) -> str:\n return code_protocol_pairs[code]\n\n self.family_definition = Family\n\n\n@functools.lru_cache(maxsize=None)\ndef generate_family_class(url: str, name: str) -> type[family.Family]:\n \"\"\"Generate a family file for a given URL and name.\n\n Args:\n url: The URL of the wiki.\n name: The short name of the wiki (customizable by the user).\n\n Returns:\n The family definition.\n\n Raises:\n ValueError: If the family definition was not generated.\n \"\"\"\n\n generator = FamilyFileGeneratorInMemory(url, name, \"Y\", \"Y\")\n generator.run()\n if generator.family_definition is None:\n raise ValueError(\"Family definition was not generated.\")\n return generator.family_definition\n\n\ndef family_class_dispatch(url: str, name: str) -> type[family.Family]:\n \"\"\"Find or generate a family class for a given URL and name.\n\n Args:\n url: The URL of the wiki.\n name: The short name of the wiki (customizable by the user).\n\n \"\"\"\n if \"wikipedia\" in url:\n import pywikibot.families.wikipedia_family # type: ignore[import-untyped]\n\n return pywikibot.families.wikipedia_family.Family\n # TODO: Support additional families pre-defined in `pywikibot.families.*_family.py` files\n return generate_family_class(url, name)\n\n\nif __name__ == \"__main__\":\n url = \"fallout.fandom.com/wiki/Fallout_Wiki\"\n name = \"falloutfandom\"\n\n categories: list[str] = []\n pages = [\"Fallout: New Vegas\"]\n recursion_depth = 1\n family_type = generate_family_class(url, name)\n\n site = pywikibot.Site(fam=family_type(), code=\"en\")\n categories = [\n pywikibot.Category(site, f\"Category:{category.replace(' ', '_')}\")\n for category in categories\n ]\n pages = [pywikibot.Page(site, page) for page in pages]\n all_pages = itertools.chain(\n pages,\n *[\n pagegenerators.CategorizedPageGenerator(category, recurse=recursion_depth)\n for category in categories\n ],\n )\n for page in all_pages:\n print(page.title())\n print(page.text[:1000])\n" }, { "path": "backend/onyx/connectors/mediawiki/wiki.py", "content": "from __future__ import annotations\n\nimport datetime\nimport itertools\nimport tempfile\nfrom collections.abc import Iterator\nfrom typing import Any\nfrom typing import cast\nfrom typing import ClassVar\n\nimport pywikibot.time # type: ignore[import-untyped]\nfrom pywikibot import pagegenerators\nfrom pywikibot import textlib\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.mediawiki.family import family_class_dispatch\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\npywikibot.config.base_dir = tempfile.TemporaryDirectory().name\n\n\ndef pywikibot_timestamp_to_utc_datetime(\n timestamp: pywikibot.time.Timestamp,\n) -> datetime.datetime:\n \"\"\"Convert a pywikibot timestamp to a datetime object in UTC.\n\n Args:\n timestamp: The pywikibot timestamp to convert.\n\n Returns:\n A datetime object in UTC.\n \"\"\"\n return datetime.datetime.astimezone(timestamp, tz=datetime.timezone.utc)\n\n\ndef get_doc_from_page(\n page: pywikibot.Page, site: pywikibot.Site | None, source_type: DocumentSource\n) -> Document:\n \"\"\"Generate Onyx Document from a MediaWiki page object.\n\n Args:\n page: Page from a MediaWiki site.\n site: MediaWiki site (used to parse the sections of the page using the site template, if available).\n source_type: Source of the document.\n\n Returns:\n Generated document.\n \"\"\"\n page_text = page.text\n sections_extracted: textlib.Content = textlib.extract_sections(page_text, site)\n\n sections = [\n TextSection(\n link=f\"{page.full_url()}#\" + section.heading.replace(\" \", \"_\"),\n text=section.title + section.content,\n )\n for section in sections_extracted.sections\n ]\n sections.append(\n TextSection(\n link=page.full_url(),\n text=sections_extracted.header,\n )\n )\n\n return Document(\n source=source_type,\n title=page.title(),\n doc_updated_at=pywikibot_timestamp_to_utc_datetime(\n page.latest_revision.timestamp\n ),\n sections=cast(list[TextSection | ImageSection], sections),\n semantic_identifier=page.title(),\n metadata={\"categories\": [category.title() for category in page.categories()]},\n id=f\"MEDIAWIKI_{page.pageid}_{page.full_url()}\",\n )\n\n\nclass MediaWikiConnector(LoadConnector, PollConnector):\n \"\"\"A connector for MediaWiki wikis.\n\n Args:\n hostname: The hostname of the wiki.\n categories: The categories to include in the index.\n pages: The pages to include in the index.\n recurse_depth: The depth to recurse into categories. -1 means unbounded recursion.\n language_code: The language code of the wiki.\n batch_size: The batch size for loading documents.\n\n Raises:\n ValueError: If `recurse_depth` is not an integer greater than or equal to -1.\n \"\"\"\n\n document_source_type: ClassVar[DocumentSource] = DocumentSource.MEDIAWIKI\n \"\"\"DocumentSource type for all documents generated by instances of this class. Can be overridden for connectors\n tailored for specific sites.\"\"\"\n\n def __init__(\n self,\n hostname: str,\n categories: list[str],\n pages: list[str],\n recurse_depth: int,\n language_code: str = \"en\",\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n if recurse_depth < -1:\n raise ValueError(\n f\"recurse_depth must be an integer greater than or equal to -1. Got {recurse_depth} instead.\"\n )\n # -1 means infinite recursion, which `pywikibot` will only do with `True`\n self.recurse_depth: bool | int = True if recurse_depth == -1 else recurse_depth\n\n self.batch_size = batch_size\n\n # short names can only have ascii letters and digits\n self.family = family_class_dispatch(hostname, \"WikipediaConnector\")()\n self.site = pywikibot.Site(fam=self.family, code=language_code)\n self.categories = [\n pywikibot.Category(\n self.site,\n (\n f\"{category.replace(' ', '_')}\"\n if category.startswith(\"Category:\")\n else f\"Category:{category.replace(' ', '_')}\"\n ),\n )\n for category in categories\n ]\n\n self.pages = []\n for page in pages:\n if not page:\n continue\n self.pages.append(pywikibot.Page(self.site, page))\n\n def load_credentials(\n self,\n credentials: dict[str, Any], # noqa: ARG002\n ) -> dict[str, Any] | None:\n \"\"\"Load credentials for a MediaWiki site.\n\n Note:\n For most read-only operations, MediaWiki API credentials are not necessary.\n This method can be overridden in the event that a particular MediaWiki site\n requires credentials.\n \"\"\"\n return None\n\n def _get_doc_batch(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> GenerateDocumentsOutput:\n \"\"\"Request batches of pages from a MediaWiki site.\n\n Args:\n start: The beginning of the time period of pages to request.\n end: The end of the time period of pages to request.\n\n Yields:\n Lists of Documents containing each parsed page in a batch.\n \"\"\"\n doc_batch: list[Document | HierarchyNode] = []\n\n # Pywikibot can handle batching for us, including only loading page contents when we finally request them.\n category_pages = [\n pagegenerators.PreloadingGenerator(\n pagegenerators.EdittimeFilterPageGenerator(\n pagegenerators.CategorizedPageGenerator(\n category, recurse=self.recurse_depth\n ),\n last_edit_start=(\n datetime.datetime.fromtimestamp(start) if start else None\n ),\n last_edit_end=datetime.datetime.fromtimestamp(end) if end else None,\n ),\n groupsize=self.batch_size,\n )\n for category in self.categories\n ]\n\n # Since we can specify both individual pages and categories, we need to iterate over all of them.\n all_pages: Iterator[pywikibot.Page] = itertools.chain(\n self.pages, *category_pages\n )\n for page in all_pages:\n logger.info(\n f\"MediaWikiConnector: title='{page.title()}' url={page.full_url()}\"\n )\n doc_batch.append(\n get_doc_from_page(page, self.site, self.document_source_type)\n )\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"Load all documents from the source.\n\n Returns:\n A generator of documents.\n \"\"\"\n return self.poll_source(None, None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch | None, end: SecondsSinceUnixEpoch | None\n ) -> GenerateDocumentsOutput:\n \"\"\"Poll the source for new documents.\n\n Args:\n start: The start of the time range to poll.\n end: The end of the time range to poll.\n\n Returns:\n A generator of documents.\n \"\"\"\n return self._get_doc_batch(start, end)\n\n\nif __name__ == \"__main__\":\n HOSTNAME = \"fallout.fandom.com\"\n test_connector = MediaWikiConnector(\n hostname=HOSTNAME,\n categories=[\"Fallout:_New_Vegas_factions\"],\n pages=[\"Fallout: New Vegas\"],\n recurse_depth=1,\n )\n\n all_docs = list(test_connector.load_from_state())\n print(\"All docs\", all_docs)\n current = datetime.datetime.now().timestamp()\n one_day_ago = current - 30 * 24 * 60 * 60 # 30 days\n\n latest_docs = list(test_connector.poll_source(one_day_ago, current))\n\n print(\"Latest docs\", latest_docs)\n" }, { "path": "backend/onyx/connectors/microsoft_graph_env.py", "content": "\"\"\"Inverse mapping from user-facing Microsoft host URLs to the SDK's AzureEnvironment.\n\nThe office365 library's GraphClient requires an ``AzureEnvironment`` string\n(e.g. ``\"Global\"``, ``\"GCC High\"``) to route requests to the correct national\ncloud. Our connectors instead expose free-text ``authority_host`` and\n``graph_api_host`` fields so the frontend doesn't need to know about SDK\ninternals.\n\nThis module bridges the gap: given the two host URLs the user configured, it\nresolves the matching ``AzureEnvironment`` value (and the implied SharePoint\ndomain suffix) so callers can pass ``environment=…`` to ``GraphClient``.\n\"\"\"\n\nfrom office365.graph_client import AzureEnvironment # type: ignore[import-untyped]\nfrom pydantic import BaseModel\n\nfrom onyx.connectors.exceptions import ConnectorValidationError\n\n\nclass MicrosoftGraphEnvironment(BaseModel):\n \"\"\"One row of the inverse mapping.\"\"\"\n\n environment: str\n graph_host: str\n authority_host: str\n sharepoint_domain_suffix: str\n\n\n_ENVIRONMENTS: list[MicrosoftGraphEnvironment] = [\n MicrosoftGraphEnvironment(\n environment=AzureEnvironment.Global,\n graph_host=\"https://graph.microsoft.com\",\n authority_host=\"https://login.microsoftonline.com\",\n sharepoint_domain_suffix=\"sharepoint.com\",\n ),\n MicrosoftGraphEnvironment(\n environment=AzureEnvironment.USGovernmentHigh,\n graph_host=\"https://graph.microsoft.us\",\n authority_host=\"https://login.microsoftonline.us\",\n sharepoint_domain_suffix=\"sharepoint.us\",\n ),\n MicrosoftGraphEnvironment(\n environment=AzureEnvironment.USGovernmentDoD,\n graph_host=\"https://dod-graph.microsoft.us\",\n authority_host=\"https://login.microsoftonline.us\",\n sharepoint_domain_suffix=\"sharepoint.us\",\n ),\n MicrosoftGraphEnvironment(\n environment=AzureEnvironment.China,\n graph_host=\"https://microsoftgraph.chinacloudapi.cn\",\n authority_host=\"https://login.chinacloudapi.cn\",\n sharepoint_domain_suffix=\"sharepoint.cn\",\n ),\n MicrosoftGraphEnvironment(\n environment=AzureEnvironment.Germany,\n graph_host=\"https://graph.microsoft.de\",\n authority_host=\"https://login.microsoftonline.de\",\n sharepoint_domain_suffix=\"sharepoint.de\",\n ),\n]\n\n_GRAPH_HOST_INDEX: dict[str, MicrosoftGraphEnvironment] = {\n env.graph_host: env for env in _ENVIRONMENTS\n}\n\n\ndef resolve_microsoft_environment(\n graph_api_host: str,\n authority_host: str,\n) -> MicrosoftGraphEnvironment:\n \"\"\"Return the ``MicrosoftGraphEnvironment`` that matches the supplied hosts.\n\n Raises ``ConnectorValidationError`` when the combination is unknown or\n internally inconsistent (e.g. a GCC-High graph host paired with a\n commercial authority host).\n \"\"\"\n graph_api_host = graph_api_host.rstrip(\"/\")\n authority_host = authority_host.rstrip(\"/\")\n\n env = _GRAPH_HOST_INDEX.get(graph_api_host)\n if env is None:\n known = \", \".join(sorted(_GRAPH_HOST_INDEX))\n raise ConnectorValidationError(\n f\"Unsupported Microsoft Graph API host '{graph_api_host}'. Recognised hosts: {known}\"\n )\n\n if env.authority_host != authority_host:\n raise ConnectorValidationError(\n f\"Authority host '{authority_host}' is inconsistent with \"\n f\"graph API host '{graph_api_host}'. \"\n f\"Expected authority host '{env.authority_host}' \"\n f\"for the {env.environment} environment.\"\n )\n\n return env\n" }, { "path": "backend/onyx/connectors/mock_connector/connector.py", "content": "from typing import Any\n\nimport httpx\nfrom pydantic import BaseModel\nfrom typing_extensions import override\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\nEXTERNAL_USER_EMAILS = {\"test@example.com\", \"admin@example.com\"}\nEXTERNAL_USER_GROUP_IDS = {\"mock-group-1\", \"mock-group-2\"}\n\n\nclass MockConnectorCheckpoint(ConnectorCheckpoint):\n last_document_id: str | None = None\n\n\nclass SingleConnectorYield(BaseModel):\n documents: list[Document]\n checkpoint: MockConnectorCheckpoint\n failures: list[ConnectorFailure]\n unhandled_exception: str | None = None\n\n\nclass MockConnector(CheckpointedConnectorWithPermSync[MockConnectorCheckpoint]):\n def __init__(\n self,\n mock_server_host: str,\n mock_server_port: int,\n ) -> None:\n self.mock_server_host = mock_server_host\n self.mock_server_port = mock_server_port\n self.client = httpx.Client(timeout=30.0)\n\n self.connector_yields: list[SingleConnectorYield] | None = None\n self.current_yield_index: int = 0\n\n def load_credentials(\n self,\n credentials: dict[str, Any], # noqa: ARG002\n ) -> dict[str, Any] | None:\n response = self.client.get(self._get_mock_server_url(\"get-documents\"))\n response.raise_for_status()\n data = response.json()\n\n self.connector_yields = [\n SingleConnectorYield(**yield_data) for yield_data in data\n ]\n return None\n\n def _get_mock_server_url(self, endpoint: str) -> str:\n return f\"http://{self.mock_server_host}:{self.mock_server_port}/{endpoint}\"\n\n def _save_checkpoint(self, checkpoint: MockConnectorCheckpoint) -> None:\n response = self.client.post(\n self._get_mock_server_url(\"add-checkpoint\"),\n json=checkpoint.model_dump(mode=\"json\"),\n )\n response.raise_for_status()\n\n def _load_from_checkpoint_common(\n self,\n start: SecondsSinceUnixEpoch, # noqa: ARG002\n end: SecondsSinceUnixEpoch, # noqa: ARG002\n checkpoint: MockConnectorCheckpoint,\n include_permissions: bool = False,\n ) -> CheckpointOutput[MockConnectorCheckpoint]:\n if self.connector_yields is None:\n raise ValueError(\"No connector yields configured\")\n\n # Save the checkpoint to the mock server\n self._save_checkpoint(checkpoint)\n\n yield_index = self.current_yield_index\n self.current_yield_index += 1\n current_yield = self.connector_yields[yield_index]\n\n # If the current yield has an unhandled exception, raise it\n # This is used to simulate an unhandled failure in the connector.\n if current_yield.unhandled_exception:\n raise RuntimeError(current_yield.unhandled_exception)\n\n # yield all documents\n for document in current_yield.documents:\n # If permissions are requested and not already set, add mock permissions\n if include_permissions and document.external_access is None:\n # Add mock permissions - make documents accessible to specific users/groups\n document.external_access = ExternalAccess(\n external_user_emails=EXTERNAL_USER_EMAILS,\n external_user_group_ids=EXTERNAL_USER_GROUP_IDS,\n is_public=False,\n )\n yield document\n\n for failure in current_yield.failures:\n yield failure\n\n return current_yield.checkpoint\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: MockConnectorCheckpoint,\n ) -> CheckpointOutput[MockConnectorCheckpoint]:\n return self._load_from_checkpoint_common(\n start, end, checkpoint, include_permissions=False\n )\n\n @override\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: MockConnectorCheckpoint,\n ) -> CheckpointOutput[MockConnectorCheckpoint]:\n return self._load_from_checkpoint_common(\n start, end, checkpoint, include_permissions=True\n )\n\n @override\n def build_dummy_checkpoint(self) -> MockConnectorCheckpoint:\n return MockConnectorCheckpoint(\n has_more=True,\n last_document_id=None,\n )\n\n def validate_checkpoint_json(self, checkpoint_json: str) -> MockConnectorCheckpoint:\n return MockConnectorCheckpoint.model_validate_json(checkpoint_json)\n" }, { "path": "backend/onyx/connectors/models.py", "content": "import sys\nfrom datetime import datetime\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom pydantic import field_validator\nfrom pydantic import model_validator\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import INDEX_SEPARATOR\nfrom onyx.configs.constants import RETURN_SEPARATOR\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.utils.text_processing import make_url_compatible\n\n\nclass InputType(str, Enum):\n LOAD_STATE = \"load_state\" # e.g. loading a current full state or a save state, such as from a file\n POLL = \"poll\" # e.g. calling an API to get all documents in the last hour\n EVENT = \"event\" # e.g. registered an endpoint as a listener, and processing connector events\n SLIM_RETRIEVAL = \"slim_retrieval\"\n\n\nclass ConnectorMissingCredentialError(PermissionError):\n def __init__(self, connector_name: str) -> None:\n connector_name = connector_name or \"Unknown\"\n super().__init__(\n f\"{connector_name} connector missing credentials, was load_credentials called?\"\n )\n\n\nclass Section(BaseModel):\n \"\"\"Base section class with common attributes\"\"\"\n\n link: str | None = None\n text: str | None = None\n image_file_id: str | None = None\n\n\nclass TextSection(Section):\n \"\"\"Section containing text content\"\"\"\n\n text: str\n\n def __sizeof__(self) -> int:\n return sys.getsizeof(self.text) + sys.getsizeof(self.link)\n\n\nclass ImageSection(Section):\n \"\"\"Section containing an image reference\"\"\"\n\n image_file_id: str\n\n def __sizeof__(self) -> int:\n return sys.getsizeof(self.image_file_id) + sys.getsizeof(self.link)\n\n\nclass BasicExpertInfo(BaseModel):\n \"\"\"Basic Information for the owner of a document, any of the fields can be left as None\n Display fallback goes as follows:\n - first_name + (optional middle_initial) + last_name\n - display_name\n - email\n - first_name\n \"\"\"\n\n display_name: str | None = None\n first_name: str | None = None\n middle_initial: str | None = None\n last_name: str | None = None\n email: str | None = None\n\n def get_semantic_name(self) -> str:\n if self.first_name and self.last_name:\n name_parts = [self.first_name]\n if self.middle_initial:\n name_parts.append(self.middle_initial + \".\")\n name_parts.append(self.last_name)\n return \" \".join([name_part.capitalize() for name_part in name_parts])\n\n if self.display_name:\n return self.display_name\n\n if self.email:\n return self.email\n\n if self.first_name:\n return self.first_name.capitalize()\n\n return \"Unknown\"\n\n def get_email(self) -> str | None:\n return self.email or None\n\n def __eq__(self, other: Any) -> bool:\n if not isinstance(other, BasicExpertInfo):\n return False\n return (\n self.display_name,\n self.first_name,\n self.middle_initial,\n self.last_name,\n self.email,\n ) == (\n other.display_name,\n other.first_name,\n other.middle_initial,\n other.last_name,\n other.email,\n )\n\n def __hash__(self) -> int:\n return hash(\n (\n self.display_name,\n self.first_name,\n self.middle_initial,\n self.last_name,\n self.email,\n )\n )\n\n def __sizeof__(self) -> int:\n size = sys.getsizeof(self.display_name)\n size += sys.getsizeof(self.first_name)\n size += sys.getsizeof(self.middle_initial)\n size += sys.getsizeof(self.last_name)\n size += sys.getsizeof(self.email)\n return size\n\n @classmethod\n def from_dict(cls, model_dict: dict[str, Any]) -> \"BasicExpertInfo\":\n\n first_name = cast(str, model_dict.get(\"FirstName\"))\n last_name = cast(str, model_dict.get(\"LastName\"))\n email = cast(str, model_dict.get(\"Email\"))\n display_name = cast(str, model_dict.get(\"Name\"))\n\n # Check if all fields are None\n if (\n first_name is None\n and last_name is None\n and email is None\n and display_name is None\n ):\n raise ValueError(\"No identifying information found for user\")\n\n return cls(\n first_name=first_name,\n last_name=last_name,\n email=email,\n display_name=display_name,\n )\n\n\nclass DocumentBase(BaseModel):\n \"\"\"Used for Onyx ingestion api, the ID is inferred before use if not provided\"\"\"\n\n id: str | None = None\n sections: list[TextSection | ImageSection]\n source: DocumentSource | None = None\n semantic_identifier: str # displayed in the UI as the main identifier for the doc\n # TODO(andrei): Ideally we could improve this to where each value is just a\n # list of strings.\n metadata: dict[str, str | list[str]]\n\n @field_validator(\"metadata\", mode=\"before\")\n @classmethod\n def _coerce_metadata_values(cls, v: dict[str, Any]) -> dict[str, str | list[str]]:\n return {\n key: [str(item) for item in val] if isinstance(val, list) else str(val)\n for key, val in v.items()\n }\n\n # UTC time\n doc_updated_at: datetime | None = None\n chunk_count: int | None = None\n\n # Owner, creator, etc.\n primary_owners: list[BasicExpertInfo] | None = None\n # Assignee, space owner, etc.\n secondary_owners: list[BasicExpertInfo] | None = None\n # title is used for search whereas semantic_identifier is used for displaying in the UI\n # different because Slack message may display as #general but general should not be part\n # of the search, at least not in the same way as a document title should be for like Confluence\n # The default title is semantic_identifier though unless otherwise specified\n title: str | None = None\n from_ingestion_api: bool = False\n # Anything else that may be useful that is specific to this particular connector type that other\n # parts of the code may need. If you're unsure, this can be left as None\n additional_info: Any = None\n\n # only filled in EE for connectors w/ permission sync enabled\n external_access: ExternalAccess | None = None\n doc_metadata: dict[str, Any] | None = None\n\n # Parent hierarchy node raw ID - the folder/space/page containing this document\n # If None, document's hierarchy position is unknown or connector doesn't support hierarchy\n parent_hierarchy_raw_node_id: str | None = None\n\n # Resolved database ID of the parent hierarchy node\n # Set during docfetching after hierarchy nodes are cached\n parent_hierarchy_node_id: int | None = None\n\n def get_title_for_document_index(\n self,\n ) -> str | None:\n # If title is explicitly empty, return a None here for embedding purposes\n if self.title == \"\":\n return None\n replace_chars = set(RETURN_SEPARATOR)\n title = self.semantic_identifier if self.title is None else self.title\n for char in replace_chars:\n title = title.replace(char, \" \")\n title = title.strip()\n return title\n\n def get_metadata_str_attributes(self) -> list[str] | None:\n if not self.metadata:\n return None\n # Combined string for the key/value for easy filtering\n return convert_metadata_dict_to_list_of_strings(self.metadata)\n\n def __sizeof__(self) -> int:\n size = sys.getsizeof(self.id)\n for section in self.sections:\n size += sys.getsizeof(section)\n size += sys.getsizeof(self.source)\n size += sys.getsizeof(self.semantic_identifier)\n size += sys.getsizeof(self.doc_updated_at)\n size += sys.getsizeof(self.chunk_count)\n\n if self.primary_owners is not None:\n for primary_owner in self.primary_owners:\n size += sys.getsizeof(primary_owner)\n else:\n size += sys.getsizeof(self.primary_owners)\n\n if self.secondary_owners is not None:\n for secondary_owner in self.secondary_owners:\n size += sys.getsizeof(secondary_owner)\n else:\n size += sys.getsizeof(self.secondary_owners)\n\n size += sys.getsizeof(self.title)\n size += sys.getsizeof(self.from_ingestion_api)\n size += sys.getsizeof(self.additional_info)\n return size\n\n def get_text_content(self) -> str:\n return \" \".join([section.text for section in self.sections if section.text])\n\n\ndef convert_metadata_dict_to_list_of_strings(\n metadata: dict[str, str | list[str]],\n) -> list[str]:\n \"\"\"Converts a metadata dict to a list of strings.\n\n Each string is a key-value pair separated by the INDEX_SEPARATOR. If a key\n points to a list of values, each value generates a unique pair.\n\n NOTE: Whatever formatting strategy is used here to generate a key-value\n string must be replicated when constructing query filters.\n\n Args:\n metadata: The metadata dict to convert where values can be either a\n string or a list of strings.\n\n Returns:\n A list of strings where each string is a key-value pair separated by the\n INDEX_SEPARATOR.\n \"\"\"\n attributes: list[str] = []\n for k, v in metadata.items():\n if isinstance(v, list):\n attributes.extend([k + INDEX_SEPARATOR + vi for vi in v])\n else:\n attributes.append(k + INDEX_SEPARATOR + v)\n return attributes\n\n\ndef convert_metadata_list_of_strings_to_dict(\n metadata_list: list[str],\n) -> dict[str, str | list[str]]:\n \"\"\"\n Converts a list of strings to a metadata dict. The inverse of\n convert_metadata_dict_to_list_of_strings.\n\n Assumes the input strings are formatted as in the output of\n convert_metadata_dict_to_list_of_strings.\n\n The schema of the output metadata dict is suboptimal yet bound to legacy\n code. Ideally each key would just point to a list of strings, where each\n list might contain just one element.\n\n Args:\n metadata_list: The list of strings to convert to a metadata dict.\n\n Returns:\n A metadata dict where values can be either a string or a list of\n strings.\n \"\"\"\n metadata: dict[str, str | list[str]] = {}\n for item in metadata_list:\n key, value = item.split(INDEX_SEPARATOR, 1)\n if key in metadata:\n # We have already seen this key therefore it must point to a list.\n if isinstance(metadata[key], list):\n cast(list[str], metadata[key]).append(value)\n else:\n metadata[key] = [cast(str, metadata[key]), value]\n else:\n metadata[key] = value\n return metadata\n\n\nclass Document(DocumentBase):\n \"\"\"Used for Onyx ingestion api, the ID is required\"\"\"\n\n id: str\n source: DocumentSource\n\n def to_short_descriptor(self) -> str:\n \"\"\"Used when logging the identity of a document\"\"\"\n return f\"ID: '{self.id}'; Semantic ID: '{self.semantic_identifier}'\"\n\n @classmethod\n def from_base(cls, base: DocumentBase) -> \"Document\":\n return cls(\n id=(\n make_url_compatible(base.id)\n if base.id\n else \"ingestion_api_\" + make_url_compatible(base.semantic_identifier)\n ),\n sections=base.sections,\n source=base.source or DocumentSource.INGESTION_API,\n semantic_identifier=base.semantic_identifier,\n metadata=base.metadata,\n doc_updated_at=base.doc_updated_at,\n primary_owners=base.primary_owners,\n secondary_owners=base.secondary_owners,\n title=base.title,\n from_ingestion_api=base.from_ingestion_api,\n )\n\n def __sizeof__(self) -> int:\n size = super().__sizeof__()\n size += sys.getsizeof(self.id)\n size += sys.getsizeof(self.source)\n return size\n\n\nclass IndexingDocument(Document):\n \"\"\"Document with processed sections for indexing\"\"\"\n\n processed_sections: list[Section] = []\n\n def get_total_char_length(self) -> int:\n \"\"\"Get the total character length of the document including processed sections\"\"\"\n title_len = len(self.title or self.semantic_identifier)\n\n # Use processed_sections if available, otherwise fall back to original sections\n if self.processed_sections:\n section_len = sum(\n len(section.text) if section.text is not None else 0\n for section in self.processed_sections\n )\n else:\n section_len = sum(\n (\n len(section.text)\n if isinstance(section, TextSection) and section.text is not None\n else 0\n )\n for section in self.sections\n )\n\n return title_len + section_len\n\n\nclass SlimDocument(BaseModel):\n id: str\n external_access: ExternalAccess | None = None\n parent_hierarchy_raw_node_id: str | None = None\n\n\nclass HierarchyNode(BaseModel):\n \"\"\"\n Hierarchy node yielded by connectors.\n\n This is the Pydantic model used by connectors, distinct from the\n SQLAlchemy HierarchyNode model in db/models.py. The connector runner\n layer converts this to the DB model when persisting to Postgres.\n \"\"\"\n\n # Raw identifier from the source system\n # e.g., \"1h7uWUR2BYZjtMfEXFt43tauj-Gp36DTPtwnsNuA665I\" for Google Drive\n raw_node_id: str\n\n # Raw ID of parent node, or None for SOURCE-level children (direct children of the source root)\n raw_parent_id: str | None = None\n\n # Human-readable name for display\n display_name: str\n\n # Link to view this node in the source system\n link: str | None = None\n\n # What kind of structural node this is (folder, space, page, etc.)\n node_type: HierarchyNodeType\n\n # If this hierarchy node represents a document (e.g., Confluence page),\n # The db model stores that doc's document_id. This gets set during docprocessing\n # after the document row is created. Matching is done by raw_node_id matching document.id.\n # so, we don't allow connectors to specify this as it would be unused\n # document_id: str | None = None\n\n # External access information for the node\n external_access: ExternalAccess | None = None\n\n\nclass IndexAttemptMetadata(BaseModel):\n connector_id: int\n credential_id: int\n batch_num: int | None = None\n attempt_id: int | None = None\n request_id: str | None = None\n\n # Work in progress: will likely contain metadata about cc pair / index attempt\n structured_id: str | None = None\n\n\nclass ConnectorCheckpoint(BaseModel):\n # TODO: maybe move this to something disk-based to handle extremely large checkpoints?\n has_more: bool\n\n def __str__(self) -> str:\n \"\"\"String representation of the checkpoint, with truncation for large checkpoint content.\"\"\"\n MAX_CHECKPOINT_CONTENT_CHARS = 1000\n\n content_str = self.model_dump_json()\n if len(content_str) > MAX_CHECKPOINT_CONTENT_CHARS:\n content_str = content_str[: MAX_CHECKPOINT_CONTENT_CHARS - 3] + \"...\"\n return content_str\n\n\nclass DocumentFailure(BaseModel):\n document_id: str\n document_link: str | None = None\n\n\nclass EntityFailure(BaseModel):\n entity_id: str\n missed_time_range: tuple[datetime, datetime] | None = None\n\n\nclass ConnectorFailure(BaseModel):\n failed_document: DocumentFailure | None = None\n failed_entity: EntityFailure | None = None\n failure_message: str\n exception: Exception | None = Field(default=None, exclude=True)\n\n model_config = {\"arbitrary_types_allowed\": True}\n\n @model_validator(mode=\"before\")\n def check_failed_fields(cls, values: dict) -> dict:\n failed_document = values.get(\"failed_document\")\n failed_entity = values.get(\"failed_entity\")\n if (failed_document is None and failed_entity is None) or (\n failed_document is not None and failed_entity is not None\n ):\n raise ValueError(\n \"Exactly one of 'failed_document' or 'failed_entity' must be specified.\"\n )\n return values\n\n\nclass ConnectorStopSignal(Exception):\n \"\"\"A custom exception used to signal a stop in processing.\"\"\"\n\n\nclass OnyxMetadata(BaseModel):\n # Careful overriding the document_id, may cause visual issues in the UI.\n # Kept here for API based use cases mostly\n document_id: str | None = None\n source_type: DocumentSource | None = None\n link: str | None = None\n file_display_name: str | None = None\n primary_owners: list[BasicExpertInfo] | None = None\n secondary_owners: list[BasicExpertInfo] | None = None\n doc_updated_at: datetime | None = None\n title: str | None = None\n\n\nclass DocExtractionContext(BaseModel):\n index_name: str\n cc_pair_id: int\n connector_id: int\n credential_id: int\n source: DocumentSource\n earliest_index_time: float\n from_beginning: bool\n is_primary: bool\n should_fetch_permissions_during_indexing: bool\n search_settings_status: IndexModelStatus\n doc_extraction_complete_batch_num: int | None\n\n\nclass DocIndexingContext(BaseModel):\n batches_done: int\n total_failures: int\n net_doc_change: int\n total_chunks: int\n" }, { "path": "backend/onyx/connectors/notion/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/notion/connector.py", "content": "import re\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\nfrom typing import Optional\nfrom urllib.parse import parse_qs\nfrom urllib.parse import urlparse\n\nimport requests\nfrom pydantic import BaseModel\nfrom retry import retry\nfrom typing_extensions import override\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import NOTION_CONNECTOR_DISABLE_RECURSIVE_PAGE_LOOKUP\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rl_requests,\n)\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import NormalizationResult\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.utils.batching import batch_generator\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_NOTION_PAGE_SIZE = 100\n_NOTION_CALL_TIMEOUT = 30 # 30 seconds\n_MAX_PAGES = 1000\n\n\n# TODO: Tables need to be ingested, Pages need to have their metadata ingested\n\n\nclass NotionPage(BaseModel):\n \"\"\"Represents a Notion Page object\"\"\"\n\n id: str\n created_time: str\n last_edited_time: str\n in_trash: bool\n properties: dict[str, Any]\n url: str\n\n database_name: str | None = None # Only applicable to the database type page (wiki)\n parent: dict[str, Any] | None = (\n None # Raw parent object from API for hierarchy tracking\n )\n\n\nclass NotionDataSource(BaseModel):\n \"\"\"Represents a Notion Data Source within a database.\"\"\"\n\n id: str\n name: str = \"\"\n\n\nclass NotionBlock(BaseModel):\n \"\"\"Represents a Notion Block object\"\"\"\n\n id: str # Used for the URL\n text: str\n # In a plaintext representation of the page, how this block should be joined\n # with the existing text up to this point, separated out from text for clarity\n prefix: str\n\n\nclass NotionSearchResponse(BaseModel):\n \"\"\"Represents the response from the Notion Search API\"\"\"\n\n results: list[dict[str, Any]]\n next_cursor: Optional[str]\n has_more: bool = False\n\n\nclass BlockReadOutput(BaseModel):\n \"\"\"Output from reading blocks of a page.\"\"\"\n\n blocks: list[NotionBlock]\n child_page_ids: list[str]\n hierarchy_nodes: list[HierarchyNode]\n\n\nclass NotionConnector(LoadConnector, PollConnector):\n \"\"\"Notion Page connector that reads all Notion pages\n this integration has been granted access to.\n\n Arguments:\n batch_size (int): Number of objects to index in a batch\n \"\"\"\n\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n recursive_index_enabled: bool = not NOTION_CONNECTOR_DISABLE_RECURSIVE_PAGE_LOOKUP,\n root_page_id: str | None = None,\n ) -> None:\n \"\"\"Initialize with parameters.\"\"\"\n self.batch_size = batch_size\n self.headers = {\n \"Content-Type\": \"application/json\",\n \"Notion-Version\": \"2026-03-11\",\n }\n self.indexed_pages: set[str] = set()\n self.root_page_id = root_page_id\n # if enabled, will recursively index child pages as they are found rather\n # relying entirely on the `search` API. We have received reports that the\n # `search` API misses many pages - in those cases, this might need to be\n # turned on. It's not currently known why/when this is required.\n # NOTE: this also removes all benefits polling, since we need to traverse\n # all pages regardless of if they are updated. If the notion workspace is\n # very large, this may not be practical.\n self.recursive_index_enabled = recursive_index_enabled or self.root_page_id\n\n # Hierarchy tracking state\n self.seen_hierarchy_node_raw_ids: set[str] = set()\n self.workspace_id: str | None = None\n self.workspace_name: str | None = None\n # Maps child page IDs to their containing page ID (discovered in _read_blocks).\n # Used to resolve block_id parent types to the actual containing page.\n self._child_page_parent_map: dict[str, str] = {}\n # Maps data_source_id -> database_id (populated in _read_pages_from_database).\n # Used to resolve data_source_id parent types back to the database.\n self._data_source_to_database_map: dict[str, str] = {}\n\n @classmethod\n @override\n def normalize_url(cls, url: str) -> NormalizationResult:\n \"\"\"Normalize a Notion URL to extract the page ID (UUID format).\"\"\"\n parsed = urlparse(url)\n netloc = parsed.netloc.lower()\n\n if not (\"notion.so\" in netloc or \"notion.site\" in netloc):\n return NormalizationResult(normalized_url=None, use_default=False)\n\n # Extract page ID from path (format: \"Title-PageID\")\n path_last = parsed.path.split(\"/\")[-1]\n candidate = path_last.split(\"-\")[-1] if \"-\" in path_last else path_last\n\n # Clean and format as UUID\n candidate = re.sub(r\"[^0-9a-fA-F-]\", \"\", candidate)\n cleaned = candidate.replace(\"-\", \"\")\n\n if len(cleaned) == 32 and re.fullmatch(r\"[0-9a-fA-F]{32}\", cleaned):\n normalized_uuid = (\n f\"{cleaned[0:8]}-{cleaned[8:12]}-{cleaned[12:16]}-{cleaned[16:20]}-{cleaned[20:]}\"\n ).lower()\n return NormalizationResult(\n normalized_url=normalized_uuid, use_default=False\n )\n\n # Try query params\n params = parse_qs(parsed.query)\n for key in (\"p\", \"page_id\"):\n if key in params and params[key]:\n candidate = params[key][0].replace(\"-\", \"\")\n if len(candidate) == 32 and re.fullmatch(r\"[0-9a-fA-F]{32}\", candidate):\n normalized_uuid = (\n f\"{candidate[0:8]}-{candidate[8:12]}-{candidate[12:16]}-{candidate[16:20]}-{candidate[20:]}\"\n ).lower()\n return NormalizationResult(\n normalized_url=normalized_uuid, use_default=False\n )\n\n return NormalizationResult(normalized_url=None, use_default=False)\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_child_blocks(\n self, block_id: str, cursor: str | None = None\n ) -> dict[str, Any] | None:\n \"\"\"Fetch all child blocks via the Notion API.\"\"\"\n logger.debug(f\"Fetching children of block with ID '{block_id}'\")\n block_url = f\"https://api.notion.com/v1/blocks/{block_id}/children\"\n query_params = None if not cursor else {\"start_cursor\": cursor}\n res = rl_requests.get(\n block_url,\n headers=self.headers,\n params=query_params,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n try:\n res.raise_for_status()\n except Exception as e:\n if res.status_code == 404:\n # this happens when a page is not shared with the integration\n # in this case, we should just ignore the page\n logger.error(\n f\"Unable to access block with ID '{block_id}'. \"\n f\"This is likely due to the block not being shared \"\n f\"with the Onyx integration. Exact exception:\\n\\n{e}\"\n )\n else:\n logger.exception(\n f\"Error fetching blocks with status code {res.status_code}: {res.json()}\"\n )\n\n # This can occasionally happen, the reason is unknown and cannot be reproduced on our internal Notion\n # Assuming this will not be a critical loss of data\n return None\n return res.json()\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_page(self, page_id: str) -> NotionPage:\n \"\"\"Fetch a page from its ID via the Notion API, retry with database if page fetch fails.\"\"\"\n logger.debug(f\"Fetching page for ID '{page_id}'\")\n page_url = f\"https://api.notion.com/v1/pages/{page_id}\"\n res = rl_requests.get(\n page_url,\n headers=self.headers,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n try:\n res.raise_for_status()\n except Exception as e:\n logger.warning(\n f\"Failed to fetch page, trying database for ID '{page_id}'. Exception: {e}\"\n )\n # Try fetching as a database if page fetch fails, this happens if the page is set to a wiki\n # it becomes a database from the notion perspective\n return self._fetch_database_as_page(page_id)\n return NotionPage(**res.json())\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_database_as_page(self, database_id: str) -> NotionPage:\n \"\"\"Attempt to fetch a database as a page.\n\n Note: As of API 2025-09-03, database objects no longer include\n `properties` (schema moved to individual data sources).\n \"\"\"\n logger.debug(f\"Fetching database for ID '{database_id}' as a page\")\n database_url = f\"https://api.notion.com/v1/databases/{database_id}\"\n res = rl_requests.get(\n database_url,\n headers=self.headers,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n try:\n res.raise_for_status()\n except Exception as e:\n logger.exception(f\"Error fetching database as page - {res.json()}\")\n raise e\n db_data = res.json()\n database_name = db_data.get(\"title\")\n database_name = (\n database_name[0].get(\"text\", {}).get(\"content\") if database_name else None\n )\n\n db_data.setdefault(\"properties\", {})\n\n return NotionPage(**db_data, database_name=database_name)\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_data_sources_for_database(\n self, database_id: str\n ) -> list[NotionDataSource]:\n \"\"\"Fetch the list of data sources for a database.\"\"\"\n logger.debug(f\"Fetching data sources for database '{database_id}'\")\n res = rl_requests.get(\n f\"https://api.notion.com/v1/databases/{database_id}\",\n headers=self.headers,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n try:\n res.raise_for_status()\n except Exception as e:\n if res.status_code in (403, 404):\n logger.error(\n f\"Unable to access database with ID '{database_id}'. \"\n f\"This is likely due to the database not being shared \"\n f\"with the Onyx integration. Exact exception:\\n{e}\"\n )\n return []\n logger.exception(f\"Error fetching database - {res.json()}\")\n raise e\n\n db_data = res.json()\n data_sources = db_data.get(\"data_sources\", [])\n return [\n NotionDataSource(id=ds[\"id\"], name=ds.get(\"name\", \"\"))\n for ds in data_sources\n if ds.get(\"id\")\n ]\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_data_source(\n self, data_source_id: str, cursor: str | None = None\n ) -> dict[str, Any]:\n \"\"\"Query a data source via POST /v1/data_sources/{id}/query.\"\"\"\n logger.debug(f\"Querying data source '{data_source_id}'\")\n url = f\"https://api.notion.com/v1/data_sources/{data_source_id}/query\"\n body = None if not cursor else {\"start_cursor\": cursor}\n res = rl_requests.post(\n url,\n headers=self.headers,\n json=body,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n try:\n res.raise_for_status()\n except Exception as e:\n if res.status_code in (403, 404):\n logger.error(\n f\"Unable to access data source with ID '{data_source_id}'. \"\n f\"This is likely due to it not being shared \"\n f\"with the Onyx integration. Exact exception:\\n{e}\"\n )\n return {\"results\": [], \"next_cursor\": None}\n logger.exception(f\"Error querying data source - {res.json()}\")\n raise e\n return res.json()\n\n @retry(tries=3, delay=1, backoff=2)\n def _fetch_workspace_info(self) -> tuple[str, str]:\n \"\"\"Fetch workspace ID and name from the bot user endpoint.\"\"\"\n res = rl_requests.get(\n \"https://api.notion.com/v1/users/me\",\n headers=self.headers,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n res.raise_for_status()\n data = res.json()\n bot = data.get(\"bot\", {})\n # workspace_id may be in bot object, fallback to user id\n workspace_id = bot.get(\"workspace_id\", data.get(\"id\"))\n workspace_name = bot.get(\"workspace_name\", \"Notion Workspace\")\n return workspace_id, workspace_name\n\n def _get_workspace_hierarchy_node(self) -> HierarchyNode | None:\n \"\"\"Get the workspace hierarchy node, fetching workspace info if needed.\n\n Returns None if the workspace node has already been yielded.\n \"\"\"\n if self.workspace_id is None:\n self.workspace_id, self.workspace_name = self._fetch_workspace_info()\n\n if self.workspace_id in self.seen_hierarchy_node_raw_ids:\n return None\n\n self.seen_hierarchy_node_raw_ids.add(self.workspace_id)\n return HierarchyNode(\n raw_node_id=self.workspace_id,\n raw_parent_id=None, # Parent is SOURCE (auto-created by system)\n display_name=self.workspace_name or \"Notion Workspace\",\n link=f\"https://notion.so/{self.workspace_id.replace('-', '')}\",\n node_type=HierarchyNodeType.WORKSPACE,\n )\n\n def _get_parent_raw_id(\n self, parent: dict[str, Any] | None, page_id: str | None = None\n ) -> str | None:\n \"\"\"Get the parent raw ID for hierarchy tracking.\n\n Returns workspace_id for top-level pages, or the direct parent ID for nested pages.\n\n Args:\n parent: The parent object from the Notion API\n page_id: The page's own ID, used to look up block_id parents in our cache\n \"\"\"\n if not parent:\n return self.workspace_id # Default to workspace if no parent info\n\n parent_type = parent.get(\"type\")\n\n if parent_type == \"workspace\":\n return self.workspace_id\n elif parent_type == \"block_id\":\n # Inline page in a block - resolve to the containing page if we discovered it\n if page_id and page_id in self._child_page_parent_map:\n return self._child_page_parent_map[page_id]\n # Fallback to workspace if we don't know the parent\n return self.workspace_id\n elif parent_type == \"data_source_id\":\n ds_id = parent.get(\"data_source_id\")\n if ds_id:\n return self._data_source_to_database_map.get(ds_id, self.workspace_id)\n elif parent_type in [\"page_id\", \"database_id\"]:\n return parent.get(parent_type)\n\n return self.workspace_id\n\n def _maybe_yield_hierarchy_node(\n self,\n raw_node_id: str,\n raw_parent_id: str | None,\n display_name: str,\n link: str | None,\n node_type: HierarchyNodeType,\n ) -> HierarchyNode | None:\n \"\"\"Create and return a hierarchy node if not already yielded.\n\n Args:\n raw_node_id: The raw ID of the node\n raw_parent_id: The raw ID of the parent node\n display_name: Human-readable name\n link: URL to the node in Notion\n node_type: Type of hierarchy node\n\n Returns:\n HierarchyNode if new, None if already yielded\n \"\"\"\n if raw_node_id in self.seen_hierarchy_node_raw_ids:\n return None\n self.seen_hierarchy_node_raw_ids.add(raw_node_id)\n return HierarchyNode(\n raw_node_id=raw_node_id,\n raw_parent_id=raw_parent_id,\n display_name=display_name,\n link=link,\n node_type=node_type,\n )\n\n @staticmethod\n def _properties_to_str(properties: dict[str, Any]) -> str:\n \"\"\"Converts Notion properties to a string\"\"\"\n\n def _recurse_list_properties(inner_list: list[Any]) -> str | None:\n list_properties: list[str | None] = []\n for item in inner_list:\n if item and isinstance(item, dict):\n list_properties.append(_recurse_properties(item))\n elif item and isinstance(item, list):\n list_properties.append(_recurse_list_properties(item))\n else:\n list_properties.append(str(item))\n return (\n \", \".join(\n [\n list_property\n for list_property in list_properties\n if list_property\n ]\n )\n or None\n )\n\n def _recurse_properties(inner_dict: dict[str, Any]) -> str | None:\n sub_inner_dict: dict[str, Any] | list[Any] | str = inner_dict\n while isinstance(sub_inner_dict, dict) and \"type\" in sub_inner_dict:\n type_name = sub_inner_dict[\"type\"]\n sub_inner_dict = sub_inner_dict[type_name]\n\n # If the innermost layer is None, the value is not set\n if not sub_inner_dict:\n return None\n\n # TODO there may be more types to handle here\n if isinstance(sub_inner_dict, list):\n return _recurse_list_properties(sub_inner_dict)\n elif isinstance(sub_inner_dict, str):\n # For some objects the innermost value could just be a string, not sure what causes this\n return sub_inner_dict\n elif isinstance(sub_inner_dict, dict):\n if \"name\" in sub_inner_dict:\n return sub_inner_dict[\"name\"]\n if \"content\" in sub_inner_dict:\n return sub_inner_dict[\"content\"]\n start = sub_inner_dict.get(\"start\")\n end = sub_inner_dict.get(\"end\")\n if start is not None:\n if end is not None:\n return f\"{start} - {end}\"\n return start\n elif end is not None:\n return f\"Until {end}\"\n\n if \"id\" in sub_inner_dict:\n # This is not useful to index, it's a reference to another Notion object\n # and this ID value in plaintext is useless outside of the Notion context\n logger.debug(\"Skipping Notion object id field property\")\n return None\n\n logger.debug(f\"Unreadable property from innermost prop: {sub_inner_dict}\")\n return None\n\n result = \"\"\n for prop_name, prop in properties.items():\n if not prop or not isinstance(prop, dict):\n continue\n\n try:\n inner_value = _recurse_properties(prop)\n except Exception as e:\n # This is not a critical failure, these properties are not the actual contents of the page\n # more similar to metadata\n logger.warning(f\"Error recursing properties for {prop_name}: {e}\")\n continue\n # Not a perfect way to format Notion database tables but there's no perfect representation\n # since this must be represented as plaintext\n if inner_value:\n result += f\"{prop_name}: {inner_value}\\t\"\n\n return result\n\n def _read_pages_from_database(\n self,\n database_id: str,\n database_parent_raw_id: str | None = None,\n database_name: str | None = None,\n ) -> BlockReadOutput:\n \"\"\"Returns blocks, page IDs, and hierarchy nodes from a database.\n\n Args:\n database_id: The ID of the database\n database_parent_raw_id: The raw ID of the database's parent (containing page or workspace)\n database_name: The name of the database (from child_database block title)\n \"\"\"\n result_blocks: list[NotionBlock] = []\n result_pages: list[str] = []\n hierarchy_nodes: list[HierarchyNode] = []\n\n # Create hierarchy node for this database if not already yielded.\n # Notion URLs omit dashes from UUIDs: https://notion.so/17ab3186873d418fb899c3f6a43f68de\n db_node = self._maybe_yield_hierarchy_node(\n raw_node_id=database_id,\n raw_parent_id=database_parent_raw_id or self.workspace_id,\n display_name=database_name or f\"Database {database_id}\",\n link=f\"https://notion.so/{database_id.replace('-', '')}\",\n node_type=HierarchyNodeType.DATABASE,\n )\n if db_node:\n hierarchy_nodes.append(db_node)\n\n # Discover all data sources under this database, then query each one.\n # Even legacy single-source databases have one entry in the array.\n data_sources = self._fetch_data_sources_for_database(database_id)\n if not data_sources:\n logger.warning(\n f\"Database '{database_id}' returned zero data sources — \"\n f\"no pages will be indexed from this database.\"\n )\n for ds in data_sources:\n self._data_source_to_database_map[ds.id] = database_id\n cursor = None\n while True:\n data = self._fetch_data_source(ds.id, cursor)\n\n for result in data[\"results\"]:\n obj_id = result[\"id\"]\n obj_type = result[\"object\"]\n text = self._properties_to_str(result.get(\"properties\", {}))\n if text:\n result_blocks.append(\n NotionBlock(id=obj_id, text=text, prefix=\"\\n\")\n )\n\n if not self.recursive_index_enabled:\n continue\n\n if obj_type == \"page\":\n logger.debug(\n f\"Found page with ID '{obj_id}' in database '{database_id}'\"\n )\n result_pages.append(result[\"id\"])\n elif obj_type == \"database\":\n logger.debug(\n f\"Found database with ID '{obj_id}' in database '{database_id}'\"\n )\n nested_db_title = result.get(\"title\", [])\n nested_db_name = None\n if nested_db_title and len(nested_db_title) > 0:\n nested_db_name = (\n nested_db_title[0].get(\"text\", {}).get(\"content\")\n )\n nested_output = self._read_pages_from_database(\n obj_id,\n database_parent_raw_id=database_id,\n database_name=nested_db_name,\n )\n result_pages.extend(nested_output.child_page_ids)\n hierarchy_nodes.extend(nested_output.hierarchy_nodes)\n\n if data[\"next_cursor\"] is None:\n break\n\n cursor = data[\"next_cursor\"]\n\n return BlockReadOutput(\n blocks=result_blocks,\n child_page_ids=result_pages,\n hierarchy_nodes=hierarchy_nodes,\n )\n\n def _read_blocks(\n self, base_block_id: str, containing_page_id: str | None = None\n ) -> BlockReadOutput:\n \"\"\"Reads all child blocks for the specified block.\n\n Args:\n base_block_id: The block ID to read children from\n containing_page_id: The ID of the page that contains this block tree.\n Used to correctly map child pages/databases to their parent page\n rather than intermediate block IDs.\n \"\"\"\n # If no containing_page_id provided, assume base_block_id is the page itself\n page_id = containing_page_id or base_block_id\n result_blocks: list[NotionBlock] = []\n child_pages: list[str] = []\n hierarchy_nodes: list[HierarchyNode] = []\n cursor = None\n while True:\n data = self._fetch_child_blocks(base_block_id, cursor)\n\n # this happens when a block is not shared with the integration\n if data is None:\n return BlockReadOutput(\n blocks=result_blocks,\n child_page_ids=child_pages,\n hierarchy_nodes=hierarchy_nodes,\n )\n\n for result in data[\"results\"]:\n logger.debug(\n f\"Found child block for block with ID '{base_block_id}': {result}\"\n )\n result_block_id = result[\"id\"]\n result_type = result[\"type\"]\n result_obj = result[result_type]\n\n if result_type == \"ai_block\":\n logger.warning(\n f\"Skipping 'ai_block' ('{result_block_id}') for base block '{base_block_id}': \"\n f\"Notion API does not currently support reading AI blocks (as of 24/02/09) \"\n f\"(discussion: https://github.com/onyx-dot-app/onyx/issues/1053)\"\n )\n continue\n\n if result_type == \"unsupported\":\n logger.warning(\n f\"Skipping unsupported block type '{result_type}' \"\n f\"('{result_block_id}') for base block '{base_block_id}': \"\n f\"(discussion: https://github.com/onyx-dot-app/onyx/issues/1230)\"\n )\n continue\n\n if result_type == \"external_object_instance_page\":\n logger.warning(\n f\"Skipping 'external_object_instance_page' ('{result_block_id}') for base block '{base_block_id}': \"\n f\"Notion API does not currently support reading external blocks (as of 24/07/03) \"\n f\"(discussion: https://github.com/onyx-dot-app/onyx/issues/1761)\"\n )\n continue\n\n cur_result_text_arr = []\n if \"rich_text\" in result_obj:\n for rich_text in result_obj[\"rich_text\"]:\n # skip if doesn't have text object\n if \"text\" in rich_text:\n text = rich_text[\"text\"][\"content\"]\n cur_result_text_arr.append(text)\n\n if result[\"has_children\"]:\n if result_type == \"child_page\":\n # Child pages will not be included at this top level, it will be a separate document.\n # Track parent page so we can resolve block_id parents later.\n # Use page_id (not base_block_id) to ensure we map to the containing page,\n # not an intermediate block like a toggle or callout.\n child_pages.append(result_block_id)\n self._child_page_parent_map[result_block_id] = page_id\n else:\n logger.debug(f\"Entering sub-block: {result_block_id}\")\n sub_output = self._read_blocks(result_block_id, page_id)\n logger.debug(f\"Finished sub-block: {result_block_id}\")\n result_blocks.extend(sub_output.blocks)\n child_pages.extend(sub_output.child_page_ids)\n hierarchy_nodes.extend(sub_output.hierarchy_nodes)\n\n if result_type == \"child_database\":\n # Extract database name from the child_database block\n db_title = result_obj.get(\"title\", \"\")\n db_output = self._read_pages_from_database(\n result_block_id,\n database_parent_raw_id=page_id, # Parent is the containing page\n database_name=db_title or None,\n )\n # A database on a page often looks like a table, we need to include it for the contents\n # of the page but the children (cells) should be processed as other Documents\n result_blocks.extend(db_output.blocks)\n hierarchy_nodes.extend(db_output.hierarchy_nodes)\n\n if self.recursive_index_enabled:\n child_pages.extend(db_output.child_page_ids)\n\n if cur_result_text_arr:\n new_block = NotionBlock(\n id=result_block_id,\n text=\"\\n\".join(cur_result_text_arr),\n prefix=\"\\n\",\n )\n result_blocks.append(new_block)\n\n if data[\"next_cursor\"] is None:\n break\n\n cursor = data[\"next_cursor\"]\n\n return BlockReadOutput(\n blocks=result_blocks,\n child_page_ids=child_pages,\n hierarchy_nodes=hierarchy_nodes,\n )\n\n def _read_page_title(self, page: NotionPage) -> str | None:\n \"\"\"Extracts the title from a Notion page\"\"\"\n page_title = None\n if hasattr(page, \"database_name\") and page.database_name:\n return page.database_name\n for _, prop in page.properties.items():\n if prop[\"type\"] == \"title\" and len(prop[\"title\"]) > 0:\n page_title = \" \".join([t[\"plain_text\"] for t in prop[\"title\"]]).strip()\n break\n\n return page_title\n\n def _read_pages(\n self,\n pages: list[NotionPage],\n ) -> Generator[Document | HierarchyNode, None, None]:\n \"\"\"Reads pages for rich text content and generates Documents and HierarchyNodes\n\n Note that a page which is turned into a \"wiki\" becomes a database but both top level pages and top level databases\n do not seem to have any properties associated with them.\n\n Pages that are part of a database can have properties which are like the values of the row in the \"database\" table\n in which they exist\n\n This is not clearly outlined in the Notion API docs but it is observable empirically.\n https://developers.notion.com/docs/working-with-page-content\n \"\"\"\n all_child_page_ids: list[str] = []\n for page in pages:\n if page.id in self.indexed_pages:\n logger.debug(f\"Already indexed page with ID '{page.id}'. Skipping.\")\n continue\n\n logger.info(f\"Reading page with ID '{page.id}', with url {page.url}\")\n block_output = self._read_blocks(page.id)\n all_child_page_ids.extend(block_output.child_page_ids)\n\n # okay to mark here since there's no way for this to not succeed\n # without a critical failure\n self.indexed_pages.add(page.id)\n\n raw_page_title = self._read_page_title(page)\n page_title = raw_page_title or f\"Untitled Page with ID {page.id}\"\n parent_raw_id = self._get_parent_raw_id(page.parent, page_id=page.id)\n\n # If this page has children (pages or databases), yield it as a hierarchy node FIRST\n # This ensures parent nodes are created before child documents reference them\n if block_output.child_page_ids or block_output.hierarchy_nodes:\n hierarchy_node = self._maybe_yield_hierarchy_node(\n raw_node_id=page.id,\n raw_parent_id=parent_raw_id,\n display_name=page_title,\n link=page.url,\n node_type=HierarchyNodeType.PAGE,\n )\n if hierarchy_node:\n yield hierarchy_node\n\n # Yield database hierarchy nodes discovered in this page's blocks\n for db_node in block_output.hierarchy_nodes:\n yield db_node\n\n if not block_output.blocks:\n if not raw_page_title:\n logger.warning(\n f\"No blocks OR title found for page with ID '{page.id}'. Skipping.\"\n )\n continue\n\n logger.debug(f\"No blocks found for page with ID '{page.id}'\")\n \"\"\"\n Something like:\n\n TITLE\n\n PROP1: PROP1_VALUE\n PROP2: PROP2_VALUE\n \"\"\"\n text = page_title\n if page.properties:\n text += \"\\n\\n\" + \"\\n\".join(\n [f\"{key}: {value}\" for key, value in page.properties.items()]\n )\n sections = [\n TextSection(\n link=f\"{page.url}\",\n text=text,\n )\n ]\n else:\n sections = [\n TextSection(\n link=f\"{page.url}#{block.id.replace('-', '')}\",\n text=block.prefix + block.text,\n )\n for block in block_output.blocks\n ]\n\n yield (\n Document(\n id=page.id,\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.NOTION,\n semantic_identifier=page_title,\n doc_updated_at=datetime.fromisoformat(\n page.last_edited_time\n ).astimezone(timezone.utc),\n metadata={},\n parent_hierarchy_raw_node_id=parent_raw_id,\n )\n )\n self.indexed_pages.add(page.id)\n\n if self.recursive_index_enabled and all_child_page_ids:\n # NOTE: checking if page_id is in self.indexed_pages to prevent extra\n # calls to `_fetch_page` for pages we've already indexed\n for child_page_batch_ids in batch_generator(\n all_child_page_ids, batch_size=INDEX_BATCH_SIZE\n ):\n child_page_batch = [\n self._fetch_page(page_id)\n for page_id in child_page_batch_ids\n if page_id not in self.indexed_pages\n ]\n yield from self._read_pages(child_page_batch)\n\n @retry(tries=3, delay=1, backoff=2)\n def _search_notion(self, query_dict: dict[str, Any]) -> NotionSearchResponse:\n \"\"\"Search for pages from a Notion database. Includes some small number of\n retries to handle misc, flakey failures.\"\"\"\n logger.debug(f\"Searching for pages in Notion with query_dict: {query_dict}\")\n res = rl_requests.post(\n \"https://api.notion.com/v1/search\",\n headers=self.headers,\n json=query_dict,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n res.raise_for_status()\n return NotionSearchResponse(**res.json())\n\n # The | Document is needed for mypy type checking\n def _yield_database_hierarchy_nodes(\n self,\n ) -> Generator[HierarchyNode | Document, None, None]:\n \"\"\"Search for all data sources and yield hierarchy nodes for their parent databases.\n\n This must be called BEFORE page indexing so that database hierarchy nodes\n exist when pages inside databases reference them as parents.\n\n With the new API, search returns data source objects instead of databases.\n Multiple data sources can share the same parent database, so we use\n database_id as the hierarchy node key and deduplicate via\n _maybe_yield_hierarchy_node.\n \"\"\"\n query_dict: dict[str, Any] = {\n \"filter\": {\"property\": \"object\", \"value\": \"data_source\"},\n \"page_size\": _NOTION_PAGE_SIZE,\n }\n pages_seen = 0\n while pages_seen < _MAX_PAGES:\n db_res = self._search_notion(query_dict)\n for ds in db_res.results:\n # Extract the parent database_id from the data source's parent\n ds_parent = ds.get(\"parent\", {})\n db_id = ds_parent.get(\"database_id\")\n if not db_id:\n continue\n\n # Populate the mapping so _get_parent_raw_id can resolve later\n ds_id = ds.get(\"id\")\n if not ds_id:\n continue\n self._data_source_to_database_map[ds_id] = db_id\n\n # Fetch the database to get its actual name and parent\n try:\n db_page = self._fetch_database_as_page(db_id)\n db_name = db_page.database_name or f\"Database {db_id}\"\n parent_raw_id = self._get_parent_raw_id(db_page.parent)\n db_url = (\n db_page.url or f\"https://notion.so/{db_id.replace('-', '')}\"\n )\n except requests.exceptions.RequestException as e:\n logger.warning(\n f\"Could not fetch database '{db_id}', \"\n f\"defaulting to workspace root. Error: {e}\"\n )\n db_name = f\"Database {db_id}\"\n parent_raw_id = self.workspace_id\n db_url = f\"https://notion.so/{db_id.replace('-', '')}\"\n\n # _maybe_yield_hierarchy_node deduplicates by raw_node_id,\n # so multiple data sources under one database produce one node.\n node = self._maybe_yield_hierarchy_node(\n raw_node_id=db_id,\n raw_parent_id=parent_raw_id or self.workspace_id,\n display_name=db_name,\n link=db_url,\n node_type=HierarchyNodeType.DATABASE,\n )\n if node:\n yield node\n\n if not db_res.has_more:\n break\n query_dict[\"start_cursor\"] = db_res.next_cursor\n pages_seen += 1\n\n def _filter_pages_by_time(\n self,\n pages: list[dict[str, Any]],\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n filter_field: str = \"last_edited_time\",\n ) -> list[NotionPage]:\n \"\"\"A helper function to filter out pages outside of a time\n range. This functionality doesn't yet exist in the Notion Search API,\n but when it does, this approach can be deprecated.\n\n Arguments:\n pages (list[dict]) - Pages to filter\n start (float) - start epoch time to filter from\n end (float) - end epoch time to filter to\n filter_field (str) - the attribute on the page to apply the filter\n \"\"\"\n filtered_pages: list[NotionPage] = []\n for page in pages:\n # Parse ISO 8601 timestamp and convert to UTC epoch time\n timestamp = page[filter_field].replace(\".000Z\", \"+00:00\")\n compare_time = datetime.fromisoformat(timestamp).timestamp()\n if compare_time > start and compare_time <= end:\n filtered_pages += [NotionPage(**page)]\n return filtered_pages\n\n def _recursive_load(self) -> GenerateDocumentsOutput:\n if self.root_page_id is None or not self.recursive_index_enabled:\n raise RuntimeError(\n \"Recursive page lookup is not enabled, but we are trying to recursively load pages. This should never happen.\"\n )\n\n # Yield workspace hierarchy node FIRST before any pages\n workspace_node = self._get_workspace_hierarchy_node()\n if workspace_node:\n yield [workspace_node]\n\n logger.info(\n f\"Recursively loading pages from Notion based on root page with ID: {self.root_page_id}\"\n )\n pages = [self._fetch_page(page_id=self.root_page_id)]\n yield from batch_generator(self._read_pages(pages), self.batch_size)\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n \"\"\"Applies integration token to headers\"\"\"\n self.headers[\"Authorization\"] = (\n f\"Bearer {credentials['notion_integration_token']}\"\n )\n return None\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"Loads all page data from a Notion workspace.\n\n Returns:\n list[Document]: list of documents.\n \"\"\"\n # TODO: remove once Notion search issue is discovered\n if self.recursive_index_enabled and self.root_page_id:\n yield from self._recursive_load()\n return\n\n # Yield workspace hierarchy node FIRST before any pages\n workspace_node = self._get_workspace_hierarchy_node()\n if workspace_node:\n yield [workspace_node]\n\n # Yield database hierarchy nodes BEFORE pages so parent references resolve\n yield from batch_generator(\n self._yield_database_hierarchy_nodes(), self.batch_size\n )\n\n query_dict: dict[str, Any] = {\n \"filter\": {\"property\": \"object\", \"value\": \"page\"},\n \"page_size\": _NOTION_PAGE_SIZE,\n }\n while True:\n db_res = self._search_notion(query_dict)\n pages = [NotionPage(**page) for page in db_res.results]\n yield from batch_generator(self._read_pages(pages), self.batch_size)\n if db_res.has_more:\n query_dict[\"start_cursor\"] = db_res.next_cursor\n else:\n break\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n \"\"\"Uses the Notion search API to fetch updated pages\n within a time period.\n Unfortunately the search API doesn't yet support filtering by times,\n so until they add that, we're just going to page through results until,\n we reach ones that are older than our search criteria.\n \"\"\"\n # TODO: remove once Notion search issue is discovered\n if self.recursive_index_enabled and self.root_page_id:\n yield from self._recursive_load()\n return\n\n # Yield workspace hierarchy node FIRST before any pages\n workspace_node = self._get_workspace_hierarchy_node()\n if workspace_node:\n yield [workspace_node]\n\n # Yield database hierarchy nodes BEFORE pages so parent references resolve.\n # We yield all databases without time filtering because a page's parent\n # database might not have been edited even if the page was.\n yield from batch_generator(\n self._yield_database_hierarchy_nodes(), self.batch_size\n )\n\n query_dict: dict[str, Any] = {\n \"page_size\": _NOTION_PAGE_SIZE,\n \"sort\": {\"timestamp\": \"last_edited_time\", \"direction\": \"descending\"},\n \"filter\": {\"property\": \"object\", \"value\": \"page\"},\n }\n while True:\n db_res = self._search_notion(query_dict)\n pages = self._filter_pages_by_time(\n db_res.results, start, end, filter_field=\"last_edited_time\"\n )\n if len(pages) > 0:\n yield from batch_generator(self._read_pages(pages), self.batch_size)\n if db_res.has_more:\n query_dict[\"start_cursor\"] = db_res.next_cursor\n else:\n break\n else:\n break\n\n def validate_connector_settings(self) -> None:\n if not self.headers.get(\"Authorization\"):\n raise ConnectorMissingCredentialError(\"Notion credentials not loaded.\")\n\n try:\n # We'll do a minimal search call (page_size=1) to confirm accessibility\n if self.root_page_id:\n # If root_page_id is set, fetch the specific page\n res = rl_requests.get(\n f\"https://api.notion.com/v1/pages/{self.root_page_id}\",\n headers=self.headers,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n else:\n # If root_page_id is not set, perform a minimal search\n test_query = {\n \"filter\": {\"property\": \"object\", \"value\": \"page\"},\n \"page_size\": 1,\n }\n res = rl_requests.post(\n \"https://api.notion.com/v1/search\",\n headers=self.headers,\n json=test_query,\n timeout=_NOTION_CALL_TIMEOUT,\n )\n res.raise_for_status()\n\n except requests.exceptions.HTTPError as http_err:\n status_code = http_err.response.status_code if http_err.response else None\n\n if status_code == 401:\n raise CredentialExpiredError(\n \"Notion credential appears to be invalid or expired (HTTP 401).\"\n )\n elif status_code == 403:\n raise InsufficientPermissionsError(\n \"Your Notion token does not have sufficient permissions (HTTP 403).\"\n )\n elif status_code == 404:\n # Typically means resource not found or not shared. Could be root_page_id is invalid.\n raise ConnectorValidationError(\n \"Notion resource not found or not shared with the integration (HTTP 404).\"\n )\n elif status_code == 429:\n raise ConnectorValidationError(\n \"Validation failed due to Notion rate-limits being exceeded (HTTP 429). Please try again later.\"\n )\n else:\n raise UnexpectedValidationError(\n f\"Unexpected Notion HTTP error (status={status_code}): {http_err}\"\n ) from http_err\n\n except Exception as exc:\n raise UnexpectedValidationError(\n f\"Unexpected error during Notion settings validation: {exc}\"\n )\n\n\nif __name__ == \"__main__\":\n import os\n\n root_page_id = os.environ.get(\"NOTION_ROOT_PAGE_ID\")\n connector = NotionConnector(root_page_id=root_page_id)\n connector.load_credentials(\n {\"notion_integration_token\": os.environ.get(\"NOTION_INTEGRATION_TOKEN\")}\n )\n document_batches = connector.load_from_state()\n for doc_batch in document_batches:\n for doc in doc_batch:\n print(doc)\n" }, { "path": "backend/onyx/connectors/outline/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/outline/client.py", "content": "from typing import Any\n\nimport requests\nfrom requests.exceptions import ConnectionError as RequestsConnectionError\nfrom requests.exceptions import RequestException\nfrom requests.exceptions import Timeout\n\nfrom onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS\n\n\nclass OutlineClientRequestFailedError(ConnectionError):\n \"\"\"Custom error class for handling failed requests to the Outline API with status code and error message\"\"\"\n\n def __init__(self, status: int, error: str) -> None:\n self.status_code = status\n self.error = error\n super().__init__(f\"Outline Client request failed with status {status}: {error}\")\n\n\nclass OutlineApiClient:\n \"\"\"Client for interacting with the Outline API. Handles authentication and making HTTP requests.\"\"\"\n\n def __init__(\n self,\n api_token: str,\n base_url: str,\n ) -> None:\n self.base_url = base_url.rstrip(\"/\")\n self.api_token = api_token\n\n def post(self, endpoint: str, data: dict[str, Any] | None = None) -> dict[str, Any]:\n if data is None:\n data = {}\n url: str = self._build_url(endpoint)\n headers = self._build_headers()\n\n try:\n response = requests.post(\n url, headers=headers, json=data, timeout=REQUEST_TIMEOUT_SECONDS\n )\n except Timeout:\n raise OutlineClientRequestFailedError(\n 408,\n f\"Request timed out - server did not respond within {REQUEST_TIMEOUT_SECONDS} seconds\",\n )\n except RequestsConnectionError as e:\n raise OutlineClientRequestFailedError(\n -1, f\"Connection error - unable to reach Outline server: {e}\"\n )\n except RequestException as e:\n raise OutlineClientRequestFailedError(-1, f\"Network error occurred: {e}\")\n\n if response.status_code >= 300:\n error = response.reason\n try:\n response_json = response.json()\n if isinstance(response_json, dict):\n response_error = response_json.get(\"error\", {}).get(\"message\", \"\")\n if response_error:\n error = response_error\n except Exception:\n # If JSON parsing fails, fall back to response.text for better debugging\n if response.text.strip():\n error = f\"{response.reason}: {response.text.strip()}\"\n raise OutlineClientRequestFailedError(response.status_code, error)\n\n try:\n return response.json()\n except Exception:\n raise OutlineClientRequestFailedError(\n response.status_code,\n f\"Response was successful but contained invalid JSON: {response.text}\",\n )\n\n def _build_headers(self) -> dict[str, str]:\n return {\n \"Authorization\": f\"Bearer {self.api_token}\",\n \"Accept\": \"application/json\",\n \"Content-Type\": \"application/json\",\n }\n\n def _build_url(self, endpoint: str) -> str:\n return self.base_url.rstrip(\"/\") + \"/api/\" + endpoint.lstrip(\"/\")\n\n def build_app_url(self, endpoint: str) -> str:\n return self.base_url.rstrip(\"/\") + \"/\" + endpoint.lstrip(\"/\")\n" }, { "path": "backend/onyx/connectors/outline/connector.py", "content": "import html\nimport time\nfrom collections.abc import Callable\nfrom typing import Any\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.outline.client import OutlineApiClient\nfrom onyx.connectors.outline.client import OutlineClientRequestFailedError\n\n\nclass OutlineConnector(LoadConnector, PollConnector):\n \"\"\"Connector for Outline knowledge base. Handles authentication, document loading and polling.\n Implements both LoadConnector for initial state loading and PollConnector for incremental updates.\n \"\"\"\n\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.batch_size = batch_size\n self.outline_client: OutlineApiClient | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n required_keys = [\"outline_api_token\", \"outline_base_url\"]\n for key in required_keys:\n if key not in credentials:\n raise ConnectorMissingCredentialError(\"Outline\")\n\n self.outline_client = OutlineApiClient(\n api_token=credentials[\"outline_api_token\"],\n base_url=credentials[\"outline_base_url\"],\n )\n return None\n\n @staticmethod\n def _get_doc_batch(\n batch_size: int,\n outline_client: OutlineApiClient,\n endpoint: str,\n transformer: Callable[[OutlineApiClient, dict], Document],\n start_ind: int,\n ) -> tuple[list[Document], int]:\n data = {\n \"limit\": batch_size,\n \"offset\": start_ind,\n }\n\n batch = outline_client.post(endpoint, data=data).get(\"data\", [])\n doc_batch = [transformer(outline_client, item) for item in batch]\n\n return doc_batch, len(batch)\n\n @staticmethod\n def _collection_to_document(\n outline_client: OutlineApiClient, collection: dict[str, Any]\n ) -> Document:\n url = outline_client.build_app_url(f\"/collection/{collection.get('id')}\")\n title = str(collection.get(\"name\", \"\"))\n name = collection.get(\"name\") or \"\"\n description = collection.get(\"description\") or \"\"\n text = name + \"\\n\" + description\n updated_at_str = (\n str(collection.get(\"updatedAt\"))\n if collection.get(\"updatedAt\") is not None\n else None\n )\n return Document(\n id=\"outline_collection__\" + str(collection.get(\"id\")),\n sections=[TextSection(link=url, text=html.unescape(text))],\n source=DocumentSource.OUTLINE,\n semantic_identifier=\"Collection: \" + title,\n title=title,\n doc_updated_at=(\n time_str_to_utc(updated_at_str) if updated_at_str is not None else None\n ),\n metadata={\"type\": \"collection\"},\n )\n\n @staticmethod\n def _document_to_document(\n outline_client: OutlineApiClient, document: dict[str, Any]\n ) -> Document:\n url = outline_client.build_app_url(f\"/doc/{document.get('id')}\")\n title = str(document.get(\"title\", \"\"))\n doc_title = document.get(\"title\") or \"\"\n doc_text = document.get(\"text\") or \"\"\n text = doc_title + \"\\n\" + doc_text\n updated_at_str = (\n str(document.get(\"updatedAt\"))\n if document.get(\"updatedAt\") is not None\n else None\n )\n return Document(\n id=\"outline_document__\" + str(document.get(\"id\")),\n sections=[TextSection(link=url, text=html.unescape(text))],\n source=DocumentSource.OUTLINE,\n semantic_identifier=\"Document: \" + title,\n title=title,\n doc_updated_at=(\n time_str_to_utc(updated_at_str) if updated_at_str is not None else None\n ),\n metadata={\"type\": \"document\"},\n )\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n if self.outline_client is None:\n raise ConnectorMissingCredentialError(\"Outline\")\n\n return self._fetch_documents()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n if self.outline_client is None:\n raise ConnectorMissingCredentialError(\"Outline\")\n\n # Outline API does not support date-based filtering natively,\n # so we implement client-side filtering after fetching documents\n def time_filter(doc: Document) -> bool:\n if doc.doc_updated_at is None:\n return False\n doc_timestamp = doc.doc_updated_at.timestamp()\n if doc_timestamp < start:\n return False\n if doc_timestamp > end:\n return False\n return True\n\n return self._fetch_documents(time_filter)\n\n def _fetch_documents(\n self, time_filter: Callable[[Document], bool] | None = None\n ) -> GenerateDocumentsOutput:\n if self.outline_client is None:\n raise ConnectorMissingCredentialError(\"Outline\")\n\n transform_by_endpoint: dict[\n str, Callable[[OutlineApiClient, dict], Document]\n ] = {\n \"documents.list\": self._document_to_document,\n \"collections.list\": self._collection_to_document,\n }\n\n for endpoint, transform in transform_by_endpoint.items():\n start_ind = 0\n while True:\n doc_batch, num_results = self._get_doc_batch(\n batch_size=self.batch_size,\n outline_client=self.outline_client,\n endpoint=endpoint,\n transformer=transform,\n start_ind=start_ind,\n )\n\n # Apply time filtering if specified\n filtered_batch: list[Document | HierarchyNode] = []\n for doc in doc_batch:\n if time_filter is None or time_filter(doc):\n filtered_batch.append(doc)\n\n start_ind += num_results\n if filtered_batch:\n yield filtered_batch\n\n if num_results < self.batch_size:\n break\n else:\n time.sleep(0.2)\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n Validate that the Outline credentials and connector settings are correct.\n Specifically checks that we can make an authenticated request to Outline.\n \"\"\"\n if not self.outline_client:\n raise ConnectorMissingCredentialError(\"Outline\")\n\n try:\n # Use auth.info endpoint for validation\n _ = self.outline_client.post(\"auth.info\", data={})\n\n except OutlineClientRequestFailedError as e:\n # Check for HTTP status codes\n if e.status_code == 401:\n raise CredentialExpiredError(\n \"Your Outline credentials appear to be invalid or expired (HTTP 401).\"\n ) from e\n elif e.status_code == 403:\n raise InsufficientPermissionsError(\n \"The configured Outline token does not have sufficient permissions (HTTP 403).\"\n ) from e\n else:\n raise ConnectorValidationError(\n f\"Unexpected Outline error (status={e.status_code}): {e}\"\n ) from e\n\n except Exception as exc:\n raise ConnectorValidationError(\n f\"Unexpected error while validating Outline connector settings: {exc}\"\n ) from exc\n" }, { "path": "backend/onyx/connectors/productboard/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/productboard/connector.py", "content": "from collections.abc import Generator\nfrom itertools import chain\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom bs4 import BeautifulSoup\nfrom dateutil import parser\nfrom retry import retry\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n_PRODUCT_BOARD_BASE_URL = \"https://api.productboard.com\"\n\n\nclass ProductboardApiError(Exception):\n pass\n\n\nclass ProductboardConnector(PollConnector):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.batch_size = batch_size\n self.access_token: str | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self.access_token = credentials[\"productboard_access_token\"]\n return None\n\n def _build_headers(self) -> dict[str, str]:\n return {\n \"Authorization\": f\"Bearer {self.access_token}\",\n \"X-Version\": \"1\",\n }\n\n @staticmethod\n def _parse_description_html(description_html: str) -> str:\n soup = BeautifulSoup(description_html, \"html.parser\")\n return soup.get_text()\n\n @staticmethod\n def _get_owner_email(productboard_obj: dict[str, Any]) -> str | None:\n owner_dict = cast(dict[str, str] | None, productboard_obj.get(\"owner\"))\n if not owner_dict:\n return None\n return owner_dict.get(\"email\")\n\n def _fetch_documents(\n self,\n initial_link: str,\n ) -> Generator[dict[str, Any], None, None]:\n headers = self._build_headers()\n\n @retry(tries=3, delay=1, backoff=2)\n def fetch(link: str) -> dict[str, Any]:\n response = requests.get(link, headers=headers)\n if not response.ok:\n # rate-limiting is at 50 requests per second.\n # The delay in this retry should handle this while this is\n # not parallelized.\n raise ProductboardApiError(\n f\"Failed to fetch from productboard - status code: {response.status_code} - response: {response.text}\"\n )\n\n return response.json()\n\n curr_link = initial_link\n while True:\n response_json = fetch(curr_link)\n for entity in response_json[\"data\"]:\n yield entity\n\n curr_link = response_json.get(\"links\", {}).get(\"next\")\n if not curr_link:\n break\n\n def _get_features(self) -> Generator[Document, None, None]:\n \"\"\"A Feature is like a ticket in Jira\"\"\"\n for feature in self._fetch_documents(\n initial_link=f\"{_PRODUCT_BOARD_BASE_URL}/features\"\n ):\n owner = self._get_owner_email(feature)\n experts = [BasicExpertInfo(email=owner)] if owner else None\n\n metadata: dict[str, str | list[str]] = {}\n entity_type = feature.get(\"type\", \"feature\")\n if entity_type:\n metadata[\"entity_type\"] = str(entity_type)\n\n status = feature.get(\"status\", {}).get(\"name\")\n if status:\n metadata[\"status\"] = str(status)\n\n yield Document(\n id=feature[\"id\"],\n sections=[\n TextSection(\n link=feature[\"links\"][\"html\"],\n text=self._parse_description_html(feature[\"description\"]),\n )\n ],\n semantic_identifier=feature[\"name\"],\n source=DocumentSource.PRODUCTBOARD,\n doc_updated_at=time_str_to_utc(feature[\"updatedAt\"]),\n primary_owners=experts,\n metadata=metadata,\n )\n\n def _get_components(self) -> Generator[Document, None, None]:\n \"\"\"A Component is like an epic in Jira. It contains Features\"\"\"\n for component in self._fetch_documents(\n initial_link=f\"{_PRODUCT_BOARD_BASE_URL}/components\"\n ):\n owner = self._get_owner_email(component)\n experts = [BasicExpertInfo(email=owner)] if owner else None\n\n yield Document(\n id=component[\"id\"],\n sections=[\n TextSection(\n link=component[\"links\"][\"html\"],\n text=self._parse_description_html(component[\"description\"]),\n )\n ],\n semantic_identifier=component[\"name\"],\n source=DocumentSource.PRODUCTBOARD,\n doc_updated_at=time_str_to_utc(component[\"updatedAt\"]),\n primary_owners=experts,\n metadata={\n \"entity_type\": \"component\",\n },\n )\n\n def _get_products(self) -> Generator[Document, None, None]:\n \"\"\"A Product is the highest level of organization.\n A Product contains components, which contains features.\"\"\"\n for product in self._fetch_documents(\n initial_link=f\"{_PRODUCT_BOARD_BASE_URL}/products\"\n ):\n owner = self._get_owner_email(product)\n experts = [BasicExpertInfo(email=owner)] if owner else None\n\n yield Document(\n id=product[\"id\"],\n sections=[\n TextSection(\n link=product[\"links\"][\"html\"],\n text=self._parse_description_html(product[\"description\"]),\n )\n ],\n semantic_identifier=product[\"name\"],\n source=DocumentSource.PRODUCTBOARD,\n doc_updated_at=time_str_to_utc(product[\"updatedAt\"]),\n primary_owners=experts,\n metadata={\n \"entity_type\": \"product\",\n },\n )\n\n def _get_objectives(self) -> Generator[Document, None, None]:\n for objective in self._fetch_documents(\n initial_link=f\"{_PRODUCT_BOARD_BASE_URL}/objectives\"\n ):\n owner = self._get_owner_email(objective)\n experts = [BasicExpertInfo(email=owner)] if owner else None\n\n metadata: dict[str, str | list[str]] = {\n \"entity_type\": \"objective\",\n }\n if objective.get(\"state\"):\n metadata[\"state\"] = str(objective[\"state\"])\n\n yield Document(\n id=objective[\"id\"],\n sections=[\n TextSection(\n link=objective[\"links\"][\"html\"],\n text=self._parse_description_html(objective[\"description\"]),\n )\n ],\n semantic_identifier=objective[\"name\"],\n source=DocumentSource.PRODUCTBOARD,\n doc_updated_at=time_str_to_utc(objective[\"updatedAt\"]),\n primary_owners=experts,\n metadata=metadata,\n )\n\n def _is_updated_at_out_of_time_range(\n self,\n document: Document,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n ) -> bool:\n updated_at = cast(str, document.metadata.get(\"updated_at\", \"\"))\n if updated_at:\n updated_at_datetime = parser.parse(updated_at)\n if (\n updated_at_datetime.timestamp() < start\n or updated_at_datetime.timestamp() > end\n ):\n return True\n else:\n logger.debug(f\"Unable to find updated_at for document '{document.id}'\")\n\n return False\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n if self.access_token is None:\n raise PermissionError(\n \"Access token is not set up, was load_credentials called?\"\n )\n\n document_batch: list[Document | HierarchyNode] = []\n\n # NOTE: there is a concept of a \"Note\" in productboard, however\n # there is no read API for it atm. Additionally, comments are not\n # included with features. Finally, \"Releases\" are not fetched atm,\n # since they do not provide an updatedAt.\n feature_documents = self._get_features()\n component_documents = self._get_components()\n product_documents = self._get_products()\n objective_documents = self._get_objectives()\n for document in chain(\n feature_documents,\n component_documents,\n product_documents,\n objective_documents,\n ):\n # skip documents that are not in the time range\n if self._is_updated_at_out_of_time_range(document, start, end):\n continue\n\n document_batch.append(document)\n if len(document_batch) >= self.batch_size:\n yield document_batch\n document_batch = []\n\n if document_batch:\n yield document_batch\n\n\nif __name__ == \"__main__\":\n import os\n import time\n\n connector = ProductboardConnector()\n connector.load_credentials(\n {\n \"productboard_access_token\": os.environ[\"PRODUCTBOARD_ACCESS_TOKEN\"],\n }\n )\n\n current = time.time()\n one_year_ago = current - 24 * 60 * 60 * 360\n latest_docs = connector.poll_source(one_year_ago, current)\n print(next(latest_docs))\n" }, { "path": "backend/onyx/connectors/registry.py", "content": "\"\"\"Registry mapping for connector classes.\"\"\"\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import DocumentSource\n\n\nclass ConnectorMapping(BaseModel):\n module_path: str\n class_name: str\n\n\n# Mapping of DocumentSource to connector details for lazy loading\nCONNECTOR_CLASS_MAP = {\n DocumentSource.WEB: ConnectorMapping(\n module_path=\"onyx.connectors.web.connector\",\n class_name=\"WebConnector\",\n ),\n DocumentSource.FILE: ConnectorMapping(\n module_path=\"onyx.connectors.file.connector\",\n class_name=\"LocalFileConnector\",\n ),\n DocumentSource.SLACK: ConnectorMapping(\n module_path=\"onyx.connectors.slack.connector\",\n class_name=\"SlackConnector\",\n ),\n DocumentSource.GITHUB: ConnectorMapping(\n module_path=\"onyx.connectors.github.connector\",\n class_name=\"GithubConnector\",\n ),\n DocumentSource.GMAIL: ConnectorMapping(\n module_path=\"onyx.connectors.gmail.connector\",\n class_name=\"GmailConnector\",\n ),\n DocumentSource.GITLAB: ConnectorMapping(\n module_path=\"onyx.connectors.gitlab.connector\",\n class_name=\"GitlabConnector\",\n ),\n DocumentSource.GITBOOK: ConnectorMapping(\n module_path=\"onyx.connectors.gitbook.connector\",\n class_name=\"GitbookConnector\",\n ),\n DocumentSource.GOOGLE_DRIVE: ConnectorMapping(\n module_path=\"onyx.connectors.google_drive.connector\",\n class_name=\"GoogleDriveConnector\",\n ),\n DocumentSource.BOOKSTACK: ConnectorMapping(\n module_path=\"onyx.connectors.bookstack.connector\",\n class_name=\"BookstackConnector\",\n ),\n DocumentSource.OUTLINE: ConnectorMapping(\n module_path=\"onyx.connectors.outline.connector\",\n class_name=\"OutlineConnector\",\n ),\n DocumentSource.CONFLUENCE: ConnectorMapping(\n module_path=\"onyx.connectors.confluence.connector\",\n class_name=\"ConfluenceConnector\",\n ),\n DocumentSource.JIRA: ConnectorMapping(\n module_path=\"onyx.connectors.jira.connector\",\n class_name=\"JiraConnector\",\n ),\n DocumentSource.PRODUCTBOARD: ConnectorMapping(\n module_path=\"onyx.connectors.productboard.connector\",\n class_name=\"ProductboardConnector\",\n ),\n DocumentSource.SLAB: ConnectorMapping(\n module_path=\"onyx.connectors.slab.connector\",\n class_name=\"SlabConnector\",\n ),\n DocumentSource.CODA: ConnectorMapping(\n module_path=\"onyx.connectors.coda.connector\",\n class_name=\"CodaConnector\",\n ),\n DocumentSource.CANVAS: ConnectorMapping(\n module_path=\"onyx.connectors.canvas.connector\",\n class_name=\"CanvasConnector\",\n ),\n DocumentSource.NOTION: ConnectorMapping(\n module_path=\"onyx.connectors.notion.connector\",\n class_name=\"NotionConnector\",\n ),\n DocumentSource.ZULIP: ConnectorMapping(\n module_path=\"onyx.connectors.zulip.connector\",\n class_name=\"ZulipConnector\",\n ),\n DocumentSource.GURU: ConnectorMapping(\n module_path=\"onyx.connectors.guru.connector\",\n class_name=\"GuruConnector\",\n ),\n DocumentSource.LINEAR: ConnectorMapping(\n module_path=\"onyx.connectors.linear.connector\",\n class_name=\"LinearConnector\",\n ),\n DocumentSource.HUBSPOT: ConnectorMapping(\n module_path=\"onyx.connectors.hubspot.connector\",\n class_name=\"HubSpotConnector\",\n ),\n DocumentSource.DOCUMENT360: ConnectorMapping(\n module_path=\"onyx.connectors.document360.connector\",\n class_name=\"Document360Connector\",\n ),\n DocumentSource.GONG: ConnectorMapping(\n module_path=\"onyx.connectors.gong.connector\",\n class_name=\"GongConnector\",\n ),\n DocumentSource.GOOGLE_SITES: ConnectorMapping(\n module_path=\"onyx.connectors.google_site.connector\",\n class_name=\"GoogleSitesConnector\",\n ),\n DocumentSource.ZENDESK: ConnectorMapping(\n module_path=\"onyx.connectors.zendesk.connector\",\n class_name=\"ZendeskConnector\",\n ),\n DocumentSource.LOOPIO: ConnectorMapping(\n module_path=\"onyx.connectors.loopio.connector\",\n class_name=\"LoopioConnector\",\n ),\n DocumentSource.DROPBOX: ConnectorMapping(\n module_path=\"onyx.connectors.dropbox.connector\",\n class_name=\"DropboxConnector\",\n ),\n DocumentSource.SHAREPOINT: ConnectorMapping(\n module_path=\"onyx.connectors.sharepoint.connector\",\n class_name=\"SharepointConnector\",\n ),\n DocumentSource.TEAMS: ConnectorMapping(\n module_path=\"onyx.connectors.teams.connector\",\n class_name=\"TeamsConnector\",\n ),\n DocumentSource.SALESFORCE: ConnectorMapping(\n module_path=\"onyx.connectors.salesforce.connector\",\n class_name=\"SalesforceConnector\",\n ),\n DocumentSource.DISCOURSE: ConnectorMapping(\n module_path=\"onyx.connectors.discourse.connector\",\n class_name=\"DiscourseConnector\",\n ),\n DocumentSource.AXERO: ConnectorMapping(\n module_path=\"onyx.connectors.axero.connector\",\n class_name=\"AxeroConnector\",\n ),\n DocumentSource.CLICKUP: ConnectorMapping(\n module_path=\"onyx.connectors.clickup.connector\",\n class_name=\"ClickupConnector\",\n ),\n DocumentSource.MEDIAWIKI: ConnectorMapping(\n module_path=\"onyx.connectors.mediawiki.wiki\",\n class_name=\"MediaWikiConnector\",\n ),\n DocumentSource.WIKIPEDIA: ConnectorMapping(\n module_path=\"onyx.connectors.wikipedia.connector\",\n class_name=\"WikipediaConnector\",\n ),\n DocumentSource.ASANA: ConnectorMapping(\n module_path=\"onyx.connectors.asana.connector\",\n class_name=\"AsanaConnector\",\n ),\n DocumentSource.S3: ConnectorMapping(\n module_path=\"onyx.connectors.blob.connector\",\n class_name=\"BlobStorageConnector\",\n ),\n DocumentSource.R2: ConnectorMapping(\n module_path=\"onyx.connectors.blob.connector\",\n class_name=\"BlobStorageConnector\",\n ),\n DocumentSource.GOOGLE_CLOUD_STORAGE: ConnectorMapping(\n module_path=\"onyx.connectors.blob.connector\",\n class_name=\"BlobStorageConnector\",\n ),\n DocumentSource.OCI_STORAGE: ConnectorMapping(\n module_path=\"onyx.connectors.blob.connector\",\n class_name=\"BlobStorageConnector\",\n ),\n DocumentSource.XENFORO: ConnectorMapping(\n module_path=\"onyx.connectors.xenforo.connector\",\n class_name=\"XenforoConnector\",\n ),\n DocumentSource.DISCORD: ConnectorMapping(\n module_path=\"onyx.connectors.discord.connector\",\n class_name=\"DiscordConnector\",\n ),\n DocumentSource.FRESHDESK: ConnectorMapping(\n module_path=\"onyx.connectors.freshdesk.connector\",\n class_name=\"FreshdeskConnector\",\n ),\n DocumentSource.FIREFLIES: ConnectorMapping(\n module_path=\"onyx.connectors.fireflies.connector\",\n class_name=\"FirefliesConnector\",\n ),\n DocumentSource.EGNYTE: ConnectorMapping(\n module_path=\"onyx.connectors.egnyte.connector\",\n class_name=\"EgnyteConnector\",\n ),\n DocumentSource.AIRTABLE: ConnectorMapping(\n module_path=\"onyx.connectors.airtable.airtable_connector\",\n class_name=\"AirtableConnector\",\n ),\n DocumentSource.HIGHSPOT: ConnectorMapping(\n module_path=\"onyx.connectors.highspot.connector\",\n class_name=\"HighspotConnector\",\n ),\n DocumentSource.DRUPAL_WIKI: ConnectorMapping(\n module_path=\"onyx.connectors.drupal_wiki.connector\",\n class_name=\"DrupalWikiConnector\",\n ),\n DocumentSource.IMAP: ConnectorMapping(\n module_path=\"onyx.connectors.imap.connector\",\n class_name=\"ImapConnector\",\n ),\n DocumentSource.BITBUCKET: ConnectorMapping(\n module_path=\"onyx.connectors.bitbucket.connector\",\n class_name=\"BitbucketConnector\",\n ),\n DocumentSource.TESTRAIL: ConnectorMapping(\n module_path=\"onyx.connectors.testrail.connector\",\n class_name=\"TestRailConnector\",\n ),\n # just for integration tests\n DocumentSource.MOCK_CONNECTOR: ConnectorMapping(\n module_path=\"onyx.connectors.mock_connector.connector\",\n class_name=\"MockConnector\",\n ),\n}\n" }, { "path": "backend/onyx/connectors/requesttracker/.gitignore", "content": ".env\n" }, { "path": "backend/onyx/connectors/requesttracker/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/requesttracker/connector.py", "content": "# from datetime import datetime\n# from datetime import timezone\n# from logging import DEBUG as LOG_LVL_DEBUG\n# from typing import Any\n# from typing import List\n# from typing import Optional\n# from rt.rest1 import ALL_QUEUES\n# from rt.rest1 import Rt\n# from onyx.configs.app_configs import INDEX_BATCH_SIZE\n# from onyx.configs.constants import DocumentSource\n# from onyx.connectors.interfaces import GenerateDocumentsOutput\n# from onyx.connectors.interfaces import PollConnector\n# from onyx.connectors.interfaces import SecondsSinceUnixEpoch\n# from onyx.connectors.models import ConnectorMissingCredentialError\n# from onyx.connectors.models import Document\n# from onyx.connectors.models import Section\n# from onyx.utils.logger import setup_logger\n# logger = setup_logger()\n# class RequestTrackerError(Exception):\n# pass\n# class RequestTrackerConnector(PollConnector):\n# def __init__(\n# self,\n# batch_size: int = INDEX_BATCH_SIZE,\n# ) -> None:\n# self.batch_size = batch_size\n# def txn_link(self, tid: int, txn: int) -> str:\n# return f\"{self.rt_base_url}/Ticket/Display.html?id={tid}&txn={txn}\"\n# def build_doc_sections_from_txn(\n# self, connection: Rt, ticket_id: int\n# ) -> List[Section]:\n# Sections: List[Section] = []\n# get_history_resp = connection.get_history(ticket_id)\n# if get_history_resp is None:\n# raise RequestTrackerError(f\"Ticket {ticket_id} cannot be found\")\n# for tx in get_history_resp:\n# Sections.append(\n# Section(\n# link=self.txn_link(ticket_id, int(tx[\"id\"])),\n# text=\"\\n\".join(\n# [\n# f\"{k}:\\n{v}\\n\" if k != \"Attachments\" else \"\"\n# for (k, v) in tx.items()\n# ]\n# ),\n# )\n# )\n# return Sections\n# def load_credentials(self, credentials: dict[str, Any]) -> Optional[dict[str, Any]]:\n# self.rt_username = credentials.get(\"requesttracker_username\")\n# self.rt_password = credentials.get(\"requesttracker_password\")\n# self.rt_base_url = credentials.get(\"requesttracker_base_url\")\n# return None\n# # This does not include RT file attachments yet.\n# def _process_tickets(\n# self, start: datetime, end: datetime\n# ) -> GenerateDocumentsOutput:\n# if any([self.rt_username, self.rt_password, self.rt_base_url]) is None:\n# raise ConnectorMissingCredentialError(\"requesttracker\")\n# Rt0 = Rt(\n# f\"{self.rt_base_url}/REST/1.0/\",\n# self.rt_username,\n# self.rt_password,\n# )\n# Rt0.login()\n# d0 = start.strftime(\"%Y-%m-%d %H:%M:%S\")\n# d1 = end.strftime(\"%Y-%m-%d %H:%M:%S\")\n# tickets = Rt0.search(\n# Queue=ALL_QUEUES,\n# raw_query=f\"Updated > '{d0}' AND Updated < '{d1}'\",\n# )\n# doc_batch: List[Document] = []\n# for ticket in tickets:\n# ticket_keys_to_omit = [\"id\", \"Subject\"]\n# tid: int = int(ticket[\"numerical_id\"])\n# ticketLink: str = f\"{self.rt_base_url}/Ticket/Display.html?id={tid}\"\n# logger.info(f\"Processing ticket {tid}\")\n# doc = Document(\n# id=ticket[\"id\"],\n# # Will add title to the first section later in processing\n# sections=[Section(link=ticketLink, text=\"\")]\n# + self.build_doc_sections_from_txn(Rt0, tid),\n# source=DocumentSource.REQUESTTRACKER,\n# semantic_identifier=ticket[\"Subject\"],\n# metadata={\n# key: value\n# for key, value in ticket.items()\n# if key not in ticket_keys_to_omit\n# },\n# )\n# doc_batch.append(doc)\n# if len(doc_batch) >= self.batch_size:\n# yield doc_batch\n# doc_batch = []\n# if doc_batch:\n# yield doc_batch\n# def poll_source(\n# self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n# ) -> GenerateDocumentsOutput:\n# # Keep query short, only look behind 1 day at maximum\n# one_day_ago: float = end - (24 * 60 * 60)\n# _start: float = start if start > one_day_ago else one_day_ago\n# start_datetime = datetime.fromtimestamp(_start, tz=timezone.utc)\n# end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)\n# yield from self._process_tickets(start_datetime, end_datetime)\n# if __name__ == \"__main__\":\n# import time\n# import os\n# from dotenv import load_dotenv\n# load_dotenv()\n# logger.setLevel(LOG_LVL_DEBUG)\n# rt_connector = RequestTrackerConnector()\n# rt_connector.load_credentials(\n# {\n# \"requesttracker_username\": os.getenv(\"RT_USERNAME\"),\n# \"requesttracker_password\": os.getenv(\"RT_PASSWORD\"),\n# \"requesttracker_base_url\": os.getenv(\"RT_BASE_URL\"),\n# }\n# )\n# current = time.time()\n# one_day_ago = current - (24 * 60 * 60) # 1 days\n# latest_docs = rt_connector.poll_source(one_day_ago, current)\n# for doc in latest_docs:\n# print(doc)\n" }, { "path": "backend/onyx/connectors/salesforce/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/salesforce/blacklist.py", "content": "# NOTE(rkuo): I can't find an actual API that allows us to distinguish\n# broken/incompatible objects from regular ones.\n# taking hints from\n# https://docs.resco.net/wiki/Salesforce_object_blacklist\n\nSALESFORCE_BLACKLISTED_PREFIXES: set[str] = set(\n [\n \"process\",\n \"aura\",\n \"app\",\n \"auth\",\n \"duplicate\",\n \"secure\",\n \"data\",\n \"listemail\",\n \"fsl__optimization\",\n \"fsl_scheduling\",\n \"feed\",\n \"chatter\",\n ]\n)\n\nSALESFORCE_BLACKLISTED_SUFFIXES: set[str] = set(\n [\n \"history\",\n \"share\",\n \"__tag\",\n \"__hd\",\n \"feed\",\n \"changeevent\",\n \"__ka\",\n \"__votestat\",\n \"__viewstat\",\n \"__kav\",\n \"__datacategoryselection\",\n \"subscription\",\n \"definition\",\n \"eventstream\",\n \"__mdt\",\n ]\n)\n\nSALESFORCE_BLACKLISTED_OBJECTS: set[str] = set(\n [\n \"acceptedeventrelation\",\n \"accountchangeevent\",\n \"accountcontactrole\",\n \"accountcontactrolechangeevent\",\n \"accounthistory\",\n \"accountshare\",\n \"actionlinkgrouptemplate\",\n \"actionlinktemplate\",\n \"activityhistory\",\n \"adminsetupevent\",\n \"aggregateresult\",\n \"announcement\",\n \"apexclass\",\n \"apexcomponent\",\n \"apexemailnotification\",\n \"apexlog\",\n \"apexpage\",\n \"apexpageinfo\",\n \"apextestqueueitem\",\n \"apextestresult\",\n \"apextestresultlimits\",\n \"apextestrunresult\",\n \"apextestsuite\",\n \"apextrigger\",\n \"apievent\",\n \"apptabmember\",\n \"assetchangeevent\",\n \"assethistory\",\n \"assetrelationshiphistory\",\n \"assettokenevent\",\n \"assignmentrule\",\n \"asyncapexjob\",\n \"backgroundoperation\",\n \"backgroundoperationresult\",\n \"batchapexerrorevent\",\n \"brandingset\",\n \"brandingsetproperty\",\n \"brandtemplate\",\n \"businessprocess\",\n \"campaignchangeevent\",\n \"campaignhistory\",\n \"campaignshare\",\n \"casechangeevent\",\n \"caseexternaldocument\",\n \"casehistory\",\n \"caseshare\",\n \"clientbrowser\",\n \"collaborationgroup\",\n \"collaborationgroupmember\",\n \"collaborationgroupmemberrequest\",\n \"collaborationinvitation\",\n \"connectedapplication\",\n \"contactchangeevent\",\n \"contacthistory\",\n \"contactrequest\",\n \"contactrequestshare\",\n \"contactshare\",\n \"contentasset\",\n \"contentbody\",\n \"contentdocumenthistory\",\n \"contenthubrepository\",\n \"contenttagsubscription\",\n \"contentusersubscription\",\n \"contentversionhistory\",\n \"contracthistory\",\n \"corswhitelistentry\",\n \"cronjobdetail\",\n \"crontrigger\",\n \"csptrustedsite\",\n \"custombrand\",\n \"custombrandasset\",\n \"customhelpmenuitem\",\n \"customhelpmenusection\",\n \"customhttpheader\",\n \"customobjectuserlicensemetrics\",\n \"custompermission\",\n \"custompermissiondependency\",\n \"dandbcompany\",\n \"dashboard\",\n \"dashboardcomponent\",\n \"digitalsignature\",\n \"documentattachmentmap\",\n \"domain\",\n \"domainsite\",\n \"emailcapture\",\n \"emaildomainfilter\",\n \"emaildomainkey\",\n \"emailrelay\",\n \"emailservicesaddress\",\n \"emailservicesfunction\",\n \"emailstatus\",\n \"emailtemplate\",\n \"embeddedservicedetail\",\n \"embeddedservicelabel\",\n \"entityparticle\",\n \"eventbussubscriber\",\n \"eventchangeevent\",\n \"eventlogfile\",\n \"eventrelationchangeevent\",\n \"expressionfilter\",\n \"expressionfiltercriteria\",\n \"externaldatasource\",\n \"externaldatauserauth\",\n \"fieldhistoryarchive\",\n \"fieldpermissions\",\n \"fieldservicemobilesettings\",\n \"filesearchactivity\",\n \"fiscalyearsettings\",\n \"flexqueueitem\",\n \"flowinterview\",\n \"flowinterviewshare\",\n \"flowrecordrelation\",\n \"flowstagerelation\",\n \"forecastingshare\",\n \"forecastshare\",\n \"fsl__criteria__c\",\n \"fsl__gantt_filter__c\",\n \"fsl__ganttpalette__c\",\n \"fsl__service_goal__c\",\n \"fsl__slr_cache__c\",\n \"fsl__territory_optimization_request__c\",\n \"goalhistory\",\n \"goalshare\",\n \"grantedbylicense\",\n \"idpeventlog\",\n \"iframewhitelisturl\",\n \"image\",\n \"imageshare\",\n \"installedmobileapp\",\n \"leadchangeevent\",\n \"leadhistory\",\n \"leadshare\",\n \"lightningexitbypagemetrics\",\n \"lightningexperiencetheme\",\n \"lightningtogglemetrics\",\n \"lightningusagebyapptypemetrics\",\n \"lightningusagebybrowsermetrics\",\n \"lightningusagebyflexipagemetrics\",\n \"lightningusagebypagemetrics\",\n \"linkedarticle\",\n \"listemailchangeevent\",\n \"listemailshare\",\n \"listview\",\n \"listviewchart\",\n \"listviewchartinstance\",\n \"listviewevent\",\n \"loginasevent\",\n \"loginevent\",\n \"logingeo\",\n \"loginhistory\",\n \"loginip\",\n \"logoutevent\",\n \"lookedupfromactivity\",\n \"macro\",\n \"macrohistory\",\n \"macroinstruction\",\n \"macroshare\",\n \"mailmergetemplate\",\n \"matchingrule\",\n \"matchingruleitem\",\n \"metricdatalinkhistory\",\n \"metrichistory\",\n \"metricshare\",\n \"mobilesettingsassignment\",\n \"mydomaindiscoverablelogin\",\n \"name\",\n \"namedcredential\",\n \"noteandattachment\",\n \"notificationmember\",\n \"oauthtoken\",\n \"objectpermissions\",\n \"onboardingmetrics\",\n \"openactivity\",\n \"opportunitychangeevent\",\n \"opportunitycontactrolechangeevent\",\n \"opportunityfieldhistory\",\n \"opportunityhistory\",\n \"opportunityshare\",\n \"orderchangeevent\",\n \"orderhistory\",\n \"orderitemchangeevent\",\n \"orderitemhistory\",\n \"ordershare\",\n \"orgdeleterequest\",\n \"orgdeleterequestshare\",\n \"orglifecyclenotification\",\n \"orgwideemailaddress\",\n \"outgoingemail\",\n \"outgoingemailrelation\",\n \"ownerchangeoptioninfo\",\n \"packagelicense\",\n \"period\",\n \"permissionsetlicense\",\n \"permissionsetlicenseassign\",\n \"permissionsettabsetting\",\n \"person\",\n \"picklistvalueinfo\",\n \"platformaction\",\n \"platformcachepartition\",\n \"platformcachepartitiontype\",\n \"platformstatusalertevent\",\n \"pricebook2history\",\n \"processinstancehistory\",\n \"product2changeevent\",\n \"product2history\",\n \"publisher\",\n \"pushtopic\",\n \"pushupgradeexcludedorg\",\n \"quicktexthistory\",\n \"quicktextshare\",\n \"quotetemplaterichtextdata\",\n \"recordaction\",\n \"recordactionhistory\",\n \"recordvisibility\",\n \"relationshipdomain\",\n \"relationshipinfo\",\n \"reportevent\",\n \"samlssoconfig\",\n \"scontrol\",\n \"searchactivity\",\n \"searchlayout\",\n \"searchpromotionrule\",\n \"securitycustombaseline\",\n \"servicereportlayout\",\n \"sessionpermsetactivation\",\n \"setupaudittrail\",\n \"setupentityaccess\",\n \"site\",\n \"sitedetail\",\n \"sitehistory\",\n \"siteiframewhitelisturl\",\n \"solutionhistory\",\n \"sosdeployment\",\n \"sossession\",\n \"sossessionactivity\",\n \"sossessionhistory\",\n \"sossessionshare\",\n \"staticresource\",\n \"streamingchannel\",\n \"streamingchannelshare\",\n \"subscriberpackage\",\n \"subscriberpackageversion\",\n \"taskchangeevent\",\n \"tenantusageentitlement\",\n \"testsuitemembership\",\n \"thirdpartyaccountlink\",\n \"todaygoal\",\n \"todaygoalshare\",\n \"transactionsecuritypolicy\",\n \"twofactorinfo\",\n \"twofactormethodsinfo\",\n \"twofactortempcode\",\n \"urievent\",\n \"userappinfo\",\n \"userappmenucustomization\",\n \"userappmenucustomizationshare\",\n \"userappmenuitem\",\n \"userchangeevent\",\n \"useremailpreferredperson\",\n \"useremailpreferredpersonshare\",\n \"userentityaccess\",\n \"userfieldaccess\",\n \"userlicense\",\n \"userlistview\",\n \"userlistviewcriterion\",\n \"userlogin\",\n \"userpackagelicense\",\n \"userpermissionaccess\",\n \"userpreference\",\n \"userprovaccount\",\n \"userprovaccountstaging\",\n \"userprovisioningconfig\",\n \"userprovisioninglog\",\n \"userprovisioningrequest\",\n \"userprovisioningrequestshare\",\n \"userprovmocktarget\",\n \"userrecordaccess\",\n \"usershare\",\n \"verificationhistory\",\n \"visibilitychangenotification\",\n \"visualforceaccessmetrics\",\n \"waveautoinstallrequest\",\n \"wavecompatibilitycheckitem\",\n \"weblink\",\n \"workcoachinghistory\",\n \"workcoachingshare\",\n \"workfeedbackhistory\",\n \"workfeedbackquestion\",\n \"workfeedbackquestionhistory\",\n \"workfeedbackquestionsethistory\",\n \"workfeedbackquestionsetshare\",\n \"workfeedbackquestionshare\",\n \"workfeedbackrequesthistory\",\n \"workfeedbackrequestshare\",\n \"workfeedbackshare\",\n \"workfeedbacktemplateshare\",\n \"workperformancecyclehistory\",\n \"workperformancecycleshare\",\n ]\n)\n" }, { "path": "backend/onyx/connectors/salesforce/connector.py", "content": "import csv\nimport gc\nimport json\nimport os\nimport sys\nimport tempfile\nimport time\nfrom collections import defaultdict\nfrom collections.abc import Callable\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import cast\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.salesforce.doc_conversion import convert_sf_object_to_doc\nfrom onyx.connectors.salesforce.doc_conversion import convert_sf_query_result_to_doc\nfrom onyx.connectors.salesforce.doc_conversion import ID_PREFIX\nfrom onyx.connectors.salesforce.onyx_salesforce import OnyxSalesforce\nfrom onyx.connectors.salesforce.salesforce_calls import fetch_all_csvs_in_parallel\nfrom onyx.connectors.salesforce.sqlite_functions import OnyxSalesforceSQLite\nfrom onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE\nfrom onyx.connectors.salesforce.utils import ID_FIELD\nfrom onyx.connectors.salesforce.utils import MODIFIED_FIELD\nfrom onyx.connectors.salesforce.utils import NAME_FIELD\nfrom onyx.connectors.salesforce.utils import USER_OBJECT_TYPE\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _convert_to_metadata_value(value: Any) -> str | list[str]:\n \"\"\"Convert a Salesforce field value to a valid metadata value.\n\n Document metadata expects str | list[str], but Salesforce returns\n various types (bool, float, int, etc.). This function ensures all\n values are properly converted to strings.\n \"\"\"\n if isinstance(value, list):\n return [str(item) for item in value]\n return str(value)\n\n\n_DEFAULT_PARENT_OBJECT_TYPES = [ACCOUNT_OBJECT_TYPE]\n\n_DEFAULT_ATTRIBUTES_TO_KEEP: dict[str, dict[str, str]] = {\n \"Opportunity\": {\n ACCOUNT_OBJECT_TYPE: \"account\",\n \"FiscalQuarter\": \"fiscal_quarter\",\n \"FiscalYear\": \"fiscal_year\",\n \"IsClosed\": \"is_closed\",\n NAME_FIELD: \"name\",\n \"StageName\": \"stage_name\",\n \"Type\": \"type\",\n \"Amount\": \"amount\",\n \"CloseDate\": \"close_date\",\n \"Probability\": \"probability\",\n \"CreatedDate\": \"created_date\",\n MODIFIED_FIELD: \"last_modified_date\",\n },\n \"Contact\": {\n ACCOUNT_OBJECT_TYPE: \"account\",\n \"CreatedDate\": \"created_date\",\n MODIFIED_FIELD: \"last_modified_date\",\n },\n}\n\n\nclass SalesforceCheckpoint(ConnectorCheckpoint):\n initial_sync_complete: bool\n current_timestamp: SecondsSinceUnixEpoch\n\n\nclass SalesforceConnectorContext:\n parent_types: set[str] = set()\n child_types: set[str] = set()\n parent_to_child_types: dict[str, set[str]] = {} # map from parent to child types\n child_to_parent_types: dict[str, set[str]] = {} # map from child to parent types\n parent_reference_fields_by_type: dict[str, dict[str, list[str]]] = {}\n type_to_queryable_fields: dict[str, set[str]] = {}\n prefix_to_type: dict[str, str] = {} # infer the object type of an id immediately\n\n parent_to_child_relationships: dict[str, set[str]] = (\n {}\n ) # map from parent to child relationships\n parent_to_relationship_queryable_fields: dict[str, dict[str, set[str]]] = (\n {}\n ) # map from relationship to queryable fields\n\n parent_child_names_to_relationships: dict[str, str] = {}\n\n\ndef _extract_fields_and_associations_from_config(\n config: dict[str, Any], object_type: str\n) -> tuple[list[str] | None, dict[str, list[str]]]:\n \"\"\"\n Extract fields and associations for a specific object type from custom config.\n\n Returns:\n tuple of (fields_list, associations_dict)\n - fields_list: List of fields to query, or None if not specified (use all)\n - associations_dict: Dict mapping association names to their config\n \"\"\"\n if object_type not in config:\n return None, {}\n\n obj_config = config[object_type]\n fields = obj_config.get(\"fields\")\n associations = obj_config.get(\"associations\", {})\n\n return fields, associations\n\n\ndef _validate_custom_query_config(config: dict[str, Any]) -> None:\n \"\"\"\n Validate the structure of the custom query configuration.\n \"\"\"\n\n for object_type, obj_config in config.items():\n if not isinstance(obj_config, dict):\n raise ValueError(\n f\"top level object {object_type} must be mapped to a dictionary\"\n )\n\n # Check if fields is a list when present\n if \"fields\" in obj_config:\n if not isinstance(obj_config[\"fields\"], list):\n raise ValueError(\"if fields key exists, value must be a list\")\n for v in obj_config[\"fields\"]:\n if not isinstance(v, str):\n raise ValueError(f\"if fields list value {v} is not a string\")\n\n # Check if associations is a dict when present\n if \"associations\" in obj_config:\n if not isinstance(obj_config[\"associations\"], dict):\n raise ValueError(\n \"if associations key exists, value must be a dictionary\"\n )\n for assoc_name, assoc_fields in obj_config[\"associations\"].items():\n if not isinstance(assoc_fields, list):\n raise ValueError(\n f\"associations list value {assoc_fields} for key {assoc_name} is not a list\"\n )\n for v in assoc_fields:\n if not isinstance(v, str):\n raise ValueError(\n f\"if associations list value {v} is not a string\"\n )\n\n\nclass SalesforceConnector(LoadConnector, PollConnector, SlimConnectorWithPermSync):\n \"\"\"Approach outline\n\n Goal\n - get data for every record of every parent object type\n - The data should consist of the parent object record and all direct child relationship objects\n\n\n Initial sync\n - Does a full sync, then indexes each parent object + children as a document via\n the local sqlite db\n\n - get the first level children object types of parent object types\n - bulk export all object types to CSV\n -- NOTE: bulk exports of an object type contain parent id's, but not child id's\n - Load all CSV's to the DB\n - generate all parent object types as documents and yield them\n\n - Initial sync's must always be for the entire dataset.\n Otherwise, you can have cases where some records relate to other records that were\n updated recently. The more recently updated records will not be pulled down in the query.\n\n Delta sync's\n - delta sync's detect changes in parent objects, then perform a full sync of\n each parent object and its children\n\n If loading the entire db, this approach is much slower. For deltas, it works well.\n\n - query all changed records (includes children and parents)\n - extrapolate all changed parent objects\n - for each parent object, construct a query and yield the result back\n\n - Delta sync's can be done object by object by identifying the parent id of any changed\n record, and querying a single record at a time to get all the updated data. In this way,\n we avoid having to keep a locally synchronized copy of the entire salesforce db.\n\n TODO: verify record to doc conversion\n figure out why sometimes the field names are missing.\n \"\"\"\n\n MAX_BATCH_BYTES = 1024 * 1024\n LOG_INTERVAL = 10.0 # how often to log stats in loop heavy parts of the connector\n\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n requested_objects: list[str] = [],\n custom_query_config: str | None = None,\n ) -> None:\n self.batch_size = batch_size\n self._sf_client: OnyxSalesforce | None = None\n\n # Validate and store custom query config\n if custom_query_config:\n config_json = json.loads(custom_query_config)\n self.custom_query_config: dict[str, Any] | None = config_json\n # If custom query config is provided, use the object types from it\n self.parent_object_list = list(config_json.keys())\n else:\n self.custom_query_config = None\n # Use the traditional requested_objects approach\n self.parent_object_list = (\n [obj.strip().capitalize() for obj in requested_objects]\n if requested_objects\n else _DEFAULT_PARENT_OBJECT_TYPES\n )\n\n def load_credentials(\n self,\n credentials: dict[str, Any],\n ) -> dict[str, Any] | None:\n domain = \"test\" if credentials.get(\"is_sandbox\") else None\n self._sf_client = OnyxSalesforce(\n username=credentials[\"sf_username\"],\n password=credentials[\"sf_password\"],\n security_token=credentials[\"sf_security_token\"],\n domain=domain,\n )\n return None\n\n @property\n def sf_client(self) -> OnyxSalesforce:\n if self._sf_client is None:\n raise ConnectorMissingCredentialError(\"Salesforce\")\n return self._sf_client\n\n @staticmethod\n def reconstruct_object_types(directory: str) -> dict[str, list[str] | None]:\n \"\"\"\n Scans the given directory for all CSV files and reconstructs the available object types.\n Assumes filenames are formatted as \"ObjectType.filename.csv\" or \"ObjectType.csv\".\n\n Args:\n directory (str): The path to the directory containing CSV files.\n\n Returns:\n dict[str, list[str]]: A dictionary mapping object types to lists of file paths.\n \"\"\"\n object_types = defaultdict(list)\n\n for filename in os.listdir(directory):\n if filename.endswith(\".csv\"):\n parts = filename.split(\".\", 1) # Split on the first period\n object_type = parts[0] # Take the first part as the object type\n object_types[object_type].append(os.path.join(directory, filename))\n\n return dict(object_types)\n\n @staticmethod\n def _download_object_csvs(\n all_types_to_filter: dict[str, bool],\n queryable_fields_by_type: dict[str, set[str]],\n directory: str,\n sf_client: OnyxSalesforce,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> None:\n # checkpoint - we've found all object types, now time to fetch the data\n logger.info(\"Fetching CSVs for all object types\")\n\n # This takes like 30 minutes first time and <2 minutes for updates\n object_type_to_csv_path = fetch_all_csvs_in_parallel(\n sf_client=sf_client,\n all_types_to_filter=all_types_to_filter,\n queryable_fields_by_type=queryable_fields_by_type,\n start=start,\n end=end,\n target_dir=directory,\n )\n\n # print useful information\n num_csvs = 0\n num_bytes = 0\n for object_type, csv_paths in object_type_to_csv_path.items():\n if not csv_paths:\n continue\n\n for csv_path in csv_paths:\n if not csv_path:\n continue\n\n file_path = Path(csv_path)\n file_size = file_path.stat().st_size\n num_csvs += 1\n num_bytes += file_size\n logger.info(\n f\"CSV download: object_type={object_type} path={csv_path} bytes={file_size}\"\n )\n\n logger.info(\n f\"CSV download total: total_csvs={num_csvs} total_bytes={num_bytes}\"\n )\n\n @staticmethod\n def _load_csvs_to_db(\n csv_directory: str, remove_ids: bool, sf_db: OnyxSalesforceSQLite\n ) -> dict[str, str]:\n \"\"\"\n Returns a dict of id to object type. Each id is a newly seen row in salesforce.\n \"\"\"\n\n updated_ids: dict[str, str] = {}\n\n object_type_to_csv_path = SalesforceConnector.reconstruct_object_types(\n csv_directory\n )\n\n # NOTE(rkuo): this timing note is meaningless without a reference point in terms\n # of number of records, etc\n # This takes like 10 seconds\n\n # This is for testing the rest of the functionality if data has\n # already been fetched and put in sqlite\n # from import onyx.connectors.salesforce.sf_db.sqlite_functions find_ids_by_type\n # for object_type in self.parent_object_list:\n # updated_ids.update(list(find_ids_by_type(object_type)))\n\n # This takes 10-70 minutes first time (idk why the range is so big)\n total_types = len(object_type_to_csv_path)\n logger.info(f\"Starting to process {total_types} object types\")\n\n for i, (object_type, csv_paths) in enumerate(\n object_type_to_csv_path.items(), 1\n ):\n logger.info(f\"Processing object type {object_type} ({i}/{total_types})\")\n # If path is None, it means it failed to fetch the csv\n if csv_paths is None:\n continue\n\n # Go through each csv path and use it to update the db\n for csv_path in csv_paths:\n num_records = 0\n\n logger.debug(\n f\"Processing CSV: object_type={object_type} \"\n f\"csv={csv_path} \"\n f\"len={Path(csv_path).stat().st_size} \"\n f\"records={num_records}\"\n )\n\n with open(csv_path, \"r\", newline=\"\", encoding=\"utf-8\") as f:\n reader = csv.DictReader(f)\n for row in reader:\n num_records += 1\n\n new_ids = sf_db.update_from_csv(\n object_type=object_type,\n csv_download_path=csv_path,\n remove_ids=remove_ids,\n )\n for new_id in new_ids:\n updated_ids[new_id] = object_type\n\n sf_db.flush()\n\n logger.debug(\n f\"Added {len(new_ids)} new/updated records for {object_type}\"\n )\n\n logger.info(\n f\"Processed CSV: object_type={object_type} \"\n f\"csv={csv_path} \"\n f\"len={Path(csv_path).stat().st_size} \"\n f\"records={num_records} \"\n f\"db_len={sf_db.file_size}\"\n )\n os.remove(csv_path)\n\n return updated_ids\n\n # @staticmethod\n # def _get_child_types(\n # parent_types: list[str], sf_client: OnyxSalesforce\n # ) -> set[str]:\n # all_types: set[str] = set(parent_types)\n\n # # Step 1 - get all object types\n # logger.info(f\"Parent object types: num={len(parent_types)} list={parent_types}\")\n\n # # This takes like 20 seconds\n # for parent_object_type in parent_types:\n # child_types = sf_client.get_children_of_sf_type(parent_object_type)\n # logger.debug(\n # f\"Found {len(child_types)} child types for {parent_object_type}\"\n # )\n\n # all_types.update(child_types.keys())\n\n # # Always want to make sure user is grabbed for permissioning purposes\n # all_types.add(USER_OBJECT_TYPE)\n # # Always want to make sure account is grabbed for reference purposes\n # all_types.add(ACCOUNT_OBJECT_TYPE)\n\n # logger.info(f\"All object types: num={len(all_types)} list={all_types}\")\n\n # # gc.collect()\n # return all_types\n\n # @staticmethod\n # def _get_all_types(parent_types: list[str], sf_client: Salesforce) -> set[str]:\n # all_types: set[str] = set(parent_types)\n\n # # Step 1 - get all object types\n # logger.info(f\"Parent object types: num={len(parent_types)} list={parent_types}\")\n\n # # This takes like 20 seconds\n # for parent_object_type in parent_types:\n # child_types = get_children_of_sf_type(sf_client, parent_object_type)\n # logger.debug(\n # f\"Found {len(child_types)} child types for {parent_object_type}\"\n # )\n\n # all_types.update(child_types)\n\n # # Always want to make sure user is grabbed for permissioning purposes\n # all_types.add(USER_OBJECT_TYPE)\n\n # logger.info(f\"All object types: num={len(all_types)} list={all_types}\")\n\n # # gc.collect()\n # return all_types\n\n def _yield_doc_batches(\n self,\n sf_db: OnyxSalesforceSQLite,\n type_to_processed: dict[str, int],\n changed_ids_to_type: dict[str, str],\n parent_types: set[str],\n increment_parents_changed: Callable[[], None],\n ) -> GenerateDocumentsOutput:\n \"\"\" \"\"\"\n docs_to_yield: list[Document | HierarchyNode] = []\n docs_to_yield_bytes = 0\n\n last_log_time = 0.0\n\n for (\n parent_type,\n parent_id,\n examined_ids,\n ) in sf_db.get_changed_parent_ids_by_type(\n changed_ids=list(changed_ids_to_type.keys()),\n parent_types=parent_types,\n ):\n now = time.monotonic()\n\n processed = examined_ids - 1\n if now - last_log_time > SalesforceConnector.LOG_INTERVAL:\n logger.info(\n f\"Processing stats: {type_to_processed} \"\n f\"file_size={sf_db.file_size} \"\n f\"processed={processed} \"\n f\"remaining={len(changed_ids_to_type) - processed}\"\n )\n last_log_time = now\n\n type_to_processed[parent_type] = type_to_processed.get(parent_type, 0) + 1\n\n parent_object = sf_db.get_record(parent_id, parent_type)\n if not parent_object:\n logger.warning(\n f\"Failed to get parent object {parent_id} for {parent_type}\"\n )\n continue\n\n # use the db to create a document we can yield\n doc = convert_sf_object_to_doc(\n sf_db,\n sf_object=parent_object,\n sf_instance=self.sf_client.sf_instance,\n )\n\n doc.metadata[\"object_type\"] = parent_type\n\n # Add default attributes to the metadata\n for (\n sf_attribute,\n canonical_attribute,\n ) in _DEFAULT_ATTRIBUTES_TO_KEEP.get(parent_type, {}).items():\n if sf_attribute in parent_object.data:\n doc.metadata[canonical_attribute] = _convert_to_metadata_value(\n parent_object.data[sf_attribute]\n )\n\n doc_sizeof = sys.getsizeof(doc)\n docs_to_yield_bytes += doc_sizeof\n docs_to_yield.append(doc)\n increment_parents_changed()\n\n # memory usage is sensitive to the input length, so we're yielding immediately\n # if the batch exceeds a certain byte length\n if (\n len(docs_to_yield) >= self.batch_size\n or docs_to_yield_bytes > SalesforceConnector.MAX_BATCH_BYTES\n ):\n yield docs_to_yield\n docs_to_yield = []\n docs_to_yield_bytes = 0\n\n # observed a memory leak / size issue with the account table if we don't gc.collect here.\n gc.collect()\n\n yield docs_to_yield\n\n def _full_sync(\n self,\n temp_dir: str,\n ) -> GenerateDocumentsOutput:\n type_to_processed: dict[str, int] = {}\n\n logger.info(\"_fetch_from_salesforce starting (full sync).\")\n if not self._sf_client:\n raise RuntimeError(\"self._sf_client is None!\")\n\n changed_ids_to_type: dict[str, str] = {}\n parents_changed = 0\n examined_ids = 0\n\n sf_db = OnyxSalesforceSQLite(os.path.join(temp_dir, \"salesforce_db.sqlite\"))\n sf_db.connect()\n\n try:\n sf_db.apply_schema()\n sf_db.log_stats()\n\n ctx = self._make_context(\n None, None, temp_dir, self.parent_object_list, self._sf_client\n )\n gc.collect()\n\n # Step 2 - load CSV's to sqlite\n object_type_to_csv_paths = SalesforceConnector.reconstruct_object_types(\n temp_dir\n )\n\n total_types = len(object_type_to_csv_paths)\n logger.info(f\"Starting to process {total_types} object types\")\n\n for i, (object_type, csv_paths) in enumerate(\n object_type_to_csv_paths.items(), 1\n ):\n logger.info(f\"Processing object type {object_type} ({i}/{total_types})\")\n # If path is None, it means it failed to fetch the csv\n if csv_paths is None:\n continue\n\n # Go through each csv path and use it to update the db\n for csv_path in csv_paths:\n num_records = 0\n with open(csv_path, \"r\", newline=\"\", encoding=\"utf-8\") as f:\n reader = csv.DictReader(f)\n for row in reader:\n num_records += 1\n\n logger.debug(\n f\"Processing CSV: object_type={object_type} \"\n f\"csv={csv_path} \"\n f\"len={Path(csv_path).stat().st_size} \"\n f\"records={num_records}\"\n )\n\n new_ids = sf_db.update_from_csv(\n object_type=object_type,\n csv_download_path=csv_path,\n )\n for new_id in new_ids:\n changed_ids_to_type[new_id] = object_type\n\n sf_db.flush()\n\n logger.debug(\n f\"Added {len(new_ids)} new/updated records for {object_type}\"\n )\n\n logger.info(\n f\"Processed CSV: object_type={object_type} \"\n f\"csv={csv_path} \"\n f\"len={Path(csv_path).stat().st_size} \"\n f\"records={num_records} \"\n f\"db_len={sf_db.file_size}\"\n )\n\n os.remove(csv_path)\n gc.collect()\n\n gc.collect()\n\n logger.info(f\"Found {len(changed_ids_to_type)} total updated records\")\n logger.info(\n f\"Starting to process parent objects of types: {ctx.parent_types}\"\n )\n\n # Step 3 - extract and index docs\n def increment_parents_changed() -> None:\n nonlocal parents_changed\n parents_changed += 1\n\n yield from self._yield_doc_batches(\n sf_db,\n type_to_processed,\n changed_ids_to_type,\n ctx.parent_types,\n increment_parents_changed,\n )\n except Exception:\n logger.exception(\"Unexpected exception\")\n raise\n finally:\n logger.info(\n f\"Final processing stats: \"\n f\"examined={examined_ids} \"\n f\"parents_changed={parents_changed} \"\n f\"remaining={len(changed_ids_to_type) - examined_ids}\"\n )\n\n logger.info(f\"Top level object types processed: {type_to_processed}\")\n\n sf_db.close()\n\n def _delta_sync(\n self,\n temp_dir: str,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n ) -> GenerateDocumentsOutput:\n type_to_processed: dict[str, int] = {}\n\n logger.info(\"_fetch_from_salesforce starting (delta sync).\")\n if not self._sf_client:\n raise RuntimeError(\"self._sf_client is None!\")\n\n changed_ids_to_type: dict[str, str] = {}\n parents_changed = 0\n processed = 0\n\n sf_db = OnyxSalesforceSQLite(os.path.join(temp_dir, \"salesforce_db.sqlite\"))\n sf_db.connect()\n\n try:\n sf_db.apply_schema()\n sf_db.log_stats()\n\n ctx = self._make_context(\n start, end, temp_dir, self.parent_object_list, self._sf_client\n )\n gc.collect()\n\n # Step 2 - load CSV's to sqlite\n changed_ids_to_type = SalesforceConnector._load_csvs_to_db(\n temp_dir, False, sf_db\n )\n gc.collect()\n\n logger.info(f\"Found {len(changed_ids_to_type)} total updated records\")\n logger.info(\n f\"Starting to process parent objects of types: {ctx.parent_types}\"\n )\n\n # Step 3 - extract and index docs\n docs_to_yield: list[Document | HierarchyNode] = []\n docs_to_yield_bytes = 0\n\n last_log_time = 0.0\n\n # this is a partial sync, so all changed parent id's must be retrieved from salesforce\n # NOTE: it may be an option to identify the object type of an id with its prefix\n # but unfortunately it's possible for an object type to not have a prefix.\n # so that would work in many important cases, but not all.\n for (\n parent_id,\n actual_parent_type,\n num_examined,\n ) in sf_db.get_changed_parent_ids_by_type_2(\n changed_ids=changed_ids_to_type,\n parent_types=ctx.parent_types,\n parent_relationship_fields_by_type=ctx.parent_reference_fields_by_type,\n prefix_to_type=ctx.prefix_to_type,\n ):\n # this yields back each changed parent record, where changed means\n # the parent record itself or a child record was updated.\n now = time.monotonic()\n\n # query salesforce for the changed parent id record\n # NOTE(rkuo): we only know the record id and its possible types,\n # so we actually need to check each type until we succeed\n # to be entirely correct\n # this may be a source of inefficiency and thinking about\n # caching the most likely parent record type might be helpful\n\n # actual_parent_type: str | None = None\n # for possible_parent_type in possible_parent_types:\n # queryable_fields = ctx.queryable_fields_by_type[\n # possible_parent_type\n # ]\n # query = _get_object_by_id_query(\n # parent_id, possible_parent_type, queryable_fields\n # )\n # result = self._sf_client.query(query)\n # if result:\n # actual_parent_type = possible_parent_type\n # print(result)\n # break\n\n # get the parent record fields\n record = self._sf_client.query_object(\n actual_parent_type, parent_id, ctx.type_to_queryable_fields\n )\n if not record:\n continue\n\n # queryable_fields = ctx.type_to_queryable_fields[\n # actual_parent_type\n # ]\n # query = get_object_by_id_query(\n # parent_id, actual_parent_type, queryable_fields\n # )\n # result = self._sf_client.query(query)\n # if not result:\n # continue\n\n # # print(result)\n # record: dict[str, Any] = {}\n\n # record_0 = result[\"records\"][0]\n # for record_key, record_value in record_0.items():\n # if record_key == \"attributes\":\n # continue\n\n # record[record_key] = record_value\n\n # for this parent type, increment the counter on the stats object\n type_to_processed[actual_parent_type] = (\n type_to_processed.get(actual_parent_type, 0) + 1\n )\n\n # get the child records\n child_relationships = ctx.parent_to_child_relationships[\n actual_parent_type\n ]\n relationship_to_queryable_fields = (\n ctx.parent_to_relationship_queryable_fields[actual_parent_type]\n )\n child_records = self.sf_client.get_child_objects_by_id(\n parent_id,\n actual_parent_type,\n list(child_relationships),\n relationship_to_queryable_fields,\n )\n\n # NOTE(rkuo): does using the parent last modified make sense if the update\n # is being triggered because a child object changed?\n primary_owner_list: list[BasicExpertInfo] | None = None\n if \"LastModifiedById\" in record:\n try:\n last_modified_by_id = record[\"LastModifiedById\"]\n user_record = self.sf_client.query_object(\n USER_OBJECT_TYPE,\n last_modified_by_id,\n ctx.type_to_queryable_fields,\n )\n if user_record:\n primary_owner = BasicExpertInfo.from_dict(user_record)\n primary_owner_list = [primary_owner]\n except Exception:\n pass\n\n # for child_record_key, child_record in child_records.items():\n # if not child_record:\n # continue\n\n # child_text_section = _extract_section(\n # child_record,\n # f\"https://{self._sf_client.sf_instance}/{child_record_key}\",\n # )\n # sections.append(child_text_section)\n\n # for parent_relationship_field in parent_relationship_fields:\n # parent_relationship_id\n # json.loads(parent_object.data)\n\n # create and yield a document from the salesforce query\n doc = convert_sf_query_result_to_doc(\n parent_id,\n record,\n child_records,\n primary_owner_list,\n self._sf_client,\n )\n\n # doc = Document(\n # id=ID_PREFIX + parent_id,\n # sections=cast(list[TextSection | ImageSection], sections),\n # source=DocumentSource.SALESFORCE,\n # semantic_identifier=parent_semantic_identifier,\n # doc_updated_at=time_str_to_utc(parent_last_modified_date),\n # primary_owners=primary_owner_list,\n # metadata={},\n # )\n\n # Add default attributes to the metadata\n for (\n sf_attribute,\n canonical_attribute,\n ) in _DEFAULT_ATTRIBUTES_TO_KEEP.get(actual_parent_type, {}).items():\n if sf_attribute in record:\n doc.metadata[canonical_attribute] = _convert_to_metadata_value(\n record[sf_attribute]\n )\n\n doc_sizeof = sys.getsizeof(doc)\n docs_to_yield_bytes += doc_sizeof\n docs_to_yield.append(doc)\n parents_changed += 1\n\n # memory usage is sensitive to the input length, so we're yielding immediately\n # if the batch exceeds a certain byte length\n if (\n len(docs_to_yield) >= self.batch_size\n or docs_to_yield_bytes > SalesforceConnector.MAX_BATCH_BYTES\n ):\n yield docs_to_yield\n docs_to_yield = []\n docs_to_yield_bytes = 0\n\n # observed a memory leak / size issue with the account table if we don't gc.collect here.\n gc.collect()\n\n processed = num_examined\n if now - last_log_time > SalesforceConnector.LOG_INTERVAL:\n logger.info(\n f\"Processing stats: {type_to_processed} \"\n f\"processed={processed} \"\n f\"remaining={len(changed_ids_to_type) - processed}\"\n )\n last_log_time = now\n\n yield docs_to_yield\n except Exception:\n logger.exception(\"Unexpected exception\")\n raise\n finally:\n logger.info(\n f\"Final processing stats: \"\n f\"processed={processed} \"\n f\"remaining={len(changed_ids_to_type) - processed} \"\n f\"parents_changed={parents_changed}\"\n )\n\n logger.info(f\"Top level object types processed: {type_to_processed}\")\n\n sf_db.close()\n\n def _make_context(\n self,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n temp_dir: str,\n parent_object_list: list[str],\n sf_client: OnyxSalesforce,\n ) -> SalesforceConnectorContext:\n \"\"\"NOTE: I suspect we're doing way too many queries here. Likely fewer queries\n and just parsing all the info we need in less passes will work.\"\"\"\n\n parent_types = set(parent_object_list)\n child_types: set[str] = set()\n parent_to_child_types: dict[str, set[str]] = (\n {}\n ) # map from parent to child types\n child_to_parent_types: dict[str, set[str]] = (\n {}\n ) # map from child to parent types\n\n parent_reference_fields_by_type: dict[str, dict[str, list[str]]] = (\n {}\n ) # for a given object, the fields reference parent objects\n type_to_queryable_fields: dict[str, set[str]] = {}\n prefix_to_type: dict[str, str] = {}\n\n parent_to_child_relationships: dict[str, set[str]] = (\n {}\n ) # map from parent to child relationships\n\n # relationship keys are formatted as \"parent__relationship\"\n # we have to do this because relationship names are not unique!\n # values are a dict of relationship names to a list of queryable fields\n parent_to_relationship_queryable_fields: dict[str, dict[str, set[str]]] = {}\n\n parent_child_names_to_relationships: dict[str, str] = {}\n\n full_sync = start is None and end is None\n\n # Step 1 - make a list of all the types to download (parent + direct child + USER_OBJECT_TYPE)\n # prefixes = {}\n\n global_description = sf_client.describe()\n if not global_description:\n raise RuntimeError(\"sf_client.describe failed\")\n\n for sobject in global_description[\"sobjects\"]:\n if sobject[\"keyPrefix\"]:\n prefix_to_type[sobject[\"keyPrefix\"]] = sobject[\"name\"]\n # prefixes[sobject['keyPrefix']] = {\n # 'object_name': sobject['name'],\n # 'label': sobject['label'],\n # 'is_custom': sobject['custom']\n # }\n\n logger.info(f\"Describe: num_prefixes={len(prefix_to_type)}\")\n\n logger.info(f\"Parent object types: num={len(parent_types)} list={parent_types}\")\n for parent_type in parent_types:\n # parent_onyx_sf_type = OnyxSalesforceType(parent_type, sf_client)\n\n custom_fields: list[str] | None = []\n associations_config: dict[str, list[str]] | None = None\n\n # Set queryable fields for parent type\n if self.custom_query_config:\n custom_fields, associations_config = (\n _extract_fields_and_associations_from_config(\n self.custom_query_config, parent_type\n )\n )\n custom_fields = custom_fields or []\n\n # Get custom fields for parent type\n field_set = set(custom_fields)\n # used during doc conversion\n # field_set.add(NAME_FIELD) # does not always exist\n field_set.add(ID_FIELD)\n field_set.add(MODIFIED_FIELD)\n\n # Use only the specified fields\n type_to_queryable_fields[parent_type] = field_set\n logger.info(f\"Using custom fields for {parent_type}: {field_set}\")\n else:\n # Use all queryable fields\n type_to_queryable_fields[parent_type] = (\n sf_client.get_queryable_fields_by_type(parent_type)\n )\n logger.info(f\"Using all fields for {parent_type}\")\n\n child_types_all = sf_client.get_children_of_sf_type(parent_type)\n logger.debug(f\"Found {len(child_types_all)} child types for {parent_type}\")\n logger.debug(f\"child types: {child_types_all}\")\n\n child_types_working = child_types_all.copy()\n if associations_config is not None:\n child_types_working = {\n k: v for k, v in child_types_all.items() if k in associations_config\n }\n any_not_found = False\n for k in associations_config:\n if k not in child_types_working:\n any_not_found = True\n logger.warning(f\"Association {k} not found in {parent_type}\")\n if any_not_found:\n queryable_fields = sf_client.get_queryable_fields_by_type(\n parent_type\n )\n raise RuntimeError(\n f\"Associations {associations_config} not found in {parent_type} \"\n \"make sure your parent-child associations are in the right order\"\n # f\"with child objects {child_types_all}\"\n # f\" and fields {queryable_fields}\"\n )\n\n parent_to_child_relationships[parent_type] = set()\n parent_to_child_types[parent_type] = set()\n parent_to_relationship_queryable_fields[parent_type] = {}\n\n for child_type, child_relationship in child_types_working.items():\n child_type = cast(str, child_type)\n\n # onyx_sf_type = OnyxSalesforceType(child_type, sf_client)\n\n # map parent name to child name\n parent_to_child_types[parent_type].add(child_type)\n\n # reverse map child name to parent name\n if child_type not in child_to_parent_types:\n child_to_parent_types[child_type] = set()\n child_to_parent_types[child_type].add(parent_type)\n\n # map parent name to child relationship\n parent_to_child_relationships[parent_type].add(child_relationship)\n\n # map relationship to queryable fields of the target table\n if config_fields := (\n associations_config and associations_config.get(child_type)\n ):\n field_set = set(config_fields)\n # these are expected and used during doc conversion\n # field_set.add(NAME_FIELD) # does not always exist\n field_set.add(ID_FIELD)\n field_set.add(MODIFIED_FIELD)\n queryable_fields = field_set\n else:\n queryable_fields = sf_client.get_queryable_fields_by_type(\n child_type\n )\n\n if child_relationship in parent_to_relationship_queryable_fields:\n raise RuntimeError(f\"{child_relationship=} already exists\")\n\n parent_to_relationship_queryable_fields[parent_type][\n child_relationship\n ] = queryable_fields\n\n type_to_queryable_fields[child_type] = queryable_fields\n\n parent_child_names_to_relationships[f\"{parent_type}__{child_type}\"] = (\n child_relationship\n )\n\n child_types.update(child_types_working.keys())\n logger.info(\n f\"Child object types: parent={parent_type} num={len(child_types_working)} list={child_types_working.keys()}\"\n )\n\n logger.info(\n f\"Final child object types: num={len(child_types)} list={child_types}\"\n )\n\n all_types: set[str] = set(parent_types)\n all_types.update(child_types)\n\n # NOTE(rkuo): should this be an implicit parent type?\n all_types.add(USER_OBJECT_TYPE) # Always add User for permissioning purposes\n all_types.add(ACCOUNT_OBJECT_TYPE) # Always add Account for reference purposes\n\n logger.info(f\"All object types: num={len(all_types)} list={all_types}\")\n\n # Ensure User and Account have queryable fields if they weren't already processed\n essential_types = [USER_OBJECT_TYPE, ACCOUNT_OBJECT_TYPE]\n for essential_type in essential_types:\n if essential_type not in type_to_queryable_fields:\n type_to_queryable_fields[essential_type] = (\n sf_client.get_queryable_fields_by_type(essential_type)\n )\n\n # 1.1 - Detect all fields in child types which reference a parent type.\n # build dicts to detect relationships between parent and child\n for child_type in child_types.union(essential_types):\n # onyx_sf_type = OnyxSalesforceType(child_type, sf_client)\n parent_reference_fields = sf_client.get_parent_reference_fields(\n child_type, parent_types\n )\n\n parent_reference_fields_by_type[child_type] = parent_reference_fields\n\n # Only add time filter if there is at least one object of the type\n # in the database. We aren't worried about partially completed object update runs\n # because this occurs after we check for existing csvs which covers this case\n # NOTE(rkuo):\n all_types_to_filter: dict[str, bool] = {}\n for sf_type in all_types:\n # onyx_sf_type = OnyxSalesforceType(sf_type, sf_client)\n\n # NOTE(rkuo): I'm not convinced it makes sense to restrict filtering at all\n # all_types_to_filter[sf_type] = sf_db.object_type_count(sf_type) > 0\n all_types_to_filter[sf_type] = not full_sync\n\n # Step 1.2 - bulk download the CSV's for each object type\n SalesforceConnector._download_object_csvs(\n all_types_to_filter,\n type_to_queryable_fields,\n temp_dir,\n sf_client,\n start,\n end,\n )\n\n return_context = SalesforceConnectorContext()\n return_context.parent_types = parent_types\n return_context.child_types = child_types\n return_context.parent_to_child_types = parent_to_child_types\n return_context.child_to_parent_types = child_to_parent_types\n return_context.parent_reference_fields_by_type = parent_reference_fields_by_type\n return_context.type_to_queryable_fields = type_to_queryable_fields\n return_context.prefix_to_type = prefix_to_type\n\n return_context.parent_to_child_relationships = parent_to_child_relationships\n return_context.parent_to_relationship_queryable_fields = (\n parent_to_relationship_queryable_fields\n )\n\n return_context.parent_child_names_to_relationships = (\n parent_child_names_to_relationships\n )\n\n return return_context\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n # Always use a temp directory for SQLite - the database is rebuilt\n # from scratch each time via CSV downloads, so there's no caching benefit\n # from persisting it. Using temp dirs also avoids collisions between\n # multiple CC pairs and eliminates stale WAL/SHM file issues.\n # TODO(evan): make this thing checkpointed and persist/load db from filestore\n with tempfile.TemporaryDirectory() as temp_dir:\n yield from self._full_sync(temp_dir)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n \"\"\"Poll source will synchronize updated parent objects one by one.\"\"\"\n # Always use a temp directory - see comment in load_from_state()\n with tempfile.TemporaryDirectory() as temp_dir:\n yield from self._delta_sync(temp_dir, start, end)\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n end: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n doc_metadata_list: list[SlimDocument | HierarchyNode] = []\n for parent_object_type in self.parent_object_list:\n query = f\"SELECT Id FROM {parent_object_type}\"\n query_result = self.sf_client.safe_query_all(query)\n doc_metadata_list.extend(\n SlimDocument(\n id=f\"{ID_PREFIX}{instance_dict.get('Id', '')}\",\n external_access=None,\n )\n for instance_dict in query_result[\"records\"]\n )\n\n yield doc_metadata_list\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n Validate that the Salesforce credentials and connector settings are correct.\n Specifically checks that we can make an authenticated request to Salesforce.\n \"\"\"\n\n try:\n # Attempt to fetch a small batch of objects (arbitrary endpoint) to verify credentials\n self.sf_client.describe()\n except Exception as e:\n raise ConnectorMissingCredentialError(\n f\"Failed to validate Salesforce credentials. Please check yourcredentials and try again. Error: {e}\"\n )\n\n if self.custom_query_config:\n try:\n _validate_custom_query_config(self.custom_query_config)\n except Exception as e:\n raise ConnectorMissingCredentialError(\n f\"Failed to validate Salesforce custom query config. Please check yourconfig and try again. Error: {e}\"\n )\n\n logger.info(\"Salesforce credentials validated successfully.\")\n\n # @override\n # def load_from_checkpoint(\n # self,\n # start: SecondsSinceUnixEpoch,\n # end: SecondsSinceUnixEpoch,\n # checkpoint: SalesforceCheckpoint,\n # ) -> CheckpointOutput[SalesforceCheckpoint]:\n # try:\n # return self._fetch_document_batches(checkpoint, start, end)\n # except Exception as e:\n # if _should_propagate_error(e) and start is not None:\n # logger.warning(\n # \"Confluence says we provided an invalid 'updated' field. This may indicate\"\n # \"a real issue, but can also appear during edge cases like daylight\"\n # f\"savings time changes. Retrying with a 1 hour offset. Error: {e}\"\n # )\n # return self._fetch_document_batches(checkpoint, start - ONE_HOUR, end)\n # raise\n\n # @override\n # def build_dummy_checkpoint(self) -> SalesforceCheckpoint:\n # return SalesforceCheckpoint(last_updated=0, has_more=True, last_seen_doc_ids=[])\n\n # @override\n # def validate_checkpoint_json(self, checkpoint_json: str) -> SalesforceCheckpoint:\n # return SalesforceCheckpoint.model_validate_json(checkpoint_json)\n\n\nif __name__ == \"__main__\":\n connector = SalesforceConnector(requested_objects=[ACCOUNT_OBJECT_TYPE])\n\n connector.load_credentials(\n {\n \"sf_username\": os.environ[\"SF_USERNAME\"],\n \"sf_password\": os.environ[\"SF_PASSWORD\"],\n \"sf_security_token\": os.environ[\"SF_SECURITY_TOKEN\"],\n }\n )\n start_time = time.monotonic()\n doc_count = 0\n section_count = 0\n text_count = 0\n for doc_batch in connector.load_from_state():\n doc_count += len(doc_batch)\n print(f\"doc_count: {doc_count}\")\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n section_count += len(doc.sections)\n for section in doc.sections:\n if isinstance(section, TextSection) and section.text is not None:\n text_count += len(section.text)\n end_time = time.monotonic()\n\n print(f\"Doc count: {doc_count}\")\n print(f\"Section count: {section_count}\")\n print(f\"Text count: {text_count}\")\n print(f\"Time taken: {end_time - start_time}\")\n" }, { "path": "backend/onyx/connectors/salesforce/doc_conversion.py", "content": "import re\nfrom typing import Any\nfrom typing import cast\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.salesforce.onyx_salesforce import OnyxSalesforce\nfrom onyx.connectors.salesforce.sqlite_functions import OnyxSalesforceSQLite\nfrom onyx.connectors.salesforce.utils import ID_FIELD\nfrom onyx.connectors.salesforce.utils import MODIFIED_FIELD\nfrom onyx.connectors.salesforce.utils import NAME_FIELD\nfrom onyx.connectors.salesforce.utils import SalesforceObject\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nID_PREFIX = \"SALESFORCE_\"\n\n# All of these types of keys are handled by specific fields in the doc\n# conversion process (E.g. URLs) or are not useful for the user (E.g. UUIDs)\n_SF_JSON_FILTER = r\"Id$|Date$|stamp$|url$\"\n\n\ndef _clean_salesforce_dict(data: dict | list) -> dict | list:\n \"\"\"Clean and transform Salesforce API response data by recursively:\n 1. Extracting records from the response if present\n 2. Merging attributes into the main dictionary\n 3. Filtering out keys matching certain patterns (Id, Date, stamp, url)\n 4. Removing '__c' suffix from custom field names\n 5. Removing None values and empty containers\n\n Args:\n data: A dictionary or list from Salesforce API response\n\n Returns:\n Cleaned dictionary or list with transformed keys and filtered values\n \"\"\"\n if isinstance(data, dict):\n if \"records\" in data.keys():\n data = data[\"records\"]\n if isinstance(data, dict):\n if \"attributes\" in data.keys():\n if isinstance(data[\"attributes\"], dict):\n data.update(data.pop(\"attributes\"))\n\n if isinstance(data, dict):\n filtered_dict = {}\n for key, value in data.items():\n if not re.search(_SF_JSON_FILTER, key, re.IGNORECASE):\n # remove the custom object indicator for display\n if \"__c\" in key:\n key = key[:-3]\n if isinstance(value, (dict, list)):\n filtered_value = _clean_salesforce_dict(value)\n # Only add non-empty dictionaries or lists\n if filtered_value:\n filtered_dict[key] = filtered_value\n elif value is not None:\n filtered_dict[key] = value\n return filtered_dict\n\n if isinstance(data, list):\n filtered_list = []\n for item in data:\n filtered_item: dict | list\n if isinstance(item, (dict, list)):\n filtered_item = _clean_salesforce_dict(item)\n # Only add non-empty dictionaries or lists\n if filtered_item:\n filtered_list.append(filtered_item)\n elif item is not None:\n filtered_list.append(item)\n return filtered_list\n\n return data\n\n\ndef _json_to_natural_language(data: dict | list, indent: int = 0) -> str:\n \"\"\"Convert a nested dictionary or list into a human-readable string format.\n\n Recursively traverses the data structure and formats it with:\n - Key-value pairs on separate lines\n - Nested structures indented for readability\n - Lists and dictionaries handled with appropriate formatting\n\n Args:\n data: The dictionary or list to convert\n indent: Number of spaces to indent (default: 0)\n\n Returns:\n A formatted string representation of the data structure\n \"\"\"\n result = []\n indent_str = \" \" * indent\n\n if isinstance(data, dict):\n for key, value in data.items():\n if isinstance(value, (dict, list)):\n result.append(f\"{indent_str}{key}:\")\n result.append(_json_to_natural_language(value, indent + 2))\n else:\n result.append(f\"{indent_str}{key}: {value}\")\n elif isinstance(data, list):\n for item in data:\n result.append(_json_to_natural_language(item, indent + 2))\n\n return \"\\n\".join(result)\n\n\ndef _extract_section(salesforce_object_data: dict[str, Any], link: str) -> TextSection:\n \"\"\"Converts a dict to a TextSection\"\"\"\n\n # Extract text from a Salesforce API response dictionary by:\n # 1. Cleaning the dictionary\n # 2. Converting the cleaned dictionary to natural language\n processed_dict = _clean_salesforce_dict(salesforce_object_data)\n natural_language_for_dict = _json_to_natural_language(processed_dict)\n\n return TextSection(\n text=natural_language_for_dict,\n link=link,\n )\n\n\ndef _extract_primary_owner(\n sf_db: OnyxSalesforceSQLite,\n sf_object: SalesforceObject,\n) -> BasicExpertInfo | None:\n object_dict = sf_object.data\n if not (last_modified_by_id := object_dict.get(\"LastModifiedById\")):\n logger.warning(f\"No LastModifiedById found for {sf_object.id}\")\n return None\n if not (last_modified_by := sf_db.get_record(last_modified_by_id)):\n logger.warning(f\"No LastModifiedBy found for {last_modified_by_id}\")\n return None\n\n user_data = last_modified_by.data\n expert_info = BasicExpertInfo(\n first_name=user_data.get(\"FirstName\"),\n last_name=user_data.get(\"LastName\"),\n email=user_data.get(\"Email\"),\n display_name=user_data.get(NAME_FIELD),\n )\n\n # Check if all fields are None\n if (\n expert_info.first_name is None\n and expert_info.last_name is None\n and expert_info.email is None\n and expert_info.display_name is None\n ):\n logger.warning(f\"No identifying information found for user {user_data}\")\n return None\n\n return expert_info\n\n\ndef convert_sf_query_result_to_doc(\n record_id: str,\n record: dict[str, Any],\n child_records: dict[str, dict[str, Any]],\n primary_owner_list: list[BasicExpertInfo] | None,\n sf_client: OnyxSalesforce,\n) -> Document:\n \"\"\"Generates a yieldable Document from query results\"\"\"\n\n base_url = f\"https://{sf_client.sf_instance}\"\n extracted_doc_updated_at = time_str_to_utc(record[MODIFIED_FIELD])\n extracted_semantic_identifier = record.get(NAME_FIELD) or record.get(\n ID_FIELD, \"Unknown Object\"\n )\n\n sections = [_extract_section(record, f\"{base_url}/{record_id}\")]\n for child_record_key, child_record in child_records.items():\n if not child_record:\n continue\n\n key_fields = child_record_key.split(\":\")\n child_record_id = key_fields[1]\n\n child_text_section = _extract_section(\n child_record,\n f\"{base_url}/{child_record_id}\",\n )\n sections.append(child_text_section)\n\n doc = Document(\n id=f\"{ID_PREFIX}{record_id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.SALESFORCE,\n semantic_identifier=extracted_semantic_identifier,\n doc_updated_at=extracted_doc_updated_at,\n primary_owners=primary_owner_list,\n metadata={},\n )\n return doc\n\n\ndef convert_sf_object_to_doc(\n sf_db: OnyxSalesforceSQLite,\n sf_object: SalesforceObject,\n sf_instance: str,\n) -> Document:\n \"\"\"Would be nice if this function was documented\"\"\"\n object_dict = sf_object.data\n salesforce_id = object_dict[ID_FIELD]\n onyx_salesforce_id = f\"{ID_PREFIX}{salesforce_id}\"\n base_url = f\"https://{sf_instance}\"\n extracted_doc_updated_at = time_str_to_utc(object_dict[MODIFIED_FIELD])\n extracted_semantic_identifier = object_dict.get(NAME_FIELD) or object_dict.get(\n ID_FIELD, \"Unknown Object\"\n )\n\n sections = [_extract_section(sf_object.data, f\"{base_url}/{sf_object.id}\")]\n for id in sf_db.get_child_ids(sf_object.id):\n if not (child_object := sf_db.get_record(id, isChild=True)):\n continue\n sections.append(\n _extract_section(child_object.data, f\"{base_url}/{child_object.id}\")\n )\n\n # NOTE(rkuo): does using the parent last modified make sense if the update\n # is being triggered because a child object changed?\n primary_owner_list: list[BasicExpertInfo] | None = None\n\n primary_owner = sf_db.make_basic_expert_info_from_record(sf_object)\n if primary_owner:\n primary_owner_list = [primary_owner]\n\n doc = Document(\n id=onyx_salesforce_id,\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.SALESFORCE,\n semantic_identifier=extracted_semantic_identifier,\n doc_updated_at=extracted_doc_updated_at,\n primary_owners=primary_owner_list,\n metadata={},\n )\n return doc\n" }, { "path": "backend/onyx/connectors/salesforce/onyx_salesforce.py", "content": "import time\nfrom typing import Any\n\nfrom simple_salesforce import Salesforce\nfrom simple_salesforce import SFType\nfrom simple_salesforce.exceptions import SalesforceRefusedRequest\n\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.salesforce.blacklist import SALESFORCE_BLACKLISTED_OBJECTS\nfrom onyx.connectors.salesforce.blacklist import SALESFORCE_BLACKLISTED_PREFIXES\nfrom onyx.connectors.salesforce.blacklist import SALESFORCE_BLACKLISTED_SUFFIXES\nfrom onyx.connectors.salesforce.salesforce_calls import get_object_by_id_query\nfrom onyx.connectors.salesforce.utils import ID_FIELD\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\n\nlogger = setup_logger()\n\n\ndef is_salesforce_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if an exception is a Salesforce rate limit error.\"\"\"\n return isinstance(\n exception, SalesforceRefusedRequest\n ) and \"REQUEST_LIMIT_EXCEEDED\" in str(exception)\n\n\nclass OnyxSalesforce(Salesforce):\n SOQL_MAX_SUBQUERIES = 20\n\n def __init__(self, *args: Any, **kwargs: Any) -> None:\n super().__init__(*args, **kwargs)\n\n self.parent_types: set[str] = set()\n self.child_types: set[str] = set()\n self.parent_to_child_types: dict[str, set[str]] = (\n {}\n ) # map from parent to child types\n self.child_to_parent_types: dict[str, set[str]] = (\n {}\n ) # map from child to parent types\n self.parent_reference_fields_by_type: dict[str, dict[str, list[str]]] = {}\n self.queryable_fields_by_type: dict[str, list[str]] = {}\n self.prefix_to_type: dict[str, str] = (\n {}\n ) # infer the object type of an id immediately\n\n def initialize(self) -> bool:\n \"\"\"Eventually cache all first run client state with this method\"\"\"\n return True\n\n def is_blacklisted(self, object_type: str) -> bool:\n \"\"\"Returns True if the object type is blacklisted.\"\"\"\n object_type_lower = object_type.lower()\n if object_type_lower in SALESFORCE_BLACKLISTED_OBJECTS:\n return True\n for prefix in SALESFORCE_BLACKLISTED_PREFIXES:\n if object_type_lower.startswith(prefix):\n return True\n\n for suffix in SALESFORCE_BLACKLISTED_SUFFIXES:\n if object_type_lower.endswith(suffix):\n return True\n\n return False\n\n @retry_builder(\n tries=6,\n delay=20,\n backoff=1.5,\n max_delay=60,\n exceptions=(SalesforceRefusedRequest,),\n )\n @rate_limit_builder(max_calls=50, period=60)\n def safe_query(self, query: str, **kwargs: Any) -> dict[str, Any]:\n \"\"\"Wrapper around the original query method with retry logic and rate limiting.\"\"\"\n try:\n return super().query(query, **kwargs)\n except SalesforceRefusedRequest as e:\n if is_salesforce_rate_limit_error(e):\n logger.warning(\n f\"Salesforce rate limit exceeded for query: {query[:100]}...\"\n )\n # Add additional delay for rate limit errors\n time.sleep(5)\n raise\n\n @retry_builder(\n tries=5,\n delay=20,\n backoff=1.5,\n max_delay=60,\n exceptions=(SalesforceRefusedRequest,),\n )\n @rate_limit_builder(max_calls=50, period=60)\n def safe_query_all(self, query: str, **kwargs: Any) -> dict[str, Any]:\n \"\"\"Wrapper around the original query_all method with retry logic and rate limiting.\"\"\"\n try:\n return super().query_all(query, **kwargs)\n except SalesforceRefusedRequest as e:\n if is_salesforce_rate_limit_error(e):\n logger.warning(\n f\"Salesforce rate limit exceeded for query_all: {query[:100]}...\"\n )\n # Add additional delay for rate limit errors\n time.sleep(5)\n raise\n\n @staticmethod\n def _make_child_objects_by_id_query(\n object_id: str,\n sf_type: str,\n child_relationships: list[str],\n relationships_to_fields: dict[str, set[str]],\n ) -> str:\n \"\"\"Returns a SOQL query given the object id, type and child relationships.\n\n object_id: the id of the parent object\n sf_type: the object name/type of the parent object\n child_relationships: a list of the child object names/types to retrieve\n relationships_to_fields: a mapping of objects to their queryable fields\n\n When the query is executed, it comes back as result.records[0][child_relationship]\n \"\"\"\n\n # supposedly the real limit is 200? But we limit to 10 for practical reasons\n SUBQUERY_LIMIT = 10\n\n query = \"SELECT \"\n for child_relationship in child_relationships:\n # TODO(rkuo): what happens if there is a very large list of child records?\n # is that possible problem?\n\n # NOTE: we actually have to list out the subqueries we want.\n # We can't use the following shortcuts:\n # FIELDS(ALL) can include binary fields, so don't use that\n # FIELDS(CUSTOM) can include aggregate queries, so don't use that\n fields = relationships_to_fields[child_relationship]\n fields_fragment = \",\".join(fields)\n query += f\"(SELECT {fields_fragment} FROM {child_relationship} LIMIT {SUBQUERY_LIMIT}), \"\n\n query = query.rstrip(\", \")\n query += f\" FROM {sf_type} WHERE Id = '{object_id}'\"\n return query\n\n def query_object(\n self,\n object_type: str,\n object_id: str,\n type_to_queryable_fields: dict[str, set[str]],\n ) -> dict[str, Any] | None:\n record: dict[str, Any] = {}\n\n queryable_fields = type_to_queryable_fields[object_type]\n query = get_object_by_id_query(object_id, object_type, queryable_fields)\n result = self.safe_query(query)\n if not result:\n return None\n\n record_0 = result[\"records\"][0]\n for record_key, record_value in record_0.items():\n if record_key == \"attributes\":\n continue\n\n record[record_key] = record_value\n\n return record\n\n def get_child_objects_by_id(\n self,\n object_id: str,\n sf_type: str,\n child_relationships: list[str],\n relationships_to_fields: dict[str, set[str]],\n ) -> dict[str, dict[str, Any]]:\n \"\"\"There's a limit on the number of subqueries we can put in a single query.\"\"\"\n child_records: dict[str, dict[str, Any]] = {}\n child_relationships_batch: list[str] = []\n remaining_child_relationships = list(child_relationships)\n\n while True:\n process_batch = False\n\n if (\n len(remaining_child_relationships) == 0\n and len(child_relationships_batch) == 0\n ):\n break\n\n if len(child_relationships_batch) >= OnyxSalesforce.SOQL_MAX_SUBQUERIES:\n process_batch = True\n\n if len(remaining_child_relationships) == 0:\n process_batch = True\n\n if process_batch:\n if len(child_relationships_batch) == 0:\n break\n\n query = OnyxSalesforce._make_child_objects_by_id_query(\n object_id,\n sf_type,\n child_relationships_batch,\n relationships_to_fields,\n )\n\n try:\n result = self.safe_query(query)\n except Exception:\n logger.exception(f\"Query failed: {query=}\")\n else:\n for child_record_key, child_result in result[\"records\"][0].items():\n if child_record_key == \"attributes\":\n continue\n\n if not child_result:\n continue\n\n for child_record in child_result[\"records\"]:\n child_record_id = child_record[ID_FIELD]\n if not child_record_id:\n logger.warning(\"Child record has no id\")\n continue\n\n child_records[f\"{child_record_key}:{child_record_id}\"] = (\n child_record\n )\n finally:\n child_relationships_batch.clear()\n\n continue\n\n if len(remaining_child_relationships) == 0:\n break\n\n child_relationship = remaining_child_relationships.pop(0)\n\n # this is binary content, skip it\n if child_relationship == \"Attachments\":\n continue\n\n child_relationships_batch.append(child_relationship)\n\n return child_records\n\n @retry_builder(\n tries=3,\n delay=1,\n backoff=2,\n exceptions=(SalesforceRefusedRequest,),\n )\n def describe_type(self, name: str) -> Any:\n sf_object = SFType(name, self.session_id, self.sf_instance)\n try:\n result = sf_object.describe()\n return result\n except SalesforceRefusedRequest as e:\n if is_salesforce_rate_limit_error(e):\n logger.warning(\n f\"Salesforce rate limit exceeded for describe_type: {name}\"\n )\n # Add additional delay for rate limit errors\n time.sleep(3)\n raise\n\n def get_queryable_fields_by_type(self, name: str) -> set[str]:\n object_description = self.describe_type(name)\n if object_description is None:\n return set()\n\n fields: list[dict[str, Any]] = object_description[\"fields\"]\n valid_fields: set[str] = set()\n field_names_to_remove: set[str] = set()\n for field in fields:\n if compound_field_name := field.get(\"compoundFieldName\"):\n # We do want to get name fields even if they are compound\n if not field.get(\"nameField\"):\n field_names_to_remove.add(compound_field_name)\n\n field_name = field.get(\"name\")\n field_type = field.get(\"type\")\n if field_type in [\"base64\", \"blob\", \"encryptedstring\"]:\n continue\n\n if field_name:\n valid_fields.add(field_name)\n\n return valid_fields - field_names_to_remove\n\n def get_children_of_sf_type(self, sf_type: str) -> dict[str, str]:\n \"\"\"Returns a dict of child object names to relationship names.\n Relationship names (not object names) are used in subqueries!\n \"\"\"\n names_to_relationships: dict[str, str] = {}\n\n object_description = self.describe_type(sf_type)\n\n index = 0\n len_relationships = len(object_description[\"childRelationships\"])\n for child_relationship in object_description[\"childRelationships\"]:\n child_name = child_relationship[\"childSObject\"]\n\n index += 1\n valid, reason = self._is_valid_child_object(child_relationship)\n if not valid:\n logger.debug(\n f\"{index}/{len_relationships} - Invalid child object: \"\n f\"parent={sf_type} child={child_name} child_field_backreference={child_relationship['field']} {reason=}\"\n )\n continue\n\n logger.debug(\n f\"{index}/{len_relationships} - Found valid child object: \"\n f\"parent={sf_type} child={child_name} child_field_backreference={child_relationship['field']}\"\n )\n\n name = child_name\n relationship = child_relationship[\"relationshipName\"]\n\n names_to_relationships[name] = relationship\n\n return names_to_relationships\n\n def _is_valid_child_object(\n self, child_relationship: dict[str, Any]\n ) -> tuple[bool, str]:\n\n if not child_relationship[\"childSObject\"]:\n return False, \"childSObject is None\"\n\n child_name = child_relationship[\"childSObject\"]\n\n if self.is_blacklisted(child_name):\n return False, f\"{child_name=} is blacklisted.\"\n\n if not child_relationship[\"relationshipName\"]:\n return False, f\"{child_name=} has no relationshipName.\"\n\n object_description = self.describe_type(child_relationship[\"childSObject\"])\n if not object_description[\"queryable\"]:\n return False, f\"{child_name=} is not queryable.\"\n\n if not child_relationship[\"field\"]:\n return False, f\"{child_name=} has no relationship field.\"\n\n if child_relationship[\"field\"] == \"RelatedToId\":\n return False, f\"{child_name=} field is RelatedToId and blacklisted.\"\n\n return True, \"\"\n\n def get_parent_reference_fields(\n self, sf_type: str, parent_types: set[str]\n ) -> dict[str, list[str]]:\n \"\"\"\n sf_type: the type in which to find parent reference fields\n parent_types: a list of parent reference field types we are actually interested in\n Other parent types will not be returned.\n\n Given an object type, returns a dict of field names to a list of referenced parent\n object types.\n (Yes, it is possible for a field to reference one of multiple object types,\n although this seems very unlikely.)\n\n Returns an empty dict if there are no parent reference fields.\n \"\"\"\n\n parent_reference_fields: dict[str, list[str]] = {}\n\n object_description = self.describe_type(sf_type)\n for field in object_description[\"fields\"]:\n if field[\"type\"] == \"reference\":\n for reference_to in field[\"referenceTo\"]:\n if reference_to in parent_types:\n if field[\"name\"] not in parent_reference_fields:\n parent_reference_fields[field[\"name\"]] = []\n parent_reference_fields[field[\"name\"]].append(\n field[\"referenceTo\"]\n )\n\n return parent_reference_fields\n" }, { "path": "backend/onyx/connectors/salesforce/salesforce_calls.py", "content": "import gc\nimport os\nimport time\nfrom concurrent.futures import ThreadPoolExecutor\nfrom datetime import datetime\n\nfrom pytz import UTC\nfrom simple_salesforce import Salesforce\nfrom simple_salesforce.bulk2 import SFBulk2Handler\nfrom simple_salesforce.bulk2 import SFBulk2Type\nfrom simple_salesforce.exceptions import SalesforceRefusedRequest\n\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.salesforce.utils import MODIFIED_FIELD\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\nlogger = setup_logger()\n\n\ndef is_salesforce_rate_limit_error(exception: Exception) -> bool:\n \"\"\"Check if an exception is a Salesforce rate limit error.\"\"\"\n return isinstance(\n exception, SalesforceRefusedRequest\n ) and \"REQUEST_LIMIT_EXCEEDED\" in str(exception)\n\n\ndef _build_last_modified_time_filter_for_salesforce(\n start: SecondsSinceUnixEpoch | None, end: SecondsSinceUnixEpoch | None\n) -> str:\n if start is None or end is None:\n return \"\"\n start_datetime = datetime.fromtimestamp(start, UTC)\n end_datetime = datetime.fromtimestamp(end, UTC)\n return f\" WHERE LastModifiedDate > {start_datetime.isoformat()} AND LastModifiedDate < {end_datetime.isoformat()}\"\n\n\ndef _build_created_date_time_filter_for_salesforce(\n start: SecondsSinceUnixEpoch | None, end: SecondsSinceUnixEpoch | None\n) -> str:\n if start is None or end is None:\n return \"\"\n start_datetime = datetime.fromtimestamp(start, UTC)\n end_datetime = datetime.fromtimestamp(end, UTC)\n return f\" WHERE CreatedDate > {start_datetime.isoformat()} AND CreatedDate < {end_datetime.isoformat()}\"\n\n\ndef _make_time_filter_for_sf_type(\n queryable_fields: set[str],\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n) -> str | None:\n\n if MODIFIED_FIELD in queryable_fields:\n return _build_last_modified_time_filter_for_salesforce(start, end)\n\n if \"CreatedDate\" in queryable_fields:\n return _build_created_date_time_filter_for_salesforce(start, end)\n\n return None\n\n\ndef _make_time_filtered_query(\n queryable_fields: set[str], sf_type: str, time_filter: str\n) -> str:\n query = f\"SELECT {', '.join(queryable_fields)} FROM {sf_type}{time_filter}\"\n return query\n\n\ndef get_object_by_id_query(\n object_id: str, sf_type: str, queryable_fields: set[str]\n) -> str:\n query = (\n f\"SELECT {', '.join(queryable_fields)} FROM {sf_type} WHERE Id = '{object_id}'\"\n )\n return query\n\n\n@retry_builder(\n tries=5,\n delay=2,\n backoff=2,\n max_delay=60,\n exceptions=(SalesforceRefusedRequest,),\n)\n@rate_limit_builder(max_calls=50, period=60)\ndef _object_type_has_api_data(\n sf_client: Salesforce, sf_type: str, time_filter: str\n) -> bool:\n \"\"\"\n Use the rest api to check to make sure the query will result in a non-empty response.\n \"\"\"\n try:\n query = f\"SELECT Count() FROM {sf_type}{time_filter} LIMIT 1\"\n result = sf_client.query(query)\n if result[\"totalSize\"] == 0:\n return False\n except SalesforceRefusedRequest as e:\n if is_salesforce_rate_limit_error(e):\n logger.warning(\n f\"Salesforce rate limit exceeded for object type check: {sf_type}\"\n )\n # Add additional delay for rate limit errors\n time.sleep(3)\n raise\n\n except Exception as e:\n if \"OPERATION_TOO_LARGE\" not in str(e):\n logger.warning(f\"Object type {sf_type} doesn't support query: {e}\")\n return False\n return True\n\n\ndef _bulk_retrieve_from_salesforce(\n sf_type: str,\n query: str,\n target_dir: str,\n sf_client: Salesforce,\n) -> tuple[str, list[str] | None]:\n \"\"\"Returns a tuple of\n 1. the salesforce object type (NOTE: seems redundant)\n 2. the list of CSV's written into the target directory\n \"\"\"\n\n bulk_2_handler: SFBulk2Handler | None = SFBulk2Handler(\n session_id=sf_client.session_id,\n bulk2_url=sf_client.bulk2_url,\n proxies=sf_client.proxies,\n session=sf_client.session,\n )\n if not bulk_2_handler:\n return sf_type, None\n\n # NOTE(rkuo): there are signs this download is allocating large\n # amounts of memory instead of streaming the results to disk.\n # we're doing a gc.collect to try and mitigate this.\n\n # see https://github.com/simple-salesforce/simple-salesforce/issues/428 for a\n # possible solution\n bulk_2_type: SFBulk2Type | None = SFBulk2Type(\n object_name=sf_type,\n bulk2_url=bulk_2_handler.bulk2_url,\n headers=bulk_2_handler.headers,\n session=bulk_2_handler.session,\n )\n if not bulk_2_type:\n return sf_type, None\n\n logger.info(f\"Downloading {sf_type}\")\n\n logger.debug(f\"Query: {query}\")\n\n try:\n # This downloads the file to a file in the target path with a random name\n results = bulk_2_type.download(\n query=query,\n path=target_dir,\n max_records=500000,\n )\n\n # prepend each downloaded csv with the object type (delimiter = '.')\n all_download_paths: list[str] = []\n for result in results:\n original_file_path = result[\"file\"]\n directory, filename = os.path.split(original_file_path)\n new_filename = f\"{sf_type}.{filename}\"\n new_file_path = os.path.join(directory, new_filename)\n os.rename(original_file_path, new_file_path)\n all_download_paths.append(new_file_path)\n except Exception as e:\n logger.error(\n f\"Failed to download salesforce csv for object type {sf_type}: {e}\"\n )\n logger.warning(f\"Exceptioning query for object type {sf_type}: {query}\")\n return sf_type, None\n finally:\n bulk_2_handler = None\n bulk_2_type = None\n gc.collect()\n\n logger.info(f\"Downloaded {sf_type} to {all_download_paths}\")\n return sf_type, all_download_paths\n\n\ndef fetch_all_csvs_in_parallel(\n sf_client: Salesforce,\n all_types_to_filter: dict[str, bool],\n queryable_fields_by_type: dict[str, set[str]],\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n target_dir: str,\n) -> dict[str, list[str] | None]:\n \"\"\"\n Fetches all the csvs in parallel for the given object types\n Returns a dict of (sf_type, full_download_path)\n\n NOTE: We can probably lift object type has api data out of here\n \"\"\"\n\n type_to_query = {}\n\n # query the available fields for each object type and determine how to filter\n for sf_type, apply_filter in all_types_to_filter.items():\n queryable_fields = queryable_fields_by_type[sf_type]\n\n time_filter = \"\"\n while True:\n if not apply_filter:\n break\n\n if start is not None and end is not None:\n time_filter_temp = _make_time_filter_for_sf_type(\n queryable_fields, start, end\n )\n if time_filter_temp is None:\n logger.warning(\n f\"Object type not filterable: type={sf_type} fields={queryable_fields}\"\n )\n time_filter = \"\"\n else:\n logger.info(\n f\"Object type filterable: type={sf_type} filter={time_filter_temp}\"\n )\n time_filter = time_filter_temp\n\n break\n\n if not _object_type_has_api_data(sf_client, sf_type, time_filter):\n logger.warning(f\"Object type skipped (no data available): type={sf_type}\")\n continue\n\n query = _make_time_filtered_query(queryable_fields, sf_type, time_filter)\n type_to_query[sf_type] = query\n\n logger.info(\n f\"Object types to query: initial={len(all_types_to_filter)} queryable={len(type_to_query)}\"\n )\n\n # Run the bulk retrieve in parallel\n # limit to 4 to help with memory usage\n with ThreadPoolExecutor(max_workers=4) as executor:\n results = executor.map(\n lambda object_type: _bulk_retrieve_from_salesforce(\n sf_type=object_type,\n query=type_to_query[object_type],\n target_dir=target_dir,\n sf_client=sf_client,\n ),\n type_to_query.keys(),\n )\n return dict(results)\n" }, { "path": "backend/onyx/connectors/salesforce/shelve_stuff/old_test_salesforce_shelves.py", "content": "import csv\nimport os\nimport shutil\n\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import find_ids_by_type\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import (\n get_affected_parent_ids_by_type,\n)\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_child_ids\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_record\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import (\n update_sf_db_with_csv,\n)\nfrom onyx.connectors.salesforce.utils import BASE_DATA_PATH\nfrom onyx.connectors.salesforce.utils import get_object_type_path\n\n_VALID_SALESFORCE_IDS = [\n \"001bm00000fd9Z3AAI\",\n \"001bm00000fdYTdAAM\",\n \"001bm00000fdYTeAAM\",\n \"001bm00000fdYTfAAM\",\n \"001bm00000fdYTgAAM\",\n \"001bm00000fdYThAAM\",\n \"001bm00000fdYTiAAM\",\n \"001bm00000fdYTjAAM\",\n \"001bm00000fdYTkAAM\",\n \"001bm00000fdYTlAAM\",\n \"001bm00000fdYTmAAM\",\n \"001bm00000fdYTnAAM\",\n \"001bm00000fdYToAAM\",\n \"500bm00000XoOxtAAF\",\n \"500bm00000XoOxuAAF\",\n \"500bm00000XoOxvAAF\",\n \"500bm00000XoOxwAAF\",\n \"500bm00000XoOxxAAF\",\n \"500bm00000XoOxyAAF\",\n \"500bm00000XoOxzAAF\",\n \"500bm00000XoOy0AAF\",\n \"500bm00000XoOy1AAF\",\n \"500bm00000XoOy2AAF\",\n \"500bm00000XoOy3AAF\",\n \"500bm00000XoOy4AAF\",\n \"500bm00000XoOy5AAF\",\n \"500bm00000XoOy6AAF\",\n \"500bm00000XoOy7AAF\",\n \"500bm00000XoOy8AAF\",\n \"500bm00000XoOy9AAF\",\n \"500bm00000XoOyAAAV\",\n \"500bm00000XoOyBAAV\",\n \"500bm00000XoOyCAAV\",\n \"500bm00000XoOyDAAV\",\n \"500bm00000XoOyEAAV\",\n \"500bm00000XoOyFAAV\",\n \"500bm00000XoOyGAAV\",\n \"500bm00000XoOyHAAV\",\n \"500bm00000XoOyIAAV\",\n \"003bm00000EjHCjAAN\",\n \"003bm00000EjHCkAAN\",\n \"003bm00000EjHClAAN\",\n \"003bm00000EjHCmAAN\",\n \"003bm00000EjHCnAAN\",\n \"003bm00000EjHCoAAN\",\n \"003bm00000EjHCpAAN\",\n \"003bm00000EjHCqAAN\",\n \"003bm00000EjHCrAAN\",\n \"003bm00000EjHCsAAN\",\n \"003bm00000EjHCtAAN\",\n \"003bm00000EjHCuAAN\",\n \"003bm00000EjHCvAAN\",\n \"003bm00000EjHCwAAN\",\n \"003bm00000EjHCxAAN\",\n \"003bm00000EjHCyAAN\",\n \"003bm00000EjHCzAAN\",\n \"003bm00000EjHD0AAN\",\n \"003bm00000EjHD1AAN\",\n \"003bm00000EjHD2AAN\",\n \"550bm00000EXc2tAAD\",\n \"006bm000006kyDpAAI\",\n \"006bm000006kyDqAAI\",\n \"006bm000006kyDrAAI\",\n \"006bm000006kyDsAAI\",\n \"006bm000006kyDtAAI\",\n \"006bm000006kyDuAAI\",\n \"006bm000006kyDvAAI\",\n \"006bm000006kyDwAAI\",\n \"006bm000006kyDxAAI\",\n \"006bm000006kyDyAAI\",\n \"006bm000006kyDzAAI\",\n \"006bm000006kyE0AAI\",\n \"006bm000006kyE1AAI\",\n \"006bm000006kyE2AAI\",\n \"006bm000006kyE3AAI\",\n \"006bm000006kyE4AAI\",\n \"006bm000006kyE5AAI\",\n \"006bm000006kyE6AAI\",\n \"006bm000006kyE7AAI\",\n \"006bm000006kyE8AAI\",\n \"006bm000006kyE9AAI\",\n \"006bm000006kyEAAAY\",\n \"006bm000006kyEBAAY\",\n \"006bm000006kyECAAY\",\n \"006bm000006kyEDAAY\",\n \"006bm000006kyEEAAY\",\n \"006bm000006kyEFAAY\",\n \"006bm000006kyEGAAY\",\n \"006bm000006kyEHAAY\",\n \"006bm000006kyEIAAY\",\n \"006bm000006kyEJAAY\",\n \"005bm000009zy0TAAQ\",\n \"005bm000009zy25AAA\",\n \"005bm000009zy26AAA\",\n \"005bm000009zy28AAA\",\n \"005bm000009zy29AAA\",\n \"005bm000009zy2AAAQ\",\n \"005bm000009zy2BAAQ\",\n]\n\n\ndef clear_sf_db() -> None:\n \"\"\"\n Clears the SF DB by deleting all files in the data directory.\n \"\"\"\n shutil.rmtree(BASE_DATA_PATH)\n\n\ndef create_csv_file(\n object_type: str, records: list[dict], filename: str = \"test_data.csv\"\n) -> None:\n \"\"\"\n Creates a CSV file for the given object type and records.\n\n Args:\n object_type: The Salesforce object type (e.g. \"Account\", \"Contact\")\n records: List of dictionaries containing the record data\n filename: Name of the CSV file to create (default: test_data.csv)\n \"\"\"\n if not records:\n return\n\n # Get all unique fields from records\n fields: set[str] = set()\n for record in records:\n fields.update(record.keys())\n fields = set(sorted(list(fields))) # Sort for consistent order\n\n # Create CSV file\n csv_path = os.path.join(get_object_type_path(object_type), filename)\n with open(csv_path, \"w\", newline=\"\", encoding=\"utf-8\") as f:\n writer = csv.DictWriter(f, fieldnames=fields)\n writer.writeheader()\n for record in records:\n writer.writerow(record)\n\n # Update the database with the CSV\n update_sf_db_with_csv(object_type, csv_path)\n\n\ndef create_csv_with_example_data() -> None:\n \"\"\"\n Creates CSV files with example data, organized by object type.\n \"\"\"\n example_data: dict[str, list[dict]] = {\n \"Account\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Acme Inc.\",\n \"BillingCity\": \"New York\",\n \"Industry\": \"Technology\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[1],\n \"Name\": \"Globex Corp\",\n \"BillingCity\": \"Los Angeles\",\n \"Industry\": \"Manufacturing\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[2],\n \"Name\": \"Initech\",\n \"BillingCity\": \"Austin\",\n \"Industry\": \"Software\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[3],\n \"Name\": \"TechCorp Solutions\",\n \"BillingCity\": \"San Francisco\",\n \"Industry\": \"Software\",\n \"AnnualRevenue\": 5000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[4],\n \"Name\": \"BioMed Research\",\n \"BillingCity\": \"Boston\",\n \"Industry\": \"Healthcare\",\n \"AnnualRevenue\": 12000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[5],\n \"Name\": \"Green Energy Co\",\n \"BillingCity\": \"Portland\",\n \"Industry\": \"Energy\",\n \"AnnualRevenue\": 8000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[6],\n \"Name\": \"DataFlow Analytics\",\n \"BillingCity\": \"Seattle\",\n \"Industry\": \"Technology\",\n \"AnnualRevenue\": 3000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[7],\n \"Name\": \"Cloud Nine Services\",\n \"BillingCity\": \"Denver\",\n \"Industry\": \"Cloud Computing\",\n \"AnnualRevenue\": 7000000,\n },\n ],\n \"Contact\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"FirstName\": \"John\",\n \"LastName\": \"Doe\",\n \"Email\": \"john.doe@acme.com\",\n \"Title\": \"CEO\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[41],\n \"FirstName\": \"Jane\",\n \"LastName\": \"Smith\",\n \"Email\": \"jane.smith@acme.com\",\n \"Title\": \"CTO\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[42],\n \"FirstName\": \"Bob\",\n \"LastName\": \"Johnson\",\n \"Email\": \"bob.j@globex.com\",\n \"Title\": \"Sales Director\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[43],\n \"FirstName\": \"Sarah\",\n \"LastName\": \"Chen\",\n \"Email\": \"sarah.chen@techcorp.com\",\n \"Title\": \"Product Manager\",\n \"Phone\": \"415-555-0101\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[44],\n \"FirstName\": \"Michael\",\n \"LastName\": \"Rodriguez\",\n \"Email\": \"m.rodriguez@biomed.com\",\n \"Title\": \"Research Director\",\n \"Phone\": \"617-555-0202\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[45],\n \"FirstName\": \"Emily\",\n \"LastName\": \"Green\",\n \"Email\": \"emily.g@greenenergy.com\",\n \"Title\": \"Sustainability Lead\",\n \"Phone\": \"503-555-0303\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[46],\n \"FirstName\": \"David\",\n \"LastName\": \"Kim\",\n \"Email\": \"david.kim@dataflow.com\",\n \"Title\": \"Data Scientist\",\n \"Phone\": \"206-555-0404\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[47],\n \"FirstName\": \"Rachel\",\n \"LastName\": \"Taylor\",\n \"Email\": \"r.taylor@cloudnine.com\",\n \"Title\": \"Cloud Architect\",\n \"Phone\": \"303-555-0505\",\n },\n ],\n \"Opportunity\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[62],\n \"Name\": \"Acme Server Upgrade\",\n \"Amount\": 50000,\n \"Stage\": \"Prospecting\",\n \"CloseDate\": \"2024-06-30\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[63],\n \"Name\": \"Globex Manufacturing Line\",\n \"Amount\": 150000,\n \"Stage\": \"Negotiation\",\n \"CloseDate\": \"2024-03-15\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[64],\n \"Name\": \"Initech Software License\",\n \"Amount\": 75000,\n \"Stage\": \"Closed Won\",\n \"CloseDate\": \"2024-01-30\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[65],\n \"Name\": \"TechCorp AI Implementation\",\n \"Amount\": 250000,\n \"Stage\": \"Needs Analysis\",\n \"CloseDate\": \"2024-08-15\",\n \"Probability\": 60,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[66],\n \"Name\": \"BioMed Lab Equipment\",\n \"Amount\": 500000,\n \"Stage\": \"Value Proposition\",\n \"CloseDate\": \"2024-09-30\",\n \"Probability\": 75,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[67],\n \"Name\": \"Green Energy Solar Project\",\n \"Amount\": 750000,\n \"Stage\": \"Proposal\",\n \"CloseDate\": \"2024-07-15\",\n \"Probability\": 80,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[68],\n \"Name\": \"DataFlow Analytics Platform\",\n \"Amount\": 180000,\n \"Stage\": \"Negotiation\",\n \"CloseDate\": \"2024-05-30\",\n \"Probability\": 90,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[69],\n \"Name\": \"Cloud Nine Infrastructure\",\n \"Amount\": 300000,\n \"Stage\": \"Qualification\",\n \"CloseDate\": \"2024-10-15\",\n \"Probability\": 40,\n },\n ],\n }\n\n # Create CSV files for each object type\n for object_type, records in example_data.items():\n create_csv_file(object_type, records)\n\n\ndef test_query() -> None:\n \"\"\"\n Tests querying functionality by verifying:\n 1. All expected Account IDs are found\n 2. Each Account's data matches what was inserted\n \"\"\"\n # Expected test data for verification\n expected_accounts: dict[str, dict[str, str | int]] = {\n _VALID_SALESFORCE_IDS[0]: {\n \"Name\": \"Acme Inc.\",\n \"BillingCity\": \"New York\",\n \"Industry\": \"Technology\",\n },\n _VALID_SALESFORCE_IDS[1]: {\n \"Name\": \"Globex Corp\",\n \"BillingCity\": \"Los Angeles\",\n \"Industry\": \"Manufacturing\",\n },\n _VALID_SALESFORCE_IDS[2]: {\n \"Name\": \"Initech\",\n \"BillingCity\": \"Austin\",\n \"Industry\": \"Software\",\n },\n _VALID_SALESFORCE_IDS[3]: {\n \"Name\": \"TechCorp Solutions\",\n \"BillingCity\": \"San Francisco\",\n \"Industry\": \"Software\",\n \"AnnualRevenue\": 5000000,\n },\n _VALID_SALESFORCE_IDS[4]: {\n \"Name\": \"BioMed Research\",\n \"BillingCity\": \"Boston\",\n \"Industry\": \"Healthcare\",\n \"AnnualRevenue\": 12000000,\n },\n _VALID_SALESFORCE_IDS[5]: {\n \"Name\": \"Green Energy Co\",\n \"BillingCity\": \"Portland\",\n \"Industry\": \"Energy\",\n \"AnnualRevenue\": 8000000,\n },\n _VALID_SALESFORCE_IDS[6]: {\n \"Name\": \"DataFlow Analytics\",\n \"BillingCity\": \"Seattle\",\n \"Industry\": \"Technology\",\n \"AnnualRevenue\": 3000000,\n },\n _VALID_SALESFORCE_IDS[7]: {\n \"Name\": \"Cloud Nine Services\",\n \"BillingCity\": \"Denver\",\n \"Industry\": \"Cloud Computing\",\n \"AnnualRevenue\": 7000000,\n },\n }\n\n # Get all Account IDs\n account_ids = find_ids_by_type(\"Account\")\n\n # Verify we found all expected accounts\n assert len(account_ids) == len(\n expected_accounts\n ), f\"Expected {len(expected_accounts)} accounts, found {len(account_ids)}\"\n assert set(account_ids) == set(\n expected_accounts.keys()\n ), \"Found account IDs don't match expected IDs\"\n\n # Verify each account's data\n for acc_id in account_ids:\n combined = get_record(acc_id)\n assert combined is not None, f\"Could not find account {acc_id}\"\n\n expected = expected_accounts[acc_id]\n\n # Verify account data matches\n for key, value in expected.items():\n value = str(value)\n assert (\n combined.data[key] == value\n ), f\"Account {acc_id} field {key} expected {value}, got {combined.data[key]}\"\n\n print(\"All query tests passed successfully!\")\n\n\ndef test_upsert() -> None:\n \"\"\"\n Tests upsert functionality by:\n 1. Updating an existing account\n 2. Creating a new account\n 3. Verifying both operations were successful\n \"\"\"\n # Create CSV for updating an existing account and adding a new one\n update_data: list[dict[str, str | int]] = [\n {\n \"Id\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Acme Inc. Updated\",\n \"BillingCity\": \"New York\",\n \"Industry\": \"Technology\",\n \"Description\": \"Updated company info\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[2],\n \"Name\": \"New Company Inc.\",\n \"BillingCity\": \"Miami\",\n \"Industry\": \"Finance\",\n \"AnnualRevenue\": 1000000,\n },\n ]\n\n create_csv_file(\"Account\", update_data, \"update_data.csv\")\n\n # Verify the update worked\n updated_record = get_record(_VALID_SALESFORCE_IDS[0])\n assert updated_record is not None, \"Updated record not found\"\n assert updated_record.data[\"Name\"] == \"Acme Inc. Updated\", \"Name not updated\"\n assert (\n updated_record.data[\"Description\"] == \"Updated company info\"\n ), \"Description not added\"\n\n # Verify the new record was created\n new_record = get_record(_VALID_SALESFORCE_IDS[2])\n assert new_record is not None, \"New record not found\"\n assert new_record.data[\"Name\"] == \"New Company Inc.\", \"New record name incorrect\"\n assert new_record.data[\"AnnualRevenue\"] == \"1000000\", \"New record revenue incorrect\"\n\n print(\"All upsert tests passed successfully!\")\n\n\ndef test_relationships() -> None:\n \"\"\"\n Tests relationship shelf updates and queries by:\n 1. Creating test data with relationships\n 2. Verifying the relationships are correctly stored\n 3. Testing relationship queries\n \"\"\"\n # Create test data for each object type\n test_data: dict[str, list[dict[str, str | int]]] = {\n \"Case\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[13],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"Subject\": \"Test Case 1\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[14],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"Subject\": \"Test Case 2\",\n },\n ],\n \"Contact\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[48],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"FirstName\": \"Test\",\n \"LastName\": \"Contact\",\n }\n ],\n \"Opportunity\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[62],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Test Opportunity\",\n \"Amount\": 100000,\n }\n ],\n }\n\n # Create and update CSV files for each object type\n for object_type, records in test_data.items():\n create_csv_file(object_type, records, \"relationship_test.csv\")\n\n # Test relationship queries\n # All these objects should be children of Acme Inc.\n child_ids = get_child_ids(_VALID_SALESFORCE_IDS[0])\n assert len(child_ids) == 4, f\"Expected 4 child objects, found {len(child_ids)}\"\n assert _VALID_SALESFORCE_IDS[13] in child_ids, \"Case 1 not found in relationship\"\n assert _VALID_SALESFORCE_IDS[14] in child_ids, \"Case 2 not found in relationship\"\n assert _VALID_SALESFORCE_IDS[48] in child_ids, \"Contact not found in relationship\"\n assert (\n _VALID_SALESFORCE_IDS[62] in child_ids\n ), \"Opportunity not found in relationship\"\n\n # Test querying relationships for a different account (should be empty)\n other_account_children = get_child_ids(_VALID_SALESFORCE_IDS[1])\n assert (\n len(other_account_children) == 0\n ), \"Expected no children for different account\"\n\n print(\"All relationship tests passed successfully!\")\n\n\ndef test_account_with_children() -> None:\n \"\"\"\n Tests querying all accounts and retrieving their child objects.\n This test verifies that:\n 1. All accounts can be retrieved\n 2. Child objects are correctly linked\n 3. Child object data is complete and accurate\n \"\"\"\n # First get all account IDs\n account_ids = find_ids_by_type(\"Account\")\n assert len(account_ids) > 0, \"No accounts found\"\n\n # For each account, get its children and verify the data\n for account_id in account_ids:\n account = get_record(account_id)\n assert account is not None, f\"Could not find account {account_id}\"\n\n # Get all child objects\n child_ids = get_child_ids(account_id)\n\n # For Acme Inc., verify specific relationships\n if account_id == _VALID_SALESFORCE_IDS[0]: # Acme Inc.\n assert (\n len(child_ids) == 4\n ), f\"Expected 4 children for Acme Inc., found {len(child_ids)}\"\n\n # Get all child records\n child_records = []\n for child_id in child_ids:\n child_record = get_record(child_id)\n if child_record is not None:\n child_records.append(child_record)\n # Verify Cases\n cases = [r for r in child_records if r.type == \"Case\"]\n assert (\n len(cases) == 2\n ), f\"Expected 2 cases for Acme Inc., found {len(cases)}\"\n case_subjects = {case.data[\"Subject\"] for case in cases}\n assert \"Test Case 1\" in case_subjects, \"Test Case 1 not found\"\n assert \"Test Case 2\" in case_subjects, \"Test Case 2 not found\"\n\n # Verify Contacts\n contacts = [r for r in child_records if r.type == \"Contact\"]\n assert (\n len(contacts) == 1\n ), f\"Expected 1 contact for Acme Inc., found {len(contacts)}\"\n contact = contacts[0]\n assert contact.data[\"FirstName\"] == \"Test\", \"Contact first name mismatch\"\n assert contact.data[\"LastName\"] == \"Contact\", \"Contact last name mismatch\"\n\n # Verify Opportunities\n opportunities = [r for r in child_records if r.type == \"Opportunity\"]\n assert (\n len(opportunities) == 1\n ), f\"Expected 1 opportunity for Acme Inc., found {len(opportunities)}\"\n opportunity = opportunities[0]\n assert (\n opportunity.data[\"Name\"] == \"Test Opportunity\"\n ), \"Opportunity name mismatch\"\n assert opportunity.data[\"Amount\"] == \"100000\", \"Opportunity amount mismatch\"\n\n print(\"All account with children tests passed successfully!\")\n\n\ndef test_relationship_updates() -> None:\n \"\"\"\n Tests that relationships are properly updated when a child object's parent reference changes.\n This test verifies:\n 1. Initial relationship is created correctly\n 2. When parent reference is updated, old relationship is removed\n 3. New relationship is created correctly\n \"\"\"\n # Create initial test data - Contact linked to Acme Inc.\n initial_contact = [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"FirstName\": \"Test\",\n \"LastName\": \"Contact\",\n }\n ]\n create_csv_file(\"Contact\", initial_contact, \"initial_contact.csv\")\n\n # Verify initial relationship\n acme_children = get_child_ids(_VALID_SALESFORCE_IDS[0])\n assert (\n _VALID_SALESFORCE_IDS[40] in acme_children\n ), \"Initial relationship not created\"\n\n # Update contact to be linked to Globex Corp instead\n updated_contact = [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"AccountId\": _VALID_SALESFORCE_IDS[1],\n \"FirstName\": \"Test\",\n \"LastName\": \"Contact\",\n }\n ]\n create_csv_file(\"Contact\", updated_contact, \"updated_contact.csv\")\n\n # Verify old relationship is removed\n acme_children = get_child_ids(_VALID_SALESFORCE_IDS[0])\n assert (\n _VALID_SALESFORCE_IDS[40] not in acme_children\n ), \"Old relationship not removed\"\n\n # Verify new relationship is created\n globex_children = get_child_ids(_VALID_SALESFORCE_IDS[1])\n assert _VALID_SALESFORCE_IDS[40] in globex_children, \"New relationship not created\"\n\n print(\"All relationship update tests passed successfully!\")\n\n\ndef test_get_affected_parent_ids() -> None:\n \"\"\"\n Tests get_affected_parent_ids functionality by verifying:\n 1. IDs that are directly in the parent_types list are included\n 2. IDs that have children in the updated_ids list are included\n 3. IDs that are neither of the above are not included\n \"\"\"\n # Create test data with relationships\n test_data = {\n \"Account\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Parent Account 1\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[1],\n \"Name\": \"Parent Account 2\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[2],\n \"Name\": \"Not Affected Account\",\n },\n ],\n \"Contact\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"FirstName\": \"Child\",\n \"LastName\": \"Contact\",\n }\n ],\n }\n\n # Create and update CSV files for test data\n for object_type, records in test_data.items():\n create_csv_file(object_type, records)\n\n # Test Case 1: Account directly in updated_ids and parent_types\n updated_ids = {_VALID_SALESFORCE_IDS[1]} # Parent Account 2\n parent_types = [\"Account\"]\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert _VALID_SALESFORCE_IDS[1] in affected_ids, \"Direct parent ID not included\"\n\n # Test Case 2: Account with child in updated_ids\n updated_ids = {_VALID_SALESFORCE_IDS[40]} # Child Contact\n parent_types = [\"Account\"]\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert (\n _VALID_SALESFORCE_IDS[0] in affected_ids\n ), \"Parent of updated child not included\"\n\n # Test Case 3: Both direct and indirect affects\n updated_ids = {_VALID_SALESFORCE_IDS[1], _VALID_SALESFORCE_IDS[40]} # Both cases\n parent_types = [\"Account\"]\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert len(affected_ids) == 2, \"Expected exactly two affected parent IDs\"\n assert _VALID_SALESFORCE_IDS[0] in affected_ids, \"Parent of child not included\"\n assert _VALID_SALESFORCE_IDS[1] in affected_ids, \"Direct parent ID not included\"\n assert (\n _VALID_SALESFORCE_IDS[2] not in affected_ids\n ), \"Unaffected ID incorrectly included\"\n\n # Test Case 4: No matches\n updated_ids = {_VALID_SALESFORCE_IDS[40]} # Child Contact\n parent_types = [\"Opportunity\"] # Wrong type\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert len(affected_ids) == 0, \"Should return empty list when no matches\"\n\n print(\"All get_affected_parent_ids tests passed successfully!\")\n\n\ndef main_build() -> None:\n clear_sf_db()\n create_csv_with_example_data()\n test_query()\n test_upsert()\n test_relationships()\n test_account_with_children()\n test_relationship_updates()\n test_get_affected_parent_ids()\n\n\nif __name__ == \"__main__\":\n main_build()\n" }, { "path": "backend/onyx/connectors/salesforce/shelve_stuff/shelve_functions.py", "content": "import csv\nimport shelve\n\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_utils import (\n get_child_to_parent_shelf_path,\n)\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_utils import get_id_type_shelf_path\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_utils import get_object_shelf_path\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_utils import (\n get_parent_to_child_shelf_path,\n)\nfrom onyx.connectors.salesforce.utils import SalesforceObject\nfrom onyx.connectors.salesforce.utils import validate_salesforce_id\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _update_relationship_shelves(\n child_id: str,\n parent_ids: set[str],\n) -> None:\n \"\"\"Update the relationship shelf when a record is updated.\"\"\"\n try:\n # Convert child_id to string once\n str_child_id = str(child_id)\n\n # First update child to parent mapping\n with shelve.open(\n get_child_to_parent_shelf_path(),\n flag=\"c\",\n protocol=None,\n writeback=True,\n ) as child_to_parent_db:\n old_parent_ids = set(child_to_parent_db.get(str_child_id, []))\n child_to_parent_db[str_child_id] = list(parent_ids)\n\n # Calculate differences outside the next context manager\n parent_ids_to_remove = old_parent_ids - parent_ids\n parent_ids_to_add = parent_ids - old_parent_ids\n\n # Only sync once at the end\n child_to_parent_db.sync()\n\n # Then update parent to child mapping in a single transaction\n if not parent_ids_to_remove and not parent_ids_to_add:\n return\n with shelve.open(\n get_parent_to_child_shelf_path(),\n flag=\"c\",\n protocol=None,\n writeback=True,\n ) as parent_to_child_db:\n # Process all removals first\n for parent_id in parent_ids_to_remove:\n str_parent_id = str(parent_id)\n existing_children = set(parent_to_child_db.get(str_parent_id, []))\n if str_child_id in existing_children:\n existing_children.remove(str_child_id)\n parent_to_child_db[str_parent_id] = list(existing_children)\n\n # Then process all additions\n for parent_id in parent_ids_to_add:\n str_parent_id = str(parent_id)\n existing_children = set(parent_to_child_db.get(str_parent_id, []))\n existing_children.add(str_child_id)\n parent_to_child_db[str_parent_id] = list(existing_children)\n\n # Single sync at the end\n parent_to_child_db.sync()\n\n except Exception as e:\n logger.error(f\"Error updating relationship shelves: {e}\")\n logger.error(f\"Child ID: {child_id}, Parent IDs: {parent_ids}\")\n raise\n\n\ndef get_child_ids(parent_id: str) -> set[str]:\n \"\"\"Get all child IDs for a given parent ID.\n\n Args:\n parent_id: The ID of the parent object\n\n Returns:\n A set of child object IDs\n \"\"\"\n with shelve.open(get_parent_to_child_shelf_path()) as parent_to_child_db:\n return set(parent_to_child_db.get(parent_id, []))\n\n\ndef update_sf_db_with_csv(\n object_type: str,\n csv_download_path: str,\n) -> list[str]:\n \"\"\"Update the SF DB with a CSV file using shelve storage.\"\"\"\n updated_ids = []\n shelf_path = get_object_shelf_path(object_type)\n\n # First read the CSV to get all the data\n with open(csv_download_path, \"r\", newline=\"\", encoding=\"utf-8\") as f:\n reader = csv.DictReader(f)\n for row in reader:\n id = row[\"Id\"]\n parent_ids = set()\n field_to_remove: set[str] = set()\n # Update relationship shelves for any parent references\n for field, value in row.items():\n if validate_salesforce_id(value) and field != \"Id\":\n parent_ids.add(value)\n field_to_remove.add(field)\n if not value:\n field_to_remove.add(field)\n _update_relationship_shelves(id, parent_ids)\n for field in field_to_remove:\n # We use this to extract the Primary Owner later\n if field != \"LastModifiedById\":\n del row[field]\n\n # Update the main object shelf\n with shelve.open(shelf_path) as object_type_db:\n object_type_db[id] = row\n # Update the ID-to-type mapping shelf\n with shelve.open(get_id_type_shelf_path()) as id_type_db:\n id_type_db[id] = object_type\n\n updated_ids.append(id)\n\n # os.remove(csv_download_path)\n return updated_ids\n\n\ndef get_type_from_id(object_id: str) -> str | None:\n \"\"\"Get the type of an object from its ID.\"\"\"\n # Look up the object type from the ID-to-type mapping\n with shelve.open(get_id_type_shelf_path()) as id_type_db:\n if object_id not in id_type_db:\n logger.warning(f\"Object ID {object_id} not found in ID-to-type mapping\")\n return None\n return id_type_db[object_id]\n\n\ndef get_record(\n object_id: str, object_type: str | None = None\n) -> SalesforceObject | None:\n \"\"\"\n Retrieve the record and return it as a SalesforceObject.\n The object type will be looked up from the ID-to-type mapping shelf.\n \"\"\"\n if object_type is None:\n if not (object_type := get_type_from_id(object_id)):\n return None\n\n shelf_path = get_object_shelf_path(object_type)\n with shelve.open(shelf_path) as db:\n if object_id not in db:\n logger.warning(f\"Object ID {object_id} not found in {shelf_path}\")\n return None\n data = db[object_id]\n return SalesforceObject(\n id=object_id,\n type=object_type,\n data=data,\n )\n\n\ndef find_ids_by_type(object_type: str) -> list[str]:\n \"\"\"\n Find all object IDs for rows of the specified type.\n \"\"\"\n shelf_path = get_object_shelf_path(object_type)\n try:\n with shelve.open(shelf_path) as db:\n return list(db.keys())\n except FileNotFoundError:\n return []\n\n\ndef get_affected_parent_ids_by_type(\n updated_ids: set[str], parent_types: list[str]\n) -> dict[str, set[str]]:\n \"\"\"Get IDs of objects that are of the specified parent types and are either in the updated_ids\n or have children in the updated_ids.\n\n Args:\n updated_ids: List of IDs that were updated\n parent_types: List of object types to filter by\n\n Returns:\n A dictionary of IDs that match the criteria\n \"\"\"\n affected_ids_by_type: dict[str, set[str]] = {}\n\n # Check each updated ID\n for updated_id in updated_ids:\n # Add the ID itself if it's of a parent type\n updated_type = get_type_from_id(updated_id)\n if updated_type in parent_types:\n affected_ids_by_type.setdefault(updated_type, set()).add(updated_id)\n continue\n\n # Get parents of this ID and add them if they're of a parent type\n with shelve.open(get_child_to_parent_shelf_path()) as child_to_parent_db:\n parent_ids = child_to_parent_db.get(updated_id, [])\n for parent_id in parent_ids:\n parent_type = get_type_from_id(parent_id)\n if parent_type in parent_types:\n affected_ids_by_type.setdefault(parent_type, set()).add(parent_id)\n\n return affected_ids_by_type\n" }, { "path": "backend/onyx/connectors/salesforce/shelve_stuff/shelve_utils.py", "content": "import os\n\nfrom onyx.connectors.salesforce.utils import BASE_DATA_PATH\nfrom onyx.connectors.salesforce.utils import get_object_type_path\n\n\ndef get_object_shelf_path(object_type: str) -> str:\n \"\"\"Get the path to the shelf file for a specific object type.\"\"\"\n base_path = get_object_type_path(object_type)\n os.makedirs(base_path, exist_ok=True)\n return os.path.join(base_path, \"data.shelf\")\n\n\ndef get_id_type_shelf_path() -> str:\n \"\"\"Get the path to the ID-to-type mapping shelf.\"\"\"\n os.makedirs(BASE_DATA_PATH, exist_ok=True)\n return os.path.join(BASE_DATA_PATH, \"id_type_mapping.shelf.4g\")\n\n\ndef get_parent_to_child_shelf_path() -> str:\n \"\"\"Get the path to the parent-to-child mapping shelf.\"\"\"\n os.makedirs(BASE_DATA_PATH, exist_ok=True)\n return os.path.join(BASE_DATA_PATH, \"parent_to_child_mapping.shelf.4g\")\n\n\ndef get_child_to_parent_shelf_path() -> str:\n \"\"\"Get the path to the child-to-parent mapping shelf.\"\"\"\n os.makedirs(BASE_DATA_PATH, exist_ok=True)\n return os.path.join(BASE_DATA_PATH, \"child_to_parent_mapping.shelf.4g\")\n" }, { "path": "backend/onyx/connectors/salesforce/shelve_stuff/test_salesforce_shelves.py", "content": "import csv\nimport os\nimport shutil\n\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import find_ids_by_type\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import (\n get_affected_parent_ids_by_type,\n)\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_child_ids\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import get_record\nfrom onyx.connectors.salesforce.shelve_stuff.shelve_functions import (\n update_sf_db_with_csv,\n)\nfrom onyx.connectors.salesforce.utils import BASE_DATA_PATH\nfrom onyx.connectors.salesforce.utils import get_object_type_path\n\n_VALID_SALESFORCE_IDS = [\n \"001bm00000fd9Z3AAI\",\n \"001bm00000fdYTdAAM\",\n \"001bm00000fdYTeAAM\",\n \"001bm00000fdYTfAAM\",\n \"001bm00000fdYTgAAM\",\n \"001bm00000fdYThAAM\",\n \"001bm00000fdYTiAAM\",\n \"001bm00000fdYTjAAM\",\n \"001bm00000fdYTkAAM\",\n \"001bm00000fdYTlAAM\",\n \"001bm00000fdYTmAAM\",\n \"001bm00000fdYTnAAM\",\n \"001bm00000fdYToAAM\",\n \"500bm00000XoOxtAAF\",\n \"500bm00000XoOxuAAF\",\n \"500bm00000XoOxvAAF\",\n \"500bm00000XoOxwAAF\",\n \"500bm00000XoOxxAAF\",\n \"500bm00000XoOxyAAF\",\n \"500bm00000XoOxzAAF\",\n \"500bm00000XoOy0AAF\",\n \"500bm00000XoOy1AAF\",\n \"500bm00000XoOy2AAF\",\n \"500bm00000XoOy3AAF\",\n \"500bm00000XoOy4AAF\",\n \"500bm00000XoOy5AAF\",\n \"500bm00000XoOy6AAF\",\n \"500bm00000XoOy7AAF\",\n \"500bm00000XoOy8AAF\",\n \"500bm00000XoOy9AAF\",\n \"500bm00000XoOyAAAV\",\n \"500bm00000XoOyBAAV\",\n \"500bm00000XoOyCAAV\",\n \"500bm00000XoOyDAAV\",\n \"500bm00000XoOyEAAV\",\n \"500bm00000XoOyFAAV\",\n \"500bm00000XoOyGAAV\",\n \"500bm00000XoOyHAAV\",\n \"500bm00000XoOyIAAV\",\n \"003bm00000EjHCjAAN\",\n \"003bm00000EjHCkAAN\",\n \"003bm00000EjHClAAN\",\n \"003bm00000EjHCmAAN\",\n \"003bm00000EjHCnAAN\",\n \"003bm00000EjHCoAAN\",\n \"003bm00000EjHCpAAN\",\n \"003bm00000EjHCqAAN\",\n \"003bm00000EjHCrAAN\",\n \"003bm00000EjHCsAAN\",\n \"003bm00000EjHCtAAN\",\n \"003bm00000EjHCuAAN\",\n \"003bm00000EjHCvAAN\",\n \"003bm00000EjHCwAAN\",\n \"003bm00000EjHCxAAN\",\n \"003bm00000EjHCyAAN\",\n \"003bm00000EjHCzAAN\",\n \"003bm00000EjHD0AAN\",\n \"003bm00000EjHD1AAN\",\n \"003bm00000EjHD2AAN\",\n \"550bm00000EXc2tAAD\",\n \"006bm000006kyDpAAI\",\n \"006bm000006kyDqAAI\",\n \"006bm000006kyDrAAI\",\n \"006bm000006kyDsAAI\",\n \"006bm000006kyDtAAI\",\n \"006bm000006kyDuAAI\",\n \"006bm000006kyDvAAI\",\n \"006bm000006kyDwAAI\",\n \"006bm000006kyDxAAI\",\n \"006bm000006kyDyAAI\",\n \"006bm000006kyDzAAI\",\n \"006bm000006kyE0AAI\",\n \"006bm000006kyE1AAI\",\n \"006bm000006kyE2AAI\",\n \"006bm000006kyE3AAI\",\n \"006bm000006kyE4AAI\",\n \"006bm000006kyE5AAI\",\n \"006bm000006kyE6AAI\",\n \"006bm000006kyE7AAI\",\n \"006bm000006kyE8AAI\",\n \"006bm000006kyE9AAI\",\n \"006bm000006kyEAAAY\",\n \"006bm000006kyEBAAY\",\n \"006bm000006kyECAAY\",\n \"006bm000006kyEDAAY\",\n \"006bm000006kyEEAAY\",\n \"006bm000006kyEFAAY\",\n \"006bm000006kyEGAAY\",\n \"006bm000006kyEHAAY\",\n \"006bm000006kyEIAAY\",\n \"006bm000006kyEJAAY\",\n \"005bm000009zy0TAAQ\",\n \"005bm000009zy25AAA\",\n \"005bm000009zy26AAA\",\n \"005bm000009zy28AAA\",\n \"005bm000009zy29AAA\",\n \"005bm000009zy2AAAQ\",\n \"005bm000009zy2BAAQ\",\n]\n\n\ndef clear_sf_db() -> None:\n \"\"\"\n Clears the SF DB by deleting all files in the data directory.\n \"\"\"\n shutil.rmtree(BASE_DATA_PATH)\n\n\ndef create_csv_file(\n object_type: str, records: list[dict], filename: str = \"test_data.csv\"\n) -> None:\n \"\"\"\n Creates a CSV file for the given object type and records.\n\n Args:\n object_type: The Salesforce object type (e.g. \"Account\", \"Contact\")\n records: List of dictionaries containing the record data\n filename: Name of the CSV file to create (default: test_data.csv)\n \"\"\"\n if not records:\n return\n\n # Get all unique fields from records\n fields: set[str] = set()\n for record in records:\n fields.update(record.keys())\n fields = set(sorted(list(fields))) # Sort for consistent order\n\n # Create CSV file\n csv_path = os.path.join(get_object_type_path(object_type), filename)\n with open(csv_path, \"w\", newline=\"\", encoding=\"utf-8\") as f:\n writer = csv.DictWriter(f, fieldnames=fields)\n writer.writeheader()\n for record in records:\n writer.writerow(record)\n\n # Update the database with the CSV\n update_sf_db_with_csv(object_type, csv_path)\n\n\ndef create_csv_with_example_data() -> None:\n \"\"\"\n Creates CSV files with example data, organized by object type.\n \"\"\"\n example_data: dict[str, list[dict]] = {\n \"Account\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Acme Inc.\",\n \"BillingCity\": \"New York\",\n \"Industry\": \"Technology\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[1],\n \"Name\": \"Globex Corp\",\n \"BillingCity\": \"Los Angeles\",\n \"Industry\": \"Manufacturing\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[2],\n \"Name\": \"Initech\",\n \"BillingCity\": \"Austin\",\n \"Industry\": \"Software\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[3],\n \"Name\": \"TechCorp Solutions\",\n \"BillingCity\": \"San Francisco\",\n \"Industry\": \"Software\",\n \"AnnualRevenue\": 5000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[4],\n \"Name\": \"BioMed Research\",\n \"BillingCity\": \"Boston\",\n \"Industry\": \"Healthcare\",\n \"AnnualRevenue\": 12000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[5],\n \"Name\": \"Green Energy Co\",\n \"BillingCity\": \"Portland\",\n \"Industry\": \"Energy\",\n \"AnnualRevenue\": 8000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[6],\n \"Name\": \"DataFlow Analytics\",\n \"BillingCity\": \"Seattle\",\n \"Industry\": \"Technology\",\n \"AnnualRevenue\": 3000000,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[7],\n \"Name\": \"Cloud Nine Services\",\n \"BillingCity\": \"Denver\",\n \"Industry\": \"Cloud Computing\",\n \"AnnualRevenue\": 7000000,\n },\n ],\n \"Contact\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"FirstName\": \"John\",\n \"LastName\": \"Doe\",\n \"Email\": \"john.doe@acme.com\",\n \"Title\": \"CEO\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[41],\n \"FirstName\": \"Jane\",\n \"LastName\": \"Smith\",\n \"Email\": \"jane.smith@acme.com\",\n \"Title\": \"CTO\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[42],\n \"FirstName\": \"Bob\",\n \"LastName\": \"Johnson\",\n \"Email\": \"bob.j@globex.com\",\n \"Title\": \"Sales Director\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[43],\n \"FirstName\": \"Sarah\",\n \"LastName\": \"Chen\",\n \"Email\": \"sarah.chen@techcorp.com\",\n \"Title\": \"Product Manager\",\n \"Phone\": \"415-555-0101\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[44],\n \"FirstName\": \"Michael\",\n \"LastName\": \"Rodriguez\",\n \"Email\": \"m.rodriguez@biomed.com\",\n \"Title\": \"Research Director\",\n \"Phone\": \"617-555-0202\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[45],\n \"FirstName\": \"Emily\",\n \"LastName\": \"Green\",\n \"Email\": \"emily.g@greenenergy.com\",\n \"Title\": \"Sustainability Lead\",\n \"Phone\": \"503-555-0303\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[46],\n \"FirstName\": \"David\",\n \"LastName\": \"Kim\",\n \"Email\": \"david.kim@dataflow.com\",\n \"Title\": \"Data Scientist\",\n \"Phone\": \"206-555-0404\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[47],\n \"FirstName\": \"Rachel\",\n \"LastName\": \"Taylor\",\n \"Email\": \"r.taylor@cloudnine.com\",\n \"Title\": \"Cloud Architect\",\n \"Phone\": \"303-555-0505\",\n },\n ],\n \"Opportunity\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[62],\n \"Name\": \"Acme Server Upgrade\",\n \"Amount\": 50000,\n \"Stage\": \"Prospecting\",\n \"CloseDate\": \"2024-06-30\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[63],\n \"Name\": \"Globex Manufacturing Line\",\n \"Amount\": 150000,\n \"Stage\": \"Negotiation\",\n \"CloseDate\": \"2024-03-15\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[64],\n \"Name\": \"Initech Software License\",\n \"Amount\": 75000,\n \"Stage\": \"Closed Won\",\n \"CloseDate\": \"2024-01-30\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[65],\n \"Name\": \"TechCorp AI Implementation\",\n \"Amount\": 250000,\n \"Stage\": \"Needs Analysis\",\n \"CloseDate\": \"2024-08-15\",\n \"Probability\": 60,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[66],\n \"Name\": \"BioMed Lab Equipment\",\n \"Amount\": 500000,\n \"Stage\": \"Value Proposition\",\n \"CloseDate\": \"2024-09-30\",\n \"Probability\": 75,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[67],\n \"Name\": \"Green Energy Solar Project\",\n \"Amount\": 750000,\n \"Stage\": \"Proposal\",\n \"CloseDate\": \"2024-07-15\",\n \"Probability\": 80,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[68],\n \"Name\": \"DataFlow Analytics Platform\",\n \"Amount\": 180000,\n \"Stage\": \"Negotiation\",\n \"CloseDate\": \"2024-05-30\",\n \"Probability\": 90,\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[69],\n \"Name\": \"Cloud Nine Infrastructure\",\n \"Amount\": 300000,\n \"Stage\": \"Qualification\",\n \"CloseDate\": \"2024-10-15\",\n \"Probability\": 40,\n },\n ],\n }\n\n # Create CSV files for each object type\n for object_type, records in example_data.items():\n create_csv_file(object_type, records)\n\n\ndef test_query() -> None:\n \"\"\"\n Tests querying functionality by verifying:\n 1. All expected Account IDs are found\n 2. Each Account's data matches what was inserted\n \"\"\"\n # Expected test data for verification\n expected_accounts: dict[str, dict[str, str | int]] = {\n _VALID_SALESFORCE_IDS[0]: {\n \"Name\": \"Acme Inc.\",\n \"BillingCity\": \"New York\",\n \"Industry\": \"Technology\",\n },\n _VALID_SALESFORCE_IDS[1]: {\n \"Name\": \"Globex Corp\",\n \"BillingCity\": \"Los Angeles\",\n \"Industry\": \"Manufacturing\",\n },\n _VALID_SALESFORCE_IDS[2]: {\n \"Name\": \"Initech\",\n \"BillingCity\": \"Austin\",\n \"Industry\": \"Software\",\n },\n _VALID_SALESFORCE_IDS[3]: {\n \"Name\": \"TechCorp Solutions\",\n \"BillingCity\": \"San Francisco\",\n \"Industry\": \"Software\",\n \"AnnualRevenue\": 5000000,\n },\n _VALID_SALESFORCE_IDS[4]: {\n \"Name\": \"BioMed Research\",\n \"BillingCity\": \"Boston\",\n \"Industry\": \"Healthcare\",\n \"AnnualRevenue\": 12000000,\n },\n _VALID_SALESFORCE_IDS[5]: {\n \"Name\": \"Green Energy Co\",\n \"BillingCity\": \"Portland\",\n \"Industry\": \"Energy\",\n \"AnnualRevenue\": 8000000,\n },\n _VALID_SALESFORCE_IDS[6]: {\n \"Name\": \"DataFlow Analytics\",\n \"BillingCity\": \"Seattle\",\n \"Industry\": \"Technology\",\n \"AnnualRevenue\": 3000000,\n },\n _VALID_SALESFORCE_IDS[7]: {\n \"Name\": \"Cloud Nine Services\",\n \"BillingCity\": \"Denver\",\n \"Industry\": \"Cloud Computing\",\n \"AnnualRevenue\": 7000000,\n },\n }\n\n # Get all Account IDs\n account_ids = find_ids_by_type(\"Account\")\n\n # Verify we found all expected accounts\n assert len(account_ids) == len(\n expected_accounts\n ), f\"Expected {len(expected_accounts)} accounts, found {len(account_ids)}\"\n assert set(account_ids) == set(\n expected_accounts.keys()\n ), \"Found account IDs don't match expected IDs\"\n\n # Verify each account's data\n for acc_id in account_ids:\n combined = get_record(acc_id)\n assert combined is not None, f\"Could not find account {acc_id}\"\n\n expected = expected_accounts[acc_id]\n\n # Verify account data matches\n for key, value in expected.items():\n value = str(value)\n assert (\n combined.data[key] == value\n ), f\"Account {acc_id} field {key} expected {value}, got {combined.data[key]}\"\n\n print(\"All query tests passed successfully!\")\n\n\ndef test_upsert() -> None:\n \"\"\"\n Tests upsert functionality by:\n 1. Updating an existing account\n 2. Creating a new account\n 3. Verifying both operations were successful\n \"\"\"\n # Create CSV for updating an existing account and adding a new one\n update_data: list[dict[str, str | int]] = [\n {\n \"Id\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Acme Inc. Updated\",\n \"BillingCity\": \"New York\",\n \"Industry\": \"Technology\",\n \"Description\": \"Updated company info\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[2],\n \"Name\": \"New Company Inc.\",\n \"BillingCity\": \"Miami\",\n \"Industry\": \"Finance\",\n \"AnnualRevenue\": 1000000,\n },\n ]\n\n create_csv_file(\"Account\", update_data, \"update_data.csv\")\n\n # Verify the update worked\n updated_record = get_record(_VALID_SALESFORCE_IDS[0])\n assert updated_record is not None, \"Updated record not found\"\n assert updated_record.data[\"Name\"] == \"Acme Inc. Updated\", \"Name not updated\"\n assert (\n updated_record.data[\"Description\"] == \"Updated company info\"\n ), \"Description not added\"\n\n # Verify the new record was created\n new_record = get_record(_VALID_SALESFORCE_IDS[2])\n assert new_record is not None, \"New record not found\"\n assert new_record.data[\"Name\"] == \"New Company Inc.\", \"New record name incorrect\"\n assert new_record.data[\"AnnualRevenue\"] == \"1000000\", \"New record revenue incorrect\"\n\n print(\"All upsert tests passed successfully!\")\n\n\ndef test_relationships() -> None:\n \"\"\"\n Tests relationship shelf updates and queries by:\n 1. Creating test data with relationships\n 2. Verifying the relationships are correctly stored\n 3. Testing relationship queries\n \"\"\"\n # Create test data for each object type\n test_data: dict[str, list[dict[str, str | int]]] = {\n \"Case\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[13],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"Subject\": \"Test Case 1\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[14],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"Subject\": \"Test Case 2\",\n },\n ],\n \"Contact\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[48],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"FirstName\": \"Test\",\n \"LastName\": \"Contact\",\n }\n ],\n \"Opportunity\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[62],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Test Opportunity\",\n \"Amount\": 100000,\n }\n ],\n }\n\n # Create and update CSV files for each object type\n for object_type, records in test_data.items():\n create_csv_file(object_type, records, \"relationship_test.csv\")\n\n # Test relationship queries\n # All these objects should be children of Acme Inc.\n child_ids = get_child_ids(_VALID_SALESFORCE_IDS[0])\n assert len(child_ids) == 4, f\"Expected 4 child objects, found {len(child_ids)}\"\n assert _VALID_SALESFORCE_IDS[13] in child_ids, \"Case 1 not found in relationship\"\n assert _VALID_SALESFORCE_IDS[14] in child_ids, \"Case 2 not found in relationship\"\n assert _VALID_SALESFORCE_IDS[48] in child_ids, \"Contact not found in relationship\"\n assert (\n _VALID_SALESFORCE_IDS[62] in child_ids\n ), \"Opportunity not found in relationship\"\n\n # Test querying relationships for a different account (should be empty)\n other_account_children = get_child_ids(_VALID_SALESFORCE_IDS[1])\n assert (\n len(other_account_children) == 0\n ), \"Expected no children for different account\"\n\n print(\"All relationship tests passed successfully!\")\n\n\ndef test_account_with_children() -> None:\n \"\"\"\n Tests querying all accounts and retrieving their child objects.\n This test verifies that:\n 1. All accounts can be retrieved\n 2. Child objects are correctly linked\n 3. Child object data is complete and accurate\n \"\"\"\n # First get all account IDs\n account_ids = find_ids_by_type(\"Account\")\n assert len(account_ids) > 0, \"No accounts found\"\n\n # For each account, get its children and verify the data\n for account_id in account_ids:\n account = get_record(account_id)\n assert account is not None, f\"Could not find account {account_id}\"\n\n # Get all child objects\n child_ids = get_child_ids(account_id)\n\n # For Acme Inc., verify specific relationships\n if account_id == _VALID_SALESFORCE_IDS[0]: # Acme Inc.\n assert (\n len(child_ids) == 4\n ), f\"Expected 4 children for Acme Inc., found {len(child_ids)}\"\n\n # Get all child records\n child_records = []\n for child_id in child_ids:\n child_record = get_record(child_id)\n if child_record is not None:\n child_records.append(child_record)\n # Verify Cases\n cases = [r for r in child_records if r.type == \"Case\"]\n assert (\n len(cases) == 2\n ), f\"Expected 2 cases for Acme Inc., found {len(cases)}\"\n case_subjects = {case.data[\"Subject\"] for case in cases}\n assert \"Test Case 1\" in case_subjects, \"Test Case 1 not found\"\n assert \"Test Case 2\" in case_subjects, \"Test Case 2 not found\"\n\n # Verify Contacts\n contacts = [r for r in child_records if r.type == \"Contact\"]\n assert (\n len(contacts) == 1\n ), f\"Expected 1 contact for Acme Inc., found {len(contacts)}\"\n contact = contacts[0]\n assert contact.data[\"FirstName\"] == \"Test\", \"Contact first name mismatch\"\n assert contact.data[\"LastName\"] == \"Contact\", \"Contact last name mismatch\"\n\n # Verify Opportunities\n opportunities = [r for r in child_records if r.type == \"Opportunity\"]\n assert (\n len(opportunities) == 1\n ), f\"Expected 1 opportunity for Acme Inc., found {len(opportunities)}\"\n opportunity = opportunities[0]\n assert (\n opportunity.data[\"Name\"] == \"Test Opportunity\"\n ), \"Opportunity name mismatch\"\n assert opportunity.data[\"Amount\"] == \"100000\", \"Opportunity amount mismatch\"\n\n print(\"All account with children tests passed successfully!\")\n\n\ndef test_relationship_updates() -> None:\n \"\"\"\n Tests that relationships are properly updated when a child object's parent reference changes.\n This test verifies:\n 1. Initial relationship is created correctly\n 2. When parent reference is updated, old relationship is removed\n 3. New relationship is created correctly\n \"\"\"\n # Create initial test data - Contact linked to Acme Inc.\n initial_contact = [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"FirstName\": \"Test\",\n \"LastName\": \"Contact\",\n }\n ]\n create_csv_file(\"Contact\", initial_contact, \"initial_contact.csv\")\n\n # Verify initial relationship\n acme_children = get_child_ids(_VALID_SALESFORCE_IDS[0])\n assert (\n _VALID_SALESFORCE_IDS[40] in acme_children\n ), \"Initial relationship not created\"\n\n # Update contact to be linked to Globex Corp instead\n updated_contact = [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"AccountId\": _VALID_SALESFORCE_IDS[1],\n \"FirstName\": \"Test\",\n \"LastName\": \"Contact\",\n }\n ]\n create_csv_file(\"Contact\", updated_contact, \"updated_contact.csv\")\n\n # Verify old relationship is removed\n acme_children = get_child_ids(_VALID_SALESFORCE_IDS[0])\n assert (\n _VALID_SALESFORCE_IDS[40] not in acme_children\n ), \"Old relationship not removed\"\n\n # Verify new relationship is created\n globex_children = get_child_ids(_VALID_SALESFORCE_IDS[1])\n assert _VALID_SALESFORCE_IDS[40] in globex_children, \"New relationship not created\"\n\n print(\"All relationship update tests passed successfully!\")\n\n\ndef test_get_affected_parent_ids() -> None:\n \"\"\"\n Tests get_affected_parent_ids functionality by verifying:\n 1. IDs that are directly in the parent_types list are included\n 2. IDs that have children in the updated_ids list are included\n 3. IDs that are neither of the above are not included\n \"\"\"\n # Create test data with relationships\n test_data = {\n \"Account\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[0],\n \"Name\": \"Parent Account 1\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[1],\n \"Name\": \"Parent Account 2\",\n },\n {\n \"Id\": _VALID_SALESFORCE_IDS[2],\n \"Name\": \"Not Affected Account\",\n },\n ],\n \"Contact\": [\n {\n \"Id\": _VALID_SALESFORCE_IDS[40],\n \"AccountId\": _VALID_SALESFORCE_IDS[0],\n \"FirstName\": \"Child\",\n \"LastName\": \"Contact\",\n }\n ],\n }\n\n # Create and update CSV files for test data\n for object_type, records in test_data.items():\n create_csv_file(object_type, records)\n\n # Test Case 1: Account directly in updated_ids and parent_types\n updated_ids = {_VALID_SALESFORCE_IDS[1]} # Parent Account 2\n parent_types = [\"Account\"]\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert _VALID_SALESFORCE_IDS[1] in affected_ids, \"Direct parent ID not included\"\n\n # Test Case 2: Account with child in updated_ids\n updated_ids = {_VALID_SALESFORCE_IDS[40]} # Child Contact\n parent_types = [\"Account\"]\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert (\n _VALID_SALESFORCE_IDS[0] in affected_ids\n ), \"Parent of updated child not included\"\n\n # Test Case 3: Both direct and indirect affects\n updated_ids = {_VALID_SALESFORCE_IDS[1], _VALID_SALESFORCE_IDS[40]} # Both cases\n parent_types = [\"Account\"]\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert len(affected_ids) == 2, \"Expected exactly two affected parent IDs\"\n assert _VALID_SALESFORCE_IDS[0] in affected_ids, \"Parent of child not included\"\n assert _VALID_SALESFORCE_IDS[1] in affected_ids, \"Direct parent ID not included\"\n assert (\n _VALID_SALESFORCE_IDS[2] not in affected_ids\n ), \"Unaffected ID incorrectly included\"\n\n # Test Case 4: No matches\n updated_ids = {_VALID_SALESFORCE_IDS[40]} # Child Contact\n parent_types = [\"Opportunity\"] # Wrong type\n affected_ids = get_affected_parent_ids_by_type(updated_ids, parent_types)\n assert len(affected_ids) == 0, \"Should return empty list when no matches\"\n\n print(\"All get_affected_parent_ids tests passed successfully!\")\n\n\ndef main_build() -> None:\n clear_sf_db()\n create_csv_with_example_data()\n test_query()\n test_upsert()\n test_relationships()\n test_account_with_children()\n test_relationship_updates()\n test_get_affected_parent_ids()\n\n\nif __name__ == \"__main__\":\n main_build()\n" }, { "path": "backend/onyx/connectors/salesforce/sqlite_functions.py", "content": "import csv\nimport json\nimport os\nimport sqlite3\nimport time\nfrom collections.abc import Iterator\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import cast\n\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE\nfrom onyx.connectors.salesforce.utils import ID_FIELD\nfrom onyx.connectors.salesforce.utils import NAME_FIELD\nfrom onyx.connectors.salesforce.utils import remove_sqlite_db_files\nfrom onyx.connectors.salesforce.utils import SalesforceObject\nfrom onyx.connectors.salesforce.utils import USER_OBJECT_TYPE\nfrom onyx.connectors.salesforce.utils import validate_salesforce_id\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.utils import batch_list\n\n\nlogger = setup_logger()\n\n\nSQLITE_DISK_IO_ERROR = \"disk I/O error\"\n\n\nclass OnyxSalesforceSQLite:\n \"\"\"Notes on context management using 'with self.conn':\n\n Does autocommit / rollback on exit.\n Does NOT close on exit! .close must be called explicitly.\n \"\"\"\n\n # NOTE(rkuo): this string could probably occur naturally. A more unique value\n # might be appropriate here.\n NULL_ID_STRING = \"N/A\"\n\n def __init__(self, filename: str, isolation_level: str | None = None):\n self.filename = filename\n self.isolation_level = isolation_level\n self._conn: sqlite3.Connection | None = None\n\n # this is only set on connection. This variable does not change\n # when a new db is initialized with this class.\n self._existing_db = True\n\n def __del__(self) -> None:\n self.close()\n\n @property\n def file_size(self) -> int:\n \"\"\"Returns -1 if the file does not exist.\"\"\"\n if not self.filename:\n return -1\n\n if not os.path.exists(self.filename):\n return -1\n\n file_path = Path(self.filename)\n return file_path.stat().st_size\n\n def connect(self) -> None:\n if self._conn is not None:\n self._conn.close()\n self._conn = None\n\n self._existing_db = os.path.exists(self.filename)\n\n # make the path if it doesn't already exist\n os.makedirs(os.path.dirname(self.filename), exist_ok=True)\n\n conn = sqlite3.connect(self.filename, timeout=60.0)\n if self.isolation_level is not None:\n conn.isolation_level = self.isolation_level\n\n self._conn = conn\n\n def close(self) -> None:\n if self._conn is None:\n return\n\n self._conn.close()\n self._conn = None\n\n def cursor(self) -> sqlite3.Cursor:\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n return self._conn.cursor()\n\n def flush(self) -> None:\n \"\"\"We're using SQLite in WAL mode sometimes. To flush to the DB we have to\n call this.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cursor = self._conn.cursor()\n cursor.execute(\"PRAGMA wal_checkpoint(FULL)\")\n\n def apply_schema(self) -> None:\n \"\"\"Initialize the SQLite database with required tables if they don't exist.\n\n Non-destructive operation. If a disk I/O error is encountered (often due\n to stale WAL/SHM files from a previous crash), this method will attempt\n to recover by removing the corrupted files and recreating the database.\n \"\"\"\n try:\n self._apply_schema_impl()\n except sqlite3.OperationalError as e:\n if SQLITE_DISK_IO_ERROR not in str(e):\n raise\n\n logger.warning(f\"SQLite disk I/O error detected, attempting recovery: {e}\")\n self._recover_from_corruption()\n self._apply_schema_impl()\n\n def _recover_from_corruption(self) -> None:\n \"\"\"Recover from SQLite corruption by removing all database files and reconnecting.\"\"\"\n logger.info(f\"Removing corrupted SQLite files: {self.filename}\")\n\n # Close existing connection\n self.close()\n\n # Remove all SQLite files (main db, WAL, SHM)\n remove_sqlite_db_files(self.filename)\n\n # Reconnect - this will create a fresh database\n self.connect()\n\n logger.info(\"SQLite recovery complete, fresh database created\")\n\n def _apply_schema_impl(self) -> None:\n \"\"\"Internal implementation of apply_schema.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n start = time.monotonic()\n\n with self._conn:\n cursor = self._conn.cursor()\n\n if self._existing_db:\n file_path = Path(self.filename)\n file_size = file_path.stat().st_size\n logger.info(f\"init_db - found existing sqlite db: len={file_size}\")\n else:\n # NOTE(rkuo): why is this only if the db doesn't exist?\n\n # Enable WAL mode for better concurrent access and write performance\n cursor.execute(\"PRAGMA journal_mode=WAL\")\n cursor.execute(\"PRAGMA synchronous=NORMAL\")\n cursor.execute(\"PRAGMA temp_store=MEMORY\")\n cursor.execute(\"PRAGMA cache_size=-2000000\") # Use 2GB memory for cache\n\n # Main table for storing Salesforce objects\n cursor.execute(\n \"\"\"\n CREATE TABLE IF NOT EXISTS salesforce_objects (\n id TEXT PRIMARY KEY,\n object_type TEXT NOT NULL,\n data TEXT NOT NULL, -- JSON serialized data\n last_modified INTEGER DEFAULT (strftime('%s', 'now')) -- Add timestamp for better cache management\n ) WITHOUT ROWID -- Optimize for primary key lookups\n \"\"\"\n )\n\n # NOTE(rkuo): this seems completely redundant with relationship_types\n # Table for parent-child relationships with covering index\n cursor.execute(\n \"\"\"\n CREATE TABLE IF NOT EXISTS relationships (\n child_id TEXT NOT NULL,\n parent_id TEXT NOT NULL,\n PRIMARY KEY (child_id, parent_id)\n ) WITHOUT ROWID -- Optimize for primary key lookups\n \"\"\"\n )\n\n # New table for caching parent-child relationships with object types\n cursor.execute(\n \"\"\"\n CREATE TABLE IF NOT EXISTS relationship_types (\n child_id TEXT NOT NULL,\n parent_id TEXT NOT NULL,\n parent_type TEXT NOT NULL,\n PRIMARY KEY (child_id, parent_id, parent_type)\n ) WITHOUT ROWID\n \"\"\"\n )\n\n # Create a table for User email to ID mapping if it doesn't exist\n cursor.execute(\n \"\"\"\n CREATE TABLE IF NOT EXISTS user_email_map (\n email TEXT PRIMARY KEY,\n user_id TEXT, -- Nullable to allow for users without IDs\n FOREIGN KEY (user_id) REFERENCES salesforce_objects(id)\n ) WITHOUT ROWID\n \"\"\"\n )\n\n # Create indexes if they don't exist (SQLite ignores IF NOT EXISTS for indexes)\n def create_index_if_not_exists(\n index_name: str, create_statement: str\n ) -> None:\n cursor.execute(\n f\"SELECT name FROM sqlite_master WHERE type='index' AND name='{index_name}'\"\n )\n if not cursor.fetchone():\n cursor.execute(create_statement)\n\n create_index_if_not_exists(\n \"idx_object_type\",\n \"\"\"\n CREATE INDEX idx_object_type\n ON salesforce_objects(object_type, id)\n WHERE object_type IS NOT NULL\n \"\"\",\n )\n\n create_index_if_not_exists(\n \"idx_parent_id\",\n \"\"\"\n CREATE INDEX idx_parent_id\n ON relationships(parent_id, child_id)\n \"\"\",\n )\n\n create_index_if_not_exists(\n \"idx_child_parent\",\n \"\"\"\n CREATE INDEX idx_child_parent\n ON relationships(child_id)\n WHERE child_id IS NOT NULL\n \"\"\",\n )\n\n create_index_if_not_exists(\n \"idx_relationship_types_lookup\",\n \"\"\"\n CREATE INDEX idx_relationship_types_lookup\n ON relationship_types(parent_type, child_id, parent_id)\n \"\"\",\n )\n\n elapsed = time.monotonic() - start\n logger.info(f\"init_db - create tables and indices: elapsed={elapsed:.2f}\")\n\n # Analyze tables to help query planner\n # NOTE(rkuo): skip ANALYZE - it takes too long and we likely don't have\n # complicated queries that need this\n # start = time.monotonic()\n # cursor.execute(\"ANALYZE relationships\")\n # cursor.execute(\"ANALYZE salesforce_objects\")\n # cursor.execute(\"ANALYZE relationship_types\")\n # cursor.execute(\"ANALYZE user_email_map\")\n # elapsed = time.monotonic() - start\n # logger.info(f\"init_db - analyze: elapsed={elapsed:.2f}\")\n\n # If database already existed but user_email_map needs to be populated\n start = time.monotonic()\n cursor.execute(\"SELECT COUNT(*) FROM user_email_map\")\n elapsed = time.monotonic() - start\n logger.info(f\"init_db - count user_email_map: elapsed={elapsed:.2f}\")\n\n start = time.monotonic()\n if cursor.fetchone()[0] == 0:\n OnyxSalesforceSQLite._update_user_email_map(cursor)\n elapsed = time.monotonic() - start\n logger.info(f\"init_db - update_user_email_map: elapsed={elapsed:.2f}\")\n\n def get_user_id_by_email(self, email: str) -> str | None:\n \"\"\"Get the Salesforce User ID for a given email address.\n\n Args:\n email: The email address to look up\n\n Returns:\n A tuple of (was_found, user_id):\n - was_found: True if the email exists in the table, False if not found\n - user_id: The Salesforce User ID if exists, None otherwise\n \"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cursor = self._conn.cursor()\n cursor.execute(\n \"SELECT user_id FROM user_email_map WHERE email = ?\", (email,)\n )\n result = cursor.fetchone()\n if result is None:\n return None\n return result[0]\n\n def update_email_to_id_table(self, email: str, id: str | None) -> None:\n \"\"\"Update the email to ID map table with a new email and ID.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n id_to_use = id or self.NULL_ID_STRING\n with self._conn:\n cursor = self._conn.cursor()\n cursor.execute(\n \"INSERT OR REPLACE INTO user_email_map (email, user_id) VALUES (?, ?)\",\n (email, id_to_use),\n )\n\n def log_stats(self) -> None:\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cache_pages = self._conn.execute(\"PRAGMA cache_size\").fetchone()[0]\n page_size = self._conn.execute(\"PRAGMA page_size\").fetchone()[0]\n if cache_pages >= 0:\n cache_bytes = cache_pages * page_size\n else:\n cache_bytes = abs(cache_pages * 1024)\n logger.info(\n f\"SQLite stats: sqlite_version={sqlite3.sqlite_version} \"\n f\"cache_pages={cache_pages} \"\n f\"page_size={page_size} \"\n f\"cache_bytes={cache_bytes}\"\n )\n\n # get_changed_parent_ids_by_type_2 replaces this\n def get_changed_parent_ids_by_type(\n self,\n changed_ids: list[str],\n parent_types: set[str],\n batch_size: int = 500,\n ) -> Iterator[tuple[str, str, int]]:\n \"\"\"Get IDs of objects that are of the specified parent types and are either in the\n updated_ids or have children in the updated_ids. Yields tuples of (parent_type, affected_ids, num_examined).\n\n NOTE(rkuo): This function used to have some interesting behavior ... it created batches of id's\n and yielded back a list once for each parent type within that batch.\n\n There's no need to expose the details of the internal batching to the caller, so\n we're now yielding once per changed parent.\n \"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n updated_parent_ids: set[str] = (\n set()\n ) # dedupes parent id's that have already been yielded\n\n # SQLite typically has a limit of 999 variables\n num_examined = 0\n updated_ids_batches = batch_list(changed_ids, batch_size)\n\n with self._conn:\n cursor = self._conn.cursor()\n\n for batch_ids in updated_ids_batches:\n num_examined += len(batch_ids)\n\n batch_ids = list(set(batch_ids) - updated_parent_ids)\n if not batch_ids:\n continue\n id_placeholders = \",\".join([\"?\" for _ in batch_ids])\n\n for parent_type in parent_types:\n affected_ids: set[str] = set()\n\n # Get directly updated objects of parent types - using index on object_type\n cursor.execute(\n f\"\"\"\n SELECT id FROM salesforce_objects\n WHERE id IN ({id_placeholders})\n AND object_type = ?\n \"\"\",\n batch_ids + [parent_type],\n )\n affected_ids.update(row[0] for row in cursor.fetchall())\n\n # Get parent objects of updated objects - using optimized relationship_types table\n cursor.execute(\n f\"\"\"\n SELECT DISTINCT parent_id\n FROM relationship_types\n INDEXED BY idx_relationship_types_lookup\n WHERE parent_type = ?\n AND child_id IN ({id_placeholders})\n \"\"\",\n [parent_type] + batch_ids,\n )\n affected_ids.update(row[0] for row in cursor.fetchall())\n\n # Remove any parent IDs that have already been processed\n newly_affected_ids = affected_ids - updated_parent_ids\n # Add the new affected IDs to the set of updated parent IDs\n if newly_affected_ids:\n # Yield each newly affected ID individually\n for parent_id in newly_affected_ids:\n yield parent_type, parent_id, num_examined\n\n updated_parent_ids.update(newly_affected_ids)\n\n def get_changed_parent_ids_by_type_2(\n self,\n changed_ids: dict[str, str],\n parent_types: set[str],\n parent_relationship_fields_by_type: dict[str, dict[str, list[str]]],\n prefix_to_type: dict[str, str],\n ) -> Iterator[tuple[str, str, int]]:\n \"\"\"\n This function yields back any changed parent id's based on\n a relationship lookup.\n\n Yields tuples of (changed_id, parent_type, num_examined)\n changed_id is the id of the changed parent record\n parent_type is the object table/type of the id (based on a prefix lookup)\n num_examined is an integer which signifies our progress through the changed_id's dict\n\n changed_ids is a list of all id's that changed, both parent and children.\n parent\n\n This is much simpler than get_changed_parent_ids_by_type.\n\n TODO(rkuo): for common entities, the first 3 chars identify the object type\n see https://help.salesforce.com/s/articleView?id=000385203&type=1\n \"\"\"\n changed_parent_ids: set[str] = (\n set()\n ) # dedupes parent id's that have already been yielded\n\n # SQLite typically has a limit of 999 variables\n num_examined = 0\n\n for changed_id, changed_type in changed_ids.items():\n num_examined += 1\n\n # if we yielded this id already, continue\n if changed_id in changed_parent_ids:\n continue\n\n # if this id is a parent type, yield it directly\n if changed_type in parent_types:\n yield changed_id, changed_type, num_examined\n changed_parent_ids.add(changed_id)\n continue\n\n # if this id is a child type, then check the columns\n # that relate it to the parent id and yield those ids\n # NOTE: Although unlikely, id's yielded in this way may not be of the\n # type we're interested in, so the caller must be prepared\n # for the id to not be present\n\n # get the child id record\n sf_object = self.get_record(changed_id, changed_type)\n if not sf_object:\n continue\n\n # get the fields that contain parent id's\n parent_relationship_fields = parent_relationship_fields_by_type[\n changed_type\n ]\n for field_name, _ in parent_relationship_fields.items():\n if field_name not in sf_object.data:\n logger.warning(f\"{field_name=} not in data for {changed_type=}!\")\n continue\n\n parent_id = cast(str, sf_object.data[field_name])\n parent_id_prefix = parent_id[:3]\n\n if parent_id_prefix not in prefix_to_type:\n logger.warning(\n f\"Could not lookup type for prefix: {parent_id_prefix=}\"\n )\n continue\n\n parent_type = prefix_to_type[parent_id_prefix]\n if parent_type not in parent_types:\n continue\n\n yield parent_id, parent_type, num_examined\n changed_parent_ids.add(parent_id)\n break\n\n def object_type_count(self, object_type: str) -> int:\n \"\"\"Check if there is at least one object of the specified type in the database.\n\n Args:\n object_type: The Salesforce object type to check\n\n Returns:\n bool: True if at least one object exists, False otherwise\n \"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cursor = self._conn.cursor()\n cursor.execute(\n \"SELECT COUNT(*) FROM salesforce_objects WHERE object_type = ?\",\n (object_type,),\n )\n count = cursor.fetchone()[0]\n return count\n\n @staticmethod\n def normalize_record(\n original_record: dict[str, Any],\n remove_ids: bool = True,\n ) -> tuple[dict[str, Any], set[str]]:\n \"\"\"Takes a dict of field names to values and removes fields\n we don't want.\n\n This means most parent id field's and any fields with null values.\n\n Return a json string and a list of parent_id's in the record.\n \"\"\"\n parent_ids: set[str] = set()\n fields_to_remove: set[str] = set()\n\n record = original_record.copy()\n\n for field, value in record.items():\n # remove empty fields\n if not value:\n fields_to_remove.add(field)\n continue\n\n if field == \"attributes\":\n fields_to_remove.add(field)\n continue\n\n # remove salesforce id's (and add to parent id set)\n if (\n field != ID_FIELD\n and isinstance(value, str)\n and validate_salesforce_id(value)\n ):\n parent_ids.add(value)\n if remove_ids:\n fields_to_remove.add(field)\n continue\n\n # this field is real data, leave it alone\n\n # Remove unwanted fields\n for field in fields_to_remove:\n if field != \"LastModifiedById\":\n del record[field]\n\n return record, parent_ids\n\n def update_from_csv(\n self, object_type: str, csv_download_path: str, remove_ids: bool = True\n ) -> list[str]:\n \"\"\"Update the SF DB with a CSV file using SQLite storage.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n # some customers need this to be larger than the default 128KB, go with 16MB\n csv.field_size_limit(16 * 1024 * 1024)\n\n updated_ids = []\n\n with self._conn:\n cursor = self._conn.cursor()\n\n with open(csv_download_path, \"r\", newline=\"\", encoding=\"utf-8\") as f:\n reader = csv.DictReader(f)\n uncommitted_rows = 0\n for row in reader:\n if ID_FIELD not in row:\n logger.warning(\n f\"Row {row} does not have an {ID_FIELD} field in {csv_download_path}\"\n )\n continue\n\n row_id = row[ID_FIELD]\n\n normalized_record, parent_ids = (\n OnyxSalesforceSQLite.normalize_record(row, remove_ids)\n )\n normalized_record_json_str = json.dumps(normalized_record)\n\n # Update main object data\n # NOTE(rkuo): looks like we take a list and dump it as json into the db\n cursor.execute(\n \"\"\"\n INSERT OR REPLACE INTO salesforce_objects (id, object_type, data)\n VALUES (?, ?, ?)\n \"\"\",\n (row_id, object_type, normalized_record_json_str),\n )\n\n # Update relationships using the same connection\n OnyxSalesforceSQLite._update_relationship_tables(\n cursor, row_id, parent_ids\n )\n updated_ids.append(row_id)\n\n # periodically commit or else memory will balloon\n uncommitted_rows += 1\n if uncommitted_rows >= 1024:\n self._conn.commit()\n uncommitted_rows = 0\n\n # If we're updating User objects, update the email map\n if object_type == USER_OBJECT_TYPE:\n OnyxSalesforceSQLite._update_user_email_map(cursor)\n\n return updated_ids\n\n def get_child_ids(self, parent_id: str) -> set[str]:\n \"\"\"Get all child IDs for a given parent ID.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cursor = self._conn.cursor()\n\n # Force index usage with INDEXED BY\n cursor.execute(\n \"SELECT child_id FROM relationships INDEXED BY idx_parent_id WHERE parent_id = ?\",\n (parent_id,),\n )\n child_ids = {row[0] for row in cursor.fetchall()}\n return child_ids\n\n def get_type_from_id(self, object_id: str) -> str | None:\n \"\"\"Get the type of an object from its ID.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cursor = self._conn.cursor()\n cursor.execute(\n \"SELECT object_type FROM salesforce_objects WHERE id = ?\", (object_id,)\n )\n result = cursor.fetchone()\n if not result:\n logger.warning(f\"Object ID {object_id} not found\")\n return None\n return result[0]\n\n def get_record(\n self, object_id: str, object_type: str | None = None, isChild: bool = False\n ) -> SalesforceObject | None:\n \"\"\"Retrieve the record and return it as a SalesforceObject.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n if object_type is None:\n object_type = self.get_type_from_id(object_id)\n if not object_type:\n return None\n\n with self._conn:\n cursor = self._conn.cursor()\n # Get the object data and account data\n if object_type == ACCOUNT_OBJECT_TYPE or isChild:\n cursor.execute(\n \"SELECT data FROM salesforce_objects WHERE id = ?\", (object_id,)\n )\n else:\n cursor.execute(\n \"SELECT pso.data, r.parent_id as parent_id, sso.object_type FROM salesforce_objects pso \\\n LEFT JOIN relationships r on r.child_id = pso.id \\\n LEFT JOIN salesforce_objects sso on r.parent_id = sso.id \\\n WHERE pso.id = ? \",\n (object_id,),\n )\n result = cursor.fetchall()\n if not result:\n logger.warning(f\"Object ID {object_id} not found\")\n return None\n\n data = json.loads(result[0][0])\n\n if object_type != ACCOUNT_OBJECT_TYPE:\n # convert any account ids of the relationships back into data fields, with name\n for row in result:\n # the following skips Account objects.\n if len(row) < 3:\n continue\n\n if row[1] and row[2] and row[2] == ACCOUNT_OBJECT_TYPE:\n data[\"AccountId\"] = row[1]\n cursor.execute(\n \"SELECT data FROM salesforce_objects WHERE id = ?\",\n (row[1],),\n )\n account_data = json.loads(cursor.fetchone()[0])\n data[ACCOUNT_OBJECT_TYPE] = account_data.get(NAME_FIELD, \"\")\n\n return SalesforceObject(id=object_id, type=object_type, data=data)\n\n def find_ids_by_type(self, object_type: str) -> list[str]:\n \"\"\"Find all object IDs for rows of the specified type.\"\"\"\n if self._conn is None:\n raise RuntimeError(\"Database connection is closed\")\n\n with self._conn:\n cursor = self._conn.cursor()\n cursor.execute(\n \"SELECT id FROM salesforce_objects WHERE object_type = ?\",\n (object_type,),\n )\n return [row[0] for row in cursor.fetchall()]\n\n @staticmethod\n def _update_relationship_tables(\n cursor: sqlite3.Cursor, child_id: str, parent_ids: set[str]\n ) -> None:\n \"\"\"Given a child id and a set of parent id's, updates the\n relationships of the child to the parents in the db and removes old relationships.\n\n Args:\n conn: The database connection to use (must be in a transaction)\n child_id: The ID of the child record\n parent_ids: Set of parent IDs to link to\n \"\"\"\n\n try:\n # Get existing parent IDs\n cursor.execute(\n \"SELECT parent_id FROM relationships WHERE child_id = ?\", (child_id,)\n )\n old_parent_ids = {row[0] for row in cursor.fetchall()}\n\n # Calculate differences\n parent_ids_to_remove = old_parent_ids - parent_ids\n parent_ids_to_add = parent_ids - old_parent_ids\n\n # Remove old relationships\n if parent_ids_to_remove:\n cursor.executemany(\n \"DELETE FROM relationships WHERE child_id = ? AND parent_id = ?\",\n [(child_id, parent_id) for parent_id in parent_ids_to_remove],\n )\n # Also remove from relationship_types\n cursor.executemany(\n \"DELETE FROM relationship_types WHERE child_id = ? AND parent_id = ?\",\n [(child_id, parent_id) for parent_id in parent_ids_to_remove],\n )\n\n # Add new relationships\n if parent_ids_to_add:\n # First add to relationships table\n cursor.executemany(\n \"INSERT INTO relationships (child_id, parent_id) VALUES (?, ?)\",\n [(child_id, parent_id) for parent_id in parent_ids_to_add],\n )\n\n # Then get the types of the parent objects and add to relationship_types\n for parent_id in parent_ids_to_add:\n cursor.execute(\n \"SELECT object_type FROM salesforce_objects WHERE id = ?\",\n (parent_id,),\n )\n result = cursor.fetchone()\n if result:\n parent_type = result[0]\n cursor.execute(\n \"\"\"\n INSERT INTO relationship_types (child_id, parent_id, parent_type)\n VALUES (?, ?, ?)\n \"\"\",\n (child_id, parent_id, parent_type),\n )\n\n except Exception:\n logger.exception(\n f\"Error updating relationship tables: child_id={child_id} parent_ids={parent_ids}\"\n )\n raise\n\n @staticmethod\n def _update_user_email_map(cursor: sqlite3.Cursor) -> None:\n \"\"\"Update the user_email_map table with current User objects.\n Called internally by update_sf_db_with_csv when User objects are updated.\n \"\"\"\n\n cursor.execute(\n \"\"\"\n INSERT OR REPLACE INTO user_email_map (email, user_id)\n SELECT json_extract(data, '$.Email'), id\n FROM salesforce_objects\n WHERE object_type = 'User'\n AND json_extract(data, '$.Email') IS NOT NULL\n \"\"\"\n )\n\n def make_basic_expert_info_from_record(\n self,\n sf_object: SalesforceObject,\n ) -> BasicExpertInfo | None:\n \"\"\"Parses record for LastModifiedById and returns BasicExpertInfo\n of the user if possible.\"\"\"\n object_dict: dict[str, Any] = sf_object.data\n if not (last_modified_by_id := object_dict.get(\"LastModifiedById\")):\n logger.warning(f\"No LastModifiedById found for {sf_object.id}\")\n return None\n if not (last_modified_by := self.get_record(last_modified_by_id)):\n logger.warning(f\"No LastModifiedBy found for {last_modified_by_id}\")\n return None\n\n try:\n expert_info = BasicExpertInfo.from_dict(last_modified_by.data)\n except Exception:\n return None\n\n return expert_info\n" }, { "path": "backend/onyx/connectors/salesforce/utils.py", "content": "import os\nfrom dataclasses import dataclass\nfrom typing import Any\n\nNAME_FIELD = \"Name\"\nMODIFIED_FIELD = \"LastModifiedDate\"\nID_FIELD = \"Id\"\nACCOUNT_OBJECT_TYPE = \"Account\"\nUSER_OBJECT_TYPE = \"User\"\n\n\n@dataclass\nclass SalesforceObject:\n id: str\n type: str\n data: dict[str, Any]\n\n def to_dict(self) -> dict[str, Any]:\n return {\n \"ID\": self.id,\n \"Type\": self.type,\n \"Data\": self.data,\n }\n\n @classmethod\n def from_dict(cls, data: dict[str, Any]) -> \"SalesforceObject\":\n return cls(\n id=data[ID_FIELD],\n type=data[\"Type\"],\n data=data,\n )\n\n\n# This defines the base path for all data files relative to this file\n# AKA BE CAREFUL WHEN MOVING THIS FILE\nBASE_DATA_PATH = os.path.join(os.path.dirname(__file__), \"data\")\n\n\ndef get_sqlite_db_path(directory: str) -> str:\n \"\"\"Get the path to the sqlite db file.\"\"\"\n return os.path.join(directory, \"salesforce_db.sqlite\")\n\n\ndef remove_sqlite_db_files(db_path: str) -> None:\n \"\"\"Remove SQLite database and all associated files (WAL, SHM).\n\n SQLite in WAL mode creates additional files:\n - .sqlite-wal: Write-ahead log\n - .sqlite-shm: Shared memory file\n\n If these files become stale (e.g., after a crash), they can cause\n 'disk I/O error' when trying to open the database. This function\n ensures all related files are removed.\n \"\"\"\n files_to_remove = [\n db_path,\n f\"{db_path}-wal\",\n f\"{db_path}-shm\",\n ]\n for file_path in files_to_remove:\n if os.path.exists(file_path):\n os.remove(file_path)\n\n\n# NOTE: only used with shelves, deprecated at this point\ndef get_object_type_path(object_type: str) -> str:\n \"\"\"Get the directory path for a specific object type.\"\"\"\n type_dir = os.path.join(BASE_DATA_PATH, object_type)\n os.makedirs(type_dir, exist_ok=True)\n return type_dir\n\n\n_CHECKSUM_CHARS = \"ABCDEFGHIJKLMNOPQRSTUVWXYZ012345\"\n_LOOKUP = {format(i, \"05b\"): _CHECKSUM_CHARS[i] for i in range(32)}\n\n\ndef validate_salesforce_id(salesforce_id: str) -> bool:\n \"\"\"Validate the checksum portion of an 18-character Salesforce ID.\n\n Args:\n salesforce_id: An 18-character Salesforce ID\n\n Returns:\n bool: True if the checksum is valid, False otherwise\n \"\"\"\n if len(salesforce_id) != 18:\n return False\n\n chunks = [salesforce_id[0:5], salesforce_id[5:10], salesforce_id[10:15]]\n\n checksum = salesforce_id[15:18]\n calculated_checksum = \"\"\n\n for chunk in chunks:\n result_string = \"\".join(\n \"1\" if char.isupper() else \"0\" for char in reversed(chunk)\n )\n calculated_checksum += _LOOKUP[result_string]\n\n return checksum == calculated_checksum\n" }, { "path": "backend/onyx/connectors/sharepoint/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/sharepoint/connector.py", "content": "import base64\nimport copy\nimport fnmatch\nimport html\nimport io\nimport os\nimport re\nimport time\nfrom collections import deque\nfrom collections.abc import Generator\nfrom collections.abc import Iterable\nfrom datetime import datetime\nfrom datetime import timezone\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\nfrom urllib.parse import quote\nfrom urllib.parse import unquote\nfrom urllib.parse import urlsplit\n\nimport msal # type: ignore[import-untyped]\nimport requests\nfrom cryptography.hazmat.primitives import hashes\nfrom cryptography.hazmat.primitives import serialization\nfrom cryptography.hazmat.primitives.serialization import pkcs12\nfrom office365.graph_client import GraphClient # type: ignore[import-untyped]\nfrom office365.onedrive.driveitems.driveItem import DriveItem # type: ignore[import-untyped]\nfrom office365.onedrive.sites.site import Site # type: ignore[import-untyped]\nfrom office365.onedrive.sites.sites_with_root import SitesWithRoot # type: ignore[import-untyped]\nfrom office365.runtime.auth.token_response import TokenResponse # type: ignore[import-untyped]\nfrom office365.runtime.client_request import ClientRequestException # type: ignore\nfrom office365.runtime.paths.resource_path import ResourcePath # type: ignore[import-untyped]\nfrom office365.runtime.queries.client_query import ClientQuery # type: ignore[import-untyped]\nfrom office365.sharepoint.client_context import ClientContext # type: ignore[import-untyped]\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom requests.exceptions import HTTPError\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import REQUEST_TIMEOUT_SECONDS\nfrom onyx.configs.app_configs import SHAREPOINT_CONNECTOR_SIZE_THRESHOLD\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import IndexingHeartbeatInterface\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.microsoft_graph_env import resolve_microsoft_environment\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import EntityFailure\nfrom onyx.connectors.models import ExternalAccess\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.sharepoint.connector_utils import get_sharepoint_external_access\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.file_processing.extract_file_text import extract_text_and_images\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.file_types import OnyxMimeTypes\nfrom onyx.file_processing.image_utils import store_image_and_create_section\nfrom onyx.utils.b64 import get_image_type_from_bytes\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\nSLIM_BATCH_SIZE = 1000\n_EPOCH = datetime.fromtimestamp(0, tz=timezone.utc)\n\n\nSHARED_DOCUMENTS_MAP = {\n \"Documents\": \"Shared Documents\",\n \"Dokumente\": \"Freigegebene Dokumente\",\n \"Documentos\": \"Documentos compartidos\",\n}\nSHARED_DOCUMENTS_MAP_REVERSE = {v: k for k, v in SHARED_DOCUMENTS_MAP.items()}\n\nASPX_EXTENSION = \".aspx\"\n\n\ndef _is_site_excluded(site_url: str, excluded_site_patterns: list[str]) -> bool:\n \"\"\"Check if a site URL matches any of the exclusion glob patterns.\"\"\"\n for pattern in excluded_site_patterns:\n if fnmatch.fnmatch(site_url, pattern) or fnmatch.fnmatch(\n site_url.rstrip(\"/\"), pattern.rstrip(\"/\")\n ):\n return True\n return False\n\n\ndef _is_path_excluded(item_path: str, excluded_path_patterns: list[str]) -> bool:\n \"\"\"Check if a drive item path matches any of the exclusion glob patterns.\n\n item_path is the relative path within a drive, e.g. \"Engineering/API/report.docx\".\n Matches are attempted against the full path and the filename alone so that\n patterns like \"*.tmp\" match files at any depth.\n \"\"\"\n filename = item_path.rsplit(\"/\", 1)[-1] if \"/\" in item_path else item_path\n for pattern in excluded_path_patterns:\n if fnmatch.fnmatch(item_path, pattern) or fnmatch.fnmatch(filename, pattern):\n return True\n return False\n\n\ndef _build_item_relative_path(parent_reference_path: str | None, item_name: str) -> str:\n \"\"\"Build the relative path of a drive item from its parentReference.path and name.\n\n Example: parentReference.path=\"/drives/abc/root:/Eng/API\", name=\"report.docx\"\n => \"Eng/API/report.docx\"\n \"\"\"\n if parent_reference_path and \"root:/\" in parent_reference_path:\n folder = unquote(parent_reference_path.split(\"root:/\", 1)[1])\n if folder:\n return f\"{folder}/{item_name}\"\n return item_name\n\n\nDEFAULT_AUTHORITY_HOST = \"https://login.microsoftonline.com\"\nDEFAULT_GRAPH_API_HOST = \"https://graph.microsoft.com\"\nDEFAULT_SHAREPOINT_DOMAIN_SUFFIX = \"sharepoint.com\"\n\nGRAPH_API_BASE = f\"{DEFAULT_GRAPH_API_HOST}/v1.0\"\nGRAPH_API_MAX_RETRIES = 5\nGRAPH_API_RETRYABLE_STATUSES = frozenset({429, 500, 502, 503, 504})\n\n\nclass DriveItemData(BaseModel):\n \"\"\"Lightweight representation of a Graph API drive item, parsed from JSON.\n\n Replaces the SDK DriveItem for fetching/listing so that we can paginate\n lazily through the Graph API without materialising every item in memory.\n \"\"\"\n\n id: str\n name: str\n web_url: str\n size: int | None = None\n mime_type: str | None = None\n download_url: str | None = None\n last_modified_datetime: datetime | None = None\n last_modified_by_display_name: str | None = None\n last_modified_by_email: str | None = None\n parent_reference_path: str | None = None\n drive_id: str | None = None\n\n @classmethod\n def from_graph_json(cls, item: dict[str, Any]) -> \"DriveItemData\":\n last_mod_raw = item.get(\"lastModifiedDateTime\")\n last_mod: datetime | None = None\n if isinstance(last_mod_raw, str):\n last_mod = datetime.fromisoformat(last_mod_raw.replace(\"Z\", \"+00:00\"))\n\n last_modified_by = item.get(\"lastModifiedBy\", {}).get(\"user\", {})\n parent_ref = item.get(\"parentReference\", {})\n\n return cls(\n id=item[\"id\"],\n name=item.get(\"name\", \"\"),\n web_url=item.get(\"webUrl\", \"\"),\n size=item.get(\"size\"),\n mime_type=item.get(\"file\", {}).get(\"mimeType\"),\n download_url=item.get(\"@microsoft.graph.downloadUrl\"),\n last_modified_datetime=last_mod,\n last_modified_by_display_name=last_modified_by.get(\"displayName\"),\n last_modified_by_email=(\n last_modified_by.get(\"email\")\n or last_modified_by.get(\"userPrincipalName\")\n ),\n parent_reference_path=parent_ref.get(\"path\"),\n drive_id=parent_ref.get(\"driveId\"),\n )\n\n def to_sdk_driveitem(self, graph_client: GraphClient) -> DriveItem:\n \"\"\"Construct a lazy SDK DriveItem for permission lookups.\"\"\"\n if not self.drive_id:\n raise ValueError(\"drive_id is required to construct SDK DriveItem\")\n path = ResourcePath(\n self.id,\n ResourcePath(\"items\", ResourcePath(self.drive_id, ResourcePath(\"drives\"))),\n )\n item = DriveItem(graph_client, path)\n item.set_property(\"id\", self.id)\n return item\n\n\n# The office365 library's ClientContext caches the access token from its\n# first request and never re-invokes the token callback. Microsoft access\n# tokens live ~60-75 minutes, so we recreate the cached ClientContext every\n# 30 minutes to let MSAL transparently handle token refresh.\n_REST_CTX_MAX_AGE_S = 30 * 60\n\n\nclass SiteDescriptor(BaseModel):\n \"\"\"Data class for storing SharePoint site information.\n\n Args:\n url: The base site URL (e.g. https://danswerai.sharepoint.com/sites/sharepoint-tests\n or https://danswerai.sharepoint.com/teams/team-name)\n drive_name: The name of the drive to access (e.g. \"Shared Documents\", \"Other Library\")\n If None, all drives will be accessed.\n folder_path: The folder path within the drive to access (e.g. \"test/nested with spaces\")\n If None, all folders will be accessed.\n \"\"\"\n\n url: str\n drive_name: str | None\n folder_path: str | None\n\n\nclass CertificateData(BaseModel):\n \"\"\"Data class for storing certificate information loaded from PFX file.\"\"\"\n\n private_key: bytes\n thumbprint: str\n\n\ndef _site_page_in_time_window(\n page: dict[str, Any],\n start: datetime | None,\n end: datetime | None,\n) -> bool:\n \"\"\"Return True if the page's lastModifiedDateTime falls within [start, end].\"\"\"\n if start is None and end is None:\n return True\n raw = page.get(\"lastModifiedDateTime\")\n if not raw:\n return True\n if not isinstance(raw, str):\n raise ValueError(f\"lastModifiedDateTime is not a string: {raw}\")\n last_modified = datetime.fromisoformat(raw.replace(\"Z\", \"+00:00\"))\n return (start is None or last_modified >= start) and (\n end is None or last_modified <= end\n )\n\n\ndef sleep_and_retry(\n query_obj: ClientQuery, method_name: str, max_retries: int = 3\n) -> Any:\n \"\"\"\n Execute a SharePoint query with retry logic for rate limiting.\n \"\"\"\n for attempt in range(max_retries + 1):\n try:\n return query_obj.execute_query()\n except ClientRequestException as e:\n status = e.response.status_code if e.response is not None else None\n\n # 429 / 503 — rate limit or transient error. Back off and retry.\n if status in (429, 503) and attempt < max_retries:\n logger.warning(\n f\"Rate limit exceeded on {method_name}, attempt {attempt + 1}/{max_retries + 1}, sleeping and retrying\"\n )\n retry_after = e.response.headers.get(\"Retry-After\")\n if retry_after:\n sleep_time = int(retry_after)\n else:\n # Exponential backoff: 2^attempt * 5 seconds\n sleep_time = min(30, (2**attempt) * 5)\n\n logger.info(f\"Sleeping for {sleep_time} seconds before retry\")\n time.sleep(sleep_time)\n continue\n\n # Non-retryable error or retries exhausted — log details and raise.\n if e.response is not None:\n logger.error(\n f\"SharePoint request failed for {method_name}: status={status}, \"\n )\n raise e\n\n\nclass SharepointConnectorCheckpoint(ConnectorCheckpoint):\n cached_site_descriptors: deque[SiteDescriptor] | None = None\n current_site_descriptor: SiteDescriptor | None = None\n\n cached_drive_names: deque[str] | None = None\n current_drive_name: str | None = None\n # Drive's web_url from the API - used as raw_node_id for DRIVE hierarchy nodes\n current_drive_web_url: str | None = None\n # Resolved drive ID — avoids re-resolving on checkpoint resume\n current_drive_id: str | None = None\n # Next delta API page URL for per-page checkpointing within a drive.\n # When set, Phase 3b fetches one page at a time so progress is persisted\n # between pages. None means BFS path or no active delta traversal.\n current_drive_delta_next_link: str | None = None\n\n process_site_pages: bool = False\n\n # Track yielded hierarchy nodes by their raw_node_id (URLs) to avoid duplicates\n seen_hierarchy_node_raw_ids: set[str] = Field(default_factory=set)\n\n # Track yielded document IDs to avoid processing the same document twice.\n # The Microsoft Graph delta API can return the same item on multiple pages.\n seen_document_ids: set[str] = Field(default_factory=set)\n\n\nclass SharepointAuthMethod(Enum):\n CLIENT_SECRET = \"client_secret\"\n CERTIFICATE = \"certificate\"\n\n\nclass SizeCapExceeded(Exception):\n \"\"\"Exception raised when the size cap is exceeded.\"\"\"\n\n\ndef _log_and_raise_for_status(response: requests.Response) -> None:\n \"\"\"Log the response text and raise for status.\"\"\"\n try:\n response.raise_for_status()\n except Exception:\n logger.error(f\"HTTP request failed: {response.text}\")\n raise\n\n\nGRAPH_INVALID_REQUEST_CODE = \"invalidRequest\"\n\n\ndef _is_graph_invalid_request(response: requests.Response) -> bool:\n \"\"\"Return True if the response body is the generic Graph API\n ``{\"error\": {\"code\": \"invalidRequest\", \"message\": \"Invalid request\"}}``\n shape. This particular error has no actionable inner error code and is\n returned by the site-pages endpoint when a page has a corrupt canvas layout\n (e.g. duplicate web-part IDs — see SharePoint/sp-dev-docs#8822).\"\"\"\n try:\n body = response.json()\n except Exception:\n return False\n error = body.get(\"error\", {})\n return error.get(\"code\") == GRAPH_INVALID_REQUEST_CODE\n\n\ndef load_certificate_from_pfx(pfx_data: bytes, password: str) -> CertificateData | None:\n \"\"\"Load certificate from .pfx file for MSAL authentication\"\"\"\n try:\n # Load the certificate and private key\n private_key, certificate, additional_certificates = (\n pkcs12.load_key_and_certificates(pfx_data, password.encode(\"utf-8\"))\n )\n\n # Validate that certificate and private key are not None\n if certificate is None or private_key is None:\n raise ValueError(\"Certificate or private key is None\")\n\n # Convert to PEM format that MSAL expects\n key_pem = private_key.private_bytes(\n encoding=serialization.Encoding.PEM,\n format=serialization.PrivateFormat.PKCS8,\n encryption_algorithm=serialization.NoEncryption(),\n )\n\n return CertificateData(\n private_key=key_pem,\n thumbprint=certificate.fingerprint(hashes.SHA1()).hex(),\n )\n except Exception as e:\n logger.error(f\"Error loading certificate: {e}\")\n return None\n\n\ndef acquire_token_for_rest(\n msal_app: msal.ConfidentialClientApplication,\n sp_tenant_domain: str,\n sharepoint_domain_suffix: str,\n) -> TokenResponse:\n token = msal_app.acquire_token_for_client(\n scopes=[f\"https://{sp_tenant_domain}.{sharepoint_domain_suffix}/.default\"]\n )\n return TokenResponse.from_json(token)\n\n\ndef _create_document_failure(\n driveitem: DriveItemData,\n error_message: str,\n exception: Exception | None = None,\n) -> ConnectorFailure:\n \"\"\"Helper method to create a ConnectorFailure for document processing errors.\"\"\"\n return ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=driveitem.id or \"unknown\",\n document_link=driveitem.web_url,\n ),\n failure_message=f\"SharePoint document '{driveitem.name or 'unknown'}': {error_message}\",\n exception=exception,\n )\n\n\ndef _create_entity_failure(\n entity_id: str,\n error_message: str,\n time_range: tuple[datetime, datetime] | None = None,\n exception: Exception | None = None,\n) -> ConnectorFailure:\n \"\"\"Helper method to create a ConnectorFailure for entity-level errors.\"\"\"\n return ConnectorFailure(\n failed_entity=EntityFailure(\n entity_id=entity_id,\n missed_time_range=time_range,\n ),\n failure_message=f\"SharePoint entity '{entity_id}': {error_message}\",\n exception=exception,\n )\n\n\ndef _probe_remote_size(url: str, timeout: int) -> int | None:\n \"\"\"Determine remote size using HEAD or a range GET probe. Returns None if unknown.\"\"\"\n try:\n head_resp = requests.head(url, timeout=timeout, allow_redirects=True)\n _log_and_raise_for_status(head_resp)\n cl = head_resp.headers.get(\"Content-Length\")\n if cl and cl.isdigit():\n return int(cl)\n except requests.RequestException:\n pass\n\n # Fallback: Range request for first byte to read total from Content-Range\n try:\n with requests.get(\n url,\n headers={\"Range\": \"bytes=0-0\"},\n timeout=timeout,\n stream=True,\n ) as range_resp:\n _log_and_raise_for_status(range_resp)\n cr = range_resp.headers.get(\"Content-Range\") # e.g., \"bytes 0-0/12345\"\n if cr and \"/\" in cr:\n total = cr.split(\"/\")[-1]\n if total.isdigit():\n return int(total)\n except requests.RequestException:\n pass\n\n # If both HEAD and a range GET failed to reveal a size, signal unknown size.\n # Callers should treat None as \"size unavailable\" and proceed with a safe\n # streaming path that enforces a hard cap to avoid excessive memory usage.\n return None\n\n\ndef _download_with_cap(url: str, timeout: int, cap: int) -> bytes:\n \"\"\"Stream download content with an upper bound on bytes read.\n\n Behavior:\n - Checks `Content-Length` first and aborts early if it exceeds `cap`.\n - Otherwise streams the body in chunks and stops once `cap` is surpassed.\n - Raises `SizeCapExceeded` when the cap would be exceeded.\n - Returns the full bytes if the content fits within `cap`.\n \"\"\"\n with requests.get(url, stream=True, timeout=timeout) as resp:\n _log_and_raise_for_status(resp)\n\n # If the server provides Content-Length, prefer an early decision.\n cl_header = resp.headers.get(\"Content-Length\")\n if cl_header and cl_header.isdigit():\n content_len = int(cl_header)\n if content_len > cap:\n logger.warning(\n f\"Content-Length {content_len} exceeds cap {cap}; skipping download.\"\n )\n raise SizeCapExceeded(\"pre_download\")\n\n buf = io.BytesIO()\n # Stream in 64KB chunks; adjust if needed for slower networks.\n for chunk in resp.iter_content(64 * 1024):\n if not chunk:\n continue\n buf.write(chunk)\n if buf.tell() > cap:\n # Avoid keeping a large partial buffer; close and signal caller to skip.\n logger.warning(\n f\"Streaming download exceeded cap {cap} bytes; aborting early.\"\n )\n raise SizeCapExceeded(\"during_download\")\n\n return buf.getvalue()\n\n\ndef _download_via_graph_api(\n access_token: str,\n drive_id: str,\n item_id: str,\n bytes_allowed: int,\n graph_api_base: str,\n) -> bytes:\n \"\"\"Download a drive item via the Graph API /content endpoint with a byte cap.\n\n Raises SizeCapExceeded if the cap is exceeded.\n \"\"\"\n url = f\"{graph_api_base}/drives/{drive_id}/items/{item_id}/content\"\n headers = {\"Authorization\": f\"Bearer {access_token}\"}\n with requests.get(\n url, headers=headers, stream=True, timeout=REQUEST_TIMEOUT_SECONDS\n ) as resp:\n _log_and_raise_for_status(resp)\n buf = io.BytesIO()\n for chunk in resp.iter_content(64 * 1024):\n if not chunk:\n continue\n buf.write(chunk)\n if buf.tell() > bytes_allowed:\n raise SizeCapExceeded(\"during_graph_api_download\")\n return buf.getvalue()\n\n\ndef _convert_driveitem_to_document_with_permissions(\n driveitem: DriveItemData,\n drive_name: str,\n ctx: ClientContext | None,\n graph_client: GraphClient,\n graph_api_base: str,\n include_permissions: bool = False,\n parent_hierarchy_raw_node_id: str | None = None,\n access_token: str | None = None,\n treat_sharing_link_as_public: bool = False,\n) -> Document | ConnectorFailure | None:\n\n if not driveitem.name or not driveitem.id:\n raise ValueError(\"DriveItem name/id is required\")\n\n if include_permissions and ctx is None:\n raise ValueError(\"ClientContext is required for permissions\")\n\n mime_type = driveitem.mime_type\n if not mime_type or mime_type in OnyxMimeTypes.EXCLUDED_IMAGE_TYPES:\n logger.debug(\n f\"Skipping malformed or excluded mime type {mime_type} for {driveitem.name}\"\n )\n return None\n\n file_size = driveitem.size\n download_url = driveitem.download_url\n\n if file_size is None and download_url:\n file_size = _probe_remote_size(download_url, REQUEST_TIMEOUT_SECONDS)\n\n if file_size is not None and file_size > SHAREPOINT_CONNECTOR_SIZE_THRESHOLD:\n logger.warning(\n f\"Skipping '{driveitem.name}' over size threshold ({file_size} > {SHAREPOINT_CONNECTOR_SIZE_THRESHOLD} bytes).\"\n )\n return None\n\n # Prefer downloadUrl streaming with size cap\n content_bytes: bytes | None = None\n if download_url:\n try:\n content_bytes = _download_with_cap(\n download_url,\n REQUEST_TIMEOUT_SECONDS,\n SHAREPOINT_CONNECTOR_SIZE_THRESHOLD,\n )\n except SizeCapExceeded as e:\n logger.warning(f\"Skipping '{driveitem.name}' exceeded size cap: {str(e)}\")\n return None\n except requests.RequestException as e:\n status = e.response.status_code if e.response is not None else -1\n logger.warning(\n f\"Failed to download via downloadUrl for '{driveitem.name}' (status={status}); falling back to Graph API.\"\n )\n\n # Fallback: download via Graph API /content endpoint\n if content_bytes is None and access_token and driveitem.drive_id:\n try:\n content_bytes = _download_via_graph_api(\n access_token,\n driveitem.drive_id,\n driveitem.id,\n SHAREPOINT_CONNECTOR_SIZE_THRESHOLD,\n graph_api_base=graph_api_base,\n )\n except SizeCapExceeded:\n logger.warning(\n f\"Skipping '{driveitem.name}' exceeded size cap during Graph API download.\"\n )\n return None\n except Exception as e:\n logger.warning(\n f\"Failed to download via Graph API for '{driveitem.name}': {e}\"\n )\n return _create_document_failure(\n driveitem, f\"Failed to download via graph api: {e}\", e\n )\n\n sections: list[TextSection | ImageSection] = []\n file_ext = get_file_ext(driveitem.name)\n\n if not content_bytes:\n logger.warning(\n f\"Zero-length content for '{driveitem.name}'. Skipping text/image extraction.\"\n )\n elif file_ext in OnyxFileExtensions.IMAGE_EXTENSIONS:\n image_section, _ = store_image_and_create_section(\n image_data=content_bytes,\n file_id=driveitem.id,\n display_name=driveitem.name,\n file_origin=FileOrigin.CONNECTOR,\n )\n image_section.link = driveitem.web_url\n sections.append(image_section)\n else:\n\n def _store_embedded_image(img_data: bytes, img_name: str) -> None:\n try:\n img_mime = get_image_type_from_bytes(img_data)\n except ValueError:\n logger.debug(\n \"Skipping embedded image with unknown format for %s\",\n driveitem.name,\n )\n return\n\n if img_mime in OnyxMimeTypes.EXCLUDED_IMAGE_TYPES:\n logger.debug(\n \"Skipping embedded image of excluded type %s for %s\",\n img_mime,\n driveitem.name,\n )\n return\n\n image_section, _ = store_image_and_create_section(\n image_data=img_data,\n file_id=f\"{driveitem.id}_img_{len(sections)}\",\n display_name=img_name or f\"{driveitem.name} - image {len(sections)}\",\n file_origin=FileOrigin.CONNECTOR,\n )\n image_section.link = driveitem.web_url\n sections.append(image_section)\n\n extraction_result = extract_text_and_images(\n file=io.BytesIO(content_bytes),\n file_name=driveitem.name,\n image_callback=_store_embedded_image,\n )\n if extraction_result.text_content:\n sections.append(\n TextSection(link=driveitem.web_url, text=extraction_result.text_content)\n )\n\n if include_permissions and ctx is not None:\n logger.info(f\"Getting external access for {driveitem.name}\")\n sdk_item = driveitem.to_sdk_driveitem(graph_client)\n external_access = get_sharepoint_external_access(\n ctx=ctx,\n graph_client=graph_client,\n drive_item=sdk_item,\n drive_name=drive_name,\n add_prefix=True,\n treat_sharing_link_as_public=treat_sharing_link_as_public,\n )\n else:\n external_access = ExternalAccess.empty()\n\n doc = Document(\n id=driveitem.id,\n sections=sections,\n source=DocumentSource.SHAREPOINT,\n semantic_identifier=driveitem.name,\n external_access=external_access,\n doc_updated_at=(\n driveitem.last_modified_datetime.replace(tzinfo=timezone.utc)\n if driveitem.last_modified_datetime\n else None\n ),\n primary_owners=[\n BasicExpertInfo(\n display_name=driveitem.last_modified_by_display_name or \"\",\n email=driveitem.last_modified_by_email or \"\",\n )\n ],\n metadata={\"drive\": drive_name},\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n )\n return doc\n\n\ndef _convert_sitepage_to_document(\n site_page: dict[str, Any],\n site_name: str | None,\n ctx: ClientContext | None,\n graph_client: GraphClient,\n include_permissions: bool = False,\n parent_hierarchy_raw_node_id: str | None = None,\n treat_sharing_link_as_public: bool = False,\n) -> Document:\n \"\"\"Convert a SharePoint site page to a Document object.\"\"\"\n # Extract text content from the site page\n page_text = \"\"\n # Get title and description\n title = cast(str, site_page.get(\"title\", \"\"))\n description = cast(str, site_page.get(\"description\", \"\"))\n\n # Build the text content\n if title:\n page_text += f\"# {title}\\n\\n\"\n if description:\n page_text += f\"{description}\\n\\n\"\n\n # Extract content from canvas layout if available\n canvas_layout = site_page.get(\"canvasLayout\", {})\n if canvas_layout:\n horizontal_sections = canvas_layout.get(\"horizontalSections\", [])\n for section in horizontal_sections:\n columns = section.get(\"columns\", [])\n for column in columns:\n webparts = column.get(\"webparts\", [])\n for webpart in webparts:\n # Extract text from different types of webparts\n webpart_type = webpart.get(\"@odata.type\", \"\")\n\n # Extract text from text webparts\n if webpart_type == \"#microsoft.graph.textWebPart\":\n inner_html = webpart.get(\"innerHtml\", \"\")\n if inner_html:\n # Basic HTML to text conversion\n # Remove HTML tags but preserve some structure\n text_content = re.sub(r\"\", \"\\n\", inner_html)\n text_content = re.sub(r\"
  • \", \"• \", text_content)\n text_content = re.sub(r\"
    • \", \"\\n\", text_content)\n text_content = re.sub(\n r\"]*>\", \"\\n## \", text_content\n )\n text_content = re.sub(r\"\", \"\\n\", text_content)\n text_content = re.sub(r\"]*>\", \"\\n\", text_content)\n text_content = re.sub(r\"

    \", \"\\n\", text_content)\n text_content = re.sub(r\"<[^>]+>\", \"\", text_content)\n # Decode HTML entities\n text_content = html.unescape(text_content)\n # Clean up extra whitespace\n text_content = re.sub(\n r\"\\n\\s*\\n\", \"\\n\\n\", text_content\n ).strip()\n if text_content:\n page_text += f\"{text_content}\\n\\n\"\n\n # Extract text from standard webparts\n elif webpart_type == \"#microsoft.graph.standardWebPart\":\n data = webpart.get(\"data\", {})\n\n # Extract from serverProcessedContent\n server_content = data.get(\"serverProcessedContent\", {})\n searchable_texts = server_content.get(\n \"searchablePlainTexts\", []\n )\n\n for text_item in searchable_texts:\n if isinstance(text_item, dict):\n key = text_item.get(\"key\", \"\")\n value = text_item.get(\"value\", \"\")\n if value:\n # Add context based on key\n if key == \"title\":\n page_text += f\"## {value}\\n\\n\"\n else:\n page_text += f\"{value}\\n\\n\"\n\n # Extract description if available\n description = data.get(\"description\", \"\")\n if description:\n page_text += f\"{description}\\n\\n\"\n\n # Extract title if available\n webpart_title = data.get(\"title\", \"\")\n if webpart_title and webpart_title != description:\n page_text += f\"## {webpart_title}\\n\\n\"\n\n page_text = page_text.strip()\n\n # If no content extracted, use the title as fallback\n if not page_text and title:\n page_text = title\n\n # Parse creation and modification info\n created_datetime = site_page.get(\"createdDateTime\")\n if created_datetime:\n if isinstance(created_datetime, str):\n created_datetime = datetime.fromisoformat(\n created_datetime.replace(\"Z\", \"+00:00\")\n )\n elif not created_datetime.tzinfo:\n created_datetime = created_datetime.replace(tzinfo=timezone.utc)\n\n last_modified_datetime = site_page.get(\"lastModifiedDateTime\")\n if last_modified_datetime:\n if isinstance(last_modified_datetime, str):\n last_modified_datetime = datetime.fromisoformat(\n last_modified_datetime.replace(\"Z\", \"+00:00\")\n )\n elif not last_modified_datetime.tzinfo:\n last_modified_datetime = last_modified_datetime.replace(tzinfo=timezone.utc)\n\n # Extract owner information\n primary_owners = []\n created_by = site_page.get(\"createdBy\", {}).get(\"user\", {})\n if created_by.get(\"displayName\"):\n primary_owners.append(\n BasicExpertInfo(\n display_name=created_by.get(\"displayName\"),\n email=created_by.get(\"email\", \"\"),\n )\n )\n\n web_url = site_page[\"webUrl\"]\n semantic_identifier = cast(str, site_page.get(\"name\", title))\n if semantic_identifier.endswith(ASPX_EXTENSION):\n semantic_identifier = semantic_identifier[: -len(ASPX_EXTENSION)]\n\n if include_permissions:\n external_access = get_sharepoint_external_access(\n ctx=ctx,\n graph_client=graph_client,\n site_page=site_page,\n add_prefix=True,\n treat_sharing_link_as_public=treat_sharing_link_as_public,\n )\n else:\n external_access = ExternalAccess.empty()\n\n doc = Document(\n id=site_page[\"id\"],\n sections=[TextSection(link=web_url, text=page_text)],\n source=DocumentSource.SHAREPOINT,\n external_access=external_access,\n semantic_identifier=semantic_identifier,\n doc_updated_at=last_modified_datetime or created_datetime,\n primary_owners=primary_owners,\n metadata=(\n {\n \"site\": site_name,\n }\n if site_name\n else {}\n ),\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n )\n return doc\n\n\ndef _convert_driveitem_to_slim_document(\n driveitem: DriveItemData,\n drive_name: str,\n ctx: ClientContext,\n graph_client: GraphClient,\n parent_hierarchy_raw_node_id: str | None = None,\n treat_sharing_link_as_public: bool = False,\n) -> SlimDocument:\n if driveitem.id is None:\n raise ValueError(\"DriveItem ID is required\")\n\n sdk_item = driveitem.to_sdk_driveitem(graph_client)\n external_access = get_sharepoint_external_access(\n ctx=ctx,\n graph_client=graph_client,\n drive_item=sdk_item,\n drive_name=drive_name,\n treat_sharing_link_as_public=treat_sharing_link_as_public,\n )\n\n return SlimDocument(\n id=driveitem.id,\n external_access=external_access,\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n )\n\n\ndef _convert_sitepage_to_slim_document(\n site_page: dict[str, Any],\n ctx: ClientContext | None,\n graph_client: GraphClient,\n parent_hierarchy_raw_node_id: str | None = None,\n treat_sharing_link_as_public: bool = False,\n) -> SlimDocument:\n \"\"\"Convert a SharePoint site page to a SlimDocument object.\"\"\"\n if site_page.get(\"id\") is None:\n raise ValueError(\"Site page ID is required\")\n\n external_access = get_sharepoint_external_access(\n ctx=ctx,\n graph_client=graph_client,\n site_page=site_page,\n treat_sharing_link_as_public=treat_sharing_link_as_public,\n )\n id = site_page.get(\"id\")\n if id is None:\n raise ValueError(\"Site page ID is required\")\n return SlimDocument(\n id=id,\n external_access=external_access,\n parent_hierarchy_raw_node_id=parent_hierarchy_raw_node_id,\n )\n\n\nclass SharepointConnector(\n SlimConnectorWithPermSync,\n CheckpointedConnectorWithPermSync[SharepointConnectorCheckpoint],\n):\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n sites: list[str] = [],\n excluded_sites: list[str] = [],\n excluded_paths: list[str] = [],\n include_site_pages: bool = True,\n include_site_documents: bool = True,\n treat_sharing_link_as_public: bool = False,\n authority_host: str = DEFAULT_AUTHORITY_HOST,\n graph_api_host: str = DEFAULT_GRAPH_API_HOST,\n sharepoint_domain_suffix: str = DEFAULT_SHAREPOINT_DOMAIN_SUFFIX,\n ) -> None:\n self.batch_size = batch_size\n self.sites = list(sites)\n self.excluded_sites = [s for p in excluded_sites if (s := p.strip())]\n self.excluded_paths = [s for p in excluded_paths if (s := p.strip())]\n self.treat_sharing_link_as_public = treat_sharing_link_as_public\n self.site_descriptors: list[SiteDescriptor] = self._extract_site_and_drive_info(\n sites\n )\n self._graph_client: GraphClient | None = None\n self.msal_app: msal.ConfidentialClientApplication | None = None\n self.include_site_pages = include_site_pages\n self.include_site_documents = include_site_documents\n self.sp_tenant_domain: str | None = None\n self._credential_json: dict[str, Any] | None = None\n self._cached_rest_ctx: ClientContext | None = None\n self._cached_rest_ctx_url: str | None = None\n self._cached_rest_ctx_created_at: float = 0.0\n\n resolved_env = resolve_microsoft_environment(graph_api_host, authority_host)\n self._azure_environment = resolved_env.environment\n self.authority_host = resolved_env.authority_host\n self.graph_api_host = resolved_env.graph_host\n self.graph_api_base = f\"{self.graph_api_host}/v1.0\"\n self.sharepoint_domain_suffix = resolved_env.sharepoint_domain_suffix\n if sharepoint_domain_suffix != resolved_env.sharepoint_domain_suffix:\n logger.warning(\n f\"Configured sharepoint_domain_suffix '{sharepoint_domain_suffix}' \"\n f\"differs from the expected suffix '{resolved_env.sharepoint_domain_suffix}' \"\n f\"for the {resolved_env.environment} environment. \"\n f\"Using '{resolved_env.sharepoint_domain_suffix}'.\"\n )\n\n def validate_connector_settings(self) -> None:\n # Validate that at least one content type is enabled\n if not self.include_site_documents and not self.include_site_pages:\n raise ConnectorValidationError(\n \"At least one content type must be enabled. \"\n \"Please check either 'Include Site Documents' or 'Include Site Pages' (or both).\"\n )\n\n # Ensure sites are sharepoint urls\n for site_url in self.sites:\n if not site_url.startswith(\"https://\") or not (\n \"/sites/\" in site_url or \"/teams/\" in site_url\n ):\n raise ConnectorValidationError(\n \"Site URLs must be full Sharepoint URLs (e.g. https://your-tenant.sharepoint.com/sites/your-site or https://your-tenant.sharepoint.com/teams/your-team)\"\n )\n\n def _extract_tenant_domain_from_sites(self) -> str | None:\n \"\"\"Extract the tenant domain from configured site URLs.\n\n Site URLs look like https://{tenant}.sharepoint.com/sites/... so the\n tenant domain is the first label of the hostname.\n \"\"\"\n for site_url in self.sites:\n try:\n hostname = urlsplit(site_url.strip()).hostname\n except ValueError:\n continue\n if not hostname:\n continue\n tenant = hostname.split(\".\")[0]\n if tenant:\n return tenant\n logger.warning(f\"No tenant domain found from {len(self.sites)} sites\")\n return None\n\n def _resolve_tenant_domain_from_root_site(self) -> str:\n \"\"\"Resolve tenant domain via GET /v1.0/sites/root which only requires\n Sites.Read.All (a permission the connector already needs).\"\"\"\n root_site = self.graph_client.sites.root.get().execute_query()\n hostname = root_site.site_collection.hostname\n if not hostname:\n raise ConnectorValidationError(\n \"Could not determine tenant domain from root site\"\n )\n tenant_domain = hostname.split(\".\")[0]\n logger.info(\n \"Resolved tenant domain '%s' from root site hostname '%s'\",\n tenant_domain,\n hostname,\n )\n return tenant_domain\n\n def _resolve_tenant_domain(self) -> str:\n \"\"\"Determine the tenant domain, preferring site URLs over a Graph API\n call to avoid needing extra permissions.\"\"\"\n from_sites = self._extract_tenant_domain_from_sites()\n if from_sites:\n logger.info(\n \"Resolved tenant domain '%s' from site URLs\",\n from_sites,\n )\n return from_sites\n\n logger.info(\"No site URLs available; resolving tenant domain from root site\")\n return self._resolve_tenant_domain_from_root_site()\n\n @property\n def graph_client(self) -> GraphClient:\n if self._graph_client is None:\n raise ConnectorMissingCredentialError(\"Sharepoint\")\n\n return self._graph_client\n\n def _create_rest_client_context(self, site_url: str) -> ClientContext:\n \"\"\"Return a ClientContext for SharePoint REST API calls, with caching.\n\n The office365 library's ClientContext caches the access token from its\n first request and never re-invokes the token callback. We cache the\n context and recreate it when the site URL changes or after\n ``_REST_CTX_MAX_AGE_S``. On recreation we also call\n ``load_credentials`` to build a fresh MSAL app with an empty token\n cache, guaranteeing a brand-new token from Azure AD.\"\"\"\n elapsed = time.monotonic() - self._cached_rest_ctx_created_at\n if (\n self._cached_rest_ctx is not None\n and self._cached_rest_ctx_url == site_url\n and elapsed <= _REST_CTX_MAX_AGE_S\n ):\n return self._cached_rest_ctx\n\n if self._credential_json:\n logger.info(\n \"Rebuilding SharePoint REST client context (elapsed=%.0fs, site_changed=%s)\",\n elapsed,\n self._cached_rest_ctx_url != site_url,\n )\n self.load_credentials(self._credential_json)\n\n if not self.msal_app or not self.sp_tenant_domain:\n raise RuntimeError(\"MSAL app or tenant domain is not set\")\n\n msal_app = self.msal_app\n sp_tenant_domain = self.sp_tenant_domain\n sp_domain_suffix = self.sharepoint_domain_suffix\n self._cached_rest_ctx = ClientContext(site_url).with_access_token(\n lambda: acquire_token_for_rest(msal_app, sp_tenant_domain, sp_domain_suffix)\n )\n self._cached_rest_ctx_url = site_url\n self._cached_rest_ctx_created_at = time.monotonic()\n return self._cached_rest_ctx\n\n @staticmethod\n def _strip_share_link_tokens(path: str) -> list[str]:\n # Share links often include a token prefix like /:f:/r/ or /:x:/r/.\n segments = [segment for segment in path.split(\"/\") if segment]\n if segments and segments[0].startswith(\":\"):\n segments = segments[1:]\n if segments and segments[0] in {\"r\", \"s\", \"g\"}:\n segments = segments[1:]\n return segments\n\n @staticmethod\n def _normalize_sharepoint_url(url: str) -> tuple[str | None, list[str]]:\n try:\n parsed = urlsplit(url)\n except ValueError:\n logger.warning(f\"Sharepoint URL '{url}' could not be parsed\")\n return None, []\n\n if not parsed.scheme or not parsed.netloc:\n logger.warning(\n f\"Sharepoint URL '{url}' is not a valid absolute URL (missing scheme or host)\"\n )\n return None, []\n\n path_segments = SharepointConnector._strip_share_link_tokens(parsed.path)\n return f\"{parsed.scheme}://{parsed.netloc}\", path_segments\n\n @staticmethod\n def _extract_site_and_drive_info(site_urls: list[str]) -> list[SiteDescriptor]:\n site_data_list = []\n for url in site_urls:\n base_url, parts = SharepointConnector._normalize_sharepoint_url(url.strip())\n if base_url is None:\n continue\n\n lower_parts = [part.lower() for part in parts]\n site_type_index = None\n for site_token in (\"sites\", \"teams\"):\n if site_token in lower_parts:\n site_type_index = lower_parts.index(site_token)\n break\n\n if site_type_index is None or len(parts) <= site_type_index + 1:\n logger.warning(\n f\"Site URL '{url}' is not a valid Sharepoint URL (must contain /sites/ or /teams/)\"\n )\n continue\n\n site_path = parts[: site_type_index + 2]\n remaining_parts = parts[site_type_index + 2 :]\n site_url = f\"{base_url}/\" + \"/\".join(site_path)\n\n # Extract drive name and folder path\n if remaining_parts:\n drive_name = unquote(remaining_parts[0])\n folder_path = (\n \"/\".join(unquote(part) for part in remaining_parts[1:])\n if len(remaining_parts) > 1\n else None\n )\n else:\n drive_name = None\n folder_path = None\n\n site_data_list.append(\n SiteDescriptor(\n url=site_url,\n drive_name=drive_name,\n folder_path=folder_path,\n )\n )\n return site_data_list\n\n def _resolve_drive(\n self,\n site_descriptor: SiteDescriptor,\n drive_name: str,\n ) -> tuple[str, str | None] | None:\n \"\"\"Find the drive ID and web_url for a given drive name on a site.\n\n Returns (drive_id, drive_web_url) or None if the drive was not found.\n Raises on auth/permission errors so callers can propagate them.\n \"\"\"\n site = self.graph_client.sites.get_by_url(site_descriptor.url)\n drives = site.drives.get().execute_query()\n logger.info(f\"Found drives: {[d.name for d in drives]}\")\n\n matched = [\n d\n for d in drives\n if (d.name and d.name.lower() == drive_name.lower())\n or (\n d.name in SHARED_DOCUMENTS_MAP\n and SHARED_DOCUMENTS_MAP[d.name] == drive_name\n )\n ]\n if not matched:\n logger.warning(f\"Drive '{drive_name}' not found\")\n return None\n\n drive = matched[0]\n drive_web_url: str | None = drive.web_url\n logger.info(f\"Found drive: {drive.name} (web_url: {drive_web_url})\")\n return cast(str, drive.id), drive_web_url\n\n def _get_drive_items_for_drive_id(\n self,\n site_descriptor: SiteDescriptor,\n drive_id: str,\n start: datetime | None = None,\n end: datetime | None = None,\n ) -> Generator[DriveItemData, None, None]:\n \"\"\"Yield drive items lazily for a given drive name.\n\n Uses the delta API for whole-drive enumeration (flat, incremental via\n timestamp token) and falls back to BFS /children traversal when a\n folder_path is configured, since delta cannot scope to a subtree\n efficiently.\n\n Returns:\n A generator of DriveItemData.\n The generator paginates through the Graph API so items are never\n all held in memory at once.\n \"\"\"\n try:\n if site_descriptor.folder_path:\n yield from self._iter_drive_items_paged(\n drive_id=drive_id,\n folder_path=site_descriptor.folder_path,\n start=start,\n end=end,\n )\n else:\n yield from self._iter_drive_items_delta(\n drive_id=drive_id,\n start=start,\n end=end,\n )\n\n except Exception as e:\n err_str = str(e)\n if (\n \"403 Client Error\" in err_str\n or \"404 Client Error\" in err_str\n or \"invalid_client\" in err_str\n ):\n raise e\n\n logger.warning(f\"Failed to process site: {site_descriptor.url} - {err_str}\")\n\n def _fetch_driveitems(\n self,\n site_descriptor: SiteDescriptor,\n start: datetime | None = None,\n end: datetime | None = None,\n ) -> Generator[tuple[DriveItemData, str, str | None], None, None]:\n \"\"\"Yield drive items lazily for all drives in a site.\n\n Yields (DriveItemData, drive_name, drive_web_url) tuples one item at\n a time, paginating through the Graph API internally.\n \"\"\"\n try:\n site = self.graph_client.sites.get_by_url(site_descriptor.url)\n drives = site.drives.get().execute_query()\n logger.debug(f\"Found drives: {[d.name for d in drives]}\")\n\n if site_descriptor.drive_name:\n drives = [\n drive\n for drive in drives\n if drive.name == site_descriptor.drive_name\n or (\n drive.name in SHARED_DOCUMENTS_MAP\n and SHARED_DOCUMENTS_MAP[drive.name]\n == site_descriptor.drive_name\n )\n ]\n if not drives:\n logger.warning(f\"Drive '{site_descriptor.drive_name}' not found\")\n return\n\n for drive in drives:\n try:\n drive_name = (\n SHARED_DOCUMENTS_MAP[drive.name]\n if drive.name in SHARED_DOCUMENTS_MAP\n else cast(str, drive.name)\n )\n drive_web_url: str | None = drive.web_url\n\n if site_descriptor.folder_path:\n item_iter = self._iter_drive_items_paged(\n drive_id=cast(str, drive.id),\n folder_path=site_descriptor.folder_path,\n start=start,\n end=end,\n )\n else:\n item_iter = self._iter_drive_items_delta(\n drive_id=cast(str, drive.id),\n start=start,\n end=end,\n )\n\n for item in item_iter:\n yield item, drive_name or \"\", drive_web_url\n\n except Exception as e:\n logger.warning(f\"Failed to process drive '{drive.name}': {str(e)}\")\n\n except Exception as e:\n err_str = str(e)\n if (\n \"403 Client Error\" in err_str\n or \"404 Client Error\" in err_str\n or \"invalid_client\" in err_str\n ):\n raise e\n\n logger.warning(f\"Failed to process site: {err_str}\")\n\n def _handle_paginated_sites(\n self, sites: SitesWithRoot\n ) -> Generator[Site, None, None]:\n while sites:\n if sites.current_page:\n yield from sites.current_page\n if not sites.has_next:\n break\n sites = sites._get_next().execute_query()\n\n def _is_driveitem_excluded(self, driveitem: DriveItemData) -> bool:\n \"\"\"Check if a drive item should be excluded based on excluded_paths patterns.\"\"\"\n if not self.excluded_paths:\n return False\n relative_path = _build_item_relative_path(\n driveitem.parent_reference_path, driveitem.name\n )\n return _is_path_excluded(relative_path, self.excluded_paths)\n\n def _filter_excluded_sites(\n self, site_descriptors: list[SiteDescriptor]\n ) -> list[SiteDescriptor]:\n \"\"\"Remove sites matching any excluded_sites glob pattern.\"\"\"\n if not self.excluded_sites:\n return site_descriptors\n result = []\n for sd in site_descriptors:\n if _is_site_excluded(sd.url, self.excluded_sites):\n logger.info(f\"Excluding site by denylist: {sd.url}\")\n continue\n result.append(sd)\n return result\n\n def fetch_sites(self) -> list[SiteDescriptor]:\n sites = self.graph_client.sites.get_all_sites().execute_query()\n\n if not sites:\n raise RuntimeError(\"No sites found in the tenant\")\n\n # OneDrive personal sites should not be indexed with SharepointConnector\n site_descriptors = [\n SiteDescriptor(\n url=site.web_url or \"\",\n drive_name=None,\n folder_path=None,\n )\n for site in self._handle_paginated_sites(sites)\n if \"-my.sharepoint\" not in site.web_url\n ]\n return self._filter_excluded_sites(site_descriptors)\n\n def _fetch_site_pages(\n self,\n site_descriptor: SiteDescriptor,\n start: datetime | None = None,\n end: datetime | None = None,\n ) -> Generator[dict[str, Any], None, None]:\n \"\"\"Yield SharePoint site pages (.aspx files) one at a time.\n\n Pages are fetched via the Graph Pages API and yielded lazily as each\n API page arrives, so memory stays bounded regardless of total page count.\n Time-window filtering is applied per-item before yielding.\n \"\"\"\n site = self.graph_client.sites.get_by_url(site_descriptor.url)\n site.execute_query()\n site_id = site.id\n\n site_pages_base = (\n f\"{self.graph_api_base}/sites/{site_id}/pages/microsoft.graph.sitePage\"\n )\n page_url: str | None = site_pages_base\n params: dict[str, str] | None = {\"$expand\": \"canvasLayout\"}\n total_yielded = 0\n yielded_ids: set[str] = set()\n\n while page_url:\n try:\n data = self._graph_api_get_json(page_url, params)\n except HTTPError as e:\n if e.response is not None and e.response.status_code == 404:\n logger.warning(f\"Site page not found: {page_url}\")\n break\n if (\n e.response is not None\n and e.response.status_code == 400\n and _is_graph_invalid_request(e.response)\n ):\n logger.warning(\n f\"$expand=canvasLayout on the LIST endpoint returned 400 \"\n f\"for site {site_descriptor.url}. Falling back to \"\n f\"per-page expansion.\"\n )\n yield from self._fetch_site_pages_individually(\n site_pages_base, start, end, skip_ids=yielded_ids\n )\n return\n raise\n\n params = None # nextLink already embeds query params\n\n for page in data.get(\"value\", []):\n if not _site_page_in_time_window(page, start, end):\n continue\n total_yielded += 1\n page_id = page.get(\"id\")\n if page_id:\n yielded_ids.add(page_id)\n yield page\n\n page_url = data.get(\"@odata.nextLink\")\n\n logger.debug(f\"Yielded {total_yielded} site pages for {site_descriptor.url}\")\n\n def _fetch_site_pages_individually(\n self,\n site_pages_base: str,\n start: datetime | None = None,\n end: datetime | None = None,\n skip_ids: set[str] | None = None,\n ) -> Generator[dict[str, Any], None, None]:\n \"\"\"Fallback for _fetch_site_pages: list pages without $expand, then\n expand canvasLayout on each page individually.\n\n The Graph API's LIST endpoint can return 400 when $expand=canvasLayout\n is used and *any* page in the site has a corrupt canvas layout (e.g.\n duplicate web part IDs — see SharePoint/sp-dev-docs#8822). Since the\n LIST expansion is all-or-nothing, a single bad page poisons the entire\n response. This method works around it by fetching metadata first, then\n expanding each page individually so only the broken page loses its\n canvas content.\n\n ``skip_ids`` contains page IDs already yielded by the caller before the\n fallback was triggered, preventing duplicates.\n \"\"\"\n page_url: str | None = site_pages_base\n total_yielded = 0\n _skip_ids = skip_ids or set()\n\n while page_url:\n try:\n data = self._graph_api_get_json(page_url)\n except HTTPError as e:\n if e.response is not None and e.response.status_code == 404:\n break\n raise\n\n for page in data.get(\"value\", []):\n if not _site_page_in_time_window(page, start, end):\n continue\n\n page_id = page.get(\"id\")\n if page_id and page_id in _skip_ids:\n continue\n\n if not page_id:\n total_yielded += 1\n yield page\n continue\n\n expanded = self._try_expand_single_page(site_pages_base, page_id, page)\n total_yielded += 1\n yield expanded\n\n page_url = data.get(\"@odata.nextLink\")\n\n logger.debug(\n f\"Yielded {total_yielded} site pages (per-page expansion fallback)\"\n )\n\n def _try_expand_single_page(\n self,\n site_pages_base: str,\n page_id: str,\n fallback_page: dict[str, Any],\n ) -> dict[str, Any]:\n \"\"\"Try to GET a single page with $expand=canvasLayout. On 400, return\n the metadata-only fallback so the page is still indexed (without canvas\n content).\"\"\"\n pages_collection = site_pages_base.removesuffix(\"/microsoft.graph.sitePage\")\n single_url = f\"{pages_collection}/{page_id}/microsoft.graph.sitePage\"\n try:\n return self._graph_api_get_json(single_url, {\"$expand\": \"canvasLayout\"})\n except HTTPError as e:\n if (\n e.response is not None\n and e.response.status_code == 400\n and _is_graph_invalid_request(e.response)\n ):\n page_name = fallback_page.get(\"name\", page_id)\n logger.warning(\n f\"$expand=canvasLayout failed for page '{page_name}' ({page_id}). Indexing metadata only.\"\n )\n return fallback_page\n raise\n\n def _acquire_token(self) -> dict[str, Any]:\n \"\"\"\n Acquire token via MSAL\n \"\"\"\n if self.msal_app is None:\n raise RuntimeError(\"MSAL app is not initialized\")\n\n token = self.msal_app.acquire_token_for_client(\n scopes=[f\"{self.graph_api_host}/.default\"]\n )\n return token\n\n def _get_graph_access_token(self) -> str:\n token_data = self._acquire_token()\n access_token = token_data.get(\"access_token\")\n if not access_token:\n raise RuntimeError(\"Failed to acquire Graph API access token\")\n return access_token\n\n def _graph_api_get_json(\n self,\n url: str,\n params: dict[str, str] | None = None,\n ) -> dict[str, Any]:\n \"\"\"Make an authenticated GET request to the Graph API with retry.\"\"\"\n access_token = self._get_graph_access_token()\n headers = {\"Authorization\": f\"Bearer {access_token}\"}\n\n for attempt in range(GRAPH_API_MAX_RETRIES + 1):\n try:\n response = requests.get(\n url,\n headers=headers,\n params=params,\n timeout=REQUEST_TIMEOUT_SECONDS,\n )\n if response.status_code in GRAPH_API_RETRYABLE_STATUSES:\n if attempt < GRAPH_API_MAX_RETRIES:\n retry_after = int(\n response.headers.get(\"Retry-After\", str(2**attempt))\n )\n wait = min(retry_after, 60)\n logger.warning(\n f\"Graph API {response.status_code} on attempt {attempt + 1}, retrying in {wait}s: {url}\"\n )\n time.sleep(wait)\n # Re-acquire token in case it expired during a long traversal\n access_token = self._get_graph_access_token()\n headers = {\"Authorization\": f\"Bearer {access_token}\"}\n continue\n _log_and_raise_for_status(response)\n return response.json()\n except (requests.ConnectionError, requests.Timeout):\n if attempt < GRAPH_API_MAX_RETRIES:\n wait = min(2**attempt, 60)\n logger.warning(\n f\"Graph API connection error on attempt {attempt + 1}, retrying in {wait}s: {url}\"\n )\n time.sleep(wait)\n continue\n raise\n\n raise RuntimeError(\n f\"Graph API request failed after {GRAPH_API_MAX_RETRIES + 1} attempts: {url}\"\n )\n\n def _iter_drive_items_paged(\n self,\n drive_id: str,\n folder_path: str | None = None,\n start: datetime | None = None,\n end: datetime | None = None,\n page_size: int = 200,\n ) -> Generator[DriveItemData, None, None]:\n \"\"\"Yield DriveItemData for every file in a drive via the Graph API.\n\n Performs BFS folder traversal manually, fetching one page of children\n at a time so that memory usage stays bounded regardless of drive size.\n \"\"\"\n base = f\"{self.graph_api_base}/drives/{drive_id}\"\n if folder_path:\n encoded_path = quote(folder_path, safe=\"/\")\n start_url = f\"{base}/root:/{encoded_path}:/children\"\n else:\n start_url = f\"{base}/root/children\"\n\n folder_queue: deque[str] = deque([start_url])\n\n while folder_queue:\n page_url: str | None = folder_queue.popleft()\n params: dict[str, str] | None = {\"$top\": str(page_size)}\n\n while page_url:\n data = self._graph_api_get_json(page_url, params)\n params = None # nextLink already embeds query params\n\n for item in data.get(\"value\", []):\n if \"folder\" in item:\n child_url = f\"{base}/items/{item['id']}/children\"\n folder_queue.append(child_url)\n continue\n\n # Skip non-file items (e.g. OneNote notebooks without a \"file\" facet)\n # but still yield them — the downstream conversion handles filtering\n # by extension / mime type.\n\n # NOTE: We are now including items without a lastModifiedDateTime,\n # and respecting when only one of start or end is set.\n if start is not None or end is not None:\n raw_ts = item.get(\"lastModifiedDateTime\")\n if raw_ts:\n mod_dt = datetime.fromisoformat(\n raw_ts.replace(\"Z\", \"+00:00\")\n )\n if start is not None and mod_dt < start:\n continue\n if end is not None and mod_dt > end:\n continue\n\n yield DriveItemData.from_graph_json(item)\n\n page_url = data.get(\"@odata.nextLink\")\n\n def _iter_drive_items_delta(\n self,\n drive_id: str,\n start: datetime | None = None,\n end: datetime | None = None,\n page_size: int = 200,\n ) -> Generator[DriveItemData, None, None]:\n \"\"\"Yield DriveItemData for every file in a drive via the Graph delta API.\n\n Uses the flat delta endpoint instead of recursive folder traversal.\n On subsequent runs (start > epoch), passes the start timestamp as a\n delta token so that only changed items are returned.\n\n Falls back to full enumeration if the API returns 410 Gone (expired token).\n \"\"\"\n use_timestamp_token = start is not None and start > _EPOCH\n\n initial_url = f\"{self.graph_api_base}/drives/{drive_id}/root/delta\"\n if use_timestamp_token:\n assert start is not None # mypy\n token = quote(start.isoformat(timespec=\"seconds\"))\n initial_url += f\"?token={token}\"\n\n yield from self._iter_delta_pages(\n initial_url=initial_url,\n drive_id=drive_id,\n start=start,\n end=end,\n page_size=page_size,\n allow_full_resync=use_timestamp_token,\n )\n\n def _iter_delta_pages(\n self,\n initial_url: str,\n drive_id: str,\n start: datetime | None,\n end: datetime | None,\n page_size: int,\n allow_full_resync: bool,\n ) -> Generator[DriveItemData, None, None]:\n \"\"\"Paginate through delta API responses, yielding file DriveItemData.\n\n If the API responds with 410 Gone and allow_full_resync is True,\n restarts with a full delta enumeration.\n \"\"\"\n page_url: str | None = initial_url\n params: dict[str, str] | None = {\"$top\": str(page_size)}\n\n while page_url:\n try:\n data = self._graph_api_get_json(page_url, params)\n except requests.HTTPError as e:\n # 410 means the delta token expired, so we need to fall back to full enumeration\n if e.response is not None and e.response.status_code == 410:\n if not allow_full_resync:\n raise\n logger.warning(\n \"Delta token expired (410 Gone) for drive '%s'. Falling back to full delta enumeration.\",\n drive_id,\n )\n yield from self._iter_delta_pages(\n initial_url=f\"{self.graph_api_base}/drives/{drive_id}/root/delta\",\n drive_id=drive_id,\n start=start,\n end=end,\n page_size=page_size,\n allow_full_resync=False,\n )\n return\n raise\n\n params = None # nextLink/deltaLink already embed query params\n\n for item in data.get(\"value\", []):\n if \"folder\" in item or \"deleted\" in item:\n continue\n\n if start is not None or end is not None:\n raw_ts = item.get(\"lastModifiedDateTime\")\n if raw_ts:\n mod_dt = datetime.fromisoformat(raw_ts.replace(\"Z\", \"+00:00\"))\n if start is not None and mod_dt < start:\n continue\n if end is not None and mod_dt > end:\n continue\n\n yield DriveItemData.from_graph_json(item)\n\n page_url = data.get(\"@odata.nextLink\")\n if not page_url:\n break\n\n def _build_delta_start_url(\n self,\n drive_id: str,\n start: datetime | None = None,\n page_size: int = 200,\n ) -> str:\n \"\"\"Build the initial delta API URL with query parameters embedded.\n\n Embeds ``$top`` (and optionally a timestamp ``token``) directly in the\n URL so that the returned string is fully self-contained and can be\n stored in a checkpoint without needing a separate params dict.\n \"\"\"\n base_url = f\"{self.graph_api_base}/drives/{drive_id}/root/delta\"\n params = [f\"$top={page_size}\"]\n if start is not None and start > _EPOCH:\n token = quote(start.isoformat(timespec=\"seconds\"))\n params.append(f\"token={token}\")\n return f\"{base_url}?{'&'.join(params)}\"\n\n def _fetch_one_delta_page(\n self,\n page_url: str,\n drive_id: str,\n start: datetime | None = None,\n end: datetime | None = None,\n page_size: int = 200,\n ) -> tuple[list[DriveItemData], str | None]:\n \"\"\"Fetch a single page of delta API results.\n\n Returns ``(items, next_page_url)``. *next_page_url* is ``None`` when\n the delta enumeration is complete (deltaLink with no nextLink).\n\n On 410 Gone (expired token) returns ``([], full_resync_url)`` so\n the caller can store the resync URL in the checkpoint and retry on\n the next cycle.\n \"\"\"\n try:\n data = self._graph_api_get_json(page_url)\n except requests.HTTPError as e:\n if e.response is not None and e.response.status_code == 410:\n logger.warning(\n \"Delta token expired (410 Gone) for drive '%s'. Will restart with full delta enumeration.\",\n drive_id,\n )\n full_url = f\"{self.graph_api_base}/drives/{drive_id}/root/delta?$top={page_size}\"\n return [], full_url\n raise\n\n items: list[DriveItemData] = []\n for item in data.get(\"value\", []):\n if \"folder\" in item or \"deleted\" in item:\n continue\n if start is not None or end is not None:\n raw_ts = item.get(\"lastModifiedDateTime\")\n if raw_ts:\n mod_dt = datetime.fromisoformat(raw_ts.replace(\"Z\", \"+00:00\"))\n if start is not None and mod_dt < start:\n continue\n if end is not None and mod_dt > end:\n continue\n items.append(DriveItemData.from_graph_json(item))\n\n next_url = data.get(\"@odata.nextLink\")\n if next_url:\n return items, next_url\n return items, None\n\n @staticmethod\n def _clear_drive_checkpoint_state(\n checkpoint: \"SharepointConnectorCheckpoint\",\n ) -> None:\n \"\"\"Reset all drive-level fields in the checkpoint.\"\"\"\n checkpoint.current_drive_name = None\n checkpoint.current_drive_id = None\n checkpoint.current_drive_web_url = None\n checkpoint.current_drive_delta_next_link = None\n checkpoint.seen_document_ids.clear()\n\n def _fetch_slim_documents_from_sharepoint(\n self,\n start: datetime | None = None,\n end: datetime | None = None,\n ) -> GenerateSlimDocumentOutput:\n site_descriptors = self._filter_excluded_sites(\n self.site_descriptors or self.fetch_sites()\n )\n\n # Create a temporary checkpoint for hierarchy node tracking\n temp_checkpoint = SharepointConnectorCheckpoint(has_more=True)\n\n # goes over all urls, converts them into SlimDocument objects and then yields them in batches\n doc_batch: list[SlimDocument | HierarchyNode] = []\n for site_descriptor in site_descriptors:\n site_url = site_descriptor.url\n\n # Yield site hierarchy node using helper\n doc_batch.extend(\n self._yield_site_hierarchy_node(site_descriptor, temp_checkpoint)\n )\n\n # Process site documents if flag is True\n if self.include_site_documents:\n for driveitem, drive_name, drive_web_url in self._fetch_driveitems(\n site_descriptor=site_descriptor,\n start=start,\n end=end,\n ):\n if self._is_driveitem_excluded(driveitem):\n logger.debug(f\"Excluding by path denylist: {driveitem.web_url}\")\n continue\n\n if drive_web_url:\n doc_batch.extend(\n self._yield_drive_hierarchy_node(\n site_url, drive_web_url, drive_name, temp_checkpoint\n )\n )\n\n folder_path = self._extract_folder_path_from_parent_reference(\n driveitem.parent_reference_path\n )\n if folder_path and drive_web_url:\n doc_batch.extend(\n self._yield_folder_hierarchy_nodes(\n site_url,\n drive_web_url,\n drive_name,\n folder_path,\n temp_checkpoint,\n )\n )\n\n parent_hierarchy_url: str | None = None\n if drive_web_url:\n parent_hierarchy_url = self._get_parent_hierarchy_url(\n site_url, drive_web_url, drive_name, driveitem\n )\n\n try:\n logger.debug(f\"Processing: {driveitem.web_url}\")\n ctx = self._create_rest_client_context(site_descriptor.url)\n doc_batch.append(\n _convert_driveitem_to_slim_document(\n driveitem,\n drive_name,\n ctx,\n self.graph_client,\n parent_hierarchy_raw_node_id=parent_hierarchy_url,\n treat_sharing_link_as_public=self.treat_sharing_link_as_public,\n )\n )\n except Exception as e:\n logger.warning(f\"Failed to process driveitem: {str(e)}\")\n\n if len(doc_batch) >= SLIM_BATCH_SIZE:\n yield doc_batch\n doc_batch = []\n\n # Process site pages if flag is True\n if self.include_site_pages:\n site_pages = self._fetch_site_pages(\n site_descriptor, start=start, end=end\n )\n for site_page in site_pages:\n logger.debug(\n f\"Processing site page: {site_page.get('webUrl', site_page.get('name', 'Unknown'))}\"\n )\n ctx = self._create_rest_client_context(site_descriptor.url)\n doc_batch.append(\n _convert_sitepage_to_slim_document(\n site_page,\n ctx,\n self.graph_client,\n parent_hierarchy_raw_node_id=site_descriptor.url,\n treat_sharing_link_as_public=self.treat_sharing_link_as_public,\n )\n )\n if len(doc_batch) >= SLIM_BATCH_SIZE:\n yield doc_batch\n doc_batch = []\n yield doc_batch\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self._credential_json = credentials\n auth_method = credentials.get(\n \"authentication_method\", SharepointAuthMethod.CLIENT_SECRET.value\n )\n sp_client_id = credentials.get(\"sp_client_id\")\n sp_client_secret = credentials.get(\"sp_client_secret\")\n sp_directory_id = credentials.get(\"sp_directory_id\")\n sp_private_key = credentials.get(\"sp_private_key\")\n sp_certificate_password = credentials.get(\"sp_certificate_password\")\n\n if not sp_client_id:\n raise ConnectorValidationError(\"Client ID is required\")\n if not sp_directory_id:\n raise ConnectorValidationError(\"Directory (tenant) ID is required\")\n\n authority_url = f\"{self.authority_host}/{sp_directory_id}\"\n\n if auth_method == SharepointAuthMethod.CERTIFICATE.value:\n logger.info(\"Using certificate authentication\")\n if not sp_private_key or not sp_certificate_password:\n raise ConnectorValidationError(\n \"Private key and certificate password are required for certificate authentication\"\n )\n\n pfx_data = base64.b64decode(sp_private_key)\n certificate_data = load_certificate_from_pfx(\n pfx_data, sp_certificate_password\n )\n if certificate_data is None:\n raise RuntimeError(\"Failed to load certificate\")\n\n logger.info(f\"Creating MSAL app with authority url {authority_url}\")\n self.msal_app = msal.ConfidentialClientApplication(\n authority=authority_url,\n client_id=sp_client_id,\n client_credential=certificate_data.model_dump(),\n )\n elif auth_method == SharepointAuthMethod.CLIENT_SECRET.value:\n logger.info(\"Using client secret authentication\")\n self.msal_app = msal.ConfidentialClientApplication(\n authority=authority_url,\n client_id=sp_client_id,\n client_credential=sp_client_secret,\n )\n else:\n raise ConnectorValidationError(\n \"Invalid authentication method or missing required credentials\"\n )\n\n def _acquire_token_for_graph() -> dict[str, Any]:\n \"\"\"\n Acquire token via MSAL\n \"\"\"\n if self.msal_app is None:\n raise ConnectorValidationError(\"MSAL app is not initialized\")\n\n token = self.msal_app.acquire_token_for_client(\n scopes=[f\"{self.graph_api_host}/.default\"]\n )\n if token is None:\n raise ConnectorValidationError(\"Failed to acquire token for graph\")\n return token\n\n self._graph_client = GraphClient(\n _acquire_token_for_graph, environment=self._azure_environment\n )\n if auth_method == SharepointAuthMethod.CERTIFICATE.value:\n self.sp_tenant_domain = self._resolve_tenant_domain()\n return None\n\n def _get_drive_names_for_site(self, site_url: str) -> list[str]:\n \"\"\"Return all library/drive names for a given SharePoint site.\"\"\"\n try:\n site = self.graph_client.sites.get_by_url(site_url)\n drives = site.drives.get_all(page_loaded=lambda _: None).execute_query()\n drive_names: list[str] = []\n for drive in drives:\n if drive.name is None:\n continue\n drive_names.append(drive.name)\n\n return drive_names\n except Exception as e:\n logger.warning(f\"Failed to fetch drives for site '{site_url}': {e}\")\n return []\n\n def _build_folder_url(\n self, site_url: str, drive_name: str, folder_path: str\n ) -> str:\n \"\"\"Build a URL for a folder to use as raw_node_id.\n\n NOTE: This constructs an approximate folder URL from components rather than\n fetching the actual webUrl from the API. The constructed URL may differ\n slightly from SharePoint's canonical webUrl (e.g., URL encoding differences),\n but it functions correctly as a unique identifier for hierarchy tracking.\n We avoid fetching folder metadata to minimize API calls.\n \"\"\"\n return f\"{site_url}/{drive_name}/{folder_path}\"\n\n def _extract_folder_path_from_parent_reference(\n self, parent_reference_path: str | None\n ) -> str | None:\n \"\"\"Extract folder path from DriveItem's parentReference.path.\n\n Example input: \"/drives/b!abc123/root:/Engineering/API\"\n Example output: \"Engineering/API\"\n\n Returns None if the item is at the root of the drive.\n \"\"\"\n if not parent_reference_path:\n return None\n\n # Path format: /drives/{drive_id}/root:/folder/path\n if \"root:/\" in parent_reference_path:\n folder_path = parent_reference_path.split(\"root:/\")[1]\n return folder_path if folder_path else None\n\n # Item is at drive root\n return None\n\n def _yield_site_hierarchy_node(\n self,\n site_descriptor: SiteDescriptor,\n checkpoint: SharepointConnectorCheckpoint,\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield a hierarchy node for a site if not already yielded.\n\n Uses site.web_url as the raw_node_id (exact URL from API).\n \"\"\"\n site_url = site_descriptor.url\n\n if site_url in checkpoint.seen_hierarchy_node_raw_ids:\n return\n\n checkpoint.seen_hierarchy_node_raw_ids.add(site_url)\n\n # Extract display name from URL (last path segment)\n display_name = site_url.rstrip(\"/\").split(\"/\")[-1]\n\n yield HierarchyNode(\n raw_node_id=site_url,\n raw_parent_id=None, # Parent is SOURCE\n display_name=display_name,\n link=site_url,\n node_type=HierarchyNodeType.SITE,\n )\n\n def _yield_drive_hierarchy_node(\n self,\n site_url: str,\n drive_web_url: str,\n drive_name: str,\n checkpoint: SharepointConnectorCheckpoint,\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield a hierarchy node for a drive if not already yielded.\n\n Uses drive.web_url as the raw_node_id (exact URL from API).\n \"\"\"\n if drive_web_url in checkpoint.seen_hierarchy_node_raw_ids:\n return\n\n checkpoint.seen_hierarchy_node_raw_ids.add(drive_web_url)\n\n yield HierarchyNode(\n raw_node_id=drive_web_url,\n raw_parent_id=site_url, # Site URL is parent\n display_name=drive_name,\n link=drive_web_url,\n node_type=HierarchyNodeType.DRIVE,\n )\n\n def _yield_folder_hierarchy_nodes(\n self,\n site_url: str,\n drive_web_url: str,\n drive_name: str,\n folder_path: str,\n checkpoint: SharepointConnectorCheckpoint,\n ) -> Generator[HierarchyNode, None, None]:\n \"\"\"Yield hierarchy nodes for all folders in a path.\n\n For path \"Engineering/API/v2\", yields nodes for:\n 1. \"Engineering\" (parent = drive)\n 2. \"Engineering/API\" (parent = \"Engineering\")\n 3. \"Engineering/API/v2\" (parent = \"Engineering/API\")\n\n Nodes are yielded in parent-to-child order.\n\n Uses constructed URLs as raw_node_id. See _build_folder_url for details\n on why we construct URLs rather than fetching them from the API.\n \"\"\"\n if not folder_path:\n return\n\n path_parts = folder_path.split(\"/\")\n\n for i, part in enumerate(path_parts):\n current_path = \"/\".join(path_parts[: i + 1])\n folder_url = self._build_folder_url(site_url, drive_name, current_path)\n\n if folder_url in checkpoint.seen_hierarchy_node_raw_ids:\n continue\n\n checkpoint.seen_hierarchy_node_raw_ids.add(folder_url)\n\n # Determine parent URL\n if i == 0:\n # First folder, parent is the drive\n parent_url = drive_web_url\n else:\n # Parent is the previous folder\n parent_path = \"/\".join(path_parts[:i])\n parent_url = self._build_folder_url(site_url, drive_name, parent_path)\n\n yield HierarchyNode(\n raw_node_id=folder_url,\n raw_parent_id=parent_url,\n display_name=part, # Just the folder name\n link=folder_url,\n node_type=HierarchyNodeType.FOLDER,\n )\n\n def _get_parent_hierarchy_url(\n self,\n site_url: str,\n drive_web_url: str,\n drive_name: str,\n driveitem: DriveItemData,\n ) -> str:\n \"\"\"Determine the parent hierarchy node URL for a document.\n\n Returns:\n - Folder URL if document is in a folder\n - Drive URL if document is at drive root\n \"\"\"\n folder_path = self._extract_folder_path_from_parent_reference(\n driveitem.parent_reference_path\n )\n\n if folder_path:\n return self._build_folder_url(site_url, drive_name, folder_path)\n\n # Document is at drive root\n return drive_web_url\n\n def _load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: SharepointConnectorCheckpoint,\n include_permissions: bool = False,\n ) -> CheckpointOutput[SharepointConnectorCheckpoint]:\n\n if self._graph_client is None:\n raise ConnectorMissingCredentialError(\"Sharepoint\")\n\n checkpoint = copy.deepcopy(checkpoint)\n\n # Phase 1: Initialize cached_site_descriptors if needed\n if (\n checkpoint.has_more\n and checkpoint.cached_site_descriptors is None\n and not checkpoint.process_site_pages\n ):\n logger.info(\"Initializing SharePoint sites for processing\")\n site_descs = self._filter_excluded_sites(\n self.site_descriptors or self.fetch_sites()\n )\n checkpoint.cached_site_descriptors = deque(site_descs)\n\n if not checkpoint.cached_site_descriptors:\n logger.warning(\n \"No SharePoint sites found or accessible - nothing to process\"\n )\n checkpoint.has_more = False\n return checkpoint\n\n logger.info(\n f\"Found {len(checkpoint.cached_site_descriptors)} sites to process\"\n )\n # Set first site and return to allow checkpoint persistence\n if checkpoint.cached_site_descriptors:\n checkpoint.current_site_descriptor = (\n checkpoint.cached_site_descriptors.popleft()\n )\n logger.info(\n f\"Starting with site: {checkpoint.current_site_descriptor.url}\"\n )\n # Yield site hierarchy node for the first site\n yield from self._yield_site_hierarchy_node(\n checkpoint.current_site_descriptor, checkpoint\n )\n return checkpoint\n\n # Phase 2: Initialize cached_drive_names for current site if needed\n if checkpoint.current_site_descriptor and checkpoint.cached_drive_names is None:\n # If site documents flag is False, set empty drive list to skip document processing\n if not self.include_site_documents:\n logger.debug(\"Documents disabled, skipping drive initialization\")\n checkpoint.cached_drive_names = deque()\n return checkpoint\n\n logger.info(\n f\"Initializing drives for site: {checkpoint.current_site_descriptor.url}\"\n )\n\n try:\n # If the user explicitly specified drive(s) for this site, honour that\n if checkpoint.current_site_descriptor.drive_name:\n logger.info(\n f\"Using explicitly specified drive: {checkpoint.current_site_descriptor.drive_name}\"\n )\n checkpoint.cached_drive_names = deque(\n [checkpoint.current_site_descriptor.drive_name]\n )\n else:\n drive_names = self._get_drive_names_for_site(\n checkpoint.current_site_descriptor.url\n )\n checkpoint.cached_drive_names = deque(drive_names)\n\n if not checkpoint.cached_drive_names:\n logger.warning(\n f\"No accessible drives found for site: {checkpoint.current_site_descriptor.url}\"\n )\n else:\n logger.info(\n f\"Found {len(checkpoint.cached_drive_names)} drives: {list(checkpoint.cached_drive_names)}\"\n )\n\n except Exception as e:\n logger.error(\n f\"Failed to initialize drives for site: {checkpoint.current_site_descriptor.url}: {e}\"\n )\n # Yield a ConnectorFailure for site-level access failures\n start_dt = datetime.fromtimestamp(start, tz=timezone.utc)\n end_dt = datetime.fromtimestamp(end, tz=timezone.utc)\n yield _create_entity_failure(\n checkpoint.current_site_descriptor.url,\n f\"Failed to access site: {str(e)}\",\n (start_dt, end_dt),\n e,\n )\n # Move to next site if available\n if (\n checkpoint.cached_site_descriptors\n and len(checkpoint.cached_site_descriptors) > 0\n ):\n checkpoint.current_site_descriptor = (\n checkpoint.cached_site_descriptors.popleft()\n )\n checkpoint.cached_drive_names = None # Reset for new site\n return checkpoint\n else:\n # No more sites - we're done\n checkpoint.has_more = False\n return checkpoint\n\n # Return checkpoint to allow persistence after drive initialization\n return checkpoint\n\n # Phase 3a: Initialize the next drive for processing\n if (\n checkpoint.current_site_descriptor\n and checkpoint.cached_drive_names\n and len(checkpoint.cached_drive_names) > 0\n and checkpoint.current_drive_name is None\n ):\n checkpoint.current_drive_name = checkpoint.cached_drive_names.popleft()\n\n start_dt = datetime.fromtimestamp(start, tz=timezone.utc)\n end_dt = datetime.fromtimestamp(end, tz=timezone.utc)\n site_descriptor = checkpoint.current_site_descriptor\n\n logger.info(\n f\"Processing drive '{checkpoint.current_drive_name}' in site: {site_descriptor.url}\"\n )\n logger.debug(f\"Time range: {start_dt} to {end_dt}\")\n\n current_drive_name = checkpoint.current_drive_name\n if current_drive_name is None:\n logger.warning(\"Current drive name is None, skipping\")\n return checkpoint\n\n try:\n logger.info(\n f\"Fetching drive items for drive name: {current_drive_name}\"\n )\n result = self._resolve_drive(site_descriptor, current_drive_name)\n if result is None:\n logger.warning(f\"Drive '{current_drive_name}' not found, skipping\")\n self._clear_drive_checkpoint_state(checkpoint)\n return checkpoint\n\n drive_id, drive_web_url = result\n checkpoint.current_drive_id = drive_id\n checkpoint.current_drive_web_url = drive_web_url\n except Exception as e:\n logger.error(\n f\"Failed to retrieve items from drive '{current_drive_name}' in site: {site_descriptor.url}: {e}\"\n )\n yield _create_entity_failure(\n f\"{site_descriptor.url}|{current_drive_name}\",\n f\"Failed to access drive '{current_drive_name}' in site '{site_descriptor.url}': {str(e)}\",\n (start_dt, end_dt),\n e,\n )\n self._clear_drive_checkpoint_state(checkpoint)\n return checkpoint\n\n display_drive_name = SHARED_DOCUMENTS_MAP.get(\n current_drive_name, current_drive_name\n )\n\n if drive_web_url:\n yield from self._yield_drive_hierarchy_node(\n site_descriptor.url,\n drive_web_url,\n display_drive_name,\n checkpoint,\n )\n\n # For non-folder-scoped drives, use delta API with per-page\n # checkpointing. Build the initial URL and fall through to 3b.\n if not site_descriptor.folder_path:\n checkpoint.current_drive_delta_next_link = self._build_delta_start_url(\n drive_id, start_dt\n )\n # else: BFS path — delta_next_link stays None;\n # Phase 3b will use _iter_drive_items_paged.\n\n # Phase 3b: Process items from the current drive\n if (\n checkpoint.current_site_descriptor\n and checkpoint.current_drive_name is not None\n and checkpoint.current_drive_id is not None\n ):\n site_descriptor = checkpoint.current_site_descriptor\n start_dt = datetime.fromtimestamp(start, tz=timezone.utc)\n end_dt = datetime.fromtimestamp(end, tz=timezone.utc)\n current_drive_name = SHARED_DOCUMENTS_MAP.get(\n checkpoint.current_drive_name, checkpoint.current_drive_name\n )\n drive_web_url = checkpoint.current_drive_web_url\n\n # --- determine item source ---\n driveitems: Iterable[DriveItemData]\n has_more_delta_pages = False\n\n if checkpoint.current_drive_delta_next_link:\n # Delta path: fetch one page at a time for checkpointing\n try:\n page_items, next_url = self._fetch_one_delta_page(\n page_url=checkpoint.current_drive_delta_next_link,\n drive_id=checkpoint.current_drive_id,\n start=start_dt,\n end=end_dt,\n )\n except Exception as e:\n logger.error(\n f\"Failed to fetch delta page for drive '{current_drive_name}': {e}\"\n )\n yield _create_entity_failure(\n f\"{site_descriptor.url}|{current_drive_name}\",\n f\"Failed to fetch delta page for drive '{current_drive_name}': {str(e)}\",\n (start_dt, end_dt),\n e,\n )\n self._clear_drive_checkpoint_state(checkpoint)\n return checkpoint\n\n driveitems = page_items\n has_more_delta_pages = next_url is not None\n if next_url:\n checkpoint.current_drive_delta_next_link = next_url\n else:\n # BFS path (folder-scoped): process all items at once\n driveitems = self._iter_drive_items_paged(\n drive_id=checkpoint.current_drive_id,\n folder_path=site_descriptor.folder_path,\n start=start_dt,\n end=end_dt,\n )\n\n item_count = 0\n for driveitem in driveitems:\n item_count += 1\n\n if self._is_driveitem_excluded(driveitem):\n logger.debug(f\"Excluding by path denylist: {driveitem.web_url}\")\n continue\n\n if driveitem.id and driveitem.id in checkpoint.seen_document_ids:\n logger.debug(\n f\"Skipping duplicate document {driveitem.id} ({driveitem.name})\"\n )\n continue\n\n driveitem_extension = get_file_ext(driveitem.name)\n if driveitem_extension not in OnyxFileExtensions.ALL_ALLOWED_EXTENSIONS:\n logger.warning(\n f\"Skipping {driveitem.web_url} as it is not a supported file type\"\n )\n continue\n\n should_yield_if_empty = (\n driveitem_extension in OnyxFileExtensions.IMAGE_EXTENSIONS\n or driveitem_extension == \".pdf\"\n )\n\n folder_path = self._extract_folder_path_from_parent_reference(\n driveitem.parent_reference_path\n )\n if folder_path and drive_web_url:\n yield from self._yield_folder_hierarchy_nodes(\n site_descriptor.url,\n drive_web_url,\n current_drive_name,\n folder_path,\n checkpoint,\n )\n\n parent_hierarchy_url: str | None = None\n if drive_web_url:\n parent_hierarchy_url = self._get_parent_hierarchy_url(\n site_descriptor.url,\n drive_web_url,\n current_drive_name,\n driveitem,\n )\n\n try:\n ctx: ClientContext | None = None\n if include_permissions:\n ctx = self._create_rest_client_context(site_descriptor.url)\n\n access_token = self._get_graph_access_token()\n doc_or_failure = _convert_driveitem_to_document_with_permissions(\n driveitem,\n current_drive_name,\n ctx,\n self.graph_client,\n include_permissions=include_permissions,\n parent_hierarchy_raw_node_id=parent_hierarchy_url,\n graph_api_base=self.graph_api_base,\n access_token=access_token,\n treat_sharing_link_as_public=self.treat_sharing_link_as_public,\n )\n\n if isinstance(doc_or_failure, Document):\n if doc_or_failure.sections:\n checkpoint.seen_document_ids.add(doc_or_failure.id)\n yield doc_or_failure\n elif should_yield_if_empty:\n doc_or_failure.sections = [\n TextSection(link=driveitem.web_url, text=\"\")\n ]\n checkpoint.seen_document_ids.add(doc_or_failure.id)\n yield doc_or_failure\n else:\n logger.warning(\n f\"Skipping {driveitem.web_url} as it is empty and not a PDF or image\"\n )\n elif isinstance(doc_or_failure, ConnectorFailure):\n yield doc_or_failure\n except Exception as e:\n logger.warning(\n f\"Failed to process driveitem {driveitem.web_url}: {e}\"\n )\n yield _create_document_failure(\n driveitem, f\"Failed to process: {str(e)}\", e\n )\n\n logger.info(f\"Processed {item_count} items in drive '{current_drive_name}'\")\n\n if has_more_delta_pages:\n return checkpoint\n\n self._clear_drive_checkpoint_state(checkpoint)\n\n # Phase 4: Progression logic - determine next step\n # If we have more drives in current site, continue with current site\n if checkpoint.cached_drive_names and len(checkpoint.cached_drive_names) > 0:\n logger.debug(\n f\"Continuing with {len(checkpoint.cached_drive_names)} remaining drives in current site\"\n )\n return checkpoint\n\n if (\n self.include_site_pages\n and not checkpoint.process_site_pages\n and checkpoint.current_site_descriptor is not None\n ):\n logger.info(\n f\"Processing site pages for site: {checkpoint.current_site_descriptor.url}\"\n )\n checkpoint.process_site_pages = True\n return checkpoint\n\n # Phase 5: Process site pages\n if (\n checkpoint.process_site_pages\n and checkpoint.current_site_descriptor is not None\n ):\n # Fetch SharePoint site pages (.aspx files)\n site_descriptor = checkpoint.current_site_descriptor\n start_dt = datetime.fromtimestamp(start, tz=timezone.utc)\n end_dt = datetime.fromtimestamp(end, tz=timezone.utc)\n site_pages = self._fetch_site_pages(\n site_descriptor, start=start_dt, end=end_dt\n )\n for site_page in site_pages:\n logger.debug(\n f\"Processing site page: {site_page.get('webUrl', site_page.get('name', 'Unknown'))}\"\n )\n client_ctx: ClientContext | None = None\n if include_permissions:\n client_ctx = self._create_rest_client_context(site_descriptor.url)\n yield (\n _convert_sitepage_to_document(\n site_page,\n site_descriptor.drive_name,\n client_ctx,\n self.graph_client,\n include_permissions=include_permissions,\n # Site pages have the site as their parent\n parent_hierarchy_raw_node_id=site_descriptor.url,\n treat_sharing_link_as_public=self.treat_sharing_link_as_public,\n )\n )\n logger.info(\n f\"Finished processing site pages for site: {site_descriptor.url}\"\n )\n\n # If no more drives, move to next site if available\n if (\n checkpoint.cached_site_descriptors\n and len(checkpoint.cached_site_descriptors) > 0\n ):\n current_site = (\n checkpoint.current_site_descriptor.url\n if checkpoint.current_site_descriptor\n else \"unknown\"\n )\n checkpoint.current_site_descriptor = (\n checkpoint.cached_site_descriptors.popleft()\n )\n checkpoint.cached_drive_names = None # Reset for new site\n checkpoint.process_site_pages = False\n logger.info(\n f\"Finished site '{current_site}', moving to next site: {checkpoint.current_site_descriptor.url}\"\n )\n logger.info(\n f\"Remaining sites to process: {len(checkpoint.cached_site_descriptors) + 1}\"\n )\n # Yield site hierarchy node for the new site\n yield from self._yield_site_hierarchy_node(\n checkpoint.current_site_descriptor, checkpoint\n )\n return checkpoint\n\n # No more sites or drives - we're done\n current_site = (\n checkpoint.current_site_descriptor.url\n if checkpoint.current_site_descriptor\n else \"unknown\"\n )\n logger.info(\n f\"SharePoint processing complete. Finished last site: {current_site}\"\n )\n checkpoint.has_more = False\n return checkpoint\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: SharepointConnectorCheckpoint,\n ) -> CheckpointOutput[SharepointConnectorCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=False\n )\n\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: SharepointConnectorCheckpoint,\n ) -> CheckpointOutput[SharepointConnectorCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=True\n )\n\n def build_dummy_checkpoint(self) -> SharepointConnectorCheckpoint:\n return SharepointConnectorCheckpoint(has_more=True)\n\n def validate_checkpoint_json(\n self, checkpoint_json: str\n ) -> SharepointConnectorCheckpoint:\n return SharepointConnectorCheckpoint.model_validate_json(checkpoint_json)\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n start_dt = (\n datetime.fromtimestamp(start, tz=timezone.utc)\n if start is not None\n else None\n )\n end_dt = (\n datetime.fromtimestamp(end, tz=timezone.utc) if end is not None else None\n )\n yield from self._fetch_slim_documents_from_sharepoint(\n start=start_dt,\n end=end_dt,\n )\n\n\nif __name__ == \"__main__\":\n from onyx.connectors.connector_runner import ConnectorRunner\n\n connector = SharepointConnector(sites=os.environ[\"SHAREPOINT_SITES\"].split(\",\"))\n\n connector.load_credentials(\n {\n \"sp_client_id\": os.environ[\"SHAREPOINT_CLIENT_ID\"],\n \"sp_client_secret\": os.environ[\"SHAREPOINT_CLIENT_SECRET\"],\n \"sp_directory_id\": os.environ[\"SHAREPOINT_CLIENT_DIRECTORY_ID\"],\n }\n )\n\n # Create a time range from epoch to now\n end_time = datetime.now(timezone.utc)\n start_time = datetime.fromtimestamp(0, tz=timezone.utc)\n time_range = (start_time, end_time)\n\n # Initialize the runner with a batch size of 10\n runner: ConnectorRunner[SharepointConnectorCheckpoint] = ConnectorRunner(\n connector, batch_size=10, include_permissions=False, time_range=time_range\n )\n\n # Get initial checkpoint\n checkpoint = connector.build_dummy_checkpoint()\n\n # Run the connector\n while checkpoint.has_more:\n for doc_batch, hierarchy_node_batch, failure, next_checkpoint in runner.run(\n checkpoint\n ):\n if doc_batch:\n print(f\"Retrieved batch of {len(doc_batch)} documents\")\n for test_doc in doc_batch:\n print(f\"Document: {test_doc.semantic_identifier}\")\n if failure:\n print(f\"Failure: {failure.failure_message}\")\n if next_checkpoint:\n checkpoint = next_checkpoint\n" }, { "path": "backend/onyx/connectors/sharepoint/connector_utils.py", "content": "from typing import Any\n\nfrom office365.graph_client import GraphClient # type: ignore[import-untyped]\nfrom office365.onedrive.driveitems.driveItem import DriveItem # type: ignore[import-untyped]\nfrom office365.sharepoint.client_context import ClientContext # type: ignore[import-untyped]\n\nfrom onyx.connectors.models import ExternalAccess\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\n\n\ndef get_sharepoint_external_access(\n ctx: ClientContext,\n graph_client: GraphClient,\n drive_item: DriveItem | None = None,\n drive_name: str | None = None,\n site_page: dict[str, Any] | None = None,\n add_prefix: bool = False,\n treat_sharing_link_as_public: bool = False,\n) -> ExternalAccess:\n if drive_item and drive_item.id is None:\n raise ValueError(\"DriveItem ID is required\")\n\n # Get external access using the EE implementation\n def noop_fallback(\n *args: Any, **kwargs: Any # noqa: ARG001\n ) -> ExternalAccess: # noqa: ARG001\n return ExternalAccess.empty()\n\n get_external_access_func = fetch_versioned_implementation_with_fallback(\n \"onyx.external_permissions.sharepoint.permission_utils\",\n \"get_external_access_from_sharepoint\",\n fallback=noop_fallback,\n )\n\n external_access = get_external_access_func(\n ctx,\n graph_client,\n drive_name,\n drive_item,\n site_page,\n add_prefix,\n treat_sharing_link_as_public,\n )\n\n return external_access\n" }, { "path": "backend/onyx/connectors/slab/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/slab/connector.py", "content": "import json\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom urllib.parse import urljoin\n\nimport requests\nfrom dateutil import parser\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n# Fairly generous retry because it's not understood why occasionally GraphQL requests fail even with timeout > 1 min\nSLAB_GRAPHQL_MAX_TRIES = 10\nSLAB_API_URL = \"https://api.slab.com/v1/graphql\"\n\n_SLIM_BATCH_SIZE = 1000\n\n\ndef run_graphql_request(\n graphql_query: dict, bot_token: str, max_tries: int = SLAB_GRAPHQL_MAX_TRIES\n) -> str:\n headers = {\"Authorization\": bot_token, \"Content-Type\": \"application/json\"}\n\n for try_count in range(max_tries):\n try:\n response = requests.post(\n SLAB_API_URL, headers=headers, json=graphql_query, timeout=60\n )\n response.raise_for_status()\n\n if response.status_code != 200:\n raise ValueError(f\"GraphQL query failed: {graphql_query}\")\n\n return response.text\n\n except (requests.exceptions.Timeout, ValueError) as e:\n if try_count < max_tries - 1:\n logger.warning(\"A Slab GraphQL error occurred. Retrying...\")\n continue\n\n if isinstance(e, requests.exceptions.Timeout):\n raise TimeoutError(\"Slab API timed out after 3 attempts\")\n else:\n raise ValueError(\"Slab GraphQL query failed after 3 attempts\")\n\n raise RuntimeError(\n \"Unexpected execution from Slab Connector. This should not happen.\"\n ) # for static checker\n\n\ndef get_all_post_ids(bot_token: str) -> list[str]:\n query = \"\"\"\n query GetAllPostIds {\n organization {\n posts {\n id\n }\n }\n }\n \"\"\"\n\n graphql_query = {\"query\": query}\n\n results = json.loads(run_graphql_request(graphql_query, bot_token))\n posts = results[\"data\"][\"organization\"][\"posts\"]\n return [post[\"id\"] for post in posts]\n\n\ndef get_post_by_id(post_id: str, bot_token: str) -> dict[str, str]:\n query = \"\"\"\n query GetPostById($postId: ID!) {\n post(id: $postId) {\n title\n content\n linkAccess\n updatedAt\n }\n }\n \"\"\"\n graphql_query = {\"query\": query, \"variables\": {\"postId\": post_id}}\n results = json.loads(run_graphql_request(graphql_query, bot_token))\n return results[\"data\"][\"post\"]\n\n\ndef iterate_post_batches(\n batch_size: int, bot_token: str\n) -> Generator[list[dict[str, str]], None, None]:\n \"\"\"This may not be safe to use, not sure if page edits will change the order of results\"\"\"\n query = \"\"\"\n query IteratePostBatches($query: String!, $first: Int, $types: [SearchType], $after: String) {\n search(query: $query, first: $first, types: $types, after: $after) {\n edges {\n node {\n ... on PostSearchResult {\n post {\n id\n title\n content\n updatedAt\n }\n }\n }\n }\n pageInfo {\n endCursor\n hasNextPage\n }\n }\n }\n \"\"\"\n pagination_start = None\n exists_more_pages = True\n while exists_more_pages:\n graphql_query = {\n \"query\": query,\n \"variables\": {\n \"query\": \"\",\n \"first\": batch_size,\n \"types\": [\"POST\"],\n \"after\": pagination_start,\n },\n }\n results = json.loads(run_graphql_request(graphql_query, bot_token))\n pagination_start = results[\"data\"][\"search\"][\"pageInfo\"][\"endCursor\"]\n hits = results[\"data\"][\"search\"][\"edges\"]\n\n posts = [hit[\"node\"] for hit in hits]\n if posts:\n yield posts\n\n exists_more_pages = results[\"data\"][\"search\"][\"pageInfo\"][\"hasNextPage\"]\n\n\ndef get_slab_url_from_title_id(base_url: str, title: str, page_id: str) -> str:\n \"\"\"This is not a documented approach but seems to be the way it works currently\n May be subject to change without notification\"\"\"\n title = (\n title.replace(\"[\", \"\")\n .replace(\"]\", \"\")\n .replace(\":\", \"\")\n .replace(\" \", \"-\")\n .lower()\n )\n url_id = title + \"-\" + page_id\n return urljoin(urljoin(base_url, \"posts/\"), url_id)\n\n\nclass SlabConnector(LoadConnector, PollConnector, SlimConnectorWithPermSync):\n def __init__(\n self,\n base_url: str,\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n self.base_url = base_url\n self.batch_size = batch_size\n self._slab_bot_token: str | None = None\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n self._slab_bot_token = credentials[\"slab_bot_token\"]\n return None\n\n @property\n def slab_bot_token(self) -> str:\n if self._slab_bot_token is None:\n raise ConnectorMissingCredentialError(\"Slab\")\n return self._slab_bot_token\n\n def _iterate_posts(\n self, time_filter: Callable[[datetime], bool] | None = None\n ) -> GenerateDocumentsOutput:\n doc_batch: list[Document | HierarchyNode] = []\n\n if self.slab_bot_token is None:\n raise ConnectorMissingCredentialError(\"Slab\")\n\n all_post_ids: list[str] = get_all_post_ids(self.slab_bot_token)\n\n for post_id in all_post_ids:\n post = get_post_by_id(post_id, self.slab_bot_token)\n last_modified = parser.parse(post[\"updatedAt\"])\n if time_filter is not None and not time_filter(last_modified):\n continue\n\n page_url = get_slab_url_from_title_id(self.base_url, post[\"title\"], post_id)\n\n content_text = \"\"\n contents = json.loads(post[\"content\"])\n for content_segment in contents:\n insert = content_segment.get(\"insert\")\n if insert and isinstance(insert, str):\n content_text += insert\n\n doc_batch.append(\n Document(\n id=post_id, # can't be url as this changes with the post title\n sections=[TextSection(link=page_url, text=content_text)],\n source=DocumentSource.SLAB,\n semantic_identifier=post[\"title\"],\n metadata={},\n )\n )\n\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n yield from self._iterate_posts()\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n start_time = datetime.fromtimestamp(start, tz=timezone.utc)\n end_time = datetime.fromtimestamp(end, tz=timezone.utc)\n\n yield from self._iterate_posts(\n time_filter=lambda t: start_time <= t <= end_time\n )\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n end: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n slim_doc_batch: list[SlimDocument | HierarchyNode] = []\n for post_id in get_all_post_ids(self.slab_bot_token):\n slim_doc_batch.append(\n SlimDocument(\n id=post_id,\n )\n )\n if len(slim_doc_batch) >= _SLIM_BATCH_SIZE:\n yield slim_doc_batch\n slim_doc_batch = []\n if slim_doc_batch:\n yield slim_doc_batch\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n Very basic validation, we could do more here\n \"\"\"\n if not self.base_url.startswith(\"https://\") and not self.base_url.startswith(\n \"http://\"\n ):\n raise ConnectorValidationError(\n \"Base URL must start with https:// or http://\"\n )\n\n try:\n get_all_post_ids(self.slab_bot_token)\n except ConnectorMissingCredentialError:\n raise\n except Exception as e:\n raise ConnectorValidationError(f\"Failed to fetch posts from Slab: {e}\")\n" }, { "path": "backend/onyx/connectors/slack/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/slack/access.py", "content": "from collections.abc import Callable\nfrom typing import cast\n\nfrom slack_sdk import WebClient\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.slack.models import ChannelType\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\n\ndef get_channel_access(\n client: WebClient,\n channel: ChannelType,\n user_cache: dict[str, BasicExpertInfo | None],\n) -> ExternalAccess | None:\n \"\"\"\n Get channel access permissions for a Slack channel.\n This functionality requires Enterprise Edition.\n\n Args:\n client: Slack WebClient instance\n channel: Slack channel object containing channel info\n user_cache: Cache of user IDs to BasicExpertInfo objects. May be updated in place.\n\n Returns:\n ExternalAccess object for the channel. None if EE is not enabled.\n \"\"\"\n # Check if EE is enabled\n if not global_version.is_ee_version():\n return None\n\n # Fetch the EE implementation\n ee_get_channel_access = cast(\n Callable[\n [WebClient, ChannelType, dict[str, BasicExpertInfo | None]],\n ExternalAccess,\n ],\n fetch_versioned_implementation(\n \"onyx.external_permissions.slack.channel_access\", \"get_channel_access\"\n ),\n )\n\n return ee_get_channel_access(client, channel, user_cache)\n" }, { "path": "backend/onyx/connectors/slack/connector.py", "content": "import contextvars\nimport copy\nimport itertools\nimport re\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import Future\nfrom concurrent.futures import ThreadPoolExecutor\nfrom datetime import datetime\nfrom datetime import timezone\nfrom enum import Enum\nfrom http.client import IncompleteRead\nfrom http.client import RemoteDisconnected\nfrom typing import Any\nfrom typing import cast\nfrom urllib.error import URLError\nfrom urllib.parse import urlparse\n\nfrom pydantic import BaseModel\nfrom redis import Redis\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\nfrom slack_sdk.http_retry import ConnectionErrorRetryHandler\nfrom slack_sdk.http_retry import RetryHandler\nfrom slack_sdk.http_retry.builtin_interval_calculators import (\n FixedValueRetryIntervalCalculator,\n)\nfrom typing_extensions import override\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.app_configs import ENABLE_EXPENSIVE_EXPERT_CALLS\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import SLACK_NUM_THREADS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import CredentialsConnector\nfrom onyx.connectors.interfaces import CredentialsProviderInterface\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import NormalizationResult\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import EntityFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.slack.access import get_channel_access\nfrom onyx.connectors.slack.models import ChannelType\nfrom onyx.connectors.slack.models import MessageType\nfrom onyx.connectors.slack.models import ThreadType\nfrom onyx.connectors.slack.onyx_retry_handler import OnyxRedisSlackRetryHandler\nfrom onyx.connectors.slack.onyx_slack_web_client import OnyxSlackWebClient\nfrom onyx.connectors.slack.utils import (\n expert_info_from_slack_id,\n)\nfrom onyx.connectors.slack.utils import get_message_link\nfrom onyx.connectors.slack.utils import make_paginated_slack_api_call\nfrom onyx.connectors.slack.utils import SlackTextCleaner\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_SLACK_LIMIT = 900\n\n\nclass SlackCheckpoint(ConnectorCheckpoint):\n channel_ids: list[str] | None # e.g. C8E6WHE2X\n\n # channel id mapped to the timestamp we want to retrieve messages up to\n # NOTE: this is usually the earliest timestamp of all the messages we have\n # since we walk backwards\n channel_completion_map: dict[str, str]\n current_channel: ChannelType | None\n current_channel_access: ExternalAccess | None\n\n seen_thread_ts: list[\n str\n ] # apparently we identify threads/messages uniquely by timestamp?\n\n\ndef _collect_paginated_channels(\n client: WebClient,\n exclude_archived: bool,\n channel_types: list[str],\n) -> list[ChannelType]:\n channels: list[ChannelType] = []\n for result in make_paginated_slack_api_call(\n client.conversations_list,\n exclude_archived=exclude_archived,\n # also get private channels the bot is added to\n types=channel_types,\n ):\n channels.extend(result[\"channels\"])\n\n return channels\n\n\ndef get_channels(\n client: WebClient,\n exclude_archived: bool = True,\n get_public: bool = True,\n get_private: bool = True,\n) -> list[ChannelType]:\n \"\"\"Get all channels in the workspace.\"\"\"\n channels: list[ChannelType] = []\n channel_types = []\n if get_public:\n channel_types.append(\"public_channel\")\n if get_private:\n channel_types.append(\"private_channel\")\n # Try fetching both public and private channels first:\n try:\n channels = _collect_paginated_channels(\n client=client,\n exclude_archived=exclude_archived,\n channel_types=channel_types,\n )\n except SlackApiError as e:\n msg = f\"Unable to fetch private channels due to: {e}.\"\n if not get_public:\n logger.warning(msg + \" Public channels are not enabled.\")\n return []\n\n logger.warning(msg + \" Trying again with public channels only.\")\n channel_types = [\"public_channel\"]\n channels = _collect_paginated_channels(\n client=client,\n exclude_archived=exclude_archived,\n channel_types=channel_types,\n )\n return channels\n\n\ndef get_channel_messages(\n client: WebClient,\n channel: ChannelType,\n oldest: str | None = None,\n latest: str | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n) -> Generator[list[MessageType], None, None]:\n \"\"\"Get all messages in a channel\"\"\"\n # join so that the bot can access messages\n if not channel[\"is_member\"]:\n client.conversations_join(\n channel=channel[\"id\"],\n is_private=channel[\"is_private\"],\n )\n logger.info(f\"Successfully joined '{channel['name']}'\")\n\n for result in make_paginated_slack_api_call(\n client.conversations_history,\n channel=channel[\"id\"],\n oldest=oldest,\n latest=latest,\n ):\n if callback:\n if callback.should_stop():\n raise RuntimeError(\"get_channel_messages: Stop signal detected\")\n\n callback.progress(\"get_channel_messages\", 0)\n yield cast(list[MessageType], result[\"messages\"])\n\n\ndef get_thread(client: WebClient, channel_id: str, thread_id: str) -> ThreadType:\n \"\"\"Get all messages in a thread\"\"\"\n threads: list[MessageType] = []\n for result in make_paginated_slack_api_call(\n client.conversations_replies, channel=channel_id, ts=thread_id\n ):\n threads.extend(result[\"messages\"])\n return threads\n\n\ndef get_latest_message_time(thread: ThreadType) -> datetime:\n max_ts = max([float(msg.get(\"ts\", 0)) for msg in thread])\n return datetime.fromtimestamp(max_ts, tz=timezone.utc)\n\n\ndef _build_doc_id(channel_id: str, thread_ts: str) -> str:\n return f\"{channel_id}__{thread_ts}\"\n\n\ndef thread_to_doc(\n channel: ChannelType,\n thread: ThreadType,\n slack_cleaner: SlackTextCleaner,\n client: WebClient,\n user_cache: dict[str, BasicExpertInfo | None],\n channel_access: ExternalAccess | None,\n) -> Document:\n channel_id = channel[\"id\"]\n\n initial_sender_expert_info = expert_info_from_slack_id(\n user_id=thread[0].get(\"user\"), client=client, user_cache=user_cache\n )\n initial_sender_name = (\n initial_sender_expert_info.get_semantic_name()\n if initial_sender_expert_info\n else \"Unknown\"\n )\n\n valid_experts = None\n if ENABLE_EXPENSIVE_EXPERT_CALLS:\n all_sender_ids = [m.get(\"user\") for m in thread]\n experts = [\n expert_info_from_slack_id(\n user_id=sender_id, client=client, user_cache=user_cache\n )\n for sender_id in all_sender_ids\n if sender_id\n ]\n valid_experts = [expert for expert in experts if expert]\n\n first_message = slack_cleaner.index_clean(cast(str, thread[0][\"text\"]))\n snippet = (\n first_message[:50].rstrip() + \"...\"\n if len(first_message) > 50\n else first_message\n )\n\n doc_sem_id = f\"{initial_sender_name} in #{channel['name']}: {snippet}\".replace(\n \"\\n\", \" \"\n )\n\n channel_name = channel[\"name\"]\n\n return Document(\n id=_build_doc_id(channel_id=channel_id, thread_ts=thread[0][\"ts\"]),\n sections=[\n TextSection(\n link=get_message_link(event=m, client=client, channel_id=channel_id),\n text=slack_cleaner.index_clean(cast(str, m[\"text\"])),\n )\n for m in thread\n ],\n source=DocumentSource.SLACK,\n semantic_identifier=doc_sem_id,\n doc_updated_at=get_latest_message_time(thread),\n primary_owners=valid_experts,\n doc_metadata={\n \"hierarchy\": {\n \"source_path\": [channel_name],\n \"channel_name\": channel_name,\n \"channel_id\": channel_id,\n }\n },\n metadata={\"Channel\": channel_name},\n external_access=channel_access,\n parent_hierarchy_raw_node_id=channel_id,\n )\n\n\n# list of subtypes can be found here: https://api.slack.com/events/message\n_DISALLOWED_MSG_SUBTYPES = {\n \"channel_join\",\n \"channel_leave\",\n \"channel_archive\",\n \"channel_unarchive\",\n \"pinned_item\",\n \"unpinned_item\",\n \"ekm_access_denied\",\n \"channel_posting_permissions\",\n \"group_join\",\n \"group_leave\",\n \"group_archive\",\n \"group_unarchive\",\n \"channel_leave\",\n \"channel_name\",\n \"channel_join\",\n}\n\n\nclass SlackMessageFilterReason(str, Enum):\n BOT = \"bot\"\n DISALLOWED = \"disallowed\"\n\n\ndef default_msg_filter(message: MessageType) -> SlackMessageFilterReason | None:\n \"\"\"Returns a filter reason if the message should be filtered out.\n Returns None if the message can be kept.\n \"\"\"\n\n # Don't keep messages from bots\n if message.get(\"bot_id\") or message.get(\"app_id\"):\n bot_profile_name = message.get(\"bot_profile\", {}).get(\"name\")\n if bot_profile_name == \"DanswerBot Testing\":\n return None\n return SlackMessageFilterReason.BOT\n\n # Uninformative\n if message.get(\"subtype\", \"\") in _DISALLOWED_MSG_SUBTYPES:\n return SlackMessageFilterReason.DISALLOWED\n\n return None\n\n\ndef _bot_inclusive_msg_filter(\n message: MessageType,\n) -> SlackMessageFilterReason | None:\n \"\"\"Like default_msg_filter but allows bot/app messages through.\n Only filters out disallowed subtypes (channel_join, channel_leave, etc.).\n \"\"\"\n if message.get(\"subtype\", \"\") in _DISALLOWED_MSG_SUBTYPES:\n return SlackMessageFilterReason.DISALLOWED\n\n return None\n\n\ndef filter_channels(\n all_channels: list[ChannelType],\n channels_to_connect: list[str] | None,\n regex_enabled: bool,\n) -> list[ChannelType]:\n if not channels_to_connect:\n return all_channels\n\n if regex_enabled:\n return [\n channel\n for channel in all_channels\n if any(\n re.fullmatch(channel_to_connect, channel[\"name\"])\n for channel_to_connect in channels_to_connect\n )\n ]\n\n # validate that all channels in `channels_to_connect` are valid\n # fail loudly in the case of an invalid channel so that the user\n # knows that one of the channels they've specified is typo'd or private\n all_channel_names = {channel[\"name\"] for channel in all_channels}\n for channel in channels_to_connect:\n if channel not in all_channel_names:\n raise ValueError(\n f\"Channel '{channel}' not found in workspace. \"\n f\"Available channels (Showing {len(all_channel_names)} of \"\n f\"{min(len(all_channel_names), SlackConnector.MAX_CHANNELS_TO_LOG)}): \"\n f\"{list(itertools.islice(all_channel_names, SlackConnector.MAX_CHANNELS_TO_LOG))}\"\n )\n\n return [\n channel for channel in all_channels if channel[\"name\"] in channels_to_connect\n ]\n\n\ndef _channel_to_hierarchy_node(\n channel: ChannelType,\n channel_access: ExternalAccess | None,\n workspace_url: str | None = None,\n) -> HierarchyNode:\n \"\"\"Convert a Slack channel to a HierarchyNode.\n\n Args:\n channel: The Slack channel object\n channel_access: External access permissions for the channel\n workspace_url: The workspace URL (e.g., https://myworkspace.slack.com)\n\n Returns:\n A HierarchyNode representing the channel\n \"\"\"\n # Link format: https://{workspace}.slack.com/archives/{channel_id}\n link = f\"{workspace_url}/archives/{channel['id']}\" if workspace_url else None\n\n return HierarchyNode(\n raw_node_id=channel[\"id\"],\n raw_parent_id=None, # Direct child of SOURCE\n display_name=f\"#{channel['name']}\",\n link=link,\n node_type=HierarchyNodeType.CHANNEL,\n external_access=channel_access,\n )\n\n\ndef _get_channel_by_id(client: WebClient, channel_id: str) -> ChannelType:\n \"\"\"Get a channel by its ID.\n\n Args:\n client: The Slack WebClient instance\n channel_id: The ID of the channel to fetch\n\n Returns:\n The channel information\n\n Raises:\n SlackApiError: If the channel cannot be fetched\n \"\"\"\n response = client.conversations_info(\n channel=channel_id,\n )\n return cast(ChannelType, response[\"channel\"])\n\n\ndef _get_messages(\n channel: ChannelType,\n client: WebClient,\n oldest: str | None = None,\n latest: str | None = None,\n limit: int = _SLACK_LIMIT,\n) -> tuple[list[MessageType], bool]:\n \"\"\"Slack goes from newest to oldest.\"\"\"\n\n # have to be in the channel in order to read messages\n if not channel[\"is_member\"]:\n try:\n client.conversations_join(\n channel=channel[\"id\"],\n is_private=channel[\"is_private\"],\n )\n except SlackApiError as e:\n if e.response[\"error\"] == \"is_archived\":\n logger.warning(f\"Channel {channel['name']} is archived. Skipping.\")\n return [], False\n\n logger.exception(f\"Error joining channel {channel['name']}\")\n raise\n logger.info(f\"Successfully joined '{channel['name']}'\")\n\n response = client.conversations_history(\n channel=channel[\"id\"],\n oldest=oldest,\n latest=latest,\n limit=limit,\n )\n response.validate()\n\n messages = cast(list[MessageType], response.get(\"messages\", []))\n\n cursor = cast(dict[str, Any], response.get(\"response_metadata\", {})).get(\n \"next_cursor\", \"\"\n )\n has_more = bool(cursor)\n return messages, has_more\n\n\ndef _message_to_doc(\n message: MessageType,\n client: WebClient,\n channel: ChannelType,\n slack_cleaner: SlackTextCleaner,\n user_cache: dict[str, BasicExpertInfo | None],\n seen_thread_ts: set[str],\n channel_access: ExternalAccess | None,\n msg_filter_func: Callable[\n [MessageType], SlackMessageFilterReason | None\n ] = default_msg_filter,\n) -> tuple[Document | None, SlackMessageFilterReason | None]:\n \"\"\"Returns a doc or None.\n If None is returned, the second element of the tuple may be a filter reason\n \"\"\"\n filtered_thread: ThreadType | None = None\n filter_reason: SlackMessageFilterReason | None = None\n thread_ts = message.get(\"thread_ts\")\n if thread_ts:\n # NOTE: if thread_ts is present, there's a thread we need to process\n # ... otherwise, we can skip it\n\n # skip threads we've already seen, since we've already processed all\n # messages in that thread\n if thread_ts in seen_thread_ts:\n return None, None\n\n thread = get_thread(\n client=client, channel_id=channel[\"id\"], thread_id=thread_ts\n )\n\n # we'll just set and use the last filter reason if\n # we bomb out later\n filtered_thread = []\n for message in thread:\n filter_reason = msg_filter_func(message)\n if filter_reason:\n continue\n\n filtered_thread.append(message)\n else:\n filter_reason = msg_filter_func(message)\n if filter_reason:\n return None, filter_reason\n\n filtered_thread = [message]\n\n # we'll just set and use the last filter reason if we get an empty list\n if not filtered_thread:\n return None, filter_reason\n\n doc = thread_to_doc(\n channel=channel,\n thread=filtered_thread,\n slack_cleaner=slack_cleaner,\n client=client,\n user_cache=user_cache,\n channel_access=channel_access,\n )\n return doc, None\n\n\ndef _get_all_doc_ids(\n client: WebClient,\n channels: list[str] | None = None,\n channel_name_regex_enabled: bool = False,\n msg_filter_func: Callable[\n [MessageType], SlackMessageFilterReason | None\n ] = default_msg_filter,\n callback: IndexingHeartbeatInterface | None = None,\n workspace_url: str | None = None,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n) -> GenerateSlimDocumentOutput:\n \"\"\"\n Get all document ids in the workspace, channel by channel\n This is pretty identical to get_all_docs, but it returns a set of ids instead of documents\n This makes it an order of magnitude faster than get_all_docs\n \"\"\"\n\n all_channels = get_channels(client)\n filtered_channels = filter_channels(\n all_channels, channels, channel_name_regex_enabled\n )\n user_cache: dict[str, BasicExpertInfo | None] = {}\n\n for channel in filtered_channels:\n channel_id = channel[\"id\"]\n # NOTE: external_access is a frozen object, so it's okay to safe to use a single\n # instance for all documents in the channel\n external_access = get_channel_access(\n client=client,\n channel=channel,\n user_cache=user_cache,\n )\n\n # Yield the channel as a HierarchyNode first (before any documents)\n yield [_channel_to_hierarchy_node(channel, external_access, workspace_url)]\n\n channel_message_batches = get_channel_messages(\n client=client,\n channel=channel,\n callback=callback,\n oldest=str(start) if start else None, # 0.0 -> None intentionally\n latest=str(end) if end is not None else None,\n )\n\n for message_batch in channel_message_batches:\n slim_doc_batch: list[SlimDocument | HierarchyNode] = []\n for message in message_batch:\n filter_reason = msg_filter_func(message)\n if filter_reason:\n continue\n\n # The document id is the channel id and the ts of the first message in the thread\n # Since we already have the first message of the thread, we dont have to\n # fetch the thread for id retrieval, saving time and API calls\n\n slim_doc_batch.append(\n SlimDocument(\n id=_build_doc_id(\n channel_id=channel_id, thread_ts=message[\"ts\"]\n ),\n external_access=external_access,\n parent_hierarchy_raw_node_id=channel_id,\n )\n )\n\n yield slim_doc_batch\n\n\nclass ProcessedSlackMessage(BaseModel):\n doc: Document | None\n # if the message is part of a thread, this is the thread_ts\n # otherwise, this is the message_ts. Either way, will be a unique identifier.\n # In the future, if the message becomes a thread, then the thread_ts\n # will be set to the message_ts.\n thread_or_message_ts: str\n\n # if doc is None, filter_reason may be populated\n filter_reason: SlackMessageFilterReason | None\n failure: ConnectorFailure | None\n\n\ndef _process_message(\n message: MessageType,\n client: WebClient,\n channel: ChannelType,\n slack_cleaner: SlackTextCleaner,\n user_cache: dict[str, BasicExpertInfo | None],\n seen_thread_ts: set[str],\n channel_access: ExternalAccess | None,\n msg_filter_func: Callable[\n [MessageType], SlackMessageFilterReason | None\n ] = default_msg_filter,\n) -> ProcessedSlackMessage:\n thread_ts = message.get(\"thread_ts\")\n thread_or_message_ts = thread_ts or message[\"ts\"]\n try:\n # causes random failures for testing checkpointing / continue on failure\n # import random\n # if random.random() > 0.95:\n # raise RuntimeError(\"Random failure :P\")\n\n doc, filter_reason = _message_to_doc(\n message=message,\n client=client,\n channel=channel,\n slack_cleaner=slack_cleaner,\n user_cache=user_cache,\n seen_thread_ts=seen_thread_ts,\n channel_access=channel_access,\n msg_filter_func=msg_filter_func,\n )\n return ProcessedSlackMessage(\n doc=doc,\n thread_or_message_ts=thread_or_message_ts,\n filter_reason=filter_reason,\n failure=None,\n )\n except Exception as e:\n logger.exception(f\"Error processing message {message['ts']}\")\n return ProcessedSlackMessage(\n doc=None,\n thread_or_message_ts=thread_or_message_ts,\n filter_reason=None,\n failure=ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=_build_doc_id(\n channel_id=channel[\"id\"], thread_ts=thread_or_message_ts\n ),\n document_link=get_message_link(message, client, channel[\"id\"]),\n ),\n failure_message=str(e),\n exception=e,\n ),\n )\n\n\nclass SlackConnector(\n SlimConnectorWithPermSync,\n CredentialsConnector,\n CheckpointedConnectorWithPermSync[SlackCheckpoint],\n):\n FAST_TIMEOUT = 1\n\n MAX_RETRIES = 7 # arbitrarily selected\n\n MAX_CHANNELS_TO_LOG = 50\n\n # *** values to use when filtering bot channels ***\n\n # the number of messages in the batch must be greater than or equal to this number\n # to consider filtering the channel\n BOT_CHANNEL_MIN_BATCH_SIZE = 256\n\n # the percentage of messages in the batch above which the channel will be considered\n # a bot channel\n BOT_CHANNEL_PERCENTAGE_THRESHOLD = 0.95\n\n def __init__(\n self,\n channels: list[str] | None = None,\n # if specified, will treat the specified channel strings as\n # regexes, and will only index channels that fully match the regexes\n channel_regex_enabled: bool = False,\n # if True, messages from bots/apps will be indexed instead of filtered out\n include_bot_messages: bool = False,\n batch_size: int = INDEX_BATCH_SIZE,\n num_threads: int = SLACK_NUM_THREADS,\n use_redis: bool = True,\n ) -> None:\n self.channels = channels\n self.channel_regex_enabled = channel_regex_enabled\n self.include_bot_messages = include_bot_messages\n self.msg_filter_func = (\n _bot_inclusive_msg_filter if include_bot_messages else default_msg_filter\n )\n self.batch_size = batch_size\n self.num_threads = num_threads\n self.client: WebClient | None = None\n self.fast_client: WebClient | None = None\n # just used for efficiency\n self.text_cleaner: SlackTextCleaner | None = None\n self.user_cache: dict[str, BasicExpertInfo | None] = {}\n self.credentials_provider: CredentialsProviderInterface | None = None\n self.credential_prefix: str | None = None\n self.use_redis: bool = use_redis\n # Workspace URL for building channel links (e.g., https://myworkspace.slack.com)\n self._workspace_url: str | None = None\n # self.delay_lock: str | None = None # the redis key for the shared lock\n # self.delay_key: str | None = None # the redis key for the shared delay\n\n @classmethod\n @override\n def normalize_url(cls, url: str) -> NormalizationResult:\n \"\"\"Normalize a Slack URL to extract channel_id__thread_ts format.\"\"\"\n parsed = urlparse(url)\n if \"slack.com\" not in parsed.netloc.lower():\n return NormalizationResult(normalized_url=None, use_default=False)\n\n # Slack document IDs are format: channel_id__thread_ts\n # Extract from URL pattern: .../archives/{channel_id}/p{timestamp}\n path_parts = parsed.path.split(\"/\")\n if \"archives\" not in path_parts:\n return NormalizationResult(normalized_url=None, use_default=False)\n\n archives_idx = path_parts.index(\"archives\")\n if archives_idx + 1 >= len(path_parts):\n return NormalizationResult(normalized_url=None, use_default=False)\n\n channel_id = path_parts[archives_idx + 1]\n if archives_idx + 2 >= len(path_parts):\n return NormalizationResult(normalized_url=None, use_default=False)\n\n thread_part = path_parts[archives_idx + 2]\n if not thread_part.startswith(\"p\"):\n return NormalizationResult(normalized_url=None, use_default=False)\n\n # Convert p1234567890123456 to 1234567890.123456 format\n timestamp_str = thread_part[1:] # Remove 'p' prefix\n if len(timestamp_str) == 16:\n # Insert dot at position 10 to match canonical format\n thread_ts = f\"{timestamp_str[:10]}.{timestamp_str[10:]}\"\n else:\n thread_ts = timestamp_str\n\n normalized = f\"{channel_id}__{thread_ts}\"\n return NormalizationResult(normalized_url=normalized, use_default=False)\n\n @staticmethod\n def make_credential_prefix(key: str) -> str:\n return f\"connector:slack:credential_{key}\"\n\n @staticmethod\n def make_delay_lock(prefix: str) -> str:\n return f\"{prefix}:delay_lock\"\n\n @staticmethod\n def make_delay_key(prefix: str) -> str:\n return f\"{prefix}:delay\"\n\n @staticmethod\n def make_slack_web_client(\n prefix: str, token: str, max_retry_count: int, r: Redis\n ) -> WebClient:\n delay_lock = SlackConnector.make_delay_lock(prefix)\n delay_key = SlackConnector.make_delay_key(prefix)\n\n # NOTE: slack has a built in RateLimitErrorRetryHandler, but it isn't designed\n # for concurrent workers. We've extended it with OnyxRedisSlackRetryHandler.\n connection_error_retry_handler = ConnectionErrorRetryHandler(\n max_retry_count=max_retry_count,\n interval_calculator=FixedValueRetryIntervalCalculator(),\n error_types=[\n URLError,\n ConnectionResetError,\n RemoteDisconnected,\n IncompleteRead,\n ],\n )\n\n onyx_rate_limit_error_retry_handler = OnyxRedisSlackRetryHandler(\n max_retry_count=max_retry_count,\n delay_key=delay_key,\n r=r,\n )\n custom_retry_handlers: list[RetryHandler] = [\n connection_error_retry_handler,\n onyx_rate_limit_error_retry_handler,\n ]\n\n client = OnyxSlackWebClient(\n delay_lock=delay_lock,\n delay_key=delay_key,\n r=r,\n token=token,\n retry_handlers=custom_retry_handlers,\n )\n return client\n\n @property\n def channels(self) -> list[str] | None:\n return self._channels\n\n @channels.setter\n def channels(self, channels: list[str] | None) -> None:\n self._channels = (\n [channel.removeprefix(\"#\") for channel in channels] if channels else None\n )\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n raise NotImplementedError(\"Use set_credentials_provider with this connector.\")\n\n def set_credentials_provider(\n self, credentials_provider: CredentialsProviderInterface\n ) -> None:\n credentials = credentials_provider.get_credentials()\n tenant_id = credentials_provider.get_tenant_id()\n if not tenant_id:\n raise ValueError(\"tenant_id cannot be None!\")\n\n bot_token = credentials[\"slack_bot_token\"]\n\n if self.use_redis:\n self.redis = get_redis_client(tenant_id=tenant_id)\n self.credential_prefix = SlackConnector.make_credential_prefix(\n credentials_provider.get_provider_key()\n )\n\n self.client = SlackConnector.make_slack_web_client(\n self.credential_prefix, bot_token, self.MAX_RETRIES, self.redis\n )\n else:\n connection_error_retry_handler = ConnectionErrorRetryHandler(\n max_retry_count=self.MAX_RETRIES,\n interval_calculator=FixedValueRetryIntervalCalculator(),\n error_types=[\n URLError,\n ConnectionResetError,\n RemoteDisconnected,\n IncompleteRead,\n ],\n )\n\n self.client = WebClient(\n token=bot_token, retry_handlers=[connection_error_retry_handler]\n )\n\n # use for requests that must return quickly (e.g. realtime flows where user is waiting)\n self.fast_client = WebClient(\n token=bot_token, timeout=SlackConnector.FAST_TIMEOUT\n )\n self.text_cleaner = SlackTextCleaner(client=self.client)\n self.credentials_provider = credentials_provider\n\n # Extract workspace URL from auth_test response for building channel links\n try:\n auth_response = self.client.auth_test()\n self._workspace_url = auth_response.get(\"url\")\n except Exception as e:\n logger.warning(f\"Failed to get workspace URL from auth_test: {e}\")\n self._workspace_url = None\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n if self.client is None:\n raise ConnectorMissingCredentialError(\"Slack\")\n\n return _get_all_doc_ids(\n client=self.client,\n channels=self.channels,\n channel_name_regex_enabled=self.channel_regex_enabled,\n msg_filter_func=self.msg_filter_func,\n callback=callback,\n workspace_url=self._workspace_url,\n start=start,\n end=end,\n )\n\n def _load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: SlackCheckpoint,\n include_permissions: bool = False,\n ) -> CheckpointOutput[SlackCheckpoint]:\n \"\"\"Rough outline:\n\n Step 1: Get all channels, yield back Checkpoint.\n Step 2: Loop through each channel. For each channel:\n Step 2.1: Get messages within the time range.\n Step 2.2: Process messages in parallel, yield back docs.\n Step 2.3: Update checkpoint with new_oldest, seen_thread_ts, and current_channel.\n Slack returns messages from newest to oldest, so we need to keep track of\n the latest message we've seen in each channel.\n Step 2.4: If there are no more messages in the channel, switch the current\n channel to the next channel.\n \"\"\"\n num_channels_remaining = 0\n\n if self.client is None or self.text_cleaner is None:\n raise ConnectorMissingCredentialError(\"Slack\")\n\n checkpoint = cast(SlackCheckpoint, copy.deepcopy(checkpoint))\n\n # if this is the very first time we've called this, need to\n # get all relevant channels and save them into the checkpoint\n if checkpoint.channel_ids is None:\n raw_channels = get_channels(self.client)\n filtered_channels = filter_channels(\n raw_channels, self.channels, self.channel_regex_enabled\n )\n logger.info(\n f\"Channels - initial checkpoint: all={len(raw_channels)} post_filtering={len(filtered_channels)}\"\n )\n\n checkpoint.channel_ids = [c[\"id\"] for c in filtered_channels]\n if len(filtered_channels) == 0:\n checkpoint.has_more = False\n return checkpoint\n\n checkpoint.current_channel = filtered_channels[0]\n if include_permissions:\n # checkpoint.current_channel is guaranteed to be non-None here since we just assigned it\n assert checkpoint.current_channel is not None\n channel_access = get_channel_access(\n client=self.client,\n channel=checkpoint.current_channel,\n user_cache=self.user_cache,\n )\n checkpoint.current_channel_access = channel_access\n checkpoint.has_more = True\n return checkpoint\n\n final_channel_ids = checkpoint.channel_ids\n for channel_id in final_channel_ids:\n if channel_id not in checkpoint.channel_completion_map:\n num_channels_remaining += 1\n\n logger.info(\n f\"Channels - current status: \"\n f\"processed={len(final_channel_ids) - num_channels_remaining} \"\n f\"remaining={num_channels_remaining} \"\n f\"total={len(final_channel_ids)}\"\n )\n\n channel = checkpoint.current_channel\n if channel is None:\n raise ValueError(\"current_channel key not set in checkpoint\")\n\n channel_id = channel[\"id\"]\n if channel_id not in final_channel_ids:\n raise ValueError(f\"Channel {channel_id} not found in checkpoint\")\n\n channel_created = channel[\"created\"]\n\n seen_thread_ts = set(checkpoint.seen_thread_ts)\n\n try:\n num_bot_filtered_messages = 0\n num_other_filtered_messages = 0\n\n oldest = str(start) if start else None\n latest = str(end)\n\n channel_message_ts = checkpoint.channel_completion_map.get(channel_id)\n if channel_message_ts:\n # Set oldest to the checkpoint timestamp to resume from where we left off\n oldest = channel_message_ts\n else:\n # First time processing this channel - yield its hierarchy node\n yield _channel_to_hierarchy_node(\n channel,\n checkpoint.current_channel_access,\n self._workspace_url,\n )\n\n logger.debug(\n f\"Getting messages for channel {channel} within range {oldest} - {latest}\"\n )\n\n message_batch, has_more_in_channel = _get_messages(\n channel, self.client, oldest, latest\n )\n\n logger.info(\n f\"Retrieved messages: {len(message_batch)=} {channel=} {oldest=} {latest=}\"\n )\n\n # message_batch[0] is the newest message (Slack returns newest to oldest)\n new_oldest = message_batch[0][\"ts\"] if message_batch else latest\n\n num_threads_start = len(seen_thread_ts)\n\n # Process messages in parallel using ThreadPoolExecutor\n with ThreadPoolExecutor(max_workers=self.num_threads) as executor:\n # NOTE(rkuo): this seems to be assuming the slack sdk is thread safe.\n # That's a very bold assumption! Haven't seen a direct issue with this\n # yet, but likely not correct to rely on.\n\n futures: list[Future[ProcessedSlackMessage]] = []\n for message in message_batch:\n # Capture the current context so that the thread gets the current tenant ID\n current_context = contextvars.copy_context()\n futures.append(\n executor.submit(\n current_context.run,\n _process_message,\n message=message,\n client=self.client,\n channel=channel,\n slack_cleaner=self.text_cleaner,\n user_cache=self.user_cache,\n seen_thread_ts=seen_thread_ts,\n channel_access=checkpoint.current_channel_access,\n msg_filter_func=self.msg_filter_func,\n )\n )\n\n for future in as_completed(futures):\n processed_slack_message = future.result()\n doc = processed_slack_message.doc\n thread_or_message_ts = processed_slack_message.thread_or_message_ts\n failure = processed_slack_message.failure\n if doc:\n # handle race conditions here since this is single\n # threaded. Multi-threaded _process_message reads from this\n # but since this is single threaded, we won't run into simul\n # writes. At worst, we can duplicate a thread, which will be\n # deduped later on.\n if thread_or_message_ts not in seen_thread_ts:\n yield doc\n\n seen_thread_ts.add(thread_or_message_ts)\n elif processed_slack_message.filter_reason:\n if (\n processed_slack_message.filter_reason\n == SlackMessageFilterReason.BOT\n ):\n num_bot_filtered_messages += 1\n else:\n num_other_filtered_messages += 1\n elif failure:\n yield failure\n\n num_threads_processed = len(seen_thread_ts) - num_threads_start\n\n # calculate a percentage progress for the current channel by determining\n # how much of the time range we've processed so far\n new_oldest_seconds_epoch = SecondsSinceUnixEpoch(new_oldest)\n range_start = start if start else max(0, channel_created)\n if new_oldest_seconds_epoch < range_start:\n range_complete = 0.0\n else:\n range_complete = new_oldest_seconds_epoch - range_start\n\n range_total = end - range_start\n if range_total <= 0:\n range_total = 1\n range_percent_complete = range_complete / range_total * 100.0\n\n num_filtered = num_bot_filtered_messages + num_other_filtered_messages\n log_func = logger.warning if num_bot_filtered_messages > 0 else logger.info\n log_func(\n f\"Message processing stats: \"\n f\"batch_len={len(message_batch)} \"\n f\"batch_yielded={num_threads_processed} \"\n f\"filtered={num_filtered} \"\n f\"(bot={num_bot_filtered_messages} other={num_other_filtered_messages}) \"\n f\"total_threads_seen={len(seen_thread_ts)}\"\n )\n\n logger.info(\n f\"Current channel processing stats: {range_start=} range_end={end} percent_complete={range_percent_complete=:.2f}\"\n )\n\n checkpoint.seen_thread_ts = list(seen_thread_ts)\n checkpoint.channel_completion_map[channel[\"id\"]] = new_oldest\n\n # bypass channels where the first set of messages seen are all\n # filtered (bots + disallowed subtypes like channel_join)\n # check at least MIN_BOT_MESSAGE_THRESHOLD messages are in the batch\n # we shouldn't skip based on a small sampling of messages\n if (\n channel_message_ts is None\n and len(message_batch) > SlackConnector.BOT_CHANNEL_MIN_BATCH_SIZE\n ):\n if (\n num_filtered\n > SlackConnector.BOT_CHANNEL_PERCENTAGE_THRESHOLD\n * len(message_batch)\n ):\n logger.warning(\n \"Bypassing this channel since it appears to be mostly bot messages\"\n )\n has_more_in_channel = False\n\n if not has_more_in_channel:\n num_channels_remaining -= 1\n\n new_channel_id = next(\n (\n channel_id\n for channel_id in final_channel_ids\n if channel_id not in checkpoint.channel_completion_map\n ),\n None,\n )\n\n if new_channel_id:\n new_channel = _get_channel_by_id(self.client, new_channel_id)\n checkpoint.current_channel = new_channel\n if include_permissions:\n channel_access = get_channel_access(\n client=self.client,\n channel=new_channel,\n user_cache=self.user_cache,\n )\n checkpoint.current_channel_access = channel_access\n else:\n checkpoint.current_channel = None\n\n checkpoint.has_more = checkpoint.current_channel is not None\n\n channels_processed = len(final_channel_ids) - num_channels_remaining\n channels_percent_complete = (\n channels_processed / len(final_channel_ids) * 100.0\n )\n logger.info(\n f\"All channels processing stats: \"\n f\"processed={len(final_channel_ids) - num_channels_remaining} \"\n f\"remaining={num_channels_remaining} \"\n f\"total={len(final_channel_ids)} \"\n f\"percent_complete={channels_percent_complete:.2f}\"\n )\n except Exception as e:\n logger.exception(f\"Error processing channel {channel['name']}\")\n yield ConnectorFailure(\n failed_entity=EntityFailure(\n entity_id=channel[\"id\"],\n missed_time_range=(\n datetime.fromtimestamp(start, tz=timezone.utc),\n datetime.fromtimestamp(end, tz=timezone.utc),\n ),\n ),\n failure_message=str(e),\n exception=e,\n )\n\n return checkpoint\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: SlackCheckpoint,\n ) -> CheckpointOutput[SlackCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=False\n )\n\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: SlackCheckpoint,\n ) -> CheckpointOutput[SlackCheckpoint]:\n return self._load_from_checkpoint(\n start, end, checkpoint, include_permissions=True\n )\n\n def validate_connector_settings(self) -> None:\n \"\"\"\n 1. Verify the bot token is valid for the workspace (via auth_test).\n 2. Ensure the bot has enough scope to list channels.\n 3. Check that every channel specified in self.channels exists (only when regex is not enabled).\n \"\"\"\n if self.fast_client is None:\n raise ConnectorMissingCredentialError(\"Slack credentials not loaded.\")\n\n try:\n # 1) Validate connection to workspace\n auth_response = self.fast_client.auth_test()\n if not auth_response.get(\"ok\", False):\n error_msg = auth_response.get(\n \"error\", \"Unknown error from Slack auth_test\"\n )\n raise ConnectorValidationError(f\"Failed Slack auth_test: {error_msg}\")\n\n # 2) Minimal test to confirm listing channels works\n test_resp = self.fast_client.conversations_list(\n limit=1, types=[\"public_channel\"]\n )\n if not test_resp.get(\"ok\", False):\n error_msg = test_resp.get(\"error\", \"Unknown error from Slack\")\n if error_msg == \"invalid_auth\":\n raise ConnectorValidationError(\n f\"Invalid Slack bot token ({error_msg}).\"\n )\n elif error_msg == \"not_authed\":\n raise CredentialExpiredError(\n f\"Invalid or expired Slack bot token ({error_msg}).\"\n )\n raise UnexpectedValidationError(\n f\"Slack API returned a failure: {error_msg}\"\n )\n\n # 3) If channels are specified and regex is not enabled, verify each is accessible\n # NOTE: removed this for now since it may be too slow for large workspaces which may\n # have some automations which create a lot of channels (100k+)\n\n # if self.channels and not self.channel_regex_enabled:\n # accessible_channels = get_channels(\n # client=self.fast_client,\n # exclude_archived=True,\n # get_public=True,\n # get_private=True,\n # )\n # # For quick lookups by name or ID, build a map:\n # accessible_channel_names = {ch[\"name\"] for ch in accessible_channels}\n # accessible_channel_ids = {ch[\"id\"] for ch in accessible_channels}\n\n # for user_channel in self.channels:\n # if (\n # user_channel not in accessible_channel_names\n # and user_channel not in accessible_channel_ids\n # ):\n # raise ConnectorValidationError(\n # f\"Channel '{user_channel}' not found or inaccessible in this workspace.\"\n # )\n\n except SlackApiError as e:\n slack_error = e.response.get(\"error\", \"\")\n if slack_error == \"ratelimited\":\n # Handle rate limiting specifically\n retry_after = int(e.response.headers.get(\"Retry-After\", 1))\n logger.warning(\n f\"Slack API rate limited during validation. Retry suggested after {retry_after} seconds. \"\n \"Proceeding with validation, but be aware that connector operations might be throttled.\"\n )\n # Continue validation without failing - the connector is likely valid but just rate limited\n return\n elif slack_error == \"missing_scope\":\n raise InsufficientPermissionsError(\n \"Slack bot token lacks the necessary scope to list/access channels. \"\n \"Please ensure your Slack app has 'channels:read' (and/or 'groups:read' for private channels).\"\n )\n elif slack_error == \"invalid_auth\":\n raise CredentialExpiredError(\n f\"Invalid Slack bot token ({slack_error}).\"\n )\n elif slack_error == \"not_authed\":\n raise CredentialExpiredError(\n f\"Invalid or expired Slack bot token ({slack_error}).\"\n )\n raise UnexpectedValidationError(\n f\"Unexpected Slack error '{slack_error}' during settings validation.\"\n )\n except ConnectorValidationError as e:\n raise e\n except Exception as e:\n raise UnexpectedValidationError(\n f\"Unexpected error during Slack settings validation: {e}\"\n )\n\n @override\n def build_dummy_checkpoint(self) -> SlackCheckpoint:\n return SlackCheckpoint(\n channel_ids=None,\n channel_completion_map={},\n current_channel=None,\n current_channel_access=None,\n seen_thread_ts=[],\n has_more=True,\n )\n\n @override\n def validate_checkpoint_json(self, checkpoint_json: str) -> SlackCheckpoint:\n return SlackCheckpoint.model_validate_json(checkpoint_json)\n\n\nif __name__ == \"__main__\":\n import os\n import time\n from onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\n from shared_configs.contextvars import get_current_tenant_id\n\n slack_channel = os.environ.get(\"SLACK_CHANNEL\")\n connector = SlackConnector(\n channels=[slack_channel] if slack_channel else None,\n )\n\n provider = OnyxStaticCredentialsProvider(\n tenant_id=get_current_tenant_id(),\n connector_name=\"slack\",\n credential_json={\n \"slack_bot_token\": os.environ[\"SLACK_BOT_TOKEN\"],\n },\n )\n connector.set_credentials_provider(provider)\n\n current = time.time()\n one_day_ago = current - 24 * 60 * 60 # 1 day\n\n checkpoint = connector.build_dummy_checkpoint()\n\n gen = connector.load_from_checkpoint(\n one_day_ago,\n current,\n cast(SlackCheckpoint, checkpoint),\n )\n try:\n for document_or_failure in gen:\n if isinstance(document_or_failure, Document):\n print(document_or_failure)\n elif isinstance(document_or_failure, ConnectorFailure):\n print(document_or_failure)\n except StopIteration as e:\n checkpoint = e.value\n print(\"Next checkpoint:\", checkpoint)\n\n print(\"Next checkpoint:\", checkpoint)\n" }, { "path": "backend/onyx/connectors/slack/models.py", "content": "from typing import NotRequired\n\nfrom typing_extensions import TypedDict\n\n\nclass ChannelTopicPurposeType(TypedDict):\n \"\"\"\n Represents the topic or purpose of a Slack channel.\n \"\"\"\n\n value: str\n creator: str\n last_set: int\n\n\nclass ChannelType(TypedDict):\n \"\"\"\n Represents a Slack channel.\n \"\"\"\n\n id: str\n name: str\n is_channel: bool\n is_group: bool\n is_im: bool\n created: int\n creator: str\n is_archived: bool\n is_general: bool\n unlinked: int\n name_normalized: str\n is_shared: bool\n is_ext_shared: bool\n is_org_shared: bool\n pending_shared: list[str]\n is_pending_ext_shared: bool\n is_member: bool\n is_private: bool\n is_mpim: bool\n updated: int\n topic: ChannelTopicPurposeType\n purpose: ChannelTopicPurposeType\n previous_names: list[str]\n num_members: int\n\n\nclass AttachmentType(TypedDict):\n \"\"\"\n Represents a Slack message attachment.\n \"\"\"\n\n service_name: NotRequired[str]\n text: NotRequired[str]\n fallback: NotRequired[str]\n thumb_url: NotRequired[str]\n thumb_width: NotRequired[int]\n thumb_height: NotRequired[int]\n id: NotRequired[int]\n\n\nclass BotProfileType(TypedDict):\n \"\"\"\n Represents a Slack bot profile.\n \"\"\"\n\n id: NotRequired[str]\n deleted: NotRequired[bool]\n name: NotRequired[str]\n updated: NotRequired[int]\n app_id: NotRequired[str]\n team_id: NotRequired[str]\n\n\nclass MessageType(TypedDict):\n \"\"\"\n Represents a Slack message.\n \"\"\"\n\n type: str\n user: str\n text: str\n ts: str\n attachments: NotRequired[list[AttachmentType]]\n # Bot-related fields\n bot_id: NotRequired[str]\n app_id: NotRequired[str]\n bot_profile: NotRequired[BotProfileType]\n # Message threading\n thread_ts: NotRequired[str]\n # Message subtype (for filtering certain message types)\n subtype: NotRequired[str]\n\n\n# list of messages in a thread\nThreadType = list[MessageType]\n" }, { "path": "backend/onyx/connectors/slack/onyx_retry_handler.py", "content": "import random\nfrom typing import cast\nfrom typing import Optional\n\nfrom redis import Redis\nfrom slack_sdk.http_retry.handler import RetryHandler\nfrom slack_sdk.http_retry.request import HttpRequest\nfrom slack_sdk.http_retry.response import HttpResponse\nfrom slack_sdk.http_retry.state import RetryState\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass OnyxRedisSlackRetryHandler(RetryHandler):\n \"\"\"\n This class uses Redis to share a rate limit among multiple threads.\n\n As currently implemented, this code is already surrounded by a lock in Redis\n via an override of _perform_urllib_http_request in OnyxSlackWebClient.\n\n This just sets the desired retry delay with TTL in redis. In conjunction with\n a custom subclass of the client, the value is read and obeyed prior to an API call\n and also serialized.\n\n Another way to do this is just to do exponential backoff. Might be easier?\n\n Adapted from slack's RateLimitErrorRetryHandler.\n \"\"\"\n\n \"\"\"RetryHandler that does retries for rate limited errors.\"\"\"\n\n def __init__(\n self,\n max_retry_count: int,\n delay_key: str,\n r: Redis,\n ):\n \"\"\"\n delay_lock: the redis key to use with RedisLock (to synchronize access to delay_key)\n delay_key: the redis key containing a shared TTL\n \"\"\"\n super().__init__(max_retry_count=max_retry_count)\n self._redis: Redis = r\n self._delay_key = delay_key\n\n def _can_retry(\n self,\n *,\n state: RetryState, # noqa: ARG002\n request: HttpRequest, # noqa: ARG002\n response: Optional[HttpResponse] = None,\n error: Optional[Exception] = None, # noqa: ARG002\n ) -> bool:\n return response is not None and response.status_code == 429\n\n def prepare_for_next_attempt(\n self,\n *,\n state: RetryState,\n request: HttpRequest, # noqa: ARG002\n response: Optional[HttpResponse] = None,\n error: Optional[Exception] = None,\n ) -> None:\n \"\"\"As initially designed by the SDK authors, this function is responsible for\n the wait to retry ... aka we actually sleep in this function.\n\n This doesn't work well with multiple clients because every thread is unaware\n of the current retry value until it actually calls the endpoint.\n\n We're combining this with an actual subclass of the slack web client so\n that the delay is used BEFORE calling an API endpoint. The subclassed client\n has already taken the lock in redis when this method is called.\n \"\"\"\n ttl_ms: int | None = None\n\n retry_after_value: str | None = None\n retry_after_header_name: Optional[str] = None\n duration_s: float = 1.0 # seconds\n\n if response is None:\n # NOTE(rkuo): this logic comes from RateLimitErrorRetryHandler.\n # This reads oddly, as if the caller itself could raise the exception.\n # We don't have the luxury of changing this.\n if error:\n raise error\n\n return\n\n state.next_attempt_requested = True # this signals the caller to retry\n\n # calculate wait duration based on retry-after + some jitter\n for k in response.headers.keys():\n if k.lower() == \"retry-after\":\n retry_after_header_name = k\n break\n\n try:\n if retry_after_header_name is None:\n # This situation usually does not arise. Just in case.\n raise ValueError(\n \"OnyxRedisSlackRetryHandler.prepare_for_next_attempt: retry-after header name is None\"\n )\n\n retry_after_header_value = response.headers.get(retry_after_header_name)\n if not retry_after_header_value:\n raise ValueError(\n \"OnyxRedisSlackRetryHandler.prepare_for_next_attempt: retry-after header value is None\"\n )\n\n # Handle case where header value might be a list\n retry_after_value = (\n retry_after_header_value[0]\n if isinstance(retry_after_header_value, list)\n else retry_after_header_value\n )\n\n retry_after_value_int = int(\n retry_after_value\n ) # will raise ValueError if somehow we can't convert to int\n jitter = retry_after_value_int * 0.25 * random.random()\n duration_s = retry_after_value_int + jitter\n except ValueError:\n duration_s += random.random()\n\n # Read and extend the ttl\n ttl_ms = cast(int, self._redis.pttl(self._delay_key))\n if ttl_ms < 0: # negative values are error status codes ... see docs\n ttl_ms = 0\n ttl_ms_new = ttl_ms + int(duration_s * 1000.0)\n self._redis.set(self._delay_key, \"1\", px=ttl_ms_new)\n\n logger.warning(\n f\"OnyxRedisSlackRetryHandler.prepare_for_next_attempt setting delay: \"\n f\"current_attempt={state.current_attempt} \"\n f\"retry-after={retry_after_value} \"\n f\"{ttl_ms_new=}\"\n )\n\n state.increment_current_attempt()\n" }, { "path": "backend/onyx/connectors/slack/onyx_slack_web_client.py", "content": "import threading\nimport time\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\nfrom urllib.request import Request\n\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom slack_sdk import WebClient\n\nfrom onyx.connectors.slack.utils import ONYX_SLACK_LOCK_BLOCKING_TIMEOUT\nfrom onyx.connectors.slack.utils import ONYX_SLACK_LOCK_TOTAL_BLOCKING_TIMEOUT\nfrom onyx.connectors.slack.utils import ONYX_SLACK_LOCK_TTL\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass OnyxSlackWebClient(WebClient):\n \"\"\"Use in combination with the Onyx Retry Handler.\n\n This client wrapper enforces a proper retry delay through redis BEFORE the api call\n so that multiple clients can synchronize and rate limit properly.\n\n The retry handler writes the correct delay value to redis so that it is can be used\n by this wrapper.\n\n \"\"\"\n\n def __init__(\n self, delay_lock: str, delay_key: str, r: Redis, *args: Any, **kwargs: Any\n ) -> None:\n super().__init__(*args, **kwargs)\n self._delay_key = delay_key\n self._delay_lock = delay_lock\n self._redis: Redis = r\n self.num_requests: int = 0\n self._lock = threading.Lock()\n\n def _perform_urllib_http_request(\n self, *, url: str, args: Dict[str, Dict[str, Any]]\n ) -> Dict[str, Any]:\n \"\"\"By locking around the base class method, we ensure that both the delay from\n Redis and parsing/writing of retry values to Redis are handled properly in\n one place\"\"\"\n # lock and extend the ttl\n lock: RedisLock = self._redis.lock(\n self._delay_lock,\n timeout=ONYX_SLACK_LOCK_TTL,\n )\n\n # try to acquire the lock\n start = time.monotonic()\n while True:\n acquired = lock.acquire(blocking_timeout=ONYX_SLACK_LOCK_BLOCKING_TIMEOUT)\n if acquired:\n break\n\n # if we couldn't acquire the lock but it exists, there's at least some activity\n # so keep trying...\n if self._redis.exists(self._delay_lock):\n continue\n\n if time.monotonic() - start > ONYX_SLACK_LOCK_TOTAL_BLOCKING_TIMEOUT:\n raise RuntimeError(\n f\"OnyxSlackWebClient._perform_urllib_http_request - \"\n f\"timed out waiting for lock: {ONYX_SLACK_LOCK_TOTAL_BLOCKING_TIMEOUT=}\"\n )\n\n try:\n result = super()._perform_urllib_http_request(url=url, args=args)\n finally:\n if lock.owned():\n lock.release()\n else:\n logger.warning(\n \"OnyxSlackWebClient._perform_urllib_http_request lock not owned on release\"\n )\n\n time.monotonic() - start\n # logger.info(\n # f\"OnyxSlackWebClient._perform_urllib_http_request: Releasing lock: {elapsed=}\"\n # )\n\n return result\n\n def _perform_urllib_http_request_internal(\n self,\n url: str,\n req: Request,\n ) -> Dict[str, Any]:\n \"\"\"Overrides the internal method which is mostly the direct call to\n urllib/urlopen ... so this is a good place to perform our delay.\"\"\"\n\n # read and execute the delay\n delay_ms = cast(int, self._redis.pttl(self._delay_key))\n if delay_ms < 0: # negative values are error status codes ... see docs\n delay_ms = 0\n\n if delay_ms > 0:\n logger.warning(\n f\"OnyxSlackWebClient._perform_urllib_http_request_internal delay: {delay_ms=} {self.num_requests=}\"\n )\n\n time.sleep(delay_ms / 1000.0)\n\n result = super()._perform_urllib_http_request_internal(url, req)\n\n with self._lock:\n self.num_requests += 1\n\n # the delay key should have naturally expired by this point\n return result\n" }, { "path": "backend/onyx/connectors/slack/utils.py", "content": "import re\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom functools import lru_cache\nfrom functools import wraps\nfrom typing import Any\nfrom typing import cast\n\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\nfrom slack_sdk.web import SlackResponse\n\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.slack.models import MessageType\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.retry_wrapper import retry_builder\n\nlogger = setup_logger()\n\n# retry after 0.1, 1.2, 3.4, 7.8, 16.6, 34.2 seconds\nbasic_retry_wrapper = retry_builder(tries=7)\n# number of messages we request per page when fetching paginated slack messages\n_SLACK_LIMIT = 900\n\n# used to serialize access to the retry TTL\nONYX_SLACK_LOCK_TTL = 1800 # how long the lock is allowed to idle before it expires\nONYX_SLACK_LOCK_BLOCKING_TIMEOUT = 60 # how long to wait for the lock per wait attempt\nONYX_SLACK_LOCK_TOTAL_BLOCKING_TIMEOUT = 3600 # how long to wait for the lock in total\n\n\n@lru_cache()\ndef get_base_url(token: str) -> str:\n \"\"\"Retrieve and cache the base URL of the Slack workspace based on the client token.\"\"\"\n client = WebClient(token=token)\n return client.auth_test()[\"url\"]\n\n\ndef get_message_link(event: MessageType, client: WebClient, channel_id: str) -> str:\n message_ts = event[\"ts\"]\n message_ts_without_dot = message_ts.replace(\".\", \"\")\n thread_ts = event.get(\"thread_ts\")\n base_url = get_base_url(client.token)\n\n link = f\"{base_url.rstrip('/')}/archives/{channel_id}/p{message_ts_without_dot}\" + (\n f\"?thread_ts={thread_ts}\" if thread_ts else \"\"\n )\n return link\n\n\ndef make_slack_api_call(\n call: Callable[..., SlackResponse], **kwargs: Any\n) -> SlackResponse:\n return call(**kwargs)\n\n\ndef make_paginated_slack_api_call(\n call: Callable[..., SlackResponse], **kwargs: Any\n) -> Generator[dict[str, Any], None, None]:\n return _make_slack_api_call_paginated(call)(**kwargs)\n\n\ndef _make_slack_api_call_paginated(\n call: Callable[..., SlackResponse],\n) -> Callable[..., Generator[dict[str, Any], None, None]]:\n \"\"\"Wraps calls to slack API so that they automatically handle pagination\"\"\"\n\n @wraps(call)\n def paginated_call(**kwargs: Any) -> Generator[dict[str, Any], None, None]:\n cursor: str | None = None\n has_more = True\n while has_more:\n response = call(cursor=cursor, limit=_SLACK_LIMIT, **kwargs)\n yield cast(dict[str, Any], response.validate())\n cursor = cast(dict[str, Any], response.get(\"response_metadata\", {})).get(\n \"next_cursor\", \"\"\n )\n has_more = bool(cursor)\n\n return paginated_call\n\n\n# NOTE(rkuo): we may not need this any more if the integrated retry handlers work as\n# expected. Do we want to keep this around?\n\n# def make_slack_api_rate_limited(\n# call: Callable[..., SlackResponse], max_retries: int = 7\n# ) -> Callable[..., SlackResponse]:\n# \"\"\"Wraps calls to slack API so that they automatically handle rate limiting\"\"\"\n\n# @wraps(call)\n# def rate_limited_call(**kwargs: Any) -> SlackResponse:\n# last_exception = None\n\n# for _ in range(max_retries):\n# try:\n# # Make the API call\n# response = call(**kwargs)\n\n# # Check for errors in the response, will raise `SlackApiError`\n# # if anything went wrong\n# response.validate()\n# return response\n\n# except SlackApiError as e:\n# last_exception = e\n# try:\n# error = e.response[\"error\"]\n# except KeyError:\n# error = \"unknown error\"\n\n# if error == \"ratelimited\":\n# # Handle rate limiting: get the 'Retry-After' header value and sleep for that duration\n# retry_after = int(e.response.headers.get(\"Retry-After\", 1))\n# logger.info(\n# f\"Slack call rate limited, retrying after {retry_after} seconds. Exception: {e}\"\n# )\n# time.sleep(retry_after)\n# elif error in [\"already_reacted\", \"no_reaction\", \"internal_error\"]:\n# # Log internal_error and return the response instead of failing\n# logger.warning(\n# f\"Slack call encountered '{error}', skipping and continuing...\"\n# )\n# return e.response\n# else:\n# # Raise the error for non-transient errors\n# raise\n\n# # If the code reaches this point, all retries have been exhausted\n# msg = f\"Max retries ({max_retries}) exceeded\"\n# if last_exception:\n# raise Exception(msg) from last_exception\n# else:\n# raise Exception(msg)\n\n# return rate_limited_call\n\n# temporarily disabling due to using a different retry approach\n# might be permanent if everything works out\n# def make_slack_api_call_w_retries(\n# call: Callable[..., SlackResponse], **kwargs: Any\n# ) -> SlackResponse:\n# return basic_retry_wrapper(call)(**kwargs)\n\n\n# def make_paginated_slack_api_call_w_retries(\n# call: Callable[..., SlackResponse], **kwargs: Any\n# ) -> Generator[dict[str, Any], None, None]:\n# return _make_slack_api_call_paginated(basic_retry_wrapper(call))(**kwargs)\n\n\ndef expert_info_from_slack_id(\n user_id: str | None,\n client: WebClient,\n user_cache: dict[str, BasicExpertInfo | None],\n) -> BasicExpertInfo | None:\n if not user_id:\n return None\n\n if user_id in user_cache:\n return user_cache[user_id]\n\n response = client.users_info(user=user_id)\n\n if not response[\"ok\"]:\n user_cache[user_id] = None\n return None\n\n user: dict = cast(dict[Any, dict], response.data).get(\"user\", {})\n profile = user.get(\"profile\", {})\n\n expert = BasicExpertInfo(\n display_name=user.get(\"real_name\") or profile.get(\"display_name\"),\n first_name=profile.get(\"first_name\"),\n last_name=profile.get(\"last_name\"),\n email=profile.get(\"email\"),\n )\n\n user_cache[user_id] = expert\n\n return expert\n\n\nclass SlackTextCleaner:\n \"\"\"Utility class to replace user IDs with usernames in a message.\n Handles caching, so the same request is not made multiple times\n for the same user ID\"\"\"\n\n def __init__(self, client: WebClient) -> None:\n self._client = client\n self._id_to_name_map: dict[str, str] = {}\n\n def _get_slack_name(self, user_id: str) -> str:\n if user_id not in self._id_to_name_map:\n try:\n response = self._client.users_info(user=user_id)\n # prefer display name if set, since that is what is shown in Slack\n self._id_to_name_map[user_id] = (\n response[\"user\"][\"profile\"][\"display_name\"]\n or response[\"user\"][\"profile\"][\"real_name\"]\n )\n except SlackApiError as e:\n logger.exception(\n f\"Error fetching data for user {user_id}: {e.response['error']}\"\n )\n raise\n\n return self._id_to_name_map[user_id]\n\n def _replace_user_ids_with_names(self, message: str) -> str:\n # Find user IDs in the message\n user_ids = re.findall(\"<@(.*?)>\", message)\n\n # Iterate over each user ID found\n for user_id in user_ids:\n try:\n if user_id in self._id_to_name_map:\n user_name = self._id_to_name_map[user_id]\n else:\n user_name = self._get_slack_name(user_id)\n\n # Replace the user ID with the username in the message\n message = message.replace(f\"<@{user_id}>\", f\"@{user_name}\")\n except Exception:\n logger.exception(\n f\"Unable to replace user ID with username for user_id '{user_id}'\"\n )\n\n return message\n\n def index_clean(self, message: str) -> str:\n \"\"\"During indexing, replace pattern sets that may cause confusion to the model\n Some special patterns are left in as they can provide information\n ie. links that contain format text|link, both the text and the link may be informative\n \"\"\"\n message = self._replace_user_ids_with_names(message)\n message = self.replace_tags_basic(message)\n message = self.replace_channels_basic(message)\n message = self.replace_special_mentions(message)\n message = self.replace_special_catchall(message)\n return message\n\n @staticmethod\n def replace_tags_basic(message: str) -> str:\n \"\"\"Simply replaces all tags with `@` in order to prevent us from\n tagging users in Slack when we don't want to\"\"\"\n # Find user IDs in the message\n user_ids = re.findall(\"<@(.*?)>\", message)\n for user_id in user_ids:\n message = message.replace(f\"<@{user_id}>\", f\"@{user_id}\")\n return message\n\n @staticmethod\n def replace_channels_basic(message: str) -> str:\n \"\"\"Simply replaces all channel mentions with `#` in order\n to make a message work as part of a link\"\"\"\n # Find user IDs in the message\n channel_matches = re.findall(r\"<#(.*?)\\|(.*?)>\", message)\n for channel_id, channel_name in channel_matches:\n message = message.replace(\n f\"<#{channel_id}|{channel_name}>\", f\"#{channel_name}\"\n )\n return message\n\n @staticmethod\n def replace_special_mentions(message: str) -> str:\n \"\"\"Simply replaces @channel, @here, and @everyone so we don't tag\n a bunch of people in Slack when we don't want to\"\"\"\n # Find user IDs in the message\n message = message.replace(\"\", \"@channel\")\n message = message.replace(\"\", \"@here\")\n message = message.replace(\"\", \"@everyone\")\n return message\n\n @staticmethod\n def replace_special_catchall(message: str) -> str:\n \"\"\"Replaces pattern of with another-thing\n This is added for but may match other cases as well\n \"\"\"\n\n pattern = r\"]+)>\"\n return re.sub(pattern, r\"\\2\", message)\n\n @staticmethod\n def add_zero_width_whitespace_after_tag(message: str) -> str:\n \"\"\"Add a 0 width whitespace after every @\"\"\"\n return message.replace(\"@\", \"@\\u200b\")\n" }, { "path": "backend/onyx/connectors/teams/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/teams/connector.py", "content": "import copy\nimport os\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nimport msal # type: ignore\nfrom office365.graph_client import GraphClient # type: ignore\nfrom office365.runtime.client_request_exception import ClientRequestException # type: ignore\nfrom office365.runtime.http.request_options import RequestOptions # type: ignore[import-untyped]\nfrom office365.teams.channels.channel import Channel # type: ignore\nfrom office365.teams.team import Team # type: ignore\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.microsoft_graph_env import resolve_microsoft_environment\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import EntityFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.teams.models import Message\nfrom onyx.connectors.teams.utils import fetch_expert_infos\nfrom onyx.connectors.teams.utils import fetch_external_access\nfrom onyx.connectors.teams.utils import fetch_messages\nfrom onyx.connectors.teams.utils import fetch_replies\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_with_timeout\n\nlogger = setup_logger()\n\n_SLIM_DOC_BATCH_SIZE = 5000\n\n\nclass TeamsCheckpoint(ConnectorCheckpoint):\n todo_team_ids: list[str] | None = None\n\n\nDEFAULT_AUTHORITY_HOST = \"https://login.microsoftonline.com\"\nDEFAULT_GRAPH_API_HOST = \"https://graph.microsoft.com\"\n\n\nclass TeamsConnector(\n CheckpointedConnectorWithPermSync[TeamsCheckpoint],\n SlimConnectorWithPermSync,\n):\n MAX_WORKERS = 10\n\n def __init__(\n self,\n # TODO: (chris) move from \"Display Names\" to IDs, since display names\n # are not necessarily guaranteed to be unique\n teams: list[str] = [],\n max_workers: int = MAX_WORKERS,\n authority_host: str = DEFAULT_AUTHORITY_HOST,\n graph_api_host: str = DEFAULT_GRAPH_API_HOST,\n ) -> None:\n self.graph_client: GraphClient | None = None\n self.msal_app: msal.ConfidentialClientApplication | None = None\n self.max_workers = max_workers\n self.requested_team_list: list[str] = teams\n\n resolved_env = resolve_microsoft_environment(graph_api_host, authority_host)\n self._azure_environment = resolved_env.environment\n self.authority_host = resolved_env.authority_host\n self.graph_api_host = resolved_env.graph_host\n\n # impls for BaseConnector\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n teams_client_id = credentials[\"teams_client_id\"]\n teams_client_secret = credentials[\"teams_client_secret\"]\n teams_directory_id = credentials[\"teams_directory_id\"]\n\n authority_url = f\"{self.authority_host}/{teams_directory_id}\"\n self.msal_app = msal.ConfidentialClientApplication(\n authority=authority_url,\n client_id=teams_client_id,\n client_credential=teams_client_secret,\n )\n\n def _acquire_token_func() -> dict[str, Any]:\n \"\"\"\n Acquire token via MSAL\n \"\"\"\n if self.msal_app is None:\n raise RuntimeError(\"MSAL app is not initialized\")\n\n token = self.msal_app.acquire_token_for_client(\n scopes=[f\"{self.graph_api_host}/.default\"]\n )\n\n if not isinstance(token, dict):\n raise RuntimeError(\"`token` instance must be of type dict\")\n\n return token\n\n self.graph_client = GraphClient(\n _acquire_token_func, environment=self._azure_environment\n )\n return None\n\n def validate_connector_settings(self) -> None:\n if self.graph_client is None:\n raise ConnectorMissingCredentialError(\"Teams credentials not loaded.\")\n\n # Check if any requested teams have special characters that need client-side filtering\n has_special_chars = _has_odata_incompatible_chars(self.requested_team_list)\n if has_special_chars:\n logger.info(\n \"Some requested team names contain special characters (&, (, )) that require \"\n \"client-side filtering during data retrieval.\"\n )\n\n # Minimal validation: just check if we can access the teams endpoint\n timeout = 10 # Short timeout for basic validation\n\n try:\n # For validation, do a lightweight check instead of full team search\n logger.info(\n f\"Requested team count: {len(self.requested_team_list) if self.requested_team_list else 0}, \"\n f\"Has special chars: {has_special_chars}\"\n )\n\n validation_query = self.graph_client.teams.get().top(1)\n run_with_timeout(\n timeout=timeout,\n func=lambda: validation_query.execute_query(),\n )\n\n logger.info(\n \"Teams validation successful - Access to teams endpoint confirmed\"\n )\n\n except TimeoutError as e:\n raise ConnectorValidationError(\n f\"Timeout while validating Teams access (waited {timeout}s). \"\n f\"This may indicate network issues or authentication problems. \"\n f\"Error: {e}\"\n )\n\n except ClientRequestException as e:\n if not e.response:\n raise RuntimeError(f\"No response provided in error; {e=}\")\n status_code = e.response.status_code\n if status_code == 401:\n raise CredentialExpiredError(\n \"Invalid or expired Microsoft Teams credentials (401 Unauthorized).\"\n )\n elif status_code == 403:\n raise InsufficientPermissionsError(\n \"Your app lacks sufficient permissions to read Teams (403 Forbidden).\"\n )\n raise UnexpectedValidationError(f\"Unexpected error retrieving teams: {e}\")\n\n except Exception as e:\n error_str = str(e).lower()\n if (\n \"unauthorized\" in error_str\n or \"401\" in error_str\n or \"invalid_grant\" in error_str\n ):\n raise CredentialExpiredError(\n \"Invalid or expired Microsoft Teams credentials.\"\n )\n elif \"forbidden\" in error_str or \"403\" in error_str:\n raise InsufficientPermissionsError(\n \"App lacks required permissions to read from Microsoft Teams.\"\n )\n raise ConnectorValidationError(\n f\"Unexpected error during Teams validation: {e}\"\n )\n\n # impls for CheckpointedConnector\n\n def build_dummy_checkpoint(self) -> TeamsCheckpoint:\n return TeamsCheckpoint(\n has_more=True,\n )\n\n def validate_checkpoint_json(self, checkpoint_json: str) -> TeamsCheckpoint:\n return TeamsCheckpoint.model_validate_json(checkpoint_json)\n\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch, # noqa: ARG002\n checkpoint: TeamsCheckpoint,\n ) -> CheckpointOutput[TeamsCheckpoint]:\n if self.graph_client is None:\n raise ConnectorMissingCredentialError(\"Teams\")\n\n checkpoint = cast(TeamsCheckpoint, copy.deepcopy(checkpoint))\n\n todos = checkpoint.todo_team_ids\n\n if todos is None:\n teams = _collect_all_teams(\n graph_client=self.graph_client,\n requested=self.requested_team_list,\n )\n todo_team_ids = [team.id for team in teams if team.id]\n return TeamsCheckpoint(\n todo_team_ids=todo_team_ids,\n has_more=bool(todo_team_ids),\n )\n\n # `todos.pop()` should always return an element. This is because if\n # `todos` was the empty list, then we would have set `has_more=False`\n # during the previous invocation of `TeamsConnector.load_from_checkpoint`,\n # meaning that this function wouldn't have been called in the first place.\n todo_team_id = todos.pop()\n team = _get_team_by_id(\n graph_client=self.graph_client,\n team_id=todo_team_id,\n )\n channels = _collect_all_channels_from_team(\n team=team,\n )\n\n # An iterator of channels, in which each channel is an iterator of docs.\n channels_docs = [\n _collect_documents_for_channel(\n graph_client=self.graph_client,\n team=team,\n channel=channel,\n start=start,\n )\n for channel in channels\n ]\n\n # Was previously `for doc in parallel_yield(gens=docs, max_workers=self.max_workers): ...`.\n # However, that lead to some weird exceptions (potentially due to non-thread-safe behaviour in the Teams library).\n # Reverting back to the non-threaded case for now.\n for channel_docs in channels_docs:\n for channel_doc in channel_docs:\n if channel_doc:\n yield channel_doc\n\n logger.info(\n f\"Processed team with id {todo_team_id}; {len(todos)} team(s) left to process\"\n )\n\n return TeamsCheckpoint(\n todo_team_ids=todos,\n has_more=bool(todos),\n )\n\n def load_from_checkpoint_with_perm_sync(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: TeamsCheckpoint,\n ) -> CheckpointOutput[TeamsCheckpoint]:\n # Teams already fetches external_access (permissions) for each document\n # in _convert_thread_to_document, so we can just delegate to load_from_checkpoint\n return self.load_from_checkpoint(start, end, checkpoint)\n\n # impls for SlimConnectorWithPermSync\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n callback: IndexingHeartbeatInterface | None = None,\n ) -> GenerateSlimDocumentOutput:\n start = start or 0\n\n teams = _collect_all_teams(\n graph_client=self.graph_client,\n requested=self.requested_team_list,\n )\n\n for team in teams:\n if not team.id:\n logger.warning(\n f\"Expected a team with an id, instead got no id: {team=}\"\n )\n continue\n\n channels = _collect_all_channels_from_team(\n team=team,\n )\n\n for channel in channels:\n if not channel.id:\n logger.warning(\n f\"Expected a channel with an id, instead got no id: {channel=}\"\n )\n continue\n\n external_access = fetch_external_access(\n graph_client=self.graph_client, channel=channel\n )\n\n messages = fetch_messages(\n graph_client=self.graph_client,\n team_id=team.id,\n channel_id=channel.id,\n start=start,\n )\n\n slim_doc_buffer: list[SlimDocument | HierarchyNode] = []\n\n for message in messages:\n slim_doc_buffer.append(\n SlimDocument(\n id=message.id,\n external_access=external_access,\n )\n )\n\n if len(slim_doc_buffer) >= _SLIM_DOC_BATCH_SIZE:\n if callback:\n if callback.should_stop():\n raise RuntimeError(\n \"retrieve_all_slim_docs_perm_sync: Stop signal detected\"\n )\n callback.progress(\"retrieve_all_slim_docs_perm_sync\", 1)\n yield slim_doc_buffer\n slim_doc_buffer = []\n\n # Flush any remaining slim documents collected for this channel\n if slim_doc_buffer:\n yield slim_doc_buffer\n slim_doc_buffer = []\n\n\ndef _escape_odata_string(name: str) -> str:\n \"\"\"Escape special characters for OData string literals.\n\n Uses proper OData v4 string literal escaping:\n - Single quotes: ' becomes ''\n - Other characters are handled by using contains() instead of eq for problematic cases\n \"\"\"\n # Escape single quotes for OData syntax (replace ' with '')\n escaped = name.replace(\"'\", \"''\")\n return escaped\n\n\ndef _has_odata_incompatible_chars(team_names: list[str] | None) -> bool:\n \"\"\"Check if any team name contains characters that break Microsoft Graph OData filters.\n\n The Microsoft Graph Teams API has limited OData support. Characters like\n &, (, and ) cause parsing errors and require client-side filtering instead.\n \"\"\"\n if not team_names:\n return False\n return any(char in name for name in team_names for char in [\"&\", \"(\", \")\"])\n\n\ndef _can_use_odata_filter(\n team_names: list[str] | None,\n) -> tuple[bool, list[str], list[str]]:\n \"\"\"Determine which teams can use OData filtering vs client-side filtering.\n\n Microsoft Graph /teams endpoint OData limitations:\n - Only supports basic 'eq' operators in filters\n - No 'contains', 'startswith', or other advanced operators\n - Special characters (&, (, )) break OData parsing\n\n Returns:\n tuple: (can_use_odata, safe_names, problematic_names)\n \"\"\"\n if not team_names:\n return False, [], []\n\n safe_names = []\n problematic_names = []\n\n for name in team_names:\n if any(char in name for char in [\"&\", \"(\", \")\"]):\n problematic_names.append(name)\n else:\n safe_names.append(name)\n\n return bool(safe_names), safe_names, problematic_names\n\n\ndef _build_simple_odata_filter(safe_names: list[str]) -> str | None:\n \"\"\"Build simple OData filter using only 'eq' operators for safe names.\"\"\"\n if not safe_names:\n return None\n\n filter_parts = []\n for name in safe_names:\n escaped_name = _escape_odata_string(name)\n filter_parts.append(f\"displayName eq '{escaped_name}'\")\n\n return \" or \".join(filter_parts)\n\n\ndef _construct_semantic_identifier(channel: Channel, top_message: Message) -> str:\n top_message_user_name: str\n\n if top_message.from_ and top_message.from_.user:\n user_display_name = top_message.from_.user.display_name\n top_message_user_name = (\n user_display_name if user_display_name else \"Unknown User\"\n )\n else:\n logger.warning(f\"Message {top_message=} has no `from.user` field\")\n top_message_user_name = \"Unknown User\"\n\n top_message_content = top_message.body.content or \"\"\n top_message_subject = top_message.subject or \"Unknown Subject\"\n channel_name = channel.properties.get(\"displayName\", \"Unknown\")\n\n try:\n snippet = parse_html_page_basic(top_message_content.rstrip())\n snippet = snippet[:50] + \"...\" if len(snippet) > 50 else snippet\n\n except Exception:\n logger.exception(\n f\"Error parsing snippet for message {top_message.id} with url {top_message.web_url}\"\n )\n snippet = \"\"\n\n semantic_identifier = (\n f\"{top_message_user_name} in {channel_name} about {top_message_subject}\"\n )\n if snippet:\n semantic_identifier += f\": {snippet}\"\n\n return semantic_identifier\n\n\ndef _convert_thread_to_document(\n graph_client: GraphClient,\n channel: Channel,\n thread: list[Message],\n) -> Document | None:\n if len(thread) == 0:\n return None\n\n most_recent_message_datetime: datetime | None = None\n top_message = thread[0]\n thread_text = \"\"\n\n sorted_thread = sorted(thread, key=lambda m: m.created_date_time, reverse=True)\n\n if sorted_thread:\n most_recent_message_datetime = sorted_thread[0].created_date_time\n\n for message in thread:\n # Add text and a newline\n if message.body.content:\n thread_text += parse_html_page_basic(message.body.content)\n\n # If it has a subject, that means its the top level post message, so grab its id, url, and subject\n if message.subject:\n top_message = message\n\n if not thread_text:\n return None\n\n semantic_string = _construct_semantic_identifier(channel, top_message)\n expert_infos = fetch_expert_infos(graph_client=graph_client, channel=channel)\n external_access = fetch_external_access(\n graph_client=graph_client, channel=channel, expert_infos=expert_infos\n )\n\n return Document(\n id=top_message.id,\n sections=[TextSection(link=top_message.web_url, text=thread_text)],\n source=DocumentSource.TEAMS,\n semantic_identifier=semantic_string,\n title=\"\", # teams threads don't really have a \"title\"\n doc_updated_at=most_recent_message_datetime,\n primary_owners=expert_infos,\n metadata={},\n external_access=external_access,\n )\n\n\ndef _update_request_url(request: RequestOptions, next_url: str) -> None:\n request.url = next_url\n\n\ndef _add_prefer_header(request: RequestOptions) -> None:\n \"\"\"Add Prefer header to work around Microsoft Graph API ampersand bug.\n See: https://developer.microsoft.com/en-us/graph/known-issues/?search=18185\n \"\"\"\n if not hasattr(request, \"headers\") or request.headers is None:\n request.headers = {}\n # Add header to handle properly encoded ampersands in filters\n request.headers[\"Prefer\"] = \"legacySearch=false\"\n\n\ndef _collect_all_teams(\n graph_client: GraphClient,\n requested: list[str] | None = None,\n) -> list[Team]:\n \"\"\"Collect teams from Microsoft Graph using appropriate filtering strategy.\n\n For teams with special characters (&, (, )), uses client-side filtering\n with paginated search. For teams without special characters, uses efficient\n OData server-side filtering.\n\n Args:\n graph_client: Authenticated Microsoft Graph client\n requested: List of team names to find, or None for all teams\n\n Returns:\n List of Team objects matching the requested names\n \"\"\"\n teams: list[Team] = []\n next_url: str | None = None\n\n # Determine filtering strategy based on Microsoft Graph limitations\n if not requested:\n # No specific teams requested - return empty list (avoid fetching all teams)\n logger.info(\"No specific teams requested - returning empty list\")\n return []\n\n _, safe_names, problematic_names = _can_use_odata_filter(requested)\n\n if problematic_names and not safe_names:\n # ALL requested teams have special characters - cannot use OData filtering\n logger.info(\n f\"All requested team names contain special characters (&, (, )) which require \"\n f\"client-side filtering. Using basic /teams endpoint with pagination. \"\n f\"Teams: {problematic_names}\"\n )\n # Use unfiltered query with pagination limit to avoid fetching too many teams\n use_client_side_filtering = True\n odata_filter = None\n elif problematic_names and safe_names:\n # Mixed scenario - need to fetch more teams to find the problematic ones\n logger.info(\n f\"Mixed team types: will use client-side filtering for all. \"\n f\"Safe names: {safe_names}, Special char names: {problematic_names}\"\n )\n use_client_side_filtering = True\n odata_filter = None\n elif safe_names:\n # All names are safe - use OData filtering\n logger.info(f\"Using OData filtering for all requested teams: {safe_names}\")\n use_client_side_filtering = False\n odata_filter = _build_simple_odata_filter(safe_names)\n else:\n # No valid names\n return []\n\n # Track pagination to avoid fetching too many teams for client-side filtering\n max_pages = 200\n page_count = 0\n\n while True:\n try:\n if use_client_side_filtering:\n # Use basic /teams endpoint with top parameter to limit results per page\n query = graph_client.teams.get().top(50) # Limit to 50 teams per page\n else:\n # Use OData filter with only 'eq' operators\n query = graph_client.teams.get().filter(odata_filter)\n\n # Add header to work around Microsoft Graph API issues\n query.before_execute(lambda req: _add_prefer_header(request=req))\n\n if next_url:\n url = next_url\n query.before_execute(\n lambda req: _update_request_url(request=req, next_url=url)\n )\n\n team_collection = query.execute_query()\n except (ClientRequestException, ValueError) as e:\n # If OData filter fails, fall back to client-side filtering\n if not use_client_side_filtering and odata_filter:\n logger.warning(\n f\"OData filter failed: {e}. Falling back to client-side filtering.\"\n )\n use_client_side_filtering = True\n odata_filter = None\n teams = []\n next_url = None\n page_count = 0\n continue\n # If client-side approach also fails, re-raise\n logger.error(f\"Teams query failed: {e}\")\n raise\n\n filtered_teams = (\n team\n for team in team_collection\n if _filter_team(team=team, requested=requested)\n )\n teams.extend(filtered_teams)\n\n # For client-side filtering, check if we found all requested teams or hit page limit\n if use_client_side_filtering:\n page_count += 1\n found_team_names = {\n team.display_name for team in teams if team.display_name\n }\n requested_set = set(requested)\n\n # Log progress every 10 pages to avoid excessive logging\n if page_count % 10 == 0:\n logger.info(\n f\"Searched {page_count} pages, found {len(found_team_names)} matching teams so far\"\n )\n\n # Stop if we found all requested teams or hit the page limit\n if requested_set.issubset(found_team_names):\n logger.info(f\"Found all requested teams after {page_count} pages\")\n break\n elif page_count >= max_pages:\n logger.warning(\n f\"Reached maximum page limit ({max_pages}) while searching for teams. \"\n f\"Found: {found_team_names & requested_set}, \"\n f\"Missing: {requested_set - found_team_names}\"\n )\n break\n\n if not team_collection.has_next:\n break\n\n if not isinstance(team_collection._next_request_url, str):\n raise ValueError(\n f\"The next request url field should be a string, instead got {type(team_collection._next_request_url)}\"\n )\n\n next_url = team_collection._next_request_url\n\n return teams\n\n\ndef _normalize_team_name(name: str) -> str:\n \"\"\"Normalize team name for flexible matching.\"\"\"\n if not name:\n return \"\"\n # Convert to lowercase and strip whitespace for case-insensitive matching\n return name.lower().strip()\n\n\ndef _matches_requested_team(\n team_display_name: str, requested: list[str] | None\n) -> bool:\n \"\"\"Check if team display name matches any of the requested team names.\n\n Uses flexible matching to handle slight variations in team names.\n \"\"\"\n if not requested or not team_display_name:\n return (\n not requested\n ) # If no teams requested, match all; if no name, don't match\n\n normalized_team_name = _normalize_team_name(team_display_name)\n\n for requested_name in requested:\n normalized_requested = _normalize_team_name(requested_name)\n\n # Exact match after normalization\n if normalized_team_name == normalized_requested:\n return True\n\n # Flexible matching - check if team name contains all significant words\n # This helps with slight variations in formatting\n team_words = set(normalized_team_name.split())\n requested_words = set(normalized_requested.split())\n\n # If the requested name has special characters, split on those too\n for char in [\"&\", \"(\", \")\"]:\n if char in normalized_requested:\n # Split on special characters and add words\n parts = normalized_requested.replace(char, \" \").split()\n requested_words.update(parts)\n\n # Remove very short words that aren't meaningful\n meaningful_requested_words = {\n word for word in requested_words if len(word) >= 3\n }\n\n # Check if team name contains most of the meaningful words\n if (\n meaningful_requested_words\n and len(meaningful_requested_words & team_words)\n >= len(meaningful_requested_words) * 0.7\n ):\n return True\n\n return False\n\n\ndef _filter_team(\n team: Team,\n requested: list[str] | None = None,\n) -> bool:\n \"\"\"\n Returns the true if:\n - Team is not expired / deleted\n - Team has a display-name and ID\n - Team display-name matches any of the requested teams (with flexible matching)\n\n Otherwise, returns false.\n \"\"\"\n\n if not team.id or not team.display_name:\n return False\n\n if not _matches_requested_team(team.display_name, requested):\n return False\n\n props = team.properties\n\n expiration = props.get(\"expirationDateTime\")\n deleted = props.get(\"deletedDateTime\")\n\n # We just check for the existence of those two fields, not their actual dates.\n # This is because if these fields do exist, they have to have occurred in the past, thus making them already\n # expired / deleted.\n return not expiration and not deleted\n\n\ndef _get_team_by_id(\n graph_client: GraphClient,\n team_id: str,\n) -> Team:\n team_collection = (\n graph_client.teams.get().filter(f\"id eq '{team_id}'\").top(1).execute_query()\n )\n\n if not team_collection:\n raise ValueError(f\"No team with {team_id=} was found\")\n elif team_collection.has_next:\n # shouldn't happen, but catching it regardless\n raise RuntimeError(f\"Multiple teams with {team_id=} were found\")\n\n return team_collection[0]\n\n\ndef _collect_all_channels_from_team(\n team: Team,\n) -> list[Channel]:\n if not team.id:\n raise RuntimeError(f\"The {team=} has an empty `id` field\")\n\n channels: list[Channel] = []\n next_url = None\n\n while True:\n query = team.channels.get_all(\n # explicitly needed because of incorrect type definitions provided by the `office365` library\n page_loaded=lambda _: None\n )\n if next_url:\n url = next_url\n query = query.before_execute(\n lambda req: _update_request_url(request=req, next_url=url)\n )\n\n channel_collection = query.execute_query()\n channels.extend(channel for channel in channel_collection if channel.id)\n\n if not channel_collection.has_next:\n break\n\n return channels\n\n\ndef _collect_documents_for_channel(\n graph_client: GraphClient,\n team: Team,\n channel: Channel,\n start: SecondsSinceUnixEpoch,\n) -> Iterator[Document | None | ConnectorFailure]:\n \"\"\"\n This function yields an iterator of `Document`s, where each `Document` corresponds to a \"thread\".\n\n A \"thread\" is the conjunction of the \"root\" message and all of its replies.\n \"\"\"\n\n for message in fetch_messages(\n graph_client=graph_client,\n team_id=team.id,\n channel_id=channel.id,\n start=start,\n ):\n try:\n replies = list(\n fetch_replies(\n graph_client=graph_client,\n team_id=team.id,\n channel_id=channel.id,\n root_message_id=message.id,\n )\n )\n\n thread = [message]\n thread.extend(replies[::-1])\n\n # Note:\n # We convert an entire *thread* (including the root message and its replies) into one, singular `Document`.\n # I.e., we don't convert each individual message and each individual reply into their own individual `Document`s.\n if doc := _convert_thread_to_document(\n graph_client=graph_client,\n channel=channel,\n thread=thread,\n ):\n yield doc\n\n except Exception as e:\n yield ConnectorFailure(\n failed_entity=EntityFailure(\n entity_id=message.id,\n ),\n failure_message=f\"Retrieval of message and its replies failed; {channel.id=} {message.id}\",\n exception=e,\n )\n\n\nif __name__ == \"__main__\":\n from tests.daily.connectors.utils import load_all_from_connector\n\n app_id = os.environ[\"TEAMS_APPLICATION_ID\"]\n dir_id = os.environ[\"TEAMS_DIRECTORY_ID\"]\n secret = os.environ[\"TEAMS_SECRET\"]\n\n teams_env_var = os.environ.get(\"TEAMS\", None)\n teams = teams_env_var.split(\",\") if teams_env_var else []\n\n teams_connector = TeamsConnector(teams=teams)\n teams_connector.load_credentials(\n {\n \"teams_client_id\": app_id,\n \"teams_directory_id\": dir_id,\n \"teams_client_secret\": secret,\n }\n )\n teams_connector.validate_connector_settings()\n\n for slim_doc in teams_connector.retrieve_all_slim_docs_perm_sync():\n ...\n\n for doc in load_all_from_connector(\n connector=teams_connector,\n start=0.0,\n end=datetime.now(tz=timezone.utc).timestamp(),\n ).documents:\n print(doc)\n" }, { "path": "backend/onyx/connectors/teams/models.py", "content": "from datetime import datetime\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import Field\nfrom pydantic.alias_generators import to_camel\n\n\nclass Body(BaseModel):\n content_type: str\n content: str | None\n\n model_config = ConfigDict(\n alias_generator=to_camel,\n populate_by_name=True,\n )\n\n\nclass User(BaseModel):\n id: str\n display_name: str\n\n model_config = ConfigDict(\n alias_generator=to_camel,\n populate_by_name=True,\n )\n\n\nclass From(BaseModel):\n user: User | None\n\n model_config = ConfigDict(\n alias_generator=to_camel,\n populate_by_name=True,\n )\n\n\nclass Message(BaseModel):\n id: str\n replyToId: str | None\n subject: str | None\n from_: From | None = Field(alias=\"from\")\n body: Body\n created_date_time: datetime\n last_modified_date_time: datetime | None\n last_edited_date_time: datetime | None\n deleted_date_time: datetime | None\n web_url: str\n\n model_config = ConfigDict(\n alias_generator=to_camel,\n populate_by_name=True,\n )\n" }, { "path": "backend/onyx/connectors/teams/utils.py", "content": "import time\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom http import HTTPStatus\n\nfrom office365.graph_client import GraphClient # type: ignore[import-untyped]\nfrom office365.teams.channels.channel import Channel # type: ignore[import-untyped]\nfrom office365.teams.channels.channel import ConversationMember\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.teams.models import Message\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n_PUBLIC_MEMBERSHIP_TYPE = \"standard\" # public teams channel\n\n\ndef _sanitize_message_user_display_name(value: dict) -> dict:\n try:\n from_obj = value.get(\"from\")\n if isinstance(from_obj, dict):\n user_obj = from_obj.get(\"user\")\n if isinstance(user_obj, dict) and user_obj.get(\"displayName\") is None:\n value = dict(value)\n from_obj = dict(from_obj)\n user_obj = dict(user_obj)\n user_obj[\"displayName\"] = \"Unknown User\"\n from_obj[\"user\"] = user_obj\n value[\"from\"] = from_obj\n except (AttributeError, TypeError, KeyError):\n pass\n return value\n\n\ndef _retry(\n graph_client: GraphClient,\n request_url: str,\n) -> dict:\n MAX_RETRIES = 10\n retry_number = 0\n\n while retry_number < MAX_RETRIES:\n response = graph_client.execute_request_direct(request_url)\n if response.ok:\n json = response.json()\n if not isinstance(json, dict):\n raise RuntimeError(f\"Expected a JSON object, instead got {json=}\")\n\n return json\n\n if response.status_code == int(HTTPStatus.TOO_MANY_REQUESTS):\n retry_number += 1\n\n cooldown = int(response.headers.get(\"Retry-After\", 10))\n time.sleep(cooldown)\n\n continue\n\n response.raise_for_status()\n\n raise RuntimeError(\n f\"Max number of retries for hitting {request_url=} exceeded; unable to fetch data\"\n )\n\n\ndef _get_next_url(\n graph_client: GraphClient,\n json_response: dict,\n) -> str | None:\n next_url = json_response.get(\"@odata.nextLink\")\n\n if not next_url:\n return None\n\n if not isinstance(next_url, str):\n raise RuntimeError(\n f\"Expected a string for the `@odata.nextUrl`, instead got {next_url=}\"\n )\n\n return next_url.removeprefix(graph_client.service_root_url()).removeprefix(\"/\")\n\n\ndef _get_or_fetch_email(\n graph_client: GraphClient,\n member: ConversationMember,\n) -> str | None:\n if email := member.properties.get(\"email\"):\n return email\n\n user_id = member.properties.get(\"userId\")\n if not user_id:\n logger.warn(f\"No user-id found for this member; {member=}\")\n return None\n\n json_data = _retry(graph_client=graph_client, request_url=f\"users/{user_id}\")\n email = json_data.get(\"userPrincipalName\")\n\n if not isinstance(email, str):\n logger.warn(f\"Expected email to be of type str, instead got {email=}\")\n return None\n\n return email\n\n\ndef _is_channel_public(channel: Channel) -> bool:\n return (\n channel.membership_type and channel.membership_type == _PUBLIC_MEMBERSHIP_TYPE\n )\n\n\ndef fetch_messages(\n graph_client: GraphClient,\n team_id: str,\n channel_id: str,\n start: SecondsSinceUnixEpoch,\n) -> Generator[Message]:\n startfmt = datetime.fromtimestamp(start, tz=timezone.utc).strftime(\n \"%Y-%m-%dT%H:%M:%SZ\"\n )\n\n initial_request_url = f\"teams/{team_id}/channels/{channel_id}/messages/delta?$filter=lastModifiedDateTime gt {startfmt}\"\n\n request_url: str | None = initial_request_url\n\n while request_url:\n json_response = _retry(graph_client=graph_client, request_url=request_url)\n\n for value in json_response.get(\"value\", []):\n yield Message(**_sanitize_message_user_display_name(value))\n\n request_url = _get_next_url(\n graph_client=graph_client, json_response=json_response\n )\n\n\ndef fetch_replies(\n graph_client: GraphClient,\n team_id: str,\n channel_id: str,\n root_message_id: str,\n) -> Generator[Message]:\n initial_request_url = (\n f\"teams/{team_id}/channels/{channel_id}/messages/{root_message_id}/replies\"\n )\n\n request_url: str | None = initial_request_url\n\n while request_url:\n json_response = _retry(graph_client=graph_client, request_url=request_url)\n\n for value in json_response.get(\"value\", []):\n yield Message(**_sanitize_message_user_display_name(value))\n\n request_url = _get_next_url(\n graph_client=graph_client, json_response=json_response\n )\n\n\ndef fetch_expert_infos(\n graph_client: GraphClient, channel: Channel\n) -> list[BasicExpertInfo]:\n members = channel.members.get_all(\n # explicitly needed because of incorrect type definitions provided by the `office365` library\n page_loaded=lambda _: None\n ).execute_query_retry()\n\n expert_infos = []\n for member in members:\n if not member.display_name:\n logger.warn(f\"Failed to grab the display-name of {member=}; skipping\")\n continue\n\n email = _get_or_fetch_email(graph_client=graph_client, member=member)\n if not email:\n logger.warn(f\"Failed to grab the email of {member=}; skipping\")\n continue\n\n expert_infos.append(\n BasicExpertInfo(\n display_name=member.display_name,\n email=email,\n )\n )\n\n return expert_infos\n\n\ndef fetch_external_access(\n graph_client: GraphClient,\n channel: Channel,\n expert_infos: list[BasicExpertInfo] | None = None,\n) -> ExternalAccess:\n is_public = _is_channel_public(channel=channel)\n\n if is_public:\n return ExternalAccess.public()\n\n expert_infos = (\n expert_infos\n if expert_infos is not None\n else fetch_expert_infos(graph_client=graph_client, channel=channel)\n )\n emails = {expert_info.email for expert_info in expert_infos if expert_info.email}\n\n return ExternalAccess(\n external_user_emails=emails,\n external_user_group_ids=set(),\n is_public=is_public,\n )\n" }, { "path": "backend/onyx/connectors/testrail/__init__.py", "content": "# Package marker for TestRail connector\n" }, { "path": "backend/onyx/connectors/testrail/connector.py", "content": "from __future__ import annotations\n\nfrom collections.abc import Iterator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import ClassVar\nfrom typing import Optional\n\nimport requests\nfrom bs4 import BeautifulSoup\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import format_document_soup\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.text_processing import remove_markdown_image_references\n\n\nlogger = setup_logger()\n\n\nclass TestRailConnector(LoadConnector, PollConnector):\n \"\"\"Connector for TestRail.\n\n Minimal implementation that indexes Test Cases per project.\n \"\"\"\n\n document_source_type: ClassVar[DocumentSource] = DocumentSource.TESTRAIL\n\n # Fields that need ID-to-label value mapping\n FIELDS_NEEDING_VALUE_MAPPING: ClassVar[set[str]] = {\n \"priority_id\",\n \"custom_automation_type\",\n \"custom_scenario_db_automation\",\n \"custom_case_golden_canvas_automation\",\n \"custom_customers\",\n \"custom_case_environments\",\n \"custom_case_overall_automation\",\n \"custom_case_team_ownership\",\n \"custom_case_unit_or_integration_automation\",\n \"custom_effort\",\n }\n\n def __init__(\n self,\n batch_size: int = INDEX_BATCH_SIZE,\n project_ids: str | list[int] | None = None,\n cases_page_size: int | None = None,\n max_pages: int | None = None,\n skip_doc_absolute_chars: int | None = None,\n ) -> None:\n self.base_url: str | None = None\n self.username: str | None = None\n self.api_key: str | None = None\n self.batch_size = batch_size\n parsed_project_ids: list[int] | None\n\n # Parse project_ids from string if needed\n # None = all projects (no filtering), [] = no projects, [1,2,3] = specific projects\n if isinstance(project_ids, str):\n if project_ids.strip():\n parsed_project_ids = [\n int(x.strip()) for x in project_ids.split(\",\") if x.strip()\n ]\n else:\n # Empty string from UI means \"all projects\"\n parsed_project_ids = None\n elif project_ids is None:\n parsed_project_ids = None\n else:\n parsed_project_ids = [int(pid) for pid in project_ids]\n\n self.project_ids: list[int] | None = parsed_project_ids\n\n # Handle empty strings from UI and convert to int with defaults\n self.cases_page_size = (\n int(cases_page_size)\n if cases_page_size and str(cases_page_size).strip()\n else 250\n )\n self.max_pages = (\n int(max_pages) if max_pages and str(max_pages).strip() else 10000\n )\n self.skip_doc_absolute_chars = (\n int(skip_doc_absolute_chars)\n if skip_doc_absolute_chars and str(skip_doc_absolute_chars).strip()\n else 200000\n )\n\n # Cache for field labels and value mappings - will be populated on first use\n self._field_labels: dict[str, str] | None = None\n self._value_maps: dict[str, dict[str, str]] | None = None\n\n # --- Rich text sanitization helpers ---\n # Note: TestRail stores some fields as HTML (e.g. shared test steps).\n # This function handles both HTML and plain text.\n @staticmethod\n def _sanitize_rich_text(value: Any) -> str:\n if value is None:\n return \"\"\n text = str(value)\n\n # Parse HTML and remove image tags\n soup = BeautifulSoup(text, \"html.parser\")\n\n # Remove all img tags and their containers\n for img_tag in soup.find_all(\"img\"):\n img_tag.decompose()\n for span in soup.find_all(\"span\", class_=\"markdown-img-container\"):\n span.decompose()\n\n # Use format_document_soup for better HTML-to-text conversion\n # This preserves document structure (paragraphs, lists, line breaks, etc.)\n text = format_document_soup(soup)\n\n # Also remove markdown-style image references (in case any remain)\n text = remove_markdown_image_references(text)\n\n return text.strip()\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n # Expected keys from UI credential JSON\n self.base_url = str(credentials[\"testrail_base_url\"]).rstrip(\"/\")\n self.username = str(credentials[\"testrail_username\"]) # email or username\n self.api_key = str(credentials[\"testrail_api_key\"]) # API key (password)\n return None\n\n def validate_connector_settings(self) -> None:\n \"\"\"Lightweight validation to surface common misconfigurations early.\"\"\"\n projects = self._list_projects()\n if not projects:\n logger.warning(\"TestRail: no projects visible to this credential.\")\n\n # ---- API helpers ----\n def _api_get(self, endpoint: str, params: Optional[dict[str, Any]] = None) -> Any:\n if not self.base_url or not self.username or not self.api_key:\n raise ConnectorMissingCredentialError(\"testrail\")\n\n # TestRail API base is typically /index.php?/api/v2/\n url = f\"{self.base_url}/index.php?/api/v2/{endpoint}\"\n try:\n response = requests.get(\n url,\n auth=(self.username, self.api_key),\n params=params,\n )\n response.raise_for_status()\n except requests.exceptions.HTTPError as e:\n status = e.response.status_code if getattr(e, \"response\", None) else None\n if status == 401:\n raise CredentialExpiredError(\n \"Invalid or expired TestRail credentials (HTTP 401).\"\n ) from e\n if status == 403:\n raise InsufficientPermissionsError(\n \"Insufficient permissions to access TestRail resources (HTTP 403).\"\n ) from e\n raise UnexpectedValidationError(\n f\"Unexpected TestRail HTTP error (status={status}).\"\n ) from e\n except requests.exceptions.RequestException as e:\n raise UnexpectedValidationError(f\"TestRail request failed: {e}\") from e\n\n try:\n return response.json()\n except ValueError as e:\n raise UnexpectedValidationError(\n \"Invalid JSON returned by TestRail API\"\n ) from e\n\n def _list_projects(self) -> list[dict[str, Any]]:\n projects = self._api_get(\"get_projects\")\n if isinstance(projects, dict):\n projects_list = projects.get(\"projects\")\n return projects_list if isinstance(projects_list, list) else []\n return []\n\n def _list_suites(self, project_id: int) -> list[dict[str, Any]]:\n \"\"\"Return suites for a project. If the project is in single-suite mode,\n some TestRail instances may return an empty list; callers should\n gracefully fallback to calling get_cases without suite_id.\n \"\"\"\n suites = self._api_get(f\"get_suites/{project_id}\")\n if isinstance(suites, dict):\n suites_list = suites.get(\"suites\")\n return suites_list if isinstance(suites_list, list) else []\n return []\n\n def _get_case_fields(self) -> list[dict[str, Any]]:\n \"\"\"Get case field definitions from TestRail API.\"\"\"\n try:\n fields = self._api_get(\"get_case_fields\")\n return fields if isinstance(fields, list) else []\n except Exception as e:\n logger.warning(f\"Failed to fetch case fields from TestRail: {e}\")\n return []\n\n def _parse_items_string(self, items_str: str) -> dict[str, str]:\n \"\"\"Parse items string from field config into ID -> label mapping.\n\n Format: \"1, Option A\\\\n2, Option B\\\\n3, Option C\"\n Returns: {\"1\": \"Option A\", \"2\": \"Option B\", \"3\": \"Option C\"}\n \"\"\"\n id_to_label: dict[str, str] = {}\n if not items_str:\n return id_to_label\n\n for line in items_str.split(\"\\n\"):\n line = line.strip()\n if not line:\n continue\n parts = line.split(\",\", 1)\n if len(parts) == 2:\n item_id = parts[0].strip()\n item_label = parts[1].strip()\n id_to_label[item_id] = item_label\n\n return id_to_label\n\n def _build_field_maps(self) -> tuple[dict[str, str], dict[str, dict[str, str]]]:\n \"\"\"Build both field labels and value mappings in one pass.\n\n Returns:\n (field_labels, value_maps) where:\n - field_labels: system_name -> label\n - value_maps: system_name -> {id -> label}\n \"\"\"\n field_labels = {}\n value_maps = {}\n\n try:\n fields = self._get_case_fields()\n for field in fields:\n system_name = field.get(\"system_name\")\n\n # Build field label map\n label = field.get(\"label\")\n if system_name and label:\n field_labels[system_name] = label\n\n # Build value map if needed\n if system_name in self.FIELDS_NEEDING_VALUE_MAPPING:\n configs = field.get(\"configs\", [])\n if configs and len(configs) > 0:\n options = configs[0].get(\"options\", {})\n items_str = options.get(\"items\")\n if items_str:\n value_maps[system_name] = self._parse_items_string(\n items_str\n )\n\n except Exception as e:\n logger.warning(f\"Failed to build field maps from TestRail: {e}\")\n\n return field_labels, value_maps\n\n def _get_field_labels(self) -> dict[str, str]:\n \"\"\"Get field labels, fetching from API if not cached.\"\"\"\n if self._field_labels is None:\n self._field_labels, self._value_maps = self._build_field_maps()\n return self._field_labels\n\n def _get_value_maps(self) -> dict[str, dict[str, str]]:\n \"\"\"Get value maps, fetching from API if not cached.\"\"\"\n if self._value_maps is None:\n self._field_labels, self._value_maps = self._build_field_maps()\n return self._value_maps\n\n def _map_field_value(self, field_name: str, field_value: Any) -> str:\n \"\"\"Map a field value using the value map if available.\n\n Examples:\n - priority_id: 2 -> \"Medium\"\n - custom_case_team_ownership: [10] -> \"Sim Platform\"\n - custom_case_environments: [1, 2] -> \"Local, Cloud\"\n \"\"\"\n if field_value is None or field_value == \"\":\n return \"\"\n\n # Get value map for this field\n value_maps = self._get_value_maps()\n value_map = value_maps.get(field_name, {})\n\n # Handle list values\n if isinstance(field_value, list):\n if not field_value:\n return \"\"\n mapped = [value_map.get(str(v), str(v)) for v in field_value]\n return \", \".join(mapped)\n\n # Handle single values\n val_str = str(field_value)\n return value_map.get(val_str, val_str)\n\n def _get_cases(\n self, project_id: int, suite_id: Optional[int], limit: int, offset: int\n ) -> list[dict[str, Any]]:\n \"\"\"Get cases for a project from the API.\"\"\"\n params: dict[str, Any] = {\"limit\": limit, \"offset\": offset}\n if suite_id is not None:\n params[\"suite_id\"] = suite_id\n cases_response = self._api_get(f\"get_cases/{project_id}\", params=params)\n cases_list: list[dict[str, Any]] = []\n if isinstance(cases_response, dict):\n cases_items = cases_response.get(\"cases\")\n if isinstance(cases_items, list):\n cases_list = cases_items\n return cases_list\n\n def _iter_cases(\n self,\n project_id: int,\n suite_id: Optional[int] = None,\n start: Optional[SecondsSinceUnixEpoch] = None,\n end: Optional[SecondsSinceUnixEpoch] = None,\n ) -> Iterator[dict[str, Any]]:\n # Pagination: TestRail supports 'limit' and 'offset' for many list endpoints\n limit = self.cases_page_size\n # Use a bounded page loop to avoid infinite loops on API anomalies\n for page_index in range(self.max_pages):\n offset = page_index * limit\n cases = self._get_cases(project_id, suite_id, limit, offset)\n\n if not cases:\n break\n\n # Filter by updated window if provided\n for case in cases:\n # 'updated_on' is unix timestamp (seconds)\n updated_on = case.get(\"updated_on\") or case.get(\"created_on\")\n if start is not None and updated_on is not None and updated_on < start:\n continue\n if end is not None and updated_on is not None and updated_on > end:\n continue\n yield case\n\n if len(cases) < limit:\n break\n\n def _build_case_link(self, project_id: int, case_id: int) -> str: # noqa: ARG002\n # Standard UI link to a case\n return f\"{self.base_url}/index.php?/cases/view/{case_id}\"\n\n def _doc_from_case(\n self,\n project: dict[str, Any],\n case: dict[str, Any],\n suite: dict[str, Any] | None = None, # noqa: ARG002\n ) -> Document | None:\n project_id = project.get(\"id\")\n if not isinstance(project_id, int):\n logger.warning(\n \"Skipping TestRail case because project id is missing or invalid: %s\",\n project_id,\n )\n return None\n\n case_id = case.get(\"id\")\n if not isinstance(case_id, int):\n logger.warning(\n \"Skipping TestRail case because case id is missing or invalid: %s\",\n case_id,\n )\n return None\n\n title = case.get(\"title\", f\"Case {case_id}\")\n case_key = f\"C{case_id}\"\n\n # Convert epoch seconds to aware datetime if available\n updated = case.get(\"updated_on\") or case.get(\"created_on\")\n updated_dt = (\n datetime.fromtimestamp(updated, tz=timezone.utc)\n if isinstance(updated, (int, float))\n else None\n )\n\n text_lines: list[str] = []\n if case.get(\"title\"):\n text_lines.append(f\"Title: {case['title']}\")\n if case_key:\n text_lines.append(f\"Case ID: {case_key}\")\n if case_id is not None:\n text_lines.append(f\"ID: {case_id}\")\n doc_link = case.get(\"custom_documentation_link\")\n if doc_link:\n text_lines.append(f\"Documentation: {doc_link}\")\n\n # Add fields that need value mapping\n field_labels = self._get_field_labels()\n for field_name in self.FIELDS_NEEDING_VALUE_MAPPING:\n field_value = case.get(field_name)\n if field_value is not None and field_value != \"\" and field_value != []:\n mapped_value = self._map_field_value(field_name, field_value)\n if mapped_value:\n # Get label from TestRail field definition\n label = field_labels.get(\n field_name, field_name.replace(\"_\", \" \").title()\n )\n text_lines.append(f\"{label}: {mapped_value}\")\n\n pre = self._sanitize_rich_text(case.get(\"custom_preconds\"))\n if pre:\n text_lines.append(f\"Preconditions: {pre}\")\n\n # Steps: use separated steps format if available\n steps_added = False\n steps_separated = case.get(\"custom_steps_separated\")\n if isinstance(steps_separated, list) and steps_separated:\n rendered_steps: list[str] = []\n for idx, step_item in enumerate(steps_separated, start=1):\n step_content = self._sanitize_rich_text(step_item.get(\"content\"))\n step_expected = self._sanitize_rich_text(step_item.get(\"expected\"))\n parts: list[str] = []\n if step_content:\n parts.append(f\"Step {idx}: {step_content}\")\n else:\n parts.append(f\"Step {idx}:\")\n if step_expected:\n parts.append(f\"Expected: {step_expected}\")\n rendered_steps.append(\"\\n\".join(parts))\n if rendered_steps:\n text_lines.append(\"Steps:\\n\" + \"\\n\".join(rendered_steps))\n steps_added = True\n\n # Fallback to custom_steps and custom_expected if no separated steps\n if not steps_added:\n custom_steps = self._sanitize_rich_text(case.get(\"custom_steps\"))\n custom_expected = self._sanitize_rich_text(case.get(\"custom_expected\"))\n if custom_steps:\n text_lines.append(f\"Steps: {custom_steps}\")\n if custom_expected:\n text_lines.append(f\"Expected: {custom_expected}\")\n\n link = self._build_case_link(project_id, case_id)\n\n # Build full text and apply size policies\n full_text = \"\\n\".join(text_lines)\n if len(full_text) > self.skip_doc_absolute_chars:\n logger.warning(\n f\"Skipping TestRail case {case_id} due to excessive size: {len(full_text)} chars\"\n )\n return None\n\n # Metadata for document identification\n metadata: dict[str, Any] = {}\n if case_key:\n metadata[\"case_key\"] = case_key\n\n # Include the human-friendly case key in identifiers for easier search\n display_title = f\"{case_key}: {title}\" if case_key else title\n\n return Document(\n id=f\"TESTRAIL_CASE_{case_id}\",\n source=DocumentSource.TESTRAIL,\n semantic_identifier=display_title,\n title=display_title,\n sections=[TextSection(link=link, text=full_text)],\n metadata=metadata,\n doc_updated_at=updated_dt,\n )\n\n def _generate_documents(\n self,\n start: Optional[SecondsSinceUnixEpoch],\n end: Optional[SecondsSinceUnixEpoch],\n ) -> GenerateDocumentsOutput:\n if not self.base_url or not self.username or not self.api_key:\n raise ConnectorMissingCredentialError(\"testrail\")\n\n doc_batch: list[Document | HierarchyNode] = []\n\n projects = self._list_projects()\n project_filter: list[int] | None = self.project_ids\n\n for project in projects:\n project_id_raw = project.get(\"id\")\n if not isinstance(project_id_raw, int):\n logger.warning(\n \"Skipping TestRail project with invalid id: %s\", project_id_raw\n )\n continue\n project_id = project_id_raw\n # None = index all, [] = index none, [1,2,3] = index only those\n if project_filter is not None and project_id not in project_filter:\n continue\n\n suites = self._list_suites(project_id)\n if suites:\n for s in suites:\n suite_id = s.get(\"id\")\n for case in self._iter_cases(project_id, suite_id, start, end):\n doc = self._doc_from_case(project, case, s)\n if doc is None:\n continue\n doc_batch.append(doc)\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n else:\n # single-suite mode fallback\n for case in self._iter_cases(project_id, None, start, end):\n doc = self._doc_from_case(project, case, None)\n if doc is None:\n continue\n doc_batch.append(doc)\n if len(doc_batch) >= self.batch_size:\n yield doc_batch\n doc_batch = []\n\n if doc_batch:\n yield doc_batch\n\n # ---- Onyx interfaces ----\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._generate_documents(start=None, end=None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n return self._generate_documents(start=start, end=end)\n\n\nif __name__ == \"__main__\":\n from onyx.configs.app_configs import (\n TESTRAIL_API_KEY,\n TESTRAIL_BASE_URL,\n TESTRAIL_USERNAME,\n )\n\n connector = TestRailConnector()\n\n connector.load_credentials(\n {\n \"testrail_base_url\": TESTRAIL_BASE_URL,\n \"testrail_username\": TESTRAIL_USERNAME,\n \"testrail_api_key\": TESTRAIL_API_KEY,\n }\n )\n\n connector.validate_connector_settings()\n\n # Probe a tiny batch from load\n total = 0\n for batch in connector.load_from_state():\n print(f\"Fetched batch: {len(batch)} docs\")\n total += len(batch)\n if total >= 10:\n break\n print(f\"Total fetched in test: {total}\")\n" }, { "path": "backend/onyx/connectors/web/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/web/connector.py", "content": "import ipaddress\nimport random\nimport socket\nimport time\nfrom datetime import datetime\nfrom datetime import timezone\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\nfrom typing import Tuple\nfrom urllib.parse import urljoin\nfrom urllib.parse import urlparse\n\nimport requests\nfrom bs4 import BeautifulSoup\nfrom oauthlib.oauth2 import BackendApplicationClient\nfrom playwright.sync_api import BrowserContext\nfrom playwright.sync_api import Playwright\nfrom playwright.sync_api import sync_playwright\nfrom playwright.sync_api import TimeoutError\nfrom requests_oauthlib import OAuth2Session # type:ignore\nfrom urllib3.exceptions import MaxRetryError\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.app_configs import WEB_CONNECTOR_OAUTH_CLIENT_ID\nfrom onyx.configs.app_configs import WEB_CONNECTOR_OAUTH_CLIENT_SECRET\nfrom onyx.configs.app_configs import WEB_CONNECTOR_OAUTH_TOKEN_URL\nfrom onyx.configs.app_configs import WEB_CONNECTOR_VALIDATE_URLS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.exceptions import UnexpectedValidationError\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import web_html_cleanup\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.sitemap import list_pages_for_site\nfrom onyx.utils.web_content import extract_pdf_text\nfrom onyx.utils.web_content import is_pdf_resource\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\nclass ScrapeSessionContext:\n \"\"\"Session level context for scraping\"\"\"\n\n def __init__(self, base_url: str, to_visit: list[str]):\n self.base_url = base_url\n self.to_visit = to_visit\n self.visited_links: set[str] = set()\n self.content_hashes: set[int] = set()\n\n self.doc_batch: list[Document | HierarchyNode] = []\n\n self.at_least_one_doc: bool = False\n self.last_error: str | None = None\n self.needs_retry: bool = False\n\n self.playwright: Playwright | None = None\n self.playwright_context: BrowserContext | None = None\n\n def initialize(self) -> None:\n self.stop()\n self.playwright, self.playwright_context = start_playwright()\n\n def stop(self) -> None:\n if self.playwright_context:\n self.playwright_context.close()\n self.playwright_context = None\n\n if self.playwright:\n self.playwright.stop()\n self.playwright = None\n\n\nclass ScrapeResult:\n doc: Document | None = None\n retry: bool = False\n\n\nWEB_CONNECTOR_MAX_SCROLL_ATTEMPTS = 20\n# Threshold for determining when to replace vs append iframe content\nIFRAME_TEXT_LENGTH_THRESHOLD = 700\n# Message indicating JavaScript is disabled, which often appears when scraping fails\nJAVASCRIPT_DISABLED_MESSAGE = \"You have JavaScript disabled in your browser\"\n# Grace period after page navigation to allow bot-detection challenges\n# and SPA content rendering to complete\nPAGE_RENDER_TIMEOUT_MS = 5000\n\n# Define common headers that mimic a real browser\nDEFAULT_USER_AGENT = \"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36\"\nDEFAULT_HEADERS = {\n \"User-Agent\": DEFAULT_USER_AGENT,\n \"Accept\": (\n \"text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,\"\n \"image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\"\n ),\n \"Accept-Language\": \"en-US,en;q=0.9\",\n # Brotli decoding has been flaky in brotlicffi/httpx for certain chunked responses;\n # stick to gzip/deflate to keep connectivity checks stable.\n \"Accept-Encoding\": \"gzip, deflate\",\n \"Connection\": \"keep-alive\",\n \"Upgrade-Insecure-Requests\": \"1\",\n \"Sec-Fetch-Dest\": \"document\",\n \"Sec-Fetch-Mode\": \"navigate\",\n \"Sec-Fetch-Site\": \"none\",\n \"Sec-Fetch-User\": \"?1\",\n \"Sec-CH-UA\": '\"Google Chrome\";v=\"123\", \"Not:A-Brand\";v=\"8\"',\n \"Sec-CH-UA-Mobile\": \"?0\",\n \"Sec-CH-UA-Platform\": '\"macOS\"',\n}\n\n\nclass WEB_CONNECTOR_VALID_SETTINGS(str, Enum):\n # Given a base site, index everything under that path\n RECURSIVE = \"recursive\"\n # Given a URL, index only the given page\n SINGLE = \"single\"\n # Given a sitemap.xml URL, parse all the pages in it\n SITEMAP = \"sitemap\"\n # Given a file upload where every line is a URL, parse all the URLs provided\n UPLOAD = \"upload\"\n\n\ndef protected_url_check(url: str) -> None:\n \"\"\"Couple considerations:\n - DNS mapping changes over time so we don't want to cache the results\n - Fetching this is assumed to be relatively fast compared to other bottlenecks like reading\n the page or embedding the contents\n - To be extra safe, all IPs associated with the URL must be global\n - This is to prevent misuse and not explicit attacks\n \"\"\"\n if not WEB_CONNECTOR_VALIDATE_URLS:\n return\n\n parse = urlparse(url)\n if parse.scheme != \"http\" and parse.scheme != \"https\":\n raise ValueError(\"URL must be of scheme https?://\")\n\n if not parse.hostname:\n raise ValueError(\"URL must include a hostname\")\n\n try:\n # This may give a large list of IP addresses for domains with extensive DNS configurations\n # such as large distributed systems of CDNs\n info = socket.getaddrinfo(parse.hostname, None)\n except socket.gaierror as e:\n raise ConnectionError(f\"DNS resolution failed for {parse.hostname}: {e}\")\n\n for address in info:\n ip = address[4][0]\n if not ipaddress.ip_address(ip).is_global:\n raise ValueError(\n f\"Non-global IP address detected: {ip}, skipping page {url}. \"\n f\"The Web Connector is not allowed to read loopback, link-local, or private ranges\"\n )\n\n\ndef check_internet_connection(url: str) -> None:\n try:\n # Use a more realistic browser-like request\n session = requests.Session()\n session.headers.update(DEFAULT_HEADERS)\n\n response = session.get(url, timeout=5, allow_redirects=True)\n\n response.raise_for_status()\n except requests.exceptions.HTTPError as e:\n # Extract status code from the response, defaulting to -1 if response is None\n status_code = e.response.status_code if e.response is not None else -1\n\n # For 403 errors, we do have internet connection, but the request is blocked by the server\n # this is usually due to bot detection. Future calls (via Playwright) will usually get\n # around this.\n if status_code == 403:\n logger.warning(\n f\"Received 403 Forbidden for {url}, will retry with browser automation\"\n )\n return\n\n error_msg = {\n 400: \"Bad Request\",\n 401: \"Unauthorized\",\n 403: \"Forbidden\",\n 404: \"Not Found\",\n 500: \"Internal Server Error\",\n 502: \"Bad Gateway\",\n 503: \"Service Unavailable\",\n 504: \"Gateway Timeout\",\n }.get(status_code, \"HTTP Error\")\n raise Exception(f\"{error_msg} ({status_code}) for {url} - {e}\")\n except requests.exceptions.SSLError as e:\n cause = (\n e.args[0].reason\n if isinstance(e.args, tuple) and isinstance(e.args[0], MaxRetryError)\n else e.args\n )\n raise Exception(f\"SSL error {str(cause)}\")\n except (requests.RequestException, ValueError) as e:\n raise Exception(f\"Unable to reach {url} - check your internet connection: {e}\")\n\n\ndef is_valid_url(url: str) -> bool:\n try:\n result = urlparse(url)\n return all([result.scheme, result.netloc])\n except ValueError:\n return False\n\n\ndef _same_site(base_url: str, candidate_url: str) -> bool:\n base, candidate = urlparse(base_url), urlparse(candidate_url)\n base_netloc = base.netloc.lower().removeprefix(\"www.\")\n candidate_netloc = candidate.netloc.lower().removeprefix(\"www.\")\n if base_netloc != candidate_netloc:\n return False\n\n base_path = (base.path or \"/\").rstrip(\"/\")\n if base_path in (\"\", \"/\"):\n return True\n\n candidate_path = candidate.path or \"/\"\n if candidate_path == base_path:\n return True\n\n boundary = f\"{base_path}/\"\n return candidate_path.startswith(boundary)\n\n\ndef get_internal_links(\n base_url: str, url: str, soup: BeautifulSoup, should_ignore_pound: bool = True\n) -> set[str]:\n internal_links = set()\n for link in cast(list[dict[str, Any]], soup.find_all(\"a\")):\n href = cast(str | None, link.get(\"href\"))\n if not href:\n continue\n\n # Account for malformed backslashes in URLs\n href = href.replace(\"\\\\\", \"/\")\n\n # \"#!\" indicates the page is using a hashbang URL, which is a client-side routing technique\n if should_ignore_pound and \"#\" in href and \"#!\" not in href:\n href = href.split(\"#\")[0]\n\n if not is_valid_url(href):\n # Relative path handling\n href = urljoin(url, href)\n\n if _same_site(base_url, href):\n internal_links.add(href)\n return internal_links\n\n\ndef start_playwright() -> Tuple[Playwright, BrowserContext]:\n playwright = sync_playwright().start()\n\n # Launch browser with more realistic settings\n browser = playwright.chromium.launch(\n headless=True,\n args=[\n \"--disable-blink-features=AutomationControlled\",\n \"--disable-features=IsolateOrigins,site-per-process\",\n \"--disable-site-isolation-trials\",\n ],\n )\n\n # Create a context with realistic browser properties\n context = browser.new_context(\n user_agent=DEFAULT_USER_AGENT,\n viewport={\"width\": 1440, \"height\": 900},\n device_scale_factor=2.0,\n locale=\"en-US\",\n timezone_id=\"America/Los_Angeles\",\n has_touch=False,\n java_script_enabled=True,\n color_scheme=\"light\",\n # Add more realistic browser properties\n bypass_csp=True,\n ignore_https_errors=True,\n )\n\n # Set additional headers to mimic a real browser\n context.set_extra_http_headers(\n {\n \"Accept\": DEFAULT_HEADERS[\"Accept\"],\n \"Accept-Language\": DEFAULT_HEADERS[\"Accept-Language\"],\n \"Sec-Fetch-Dest\": DEFAULT_HEADERS[\"Sec-Fetch-Dest\"],\n \"Sec-Fetch-Mode\": DEFAULT_HEADERS[\"Sec-Fetch-Mode\"],\n \"Sec-Fetch-Site\": DEFAULT_HEADERS[\"Sec-Fetch-Site\"],\n \"Sec-Fetch-User\": DEFAULT_HEADERS[\"Sec-Fetch-User\"],\n \"Sec-CH-UA\": DEFAULT_HEADERS[\"Sec-CH-UA\"],\n \"Sec-CH-UA-Mobile\": DEFAULT_HEADERS[\"Sec-CH-UA-Mobile\"],\n \"Sec-CH-UA-Platform\": DEFAULT_HEADERS[\"Sec-CH-UA-Platform\"],\n \"Cache-Control\": \"max-age=0\",\n \"DNT\": \"1\",\n }\n )\n\n # Add a script to modify navigator properties to avoid detection\n context.add_init_script(\n \"\"\"\n Object.defineProperty(navigator, 'webdriver', {\n get: () => undefined\n });\n Object.defineProperty(navigator, 'plugins', {\n get: () => [1, 2, 3, 4, 5]\n });\n Object.defineProperty(navigator, 'languages', {\n get: () => ['en-US', 'en']\n });\n \"\"\"\n )\n\n if (\n WEB_CONNECTOR_OAUTH_CLIENT_ID\n and WEB_CONNECTOR_OAUTH_CLIENT_SECRET\n and WEB_CONNECTOR_OAUTH_TOKEN_URL\n ):\n client = BackendApplicationClient(client_id=WEB_CONNECTOR_OAUTH_CLIENT_ID)\n oauth = OAuth2Session(client=client)\n token = oauth.fetch_token(\n token_url=WEB_CONNECTOR_OAUTH_TOKEN_URL,\n client_id=WEB_CONNECTOR_OAUTH_CLIENT_ID,\n client_secret=WEB_CONNECTOR_OAUTH_CLIENT_SECRET,\n )\n context.set_extra_http_headers(\n {\"Authorization\": \"Bearer {}\".format(token[\"access_token\"])}\n )\n\n return playwright, context\n\n\ndef extract_urls_from_sitemap(sitemap_url: str) -> list[str]:\n # requests should handle brotli compression automatically\n # as long as the brotli package is available in the venv. Leaving this line here to avoid\n # a regression as someone says \"Ah, looks like this brotli package isn't used anywhere, let's remove it\"\n # import brotli\n try:\n response = requests.get(sitemap_url, headers=DEFAULT_HEADERS)\n response.raise_for_status()\n soup = BeautifulSoup(response.content, \"html.parser\")\n urls = [\n _ensure_absolute_url(sitemap_url, loc_tag.text)\n for loc_tag in soup.find_all(\"loc\")\n ]\n\n if len(urls) == 0 and len(soup.find_all(\"urlset\")) == 0:\n # the given url doesn't look like a sitemap, let's try to find one\n urls = list_pages_for_site(sitemap_url)\n\n if len(urls) == 0:\n raise ValueError(\n f\"No URLs found in sitemap {sitemap_url}. Try using the 'single' or 'recursive' scraping options instead.\"\n )\n\n return urls\n except requests.RequestException as e:\n raise RuntimeError(f\"Failed to fetch sitemap from {sitemap_url}: {e}\")\n except ValueError as e:\n raise RuntimeError(f\"Error processing sitemap {sitemap_url}: {e}\")\n except Exception as e:\n raise RuntimeError(\n f\"Unexpected error while processing sitemap {sitemap_url}: {e}\"\n )\n\n\ndef _ensure_absolute_url(source_url: str, maybe_relative_url: str) -> str:\n if not urlparse(maybe_relative_url).netloc:\n return urljoin(source_url, maybe_relative_url)\n return maybe_relative_url\n\n\ndef _ensure_valid_url(url: str) -> str:\n if \"://\" not in url:\n return \"https://\" + url\n return url\n\n\ndef _read_urls_file(location: str) -> list[str]:\n with open(location, \"r\") as f:\n urls = [_ensure_valid_url(line.strip()) for line in f if line.strip()]\n return urls\n\n\ndef _get_datetime_from_last_modified_header(last_modified: str) -> datetime | None:\n try:\n return datetime.strptime(last_modified, \"%a, %d %b %Y %H:%M:%S %Z\").replace(\n tzinfo=timezone.utc\n )\n except (ValueError, TypeError):\n return None\n\n\ndef _handle_cookies(context: BrowserContext, url: str) -> None:\n \"\"\"Handle cookies for the given URL to help with bot detection\"\"\"\n try:\n # Parse the URL to get the domain\n parsed_url = urlparse(url)\n domain = parsed_url.netloc\n\n # Add some common cookies that might help with bot detection\n cookies: list[dict[str, str]] = [\n {\n \"name\": \"cookieconsent\",\n \"value\": \"accepted\",\n \"domain\": domain,\n \"path\": \"/\",\n },\n {\n \"name\": \"consent\",\n \"value\": \"true\",\n \"domain\": domain,\n \"path\": \"/\",\n },\n {\n \"name\": \"session\",\n \"value\": \"random_session_id\",\n \"domain\": domain,\n \"path\": \"/\",\n },\n ]\n\n # Add cookies to the context\n for cookie in cookies:\n try:\n context.add_cookies([cookie]) # type: ignore\n except Exception as e:\n logger.debug(f\"Failed to add cookie {cookie['name']} for {domain}: {e}\")\n except Exception:\n logger.exception(\n f\"Unexpected error while handling cookies for Web Connector with URL {url}\"\n )\n\n\nclass WebConnector(LoadConnector):\n MAX_RETRIES = 3\n\n def __init__(\n self,\n base_url: str, # Can't change this without disrupting existing users\n web_connector_type: str = WEB_CONNECTOR_VALID_SETTINGS.RECURSIVE.value,\n mintlify_cleanup: bool = True, # Mostly ok to apply to other websites as well\n batch_size: int = INDEX_BATCH_SIZE,\n scroll_before_scraping: bool = False,\n **kwargs: Any, # noqa: ARG002\n ) -> None:\n self.mintlify_cleanup = mintlify_cleanup\n self.batch_size = batch_size\n self.recursive = False\n self.scroll_before_scraping = scroll_before_scraping\n self.web_connector_type = web_connector_type\n if web_connector_type == WEB_CONNECTOR_VALID_SETTINGS.RECURSIVE.value:\n self.recursive = True\n self.to_visit_list = [_ensure_valid_url(base_url)]\n return\n\n elif web_connector_type == WEB_CONNECTOR_VALID_SETTINGS.SINGLE.value:\n self.to_visit_list = [_ensure_valid_url(base_url)]\n\n elif web_connector_type == WEB_CONNECTOR_VALID_SETTINGS.SITEMAP:\n self.to_visit_list = extract_urls_from_sitemap(_ensure_valid_url(base_url))\n\n elif web_connector_type == WEB_CONNECTOR_VALID_SETTINGS.UPLOAD:\n # Explicitly check if running in multi-tenant mode to prevent potential security risks\n if MULTI_TENANT:\n raise ValueError(\n \"Upload input for web connector is not supported in cloud environments\"\n )\n\n logger.warning(\n \"This is not a UI supported Web Connector flow, are you sure you want to do this?\"\n )\n self.to_visit_list = _read_urls_file(base_url)\n\n else:\n raise ValueError(\n \"Invalid Web Connector Config, must choose a valid type between: \"\n )\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n if credentials:\n logger.warning(\"Unexpected credentials provided for Web Connector\")\n return None\n\n def _do_scrape(\n self,\n index: int,\n initial_url: str,\n session_ctx: ScrapeSessionContext,\n ) -> ScrapeResult:\n \"\"\"Returns a ScrapeResult object with a doc and retry flag.\"\"\"\n\n if session_ctx.playwright is None:\n raise RuntimeError(\"scrape_context.playwright is None\")\n\n if session_ctx.playwright_context is None:\n raise RuntimeError(\"scrape_context.playwright_context is None\")\n\n result = ScrapeResult()\n\n # Handle cookies for the URL\n _handle_cookies(session_ctx.playwright_context, initial_url)\n\n # First do a HEAD request to check content type without downloading the entire content\n head_response = requests.head(\n initial_url, headers=DEFAULT_HEADERS, allow_redirects=True\n )\n content_type = head_response.headers.get(\"content-type\")\n is_pdf = is_pdf_resource(initial_url, content_type)\n\n if is_pdf:\n # PDF files are not checked for links\n response = requests.get(initial_url, headers=DEFAULT_HEADERS)\n page_text, metadata = extract_pdf_text(response.content)\n last_modified = response.headers.get(\"Last-Modified\")\n\n result.doc = Document(\n id=initial_url,\n sections=[TextSection(link=initial_url, text=page_text)],\n source=DocumentSource.WEB,\n semantic_identifier=initial_url.rstrip(\"/\").split(\"/\")[-1]\n or initial_url,\n metadata=metadata,\n doc_updated_at=(\n _get_datetime_from_last_modified_header(last_modified)\n if last_modified\n else None\n ),\n )\n\n return result\n\n page = session_ctx.playwright_context.new_page()\n try:\n # Use \"commit\" instead of \"domcontentloaded\" to avoid hanging on bot-detection pages\n # that may never fire domcontentloaded. \"commit\" waits only for navigation to be\n # committed (response received), then we add a short wait for initial rendering.\n page_response = page.goto(\n initial_url,\n timeout=30000, # 30 seconds\n wait_until=\"commit\", # Wait for navigation to commit\n )\n # Give the page a moment to start rendering after navigation commits.\n # Allows CloudFlare and other bot-detection challenges to complete.\n page.wait_for_timeout(PAGE_RENDER_TIMEOUT_MS)\n\n # Wait for network activity to settle so SPAs that fetch content\n # asynchronously after the initial JS bundle have time to render.\n try:\n # A bit of extra time to account for long-polling, websockets, etc.\n page.wait_for_load_state(\"networkidle\", timeout=PAGE_RENDER_TIMEOUT_MS)\n except TimeoutError:\n pass\n\n last_modified = (\n page_response.header_value(\"Last-Modified\") if page_response else None\n )\n final_url = page.url\n if final_url != initial_url:\n protected_url_check(final_url)\n initial_url = final_url\n if initial_url in session_ctx.visited_links:\n logger.info(\n f\"{index}: {initial_url} redirected to {final_url} - already indexed\"\n )\n page.close()\n return result\n\n logger.info(f\"{index}: {initial_url} redirected to {final_url}\")\n session_ctx.visited_links.add(initial_url)\n\n # If we got here, the request was successful\n if self.scroll_before_scraping:\n scroll_attempts = 0\n previous_height = page.evaluate(\"document.body.scrollHeight\")\n while scroll_attempts < WEB_CONNECTOR_MAX_SCROLL_ATTEMPTS:\n page.evaluate(\"window.scrollTo(0, document.body.scrollHeight)\")\n # Wait for content to load, but catch timeout if page never reaches networkidle\n # (e.g., CloudFlare protection keeps making requests)\n try:\n page.wait_for_load_state(\n \"networkidle\", timeout=PAGE_RENDER_TIMEOUT_MS\n )\n except TimeoutError:\n # If networkidle times out, just give it a moment for content to render\n time.sleep(1)\n time.sleep(0.5) # let javascript run\n\n new_height = page.evaluate(\"document.body.scrollHeight\")\n if new_height == previous_height:\n break # Stop scrolling when no more content is loaded\n previous_height = new_height\n scroll_attempts += 1\n\n content = page.content()\n soup = BeautifulSoup(content, \"html.parser\")\n\n if self.recursive:\n internal_links = get_internal_links(\n session_ctx.base_url, initial_url, soup\n )\n for link in internal_links:\n if link not in session_ctx.visited_links:\n session_ctx.to_visit.append(link)\n\n if page_response and str(page_response.status)[0] in (\"4\", \"5\"):\n session_ctx.last_error = f\"Skipped indexing {initial_url} due to HTTP {page_response.status} response\"\n logger.info(session_ctx.last_error)\n result.retry = True\n return result\n\n # after this point, we don't need the caller to retry\n parsed_html = web_html_cleanup(soup, self.mintlify_cleanup)\n\n \"\"\"For websites containing iframes that need to be scraped,\n the code below can extract text from within these iframes.\n \"\"\"\n logger.debug(\n f\"{index}: Length of cleaned text {len(parsed_html.cleaned_text)}\"\n )\n if JAVASCRIPT_DISABLED_MESSAGE in parsed_html.cleaned_text:\n iframe_count = page.frame_locator(\"iframe\").locator(\"html\").count()\n if iframe_count > 0:\n iframe_texts = (\n page.frame_locator(\"iframe\").locator(\"html\").all_inner_texts()\n )\n document_text = \"\\n\".join(iframe_texts)\n \"\"\" 700 is the threshold value for the length of the text extracted\n from the iframe based on the issue faced \"\"\"\n if len(parsed_html.cleaned_text) < IFRAME_TEXT_LENGTH_THRESHOLD:\n parsed_html.cleaned_text = document_text\n else:\n parsed_html.cleaned_text += \"\\n\" + document_text\n\n # Sometimes pages with #! will serve duplicate content\n # There are also just other ways this can happen\n hashed_text = hash((parsed_html.title, parsed_html.cleaned_text))\n if hashed_text in session_ctx.content_hashes:\n logger.info(\n f\"{index}: Skipping duplicate title + content for {initial_url}\"\n )\n return result\n\n session_ctx.content_hashes.add(hashed_text)\n\n result.doc = Document(\n id=initial_url,\n sections=[TextSection(link=initial_url, text=parsed_html.cleaned_text)],\n source=DocumentSource.WEB,\n semantic_identifier=parsed_html.title or initial_url,\n metadata={},\n doc_updated_at=(\n _get_datetime_from_last_modified_header(last_modified)\n if last_modified\n else None\n ),\n )\n finally:\n page.close()\n\n return result\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n \"\"\"Traverses through all pages found on the website\n and converts them into documents\"\"\"\n\n if not self.to_visit_list:\n raise ValueError(\"No URLs to visit\")\n\n base_url = self.to_visit_list[0] # For the recursive case\n check_internet_connection(base_url) # make sure we can connect to the base url\n\n session_ctx = ScrapeSessionContext(base_url, self.to_visit_list)\n session_ctx.initialize()\n\n while session_ctx.to_visit:\n initial_url = session_ctx.to_visit.pop()\n if initial_url in session_ctx.visited_links:\n continue\n session_ctx.visited_links.add(initial_url)\n\n try:\n protected_url_check(initial_url)\n except Exception as e:\n session_ctx.last_error = f\"Invalid URL {initial_url} due to {e}\"\n logger.warning(session_ctx.last_error)\n continue\n\n index = len(session_ctx.visited_links)\n logger.info(f\"{index}: Visiting {initial_url}\")\n\n # Add retry mechanism with exponential backoff\n retry_count = 0\n\n while retry_count < self.MAX_RETRIES:\n if retry_count > 0:\n # Add a random delay between retries (exponential backoff)\n delay = min(2**retry_count + random.uniform(0, 1), 10)\n logger.info(\n f\"Retry {retry_count}/{self.MAX_RETRIES} for {initial_url} after {delay:.2f}s delay\"\n )\n time.sleep(delay)\n\n try:\n result = self._do_scrape(index, initial_url, session_ctx)\n if result.retry:\n continue\n\n if result.doc:\n session_ctx.doc_batch.append(result.doc)\n except Exception as e:\n session_ctx.last_error = f\"Failed to fetch '{initial_url}': {e}\"\n logger.exception(session_ctx.last_error)\n session_ctx.initialize()\n continue\n finally:\n retry_count += 1\n\n break # success / don't retry\n\n if len(session_ctx.doc_batch) >= self.batch_size:\n session_ctx.initialize()\n session_ctx.at_least_one_doc = True\n yield session_ctx.doc_batch\n session_ctx.doc_batch = []\n\n if session_ctx.doc_batch:\n session_ctx.stop()\n session_ctx.at_least_one_doc = True\n yield session_ctx.doc_batch\n\n if not session_ctx.at_least_one_doc:\n if session_ctx.last_error:\n raise RuntimeError(session_ctx.last_error)\n raise RuntimeError(\"No valid pages found.\")\n\n session_ctx.stop()\n\n def validate_connector_settings(self) -> None:\n # Make sure we have at least one valid URL to check\n if not self.to_visit_list:\n raise ConnectorValidationError(\n \"No URL configured. Please provide at least one valid URL.\"\n )\n\n if (\n self.web_connector_type == WEB_CONNECTOR_VALID_SETTINGS.SITEMAP.value\n or self.web_connector_type == WEB_CONNECTOR_VALID_SETTINGS.RECURSIVE.value\n ):\n return None\n\n # We'll just test the first URL for connectivity and correctness\n test_url = self.to_visit_list[0]\n\n # Check that the URL is allowed and well-formed\n try:\n protected_url_check(test_url)\n except ValueError as e:\n raise ConnectorValidationError(\n f\"Protected URL check failed for '{test_url}': {e}\"\n )\n except ConnectionError as e:\n # Typically DNS or other network issues\n raise ConnectorValidationError(str(e))\n\n # Make a quick request to see if we get a valid response\n try:\n check_internet_connection(test_url)\n except Exception as e:\n err_str = str(e)\n if \"401\" in err_str:\n raise CredentialExpiredError(\n f\"Unauthorized access to '{test_url}': {e}\"\n )\n elif \"403\" in err_str:\n raise InsufficientPermissionsError(\n f\"Forbidden access to '{test_url}': {e}\"\n )\n elif \"404\" in err_str:\n raise ConnectorValidationError(f\"Page not found for '{test_url}': {e}\")\n elif \"Max retries exceeded\" in err_str and \"NameResolutionError\" in err_str:\n raise ConnectorValidationError(\n f\"Unable to resolve hostname for '{test_url}'. Please check the URL and your internet connection.\"\n )\n else:\n # Could be a 5xx or another error, treat as unexpected\n raise UnexpectedValidationError(\n f\"Unexpected error validating '{test_url}': {e}\"\n )\n\n\nif __name__ == \"__main__\":\n connector = WebConnector(\"https://docs.onyx.app/\")\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/wikipedia/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/wikipedia/connector.py", "content": "from typing import ClassVar\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.mediawiki import wiki\n\n\nclass WikipediaConnector(wiki.MediaWikiConnector):\n \"\"\"Connector for Wikipedia.\"\"\"\n\n document_source_type: ClassVar[DocumentSource] = DocumentSource.WIKIPEDIA\n\n def __init__(\n self,\n categories: list[str],\n pages: list[str],\n recurse_depth: int,\n language_code: str = \"en\",\n batch_size: int = INDEX_BATCH_SIZE,\n ) -> None:\n super().__init__(\n hostname=\"wikipedia.org\",\n categories=categories,\n pages=pages,\n recurse_depth=recurse_depth,\n language_code=language_code,\n batch_size=batch_size,\n )\n" }, { "path": "backend/onyx/connectors/xenforo/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/xenforo/connector.py", "content": "\"\"\"\nThis is the XenforoConnector class. It is used to connect to a Xenforo forum and load or update documents from the forum.\n\nTo use this class, you need to provide the URL of the Xenforo forum board you want to connect to when creating an instance\nof the class. The URL should be a string that starts with 'http://' or 'https://', followed by the domain name of the\nforum, followed by the board name. For example:\n\n base_url = 'https://www.example.com/forum/boards/some-topic/'\n\nThe `load_from_state` method is used to load documents from the forum. It takes an optional `state` parameter, which\ncan be used to specify a state from which to start loading documents.\n\"\"\"\n\nimport re\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom urllib.parse import urlparse\n\nimport pytz\nimport requests\nfrom bs4 import BeautifulSoup\nfrom bs4 import Tag\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import datetime_to_utc\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef get_title(soup: BeautifulSoup) -> str:\n el = soup.find(\"h1\", \"p-title-value\")\n if not el:\n return \"\"\n title = el.text\n for char in (\";\", \":\", \"!\", \"*\", \"/\", \"\\\\\", \"?\", '\"', \"<\", \">\", \"|\"):\n title = title.replace(char, \"_\")\n return title\n\n\ndef get_pages(soup: BeautifulSoup, url: str) -> list[str]:\n page_tags = soup.select(\"li.pageNav-page\")\n page_numbers = []\n for button in page_tags:\n if re.match(r\"^\\d+$\", button.text):\n page_numbers.append(button.text)\n\n max_pages = int(max(page_numbers, key=int)) if page_numbers else 1\n\n all_pages = []\n for x in range(1, int(max_pages) + 1):\n all_pages.append(f\"{url}page-{x}\")\n return all_pages\n\n\ndef parse_post_date(post_element: BeautifulSoup) -> datetime:\n el = post_element.find(\"time\")\n if not isinstance(el, Tag) or \"datetime\" not in el.attrs:\n return datetime.utcfromtimestamp(0).replace(tzinfo=timezone.utc)\n\n date_value = el[\"datetime\"]\n\n # Ensure date_value is a string (if it's a list, take the first element)\n if isinstance(date_value, list):\n date_value = date_value[0]\n\n post_date = datetime.strptime(date_value, \"%Y-%m-%dT%H:%M:%S%z\")\n return datetime_to_utc(post_date)\n\n\ndef scrape_page_posts(\n soup: BeautifulSoup,\n page_index: int,\n url: str,\n initial_run: bool,\n start_time: datetime,\n) -> list:\n title = get_title(soup)\n\n documents = []\n for post in soup.find_all(\"div\", class_=\"message-inner\"):\n post_date = parse_post_date(post)\n if initial_run or post_date > start_time:\n el = post.find(\"div\", class_=\"bbWrapper\")\n if not el:\n continue\n post_text = el.get_text(strip=True) + \"\\n\"\n author_tag = post.find(\"a\", class_=\"username\")\n if author_tag is None:\n author_tag = post.find(\"span\", class_=\"username\")\n author = author_tag.get_text(strip=True) if author_tag else \"Deleted author\"\n formatted_time = post_date.strftime(\"%Y-%m-%d %H:%M:%S\")\n\n # TODO: if a caller calls this for each page of a thread, it may see the\n # same post multiple times if there is a sticky post\n # that appears on each page of a thread.\n # it's important to generate unique doc id's, so page index is part of the\n # id. We may want to de-dupe this stuff inside the indexing service.\n document = Document(\n id=f\"{DocumentSource.XENFORO.value}_{title}_{page_index}_{formatted_time}\",\n sections=[TextSection(link=url, text=post_text)],\n title=title,\n source=DocumentSource.XENFORO,\n semantic_identifier=title,\n primary_owners=[BasicExpertInfo(display_name=author)],\n metadata={\n \"type\": \"post\",\n \"author\": author,\n \"time\": formatted_time,\n },\n doc_updated_at=post_date,\n )\n\n documents.append(document)\n return documents\n\n\nclass XenforoConnector(LoadConnector):\n # Class variable to track if the connector has been run before\n has_been_run_before = False\n\n def __init__(self, base_url: str) -> None:\n self.base_url = base_url\n self.initial_run = not XenforoConnector.has_been_run_before\n self.start = datetime.utcnow().replace(tzinfo=pytz.utc) - timedelta(days=1)\n self.cookies: dict[str, str] = {}\n # mimic user browser to avoid being blocked by the website (see: https://www.useragents.me/)\n self.headers = {\n \"User-Agent\": \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) \"\n \"AppleWebKit/537.36 (KHTML, like Gecko) \"\n \"Chrome/121.0.0.0 Safari/537.36\"\n }\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n if credentials:\n logger.warning(\"Unexpected credentials provided for Xenforo Connector\")\n return None\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n # Standardize URL to always end in /.\n if self.base_url[-1] != \"/\":\n self.base_url += \"/\"\n\n # Remove all extra parameters from the end such as page, post.\n matches = (\"threads/\", \"boards/\", \"forums/\")\n for each in matches:\n if each in self.base_url:\n try:\n self.base_url = self.base_url[\n 0 : self.base_url.index(\n \"/\", self.base_url.index(each) + len(each)\n )\n + 1\n ]\n except ValueError:\n pass\n\n doc_batch: list[Document | HierarchyNode] = []\n all_threads = []\n\n # If the URL contains \"boards/\" or \"forums/\", find all threads.\n if \"boards/\" in self.base_url or \"forums/\" in self.base_url:\n pages = get_pages(self.requestsite(self.base_url), self.base_url)\n\n # Get all pages on thread_list_page\n for pre_count, thread_list_page in enumerate(pages, start=1):\n logger.info(\n f\"Getting pages from thread_list_page.. Current: {pre_count}/{len(pages)}\\r\"\n )\n all_threads += self.get_threads(thread_list_page)\n # If the URL contains \"threads/\", add the thread to the list.\n elif \"threads/\" in self.base_url:\n all_threads.append(self.base_url)\n\n # Process all threads\n for thread_count, thread_url in enumerate(all_threads, start=1):\n soup = self.requestsite(thread_url)\n if soup is None:\n logger.error(f\"Failed to load page: {self.base_url}\")\n continue\n pages = get_pages(soup, thread_url)\n # Getting all pages for all threads\n for page_index, page in enumerate(pages, start=1):\n logger.info(\n f\"Progress: Page {page_index}/{len(pages)} - Thread {thread_count}/{len(all_threads)}\\r\"\n )\n soup_page = self.requestsite(page)\n doc_batch.extend(\n scrape_page_posts(\n soup_page, page_index, thread_url, self.initial_run, self.start\n )\n )\n if doc_batch:\n yield doc_batch\n\n # Mark the initial run finished after all threads and pages have been processed\n XenforoConnector.has_been_run_before = True\n\n def get_threads(self, url: str) -> list[str]:\n soup = self.requestsite(url)\n thread_tags = soup.find_all(class_=\"structItem-title\")\n base_url = \"{uri.scheme}://{uri.netloc}\".format(uri=urlparse(url))\n threads = []\n for x in thread_tags:\n y = x.find_all(href=True)\n for element in y:\n link = element[\"href\"]\n if \"threads/\" in link:\n stripped = link[0 : link.rfind(\"/\") + 1]\n if base_url + stripped not in threads:\n threads.append(base_url + stripped)\n return threads\n\n def requestsite(self, url: str) -> BeautifulSoup:\n try:\n response = requests.get(\n url, cookies=self.cookies, headers=self.headers, timeout=10\n )\n if response.status_code != 200:\n logger.error(\n f\"<{url}> Request Error: {response.status_code} - {response.reason}\"\n )\n return BeautifulSoup(response.text, \"html.parser\")\n except TimeoutError:\n logger.error(\"Timed out Error.\")\n except Exception as e:\n logger.error(f\"Error on {url}\")\n logger.exception(e)\n return BeautifulSoup(\"\", \"html.parser\")\n\n\nif __name__ == \"__main__\":\n connector = XenforoConnector(\n # base_url=\"https://cassiopaea.org/forum/threads/how-to-change-your-emotional-state.41381/\"\n base_url=\"https://xenforo.com/community/threads/whats-new-with-enhanced-search-resource-manager-and-media-gallery-in-xenforo-2-3.220935/\"\n )\n document_batches = connector.load_from_state()\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/zendesk/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/zendesk/connector.py", "content": "import copy\nimport time\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom typing import Any\nfrom typing import cast\n\nimport requests\nfrom pydantic import BaseModel\nfrom requests.exceptions import HTTPError\nfrom typing_extensions import override\n\nfrom onyx.configs.app_configs import ZENDESK_CONNECTOR_SKIP_ARTICLE_LABELS\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n time_str_to_utc,\n)\nfrom onyx.connectors.cross_connector_utils.rate_limit_wrapper import (\n rate_limit_builder,\n)\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.exceptions import InsufficientPermissionsError\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CheckpointOutput\nfrom onyx.connectors.interfaces import ConnectorFailure\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import BasicExpertInfo\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.utils.retry_wrapper import retry_builder\n\n\nMAX_PAGE_SIZE = 30 # Zendesk API maximum\nMAX_AUTHOR_MAP_SIZE = 50_000 # Reset author map cache if it gets too large\n_SLIM_BATCH_SIZE = 1000\n\n\nclass ZendeskCredentialsNotSetUpError(PermissionError):\n def __init__(self) -> None:\n super().__init__(\n \"Zendesk Credentials are not set up, was load_credentials called?\"\n )\n\n\nclass ZendeskClient:\n def __init__(\n self,\n subdomain: str,\n email: str,\n token: str,\n calls_per_minute: int | None = None,\n ):\n self.base_url = f\"https://{subdomain}.zendesk.com/api/v2\"\n self.auth = (f\"{email}/token\", token)\n self.make_request = request_with_rate_limit(self, calls_per_minute)\n\n\ndef request_with_rate_limit(\n client: ZendeskClient, max_calls_per_minute: int | None = None\n) -> Callable[[str, dict[str, Any]], dict[str, Any]]:\n @retry_builder()\n @(\n rate_limit_builder(max_calls=max_calls_per_minute, period=60)\n if max_calls_per_minute\n else lambda x: x\n )\n def make_request(endpoint: str, params: dict[str, Any]) -> dict[str, Any]:\n response = requests.get(\n f\"{client.base_url}/{endpoint}\", auth=client.auth, params=params\n )\n\n if response.status_code == 429:\n retry_after = response.headers.get(\"Retry-After\")\n if retry_after is not None:\n # Sleep for the duration indicated by the Retry-After header\n time.sleep(int(retry_after))\n\n elif (\n response.status_code == 403\n and response.json().get(\"error\") == \"SupportProductInactive\"\n ):\n return response.json()\n\n response.raise_for_status()\n return response.json()\n\n return make_request\n\n\nclass ZendeskPageResponse(BaseModel):\n data: list[dict[str, Any]]\n meta: dict[str, Any]\n has_more: bool\n\n\ndef _get_content_tag_mapping(client: ZendeskClient) -> dict[str, str]:\n content_tags: dict[str, str] = {}\n params = {\"page[size]\": MAX_PAGE_SIZE}\n\n try:\n while True:\n data = client.make_request(\"guide/content_tags\", params)\n\n for tag in data.get(\"records\", []):\n content_tags[tag[\"id\"]] = tag[\"name\"]\n\n # Check if there are more pages\n if data.get(\"meta\", {}).get(\"has_more\", False):\n params[\"page[after]\"] = data[\"meta\"][\"after_cursor\"]\n else:\n break\n\n return content_tags\n except Exception as e:\n raise Exception(f\"Error fetching content tags: {str(e)}\")\n\n\ndef _get_articles(\n client: ZendeskClient, start_time: int | None = None, page_size: int = MAX_PAGE_SIZE\n) -> Iterator[dict[str, Any]]:\n params = {\"page[size]\": page_size, \"sort_by\": \"updated_at\", \"sort_order\": \"asc\"}\n if start_time is not None:\n params[\"start_time\"] = start_time\n\n while True:\n data = client.make_request(\"help_center/articles\", params)\n for article in data[\"articles\"]:\n yield article\n\n if not data.get(\"meta\", {}).get(\"has_more\"):\n break\n params[\"page[after]\"] = data[\"meta\"][\"after_cursor\"]\n\n\ndef _get_article_page(\n client: ZendeskClient,\n start_time: int | None = None,\n after_cursor: str | None = None,\n page_size: int = MAX_PAGE_SIZE,\n) -> ZendeskPageResponse:\n params = {\"page[size]\": page_size, \"sort_by\": \"updated_at\", \"sort_order\": \"asc\"}\n if start_time is not None:\n params[\"start_time\"] = start_time\n if after_cursor is not None:\n params[\"page[after]\"] = after_cursor\n\n data = client.make_request(\"help_center/articles\", params)\n return ZendeskPageResponse(\n data=data[\"articles\"],\n meta=data[\"meta\"],\n has_more=bool(data[\"meta\"].get(\"has_more\", False)),\n )\n\n\ndef _get_tickets(\n client: ZendeskClient, start_time: int | None = None\n) -> Iterator[dict[str, Any]]:\n params = {\"start_time\": start_time or 0}\n\n while True:\n data = client.make_request(\"incremental/tickets.json\", params)\n for ticket in data[\"tickets\"]:\n yield ticket\n\n if not data.get(\"end_of_stream\", False):\n params[\"start_time\"] = data[\"end_time\"]\n else:\n break\n\n\n# TODO: maybe these don't need to be their own functions?\ndef _get_tickets_page(\n client: ZendeskClient, start_time: int | None = None\n) -> ZendeskPageResponse:\n params = {\"start_time\": start_time or 0}\n\n # NOTE: for some reason zendesk doesn't seem to be respecting the start_time param\n # in my local testing with very few tickets. We'll look into it if this becomes an\n # issue in larger deployments\n data = client.make_request(\"incremental/tickets.json\", params)\n if data.get(\"error\") == \"SupportProductInactive\":\n raise ValueError(\n \"Zendesk Support Product is not active for this account, No tickets to index\"\n )\n return ZendeskPageResponse(\n data=data[\"tickets\"],\n meta={\"end_time\": data[\"end_time\"]},\n has_more=not bool(data.get(\"end_of_stream\", False)),\n )\n\n\ndef _fetch_author(\n client: ZendeskClient, author_id: str | int\n) -> BasicExpertInfo | None:\n # Skip fetching if author_id is invalid\n # cast to str to avoid issues with zendesk changing their types\n if not author_id or str(author_id) == \"-1\":\n return None\n\n try:\n author_data = client.make_request(f\"users/{author_id}\", {})\n user = author_data.get(\"user\")\n return (\n BasicExpertInfo(display_name=user.get(\"name\"), email=user.get(\"email\"))\n if user and user.get(\"name\") and user.get(\"email\")\n else None\n )\n except requests.exceptions.HTTPError:\n # Handle any API errors gracefully\n return None\n\n\ndef _article_to_document(\n article: dict[str, Any],\n content_tags: dict[str, str],\n author_map: dict[str, BasicExpertInfo],\n client: ZendeskClient,\n) -> tuple[dict[str, BasicExpertInfo] | None, Document]:\n author_id = article.get(\"author_id\")\n if not author_id:\n author = None\n else:\n author = (\n author_map.get(author_id)\n if author_id in author_map\n else _fetch_author(client, author_id)\n )\n\n new_author_mapping = {author_id: author} if author_id and author else None\n\n updated_at = article.get(\"updated_at\")\n update_time = time_str_to_utc(updated_at) if updated_at else None\n\n # Build metadata\n metadata: dict[str, str | list[str]] = {\n \"labels\": [str(label) for label in article.get(\"label_names\", []) if label],\n \"content_tags\": [\n content_tags[tag_id]\n for tag_id in article.get(\"content_tag_ids\", [])\n if tag_id in content_tags\n ],\n }\n\n # Remove empty values\n metadata = {k: v for k, v in metadata.items() if v}\n\n return new_author_mapping, Document(\n id=f\"article:{article['id']}\",\n sections=[\n TextSection(\n link=cast(str, article.get(\"html_url\")),\n text=parse_html_page_basic(article[\"body\"]),\n )\n ],\n source=DocumentSource.ZENDESK,\n semantic_identifier=article[\"title\"],\n doc_updated_at=update_time,\n primary_owners=[author] if author else None,\n metadata=metadata,\n )\n\n\ndef _get_comment_text(\n comment: dict[str, Any],\n author_map: dict[str, BasicExpertInfo],\n client: ZendeskClient,\n) -> tuple[dict[str, BasicExpertInfo] | None, str]:\n author_id = comment.get(\"author_id\")\n if not author_id:\n author = None\n else:\n author = (\n author_map.get(author_id)\n if author_id in author_map\n else _fetch_author(client, author_id)\n )\n\n new_author_mapping = {author_id: author} if author_id and author else None\n\n comment_text = f\"Comment{' by ' + author.display_name if author and author.display_name else ''}\"\n comment_text += f\"{' at ' + comment['created_at'] if comment.get('created_at') else ''}:\\n{comment['body']}\"\n\n return new_author_mapping, comment_text\n\n\ndef _ticket_to_document(\n ticket: dict[str, Any],\n author_map: dict[str, BasicExpertInfo],\n client: ZendeskClient,\n default_subdomain: str,\n) -> tuple[dict[str, BasicExpertInfo] | None, Document]:\n submitter_id = ticket.get(\"submitter\")\n if not submitter_id:\n submitter = None\n else:\n submitter = (\n author_map.get(submitter_id)\n if submitter_id in author_map\n else _fetch_author(client, submitter_id)\n )\n\n new_author_mapping = (\n {submitter_id: submitter} if submitter_id and submitter else None\n )\n\n updated_at = ticket.get(\"updated_at\")\n update_time = time_str_to_utc(updated_at) if updated_at else None\n\n metadata: dict[str, str | list[str]] = {}\n if status := ticket.get(\"status\"):\n metadata[\"status\"] = status\n if priority := ticket.get(\"priority\"):\n metadata[\"priority\"] = priority\n if tags := ticket.get(\"tags\"):\n metadata[\"tags\"] = tags\n if ticket_type := ticket.get(\"type\"):\n metadata[\"ticket_type\"] = ticket_type\n\n # Fetch comments for the ticket\n comments_data = client.make_request(f\"tickets/{ticket.get('id')}/comments\", {})\n comments = comments_data.get(\"comments\", [])\n\n comment_texts = []\n for comment in comments:\n new_author_mapping, comment_text = _get_comment_text(\n comment, author_map, client\n )\n if new_author_mapping:\n author_map.update(new_author_mapping)\n comment_texts.append(comment_text)\n\n comments_text = \"\\n\\n\".join(comment_texts)\n\n subject = ticket.get(\"subject\")\n full_text = f\"Ticket Subject:\\n{subject}\\n\\nComments:\\n{comments_text}\"\n\n ticket_url = ticket.get(\"url\")\n subdomain = (\n ticket_url.split(\"//\")[1].split(\".zendesk.com\")[0]\n if ticket_url\n else default_subdomain\n )\n\n ticket_display_url = (\n f\"https://{subdomain}.zendesk.com/agent/tickets/{ticket.get('id')}\"\n )\n\n return new_author_mapping, Document(\n id=f\"zendesk_ticket_{ticket['id']}\",\n sections=[TextSection(link=ticket_display_url, text=full_text)],\n source=DocumentSource.ZENDESK,\n semantic_identifier=f\"Ticket #{ticket['id']}: {subject or 'No Subject'}\",\n doc_updated_at=update_time,\n primary_owners=[submitter] if submitter else None,\n metadata=metadata,\n )\n\n\nclass ZendeskConnectorCheckpoint(ConnectorCheckpoint):\n # We use cursor-based paginated retrieval for articles\n after_cursor_articles: str | None\n\n # We use timestamp-based paginated retrieval for tickets\n next_start_time_tickets: int | None\n\n cached_author_map: dict[str, BasicExpertInfo] | None\n cached_content_tags: dict[str, str] | None\n\n\nclass ZendeskConnector(\n SlimConnectorWithPermSync, CheckpointedConnector[ZendeskConnectorCheckpoint]\n):\n def __init__(\n self,\n content_type: str = \"articles\",\n calls_per_minute: int | None = None,\n ) -> None:\n self.content_type = content_type\n self.subdomain = \"\"\n # Fetch all tags ahead of time\n self.content_tags: dict[str, str] = {}\n self.calls_per_minute = calls_per_minute\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n # Subdomain is actually the whole URL\n subdomain = (\n credentials[\"zendesk_subdomain\"]\n .replace(\"https://\", \"\")\n .split(\".zendesk.com\")[0]\n )\n self.subdomain = subdomain\n\n self.client = ZendeskClient(\n subdomain,\n credentials[\"zendesk_email\"],\n credentials[\"zendesk_token\"],\n calls_per_minute=self.calls_per_minute,\n )\n return None\n\n @override\n def load_from_checkpoint(\n self,\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n checkpoint: ZendeskConnectorCheckpoint,\n ) -> CheckpointOutput[ZendeskConnectorCheckpoint]:\n if self.client is None:\n raise ZendeskCredentialsNotSetUpError()\n\n if checkpoint.cached_content_tags is None:\n checkpoint.cached_content_tags = _get_content_tag_mapping(self.client)\n return checkpoint # save the content tags to the checkpoint\n self.content_tags = checkpoint.cached_content_tags\n\n if self.content_type == \"articles\":\n checkpoint = yield from self._retrieve_articles(start, end, checkpoint)\n return checkpoint\n elif self.content_type == \"tickets\":\n checkpoint = yield from self._retrieve_tickets(start, end, checkpoint)\n return checkpoint\n else:\n raise ValueError(f\"Unsupported content_type: {self.content_type}\")\n\n def _retrieve_articles(\n self,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n checkpoint: ZendeskConnectorCheckpoint,\n ) -> CheckpointOutput[ZendeskConnectorCheckpoint]:\n checkpoint = copy.deepcopy(checkpoint)\n # This one is built on the fly as there may be more many more authors than tags\n author_map: dict[str, BasicExpertInfo] = checkpoint.cached_author_map or {}\n after_cursor = checkpoint.after_cursor_articles\n doc_batch: list[Document] = []\n\n response = _get_article_page(\n self.client,\n start_time=int(start) if start else None,\n after_cursor=after_cursor,\n )\n articles = response.data\n has_more = response.has_more\n after_cursor = response.meta.get(\"after_cursor\")\n for article in articles:\n if (\n article.get(\"body\") is None\n or article.get(\"draft\")\n or any(\n label in ZENDESK_CONNECTOR_SKIP_ARTICLE_LABELS\n for label in article.get(\"label_names\", [])\n )\n ):\n continue\n\n try:\n new_author_map, document = _article_to_document(\n article, self.content_tags, author_map, self.client\n )\n except Exception as e:\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=f\"{article.get('id')}\",\n document_link=article.get(\"html_url\", \"\"),\n ),\n failure_message=str(e),\n exception=e,\n )\n continue\n\n if new_author_map:\n author_map.update(new_author_map)\n\n doc_batch.append(document)\n\n if not has_more:\n yield from doc_batch\n checkpoint.has_more = False\n return checkpoint\n\n # Sometimes no documents are retrieved, but the cursor\n # is still updated so the connector makes progress.\n yield from doc_batch\n checkpoint.after_cursor_articles = after_cursor\n\n last_doc_updated_at = doc_batch[-1].doc_updated_at if doc_batch else None\n checkpoint.has_more = bool(\n end is None\n or last_doc_updated_at is None\n or last_doc_updated_at.timestamp() <= end\n )\n checkpoint.cached_author_map = (\n author_map if len(author_map) <= MAX_AUTHOR_MAP_SIZE else None\n )\n return checkpoint\n\n def _retrieve_tickets(\n self,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None,\n checkpoint: ZendeskConnectorCheckpoint,\n ) -> CheckpointOutput[ZendeskConnectorCheckpoint]:\n checkpoint = copy.deepcopy(checkpoint)\n if self.client is None:\n raise ZendeskCredentialsNotSetUpError()\n\n author_map: dict[str, BasicExpertInfo] = checkpoint.cached_author_map or {}\n\n doc_batch: list[Document] = []\n next_start_time = int(checkpoint.next_start_time_tickets or start or 0)\n ticket_response = _get_tickets_page(self.client, start_time=next_start_time)\n tickets = ticket_response.data\n has_more = ticket_response.has_more\n next_start_time = ticket_response.meta[\"end_time\"]\n for ticket in tickets:\n if ticket.get(\"status\") == \"deleted\":\n continue\n\n try:\n new_author_map, document = _ticket_to_document(\n ticket=ticket,\n author_map=author_map,\n client=self.client,\n default_subdomain=self.subdomain,\n )\n except Exception as e:\n yield ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=f\"{ticket.get('id')}\",\n document_link=ticket.get(\"url\", \"\"),\n ),\n failure_message=str(e),\n exception=e,\n )\n continue\n\n if new_author_map:\n author_map.update(new_author_map)\n\n doc_batch.append(document)\n\n if not has_more:\n yield from doc_batch\n checkpoint.has_more = False\n return checkpoint\n\n yield from doc_batch\n checkpoint.next_start_time_tickets = next_start_time\n last_doc_updated_at = doc_batch[-1].doc_updated_at if doc_batch else None\n checkpoint.has_more = bool(\n end is None\n or last_doc_updated_at is None\n or last_doc_updated_at.timestamp() <= end\n )\n checkpoint.cached_author_map = (\n author_map if len(author_map) <= MAX_AUTHOR_MAP_SIZE else None\n )\n return checkpoint\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None,\n end: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n slim_doc_batch: list[SlimDocument | HierarchyNode] = []\n if self.content_type == \"articles\":\n articles = _get_articles(\n self.client, start_time=int(start) if start else None\n )\n for article in articles:\n slim_doc_batch.append(\n SlimDocument(\n id=f\"article:{article['id']}\",\n )\n )\n if len(slim_doc_batch) >= _SLIM_BATCH_SIZE:\n yield slim_doc_batch\n slim_doc_batch = []\n elif self.content_type == \"tickets\":\n tickets = _get_tickets(\n self.client, start_time=int(start) if start else None\n )\n for ticket in tickets:\n slim_doc_batch.append(\n SlimDocument(\n id=f\"zendesk_ticket_{ticket['id']}\",\n )\n )\n if len(slim_doc_batch) >= _SLIM_BATCH_SIZE:\n yield slim_doc_batch\n slim_doc_batch = []\n else:\n raise ValueError(f\"Unsupported content_type: {self.content_type}\")\n if slim_doc_batch:\n yield slim_doc_batch\n\n @override\n def validate_connector_settings(self) -> None:\n if self.client is None:\n raise ZendeskCredentialsNotSetUpError()\n\n try:\n _get_article_page(self.client, start_time=0)\n except HTTPError as e:\n # Check for HTTP status codes\n if e.response.status_code == 401:\n raise CredentialExpiredError(\n \"Your Zendesk credentials appear to be invalid or expired (HTTP 401).\"\n ) from e\n elif e.response.status_code == 403:\n raise InsufficientPermissionsError(\n \"Your Zendesk token does not have sufficient permissions (HTTP 403).\"\n ) from e\n elif e.response.status_code == 404:\n raise ConnectorValidationError(\n \"Zendesk resource not found (HTTP 404).\"\n ) from e\n else:\n raise ConnectorValidationError(\n f\"Unexpected Zendesk error (status={e.response.status_code}): {e}\"\n ) from e\n\n @override\n def validate_checkpoint_json(\n self, checkpoint_json: str\n ) -> ZendeskConnectorCheckpoint:\n return ZendeskConnectorCheckpoint.model_validate_json(checkpoint_json)\n\n @override\n def build_dummy_checkpoint(self) -> ZendeskConnectorCheckpoint:\n return ZendeskConnectorCheckpoint(\n after_cursor_articles=None,\n next_start_time_tickets=None,\n cached_author_map=None,\n cached_content_tags=None,\n has_more=True,\n )\n\n\nif __name__ == \"__main__\":\n import os\n\n connector = ZendeskConnector()\n connector.load_credentials(\n {\n \"zendesk_subdomain\": os.environ[\"ZENDESK_SUBDOMAIN\"],\n \"zendesk_email\": os.environ[\"ZENDESK_EMAIL\"],\n \"zendesk_token\": os.environ[\"ZENDESK_TOKEN\"],\n }\n )\n\n current = time.time()\n one_day_ago = current - 24 * 60 * 60 # 1 day\n document_batches = connector.load_from_checkpoint(\n one_day_ago,\n current,\n connector.build_dummy_checkpoint(),\n )\n\n print(next(document_batches))\n" }, { "path": "backend/onyx/connectors/zulip/__init__.py", "content": "" }, { "path": "backend/onyx/connectors/zulip/connector.py", "content": "import os\nimport tempfile\nimport urllib.parse\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import Dict\nfrom typing import List\nfrom typing import Tuple\nfrom typing import Union\n\nfrom zulip import Client\n\nfrom onyx.configs.app_configs import INDEX_BATCH_SIZE\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateDocumentsOutput\nfrom onyx.connectors.interfaces import LoadConnector\nfrom onyx.connectors.interfaces import PollConnector\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.connectors.zulip.schemas import GetMessagesResponse\nfrom onyx.connectors.zulip.schemas import Message\nfrom onyx.connectors.zulip.utils import build_search_narrow\nfrom onyx.connectors.zulip.utils import call_api\nfrom onyx.connectors.zulip.utils import encode_zulip_narrow_operand\nfrom onyx.utils.logger import setup_logger\n\n# Potential improvements\n# 1. Group documents messages into topics, make 1 document per topic per week\n# 2. Add end date support once https://github.com/zulip/zulip/issues/25436 is solved\n\nlogger = setup_logger()\n\n\nclass ZulipConnector(LoadConnector, PollConnector):\n def __init__(\n self, realm_name: str, realm_url: str, batch_size: int = INDEX_BATCH_SIZE\n ) -> None:\n self.batch_size = batch_size\n self.realm_name = realm_name\n\n # Clean and normalize the URL\n realm_url = realm_url.strip().lower()\n\n # Remove any trailing slashes\n realm_url = realm_url.rstrip(\"/\")\n\n # Ensure the URL has a scheme\n if not realm_url.startswith((\"http://\", \"https://\")):\n realm_url = f\"https://{realm_url}\"\n\n try:\n parsed = urllib.parse.urlparse(realm_url)\n\n # Extract the base domain without any paths or ports\n netloc = parsed.netloc.split(\":\")[0] # Remove port if present\n\n if not netloc:\n raise ValueError(\n f\"Invalid realm URL format: {realm_url}. URL must include a valid domain name.\"\n )\n\n # Always use HTTPS for security\n self.base_url = f\"https://{netloc}\"\n self.client: Client | None = None\n\n except Exception as e:\n raise ValueError(\n f\"Failed to parse Zulip realm URL: {realm_url}. \"\n f\"Please provide a URL in the format: domain.com or https://domain.com. \"\n f\"Error: {str(e)}\"\n )\n\n def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:\n contents = credentials[\"zuliprc_content\"]\n # The input field converts newlines to spaces in the provided\n # zuliprc file. This reverts them back to newlines.\n contents_spaces_to_newlines = contents.replace(\" \", \"\\n\")\n # create a temporary zuliprc file\n tempdir = tempfile.tempdir\n if tempdir is None:\n raise Exception(\"Could not determine tempfile directory\")\n config_file = os.path.join(tempdir, f\"zuliprc-{self.realm_name}\")\n with open(config_file, \"w\") as f:\n f.write(contents_spaces_to_newlines)\n self.client = Client(config_file=config_file)\n return None\n\n def _message_to_narrow_link(self, m: Message) -> str:\n try:\n stream_name = m.display_recipient # assume str\n stream_operand = encode_zulip_narrow_operand(f\"{m.stream_id}-{stream_name}\")\n topic_operand = encode_zulip_narrow_operand(m.subject)\n\n narrow_link = f\"{self.base_url}#narrow/stream/{stream_operand}/topic/{topic_operand}/near/{m.id}\"\n return narrow_link\n except Exception as e:\n logger.error(f\"Error generating Zulip message link: {e}\")\n # Fallback to a basic link that at least includes the base URL\n return f\"{self.base_url}#narrow/id/{m.id}\"\n\n def _get_message_batch(self, anchor: str) -> Tuple[bool, List[Message]]:\n if self.client is None:\n raise ConnectorMissingCredentialError(\"Zulip\")\n\n logger.info(f\"Fetching messages starting with anchor={anchor}\")\n request = build_search_narrow(\n limit=INDEX_BATCH_SIZE, anchor=anchor, apply_md=False\n )\n response = GetMessagesResponse(**call_api(self.client.get_messages, request))\n\n end = False\n if len(response.messages) == 0 or response.found_oldest:\n end = True\n\n # reverse, so that the last message is the new anchor\n # and the order is from newest to oldest\n return end, response.messages[::-1]\n\n def _message_to_doc(self, message: Message) -> Document:\n text = f\"{message.sender_full_name}: {message.content}\"\n\n try:\n # Convert timestamps to UTC datetime objects\n post_time = datetime.fromtimestamp(message.timestamp, tz=timezone.utc)\n edit_time = (\n datetime.fromtimestamp(message.last_edit_timestamp, tz=timezone.utc)\n if message.last_edit_timestamp is not None\n else None\n )\n\n # Use the most recent edit time if available, otherwise use post time\n doc_time = edit_time if edit_time is not None else post_time\n\n except (ValueError, TypeError) as e:\n logger.warning(f\"Failed to parse timestamp for message {message.id}: {e}\")\n post_time = None\n edit_time = None\n doc_time = None\n\n metadata: Dict[str, Union[str, List[str]]] = {\n \"stream_name\": str(message.display_recipient),\n \"topic\": str(message.subject),\n \"sender_name\": str(message.sender_full_name),\n \"sender_email\": str(message.sender_email),\n \"message_timestamp\": str(message.timestamp),\n \"message_id\": str(message.id),\n \"stream_id\": str(message.stream_id),\n \"has_reactions\": str(len(message.reactions) > 0),\n \"content_type\": str(message.content_type or \"text\"),\n }\n\n # Always include edit timestamp in metadata when available\n if edit_time is not None:\n metadata[\"edit_timestamp\"] = str(message.last_edit_timestamp)\n\n return Document(\n id=f\"{message.stream_id}__{message.id}\",\n sections=[\n TextSection(\n link=self._message_to_narrow_link(message),\n text=text,\n )\n ],\n source=DocumentSource.ZULIP,\n semantic_identifier=f\"{message.display_recipient} > {message.subject}\",\n metadata=metadata,\n doc_updated_at=doc_time, # Use most recent edit time or post time\n )\n\n def _get_docs(\n self, anchor: str, start: SecondsSinceUnixEpoch | None = None\n ) -> Generator[Document, None, None]:\n message: Message | None = None\n while True:\n end, message_batch = self._get_message_batch(anchor)\n\n for message in message_batch:\n if start is not None and float(message.timestamp) < start:\n return\n yield self._message_to_doc(message)\n\n if end or message is None:\n return\n\n # Last message is oldest, use as next anchor\n anchor = str(message.id)\n\n def _poll_source(\n self,\n start: SecondsSinceUnixEpoch | None,\n end: SecondsSinceUnixEpoch | None, # noqa: ARG002\n ) -> GenerateDocumentsOutput:\n # Since Zulip doesn't support searching by timestamp,\n # we have to always start from the newest message\n # and go backwards.\n anchor = \"newest\"\n\n docs: list[Document | HierarchyNode] = []\n for doc in self._get_docs(anchor=anchor, start=start):\n docs.append(doc)\n if len(docs) == self.batch_size:\n yield docs\n docs = []\n if docs:\n yield docs\n\n def load_from_state(self) -> GenerateDocumentsOutput:\n return self._poll_source(start=None, end=None)\n\n def poll_source(\n self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch\n ) -> GenerateDocumentsOutput:\n return self._poll_source(start, end)\n" }, { "path": "backend/onyx/connectors/zulip/schemas.py", "content": "from typing import Any\nfrom typing import List\nfrom typing import Optional\nfrom typing import Union\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\n\nclass Message(BaseModel):\n id: int\n sender_id: int\n content: str\n recipient_id: int\n timestamp: int\n client: str\n is_me_message: bool\n sender_full_name: str\n sender_email: str\n sender_realm_str: str\n subject: str\n topic_links: Optional[List[Any]] = None\n last_edit_timestamp: Optional[int] = None\n edit_history: Any = None\n reactions: List[Any]\n submessages: List[Any]\n flags: List[str] = Field(default_factory=list)\n display_recipient: Optional[str] = None\n type: Optional[str] = None\n stream_id: int\n avatar_url: Optional[str]\n content_type: Optional[str]\n rendered_content: Optional[str] = None\n\n\nclass GetMessagesResponse(BaseModel):\n result: str\n msg: str\n found_anchor: Optional[bool] = None\n found_oldest: Optional[bool] = None\n found_newest: Optional[bool] = None\n history_limited: Optional[bool] = None\n anchor: Optional[Union[str, int]] = None\n messages: List[Message] = Field(default_factory=list)\n" }, { "path": "backend/onyx/connectors/zulip/utils.py", "content": "import time\nfrom collections.abc import Callable\nfrom typing import Any\nfrom typing import Dict\nfrom typing import Optional\nfrom urllib.parse import quote\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass ZulipAPIError(Exception):\n def __init__(self, code: Any = None, msg: str | None = None) -> None:\n self.code = code\n self.msg = msg\n\n def __str__(self) -> str:\n return (\n f\"Error occurred during Zulip API call: {self.msg}\" + \"\"\n if self.code is None\n else f\" ({self.code})\"\n )\n\n\nclass ZulipHTTPError(ZulipAPIError):\n def __init__(self, msg: str | None = None, status_code: Any = None) -> None:\n super().__init__(code=None, msg=msg)\n self.status_code = status_code\n\n def __str__(self) -> str:\n return f\"HTTP error {self.status_code} occurred during Zulip API call\"\n\n\ndef __call_with_retry(fun: Callable, *args: Any, **kwargs: Any) -> Dict[str, Any]:\n result = fun(*args, **kwargs)\n if result.get(\"result\") == \"error\":\n if result.get(\"code\") == \"RATE_LIMIT_HIT\":\n retry_after = float(result[\"retry-after\"]) + 1\n logger.warn(f\"Rate limit hit, retrying after {retry_after} seconds\")\n time.sleep(retry_after)\n return __call_with_retry(fun, *args)\n return result\n\n\ndef __raise_if_error(response: dict[str, Any]) -> None:\n if response.get(\"result\") == \"error\":\n raise ZulipAPIError(\n code=response.get(\"code\"),\n msg=response.get(\"msg\"),\n )\n elif response.get(\"result\") == \"http-error\":\n raise ZulipHTTPError(\n msg=response.get(\"msg\"), status_code=response.get(\"status_code\")\n )\n\n\ndef call_api(fun: Callable, *args: Any, **kwargs: Any) -> Dict[str, Any]:\n response = __call_with_retry(fun, *args, **kwargs)\n __raise_if_error(response)\n return response\n\n\ndef build_search_narrow(\n *,\n stream: Optional[str] = None,\n topic: Optional[str] = None,\n limit: int = 100,\n content: Optional[str] = None,\n apply_md: bool = False,\n anchor: str = \"newest\",\n) -> Dict[str, Any]:\n narrow_filters = []\n\n if stream:\n narrow_filters.append({\"operator\": \"stream\", \"operand\": stream})\n\n if topic:\n narrow_filters.append({\"operator\": \"topic\", \"operand\": topic})\n\n if content:\n narrow_filters.append({\"operator\": \"has\", \"operand\": content})\n\n if not stream and not topic and not content:\n narrow_filters.append({\"operator\": \"streams\", \"operand\": \"public\"})\n\n narrow = {\n \"anchor\": anchor,\n \"num_before\": limit,\n \"num_after\": 0,\n \"narrow\": narrow_filters,\n }\n narrow[\"apply_markdown\"] = apply_md\n\n return narrow\n\n\ndef encode_zulip_narrow_operand(value: str) -> str:\n # like https://github.com/zulip/zulip/blob/1577662a6/static/js/hash_util.js#L18-L25\n # safe characters necessary to make Python match Javascript's escaping behaviour,\n # see: https://stackoverflow.com/a/74439601\n return quote(value, safe=\"!~*'()\").replace(\".\", \"%2E\").replace(\"%\", \".\")\n" }, { "path": "backend/onyx/context/search/__init__.py", "content": "" }, { "path": "backend/onyx/context/search/enums.py", "content": "\"\"\"NOTE: this needs to be separate from models.py because of circular imports.\nBoth search/models.py and db/models.py import enums from this file AND\nsearch/models.py imports from db/models.py.\"\"\"\n\nfrom enum import Enum\n\n\nclass RecencyBiasSetting(str, Enum):\n FAVOR_RECENT = \"favor_recent\" # 2x decay rate\n BASE_DECAY = \"base_decay\"\n NO_DECAY = \"no_decay\"\n # Determine based on query if to use base_decay or favor_recent\n AUTO = \"auto\"\n\n\nclass QueryType(str, Enum):\n \"\"\"\n The type of first-pass query to use for hybrid search.\n\n The values of this enum are injected into the ranking profile name which\n should match the name in the schema.\n \"\"\"\n\n KEYWORD = \"keyword\"\n SEMANTIC = \"semantic\"\n\n\nclass SearchType(str, Enum):\n KEYWORD = \"keyword\"\n SEMANTIC = \"semantic\"\n INTERNET = \"internet\"\n" }, { "path": "backend/onyx/context/search/federated/models.py", "content": "from datetime import datetime\nfrom typing import TypedDict\n\nfrom pydantic import BaseModel\n\nfrom onyx.onyxbot.slack.models import ChannelType\n\n\nclass ChannelMetadata(TypedDict):\n \"\"\"Type definition for cached channel metadata.\"\"\"\n\n name: str\n type: ChannelType\n is_private: bool\n is_member: bool\n\n\nclass SlackMessage(BaseModel):\n document_id: str\n channel_id: str\n message_id: str\n thread_id: str | None\n link: str\n metadata: dict[str, str | list[str]]\n timestamp: datetime\n recency_bias: float\n semantic_identifier: str\n text: str\n highlighted_texts: set[str]\n slack_score: float\n" }, { "path": "backend/onyx/context/search/federated/slack_search.py", "content": "import json\nimport re\nimport time\nfrom datetime import datetime\nfrom typing import Any\nfrom typing import cast\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import ValidationError\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import ENABLE_CONTEXTUAL_RAG\nfrom onyx.configs.app_configs import MAX_SLACK_THREAD_CONTEXT_MESSAGES\nfrom onyx.configs.app_configs import SLACK_THREAD_CONTEXT_BATCH_SIZE\nfrom onyx.configs.chat_configs import DOC_TIME_DECAY\nfrom onyx.connectors.models import IndexingDocument\nfrom onyx.connectors.models import TextSection\nfrom onyx.context.search.federated.models import ChannelMetadata\nfrom onyx.context.search.federated.models import SlackMessage\nfrom onyx.context.search.federated.slack_search_utils import ALL_CHANNEL_TYPES\nfrom onyx.context.search.federated.slack_search_utils import build_channel_query_filter\nfrom onyx.context.search.federated.slack_search_utils import build_slack_queries\nfrom onyx.context.search.federated.slack_search_utils import get_channel_type\nfrom onyx.context.search.federated.slack_search_utils import (\n get_channel_type_for_missing_scope,\n)\nfrom onyx.context.search.federated.slack_search_utils import is_recency_query\nfrom onyx.context.search.federated.slack_search_utils import should_include_message\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.document import DocumentSource\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.document_index.document_index_utils import (\n get_multipass_config,\n)\nfrom onyx.federated_connectors.slack.models import SlackEntities\nfrom onyx.indexing.chunker import Chunker\nfrom onyx.indexing.embedder import DefaultIndexingEmbedder\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.onyxbot.slack.models import ChannelType\nfrom onyx.onyxbot.slack.models import SlackContext\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.federated.models import FederatedConnectorDetail\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE\n\nlogger = setup_logger()\n\nHIGHLIGHT_START_CHAR = \"\\ue000\"\nHIGHLIGHT_END_CHAR = \"\\ue001\"\n\nCHANNEL_METADATA_CACHE_TTL = 60 * 60 * 24 # 24 hours\nUSER_PROFILE_CACHE_TTL = 60 * 60 * 24 # 24 hours\nSLACK_THREAD_CONTEXT_WINDOW = 3 # Number of messages before matched message to include\nCHANNEL_METADATA_MAX_RETRIES = 3 # Maximum retry attempts for channel metadata fetching\nCHANNEL_METADATA_RETRY_DELAY = 1 # Initial retry delay in seconds (exponential backoff)\n\n\ndef fetch_and_cache_channel_metadata(\n access_token: str, team_id: str, include_private: bool = True\n) -> dict[str, ChannelMetadata]:\n \"\"\"\n Fetch ALL channel metadata in one API call and cache it.\n\n Returns a dict mapping channel_id -> metadata including name, type, etc.\n This replaces multiple conversations.info calls with a single conversations.list.\n\n Note: We ALWAYS fetch all channel types (including private) and cache them together.\n This ensures a single cache entry per team, avoiding duplicate API calls.\n \"\"\"\n # Use tenant-specific Redis client\n redis_client = get_redis_client()\n # (tenant_id prefix is added automatically by TenantRedis)\n cache_key = f\"slack_federated_search:{team_id}:channels:metadata\"\n\n try:\n cached = redis_client.get(cache_key)\n if cached:\n logger.debug(f\"Channel metadata cache HIT for team {team_id}\")\n cached_str: str = (\n cached.decode(\"utf-8\") if isinstance(cached, bytes) else str(cached)\n )\n cached_data = cast(dict[str, ChannelMetadata], json.loads(cached_str))\n logger.debug(f\"Loaded {len(cached_data)} channels from cache\")\n if not include_private:\n filtered: dict[str, ChannelMetadata] = {\n k: v\n for k, v in cached_data.items()\n if v.get(\"type\") != ChannelType.PRIVATE_CHANNEL.value\n }\n logger.debug(f\"Filtered to {len(filtered)} channels (exclude private)\")\n return filtered\n return cached_data\n except Exception as e:\n logger.warning(f\"Error reading from channel metadata cache: {e}\")\n\n # Cache miss - fetch from Slack API with retry logic\n logger.debug(f\"Channel metadata cache MISS for team {team_id} - fetching from API\")\n slack_client = WebClient(token=access_token)\n channel_metadata: dict[str, ChannelMetadata] = {}\n\n # Retry logic with exponential backoff\n last_exception = None\n available_channel_types = ALL_CHANNEL_TYPES.copy()\n\n for attempt in range(CHANNEL_METADATA_MAX_RETRIES):\n try:\n # Use available channel types (may be reduced if scopes are missing)\n channel_types = \",\".join(available_channel_types)\n\n # Fetch all channels in one call\n cursor = None\n channel_count = 0\n while True:\n response = slack_client.conversations_list(\n types=channel_types,\n exclude_archived=True,\n limit=1000,\n cursor=cursor,\n )\n response.validate()\n\n # Cast response.data to dict for type checking\n response_data: dict[str, Any] = response.data # type: ignore\n for ch in response_data.get(\"channels\", []):\n channel_id = ch.get(\"id\")\n if not channel_id:\n continue\n\n # Determine channel type\n channel_type_enum = get_channel_type(channel_info=ch)\n channel_type = ChannelType(channel_type_enum.value)\n\n channel_metadata[channel_id] = {\n \"name\": ch.get(\"name\", \"\"),\n \"type\": channel_type,\n \"is_private\": ch.get(\"is_private\", False),\n \"is_member\": ch.get(\"is_member\", False),\n }\n channel_count += 1\n\n cursor = response_data.get(\"response_metadata\", {}).get(\"next_cursor\")\n if not cursor:\n break\n\n logger.info(f\"Fetched {channel_count} channels for team {team_id}\")\n\n # Cache the results\n try:\n redis_client.set(\n cache_key,\n json.dumps(channel_metadata),\n ex=CHANNEL_METADATA_CACHE_TTL,\n )\n logger.info(\n f\"Cached {channel_count} channels for team {team_id} (TTL: {CHANNEL_METADATA_CACHE_TTL}s, key: {cache_key})\"\n )\n except Exception as e:\n logger.warning(f\"Error caching channel metadata: {e}\")\n\n return channel_metadata\n\n except SlackApiError as e:\n last_exception = e\n\n # Extract all needed fields from response upfront\n if e.response:\n error_response = e.response.get(\"error\", \"\")\n needed_scope = e.response.get(\"needed\", \"\")\n else:\n error_response = \"\"\n needed_scope = \"\"\n\n # Check if this is a missing_scope error\n if error_response == \"missing_scope\":\n # Get the channel type that requires this scope\n missing_channel_type = get_channel_type_for_missing_scope(needed_scope)\n\n if (\n missing_channel_type\n and missing_channel_type in available_channel_types\n ):\n # Remove the problematic channel type and retry\n available_channel_types.remove(missing_channel_type)\n logger.warning(\n f\"Missing scope '{needed_scope}' for channel type '{missing_channel_type}'. \"\n f\"Continuing with reduced channel types: {available_channel_types}\"\n )\n # Don't count this as a retry attempt, just try again with fewer types\n if available_channel_types: # Only continue if we have types left\n continue\n # Otherwise fall through to retry logic\n else:\n logger.error(\n f\"Missing scope '{needed_scope}' but could not map to channel type or already removed. \"\n f\"Response: {e.response}\"\n )\n\n # For other errors, use retry logic\n if attempt < CHANNEL_METADATA_MAX_RETRIES - 1:\n retry_delay = CHANNEL_METADATA_RETRY_DELAY * (2**attempt)\n logger.warning(\n f\"Failed to fetch channel metadata (attempt {attempt + 1}/{CHANNEL_METADATA_MAX_RETRIES}): {e}. \"\n f\"Retrying in {retry_delay}s...\"\n )\n time.sleep(retry_delay)\n else:\n logger.error(\n f\"Failed to fetch channel metadata after {CHANNEL_METADATA_MAX_RETRIES} attempts: {e}\"\n )\n\n # If we have some channel metadata despite errors, return it with a warning\n if channel_metadata:\n logger.warning(\n f\"Returning partial channel metadata ({len(channel_metadata)} channels) despite errors. Last error: {last_exception}\"\n )\n return channel_metadata\n\n # If we exhausted all retries and have no data, raise the last exception\n if last_exception:\n raise SlackApiError(\n f\"Channel metadata fetching failed after {CHANNEL_METADATA_MAX_RETRIES} attempts\",\n last_exception.response,\n )\n\n return {}\n\n\ndef get_available_channels(\n access_token: str, team_id: str, include_private: bool = False\n) -> list[str]:\n \"\"\"Fetch list of available channel names using cached metadata.\"\"\"\n metadata = fetch_and_cache_channel_metadata(access_token, team_id, include_private)\n return [meta[\"name\"] for meta in metadata.values() if meta[\"name\"]]\n\n\ndef get_cached_user_profile(\n access_token: str, team_id: str, user_id: str\n) -> str | None:\n \"\"\"\n Get a user's display name from cache or fetch from Slack API.\n\n Uses Redis caching to avoid repeated API calls and rate limiting.\n Returns the user's real_name or email, or None if not found.\n \"\"\"\n redis_client = get_redis_client()\n cache_key = f\"slack_federated_search:{team_id}:user:{user_id}\"\n\n # Check cache first\n try:\n cached = redis_client.get(cache_key)\n if cached is not None:\n cached_str = (\n cached.decode(\"utf-8\") if isinstance(cached, bytes) else str(cached)\n )\n # Empty string means user was not found previously\n return cached_str if cached_str else None\n except Exception as e:\n logger.debug(f\"Error reading user profile cache: {e}\")\n\n # Cache miss - fetch from Slack API\n slack_client = WebClient(token=access_token)\n try:\n response = slack_client.users_profile_get(user=user_id)\n response.validate()\n profile: dict[str, Any] = response.get(\"profile\", {})\n name: str | None = profile.get(\"real_name\") or profile.get(\"email\")\n\n # Cache the result (empty string for not found)\n try:\n redis_client.set(\n cache_key,\n name or \"\",\n ex=USER_PROFILE_CACHE_TTL,\n )\n except Exception as e:\n logger.debug(f\"Error caching user profile: {e}\")\n\n return name\n\n except SlackApiError as e:\n error_str = str(e)\n if \"user_not_found\" in error_str:\n logger.debug(\n f\"User {user_id} not found in Slack workspace (likely deleted/deactivated)\"\n )\n elif \"ratelimited\" in error_str:\n # Don't cache rate limit errors - we'll retry later\n logger.debug(f\"Rate limited fetching user {user_id}, will retry later\")\n return None\n else:\n logger.warning(f\"Could not fetch profile for user {user_id}: {e}\")\n\n # Cache negative result to avoid repeated lookups for missing users\n try:\n redis_client.set(cache_key, \"\", ex=USER_PROFILE_CACHE_TTL)\n except Exception:\n pass\n\n return None\n\n\ndef batch_get_user_profiles(\n access_token: str, team_id: str, user_ids: set[str]\n) -> dict[str, str]:\n \"\"\"\n Batch fetch user profiles with caching.\n\n Returns a dict mapping user_id -> display_name for users that were found.\n \"\"\"\n result: dict[str, str] = {}\n\n for user_id in user_ids:\n name = get_cached_user_profile(access_token, team_id, user_id)\n if name:\n result[user_id] = name\n\n return result\n\n\ndef _extract_channel_data_from_entities(\n entities: dict[str, Any] | None,\n channel_metadata_dict: dict[str, ChannelMetadata] | None,\n) -> list[str] | None:\n \"\"\"Extract available channels list from metadata based on entity configuration.\n\n Args:\n entities: Entity filter configuration dict\n channel_metadata_dict: Pre-fetched channel metadata dictionary\n\n Returns:\n List of available channel names, or None if not needed\n \"\"\"\n if not entities or not channel_metadata_dict:\n return None\n\n try:\n parsed_entities = SlackEntities(**entities)\n # Only extract if we have exclusions or channel filters\n if parsed_entities.exclude_channels or parsed_entities.channels:\n # Extract channel names from metadata dict\n return [\n meta[\"name\"]\n for meta in channel_metadata_dict.values()\n if meta[\"name\"]\n and (\n parsed_entities.include_private_channels\n or meta.get(\"type\") != ChannelType.PRIVATE_CHANNEL.value\n )\n ]\n except ValidationError:\n logger.debug(\"Failed to parse entities for channel data extraction\")\n\n return None\n\n\ndef _should_skip_channel(\n channel_id: str,\n allowed_private_channel: str | None,\n bot_token: str | None,\n access_token: str,\n include_dm: bool,\n channel_metadata_dict: dict[str, ChannelMetadata] | None = None,\n) -> bool:\n \"\"\"Bot context filtering: skip private channels unless explicitly allowed.\n\n Uses pre-fetched channel metadata when available to avoid API calls.\n \"\"\"\n if bot_token and not include_dm:\n try:\n # First try to use pre-fetched metadata from cache\n if channel_metadata_dict and channel_id in channel_metadata_dict:\n channel_meta = channel_metadata_dict[channel_id]\n channel_type_str = channel_meta.get(\"type\", \"\")\n is_private_or_dm = channel_type_str in [\n ChannelType.PRIVATE_CHANNEL.value,\n ChannelType.IM.value,\n ChannelType.MPIM.value,\n ]\n if is_private_or_dm and channel_id != allowed_private_channel:\n return True\n return False\n\n # Fallback: API call only if not in cache (should be rare)\n token_to_use = bot_token or access_token\n channel_client = WebClient(token=token_to_use)\n channel_info = channel_client.conversations_info(channel=channel_id)\n\n if isinstance(channel_info.data, dict):\n channel_data = channel_info.data.get(\"channel\", {})\n channel_type = get_channel_type(channel_info=channel_data)\n is_private_or_dm = channel_type in [\n ChannelType.PRIVATE_CHANNEL,\n ChannelType.IM,\n ChannelType.MPIM,\n ]\n\n if is_private_or_dm and channel_id != allowed_private_channel:\n return True\n except Exception as e:\n logger.warning(\n f\"Could not determine channel type for {channel_id}, filtering out: {e}\"\n )\n return True\n return False\n\n\nclass SlackQueryResult(BaseModel):\n \"\"\"Result from a single Slack query including stats.\"\"\"\n\n model_config = ConfigDict(arbitrary_types_allowed=True)\n\n messages: list[SlackMessage]\n filtered_channels: list[str] # Channels filtered out during this query\n\n\ndef query_slack(\n query_string: str,\n access_token: str,\n limit: int | None = None,\n allowed_private_channel: str | None = None,\n bot_token: str | None = None,\n include_dm: bool = False,\n entities: dict[str, Any] | None = None,\n available_channels: list[str] | None = None,\n channel_metadata_dict: dict[str, ChannelMetadata] | None = None,\n) -> SlackQueryResult:\n\n # Check if query has channel override (user specified channels in query)\n has_channel_override = query_string.startswith(\"__CHANNEL_OVERRIDE__\")\n\n if has_channel_override:\n # Remove the marker and use the query as-is (already has channel filters)\n final_query = query_string.replace(\"__CHANNEL_OVERRIDE__\", \"\").strip()\n else:\n # Normal flow: build channel filters from entity config\n channel_filter = \"\"\n if entities:\n channel_filter = build_channel_query_filter(entities, available_channels)\n\n final_query = query_string\n if channel_filter:\n # Add channel filter to query\n final_query = f\"{query_string} {channel_filter}\"\n\n logger.info(f\"Final query to slack: {final_query}\")\n\n # Detect if query asks for most recent results\n sort_by_time = is_recency_query(query_string)\n\n slack_client = WebClient(token=access_token)\n try:\n search_params: dict[str, Any] = {\n \"query\": final_query,\n \"count\": limit,\n \"highlight\": True,\n }\n\n # Sort by timestamp for recency-focused queries, otherwise by relevance\n if sort_by_time:\n search_params[\"sort\"] = \"timestamp\"\n search_params[\"sort_dir\"] = \"desc\"\n\n response = slack_client.search_messages(**search_params)\n response.validate()\n\n messages: dict[str, Any] = response.get(\"messages\", {})\n matches: list[dict[str, Any]] = messages.get(\"matches\", [])\n\n logger.info(f\"Slack search found {len(matches)} messages\")\n except SlackApiError as slack_error:\n logger.error(f\"Slack API error in search_messages: {slack_error}\")\n logger.error(\n f\"Slack API error details: status={slack_error.response.status_code}, error={slack_error.response.get('error')}\"\n )\n if \"not_allowed_token_type\" in str(slack_error):\n # Log token type prefix\n token_prefix = access_token[:4] if len(access_token) >= 4 else \"unknown\"\n logger.error(f\"TOKEN TYPE ERROR: access_token type: {token_prefix}...\")\n return SlackQueryResult(messages=[], filtered_channels=[])\n\n # convert matches to slack messages\n slack_messages: list[SlackMessage] = []\n filtered_channels: list[str] = []\n for match in matches:\n text: str | None = match.get(\"text\")\n permalink: str | None = match.get(\"permalink\")\n message_id: str | None = match.get(\"ts\")\n channel_id: str | None = match.get(\"channel\", {}).get(\"id\")\n channel_name: str | None = match.get(\"channel\", {}).get(\"name\")\n username: str | None = match.get(\"username\")\n if not username:\n # Fallback: try to get from user field if username is missing\n user_info = match.get(\"user\", \"\")\n if isinstance(user_info, str) and user_info:\n username = user_info # Use user ID as fallback\n else:\n username = \"unknown_user\"\n score: float = match.get(\"score\", 0.0)\n if ( # can't use any() because of type checking :(\n not text\n or not permalink\n or not message_id\n or not channel_id\n or not channel_name\n or not username\n ):\n continue\n\n # Apply channel filtering if needed\n if _should_skip_channel(\n channel_id,\n allowed_private_channel,\n bot_token,\n access_token,\n include_dm,\n channel_metadata_dict,\n ):\n filtered_channels.append(f\"{channel_name}({channel_id})\")\n continue\n\n # generate thread id and document id\n thread_id = (\n permalink.split(\"?thread_ts=\", 1)[1] if \"?thread_ts=\" in permalink else None\n )\n document_id = f\"{channel_id}_{message_id}\"\n\n decay_factor = DOC_TIME_DECAY\n doc_time = datetime.fromtimestamp(float(message_id))\n doc_age_years = (datetime.now() - doc_time).total_seconds() / (\n 365 * 24 * 60 * 60\n )\n recency_bias = max(1 / (1 + decay_factor * doc_age_years), 0.75)\n metadata: dict[str, str | list[str]] = {\n \"channel\": channel_name,\n \"time\": doc_time.isoformat(),\n }\n\n # extract out the highlighted texts\n highlighted_texts = set(\n re.findall(\n rf\"{re.escape(HIGHLIGHT_START_CHAR)}(.*?){re.escape(HIGHLIGHT_END_CHAR)}\",\n text,\n )\n )\n cleaned_text = text.replace(HIGHLIGHT_START_CHAR, \"\").replace(\n HIGHLIGHT_END_CHAR, \"\"\n )\n\n # get the semantic identifier\n snippet = (\n cleaned_text[:50].rstrip() + \"...\" if len(cleaned_text) > 50 else text\n ).replace(\"\\n\", \" \")\n doc_sem_id = f\"{username} in #{channel_name}: {snippet}\"\n\n slack_messages.append(\n SlackMessage(\n document_id=document_id,\n channel_id=channel_id,\n message_id=message_id,\n thread_id=thread_id,\n link=permalink,\n metadata=metadata,\n timestamp=doc_time,\n recency_bias=recency_bias,\n semantic_identifier=doc_sem_id,\n text=f\"{username}: {cleaned_text}\",\n highlighted_texts=highlighted_texts,\n slack_score=score,\n )\n )\n\n return SlackQueryResult(\n messages=slack_messages, filtered_channels=filtered_channels\n )\n\n\ndef merge_slack_messages(\n query_results: list[SlackQueryResult],\n) -> tuple[list[SlackMessage], dict[str, SlackMessage], set[str]]:\n \"\"\"Merge messages from multiple query results, deduplicating by document_id.\n\n Returns:\n Tuple of (merged_messages, docid_to_message, all_filtered_channels)\n \"\"\"\n merged_messages: list[SlackMessage] = []\n docid_to_message: dict[str, SlackMessage] = {}\n all_filtered_channels: set[str] = set()\n\n for result in query_results:\n # Collect filtered channels from all queries\n all_filtered_channels.update(result.filtered_channels)\n\n for message in result.messages:\n if message.document_id in docid_to_message:\n # update the score and highlighted texts, rest should be identical\n docid_to_message[message.document_id].slack_score = max(\n docid_to_message[message.document_id].slack_score,\n message.slack_score,\n )\n docid_to_message[message.document_id].highlighted_texts.update(\n message.highlighted_texts\n )\n continue\n\n # add the message to the list\n docid_to_message[message.document_id] = message\n merged_messages.append(message)\n\n # re-sort by score\n merged_messages.sort(key=lambda x: x.slack_score, reverse=True)\n\n return merged_messages, docid_to_message, all_filtered_channels\n\n\nclass SlackRateLimitError(Exception):\n \"\"\"Raised when Slack API returns a rate limit error (429).\"\"\"\n\n\nclass ThreadContextResult:\n \"\"\"Result wrapper for thread context fetch that captures error type.\"\"\"\n\n __slots__ = (\"text\", \"is_rate_limited\", \"is_error\")\n\n def __init__(\n self, text: str, is_rate_limited: bool = False, is_error: bool = False\n ):\n self.text = text\n self.is_rate_limited = is_rate_limited\n self.is_error = is_error\n\n @classmethod\n def success(cls, text: str) -> \"ThreadContextResult\":\n return cls(text)\n\n @classmethod\n def rate_limited(cls, original_text: str) -> \"ThreadContextResult\":\n return cls(original_text, is_rate_limited=True)\n\n @classmethod\n def error(cls, original_text: str) -> \"ThreadContextResult\":\n return cls(original_text, is_error=True)\n\n\ndef _fetch_thread_context(\n message: SlackMessage, access_token: str, team_id: str | None = None\n) -> ThreadContextResult:\n \"\"\"\n Fetch thread context for a message, returning a result object.\n\n Returns ThreadContextResult with:\n - success: enriched thread text\n - rate_limited: original text + flag indicating we should stop\n - error: original text for other failures (graceful degradation)\n \"\"\"\n channel_id = message.channel_id\n thread_id = message.thread_id\n message_id = message.message_id\n\n # If not a thread, return original text as success\n if thread_id is None:\n return ThreadContextResult.success(message.text)\n\n slack_client = WebClient(token=access_token, timeout=30)\n try:\n response = slack_client.conversations_replies(\n channel=channel_id,\n ts=thread_id,\n )\n response.validate()\n messages: list[dict[str, Any]] = response.get(\"messages\", [])\n except SlackApiError as e:\n # Check for rate limit error specifically\n if e.response and e.response.status_code == 429:\n logger.warning(\n f\"Slack rate limit hit while fetching thread context for {channel_id}/{thread_id}\"\n )\n return ThreadContextResult.rate_limited(message.text)\n # For other Slack errors, log and return original text\n logger.error(f\"Slack API error in thread context fetch: {e}\")\n return ThreadContextResult.error(message.text)\n except Exception as e:\n # Network errors, timeouts, etc - treat as recoverable error\n logger.error(f\"Unexpected error in thread context fetch: {e}\")\n return ThreadContextResult.error(message.text)\n\n # If empty response or single message (not a thread), return original text\n if len(messages) <= 1:\n return ThreadContextResult.success(message.text)\n\n # Build thread text from thread starter + context window around matched message\n thread_text = _build_thread_text(\n messages, message_id, thread_id, access_token, team_id, slack_client\n )\n return ThreadContextResult.success(thread_text)\n\n\ndef _build_thread_text(\n messages: list[dict[str, Any]],\n message_id: str,\n thread_id: str,\n access_token: str,\n team_id: str | None,\n slack_client: WebClient,\n) -> str:\n \"\"\"Build the thread text from messages.\"\"\"\n msg_text = messages[0].get(\"text\", \"\")\n msg_sender = messages[0].get(\"user\", \"\")\n thread_text = f\"<@{msg_sender}>: {msg_text}\"\n\n thread_text += \"\\n\\nReplies:\"\n if thread_id == message_id:\n message_id_idx = 0\n else:\n message_id_idx = next(\n (i for i, msg in enumerate(messages) if msg.get(\"ts\") == message_id), 0\n )\n if not message_id_idx:\n return thread_text\n\n start_idx = max(1, message_id_idx - SLACK_THREAD_CONTEXT_WINDOW)\n\n if start_idx > 1:\n thread_text += \"\\n...\"\n\n for i in range(start_idx, message_id_idx):\n msg_text = messages[i].get(\"text\", \"\")\n msg_sender = messages[i].get(\"user\", \"\")\n thread_text += f\"\\n\\n<@{msg_sender}>: {msg_text}\"\n\n msg_text = messages[message_id_idx].get(\"text\", \"\")\n msg_sender = messages[message_id_idx].get(\"user\", \"\")\n thread_text += f\"\\n\\n<@{msg_sender}>: {msg_text}\"\n\n # Add following replies\n len_replies = 0\n for msg in messages[message_id_idx + 1 :]:\n msg_text = msg.get(\"text\", \"\")\n msg_sender = msg.get(\"user\", \"\")\n reply = f\"\\n\\n<@{msg_sender}>: {msg_text}\"\n thread_text += reply\n\n len_replies += len(reply)\n if len_replies >= DOC_EMBEDDING_CONTEXT_SIZE * 4:\n thread_text += \"\\n...\"\n break\n\n # Replace user IDs with names using cached lookups\n userids: set[str] = set(re.findall(r\"<@([A-Z0-9]+)>\", thread_text))\n\n if team_id:\n user_profiles = batch_get_user_profiles(access_token, team_id, userids)\n for userid, name in user_profiles.items():\n thread_text = thread_text.replace(f\"<@{userid}>\", name)\n else:\n for userid in userids:\n try:\n response = slack_client.users_profile_get(user=userid)\n response.validate()\n profile: dict[str, Any] = response.get(\"profile\", {})\n user_name: str | None = profile.get(\"real_name\") or profile.get(\"email\")\n except SlackApiError as e:\n if \"user_not_found\" in str(e):\n logger.debug(\n f\"User {userid} not found (likely deleted/deactivated)\"\n )\n else:\n logger.warning(f\"Could not fetch profile for user {userid}: {e}\")\n continue\n if not user_name:\n continue\n thread_text = thread_text.replace(f\"<@{userid}>\", user_name)\n\n return thread_text\n\n\ndef fetch_thread_contexts_with_rate_limit_handling(\n slack_messages: list[SlackMessage],\n access_token: str,\n team_id: str | None,\n batch_size: int = SLACK_THREAD_CONTEXT_BATCH_SIZE,\n max_messages: int | None = MAX_SLACK_THREAD_CONTEXT_MESSAGES,\n) -> list[str]:\n \"\"\"\n Fetch thread contexts in controlled batches, stopping on rate limit.\n\n Distinguishes between error types:\n - Rate limit (429): Stop processing further batches\n - Other errors: Continue processing (graceful degradation)\n\n Args:\n slack_messages: Messages to fetch thread context for (should be sorted by relevance)\n access_token: Slack OAuth token\n team_id: Slack team ID for user profile caching\n batch_size: Number of concurrent API calls per batch\n max_messages: Maximum messages to fetch thread context for (None = no limit)\n\n Returns:\n List of thread texts, one per input message.\n Messages beyond max_messages or after rate limit get their original text.\n \"\"\"\n if not slack_messages:\n return []\n\n # Limit how many messages we fetch thread context for (if max_messages is set)\n if max_messages and max_messages < len(slack_messages):\n messages_for_context = slack_messages[:max_messages]\n messages_without_context = slack_messages[max_messages:]\n else:\n messages_for_context = slack_messages\n messages_without_context = []\n\n logger.info(\n f\"Fetching thread context for {len(messages_for_context)} of {len(slack_messages)} messages \"\n f\"(batch_size={batch_size}, max={max_messages or 'unlimited'})\"\n )\n\n results: list[str] = []\n rate_limited = False\n total_batches = (len(messages_for_context) + batch_size - 1) // batch_size\n rate_limit_batch = 0\n\n # Process in batches\n for i in range(0, len(messages_for_context), batch_size):\n current_batch = i // batch_size + 1\n\n if rate_limited:\n # Skip remaining batches, use original message text\n remaining = messages_for_context[i:]\n skipped_batches = total_batches - rate_limit_batch\n logger.warning(\n f\"Slack rate limit: skipping {len(remaining)} remaining messages \"\n f\"({skipped_batches} of {total_batches} batches). \"\n f\"Successfully enriched {len(results)} messages before rate limit.\"\n )\n results.extend([msg.text for msg in remaining])\n break\n\n batch = messages_for_context[i : i + batch_size]\n\n # _fetch_thread_context returns ThreadContextResult (never raises)\n # allow_failures=True is a safety net for any unexpected exceptions\n batch_results: list[ThreadContextResult | None] = (\n run_functions_tuples_in_parallel(\n [\n (\n _fetch_thread_context,\n (msg, access_token, team_id),\n )\n for msg in batch\n ],\n allow_failures=True,\n max_workers=batch_size,\n )\n )\n\n # Process results - ThreadContextResult tells us exactly what happened\n for j, result in enumerate(batch_results):\n if result is None:\n # Unexpected exception (shouldn't happen) - use original text, stop\n logger.error(f\"Unexpected None result for message {j} in batch\")\n results.append(batch[j].text)\n rate_limited = True\n rate_limit_batch = current_batch\n elif result.is_rate_limited:\n # Rate limit hit - use original text, stop further batches\n results.append(result.text)\n rate_limited = True\n rate_limit_batch = current_batch\n else:\n # Success or recoverable error - use the text (enriched or original)\n results.append(result.text)\n\n if rate_limited:\n logger.warning(\n f\"Slack rate limit (429) hit at batch {current_batch}/{total_batches} \"\n f\"while fetching thread context. Stopping further API calls.\"\n )\n\n # Add original text for messages we didn't fetch context for\n results.extend([msg.text for msg in messages_without_context])\n\n return results\n\n\ndef convert_slack_score(slack_score: float) -> float:\n \"\"\"\n Convert slack score to a score between 0 and 1.\n Will affect UI ordering and LLM ordering, but not the pruning.\n I.e., should have very little effect on the search/answer quality.\n \"\"\"\n return max(0.0, min(1.0, slack_score / 90_000))\n\n\n@log_function_time(print_only=True)\ndef slack_retrieval(\n query: ChunkIndexRequest,\n access_token: str,\n db_session: Session | None = None,\n connector: FederatedConnectorDetail | None = None, # noqa: ARG001\n entities: dict[str, Any] | None = None,\n limit: int | None = None,\n slack_event_context: SlackContext | None = None,\n bot_token: str | None = None, # Add bot token parameter\n team_id: str | None = None,\n # Pre-fetched data — when provided, avoids DB query (no session needed)\n search_settings: SearchSettings | None = None,\n) -> list[InferenceChunk]:\n \"\"\"\n Main entry point for Slack federated search with entity filtering.\n\n Applies entity filtering including:\n - Channel selection and exclusion\n - Date range extraction and enforcement\n - DM/private channel filtering\n - Multi-layer caching\n\n Args:\n query: Search query object\n access_token: User OAuth access token\n db_session: Database session (optional if search_settings provided)\n connector: Federated connector detail (unused, kept for backwards compat)\n entities: Connector-level config (entity filtering configuration)\n limit: Maximum number of results\n slack_event_context: Context when called from Slack bot\n bot_token: Bot token for enhanced permissions\n team_id: Slack team/workspace ID\n\n Returns:\n List of InferenceChunk objects\n \"\"\"\n # Use connector-level config\n entities = entities or {}\n\n if not entities:\n logger.debug(\"No entity configuration found, using defaults\")\n else:\n logger.debug(f\"Using entity configuration: {entities}\")\n\n # Extract limit from entity config if not explicitly provided\n query_limit = limit\n if entities:\n try:\n parsed_entities = SlackEntities(**entities)\n if limit is None:\n query_limit = parsed_entities.max_messages_per_query\n logger.debug(f\"Using max_messages_per_query from config: {query_limit}\")\n except Exception as e:\n logger.warning(f\"Error parsing entities for limit: {e}\")\n if limit is None:\n query_limit = 100 # Fallback default\n elif limit is None:\n query_limit = 100 # Default when no entities and no limit provided\n\n # Pre-fetch channel metadata from Redis cache and extract available channels\n # This avoids repeated Redis lookups during parallel search execution\n available_channels = None\n channel_metadata_dict = None\n if team_id:\n # Always fetch all channel types (include_private=True) to ensure single cache entry\n channel_metadata_dict = fetch_and_cache_channel_metadata(\n access_token, team_id, include_private=True\n )\n\n # Extract available channels list if needed for pattern matching\n available_channels = _extract_channel_data_from_entities(\n entities, channel_metadata_dict\n )\n\n # Query slack with entity filtering\n llm = get_default_llm()\n query_strings = build_slack_queries(query, llm, entities, available_channels)\n\n # Determine filtering based on entities OR context (bot)\n include_dm = False\n allowed_private_channel = None\n\n # Bot context overrides (if entities not specified)\n if slack_event_context and not entities:\n channel_type = slack_event_context.channel_type\n if channel_type == ChannelType.IM: # DM with user\n include_dm = True\n if channel_type == ChannelType.PRIVATE_CHANNEL:\n allowed_private_channel = slack_event_context.channel_id\n logger.debug(\n f\"Private channel context: will only allow messages from {allowed_private_channel} + public channels\"\n )\n\n # Build search tasks\n search_tasks = [\n (\n query_slack,\n (\n query_string,\n access_token,\n query_limit,\n allowed_private_channel,\n bot_token,\n include_dm,\n entities,\n available_channels,\n channel_metadata_dict,\n ),\n )\n for query_string in query_strings\n ]\n\n # If include_dm is True AND we're not already searching all channels,\n # add additional searches without channel filters.\n # This allows searching DMs/group DMs while still searching the specified channels.\n # Skip this if search_all_channels is already True (would be duplicate queries).\n if (\n entities\n and entities.get(\"include_dm\")\n and not entities.get(\"search_all_channels\")\n ):\n # Create a minimal entities dict that won't add channel filters\n # This ensures we search ALL conversations (DMs, group DMs, private channels)\n # BUT we still want to exclude channels specified in exclude_channels\n dm_entities = {\n \"include_dm\": True,\n \"include_private_channels\": entities.get(\"include_private_channels\", False),\n \"default_search_days\": entities.get(\"default_search_days\", 30),\n \"search_all_channels\": True,\n \"channels\": None,\n \"exclude_channels\": entities.get(\n \"exclude_channels\"\n ), # ALWAYS apply exclude_channels\n }\n\n for query_string in query_strings:\n search_tasks.append(\n (\n query_slack,\n (\n query_string,\n access_token,\n query_limit,\n allowed_private_channel,\n bot_token,\n include_dm,\n dm_entities,\n available_channels,\n channel_metadata_dict,\n ),\n )\n )\n\n # Execute searches in parallel\n results = run_functions_tuples_in_parallel(search_tasks)\n\n # Calculate stats for consolidated logging\n total_raw_messages = sum(len(r.messages) for r in results)\n\n # Merge and post-filter results\n slack_messages, docid_to_message, query_filtered_channels = merge_slack_messages(\n results\n )\n messages_after_dedup = len(slack_messages)\n\n # Post-filter by channel type (DM, private channel, etc.)\n # NOTE: We must post-filter because Slack's search.messages API only supports\n # filtering by channel NAME (via in:#channel syntax), not by channel TYPE.\n # There's no way to specify \"only public channels\" or \"exclude DMs\" in the query.\n # Start with channels filtered during query execution, then add post-filter channels\n filtered_out_channels: set[str] = set(query_filtered_channels)\n if entities and team_id:\n # Use pre-fetched channel metadata to avoid cache misses\n # Pass it directly instead of relying on Redis cache\n\n filtered_messages = []\n for msg in slack_messages:\n # Pass pre-fetched metadata to avoid cache lookups\n channel_type = get_channel_type(\n channel_id=msg.channel_id,\n channel_metadata=channel_metadata_dict,\n )\n if should_include_message(channel_type, entities):\n filtered_messages.append(msg)\n else:\n # Track unique channel name for summary\n channel_name = msg.metadata.get(\"channel\", msg.channel_id)\n filtered_out_channels.add(f\"{channel_name}({msg.channel_id})\")\n\n slack_messages = filtered_messages\n\n slack_messages = slack_messages[: limit or len(slack_messages)]\n\n # Log consolidated summary with request ID for correlation\n request_id = (\n slack_event_context.message_ts[:10]\n if slack_event_context and slack_event_context.message_ts\n else \"no-ctx\"\n )\n logger.info(\n f\"[req:{request_id}] Slack federated search: {len(search_tasks)} queries, \"\n f\"{total_raw_messages} raw msgs -> {messages_after_dedup} after dedup -> \"\n f\"{len(slack_messages)} final\"\n + (\n f\", filtered channels: {sorted(filtered_out_channels)}\"\n if filtered_out_channels\n else \"\"\n )\n )\n\n if not slack_messages:\n return []\n\n # Fetch thread context with rate limit handling and message limiting\n # Messages are already sorted by relevance (slack_score), so top N get full context\n thread_texts = fetch_thread_contexts_with_rate_limit_handling(\n slack_messages=slack_messages,\n access_token=access_token,\n team_id=team_id,\n )\n for slack_message, thread_text in zip(slack_messages, thread_texts):\n slack_message.text = thread_text\n\n # get the highlighted texts from shortest to longest\n highlighted_texts: set[str] = set()\n for slack_message in slack_messages:\n highlighted_texts.update(slack_message.highlighted_texts)\n sorted_highlighted_texts = sorted(highlighted_texts, key=len)\n\n # For queries without highlights (e.g., empty recency queries), we should keep all chunks\n has_highlights = len(sorted_highlighted_texts) > 0\n\n # convert slack messages to index documents\n index_docs: list[IndexingDocument] = []\n for slack_message in slack_messages:\n section: TextSection = TextSection(\n text=slack_message.text, link=slack_message.link\n )\n index_docs.append(\n IndexingDocument(\n id=slack_message.document_id,\n sections=[section],\n processed_sections=[section],\n source=DocumentSource.SLACK,\n title=slack_message.semantic_identifier,\n semantic_identifier=slack_message.semantic_identifier,\n metadata=slack_message.metadata,\n doc_updated_at=slack_message.timestamp,\n )\n )\n\n # chunk index docs into doc aware chunks\n # a single index doc can get split into multiple chunks\n if search_settings is None:\n if db_session is None:\n raise ValueError(\"Either db_session or search_settings must be provided\")\n search_settings = get_current_search_settings(db_session)\n embedder = DefaultIndexingEmbedder.from_db_search_settings(\n search_settings=search_settings\n )\n multipass_config = get_multipass_config(search_settings)\n enable_contextual_rag = (\n search_settings.enable_contextual_rag or ENABLE_CONTEXTUAL_RAG\n )\n chunker = Chunker(\n tokenizer=embedder.embedding_model.tokenizer,\n enable_multipass=multipass_config.multipass_indexing,\n enable_large_chunks=multipass_config.enable_large_chunks,\n enable_contextual_rag=enable_contextual_rag,\n )\n chunks = chunker.chunk(index_docs)\n\n # prune chunks without any highlighted texts\n # BUT: for recency queries without keywords, keep all chunks\n relevant_chunks: list[DocAwareChunk] = []\n chunkid_to_match_highlight: dict[str, str] = {}\n\n if not has_highlights:\n # No highlighted terms - keep all chunks (recency query)\n for chunk in chunks:\n chunk_id = f\"{chunk.source_document.id}__{chunk.chunk_id}\"\n relevant_chunks.append(chunk)\n chunkid_to_match_highlight[chunk_id] = chunk.content # No highlighting\n if limit and len(relevant_chunks) >= limit:\n break\n else:\n # Prune chunks that don't contain highlighted terms\n for chunk in chunks:\n match_highlight = chunk.content\n for highlight in sorted_highlighted_texts: # faster than re sub\n match_highlight = match_highlight.replace(\n highlight, f\"{highlight}\"\n )\n\n # if nothing got replaced, the chunk is irrelevant\n if len(match_highlight) == len(chunk.content):\n continue\n\n chunk_id = f\"{chunk.source_document.id}__{chunk.chunk_id}\"\n relevant_chunks.append(chunk)\n chunkid_to_match_highlight[chunk_id] = match_highlight\n if limit and len(relevant_chunks) >= limit:\n break\n\n # convert to inference chunks\n top_chunks: list[InferenceChunk] = []\n for chunk in relevant_chunks:\n document_id = chunk.source_document.id\n chunk_id = f\"{document_id}__{chunk.chunk_id}\"\n\n top_chunks.append(\n InferenceChunk(\n chunk_id=chunk.chunk_id,\n blurb=chunk.blurb,\n content=chunk.content,\n source_links=chunk.source_links,\n image_file_id=chunk.image_file_id,\n section_continuation=chunk.section_continuation,\n semantic_identifier=docid_to_message[document_id].semantic_identifier,\n document_id=document_id,\n source_type=DocumentSource.SLACK,\n title=chunk.title_prefix,\n boost=0,\n score=convert_slack_score(docid_to_message[document_id].slack_score),\n hidden=False,\n is_relevant=None,\n relevance_explanation=\"\",\n metadata=docid_to_message[document_id].metadata,\n match_highlights=[chunkid_to_match_highlight[chunk_id]],\n doc_summary=\"\",\n chunk_context=\"\",\n updated_at=docid_to_message[document_id].timestamp,\n is_federated=True,\n )\n )\n\n return top_chunks\n" }, { "path": "backend/onyx/context/search/federated/slack_search_utils.py", "content": "import fnmatch\nimport json\nimport re\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\n\nfrom pydantic import ValidationError\n\nfrom onyx.configs.app_configs import MAX_SLACK_QUERY_EXPANSIONS\nfrom onyx.context.search.federated.models import ChannelMetadata\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.federated_connectors.slack.models import SlackEntities\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.natural_language_processing.english_stopwords import ENGLISH_STOPWORDS_SET\nfrom onyx.onyxbot.slack.models import ChannelType\nfrom onyx.prompts.federated_search import SLACK_DATE_EXTRACTION_PROMPT\nfrom onyx.prompts.federated_search import SLACK_QUERY_EXPANSION_PROMPT\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Constants for date extraction heuristics\nDEFAULT_RECENCY_DAYS = 7\nDEFAULT_LATELY_DAYS = 14\nDAYS_PER_WEEK = 7\nDAYS_PER_MONTH = 30\nMAX_CONTENT_WORDS = 3\n\n# Punctuation to strip from words during analysis\nWORD_PUNCTUATION = \".,!?;:\\\"'#\"\n\nRECENCY_KEYWORDS = [\"recent\", \"latest\", \"newest\", \"last\"]\n\n# All Slack channel types for fetching metadata\nALL_CHANNEL_TYPES = [\n ChannelType.PUBLIC_CHANNEL.value,\n ChannelType.IM.value,\n ChannelType.MPIM.value,\n ChannelType.PRIVATE_CHANNEL.value,\n]\n\n# Map Slack API scopes to their corresponding channel types\n# This is used for graceful degradation when scopes are missing\nSCOPE_TO_CHANNEL_TYPE_MAP = {\n \"mpim:read\": ChannelType.MPIM.value,\n \"mpim:history\": ChannelType.MPIM.value,\n \"im:read\": ChannelType.IM.value,\n \"im:history\": ChannelType.IM.value,\n \"groups:read\": ChannelType.PRIVATE_CHANNEL.value,\n \"groups:history\": ChannelType.PRIVATE_CHANNEL.value,\n \"channels:read\": ChannelType.PUBLIC_CHANNEL.value,\n \"channels:history\": ChannelType.PUBLIC_CHANNEL.value,\n}\n\n\ndef get_channel_type_for_missing_scope(scope: str) -> str | None:\n \"\"\"Get the channel type that requires a specific Slack scope.\n\n Args:\n scope: The Slack API scope (e.g., 'mpim:read', 'im:history')\n\n Returns:\n The channel type string if scope is recognized, None otherwise\n\n Examples:\n >>> get_channel_type_for_missing_scope('mpim:read')\n 'mpim'\n >>> get_channel_type_for_missing_scope('im:read')\n 'im'\n >>> get_channel_type_for_missing_scope('unknown:scope')\n None\n \"\"\"\n return SCOPE_TO_CHANNEL_TYPE_MAP.get(scope)\n\n\ndef _parse_llm_code_block_response(response: str) -> str:\n \"\"\"Remove code block markers from LLM response if present.\n\n Handles responses wrapped in triple backticks (```) by removing\n the opening and closing markers.\n\n Args:\n response: Raw LLM response string\n\n Returns:\n Cleaned response with code block markers removed\n \"\"\"\n response_clean = response.strip()\n if response_clean.startswith(\"```\"):\n lines = response_clean.split(\"\\n\")\n lines = lines[1:]\n if lines and lines[-1].strip() == \"```\":\n lines = lines[:-1]\n response_clean = \"\\n\".join(lines)\n return response_clean\n\n\ndef is_recency_query(query: str) -> bool:\n \"\"\"Check if a query is primarily about recency (not content + recency).\n\n Returns True only for pure recency queries like \"recent messages\" or \"latest updates\",\n but False for queries with content + recency like \"golf scores last saturday\".\n \"\"\"\n # Check if query contains recency keywords\n has_recency_keyword = any(\n re.search(rf\"\\b{re.escape(keyword)}\\b\", query, flags=re.IGNORECASE)\n for keyword in RECENCY_KEYWORDS\n )\n\n if not has_recency_keyword:\n return False\n\n # Get combined stop words (English + Slack-specific)\n all_stop_words = _get_combined_stop_words()\n\n # Extract content words (excluding stop words)\n query_lower = query.lower()\n words = query_lower.split()\n\n # Count content words (not stop words, length > 2)\n content_word_count = 0\n for word in words:\n clean_word = word.strip(WORD_PUNCTUATION)\n if clean_word and len(clean_word) > 2 and clean_word not in all_stop_words:\n content_word_count += 1\n\n # If query has significant content words (>= 2), it's not a pure recency query\n # Examples:\n # - \"recent messages\" -> content_word_count = 0 -> pure recency\n # - \"golf scores last saturday\" -> content_word_count = 3 (golf, scores, saturday) -> not pure recency\n return content_word_count < 2\n\n\ndef extract_date_range_from_query(\n query: str,\n llm: LLM,\n default_search_days: int,\n) -> int:\n query_lower = query.lower()\n\n if re.search(r\"\\btoday(?:\\'?s)?\\b\", query_lower):\n return 0\n\n if re.search(r\"\\byesterday\\b\", query_lower):\n return min(1, default_search_days)\n\n # Handle \"last [day of week]\" - e.g., \"last monday\", \"last saturday\"\n days_of_week = [\n \"monday\",\n \"tuesday\",\n \"wednesday\",\n \"thursday\",\n \"friday\",\n \"saturday\",\n \"sunday\",\n ]\n for day in days_of_week:\n if re.search(rf\"\\b(?:last|this)\\s+{day}\\b\", query_lower):\n # Assume last occurrence of that day was within the past week\n return min(DAYS_PER_WEEK, default_search_days)\n\n match = re.search(r\"\\b(?:last|past)\\s+(\\d+)\\s+days?\\b\", query_lower)\n if match:\n days = int(match.group(1))\n return min(days, default_search_days)\n\n if re.search(r\"\\b(?:last|past|this)\\s+week\\b\", query_lower):\n return min(DAYS_PER_WEEK, default_search_days)\n\n match = re.search(r\"\\b(?:last|past)\\s+(\\d+)\\s+weeks?\\b\", query_lower)\n if match:\n weeks = int(match.group(1))\n return min(weeks * DAYS_PER_WEEK, default_search_days)\n\n if re.search(r\"\\b(?:last|past|this)\\s+month\\b\", query_lower):\n return min(DAYS_PER_MONTH, default_search_days)\n\n match = re.search(r\"\\b(?:last|past)\\s+(\\d+)\\s+months?\\b\", query_lower)\n if match:\n months = int(match.group(1))\n return min(months * DAYS_PER_MONTH, default_search_days)\n\n if re.search(r\"\\brecent(?:ly)?\\b\", query_lower):\n return min(DEFAULT_RECENCY_DAYS, default_search_days)\n\n if re.search(r\"\\blately\\b\", query_lower):\n return min(DEFAULT_LATELY_DAYS, default_search_days)\n\n try:\n prompt = SLACK_DATE_EXTRACTION_PROMPT.format(query=query)\n prompt_msg = UserMessage(content=prompt)\n\n # Call LLM with Braintrust tracing\n with llm_generation_span(\n llm=llm, flow=\"slack_date_extraction\", input_messages=[prompt_msg]\n ) as span_generation:\n llm_response = llm.invoke(prompt_msg)\n record_llm_response(span_generation, llm_response)\n response = llm_response_to_string(llm_response)\n\n response_clean = _parse_llm_code_block_response(response)\n\n try:\n data = json.loads(response_clean)\n if not isinstance(data, dict):\n logger.debug(\n f\"LLM date extraction returned non-dict response for query: \"\n f\"'{query}', using default: {default_search_days} days\"\n )\n return default_search_days\n\n days_back = data.get(\"days_back\")\n if days_back is None:\n logger.debug(\n f\"LLM date extraction returned null for query: '{query}', using default: {default_search_days} days\"\n )\n return default_search_days\n\n if not isinstance(days_back, (int, float)):\n logger.debug(\n f\"LLM date extraction returned non-numeric days_back for \"\n f\"query: '{query}', using default: {default_search_days} days\"\n )\n return default_search_days\n\n except json.JSONDecodeError:\n logger.debug(\n f\"Failed to parse LLM date extraction response for query: '{query}' \"\n f\"(response: '{response_clean}'), \"\n f\"using default: {default_search_days} days\"\n )\n return default_search_days\n\n return min(int(days_back), default_search_days)\n\n except Exception as e:\n logger.warning(f\"Error extracting date range with LLM for query '{query}': {e}\")\n return default_search_days\n\n\ndef matches_exclude_pattern(channel_name: str, patterns: list[str]) -> bool:\n if not patterns:\n return False\n\n channel_norm = channel_name.lower().strip().lstrip(\"#\")\n\n for pattern in patterns:\n pattern_norm = pattern.lower().strip().lstrip(\"#\")\n if fnmatch.fnmatch(channel_norm, pattern_norm):\n return True\n\n return False\n\n\ndef build_channel_query_filter(\n parsed_entities: SlackEntities | dict[str, Any],\n available_channels: list[str] | None = None,\n) -> str:\n # Parse entities if dict\n try:\n if isinstance(parsed_entities, dict):\n entities = SlackEntities(**parsed_entities)\n else:\n entities = parsed_entities\n except ValidationError:\n return \"\"\n\n search_all_channels = entities.search_all_channels\n\n if search_all_channels:\n if not entities.exclude_channels:\n return \"\"\n\n # Can't apply exclusions without available_channels\n if not available_channels:\n return \"\"\n\n excluded_channels = [\n ch\n for ch in available_channels\n if matches_exclude_pattern(ch, entities.exclude_channels)\n ]\n normalized_excluded = [ch.lstrip(\"#\") for ch in excluded_channels]\n\n exclusion_filters = [f\"-in:#{channel}\" for channel in normalized_excluded]\n return \" \".join(exclusion_filters)\n\n if not entities.channels:\n return \"\"\n\n included_channels: list[str] = []\n for pattern in entities.channels:\n pattern_norm = pattern.lstrip(\"#\")\n if \"*\" in pattern_norm or \"?\" in pattern_norm:\n # Glob patterns require available_channels\n if available_channels:\n matching = [\n ch\n for ch in available_channels\n if fnmatch.fnmatch(ch.lstrip(\"#\").lower(), pattern_norm.lower())\n ]\n included_channels.extend(matching)\n else:\n # Exact match: use directly or verify against available_channels\n if not available_channels or pattern_norm in [\n ch.lstrip(\"#\") for ch in available_channels\n ]:\n included_channels.append(pattern_norm)\n\n # Apply exclusions to included channels\n if entities.exclude_channels:\n included_channels = [\n ch\n for ch in included_channels\n if not matches_exclude_pattern(ch, entities.exclude_channels)\n ]\n\n if not included_channels:\n return \"\"\n\n normalized_channels = [ch.lstrip(\"#\") for ch in included_channels]\n filters = [f\"in:#{channel}\" for channel in normalized_channels]\n return \" \".join(filters)\n\n\ndef get_channel_type(\n channel_info: dict[str, Any] | None = None,\n channel_id: str | None = None,\n channel_metadata: dict[str, ChannelMetadata] | None = None,\n) -> ChannelType:\n \"\"\"\n Determine channel type from channel info dict or by looking up channel_id.\n\n Args:\n channel_info: Channel info dict from Slack API (direct mode)\n channel_id: Channel ID to look up (lookup mode)\n channel_metadata: Pre-fetched metadata dict (for lookup mode)\n\n Returns:\n ChannelType enum\n \"\"\"\n if channel_info is not None:\n if channel_info.get(\"is_im\"):\n return ChannelType.IM\n if channel_info.get(\"is_mpim\"):\n return ChannelType.MPIM\n if channel_info.get(\"is_private\"):\n return ChannelType.PRIVATE_CHANNEL\n return ChannelType.PUBLIC_CHANNEL\n\n # Lookup mode: get type from pre-fetched metadata\n if channel_id and channel_metadata:\n ch_meta = channel_metadata.get(channel_id)\n if ch_meta:\n type_str = ch_meta.get(\"type\")\n if type_str == ChannelType.IM.value:\n return ChannelType.IM\n elif type_str == ChannelType.MPIM.value:\n return ChannelType.MPIM\n elif type_str == ChannelType.PRIVATE_CHANNEL.value:\n return ChannelType.PRIVATE_CHANNEL\n return ChannelType.PUBLIC_CHANNEL\n\n return ChannelType.PUBLIC_CHANNEL\n\n\ndef should_include_message(channel_type: ChannelType, entities: dict[str, Any]) -> bool:\n include_dm = entities.get(\"include_dm\", False)\n include_group_dm = entities.get(\"include_group_dm\", False)\n include_private = entities.get(\"include_private_channels\", False)\n\n if channel_type == ChannelType.IM:\n return include_dm\n if channel_type == ChannelType.MPIM:\n return include_group_dm\n if channel_type == ChannelType.PRIVATE_CHANNEL:\n return include_private\n return True\n\n\ndef extract_channel_references_from_query(query_text: str) -> set[str]:\n \"\"\"Extract channel names referenced in the query text.\n\n Only matches explicit channel references with prepositions or # symbols:\n - \"in the office channel\"\n - \"from the office channel\"\n - \"in #office\"\n - \"from #office\"\n\n Does NOT match generic phrases like \"slack discussions\" or \"team channel\".\n\n Args:\n query_text: The user's query text\n\n Returns:\n Set of channel names (without # prefix)\n \"\"\"\n channel_references = set()\n query_lower = query_text.lower()\n\n # Only match channels with explicit prepositions (in/from) or # prefix\n # This prevents false positives like \"slack discussions\" being interpreted as channel \"slack\"\n channel_patterns = [\n r\"\\bin\\s+(?:the\\s+)?([a-z0-9_-]+)\\s+(?:slack\\s+)?channels?\\b\", # \"in the office channel\"\n r\"\\bfrom\\s+(?:the\\s+)?([a-z0-9_-]+)\\s+(?:slack\\s+)?channels?\\b\", # \"from the office channel\"\n r\"\\bin[:\\s]*#([a-z0-9_-]+)\\b\", # \"in #office\" or \"in:#office\"\n r\"\\bfrom[:\\s]*#([a-z0-9_-]+)\\b\", # \"from #office\" or \"from:#office\"\n ]\n\n for pattern in channel_patterns:\n matches = re.finditer(pattern, query_lower)\n for match in matches:\n channel_references.add(match.group(1))\n\n return channel_references\n\n\ndef validate_channel_references(\n channel_references: set[str],\n entities: dict[str, Any],\n available_channels: list[str] | None,\n) -> None:\n \"\"\"Validate that referenced channels exist and are allowed by entity config.\n\n Args:\n channel_references: Set of channel names extracted from query\n entities: Entity configuration dict\n available_channels: List of available channel names in workspace\n\n Raises:\n ValueError: If channel doesn't exist, is excluded, or not in inclusion list\n \"\"\"\n if not channel_references or not entities:\n return\n\n try:\n parsed_entities = SlackEntities(**entities)\n\n for channel_name in channel_references:\n # Check if channel exists\n if available_channels is not None:\n # Normalize for comparison (available_channels may or may not have #)\n normalized_available = [\n ch.lstrip(\"#\").lower() for ch in available_channels\n ]\n if channel_name.lower() not in normalized_available:\n raise ValueError(\n f\"Channel '{channel_name}' does not exist in your Slack workspace. \"\n f\"Please check the channel name and try again.\"\n )\n\n # Check if channel is in exclusion list\n if parsed_entities.exclude_channels:\n if matches_exclude_pattern(\n channel_name, parsed_entities.exclude_channels\n ):\n raise ValueError(\n f\"Channel '{channel_name}' is excluded from search by your configuration. \"\n f\"Please update your connector settings to search this channel.\"\n )\n\n # Check if channel is in inclusion list (when search_all_channels is False)\n if not parsed_entities.search_all_channels:\n if parsed_entities.channels:\n # Normalize channel lists for comparison\n normalized_channels = [\n ch.lstrip(\"#\").lower() for ch in parsed_entities.channels\n ]\n if channel_name.lower() not in normalized_channels:\n raise ValueError(\n f\"Channel '{channel_name}' is not in your configured channel list. \"\n f\"Please update your connector settings to include this channel.\"\n )\n\n except ValidationError:\n # If entities are malformed, skip validation\n pass\n\n\ndef build_channel_override_query(channel_references: set[str], time_filter: str) -> str:\n \"\"\"Build a Slack query with ONLY channel filters and time filter (no keywords).\n\n Args:\n channel_references: Set of channel names to search\n time_filter: Time filter string (e.g., \" after:2025-11-07\")\n\n Returns:\n Query string with __CHANNEL_OVERRIDE__ marker\n \"\"\"\n normalized_channels = [ch.lstrip(\"#\") for ch in channel_references]\n channel_filter = \" \".join([f\"in:#{channel}\" for channel in normalized_channels])\n return f\"__CHANNEL_OVERRIDE__ {channel_filter}{time_filter}\"\n\n\n# Slack-specific stop words (in addition to standard English stop words)\n# These include Slack-specific terms and temporal/recency keywords\nSLACK_SPECIFIC_STOP_WORDS = frozenset(\n RECENCY_KEYWORDS\n + [\n \"dm\",\n \"dms\",\n \"message\",\n \"messages\",\n \"channel\",\n \"channels\",\n \"slack\",\n \"post\",\n \"posted\",\n \"posting\",\n \"sent\",\n ]\n)\n\n\ndef _get_combined_stop_words() -> frozenset[str]:\n \"\"\"Get combined English + Slack-specific stop words.\n\n Returns a frozenset of stop words for filtering content words.\n\n Note: Currently only supports English stop words. Non-English queries\n may have suboptimal content word extraction. Future enhancement could\n detect query language and load appropriate stop words.\n \"\"\"\n return ENGLISH_STOPWORDS_SET | SLACK_SPECIFIC_STOP_WORDS\n\n\ndef extract_content_words_from_recency_query(\n query_text: str, channel_references: set[str]\n) -> list[str]:\n \"\"\"Extract meaningful content words from a recency query.\n\n Filters out English stop words, Slack-specific terms, channel references, and proper nouns.\n\n Args:\n query_text: The user's query text\n channel_references: Channel names to exclude from content words\n\n Returns:\n List of content words (up to MAX_CONTENT_WORDS)\n \"\"\"\n # Get combined stop words (English + Slack-specific)\n all_stop_words = _get_combined_stop_words()\n\n words = query_text.split()\n content_words = []\n\n for word in words:\n clean_word = word.lower().strip(WORD_PUNCTUATION)\n # Skip if it's a channel reference or a stop word\n if clean_word in channel_references:\n continue\n if clean_word and clean_word not in all_stop_words and len(clean_word) > 2:\n clean_word_orig = word.strip(WORD_PUNCTUATION)\n if clean_word_orig.lower() not in all_stop_words:\n content_words.append(clean_word_orig)\n\n # Filter out proper nouns (capitalized words)\n content_words_filtered = [word for word in content_words if not word[0].isupper()]\n\n return content_words_filtered[:MAX_CONTENT_WORDS]\n\n\ndef _is_valid_keyword_query(line: str) -> bool:\n \"\"\"Check if a line looks like a valid keyword query vs explanatory text.\n\n Returns False for lines that appear to be LLM explanations rather than keywords.\n \"\"\"\n # Reject lines that start with parentheses (explanatory notes)\n if line.startswith(\"(\"):\n return False\n\n # Reject lines that are too long (likely sentences, not keywords)\n # Keywords should be short - reject if > 50 chars or > 6 words\n if len(line) > 50 or len(line.split()) > 6:\n return False\n\n return True\n\n\ndef expand_query_with_llm(query_text: str, llm: LLM) -> list[str]:\n \"\"\"Use LLM to expand query into multiple search variations.\n\n Args:\n query_text: The user's original query\n llm: LLM instance to use for expansion\n\n Returns:\n List of rephrased query strings (up to MAX_SLACK_QUERY_EXPANSIONS)\n \"\"\"\n prompt = UserMessage(\n content=SLACK_QUERY_EXPANSION_PROMPT.format(\n query=query_text, max_queries=MAX_SLACK_QUERY_EXPANSIONS\n )\n )\n\n try:\n # Call LLM with Braintrust tracing\n with llm_generation_span(\n llm=llm, flow=\"slack_query_expansion\", input_messages=[prompt]\n ) as span_generation:\n llm_response = llm.invoke(prompt)\n record_llm_response(span_generation, llm_response)\n response = llm_response_to_string(llm_response)\n\n response_clean = _parse_llm_code_block_response(response)\n\n # Split into lines and filter out empty lines\n raw_queries = [\n line.strip() for line in response_clean.split(\"\\n\") if line.strip()\n ]\n\n # Filter out lines that look like explanatory text rather than keywords\n rephrased_queries = [q for q in raw_queries if _is_valid_keyword_query(q)]\n\n # Log if we filtered out garbage\n if len(raw_queries) != len(rephrased_queries):\n filtered_out = set(raw_queries) - set(rephrased_queries)\n logger.warning(f\"Filtered out non-keyword LLM responses: {filtered_out}\")\n\n # If no queries generated, use empty query\n if not rephrased_queries:\n logger.debug(\"No content keywords extracted from query expansion\")\n return [\"\"]\n\n logger.debug(\n f\"Expanded query into {len(rephrased_queries)} queries: {rephrased_queries}\"\n )\n return rephrased_queries[:MAX_SLACK_QUERY_EXPANSIONS]\n\n except Exception as e:\n logger.error(f\"Error expanding query: {e}\")\n return [query_text]\n\n\ndef build_slack_queries(\n query: ChunkIndexRequest,\n llm: LLM,\n entities: dict[str, Any] | None = None,\n available_channels: list[str] | None = None,\n) -> list[str]:\n \"\"\"Build Slack query strings with date filtering and query expansion.\"\"\"\n default_search_days = 30\n if entities:\n try:\n parsed_entities = SlackEntities(**entities)\n default_search_days = parsed_entities.default_search_days\n except ValidationError as e:\n logger.warning(f\"Invalid entities in build_slack_queries: {e}\")\n\n days_back = extract_date_range_from_query(\n query=query.query,\n llm=llm,\n default_search_days=default_search_days,\n )\n\n # get time filter\n time_filter = \"\"\n if days_back is not None and days_back >= 0:\n if days_back == 0:\n time_filter = \" on:today\"\n else:\n cutoff_date = datetime.now(timezone.utc) - timedelta(days=days_back)\n time_filter = f\" after:{cutoff_date.strftime('%Y-%m-%d')}\"\n\n # ALWAYS extract channel references from the query (not just for recency queries)\n channel_references = extract_channel_references_from_query(query.query)\n\n # Validate channel references against available channels and entity config\n # This will raise ValueError if channels are invalid\n if channel_references and entities:\n try:\n validate_channel_references(\n channel_references, entities, available_channels\n )\n logger.info(\n f\"Detected and validated channel references: {channel_references}\"\n )\n\n # If valid channels detected, use ONLY those channels with NO keywords\n # Return query with ONLY time filter + channel filter (no keywords)\n return [build_channel_override_query(channel_references, time_filter)]\n except ValueError as e:\n # If validation fails, log the error and continue with normal flow\n logger.warning(f\"Channel reference validation failed: {e}\")\n channel_references = set()\n\n # use llm to generate slack queries (use original query to use same keywords as the user)\n if is_recency_query(query.query):\n # For recency queries, extract content words (excluding channel names and stop words)\n content_words = extract_content_words_from_recency_query(\n query.query, channel_references\n )\n rephrased_queries = [\" \".join(content_words)] if content_words else [\"\"]\n else:\n # For other queries, use LLM to expand into multiple variations\n rephrased_queries = expand_query_with_llm(query.query, llm)\n\n # Build final query strings with time filters\n return [\n rephrased_query.strip() + time_filter\n for rephrased_query in rephrased_queries[:MAX_SLACK_QUERY_EXPANSIONS]\n ]\n" }, { "path": "backend/onyx/context/search/models.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\nfrom enum import Enum\nfrom typing import Any\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import SearchSettings\nfrom onyx.indexing.models import BaseChunk\nfrom onyx.indexing.models import IndexingSetting\nfrom onyx.tools.tool_implementations.web_search.models import WEB_SEARCH_PREFIX\n\n\nclass QueryExpansions(BaseModel):\n keywords_expansions: list[str] | None = None\n semantic_expansions: list[str] | None = None\n\n\nclass QueryExpansionType(Enum):\n KEYWORD = \"keyword\"\n SEMANTIC = \"semantic\"\n\n\nclass SearchSettingsCreationRequest(IndexingSetting):\n @classmethod\n def from_db_model(\n cls, search_settings: SearchSettings\n ) -> \"SearchSettingsCreationRequest\":\n indexing_setting = IndexingSetting.from_db_model(search_settings)\n return cls(**indexing_setting.model_dump())\n\n\nclass SavedSearchSettings(IndexingSetting):\n # Previously this contained also Inference time settings. Keeping this wrapper class around\n # as there may again be inference time settings that may get added.\n @classmethod\n def from_db_model(cls, search_settings: SearchSettings) -> \"SavedSearchSettings\":\n return cls(\n # Indexing Setting\n model_name=search_settings.model_name,\n model_dim=search_settings.model_dim,\n normalize=search_settings.normalize,\n query_prefix=search_settings.query_prefix,\n passage_prefix=search_settings.passage_prefix,\n provider_type=search_settings.provider_type,\n index_name=search_settings.index_name,\n multipass_indexing=search_settings.multipass_indexing,\n embedding_precision=search_settings.embedding_precision,\n reduced_dimension=search_settings.reduced_dimension,\n switchover_type=search_settings.switchover_type,\n enable_contextual_rag=search_settings.enable_contextual_rag,\n contextual_rag_llm_name=search_settings.contextual_rag_llm_name,\n contextual_rag_llm_provider=search_settings.contextual_rag_llm_provider,\n )\n\n\nclass Tag(BaseModel):\n tag_key: str\n tag_value: str\n\n\nclass BaseFilters(BaseModel):\n source_type: list[DocumentSource] | None = None\n document_set: list[str] | None = None\n time_cutoff: datetime | None = None\n tags: list[Tag] | None = None\n\n\nclass UserFileFilters(BaseModel):\n # Scopes search to user files tagged with a given project/persona in Vespa.\n # These are NOT simply the IDs of the current project or persona — they are\n # only set when the persona's/project's user files overflowed the LLM\n # context window and must be searched via vector DB instead of being loaded\n # directly into the prompt.\n project_id_filter: int | None = None\n persona_id_filter: int | None = None\n\n\nclass AssistantKnowledgeFilters(BaseModel):\n \"\"\"Filters for knowledge attached to an assistant (persona).\n\n These filters scope search to documents/folders explicitly attached\n to the assistant. When present, only documents matching these criteria\n are searched (in addition to ACL filtering).\n \"\"\"\n\n # Document IDs explicitly attached to the assistant\n attached_document_ids: list[str] | None = None\n # Hierarchy node IDs (folders/spaces) attached to the assistant.\n # Matches chunks where ancestor_hierarchy_node_ids contains any of these.\n hierarchy_node_ids: list[int] | None = None\n\n\nclass IndexFilters(BaseFilters, UserFileFilters, AssistantKnowledgeFilters):\n # NOTE: These strings must be formatted in the same way as the output of\n # DocumentAccess::to_acl.\n access_control_list: list[str] | None\n tenant_id: str | None = None\n\n\nclass BasicChunkRequest(BaseModel):\n query: str\n\n # In case the caller wants to override the weighting between semantic and keyword search.\n hybrid_alpha: float | None = None\n\n # In case some queries favor recency more than other queries.\n recency_bias_multiplier: float = 1.0\n\n limit: int | None = None\n\n\nclass ChunkSearchRequest(BasicChunkRequest):\n # Final filters are calculated from these\n user_selected_filters: BaseFilters | None = None\n\n # Use with caution!\n bypass_acl: bool = False\n\n\n# From the Chat Session we know what project (if any) this search should include\n# From the user uploads and persona uploaded files, we know which of those to include\nclass ChunkIndexRequest(BasicChunkRequest):\n # Calculated final filters\n filters: IndexFilters\n\n query_keywords: list[str] | None = None\n\n\nclass ContextExpansionType(str, Enum):\n NOT_RELEVANT = \"not_relevant\"\n MAIN_SECTION_ONLY = \"main_section_only\"\n INCLUDE_ADJACENT_SECTIONS = \"include_adjacent_sections\"\n FULL_DOCUMENT = \"full_document\"\n\n\nclass InferenceChunk(BaseChunk):\n document_id: str\n source_type: DocumentSource\n semantic_identifier: str\n title: str | None # Separate from Semantic Identifier though often same\n boost: int\n score: float | None\n hidden: bool\n is_relevant: bool | None = None\n relevance_explanation: str | None = None\n # TODO(andrei): Ideally we could improve this to where each value is just a\n # list of strings.\n metadata: dict[str, str | list[str]]\n # Matched sections in the chunk. Uses Vespa syntax e.g. TEXT\n # to specify that a set of words should be highlighted. For example:\n # [\"the answer is 42\", \"he couldn't find an answer\"]\n match_highlights: list[str]\n doc_summary: str\n chunk_context: str\n\n # when the doc was last updated\n updated_at: datetime | None\n primary_owners: list[str] | None = None\n secondary_owners: list[str] | None = None\n large_chunk_reference_ids: list[int] = Field(default_factory=list)\n\n is_federated: bool = False\n\n @property\n def unique_id(self) -> str:\n return f\"{self.document_id}__{self.chunk_id}\"\n\n def __repr__(self) -> str:\n blurb_words = self.blurb.split()\n short_blurb = \"\"\n for word in blurb_words:\n if not short_blurb:\n short_blurb = word\n continue\n if len(short_blurb) > 25:\n break\n short_blurb += \" \" + word\n return f\"Inference Chunk: {self.document_id} - {short_blurb}...\"\n\n def __eq__(self, other: Any) -> bool:\n if not isinstance(other, InferenceChunk):\n return False\n return (self.document_id, self.chunk_id) == (other.document_id, other.chunk_id)\n\n def __hash__(self) -> int:\n return hash((self.document_id, self.chunk_id))\n\n def __lt__(self, other: Any) -> bool:\n if not isinstance(other, InferenceChunk):\n return NotImplemented\n if self.score is None:\n if other.score is None:\n return self.chunk_id > other.chunk_id\n return True\n if other.score is None:\n return False\n if self.score == other.score:\n return self.chunk_id > other.chunk_id\n return self.score < other.score\n\n def __gt__(self, other: Any) -> bool:\n if not isinstance(other, InferenceChunk):\n return NotImplemented\n if self.score is None:\n return False\n if other.score is None:\n return True\n if self.score == other.score:\n return self.chunk_id < other.chunk_id\n return self.score > other.score\n\n\nclass InferenceChunkUncleaned(InferenceChunk):\n metadata_suffix: str | None\n\n def to_inference_chunk(self) -> InferenceChunk:\n # Create a dict of all fields except 'metadata_suffix'\n # Assumes the cleaning has already been applied and just needs to translate to the right type\n inference_chunk_data = {\n k: v\n for k, v in self.model_dump().items()\n if k\n not in [\"metadata_suffix\"] # May be other fields to throw out in the future\n }\n return InferenceChunk(**inference_chunk_data)\n\n\nclass InferenceSection(BaseModel):\n \"\"\"Section list of chunks with a combined content. A section could be a single chunk, several\n chunks from the same document or the entire document.\"\"\"\n\n center_chunk: InferenceChunk\n chunks: list[InferenceChunk]\n combined_content: str\n\n\nclass SearchDoc(BaseModel):\n document_id: str\n chunk_ind: int\n semantic_identifier: str\n link: str | None = None\n blurb: str\n source_type: DocumentSource\n boost: int\n # Whether the document is hidden when doing a standard search\n # since a standard search will never find a hidden doc, this can only ever\n # be `True` when doing an admin search\n hidden: bool\n metadata: dict[str, str | list[str]]\n score: float | None = None\n is_relevant: bool | None = None\n relevance_explanation: str | None = None\n # Matched sections in the doc. Uses Vespa syntax e.g. TEXT\n # to specify that a set of words should be highlighted. For example:\n # [\"the answer is 42\", \"the answer is 42\"\"]\n match_highlights: list[str]\n # when the doc was last updated\n updated_at: datetime | None = None\n primary_owners: list[str] | None = None\n secondary_owners: list[str] | None = None\n is_internet: bool = False\n\n @classmethod\n def from_chunks_or_sections(\n cls,\n items: \"Sequence[InferenceChunk | InferenceSection] | None\",\n ) -> list[\"SearchDoc\"]:\n \"\"\"Convert a sequence of InferenceChunk or InferenceSection objects to SearchDoc objects.\"\"\"\n if not items:\n return []\n\n search_docs = [\n cls(\n document_id=(\n chunk := (\n item.center_chunk\n if isinstance(item, InferenceSection)\n else item\n )\n ).document_id,\n chunk_ind=chunk.chunk_id,\n semantic_identifier=chunk.semantic_identifier or \"Unknown\",\n link=chunk.source_links[0] if chunk.source_links else None,\n blurb=chunk.blurb,\n source_type=chunk.source_type,\n boost=chunk.boost,\n hidden=chunk.hidden,\n metadata=chunk.metadata,\n score=chunk.score,\n match_highlights=chunk.match_highlights,\n updated_at=chunk.updated_at,\n primary_owners=chunk.primary_owners,\n secondary_owners=chunk.secondary_owners,\n is_internet=False,\n )\n for item in items\n ]\n\n return search_docs\n\n # TODO - there is likely a way to clean this all up and not have the switch between these\n @classmethod\n def from_saved_search_doc(cls, saved_search_doc: \"SavedSearchDoc\") -> \"SearchDoc\":\n \"\"\"Convert a SavedSearchDoc to SearchDoc by dropping the db_doc_id field.\"\"\"\n saved_search_doc_data = saved_search_doc.model_dump()\n # Remove db_doc_id as it's not part of SearchDoc\n saved_search_doc_data.pop(\"db_doc_id\", None)\n return cls(**saved_search_doc_data)\n\n @classmethod\n def from_saved_search_docs(\n cls, saved_search_docs: list[\"SavedSearchDoc\"]\n ) -> list[\"SearchDoc\"]:\n return [\n cls.from_saved_search_doc(saved_search_doc)\n for saved_search_doc in saved_search_docs\n ]\n\n def model_dump(self, *args: list, **kwargs: dict[str, Any]) -> dict[str, Any]: # type: ignore\n initial_dict = super().model_dump(*args, **kwargs) # type: ignore\n initial_dict[\"updated_at\"] = (\n self.updated_at.isoformat() if self.updated_at else None\n )\n return initial_dict\n\n\nclass SearchDocsResponse(BaseModel):\n search_docs: list[SearchDoc]\n # Maps the citation number to the document id\n # Since these are no longer just links on the frontend but instead document cards, mapping it to the\n # document id is the most staightforward way.\n citation_mapping: dict[int, str]\n\n # For cases where the frontend only needs to display a subset of the search docs\n # The whole list is typically still needed for later steps but this set should be saved separately\n displayed_docs: list[SearchDoc] | None = None\n\n\nclass SavedSearchDoc(SearchDoc):\n db_doc_id: int\n score: float | None = 0.0\n\n @classmethod\n def from_search_doc(\n cls, search_doc: SearchDoc, db_doc_id: int = 0\n ) -> \"SavedSearchDoc\":\n \"\"\"IMPORTANT: careful using this and not providing a db_doc_id If db_doc_id is not\n provided, it won't be able to actually fetch the saved doc and info later on. So only skip\n providing this if the SavedSearchDoc will not be used in the future\"\"\"\n search_doc_data = search_doc.model_dump()\n search_doc_data[\"score\"] = search_doc_data.get(\"score\") or 0.0\n return cls(**search_doc_data, db_doc_id=db_doc_id)\n\n @classmethod\n def from_dict(cls, data: dict[str, Any]) -> \"SavedSearchDoc\":\n \"\"\"Create SavedSearchDoc from serialized dictionary data (e.g., from database JSON)\"\"\"\n return cls(**data)\n\n @classmethod\n def from_url(cls, url: str) -> \"SavedSearchDoc\":\n \"\"\"Create a SavedSearchDoc from a URL for internet search documents.\n\n Uses the INTERNET_SEARCH_DOC_ prefix for document_id to match the format\n used by inference sections created from internet content.\n \"\"\"\n return cls(\n # db_doc_id can be a filler value since these docs are not saved to the database.\n db_doc_id=0,\n document_id=WEB_SEARCH_PREFIX + url,\n chunk_ind=0,\n semantic_identifier=url,\n link=url,\n blurb=\"\",\n source_type=DocumentSource.WEB,\n boost=1,\n hidden=False,\n metadata={},\n score=0.0,\n is_relevant=None,\n relevance_explanation=None,\n match_highlights=[],\n updated_at=None,\n primary_owners=None,\n secondary_owners=None,\n is_internet=True,\n )\n\n def __lt__(self, other: Any) -> bool:\n if not isinstance(other, SavedSearchDoc):\n return NotImplemented\n self_score = self.score if self.score is not None else 0.0\n other_score = other.score if other.score is not None else 0.0\n return self_score < other_score\n\n\nclass SavedSearchDocWithContent(SavedSearchDoc):\n \"\"\"Used for endpoints that need to return the actual contents of the retrieved\n section in addition to the match_highlights.\"\"\"\n\n content: str\n\n\nclass PersonaSearchInfo(BaseModel):\n \"\"\"Snapshot of persona data needed by the search pipeline.\n\n Extracted from the ORM Persona before the DB session is released so that\n SearchTool and search_pipeline never lazy-load relationships post-commit.\n \"\"\"\n\n document_set_names: list[str]\n search_start_date: datetime | None\n attached_document_ids: list[str]\n hierarchy_node_ids: list[int]\n" }, { "path": "backend/onyx/context/search/pipeline.py", "content": "from collections import defaultdict\nfrom datetime import datetime\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.context.search.models import BaseFilters\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.context.search.models import ChunkSearchRequest\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceSection\nfrom onyx.context.search.models import PersonaSearchInfo\nfrom onyx.context.search.preprocessing.access_filters import (\n build_access_filters_for_user,\n)\nfrom onyx.context.search.retrieval.search_runner import search_chunks\nfrom onyx.context.search.utils import inference_section_from_chunks\nfrom onyx.db.models import User\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo\nfrom onyx.llm.interfaces import LLM\nfrom onyx.natural_language_processing.english_stopwords import strip_stopwords\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.secondary_llm_flows.source_filter import extract_source_filter\nfrom onyx.secondary_llm_flows.time_filter import extract_time_filter\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import FunctionCall\nfrom onyx.utils.threadpool_concurrency import run_functions_in_parallel\nfrom onyx.utils.timing import log_function_time\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\n@log_function_time(print_only=True)\ndef _build_index_filters(\n user_provided_filters: BaseFilters | None,\n user: User, # Used for ACLs, anonymous users only see public docs\n project_id_filter: int | None,\n persona_id_filter: int | None,\n persona_document_sets: list[str] | None,\n persona_time_cutoff: datetime | None,\n db_session: Session | None = None,\n auto_detect_filters: bool = False,\n query: str | None = None,\n llm: LLM | None = None,\n bypass_acl: bool = False,\n # Assistant knowledge filters\n attached_document_ids: list[str] | None = None,\n hierarchy_node_ids: list[int] | None = None,\n # Pre-fetched ACL filters (skips DB query when provided)\n acl_filters: list[str] | None = None,\n) -> IndexFilters:\n if auto_detect_filters and (llm is None or query is None):\n raise RuntimeError(\"LLM and query are required for auto detect filters\")\n\n base_filters = user_provided_filters or BaseFilters()\n\n document_set_filter = (\n base_filters.document_set\n if base_filters.document_set is not None\n else persona_document_sets\n )\n\n time_filter = base_filters.time_cutoff or persona_time_cutoff\n source_filter = base_filters.source_type\n\n detected_time_filter = None\n detected_source_filter = None\n if auto_detect_filters:\n time_filter_fnc = FunctionCall(extract_time_filter, (query, llm), {})\n if not source_filter:\n source_filter_fnc = FunctionCall(\n extract_source_filter, (query, llm, db_session), {}\n )\n else:\n source_filter_fnc = None\n\n functions_to_run = [fn for fn in [time_filter_fnc, source_filter_fnc] if fn]\n parallel_results = run_functions_in_parallel(functions_to_run)\n # Detected favor recent is not used for now\n detected_time_filter, _detected_favor_recent = parallel_results[\n time_filter_fnc.result_id\n ]\n if source_filter_fnc:\n detected_source_filter = parallel_results[source_filter_fnc.result_id]\n\n # If the detected time filter is more recent, use that one\n if time_filter and detected_time_filter and detected_time_filter > time_filter:\n time_filter = detected_time_filter\n\n # If the user has explicitly set a source filter, use that one\n if not source_filter and detected_source_filter:\n source_filter = detected_source_filter\n\n if bypass_acl:\n user_acl_filters = None\n elif acl_filters is not None:\n user_acl_filters = acl_filters\n else:\n if db_session is None:\n raise ValueError(\"Either db_session or acl_filters must be provided\")\n user_acl_filters = build_access_filters_for_user(user, db_session)\n\n final_filters = IndexFilters(\n project_id_filter=project_id_filter,\n persona_id_filter=persona_id_filter,\n source_type=source_filter,\n document_set=document_set_filter,\n time_cutoff=time_filter,\n tags=base_filters.tags,\n access_control_list=user_acl_filters,\n tenant_id=get_current_tenant_id() if MULTI_TENANT else None,\n # Assistant knowledge filters\n attached_document_ids=attached_document_ids,\n hierarchy_node_ids=hierarchy_node_ids,\n )\n\n return final_filters\n\n\ndef merge_individual_chunks(\n chunks: list[InferenceChunk],\n) -> list[InferenceSection]:\n \"\"\"Merge adjacent chunks from the same document into sections.\n\n Chunks are considered adjacent if their chunk_ids differ by 1 and they\n are from the same document. The section maintains the position of the\n first chunk in the original list.\n \"\"\"\n if not chunks:\n return []\n\n # Create a mapping from (document_id, chunk_id) to original index\n # This helps us find the chunk that appears first in the original list\n chunk_to_original_index: dict[tuple[str, int], int] = {}\n for idx, chunk in enumerate(chunks):\n chunk_to_original_index[(chunk.document_id, chunk.chunk_id)] = idx\n\n # Group chunks by document_id\n doc_chunks: dict[str, list[InferenceChunk]] = defaultdict(list)\n for chunk in chunks:\n doc_chunks[chunk.document_id].append(chunk)\n\n # For each document, sort chunks by chunk_id to identify adjacent chunks\n for doc_id in doc_chunks:\n doc_chunks[doc_id].sort(key=lambda c: c.chunk_id)\n\n # Create a mapping from (document_id, chunk_id) to the section it belongs to\n # This helps us maintain the original order\n chunk_to_section: dict[tuple[str, int], InferenceSection] = {}\n\n # Process each document's chunks\n for doc_id, doc_chunk_list in doc_chunks.items():\n if not doc_chunk_list:\n continue\n\n # Group adjacent chunks into sections\n current_section_chunks = [doc_chunk_list[0]]\n\n for i in range(1, len(doc_chunk_list)):\n prev_chunk = doc_chunk_list[i - 1]\n curr_chunk = doc_chunk_list[i]\n\n # Check if chunks are adjacent (chunk_id difference is 1)\n if curr_chunk.chunk_id == prev_chunk.chunk_id + 1:\n # Add to current section\n current_section_chunks.append(curr_chunk)\n else:\n # Create section from previous chunks\n # Find the chunk that appears first in the original list\n center_chunk = min(\n current_section_chunks,\n key=lambda c: chunk_to_original_index.get(\n (c.document_id, c.chunk_id), float(\"inf\")\n ),\n )\n section = inference_section_from_chunks(\n center_chunk=center_chunk,\n chunks=current_section_chunks.copy(),\n )\n if section:\n for chunk in current_section_chunks:\n chunk_to_section[(chunk.document_id, chunk.chunk_id)] = section\n\n # Start new section\n current_section_chunks = [curr_chunk]\n\n # Create section for the last group\n if current_section_chunks:\n # Find the chunk that appears first in the original list\n center_chunk = min(\n current_section_chunks,\n key=lambda c: chunk_to_original_index.get(\n (c.document_id, c.chunk_id), float(\"inf\")\n ),\n )\n section = inference_section_from_chunks(\n center_chunk=center_chunk,\n chunks=current_section_chunks.copy(),\n )\n if section:\n for chunk in current_section_chunks:\n chunk_to_section[(chunk.document_id, chunk.chunk_id)] = section\n\n # Build result list maintaining original order\n # Use (document_id, chunk_id) of center_chunk as unique identifier for sections\n seen_section_ids: set[tuple[str, int]] = set()\n result: list[InferenceSection] = []\n\n for chunk in chunks:\n section = chunk_to_section.get((chunk.document_id, chunk.chunk_id))\n if section:\n section_id = (\n section.center_chunk.document_id,\n section.center_chunk.chunk_id,\n )\n if section_id not in seen_section_ids:\n seen_section_ids.add(section_id)\n result.append(section)\n else:\n # Chunk wasn't part of any merged section, create a single-chunk section\n single_section = inference_section_from_chunks(\n center_chunk=chunk,\n chunks=[chunk],\n )\n if single_section:\n single_section_id = (\n single_section.center_chunk.document_id,\n single_section.center_chunk.chunk_id,\n )\n if single_section_id not in seen_section_ids:\n seen_section_ids.add(single_section_id)\n result.append(single_section)\n\n return result\n\n\n@log_function_time(print_only=True, debug_only=True)\ndef search_pipeline(\n # Query and settings\n chunk_search_request: ChunkSearchRequest,\n # Document index to search over\n # Note that federated sources will also be used (not related to this arg)\n document_index: DocumentIndex,\n # Used for ACLs and federated search, anonymous users only see public docs\n user: User,\n # Pre-extracted persona search configuration (None when no persona)\n persona_search_info: PersonaSearchInfo | None,\n db_session: Session | None = None,\n auto_detect_filters: bool = False,\n llm: LLM | None = None,\n # Vespa metadata filters for overflowing user files. NOT the raw IDs\n # of the current project/persona — only set when user files couldn't fit\n # in the LLM context and need to be searched via vector DB.\n project_id_filter: int | None = None,\n persona_id_filter: int | None = None,\n # Pre-fetched data — when provided, avoids DB queries (no session needed)\n acl_filters: list[str] | None = None,\n embedding_model: EmbeddingModel | None = None,\n prefetched_federated_retrieval_infos: list[FederatedRetrievalInfo] | None = None,\n) -> list[InferenceChunk]:\n persona_document_sets: list[str] | None = (\n persona_search_info.document_set_names if persona_search_info else None\n )\n persona_time_cutoff: datetime | None = (\n persona_search_info.search_start_date if persona_search_info else None\n )\n attached_document_ids: list[str] | None = (\n persona_search_info.attached_document_ids or None\n if persona_search_info\n else None\n )\n hierarchy_node_ids: list[int] | None = (\n persona_search_info.hierarchy_node_ids or None if persona_search_info else None\n )\n\n filters = _build_index_filters(\n user_provided_filters=chunk_search_request.user_selected_filters,\n user=user,\n project_id_filter=project_id_filter,\n persona_id_filter=persona_id_filter,\n persona_document_sets=persona_document_sets,\n persona_time_cutoff=persona_time_cutoff,\n db_session=db_session,\n auto_detect_filters=auto_detect_filters,\n query=chunk_search_request.query,\n llm=llm,\n bypass_acl=chunk_search_request.bypass_acl,\n attached_document_ids=attached_document_ids,\n hierarchy_node_ids=hierarchy_node_ids,\n acl_filters=acl_filters,\n )\n\n query_keywords = strip_stopwords(chunk_search_request.query)\n\n query_request = ChunkIndexRequest(\n query=chunk_search_request.query,\n hybrid_alpha=chunk_search_request.hybrid_alpha,\n recency_bias_multiplier=chunk_search_request.recency_bias_multiplier,\n query_keywords=query_keywords,\n filters=filters,\n limit=chunk_search_request.limit,\n )\n\n retrieved_chunks = search_chunks(\n query_request=query_request,\n user_id=user.id if user else None,\n document_index=document_index,\n db_session=db_session,\n embedding_model=embedding_model,\n prefetched_federated_retrieval_infos=prefetched_federated_retrieval_infos,\n )\n\n # For some specific connectors like Salesforce, a user that has access to an object doesn't mean\n # that they have access to all of the fields of the object.\n censored_chunks: list[InferenceChunk] = fetch_ee_implementation_or_noop(\n \"onyx.external_permissions.post_query_censoring\",\n \"_post_query_chunk_censoring\",\n retrieved_chunks,\n )(\n chunks=retrieved_chunks,\n user=user,\n )\n\n return censored_chunks\n" }, { "path": "backend/onyx/context/search/preprocessing/access_filters.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.access.access import get_acl_for_user\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.db.models import User\n\n\ndef build_access_filters_for_user(user: User, session: Session) -> list[str]:\n user_acl = get_acl_for_user(user, session)\n return list(user_acl)\n\n\ndef build_user_only_filters(user: User, db_session: Session) -> IndexFilters:\n user_acl_filters = build_access_filters_for_user(user, db_session)\n return IndexFilters(\n source_type=None,\n document_set=None,\n time_cutoff=None,\n tags=None,\n access_control_list=user_acl_filters,\n )\n" }, { "path": "backend/onyx/context/search/retrieval/search_runner.py", "content": "from collections.abc import Callable\nfrom uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.chat_configs import HYBRID_ALPHA\nfrom onyx.configs.chat_configs import NUM_RETURNED_HITS\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceSection\nfrom onyx.context.search.models import QueryExpansionType\nfrom onyx.context.search.utils import get_query_embedding\nfrom onyx.context.search.utils import inference_section_from_chunks\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces_new import DocumentIndex as NewDocumentIndex\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchOldDocumentIndex,\n)\nfrom onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo\nfrom onyx.federated_connectors.federated_retrieval import (\n get_federated_retrieval_functions,\n)\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\n\nlogger = setup_logger()\n\n\ndef combine_retrieval_results(\n chunk_sets: list[list[InferenceChunk]],\n) -> list[InferenceChunk]:\n all_chunks = [chunk for chunk_set in chunk_sets for chunk in chunk_set]\n\n unique_chunks: dict[tuple[str, int], InferenceChunk] = {}\n for chunk in all_chunks:\n key = (chunk.document_id, chunk.chunk_id)\n if key not in unique_chunks:\n unique_chunks[key] = chunk\n continue\n\n stored_chunk_score = unique_chunks[key].score or 0\n this_chunk_score = chunk.score or 0\n if stored_chunk_score < this_chunk_score:\n unique_chunks[key] = chunk\n\n sorted_chunks = sorted(\n unique_chunks.values(), key=lambda x: x.score or 0, reverse=True\n )\n\n return sorted_chunks\n\n\ndef _embed_and_hybrid_search(\n query_request: ChunkIndexRequest,\n document_index: DocumentIndex,\n db_session: Session | None = None,\n embedding_model: EmbeddingModel | None = None,\n) -> list[InferenceChunk]:\n query_embedding = get_query_embedding(\n query_request.query,\n db_session=db_session,\n embedding_model=embedding_model,\n )\n\n hybrid_alpha = query_request.hybrid_alpha or HYBRID_ALPHA\n\n top_chunks = document_index.hybrid_retrieval(\n query=query_request.query,\n query_embedding=query_embedding,\n final_keywords=query_request.query_keywords,\n filters=query_request.filters,\n hybrid_alpha=hybrid_alpha,\n time_decay_multiplier=query_request.recency_bias_multiplier,\n num_to_retrieve=query_request.limit or NUM_RETURNED_HITS,\n ranking_profile_type=(\n QueryExpansionType.KEYWORD\n if hybrid_alpha <= 0.3\n else QueryExpansionType.SEMANTIC\n ),\n )\n\n return top_chunks\n\n\ndef _keyword_search(\n query_request: ChunkIndexRequest,\n document_index: NewDocumentIndex,\n) -> list[InferenceChunk]:\n return document_index.keyword_retrieval(\n query=query_request.query,\n filters=query_request.filters,\n num_to_retrieve=query_request.limit or NUM_RETURNED_HITS,\n )\n\n\ndef search_chunks(\n query_request: ChunkIndexRequest,\n user_id: UUID | None,\n document_index: DocumentIndex,\n db_session: Session | None = None,\n embedding_model: EmbeddingModel | None = None,\n prefetched_federated_retrieval_infos: list[FederatedRetrievalInfo] | None = None,\n) -> list[InferenceChunk]:\n run_queries: list[tuple[Callable, tuple]] = []\n\n source_filters = (\n set(query_request.filters.source_type)\n if query_request.filters.source_type\n else None\n )\n\n # Federated retrieval — use pre-fetched if available, otherwise query DB\n if prefetched_federated_retrieval_infos is not None:\n federated_retrieval_infos = prefetched_federated_retrieval_infos\n else:\n if db_session is None:\n raise ValueError(\n \"Either db_session or prefetched_federated_retrieval_infos must be provided\"\n )\n federated_retrieval_infos = get_federated_retrieval_functions(\n db_session=db_session,\n user_id=user_id,\n source_types=list(source_filters) if source_filters else None,\n document_set_names=query_request.filters.document_set,\n )\n\n federated_sources = set(\n federated_retrieval_info.source.to_non_federated_source()\n for federated_retrieval_info in federated_retrieval_infos\n )\n for federated_retrieval_info in federated_retrieval_infos:\n run_queries.append(\n (federated_retrieval_info.retrieval_function, (query_request,))\n )\n\n # Don't run normal hybrid search if there are no indexed sources to\n # search over\n normal_search_enabled = (source_filters is None) or (\n len(set(source_filters) - federated_sources) > 0\n )\n\n if normal_search_enabled:\n if (\n query_request.hybrid_alpha is not None\n and query_request.hybrid_alpha == 0.0\n and isinstance(document_index, OpenSearchOldDocumentIndex)\n ):\n # If hybrid alpha is explicitly set to keyword only, do pure keyword\n # search without generating an embedding. This is currently only\n # supported with OpenSearchDocumentIndex.\n opensearch_new_document_index: NewDocumentIndex = document_index._real_index\n run_queries.append(\n (\n lambda: _keyword_search(\n query_request, opensearch_new_document_index\n ),\n (),\n )\n )\n else:\n run_queries.append(\n (\n _embed_and_hybrid_search,\n (query_request, document_index, db_session, embedding_model),\n )\n )\n\n parallel_search_results = run_functions_tuples_in_parallel(run_queries)\n top_chunks = combine_retrieval_results(parallel_search_results)\n\n if not top_chunks:\n logger.debug(\n f\"Search returned no results for query: {query_request.query} with filters: {query_request.filters}.\"\n )\n\n return top_chunks\n\n\n# TODO: This is unused code.\ndef inference_sections_from_ids(\n doc_identifiers: list[tuple[str, int]],\n document_index: DocumentIndex,\n) -> list[InferenceSection]:\n # Currently only fetches whole docs\n doc_ids_set = set(doc_id for doc_id, _ in doc_identifiers)\n\n chunk_requests: list[VespaChunkRequest] = [\n VespaChunkRequest(document_id=doc_id) for doc_id in doc_ids_set\n ]\n\n # No need for ACL here because the doc ids were validated beforehand\n filters = IndexFilters(access_control_list=None)\n\n retrieved_chunks = document_index.id_based_retrieval(\n chunk_requests=chunk_requests,\n filters=filters,\n )\n\n if not retrieved_chunks:\n return []\n\n # Group chunks by document ID\n chunks_by_doc_id: dict[str, list[InferenceChunk]] = {}\n for chunk in retrieved_chunks:\n chunks_by_doc_id.setdefault(chunk.document_id, []).append(chunk)\n\n inference_sections = [\n section\n for chunks in chunks_by_doc_id.values()\n if chunks\n and (\n section := inference_section_from_chunks(\n # The scores will always be 0 because the fetching by id gives back\n # no search scores. This is not needed though if the user is explicitly\n # selecting a document.\n center_chunk=chunks[0],\n chunks=chunks,\n )\n )\n ]\n\n return inference_sections\n" }, { "path": "backend/onyx/context/search/utils.py", "content": "from typing import TypeVar\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceSection\nfrom onyx.context.search.models import SavedSearchDoc\nfrom onyx.context.search.models import SavedSearchDocWithContent\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.configs import MODEL_SERVER_HOST\nfrom shared_configs.configs import MODEL_SERVER_PORT\nfrom shared_configs.enums import EmbedTextType\nfrom shared_configs.model_server_models import Embedding\n\nlogger = setup_logger()\n\n\nT = TypeVar(\n \"T\",\n InferenceSection,\n InferenceChunk,\n SearchDoc,\n SavedSearchDoc,\n SavedSearchDocWithContent,\n)\n\nTSection = TypeVar(\n \"TSection\",\n InferenceSection,\n SearchDoc,\n SavedSearchDoc,\n SavedSearchDocWithContent,\n)\n\n\ndef inference_section_from_chunks(\n center_chunk: InferenceChunk,\n chunks: list[InferenceChunk],\n) -> InferenceSection | None:\n if not chunks:\n return None\n\n combined_content = \"\\n\".join([chunk.content for chunk in chunks])\n\n return InferenceSection(\n center_chunk=center_chunk,\n chunks=chunks,\n combined_content=combined_content,\n )\n\n\n# If it should be a real section, don't use this one\ndef inference_section_from_single_chunk(\n chunk: InferenceChunk,\n) -> InferenceSection:\n return InferenceSection(\n center_chunk=chunk,\n chunks=[chunk],\n combined_content=chunk.content,\n )\n\n\ndef get_query_embeddings(\n queries: list[str],\n db_session: Session | None = None,\n embedding_model: EmbeddingModel | None = None,\n) -> list[Embedding]:\n if embedding_model is None:\n if db_session is None:\n raise ValueError(\"Either db_session or embedding_model must be provided\")\n search_settings = get_current_search_settings(db_session)\n embedding_model = EmbeddingModel.from_db_model(\n search_settings=search_settings,\n server_host=MODEL_SERVER_HOST,\n server_port=MODEL_SERVER_PORT,\n )\n\n query_embedding = embedding_model.encode(queries, text_type=EmbedTextType.QUERY)\n return query_embedding\n\n\n@log_function_time(print_only=True, debug_only=True)\ndef get_query_embedding(\n query: str,\n db_session: Session | None = None,\n embedding_model: EmbeddingModel | None = None,\n) -> Embedding:\n return get_query_embeddings(\n [query], db_session=db_session, embedding_model=embedding_model\n )[0]\n\n\ndef convert_inference_sections_to_search_docs(\n inference_sections: list[InferenceSection],\n is_internet: bool = False,\n) -> list[SearchDoc]:\n search_docs = SearchDoc.from_chunks_or_sections(inference_sections)\n for search_doc in search_docs:\n search_doc.is_internet = is_internet\n return search_docs\n" }, { "path": "backend/onyx/db/README.md", "content": "An explanation of how the history of messages, tool calls, and docs are stored in the database:\n\nMessages are grouped by a chat session, a tree structured is used to allow edits and for the\nuser to switch between branches. Each ChatMessage is either a user message or an assistant message.\nIt should always alternate between the two, System messages, custom agent prompt injections, and\nreminder messages are injected dynamically after the chat session is loaded into memory. The user\nand assistant messages are stored in pairs, though it is ok if the user message is stored and the\nassistant message fails.\n\nThe user chat message is relatively simple and includes the user prompt and any attached documents.\nThe assistant message includes the response, tool calls, feedback, citations, etc.\nThings provided as input are part of the user message, things that happen during the inference and\nLLM loop are part of the assistant message.\n\nReasoning is part of the message or tool call that occured after the reasoning. Really the reasoning\nshould be part of the previous message / tool call because if it branches afterwards as a result of\nthe reasoning, this is somewhat unintuitive. But to not include reasoning as part of the user message,\nit is instead included with the following message or tool call. With parallel tool calls, the reasoning\nwill be included with each of the tool calls.\n\nTool calls are stored in the ToolCall table and can represent all of the following:\n- Parallel tool calls, these will have the same turn number and parent tool call id\n- Sequential tool calls, these will have a different turn number and parent tool call id\n- Tool calls attached to the ChatMessage are top level tool calls directly triggered by the LLM\n- Tool calls that are instead attached to other ToolCalls are tool calls that happen as part of an\n agent that has been called. The top level tool call is the agent call and the tool calls that have\n the agent call as a parent are the tool calls that happen as part of the agent.\n\nThe different branches are generated by sending a new search query to an existing parent.\n```\n [Empty Root Message] (This allows the first message to be branched/edited as well)\n / | \\\n[First Message] [First Message Edit 1] [First Message Edit 2]\n | |\n[Second Message] [Second Message of Edit 1 Branch]\n```\n" }, { "path": "backend/onyx/db/__init__.py", "content": "" }, { "path": "backend/onyx/db/_deprecated/pg_file_store.py", "content": "\"\"\"Kept around since it's used in the migration to move to S3/MinIO\"\"\"\n\nimport tempfile\nfrom io import BytesIO\nfrom typing import IO\n\nfrom psycopg2.extensions import connection\nfrom sqlalchemy import text # NEW: for SQL large-object helpers\nfrom sqlalchemy.orm import Session\n\nfrom onyx.file_store.constants import MAX_IN_MEMORY_SIZE\nfrom onyx.file_store.constants import STANDARD_CHUNK_SIZE\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef get_pg_conn_from_session(db_session: Session) -> connection:\n return db_session.connection().connection.connection # type: ignore\n\n\ndef create_populate_lobj(\n content: IO,\n db_session: Session,\n) -> int:\n \"\"\"Create a PostgreSQL large object from *content* and return its OID.\n\n Preferred approach is to use the psycopg2 ``lobject`` API, but if that is\n unavailable (e.g. when the underlying connection is an asyncpg adapter)\n we fall back to PostgreSQL helper functions such as ``lo_from_bytea``.\n\n NOTE: this function intentionally *does not* commit the surrounding\n transaction – that is handled by the caller so all work stays atomic.\n \"\"\"\n\n pg_conn = None\n try:\n pg_conn = get_pg_conn_from_session(db_session)\n # ``AsyncAdapt_asyncpg_connection`` (asyncpg) has no ``lobject``\n if not hasattr(pg_conn, \"lobject\"):\n raise AttributeError # will be handled by fallback below\n\n large_object = pg_conn.lobject()\n\n # write in multiple chunks to avoid loading the whole file into memory\n while True:\n chunk = content.read(STANDARD_CHUNK_SIZE)\n if not chunk:\n break\n large_object.write(chunk)\n\n large_object.close()\n\n return large_object.oid\n\n except AttributeError:\n # Fall back to SQL helper functions – read the full content into memory\n # (acceptable for the limited number and size of files handled during\n # migrations). ``lo_from_bytea`` returns the new OID.\n byte_data = content.read()\n result = db_session.execute(\n text(\"SELECT lo_from_bytea(0, :data) AS oid\"),\n {\"data\": byte_data},\n )\n # ``scalar_one`` is 2.0-style; ``scalar`` works on both 1.4/2.0.\n lobj_oid = result.scalar()\n if lobj_oid is None:\n raise RuntimeError(\"Failed to create large object\")\n return int(lobj_oid)\n\n\ndef read_lobj(\n lobj_oid: int,\n db_session: Session,\n mode: str | None = None,\n use_tempfile: bool = False,\n) -> IO:\n \"\"\"Read a PostgreSQL large object identified by *lobj_oid*.\n\n Attempts to use the native ``lobject`` API first; if unavailable falls back\n to ``lo_get`` which returns the large object's contents as *bytea*.\n \"\"\"\n\n pg_conn = None\n try:\n pg_conn = get_pg_conn_from_session(db_session)\n if not hasattr(pg_conn, \"lobject\"):\n raise AttributeError\n\n # Ensure binary mode by default\n if mode is None:\n mode = \"rb\"\n large_object = (\n pg_conn.lobject(lobj_oid, mode=mode) if mode else pg_conn.lobject(lobj_oid)\n )\n\n if use_tempfile:\n temp_file = tempfile.SpooledTemporaryFile(max_size=MAX_IN_MEMORY_SIZE)\n while True:\n chunk = large_object.read(STANDARD_CHUNK_SIZE)\n if not chunk:\n break\n temp_file.write(chunk)\n temp_file.seek(0)\n return temp_file\n else:\n return BytesIO(large_object.read())\n\n except AttributeError:\n # Fallback path using ``lo_get``\n result = db_session.execute(\n text(\"SELECT lo_get(:oid) AS data\"),\n {\"oid\": lobj_oid},\n )\n byte_data = result.scalar()\n if byte_data is None:\n raise RuntimeError(\"Failed to read large object\")\n\n if use_tempfile:\n temp_file = tempfile.SpooledTemporaryFile(max_size=MAX_IN_MEMORY_SIZE)\n temp_file.write(byte_data)\n temp_file.seek(0)\n return temp_file\n return BytesIO(byte_data)\n\n\ndef delete_lobj_by_id(\n lobj_oid: int,\n db_session: Session,\n) -> None:\n \"\"\"Remove a large object by OID, regardless of driver implementation.\"\"\"\n\n try:\n pg_conn = get_pg_conn_from_session(db_session)\n if hasattr(pg_conn, \"lobject\"):\n pg_conn.lobject(lobj_oid).unlink()\n return\n raise AttributeError\n except AttributeError:\n # Fallback for drivers without ``lobject`` support\n db_session.execute(text(\"SELECT lo_unlink(:oid)\"), {\"oid\": lobj_oid})\n # No explicit result expected\n" }, { "path": "backend/onyx/db/api_key.py", "content": "import uuid\n\nfrom fastapi_users.password import PasswordHelper\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.api_key import ApiKeyDescriptor\nfrom onyx.auth.api_key import build_displayable_api_key\nfrom onyx.auth.api_key import generate_api_key\nfrom onyx.auth.api_key import hash_api_key\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN\nfrom onyx.configs.constants import DANSWER_API_KEY_PREFIX\nfrom onyx.configs.constants import UNNAMED_KEY_PLACEHOLDER\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import ApiKey\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.db.permissions import recompute_user_permissions__no_commit\nfrom onyx.db.users import assign_user_to_default_groups__no_commit\nfrom onyx.server.api_key.models import APIKeyArgs\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\ndef get_api_key_email_pattern() -> str:\n return DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN\n\n\ndef is_api_key_email_address(email: str) -> bool:\n return email.endswith(get_api_key_email_pattern())\n\n\ndef fetch_api_keys(db_session: Session) -> list[ApiKeyDescriptor]:\n api_keys = (\n db_session.scalars(select(ApiKey).options(joinedload(ApiKey.user)))\n .unique()\n .all()\n )\n return [\n ApiKeyDescriptor(\n api_key_id=api_key.id,\n api_key_role=api_key.user.role,\n api_key_display=api_key.api_key_display,\n api_key_name=api_key.name,\n user_id=api_key.user_id,\n )\n for api_key in api_keys\n ]\n\n\nasync def fetch_user_for_api_key(\n hashed_api_key: str, async_db_session: AsyncSession\n) -> User | None:\n \"\"\"NOTE: this is async, since it's used during auth\n (which is necessarily async due to FastAPI Users)\"\"\"\n return await async_db_session.scalar(\n select(User)\n .join(ApiKey, ApiKey.user_id == User.id)\n .where(ApiKey.hashed_api_key == hashed_api_key)\n )\n\n\ndef get_api_key_fake_email(\n name: str,\n unique_id: str,\n) -> str:\n return f\"{DANSWER_API_KEY_PREFIX}{name}@{unique_id}{DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN}\"\n\n\ndef insert_api_key(\n db_session: Session, api_key_args: APIKeyArgs, user_id: uuid.UUID | None\n) -> ApiKeyDescriptor:\n std_password_helper = PasswordHelper()\n\n # Get tenant_id from context var (will be default schema for single tenant)\n tenant_id = get_current_tenant_id()\n\n api_key = generate_api_key(tenant_id)\n api_key_user_id = uuid.uuid4()\n\n display_name = api_key_args.name or UNNAMED_KEY_PLACEHOLDER\n api_key_user_row = User(\n id=api_key_user_id,\n email=get_api_key_fake_email(display_name, str(api_key_user_id)),\n # a random password for the \"user\"\n hashed_password=std_password_helper.hash(std_password_helper.generate()),\n is_active=True,\n is_superuser=False,\n is_verified=True,\n role=api_key_args.role,\n account_type=AccountType.SERVICE_ACCOUNT,\n )\n db_session.add(api_key_user_row)\n\n api_key_row = ApiKey(\n name=api_key_args.name,\n hashed_api_key=hash_api_key(api_key),\n api_key_display=build_displayable_api_key(api_key),\n user_id=api_key_user_id,\n owner_id=user_id,\n )\n db_session.add(api_key_row)\n\n # Assign the API key virtual user to the appropriate default group\n # before commit so everything is atomic.\n # LIMITED role service accounts should have no group membership.\n if api_key_args.role != UserRole.LIMITED:\n assign_user_to_default_groups__no_commit(\n db_session,\n api_key_user_row,\n is_admin=(api_key_args.role == UserRole.ADMIN),\n )\n\n db_session.commit()\n\n return ApiKeyDescriptor(\n api_key_id=api_key_row.id,\n api_key_role=api_key_user_row.role,\n api_key_display=api_key_row.api_key_display,\n api_key=api_key,\n api_key_name=api_key_args.name,\n user_id=api_key_user_id,\n )\n\n\ndef update_api_key(\n db_session: Session, api_key_id: int, api_key_args: APIKeyArgs\n) -> ApiKeyDescriptor:\n existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id))\n if existing_api_key is None:\n raise ValueError(f\"API key with id {api_key_id} does not exist\")\n\n existing_api_key.name = api_key_args.name\n api_key_user = db_session.scalar(\n select(User).where(User.id == existing_api_key.user_id) # type: ignore\n )\n if api_key_user is None:\n raise RuntimeError(\"API Key does not have associated user.\")\n\n email_name = api_key_args.name or UNNAMED_KEY_PLACEHOLDER\n api_key_user.email = get_api_key_fake_email(email_name, str(api_key_user.id))\n\n old_role = api_key_user.role\n api_key_user.role = api_key_args.role\n\n # Reconcile default-group membership when the role changes.\n if old_role != api_key_args.role:\n # Remove from all default groups first.\n delete_stmt = delete(User__UserGroup).where(\n User__UserGroup.user_id == api_key_user.id,\n User__UserGroup.user_group_id.in_(\n select(UserGroup.id).where(UserGroup.is_default.is_(True))\n ),\n )\n db_session.execute(delete_stmt)\n\n # Re-assign to the correct default group (skip for LIMITED).\n if api_key_args.role != UserRole.LIMITED:\n assign_user_to_default_groups__no_commit(\n db_session,\n api_key_user,\n is_admin=(api_key_args.role == UserRole.ADMIN),\n )\n else:\n # No group assigned for LIMITED, but we still need to recompute\n # since we just removed the old default-group membership above.\n recompute_user_permissions__no_commit(api_key_user.id, db_session)\n\n db_session.commit()\n\n return ApiKeyDescriptor(\n api_key_id=existing_api_key.id,\n api_key_display=existing_api_key.api_key_display,\n api_key_name=api_key_args.name,\n api_key_role=api_key_user.role,\n user_id=existing_api_key.user_id,\n )\n\n\ndef regenerate_api_key(db_session: Session, api_key_id: int) -> ApiKeyDescriptor:\n \"\"\"NOTE: currently, any admin can regenerate any API key.\"\"\"\n existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id))\n if existing_api_key is None:\n raise ValueError(f\"API key with id {api_key_id} does not exist\")\n\n api_key_user = db_session.scalar(\n select(User).where(User.id == existing_api_key.user_id) # type: ignore\n )\n if api_key_user is None:\n raise RuntimeError(\"API Key does not have associated user.\")\n\n # Get tenant_id from context var (will be default schema for single tenant)\n tenant_id = get_current_tenant_id()\n\n new_api_key = generate_api_key(tenant_id)\n existing_api_key.hashed_api_key = hash_api_key(new_api_key)\n existing_api_key.api_key_display = build_displayable_api_key(new_api_key)\n db_session.commit()\n\n return ApiKeyDescriptor(\n api_key_id=existing_api_key.id,\n api_key_display=existing_api_key.api_key_display,\n api_key=new_api_key,\n api_key_name=existing_api_key.name,\n api_key_role=api_key_user.role,\n user_id=existing_api_key.user_id,\n )\n\n\ndef remove_api_key(db_session: Session, api_key_id: int) -> None:\n existing_api_key = db_session.scalar(select(ApiKey).where(ApiKey.id == api_key_id))\n if existing_api_key is None:\n raise ValueError(f\"API key with id {api_key_id} does not exist\")\n\n user_associated_with_key = db_session.scalar(\n select(User).where(User.id == existing_api_key.user_id) # type: ignore\n )\n if user_associated_with_key is None:\n raise ValueError(\n f\"User associated with API key with id {api_key_id} does not exist. This should not happen.\"\n )\n\n db_session.delete(existing_api_key)\n db_session.delete(user_associated_with_key)\n db_session.commit()\n" }, { "path": "backend/onyx/db/auth.py", "content": "from collections.abc import AsyncGenerator\nfrom collections.abc import Callable\nfrom typing import Any\nfrom typing import Dict\nfrom typing import TypeVar\n\nfrom fastapi import Depends\nfrom fastapi_users.models import ID\nfrom fastapi_users.models import UP\nfrom fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase\nfrom fastapi_users_db_sqlalchemy.access_token import SQLAlchemyAccessTokenDatabase\nfrom sqlalchemy import func\nfrom sqlalchemy import Select\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.future import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import ANONYMOUS_USER_EMAIL\nfrom onyx.configs.constants import NO_AUTH_PLACEHOLDER_USER_EMAIL\nfrom onyx.db.api_key import get_api_key_email_pattern\nfrom onyx.db.engine.async_sql_engine import get_async_session\nfrom onyx.db.engine.async_sql_engine import get_async_session_context_manager\nfrom onyx.db.models import AccessToken\nfrom onyx.db.models import OAuthAccount\nfrom onyx.db.models import User\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\n\nT = TypeVar(\"T\", bound=tuple[Any, ...])\n\n\ndef get_default_admin_user_emails() -> list[str]:\n \"\"\"Returns a list of emails who should default to Admin role.\n Only used in the EE version. For MIT, just return empty list.\"\"\"\n get_default_admin_user_emails_fn: Callable[[], list[str]] = (\n fetch_versioned_implementation_with_fallback(\n \"onyx.auth.users\", \"get_default_admin_user_emails_\", lambda: list[str]()\n )\n )\n return get_default_admin_user_emails_fn()\n\n\ndef _add_live_user_count_where_clause(\n select_stmt: Select[T],\n only_admin_users: bool,\n) -> Select[T]:\n \"\"\"\n Builds a SQL column expression that can be used to filter out\n users who should not be included in the live user count.\n\n Excludes:\n - API key users (by email pattern)\n - System users (anonymous user, no-auth placeholder)\n - External permission users (unless only_admin_users is True)\n \"\"\"\n select_stmt = select_stmt.where(~User.email.endswith(get_api_key_email_pattern())) # type: ignore\n\n # Exclude system users (anonymous user, no-auth placeholder)\n select_stmt = select_stmt.where(User.email != ANONYMOUS_USER_EMAIL) # type: ignore\n select_stmt = select_stmt.where(User.email != NO_AUTH_PLACEHOLDER_USER_EMAIL) # type: ignore\n\n if only_admin_users:\n return select_stmt.where(User.role == UserRole.ADMIN)\n\n return select_stmt.where(\n User.role != UserRole.EXT_PERM_USER,\n )\n\n\ndef get_live_users_count(db_session: Session) -> int:\n \"\"\"\n Returns the number of users in the system.\n This does NOT include invited users, \"users\" pulled in\n from external connectors, or API keys.\n \"\"\"\n count_stmt = func.count(User.id)\n select_stmt = select(count_stmt)\n select_stmt_w_filters = _add_live_user_count_where_clause(select_stmt, False)\n user_count = db_session.scalar(select_stmt_w_filters)\n if user_count is None:\n raise RuntimeError(\"Was not able to fetch the user count.\")\n return user_count\n\n\nasync def get_user_count(only_admin_users: bool = False) -> int:\n async with get_async_session_context_manager() as session:\n count_stmt = func.count(User.id)\n stmt = select(count_stmt)\n stmt_w_filters = _add_live_user_count_where_clause(stmt, only_admin_users)\n user_count = await session.scalar(stmt_w_filters)\n if user_count is None:\n raise RuntimeError(\"Was not able to fetch the user count.\")\n return user_count\n\n\n# Need to override this because FastAPI Users doesn't give flexibility for backend field creation logic in OAuth flow\nclass SQLAlchemyUserAdminDB(SQLAlchemyUserDatabase[UP, ID]):\n async def create(\n self,\n create_dict: Dict[str, Any],\n ) -> UP:\n user_count = await get_user_count()\n if user_count == 0 or create_dict[\"email\"] in get_default_admin_user_emails():\n create_dict[\"role\"] = UserRole.ADMIN\n else:\n create_dict[\"role\"] = UserRole.BASIC\n return await super().create(create_dict)\n\n\nasync def get_user_db(\n session: AsyncSession = Depends(get_async_session),\n) -> AsyncGenerator[SQLAlchemyUserAdminDB, None]:\n yield SQLAlchemyUserAdminDB(session, User, OAuthAccount)\n\n\nasync def get_access_token_db(\n session: AsyncSession = Depends(get_async_session),\n) -> AsyncGenerator[SQLAlchemyAccessTokenDatabase, None]:\n yield SQLAlchemyAccessTokenDatabase(session, AccessToken)\n" }, { "path": "backend/onyx/db/background_error.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.db.models import BackgroundError\n\n\ndef create_background_error(\n db_session: Session, message: str, cc_pair_id: int | None\n) -> None:\n db_session.add(BackgroundError(message=message, cc_pair_id=cc_pair_id))\n db_session.commit()\n" }, { "path": "backend/onyx/db/chat.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Tuple\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import delete\nfrom sqlalchemy import desc\nfrom sqlalchemy import func\nfrom sqlalchemy import nullsfirst\nfrom sqlalchemy import or_\nfrom sqlalchemy import Row\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.exc import MultipleResultsFound\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.chat_configs import HARD_DELETE_CHATS\nfrom onyx.configs.constants import MessageType\nfrom onyx.context.search.models import InferenceSection\nfrom onyx.context.search.models import SavedSearchDoc\nfrom onyx.context.search.models import SearchDoc as ServerSearchDoc\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatMessage__SearchDoc\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import ChatSessionSharedStatus\nfrom onyx.db.models import Persona\nfrom onyx.db.models import SearchDoc as DBSearchDoc\nfrom onyx.db.models import ToolCall\nfrom onyx.db.models import User\nfrom onyx.db.persona import get_best_persona_id_for_user\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.llm.override_models import PromptOverride\nfrom onyx.server.query_and_chat.models import ChatMessageDetail\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.postgres_sanitization import sanitize_string\n\n\nlogger = setup_logger()\n\n\n# Note: search/streaming packet helpers moved to streaming_utils.py\n\n\ndef get_chat_session_by_id(\n chat_session_id: UUID,\n user_id: UUID | None,\n db_session: Session,\n include_deleted: bool = False,\n is_shared: bool = False,\n eager_load_persona: bool = False,\n) -> ChatSession:\n stmt = select(ChatSession).where(ChatSession.id == chat_session_id)\n\n if eager_load_persona:\n stmt = stmt.options(\n joinedload(ChatSession.persona).options(\n selectinload(Persona.tools),\n selectinload(Persona.user_files),\n selectinload(Persona.document_sets),\n selectinload(Persona.attached_documents),\n selectinload(Persona.hierarchy_nodes),\n ),\n joinedload(ChatSession.project),\n )\n\n if is_shared:\n stmt = stmt.where(ChatSession.shared_status == ChatSessionSharedStatus.PUBLIC)\n else:\n # if user_id is None, assume this is an admin who should be able\n # to view all chat sessions\n if user_id is not None:\n stmt = stmt.where(\n or_(ChatSession.user_id == user_id, ChatSession.user_id.is_(None))\n )\n\n result = db_session.execute(stmt)\n chat_session = result.scalar_one_or_none()\n\n if not chat_session:\n raise ValueError(\"Invalid Chat Session ID provided\")\n\n if not include_deleted and chat_session.deleted:\n raise ValueError(\"Chat session has been deleted\")\n\n return chat_session\n\n\ndef get_chat_sessions_by_slack_thread_id(\n slack_thread_id: str,\n user_id: UUID | None,\n db_session: Session,\n) -> Sequence[ChatSession]:\n stmt = select(ChatSession).where(ChatSession.slack_thread_id == slack_thread_id)\n if user_id is not None:\n stmt = stmt.where(\n or_(ChatSession.user_id == user_id, ChatSession.user_id.is_(None))\n )\n return db_session.scalars(stmt).all()\n\n\n# Retrieves chat sessions by user\n# Chat sessions do not include onyxbot flows\ndef get_chat_sessions_by_user(\n user_id: UUID | None,\n deleted: bool | None,\n db_session: Session,\n include_onyxbot_flows: bool = False,\n limit: int = 50,\n before: datetime | None = None,\n project_id: int | None = None,\n only_non_project_chats: bool = False,\n include_failed_chats: bool = False,\n) -> list[ChatSession]:\n stmt = select(ChatSession).where(ChatSession.user_id == user_id)\n\n if not include_onyxbot_flows:\n stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))\n\n stmt = stmt.order_by(desc(ChatSession.time_updated))\n\n if deleted is not None:\n stmt = stmt.where(ChatSession.deleted == deleted)\n\n if before is not None:\n stmt = stmt.where(ChatSession.time_updated < before)\n\n if project_id is not None:\n stmt = stmt.where(ChatSession.project_id == project_id)\n elif only_non_project_chats:\n stmt = stmt.where(ChatSession.project_id.is_(None))\n\n # When filtering out failed chats, we apply the limit in Python after\n # filtering rather than in SQL, since the post-filter may remove rows.\n if limit and include_failed_chats:\n stmt = stmt.limit(limit)\n\n result = db_session.execute(stmt)\n chat_sessions = list(result.scalars().all())\n\n if not include_failed_chats and chat_sessions:\n # Filter out \"failed\" sessions (those with only SYSTEM messages)\n # using a separate efficient query instead of a correlated EXISTS\n # subquery, which causes full sequential scans of chat_message.\n leeway = datetime.now(timezone.utc) - timedelta(minutes=5)\n session_ids = [cs.id for cs in chat_sessions if cs.time_created < leeway]\n\n if session_ids:\n valid_session_ids_stmt = (\n select(ChatMessage.chat_session_id)\n .where(ChatMessage.chat_session_id.in_(session_ids))\n .where(ChatMessage.message_type != MessageType.SYSTEM)\n .distinct()\n )\n valid_session_ids = set(\n db_session.execute(valid_session_ids_stmt).scalars().all()\n )\n\n chat_sessions = [\n cs\n for cs in chat_sessions\n if cs.time_created >= leeway or cs.id in valid_session_ids\n ]\n\n if limit:\n chat_sessions = chat_sessions[:limit]\n\n return chat_sessions\n\n\ndef delete_orphaned_search_docs(db_session: Session) -> None:\n orphaned_docs = (\n db_session.query(DBSearchDoc)\n .outerjoin(ChatMessage__SearchDoc)\n .filter(ChatMessage__SearchDoc.chat_message_id.is_(None))\n .all()\n )\n for doc in orphaned_docs:\n db_session.delete(doc)\n db_session.commit()\n\n\ndef delete_messages_and_files_from_chat_session(\n chat_session_id: UUID, db_session: Session\n) -> None:\n # Select messages older than cutoff_time with files\n messages_with_files = db_session.execute(\n select(ChatMessage.id, ChatMessage.files).where(\n ChatMessage.chat_session_id == chat_session_id,\n )\n ).fetchall()\n\n for _, files in messages_with_files:\n file_store = get_default_file_store()\n for file_info in files or []:\n file_store.delete_file(file_id=file_info.get(\"id\"))\n\n # Delete ChatMessage records - CASCADE constraints will automatically handle:\n # - ChatMessage__StandardAnswer relationship records\n db_session.execute(\n delete(ChatMessage).where(ChatMessage.chat_session_id == chat_session_id)\n )\n db_session.commit()\n\n delete_orphaned_search_docs(db_session)\n\n\ndef create_chat_session(\n db_session: Session,\n description: str | None,\n user_id: UUID | None,\n persona_id: int | None, # Can be none if temporary persona is used\n llm_override: LLMOverride | None = None,\n prompt_override: PromptOverride | None = None,\n onyxbot_flow: bool = False,\n slack_thread_id: str | None = None,\n project_id: int | None = None,\n) -> ChatSession:\n chat_session = ChatSession(\n user_id=user_id,\n persona_id=persona_id,\n description=description,\n llm_override=llm_override,\n prompt_override=prompt_override,\n onyxbot_flow=onyxbot_flow,\n slack_thread_id=slack_thread_id,\n project_id=project_id,\n )\n\n db_session.add(chat_session)\n db_session.commit()\n\n return chat_session\n\n\ndef duplicate_chat_session_for_user_from_slack(\n db_session: Session,\n user: User,\n chat_session_id: UUID,\n) -> ChatSession:\n \"\"\"\n This takes a chat session id for a session in Slack and:\n - Creates a new chat session in the DB\n - Tries to copy the persona from the original chat session\n (if it is available to the user clicking the button)\n - Sets the user to the given user (if provided)\n \"\"\"\n chat_session = get_chat_session_by_id(\n chat_session_id=chat_session_id,\n user_id=None, # Ignore user permissions for this\n db_session=db_session,\n )\n if not chat_session:\n raise HTTPException(status_code=400, detail=\"Invalid Chat Session ID provided\")\n\n # This enforces permissions and sets a default\n new_persona_id = get_best_persona_id_for_user(\n db_session=db_session,\n user=user,\n persona_id=chat_session.persona_id,\n )\n\n return create_chat_session(\n db_session=db_session,\n user_id=user.id,\n persona_id=new_persona_id,\n # Set this to empty string so the frontend will force a rename\n description=\"\",\n llm_override=chat_session.llm_override,\n prompt_override=chat_session.prompt_override,\n # Chat is in UI now so this is false\n onyxbot_flow=False,\n # Maybe we want this in the future to track if it was created from Slack\n slack_thread_id=None,\n )\n\n\ndef update_chat_session(\n db_session: Session,\n user_id: UUID | None,\n chat_session_id: UUID,\n description: str | None = None,\n sharing_status: ChatSessionSharedStatus | None = None,\n) -> ChatSession:\n chat_session = get_chat_session_by_id(\n chat_session_id=chat_session_id, user_id=user_id, db_session=db_session\n )\n\n if chat_session.deleted:\n raise ValueError(\"Trying to rename a deleted chat session\")\n\n if description is not None:\n chat_session.description = description\n if sharing_status is not None:\n chat_session.shared_status = sharing_status\n\n db_session.commit()\n\n return chat_session\n\n\ndef delete_all_chat_sessions_for_user(\n user: User, db_session: Session, hard_delete: bool = HARD_DELETE_CHATS\n) -> None:\n user_id = user.id\n\n chat_sessions = (\n db_session.query(ChatSession)\n .filter(ChatSession.user_id == user_id, ChatSession.onyxbot_flow.is_(False))\n .all()\n )\n\n if hard_delete:\n for chat_session in chat_sessions:\n delete_messages_and_files_from_chat_session(chat_session.id, db_session)\n db_session.execute(\n delete(ChatSession).where(\n ChatSession.user_id == user_id, ChatSession.onyxbot_flow.is_(False)\n )\n )\n else:\n db_session.execute(\n update(ChatSession)\n .where(ChatSession.user_id == user_id, ChatSession.onyxbot_flow.is_(False))\n .values(deleted=True)\n )\n\n db_session.commit()\n\n\ndef delete_chat_session(\n user_id: UUID | None,\n chat_session_id: UUID,\n db_session: Session,\n include_deleted: bool = False,\n hard_delete: bool = HARD_DELETE_CHATS,\n) -> None:\n chat_session = get_chat_session_by_id(\n chat_session_id=chat_session_id,\n user_id=user_id,\n db_session=db_session,\n include_deleted=include_deleted,\n )\n\n if chat_session.deleted and not include_deleted:\n raise ValueError(\"Cannot delete an already deleted chat session\")\n\n if hard_delete:\n delete_messages_and_files_from_chat_session(chat_session_id, db_session)\n db_session.execute(delete(ChatSession).where(ChatSession.id == chat_session_id))\n else:\n chat_session = get_chat_session_by_id(\n chat_session_id=chat_session_id, user_id=user_id, db_session=db_session\n )\n chat_session.deleted = True\n\n db_session.commit()\n\n\ndef get_chat_sessions_older_than(\n days_old: int, db_session: Session\n) -> list[tuple[UUID | None, UUID]]:\n \"\"\"\n Retrieves chat sessions older than a specified number of days.\n\n Args:\n days_old: The number of days to consider as \"old\".\n db_session: The database session.\n\n Returns:\n A list of tuples, where each tuple contains the user_id (can be None) and the chat_session_id of an old chat session.\n \"\"\"\n\n cutoff_time = datetime.utcnow() - timedelta(days=days_old)\n old_sessions: Sequence[Row[Tuple[UUID | None, UUID]]] = db_session.execute(\n select(ChatSession.user_id, ChatSession.id).where(\n ChatSession.time_created < cutoff_time\n )\n ).fetchall()\n\n # convert old_sessions to a conventional list of tuples\n returned_sessions: list[tuple[UUID | None, UUID]] = [\n (user_id, session_id) for user_id, session_id in old_sessions\n ]\n\n return returned_sessions\n\n\ndef get_chat_message(\n chat_message_id: int,\n user_id: UUID | None,\n db_session: Session,\n) -> ChatMessage:\n stmt = select(ChatMessage).where(ChatMessage.id == chat_message_id)\n\n result = db_session.execute(stmt)\n chat_message = result.scalar_one_or_none()\n\n if not chat_message:\n raise ValueError(\"Invalid Chat Message specified\")\n\n chat_user = chat_message.chat_session.user\n expected_user_id = chat_user.id if chat_user is not None else None\n\n if expected_user_id != user_id:\n logger.error(\n f\"User {user_id} tried to fetch a chat message that does not belong to them\"\n )\n raise ValueError(\"Chat message does not belong to user\")\n\n return chat_message\n\n\ndef get_chat_session_by_message_id(\n db_session: Session,\n message_id: int,\n) -> ChatSession:\n \"\"\"\n Should only be used for Slack\n Get the chat session associated with a specific message ID\n Note: this ignores permission checks.\n \"\"\"\n stmt = select(ChatMessage).where(ChatMessage.id == message_id)\n\n result = db_session.execute(stmt)\n chat_message = result.scalar_one_or_none()\n\n if chat_message is None:\n raise ValueError(\n f\"Unable to find chat session associated with message ID: {message_id}\"\n )\n\n return chat_message.chat_session\n\n\ndef get_chat_messages_by_sessions(\n chat_session_ids: list[UUID],\n user_id: UUID | None,\n db_session: Session,\n skip_permission_check: bool = False,\n) -> Sequence[ChatMessage]:\n if not skip_permission_check:\n for chat_session_id in chat_session_ids:\n get_chat_session_by_id(\n chat_session_id=chat_session_id, user_id=user_id, db_session=db_session\n )\n stmt = (\n select(ChatMessage)\n .where(ChatMessage.chat_session_id.in_(chat_session_ids))\n .order_by(nullsfirst(ChatMessage.parent_message_id))\n )\n return db_session.execute(stmt).scalars().all()\n\n\ndef add_chats_to_session_from_slack_thread(\n db_session: Session,\n slack_chat_session_id: UUID,\n new_chat_session_id: UUID,\n) -> None:\n new_root_message = get_or_create_root_message(\n chat_session_id=new_chat_session_id,\n db_session=db_session,\n )\n\n for chat_message in get_chat_messages_by_sessions(\n chat_session_ids=[slack_chat_session_id],\n user_id=None, # Ignore user permissions for this\n db_session=db_session,\n skip_permission_check=True,\n ):\n if chat_message.message_type == MessageType.SYSTEM:\n continue\n # Duplicate the message\n new_root_message = create_new_chat_message(\n db_session=db_session,\n chat_session_id=new_chat_session_id,\n parent_message=new_root_message,\n message=chat_message.message,\n files=chat_message.files,\n error=chat_message.error,\n token_count=chat_message.token_count,\n message_type=chat_message.message_type,\n reasoning_tokens=chat_message.reasoning_tokens,\n )\n\n\ndef add_search_docs_to_chat_message(\n chat_message_id: int, search_doc_ids: list[int], db_session: Session\n) -> None:\n \"\"\"\n Link SearchDocs to a ChatMessage by creating entries in the chat_message__search_doc junction table.\n\n Args:\n chat_message_id: The ID of the chat message\n search_doc_ids: List of search document IDs to link\n db_session: The database session\n \"\"\"\n for search_doc_id in search_doc_ids:\n chat_message_search_doc = ChatMessage__SearchDoc(\n chat_message_id=chat_message_id, search_doc_id=search_doc_id\n )\n db_session.add(chat_message_search_doc)\n\n\ndef add_search_docs_to_tool_call(\n tool_call_id: int, search_doc_ids: list[int], db_session: Session\n) -> None:\n \"\"\"\n Link SearchDocs to a ToolCall by creating entries in the tool_call__search_doc junction table.\n\n Args:\n tool_call_id: The ID of the tool call\n search_doc_ids: List of search document IDs to link\n db_session: The database session\n \"\"\"\n from onyx.db.models import ToolCall__SearchDoc\n\n for search_doc_id in search_doc_ids:\n tool_call_search_doc = ToolCall__SearchDoc(\n tool_call_id=tool_call_id, search_doc_id=search_doc_id\n )\n db_session.add(tool_call_search_doc)\n\n\ndef get_chat_messages_by_session(\n chat_session_id: UUID,\n user_id: UUID | None,\n db_session: Session,\n skip_permission_check: bool = False,\n prefetch_top_two_level_tool_calls: bool = True,\n) -> list[ChatMessage]:\n if not skip_permission_check:\n # bug if we ever call this expecting the permission check to not be skipped\n get_chat_session_by_id(\n chat_session_id=chat_session_id, user_id=user_id, db_session=db_session\n )\n\n stmt = (\n select(ChatMessage)\n .where(ChatMessage.chat_session_id == chat_session_id)\n .order_by(nullsfirst(ChatMessage.parent_message_id))\n )\n\n # This should handle both the top level tool calls and deep research\n # If there are future nested agents, this can be extended.\n if prefetch_top_two_level_tool_calls:\n # Load tool_calls and their direct children (one level deep)\n stmt = stmt.options(\n selectinload(ChatMessage.tool_calls).selectinload(\n ToolCall.tool_call_children\n )\n )\n result = db_session.scalars(stmt).unique().all()\n else:\n result = db_session.scalars(stmt).all()\n\n return list(result)\n\n\ndef get_or_create_root_message(\n chat_session_id: UUID,\n db_session: Session,\n) -> ChatMessage:\n try:\n root_message: ChatMessage | None = (\n db_session.query(ChatMessage)\n .filter(\n ChatMessage.chat_session_id == chat_session_id,\n ChatMessage.parent_message_id.is_(None),\n )\n .one_or_none()\n )\n except MultipleResultsFound:\n raise Exception(\n \"Multiple root messages found for chat session. Data inconsistency detected.\"\n )\n\n if root_message is not None:\n return root_message\n else:\n new_root_message = ChatMessage(\n chat_session_id=chat_session_id,\n parent_message_id=None,\n latest_child_message_id=None,\n message=\"\",\n token_count=0,\n message_type=MessageType.SYSTEM,\n )\n db_session.add(new_root_message)\n db_session.commit()\n return new_root_message\n\n\ndef reserve_message_id(\n db_session: Session,\n chat_session_id: UUID,\n parent_message: int,\n message_type: MessageType = MessageType.ASSISTANT,\n) -> ChatMessage:\n # Create an temporary holding chat message to the updated and saved at the end\n empty_message = ChatMessage(\n chat_session_id=chat_session_id,\n parent_message_id=parent_message,\n latest_child_message_id=None,\n message=\"Response was terminated prior to completion, try regenerating.\",\n token_count=15,\n message_type=message_type,\n )\n\n # Add the empty message to the session\n db_session.add(empty_message)\n db_session.flush()\n\n # Get the parent message and set its child pointer to the current message\n parent_chat_message = (\n db_session.query(ChatMessage).filter(ChatMessage.id == parent_message).first()\n )\n if parent_chat_message:\n parent_chat_message.latest_child_message_id = empty_message.id\n\n # Committing because it's ok to recover this state. More clear to the user than it is now.\n # Ideally there's a special UI for a case like this with a regenerate button but not needed for now.\n db_session.commit()\n\n return empty_message\n\n\ndef reserve_multi_model_message_ids(\n db_session: Session,\n chat_session_id: UUID,\n parent_message_id: int,\n model_display_names: list[str],\n) -> list[ChatMessage]:\n \"\"\"Reserve N assistant message placeholders for multi-model parallel streaming.\n\n All messages share the same parent (the user message). The parent's\n latest_child_message_id points to the LAST reserved message so that the\n default history-chain walker picks it up.\n \"\"\"\n reserved: list[ChatMessage] = []\n for display_name in model_display_names:\n msg = ChatMessage(\n chat_session_id=chat_session_id,\n parent_message_id=parent_message_id,\n latest_child_message_id=None,\n message=\"Response was terminated prior to completion, try regenerating.\",\n token_count=15, # placeholder; updated on completion by llm_loop_completion_handle\n message_type=MessageType.ASSISTANT,\n model_display_name=display_name,\n )\n db_session.add(msg)\n reserved.append(msg)\n\n # Flush to assign IDs without committing yet\n db_session.flush()\n\n # Point parent's latest_child to the last reserved message\n parent = (\n db_session.query(ChatMessage)\n .filter(ChatMessage.id == parent_message_id)\n .first()\n )\n if parent:\n parent.latest_child_message_id = reserved[-1].id\n\n db_session.commit()\n return reserved\n\n\ndef set_preferred_response(\n db_session: Session,\n user_message_id: int,\n preferred_assistant_message_id: int,\n) -> None:\n \"\"\"Mark one assistant response as the user's preferred choice in a multi-model turn.\n\n Also advances ``latest_child_message_id`` so the preferred response becomes\n the active branch for any subsequent messages in the conversation.\n\n Args:\n db_session: Active database session.\n user_message_id: Primary key of the ``USER``-type ``ChatMessage`` whose\n preferred response is being set.\n preferred_assistant_message_id: Primary key of the ``ASSISTANT``-type\n ``ChatMessage`` to prefer. Must be a direct child of ``user_message_id``.\n\n Raises:\n ValueError: If either message is not found, if ``user_message_id`` does not\n refer to a USER message, or if the assistant message is not a direct child\n of the user message.\n \"\"\"\n user_msg = db_session.get(ChatMessage, user_message_id)\n if user_msg is None:\n raise ValueError(f\"User message {user_message_id} not found\")\n if user_msg.message_type != MessageType.USER:\n raise ValueError(f\"Message {user_message_id} is not a user message\")\n\n assistant_msg = db_session.get(ChatMessage, preferred_assistant_message_id)\n if assistant_msg is None:\n raise ValueError(\n f\"Assistant message {preferred_assistant_message_id} not found\"\n )\n if assistant_msg.parent_message_id != user_message_id:\n raise ValueError(\n f\"Assistant message {preferred_assistant_message_id} is not a child of user message {user_message_id}\"\n )\n\n user_msg.preferred_response_id = preferred_assistant_message_id\n user_msg.latest_child_message_id = preferred_assistant_message_id\n db_session.commit()\n\n\ndef create_new_chat_message(\n chat_session_id: UUID,\n parent_message: ChatMessage,\n message: str,\n token_count: int,\n message_type: MessageType,\n db_session: Session,\n files: list[FileDescriptor] | None = None,\n error: str | None = None,\n commit: bool = True,\n reserved_message_id: int | None = None,\n reasoning_tokens: str | None = None,\n) -> ChatMessage:\n if reserved_message_id is not None:\n # Edit existing message\n existing_message = db_session.query(ChatMessage).get(reserved_message_id)\n if existing_message is None:\n raise ValueError(f\"No message found with id {reserved_message_id}\")\n\n existing_message.chat_session_id = chat_session_id\n existing_message.parent_message_id = parent_message.id\n existing_message.message = message\n existing_message.token_count = token_count\n existing_message.message_type = message_type\n existing_message.files = files\n existing_message.error = error\n existing_message.reasoning_tokens = reasoning_tokens\n new_chat_message = existing_message\n else:\n # Create new message\n new_chat_message = ChatMessage(\n chat_session_id=chat_session_id,\n parent_message_id=parent_message.id,\n latest_child_message_id=None,\n message=message,\n token_count=token_count,\n message_type=message_type,\n files=files,\n error=error,\n reasoning_tokens=reasoning_tokens,\n )\n db_session.add(new_chat_message)\n\n # Flush the session to get an ID for the new chat message\n db_session.flush()\n\n parent_message.latest_child_message_id = new_chat_message.id\n if commit:\n db_session.commit()\n\n return new_chat_message\n\n\ndef set_as_latest_chat_message(\n chat_message: ChatMessage,\n user_id: UUID | None,\n db_session: Session,\n) -> None:\n parent_message_id = chat_message.parent_message_id\n\n if parent_message_id is None:\n raise RuntimeError(\n f\"Trying to set a latest message without parent, message id: {chat_message.id}\"\n )\n\n parent_message = get_chat_message(\n chat_message_id=parent_message_id, user_id=user_id, db_session=db_session\n )\n\n parent_message.latest_child_message_id = chat_message.id\n\n db_session.commit()\n\n\ndef create_db_search_doc(\n server_search_doc: ServerSearchDoc,\n db_session: Session,\n commit: bool = True,\n) -> DBSearchDoc:\n db_search_doc = DBSearchDoc(\n document_id=sanitize_string(server_search_doc.document_id),\n chunk_ind=server_search_doc.chunk_ind,\n semantic_id=sanitize_string(server_search_doc.semantic_identifier),\n link=(\n sanitize_string(server_search_doc.link)\n if server_search_doc.link is not None\n else None\n ),\n blurb=sanitize_string(server_search_doc.blurb),\n source_type=server_search_doc.source_type,\n boost=server_search_doc.boost,\n hidden=server_search_doc.hidden,\n doc_metadata=server_search_doc.metadata,\n is_relevant=server_search_doc.is_relevant,\n relevance_explanation=(\n sanitize_string(server_search_doc.relevance_explanation)\n if server_search_doc.relevance_explanation is not None\n else None\n ),\n score=server_search_doc.score or 0.0,\n match_highlights=[\n sanitize_string(h) for h in server_search_doc.match_highlights\n ],\n updated_at=server_search_doc.updated_at,\n primary_owners=(\n [sanitize_string(o) for o in server_search_doc.primary_owners]\n if server_search_doc.primary_owners is not None\n else None\n ),\n secondary_owners=(\n [sanitize_string(o) for o in server_search_doc.secondary_owners]\n if server_search_doc.secondary_owners is not None\n else None\n ),\n is_internet=server_search_doc.is_internet,\n )\n\n db_session.add(db_search_doc)\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n return db_search_doc\n\n\ndef get_db_search_doc_by_id(doc_id: int, db_session: Session) -> DBSearchDoc | None:\n \"\"\"There are no safety checks here like user permission etc., use with caution\"\"\"\n search_doc = db_session.query(DBSearchDoc).filter(DBSearchDoc.id == doc_id).first()\n return search_doc\n\n\ndef get_db_search_doc_by_document_id(\n document_id: str, db_session: Session\n) -> DBSearchDoc | None:\n \"\"\"Get SearchDoc by document_id field. There are no safety checks here like user permission etc., use with caution\"\"\"\n search_doc = (\n db_session.query(DBSearchDoc)\n .filter(DBSearchDoc.document_id == document_id)\n .first()\n )\n return search_doc\n\n\ndef translate_db_search_doc_to_saved_search_doc(\n db_search_doc: DBSearchDoc,\n remove_doc_content: bool = False,\n) -> SavedSearchDoc:\n return SavedSearchDoc(\n db_doc_id=db_search_doc.id,\n score=db_search_doc.score,\n document_id=db_search_doc.document_id,\n chunk_ind=db_search_doc.chunk_ind,\n semantic_identifier=db_search_doc.semantic_id,\n link=db_search_doc.link,\n blurb=db_search_doc.blurb if not remove_doc_content else \"\",\n source_type=db_search_doc.source_type,\n boost=db_search_doc.boost,\n hidden=db_search_doc.hidden,\n metadata=db_search_doc.doc_metadata if not remove_doc_content else {},\n match_highlights=(\n db_search_doc.match_highlights if not remove_doc_content else []\n ),\n relevance_explanation=db_search_doc.relevance_explanation,\n is_relevant=db_search_doc.is_relevant,\n updated_at=db_search_doc.updated_at if not remove_doc_content else None,\n primary_owners=db_search_doc.primary_owners if not remove_doc_content else [],\n secondary_owners=(\n db_search_doc.secondary_owners if not remove_doc_content else []\n ),\n is_internet=db_search_doc.is_internet,\n )\n\n\ndef translate_db_message_to_chat_message_detail(\n chat_message: ChatMessage,\n remove_doc_content: bool = False,\n) -> ChatMessageDetail:\n # Get current feedback if any\n current_feedback = None\n if chat_message.chat_message_feedbacks:\n latest_feedback = chat_message.chat_message_feedbacks[-1]\n if latest_feedback.is_positive is not None:\n current_feedback = \"like\" if latest_feedback.is_positive else \"dislike\"\n\n # Convert citations from {citation_num: db_doc_id} to {citation_num: document_id}\n converted_citations = None\n if chat_message.citations and chat_message.search_docs:\n # Build lookup map: db_doc_id -> document_id\n db_doc_id_to_document_id = {\n doc.id: doc.document_id for doc in chat_message.search_docs\n }\n\n converted_citations = {}\n for citation_num, db_doc_id in chat_message.citations.items():\n document_id = db_doc_id_to_document_id.get(db_doc_id)\n if document_id:\n converted_citations[citation_num] = document_id\n\n top_documents = [\n translate_db_search_doc_to_saved_search_doc(\n db_doc, remove_doc_content=remove_doc_content\n )\n for db_doc in chat_message.search_docs\n ]\n top_documents = sorted(\n top_documents, key=lambda doc: doc.score or 0.0, reverse=True\n )\n chat_msg_detail = ChatMessageDetail(\n chat_session_id=chat_message.chat_session_id,\n message_id=chat_message.id,\n parent_message=chat_message.parent_message_id,\n latest_child_message=chat_message.latest_child_message_id,\n message=chat_message.message,\n reasoning_tokens=chat_message.reasoning_tokens,\n message_type=chat_message.message_type,\n context_docs=top_documents,\n citations=converted_citations,\n time_sent=chat_message.time_sent,\n files=chat_message.files or [],\n error=chat_message.error,\n current_feedback=current_feedback,\n processing_duration_seconds=chat_message.processing_duration_seconds,\n preferred_response_id=chat_message.preferred_response_id,\n model_display_name=chat_message.model_display_name,\n )\n\n return chat_msg_detail\n\n\ndef update_chat_session_updated_at_timestamp(\n chat_session_id: UUID, db_session: Session\n) -> None:\n \"\"\"\n Explicitly update the timestamp on a chat session without modifying other fields.\n This is useful when adding messages to a chat session to reflect recent activity.\n \"\"\"\n\n # Direct SQL update to avoid loading the entire object if it's not already loaded\n db_session.execute(\n update(ChatSession)\n .where(ChatSession.id == chat_session_id)\n .values(time_updated=func.now())\n )\n # No commit - the caller is responsible for committing the transaction\n\n\ndef create_search_doc_from_inference_section(\n inference_section: InferenceSection,\n is_internet: bool,\n db_session: Session,\n score: float = 0.0,\n is_relevant: bool | None = None,\n relevance_explanation: str | None = None,\n commit: bool = False,\n) -> DBSearchDoc:\n \"\"\"Create a SearchDoc in the database from an InferenceSection.\"\"\"\n\n db_search_doc = DBSearchDoc(\n document_id=inference_section.center_chunk.document_id,\n chunk_ind=inference_section.center_chunk.chunk_id,\n semantic_id=inference_section.center_chunk.semantic_identifier,\n link=(\n inference_section.center_chunk.source_links.get(0)\n if inference_section.center_chunk.source_links\n else None\n ),\n blurb=inference_section.center_chunk.blurb,\n source_type=inference_section.center_chunk.source_type,\n boost=inference_section.center_chunk.boost,\n hidden=inference_section.center_chunk.hidden,\n doc_metadata=inference_section.center_chunk.metadata,\n score=score,\n is_relevant=is_relevant,\n relevance_explanation=relevance_explanation,\n match_highlights=inference_section.center_chunk.match_highlights,\n updated_at=inference_section.center_chunk.updated_at,\n primary_owners=inference_section.center_chunk.primary_owners or [],\n secondary_owners=inference_section.center_chunk.secondary_owners or [],\n is_internet=is_internet,\n )\n\n db_session.add(db_search_doc)\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n return db_search_doc\n\n\ndef create_search_doc_from_saved_search_doc(\n saved_search_doc: SavedSearchDoc,\n) -> DBSearchDoc:\n \"\"\"Convert SavedSearchDoc (server model) into DB SearchDoc with correct field mapping.\"\"\"\n return DBSearchDoc(\n document_id=saved_search_doc.document_id,\n chunk_ind=saved_search_doc.chunk_ind,\n # Map Pydantic semantic_identifier -> DB semantic_id; ensure non-null\n semantic_id=saved_search_doc.semantic_identifier or \"Unknown\",\n link=saved_search_doc.link,\n blurb=saved_search_doc.blurb,\n source_type=saved_search_doc.source_type,\n boost=saved_search_doc.boost,\n hidden=saved_search_doc.hidden,\n # Map metadata -> doc_metadata (DB column name)\n doc_metadata=saved_search_doc.metadata,\n # SavedSearchDoc.score exists and defaults to 0.0\n score=saved_search_doc.score or 0.0,\n match_highlights=saved_search_doc.match_highlights,\n updated_at=saved_search_doc.updated_at,\n primary_owners=saved_search_doc.primary_owners,\n secondary_owners=saved_search_doc.secondary_owners,\n is_internet=saved_search_doc.is_internet,\n is_relevant=saved_search_doc.is_relevant,\n relevance_explanation=saved_search_doc.relevance_explanation,\n )\n\n\ndef update_db_session_with_messages(\n db_session: Session,\n chat_message_id: int,\n chat_session_id: UUID,\n message: str | None = None,\n message_type: str | None = None,\n token_count: int | None = None,\n error: str | None = None,\n update_parent_message: bool = True,\n files: list[FileDescriptor] | None = None,\n reasoning_tokens: str | None = None,\n commit: bool = False,\n) -> ChatMessage:\n chat_message = (\n db_session.query(ChatMessage)\n .filter(\n ChatMessage.id == chat_message_id,\n ChatMessage.chat_session_id == chat_session_id,\n )\n .first()\n )\n if not chat_message:\n raise ValueError(\"Chat message with id not found\") # should never happen\n\n if message:\n chat_message.message = message\n if message_type:\n chat_message.message_type = MessageType(message_type)\n if token_count:\n chat_message.token_count = token_count\n if error:\n chat_message.error = error\n if files is not None:\n chat_message.files = files\n if reasoning_tokens is not None:\n chat_message.reasoning_tokens = reasoning_tokens\n\n if update_parent_message:\n parent_chat_message = (\n db_session.query(ChatMessage)\n .filter(ChatMessage.id == chat_message.parent_message_id)\n .first()\n )\n if parent_chat_message:\n parent_chat_message.latest_child_message_id = chat_message.id\n\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n return chat_message\n" }, { "path": "backend/onyx/db/chat_search.py", "content": "from typing import List\nfrom typing import Optional\nfrom typing import Tuple\nfrom uuid import UUID\n\nfrom sqlalchemy import column\nfrom sqlalchemy import desc\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql.expression import ColumnClause\n\nfrom onyx.db.models import ChatMessage\nfrom onyx.db.models import ChatSession\n\n\ndef search_chat_sessions(\n user_id: UUID | None,\n db_session: Session,\n query: Optional[str] = None,\n page: int = 1,\n page_size: int = 10,\n include_deleted: bool = False,\n include_onyxbot_flows: bool = False,\n) -> Tuple[List[ChatSession], bool]:\n \"\"\"\n Fast full-text search on ChatSession + ChatMessage using tsvectors.\n\n If no query is provided, returns the most recent chat sessions.\n Otherwise, searches both chat messages and session descriptions.\n\n Returns a tuple of (sessions, has_more) where has_more indicates if\n there are additional results beyond the requested page.\n \"\"\"\n offset_val = (page - 1) * page_size\n\n # If no query, just return the most recent sessions\n if not query or not query.strip():\n stmt = (\n select(ChatSession)\n .order_by(desc(ChatSession.time_created))\n .offset(offset_val)\n .limit(page_size + 1)\n )\n if user_id is not None:\n stmt = stmt.where(ChatSession.user_id == user_id)\n if not include_onyxbot_flows:\n stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))\n if not include_deleted:\n stmt = stmt.where(ChatSession.deleted.is_(False))\n\n result = db_session.execute(stmt.options(joinedload(ChatSession.persona)))\n sessions = result.scalars().all()\n\n has_more = len(sessions) > page_size\n if has_more:\n sessions = sessions[:page_size]\n\n return list(sessions), has_more\n\n # Otherwise, proceed with full-text search\n query = query.strip()\n\n base_conditions = []\n if user_id is not None:\n base_conditions.append(ChatSession.user_id == user_id)\n if not include_onyxbot_flows:\n base_conditions.append(ChatSession.onyxbot_flow.is_(False))\n if not include_deleted:\n base_conditions.append(ChatSession.deleted.is_(False))\n\n message_tsv: ColumnClause = column(\"message_tsv\")\n description_tsv: ColumnClause = column(\"description_tsv\")\n\n ts_query = func.plainto_tsquery(\"english\", query)\n\n description_session_ids = (\n select(ChatSession.id)\n .where(*base_conditions)\n .where(description_tsv.op(\"@@\")(ts_query))\n )\n\n message_session_ids = (\n select(ChatMessage.chat_session_id)\n .join(ChatSession, ChatMessage.chat_session_id == ChatSession.id)\n .where(*base_conditions)\n .where(message_tsv.op(\"@@\")(ts_query))\n )\n\n combined_ids = description_session_ids.union(message_session_ids).alias(\n \"combined_ids\"\n )\n\n final_stmt = (\n select(ChatSession)\n .join(combined_ids, ChatSession.id == combined_ids.c.id)\n .order_by(desc(ChatSession.time_created))\n .distinct()\n .offset(offset_val)\n .limit(page_size + 1)\n .options(joinedload(ChatSession.persona))\n )\n\n session_objs = db_session.execute(final_stmt).scalars().all()\n\n has_more = len(session_objs) > page_size\n if has_more:\n session_objs = session_objs[:page_size]\n\n return list(session_objs), has_more\n" }, { "path": "backend/onyx/db/chunk.py", "content": "from datetime import datetime\nfrom datetime import timezone\n\nfrom sqlalchemy import delete\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import ChunkStats\nfrom onyx.indexing.models import UpdatableChunkData\n\n\ndef update_chunk_boost_components__no_commit(\n chunk_data: list[UpdatableChunkData],\n db_session: Session,\n) -> None:\n \"\"\"Updates the chunk_boost_components for chunks in the database.\n\n Args:\n chunk_data: List of dicts containing chunk_id, document_id, and boost_score\n db_session: SQLAlchemy database session\n \"\"\"\n if not chunk_data:\n return\n\n for data in chunk_data:\n chunk_in_doc_id = int(data.chunk_id)\n if chunk_in_doc_id < 0:\n raise ValueError(f\"Chunk ID is empty for chunk {data}\")\n\n chunk_document_id = f\"{data.document_id}__{chunk_in_doc_id}\"\n chunk_stats = (\n db_session.query(ChunkStats)\n .filter(\n ChunkStats.id == chunk_document_id,\n )\n .first()\n )\n\n score = data.boost_score\n\n if chunk_stats:\n chunk_stats.information_content_boost = score\n chunk_stats.last_modified = datetime.now(timezone.utc)\n db_session.add(chunk_stats)\n else:\n # do not save new chunks with a neutral boost score\n if score == 1.0:\n continue\n # Create new record\n chunk_stats = ChunkStats(\n document_id=data.document_id,\n chunk_in_doc_id=chunk_in_doc_id,\n information_content_boost=score,\n )\n db_session.add(chunk_stats)\n\n\ndef delete_chunk_stats_by_connector_credential_pair__no_commit(\n db_session: Session, document_ids: list[str]\n) -> None:\n \"\"\"This deletes just chunk stats in postgres.\"\"\"\n stmt = delete(ChunkStats).where(ChunkStats.document_id.in_(document_ids))\n\n db_session.execute(stmt)\n" }, { "path": "backend/onyx/db/code_interpreter.py", "content": "from sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import CodeInterpreterServer\n\n\ndef fetch_code_interpreter_server(\n db_session: Session,\n) -> CodeInterpreterServer:\n server = db_session.scalars(select(CodeInterpreterServer)).one()\n return server\n\n\ndef update_code_interpreter_server_enabled(\n db_session: Session,\n enabled: bool,\n) -> CodeInterpreterServer:\n server = db_session.scalars(select(CodeInterpreterServer)).one()\n server.server_enabled = enabled\n db_session.commit()\n return server\n" }, { "path": "backend/onyx/db/connector.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom typing import cast\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import exists\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DEFAULT_PRUNING_FREQ\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import IndexingMode\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import FederatedConnector\nfrom onyx.db.models import IndexAttempt\nfrom onyx.kg.models import KGConnectorData\nfrom onyx.server.documents.models import ConnectorBase\nfrom onyx.server.documents.models import ObjectCreationIdResponse\nfrom onyx.server.models import StatusResponse\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef check_federated_connectors_exist(db_session: Session) -> bool:\n stmt = select(exists(FederatedConnector))\n result = db_session.execute(stmt)\n return result.scalar() or False\n\n\ndef check_connectors_exist(db_session: Session) -> bool:\n # Connector 0 is created on server startup as a default for ingestion\n # it will always exist and we don't need to count it for this\n stmt = select(exists(Connector).where(Connector.id > 0))\n result = db_session.execute(stmt)\n return result.scalar() or False\n\n\ndef check_user_files_exist(db_session: Session) -> bool:\n \"\"\"Check if any user files exist in the system.\n\n This is used to determine if the search tool should be available\n when there are no regular connectors but there are user files\n (User Knowledge mode).\n \"\"\"\n from onyx.db.models import UserFile\n from onyx.db.enums import UserFileStatus\n\n stmt = select(exists(UserFile).where(UserFile.status == UserFileStatus.COMPLETED))\n result = db_session.execute(stmt)\n return result.scalar() or False\n\n\ndef fetch_connectors(\n db_session: Session,\n sources: list[DocumentSource] | None = None,\n input_types: list[InputType] | None = None,\n) -> list[Connector]:\n stmt = select(Connector)\n if sources is not None:\n stmt = stmt.where(Connector.source.in_(sources))\n if input_types is not None:\n stmt = stmt.where(Connector.input_type.in_(input_types))\n results = db_session.scalars(stmt)\n return list(results.all())\n\n\ndef connector_by_name_source_exists(\n connector_name: str, source: DocumentSource, db_session: Session\n) -> bool:\n stmt = select(Connector).where(\n Connector.name == connector_name, Connector.source == source\n )\n result = db_session.execute(stmt)\n connector = result.scalar_one_or_none()\n return connector is not None\n\n\ndef fetch_connector_by_id(connector_id: int, db_session: Session) -> Connector | None:\n stmt = select(Connector).where(Connector.id == connector_id)\n result = db_session.execute(stmt)\n connector = result.scalar_one_or_none()\n return connector\n\n\ndef fetch_ingestion_connector_by_name(\n connector_name: str, db_session: Session\n) -> Connector | None:\n stmt = (\n select(Connector)\n .where(Connector.name == connector_name)\n .where(Connector.source == DocumentSource.INGESTION_API)\n )\n result = db_session.execute(stmt)\n connector = result.scalar_one_or_none()\n return connector\n\n\ndef create_connector(\n db_session: Session,\n connector_data: ConnectorBase,\n) -> ObjectCreationIdResponse:\n if connector_by_name_source_exists(\n connector_data.name, connector_data.source, db_session\n ):\n raise ValueError(\n \"Connector by this name already exists, duplicate naming not allowed.\"\n )\n\n connector = Connector(\n name=connector_data.name,\n source=connector_data.source,\n input_type=connector_data.input_type,\n connector_specific_config=connector_data.connector_specific_config,\n refresh_freq=connector_data.refresh_freq,\n indexing_start=connector_data.indexing_start,\n prune_freq=connector_data.prune_freq,\n )\n db_session.add(connector)\n db_session.commit()\n\n return ObjectCreationIdResponse(id=connector.id)\n\n\ndef update_connector(\n connector_id: int,\n connector_data: ConnectorBase,\n db_session: Session,\n) -> Connector | None:\n connector = fetch_connector_by_id(connector_id, db_session)\n if connector is None:\n return None\n\n if connector_data.name != connector.name and connector_by_name_source_exists(\n connector_data.name, connector_data.source, db_session\n ):\n raise ValueError(\n \"Connector by this name already exists, duplicate naming not allowed.\"\n )\n\n connector.name = connector_data.name\n connector.source = connector_data.source\n connector.input_type = connector_data.input_type\n connector.connector_specific_config = connector_data.connector_specific_config\n connector.refresh_freq = connector_data.refresh_freq\n connector.prune_freq = (\n connector_data.prune_freq\n if connector_data.prune_freq is not None\n else DEFAULT_PRUNING_FREQ\n )\n\n db_session.commit()\n return connector\n\n\ndef delete_connector(\n db_session: Session,\n connector_id: int,\n) -> StatusResponse[int]:\n \"\"\"Only used in special cases (e.g. a connector is in a bad state and we need to delete it).\n Be VERY careful using this, as it could lead to a bad state if not used correctly.\n \"\"\"\n connector = fetch_connector_by_id(connector_id, db_session)\n if connector is None:\n return StatusResponse(\n success=True, message=\"Connector was already deleted\", data=connector_id\n )\n\n db_session.delete(connector)\n return StatusResponse(\n success=True, message=\"Connector deleted successfully\", data=connector_id\n )\n\n\ndef get_connector_credential_ids(\n connector_id: int,\n db_session: Session,\n) -> list[int]:\n connector = fetch_connector_by_id(connector_id, db_session)\n if connector is None:\n raise ValueError(f\"Connector by id {connector_id} does not exist\")\n\n return [association.credential.id for association in connector.credentials]\n\n\ndef fetch_latest_index_attempt_by_connector(\n db_session: Session,\n source: DocumentSource | None = None,\n) -> list[IndexAttempt]:\n latest_index_attempts: list[IndexAttempt] = []\n\n if source:\n connectors = fetch_connectors(db_session, sources=[source])\n else:\n connectors = fetch_connectors(db_session)\n\n if not connectors:\n return []\n\n for connector in connectors:\n latest_index_attempt = (\n db_session.query(IndexAttempt)\n .join(ConnectorCredentialPair)\n .filter(ConnectorCredentialPair.connector_id == connector.id)\n .order_by(IndexAttempt.time_updated.desc())\n .first()\n )\n\n if latest_index_attempt is not None:\n latest_index_attempts.append(latest_index_attempt)\n\n return latest_index_attempts\n\n\ndef fetch_latest_index_attempts_by_status(\n db_session: Session,\n) -> list[IndexAttempt]:\n subquery = (\n db_session.query(\n IndexAttempt.connector_credential_pair_id,\n IndexAttempt.status,\n func.max(IndexAttempt.time_updated).label(\"time_updated\"),\n )\n .group_by(IndexAttempt.connector_credential_pair_id)\n .group_by(IndexAttempt.status)\n .subquery()\n )\n\n alias = aliased(IndexAttempt, subquery)\n\n query = db_session.query(IndexAttempt).join(\n alias,\n and_(\n IndexAttempt.connector_credential_pair_id\n == alias.connector_credential_pair_id,\n IndexAttempt.status == alias.status,\n IndexAttempt.time_updated == alias.time_updated,\n ),\n )\n\n return cast(list[IndexAttempt], query.all())\n\n\ndef fetch_unique_document_sources(db_session: Session) -> list[DocumentSource]:\n distinct_sources = db_session.query(Connector.source).distinct().all()\n\n sources = [\n source[0]\n for source in distinct_sources\n if source[0] != DocumentSource.INGESTION_API\n ]\n\n return sources\n\n\ndef create_initial_default_connector(db_session: Session) -> None:\n default_connector_id = 0\n default_connector = fetch_connector_by_id(default_connector_id, db_session)\n if default_connector is not None:\n if (\n default_connector.source != DocumentSource.INGESTION_API\n or default_connector.input_type != InputType.LOAD_STATE\n or default_connector.refresh_freq is not None\n or default_connector.name != \"Ingestion API\"\n or default_connector.connector_specific_config != {}\n or default_connector.prune_freq is not None\n ):\n logger.warning(\n \"Default connector does not have expected values. Updating to proper state.\"\n )\n # Ensure default connector has correct values\n default_connector.source = DocumentSource.INGESTION_API\n default_connector.input_type = InputType.LOAD_STATE\n default_connector.refresh_freq = None\n default_connector.name = \"Ingestion API\"\n default_connector.connector_specific_config = {}\n default_connector.prune_freq = None\n db_session.commit()\n return\n\n # Create a new default connector if it doesn't exist\n connector = Connector(\n id=default_connector_id,\n name=\"Ingestion API\",\n source=DocumentSource.INGESTION_API,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={},\n refresh_freq=None,\n prune_freq=None,\n )\n db_session.add(connector)\n db_session.commit()\n\n\ndef mark_ccpair_as_pruned(cc_pair_id: int, db_session: Session) -> None:\n stmt = select(ConnectorCredentialPair).where(\n ConnectorCredentialPair.id == cc_pair_id\n )\n cc_pair = db_session.scalar(stmt)\n if cc_pair is None:\n raise ValueError(f\"No cc_pair with ID: {cc_pair_id}\")\n\n cc_pair.last_pruned = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef mark_cc_pair_as_hierarchy_fetched(db_session: Session, cc_pair_id: int) -> None:\n stmt = select(ConnectorCredentialPair).where(\n ConnectorCredentialPair.id == cc_pair_id\n )\n cc_pair = db_session.scalar(stmt)\n if cc_pair is None:\n raise ValueError(f\"No cc_pair with ID: {cc_pair_id}\")\n\n cc_pair.last_time_hierarchy_fetch = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef mark_cc_pair_as_permissions_synced(\n db_session: Session, cc_pair_id: int, start_time: datetime | None\n) -> None:\n stmt = select(ConnectorCredentialPair).where(\n ConnectorCredentialPair.id == cc_pair_id\n )\n cc_pair = db_session.scalar(stmt)\n if cc_pair is None:\n raise ValueError(f\"No cc_pair with ID: {cc_pair_id}\")\n\n cc_pair.last_time_perm_sync = start_time\n db_session.commit()\n\n\ndef mark_cc_pair_as_external_group_synced(db_session: Session, cc_pair_id: int) -> None:\n stmt = select(ConnectorCredentialPair).where(\n ConnectorCredentialPair.id == cc_pair_id\n )\n cc_pair = db_session.scalar(stmt)\n if cc_pair is None:\n raise ValueError(f\"No cc_pair with ID: {cc_pair_id}\")\n\n # The sync time can be marked after it ran because all group syncs\n # are run in full, not polling for changes.\n # If this changes, we need to update this function.\n cc_pair.last_time_external_group_sync = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef mark_ccpair_with_indexing_trigger(\n cc_pair_id: int, indexing_mode: IndexingMode | None, db_session: Session\n) -> None:\n \"\"\"indexing_mode sets a field which will be picked up by a background task\n to trigger indexing. Set to None to disable the trigger.\"\"\"\n try:\n cc_pair = db_session.execute(\n select(ConnectorCredentialPair)\n .where(ConnectorCredentialPair.id == cc_pair_id)\n .with_for_update()\n ).scalar_one()\n\n if cc_pair is None:\n raise ValueError(f\"No cc_pair with ID: {cc_pair_id}\")\n\n cc_pair.indexing_trigger = indexing_mode\n db_session.commit()\n except Exception:\n db_session.rollback()\n raise\n\n\ndef get_kg_enabled_connectors(db_session: Session) -> list[KGConnectorData]:\n \"\"\"\n Retrieves a list of connector IDs that have not been KG processed for a given tenant.\n Args:\n db_session (Session): The database session to use\n Returns:\n list[KGConnectorData]: List of connector IDs with KG extraction enabled but have unprocessed documents\n \"\"\"\n try:\n stmt = select(Connector.id, Connector.source, Connector.kg_coverage_days).where(\n Connector.kg_processing_enabled\n )\n result = db_session.execute(stmt)\n\n connector_results = [\n KGConnectorData(id=row[0], source=row[1].lower(), kg_coverage_days=row[2])\n for row in result.fetchall()\n ]\n\n return connector_results\n\n except Exception as e:\n logger.error(f\"Error fetching unprocessed connector IDs: {str(e)}\")\n raise e\n" }, { "path": "backend/onyx/db/connector_credential_pair.py", "content": "from datetime import datetime\nfrom enum import Enum\nfrom typing import TypeVarTuple\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import delete\nfrom sqlalchemy import desc\nfrom sqlalchemy import exists\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.connector import fetch_connector_by_id\nfrom onyx.db.credentials import fetch_credential_by_id\nfrom onyx.db.credentials import fetch_credential_by_id_for_user\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import ProcessingMode\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexingStatus\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup__ConnectorCredentialPair\nfrom onyx.db.models import UserRole\nfrom onyx.server.models import StatusResponse\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\nlogger = setup_logger()\n\nR = TypeVarTuple(\"R\")\n\n\nclass ConnectorType(str, Enum):\n STANDARD = \"standard\"\n USER_FILE = \"user_file\"\n\n\ndef _add_user_filters(\n stmt: Select[tuple[*R]], user: User, get_editable: bool = True\n) -> Select[tuple[*R]]:\n if user.role == UserRole.ADMIN:\n return stmt\n\n # If anonymous user, only show public cc_pairs\n if user.is_anonymous:\n where_clause = ConnectorCredentialPair.access_type == AccessType.PUBLIC\n return stmt.where(where_clause)\n\n stmt = stmt.distinct()\n UG__CCpair = aliased(UserGroup__ConnectorCredentialPair)\n User__UG = aliased(User__UserGroup)\n\n \"\"\"\n Here we select cc_pairs by relation:\n User -> User__UserGroup -> UserGroup__ConnectorCredentialPair ->\n ConnectorCredentialPair\n \"\"\"\n stmt = stmt.outerjoin(UG__CCpair).outerjoin(\n User__UG,\n User__UG.user_group_id == UG__CCpair.user_group_id,\n )\n\n \"\"\"\n Filter cc_pairs by:\n - if the user is in the user_group that owns the cc_pair\n - if the user is not a global_curator, they must also have a curator relationship\n to the user_group\n - if editing is being done, we also filter out cc_pairs that are owned by groups\n that the user isn't a curator for\n - if we are not editing, we show all cc_pairs in the groups the user is a curator\n for (as well as public cc_pairs)\n \"\"\"\n\n where_clause = User__UG.user_id == user.id\n if user.role == UserRole.CURATOR and get_editable:\n where_clause &= User__UG.is_curator == True # noqa: E712\n if get_editable:\n user_groups = select(User__UG.user_group_id).where(User__UG.user_id == user.id)\n if user.role == UserRole.CURATOR:\n user_groups = user_groups.where(\n User__UserGroup.is_curator == True # noqa: E712\n )\n where_clause &= (\n ~exists()\n .where(UG__CCpair.cc_pair_id == ConnectorCredentialPair.id)\n .where(~UG__CCpair.user_group_id.in_(user_groups))\n .correlate(ConnectorCredentialPair)\n )\n where_clause |= ConnectorCredentialPair.creator_id == user.id\n else:\n where_clause |= ConnectorCredentialPair.access_type == AccessType.PUBLIC\n where_clause |= ConnectorCredentialPair.access_type == AccessType.SYNC\n\n return stmt.where(where_clause)\n\n\ndef get_connector_credential_pairs_for_user(\n db_session: Session,\n user: User,\n get_editable: bool = True,\n ids: list[int] | None = None,\n eager_load_connector: bool = False,\n eager_load_credential: bool = False,\n eager_load_user: bool = False,\n order_by_desc: bool = False,\n source: DocumentSource | None = None,\n processing_mode: ProcessingMode | None = ProcessingMode.REGULAR,\n defer_connector_config: bool = False,\n) -> list[ConnectorCredentialPair]:\n \"\"\"Get connector credential pairs for a user.\n\n Args:\n processing_mode: Filter by processing mode. Defaults to REGULAR to hide\n FILE_SYSTEM connectors from standard admin UI. Pass None to get all.\n defer_connector_config: If True, skips loading Connector.connector_specific_config\n to avoid fetching large JSONB blobs when they aren't needed.\n \"\"\"\n if eager_load_user:\n assert (\n eager_load_credential\n ), \"eager_load_credential must be True if eager_load_user is True\"\n stmt = select(ConnectorCredentialPair).distinct()\n\n if eager_load_connector:\n connector_load = selectinload(ConnectorCredentialPair.connector)\n if defer_connector_config:\n connector_load = connector_load.defer(Connector.connector_specific_config)\n stmt = stmt.options(connector_load)\n\n if eager_load_credential:\n load_opts = selectinload(ConnectorCredentialPair.credential)\n if eager_load_user:\n load_opts = load_opts.joinedload(Credential.user)\n stmt = stmt.options(load_opts)\n\n stmt = _add_user_filters(stmt, user, get_editable)\n\n if source:\n stmt = stmt.join(ConnectorCredentialPair.connector).where(\n Connector.source == source.value\n )\n\n if ids:\n stmt = stmt.where(ConnectorCredentialPair.id.in_(ids))\n\n if processing_mode is not None:\n stmt = stmt.where(ConnectorCredentialPair.processing_mode == processing_mode)\n\n if order_by_desc:\n stmt = stmt.order_by(desc(ConnectorCredentialPair.id))\n\n return list(db_session.scalars(stmt).unique().all())\n\n\n# For use with our thread-level parallelism utils. Note that any relationships\n# you wish to use MUST be eagerly loaded, as the session will not be available\n# after this function to allow lazy loading.\ndef get_connector_credential_pairs_for_user_parallel(\n user: User,\n get_editable: bool = True,\n ids: list[int] | None = None,\n eager_load_connector: bool = False,\n eager_load_credential: bool = False,\n eager_load_user: bool = False,\n order_by_desc: bool = False,\n source: DocumentSource | None = None,\n processing_mode: ProcessingMode | None = ProcessingMode.REGULAR,\n defer_connector_config: bool = False,\n) -> list[ConnectorCredentialPair]:\n with get_session_with_current_tenant() as db_session:\n return get_connector_credential_pairs_for_user(\n db_session=db_session,\n user=user,\n get_editable=get_editable,\n ids=ids,\n eager_load_connector=eager_load_connector,\n eager_load_credential=eager_load_credential,\n eager_load_user=eager_load_user,\n order_by_desc=order_by_desc,\n source=source,\n processing_mode=processing_mode,\n defer_connector_config=defer_connector_config,\n )\n\n\ndef get_connector_credential_pairs(\n db_session: Session, ids: list[int] | None = None\n) -> list[ConnectorCredentialPair]:\n stmt = select(ConnectorCredentialPair).distinct()\n\n if ids:\n stmt = stmt.where(ConnectorCredentialPair.id.in_(ids))\n\n return list(db_session.scalars(stmt).all())\n\n\ndef add_deletion_failure_message(\n db_session: Session,\n cc_pair_id: int,\n failure_message: str,\n) -> None:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n return\n cc_pair.deletion_failure_message = failure_message\n db_session.commit()\n\n\ndef get_cc_pair_groups_for_ids(\n db_session: Session,\n cc_pair_ids: list[int],\n) -> list[UserGroup__ConnectorCredentialPair]:\n stmt = select(UserGroup__ConnectorCredentialPair).distinct()\n stmt = stmt.outerjoin(\n ConnectorCredentialPair,\n UserGroup__ConnectorCredentialPair.cc_pair_id == ConnectorCredentialPair.id,\n )\n stmt = stmt.where(UserGroup__ConnectorCredentialPair.cc_pair_id.in_(cc_pair_ids))\n return list(db_session.scalars(stmt).all())\n\n\n# For use with our thread-level parallelism utils. Note that any relationships\n# you wish to use MUST be eagerly loaded, as the session will not be available\n# after this function to allow lazy loading.\ndef get_cc_pair_groups_for_ids_parallel(\n cc_pair_ids: list[int],\n) -> list[UserGroup__ConnectorCredentialPair]:\n with get_session_with_current_tenant() as db_session:\n return get_cc_pair_groups_for_ids(db_session, cc_pair_ids)\n\n\ndef get_connector_credential_pair_for_user(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n user: User,\n get_editable: bool = True,\n) -> ConnectorCredentialPair | None:\n stmt = select(ConnectorCredentialPair)\n stmt = _add_user_filters(stmt, user, get_editable)\n stmt = stmt.where(ConnectorCredentialPair.connector_id == connector_id)\n stmt = stmt.where(ConnectorCredentialPair.credential_id == credential_id)\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef get_connector_credential_pair(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n) -> ConnectorCredentialPair | None:\n stmt = select(ConnectorCredentialPair)\n stmt = stmt.where(ConnectorCredentialPair.connector_id == connector_id)\n stmt = stmt.where(ConnectorCredentialPair.credential_id == credential_id)\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef get_connector_credential_pair_from_id_for_user(\n cc_pair_id: int,\n db_session: Session,\n user: User,\n get_editable: bool = True,\n) -> ConnectorCredentialPair | None:\n stmt = select(ConnectorCredentialPair).distinct()\n stmt = _add_user_filters(stmt, user, get_editable)\n stmt = stmt.where(ConnectorCredentialPair.id == cc_pair_id)\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef verify_user_has_access_to_cc_pair(\n cc_pair_id: int,\n db_session: Session,\n user: User,\n get_editable: bool = True,\n) -> bool:\n stmt = select(ConnectorCredentialPair.id)\n stmt = _add_user_filters(stmt, user, get_editable)\n stmt = stmt.where(ConnectorCredentialPair.id == cc_pair_id)\n result = db_session.execute(stmt)\n return result.scalars().first() is not None\n\n\ndef get_connector_credential_pair_from_id(\n db_session: Session,\n cc_pair_id: int,\n eager_load_connector: bool = False,\n eager_load_credential: bool = False,\n) -> ConnectorCredentialPair | None:\n stmt = select(ConnectorCredentialPair).distinct()\n stmt = stmt.where(ConnectorCredentialPair.id == cc_pair_id)\n\n if eager_load_credential:\n stmt = stmt.options(joinedload(ConnectorCredentialPair.credential))\n if eager_load_connector:\n stmt = stmt.options(joinedload(ConnectorCredentialPair.connector))\n\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef get_connector_credential_pairs_for_source(\n db_session: Session,\n source: DocumentSource,\n) -> list[ConnectorCredentialPair]:\n stmt = (\n select(ConnectorCredentialPair)\n .join(ConnectorCredentialPair.connector)\n .where(Connector.source == source)\n )\n return list(db_session.scalars(stmt).unique().all())\n\n\ndef get_last_successful_attempt_poll_range_end(\n cc_pair_id: int,\n earliest_index: float,\n search_settings: SearchSettings,\n db_session: Session,\n) -> float:\n \"\"\"Used to get the latest `poll_range_end` for a given connector and credential.\n\n This can be used to determine the next \"start\" time for a new index attempt.\n\n Note that the attempts time_started is not necessarily correct - that gets set\n separately and is similar but not exactly the same as the `poll_range_end`.\n \"\"\"\n latest_successful_index_attempt = (\n db_session.query(IndexAttempt)\n .join(\n ConnectorCredentialPair,\n IndexAttempt.connector_credential_pair_id == ConnectorCredentialPair.id,\n )\n .filter(\n ConnectorCredentialPair.id == cc_pair_id,\n IndexAttempt.search_settings_id == search_settings.id,\n IndexAttempt.status == IndexingStatus.SUCCESS,\n )\n .order_by(IndexAttempt.poll_range_end.desc())\n .first()\n )\n if (\n not latest_successful_index_attempt\n or not latest_successful_index_attempt.poll_range_end\n ):\n return earliest_index\n\n return latest_successful_index_attempt.poll_range_end.timestamp()\n\n\n\"\"\"Updates\"\"\"\n\n\ndef _update_connector_credential_pair(\n db_session: Session,\n cc_pair: ConnectorCredentialPair,\n status: ConnectorCredentialPairStatus | None = None,\n net_docs: int | None = None,\n run_dt: datetime | None = None,\n) -> None:\n # simply don't update last_successful_index_time if run_dt is not specified\n # at worst, this would result in re-indexing documents that were already indexed\n if run_dt is not None:\n cc_pair.last_successful_index_time = run_dt\n if net_docs is not None:\n cc_pair.total_docs_indexed += net_docs\n if status is not None:\n cc_pair.status = status\n\n db_session.commit()\n\n\ndef update_connector_credential_pair_from_id(\n db_session: Session,\n cc_pair_id: int,\n status: ConnectorCredentialPairStatus | None = None,\n net_docs: int | None = None,\n run_dt: datetime | None = None,\n) -> None:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n logger.warning(\n f\"Attempted to update pair for Connector Credential Pair '{cc_pair_id}' but it does not exist\"\n )\n return\n\n _update_connector_credential_pair(\n db_session=db_session,\n cc_pair=cc_pair,\n status=status,\n net_docs=net_docs,\n run_dt=run_dt,\n )\n\n\ndef update_connector_credential_pair(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n status: ConnectorCredentialPairStatus | None = None,\n net_docs: int | None = None,\n run_dt: datetime | None = None,\n) -> None:\n cc_pair = get_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n )\n if not cc_pair:\n logger.warning(\n f\"Attempted to update pair for connector id {connector_id} and credential id {credential_id}\"\n )\n return\n\n _update_connector_credential_pair(\n db_session=db_session,\n cc_pair=cc_pair,\n status=status,\n net_docs=net_docs,\n run_dt=run_dt,\n )\n\n\ndef set_cc_pair_repeated_error_state(\n db_session: Session,\n cc_pair_id: int,\n in_repeated_error_state: bool,\n) -> None:\n stmt = (\n update(ConnectorCredentialPair)\n .where(ConnectorCredentialPair.id == cc_pair_id)\n .values(in_repeated_error_state=in_repeated_error_state)\n )\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef delete_connector_credential_pair__no_commit(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n) -> None:\n stmt = delete(ConnectorCredentialPair).where(\n ConnectorCredentialPair.connector_id == connector_id,\n ConnectorCredentialPair.credential_id == credential_id,\n )\n db_session.execute(stmt)\n\n\ndef associate_default_cc_pair(db_session: Session) -> None:\n existing_association = (\n db_session.query(ConnectorCredentialPair)\n .filter(\n ConnectorCredentialPair.connector_id == 0,\n ConnectorCredentialPair.credential_id == 0,\n )\n .one_or_none()\n )\n if existing_association is not None:\n return\n\n # DefaultCCPair has id 1 since it is the first CC pair created\n # It is DEFAULT_CC_PAIR_ID, but can't set it explicitly because it messed with the\n # auto-incrementing id\n association = ConnectorCredentialPair(\n connector_id=0,\n credential_id=0,\n access_type=AccessType.PUBLIC,\n name=\"DefaultCCPair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n )\n db_session.add(association)\n db_session.commit()\n\n\ndef _relate_groups_to_cc_pair__no_commit(\n db_session: Session,\n cc_pair_id: int,\n user_group_ids: list[int] | None = None,\n) -> None:\n if not user_group_ids:\n return\n\n for group_id in user_group_ids:\n db_session.add(\n UserGroup__ConnectorCredentialPair(\n user_group_id=group_id, cc_pair_id=cc_pair_id\n )\n )\n\n\ndef add_credential_to_connector(\n db_session: Session,\n user: User,\n connector_id: int,\n credential_id: int,\n cc_pair_name: str,\n access_type: AccessType,\n groups: list[int] | None,\n auto_sync_options: dict | None = None,\n initial_status: ConnectorCredentialPairStatus = ConnectorCredentialPairStatus.SCHEDULED,\n last_successful_index_time: datetime | None = None,\n seeding_flow: bool = False,\n processing_mode: ProcessingMode = ProcessingMode.REGULAR,\n) -> StatusResponse:\n connector = fetch_connector_by_id(connector_id, db_session)\n\n # If we are in the seeding flow, we shouldn't need to check if the credential belongs to the user\n if seeding_flow:\n credential = fetch_credential_by_id(\n credential_id=credential_id,\n db_session=db_session,\n )\n else:\n credential = fetch_credential_by_id_for_user(\n credential_id,\n user,\n db_session,\n get_editable=False,\n )\n\n if connector is None:\n raise HTTPException(status_code=404, detail=\"Connector does not exist\")\n\n if access_type == AccessType.SYNC:\n if not fetch_ee_implementation_or_noop(\n \"onyx.external_permissions.sync_params\",\n \"check_if_valid_sync_source\",\n noop_return_value=True,\n )(connector.source):\n raise HTTPException(\n status_code=400,\n detail=f\"Connector of type {connector.source} does not support SYNC access type\",\n )\n\n if credential is None:\n error_msg = (\n f\"Credential {credential_id} does not exist or does not belong to user\"\n )\n logger.error(error_msg)\n raise HTTPException(\n status_code=401,\n detail=error_msg,\n )\n\n existing_association = (\n db_session.query(ConnectorCredentialPair)\n .filter(\n ConnectorCredentialPair.connector_id == connector_id,\n ConnectorCredentialPair.credential_id == credential_id,\n )\n .one_or_none()\n )\n if existing_association is not None:\n return StatusResponse(\n success=False,\n message=f\"Connector {connector_id} already has Credential {credential_id}\",\n data=connector_id,\n )\n\n association = ConnectorCredentialPair(\n creator_id=user.id,\n connector_id=connector_id,\n credential_id=credential_id,\n name=cc_pair_name,\n status=initial_status,\n access_type=access_type,\n auto_sync_options=auto_sync_options,\n last_successful_index_time=last_successful_index_time,\n processing_mode=processing_mode,\n )\n db_session.add(association)\n db_session.flush() # make sure the association has an id\n db_session.refresh(association)\n\n _relate_groups_to_cc_pair__no_commit(\n db_session=db_session,\n cc_pair_id=association.id,\n user_group_ids=groups,\n )\n\n db_session.commit()\n\n return StatusResponse(\n success=True,\n message=f\"Creating new association between Connector {connector_id} and Credential {credential_id}\",\n data=association.id,\n )\n\n\ndef remove_credential_from_connector(\n connector_id: int,\n credential_id: int,\n user: User,\n db_session: Session,\n) -> StatusResponse[int]:\n connector = fetch_connector_by_id(connector_id, db_session)\n credential = fetch_credential_by_id_for_user(\n credential_id,\n user,\n db_session,\n get_editable=False,\n )\n\n if connector is None:\n raise HTTPException(status_code=404, detail=\"Connector does not exist\")\n\n if credential is None:\n raise HTTPException(\n status_code=404,\n detail=\"Credential does not exist or does not belong to user\",\n )\n\n association = get_connector_credential_pair_for_user(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n user=user,\n get_editable=True,\n )\n\n if association is not None:\n fetch_ee_implementation_or_noop(\n \"onyx.db.external_perm\",\n \"delete_user__ext_group_for_cc_pair__no_commit\",\n )(\n db_session=db_session,\n cc_pair_id=association.id,\n )\n db_session.delete(association)\n db_session.commit()\n return StatusResponse(\n success=True,\n message=f\"Credential {credential_id} removed from Connector\",\n data=connector_id,\n )\n\n return StatusResponse(\n success=False,\n message=f\"Connector already does not have Credential {credential_id}\",\n data=connector_id,\n )\n\n\ndef fetch_indexable_standard_connector_credential_pair_ids(\n db_session: Session,\n active_cc_pairs_only: bool = True,\n limit: int | None = None,\n) -> list[int]:\n stmt = select(ConnectorCredentialPair.id)\n\n # For regular indexing checks\n if active_cc_pairs_only:\n stmt = stmt.where(\n ConnectorCredentialPair.status.in_(\n ConnectorCredentialPairStatus.active_statuses()\n )\n )\n else:\n # For embedding swap checks, include PAUSED and exclude DELETING or INVALID\n stmt = stmt.where(\n ConnectorCredentialPair.status.in_(\n ConnectorCredentialPairStatus.indexable_statuses()\n )\n )\n\n if limit:\n stmt = stmt.limit(limit)\n\n return list(db_session.scalars(stmt))\n\n\ndef fetch_connector_credential_pair_for_connector(\n db_session: Session,\n connector_id: int,\n) -> ConnectorCredentialPair | None:\n stmt = select(ConnectorCredentialPair).where(\n ConnectorCredentialPair.connector_id == connector_id,\n )\n return db_session.scalar(stmt)\n\n\ndef resync_cc_pair(\n cc_pair: ConnectorCredentialPair,\n search_settings_id: int,\n db_session: Session,\n) -> None:\n \"\"\"\n Updates state stored in the connector_credential_pair table based on the\n latest index attempt for the given search settings.\n\n Args:\n cc_pair: ConnectorCredentialPair to resync\n search_settings_id: SearchSettings to use for resync\n db_session: Database session\n \"\"\"\n\n def find_latest_index_attempt(\n connector_id: int,\n credential_id: int,\n only_include_success: bool,\n db_session: Session,\n ) -> IndexAttempt | None:\n query = (\n db_session.query(IndexAttempt)\n .join(\n ConnectorCredentialPair,\n IndexAttempt.connector_credential_pair_id == ConnectorCredentialPair.id,\n )\n .filter(\n ConnectorCredentialPair.connector_id == connector_id,\n ConnectorCredentialPair.credential_id == credential_id,\n IndexAttempt.search_settings_id == search_settings_id,\n )\n )\n\n if only_include_success:\n query = query.filter(IndexAttempt.status == IndexingStatus.SUCCESS)\n\n latest_index_attempt = query.order_by(desc(IndexAttempt.time_started)).first()\n\n return latest_index_attempt\n\n last_success = find_latest_index_attempt(\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n only_include_success=True,\n db_session=db_session,\n )\n\n cc_pair.last_successful_index_time = (\n last_success.time_started if last_success else None\n )\n\n db_session.commit()\n\n\n# ── Metrics query helpers ──────────────────────────────────────────────\n\n\ndef get_connector_health_for_metrics(\n db_session: Session,\n) -> list: # Returns list of Row tuples\n \"\"\"Return connector health data for Prometheus metrics.\n\n Each row is (cc_pair_id, status, in_repeated_error_state,\n last_successful_index_time, name, source).\n \"\"\"\n return (\n db_session.query(\n ConnectorCredentialPair.id,\n ConnectorCredentialPair.status,\n ConnectorCredentialPair.in_repeated_error_state,\n ConnectorCredentialPair.last_successful_index_time,\n ConnectorCredentialPair.name,\n Connector.source,\n )\n .join(\n Connector,\n ConnectorCredentialPair.connector_id == Connector.id,\n )\n .all()\n )\n" }, { "path": "backend/onyx/db/constants.py", "content": "SLACK_BOT_PERSONA_PREFIX = \"__slack_bot_persona__\"\nDEFAULT_PERSONA_SLACK_CHANNEL_NAME = \"DEFAULT_SLACK_CHANNEL\"\n\nCONNECTOR_VALIDATION_ERROR_MESSAGE_PREFIX = \"ConnectorValidationError:\"\n\n\n# Sentinel value to distinguish between \"not provided\" and \"explicitly set to None\"\nclass UnsetType:\n def __repr__(self) -> str:\n return \"\"\n\n\nUNSET = UnsetType()\n" }, { "path": "backend/onyx/db/credentials.py", "content": "from typing import Any\n\nfrom sqlalchemy import exists\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql.expression import and_\nfrom sqlalchemy.sql.expression import or_\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,\n)\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import Credential__UserGroup\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n# The credentials for these sources are not real so\n# permissions are not enforced for them\nCREDENTIAL_PERMISSIONS_TO_IGNORE = {\n DocumentSource.FILE,\n DocumentSource.WEB,\n DocumentSource.NOT_APPLICABLE,\n DocumentSource.GOOGLE_SITES,\n DocumentSource.WIKIPEDIA,\n DocumentSource.MEDIAWIKI,\n}\n\nPUBLIC_CREDENTIAL_ID = 0\n\n\ndef _add_user_filters(\n stmt: Select,\n user: User,\n get_editable: bool = True,\n) -> Select:\n \"\"\"Attaches filters to the statement to ensure that the user can only\n access the appropriate credentials\"\"\"\n if user.is_anonymous:\n raise ValueError(\"Anonymous users are not allowed to access credentials\")\n\n if user.role == UserRole.ADMIN:\n # Admins can access all credentials that are public or owned by them\n # or are not associated with any user\n return stmt.where(\n or_(\n Credential.user_id == user.id,\n Credential.user_id.is_(None),\n Credential.admin_public == True, # noqa: E712\n Credential.source.in_(CREDENTIAL_PERMISSIONS_TO_IGNORE),\n )\n )\n if user.role == UserRole.BASIC:\n # Basic users can only access credentials that are owned by them\n return stmt.where(Credential.user_id == user.id)\n\n stmt = stmt.distinct()\n \"\"\"\n THIS PART IS FOR CURATORS AND GLOBAL CURATORS\n Here we select cc_pairs by relation:\n User -> User__UserGroup -> Credential__UserGroup -> Credential\n \"\"\"\n stmt = stmt.outerjoin(Credential__UserGroup).outerjoin(\n User__UserGroup,\n User__UserGroup.user_group_id == Credential__UserGroup.user_group_id,\n )\n \"\"\"\n Filter Credentials by:\n - if the user is in the user_group that owns the Credential\n - if the user is a curator, they must also have a curator relationship\n to the user_group\n - if editing is being done, we also filter out Credentials that are owned by groups\n that the user isn't a curator for\n - if we are not editing, we show all Credentials in the groups the user is a curator\n for (as well as public Credentials)\n - if we are not editing, we return all Credentials directly connected to the user\n \"\"\"\n where_clause = User__UserGroup.user_id == user.id\n if user.role == UserRole.CURATOR:\n where_clause &= User__UserGroup.is_curator == True # noqa: E712\n\n if get_editable:\n user_groups = select(User__UserGroup.user_group_id).where(\n User__UserGroup.user_id == user.id\n )\n if user.role == UserRole.CURATOR:\n user_groups = user_groups.where(\n User__UserGroup.is_curator == True # noqa: E712\n )\n where_clause &= (\n ~exists()\n .where(Credential__UserGroup.credential_id == Credential.id)\n .where(~Credential__UserGroup.user_group_id.in_(user_groups))\n .correlate(Credential)\n )\n else:\n where_clause |= Credential.curator_public == True # noqa: E712\n where_clause |= Credential.user_id == user.id # noqa: E712\n\n where_clause |= Credential.source.in_(CREDENTIAL_PERMISSIONS_TO_IGNORE)\n\n return stmt.where(where_clause)\n\n\ndef _relate_credential_to_user_groups__no_commit(\n db_session: Session,\n credential_id: int,\n user_group_ids: list[int],\n) -> None:\n credential_user_groups = []\n for group_id in user_group_ids:\n credential_user_groups.append(\n Credential__UserGroup(\n credential_id=credential_id,\n user_group_id=group_id,\n )\n )\n db_session.add_all(credential_user_groups)\n\n\ndef fetch_credentials_for_user(\n db_session: Session,\n user: User,\n get_editable: bool = True,\n) -> list[Credential]:\n stmt = select(Credential)\n stmt = _add_user_filters(stmt, user, get_editable=get_editable)\n results = db_session.scalars(stmt)\n return list(results.all())\n\n\ndef fetch_credential_by_id_for_user(\n credential_id: int,\n user: User,\n db_session: Session,\n get_editable: bool = True,\n) -> Credential | None:\n stmt = select(Credential).distinct()\n stmt = stmt.where(Credential.id == credential_id)\n stmt = _add_user_filters(\n stmt=stmt,\n user=user,\n get_editable=get_editable,\n )\n result = db_session.execute(stmt)\n credential = result.scalar_one_or_none()\n return credential\n\n\ndef fetch_credential_by_id(\n credential_id: int,\n db_session: Session,\n) -> Credential | None:\n stmt = select(Credential).distinct()\n stmt = stmt.where(Credential.id == credential_id)\n result = db_session.execute(stmt)\n credential = result.scalar_one_or_none()\n return credential\n\n\ndef fetch_credentials_by_source_for_user(\n db_session: Session,\n user: User,\n document_source: DocumentSource | None = None,\n get_editable: bool = True,\n) -> list[Credential]:\n base_query = select(Credential).where(Credential.source == document_source)\n base_query = _add_user_filters(base_query, user, get_editable=get_editable)\n credentials = db_session.execute(base_query).scalars().all()\n return list(credentials)\n\n\ndef fetch_credentials_by_source(\n db_session: Session,\n document_source: DocumentSource | None = None,\n) -> list[Credential]:\n base_query = select(Credential).where(Credential.source == document_source)\n credentials = db_session.execute(base_query).scalars().all()\n return list(credentials)\n\n\ndef swap_credentials_connector(\n new_credential_id: int, connector_id: int, user: User, db_session: Session\n) -> ConnectorCredentialPair:\n # Check if the user has permission to use the new credential\n new_credential = fetch_credential_by_id_for_user(\n new_credential_id, user, db_session\n )\n if not new_credential:\n raise ValueError(\n f\"No Credential found with id {new_credential_id} or user doesn't have permission to use it\"\n )\n\n # Existing pair\n existing_pair = db_session.execute(\n select(ConnectorCredentialPair).where(\n ConnectorCredentialPair.connector_id == connector_id\n )\n ).scalar_one_or_none()\n\n if not existing_pair:\n raise ValueError(\n f\"No ConnectorCredentialPair found for connector_id {connector_id}\"\n )\n\n # Check if the new credential is compatible with the connector\n if new_credential.source != existing_pair.connector.source:\n raise ValueError(\n f\"New credential source {new_credential.source} does not match connector source {existing_pair.connector.source}\"\n )\n\n db_session.execute(\n update(DocumentByConnectorCredentialPair)\n .where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == existing_pair.credential_id,\n )\n )\n .values(credential_id=new_credential_id)\n )\n\n # Update the existing pair with the new credential\n existing_pair.credential_id = new_credential_id\n existing_pair.credential = new_credential\n\n # Update ccpair status if it's in INVALID state\n if existing_pair.status == ConnectorCredentialPairStatus.INVALID:\n existing_pair.status = ConnectorCredentialPairStatus.ACTIVE\n\n # Commit the changes\n db_session.commit()\n\n # Refresh the object to ensure all relationships are up-to-date\n db_session.refresh(existing_pair)\n return existing_pair\n\n\ndef create_credential(\n credential_data: CredentialBase,\n user: User,\n db_session: Session,\n) -> Credential:\n credential = Credential(\n credential_json=credential_data.credential_json,\n user_id=user.id,\n admin_public=credential_data.admin_public,\n source=credential_data.source,\n name=credential_data.name,\n curator_public=credential_data.curator_public,\n )\n db_session.add(credential)\n db_session.flush() # This ensures the credential gets an ID\n _relate_credential_to_user_groups__no_commit(\n db_session=db_session,\n credential_id=credential.id,\n user_group_ids=credential_data.groups,\n )\n\n db_session.commit()\n # Expire to ensure credential_json is reloaded as SensitiveValue from DB\n db_session.expire(credential)\n return credential\n\n\ndef _cleanup_credential__user_group_relationships__no_commit(\n db_session: Session, credential_id: int\n) -> None:\n \"\"\"NOTE: does not commit the transaction.\"\"\"\n db_session.query(Credential__UserGroup).filter(\n Credential__UserGroup.credential_id == credential_id\n ).delete(synchronize_session=False)\n\n\ndef alter_credential(\n credential_id: int,\n name: str,\n credential_json: dict[str, Any],\n user: User,\n db_session: Session,\n) -> Credential | None:\n # TODO: add user group relationship update\n credential = fetch_credential_by_id_for_user(credential_id, user, db_session)\n\n if credential is None:\n return None\n\n credential.name = name\n\n # Get existing credential_json and merge with new values\n existing_json = (\n credential.credential_json.get_value(apply_mask=False)\n if credential.credential_json\n else {}\n )\n credential.credential_json = { # type: ignore[assignment]\n **existing_json,\n **credential_json,\n }\n\n credential.user_id = user.id\n db_session.commit()\n # Expire to ensure credential_json is reloaded as SensitiveValue from DB\n db_session.expire(credential)\n return credential\n\n\ndef update_credential(\n credential_id: int,\n credential_data: CredentialBase,\n user: User,\n db_session: Session,\n) -> Credential | None:\n credential = fetch_credential_by_id_for_user(credential_id, user, db_session)\n if credential is None:\n return None\n\n credential.credential_json = credential_data.credential_json # type: ignore[assignment]\n credential.user_id = user.id if user is not None else None\n\n db_session.commit()\n # Expire to ensure credential_json is reloaded as SensitiveValue from DB\n db_session.expire(credential)\n return credential\n\n\ndef update_credential_json(\n credential_id: int,\n credential_json: dict[str, Any],\n user: User,\n db_session: Session,\n) -> Credential | None:\n credential = fetch_credential_by_id_for_user(credential_id, user, db_session)\n if credential is None:\n return None\n\n credential.credential_json = credential_json # type: ignore[assignment]\n db_session.commit()\n # Expire to ensure credential_json is reloaded as SensitiveValue from DB\n db_session.expire(credential)\n return credential\n\n\ndef backend_update_credential_json(\n credential: Credential,\n credential_json: dict[str, Any],\n db_session: Session,\n) -> None:\n \"\"\"This should not be used in any flows involving the frontend or users\"\"\"\n credential.credential_json = credential_json # type: ignore[assignment]\n db_session.commit()\n\n\ndef _delete_credential_internal(\n credential: Credential,\n credential_id: int,\n db_session: Session,\n force: bool = False,\n) -> None:\n \"\"\"Internal utility function to handle the actual deletion of a credential\"\"\"\n associated_connectors = (\n db_session.query(ConnectorCredentialPair)\n .filter(ConnectorCredentialPair.credential_id == credential_id)\n .all()\n )\n\n associated_doc_cc_pairs = (\n db_session.query(DocumentByConnectorCredentialPair)\n .filter(DocumentByConnectorCredentialPair.credential_id == credential_id)\n .all()\n )\n\n if associated_connectors or associated_doc_cc_pairs:\n if force:\n logger.warning(\n f\"Force deleting credential {credential_id} and its associated records\"\n )\n\n # Delete DocumentByConnectorCredentialPair records first\n for doc_cc_pair in associated_doc_cc_pairs:\n db_session.delete(doc_cc_pair)\n\n # Then delete ConnectorCredentialPair records\n for connector in associated_connectors:\n db_session.delete(connector)\n\n # Commit these deletions before deleting the credential\n db_session.flush()\n else:\n raise ValueError(\n f\"Cannot delete credential as it is still associated with \"\n f\"{len(associated_connectors)} connector(s) and {len(associated_doc_cc_pairs)} document(s). \"\n )\n\n if force:\n logger.warning(f\"Force deleting credential {credential_id}\")\n else:\n logger.notice(f\"Deleting credential {credential_id}\")\n\n _cleanup_credential__user_group_relationships__no_commit(db_session, credential_id)\n db_session.delete(credential)\n db_session.commit()\n\n\ndef delete_credential_for_user(\n credential_id: int,\n user: User,\n db_session: Session,\n force: bool = False,\n) -> None:\n \"\"\"Delete a credential that belongs to a specific user\"\"\"\n credential = fetch_credential_by_id_for_user(credential_id, user, db_session)\n if credential is None:\n raise ValueError(\n f\"Credential by provided id {credential_id} does not exist or does not belong to user\"\n )\n\n _delete_credential_internal(credential, credential_id, db_session, force)\n\n\ndef delete_credential(\n credential_id: int,\n db_session: Session,\n force: bool = False,\n) -> None:\n \"\"\"Delete a credential regardless of ownership (admin function)\"\"\"\n credential = fetch_credential_by_id(credential_id, db_session)\n if credential is None:\n raise ValueError(f\"Credential by provided id {credential_id} does not exist\")\n\n _delete_credential_internal(credential, credential_id, db_session, force)\n\n\ndef create_initial_public_credential(db_session: Session) -> None:\n error_msg = (\n \"DB is not in a valid initial state.\"\n \"There must exist an empty public credential for data connectors that do not require additional Auth.\"\n )\n first_credential = fetch_credential_by_id(\n credential_id=PUBLIC_CREDENTIAL_ID,\n db_session=db_session,\n )\n\n if first_credential is not None:\n credential_json_value = (\n first_credential.credential_json.get_value(apply_mask=False)\n if first_credential.credential_json\n else {}\n )\n if credential_json_value != {} or first_credential.user is not None:\n raise ValueError(error_msg)\n return\n\n credential = Credential(\n id=PUBLIC_CREDENTIAL_ID,\n credential_json={},\n user_id=None,\n )\n db_session.add(credential)\n db_session.commit()\n\n\ndef cleanup_gmail_credentials(db_session: Session) -> None:\n gmail_credentials = fetch_credentials_by_source(\n db_session=db_session, document_source=DocumentSource.GMAIL\n )\n for credential in gmail_credentials:\n db_session.delete(credential)\n db_session.commit()\n\n\ndef cleanup_google_drive_credentials(db_session: Session) -> None:\n google_drive_credentials = fetch_credentials_by_source(\n db_session=db_session, document_source=DocumentSource.GOOGLE_DRIVE\n )\n for credential in google_drive_credentials:\n db_session.delete(credential)\n db_session.commit()\n\n\ndef delete_service_account_credentials(\n user: User, db_session: Session, source: DocumentSource\n) -> None:\n credentials = fetch_credentials_for_user(db_session=db_session, user=user)\n for credential in credentials:\n credential_json = (\n credential.credential_json.get_value(apply_mask=False)\n if credential.credential_json\n else {}\n )\n if (\n credential_json.get(DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY)\n and credential.source == source\n ):\n db_session.delete(credential)\n\n db_session.commit()\n" }, { "path": "backend/onyx/db/dal.py", "content": "\"\"\"Base Data Access Layer (DAL) for database operations.\n\nThe DAL pattern groups related database operations into cohesive classes\nwith explicit session management. It supports two usage modes:\n\n 1. **External session** (FastAPI endpoints) — the caller provides a session\n whose lifecycle is managed by FastAPI's dependency injection.\n\n 2. **Self-managed session** (Celery tasks, scripts) — the DAL creates its\n own session via the tenant-aware session factory.\n\nSubclasses add domain-specific query methods while inheriting session\nmanagement. See ``ee.onyx.db.scim.ScimDAL`` for a concrete example.\n\nExample (FastAPI)::\n\n def get_scim_dal(db_session: Session = Depends(get_session)) -> ScimDAL:\n return ScimDAL(db_session)\n\n @router.get(\"/users\")\n def list_users(dal: ScimDAL = Depends(get_scim_dal)) -> ...:\n return dal.list_user_mappings(...)\n\nExample (Celery)::\n\n with ScimDAL.from_tenant(\"tenant_abc\") as dal:\n dal.create_user_mapping(...)\n dal.commit()\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\n\n\nclass DAL:\n \"\"\"Base Data Access Layer.\n\n Holds a SQLAlchemy session and provides transaction control helpers.\n Subclasses add domain-specific query methods.\n \"\"\"\n\n def __init__(self, db_session: Session) -> None:\n self._session = db_session\n\n @property\n def session(self) -> Session:\n \"\"\"Direct access to the underlying session for advanced use cases.\"\"\"\n return self._session\n\n def commit(self) -> None:\n self._session.commit()\n\n def flush(self) -> None:\n self._session.flush()\n\n def rollback(self) -> None:\n self._session.rollback()\n\n @classmethod\n @contextmanager\n def from_tenant(cls, tenant_id: str) -> Generator[\"DAL\", None, None]:\n \"\"\"Create a DAL with a self-managed session for the given tenant.\n\n The session is automatically closed when the context manager exits.\n The caller must explicitly call ``commit()`` to persist changes.\n \"\"\"\n with get_session_with_tenant(tenant_id=tenant_id) as session:\n yield cls(session)\n" }, { "path": "backend/onyx/db/deletion_attempt.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.db.index_attempt import get_last_attempt\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import IndexingStatus\nfrom onyx.db.search_settings import get_current_search_settings\n\n\ndef check_deletion_attempt_is_allowed(\n connector_credential_pair: ConnectorCredentialPair,\n db_session: Session,\n allow_scheduled: bool = False,\n) -> str | None:\n \"\"\"\n To be deletable:\n (1) connector should be paused\n (2) there should be no in-progress/planned index attempts\n\n Returns an error message if the deletion attempt is not allowed, otherwise None.\n \"\"\"\n base_error_msg = (\n f\"Connector with ID '{connector_credential_pair.connector_id}' and credential ID \"\n f\"'{connector_credential_pair.credential_id}' is not deletable.\"\n )\n\n if connector_credential_pair.status.is_active():\n return base_error_msg + \" Connector must be paused.\"\n\n connector_id = connector_credential_pair.connector_id\n credential_id = connector_credential_pair.credential_id\n search_settings = get_current_search_settings(db_session)\n\n last_indexing = get_last_attempt(\n connector_id=connector_id,\n credential_id=credential_id,\n search_settings_id=search_settings.id,\n db_session=db_session,\n )\n\n if not last_indexing:\n return None\n\n if last_indexing.status == IndexingStatus.IN_PROGRESS or (\n last_indexing.status == IndexingStatus.NOT_STARTED and not allow_scheduled\n ):\n return (\n base_error_msg\n + \" There is an ongoing / planned indexing attempt. \"\n + \"The indexing attempt must be completed or cancelled before deletion.\"\n )\n\n return None\n" }, { "path": "backend/onyx/db/discord_bot.py", "content": "\"\"\"CRUD operations for Discord bot models.\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.api_key import build_displayable_api_key\nfrom onyx.auth.api_key import generate_api_key\nfrom onyx.auth.api_key import hash_api_key\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import DISCORD_SERVICE_API_KEY_NAME\nfrom onyx.db.api_key import insert_api_key\nfrom onyx.db.models import ApiKey\nfrom onyx.db.models import DiscordBotConfig\nfrom onyx.db.models import DiscordChannelConfig\nfrom onyx.db.models import DiscordGuildConfig\nfrom onyx.db.models import User\nfrom onyx.db.utils import DiscordChannelView\nfrom onyx.server.api_key.models import APIKeyArgs\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# === DiscordBotConfig ===\n\n\ndef get_discord_bot_config(db_session: Session) -> DiscordBotConfig | None:\n \"\"\"Get the Discord bot config for this tenant (at most one).\"\"\"\n return db_session.scalar(select(DiscordBotConfig).limit(1))\n\n\ndef create_discord_bot_config(\n db_session: Session,\n bot_token: str,\n) -> DiscordBotConfig:\n \"\"\"Create the Discord bot config. Raises ValueError if already exists.\n\n The check constraint on id='SINGLETON' ensures only one config per tenant.\n \"\"\"\n existing = get_discord_bot_config(db_session)\n if existing:\n raise ValueError(\"Discord bot config already exists\")\n\n config = DiscordBotConfig(bot_token=bot_token)\n db_session.add(config)\n try:\n db_session.flush()\n except IntegrityError:\n # Race condition: another request created the config concurrently\n db_session.rollback()\n raise ValueError(\"Discord bot config already exists\")\n return config\n\n\ndef delete_discord_bot_config(db_session: Session) -> bool:\n \"\"\"Delete the Discord bot config. Returns True if deleted.\"\"\"\n result = db_session.execute(delete(DiscordBotConfig))\n db_session.flush()\n return result.rowcount > 0 # type: ignore[attr-defined]\n\n\n# === Discord Service API Key ===\n\n\ndef get_discord_service_api_key(db_session: Session) -> ApiKey | None:\n \"\"\"Get the Discord service API key if it exists.\"\"\"\n return db_session.scalar(\n select(ApiKey).where(ApiKey.name == DISCORD_SERVICE_API_KEY_NAME)\n )\n\n\ndef get_or_create_discord_service_api_key(\n db_session: Session,\n tenant_id: str,\n) -> str:\n \"\"\"Get existing Discord service API key or create one.\n\n The API key is used by the Discord bot to authenticate with the\n Onyx API pods when sending chat requests.\n\n Args:\n db_session: Database session for the tenant.\n tenant_id: The tenant ID (used for logging/context).\n\n Returns:\n The raw API key string (not hashed).\n\n Raises:\n RuntimeError: If API key creation fails.\n \"\"\"\n # Check for existing key\n existing = get_discord_service_api_key(db_session)\n if existing:\n # Database only stores the hash, so we must regenerate to get the raw key.\n # This is safe since the Discord bot is the only consumer of this key.\n logger.debug(\n f\"Found existing Discord service API key for tenant {tenant_id} that isn't in cache, regenerating to update cache\"\n )\n new_api_key = generate_api_key(tenant_id)\n existing.hashed_api_key = hash_api_key(new_api_key)\n existing.api_key_display = build_displayable_api_key(new_api_key)\n db_session.flush()\n return new_api_key\n\n # Create new API key\n logger.info(f\"Creating Discord service API key for tenant {tenant_id}\")\n api_key_args = APIKeyArgs(\n name=DISCORD_SERVICE_API_KEY_NAME,\n role=UserRole.LIMITED, # Limited role is sufficient for chat requests\n )\n api_key_descriptor = insert_api_key(\n db_session=db_session,\n api_key_args=api_key_args,\n user_id=None, # Service account, no owner\n )\n\n if not api_key_descriptor.api_key:\n raise RuntimeError(\n f\"Failed to create Discord service API key for tenant {tenant_id}\"\n )\n\n return api_key_descriptor.api_key\n\n\ndef delete_discord_service_api_key(db_session: Session) -> bool:\n \"\"\"Delete the Discord service API key for a tenant.\n\n Called when:\n - Bot config is deleted (self-hosted)\n - All guild configs are deleted (Cloud)\n\n Args:\n db_session: Database session for the tenant.\n\n Returns:\n True if the key was deleted, False if it didn't exist.\n \"\"\"\n existing_key = get_discord_service_api_key(db_session)\n if not existing_key:\n return False\n\n # Also delete the associated user\n api_key_user = db_session.scalar(\n select(User).where(User.id == existing_key.user_id) # type: ignore[arg-type]\n )\n\n db_session.delete(existing_key)\n if api_key_user:\n db_session.delete(api_key_user)\n\n db_session.flush()\n logger.info(\"Deleted Discord service API key\")\n return True\n\n\n# === DiscordGuildConfig ===\n\n\ndef get_guild_configs(\n db_session: Session,\n include_channels: bool = False,\n) -> list[DiscordGuildConfig]:\n \"\"\"Get all guild configs for this tenant.\"\"\"\n stmt = select(DiscordGuildConfig)\n if include_channels:\n stmt = stmt.options(joinedload(DiscordGuildConfig.channels))\n return list(db_session.scalars(stmt).unique().all())\n\n\ndef get_guild_config_by_internal_id(\n db_session: Session,\n internal_id: int,\n) -> DiscordGuildConfig | None:\n \"\"\"Get a specific guild config by its ID.\"\"\"\n return db_session.scalar(\n select(DiscordGuildConfig).where(DiscordGuildConfig.id == internal_id)\n )\n\n\ndef get_guild_config_by_discord_id(\n db_session: Session,\n guild_id: int,\n) -> DiscordGuildConfig | None:\n \"\"\"Get a guild config by Discord guild ID.\"\"\"\n return db_session.scalar(\n select(DiscordGuildConfig).where(DiscordGuildConfig.guild_id == guild_id)\n )\n\n\ndef get_guild_config_by_registration_key(\n db_session: Session,\n registration_key: str,\n) -> DiscordGuildConfig | None:\n \"\"\"Get a guild config by its registration key.\"\"\"\n return db_session.scalar(\n select(DiscordGuildConfig).where(\n DiscordGuildConfig.registration_key == registration_key\n )\n )\n\n\ndef create_guild_config(\n db_session: Session,\n registration_key: str,\n) -> DiscordGuildConfig:\n \"\"\"Create a new guild config with a registration key (guild_id=NULL).\"\"\"\n config = DiscordGuildConfig(registration_key=registration_key)\n db_session.add(config)\n db_session.flush()\n return config\n\n\ndef register_guild(\n db_session: Session,\n config: DiscordGuildConfig,\n guild_id: int,\n guild_name: str,\n) -> DiscordGuildConfig:\n \"\"\"Complete registration by setting guild_id and guild_name.\"\"\"\n config.guild_id = guild_id\n config.guild_name = guild_name\n config.registered_at = datetime.now(timezone.utc)\n db_session.flush()\n return config\n\n\ndef update_guild_config(\n db_session: Session,\n config: DiscordGuildConfig,\n enabled: bool,\n default_persona_id: int | None = None,\n) -> DiscordGuildConfig:\n \"\"\"Update guild config fields.\"\"\"\n config.enabled = enabled\n config.default_persona_id = default_persona_id\n db_session.flush()\n return config\n\n\ndef delete_guild_config(\n db_session: Session,\n internal_id: int,\n) -> bool:\n \"\"\"Delete guild config (cascades to channel configs). Returns True if deleted.\"\"\"\n result = db_session.execute(\n delete(DiscordGuildConfig).where(DiscordGuildConfig.id == internal_id)\n )\n db_session.flush()\n return result.rowcount > 0 # type: ignore[attr-defined]\n\n\n# === DiscordChannelConfig ===\n\n\ndef get_channel_configs(\n db_session: Session,\n guild_config_id: int,\n) -> list[DiscordChannelConfig]:\n \"\"\"Get all channel configs for a guild.\"\"\"\n return list(\n db_session.scalars(\n select(DiscordChannelConfig).where(\n DiscordChannelConfig.guild_config_id == guild_config_id\n )\n ).all()\n )\n\n\ndef get_channel_config_by_discord_ids(\n db_session: Session,\n guild_id: int,\n channel_id: int,\n) -> DiscordChannelConfig | None:\n \"\"\"Get a specific channel config by guild_id and channel_id.\"\"\"\n return db_session.scalar(\n select(DiscordChannelConfig)\n .join(DiscordGuildConfig)\n .where(\n DiscordGuildConfig.guild_id == guild_id,\n DiscordChannelConfig.channel_id == channel_id,\n )\n )\n\n\ndef get_channel_config_by_internal_ids(\n db_session: Session,\n guild_config_id: int,\n channel_config_id: int,\n) -> DiscordChannelConfig | None:\n \"\"\"Get a specific channel config by guild_config_id and channel_config_id\"\"\"\n return db_session.scalar(\n select(DiscordChannelConfig).where(\n DiscordChannelConfig.guild_config_id == guild_config_id,\n DiscordChannelConfig.id == channel_config_id,\n )\n )\n\n\ndef update_discord_channel_config(\n db_session: Session,\n config: DiscordChannelConfig,\n channel_name: str,\n thread_only_mode: bool,\n require_bot_invocation: bool,\n enabled: bool,\n persona_override_id: int | None = None,\n) -> DiscordChannelConfig:\n \"\"\"Update channel config fields.\"\"\"\n config.channel_name = channel_name\n config.require_bot_invocation = require_bot_invocation\n config.persona_override_id = persona_override_id\n config.enabled = enabled\n config.thread_only_mode = thread_only_mode\n db_session.flush()\n return config\n\n\ndef delete_discord_channel_config(\n db_session: Session,\n guild_config_id: int,\n channel_config_id: int,\n) -> bool:\n \"\"\"Delete a channel config. Returns True if deleted.\"\"\"\n result = db_session.execute(\n delete(DiscordChannelConfig).where(\n DiscordChannelConfig.guild_config_id == guild_config_id,\n DiscordChannelConfig.id == channel_config_id,\n )\n )\n db_session.flush()\n return result.rowcount > 0 # type: ignore[attr-defined]\n\n\ndef create_channel_config(\n db_session: Session,\n guild_config_id: int,\n channel_view: DiscordChannelView,\n) -> DiscordChannelConfig:\n \"\"\"Create a new channel config with default settings (disabled by default, admin enables via UI).\"\"\"\n config = DiscordChannelConfig(\n guild_config_id=guild_config_id,\n channel_id=channel_view.channel_id,\n channel_name=channel_view.channel_name,\n channel_type=channel_view.channel_type,\n is_private=channel_view.is_private,\n )\n db_session.add(config)\n db_session.flush()\n return config\n\n\ndef bulk_create_channel_configs(\n db_session: Session,\n guild_config_id: int,\n channels: list[DiscordChannelView],\n) -> list[DiscordChannelConfig]:\n \"\"\"Create multiple channel configs at once. Skips existing channels.\"\"\"\n # Get existing channel IDs for this guild\n existing_channel_ids = set(\n db_session.scalars(\n select(DiscordChannelConfig.channel_id).where(\n DiscordChannelConfig.guild_config_id == guild_config_id\n )\n ).all()\n )\n\n # Create configs for new channels only\n new_configs = []\n for channel_view in channels:\n if channel_view.channel_id not in existing_channel_ids:\n config = DiscordChannelConfig(\n guild_config_id=guild_config_id,\n channel_id=channel_view.channel_id,\n channel_name=channel_view.channel_name,\n channel_type=channel_view.channel_type,\n is_private=channel_view.is_private,\n )\n db_session.add(config)\n new_configs.append(config)\n\n db_session.flush()\n return new_configs\n\n\ndef sync_channel_configs(\n db_session: Session,\n guild_config_id: int,\n current_channels: list[DiscordChannelView],\n) -> tuple[int, int, int]:\n \"\"\"Sync channel configs with current Discord channels.\n\n - Creates configs for new channels (disabled by default)\n - Removes configs for deleted channels\n - Updates names and types for existing channels if changed\n\n Returns: (added_count, removed_count, updated_count)\n \"\"\"\n current_channel_map = {\n channel_view.channel_id: channel_view for channel_view in current_channels\n }\n current_channel_ids = set(current_channel_map.keys())\n\n # Get existing configs\n existing_configs = get_channel_configs(db_session, guild_config_id)\n existing_channel_ids = {c.channel_id for c in existing_configs}\n\n # Find channels to add, remove, and potentially update\n to_add = current_channel_ids - existing_channel_ids\n to_remove = existing_channel_ids - current_channel_ids\n\n # Add new channels\n added_count = 0\n for channel_id in to_add:\n channel_view = current_channel_map[channel_id]\n create_channel_config(db_session, guild_config_id, channel_view)\n added_count += 1\n\n # Remove deleted channels\n removed_count = 0\n for config in existing_configs:\n if config.channel_id in to_remove:\n db_session.delete(config)\n removed_count += 1\n\n # Update names, types, and privacy for existing channels if changed\n updated_count = 0\n for config in existing_configs:\n if config.channel_id in current_channel_ids:\n channel_view = current_channel_map[config.channel_id]\n changed = False\n if config.channel_name != channel_view.channel_name:\n config.channel_name = channel_view.channel_name\n changed = True\n if config.channel_type != channel_view.channel_type:\n config.channel_type = channel_view.channel_type\n changed = True\n if config.is_private != channel_view.is_private:\n config.is_private = channel_view.is_private\n changed = True\n if changed:\n updated_count += 1\n\n db_session.flush()\n return added_count, removed_count, updated_count\n" }, { "path": "backend/onyx/db/document.py", "content": "import contextlib\nimport time\nfrom collections.abc import Generator\nfrom collections.abc import Iterable\nfrom collections.abc import Sequence\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import delete\nfrom sqlalchemy import exists\nfrom sqlalchemy import func\nfrom sqlalchemy import or_\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import tuple_\nfrom sqlalchemy import update\nfrom sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.engine.util import TransactionalContext\nfrom sqlalchemy.exc import OperationalError\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql.expression import null\n\nfrom onyx.configs.constants import DEFAULT_BOOST\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.kg_configs import KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES\nfrom onyx.db.chunk import delete_chunk_stats_by_connector_credential_pair__no_commit\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.document_access import apply_document_access_filter\nfrom onyx.db.entities import delete_from_kg_entities__no_commit\nfrom onyx.db.entities import delete_from_kg_entities_extraction_staging__no_commit\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.feedback import delete_document_feedback_for_documents__no_commit\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import Document as DbDocument\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import KGEntity\nfrom onyx.db.models import KGRelationship\nfrom onyx.db.models import User\nfrom onyx.db.relationships import delete_from_kg_relationships__no_commit\nfrom onyx.db.relationships import (\n delete_from_kg_relationships_extraction_staging__no_commit,\n)\nfrom onyx.db.tag import delete_document_tags_for_documents__no_commit\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import model_to_dict\nfrom onyx.db.utils import SortOrder\nfrom onyx.document_index.interfaces import DocumentMetadata\nfrom onyx.kg.models import KGStage\nfrom onyx.server.documents.models import ConnectorCredentialPairIdentifier\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nONE_HOUR_IN_SECONDS = 60 * 60\n\n\ndef check_docs_exist(db_session: Session) -> bool:\n stmt = select(exists(DbDocument))\n result = db_session.execute(stmt)\n return result.scalar() or False\n\n\ndef count_documents_by_needs_sync(session: Session) -> int:\n \"\"\"Get the count of all documents where:\n 1. last_modified is newer than last_synced\n 2. last_synced is null (meaning we've never synced)\n AND the document has a relationship with a connector/credential pair\n\n TODO: The documents without a relationship with a connector/credential pair\n should be cleaned up somehow eventually.\n\n This function executes the query and returns the count of\n documents matching the criteria.\"\"\"\n\n return (\n session.query(DbDocument.id)\n .filter(\n or_(\n DbDocument.last_modified > DbDocument.last_synced,\n DbDocument.last_synced.is_(None),\n )\n )\n .count()\n )\n\n\ndef construct_document_id_select_by_needs_sync() -> Select:\n \"\"\"Get all document IDs that need syncing across all connector credential pairs.\n\n Returns a Select statement for documents where:\n 1. last_modified is newer than last_synced\n 2. last_synced is null (meaning we've never synced)\n AND the document has a relationship with a connector/credential pair\n \"\"\"\n return select(DbDocument.id).where(\n or_(\n DbDocument.last_modified > DbDocument.last_synced,\n DbDocument.last_synced.is_(None),\n )\n )\n\n\ndef construct_document_id_select_for_connector_credential_pair(\n connector_id: int, credential_id: int | None = None\n) -> Select:\n initial_doc_ids_stmt = select(DocumentByConnectorCredentialPair.id).where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n )\n stmt = (\n select(DbDocument.id).where(DbDocument.id.in_(initial_doc_ids_stmt)).distinct()\n )\n return stmt\n\n\ndef construct_document_select_for_connector_credential_pair(\n connector_id: int, credential_id: int | None = None\n) -> Select:\n initial_doc_ids_stmt = select(DocumentByConnectorCredentialPair.id).where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n )\n stmt = select(DbDocument).where(DbDocument.id.in_(initial_doc_ids_stmt)).distinct()\n return stmt\n\n\ndef get_documents_for_cc_pair(\n db_session: Session,\n cc_pair_id: int,\n) -> list[DbDocument]:\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n if not cc_pair:\n raise ValueError(f\"No CC pair found with ID: {cc_pair_id}\")\n stmt = construct_document_select_for_connector_credential_pair(\n connector_id=cc_pair.connector_id, credential_id=cc_pair.credential_id\n )\n return list(db_session.scalars(stmt).all())\n\n\ndef get_document_ids_for_connector_credential_pair(\n db_session: Session, connector_id: int, credential_id: int\n) -> list[str]:\n doc_ids_stmt = select(DocumentByConnectorCredentialPair.id).where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n )\n return list(db_session.execute(doc_ids_stmt).scalars().all())\n\n\ndef get_documents_for_connector_credential_pair_limited_columns(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n sort_order: SortOrder | None = None,\n) -> Sequence[DocumentRow]:\n\n doc_ids_subquery = select(DocumentByConnectorCredentialPair.id).where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n )\n doc_ids_subquery = doc_ids_subquery.join(\n DbDocument, DocumentByConnectorCredentialPair.id == DbDocument.id\n )\n\n stmt = select(\n DbDocument.id, DbDocument.doc_metadata, DbDocument.external_user_group_ids\n )\n\n stmt = stmt.where(DbDocument.id.in_(doc_ids_subquery))\n\n if sort_order == SortOrder.ASC:\n stmt = stmt.order_by(DbDocument.last_modified.asc())\n elif sort_order == SortOrder.DESC:\n stmt = stmt.order_by(DbDocument.last_modified.desc())\n\n rows = db_session.execute(stmt).mappings().all()\n\n doc_rows: list[DocumentRow] = []\n for row in rows:\n doc_row = DocumentRow(\n id=row.id,\n doc_metadata=row.doc_metadata,\n external_user_group_ids=row.external_user_group_ids or [],\n )\n doc_rows.append(doc_row)\n return doc_rows\n\n\ndef get_documents_for_connector_credential_pair(\n db_session: Session, connector_id: int, credential_id: int, limit: int | None = None\n) -> Sequence[DbDocument]:\n initial_doc_ids_stmt = select(DocumentByConnectorCredentialPair.id).where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n )\n stmt = select(DbDocument).where(DbDocument.id.in_(initial_doc_ids_stmt)).distinct()\n if limit:\n stmt = stmt.limit(limit)\n return db_session.scalars(stmt).all()\n\n\ndef get_documents_by_ids(\n db_session: Session,\n document_ids: list[str],\n) -> list[DbDocument]:\n stmt = select(DbDocument).where(DbDocument.id.in_(document_ids))\n documents = db_session.execute(stmt).scalars().all()\n return list(documents)\n\n\ndef get_documents_by_source(\n db_session: Session,\n source: DocumentSource,\n creator_id: UUID | None = None,\n) -> list[DbDocument]:\n \"\"\"Get all documents associated with a specific source type.\n\n This queries through the connector relationship to find all documents\n that were indexed by connectors of the given source type.\n\n Args:\n db_session: Database session\n source: The document source type to filter by\n creator_id: If provided, only return documents from connectors\n created by this user. Filters via ConnectorCredentialPair.\n \"\"\"\n stmt = (\n select(DbDocument)\n .join(\n DocumentByConnectorCredentialPair,\n DbDocument.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n Connector,\n ConnectorCredentialPair.connector_id == Connector.id,\n )\n .where(Connector.source == source)\n )\n if creator_id is not None:\n stmt = stmt.where(ConnectorCredentialPair.creator_id == creator_id)\n stmt = stmt.distinct()\n documents = db_session.execute(stmt).scalars().all()\n return list(documents)\n\n\ndef _apply_last_updated_cursor_filter(\n stmt: Select,\n cursor_last_modified: datetime | None,\n cursor_last_synced: datetime | None,\n cursor_document_id: str | None,\n is_ascending: bool,\n) -> Select:\n \"\"\"Apply cursor filter for last_updated sorting.\n\n ASC uses nulls_first (NULLs at start), DESC uses nulls_last (NULLs at end).\n This affects which extra clauses are needed when the cursor has NULL last_synced\n vs non-NULL last_synced.\n \"\"\"\n if not cursor_last_modified or not cursor_document_id:\n return stmt\n\n # Pick comparison operators based on sort direction\n if is_ascending:\n modified_cmp = DbDocument.last_modified > cursor_last_modified\n synced_cmp = DbDocument.last_synced > cursor_last_synced\n id_cmp = DbDocument.id > cursor_document_id\n else:\n modified_cmp = DbDocument.last_modified < cursor_last_modified\n synced_cmp = DbDocument.last_synced < cursor_last_synced\n id_cmp = DbDocument.id < cursor_document_id\n\n if cursor_last_synced is None:\n # Cursor has NULL last_synced\n # ASC (nulls_first): NULL is at start, so non-NULL values come after\n # DESC (nulls_last): NULL is at end, so nothing with non-NULL comes after\n base_clauses = [\n modified_cmp,\n and_(\n DbDocument.last_modified == cursor_last_modified,\n DbDocument.last_synced.is_(None),\n id_cmp,\n ),\n ]\n if is_ascending:\n # Any non-NULL last_synced comes after NULL when nulls_first\n base_clauses.append(\n and_(\n DbDocument.last_modified == cursor_last_modified,\n DbDocument.last_synced.is_not(None),\n )\n )\n return stmt.where(or_(*base_clauses))\n\n # Cursor has non-NULL last_synced\n # ASC (nulls_first): NULLs came before, so no NULL clause needed\n # DESC (nulls_last): NULLs come after non-NULL values\n synced_clauses = [\n synced_cmp,\n and_(DbDocument.last_synced == cursor_last_synced, id_cmp),\n ]\n if not is_ascending:\n # NULLs come after all non-NULL values when nulls_last\n synced_clauses.append(DbDocument.last_synced.is_(None))\n\n return stmt.where(\n or_(\n modified_cmp,\n and_(\n DbDocument.last_modified == cursor_last_modified,\n or_(*synced_clauses),\n ),\n )\n )\n\n\ndef _apply_name_cursor_filter_asc(\n stmt: Select,\n cursor_name: str | None,\n cursor_document_id: str | None,\n) -> Select:\n \"\"\"Apply cursor filter for name ASC sorting.\"\"\"\n if not cursor_name or not cursor_document_id:\n return stmt\n return stmt.where(\n or_(\n DbDocument.semantic_id > cursor_name,\n and_(\n DbDocument.semantic_id == cursor_name,\n DbDocument.id > cursor_document_id,\n ),\n )\n )\n\n\ndef _apply_name_cursor_filter_desc(\n stmt: Select,\n cursor_name: str | None,\n cursor_document_id: str | None,\n) -> Select:\n \"\"\"Apply cursor filter for name DESC sorting.\"\"\"\n if not cursor_name or not cursor_document_id:\n return stmt\n return stmt.where(\n or_(\n DbDocument.semantic_id < cursor_name,\n and_(\n DbDocument.semantic_id == cursor_name,\n DbDocument.id < cursor_document_id,\n ),\n )\n )\n\n\ndef get_accessible_documents_for_hierarchy_node_paginated(\n db_session: Session,\n parent_hierarchy_node_id: int,\n user_email: str | None,\n external_group_ids: list[str],\n limit: int,\n # Sort options\n sort_by_name: bool = False,\n sort_ascending: bool = False,\n # Cursor fields for last_updated sorting\n cursor_last_modified: datetime | None = None,\n cursor_last_synced: datetime | None = None,\n # Cursor field for name sorting\n cursor_name: str | None = None,\n # Document ID for tie-breaking (used by both sort types)\n cursor_document_id: str | None = None,\n) -> list[DbDocument]:\n stmt = select(DbDocument).where(\n DbDocument.parent_hierarchy_node_id == parent_hierarchy_node_id\n )\n stmt = apply_document_access_filter(stmt, user_email, external_group_ids)\n\n # Apply cursor filter based on sort type and direction\n if sort_by_name:\n if sort_ascending:\n stmt = _apply_name_cursor_filter_asc(stmt, cursor_name, cursor_document_id)\n stmt = stmt.order_by(DbDocument.semantic_id.asc(), DbDocument.id.asc())\n else:\n stmt = _apply_name_cursor_filter_desc(stmt, cursor_name, cursor_document_id)\n stmt = stmt.order_by(DbDocument.semantic_id.desc(), DbDocument.id.desc())\n else:\n # Sort by last_updated\n if sort_ascending:\n stmt = _apply_last_updated_cursor_filter(\n stmt,\n cursor_last_modified,\n cursor_last_synced,\n cursor_document_id,\n is_ascending=True,\n )\n stmt = stmt.order_by(\n DbDocument.last_modified.asc(),\n DbDocument.last_synced.asc().nulls_first(),\n DbDocument.id.asc(),\n )\n else:\n stmt = _apply_last_updated_cursor_filter(\n stmt,\n cursor_last_modified,\n cursor_last_synced,\n cursor_document_id,\n is_ascending=False,\n )\n stmt = stmt.order_by(\n DbDocument.last_modified.desc(),\n DbDocument.last_synced.desc().nulls_last(),\n DbDocument.id.desc(),\n )\n\n # Use distinct to avoid duplicates when a document belongs to multiple cc_pairs\n stmt = stmt.distinct()\n stmt = stmt.limit(limit)\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef filter_existing_document_ids(\n db_session: Session,\n document_ids: list[str],\n) -> set[str]:\n \"\"\"Filter a list of document IDs to only those that exist in the database.\n\n Args:\n db_session: Database session\n document_ids: List of document IDs to check for existence\n\n Returns:\n Set of document IDs from the input list that exist in the database\n \"\"\"\n if not document_ids:\n return set()\n stmt = select(DbDocument.id).where(DbDocument.id.in_(document_ids))\n return set(db_session.execute(stmt).scalars().all())\n\n\ndef fetch_document_ids_by_links(\n db_session: Session,\n links: list[str],\n) -> dict[str, str]:\n \"\"\"Fetch document IDs for documents whose link matches any of the provided values.\"\"\"\n if not links:\n return {}\n\n stmt = select(DbDocument.link, DbDocument.id).where(DbDocument.link.in_(links))\n rows = db_session.execute(stmt).all()\n return {link: doc_id for link, doc_id in rows if link}\n\n\ndef get_document_connector_count(\n db_session: Session,\n document_id: str,\n) -> int:\n results = get_document_connector_counts(db_session, [document_id])\n if not results or len(results) == 0:\n return 0\n\n return results[0][1]\n\n\ndef get_document_connector_counts(\n db_session: Session,\n document_ids: list[str],\n) -> Sequence[tuple[str, int]]:\n stmt = (\n select(\n DocumentByConnectorCredentialPair.id,\n func.count(),\n )\n .where(DocumentByConnectorCredentialPair.id.in_(document_ids))\n .group_by(DocumentByConnectorCredentialPair.id)\n )\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef get_document_counts_for_cc_pairs(\n db_session: Session, cc_pairs: list[ConnectorCredentialPairIdentifier]\n) -> Sequence[tuple[int, int, int]]:\n \"\"\"Returns a sequence of tuples of (connector_id, credential_id, document count)\"\"\"\n\n if not cc_pairs:\n return []\n\n # Prepare a list of (connector_id, credential_id) tuples\n cc_ids = [(x.connector_id, x.credential_id) for x in cc_pairs]\n\n # Batch to avoid generating extremely large IN clauses that can blow Postgres stack depth\n batch_size = 1000\n aggregated_counts: dict[tuple[int, int], int] = {}\n\n for start_idx in range(0, len(cc_ids), batch_size):\n batch = cc_ids[start_idx : start_idx + batch_size]\n\n stmt = (\n select(\n DocumentByConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id,\n func.count(),\n )\n .where(\n and_(\n tuple_(\n DocumentByConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id,\n ).in_(batch),\n DocumentByConnectorCredentialPair.has_been_indexed.is_(True),\n )\n )\n .group_by(\n DocumentByConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id,\n )\n )\n\n for connector_id, credential_id, cnt in db_session.execute(stmt).all():\n aggregated_counts[(connector_id, credential_id)] = cnt\n\n # Convert aggregated results back to the expected sequence of tuples\n return [\n (connector_id, credential_id, cnt)\n for (connector_id, credential_id), cnt in aggregated_counts.items()\n ]\n\n\ndef get_document_counts_for_all_cc_pairs(\n db_session: Session,\n) -> Sequence[tuple[int, int, int]]:\n \"\"\"Return (connector_id, credential_id, count) for ALL CC pairs with indexed docs.\n\n Executes a single grouped query so Postgres can fully leverage indexes,\n avoiding large batched IN-lists.\n \"\"\"\n stmt = (\n select(\n DocumentByConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id,\n func.count(),\n )\n .where(DocumentByConnectorCredentialPair.has_been_indexed.is_(True))\n .group_by(\n DocumentByConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id,\n )\n )\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef get_access_info_for_document(\n db_session: Session,\n document_id: str,\n) -> tuple[str, list[str | None], bool] | None:\n \"\"\"Gets access info for a single document by calling the get_access_info_for_documents function\n and passing a list with a single document ID.\n Args:\n db_session (Session): The database session to use.\n document_id (str): The document ID to fetch access info for.\n Returns:\n Optional[Tuple[str, List[str | None], bool]]: A tuple containing the document ID, a list of user emails,\n and a boolean indicating if the document is globally public, or None if no results are found.\n \"\"\"\n results = get_access_info_for_documents(db_session, [document_id])\n if not results:\n return None\n\n return results[0]\n\n\ndef get_access_info_for_documents(\n db_session: Session,\n document_ids: list[str],\n) -> Sequence[tuple[str, list[str | None], bool]]:\n \"\"\"Gets back all relevant access info for the given documents. This includes\n the user_ids for cc pairs that the document is associated with + whether any\n of the associated cc pairs are intending to make the document globally public.\n Returns the list where each element contains:\n - Document ID (which is also the ID of the DocumentByConnectorCredentialPair)\n - List of emails of Onyx users with direct access to the doc (includes a \"None\" element if\n the connector was set up by an admin when auth was off\n - bool for whether the document is public (the document later can also be marked public by\n automatic permission sync step)\n \"\"\"\n stmt = select(\n DocumentByConnectorCredentialPair.id,\n func.array_agg(func.coalesce(User.email, null())).label(\"user_emails\"),\n func.bool_or(ConnectorCredentialPair.access_type == AccessType.PUBLIC).label(\n \"public_doc\"\n ),\n ).where(DocumentByConnectorCredentialPair.id.in_(document_ids))\n\n stmt = (\n stmt.join(\n Credential,\n DocumentByConnectorCredentialPair.credential_id == Credential.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .outerjoin(\n User,\n and_(\n Credential.user_id == User.id,\n ConnectorCredentialPair.access_type != AccessType.SYNC,\n ),\n )\n # don't include CC pairs that are being deleted\n # NOTE: CC pairs can never go from DELETING to any other state -> it's safe to ignore them\n .where(ConnectorCredentialPair.status != ConnectorCredentialPairStatus.DELETING)\n .group_by(DocumentByConnectorCredentialPair.id)\n )\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef upsert_documents(\n db_session: Session,\n document_metadata_batch: list[DocumentMetadata],\n initial_boost: int = DEFAULT_BOOST,\n) -> None:\n \"\"\"NOTE: this function is Postgres specific. Not all DBs support the ON CONFLICT clause.\n Also note, this function should not be used for updating documents, only creating and\n ensuring that it exists. It IGNORES the doc_updated_at field\"\"\"\n seen_documents: dict[str, DocumentMetadata] = {}\n for document_metadata in document_metadata_batch:\n doc_id = document_metadata.document_id\n if doc_id not in seen_documents:\n seen_documents[doc_id] = document_metadata\n\n if not seen_documents:\n logger.info(\"No documents to upsert. Skipping.\")\n return\n\n includes_permissions = any(doc.external_access for doc in seen_documents.values())\n\n insert_stmt = insert(DbDocument).values(\n [\n model_to_dict(\n DbDocument(\n id=doc.document_id,\n from_ingestion_api=doc.from_ingestion_api,\n boost=initial_boost,\n hidden=False,\n semantic_id=doc.semantic_identifier,\n link=doc.first_link,\n doc_updated_at=None, # this is intentional\n last_modified=datetime.now(timezone.utc),\n primary_owners=doc.primary_owners,\n secondary_owners=doc.secondary_owners,\n kg_stage=KGStage.NOT_STARTED,\n parent_hierarchy_node_id=doc.parent_hierarchy_node_id,\n **(\n {\n \"external_user_emails\": list(\n doc.external_access.external_user_emails\n ),\n \"external_user_group_ids\": list(\n doc.external_access.external_user_group_ids\n ),\n \"is_public\": doc.external_access.is_public,\n }\n if doc.external_access\n else {}\n ),\n doc_metadata=doc.doc_metadata,\n )\n )\n for doc in seen_documents.values()\n ]\n )\n\n update_set = {\n \"from_ingestion_api\": insert_stmt.excluded.from_ingestion_api,\n \"boost\": insert_stmt.excluded.boost,\n \"hidden\": insert_stmt.excluded.hidden,\n \"semantic_id\": insert_stmt.excluded.semantic_id,\n \"link\": insert_stmt.excluded.link,\n \"primary_owners\": insert_stmt.excluded.primary_owners,\n \"secondary_owners\": insert_stmt.excluded.secondary_owners,\n \"doc_metadata\": insert_stmt.excluded.doc_metadata,\n \"parent_hierarchy_node_id\": insert_stmt.excluded.parent_hierarchy_node_id,\n }\n if includes_permissions:\n # Use COALESCE to preserve existing permissions when new values are NULL.\n # This prevents subsequent indexing runs (which don't fetch permissions)\n # from overwriting permissions set by permission sync jobs.\n update_set.update(\n {\n \"external_user_emails\": func.coalesce(\n insert_stmt.excluded.external_user_emails,\n DbDocument.external_user_emails,\n ),\n \"external_user_group_ids\": func.coalesce(\n insert_stmt.excluded.external_user_group_ids,\n DbDocument.external_user_group_ids,\n ),\n \"is_public\": func.coalesce(\n insert_stmt.excluded.is_public,\n DbDocument.is_public,\n ),\n }\n )\n on_conflict_stmt = insert_stmt.on_conflict_do_update(\n index_elements=[\"id\"],\n set_=update_set, # Conflict target\n )\n db_session.execute(on_conflict_stmt)\n db_session.commit()\n\n\ndef upsert_document_by_connector_credential_pair(\n db_session: Session, connector_id: int, credential_id: int, document_ids: list[str]\n) -> None:\n \"\"\"NOTE: this function is Postgres specific. Not all DBs support the ON CONFLICT clause.\"\"\"\n if not document_ids:\n logger.info(\"`document_ids` is empty. Skipping.\")\n return\n\n insert_stmt = insert(DocumentByConnectorCredentialPair).values(\n [\n model_to_dict(\n DocumentByConnectorCredentialPair(\n id=doc_id,\n connector_id=connector_id,\n credential_id=credential_id,\n has_been_indexed=False,\n )\n )\n for doc_id in document_ids\n ]\n )\n # this must be `on_conflict_do_nothing` rather than `on_conflict_do_update`\n # since we don't want to update the `has_been_indexed` field for documents\n # that already exist\n on_conflict_stmt = insert_stmt.on_conflict_do_nothing()\n db_session.execute(on_conflict_stmt)\n db_session.commit()\n\n\ndef mark_document_as_indexed_for_cc_pair__no_commit(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n document_ids: Iterable[str],\n) -> None:\n \"\"\"Should be called only after a successful index operation for a batch.\"\"\"\n db_session.execute(\n update(DocumentByConnectorCredentialPair)\n .where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n DocumentByConnectorCredentialPair.id.in_(document_ids),\n )\n )\n .values(has_been_indexed=True)\n )\n\n\ndef update_docs_updated_at__no_commit(\n ids_to_new_updated_at: dict[str, datetime],\n db_session: Session,\n) -> None:\n doc_ids = list(ids_to_new_updated_at.keys())\n documents_to_update = (\n db_session.query(DbDocument).filter(DbDocument.id.in_(doc_ids)).all()\n )\n\n for document in documents_to_update:\n document.doc_updated_at = ids_to_new_updated_at[document.id]\n\n\ndef update_docs_last_modified__no_commit(\n document_ids: list[str],\n db_session: Session,\n) -> None:\n documents_to_update = (\n db_session.query(DbDocument).filter(DbDocument.id.in_(document_ids)).all()\n )\n\n now = datetime.now(timezone.utc)\n for doc in documents_to_update:\n doc.last_modified = now\n\n\ndef update_docs_chunk_count__no_commit(\n document_ids: list[str],\n doc_id_to_chunk_count: dict[str, int],\n db_session: Session,\n) -> None:\n documents_to_update = (\n db_session.query(DbDocument).filter(DbDocument.id.in_(document_ids)).all()\n )\n for doc in documents_to_update:\n doc.chunk_count = doc_id_to_chunk_count[doc.id]\n\n\ndef mark_document_as_modified(\n document_id: str,\n db_session: Session,\n) -> None:\n stmt = select(DbDocument).where(DbDocument.id == document_id)\n doc = db_session.scalar(stmt)\n if doc is None:\n raise ValueError(f\"No document with ID: {document_id}\")\n\n # update last_synced\n doc.last_modified = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef mark_document_as_synced(document_id: str, db_session: Session) -> None:\n stmt = select(DbDocument).where(DbDocument.id == document_id)\n doc = db_session.scalar(stmt)\n if doc is None:\n raise ValueError(f\"No document with ID: {document_id}\")\n\n # update last_synced\n doc.last_synced = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef delete_document_by_connector_credential_pair__no_commit(\n db_session: Session,\n document_id: str,\n connector_credential_pair_identifier: (\n ConnectorCredentialPairIdentifier | None\n ) = None,\n) -> None:\n \"\"\"Deletes a single document by cc pair relationship entry.\n Foreign key rows are left in place.\n The implicit assumption is that the document itself still has other cc_pair\n references and needs to continue existing.\n \"\"\"\n delete_documents_by_connector_credential_pair__no_commit(\n db_session=db_session,\n document_ids=[document_id],\n connector_credential_pair_identifier=connector_credential_pair_identifier,\n )\n\n\ndef delete_documents_by_connector_credential_pair__no_commit(\n db_session: Session,\n document_ids: list[str],\n connector_credential_pair_identifier: (\n ConnectorCredentialPairIdentifier | None\n ) = None,\n) -> None:\n \"\"\"This deletes just the document by cc pair entries for a particular cc pair.\n Foreign key rows are left in place.\n The implicit assumption is that the document itself still has other cc_pair\n references and needs to continue existing.\n \"\"\"\n stmt = delete(DocumentByConnectorCredentialPair).where(\n DocumentByConnectorCredentialPair.id.in_(document_ids)\n )\n if connector_credential_pair_identifier:\n stmt = stmt.where(\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == connector_credential_pair_identifier.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == connector_credential_pair_identifier.credential_id,\n )\n )\n db_session.execute(stmt)\n\n\ndef delete_all_documents_by_connector_credential_pair__no_commit(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n) -> None:\n \"\"\"Deletes all document by connector credential pair entries for a specific connector and credential.\n This is primarily used during connector deletion to ensure all references are removed\n before deleting the connector itself. This is crucial because connector_id is part of the\n primary key in DocumentByConnectorCredentialPair, and attempting to delete the Connector\n would otherwise try to set the foreign key to NULL, which fails for primary keys.\n\n NOTE: Does not commit the transaction, this must be done by the caller.\n \"\"\"\n stmt = delete(DocumentByConnectorCredentialPair).where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n )\n db_session.execute(stmt)\n\n\ndef delete_documents__no_commit(db_session: Session, document_ids: list[str]) -> None:\n db_session.execute(delete(DbDocument).where(DbDocument.id.in_(document_ids)))\n\n\ndef delete_documents_complete__no_commit(\n db_session: Session, document_ids: list[str]\n) -> None:\n \"\"\"This completely deletes the documents from the db, including all foreign key relationships\"\"\"\n\n # Start with the kg references\n\n delete_from_kg_relationships__no_commit(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n delete_from_kg_entities__no_commit(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n delete_from_kg_relationships_extraction_staging__no_commit(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n delete_from_kg_entities_extraction_staging__no_commit(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n # Continue with deleting the chunk stats for the documents\n delete_chunk_stats_by_connector_credential_pair__no_commit(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n delete_documents_by_connector_credential_pair__no_commit(db_session, document_ids)\n delete_document_feedback_for_documents__no_commit(\n document_ids=document_ids, db_session=db_session\n )\n delete_document_tags_for_documents__no_commit(\n document_ids=document_ids, db_session=db_session\n )\n delete_documents__no_commit(db_session, document_ids)\n\n\ndef delete_all_documents_for_connector_credential_pair(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n timeout: int = ONE_HOUR_IN_SECONDS,\n) -> None:\n \"\"\"Delete all documents for a given connector credential pair.\n This will delete all documents and their associated data (chunks, feedback, tags, etc.)\n\n NOTE: a bit inefficient, but it's not a big deal since this is done rarely - only during\n an index swap. If we wanted to make this more efficient, we could use a single delete\n statement + cascade.\n \"\"\"\n batch_size = 1000\n start_time = time.monotonic()\n\n while True:\n # Get document IDs in batches\n stmt = (\n select(DocumentByConnectorCredentialPair.id)\n .where(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DocumentByConnectorCredentialPair.credential_id == credential_id,\n )\n .limit(batch_size)\n )\n document_ids = db_session.scalars(stmt).all()\n\n if not document_ids:\n break\n\n delete_documents_complete__no_commit(\n db_session=db_session, document_ids=list(document_ids)\n )\n db_session.commit()\n\n if time.monotonic() - start_time > timeout:\n raise RuntimeError(\"Timeout reached while deleting documents\")\n\n\ndef acquire_document_locks(db_session: Session, document_ids: list[str]) -> bool:\n \"\"\"Acquire locks for the specified documents. Ideally this shouldn't be\n called with large list of document_ids (an exception could be made if the\n length of holding the lock is very short).\n\n Will simply raise an exception if any of the documents are already locked.\n This prevents deadlocks (assuming that the caller passes in all required\n document IDs in a single call).\n \"\"\"\n stmt = (\n select(DbDocument.id)\n .where(DbDocument.id.in_(document_ids))\n .with_for_update(nowait=True)\n )\n # will raise exception if any of the documents are already locked\n documents = db_session.scalars(stmt).all()\n\n # make sure we found every document\n if len(documents) != len(set(document_ids)):\n logger.warning(\"Didn't find row for all specified document IDs. Aborting.\")\n return False\n\n return True\n\n\n_NUM_LOCK_ATTEMPTS = 10\n_LOCK_RETRY_DELAY = 10\n\n\n@contextlib.contextmanager\ndef prepare_to_modify_documents(\n db_session: Session, document_ids: list[str], retry_delay: int = _LOCK_RETRY_DELAY\n) -> Generator[TransactionalContext, None, None]:\n \"\"\"Try and acquire locks for the documents to prevent other jobs from\n modifying them at the same time (e.g. avoid race conditions). This should be\n called ahead of any modification to Vespa. Locks should be released by the\n caller as soon as updates are complete by finishing the transaction.\n\n NOTE: only one commit is allowed within the context manager returned by this function.\n Multiple commits will result in a sqlalchemy.exc.InvalidRequestError.\n NOTE: this function will commit any existing transaction.\n \"\"\"\n\n db_session.commit() # ensure that we're not in a transaction\n\n lock_acquired = False\n for i in range(_NUM_LOCK_ATTEMPTS):\n try:\n with db_session.begin() as transaction:\n lock_acquired = acquire_document_locks(\n db_session=db_session, document_ids=document_ids\n )\n if lock_acquired:\n yield transaction\n break\n except OperationalError as e:\n logger.warning(\n f\"Failed to acquire locks for documents on attempt {i}, retrying. Error: {e}\"\n )\n\n time.sleep(retry_delay)\n\n if not lock_acquired:\n raise RuntimeError(\n f\"Failed to acquire locks after {_NUM_LOCK_ATTEMPTS} attempts for documents: {document_ids}\"\n )\n\n\ndef get_ingestion_documents(\n db_session: Session,\n) -> list[DbDocument]:\n # TODO add the option to filter by DocumentSource\n stmt = select(DbDocument).where(DbDocument.from_ingestion_api.is_(True))\n documents = db_session.execute(stmt).scalars().all()\n return list(documents)\n\n\ndef get_documents_by_cc_pair(\n cc_pair_id: int,\n db_session: Session,\n) -> list[DbDocument]:\n return (\n db_session.query(DbDocument)\n .join(\n DocumentByConnectorCredentialPair,\n DbDocument.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .filter(ConnectorCredentialPair.id == cc_pair_id)\n .all()\n )\n\n\ndef get_document(\n document_id: str,\n db_session: Session,\n) -> DbDocument | None:\n stmt = select(DbDocument).where(DbDocument.id == document_id)\n doc: DbDocument | None = db_session.execute(stmt).scalar_one_or_none()\n return doc\n\n\ndef get_cc_pairs_for_document(\n db_session: Session,\n document_id: str,\n) -> list[ConnectorCredentialPair]:\n stmt = (\n select(ConnectorCredentialPair)\n .join(\n DocumentByConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .where(DocumentByConnectorCredentialPair.id == document_id)\n )\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef get_document_sources(\n db_session: Session,\n document_ids: list[str],\n) -> dict[str, DocumentSource]:\n \"\"\"Gets the sources for a list of document IDs.\n Returns a dictionary mapping document ID to its source.\n If a document has multiple sources (multiple CC pairs), returns the first one found.\n \"\"\"\n stmt = (\n select(\n DocumentByConnectorCredentialPair.id,\n Connector.source,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n Connector,\n ConnectorCredentialPair.connector_id == Connector.id,\n )\n .where(DocumentByConnectorCredentialPair.id.in_(document_ids))\n .distinct()\n )\n\n results = db_session.execute(stmt).all()\n return {doc_id: source for doc_id, source in results}\n\n\ndef fetch_chunk_counts_for_documents(\n document_ids: list[str],\n db_session: Session,\n) -> list[tuple[str, int]]:\n \"\"\"\n Return a list of (document_id, chunk_count) tuples.\n If a document_id is not found in the database, it will be returned with a chunk_count of 0.\n \"\"\"\n stmt = select(DbDocument.id, DbDocument.chunk_count).where(\n DbDocument.id.in_(document_ids)\n )\n\n results = db_session.execute(stmt).all()\n\n # Create a dictionary of document_id to chunk_count\n chunk_counts = {str(row.id): row.chunk_count or 0 for row in results}\n\n # Return a list of tuples, preserving `None` for documents not found or with\n # an unknown chunk count. Callers should handle the `None` case and fall\n # back to an existence check against the vector DB if necessary.\n return [(doc_id, chunk_counts.get(doc_id, 0)) for doc_id in document_ids]\n\n\ndef fetch_chunk_count_for_document(\n document_id: str,\n db_session: Session,\n) -> int | None:\n stmt = select(DbDocument.chunk_count).where(DbDocument.id == document_id)\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef get_unprocessed_kg_document_batch_for_connector(\n db_session: Session,\n connector_id: int,\n kg_coverage_start: datetime,\n kg_max_coverage_days: int,\n batch_size: int = 100,\n) -> list[DbDocument]:\n \"\"\"\n Retrieves a batch of documents that have not been processed for knowledge graph extraction.\n Args:\n db_session (Session): The database session to use\n connector_id (int): The ID of the connector to get documents for\n batch_size (int): The maximum number of documents to retrieve\n Returns:\n list[DbDocument]: List of documents that need KG processing\n \"\"\"\n\n stmt = (\n select(DbDocument)\n .join(\n DocumentByConnectorCredentialPair,\n DbDocument.id == DocumentByConnectorCredentialPair.id,\n )\n .where(\n and_(\n DocumentByConnectorCredentialPair.connector_id == connector_id,\n DbDocument.doc_updated_at\n >= max(\n kg_coverage_start,\n datetime.now() - timedelta(days=kg_max_coverage_days),\n ),\n or_(\n DbDocument.kg_stage.is_(None),\n DbDocument.kg_stage == KGStage.NOT_STARTED,\n DbDocument.doc_updated_at > DbDocument.kg_processing_time,\n ),\n )\n )\n .distinct()\n .limit(batch_size)\n )\n\n documents = db_session.scalars(stmt).all()\n db_session.flush()\n\n return list(documents)\n\n\ndef get_kg_extracted_document_ids(db_session: Session) -> list[str]:\n \"\"\"\n Retrieves all document IDs where kg_stage is EXTRACTED.\n Args:\n db_session (Session): The database session to use\n Returns:\n list[str]: List of document IDs that have been KG processed\n \"\"\"\n stmt = select(DbDocument.id).where(DbDocument.kg_stage == KGStage.EXTRACTED)\n\n return list(db_session.scalars(stmt).all())\n\n\ndef update_document_kg_info(\n db_session: Session, document_id: str, kg_stage: KGStage\n) -> None:\n \"\"\"Updates the knowledge graph related information for a document.\n Args:\n db_session (Session): The database session to use\n document_id (str): The ID of the document to update\n kg_stage (KGStage): The stage of the knowledge graph processing for the document\n Raises:\n ValueError: If the document with the given ID is not found\n \"\"\"\n stmt = (\n update(DbDocument)\n .where(DbDocument.id == document_id)\n .values(\n kg_stage=kg_stage,\n kg_processing_time=datetime.now(timezone.utc),\n )\n )\n db_session.execute(stmt)\n\n\ndef update_document_kg_stage(\n db_session: Session,\n document_id: str,\n kg_stage: KGStage,\n) -> None:\n stmt = (\n update(DbDocument).where(DbDocument.id == document_id).values(kg_stage=kg_stage)\n )\n db_session.execute(stmt)\n db_session.flush()\n\n\ndef get_all_kg_extracted_documents_info(\n db_session: Session,\n) -> list[str]:\n \"\"\"Retrieves the knowledge graph data for all documents that have been processed.\n Args:\n db_session (Session): The database session to use\n Returns:\n List[Tuple[str, dict]]: A list of tuples containing:\n - str: The document ID\n - dict: The KG data containing 'entities', 'relationships', and 'terms'\n Only returns documents where kg_stage is EXTRACTED\n \"\"\"\n stmt = (\n select(DbDocument.id)\n .where(DbDocument.kg_stage == KGStage.EXTRACTED)\n .order_by(DbDocument.id)\n )\n\n results = db_session.execute(stmt).all()\n return [str(doc_id) for doc_id in results]\n\n\ndef get_base_llm_doc_information(\n db_session: Session, document_ids: list[str]\n) -> list[str]:\n stmt = select(DbDocument).where(DbDocument.id.in_(document_ids))\n results = db_session.execute(stmt).all()\n\n documents = []\n\n for doc_nr, doc in enumerate(results):\n bare_doc = doc[0]\n documents.append(\n f\"\"\"* [{bare_doc.semantic_id}]({bare_doc.link}) ({bare_doc.doc_updated_at})\"\"\"\n )\n\n return documents[:KG_SIMPLE_ANSWER_MAX_DISPLAYED_SOURCES]\n\n\ndef get_document_updated_at(\n document_id: str,\n db_session: Session,\n) -> datetime | None:\n \"\"\"Retrieves the doc_updated_at timestamp for a given document ID.\n Args:\n document_id (str): The ID of the document to query\n db_session (Session): The database session to use\n Returns:\n Optional[datetime]: The doc_updated_at timestamp if found, None if document doesn't exist\n \"\"\"\n\n stmt = select(DbDocument.doc_updated_at).where(DbDocument.id == document_id)\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef reset_all_document_kg_stages(db_session: Session) -> int:\n \"\"\"Reset the KG stage of all documents that are not in NOT_STARTED state to NOT_STARTED.\n\n Args:\n db_session (Session): The database session to use\n\n Returns:\n int: Number of documents that were reset\n \"\"\"\n stmt = (\n update(DbDocument)\n .where(DbDocument.kg_stage != KGStage.NOT_STARTED)\n .values(kg_stage=KGStage.NOT_STARTED)\n )\n result = db_session.execute(stmt)\n\n # The hasattr check is needed for type checking, even though rowcount\n # is guaranteed to exist at runtime for UPDATE operations\n return result.rowcount if hasattr(result, \"rowcount\") else 0\n\n\ndef update_document_kg_stages(\n db_session: Session, source_stage: KGStage, target_stage: KGStage\n) -> int:\n \"\"\"Reset the KG stage only of documents back to NOT_STARTED.\n Part of reset flow for documents that have been extracted but not clustered.\n\n Args:\n db_session (Session): The database session to use\n\n Returns:\n int: Number of documents that were reset\n \"\"\"\n stmt = (\n update(DbDocument)\n .where(DbDocument.kg_stage == source_stage)\n .values(kg_stage=target_stage)\n )\n result = db_session.execute(stmt)\n # The hasattr check is needed for type checking, even though rowcount\n # is guaranteed to exist at runtime for UPDATE operations\n return result.rowcount if hasattr(result, \"rowcount\") else 0\n\n\ndef get_skipped_kg_documents(db_session: Session) -> list[str]:\n \"\"\"\n Retrieves all document IDs where kg_stage is SKIPPED.\n Args:\n db_session (Session): The database session to use\n Returns:\n list[str]: List of document IDs that have been skipped in KG processing\n \"\"\"\n stmt = select(DbDocument.id).where(DbDocument.kg_stage == KGStage.SKIPPED)\n\n return list(db_session.scalars(stmt).all())\n\n\n# def get_kg_doc_info_for_entity_name(\n# db_session: Session, document_id: str, entity_type: str\n# ) -> KGEntityDocInfo:\n# \"\"\"\n# Get the semantic ID and the link for an entity name.\n# \"\"\"\n\n# result = (\n# db_session.query(Document.semantic_id, Document.link)\n# .filter(Document.id == document_id)\n# .first()\n# )\n\n# if result is None:\n# return KGEntityDocInfo(\n# doc_id=None,\n# doc_semantic_id=None,\n# doc_link=None,\n# semantic_entity_name=f\"{entity_type}:{document_id}\",\n# semantic_linked_entity_name=f\"{entity_type}:{document_id}\",\n# )\n\n# return KGEntityDocInfo(\n# doc_id=document_id,\n# doc_semantic_id=result[0],\n# doc_link=result[1],\n# semantic_entity_name=f\"{entity_type.upper()}:{result[0]}\",\n# semantic_linked_entity_name=f\"[{entity_type.upper()}:{result[0]}]({result[1]})\",\n# )\n\n\ndef check_for_documents_needing_kg_processing(\n db_session: Session, kg_coverage_start: datetime, kg_max_coverage_days: int\n) -> bool:\n \"\"\"Check if there are any documents that need KG processing.\n\n A document needs KG processing if:\n 1. It is associated with a connector that has kg_processing_enabled = true\n 2. AND either:\n - Its kg_stage is NOT_STARTED or NULL\n - OR its last_updated timestamp is greater than its kg_processing_time\n\n Args:\n db_session (Session): The database session to use\n\n Returns:\n bool: True if there are any documents needing KG processing, False otherwise\n \"\"\"\n\n stmt = (\n select(1)\n .select_from(DbDocument)\n .join(\n DocumentByConnectorCredentialPair,\n DbDocument.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n Connector,\n DocumentByConnectorCredentialPair.connector_id == Connector.id,\n )\n .where(\n and_(\n Connector.kg_processing_enabled.is_(True),\n DbDocument.doc_updated_at\n >= max(\n kg_coverage_start,\n datetime.now() - timedelta(days=kg_max_coverage_days),\n ),\n or_(\n DbDocument.kg_stage.is_(None),\n DbDocument.kg_stage == KGStage.NOT_STARTED,\n DbDocument.doc_updated_at > DbDocument.kg_processing_time,\n ),\n )\n )\n .exists()\n )\n\n return db_session.execute(select(stmt)).scalar() or False\n\n\ndef check_for_documents_needing_kg_clustering(db_session: Session) -> bool:\n \"\"\"Check if there are any documents that need KG clustering.\n\n A document needs KG clustering if:\n 1. It is associated with a connector that has kg_processing_enabled = true\n 2. AND either:\n - Its kg_stage is EXTRACTED\n - OR its last_updated timestamp is greater than its kg_processing_time\n\n Args:\n db_session (Session): The database session to use\n\n Returns:\n bool: True if there are any documents needing KG clustering, False otherwise\n \"\"\"\n stmt = (\n select(1)\n .select_from(DbDocument)\n .join(\n DocumentByConnectorCredentialPair,\n DbDocument.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n Connector,\n ConnectorCredentialPair.connector_id == Connector.id,\n )\n .where(\n and_(\n Connector.kg_processing_enabled.is_(True),\n ConnectorCredentialPair.status\n != ConnectorCredentialPairStatus.DELETING,\n or_(\n DbDocument.kg_stage == KGStage.EXTRACTED,\n DbDocument.last_modified > DbDocument.kg_processing_time,\n ),\n )\n )\n .exists()\n )\n\n return db_session.execute(select(stmt)).scalar() or False\n\n\ndef get_document_kg_entities_and_relationships(\n db_session: Session, document_id: str\n) -> tuple[list[KGEntity], list[KGRelationship]]:\n \"\"\"\n Get the KG entities and relationships that references the document.\n \"\"\"\n entities = (\n db_session.query(KGEntity).filter(KGEntity.document_id == document_id).all()\n )\n if not entities:\n return [], []\n entity_id_names = [entity.id_name for entity in entities]\n\n relationships = (\n db_session.query(KGRelationship)\n .filter(\n or_(\n KGRelationship.source_node.in_(entity_id_names),\n KGRelationship.target_node.in_(entity_id_names),\n KGRelationship.source_document == document_id,\n )\n )\n .all()\n )\n return entities, relationships\n\n\ndef get_num_chunks_for_document(db_session: Session, document_id: str) -> int:\n stmt = select(DbDocument.chunk_count).where(DbDocument.id == document_id)\n return db_session.execute(stmt).scalar_one_or_none() or 0\n\n\ndef update_document_metadata__no_commit(\n db_session: Session,\n document_id: str,\n doc_metadata: dict[str, Any],\n) -> None:\n \"\"\"Update the doc_metadata field for a document.\n\n Note: Does not commit. Caller is responsible for committing.\n\n Args:\n db_session: Database session\n document_id: The ID of the document to update\n doc_metadata: The new metadata dictionary to set\n \"\"\"\n stmt = (\n update(DbDocument)\n .where(DbDocument.id == document_id)\n .values(doc_metadata=doc_metadata)\n )\n db_session.execute(stmt)\n\n\ndef delete_document_by_id__no_commit(\n db_session: Session,\n document_id: str,\n) -> None:\n \"\"\"Delete a single document and its connector credential pair relationships.\n\n Note: Does not commit. Caller is responsible for committing.\n\n This uses delete_documents_complete__no_commit which handles\n all foreign key relationships (KG entities, relationships, chunk stats,\n cc pair associations, feedback, tags).\n \"\"\"\n delete_documents_complete__no_commit(db_session, [document_id])\n" }, { "path": "backend/onyx/db/document_access.py", "content": "\"\"\"\nDocument access filtering utilities.\n\nThis module provides reusable access filtering logic for documents based on:\n- Connector access type (PUBLIC vs SYNC)\n- Document-level public flag\n- User email matching external_user_emails\n- User group overlap with external_user_group_ids\n\nThis is a standalone module to avoid circular imports between document.py and persona.py.\n\"\"\"\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import any_\nfrom sqlalchemy import cast\nfrom sqlalchemy import or_\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import String\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql.elements import ColumnElement\n\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Document\nfrom onyx.db.models import DocumentByConnectorCredentialPair\n\n\ndef apply_document_access_filter(\n stmt: Select,\n user_email: str | None,\n external_group_ids: list[str],\n) -> Select:\n \"\"\"\n Apply document access filtering to a query.\n\n This joins with DocumentByConnectorCredentialPair and ConnectorCredentialPair to:\n 1. Check if the document is from a PUBLIC connector (access_type = PUBLIC)\n 2. Check document-level permissions (is_public, external_user_emails, external_user_group_ids)\n 3. Exclude documents from cc_pairs that are being deleted\n\n Args:\n stmt: The SELECT statement to modify (must be selecting from Document)\n user_email: The user's email for permission checking\n external_group_ids: List of external group IDs the user belongs to\n\n Returns:\n Modified SELECT statement with access filtering applied\n \"\"\"\n # Join to get cc_pair info for each document\n stmt = stmt.join(\n DocumentByConnectorCredentialPair,\n Document.id == DocumentByConnectorCredentialPair.id,\n ).join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n\n # Exclude documents from cc_pairs that are being deleted\n stmt = stmt.where(\n ConnectorCredentialPair.status != ConnectorCredentialPairStatus.DELETING\n )\n\n # Build access filters\n access_filters: list[ColumnElement[bool]] = [\n # Document is from a PUBLIC connector\n ConnectorCredentialPair.access_type == AccessType.PUBLIC,\n # Document is marked as public (e.g., \"Anyone with link\" in source)\n Document.is_public.is_(True),\n ]\n if user_email:\n access_filters.append(any_(Document.external_user_emails) == user_email)\n if external_group_ids:\n access_filters.append(\n Document.external_user_group_ids.overlap(\n cast(postgresql.array(external_group_ids), postgresql.ARRAY(String))\n )\n )\n\n stmt = stmt.where(or_(*access_filters))\n return stmt\n\n\ndef get_accessible_documents_by_ids(\n db_session: Session,\n document_ids: list[str],\n user_email: str | None,\n external_group_ids: list[str],\n) -> list[Document]:\n \"\"\"\n Fetch documents by IDs, filtering to only those the user has access to.\n\n Uses the same access filtering logic as other document queries:\n - Documents from PUBLIC connectors\n - Documents marked as public (e.g., \"Anyone with link\")\n - Documents where user email matches external_user_emails\n - Documents where user's groups overlap with external_user_group_ids\n\n Args:\n db_session: Database session\n document_ids: List of document IDs to fetch\n user_email: User's email for permission checking\n external_group_ids: List of external group IDs the user belongs to\n\n Returns:\n List of Document objects from the input that the user has access to\n \"\"\"\n if not document_ids:\n return []\n\n stmt = select(Document).where(Document.id.in_(document_ids))\n stmt = apply_document_access_filter(stmt, user_email, external_group_ids)\n # Use distinct to avoid duplicates when a document belongs to multiple cc_pairs\n stmt = stmt.distinct()\n return list(db_session.execute(stmt).scalars().all())\n" }, { "path": "backend/onyx/db/document_set.py", "content": "from collections.abc import Sequence\nfrom typing import cast\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import delete\nfrom sqlalchemy import exists\nfrom sqlalchemy import func\nfrom sqlalchemy import or_\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.db.connector_credential_pair import get_cc_pair_groups_for_ids\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.federated import create_federated_connector_document_set_mapping\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Document\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import DocumentSet as DocumentSetDBModel\nfrom onyx.db.models import DocumentSet__ConnectorCredentialPair\nfrom onyx.db.models import DocumentSet__UserGroup\nfrom onyx.db.models import FederatedConnector__DocumentSet\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserRole\nfrom onyx.server.features.document_set.models import DocumentSetCreationRequest\nfrom onyx.server.features.document_set.models import DocumentSetUpdateRequest\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\nlogger = setup_logger()\n\n\ndef _add_user_filters(stmt: Select, user: User, get_editable: bool = True) -> Select:\n if user.role == UserRole.ADMIN:\n return stmt\n\n stmt = stmt.distinct()\n DocumentSet__UG = aliased(DocumentSet__UserGroup)\n User__UG = aliased(User__UserGroup)\n \"\"\"\n Here we select cc_pairs by relation:\n User -> User__UserGroup -> DocumentSet__UserGroup -> DocumentSet\n \"\"\"\n stmt = stmt.outerjoin(DocumentSet__UG).outerjoin(\n User__UserGroup,\n User__UserGroup.user_group_id == DocumentSet__UG.user_group_id,\n )\n \"\"\"\n Filter DocumentSets by:\n - if the user is in the user_group that owns the DocumentSet\n - if the user is not a global_curator, they must also have a curator relationship\n to the user_group\n - if editing is being done, we also filter out DocumentSets that are owned by groups\n that the user isn't a curator for\n - if we are not editing, we show all DocumentSets in the groups the user is a curator\n for (as well as public DocumentSets)\n \"\"\"\n\n # Anonymous users only see public DocumentSets\n if user.is_anonymous:\n where_clause = DocumentSetDBModel.is_public == True # noqa: E712\n return stmt.where(where_clause)\n\n where_clause = User__UserGroup.user_id == user.id\n if user.role == UserRole.CURATOR and get_editable:\n where_clause &= User__UserGroup.is_curator == True # noqa: E712\n if get_editable:\n user_groups = select(User__UG.user_group_id).where(User__UG.user_id == user.id)\n if user.role == UserRole.CURATOR:\n user_groups = user_groups.where(User__UG.is_curator == True) # noqa: E712\n where_clause &= (\n ~exists()\n .where(DocumentSet__UG.document_set_id == DocumentSetDBModel.id)\n .where(~DocumentSet__UG.user_group_id.in_(user_groups))\n .correlate(DocumentSetDBModel)\n )\n where_clause |= DocumentSetDBModel.user_id == user.id\n else:\n where_clause |= DocumentSetDBModel.is_public == True # noqa: E712\n\n return stmt.where(where_clause)\n\n\ndef _delete_document_set_cc_pairs__no_commit(\n db_session: Session, document_set_id: int, is_current: bool | None = None\n) -> None:\n \"\"\"NOTE: does not commit transaction, this must be done by the caller\"\"\"\n stmt = delete(DocumentSet__ConnectorCredentialPair).where(\n DocumentSet__ConnectorCredentialPair.document_set_id == document_set_id\n )\n if is_current is not None:\n stmt = stmt.where(DocumentSet__ConnectorCredentialPair.is_current == is_current)\n db_session.execute(stmt)\n\n\ndef _mark_document_set_cc_pairs_as_outdated__no_commit(\n db_session: Session, document_set_id: int\n) -> None:\n \"\"\"NOTE: does not commit transaction, this must be done by the caller\"\"\"\n stmt = select(DocumentSet__ConnectorCredentialPair).where(\n DocumentSet__ConnectorCredentialPair.document_set_id == document_set_id\n )\n for row in db_session.scalars(stmt):\n row.is_current = False\n\n\ndef delete_document_set_privacy__no_commit(\n document_set_id: int, db_session: Session\n) -> None:\n \"\"\"No private document sets in Onyx MIT\"\"\"\n\n\ndef get_document_set_by_id_for_user(\n db_session: Session,\n document_set_id: int,\n user: User,\n get_editable: bool = True,\n) -> DocumentSetDBModel | None:\n stmt = (\n select(DocumentSetDBModel)\n .distinct()\n .options(selectinload(DocumentSetDBModel.federated_connectors))\n )\n stmt = stmt.where(DocumentSetDBModel.id == document_set_id)\n stmt = _add_user_filters(stmt=stmt, user=user, get_editable=get_editable)\n return db_session.scalar(stmt)\n\n\ndef get_document_set_by_id(\n db_session: Session,\n document_set_id: int,\n) -> DocumentSetDBModel | None:\n stmt = select(DocumentSetDBModel).distinct()\n stmt = stmt.where(DocumentSetDBModel.id == document_set_id)\n return db_session.scalar(stmt)\n\n\ndef get_document_set_by_name(\n db_session: Session, document_set_name: str\n) -> DocumentSetDBModel | None:\n return db_session.scalar(\n select(DocumentSetDBModel).where(DocumentSetDBModel.name == document_set_name)\n )\n\n\ndef get_document_sets_by_name(\n db_session: Session, document_set_names: list[str]\n) -> Sequence[DocumentSetDBModel]:\n return db_session.scalars(\n select(DocumentSetDBModel).where(\n DocumentSetDBModel.name.in_(document_set_names)\n )\n ).all()\n\n\ndef get_document_sets_by_ids(\n db_session: Session, document_set_ids: list[int]\n) -> Sequence[DocumentSetDBModel]:\n if not document_set_ids:\n return []\n return db_session.scalars(\n select(DocumentSetDBModel).where(DocumentSetDBModel.id.in_(document_set_ids))\n ).all()\n\n\ndef make_doc_set_private(\n document_set_id: int, # noqa: ARG001\n user_ids: list[UUID] | None,\n group_ids: list[int] | None,\n db_session: Session, # noqa: ARG001\n) -> None:\n # May cause error if someone switches down to MIT from EE\n if user_ids or group_ids:\n raise NotImplementedError(\"Onyx MIT does not support private Document Sets\")\n\n\ndef _check_if_cc_pairs_are_owned_by_groups(\n db_session: Session,\n cc_pair_ids: list[int],\n group_ids: list[int],\n) -> None:\n \"\"\"\n This function checks if the CC pairs are owned by the specified groups or public.\n If not, it raises a ValueError.\n \"\"\"\n group_cc_pair_relationships = get_cc_pair_groups_for_ids(\n db_session=db_session,\n cc_pair_ids=cc_pair_ids,\n )\n\n group_cc_pair_relationships_set = {\n (relationship.cc_pair_id, relationship.user_group_id)\n for relationship in group_cc_pair_relationships\n }\n\n missing_cc_pair_ids = []\n for cc_pair_id in cc_pair_ids:\n for group_id in group_ids:\n if (cc_pair_id, group_id) not in group_cc_pair_relationships_set:\n missing_cc_pair_ids.append(cc_pair_id)\n break\n\n if missing_cc_pair_ids:\n cc_pairs = get_connector_credential_pairs(\n db_session=db_session,\n ids=missing_cc_pair_ids,\n )\n for cc_pair in cc_pairs:\n if cc_pair.access_type == AccessType.PRIVATE:\n raise ValueError(\n f\"Connector Credential Pair with ID: '{cc_pair.id}' is not owned by the specified groups\"\n )\n\n\ndef insert_document_set(\n document_set_creation_request: DocumentSetCreationRequest,\n user_id: UUID | None,\n db_session: Session,\n) -> tuple[DocumentSetDBModel, list[DocumentSet__ConnectorCredentialPair]]:\n # Check if we have either CC pairs or federated connectors (or both)\n if (\n not document_set_creation_request.cc_pair_ids\n and not document_set_creation_request.federated_connectors\n ):\n raise ValueError(\"Cannot create a document set with no connectors\")\n\n if not document_set_creation_request.is_public:\n _check_if_cc_pairs_are_owned_by_groups(\n db_session=db_session,\n cc_pair_ids=document_set_creation_request.cc_pair_ids,\n group_ids=document_set_creation_request.groups or [],\n )\n\n new_document_set_row: DocumentSetDBModel\n ds_cc_pairs: list[DocumentSet__ConnectorCredentialPair]\n try:\n new_document_set_row = DocumentSetDBModel(\n name=document_set_creation_request.name,\n description=document_set_creation_request.description,\n user_id=user_id,\n is_public=document_set_creation_request.is_public,\n is_up_to_date=DISABLE_VECTOR_DB,\n time_last_modified_by_user=func.now(),\n )\n db_session.add(new_document_set_row)\n db_session.flush() # ensure the new document set gets assigned an ID\n\n # Create CC pair mappings\n ds_cc_pairs = [\n DocumentSet__ConnectorCredentialPair(\n document_set_id=new_document_set_row.id,\n connector_credential_pair_id=cc_pair_id,\n is_current=True,\n )\n for cc_pair_id in document_set_creation_request.cc_pair_ids\n ]\n db_session.add_all(ds_cc_pairs)\n\n # Create federated connector mappings\n from onyx.db.federated import create_federated_connector_document_set_mapping\n\n for fc_config in document_set_creation_request.federated_connectors:\n create_federated_connector_document_set_mapping(\n db_session=db_session,\n federated_connector_id=fc_config.federated_connector_id,\n document_set_id=new_document_set_row.id,\n entities=fc_config.entities,\n )\n\n versioned_private_doc_set_fn = fetch_versioned_implementation(\n \"onyx.db.document_set\", \"make_doc_set_private\"\n )\n\n # Private Document Sets\n versioned_private_doc_set_fn(\n document_set_id=new_document_set_row.id,\n user_ids=document_set_creation_request.users,\n group_ids=document_set_creation_request.groups,\n db_session=db_session,\n )\n\n db_session.commit()\n except Exception as e:\n db_session.rollback()\n logger.error(f\"Error creating document set: {e}\")\n raise\n\n return new_document_set_row, ds_cc_pairs\n\n\ndef update_document_set(\n db_session: Session,\n document_set_update_request: DocumentSetUpdateRequest,\n user: User,\n) -> tuple[DocumentSetDBModel, list[DocumentSet__ConnectorCredentialPair]]:\n \"\"\"If successful, this sets document_set_row.is_up_to_date = False.\n That will be processed via Celery in check_for_vespa_sync_task\n and trigger a long running background sync to Vespa.\n \"\"\"\n # Check if we have either CC pairs or federated connectors (or both)\n if (\n not document_set_update_request.cc_pair_ids\n and not document_set_update_request.federated_connectors\n ):\n raise ValueError(\"Cannot update a document set with no connectors\")\n\n if not document_set_update_request.is_public:\n _check_if_cc_pairs_are_owned_by_groups(\n db_session=db_session,\n cc_pair_ids=document_set_update_request.cc_pair_ids,\n group_ids=document_set_update_request.groups,\n )\n\n try:\n # update the description\n document_set_row = get_document_set_by_id_for_user(\n db_session=db_session,\n document_set_id=document_set_update_request.id,\n user=user,\n get_editable=True,\n )\n if document_set_row is None:\n raise ValueError(\n f\"No document set with ID '{document_set_update_request.id}'\"\n )\n if not document_set_row.is_up_to_date:\n raise ValueError(\n \"Cannot update document set while it is syncing. Please wait for it to finish syncing, and then try again.\"\n )\n\n document_set_row.description = document_set_update_request.description\n if not DISABLE_VECTOR_DB:\n document_set_row.is_up_to_date = False\n document_set_row.is_public = document_set_update_request.is_public\n document_set_row.time_last_modified_by_user = func.now()\n versioned_private_doc_set_fn = fetch_versioned_implementation(\n \"onyx.db.document_set\", \"make_doc_set_private\"\n )\n\n # Private Document Sets\n versioned_private_doc_set_fn(\n document_set_id=document_set_row.id,\n user_ids=document_set_update_request.users,\n group_ids=document_set_update_request.groups,\n db_session=db_session,\n )\n\n # update the attached CC pairs\n # first, mark all existing CC pairs as not current\n _mark_document_set_cc_pairs_as_outdated__no_commit(\n db_session=db_session, document_set_id=document_set_row.id\n )\n # add in rows for the new CC pairs\n ds_cc_pairs = [\n DocumentSet__ConnectorCredentialPair(\n document_set_id=document_set_update_request.id,\n connector_credential_pair_id=cc_pair_id,\n is_current=True,\n )\n for cc_pair_id in document_set_update_request.cc_pair_ids\n ]\n db_session.add_all(ds_cc_pairs)\n\n # Update federated connector mappings\n # Delete existing federated connector mappings for this document set\n delete_stmt = delete(FederatedConnector__DocumentSet).where(\n FederatedConnector__DocumentSet.document_set_id == document_set_row.id\n )\n db_session.execute(delete_stmt)\n\n # Create new federated connector mappings\n for fc_config in document_set_update_request.federated_connectors:\n create_federated_connector_document_set_mapping(\n db_session=db_session,\n federated_connector_id=fc_config.federated_connector_id,\n document_set_id=document_set_row.id,\n entities=fc_config.entities,\n )\n\n db_session.commit()\n except:\n db_session.rollback()\n raise\n\n return document_set_row, ds_cc_pairs\n\n\ndef mark_document_set_as_synced(document_set_id: int, db_session: Session) -> None:\n stmt = select(DocumentSetDBModel).where(DocumentSetDBModel.id == document_set_id)\n document_set = db_session.scalar(stmt)\n if document_set is None:\n raise ValueError(f\"No document set with ID: {document_set_id}\")\n\n # mark as up to date\n document_set.is_up_to_date = True\n # delete outdated relationship table rows\n _delete_document_set_cc_pairs__no_commit(\n db_session=db_session, document_set_id=document_set_id, is_current=False\n )\n db_session.commit()\n\n\ndef delete_document_set(\n document_set_row: DocumentSetDBModel, db_session: Session\n) -> None:\n # delete all relationships to CC pairs\n _delete_document_set_cc_pairs__no_commit(\n db_session=db_session, document_set_id=document_set_row.id\n )\n db_session.delete(document_set_row)\n db_session.commit()\n\n\ndef mark_document_set_as_to_be_deleted(\n db_session: Session,\n document_set_id: int,\n user: User,\n) -> None:\n \"\"\"Cleans up all document_set -> cc_pair relationships and marks the document set\n as needing an update. The actual document set row will be deleted by the background\n job which syncs these changes to Vespa.\"\"\"\n\n try:\n document_set_row = get_document_set_by_id_for_user(\n db_session=db_session,\n document_set_id=document_set_id,\n user=user,\n get_editable=True,\n )\n if document_set_row is None:\n error_msg = f\"Document set with ID: '{document_set_id}' does not exist \"\n if user is not None:\n error_msg += f\"or is not editable by user with email: '{user.email}'\"\n raise ValueError(error_msg)\n if not document_set_row.is_up_to_date:\n raise ValueError(\n \"Cannot delete document set while it is syncing. Please wait for it to finish syncing, and then try again.\"\n )\n\n # delete all relationships to CC pairs\n _delete_document_set_cc_pairs__no_commit(\n db_session=db_session, document_set_id=document_set_id\n )\n\n # delete all federated connector mappings so the cleanup task can fully\n # remove the document set once the Vespa sync completes\n delete_stmt = delete(FederatedConnector__DocumentSet).where(\n FederatedConnector__DocumentSet.document_set_id == document_set_id\n )\n db_session.execute(delete_stmt)\n\n # delete all private document set information\n versioned_delete_private_fn = fetch_versioned_implementation(\n \"onyx.db.document_set\", \"delete_document_set_privacy__no_commit\"\n )\n versioned_delete_private_fn(\n document_set_id=document_set_id, db_session=db_session\n )\n\n # mark the row as needing a sync, it will be deleted there since there\n # are no more relationships to cc pairs\n document_set_row.is_up_to_date = False\n db_session.commit()\n except:\n db_session.rollback()\n raise\n\n\ndef delete_document_set_cc_pair_relationship__no_commit(\n connector_id: int, credential_id: int, db_session: Session\n) -> int:\n \"\"\"Deletes all rows from DocumentSet__ConnectorCredentialPair where the\n connector_credential_pair_id matches the given cc_pair_id.\"\"\"\n delete_stmt = delete(DocumentSet__ConnectorCredentialPair).where(\n and_(\n ConnectorCredentialPair.connector_id == connector_id,\n ConnectorCredentialPair.credential_id == credential_id,\n DocumentSet__ConnectorCredentialPair.connector_credential_pair_id\n == ConnectorCredentialPair.id,\n )\n )\n result = db_session.execute(delete_stmt)\n return result.rowcount # type: ignore\n\n\ndef fetch_document_sets(\n user_id: UUID | None, # noqa: ARG001\n db_session: Session,\n include_outdated: bool = False,\n) -> list[tuple[DocumentSetDBModel, list[ConnectorCredentialPair]]]:\n \"\"\"Return is a list where each element contains a tuple of:\n 1. The document set itself\n 2. All CC pairs associated with the document set\"\"\"\n stmt = (\n select(DocumentSetDBModel, ConnectorCredentialPair)\n .join(\n DocumentSet__ConnectorCredentialPair,\n DocumentSetDBModel.id\n == DocumentSet__ConnectorCredentialPair.document_set_id,\n isouter=True, # outer join is needed to also fetch document sets with no cc pairs\n )\n .join(\n ConnectorCredentialPair,\n ConnectorCredentialPair.id\n == DocumentSet__ConnectorCredentialPair.connector_credential_pair_id,\n isouter=True, # outer join is needed to also fetch document sets with no cc pairs\n )\n )\n if not include_outdated:\n stmt = stmt.where(\n or_(\n DocumentSet__ConnectorCredentialPair.is_current == True, # noqa: E712\n # `None` handles case where no CC Pairs exist for a Document Set\n DocumentSet__ConnectorCredentialPair.is_current.is_(None),\n )\n )\n\n results = cast(\n list[tuple[DocumentSetDBModel, ConnectorCredentialPair | None]],\n db_session.execute(stmt).all(),\n )\n\n aggregated_results: dict[\n int, tuple[DocumentSetDBModel, list[ConnectorCredentialPair]]\n ] = {}\n for document_set, cc_pair in results:\n if document_set.id not in aggregated_results:\n aggregated_results[document_set.id] = (\n document_set,\n [cc_pair] if cc_pair else [],\n )\n else:\n if cc_pair:\n aggregated_results[document_set.id][1].append(cc_pair)\n\n return [\n (document_set, cc_pairs)\n for document_set, cc_pairs in aggregated_results.values()\n ]\n\n\ndef fetch_all_document_sets_for_user(\n db_session: Session,\n user: User,\n get_editable: bool = True,\n) -> Sequence[DocumentSetDBModel]:\n stmt = (\n select(DocumentSetDBModel)\n .distinct()\n .options(\n selectinload(DocumentSetDBModel.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSetDBModel.users),\n selectinload(DocumentSetDBModel.groups),\n selectinload(DocumentSetDBModel.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n )\n )\n stmt = _add_user_filters(stmt, user, get_editable=get_editable)\n return db_session.scalars(stmt).unique().all()\n\n\ndef fetch_documents_for_document_set_paginated(\n document_set_id: int,\n db_session: Session,\n current_only: bool = True,\n last_document_id: str | None = None,\n limit: int = 100,\n) -> tuple[Sequence[Document], str | None]:\n stmt = (\n select(Document)\n .join(\n DocumentByConnectorCredentialPair,\n DocumentByConnectorCredentialPair.id == Document.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n ConnectorCredentialPair.connector_id\n == DocumentByConnectorCredentialPair.connector_id,\n ConnectorCredentialPair.credential_id\n == DocumentByConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n DocumentSet__ConnectorCredentialPair,\n DocumentSet__ConnectorCredentialPair.connector_credential_pair_id\n == ConnectorCredentialPair.id,\n )\n .join(\n DocumentSetDBModel,\n DocumentSetDBModel.id\n == DocumentSet__ConnectorCredentialPair.document_set_id,\n )\n .where(DocumentSetDBModel.id == document_set_id)\n .order_by(Document.id)\n .limit(limit)\n )\n if last_document_id is not None:\n stmt = stmt.where(Document.id > last_document_id)\n if current_only:\n stmt = stmt.where(\n DocumentSet__ConnectorCredentialPair.is_current == True # noqa: E712\n )\n stmt = stmt.distinct()\n\n documents = db_session.scalars(stmt).all()\n return documents, documents[-1].id if documents else None\n\n\ndef construct_document_id_select_by_docset(\n document_set_id: int,\n current_only: bool = True,\n) -> Select:\n \"\"\"This returns a statement that should be executed using\n .yield_per() to minimize overhead. The primary consumers of this function\n are background processing task generators.\"\"\"\n\n stmt = (\n select(Document.id)\n .join(\n DocumentByConnectorCredentialPair,\n DocumentByConnectorCredentialPair.id == Document.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n ConnectorCredentialPair.connector_id\n == DocumentByConnectorCredentialPair.connector_id,\n ConnectorCredentialPair.credential_id\n == DocumentByConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n DocumentSet__ConnectorCredentialPair,\n DocumentSet__ConnectorCredentialPair.connector_credential_pair_id\n == ConnectorCredentialPair.id,\n )\n .join(\n DocumentSetDBModel,\n DocumentSetDBModel.id\n == DocumentSet__ConnectorCredentialPair.document_set_id,\n )\n .where(DocumentSetDBModel.id == document_set_id)\n .order_by(Document.id)\n )\n\n if current_only:\n stmt = stmt.where(\n DocumentSet__ConnectorCredentialPair.is_current == True # noqa: E712\n )\n\n stmt = stmt.distinct()\n return stmt\n\n\ndef fetch_document_sets_for_document(\n document_id: str,\n db_session: Session,\n) -> list[str]:\n \"\"\"\n Fetches the document set names for a single document ID.\n\n :param document_id: The ID of the document to fetch sets for.\n :param db_session: The SQLAlchemy session to use for the query.\n :return: A list of document set names, or None if no result is found.\n \"\"\"\n result = fetch_document_sets_for_documents([document_id], db_session)\n if not result:\n return []\n return result[0][1]\n\n\ndef fetch_document_sets_for_documents(\n document_ids: list[str],\n db_session: Session,\n) -> Sequence[tuple[str, list[str]]]:\n \"\"\"Gives back a list of (document_id, list[document_set_names]) tuples\"\"\"\n\n \"\"\"Building subqueries\"\"\"\n # NOTE: have to build these subqueries first in order to guarantee that we get one\n # returned row for each specified document_id. Basically, we want to do the filters first,\n # then the outer joins.\n\n # don't include CC pairs that are being deleted\n # NOTE: CC pairs can never go from DELETING to any other state -> it's safe to ignore them\n # as we can assume their document sets are no longer relevant\n valid_cc_pairs_subquery = aliased(\n ConnectorCredentialPair,\n select(ConnectorCredentialPair)\n .where(\n ConnectorCredentialPair.status != ConnectorCredentialPairStatus.DELETING\n ) # noqa: E712\n .subquery(),\n )\n\n valid_document_set__cc_pairs_subquery = aliased(\n DocumentSet__ConnectorCredentialPair,\n select(DocumentSet__ConnectorCredentialPair)\n .where(DocumentSet__ConnectorCredentialPair.is_current == True) # noqa: E712\n .subquery(),\n )\n \"\"\"End building subqueries\"\"\"\n\n stmt = (\n select(\n Document.id,\n func.coalesce(\n func.array_remove(func.array_agg(DocumentSetDBModel.name), None), []\n ).label(\"document_set_names\"),\n )\n # Here we select document sets by relation:\n # Document -> DocumentByConnectorCredentialPair -> ConnectorCredentialPair ->\n # DocumentSet__ConnectorCredentialPair -> DocumentSet\n .outerjoin(\n DocumentByConnectorCredentialPair,\n Document.id == DocumentByConnectorCredentialPair.id,\n )\n .outerjoin(\n valid_cc_pairs_subquery,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == valid_cc_pairs_subquery.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == valid_cc_pairs_subquery.credential_id,\n ),\n )\n .outerjoin(\n valid_document_set__cc_pairs_subquery,\n valid_cc_pairs_subquery.id\n == valid_document_set__cc_pairs_subquery.connector_credential_pair_id,\n )\n .outerjoin(\n DocumentSetDBModel,\n DocumentSetDBModel.id\n == valid_document_set__cc_pairs_subquery.document_set_id,\n )\n .where(Document.id.in_(document_ids))\n .group_by(Document.id)\n )\n return db_session.execute(stmt).all() # type: ignore\n\n\ndef get_or_create_document_set_by_name(\n db_session: Session,\n document_set_name: str,\n document_set_description: str = \"Default Persona created Document-Set, please update description\",\n) -> DocumentSetDBModel:\n \"\"\"This is used by the default personas which need to attach to document sets\n on server startup\"\"\"\n doc_set = get_document_set_by_name(db_session, document_set_name)\n if doc_set is not None:\n return doc_set\n\n new_doc_set = DocumentSetDBModel(\n name=document_set_name,\n description=document_set_description,\n user_id=None,\n is_up_to_date=True,\n )\n\n db_session.add(new_doc_set)\n db_session.commit()\n\n return new_doc_set\n\n\ndef check_document_sets_are_public(\n db_session: Session,\n document_set_ids: list[int],\n) -> bool:\n \"\"\"Checks if any of the CC-Pairs are Non Public (meaning that some documents in this document\n set is not Public\"\"\"\n connector_credential_pair_ids = (\n db_session.query(\n DocumentSet__ConnectorCredentialPair.connector_credential_pair_id\n )\n .filter(\n DocumentSet__ConnectorCredentialPair.document_set_id.in_(document_set_ids)\n )\n .subquery()\n )\n\n not_public_exists = (\n db_session.query(ConnectorCredentialPair.id)\n .filter(\n ConnectorCredentialPair.id.in_(\n connector_credential_pair_ids # type:ignore\n ),\n ConnectorCredentialPair.access_type != AccessType.PUBLIC,\n )\n .limit(1)\n .first()\n is not None\n )\n\n return not not_public_exists\n" }, { "path": "backend/onyx/db/engine/__init__.py", "content": "" }, { "path": "backend/onyx/db/engine/async_sql_engine.py", "content": "from collections.abc import AsyncGenerator\nfrom contextlib import asynccontextmanager\nfrom typing import Any\nfrom typing import AsyncContextManager\n\nimport asyncpg # type: ignore\nfrom fastapi import HTTPException\nfrom sqlalchemy import event\nfrom sqlalchemy import pool\nfrom sqlalchemy.ext.asyncio import AsyncEngine\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.ext.asyncio import create_async_engine\n\nfrom onyx.configs.app_configs import AWS_REGION_NAME\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE\nfrom onyx.configs.app_configs import POSTGRES_DB\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_POOL_PRE_PING\nfrom onyx.configs.app_configs import POSTGRES_POOL_RECYCLE\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USE_NULL_POOL\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.db.engine.iam_auth import create_ssl_context_if_iam\nfrom onyx.db.engine.iam_auth import get_iam_auth_token\nfrom onyx.db.engine.sql_engine import ASYNC_DB_API\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.engine.sql_engine import is_valid_schema_name\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.engine.sql_engine import USE_IAM_AUTH\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\n# Global so we don't create more than one engine per process\n_ASYNC_ENGINE: AsyncEngine | None = None\n\n\nasync def get_async_connection() -> Any:\n \"\"\"\n Custom connection function for async engine when using IAM auth.\n \"\"\"\n host = POSTGRES_HOST\n port = POSTGRES_PORT\n user = POSTGRES_USER\n db = POSTGRES_DB\n token = get_iam_auth_token(host, port, user, AWS_REGION_NAME)\n\n # asyncpg requires 'ssl=\"require\"' if SSL needed\n return await asyncpg.connect(\n user=user, password=token, host=host, port=int(port), database=db, ssl=\"require\"\n )\n\n\ndef get_sqlalchemy_async_engine() -> AsyncEngine:\n global _ASYNC_ENGINE\n if _ASYNC_ENGINE is None:\n app_name = SqlEngine.get_app_name() + \"_async\"\n connection_string = build_connection_string(\n db_api=ASYNC_DB_API,\n use_iam_auth=USE_IAM_AUTH,\n )\n\n connect_args: dict[str, Any] = {}\n if app_name:\n connect_args[\"server_settings\"] = {\"application_name\": app_name}\n\n connect_args[\"ssl\"] = create_ssl_context_if_iam()\n\n engine_kwargs = {\n \"connect_args\": connect_args,\n \"pool_pre_ping\": POSTGRES_POOL_PRE_PING,\n \"pool_recycle\": POSTGRES_POOL_RECYCLE,\n }\n\n if POSTGRES_USE_NULL_POOL:\n engine_kwargs[\"poolclass\"] = pool.NullPool\n else:\n engine_kwargs[\"pool_size\"] = POSTGRES_API_SERVER_POOL_SIZE\n engine_kwargs[\"max_overflow\"] = POSTGRES_API_SERVER_POOL_OVERFLOW\n\n _ASYNC_ENGINE = create_async_engine(\n connection_string,\n **engine_kwargs,\n )\n\n if USE_IAM_AUTH:\n\n @event.listens_for(_ASYNC_ENGINE.sync_engine, \"do_connect\")\n def provide_iam_token_async(\n dialect: Any, # noqa: ARG001\n conn_rec: Any, # noqa: ARG001\n cargs: Any, # noqa: ARG001\n cparams: Any,\n ) -> None:\n # For async engine using asyncpg, we still need to set the IAM token here.\n host = POSTGRES_HOST\n port = POSTGRES_PORT\n user = POSTGRES_USER\n token = get_iam_auth_token(host, port, user, AWS_REGION_NAME)\n cparams[\"password\"] = token\n cparams[\"ssl\"] = create_ssl_context_if_iam()\n\n return _ASYNC_ENGINE\n\n\nasync def get_async_session(\n tenant_id: str | None = None,\n) -> AsyncGenerator[AsyncSession, None]:\n \"\"\"For use w/ Depends for *async* FastAPI endpoints.\n\n For standard `async with ... as ...` use, use get_async_session_context_manager.\n \"\"\"\n\n if tenant_id is None:\n tenant_id = get_current_tenant_id()\n\n if not is_valid_schema_name(tenant_id):\n raise HTTPException(status_code=400, detail=\"Invalid tenant ID\")\n\n engine = get_sqlalchemy_async_engine()\n\n # no need to use the schema translation map for self-hosted + default schema\n if not MULTI_TENANT and tenant_id == POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE:\n async with AsyncSession(bind=engine, expire_on_commit=False) as session:\n yield session\n return\n\n # Create connection with schema translation to handle querying the right schema\n schema_translate_map = {None: tenant_id}\n async with engine.connect() as connection:\n connection = await connection.execution_options(\n schema_translate_map=schema_translate_map\n )\n async with AsyncSession(\n bind=connection, expire_on_commit=False\n ) as async_session:\n yield async_session\n\n\ndef get_async_session_context_manager(\n tenant_id: str | None = None,\n) -> AsyncContextManager[AsyncSession]:\n return asynccontextmanager(get_async_session)(tenant_id)\n" }, { "path": "backend/onyx/db/engine/connection_warmup.py", "content": "from sqlalchemy import text\n\nfrom onyx.db.engine.async_sql_engine import get_sqlalchemy_async_engine\nfrom onyx.db.engine.sql_engine import get_sqlalchemy_engine\n\n\nasync def warm_up_connections(\n sync_connections_to_warm_up: int = 20, async_connections_to_warm_up: int = 20\n) -> None:\n sync_postgres_engine = get_sqlalchemy_engine()\n connections = [\n sync_postgres_engine.connect() for _ in range(sync_connections_to_warm_up)\n ]\n for conn in connections:\n conn.execute(text(\"SELECT 1\"))\n for conn in connections:\n conn.close()\n\n async_postgres_engine = get_sqlalchemy_async_engine()\n async_connections = [\n await async_postgres_engine.connect()\n for _ in range(async_connections_to_warm_up)\n ]\n for async_conn in async_connections:\n await async_conn.execute(text(\"SELECT 1\"))\n for async_conn in async_connections:\n await async_conn.close()\n" }, { "path": "backend/onyx/db/engine/iam_auth.py", "content": "import functools\nimport os\nimport ssl\nfrom typing import Any\n\nimport boto3\n\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.configs.app_configs import USE_IAM_AUTH\nfrom onyx.configs.constants import SSL_CERT_FILE\n\n\ndef get_iam_auth_token(\n host: str, port: str, user: str, region: str = \"us-east-2\"\n) -> str:\n \"\"\"\n Generate an IAM authentication token using boto3.\n \"\"\"\n client = boto3.client(\"rds\", region_name=region)\n token = client.generate_db_auth_token(\n DBHostname=host, Port=int(port), DBUsername=user\n )\n return token\n\n\ndef configure_psycopg2_iam_auth(\n cparams: dict[str, Any], host: str, port: str, user: str, region: str\n) -> None:\n \"\"\"\n Configure cparams for psycopg2 with IAM token and SSL.\n \"\"\"\n token = get_iam_auth_token(host, port, user, region)\n cparams[\"password\"] = token\n cparams[\"sslmode\"] = \"require\"\n cparams[\"sslrootcert\"] = SSL_CERT_FILE\n\n\ndef provide_iam_token(\n dialect: Any, # noqa: ARG001\n conn_rec: Any, # noqa: ARG001\n cargs: Any, # noqa: ARG001\n cparams: Any,\n) -> None:\n if USE_IAM_AUTH:\n host = POSTGRES_HOST\n port = POSTGRES_PORT\n user = POSTGRES_USER\n region = os.getenv(\"AWS_REGION_NAME\", \"us-east-2\")\n # Configure for psycopg2 with IAM token\n configure_psycopg2_iam_auth(cparams, host, port, user, region)\n\n\n@functools.cache\ndef create_ssl_context_if_iam() -> ssl.SSLContext | None:\n \"\"\"Create an SSL context if IAM authentication is enabled, else return None.\"\"\"\n if USE_IAM_AUTH:\n return ssl.create_default_context(cafile=SSL_CERT_FILE)\n return None\n" }, { "path": "backend/onyx/db/engine/sql_engine.py", "content": "import os\nimport re\nimport threading\nimport time\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import event\nfrom sqlalchemy import pool\nfrom sqlalchemy.engine import create_engine\nfrom sqlalchemy.engine import Engine\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DB_READONLY_PASSWORD\nfrom onyx.configs.app_configs import DB_READONLY_USER\nfrom onyx.configs.app_configs import LOG_POSTGRES_CONN_COUNTS\nfrom onyx.configs.app_configs import LOG_POSTGRES_LATENCY\nfrom onyx.configs.app_configs import POSTGRES_DB\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PASSWORD\nfrom onyx.configs.app_configs import POSTGRES_POOL_PRE_PING\nfrom onyx.configs.app_configs import POSTGRES_POOL_RECYCLE\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USE_NULL_POOL\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.configs.constants import POSTGRES_UNKNOWN_APP_NAME\nfrom onyx.db.engine.iam_auth import provide_iam_token\nfrom onyx.server.utils import BasicAuthenticationError\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.contextvars import get_current_tenant_id\n\n# Moved is_valid_schema_name here to avoid circular import\n\n\nlogger = setup_logger()\n\n\n# Schema name validation (moved here to avoid circular import)\nSCHEMA_NAME_REGEX = re.compile(r\"^[a-zA-Z0-9_-]+$\")\n\n\ndef is_valid_schema_name(name: str) -> bool:\n return SCHEMA_NAME_REGEX.match(name) is not None\n\n\nSYNC_DB_API = \"psycopg2\"\nASYNC_DB_API = \"asyncpg\"\n\n# why isn't this in configs?\nUSE_IAM_AUTH = os.getenv(\"USE_IAM_AUTH\", \"False\").lower() == \"true\"\n\n\ndef build_connection_string(\n *,\n db_api: str = ASYNC_DB_API,\n user: str = POSTGRES_USER,\n password: str = POSTGRES_PASSWORD,\n host: str = POSTGRES_HOST,\n port: str = POSTGRES_PORT,\n db: str = POSTGRES_DB,\n app_name: str | None = None,\n use_iam_auth: bool = USE_IAM_AUTH,\n region: str = \"us-west-2\", # noqa: ARG001\n) -> str:\n if use_iam_auth:\n base_conn_str = f\"postgresql+{db_api}://{user}@{host}:{port}/{db}\"\n else:\n base_conn_str = f\"postgresql+{db_api}://{user}:{password}@{host}:{port}/{db}\"\n\n # For asyncpg, do not include application_name in the connection string\n if app_name and db_api != \"asyncpg\":\n if \"?\" in base_conn_str:\n return f\"{base_conn_str}&application_name={app_name}\"\n else:\n return f\"{base_conn_str}?application_name={app_name}\"\n return base_conn_str\n\n\nif LOG_POSTGRES_LATENCY:\n\n @event.listens_for(Engine, \"before_cursor_execute\")\n def before_cursor_execute( # type: ignore\n conn,\n cursor, # noqa: ARG001\n statement, # noqa: ARG001\n parameters, # noqa: ARG001\n context, # noqa: ARG001\n executemany, # noqa: ARG001\n ):\n conn.info[\"query_start_time\"] = time.time()\n\n @event.listens_for(Engine, \"after_cursor_execute\")\n def after_cursor_execute( # type: ignore\n conn,\n cursor, # noqa: ARG001\n statement,\n parameters, # noqa: ARG001\n context, # noqa: ARG001\n executemany, # noqa: ARG001\n ):\n total_time = time.time() - conn.info[\"query_start_time\"]\n if total_time > 0.1:\n logger.debug(\n f\"Query Complete: {statement}\\n\\nTotal Time: {total_time:.4f} seconds\"\n )\n\n\nif LOG_POSTGRES_CONN_COUNTS:\n checkout_count = 0\n checkin_count = 0\n\n @event.listens_for(Engine, \"checkout\")\n def log_checkout(dbapi_connection, connection_record, connection_proxy): # type: ignore # noqa: ARG001\n global checkout_count\n checkout_count += 1\n\n active_connections = connection_proxy._pool.checkedout()\n idle_connections = connection_proxy._pool.checkedin()\n pool_size = connection_proxy._pool.size()\n logger.debug(\n \"Connection Checkout\\n\"\n f\"Active Connections: {active_connections};\\n\"\n f\"Idle: {idle_connections};\\n\"\n f\"Pool Size: {pool_size};\\n\"\n f\"Total connection checkouts: {checkout_count}\"\n )\n\n @event.listens_for(Engine, \"checkin\")\n def log_checkin(dbapi_connection, connection_record): # type: ignore # noqa: ARG001\n global checkin_count\n checkin_count += 1\n logger.debug(f\"Total connection checkins: {checkin_count}\")\n\n\nclass SqlEngine:\n _engine: Engine | None = None\n _readonly_engine: Engine | None = None\n _lock: threading.Lock = threading.Lock()\n _readonly_lock: threading.Lock = threading.Lock()\n _app_name: str = POSTGRES_UNKNOWN_APP_NAME\n\n @classmethod\n def init_engine(\n cls,\n pool_size: int,\n # is really `pool_max_overflow`, but calling it `max_overflow` to stay consistent with SQLAlchemy\n max_overflow: int,\n app_name: str | None = None, # noqa: ARG003\n db_api: str = SYNC_DB_API,\n use_iam: bool = USE_IAM_AUTH,\n connection_string: str | None = None,\n **extra_engine_kwargs: Any,\n ) -> None:\n \"\"\"NOTE: enforce that pool_size and pool_max_overflow are passed in. These are\n important args, and if incorrectly specified, we have run into hitting the pool\n limit / using too many connections and overwhelming the database.\n\n Specifying connection_string directly will cause some of the other parameters\n to be ignored.\n \"\"\"\n with cls._lock:\n if cls._engine:\n return\n\n if not connection_string:\n connection_string = build_connection_string(\n db_api=db_api,\n app_name=cls._app_name + \"_sync\",\n use_iam_auth=use_iam,\n )\n\n # Start with base kwargs that are valid for all pool types\n final_engine_kwargs: dict[str, Any] = {}\n\n if POSTGRES_USE_NULL_POOL:\n # if null pool is specified, then we need to make sure that\n # we remove any passed in kwargs related to pool size that would\n # cause the initialization to fail\n final_engine_kwargs.update(extra_engine_kwargs)\n\n final_engine_kwargs[\"poolclass\"] = pool.NullPool\n if \"pool_size\" in final_engine_kwargs:\n del final_engine_kwargs[\"pool_size\"]\n if \"max_overflow\" in final_engine_kwargs:\n del final_engine_kwargs[\"max_overflow\"]\n else:\n final_engine_kwargs[\"pool_size\"] = pool_size\n final_engine_kwargs[\"max_overflow\"] = max_overflow\n final_engine_kwargs[\"pool_pre_ping\"] = POSTGRES_POOL_PRE_PING\n final_engine_kwargs[\"pool_recycle\"] = POSTGRES_POOL_RECYCLE\n\n # any passed in kwargs override the defaults\n final_engine_kwargs.update(extra_engine_kwargs)\n\n logger.info(f\"Creating engine with kwargs: {final_engine_kwargs}\")\n # echo=True here for inspecting all emitted db queries\n engine = create_engine(connection_string, **final_engine_kwargs)\n\n if use_iam:\n event.listen(engine, \"do_connect\", provide_iam_token)\n\n cls._engine = engine\n\n @classmethod\n def init_readonly_engine(\n cls,\n pool_size: int,\n # is really `pool_max_overflow`, but calling it `max_overflow` to stay consistent with SQLAlchemy\n max_overflow: int,\n **extra_engine_kwargs: Any,\n ) -> None:\n \"\"\"NOTE: enforce that pool_size and pool_max_overflow are passed in. These are\n important args, and if incorrectly specified, we have run into hitting the pool\n limit / using too many connections and overwhelming the database.\"\"\"\n with cls._readonly_lock:\n if cls._readonly_engine:\n return\n\n if not DB_READONLY_USER or not DB_READONLY_PASSWORD:\n raise ValueError(\n \"Custom database user credentials not configured in environment variables\"\n )\n\n # Build connection string with custom user\n connection_string = build_connection_string(\n user=DB_READONLY_USER,\n password=DB_READONLY_PASSWORD,\n use_iam_auth=False, # Custom users typically don't use IAM auth\n db_api=SYNC_DB_API, # Explicitly use sync DB API\n )\n\n # Start with base kwargs that are valid for all pool types\n final_engine_kwargs: dict[str, Any] = {}\n\n if POSTGRES_USE_NULL_POOL:\n # if null pool is specified, then we need to make sure that\n # we remove any passed in kwargs related to pool size that would\n # cause the initialization to fail\n final_engine_kwargs.update(extra_engine_kwargs)\n\n final_engine_kwargs[\"poolclass\"] = pool.NullPool\n if \"pool_size\" in final_engine_kwargs:\n del final_engine_kwargs[\"pool_size\"]\n if \"max_overflow\" in final_engine_kwargs:\n del final_engine_kwargs[\"max_overflow\"]\n else:\n final_engine_kwargs[\"pool_size\"] = pool_size\n final_engine_kwargs[\"max_overflow\"] = max_overflow\n final_engine_kwargs[\"pool_pre_ping\"] = POSTGRES_POOL_PRE_PING\n final_engine_kwargs[\"pool_recycle\"] = POSTGRES_POOL_RECYCLE\n\n # any passed in kwargs override the defaults\n final_engine_kwargs.update(extra_engine_kwargs)\n\n logger.info(f\"Creating engine with kwargs: {final_engine_kwargs}\")\n # echo=True here for inspecting all emitted db queries\n engine = create_engine(connection_string, **final_engine_kwargs)\n\n if USE_IAM_AUTH:\n event.listen(engine, \"do_connect\", provide_iam_token)\n\n cls._readonly_engine = engine\n\n @classmethod\n def get_engine(cls) -> Engine:\n if not cls._engine:\n raise RuntimeError(\"Engine not initialized. Must call init_engine first.\")\n return cls._engine\n\n @classmethod\n def get_readonly_engine(cls) -> Engine:\n if not cls._readonly_engine:\n raise RuntimeError(\n \"Readonly engine not initialized. Must call init_readonly_engine first.\"\n )\n return cls._readonly_engine\n\n @classmethod\n def set_app_name(cls, app_name: str) -> None:\n cls._app_name = app_name\n\n @classmethod\n def get_app_name(cls) -> str:\n if not cls._app_name:\n return \"\"\n return cls._app_name\n\n @classmethod\n def reset_engine(cls) -> None:\n with cls._lock:\n if cls._engine:\n cls._engine.dispose()\n cls._engine = None\n\n @classmethod\n @contextmanager\n def scoped_engine(cls, **init_kwargs: Any) -> Generator[None, None, None]:\n \"\"\"Context manager that initializes the engine and guarantees cleanup.\"\"\"\n cls.init_engine(**init_kwargs)\n try:\n yield\n finally:\n cls.reset_engine()\n\n\ndef get_sqlalchemy_engine() -> Engine:\n return SqlEngine.get_engine()\n\n\ndef get_readonly_sqlalchemy_engine() -> Engine:\n return SqlEngine.get_readonly_engine()\n\n\n@contextmanager\ndef get_session_with_current_tenant() -> Generator[Session, None, None]:\n \"\"\"Standard way to get a DB session.\"\"\"\n tenant_id = get_current_tenant_id()\n with get_session_with_tenant(tenant_id=tenant_id) as session:\n yield session\n\n\n@contextmanager\ndef get_session_with_current_tenant_if_none(\n session: Session | None,\n) -> Generator[Session, None, None]:\n if session is None:\n tenant_id = get_current_tenant_id()\n with get_session_with_tenant(tenant_id=tenant_id) as session:\n yield session\n else:\n yield session\n\n\n# Used in multi tenant mode when need to refer to the shared `public` schema\n@contextmanager\ndef get_session_with_shared_schema() -> Generator[Session, None, None]:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA)\n with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as session:\n yield session\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\n@contextmanager\ndef get_session_with_tenant(*, tenant_id: str) -> Generator[Session, None, None]:\n \"\"\"\n Generate a database session for a specific tenant.\n \"\"\"\n engine = get_sqlalchemy_engine()\n\n if not is_valid_schema_name(tenant_id):\n raise HTTPException(status_code=400, detail=\"Invalid tenant ID\")\n\n # no need to use the schema translation map for self-hosted + default schema\n if not MULTI_TENANT and tenant_id == POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE:\n with Session(bind=engine, expire_on_commit=False) as session:\n yield session\n return\n\n # Create connection with schema translation to handle querying the right schema\n schema_translate_map = {None: tenant_id}\n with engine.connect().execution_options(\n schema_translate_map=schema_translate_map\n ) as connection:\n with Session(bind=connection, expire_on_commit=False) as session:\n yield session\n\n\ndef get_session() -> Generator[Session, None, None]:\n \"\"\"For use w/ Depends for FastAPI endpoints.\n\n Has some additional validation, and likely should be merged\n with get_session_with_current_tenant in the future.\"\"\"\n tenant_id = get_current_tenant_id()\n if tenant_id == POSTGRES_DEFAULT_SCHEMA and MULTI_TENANT:\n raise BasicAuthenticationError(detail=\"User must authenticate\")\n\n if not is_valid_schema_name(tenant_id):\n raise HTTPException(status_code=400, detail=\"Invalid tenant ID\")\n\n with get_session_with_current_tenant() as db_session:\n yield db_session\n\n\n@contextmanager\ndef get_db_readonly_user_session_with_current_tenant() -> (\n Generator[Session, None, None]\n):\n \"\"\"\n Generate a database session using a custom database user for the current tenant.\n The custom user credentials are obtained from environment variables.\n \"\"\"\n tenant_id = get_current_tenant_id()\n\n readonly_engine = get_readonly_sqlalchemy_engine()\n\n if not is_valid_schema_name(tenant_id):\n raise HTTPException(status_code=400, detail=\"Invalid tenant ID\")\n\n # no need to use the schema translation map for self-hosted + default schema\n if not MULTI_TENANT and tenant_id == POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE:\n with Session(readonly_engine, expire_on_commit=False) as session:\n yield session\n return\n\n schema_translate_map = {None: tenant_id}\n with readonly_engine.connect().execution_options(\n schema_translate_map=schema_translate_map\n ) as connection:\n with Session(bind=connection, expire_on_commit=False) as session:\n yield session\n" }, { "path": "backend/onyx/db/engine/tenant_utils.py", "content": "from sqlalchemy import text\n\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import TENANT_ID_PREFIX\n\n\ndef get_schemas_needing_migration(\n tenant_schemas: list[str], head_rev: str\n) -> list[str]:\n \"\"\"Return only schemas whose current alembic version is not at head.\n\n Uses a server-side PL/pgSQL loop to collect each schema's alembic version\n into a temp table one at a time. This avoids building a massive UNION ALL\n query (which locks the DB and times out at 17k+ schemas) and instead\n acquires locks sequentially, one schema per iteration.\n \"\"\"\n if not tenant_schemas:\n return []\n\n engine = SqlEngine.get_engine()\n\n with engine.connect() as conn:\n # Populate a temp input table with exactly the schemas we care about.\n # The DO block reads from this table so it only iterates the requested\n # schemas instead of every tenant_% schema in the database.\n conn.execute(text(\"DROP TABLE IF EXISTS _alembic_version_snapshot\"))\n conn.execute(text(\"DROP TABLE IF EXISTS _tenant_schemas_input\"))\n conn.execute(text(\"CREATE TEMP TABLE _tenant_schemas_input (schema_name text)\"))\n conn.execute(\n text(\n \"INSERT INTO _tenant_schemas_input (schema_name) SELECT unnest(CAST(:schemas AS text[]))\"\n ),\n {\"schemas\": tenant_schemas},\n )\n conn.execute(\n text(\n \"CREATE TEMP TABLE _alembic_version_snapshot (schema_name text, version_num text)\"\n )\n )\n\n conn.execute(\n text(\n \"\"\"\n DO $$\n DECLARE\n s text;\n schemas text[];\n BEGIN\n SELECT array_agg(schema_name) INTO schemas\n FROM _tenant_schemas_input;\n\n IF schemas IS NULL THEN\n RAISE NOTICE 'No tenant schemas found.';\n RETURN;\n END IF;\n\n FOREACH s IN ARRAY schemas LOOP\n BEGIN\n EXECUTE format(\n 'INSERT INTO _alembic_version_snapshot\n SELECT %L, version_num FROM %I.alembic_version',\n s, s\n );\n EXCEPTION\n -- undefined_table: schema exists but has no alembic_version\n -- table yet (new tenant, not yet migrated).\n -- invalid_schema_name: tenant is registered but its\n -- PostgreSQL schema does not exist yet (e.g. provisioning\n -- incomplete). Both cases mean no version is available and\n -- the schema will be included in the migration list.\n WHEN undefined_table THEN NULL;\n WHEN invalid_schema_name THEN NULL;\n END;\n END LOOP;\n END;\n $$\n \"\"\"\n )\n )\n\n rows = conn.execute(\n text(\"SELECT schema_name, version_num FROM _alembic_version_snapshot\")\n )\n version_by_schema = {row[0]: row[1] for row in rows}\n\n conn.execute(text(\"DROP TABLE IF EXISTS _alembic_version_snapshot\"))\n conn.execute(text(\"DROP TABLE IF EXISTS _tenant_schemas_input\"))\n\n # Schemas missing from the snapshot have no alembic_version table yet and\n # also need migration. version_by_schema.get(s) returns None for those,\n # and None != head_rev, so they are included automatically.\n return [s for s in tenant_schemas if version_by_schema.get(s) != head_rev]\n\n\ndef get_all_tenant_ids() -> list[str]:\n \"\"\"Returning [None] means the only tenant is the 'public' or self hosted tenant.\"\"\"\n\n tenant_ids: list[str]\n\n if not MULTI_TENANT:\n return [POSTGRES_DEFAULT_SCHEMA]\n\n with get_session_with_shared_schema() as session:\n result = session.execute(\n text(\n f\"\"\"\n SELECT schema_name\n FROM information_schema.schemata\n WHERE schema_name NOT IN ('pg_catalog', 'information_schema', '{POSTGRES_DEFAULT_SCHEMA}')\"\"\"\n )\n )\n tenant_ids = [row[0] for row in result]\n\n valid_tenants = [\n tenant\n for tenant in tenant_ids\n if tenant is None or tenant.startswith(TENANT_ID_PREFIX)\n ]\n return valid_tenants\n" }, { "path": "backend/onyx/db/engine/time_utils.py", "content": "from datetime import datetime\n\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\n\n\ndef get_db_current_time(db_session: Session) -> datetime:\n result = db_session.execute(text(\"SELECT NOW()\")).scalar()\n if result is None:\n raise ValueError(\"Database did not return a time\")\n return result\n" }, { "path": "backend/onyx/db/entities.py", "content": "import uuid\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import List\n\nfrom sqlalchemy import func\nfrom sqlalchemy import literal\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.dialects.postgresql import JSONB\nfrom sqlalchemy.orm import Session\n\nimport onyx.db.document as dbdocument\nfrom onyx.db.entity_type import UNGROUNDED_SOURCE_NAME\nfrom onyx.db.models import Document\nfrom onyx.db.models import KGEntity\nfrom onyx.db.models import KGEntityExtractionStaging\nfrom onyx.db.models import KGEntityType\nfrom onyx.kg.models import KGGroundingType\nfrom onyx.kg.models import KGStage\nfrom onyx.kg.utils.formatting_utils import make_entity_id\n\n\ndef upsert_staging_entity(\n db_session: Session,\n name: str,\n entity_type: str,\n document_id: str | None = None,\n occurrences: int = 1,\n attributes: dict[str, str] | None = None,\n event_time: datetime | None = None,\n) -> KGEntityExtractionStaging:\n \"\"\"Add or update a new staging entity to the database.\n\n Args:\n db_session: SQLAlchemy session\n name: Name of the entity\n entity_type: Type of the entity (must match an existing KGEntityType)\n document_id: ID of the document the entity belongs to\n occurrences: Number of times this entity has been found\n attributes: Attributes of the entity\n event_time: Time the entity was added to the database\n\n Returns:\n KGEntityExtractionStaging: The created entity\n \"\"\"\n entity_type = entity_type.upper()\n name = name.title()\n id_name = make_entity_id(entity_type, name)\n attributes = attributes or {}\n\n entity_key = attributes.get(\"key\")\n entity_parent = attributes.get(\"parent\")\n\n keep_attributes = {\n attr_key: attr_val\n for attr_key, attr_val in attributes.items()\n if attr_key not in (\"key\", \"parent\")\n }\n\n # Create new entity\n stmt = (\n pg_insert(KGEntityExtractionStaging)\n .values(\n id_name=id_name,\n name=name,\n entity_type_id_name=entity_type,\n entity_key=entity_key,\n parent_key=entity_parent,\n document_id=document_id,\n occurrences=occurrences,\n attributes=keep_attributes,\n event_time=event_time,\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\"],\n set_=dict(\n occurrences=KGEntityExtractionStaging.occurrences + occurrences,\n ),\n )\n .returning(KGEntityExtractionStaging)\n )\n\n result = db_session.execute(stmt).scalar()\n if result is None:\n raise RuntimeError(\n f\"Failed to create or increment staging entity with id_name: {id_name}\"\n )\n\n # Update the document's kg_stage if document_id is provided\n if document_id is not None:\n db_session.query(Document).filter(Document.id == document_id).update(\n {\n \"kg_stage\": KGStage.EXTRACTED,\n \"kg_processing_time\": datetime.now(timezone.utc),\n }\n )\n db_session.flush()\n\n return result\n\n\ndef transfer_entity(\n db_session: Session,\n entity: KGEntityExtractionStaging,\n) -> KGEntity:\n \"\"\"Transfer an entity from the extraction staging table to the normalized table.\n\n Args:\n db_session: SQLAlchemy session\n entity: Entity to transfer\n\n Returns:\n KGEntity: The transferred entity\n \"\"\"\n # Create the transferred entity\n stmt = (\n pg_insert(KGEntity)\n .values(\n id_name=make_entity_id(entity.entity_type_id_name, uuid.uuid4().hex[:20]),\n name=entity.name.casefold(),\n entity_key=entity.entity_key,\n parent_key=entity.parent_key,\n alternative_names=entity.alternative_names or [],\n entity_type_id_name=entity.entity_type_id_name,\n document_id=entity.document_id,\n occurrences=entity.occurrences,\n attributes=entity.attributes,\n event_time=entity.event_time,\n )\n .on_conflict_do_update(\n index_elements=[\"name\", \"entity_type_id_name\", \"document_id\"],\n set_=dict(\n occurrences=KGEntity.occurrences + entity.occurrences,\n attributes=KGEntity.attributes.op(\"||\")(\n literal(entity.attributes, JSONB)\n ),\n entity_key=func.coalesce(KGEntity.entity_key, entity.entity_key),\n parent_key=func.coalesce(KGEntity.parent_key, entity.parent_key),\n event_time=entity.event_time,\n time_updated=datetime.now(),\n ),\n )\n .returning(KGEntity)\n )\n new_entity = db_session.execute(stmt).scalar()\n if new_entity is None:\n raise RuntimeError(f\"Failed to transfer entity with id_name: {entity.id_name}\")\n\n # Update the document's kg_stage if document_id is provided\n if entity.document_id is not None:\n dbdocument.update_document_kg_info(\n db_session,\n document_id=entity.document_id,\n kg_stage=KGStage.NORMALIZED,\n )\n\n # Update transferred\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.id_name == entity.id_name\n ).update({\"transferred_id_name\": new_entity.id_name})\n db_session.flush()\n\n return new_entity\n\n\ndef merge_entities(\n db_session: Session, parent: KGEntity, child: KGEntityExtractionStaging\n) -> KGEntity:\n \"\"\"Merge an entity from the extraction staging table into\n an existing entity in the normalized table.\n\n Args:\n db_session: SQLAlchemy session\n parent: Parent entity to merge into\n child: Child staging entity to merge\n\n Returns:\n KGEntity: The merged entity\n \"\"\"\n # check we're not merging two entities with different document_ids\n if (\n parent.document_id is not None\n and child.document_id is not None\n and parent.document_id != child.document_id\n ):\n raise ValueError(\n \"Overwriting the document_id of an entity with a document_id already is not allowed\"\n )\n\n # update the parent entity (only document_id, alternative_names, occurrences)\n setting_doc = parent.document_id is None and child.document_id is not None\n document_id = child.document_id if setting_doc else parent.document_id\n alternative_names = set(parent.alternative_names or [])\n alternative_names.update(child.alternative_names or [])\n alternative_names.add(child.name.lower())\n alternative_names.discard(parent.name)\n\n stmt = (\n update(KGEntity)\n .where(KGEntity.id_name == parent.id_name)\n .values(\n document_id=document_id,\n alternative_names=list(alternative_names),\n occurrences=parent.occurrences + child.occurrences,\n attributes=parent.attributes | child.attributes,\n entity_key=parent.entity_key or child.entity_key,\n parent_key=parent.parent_key or child.parent_key,\n )\n .returning(KGEntity)\n )\n\n result = db_session.execute(stmt).scalar()\n if result is None:\n raise RuntimeError(f\"Failed to merge entities with id_name: {parent.id_name}\")\n\n # Update the document's kg_stage if document_id is set\n if setting_doc and child.document_id is not None:\n dbdocument.update_document_kg_info(\n db_session,\n document_id=child.document_id,\n kg_stage=KGStage.NORMALIZED,\n )\n\n # Update transferred\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.id_name == child.id_name\n ).update({\"transferred_id_name\": parent.id_name})\n db_session.flush()\n\n return result\n\n\ndef get_kg_entity_by_document(db: Session, document_id: str) -> KGEntity | None:\n \"\"\"\n Check if a document_id exists in the kg_entities table and return its id_name if found.\n\n Args:\n db: SQLAlchemy database session\n document_id: The document ID to search for\n\n Returns:\n The id_name of the matching KGEntity if found, None otherwise\n \"\"\"\n query = select(KGEntity).where(KGEntity.document_id == document_id)\n result = db.execute(query).scalar()\n return result\n\n\ndef get_grounded_entities_by_types(\n db_session: Session, entity_types: List[str], grounding: KGGroundingType\n) -> List[KGEntity]:\n \"\"\"Get all entities matching an entity_type.\n\n Args:\n db_session: SQLAlchemy session\n entity_types: List of entity types to filter by\n\n Returns:\n List of KGEntity objects belonging to the specified entity types\n \"\"\"\n return (\n db_session.query(KGEntity)\n .join(KGEntityType, KGEntity.entity_type_id_name == KGEntityType.id_name)\n .filter(KGEntity.entity_type_id_name.in_(entity_types))\n .filter(KGEntityType.grounding == grounding)\n .all()\n )\n\n\ndef get_document_id_for_entity(db_session: Session, entity_id_name: str) -> str | None:\n \"\"\"Get the document ID associated with an entity.\n\n Args:\n db_session: SQLAlchemy database session\n entity_id_name: The entity id_name to look up\n\n Returns:\n The document ID if found, None otherwise\n \"\"\"\n entity = (\n db_session.query(KGEntity).filter(KGEntity.id_name == entity_id_name).first()\n )\n return entity.document_id if entity else None\n\n\ndef delete_from_kg_entities_extraction_staging__no_commit(\n db_session: Session, document_ids: list[str]\n) -> None:\n \"\"\"Delete entities from the extraction staging table.\"\"\"\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.document_id.in_(document_ids)\n ).delete(synchronize_session=False)\n\n\ndef delete_from_kg_entities__no_commit(\n db_session: Session, document_ids: list[str]\n) -> None:\n \"\"\"Delete entities from the normalized table.\"\"\"\n db_session.query(KGEntity).filter(KGEntity.document_id.in_(document_ids)).delete(\n synchronize_session=False\n )\n\n\ndef get_entity_name(db_session: Session, entity_id_name: str) -> str | None:\n \"\"\"Get the name of an entity.\"\"\"\n entity = (\n db_session.query(KGEntity).filter(KGEntity.id_name == entity_id_name).first()\n )\n return entity.name if entity else None\n\n\ndef get_entity_stats_by_grounded_source_name(\n db_session: Session,\n) -> dict[str, tuple[datetime, int]]:\n \"\"\"\n Returns a dict mapping each grounded_source_name to a tuple in which:\n - the first element is the latest update time across all entities with the same entity-type\n - the second element is the count of `KGEntity`s\n \"\"\"\n results = (\n db_session.query(\n KGEntityType.grounded_source_name,\n func.count(KGEntity.id_name).label(\"entities_count\"),\n func.max(KGEntity.time_updated).label(\"last_updated\"),\n )\n .join(KGEntityType, KGEntity.entity_type_id_name == KGEntityType.id_name)\n .group_by(KGEntityType.grounded_source_name)\n .all()\n )\n\n # `row.grounded_source_name` is NULLABLE in the database schema.\n # Thus, for all \"ungrounded\" entity-types, we use a default name.\n return {\n (row.grounded_source_name or UNGROUNDED_SOURCE_NAME): (\n row.last_updated,\n row.entities_count,\n )\n for row in results\n }\n" }, { "path": "backend/onyx/db/entity_type.py", "content": "from collections import defaultdict\n\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.connector import fetch_unique_document_sources\nfrom onyx.db.document import DocumentSource\nfrom onyx.db.models import Connector\nfrom onyx.db.models import KGEntityType\nfrom onyx.kg.models import KGAttributeEntityOption\nfrom onyx.server.kg.models import EntityType\n\n\nUNGROUNDED_SOURCE_NAME = \"Ungrounded\"\n\n\ndef get_entity_types_with_grounded_source_name(\n db_session: Session,\n) -> list[KGEntityType]:\n \"\"\"Get all entity types that have non-null grounded_source_name.\n\n Args:\n db_session: SQLAlchemy session\n\n Returns:\n List of KGEntityType objects that have grounded_source_name defined\n \"\"\"\n return (\n db_session.query(KGEntityType)\n .filter(KGEntityType.grounded_source_name.isnot(None))\n .all()\n )\n\n\ndef get_entity_types(\n db_session: Session,\n active: bool | None = True,\n) -> list[KGEntityType]:\n # Query the database for all distinct entity types\n\n if active is None:\n return db_session.query(KGEntityType).order_by(KGEntityType.id_name).all()\n\n else:\n return (\n db_session.query(KGEntityType)\n .filter(KGEntityType.active == active)\n .order_by(KGEntityType.id_name)\n .all()\n )\n\n\ndef get_configured_entity_types(db_session: Session) -> dict[str, list[KGEntityType]]:\n # get entity types from configured sources\n configured_connector_sources = {\n source.value.lower()\n for source in fetch_unique_document_sources(db_session=db_session)\n }\n entity_types = (\n db_session.query(KGEntityType)\n .filter(KGEntityType.grounded_source_name.in_(configured_connector_sources))\n .all()\n )\n entity_type_set = {et.id_name for et in entity_types}\n\n # get implied entity types from those entity types\n for et in entity_types:\n for prop in et.parsed_attributes.metadata_attribute_conversion.values():\n if prop.implication_property is None:\n continue\n\n implied_et = prop.implication_property.implied_entity_type\n if implied_et == KGAttributeEntityOption.FROM_EMAIL:\n if \"ACCOUNT\" not in entity_type_set:\n entity_type_set.add(\"ACCOUNT\")\n if \"EMPLOYEE\" not in entity_type_set:\n entity_type_set.add(\"EMPLOYEE\")\n elif isinstance(implied_et, str):\n if implied_et not in entity_type_set:\n entity_type_set.add(implied_et)\n\n ets = (\n db_session.query(KGEntityType)\n .filter(KGEntityType.id_name.in_(entity_type_set))\n .all()\n )\n\n et_map = defaultdict(list)\n for et in ets:\n key = et.grounded_source_name or UNGROUNDED_SOURCE_NAME\n et_map[key].append(et)\n\n return et_map\n\n\ndef update_entity_types_and_related_connectors__commit(\n db_session: Session, updates: list[EntityType]\n) -> None:\n for upd in updates:\n db_session.execute(\n update(KGEntityType)\n .where(KGEntityType.id_name == upd.name)\n .values(\n description=upd.description,\n active=upd.active,\n )\n )\n db_session.flush()\n\n # Update connector sources\n\n configured_entity_types = get_configured_entity_types(db_session=db_session)\n\n active_entity_type_sources = {\n et.grounded_source_name\n for ets in configured_entity_types.values()\n for et in ets\n if et.active\n }\n\n # Update connectors that should be enabled\n db_session.execute(\n update(Connector)\n .where(\n Connector.source.in_(\n [\n source\n for source in DocumentSource\n if source.value.lower() in active_entity_type_sources\n ]\n )\n )\n .where(~Connector.kg_processing_enabled)\n .values(kg_processing_enabled=True)\n )\n\n # Update connectors that should be disabled\n db_session.execute(\n update(Connector)\n .where(\n Connector.source.in_(\n [\n source\n for source in DocumentSource\n if source.value.lower() not in active_entity_type_sources\n ]\n )\n )\n .where(Connector.kg_processing_enabled)\n .values(kg_processing_enabled=False)\n )\n\n db_session.commit()\n" }, { "path": "backend/onyx/db/enums.py", "content": "from __future__ import annotations\n\nfrom enum import Enum as PyEnum\nfrom typing import ClassVar\n\n\nclass AccountType(str, PyEnum):\n \"\"\"\n What kind of account this is — determines whether the user\n enters the group-based permission system.\n\n STANDARD + SERVICE_ACCOUNT → participate in group system\n BOT, EXT_PERM_USER, ANONYMOUS → fixed behavior\n \"\"\"\n\n STANDARD = \"STANDARD\"\n BOT = \"BOT\"\n EXT_PERM_USER = \"EXT_PERM_USER\"\n SERVICE_ACCOUNT = \"SERVICE_ACCOUNT\"\n ANONYMOUS = \"ANONYMOUS\"\n\n def is_web_login(self) -> bool:\n \"\"\"Whether this account type supports interactive web login.\"\"\"\n return self not in (\n AccountType.BOT,\n AccountType.EXT_PERM_USER,\n )\n\n\nclass GrantSource(str, PyEnum):\n \"\"\"How a permission grant was created.\"\"\"\n\n USER = \"USER\"\n SCIM = \"SCIM\"\n SYSTEM = \"SYSTEM\"\n\n\nclass IndexingStatus(str, PyEnum):\n NOT_STARTED = \"not_started\"\n IN_PROGRESS = \"in_progress\"\n SUCCESS = \"success\"\n CANCELED = \"canceled\"\n FAILED = \"failed\"\n COMPLETED_WITH_ERRORS = \"completed_with_errors\"\n\n def is_terminal(self) -> bool:\n terminal_states = {\n IndexingStatus.SUCCESS,\n IndexingStatus.COMPLETED_WITH_ERRORS,\n IndexingStatus.CANCELED,\n IndexingStatus.FAILED,\n }\n return self in terminal_states\n\n def is_successful(self) -> bool:\n return (\n self == IndexingStatus.SUCCESS\n or self == IndexingStatus.COMPLETED_WITH_ERRORS\n )\n\n\nclass PermissionSyncStatus(str, PyEnum):\n \"\"\"Status enum for permission sync attempts\"\"\"\n\n NOT_STARTED = \"not_started\"\n IN_PROGRESS = \"in_progress\"\n SUCCESS = \"success\"\n CANCELED = \"canceled\"\n FAILED = \"failed\"\n COMPLETED_WITH_ERRORS = \"completed_with_errors\"\n\n def is_terminal(self) -> bool:\n terminal_states = {\n PermissionSyncStatus.SUCCESS,\n PermissionSyncStatus.COMPLETED_WITH_ERRORS,\n PermissionSyncStatus.CANCELED,\n PermissionSyncStatus.FAILED,\n }\n return self in terminal_states\n\n def is_successful(self) -> bool:\n return (\n self == PermissionSyncStatus.SUCCESS\n or self == PermissionSyncStatus.COMPLETED_WITH_ERRORS\n )\n\n\nclass IndexingMode(str, PyEnum):\n UPDATE = \"update\"\n REINDEX = \"reindex\"\n\n\nclass ProcessingMode(str, PyEnum):\n \"\"\"Determines how documents are processed after fetching.\"\"\"\n\n REGULAR = \"REGULAR\" # Full pipeline: chunk → embed → Vespa\n FILE_SYSTEM = \"FILE_SYSTEM\" # Write to file system only (JSON documents)\n RAW_BINARY = \"RAW_BINARY\" # Write raw binary to S3 (no text extraction)\n\n\nclass SyncType(str, PyEnum):\n DOCUMENT_SET = \"document_set\"\n USER_GROUP = \"user_group\"\n CONNECTOR_DELETION = \"connector_deletion\"\n PRUNING = \"pruning\" # not really a sync, but close enough\n EXTERNAL_PERMISSIONS = \"external_permissions\"\n EXTERNAL_GROUP = \"external_group\"\n\n def __str__(self) -> str:\n return self.value\n\n\nclass SyncStatus(str, PyEnum):\n IN_PROGRESS = \"in_progress\"\n SUCCESS = \"success\"\n FAILED = \"failed\"\n CANCELED = \"canceled\"\n\n def is_terminal(self) -> bool:\n terminal_states = {\n SyncStatus.SUCCESS,\n SyncStatus.FAILED,\n }\n return self in terminal_states\n\n\nclass MCPAuthenticationType(str, PyEnum):\n NONE = \"NONE\"\n API_TOKEN = \"API_TOKEN\"\n OAUTH = \"OAUTH\"\n PT_OAUTH = \"PT_OAUTH\" # Pass-Through OAuth\n\n\nclass MCPTransport(str, PyEnum):\n \"\"\"MCP transport types\"\"\"\n\n STDIO = \"STDIO\" # TODO: currently unsupported, need to add a user guide for setup\n SSE = \"SSE\" # Server-Sent Events (deprecated but still used)\n STREAMABLE_HTTP = \"STREAMABLE_HTTP\" # Modern HTTP streaming\n\n\nclass MCPAuthenticationPerformer(str, PyEnum):\n ADMIN = \"ADMIN\"\n PER_USER = \"PER_USER\"\n\n\nclass MCPServerStatus(str, PyEnum):\n CREATED = \"CREATED\" # Server created, needs auth configuration\n AWAITING_AUTH = \"AWAITING_AUTH\" # Auth configured, pending user authentication\n FETCHING_TOOLS = \"FETCHING_TOOLS\" # Auth complete, fetching tools\n CONNECTED = \"CONNECTED\" # Fully configured and connected\n DISCONNECTED = \"DISCONNECTED\" # Server disconnected, but not deleted\n\n\n# Consistent with Celery task statuses\nclass TaskStatus(str, PyEnum):\n PENDING = \"PENDING\"\n STARTED = \"STARTED\"\n SUCCESS = \"SUCCESS\"\n FAILURE = \"FAILURE\"\n\n\nclass IndexModelStatus(str, PyEnum):\n PAST = \"PAST\"\n PRESENT = \"PRESENT\"\n FUTURE = \"FUTURE\"\n\n def is_current(self) -> bool:\n return self == IndexModelStatus.PRESENT\n\n def is_future(self) -> bool:\n return self == IndexModelStatus.FUTURE\n\n\nclass ChatSessionSharedStatus(str, PyEnum):\n PUBLIC = \"public\"\n PRIVATE = \"private\"\n\n\nclass ConnectorCredentialPairStatus(str, PyEnum):\n SCHEDULED = \"SCHEDULED\"\n INITIAL_INDEXING = \"INITIAL_INDEXING\"\n ACTIVE = \"ACTIVE\"\n PAUSED = \"PAUSED\"\n DELETING = \"DELETING\"\n INVALID = \"INVALID\"\n\n @classmethod\n def active_statuses(cls) -> list[\"ConnectorCredentialPairStatus\"]:\n return [\n ConnectorCredentialPairStatus.ACTIVE,\n ConnectorCredentialPairStatus.SCHEDULED,\n ConnectorCredentialPairStatus.INITIAL_INDEXING,\n ]\n\n @classmethod\n def indexable_statuses(self) -> list[\"ConnectorCredentialPairStatus\"]:\n # Superset of active statuses for indexing model swaps\n return self.active_statuses() + [\n ConnectorCredentialPairStatus.PAUSED,\n ]\n\n def is_active(self) -> bool:\n return self in self.active_statuses()\n\n\nclass AccessType(str, PyEnum):\n PUBLIC = \"public\"\n PRIVATE = \"private\"\n SYNC = \"sync\"\n\n\nclass EmbeddingPrecision(str, PyEnum):\n # matches vespa tensor type\n # only support float / bfloat16 for now, since there's not a\n # good reason to specify anything else\n BFLOAT16 = \"bfloat16\"\n FLOAT = \"float\"\n\n\nclass UserFileStatus(str, PyEnum):\n PROCESSING = \"PROCESSING\"\n INDEXING = \"INDEXING\"\n COMPLETED = \"COMPLETED\"\n SKIPPED = \"SKIPPED\"\n FAILED = \"FAILED\"\n CANCELED = \"CANCELED\"\n DELETING = \"DELETING\"\n\n\nclass ThemePreference(str, PyEnum):\n LIGHT = \"light\"\n DARK = \"dark\"\n SYSTEM = \"system\"\n\n\nclass DefaultAppMode(str, PyEnum):\n AUTO = \"AUTO\"\n CHAT = \"CHAT\"\n SEARCH = \"SEARCH\"\n\n\nclass SwitchoverType(str, PyEnum):\n REINDEX = \"reindex\"\n ACTIVE_ONLY = \"active_only\"\n INSTANT = \"instant\"\n\n\nclass OpenSearchDocumentMigrationStatus(str, PyEnum):\n \"\"\"Status for Vespa to OpenSearch migration per document.\"\"\"\n\n PENDING = \"pending\"\n COMPLETED = \"completed\"\n FAILED = \"failed\"\n PERMANENTLY_FAILED = \"permanently_failed\"\n\n\nclass OpenSearchTenantMigrationStatus(str, PyEnum):\n \"\"\"Status for tenant-level OpenSearch migration.\"\"\"\n\n PENDING = \"pending\"\n COMPLETED = \"completed\"\n\n\n# Onyx Build Mode Enums\nclass BuildSessionStatus(str, PyEnum):\n ACTIVE = \"active\"\n IDLE = \"idle\"\n\n\nclass SharingScope(str, PyEnum):\n PRIVATE = \"private\"\n PUBLIC_ORG = \"public_org\"\n PUBLIC_GLOBAL = \"public_global\"\n\n\nclass SandboxStatus(str, PyEnum):\n PROVISIONING = \"provisioning\"\n RUNNING = \"running\"\n SLEEPING = \"sleeping\" # Pod terminated, snapshots saved to S3\n TERMINATED = \"terminated\"\n FAILED = \"failed\"\n\n def is_active(self) -> bool:\n \"\"\"Check if sandbox is in an active state (running).\"\"\"\n return self == SandboxStatus.RUNNING\n\n def is_terminal(self) -> bool:\n \"\"\"Check if sandbox is in a terminal state.\"\"\"\n return self in (SandboxStatus.TERMINATED, SandboxStatus.FAILED)\n\n def is_sleeping(self) -> bool:\n \"\"\"Check if sandbox is sleeping (pod terminated but can be restored).\"\"\"\n return self == SandboxStatus.SLEEPING\n\n\nclass ArtifactType(str, PyEnum):\n WEB_APP = \"web_app\"\n PPTX = \"pptx\"\n DOCX = \"docx\"\n IMAGE = \"image\"\n MARKDOWN = \"markdown\"\n EXCEL = \"excel\"\n\n\nclass HierarchyNodeType(str, PyEnum):\n \"\"\"Types of hierarchy nodes across different sources\"\"\"\n\n # Generic\n FOLDER = \"folder\"\n\n # Root-level type\n SOURCE = \"source\" # Root node for a source (e.g., \"Google Drive\")\n\n # Google Drive\n SHARED_DRIVE = \"shared_drive\"\n MY_DRIVE = \"my_drive\"\n\n # Confluence\n SPACE = \"space\"\n PAGE = \"page\" # Confluence pages can be both hierarchy nodes AND documents\n\n # Jira\n PROJECT = \"project\"\n\n # Notion\n DATABASE = \"database\"\n WORKSPACE = \"workspace\"\n\n # Sharepoint\n SITE = \"site\"\n DRIVE = \"drive\" # Document library within a site\n\n # Slack\n CHANNEL = \"channel\"\n\n\nclass LLMModelFlowType(str, PyEnum):\n CHAT = \"chat\"\n VISION = \"vision\"\n CONTEXTUAL_RAG = \"contextual_rag\"\n\n\nclass HookPoint(str, PyEnum):\n DOCUMENT_INGESTION = \"document_ingestion\"\n QUERY_PROCESSING = \"query_processing\"\n\n\nclass HookFailStrategy(str, PyEnum):\n HARD = \"hard\" # exception propagates, pipeline aborts\n SOFT = \"soft\" # log error, return original input, pipeline continues\n\n\nclass Permission(str, PyEnum):\n \"\"\"\n Permission tokens for group-based authorization.\n 19 tokens total. full_admin_panel_access is an override —\n if present, any permission check passes.\n \"\"\"\n\n # Basic (auto-granted to every new group)\n BASIC_ACCESS = \"basic\"\n\n # Read tokens — implied only, never granted directly\n READ_CONNECTORS = \"read:connectors\"\n READ_DOCUMENT_SETS = \"read:document_sets\"\n READ_AGENTS = \"read:agents\"\n READ_USERS = \"read:users\"\n\n # Add / Manage pairs\n ADD_AGENTS = \"add:agents\"\n MANAGE_AGENTS = \"manage:agents\"\n MANAGE_DOCUMENT_SETS = \"manage:document_sets\"\n ADD_CONNECTORS = \"add:connectors\"\n MANAGE_CONNECTORS = \"manage:connectors\"\n MANAGE_LLMS = \"manage:llms\"\n\n # Toggle tokens\n READ_AGENT_ANALYTICS = \"read:agent_analytics\"\n MANAGE_ACTIONS = \"manage:actions\"\n READ_QUERY_HISTORY = \"read:query_history\"\n MANAGE_USER_GROUPS = \"manage:user_groups\"\n CREATE_USER_API_KEYS = \"create:user_api_keys\"\n CREATE_SERVICE_ACCOUNT_API_KEYS = \"create:service_account_api_keys\"\n CREATE_SLACK_DISCORD_BOTS = \"create:slack_discord_bots\"\n\n # Override — any permission check passes\n FULL_ADMIN_PANEL_ACCESS = \"admin\"\n\n # Permissions that are implied by other grants and must never be stored\n # directly in the permission_grant table.\n IMPLIED: ClassVar[frozenset[Permission]]\n\n\nPermission.IMPLIED = frozenset(\n {\n Permission.READ_CONNECTORS,\n Permission.READ_DOCUMENT_SETS,\n Permission.READ_AGENTS,\n Permission.READ_USERS,\n }\n)\n" }, { "path": "backend/onyx/db/federated.py", "content": "from datetime import datetime\nfrom typing import Any\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import FederatedConnectorSource\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import FederatedConnector\nfrom onyx.db.models import FederatedConnector__DocumentSet\nfrom onyx.db.models import FederatedConnectorOAuthToken\nfrom onyx.federated_connectors.factory import get_federated_connector\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef fetch_federated_connector_by_id(\n federated_connector_id: int, db_session: Session\n) -> FederatedConnector | None:\n \"\"\"Fetch a federated connector by its ID.\"\"\"\n stmt = select(FederatedConnector).where(\n FederatedConnector.id == federated_connector_id\n )\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef fetch_all_federated_connectors(db_session: Session) -> list[FederatedConnector]:\n \"\"\"Fetch all federated connectors with their OAuth tokens and document sets.\"\"\"\n stmt = select(FederatedConnector).options(\n selectinload(FederatedConnector.oauth_tokens),\n selectinload(FederatedConnector.document_sets),\n )\n result = db_session.execute(stmt)\n return list(result.scalars().all())\n\n\ndef fetch_all_federated_connectors_parallel() -> list[FederatedConnector]:\n with get_session_with_current_tenant() as db_session:\n return fetch_all_federated_connectors(db_session)\n\n\ndef validate_federated_connector_credentials(\n source: FederatedConnectorSource,\n credentials: dict[str, Any],\n) -> bool:\n \"\"\"Validate credentials for a federated connector using the connector's validation logic.\"\"\"\n try:\n # the initialization will fail if the credentials are invalid\n get_federated_connector(source, credentials)\n return True\n except Exception as e:\n logger.error(f\"Error validating credentials for source {source}: {e}\")\n return False\n\n\ndef create_federated_connector(\n db_session: Session,\n source: FederatedConnectorSource,\n credentials: dict[str, Any],\n config: dict[str, Any] | None = None,\n) -> FederatedConnector:\n \"\"\"Create a new federated connector with credential and config validation.\"\"\"\n # Validate credentials before creating\n if not validate_federated_connector_credentials(source, credentials):\n raise ValueError(\n f\"Invalid credentials for federated connector source: {source}\"\n )\n\n # Validate config using connector-specific validation\n if config:\n try:\n # Get connector instance to access validate_config method\n connector = get_federated_connector(source, credentials)\n if not connector.validate_config(config):\n raise ValueError(\n f\"Invalid config for federated connector source: {source}\"\n )\n except Exception as e:\n raise ValueError(f\"Config validation failed for {source}: {str(e)}\")\n\n federated_connector = FederatedConnector(\n source=source,\n credentials=credentials,\n config=config or {},\n )\n db_session.add(federated_connector)\n db_session.commit()\n return federated_connector\n\n\ndef update_federated_connector_oauth_token(\n db_session: Session,\n federated_connector_id: int,\n user_id: UUID,\n token: str,\n expires_at: datetime | None = None,\n) -> FederatedConnectorOAuthToken:\n \"\"\"Update or create OAuth token for a federated connector and user.\"\"\"\n # First, try to find existing token for this user and connector\n stmt = select(FederatedConnectorOAuthToken).where(\n FederatedConnectorOAuthToken.federated_connector_id == federated_connector_id,\n FederatedConnectorOAuthToken.user_id == user_id,\n )\n existing_token = db_session.execute(stmt).scalar_one_or_none()\n\n if existing_token:\n # Update existing token\n existing_token.token = token # type: ignore[assignment]\n existing_token.expires_at = expires_at\n db_session.commit()\n return existing_token\n else:\n # Create new token\n oauth_token = FederatedConnectorOAuthToken(\n federated_connector_id=federated_connector_id,\n user_id=user_id,\n token=token,\n expires_at=expires_at,\n )\n db_session.add(oauth_token)\n db_session.commit()\n return oauth_token\n\n\ndef get_federated_connector_oauth_token(\n db_session: Session,\n federated_connector_id: int,\n user_id: UUID,\n) -> FederatedConnectorOAuthToken | None:\n \"\"\"Get OAuth token for a federated connector and user.\"\"\"\n stmt = select(FederatedConnectorOAuthToken).where(\n FederatedConnectorOAuthToken.federated_connector_id == federated_connector_id,\n FederatedConnectorOAuthToken.user_id == user_id,\n )\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef list_federated_connector_oauth_tokens(\n db_session: Session,\n user_id: UUID,\n) -> list[FederatedConnectorOAuthToken]:\n \"\"\"List all OAuth tokens for all federated connectors.\"\"\"\n stmt = (\n select(FederatedConnectorOAuthToken)\n .where(\n FederatedConnectorOAuthToken.user_id == user_id,\n )\n .options(\n joinedload(FederatedConnectorOAuthToken.federated_connector),\n )\n )\n result = db_session.scalars(stmt)\n return list(result)\n\n\ndef create_federated_connector_document_set_mapping(\n db_session: Session,\n federated_connector_id: int,\n document_set_id: int,\n entities: dict[str, Any],\n) -> FederatedConnector__DocumentSet:\n \"\"\"Create a mapping between federated connector and document set with entities.\"\"\"\n mapping = FederatedConnector__DocumentSet(\n federated_connector_id=federated_connector_id,\n document_set_id=document_set_id,\n entities=entities,\n )\n db_session.add(mapping)\n db_session.commit()\n return mapping\n\n\ndef update_federated_connector_document_set_entities(\n db_session: Session,\n federated_connector_id: int,\n document_set_id: int,\n entities: dict[str, Any],\n) -> FederatedConnector__DocumentSet | None:\n \"\"\"Update entities for a federated connector document set mapping.\"\"\"\n stmt = select(FederatedConnector__DocumentSet).where(\n FederatedConnector__DocumentSet.federated_connector_id\n == federated_connector_id,\n FederatedConnector__DocumentSet.document_set_id == document_set_id,\n )\n mapping = db_session.execute(stmt).scalar_one_or_none()\n\n if mapping:\n mapping.entities = entities\n db_session.commit()\n return mapping\n\n return None\n\n\ndef get_federated_connector_document_set_mappings(\n db_session: Session,\n federated_connector_id: int,\n) -> list[FederatedConnector__DocumentSet]:\n \"\"\"Get all document set mappings for a federated connector.\"\"\"\n stmt = select(FederatedConnector__DocumentSet).where(\n FederatedConnector__DocumentSet.federated_connector_id == federated_connector_id\n )\n result = db_session.execute(stmt)\n return list(result.scalars().all())\n\n\ndef delete_federated_connector_document_set_mapping(\n db_session: Session,\n federated_connector_id: int,\n document_set_id: int,\n) -> bool:\n \"\"\"Delete a federated connector document set mapping.\"\"\"\n stmt = select(FederatedConnector__DocumentSet).where(\n FederatedConnector__DocumentSet.federated_connector_id\n == federated_connector_id,\n FederatedConnector__DocumentSet.document_set_id == document_set_id,\n )\n mapping = db_session.execute(stmt).scalar_one_or_none()\n\n if mapping:\n db_session.delete(mapping)\n db_session.commit()\n return True\n\n return False\n\n\ndef get_federated_connector_document_set_mappings_by_document_set_names(\n db_session: Session,\n document_set_names: list[str],\n) -> list[FederatedConnector__DocumentSet]:\n \"\"\"Get all document set mappings for a federated connector by document set names.\"\"\"\n stmt = (\n select(FederatedConnector__DocumentSet)\n .join(\n DocumentSet,\n FederatedConnector__DocumentSet.document_set_id == DocumentSet.id,\n )\n .options(joinedload(FederatedConnector__DocumentSet.federated_connector))\n .where(DocumentSet.name.in_(document_set_names))\n )\n result = db_session.scalars(stmt)\n # Use unique() because joinedload can cause duplicate rows\n return list(result.unique())\n\n\ndef update_federated_connector(\n db_session: Session,\n federated_connector_id: int,\n credentials: dict[str, Any] | None = None,\n config: dict[str, Any] | None = None,\n) -> FederatedConnector | None:\n \"\"\"Update a federated connector with credential and config validation.\"\"\"\n federated_connector = fetch_federated_connector_by_id(\n federated_connector_id, db_session\n )\n if not federated_connector:\n return None\n\n # Use provided credentials if updating them, otherwise use existing credentials\n # This is needed to instantiate the connector for config validation when only config is being updated\n creds_to_use = (\n credentials\n if credentials is not None\n else (\n federated_connector.credentials.get_value(apply_mask=False)\n if federated_connector.credentials\n else {}\n )\n )\n\n if credentials is not None:\n # Validate credentials before updating\n if not validate_federated_connector_credentials(\n federated_connector.source, credentials\n ):\n raise ValueError(\n f\"Invalid credentials for federated connector source: {federated_connector.source}\"\n )\n federated_connector.credentials = credentials # type: ignore[assignment]\n\n if config is not None:\n # Validate config using connector-specific validation\n try:\n # Get connector instance to access validate_config method\n connector = get_federated_connector(\n federated_connector.source, creds_to_use\n )\n if not connector.validate_config(config):\n raise ValueError(\n f\"Invalid config for federated connector source: {federated_connector.source}\"\n )\n except Exception as e:\n raise ValueError(\n f\"Config validation failed for {federated_connector.source}: {str(e)}\"\n )\n federated_connector.config = config\n\n db_session.commit()\n return federated_connector\n\n\ndef delete_federated_connector(\n db_session: Session,\n federated_connector_id: int,\n) -> bool:\n \"\"\"Delete a federated connector and all its related data.\"\"\"\n federated_connector = fetch_federated_connector_by_id(\n federated_connector_id, db_session\n )\n if not federated_connector:\n return False\n\n # Delete related OAuth tokens (cascade should handle this)\n # Delete related document set mappings (cascade should handle this)\n db_session.delete(federated_connector)\n db_session.commit()\n return True\n" }, { "path": "backend/onyx/db/feedback.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import and_\nfrom sqlalchemy import asc\nfrom sqlalchemy import delete\nfrom sqlalchemy import desc\nfrom sqlalchemy import exists\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import SearchFeedbackType\nfrom onyx.db.chat import get_chat_message\nfrom onyx.db.enums import AccessType\nfrom onyx.db.models import ChatMessageFeedback\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Document as DbDocument\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import DocumentRetrievalFeedback\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup__ConnectorCredentialPair\nfrom onyx.db.models import UserRole\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _fetch_db_doc_by_id(doc_id: str, db_session: Session) -> DbDocument:\n stmt = select(DbDocument).where(DbDocument.id == doc_id)\n result = db_session.execute(stmt)\n doc = result.scalar_one_or_none()\n\n if not doc:\n raise ValueError(\"Invalid Document ID Provided\")\n\n return doc\n\n\ndef _add_user_filters(stmt: Select, user: User, get_editable: bool = True) -> Select:\n if user.role == UserRole.ADMIN:\n return stmt\n\n stmt = stmt.distinct()\n DocByCC = aliased(DocumentByConnectorCredentialPair)\n CCPair = aliased(ConnectorCredentialPair)\n UG__CCpair = aliased(UserGroup__ConnectorCredentialPair)\n User__UG = aliased(User__UserGroup)\n\n \"\"\"\n Here we select documents by relation:\n User -> User__UserGroup -> UserGroup__ConnectorCredentialPair ->\n ConnectorCredentialPair -> DocumentByConnectorCredentialPair -> Document\n \"\"\"\n stmt = (\n stmt.outerjoin(DocByCC, DocByCC.id == DbDocument.id)\n .outerjoin(\n CCPair,\n and_(\n CCPair.connector_id == DocByCC.connector_id,\n CCPair.credential_id == DocByCC.credential_id,\n ),\n )\n .outerjoin(UG__CCpair, UG__CCpair.cc_pair_id == CCPair.id)\n .outerjoin(User__UG, User__UG.user_group_id == UG__CCpair.user_group_id)\n )\n\n \"\"\"\n Filter Documents by:\n - if the user is in the user_group that owns the object\n - if the user is not a global_curator, they must also have a curator relationship\n to the user_group\n - if editing is being done, we also filter out objects that are owned by groups\n that the user isn't a curator for\n - if we are not editing, we show all objects in the groups the user is a curator\n for (as well as public objects as well)\n \"\"\"\n\n # Anonymous users only see public documents\n if user.is_anonymous:\n where_clause = CCPair.access_type == AccessType.PUBLIC\n return stmt.where(where_clause)\n\n where_clause = User__UG.user_id == user.id\n if user.role == UserRole.CURATOR and get_editable:\n where_clause &= User__UG.is_curator == True # noqa: E712\n if get_editable:\n user_groups = select(User__UG.user_group_id).where(User__UG.user_id == user.id)\n where_clause &= (\n ~exists()\n .where(UG__CCpair.cc_pair_id == CCPair.id)\n .where(~UG__CCpair.user_group_id.in_(user_groups))\n .correlate(CCPair)\n )\n else:\n where_clause |= CCPair.access_type == AccessType.PUBLIC\n\n return stmt.where(where_clause)\n\n\ndef fetch_docs_ranked_by_boost_for_user(\n db_session: Session,\n user: User,\n ascending: bool = False,\n limit: int = 100,\n) -> list[DbDocument]:\n order_func = asc if ascending else desc\n stmt = select(DbDocument)\n\n stmt = _add_user_filters(stmt=stmt, user=user, get_editable=False)\n\n stmt = stmt.order_by(\n order_func(DbDocument.boost), order_func(DbDocument.semantic_id)\n )\n stmt = stmt.limit(limit)\n result = db_session.execute(stmt)\n doc_list = result.scalars().all()\n\n return list(doc_list)\n\n\ndef update_document_boost_for_user(\n db_session: Session,\n document_id: str,\n boost: int,\n user: User,\n) -> None:\n stmt = select(DbDocument).where(DbDocument.id == document_id)\n stmt = _add_user_filters(stmt, user, get_editable=True)\n result: DbDocument | None = db_session.execute(stmt).scalar_one_or_none()\n if result is None:\n raise HTTPException(\n status_code=400, detail=\"Document is not editable by this user\"\n )\n\n result.boost = boost\n\n # updating last_modified triggers sync\n # TODO: Should this submit to the queue directly so that the UI can update?\n result.last_modified = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef update_document_hidden_for_user(\n db_session: Session,\n document_id: str,\n hidden: bool,\n user: User,\n) -> None:\n stmt = select(DbDocument).where(DbDocument.id == document_id)\n stmt = _add_user_filters(stmt, user, get_editable=True)\n result = db_session.execute(stmt).scalar_one_or_none()\n if result is None:\n raise HTTPException(\n status_code=400, detail=\"Document is not editable by this user\"\n )\n\n result.hidden = hidden\n\n # updating last_modified triggers sync\n # TODO: Should this submit to the queue directly so that the UI can update?\n result.last_modified = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef create_doc_retrieval_feedback(\n message_id: int,\n document_id: str,\n document_rank: int,\n db_session: Session,\n clicked: bool = False,\n feedback: SearchFeedbackType | None = None,\n) -> None:\n \"\"\"Creates a new Document feedback row and updates the boost value in Postgres and Vespa\"\"\"\n db_doc = _fetch_db_doc_by_id(document_id, db_session)\n\n retrieval_feedback = DocumentRetrievalFeedback(\n chat_message_id=message_id,\n document_id=document_id,\n document_rank=document_rank,\n clicked=clicked,\n feedback=feedback,\n )\n\n if feedback is not None:\n if feedback == SearchFeedbackType.ENDORSE:\n db_doc.boost += 1\n elif feedback == SearchFeedbackType.REJECT:\n db_doc.boost -= 1\n elif feedback == SearchFeedbackType.HIDE:\n db_doc.hidden = True\n elif feedback == SearchFeedbackType.UNHIDE:\n db_doc.hidden = False\n else:\n raise ValueError(\"Unhandled document feedback type\")\n\n if feedback in [\n SearchFeedbackType.ENDORSE,\n SearchFeedbackType.REJECT,\n SearchFeedbackType.HIDE,\n ]:\n # updating last_modified triggers sync\n # TODO: Should this submit to the queue directly so that the UI can update?\n db_doc.last_modified = datetime.now(timezone.utc)\n\n db_session.add(retrieval_feedback)\n db_session.commit()\n\n\ndef delete_document_feedback_for_documents__no_commit(\n document_ids: list[str], db_session: Session\n) -> None:\n \"\"\"NOTE: does not commit transaction so that this can be used as part of a\n larger transaction block.\"\"\"\n stmt = delete(DocumentRetrievalFeedback).where(\n DocumentRetrievalFeedback.document_id.in_(document_ids)\n )\n db_session.execute(stmt)\n\n\ndef create_chat_message_feedback(\n is_positive: bool | None,\n feedback_text: str | None,\n chat_message_id: int,\n user_id: UUID | None,\n db_session: Session,\n # Slack user requested help from human\n required_followup: bool | None = None,\n predefined_feedback: str | None = None, # Added predefined_feedback parameter\n) -> None:\n if (\n is_positive is None\n and feedback_text is None\n and required_followup is None\n and predefined_feedback is None\n ):\n raise ValueError(\"No feedback provided\")\n\n chat_message = get_chat_message(\n chat_message_id=chat_message_id, user_id=user_id, db_session=db_session\n )\n\n if chat_message.message_type != MessageType.ASSISTANT:\n raise ValueError(\"Can only provide feedback on LLM Outputs\")\n\n message_feedback = ChatMessageFeedback(\n chat_message_id=chat_message_id,\n is_positive=is_positive,\n feedback_text=feedback_text,\n required_followup=required_followup,\n predefined_feedback=predefined_feedback,\n )\n\n db_session.add(message_feedback)\n db_session.commit()\n\n\ndef remove_chat_message_feedback(\n chat_message_id: int,\n user_id: UUID | None,\n db_session: Session,\n) -> None:\n \"\"\"Remove all feedback for a chat message.\"\"\"\n chat_message = get_chat_message(\n chat_message_id=chat_message_id, user_id=user_id, db_session=db_session\n )\n\n if chat_message.message_type != MessageType.ASSISTANT:\n raise ValueError(\"Can only remove feedback from LLM Outputs\")\n\n # Delete all feedback for this message\n db_session.query(ChatMessageFeedback).filter(\n ChatMessageFeedback.chat_message_id == chat_message_id\n ).delete()\n\n db_session.commit()\n" }, { "path": "backend/onyx/db/file_content.py", "content": "from sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import FileContent\n\n\ndef get_file_content_by_file_id(\n file_id: str,\n db_session: Session,\n) -> FileContent:\n record = db_session.query(FileContent).filter_by(file_id=file_id).first()\n if not record:\n raise RuntimeError(\n f\"File content for file_id {file_id} does not exist or was deleted\"\n )\n return record\n\n\ndef get_file_content_by_file_id_optional(\n file_id: str,\n db_session: Session,\n) -> FileContent | None:\n return db_session.query(FileContent).filter_by(file_id=file_id).first()\n\n\ndef upsert_file_content(\n file_id: str,\n lobj_oid: int,\n file_size: int,\n db_session: Session,\n) -> FileContent:\n \"\"\"Atomic upsert using INSERT ... ON CONFLICT DO UPDATE to avoid\n race conditions when concurrent calls target the same file_id.\"\"\"\n stmt = insert(FileContent).values(\n file_id=file_id,\n lobj_oid=lobj_oid,\n file_size=file_size,\n )\n stmt = stmt.on_conflict_do_update(\n index_elements=[FileContent.file_id],\n set_={\n \"lobj_oid\": stmt.excluded.lobj_oid,\n \"file_size\": stmt.excluded.file_size,\n },\n )\n db_session.execute(stmt)\n\n # Return the merged ORM instance so callers can inspect the result\n return db_session.get(FileContent, file_id) # type: ignore[return-value]\n\n\ndef transfer_file_content_file_id(\n old_file_id: str,\n new_file_id: str,\n db_session: Session,\n) -> None:\n \"\"\"Move a file_content row from old_file_id to new_file_id in-place.\n\n This avoids creating a duplicate row that shares the same Large Object OID,\n keeping OID ownership unique at all times. The caller must ensure that\n new_file_id already exists in file_record (FK target).\"\"\"\n rows = (\n db_session.query(FileContent)\n .filter_by(file_id=old_file_id)\n .update({\"file_id\": new_file_id})\n )\n if not rows:\n raise RuntimeError(\n f\"File content for file_id {old_file_id} does not exist or was deleted\"\n )\n\n\ndef delete_file_content_by_file_id(\n file_id: str,\n db_session: Session,\n) -> None:\n db_session.query(FileContent).filter_by(file_id=file_id).delete()\n" }, { "path": "backend/onyx/db/file_record.py", "content": "from sqlalchemy import and_\nfrom sqlalchemy import select\nfrom sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.task_utils import QUERY_REPORT_NAME_PREFIX\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import FileType\nfrom onyx.db.models import FileRecord\n\n\ndef get_query_history_export_files(\n db_session: Session,\n) -> list[FileRecord]:\n return list(\n db_session.scalars(\n select(FileRecord).where(\n and_(\n FileRecord.file_id.like(f\"{QUERY_REPORT_NAME_PREFIX}-%\"),\n FileRecord.file_type == FileType.CSV,\n FileRecord.file_origin == FileOrigin.QUERY_HISTORY_CSV,\n )\n )\n )\n )\n\n\ndef get_filerecord_by_file_id_optional(\n file_id: str,\n db_session: Session,\n) -> FileRecord | None:\n return db_session.query(FileRecord).filter_by(file_id=file_id).first()\n\n\ndef get_filerecord_by_file_id(\n file_id: str,\n db_session: Session,\n) -> FileRecord:\n filestore = db_session.query(FileRecord).filter_by(file_id=file_id).first()\n\n if not filestore:\n raise RuntimeError(f\"File by id {file_id} does not exist or was deleted\")\n\n return filestore\n\n\ndef get_filerecord_by_prefix(\n prefix: str,\n db_session: Session,\n) -> list[FileRecord]:\n if not prefix:\n return db_session.query(FileRecord).all()\n return (\n db_session.query(FileRecord).filter(FileRecord.file_id.like(f\"{prefix}%\")).all()\n )\n\n\ndef delete_filerecord_by_file_id(\n file_id: str,\n db_session: Session,\n) -> None:\n db_session.query(FileRecord).filter_by(file_id=file_id).delete()\n\n\ndef upsert_filerecord(\n file_id: str,\n display_name: str,\n file_origin: FileOrigin,\n file_type: str,\n bucket_name: str,\n object_key: str,\n db_session: Session,\n file_metadata: dict | None = None,\n) -> FileRecord:\n \"\"\"Atomic upsert using INSERT ... ON CONFLICT DO UPDATE to avoid\n race conditions when concurrent calls target the same file_id.\"\"\"\n stmt = insert(FileRecord).values(\n file_id=file_id,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_metadata=file_metadata,\n bucket_name=bucket_name,\n object_key=object_key,\n )\n stmt = stmt.on_conflict_do_update(\n index_elements=[FileRecord.file_id],\n set_={\n \"display_name\": stmt.excluded.display_name,\n \"file_origin\": stmt.excluded.file_origin,\n \"file_type\": stmt.excluded.file_type,\n \"file_metadata\": stmt.excluded.file_metadata,\n \"bucket_name\": stmt.excluded.bucket_name,\n \"object_key\": stmt.excluded.object_key,\n },\n )\n db_session.execute(stmt)\n\n return db_session.get(FileRecord, file_id) # type: ignore[return-value]\n" }, { "path": "backend/onyx/db/hierarchy.py", "content": "\"\"\"CRUD operations for HierarchyNode.\"\"\"\n\nfrom collections import defaultdict\n\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.engine import CursorResult\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import HierarchyNode as PydanticHierarchyNode\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.db.models import Document\nfrom onyx.db.models import HierarchyNode\nfrom onyx.db.models import HierarchyNodeByConnectorCredentialPair\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\nlogger = setup_logger()\n\n# Sources where hierarchy nodes can also be documents.\n# For these sources, pages/items can be both a hierarchy node (with children)\n# AND a document with indexed content. For example:\n# - Notion: Pages with child pages are hierarchy nodes, but also documents\n# - Confluence: Pages can have child pages and also contain content\n# Other sources like Google Drive have folders as hierarchy nodes, but folders\n# are not documents themselves.\nSOURCES_WITH_HIERARCHY_NODE_DOCUMENTS: set[DocumentSource] = {\n DocumentSource.NOTION,\n DocumentSource.CONFLUENCE,\n}\n\n\ndef _get_source_display_name(source: DocumentSource) -> str:\n \"\"\"Get a human-readable display name for a source type.\"\"\"\n return source.value.replace(\"_\", \" \").title()\n\n\ndef get_hierarchy_node_by_raw_id(\n db_session: Session,\n raw_node_id: str,\n source: DocumentSource,\n) -> HierarchyNode | None:\n \"\"\"Get a hierarchy node by its raw ID and source.\"\"\"\n stmt = select(HierarchyNode).where(\n HierarchyNode.raw_node_id == raw_node_id,\n HierarchyNode.source == source,\n )\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef get_source_hierarchy_node(\n db_session: Session,\n source: DocumentSource,\n) -> HierarchyNode | None:\n \"\"\"Get the SOURCE-type root node for a given source.\"\"\"\n stmt = select(HierarchyNode).where(\n HierarchyNode.source == source,\n HierarchyNode.node_type == HierarchyNodeType.SOURCE,\n )\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef ensure_source_node_exists(\n db_session: Session,\n source: DocumentSource,\n commit: bool = True,\n) -> HierarchyNode:\n \"\"\"\n Ensure that a SOURCE-type root node exists for the given source.\n\n This function is idempotent - it will return the existing SOURCE node if one\n exists, or create a new one if not.\n\n The SOURCE node is the root of the hierarchy tree for a given source type\n (e.g., \"Google Drive\", \"Confluence\"). All other hierarchy nodes for that\n source should ultimately have this node as an ancestor.\n\n For the SOURCE node:\n - raw_node_id is set to the source name (e.g., \"google_drive\")\n - parent_id is None (it's the root)\n - display_name is a human-readable version (e.g., \"Google Drive\")\n\n Args:\n db_session: SQLAlchemy session\n source: The document source type\n commit: Whether to commit the transaction\n\n Returns:\n The existing or newly created SOURCE-type HierarchyNode\n \"\"\"\n # Try to get existing SOURCE node first\n existing_node = get_source_hierarchy_node(db_session, source)\n if existing_node:\n return existing_node\n\n # Create the SOURCE node\n display_name = _get_source_display_name(source)\n\n source_node = HierarchyNode(\n raw_node_id=source.value, # Use source name as raw_node_id\n display_name=display_name,\n link=None,\n source=source,\n node_type=HierarchyNodeType.SOURCE,\n document_id=None,\n parent_id=None, # SOURCE nodes have no parent\n )\n\n db_session.add(source_node)\n\n # Flush to get the ID and detect any race conditions\n try:\n db_session.flush()\n except Exception:\n # Race condition - another worker created it. Roll back and fetch.\n db_session.rollback()\n existing_node = get_source_hierarchy_node(db_session, source)\n if existing_node:\n return existing_node\n # If still not found, re-raise the original exception\n raise\n\n if commit:\n db_session.commit()\n\n logger.info(\n f\"Created SOURCE hierarchy node for {source.value}: id={source_node.id}, display_name={display_name}\"\n )\n\n return source_node\n\n\ndef resolve_parent_hierarchy_node_id(\n db_session: Session,\n raw_parent_id: str | None,\n source: DocumentSource,\n) -> int | None:\n \"\"\"\n Resolve a raw_parent_id to a database HierarchyNode ID.\n\n If raw_parent_id is None, returns the SOURCE node ID for backward compatibility.\n If the parent node doesn't exist, returns the SOURCE node ID as fallback.\n \"\"\"\n if raw_parent_id is None:\n # No parent specified - use the SOURCE node\n source_node = get_source_hierarchy_node(db_session, source)\n return source_node.id if source_node else None\n\n parent_node = get_hierarchy_node_by_raw_id(db_session, raw_parent_id, source)\n if parent_node:\n return parent_node.id\n\n # Parent not found - fall back to SOURCE node\n logger.warning(\n f\"Parent hierarchy node not found: raw_id={raw_parent_id}, source={source}. Falling back to SOURCE node.\"\n )\n source_node = get_source_hierarchy_node(db_session, source)\n return source_node.id if source_node else None\n\n\ndef upsert_parents(\n db_session: Session,\n node: PydanticHierarchyNode,\n source: DocumentSource,\n node_by_id: dict[str, PydanticHierarchyNode],\n done_ids: set[str],\n is_connector_public: bool = False,\n) -> None:\n \"\"\"\n Upsert the parents of a hierarchy node.\n \"\"\"\n if (\n node.node_type == HierarchyNodeType.SOURCE\n or (node.raw_parent_id not in node_by_id)\n or (node.raw_parent_id in done_ids)\n ):\n return\n parent_node = node_by_id[node.raw_parent_id]\n upsert_parents(\n db_session,\n parent_node,\n source,\n node_by_id,\n done_ids,\n is_connector_public=is_connector_public,\n )\n upsert_hierarchy_node(\n db_session,\n parent_node,\n source,\n commit=False,\n is_connector_public=is_connector_public,\n )\n done_ids.add(parent_node.raw_node_id)\n\n\ndef upsert_hierarchy_node(\n db_session: Session,\n node: PydanticHierarchyNode,\n source: DocumentSource,\n commit: bool = True,\n is_connector_public: bool = False,\n) -> HierarchyNode:\n \"\"\"\n Upsert a hierarchy node from a Pydantic model.\n\n If a node with the same raw_node_id and source exists, updates it.\n Otherwise, creates a new node.\n\n Args:\n db_session: SQLAlchemy session\n node: The Pydantic hierarchy node to upsert\n source: Document source type\n commit: Whether to commit the transaction\n is_connector_public: If True, the connector is public (organization-wide access)\n and all hierarchy nodes should be marked as public regardless of their\n external_access settings. This ensures nodes from public connectors are\n accessible to all users.\n \"\"\"\n # Resolve parent_id from raw_parent_id\n parent_id = (\n None\n if node.node_type == HierarchyNodeType.SOURCE\n else resolve_parent_hierarchy_node_id(db_session, node.raw_parent_id, source)\n )\n\n # For public connectors, all nodes are public\n # Otherwise, extract permission fields from external_access if present\n if is_connector_public:\n is_public = True\n external_user_emails: list[str] | None = None\n external_user_group_ids: list[str] | None = None\n elif node.external_access:\n is_public = node.external_access.is_public\n external_user_emails = (\n list(node.external_access.external_user_emails)\n if node.external_access.external_user_emails\n else None\n )\n external_user_group_ids = (\n list(node.external_access.external_user_group_ids)\n if node.external_access.external_user_group_ids\n else None\n )\n else:\n is_public = False\n external_user_emails = None\n external_user_group_ids = None\n\n # Check if node already exists\n existing_node = get_hierarchy_node_by_raw_id(db_session, node.raw_node_id, source)\n\n if existing_node:\n # Update existing node\n existing_node.display_name = node.display_name\n existing_node.link = node.link\n existing_node.node_type = node.node_type\n existing_node.parent_id = parent_id\n # Update permission fields\n existing_node.is_public = is_public\n existing_node.external_user_emails = external_user_emails\n existing_node.external_user_group_ids = external_user_group_ids\n hierarchy_node = existing_node\n else:\n # Create new node\n hierarchy_node = HierarchyNode(\n raw_node_id=node.raw_node_id,\n display_name=node.display_name,\n link=node.link,\n source=source,\n node_type=node.node_type,\n parent_id=parent_id,\n is_public=is_public,\n external_user_emails=external_user_emails,\n external_user_group_ids=external_user_group_ids,\n )\n db_session.add(hierarchy_node)\n\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n return hierarchy_node\n\n\ndef upsert_hierarchy_nodes_batch(\n db_session: Session,\n nodes: list[PydanticHierarchyNode],\n source: DocumentSource,\n commit: bool = True,\n is_connector_public: bool = False,\n) -> list[HierarchyNode]:\n \"\"\"\n Batch upsert hierarchy nodes.\n\n Note: This function requires that for each node passed in, all\n its ancestors exist in either the database or elsewhere in the nodes list.\n This function handles parent dependencies for you as long as that condition is met\n (so you don't need to worry about parent nodes appearing before their children in the list).\n\n Args:\n db_session: SQLAlchemy session\n nodes: List of Pydantic hierarchy nodes to upsert\n source: Document source type\n commit: Whether to commit the transaction\n is_connector_public: If True, the connector is public (organization-wide access)\n and all hierarchy nodes should be marked as public regardless of their\n external_access settings.\n \"\"\"\n node_by_id = {}\n for node in nodes:\n if node.node_type != HierarchyNodeType.SOURCE:\n node_by_id[node.raw_node_id] = node\n done_ids = set[str]()\n\n results = []\n for node in nodes:\n if node.raw_node_id in done_ids:\n continue\n upsert_parents(\n db_session,\n node,\n source,\n node_by_id,\n done_ids,\n is_connector_public=is_connector_public,\n )\n hierarchy_node = upsert_hierarchy_node(\n db_session,\n node,\n source,\n commit=False,\n is_connector_public=is_connector_public,\n )\n done_ids.add(node.raw_node_id)\n results.append(hierarchy_node)\n\n if commit:\n db_session.commit()\n\n return results\n\n\ndef link_hierarchy_nodes_to_documents(\n db_session: Session,\n document_ids: list[str],\n source: DocumentSource,\n commit: bool = True,\n) -> int:\n \"\"\"\n Link hierarchy nodes to their corresponding documents.\n\n For connectors like Notion and Confluence where pages can be both hierarchy nodes\n AND documents, we need to set the document_id field on hierarchy nodes after the\n documents are created. This is because hierarchy nodes are processed before documents,\n and the FK constraint on document_id requires the document to exist first.\n\n Args:\n db_session: SQLAlchemy session\n document_ids: List of document IDs that were just created/updated\n source: The document source (e.g., NOTION, CONFLUENCE)\n commit: Whether to commit the transaction\n\n Returns:\n Number of hierarchy nodes that were linked to documents\n \"\"\"\n # Skip for sources where hierarchy nodes cannot also be documents\n if source not in SOURCES_WITH_HIERARCHY_NODE_DOCUMENTS:\n return 0\n\n if not document_ids:\n return 0\n\n # Find hierarchy nodes where raw_node_id matches a document_id\n # These are pages that are both hierarchy nodes and documents\n stmt = select(HierarchyNode).where(\n HierarchyNode.source == source,\n HierarchyNode.raw_node_id.in_(document_ids),\n HierarchyNode.document_id.is_(None), # Only update if not already linked\n )\n nodes_to_update = list(db_session.execute(stmt).scalars().all())\n\n # Update document_id for each matching node\n for node in nodes_to_update:\n node.document_id = node.raw_node_id\n\n if commit:\n db_session.commit()\n\n if nodes_to_update:\n logger.debug(\n f\"Linked {len(nodes_to_update)} hierarchy nodes to documents for source {source.value}\"\n )\n\n return len(nodes_to_update)\n\n\ndef get_hierarchy_node_children(\n db_session: Session,\n parent_id: int,\n limit: int = 100,\n offset: int = 0,\n) -> list[HierarchyNode]:\n \"\"\"Get children of a hierarchy node, paginated.\"\"\"\n stmt = (\n select(HierarchyNode)\n .where(HierarchyNode.parent_id == parent_id)\n .order_by(HierarchyNode.display_name)\n .limit(limit)\n .offset(offset)\n )\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef get_hierarchy_node_by_id(\n db_session: Session,\n node_id: int,\n) -> HierarchyNode | None:\n \"\"\"Get a hierarchy node by its database ID.\"\"\"\n return db_session.get(HierarchyNode, node_id)\n\n\ndef get_root_hierarchy_nodes_for_source(\n db_session: Session,\n source: DocumentSource,\n) -> list[HierarchyNode]:\n \"\"\"Get all root-level hierarchy nodes for a source (children of SOURCE node).\"\"\"\n source_node = get_source_hierarchy_node(db_session, source)\n if not source_node:\n return []\n\n return get_hierarchy_node_children(db_session, source_node.id)\n\n\ndef get_all_hierarchy_nodes_for_source(\n db_session: Session,\n source: DocumentSource,\n) -> list[HierarchyNode]:\n \"\"\"\n Get ALL hierarchy nodes for a given source.\n\n This is used to populate the Redis cache. Returns all nodes including\n the SOURCE-type root node.\n\n Args:\n db_session: SQLAlchemy session\n source: The document source to get nodes for\n\n Returns:\n List of all HierarchyNode objects for the source\n \"\"\"\n stmt = select(HierarchyNode).where(HierarchyNode.source == source)\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef _get_accessible_hierarchy_nodes_for_source(\n db_session: Session,\n source: DocumentSource,\n user_email: str, # noqa: ARG001\n external_group_ids: list[str], # noqa: ARG001\n) -> list[HierarchyNode]:\n \"\"\"\n MIT version: Returns all hierarchy nodes for the source without permission filtering.\n\n In the MIT version, permission checks are not performed on hierarchy nodes.\n The EE version overrides this to apply permission filtering based on user\n email and external group IDs.\n\n Args:\n db_session: SQLAlchemy session\n source: Document source type\n user_email: User's email (unused in MIT version)\n external_group_ids: User's external group IDs (unused in MIT version)\n\n Returns:\n List of all HierarchyNode objects for the source\n \"\"\"\n stmt = select(HierarchyNode).where(HierarchyNode.source == source)\n stmt = stmt.order_by(HierarchyNode.display_name)\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef get_accessible_hierarchy_nodes_for_source(\n db_session: Session,\n source: DocumentSource,\n user_email: str,\n external_group_ids: list[str],\n) -> list[HierarchyNode]:\n \"\"\"\n Get hierarchy nodes for a source that are accessible to the user.\n\n Uses fetch_versioned_implementation to get the appropriate version:\n - MIT version: Returns all nodes (no permission filtering)\n - EE version: Filters based on user email and external group IDs\n \"\"\"\n versioned_fn = fetch_versioned_implementation(\n \"onyx.db.hierarchy\", \"_get_accessible_hierarchy_nodes_for_source\"\n )\n return versioned_fn(db_session, source, user_email, external_group_ids)\n\n\ndef get_document_parent_hierarchy_node_ids(\n db_session: Session,\n document_ids: list[str],\n) -> dict[str, int | None]:\n \"\"\"\n Get the parent_hierarchy_node_id for multiple documents in a single query.\n\n Args:\n db_session: SQLAlchemy session\n document_ids: List of document IDs to look up\n\n Returns:\n Dict mapping document_id -> parent_hierarchy_node_id (or None if not set)\n \"\"\"\n\n if not document_ids:\n return {}\n\n stmt = select(Document.id, Document.parent_hierarchy_node_id).where(\n Document.id.in_(document_ids)\n )\n results = db_session.execute(stmt).all()\n\n return {doc_id: parent_id for doc_id, parent_id in results}\n\n\ndef update_document_parent_hierarchy_nodes(\n db_session: Session,\n doc_parent_map: dict[str, int | None],\n commit: bool = True,\n) -> int:\n \"\"\"Bulk-update Document.parent_hierarchy_node_id for multiple documents.\n\n Only updates rows whose current value differs from the desired value to\n avoid unnecessary writes.\n\n Args:\n db_session: SQLAlchemy session\n doc_parent_map: Mapping of document_id → desired parent_hierarchy_node_id\n commit: Whether to commit the transaction\n\n Returns:\n Number of documents actually updated\n \"\"\"\n if not doc_parent_map:\n return 0\n\n doc_ids = list(doc_parent_map.keys())\n existing = get_document_parent_hierarchy_node_ids(db_session, doc_ids)\n\n by_parent: dict[int | None, list[str]] = defaultdict(list)\n for doc_id, desired_parent_id in doc_parent_map.items():\n current = existing.get(doc_id)\n if current == desired_parent_id or doc_id not in existing:\n continue\n by_parent[desired_parent_id].append(doc_id)\n\n updated = 0\n for desired_parent_id, ids in by_parent.items():\n db_session.query(Document).filter(Document.id.in_(ids)).update(\n {Document.parent_hierarchy_node_id: desired_parent_id},\n synchronize_session=False,\n )\n updated += len(ids)\n\n if commit:\n db_session.commit()\n elif updated:\n db_session.flush()\n\n return updated\n\n\ndef update_hierarchy_node_permissions(\n db_session: Session,\n raw_node_id: str,\n source: DocumentSource,\n is_public: bool,\n external_user_emails: list[str] | None,\n external_user_group_ids: list[str] | None,\n commit: bool = True,\n) -> bool:\n \"\"\"\n Update permissions for an existing hierarchy node.\n\n This is used during permission sync to update folder permissions\n without needing the full Pydantic HierarchyNode model.\n\n Args:\n db_session: SQLAlchemy session\n raw_node_id: Raw node ID from the source system\n source: Document source type\n is_public: Whether the node is public\n external_user_emails: List of user emails with access\n external_user_group_ids: List of group IDs with access\n commit: Whether to commit the transaction\n\n Returns:\n True if the node was found and updated, False if not found\n \"\"\"\n existing_node = get_hierarchy_node_by_raw_id(db_session, raw_node_id, source)\n\n if not existing_node:\n logger.warning(\n f\"Hierarchy node not found for permission update: raw_node_id={raw_node_id}, source={source}\"\n )\n return False\n\n existing_node.is_public = is_public\n existing_node.external_user_emails = external_user_emails\n existing_node.external_user_group_ids = external_user_group_ids\n\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n return True\n\n\ndef upsert_hierarchy_node_cc_pair_entries(\n db_session: Session,\n hierarchy_node_ids: list[int],\n connector_id: int,\n credential_id: int,\n commit: bool = True,\n) -> None:\n \"\"\"Insert rows into HierarchyNodeByConnectorCredentialPair, ignoring conflicts.\n\n This records that the given cc_pair \"owns\" these hierarchy nodes. Used by\n indexing, pruning, and hierarchy-fetching paths.\n \"\"\"\n if not hierarchy_node_ids:\n return\n\n _M = HierarchyNodeByConnectorCredentialPair\n stmt = pg_insert(_M).values(\n [\n {\n _M.hierarchy_node_id: node_id,\n _M.connector_id: connector_id,\n _M.credential_id: credential_id,\n }\n for node_id in hierarchy_node_ids\n ]\n )\n stmt = stmt.on_conflict_do_nothing()\n db_session.execute(stmt)\n\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n\ndef remove_stale_hierarchy_node_cc_pair_entries(\n db_session: Session,\n connector_id: int,\n credential_id: int,\n live_hierarchy_node_ids: set[int],\n commit: bool = True,\n) -> int:\n \"\"\"Delete join-table rows for this cc_pair that are NOT in the live set.\n\n If ``live_hierarchy_node_ids`` is empty ALL rows for the cc_pair are deleted\n (i.e. the connector no longer has any hierarchy nodes). Callers that want a\n no-op when there are no live nodes must guard before calling.\n\n Returns the number of deleted rows.\n \"\"\"\n stmt = delete(HierarchyNodeByConnectorCredentialPair).where(\n HierarchyNodeByConnectorCredentialPair.connector_id == connector_id,\n HierarchyNodeByConnectorCredentialPair.credential_id == credential_id,\n )\n if live_hierarchy_node_ids:\n stmt = stmt.where(\n HierarchyNodeByConnectorCredentialPair.hierarchy_node_id.notin_(\n live_hierarchy_node_ids\n )\n )\n\n result: CursorResult = db_session.execute(stmt) # type: ignore[assignment]\n deleted = result.rowcount\n\n if commit:\n db_session.commit()\n elif deleted:\n db_session.flush()\n\n return deleted\n\n\ndef delete_orphaned_hierarchy_nodes(\n db_session: Session,\n source: DocumentSource,\n commit: bool = True,\n) -> list[str]:\n \"\"\"Delete hierarchy nodes for a source that have zero cc_pair associations.\n\n SOURCE-type nodes are excluded (they are synthetic roots).\n\n Returns the list of raw_node_ids that were deleted (for cache eviction).\n \"\"\"\n # Find orphaned nodes: no rows in the join table\n orphan_stmt = (\n select(HierarchyNode.id, HierarchyNode.raw_node_id)\n .outerjoin(\n HierarchyNodeByConnectorCredentialPair,\n HierarchyNode.id\n == HierarchyNodeByConnectorCredentialPair.hierarchy_node_id,\n )\n .where(\n HierarchyNode.source == source,\n HierarchyNode.node_type != HierarchyNodeType.SOURCE,\n HierarchyNodeByConnectorCredentialPair.hierarchy_node_id.is_(None),\n )\n )\n orphans = db_session.execute(orphan_stmt).all()\n if not orphans:\n return []\n\n orphan_ids = [row[0] for row in orphans]\n deleted_raw_ids = [row[1] for row in orphans]\n\n db_session.execute(delete(HierarchyNode).where(HierarchyNode.id.in_(orphan_ids)))\n\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n return deleted_raw_ids\n\n\ndef reparent_orphaned_hierarchy_nodes(\n db_session: Session,\n source: DocumentSource,\n commit: bool = True,\n) -> list[HierarchyNode]:\n \"\"\"Re-parent hierarchy nodes whose parent_id is NULL to the SOURCE node.\n\n After pruning deletes stale nodes, their former children get parent_id=NULL\n via the SET NULL cascade. This function points them back to the SOURCE root.\n\n Returns the reparented HierarchyNode objects (with updated parent_id)\n so callers can refresh downstream caches.\n \"\"\"\n source_node = get_source_hierarchy_node(db_session, source)\n if not source_node:\n return []\n\n stmt = select(HierarchyNode).where(\n HierarchyNode.source == source,\n HierarchyNode.parent_id.is_(None),\n HierarchyNode.node_type != HierarchyNodeType.SOURCE,\n )\n orphans = list(db_session.execute(stmt).scalars().all())\n if not orphans:\n return []\n\n for node in orphans:\n node.parent_id = source_node.id\n\n if commit:\n db_session.commit()\n else:\n db_session.flush()\n\n return orphans\n" }, { "path": "backend/onyx/db/hook.py", "content": "import datetime\nfrom uuid import UUID\n\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy.engine import CursorResult\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.constants import UNSET\nfrom onyx.db.constants import UnsetType\nfrom onyx.db.enums import HookFailStrategy\nfrom onyx.db.enums import HookPoint\nfrom onyx.db.models import Hook\nfrom onyx.db.models import HookExecutionLog\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\n\n\n# ── Hook CRUD ────────────────────────────────────────────────────────────\n\n\ndef get_hook_by_id(\n *,\n db_session: Session,\n hook_id: int,\n include_deleted: bool = False,\n include_creator: bool = False,\n) -> Hook | None:\n stmt = select(Hook).where(Hook.id == hook_id)\n if not include_deleted:\n stmt = stmt.where(Hook.deleted.is_(False))\n if include_creator:\n stmt = stmt.options(selectinload(Hook.creator))\n return db_session.scalar(stmt)\n\n\ndef get_non_deleted_hook_by_hook_point(\n *,\n db_session: Session,\n hook_point: HookPoint,\n include_creator: bool = False,\n) -> Hook | None:\n stmt = (\n select(Hook).where(Hook.hook_point == hook_point).where(Hook.deleted.is_(False))\n )\n if include_creator:\n stmt = stmt.options(selectinload(Hook.creator))\n return db_session.scalar(stmt)\n\n\ndef get_hooks(\n *,\n db_session: Session,\n include_deleted: bool = False,\n include_creator: bool = False,\n) -> list[Hook]:\n stmt = select(Hook)\n if not include_deleted:\n stmt = stmt.where(Hook.deleted.is_(False))\n if include_creator:\n stmt = stmt.options(selectinload(Hook.creator))\n stmt = stmt.order_by(Hook.hook_point, Hook.created_at.desc())\n return list(db_session.scalars(stmt).all())\n\n\ndef create_hook__no_commit(\n *,\n db_session: Session,\n name: str,\n hook_point: HookPoint,\n endpoint_url: str | None = None,\n api_key: str | None = None,\n fail_strategy: HookFailStrategy,\n timeout_seconds: float,\n is_active: bool = False,\n is_reachable: bool | None = None,\n creator_id: UUID | None = None,\n) -> Hook:\n \"\"\"Create a new hook for the given hook point.\n\n At most one non-deleted hook per hook point is allowed. Raises\n OnyxError(CONFLICT) if a hook already exists, including under concurrent\n duplicate creates where the partial unique index fires an IntegrityError.\n \"\"\"\n existing = get_non_deleted_hook_by_hook_point(\n db_session=db_session, hook_point=hook_point\n )\n if existing:\n raise OnyxError(\n OnyxErrorCode.CONFLICT,\n f\"A hook for '{hook_point.value}' already exists (id={existing.id}).\",\n )\n\n hook = Hook(\n name=name,\n hook_point=hook_point,\n endpoint_url=endpoint_url,\n api_key=api_key,\n fail_strategy=fail_strategy,\n timeout_seconds=timeout_seconds,\n is_active=is_active,\n is_reachable=is_reachable,\n creator_id=creator_id,\n )\n # Use a savepoint so that a failed insert only rolls back this operation,\n # not the entire outer transaction.\n savepoint = db_session.begin_nested()\n try:\n db_session.add(hook)\n savepoint.commit()\n except IntegrityError as exc:\n savepoint.rollback()\n if \"ix_hook_one_non_deleted_per_point\" in str(exc.orig):\n raise OnyxError(\n OnyxErrorCode.CONFLICT,\n f\"A hook for '{hook_point.value}' already exists.\",\n )\n raise # re-raise unrelated integrity errors (FK violations, etc.)\n return hook\n\n\ndef update_hook__no_commit(\n *,\n db_session: Session,\n hook_id: int,\n name: str | None = None,\n endpoint_url: str | None | UnsetType = UNSET,\n api_key: str | None | UnsetType = UNSET,\n fail_strategy: HookFailStrategy | None = None,\n timeout_seconds: float | None = None,\n is_active: bool | None = None,\n is_reachable: bool | None = None,\n include_creator: bool = False,\n) -> Hook:\n \"\"\"Update hook fields.\n\n Sentinel conventions:\n - endpoint_url, api_key: pass UNSET to leave unchanged; pass None to clear.\n - name, fail_strategy, timeout_seconds, is_active, is_reachable: pass None to leave unchanged.\n \"\"\"\n hook = get_hook_by_id(\n db_session=db_session, hook_id=hook_id, include_creator=include_creator\n )\n if hook is None:\n raise OnyxError(OnyxErrorCode.NOT_FOUND, f\"Hook with id {hook_id} not found.\")\n\n if name is not None:\n hook.name = name\n if not isinstance(endpoint_url, UnsetType):\n hook.endpoint_url = endpoint_url\n if not isinstance(api_key, UnsetType):\n hook.api_key = api_key # type: ignore[assignment] # EncryptedString coerces str → SensitiveValue at the ORM level\n if fail_strategy is not None:\n hook.fail_strategy = fail_strategy\n if timeout_seconds is not None:\n hook.timeout_seconds = timeout_seconds\n if is_active is not None:\n hook.is_active = is_active\n if is_reachable is not None:\n hook.is_reachable = is_reachable\n\n db_session.flush()\n return hook\n\n\ndef delete_hook__no_commit(\n *,\n db_session: Session,\n hook_id: int,\n) -> None:\n hook = get_hook_by_id(db_session=db_session, hook_id=hook_id)\n if hook is None:\n raise OnyxError(OnyxErrorCode.NOT_FOUND, f\"Hook with id {hook_id} not found.\")\n\n hook.deleted = True\n hook.is_active = False\n db_session.flush()\n\n\n# ── HookExecutionLog CRUD ────────────────────────────────────────────────\n\n\ndef create_hook_execution_log__no_commit(\n *,\n db_session: Session,\n hook_id: int,\n is_success: bool,\n error_message: str | None = None,\n status_code: int | None = None,\n duration_ms: int | None = None,\n) -> HookExecutionLog:\n log = HookExecutionLog(\n hook_id=hook_id,\n is_success=is_success,\n error_message=error_message,\n status_code=status_code,\n duration_ms=duration_ms,\n )\n db_session.add(log)\n db_session.flush()\n return log\n\n\ndef get_hook_execution_logs(\n *,\n db_session: Session,\n hook_id: int,\n limit: int,\n) -> list[HookExecutionLog]:\n stmt = (\n select(HookExecutionLog)\n .where(HookExecutionLog.hook_id == hook_id)\n .order_by(HookExecutionLog.created_at.desc())\n .limit(limit)\n )\n return list(db_session.scalars(stmt).all())\n\n\ndef cleanup_old_execution_logs__no_commit(\n *,\n db_session: Session,\n max_age_days: int,\n) -> int:\n \"\"\"Delete execution logs older than max_age_days. Returns the number of rows deleted.\"\"\"\n cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(\n days=max_age_days\n )\n result: CursorResult = db_session.execute( # type: ignore[assignment]\n delete(HookExecutionLog)\n .where(HookExecutionLog.created_at < cutoff)\n .execution_options(synchronize_session=False)\n )\n return result.rowcount\n" }, { "path": "backend/onyx/db/image_generation.py", "content": "from sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import ImageGenerationConfig\nfrom onyx.db.models import LLMProvider\nfrom onyx.db.models import ModelConfiguration\nfrom onyx.llm.utils import get_max_input_tokens\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Default image generation config constants\nDEFAULT_IMAGE_PROVIDER_ID = \"openai_gpt_image_1\"\nDEFAULT_IMAGE_MODEL_NAME = \"gpt-image-1\"\nDEFAULT_IMAGE_PROVIDER = \"openai\"\n\n\ndef create_image_generation_config__no_commit(\n db_session: Session,\n image_provider_id: str,\n model_configuration_id: int,\n is_default: bool = False,\n) -> ImageGenerationConfig:\n \"\"\"Create a new image generation config.\"\"\"\n # If setting as default, clear ALL existing defaults in a single atomic update\n # This is more atomic than select-then-update pattern\n if is_default:\n db_session.execute(\n update(ImageGenerationConfig)\n .where(ImageGenerationConfig.is_default.is_(True))\n .values(is_default=False)\n )\n\n new_config = ImageGenerationConfig(\n image_provider_id=image_provider_id,\n model_configuration_id=model_configuration_id,\n is_default=is_default,\n )\n db_session.add(new_config)\n db_session.flush()\n return new_config\n\n\ndef get_all_image_generation_configs(\n db_session: Session,\n) -> list[ImageGenerationConfig]:\n \"\"\"Get all image generation configs.\n\n Returns:\n List of all ImageGenerationConfig objects\n \"\"\"\n stmt = select(ImageGenerationConfig)\n return list(db_session.scalars(stmt).all())\n\n\ndef get_image_generation_config(\n db_session: Session,\n image_provider_id: str,\n) -> ImageGenerationConfig | None:\n \"\"\"Get a single image generation config by image_provider_id with relationships loaded.\n\n Args:\n db_session: Database session\n image_provider_id: The image provider ID (primary key)\n\n Returns:\n The ImageGenerationConfig or None if not found\n \"\"\"\n stmt = (\n select(ImageGenerationConfig)\n .where(ImageGenerationConfig.image_provider_id == image_provider_id)\n .options(\n selectinload(ImageGenerationConfig.model_configuration).selectinload(\n ModelConfiguration.llm_provider\n )\n )\n )\n return db_session.scalar(stmt)\n\n\ndef get_default_image_generation_config(\n db_session: Session,\n) -> ImageGenerationConfig | None:\n \"\"\"Get the default image generation config.\n\n Returns:\n The default ImageGenerationConfig or None if not set\n \"\"\"\n stmt = (\n select(ImageGenerationConfig)\n .where(ImageGenerationConfig.is_default.is_(True))\n .options(\n selectinload(ImageGenerationConfig.model_configuration).selectinload(\n ModelConfiguration.llm_provider\n )\n )\n )\n return db_session.scalar(stmt)\n\n\ndef set_default_image_generation_config(\n db_session: Session,\n image_provider_id: str,\n) -> None:\n \"\"\"Set a config as the default (clears previous default).\n\n Args:\n db_session: Database session\n image_provider_id: The image provider ID to set as default\n\n Raises:\n ValueError: If config not found\n \"\"\"\n # Get the config to set as default\n new_default = db_session.get(ImageGenerationConfig, image_provider_id)\n if not new_default:\n raise ValueError(\n f\"ImageGenerationConfig with image_provider_id {image_provider_id} not found\"\n )\n\n # Clear ALL existing defaults in a single atomic update\n # This is more atomic than select-then-update pattern\n db_session.execute(\n update(ImageGenerationConfig)\n .where(\n ImageGenerationConfig.is_default.is_(True),\n ImageGenerationConfig.image_provider_id != image_provider_id,\n )\n .values(is_default=False)\n )\n\n # Set new default\n new_default.is_default = True\n db_session.commit()\n\n\ndef unset_default_image_generation_config(\n db_session: Session,\n image_provider_id: str,\n) -> None:\n \"\"\"Unset a config as the default.\"\"\"\n config = db_session.get(ImageGenerationConfig, image_provider_id)\n if not config:\n raise ValueError(\n f\"ImageGenerationConfig with image_provider_id {image_provider_id} not found\"\n )\n config.is_default = False\n db_session.commit()\n\n\ndef delete_image_generation_config__no_commit(\n db_session: Session,\n image_provider_id: str,\n) -> None:\n \"\"\"Delete an image generation config by image_provider_id.\"\"\"\n config = db_session.get(ImageGenerationConfig, image_provider_id)\n if not config:\n raise ValueError(\n f\"ImageGenerationConfig with image_provider_id {image_provider_id} not found\"\n )\n\n db_session.delete(config)\n db_session.flush()\n\n\ndef create_default_image_gen_config_from_api_key(\n db_session: Session,\n api_key: str,\n provider: str = DEFAULT_IMAGE_PROVIDER,\n image_provider_id: str = DEFAULT_IMAGE_PROVIDER_ID,\n model_name: str = DEFAULT_IMAGE_MODEL_NAME,\n) -> ImageGenerationConfig | None:\n \"\"\"Create default image gen config using an API key directly.\n\n This function is used during tenant provisioning to automatically create\n a default image generation config when an OpenAI provider is configured.\n\n Args:\n db_session: Database session\n api_key: API key for the LLM provider\n provider: Provider name (default: openai)\n image_provider_id: Static unique key for the config (default: openai_gpt_image_1)\n model_name: Model name for image generation (default: gpt-image-1)\n\n Returns:\n The created ImageGenerationConfig, or None if:\n - image_generation_config table already has records\n \"\"\"\n # Check if any image generation configs already exist (optimization to avoid work)\n existing_configs = get_all_image_generation_configs(db_session)\n if existing_configs:\n logger.info(\"Image generation config already exists, skipping default creation\")\n return None\n\n try:\n # Create new LLM provider for image generation\n new_provider = LLMProvider(\n name=f\"Image Gen - {image_provider_id}\",\n provider=provider,\n api_key=api_key,\n api_base=None,\n api_version=None,\n deployment_name=None,\n is_public=True,\n )\n db_session.add(new_provider)\n db_session.flush()\n\n # Create model configuration\n max_input_tokens = get_max_input_tokens(\n model_name=model_name,\n model_provider=provider,\n )\n\n model_config = ModelConfiguration(\n llm_provider_id=new_provider.id,\n name=model_name,\n is_visible=True,\n max_input_tokens=max_input_tokens,\n )\n db_session.add(model_config)\n db_session.flush()\n\n # Create image generation config\n config = create_image_generation_config__no_commit(\n db_session=db_session,\n image_provider_id=image_provider_id,\n model_configuration_id=model_config.id,\n is_default=True,\n )\n\n db_session.commit()\n\n logger.info(f\"Created default image generation config: {image_provider_id}\")\n\n return config\n\n except Exception:\n db_session.rollback()\n logger.exception(\n f\"Failed to create default image generation config {image_provider_id}\"\n )\n return None\n" }, { "path": "backend/onyx/db/index_attempt.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import NamedTuple\nfrom typing import TYPE_CHECKING\nfrom typing import TypeVarTuple\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import delete\nfrom sqlalchemy import desc\nfrom sqlalchemy import func\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexAttemptError\nfrom onyx.db.models import SearchSettings\nfrom onyx.server.documents.models import ConnectorCredentialPairIdentifier\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\n\nif TYPE_CHECKING:\n from onyx.configs.constants import DocumentSource\n\n# from sqlalchemy.sql.selectable import Select\n\n# Comment out unused imports that cause mypy errors\n# from onyx.auth.models import UserRole\n# from onyx.configs.constants import MAX_LAST_VALID_CHECKPOINT_AGE_SECONDS\n# from onyx.db.connector_credential_pair import ConnectorCredentialPairIdentifier\n# from onyx.db.engine import async_query_for_dms\n\nlogger = setup_logger()\n\n\ndef get_last_attempt_for_cc_pair(\n cc_pair_id: int,\n search_settings_id: int,\n db_session: Session,\n) -> IndexAttempt | None:\n return (\n db_session.query(IndexAttempt)\n .filter(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n IndexAttempt.search_settings_id == search_settings_id,\n )\n .order_by(IndexAttempt.time_updated.desc())\n .first()\n )\n\n\ndef get_recent_completed_attempts_for_cc_pair(\n cc_pair_id: int,\n search_settings_id: int,\n limit: int,\n db_session: Session,\n) -> list[IndexAttempt]:\n \"\"\"Most recent to least recent.\"\"\"\n return (\n db_session.query(IndexAttempt)\n .filter(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n IndexAttempt.search_settings_id == search_settings_id,\n IndexAttempt.status.notin_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n ),\n )\n .order_by(IndexAttempt.time_updated.desc())\n .limit(limit)\n .all()\n )\n\n\ndef get_recent_attempts_for_cc_pair(\n cc_pair_id: int,\n search_settings_id: int,\n limit: int,\n db_session: Session,\n) -> list[IndexAttempt]:\n \"\"\"Most recent to least recent.\"\"\"\n return (\n db_session.query(IndexAttempt)\n .filter(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n IndexAttempt.search_settings_id == search_settings_id,\n )\n .order_by(IndexAttempt.time_updated.desc())\n .limit(limit)\n .all()\n )\n\n\ndef get_index_attempt(\n db_session: Session,\n index_attempt_id: int,\n eager_load_cc_pair: bool = False,\n eager_load_search_settings: bool = False,\n) -> IndexAttempt | None:\n stmt = select(IndexAttempt).where(IndexAttempt.id == index_attempt_id)\n if eager_load_cc_pair:\n stmt = stmt.options(\n joinedload(IndexAttempt.connector_credential_pair).joinedload(\n ConnectorCredentialPair.connector\n )\n )\n stmt = stmt.options(\n joinedload(IndexAttempt.connector_credential_pair).joinedload(\n ConnectorCredentialPair.credential\n )\n )\n if eager_load_search_settings:\n stmt = stmt.options(joinedload(IndexAttempt.search_settings))\n return db_session.scalars(stmt).first()\n\n\ndef count_error_rows_for_index_attempt(\n index_attempt_id: int,\n db_session: Session,\n) -> int:\n return (\n db_session.query(IndexAttemptError)\n .filter(IndexAttemptError.index_attempt_id == index_attempt_id)\n .count()\n )\n\n\ndef create_index_attempt(\n connector_credential_pair_id: int,\n search_settings_id: int,\n db_session: Session,\n from_beginning: bool = False,\n celery_task_id: str | None = None,\n) -> int:\n new_attempt = IndexAttempt(\n connector_credential_pair_id=connector_credential_pair_id,\n search_settings_id=search_settings_id,\n from_beginning=from_beginning,\n status=IndexingStatus.NOT_STARTED,\n celery_task_id=celery_task_id,\n )\n db_session.add(new_attempt)\n db_session.commit()\n\n return new_attempt.id\n\n\ndef delete_index_attempt(db_session: Session, index_attempt_id: int) -> None:\n index_attempt = get_index_attempt(db_session, index_attempt_id)\n if index_attempt:\n db_session.delete(index_attempt)\n db_session.commit()\n\n\ndef mock_successful_index_attempt(\n connector_credential_pair_id: int,\n search_settings_id: int,\n docs_indexed: int,\n db_session: Session,\n) -> int:\n \"\"\"Should not be used in any user triggered flows\"\"\"\n db_time = func.now()\n new_attempt = IndexAttempt(\n connector_credential_pair_id=connector_credential_pair_id,\n search_settings_id=search_settings_id,\n from_beginning=True,\n status=IndexingStatus.SUCCESS,\n total_docs_indexed=docs_indexed,\n new_docs_indexed=docs_indexed,\n # Need this to be some convincing random looking value and it can't be 0\n # or the indexing rate would calculate out to infinity\n time_started=db_time - timedelta(seconds=1.92),\n time_updated=db_time,\n )\n db_session.add(new_attempt)\n db_session.commit()\n\n return new_attempt.id\n\n\ndef get_in_progress_index_attempts(\n connector_id: int | None,\n db_session: Session,\n) -> list[IndexAttempt]:\n stmt = select(IndexAttempt)\n if connector_id is not None:\n stmt = stmt.where(\n IndexAttempt.connector_credential_pair.has(connector_id=connector_id)\n )\n stmt = stmt.where(IndexAttempt.status == IndexingStatus.IN_PROGRESS)\n\n incomplete_attempts = db_session.scalars(stmt)\n return list(incomplete_attempts.all())\n\n\ndef get_all_index_attempts_by_status(\n status: IndexingStatus, db_session: Session\n) -> list[IndexAttempt]:\n \"\"\"Returns index attempts with the given status.\n Only recommend calling this with non-terminal states as the full list of\n terminal statuses may be quite large.\n\n Results are ordered by time_created (oldest to newest).\"\"\"\n stmt = select(IndexAttempt)\n stmt = stmt.where(IndexAttempt.status == status)\n stmt = stmt.order_by(IndexAttempt.time_created)\n new_attempts = db_session.scalars(stmt)\n return list(new_attempts.all())\n\n\ndef transition_attempt_to_in_progress(\n index_attempt_id: int,\n db_session: Session,\n) -> IndexAttempt:\n \"\"\"Locks the row when we try to update\"\"\"\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if attempt is None:\n raise RuntimeError(\n f\"Unable to find IndexAttempt for ID '{index_attempt_id}'\"\n )\n\n if attempt.status != IndexingStatus.NOT_STARTED:\n raise RuntimeError(\n f\"Indexing attempt with ID '{index_attempt_id}' is not in NOT_STARTED status. \"\n f\"Current status is '{attempt.status}'.\"\n )\n\n attempt.status = IndexingStatus.IN_PROGRESS\n attempt.time_started = attempt.time_started or func.now() # type: ignore\n db_session.commit()\n return attempt\n except Exception:\n db_session.rollback()\n logger.exception(\"transition_attempt_to_in_progress exceptioned.\")\n raise\n\n\ndef mark_attempt_in_progress(\n index_attempt: IndexAttempt,\n db_session: Session,\n) -> None:\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt.id)\n .with_for_update()\n ).scalar_one()\n\n attempt.status = IndexingStatus.IN_PROGRESS\n attempt.time_started = index_attempt.time_started or func.now() # type: ignore\n db_session.commit()\n\n # Add telemetry for index attempt status change\n optional_telemetry(\n record_type=RecordType.INDEX_ATTEMPT_STATUS,\n data={\n \"index_attempt_id\": index_attempt.id,\n \"status\": IndexingStatus.IN_PROGRESS.value,\n \"cc_pair_id\": index_attempt.connector_credential_pair_id,\n },\n )\n except Exception:\n db_session.rollback()\n raise\n\n\ndef mark_attempt_succeeded(\n index_attempt_id: int,\n db_session: Session,\n) -> IndexAttempt:\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update()\n ).scalar_one()\n\n attempt.status = IndexingStatus.SUCCESS\n attempt.celery_task_id = None\n db_session.commit()\n\n # Add telemetry for index attempt status change\n optional_telemetry(\n record_type=RecordType.INDEX_ATTEMPT_STATUS,\n data={\n \"index_attempt_id\": index_attempt_id,\n \"status\": IndexingStatus.SUCCESS.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n return attempt\n except Exception:\n db_session.rollback()\n raise\n\n\ndef mark_attempt_partially_succeeded(\n index_attempt_id: int,\n db_session: Session,\n) -> IndexAttempt:\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update()\n ).scalar_one()\n\n attempt.status = IndexingStatus.COMPLETED_WITH_ERRORS\n attempt.celery_task_id = None\n db_session.commit()\n\n # Add telemetry for index attempt status change\n optional_telemetry(\n record_type=RecordType.INDEX_ATTEMPT_STATUS,\n data={\n \"index_attempt_id\": index_attempt_id,\n \"status\": IndexingStatus.COMPLETED_WITH_ERRORS.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n return attempt\n except Exception:\n db_session.rollback()\n raise\n\n\ndef mark_attempt_canceled(\n index_attempt_id: int,\n db_session: Session,\n reason: str = \"Unknown\",\n) -> None:\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if not attempt.time_started:\n attempt.time_started = datetime.now(timezone.utc)\n attempt.status = IndexingStatus.CANCELED\n attempt.error_msg = reason\n db_session.commit()\n\n # Add telemetry for index attempt status change\n optional_telemetry(\n record_type=RecordType.INDEX_ATTEMPT_STATUS,\n data={\n \"index_attempt_id\": index_attempt_id,\n \"status\": IndexingStatus.CANCELED.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n except Exception:\n db_session.rollback()\n raise\n\n\ndef mark_attempt_failed(\n index_attempt_id: int,\n db_session: Session,\n failure_reason: str = \"Unknown\",\n full_exception_trace: str | None = None,\n) -> None:\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if not attempt.time_started:\n attempt.time_started = datetime.now(timezone.utc)\n attempt.status = IndexingStatus.FAILED\n attempt.error_msg = failure_reason\n attempt.full_exception_trace = full_exception_trace\n attempt.celery_task_id = None\n db_session.commit()\n\n # Add telemetry for index attempt status change\n optional_telemetry(\n record_type=RecordType.INDEX_ATTEMPT_STATUS,\n data={\n \"index_attempt_id\": index_attempt_id,\n \"status\": IndexingStatus.FAILED.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n except Exception:\n db_session.rollback()\n raise\n\n\ndef update_docs_indexed(\n db_session: Session,\n index_attempt_id: int,\n total_docs_indexed: int,\n new_docs_indexed: int,\n docs_removed_from_index: int,\n) -> None:\n \"\"\"Updates the docs_indexed and new_docs_indexed fields of an index attempt.\n Adds the given values to the current values in the db\"\"\"\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update() # Locks the row when we try to update\n ).scalar_one()\n\n attempt.total_docs_indexed = (\n attempt.total_docs_indexed or 0\n ) + total_docs_indexed\n attempt.new_docs_indexed = (attempt.new_docs_indexed or 0) + new_docs_indexed\n attempt.docs_removed_from_index = (\n attempt.docs_removed_from_index or 0\n ) + docs_removed_from_index\n db_session.commit()\n except Exception:\n db_session.rollback()\n logger.exception(\"update_docs_indexed exceptioned.\")\n raise\n\n\ndef get_last_attempt(\n connector_id: int,\n credential_id: int,\n search_settings_id: int | None,\n db_session: Session,\n) -> IndexAttempt | None:\n stmt = (\n select(IndexAttempt)\n .join(ConnectorCredentialPair)\n .where(\n ConnectorCredentialPair.connector_id == connector_id,\n ConnectorCredentialPair.credential_id == credential_id,\n IndexAttempt.search_settings_id == search_settings_id,\n )\n )\n\n # Note, the below is using time_created instead of time_updated\n stmt = stmt.order_by(desc(IndexAttempt.time_created))\n\n return db_session.execute(stmt).scalars().first()\n\n\ndef get_latest_index_attempts_by_status(\n secondary_index: bool,\n db_session: Session,\n status: IndexingStatus,\n) -> Sequence[IndexAttempt]:\n \"\"\"\n Retrieves the most recent index attempt with the specified status for each connector_credential_pair.\n Filters attempts based on the secondary_index flag to get either future or present index attempts.\n Returns a sequence of IndexAttempt objects, one for each unique connector_credential_pair.\n \"\"\"\n latest_failed_attempts = (\n select(\n IndexAttempt.connector_credential_pair_id,\n func.max(IndexAttempt.id).label(\"max_failed_id\"),\n )\n .join(SearchSettings, IndexAttempt.search_settings_id == SearchSettings.id)\n .where(\n SearchSettings.status\n == (\n IndexModelStatus.FUTURE if secondary_index else IndexModelStatus.PRESENT\n ),\n IndexAttempt.status == status,\n )\n .group_by(IndexAttempt.connector_credential_pair_id)\n .subquery()\n )\n\n stmt = select(IndexAttempt).join(\n latest_failed_attempts,\n (\n IndexAttempt.connector_credential_pair_id\n == latest_failed_attempts.c.connector_credential_pair_id\n )\n & (IndexAttempt.id == latest_failed_attempts.c.max_failed_id),\n )\n\n return db_session.execute(stmt).scalars().all()\n\n\nT = TypeVarTuple(\"T\")\n\n\ndef _add_only_finished_clause(stmt: Select[tuple[*T]]) -> Select[tuple[*T]]:\n return stmt.where(\n IndexAttempt.status.not_in(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n ),\n )\n\n\ndef get_latest_index_attempts(\n secondary_index: bool,\n db_session: Session,\n eager_load_cc_pair: bool = False,\n only_finished: bool = False,\n) -> Sequence[IndexAttempt]:\n ids_stmt = select(\n IndexAttempt.connector_credential_pair_id,\n func.max(IndexAttempt.id).label(\"max_id\"),\n ).join(SearchSettings, IndexAttempt.search_settings_id == SearchSettings.id)\n\n status = IndexModelStatus.FUTURE if secondary_index else IndexModelStatus.PRESENT\n ids_stmt = ids_stmt.where(SearchSettings.status == status)\n\n if only_finished:\n ids_stmt = _add_only_finished_clause(ids_stmt)\n\n ids_stmt = ids_stmt.group_by(IndexAttempt.connector_credential_pair_id)\n ids_subquery = ids_stmt.subquery()\n\n stmt = (\n select(IndexAttempt)\n .join(\n ids_subquery,\n IndexAttempt.connector_credential_pair_id\n == ids_subquery.c.connector_credential_pair_id,\n )\n .where(IndexAttempt.id == ids_subquery.c.max_id)\n )\n\n if only_finished:\n stmt = _add_only_finished_clause(stmt)\n\n if eager_load_cc_pair:\n stmt = stmt.options(\n joinedload(IndexAttempt.connector_credential_pair),\n joinedload(IndexAttempt.error_rows),\n )\n\n return db_session.execute(stmt).scalars().unique().all()\n\n\n# For use with our thread-level parallelism utils. Note that any relationships\n# you wish to use MUST be eagerly loaded, as the session will not be available\n# after this function to allow lazy loading.\ndef get_latest_index_attempts_parallel(\n secondary_index: bool,\n eager_load_cc_pair: bool = False,\n only_finished: bool = False,\n) -> Sequence[IndexAttempt]:\n with get_session_with_current_tenant() as db_session:\n return get_latest_index_attempts(\n secondary_index,\n db_session,\n eager_load_cc_pair,\n only_finished,\n )\n\n\ndef get_latest_index_attempt_for_cc_pair_id(\n db_session: Session,\n connector_credential_pair_id: int,\n secondary_index: bool,\n only_finished: bool = True,\n) -> IndexAttempt | None:\n stmt = select(IndexAttempt)\n stmt = stmt.where(\n IndexAttempt.connector_credential_pair_id == connector_credential_pair_id,\n )\n if only_finished:\n stmt = _add_only_finished_clause(stmt)\n\n status = IndexModelStatus.FUTURE if secondary_index else IndexModelStatus.PRESENT\n stmt = stmt.join(SearchSettings).where(SearchSettings.status == status)\n stmt = stmt.order_by(desc(IndexAttempt.time_created))\n stmt = stmt.limit(1)\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef get_latest_successful_index_attempt_for_cc_pair_id(\n db_session: Session,\n connector_credential_pair_id: int,\n secondary_index: bool = False,\n) -> IndexAttempt | None:\n \"\"\"Returns the most recent successful index attempt for the given cc pair,\n filtered to the current (or future) search settings.\n Uses MAX(id) semantics to match get_latest_index_attempts_by_status.\"\"\"\n status = IndexModelStatus.FUTURE if secondary_index else IndexModelStatus.PRESENT\n stmt = (\n select(IndexAttempt)\n .where(\n IndexAttempt.connector_credential_pair_id == connector_credential_pair_id,\n IndexAttempt.status.in_(\n [IndexingStatus.SUCCESS, IndexingStatus.COMPLETED_WITH_ERRORS]\n ),\n )\n .join(SearchSettings)\n .where(SearchSettings.status == status)\n .order_by(desc(IndexAttempt.id))\n .limit(1)\n )\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef get_latest_successful_index_attempts_parallel(\n secondary_index: bool = False,\n) -> Sequence[IndexAttempt]:\n \"\"\"Batch version: returns the latest successful index attempt per cc pair.\n Covers both SUCCESS and COMPLETED_WITH_ERRORS (matching is_successful()).\"\"\"\n model_status = (\n IndexModelStatus.FUTURE if secondary_index else IndexModelStatus.PRESENT\n )\n with get_session_with_current_tenant() as db_session:\n latest_ids = (\n select(\n IndexAttempt.connector_credential_pair_id,\n func.max(IndexAttempt.id).label(\"max_id\"),\n )\n .join(SearchSettings, IndexAttempt.search_settings_id == SearchSettings.id)\n .where(\n SearchSettings.status == model_status,\n IndexAttempt.status.in_(\n [IndexingStatus.SUCCESS, IndexingStatus.COMPLETED_WITH_ERRORS]\n ),\n )\n .group_by(IndexAttempt.connector_credential_pair_id)\n .subquery()\n )\n\n stmt = select(IndexAttempt).join(\n latest_ids,\n (\n IndexAttempt.connector_credential_pair_id\n == latest_ids.c.connector_credential_pair_id\n )\n & (IndexAttempt.id == latest_ids.c.max_id),\n )\n return db_session.execute(stmt).scalars().all()\n\n\ndef count_index_attempts_for_cc_pair(\n db_session: Session,\n cc_pair_id: int,\n only_current: bool = True,\n disinclude_finished: bool = False,\n) -> int:\n stmt = select(IndexAttempt).where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id\n )\n if disinclude_finished:\n stmt = stmt.where(\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n )\n )\n if only_current:\n stmt = stmt.join(SearchSettings).where(\n SearchSettings.status == IndexModelStatus.PRESENT\n )\n # Count total items for pagination\n count_stmt = stmt.with_only_columns(func.count()).order_by(None)\n total_count = db_session.execute(count_stmt).scalar_one()\n return total_count\n\n\ndef get_paginated_index_attempts_for_cc_pair_id(\n db_session: Session,\n cc_pair_id: int,\n page: int,\n page_size: int,\n only_current: bool = True,\n disinclude_finished: bool = False,\n) -> list[IndexAttempt]:\n stmt = select(IndexAttempt).where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id\n )\n if disinclude_finished:\n stmt = stmt.where(\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n )\n )\n if only_current:\n stmt = stmt.join(SearchSettings).where(\n SearchSettings.status == IndexModelStatus.PRESENT\n )\n\n stmt = stmt.order_by(IndexAttempt.time_started.desc())\n\n # Apply pagination\n stmt = stmt.offset(page * page_size).limit(page_size)\n\n return list(db_session.execute(stmt).scalars().unique().all())\n\n\ndef get_index_attempts_for_cc_pair(\n db_session: Session,\n cc_pair_identifier: ConnectorCredentialPairIdentifier,\n only_current: bool = True,\n disinclude_finished: bool = False,\n) -> Sequence[IndexAttempt]:\n stmt = (\n select(IndexAttempt)\n .join(ConnectorCredentialPair)\n .where(\n and_(\n ConnectorCredentialPair.connector_id == cc_pair_identifier.connector_id,\n ConnectorCredentialPair.credential_id\n == cc_pair_identifier.credential_id,\n )\n )\n )\n if disinclude_finished:\n stmt = stmt.where(\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n )\n )\n if only_current:\n stmt = stmt.join(SearchSettings).where(\n SearchSettings.status == IndexModelStatus.PRESENT\n )\n\n stmt = stmt.order_by(IndexAttempt.time_created.desc())\n return db_session.execute(stmt).scalars().all()\n\n\ndef delete_index_attempts(\n cc_pair_id: int,\n db_session: Session,\n) -> None:\n # First, delete related entries in IndexAttemptErrors\n stmt_errors = delete(IndexAttemptError).where(\n IndexAttemptError.index_attempt_id.in_(\n select(IndexAttempt.id).where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id\n )\n )\n )\n db_session.execute(stmt_errors)\n\n stmt = delete(IndexAttempt).where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n )\n\n db_session.execute(stmt)\n\n\ndef expire_index_attempts(\n search_settings_id: int,\n db_session: Session,\n) -> None:\n not_started_query = (\n update(IndexAttempt)\n .where(IndexAttempt.search_settings_id == search_settings_id)\n .where(IndexAttempt.status == IndexingStatus.NOT_STARTED)\n .values(\n status=IndexingStatus.CANCELED,\n error_msg=\"Canceled, likely due to model swap\",\n )\n )\n db_session.execute(not_started_query)\n\n update_query = (\n update(IndexAttempt)\n .where(IndexAttempt.search_settings_id == search_settings_id)\n .where(IndexAttempt.status != IndexingStatus.SUCCESS)\n .values(\n status=IndexingStatus.FAILED,\n error_msg=\"Canceled due to embedding model swap\",\n )\n )\n db_session.execute(update_query)\n\n db_session.commit()\n\n\ndef cancel_indexing_attempts_for_ccpair(\n cc_pair_id: int,\n db_session: Session,\n include_secondary_index: bool = False,\n) -> None:\n stmt = (\n update(IndexAttempt)\n .where(IndexAttempt.connector_credential_pair_id == cc_pair_id)\n .where(IndexAttempt.status == IndexingStatus.NOT_STARTED)\n .values(\n status=IndexingStatus.CANCELED,\n error_msg=\"Canceled by user\",\n time_started=datetime.now(timezone.utc),\n )\n )\n\n if not include_secondary_index:\n subquery = select(SearchSettings.id).where(\n SearchSettings.status != IndexModelStatus.FUTURE\n )\n stmt = stmt.where(IndexAttempt.search_settings_id.in_(subquery))\n\n db_session.execute(stmt)\n\n\ndef cancel_indexing_attempts_past_model(\n db_session: Session,\n) -> None:\n \"\"\"Stops all indexing attempts that are in progress or not started for\n any embedding model that not present/future\"\"\"\n\n db_session.execute(\n update(IndexAttempt)\n .where(\n IndexAttempt.status.in_(\n [IndexingStatus.IN_PROGRESS, IndexingStatus.NOT_STARTED]\n ),\n IndexAttempt.search_settings_id == SearchSettings.id,\n SearchSettings.status == IndexModelStatus.PAST,\n )\n .values(status=IndexingStatus.FAILED)\n )\n\n\ndef cancel_indexing_attempts_for_search_settings(\n search_settings_id: int,\n db_session: Session,\n) -> None:\n \"\"\"Stops all indexing attempts that are in progress or not started for\n the specified search settings.\"\"\"\n\n db_session.execute(\n update(IndexAttempt)\n .where(\n IndexAttempt.status.in_(\n [IndexingStatus.IN_PROGRESS, IndexingStatus.NOT_STARTED]\n ),\n IndexAttempt.search_settings_id == search_settings_id,\n )\n .values(status=IndexingStatus.FAILED)\n )\n\n\ndef count_unique_cc_pairs_with_successful_index_attempts(\n search_settings_id: int | None,\n db_session: Session,\n) -> int:\n \"\"\"Collect all of the Index Attempts that are successful and for the specified embedding model\n Then do distinct by connector_id and credential_id which is equivalent to the cc-pair. Finally,\n do a count to get the total number of unique cc-pairs with successful attempts\"\"\"\n unique_pairs_count = (\n db_session.query(IndexAttempt.connector_credential_pair_id)\n .join(ConnectorCredentialPair)\n .filter(\n IndexAttempt.search_settings_id == search_settings_id,\n IndexAttempt.status == IndexingStatus.SUCCESS,\n )\n .distinct()\n .count()\n )\n\n return unique_pairs_count\n\n\ndef count_unique_active_cc_pairs_with_successful_index_attempts(\n search_settings_id: int | None,\n db_session: Session,\n) -> int:\n \"\"\"Collect all of the Index Attempts that are successful and for the specified embedding model,\n but only for non-paused connector-credential pairs. Then do distinct by connector_id and credential_id\n which is equivalent to the cc-pair. Finally, do a count to get the total number of unique non-paused\n cc-pairs with successful attempts.\"\"\"\n unique_pairs_count = (\n db_session.query(IndexAttempt.connector_credential_pair_id)\n .join(ConnectorCredentialPair)\n .filter(\n IndexAttempt.search_settings_id == search_settings_id,\n IndexAttempt.status == IndexingStatus.SUCCESS,\n ConnectorCredentialPair.status != ConnectorCredentialPairStatus.PAUSED,\n )\n .distinct()\n .count()\n )\n\n return unique_pairs_count\n\n\ndef create_index_attempt_error(\n index_attempt_id: int | None,\n connector_credential_pair_id: int,\n failure: ConnectorFailure,\n db_session: Session,\n) -> int:\n new_error = IndexAttemptError(\n index_attempt_id=index_attempt_id,\n connector_credential_pair_id=connector_credential_pair_id,\n document_id=(\n failure.failed_document.document_id if failure.failed_document else None\n ),\n document_link=(\n failure.failed_document.document_link if failure.failed_document else None\n ),\n entity_id=(failure.failed_entity.entity_id if failure.failed_entity else None),\n failed_time_range_start=(\n failure.failed_entity.missed_time_range[0]\n if failure.failed_entity and failure.failed_entity.missed_time_range\n else None\n ),\n failed_time_range_end=(\n failure.failed_entity.missed_time_range[1]\n if failure.failed_entity and failure.failed_entity.missed_time_range\n else None\n ),\n failure_message=failure.failure_message,\n is_resolved=False,\n )\n db_session.add(new_error)\n db_session.commit()\n\n return new_error.id\n\n\ndef get_index_attempt_errors(\n index_attempt_id: int,\n db_session: Session,\n) -> list[IndexAttemptError]:\n stmt = select(IndexAttemptError).where(\n IndexAttemptError.index_attempt_id == index_attempt_id\n )\n\n errors = db_session.scalars(stmt)\n return list(errors.all())\n\n\ndef count_index_attempt_errors_for_cc_pair(\n cc_pair_id: int,\n unresolved_only: bool,\n db_session: Session,\n) -> int:\n stmt = (\n select(func.count())\n .select_from(IndexAttemptError)\n .where(IndexAttemptError.connector_credential_pair_id == cc_pair_id)\n )\n if unresolved_only:\n stmt = stmt.where(IndexAttemptError.is_resolved.is_(False))\n\n result = db_session.scalar(stmt)\n return 0 if result is None else result\n\n\ndef get_index_attempt_errors_for_cc_pair(\n cc_pair_id: int,\n unresolved_only: bool,\n db_session: Session,\n page: int | None = None,\n page_size: int | None = None,\n) -> list[IndexAttemptError]:\n stmt = select(IndexAttemptError).where(\n IndexAttemptError.connector_credential_pair_id == cc_pair_id\n )\n if unresolved_only:\n stmt = stmt.where(IndexAttemptError.is_resolved.is_(False))\n\n # Order by most recent first\n stmt = stmt.order_by(desc(IndexAttemptError.time_created))\n\n if page is not None and page_size is not None:\n stmt = stmt.offset(page * page_size).limit(page_size)\n\n return list(db_session.scalars(stmt).all())\n\n\n# ── Metrics query helpers ──────────────────────────────────────────────\n\n\nclass ActiveIndexAttemptMetric(NamedTuple):\n \"\"\"Row returned by get_active_index_attempts_for_metrics.\"\"\"\n\n status: IndexingStatus\n source: \"DocumentSource\"\n cc_pair_id: int\n cc_pair_name: str | None\n attempt_count: int\n\n\ndef get_active_index_attempts_for_metrics(\n db_session: Session,\n) -> list[ActiveIndexAttemptMetric]:\n \"\"\"Return non-terminal index attempts grouped by status, source, and connector.\n\n Each row is (status, source, cc_pair_id, cc_pair_name, attempt_count).\n \"\"\"\n from onyx.db.models import Connector\n\n terminal_statuses = [s for s in IndexingStatus if s.is_terminal()]\n rows = (\n db_session.query(\n IndexAttempt.status,\n Connector.source,\n ConnectorCredentialPair.id,\n ConnectorCredentialPair.name,\n func.count(),\n )\n .join(\n ConnectorCredentialPair,\n IndexAttempt.connector_credential_pair_id == ConnectorCredentialPair.id,\n )\n .join(\n Connector,\n ConnectorCredentialPair.connector_id == Connector.id,\n )\n .filter(IndexAttempt.status.notin_(terminal_statuses))\n .group_by(\n IndexAttempt.status,\n Connector.source,\n ConnectorCredentialPair.id,\n ConnectorCredentialPair.name,\n )\n .all()\n )\n return [ActiveIndexAttemptMetric(*row) for row in rows]\n\n\ndef get_failed_attempt_counts_by_cc_pair(\n db_session: Session,\n since: datetime | None = None,\n) -> dict[int, int]:\n \"\"\"Return {cc_pair_id: failed_attempt_count} for all connectors.\n\n When ``since`` is provided, only attempts created after that timestamp\n are counted. Defaults to the last 90 days to avoid unbounded historical\n aggregation.\n \"\"\"\n if since is None:\n since = datetime.now(timezone.utc) - timedelta(days=90)\n\n rows = (\n db_session.query(\n IndexAttempt.connector_credential_pair_id,\n func.count(),\n )\n .filter(IndexAttempt.status == IndexingStatus.FAILED)\n .filter(IndexAttempt.time_created >= since)\n .group_by(IndexAttempt.connector_credential_pair_id)\n .all()\n )\n return {cc_id: count for cc_id, count in rows}\n\n\ndef get_docs_indexed_by_cc_pair(\n db_session: Session,\n since: datetime | None = None,\n) -> dict[int, int]:\n \"\"\"Return {cc_pair_id: total_new_docs_indexed} across successful attempts.\n\n Only counts attempts with status SUCCESS to avoid inflating counts with\n partial results from failed attempts. When ``since`` is provided, only\n attempts created after that timestamp are included.\n \"\"\"\n if since is None:\n since = datetime.now(timezone.utc) - timedelta(days=90)\n\n query = (\n db_session.query(\n IndexAttempt.connector_credential_pair_id,\n func.sum(func.coalesce(IndexAttempt.new_docs_indexed, 0)),\n )\n .filter(IndexAttempt.status == IndexingStatus.SUCCESS)\n .filter(IndexAttempt.time_created >= since)\n .group_by(IndexAttempt.connector_credential_pair_id)\n )\n rows = query.all()\n return {cc_id: int(total or 0) for cc_id, total in rows}\n" }, { "path": "backend/onyx/db/indexing_coordination.py", "content": "\"\"\"Database-based indexing coordination to replace Redis fencing.\"\"\"\n\nfrom pydantic import BaseModel\nfrom sqlalchemy import select\nfrom sqlalchemy.exc import SQLAlchemyError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.index_attempt import count_error_rows_for_index_attempt\nfrom onyx.db.index_attempt import create_index_attempt\nfrom onyx.db.index_attempt import get_index_attempt\nfrom onyx.db.models import IndexAttempt\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nINDEXING_PROGRESS_TIMEOUT_HOURS = 6\n\n\nclass CoordinationStatus(BaseModel):\n \"\"\"Status of an indexing attempt's coordination.\"\"\"\n\n found: bool\n total_batches: int | None\n completed_batches: int\n total_failures: int\n total_docs: int\n total_chunks: int\n status: IndexingStatus | None = None\n cancellation_requested: bool = False\n\n\nclass IndexingCoordination:\n \"\"\"Database-based coordination for indexing tasks, replacing Redis fencing.\"\"\"\n\n @staticmethod\n def try_create_index_attempt(\n db_session: Session,\n cc_pair_id: int,\n search_settings_id: int,\n celery_task_id: str,\n from_beginning: bool = False,\n ) -> int | None:\n \"\"\"\n Try to create a new index attempt for the given CC pair and search settings.\n Returns the index_attempt_id if successful, None if another attempt is already running.\n\n This replaces the Redis fencing mechanism by using database constraints\n and transactions to prevent duplicate attempts.\n \"\"\"\n try:\n # Check for existing active attempts (this is the \"fence\" check)\n existing_attempt = db_session.execute(\n select(IndexAttempt)\n .where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n IndexAttempt.search_settings_id == search_settings_id,\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n ),\n )\n .with_for_update(nowait=True)\n ).first()\n\n if existing_attempt:\n logger.info(\n f\"Indexing already in progress: \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id} \"\n f\"existing_attempt={existing_attempt[0].id}\"\n )\n return None\n\n # Create new index attempt (this is setting the \"fence\")\n attempt_id = create_index_attempt(\n connector_credential_pair_id=cc_pair_id,\n search_settings_id=search_settings_id,\n from_beginning=from_beginning,\n db_session=db_session,\n celery_task_id=celery_task_id,\n )\n\n logger.info(\n f\"Created Index Attempt: \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id} \"\n f\"attempt_id={attempt_id} \"\n f\"celery_task_id={celery_task_id}\"\n )\n\n return attempt_id\n\n except SQLAlchemyError as e:\n logger.info(\n f\"Failed to create index attempt (likely race condition): \"\n f\"cc_pair={cc_pair_id} \"\n f\"search_settings={search_settings_id} \"\n f\"error={str(e)}\"\n )\n db_session.rollback()\n return None\n\n @staticmethod\n def check_cancellation_requested(\n db_session: Session,\n index_attempt_id: int,\n ) -> bool:\n \"\"\"\n Check if cancellation has been requested for this indexing attempt.\n This replaces Redis termination signals.\n \"\"\"\n attempt = get_index_attempt(db_session, index_attempt_id)\n return attempt.cancellation_requested if attempt else False\n\n @staticmethod\n def request_cancellation(\n db_session: Session,\n index_attempt_id: int,\n ) -> None:\n \"\"\"\n Request cancellation of an indexing attempt.\n This replaces Redis termination signals.\n \"\"\"\n attempt = get_index_attempt(db_session, index_attempt_id)\n if attempt:\n attempt.cancellation_requested = True\n db_session.commit()\n\n logger.info(f\"Requested cancellation for attempt {index_attempt_id}\")\n\n @staticmethod\n def set_total_batches(\n db_session: Session,\n index_attempt_id: int,\n total_batches: int,\n ) -> None:\n \"\"\"\n Set the total number of batches for this indexing attempt.\n Called by docfetching when extraction is complete.\n \"\"\"\n attempt = get_index_attempt(db_session, index_attempt_id)\n if attempt:\n attempt.total_batches = total_batches\n db_session.commit()\n\n logger.info(\n f\"Set total batches: attempt={index_attempt_id} total={total_batches}\"\n )\n\n @staticmethod\n def update_batch_completion_and_docs(\n db_session: Session,\n index_attempt_id: int,\n total_docs_indexed: int,\n new_docs_indexed: int,\n total_chunks: int,\n ) -> tuple[int, int | None]:\n \"\"\"\n Update batch completion and document counts atomically.\n Returns (completed_batches, total_batches).\n This extends the existing update_docs_indexed pattern.\n \"\"\"\n try:\n attempt = db_session.execute(\n select(IndexAttempt)\n .where(IndexAttempt.id == index_attempt_id)\n .with_for_update() # Same pattern as existing update_docs_indexed\n ).scalar_one()\n\n # Existing document count updates\n attempt.total_docs_indexed = (\n attempt.total_docs_indexed or 0\n ) + total_docs_indexed\n attempt.new_docs_indexed = (\n attempt.new_docs_indexed or 0\n ) + new_docs_indexed\n\n # New coordination updates\n attempt.completed_batches = (attempt.completed_batches or 0) + 1\n attempt.total_chunks = (attempt.total_chunks or 0) + total_chunks\n\n db_session.commit()\n\n logger.info(\n f\"Updated batch completion: \"\n f\"attempt={index_attempt_id} \"\n f\"completed={attempt.completed_batches} \"\n f\"total={attempt.total_batches} \"\n f\"docs={total_docs_indexed} \"\n )\n\n return attempt.completed_batches, attempt.total_batches\n\n except Exception:\n db_session.rollback()\n logger.exception(\n f\"Failed to update batch completion for attempt {index_attempt_id}\"\n )\n raise\n\n @staticmethod\n def get_coordination_status(\n db_session: Session,\n index_attempt_id: int,\n ) -> CoordinationStatus:\n \"\"\"\n Get the current coordination status for an indexing attempt.\n This replaces reading FileStore state files.\n \"\"\"\n attempt = get_index_attempt(db_session, index_attempt_id)\n if not attempt:\n return CoordinationStatus(\n found=False,\n total_batches=None,\n completed_batches=0,\n total_failures=0,\n total_docs=0,\n total_chunks=0,\n status=None,\n cancellation_requested=False,\n )\n\n return CoordinationStatus(\n found=True,\n total_batches=attempt.total_batches,\n completed_batches=attempt.completed_batches,\n total_failures=count_error_rows_for_index_attempt(\n index_attempt_id, db_session\n ),\n total_docs=attempt.total_docs_indexed or 0,\n total_chunks=attempt.total_chunks,\n status=attempt.status,\n cancellation_requested=attempt.cancellation_requested,\n )\n\n @staticmethod\n def get_orphaned_index_attempt_ids(db_session: Session) -> list[int]:\n \"\"\"\n Gets a list of potentially orphaned index attempts.\n These are attempts in non-terminal state that have task IDs but may have died.\n\n This replaces the old get_unfenced_index_attempt_ids function.\n The actual orphan detection requires checking with Celery, which should be\n done by the caller.\n \"\"\"\n # Find attempts that are active and have task IDs\n # The caller needs to check each one with Celery to confirm orphaned status\n active_attempts = (\n db_session.execute(\n select(IndexAttempt).where(\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n ),\n IndexAttempt.celery_task_id.isnot(None),\n )\n )\n .scalars()\n .all()\n )\n\n return [attempt.id for attempt in active_attempts]\n\n @staticmethod\n def update_progress_tracking(\n db_session: Session,\n index_attempt_id: int,\n current_batches_completed: int,\n timeout_hours: int = INDEXING_PROGRESS_TIMEOUT_HOURS,\n force_update_progress: bool = False,\n ) -> bool:\n \"\"\"\n Update progress tracking for stall detection.\n Returns True if sufficient progress was made, False if stalled.\n \"\"\"\n\n attempt = get_index_attempt(db_session, index_attempt_id)\n if not attempt:\n logger.error(f\"Index attempt {index_attempt_id} not found in database\")\n return False\n\n current_time = get_db_current_time(db_session)\n\n # No progress - check if this is the first time tracking\n # or if the caller wants to simulate guaranteed progress\n if attempt.last_progress_time is None or force_update_progress:\n # First time tracking - initialize\n attempt.last_progress_time = current_time\n attempt.last_batches_completed_count = current_batches_completed\n db_session.commit()\n return True\n\n time_elapsed = (current_time - attempt.last_progress_time).total_seconds()\n # only actually write to db every timeout_hours/2\n # this ensure thats at most timeout_hours will pass with no activity\n if time_elapsed < timeout_hours * 1800:\n return True\n\n # Check if progress has been made\n if current_batches_completed <= attempt.last_batches_completed_count:\n # if between timeout_hours/2 and timeout_hours has passed\n # without an update, we consider the attempt stalled\n return False\n\n # Progress made - update tracking\n attempt.last_progress_time = current_time\n attempt.last_batches_completed_count = current_batches_completed\n db_session.commit()\n return True\n" }, { "path": "backend/onyx/db/input_prompt.py", "content": "from uuid import UUID\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import InputPrompt\nfrom onyx.db.models import InputPrompt__User\nfrom onyx.db.models import User\nfrom onyx.server.features.input_prompt.models import InputPromptSnapshot\nfrom onyx.server.manage.models import UserInfo\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef insert_input_prompt(\n prompt: str,\n content: str,\n is_public: bool,\n user: User | None,\n db_session: Session,\n) -> InputPrompt:\n user_id = user.id if user else None\n\n # Use atomic INSERT ... ON CONFLICT DO NOTHING with RETURNING\n # to avoid race conditions with the uniqueness check\n stmt = pg_insert(InputPrompt).values(\n prompt=prompt,\n content=content,\n active=True,\n is_public=is_public,\n user_id=user_id,\n )\n\n # Use the appropriate constraint based on whether this is a user-owned or public prompt\n if user_id is not None:\n stmt = stmt.on_conflict_do_nothing(constraint=\"uq_inputprompt_prompt_user_id\")\n else:\n # Partial unique indexes cannot be targeted by constraint name;\n # must use index_elements + index_where\n stmt = stmt.on_conflict_do_nothing(\n index_elements=[InputPrompt.prompt],\n index_where=InputPrompt.user_id.is_(None),\n )\n\n stmt = stmt.returning(InputPrompt)\n\n result = db_session.execute(stmt)\n input_prompt = result.scalar_one_or_none()\n\n if input_prompt is None:\n raise HTTPException(\n status_code=409,\n detail=f\"A prompt shortcut with the name '{prompt}' already exists\",\n )\n\n db_session.commit()\n return input_prompt\n\n\ndef update_input_prompt(\n user: User,\n input_prompt_id: int,\n prompt: str,\n content: str,\n active: bool,\n db_session: Session,\n) -> InputPrompt:\n input_prompt = db_session.scalar(\n select(InputPrompt).where(InputPrompt.id == input_prompt_id)\n )\n if input_prompt is None:\n raise ValueError(f\"No input prompt with id {input_prompt_id}\")\n\n if not validate_user_prompt_authorization(user, input_prompt):\n raise HTTPException(status_code=401, detail=\"You don't own this prompt\")\n\n input_prompt.prompt = prompt\n input_prompt.content = content\n input_prompt.active = active\n\n try:\n db_session.commit()\n except IntegrityError:\n db_session.rollback()\n raise HTTPException(\n status_code=409,\n detail=f\"A prompt shortcut with the name '{prompt}' already exists\",\n )\n\n return input_prompt\n\n\ndef validate_user_prompt_authorization(user: User, input_prompt: InputPrompt) -> bool:\n prompt = InputPromptSnapshot.from_model(input_prompt=input_prompt)\n\n # Public prompts cannot be modified via the user API (only admins via admin endpoints)\n if prompt.is_public or prompt.user_id is None:\n return False\n\n # Anonymous users cannot modify user-owned prompts\n if user.is_anonymous:\n return False\n\n # User must own the prompt\n user_details = UserInfo.from_model(user)\n return str(user_details.id) == str(prompt.user_id)\n\n\ndef remove_public_input_prompt(input_prompt_id: int, db_session: Session) -> None:\n input_prompt = db_session.scalar(\n select(InputPrompt).where(InputPrompt.id == input_prompt_id)\n )\n\n if input_prompt is None:\n raise ValueError(f\"No input prompt with id {input_prompt_id}\")\n\n if not input_prompt.is_public:\n raise HTTPException(status_code=400, detail=\"This prompt is not public\")\n\n db_session.delete(input_prompt)\n db_session.commit()\n\n\ndef remove_input_prompt(\n user: User,\n input_prompt_id: int,\n db_session: Session,\n delete_public: bool = False,\n) -> None:\n input_prompt = db_session.scalar(\n select(InputPrompt).where(InputPrompt.id == input_prompt_id)\n )\n if input_prompt is None:\n raise ValueError(f\"No input prompt with id {input_prompt_id}\")\n\n if input_prompt.is_public and not delete_public:\n raise HTTPException(\n status_code=400, detail=\"Cannot delete public prompts with this method\"\n )\n\n if not validate_user_prompt_authorization(user, input_prompt):\n raise HTTPException(status_code=401, detail=\"You do not own this prompt\")\n\n db_session.delete(input_prompt)\n db_session.commit()\n\n\ndef fetch_input_prompt_by_id(\n id: int, user_id: UUID | None, db_session: Session\n) -> InputPrompt:\n query = select(InputPrompt).where(InputPrompt.id == id)\n\n if user_id:\n query = query.where(\n (InputPrompt.user_id == user_id) | (InputPrompt.user_id is None)\n )\n else:\n # If no user_id is provided, only fetch prompts without a user_id (aka public)\n query = query.where(InputPrompt.user_id == None) # noqa\n\n result = db_session.scalar(query)\n\n if result is None:\n raise HTTPException(422, \"No input prompt found\")\n\n return result\n\n\ndef fetch_public_input_prompts(\n db_session: Session,\n) -> list[InputPrompt]:\n query = select(InputPrompt).where(InputPrompt.is_public)\n return list(db_session.scalars(query).all())\n\n\ndef fetch_input_prompts_by_user(\n db_session: Session,\n user_id: UUID | None,\n active: bool | None = None,\n include_public: bool = False,\n) -> list[InputPrompt]:\n \"\"\"\n Returns all prompts belonging to the user or public prompts,\n excluding those the user has specifically disabled.\n \"\"\"\n\n query = select(InputPrompt)\n\n if user_id is not None:\n # If we have a user, left join to InputPrompt__User to check \"disabled\"\n IPU = aliased(InputPrompt__User)\n query = query.join(\n IPU,\n (IPU.input_prompt_id == InputPrompt.id) & (IPU.user_id == user_id),\n isouter=True,\n )\n\n # Exclude disabled prompts\n query = query.where(or_(IPU.disabled.is_(None), IPU.disabled.is_(False)))\n\n if include_public:\n # Return both user-owned and public prompts\n query = query.where(\n or_(\n InputPrompt.user_id == user_id,\n InputPrompt.is_public,\n )\n )\n else:\n # Return only user-owned prompts\n query = query.where(InputPrompt.user_id == user_id)\n\n else:\n # user_id is None - anonymous usage\n if include_public:\n query = query.where(InputPrompt.is_public)\n else:\n # No user and not requesting public prompts - return nothing\n return []\n\n if active is not None:\n query = query.where(InputPrompt.active == active)\n\n return list(db_session.scalars(query).all())\n\n\ndef disable_input_prompt_for_user(\n input_prompt_id: int,\n user_id: UUID,\n db_session: Session,\n) -> None:\n \"\"\"\n Sets (or creates) a record in InputPrompt__User with disabled=True\n so that this prompt is hidden for the user.\n \"\"\"\n ipu = (\n db_session.query(InputPrompt__User)\n .filter_by(input_prompt_id=input_prompt_id, user_id=user_id)\n .first()\n )\n\n if ipu is None:\n # Create a new association row\n ipu = InputPrompt__User(\n input_prompt_id=input_prompt_id, user_id=user_id, disabled=True\n )\n db_session.add(ipu)\n else:\n # Just update the existing record\n ipu.disabled = True\n\n db_session.commit()\n" }, { "path": "backend/onyx/db/kg_config.py", "content": "from onyx.configs.constants import KV_KG_CONFIG_KEY\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.kg.models import KGConfigSettings\nfrom onyx.server.kg.models import EnableKGConfigRequest\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef set_kg_config_settings(kg_config_settings: KGConfigSettings) -> None:\n kv_store = get_kv_store()\n kv_store.store(KV_KG_CONFIG_KEY, kg_config_settings.model_dump())\n\n\ndef get_kg_config_settings() -> KGConfigSettings:\n kv_store = get_kv_store()\n try:\n # refresh cache True until beta is over as we may manually update the config in the db\n stored_config = kv_store.load(KV_KG_CONFIG_KEY, refresh_cache=True)\n return KGConfigSettings.model_validate(stored_config or {})\n except KvKeyNotFoundError:\n # Default to empty kg config if no config have been set yet\n logger.debug(f\"No kg config found in KV store for key: {KV_KG_CONFIG_KEY}\")\n return KGConfigSettings()\n except Exception as e:\n logger.error(f\"Error loading kg config from KV store: {str(e)}\")\n return KGConfigSettings()\n\n\ndef validate_kg_settings(kg_config_settings: KGConfigSettings) -> None:\n if not kg_config_settings.KG_ENABLED:\n raise ValueError(\"KG is not enabled\")\n if not kg_config_settings.KG_VENDOR:\n raise ValueError(\"KG_VENDOR is not set\")\n if not kg_config_settings.KG_VENDOR_DOMAINS:\n raise ValueError(\"KG_VENDOR_DOMAINS is not set\")\n\n\ndef is_kg_config_settings_enabled_valid(kg_config_settings: KGConfigSettings) -> bool:\n try:\n validate_kg_settings(kg_config_settings)\n return True\n except Exception:\n return False\n\n\ndef enable_kg(enable_req: EnableKGConfigRequest) -> None:\n kg_config_settings = get_kg_config_settings()\n kg_config_settings.KG_ENABLED = True\n kg_config_settings.KG_VENDOR = enable_req.vendor\n kg_config_settings.KG_VENDOR_DOMAINS = enable_req.vendor_domains\n kg_config_settings.KG_IGNORE_EMAIL_DOMAINS = enable_req.ignore_domains\n kg_config_settings.KG_COVERAGE_START = enable_req.coverage_start.strftime(\n \"%Y-%m-%d\"\n )\n kg_config_settings.KG_MAX_COVERAGE_DAYS = 10000 # TODO: revisit after public beta\n\n validate_kg_settings(kg_config_settings)\n set_kg_config_settings(kg_config_settings)\n\n\ndef disable_kg() -> None:\n kg_config_settings = get_kg_config_settings()\n kg_config_settings.KG_ENABLED = False\n set_kg_config_settings(kg_config_settings)\n" }, { "path": "backend/onyx/db/kg_temp_view.py", "content": "# import random\n\n# from sqlalchemy import text\n# from sqlalchemy.ext.declarative import declarative_base\n# from sqlalchemy.orm import Session\n\n# from onyx.agents.agent_search.kb_search.models import KGViewNames\n# from onyx.configs.app_configs import DB_READONLY_USER\n# from onyx.configs.kg_configs import KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX\n# from onyx.configs.kg_configs import KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX\n# from onyx.configs.kg_configs import KG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX\n# from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n\n# Base = declarative_base()\n\n\n# def get_user_view_names(\n# user_email: str, tenant_id: str\n# ) -> KGViewNames:\n# user_email_cleaned = (\n# user_email.replace(\"@\", \"__\")\n# .replace(\".\", \"_\")\n# .replace(\"+\", \"_\")\n# )\n# random_suffix_str = str(\n# random.randint(1000000, 9999999)\n# )\n# return KGViewNames(\n# allowed_docs_view_name=(\n# f'\"{tenant_id}\".'\n# f\"{KG_TEMP_ALLOWED_DOCS_VIEW_NAME_PREFIX}_\"\n# f\"{user_email_cleaned}_{random_suffix_str}\"\n# ),\n# kg_relationships_view_name=(\n# f'\"{tenant_id}\".'\n# f\"{KG_TEMP_KG_RELATIONSHIPS_VIEW_NAME_PREFIX}_\"\n# f\"{user_email_cleaned}_{random_suffix_str}\"\n# ),\n# kg_entity_view_name=(\n# f'\"{tenant_id}\".'\n# f\"{KG_TEMP_KG_ENTITIES_VIEW_NAME_PREFIX}_\"\n# f\"{user_email_cleaned}_{random_suffix_str}\"\n# ),\n# )\n\n\n# # First, create the view definition\n# def create_views(\n# db_session: Session,\n# tenant_id: str,\n# user_email: str,\n# allowed_docs_view_name: str,\n# kg_relationships_view_name: str,\n# kg_entity_view_name: str,\n# ) -> None:\n\n# # Create ALLOWED_DOCS view\n# allowed_docs_view = text(\n# f\"\"\"\n# CREATE OR REPLACE VIEW {allowed_docs_view_name} AS\n# WITH kg_used_docs AS (\n# SELECT document_id as kg_used_doc_id\n# FROM \"{tenant_id}\".kg_entity d\n# WHERE document_id IS NOT NULL\n# ),\n\n# base_public_docs AS (\n# SELECT d.id as allowed_doc_id\n# FROM \"{tenant_id}\".document d\n# INNER JOIN kg_used_docs kud ON kud.kg_used_doc_id = d.id\n# WHERE d.is_public\n# ),\n# user_owned_and_public_docs AS (\n# SELECT d.id as allowed_doc_id\n# FROM \"{tenant_id}\".document_by_connector_credential_pair d\n# JOIN \"{tenant_id}\".credential c ON d.credential_id = c.id\n# JOIN \"{tenant_id}\".connector_credential_pair ccp ON\n# d.connector_id = ccp.connector_id AND\n# d.credential_id = ccp.credential_id\n# JOIN \"{tenant_id}\".user u ON c.user_id = u.id\n# INNER JOIN kg_used_docs kud ON kud.kg_used_doc_id = d.id\n# WHERE ccp.status != 'DELETING'\n# AND ccp.access_type != 'SYNC'\n# AND (u.email = :user_email or ccp.access_type::text = 'PUBLIC')\n# ),\n# user_group_accessible_docs AS (\n# SELECT d.id as allowed_doc_id\n# FROM \"{tenant_id}\".document_by_connector_credential_pair d\n# JOIN \"{tenant_id}\".connector_credential_pair ccp ON\n# d.connector_id = ccp.connector_id AND\n# d.credential_id = ccp.credential_id\n# JOIN \"{tenant_id}\".user_group__connector_credential_pair ugccp ON\n# ccp.id = ugccp.cc_pair_id\n# JOIN \"{tenant_id}\".user__user_group uug ON\n# uug.user_group_id = ugccp.user_group_id\n# JOIN \"{tenant_id}\".user u ON uug.user_id = u.id\n# INNER JOIN kg_used_docs kud ON kud.kg_used_doc_id = d.id\n# WHERE kud.kg_used_doc_id IS NOT NULL\n# AND ccp.status != 'DELETING'\n# AND ccp.access_type != 'SYNC'\n# AND u.email = :user_email\n# ),\n# external_user_docs AS (\n# SELECT d.id as allowed_doc_id\n# FROM \"{tenant_id}\".document d\n# INNER JOIN kg_used_docs kud ON kud.kg_used_doc_id = d.id\n# WHERE kud.kg_used_doc_id IS NOT NULL\n# AND :user_email = ANY(external_user_emails)\n# ),\n# external_group_docs AS (\n# SELECT d.id as allowed_doc_id\n# FROM \"{tenant_id}\".document d\n# INNER JOIN kg_used_docs kud ON kud.kg_used_doc_id = d.id\n# JOIN \"{tenant_id}\".user__external_user_group_id ueg ON ueg.external_user_group_id = ANY(d.external_user_group_ids)\n# JOIN \"{tenant_id}\".user u ON ueg.user_id = u.id\n# WHERE kud.kg_used_doc_id IS NOT NULL\n# AND u.email = :user_email\n# )\n# SELECT DISTINCT allowed_doc_id FROM (\n# SELECT allowed_doc_id FROM base_public_docs\n# UNION\n# SELECT allowed_doc_id FROM user_owned_and_public_docs\n# UNION\n# SELECT allowed_doc_id FROM user_group_accessible_docs\n# UNION\n# SELECT allowed_doc_id FROM external_user_docs\n# UNION\n# SELECT allowed_doc_id FROM external_group_docs\n# ) combined_docs\n# \"\"\"\n# ).bindparams(user_email=user_email)\n\n# # Create the main view that uses ALLOWED_DOCS for Relationships\n# kg_relationships_view = text(\n# f\"\"\"\n# CREATE OR REPLACE VIEW {kg_relationships_view_name} AS\n# SELECT kgr.id_name as relationship,\n# kgr.source_node as source_entity,\n# kgr.target_node as target_entity,\n# kgr.source_node_type as source_entity_type,\n# kgr.target_node_type as target_entity_type,\n# kgr.type as relationship_description,\n# kgr.relationship_type_id_name as relationship_type,\n# kgr.source_document as source_document,\n# d.doc_updated_at as source_date,\n# se.attributes as source_entity_attributes,\n# te.attributes as target_entity_attributes\n# FROM \"{tenant_id}\".kg_relationship kgr\n# INNER JOIN {allowed_docs_view_name} AD on AD.allowed_doc_id = kgr.source_document\n# JOIN \"{tenant_id}\".document d on d.id = kgr.source_document\n# JOIN \"{tenant_id}\".kg_entity se on se.id_name = kgr.source_node\n# JOIN \"{tenant_id}\".kg_entity te on te.id_name = kgr.target_node\n# \"\"\"\n# )\n\n# # Create the main view that uses ALLOWED_DOCS for Entities\n# kg_entity_view = text(\n# f\"\"\"\n# CREATE OR REPLACE VIEW {kg_entity_view_name} AS\n# SELECT kge.id_name as entity,\n# kge.entity_type_id_name as entity_type,\n# kge.attributes as entity_attributes,\n# kge.document_id as source_document,\n# d.doc_updated_at as source_date\n# FROM \"{tenant_id}\".kg_entity kge\n# INNER JOIN {allowed_docs_view_name} AD on AD.allowed_doc_id = kge.document_id\n# JOIN \"{tenant_id}\".document d on d.id = kge.document_id\n# \"\"\"\n# )\n\n# # Execute the views using the session\n# db_session.execute(allowed_docs_view)\n# db_session.execute(kg_relationships_view)\n# db_session.execute(kg_entity_view)\n\n# # Grant permissions on view to readonly user\n\n# db_session.execute(\n# text(f\"GRANT SELECT ON {kg_relationships_view_name} TO {DB_READONLY_USER}\")\n# )\n# db_session.execute(\n# text(f\"GRANT SELECT ON {kg_entity_view_name} TO {DB_READONLY_USER}\")\n# )\n\n# db_session.commit()\n\n# return None\n\n\n# def drop_views(\n# allowed_docs_view_name: str | None = None,\n# kg_relationships_view_name: str | None = None,\n# kg_entity_view_name: str | None = None,\n# ) -> None:\n# \"\"\"\n# Drops the temporary views created by create_views.\n\n# Args:\n# db_session: SQLAlchemy session\n# allowed_docs_view_name: Name of the allowed_docs view\n# kg_relationships_view_name: Name of the allowed kg_relationships view\n# kg_entity_view_name: Name of the allowed kg_entity view\n# \"\"\"\n\n# with get_session_with_current_tenant() as db_drop_session:\n# if kg_relationships_view_name:\n# revoke_kg_relationships = text(\n# f\"REVOKE SELECT ON {kg_relationships_view_name} FROM {DB_READONLY_USER}\"\n# )\n# db_drop_session.execute(revoke_kg_relationships)\n# drop_kg_relationships = text(\n# f\"DROP VIEW IF EXISTS {kg_relationships_view_name}\"\n# )\n# db_drop_session.execute(drop_kg_relationships)\n\n# if kg_entity_view_name:\n# revoke_kg_entities = text(\n# f\"REVOKE SELECT ON {kg_entity_view_name} FROM {DB_READONLY_USER}\"\n# )\n# db_drop_session.execute(revoke_kg_entities)\n# drop_kg_entities = text(f\"DROP VIEW IF EXISTS {kg_entity_view_name}\")\n# db_drop_session.execute(drop_kg_entities)\n\n# if allowed_docs_view_name:\n# drop_allowed_docs = text(f\"DROP VIEW IF EXISTS {allowed_docs_view_name}\")\n# db_drop_session.execute(drop_allowed_docs)\n\n# db_drop_session.commit()\n# return None\n" }, { "path": "backend/onyx/db/llm.py", "content": "from sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import LLMModelFlowType\nfrom onyx.db.models import CloudEmbeddingProvider as CloudEmbeddingProviderModel\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import ImageGenerationConfig\nfrom onyx.db.models import LLMModelFlow\nfrom onyx.db.models import LLMProvider as LLMProviderModel\nfrom onyx.db.models import LLMProvider__Persona\nfrom onyx.db.models import LLMProvider__UserGroup\nfrom onyx.db.models import ModelConfiguration\nfrom onyx.db.models import Persona\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.models import Tool as ToolModel\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.llm.utils import model_supports_image_input\nfrom onyx.llm.well_known_providers.auto_update_models import LLMRecommendations\nfrom onyx.server.manage.embedding.models import CloudEmbeddingProvider\nfrom onyx.server.manage.embedding.models import CloudEmbeddingProviderCreationRequest\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.server.manage.llm.models import SyncModelEntry\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.enums import EmbeddingProvider\n\nlogger = setup_logger()\n\n\ndef update_group_llm_provider_relationships__no_commit(\n llm_provider_id: int,\n group_ids: list[int] | None,\n db_session: Session,\n) -> None:\n # Delete existing relationships\n db_session.query(LLMProvider__UserGroup).filter(\n LLMProvider__UserGroup.llm_provider_id == llm_provider_id\n ).delete(synchronize_session=\"fetch\")\n\n # Add new relationships from given group_ids\n if group_ids:\n new_relationships = [\n LLMProvider__UserGroup(\n llm_provider_id=llm_provider_id,\n user_group_id=group_id,\n )\n for group_id in group_ids\n ]\n db_session.add_all(new_relationships)\n\n\ndef update_llm_provider_persona_relationships__no_commit(\n db_session: Session,\n llm_provider_id: int,\n persona_ids: list[int] | None,\n) -> None:\n \"\"\"Replace the persona restrictions for a provider within an open transaction.\"\"\"\n db_session.execute(\n delete(LLMProvider__Persona).where(\n LLMProvider__Persona.llm_provider_id == llm_provider_id\n )\n )\n\n if persona_ids:\n db_session.add_all(\n LLMProvider__Persona(\n llm_provider_id=llm_provider_id,\n persona_id=persona_id,\n )\n for persona_id in persona_ids\n )\n\n\ndef fetch_user_group_ids(db_session: Session, user: User) -> set[int]:\n \"\"\"Fetch the set of user group IDs for a given user.\n\n Args:\n db_session: Database session\n user: User to fetch groups for\n\n Returns:\n Set of user group IDs. Empty set for anonymous users.\n \"\"\"\n if user.is_anonymous:\n return set()\n\n return set(\n db_session.scalars(\n select(User__UserGroup.user_group_id).where(\n User__UserGroup.user_id == user.id\n )\n ).all()\n )\n\n\ndef can_user_access_llm_provider(\n provider: LLMProviderModel,\n user_group_ids: set[int],\n persona: Persona | None,\n is_admin: bool = False,\n) -> bool:\n \"\"\"Check if a user may use an LLM provider.\n\n Args:\n provider: The LLM provider to check access for\n user_group_ids: Set of user group IDs the user belongs to\n persona: The persona being used (if any)\n is_admin: If True, bypass user group restrictions but still respect persona restrictions\n\n Access logic:\n - is_public controls USER access (group bypass): when True, all users can access\n regardless of group membership. When False, user must be in a whitelisted group\n (or be admin).\n - Persona restrictions are ALWAYS enforced when set, regardless of is_public.\n This allows admins to make a provider available to all users while still\n restricting which personas (assistants) can use it.\n\n Decision matrix:\n 1. is_public=True, no personas set → everyone has access\n 2. is_public=True, personas set → all users, but only whitelisted personas\n 3. is_public=False, groups+personas set → must satisfy BOTH (admins bypass groups)\n 4. is_public=False, only groups set → must be in group (admins bypass)\n 5. is_public=False, only personas set → must use whitelisted persona\n 6. is_public=False, neither set → admin-only (locked)\n \"\"\"\n provider_group_ids = {g.id for g in (provider.groups or [])}\n provider_persona_ids = {p.id for p in (provider.personas or [])}\n has_groups = bool(provider_group_ids)\n has_personas = bool(provider_persona_ids)\n\n # Persona restrictions are always enforced when set, regardless of is_public\n if has_personas and not (persona and persona.id in provider_persona_ids):\n return False\n\n if provider.is_public:\n return True\n\n if has_groups:\n return is_admin or bool(user_group_ids & provider_group_ids)\n\n # No groups: either persona-whitelisted (already passed) or admin-only if locked\n return has_personas or is_admin\n\n\ndef validate_persona_ids_exist(\n db_session: Session, persona_ids: list[int]\n) -> tuple[set[int], list[int]]:\n \"\"\"Validate that persona IDs exist in the database.\n\n Returns:\n Tuple of (fetched_persona_ids, missing_personas)\n \"\"\"\n fetched_persona_ids = set(\n db_session.scalars(select(Persona.id).where(Persona.id.in_(persona_ids))).all()\n )\n missing_personas = sorted(set(persona_ids) - fetched_persona_ids)\n return fetched_persona_ids, missing_personas\n\n\ndef get_personas_using_provider(\n db_session: Session, provider_name: str\n) -> list[Persona]:\n \"\"\"Get all non-deleted personas that use a specific LLM provider.\"\"\"\n return list(\n db_session.scalars(\n select(Persona).where(\n Persona.llm_model_provider_override == provider_name,\n Persona.deleted == False, # noqa: E712\n )\n ).all()\n )\n\n\ndef fetch_persona_with_groups(db_session: Session, persona_id: int) -> Persona | None:\n \"\"\"Fetch a persona with its groups eagerly loaded.\"\"\"\n return db_session.scalar(\n select(Persona)\n .options(selectinload(Persona.groups))\n .where(Persona.id == persona_id, Persona.deleted == False) # noqa: E712\n )\n\n\ndef upsert_cloud_embedding_provider(\n db_session: Session, provider: CloudEmbeddingProviderCreationRequest\n) -> CloudEmbeddingProvider:\n existing_provider = (\n db_session.query(CloudEmbeddingProviderModel)\n .filter_by(provider_type=provider.provider_type)\n .first()\n )\n if existing_provider:\n for key, value in provider.model_dump().items():\n setattr(existing_provider, key, value)\n else:\n new_provider = CloudEmbeddingProviderModel(**provider.model_dump())\n\n db_session.add(new_provider)\n existing_provider = new_provider\n db_session.commit()\n db_session.refresh(existing_provider)\n return CloudEmbeddingProvider.from_request(existing_provider)\n\n\ndef upsert_llm_provider(\n llm_provider_upsert_request: LLMProviderUpsertRequest,\n db_session: Session,\n) -> LLMProviderView:\n existing_llm_provider: LLMProviderModel | None = None\n if llm_provider_upsert_request.id:\n existing_llm_provider = fetch_existing_llm_provider_by_id(\n id=llm_provider_upsert_request.id, db_session=db_session\n )\n if not existing_llm_provider:\n raise ValueError(\n f\"LLM provider with id {llm_provider_upsert_request.id} not found\"\n )\n\n if existing_llm_provider.name != llm_provider_upsert_request.name:\n raise ValueError(\n f\"LLM provider with id {llm_provider_upsert_request.id} name change not allowed\"\n )\n else:\n existing_llm_provider = fetch_existing_llm_provider(\n name=llm_provider_upsert_request.name, db_session=db_session\n )\n if existing_llm_provider:\n raise ValueError(\n f\"LLM provider with name '{llm_provider_upsert_request.name}' already exists\"\n )\n existing_llm_provider = LLMProviderModel(name=llm_provider_upsert_request.name)\n db_session.add(existing_llm_provider)\n\n # Filter out empty strings and None values from custom_config to allow\n # providers like Bedrock to fall back to IAM roles when credentials are not provided\n custom_config = llm_provider_upsert_request.custom_config\n if custom_config:\n custom_config = {\n k: v for k, v in custom_config.items() if v is not None and v.strip() != \"\"\n }\n # Set to None if the dict is empty after filtering\n custom_config = custom_config or None\n\n api_base = llm_provider_upsert_request.api_base or None\n existing_llm_provider.provider = llm_provider_upsert_request.provider\n # EncryptedString accepts str for writes, returns SensitiveValue for reads\n existing_llm_provider.api_key = llm_provider_upsert_request.api_key # type: ignore[assignment]\n existing_llm_provider.api_base = api_base\n existing_llm_provider.api_version = llm_provider_upsert_request.api_version\n existing_llm_provider.custom_config = custom_config\n\n existing_llm_provider.is_public = llm_provider_upsert_request.is_public\n existing_llm_provider.is_auto_mode = llm_provider_upsert_request.is_auto_mode\n existing_llm_provider.deployment_name = llm_provider_upsert_request.deployment_name\n\n if not existing_llm_provider.id:\n # If its not already in the db, we need to generate an ID by flushing\n db_session.flush()\n\n # Build a lookup of existing model configurations by name (single iteration)\n existing_by_name = {\n mc.name: mc for mc in existing_llm_provider.model_configurations\n }\n\n models_to_exist = {\n mc.name for mc in llm_provider_upsert_request.model_configurations\n }\n\n # Build a lookup of requested visibility by model name\n requested_visibility = {\n mc.name: mc.is_visible\n for mc in llm_provider_upsert_request.model_configurations\n }\n\n # Delete removed models\n removed_ids = [\n mc.id for name, mc in existing_by_name.items() if name not in models_to_exist\n ]\n\n default_model = fetch_default_llm_model(db_session)\n\n # Prevent removing and hiding the default model\n if default_model:\n for name, mc in existing_by_name.items():\n if mc.id == default_model.id:\n if default_model.id in removed_ids:\n raise ValueError(\n f\"Cannot remove the default model '{name}'. Please change the default model before removing.\"\n )\n if not requested_visibility.get(name, True):\n raise ValueError(\n f\"Cannot hide the default model '{name}'. Please change the default model before hiding.\"\n )\n break\n\n if removed_ids:\n db_session.query(ModelConfiguration).filter(\n ModelConfiguration.id.in_(removed_ids)\n ).delete(synchronize_session=\"fetch\")\n db_session.flush()\n\n # Import here to avoid circular imports\n from onyx.llm.utils import get_max_input_tokens\n\n for model_config in llm_provider_upsert_request.model_configurations:\n max_input_tokens = model_config.max_input_tokens\n if max_input_tokens is None:\n max_input_tokens = get_max_input_tokens(\n model_name=model_config.name,\n model_provider=llm_provider_upsert_request.provider,\n )\n\n supported_flows = [LLMModelFlowType.CHAT]\n if model_config.supports_image_input:\n supported_flows.append(LLMModelFlowType.VISION)\n\n existing = existing_by_name.get(model_config.name)\n if existing:\n update_model_configuration__no_commit(\n db_session=db_session,\n model_configuration_id=existing.id,\n supported_flows=supported_flows,\n is_visible=model_config.is_visible,\n max_input_tokens=max_input_tokens,\n display_name=model_config.display_name,\n )\n else:\n insert_new_model_configuration__no_commit(\n db_session=db_session,\n llm_provider_id=existing_llm_provider.id,\n model_name=model_config.name,\n supported_flows=supported_flows,\n is_visible=model_config.is_visible,\n max_input_tokens=max_input_tokens,\n display_name=model_config.display_name,\n )\n\n # Make sure the relationship table stays up to date\n update_group_llm_provider_relationships__no_commit(\n llm_provider_id=existing_llm_provider.id,\n group_ids=llm_provider_upsert_request.groups,\n db_session=db_session,\n )\n update_llm_provider_persona_relationships__no_commit(\n db_session=db_session,\n llm_provider_id=existing_llm_provider.id,\n persona_ids=llm_provider_upsert_request.personas,\n )\n\n db_session.flush()\n db_session.refresh(existing_llm_provider)\n\n try:\n db_session.commit()\n except Exception as e:\n db_session.rollback()\n raise ValueError(f\"Failed to save LLM provider: {str(e)}\") from e\n\n full_llm_provider = LLMProviderView.from_model(existing_llm_provider)\n return full_llm_provider\n\n\ndef sync_model_configurations(\n db_session: Session,\n provider_name: str,\n models: list[SyncModelEntry],\n) -> int:\n \"\"\"Sync model configurations for a dynamic provider (OpenRouter, Bedrock, Ollama, etc.).\n\n This inserts NEW models from the source API without overwriting existing ones.\n User preferences (is_visible, max_input_tokens) are preserved for existing models.\n\n Args:\n db_session: Database session\n provider_name: Name of the LLM provider\n models: List of SyncModelEntry objects describing the fetched models\n\n Returns:\n Number of new models added\n \"\"\"\n provider = fetch_existing_llm_provider(name=provider_name, db_session=db_session)\n if not provider:\n raise ValueError(f\"LLM Provider '{provider_name}' not found\")\n\n # Get existing model names to count new additions\n existing_names = {mc.name for mc in provider.model_configurations}\n\n new_count = 0\n for model in models:\n if model.name not in existing_names:\n # Insert new model with is_visible=False (user must explicitly enable)\n supported_flows = [LLMModelFlowType.CHAT]\n if model.supports_image_input:\n supported_flows.append(LLMModelFlowType.VISION)\n\n insert_new_model_configuration__no_commit(\n db_session=db_session,\n llm_provider_id=provider.id,\n model_name=model.name,\n supported_flows=supported_flows,\n is_visible=False,\n max_input_tokens=model.max_input_tokens,\n display_name=model.display_name,\n )\n new_count += 1\n\n if new_count > 0:\n db_session.commit()\n\n return new_count\n\n\ndef fetch_existing_embedding_providers(\n db_session: Session,\n) -> list[CloudEmbeddingProviderModel]:\n return list(db_session.scalars(select(CloudEmbeddingProviderModel)).all())\n\n\ndef fetch_existing_doc_sets(\n db_session: Session, doc_ids: list[int]\n) -> list[DocumentSet]:\n return list(\n db_session.scalars(select(DocumentSet).where(DocumentSet.id.in_(doc_ids))).all()\n )\n\n\ndef fetch_existing_tools(db_session: Session, tool_ids: list[int]) -> list[ToolModel]:\n return list(\n db_session.scalars(select(ToolModel).where(ToolModel.id.in_(tool_ids))).all()\n )\n\n\ndef fetch_existing_models(\n db_session: Session,\n flow_types: list[LLMModelFlowType],\n) -> list[ModelConfiguration]:\n models = (\n select(ModelConfiguration)\n .join(LLMModelFlow)\n .where(LLMModelFlow.llm_model_flow_type.in_(flow_types))\n .options(\n selectinload(ModelConfiguration.llm_provider),\n selectinload(ModelConfiguration.llm_model_flows),\n )\n )\n\n return list(db_session.scalars(models).all())\n\n\ndef fetch_existing_llm_providers(\n db_session: Session,\n flow_type_filter: list[LLMModelFlowType],\n only_public: bool = False,\n exclude_image_generation_providers: bool = True,\n) -> list[LLMProviderModel]:\n \"\"\"Fetch all LLM providers with optional filtering.\n\n Args:\n db_session: Database session\n flow_type_filter: List of flow types to filter by, empty list for no filter\n only_public: If True, only return public providers\n exclude_image_generation_providers: If True, exclude providers that are\n used for image generation configs\n \"\"\"\n stmt = select(LLMProviderModel)\n\n if flow_type_filter:\n providers_with_flows = (\n select(ModelConfiguration.llm_provider_id)\n .join(LLMModelFlow)\n .where(LLMModelFlow.llm_model_flow_type.in_(flow_type_filter))\n .distinct()\n )\n stmt = stmt.where(LLMProviderModel.id.in_(providers_with_flows))\n\n if exclude_image_generation_providers:\n image_gen_provider_ids = select(ModelConfiguration.llm_provider_id).join(\n ImageGenerationConfig\n )\n stmt = stmt.where(~LLMProviderModel.id.in_(image_gen_provider_ids))\n\n stmt = stmt.options(\n selectinload(LLMProviderModel.model_configurations),\n selectinload(LLMProviderModel.groups),\n selectinload(LLMProviderModel.personas),\n )\n\n providers = list(db_session.scalars(stmt).all())\n if only_public:\n return [provider for provider in providers if provider.is_public]\n return providers\n\n\ndef fetch_existing_llm_provider(\n name: str, db_session: Session\n) -> LLMProviderModel | None:\n provider_model = db_session.scalar(\n select(LLMProviderModel)\n .where(LLMProviderModel.name == name)\n .options(\n selectinload(LLMProviderModel.model_configurations),\n selectinload(LLMProviderModel.groups),\n selectinload(LLMProviderModel.personas),\n )\n )\n\n return provider_model\n\n\ndef fetch_existing_llm_provider_by_id(\n id: int, db_session: Session\n) -> LLMProviderModel | None:\n provider_model = db_session.scalar(\n select(LLMProviderModel)\n .where(LLMProviderModel.id == id)\n .options(\n selectinload(LLMProviderModel.model_configurations),\n selectinload(LLMProviderModel.groups),\n selectinload(LLMProviderModel.personas),\n )\n )\n\n return provider_model\n\n\ndef fetch_embedding_provider(\n db_session: Session, provider_type: EmbeddingProvider\n) -> CloudEmbeddingProviderModel | None:\n return db_session.scalar(\n select(CloudEmbeddingProviderModel).where(\n CloudEmbeddingProviderModel.provider_type == provider_type\n )\n )\n\n\ndef fetch_default_llm_model(db_session: Session) -> ModelConfiguration | None:\n return fetch_default_model(db_session, LLMModelFlowType.CHAT)\n\n\ndef fetch_default_vision_model(db_session: Session) -> ModelConfiguration | None:\n return fetch_default_model(db_session, LLMModelFlowType.VISION)\n\n\ndef fetch_default_contextual_rag_model(\n db_session: Session,\n) -> ModelConfiguration | None:\n return fetch_default_model(db_session, LLMModelFlowType.CONTEXTUAL_RAG)\n\n\ndef fetch_default_model(\n db_session: Session,\n flow_type: LLMModelFlowType,\n) -> ModelConfiguration | None:\n model_config = db_session.scalar(\n select(ModelConfiguration)\n .options(selectinload(ModelConfiguration.llm_provider))\n .join(LLMModelFlow)\n .where(\n LLMModelFlow.llm_model_flow_type == flow_type,\n LLMModelFlow.is_default == True, # noqa: E712\n )\n )\n\n return model_config\n\n\ndef fetch_llm_provider_view(\n db_session: Session, provider_name: str\n) -> LLMProviderView | None:\n provider_model = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n if not provider_model:\n return None\n return LLMProviderView.from_model(provider_model)\n\n\ndef remove_embedding_provider(\n db_session: Session, provider_type: EmbeddingProvider\n) -> None:\n db_session.execute(\n delete(SearchSettings).where(SearchSettings.provider_type == provider_type)\n )\n\n # Delete the embedding provider\n db_session.execute(\n delete(CloudEmbeddingProviderModel).where(\n CloudEmbeddingProviderModel.provider_type == provider_type\n )\n )\n\n db_session.commit()\n\n\ndef remove_llm_provider(db_session: Session, provider_id: int) -> None:\n provider = db_session.get(LLMProviderModel, provider_id)\n if not provider:\n raise ValueError(\"LLM Provider not found\")\n\n # Clear the provider override from any personas using it\n # This causes them to fall back to the default provider\n personas_using_provider = get_personas_using_provider(db_session, provider.name)\n for persona in personas_using_provider:\n persona.llm_model_provider_override = None\n\n db_session.execute(\n delete(LLMProvider__UserGroup).where(\n LLMProvider__UserGroup.llm_provider_id == provider_id\n )\n )\n # Remove LLMProvider\n db_session.execute(\n delete(LLMProviderModel).where(LLMProviderModel.id == provider_id)\n )\n db_session.commit()\n\n\ndef remove_llm_provider__no_commit(db_session: Session, provider_id: int) -> None:\n \"\"\"Remove LLM provider.\"\"\"\n provider = db_session.get(LLMProviderModel, provider_id)\n if not provider:\n raise ValueError(\"LLM Provider not found\")\n\n # Clear the provider override from any personas using it\n # This causes them to fall back to the default provider\n personas_using_provider = get_personas_using_provider(db_session, provider.name)\n for persona in personas_using_provider:\n persona.llm_model_provider_override = None\n\n db_session.execute(\n delete(LLMProvider__UserGroup).where(\n LLMProvider__UserGroup.llm_provider_id == provider_id\n )\n )\n # Remove LLMProvider\n db_session.execute(\n delete(LLMProviderModel).where(LLMProviderModel.id == provider_id)\n )\n db_session.flush()\n\n\ndef update_default_provider(\n provider_id: int, model_name: str, db_session: Session\n) -> None:\n _update_default_model(\n db_session,\n provider_id,\n model_name,\n LLMModelFlowType.CHAT,\n )\n\n\ndef update_default_vision_provider(\n provider_id: int, vision_model: str, db_session: Session\n) -> None:\n provider = db_session.scalar(\n select(LLMProviderModel).where(\n LLMProviderModel.id == provider_id,\n )\n )\n\n if provider is None:\n raise ValueError(f\"LLM Provider with id={provider_id} does not exist\")\n\n if not model_supports_image_input(vision_model, provider.provider):\n raise ValueError(\n f\"Model '{vision_model}' for provider '{provider.provider} does not support image input\"\n )\n\n _update_default_model(\n db_session=db_session,\n provider_id=provider_id,\n model=vision_model,\n flow_type=LLMModelFlowType.VISION,\n )\n\n\ndef update_no_default_contextual_rag_provider(\n db_session: Session,\n) -> None:\n db_session.execute(\n update(LLMModelFlow)\n .where(\n LLMModelFlow.llm_model_flow_type == LLMModelFlowType.CONTEXTUAL_RAG,\n LLMModelFlow.is_default == True, # noqa: E712\n )\n .values(is_default=False)\n )\n db_session.commit()\n\n\ndef update_default_contextual_model(\n db_session: Session,\n enable_contextual_rag: bool,\n contextual_rag_llm_provider: str | None,\n contextual_rag_llm_name: str | None,\n) -> None:\n \"\"\"Sets or clears the default contextual RAG model.\n\n Should be called whenever the PRESENT search settings change\n (e.g. inline update or FUTURE → PRESENT swap).\n \"\"\"\n if (\n not enable_contextual_rag\n or not contextual_rag_llm_name\n or not contextual_rag_llm_provider\n ):\n update_no_default_contextual_rag_provider(db_session=db_session)\n return\n\n provider = fetch_existing_llm_provider(\n name=contextual_rag_llm_provider, db_session=db_session\n )\n if not provider:\n raise ValueError(f\"Provider '{contextual_rag_llm_provider}' not found\")\n\n model_config = next(\n (\n mc\n for mc in provider.model_configurations\n if mc.name == contextual_rag_llm_name\n ),\n None,\n )\n if not model_config:\n raise ValueError(\n f\"Model '{contextual_rag_llm_name}' not found for provider '{contextual_rag_llm_provider}'\"\n )\n\n add_model_to_flow(\n db_session=db_session,\n model_configuration_id=model_config.id,\n flow_type=LLMModelFlowType.CONTEXTUAL_RAG,\n )\n _update_default_model(\n db_session=db_session,\n provider_id=provider.id,\n model=contextual_rag_llm_name,\n flow_type=LLMModelFlowType.CONTEXTUAL_RAG,\n )\n\n return\n\n\ndef fetch_auto_mode_providers(db_session: Session) -> list[LLMProviderModel]:\n \"\"\"Fetch all LLM providers that are in Auto mode.\"\"\"\n query = (\n select(LLMProviderModel)\n .where(LLMProviderModel.is_auto_mode.is_(True))\n .options(selectinload(LLMProviderModel.model_configurations))\n )\n return list(db_session.scalars(query).all())\n\n\ndef sync_auto_mode_models(\n db_session: Session,\n provider: LLMProviderModel,\n llm_recommendations: LLMRecommendations,\n) -> int:\n \"\"\"Sync models from GitHub config to a provider in Auto mode.\n\n In Auto mode, the model list and default are controlled by GitHub config.\n The schema has:\n - default_model: The default model config (always visible)\n - additional_visible_models: List of additional visible models\n\n Admin only provides API credentials.\n\n Args:\n db_session: Database session\n provider: LLM provider in Auto mode\n github_config: Configuration from GitHub\n\n Returns:\n The number of changes made.\n \"\"\"\n changes = 0\n\n # Build the list of all visible models from the config\n # All models in the config are visible (default + additional_visible_models)\n recommended_visible_models = llm_recommendations.get_visible_models(\n provider.provider\n )\n recommended_visible_model_names = [\n model.name for model in recommended_visible_models\n ]\n\n # Get existing models\n existing_models: dict[str, ModelConfiguration] = {\n mc.name: mc\n for mc in db_session.scalars(\n select(ModelConfiguration).where(\n ModelConfiguration.llm_provider_id == provider.id\n )\n ).all()\n }\n\n # Mark models that are no longer in GitHub config as not visible\n for model_name, model in existing_models.items():\n if model_name not in recommended_visible_model_names:\n if model.is_visible:\n model.is_visible = False\n changes += 1\n\n # Add or update models from GitHub config\n for model_config in recommended_visible_models:\n if model_config.name in existing_models:\n # Update existing model\n existing = existing_models[model_config.name]\n # Check each field for changes\n updated = False\n if existing.display_name != model_config.display_name:\n existing.display_name = model_config.display_name\n updated = True\n # All models in the config are visible\n if not existing.is_visible:\n existing.is_visible = True\n updated = True\n if updated:\n changes += 1\n else:\n # Add new model - all models from GitHub config are visible\n insert_new_model_configuration__no_commit(\n db_session=db_session,\n llm_provider_id=provider.id,\n model_name=model_config.name,\n supported_flows=[LLMModelFlowType.CHAT],\n is_visible=True,\n max_input_tokens=None,\n display_name=model_config.display_name,\n )\n changes += 1\n\n # Update the default if this provider currently holds the global CHAT default.\n # We flush (but don't commit) so that _update_default_model can see the new\n # model rows, then commit everything atomically to avoid a window where the\n # old default is invisible but still pointed-to.\n db_session.flush()\n\n recommended_default = llm_recommendations.get_default_model(provider.provider)\n if recommended_default:\n current_default = fetch_default_llm_model(db_session)\n\n if (\n current_default\n and current_default.llm_provider_id == provider.id\n and current_default.name != recommended_default.name\n ):\n _update_default_model__no_commit(\n db_session=db_session,\n provider_id=provider.id,\n model=recommended_default.name,\n flow_type=LLMModelFlowType.CHAT,\n )\n changes += 1\n\n db_session.commit()\n return changes\n\n\ndef create_new_flow_mapping__no_commit(\n db_session: Session,\n model_configuration_id: int,\n flow_type: LLMModelFlowType,\n) -> LLMModelFlow:\n result = db_session.execute(\n insert(LLMModelFlow)\n .values(\n model_configuration_id=model_configuration_id,\n llm_model_flow_type=flow_type,\n is_default=False,\n )\n .on_conflict_do_nothing()\n .returning(LLMModelFlow)\n )\n\n flow = result.scalar()\n if not flow:\n # Row already exists — fetch it\n flow = db_session.scalar(\n select(LLMModelFlow).where(\n LLMModelFlow.model_configuration_id == model_configuration_id,\n LLMModelFlow.llm_model_flow_type == flow_type,\n )\n )\n if not flow:\n raise ValueError(\n f\"Failed to create or find flow mapping for model_configuration_id={model_configuration_id} and flow_type={flow_type}\"\n )\n\n return flow\n\n\ndef insert_new_model_configuration__no_commit(\n db_session: Session,\n llm_provider_id: int,\n model_name: str,\n supported_flows: list[LLMModelFlowType],\n is_visible: bool,\n max_input_tokens: int | None,\n display_name: str | None,\n) -> int | None:\n result = db_session.execute(\n insert(ModelConfiguration)\n .values(\n llm_provider_id=llm_provider_id,\n name=model_name,\n is_visible=is_visible,\n max_input_tokens=max_input_tokens,\n display_name=display_name,\n supports_image_input=LLMModelFlowType.VISION in supported_flows,\n )\n .on_conflict_do_nothing()\n .returning(ModelConfiguration.id)\n )\n\n model_config_id = result.scalar()\n\n if not model_config_id:\n return None\n\n for flow_type in supported_flows:\n create_new_flow_mapping__no_commit(\n db_session=db_session,\n model_configuration_id=model_config_id,\n flow_type=flow_type,\n )\n\n return model_config_id\n\n\ndef update_model_configuration__no_commit(\n db_session: Session,\n model_configuration_id: int,\n supported_flows: list[LLMModelFlowType],\n is_visible: bool,\n max_input_tokens: int | None,\n display_name: str | None,\n) -> None:\n result = db_session.execute(\n update(ModelConfiguration)\n .values(\n is_visible=is_visible,\n max_input_tokens=max_input_tokens,\n display_name=display_name,\n supports_image_input=LLMModelFlowType.VISION in supported_flows,\n )\n .where(ModelConfiguration.id == model_configuration_id)\n .returning(ModelConfiguration)\n )\n\n model_configuration = result.scalar()\n if not model_configuration:\n raise ValueError(\n f\"Failed to update model configuration with id={model_configuration_id}\"\n )\n\n new_flows = {\n flow_type\n for flow_type in supported_flows\n if flow_type not in model_configuration.llm_model_flow_types\n }\n removed_flows = {\n flow_type\n for flow_type in model_configuration.llm_model_flow_types\n if flow_type not in supported_flows\n }\n\n for flow_type in new_flows:\n create_new_flow_mapping__no_commit(\n db_session=db_session,\n model_configuration_id=model_configuration_id,\n flow_type=flow_type,\n )\n\n for flow_type in removed_flows:\n db_session.execute(\n delete(LLMModelFlow).where(\n LLMModelFlow.model_configuration_id == model_configuration_id,\n LLMModelFlow.llm_model_flow_type == flow_type,\n )\n )\n\n db_session.flush()\n\n\ndef _update_default_model__no_commit(\n db_session: Session,\n provider_id: int,\n model: str,\n flow_type: LLMModelFlowType,\n) -> None:\n result = db_session.execute(\n select(ModelConfiguration, LLMModelFlow)\n .join(\n LLMModelFlow, LLMModelFlow.model_configuration_id == ModelConfiguration.id\n )\n .where(\n ModelConfiguration.llm_provider_id == provider_id,\n ModelConfiguration.name == model,\n LLMModelFlow.llm_model_flow_type == flow_type,\n )\n ).first()\n\n if not result:\n raise ValueError(\n f\"Model '{model}' is not a valid model for provider_id={provider_id}\"\n )\n\n model_config, new_default = result\n\n # Clear existing default and set in an atomic operation\n db_session.execute(\n update(LLMModelFlow)\n .where(\n LLMModelFlow.llm_model_flow_type == flow_type,\n LLMModelFlow.is_default == True, # noqa: E712\n )\n .values(is_default=False)\n )\n\n new_default.is_default = True\n model_config.is_visible = True\n\n\ndef _update_default_model(\n db_session: Session,\n provider_id: int,\n model: str,\n flow_type: LLMModelFlowType,\n) -> None:\n _update_default_model__no_commit(db_session, provider_id, model, flow_type)\n db_session.commit()\n\n\ndef add_model_to_flow(\n db_session: Session,\n model_configuration_id: int,\n flow_type: LLMModelFlowType,\n) -> None:\n # Function does nothing on conflict\n create_new_flow_mapping__no_commit(\n db_session=db_session,\n model_configuration_id=model_configuration_id,\n flow_type=flow_type,\n )\n\n db_session.commit()\n" }, { "path": "backend/onyx/db/mcp.py", "content": "import datetime\nfrom typing import cast\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.orm.attributes import flag_modified\n\nfrom onyx.db.enums import MCPAuthenticationPerformer\nfrom onyx.db.enums import MCPServerStatus\nfrom onyx.db.enums import MCPTransport\nfrom onyx.db.models import MCPAuthenticationType\nfrom onyx.db.models import MCPConnectionConfig\nfrom onyx.db.models import MCPServer\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Tool\nfrom onyx.db.models import User\nfrom onyx.server.features.mcp.models import MCPConnectionData\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.sensitive import SensitiveValue\n\nlogger = setup_logger()\n\n\n# MCPServer operations\ndef get_all_mcp_servers(db_session: Session) -> list[MCPServer]:\n \"\"\"Get all MCP servers\"\"\"\n return list(\n db_session.scalars(select(MCPServer).order_by(MCPServer.created_at)).all()\n )\n\n\ndef get_mcp_server_by_id(server_id: int, db_session: Session) -> MCPServer:\n \"\"\"Get MCP server by ID\"\"\"\n server = db_session.scalar(select(MCPServer).where(MCPServer.id == server_id))\n if not server:\n raise ValueError(\"MCP server by specified id does not exist\")\n return server\n\n\ndef get_mcp_servers_by_owner(owner_email: str, db_session: Session) -> list[MCPServer]:\n \"\"\"Get all MCP servers owned by a specific user\"\"\"\n return list(\n db_session.scalars(\n select(MCPServer).where(MCPServer.owner == owner_email)\n ).all()\n )\n\n\ndef get_mcp_servers_for_persona(\n persona_id: int,\n db_session: Session,\n user: User, # noqa: ARG001\n) -> list[MCPServer]:\n \"\"\"Get all MCP servers associated with a persona via its tools\"\"\"\n # Get the persona and its tools\n persona = db_session.query(Persona).filter(Persona.id == persona_id).first()\n if not persona:\n return []\n\n # Collect unique MCP server IDs from the persona's tools\n mcp_server_ids = set()\n for tool in persona.tools:\n if tool.mcp_server_id:\n mcp_server_ids.add(tool.mcp_server_id)\n\n if not mcp_server_ids:\n return []\n\n # Fetch the MCP servers\n mcp_servers = (\n db_session.query(MCPServer).filter(MCPServer.id.in_(mcp_server_ids)).all()\n )\n\n return list(mcp_servers)\n\n\ndef get_mcp_servers_accessible_to_user(\n user_id: UUID, db_session: Session\n) -> list[MCPServer]:\n \"\"\"Get all MCP servers accessible to a user (directly or through groups)\"\"\"\n user = db_session.scalar(select(User).where(User.id == user_id)) # type: ignore\n if not user:\n return []\n user = cast(User, user)\n # Get servers accessible directly to user\n user_servers = list(user.accessible_mcp_servers)\n\n # TODO: Add group-based access once relationships are fully implemented\n # For now, just return direct user access\n return user_servers\n\n\ndef create_mcp_server__no_commit(\n owner_email: str,\n name: str,\n description: str | None,\n server_url: str,\n auth_type: MCPAuthenticationType | None,\n transport: MCPTransport | None,\n auth_performer: MCPAuthenticationPerformer | None,\n db_session: Session,\n admin_connection_config_id: int | None = None,\n) -> MCPServer:\n \"\"\"Create a new MCP server\"\"\"\n new_server = MCPServer(\n owner=owner_email,\n name=name,\n description=description,\n server_url=server_url,\n transport=transport,\n auth_type=auth_type,\n auth_performer=auth_performer,\n admin_connection_config_id=admin_connection_config_id,\n )\n db_session.add(new_server)\n db_session.flush() # Get the ID without committing\n return new_server\n\n\ndef update_mcp_server__no_commit(\n server_id: int,\n db_session: Session,\n name: str | None = None,\n description: str | None = None,\n server_url: str | None = None,\n auth_type: MCPAuthenticationType | None = None,\n admin_connection_config_id: int | None = None,\n auth_performer: MCPAuthenticationPerformer | None = None,\n transport: MCPTransport | None = None,\n status: MCPServerStatus | None = None,\n last_refreshed_at: datetime.datetime | None = None,\n) -> MCPServer:\n \"\"\"Update an existing MCP server\"\"\"\n server = get_mcp_server_by_id(server_id, db_session)\n\n if name is not None:\n server.name = name\n if description is not None:\n server.description = description\n if server_url is not None:\n server.server_url = server_url\n if auth_type is not None:\n server.auth_type = auth_type\n if admin_connection_config_id is not None:\n server.admin_connection_config_id = admin_connection_config_id\n if auth_performer is not None:\n server.auth_performer = auth_performer\n if transport is not None:\n server.transport = transport\n if status is not None:\n server.status = status\n if last_refreshed_at is not None:\n server.last_refreshed_at = last_refreshed_at\n\n db_session.flush() # Don't commit yet, let caller decide when to commit\n return server\n\n\ndef delete_mcp_server(server_id: int, db_session: Session) -> None:\n \"\"\"Delete an MCP server and all associated tools (via CASCADE)\"\"\"\n server = get_mcp_server_by_id(server_id, db_session)\n\n # Count tools that will be deleted\n tools_count = db_session.query(Tool).filter(Tool.mcp_server_id == server_id).count()\n logger.info(f\"Deleting MCP server {server_id} with {tools_count} associated tools\")\n\n db_session.delete(server)\n db_session.commit()\n\n logger.info(f\"Successfully deleted MCP server {server_id} and its tools\")\n\n\ndef get_all_mcp_tools_for_server(server_id: int, db_session: Session) -> list[Tool]:\n \"\"\"Get all MCP tools for a server\"\"\"\n return list(\n db_session.scalars(select(Tool).where(Tool.mcp_server_id == server_id)).all()\n )\n\n\ndef add_user_to_mcp_server(server_id: int, user_id: UUID, db_session: Session) -> None:\n \"\"\"Grant a user access to an MCP server\"\"\"\n server = get_mcp_server_by_id(server_id, db_session)\n user = db_session.scalar(select(User).where(User.id == user_id)) # type: ignore\n if not user:\n raise ValueError(\"User not found\")\n\n if user not in server.users:\n server.users.append(user)\n db_session.commit()\n\n\ndef remove_user_from_mcp_server(\n server_id: int, user_id: UUID, db_session: Session\n) -> None:\n \"\"\"Remove a user's access to an MCP server\"\"\"\n server = get_mcp_server_by_id(server_id, db_session)\n user = db_session.scalar(select(User).where(User.id == user_id)) # type: ignore\n if not user:\n raise ValueError(\"User not found\")\n\n if user in server.users:\n server.users.remove(user)\n db_session.commit()\n\n\n# MCPConnectionConfig operations\ndef extract_connection_data(\n config: MCPConnectionConfig | None, apply_mask: bool = False\n) -> MCPConnectionData:\n \"\"\"Extract MCPConnectionData from a connection config, with proper typing.\n\n This helper encapsulates the cast from the JSON column's dict[str, Any]\n to the typed MCPConnectionData structure.\n \"\"\"\n if config is None or config.config is None:\n return MCPConnectionData(headers={})\n if isinstance(config.config, SensitiveValue):\n return cast(MCPConnectionData, config.config.get_value(apply_mask=apply_mask))\n return cast(MCPConnectionData, config.config)\n\n\ndef get_connection_config_by_id(\n config_id: int, db_session: Session\n) -> MCPConnectionConfig:\n \"\"\"Get connection config by ID\"\"\"\n config = db_session.scalar(\n select(MCPConnectionConfig).where(MCPConnectionConfig.id == config_id)\n )\n if not config:\n raise ValueError(\"Connection config by specified id does not exist\")\n return config\n\n\ndef get_user_connection_config(\n server_id: int, user_email: str, db_session: Session\n) -> MCPConnectionConfig | None:\n \"\"\"Get a user's connection config for a specific MCP server\"\"\"\n return db_session.scalar(\n select(MCPConnectionConfig).where(\n and_(\n MCPConnectionConfig.mcp_server_id == server_id,\n MCPConnectionConfig.user_email == user_email,\n )\n )\n )\n\n\ndef get_user_connection_configs_for_server(\n server_id: int, db_session: Session\n) -> list[MCPConnectionConfig]:\n \"\"\"Get all user connection configs for a specific MCP server\"\"\"\n return list(\n db_session.scalars(\n select(MCPConnectionConfig).where(\n MCPConnectionConfig.mcp_server_id == server_id\n )\n ).all()\n )\n\n\ndef create_connection_config(\n config_data: MCPConnectionData,\n db_session: Session,\n mcp_server_id: int | None = None,\n user_email: str = \"\",\n) -> MCPConnectionConfig:\n \"\"\"Create a new connection config\"\"\"\n new_config = MCPConnectionConfig(\n mcp_server_id=mcp_server_id,\n user_email=user_email,\n config=config_data,\n )\n db_session.add(new_config)\n db_session.flush() # Don't commit yet, let caller decide when to commit\n return new_config\n\n\ndef update_connection_config(\n config_id: int,\n db_session: Session,\n config_data: MCPConnectionData | None = None,\n) -> MCPConnectionConfig:\n \"\"\"Update an existing connection config\"\"\"\n config = get_connection_config_by_id(config_id, db_session)\n\n if config_data is not None:\n config.config = config_data # type: ignore[assignment]\n # Force SQLAlchemy to detect the change by marking the field as modified\n flag_modified(config, \"config\")\n\n db_session.commit()\n return config\n\n\ndef upsert_user_connection_config(\n server_id: int,\n user_email: str,\n config_data: MCPConnectionData,\n db_session: Session,\n) -> MCPConnectionConfig:\n \"\"\"Create or update a user's connection config for an MCP server\"\"\"\n existing_config = get_user_connection_config(server_id, user_email, db_session)\n\n if existing_config:\n existing_config.config = config_data # type: ignore[assignment]\n db_session.flush() # Don't commit yet, let caller decide when to commit\n return existing_config\n else:\n return create_connection_config(\n config_data=config_data,\n mcp_server_id=server_id,\n user_email=user_email,\n db_session=db_session,\n )\n\n\n# TODO: do this in one db call\ndef get_server_auth_template(\n server_id: int, db_session: Session\n) -> MCPConnectionConfig | None:\n \"\"\"Get the authentication template for a server (from the admin connection config)\"\"\"\n server = get_mcp_server_by_id(server_id, db_session)\n if not server.admin_connection_config_id:\n return None\n\n if server.auth_performer == MCPAuthenticationPerformer.ADMIN:\n return None # admin server implies no template\n return server.admin_connection_config\n\n\ndef delete_connection_config(config_id: int, db_session: Session) -> None:\n \"\"\"Delete a connection config\"\"\"\n config = get_connection_config_by_id(config_id, db_session)\n db_session.delete(config)\n db_session.flush() # Don't commit yet, let caller decide when to commit\n\n\ndef delete_user_connection_configs_for_server(\n server_id: int, user_email: str, db_session: Session\n) -> None:\n \"\"\"Delete all connection configs for a user on a specific server\"\"\"\n configs = db_session.scalars(\n select(MCPConnectionConfig).where(\n and_(\n MCPConnectionConfig.mcp_server_id == server_id,\n MCPConnectionConfig.user_email == user_email,\n )\n )\n ).all()\n\n for config in configs:\n db_session.delete(config)\n\n db_session.commit()\n\n\ndef delete_all_user_connection_configs_for_server_no_commit(\n server_id: int, db_session: Session\n) -> None:\n \"\"\"Delete all user connection configs for a specific MCP server\"\"\"\n db_session.execute(\n delete(MCPConnectionConfig).where(\n MCPConnectionConfig.mcp_server_id == server_id\n )\n )\n db_session.flush() # Don't commit yet, let caller decide when to commit\n" }, { "path": "backend/onyx/db/memory.py", "content": "from uuid import UUID\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import Memory\nfrom onyx.db.models import User\n\nMAX_MEMORIES_PER_USER = 10\n\n\nclass UserInfo(BaseModel):\n name: str | None = None\n role: str | None = None\n email: str | None = None\n\n def to_dict(self) -> dict:\n return {\n \"name\": self.name,\n \"role\": self.role,\n \"email\": self.email,\n }\n\n\nclass UserMemoryContext(BaseModel):\n model_config = ConfigDict(frozen=True)\n\n user_id: UUID | None = None\n user_info: UserInfo\n user_preferences: str | None = None\n memories: tuple[str, ...] = ()\n\n def without_memories(self) -> \"UserMemoryContext\":\n \"\"\"Return a copy with memories cleared but user info/preferences intact.\"\"\"\n return UserMemoryContext(\n user_id=self.user_id,\n user_info=self.user_info,\n user_preferences=self.user_preferences,\n memories=(),\n )\n\n def as_formatted_list(self) -> list[str]:\n \"\"\"Returns combined list of user info, preferences, and memories.\"\"\"\n result = []\n if self.user_info.name:\n result.append(f\"User's name: {self.user_info.name}\")\n if self.user_info.role:\n result.append(f\"User's role: {self.user_info.role}\")\n if self.user_info.email:\n result.append(f\"User's email: {self.user_info.email}\")\n if self.user_preferences:\n result.append(f\"User preferences: {self.user_preferences}\")\n result.extend(self.memories)\n return result\n\n\ndef get_memories(user: User, db_session: Session) -> UserMemoryContext:\n user_info = UserInfo(\n name=user.personal_name,\n role=user.personal_role,\n email=user.email,\n )\n\n user_preferences = None\n if user.user_preferences:\n user_preferences = user.user_preferences\n\n memory_rows = db_session.scalars(\n select(Memory).where(Memory.user_id == user.id).order_by(Memory.id.asc())\n ).all()\n memories = tuple(memory.memory_text for memory in memory_rows if memory.memory_text)\n\n return UserMemoryContext(\n user_id=user.id,\n user_info=user_info,\n user_preferences=user_preferences,\n memories=memories,\n )\n\n\ndef add_memory(\n user_id: UUID,\n memory_text: str,\n db_session: Session,\n) -> Memory:\n \"\"\"Insert a new Memory row for the given user.\n\n If the user already has MAX_MEMORIES_PER_USER memories, the oldest\n one (lowest id) is deleted before inserting the new one.\n \"\"\"\n existing = db_session.scalars(\n select(Memory).where(Memory.user_id == user_id).order_by(Memory.id.asc())\n ).all()\n\n if len(existing) >= MAX_MEMORIES_PER_USER:\n db_session.delete(existing[0])\n\n memory = Memory(\n user_id=user_id,\n memory_text=memory_text,\n )\n db_session.add(memory)\n db_session.commit()\n return memory\n\n\ndef update_memory_at_index(\n user_id: UUID,\n index: int,\n new_text: str,\n db_session: Session,\n) -> Memory | None:\n \"\"\"Update the memory at the given 0-based index (ordered by id ASC, matching get_memories()).\n\n Returns the updated Memory row, or None if the index is out of range.\n \"\"\"\n memory_rows = db_session.scalars(\n select(Memory).where(Memory.user_id == user_id).order_by(Memory.id.asc())\n ).all()\n\n if index < 0 or index >= len(memory_rows):\n return None\n\n memory = memory_rows[index]\n memory.memory_text = new_text\n db_session.commit()\n return memory\n" }, { "path": "backend/onyx/db/models.py", "content": "import datetime\nimport json\nfrom typing import Any\nfrom typing import Literal\nfrom typing import NotRequired\nfrom uuid import uuid4\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import validates\n\nfrom typing_extensions import TypedDict # noreorder\nfrom uuid import UUID\nfrom pydantic import ValidationError\n\nfrom sqlalchemy.dialects.postgresql import JSONB as PGJSONB\nfrom sqlalchemy.dialects.postgresql import UUID as PGUUID\n\nfrom fastapi_users_db_sqlalchemy import SQLAlchemyBaseOAuthAccountTableUUID\nfrom fastapi_users_db_sqlalchemy import SQLAlchemyBaseUserTableUUID\nfrom fastapi_users_db_sqlalchemy.access_token import SQLAlchemyBaseAccessTokenTableUUID\nfrom fastapi_users_db_sqlalchemy.generics import TIMESTAMPAware\nfrom sqlalchemy import Boolean\nfrom sqlalchemy import DateTime\nfrom sqlalchemy import desc\nfrom sqlalchemy import Enum\nfrom sqlalchemy import Float\nfrom sqlalchemy import ForeignKey\nfrom sqlalchemy import ForeignKeyConstraint\nfrom sqlalchemy import func\nfrom sqlalchemy import Index\nfrom sqlalchemy import Integer\nfrom sqlalchemy import BigInteger\n\nfrom sqlalchemy import Sequence\nfrom sqlalchemy import String\nfrom sqlalchemy import Text\nfrom sqlalchemy import text\nfrom sqlalchemy import UniqueConstraint\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy import event\nfrom sqlalchemy.engine.interfaces import Dialect\nfrom sqlalchemy.orm import DeclarativeBase\nfrom sqlalchemy.orm import Mapped\nfrom sqlalchemy.orm import Mapper\nfrom sqlalchemy.orm import mapped_column\nfrom sqlalchemy.orm import relationship\nfrom sqlalchemy.types import LargeBinary\nfrom sqlalchemy.types import TypeDecorator\nfrom sqlalchemy import PrimaryKeyConstraint\n\nfrom onyx.db.enums import AccountType\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import (\n ANONYMOUS_USER_UUID,\n DEFAULT_BOOST,\n FederatedConnectorSource,\n MilestoneRecordType,\n)\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.enums import (\n AccessType,\n ArtifactType,\n BuildSessionStatus,\n EmbeddingPrecision,\n HierarchyNodeType,\n HookFailStrategy,\n HookPoint,\n IndexingMode,\n OpenSearchDocumentMigrationStatus,\n OpenSearchTenantMigrationStatus,\n ProcessingMode,\n SandboxStatus,\n SyncType,\n SyncStatus,\n MCPAuthenticationType,\n UserFileStatus,\n MCPAuthenticationPerformer,\n MCPTransport,\n MCPServerStatus,\n Permission,\n GrantSource,\n LLMModelFlowType,\n ThemePreference,\n DefaultAppMode,\n SwitchoverType,\n SharingScope,\n)\nfrom onyx.configs.constants import NotificationType\nfrom onyx.configs.constants import SearchFeedbackType\nfrom onyx.configs.constants import TokenRateLimitScope\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import ChatSessionSharedStatus\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.enums import TaskStatus\nfrom onyx.db.pydantic_type import PydanticListType, PydanticType\nfrom onyx.kg.models import KGEntityTypeAttributes\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.special_types import JSON_ro\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.llm.override_models import PromptOverride\nfrom onyx.kg.models import KGStage\nfrom onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig\nfrom onyx.utils.encryption import decrypt_bytes_to_string\nfrom onyx.utils.encryption import encrypt_string_to_bytes\nfrom onyx.utils.sensitive import SensitiveValue\nfrom onyx.utils.headers import HeaderItemDict\nfrom shared_configs.enums import EmbeddingProvider\n\n# TODO: After anonymous user migration has been deployed, make user_id columns NOT NULL\n# and update Mapped[User | None] relationships to Mapped[User] where needed.\n\n\nlogger = setup_logger()\n\nPROMPT_LENGTH = 5_000_000\n\n\nclass Base(DeclarativeBase):\n __abstract__ = True\n\n\nclass _EncryptedBase(TypeDecorator):\n \"\"\"Base for encrypted column types that wrap values in SensitiveValue.\"\"\"\n\n impl = LargeBinary\n cache_ok = True\n _is_json: bool = False\n\n def wrap_raw(self, value: Any) -> SensitiveValue:\n \"\"\"Encrypt a raw value and wrap it in SensitiveValue.\n\n Called by the attribute set event so the Python-side type is always\n SensitiveValue, regardless of whether the value was loaded from the DB\n or assigned in application code.\n \"\"\"\n if self._is_json:\n if not isinstance(value, dict):\n raise TypeError(\n f\"EncryptedJson column expected dict, got {type(value).__name__}\"\n )\n raw_str = json.dumps(value)\n else:\n if not isinstance(value, str):\n raise TypeError(\n f\"EncryptedString column expected str, got {type(value).__name__}\"\n )\n raw_str = value\n return SensitiveValue(\n encrypted_bytes=encrypt_string_to_bytes(raw_str),\n decrypt_fn=decrypt_bytes_to_string,\n is_json=self._is_json,\n )\n\n def compare_values(self, x: Any, y: Any) -> bool:\n if x is None or y is None:\n return x == y\n if isinstance(x, SensitiveValue):\n x = x.get_value(apply_mask=False)\n if isinstance(y, SensitiveValue):\n y = y.get_value(apply_mask=False)\n return x == y\n\n\nclass EncryptedString(_EncryptedBase):\n # Must redeclare cache_ok in this child class since we explicitly redeclare _is_json\n cache_ok = True\n _is_json: bool = False\n\n def process_bind_param(\n self,\n value: str | SensitiveValue[str] | None,\n dialect: Dialect, # noqa: ARG002\n ) -> bytes | None:\n if value is not None:\n # Handle both raw strings and SensitiveValue wrappers\n if isinstance(value, SensitiveValue):\n # Get raw value for storage\n value = value.get_value(apply_mask=False)\n return encrypt_string_to_bytes(value)\n return value\n\n def process_result_value(\n self,\n value: bytes | None,\n dialect: Dialect, # noqa: ARG002\n ) -> SensitiveValue[str] | None:\n if value is not None:\n return SensitiveValue(\n encrypted_bytes=value,\n decrypt_fn=decrypt_bytes_to_string,\n is_json=False,\n )\n return None\n\n\nclass EncryptedJson(_EncryptedBase):\n cache_ok = True\n _is_json: bool = True\n\n def process_bind_param(\n self,\n value: dict[str, Any] | SensitiveValue[dict[str, Any]] | None,\n dialect: Dialect, # noqa: ARG002\n ) -> bytes | None:\n if value is not None:\n if isinstance(value, SensitiveValue):\n value = value.get_value(apply_mask=False)\n json_str = json.dumps(value)\n return encrypt_string_to_bytes(json_str)\n return value\n\n def process_result_value(\n self,\n value: bytes | None,\n dialect: Dialect, # noqa: ARG002\n ) -> SensitiveValue[dict[str, Any]] | None:\n if value is not None:\n return SensitiveValue(\n encrypted_bytes=value,\n decrypt_fn=decrypt_bytes_to_string,\n is_json=True,\n )\n return None\n\n\n_REGISTERED_ATTRS: set[str] = set()\n\n\n@event.listens_for(Mapper, \"mapper_configured\")\ndef _register_sensitive_value_set_events(\n mapper: Mapper,\n class_: type,\n) -> None:\n \"\"\"Auto-wrap raw values in SensitiveValue when assigned to encrypted columns.\"\"\"\n for prop in mapper.column_attrs:\n for col in prop.columns:\n if isinstance(col.type, _EncryptedBase):\n col_type = col.type\n attr = getattr(class_, prop.key)\n\n # Guard against double-registration (e.g. if mapper is\n # re-configured in test setups)\n attr_key = f\"{class_.__qualname__}.{prop.key}\"\n if attr_key in _REGISTERED_ATTRS:\n continue\n _REGISTERED_ATTRS.add(attr_key)\n\n @event.listens_for(attr, \"set\", retval=True)\n def _wrap_value(\n target: Any, # noqa: ARG001\n value: Any,\n oldvalue: Any, # noqa: ARG001\n initiator: Any, # noqa: ARG001\n _col_type: _EncryptedBase = col_type,\n ) -> Any:\n if value is not None and not isinstance(value, SensitiveValue):\n return _col_type.wrap_raw(value)\n return value\n\n\nclass NullFilteredString(TypeDecorator):\n impl = String\n # This type's behavior is fully deterministic and doesn't depend on any external factors.\n cache_ok = True\n\n def process_bind_param(\n self,\n value: str | None,\n dialect: Dialect, # noqa: ARG002\n ) -> str | None:\n if value is not None and \"\\x00\" in value:\n logger.warning(f\"NUL characters found in value: {value}\")\n return value.replace(\"\\x00\", \"\")\n return value\n\n def process_result_value(\n self,\n value: str | None,\n dialect: Dialect, # noqa: ARG002\n ) -> str | None:\n return value\n\n\n\"\"\"\nAuth/Authz (users, permissions, access) Tables\n\"\"\"\n\n\nclass OAuthAccount(SQLAlchemyBaseOAuthAccountTableUUID, Base):\n # even an almost empty token from keycloak will not fit the default 1024 bytes\n access_token: Mapped[str] = mapped_column(Text, nullable=False) # type: ignore\n refresh_token: Mapped[str] = mapped_column(Text, nullable=False) # type: ignore\n\n\nclass User(SQLAlchemyBaseUserTableUUID, Base):\n oauth_accounts: Mapped[list[OAuthAccount]] = relationship(\n \"OAuthAccount\", lazy=\"joined\", cascade=\"all, delete-orphan\"\n )\n role: Mapped[UserRole] = mapped_column(\n Enum(UserRole, native_enum=False, default=UserRole.BASIC)\n )\n account_type: Mapped[AccountType] = mapped_column(\n Enum(AccountType, native_enum=False),\n nullable=False,\n default=AccountType.STANDARD,\n server_default=\"STANDARD\",\n )\n\n \"\"\"\n Preferences probably should be in a separate table at some point, but for now\n putting here for simpicity\n \"\"\"\n\n temperature_override_enabled: Mapped[bool | None] = mapped_column(\n Boolean, default=None\n )\n auto_scroll: Mapped[bool | None] = mapped_column(Boolean, default=None)\n shortcut_enabled: Mapped[bool] = mapped_column(Boolean, default=False)\n theme_preference: Mapped[ThemePreference | None] = mapped_column(\n Enum(ThemePreference, native_enum=False),\n nullable=True,\n default=None,\n )\n chat_background: Mapped[str | None] = mapped_column(String, nullable=True)\n default_app_mode: Mapped[DefaultAppMode] = mapped_column(\n Enum(DefaultAppMode, native_enum=False),\n nullable=False,\n default=DefaultAppMode.CHAT,\n )\n # personalization fields are exposed via the chat user settings \"Personalization\" tab\n personal_name: Mapped[str | None] = mapped_column(String, nullable=True)\n personal_role: Mapped[str | None] = mapped_column(String, nullable=True)\n use_memories: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n enable_memory_tool: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=True\n )\n user_preferences: Mapped[str | None] = mapped_column(Text, nullable=True)\n\n chosen_assistants: Mapped[list[int] | None] = mapped_column(\n postgresql.JSONB(), nullable=True, default=None\n )\n visible_assistants: Mapped[list[int]] = mapped_column(\n postgresql.JSONB(), nullable=False, default=[]\n )\n hidden_assistants: Mapped[list[int]] = mapped_column(\n postgresql.JSONB(), nullable=False, default=[]\n )\n\n pinned_assistants: Mapped[list[int] | None] = mapped_column(\n postgresql.JSONB(), nullable=True, default=None\n )\n\n effective_permissions: Mapped[list[str]] = mapped_column(\n postgresql.JSONB(),\n nullable=False,\n default=list,\n server_default=text(\"'[]'::jsonb\"),\n )\n\n oidc_expiry: Mapped[datetime.datetime] = mapped_column(\n TIMESTAMPAware(timezone=True), nullable=True\n )\n\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n default_model: Mapped[str] = mapped_column(Text, nullable=True)\n # organized in typical structured fashion\n # formatted as `displayName__provider__modelName`\n\n # Voice preferences\n voice_auto_send: Mapped[bool] = mapped_column(Boolean, default=False)\n voice_auto_playback: Mapped[bool] = mapped_column(Boolean, default=False)\n voice_playback_speed: Mapped[float] = mapped_column(Float, default=1.0)\n\n # relationships\n credentials: Mapped[list[\"Credential\"]] = relationship(\n \"Credential\", back_populates=\"user\"\n )\n chat_sessions: Mapped[list[\"ChatSession\"]] = relationship(\n \"ChatSession\", back_populates=\"user\"\n )\n\n input_prompts: Mapped[list[\"InputPrompt\"]] = relationship(\n \"InputPrompt\", back_populates=\"user\"\n )\n # Personas owned by this user\n personas: Mapped[list[\"Persona\"]] = relationship(\"Persona\", back_populates=\"user\")\n # Custom tools created by this user\n custom_tools: Mapped[list[\"Tool\"]] = relationship(\"Tool\", back_populates=\"user\")\n # Notifications for the UI\n notifications: Mapped[list[\"Notification\"]] = relationship(\n \"Notification\", back_populates=\"user\"\n )\n cc_pairs: Mapped[list[\"ConnectorCredentialPair\"]] = relationship(\n \"ConnectorCredentialPair\",\n back_populates=\"creator\",\n primaryjoin=\"User.id == foreign(ConnectorCredentialPair.creator_id)\",\n )\n projects: Mapped[list[\"UserProject\"]] = relationship(\n \"UserProject\", back_populates=\"user\"\n )\n files: Mapped[list[\"UserFile\"]] = relationship(\"UserFile\", back_populates=\"user\")\n # MCP servers accessible to this user\n accessible_mcp_servers: Mapped[list[\"MCPServer\"]] = relationship(\n \"MCPServer\", secondary=\"mcp_server__user\", back_populates=\"users\"\n )\n memories: Mapped[list[\"Memory\"]] = relationship(\n \"Memory\",\n back_populates=\"user\",\n cascade=\"all, delete-orphan\",\n order_by=\"desc(Memory.id)\",\n )\n oauth_user_tokens: Mapped[list[\"OAuthUserToken\"]] = relationship(\n \"OAuthUserToken\",\n back_populates=\"user\",\n cascade=\"all, delete-orphan\",\n )\n\n @validates(\"email\")\n def validate_email(self, key: str, value: str) -> str: # noqa: ARG002\n return value.lower() if value else value\n\n @property\n def password_configured(self) -> bool:\n \"\"\"\n Returns True if the user has at least one OAuth (or OIDC) account.\n \"\"\"\n return not bool(self.oauth_accounts)\n\n @property\n def is_anonymous(self) -> bool:\n \"\"\"Returns True if this is the anonymous user.\"\"\"\n return str(self.id) == ANONYMOUS_USER_UUID\n\n\nclass AccessToken(SQLAlchemyBaseAccessTokenTableUUID, Base):\n pass\n\n\nclass Memory(Base):\n __tablename__ = \"memory\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)\n user_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=False\n )\n memory_text: Mapped[str] = mapped_column(Text, nullable=False)\n conversation_id: Mapped[UUID | None] = mapped_column(\n PGUUID(as_uuid=True), nullable=True\n )\n message_id: Mapped[int | None] = mapped_column(Integer, nullable=True)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n user: Mapped[\"User\"] = relationship(\"User\", back_populates=\"memories\")\n\n\nclass ApiKey(Base):\n __tablename__ = \"api_key\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str | None] = mapped_column(String, nullable=True)\n hashed_api_key: Mapped[str] = mapped_column(String, unique=True)\n api_key_display: Mapped[str] = mapped_column(String, unique=True)\n # the ID of the \"user\" who represents the access credentials for the API key\n user_id: Mapped[UUID] = mapped_column(ForeignKey(\"user.id\"), nullable=False)\n # the ID of the user who owns the key\n owner_id: Mapped[UUID | None] = mapped_column(ForeignKey(\"user.id\"), nullable=True)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n # Add this relationship to access the User object via user_id\n user: Mapped[\"User\"] = relationship(\"User\", foreign_keys=[user_id])\n\n\nclass PersonalAccessToken(Base):\n __tablename__ = \"personal_access_token\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, nullable=False) # User-provided label\n hashed_token: Mapped[str] = mapped_column(\n String(64), unique=True, nullable=False\n ) # SHA256 = 64 hex chars\n token_display: Mapped[str] = mapped_column(String, nullable=False)\n\n user_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=False\n )\n\n expires_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True, index=True\n ) # NULL = no expiration. Revocation sets this to NOW() for immediate expiry.\n\n # Audit fields\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n last_used_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n is_revoked: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"false\"), nullable=False\n ) # True if user explicitly revoked (vs naturally expired)\n\n user: Mapped[\"User\"] = relationship(\"User\", foreign_keys=[user_id])\n\n # Indexes for performance\n __table_args__ = (\n Index(\n \"ix_pat_user_created\", user_id, created_at.desc()\n ), # Fast user token listing\n )\n\n\nclass Notification(Base):\n __tablename__ = \"notification\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n notif_type: Mapped[NotificationType] = mapped_column(\n Enum(NotificationType, native_enum=False)\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n dismissed: Mapped[bool] = mapped_column(Boolean, default=False)\n last_shown: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))\n first_shown: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))\n title: Mapped[str] = mapped_column(String)\n description: Mapped[str | None] = mapped_column(String, nullable=True)\n\n user: Mapped[User] = relationship(\"User\", back_populates=\"notifications\")\n additional_data: Mapped[dict | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Unique constraint ix_notification_user_type_data on (user_id, notif_type, additional_data)\n # ensures notification deduplication for batch inserts. Defined in migration 8405ca81cc83.\n __table_args__ = (\n Index(\n \"ix_notification_user_sort\",\n \"user_id\",\n \"dismissed\",\n desc(\"first_shown\"),\n ),\n )\n\n\n\"\"\"\nAssociation Tables\nNOTE: must be at the top since they are referenced by other tables\n\"\"\"\n\n\nclass Persona__DocumentSet(Base):\n __tablename__ = \"persona__document_set\"\n\n persona_id: Mapped[int] = mapped_column(ForeignKey(\"persona.id\"), primary_key=True)\n document_set_id: Mapped[int] = mapped_column(\n ForeignKey(\"document_set.id\"), primary_key=True\n )\n\n\nclass Persona__User(Base):\n __tablename__ = \"persona__user\"\n\n persona_id: Mapped[int] = mapped_column(ForeignKey(\"persona.id\"), primary_key=True)\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), primary_key=True, nullable=True\n )\n\n\nclass DocumentSet__User(Base):\n __tablename__ = \"document_set__user\"\n\n document_set_id: Mapped[int] = mapped_column(\n ForeignKey(\"document_set.id\"), primary_key=True\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), primary_key=True, nullable=True\n )\n\n\nclass DocumentSet__ConnectorCredentialPair(Base):\n __tablename__ = \"document_set__connector_credential_pair\"\n\n document_set_id: Mapped[int] = mapped_column(\n ForeignKey(\"document_set.id\"), primary_key=True\n )\n connector_credential_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"), primary_key=True\n )\n # if `True`, then is part of the current state of the document set\n # if `False`, then is a part of the prior state of the document set\n # rows with `is_current=False` should be deleted when the document\n # set is updated and should not exist for a given document set if\n # `DocumentSet.is_up_to_date == True`\n is_current: Mapped[bool] = mapped_column(\n Boolean,\n nullable=False,\n default=True,\n primary_key=True,\n )\n\n document_set: Mapped[\"DocumentSet\"] = relationship(\"DocumentSet\")\n\n\nclass ChatMessage__SearchDoc(Base):\n __tablename__ = \"chat_message__search_doc\"\n\n chat_message_id: Mapped[int] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n search_doc_id: Mapped[int] = mapped_column(\n ForeignKey(\"search_doc.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass ToolCall__SearchDoc(Base):\n __tablename__ = \"tool_call__search_doc\"\n\n tool_call_id: Mapped[int] = mapped_column(\n ForeignKey(\"tool_call.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n search_doc_id: Mapped[int] = mapped_column(\n ForeignKey(\"search_doc.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass Document__Tag(Base):\n __tablename__ = \"document__tag\"\n\n document_id: Mapped[str] = mapped_column(\n ForeignKey(\"document.id\"), primary_key=True\n )\n tag_id: Mapped[int] = mapped_column(\n ForeignKey(\"tag.id\"), primary_key=True, index=True\n )\n\n\nclass Persona__Tool(Base):\n \"\"\"An entry in this table represents a tool that is **available** to a persona.\n It does NOT necessarily mean that the tool is actually usable to the persona.\n\n For example, a persona may have the image generation tool attached to it, even though\n the image generation tool is not set up / enabled. In this case, the tool should not\n show up in the UI for the persona + it should not be usable by the persona in chat.\n \"\"\"\n\n __tablename__ = \"persona__tool\"\n\n persona_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n tool_id: Mapped[int] = mapped_column(\n ForeignKey(\"tool.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass StandardAnswer__StandardAnswerCategory(Base):\n __tablename__ = \"standard_answer__standard_answer_category\"\n\n standard_answer_id: Mapped[int] = mapped_column(\n ForeignKey(\"standard_answer.id\"), primary_key=True\n )\n standard_answer_category_id: Mapped[int] = mapped_column(\n ForeignKey(\"standard_answer_category.id\"), primary_key=True\n )\n\n\nclass SlackChannelConfig__StandardAnswerCategory(Base):\n __tablename__ = \"slack_channel_config__standard_answer_category\"\n\n slack_channel_config_id: Mapped[int] = mapped_column(\n ForeignKey(\"slack_channel_config.id\"), primary_key=True\n )\n standard_answer_category_id: Mapped[int] = mapped_column(\n ForeignKey(\"standard_answer_category.id\"), primary_key=True\n )\n\n\nclass ChatMessage__StandardAnswer(Base):\n __tablename__ = \"chat_message__standard_answer\"\n\n chat_message_id: Mapped[int] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n standard_answer_id: Mapped[int] = mapped_column(\n ForeignKey(\"standard_answer.id\"), primary_key=True\n )\n\n\n\"\"\"\nDocuments/Indexing Tables\n\"\"\"\n\n\nclass ConnectorCredentialPair(Base):\n \"\"\"Connectors and Credentials can have a many-to-many relationship\n I.e. A Confluence Connector may have multiple admin users who can run it with their own credentials\n I.e. An admin user may use the same credential to index multiple Confluence Spaces\n \"\"\"\n\n __tablename__ = \"connector_credential_pair\"\n # NOTE: this `id` column has to use `Sequence` instead of `autoincrement=True`\n # due to some SQLAlchemy quirks + this not being a primary key column\n id: Mapped[int] = mapped_column(\n Integer,\n Sequence(\"connector_credential_pair_id_seq\"),\n unique=True,\n nullable=False,\n )\n name: Mapped[str] = mapped_column(String, nullable=False)\n status: Mapped[ConnectorCredentialPairStatus] = mapped_column(\n Enum(ConnectorCredentialPairStatus, native_enum=False), nullable=False\n )\n # this is separate from the `status` above, since a connector can be `INITIAL_INDEXING`, `ACTIVE`,\n # or `PAUSED` and still be in a repeated error state.\n in_repeated_error_state: Mapped[bool] = mapped_column(Boolean, default=False)\n connector_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector.id\"), primary_key=True\n )\n\n deletion_failure_message: Mapped[str | None] = mapped_column(String, nullable=True)\n\n credential_id: Mapped[int] = mapped_column(\n ForeignKey(\"credential.id\"), primary_key=True\n )\n # controls whether the documents indexed by this CC pair are visible to all\n # or if they are only visible to those with that are given explicit access\n # (e.g. via owning the credential or being a part of a group that is given access)\n access_type: Mapped[AccessType] = mapped_column(\n Enum(AccessType, native_enum=False), nullable=False\n )\n\n # special info needed for the auto-sync feature. The exact structure depends on the\n\n # source type (defined in the connector's `source` field)\n # E.g. for google_drive perm sync:\n # {\"customer_id\": \"123567\", \"company_domain\": \"@onyx.app\"}\n auto_sync_options: Mapped[dict[str, Any] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n last_time_perm_sync: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n last_time_external_group_sync: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n # Time finished, not used for calculating backend jobs which uses time started (created)\n last_successful_index_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n\n # last successful prune\n last_pruned: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True, index=True\n )\n\n # last successful hierarchy fetch\n last_time_hierarchy_fetch: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n total_docs_indexed: Mapped[int] = mapped_column(Integer, default=0)\n\n indexing_trigger: Mapped[IndexingMode | None] = mapped_column(\n Enum(IndexingMode, native_enum=False), nullable=True\n )\n\n # Determines how documents are processed after fetching:\n # REGULAR: Full pipeline (chunk → embed → Vespa)\n # FILE_SYSTEM: Write to file system only (for CLI agent sandbox)\n processing_mode: Mapped[ProcessingMode] = mapped_column(\n Enum(ProcessingMode, native_enum=False),\n nullable=False,\n default=ProcessingMode.REGULAR,\n server_default=\"REGULAR\",\n )\n\n connector: Mapped[\"Connector\"] = relationship(\n \"Connector\", back_populates=\"credentials\"\n )\n credential: Mapped[\"Credential\"] = relationship(\n \"Credential\", back_populates=\"connectors\"\n )\n document_sets: Mapped[list[\"DocumentSet\"]] = relationship(\n \"DocumentSet\",\n secondary=DocumentSet__ConnectorCredentialPair.__table__,\n primaryjoin=(\n (DocumentSet__ConnectorCredentialPair.connector_credential_pair_id == id)\n & (DocumentSet__ConnectorCredentialPair.is_current.is_(True))\n ),\n back_populates=\"connector_credential_pairs\",\n overlaps=\"document_set\",\n )\n index_attempts: Mapped[list[\"IndexAttempt\"]] = relationship(\n \"IndexAttempt\", back_populates=\"connector_credential_pair\"\n )\n\n # the user id of the user that created this cc pair\n creator_id: Mapped[UUID | None] = mapped_column(nullable=True)\n creator: Mapped[\"User\"] = relationship(\n \"User\",\n back_populates=\"cc_pairs\",\n primaryjoin=\"foreign(ConnectorCredentialPair.creator_id) == remote(User.id)\",\n )\n\n background_errors: Mapped[list[\"BackgroundError\"]] = relationship(\n \"BackgroundError\", back_populates=\"cc_pair\", cascade=\"all, delete-orphan\"\n )\n\n\nclass HierarchyNode(Base):\n \"\"\"\n Represents a structural node in a connected source's hierarchy.\n Examples: folders, drives, spaces, projects, channels.\n\n Stores hierarchy structure WITH permission information, using the same\n permission model as Documents (external_user_emails, external_user_group_ids,\n is_public). This enables user-scoped hierarchy browsing in the UI.\n\n Some hierarchy nodes (e.g., Confluence pages) can also be documents.\n In these cases, `document_id` will be set.\n \"\"\"\n\n __tablename__ = \"hierarchy_node\"\n\n # Primary key - Integer for simplicity\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n\n # Raw identifier from the source system\n # e.g., \"1h7uWUR2BYZjtMfEXFt43tauj-Gp36DTPtwnsNuA665I\" for Google Drive\n # For SOURCE nodes, this is the source name (e.g., \"google_drive\")\n raw_node_id: Mapped[str] = mapped_column(String, nullable=False)\n\n # Human-readable name for display\n # e.g., \"Engineering\", \"Q4 Planning\", \"Google Drive\"\n display_name: Mapped[str] = mapped_column(String, nullable=False)\n\n # Link to view this node in the source system\n link: Mapped[str | None] = mapped_column(NullFilteredString, nullable=True)\n\n # Source type (google_drive, confluence, etc.)\n source: Mapped[DocumentSource] = mapped_column(\n Enum(DocumentSource, native_enum=False), nullable=False\n )\n\n # What kind of structural node this is\n node_type: Mapped[HierarchyNodeType] = mapped_column(\n Enum(HierarchyNodeType, native_enum=False), nullable=False\n )\n\n # ============= PERMISSION FIELDS (same pattern as Document) =============\n # Email addresses of external users with access to this node in the source system\n external_user_emails: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n # External group IDs with access (prefixed by source type)\n external_user_group_ids: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n # Whether this node is publicly accessible (org-wide or world-public)\n # SOURCE nodes are always public. Other nodes get this from source permissions.\n is_public: Mapped[bool] = mapped_column(Boolean, default=False)\n # ==========================================================================\n\n # Foreign keys\n # For hierarchy nodes that are also documents (e.g., Confluence pages)\n # SET NULL when document is deleted - node can exist without its document\n document_id: Mapped[str | None] = mapped_column(\n ForeignKey(\"document.id\", ondelete=\"SET NULL\"), nullable=True\n )\n\n # Self-referential FK for tree structure\n # SET NULL when parent is deleted - orphan children for cleanup via pruning\n parent_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"hierarchy_node.id\", ondelete=\"SET NULL\"), nullable=True, index=True\n )\n\n # Relationships\n document: Mapped[\"Document | None\"] = relationship(\n \"Document\", back_populates=\"hierarchy_node\", foreign_keys=[document_id]\n )\n parent: Mapped[\"HierarchyNode | None\"] = relationship(\n \"HierarchyNode\", remote_side=[id], back_populates=\"children\"\n )\n children: Mapped[list[\"HierarchyNode\"]] = relationship(\n \"HierarchyNode\", back_populates=\"parent\", passive_deletes=True\n )\n child_documents: Mapped[list[\"Document\"]] = relationship(\n \"Document\",\n back_populates=\"parent_hierarchy_node\",\n foreign_keys=\"Document.parent_hierarchy_node_id\",\n passive_deletes=True,\n )\n # Personas that have this hierarchy node attached for scoped search\n personas: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=\"persona__hierarchy_node\",\n back_populates=\"hierarchy_nodes\",\n viewonly=True,\n )\n\n __table_args__ = (\n # Unique constraint: same raw_node_id + source should not exist twice\n UniqueConstraint(\n \"raw_node_id\", \"source\", name=\"uq_hierarchy_node_raw_id_source\"\n ),\n Index(\"ix_hierarchy_node_source_type\", source, node_type),\n )\n\n\nclass Document(Base):\n __tablename__ = \"document\"\n # NOTE: if more sensitive data is added here for display, make sure to add user/group permission\n\n # this should correspond to the ID of the document\n # (as is passed around in Onyx)\n id: Mapped[str] = mapped_column(NullFilteredString, primary_key=True)\n from_ingestion_api: Mapped[bool] = mapped_column(\n Boolean, default=False, nullable=True\n )\n # 0 for neutral, positive for mostly endorse, negative for mostly reject\n boost: Mapped[int] = mapped_column(Integer, default=DEFAULT_BOOST)\n hidden: Mapped[bool] = mapped_column(Boolean, default=False)\n semantic_id: Mapped[str] = mapped_column(NullFilteredString)\n # First Section's link\n link: Mapped[str | None] = mapped_column(NullFilteredString, nullable=True)\n\n # The updated time is also used as a measure of the last successful state of the doc\n # pulled from the source (to help skip reindexing already updated docs in case of\n # connector retries)\n # TODO: rename this column because it conflates the time of the source doc\n # with the local last modified time of the doc and any associated metadata\n # it should just be the server timestamp of the source doc\n doc_updated_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n # Number of chunks in the document (in Vespa)\n # Only null for documents indexed prior to this change\n chunk_count: Mapped[int | None] = mapped_column(Integer, nullable=True)\n\n # last time any vespa relevant row metadata or the doc changed.\n # does not include last_synced\n last_modified: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=False, index=True, default=func.now()\n )\n\n # last successful sync to vespa\n last_synced: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True, index=True\n )\n # The following are not attached to User because the account/email may not be known\n # within Onyx\n # Something like the document creator\n primary_owners: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n secondary_owners: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n # Permission sync columns\n # Email addresses are saved at the document level for externally synced permissions\n # This is becuase the normal flow of assigning permissions is through the cc_pair\n # doesn't apply here\n external_user_emails: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n # These group ids have been prefixed by the source type\n external_user_group_ids: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n is_public: Mapped[bool] = mapped_column(Boolean, default=False)\n\n # Reference to parent hierarchy node (the folder/space containing this doc)\n # If None, document's hierarchy position is unknown or connector doesn't support hierarchy\n # SET NULL when hierarchy node is deleted - document should not be blocked by node deletion\n parent_hierarchy_node_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"hierarchy_node.id\", ondelete=\"SET NULL\"), nullable=True, index=True\n )\n\n # tables for the knowledge graph data\n kg_stage: Mapped[KGStage] = mapped_column(\n Enum(KGStage, native_enum=False),\n comment=\"Status of knowledge graph extraction for this document\",\n index=True,\n )\n\n kg_processing_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n retrieval_feedbacks: Mapped[list[\"DocumentRetrievalFeedback\"]] = relationship(\n \"DocumentRetrievalFeedback\", back_populates=\"document\"\n )\n\n doc_metadata: Mapped[dict[str, Any] | None] = mapped_column(\n postgresql.JSONB(), nullable=True, default=None\n )\n tags = relationship(\n \"Tag\",\n secondary=Document__Tag.__table__,\n back_populates=\"documents\",\n )\n\n # Relationship to parent hierarchy node (the folder/space containing this doc)\n parent_hierarchy_node: Mapped[\"HierarchyNode | None\"] = relationship(\n \"HierarchyNode\",\n back_populates=\"child_documents\",\n foreign_keys=[parent_hierarchy_node_id],\n )\n\n # For documents that ARE hierarchy nodes (e.g., Confluence pages with children)\n hierarchy_node: Mapped[\"HierarchyNode | None\"] = relationship(\n \"HierarchyNode\",\n back_populates=\"document\",\n foreign_keys=\"HierarchyNode.document_id\",\n passive_deletes=True,\n )\n # Personas that have this document directly attached for scoped search\n attached_personas: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=\"persona__document\",\n back_populates=\"attached_documents\",\n viewonly=True,\n )\n\n __table_args__ = (\n Index(\n \"ix_document_sync_status\",\n last_modified,\n last_synced,\n ),\n )\n\n\nclass OpenSearchDocumentMigrationRecord(Base):\n \"\"\"Tracks the migration status of documents from Vespa to OpenSearch.\n\n This table can be dropped when the migration is complete for all Onyx\n instances.\n \"\"\"\n\n __tablename__ = \"opensearch_document_migration_record\"\n\n document_id: Mapped[str] = mapped_column(\n String,\n ForeignKey(\"document.id\", ondelete=\"CASCADE\"),\n primary_key=True,\n nullable=False,\n index=True,\n )\n status: Mapped[OpenSearchDocumentMigrationStatus] = mapped_column(\n Enum(OpenSearchDocumentMigrationStatus, native_enum=False),\n default=OpenSearchDocumentMigrationStatus.PENDING,\n nullable=False,\n index=True,\n )\n error_message: Mapped[str | None] = mapped_column(Text, nullable=True)\n attempts_count: Mapped[int] = mapped_column(\n Integer, default=0, nullable=False, index=True\n )\n last_attempt_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n nullable=False,\n index=True,\n )\n\n document: Mapped[\"Document\"] = relationship(\"Document\")\n\n\nclass OpenSearchTenantMigrationRecord(Base):\n \"\"\"Tracks the state of the OpenSearch migration for a tenant.\n\n Should only contain one row.\n\n This table can be dropped when the migration is complete for all Onyx\n instances.\n \"\"\"\n\n __tablename__ = \"opensearch_tenant_migration_record\"\n __table_args__ = (\n # Singleton pattern - unique index on constant ensures only one row.\n Index(\"idx_opensearch_tenant_migration_singleton\", text(\"(true)\"), unique=True),\n )\n\n id: Mapped[int] = mapped_column(primary_key=True, nullable=False)\n document_migration_record_table_population_status: Mapped[\n OpenSearchTenantMigrationStatus\n ] = mapped_column(\n Enum(OpenSearchTenantMigrationStatus, native_enum=False),\n default=OpenSearchTenantMigrationStatus.PENDING,\n nullable=False,\n )\n num_times_observed_no_additional_docs_to_populate_migration_table: Mapped[int] = (\n mapped_column(Integer, default=0, nullable=False)\n )\n overall_document_migration_status: Mapped[OpenSearchTenantMigrationStatus] = (\n mapped_column(\n Enum(OpenSearchTenantMigrationStatus, native_enum=False),\n default=OpenSearchTenantMigrationStatus.PENDING,\n nullable=False,\n )\n )\n num_times_observed_no_additional_docs_to_migrate: Mapped[int] = mapped_column(\n Integer,\n default=0,\n nullable=False,\n )\n last_updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n # Opaque continuation token from Vespa's Visit API.\n # NULL means \"not started\".\n # Otherwise contains a serialized mapping between slice ID and continuation\n # token for that slice.\n vespa_visit_continuation_token: Mapped[str | None] = mapped_column(\n Text, nullable=True\n )\n total_chunks_migrated: Mapped[int] = mapped_column(\n Integer, default=0, nullable=False\n )\n total_chunks_errored: Mapped[int] = mapped_column(\n Integer, default=0, nullable=False\n )\n total_chunks_in_vespa: Mapped[int] = mapped_column(\n Integer, default=0, nullable=False\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n nullable=False,\n )\n migration_completed_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n enable_opensearch_retrieval: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False\n )\n approx_chunk_count_in_vespa: Mapped[int | None] = mapped_column(\n Integer, nullable=True\n )\n\n\nclass KGEntityType(Base):\n __tablename__ = \"kg_entity_type\"\n\n # Primary identifier\n id_name: Mapped[str] = mapped_column(\n String, primary_key=True, nullable=False, index=True\n )\n\n description: Mapped[str | None] = mapped_column(NullFilteredString, nullable=True)\n\n grounding: Mapped[str] = mapped_column(\n NullFilteredString, nullable=False, index=False\n )\n\n attributes: Mapped[dict | None] = mapped_column(\n postgresql.JSONB,\n nullable=True,\n default=dict,\n server_default=\"{}\",\n comment=\"Filtering based on document attribute\",\n )\n\n @property\n def parsed_attributes(self) -> KGEntityTypeAttributes:\n if self.attributes is None:\n return KGEntityTypeAttributes()\n\n try:\n return KGEntityTypeAttributes(**self.attributes)\n except ValidationError:\n return KGEntityTypeAttributes()\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n deep_extraction: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False\n )\n\n # Tracking fields\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n grounded_source_name: Mapped[str | None] = mapped_column(\n NullFilteredString, nullable=True, index=False\n )\n\n entity_values: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=True, default=None\n )\n\n clustering: Mapped[dict] = mapped_column(\n postgresql.JSONB,\n nullable=False,\n default=dict,\n server_default=\"{}\",\n comment=\"Clustering information for this entity type\",\n )\n\n\nclass KGRelationshipType(Base):\n __tablename__ = \"kg_relationship_type\"\n\n # Primary identifier\n id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n primary_key=True,\n nullable=False,\n index=True,\n )\n\n name: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n source_entity_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n target_entity_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n definition: Mapped[bool] = mapped_column(\n Boolean,\n nullable=False,\n default=False,\n comment=\"Whether this relationship type represents a definition\",\n )\n\n clustering: Mapped[dict] = mapped_column(\n postgresql.JSONB,\n nullable=False,\n default=dict,\n server_default=\"{}\",\n comment=\"Clustering information for this relationship type\",\n )\n\n type: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n # Tracking fields\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n # Relationships to EntityType\n source_type: Mapped[\"KGEntityType\"] = relationship(\n \"KGEntityType\",\n foreign_keys=[source_entity_type_id_name],\n backref=\"source_relationship_type\",\n )\n target_type: Mapped[\"KGEntityType\"] = relationship(\n \"KGEntityType\",\n foreign_keys=[target_entity_type_id_name],\n backref=\"target_relationship_type\",\n )\n\n\nclass KGRelationshipTypeExtractionStaging(Base):\n __tablename__ = \"kg_relationship_type_extraction_staging\"\n\n # Primary identifier\n id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n primary_key=True,\n nullable=False,\n index=True,\n )\n\n name: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n source_entity_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n target_entity_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n definition: Mapped[bool] = mapped_column(\n Boolean,\n nullable=False,\n default=False,\n comment=\"Whether this relationship type represents a definition\",\n )\n\n clustering: Mapped[dict] = mapped_column(\n postgresql.JSONB,\n nullable=False,\n default=dict,\n server_default=\"{}\",\n comment=\"Clustering information for this relationship type\",\n )\n\n type: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n transferred: Mapped[bool] = mapped_column(\n Boolean,\n nullable=False,\n default=False,\n )\n\n # Tracking fields\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n # Relationships to EntityType\n source_type: Mapped[\"KGEntityType\"] = relationship(\n \"KGEntityType\",\n foreign_keys=[source_entity_type_id_name],\n backref=\"source_relationship_type_staging\",\n )\n target_type: Mapped[\"KGEntityType\"] = relationship(\n \"KGEntityType\",\n foreign_keys=[target_entity_type_id_name],\n backref=\"target_relationship_type_staging\",\n )\n\n\nclass KGEntity(Base):\n __tablename__ = \"kg_entity\"\n\n # Primary identifier\n id_name: Mapped[str] = mapped_column(\n NullFilteredString, primary_key=True, index=True\n )\n\n # Basic entity information\n name: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n entity_key: Mapped[str] = mapped_column(\n NullFilteredString, nullable=True, index=True\n )\n parent_key: Mapped[str | None] = mapped_column(\n NullFilteredString, nullable=True, index=True\n )\n\n name_trigrams: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String(3)),\n nullable=True,\n )\n\n attributes: Mapped[dict] = mapped_column(\n postgresql.JSONB,\n nullable=False,\n default=dict,\n server_default=\"{}\",\n comment=\"Attributes for this entity\",\n )\n\n document_id: Mapped[str | None] = mapped_column(\n NullFilteredString, nullable=True, index=True\n )\n\n alternative_names: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n # Reference to KGEntityType\n entity_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n # Relationship to KGEntityType\n entity_type: Mapped[\"KGEntityType\"] = relationship(\"KGEntityType\", backref=\"entity\")\n\n description: Mapped[str | None] = mapped_column(String, nullable=True)\n\n keywords: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n # Access control\n acl: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n # Boosts - using JSON for flexibility\n boosts: Mapped[dict] = mapped_column(postgresql.JSONB, nullable=False, default=dict)\n\n event_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True),\n nullable=True,\n comment=\"Time of the event being processed\",\n )\n\n # Tracking fields\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n __table_args__ = (\n # Fixed column names in indexes\n Index(\"ix_entity_type_acl\", entity_type_id_name, acl),\n Index(\"ix_entity_name_search\", name, entity_type_id_name),\n )\n\n\nclass KGEntityExtractionStaging(Base):\n __tablename__ = \"kg_entity_extraction_staging\"\n\n # Primary identifier\n id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n primary_key=True,\n nullable=False,\n index=True,\n )\n\n # Basic entity information\n name: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n attributes: Mapped[dict] = mapped_column(\n postgresql.JSONB,\n nullable=False,\n default=dict,\n server_default=\"{}\",\n comment=\"Attributes for this entity\",\n )\n\n document_id: Mapped[str | None] = mapped_column(\n NullFilteredString, nullable=True, index=True\n )\n\n alternative_names: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n # Reference to KGEntityType\n entity_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n # Relationship to KGEntityType\n entity_type: Mapped[\"KGEntityType\"] = relationship(\n \"KGEntityType\", backref=\"entity_staging\"\n )\n\n description: Mapped[str | None] = mapped_column(String, nullable=True)\n\n keywords: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n # Access control\n acl: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n # Boosts - using JSON for flexibility\n boosts: Mapped[dict] = mapped_column(postgresql.JSONB, nullable=False, default=dict)\n\n transferred_id_name: Mapped[str | None] = mapped_column(\n NullFilteredString,\n nullable=True,\n )\n\n # Parent Child Information\n entity_key: Mapped[str] = mapped_column(\n NullFilteredString, nullable=True, index=True\n )\n parent_key: Mapped[str | None] = mapped_column(\n NullFilteredString, nullable=True, index=True\n )\n\n event_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True),\n nullable=True,\n comment=\"Time of the event being processed\",\n )\n\n # Tracking fields\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n __table_args__ = (\n # Fixed column names in indexes\n Index(\"ix_entity_type_acl\", entity_type_id_name, acl),\n Index(\"ix_entity_name_search\", name, entity_type_id_name),\n )\n\n\nclass KGRelationship(Base):\n __tablename__ = \"kg_relationship\"\n\n # Primary identifier - now part of composite key\n id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n nullable=False,\n index=True,\n )\n\n source_document: Mapped[str | None] = mapped_column(\n NullFilteredString, ForeignKey(\"document.id\"), nullable=True, index=True\n )\n\n # Source and target nodes (foreign keys to Entity table)\n source_node: Mapped[str] = mapped_column(\n NullFilteredString, ForeignKey(\"kg_entity.id_name\"), nullable=False, index=True\n )\n\n target_node: Mapped[str] = mapped_column(\n NullFilteredString, ForeignKey(\"kg_entity.id_name\"), nullable=False, index=True\n )\n\n source_node_type: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n target_node_type: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n # Relationship type\n type: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n # Add new relationship type reference\n relationship_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_relationship_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n # Add the SQLAlchemy relationship property\n relationship_type: Mapped[\"KGRelationshipType\"] = relationship(\n \"KGRelationshipType\", backref=\"relationship\"\n )\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n # Tracking fields\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n # Relationships to Entity table\n source: Mapped[\"KGEntity\"] = relationship(\"KGEntity\", foreign_keys=[source_node])\n target: Mapped[\"KGEntity\"] = relationship(\"KGEntity\", foreign_keys=[target_node])\n document: Mapped[\"Document\"] = relationship(\n \"Document\", foreign_keys=[source_document]\n )\n\n __table_args__ = (\n # Composite primary key\n PrimaryKeyConstraint(\"id_name\", \"source_document\"),\n # Index for querying relationships by type\n Index(\"ix_kg_relationship_type\", type),\n # Composite index for source/target queries\n Index(\"ix_kg_relationship_nodes\", source_node, target_node),\n # Ensure unique relationships between nodes of a specific type\n UniqueConstraint(\n \"source_node\",\n \"target_node\",\n \"type\",\n name=\"uq_kg_relationship_source_target_type\",\n ),\n )\n\n\nclass KGRelationshipExtractionStaging(Base):\n __tablename__ = \"kg_relationship_extraction_staging\"\n\n # Primary identifier - now part of composite key\n id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n nullable=False,\n index=True,\n )\n\n source_document: Mapped[str | None] = mapped_column(\n NullFilteredString, ForeignKey(\"document.id\"), nullable=True, index=True\n )\n\n # Source and target nodes (foreign keys to Entity table)\n source_node: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_extraction_staging.id_name\"),\n nullable=False,\n index=True,\n )\n\n target_node: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_extraction_staging.id_name\"),\n nullable=False,\n index=True,\n )\n\n source_node_type: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n target_node_type: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_entity_type.id_name\"),\n nullable=False,\n index=True,\n )\n\n # Relationship type\n type: Mapped[str] = mapped_column(NullFilteredString, nullable=False, index=True)\n\n # Add new relationship type reference\n relationship_type_id_name: Mapped[str] = mapped_column(\n NullFilteredString,\n ForeignKey(\"kg_relationship_type_extraction_staging.id_name\"),\n nullable=False,\n index=True,\n )\n\n # Add the SQLAlchemy relationship property\n relationship_type: Mapped[\"KGRelationshipTypeExtractionStaging\"] = relationship(\n \"KGRelationshipTypeExtractionStaging\", backref=\"relationship_staging\"\n )\n\n occurrences: Mapped[int] = mapped_column(Integer, nullable=False, default=1)\n\n transferred: Mapped[bool] = mapped_column(\n Boolean,\n nullable=False,\n default=False,\n )\n\n # Tracking fields\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n # Relationships to Entity table\n source: Mapped[\"KGEntityExtractionStaging\"] = relationship(\n \"KGEntityExtractionStaging\", foreign_keys=[source_node]\n )\n target: Mapped[\"KGEntityExtractionStaging\"] = relationship(\n \"KGEntityExtractionStaging\", foreign_keys=[target_node]\n )\n document: Mapped[\"Document\"] = relationship(\n \"Document\", foreign_keys=[source_document]\n )\n\n __table_args__ = (\n # Composite primary key\n PrimaryKeyConstraint(\"id_name\", \"source_document\"),\n # Index for querying relationships by type\n Index(\"ix_kg_relationship_type\", type),\n # Composite index for source/target queries\n Index(\"ix_kg_relationship_nodes\", source_node, target_node),\n # Ensure unique relationships between nodes of a specific type\n UniqueConstraint(\n \"source_node\",\n \"target_node\",\n \"type\",\n name=\"uq_kg_relationship_source_target_type\",\n ),\n )\n\n\nclass KGTerm(Base):\n __tablename__ = \"kg_term\"\n\n # Make id_term the primary key\n id_term: Mapped[str] = mapped_column(\n NullFilteredString, primary_key=True, nullable=False, index=True\n )\n\n # List of entity types this term applies to\n entity_types: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), nullable=False, default=list\n )\n\n # Tracking fields\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n __table_args__ = (\n # Index for searching terms with specific entity types\n Index(\"ix_search_term_entities\", entity_types),\n # Index for term lookups\n Index(\"ix_search_term_term\", id_term),\n )\n\n\nclass ChunkStats(Base):\n __tablename__ = \"chunk_stats\"\n # NOTE: if more sensitive data is added here for display, make sure to add user/group permission\n\n # this should correspond to the ID of the document\n # (as is passed around in Onyx)x\n id: Mapped[str] = mapped_column(\n NullFilteredString,\n primary_key=True,\n default=lambda context: (\n f\"{context.get_current_parameters()['document_id']}__{context.get_current_parameters()['chunk_in_doc_id']}\"\n ),\n index=True,\n )\n\n # Reference to parent document\n document_id: Mapped[str] = mapped_column(\n NullFilteredString, ForeignKey(\"document.id\"), nullable=False, index=True\n )\n\n chunk_in_doc_id: Mapped[int] = mapped_column(\n Integer,\n nullable=False,\n )\n\n information_content_boost: Mapped[float | None] = mapped_column(\n Float, nullable=True\n )\n\n last_modified: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=False, index=True, default=func.now()\n )\n last_synced: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True, index=True\n )\n\n __table_args__ = (\n Index(\n \"ix_chunk_sync_status\",\n last_modified,\n last_synced,\n ),\n UniqueConstraint(\n \"document_id\", \"chunk_in_doc_id\", name=\"uq_chunk_stats_doc_chunk\"\n ),\n )\n\n\nclass Tag(Base):\n __tablename__ = \"tag\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n tag_key: Mapped[str] = mapped_column(String)\n tag_value: Mapped[str] = mapped_column(String)\n source: Mapped[DocumentSource] = mapped_column(\n Enum(DocumentSource, native_enum=False)\n )\n is_list: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n documents = relationship(\n \"Document\",\n secondary=Document__Tag.__table__,\n back_populates=\"tags\",\n )\n\n __table_args__ = (\n UniqueConstraint(\n \"tag_key\",\n \"tag_value\",\n \"source\",\n \"is_list\",\n name=\"_tag_key_value_source_list_uc\",\n ),\n )\n\n\nclass Connector(Base):\n __tablename__ = \"connector\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n name: Mapped[str] = mapped_column(String)\n source: Mapped[DocumentSource] = mapped_column(\n Enum(DocumentSource, native_enum=False)\n )\n input_type = mapped_column(Enum(InputType, native_enum=False))\n connector_specific_config: Mapped[dict[str, Any]] = mapped_column(\n postgresql.JSONB()\n )\n indexing_start: Mapped[datetime.datetime | None] = mapped_column(\n DateTime, nullable=True\n )\n\n kg_processing_enabled: Mapped[bool] = mapped_column(\n Boolean,\n nullable=False,\n default=False,\n comment=\"Whether this connector should extract knowledge graph entities\",\n )\n\n kg_coverage_days: Mapped[int | None] = mapped_column(Integer, nullable=True)\n\n refresh_freq: Mapped[int | None] = mapped_column(Integer, nullable=True)\n prune_freq: Mapped[int | None] = mapped_column(Integer, nullable=True)\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n credentials: Mapped[list[\"ConnectorCredentialPair\"]] = relationship(\n \"ConnectorCredentialPair\",\n back_populates=\"connector\",\n cascade=\"all, delete-orphan\",\n )\n documents_by_connector: Mapped[list[\"DocumentByConnectorCredentialPair\"]] = (\n relationship(\n \"DocumentByConnectorCredentialPair\",\n back_populates=\"connector\",\n passive_deletes=True,\n )\n )\n\n # synchronize this validation logic with RefreshFrequencySchema etc on front end\n # until we have a centralized validation schema\n\n # TODO(rkuo): experiment with SQLAlchemy validators rather than manual checks\n # https://docs.sqlalchemy.org/en/20/orm/mapped_attributes.html\n def validate_refresh_freq(self) -> None:\n if self.refresh_freq is not None:\n if self.refresh_freq < 60:\n raise ValueError(\n \"refresh_freq must be greater than or equal to 1 minute.\"\n )\n\n def validate_prune_freq(self) -> None:\n if self.prune_freq is not None:\n if self.prune_freq < 300:\n raise ValueError(\n \"prune_freq must be greater than or equal to 5 minutes.\"\n )\n\n\nclass Credential(Base):\n __tablename__ = \"credential\"\n\n name: Mapped[str] = mapped_column(String, nullable=True)\n\n source: Mapped[DocumentSource] = mapped_column(\n Enum(DocumentSource, native_enum=False)\n )\n\n id: Mapped[int] = mapped_column(primary_key=True)\n credential_json: Mapped[SensitiveValue[dict[str, Any]] | None] = mapped_column(\n EncryptedJson()\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # if `true`, then all Admins will have access to the credential\n admin_public: Mapped[bool] = mapped_column(Boolean, default=True)\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n curator_public: Mapped[bool] = mapped_column(Boolean, default=False)\n\n connectors: Mapped[list[\"ConnectorCredentialPair\"]] = relationship(\n \"ConnectorCredentialPair\",\n back_populates=\"credential\",\n cascade=\"all, delete-orphan\",\n )\n documents_by_credential: Mapped[list[\"DocumentByConnectorCredentialPair\"]] = (\n relationship(\n \"DocumentByConnectorCredentialPair\",\n back_populates=\"credential\",\n passive_deletes=True,\n )\n )\n\n user: Mapped[User | None] = relationship(\"User\", back_populates=\"credentials\")\n\n\nclass FederatedConnector(Base):\n __tablename__ = \"federated_connector\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n source: Mapped[FederatedConnectorSource] = mapped_column(\n Enum(FederatedConnectorSource, native_enum=False)\n )\n credentials: Mapped[SensitiveValue[dict[str, Any]] | None] = mapped_column(\n EncryptedJson(), nullable=False\n )\n config: Mapped[dict[str, Any]] = mapped_column(\n postgresql.JSONB(), default=dict, nullable=False, server_default=\"{}\"\n )\n\n oauth_tokens: Mapped[list[\"FederatedConnectorOAuthToken\"]] = relationship(\n \"FederatedConnectorOAuthToken\",\n back_populates=\"federated_connector\",\n cascade=\"all, delete-orphan\",\n )\n document_sets: Mapped[list[\"FederatedConnector__DocumentSet\"]] = relationship(\n \"FederatedConnector__DocumentSet\",\n back_populates=\"federated_connector\",\n cascade=\"all, delete-orphan\",\n )\n\n\nclass FederatedConnectorOAuthToken(Base):\n __tablename__ = \"federated_connector_oauth_token\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n federated_connector_id: Mapped[int] = mapped_column(\n ForeignKey(\"federated_connector.id\", ondelete=\"CASCADE\"), nullable=False\n )\n user_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=False\n )\n token: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=False\n )\n expires_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime, nullable=True\n )\n\n federated_connector: Mapped[\"FederatedConnector\"] = relationship(\n \"FederatedConnector\", back_populates=\"oauth_tokens\"\n )\n user: Mapped[\"User\"] = relationship(\"User\")\n\n\nclass FederatedConnector__DocumentSet(Base):\n __tablename__ = \"federated_connector__document_set\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n federated_connector_id: Mapped[int] = mapped_column(\n ForeignKey(\"federated_connector.id\", ondelete=\"CASCADE\"), nullable=False\n )\n document_set_id: Mapped[int] = mapped_column(\n ForeignKey(\"document_set.id\", ondelete=\"CASCADE\"), nullable=False\n )\n # unique per source type. Validated before insertion.\n entities: Mapped[dict[str, Any]] = mapped_column(postgresql.JSONB(), nullable=False)\n\n federated_connector: Mapped[\"FederatedConnector\"] = relationship(\n \"FederatedConnector\", back_populates=\"document_sets\"\n )\n document_set: Mapped[\"DocumentSet\"] = relationship(\n \"DocumentSet\", back_populates=\"federated_connectors\"\n )\n\n __table_args__ = (\n UniqueConstraint(\n \"federated_connector_id\",\n \"document_set_id\",\n name=\"uq_federated_connector_document_set\",\n ),\n )\n\n\nclass SearchSettings(Base):\n __tablename__ = \"search_settings\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n model_name: Mapped[str] = mapped_column(String)\n model_dim: Mapped[int] = mapped_column(Integer)\n normalize: Mapped[bool] = mapped_column(Boolean)\n query_prefix: Mapped[str | None] = mapped_column(String, nullable=True)\n passage_prefix: Mapped[str | None] = mapped_column(String, nullable=True)\n\n status: Mapped[IndexModelStatus] = mapped_column(\n Enum(IndexModelStatus, native_enum=False)\n )\n index_name: Mapped[str] = mapped_column(String)\n provider_type: Mapped[EmbeddingProvider | None] = mapped_column(\n ForeignKey(\"embedding_provider.provider_type\"), nullable=True\n )\n\n # Type of switchover to perform when switching embedding models\n # REINDEX: waits for all connectors to complete\n # ACTIVE_ONLY: waits for only non-paused connectors to complete\n # INSTANT: swaps immediately without waiting\n switchover_type: Mapped[SwitchoverType] = mapped_column(\n Enum(SwitchoverType, native_enum=False), default=SwitchoverType.REINDEX\n )\n\n # allows for quantization -> less memory usage for a small performance hit\n embedding_precision: Mapped[EmbeddingPrecision] = mapped_column(\n Enum(EmbeddingPrecision, native_enum=False)\n )\n\n # can be used to reduce dimensionality of vectors and save memory with\n # a small performance hit. More details in the `Reducing embedding dimensions`\n # section here:\n # https://platform.openai.com/docs/guides/embeddings#embedding-models\n # If not specified, will just use the model_dim without any reduction.\n # NOTE: this is only currently available for OpenAI models\n reduced_dimension: Mapped[int | None] = mapped_column(Integer, nullable=True)\n\n # Mini and Large Chunks (large chunk also checks for model max context)\n multipass_indexing: Mapped[bool] = mapped_column(Boolean, default=True)\n\n # Contextual RAG\n enable_contextual_rag: Mapped[bool] = mapped_column(Boolean, default=False)\n\n # Contextual RAG LLM\n contextual_rag_llm_name: Mapped[str | None] = mapped_column(String, nullable=True)\n contextual_rag_llm_provider: Mapped[str | None] = mapped_column(\n String, nullable=True\n )\n\n multilingual_expansion: Mapped[list[str]] = mapped_column(\n postgresql.ARRAY(String), default=[]\n )\n\n cloud_provider: Mapped[\"CloudEmbeddingProvider\"] = relationship(\n \"CloudEmbeddingProvider\",\n back_populates=\"search_settings\",\n foreign_keys=[provider_type],\n )\n\n index_attempts: Mapped[list[\"IndexAttempt\"]] = relationship(\n \"IndexAttempt\", back_populates=\"search_settings\"\n )\n\n __table_args__ = (\n Index(\n \"ix_embedding_model_present_unique\",\n \"status\",\n unique=True,\n postgresql_where=(status == IndexModelStatus.PRESENT),\n ),\n Index(\n \"ix_embedding_model_future_unique\",\n \"status\",\n unique=True,\n postgresql_where=(status == IndexModelStatus.FUTURE),\n ),\n )\n\n def __repr__(self) -> str:\n return f\"\"\n\n @property\n def api_version(self) -> str | None:\n return (\n self.cloud_provider.api_version if self.cloud_provider is not None else None\n )\n\n @property\n def deployment_name(self) -> str | None:\n return (\n self.cloud_provider.deployment_name\n if self.cloud_provider is not None\n else None\n )\n\n @property\n def api_url(self) -> str | None:\n return self.cloud_provider.api_url if self.cloud_provider is not None else None\n\n @property\n def api_key(self) -> str | None:\n if self.cloud_provider is None or self.cloud_provider.api_key is None:\n return None\n return self.cloud_provider.api_key.get_value(apply_mask=False)\n\n @property\n def large_chunks_enabled(self) -> bool:\n \"\"\"\n Given multipass usage and an embedder, decides whether large chunks are allowed\n based on model/provider constraints.\n \"\"\"\n # Only local models that support a larger context are from Nomic\n # Cohere does not support larger contexts (they recommend not going above ~512 tokens)\n return SearchSettings.can_use_large_chunks(\n self.multipass_indexing, self.model_name, self.provider_type\n )\n\n @property\n def final_embedding_dim(self) -> int:\n return self.reduced_dimension or self.model_dim\n\n @staticmethod\n def can_use_large_chunks(\n multipass: bool, model_name: str, provider_type: EmbeddingProvider | None\n ) -> bool:\n \"\"\"\n Given multipass usage and an embedder, decides whether large chunks are allowed\n based on model/provider constraints.\n \"\"\"\n # Only local models that support a larger context are from Nomic\n # Cohere does not support larger contexts (they recommend not going above ~512 tokens)\n return (\n multipass\n and model_name.startswith(\"nomic-ai\")\n and provider_type != EmbeddingProvider.COHERE\n )\n\n\nclass IndexAttempt(Base):\n \"\"\"\n Represents an attempt to index a group of 0 or more documents from a\n source. For example, a single pull from Google Drive, a single event from\n slack event API, or a single website crawl.\n \"\"\"\n\n __tablename__ = \"index_attempt\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n connector_credential_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"),\n nullable=False,\n )\n\n # Some index attempts that run from beginning will still have this as False\n # This is only for attempts that are explicitly marked as from the start via\n # the run once API\n from_beginning: Mapped[bool] = mapped_column(Boolean)\n status: Mapped[IndexingStatus] = mapped_column(\n Enum(IndexingStatus, native_enum=False, index=True)\n )\n # The two below may be slightly out of sync if user switches Embedding Model\n new_docs_indexed: Mapped[int | None] = mapped_column(Integer, default=0)\n total_docs_indexed: Mapped[int | None] = mapped_column(Integer, default=0)\n docs_removed_from_index: Mapped[int | None] = mapped_column(Integer, default=0)\n # only filled if status = \"failed\"\n error_msg: Mapped[str | None] = mapped_column(Text, default=None)\n # only filled if status = \"failed\" AND an unhandled exception caused the failure\n full_exception_trace: Mapped[str | None] = mapped_column(Text, default=None)\n # Nullable because in the past, we didn't allow swapping out embedding models live\n search_settings_id: Mapped[int] = mapped_column(\n ForeignKey(\"search_settings.id\", ondelete=\"SET NULL\"),\n nullable=True,\n )\n\n # for polling connectors, the start and end time of the poll window\n # will be set when the index attempt starts\n poll_range_start: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True, default=None\n )\n poll_range_end: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True, default=None\n )\n\n # Points to the last checkpoint that was saved for this run. The pointer here\n # can be taken to the FileStore to grab the actual checkpoint value\n checkpoint_pointer: Mapped[str | None] = mapped_column(String, nullable=True)\n\n # Database-based coordination fields (replacing Redis fencing)\n celery_task_id: Mapped[str | None] = mapped_column(String, nullable=True)\n cancellation_requested: Mapped[bool] = mapped_column(Boolean, default=False)\n\n # Batch coordination fields\n # Once this is set, docfetching has completed\n total_batches: Mapped[int | None] = mapped_column(Integer, nullable=True)\n # batches that are fully indexed (i.e. have completed docfetching and docprocessing)\n completed_batches: Mapped[int] = mapped_column(Integer, default=0)\n # TODO: unused, remove this column\n total_failures_batch_level: Mapped[int] = mapped_column(Integer, default=0)\n total_chunks: Mapped[int] = mapped_column(Integer, default=0)\n\n # Progress tracking for stall detection\n last_progress_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n last_batches_completed_count: Mapped[int] = mapped_column(Integer, default=0)\n\n # Heartbeat tracking for worker liveness detection\n heartbeat_counter: Mapped[int] = mapped_column(Integer, default=0)\n last_heartbeat_value: Mapped[int] = mapped_column(Integer, default=0)\n last_heartbeat_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n index=True,\n )\n # when the actual indexing run began\n # NOTE: will use the api_server clock rather than DB server clock\n time_started: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n\n connector_credential_pair: Mapped[ConnectorCredentialPair] = relationship(\n \"ConnectorCredentialPair\", back_populates=\"index_attempts\"\n )\n\n search_settings: Mapped[SearchSettings | None] = relationship(\n \"SearchSettings\", back_populates=\"index_attempts\"\n )\n\n error_rows = relationship(\n \"IndexAttemptError\",\n back_populates=\"index_attempt\",\n cascade=\"all, delete-orphan\",\n )\n\n __table_args__ = (\n Index(\n \"ix_index_attempt_latest_for_connector_credential_pair\",\n \"connector_credential_pair_id\",\n \"time_created\",\n ),\n Index(\n \"ix_index_attempt_ccpair_search_settings_time_updated\",\n \"connector_credential_pair_id\",\n \"search_settings_id\",\n desc(\"time_updated\"),\n unique=False,\n ),\n Index(\n \"ix_index_attempt_cc_pair_settings_poll\",\n \"connector_credential_pair_id\",\n \"search_settings_id\",\n \"status\",\n desc(\"time_updated\"),\n ),\n # NEW: Index for coordination queries\n Index(\n \"ix_index_attempt_active_coordination\",\n \"connector_credential_pair_id\",\n \"search_settings_id\",\n \"status\",\n ),\n )\n\n def __repr__(self) -> str:\n return (\n f\"\"\n f\"time_created={self.time_created!r}, \"\n f\"time_updated={self.time_updated!r}, \"\n )\n\n def is_finished(self) -> bool:\n return self.status.is_terminal()\n\n def is_coordination_complete(self) -> bool:\n \"\"\"Check if all batches have been processed\"\"\"\n return (\n self.total_batches is not None\n and self.completed_batches >= self.total_batches\n )\n\n\nclass HierarchyFetchAttempt(Base):\n \"\"\"Tracks attempts to fetch hierarchy nodes from a source\"\"\"\n\n __tablename__ = \"hierarchy_fetch_attempt\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n\n connector_credential_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n\n status: Mapped[IndexingStatus] = mapped_column(\n Enum(IndexingStatus, native_enum=False), nullable=False, index=True\n )\n\n # Statistics\n nodes_fetched: Mapped[int | None] = mapped_column(Integer, default=0)\n nodes_updated: Mapped[int | None] = mapped_column(Integer, default=0)\n\n # Error information (only filled if status = \"failed\")\n error_msg: Mapped[str | None] = mapped_column(Text, default=None)\n full_exception_trace: Mapped[str | None] = mapped_column(Text, default=None)\n\n # Timestamps\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n index=True,\n )\n time_started: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n\n # Relationships\n connector_credential_pair: Mapped[\"ConnectorCredentialPair\"] = relationship(\n \"ConnectorCredentialPair\"\n )\n\n __table_args__ = (\n Index(\n \"ix_hierarchy_fetch_attempt_cc_pair\",\n connector_credential_pair_id,\n ),\n )\n\n\nclass IndexAttemptError(Base):\n __tablename__ = \"index_attempt_errors\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n index_attempt_id: Mapped[int] = mapped_column(\n ForeignKey(\"index_attempt.id\"),\n nullable=False,\n )\n connector_credential_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"),\n nullable=False,\n )\n\n document_id: Mapped[str | None] = mapped_column(String, nullable=True)\n document_link: Mapped[str | None] = mapped_column(String, nullable=True)\n\n entity_id: Mapped[str | None] = mapped_column(String, nullable=True)\n failed_time_range_start: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n failed_time_range_end: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n failure_message: Mapped[str] = mapped_column(Text)\n is_resolved: Mapped[bool] = mapped_column(Boolean, default=False)\n\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n )\n\n # This is the reverse side of the relationship\n index_attempt = relationship(\"IndexAttempt\", back_populates=\"error_rows\")\n\n\nclass SyncRecord(Base):\n \"\"\"\n Represents the status of a \"sync\" operation (e.g. document set, user group, deletion).\n\n A \"sync\" operation is an operation which needs to update a set of documents within\n Vespa, usually to match the state of Postgres.\n \"\"\"\n\n __tablename__ = \"sync_record\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n # document set id, user group id, or deletion id\n entity_id: Mapped[int] = mapped_column(Integer)\n\n sync_type: Mapped[SyncType] = mapped_column(Enum(SyncType, native_enum=False))\n sync_status: Mapped[SyncStatus] = mapped_column(Enum(SyncStatus, native_enum=False))\n\n num_docs_synced: Mapped[int] = mapped_column(Integer, default=0)\n\n sync_start_time: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))\n sync_end_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n __table_args__ = (\n Index(\n \"ix_sync_record_entity_id_sync_type_sync_start_time\",\n \"entity_id\",\n \"sync_type\",\n \"sync_start_time\",\n ),\n Index(\n \"ix_sync_record_entity_id_sync_type_sync_status\",\n \"entity_id\",\n \"sync_type\",\n \"sync_status\",\n ),\n )\n\n\nclass HierarchyNodeByConnectorCredentialPair(Base):\n \"\"\"Tracks which cc_pairs reference each hierarchy node.\n\n During pruning, stale entries are removed for the current cc_pair.\n Hierarchy nodes with zero remaining entries are then deleted.\n \"\"\"\n\n __tablename__ = \"hierarchy_node_by_connector_credential_pair\"\n\n hierarchy_node_id: Mapped[int] = mapped_column(\n ForeignKey(\"hierarchy_node.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n connector_id: Mapped[int] = mapped_column(primary_key=True)\n credential_id: Mapped[int] = mapped_column(primary_key=True)\n\n __table_args__ = (\n ForeignKeyConstraint(\n [\"connector_id\", \"credential_id\"],\n [\n \"connector_credential_pair.connector_id\",\n \"connector_credential_pair.credential_id\",\n ],\n ondelete=\"CASCADE\",\n ),\n Index(\n \"ix_hierarchy_node_cc_pair_connector_credential\",\n \"connector_id\",\n \"credential_id\",\n ),\n )\n\n\nclass DocumentByConnectorCredentialPair(Base):\n \"\"\"Represents an indexing of a document by a specific connector / credential pair\"\"\"\n\n __tablename__ = \"document_by_connector_credential_pair\"\n\n id: Mapped[str] = mapped_column(ForeignKey(\"document.id\"), primary_key=True)\n # TODO: transition this to use the ConnectorCredentialPair id directly\n connector_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n credential_id: Mapped[int] = mapped_column(\n ForeignKey(\"credential.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n # used to better keep track of document counts at a connector level\n # e.g. if a document is added as part of permission syncing, it should\n # not be counted as part of the connector's document count until\n # the actual indexing is complete\n has_been_indexed: Mapped[bool] = mapped_column(Boolean)\n\n connector: Mapped[Connector] = relationship(\n \"Connector\", back_populates=\"documents_by_connector\", passive_deletes=True\n )\n credential: Mapped[Credential] = relationship(\n \"Credential\", back_populates=\"documents_by_credential\", passive_deletes=True\n )\n\n __table_args__ = (\n Index(\n \"idx_document_cc_pair_connector_credential\",\n \"connector_id\",\n \"credential_id\",\n unique=False,\n ),\n # Index to optimize get_document_counts_for_cc_pairs query pattern\n Index(\n \"idx_document_cc_pair_counts\",\n \"connector_id\",\n \"credential_id\",\n \"has_been_indexed\",\n unique=False,\n ),\n )\n\n\n\"\"\"\nMessages Tables\n\"\"\"\n\n\nclass ChatSession(Base):\n __tablename__ = \"chat_session\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n persona_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"persona.id\"), nullable=True\n )\n description: Mapped[str | None] = mapped_column(Text, nullable=True)\n # This chat created by OnyxBot\n onyxbot_flow: Mapped[bool] = mapped_column(Boolean, default=False)\n # Only ever set to True if system is set to not hard-delete chats\n deleted: Mapped[bool] = mapped_column(Boolean, default=False)\n # controls whether or not this conversation is viewable by others\n shared_status: Mapped[ChatSessionSharedStatus] = mapped_column(\n Enum(ChatSessionSharedStatus, native_enum=False),\n default=ChatSessionSharedStatus.PRIVATE,\n )\n\n current_alternate_model: Mapped[str | None] = mapped_column(String, default=None)\n\n slack_thread_id: Mapped[str | None] = mapped_column(\n String, nullable=True, default=None\n )\n\n project_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"user_project.id\"), nullable=True\n )\n\n project: Mapped[\"UserProject\"] = relationship(\n \"UserProject\", back_populates=\"chat_sessions\", foreign_keys=[project_id]\n )\n\n # the latest \"overrides\" specified by the user. These take precedence over\n # the attached persona. However, overrides specified directly in the\n # `send-message` call will take precedence over these.\n # NOTE: currently only used by the chat seeding flow, will be used in the\n # future once we allow users to override default values via the Chat UI\n # itself\n llm_override: Mapped[LLMOverride | None] = mapped_column(\n PydanticType(LLMOverride), nullable=True\n )\n\n # The latest temperature override specified by the user\n temperature_override: Mapped[float | None] = mapped_column(Float, nullable=True)\n\n prompt_override: Mapped[PromptOverride | None] = mapped_column(\n PydanticType(PromptOverride), nullable=True\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n user: Mapped[User] = relationship(\"User\", back_populates=\"chat_sessions\")\n messages: Mapped[list[\"ChatMessage\"]] = relationship(\n \"ChatMessage\",\n back_populates=\"chat_session\",\n cascade=\"all, delete-orphan\",\n foreign_keys=\"ChatMessage.chat_session_id\",\n )\n persona: Mapped[\"Persona\"] = relationship(\"Persona\")\n\n\nclass ChatMessage(Base):\n \"\"\"Note, the first message in a chain has no contents, it's a workaround to allow edits\n on the first message of a session, an empty root node basically\n\n Since every user message is followed by a LLM response, chat messages generally come in pairs.\n Keeping them as separate messages however for future Agentification extensions\n Fields will be largely duplicated in the pair.\n \"\"\"\n\n __tablename__ = \"chat_message\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n # Where is this message located\n chat_session_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), ForeignKey(\"chat_session.id\")\n )\n\n # Parent message pointer for the tree structure, nullable because the first message is\n # an empty root node to allow edits on the first message of a session.\n parent_message_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\"), nullable=True\n )\n # This only maps to the latest because only that message chain is needed.\n # It can be updated as needed to trace other branches.\n latest_child_message_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\"), nullable=True\n )\n\n # Only set on summary messages - the ID of the last message included in this summary\n # Used for chat history compression\n last_summarized_message_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"SET NULL\"),\n nullable=True,\n )\n\n # For multi-model turns: the user message points to which assistant response\n # was selected as the preferred one to continue the conversation with.\n preferred_response_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"SET NULL\"), nullable=True\n )\n\n # The display name of the model that generated this assistant message\n model_display_name: Mapped[str | None] = mapped_column(String, nullable=True)\n\n # What does this message contain\n reasoning_tokens: Mapped[str | None] = mapped_column(Text, nullable=True)\n message: Mapped[str] = mapped_column(Text)\n token_count: Mapped[int] = mapped_column(Integer)\n message_type: Mapped[MessageType] = mapped_column(\n Enum(MessageType, native_enum=False)\n )\n # Files attached to the message, when parsed into history, it becomes a separate message\n files: Mapped[list[FileDescriptor] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Maps the citation numbers to a SearchDoc id\n citations: Mapped[dict[int, int] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Metadata\n error: Mapped[str | None] = mapped_column(Text, nullable=True)\n time_sent: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n # True if this assistant message is a clarification question (deep research flow)\n is_clarification: Mapped[bool] = mapped_column(Boolean, default=False)\n # Duration in seconds for processing this message (assistant messages only)\n processing_duration_seconds: Mapped[float | None] = mapped_column(\n Float, nullable=True\n )\n\n # Relationships\n chat_session: Mapped[ChatSession] = relationship(\n \"ChatSession\",\n back_populates=\"messages\",\n foreign_keys=[chat_session_id],\n )\n\n chat_message_feedbacks: Mapped[list[\"ChatMessageFeedback\"]] = relationship(\n \"ChatMessageFeedback\",\n back_populates=\"chat_message\",\n )\n\n document_feedbacks: Mapped[list[\"DocumentRetrievalFeedback\"]] = relationship(\n \"DocumentRetrievalFeedback\",\n back_populates=\"chat_message\",\n )\n\n # Even though search docs come from tool calls, the answer has a final set of saved search docs that we will show\n search_docs: Mapped[list[\"SearchDoc\"]] = relationship(\n \"SearchDoc\",\n secondary=ChatMessage__SearchDoc.__table__,\n back_populates=\"chat_messages\",\n cascade=\"all, delete-orphan\",\n single_parent=True,\n )\n\n parent_message: Mapped[\"ChatMessage | None\"] = relationship(\n \"ChatMessage\",\n foreign_keys=[parent_message_id],\n remote_side=\"ChatMessage.id\",\n )\n\n latest_child_message: Mapped[\"ChatMessage | None\"] = relationship(\n \"ChatMessage\",\n foreign_keys=[latest_child_message_id],\n remote_side=\"ChatMessage.id\",\n )\n\n preferred_response: Mapped[\"ChatMessage | None\"] = relationship(\n \"ChatMessage\",\n foreign_keys=[preferred_response_id],\n remote_side=\"ChatMessage.id\",\n )\n\n # Chat messages only need to know their immediate tool call children\n # If there are nested tool calls, they are stored in the tool_call_children relationship.\n tool_calls: Mapped[list[\"ToolCall\"] | None] = relationship(\n \"ToolCall\",\n back_populates=\"chat_message\",\n )\n\n standard_answers: Mapped[list[\"StandardAnswer\"]] = relationship(\n \"StandardAnswer\",\n secondary=ChatMessage__StandardAnswer.__table__,\n back_populates=\"chat_messages\",\n )\n\n\nclass ToolCall(Base):\n \"\"\"Represents a Tool Call and Tool Response\"\"\"\n\n __tablename__ = \"tool_call\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n chat_session_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), ForeignKey(\"chat_session.id\", ondelete=\"CASCADE\")\n )\n\n # If this is not None, it's a top level tool call from the user message\n # If this is None, it's a lower level call from another tool/agent\n parent_chat_message_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # If this is not None, this tool call is a child of another tool call\n parent_tool_call_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"tool_call.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # The tools with the same turn number (and parent) were called in parallel\n # Ones with different turn numbers (and same parent) were called sequentially\n turn_number: Mapped[int] = mapped_column(Integer)\n # Index order of tool calls from the LLM for parallel tool calls\n tab_index: Mapped[int] = mapped_column(Integer, default=0)\n\n # Not a FK because we want to be able to delete the tool without deleting\n # this entry\n tool_id: Mapped[int] = mapped_column(Integer())\n # This is needed because LLMs expect the tool call and the response to have matching IDs\n # This is better than just regenerating one randomly\n tool_call_id: Mapped[str] = mapped_column(String())\n # Preceeding reasoning tokens for this tool call, not included in the history\n reasoning_tokens: Mapped[str | None] = mapped_column(Text, nullable=True)\n # For \"Agents\" like the Research Agent for Deep Research -\n # the argument and final report are stored as the argument and response.\n tool_call_arguments: Mapped[dict[str, JSON_ro]] = mapped_column(postgresql.JSONB())\n tool_call_response: Mapped[str] = mapped_column(Text)\n # This just counts the number of tokens in the arg because it's all that's kept for the history\n # Only the top level tools (the ones with a parent_chat_message_id) have token counts that are counted\n # towards the session total.\n tool_call_tokens: Mapped[int] = mapped_column(Integer())\n # For image generation tool - stores GeneratedImage objects for replay\n generated_images: Mapped[list[dict] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Relationships\n chat_session: Mapped[ChatSession] = relationship(\"ChatSession\")\n\n chat_message: Mapped[\"ChatMessage | None\"] = relationship(\n \"ChatMessage\",\n foreign_keys=[parent_chat_message_id],\n back_populates=\"tool_calls\",\n )\n parent_tool_call: Mapped[\"ToolCall | None\"] = relationship(\n \"ToolCall\",\n foreign_keys=[parent_tool_call_id],\n remote_side=\"ToolCall.id\",\n )\n tool_call_children: Mapped[list[\"ToolCall\"]] = relationship(\n \"ToolCall\",\n foreign_keys=[parent_tool_call_id],\n back_populates=\"parent_tool_call\",\n )\n # Other tools may need to save other things, might need to figure out a more generic way to store\n # rich tool returns\n search_docs: Mapped[list[\"SearchDoc\"]] = relationship(\n \"SearchDoc\",\n secondary=ToolCall__SearchDoc.__table__,\n back_populates=\"tool_calls\",\n cascade=\"all, delete-orphan\",\n single_parent=True,\n )\n\n\nclass SearchDoc(Base):\n \"\"\"Different from Document table. This one stores the state of a document from a retrieval.\n This allows chat sessions to be replayed with the searched docs\n\n Notably, this does not include the contents of the Document/Chunk, during inference if a stored\n SearchDoc is selected, an inference must be remade to retrieve the contents\n \"\"\"\n\n __tablename__ = \"search_doc\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n document_id: Mapped[str] = mapped_column(String)\n chunk_ind: Mapped[int] = mapped_column(Integer)\n semantic_id: Mapped[str] = mapped_column(String)\n link: Mapped[str | None] = mapped_column(String, nullable=True)\n blurb: Mapped[str] = mapped_column(String)\n boost: Mapped[int] = mapped_column(Integer)\n source_type: Mapped[DocumentSource] = mapped_column(\n Enum(DocumentSource, native_enum=False)\n )\n hidden: Mapped[bool] = mapped_column(Boolean)\n doc_metadata: Mapped[dict[str, str | list[str]]] = mapped_column(postgresql.JSONB())\n score: Mapped[float] = mapped_column(Float)\n match_highlights: Mapped[list[str]] = mapped_column(postgresql.ARRAY(String))\n # This is for the document, not this row in the table\n updated_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n primary_owners: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n secondary_owners: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n is_internet: Mapped[bool] = mapped_column(Boolean, default=False, nullable=True)\n\n is_relevant: Mapped[bool | None] = mapped_column(Boolean, nullable=True)\n relevance_explanation: Mapped[str | None] = mapped_column(String, nullable=True)\n\n chat_messages: Mapped[list[\"ChatMessage\"]] = relationship(\n \"ChatMessage\",\n secondary=ChatMessage__SearchDoc.__table__,\n back_populates=\"search_docs\",\n )\n\n tool_calls: Mapped[list[\"ToolCall\"]] = relationship(\n \"ToolCall\",\n secondary=ToolCall__SearchDoc.__table__,\n back_populates=\"search_docs\",\n )\n\n\nclass SearchQuery(Base):\n # This table contains search queries for the Search UI. There are no followups and less is stored because the reply\n # functionality is simply to rerun the search query again as things may have changed and this is more common for search.\n __tablename__ = \"search_query\"\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n user_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), ForeignKey(\"user.id\", ondelete=\"CASCADE\")\n )\n query: Mapped[str] = mapped_column(String)\n query_expansions: Mapped[list[str] | None] = mapped_column(\n postgresql.ARRAY(String), nullable=True\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n\n\"\"\"\nFeedback, Logging, Metrics Tables\n\"\"\"\n\n\nclass DocumentRetrievalFeedback(Base):\n __tablename__ = \"document_retrieval_feedback\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n chat_message_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"SET NULL\"), nullable=True\n )\n document_id: Mapped[str] = mapped_column(ForeignKey(\"document.id\"))\n # How high up this document is in the results, 1 for first\n document_rank: Mapped[int] = mapped_column(Integer)\n clicked: Mapped[bool] = mapped_column(Boolean, default=False)\n feedback: Mapped[SearchFeedbackType | None] = mapped_column(\n Enum(SearchFeedbackType, native_enum=False), nullable=True\n )\n\n chat_message: Mapped[ChatMessage] = relationship(\n \"ChatMessage\",\n back_populates=\"document_feedbacks\",\n foreign_keys=[chat_message_id],\n )\n document: Mapped[Document] = relationship(\n \"Document\", back_populates=\"retrieval_feedbacks\"\n )\n\n\nclass ChatMessageFeedback(Base):\n __tablename__ = \"chat_feedback\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n chat_message_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"chat_message.id\", ondelete=\"SET NULL\"), nullable=True\n )\n is_positive: Mapped[bool | None] = mapped_column(Boolean, nullable=True)\n required_followup: Mapped[bool | None] = mapped_column(Boolean, nullable=True)\n feedback_text: Mapped[str | None] = mapped_column(Text, nullable=True)\n predefined_feedback: Mapped[str | None] = mapped_column(String, nullable=True)\n\n chat_message: Mapped[ChatMessage] = relationship(\n \"ChatMessage\",\n back_populates=\"chat_message_feedbacks\",\n foreign_keys=[chat_message_id],\n )\n\n\nclass LLMProvider(Base):\n __tablename__ = \"llm_provider\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True)\n provider: Mapped[str] = mapped_column(String)\n api_key: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=True\n )\n api_base: Mapped[str | None] = mapped_column(String, nullable=True)\n api_version: Mapped[str | None] = mapped_column(String, nullable=True)\n # custom configs that should be passed to the LLM provider at inference time\n # (e.g. `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, etc. for bedrock)\n custom_config: Mapped[dict[str, str] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Deprecated: use LLMModelFlow with CHAT flow type instead\n default_model_name: Mapped[str | None] = mapped_column(String, nullable=True)\n\n deployment_name: Mapped[str | None] = mapped_column(String, nullable=True)\n\n # Deprecated: use LLMModelFlow.is_default with CHAT flow type instead\n is_default_provider: Mapped[bool | None] = mapped_column(Boolean, nullable=True)\n # Deprecated: use LLMModelFlow.is_default with VISION flow type instead\n is_default_vision_provider: Mapped[bool | None] = mapped_column(Boolean)\n # Deprecated: use LLMModelFlow with VISION flow type instead\n default_vision_model: Mapped[str | None] = mapped_column(String, nullable=True)\n # EE only\n is_public: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n # Auto mode: models, visibility, and defaults are managed by GitHub config\n is_auto_mode: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n groups: Mapped[list[\"UserGroup\"]] = relationship(\n \"UserGroup\",\n secondary=\"llm_provider__user_group\",\n viewonly=True,\n )\n personas: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=\"llm_provider__persona\",\n back_populates=\"allowed_by_llm_providers\",\n viewonly=True,\n )\n model_configurations: Mapped[list[\"ModelConfiguration\"]] = relationship(\n \"ModelConfiguration\",\n back_populates=\"llm_provider\",\n foreign_keys=\"ModelConfiguration.llm_provider_id\",\n )\n\n\nclass ModelConfiguration(Base):\n __tablename__ = \"model_configuration\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n llm_provider_id: Mapped[int] = mapped_column(\n ForeignKey(\"llm_provider.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n name: Mapped[str] = mapped_column(String, nullable=False)\n\n # Represents whether or not a given model will be usable by the end user or not.\n # This field is primarily used for \"Well Known LLM Providers\", since for them,\n # we have a pre-defined list of LLM models that we allow them to choose from.\n # For example, for OpenAI, we allow the end-user to choose multiple models from\n # `[\"gpt-4\", \"gpt-4o\", etc.]`. Once they make their selections, we set each\n # selected model to `is_visible = True`.\n #\n # For \"Custom LLM Providers\", we don't provide a comprehensive list of models\n # for the end-user to choose from; *they provide it themselves*. Therefore,\n # for Custom LLM Providers, `is_visible` will always be True.\n is_visible: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n # Max input tokens can be null when:\n # - The end-user configures models through a \"Well Known LLM Provider\".\n # - The end-user is configuring a model and chooses not to set a max-input-tokens limit.\n max_input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)\n\n # Deprecated: use LLMModelFlow with VISION flow type instead\n supports_image_input: Mapped[bool | None] = mapped_column(Boolean, nullable=True)\n\n # Human-readable display name for the model.\n # For dynamic providers (OpenRouter, Bedrock, Ollama), this comes from the source API.\n # For static providers (OpenAI, Anthropic), this may be null and will fall back to LiteLLM.\n display_name: Mapped[str | None] = mapped_column(String, nullable=True)\n\n llm_provider: Mapped[\"LLMProvider\"] = relationship(\n \"LLMProvider\",\n back_populates=\"model_configurations\",\n )\n\n llm_model_flows: Mapped[list[\"LLMModelFlow\"]] = relationship(\n \"LLMModelFlow\",\n back_populates=\"model_configuration\",\n cascade=\"all, delete-orphan\",\n passive_deletes=True,\n )\n\n @property\n def llm_model_flow_types(self) -> list[LLMModelFlowType]:\n return [flow.llm_model_flow_type for flow in self.llm_model_flows]\n\n\nclass LLMModelFlow(Base):\n __tablename__ = \"llm_model_flow\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n\n llm_model_flow_type: Mapped[LLMModelFlowType] = mapped_column(\n Enum(LLMModelFlowType, native_enum=False), nullable=False\n )\n model_configuration_id: Mapped[int] = mapped_column(\n ForeignKey(\"model_configuration.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n is_default: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n model_configuration: Mapped[\"ModelConfiguration\"] = relationship(\n \"ModelConfiguration\",\n back_populates=\"llm_model_flows\",\n )\n\n __table_args__ = (\n UniqueConstraint(\n \"llm_model_flow_type\",\n \"model_configuration_id\",\n name=\"uq_model_config_per_llm_model_flow_type\",\n ),\n Index(\n \"ix_one_default_per_llm_model_flow\",\n \"llm_model_flow_type\",\n unique=True,\n postgresql_where=(is_default == True), # noqa: E712\n ),\n )\n\n\nclass ImageGenerationConfig(Base):\n __tablename__ = \"image_generation_config\"\n\n image_provider_id: Mapped[str] = mapped_column(String, primary_key=True)\n model_configuration_id: Mapped[int] = mapped_column(\n ForeignKey(\"model_configuration.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n is_default: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n model_configuration: Mapped[\"ModelConfiguration\"] = relationship(\n \"ModelConfiguration\"\n )\n\n __table_args__ = (\n Index(\"ix_image_generation_config_is_default\", \"is_default\"),\n Index(\n \"ix_image_generation_config_model_configuration_id\",\n \"model_configuration_id\",\n ),\n )\n\n\nclass VoiceProvider(Base):\n \"\"\"Configuration for voice services (STT and TTS).\"\"\"\n\n __tablename__ = \"voice_provider\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True)\n provider_type: Mapped[str] = mapped_column(\n String\n ) # \"openai\", \"azure\", \"elevenlabs\"\n api_key: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=True\n )\n api_base: Mapped[str | None] = mapped_column(String, nullable=True)\n custom_config: Mapped[dict[str, Any] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Model/voice configuration\n stt_model: Mapped[str | None] = mapped_column(\n String, nullable=True\n ) # e.g., \"whisper-1\"\n tts_model: Mapped[str | None] = mapped_column(\n String, nullable=True\n ) # e.g., \"tts-1\", \"tts-1-hd\"\n default_voice: Mapped[str | None] = mapped_column(\n String, nullable=True\n ) # e.g., \"alloy\", \"echo\"\n\n # STT and TTS can use different providers - only one provider per type\n is_default_stt: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n is_default_tts: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n # Enforce only one default STT provider and one default TTS provider at DB level\n __table_args__ = (\n Index(\n \"ix_voice_provider_one_default_stt\",\n \"is_default_stt\",\n unique=True,\n postgresql_where=(is_default_stt == True), # noqa: E712\n ),\n Index(\n \"ix_voice_provider_one_default_tts\",\n \"is_default_tts\",\n unique=True,\n postgresql_where=(is_default_tts == True), # noqa: E712\n ),\n )\n\n\nclass CloudEmbeddingProvider(Base):\n __tablename__ = \"embedding_provider\"\n\n provider_type: Mapped[EmbeddingProvider] = mapped_column(\n Enum(EmbeddingProvider), primary_key=True\n )\n api_url: Mapped[str | None] = mapped_column(String, nullable=True)\n api_key: Mapped[SensitiveValue[str] | None] = mapped_column(EncryptedString())\n api_version: Mapped[str | None] = mapped_column(String, nullable=True)\n deployment_name: Mapped[str | None] = mapped_column(String, nullable=True)\n\n search_settings: Mapped[list[\"SearchSettings\"]] = relationship(\n \"SearchSettings\",\n back_populates=\"cloud_provider\",\n )\n\n def __repr__(self) -> str:\n return f\"\"\n\n\nclass InternetSearchProvider(Base):\n __tablename__ = \"internet_search_provider\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True, nullable=False)\n provider_type: Mapped[str] = mapped_column(String, nullable=False)\n api_key: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=True\n )\n config: Mapped[dict[str, str] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n def __repr__(self) -> str:\n return f\"\"\n\n\nclass InternetContentProvider(Base):\n __tablename__ = \"internet_content_provider\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True, nullable=False)\n provider_type: Mapped[str] = mapped_column(String, nullable=False)\n api_key: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=True\n )\n config: Mapped[WebContentProviderConfig | None] = mapped_column(\n PydanticType(WebContentProviderConfig), nullable=True\n )\n is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n time_updated: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n def __repr__(self) -> str:\n return f\"\"\n\n\nclass DocumentSet(Base):\n __tablename__ = \"document_set\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True)\n description: Mapped[str | None] = mapped_column(String)\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # Whether changes to the document set have been propagated\n is_up_to_date: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n # If `False`, then the document set is not visible to users who are not explicitly\n # given access to it either via the `users` or `groups` relationships\n is_public: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n\n # Last time a user updated this document set\n time_last_modified_by_user: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n connector_credential_pairs: Mapped[list[ConnectorCredentialPair]] = relationship(\n \"ConnectorCredentialPair\",\n secondary=DocumentSet__ConnectorCredentialPair.__table__,\n primaryjoin=(\n (DocumentSet__ConnectorCredentialPair.document_set_id == id)\n & (DocumentSet__ConnectorCredentialPair.is_current.is_(True))\n ),\n secondaryjoin=(\n DocumentSet__ConnectorCredentialPair.connector_credential_pair_id\n == ConnectorCredentialPair.id\n ),\n back_populates=\"document_sets\",\n overlaps=\"document_set\",\n )\n personas: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=Persona__DocumentSet.__table__,\n back_populates=\"document_sets\",\n )\n # Other users with access\n users: Mapped[list[User]] = relationship(\n \"User\",\n secondary=DocumentSet__User.__table__,\n viewonly=True,\n )\n # EE only\n groups: Mapped[list[\"UserGroup\"]] = relationship(\n \"UserGroup\",\n secondary=\"document_set__user_group\",\n viewonly=True,\n )\n federated_connectors: Mapped[list[\"FederatedConnector__DocumentSet\"]] = (\n relationship(\n \"FederatedConnector__DocumentSet\",\n back_populates=\"document_set\",\n cascade=\"all, delete-orphan\",\n )\n )\n\n\nclass Tool(Base):\n __tablename__ = \"tool\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n # The name of the tool that the LLM will see\n name: Mapped[str] = mapped_column(String, nullable=False)\n description: Mapped[str] = mapped_column(Text, nullable=True)\n # ID of the tool in the codebase, only applies for in-code tools.\n # tools defined via the UI will have this as None\n in_code_tool_id: Mapped[str | None] = mapped_column(String, nullable=True)\n display_name: Mapped[str] = mapped_column(String, nullable=True)\n\n # OpenAPI scheme for the tool. Only applies to tools defined via the UI.\n openapi_schema: Mapped[dict[str, Any] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n # MCP tool input schema. Only applies to MCP tools.\n mcp_input_schema: Mapped[dict[str, Any] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n custom_headers: Mapped[list[HeaderItemDict] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n # user who created / owns the tool. Will be None for built-in tools.\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # whether to pass through the user's OAuth token as Authorization header\n passthrough_auth: Mapped[bool] = mapped_column(Boolean, default=False)\n # MCP server this tool is associated with (null for non-MCP tools)\n mcp_server_id: Mapped[int | None] = mapped_column(\n Integer, ForeignKey(\"mcp_server.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # OAuth configuration for this tool (null for tools without OAuth)\n oauth_config_id: Mapped[int | None] = mapped_column(\n Integer, ForeignKey(\"oauth_config.id\", ondelete=\"SET NULL\"), nullable=True\n )\n enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n\n user: Mapped[User | None] = relationship(\"User\", back_populates=\"custom_tools\")\n oauth_config: Mapped[\"OAuthConfig | None\"] = relationship(\n \"OAuthConfig\", back_populates=\"tools\"\n )\n # Relationship to Persona through the association table\n personas: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=Persona__Tool.__table__,\n back_populates=\"tools\",\n )\n # MCP server relationship\n mcp_server: Mapped[\"MCPServer | None\"] = relationship(\n \"MCPServer\", back_populates=\"current_actions\"\n )\n\n\nclass OAuthConfig(Base):\n \"\"\"OAuth provider configuration that can be shared across multiple tools\"\"\"\n\n __tablename__ = \"oauth_config\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True, nullable=False)\n\n # OAuth provider endpoints\n authorization_url: Mapped[str] = mapped_column(Text, nullable=False)\n token_url: Mapped[str] = mapped_column(Text, nullable=False)\n\n # Client credentials (encrypted)\n client_id: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=False\n )\n client_secret: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=False\n )\n\n # Optional configurations\n scopes: Mapped[list[str] | None] = mapped_column(postgresql.JSONB(), nullable=True)\n additional_params: Mapped[dict[str, Any] | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n\n # Metadata\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n # Relationships\n tools: Mapped[list[\"Tool\"]] = relationship(\"Tool\", back_populates=\"oauth_config\")\n user_tokens: Mapped[list[\"OAuthUserToken\"]] = relationship(\n \"OAuthUserToken\", back_populates=\"oauth_config\", cascade=\"all, delete-orphan\"\n )\n\n\nclass OAuthUserToken(Base):\n \"\"\"Per-user OAuth tokens for a specific OAuth configuration\"\"\"\n\n __tablename__ = \"oauth_user_token\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n oauth_config_id: Mapped[int] = mapped_column(\n ForeignKey(\"oauth_config.id\", ondelete=\"CASCADE\"), nullable=False\n )\n user_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=False\n )\n\n # Token data (encrypted)\n # Structure: {\n # \"access_token\": \"...\",\n # \"refresh_token\": \"...\", # Optional\n # \"token_type\": \"Bearer\",\n # \"expires_at\": 1234567890, # Unix timestamp, optional\n # \"scope\": \"repo user\" # Optional\n # }\n token_data: Mapped[SensitiveValue[dict[str, Any]] | None] = mapped_column(\n EncryptedJson(), nullable=False\n )\n\n # Metadata\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n # Relationships\n oauth_config: Mapped[\"OAuthConfig\"] = relationship(\n \"OAuthConfig\", back_populates=\"user_tokens\"\n )\n user: Mapped[\"User\"] = relationship(\"User\")\n\n # Unique constraint: One token per user per OAuth config\n __table_args__ = (\n UniqueConstraint(\"oauth_config_id\", \"user_id\", name=\"uq_oauth_user_token\"),\n )\n\n\nclass StarterMessage(BaseModel):\n \"\"\"Starter message for a persona.\"\"\"\n\n name: str\n message: str\n\n\nclass Persona__PersonaLabel(Base):\n __tablename__ = \"persona__persona_label\"\n\n persona_id: Mapped[int] = mapped_column(ForeignKey(\"persona.id\"), primary_key=True)\n persona_label_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona_label.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass Persona(Base):\n __tablename__ = \"persona\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n name: Mapped[str] = mapped_column(String)\n description: Mapped[str] = mapped_column(String)\n\n # Allows the persona to specify a specific default LLM model\n # NOTE: only is applied on the actual response generation - is not used for things like\n # auto-detected time filters, relevance filters, etc.\n llm_model_provider_override: Mapped[str | None] = mapped_column(\n String, nullable=True\n )\n llm_model_version_override: Mapped[str | None] = mapped_column(\n String, nullable=True\n )\n default_model_configuration_id: Mapped[int | None] = mapped_column(\n Integer,\n ForeignKey(\"model_configuration.id\", ondelete=\"SET NULL\"),\n nullable=True,\n )\n\n starter_messages: Mapped[list[StarterMessage] | None] = mapped_column(\n PydanticListType(StarterMessage), nullable=True\n )\n search_start_date: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n # Built-in personas are configured via backend during deployment\n # Treated specially (cannot be user edited etc.)\n builtin_persona: Mapped[bool] = mapped_column(Boolean, default=False)\n\n # Featured personas are highlighted in the UI\n is_featured: Mapped[bool] = mapped_column(Boolean, default=False)\n # controls whether the persona is listed in user-facing agent lists\n is_listed: Mapped[bool] = mapped_column(Boolean, default=True)\n # controls the ordering of personas in the UI\n # higher priority personas are displayed first, ties are resolved by the ID,\n # where lower value IDs (e.g. created earlier) are displayed first\n display_priority: Mapped[int | None] = mapped_column(\n Integer, nullable=True, default=None\n )\n deleted: Mapped[bool] = mapped_column(Boolean, default=False)\n\n # Custom Agent Prompt\n system_prompt: Mapped[str | None] = mapped_column(\n String(length=PROMPT_LENGTH), nullable=True\n )\n replace_base_system_prompt: Mapped[bool] = mapped_column(Boolean, default=False)\n task_prompt: Mapped[str | None] = mapped_column(\n String(length=PROMPT_LENGTH), nullable=True\n )\n datetime_aware: Mapped[bool] = mapped_column(Boolean, default=True)\n\n uploaded_image_id: Mapped[str | None] = mapped_column(String, nullable=True)\n icon_name: Mapped[str | None] = mapped_column(String, nullable=True)\n\n # These are only defaults, users can select from all if desired\n document_sets: Mapped[list[DocumentSet]] = relationship(\n \"DocumentSet\",\n secondary=Persona__DocumentSet.__table__,\n back_populates=\"personas\",\n )\n tools: Mapped[list[Tool]] = relationship(\n \"Tool\",\n secondary=Persona__Tool.__table__,\n back_populates=\"personas\",\n )\n # Owner\n user: Mapped[User | None] = relationship(\"User\", back_populates=\"personas\")\n # Other users with access\n users: Mapped[list[User]] = relationship(\n \"User\",\n secondary=Persona__User.__table__,\n viewonly=True,\n )\n # EE only\n is_public: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n groups: Mapped[list[\"UserGroup\"]] = relationship(\n \"UserGroup\",\n secondary=\"persona__user_group\",\n viewonly=True,\n )\n allowed_by_llm_providers: Mapped[list[\"LLMProvider\"]] = relationship(\n \"LLMProvider\",\n secondary=\"llm_provider__persona\",\n back_populates=\"personas\",\n viewonly=True,\n )\n # Relationship to UserFile\n user_files: Mapped[list[\"UserFile\"]] = relationship(\n \"UserFile\",\n secondary=\"persona__user_file\",\n back_populates=\"assistants\",\n )\n labels: Mapped[list[\"PersonaLabel\"]] = relationship(\n \"PersonaLabel\",\n secondary=Persona__PersonaLabel.__table__,\n back_populates=\"personas\",\n )\n # Hierarchy nodes attached to this persona for scoped search\n hierarchy_nodes: Mapped[list[\"HierarchyNode\"]] = relationship(\n \"HierarchyNode\",\n secondary=\"persona__hierarchy_node\",\n back_populates=\"personas\",\n )\n # Individual documents attached to this persona for scoped search\n attached_documents: Mapped[list[\"Document\"]] = relationship(\n \"Document\",\n secondary=\"persona__document\",\n back_populates=\"attached_personas\",\n )\n\n # Default personas loaded via yaml cannot have the same name\n __table_args__ = (\n Index(\n \"_builtin_persona_name_idx\",\n \"name\",\n unique=True,\n postgresql_where=(builtin_persona == True), # noqa: E712\n ),\n )\n\n\nclass Persona__UserFile(Base):\n __tablename__ = \"persona__user_file\"\n\n persona_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n user_file_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user_file.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass Persona__HierarchyNode(Base):\n \"\"\"Association table linking personas to hierarchy nodes.\n\n This allows assistants to be configured with specific hierarchy nodes\n (folders, spaces, channels, etc.) for scoped search/retrieval.\n \"\"\"\n\n __tablename__ = \"persona__hierarchy_node\"\n\n persona_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n hierarchy_node_id: Mapped[int] = mapped_column(\n ForeignKey(\"hierarchy_node.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass Persona__Document(Base):\n \"\"\"Association table linking personas to individual documents.\n\n This allows assistants to be configured with specific documents\n for scoped search/retrieval. Complements hierarchy_nodes which\n allow attaching folders/spaces.\n \"\"\"\n\n __tablename__ = \"persona__document\"\n\n persona_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n document_id: Mapped[str] = mapped_column(\n ForeignKey(\"document.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass PersonaLabel(Base):\n __tablename__ = \"persona_label\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True)\n personas: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=Persona__PersonaLabel.__table__,\n back_populates=\"labels\",\n )\n\n\nclass Assistant__UserSpecificConfig(Base):\n __tablename__ = \"assistant__user_specific_config\"\n\n assistant_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n user_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n disabled_tool_ids: Mapped[list[int]] = mapped_column(\n postgresql.ARRAY(Integer), nullable=False\n )\n\n\nAllowedAnswerFilters = (\n Literal[\"well_answered_postfilter\"] | Literal[\"questionmark_prefilter\"]\n)\n\n\nclass ChannelConfig(TypedDict):\n \"\"\"NOTE: is a `TypedDict` so it can be used as a type hint for a JSONB column\n in Postgres\"\"\"\n\n channel_name: str | None # None for default channel config\n respond_tag_only: NotRequired[bool] # defaults to False\n respond_to_bots: NotRequired[bool] # defaults to False\n is_ephemeral: NotRequired[bool] # defaults to False\n respond_member_group_list: NotRequired[list[str]]\n answer_filters: NotRequired[list[AllowedAnswerFilters]]\n # If None then no follow up\n # If empty list, follow up with no tags\n follow_up_tags: NotRequired[list[str]]\n show_continue_in_web_ui: NotRequired[bool] # defaults to False\n disabled: NotRequired[bool] # defaults to False\n\n\nclass SlackChannelConfig(Base):\n __tablename__ = \"slack_channel_config\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n slack_bot_id: Mapped[int] = mapped_column(\n ForeignKey(\"slack_bot.id\"), nullable=False\n )\n persona_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"persona.id\"), nullable=True\n )\n channel_config: Mapped[ChannelConfig] = mapped_column(\n postgresql.JSONB(), nullable=False\n )\n\n enable_auto_filters: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False\n )\n\n is_default: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n persona: Mapped[Persona | None] = relationship(\"Persona\")\n\n slack_bot: Mapped[\"SlackBot\"] = relationship(\n \"SlackBot\",\n back_populates=\"slack_channel_configs\",\n )\n standard_answer_categories: Mapped[list[\"StandardAnswerCategory\"]] = relationship(\n \"StandardAnswerCategory\",\n secondary=SlackChannelConfig__StandardAnswerCategory.__table__,\n back_populates=\"slack_channel_configs\",\n )\n\n __table_args__ = (\n UniqueConstraint(\n \"slack_bot_id\",\n \"is_default\",\n name=\"uq_slack_channel_config_slack_bot_id_default\",\n ),\n Index(\n \"ix_slack_channel_config_slack_bot_id_default\",\n \"slack_bot_id\",\n \"is_default\",\n unique=True,\n postgresql_where=(is_default is True),\n ),\n )\n\n\nclass SlackBot(Base):\n __tablename__ = \"slack_bot\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n name: Mapped[str] = mapped_column(String)\n enabled: Mapped[bool] = mapped_column(Boolean, default=True)\n\n bot_token: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), unique=True\n )\n app_token: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), unique=True\n )\n user_token: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=True\n )\n\n slack_channel_configs: Mapped[list[SlackChannelConfig]] = relationship(\n \"SlackChannelConfig\",\n back_populates=\"slack_bot\",\n cascade=\"all, delete-orphan\",\n )\n\n\nclass DiscordBotConfig(Base):\n \"\"\"Global Discord bot configuration (one per tenant).\n\n Stores the bot token when not provided via DISCORD_BOT_TOKEN env var.\n Uses a fixed ID with check constraint to enforce only one row per tenant.\n \"\"\"\n\n __tablename__ = \"discord_bot_config\"\n\n id: Mapped[str] = mapped_column(\n String, primary_key=True, server_default=text(\"'SINGLETON'\")\n )\n bot_token: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=False\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n\n\nclass DiscordGuildConfig(Base):\n \"\"\"Configuration for a Discord guild (server) connected to this tenant.\n\n registration_key is a one-time key used to link a Discord server to this tenant.\n Format: discord_.\n guild_id is NULL until the Discord admin runs !register with the key.\n \"\"\"\n\n __tablename__ = \"discord_guild_config\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n # Discord snowflake - NULL until registered via command in Discord\n guild_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True, unique=True)\n guild_name: Mapped[str | None] = mapped_column(String(256), nullable=True)\n\n # One-time registration key: discord_.\n registration_key: Mapped[str] = mapped_column(String, unique=True, nullable=False)\n\n registered_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n # Configuration\n default_persona_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"SET NULL\"), nullable=True\n )\n enabled: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"true\"), nullable=False\n )\n\n # Relationships\n default_persona: Mapped[\"Persona | None\"] = relationship(\n \"Persona\", foreign_keys=[default_persona_id]\n )\n channels: Mapped[list[\"DiscordChannelConfig\"]] = relationship(\n back_populates=\"guild_config\", cascade=\"all, delete-orphan\"\n )\n\n\nclass DiscordChannelConfig(Base):\n \"\"\"Per-channel configuration for Discord bot behavior.\n\n Used to whitelist specific channels and configure per-channel behavior.\n \"\"\"\n\n __tablename__ = \"discord_channel_config\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n guild_config_id: Mapped[int] = mapped_column(\n ForeignKey(\"discord_guild_config.id\", ondelete=\"CASCADE\"), nullable=False\n )\n\n # Discord snowflake\n channel_id: Mapped[int] = mapped_column(BigInteger, nullable=False)\n channel_name: Mapped[str] = mapped_column(String(), nullable=False)\n\n # Channel type from Discord (text, forum)\n channel_type: Mapped[str] = mapped_column(\n String(20), server_default=text(\"'text'\"), nullable=False\n )\n\n # True if @everyone cannot view the channel\n is_private: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"false\"), nullable=False\n )\n\n # If true, bot only responds to messages in threads\n # Otherwise, will reply in channel\n thread_only_mode: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"false\"), nullable=False\n )\n\n # If true (default), bot only responds when @mentioned\n # If false, bot responds to ALL messages in this channel\n require_bot_invocation: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"true\"), nullable=False\n )\n\n # Override the guild's default persona for this channel\n persona_override_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"SET NULL\"), nullable=True\n )\n\n enabled: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"false\"), nullable=False\n )\n\n # Relationships\n guild_config: Mapped[\"DiscordGuildConfig\"] = relationship(back_populates=\"channels\")\n persona_override: Mapped[\"Persona | None\"] = relationship()\n\n # Constraints\n __table_args__ = (\n UniqueConstraint(\n \"guild_config_id\", \"channel_id\", name=\"uq_discord_channel_guild_channel\"\n ),\n )\n\n\nclass Milestone(Base):\n # This table is used to track significant events for a deployment towards finding value\n # The table is currently not used for features but it may be used in the future to inform\n # users about the product features and encourage usage/exploration.\n __tablename__ = \"milestone\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n event_type: Mapped[MilestoneRecordType] = mapped_column(String)\n # Need to track counts and specific ids of certain events to know if the Milestone has been reached\n event_tracker: Mapped[dict | None] = mapped_column(\n postgresql.JSONB(), nullable=True\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n user: Mapped[User | None] = relationship(\"User\")\n\n __table_args__ = (UniqueConstraint(\"event_type\", name=\"uq_milestone_event_type\"),)\n\n\nclass TaskQueueState(Base):\n # Currently refers to Celery Tasks\n __tablename__ = \"task_queue_jobs\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n # Celery task id. currently only for readability/diagnostics\n task_id: Mapped[str] = mapped_column(String)\n # For any job type, this would be the same\n task_name: Mapped[str] = mapped_column(String)\n # Note that if the task dies, this won't necessarily be marked FAILED correctly\n status: Mapped[TaskStatus] = mapped_column(Enum(TaskStatus, native_enum=False))\n start_time: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True)\n )\n register_time: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n\nclass KVStore(Base):\n __tablename__ = \"key_value_store\"\n\n key: Mapped[str] = mapped_column(String, primary_key=True)\n value: Mapped[JSON_ro] = mapped_column(postgresql.JSONB(), nullable=True)\n encrypted_value: Mapped[SensitiveValue[dict[str, Any]] | None] = mapped_column(\n EncryptedJson(), nullable=True\n )\n\n\nclass FileRecord(Base):\n __tablename__ = \"file_record\"\n\n # Internal file ID, must be unique across all files.\n file_id: Mapped[str] = mapped_column(String, primary_key=True)\n\n display_name: Mapped[str] = mapped_column(String, nullable=True)\n file_origin: Mapped[FileOrigin] = mapped_column(Enum(FileOrigin, native_enum=False))\n file_type: Mapped[str] = mapped_column(String, default=\"text/plain\")\n file_metadata: Mapped[JSON_ro] = mapped_column(postgresql.JSONB(), nullable=True)\n\n # External storage support (S3, MinIO, Azure Blob, etc.)\n bucket_name: Mapped[str] = mapped_column(String)\n object_key: Mapped[str] = mapped_column(String)\n\n # Timestamps for external storage\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n\nclass FileContent(Base):\n \"\"\"Stores file content in PostgreSQL using Large Objects.\n Used when FILE_STORE_BACKEND=postgres to avoid needing S3/MinIO.\"\"\"\n\n __tablename__ = \"file_content\"\n\n file_id: Mapped[str] = mapped_column(\n String,\n ForeignKey(\"file_record.file_id\", ondelete=\"CASCADE\"),\n primary_key=True,\n )\n # PostgreSQL Large Object OID referencing pg_largeobject\n lobj_oid: Mapped[int] = mapped_column(BigInteger, nullable=False)\n file_size: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)\n\n\n\"\"\"\n************************************************************************\nEnterprise Edition Models\n************************************************************************\n\nThese models are only used in Enterprise Edition only features in Onyx.\nThey are kept here to simplify the codebase and avoid having different assumptions\non the shape of data being passed around between the MIT and EE versions of Onyx.\n\nIn the MIT version of Onyx, assume these tables are always empty.\n\"\"\"\n\n\nclass SamlAccount(Base):\n __tablename__ = \"saml\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n user_id: Mapped[int] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), unique=True\n )\n encrypted_cookie: Mapped[str] = mapped_column(Text, unique=True)\n expires_at: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n user: Mapped[User] = relationship(\"User\")\n\n\nclass User__UserGroup(Base):\n __tablename__ = \"user__user_group\"\n\n __table_args__ = (Index(\"ix_user__user_group_user_id\", \"user_id\"),)\n\n is_curator: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), primary_key=True, nullable=True\n )\n\n\nclass PermissionGrant(Base):\n __tablename__ = \"permission_grant\"\n\n __table_args__ = (\n UniqueConstraint(\n \"group_id\", \"permission\", name=\"uq_permission_grant_group_permission\"\n ),\n )\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)\n group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\", ondelete=\"CASCADE\"), nullable=False\n )\n permission: Mapped[Permission] = mapped_column(\n Enum(\n Permission,\n native_enum=False,\n values_callable=lambda x: [e.value for e in x],\n ),\n nullable=False,\n )\n grant_source: Mapped[GrantSource] = mapped_column(\n Enum(GrantSource, native_enum=False), nullable=False\n )\n granted_by: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"SET NULL\"), nullable=True\n )\n granted_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n is_deleted: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False, server_default=text(\"false\")\n )\n\n group: Mapped[\"UserGroup\"] = relationship(\n \"UserGroup\", back_populates=\"permission_grants\"\n )\n\n @validates(\"permission\")\n def _validate_permission(self, _key: str, value: Permission) -> Permission:\n if value in Permission.IMPLIED:\n raise ValueError(\n f\"{value!r} is an implied permission and cannot be granted directly\"\n )\n return value\n\n\nclass UserGroup__ConnectorCredentialPair(Base):\n __tablename__ = \"user_group__connector_credential_pair\"\n\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n cc_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"), primary_key=True\n )\n # if `True`, then is part of the current state of the UserGroup\n # if `False`, then is a part of the prior state of the UserGroup\n # rows with `is_current=False` should be deleted when the UserGroup\n # is updated and should not exist for a given UserGroup if\n # `UserGroup.is_up_to_date == True`\n is_current: Mapped[bool] = mapped_column(\n Boolean,\n default=True,\n primary_key=True,\n )\n\n cc_pair: Mapped[ConnectorCredentialPair] = relationship(\n \"ConnectorCredentialPair\",\n )\n\n\nclass Persona__UserGroup(Base):\n __tablename__ = \"persona__user_group\"\n\n persona_id: Mapped[int] = mapped_column(ForeignKey(\"persona.id\"), primary_key=True)\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n\n\nclass LLMProvider__Persona(Base):\n \"\"\"Association table restricting LLM providers to specific personas.\n\n If no such rows exist for a given LLM provider, then it is accessible by all personas.\n \"\"\"\n\n __tablename__ = \"llm_provider__persona\"\n\n llm_provider_id: Mapped[int] = mapped_column(\n ForeignKey(\"llm_provider.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n persona_id: Mapped[int] = mapped_column(\n ForeignKey(\"persona.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass LLMProvider__UserGroup(Base):\n __tablename__ = \"llm_provider__user_group\"\n\n llm_provider_id: Mapped[int] = mapped_column(\n ForeignKey(\"llm_provider.id\"), primary_key=True\n )\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n\n\nclass DocumentSet__UserGroup(Base):\n __tablename__ = \"document_set__user_group\"\n\n document_set_id: Mapped[int] = mapped_column(\n ForeignKey(\"document_set.id\"), primary_key=True\n )\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n\n\nclass Credential__UserGroup(Base):\n __tablename__ = \"credential__user_group\"\n\n credential_id: Mapped[int] = mapped_column(\n ForeignKey(\"credential.id\"), primary_key=True\n )\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n\n\nclass UserGroup(Base):\n __tablename__ = \"user_group\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True)\n # whether or not changes to the UserGroup have been propagated to Vespa\n is_up_to_date: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n # tell the sync job to clean up the group\n is_up_for_deletion: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False\n )\n # whether this is a default group (e.g. \"Basic\", \"Admins\") that cannot be deleted\n is_default: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n # Last time a user updated this user group\n time_last_modified_by_user: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n users: Mapped[list[User]] = relationship(\n \"User\",\n secondary=User__UserGroup.__table__,\n )\n user_group_relationships: Mapped[list[User__UserGroup]] = relationship(\n \"User__UserGroup\",\n viewonly=True,\n )\n cc_pairs: Mapped[list[ConnectorCredentialPair]] = relationship(\n \"ConnectorCredentialPair\",\n secondary=UserGroup__ConnectorCredentialPair.__table__,\n viewonly=True,\n )\n cc_pair_relationships: Mapped[list[UserGroup__ConnectorCredentialPair]] = (\n relationship(\n \"UserGroup__ConnectorCredentialPair\",\n viewonly=True,\n )\n )\n personas: Mapped[list[Persona]] = relationship(\n \"Persona\",\n secondary=Persona__UserGroup.__table__,\n viewonly=True,\n )\n document_sets: Mapped[list[DocumentSet]] = relationship(\n \"DocumentSet\",\n secondary=DocumentSet__UserGroup.__table__,\n viewonly=True,\n )\n credentials: Mapped[list[Credential]] = relationship(\n \"Credential\",\n secondary=Credential__UserGroup.__table__,\n )\n # MCP servers accessible to this user group\n accessible_mcp_servers: Mapped[list[\"MCPServer\"]] = relationship(\n \"MCPServer\", secondary=\"mcp_server__user_group\", back_populates=\"user_groups\"\n )\n permission_grants: Mapped[list[\"PermissionGrant\"]] = relationship(\n \"PermissionGrant\", back_populates=\"group\", cascade=\"all, delete-orphan\"\n )\n\n\n\"\"\"Tables related to Token Rate Limiting\nNOTE: `TokenRateLimit` is partially an MIT feature (global rate limit)\n\"\"\"\n\n\nclass TokenRateLimit(Base):\n __tablename__ = \"token_rate_limit\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n token_budget: Mapped[int] = mapped_column(Integer, nullable=False)\n period_hours: Mapped[int] = mapped_column(Integer, nullable=False)\n scope: Mapped[TokenRateLimitScope] = mapped_column(\n Enum(TokenRateLimitScope, native_enum=False)\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n\nclass TokenRateLimit__UserGroup(Base):\n __tablename__ = \"token_rate_limit__user_group\"\n\n rate_limit_id: Mapped[int] = mapped_column(\n ForeignKey(\"token_rate_limit.id\"), primary_key=True\n )\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n\n\nclass StandardAnswerCategory(Base):\n __tablename__ = \"standard_answer_category\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n name: Mapped[str] = mapped_column(String, unique=True)\n standard_answers: Mapped[list[\"StandardAnswer\"]] = relationship(\n \"StandardAnswer\",\n secondary=StandardAnswer__StandardAnswerCategory.__table__,\n back_populates=\"categories\",\n )\n slack_channel_configs: Mapped[list[\"SlackChannelConfig\"]] = relationship(\n \"SlackChannelConfig\",\n secondary=SlackChannelConfig__StandardAnswerCategory.__table__,\n back_populates=\"standard_answer_categories\",\n )\n\n\nclass StandardAnswer(Base):\n __tablename__ = \"standard_answer\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n keyword: Mapped[str] = mapped_column(String)\n answer: Mapped[str] = mapped_column(String)\n active: Mapped[bool] = mapped_column(Boolean)\n match_regex: Mapped[bool] = mapped_column(Boolean)\n match_any_keywords: Mapped[bool] = mapped_column(Boolean)\n\n __table_args__ = (\n Index(\n \"unique_keyword_active\",\n keyword,\n active,\n unique=True,\n postgresql_where=(active == True), # noqa: E712\n ),\n )\n\n categories: Mapped[list[StandardAnswerCategory]] = relationship(\n \"StandardAnswerCategory\",\n secondary=StandardAnswer__StandardAnswerCategory.__table__,\n back_populates=\"standard_answers\",\n )\n chat_messages: Mapped[list[ChatMessage]] = relationship(\n \"ChatMessage\",\n secondary=ChatMessage__StandardAnswer.__table__,\n back_populates=\"standard_answers\",\n )\n\n\nclass BackgroundError(Base):\n \"\"\"Important background errors. Serves to:\n 1. Ensure that important logs are kept around and not lost on rotation/container restarts\n 2. A trail for high-signal events so that the debugger doesn't need to remember/know every\n possible relevant log line.\n \"\"\"\n\n __tablename__ = \"background_error\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n message: Mapped[str] = mapped_column(String)\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n # option to link the error to a specific CC Pair\n cc_pair_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\", ondelete=\"CASCADE\"), nullable=True\n )\n\n cc_pair: Mapped[\"ConnectorCredentialPair | None\"] = relationship(\n \"ConnectorCredentialPair\", back_populates=\"background_errors\"\n )\n\n\n\"\"\"Tables related to Permission Sync\"\"\"\n\n\nclass User__ExternalUserGroupId(Base):\n \"\"\"Maps user info both internal and external to the name of the external group\n This maps the user to all of their external groups so that the external group name can be\n attached to the ACL list matching during query time. User level permissions can be handled by\n directly adding the Onyx user to the doc ACL list\"\"\"\n\n __tablename__ = \"user__external_user_group_id\"\n\n user_id: Mapped[UUID] = mapped_column(ForeignKey(\"user.id\"), primary_key=True)\n # These group ids have been prefixed by the source type\n external_user_group_id: Mapped[str] = mapped_column(String, primary_key=True)\n cc_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"), primary_key=True\n )\n\n # Signifies whether or not the group should be cleaned up at the end of a\n # group sync run.\n stale: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n __table_args__ = (\n Index(\n \"ix_user_external_group_cc_pair_stale\",\n \"cc_pair_id\",\n \"stale\",\n ),\n Index(\n \"ix_user_external_group_stale\",\n \"stale\",\n ),\n )\n\n\nclass PublicExternalUserGroup(Base):\n \"\"\"Stores all public external user \"groups\".\n\n For example, things like Google Drive folders that are marked\n as `Anyone with the link` or `Anyone in the domain`\n \"\"\"\n\n __tablename__ = \"public_external_user_group\"\n\n external_user_group_id: Mapped[str] = mapped_column(String, primary_key=True)\n cc_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n # Signifies whether or not the group should be cleaned up at the end of a\n # group sync run.\n stale: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n __table_args__ = (\n Index(\n \"ix_public_external_group_cc_pair_stale\",\n \"cc_pair_id\",\n \"stale\",\n ),\n Index(\n \"ix_public_external_group_stale\",\n \"stale\",\n ),\n )\n\n\nclass UsageReport(Base):\n \"\"\"This stores metadata about usage reports generated by admin including user who generated\n them as well as the period they cover. The actual zip file of the report is stored as a lo\n using the FileRecord\n \"\"\"\n\n __tablename__ = \"usage_reports\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n report_name: Mapped[str] = mapped_column(ForeignKey(\"file_record.file_id\"))\n\n # if None, report was auto-generated\n requestor_user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n period_from: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True)\n )\n period_to: Mapped[datetime.datetime | None] = mapped_column(DateTime(timezone=True))\n\n requestor = relationship(\"User\")\n file = relationship(\"FileRecord\")\n\n\nclass InputPrompt(Base):\n __tablename__ = \"inputprompt\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)\n prompt: Mapped[str] = mapped_column(String)\n content: Mapped[str] = mapped_column(String)\n active: Mapped[bool] = mapped_column(Boolean)\n user: Mapped[User | None] = relationship(\"User\", back_populates=\"input_prompts\")\n is_public: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n\n __table_args__ = (\n # Unique constraint on (prompt, user_id) for user-owned prompts\n UniqueConstraint(\"prompt\", \"user_id\", name=\"uq_inputprompt_prompt_user_id\"),\n # Partial unique index for public prompts (user_id IS NULL)\n Index(\n \"uq_inputprompt_prompt_public\",\n \"prompt\",\n unique=True,\n postgresql_where=text(\"user_id IS NULL\"),\n ),\n )\n\n\nclass InputPrompt__User(Base):\n __tablename__ = \"inputprompt__user\"\n\n input_prompt_id: Mapped[int] = mapped_column(\n ForeignKey(\"inputprompt.id\"), primary_key=True\n )\n user_id: Mapped[UUID | None] = mapped_column(\n ForeignKey(\"user.id\"), primary_key=True\n )\n disabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n\n\nclass Project__UserFile(Base):\n __tablename__ = \"project__user_file\"\n\n project_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_project.id\"), primary_key=True\n )\n user_file_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user_file.id\"), primary_key=True\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n\n __table_args__ = (\n Index(\n \"ix_project__user_file_project_id_created_at\",\n project_id,\n created_at.desc(),\n ),\n )\n\n\nclass UserProject(Base):\n __tablename__ = \"user_project\"\n\n id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)\n user_id: Mapped[UUID | None] = mapped_column(ForeignKey(\"user.id\"), nullable=False)\n name: Mapped[str] = mapped_column(nullable=False)\n description: Mapped[str] = mapped_column(nullable=True)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n user: Mapped[\"User\"] = relationship(back_populates=\"projects\")\n user_files: Mapped[list[\"UserFile\"]] = relationship(\n \"UserFile\",\n secondary=Project__UserFile.__table__,\n back_populates=\"projects\",\n )\n chat_sessions: Mapped[list[\"ChatSession\"]] = relationship(\n \"ChatSession\", back_populates=\"project\", lazy=\"selectin\"\n )\n instructions: Mapped[str] = mapped_column(String)\n\n\nclass UserDocument(str, Enum):\n CHAT = \"chat\"\n RECENT = \"recent\"\n FILE = \"file\"\n\n\nclass UserFile(Base):\n __tablename__ = \"user_file\"\n\n id: Mapped[UUID] = mapped_column(PGUUID(as_uuid=True), primary_key=True)\n user_id: Mapped[UUID | None] = mapped_column(ForeignKey(\"user.id\"), nullable=False)\n assistants: Mapped[list[\"Persona\"]] = relationship(\n \"Persona\",\n secondary=Persona__UserFile.__table__,\n back_populates=\"user_files\",\n )\n file_id: Mapped[str] = mapped_column(nullable=False)\n name: Mapped[str] = mapped_column(nullable=False)\n created_at: Mapped[datetime.datetime] = mapped_column(\n default=datetime.datetime.utcnow\n )\n user: Mapped[\"User\"] = relationship(back_populates=\"files\")\n token_count: Mapped[int | None] = mapped_column(Integer, nullable=True)\n\n file_type: Mapped[str] = mapped_column(String, nullable=False)\n\n status: Mapped[UserFileStatus] = mapped_column(\n Enum(UserFileStatus, native_enum=False, name=\"userfilestatus\"),\n nullable=False,\n default=UserFileStatus.PROCESSING,\n )\n needs_project_sync: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False\n )\n needs_persona_sync: Mapped[bool] = mapped_column(\n Boolean, nullable=False, default=False\n )\n last_project_sync_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n chunk_count: Mapped[int | None] = mapped_column(Integer, nullable=True)\n last_accessed_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n link_url: Mapped[str | None] = mapped_column(String, nullable=True)\n content_type: Mapped[str | None] = mapped_column(String, nullable=True)\n\n projects: Mapped[list[\"UserProject\"]] = relationship(\n \"UserProject\",\n secondary=Project__UserFile.__table__,\n back_populates=\"user_files\",\n lazy=\"selectin\",\n )\n\n\n\"\"\"\nMulti-tenancy related tables\n\"\"\"\n\n\nclass PublicBase(DeclarativeBase):\n __abstract__ = True\n\n\n# Strictly keeps track of the tenant that a given user will authenticate to.\nclass UserTenantMapping(Base):\n __tablename__ = \"user_tenant_mapping\"\n __table_args__ = ({\"schema\": \"public\"},)\n\n email: Mapped[str] = mapped_column(String, nullable=False, primary_key=True)\n tenant_id: Mapped[str] = mapped_column(String, nullable=False, primary_key=True)\n active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n\n @validates(\"email\")\n def validate_email(self, key: str, value: str) -> str: # noqa: ARG002\n return value.lower() if value else value\n\n\nclass AvailableTenant(Base):\n __tablename__ = \"available_tenant\"\n \"\"\"\n These entries will only exist ephemerally and are meant to be picked up by new users on registration.\n \"\"\"\n\n tenant_id: Mapped[str] = mapped_column(String, primary_key=True, nullable=False)\n alembic_version: Mapped[str] = mapped_column(String, nullable=False)\n date_created: Mapped[datetime.datetime] = mapped_column(DateTime, nullable=False)\n\n\n# This is a mapping from tenant IDs to anonymous user paths\nclass TenantAnonymousUserPath(Base):\n __tablename__ = \"tenant_anonymous_user_path\"\n\n tenant_id: Mapped[str] = mapped_column(String, primary_key=True, nullable=False)\n anonymous_user_path: Mapped[str] = mapped_column(\n String, nullable=False, unique=True\n )\n\n\nclass MCPServer(Base):\n \"\"\"Model for storing MCP server configurations\"\"\"\n\n __tablename__ = \"mcp_server\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n # Owner email of user who configured this server\n owner: Mapped[str] = mapped_column(String, nullable=False)\n name: Mapped[str] = mapped_column(String, nullable=False)\n description: Mapped[str | None] = mapped_column(String, nullable=True)\n server_url: Mapped[str] = mapped_column(String, nullable=False)\n # Transport type for connecting to the MCP server\n transport: Mapped[MCPTransport | None] = mapped_column(\n Enum(MCPTransport, native_enum=False), nullable=True\n )\n # Auth type: \"none\", \"api_token\", or \"oauth\"\n auth_type: Mapped[MCPAuthenticationType | None] = mapped_column(\n Enum(MCPAuthenticationType, native_enum=False), nullable=True\n )\n # Who performs authentication for this server (ADMIN or PER_USER)\n auth_performer: Mapped[MCPAuthenticationPerformer | None] = mapped_column(\n Enum(MCPAuthenticationPerformer, native_enum=False), nullable=True\n )\n # Status tracking for configuration flow\n status: Mapped[MCPServerStatus] = mapped_column(\n Enum(MCPServerStatus, native_enum=False),\n nullable=False,\n server_default=\"CREATED\",\n )\n # Admin connection config - used for the config page\n # and (when applicable) admin-managed auth\n # and (when applicable) per-user auth\n admin_connection_config_id: Mapped[int | None] = mapped_column(\n Integer,\n ForeignKey(\"mcp_connection_config.id\", ondelete=\"SET NULL\"),\n nullable=True,\n )\n\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n last_refreshed_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n # Relationships\n admin_connection_config: Mapped[\"MCPConnectionConfig | None\"] = relationship(\n \"MCPConnectionConfig\",\n foreign_keys=[admin_connection_config_id],\n back_populates=\"admin_servers\",\n )\n\n user_connection_configs: Mapped[list[\"MCPConnectionConfig\"]] = relationship(\n \"MCPConnectionConfig\",\n foreign_keys=\"MCPConnectionConfig.mcp_server_id\",\n back_populates=\"mcp_server\",\n passive_deletes=True,\n )\n current_actions: Mapped[list[\"Tool\"]] = relationship(\n \"Tool\", back_populates=\"mcp_server\", cascade=\"all, delete-orphan\"\n )\n # Many-to-many relationships for access control\n users: Mapped[list[\"User\"]] = relationship(\n \"User\", secondary=\"mcp_server__user\", back_populates=\"accessible_mcp_servers\"\n )\n user_groups: Mapped[list[\"UserGroup\"]] = relationship(\n \"UserGroup\",\n secondary=\"mcp_server__user_group\",\n back_populates=\"accessible_mcp_servers\",\n )\n\n\nclass MCPServer__User(Base):\n __tablename__ = \"mcp_server__user\"\n mcp_server_id: Mapped[int] = mapped_column(\n ForeignKey(\"mcp_server.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n user_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), primary_key=True\n )\n\n\nclass MCPServer__UserGroup(Base):\n __tablename__ = \"mcp_server__user_group\"\n mcp_server_id: Mapped[int] = mapped_column(\n ForeignKey(\"mcp_server.id\"), primary_key=True\n )\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\"), primary_key=True\n )\n\n\nclass MCPConnectionConfig(Base):\n \"\"\"Model for storing MCP connection configurations (credentials, auth data)\"\"\"\n\n __tablename__ = \"mcp_connection_config\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n # Server this config is for (nullable for template configs)\n mcp_server_id: Mapped[int | None] = mapped_column(\n Integer, ForeignKey(\"mcp_server.id\", ondelete=\"CASCADE\"), nullable=True\n )\n # User email this config is for (empty for admin configs and templates)\n user_email: Mapped[str] = mapped_column(String, nullable=False, default=\"\")\n # Config data stored as JSON\n # Format: {\n # \"refresh_token\": \"\", # OAuth only\n # \"access_token\": \"\", # OAuth only\n # \"headers\": {\"key\": \"value\", \"key2\": \"value2\"},\n # \"header_substitutions\": {\"\": \"\"}, # stored header template substitutions\n # \"request_body\": [\"path/in/body:value\", \"path2/in2/body2:value2\"] # TBD\n # \"client_id\": \"\", # For dynamically registered OAuth clients\n # \"client_secret\": \"\", # For confidential clients\n # \"registration_access_token\": \"\", # For managing registration\n # \"registration_client_uri\": \"\", # For managing registration\n # }\n config: Mapped[SensitiveValue[dict[str, Any]] | None] = mapped_column(\n EncryptedJson(), nullable=False, default=dict\n )\n\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n # Relationships\n mcp_server: Mapped[\"MCPServer | None\"] = relationship(\n \"MCPServer\",\n foreign_keys=[mcp_server_id],\n back_populates=\"user_connection_configs\",\n )\n admin_servers: Mapped[list[\"MCPServer\"]] = relationship(\n \"MCPServer\",\n foreign_keys=\"MCPServer.admin_connection_config_id\",\n back_populates=\"admin_connection_config\",\n )\n\n __table_args__ = (\n Index(\"ix_mcp_connection_config_user_email\", \"user_email\"),\n Index(\"ix_mcp_connection_config_server_user\", \"mcp_server_id\", \"user_email\"),\n )\n\n\n\"\"\"\nPermission Sync Tables\n\"\"\"\n\n\nclass DocPermissionSyncAttempt(Base):\n \"\"\"\n Represents an attempt to sync document permissions for a connector credential pair.\n Similar to IndexAttempt but specifically for document permission syncing operations.\n \"\"\"\n\n __tablename__ = \"doc_permission_sync_attempt\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n connector_credential_pair_id: Mapped[int] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"),\n nullable=False,\n )\n\n # Status of the sync attempt\n status: Mapped[PermissionSyncStatus] = mapped_column(\n Enum(PermissionSyncStatus, native_enum=False, index=True)\n )\n\n # Counts for tracking progress\n total_docs_synced: Mapped[int | None] = mapped_column(Integer, default=0)\n docs_with_permission_errors: Mapped[int | None] = mapped_column(Integer, default=0)\n\n # Error message if sync fails\n error_message: Mapped[str | None] = mapped_column(Text, default=None)\n\n # Timestamps\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n index=True,\n )\n time_started: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n time_finished: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n\n # Relationships\n connector_credential_pair: Mapped[ConnectorCredentialPair] = relationship(\n \"ConnectorCredentialPair\"\n )\n\n __table_args__ = (\n Index(\n \"ix_permission_sync_attempt_latest_for_cc_pair\",\n \"connector_credential_pair_id\",\n \"time_created\",\n ),\n Index(\n \"ix_permission_sync_attempt_status_time\",\n \"status\",\n desc(\"time_finished\"),\n ),\n )\n\n def __repr__(self) -> str:\n return f\"\"\n\n def is_finished(self) -> bool:\n return self.status.is_terminal()\n\n\nclass ExternalGroupPermissionSyncAttempt(Base):\n \"\"\"\n Represents an attempt to sync external group memberships for users.\n This tracks the syncing of user-to-external-group mappings across connectors.\n \"\"\"\n\n __tablename__ = \"external_group_permission_sync_attempt\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n # Can be tied to a specific connector or be a global group sync\n connector_credential_pair_id: Mapped[int | None] = mapped_column(\n ForeignKey(\"connector_credential_pair.id\"),\n nullable=True, # Nullable for global group syncs across all connectors\n )\n\n # Status of the group sync attempt\n status: Mapped[PermissionSyncStatus] = mapped_column(\n Enum(PermissionSyncStatus, native_enum=False, index=True)\n )\n\n # Counts for tracking progress\n total_users_processed: Mapped[int | None] = mapped_column(Integer, default=0)\n total_groups_processed: Mapped[int | None] = mapped_column(Integer, default=0)\n total_group_memberships_synced: Mapped[int | None] = mapped_column(\n Integer, default=0\n )\n\n # Error message if sync fails\n error_message: Mapped[str | None] = mapped_column(Text, default=None)\n\n # Timestamps\n time_created: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n index=True,\n )\n time_started: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n time_finished: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), default=None\n )\n\n # Relationships\n connector_credential_pair: Mapped[ConnectorCredentialPair | None] = relationship(\n \"ConnectorCredentialPair\"\n )\n\n __table_args__ = (\n Index(\n \"ix_group_sync_attempt_cc_pair_time\",\n \"connector_credential_pair_id\",\n \"time_created\",\n ),\n Index(\n \"ix_group_sync_attempt_status_time\",\n \"status\",\n desc(\"time_finished\"),\n ),\n )\n\n def __repr__(self) -> str:\n return f\"\"\n\n def is_finished(self) -> bool:\n return self.status.is_terminal()\n\n\nclass License(Base):\n \"\"\"Stores the signed license blob (singleton pattern - only one row).\"\"\"\n\n __tablename__ = \"license\"\n __table_args__ = (\n # Singleton pattern - unique index on constant ensures only one row\n Index(\"idx_license_singleton\", text(\"(true)\"), unique=True),\n )\n\n id: Mapped[int] = mapped_column(primary_key=True)\n license_data: Mapped[str] = mapped_column(Text, nullable=False)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now()\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n\nclass TenantUsage(Base):\n \"\"\"\n Tracks per-tenant usage statistics within a time window for cloud usage limits.\n\n Each row represents usage for a specific tenant during a specific time window.\n A new row is created when the window rolls over (typically weekly).\n \"\"\"\n\n __tablename__ = \"tenant_usage\"\n\n id: Mapped[int] = mapped_column(primary_key=True)\n\n # The start of the usage tracking window (e.g., start of the week in UTC)\n window_start: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), nullable=False, index=True\n )\n\n # Cumulative LLM usage cost in cents for the window\n llm_cost_cents: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)\n\n # Number of chunks indexed during the window\n chunks_indexed: Mapped[int] = mapped_column(Integer, nullable=False, default=0)\n\n # Number of API calls using API keys or Personal Access Tokens\n api_calls: Mapped[int] = mapped_column(Integer, nullable=False, default=0)\n\n # Number of non-streaming API calls (more expensive operations)\n non_streaming_api_calls: Mapped[int] = mapped_column(\n Integer, nullable=False, default=0\n )\n\n # Last updated timestamp for tracking freshness\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), onupdate=func.now()\n )\n\n __table_args__ = (\n # Ensure only one row per window start (tenant_id is in the schema name)\n UniqueConstraint(\"window_start\", name=\"uq_tenant_usage_window\"),\n )\n\n\n\"\"\"Tables related to Build Mode (CLI Agent Platform)\"\"\"\n\n\nclass BuildSession(Base):\n \"\"\"Stores metadata about CLI agent build sessions.\"\"\"\n\n __tablename__ = \"build_session\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n user_id: Mapped[UUID | None] = mapped_column(\n PGUUID(as_uuid=True), ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=True\n )\n name: Mapped[str | None] = mapped_column(String, nullable=True)\n status: Mapped[BuildSessionStatus] = mapped_column(\n Enum(BuildSessionStatus, native_enum=False, name=\"buildsessionstatus\"),\n nullable=False,\n default=BuildSessionStatus.ACTIVE,\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n last_activity_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n nextjs_port: Mapped[int | None] = mapped_column(Integer, nullable=True)\n demo_data_enabled: Mapped[bool] = mapped_column(\n Boolean, nullable=False, server_default=text(\"true\")\n )\n sharing_scope: Mapped[SharingScope] = mapped_column(\n String,\n nullable=False,\n default=SharingScope.PRIVATE,\n server_default=\"private\",\n )\n\n # Relationships\n user: Mapped[User | None] = relationship(\"User\", foreign_keys=[user_id])\n artifacts: Mapped[list[\"Artifact\"]] = relationship(\n \"Artifact\", back_populates=\"session\", cascade=\"all, delete-orphan\"\n )\n messages: Mapped[list[\"BuildMessage\"]] = relationship(\n \"BuildMessage\", back_populates=\"session\", cascade=\"all, delete-orphan\"\n )\n snapshots: Mapped[list[\"Snapshot\"]] = relationship(\n \"Snapshot\", back_populates=\"session\", cascade=\"all, delete-orphan\"\n )\n\n __table_args__ = (\n Index(\"ix_build_session_user_created\", \"user_id\", desc(\"created_at\")),\n Index(\"ix_build_session_status\", \"status\"),\n )\n\n\nclass Sandbox(Base):\n \"\"\"Stores sandbox container metadata for users (one sandbox per user).\"\"\"\n\n __tablename__ = \"sandbox\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n user_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True),\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"),\n nullable=False,\n unique=True,\n )\n container_id: Mapped[str | None] = mapped_column(String, nullable=True)\n status: Mapped[SandboxStatus] = mapped_column(\n Enum(SandboxStatus, native_enum=False, name=\"sandboxstatus\"),\n nullable=False,\n default=SandboxStatus.PROVISIONING,\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n last_heartbeat: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n # Relationships\n user: Mapped[User] = relationship(\"User\")\n\n __table_args__ = (\n Index(\"ix_sandbox_status\", \"status\"),\n Index(\"ix_sandbox_container_id\", \"container_id\"),\n )\n\n\nclass Artifact(Base):\n \"\"\"Stores metadata about artifacts generated by CLI agents.\"\"\"\n\n __tablename__ = \"artifact\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n session_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True),\n ForeignKey(\"build_session.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n type: Mapped[ArtifactType] = mapped_column(\n Enum(ArtifactType, native_enum=False, name=\"artifacttype\"), nullable=False\n )\n # path of artifact in sandbox relative to outputs/\n path: Mapped[str] = mapped_column(String, nullable=False)\n name: Mapped[str] = mapped_column(String, nullable=False)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n # Relationships\n session: Mapped[BuildSession] = relationship(\n \"BuildSession\", back_populates=\"artifacts\"\n )\n\n __table_args__ = (\n Index(\"ix_artifact_session_created\", \"session_id\", desc(\"created_at\")),\n Index(\"ix_artifact_type\", \"type\"),\n )\n\n\nclass Snapshot(Base):\n \"\"\"Stores metadata about session output snapshots.\"\"\"\n\n __tablename__ = \"snapshot\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n session_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True),\n ForeignKey(\"build_session.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n storage_path: Mapped[str] = mapped_column(String, nullable=False)\n size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n\n # Relationships\n session: Mapped[BuildSession] = relationship(\n \"BuildSession\", back_populates=\"snapshots\"\n )\n\n __table_args__ = (\n Index(\"ix_snapshot_session_created\", \"session_id\", desc(\"created_at\")),\n )\n\n\nclass BuildMessage(Base):\n \"\"\"Stores messages exchanged in build sessions.\n\n All message data is stored in message_metadata as JSON (the raw ACP packet).\n The turn_index groups all assistant responses under the user prompt they respond to.\n\n Packet types stored in message_metadata:\n - user_message: {type: \"user_message\", content: {...}}\n - agent_message: {type: \"agent_message\", content: {...}} (accumulated from chunks)\n - agent_thought: {type: \"agent_thought\", content: {...}} (accumulated from chunks)\n - tool_call_progress: {type: \"tool_call_progress\", status: \"completed\", ...} (only completed)\n - agent_plan_update: {type: \"agent_plan_update\", entries: [...]} (upserted, latest only)\n \"\"\"\n\n __tablename__ = \"build_message\"\n\n id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True), primary_key=True, default=uuid4\n )\n session_id: Mapped[UUID] = mapped_column(\n PGUUID(as_uuid=True),\n ForeignKey(\"build_session.id\", ondelete=\"CASCADE\"),\n nullable=False,\n )\n turn_index: Mapped[int] = mapped_column(Integer, nullable=False)\n type: Mapped[MessageType] = mapped_column(\n Enum(MessageType, native_enum=False, name=\"messagetype\"), nullable=False\n )\n message_metadata: Mapped[dict[str, Any]] = mapped_column(PGJSONB, nullable=False)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n\n # Relationships\n session: Mapped[BuildSession] = relationship(\n \"BuildSession\", back_populates=\"messages\"\n )\n\n __table_args__ = (\n Index(\n \"ix_build_message_session_turn\", \"session_id\", \"turn_index\", \"created_at\"\n ),\n )\n\n\n\"\"\"\nSCIM 2.0 Provisioning Models (Enterprise Edition only)\nUsed for automated user/group provisioning from identity providers (Okta, Azure AD).\n\"\"\"\n\n\nclass ScimToken(Base):\n \"\"\"Bearer tokens for IdP SCIM authentication.\"\"\"\n\n __tablename__ = \"scim_token\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, nullable=False)\n hashed_token: Mapped[str] = mapped_column(\n String(64), unique=True, nullable=False\n ) # SHA256 = 64 hex chars\n token_display: Mapped[str] = mapped_column(\n String, nullable=False\n ) # Last 4 chars for UI identification\n\n created_by_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), nullable=False\n )\n\n is_active: Mapped[bool] = mapped_column(\n Boolean, server_default=text(\"true\"), nullable=False\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n last_used_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n created_by: Mapped[User] = relationship(\"User\", foreign_keys=[created_by_id])\n\n\nclass ScimUserMapping(Base):\n \"\"\"Maps SCIM externalId from the IdP to an Onyx User.\"\"\"\n\n __tablename__ = \"scim_user_mapping\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n external_id: Mapped[str | None] = mapped_column(\n String, unique=True, index=True, nullable=True\n )\n user_id: Mapped[UUID] = mapped_column(\n ForeignKey(\"user.id\", ondelete=\"CASCADE\"), unique=True, nullable=False\n )\n scim_username: Mapped[str | None] = mapped_column(String, nullable=True)\n department: Mapped[str | None] = mapped_column(String, nullable=True)\n manager: Mapped[str | None] = mapped_column(String, nullable=True)\n given_name: Mapped[str | None] = mapped_column(String, nullable=True)\n family_name: Mapped[str | None] = mapped_column(String, nullable=True)\n scim_emails_json: Mapped[str | None] = mapped_column(Text, nullable=True)\n\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n user: Mapped[User] = relationship(\"User\", foreign_keys=[user_id])\n\n\nclass ScimGroupMapping(Base):\n \"\"\"Maps SCIM externalId from the IdP to an Onyx UserGroup.\"\"\"\n\n __tablename__ = \"scim_group_mapping\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n external_id: Mapped[str] = mapped_column(String, unique=True, index=True)\n user_group_id: Mapped[int] = mapped_column(\n ForeignKey(\"user_group.id\", ondelete=\"CASCADE\"), unique=True, nullable=False\n )\n\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n user_group: Mapped[UserGroup] = relationship(\n \"UserGroup\", foreign_keys=[user_group_id]\n )\n\n\nclass CodeInterpreterServer(Base):\n \"\"\"Details about the code interpreter server\"\"\"\n\n __tablename__ = \"code_interpreter_server\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n server_enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)\n\n\nclass CacheStore(Base):\n \"\"\"Key-value cache table used by ``PostgresCacheBackend``.\n\n Replaces Redis for simple KV caching, locks, and list operations\n when ``CACHE_BACKEND=postgres`` (NO_VECTOR_DB deployments).\n\n Intentionally separate from ``KVStore``:\n - Stores raw bytes (LargeBinary) vs JSONB, matching Redis semantics.\n - Has ``expires_at`` for TTL; rows are periodically garbage-collected.\n - Holds ephemeral data (tokens, stop signals, lock state) not\n persistent application config, so cleanup can be aggressive.\n \"\"\"\n\n __tablename__ = \"cache_store\"\n\n key: Mapped[str] = mapped_column(String, primary_key=True)\n value: Mapped[bytes | None] = mapped_column(LargeBinary, nullable=True)\n expires_at: Mapped[datetime.datetime | None] = mapped_column(\n DateTime(timezone=True), nullable=True\n )\n\n\nclass Hook(Base):\n \"\"\"Pairs a HookPoint with a customer-provided API endpoint.\n\n At most one non-deleted Hook per HookPoint is allowed, enforced by a\n partial unique index on (hook_point) where deleted=false.\n \"\"\"\n\n __tablename__ = \"hook\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n name: Mapped[str] = mapped_column(String, nullable=False)\n hook_point: Mapped[HookPoint] = mapped_column(\n Enum(HookPoint, native_enum=False), nullable=False\n )\n endpoint_url: Mapped[str | None] = mapped_column(Text, nullable=True)\n api_key: Mapped[SensitiveValue[str] | None] = mapped_column(\n EncryptedString(), nullable=True\n )\n is_reachable: Mapped[bool | None] = mapped_column(\n Boolean, nullable=True, default=None\n ) # null = never validated, true = last check passed, false = last check failed\n fail_strategy: Mapped[HookFailStrategy] = mapped_column(\n Enum(HookFailStrategy, native_enum=False),\n nullable=False,\n default=HookFailStrategy.HARD,\n )\n timeout_seconds: Mapped[float] = mapped_column(Float, nullable=False, default=30.0)\n is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n deleted: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n creator_id: Mapped[UUID | None] = mapped_column(\n PGUUID(as_uuid=True),\n ForeignKey(\"user.id\", ondelete=\"SET NULL\"),\n nullable=True,\n )\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False\n )\n updated_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True),\n server_default=func.now(),\n onupdate=func.now(),\n nullable=False,\n )\n\n creator: Mapped[\"User | None\"] = relationship(\"User\", foreign_keys=[creator_id])\n execution_logs: Mapped[list[\"HookExecutionLog\"]] = relationship(\n \"HookExecutionLog\", back_populates=\"hook\", cascade=\"all, delete-orphan\"\n )\n\n __table_args__ = (\n Index(\n \"ix_hook_one_non_deleted_per_point\",\n \"hook_point\",\n unique=True,\n postgresql_where=(deleted == False), # noqa: E712\n ),\n )\n\n\nclass HookExecutionLog(Base):\n \"\"\"Records hook executions for health monitoring and debugging.\n\n Currently only failures are logged; the is_success column exists so\n success logging can be added later without a schema change.\n Retention: rows older than 30 days are deleted by a nightly Celery task.\n \"\"\"\n\n __tablename__ = \"hook_execution_log\"\n\n id: Mapped[int] = mapped_column(Integer, primary_key=True)\n hook_id: Mapped[int] = mapped_column(\n Integer,\n ForeignKey(\"hook.id\", ondelete=\"CASCADE\"),\n nullable=False,\n index=True,\n )\n is_success: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)\n error_message: Mapped[str | None] = mapped_column(Text, nullable=True)\n status_code: Mapped[int | None] = mapped_column(Integer, nullable=True)\n duration_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)\n created_at: Mapped[datetime.datetime] = mapped_column(\n DateTime(timezone=True), server_default=func.now(), nullable=False, index=True\n )\n\n hook: Mapped[\"Hook\"] = relationship(\"Hook\", back_populates=\"execution_logs\")\n" }, { "path": "backend/onyx/db/notification.py", "content": "from datetime import datetime\nfrom datetime import timezone\nfrom uuid import UUID\n\nfrom sqlalchemy import cast\nfrom sqlalchemy import select\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql import func\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import NotificationType\nfrom onyx.db.models import Notification\nfrom onyx.db.models import User\n\n\ndef create_notification(\n user_id: UUID | None,\n notif_type: NotificationType,\n db_session: Session,\n title: str,\n description: str | None = None,\n additional_data: dict | None = None,\n autocommit: bool = True,\n) -> Notification:\n # Previously, we only matched the first identical, undismissed notification\n # Now, we assume some uniqueness to notifications\n # If we previously issued a notification that was dismissed, we no longer issue a new one\n\n # Normalize additional_data to match the unique index behavior\n # The index uses COALESCE(additional_data, '{}'::jsonb)\n # We need to match this logic in our query\n additional_data_normalized = additional_data if additional_data is not None else {}\n\n existing_notification = (\n db_session.query(Notification)\n .filter_by(user_id=user_id, notif_type=notif_type)\n .filter(\n func.coalesce(Notification.additional_data, cast({}, postgresql.JSONB))\n == additional_data_normalized\n )\n .first()\n )\n\n if existing_notification:\n # Update the last_shown timestamp if the notification is not dismissed\n if not existing_notification.dismissed:\n existing_notification.last_shown = func.now()\n if autocommit:\n db_session.commit()\n return existing_notification\n\n # Create a new notification if none exists\n notification = Notification(\n user_id=user_id,\n notif_type=notif_type,\n title=title,\n description=description,\n dismissed=False,\n last_shown=func.now(),\n first_shown=func.now(),\n additional_data=additional_data,\n )\n db_session.add(notification)\n if autocommit:\n db_session.commit()\n return notification\n\n\ndef get_notification_by_id(\n notification_id: int, user: User, db_session: Session\n) -> Notification:\n user_id = user.id\n notif = db_session.get(Notification, notification_id)\n if not notif:\n raise ValueError(f\"No notification found with id {notification_id}\")\n if notif.user_id != user_id and not (\n notif.user_id is None and user is not None and user.role == UserRole.ADMIN\n ):\n raise PermissionError(\n f\"User {user_id} is not authorized to access notification {notification_id}\"\n )\n return notif\n\n\ndef get_notifications(\n user: User | None,\n db_session: Session,\n notif_type: NotificationType | None = None,\n include_dismissed: bool = True,\n) -> list[Notification]:\n query = select(Notification).where(\n Notification.user_id == user.id if user else Notification.user_id.is_(None)\n )\n if not include_dismissed:\n query = query.where(Notification.dismissed.is_(False))\n if notif_type:\n query = query.where(Notification.notif_type == notif_type)\n # Sort: undismissed first, then by date (newest first)\n query = query.order_by(\n Notification.dismissed.asc(),\n Notification.first_shown.desc(),\n )\n return list(db_session.execute(query).scalars().all())\n\n\ndef dismiss_all_notifications(\n notif_type: NotificationType,\n db_session: Session,\n) -> None:\n db_session.query(Notification).filter(Notification.notif_type == notif_type).update(\n {\"dismissed\": True}\n )\n db_session.commit()\n\n\ndef dismiss_notification(notification: Notification, db_session: Session) -> None:\n notification.dismissed = True\n db_session.commit()\n\n\ndef batch_dismiss_notifications(\n notifications: list[Notification],\n db_session: Session,\n) -> None:\n for notification in notifications:\n notification.dismissed = True\n db_session.commit()\n\n\ndef batch_create_notifications(\n user_ids: list[UUID],\n notif_type: NotificationType,\n db_session: Session,\n title: str,\n description: str | None = None,\n additional_data: dict | None = None,\n) -> int:\n \"\"\"\n Create notifications for multiple users in a single batch operation.\n Uses ON CONFLICT DO NOTHING for atomic idempotent inserts - if a user already\n has a notification with the same (user_id, notif_type, additional_data), the\n insert is silently skipped.\n\n Returns the number of notifications created.\n\n Relies on unique index on (user_id, notif_type, COALESCE(additional_data, '{}'))\n \"\"\"\n if not user_ids:\n return 0\n\n now = datetime.now(timezone.utc)\n # Use empty dict instead of None to match COALESCE behavior in the unique index\n additional_data_normalized = additional_data if additional_data is not None else {}\n\n values = [\n {\n \"user_id\": uid,\n \"notif_type\": notif_type.value,\n \"title\": title,\n \"description\": description,\n \"dismissed\": False,\n \"last_shown\": now,\n \"first_shown\": now,\n \"additional_data\": additional_data_normalized,\n }\n for uid in user_ids\n ]\n\n stmt = insert(Notification).values(values).on_conflict_do_nothing()\n result = db_session.execute(stmt)\n db_session.commit()\n\n # rowcount returns number of rows inserted (excludes conflicts)\n # CursorResult has rowcount but session.execute type hints are too broad\n return result.rowcount if result.rowcount >= 0 else 0 # type: ignore[attr-defined]\n\n\ndef update_notification_last_shown(\n notification: Notification, db_session: Session\n) -> None:\n notification.last_shown = func.now()\n db_session.commit()\n" }, { "path": "backend/onyx/db/oauth_config.py", "content": "from typing import Any\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import OAuthConfig\nfrom onyx.db.models import OAuthUserToken\nfrom onyx.db.models import Tool\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n# OAuth Config CRUD operations\n\n\ndef create_oauth_config(\n name: str,\n authorization_url: str,\n token_url: str,\n client_id: str,\n client_secret: str,\n scopes: list[str] | None,\n additional_params: dict[str, str] | None,\n db_session: Session,\n) -> OAuthConfig:\n \"\"\"Create a new OAuth configuration\"\"\"\n oauth_config = OAuthConfig(\n name=name,\n authorization_url=authorization_url,\n token_url=token_url,\n client_id=client_id,\n client_secret=client_secret,\n scopes=scopes,\n additional_params=additional_params,\n )\n db_session.add(oauth_config)\n db_session.commit()\n return oauth_config\n\n\ndef get_oauth_config(oauth_config_id: int, db_session: Session) -> OAuthConfig | None:\n \"\"\"Get OAuth configuration by ID\"\"\"\n return db_session.scalar(\n select(OAuthConfig).where(OAuthConfig.id == oauth_config_id)\n )\n\n\ndef get_oauth_configs(db_session: Session) -> list[OAuthConfig]:\n \"\"\"Get all OAuth configurations\"\"\"\n return list(db_session.scalars(select(OAuthConfig)).all())\n\n\ndef update_oauth_config(\n oauth_config_id: int,\n db_session: Session,\n name: str | None = None,\n authorization_url: str | None = None,\n token_url: str | None = None,\n client_id: str | None = None,\n client_secret: str | None = None,\n scopes: list[str] | None = None,\n additional_params: dict[str, Any] | None = None,\n clear_client_id: bool = False,\n clear_client_secret: bool = False,\n) -> OAuthConfig:\n \"\"\"\n Update OAuth configuration.\n\n NOTE: If client_id or client_secret are None, existing values are preserved.\n To clear these values, set clear_client_id or clear_client_secret to True.\n This allows partial updates without re-entering secrets.\n \"\"\"\n oauth_config = db_session.scalar(\n select(OAuthConfig).where(OAuthConfig.id == oauth_config_id)\n )\n if oauth_config is None:\n raise ValueError(f\"OAuth config with id {oauth_config_id} does not exist\")\n\n # Update only provided fields\n if name is not None:\n oauth_config.name = name\n if authorization_url is not None:\n oauth_config.authorization_url = authorization_url\n if token_url is not None:\n oauth_config.token_url = token_url\n if clear_client_id:\n oauth_config.client_id = \"\" # type: ignore[assignment]\n elif client_id is not None:\n oauth_config.client_id = client_id # type: ignore[assignment]\n if clear_client_secret:\n oauth_config.client_secret = \"\" # type: ignore[assignment]\n elif client_secret is not None:\n oauth_config.client_secret = client_secret # type: ignore[assignment]\n if scopes is not None:\n oauth_config.scopes = scopes\n if additional_params is not None:\n oauth_config.additional_params = additional_params\n\n db_session.commit()\n return oauth_config\n\n\ndef delete_oauth_config(oauth_config_id: int, db_session: Session) -> None:\n \"\"\"\n Delete OAuth configuration.\n\n Sets oauth_config_id to NULL for associated tools due to SET NULL foreign key.\n Cascades delete to user tokens.\n \"\"\"\n oauth_config = db_session.scalar(\n select(OAuthConfig).where(OAuthConfig.id == oauth_config_id)\n )\n if oauth_config is None:\n raise ValueError(f\"OAuth config with id {oauth_config_id} does not exist\")\n\n db_session.delete(oauth_config)\n db_session.commit()\n\n\n# User Token operations\n\n\ndef get_user_oauth_token(\n oauth_config_id: int, user_id: UUID, db_session: Session\n) -> OAuthUserToken | None:\n \"\"\"Get user's OAuth token for a specific configuration\"\"\"\n return db_session.scalar(\n select(OAuthUserToken).where(\n OAuthUserToken.oauth_config_id == oauth_config_id,\n OAuthUserToken.user_id == user_id,\n )\n )\n\n\ndef get_all_user_oauth_tokens(\n user_id: UUID, db_session: Session\n) -> list[OAuthUserToken]:\n \"\"\"\n Get all user OAuth tokens.\n \"\"\"\n stmt = select(OAuthUserToken).where(OAuthUserToken.user_id == user_id)\n\n return list(db_session.scalars(stmt).all())\n\n\ndef upsert_user_oauth_token(\n oauth_config_id: int, user_id: UUID, token_data: dict, db_session: Session\n) -> OAuthUserToken:\n \"\"\"Insert or update user's OAuth token for a specific configuration\"\"\"\n existing_token = get_user_oauth_token(oauth_config_id, user_id, db_session)\n\n if existing_token:\n # Update existing token\n existing_token.token_data = token_data # type: ignore[assignment]\n db_session.commit()\n return existing_token\n else:\n # Create new token\n new_token = OAuthUserToken(\n oauth_config_id=oauth_config_id,\n user_id=user_id,\n token_data=token_data,\n )\n db_session.add(new_token)\n db_session.commit()\n return new_token\n\n\ndef delete_user_oauth_token(\n oauth_config_id: int, user_id: UUID, db_session: Session\n) -> None:\n \"\"\"Delete user's OAuth token for a specific configuration\"\"\"\n user_token = get_user_oauth_token(oauth_config_id, user_id, db_session)\n if user_token is None:\n raise ValueError(\n f\"OAuth token for user {user_id} and config {oauth_config_id} does not exist\"\n )\n\n db_session.delete(user_token)\n db_session.commit()\n\n\n# Helper operations\n\n\ndef get_tools_by_oauth_config(oauth_config_id: int, db_session: Session) -> list[Tool]:\n \"\"\"Get all tools that use a specific OAuth configuration\"\"\"\n return list(\n db_session.scalars(\n select(Tool).where(Tool.oauth_config_id == oauth_config_id)\n ).all()\n )\n" }, { "path": "backend/onyx/db/opensearch_migration.py", "content": "\"\"\"Database operations for OpenSearch migration tracking.\n\nThis module provides functions to track the progress of migrating documents\nfrom Vespa to OpenSearch.\n\"\"\"\n\nimport json\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom sqlalchemy import select\nfrom sqlalchemy import text\nfrom sqlalchemy.dialects.postgresql import insert\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n GET_VESPA_CHUNKS_SLICE_COUNT,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n TOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE,\n)\nfrom onyx.configs.app_configs import ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX\nfrom onyx.db.enums import OpenSearchDocumentMigrationStatus\nfrom onyx.db.models import Document\nfrom onyx.db.models import OpenSearchDocumentMigrationRecord\nfrom onyx.db.models import OpenSearchTenantMigrationRecord\nfrom onyx.document_index.vespa.shared_utils.utils import (\n replace_invalid_doc_id_characters,\n)\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef get_paginated_document_batch(\n db_session: Session,\n limit: int,\n prev_ending_document_id: str | None = None,\n) -> list[str]:\n \"\"\"Gets a paginated batch of document IDs from the Document table.\n\n We need some deterministic ordering to ensure that we don't miss any\n documents when paginating. This function uses the document ID. It is\n possible a document is inserted above a spot this function has already\n passed. In that event we assume that the document will be indexed into\n OpenSearch anyway and we don't need to migrate.\n TODO(andrei): Consider ordering on last_modified in addition to ID to better\n match get_opensearch_migration_records_needing_migration.\n\n Args:\n db_session: SQLAlchemy session.\n limit: Number of document IDs to fetch.\n prev_ending_document_id: Document ID to start after (for pagination). If\n None, returns the first batch of documents. If not None, this should\n be the last ordered ID which was fetched in a previous batch.\n Defaults to None.\n\n Returns:\n List of document IDs.\n \"\"\"\n stmt = select(Document.id).order_by(Document.id.asc()).limit(limit)\n if prev_ending_document_id is not None:\n stmt = stmt.where(Document.id > prev_ending_document_id)\n return list(db_session.scalars(stmt).all())\n\n\ndef get_last_opensearch_migration_document_id(\n db_session: Session,\n) -> str | None:\n \"\"\"\n Gets the last document ID in the OpenSearchDocumentMigrationRecord table.\n\n Returns None if no records are found.\n \"\"\"\n stmt = (\n select(OpenSearchDocumentMigrationRecord.document_id)\n .order_by(OpenSearchDocumentMigrationRecord.document_id.desc())\n .limit(1)\n )\n return db_session.scalars(stmt).first()\n\n\ndef create_opensearch_migration_records_with_commit(\n db_session: Session,\n document_ids: list[str],\n) -> None:\n \"\"\"Creates new OpenSearchDocumentMigrationRecord records.\n\n Silently skips any document IDs that already have records.\n \"\"\"\n if not document_ids:\n return\n\n values = [\n {\n \"document_id\": document_id,\n \"status\": OpenSearchDocumentMigrationStatus.PENDING,\n }\n for document_id in document_ids\n ]\n\n stmt = insert(OpenSearchDocumentMigrationRecord).values(values)\n stmt = stmt.on_conflict_do_nothing(index_elements=[\"document_id\"])\n\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef get_opensearch_migration_records_needing_migration(\n db_session: Session,\n limit: int,\n) -> list[OpenSearchDocumentMigrationRecord]:\n \"\"\"Gets records of documents that need to be migrated.\n\n Properties:\n - First tries documents with status PENDING.\n - Of these, orders documents with the oldest last_modified to prioritize\n documents that were modified a long time ago, as they are presumed to be\n stable. This column is modified in many flows so is not a guarantee of the\n document having been indexed.\n - Then if there's room in the result, tries documents with status FAILED.\n - Of these, first orders documents on the least attempts_count so as to have\n a backoff for recently-failed docs. Then orders on last_modified as\n before.\n \"\"\"\n result: list[OpenSearchDocumentMigrationRecord] = []\n\n # Step 1: Fetch as many PENDING status records as possible ordered by\n # last_modified (oldest first). last_modified lives on Document, so we join.\n stmt_pending = (\n select(OpenSearchDocumentMigrationRecord)\n .join(Document, OpenSearchDocumentMigrationRecord.document_id == Document.id)\n .where(\n OpenSearchDocumentMigrationRecord.status\n == OpenSearchDocumentMigrationStatus.PENDING\n )\n .order_by(Document.last_modified.asc())\n .limit(limit)\n )\n result.extend(list(db_session.scalars(stmt_pending).all()))\n remaining = limit - len(result)\n\n # Step 2: If more are needed, fetch records with status FAILED, ordered by\n # attempts_count (lowest first), then last_modified (oldest first).\n if remaining > 0:\n stmt_failed = (\n select(OpenSearchDocumentMigrationRecord)\n .join(\n Document,\n OpenSearchDocumentMigrationRecord.document_id == Document.id,\n )\n .where(\n OpenSearchDocumentMigrationRecord.status\n == OpenSearchDocumentMigrationStatus.FAILED\n )\n .order_by(\n OpenSearchDocumentMigrationRecord.attempts_count.asc(),\n Document.last_modified.asc(),\n )\n .limit(remaining)\n )\n result.extend(list(db_session.scalars(stmt_failed).all()))\n\n return result\n\n\ndef get_total_opensearch_migration_record_count(\n db_session: Session,\n) -> int:\n \"\"\"Gets the total number of OpenSearch migration records.\n\n Used to check whether every document has been tracked for migration.\n \"\"\"\n return db_session.query(OpenSearchDocumentMigrationRecord).count()\n\n\ndef get_total_document_count(db_session: Session) -> int:\n \"\"\"Gets the total number of documents.\n\n Used to check whether every document has been tracked for migration.\n \"\"\"\n return db_session.query(Document).count()\n\n\ndef try_insert_opensearch_tenant_migration_record_with_commit(\n db_session: Session,\n) -> None:\n \"\"\"Tries to insert the singleton row on OpenSearchTenantMigrationRecord.\n\n Does nothing if the row already exists.\n \"\"\"\n stmt = insert(OpenSearchTenantMigrationRecord).on_conflict_do_nothing(\n index_elements=[text(\"(true)\")]\n )\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef increment_num_times_observed_no_additional_docs_to_migrate_with_commit(\n db_session: Session,\n) -> None:\n \"\"\"Increments the number of times observed no additional docs to migrate.\n\n Requires the OpenSearchTenantMigrationRecord to exist.\n\n Used to track when to stop the migration task.\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n raise RuntimeError(\"OpenSearchTenantMigrationRecord not found.\")\n record.num_times_observed_no_additional_docs_to_migrate += 1\n db_session.commit()\n\n\ndef increment_num_times_observed_no_additional_docs_to_populate_migration_table_with_commit(\n db_session: Session,\n) -> None:\n \"\"\"\n Increments the number of times observed no additional docs to populate the\n migration table.\n\n Requires the OpenSearchTenantMigrationRecord to exist.\n\n Used to track when to stop the migration check task.\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n raise RuntimeError(\"OpenSearchTenantMigrationRecord not found.\")\n record.num_times_observed_no_additional_docs_to_populate_migration_table += 1\n db_session.commit()\n\n\ndef should_document_migration_be_permanently_failed(\n opensearch_document_migration_record: OpenSearchDocumentMigrationRecord,\n) -> bool:\n return (\n opensearch_document_migration_record.status\n == OpenSearchDocumentMigrationStatus.PERMANENTLY_FAILED\n or (\n opensearch_document_migration_record.status\n == OpenSearchDocumentMigrationStatus.FAILED\n and opensearch_document_migration_record.attempts_count\n >= TOTAL_ALLOWABLE_DOC_MIGRATION_ATTEMPTS_BEFORE_PERMANENT_FAILURE\n )\n )\n\n\ndef get_vespa_visit_state(\n db_session: Session,\n) -> tuple[dict[int, str | None], int]:\n \"\"\"Gets the current Vespa migration state from the tenant migration record.\n\n Requires the OpenSearchTenantMigrationRecord to exist.\n\n Returns:\n Tuple of (continuation_token_map, total_chunks_migrated).\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n raise RuntimeError(\"OpenSearchTenantMigrationRecord not found.\")\n if record.vespa_visit_continuation_token is None:\n continuation_token_map: dict[int, str | None] = {\n slice_id: None for slice_id in range(GET_VESPA_CHUNKS_SLICE_COUNT)\n }\n else:\n json_loaded_continuation_token_map = json.loads(\n record.vespa_visit_continuation_token\n )\n continuation_token_map = {\n int(key): value for key, value in json_loaded_continuation_token_map.items()\n }\n return continuation_token_map, record.total_chunks_migrated\n\n\ndef update_vespa_visit_progress_with_commit(\n db_session: Session,\n continuation_token_map: dict[int, str | None],\n chunks_processed: int,\n chunks_errored: int,\n approx_chunk_count_in_vespa: int | None,\n) -> None:\n \"\"\"Updates the Vespa migration progress and commits.\n\n Requires the OpenSearchTenantMigrationRecord to exist.\n\n Args:\n db_session: SQLAlchemy session.\n continuation_token_map: The new continuation token map. None entry means\n the visit is complete for that slice.\n chunks_processed: Number of chunks processed in this batch (added to\n the running total).\n chunks_errored: Number of chunks errored in this batch (added to the\n running errored total).\n approx_chunk_count_in_vespa: Approximate number of chunks in Vespa. If\n None, the existing value is used.\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n raise RuntimeError(\"OpenSearchTenantMigrationRecord not found.\")\n record.vespa_visit_continuation_token = json.dumps(continuation_token_map)\n record.total_chunks_migrated += chunks_processed\n record.total_chunks_errored += chunks_errored\n record.approx_chunk_count_in_vespa = (\n approx_chunk_count_in_vespa\n if approx_chunk_count_in_vespa is not None\n else record.approx_chunk_count_in_vespa\n )\n db_session.commit()\n\n\ndef mark_migration_completed_time_if_not_set_with_commit(\n db_session: Session,\n) -> None:\n \"\"\"Marks the migration completed time if not set.\n\n Requires the OpenSearchTenantMigrationRecord to exist.\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n raise RuntimeError(\"OpenSearchTenantMigrationRecord not found.\")\n if record.migration_completed_at is not None:\n return\n record.migration_completed_at = datetime.now(timezone.utc)\n db_session.commit()\n\n\ndef is_migration_completed(db_session: Session) -> bool:\n \"\"\"Returns True if the migration is completed.\n\n Can be run even if the migration record does not exist.\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n return record is not None and record.migration_completed_at is not None\n\n\ndef build_sanitized_to_original_doc_id_mapping(\n db_session: Session,\n) -> dict[str, str]:\n \"\"\"Pre-computes a mapping of sanitized -> original document IDs.\n\n Only includes documents whose ID contains single quotes (the only character\n that gets sanitized by replace_invalid_doc_id_characters). For all other\n documents, sanitized == original and no mapping entry is needed.\n\n Scans over all documents.\n\n Checks if the sanitized ID already exists as a genuine separate document in\n the Document table. If so, raises as there is no way of resolving the\n conflict in the migration. The user will need to reindex.\n\n Args:\n db_session: SQLAlchemy session.\n\n Returns:\n Dict mapping sanitized_id -> original_id, only for documents where\n the IDs differ. Empty dict means no documents have single quotes\n in their IDs.\n \"\"\"\n # Find all documents with single quotes in their ID.\n stmt = select(Document.id).where(Document.id.contains(\"'\"))\n ids_with_quotes = list(db_session.scalars(stmt).all())\n\n result: dict[str, str] = {}\n for original_id in ids_with_quotes:\n sanitized_id = replace_invalid_doc_id_characters(original_id)\n if sanitized_id != original_id:\n result[sanitized_id] = original_id\n\n # See if there are any documents whose ID is a sanitized ID of another\n # document. If there is even one match, we cannot proceed.\n stmt = select(Document.id).where(Document.id.in_(result.keys()))\n ids_with_matches = list(db_session.scalars(stmt).all())\n if ids_with_matches:\n raise RuntimeError(\n f\"Documents with IDs {ids_with_matches} have sanitized IDs that match other documents. \"\n \"This is not supported and the user will need to reindex.\"\n )\n\n return result\n\n\ndef get_opensearch_migration_state(\n db_session: Session,\n) -> tuple[int, datetime | None, datetime | None, int | None]:\n \"\"\"Returns the state of the Vespa to OpenSearch migration.\n\n If the tenant migration record is not found, returns defaults of 0, None,\n None, None.\n\n Args:\n db_session: SQLAlchemy session.\n\n Returns:\n Tuple of (total_chunks_migrated, created_at, migration_completed_at,\n approx_chunk_count_in_vespa).\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n return 0, None, None, None\n return (\n record.total_chunks_migrated,\n record.created_at,\n record.migration_completed_at,\n record.approx_chunk_count_in_vespa,\n )\n\n\ndef get_opensearch_retrieval_state(\n db_session: Session,\n) -> bool:\n \"\"\"Returns the state of the OpenSearch retrieval.\n\n If the tenant migration record is not found, defaults to\n ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX.\n \"\"\"\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n return ENABLE_OPENSEARCH_RETRIEVAL_FOR_ONYX\n return record.enable_opensearch_retrieval\n\n\ndef set_enable_opensearch_retrieval_with_commit(\n db_session: Session,\n enable: bool,\n) -> None:\n \"\"\"Sets the enable_opensearch_retrieval flag on the singleton record.\n\n Creates the record if it doesn't exist yet.\n \"\"\"\n try_insert_opensearch_tenant_migration_record_with_commit(db_session)\n record = db_session.query(OpenSearchTenantMigrationRecord).first()\n if record is None:\n raise RuntimeError(\"OpenSearchTenantMigrationRecord not found.\")\n record.enable_opensearch_retrieval = enable\n db_session.commit()\n" }, { "path": "backend/onyx/db/pat.py", "content": "\"\"\"Database operations for Personal Access Tokens.\"\"\"\n\nimport asyncio\nfrom datetime import datetime\nfrom datetime import timezone\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.pat import build_displayable_pat\nfrom onyx.auth.pat import calculate_expiration\nfrom onyx.auth.pat import generate_pat\nfrom onyx.auth.pat import hash_pat\nfrom onyx.db.engine.async_sql_engine import get_async_session_context_manager\nfrom onyx.db.models import PersonalAccessToken\nfrom onyx.db.models import User\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\nlogger = setup_logger()\n\n\nasync def fetch_user_for_pat(\n hashed_token: str, async_db_session: AsyncSession\n) -> User | None:\n \"\"\"Fetch user associated with PAT. Returns None if invalid, expired, or inactive user.\n\n NOTE: This is async since it's used during auth (which is necessarily async due to FastAPI Users).\n NOTE: Expired includes both naturally expired and user-revoked tokens (revocation sets expires_at=NOW()).\n\n Uses select(User) as primary entity so that joined-eager relationships (e.g. oauth_accounts)\n are loaded correctly — matching the pattern in fetch_user_for_api_key.\n \"\"\"\n now = datetime.now(timezone.utc)\n\n user = await async_db_session.scalar(\n select(User)\n .join(PersonalAccessToken, PersonalAccessToken.user_id == User.id)\n .where(PersonalAccessToken.hashed_token == hashed_token)\n .where(User.is_active) # type: ignore\n .where(\n (PersonalAccessToken.expires_at.is_(None))\n | (PersonalAccessToken.expires_at > now)\n )\n )\n if not user:\n return None\n\n _schedule_pat_last_used_update(hashed_token, now)\n return user\n\n\ndef _schedule_pat_last_used_update(hashed_token: str, now: datetime) -> None:\n \"\"\"Fire-and-forget update of last_used_at, throttled to 5-minute granularity.\"\"\"\n\n async def _update() -> None:\n try:\n tenant_id = get_current_tenant_id()\n async with get_async_session_context_manager(tenant_id) as session:\n pat = await session.scalar(\n select(PersonalAccessToken).where(\n PersonalAccessToken.hashed_token == hashed_token\n )\n )\n if not pat:\n return\n if (\n pat.last_used_at is not None\n and (now - pat.last_used_at).total_seconds() <= 300\n ):\n return\n await session.execute(\n update(PersonalAccessToken)\n .where(PersonalAccessToken.hashed_token == hashed_token)\n .values(last_used_at=now)\n )\n await session.commit()\n except Exception as e:\n logger.warning(f\"Failed to update last_used_at for PAT: {e}\")\n\n asyncio.create_task(_update())\n\n\ndef create_pat(\n db_session: Session,\n user_id: UUID,\n name: str,\n expiration_days: int | None,\n) -> tuple[PersonalAccessToken, str]:\n \"\"\"Create new PAT. Returns (db_record, raw_token).\n\n Raises ValueError if user is inactive or not found.\n \"\"\"\n user = db_session.scalar(select(User).where(User.id == user_id)) # type: ignore\n if not user or not user.is_active:\n raise ValueError(\"Cannot create PAT for inactive or non-existent user\")\n\n tenant_id = get_current_tenant_id()\n raw_token = generate_pat(tenant_id)\n\n pat = PersonalAccessToken(\n name=name,\n hashed_token=hash_pat(raw_token),\n token_display=build_displayable_pat(raw_token),\n user_id=user_id,\n expires_at=calculate_expiration(expiration_days),\n )\n db_session.add(pat)\n db_session.commit()\n\n return pat, raw_token\n\n\ndef list_user_pats(db_session: Session, user_id: UUID) -> list[PersonalAccessToken]:\n \"\"\"List all active (non-expired) PATs for a user.\"\"\"\n return list(\n db_session.scalars(\n select(PersonalAccessToken)\n .where(PersonalAccessToken.user_id == user_id)\n .where(\n (PersonalAccessToken.expires_at.is_(None))\n | (PersonalAccessToken.expires_at > datetime.now(timezone.utc))\n )\n .order_by(PersonalAccessToken.created_at.desc())\n ).all()\n )\n\n\ndef revoke_pat(db_session: Session, pat_id: int, user_id: UUID) -> bool:\n \"\"\"Revoke PAT by setting expires_at=NOW() for immediate expiry.\n\n Returns True if revoked, False if not found, not owned by user, or already expired.\n \"\"\"\n now = datetime.now(timezone.utc)\n pat = db_session.scalar(\n select(PersonalAccessToken)\n .where(PersonalAccessToken.id == pat_id)\n .where(PersonalAccessToken.user_id == user_id)\n .where(\n (PersonalAccessToken.expires_at.is_(None))\n | (PersonalAccessToken.expires_at > now)\n ) # Only revoke active (non-expired) tokens\n )\n if not pat:\n return False\n\n # Revoke by setting expires_at to NOW() and marking as revoked for audit trail\n pat.expires_at = now\n pat.is_revoked = True\n db_session.commit()\n return True\n" }, { "path": "backend/onyx/db/permission_sync_attempt.py", "content": "\"\"\"Permission sync attempt CRUD operations and utilities.\n\nThis module contains all CRUD operations for both DocPermissionSyncAttempt\nand ExternalGroupPermissionSyncAttempt models, along with shared utilities.\n\"\"\"\n\nfrom typing import Any\nfrom typing import cast\n\nfrom sqlalchemy import delete\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.engine.cursor import CursorResult\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import DocPermissionSyncAttempt\nfrom onyx.db.models import ExternalGroupPermissionSyncAttempt\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\n\nlogger = setup_logger()\n\n\n# =============================================================================\n# DOC PERMISSION SYNC ATTEMPT CRUD\n# =============================================================================\n\n\ndef create_doc_permission_sync_attempt(\n connector_credential_pair_id: int,\n db_session: Session,\n) -> int:\n \"\"\"Create a new doc permission sync attempt.\n\n Args:\n connector_credential_pair_id: The ID of the connector credential pair\n db_session: The database session\n\n Returns:\n The ID of the created attempt\n \"\"\"\n attempt = DocPermissionSyncAttempt(\n connector_credential_pair_id=connector_credential_pair_id,\n status=PermissionSyncStatus.NOT_STARTED,\n )\n db_session.add(attempt)\n db_session.commit()\n\n return attempt.id\n\n\ndef get_doc_permission_sync_attempt(\n db_session: Session,\n attempt_id: int,\n eager_load_connector: bool = False,\n) -> DocPermissionSyncAttempt | None:\n \"\"\"Get a doc permission sync attempt by ID.\n\n Args:\n db_session: The database session\n attempt_id: The ID of the attempt\n eager_load_connector: If True, eagerly loads the connector and cc_pair relationships\n\n Returns:\n The attempt if found, None otherwise\n \"\"\"\n stmt = select(DocPermissionSyncAttempt).where(\n DocPermissionSyncAttempt.id == attempt_id\n )\n\n if eager_load_connector:\n stmt = stmt.options(\n joinedload(DocPermissionSyncAttempt.connector_credential_pair).joinedload(\n ConnectorCredentialPair.connector\n )\n )\n\n return db_session.scalars(stmt).first()\n\n\ndef get_latest_doc_permission_sync_attempt_for_cc_pair(\n db_session: Session,\n connector_credential_pair_id: int,\n) -> DocPermissionSyncAttempt | None:\n \"\"\"Get the latest doc permission sync attempt for a connector credential pair.\"\"\"\n return db_session.execute(\n select(DocPermissionSyncAttempt)\n .where(\n DocPermissionSyncAttempt.connector_credential_pair_id\n == connector_credential_pair_id\n )\n .order_by(DocPermissionSyncAttempt.time_created.desc())\n .limit(1)\n ).scalar_one_or_none()\n\n\ndef get_recent_doc_permission_sync_attempts_for_cc_pair(\n cc_pair_id: int,\n limit: int,\n db_session: Session,\n) -> list[DocPermissionSyncAttempt]:\n \"\"\"Get recent doc permission sync attempts for a cc pair, most recent first.\"\"\"\n return list(\n db_session.execute(\n select(DocPermissionSyncAttempt)\n .where(DocPermissionSyncAttempt.connector_credential_pair_id == cc_pair_id)\n .order_by(DocPermissionSyncAttempt.time_created.desc())\n .limit(limit)\n ).scalars()\n )\n\n\ndef mark_doc_permission_sync_attempt_in_progress(\n attempt_id: int,\n db_session: Session,\n) -> DocPermissionSyncAttempt:\n \"\"\"Mark a doc permission sync attempt as IN_PROGRESS.\n Locks the row during update.\"\"\"\n try:\n attempt = db_session.execute(\n select(DocPermissionSyncAttempt)\n .where(DocPermissionSyncAttempt.id == attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if attempt.status != PermissionSyncStatus.NOT_STARTED:\n raise RuntimeError(\n f\"Doc permission sync attempt with ID '{attempt_id}' is not in NOT_STARTED status. \"\n f\"Current status is '{attempt.status}'.\"\n )\n\n attempt.status = PermissionSyncStatus.IN_PROGRESS\n attempt.time_started = func.now() # type: ignore\n db_session.commit()\n return attempt\n except Exception:\n db_session.rollback()\n logger.exception(\"mark_doc_permission_sync_attempt_in_progress exceptioned.\")\n raise\n\n\ndef mark_doc_permission_sync_attempt_failed(\n attempt_id: int,\n db_session: Session,\n error_message: str,\n) -> None:\n \"\"\"Mark a doc permission sync attempt as failed.\"\"\"\n try:\n attempt = db_session.execute(\n select(DocPermissionSyncAttempt)\n .where(DocPermissionSyncAttempt.id == attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if not attempt.time_started:\n attempt.time_started = func.now() # type: ignore\n attempt.status = PermissionSyncStatus.FAILED\n attempt.time_finished = func.now() # type: ignore\n attempt.error_message = error_message\n db_session.commit()\n\n # Add telemetry for permission sync attempt status change\n optional_telemetry(\n record_type=RecordType.PERMISSION_SYNC_COMPLETE,\n data={\n \"doc_permission_sync_attempt_id\": attempt_id,\n \"status\": PermissionSyncStatus.FAILED.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n except Exception:\n db_session.rollback()\n raise\n\n\ndef complete_doc_permission_sync_attempt(\n db_session: Session,\n attempt_id: int,\n total_docs_synced: int,\n docs_with_permission_errors: int,\n) -> DocPermissionSyncAttempt:\n \"\"\"Complete a doc permission sync attempt by updating progress and setting final status.\n\n This combines the progress update and final status marking into a single operation.\n If there were permission errors, the attempt is marked as COMPLETED_WITH_ERRORS,\n otherwise it's marked as SUCCESS.\n\n Args:\n db_session: The database session\n attempt_id: The ID of the attempt\n total_docs_synced: Total number of documents synced\n docs_with_permission_errors: Number of documents that had permission errors\n\n Returns:\n The completed attempt\n \"\"\"\n try:\n attempt = db_session.execute(\n select(DocPermissionSyncAttempt)\n .where(DocPermissionSyncAttempt.id == attempt_id)\n .with_for_update()\n ).scalar_one()\n\n # Update progress counters\n attempt.total_docs_synced = (attempt.total_docs_synced or 0) + total_docs_synced\n attempt.docs_with_permission_errors = (\n attempt.docs_with_permission_errors or 0\n ) + docs_with_permission_errors\n\n # Set final status based on whether there were errors\n if docs_with_permission_errors > 0:\n attempt.status = PermissionSyncStatus.COMPLETED_WITH_ERRORS\n else:\n attempt.status = PermissionSyncStatus.SUCCESS\n\n attempt.time_finished = func.now() # type: ignore\n db_session.commit()\n\n # Add telemetry\n optional_telemetry(\n record_type=RecordType.PERMISSION_SYNC_COMPLETE,\n data={\n \"doc_permission_sync_attempt_id\": attempt_id,\n \"status\": attempt.status.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n return attempt\n except Exception:\n db_session.rollback()\n logger.exception(\"complete_doc_permission_sync_attempt exceptioned.\")\n raise\n\n\n# =============================================================================\n# EXTERNAL GROUP PERMISSION SYNC ATTEMPT CRUD\n# =============================================================================\n\n\ndef create_external_group_sync_attempt(\n connector_credential_pair_id: int | None,\n db_session: Session,\n) -> int:\n \"\"\"Create a new external group sync attempt.\n\n Args:\n connector_credential_pair_id: The ID of the connector credential pair, or None for global syncs\n db_session: The database session\n\n Returns:\n The ID of the created attempt\n \"\"\"\n attempt = ExternalGroupPermissionSyncAttempt(\n connector_credential_pair_id=connector_credential_pair_id,\n status=PermissionSyncStatus.NOT_STARTED,\n )\n db_session.add(attempt)\n db_session.commit()\n\n return attempt.id\n\n\ndef get_external_group_sync_attempt(\n db_session: Session,\n attempt_id: int,\n eager_load_connector: bool = False,\n) -> ExternalGroupPermissionSyncAttempt | None:\n \"\"\"Get an external group sync attempt by ID.\n\n Args:\n db_session: The database session\n attempt_id: The ID of the attempt\n eager_load_connector: If True, eagerly loads the connector and cc_pair relationships\n\n Returns:\n The attempt if found, None otherwise\n \"\"\"\n stmt = select(ExternalGroupPermissionSyncAttempt).where(\n ExternalGroupPermissionSyncAttempt.id == attempt_id\n )\n\n if eager_load_connector:\n stmt = stmt.options(\n joinedload(\n ExternalGroupPermissionSyncAttempt.connector_credential_pair\n ).joinedload(ConnectorCredentialPair.connector)\n )\n\n return db_session.scalars(stmt).first()\n\n\ndef get_recent_external_group_sync_attempts_for_cc_pair(\n cc_pair_id: int | None,\n limit: int,\n db_session: Session,\n) -> list[ExternalGroupPermissionSyncAttempt]:\n \"\"\"Get recent external group sync attempts for a cc pair, most recent first.\n If cc_pair_id is None, gets global group sync attempts.\"\"\"\n stmt = select(ExternalGroupPermissionSyncAttempt)\n\n if cc_pair_id is not None:\n stmt = stmt.where(\n ExternalGroupPermissionSyncAttempt.connector_credential_pair_id\n == cc_pair_id\n )\n else:\n stmt = stmt.where(\n ExternalGroupPermissionSyncAttempt.connector_credential_pair_id.is_(None)\n )\n\n return list(\n db_session.execute(\n stmt.order_by(ExternalGroupPermissionSyncAttempt.time_created.desc()).limit(\n limit\n )\n ).scalars()\n )\n\n\ndef mark_external_group_sync_attempt_in_progress(\n attempt_id: int,\n db_session: Session,\n) -> ExternalGroupPermissionSyncAttempt:\n \"\"\"Mark an external group sync attempt as IN_PROGRESS.\n Locks the row during update.\"\"\"\n try:\n attempt = db_session.execute(\n select(ExternalGroupPermissionSyncAttempt)\n .where(ExternalGroupPermissionSyncAttempt.id == attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if attempt.status != PermissionSyncStatus.NOT_STARTED:\n raise RuntimeError(\n f\"External group sync attempt with ID '{attempt_id}' is not in NOT_STARTED status. \"\n f\"Current status is '{attempt.status}'.\"\n )\n\n attempt.status = PermissionSyncStatus.IN_PROGRESS\n attempt.time_started = func.now() # type: ignore\n db_session.commit()\n return attempt\n except Exception:\n db_session.rollback()\n logger.exception(\"mark_external_group_sync_attempt_in_progress exceptioned.\")\n raise\n\n\ndef mark_external_group_sync_attempt_failed(\n attempt_id: int,\n db_session: Session,\n error_message: str,\n) -> None:\n \"\"\"Mark an external group sync attempt as failed.\"\"\"\n try:\n attempt = db_session.execute(\n select(ExternalGroupPermissionSyncAttempt)\n .where(ExternalGroupPermissionSyncAttempt.id == attempt_id)\n .with_for_update()\n ).scalar_one()\n\n if not attempt.time_started:\n attempt.time_started = func.now() # type: ignore\n attempt.status = PermissionSyncStatus.FAILED\n attempt.time_finished = func.now() # type: ignore\n attempt.error_message = error_message\n db_session.commit()\n\n # Add telemetry for permission sync attempt status change\n optional_telemetry(\n record_type=RecordType.PERMISSION_SYNC_COMPLETE,\n data={\n \"external_group_sync_attempt_id\": attempt_id,\n \"status\": PermissionSyncStatus.FAILED.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n except Exception:\n db_session.rollback()\n raise\n\n\ndef complete_external_group_sync_attempt(\n db_session: Session,\n attempt_id: int,\n total_users_processed: int,\n total_groups_processed: int,\n total_group_memberships_synced: int,\n errors_encountered: int = 0,\n) -> ExternalGroupPermissionSyncAttempt:\n \"\"\"Complete an external group sync attempt by updating progress and setting final status.\n\n This combines the progress update and final status marking into a single operation.\n If there were errors, the attempt is marked as COMPLETED_WITH_ERRORS,\n otherwise it's marked as SUCCESS.\n\n Args:\n db_session: The database session\n attempt_id: The ID of the attempt\n total_users_processed: Total users processed\n total_groups_processed: Total groups processed\n total_group_memberships_synced: Total group memberships synced\n errors_encountered: Number of errors encountered (determines if COMPLETED_WITH_ERRORS)\n\n Returns:\n The completed attempt\n \"\"\"\n try:\n attempt = db_session.execute(\n select(ExternalGroupPermissionSyncAttempt)\n .where(ExternalGroupPermissionSyncAttempt.id == attempt_id)\n .with_for_update()\n ).scalar_one()\n\n # Update progress counters\n attempt.total_users_processed = (\n attempt.total_users_processed or 0\n ) + total_users_processed\n attempt.total_groups_processed = (\n attempt.total_groups_processed or 0\n ) + total_groups_processed\n attempt.total_group_memberships_synced = (\n attempt.total_group_memberships_synced or 0\n ) + total_group_memberships_synced\n\n # Set final status based on whether there were errors\n if errors_encountered > 0:\n attempt.status = PermissionSyncStatus.COMPLETED_WITH_ERRORS\n else:\n attempt.status = PermissionSyncStatus.SUCCESS\n\n attempt.time_finished = func.now() # type: ignore\n db_session.commit()\n\n # Add telemetry\n optional_telemetry(\n record_type=RecordType.PERMISSION_SYNC_COMPLETE,\n data={\n \"external_group_sync_attempt_id\": attempt_id,\n \"status\": attempt.status.value,\n \"cc_pair_id\": attempt.connector_credential_pair_id,\n },\n )\n return attempt\n except Exception:\n db_session.rollback()\n logger.exception(\"complete_external_group_sync_attempt exceptioned.\")\n raise\n\n\n# =============================================================================\n# DELETION FUNCTIONS\n# =============================================================================\n\n\ndef delete_doc_permission_sync_attempts__no_commit(\n db_session: Session,\n cc_pair_id: int,\n) -> int:\n \"\"\"Delete all doc permission sync attempts for a connector credential pair.\n\n This does not commit the transaction. It should be used within an existing transaction.\n\n Args:\n db_session: The database session\n cc_pair_id: The connector credential pair ID\n\n Returns:\n The number of attempts deleted\n \"\"\"\n stmt = delete(DocPermissionSyncAttempt).where(\n DocPermissionSyncAttempt.connector_credential_pair_id == cc_pair_id\n )\n result = cast(CursorResult[Any], db_session.execute(stmt))\n return result.rowcount or 0\n\n\ndef delete_external_group_permission_sync_attempts__no_commit(\n db_session: Session,\n cc_pair_id: int,\n) -> int:\n \"\"\"Delete all external group permission sync attempts for a connector credential pair.\n\n This does not commit the transaction. It should be used within an existing transaction.\n\n Args:\n db_session: The database session\n cc_pair_id: The connector credential pair ID\n\n Returns:\n The number of attempts deleted\n \"\"\"\n stmt = delete(ExternalGroupPermissionSyncAttempt).where(\n ExternalGroupPermissionSyncAttempt.connector_credential_pair_id == cc_pair_id\n )\n result = cast(CursorResult[Any], db_session.execute(stmt))\n return result.rowcount or 0\n" }, { "path": "backend/onyx/db/permissions.py", "content": "\"\"\"\nDB operations for recomputing user effective_permissions.\n\nThese live in onyx/db/ (not onyx/auth/) because they are pure DB operations\nthat query PermissionGrant rows and update the User.effective_permissions\nJSONB column. Keeping them here avoids circular imports when called from\nother onyx/db/ modules such as users.py.\n\"\"\"\n\nfrom collections import defaultdict\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import PermissionGrant\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\n\n\ndef recompute_user_permissions__no_commit(\n user_ids: UUID | str | list[UUID] | list[str], db_session: Session\n) -> None:\n \"\"\"Recompute granted permissions for one or more users.\n\n Accepts a single UUID or a list. Uses a single query regardless of\n how many users are passed, avoiding N+1 issues.\n\n Stores only directly granted permissions — implication expansion\n happens at read time via get_effective_permissions().\n\n Does NOT commit — caller must commit the session.\n \"\"\"\n if isinstance(user_ids, (UUID, str)):\n uid_list = [user_ids]\n else:\n uid_list = list(user_ids)\n\n if not uid_list:\n return\n\n # Single query to fetch ALL permissions for these users across ALL their\n # groups (a user may belong to multiple groups with different grants).\n rows = db_session.execute(\n select(User__UserGroup.user_id, PermissionGrant.permission)\n .join(\n PermissionGrant,\n PermissionGrant.group_id == User__UserGroup.user_group_id,\n )\n .where(\n User__UserGroup.user_id.in_(uid_list),\n PermissionGrant.is_deleted.is_(False),\n )\n ).all()\n\n # Group permissions by user; users with no grants get an empty set.\n perms_by_user: dict[UUID | str, set[str]] = defaultdict(set)\n for uid in uid_list:\n perms_by_user[uid] # ensure every user has an entry\n for uid, perm in rows:\n perms_by_user[uid].add(perm.value)\n\n for uid, perms in perms_by_user.items():\n db_session.execute(\n update(User)\n .where(User.id == uid) # type: ignore[arg-type]\n .values(effective_permissions=sorted(perms))\n )\n\n\ndef recompute_permissions_for_group__no_commit(\n group_id: int, db_session: Session\n) -> None:\n \"\"\"Recompute granted permissions for all users in a group.\n\n Does NOT commit — caller must commit the session.\n \"\"\"\n user_ids: list[UUID] = [\n uid\n for uid in db_session.execute(\n select(User__UserGroup.user_id).where(\n User__UserGroup.user_group_id == group_id,\n User__UserGroup.user_id.isnot(None),\n )\n )\n .scalars()\n .all()\n if uid is not None\n ]\n\n if not user_ids:\n return\n\n recompute_user_permissions__no_commit(user_ids, db_session)\n" }, { "path": "backend/onyx/db/persona.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\nfrom enum import Enum\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom sqlalchemy import exists\nfrom sqlalchemy import func\nfrom sqlalchemy import not_\nfrom sqlalchemy import or_\nfrom sqlalchemy import Select\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import aliased\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.hierarchy_access import get_user_external_group_ids\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.app_configs import CURATORS_CANNOT_VIEW_OR_EDIT_NON_OWNED_ASSISTANTS\nfrom onyx.configs.constants import DEFAULT_PERSONA_ID\nfrom onyx.configs.constants import NotificationType\nfrom onyx.db.constants import SLACK_BOT_PERSONA_PREFIX\nfrom onyx.db.document_access import get_accessible_documents_by_ids\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Document\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import FederatedConnector__DocumentSet\nfrom onyx.db.models import HierarchyNode\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__User\nfrom onyx.db.models import Persona__UserGroup\nfrom onyx.db.models import PersonaLabel\nfrom onyx.db.models import StarterMessage\nfrom onyx.db.models import Tool\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserFile\nfrom onyx.db.models import UserGroup\nfrom onyx.db.notification import create_notification\nfrom onyx.server.features.persona.models import FullPersonaSnapshot\nfrom onyx.server.features.persona.models import MinimalPersonaSnapshot\nfrom onyx.server.features.persona.models import PersonaSharedNotificationData\nfrom onyx.server.features.persona.models import PersonaSnapshot\nfrom onyx.server.features.persona.models import PersonaUpsertRequest\nfrom onyx.server.features.tool.tool_visibility import should_expose_tool_to_fe\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\nlogger = setup_logger()\n\n\ndef get_default_behavior_persona(\n db_session: Session,\n eager_load_for_tools: bool = False,\n) -> Persona | None:\n stmt = select(Persona).where(Persona.id == DEFAULT_PERSONA_ID)\n if eager_load_for_tools:\n stmt = stmt.options(\n selectinload(Persona.tools),\n selectinload(Persona.document_sets),\n selectinload(Persona.attached_documents),\n selectinload(Persona.hierarchy_nodes),\n )\n return db_session.scalars(stmt).first()\n\n\nclass PersonaLoadType(Enum):\n NONE = \"none\"\n MINIMAL = \"minimal\"\n FULL = \"full\"\n\n\ndef _add_user_filters(\n stmt: Select[tuple[Persona]], user: User, get_editable: bool = True\n) -> Select[tuple[Persona]]:\n if user.role == UserRole.ADMIN:\n return stmt\n\n stmt = stmt.distinct()\n Persona__UG = aliased(Persona__UserGroup)\n User__UG = aliased(User__UserGroup)\n \"\"\"\n Here we select cc_pairs by relation:\n User -> User__UserGroup -> Persona__UserGroup -> Persona\n \"\"\"\n stmt = (\n stmt.outerjoin(Persona__UG)\n .outerjoin(\n User__UserGroup,\n User__UserGroup.user_group_id == Persona__UG.user_group_id,\n )\n .outerjoin(\n Persona__User,\n Persona__User.persona_id == Persona.id,\n )\n )\n \"\"\"\n Filter Personas by:\n - if the user is in the user_group that owns the Persona\n - if the user is not a global_curator, they must also have a curator relationship\n to the user_group\n - if editing is being done, we also filter out Personas that are owned by groups\n that the user isn't a curator for\n - if we are not editing, we show all Personas in the groups the user is a curator\n for (as well as public Personas)\n - if we are not editing, we return all Personas directly connected to the user\n \"\"\"\n\n # Anonymous users only see public Personas\n if user.is_anonymous:\n where_clause = Persona.is_public == True # noqa: E712\n return stmt.where(where_clause)\n\n # If curator ownership restriction is enabled, curators can only access their own assistants\n if CURATORS_CANNOT_VIEW_OR_EDIT_NON_OWNED_ASSISTANTS and user.role in [\n UserRole.CURATOR,\n UserRole.GLOBAL_CURATOR,\n ]:\n where_clause = (Persona.user_id == user.id) | (Persona.user_id.is_(None))\n return stmt.where(where_clause)\n\n where_clause = User__UserGroup.user_id == user.id\n if user.role == UserRole.CURATOR and get_editable:\n where_clause &= User__UserGroup.is_curator == True # noqa: E712\n if get_editable:\n user_groups = select(User__UG.user_group_id).where(User__UG.user_id == user.id)\n if user.role == UserRole.CURATOR:\n user_groups = user_groups.where(User__UG.is_curator == True) # noqa: E712\n where_clause &= (\n ~exists()\n .where(Persona__UG.persona_id == Persona.id)\n .where(~Persona__UG.user_group_id.in_(user_groups))\n .correlate(Persona)\n )\n else:\n # Group the public persona conditions\n public_condition = (Persona.is_public == True) & ( # noqa: E712\n Persona.is_listed == True # noqa: E712\n )\n\n where_clause |= public_condition\n where_clause |= Persona__User.user_id == user.id\n\n where_clause |= Persona.user_id == user.id\n\n return stmt.where(where_clause)\n\n\ndef fetch_persona_by_id_for_user(\n db_session: Session, persona_id: int, user: User, get_editable: bool = True\n) -> Persona:\n stmt = select(Persona).where(Persona.id == persona_id).distinct()\n stmt = _add_user_filters(stmt=stmt, user=user, get_editable=get_editable)\n persona = db_session.scalars(stmt).one_or_none()\n if not persona:\n raise HTTPException(\n status_code=403,\n detail=f\"Persona with ID {persona_id} does not exist or user is not authorized to access it\",\n )\n return persona\n\n\ndef get_best_persona_id_for_user(\n db_session: Session, user: User, persona_id: int | None = None\n) -> int | None:\n if persona_id is not None:\n stmt = select(Persona).where(Persona.id == persona_id).distinct()\n stmt = _add_user_filters(\n stmt=stmt,\n user=user,\n # We don't want to filter by editable here, we just want to see if the\n # persona is usable by the user\n get_editable=False,\n )\n persona = db_session.scalars(stmt).one_or_none()\n if persona:\n return persona.id\n\n # If the persona is not found, or the slack bot is using doc sets instead of personas,\n # we need to find the best persona for the user\n # This is the persona with the highest display priority that the user has access to\n stmt = select(Persona).order_by(Persona.display_priority.desc()).distinct()\n stmt = _add_user_filters(stmt=stmt, user=user, get_editable=True)\n persona = db_session.scalars(stmt).one_or_none()\n return persona.id if persona else None\n\n\ndef _get_persona_by_name(\n persona_name: str, user: User | None, db_session: Session\n) -> Persona | None:\n \"\"\"Fetch a persona by name with access control.\n\n Access rules:\n - user=None (system operations): can see all personas\n - Admin users: can see all personas\n - Non-admin users: can only see their own personas\n \"\"\"\n stmt = select(Persona).where(Persona.name == persona_name)\n if user and user.role != UserRole.ADMIN:\n stmt = stmt.where(Persona.user_id == user.id)\n result = db_session.execute(stmt).scalar_one_or_none()\n return result\n\n\ndef update_persona_access(\n persona_id: int,\n creator_user_id: UUID | None,\n db_session: Session,\n is_public: bool | None = None,\n user_ids: list[UUID] | None = None,\n group_ids: list[int] | None = None,\n) -> None:\n \"\"\"Updates the access settings for a persona including public status and user shares.\n\n NOTE: Callers are responsible for committing.\"\"\"\n\n needs_sync = False\n if is_public is not None:\n needs_sync = True\n persona = db_session.query(Persona).filter(Persona.id == persona_id).first()\n if persona:\n persona.is_public = is_public\n\n # NOTE: For user-ids and group-ids, `None` means \"leave unchanged\", `[]` means \"clear all shares\",\n # and a non-empty list means \"replace with these shares\".\n if user_ids is not None:\n needs_sync = True\n db_session.query(Persona__User).filter(\n Persona__User.persona_id == persona_id\n ).delete(synchronize_session=\"fetch\")\n\n for user_uuid in user_ids:\n db_session.add(Persona__User(persona_id=persona_id, user_id=user_uuid))\n if user_uuid != creator_user_id:\n create_notification(\n user_id=user_uuid,\n notif_type=NotificationType.PERSONA_SHARED,\n title=\"A new agent was shared with you!\",\n db_session=db_session,\n additional_data=PersonaSharedNotificationData(\n persona_id=persona_id,\n ).model_dump(),\n )\n\n # MIT doesn't support group-based sharing, so we allow clearing (no-op since\n # there shouldn't be any) but raise an error if trying to add actual groups.\n if group_ids is not None:\n needs_sync = True\n db_session.query(Persona__UserGroup).filter(\n Persona__UserGroup.persona_id == persona_id\n ).delete(synchronize_session=\"fetch\")\n\n if group_ids:\n raise NotImplementedError(\"Onyx MIT does not support group-based sharing\")\n\n # When sharing changes, user file ACLs need to be updated in the vector DB\n if needs_sync:\n mark_persona_user_files_for_sync(persona_id, db_session)\n\n\ndef create_update_persona(\n persona_id: int | None,\n create_persona_request: PersonaUpsertRequest,\n user: User,\n db_session: Session,\n) -> FullPersonaSnapshot:\n \"\"\"Higher level function than upsert_persona, although either is valid to use.\"\"\"\n # Permission to actually use these is checked later\n\n try:\n # Featured persona validation\n if create_persona_request.is_featured:\n # Curators can edit featured personas, but not make them\n # TODO this will be reworked soon with RBAC permissions feature\n if user.role == UserRole.CURATOR or user.role == UserRole.GLOBAL_CURATOR:\n pass\n elif user.role != UserRole.ADMIN:\n raise ValueError(\"Only admins can make a featured persona\")\n\n # Convert incoming string UUIDs to UUID objects for DB operations\n converted_user_file_ids = None\n if create_persona_request.user_file_ids is not None:\n try:\n converted_user_file_ids = [\n UUID(str_id) for str_id in create_persona_request.user_file_ids\n ]\n except Exception:\n raise ValueError(\"Invalid user_file_ids; must be UUID strings\")\n\n persona = upsert_persona(\n persona_id=persona_id,\n user=user,\n db_session=db_session,\n description=create_persona_request.description,\n name=create_persona_request.name,\n document_set_ids=create_persona_request.document_set_ids,\n tool_ids=create_persona_request.tool_ids,\n is_public=create_persona_request.is_public,\n llm_model_provider_override=create_persona_request.llm_model_provider_override,\n llm_model_version_override=create_persona_request.llm_model_version_override,\n starter_messages=create_persona_request.starter_messages,\n system_prompt=create_persona_request.system_prompt,\n task_prompt=create_persona_request.task_prompt,\n datetime_aware=create_persona_request.datetime_aware,\n replace_base_system_prompt=create_persona_request.replace_base_system_prompt,\n uploaded_image_id=create_persona_request.uploaded_image_id,\n icon_name=create_persona_request.icon_name,\n display_priority=create_persona_request.display_priority,\n remove_image=create_persona_request.remove_image,\n search_start_date=create_persona_request.search_start_date,\n label_ids=create_persona_request.label_ids,\n is_featured=create_persona_request.is_featured,\n user_file_ids=converted_user_file_ids,\n commit=False,\n hierarchy_node_ids=create_persona_request.hierarchy_node_ids,\n document_ids=create_persona_request.document_ids,\n )\n\n versioned_update_persona_access = fetch_versioned_implementation(\n \"onyx.db.persona\", \"update_persona_access\"\n )\n\n versioned_update_persona_access(\n persona_id=persona.id,\n creator_user_id=user.id,\n db_session=db_session,\n user_ids=create_persona_request.users,\n group_ids=create_persona_request.groups,\n )\n db_session.commit()\n\n except ValueError as e:\n logger.exception(\"Failed to create persona\")\n raise HTTPException(status_code=400, detail=str(e))\n\n return FullPersonaSnapshot.from_model(persona)\n\n\ndef update_persona_shared(\n persona_id: int,\n user_ids: list[UUID] | None,\n user: User,\n db_session: Session,\n group_ids: list[int] | None = None,\n is_public: bool | None = None,\n label_ids: list[int] | None = None,\n) -> None:\n \"\"\"Simplified version of `create_update_persona` which only touches the\n accessibility rather than any of the logic (e.g. prompt, connected data sources,\n etc.).\"\"\"\n persona = fetch_persona_by_id_for_user(\n db_session=db_session, persona_id=persona_id, user=user, get_editable=True\n )\n\n if user and user.role != UserRole.ADMIN and persona.user_id != user.id:\n raise PermissionError(\"You don't have permission to modify this persona\")\n\n versioned_update_persona_access = fetch_versioned_implementation(\n \"onyx.db.persona\", \"update_persona_access\"\n )\n versioned_update_persona_access(\n persona_id=persona_id,\n creator_user_id=user.id,\n db_session=db_session,\n is_public=is_public,\n user_ids=user_ids,\n group_ids=group_ids,\n )\n\n if label_ids is not None:\n labels = (\n db_session.query(PersonaLabel).filter(PersonaLabel.id.in_(label_ids)).all()\n )\n if len(labels) != len(label_ids):\n raise ValueError(\"Some label IDs were not found in the database\")\n persona.labels.clear()\n persona.labels = labels\n\n db_session.commit()\n\n\ndef update_persona_public_status(\n persona_id: int,\n is_public: bool,\n db_session: Session,\n user: User,\n) -> None:\n persona = fetch_persona_by_id_for_user(\n db_session=db_session, persona_id=persona_id, user=user, get_editable=True\n )\n if user.role != UserRole.ADMIN and persona.user_id != user.id:\n raise ValueError(\"You don't have permission to modify this persona\")\n\n persona.is_public = is_public\n db_session.commit()\n\n\ndef _build_persona_filters(\n stmt: Select[tuple[Persona]],\n include_default: bool,\n include_slack_bot_personas: bool,\n include_deleted: bool,\n) -> Select[tuple[Persona]]:\n \"\"\"Filters which Personas are included in the query.\n\n Args:\n stmt: The base query to filter.\n include_default: If True, includes builtin/default personas.\n include_slack_bot_personas: If True, includes Slack bot personas.\n include_deleted: If True, includes deleted personas.\n\n Returns:\n The modified query with the filters applied.\n \"\"\"\n if not include_default:\n stmt = stmt.where(Persona.builtin_persona.is_(False))\n if not include_slack_bot_personas:\n stmt = stmt.where(not_(Persona.name.startswith(SLACK_BOT_PERSONA_PREFIX)))\n if not include_deleted:\n stmt = stmt.where(Persona.deleted.is_(False))\n return stmt\n\n\ndef get_minimal_persona_snapshots_for_user(\n user: User,\n db_session: Session,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> list[MinimalPersonaSnapshot]:\n stmt = select(Persona)\n stmt = _add_user_filters(stmt, user, get_editable)\n stmt = _build_persona_filters(\n stmt, include_default, include_slack_bot_personas, include_deleted\n )\n stmt = stmt.options(\n selectinload(Persona.tools),\n selectinload(Persona.labels),\n selectinload(Persona.document_sets).options(\n selectinload(DocumentSet.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSet.users),\n selectinload(DocumentSet.groups),\n selectinload(DocumentSet.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n ),\n selectinload(Persona.hierarchy_nodes),\n selectinload(Persona.attached_documents).selectinload(\n Document.parent_hierarchy_node\n ),\n selectinload(Persona.user),\n )\n results = db_session.scalars(stmt).all()\n return [MinimalPersonaSnapshot.from_model(persona) for persona in results]\n\n\ndef get_persona_snapshots_for_user(\n user: User,\n db_session: Session,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> list[PersonaSnapshot]:\n stmt = select(Persona)\n stmt = _add_user_filters(stmt, user, get_editable)\n stmt = _build_persona_filters(\n stmt, include_default, include_slack_bot_personas, include_deleted\n )\n stmt = stmt.options(\n selectinload(Persona.tools),\n selectinload(Persona.hierarchy_nodes),\n selectinload(Persona.attached_documents).selectinload(\n Document.parent_hierarchy_node\n ),\n selectinload(Persona.labels),\n selectinload(Persona.document_sets).options(\n selectinload(DocumentSet.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSet.users),\n selectinload(DocumentSet.groups),\n selectinload(DocumentSet.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n ),\n selectinload(Persona.user),\n selectinload(Persona.user_files),\n selectinload(Persona.users),\n selectinload(Persona.groups),\n )\n\n results = db_session.scalars(stmt).all()\n return [PersonaSnapshot.from_model(persona) for persona in results]\n\n\ndef get_persona_count_for_user(\n user: User,\n db_session: Session,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> int:\n \"\"\"Counts the total number of personas accessible to the user.\n\n Args:\n user: The user to filter personas for. If None and auth is disabled,\n assumes the user is an admin. Otherwise, if None shows only public\n personas.\n db_session: Database session for executing queries.\n get_editable: If True, only returns personas the user can edit.\n include_default: If True, includes builtin/default personas.\n include_slack_bot_personas: If True, includes Slack bot personas.\n include_deleted: If True, includes deleted personas.\n\n Returns:\n Total count of personas matching the filters and user permissions.\n \"\"\"\n stmt = _build_persona_base_query(\n user=user,\n get_editable=get_editable,\n include_default=include_default,\n include_slack_bot_personas=include_slack_bot_personas,\n include_deleted=include_deleted,\n )\n # Convert to count query.\n count_stmt = stmt.with_only_columns(func.count(func.distinct(Persona.id))).order_by(\n None\n )\n return db_session.scalar(count_stmt) or 0\n\n\ndef get_minimal_persona_snapshots_paginated(\n user: User,\n db_session: Session,\n page_num: int,\n page_size: int,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> list[MinimalPersonaSnapshot]:\n \"\"\"Gets a single page of minimal persona snapshots with ordering.\n\n Personas are ordered by display_priority (ASC, nulls last) then by ID (ASC\n distance from 0).\n\n Args:\n user: The user to filter personas for. If None and auth is disabled,\n assumes the user is an admin. Otherwise, if None shows only public\n personas.\n db_session: Database session for executing queries.\n page_num: Zero-indexed page number (e.g., 0 for the first page).\n page_size: Number of items per page.\n get_editable: If True, only returns personas the user can edit.\n include_default: If True, includes builtin/default personas.\n include_slack_bot_personas: If True, includes Slack bot personas.\n include_deleted: If True, includes deleted personas.\n\n Returns:\n List of MinimalPersonaSnapshot objects for the requested page, ordered\n by display_priority (nulls last) then ID.\n \"\"\"\n stmt = _get_paginated_persona_query(\n user,\n page_num,\n page_size,\n get_editable,\n include_default,\n include_slack_bot_personas,\n include_deleted,\n )\n # Do eager loading of columns we know MinimalPersonaSnapshot.from_model will\n # need.\n stmt = stmt.options(\n selectinload(Persona.tools),\n selectinload(Persona.hierarchy_nodes),\n selectinload(Persona.attached_documents).selectinload(\n Document.parent_hierarchy_node\n ),\n selectinload(Persona.labels),\n selectinload(Persona.document_sets).options(\n selectinload(DocumentSet.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSet.users),\n selectinload(DocumentSet.groups),\n selectinload(DocumentSet.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n ),\n selectinload(Persona.user),\n )\n\n results = db_session.scalars(stmt).all()\n return [MinimalPersonaSnapshot.from_model(persona) for persona in results]\n\n\ndef get_persona_snapshots_paginated(\n user: User,\n db_session: Session,\n page_num: int,\n page_size: int,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> list[PersonaSnapshot]:\n \"\"\"Gets a single page of persona snapshots (admin view) with ordering.\n\n Personas are ordered by display_priority (ASC, nulls last) then by ID (ASC\n distance from 0).\n\n This function returns PersonaSnapshot objects which contain more detailed\n information than MinimalPersonaSnapshot, used for admin views.\n\n Args:\n user: The user to filter personas for. If None and auth is disabled,\n assumes the user is an admin. Otherwise, if None shows only public\n personas.\n db_session: Database session for executing queries.\n page_num: Zero-indexed page number (e.g., 0 for the first page).\n page_size: Number of items per page.\n get_editable: If True, only returns personas the user can edit.\n include_default: If True, includes builtin/default personas.\n include_slack_bot_personas: If True, includes Slack bot personas.\n include_deleted: If True, includes deleted personas.\n\n Returns:\n List of PersonaSnapshot objects for the requested page, ordered by\n display_priority (nulls last) then ID.\n \"\"\"\n stmt = _get_paginated_persona_query(\n user,\n page_num,\n page_size,\n get_editable,\n include_default,\n include_slack_bot_personas,\n include_deleted,\n )\n # Do eager loading of columns we know PersonaSnapshot.from_model will need.\n stmt = stmt.options(\n selectinload(Persona.tools),\n selectinload(Persona.hierarchy_nodes),\n selectinload(Persona.attached_documents).selectinload(\n Document.parent_hierarchy_node\n ),\n selectinload(Persona.labels),\n selectinload(Persona.document_sets).options(\n selectinload(DocumentSet.connector_credential_pairs).selectinload(\n ConnectorCredentialPair.connector\n ),\n selectinload(DocumentSet.users),\n selectinload(DocumentSet.groups),\n selectinload(DocumentSet.federated_connectors).selectinload(\n FederatedConnector__DocumentSet.federated_connector\n ),\n ),\n selectinload(Persona.user),\n selectinload(Persona.user_files),\n selectinload(Persona.users),\n selectinload(Persona.groups),\n )\n\n results = db_session.scalars(stmt).all()\n return [PersonaSnapshot.from_model(persona) for persona in results]\n\n\ndef _get_paginated_persona_query(\n user: User,\n page_num: int,\n page_size: int,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> Select[tuple[Persona]]:\n \"\"\"Builds a paginated query on personas ordered on display_priority and id.\n\n Personas are ordered by display_priority (ASC, nulls last) then by ID (ASC\n distance from 0) to match the frontend personaComparator() logic.\n\n Args:\n user: The user to filter personas for. If None and auth is disabled,\n assumes the user is an admin. Otherwise, if None shows only public\n personas.\n page_num: Zero-indexed page number (e.g., 0 for the first page).\n page_size: Number of items per page.\n get_editable: If True, only returns personas the user can edit.\n include_default: If True, includes builtin/default personas.\n include_slack_bot_personas: If True, includes Slack bot personas.\n include_deleted: If True, includes deleted personas.\n\n Returns:\n SQLAlchemy Select statement with all filters, ordering, and pagination\n applied.\n \"\"\"\n stmt = _build_persona_base_query(\n user=user,\n get_editable=get_editable,\n include_default=include_default,\n include_slack_bot_personas=include_slack_bot_personas,\n include_deleted=include_deleted,\n )\n # Add the abs(id) expression to the SELECT list (required for DISTINCT +\n # ORDER BY).\n stmt = stmt.add_columns(func.abs(Persona.id).label(\"abs_id\"))\n # Apply ordering.\n stmt = stmt.order_by(\n Persona.display_priority.asc().nullslast(),\n func.abs(Persona.id).asc(),\n )\n # Apply pagination.\n stmt = stmt.offset(page_num * page_size).limit(page_size)\n return stmt\n\n\ndef _build_persona_base_query(\n user: User,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> Select[tuple[Persona]]:\n \"\"\"Builds a base persona query with all user and persona filters applied.\n\n This helper constructs a filtered query that can then be customized for\n counting, pagination, or full retrieval.\n\n Args:\n user: The user to filter personas for. If None and auth is disabled,\n assumes the user is an admin. Otherwise, if None shows only public\n personas.\n get_editable: If True, only returns personas the user can edit.\n include_default: If True, includes builtin/default personas.\n include_slack_bot_personas: If True, includes Slack bot personas.\n include_deleted: If True, includes deleted personas.\n\n Returns:\n SQLAlchemy Select statement with all filters applied.\n \"\"\"\n stmt = select(Persona)\n stmt = _add_user_filters(stmt, user, get_editable)\n stmt = _build_persona_filters(\n stmt, include_default, include_slack_bot_personas, include_deleted\n )\n return stmt\n\n\ndef get_raw_personas_for_user(\n user: User,\n db_session: Session,\n get_editable: bool = True,\n include_default: bool = True,\n include_slack_bot_personas: bool = False,\n include_deleted: bool = False,\n) -> Sequence[Persona]:\n stmt = _build_persona_base_query(\n user, get_editable, include_default, include_slack_bot_personas, include_deleted\n )\n return db_session.scalars(stmt).all()\n\n\ndef get_personas(db_session: Session) -> Sequence[Persona]:\n \"\"\"WARNING: Unsafe, can fetch personas from all users.\"\"\"\n stmt = select(Persona).distinct()\n stmt = stmt.where(not_(Persona.name.startswith(SLACK_BOT_PERSONA_PREFIX)))\n stmt = stmt.where(Persona.deleted.is_(False))\n return db_session.execute(stmt).unique().scalars().all()\n\n\ndef mark_persona_as_deleted(\n persona_id: int,\n user: User,\n db_session: Session,\n) -> None:\n persona = get_persona_by_id(persona_id=persona_id, user=user, db_session=db_session)\n persona.deleted = True\n affected_file_ids = [uf.id for uf in persona.user_files]\n if affected_file_ids:\n _mark_files_need_persona_sync(db_session, affected_file_ids)\n db_session.commit()\n\n\ndef mark_persona_as_not_deleted(\n persona_id: int,\n user: User,\n db_session: Session,\n) -> None:\n persona = get_persona_by_id(\n persona_id=persona_id, user=user, db_session=db_session, include_deleted=True\n )\n if not persona.deleted:\n raise ValueError(f\"Persona with ID {persona_id} is not deleted.\")\n persona.deleted = False\n affected_file_ids = [uf.id for uf in persona.user_files]\n if affected_file_ids:\n _mark_files_need_persona_sync(db_session, affected_file_ids)\n db_session.commit()\n\n\ndef mark_delete_persona_by_name(\n persona_name: str, db_session: Session, is_default: bool = True\n) -> None:\n stmt = (\n update(Persona)\n .where(Persona.name == persona_name, Persona.builtin_persona == is_default)\n .values(deleted=True)\n )\n\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef update_personas_display_priority(\n display_priority_map: dict[int, int],\n db_session: Session,\n user: User,\n commit_db_txn: bool = False,\n) -> None:\n \"\"\"Updates the display priorities of the specified Personas.\n\n Args:\n display_priority_map: A map of persona IDs to intended display\n priorities.\n db_session: Database session for executing queries.\n user: The user to filter personas for. If None and auth is disabled,\n assumes the user is an admin. Otherwise, if None shows only public\n personas.\n commit_db_txn: If True, commits the database transaction after\n updating the display priorities. Defaults to False.\n\n Raises:\n ValueError: The caller tried to update a persona for which the user does\n not have access.\n \"\"\"\n # No-op to save a query if it is not necessary.\n if len(display_priority_map) == 0:\n return\n\n personas = get_raw_personas_for_user(\n user,\n db_session,\n get_editable=False,\n include_default=True,\n include_slack_bot_personas=True,\n include_deleted=True,\n )\n available_personas_map: dict[int, Persona] = {\n persona.id: persona for persona in personas\n }\n\n for persona_id, priority in display_priority_map.items():\n if persona_id not in available_personas_map:\n raise ValueError(\n f\"Invalid persona ID provided: Persona with ID {persona_id} was not found for this user.\"\n )\n\n available_personas_map[persona_id].display_priority = priority\n\n if commit_db_txn:\n db_session.commit()\n\n\ndef mark_persona_user_files_for_sync(\n persona_id: int,\n db_session: Session,\n) -> None:\n \"\"\"When persona sharing changes, mark all of its user files for sync\n so that their ACLs get updated in the vector DB.\"\"\"\n persona = (\n db_session.query(Persona)\n .options(selectinload(Persona.user_files))\n .filter(Persona.id == persona_id)\n .first()\n )\n if not persona:\n return\n file_ids = [uf.id for uf in persona.user_files]\n _mark_files_need_persona_sync(db_session, file_ids)\n\n\ndef _mark_files_need_persona_sync(\n db_session: Session,\n user_file_ids: list[UUID],\n) -> None:\n \"\"\"Flag the given UserFile rows so the background sync task picks them up\n and updates their persona metadata in the vector DB.\"\"\"\n if not user_file_ids:\n return\n db_session.query(UserFile).filter(UserFile.id.in_(user_file_ids)).update(\n {UserFile.needs_persona_sync: True},\n synchronize_session=False,\n )\n\n\ndef upsert_persona(\n user: User | None,\n name: str,\n description: str,\n llm_model_provider_override: str | None,\n llm_model_version_override: str | None,\n starter_messages: list[StarterMessage] | None,\n # Embedded prompt fields\n system_prompt: str | None,\n task_prompt: str | None,\n datetime_aware: bool | None,\n is_public: bool,\n db_session: Session,\n document_set_ids: list[int] | None = None,\n tool_ids: list[int] | None = None,\n persona_id: int | None = None,\n commit: bool = True,\n uploaded_image_id: str | None = None,\n icon_name: str | None = None,\n display_priority: int | None = None,\n is_listed: bool = True,\n remove_image: bool | None = None,\n search_start_date: datetime | None = None,\n builtin_persona: bool = False,\n is_featured: bool | None = None,\n label_ids: list[int] | None = None,\n user_file_ids: list[UUID] | None = None,\n hierarchy_node_ids: list[int] | None = None,\n document_ids: list[str] | None = None,\n replace_base_system_prompt: bool = False,\n) -> Persona:\n \"\"\"\n NOTE: This operation cannot update persona configuration options that\n are core to the persona, such as its display priority and\n whether or not the assistant is a built-in / default assistant\n \"\"\"\n\n if persona_id is not None:\n existing_persona = db_session.query(Persona).filter_by(id=persona_id).first()\n else:\n existing_persona = _get_persona_by_name(\n persona_name=name, user=user, db_session=db_session\n )\n\n # Check for duplicate names when creating new personas\n # Deleted personas are allowed to be overwritten\n if existing_persona and not existing_persona.deleted:\n raise ValueError(\n f\"Assistant with name '{name}' already exists. Please rename your assistant.\"\n )\n\n if existing_persona and user:\n # this checks if the user has permission to edit the persona\n # will raise an Exception if the user does not have permission\n # Skip check if user is None (system/admin operation)\n existing_persona = fetch_persona_by_id_for_user(\n db_session=db_session,\n persona_id=existing_persona.id,\n user=user,\n get_editable=True,\n )\n\n # Fetch and attach tools by IDs\n tools = None\n if tool_ids is not None:\n tools = db_session.query(Tool).filter(Tool.id.in_(tool_ids)).all()\n if not tools and tool_ids:\n raise ValueError(\"Tools not found\")\n\n # Fetch and attach document_sets by IDs\n document_sets = None\n if document_set_ids is not None:\n document_sets = (\n db_session.query(DocumentSet)\n .filter(DocumentSet.id.in_(document_set_ids))\n .all()\n )\n if not document_sets and document_set_ids:\n raise ValueError(\"document_sets not found\")\n\n # Fetch and attach user_files by IDs\n user_files = None\n if user_file_ids is not None:\n user_files = (\n db_session.query(UserFile).filter(UserFile.id.in_(user_file_ids)).all()\n )\n if not user_files and user_file_ids:\n raise ValueError(\"user_files not found\")\n\n labels = None\n if label_ids is not None:\n labels = (\n db_session.query(PersonaLabel).filter(PersonaLabel.id.in_(label_ids)).all()\n )\n if len(labels) != len(label_ids):\n raise ValueError(\"Some label IDs were not found in the database\")\n\n # Fetch and attach hierarchy_nodes by IDs\n hierarchy_nodes = None\n if hierarchy_node_ids:\n hierarchy_nodes = (\n db_session.query(HierarchyNode)\n .filter(HierarchyNode.id.in_(hierarchy_node_ids))\n .all()\n )\n if not hierarchy_nodes and hierarchy_node_ids:\n raise ValueError(\"hierarchy_nodes not found\")\n\n # Fetch and attach documents by IDs, filtering for access permissions\n attached_documents = None\n if document_ids is not None:\n user_email = user.email if user else None\n external_group_ids = (\n get_user_external_group_ids(db_session, user) if user else []\n )\n attached_documents = get_accessible_documents_by_ids(\n db_session=db_session,\n document_ids=document_ids,\n user_email=user_email,\n external_group_ids=external_group_ids,\n )\n if not attached_documents and document_ids:\n raise ValueError(\"documents not found or not accessible\")\n\n # ensure all specified tools are valid\n if tools:\n validate_persona_tools(tools, db_session)\n\n if existing_persona:\n # Built-in personas can only be updated through YAML configuration.\n # This ensures that core system personas are not modified unintentionally.\n if existing_persona.builtin_persona and not builtin_persona:\n raise ValueError(\"Cannot update builtin persona with non-builtin.\")\n\n # The following update excludes `default`, `built-in`, and display priority.\n # Display priority is handled separately in the `display-priority` endpoint.\n # `default` and `built-in` properties can only be set when creating a persona.\n existing_persona.name = name\n existing_persona.description = description\n existing_persona.llm_model_provider_override = llm_model_provider_override\n existing_persona.llm_model_version_override = llm_model_version_override\n existing_persona.starter_messages = starter_messages\n existing_persona.deleted = False # Un-delete if previously deleted\n existing_persona.is_public = is_public\n if remove_image or uploaded_image_id:\n existing_persona.uploaded_image_id = uploaded_image_id\n existing_persona.icon_name = icon_name\n existing_persona.is_listed = is_listed\n existing_persona.search_start_date = search_start_date\n if label_ids is not None:\n existing_persona.labels.clear()\n existing_persona.labels = labels or []\n existing_persona.is_featured = (\n is_featured if is_featured is not None else existing_persona.is_featured\n )\n # Update embedded prompt fields if provided\n if system_prompt is not None:\n existing_persona.system_prompt = system_prompt\n if task_prompt is not None:\n existing_persona.task_prompt = task_prompt\n if datetime_aware is not None:\n existing_persona.datetime_aware = datetime_aware\n existing_persona.replace_base_system_prompt = replace_base_system_prompt\n\n # Do not delete any associations manually added unless\n # a new updated list is provided\n if document_sets is not None:\n existing_persona.document_sets.clear()\n existing_persona.document_sets = document_sets or []\n\n # Note: prompts are now embedded in personas - no separate prompts relationship\n\n if tools is not None:\n existing_persona.tools = tools or []\n\n if user_file_ids is not None:\n old_file_ids = {uf.id for uf in existing_persona.user_files}\n new_file_ids = {uf.id for uf in (user_files or [])}\n affected_file_ids = old_file_ids | new_file_ids\n existing_persona.user_files.clear()\n existing_persona.user_files = user_files or []\n if affected_file_ids:\n _mark_files_need_persona_sync(db_session, list(affected_file_ids))\n\n if hierarchy_node_ids is not None:\n existing_persona.hierarchy_nodes.clear()\n existing_persona.hierarchy_nodes = hierarchy_nodes or []\n\n if document_ids is not None:\n existing_persona.attached_documents.clear()\n existing_persona.attached_documents = attached_documents or []\n\n # We should only update display priority if it is not already set\n if existing_persona.display_priority is None:\n existing_persona.display_priority = display_priority\n\n persona = existing_persona\n\n else:\n # Create new persona - prompt configuration will be set separately if needed\n new_persona = Persona(\n id=persona_id,\n user_id=user.id if user else None,\n is_public=is_public,\n name=name,\n description=description,\n builtin_persona=builtin_persona,\n system_prompt=system_prompt or \"\",\n task_prompt=task_prompt or \"\",\n datetime_aware=(datetime_aware if datetime_aware is not None else True),\n replace_base_system_prompt=replace_base_system_prompt,\n document_sets=document_sets or [],\n llm_model_provider_override=llm_model_provider_override,\n llm_model_version_override=llm_model_version_override,\n starter_messages=starter_messages,\n tools=tools or [],\n uploaded_image_id=uploaded_image_id,\n icon_name=icon_name,\n display_priority=display_priority,\n is_listed=is_listed,\n search_start_date=search_start_date,\n is_featured=(is_featured if is_featured is not None else False),\n user_files=user_files or [],\n labels=labels or [],\n hierarchy_nodes=hierarchy_nodes or [],\n attached_documents=attached_documents or [],\n )\n db_session.add(new_persona)\n if user_files:\n _mark_files_need_persona_sync(db_session, [uf.id for uf in user_files])\n persona = new_persona\n if commit:\n db_session.commit()\n else:\n # flush the session so that the persona has an ID\n db_session.flush()\n\n return persona\n\n\ndef delete_old_default_personas(\n db_session: Session,\n) -> None:\n \"\"\"Note, this locks out the Summarize and Paraphrase personas for now\n Need a more graceful fix later or those need to never have IDs.\n\n This function is idempotent, so it can be run multiple times without issue.\n \"\"\"\n OLD_SUFFIX = \"_old\"\n stmt = (\n update(Persona)\n .where(\n Persona.builtin_persona,\n Persona.id > 0,\n or_(\n Persona.deleted.is_(False),\n not_(Persona.name.endswith(OLD_SUFFIX)),\n ),\n )\n .values(deleted=True, name=func.concat(Persona.name, OLD_SUFFIX))\n )\n\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef update_persona_featured(\n persona_id: int,\n is_featured: bool,\n db_session: Session,\n user: User,\n) -> None:\n persona = fetch_persona_by_id_for_user(\n db_session=db_session, persona_id=persona_id, user=user, get_editable=True\n )\n\n persona.is_featured = is_featured\n db_session.commit()\n\n\ndef update_persona_visibility(\n persona_id: int,\n is_listed: bool,\n db_session: Session,\n user: User,\n) -> None:\n persona = fetch_persona_by_id_for_user(\n db_session=db_session, persona_id=persona_id, user=user, get_editable=True\n )\n\n persona.is_listed = is_listed\n db_session.commit()\n\n\ndef validate_persona_tools(tools: list[Tool], db_session: Session) -> None:\n # local import to avoid circular import. DB layer should not depend on tools layer.\n from onyx.tools.built_in_tools import get_built_in_tool_by_id\n\n for tool in tools:\n if tool.in_code_tool_id is not None:\n tool_cls = get_built_in_tool_by_id(tool.in_code_tool_id)\n if not tool_cls.is_available(db_session):\n raise ValueError(f\"Tool {tool.in_code_tool_id} is not available\")\n\n\n# TODO: since this gets called with every chat message, could it be more efficient to pregenerate\n# a direct mapping indicating whether a user has access to a specific persona?\ndef get_persona_by_id(\n persona_id: int,\n user: User | None,\n db_session: Session,\n include_deleted: bool = False,\n is_for_edit: bool = True, # NOTE: assume true for safety\n) -> Persona:\n persona_stmt = (\n select(Persona)\n .distinct()\n .outerjoin(Persona.groups)\n .outerjoin(Persona.users)\n .outerjoin(UserGroup.user_group_relationships)\n .where(Persona.id == persona_id)\n )\n\n if not include_deleted:\n persona_stmt = persona_stmt.where(Persona.deleted.is_(False))\n\n if not user or user.role == UserRole.ADMIN:\n result = db_session.execute(persona_stmt)\n persona = result.scalar_one_or_none()\n if persona is None:\n raise ValueError(f\"Persona with ID {persona_id} does not exist\")\n return persona\n\n # or check if user owns persona\n or_conditions = Persona.user_id == user.id\n # allow access if persona user id is None\n or_conditions |= Persona.user_id == None # noqa: E711\n if not is_for_edit:\n # if the user is in a group related to the persona\n or_conditions |= User__UserGroup.user_id == user.id\n # if the user is in the .users of the persona\n or_conditions |= User.id == user.id\n or_conditions |= Persona.is_public == True # noqa: E712\n elif user.role == UserRole.GLOBAL_CURATOR:\n # global curators can edit personas for the groups they are in\n or_conditions |= User__UserGroup.user_id == user.id\n elif user.role == UserRole.CURATOR:\n # curators can edit personas for the groups they are curators of\n or_conditions |= (User__UserGroup.user_id == user.id) & (\n User__UserGroup.is_curator == True # noqa: E712\n )\n\n persona_stmt = persona_stmt.where(or_conditions)\n result = db_session.execute(persona_stmt)\n persona = result.scalar_one_or_none()\n if persona is None:\n raise ValueError(\n f\"Persona with ID {persona_id} does not exist or does not belong to user\"\n )\n return persona\n\n\ndef get_personas_by_ids(\n persona_ids: list[int], db_session: Session\n) -> Sequence[Persona]:\n \"\"\"WARNING: Unsafe, can fetch personas from all users.\"\"\"\n if not persona_ids:\n return []\n personas = db_session.scalars(\n select(Persona).where(Persona.id.in_(persona_ids))\n ).all()\n\n return personas\n\n\ndef delete_persona_by_name(\n persona_name: str, db_session: Session, is_default: bool = True\n) -> None:\n stmt = (\n update(Persona)\n .where(Persona.name == persona_name, Persona.builtin_persona == is_default)\n .values(deleted=True)\n )\n\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef get_assistant_labels(db_session: Session) -> list[PersonaLabel]:\n return db_session.query(PersonaLabel).all()\n\n\ndef create_assistant_label(db_session: Session, name: str) -> PersonaLabel:\n label = PersonaLabel(name=name)\n db_session.add(label)\n db_session.commit()\n return label\n\n\ndef update_persona_label(\n label_id: int,\n label_name: str,\n db_session: Session,\n) -> None:\n persona_label = (\n db_session.query(PersonaLabel).filter(PersonaLabel.id == label_id).one_or_none()\n )\n if persona_label is None:\n raise ValueError(f\"Persona label with ID {label_id} does not exist\")\n persona_label.name = label_name\n db_session.commit()\n\n\ndef delete_persona_label(label_id: int, db_session: Session) -> None:\n db_session.query(PersonaLabel).filter(PersonaLabel.id == label_id).delete()\n db_session.commit()\n\n\ndef persona_has_search_tool(persona_id: int, db_session: Session) -> bool:\n persona = (\n db_session.query(Persona)\n .options(selectinload(Persona.tools))\n .filter(Persona.id == persona_id)\n .one_or_none()\n )\n if persona is None:\n raise ValueError(f\"Persona with ID {persona_id} does not exist\")\n return any(tool.in_code_tool_id == \"run_search\" for tool in persona.tools)\n\n\ndef get_default_assistant(db_session: Session) -> Persona | None:\n \"\"\"Fetch the default assistant (persona with builtin_persona=True).\"\"\"\n return (\n db_session.query(Persona)\n .options(selectinload(Persona.tools))\n .filter(Persona.builtin_persona.is_(True))\n # NOTE: need to add this since we had prior builtin personas\n # that have since been deleted\n .filter(Persona.deleted.is_(False))\n .one_or_none()\n )\n\n\ndef update_default_assistant_configuration(\n db_session: Session,\n tool_ids: list[int] | None = None,\n system_prompt: str | None = None,\n update_system_prompt: bool = False,\n) -> Persona:\n \"\"\"Update only tools and system_prompt for the default assistant.\n\n Args:\n db_session: Database session\n tool_ids: List of tool IDs to enable (if None, tools are not updated)\n system_prompt: New system prompt value (None means use default)\n update_system_prompt: If True, update the system_prompt field (allows setting to None)\n\n Returns:\n Updated Persona object\n\n Raises:\n ValueError: If default assistant not found or invalid tool IDs provided\n \"\"\"\n # Get the default assistant\n persona = get_default_assistant(db_session)\n if not persona:\n raise ValueError(\"Default assistant not found\")\n\n # Update system prompt if explicitly requested\n if update_system_prompt:\n persona.system_prompt = system_prompt\n\n # Update tools if provided\n if tool_ids is not None:\n # Clear existing tool associations\n persona.tools = []\n\n # Add new tool associations\n for tool_id in tool_ids:\n tool = db_session.query(Tool).filter(Tool.id == tool_id).one_or_none()\n if not tool:\n raise ValueError(f\"Tool with ID {tool_id} not found\")\n\n if not should_expose_tool_to_fe(tool):\n raise ValueError(f\"Tool with ID {tool_id} cannot be assigned\")\n\n if not tool.enabled:\n raise ValueError(\n f\"Enable tool {tool.display_name or tool.name} before assigning it\"\n )\n\n persona.tools.append(tool)\n\n db_session.commit()\n return persona\n\n\ndef user_can_access_persona(\n db_session: Session, persona_id: int, user: User, get_editable: bool = False\n) -> bool:\n \"\"\"Check if a user has access to a specific persona.\n\n Args:\n db_session: Database session\n persona_id: ID of the persona to check\n user: User to check access for\n get_editable: If True, check for edit access; if False, check for view access\n\n Returns:\n True if user can access the persona, False otherwise\n \"\"\"\n stmt = select(Persona).where(Persona.id == persona_id, Persona.deleted.is_(False))\n stmt = _add_user_filters(stmt, user, get_editable=get_editable)\n return db_session.scalar(stmt) is not None\n" }, { "path": "backend/onyx/db/projects.py", "content": "import datetime\nimport uuid\nfrom typing import List\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom fastapi import UploadFile\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import Field\nfrom sqlalchemy import func\nfrom sqlalchemy.orm import Session\nfrom starlette.background import BackgroundTasks\n\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import Project__UserFile\nfrom onyx.db.models import User\nfrom onyx.db.models import UserFile\nfrom onyx.db.models import UserProject\nfrom onyx.server.documents.connector import upload_files\nfrom onyx.server.features.projects.projects_file_utils import categorize_uploaded_files\nfrom onyx.server.features.projects.projects_file_utils import RejectedFile\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\nclass CategorizedFilesResult(BaseModel):\n user_files: list[UserFile]\n rejected_files: list[RejectedFile]\n id_to_temp_id: dict[str, str]\n # Filenames that should be stored but not indexed.\n skip_indexing_filenames: set[str] = Field(default_factory=set)\n # Allow SQLAlchemy ORM models inside this result container\n model_config = ConfigDict(arbitrary_types_allowed=True)\n\n @property\n def indexable_files(self) -> list[UserFile]:\n return [\n uf\n for uf in self.user_files\n if (uf.name or \"\") not in self.skip_indexing_filenames\n ]\n\n\ndef build_hashed_file_key(file: UploadFile) -> str:\n name_prefix = (file.filename or \"\")[:50]\n return f\"{file.size}|{name_prefix}\"\n\n\ndef create_user_files(\n files: List[UploadFile],\n project_id: int | None,\n user: User,\n db_session: Session,\n link_url: str | None = None,\n temp_id_map: dict[str, str] | None = None,\n) -> CategorizedFilesResult:\n\n # Categorize the files\n categorized_files = categorize_uploaded_files(files, db_session)\n # NOTE: At the moment, zip metadata is not used for user files.\n # Should revisit to decide whether this should be a feature.\n upload_response = upload_files(categorized_files.acceptable, FileOrigin.USER_FILE)\n user_files = []\n rejected_files = categorized_files.rejected\n id_to_temp_id: dict[str, str] = {}\n # Pair returned storage paths with the same set of acceptable files we uploaded\n for file_path, file in zip(\n upload_response.file_paths, categorized_files.acceptable\n ):\n new_id = uuid.uuid4()\n new_temp_id = (\n temp_id_map.get(build_hashed_file_key(file)) if temp_id_map else None\n )\n if new_temp_id is not None:\n id_to_temp_id[str(new_id)] = new_temp_id\n should_skip = (file.filename or \"\") in categorized_files.skip_indexing\n new_file = UserFile(\n id=new_id,\n user_id=user.id,\n file_id=file_path,\n name=file.filename,\n token_count=categorized_files.acceptable_file_to_token_count[\n file.filename or \"\"\n ],\n link_url=link_url,\n content_type=file.content_type,\n file_type=file.content_type,\n status=UserFileStatus.SKIPPED if should_skip else UserFileStatus.PROCESSING,\n last_accessed_at=datetime.datetime.now(datetime.timezone.utc),\n )\n # Persist the UserFile first to satisfy FK constraints for association table\n db_session.add(new_file)\n db_session.flush()\n if project_id:\n project_to_user_file = Project__UserFile(\n project_id=project_id,\n user_file_id=new_file.id,\n )\n db_session.add(project_to_user_file)\n user_files.append(new_file)\n db_session.commit()\n return CategorizedFilesResult(\n user_files=user_files,\n rejected_files=rejected_files,\n id_to_temp_id=id_to_temp_id,\n skip_indexing_filenames=categorized_files.skip_indexing,\n )\n\n\ndef upload_files_to_user_files_with_indexing(\n files: List[UploadFile],\n project_id: int | None,\n user: User,\n temp_id_map: dict[str, str] | None,\n db_session: Session,\n background_tasks: BackgroundTasks | None = None,\n) -> CategorizedFilesResult:\n if project_id is not None and user is not None:\n if not check_project_ownership(project_id, user.id, db_session):\n raise HTTPException(status_code=404, detail=\"Project not found\")\n\n categorized_files_result = create_user_files(\n files,\n project_id,\n user,\n db_session,\n temp_id_map=temp_id_map,\n )\n user_files = categorized_files_result.user_files\n rejected_files = categorized_files_result.rejected_files\n id_to_temp_id = categorized_files_result.id_to_temp_id\n indexable_files = categorized_files_result.indexable_files\n # Trigger per-file processing immediately for the current tenant\n tenant_id = get_current_tenant_id()\n for rejected_file in rejected_files:\n logger.warning(\n f\"File {rejected_file.filename} rejected for {rejected_file.reason}\"\n )\n\n if DISABLE_VECTOR_DB and background_tasks is not None:\n from onyx.background.task_utils import drain_processing_loop\n\n background_tasks.add_task(drain_processing_loop, tenant_id)\n for user_file in indexable_files:\n logger.info(f\"Queued in-process processing for user_file_id={user_file.id}\")\n else:\n from onyx.background.celery.versioned_apps.client import app as client_app\n\n for user_file in indexable_files:\n task = client_app.send_task(\n OnyxCeleryTask.PROCESS_SINGLE_USER_FILE,\n kwargs={\"user_file_id\": user_file.id, \"tenant_id\": tenant_id},\n queue=OnyxCeleryQueues.USER_FILE_PROCESSING,\n priority=OnyxCeleryPriority.HIGH,\n expires=CELERY_USER_FILE_PROCESSING_TASK_EXPIRES,\n )\n logger.info(\n f\"Triggered indexing for user_file_id={user_file.id} with task_id={task.id}\"\n )\n\n return CategorizedFilesResult(\n user_files=user_files,\n rejected_files=rejected_files,\n id_to_temp_id=id_to_temp_id,\n skip_indexing_filenames=categorized_files_result.skip_indexing_filenames,\n )\n\n\ndef check_project_ownership(\n project_id: int, user_id: UUID | None, db_session: Session\n) -> bool:\n # In no-auth mode, all projects are accessible\n if user_id is None:\n # Verify project exists\n return (\n db_session.query(UserProject).filter(UserProject.id == project_id).first()\n is not None\n )\n\n return (\n db_session.query(UserProject)\n .filter(UserProject.id == project_id, UserProject.user_id == user_id)\n .first()\n is not None\n )\n\n\ndef get_user_files_from_project(\n project_id: int, user_id: UUID | None, db_session: Session\n) -> list[UserFile]:\n # First check if the user owns the project\n if not check_project_ownership(project_id, user_id, db_session):\n return []\n\n return (\n db_session.query(UserFile)\n .join(Project__UserFile)\n .filter(Project__UserFile.project_id == project_id)\n .all()\n )\n\n\ndef get_project_instructions(db_session: Session, project_id: int | None) -> str | None:\n \"\"\"Return the project's instruction text from the project, else None.\n\n Safe helper that swallows DB errors and returns None on any failure.\n \"\"\"\n if not project_id:\n return None\n try:\n project = (\n db_session.query(UserProject)\n .filter(UserProject.id == project_id)\n .one_or_none()\n )\n if not project or not project.instructions:\n return None\n instructions = project.instructions.strip()\n return instructions or None\n except Exception:\n return None\n\n\ndef get_project_token_count(\n project_id: int | None,\n user_id: UUID | None,\n db_session: Session,\n) -> int:\n \"\"\"Return sum of token_count for all user files in the given project.\n\n If project_id is None, returns 0.\n \"\"\"\n if project_id is None:\n return 0\n\n total_tokens = (\n db_session.query(func.coalesce(func.sum(UserFile.token_count), 0))\n .filter(\n UserFile.user_id == user_id,\n UserFile.projects.any(id=project_id),\n )\n .scalar()\n or 0\n )\n\n return int(total_tokens)\n" }, { "path": "backend/onyx/db/pydantic_type.py", "content": "import json\nfrom typing import Any\nfrom typing import Optional\nfrom typing import Type\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.dialects.postgresql import JSONB\nfrom sqlalchemy.types import TypeDecorator\n\n\nclass PydanticType(TypeDecorator):\n impl = JSONB\n\n def __init__(\n self, pydantic_model: Type[BaseModel], *args: Any, **kwargs: Any\n ) -> None:\n super().__init__(*args, **kwargs)\n self.pydantic_model = pydantic_model\n\n def process_bind_param(\n self,\n value: Optional[BaseModel],\n dialect: Any, # noqa: ARG002\n ) -> Optional[dict]:\n if value is not None:\n return json.loads(value.json())\n return None\n\n def process_result_value(\n self,\n value: Optional[dict],\n dialect: Any, # noqa: ARG002\n ) -> Optional[BaseModel]:\n if value is not None:\n return self.pydantic_model.parse_obj(value)\n return None\n\n\nclass PydanticListType(TypeDecorator):\n impl = JSONB\n\n def __init__(\n self, pydantic_model: Type[BaseModel], *args: Any, **kwargs: Any\n ) -> None:\n super().__init__(*args, **kwargs)\n self.pydantic_model = pydantic_model\n\n def process_bind_param(\n self,\n value: Optional[list[BaseModel]],\n dialect: Any, # noqa: ARG002\n ) -> Optional[list[dict]]:\n if value is not None:\n return [json.loads(item.model_dump_json()) for item in value]\n return None\n\n def process_result_value(\n self,\n value: Optional[list[dict]],\n dialect: Any, # noqa: ARG002\n ) -> Optional[list[BaseModel]]:\n if value is not None:\n return [self.pydantic_model.model_validate(item) for item in value]\n return None\n" }, { "path": "backend/onyx/db/relationships.py", "content": "from typing import List\n\nfrom sqlalchemy import or_\nfrom sqlalchemy.dialects import postgresql\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.orm import Session\n\nimport onyx.db.document as dbdocument\nfrom onyx.db.models import KGEntity\nfrom onyx.db.models import KGEntityExtractionStaging\nfrom onyx.db.models import KGRelationship\nfrom onyx.db.models import KGRelationshipExtractionStaging\nfrom onyx.db.models import KGRelationshipType\nfrom onyx.db.models import KGRelationshipTypeExtractionStaging\nfrom onyx.db.models import KGStage\nfrom onyx.kg.utils.formatting_utils import extract_relationship_type_id\nfrom onyx.kg.utils.formatting_utils import format_relationship_id\nfrom onyx.kg.utils.formatting_utils import get_entity_type\nfrom onyx.kg.utils.formatting_utils import make_relationship_id\nfrom onyx.kg.utils.formatting_utils import make_relationship_type_id\nfrom onyx.kg.utils.formatting_utils import split_relationship_id\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef upsert_staging_relationship(\n db_session: Session,\n relationship_id_name: str,\n source_document_id: str | None,\n occurrences: int = 1,\n) -> KGRelationshipExtractionStaging:\n \"\"\"\n Add or update a new staging relationship to the database.\n\n Args:\n db_session: SQLAlchemy database session\n relationship_id_name: The ID name of the relationship in format \"source__relationship__target\"\n source_document_id: ID of the source document\n occurrences: Number of times this relationship has been found\n Returns:\n The created or updated KGRelationshipExtractionStaging object\n\n Raises:\n sqlalchemy.exc.IntegrityError: If there's an error with the database operation\n \"\"\"\n # Generate a unique ID for the relationship\n relationship_id_name = format_relationship_id(relationship_id_name)\n (\n source_entity_id_name,\n relationship_string,\n target_entity_id_name,\n ) = split_relationship_id(relationship_id_name)\n\n source_entity_type = get_entity_type(source_entity_id_name)\n target_entity_type = get_entity_type(target_entity_id_name)\n relationship_type = extract_relationship_type_id(relationship_id_name)\n\n # Insert the new relationship\n stmt = (\n postgresql.insert(KGRelationshipExtractionStaging)\n .values(\n {\n \"id_name\": relationship_id_name,\n \"source_node\": source_entity_id_name,\n \"target_node\": target_entity_id_name,\n \"source_node_type\": source_entity_type,\n \"target_node_type\": target_entity_type,\n \"type\": relationship_string.lower(),\n \"relationship_type_id_name\": relationship_type,\n \"source_document\": source_document_id,\n \"occurrences\": occurrences,\n }\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\", \"source_document\"],\n set_=dict(\n occurrences=KGRelationshipExtractionStaging.occurrences + occurrences,\n ),\n )\n .returning(KGRelationshipExtractionStaging)\n )\n\n result = db_session.execute(stmt).scalar()\n if result is None:\n raise RuntimeError(\n f\"Failed to create or increment staging relationship with id_name: {relationship_id_name}\"\n )\n\n # Update the document's kg_stage if source_document is provided\n if source_document_id is not None:\n dbdocument.update_document_kg_info(\n db_session,\n document_id=source_document_id,\n kg_stage=KGStage.EXTRACTED,\n )\n db_session.flush() # Flush to get any DB errors early\n\n return result\n\n\ndef upsert_relationship(\n db_session: Session,\n relationship_id_name: str,\n source_document_id: str | None,\n occurrences: int = 1,\n) -> KGRelationship:\n \"\"\"\n Upsert a new relationship directly to the database.\n\n Args:\n db_session: SQLAlchemy database session\n relationship_id_name: The ID name of the relationship in format \"source__relationship__target\"\n source_document_id: ID of the source document\n occurrences: Number of times this relationship has been found\n Returns:\n The created or updated KGRelationship object\n\n Raises:\n sqlalchemy.exc.IntegrityError: If there's an error with the database operation\n \"\"\"\n # Generate a unique ID for the relationship\n relationship_id_name = format_relationship_id(relationship_id_name)\n (\n source_entity_id_name,\n relationship_string,\n target_entity_id_name,\n ) = split_relationship_id(relationship_id_name)\n\n source_entity_type = get_entity_type(source_entity_id_name)\n target_entity_type = get_entity_type(target_entity_id_name)\n relationship_type = extract_relationship_type_id(relationship_id_name)\n\n # Insert the new relationship\n stmt = (\n postgresql.insert(KGRelationship)\n .values(\n {\n \"id_name\": relationship_id_name,\n \"source_node\": source_entity_id_name,\n \"target_node\": target_entity_id_name,\n \"source_node_type\": source_entity_type,\n \"target_node_type\": target_entity_type,\n \"type\": relationship_string.lower(),\n \"relationship_type_id_name\": relationship_type,\n \"source_document\": source_document_id,\n \"occurrences\": occurrences,\n }\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\", \"source_document\"],\n set_=dict(\n occurrences=KGRelationship.occurrences + occurrences,\n ),\n )\n .returning(KGRelationship)\n )\n\n new_relationship = db_session.execute(stmt).scalar()\n if new_relationship is None:\n raise RuntimeError(\n f\"Failed to upsert relationship with id_name: {relationship_id_name}\"\n )\n db_session.flush()\n return new_relationship\n\n\ndef transfer_relationship(\n db_session: Session,\n relationship: KGRelationshipExtractionStaging,\n entity_translations: dict[str, str],\n) -> KGRelationship:\n \"\"\"\n Transfer a relationship from the staging table to the normalized table.\n \"\"\"\n # Translate the source and target nodes\n source_node = entity_translations[relationship.source_node]\n target_node = entity_translations[relationship.target_node]\n relationship_id_name = make_relationship_id(\n source_node, relationship.type, target_node\n )\n\n # Create the transferred relationship\n stmt = (\n pg_insert(KGRelationship)\n .values(\n id_name=relationship_id_name,\n source_node=source_node,\n target_node=target_node,\n source_node_type=relationship.source_node_type,\n target_node_type=relationship.target_node_type,\n type=relationship.type,\n relationship_type_id_name=relationship.relationship_type_id_name,\n source_document=relationship.source_document,\n occurrences=relationship.occurrences,\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\", \"source_document\"],\n set_=dict(\n occurrences=KGRelationship.occurrences + relationship.occurrences,\n ),\n )\n .returning(KGRelationship)\n )\n\n new_relationship = db_session.execute(stmt).scalar()\n if new_relationship is None:\n raise RuntimeError(\n f\"Failed to transfer relationship with id_name: {relationship.id_name}\"\n )\n\n # Update transferred\n db_session.query(KGRelationshipExtractionStaging).filter(\n KGRelationshipExtractionStaging.id_name == relationship.id_name,\n KGRelationshipExtractionStaging.source_document == relationship.source_document,\n ).update({\"transferred\": True})\n db_session.flush()\n\n return new_relationship\n\n\ndef upsert_staging_relationship_type(\n db_session: Session,\n source_entity_type: str,\n relationship_type: str,\n target_entity_type: str,\n definition: bool = False,\n extraction_count: int = 1,\n) -> KGRelationshipTypeExtractionStaging:\n \"\"\"\n Add a new relationship type to the database.\n\n Args:\n db_session: SQLAlchemy session\n source_entity_type: Type of the source entity\n relationship_type: Type of relationship\n target_entity_type: Type of the target entity\n definition: Whether this relationship type represents a definition (default False)\n\n Returns:\n The created KGRelationshipTypeExtractionStaging object\n \"\"\"\n\n id_name = make_relationship_type_id(\n source_entity_type, relationship_type, target_entity_type\n )\n\n # Create new relationship type\n stmt = (\n postgresql.insert(KGRelationshipTypeExtractionStaging)\n .values(\n {\n \"id_name\": id_name,\n \"name\": relationship_type,\n \"source_entity_type_id_name\": source_entity_type.upper(),\n \"target_entity_type_id_name\": target_entity_type.upper(),\n \"definition\": definition,\n \"occurrences\": extraction_count,\n \"type\": relationship_type, # Using the relationship_type as the type\n \"active\": True, # Setting as active by default\n }\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\"],\n set_=dict(\n occurrences=KGRelationshipTypeExtractionStaging.occurrences\n + extraction_count,\n ),\n )\n .returning(KGRelationshipTypeExtractionStaging)\n )\n\n result = db_session.execute(stmt).scalar()\n if result is None:\n raise RuntimeError(\n f\"Failed to create or increment staging relationship type with id_name: {id_name}\"\n )\n db_session.flush() # Flush to get any DB errors early\n\n return result\n\n\ndef upsert_relationship_type(\n db_session: Session,\n source_entity_type: str,\n relationship_type: str,\n target_entity_type: str,\n definition: bool = False,\n extraction_count: int = 1,\n) -> KGRelationshipType:\n \"\"\"\n Upsert a new relationship type directly to the database.\n\n Args:\n db_session: SQLAlchemy session\n source_entity_type: Type of the source entity\n relationship_type: Type of relationship\n target_entity_type: Type of the target entity\n definition: Whether this relationship type represents a definition (default False)\n\n Returns:\n The created KGRelationshipType object\n \"\"\"\n\n id_name = make_relationship_type_id(\n source_entity_type, relationship_type, target_entity_type\n )\n\n # Create new relationship type\n stmt = (\n postgresql.insert(KGRelationshipType)\n .values(\n {\n \"id_name\": id_name,\n \"name\": relationship_type,\n \"source_entity_type_id_name\": source_entity_type.upper(),\n \"target_entity_type_id_name\": target_entity_type.upper(),\n \"definition\": definition,\n \"occurrences\": extraction_count,\n \"type\": relationship_type, # Using the relationship_type as the type\n \"active\": True, # Setting as active by default\n }\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\"],\n set_=dict(\n occurrences=KGRelationshipType.occurrences + extraction_count,\n ),\n )\n .returning(KGRelationshipType)\n )\n\n new_relationship_type = db_session.execute(stmt).scalar()\n if new_relationship_type is None:\n raise RuntimeError(\n f\"Failed to upsert relationship type with id_name: {id_name}\"\n )\n db_session.flush()\n return new_relationship_type\n\n\ndef transfer_relationship_type(\n db_session: Session,\n relationship_type: KGRelationshipTypeExtractionStaging,\n) -> KGRelationshipType:\n \"\"\"\n Transfer a relationship type from the staging table to the normalized table.\n \"\"\"\n stmt = (\n pg_insert(KGRelationshipType)\n .values(\n id_name=relationship_type.id_name,\n name=relationship_type.name,\n source_entity_type_id_name=relationship_type.source_entity_type_id_name,\n target_entity_type_id_name=relationship_type.target_entity_type_id_name,\n definition=relationship_type.definition,\n occurrences=relationship_type.occurrences,\n type=relationship_type.type,\n active=relationship_type.active,\n )\n .on_conflict_do_update(\n index_elements=[\"id_name\"],\n set_=dict(\n occurrences=KGRelationshipType.occurrences\n + relationship_type.occurrences,\n ),\n )\n .returning(KGRelationshipType)\n )\n\n new_relationship_type = db_session.execute(stmt).scalar()\n if new_relationship_type is None:\n raise RuntimeError(\n f\"Failed to transfer relationship type with id_name: {relationship_type.id_name}\"\n )\n\n # Update transferred\n db_session.query(KGRelationshipTypeExtractionStaging).filter(\n KGRelationshipTypeExtractionStaging.id_name == relationship_type.id_name\n ).update({\"transferred\": True})\n db_session.flush()\n\n return new_relationship_type\n\n\ndef delete_relationships_by_id_names(\n db_session: Session, id_names: list[str], kg_stage: KGStage\n) -> int:\n \"\"\"\n Delete relationships from the database based on a list of id_names.\n\n Args:\n db_session: SQLAlchemy database session\n id_names: List of relationship id_names to delete\n\n Returns:\n Number of relationships deleted\n\n Raises:\n sqlalchemy.exc.SQLAlchemyError: If there's an error during deletion\n \"\"\"\n\n deleted_count = 0\n\n if kg_stage == KGStage.EXTRACTED:\n deleted_count = (\n db_session.query(KGRelationshipExtractionStaging)\n .filter(KGRelationshipExtractionStaging.id_name.in_(id_names))\n .delete(synchronize_session=False)\n )\n elif kg_stage == KGStage.NORMALIZED:\n deleted_count = (\n db_session.query(KGRelationship)\n .filter(KGRelationship.id_name.in_(id_names))\n .delete(synchronize_session=False)\n )\n\n db_session.flush() # Flush to ensure deletion is processed\n return deleted_count\n\n\ndef delete_relationship_types_by_id_names(\n db_session: Session, id_names: list[str], kg_stage: KGStage\n) -> int:\n \"\"\"\n Delete relationship types from the database based on a list of id_names.\n\n Args:\n db_session: SQLAlchemy database session\n id_names: List of relationship type id_names to delete\n\n Returns:\n Number of relationship types deleted\n\n Raises:\n sqlalchemy.exc.SQLAlchemyError: If there's an error during deletion\n \"\"\"\n deleted_count = 0\n\n if kg_stage == KGStage.EXTRACTED:\n deleted_count = (\n db_session.query(KGRelationshipTypeExtractionStaging)\n .filter(KGRelationshipTypeExtractionStaging.id_name.in_(id_names))\n .delete(synchronize_session=False)\n )\n elif kg_stage == KGStage.NORMALIZED:\n deleted_count = (\n db_session.query(KGRelationshipType)\n .filter(KGRelationshipType.id_name.in_(id_names))\n .delete(synchronize_session=False)\n )\n\n db_session.flush() # Flush to ensure deletion is processed\n return deleted_count\n\n\ndef get_relationships_for_entity_type_pairs(\n db_session: Session, entity_type_pairs: list[tuple[str, str]]\n) -> list[\"KGRelationshipType\"]:\n \"\"\"\n Get relationship types from the database based on a list of entity type pairs.\n\n Args:\n db_session: SQLAlchemy database session\n entity_type_pairs: List of tuples where each tuple contains (source_entity_type, target_entity_type)\n\n Returns:\n List of KGRelationshipType objects where source and target types match the provided pairs\n \"\"\"\n\n conditions = [\n (\n (KGRelationshipType.source_entity_type_id_name == source_type)\n & (KGRelationshipType.target_entity_type_id_name == target_type)\n )\n for source_type, target_type in entity_type_pairs\n ]\n\n return db_session.query(KGRelationshipType).filter(or_(*conditions)).all()\n\n\ndef get_allowed_relationship_type_pairs(\n db_session: Session, entities: list[str]\n) -> list[str]:\n \"\"\"\n Get the allowed relationship pairs for the given entities.\n\n Args:\n db_session: SQLAlchemy database session\n entities: List of entity type ID names to filter by\n\n Returns:\n List of id_names from KGRelationshipType where source or target entity types\n are in the provided entities list. We also filter out for now the catch-all\n relationship types 'VENDOR____'\n \"\"\"\n\n entity_types = list({get_entity_type(entity) for entity in entities})\n\n return [\n row[0]\n for row in (\n db_session.query(KGRelationshipType.id_name)\n .filter(\n or_(\n KGRelationshipType.source_entity_type_id_name.in_(entity_types),\n KGRelationshipType.target_entity_type_id_name.in_(entity_types),\n )\n )\n .filter(~KGRelationshipType.source_entity_type_id_name.like(\"VENDOR::%\"))\n .distinct()\n .all()\n )\n ]\n\n\ndef get_relationships_of_entity(db_session: Session, entity_id: str) -> List[str]:\n \"\"\"Get all relationship ID names where the given entity is either the source or target node.\n\n Args:\n db_session: SQLAlchemy session\n entity_id: ID of the entity to find relationships for\n\n Returns:\n List of relationship ID names where the entity is either source or target\n \"\"\"\n return [\n row[0]\n for row in (\n db_session.query(KGRelationship.id_name)\n .filter(\n or_(\n KGRelationship.source_node == entity_id,\n KGRelationship.target_node == entity_id,\n )\n )\n .all()\n )\n ]\n\n\ndef get_relationship_types_of_entity_types(\n db_session: Session, entity_types_id: str\n) -> List[str]:\n \"\"\"Get all relationship ID names where the given entity is either the source or target node.\n\n Args:\n db_session: SQLAlchemy session\n entity_types_id: ID of the entity to find relationships for\n\n Returns:\n List of relationship ID names where the entity is either source or target\n \"\"\"\n\n if entity_types_id.endswith(\":*\"):\n entity_types_id = entity_types_id[:-2]\n\n return [\n row[0]\n for row in (\n db_session.query(KGRelationshipType.id_name)\n .filter(\n or_(\n KGRelationshipType.source_entity_type_id_name == entity_types_id,\n KGRelationshipType.target_entity_type_id_name == entity_types_id,\n )\n )\n .all()\n )\n ]\n\n\ndef delete_document_references_from_kg(db_session: Session, document_id: str) -> None:\n # Delete relationships from normalized stage\n db_session.query(KGRelationship).filter(\n KGRelationship.source_document == document_id\n ).delete(synchronize_session=False)\n\n # Delete relationships from extraction staging\n db_session.query(KGRelationshipExtractionStaging).filter(\n KGRelationshipExtractionStaging.source_document == document_id\n ).delete(synchronize_session=False)\n\n # Delete entities from normalized stage\n db_session.query(KGEntity).filter(KGEntity.document_id == document_id).delete(\n synchronize_session=False\n )\n\n # Delete entities from extraction staging\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.document_id == document_id\n ).delete(synchronize_session=False)\n\n db_session.flush()\n\n\ndef delete_from_kg_relationships_extraction_staging__no_commit(\n db_session: Session, document_ids: list[str]\n) -> None:\n \"\"\"Delete relationships from the extraction staging table.\"\"\"\n db_session.query(KGRelationshipExtractionStaging).filter(\n KGRelationshipExtractionStaging.source_document.in_(document_ids)\n ).delete(synchronize_session=False)\n\n\ndef delete_from_kg_relationships__no_commit(\n db_session: Session, document_ids: list[str]\n) -> None:\n \"\"\"Delete relationships from the normalized table.\"\"\"\n db_session.query(KGRelationship).filter(\n KGRelationship.source_document.in_(document_ids)\n ).delete(synchronize_session=False)\n" }, { "path": "backend/onyx/db/release_notes.py", "content": "\"\"\"Database functions for release notes functionality.\"\"\"\n\nfrom urllib.parse import urlencode\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import INSTANCE_TYPE\nfrom onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN\nfrom onyx.configs.constants import NotificationType\nfrom onyx.configs.constants import ONYX_UTM_SOURCE\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import User\nfrom onyx.db.notification import batch_create_notifications\nfrom onyx.server.features.release_notes.constants import DOCS_CHANGELOG_BASE_URL\nfrom onyx.server.features.release_notes.models import ReleaseNoteEntry\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef create_release_notifications_for_versions(\n db_session: Session,\n release_note_entries: list[ReleaseNoteEntry],\n) -> int:\n \"\"\"\n Create release notes notifications for each release note entry.\n Uses batch_create_notifications for efficient bulk insertion.\n\n If a user already has a notification for a specific version (dismissed or not),\n no new one is created (handled by unique constraint on additional_data).\n\n Note: Entries should already be filtered by app_version before calling this\n function. The filtering happens in _parse_mdx_to_release_note_entries().\n\n Args:\n db_session: Database session\n release_note_entries: List of release note entries to notify about (pre-filtered)\n\n Returns:\n Total number of notifications created across all versions.\n \"\"\"\n if not release_note_entries:\n logger.debug(\"No release note entries to notify about\")\n return 0\n\n # Get active users and exclude API key users\n user_ids = list(\n db_session.scalars(\n select(User.id).where( # type: ignore\n User.is_active == True, # noqa: E712\n User.account_type.notin_([AccountType.BOT, AccountType.EXT_PERM_USER]),\n User.email.endswith(DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN).is_(False), # type: ignore[attr-defined]\n )\n ).all()\n )\n\n total_created = 0\n for entry in release_note_entries:\n # Convert version to anchor format for external docs links\n # v2.7.0 -> v2-7-0\n version_anchor = entry.version.replace(\".\", \"-\")\n\n # Build UTM parameters for tracking\n utm_params = {\n \"utm_source\": ONYX_UTM_SOURCE,\n \"utm_medium\": \"notification\",\n \"utm_campaign\": INSTANCE_TYPE,\n \"utm_content\": f\"release_notes-{entry.version}\",\n }\n\n link = f\"{DOCS_CHANGELOG_BASE_URL}#{version_anchor}?{urlencode(utm_params)}\"\n\n additional_data: dict[str, str] = {\n \"version\": entry.version,\n \"link\": link,\n }\n\n created_count = batch_create_notifications(\n user_ids,\n NotificationType.RELEASE_NOTES,\n db_session,\n title=entry.title,\n description=f\"Check out what's new in {entry.version}\",\n additional_data=additional_data,\n )\n total_created += created_count\n\n logger.debug(\n f\"Created {created_count} release notes notifications (version {entry.version}, {len(user_ids)} eligible users)\"\n )\n\n return total_created\n" }, { "path": "backend/onyx/db/rotate_encryption_key.py", "content": "\"\"\"Rotate encryption key for all encrypted columns.\n\nDynamically discovers all columns using EncryptedString / EncryptedJson,\ndecrypts each value with the old key, and re-encrypts with the current\nENCRYPTION_KEY_SECRET.\n\nThe operation is idempotent: rows already encrypted with the current key\nare skipped. Commits are made in batches so a crash mid-rotation can be\nsafely resumed by re-running.\n\"\"\"\n\nimport json\nfrom typing import Any\n\nfrom sqlalchemy import LargeBinary\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import ENCRYPTION_KEY_SECRET\nfrom onyx.db.models import Base\nfrom onyx.db.models import EncryptedJson\nfrom onyx.db.models import EncryptedString\nfrom onyx.utils.encryption import decrypt_bytes_to_string\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import global_version\n\nlogger = setup_logger()\n\n_BATCH_SIZE = 500\n\n\ndef _can_decrypt_with_current_key(data: bytes) -> bool:\n \"\"\"Check if data is already encrypted with the current key.\n\n Passes the key explicitly so the fallback-to-raw-decode path in\n _decrypt_bytes is NOT triggered — a clean success/failure signal.\n \"\"\"\n try:\n decrypt_bytes_to_string(data, key=ENCRYPTION_KEY_SECRET)\n return True\n except Exception:\n return False\n\n\ndef _discover_encrypted_columns() -> list[tuple[type, str, list[str], bool]]:\n \"\"\"Walk all ORM models and find columns using EncryptedString/EncryptedJson.\n\n Returns list of (ModelClass, column_attr_name, [pk_attr_names], is_json).\n \"\"\"\n results: list[tuple[type, str, list[str], bool]] = []\n\n for mapper in Base.registry.mappers:\n model_cls = mapper.class_\n pk_names = [col.key for col in mapper.primary_key]\n\n for prop in mapper.column_attrs:\n for col in prop.columns:\n if isinstance(col.type, EncryptedJson):\n results.append((model_cls, prop.key, pk_names, True))\n elif isinstance(col.type, EncryptedString):\n results.append((model_cls, prop.key, pk_names, False))\n\n return results\n\n\ndef rotate_encryption_key(\n db_session: Session,\n old_key: str | None,\n dry_run: bool = False,\n) -> dict[str, int]:\n \"\"\"Decrypt all encrypted columns with old_key and re-encrypt with the current key.\n\n Args:\n db_session: Active database session.\n old_key: The previous encryption key. Pass None or \"\" if values were\n not previously encrypted with a key.\n dry_run: If True, count rows that need rotation without modifying data.\n\n Returns:\n Dict of \"table.column\" -> number of rows re-encrypted (or would be).\n\n Commits every _BATCH_SIZE rows so that locks are held briefly and progress\n is preserved on crash. Already-rotated rows are detected and skipped,\n making the operation safe to re-run.\n \"\"\"\n if not global_version.is_ee_version():\n raise RuntimeError(\"EE mode is not enabled — rotation requires EE encryption.\")\n\n if not ENCRYPTION_KEY_SECRET:\n raise RuntimeError(\n \"ENCRYPTION_KEY_SECRET is not set — cannot rotate. Set the target encryption key in the environment before running.\"\n )\n\n encrypted_columns = _discover_encrypted_columns()\n totals: dict[str, int] = {}\n\n for model_cls, col_name, pk_names, is_json in encrypted_columns:\n table_name: str = model_cls.__tablename__ # type: ignore[attr-defined]\n col_attr = getattr(model_cls, col_name)\n pk_attrs = [getattr(model_cls, pk) for pk in pk_names]\n\n # Read raw bytes directly, bypassing the TypeDecorator\n raw_col = col_attr.property.columns[0]\n\n stmt = select(*pk_attrs, raw_col.cast(LargeBinary)).where(col_attr.is_not(None))\n rows = db_session.execute(stmt).all()\n\n reencrypted = 0\n batch_pending = 0\n for row in rows:\n raw_bytes: bytes | None = row[-1]\n if raw_bytes is None:\n continue\n\n if _can_decrypt_with_current_key(raw_bytes):\n continue\n\n try:\n if not old_key:\n decrypted_str = raw_bytes.decode(\"utf-8\")\n else:\n decrypted_str = decrypt_bytes_to_string(raw_bytes, key=old_key)\n\n # For EncryptedJson, parse back to dict so the TypeDecorator\n # can json.dumps() it cleanly (avoids double-encoding).\n value: Any = json.loads(decrypted_str) if is_json else decrypted_str\n except (ValueError, UnicodeDecodeError) as e:\n pk_vals = [row[i] for i in range(len(pk_names))]\n logger.warning(\n f\"Could not decrypt/parse {table_name}.{col_name} row {pk_vals} — skipping: {e}\"\n )\n continue\n\n if not dry_run:\n pk_filters = [pk_attr == row[i] for i, pk_attr in enumerate(pk_attrs)]\n update_stmt = (\n update(model_cls).where(*pk_filters).values({col_name: value})\n )\n db_session.execute(update_stmt)\n batch_pending += 1\n\n if batch_pending >= _BATCH_SIZE:\n db_session.commit()\n batch_pending = 0\n reencrypted += 1\n\n # Flush remaining rows in this column\n if batch_pending > 0:\n db_session.commit()\n\n if reencrypted > 0:\n totals[f\"{table_name}.{col_name}\"] = reencrypted\n logger.info(\n f\"{'[DRY RUN] Would re-encrypt' if dry_run else 'Re-encrypted'} {reencrypted} value(s) in {table_name}.{col_name}\"\n )\n\n return totals\n" }, { "path": "backend/onyx/db/saml.py", "content": "import datetime\nfrom typing import cast\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.ext.asyncio import AsyncSession\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS\nfrom onyx.db.models import SamlAccount\n\n\ndef upsert_saml_account(\n user_id: UUID,\n cookie: str,\n db_session: Session,\n expiration_offset: int = SESSION_EXPIRE_TIME_SECONDS,\n) -> datetime.datetime:\n expires_at = func.now() + datetime.timedelta(seconds=expiration_offset)\n\n existing_saml_acc = (\n db_session.query(SamlAccount)\n .filter(SamlAccount.user_id == user_id)\n .one_or_none()\n )\n\n if existing_saml_acc:\n existing_saml_acc.encrypted_cookie = cookie\n existing_saml_acc.expires_at = cast(datetime.datetime, expires_at)\n existing_saml_acc.updated_at = func.now()\n saml_acc = existing_saml_acc\n else:\n saml_acc = SamlAccount(\n user_id=user_id,\n encrypted_cookie=cookie,\n expires_at=expires_at,\n )\n db_session.add(saml_acc)\n\n db_session.commit()\n\n return saml_acc.expires_at\n\n\nasync def get_saml_account(\n cookie: str, async_db_session: AsyncSession\n) -> SamlAccount | None:\n \"\"\"NOTE: this is async, since it's used during auth\n (which is necessarily async due to FastAPI Users)\"\"\"\n stmt = (\n select(SamlAccount)\n .options(selectinload(SamlAccount.user)) # Use selectinload for collections\n .where(\n and_(\n SamlAccount.encrypted_cookie == cookie,\n SamlAccount.expires_at > func.now(),\n )\n )\n )\n\n result = await async_db_session.execute(stmt)\n return result.scalars().unique().one_or_none()\n\n\nasync def expire_saml_account(\n saml_account: SamlAccount, async_db_session: AsyncSession\n) -> None:\n saml_account.expires_at = func.now()\n await async_db_session.commit()\n" }, { "path": "backend/onyx/db/search_settings.py", "content": "from sqlalchemy import and_\nfrom sqlalchemy import delete\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.model_configs import DEFAULT_DOCUMENT_ENCODER_MODEL\nfrom onyx.configs.model_configs import DOCUMENT_ENCODER_MODEL\nfrom onyx.context.search.models import SavedSearchSettings\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.llm import fetch_embedding_provider\nfrom onyx.db.models import CloudEmbeddingProvider\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexModelStatus\nfrom onyx.db.models import SearchSettings\nfrom onyx.server.manage.embedding.models import (\n CloudEmbeddingProvider as ServerCloudEmbeddingProvider,\n)\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import PRESERVED_SEARCH_FIELDS\nfrom shared_configs.enums import EmbeddingProvider\n\n\nlogger = setup_logger()\n\n\nclass ActiveSearchSettings:\n primary: SearchSettings\n secondary: SearchSettings | None\n\n def __init__(\n self, primary: SearchSettings, secondary: SearchSettings | None\n ) -> None:\n self.primary = primary\n self.secondary = secondary\n\n\ndef create_search_settings(\n search_settings: SavedSearchSettings,\n db_session: Session,\n status: IndexModelStatus = IndexModelStatus.FUTURE,\n) -> SearchSettings:\n embedding_model = SearchSettings(\n model_name=search_settings.model_name,\n model_dim=search_settings.model_dim,\n normalize=search_settings.normalize,\n query_prefix=search_settings.query_prefix,\n passage_prefix=search_settings.passage_prefix,\n status=status,\n index_name=search_settings.index_name,\n provider_type=search_settings.provider_type,\n multipass_indexing=search_settings.multipass_indexing,\n embedding_precision=search_settings.embedding_precision,\n reduced_dimension=search_settings.reduced_dimension,\n enable_contextual_rag=search_settings.enable_contextual_rag,\n contextual_rag_llm_name=search_settings.contextual_rag_llm_name,\n contextual_rag_llm_provider=search_settings.contextual_rag_llm_provider,\n switchover_type=search_settings.switchover_type,\n )\n\n db_session.add(embedding_model)\n db_session.commit()\n\n return embedding_model\n\n\ndef get_embedding_provider_from_provider_type(\n db_session: Session, provider_type: EmbeddingProvider\n) -> CloudEmbeddingProvider | None:\n query = select(CloudEmbeddingProvider).where(\n CloudEmbeddingProvider.provider_type == provider_type\n )\n provider = db_session.execute(query).scalars().first()\n return provider if provider else None\n\n\ndef get_current_db_embedding_provider(\n db_session: Session,\n) -> ServerCloudEmbeddingProvider | None:\n search_settings = get_current_search_settings(db_session=db_session)\n\n if search_settings.provider_type is None:\n return None\n\n embedding_provider = fetch_embedding_provider(\n db_session=db_session,\n provider_type=search_settings.provider_type,\n )\n if embedding_provider is None:\n raise RuntimeError(\"No embedding provider exists for this model.\")\n\n current_embedding_provider = ServerCloudEmbeddingProvider.from_request(\n cloud_provider_model=embedding_provider\n )\n\n return current_embedding_provider\n\n\ndef delete_search_settings(db_session: Session, search_settings_id: int) -> None:\n current_settings = get_current_search_settings(db_session)\n\n if current_settings.id == search_settings_id:\n raise ValueError(\"Cannot delete currently active search settings\")\n\n # First, delete associated index attempts\n index_attempts_query = delete(IndexAttempt).where(\n IndexAttempt.search_settings_id == search_settings_id\n )\n db_session.execute(index_attempts_query)\n\n # Then, delete the search settings\n search_settings_query = delete(SearchSettings).where(\n and_(\n SearchSettings.id == search_settings_id,\n SearchSettings.status != IndexModelStatus.PRESENT,\n )\n )\n\n db_session.execute(search_settings_query)\n db_session.commit()\n\n\ndef get_current_search_settings(db_session: Session) -> SearchSettings:\n query = (\n select(SearchSettings)\n .where(SearchSettings.status == IndexModelStatus.PRESENT)\n .order_by(SearchSettings.id.desc())\n )\n result = db_session.execute(query)\n latest_settings = result.scalars().first()\n\n if not latest_settings:\n raise RuntimeError(\"No search settings specified; DB is not in a valid state.\")\n return latest_settings\n\n\ndef get_secondary_search_settings(db_session: Session) -> SearchSettings | None:\n query = (\n select(SearchSettings)\n .where(SearchSettings.status == IndexModelStatus.FUTURE)\n .order_by(SearchSettings.id.desc())\n )\n result = db_session.execute(query)\n latest_settings = result.scalars().first()\n\n return latest_settings\n\n\ndef get_active_search_settings(db_session: Session) -> ActiveSearchSettings:\n \"\"\"Returns active search settings. Secondary search settings may be None.\"\"\"\n\n # Get the primary and secondary search settings\n primary_search_settings = get_current_search_settings(db_session)\n secondary_search_settings = get_secondary_search_settings(db_session)\n return ActiveSearchSettings(\n primary=primary_search_settings, secondary=secondary_search_settings\n )\n\n\ndef get_active_search_settings_list(db_session: Session) -> list[SearchSettings]:\n \"\"\"Returns active search settings as a list. Primary settings are the first element,\n and if secondary search settings exist, they will be the second element.\"\"\"\n\n search_settings_list: list[SearchSettings] = []\n\n active_search_settings = get_active_search_settings(db_session)\n search_settings_list.append(active_search_settings.primary)\n if active_search_settings.secondary:\n search_settings_list.append(active_search_settings.secondary)\n\n return search_settings_list\n\n\ndef get_all_search_settings(db_session: Session) -> list[SearchSettings]:\n query = select(SearchSettings).order_by(SearchSettings.id.desc())\n result = db_session.execute(query)\n all_settings = result.scalars().all()\n return list(all_settings)\n\n\ndef get_multilingual_expansion(db_session: Session | None = None) -> list[str]:\n if db_session is None:\n with get_session_with_current_tenant() as db_session:\n search_settings = get_current_search_settings(db_session)\n else:\n search_settings = get_current_search_settings(db_session)\n if not search_settings:\n return []\n return search_settings.multilingual_expansion\n\n\ndef update_search_settings(\n current_settings: SearchSettings,\n updated_settings: SavedSearchSettings,\n preserved_fields: list[str],\n) -> None:\n for field, value in updated_settings.dict().items():\n if field not in preserved_fields:\n setattr(current_settings, field, value)\n\n\ndef update_current_search_settings(\n db_session: Session,\n search_settings: SavedSearchSettings,\n preserved_fields: list[str] = PRESERVED_SEARCH_FIELDS,\n) -> None:\n current_settings = get_current_search_settings(db_session)\n if not current_settings:\n logger.warning(\"No current search settings found to update\")\n return\n\n update_search_settings(current_settings, search_settings, preserved_fields)\n db_session.commit()\n logger.info(\"Current search settings updated successfully\")\n\n\ndef update_secondary_search_settings(\n db_session: Session,\n search_settings: SavedSearchSettings,\n preserved_fields: list[str] = PRESERVED_SEARCH_FIELDS,\n) -> None:\n secondary_settings = get_secondary_search_settings(db_session)\n if not secondary_settings:\n logger.warning(\"No secondary search settings found to update\")\n return\n\n preserved_fields = PRESERVED_SEARCH_FIELDS\n update_search_settings(secondary_settings, search_settings, preserved_fields)\n\n db_session.commit()\n logger.info(\"Secondary search settings updated successfully\")\n\n\ndef update_search_settings_status(\n search_settings: SearchSettings, new_status: IndexModelStatus, db_session: Session\n) -> None:\n search_settings.status = new_status\n db_session.commit()\n\n\ndef user_has_overridden_embedding_model() -> bool:\n return DOCUMENT_ENCODER_MODEL != DEFAULT_DOCUMENT_ENCODER_MODEL\n" }, { "path": "backend/onyx/db/seeding/chat_history_seeding.py", "content": "import random\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom logging import getLogger\nfrom uuid import UUID\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.chat import create_new_chat_message\nfrom onyx.db.chat import get_or_create_root_message\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import ChatSession\n\nlogger = getLogger(__name__)\n\n\ndef seed_chat_history(\n num_sessions: int,\n num_messages: int,\n days: int,\n user_id: UUID | None = None,\n persona_id: int | None = None,\n) -> None:\n \"\"\"Utility function to seed chat history for testing.\n\n num_sessions: the number of sessions to seed\n num_messages: the number of messages to seed per sessions\n days: the number of days looking backwards from the current time over which to randomize\n the times.\n user_id: optional user to associate with sessions\n persona_id: optional persona/assistant to associate with sessions\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n logger.info(f\"Seeding {num_sessions} sessions.\")\n for y in range(0, num_sessions):\n create_chat_session(db_session, f\"pytest_session_{y}\", user_id, persona_id)\n\n # randomize all session times\n logger.info(f\"Seeding {num_messages} messages per session.\")\n rows = db_session.query(ChatSession).all()\n for x in range(0, len(rows)):\n if x % 1024 == 0:\n logger.info(f\"Seeded messages for {x} sessions so far.\")\n\n row = rows[x]\n row.time_created = datetime.utcnow() - timedelta(\n days=random.randint(0, days)\n )\n row.time_updated = row.time_created + timedelta(\n minutes=random.randint(0, 10)\n )\n\n root_message = get_or_create_root_message(row.id, db_session)\n\n current_message_type = MessageType.USER\n parent_message = root_message\n for x in range(0, num_messages):\n if current_message_type == MessageType.USER:\n msg = f\"pytest_message_user_{x}\"\n else:\n msg = f\"pytest_message_assistant_{x}\"\n\n chat_message = create_new_chat_message(\n chat_session_id=row.id,\n parent_message=parent_message,\n message=msg,\n token_count=0,\n message_type=current_message_type,\n commit=False,\n db_session=db_session,\n )\n\n chat_message.time_sent = row.time_created + timedelta(\n minutes=random.randint(0, 10)\n )\n\n db_session.commit()\n\n current_message_type = (\n MessageType.ASSISTANT\n if current_message_type == MessageType.USER\n else MessageType.USER\n )\n parent_message = chat_message\n\n db_session.commit()\n\n logger.info(f\"Seeded messages for {len(rows)} sessions. Finished.\")\n" }, { "path": "backend/onyx/db/slack_bot.py", "content": "from collections.abc import Sequence\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import SlackBot\n\n\ndef insert_slack_bot(\n db_session: Session,\n name: str,\n enabled: bool,\n bot_token: str,\n app_token: str,\n user_token: str | None = None,\n) -> SlackBot:\n slack_bot = SlackBot(\n name=name,\n enabled=enabled,\n bot_token=bot_token,\n app_token=app_token,\n user_token=user_token,\n )\n db_session.add(slack_bot)\n db_session.commit()\n\n return slack_bot\n\n\ndef update_slack_bot(\n db_session: Session,\n slack_bot_id: int,\n name: str,\n enabled: bool,\n bot_token: str,\n app_token: str,\n user_token: str | None = None,\n) -> SlackBot:\n slack_bot = db_session.scalar(select(SlackBot).where(SlackBot.id == slack_bot_id))\n if slack_bot is None:\n raise ValueError(f\"Unable to find Slack Bot with ID {slack_bot_id}\")\n\n # update the app\n slack_bot.name = name\n slack_bot.enabled = enabled\n slack_bot.bot_token = bot_token # type: ignore[assignment]\n slack_bot.app_token = app_token # type: ignore[assignment]\n slack_bot.user_token = user_token # type: ignore[assignment]\n\n db_session.commit()\n\n return slack_bot\n\n\ndef fetch_slack_bot(\n db_session: Session,\n slack_bot_id: int,\n) -> SlackBot:\n slack_bot = db_session.scalar(select(SlackBot).where(SlackBot.id == slack_bot_id))\n if slack_bot is None:\n raise ValueError(f\"Unable to find Slack Bot with ID {slack_bot_id}\")\n\n return slack_bot\n\n\ndef remove_slack_bot(\n db_session: Session,\n slack_bot_id: int,\n) -> None:\n slack_bot = fetch_slack_bot(\n db_session=db_session,\n slack_bot_id=slack_bot_id,\n )\n\n db_session.delete(slack_bot)\n db_session.commit()\n\n\ndef fetch_slack_bots(db_session: Session) -> Sequence[SlackBot]:\n return db_session.scalars(select(SlackBot)).all()\n" }, { "path": "backend/onyx/db/slack_channel_config.py", "content": "from collections.abc import Sequence\nfrom typing import Any\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.constants import DEFAULT_PERSONA_SLACK_CHANNEL_NAME\nfrom onyx.db.constants import SLACK_BOT_PERSONA_PREFIX\nfrom onyx.db.models import ChannelConfig\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__DocumentSet\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.db.models import User\nfrom onyx.db.persona import mark_persona_as_deleted\nfrom onyx.db.persona import upsert_persona\nfrom onyx.db.tools import get_builtin_tool\nfrom onyx.tools.tool_implementations.search.search_tool import SearchTool\nfrom onyx.utils.errors import EERequiredError\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\n\n\ndef _build_persona_name(channel_name: str | None) -> str:\n return f\"{SLACK_BOT_PERSONA_PREFIX}{channel_name if channel_name else DEFAULT_PERSONA_SLACK_CHANNEL_NAME}\"\n\n\ndef _cleanup_relationships(db_session: Session, persona_id: int) -> None:\n \"\"\"NOTE: does not commit changes\"\"\"\n # delete existing persona-document_set relationships\n existing_relationships = db_session.scalars(\n select(Persona__DocumentSet).where(\n Persona__DocumentSet.persona_id == persona_id\n )\n )\n for rel in existing_relationships:\n db_session.delete(rel)\n\n\ndef create_slack_channel_persona(\n db_session: Session,\n channel_name: str | None,\n document_set_ids: list[int],\n existing_persona_id: int | None = None,\n) -> Persona:\n \"\"\"NOTE: does not commit changes\"\"\"\n\n search_tool = get_builtin_tool(db_session=db_session, tool_type=SearchTool)\n\n # create/update persona associated with the Slack channel\n persona_name = _build_persona_name(channel_name)\n persona_id_to_update = existing_persona_id\n if persona_id_to_update is None:\n # Reuse any previous Slack persona for this channel (even if the config was\n # temporarily switched to a different persona) so we don't trip duplicate name\n # validation inside `upsert_persona`.\n existing_persona = db_session.scalar(\n select(Persona).where(Persona.name == persona_name)\n )\n if existing_persona:\n persona_id_to_update = existing_persona.id\n\n persona = upsert_persona(\n user=None, # Slack channel Personas are not attached to users\n persona_id=persona_id_to_update,\n name=persona_name,\n description=\"\",\n system_prompt=\"\",\n task_prompt=\"\",\n datetime_aware=True,\n tool_ids=[search_tool.id],\n document_set_ids=document_set_ids,\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n is_public=True,\n is_featured=False,\n db_session=db_session,\n commit=False,\n )\n\n return persona\n\n\ndef _no_ee_standard_answer_categories(\n *args: Any, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n) -> list:\n return []\n\n\ndef insert_slack_channel_config(\n db_session: Session,\n slack_bot_id: int,\n persona_id: int | None,\n channel_config: ChannelConfig,\n standard_answer_category_ids: list[int],\n enable_auto_filters: bool,\n is_default: bool = False,\n) -> SlackChannelConfig:\n versioned_fetch_standard_answer_categories_by_ids = (\n fetch_versioned_implementation_with_fallback(\n \"onyx.db.standard_answer\",\n \"fetch_standard_answer_categories_by_ids\",\n _no_ee_standard_answer_categories,\n )\n )\n existing_standard_answer_categories = (\n versioned_fetch_standard_answer_categories_by_ids(\n standard_answer_category_ids=standard_answer_category_ids,\n db_session=db_session,\n )\n )\n\n if len(existing_standard_answer_categories) != len(standard_answer_category_ids):\n if len(existing_standard_answer_categories) == 0:\n raise EERequiredError(\n \"Standard answers are a paid Enterprise Edition feature - enable EE or remove standard answer categories\"\n )\n else:\n raise ValueError(\n f\"Some or all categories with ids {standard_answer_category_ids} do not exist\"\n )\n\n if is_default:\n existing_default = db_session.scalar(\n select(SlackChannelConfig).where(\n SlackChannelConfig.slack_bot_id == slack_bot_id,\n SlackChannelConfig.is_default is True, # type: ignore\n )\n )\n if existing_default:\n raise ValueError(\"A default config already exists for this Slack bot.\")\n else:\n if \"channel_name\" not in channel_config:\n raise ValueError(\"Channel name is required for non-default configs.\")\n\n slack_channel_config = SlackChannelConfig(\n slack_bot_id=slack_bot_id,\n persona_id=persona_id,\n channel_config=channel_config,\n standard_answer_categories=existing_standard_answer_categories,\n enable_auto_filters=enable_auto_filters,\n is_default=is_default,\n )\n db_session.add(slack_channel_config)\n db_session.commit()\n\n return slack_channel_config\n\n\ndef update_slack_channel_config(\n db_session: Session,\n slack_channel_config_id: int,\n persona_id: int | None,\n channel_config: ChannelConfig,\n standard_answer_category_ids: list[int],\n enable_auto_filters: bool,\n disabled: bool, # noqa: ARG001\n) -> SlackChannelConfig:\n slack_channel_config = db_session.scalar(\n select(SlackChannelConfig).where(\n SlackChannelConfig.id == slack_channel_config_id\n )\n )\n if slack_channel_config is None:\n raise ValueError(\n f\"Unable to find Slack channel config with ID {slack_channel_config_id}\"\n )\n\n versioned_fetch_standard_answer_categories_by_ids = (\n fetch_versioned_implementation_with_fallback(\n \"onyx.db.standard_answer\",\n \"fetch_standard_answer_categories_by_ids\",\n _no_ee_standard_answer_categories,\n )\n )\n existing_standard_answer_categories = (\n versioned_fetch_standard_answer_categories_by_ids(\n standard_answer_category_ids=standard_answer_category_ids,\n db_session=db_session,\n )\n )\n if len(existing_standard_answer_categories) != len(standard_answer_category_ids):\n raise ValueError(\n f\"Some or all categories with ids {standard_answer_category_ids} do not exist\"\n )\n\n # update the config\n slack_channel_config.persona_id = persona_id\n slack_channel_config.channel_config = channel_config\n slack_channel_config.standard_answer_categories = list(\n existing_standard_answer_categories\n )\n slack_channel_config.enable_auto_filters = enable_auto_filters\n\n db_session.commit()\n\n return slack_channel_config\n\n\ndef remove_slack_channel_config(\n db_session: Session,\n slack_channel_config_id: int,\n user: User,\n) -> None:\n slack_channel_config = db_session.scalar(\n select(SlackChannelConfig).where(\n SlackChannelConfig.id == slack_channel_config_id\n )\n )\n if slack_channel_config is None:\n raise ValueError(\n f\"Unable to find Slack channel config with ID {slack_channel_config_id}\"\n )\n\n existing_persona_id = slack_channel_config.persona_id\n if existing_persona_id:\n existing_persona = db_session.scalar(\n select(Persona).where(Persona.id == existing_persona_id)\n )\n # if the existing persona was one created just for use with this Slack channel,\n # then clean it up\n if existing_persona and existing_persona.name.startswith(\n SLACK_BOT_PERSONA_PREFIX\n ):\n _cleanup_relationships(\n db_session=db_session, persona_id=existing_persona_id\n )\n mark_persona_as_deleted(\n persona_id=existing_persona_id, user=user, db_session=db_session\n )\n\n db_session.delete(slack_channel_config)\n db_session.commit()\n\n\ndef fetch_slack_channel_configs(\n db_session: Session, slack_bot_id: int | None = None\n) -> Sequence[SlackChannelConfig]:\n if not slack_bot_id:\n return db_session.scalars(select(SlackChannelConfig)).all()\n\n return db_session.scalars(\n select(SlackChannelConfig).where(\n SlackChannelConfig.slack_bot_id == slack_bot_id\n )\n ).all()\n\n\ndef fetch_slack_channel_config(\n db_session: Session, slack_channel_config_id: int\n) -> SlackChannelConfig | None:\n return db_session.scalar(\n select(SlackChannelConfig).where(\n SlackChannelConfig.id == slack_channel_config_id\n )\n )\n\n\ndef fetch_slack_channel_config_for_channel_or_default(\n db_session: Session, slack_bot_id: int, channel_name: str | None\n) -> SlackChannelConfig | None:\n # attempt to find channel-specific config first\n if channel_name is not None:\n sc_config = db_session.scalar(\n select(SlackChannelConfig)\n .options(joinedload(SlackChannelConfig.persona))\n .where(\n SlackChannelConfig.slack_bot_id == slack_bot_id,\n SlackChannelConfig.channel_config[\"channel_name\"].astext\n == channel_name,\n )\n )\n else:\n sc_config = None\n\n if sc_config:\n return sc_config\n\n # if none found, see if there is a default\n default_sc = db_session.scalar(\n select(SlackChannelConfig)\n .options(joinedload(SlackChannelConfig.persona))\n .where(\n SlackChannelConfig.slack_bot_id == slack_bot_id,\n SlackChannelConfig.is_default == True, # noqa: E712\n )\n )\n\n return default_sc\n" }, { "path": "backend/onyx/db/swap_index.py", "content": "import time\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import VESPA_NUM_ATTEMPTS_ON_STARTUP\nfrom onyx.configs.constants import KV_REINDEX_KEY\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs\nfrom onyx.db.connector_credential_pair import resync_cc_pair\nfrom onyx.db.document import delete_all_documents_for_connector_credential_pair\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.enums import SwitchoverType\nfrom onyx.db.index_attempt import cancel_indexing_attempts_for_search_settings\nfrom onyx.db.index_attempt import (\n count_unique_active_cc_pairs_with_successful_index_attempts,\n)\nfrom onyx.db.index_attempt import count_unique_cc_pairs_with_successful_index_attempts\nfrom onyx.db.llm import update_default_contextual_model\nfrom onyx.db.llm import update_no_default_contextual_rag_provider\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.db.search_settings import get_secondary_search_settings\nfrom onyx.db.search_settings import update_search_settings_status\nfrom onyx.document_index.factory import get_all_document_indices\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _perform_index_swap(\n db_session: Session,\n new_search_settings: SearchSettings,\n all_cc_pairs: list[ConnectorCredentialPair],\n cleanup_documents: bool = False,\n) -> SearchSettings | None:\n \"\"\"Swap the indices and expire the old one.\n\n Returns the old search settings if the swap was successful, otherwise None.\n \"\"\"\n current_search_settings = get_current_search_settings(db_session)\n if len(all_cc_pairs) > 0:\n kv_store = get_kv_store()\n kv_store.store(KV_REINDEX_KEY, False)\n\n # Expire jobs for the now past index/embedding model\n cancel_indexing_attempts_for_search_settings(\n search_settings_id=current_search_settings.id,\n db_session=db_session,\n )\n\n # Recount aggregates\n for cc_pair in all_cc_pairs:\n resync_cc_pair(\n cc_pair=cc_pair,\n # sync based on the new search settings\n search_settings_id=new_search_settings.id,\n db_session=db_session,\n )\n\n if cleanup_documents:\n # clean up all DocumentByConnectorCredentialPair / Document rows, since we're\n # doing an instant swap and no documents will exist in the new index.\n for cc_pair in all_cc_pairs:\n delete_all_documents_for_connector_credential_pair(\n db_session=db_session,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n )\n\n # swap over search settings\n update_search_settings_status(\n search_settings=current_search_settings,\n new_status=IndexModelStatus.PAST,\n db_session=db_session,\n )\n update_search_settings_status(\n search_settings=new_search_settings,\n new_status=IndexModelStatus.PRESENT,\n db_session=db_session,\n )\n\n # Update the default contextual model to match the newly promoted settings\n try:\n update_default_contextual_model(\n db_session=db_session,\n enable_contextual_rag=new_search_settings.enable_contextual_rag,\n contextual_rag_llm_provider=new_search_settings.contextual_rag_llm_provider,\n contextual_rag_llm_name=new_search_settings.contextual_rag_llm_name,\n )\n except ValueError as e:\n logger.error(f\"Model not found, defaulting to no contextual model: {e}\")\n update_no_default_contextual_rag_provider(\n db_session=db_session,\n )\n new_search_settings.enable_contextual_rag = False\n new_search_settings.contextual_rag_llm_provider = None\n new_search_settings.contextual_rag_llm_name = None\n db_session.commit()\n\n # This flow is for checking and possibly creating an index so we get all\n # indices.\n document_indices = get_all_document_indices(new_search_settings, None, None)\n\n WAIT_SECONDS = 5\n\n for document_index in document_indices:\n success = False\n for x in range(VESPA_NUM_ATTEMPTS_ON_STARTUP):\n try:\n logger.notice(\n f\"Document index {document_index.__class__.__name__} swap (attempt {x + 1}/{VESPA_NUM_ATTEMPTS_ON_STARTUP})...\"\n )\n document_index.ensure_indices_exist(\n primary_embedding_dim=new_search_settings.final_embedding_dim,\n primary_embedding_precision=new_search_settings.embedding_precision,\n # just finished swap, no more secondary index\n secondary_index_embedding_dim=None,\n secondary_index_embedding_precision=None,\n )\n\n logger.notice(\"Document index swap complete.\")\n success = True\n break\n except Exception:\n logger.exception(\n f\"Document index swap for {document_index.__class__.__name__} did not succeed. \"\n f\"The document index services may not be ready yet. Retrying in {WAIT_SECONDS} seconds.\"\n )\n time.sleep(WAIT_SECONDS)\n\n if not success:\n logger.error(\n f\"Document index swap for {document_index.__class__.__name__} did not succeed. \"\n f\"Attempt limit reached. ({VESPA_NUM_ATTEMPTS_ON_STARTUP})\"\n )\n return None\n\n return current_search_settings\n\n\ndef check_and_perform_index_swap(db_session: Session) -> SearchSettings | None:\n \"\"\"Get count of cc-pairs and count of successful index_attempts for the\n new model grouped by connector + credential, if it's the same, then assume\n new index is done building. If so, swap the indices and expire the old one.\n\n Returns None if search settings did not change, or the old search settings if they\n did change.\n \"\"\"\n if DISABLE_VECTOR_DB:\n return None\n\n # Default CC-pair created for Ingestion API unused here\n all_cc_pairs = get_connector_credential_pairs(db_session)\n cc_pair_count = max(len(all_cc_pairs) - 1, 0)\n new_search_settings = get_secondary_search_settings(db_session)\n\n if not new_search_settings:\n return None\n\n # Handle switchover based on switchover_type\n switchover_type = new_search_settings.switchover_type\n\n # INSTANT: Swap immediately without waiting\n if switchover_type == SwitchoverType.INSTANT:\n return _perform_index_swap(\n db_session=db_session,\n new_search_settings=new_search_settings,\n all_cc_pairs=all_cc_pairs,\n # clean up all DocumentByConnectorCredentialPair / Document rows, since we're\n # doing an instant swap.\n cleanup_documents=True,\n )\n\n # REINDEX: Wait for all connectors to complete\n elif switchover_type == SwitchoverType.REINDEX:\n unique_cc_indexings = count_unique_cc_pairs_with_successful_index_attempts(\n search_settings_id=new_search_settings.id, db_session=db_session\n )\n\n # Index Attempts are cleaned up as well when the cc-pair is deleted so the logic in this\n # function is correct. The unique_cc_indexings are specifically for the existing cc-pairs\n if unique_cc_indexings > cc_pair_count:\n logger.error(\"More unique indexings than cc pairs, should not occur\")\n\n if cc_pair_count == 0 or cc_pair_count == unique_cc_indexings:\n # Swap indices\n return _perform_index_swap(\n db_session=db_session,\n new_search_settings=new_search_settings,\n all_cc_pairs=all_cc_pairs,\n )\n\n return None\n\n # ACTIVE_ONLY: Wait for only non-paused connectors to complete\n elif switchover_type == SwitchoverType.ACTIVE_ONLY:\n # Count non-paused cc_pairs (excluding the default Ingestion API cc_pair)\n active_cc_pairs = [\n cc_pair\n for cc_pair in all_cc_pairs\n if cc_pair.status != ConnectorCredentialPairStatus.PAUSED\n ]\n active_cc_pair_count = max(len(active_cc_pairs) - 1, 0)\n\n unique_active_cc_indexings = (\n count_unique_active_cc_pairs_with_successful_index_attempts(\n search_settings_id=new_search_settings.id, db_session=db_session\n )\n )\n\n if unique_active_cc_indexings > active_cc_pair_count:\n logger.error(\n \"More unique active indexings than active cc pairs, should not occur\"\n )\n\n if (\n active_cc_pair_count == 0\n or active_cc_pair_count == unique_active_cc_indexings\n ):\n # Swap indices\n return _perform_index_swap(\n db_session=db_session,\n new_search_settings=new_search_settings,\n all_cc_pairs=all_cc_pairs,\n )\n\n return None\n\n # Should not reach here, but handle gracefully\n logger.error(f\"Unknown switchover_type: {switchover_type}\")\n return None\n" }, { "path": "backend/onyx/db/sync_record.py", "content": "from sqlalchemy import and_\nfrom sqlalchemy import desc\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import SyncStatus\nfrom onyx.db.enums import SyncType\nfrom onyx.db.models import SyncRecord\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef insert_sync_record(\n db_session: Session,\n entity_id: int,\n sync_type: SyncType,\n) -> SyncRecord:\n \"\"\"Insert a new sync record into the database, cancelling any existing in-progress records.\n\n Args:\n db_session: The database session to use\n entity_id: The ID of the entity being synced (document set ID, user group ID, etc.)\n sync_type: The type of sync operation\n \"\"\"\n # If an existing in-progress sync record exists, mark as cancelled\n existing_in_progress_sync_record = fetch_latest_sync_record(\n db_session, entity_id, sync_type, sync_status=SyncStatus.IN_PROGRESS\n )\n\n if existing_in_progress_sync_record is not None:\n logger.info(\n f\"Cancelling existing in-progress sync record {existing_in_progress_sync_record.id} \"\n f\"for entity_id={entity_id} sync_type={sync_type}\"\n )\n mark_sync_records_as_cancelled(db_session, entity_id, sync_type)\n\n return _create_sync_record(db_session, entity_id, sync_type)\n\n\ndef mark_sync_records_as_cancelled(\n db_session: Session,\n entity_id: int | None,\n sync_type: SyncType,\n) -> None:\n stmt = (\n update(SyncRecord)\n .where(\n and_(\n SyncRecord.entity_id == entity_id,\n SyncRecord.sync_type == sync_type,\n SyncRecord.sync_status == SyncStatus.IN_PROGRESS,\n )\n )\n .values(sync_status=SyncStatus.CANCELED)\n )\n db_session.execute(stmt)\n db_session.commit()\n\n\ndef _create_sync_record(\n db_session: Session,\n entity_id: int | None,\n sync_type: SyncType,\n) -> SyncRecord:\n \"\"\"Create and insert a new sync record into the database.\"\"\"\n sync_record = SyncRecord(\n entity_id=entity_id,\n sync_type=sync_type,\n sync_status=SyncStatus.IN_PROGRESS,\n num_docs_synced=0,\n sync_start_time=func.now(),\n )\n db_session.add(sync_record)\n db_session.commit()\n\n return sync_record\n\n\ndef fetch_latest_sync_record(\n db_session: Session,\n entity_id: int,\n sync_type: SyncType,\n sync_status: SyncStatus | None = None,\n) -> SyncRecord | None:\n \"\"\"Fetch the most recent sync record for a given entity ID and status.\n\n Args:\n db_session: The database session to use\n entity_id: The ID of the entity to fetch sync record for\n sync_type: The type of sync operation\n \"\"\"\n stmt = (\n select(SyncRecord)\n .where(\n and_(\n SyncRecord.entity_id == entity_id,\n SyncRecord.sync_type == sync_type,\n )\n )\n .order_by(desc(SyncRecord.sync_start_time))\n .limit(1)\n )\n\n if sync_status is not None:\n stmt = stmt.where(SyncRecord.sync_status == sync_status)\n\n result = db_session.execute(stmt)\n return result.scalar_one_or_none()\n\n\ndef update_sync_record_status(\n db_session: Session,\n entity_id: int,\n sync_type: SyncType,\n sync_status: SyncStatus,\n num_docs_synced: int | None = None,\n) -> None:\n \"\"\"Update the status of a sync record.\n\n Args:\n db_session: The database session to use\n entity_id: The ID of the entity being synced\n sync_type: The type of sync operation\n sync_status: The new status to set\n num_docs_synced: Optional number of documents synced to update\n \"\"\"\n sync_record = fetch_latest_sync_record(db_session, entity_id, sync_type)\n if sync_record is None:\n raise ValueError(\n f\"No sync record found for entity_id={entity_id} sync_type={sync_type}\"\n )\n\n sync_record.sync_status = sync_status\n if num_docs_synced is not None:\n sync_record.num_docs_synced = num_docs_synced\n\n if sync_status.is_terminal():\n sync_record.sync_end_time = func.now() # type: ignore\n\n db_session.commit()\n\n\ndef cleanup_sync_records(\n db_session: Session, entity_id: int, sync_type: SyncType\n) -> None:\n \"\"\"Cleanup sync records for a given entity ID and sync type by marking them as failed.\"\"\"\n stmt = (\n update(SyncRecord)\n .where(SyncRecord.entity_id == entity_id)\n .where(SyncRecord.sync_type == sync_type)\n .where(SyncRecord.sync_status == SyncStatus.IN_PROGRESS)\n .values(sync_status=SyncStatus.CANCELED, sync_end_time=func.now())\n )\n db_session.execute(stmt)\n db_session.commit()\n" }, { "path": "backend/onyx/db/tag.py", "content": "from typing import Any\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import delete\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import Document\nfrom onyx.db.models import Document__Tag\nfrom onyx.db.models import Tag\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef check_tag_validity(tag_key: str, tag_value: str) -> bool:\n \"\"\"If a tag is too long, it should not be used (it will cause an error in Postgres\n as the unique constraint can only apply to entries that are less than 2704 bytes).\n\n Additionally, extremely long tags are not really usable / useful.\"\"\"\n if len(tag_key) + len(tag_value) > 255:\n logger.error(\n f\"Tag with key '{tag_key}' and value '{tag_value}' is too long, cannot be used\"\n )\n return False\n\n return True\n\n\ndef create_or_add_document_tag(\n tag_key: str,\n tag_value: str,\n source: DocumentSource,\n document_id: str,\n db_session: Session,\n) -> Tag | None:\n if not check_tag_validity(tag_key, tag_value):\n return None\n\n document = db_session.get(Document, document_id)\n if not document:\n raise ValueError(\"Invalid Document, cannot attach Tags\")\n\n # Use upsert to avoid race condition when multiple workers try to create the same tag\n insert_stmt = pg_insert(Tag).values(\n tag_key=tag_key,\n tag_value=tag_value,\n source=source,\n is_list=False,\n )\n insert_stmt = insert_stmt.on_conflict_do_nothing(\n constraint=\"_tag_key_value_source_list_uc\"\n )\n db_session.execute(insert_stmt)\n\n # Now fetch the tag (either just inserted or already existed)\n tag_stmt = select(Tag).where(\n Tag.tag_key == tag_key,\n Tag.tag_value == tag_value,\n Tag.source == source,\n Tag.is_list.is_(False),\n )\n tag = db_session.execute(tag_stmt).scalar_one()\n\n if tag not in document.tags:\n document.tags.append(tag)\n\n db_session.commit()\n return tag\n\n\ndef create_or_add_document_tag_list(\n tag_key: str,\n tag_values: list[str],\n source: DocumentSource,\n document_id: str,\n db_session: Session,\n) -> list[Tag]:\n valid_tag_values = [\n tag_value for tag_value in tag_values if check_tag_validity(tag_key, tag_value)\n ]\n if not valid_tag_values:\n return []\n\n document = db_session.get(Document, document_id)\n if not document:\n raise ValueError(\"Invalid Document, cannot attach Tags\")\n\n # Use upsert to avoid race condition when multiple workers try to create the same tags\n for tag_value in valid_tag_values:\n insert_stmt = pg_insert(Tag).values(\n tag_key=tag_key,\n tag_value=tag_value,\n source=source,\n is_list=True,\n )\n insert_stmt = insert_stmt.on_conflict_do_nothing(\n constraint=\"_tag_key_value_source_list_uc\"\n )\n db_session.execute(insert_stmt)\n\n # Now fetch all tags (either just inserted or already existed)\n all_tags_stmt = select(Tag).where(\n Tag.tag_key == tag_key,\n Tag.tag_value.in_(valid_tag_values),\n Tag.source == source,\n Tag.is_list.is_(True),\n )\n all_tags = list(db_session.execute(all_tags_stmt).scalars().all())\n\n for tag in all_tags:\n if tag not in document.tags:\n document.tags.append(tag)\n\n db_session.commit()\n return all_tags\n\n\ndef upsert_document_tags(\n document_id: str,\n source: DocumentSource,\n metadata: dict[str, str | list[str]],\n db_session: Session,\n) -> list[Tag]:\n document = db_session.get(Document, document_id)\n if not document:\n raise ValueError(\"Invalid Document, cannot attach Tags\")\n\n old_tag_ids: set[int] = {tag.id for tag in document.tags}\n\n new_tags: list[Tag] = []\n new_tag_ids: set[int] = set()\n for k, v in metadata.items():\n if isinstance(v, list):\n new_tags.extend(\n create_or_add_document_tag_list(k, v, source, document_id, db_session)\n )\n new_tag_ids.update({tag.id for tag in new_tags})\n continue\n\n new_tag = create_or_add_document_tag(k, v, source, document_id, db_session)\n if new_tag:\n new_tag_ids.add(new_tag.id)\n new_tags.append(new_tag)\n\n delete_tags = old_tag_ids - new_tag_ids\n if delete_tags:\n delete_stmt = delete(Document__Tag).where(\n Document__Tag.document_id == document_id,\n Document__Tag.tag_id.in_(delete_tags),\n )\n db_session.execute(delete_stmt)\n db_session.commit()\n\n return new_tags\n\n\ndef find_tags(\n tag_key_prefix: str | None,\n tag_value_prefix: str | None,\n sources: list[DocumentSource] | None,\n limit: int | None,\n db_session: Session,\n # if set, both tag_key_prefix and tag_value_prefix must be a match\n require_both_to_match: bool = False,\n) -> list[Tag]:\n query = select(Tag)\n\n if tag_key_prefix or tag_value_prefix:\n conditions = []\n if tag_key_prefix:\n conditions.append(Tag.tag_key.ilike(f\"{tag_key_prefix}%\"))\n if tag_value_prefix:\n conditions.append(Tag.tag_value.ilike(f\"{tag_value_prefix}%\"))\n\n final_prefix_condition = (\n and_(*conditions) if require_both_to_match else or_(*conditions)\n )\n query = query.where(final_prefix_condition)\n\n if sources:\n query = query.where(Tag.source.in_(sources))\n\n if limit:\n query = query.limit(limit)\n\n result = db_session.execute(query)\n\n tags = result.scalars().all()\n return list(tags)\n\n\ndef get_structured_tags_for_document(\n document_id: str, db_session: Session\n) -> dict[str, str | list[str]]:\n \"\"\"Essentially returns the document metadata from postgres.\"\"\"\n document = db_session.get(Document, document_id)\n if not document:\n raise ValueError(\"Invalid Document, cannot find tags\")\n\n document_metadata: dict[str, Any] = {}\n for tag in document.tags:\n if tag.is_list:\n document_metadata.setdefault(tag.tag_key, [])\n # should always be a list (if tag.is_list is always True for this key), but just in case\n if not isinstance(document_metadata[tag.tag_key], list):\n logger.warning(\n \"Inconsistent is_list for document %s, tag_key %s\",\n document_id,\n tag.tag_key,\n )\n document_metadata[tag.tag_key] = [document_metadata[tag.tag_key]]\n document_metadata[tag.tag_key].append(tag.tag_value)\n continue\n\n # set value (ignore duplicate keys, though there should be none)\n document_metadata.setdefault(tag.tag_key, tag.tag_value)\n\n # should always be a value, but just in case (treat it as a list in this case)\n if isinstance(document_metadata[tag.tag_key], list):\n logger.warning(\n \"Inconsistent is_list for document %s, tag_key %s\",\n document_id,\n tag.tag_key,\n )\n document_metadata[tag.tag_key] = [document_metadata[tag.tag_key]]\n return document_metadata\n\n\ndef delete_document_tags_for_documents__no_commit(\n document_ids: list[str], db_session: Session\n) -> None:\n stmt = delete(Document__Tag).where(Document__Tag.document_id.in_(document_ids))\n db_session.execute(stmt)\n\n\ndef delete_orphan_tags__no_commit(db_session: Session) -> None:\n orphan_tags_query = select(Tag.id).where(\n ~db_session.query(Document__Tag.tag_id)\n .filter(Document__Tag.tag_id == Tag.id)\n .exists()\n )\n\n orphan_tags = db_session.execute(orphan_tags_query).scalars().all()\n\n if orphan_tags:\n delete_orphan_tags_stmt = delete(Tag).where(Tag.id.in_(orphan_tags))\n db_session.execute(delete_orphan_tags_stmt)\n" }, { "path": "backend/onyx/db/tasks.py", "content": "from datetime import datetime\n\nfrom sqlalchemy import desc\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql import delete\n\nfrom onyx.configs.app_configs import JOB_TIMEOUT\nfrom onyx.db.engine.time_utils import get_db_current_time\nfrom onyx.db.models import TaskQueueState\nfrom onyx.db.models import TaskStatus\n\n\ndef get_latest_task(\n task_name: str,\n db_session: Session,\n) -> TaskQueueState | None:\n stmt = (\n select(TaskQueueState)\n .where(TaskQueueState.task_name == task_name)\n .order_by(desc(TaskQueueState.id))\n .limit(1)\n )\n\n result = db_session.execute(stmt)\n latest_task = result.scalars().first()\n\n return latest_task\n\n\ndef get_latest_task_by_type(\n task_name: str,\n db_session: Session,\n) -> TaskQueueState | None:\n stmt = (\n select(TaskQueueState)\n .where(TaskQueueState.task_name.like(f\"%{task_name}%\"))\n .order_by(desc(TaskQueueState.id))\n .limit(1)\n )\n\n result = db_session.execute(stmt)\n latest_task = result.scalars().first()\n\n return latest_task\n\n\ndef register_task(\n task_name: str,\n db_session: Session,\n task_id: str = \"\",\n status: TaskStatus = TaskStatus.PENDING,\n start_time: datetime | None = None,\n) -> TaskQueueState:\n new_task = TaskQueueState(\n task_id=task_id,\n task_name=task_name,\n status=status,\n start_time=start_time,\n )\n\n db_session.add(new_task)\n db_session.commit()\n\n return new_task\n\n\ndef get_task_with_id(\n db_session: Session,\n task_id: str,\n) -> TaskQueueState | None:\n return db_session.scalar(\n select(TaskQueueState).where(TaskQueueState.task_id == task_id)\n )\n\n\ndef delete_task_with_id(\n db_session: Session,\n task_id: str,\n) -> None:\n db_session.execute(delete(TaskQueueState).where(TaskQueueState.task_id == task_id))\n db_session.commit()\n\n\ndef get_all_tasks_with_prefix(\n db_session: Session, task_name_prefix: str\n) -> list[TaskQueueState]:\n return list(\n db_session.scalars(\n select(TaskQueueState).where(\n TaskQueueState.task_name.like(f\"{task_name_prefix}_%\")\n )\n )\n )\n\n\ndef mark_task_as_started_with_id(\n db_session: Session,\n task_id: str,\n) -> None:\n task = get_task_with_id(db_session=db_session, task_id=task_id)\n if not task:\n raise RuntimeError(f\"A task with the task-id {task_id=} does not exist\")\n\n task.status = TaskStatus.STARTED\n db_session.commit()\n\n\ndef mark_task_as_finished_with_id(\n db_session: Session,\n task_id: str,\n success: bool = True,\n) -> None:\n task = get_task_with_id(db_session=db_session, task_id=task_id)\n if not task:\n raise RuntimeError(f\"A task with the task-id {task_id=} does not exist\")\n\n task.status = TaskStatus.SUCCESS if success else TaskStatus.FAILURE\n db_session.commit()\n\n\ndef mark_task_start(\n task_name: str,\n db_session: Session,\n) -> None:\n task = get_latest_task(task_name, db_session)\n if not task:\n raise ValueError(f\"No task found with name {task_name}\")\n\n task.start_time = func.now() # type: ignore\n db_session.commit()\n\n\ndef mark_task_finished(\n task_name: str,\n db_session: Session,\n success: bool = True,\n) -> None:\n latest_task = get_latest_task(task_name, db_session)\n if latest_task is None:\n raise ValueError(f\"tasks for {task_name} do not exist\")\n\n latest_task.status = TaskStatus.SUCCESS if success else TaskStatus.FAILURE\n db_session.commit()\n\n\ndef check_task_is_live_and_not_timed_out(\n task: TaskQueueState,\n db_session: Session,\n timeout: int = JOB_TIMEOUT,\n) -> bool:\n # We only care for live tasks to not create new periodic tasks\n if task.status in [TaskStatus.SUCCESS, TaskStatus.FAILURE]:\n return False\n\n current_db_time = get_db_current_time(db_session=db_session)\n\n last_update_time = task.register_time\n if task.start_time:\n last_update_time = max(task.register_time, task.start_time)\n\n time_elapsed = current_db_time - last_update_time\n return time_elapsed.total_seconds() < timeout\n" }, { "path": "backend/onyx/db/token_limit.py", "content": "from collections.abc import Sequence\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import TokenRateLimitScope\nfrom onyx.db.models import TokenRateLimit\nfrom onyx.db.models import TokenRateLimit__UserGroup\nfrom onyx.server.token_rate_limits.models import TokenRateLimitArgs\n\n\ndef fetch_all_user_token_rate_limits(\n db_session: Session,\n enabled_only: bool = False,\n ordered: bool = True,\n) -> Sequence[TokenRateLimit]:\n query = select(TokenRateLimit).where(\n TokenRateLimit.scope == TokenRateLimitScope.USER\n )\n\n if enabled_only:\n query = query.where(TokenRateLimit.enabled.is_(True))\n\n if ordered:\n query = query.order_by(TokenRateLimit.created_at.desc())\n\n return db_session.scalars(query).all()\n\n\ndef fetch_all_global_token_rate_limits(\n db_session: Session,\n enabled_only: bool = False,\n ordered: bool = True,\n) -> Sequence[TokenRateLimit]:\n query = select(TokenRateLimit).where(\n TokenRateLimit.scope == TokenRateLimitScope.GLOBAL\n )\n\n if enabled_only:\n query = query.where(TokenRateLimit.enabled.is_(True))\n\n if ordered:\n query = query.order_by(TokenRateLimit.created_at.desc())\n\n token_rate_limits = db_session.scalars(query).all()\n return token_rate_limits\n\n\ndef insert_user_token_rate_limit(\n db_session: Session,\n token_rate_limit_settings: TokenRateLimitArgs,\n) -> TokenRateLimit:\n token_limit = TokenRateLimit(\n enabled=token_rate_limit_settings.enabled,\n token_budget=token_rate_limit_settings.token_budget,\n period_hours=token_rate_limit_settings.period_hours,\n scope=TokenRateLimitScope.USER,\n )\n db_session.add(token_limit)\n db_session.commit()\n\n return token_limit\n\n\ndef insert_global_token_rate_limit(\n db_session: Session,\n token_rate_limit_settings: TokenRateLimitArgs,\n) -> TokenRateLimit:\n token_limit = TokenRateLimit(\n enabled=token_rate_limit_settings.enabled,\n token_budget=token_rate_limit_settings.token_budget,\n period_hours=token_rate_limit_settings.period_hours,\n scope=TokenRateLimitScope.GLOBAL,\n )\n db_session.add(token_limit)\n db_session.commit()\n\n return token_limit\n\n\ndef update_token_rate_limit(\n db_session: Session,\n token_rate_limit_id: int,\n token_rate_limit_settings: TokenRateLimitArgs,\n) -> TokenRateLimit:\n token_limit = db_session.get(TokenRateLimit, token_rate_limit_id)\n if token_limit is None:\n raise ValueError(f\"TokenRateLimit with id '{token_rate_limit_id}' not found\")\n\n token_limit.enabled = token_rate_limit_settings.enabled\n token_limit.token_budget = token_rate_limit_settings.token_budget\n token_limit.period_hours = token_rate_limit_settings.period_hours\n db_session.commit()\n\n return token_limit\n\n\ndef delete_token_rate_limit(\n db_session: Session,\n token_rate_limit_id: int,\n) -> None:\n token_limit = db_session.get(TokenRateLimit, token_rate_limit_id)\n if token_limit is None:\n raise ValueError(f\"TokenRateLimit with id '{token_rate_limit_id}' not found\")\n\n db_session.query(TokenRateLimit__UserGroup).filter(\n TokenRateLimit__UserGroup.rate_limit_id == token_rate_limit_id\n ).delete()\n\n db_session.delete(token_limit)\n db_session.commit()\n" }, { "path": "backend/onyx/db/tools.py", "content": "from typing import Any\nfrom typing import cast\nfrom typing import Type\nfrom typing import TYPE_CHECKING\nfrom uuid import UUID\n\nfrom sqlalchemy import func\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.constants import UNSET\nfrom onyx.db.constants import UnsetType\nfrom onyx.db.enums import MCPServerStatus\nfrom onyx.db.models import MCPServer\nfrom onyx.db.models import OAuthConfig\nfrom onyx.db.models import Tool\nfrom onyx.db.models import ToolCall\nfrom onyx.server.features.tool.models import Header\nfrom onyx.tools.built_in_tools import BUILT_IN_TOOL_TYPES\nfrom onyx.utils.headers import HeaderItemDict\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.postgres_sanitization import sanitize_json_like\nfrom onyx.utils.postgres_sanitization import sanitize_string\n\nif TYPE_CHECKING:\n pass\n\nlogger = setup_logger()\n\n\ndef get_tools(\n db_session: Session,\n *,\n only_enabled: bool = False,\n only_connected_mcp: bool = False,\n only_openapi: bool = False,\n) -> list[Tool]:\n query = select(Tool)\n\n if only_connected_mcp:\n # Keep tools that either:\n # 1. Don't have an MCP server (mcp_server_id IS NULL) - Non-MCP tools\n # 2. Have an MCP server that is connected - Connected MCP tools\n query = query.outerjoin(MCPServer, Tool.mcp_server_id == MCPServer.id).where(\n or_(\n Tool.mcp_server_id.is_(None), # Non-MCP tools (built-in, custom)\n MCPServer.status == MCPServerStatus.CONNECTED, # MCP tools connected\n )\n )\n\n if only_enabled:\n query = query.where(Tool.enabled.is_(True))\n\n if only_openapi:\n query = query.where(\n Tool.openapi_schema.is_not(None),\n # To avoid showing rows that have JSON literal `null` stored in the column to the user.\n # tools from mcp servers will not have an openapi schema but it has `null`, so we need to exclude them.\n func.jsonb_typeof(Tool.openapi_schema) == \"object\",\n # Exclude built-in tools that happen to have an openapi_schema\n Tool.in_code_tool_id.is_(None),\n )\n\n return list(db_session.scalars(query).all())\n\n\ndef get_tools_by_mcp_server_id(\n mcp_server_id: int,\n db_session: Session,\n *,\n only_enabled: bool = False,\n order_by_id: bool = False,\n) -> list[Tool]:\n query = select(Tool).where(Tool.mcp_server_id == mcp_server_id)\n if only_enabled:\n query = query.where(Tool.enabled.is_(True))\n if order_by_id:\n query = query.order_by(Tool.id)\n return list(db_session.scalars(query).all())\n\n\ndef get_tools_by_ids(tool_ids: list[int], db_session: Session) -> list[Tool]:\n if not tool_ids:\n return []\n stmt = select(Tool).where(Tool.id.in_(tool_ids))\n return list(db_session.scalars(stmt).all())\n\n\ndef get_tool_by_id(tool_id: int, db_session: Session) -> Tool:\n tool = db_session.scalar(select(Tool).where(Tool.id == tool_id))\n if not tool:\n raise ValueError(\"Tool by specified id does not exist\")\n return tool\n\n\ndef get_tool_by_name(tool_name: str, db_session: Session) -> Tool:\n tool = db_session.scalar(select(Tool).where(Tool.name == tool_name))\n if not tool:\n raise ValueError(\"Tool by specified name does not exist\")\n return tool\n\n\ndef create_tool__no_commit(\n name: str,\n description: str | None,\n openapi_schema: dict[str, Any] | None,\n custom_headers: list[Header] | None,\n user_id: UUID | None,\n db_session: Session,\n passthrough_auth: bool,\n *,\n mcp_server_id: int | None = None,\n oauth_config_id: int | None = None,\n enabled: bool = True,\n) -> Tool:\n new_tool = Tool(\n name=name,\n description=description,\n in_code_tool_id=None,\n openapi_schema=openapi_schema,\n custom_headers=(\n [header.model_dump() for header in custom_headers] if custom_headers else []\n ),\n user_id=user_id,\n passthrough_auth=passthrough_auth,\n mcp_server_id=mcp_server_id,\n oauth_config_id=oauth_config_id,\n enabled=enabled,\n )\n db_session.add(new_tool)\n db_session.flush() # Don't commit yet, let caller decide when to commit\n return new_tool\n\n\ndef update_tool(\n tool_id: int,\n name: str | None,\n description: str | None,\n openapi_schema: dict[str, Any] | None,\n custom_headers: list[Header] | None,\n user_id: UUID | None,\n db_session: Session,\n passthrough_auth: bool | None,\n oauth_config_id: int | None | UnsetType = UNSET,\n) -> Tool:\n tool = get_tool_by_id(tool_id, db_session)\n if tool is None:\n raise ValueError(f\"Tool with ID {tool_id} does not exist\")\n\n if name is not None:\n tool.name = name\n if description is not None:\n tool.description = description\n if openapi_schema is not None:\n tool.openapi_schema = openapi_schema\n if user_id is not None:\n tool.user_id = user_id\n if custom_headers is not None:\n tool.custom_headers = [\n cast(HeaderItemDict, header.model_dump()) for header in custom_headers\n ]\n if passthrough_auth is not None:\n tool.passthrough_auth = passthrough_auth\n old_oauth_config_id = tool.oauth_config_id\n if not isinstance(oauth_config_id, UnsetType):\n tool.oauth_config_id = oauth_config_id\n db_session.flush()\n\n # Clean up orphaned OAuthConfig if the oauth_config_id was changed\n if (\n old_oauth_config_id is not None\n and not isinstance(oauth_config_id, UnsetType)\n and old_oauth_config_id != oauth_config_id\n ):\n other_tools = db_session.scalars(\n select(Tool).where(Tool.oauth_config_id == old_oauth_config_id)\n ).all()\n if not other_tools:\n oauth_config = db_session.get(OAuthConfig, old_oauth_config_id)\n if oauth_config:\n db_session.delete(oauth_config)\n\n db_session.commit()\n return tool\n\n\ndef delete_tool__no_commit(tool_id: int, db_session: Session) -> None:\n tool = get_tool_by_id(tool_id, db_session)\n if tool is None:\n raise ValueError(f\"Tool with ID {tool_id} does not exist\")\n\n oauth_config_id = tool.oauth_config_id\n\n db_session.delete(tool)\n db_session.flush()\n\n # Clean up orphaned OAuthConfig if no other tools reference it\n if oauth_config_id is not None:\n other_tools = db_session.scalars(\n select(Tool).where(Tool.oauth_config_id == oauth_config_id)\n ).all()\n if not other_tools:\n oauth_config = db_session.get(OAuthConfig, oauth_config_id)\n if oauth_config:\n db_session.delete(oauth_config)\n db_session.flush()\n\n\ndef get_builtin_tool(\n db_session: Session,\n tool_type: Type[BUILT_IN_TOOL_TYPES],\n) -> Tool:\n \"\"\"\n Retrieves a built-in tool from the database based on the tool type.\n \"\"\"\n # local import to avoid circular import. DB layer should not depend on tools layer.\n from onyx.tools.built_in_tools import BUILT_IN_TOOL_MAP\n\n tool_id = next(\n (\n in_code_tool_id\n for in_code_tool_id, tool_cls in BUILT_IN_TOOL_MAP.items()\n if tool_cls.__name__ == tool_type.__name__\n ),\n None,\n )\n\n if not tool_id:\n raise RuntimeError(\n f\"Tool type {tool_type.__name__} not found in the BUILT_IN_TOOLS list.\"\n )\n\n db_tool = db_session.execute(\n select(Tool).where(Tool.in_code_tool_id == tool_id)\n ).scalar_one_or_none()\n\n if not db_tool:\n raise RuntimeError(f\"Tool type {tool_type.__name__} not found in the database.\")\n\n return db_tool\n\n\ndef create_tool_call_no_commit(\n chat_session_id: UUID,\n parent_chat_message_id: int | None,\n turn_number: int,\n tool_id: int,\n tool_call_id: str,\n tool_call_arguments: dict[str, Any],\n tool_call_response: Any,\n tool_call_tokens: int,\n db_session: Session,\n *,\n parent_tool_call_id: int | None = None,\n reasoning_tokens: str | None = None,\n generated_images: list[dict] | None = None,\n tab_index: int = 0,\n add_only: bool = True,\n) -> ToolCall:\n \"\"\"\n Create a ToolCall entry in the database.\n\n Args:\n chat_session_id: The chat session ID\n parent_chat_message_id: The parent chat message ID\n turn_number: The turn number for this tool call\n tool_id: The tool ID\n tool_call_id: The tool call ID (string identifier from LLM)\n tool_call_arguments: The tool call arguments\n tool_call_response: The tool call response\n tool_call_tokens: The number of tokens in the tool call arguments\n db_session: The database session\n parent_tool_call_id: Optional parent tool call ID (for nested tool calls)\n reasoning_tokens: Optional reasoning tokens\n generated_images: Optional list of generated image metadata for replay\n tab_index: Index order of tool calls from the LLM for parallel tool calls\n commit: If True, commit the transaction; if False, flush only\n\n Returns:\n The created ToolCall object\n \"\"\"\n tool_call = ToolCall(\n chat_session_id=chat_session_id,\n parent_chat_message_id=parent_chat_message_id,\n parent_tool_call_id=parent_tool_call_id,\n turn_number=turn_number,\n tab_index=tab_index,\n tool_id=tool_id,\n tool_call_id=tool_call_id,\n reasoning_tokens=(\n sanitize_string(reasoning_tokens) if reasoning_tokens else reasoning_tokens\n ),\n tool_call_arguments=sanitize_json_like(tool_call_arguments),\n tool_call_response=sanitize_json_like(tool_call_response),\n tool_call_tokens=tool_call_tokens,\n generated_images=sanitize_json_like(generated_images),\n )\n\n db_session.add(tool_call)\n if not add_only:\n db_session.add(tool_call)\n else:\n db_session.flush()\n return tool_call\n" }, { "path": "backend/onyx/db/usage.py", "content": "\"\"\"Database interactions for tenant usage tracking (cloud usage limits).\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\nfrom enum import Enum\n\nfrom pydantic import BaseModel\nfrom sqlalchemy import select\nfrom sqlalchemy.dialects.postgresql import insert as pg_insert\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import TenantUsage\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import USAGE_LIMIT_WINDOW_SECONDS\n\nlogger = setup_logger()\n\n\nclass UsageType(str, Enum):\n \"\"\"Types of usage that can be tracked and limited.\"\"\"\n\n LLM_COST = \"llm_cost_cents\"\n CHUNKS_INDEXED = \"chunks_indexed\"\n API_CALLS = \"api_calls\"\n NON_STREAMING_API_CALLS = \"non_streaming_api_calls\"\n\n\nclass TenantUsageStats(BaseModel):\n \"\"\"Current usage statistics for a tenant.\"\"\"\n\n window_start: datetime\n llm_cost_cents: float\n chunks_indexed: int\n api_calls: int\n non_streaming_api_calls: int\n\n\nclass UsageLimitExceededError(Exception):\n \"\"\"Raised when a tenant exceeds their usage limit.\"\"\"\n\n def __init__(self, usage_type: UsageType, current: float, limit: float):\n self.usage_type = usage_type\n self.current = current\n self.limit = limit\n super().__init__(\n f\"Usage limit exceeded for {usage_type.value}: current usage {current}, limit {limit}\"\n )\n\n\ndef get_current_window_start() -> datetime:\n \"\"\"\n Calculate the start of the current usage window.\n\n Uses fixed windows aligned to Monday 00:00 UTC for predictability.\n The window duration is configured via USAGE_LIMIT_WINDOW_SECONDS.\n \"\"\"\n now = datetime.now(timezone.utc)\n # For weekly windows (default), align to Monday 00:00 UTC\n if USAGE_LIMIT_WINDOW_SECONDS == 604800: # 1 week\n # Get the start of the current week (Monday)\n days_since_monday = now.weekday()\n window_start = now.replace(\n hour=0, minute=0, second=0, microsecond=0\n ) - __import__(\"datetime\").timedelta(days=days_since_monday)\n return window_start\n\n # For other window sizes, use epoch-aligned windows\n epoch = datetime(1970, 1, 1, tzinfo=timezone.utc)\n seconds_since_epoch = int((now - epoch).total_seconds())\n window_number = seconds_since_epoch // USAGE_LIMIT_WINDOW_SECONDS\n window_start_seconds = window_number * USAGE_LIMIT_WINDOW_SECONDS\n return epoch + __import__(\"datetime\").timedelta(seconds=window_start_seconds)\n\n\ndef get_or_create_tenant_usage(\n db_session: Session,\n window_start: datetime | None = None,\n) -> TenantUsage:\n \"\"\"\n Get or create the usage record for the current window.\n\n Uses INSERT ... ON CONFLICT DO UPDATE to atomically create or get the record,\n avoiding TOCTOU race conditions where two concurrent requests could both\n attempt to insert a new record.\n \"\"\"\n if window_start is None:\n window_start = get_current_window_start()\n\n # Atomic upsert: insert if not exists, or update a field to itself if exists\n # This ensures we always get back a valid row without race conditions\n stmt = (\n pg_insert(TenantUsage)\n .values(\n window_start=window_start,\n llm_cost_cents=0.0,\n chunks_indexed=0,\n api_calls=0,\n non_streaming_api_calls=0,\n )\n .on_conflict_do_update(\n index_elements=[\"window_start\"],\n # No-op update: just set a field to its current value\n # This ensures the row is returned even on conflict\n set_={\"llm_cost_cents\": TenantUsage.llm_cost_cents},\n )\n .returning(TenantUsage)\n )\n\n result = db_session.execute(stmt).scalar_one()\n db_session.flush()\n\n return result\n\n\ndef get_tenant_usage_stats(\n db_session: Session,\n window_start: datetime | None = None,\n) -> TenantUsageStats:\n \"\"\"Get the current usage statistics for the tenant (read-only, no lock).\"\"\"\n if window_start is None:\n window_start = get_current_window_start()\n\n usage = db_session.execute(\n select(TenantUsage).where(TenantUsage.window_start == window_start)\n ).scalar_one_or_none()\n\n if usage is None:\n # No usage recorded yet for this window\n return TenantUsageStats(\n window_start=window_start,\n llm_cost_cents=0.0,\n chunks_indexed=0,\n api_calls=0,\n non_streaming_api_calls=0,\n )\n\n return TenantUsageStats(\n window_start=usage.window_start,\n llm_cost_cents=usage.llm_cost_cents,\n chunks_indexed=usage.chunks_indexed,\n api_calls=usage.api_calls,\n non_streaming_api_calls=usage.non_streaming_api_calls,\n )\n\n\ndef increment_usage(\n db_session: Session,\n usage_type: UsageType,\n amount: float | int,\n) -> None:\n \"\"\"\n Atomically increment a usage counter.\n\n Uses row-level locking to prevent race conditions.\n The caller should handle the transaction commit.\n \"\"\"\n usage = get_or_create_tenant_usage(db_session)\n\n if usage_type == UsageType.LLM_COST:\n usage.llm_cost_cents += float(amount)\n elif usage_type == UsageType.CHUNKS_INDEXED:\n usage.chunks_indexed += int(amount)\n elif usage_type == UsageType.API_CALLS:\n usage.api_calls += int(amount)\n elif usage_type == UsageType.NON_STREAMING_API_CALLS:\n usage.non_streaming_api_calls += int(amount)\n\n db_session.flush()\n\n\ndef check_usage_limit(\n db_session: Session,\n usage_type: UsageType,\n limit: float | int,\n pending_amount: float | int = 0,\n) -> None:\n \"\"\"\n Check if the current usage plus pending amount would exceed the limit.\n\n Args:\n db_session: Database session\n usage_type: Type of usage to check\n limit: The maximum allowed usage\n pending_amount: Amount about to be used (to check before committing)\n\n Raises:\n UsageLimitExceededError: If usage would exceed the limit\n \"\"\"\n stats = get_tenant_usage_stats(db_session)\n\n current_value: float\n if usage_type == UsageType.LLM_COST:\n current_value = stats.llm_cost_cents\n elif usage_type == UsageType.CHUNKS_INDEXED:\n current_value = float(stats.chunks_indexed)\n elif usage_type == UsageType.API_CALLS:\n current_value = float(stats.api_calls)\n elif usage_type == UsageType.NON_STREAMING_API_CALLS:\n current_value = float(stats.non_streaming_api_calls)\n else:\n current_value = 0.0\n\n if current_value + pending_amount > limit:\n raise UsageLimitExceededError(\n usage_type=usage_type,\n current=current_value + pending_amount,\n limit=float(limit),\n )\n" }, { "path": "backend/onyx/db/user_file.py", "content": "import datetime\nfrom uuid import UUID\n\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import joinedload\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Project__UserFile\nfrom onyx.db.models import UserFile\n\n\ndef fetch_chunk_counts_for_user_files(\n user_file_ids: list[str],\n db_session: Session,\n) -> list[tuple[str, int]]:\n \"\"\"\n Return a list of (user_file_id, chunk_count) tuples.\n If a user_file_id is not found in the database, it will be returned with a chunk_count of 0.\n \"\"\"\n stmt = select(UserFile.id, UserFile.chunk_count).where(\n UserFile.id.in_(user_file_ids)\n )\n\n results = db_session.execute(stmt).all()\n\n # Create a dictionary of user_file_id to chunk_count\n chunk_counts = {str(row.id): row.chunk_count or 0 for row in results}\n\n # Return a list of tuples, preserving `None` for documents not found or with\n # an unknown chunk count. Callers should handle the `None` case and fall\n # back to an existence check against the vector DB if necessary.\n return [\n (user_file_id, chunk_counts.get(user_file_id, 0))\n for user_file_id in user_file_ids\n ]\n\n\ndef calculate_user_files_token_count(file_ids: list[UUID], db_session: Session) -> int:\n \"\"\"Calculate total token count for specified files\"\"\"\n total_tokens = 0\n\n # Get tokens from individual files\n if file_ids:\n file_tokens = (\n db_session.query(func.sum(UserFile.token_count))\n .filter(UserFile.id.in_(file_ids))\n .scalar()\n or 0\n )\n total_tokens += file_tokens\n\n return total_tokens\n\n\ndef fetch_user_project_ids_for_user_files(\n user_file_ids: list[str],\n db_session: Session,\n) -> dict[str, list[int]]:\n \"\"\"Fetch user project ids for specified user files\"\"\"\n user_file_uuid_ids = [UUID(user_file_id) for user_file_id in user_file_ids]\n stmt = select(Project__UserFile.user_file_id, Project__UserFile.project_id).where(\n Project__UserFile.user_file_id.in_(user_file_uuid_ids)\n )\n rows = db_session.execute(stmt).all()\n\n user_file_id_to_project_ids: dict[str, list[int]] = {\n user_file_id: [] for user_file_id in user_file_ids\n }\n for user_file_id, project_id in rows:\n user_file_id_to_project_ids[str(user_file_id)].append(project_id)\n\n return user_file_id_to_project_ids\n\n\ndef fetch_persona_ids_for_user_files(\n user_file_ids: list[str],\n db_session: Session,\n) -> dict[str, list[int]]:\n \"\"\"Fetch persona (assistant) ids for specified user files.\"\"\"\n stmt = (\n select(UserFile)\n .where(UserFile.id.in_(user_file_ids))\n .options(selectinload(UserFile.assistants))\n )\n results = db_session.execute(stmt).scalars().all()\n return {\n str(user_file.id): [persona.id for persona in user_file.assistants]\n for user_file in results\n }\n\n\ndef update_last_accessed_at_for_user_files(\n user_file_ids: list[UUID],\n db_session: Session,\n) -> None:\n \"\"\"Update `last_accessed_at` to now (UTC) for the given user files.\"\"\"\n if not user_file_ids:\n return\n now = datetime.datetime.now(datetime.timezone.utc)\n (\n db_session.query(UserFile)\n .filter(UserFile.id.in_(user_file_ids))\n .update({UserFile.last_accessed_at: now}, synchronize_session=False)\n )\n db_session.commit()\n\n\ndef get_file_id_by_user_file_id(user_file_id: str, db_session: Session) -> str | None:\n user_file = db_session.query(UserFile).filter(UserFile.id == user_file_id).first()\n if user_file:\n return user_file.file_id\n return None\n\n\ndef get_file_ids_by_user_file_ids(\n user_file_ids: list[UUID], db_session: Session\n) -> list[str]:\n user_files = db_session.query(UserFile).filter(UserFile.id.in_(user_file_ids)).all()\n return [user_file.file_id for user_file in user_files]\n\n\ndef fetch_user_files_with_access_relationships(\n user_file_ids: list[str],\n db_session: Session,\n eager_load_groups: bool = False,\n) -> list[UserFile]:\n \"\"\"Fetch user files with the owner and assistant relationships\n eagerly loaded (needed for computing access control).\n\n When eager_load_groups is True, Persona.groups is also loaded so that\n callers can extract user-group names without a second DB round-trip.\"\"\"\n persona_sub_options = [\n selectinload(Persona.users),\n selectinload(Persona.user),\n ]\n if eager_load_groups:\n persona_sub_options.append(selectinload(Persona.groups))\n\n return (\n db_session.query(UserFile)\n .options(\n joinedload(UserFile.user),\n selectinload(UserFile.assistants).options(*persona_sub_options),\n )\n .filter(UserFile.id.in_(user_file_ids))\n .all()\n )\n" }, { "path": "backend/onyx/db/user_preferences.py", "content": "from collections.abc import Sequence\nfrom uuid import UUID\n\nfrom sqlalchemy import Column\nfrom sqlalchemy import delete\nfrom sqlalchemy import desc\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.db.enums import AccountType\nfrom onyx.db.enums import DefaultAppMode\nfrom onyx.db.enums import ThemePreference\nfrom onyx.db.models import AccessToken\nfrom onyx.db.models import Assistant__UserSpecificConfig\nfrom onyx.db.models import Memory\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.db.permissions import recompute_user_permissions__no_commit\nfrom onyx.db.users import assign_user_to_default_groups__no_commit\nfrom onyx.server.manage.models import MemoryItem\nfrom onyx.server.manage.models import UserSpecificAssistantPreference\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n_ROLE_TO_ACCOUNT_TYPE: dict[UserRole, AccountType] = {\n UserRole.SLACK_USER: AccountType.BOT,\n UserRole.EXT_PERM_USER: AccountType.EXT_PERM_USER,\n}\n\n\ndef update_user_role(\n user: User,\n new_role: UserRole,\n db_session: Session,\n) -> None:\n \"\"\"Update a user's role in the database.\n Dual-writes account_type to keep it in sync with role and\n reconciles default-group membership (Admin / Basic).\"\"\"\n old_role = user.role\n user.role = new_role\n # Note: setting account_type to BOT or EXT_PERM_USER causes\n # assign_user_to_default_groups__no_commit to early-return, which is\n # intentional — these account types should not be in default groups.\n if new_role in _ROLE_TO_ACCOUNT_TYPE:\n user.account_type = _ROLE_TO_ACCOUNT_TYPE[new_role]\n elif user.account_type in (AccountType.BOT, AccountType.EXT_PERM_USER):\n # Upgrading from a non-web-login account type to a web role\n user.account_type = AccountType.STANDARD\n\n # Reconcile default-group membership when the role changes.\n if old_role != new_role:\n # Remove from all default groups first.\n db_session.execute(\n delete(User__UserGroup).where(\n User__UserGroup.user_id == user.id,\n User__UserGroup.user_group_id.in_(\n select(UserGroup.id).where(UserGroup.is_default.is_(True))\n ),\n )\n )\n\n # Re-assign to the correct default group (skip for LIMITED).\n if new_role != UserRole.LIMITED:\n assign_user_to_default_groups__no_commit(\n db_session,\n user,\n is_admin=(new_role == UserRole.ADMIN),\n )\n\n recompute_user_permissions__no_commit(user.id, db_session)\n\n db_session.commit()\n\n\ndef deactivate_user(\n user: User,\n db_session: Session,\n) -> None:\n \"\"\"Deactivate a user by setting is_active to False.\"\"\"\n user.is_active = False\n db_session.add(user)\n db_session.commit()\n\n\ndef activate_user(\n user: User,\n db_session: Session,\n) -> None:\n \"\"\"Activate a user by setting is_active to True.\n\n Also reconciles default-group membership — the user may have been\n created while inactive or deactivated before the backfill migration.\n \"\"\"\n user.is_active = True\n if user.role != UserRole.LIMITED:\n assign_user_to_default_groups__no_commit(\n db_session, user, is_admin=(user.role == UserRole.ADMIN)\n )\n db_session.add(user)\n db_session.commit()\n\n\ndef get_latest_access_token_for_user(\n user_id: UUID,\n db_session: Session,\n) -> AccessToken | None:\n \"\"\"Get the most recent access token for a user.\"\"\"\n try:\n result = db_session.execute(\n select(AccessToken)\n .where(AccessToken.user_id == user_id) # type: ignore\n .order_by(desc(Column(\"created_at\")))\n .limit(1)\n )\n return result.scalar_one_or_none()\n except Exception as e:\n logger.error(f\"Error fetching AccessToken: {e}\")\n return None\n\n\ndef update_user_temperature_override_enabled(\n user_id: UUID,\n temperature_override_enabled: bool,\n db_session: Session,\n) -> None:\n \"\"\"Update user's temperature override enabled setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(temperature_override_enabled=temperature_override_enabled)\n )\n db_session.commit()\n\n\ndef update_user_shortcut_enabled(\n user_id: UUID,\n shortcut_enabled: bool,\n db_session: Session,\n) -> None:\n \"\"\"Update user's shortcut enabled setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(shortcut_enabled=shortcut_enabled)\n )\n db_session.commit()\n\n\ndef update_user_auto_scroll(\n user_id: UUID,\n auto_scroll: bool | None,\n db_session: Session,\n) -> None:\n \"\"\"Update user's auto scroll setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(auto_scroll=auto_scroll)\n )\n db_session.commit()\n\n\ndef update_user_default_model(\n user_id: UUID,\n default_model: str | None,\n db_session: Session,\n) -> None:\n \"\"\"Update user's default model setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(default_model=default_model)\n )\n db_session.commit()\n\n\ndef update_user_theme_preference(\n user_id: UUID,\n theme_preference: ThemePreference,\n db_session: Session,\n) -> None:\n \"\"\"Update user's theme preference setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(theme_preference=theme_preference)\n )\n db_session.commit()\n\n\ndef update_user_chat_background(\n user_id: UUID,\n chat_background: str | None,\n db_session: Session,\n) -> None:\n \"\"\"Update user's chat background setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(chat_background=chat_background)\n )\n db_session.commit()\n\n\ndef update_user_default_app_mode(\n user_id: UUID,\n default_app_mode: DefaultAppMode,\n db_session: Session,\n) -> None:\n \"\"\"Update user's default app mode setting.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(default_app_mode=default_app_mode)\n )\n db_session.commit()\n\n\ndef update_user_personalization(\n user_id: UUID,\n *,\n personal_name: str | None,\n personal_role: str | None,\n use_memories: bool,\n enable_memory_tool: bool,\n memories: list[MemoryItem],\n user_preferences: str | None,\n db_session: Session,\n) -> None:\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(\n personal_name=personal_name,\n personal_role=personal_role,\n use_memories=use_memories,\n enable_memory_tool=enable_memory_tool,\n user_preferences=user_preferences,\n )\n )\n\n # ID-based upsert: use real DB IDs from the frontend to match memories.\n incoming_ids = {m.id for m in memories if m.id is not None}\n\n # Delete existing rows not in the incoming set (scoped to user_id)\n existing_memories = list(\n db_session.scalars(select(Memory).where(Memory.user_id == user_id)).all()\n )\n existing_ids = {mem.id for mem in existing_memories}\n ids_to_delete = existing_ids - incoming_ids\n if ids_to_delete:\n db_session.execute(\n delete(Memory).where(\n Memory.id.in_(ids_to_delete),\n Memory.user_id == user_id,\n )\n )\n\n # Update existing rows whose IDs match\n existing_by_id = {mem.id: mem for mem in existing_memories}\n for item in memories:\n if item.id is not None and item.id in existing_by_id:\n existing_by_id[item.id].memory_text = item.content\n\n # Create new rows for items without an ID\n new_items = [m for m in memories if m.id is None]\n if new_items:\n db_session.add_all(\n [Memory(user_id=user_id, memory_text=item.content) for item in new_items]\n )\n\n db_session.commit()\n\n\ndef get_memories_for_user(\n user_id: UUID,\n db_session: Session,\n) -> Sequence[Memory]:\n return db_session.scalars(\n select(Memory).where(Memory.user_id == user_id).order_by(Memory.id.desc())\n ).all()\n\n\ndef update_user_pinned_assistants(\n user_id: UUID,\n pinned_assistants: list[int],\n db_session: Session,\n) -> None:\n \"\"\"Update user's pinned assistants list.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(pinned_assistants=pinned_assistants)\n )\n db_session.commit()\n\n\ndef update_user_assistant_visibility(\n user_id: UUID,\n hidden_assistants: list[int] | None,\n visible_assistants: list[int] | None,\n chosen_assistants: list[int] | None,\n db_session: Session,\n) -> None:\n \"\"\"Update user's assistant visibility settings.\"\"\"\n db_session.execute(\n update(User)\n .where(User.id == user_id) # type: ignore\n .values(\n hidden_assistants=hidden_assistants,\n visible_assistants=visible_assistants,\n chosen_assistants=chosen_assistants,\n )\n )\n db_session.commit()\n\n\ndef get_all_user_assistant_specific_configs(\n user_id: UUID,\n db_session: Session,\n) -> Sequence[Assistant__UserSpecificConfig]:\n \"\"\"Get the full user assistant specific config for a specific assistant and user.\"\"\"\n return db_session.scalars(\n select(Assistant__UserSpecificConfig).where(\n Assistant__UserSpecificConfig.user_id == user_id\n )\n ).all()\n\n\ndef update_assistant_preferences(\n assistant_id: int,\n user_id: UUID,\n new_assistant_preference: UserSpecificAssistantPreference,\n db_session: Session,\n) -> None:\n \"\"\"Update the disabled tools for a specific assistant for a specific user.\"\"\"\n # First check if a config already exists\n result = db_session.execute(\n select(Assistant__UserSpecificConfig)\n .where(Assistant__UserSpecificConfig.assistant_id == assistant_id)\n .where(Assistant__UserSpecificConfig.user_id == user_id)\n )\n config = result.scalar_one_or_none()\n\n if config:\n # Update existing config\n config.disabled_tool_ids = new_assistant_preference.disabled_tool_ids\n else:\n # Create new config\n config = Assistant__UserSpecificConfig(\n assistant_id=assistant_id,\n user_id=user_id,\n disabled_tool_ids=new_assistant_preference.disabled_tool_ids,\n )\n db_session.add(config)\n\n db_session.commit()\n" }, { "path": "backend/onyx/db/users.py", "content": "from collections.abc import Sequence\nfrom typing import Any\nfrom uuid import UUID\n\nfrom fastapi import HTTPException\nfrom fastapi_users.password import PasswordHelper\nfrom sqlalchemy import case\nfrom sqlalchemy import func\nfrom sqlalchemy import select\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.sql import expression\nfrom sqlalchemy.sql.elements import ColumnElement\nfrom sqlalchemy.sql.elements import KeyedColumnElement\nfrom sqlalchemy.sql.expression import or_\n\nfrom onyx.auth.invited_users import remove_user_from_invited_users\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import ANONYMOUS_USER_EMAIL\nfrom onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN\nfrom onyx.configs.constants import NO_AUTH_PLACEHOLDER_USER_EMAIL\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import DocumentSet__User\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__User\nfrom onyx.db.models import SamlAccount\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\nlogger = setup_logger()\n\n\ndef validate_user_role_update(\n requested_role: UserRole,\n current_role: UserRole,\n current_account_type: AccountType,\n explicit_override: bool = False,\n) -> None:\n \"\"\"\n Validate that a user role update is valid.\n Assumed only admins can hit this endpoint.\n raise if:\n - requested role is a curator\n - requested role is a slack user\n - requested role is an external permissioned user\n - requested role is a limited user\n - current account type is BOT (slack user)\n - current account type is EXT_PERM_USER\n - current role is a limited user\n \"\"\"\n\n if current_account_type == AccountType.BOT:\n raise HTTPException(\n status_code=400,\n detail=\"To change a Slack User's role, they must first login to Onyx via the web app.\",\n )\n\n if current_account_type == AccountType.EXT_PERM_USER:\n raise HTTPException(\n status_code=400,\n detail=\"To change an External Permissioned User's role, they must first login to Onyx via the web app.\",\n )\n\n if current_role == UserRole.LIMITED:\n raise HTTPException(\n status_code=400,\n detail=\"To change a Limited User's role, they must first login to Onyx via the web app.\",\n )\n\n if explicit_override:\n return\n\n if requested_role == UserRole.CURATOR:\n # This shouldn't happen, but just in case\n raise HTTPException(\n status_code=400,\n detail=\"Curator role must be set via the User Group Menu\",\n )\n\n if requested_role == UserRole.LIMITED:\n # This shouldn't happen, but just in case\n raise HTTPException(\n status_code=400,\n detail=(\n \"A user cannot be set to a Limited User role. \"\n \"This role is automatically assigned to users through certain endpoints in the API.\"\n ),\n )\n\n if requested_role == UserRole.SLACK_USER:\n # This shouldn't happen, but just in case\n raise HTTPException(\n status_code=400,\n detail=(\n \"A user cannot be set to a Slack User role. \"\n \"This role is automatically assigned to users who only use Onyx via Slack.\"\n ),\n )\n\n if requested_role == UserRole.EXT_PERM_USER:\n # This shouldn't happen, but just in case\n raise HTTPException(\n status_code=400,\n detail=(\n \"A user cannot be set to an External Permissioned User role. \"\n \"This role is automatically assigned to users who have been \"\n \"pulled in to the system via an external permissions system.\"\n ),\n )\n\n\ndef get_all_users(\n db_session: Session,\n email_filter_string: str | None = None,\n include_external: bool = False,\n) -> Sequence[User]:\n \"\"\"List all users. No pagination as of now, as the # of users\n is assumed to be relatively small (<< 1 million)\"\"\"\n stmt = select(User)\n\n # Exclude system users (anonymous user, no-auth placeholder)\n stmt = stmt.where(User.email != ANONYMOUS_USER_EMAIL) # type: ignore\n stmt = stmt.where(User.email != NO_AUTH_PLACEHOLDER_USER_EMAIL) # type: ignore\n\n if not include_external:\n stmt = stmt.where(User.role != UserRole.EXT_PERM_USER)\n\n if email_filter_string is not None:\n stmt = stmt.where(User.email.ilike(f\"%{email_filter_string}%\")) # type: ignore\n\n return db_session.scalars(stmt).unique().all()\n\n\ndef _get_accepted_user_where_clause(\n email_filter_string: str | None = None,\n roles_filter: list[UserRole] = [],\n include_external: bool = False,\n is_active_filter: bool | None = None,\n) -> list[ColumnElement[bool]]:\n \"\"\"\n Generates a SQLAlchemy where clause for filtering users based on the provided parameters.\n This is used to build the filters for the function that retrieves the users for the users table in the admin panel.\n\n Parameters:\n - email_filter_string: A substring to filter user emails. Only users whose emails contain this substring will be included.\n - is_active_filter: When True, only active users will be included. When False, only inactive users will be included.\n - roles_filter: A list of user roles to filter by. Only users with roles in this list will be included.\n - include_external: If False, external permissioned users will be excluded.\n\n Returns:\n - list: A list of conditions to be used in a SQLAlchemy query to filter users.\n \"\"\"\n\n # Access table columns directly via __table__.c to get proper SQLAlchemy column types\n # This ensures type checking works correctly for SQL operations like ilike, endswith, and is_\n email_col: KeyedColumnElement[Any] = User.__table__.c.email\n is_active_col: KeyedColumnElement[Any] = User.__table__.c.is_active\n\n where_clause: list[ColumnElement[bool]] = [\n expression.not_(email_col.endswith(DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN)),\n # Exclude system users (anonymous user, no-auth placeholder)\n email_col != ANONYMOUS_USER_EMAIL,\n email_col != NO_AUTH_PLACEHOLDER_USER_EMAIL,\n ]\n\n if not include_external:\n where_clause.append(User.role != UserRole.EXT_PERM_USER)\n\n if email_filter_string is not None:\n personal_name_col: KeyedColumnElement[Any] = User.__table__.c.personal_name\n where_clause.append(\n or_(\n email_col.ilike(f\"%{email_filter_string}%\"),\n personal_name_col.ilike(f\"%{email_filter_string}%\"),\n )\n )\n\n if roles_filter:\n where_clause.append(User.role.in_(roles_filter))\n\n if is_active_filter is not None:\n where_clause.append(is_active_col.is_(is_active_filter))\n\n return where_clause\n\n\ndef get_all_accepted_users(\n db_session: Session,\n include_external: bool = False,\n) -> Sequence[User]:\n \"\"\"Returns all accepted users without pagination.\n Uses the same filtering as the paginated endpoint but without\n search, role, or active filters.\"\"\"\n stmt = select(User)\n where_clause = _get_accepted_user_where_clause(\n include_external=include_external,\n )\n stmt = stmt.where(*where_clause).order_by(User.email)\n return db_session.scalars(stmt).unique().all()\n\n\ndef get_page_of_filtered_users(\n db_session: Session,\n page_size: int,\n page_num: int,\n email_filter_string: str | None = None,\n is_active_filter: bool | None = None,\n roles_filter: list[UserRole] = [],\n include_external: bool = False,\n) -> Sequence[User]:\n users_stmt = select(User)\n\n where_clause = _get_accepted_user_where_clause(\n email_filter_string=email_filter_string,\n roles_filter=roles_filter,\n include_external=include_external,\n is_active_filter=is_active_filter,\n )\n # Apply pagination\n users_stmt = users_stmt.offset((page_num) * page_size).limit(page_size)\n # Apply filtering\n users_stmt = users_stmt.where(*where_clause)\n\n return db_session.scalars(users_stmt).unique().all()\n\n\ndef get_total_filtered_users_count(\n db_session: Session,\n email_filter_string: str | None = None,\n is_active_filter: bool | None = None,\n roles_filter: list[UserRole] = [],\n include_external: bool = False,\n) -> int:\n where_clause = _get_accepted_user_where_clause(\n email_filter_string=email_filter_string,\n roles_filter=roles_filter,\n include_external=include_external,\n is_active_filter=is_active_filter,\n )\n total_count_stmt = select(func.count()).select_from(User)\n # Apply filtering\n total_count_stmt = total_count_stmt.where(*where_clause)\n\n return db_session.scalar(total_count_stmt) or 0\n\n\ndef get_user_counts_by_role_and_status(\n db_session: Session,\n) -> dict[str, dict[str, int]]:\n \"\"\"Returns user counts grouped by role and by active/inactive status.\n\n Excludes API key users, anonymous users, and no-auth placeholder users.\n Uses a single query with conditional aggregation.\n \"\"\"\n base_where = _get_accepted_user_where_clause()\n role_col = User.__table__.c.role\n is_active_col = User.__table__.c.is_active\n\n stmt = (\n select(\n role_col,\n func.count().label(\"total\"),\n func.sum(case((is_active_col.is_(True), 1), else_=0)).label(\"active\"),\n func.sum(case((is_active_col.is_(False), 1), else_=0)).label(\"inactive\"),\n )\n .where(*base_where)\n .group_by(role_col)\n )\n\n role_counts: dict[str, int] = {}\n status_counts: dict[str, int] = {\"active\": 0, \"inactive\": 0}\n\n for role_val, total, active, inactive in db_session.execute(stmt).all():\n key = role_val.value if hasattr(role_val, \"value\") else str(role_val)\n role_counts[key] = total\n status_counts[\"active\"] += active or 0\n status_counts[\"inactive\"] += inactive or 0\n\n return {\"role_counts\": role_counts, \"status_counts\": status_counts}\n\n\ndef get_user_by_email(email: str, db_session: Session) -> User | None:\n user = (\n db_session.query(User)\n .filter(func.lower(User.email) == func.lower(email))\n .first()\n )\n return user\n\n\ndef fetch_user_by_id(db_session: Session, user_id: UUID) -> User | None:\n return db_session.query(User).filter(User.id == user_id).first() # type: ignore\n\n\ndef _generate_slack_user(email: str) -> User:\n fastapi_users_pw_helper = PasswordHelper()\n password = fastapi_users_pw_helper.generate()\n hashed_pass = fastapi_users_pw_helper.hash(password)\n return User(\n email=email,\n hashed_password=hashed_pass,\n role=UserRole.SLACK_USER,\n account_type=AccountType.BOT,\n )\n\n\ndef add_slack_user_if_not_exists(db_session: Session, email: str) -> User:\n email = email.lower()\n user = get_user_by_email(email, db_session)\n if user is not None:\n # If the user is an external permissioned user, we update it to a slack user\n if user.account_type == AccountType.EXT_PERM_USER:\n user.role = UserRole.SLACK_USER\n user.account_type = AccountType.BOT\n db_session.commit()\n return user\n\n user = _generate_slack_user(email=email)\n db_session.add(user)\n db_session.commit()\n return user\n\n\ndef _get_users_by_emails(\n db_session: Session, lower_emails: list[str]\n) -> tuple[list[User], list[str]]:\n \"\"\"given a list of lowercase emails,\n returns a list[User] of Users whose emails match and a list[str]\n the missing emails that had no User\"\"\"\n stmt = select(User).filter(func.lower(User.email).in_(lower_emails))\n found_users = list(db_session.scalars(stmt).unique().all()) # Convert to list\n\n # Extract found emails and convert to lowercase to avoid case sensitivity issues\n found_users_emails = [user.email.lower() for user in found_users]\n\n # Separate emails for users that were not found\n missing_user_emails = [\n email for email in lower_emails if email not in found_users_emails\n ]\n return found_users, missing_user_emails\n\n\ndef _generate_ext_permissioned_user(email: str) -> User:\n fastapi_users_pw_helper = PasswordHelper()\n password = fastapi_users_pw_helper.generate()\n hashed_pass = fastapi_users_pw_helper.hash(password)\n return User(\n email=email,\n hashed_password=hashed_pass,\n role=UserRole.EXT_PERM_USER,\n account_type=AccountType.EXT_PERM_USER,\n )\n\n\ndef batch_add_ext_perm_user_if_not_exists(\n db_session: Session, emails: list[str], continue_on_error: bool = False\n) -> list[User]:\n lower_emails = [email.lower() for email in emails]\n found_users, missing_lower_emails = _get_users_by_emails(db_session, lower_emails)\n\n # Use savepoints (begin_nested) so that a failed insert only rolls back\n # that single user, not the entire transaction. A plain rollback() would\n # discard all previously flushed users in the same transaction.\n # We also avoid add_all() because SQLAlchemy 2.0's insertmanyvalues\n # batch path hits a UUID sentinel mismatch with server_default columns.\n for email in missing_lower_emails:\n user = _generate_ext_permissioned_user(email=email)\n savepoint = db_session.begin_nested()\n try:\n db_session.add(user)\n savepoint.commit()\n except IntegrityError:\n savepoint.rollback()\n if not continue_on_error:\n raise\n\n db_session.commit()\n # Fetch all users again to ensure we have the most up-to-date list\n all_users, _ = _get_users_by_emails(db_session, lower_emails)\n return all_users\n\n\ndef assign_user_to_default_groups__no_commit(\n db_session: Session,\n user: User,\n is_admin: bool = False,\n) -> None:\n \"\"\"Assign a newly created user to the appropriate default group.\n\n Does NOT commit — callers must commit the session themselves so that\n group assignment can be part of the same transaction as user creation.\n\n Args:\n is_admin: If True, assign to Admin default group; otherwise Basic.\n Callers determine this from their own context (e.g. user_count,\n admin email list, explicit choice). Defaults to False (Basic).\n \"\"\"\n if user.account_type in (\n AccountType.BOT,\n AccountType.EXT_PERM_USER,\n AccountType.ANONYMOUS,\n ):\n return\n\n target_group_name = \"Admin\" if is_admin else \"Basic\"\n\n default_group = (\n db_session.query(UserGroup)\n .filter(\n UserGroup.name == target_group_name,\n UserGroup.is_default.is_(True),\n )\n .first()\n )\n\n if default_group is None:\n raise RuntimeError(\n f\"Default group '{target_group_name}' not found. \"\n f\"Cannot assign user {user.email} to a group. \"\n f\"Ensure the seed_default_groups migration has run.\"\n )\n\n # Check if the user is already in the group\n existing = (\n db_session.query(User__UserGroup)\n .filter(\n User__UserGroup.user_id == user.id,\n User__UserGroup.user_group_id == default_group.id,\n )\n .first()\n )\n if existing is not None:\n return\n\n savepoint = db_session.begin_nested()\n try:\n db_session.add(\n User__UserGroup(\n user_id=user.id,\n user_group_id=default_group.id,\n )\n )\n db_session.flush()\n except IntegrityError:\n # Race condition: another transaction inserted this membership\n # between our SELECT and INSERT. The savepoint isolates the failure\n # so the outer transaction (user creation) stays intact.\n savepoint.rollback()\n return\n\n from onyx.db.permissions import recompute_user_permissions__no_commit\n\n recompute_user_permissions__no_commit(user.id, db_session)\n\n logger.info(f\"Assigned user {user.email} to default group '{default_group.name}'\")\n\n\ndef delete_user_from_db(\n user_to_delete: User,\n db_session: Session,\n) -> None:\n for oauth_account in user_to_delete.oauth_accounts:\n db_session.delete(oauth_account)\n\n fetch_ee_implementation_or_noop(\n \"onyx.db.external_perm\",\n \"delete_user__ext_group_for_user__no_commit\",\n )(\n db_session=db_session,\n user_id=user_to_delete.id,\n )\n db_session.query(SamlAccount).filter(\n SamlAccount.user_id == user_to_delete.id\n ).delete()\n # Null out ownership on document sets and personas so they're\n # preserved for other users instead of being cascade-deleted\n db_session.query(DocumentSet).filter(\n DocumentSet.user_id == user_to_delete.id\n ).update({DocumentSet.user_id: None})\n db_session.query(Persona).filter(Persona.user_id == user_to_delete.id).update(\n {Persona.user_id: None}\n )\n\n db_session.query(DocumentSet__User).filter(\n DocumentSet__User.user_id == user_to_delete.id\n ).delete()\n db_session.query(Persona__User).filter(\n Persona__User.user_id == user_to_delete.id\n ).delete()\n db_session.query(User__UserGroup).filter(\n User__UserGroup.user_id == user_to_delete.id\n ).delete()\n db_session.delete(user_to_delete)\n db_session.commit()\n\n # NOTE: edge case may exist with race conditions\n # with this `invited user` scheme generally.\n remove_user_from_invited_users(user_to_delete.email)\n\n\ndef batch_get_user_groups(\n db_session: Session,\n user_ids: list[UUID],\n include_default: bool = False,\n) -> dict[UUID, list[tuple[int, str]]]:\n \"\"\"Fetch group memberships for a batch of users in a single query.\n Returns a mapping of user_id -> list of (group_id, group_name) tuples.\"\"\"\n if not user_ids:\n return {}\n\n stmt = (\n select(\n User__UserGroup.user_id,\n UserGroup.id,\n UserGroup.name,\n )\n .join(UserGroup, UserGroup.id == User__UserGroup.user_group_id)\n .where(User__UserGroup.user_id.in_(user_ids))\n )\n if not include_default:\n stmt = stmt.where(UserGroup.is_default == False) # noqa: E712\n\n rows = db_session.execute(stmt).all()\n\n result: dict[UUID, list[tuple[int, str]]] = {uid: [] for uid in user_ids}\n for user_id, group_id, group_name in rows:\n result[user_id].append((group_id, group_name))\n return result\n" }, { "path": "backend/onyx/db/utils.py", "content": "from enum import Enum\nfrom typing import Any\n\nfrom psycopg2 import errorcodes\nfrom psycopg2 import OperationalError\nfrom pydantic import BaseModel\nfrom sqlalchemy import inspect\n\nfrom onyx.db.models import Base\n\n\ndef model_to_dict(model: Base) -> dict[str, Any]:\n return {c.key: getattr(model, c.key) for c in inspect(model).mapper.column_attrs} # type: ignore\n\n\nRETRYABLE_PG_CODES = {\n errorcodes.SERIALIZATION_FAILURE, # '40001'\n errorcodes.DEADLOCK_DETECTED, # '40P01'\n errorcodes.CONNECTION_EXCEPTION, # '08000'\n errorcodes.CONNECTION_DOES_NOT_EXIST, # '08003'\n errorcodes.CONNECTION_FAILURE, # '08006'\n errorcodes.TRANSACTION_ROLLBACK, # '40000'\n}\n\n\ndef is_retryable_sqlalchemy_error(exc: BaseException) -> bool:\n \"\"\"Helper function for use with tenacity's retry_if_exception as the callback\"\"\"\n if isinstance(exc, OperationalError):\n pgcode = getattr(getattr(exc, \"orig\", None), \"pgcode\", None)\n return pgcode in RETRYABLE_PG_CODES\n return False\n\n\nclass DocumentRow(BaseModel):\n id: str\n doc_metadata: dict[str, Any]\n external_user_group_ids: list[str]\n\n\nclass SortOrder(str, Enum):\n ASC = \"asc\"\n DESC = \"desc\"\n\n\nclass DiscordChannelView(BaseModel):\n channel_id: int\n channel_name: str\n channel_type: str = \"text\" # text, forum\n is_private: bool = False # True if @everyone cannot view the channel\n" }, { "path": "backend/onyx/db/voice.py", "content": "from typing import Any\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import User\nfrom onyx.db.models import VoiceProvider\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\n\nMIN_VOICE_PLAYBACK_SPEED = 0.5\nMAX_VOICE_PLAYBACK_SPEED = 2.0\n\n\ndef fetch_voice_providers(db_session: Session) -> list[VoiceProvider]:\n \"\"\"Fetch all voice providers.\"\"\"\n return list(\n db_session.scalars(select(VoiceProvider).order_by(VoiceProvider.name)).all()\n )\n\n\ndef fetch_voice_provider_by_id(\n db_session: Session, provider_id: int\n) -> VoiceProvider | None:\n \"\"\"Fetch a voice provider by ID.\"\"\"\n return db_session.scalar(\n select(VoiceProvider).where(VoiceProvider.id == provider_id)\n )\n\n\ndef fetch_default_stt_provider(db_session: Session) -> VoiceProvider | None:\n \"\"\"Fetch the default STT provider.\"\"\"\n return db_session.scalar(\n select(VoiceProvider).where(VoiceProvider.is_default_stt.is_(True))\n )\n\n\ndef fetch_default_tts_provider(db_session: Session) -> VoiceProvider | None:\n \"\"\"Fetch the default TTS provider.\"\"\"\n return db_session.scalar(\n select(VoiceProvider).where(VoiceProvider.is_default_tts.is_(True))\n )\n\n\ndef fetch_voice_provider_by_type(\n db_session: Session, provider_type: str\n) -> VoiceProvider | None:\n \"\"\"Fetch a voice provider by type.\"\"\"\n return db_session.scalar(\n select(VoiceProvider).where(VoiceProvider.provider_type == provider_type)\n )\n\n\ndef upsert_voice_provider(\n *,\n db_session: Session,\n provider_id: int | None,\n name: str,\n provider_type: str,\n api_key: str | None,\n api_key_changed: bool,\n api_base: str | None = None,\n custom_config: dict[str, Any] | None = None,\n stt_model: str | None = None,\n tts_model: str | None = None,\n default_voice: str | None = None,\n activate_stt: bool = False,\n activate_tts: bool = False,\n) -> VoiceProvider:\n \"\"\"Create or update a voice provider.\"\"\"\n provider: VoiceProvider | None = None\n\n if provider_id is not None:\n provider = fetch_voice_provider_by_id(db_session, provider_id)\n if provider is None:\n raise OnyxError(\n OnyxErrorCode.NOT_FOUND,\n f\"No voice provider with id {provider_id} exists.\",\n )\n else:\n provider = VoiceProvider()\n db_session.add(provider)\n\n # Apply updates\n provider.name = name\n provider.provider_type = provider_type\n provider.api_base = api_base\n provider.custom_config = custom_config\n provider.stt_model = stt_model\n provider.tts_model = tts_model\n provider.default_voice = default_voice\n\n # Only update API key if explicitly changed or if provider has no key\n if api_key_changed or provider.api_key is None:\n provider.api_key = api_key # type: ignore[assignment]\n\n db_session.flush()\n\n if activate_stt:\n set_default_stt_provider(db_session=db_session, provider_id=provider.id)\n if activate_tts:\n set_default_tts_provider(db_session=db_session, provider_id=provider.id)\n\n db_session.refresh(provider)\n return provider\n\n\ndef delete_voice_provider(db_session: Session, provider_id: int) -> None:\n \"\"\"Delete a voice provider by ID.\"\"\"\n provider = fetch_voice_provider_by_id(db_session, provider_id)\n if provider:\n db_session.delete(provider)\n db_session.flush()\n\n\ndef set_default_stt_provider(*, db_session: Session, provider_id: int) -> VoiceProvider:\n \"\"\"Set a voice provider as the default STT provider.\"\"\"\n provider = fetch_voice_provider_by_id(db_session, provider_id)\n if provider is None:\n raise OnyxError(\n OnyxErrorCode.NOT_FOUND,\n f\"No voice provider with id {provider_id} exists.\",\n )\n\n # Deactivate all other STT providers\n db_session.execute(\n update(VoiceProvider)\n .where(\n VoiceProvider.is_default_stt.is_(True),\n VoiceProvider.id != provider_id,\n )\n .values(is_default_stt=False)\n )\n\n # Activate this provider\n provider.is_default_stt = True\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef set_default_tts_provider(\n *, db_session: Session, provider_id: int, tts_model: str | None = None\n) -> VoiceProvider:\n \"\"\"Set a voice provider as the default TTS provider.\"\"\"\n provider = fetch_voice_provider_by_id(db_session, provider_id)\n if provider is None:\n raise OnyxError(\n OnyxErrorCode.NOT_FOUND,\n f\"No voice provider with id {provider_id} exists.\",\n )\n\n # Deactivate all other TTS providers\n db_session.execute(\n update(VoiceProvider)\n .where(\n VoiceProvider.is_default_tts.is_(True),\n VoiceProvider.id != provider_id,\n )\n .values(is_default_tts=False)\n )\n\n # Activate this provider\n provider.is_default_tts = True\n\n # Update the TTS model if specified\n if tts_model is not None:\n provider.tts_model = tts_model\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef deactivate_stt_provider(*, db_session: Session, provider_id: int) -> VoiceProvider:\n \"\"\"Remove the default STT status from a voice provider.\"\"\"\n provider = fetch_voice_provider_by_id(db_session, provider_id)\n if provider is None:\n raise OnyxError(\n OnyxErrorCode.NOT_FOUND,\n f\"No voice provider with id {provider_id} exists.\",\n )\n\n provider.is_default_stt = False\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef deactivate_tts_provider(*, db_session: Session, provider_id: int) -> VoiceProvider:\n \"\"\"Remove the default TTS status from a voice provider.\"\"\"\n provider = fetch_voice_provider_by_id(db_session, provider_id)\n if provider is None:\n raise OnyxError(\n OnyxErrorCode.NOT_FOUND,\n f\"No voice provider with id {provider_id} exists.\",\n )\n\n provider.is_default_tts = False\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\n# User voice preferences\n\n\ndef update_user_voice_settings(\n db_session: Session,\n user_id: UUID,\n auto_send: bool | None = None,\n auto_playback: bool | None = None,\n playback_speed: float | None = None,\n) -> None:\n \"\"\"Update user's voice settings.\n\n For all fields, None means \"don't update this field\".\n \"\"\"\n values: dict[str, bool | float] = {}\n\n if auto_send is not None:\n values[\"voice_auto_send\"] = auto_send\n if auto_playback is not None:\n values[\"voice_auto_playback\"] = auto_playback\n if playback_speed is not None:\n values[\"voice_playback_speed\"] = max(\n MIN_VOICE_PLAYBACK_SPEED, min(MAX_VOICE_PLAYBACK_SPEED, playback_speed)\n )\n\n if values:\n db_session.execute(update(User).where(User.id == user_id).values(**values)) # type: ignore[arg-type]\n db_session.flush()\n" }, { "path": "backend/onyx/db/web_search.py", "content": "from __future__ import annotations\n\nfrom sqlalchemy import select\nfrom sqlalchemy import update\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import InternetContentProvider\nfrom onyx.db.models import InternetSearchProvider\nfrom onyx.tools.tool_implementations.web_search.models import WebContentProviderConfig\nfrom shared_configs.enums import WebContentProviderType\nfrom shared_configs.enums import WebSearchProviderType\n\n\ndef fetch_web_search_providers(db_session: Session) -> list[InternetSearchProvider]:\n stmt = select(InternetSearchProvider).order_by(InternetSearchProvider.id.asc())\n return list(db_session.scalars(stmt).all())\n\n\ndef fetch_web_content_providers(db_session: Session) -> list[InternetContentProvider]:\n stmt = select(InternetContentProvider).order_by(InternetContentProvider.id.asc())\n return list(db_session.scalars(stmt).all())\n\n\ndef fetch_active_web_search_provider(\n db_session: Session,\n) -> InternetSearchProvider | None:\n stmt = select(InternetSearchProvider).where(\n InternetSearchProvider.is_active.is_(True)\n )\n return db_session.scalars(stmt).first()\n\n\ndef fetch_web_search_provider_by_id(\n provider_id: int, db_session: Session\n) -> InternetSearchProvider | None:\n return db_session.get(InternetSearchProvider, provider_id)\n\n\ndef fetch_web_search_provider_by_name(\n name: str, db_session: Session\n) -> InternetSearchProvider | None:\n stmt = select(InternetSearchProvider).where(InternetSearchProvider.name.ilike(name))\n return db_session.scalars(stmt).first()\n\n\ndef fetch_web_search_provider_by_type(\n provider_type: WebSearchProviderType, db_session: Session\n) -> InternetSearchProvider | None:\n stmt = select(InternetSearchProvider).where(\n InternetSearchProvider.provider_type == provider_type.value\n )\n return db_session.scalars(stmt).first()\n\n\ndef _ensure_unique_search_name(\n name: str, provider_id: int | None, db_session: Session\n) -> None:\n existing = fetch_web_search_provider_by_name(name=name, db_session=db_session)\n if existing and existing.id != provider_id:\n raise ValueError(f\"A web search provider named '{name}' already exists.\")\n\n\ndef _apply_search_provider_updates(\n provider: InternetSearchProvider,\n *,\n name: str,\n provider_type: WebSearchProviderType,\n api_key: str | None,\n api_key_changed: bool,\n config: dict[str, str] | None,\n) -> None:\n provider.name = name\n provider.provider_type = provider_type.value\n provider.config = config\n if api_key_changed or provider.api_key is None:\n # EncryptedString accepts str for writes, returns SensitiveValue for reads\n provider.api_key = api_key # type: ignore[assignment]\n\n\ndef upsert_web_search_provider(\n *,\n provider_id: int | None,\n name: str,\n provider_type: WebSearchProviderType,\n api_key: str | None,\n api_key_changed: bool,\n config: dict[str, str] | None,\n activate: bool,\n db_session: Session,\n) -> InternetSearchProvider:\n _ensure_unique_search_name(\n name=name, provider_id=provider_id, db_session=db_session\n )\n\n provider: InternetSearchProvider | None = None\n if provider_id is not None:\n provider = fetch_web_search_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web search provider with id {provider_id} exists.\")\n else:\n provider = InternetSearchProvider()\n db_session.add(provider)\n\n _apply_search_provider_updates(\n provider,\n name=name,\n provider_type=provider_type,\n api_key=api_key,\n api_key_changed=api_key_changed,\n config=config,\n )\n\n db_session.flush()\n\n if activate:\n set_active_web_search_provider(provider_id=provider.id, db_session=db_session)\n\n db_session.refresh(provider)\n return provider\n\n\ndef set_active_web_search_provider(\n *, provider_id: int | None, db_session: Session\n) -> InternetSearchProvider:\n if provider_id is None:\n raise ValueError(\"Cannot activate a provider without an id.\")\n\n provider = fetch_web_search_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web search provider with id {provider_id} exists.\")\n\n db_session.execute(\n update(InternetSearchProvider)\n .where(\n InternetSearchProvider.is_active.is_(True),\n InternetSearchProvider.id != provider_id,\n )\n .values(is_active=False)\n )\n provider.is_active = True\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef deactivate_web_search_provider(\n *, provider_id: int | None, db_session: Session\n) -> InternetSearchProvider:\n if provider_id is None:\n raise ValueError(\"Cannot deactivate a provider without an id.\")\n\n provider = fetch_web_search_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web search provider with id {provider_id} exists.\")\n\n provider.is_active = False\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef delete_web_search_provider(provider_id: int, db_session: Session) -> None:\n provider = fetch_web_search_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web search provider with id {provider_id} exists.\")\n\n db_session.delete(provider)\n db_session.flush()\n\n db_session.commit()\n\n\n# Content provider helpers\n\n\ndef fetch_active_web_content_provider(\n db_session: Session,\n) -> InternetContentProvider | None:\n stmt = select(InternetContentProvider).where(\n InternetContentProvider.is_active.is_(True)\n )\n return db_session.scalars(stmt).first()\n\n\ndef fetch_web_content_provider_by_id(\n provider_id: int, db_session: Session\n) -> InternetContentProvider | None:\n return db_session.get(InternetContentProvider, provider_id)\n\n\ndef fetch_web_content_provider_by_name(\n name: str, db_session: Session\n) -> InternetContentProvider | None:\n stmt = select(InternetContentProvider).where(\n InternetContentProvider.name.ilike(name)\n )\n return db_session.scalars(stmt).first()\n\n\ndef fetch_web_content_provider_by_type(\n provider_type: WebContentProviderType, db_session: Session\n) -> InternetContentProvider | None:\n stmt = select(InternetContentProvider).where(\n InternetContentProvider.provider_type == provider_type.value\n )\n return db_session.scalars(stmt).first()\n\n\ndef _ensure_unique_content_name(\n name: str, provider_id: int | None, db_session: Session\n) -> None:\n existing = fetch_web_content_provider_by_name(name=name, db_session=db_session)\n if existing and existing.id != provider_id:\n raise ValueError(f\"A web content provider named '{name}' already exists.\")\n\n\ndef _apply_content_provider_updates(\n provider: InternetContentProvider,\n *,\n name: str,\n provider_type: WebContentProviderType,\n api_key: str | None,\n api_key_changed: bool,\n config: WebContentProviderConfig | None,\n) -> None:\n provider.name = name\n provider.provider_type = provider_type.value\n provider.config = config\n if api_key_changed or provider.api_key is None:\n # EncryptedString accepts str for writes, returns SensitiveValue for reads\n provider.api_key = api_key # type: ignore[assignment]\n\n\ndef upsert_web_content_provider(\n *,\n provider_id: int | None,\n name: str,\n provider_type: WebContentProviderType,\n api_key: str | None,\n api_key_changed: bool,\n config: WebContentProviderConfig | None,\n activate: bool,\n db_session: Session,\n) -> InternetContentProvider:\n _ensure_unique_content_name(\n name=name, provider_id=provider_id, db_session=db_session\n )\n\n provider: InternetContentProvider | None = None\n if provider_id is not None:\n provider = fetch_web_content_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web content provider with id {provider_id} exists.\")\n else:\n provider = InternetContentProvider()\n db_session.add(provider)\n\n _apply_content_provider_updates(\n provider,\n name=name,\n provider_type=provider_type,\n api_key=api_key,\n api_key_changed=api_key_changed,\n config=config,\n )\n\n db_session.flush()\n\n if activate:\n set_active_web_content_provider(provider_id=provider.id, db_session=db_session)\n\n db_session.refresh(provider)\n return provider\n\n\ndef set_active_web_content_provider(\n *, provider_id: int | None, db_session: Session\n) -> InternetContentProvider:\n if provider_id is None:\n raise ValueError(\"Cannot activate a provider without an id.\")\n\n provider = fetch_web_content_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web content provider with id {provider_id} exists.\")\n\n db_session.execute(\n update(InternetContentProvider)\n .where(\n InternetContentProvider.is_active.is_(True),\n InternetContentProvider.id != provider_id,\n )\n .values(is_active=False)\n )\n provider.is_active = True\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef deactivate_web_content_provider(\n *, provider_id: int | None, db_session: Session\n) -> InternetContentProvider:\n if provider_id is None:\n raise ValueError(\"Cannot deactivate a provider without an id.\")\n\n provider = fetch_web_content_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web content provider with id {provider_id} exists.\")\n\n provider.is_active = False\n\n db_session.flush()\n db_session.refresh(provider)\n return provider\n\n\ndef delete_web_content_provider(provider_id: int, db_session: Session) -> None:\n provider = fetch_web_content_provider_by_id(provider_id, db_session)\n if provider is None:\n raise ValueError(f\"No web content provider with id {provider_id} exists.\")\n\n db_session.delete(provider)\n db_session.flush()\n\n db_session.commit()\n" }, { "path": "backend/onyx/deep_research/__init__.py", "content": "" }, { "path": "backend/onyx/deep_research/dr_loop.py", "content": "# TODO: Notes for potential extensions and future improvements:\n# 1. Allow tools that aren't search specific tools\n# 2. Use user provided custom prompts\n# 3. Save the plan for replay\n\nimport time\nfrom collections.abc import Callable\nfrom typing import cast\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.chat_state import ChatStateContainer\nfrom onyx.chat.citation_processor import CitationMapping\nfrom onyx.chat.citation_processor import DynamicCitationProcessor\nfrom onyx.chat.emitter import Emitter\nfrom onyx.chat.llm_loop import construct_message_history\nfrom onyx.chat.llm_step import run_llm_step\nfrom onyx.chat.llm_step import run_llm_step_pkt_generator\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.chat.models import FileToolMetadata\nfrom onyx.chat.models import LlmStepResult\nfrom onyx.chat.models import ToolCallSimple\nfrom onyx.configs.chat_configs import SKIP_DEEP_RESEARCH_CLARIFICATION\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.tools import get_tool_by_name\nfrom onyx.deep_research.dr_mock_tools import get_clarification_tool_definitions\nfrom onyx.deep_research.dr_mock_tools import get_orchestrator_tools\nfrom onyx.deep_research.dr_mock_tools import RESEARCH_AGENT_TOOL_NAME\nfrom onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_MESSAGE\nfrom onyx.deep_research.dr_mock_tools import THINK_TOOL_RESPONSE_TOKEN_COUNT\nfrom onyx.deep_research.utils import check_special_tool_calls\nfrom onyx.deep_research.utils import create_think_tool_token_processor\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.models import ToolChoiceOptions\nfrom onyx.llm.utils import model_is_reasoning_model\nfrom onyx.prompts.deep_research.orchestration_layer import CLARIFICATION_PROMPT\nfrom onyx.prompts.deep_research.orchestration_layer import FINAL_REPORT_PROMPT\nfrom onyx.prompts.deep_research.orchestration_layer import FIRST_CYCLE_REMINDER\nfrom onyx.prompts.deep_research.orchestration_layer import FIRST_CYCLE_REMINDER_TOKENS\nfrom onyx.prompts.deep_research.orchestration_layer import (\n INTERNAL_SEARCH_CLARIFICATION_GUIDANCE,\n)\nfrom onyx.prompts.deep_research.orchestration_layer import (\n INTERNAL_SEARCH_RESEARCH_TASK_GUIDANCE,\n)\nfrom onyx.prompts.deep_research.orchestration_layer import ORCHESTRATOR_PROMPT\nfrom onyx.prompts.deep_research.orchestration_layer import ORCHESTRATOR_PROMPT_REASONING\nfrom onyx.prompts.deep_research.orchestration_layer import RESEARCH_PLAN_PROMPT\nfrom onyx.prompts.deep_research.orchestration_layer import RESEARCH_PLAN_REMINDER\nfrom onyx.prompts.deep_research.orchestration_layer import USER_FINAL_REPORT_QUERY\nfrom onyx.prompts.prompt_utils import get_current_llm_day_time\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import DeepResearchPlanDelta\nfrom onyx.server.query_and_chat.streaming_models import DeepResearchPlanStart\nfrom onyx.server.query_and_chat.streaming_models import OverallStop\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import SectionEnd\nfrom onyx.server.query_and_chat.streaming_models import TopLevelBranching\nfrom onyx.tools.fake_tools.research_agent import run_research_agent_calls\nfrom onyx.tools.interface import Tool\nfrom onyx.tools.models import ToolCallInfo\nfrom onyx.tools.models import ToolCallKickoff\nfrom onyx.tools.tool_implementations.open_url.open_url_tool import OpenURLTool\nfrom onyx.tools.tool_implementations.search.search_tool import SearchTool\nfrom onyx.tools.tool_implementations.web_search.web_search_tool import WebSearchTool\nfrom onyx.tracing.framework.create import function_span\nfrom onyx.tracing.framework.create import trace\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nMAX_USER_MESSAGES_FOR_CONTEXT = 5\nMAX_FINAL_REPORT_TOKENS = 20000\n\n# 30 minute timeout before forcing final report generation\n# NOTE: The overall execution may be much longer still because it could run a research cycle at minute 29\n# and that runs for another nearly 30 minutes.\nDEEP_RESEARCH_FORCE_REPORT_SECONDS = 30 * 60\n\n# Might be something like (this gives a lot of leeway for change but typically the models don't do this):\n# 0. Research topics 1-3\n# 1. Think\n# 2. Research topics 4-5\n# 3. Think\n# 4. Research topics 6 + something new or different from the plan\n# 5. Think\n# 6. Research, possibly something new or different from the plan\n# 7. Think\n# 8. Generate report\nMAX_ORCHESTRATOR_CYCLES = 8\n\n# Similar but without the 4 thinking tool calls\nMAX_ORCHESTRATOR_CYCLES_REASONING = 4\n\n\ndef generate_final_report(\n history: list[ChatMessageSimple],\n research_plan: str,\n llm: LLM,\n token_counter: Callable[[str], int],\n state_container: ChatStateContainer,\n emitter: Emitter,\n turn_index: int,\n citation_mapping: CitationMapping,\n user_identity: LLMUserIdentity | None,\n saved_reasoning: str | None = None,\n pre_answer_processing_time: float | None = None,\n all_injected_file_metadata: dict[str, FileToolMetadata] | None = None,\n) -> bool:\n \"\"\"Generate the final research report.\n\n Returns:\n bool: True if reasoning occurred during report generation (turn_index was incremented),\n False otherwise.\n \"\"\"\n with function_span(\"generate_report\") as span:\n span.span_data.input = f\"history_length={len(history)}, turn_index={turn_index}\"\n final_report_prompt = FINAL_REPORT_PROMPT.format(\n current_datetime=get_current_llm_day_time(full_sentence=False),\n )\n system_prompt = ChatMessageSimple(\n message=final_report_prompt,\n token_count=token_counter(final_report_prompt),\n message_type=MessageType.SYSTEM,\n )\n final_reminder = USER_FINAL_REPORT_QUERY.format(research_plan=research_plan)\n reminder_message = ChatMessageSimple(\n message=final_reminder,\n token_count=token_counter(final_reminder),\n message_type=MessageType.USER_REMINDER,\n )\n final_report_history = construct_message_history(\n system_prompt=system_prompt,\n custom_agent_prompt=None,\n simple_chat_history=history,\n reminder_message=reminder_message,\n context_files=None,\n available_tokens=llm.config.max_input_tokens,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n\n citation_processor = DynamicCitationProcessor()\n citation_processor.update_citation_mapping(citation_mapping)\n\n # Only passing in the cited documents as the whole list would be too long\n final_documents = list(citation_processor.citation_to_doc.values())\n\n llm_step_result, has_reasoned = run_llm_step(\n emitter=emitter,\n history=final_report_history,\n tool_definitions=[],\n tool_choice=ToolChoiceOptions.NONE,\n llm=llm,\n placement=Placement(turn_index=turn_index),\n citation_processor=citation_processor,\n state_container=state_container,\n final_documents=final_documents,\n user_identity=user_identity,\n max_tokens=MAX_FINAL_REPORT_TOKENS,\n is_deep_research=True,\n pre_answer_processing_time=pre_answer_processing_time,\n timeout_override=300, # 5 minute read timeout for long report generation\n )\n\n # Save citation mapping to state_container so citations are persisted\n state_container.set_citation_mapping(citation_processor.citation_to_doc)\n\n final_report = llm_step_result.answer\n if final_report is None:\n raise ValueError(\"LLM failed to generate the final deep research report\")\n\n if saved_reasoning:\n # The reasoning we want to save with the message is more about calling this\n # generate report and why it's done. Also some models don't have reasoning\n # but we'd still want to capture the reasoning from the think_tool of theprevious turn.\n state_container.set_reasoning_tokens(saved_reasoning)\n\n span.span_data.output = final_report if final_report else None\n return has_reasoned\n\n\n@log_function_time(print_only=True)\ndef run_deep_research_llm_loop(\n emitter: Emitter,\n state_container: ChatStateContainer,\n simple_chat_history: list[ChatMessageSimple],\n tools: list[Tool],\n custom_agent_prompt: str | None, # noqa: ARG001\n llm: LLM,\n token_counter: Callable[[str], int],\n db_session: Session,\n skip_clarification: bool = False,\n user_identity: LLMUserIdentity | None = None,\n chat_session_id: str | None = None,\n all_injected_file_metadata: dict[str, FileToolMetadata] | None = None,\n) -> None:\n with trace(\n \"run_deep_research_llm_loop\",\n group_id=chat_session_id,\n metadata={\n \"tenant_id\": get_current_tenant_id(),\n \"chat_session_id\": chat_session_id,\n },\n ):\n # Here for lazy load LiteLLM\n from onyx.llm.litellm_singleton.config import initialize_litellm\n\n # An approximate limit. In extreme cases it may still fail but this should allow deep research\n # to work in most cases.\n if llm.config.max_input_tokens < 50000:\n raise RuntimeError(\n \"Cannot run Deep Research with an LLM that has less than 50,000 max input tokens\"\n )\n\n initialize_litellm()\n\n # Track processing start time for tool duration calculation\n processing_start_time = time.monotonic()\n\n available_tokens = llm.config.max_input_tokens\n\n llm_step_result: LlmStepResult | None = None\n\n # Filter tools to only allow web search, internal search, and open URL\n allowed_tool_names = {SearchTool.NAME, WebSearchTool.NAME, OpenURLTool.NAME}\n allowed_tools = [tool for tool in tools if tool.name in allowed_tool_names]\n include_internal_search_tunings = SearchTool.NAME in allowed_tool_names\n orchestrator_start_turn_index = 1\n\n #########################################################\n # CLARIFICATION STEP (optional)\n #########################################################\n internal_search_clarification_guidance = (\n INTERNAL_SEARCH_CLARIFICATION_GUIDANCE\n if include_internal_search_tunings\n else \"\"\n )\n if not SKIP_DEEP_RESEARCH_CLARIFICATION and not skip_clarification:\n with function_span(\"clarification_step\") as span:\n clarification_prompt = CLARIFICATION_PROMPT.format(\n current_datetime=get_current_llm_day_time(full_sentence=False),\n internal_search_clarification_guidance=internal_search_clarification_guidance,\n )\n system_prompt = ChatMessageSimple(\n message=clarification_prompt,\n token_count=300, # Skips the exact token count but has enough leeway\n message_type=MessageType.SYSTEM,\n )\n\n truncated_message_history = construct_message_history(\n system_prompt=system_prompt,\n custom_agent_prompt=None,\n simple_chat_history=simple_chat_history,\n reminder_message=None,\n context_files=None,\n available_tokens=available_tokens,\n last_n_user_messages=MAX_USER_MESSAGES_FOR_CONTEXT,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n\n # Calculate tool processing duration for clarification step\n # (used if the LLM emits a clarification question instead of calling tools)\n clarification_tool_duration = time.monotonic() - processing_start_time\n llm_step_result, _ = run_llm_step(\n emitter=emitter,\n history=truncated_message_history,\n tool_definitions=get_clarification_tool_definitions(),\n tool_choice=ToolChoiceOptions.AUTO,\n llm=llm,\n placement=Placement(turn_index=0),\n # No citations in this step, it should just pass through all\n # tokens directly so initialized as an empty citation processor\n citation_processor=None,\n state_container=state_container,\n final_documents=None,\n user_identity=user_identity,\n is_deep_research=True,\n pre_answer_processing_time=clarification_tool_duration,\n )\n\n if not llm_step_result.tool_calls:\n # Mark this turn as a clarification question\n state_container.set_is_clarification(True)\n span.span_data.output = \"clarification_required\"\n\n emitter.emit(\n Packet(\n placement=Placement(turn_index=0),\n obj=OverallStop(type=\"stop\"),\n )\n )\n\n # If a clarification is asked, we need to end this turn and wait on user input\n return\n\n #########################################################\n # RESEARCH PLAN STEP\n #########################################################\n with function_span(\"research_plan_step\") as span:\n system_prompt = ChatMessageSimple(\n message=RESEARCH_PLAN_PROMPT.format(\n current_datetime=get_current_llm_day_time(full_sentence=False)\n ),\n token_count=300,\n message_type=MessageType.SYSTEM,\n )\n # Note this is fine to use a USER message type here as it can just be interpretered as a\n # user's message directly to the LLM.\n reminder_message = ChatMessageSimple(\n message=RESEARCH_PLAN_REMINDER,\n token_count=100,\n message_type=MessageType.USER,\n )\n truncated_message_history = construct_message_history(\n system_prompt=system_prompt,\n custom_agent_prompt=None,\n simple_chat_history=simple_chat_history + [reminder_message],\n reminder_message=None,\n context_files=None,\n available_tokens=available_tokens,\n last_n_user_messages=MAX_USER_MESSAGES_FOR_CONTEXT + 1,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n\n research_plan_generator = run_llm_step_pkt_generator(\n history=truncated_message_history,\n tool_definitions=[],\n tool_choice=ToolChoiceOptions.NONE,\n llm=llm,\n placement=Placement(turn_index=0),\n citation_processor=None,\n state_container=state_container,\n final_documents=None,\n user_identity=user_identity,\n is_deep_research=True,\n )\n\n while True:\n try:\n packet = next(research_plan_generator)\n # Translate AgentResponseStart/Delta packets to DeepResearchPlanStart/Delta\n # The LLM response from this prompt is the research plan\n if isinstance(packet.obj, AgentResponseStart):\n emitter.emit(\n Packet(\n placement=packet.placement,\n obj=DeepResearchPlanStart(),\n )\n )\n elif isinstance(packet.obj, AgentResponseDelta):\n emitter.emit(\n Packet(\n placement=packet.placement,\n obj=DeepResearchPlanDelta(content=packet.obj.content),\n )\n )\n else:\n # Pass through other packet types (e.g., ReasoningStart, ReasoningDelta, etc.)\n emitter.emit(packet)\n except StopIteration as e:\n llm_step_result, reasoned = e.value\n emitter.emit(\n Packet(\n # Marks the last turn end which should be the plan generation\n placement=Placement(\n turn_index=1 if reasoned else 0,\n ),\n obj=SectionEnd(),\n )\n )\n if reasoned:\n orchestrator_start_turn_index += 1\n break\n llm_step_result = cast(LlmStepResult, llm_step_result)\n\n research_plan = llm_step_result.answer\n if research_plan is None:\n raise RuntimeError(\"Deep Research failed to generate a research plan\")\n span.span_data.output = research_plan if research_plan else None\n\n #########################################################\n # RESEARCH EXECUTION STEP\n #########################################################\n with function_span(\"research_execution_step\") as span:\n is_reasoning_model = model_is_reasoning_model(\n llm.config.model_name, llm.config.model_provider\n )\n\n max_orchestrator_cycles = (\n MAX_ORCHESTRATOR_CYCLES\n if not is_reasoning_model\n else MAX_ORCHESTRATOR_CYCLES_REASONING\n )\n\n orchestrator_prompt_template = (\n ORCHESTRATOR_PROMPT\n if not is_reasoning_model\n else ORCHESTRATOR_PROMPT_REASONING\n )\n\n internal_search_research_task_guidance = (\n INTERNAL_SEARCH_RESEARCH_TASK_GUIDANCE\n if include_internal_search_tunings\n else \"\"\n )\n token_count_prompt = orchestrator_prompt_template.format(\n current_datetime=get_current_llm_day_time(full_sentence=False),\n current_cycle_count=1,\n max_cycles=max_orchestrator_cycles,\n research_plan=research_plan,\n internal_search_research_task_guidance=internal_search_research_task_guidance,\n )\n orchestration_tokens = token_counter(token_count_prompt)\n\n reasoning_cycles = 0\n most_recent_reasoning: str | None = None\n citation_mapping: CitationMapping = {}\n final_turn_index: int = (\n orchestrator_start_turn_index # Track the final turn_index for stop packet\n )\n for cycle in range(max_orchestrator_cycles):\n # Check if we've exceeded the time limit or reached the last cycle\n # - if so, skip LLM and generate final report\n elapsed_seconds = time.monotonic() - processing_start_time\n timed_out = elapsed_seconds > DEEP_RESEARCH_FORCE_REPORT_SECONDS\n is_last_cycle = cycle == max_orchestrator_cycles - 1\n\n if timed_out or is_last_cycle:\n if timed_out:\n logger.info(\n f\"Deep research exceeded {DEEP_RESEARCH_FORCE_REPORT_SECONDS}s \"\n f\"(elapsed: {elapsed_seconds:.1f}s), forcing final report generation\"\n )\n report_turn_index = (\n orchestrator_start_turn_index + cycle + reasoning_cycles\n )\n report_reasoned = generate_final_report(\n history=simple_chat_history,\n research_plan=research_plan,\n llm=llm,\n token_counter=token_counter,\n state_container=state_container,\n emitter=emitter,\n turn_index=report_turn_index,\n citation_mapping=citation_mapping,\n user_identity=user_identity,\n pre_answer_processing_time=elapsed_seconds,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n final_turn_index = report_turn_index + (1 if report_reasoned else 0)\n break\n\n if cycle == 1:\n first_cycle_reminder_message = ChatMessageSimple(\n message=FIRST_CYCLE_REMINDER,\n token_count=FIRST_CYCLE_REMINDER_TOKENS,\n message_type=MessageType.USER_REMINDER,\n )\n else:\n first_cycle_reminder_message = None\n\n research_agent_calls: list[ToolCallKickoff] = []\n\n orchestrator_prompt = orchestrator_prompt_template.format(\n current_datetime=get_current_llm_day_time(full_sentence=False),\n current_cycle_count=cycle,\n max_cycles=max_orchestrator_cycles,\n research_plan=research_plan,\n internal_search_research_task_guidance=internal_search_research_task_guidance,\n )\n\n system_prompt = ChatMessageSimple(\n message=orchestrator_prompt,\n token_count=orchestration_tokens,\n message_type=MessageType.SYSTEM,\n )\n\n truncated_message_history = construct_message_history(\n system_prompt=system_prompt,\n custom_agent_prompt=None,\n simple_chat_history=simple_chat_history,\n reminder_message=first_cycle_reminder_message,\n context_files=None,\n available_tokens=available_tokens,\n last_n_user_messages=MAX_USER_MESSAGES_FOR_CONTEXT,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n\n # Use think tool processor for non-reasoning models to convert\n # think_tool calls to reasoning content\n custom_processor = (\n create_think_tool_token_processor()\n if not is_reasoning_model\n else None\n )\n\n llm_step_result, has_reasoned = run_llm_step(\n emitter=emitter,\n history=truncated_message_history,\n tool_definitions=get_orchestrator_tools(\n include_think_tool=not is_reasoning_model\n ),\n tool_choice=ToolChoiceOptions.REQUIRED,\n llm=llm,\n placement=Placement(\n turn_index=orchestrator_start_turn_index\n + cycle\n + reasoning_cycles\n ),\n # No citations in this step, it should just pass through all\n # tokens directly so initialized as an empty citation processor\n citation_processor=DynamicCitationProcessor(),\n state_container=state_container,\n final_documents=None,\n user_identity=user_identity,\n custom_token_processor=custom_processor,\n is_deep_research=True,\n # Even for the reasoning tool, this should be plenty\n # The generation here should never be very long as it's just the tool calls.\n # This prevents timeouts where the model gets into an endless loop of null or bad tokens.\n max_tokens=1024,\n )\n if has_reasoned:\n reasoning_cycles += 1\n\n tool_calls = llm_step_result.tool_calls or []\n\n if not tool_calls and cycle == 0:\n raise RuntimeError(\n \"Deep Research failed to generate any research tasks for the agents.\"\n )\n\n if not tool_calls:\n # Basically hope that this is an infrequent occurence and hopefully multiple research\n # cycles have already ran\n logger.warning(\"No tool calls found, this should not happen.\")\n report_turn_index = (\n orchestrator_start_turn_index + cycle + reasoning_cycles\n )\n report_reasoned = generate_final_report(\n history=simple_chat_history,\n research_plan=research_plan,\n llm=llm,\n token_counter=token_counter,\n state_container=state_container,\n emitter=emitter,\n turn_index=report_turn_index,\n citation_mapping=citation_mapping,\n user_identity=user_identity,\n pre_answer_processing_time=time.monotonic()\n - processing_start_time,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n final_turn_index = report_turn_index + (1 if report_reasoned else 0)\n break\n\n special_tool_calls = check_special_tool_calls(tool_calls=tool_calls)\n\n if special_tool_calls.generate_report_tool_call:\n report_turn_index = (\n special_tool_calls.generate_report_tool_call.placement.turn_index\n )\n report_reasoned = generate_final_report(\n history=simple_chat_history,\n research_plan=research_plan,\n llm=llm,\n token_counter=token_counter,\n state_container=state_container,\n emitter=emitter,\n turn_index=report_turn_index,\n citation_mapping=citation_mapping,\n user_identity=user_identity,\n saved_reasoning=most_recent_reasoning,\n pre_answer_processing_time=time.monotonic()\n - processing_start_time,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n final_turn_index = report_turn_index + (1 if report_reasoned else 0)\n break\n elif special_tool_calls.think_tool_call:\n think_tool_call = special_tool_calls.think_tool_call\n # Only process the THINK_TOOL and skip all other tool calls\n # This will not actually get saved to the db as a tool call but we'll attach it to the tool(s) called after\n # it as if it were just a reasoning model doing it. In the chat history, because it happens in 2 steps,\n # we will show it as a separate message.\n # NOTE: This does not need to increment the reasoning cycles because the custom token processor causes\n # the LLM step to handle this\n with function_span(\"think_tool\") as span:\n span.span_data.input = str(think_tool_call.tool_args)\n most_recent_reasoning = state_container.reasoning_tokens\n tool_call_message = think_tool_call.to_msg_str()\n tool_call_token_count = token_counter(tool_call_message)\n\n # Create ASSISTANT message with tool_calls (OpenAI parallel format)\n think_tool_simple = ToolCallSimple(\n tool_call_id=think_tool_call.tool_call_id,\n tool_name=think_tool_call.tool_name,\n tool_arguments=think_tool_call.tool_args,\n token_count=tool_call_token_count,\n )\n think_assistant_msg = ChatMessageSimple(\n message=\"\",\n token_count=tool_call_token_count,\n message_type=MessageType.ASSISTANT,\n tool_calls=[think_tool_simple],\n image_files=None,\n )\n simple_chat_history.append(think_assistant_msg)\n\n think_tool_response_msg = ChatMessageSimple(\n message=THINK_TOOL_RESPONSE_MESSAGE,\n token_count=THINK_TOOL_RESPONSE_TOKEN_COUNT,\n message_type=MessageType.TOOL_CALL_RESPONSE,\n tool_call_id=think_tool_call.tool_call_id,\n image_files=None,\n )\n simple_chat_history.append(think_tool_response_msg)\n span.span_data.output = THINK_TOOL_RESPONSE_MESSAGE\n continue\n else:\n for tool_call in tool_calls:\n if tool_call.tool_name != RESEARCH_AGENT_TOOL_NAME:\n logger.warning(\n f\"Unexpected tool call: {tool_call.tool_name}\"\n )\n continue\n\n research_agent_calls.append(tool_call)\n\n if not research_agent_calls:\n logger.warning(\n \"No research agent tool calls found, this should not happen.\"\n )\n report_turn_index = (\n orchestrator_start_turn_index + cycle + reasoning_cycles\n )\n report_reasoned = generate_final_report(\n history=simple_chat_history,\n research_plan=research_plan,\n llm=llm,\n token_counter=token_counter,\n state_container=state_container,\n emitter=emitter,\n turn_index=report_turn_index,\n citation_mapping=citation_mapping,\n user_identity=user_identity,\n pre_answer_processing_time=time.monotonic()\n - processing_start_time,\n all_injected_file_metadata=all_injected_file_metadata,\n )\n final_turn_index = report_turn_index + (\n 1 if report_reasoned else 0\n )\n break\n\n if len(research_agent_calls) > 1:\n emitter.emit(\n Packet(\n placement=Placement(\n turn_index=research_agent_calls[\n 0\n ].placement.turn_index\n ),\n obj=TopLevelBranching(\n num_parallel_branches=len(research_agent_calls)\n ),\n )\n )\n\n research_results = run_research_agent_calls(\n # The tool calls here contain the placement information\n research_agent_calls=research_agent_calls,\n parent_tool_call_ids=[\n tool_call.tool_call_id for tool_call in tool_calls\n ],\n tools=allowed_tools,\n emitter=emitter,\n state_container=state_container,\n llm=llm,\n is_reasoning_model=is_reasoning_model,\n token_counter=token_counter,\n citation_mapping=citation_mapping,\n user_identity=user_identity,\n )\n\n citation_mapping = research_results.citation_mapping\n\n # Build ONE ASSISTANT message with all tool calls (OpenAI parallel format)\n tool_calls_simple: list[ToolCallSimple] = []\n for current_tool_call in research_agent_calls:\n tool_call_message = current_tool_call.to_msg_str()\n tool_call_token_count = token_counter(tool_call_message)\n tool_calls_simple.append(\n ToolCallSimple(\n tool_call_id=current_tool_call.tool_call_id,\n tool_name=current_tool_call.tool_name,\n tool_arguments=current_tool_call.tool_args,\n token_count=tool_call_token_count,\n )\n )\n\n total_tool_call_tokens = sum(\n tc.token_count for tc in tool_calls_simple\n )\n assistant_with_tools = ChatMessageSimple(\n message=\"\",\n token_count=total_tool_call_tokens,\n message_type=MessageType.ASSISTANT,\n tool_calls=tool_calls_simple,\n image_files=None,\n )\n simple_chat_history.append(assistant_with_tools)\n\n # Now add TOOL_CALL_RESPONSE messages and tool call info for each result\n for tab_index, report in enumerate(\n research_results.intermediate_reports\n ):\n if report is None:\n # The LLM will not see that this research was even attempted, it may try\n # something similar again but this is not bad.\n logger.error(\n f\"Research agent call at tab_index {tab_index} failed, skipping\"\n )\n continue\n\n current_tool_call = research_agent_calls[tab_index]\n tool_call_info = ToolCallInfo(\n parent_tool_call_id=None,\n turn_index=orchestrator_start_turn_index\n + cycle\n + reasoning_cycles,\n tab_index=tab_index,\n tool_name=current_tool_call.tool_name,\n tool_call_id=current_tool_call.tool_call_id,\n tool_id=get_tool_by_name(\n tool_name=RESEARCH_AGENT_TOOL_NAME,\n db_session=db_session,\n ).id,\n reasoning_tokens=llm_step_result.reasoning\n or most_recent_reasoning,\n tool_call_arguments=current_tool_call.tool_args,\n tool_call_response=report,\n search_docs=None, # Intermediate docs are not saved/shown\n generated_images=None,\n )\n state_container.add_tool_call(tool_call_info)\n\n tool_call_response_msg = ChatMessageSimple(\n message=report,\n token_count=token_counter(report),\n message_type=MessageType.TOOL_CALL_RESPONSE,\n tool_call_id=current_tool_call.tool_call_id,\n image_files=None,\n )\n simple_chat_history.append(tool_call_response_msg)\n\n # If it reached this point, it did not call reasoning, so here we wipe it to not save it to multiple turns\n most_recent_reasoning = None\n\n emitter.emit(\n Packet(\n placement=Placement(turn_index=final_turn_index),\n obj=OverallStop(type=\"stop\"),\n )\n )\n" }, { "path": "backend/onyx/deep_research/dr_mock_tools.py", "content": "GENERATE_PLAN_TOOL_NAME = \"generate_plan\"\n\nRESEARCH_AGENT_IN_CODE_ID = \"ResearchAgent\"\nRESEARCH_AGENT_TOOL_NAME = \"research_agent\"\nRESEARCH_AGENT_TASK_KEY = \"task\"\n\nGENERATE_REPORT_TOOL_NAME = \"generate_report\"\n\nTHINK_TOOL_NAME = \"think_tool\"\n\n\n# ruff: noqa: E501, W605 start\nGENERATE_PLAN_TOOL_DESCRIPTION = {\n \"type\": \"function\",\n \"function\": {\n \"name\": GENERATE_PLAN_TOOL_NAME,\n \"description\": \"No clarification needed, generate a research plan for the user's query.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {},\n \"required\": [],\n },\n },\n}\n\n\nRESEARCH_AGENT_TOOL_DESCRIPTION = {\n \"type\": \"function\",\n \"function\": {\n \"name\": RESEARCH_AGENT_TOOL_NAME,\n \"description\": \"Conduct research on a specific topic.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n RESEARCH_AGENT_TASK_KEY: {\n \"type\": \"string\",\n \"description\": \"The research task to investigate, should be 1-2 descriptive sentences outlining the direction of investigation.\",\n }\n },\n \"required\": [RESEARCH_AGENT_TASK_KEY],\n },\n },\n}\n\n\nGENERATE_REPORT_TOOL_DESCRIPTION = {\n \"type\": \"function\",\n \"function\": {\n \"name\": GENERATE_REPORT_TOOL_NAME,\n \"description\": \"Generate the final research report from all of the findings. Should be called when all aspects of the user's query have been researched, or maximum cycles are reached.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {},\n \"required\": [],\n },\n },\n}\n\n\nTHINK_TOOL_DESCRIPTION = {\n \"type\": \"function\",\n \"function\": {\n \"name\": THINK_TOOL_NAME,\n \"description\": \"Use this for reasoning between research_agent calls and before calling generate_report. Think deeply about key results, identify knowledge gaps, and plan next steps.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"reasoning\": {\n \"type\": \"string\",\n \"description\": \"Your chain of thought reasoning, use paragraph format, no lists.\",\n }\n },\n \"required\": [\"reasoning\"],\n },\n },\n}\n\n\nRESEARCH_AGENT_THINK_TOOL_DESCRIPTION = {\n \"type\": \"function\",\n \"function\": {\n \"name\": \"think_tool\",\n \"description\": \"Use this for reasoning between research steps. Think deeply about key results, identify knowledge gaps, and plan next steps.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"reasoning\": {\n \"type\": \"string\",\n \"description\": \"Your chain of thought reasoning, can be as long as a lengthy paragraph.\",\n }\n },\n \"required\": [\"reasoning\"],\n },\n },\n}\n\n\nRESEARCH_AGENT_GENERATE_REPORT_TOOL_DESCRIPTION = {\n \"type\": \"function\",\n \"function\": {\n \"name\": \"generate_report\",\n \"description\": \"Generate the final research report from all findings. Should be called when research is complete.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {},\n \"required\": [],\n },\n },\n}\n\n\nTHINK_TOOL_RESPONSE_MESSAGE = \"Acknowledged, please continue.\"\nTHINK_TOOL_RESPONSE_TOKEN_COUNT = 10\n\n\ndef get_clarification_tool_definitions() -> list[dict]:\n return [GENERATE_PLAN_TOOL_DESCRIPTION]\n\n\ndef get_orchestrator_tools(include_think_tool: bool) -> list[dict]:\n tools = [\n RESEARCH_AGENT_TOOL_DESCRIPTION,\n GENERATE_REPORT_TOOL_DESCRIPTION,\n ]\n if include_think_tool:\n tools.append(THINK_TOOL_DESCRIPTION)\n return tools\n\n\ndef get_research_agent_additional_tool_definitions(\n include_think_tool: bool,\n) -> list[dict]:\n tools = [GENERATE_REPORT_TOOL_DESCRIPTION]\n if include_think_tool:\n tools.append(RESEARCH_AGENT_THINK_TOOL_DESCRIPTION)\n return tools\n\n\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/deep_research/models.py", "content": "from pydantic import BaseModel\n\nfrom onyx.chat.citation_processor import CitationMapping\nfrom onyx.tools.models import ToolCallKickoff\n\n\nclass SpecialToolCalls(BaseModel):\n think_tool_call: ToolCallKickoff | None = None\n generate_report_tool_call: ToolCallKickoff | None = None\n\n\nclass ResearchAgentCallResult(BaseModel):\n intermediate_report: str\n citation_mapping: CitationMapping\n\n\nclass CombinedResearchAgentCallResult(BaseModel):\n # The None is needed here to keep the mappings consistent\n # we later skip the failed research results but we need to know\n # which ones failed\n intermediate_reports: list[str | None]\n citation_mapping: CitationMapping\n" }, { "path": "backend/onyx/deep_research/utils.py", "content": "from collections.abc import Callable\nfrom typing import Any\n\nfrom pydantic import BaseModel\n\nfrom onyx.deep_research.dr_mock_tools import GENERATE_REPORT_TOOL_NAME\nfrom onyx.deep_research.dr_mock_tools import THINK_TOOL_NAME\nfrom onyx.deep_research.models import SpecialToolCalls\nfrom onyx.llm.model_response import ChatCompletionDeltaToolCall\nfrom onyx.llm.model_response import Delta\nfrom onyx.llm.model_response import FunctionCall\nfrom onyx.tools.models import ToolCallKickoff\n\n\n# JSON prefixes to detect in think_tool arguments\n# The schema is: {\"reasoning\": \"...content...\"}\nJSON_PREFIX_WITH_SPACE = '{\"reasoning\": \"'\nJSON_PREFIX_NO_SPACE = '{\"reasoning\":\"'\n\n\nclass ThinkToolProcessorState(BaseModel):\n \"\"\"State for tracking think tool processing across streaming deltas.\"\"\"\n\n think_tool_found: bool = False\n think_tool_index: int | None = None\n think_tool_id: str | None = None\n full_arguments: str = \"\" # Full accumulated arguments for final tool call\n accumulated_args: str = \"\" # Working buffer for JSON parsing\n json_prefix_stripped: bool = False\n # Buffer holds content that might be the JSON suffix \"}\n # We hold back 2 chars to avoid emitting the closing \"}\n buffer: str = \"\"\n\n\ndef _unescape_json_string(s: str) -> str:\n \"\"\"\n Unescape JSON string escape sequences.\n\n JSON strings use backslash escapes like \\\\n for newlines, \\\\t for tabs, etc.\n When we extract content from JSON by string manipulation (without json.loads),\n we need to manually decode these escape sequences.\n\n Note: We use a placeholder approach to handle escaped backslashes correctly.\n For example, \"\\\\\\\\n\" (escaped backslash + n) should become \"\\\\n\" (literal backslash + n),\n not a newline character.\n \"\"\"\n # First, protect escaped backslashes with a placeholder\n placeholder = \"\\x00ESCAPED_BACKSLASH\\x00\"\n result = s.replace(\"\\\\\\\\\", placeholder)\n\n # Now unescape common JSON escape sequences\n result = result.replace(\"\\\\n\", \"\\n\")\n result = result.replace(\"\\\\r\", \"\\r\")\n result = result.replace(\"\\\\t\", \"\\t\")\n result = result.replace('\\\\\"', '\"')\n\n # Finally, restore escaped backslashes as single backslashes\n result = result.replace(placeholder, \"\\\\\")\n\n return result\n\n\ndef _extract_reasoning_chunk(state: ThinkToolProcessorState) -> str | None:\n \"\"\"\n Extract reasoning content from accumulated arguments, stripping JSON wrapper.\n\n Returns the next chunk of reasoning to emit, or None if nothing to emit yet.\n \"\"\"\n # If we haven't found the JSON prefix yet, look for it\n if not state.json_prefix_stripped:\n # Try both prefix variants\n for prefix in [JSON_PREFIX_WITH_SPACE, JSON_PREFIX_NO_SPACE]:\n prefix_pos = state.accumulated_args.find(prefix)\n if prefix_pos != -1:\n # Found prefix - extract content after it\n content_start = prefix_pos + len(prefix)\n state.buffer = state.accumulated_args[content_start:]\n state.accumulated_args = \"\"\n state.json_prefix_stripped = True\n break\n\n if not state.json_prefix_stripped:\n # Haven't seen full prefix yet, keep accumulating\n return None\n else:\n # Already stripped prefix, add new content to buffer\n state.buffer += state.accumulated_args\n state.accumulated_args = \"\"\n\n # Hold back enough chars to avoid splitting escape sequences AND the JSON suffix \"}\n # We need at least 2 for the suffix, but we also need to ensure escape sequences\n # like \\n, \\t, \\\\, \\\" don't get split. The longest escape is \\\\ (2 chars).\n # So we hold back 3 chars to be safe: if the last char is \\, we don't want to\n # emit it without knowing what follows.\n holdback = 3\n if len(state.buffer) <= holdback:\n return None\n\n # Check if there's a trailing backslash that could be part of an escape sequence\n # If so, hold back one more character to avoid splitting the escape\n to_emit = state.buffer[:-holdback]\n remaining = state.buffer[-holdback:]\n\n # If to_emit ends with a backslash, it might be the start of an escape sequence\n # Move it to the remaining buffer to process with the next chunk\n # If to_emit ends with a backslash, it might be the start of an escape sequence\n # Move it to the remaining buffer to process with the next chunk\n if to_emit and to_emit[-1] == \"\\\\\":\n remaining = to_emit[-1] + remaining\n to_emit = to_emit[:-1]\n\n state.buffer = remaining\n\n # Unescape JSON escape sequences (e.g., \\\\n -> \\n)\n if to_emit:\n to_emit = _unescape_json_string(to_emit)\n\n return to_emit if to_emit else None\n\n\ndef create_think_tool_token_processor() -> (\n Callable[[Delta | None, Any], tuple[Delta | None, Any]]\n):\n \"\"\"\n Create a custom token processor that converts think_tool calls to reasoning content.\n\n When the think_tool is detected:\n - Tool call arguments are converted to reasoning_content (JSON wrapper stripped)\n - All other deltas (content, other tool calls) are dropped\n\n This allows non-reasoning models to emit chain-of-thought via the think_tool,\n which gets displayed as reasoning tokens in the UI.\n\n Returns:\n A function compatible with run_llm_step_pkt_generator's custom_token_processor parameter.\n The function takes (Delta, state) and returns (modified Delta | None, new state).\n \"\"\"\n\n def process_token(delta: Delta | None, state: Any) -> tuple[Delta | None, Any]:\n if state is None:\n state = ThinkToolProcessorState()\n\n # Handle flush signal (delta=None) - emit the complete tool call\n if delta is None:\n if state.think_tool_found and state.think_tool_id:\n # Return the complete think tool call\n complete_tool_call = ChatCompletionDeltaToolCall(\n id=state.think_tool_id,\n index=state.think_tool_index or 0,\n type=\"function\",\n function=FunctionCall(\n name=THINK_TOOL_NAME,\n arguments=state.full_arguments,\n ),\n )\n return Delta(tool_calls=[complete_tool_call]), state\n return None, state\n\n # Check for think tool in tool_calls\n if delta.tool_calls:\n for tool_call in delta.tool_calls:\n # Detect think tool by name\n if tool_call.function and tool_call.function.name == THINK_TOOL_NAME:\n state.think_tool_found = True\n state.think_tool_index = tool_call.index\n\n # Capture tool call id when available\n if (\n state.think_tool_found\n and tool_call.index == state.think_tool_index\n and tool_call.id\n ):\n state.think_tool_id = tool_call.id\n\n # Accumulate arguments for the think tool\n if (\n state.think_tool_found\n and tool_call.index == state.think_tool_index\n and tool_call.function\n and tool_call.function.arguments\n ):\n # Track full arguments for final tool call\n state.full_arguments += tool_call.function.arguments\n # Also accumulate for JSON parsing\n state.accumulated_args += tool_call.function.arguments\n\n # Try to extract reasoning content\n reasoning_chunk = _extract_reasoning_chunk(state)\n if reasoning_chunk:\n # Return delta with reasoning_content to trigger reasoning streaming\n return Delta(reasoning_content=reasoning_chunk), state\n\n # If think tool found, drop all other content\n if state.think_tool_found:\n return None, state\n\n # No think tool detected, pass through original delta\n return delta, state\n\n return process_token\n\n\ndef check_special_tool_calls(tool_calls: list[ToolCallKickoff]) -> SpecialToolCalls:\n think_tool_call: ToolCallKickoff | None = None\n generate_report_tool_call: ToolCallKickoff | None = None\n\n for tool_call in tool_calls:\n if tool_call.tool_name == THINK_TOOL_NAME:\n think_tool_call = tool_call\n elif tool_call.tool_name == GENERATE_REPORT_TOOL_NAME:\n generate_report_tool_call = tool_call\n\n return SpecialToolCalls(\n think_tool_call=think_tool_call,\n generate_report_tool_call=generate_report_tool_call,\n )\n" }, { "path": "backend/onyx/document_index/FILTER_SEMANTICS.md", "content": "# Vector DB Filter Semantics\n\nHow `IndexFilters` fields combine into the final query filter. Applies to both Vespa and OpenSearch.\n\n## Filter categories\n\n| Category | Fields | Join logic |\n|---|---|---|\n| **Visibility** | `hidden` | Always applied (unless `include_hidden`) |\n| **Tenant** | `tenant_id` | AND (multi-tenant only) |\n| **ACL** | `access_control_list` | OR within, AND with rest |\n| **Narrowing** | `source_type`, `tags`, `time_cutoff` | Each OR within, AND with rest |\n| **Knowledge scope** | `document_set`, `attached_document_ids`, `hierarchy_node_ids`, `persona_id_filter` | OR within group, AND with rest |\n| **Additive scope** | `project_id_filter` | OR'd into knowledge scope **only when** a knowledge scope filter already exists |\n\n## How filters combine\n\nAll categories are AND'd together. Within the knowledge scope category, individual filters are OR'd.\n\n```\nNOT hidden\nAND tenant = T -- if multi-tenant\nAND (acl contains A1 OR acl contains A2)\nAND (source_type = S1 OR ...) -- if set\nAND (tag = T1 OR ...) -- if set\nAND -- see below\nAND time >= cutoff -- if set\n```\n\n## Knowledge scope rules\n\nThe knowledge scope filter controls **what knowledge an assistant can access**.\n\n### Primary vs additive triggers\n\n- **`persona_id_filter`** is a **primary** trigger. A persona with user files IS explicit\n knowledge, so `persona_id_filter` alone can start a knowledge scope. Note: this is\n NOT the raw ID of the persona being used — it is only set when the persona's\n user files overflowed the LLM context window.\n- **`project_id_filter`** is **additive**. It widens an existing scope to include project\n files but never restricts on its own — a chat inside a project should still search\n team knowledge when no other knowledge is attached.\n\n### No explicit knowledge attached\n\nWhen `document_set`, `attached_document_ids`, `hierarchy_node_ids`, and `persona_id_filter` are all empty/None:\n\n- **No knowledge scope filter is applied.** The assistant can see everything (subject to ACL).\n- `project_id_filter` is ignored — it never restricts on its own.\n\n### One explicit knowledge type\n\n```\n-- Only document sets\nAND (document_sets contains \"Engineering\" OR document_sets contains \"Legal\")\n\n-- Only persona user files (overflowed context)\nAND (personas contains 42)\n```\n\n### Multiple explicit knowledge types (OR'd)\n\n```\n-- Document sets + persona user files\nAND (\n document_sets contains \"Engineering\"\n OR personas contains 42\n)\n```\n\n### Explicit knowledge + overflowing project files\n\nWhen an explicit knowledge restriction is in effect **and** `project_id_filter` is set (project files overflowed the LLM context window), `project_id_filter` widens the filter:\n\n```\n-- Document sets + project files overflowed\nAND (\n document_sets contains \"Engineering\"\n OR user_project contains 7\n)\n\n-- Persona user files + project files (won't happen in practice;\n-- custom personas ignore project files per the precedence rule)\nAND (\n personas contains 42\n OR user_project contains 7\n)\n```\n\n### Only project_id_filter (no explicit knowledge)\n\nNo knowledge scope filter. The assistant searches everything.\n\n```\n-- Just ACL, no restriction\nNOT hidden\nAND (acl contains ...)\n```\n\n## Field reference\n\n| Filter field | Vespa field | Vespa type | Purpose |\n|---|---|---|---|\n| `document_set` | `document_sets` | `weightedset` | Connector doc sets attached to assistant |\n| `attached_document_ids` | `document_id` | `string` | Documents explicitly attached (OpenSearch only) |\n| `hierarchy_node_ids` | `ancestor_hierarchy_node_ids` | `array` | Folder/space nodes (OpenSearch only) |\n| `persona_id_filter` | `personas` | `array` | Persona tag for overflowing user files (**primary** trigger) |\n| `project_id_filter` | `user_project` | `array` | Project tag for overflowing project files (**additive** only) |\n| `access_control_list` | `access_control_list` | `weightedset` | ACL entries for the requesting user |\n| `source_type` | `source_type` | `string` | Connector source type (e.g. `web`, `jira`) |\n| `tags` | `metadata_list` | `array` | Document metadata tags |\n| `time_cutoff` | `doc_updated_at` | `long` | Minimum document update timestamp |\n| `tenant_id` | `tenant_id` | `string` | Tenant isolation (multi-tenant) |\n" }, { "path": "backend/onyx/document_index/__init__.py", "content": "" }, { "path": "backend/onyx/document_index/chunk_content_enrichment.py", "content": "from onyx.configs.app_configs import BLURB_SIZE\nfrom onyx.configs.constants import RETURN_SEPARATOR\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceChunkUncleaned\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\n\n\ndef generate_enriched_content_for_chunk_text(chunk: DocMetadataAwareIndexChunk) -> str:\n return f\"{chunk.title_prefix}{chunk.doc_summary}{chunk.content}{chunk.chunk_context}{chunk.metadata_suffix_keyword}\"\n\n\ndef generate_enriched_content_for_chunk_embedding(chunk: DocAwareChunk) -> str:\n return f\"{chunk.title_prefix}{chunk.doc_summary}{chunk.content}{chunk.chunk_context}{chunk.metadata_suffix_semantic}\"\n\n\ndef cleanup_content_for_chunks(\n chunks: list[InferenceChunkUncleaned],\n) -> list[InferenceChunk]:\n \"\"\"\n Removes indexing-time content additions from chunks. Inverse of\n generate_enriched_content_for_chunk.\n\n During indexing, chunks are augmented with additional text to improve search\n quality:\n - Title prepended to content (for better keyword/semantic matching)\n - Metadata suffix appended to content\n - Contextual RAG: doc_summary (beginning) and chunk_context (end)\n\n This function strips these additions before returning chunks to users,\n restoring the original document content. Cleaning is applied in sequence:\n 1. Title removal:\n - Full match: Strips exact title from beginning\n - Partial match: If content starts with title[:BLURB_SIZE], splits on\n RETURN_SEPARATOR to remove title section\n 2. Metadata suffix removal:\n - Strips metadata_suffix from end, plus trailing RETURN_SEPARATOR\n 3. Contextual RAG removal:\n - Strips doc_summary from beginning (if present)\n - Strips chunk_context from end (if present)\n\n TODO(andrei): This entire function is not that fantastic, clean it up during\n QA before rolling out OpenSearch.\n\n Args:\n chunks: Chunks as retrieved from the document index with indexing\n augmentations intact.\n\n Returns:\n Clean InferenceChunk objects with augmentations removed, containing only\n the original document content that should be shown to users.\n \"\"\"\n\n def _remove_title(chunk: InferenceChunkUncleaned) -> str:\n # TODO(andrei): This was ported over from\n # backend/onyx/document_index/vespa/vespa_document_index.py but I don't\n # think this logic is correct. In Vespa at least we set the title field\n # from the output of get_title_for_document_index, which is not\n # necessarily the same data that is prepended to the content; that comes\n # from title_prefix.\n # This was added in\n # https://github.com/onyx-dot-app/onyx/commit/e90c66c1b61c5b7da949652d703f7c906863e6e4#diff-2a2a29d5929de75cdaea77867a397934d9f8b785ce40a861c0d704033e3663ab,\n # see postprocessing.py. At that time the content enrichment logic was\n # also added in that commit, see\n # https://github.com/onyx-dot-app/onyx/commit/e90c66c1b61c5b7da949652d703f7c906863e6e4#diff-d807718aa263a15c1d991a4ab063c360c8419eaad210b4ba70e1e9f47d2aa6d2R77\n # chunker.py.\n if not chunk.title or not chunk.content:\n return chunk.content\n\n if chunk.content.startswith(chunk.title):\n return chunk.content[len(chunk.title) :].lstrip()\n\n # BLURB SIZE is by token instead of char but each token is at least 1 char\n # If this prefix matches the content, it's assumed the title was prepended\n if chunk.content.startswith(chunk.title[:BLURB_SIZE]):\n return (\n chunk.content.split(RETURN_SEPARATOR, 1)[-1]\n if RETURN_SEPARATOR in chunk.content\n else chunk.content\n )\n return chunk.content\n\n def _remove_metadata_suffix(chunk: InferenceChunkUncleaned) -> str:\n if not chunk.metadata_suffix:\n return chunk.content\n return chunk.content.removesuffix(chunk.metadata_suffix).rstrip(\n RETURN_SEPARATOR\n )\n\n def _remove_contextual_rag(chunk: InferenceChunkUncleaned) -> str:\n # remove document summary\n if chunk.doc_summary and chunk.content.startswith(chunk.doc_summary):\n chunk.content = chunk.content[len(chunk.doc_summary) :].lstrip()\n # remove chunk context\n if chunk.chunk_context and chunk.content.endswith(chunk.chunk_context):\n chunk.content = chunk.content[\n : len(chunk.content) - len(chunk.chunk_context)\n ].rstrip()\n return chunk.content\n\n for chunk in chunks:\n chunk.content = _remove_title(chunk)\n chunk.content = _remove_metadata_suffix(chunk)\n chunk.content = _remove_contextual_rag(chunk)\n\n return [chunk.to_inference_chunk() for chunk in chunks]\n" }, { "path": "backend/onyx/document_index/disabled.py", "content": "\"\"\"A DocumentIndex implementation that raises on every operation.\n\nUsed as a safety net when DISABLE_VECTOR_DB is True. Any code path that\naccidentally reaches the vector DB layer will fail loudly instead of timing\nout against a nonexistent Vespa/OpenSearch instance.\n\"\"\"\n\nfrom collections.abc import Iterable\nfrom typing import Any\n\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import QueryExpansionType\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import DocumentInsertionRecord\nfrom onyx.document_index.interfaces import IndexBatchParams\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom shared_configs.model_server_models import Embedding\n\nVECTOR_DB_DISABLED_ERROR = \"Vector DB is disabled (DISABLE_VECTOR_DB=true). This operation requires a vector database.\"\n\n\nclass DisabledDocumentIndex(DocumentIndex):\n \"\"\"A DocumentIndex where every method raises RuntimeError.\n\n Returned by the factory when DISABLE_VECTOR_DB is True so that any\n accidental vector-DB call surfaces immediately.\n \"\"\"\n\n def __init__(\n self,\n index_name: str = \"disabled\",\n secondary_index_name: str | None = None,\n *args: Any, # noqa: ARG002\n **kwargs: Any, # noqa: ARG002\n ) -> None:\n self.index_name = index_name\n self.secondary_index_name = secondary_index_name\n\n # ------------------------------------------------------------------\n # Verifiable\n # ------------------------------------------------------------------\n def ensure_indices_exist(\n self,\n primary_embedding_dim: int, # noqa: ARG002\n primary_embedding_precision: EmbeddingPrecision, # noqa: ARG002\n secondary_index_embedding_dim: int | None, # noqa: ARG002\n secondary_index_embedding_precision: EmbeddingPrecision | None, # noqa: ARG002\n ) -> None:\n # No-op: there are no indices to create when the vector DB is disabled.\n pass\n\n @staticmethod\n def register_multitenant_indices(\n indices: list[str], # noqa: ARG002, ARG004\n embedding_dims: list[int], # noqa: ARG002, ARG004\n embedding_precisions: list[EmbeddingPrecision], # noqa: ARG002, ARG004\n ) -> None:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # Indexable\n # ------------------------------------------------------------------\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk], # noqa: ARG002\n index_batch_params: IndexBatchParams, # noqa: ARG002\n ) -> set[DocumentInsertionRecord]:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # Deletable\n # ------------------------------------------------------------------\n def delete_single(\n self,\n doc_id: str, # noqa: ARG002\n *,\n tenant_id: str, # noqa: ARG002\n chunk_count: int | None, # noqa: ARG002\n ) -> int:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # Updatable\n # ------------------------------------------------------------------\n def update_single(\n self,\n doc_id: str, # noqa: ARG002\n *,\n tenant_id: str, # noqa: ARG002\n chunk_count: int | None, # noqa: ARG002\n fields: VespaDocumentFields | None, # noqa: ARG002\n user_fields: VespaDocumentUserFields | None, # noqa: ARG002\n ) -> None:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # IdRetrievalCapable\n # ------------------------------------------------------------------\n def id_based_retrieval(\n self,\n chunk_requests: list[VespaChunkRequest], # noqa: ARG002\n filters: IndexFilters, # noqa: ARG002\n batch_retrieval: bool = False, # noqa: ARG002\n ) -> list[InferenceChunk]:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # HybridCapable\n # ------------------------------------------------------------------\n def hybrid_retrieval(\n self,\n query: str, # noqa: ARG002\n query_embedding: Embedding, # noqa: ARG002\n final_keywords: list[str] | None, # noqa: ARG002\n filters: IndexFilters, # noqa: ARG002\n hybrid_alpha: float, # noqa: ARG002\n time_decay_multiplier: float, # noqa: ARG002\n num_to_retrieve: int, # noqa: ARG002\n ranking_profile_type: QueryExpansionType, # noqa: ARG002\n title_content_ratio: float | None = None, # noqa: ARG002\n ) -> list[InferenceChunk]:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # AdminCapable\n # ------------------------------------------------------------------\n def admin_retrieval(\n self,\n query: str, # noqa: ARG002\n query_embedding: Embedding, # noqa: ARG002\n filters: IndexFilters, # noqa: ARG002\n num_to_retrieve: int = 10, # noqa: ARG002\n ) -> list[InferenceChunk]:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n\n # ------------------------------------------------------------------\n # RandomCapable\n # ------------------------------------------------------------------\n def random_retrieval(\n self,\n filters: IndexFilters, # noqa: ARG002\n num_to_retrieve: int = 10, # noqa: ARG002\n ) -> list[InferenceChunk]:\n raise RuntimeError(VECTOR_DB_DISABLED_ERROR)\n" }, { "path": "backend/onyx/document_index/document_index_utils.py", "content": "import math\nimport uuid\nfrom uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import ENABLE_MULTIPASS_INDEXING\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.db.search_settings import get_secondary_search_settings\nfrom onyx.document_index.interfaces import EnrichedDocumentIndexingInfo\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.indexing.models import MultipassConfig\nfrom shared_configs.configs import MULTI_TENANT\n\nDEFAULT_BATCH_SIZE = 30\nDEFAULT_INDEX_NAME = \"danswer_chunk\"\n\n\ndef should_use_multipass(search_settings: SearchSettings | None) -> bool:\n \"\"\"\n Determines whether multipass should be used based on the search settings\n or the default config if settings are unavailable.\n \"\"\"\n if search_settings is not None:\n return search_settings.multipass_indexing\n return ENABLE_MULTIPASS_INDEXING\n\n\ndef get_multipass_config(search_settings: SearchSettings) -> MultipassConfig:\n \"\"\"\n Determines whether to enable multipass and large chunks by examining\n the current search settings and the embedder configuration.\n \"\"\"\n multipass = should_use_multipass(search_settings)\n enable_large_chunks = SearchSettings.can_use_large_chunks(\n multipass, search_settings.model_name, search_settings.provider_type\n )\n return MultipassConfig(\n multipass_indexing=multipass, enable_large_chunks=enable_large_chunks\n )\n\n\ndef get_both_index_properties(\n db_session: Session,\n) -> tuple[str, str | None, bool, bool | None]:\n search_settings = get_current_search_settings(db_session)\n config_1 = get_multipass_config(search_settings)\n\n search_settings_new = get_secondary_search_settings(db_session)\n if not search_settings_new:\n return search_settings.index_name, None, config_1.enable_large_chunks, None\n\n config_2 = get_multipass_config(search_settings)\n return (\n search_settings.index_name,\n search_settings_new.index_name,\n config_1.enable_large_chunks,\n config_2.enable_large_chunks,\n )\n\n\ndef translate_boost_count_to_multiplier(boost: int) -> float:\n \"\"\"Mapping boost integer values to a multiplier according to a sigmoid curve\n Piecewise such that at many downvotes, its 0.5x the score and with many upvotes\n it is 2x the score. This should be in line with the Vespa calculation.\"\"\"\n # 3 in the equation below stretches it out to hit asymptotes slower\n if boost < 0:\n # 0.5 + sigmoid -> range of 0.5 to 1\n return 0.5 + (1 / (1 + math.exp(-1 * boost / 3)))\n\n # 2 x sigmoid -> range of 1 to 2\n return 2 / (1 + math.exp(-1 * boost / 3))\n\n\n# Assembles a list of Vespa chunk IDs for a document\n# given the required context. This can be used to directly query\n# Vespa's Document API.\ndef get_document_chunk_ids(\n enriched_document_info_list: list[EnrichedDocumentIndexingInfo],\n tenant_id: str,\n large_chunks_enabled: bool,\n) -> list[UUID]:\n doc_chunk_ids = []\n\n for enriched_document_info in enriched_document_info_list:\n for chunk_index in range(\n enriched_document_info.chunk_start_index,\n enriched_document_info.chunk_end_index,\n ):\n if not enriched_document_info.old_version:\n doc_chunk_ids.append(\n get_uuid_from_chunk_info(\n document_id=enriched_document_info.doc_id,\n chunk_id=chunk_index,\n tenant_id=tenant_id,\n )\n )\n else:\n doc_chunk_ids.append(\n get_uuid_from_chunk_info_old(\n document_id=enriched_document_info.doc_id,\n chunk_id=chunk_index,\n )\n )\n\n if large_chunks_enabled and chunk_index % 4 == 0:\n large_chunk_id = int(chunk_index / 4)\n large_chunk_reference_ids = [\n large_chunk_id + i\n for i in range(4)\n if large_chunk_id + i < enriched_document_info.chunk_end_index\n ]\n if enriched_document_info.old_version:\n doc_chunk_ids.append(\n get_uuid_from_chunk_info_old(\n document_id=enriched_document_info.doc_id,\n chunk_id=large_chunk_id,\n large_chunk_reference_ids=large_chunk_reference_ids,\n )\n )\n else:\n doc_chunk_ids.append(\n get_uuid_from_chunk_info(\n document_id=enriched_document_info.doc_id,\n chunk_id=large_chunk_id,\n tenant_id=tenant_id,\n large_chunk_id=large_chunk_id,\n )\n )\n\n return doc_chunk_ids\n\n\ndef get_uuid_from_chunk_info(\n *,\n document_id: str,\n chunk_id: int,\n tenant_id: str,\n large_chunk_id: int | None = None,\n) -> UUID:\n \"\"\"NOTE: be VERY carefuly about changing this function. If changed without a migration,\n this can cause deletion/update/insertion to function incorrectly.\"\"\"\n doc_str = document_id\n\n # Web parsing URL duplicate catching\n if doc_str and doc_str[-1] == \"/\":\n doc_str = doc_str[:-1]\n\n chunk_index = (\n \"large_\" + str(large_chunk_id) if large_chunk_id is not None else str(chunk_id)\n )\n unique_identifier_string = \"_\".join([doc_str, chunk_index])\n if MULTI_TENANT:\n unique_identifier_string += \"_\" + tenant_id\n\n uuid_value = uuid.uuid5(uuid.NAMESPACE_X500, unique_identifier_string)\n return uuid_value\n\n\ndef get_uuid_from_chunk_info_old(\n *, document_id: str, chunk_id: int, large_chunk_reference_ids: list[int] = []\n) -> UUID:\n doc_str = document_id\n\n # Web parsing URL duplicate catching\n if doc_str and doc_str[-1] == \"/\":\n doc_str = doc_str[:-1]\n unique_identifier_string = \"_\".join([doc_str, str(chunk_id), \"0\"])\n if large_chunk_reference_ids:\n unique_identifier_string += \"_large\" + \"_\".join(\n [\n str(referenced_chunk_id)\n for referenced_chunk_id in large_chunk_reference_ids\n ]\n )\n return uuid.uuid5(uuid.NAMESPACE_X500, unique_identifier_string)\n\n\ndef get_uuid_from_chunk(chunk: DocMetadataAwareIndexChunk) -> uuid.UUID:\n return get_uuid_from_chunk_info(\n document_id=chunk.source_document.id,\n chunk_id=chunk.chunk_id,\n tenant_id=chunk.tenant_id,\n large_chunk_id=chunk.large_chunk_id,\n )\n\n\ndef get_uuid_from_chunk_old(\n chunk: DocMetadataAwareIndexChunk, large_chunk_reference_ids: list[int] = []\n) -> UUID:\n return get_uuid_from_chunk_info_old(\n document_id=chunk.source_document.id,\n chunk_id=chunk.chunk_id,\n large_chunk_reference_ids=large_chunk_reference_ids,\n )\n" }, { "path": "backend/onyx/document_index/factory.py", "content": "import httpx\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.opensearch_migration import get_opensearch_retrieval_state\nfrom onyx.document_index.disabled import DisabledDocumentIndex\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchOldDocumentIndex,\n)\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.indexing.models import IndexingSetting\nfrom shared_configs.configs import MULTI_TENANT\n\n\ndef get_default_document_index(\n search_settings: SearchSettings,\n secondary_search_settings: SearchSettings | None,\n db_session: Session,\n httpx_client: httpx.Client | None = None,\n) -> DocumentIndex:\n \"\"\"Gets the default document index from env vars.\n\n To be used for retrieval only. Indexing should be done through both indices\n until Vespa is deprecated.\n\n Primary index is the index that is used for querying/updating etc. Secondary\n index is for when both the currently used index and the upcoming index both\n need to be updated. Updates are applied to both indices.\n WARNING: In that case, get_all_document_indices should be used.\n \"\"\"\n if DISABLE_VECTOR_DB:\n return DisabledDocumentIndex(\n index_name=search_settings.index_name,\n secondary_index_name=(\n secondary_search_settings.index_name\n if secondary_search_settings\n else None\n ),\n )\n\n secondary_index_name: str | None = None\n secondary_large_chunks_enabled: bool | None = None\n if secondary_search_settings:\n secondary_index_name = secondary_search_settings.index_name\n secondary_large_chunks_enabled = secondary_search_settings.large_chunks_enabled\n\n opensearch_retrieval_enabled = get_opensearch_retrieval_state(db_session)\n if opensearch_retrieval_enabled:\n indexing_setting = IndexingSetting.from_db_model(search_settings)\n secondary_indexing_setting = (\n IndexingSetting.from_db_model(secondary_search_settings)\n if secondary_search_settings\n else None\n )\n return OpenSearchOldDocumentIndex(\n index_name=search_settings.index_name,\n embedding_dim=indexing_setting.final_embedding_dim,\n embedding_precision=indexing_setting.embedding_precision,\n secondary_index_name=secondary_index_name,\n secondary_embedding_dim=(\n secondary_indexing_setting.final_embedding_dim\n if secondary_indexing_setting\n else None\n ),\n secondary_embedding_precision=(\n secondary_indexing_setting.embedding_precision\n if secondary_indexing_setting\n else None\n ),\n large_chunks_enabled=search_settings.large_chunks_enabled,\n secondary_large_chunks_enabled=secondary_large_chunks_enabled,\n multitenant=MULTI_TENANT,\n httpx_client=httpx_client,\n )\n else:\n return VespaIndex(\n index_name=search_settings.index_name,\n secondary_index_name=secondary_index_name,\n large_chunks_enabled=search_settings.large_chunks_enabled,\n secondary_large_chunks_enabled=secondary_large_chunks_enabled,\n multitenant=MULTI_TENANT,\n httpx_client=httpx_client,\n )\n\n\ndef get_all_document_indices(\n search_settings: SearchSettings,\n secondary_search_settings: SearchSettings | None,\n httpx_client: httpx.Client | None = None,\n) -> list[DocumentIndex]:\n \"\"\"Gets all document indices.\n\n NOTE: Will only return an OpenSearch index interface if\n ENABLE_OPENSEARCH_INDEXING_FOR_ONYX is True. This is so we don't break flows\n where we know it won't be enabled.\n\n Used for indexing only. Until Vespa is deprecated we will index into both\n document indices. Retrieval is done through only one index however.\n\n Large chunks are not currently supported so we hardcode appropriate values.\n\n NOTE: Make sure the Vespa index object is returned first. In the rare event\n that there is some conflict between indexing and the migration task, it is\n assumed that the state of Vespa is more up-to-date than the state of\n OpenSearch.\n \"\"\"\n if DISABLE_VECTOR_DB:\n return [\n DisabledDocumentIndex(\n index_name=search_settings.index_name,\n secondary_index_name=(\n secondary_search_settings.index_name\n if secondary_search_settings\n else None\n ),\n )\n ]\n\n vespa_document_index = VespaIndex(\n index_name=search_settings.index_name,\n secondary_index_name=(\n secondary_search_settings.index_name if secondary_search_settings else None\n ),\n large_chunks_enabled=search_settings.large_chunks_enabled,\n secondary_large_chunks_enabled=(\n secondary_search_settings.large_chunks_enabled\n if secondary_search_settings\n else None\n ),\n multitenant=MULTI_TENANT,\n httpx_client=httpx_client,\n )\n opensearch_document_index: OpenSearchOldDocumentIndex | None = None\n if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX:\n indexing_setting = IndexingSetting.from_db_model(search_settings)\n secondary_indexing_setting = (\n IndexingSetting.from_db_model(secondary_search_settings)\n if secondary_search_settings\n else None\n )\n opensearch_document_index = OpenSearchOldDocumentIndex(\n index_name=search_settings.index_name,\n embedding_dim=indexing_setting.final_embedding_dim,\n embedding_precision=indexing_setting.embedding_precision,\n secondary_index_name=(\n secondary_search_settings.index_name\n if secondary_search_settings\n else None\n ),\n secondary_embedding_dim=(\n secondary_indexing_setting.final_embedding_dim\n if secondary_indexing_setting\n else None\n ),\n secondary_embedding_precision=(\n secondary_indexing_setting.embedding_precision\n if secondary_indexing_setting\n else None\n ),\n large_chunks_enabled=search_settings.large_chunks_enabled,\n secondary_large_chunks_enabled=(\n secondary_search_settings.large_chunks_enabled\n if secondary_search_settings\n else None\n ),\n multitenant=MULTI_TENANT,\n httpx_client=httpx_client,\n )\n result: list[DocumentIndex] = [vespa_document_index]\n if opensearch_document_index:\n result.append(opensearch_document_index)\n return result\n" }, { "path": "backend/onyx/document_index/interfaces.py", "content": "import abc\nfrom collections.abc import Iterable\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom typing import Any\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.access.models import ExternalAccess\nfrom onyx.configs.chat_configs import NUM_RETURNED_HITS\nfrom onyx.configs.chat_configs import TITLE_CONTENT_RATIO\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import QueryExpansionType\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom shared_configs.model_server_models import Embedding\n\n\n@dataclass(frozen=True)\nclass DocumentInsertionRecord:\n document_id: str\n already_existed: bool\n\n\n@dataclass(frozen=True)\nclass VespaChunkRequest:\n document_id: str\n min_chunk_ind: int | None = None\n max_chunk_ind: int | None = None\n\n @property\n def is_capped(self) -> bool:\n # If the max chunk index is not None, then the chunk request is capped\n # If the min chunk index is None, we can assume the min is 0\n return self.max_chunk_ind is not None\n\n @property\n def range(self) -> int | None:\n if self.max_chunk_ind is not None:\n return (self.max_chunk_ind - (self.min_chunk_ind or 0)) + 1\n return None\n\n\n@dataclass\nclass IndexBatchParams:\n \"\"\"\n Information necessary for efficiently indexing a batch of documents\n \"\"\"\n\n doc_id_to_previous_chunk_cnt: dict[str, int]\n doc_id_to_new_chunk_cnt: dict[str, int]\n tenant_id: str\n large_chunks_enabled: bool\n\n\n@dataclass\nclass MinimalDocumentIndexingInfo:\n \"\"\"\n Minimal information necessary for indexing a document\n \"\"\"\n\n doc_id: str\n chunk_start_index: int\n\n\n@dataclass\nclass EnrichedDocumentIndexingInfo(MinimalDocumentIndexingInfo):\n \"\"\"\n Enriched information necessary for indexing a document, including version and chunk range.\n \"\"\"\n\n old_version: bool\n chunk_end_index: int\n\n\n@dataclass\nclass DocumentMetadata:\n \"\"\"\n Document information that needs to be inserted into Postgres on first time encountering this\n document during indexing across any of the connectors.\n \"\"\"\n\n connector_id: int\n credential_id: int\n document_id: str\n semantic_identifier: str\n first_link: str\n doc_updated_at: datetime | None = None\n # Emails, not necessarily attached to users\n # Users may not be in Onyx\n primary_owners: list[str] | None = None\n secondary_owners: list[str] | None = None\n from_ingestion_api: bool = False\n\n external_access: ExternalAccess | None = None\n doc_metadata: dict[str, Any] | None = None\n\n # The resolved database ID of the parent hierarchy node (folder/container)\n parent_hierarchy_node_id: int | None = None\n\n\n@dataclass\nclass VespaDocumentFields:\n \"\"\"\n Specifies fields in Vespa for a document. Fields set to None will be ignored.\n Perhaps we should name this in an implementation agnostic fashion, but it's more\n understandable like this for now.\n \"\"\"\n\n # all other fields except these 4 will always be left alone by the update request\n access: DocumentAccess | None = None\n document_sets: set[str] | None = None\n boost: float | None = None\n hidden: bool | None = None\n aggregated_chunk_boost_factor: float | None = None\n\n\n@dataclass\nclass VespaDocumentUserFields:\n \"\"\"\n Fields that are specific to the user who is indexing the document.\n \"\"\"\n\n user_projects: list[int] | None = None\n personas: list[int] | None = None\n\n\n@dataclass\nclass UpdateRequest:\n \"\"\"\n For all document_ids, update the allowed_users and the boost to the new values\n Does not update any of the None fields\n \"\"\"\n\n minimal_document_indexing_info: list[MinimalDocumentIndexingInfo]\n # all other fields except these 4 will always be left alone by the update request\n access: DocumentAccess | None = None\n document_sets: set[str] | None = None\n boost: float | None = None\n hidden: bool | None = None\n\n\nclass Verifiable(abc.ABC):\n \"\"\"\n Class must implement document index schema verification. For example, verify that all of the\n necessary attributes for indexing, querying, filtering, and fields to return from search are\n all valid in the schema.\n\n Parameters:\n - index_name: The name of the primary index currently used for querying\n - secondary_index_name: The name of the secondary index being built in the background, if it\n currently exists. Some functions on the document index act on both the primary and\n secondary index, some act on just one.\n \"\"\"\n\n @abc.abstractmethod\n def __init__(\n self,\n index_name: str,\n secondary_index_name: str | None,\n *args: Any,\n **kwargs: Any,\n ) -> None:\n super().__init__(*args, **kwargs)\n self.index_name = index_name\n self.secondary_index_name = secondary_index_name\n\n @abc.abstractmethod\n def ensure_indices_exist(\n self,\n primary_embedding_dim: int,\n primary_embedding_precision: EmbeddingPrecision,\n secondary_index_embedding_dim: int | None,\n secondary_index_embedding_precision: EmbeddingPrecision | None,\n ) -> None:\n \"\"\"\n Verify that the document index exists and is consistent with the expectations in the code.\n\n Parameters:\n - primary_embedding_dim: Vector dimensionality for the vector similarity part of the search\n - primary_embedding_precision: Precision of the vector similarity part of the search\n - secondary_index_embedding_dim: Vector dimensionality of the secondary index being built\n behind the scenes. The secondary index should only be built when switching\n embedding models therefore this dim should be different from the primary index.\n - secondary_index_embedding_precision: Precision of the vector similarity part of the secondary index\n \"\"\"\n raise NotImplementedError\n\n @staticmethod\n @abc.abstractmethod\n def register_multitenant_indices(\n indices: list[str],\n embedding_dims: list[int],\n embedding_precisions: list[EmbeddingPrecision],\n ) -> None:\n \"\"\"\n Register multitenant indices with the document index.\n \"\"\"\n raise NotImplementedError\n\n\nclass Indexable(abc.ABC):\n \"\"\"\n Class must implement the ability to index document chunks\n \"\"\"\n\n @abc.abstractmethod\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk],\n index_batch_params: IndexBatchParams,\n ) -> set[DocumentInsertionRecord]:\n \"\"\"\n Takes a list of document chunks and indexes them in the document index\n\n NOTE: When a document is reindexed/updated here, it must clear all of the existing document\n chunks before reindexing. This is because the document may have gotten shorter since the\n last run. Therefore, upserting the first 0 through n chunks may leave some old chunks that\n have not been written over.\n\n NOTE: The chunks of a document are never separated into separate index() calls. So there is\n no worry of receiving the first 0 through n chunks in one index call and the next n through\n m chunks of a docu in the next index call.\n\n NOTE: Due to some asymmetry between the primary and secondary indexing logic, this function\n only needs to index chunks into the PRIMARY index. Do not update the secondary index here,\n it is done automatically outside of this code.\n\n Parameters:\n - chunks: Document chunks with all of the information needed for\n indexing to the document index.\n - tenant_id: The tenant id of the user whose chunks are being indexed\n - large_chunks_enabled: Whether large chunks are enabled\n\n Returns:\n List of document ids which map to unique documents and are used for deduping chunks\n when updating, as well as if the document is newly indexed or already existed and\n just updated\n \"\"\"\n raise NotImplementedError\n\n\nclass Deletable(abc.ABC):\n \"\"\"\n Class must implement the ability to delete document by a given unique document id.\n \"\"\"\n\n @abc.abstractmethod\n def delete_single(\n self,\n doc_id: str,\n *,\n tenant_id: str,\n chunk_count: int | None,\n ) -> int:\n \"\"\"\n Given a single document id, hard delete it from the document index\n\n Parameters:\n - doc_id: document id as specified by the connector\n \"\"\"\n raise NotImplementedError\n\n\nclass Updatable(abc.ABC):\n \"\"\"\n Class must implement the ability to update certain attributes of a document without needing to\n update all of the fields. Specifically, needs to be able to update:\n - Access Control List\n - Document-set membership\n - Boost value (learning from feedback mechanism)\n - Whether the document is hidden or not, hidden documents are not returned from search\n \"\"\"\n\n @abc.abstractmethod\n def update_single(\n self,\n doc_id: str,\n *,\n tenant_id: str,\n chunk_count: int | None,\n fields: VespaDocumentFields | None,\n user_fields: VespaDocumentUserFields | None,\n ) -> None:\n \"\"\"\n Updates all chunks for a document with the specified fields.\n None values mean that the field does not need an update.\n\n The rationale for a single update function is that it allows retries and parallelism\n to happen at a higher / more strategic level, is simpler to read, and allows\n us to individually handle error conditions per document.\n\n Parameters:\n - fields: the fields to update in the document. Any field set to None will not be changed.\n\n Return:\n None\n \"\"\"\n raise NotImplementedError\n\n\nclass IdRetrievalCapable(abc.ABC):\n \"\"\"\n Class must implement the ability to retrieve either:\n - all of the chunks of a document IN ORDER given a document id.\n - a specific chunk given a document id and a chunk index (0 based)\n \"\"\"\n\n @abc.abstractmethod\n def id_based_retrieval(\n self,\n chunk_requests: list[VespaChunkRequest],\n filters: IndexFilters,\n batch_retrieval: bool = False,\n ) -> list[InferenceChunk]:\n \"\"\"\n Fetch chunk(s) based on document id\n\n NOTE: This is used to reconstruct a full document or an extended (multi-chunk) section\n of a document. Downstream currently assumes that the chunking does not introduce overlaps\n between the chunks. If there are overlaps for the chunks, then the reconstructed document\n or extended section will have duplicate segments.\n\n Parameters:\n - chunk_requests: requests containing the document id and the chunk range to retrieve\n - filters: Filters to apply to retrieval\n - batch_retrieval: If True, perform a batch retrieval\n\n Returns:\n list of chunks for the document id or the specific chunk by the specified chunk index\n and document id\n \"\"\"\n raise NotImplementedError\n\n\nclass HybridCapable(abc.ABC):\n \"\"\"\n Class must implement hybrid (keyword + vector) search functionality\n \"\"\"\n\n @abc.abstractmethod\n def hybrid_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n final_keywords: list[str] | None,\n filters: IndexFilters,\n hybrid_alpha: float,\n time_decay_multiplier: float,\n num_to_retrieve: int,\n ranking_profile_type: QueryExpansionType,\n title_content_ratio: float | None = TITLE_CONTENT_RATIO,\n ) -> list[InferenceChunk]:\n \"\"\"\n Run hybrid search and return a list of inference chunks.\n\n NOTE: the query passed in here is the unprocessed plain text query. Preprocessing is\n expected to be handled by this function as it may depend on the index implementation.\n Things like query expansion, synonym injection, stop word removal, lemmatization, etc. are\n done here.\n\n Parameters:\n - query: unmodified user query. This is needed for getting the matching highlighted\n keywords\n - query_embedding: vector representation of the query, must be of the correct\n dimensionality for the primary index\n - final_keywords: Final keywords to be used from the query, defaults to query if not set\n - filters: standard filter object\n - hybrid_alpha: weighting between the keyword and vector search results. It is important\n that the two scores are normalized to the same range so that a meaningful\n comparison can be made. 1 for 100% weighting on vector score, 0 for 100% weighting\n on keyword score.\n - time_decay_multiplier: how much to decay the document scores as they age. Some queries\n based on the persona settings, will have this be a 2x or 3x of the default\n - num_to_retrieve: number of highest matching chunks to return\n\n Returns:\n best matching chunks based on weighted sum of keyword and vector/semantic search scores\n \"\"\"\n raise NotImplementedError\n\n\nclass AdminCapable(abc.ABC):\n \"\"\"\n Class must implement a search for the admin \"Explorer\" page. The assumption here is that the\n admin is not \"searching\" for knowledge but has some document already in mind. They are either\n looking to positively boost it because they know it's a good reference document, looking to\n negatively boost it as a way of \"deprecating\", or hiding the document.\n\n Assuming the admin knows the document name, this search has high emphasis on the title match.\n\n Suggested implementation:\n Keyword only, BM25 search with 5x weighting on the title field compared to the contents\n \"\"\"\n\n @abc.abstractmethod\n def admin_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n filters: IndexFilters,\n num_to_retrieve: int = NUM_RETURNED_HITS,\n ) -> list[InferenceChunk]:\n \"\"\"\n Run the special search for the admin document explorer page\n\n Parameters:\n - query: unmodified user query. Though in this flow probably unmodified is best\n - filters: standard filter object\n - num_to_retrieve: number of highest matching chunks to return\n\n Returns:\n list of best matching chunks for the explorer page query\n \"\"\"\n raise NotImplementedError\n\n\nclass RandomCapable(abc.ABC):\n \"\"\"Class must implement random document retrieval capability\"\"\"\n\n @abc.abstractmethod\n def random_retrieval(\n self,\n filters: IndexFilters,\n num_to_retrieve: int = 10,\n ) -> list[InferenceChunk]:\n \"\"\"Retrieve random chunks matching the filters\"\"\"\n raise NotImplementedError\n\n\nclass BaseIndex(\n Verifiable,\n Indexable,\n Updatable,\n Deletable,\n AdminCapable,\n IdRetrievalCapable,\n RandomCapable,\n abc.ABC,\n):\n \"\"\"\n All basic document index functionalities excluding the actual querying approach.\n\n As a summary, document indices need to be able to\n - Verify the schema definition is valid\n - Index new documents\n - Update specific attributes of existing documents\n - Delete documents\n - Provide a search for the admin document explorer page\n - Retrieve documents based on document id\n \"\"\"\n\n\nclass DocumentIndex(HybridCapable, BaseIndex, abc.ABC):\n \"\"\"\n A valid document index that can plug into all Onyx flows must implement all of these\n functionalities, though \"technically\" it does not need to be keyword or vector capable as\n currently all default search flows use Hybrid Search.\n \"\"\"\n" }, { "path": "backend/onyx/document_index/interfaces_new.py", "content": "import abc\nfrom collections.abc import Iterable\nfrom typing import Self\n\nfrom pydantic import BaseModel\nfrom pydantic import model_validator\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\nfrom onyx.context.search.enums import QueryType\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom shared_configs.model_server_models import Embedding\n\n# NOTE: \"Document\" in the naming convention is used to refer to the entire\n# document as represented in Onyx. What is actually stored in the index is the\n# document chunks. By the terminology of most search engines / vector databases,\n# the individual objects stored are called documents, but in this case it refers\n# to a chunk.\n\n\n__all__ = [\n # Main interfaces - these are what you should inherit from\n \"DocumentIndex\",\n # Data models - used in method signatures\n \"DocumentInsertionRecord\",\n \"DocumentSectionRequest\",\n \"IndexingMetadata\",\n \"MetadataUpdateRequest\",\n # Capability mixins - for custom compositions or type checking\n \"SchemaVerifiable\",\n \"Indexable\",\n \"Deletable\",\n \"Updatable\",\n \"IdRetrievalCapable\",\n \"HybridCapable\",\n \"RandomCapable\",\n]\n\n\nclass TenantState(BaseModel):\n \"\"\"\n Captures the tenant-related state for an instance of DocumentIndex.\n\n NOTE: Tenant ID must be set in multitenant mode.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n tenant_id: str\n multitenant: bool\n\n def __str__(self) -> str:\n return (\n f\"TenantState(tenant_id={self.tenant_id}, multitenant={self.multitenant})\"\n )\n\n @model_validator(mode=\"after\")\n def check_tenant_id_is_set_in_multitenant_mode(self) -> Self:\n if self.multitenant and not self.tenant_id:\n raise ValueError(\"Bug: Tenant ID must be set in multitenant mode.\")\n return self\n\n\nclass DocumentInsertionRecord(BaseModel):\n \"\"\"\n Result of indexing a document.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n document_id: str\n already_existed: bool\n\n\nclass DocumentSectionRequest(BaseModel):\n \"\"\"Request for a document section or whole document.\n\n If no min_chunk_ind is provided it should start at the beginning of the\n document.\n If no max_chunk_ind is provided it should go to the end of the document.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n document_id: str\n min_chunk_ind: int | None = None\n max_chunk_ind: int | None = None\n # A given document can have multiple chunking strategies.\n max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE\n\n @model_validator(mode=\"after\")\n def check_chunk_index_range_is_valid(self) -> Self:\n if (\n self.min_chunk_ind is not None\n and self.max_chunk_ind is not None\n and self.min_chunk_ind > self.max_chunk_ind\n ):\n raise ValueError(\n \"Bug: Min chunk index must be less than or equal to max chunk index.\"\n )\n return self\n\n\nclass IndexingMetadata(BaseModel):\n \"\"\"\n Information about chunk counts for efficient cleaning / updating of document\n chunks.\n\n A common pattern to ensure that no chunks are left over is to delete all of\n the chunks for a document and then re-index the document. This information\n allows us to only delete the extra \"tail\" chunks when the document has\n gotten shorter.\n \"\"\"\n\n class ChunkCounts(BaseModel):\n model_config = {\"frozen\": True}\n\n old_chunk_cnt: int\n new_chunk_cnt: int\n\n model_config = {\"frozen\": True}\n\n doc_id_to_chunk_cnt_diff: dict[str, ChunkCounts]\n\n\nclass MetadataUpdateRequest(BaseModel):\n \"\"\"\n Updates to the documents that can happen without there being an update to\n the contents of the document.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n document_ids: list[str]\n # Passed in to help with potential optimizations of the implementation. The\n # keys should be redundant with document_ids.\n # NOTE: Generally the chunk count should always be known, however for\n # documents still using the legacy chunk ID system it may not be. Any chunk\n # count value < 0 should represent an unknown chunk count.\n doc_id_to_chunk_cnt: dict[str, int]\n # For the ones that are None, there is no update required to that field.\n access: DocumentAccess | None = None\n document_sets: set[str] | None = None\n boost: float | None = None\n hidden: bool | None = None\n secondary_index_updated: bool | None = None\n project_ids: set[int] | None = None\n persona_ids: set[int] | None = None\n\n\nclass IndexRetrievalFilters(BaseModel):\n \"\"\"\n Filters for retrieving chunks from the index.\n\n Used to filter on permissions and other Onyx-specific metadata rather than\n chunk content. Should be passed in for every retrieval method.\n\n TODO(andrei): Currently unused, use this when making retrieval methods more\n strict.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n # frozenset gets around the issue of python's mutable defaults.\n # WARNING: Falls back to only public docs as default for security. If\n # callers want no access filtering they must explicitly supply an empty set.\n # Doing so should be done sparingly.\n access_control_list: frozenset[str] = frozenset({PUBLIC_DOC_PAT})\n\n\nclass SchemaVerifiable(abc.ABC):\n \"\"\"\n Class must implement document index schema verification. For example, verify\n that all of the necessary attributes for indexing, querying, filtering, and\n fields to return from search are all valid in the schema.\n \"\"\"\n\n @abc.abstractmethod\n def verify_and_create_index_if_necessary(\n self,\n embedding_dim: int,\n embedding_precision: EmbeddingPrecision,\n ) -> None:\n \"\"\"\n Verifies that the document index exists and is consistent with the\n expectations in the code.\n\n For certain search engines, the schema needs to be created before\n indexing can happen. This call should create the schema if it does not\n exist.\n\n Args:\n embedding_dim: Vector dimensionality for the vector similarity part\n of the search.\n embedding_precision: Precision of the values of the vectors for the\n similarity part of the search.\n \"\"\"\n raise NotImplementedError\n\n\nclass Indexable(abc.ABC):\n \"\"\"\n Class must implement the ability to index document chunks.\n \"\"\"\n\n @abc.abstractmethod\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk],\n indexing_metadata: IndexingMetadata,\n ) -> list[DocumentInsertionRecord]:\n \"\"\"Indexes an iterable of document chunks into the document index.\n\n This is often a batch operation including chunks from multiple\n documents.\n\n NOTE: When a document is reindexed/updated here and has gotten shorter,\n it is important to delete the extra chunks at the end to ensure there\n are no stale chunks in the index. The implementation should do this.\n\n NOTE: The chunks of a document are never separated into separate index()\n calls. So there is no worry of receiving the first 0 through n chunks in\n one index call and the next n through m chunks of a document in the next\n index call.\n\n Args:\n chunks: Document chunks with all of the information needed for\n indexing to the document index.\n indexing_metadata: Information about chunk counts for efficient\n cleaning / updating.\n\n Returns:\n List of document IDs which map to unique documents as well as if the\n document is newly indexed or had already existed and was just\n updated.\n \"\"\"\n raise NotImplementedError\n\n\nclass Deletable(abc.ABC):\n \"\"\"\n Class must implement the ability to delete a document by a given unique\n document ID.\n \"\"\"\n\n @abc.abstractmethod\n def delete(\n self,\n # TODO(andrei): Fine for now but this can probably be a batch operation\n # that takes in a list of IDs.\n document_id: str,\n chunk_count: int | None = None,\n # TODO(andrei): Shouldn't this also have some acl filtering at minimum?\n ) -> int:\n \"\"\"\n Hard deletes all of the chunks for the corresponding document in the\n document index.\n\n TODO(andrei): Not a pressing issue now but think about what we want the\n contract of this method to be in the event the specified document ID\n does not exist.\n\n Args:\n document_id: The unique identifier for the document as represented\n in Onyx, not necessarily in the document index.\n chunk_count: The number of chunks in the document. May be useful for\n improving the efficiency of the delete operation. Defaults to\n None.\n\n Returns:\n The number of chunks deleted.\n \"\"\"\n raise NotImplementedError\n\n\nclass Updatable(abc.ABC):\n \"\"\"\n Class must implement the ability to update certain attributes of a document\n without needing to update all of the fields. Specifically, needs to be able\n to update:\n - Access Control List\n - Document-set membership\n - Boost value (learning from feedback mechanism)\n - Whether the document is hidden or not; hidden documents are not returned\n from search\n - Which Projects the document is a part of\n \"\"\"\n\n @abc.abstractmethod\n def update(\n self,\n update_requests: list[MetadataUpdateRequest],\n ) -> None:\n \"\"\"Updates some set of chunks.\n\n The document and fields to update are specified in the update requests.\n Each update request in the list applies its changes to a list of\n document IDs. None values mean that the field does not need an update.\n\n Args:\n update_requests: A list of update requests, each containing a list\n of document IDs and the fields to update. The field updates\n apply to all of the specified documents in each update request.\n \"\"\"\n raise NotImplementedError\n\n\nclass IdRetrievalCapable(abc.ABC):\n \"\"\"\n Class must implement the ability to retrieve either:\n - All of the chunks of a document IN ORDER given a document ID.\n - A specific section (continuous set of chunks) for some document.\n \"\"\"\n\n @abc.abstractmethod\n def id_based_retrieval(\n self,\n chunk_requests: list[DocumentSectionRequest],\n # TODO(andrei): Make this more strict w.r.t. acl, temporary for now.\n filters: IndexFilters,\n # TODO(andrei): This is temporary, we will not expose this in the long\n # run.\n batch_retrieval: bool = False,\n # TODO(andrei): Add a param for whether to retrieve hidden docs.\n ) -> list[InferenceChunk]:\n \"\"\"Fetches chunk(s) based on document ID.\n\n NOTE: This is used to reconstruct a full document or an extended\n (multi-chunk) section of a document. Downstream currently assumes that\n the chunking does not introduce overlaps between the chunks. If there\n are overlaps for the chunks, then the reconstructed document or extended\n section will have duplicate segments.\n\n Args:\n chunk_requests: Requests containing the document ID and the chunk\n range to retrieve.\n\n Returns:\n List of sections from the documents specified.\n \"\"\"\n raise NotImplementedError\n\n\nclass HybridCapable(abc.ABC):\n \"\"\"\n Class must implement hybrid (keyword + vector) search functionality.\n \"\"\"\n\n @abc.abstractmethod\n def hybrid_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n # TODO(andrei): This param is not great design, get rid of it.\n final_keywords: list[str] | None,\n query_type: QueryType,\n # TODO(andrei): Make this more strict w.r.t. acl, temporary for now.\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n \"\"\"Runs hybrid search and returns a list of inference chunks.\n\n Args:\n query: Unmodified user query. This may be needed for getting the\n matching highlighted keywords or for logging purposes.\n query_embedding: Vector representation of the query. Must be of the\n correct dimensionality for the primary index.\n final_keywords: Final keywords to be used from the query; defaults\n to query if not set.\n query_type: Semantic or keyword type query; may use different\n scoring logic for each.\n filters: Filters for things like permissions, source type, time,\n etc.\n num_to_retrieve: Number of highest matching chunks to return.\n\n Returns:\n Score-ranked (highest first) list of highest matching chunks.\n \"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def keyword_retrieval(\n self,\n query: str,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n \"\"\"Runs keyword-only search and returns a list of inference chunks.\n\n Args:\n query: User query.\n filters: Filters for things like permissions, source type, time,\n etc.\n num_to_retrieve: Number of highest matching chunks to return.\n\n Returns:\n Score-ranked (highest first) list of highest matching chunks.\n \"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def semantic_retrieval(\n self,\n query_embedding: Embedding,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n \"\"\"Runs semantic-only search and returns a list of inference chunks.\n\n Args:\n query_embedding: Vector representation of the query. Must be of the\n correct dimensionality for the primary index.\n filters: Filters for things like permissions, source type, time,\n etc.\n num_to_retrieve: Number of highest matching chunks to return.\n\n Returns:\n Score-ranked (highest first) list of highest matching chunks.\n \"\"\"\n raise NotImplementedError\n\n\nclass RandomCapable(abc.ABC):\n \"\"\"\n Class must implement random document retrieval.\n \"\"\"\n\n @abc.abstractmethod\n def random_retrieval(\n self,\n # TODO(andrei): Make this more strict w.r.t. acl, temporary for now.\n filters: IndexFilters,\n num_to_retrieve: int = 10,\n dirty: bool | None = None,\n ) -> list[InferenceChunk]:\n \"\"\"Retrieves random chunks matching the filters.\n\n Args:\n filters: Filters for things like permissions, source type, time,\n etc.\n num_to_retrieve: Number of chunks to retrieve. Defaults to 10.\n dirty: If set, retrieve chunks whose \"dirty\" flag matches this\n argument. If None, there is no restriction on retrieved chunks\n with respect to that flag. A chunk is considered dirty if there\n is a secondary index but the chunk's state has not been ported\n over to it yet. Defaults to None.\n\n Returns:\n List of chunks matching the filters.\n \"\"\"\n raise NotImplementedError\n\n\nclass DocumentIndex(\n SchemaVerifiable,\n Indexable,\n Updatable,\n Deletable,\n HybridCapable,\n IdRetrievalCapable,\n RandomCapable,\n abc.ABC,\n):\n \"\"\"\n A valid document index that can plug into all Onyx flows must implement all\n of these functionalities.\n\n As a high-level summary, document indices need to be able to:\n - Verify the schema definition is valid\n - Index new documents\n - Update specific attributes of existing documents\n - Delete documents\n - Run hybrid search\n - Retrieve document or sections of documents based on document id\n - Retrieve sets of random documents\n \"\"\"\n" }, { "path": "backend/onyx/document_index/opensearch/README.md", "content": "# Opensearch Idiosyncrasies\n\n## How it works at a high level\nOpensearch has 2 phases, a `Search` phase and a `Fetch` phase. The `Search` phase works by getting the document scores on each\nshard separately, then typically a fetch phase grabs all of the relevant fields/data for returning to the user. There is also\nan intermediate phase (seemingly built specifically to handle hybrid search queries) which can run in between as a processor.\nReferences:\nhttps://docs.opensearch.org/latest/search-plugins/search-pipelines/search-processors/\nhttps://docs.opensearch.org/latest/search-plugins/search-pipelines/normalization-processor/\nhttps://docs.opensearch.org/latest/query-dsl/compound/hybrid/\n\n## How Hybrid queries work\nHybrid queries are basically parallel queries that each run through their own `Search` phase and do not interact in any way.\nThey also run across all the shards. It is not entirely clear what happens if a combination pipeline is not specified for them,\nperhaps the scores are just summed.\n\nWhen the normalization processor is applied to keyword/vector hybrid searches, documents that show up due to keyword match may\nnot also have showed up in the vector search and vice versa. In these situations, it just receives a 0 score for the missing\nquery component. Opensearch does not run another phase to recapture those missing values. The impact of this is that after\nnormalizing, the missing scores are 0 but this is a higher score than if it actually received a non-zero score.\n\nThis may not be immediately obvious so an explanation is included here. If it got a non-zero score instead, it must be lower\nthan all of the other scores of the list (otherwise it would have shown up). Therefore it would impact the normalization and\npush the other scores higher so that it's not only the lowest score still, but now it's a differentiated lowest score. This is\nnot strictly the case in a multi-node setup but the high level concept approximately holds. So basically the 0 score is a form\nof \"minimum value clipping\".\n\n## On time decay and boosting\nEmbedding models do not have a uniform distribution from 0 to 1. The values typically cluster strongly around 0.6 to 0.8 but also\nvaries between models and even the query. It is not a safe assumption to pre-normalize the scores so we also cannot apply any\nadditive or multiplicative boost to it. i.e. if results of a doc cluster around 0.6 to 0.8 and I give a 50% penalty to the score,\nit doesn't bring a result from the top of the range to 50th percentile, it brings it under the 0.6 and is now the worst match.\nSame logic applies to additive boosting.\n\nSo these boosts can only be applied after normalization. Unfortunately with Opensearch, the normalization processor runs last\nand only applies to the results of the completely independent `Search` phase queries. So if a time based boost (a separate\nquery which filters on recently updated documents) is added, it would not be able to introduce any new documents\nto the set (since the new documents would have no keyword/vector score or already be present) since the 0 scores on keyword\nand vector would make the docs which only came because of time filter very low scoring. This can however make some of the lower\nscored documents from the union of all the `Search` phase documents to show up higher and potentially not get dropped before\nbeing fetched and returned to the user. But there are other issues of including these:\n- There is no way to sort by this field, only a filter, so there's no way to guarantee the best docs even irrespective of the\ncontents. If there are lots of updates, this may miss.\n- There is not a good way to normalize this field, the best is to clip it on the bottom.\n- This would require using min-max norm but z-score norm is better for the other functions due to things like it being less\nsensitive to outliers, better handles distribution drifts (min-max assumes stable meaningful ranges), better for comparing\n\"unusual-ness\" across distributions.\n\nSo while it is possible to apply time based boosting at the normalization stage (or specifically to the keyword score), we have\ndecided it is better to not apply it during the OpenSearch query.\n\nBecause of these limitations, Onyx in code applies further refinements, boostings, etc. based on OpenSearch providing an initial\nfiltering. The impact of time decay and boost should not be so big that we would need orders of magnitude more results back\nfrom OpenSearch.\n\n## Other concepts to be aware of\nWithin the `Search` phase, there are optional steps like Rescore but these are not useful for the combination/normalization\nwork that is relevant for the hybrid search. Since the Rescore happens prior to normalization, it's not able to provide any\nmeaningful operations to the query for our usage.\n\nBecause the Title is included in the Contents for both embedding and keyword searches, the Title scores are very low relative to\nthe actual full contents scoring. It is seen as a boost rather than a core scoring component. Time decay works similarly.\n" }, { "path": "backend/onyx/document_index/opensearch/client.py", "content": "import json\nimport logging\nimport time\nfrom contextlib import AbstractContextManager\nfrom contextlib import nullcontext\nfrom typing import Any\nfrom typing import Generic\nfrom typing import TypeVar\n\nfrom opensearchpy import OpenSearch\nfrom opensearchpy import TransportError\nfrom opensearchpy.helpers import bulk\nfrom pydantic import BaseModel\n\nfrom onyx.configs.app_configs import DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S\nfrom onyx.configs.app_configs import OPENSEARCH_ADMIN_PASSWORD\nfrom onyx.configs.app_configs import OPENSEARCH_ADMIN_USERNAME\nfrom onyx.configs.app_configs import OPENSEARCH_HOST\nfrom onyx.configs.app_configs import OPENSEARCH_REST_API_PORT\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.constants import OpenSearchSearchType\nfrom onyx.document_index.opensearch.schema import DocumentChunk\nfrom onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors\nfrom onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id\nfrom onyx.document_index.opensearch.search import DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW\nfrom onyx.server.metrics.opensearch_search import observe_opensearch_search\nfrom onyx.server.metrics.opensearch_search import track_opensearch_search_in_progress\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.timing import log_function_time\n\n\nCLIENT_THRESHOLD_TO_LOG_SLOW_SEARCH_MS = 2000\n\n\nlogger = setup_logger(__name__)\n# Set the logging level to WARNING to ignore INFO and DEBUG logs from\n# opensearch. By default it emits INFO-level logs for every request.\n# The opensearch-py library uses \"opensearch\" as the logger name for HTTP\n# requests (see opensearchpy/connection/base.py)\nopensearch_logger = logging.getLogger(\"opensearch\")\nopensearch_logger.setLevel(logging.WARNING)\n\n\nSchemaDocumentModel = TypeVar(\"SchemaDocumentModel\")\n\n\nclass SearchHit(BaseModel, Generic[SchemaDocumentModel]):\n \"\"\"Represents a hit from OpenSearch in response to a query.\n\n Templated on the specific document model as defined by a schema.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n # The document chunk source retrieved from OpenSearch.\n document_chunk: SchemaDocumentModel\n # The match score for the document chunk as calculated by OpenSearch. Only\n # relevant for \"fuzzy searches\"; this will be None for direct queries where\n # score is not relevant like direct retrieval on ID.\n score: float | None = None\n # Maps schema property name to a list of highlighted snippets with match\n # terms wrapped in tags (e.g. \"something keyword other thing\").\n match_highlights: dict[str, list[str]] = {}\n # Score explanation from OpenSearch when \"explain\": true is set in the\n # query. Contains detailed breakdown of how the score was calculated.\n explanation: dict[str, Any] | None = None\n\n\nclass IndexInfo(BaseModel):\n \"\"\"\n Represents information about an OpenSearch index.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n name: str\n health: str\n status: str\n num_primary_shards: str\n num_replica_shards: str\n docs_count: str\n docs_deleted: str\n created_at: str\n total_size: str\n primary_shards_size: str\n\n\ndef get_new_body_without_vectors(body: dict[str, Any]) -> dict[str, Any]:\n \"\"\"Recursively replaces vectors in the body with their length.\n\n TODO(andrei): Do better.\n\n Args:\n body: The body to replace the vectors.\n\n Returns:\n A copy of body with vectors replaced with their length.\n \"\"\"\n new_body: dict[str, Any] = {}\n for k, v in body.items():\n if k == \"vector\":\n new_body[k] = len(v)\n elif isinstance(v, dict):\n new_body[k] = get_new_body_without_vectors(v)\n elif isinstance(v, list) and len(v) > 0 and isinstance(v[0], dict):\n new_body[k] = [get_new_body_without_vectors(item) for item in v]\n else:\n new_body[k] = v\n return new_body\n\n\nclass OpenSearchClient(AbstractContextManager):\n \"\"\"Client for interacting with OpenSearch for cluster-level operations.\n\n Args:\n host: The host of the OpenSearch cluster.\n port: The port of the OpenSearch cluster.\n auth: The authentication credentials for the OpenSearch cluster. A tuple\n of (username, password).\n use_ssl: Whether to use SSL for the OpenSearch cluster. Defaults to\n True.\n verify_certs: Whether to verify the SSL certificates for the OpenSearch\n cluster. Defaults to False.\n ssl_show_warn: Whether to show warnings for SSL certificates. Defaults\n to False.\n timeout: The timeout for the OpenSearch cluster. Defaults to\n DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S.\n \"\"\"\n\n def __init__(\n self,\n host: str = OPENSEARCH_HOST,\n port: int = OPENSEARCH_REST_API_PORT,\n auth: tuple[str, str] = (OPENSEARCH_ADMIN_USERNAME, OPENSEARCH_ADMIN_PASSWORD),\n use_ssl: bool = True,\n verify_certs: bool = False,\n ssl_show_warn: bool = False,\n timeout: int = DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S,\n ):\n logger.debug(\n f\"Creating OpenSearch client with host {host}, port {port} and timeout {timeout} seconds.\"\n )\n self._client = OpenSearch(\n hosts=[{\"host\": host, \"port\": port}],\n http_auth=auth,\n use_ssl=use_ssl,\n verify_certs=verify_certs,\n ssl_show_warn=ssl_show_warn,\n # NOTE: This timeout applies to all requests the client makes,\n # including bulk indexing. When exceeded, the client will raise a\n # ConnectionTimeout and return no useful results. The OpenSearch\n # server will log that the client cancelled the request. To get\n # partial results from OpenSearch, pass in a timeout parameter to\n # your request body that is less than this value.\n timeout=timeout,\n )\n\n def __exit__(self, *_: Any) -> None:\n self.close()\n\n def __del__(self) -> None:\n try:\n self.close()\n except Exception:\n pass\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def create_search_pipeline(\n self,\n pipeline_id: str,\n pipeline_body: dict[str, Any],\n ) -> None:\n \"\"\"Creates a search pipeline.\n\n See the OpenSearch documentation for more information on the search\n pipeline body.\n https://docs.opensearch.org/latest/search-plugins/search-pipelines/index/\n\n Args:\n pipeline_id: The ID of the search pipeline to create.\n pipeline_body: The body of the search pipeline to create.\n\n Raises:\n Exception: There was an error creating the search pipeline.\n \"\"\"\n response = self._client.search_pipeline.put(id=pipeline_id, body=pipeline_body)\n if not response.get(\"acknowledged\", False):\n raise RuntimeError(f\"Failed to create search pipeline {pipeline_id}.\")\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def delete_search_pipeline(self, pipeline_id: str) -> None:\n \"\"\"Deletes a search pipeline.\n\n Args:\n pipeline_id: The ID of the search pipeline to delete.\n\n Raises:\n Exception: There was an error deleting the search pipeline.\n \"\"\"\n response = self._client.search_pipeline.delete(id=pipeline_id)\n if not response.get(\"acknowledged\", False):\n raise RuntimeError(f\"Failed to delete search pipeline {pipeline_id}.\")\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def put_cluster_settings(self, settings: dict[str, Any]) -> bool:\n \"\"\"Puts cluster settings.\n\n Args:\n settings: The settings to put.\n\n Raises:\n Exception: There was an error putting the cluster settings.\n\n Returns:\n True if the settings were put successfully, False otherwise.\n \"\"\"\n response = self._client.cluster.put_settings(body=settings)\n if response.get(\"acknowledged\", False):\n logger.info(\"Successfully put cluster settings.\")\n return True\n else:\n logger.error(f\"Failed to put cluster settings: {response}.\")\n return False\n\n @log_function_time(print_only=True, debug_only=True)\n def list_indices_with_info(self) -> list[IndexInfo]:\n \"\"\"\n Lists the indices in the OpenSearch cluster with information about each\n index.\n\n Returns:\n A list of IndexInfo objects for each index.\n \"\"\"\n response = self._client.cat.indices(format=\"json\")\n indices: list[IndexInfo] = []\n for raw_index_info in response:\n indices.append(\n IndexInfo(\n name=raw_index_info.get(\"index\", \"\"),\n health=raw_index_info.get(\"health\", \"\"),\n status=raw_index_info.get(\"status\", \"\"),\n num_primary_shards=raw_index_info.get(\"pri\", \"\"),\n num_replica_shards=raw_index_info.get(\"rep\", \"\"),\n docs_count=raw_index_info.get(\"docs.count\", \"\"),\n docs_deleted=raw_index_info.get(\"docs.deleted\", \"\"),\n created_at=raw_index_info.get(\"creation.date.string\", \"\"),\n total_size=raw_index_info.get(\"store.size\", \"\"),\n primary_shards_size=raw_index_info.get(\"pri.store.size\", \"\"),\n )\n )\n return indices\n\n @log_function_time(print_only=True, debug_only=True)\n def ping(self) -> bool:\n \"\"\"Pings the OpenSearch cluster.\n\n Returns:\n True if OpenSearch could be reached, False if it could not.\n \"\"\"\n return self._client.ping()\n\n def close(self) -> None:\n \"\"\"Closes the client.\n\n Raises:\n Exception: There was an error closing the client.\n \"\"\"\n self._client.close()\n\n\nclass OpenSearchIndexClient(OpenSearchClient):\n \"\"\"Client for interacting with OpenSearch for index-level operations.\n\n OpenSearch's Python module has pretty bad typing support so this client\n attempts to protect the rest of the codebase from this. As a consequence,\n most methods here return the minimum data needed for the rest of Onyx, and\n tend to rely on Exceptions to handle errors.\n\n TODO(andrei): This class currently assumes the structure of the database\n schema when it returns a DocumentChunk. Make the class, or at least the\n search method, templated on the structure the caller can expect.\n\n Args:\n index_name: The name of the index to interact with.\n host: The host of the OpenSearch cluster.\n port: The port of the OpenSearch cluster.\n auth: The authentication credentials for the OpenSearch cluster. A tuple\n of (username, password).\n use_ssl: Whether to use SSL for the OpenSearch cluster. Defaults to\n True.\n verify_certs: Whether to verify the SSL certificates for the OpenSearch\n cluster. Defaults to False.\n ssl_show_warn: Whether to show warnings for SSL certificates. Defaults\n to False.\n timeout: The timeout for the OpenSearch cluster. Defaults to\n DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S.\n \"\"\"\n\n def __init__(\n self,\n index_name: str,\n host: str = OPENSEARCH_HOST,\n port: int = OPENSEARCH_REST_API_PORT,\n auth: tuple[str, str] = (OPENSEARCH_ADMIN_USERNAME, OPENSEARCH_ADMIN_PASSWORD),\n use_ssl: bool = True,\n verify_certs: bool = False,\n ssl_show_warn: bool = False,\n timeout: int = DEFAULT_OPENSEARCH_CLIENT_TIMEOUT_S,\n emit_metrics: bool = True,\n ):\n super().__init__(\n host=host,\n port=port,\n auth=auth,\n use_ssl=use_ssl,\n verify_certs=verify_certs,\n ssl_show_warn=ssl_show_warn,\n timeout=timeout,\n )\n self._index_name = index_name\n self._emit_metrics = emit_metrics\n logger.debug(\n f\"OpenSearch client created successfully for index {self._index_name}.\"\n )\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def create_index(self, mappings: dict[str, Any], settings: dict[str, Any]) -> None:\n \"\"\"Creates the index.\n\n See the OpenSearch documentation for more information on mappings and\n settings.\n\n Args:\n mappings: The mappings for the index to create.\n settings: The settings for the index to create.\n\n Raises:\n Exception: There was an error creating the index.\n \"\"\"\n body: dict[str, Any] = {\n \"mappings\": mappings,\n \"settings\": settings,\n }\n logger.debug(f\"Creating index {self._index_name} with body {body}.\")\n response = self._client.indices.create(index=self._index_name, body=body)\n if not response.get(\"acknowledged\", False):\n raise RuntimeError(f\"Failed to create index {self._index_name}.\")\n response_index = response.get(\"index\", \"\")\n if response_index != self._index_name:\n raise RuntimeError(\n f\"OpenSearch responded with index name {response_index} when creating index {self._index_name}.\"\n )\n logger.debug(f\"Index {self._index_name} created successfully.\")\n\n @log_function_time(print_only=True, debug_only=True)\n def delete_index(self) -> bool:\n \"\"\"Deletes the index.\n\n Raises:\n Exception: There was an error deleting the index.\n\n Returns:\n True if the index was deleted, False if it did not exist.\n \"\"\"\n if not self._client.indices.exists(index=self._index_name):\n logger.warning(\n f\"Tried to delete index {self._index_name} but it does not exist.\"\n )\n return False\n\n logger.debug(f\"Deleting index {self._index_name}.\")\n response = self._client.indices.delete(index=self._index_name)\n if not response.get(\"acknowledged\", False):\n raise RuntimeError(f\"Failed to delete index {self._index_name}.\")\n return True\n\n @log_function_time(print_only=True, debug_only=True)\n def index_exists(self) -> bool:\n \"\"\"Checks if the index exists.\n\n Raises:\n Exception: There was an error checking if the index exists.\n\n Returns:\n True if the index exists, False if it does not.\n \"\"\"\n return self._client.indices.exists(index=self._index_name)\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def put_mapping(self, mappings: dict[str, Any]) -> None:\n \"\"\"Updates the index mapping in an idempotent manner.\n\n - Existing fields with the same definition: No-op (succeeds silently).\n - New fields: Added to the index.\n - Existing fields with different types: Raises exception (requires\n reindex).\n\n See the OpenSearch documentation for more information:\n https://docs.opensearch.org/latest/api-reference/index-apis/put-mapping/\n\n Args:\n mappings: The complete mapping definition to apply. This will be\n merged with existing mappings in the index.\n\n Raises:\n Exception: There was an error updating the mappings, such as\n attempting to change the type of an existing field.\n \"\"\"\n logger.debug(\n f\"Putting mappings for index {self._index_name} with mappings {mappings}.\"\n )\n response = self._client.indices.put_mapping(\n index=self._index_name, body=mappings\n )\n if not response.get(\"acknowledged\", False):\n raise RuntimeError(\n f\"Failed to put the mapping update for index {self._index_name}.\"\n )\n logger.debug(f\"Successfully put mappings for index {self._index_name}.\")\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def validate_index(self, expected_mappings: dict[str, Any]) -> bool:\n \"\"\"Validates the index.\n\n Short-circuit returns False on the first mismatch. Logs the mismatch.\n\n See the OpenSearch documentation for more information on the index\n mappings.\n https://docs.opensearch.org/latest/mappings/\n\n Args:\n mappings: The expected mappings of the index to validate.\n\n Raises:\n Exception: There was an error validating the index.\n\n Returns:\n True if the index is valid, False if it is not based on the mappings\n supplied.\n \"\"\"\n # OpenSearch's documentation makes no mention of what happens when you\n # invoke client.indices.get on an index that does not exist, so we check\n # for existence explicitly just to be sure.\n exists_response = self.index_exists()\n if not exists_response:\n logger.warning(\n f\"Tried to validate index {self._index_name} but it does not exist.\"\n )\n return False\n logger.debug(\n f\"Validating index {self._index_name} with expected mappings {expected_mappings}.\"\n )\n\n get_result = self._client.indices.get(index=self._index_name)\n index_info: dict[str, Any] = get_result.get(self._index_name, {})\n if not index_info:\n raise ValueError(\n f\"Bug: OpenSearch did not return any index info for index {self._index_name}, \"\n \"even though it confirmed that the index exists.\"\n )\n index_mapping_properties: dict[str, Any] = index_info.get(\"mappings\", {}).get(\n \"properties\", {}\n )\n expected_mapping_properties: dict[str, Any] = expected_mappings.get(\n \"properties\", {}\n )\n assert (\n expected_mapping_properties\n ), \"Bug: No properties were found in the provided expected mappings.\"\n\n for property in expected_mapping_properties:\n if property not in index_mapping_properties:\n logger.warning(\n f'The field \"{property}\" was not found in the index {self._index_name}.'\n )\n return False\n\n expected_property_type = expected_mapping_properties[property].get(\n \"type\", \"\"\n )\n assert (\n expected_property_type\n ), f'Bug: The field \"{property}\" in the supplied expected schema mappings has no type.'\n\n index_property_type = index_mapping_properties[property].get(\"type\", \"\")\n if expected_property_type != index_property_type:\n logger.warning(\n f'The field \"{property}\" in the index {self._index_name} has type {index_property_type} '\n f\"but the expected type is {expected_property_type}.\"\n )\n return False\n\n logger.debug(f\"Index {self._index_name} validated successfully.\")\n return True\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def update_settings(self, settings: dict[str, Any]) -> None:\n \"\"\"Updates the settings of the index.\n\n See the OpenSearch documentation for more information on the index\n settings.\n https://docs.opensearch.org/latest/install-and-configure/configuring-opensearch/index-settings/\n\n Args:\n settings: The settings to update the index with.\n\n Raises:\n Exception: There was an error updating the settings of the index.\n \"\"\"\n # TODO(andrei): Implement this.\n raise NotImplementedError\n\n @log_function_time(\n print_only=True,\n debug_only=True,\n include_args_subset={\n \"document\": str,\n \"tenant_state\": str,\n \"update_if_exists\": str,\n },\n )\n def index_document(\n self,\n document: DocumentChunk,\n tenant_state: TenantState,\n update_if_exists: bool = False,\n ) -> None:\n \"\"\"Indexes a document.\n\n Args:\n document: The document to index. In Onyx this is a chunk of a\n document, OpenSearch simply refers to this as a document as\n well.\n tenant_state: The tenant state of the caller.\n update_if_exists: Whether to update the document if it already\n exists. If False, will raise an exception if the document\n already exists. Defaults to False.\n\n Raises:\n Exception: There was an error indexing the document. This includes\n the case where a document with the same ID already exists if\n update_if_exists is False.\n \"\"\"\n logger.debug(\n f\"Trying to index document ID {document.document_id} for tenant {tenant_state.tenant_id}. \"\n f\"update_if_exists={update_if_exists}.\"\n )\n document_chunk_id: str = get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=document.document_id,\n chunk_index=document.chunk_index,\n max_chunk_size=document.max_chunk_size,\n )\n body: dict[str, Any] = document.model_dump(exclude_none=True)\n # client.create will raise if a doc with the same ID exists.\n # client.index does not do this.\n if update_if_exists:\n result = self._client.index(\n index=self._index_name, id=document_chunk_id, body=body\n )\n else:\n result = self._client.create(\n index=self._index_name, id=document_chunk_id, body=body\n )\n result_id = result.get(\"_id\", \"\")\n # Sanity check.\n if result_id != document_chunk_id:\n raise RuntimeError(\n f'Upon trying to index a document, OpenSearch responded with ID \"{result_id}\" '\n f'instead of \"{document_chunk_id}\" which is the ID it was given.'\n )\n result_string: str = result.get(\"result\", \"\")\n match result_string:\n # Sanity check.\n case \"created\":\n pass\n case \"updated\":\n if not update_if_exists:\n raise RuntimeError(\n f'The OpenSearch client returned result \"updated\" for indexing document chunk \"{document_chunk_id}\". '\n \"This indicates that a document chunk with that ID already exists, which is not expected.\"\n )\n case _:\n raise RuntimeError(\n f'Unknown OpenSearch indexing result: \"{result_string}\".'\n )\n logger.debug(f\"Successfully indexed {document_chunk_id}.\")\n\n @log_function_time(\n print_only=True,\n debug_only=True,\n include_args_subset={\n \"documents\": len,\n \"tenant_state\": str,\n \"update_if_exists\": str,\n },\n )\n def bulk_index_documents(\n self,\n documents: list[DocumentChunk],\n tenant_state: TenantState,\n update_if_exists: bool = False,\n ) -> None:\n \"\"\"Bulk indexes documents.\n\n Raises if there are any errors during the bulk index. It should be\n assumed that no documents in the batch were indexed successfully if\n there is an error.\n\n Retries on 429 too many requests.\n\n Args:\n documents: The documents to index. In Onyx this is a chunk of a\n document, OpenSearch simply refers to this as a document as\n well.\n tenant_state: The tenant state of the caller.\n update_if_exists: Whether to update the document if it already\n exists. If False, will raise an exception if the document\n already exists. Defaults to False.\n\n Raises:\n Exception: There was an error during the bulk index. This\n includes the case where a document with the same ID already\n exists if update_if_exists is False.\n \"\"\"\n if not documents:\n return\n logger.debug(\n f\"Bulk indexing {len(documents)} documents for tenant {tenant_state.tenant_id}. update_if_exists={update_if_exists}.\"\n )\n data = []\n for document in documents:\n document_chunk_id: str = get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=document.document_id,\n chunk_index=document.chunk_index,\n max_chunk_size=document.max_chunk_size,\n )\n body: dict[str, Any] = document.model_dump(exclude_none=True)\n data_for_document: dict[str, Any] = {\n \"_index\": self._index_name,\n \"_id\": document_chunk_id,\n \"_op_type\": \"index\" if update_if_exists else \"create\",\n \"_source\": body,\n }\n data.append(data_for_document)\n # max_retries is the number of times to retry a request if we get a 429.\n success, errors = bulk(self._client, data, max_retries=3)\n if errors:\n raise RuntimeError(\n f\"Failed to bulk index documents for index {self._index_name}. Errors: {errors}\"\n )\n if success != len(documents):\n raise RuntimeError(\n f\"OpenSearch reported no errors during bulk index but the number of successful operations \"\n f\"({success}) does not match the number of documents ({len(documents)}).\"\n )\n logger.debug(f\"Successfully bulk indexed {len(documents)} documents.\")\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def delete_document(self, document_chunk_id: str) -> bool:\n \"\"\"Deletes a document.\n\n Args:\n document_chunk_id: The OpenSearch ID of the document chunk to\n delete.\n\n Raises:\n Exception: There was an error deleting the document.\n\n Returns:\n True if the document was deleted, False if it was not found.\n \"\"\"\n try:\n logger.debug(\n f\"Trying to delete document chunk {document_chunk_id} from index {self._index_name}.\"\n )\n result = self._client.delete(index=self._index_name, id=document_chunk_id)\n except TransportError as e:\n if e.status_code == 404:\n logger.debug(\n f\"Document chunk {document_chunk_id} not found in index {self._index_name}.\"\n )\n return False\n else:\n raise e\n\n result_string: str = result.get(\"result\", \"\")\n match result_string:\n case \"deleted\":\n logger.debug(\n f\"Successfully deleted document chunk {document_chunk_id} from index {self._index_name}.\"\n )\n return True\n case \"not_found\":\n logger.debug(\n f\"Document chunk {document_chunk_id} not found in index {self._index_name}.\"\n )\n return False\n case _:\n raise RuntimeError(\n f'Unknown OpenSearch deletion result: \"{result_string}\".'\n )\n\n @log_function_time(print_only=True, debug_only=True)\n def delete_by_query(self, query_body: dict[str, Any]) -> int:\n \"\"\"Deletes documents by a query.\n\n Args:\n query_body: The body of the query to delete documents by.\n\n Raises:\n Exception: There was an error deleting the documents.\n\n Returns:\n The number of documents deleted.\n \"\"\"\n logger.debug(\n f\"Trying to delete documents by query for index {self._index_name}.\"\n )\n result = self._client.delete_by_query(index=self._index_name, body=query_body)\n if result.get(\"timed_out\", False):\n raise RuntimeError(\n f\"Delete by query timed out for index {self._index_name}.\"\n )\n if len(result.get(\"failures\", [])) > 0:\n raise RuntimeError(\n f\"Failed to delete some or all of the documents for index {self._index_name}.\"\n )\n\n num_deleted = result.get(\"deleted\", 0)\n num_processed = result.get(\"total\", 0)\n if num_deleted != num_processed:\n raise RuntimeError(\n f\"Failed to delete some or all of the documents for index {self._index_name}. \"\n f\"{num_deleted} documents were deleted out of {num_processed} documents that were processed.\"\n )\n\n logger.debug(\n f\"Successfully deleted {num_deleted} documents by query for index {self._index_name}.\"\n )\n return num_deleted\n\n @log_function_time(\n print_only=True,\n debug_only=True,\n include_args_subset={\n \"document_chunk_id\": str,\n \"properties_to_update\": lambda x: x.keys(),\n },\n )\n def update_document(\n self, document_chunk_id: str, properties_to_update: dict[str, Any]\n ) -> None:\n \"\"\"Updates an OpenSearch document chunk's properties.\n\n Args:\n document_chunk_id: The OpenSearch ID of the document chunk to\n update.\n properties_to_update: The properties of the document to update. Each\n property should exist in the schema.\n\n Raises:\n Exception: There was an error updating the document.\n \"\"\"\n logger.debug(\n f\"Trying to update document chunk {document_chunk_id} for index {self._index_name}.\"\n )\n update_body: dict[str, Any] = {\"doc\": properties_to_update}\n result = self._client.update(\n index=self._index_name,\n id=document_chunk_id,\n body=update_body,\n _source=False,\n )\n result_id = result.get(\"_id\", \"\")\n # Sanity check.\n if result_id != document_chunk_id:\n raise RuntimeError(\n f'Upon trying to update a document, OpenSearch responded with ID \"{result_id}\" '\n f'instead of \"{document_chunk_id}\" which is the ID it was given.'\n )\n result_string: str = result.get(\"result\", \"\")\n match result_string:\n # Sanity check.\n case \"updated\":\n logger.debug(\n f\"Successfully updated document chunk {document_chunk_id} for index {self._index_name}.\"\n )\n return\n case \"noop\":\n logger.warning(\n f'OpenSearch reported a no-op when trying to update document with ID \"{document_chunk_id}\".'\n )\n return\n case _:\n raise RuntimeError(\n f'The OpenSearch client returned result \"{result_string}\" for updating document chunk \"{document_chunk_id}\". '\n \"This is unexpected.\"\n )\n\n @log_function_time(print_only=True, debug_only=True, include_args=True)\n def get_document(self, document_chunk_id: str) -> DocumentChunk:\n \"\"\"Gets an OpenSearch document chunk.\n\n Will raise an exception if the document chunk is not found.\n\n Args:\n document_chunk_id: The OpenSearch ID of the document chunk to get.\n\n Raises:\n Exception: There was an error getting the document. This includes\n the case where the document is not found.\n\n Returns:\n The document chunk.\n \"\"\"\n logger.debug(\n f\"Trying to get document chunk {document_chunk_id} from index {self._index_name}.\"\n )\n result = self._client.get(index=self._index_name, id=document_chunk_id)\n found_result: bool = result.get(\"found\", False)\n if not found_result:\n raise RuntimeError(\n f'Document chunk with ID \"{document_chunk_id}\" was not found.'\n )\n\n document_chunk_source: dict[str, Any] | None = result.get(\"_source\")\n if not document_chunk_source:\n raise RuntimeError(\n f'Document chunk with ID \"{document_chunk_id}\" has no data.'\n )\n\n logger.debug(\n f\"Successfully got document chunk {document_chunk_id} from index {self._index_name}.\"\n )\n return DocumentChunk.model_validate(document_chunk_source)\n\n @log_function_time(print_only=True, debug_only=True)\n def search(\n self,\n body: dict[str, Any],\n search_pipeline_id: str | None,\n search_type: OpenSearchSearchType = OpenSearchSearchType.UNKNOWN,\n ) -> list[SearchHit[DocumentChunkWithoutVectors]]:\n \"\"\"Searches the index.\n\n NOTE: Does not return vector fields. In order to take advantage of\n performance benefits, the search body should exclude the schema's vector\n fields.\n\n TODO(andrei): Ideally we could check that every field in the body is\n present in the index, to avoid a class of runtime bugs that could easily\n be caught during development. Or change the function signature to accept\n a predefined pydantic model of allowed fields.\n\n Args:\n body: The body of the search request. See the OpenSearch\n documentation for more information on search request bodies.\n search_pipeline_id: The ID of the search pipeline to use. If None,\n the default search pipeline will be used.\n search_type: Label for Prometheus metrics. Does not affect search\n behavior.\n\n Raises:\n Exception: There was an error searching the index.\n\n Returns:\n List of search hits that match the search request.\n \"\"\"\n logger.debug(\n f\"Trying to search index {self._index_name} with search pipeline {search_pipeline_id}.\"\n )\n result: dict[str, Any]\n params = {\"phase_took\": \"true\"}\n ctx = self._get_emit_metrics_context_manager(search_type)\n t0 = time.perf_counter()\n with ctx:\n if search_pipeline_id:\n result = self._client.search(\n index=self._index_name,\n search_pipeline=search_pipeline_id,\n body=body,\n params=params,\n )\n else:\n result = self._client.search(\n index=self._index_name, body=body, params=params\n )\n client_duration_s = time.perf_counter() - t0\n\n hits, time_took, timed_out, phase_took, profile = (\n self._get_hits_and_profile_from_search_result(result)\n )\n if self._emit_metrics:\n observe_opensearch_search(search_type, client_duration_s, time_took)\n self._log_search_result_perf(\n time_took=time_took,\n timed_out=timed_out,\n phase_took=phase_took,\n profile=profile,\n body=body,\n search_pipeline_id=search_pipeline_id,\n raise_on_timeout=True,\n )\n\n search_hits: list[SearchHit[DocumentChunkWithoutVectors]] = []\n for hit in hits:\n document_chunk_source: dict[str, Any] | None = hit.get(\"_source\")\n if not document_chunk_source:\n raise RuntimeError(\n f'Document chunk with ID \"{hit.get(\"_id\", \"\")}\" has no data.'\n )\n document_chunk_score = hit.get(\"_score\", None)\n match_highlights: dict[str, list[str]] = hit.get(\"highlight\", {})\n explanation: dict[str, Any] | None = hit.get(\"_explanation\", None)\n search_hit = SearchHit[DocumentChunkWithoutVectors](\n document_chunk=DocumentChunkWithoutVectors.model_validate(\n document_chunk_source\n ),\n score=document_chunk_score,\n match_highlights=match_highlights,\n explanation=explanation,\n )\n search_hits.append(search_hit)\n logger.debug(\n f\"Successfully searched index {self._index_name} and got {len(search_hits)} hits.\"\n )\n return search_hits\n\n @log_function_time(print_only=True, debug_only=True)\n def search_for_document_ids(\n self,\n body: dict[str, Any],\n search_type: OpenSearchSearchType = OpenSearchSearchType.UNKNOWN,\n ) -> list[str]:\n \"\"\"Searches the index and returns only document chunk IDs.\n\n In order to take advantage of the performance benefits of only returning\n IDs, the body should have a key, value pair of \"_source\": False.\n Otherwise, OpenSearch will return the entire document body and this\n method's performance will be the same as the search method's.\n\n TODO(andrei): Ideally we could check that every field in the body is\n present in the index, to avoid a class of runtime bugs that could easily\n be caught during development.\n\n Args:\n body: The body of the search request. See the OpenSearch\n documentation for more information on search request bodies.\n TODO(andrei): Make this a more deep interface; callers shouldn't\n need to know to set _source: False for example.\n search_type: Label for Prometheus metrics. Does not affect search\n behavior.\n\n Raises:\n Exception: There was an error searching the index.\n\n Returns:\n List of document chunk IDs that match the search request.\n \"\"\"\n logger.debug(\n f\"Trying to search for document chunk IDs in index {self._index_name}.\"\n )\n if \"_source\" not in body or body[\"_source\"] is not False:\n logger.warning(\n \"The body of the search request for document chunk IDs is missing the key, value pair of \"\n '\"_source\": False. This query will therefore be inefficient.'\n )\n\n params = {\"phase_took\": \"true\"}\n ctx = self._get_emit_metrics_context_manager(search_type)\n t0 = time.perf_counter()\n with ctx:\n result: dict[str, Any] = self._client.search(\n index=self._index_name, body=body, params=params\n )\n client_duration_s = time.perf_counter() - t0\n\n hits, time_took, timed_out, phase_took, profile = (\n self._get_hits_and_profile_from_search_result(result)\n )\n if self._emit_metrics:\n observe_opensearch_search(search_type, client_duration_s, time_took)\n self._log_search_result_perf(\n time_took=time_took,\n timed_out=timed_out,\n phase_took=phase_took,\n profile=profile,\n body=body,\n raise_on_timeout=True,\n )\n\n # TODO(andrei): Implement scroll/point in time for results so that we\n # can return arbitrarily-many IDs.\n if len(hits) == DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW:\n logger.warning(\n \"The search request for document chunk IDs returned the maximum number of results. \"\n \"It is extremely likely that there are more hits in OpenSearch than the returned results.\"\n )\n\n # Extract only the _id field from each hit.\n document_chunk_ids: list[str] = []\n for hit in hits:\n document_chunk_id = hit.get(\"_id\")\n if not document_chunk_id:\n raise RuntimeError(\n \"Received a hit from OpenSearch but the _id field is missing.\"\n )\n document_chunk_ids.append(document_chunk_id)\n logger.debug(\n f\"Successfully searched for document chunk IDs in index {self._index_name} and got {len(document_chunk_ids)} hits.\"\n )\n return document_chunk_ids\n\n @log_function_time(print_only=True, debug_only=True)\n def refresh_index(self) -> None:\n \"\"\"Refreshes the index to make recent changes searchable.\n\n In OpenSearch, documents are not immediately searchable after indexing.\n This method forces a refresh to make them available for search.\n\n Raises:\n Exception: There was an error refreshing the index.\n \"\"\"\n self._client.indices.refresh(index=self._index_name)\n\n def _get_hits_and_profile_from_search_result(\n self, result: dict[str, Any]\n ) -> tuple[list[Any], int | None, bool | None, dict[str, Any], dict[str, Any]]:\n \"\"\"Extracts the hits and profiling information from a search result.\n\n Args:\n result: The search result to extract the hits from.\n\n Raises:\n Exception: There was an error extracting the hits from the search\n result.\n\n Returns:\n A tuple containing the hits from the search result, the time taken\n to execute the search in milliseconds, whether the search timed\n out, the time taken to execute each phase of the search, and the\n profile.\n \"\"\"\n time_took: int | None = result.get(\"took\")\n timed_out: bool | None = result.get(\"timed_out\")\n phase_took: dict[str, Any] = result.get(\"phase_took\", {})\n profile: dict[str, Any] = result.get(\"profile\", {})\n\n hits_first_layer: dict[str, Any] = result.get(\"hits\", {})\n if not hits_first_layer:\n raise RuntimeError(\n f\"Hits field missing from response when trying to search index {self._index_name}.\"\n )\n hits_second_layer: list[Any] = hits_first_layer.get(\"hits\", [])\n\n return hits_second_layer, time_took, timed_out, phase_took, profile\n\n def _log_search_result_perf(\n self,\n time_took: int | None,\n timed_out: bool | None,\n phase_took: dict[str, Any],\n profile: dict[str, Any],\n body: dict[str, Any],\n search_pipeline_id: str | None = None,\n raise_on_timeout: bool = False,\n ) -> None:\n \"\"\"Logs the performance of a search result.\n\n Args:\n time_took: The time taken to execute the search in milliseconds.\n timed_out: Whether the search timed out.\n phase_took: The time taken to execute each phase of the search.\n profile: The profile for the search.\n body: The body of the search request for logging.\n search_pipeline_id: The ID of the search pipeline used for the\n search, if any, for logging. Defaults to None.\n raise_on_timeout: Whether to raise an exception if the search timed\n out. Note that the result may still contain useful partial\n results. Defaults to False.\n\n Raises:\n Exception: If raise_on_timeout is True and the search timed out.\n \"\"\"\n if time_took and time_took > CLIENT_THRESHOLD_TO_LOG_SLOW_SEARCH_MS:\n logger.warning(\n f\"OpenSearch client warning: Search for index {self._index_name} took {time_took} milliseconds.\\n\"\n f\"Body: {get_new_body_without_vectors(body)}\\n\"\n f\"Search pipeline ID: {search_pipeline_id}\\n\"\n f\"Phase took: {phase_took}\\n\"\n f\"Profile: {json.dumps(profile, indent=2)}\\n\"\n )\n if timed_out:\n error_str = f\"OpenSearch client error: Search timed out for index {self._index_name}.\"\n logger.error(error_str)\n if raise_on_timeout:\n raise RuntimeError(error_str)\n\n def _get_emit_metrics_context_manager(\n self, search_type: OpenSearchSearchType\n ) -> AbstractContextManager[None]:\n \"\"\"\n Returns a context manager that tracks in-flight OpenSearch searches via\n a Gauge if emit_metrics is True, otherwise returns a null context\n manager.\n \"\"\"\n return (\n track_opensearch_search_in_progress(search_type)\n if self._emit_metrics\n else nullcontext()\n )\n\n\ndef wait_for_opensearch_with_timeout(\n wait_interval_s: int = 5,\n wait_limit_s: int = 60,\n client: OpenSearchClient | None = None,\n) -> bool:\n \"\"\"Waits for OpenSearch to become ready subject to a timeout.\n\n Will create a new dummy client if no client is provided. Will close this\n client at the end of the function. Will not close the client if it was\n supplied.\n\n Args:\n wait_interval_s: The interval in seconds to wait between checks.\n Defaults to 5.\n wait_limit_s: The total timeout in seconds to wait for OpenSearch to\n become ready. Defaults to 60.\n client: The OpenSearch client to use for pinging. If None, a new dummy\n client will be created. Defaults to None.\n\n Returns:\n True if OpenSearch is ready, False otherwise.\n \"\"\"\n with nullcontext(client) if client else OpenSearchClient() as client:\n time_start = time.monotonic()\n while True:\n if client.ping():\n logger.info(\"[OpenSearch] Readiness probe succeeded. Continuing...\")\n return True\n time_elapsed = time.monotonic() - time_start\n if time_elapsed > wait_limit_s:\n logger.info(\n f\"[OpenSearch] Readiness probe did not succeed within the timeout ({wait_limit_s} seconds).\"\n )\n return False\n logger.info(\n f\"[OpenSearch] Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={wait_limit_s:.1f}\"\n )\n time.sleep(wait_interval_s)\n" }, { "path": "backend/onyx/document_index/opensearch/cluster_settings.py", "content": "from typing import Any\n\nOPENSEARCH_CLUSTER_SETTINGS: dict[str, Any] = {\n \"persistent\": {\n # By default, when you index a document to a non-existent index,\n # OpenSearch will automatically create the index. This behavior is\n # undesirable so this function exposes the ability to disable it.\n # See\n # https://docs.opensearch.org/latest/install-and-configure/configuring-opensearch/index/#updating-cluster-settings-using-the-api\n \"action.auto_create_index\": False,\n # Thresholds for OpenSearch to log slow queries at the server level.\n \"cluster.search.request.slowlog.level\": \"INFO\",\n \"cluster.search.request.slowlog.threshold.warn\": \"5s\",\n \"cluster.search.request.slowlog.threshold.info\": \"2s\",\n \"cluster.search.request.slowlog.threshold.debug\": \"1s\",\n \"cluster.search.request.slowlog.threshold.trace\": \"500ms\",\n }\n}\n" }, { "path": "backend/onyx/document_index/opensearch/constants.py", "content": "# Default value for the maximum number of tokens a chunk can hold, if none is\n# specified when creating an index.\nimport os\nfrom enum import Enum\n\n\nDEFAULT_MAX_CHUNK_SIZE = 512\n\n\n# By default OpenSearch will only return a maximum of this many results in a\n# given search. This value is configurable in the index settings.\nDEFAULT_OPENSEARCH_MAX_RESULT_WINDOW = 10_000\n\n\n# For documents which do not have a value for LAST_UPDATED_FIELD_NAME, we assume\n# that the document was last updated this many days ago for the purpose of time\n# cutoff filtering during retrieval.\nASSUMED_DOCUMENT_AGE_DAYS = 90\n\n\n# Size of the dynamic list used to consider elements during kNN graph creation.\n# Higher values improve search quality but increase indexing time. Values\n# typically range between 100 - 512.\nEF_CONSTRUCTION = 256\n# Number of bi-directional links per element. Higher values improve search\n# quality but increase memory footprint. Values typically range between 12 - 48.\nM = 32 # Set relatively high for better accuracy.\n\n# When performing hybrid search, we need to consider more candidates than the\n# number of results to be returned. This is because the scoring is hybrid and\n# the results are reordered due to the hybrid scoring. Higher = more candidates\n# for hybrid fusion = better retrieval accuracy, but results in more computation\n# per query. Imagine a simple case with a single keyword query and a single\n# vector query and we want 10 final docs. If we only fetch 10 candidates from\n# each of keyword and vector, they would have to have perfect overlap to get a\n# good hybrid ranking for the 10 results. If we fetch 1000 candidates from each,\n# we have a much higher chance of all 10 of the final desired docs showing up\n# and getting scored. In worse situations, the final 10 docs don't even show up\n# as the final 10 (worse than just a miss at the reranking step).\n# Defaults to 500 for now. Initially this defaulted to 750 but we were seeing\n# poor search performance; bumped from 100 to 500 to improve recall.\nDEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES = int(\n os.environ.get(\"DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES\", 500)\n)\n\n# Number of vectors to examine to decide the top k neighbors for the HNSW\n# method.\n# NOTE: \"When creating a search query, you must specify k. If you provide both k\n# and ef_search, then the larger value is passed to the engine. If ef_search is\n# larger than k, you can provide the size parameter to limit the final number of\n# results to k.\" from\n# https://docs.opensearch.org/latest/query-dsl/specialized/k-nn/index/#ef_search\nEF_SEARCH = DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES\n\n\nclass OpenSearchSearchType(str, Enum):\n \"\"\"Search type label used for Prometheus metrics.\"\"\"\n\n HYBRID = \"hybrid\"\n KEYWORD = \"keyword\"\n SEMANTIC = \"semantic\"\n RANDOM = \"random\"\n DOC_ID_RETRIEVAL = \"doc_id_retrieval\"\n UNKNOWN = \"unknown\"\n\n\nclass HybridSearchSubqueryConfiguration(Enum):\n TITLE_VECTOR_CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD = 1\n # Current default.\n CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD = 2\n\n\n# Will raise and block application start if HYBRID_SEARCH_SUBQUERY_CONFIGURATION\n# is set but not a valid value. If not set, defaults to\n# CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD.\nHYBRID_SEARCH_SUBQUERY_CONFIGURATION: HybridSearchSubqueryConfiguration = (\n HybridSearchSubqueryConfiguration(\n int(os.environ[\"HYBRID_SEARCH_SUBQUERY_CONFIGURATION\"])\n )\n if os.environ.get(\"HYBRID_SEARCH_SUBQUERY_CONFIGURATION\", None) is not None\n else HybridSearchSubqueryConfiguration.CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD\n)\n\n\nclass HybridSearchNormalizationPipeline(Enum):\n # Current default.\n MIN_MAX = 1\n # NOTE: Using z-score normalization is better for hybrid search from a\n # theoretical standpoint. Empirically on a small dataset of up to 10K docs,\n # it's not very different. Likely more impactful at scale.\n # https://opensearch.org/blog/introducing-the-z-score-normalization-technique-for-hybrid-search/\n ZSCORE = 2\n\n\n# Will raise and block application start if HYBRID_SEARCH_NORMALIZATION_PIPELINE\n# is set but not a valid value. If not set, defaults to MIN_MAX.\nHYBRID_SEARCH_NORMALIZATION_PIPELINE: HybridSearchNormalizationPipeline = (\n HybridSearchNormalizationPipeline(\n int(os.environ[\"HYBRID_SEARCH_NORMALIZATION_PIPELINE\"])\n )\n if os.environ.get(\"HYBRID_SEARCH_NORMALIZATION_PIPELINE\", None) is not None\n else HybridSearchNormalizationPipeline.MIN_MAX\n)\n" }, { "path": "backend/onyx/document_index/opensearch/opensearch_document_index.py", "content": "import json\nfrom collections.abc import Iterable\nfrom typing import Any\n\nimport httpx\nfrom opensearchpy import NotFoundError\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.configs.app_configs import MAX_CHUNKS_PER_DOC_BATCH\nfrom onyx.configs.app_configs import VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT\nfrom onyx.configs.chat_configs import NUM_RETURNED_HITS\nfrom onyx.configs.chat_configs import TITLE_CONTENT_RATIO\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n get_experts_stores_representations,\n)\nfrom onyx.connectors.models import convert_metadata_list_of_strings_to_dict\nfrom onyx.context.search.enums import QueryType\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceChunkUncleaned\nfrom onyx.context.search.models import QueryExpansionType\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.db.models import DocumentSource\nfrom onyx.document_index.chunk_content_enrichment import cleanup_content_for_chunks\nfrom onyx.document_index.chunk_content_enrichment import (\n generate_enriched_content_for_chunk_text,\n)\nfrom onyx.document_index.interfaces import DocumentIndex as OldDocumentIndex\nfrom onyx.document_index.interfaces import (\n DocumentInsertionRecord as OldDocumentInsertionRecord,\n)\nfrom onyx.document_index.interfaces import IndexBatchParams\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\nfrom onyx.document_index.interfaces_new import DocumentIndex\nfrom onyx.document_index.interfaces_new import DocumentInsertionRecord\nfrom onyx.document_index.interfaces_new import DocumentSectionRequest\nfrom onyx.document_index.interfaces_new import IndexingMetadata\nfrom onyx.document_index.interfaces_new import MetadataUpdateRequest\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.client import OpenSearchClient\nfrom onyx.document_index.opensearch.client import OpenSearchIndexClient\nfrom onyx.document_index.opensearch.client import SearchHit\nfrom onyx.document_index.opensearch.cluster_settings import OPENSEARCH_CLUSTER_SETTINGS\nfrom onyx.document_index.opensearch.constants import OpenSearchSearchType\nfrom onyx.document_index.opensearch.schema import ACCESS_CONTROL_LIST_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import CONTENT_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DOCUMENT_SETS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DocumentChunk\nfrom onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors\nfrom onyx.document_index.opensearch.schema import DocumentSchema\nfrom onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id\nfrom onyx.document_index.opensearch.schema import GLOBAL_BOOST_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import HIDDEN_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import PERSONAS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import USER_PROJECTS_FIELD_NAME\nfrom onyx.document_index.opensearch.search import DocumentQuery\nfrom onyx.document_index.opensearch.search import (\n get_min_max_normalization_pipeline_name_and_config,\n)\nfrom onyx.document_index.opensearch.search import (\n get_normalization_pipeline_name_and_config,\n)\nfrom onyx.document_index.opensearch.search import (\n get_zscore_normalization_pipeline_name_and_config,\n)\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.indexing.models import Document\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.text_processing import remove_invalid_unicode_chars\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom shared_configs.model_server_models import Embedding\n\n\nlogger = setup_logger(__name__)\n\n\nclass ChunkCountNotFoundError(ValueError):\n \"\"\"Raised when a document has no chunk count.\"\"\"\n\n\ndef generate_opensearch_filtered_access_control_list(\n access: DocumentAccess,\n) -> list[str]:\n \"\"\"Generates an access control list with PUBLIC_DOC_PAT removed.\n\n In the OpenSearch schema this is represented by PUBLIC_FIELD_NAME.\n \"\"\"\n access_control_list = access.to_acl()\n access_control_list.discard(PUBLIC_DOC_PAT)\n return list(access_control_list)\n\n\ndef set_cluster_state(client: OpenSearchClient) -> None:\n if not client.put_cluster_settings(settings=OPENSEARCH_CLUSTER_SETTINGS):\n logger.error(\n \"Failed to put cluster settings. If the settings have never been set before, \"\n \"this may cause unexpected index creation when indexing documents into an \"\n \"index that does not exist, or may cause expected logs to not appear. If this \"\n \"is not the first time running Onyx against this instance of OpenSearch, these \"\n \"settings have likely already been set. Not taking any further action...\"\n )\n min_max_normalization_pipeline_name, min_max_normalization_pipeline_config = (\n get_min_max_normalization_pipeline_name_and_config()\n )\n zscore_normalization_pipeline_name, zscore_normalization_pipeline_config = (\n get_zscore_normalization_pipeline_name_and_config()\n )\n client.create_search_pipeline(\n pipeline_id=min_max_normalization_pipeline_name,\n pipeline_body=min_max_normalization_pipeline_config,\n )\n client.create_search_pipeline(\n pipeline_id=zscore_normalization_pipeline_name,\n pipeline_body=zscore_normalization_pipeline_config,\n )\n\n\ndef _convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned(\n chunk: DocumentChunkWithoutVectors,\n score: float | None,\n highlights: dict[str, list[str]],\n) -> InferenceChunkUncleaned:\n \"\"\"\n Generates an inference chunk from an OpenSearch document chunk, its score,\n and its match highlights.\n\n Args:\n chunk: The document chunk returned by OpenSearch.\n score: The document chunk match score as calculated by OpenSearch. Only\n relevant for searches like hybrid search. It is acceptable for this\n value to be None for results from other queries like ID-based\n retrieval as a match score makes no sense in those contexts.\n highlights: Maps schema property name to a list of highlighted snippets\n with match terms wrapped in tags (e.g. \"something keyword\n other thing\").\n\n Returns:\n An Onyx inference chunk representation.\n \"\"\"\n return InferenceChunkUncleaned(\n chunk_id=chunk.chunk_index,\n blurb=chunk.blurb,\n # Includes extra content prepended/appended during indexing.\n content=chunk.content,\n # When we read a string and turn it into a dict the keys will be\n # strings, but in this case they need to be ints.\n source_links=(\n {int(k): v for k, v in json.loads(chunk.source_links).items()}\n if chunk.source_links\n else None\n ),\n image_file_id=chunk.image_file_id,\n # Deprecated. Fill in some reasonable default.\n section_continuation=False,\n document_id=chunk.document_id,\n source_type=DocumentSource(chunk.source_type),\n semantic_identifier=chunk.semantic_identifier,\n title=chunk.title,\n boost=chunk.global_boost,\n score=score,\n hidden=chunk.hidden,\n metadata=(\n convert_metadata_list_of_strings_to_dict(chunk.metadata_list)\n if chunk.metadata_list\n else {}\n ),\n # Extract highlighted snippets from the content field, if available. In\n # the future we may want to match on other fields too, currently we only\n # use the content field.\n match_highlights=highlights.get(CONTENT_FIELD_NAME, []),\n # TODO(andrei) Consider storing a chunk content index instead of a full\n # string when working on chunk content augmentation.\n doc_summary=chunk.doc_summary,\n # TODO(andrei) Same thing as above.\n chunk_context=chunk.chunk_context,\n updated_at=chunk.last_updated,\n primary_owners=chunk.primary_owners,\n secondary_owners=chunk.secondary_owners,\n # TODO(andrei) Same thing as chunk_context above.\n metadata_suffix=chunk.metadata_suffix,\n )\n\n\ndef _convert_onyx_chunk_to_opensearch_document(\n chunk: DocMetadataAwareIndexChunk,\n) -> DocumentChunk:\n filtered_blurb = remove_invalid_unicode_chars(chunk.blurb)\n _title = chunk.source_document.get_title_for_document_index()\n filtered_title = remove_invalid_unicode_chars(_title) if _title else None\n filtered_content = remove_invalid_unicode_chars(\n generate_enriched_content_for_chunk_text(chunk)\n )\n filtered_semantic_identifier = remove_invalid_unicode_chars(\n chunk.source_document.semantic_identifier\n )\n filtered_metadata_suffix = remove_invalid_unicode_chars(\n chunk.metadata_suffix_keyword\n )\n _metadata_list = chunk.source_document.get_metadata_str_attributes()\n filtered_metadata_list = (\n [remove_invalid_unicode_chars(metadata) for metadata in _metadata_list]\n if _metadata_list\n else None\n )\n return DocumentChunk(\n document_id=chunk.source_document.id,\n chunk_index=chunk.chunk_id,\n # Use get_title_for_document_index to match the logic used when creating\n # the title_embedding in the embedder. This method falls back to\n # semantic_identifier when title is None (but not empty string).\n title=filtered_title,\n title_vector=chunk.title_embedding,\n content=filtered_content,\n content_vector=chunk.embeddings.full_embedding,\n source_type=chunk.source_document.source.value,\n metadata_list=filtered_metadata_list,\n metadata_suffix=filtered_metadata_suffix,\n last_updated=chunk.source_document.doc_updated_at,\n public=chunk.access.is_public,\n access_control_list=generate_opensearch_filtered_access_control_list(\n chunk.access\n ),\n global_boost=chunk.boost,\n semantic_identifier=filtered_semantic_identifier,\n image_file_id=chunk.image_file_id,\n # Small optimization, if this list is empty we can supply None to\n # OpenSearch and it will not store any data at all for this field, which\n # is different from supplying an empty list.\n source_links=json.dumps(chunk.source_links) if chunk.source_links else None,\n blurb=filtered_blurb,\n doc_summary=chunk.doc_summary,\n chunk_context=chunk.chunk_context,\n # Small optimization, if this list is empty we can supply None to\n # OpenSearch and it will not store any data at all for this field, which\n # is different from supplying an empty list.\n document_sets=list(chunk.document_sets) if chunk.document_sets else None,\n # Small optimization, if this list is empty we can supply None to\n # OpenSearch and it will not store any data at all for this field, which\n # is different from supplying an empty list.\n user_projects=chunk.user_project or None,\n personas=chunk.personas or None,\n primary_owners=get_experts_stores_representations(\n chunk.source_document.primary_owners\n ),\n secondary_owners=get_experts_stores_representations(\n chunk.source_document.secondary_owners\n ),\n # TODO(andrei): Consider not even getting this from\n # DocMetadataAwareIndexChunk and instead using OpenSearchDocumentIndex's\n # instance variable. One source of truth -> less chance of a very bad\n # bug in prod.\n tenant_id=TenantState(tenant_id=chunk.tenant_id, multitenant=MULTI_TENANT),\n # Store ancestor hierarchy node IDs for hierarchy-based filtering.\n ancestor_hierarchy_node_ids=chunk.ancestor_hierarchy_node_ids or None,\n )\n\n\nclass OpenSearchOldDocumentIndex(OldDocumentIndex):\n \"\"\"\n Wrapper for OpenSearch to adapt the new DocumentIndex interface with\n invocations to the old DocumentIndex interface in the hotpath.\n\n The analogous class for Vespa is VespaIndex which calls to\n VespaDocumentIndex.\n\n TODO(andrei): This is very dumb and purely temporary until there are no more\n references to the old interface in the hotpath.\n \"\"\"\n\n def __init__(\n self,\n index_name: str,\n embedding_dim: int,\n embedding_precision: EmbeddingPrecision,\n secondary_index_name: str | None,\n secondary_embedding_dim: int | None,\n secondary_embedding_precision: EmbeddingPrecision | None,\n # NOTE: We do not support large chunks right now.\n large_chunks_enabled: bool, # noqa: ARG002\n secondary_large_chunks_enabled: bool | None, # noqa: ARG002\n multitenant: bool = False,\n httpx_client: httpx.Client | None = None, # noqa: ARG002\n ) -> None:\n super().__init__(\n index_name=index_name,\n secondary_index_name=secondary_index_name,\n )\n if multitenant != MULTI_TENANT:\n raise ValueError(\n \"Bug: Multitenant mismatch when initializing an OpenSearchDocumentIndex. \"\n f\"Expected {MULTI_TENANT}, got {multitenant}.\"\n )\n tenant_id = get_current_tenant_id()\n tenant_state = TenantState(tenant_id=tenant_id, multitenant=multitenant)\n self._real_index = OpenSearchDocumentIndex(\n tenant_state=tenant_state,\n index_name=index_name,\n embedding_dim=embedding_dim,\n embedding_precision=embedding_precision,\n )\n self._secondary_real_index: OpenSearchDocumentIndex | None = None\n if self.secondary_index_name:\n if secondary_embedding_dim is None or secondary_embedding_precision is None:\n raise ValueError(\n \"Bug: Secondary index embedding dimension and precision are not set.\"\n )\n self._secondary_real_index = OpenSearchDocumentIndex(\n tenant_state=tenant_state,\n index_name=self.secondary_index_name,\n embedding_dim=secondary_embedding_dim,\n embedding_precision=secondary_embedding_precision,\n )\n\n @staticmethod\n def register_multitenant_indices(\n indices: list[str],\n embedding_dims: list[int],\n embedding_precisions: list[EmbeddingPrecision],\n ) -> None:\n raise NotImplementedError(\n \"Bug: Multitenant index registration is not supported for OpenSearch.\"\n )\n\n def ensure_indices_exist(\n self,\n primary_embedding_dim: int,\n primary_embedding_precision: EmbeddingPrecision,\n secondary_index_embedding_dim: int | None,\n secondary_index_embedding_precision: EmbeddingPrecision | None,\n ) -> None:\n self._real_index.verify_and_create_index_if_necessary(\n primary_embedding_dim, primary_embedding_precision\n )\n if self.secondary_index_name:\n if (\n secondary_index_embedding_dim is None\n or secondary_index_embedding_precision is None\n ):\n raise ValueError(\n \"Bug: Secondary index embedding dimension and precision are not set.\"\n )\n assert (\n self._secondary_real_index is not None\n ), \"Bug: Secondary index is not initialized.\"\n self._secondary_real_index.verify_and_create_index_if_necessary(\n secondary_index_embedding_dim, secondary_index_embedding_precision\n )\n\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk],\n index_batch_params: IndexBatchParams,\n ) -> set[OldDocumentInsertionRecord]:\n \"\"\"\n NOTE: Do NOT consider the secondary index here. A separate indexing\n pipeline will be responsible for indexing to the secondary index. This\n design is not ideal and we should reconsider this when revamping index\n swapping.\n \"\"\"\n # Convert IndexBatchParams to IndexingMetadata.\n chunk_counts: dict[str, IndexingMetadata.ChunkCounts] = {}\n for doc_id in index_batch_params.doc_id_to_new_chunk_cnt:\n old_count = index_batch_params.doc_id_to_previous_chunk_cnt[doc_id]\n new_count = index_batch_params.doc_id_to_new_chunk_cnt[doc_id]\n chunk_counts[doc_id] = IndexingMetadata.ChunkCounts(\n old_chunk_cnt=old_count,\n new_chunk_cnt=new_count,\n )\n\n indexing_metadata = IndexingMetadata(doc_id_to_chunk_cnt_diff=chunk_counts)\n\n results = self._real_index.index(chunks, indexing_metadata)\n\n # Convert list[DocumentInsertionRecord] to\n # set[OldDocumentInsertionRecord].\n return {\n OldDocumentInsertionRecord(\n document_id=record.document_id,\n already_existed=record.already_existed,\n )\n for record in results\n }\n\n def delete_single(\n self,\n doc_id: str,\n *,\n tenant_id: str, # noqa: ARG002\n chunk_count: int | None,\n ) -> int:\n \"\"\"\n NOTE: Remember to handle the secondary index here. There is no separate\n pipeline for deleting chunks in the secondary index. This design is not\n ideal and we should reconsider this when revamping index swapping.\n \"\"\"\n total_chunks_deleted = self._real_index.delete(doc_id, chunk_count)\n if self.secondary_index_name:\n assert (\n self._secondary_real_index is not None\n ), \"Bug: Secondary index is not initialized.\"\n total_chunks_deleted += self._secondary_real_index.delete(\n doc_id, chunk_count\n )\n return total_chunks_deleted\n\n def update_single(\n self,\n doc_id: str,\n *,\n tenant_id: str, # noqa: ARG002\n chunk_count: int | None,\n fields: VespaDocumentFields | None,\n user_fields: VespaDocumentUserFields | None,\n ) -> None:\n \"\"\"\n NOTE: Remember to handle the secondary index here. There is no separate\n pipeline for updating chunks in the secondary index. This design is not\n ideal and we should reconsider this when revamping index swapping.\n \"\"\"\n if fields is None and user_fields is None:\n logger.warning(\n f\"Tried to update document {doc_id} with no updated fields or user fields.\"\n )\n return\n\n # Convert VespaDocumentFields to MetadataUpdateRequest.\n update_request = MetadataUpdateRequest(\n document_ids=[doc_id],\n doc_id_to_chunk_cnt={\n doc_id: chunk_count if chunk_count is not None else -1\n },\n access=fields.access if fields else None,\n document_sets=fields.document_sets if fields else None,\n boost=fields.boost if fields else None,\n hidden=fields.hidden if fields else None,\n project_ids=(\n set(user_fields.user_projects)\n # NOTE: Empty user_projects is semantically different from None\n # user_projects.\n if user_fields and user_fields.user_projects is not None\n else None\n ),\n persona_ids=(\n set(user_fields.personas)\n # NOTE: Empty personas is semantically different from None\n # personas.\n if user_fields and user_fields.personas is not None\n else None\n ),\n )\n\n try:\n self._real_index.update([update_request])\n if self.secondary_index_name:\n assert (\n self._secondary_real_index is not None\n ), \"Bug: Secondary index is not initialized.\"\n self._secondary_real_index.update([update_request])\n except NotFoundError:\n logger.exception(\n f\"Tried to update document {doc_id} but at least one of its chunks was not found in OpenSearch. \"\n \"This is likely due to it not having been indexed yet. Skipping update for now...\"\n )\n return\n except ChunkCountNotFoundError:\n logger.exception(\n f\"Tried to update document {doc_id} but its chunk count is not known. We tolerate this for now \"\n \"but this will not be an acceptable state once OpenSearch is the primary document index and the \"\n \"indexing/updating race condition is fixed.\"\n )\n return\n\n def id_based_retrieval(\n self,\n chunk_requests: list[VespaChunkRequest],\n filters: IndexFilters,\n batch_retrieval: bool = False,\n get_large_chunks: bool = False, # noqa: ARG002\n ) -> list[InferenceChunk]:\n section_requests = [\n DocumentSectionRequest(\n document_id=req.document_id,\n min_chunk_ind=req.min_chunk_ind,\n max_chunk_ind=req.max_chunk_ind,\n )\n for req in chunk_requests\n ]\n\n return self._real_index.id_based_retrieval(\n section_requests, filters, batch_retrieval\n )\n\n def hybrid_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n final_keywords: list[str] | None,\n filters: IndexFilters,\n hybrid_alpha: float,\n time_decay_multiplier: float, # noqa: ARG002\n num_to_retrieve: int,\n ranking_profile_type: QueryExpansionType = QueryExpansionType.SEMANTIC, # noqa: ARG002\n title_content_ratio: float | None = TITLE_CONTENT_RATIO, # noqa: ARG002\n ) -> list[InferenceChunk]:\n # Determine query type based on hybrid_alpha.\n if hybrid_alpha >= 0.8:\n query_type = QueryType.SEMANTIC\n elif hybrid_alpha <= 0.2:\n query_type = QueryType.KEYWORD\n else:\n query_type = QueryType.SEMANTIC # Default to semantic for hybrid.\n\n return self._real_index.hybrid_retrieval(\n query=query,\n query_embedding=query_embedding,\n final_keywords=final_keywords,\n query_type=query_type,\n filters=filters,\n num_to_retrieve=num_to_retrieve,\n )\n\n def admin_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n filters: IndexFilters,\n num_to_retrieve: int = NUM_RETURNED_HITS,\n ) -> list[InferenceChunk]:\n return self._real_index.hybrid_retrieval(\n query=query,\n query_embedding=query_embedding,\n final_keywords=None,\n query_type=QueryType.KEYWORD,\n filters=filters,\n num_to_retrieve=num_to_retrieve,\n )\n\n def random_retrieval(\n self,\n filters: IndexFilters,\n num_to_retrieve: int = 10,\n ) -> list[InferenceChunk]:\n return self._real_index.random_retrieval(\n filters=filters,\n num_to_retrieve=num_to_retrieve,\n dirty=None,\n )\n\n\nclass OpenSearchDocumentIndex(DocumentIndex):\n \"\"\"OpenSearch-specific implementation of the DocumentIndex interface.\n\n This class provides document indexing, retrieval, and management operations\n for an OpenSearch search engine instance. It handles the complete lifecycle\n of document chunks within a specific OpenSearch index/schema.\n\n Each kind of embedding used should correspond to a different instance of\n this class, and therefore a different index in OpenSearch.\n\n If in a multitenant environment and\n VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT, will verify and create the index\n if necessary on initialization. This is because there is no logic which runs\n on cluster restart which scans through all search settings over all tenants\n and creates the relevant indices.\n\n Args:\n tenant_state: The tenant state of the caller.\n index_name: The name of the index to interact with.\n embedding_dim: The dimensionality of the embeddings used for the index.\n embedding_precision: The precision of the embeddings used for the index.\n \"\"\"\n\n def __init__(\n self,\n tenant_state: TenantState,\n index_name: str,\n embedding_dim: int,\n embedding_precision: EmbeddingPrecision,\n ) -> None:\n self._index_name: str = index_name\n self._tenant_state: TenantState = tenant_state\n self._client = OpenSearchIndexClient(index_name=self._index_name)\n\n if self._tenant_state.multitenant and VERIFY_CREATE_OPENSEARCH_INDEX_ON_INIT_MT:\n self.verify_and_create_index_if_necessary(\n embedding_dim=embedding_dim, embedding_precision=embedding_precision\n )\n\n def verify_and_create_index_if_necessary(\n self,\n embedding_dim: int,\n embedding_precision: EmbeddingPrecision, # noqa: ARG002\n ) -> None:\n \"\"\"Verifies and creates the index if necessary.\n\n Also puts the desired cluster settings if not in a multitenant\n environment.\n\n Also puts the desired search pipeline state if not in a multitenant\n environment, creating the pipelines if they do not exist and updating\n them otherwise.\n\n In a multitenant environment, the above steps happen explicitly on\n setup.\n\n Args:\n embedding_dim: Vector dimensionality for the vector similarity part\n of the search.\n embedding_precision: Precision of the values of the vectors for the\n similarity part of the search.\n\n Raises:\n Exception: There was an error verifying or creating the index or\n search pipelines.\n \"\"\"\n logger.debug(\n f\"[OpenSearchDocumentIndex] Verifying and creating index {self._index_name} if \"\n f\"necessary, with embedding dimension {embedding_dim}.\"\n )\n\n if not self._tenant_state.multitenant:\n set_cluster_state(self._client)\n\n expected_mappings = DocumentSchema.get_document_schema(\n embedding_dim, self._tenant_state.multitenant\n )\n\n if not self._client.index_exists():\n index_settings = DocumentSchema.get_index_settings_based_on_environment()\n self._client.create_index(\n mappings=expected_mappings,\n settings=index_settings,\n )\n else:\n # Ensure schema is up to date by applying the current mappings.\n try:\n self._client.put_mapping(expected_mappings)\n except Exception as e:\n logger.error(\n f\"Failed to update mappings for index {self._index_name}. This likely means a \"\n f\"field type was changed which requires reindexing. Error: {e}\"\n )\n raise\n\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk],\n indexing_metadata: IndexingMetadata,\n ) -> list[DocumentInsertionRecord]:\n \"\"\"Indexes an iterable of document chunks into the document index.\n\n Groups chunks by document ID and for each document, deletes existing\n chunks and indexes the new chunks in bulk.\n\n NOTE: It is assumed that chunks for a given document are not spread out\n over multiple index() calls.\n\n Args:\n chunks: Document chunks with all of the information needed for\n indexing to the document index.\n indexing_metadata: Information about chunk counts for efficient\n cleaning / updating.\n\n Raises:\n Exception: Failed to index some or all of the chunks for the\n specified documents.\n\n Returns:\n List of document IDs which map to unique documents as well as if the\n document is newly indexed or had already existed and was just\n updated.\n \"\"\"\n total_chunks = sum(\n cc.new_chunk_cnt\n for cc in indexing_metadata.doc_id_to_chunk_cnt_diff.values()\n )\n logger.debug(\n f\"[OpenSearchDocumentIndex] Indexing {total_chunks} chunks from {len(indexing_metadata.doc_id_to_chunk_cnt_diff)} \"\n f\"documents for index {self._index_name}.\"\n )\n\n document_indexing_results: list[DocumentInsertionRecord] = []\n deleted_doc_ids: set[str] = set()\n # Buffer chunks per document as they arrive from the iterable.\n # When the document ID changes flush the buffered chunks.\n current_doc_id: str | None = None\n current_chunks: list[DocMetadataAwareIndexChunk] = []\n\n def _flush_chunks(doc_chunks: list[DocMetadataAwareIndexChunk]) -> None:\n assert len(doc_chunks) > 0, \"doc_chunks is empty\"\n\n # Create a batch of OpenSearch-formatted chunks for bulk insertion.\n # Since we are doing this in batches, an error occurring midway\n # can result in a state where chunks are deleted and not all the\n # new chunks have been indexed.\n chunk_batch: list[DocumentChunk] = [\n _convert_onyx_chunk_to_opensearch_document(chunk)\n for chunk in doc_chunks\n ]\n onyx_document: Document = doc_chunks[0].source_document\n # First delete the doc's chunks from the index. This is so that\n # there are no dangling chunks in the index, in the event that the\n # new document's content contains fewer chunks than the previous\n # content.\n # TODO(andrei): This can possibly be made more efficient by checking\n # if the chunk count has actually decreased. This assumes that\n # overlapping chunks are perfectly overwritten. If we can't\n # guarantee that then we need the code as-is.\n if onyx_document.id not in deleted_doc_ids:\n num_chunks_deleted = self.delete(\n onyx_document.id, onyx_document.chunk_count\n )\n deleted_doc_ids.add(onyx_document.id)\n # If we see that chunks were deleted we assume the doc already\n # existed. We record the result before bulk_index_documents\n # runs. If indexing raises, this entire result list is discarded\n # by the caller's retry logic, so early recording is safe.\n document_indexing_results.append(\n DocumentInsertionRecord(\n document_id=onyx_document.id,\n already_existed=num_chunks_deleted > 0,\n )\n )\n # Now index. This will raise if a chunk of the same ID exists, which\n # we do not expect because we should have deleted all chunks.\n self._client.bulk_index_documents(\n documents=chunk_batch,\n tenant_state=self._tenant_state,\n )\n\n for chunk in chunks:\n doc_id = chunk.source_document.id\n if doc_id != current_doc_id:\n if current_chunks:\n _flush_chunks(current_chunks)\n current_doc_id = doc_id\n current_chunks = [chunk]\n elif len(current_chunks) >= MAX_CHUNKS_PER_DOC_BATCH:\n _flush_chunks(current_chunks)\n current_chunks = [chunk]\n else:\n current_chunks.append(chunk)\n\n if current_chunks:\n _flush_chunks(current_chunks)\n\n return document_indexing_results\n\n def delete(\n self,\n document_id: str,\n chunk_count: int | None = None, # noqa: ARG002\n ) -> int:\n \"\"\"Deletes all chunks for a given document.\n\n Does nothing if the specified document ID does not exist.\n\n TODO(andrei): Consider implementing this method to delete on document\n chunk IDs vs querying for matching document chunks. Unclear if this is\n any better though.\n\n Args:\n document_id: The unique identifier for the document as represented\n in Onyx, not necessarily in the document index.\n chunk_count: The number of chunks in OpenSearch for the document.\n Defaults to None.\n\n Raises:\n Exception: Failed to delete some or all of the chunks for the\n document.\n\n Returns:\n The number of chunks successfully deleted.\n \"\"\"\n logger.debug(\n f\"[OpenSearchDocumentIndex] Deleting document {document_id} from index {self._index_name}.\"\n )\n query_body = DocumentQuery.delete_from_document_id_query(\n document_id=document_id,\n tenant_state=self._tenant_state,\n )\n\n return self._client.delete_by_query(query_body)\n\n def update(\n self,\n update_requests: list[MetadataUpdateRequest],\n ) -> None:\n \"\"\"Updates some set of chunks.\n\n NOTE: Will raise if one of the specified document chunks do not exist.\n This may be due to a concurrent ongoing indexing operation. In that\n event callers are expected to retry after a bit once the state of the\n document index is updated.\n NOTE: Requires document chunk count be known; will raise if it is not.\n This may be caused by the same situation outlined above.\n NOTE: Will no-op if an update request has no fields to update.\n\n TODO(andrei): Consider exploring a batch API for OpenSearch for this\n operation.\n\n Args:\n update_requests: A list of update requests, each containing a list\n of document IDs and the fields to update. The field updates\n apply to all of the specified documents in each update request.\n\n Raises:\n Exception: Failed to update some or all of the chunks for the\n specified documents.\n \"\"\"\n logger.debug(\n f\"[OpenSearchDocumentIndex] Updating {len(update_requests)} chunks for index {self._index_name}.\"\n )\n for update_request in update_requests:\n properties_to_update: dict[str, Any] = dict()\n # TODO(andrei): Nit but consider if we can use DocumentChunk\n # here so we don't have to think about passing in the\n # appropriate types into this dict.\n if update_request.access is not None:\n properties_to_update[ACCESS_CONTROL_LIST_FIELD_NAME] = (\n generate_opensearch_filtered_access_control_list(\n update_request.access\n )\n )\n if update_request.document_sets is not None:\n properties_to_update[DOCUMENT_SETS_FIELD_NAME] = list(\n update_request.document_sets\n )\n if update_request.boost is not None:\n properties_to_update[GLOBAL_BOOST_FIELD_NAME] = int(\n update_request.boost\n )\n if update_request.hidden is not None:\n properties_to_update[HIDDEN_FIELD_NAME] = update_request.hidden\n if update_request.project_ids is not None:\n properties_to_update[USER_PROJECTS_FIELD_NAME] = list(\n update_request.project_ids\n )\n if update_request.persona_ids is not None:\n properties_to_update[PERSONAS_FIELD_NAME] = list(\n update_request.persona_ids\n )\n\n if not properties_to_update:\n if len(update_request.document_ids) > 1:\n update_string = f\"{len(update_request.document_ids)} documents\"\n else:\n update_string = f\"document {update_request.document_ids[0]}\"\n logger.warning(\n f\"[OpenSearchDocumentIndex] Tried to update {update_string} \"\n \"with no specified update fields. This will be a no-op.\"\n )\n continue\n\n for doc_id in update_request.document_ids:\n doc_chunk_count = update_request.doc_id_to_chunk_cnt.get(doc_id, -1)\n if doc_chunk_count < 0:\n # This means the chunk count is not known. This is due to a\n # race condition between doc indexing and updating steps\n # which run concurrently when a doc is indexed. The indexing\n # step should update chunk count shortly. This could also\n # have been due to an older version of the indexing pipeline\n # which did not compute chunk count, but that codepath has\n # since been deprecated and should no longer be the case\n # here.\n # TODO(andrei): Fix the aforementioned race condition.\n raise ChunkCountNotFoundError(\n f\"Tried to update document {doc_id} but its chunk count is not known. \"\n \"Older versions of the application used to permit this but is not a \"\n \"supported state for a document when using OpenSearch. The document was \"\n \"likely just added to the indexing pipeline and the chunk count will be \"\n \"updated shortly.\"\n )\n if doc_chunk_count == 0:\n raise ValueError(\n f\"Bug: Tried to update document {doc_id} but its chunk count was 0.\"\n )\n\n for chunk_index in range(doc_chunk_count):\n document_chunk_id = get_opensearch_doc_chunk_id(\n tenant_state=self._tenant_state,\n document_id=doc_id,\n chunk_index=chunk_index,\n )\n self._client.update_document(\n document_chunk_id=document_chunk_id,\n properties_to_update=properties_to_update,\n )\n\n def id_based_retrieval(\n self,\n chunk_requests: list[DocumentSectionRequest],\n filters: IndexFilters,\n # TODO(andrei): Remove this from the new interface at some point; we\n # should not be exposing this.\n batch_retrieval: bool = False, # noqa: ARG002\n # TODO(andrei): Add a param for whether to retrieve hidden docs.\n ) -> list[InferenceChunk]:\n \"\"\"\n TODO(andrei): Consider implementing this method to retrieve on document\n chunk IDs vs querying for matching document chunks.\n \"\"\"\n logger.debug(\n f\"[OpenSearchDocumentIndex] Retrieving {len(chunk_requests)} chunks for index {self._index_name}.\"\n )\n results: list[InferenceChunk] = []\n for chunk_request in chunk_requests:\n search_hits: list[SearchHit[DocumentChunkWithoutVectors]] = []\n query_body = DocumentQuery.get_from_document_id_query(\n document_id=chunk_request.document_id,\n tenant_state=self._tenant_state,\n # NOTE: Index filters includes metadata tags which were filtered\n # for invalid unicode at indexing time. In theory it would be\n # ideal to do filtering here as well, in practice we never did\n # that in the Vespa codepath and have not seen issues in\n # production, so we deliberately conform to the existing logic\n # in order to not unknowningly introduce a possible bug.\n index_filters=filters,\n include_hidden=False,\n max_chunk_size=chunk_request.max_chunk_size,\n min_chunk_index=chunk_request.min_chunk_ind,\n max_chunk_index=chunk_request.max_chunk_ind,\n )\n search_hits = self._client.search(\n body=query_body,\n search_pipeline_id=None,\n search_type=OpenSearchSearchType.DOC_ID_RETRIEVAL,\n )\n inference_chunks_uncleaned: list[InferenceChunkUncleaned] = [\n _convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned(\n search_hit.document_chunk, None, {}\n )\n for search_hit in search_hits\n ]\n inference_chunks: list[InferenceChunk] = cleanup_content_for_chunks(\n inference_chunks_uncleaned\n )\n results.extend(inference_chunks)\n return results\n\n def hybrid_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n # TODO(andrei): This param is not great design, get rid of it.\n final_keywords: list[str] | None,\n query_type: QueryType, # noqa: ARG002\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n # TODO(andrei): There is some duplicated logic in this function with\n # others in this file.\n logger.debug(\n f\"[OpenSearchDocumentIndex] Hybrid retrieving {num_to_retrieve} chunks for index {self._index_name}.\"\n )\n # TODO(andrei): This could be better, the caller should just make this\n # decision when passing in the query param. See the above comment in the\n # function signature.\n final_query = \" \".join(final_keywords) if final_keywords else query\n query_body = DocumentQuery.get_hybrid_search_query(\n query_text=final_query,\n query_vector=query_embedding,\n num_hits=num_to_retrieve,\n tenant_state=self._tenant_state,\n # NOTE: Index filters includes metadata tags which were filtered\n # for invalid unicode at indexing time. In theory it would be\n # ideal to do filtering here as well, in practice we never did\n # that in the Vespa codepath and have not seen issues in\n # production, so we deliberately conform to the existing logic\n # in order to not unknowningly introduce a possible bug.\n index_filters=filters,\n include_hidden=False,\n )\n normalization_pipeline_name, _ = get_normalization_pipeline_name_and_config()\n search_hits: list[SearchHit[DocumentChunkWithoutVectors]] = self._client.search(\n body=query_body,\n search_pipeline_id=normalization_pipeline_name,\n search_type=OpenSearchSearchType.HYBRID,\n )\n\n # Good place for a breakpoint to inspect the search hits if you have\n # \"explain\" enabled.\n inference_chunks_uncleaned: list[InferenceChunkUncleaned] = [\n _convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned(\n search_hit.document_chunk, search_hit.score, search_hit.match_highlights\n )\n for search_hit in search_hits\n ]\n inference_chunks: list[InferenceChunk] = cleanup_content_for_chunks(\n inference_chunks_uncleaned\n )\n\n return inference_chunks\n\n def keyword_retrieval(\n self,\n query: str,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n # TODO(andrei): There is some duplicated logic in this function with\n # others in this file.\n logger.debug(\n f\"[OpenSearchDocumentIndex] Keyword retrieving {num_to_retrieve} chunks for index {self._index_name}.\"\n )\n query_body = DocumentQuery.get_keyword_search_query(\n query_text=query,\n num_hits=num_to_retrieve,\n tenant_state=self._tenant_state,\n # NOTE: Index filters includes metadata tags which were filtered\n # for invalid unicode at indexing time. In theory it would be\n # ideal to do filtering here as well, in practice we never did\n # that in the Vespa codepath and have not seen issues in\n # production, so we deliberately conform to the existing logic\n # in order to not unknowningly introduce a possible bug.\n index_filters=filters,\n include_hidden=False,\n )\n search_hits: list[SearchHit[DocumentChunkWithoutVectors]] = self._client.search(\n body=query_body,\n search_pipeline_id=None,\n search_type=OpenSearchSearchType.KEYWORD,\n )\n\n inference_chunks_uncleaned: list[InferenceChunkUncleaned] = [\n _convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned(\n search_hit.document_chunk, search_hit.score, search_hit.match_highlights\n )\n for search_hit in search_hits\n ]\n inference_chunks: list[InferenceChunk] = cleanup_content_for_chunks(\n inference_chunks_uncleaned\n )\n\n return inference_chunks\n\n def semantic_retrieval(\n self,\n query_embedding: Embedding,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n # TODO(andrei): There is some duplicated logic in this function with\n # others in this file.\n logger.debug(\n f\"[OpenSearchDocumentIndex] Semantic retrieving {num_to_retrieve} chunks for index {self._index_name}.\"\n )\n query_body = DocumentQuery.get_semantic_search_query(\n query_embedding=query_embedding,\n num_hits=num_to_retrieve,\n tenant_state=self._tenant_state,\n # NOTE: Index filters includes metadata tags which were filtered\n # for invalid unicode at indexing time. In theory it would be\n # ideal to do filtering here as well, in practice we never did\n # that in the Vespa codepath and have not seen issues in\n # production, so we deliberately conform to the existing logic\n # in order to not unknowningly introduce a possible bug.\n index_filters=filters,\n include_hidden=False,\n )\n search_hits: list[SearchHit[DocumentChunkWithoutVectors]] = self._client.search(\n body=query_body,\n search_pipeline_id=None,\n search_type=OpenSearchSearchType.SEMANTIC,\n )\n\n inference_chunks_uncleaned: list[InferenceChunkUncleaned] = [\n _convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned(\n search_hit.document_chunk, search_hit.score, search_hit.match_highlights\n )\n for search_hit in search_hits\n ]\n inference_chunks: list[InferenceChunk] = cleanup_content_for_chunks(\n inference_chunks_uncleaned\n )\n\n return inference_chunks\n\n def random_retrieval(\n self,\n filters: IndexFilters,\n num_to_retrieve: int = 10,\n dirty: bool | None = None, # noqa: ARG002\n ) -> list[InferenceChunk]:\n logger.debug(\n f\"[OpenSearchDocumentIndex] Randomly retrieving {num_to_retrieve} chunks for index {self._index_name}.\"\n )\n query_body = DocumentQuery.get_random_search_query(\n tenant_state=self._tenant_state,\n index_filters=filters,\n num_to_retrieve=num_to_retrieve,\n )\n search_hits: list[SearchHit[DocumentChunkWithoutVectors]] = self._client.search(\n body=query_body,\n search_pipeline_id=None,\n search_type=OpenSearchSearchType.RANDOM,\n )\n inference_chunks_uncleaned: list[InferenceChunkUncleaned] = [\n _convert_retrieved_opensearch_chunk_to_inference_chunk_uncleaned(\n search_hit.document_chunk, search_hit.score, search_hit.match_highlights\n )\n for search_hit in search_hits\n ]\n inference_chunks: list[InferenceChunk] = cleanup_content_for_chunks(\n inference_chunks_uncleaned\n )\n\n return inference_chunks\n\n def index_raw_chunks(self, chunks: list[DocumentChunk]) -> None:\n \"\"\"Indexes raw document chunks into OpenSearch.\n\n Used in the Vespa migration task. Can be deleted after migrations are\n complete.\n \"\"\"\n logger.debug(\n f\"[OpenSearchDocumentIndex] Indexing {len(chunks)} raw chunks for index {self._index_name}.\"\n )\n # Do not raise if the document already exists, just update. This is\n # because the document may already have been indexed during the\n # OpenSearch transition period.\n self._client.bulk_index_documents(\n documents=chunks, tenant_state=self._tenant_state, update_if_exists=True\n )\n" }, { "path": "backend/onyx/document_index/opensearch/schema.py", "content": "import hashlib\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import Self\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom pydantic import field_serializer\nfrom pydantic import field_validator\nfrom pydantic import model_serializer\nfrom pydantic import model_validator\nfrom pydantic import SerializerFunctionWrapHandler\n\nfrom onyx.configs.app_configs import OPENSEARCH_INDEX_NUM_REPLICAS\nfrom onyx.configs.app_configs import OPENSEARCH_INDEX_NUM_SHARDS\nfrom onyx.configs.app_configs import OPENSEARCH_TEXT_ANALYZER\nfrom onyx.configs.app_configs import USING_AWS_MANAGED_OPENSEARCH\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE\nfrom onyx.document_index.opensearch.constants import EF_CONSTRUCTION\nfrom onyx.document_index.opensearch.constants import EF_SEARCH\nfrom onyx.document_index.opensearch.constants import M\nfrom onyx.document_index.opensearch.string_filtering import DocumentIDTooLongError\nfrom onyx.document_index.opensearch.string_filtering import (\n filter_and_validate_document_id,\n)\nfrom onyx.document_index.opensearch.string_filtering import (\n MAX_DOCUMENT_ID_ENCODED_LENGTH,\n)\nfrom onyx.utils.tenant import get_tenant_id_short_string\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\nTITLE_FIELD_NAME = \"title\"\nTITLE_VECTOR_FIELD_NAME = \"title_vector\"\nCONTENT_FIELD_NAME = \"content\"\nCONTENT_VECTOR_FIELD_NAME = \"content_vector\"\nSOURCE_TYPE_FIELD_NAME = \"source_type\"\nMETADATA_LIST_FIELD_NAME = \"metadata_list\"\nLAST_UPDATED_FIELD_NAME = \"last_updated\"\nPUBLIC_FIELD_NAME = \"public\"\nACCESS_CONTROL_LIST_FIELD_NAME = \"access_control_list\"\nHIDDEN_FIELD_NAME = \"hidden\"\nGLOBAL_BOOST_FIELD_NAME = \"global_boost\"\nSEMANTIC_IDENTIFIER_FIELD_NAME = \"semantic_identifier\"\nIMAGE_FILE_ID_FIELD_NAME = \"image_file_id\"\nSOURCE_LINKS_FIELD_NAME = \"source_links\"\nDOCUMENT_SETS_FIELD_NAME = \"document_sets\"\nUSER_PROJECTS_FIELD_NAME = \"user_projects\"\nPERSONAS_FIELD_NAME = \"personas\"\nDOCUMENT_ID_FIELD_NAME = \"document_id\"\nCHUNK_INDEX_FIELD_NAME = \"chunk_index\"\nMAX_CHUNK_SIZE_FIELD_NAME = \"max_chunk_size\"\nTENANT_ID_FIELD_NAME = \"tenant_id\"\nBLURB_FIELD_NAME = \"blurb\"\nDOC_SUMMARY_FIELD_NAME = \"doc_summary\"\nCHUNK_CONTEXT_FIELD_NAME = \"chunk_context\"\nMETADATA_SUFFIX_FIELD_NAME = \"metadata_suffix\"\nPRIMARY_OWNERS_FIELD_NAME = \"primary_owners\"\nSECONDARY_OWNERS_FIELD_NAME = \"secondary_owners\"\n# Hierarchy filtering - list of ancestor hierarchy node IDs\nANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME = \"ancestor_hierarchy_node_ids\"\n\n\n# Faiss was also tried but it didn't have any benefits\n# NMSLIB is deprecated, not recommended\nOPENSEARCH_KNN_ENGINE = \"lucene\"\n\n\ndef get_opensearch_doc_chunk_id(\n tenant_state: TenantState,\n document_id: str,\n chunk_index: int,\n max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE,\n) -> str:\n \"\"\"\n Returns a unique identifier for the chunk.\n\n This will be the string used to identify the chunk in OpenSearch. Any direct\n chunk queries should use this function.\n\n If the document ID is too long, a hash of the ID is used instead.\n \"\"\"\n opensearch_doc_chunk_id_suffix: str = f\"__{max_chunk_size}__{chunk_index}\"\n encoded_suffix_length: int = len(opensearch_doc_chunk_id_suffix.encode(\"utf-8\"))\n max_encoded_permissible_doc_id_length: int = (\n MAX_DOCUMENT_ID_ENCODED_LENGTH - encoded_suffix_length\n )\n opensearch_doc_chunk_id_tenant_prefix: str = \"\"\n if tenant_state.multitenant:\n short_tenant_id: str = get_tenant_id_short_string(tenant_state.tenant_id)\n # Use tenant ID because in multitenant mode each tenant has its own\n # Documents table, so there is a very small chance that doc IDs are not\n # actually unique across all tenants.\n opensearch_doc_chunk_id_tenant_prefix = f\"{short_tenant_id}__\"\n encoded_prefix_length: int = len(\n opensearch_doc_chunk_id_tenant_prefix.encode(\"utf-8\")\n )\n max_encoded_permissible_doc_id_length -= encoded_prefix_length\n\n try:\n sanitized_document_id: str = filter_and_validate_document_id(\n document_id, max_encoded_length=max_encoded_permissible_doc_id_length\n )\n except DocumentIDTooLongError:\n # If the document ID is too long, use a hash instead.\n # We use blake2b because it is faster and equally secure as SHA256, and\n # accepts digest_size which controls the number of bytes returned in the\n # hash.\n # digest_size is the size of the returned hash in bytes. Since we're\n # decoding the hash bytes as a hex string, the digest_size should be\n # half the max target size of the hash string.\n # Subtract 1 because filter_and_validate_document_id compares on >= on\n # max_encoded_length.\n # 64 is the max digest_size blake2b returns.\n digest_size: int = min((max_encoded_permissible_doc_id_length - 1) // 2, 64)\n sanitized_document_id = hashlib.blake2b(\n document_id.encode(\"utf-8\"), digest_size=digest_size\n ).hexdigest()\n\n opensearch_doc_chunk_id: str = (\n f\"{opensearch_doc_chunk_id_tenant_prefix}{sanitized_document_id}{opensearch_doc_chunk_id_suffix}\"\n )\n\n # Do one more validation to ensure we haven't exceeded the max length.\n opensearch_doc_chunk_id = filter_and_validate_document_id(opensearch_doc_chunk_id)\n return opensearch_doc_chunk_id\n\n\ndef set_or_convert_timezone_to_utc(value: datetime) -> datetime:\n if value.tzinfo is None:\n # astimezone will raise if value does not have a timezone set.\n value = value.replace(tzinfo=timezone.utc)\n else:\n # Does appropriate time conversion if value was set in a different\n # timezone.\n value = value.astimezone(timezone.utc)\n return value\n\n\nclass DocumentChunkWithoutVectors(BaseModel):\n \"\"\"\n Represents a chunk of a document in the OpenSearch index without vectors.\n\n The names of these fields are based on the OpenSearch schema. Changes to the\n schema require changes here. See get_document_schema.\n\n WARNING: Relies on MULTI_TENANT which is global state. Also uses\n get_current_tenant_id. Generally relying on global state is bad, in this\n case we accept it because of the importance of validating tenant logic.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n document_id: str\n chunk_index: int\n # The maximum number of tokens this chunk's content can hold. Previously\n # there was a concept of large chunks, this is a generic concept of that. We\n # can choose to have any size of chunks in the index and they should be\n # distinct from one another.\n max_chunk_size: int = DEFAULT_MAX_CHUNK_SIZE\n\n # Either both should be None or both should be non-None.\n title: str | None = None\n content: str\n\n source_type: str\n # A list of key-value pairs separated by INDEX_SEPARATOR. See\n # convert_metadata_dict_to_list_of_strings.\n metadata_list: list[str] | None = None\n # If it exists, time zone should always be UTC.\n last_updated: datetime | None = None\n\n public: bool\n access_control_list: list[str]\n # Defaults to False, currently gets written during update not index.\n hidden: bool = False\n\n global_boost: int\n\n semantic_identifier: str\n image_file_id: str | None = None\n # Contains a string representation of a dict which maps offset into the raw\n # chunk text to the link corresponding to that point.\n source_links: str | None = None\n blurb: str\n # doc_summary, chunk_context, and metadata_suffix are all stored simply to\n # reverse the augmentations to content. Ideally these would just be start\n # and stop indices into the content string. For legacy reasons they are not\n # right now.\n doc_summary: str\n chunk_context: str\n metadata_suffix: str | None = None\n\n document_sets: list[str] | None = None\n user_projects: list[int] | None = None\n personas: list[int] | None = None\n primary_owners: list[str] | None = None\n secondary_owners: list[str] | None = None\n\n # List of ancestor hierarchy node IDs for hierarchy-based filtering.\n # None means no hierarchy info (document will be excluded from\n # hierarchy-filtered searches).\n ancestor_hierarchy_node_ids: list[int] | None = None\n\n tenant_id: TenantState = Field(\n default_factory=lambda: TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n )\n\n def __str__(self) -> str:\n return (\n f\"DocumentChunk(document_id={self.document_id}, chunk_index={self.chunk_index}, \"\n f\"content length={len(self.content)}, tenant_id={self.tenant_id.tenant_id}).\"\n )\n\n @model_serializer(mode=\"wrap\")\n def serialize_model(\n self, handler: SerializerFunctionWrapHandler\n ) -> dict[str, object]:\n \"\"\"Invokes pydantic's serialization logic, then excludes Nones.\n\n We do this because .model_dump(exclude_none=True) does not work after\n @field_serializer logic, so for some field serializers which return None\n and which we would like to exclude from the final dump, they would be\n included without this.\n\n Args:\n handler: Callable from pydantic which takes the instance of the\n model as an argument and performs standard serialization.\n\n Returns:\n The return of handler but with None items excluded.\n \"\"\"\n serialized: dict[str, object] = handler(self)\n serialized_exclude_none = {k: v for k, v in serialized.items() if v is not None}\n return serialized_exclude_none\n\n @field_serializer(\"last_updated\", mode=\"wrap\")\n def serialize_datetime_fields_to_epoch_seconds(\n self,\n value: datetime | None,\n handler: SerializerFunctionWrapHandler, # noqa: ARG002\n ) -> int | None:\n \"\"\"\n Serializes datetime fields to seconds since the Unix epoch.\n\n If there is no datetime, returns None.\n \"\"\"\n if value is None:\n return None\n value = set_or_convert_timezone_to_utc(value)\n return int(value.timestamp())\n\n @field_validator(\"last_updated\", mode=\"before\")\n @classmethod\n def parse_epoch_seconds_to_datetime(cls, value: Any) -> datetime | None:\n \"\"\"Parses seconds since the Unix epoch to a datetime object.\n\n If the input is None, returns None.\n\n The datetime returned will be in UTC.\n \"\"\"\n if value is None:\n return None\n if isinstance(value, datetime):\n value = set_or_convert_timezone_to_utc(value)\n return value\n if not isinstance(value, int):\n raise ValueError(\n f\"Bug: Expected an int for the last_updated property from OpenSearch, got {type(value)} instead.\"\n )\n return datetime.fromtimestamp(value, tz=timezone.utc)\n\n @field_serializer(\"tenant_id\", mode=\"wrap\")\n def serialize_tenant_state(\n self,\n value: TenantState,\n handler: SerializerFunctionWrapHandler, # noqa: ARG002\n ) -> str | None:\n \"\"\"\n Serializes tenant_state to the tenant str if multitenant, or None if\n not.\n\n The idea is that in single tenant mode, the schema does not have a\n tenant_id field, so we don't want to supply it in our serialized\n DocumentChunk. This assumes the final serialized model excludes None\n fields, which serialize_model should enforce.\n \"\"\"\n if not value.multitenant:\n return None\n else:\n return value.tenant_id\n\n @field_validator(\"tenant_id\", mode=\"before\")\n @classmethod\n def parse_tenant_id(cls, value: Any) -> TenantState:\n \"\"\"\n Generates a TenantState from OpenSearch's tenant_id if it exists, or\n generates a default state if it does not (implies we are in single\n tenant mode).\n \"\"\"\n if value is None:\n if MULTI_TENANT:\n raise ValueError(\n \"Bug: No tenant_id was supplied but multi-tenant mode is enabled.\"\n )\n return TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n elif isinstance(value, TenantState):\n if MULTI_TENANT != value.multitenant:\n raise ValueError(\n f\"Bug: An existing TenantState object was supplied to the DocumentChunk model \"\n f\"but its multi-tenant mode ({value.multitenant}) does not match the program's \"\n \"current global tenancy state.\"\n )\n return value\n elif not isinstance(value, str):\n raise ValueError(\n f\"Bug: Expected a str for the tenant_id property from OpenSearch, got {type(value)} instead.\"\n )\n else:\n if not MULTI_TENANT:\n raise ValueError(\n \"Bug: Got a non-null str for the tenant_id property from OpenSearch but \"\n \"multi-tenant mode is not enabled. This is unexpected because in single-tenant \"\n \"mode we don't expect to see a tenant_id.\"\n )\n return TenantState(tenant_id=value, multitenant=MULTI_TENANT)\n\n\nclass DocumentChunk(DocumentChunkWithoutVectors):\n \"\"\"Represents a chunk of a document in the OpenSearch index.\n\n The names of these fields are based on the OpenSearch schema. Changes to the\n schema require changes here. See get_document_schema.\n \"\"\"\n\n model_config = {\"frozen\": True}\n\n title_vector: list[float] | None = None\n content_vector: list[float]\n\n def __str__(self) -> str:\n return (\n f\"DocumentChunk(document_id={self.document_id}, chunk_index={self.chunk_index}, \"\n f\"content length={len(self.content)}, content vector length={len(self.content_vector)}, \"\n f\"tenant_id={self.tenant_id.tenant_id})\"\n )\n\n @model_validator(mode=\"after\")\n def check_title_and_title_vector_are_consistent(self) -> Self:\n # title and title_vector should both either be None or not.\n if self.title is not None and self.title_vector is None:\n raise ValueError(\"Bug: Title vector must not be None if title is not None.\")\n if self.title_vector is not None and self.title is None:\n raise ValueError(\"Bug: Title must not be None if title vector is not None.\")\n return self\n\n\nclass DocumentSchema:\n \"\"\"\n Represents the schema and indexing strategies of the OpenSearch index.\n\n TODO(andrei): Implement multi-phase indexing strategies.\n \"\"\"\n\n @staticmethod\n def get_document_schema(vector_dimension: int, multitenant: bool) -> dict[str, Any]:\n \"\"\"Returns the document schema for the OpenSearch index.\n\n WARNING: Changes / additions to field names here require changes to the\n DocumentChunk class above.\n\n Notes:\n - By default all fields have indexing enabled.\n - By default almost all fields except text fields have doc_values\n enabled, enabling operations like sorting and aggregations.\n - By default all fields are nullable.\n - \"type\": \"keyword\" fields are stored as-is, used for exact matches,\n filtering, etc.\n - \"type\": \"text\" fields are OpenSearch-processed strings, used for\n full-text searches.\n - \"store\": True fields are stored and can be returned on their own,\n independent of the parent document.\n - \"index\": True fields can be queried on.\n - \"doc_values\": True fields can be sorted and aggregated efficiently.\n Not supported for \"text\" type fields.\n - \"store\": True fields are stored separately from the source document\n and can thus be returned from a query separately from _source.\n Generally this is not necessary.\n\n Args:\n vector_dimension: The dimension of vector embeddings. Must be a\n positive integer.\n multitenant: Whether the index is multitenant.\n\n Returns:\n A dictionary representing the document schema, to be supplied to the\n OpenSearch client. The structure of this dictionary is\n determined by OpenSearch documentation.\n \"\"\"\n schema: dict[str, Any] = {\n # By default OpenSearch allows dynamically adding new properties\n # based on indexed documents. This is awful and we disable it here.\n # An exception will be raised if you try to index a new doc which\n # contains unexpected fields.\n \"dynamic\": \"strict\",\n \"properties\": {\n TITLE_FIELD_NAME: {\n \"type\": \"text\",\n # Language analyzer (e.g. english) stems at index and search\n # time for variant matching. Configure via\n # OPENSEARCH_TEXT_ANALYZER. Existing indices need reindexing\n # after a change.\n \"analyzer\": OPENSEARCH_TEXT_ANALYZER,\n \"fields\": {\n # Subfield accessed as title.keyword. Not indexed for\n # values longer than 256 chars.\n # TODO(andrei): Ask Yuhong do we want this?\n \"keyword\": {\"type\": \"keyword\", \"ignore_above\": 256}\n },\n # This makes highlighting text during queries more efficient\n # at the cost of disk space. See\n # https://docs.opensearch.org/latest/search-plugins/searching-data/highlight/#methods-of-obtaining-offsets\n \"index_options\": \"offsets\",\n },\n CONTENT_FIELD_NAME: {\n \"type\": \"text\",\n \"store\": True,\n \"analyzer\": OPENSEARCH_TEXT_ANALYZER,\n \"index_options\": \"offsets\",\n },\n TITLE_VECTOR_FIELD_NAME: {\n \"type\": \"knn_vector\",\n \"dimension\": vector_dimension,\n \"method\": {\n \"name\": \"hnsw\",\n \"space_type\": \"cosinesimil\",\n \"engine\": OPENSEARCH_KNN_ENGINE,\n \"parameters\": {\"ef_construction\": EF_CONSTRUCTION, \"m\": M},\n },\n },\n # TODO(andrei): This is a tensor in Vespa. Also look at feature\n # parity for these other method fields.\n CONTENT_VECTOR_FIELD_NAME: {\n \"type\": \"knn_vector\",\n \"dimension\": vector_dimension,\n \"method\": {\n \"name\": \"hnsw\",\n \"space_type\": \"cosinesimil\",\n \"engine\": OPENSEARCH_KNN_ENGINE,\n \"parameters\": {\"ef_construction\": EF_CONSTRUCTION, \"m\": M},\n },\n },\n SOURCE_TYPE_FIELD_NAME: {\"type\": \"keyword\"},\n METADATA_LIST_FIELD_NAME: {\"type\": \"keyword\"},\n LAST_UPDATED_FIELD_NAME: {\n \"type\": \"date\",\n \"format\": \"epoch_second\",\n # For some reason date defaults to False, even though it\n # would make sense to sort by date.\n \"doc_values\": True,\n },\n # Access control fields.\n # Whether the doc is public. Could have fallen under access\n # control list but is such a broad and critical filter that it\n # is its own field. If true, ACCESS_CONTROL_LIST_FIELD_NAME\n # should have no effect on queries.\n PUBLIC_FIELD_NAME: {\"type\": \"boolean\"},\n # Access control list for the doc, excluding public access,\n # which is covered above.\n # If a user's access set contains at least one entry from this\n # set, the user should be able to retrieve this document. This\n # only applies if public is set to false; public non-hidden\n # documents are always visible to anyone in a given tenancy\n # regardless of this field.\n ACCESS_CONTROL_LIST_FIELD_NAME: {\"type\": \"keyword\"},\n # Whether the doc is hidden from search results.\n # Should clobber all other access search filters, namely\n # PUBLIC_FIELD_NAME and ACCESS_CONTROL_LIST_FIELD_NAME; up to\n # search implementations to guarantee this.\n HIDDEN_FIELD_NAME: {\"type\": \"boolean\"},\n GLOBAL_BOOST_FIELD_NAME: {\"type\": \"integer\"},\n # This field is only used for displaying a useful name for the\n # doc in the UI and is not used for searching. Disabling these\n # features to increase perf. This field is therefore essentially\n # just metadata.\n SEMANTIC_IDENTIFIER_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n # Generally False by default; just making sure.\n \"store\": False,\n },\n # Same as above; used to display an image along with the doc.\n IMAGE_FILE_ID_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n # Generally False by default; just making sure.\n \"store\": False,\n },\n # Same as above; used to link to the source doc.\n SOURCE_LINKS_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n # Generally False by default; just making sure.\n \"store\": False,\n },\n # Same as above; used to quickly summarize the doc in the UI.\n BLURB_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n # Generally False by default; just making sure.\n \"store\": False,\n },\n # Same as above.\n # TODO(andrei): If we want to search on this this needs to be\n # changed.\n DOC_SUMMARY_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n # Generally False by default; just making sure.\n \"store\": False,\n },\n # Same as above.\n # TODO(andrei): If we want to search on this this needs to be\n # changed.\n CHUNK_CONTEXT_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n # Generally False by default; just making sure.\n \"store\": False,\n },\n # Same as above.\n METADATA_SUFFIX_FIELD_NAME: {\n \"type\": \"keyword\",\n \"index\": False,\n \"doc_values\": False,\n \"store\": False,\n },\n # Product-specific fields.\n DOCUMENT_SETS_FIELD_NAME: {\"type\": \"keyword\"},\n USER_PROJECTS_FIELD_NAME: {\"type\": \"integer\"},\n PERSONAS_FIELD_NAME: {\"type\": \"integer\"},\n PRIMARY_OWNERS_FIELD_NAME: {\"type\": \"keyword\"},\n SECONDARY_OWNERS_FIELD_NAME: {\"type\": \"keyword\"},\n # OpenSearch metadata fields.\n DOCUMENT_ID_FIELD_NAME: {\"type\": \"keyword\"},\n CHUNK_INDEX_FIELD_NAME: {\"type\": \"integer\"},\n # The maximum number of tokens this chunk's content can hold.\n MAX_CHUNK_SIZE_FIELD_NAME: {\"type\": \"integer\"},\n # Hierarchy filtering - list of ancestor hierarchy node IDs.\n # Used for scoped search within folder/space hierarchies.\n # OpenSearch's terms query with value_type: \"bitmap\" can\n # efficiently check if any value in this array matches a\n # query bitmap.\n ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME: {\"type\": \"integer\"},\n },\n }\n\n if multitenant:\n schema[\"properties\"][TENANT_ID_FIELD_NAME] = {\"type\": \"keyword\"}\n\n return schema\n\n @staticmethod\n def get_index_settings_based_on_environment() -> dict[str, Any]:\n \"\"\"\n Returns the index settings based on the environment.\n \"\"\"\n if USING_AWS_MANAGED_OPENSEARCH:\n # NOTE: The number of data copies, including the primary (not a\n # replica) copy, must be divisible by the number of AZs.\n if MULTI_TENANT:\n number_of_shards = 324\n number_of_replicas = 2\n else:\n number_of_shards = 3\n number_of_replicas = 2\n else:\n number_of_shards = 1\n number_of_replicas = 1\n\n if OPENSEARCH_INDEX_NUM_SHARDS is not None:\n number_of_shards = OPENSEARCH_INDEX_NUM_SHARDS\n if OPENSEARCH_INDEX_NUM_REPLICAS is not None:\n number_of_replicas = OPENSEARCH_INDEX_NUM_REPLICAS\n\n return {\n \"index\": {\n \"number_of_shards\": number_of_shards,\n \"number_of_replicas\": number_of_replicas,\n # Required for vector search.\n \"knn\": True,\n \"knn.algo_param.ef_search\": EF_SEARCH,\n }\n }\n" }, { "path": "backend/onyx/document_index/opensearch/search.py", "content": "import random\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import TypeAlias\nfrom typing import TypeVar\n\nfrom onyx.configs.app_configs import DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S\nfrom onyx.configs.app_configs import OPENSEARCH_EXPLAIN_ENABLED\nfrom onyx.configs.app_configs import OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED\nfrom onyx.configs.app_configs import OPENSEARCH_PROFILING_DISABLED\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import INDEX_SEPARATOR\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import Tag\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.constants import ASSUMED_DOCUMENT_AGE_DAYS\nfrom onyx.document_index.opensearch.constants import (\n DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES,\n)\nfrom onyx.document_index.opensearch.constants import (\n DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW,\n)\nfrom onyx.document_index.opensearch.constants import (\n HYBRID_SEARCH_NORMALIZATION_PIPELINE,\n)\nfrom onyx.document_index.opensearch.constants import (\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION,\n)\nfrom onyx.document_index.opensearch.constants import HybridSearchNormalizationPipeline\nfrom onyx.document_index.opensearch.constants import HybridSearchSubqueryConfiguration\nfrom onyx.document_index.opensearch.schema import ACCESS_CONTROL_LIST_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import CHUNK_INDEX_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import CONTENT_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import CONTENT_VECTOR_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DOCUMENT_ID_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DOCUMENT_SETS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import HIDDEN_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import LAST_UPDATED_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import MAX_CHUNK_SIZE_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import METADATA_LIST_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import PERSONAS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import PUBLIC_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import set_or_convert_timezone_to_utc\nfrom onyx.document_index.opensearch.schema import SOURCE_TYPE_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import TENANT_ID_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import TITLE_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import TITLE_VECTOR_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import USER_PROJECTS_FIELD_NAME\n\n# See https://docs.opensearch.org/latest/query-dsl/term/terms/.\nMAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY = 65_536\n\n\n_T = TypeVar(\"_T\")\nTermsQuery: TypeAlias = dict[str, dict[str, list[_T]]]\nTermQuery: TypeAlias = dict[str, dict[str, dict[str, _T]]]\n\n\n# TODO(andrei): Turn all magic dictionaries to pydantic models.\n\n\n# Normalization pipelines combine document scores from multiple query clauses.\n# The number and ordering of weights should match the query clauses. The values\n# of the weights should sum to 1.\ndef _get_hybrid_search_normalization_weights() -> list[float]:\n if (\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION\n is HybridSearchSubqueryConfiguration.TITLE_VECTOR_CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD\n ):\n # Since the titles are included in the contents, the embedding matches\n # are heavily downweighted as they act as a boost rather than an\n # independent scoring component.\n search_title_vector_weight = 0.1\n search_content_vector_weight = 0.45\n # Single keyword weight for both title and content (merged from former\n # title keyword + content keyword).\n search_keyword_weight = 0.45\n\n # NOTE: It is critical that the order of these weights matches the order\n # of the sub-queries in the hybrid search.\n hybrid_search_normalization_weights = [\n search_title_vector_weight,\n search_content_vector_weight,\n search_keyword_weight,\n ]\n elif (\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION\n is HybridSearchSubqueryConfiguration.CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD\n ):\n search_content_vector_weight = 0.5\n # Single keyword weight for both title and content (merged from former\n # title keyword + content keyword).\n search_keyword_weight = 0.5\n\n # NOTE: It is critical that the order of these weights matches the order\n # of the sub-queries in the hybrid search.\n hybrid_search_normalization_weights = [\n search_content_vector_weight,\n search_keyword_weight,\n ]\n else:\n raise ValueError(\n f\"Bug: Unhandled hybrid search subquery configuration: {HYBRID_SEARCH_SUBQUERY_CONFIGURATION}.\"\n )\n\n assert (\n sum(hybrid_search_normalization_weights) == 1.0\n ), \"Bug: Hybrid search normalization weights do not sum to 1.0.\"\n\n return hybrid_search_normalization_weights\n\n\ndef get_min_max_normalization_pipeline_name_and_config() -> tuple[str, dict[str, Any]]:\n min_max_normalization_pipeline_name = \"normalization_pipeline_min_max\"\n min_max_normalization_pipeline_config: dict[str, Any] = {\n \"description\": \"Normalization for keyword and vector scores using min-max\",\n \"phase_results_processors\": [\n {\n # https://docs.opensearch.org/latest/search-plugins/search-pipelines/normalization-processor/\n \"normalization-processor\": {\n \"normalization\": {\"technique\": \"min_max\"},\n \"combination\": {\n \"technique\": \"arithmetic_mean\",\n \"parameters\": {\n \"weights\": _get_hybrid_search_normalization_weights()\n },\n },\n }\n }\n ],\n }\n return min_max_normalization_pipeline_name, min_max_normalization_pipeline_config\n\n\ndef get_zscore_normalization_pipeline_name_and_config() -> tuple[str, dict[str, Any]]:\n zscore_normalization_pipeline_name = \"normalization_pipeline_zscore\"\n zscore_normalization_pipeline_config: dict[str, Any] = {\n \"description\": \"Normalization for keyword and vector scores using z-score\",\n \"phase_results_processors\": [\n {\n # https://docs.opensearch.org/latest/search-plugins/search-pipelines/normalization-processor/\n \"normalization-processor\": {\n \"normalization\": {\"technique\": \"z_score\"},\n \"combination\": {\n \"technique\": \"arithmetic_mean\",\n \"parameters\": {\n \"weights\": _get_hybrid_search_normalization_weights()\n },\n },\n }\n }\n ],\n }\n return zscore_normalization_pipeline_name, zscore_normalization_pipeline_config\n\n\ndef get_normalization_pipeline_name_and_config() -> tuple[str, dict[str, Any]]:\n if (\n HYBRID_SEARCH_NORMALIZATION_PIPELINE\n is HybridSearchNormalizationPipeline.MIN_MAX\n ):\n return get_min_max_normalization_pipeline_name_and_config()\n elif (\n HYBRID_SEARCH_NORMALIZATION_PIPELINE is HybridSearchNormalizationPipeline.ZSCORE\n ):\n return get_zscore_normalization_pipeline_name_and_config()\n else:\n raise ValueError(\n f\"Bug: Unhandled hybrid search normalization pipeline: {HYBRID_SEARCH_NORMALIZATION_PIPELINE}.\"\n )\n\n\nclass DocumentQuery:\n \"\"\"\n TODO(andrei): Implement multi-phase search strategies.\n TODO(andrei): Implement document boost.\n TODO(andrei): Implement document age.\n \"\"\"\n\n @staticmethod\n def get_from_document_id_query(\n document_id: str,\n tenant_state: TenantState,\n index_filters: IndexFilters,\n include_hidden: bool,\n max_chunk_size: int,\n min_chunk_index: int | None,\n max_chunk_index: int | None,\n get_full_document: bool = True,\n ) -> dict[str, Any]:\n \"\"\"\n Returns a final search query which gets chunks from a given document ID.\n\n This query can be directly supplied to the OpenSearch client.\n\n TODO(andrei): Currently capped at 10k results. Implement scroll/point in\n time for results so that we can return arbitrarily-many IDs.\n\n Args:\n document_id: Onyx document ID. Notably not an OpenSearch document\n ID, which points to what Onyx would refer to as a chunk.\n tenant_state: Tenant state containing the tenant ID.\n index_filters: Filters for the document retrieval query.\n include_hidden: Whether to include hidden documents.\n max_chunk_size: Document chunks are categorized by the maximum\n number of tokens they can hold. This parameter specifies the\n maximum size category of document chunks to retrieve.\n min_chunk_index: The minimum chunk index to retrieve, inclusive. If\n None, no minimum chunk index will be applied.\n max_chunk_index: The maximum chunk index to retrieve, inclusive. If\n None, no maximum chunk index will be applied.\n get_full_document: Whether to get the full document body. If False,\n OpenSearch will only return the matching document chunk IDs plus\n metadata; the source data will be omitted from the response. Use\n this for performance optimization if OpenSearch IDs are\n sufficient. Defaults to True.\n\n Returns:\n A dictionary representing the final ID search query.\n \"\"\"\n filter_clauses = DocumentQuery._get_search_filters(\n tenant_state=tenant_state,\n include_hidden=include_hidden,\n access_control_list=index_filters.access_control_list,\n source_types=index_filters.source_type or [],\n tags=index_filters.tags or [],\n document_sets=index_filters.document_set or [],\n project_id_filter=index_filters.project_id_filter,\n persona_id_filter=index_filters.persona_id_filter,\n time_cutoff=index_filters.time_cutoff,\n min_chunk_index=min_chunk_index,\n max_chunk_index=max_chunk_index,\n max_chunk_size=max_chunk_size,\n document_id=document_id,\n attached_document_ids=index_filters.attached_document_ids,\n hierarchy_node_ids=index_filters.hierarchy_node_ids,\n )\n final_get_ids_query: dict[str, Any] = {\n \"query\": {\"bool\": {\"filter\": filter_clauses}},\n # We include this to make sure OpenSearch does not revert to\n # returning some number of results less than the index max allowed\n # return size.\n \"size\": DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW,\n # By default exclude retrieving the vector fields in order to save\n # on retrieval cost as we don't need them upstream.\n \"_source\": {\n \"excludes\": [TITLE_VECTOR_FIELD_NAME, CONTENT_VECTOR_FIELD_NAME]\n },\n \"timeout\": f\"{DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S}s\",\n }\n if not get_full_document:\n # If we explicitly do not want the underlying document, we will only\n # retrieve IDs.\n final_get_ids_query[\"_source\"] = False\n if not OPENSEARCH_PROFILING_DISABLED:\n final_get_ids_query[\"profile\"] = True\n\n return final_get_ids_query\n\n @staticmethod\n def delete_from_document_id_query(\n document_id: str,\n tenant_state: TenantState,\n ) -> dict[str, Any]:\n \"\"\"\n Returns a final search query which deletes chunks from a given document\n ID.\n\n This query can be directly supplied to the OpenSearch client.\n\n Intended to be supplied to the OpenSearch client's delete_by_query\n method.\n\n TODO(andrei): There is no limit to the number of document chunks that\n can be deleted by this query. This could get expensive. Consider\n implementing batching.\n\n Args:\n document_id: Onyx document ID. Notably not an OpenSearch document\n ID, which points to what Onyx would refer to as a chunk.\n tenant_state: Tenant state containing the tenant ID.\n\n Returns:\n A dictionary representing the final delete query.\n \"\"\"\n filter_clauses = DocumentQuery._get_search_filters(\n tenant_state=tenant_state,\n # Delete hidden docs too.\n include_hidden=True,\n access_control_list=None,\n source_types=[],\n tags=[],\n document_sets=[],\n project_id_filter=None,\n persona_id_filter=None,\n time_cutoff=None,\n min_chunk_index=None,\n max_chunk_index=None,\n max_chunk_size=None,\n document_id=document_id,\n )\n final_delete_query: dict[str, Any] = {\n \"query\": {\"bool\": {\"filter\": filter_clauses}},\n \"timeout\": f\"{DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S}s\",\n }\n if not OPENSEARCH_PROFILING_DISABLED:\n final_delete_query[\"profile\"] = True\n\n return final_delete_query\n\n @staticmethod\n def get_hybrid_search_query(\n query_text: str,\n query_vector: list[float],\n num_hits: int,\n tenant_state: TenantState,\n index_filters: IndexFilters,\n include_hidden: bool,\n ) -> dict[str, Any]:\n \"\"\"Returns a final hybrid search query.\n\n NOTE: This query can be directly supplied to the OpenSearch client, but\n it MUST be supplied in addition to a search pipeline. The results from\n hybrid search are not meaningful without that step.\n\n TODO(andrei): There is some duplicated logic in this function with\n others in this file.\n\n Args:\n query_text: The text to query for.\n query_vector: The vector embedding of the text to query for.\n num_hits: The final number of hits to return.\n tenant_state: Tenant state containing the tenant ID.\n index_filters: Filters for the hybrid search query.\n include_hidden: Whether to include hidden documents.\n\n Returns:\n A dictionary representing the final hybrid search query.\n \"\"\"\n # WARNING: Profiling does not work with hybrid search; do not add it at\n # this level. See https://github.com/opensearch-project/neural-search/issues/1255\n\n if num_hits > DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW:\n raise ValueError(\n f\"Bug: num_hits ({num_hits}) is greater than the current maximum allowed \"\n f\"result window ({DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW}).\"\n )\n\n # TODO(andrei, yuhong): We can tune this more dynamically based on\n # num_hits.\n max_results_per_subquery = DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES\n\n hybrid_search_subqueries = DocumentQuery._get_hybrid_search_subqueries(\n query_text, query_vector, vector_candidates=max_results_per_subquery\n )\n hybrid_search_filters = DocumentQuery._get_search_filters(\n tenant_state=tenant_state,\n include_hidden=include_hidden,\n # TODO(andrei): We've done no filtering for PUBLIC_DOC_PAT up to\n # now. This should not cause any issues but it can introduce\n # redundant filters in queries that may affect performance.\n access_control_list=index_filters.access_control_list,\n source_types=index_filters.source_type or [],\n tags=index_filters.tags or [],\n document_sets=index_filters.document_set or [],\n project_id_filter=index_filters.project_id_filter,\n persona_id_filter=index_filters.persona_id_filter,\n time_cutoff=index_filters.time_cutoff,\n min_chunk_index=None,\n max_chunk_index=None,\n attached_document_ids=index_filters.attached_document_ids,\n hierarchy_node_ids=index_filters.hierarchy_node_ids,\n )\n\n # See https://docs.opensearch.org/latest/query-dsl/compound/hybrid/\n hybrid_search_query: dict[str, Any] = {\n \"hybrid\": {\n \"queries\": hybrid_search_subqueries,\n # Max results per subquery per shard before aggregation. Ensures\n # keyword and vector subqueries contribute equally to the\n # candidate pool for hybrid fusion.\n # Sources:\n # https://docs.opensearch.org/latest/vector-search/ai-search/hybrid-search/pagination/\n # https://opensearch.org/blog/navigating-pagination-in-hybrid-queries-with-the-pagination_depth-parameter/\n \"pagination_depth\": max_results_per_subquery,\n # Applied to all the sub-queries independently (this avoids\n # subqueries having a lot of results thrown out during\n # aggregation).\n # Sources:\n # https://docs.opensearch.org/latest/query-dsl/compound/hybrid/\n # https://opensearch.org/blog/introducing-common-filter-support-for-hybrid-search-queries\n # Does AND for each filter in the list.\n \"filter\": {\"bool\": {\"filter\": hybrid_search_filters}},\n }\n }\n\n final_hybrid_search_body: dict[str, Any] = {\n \"query\": hybrid_search_query,\n \"size\": num_hits,\n \"timeout\": f\"{DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S}s\",\n # Exclude retrieving the vector fields in order to save on\n # retrieval cost as we don't need them upstream.\n \"_source\": {\n \"excludes\": [TITLE_VECTOR_FIELD_NAME, CONTENT_VECTOR_FIELD_NAME]\n },\n }\n\n if not OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED:\n final_hybrid_search_body[\"highlight\"] = (\n DocumentQuery._get_match_highlights_configuration()\n )\n\n # Explain is for scoring breakdowns. Setting this significantly\n # increases query latency.\n if OPENSEARCH_EXPLAIN_ENABLED:\n final_hybrid_search_body[\"explain\"] = True\n\n return final_hybrid_search_body\n\n @staticmethod\n def get_keyword_search_query(\n query_text: str,\n num_hits: int,\n tenant_state: TenantState,\n index_filters: IndexFilters,\n include_hidden: bool,\n ) -> dict[str, Any]:\n \"\"\"Returns a final keyword search query.\n\n This query can be directly supplied to the OpenSearch client.\n\n TODO(andrei): There is some duplicated logic in this function with\n others in this file.\n\n Args:\n query_text: The text to query for.\n num_hits: The final number of hits to return.\n tenant_state: Tenant state containing the tenant ID.\n index_filters: Filters for the keyword search query.\n include_hidden: Whether to include hidden documents.\n\n Returns:\n A dictionary representing the final keyword search query.\n \"\"\"\n if num_hits > DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW:\n raise ValueError(\n f\"Bug: num_hits ({num_hits}) is greater than the current maximum allowed \"\n f\"result window ({DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW}).\"\n )\n\n keyword_search_filters = DocumentQuery._get_search_filters(\n tenant_state=tenant_state,\n include_hidden=include_hidden,\n # TODO(andrei): We've done no filtering for PUBLIC_DOC_PAT up to\n # now. This should not cause any issues but it can introduce\n # redundant filters in queries that may affect performance.\n access_control_list=index_filters.access_control_list,\n source_types=index_filters.source_type or [],\n tags=index_filters.tags or [],\n document_sets=index_filters.document_set or [],\n project_id_filter=index_filters.project_id_filter,\n persona_id_filter=index_filters.persona_id_filter,\n time_cutoff=index_filters.time_cutoff,\n min_chunk_index=None,\n max_chunk_index=None,\n attached_document_ids=index_filters.attached_document_ids,\n hierarchy_node_ids=index_filters.hierarchy_node_ids,\n )\n\n keyword_search_query = (\n DocumentQuery._get_title_content_combined_keyword_search_query(\n query_text, search_filters=keyword_search_filters\n )\n )\n\n final_keyword_search_query: dict[str, Any] = {\n \"query\": keyword_search_query,\n \"size\": num_hits,\n \"timeout\": f\"{DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S}s\",\n # Exclude retrieving the vector fields in order to save on\n # retrieval cost as we don't need them upstream.\n \"_source\": {\n \"excludes\": [TITLE_VECTOR_FIELD_NAME, CONTENT_VECTOR_FIELD_NAME]\n },\n }\n\n if not OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED:\n final_keyword_search_query[\"highlight\"] = (\n DocumentQuery._get_match_highlights_configuration()\n )\n\n if not OPENSEARCH_PROFILING_DISABLED:\n final_keyword_search_query[\"profile\"] = True\n\n # Explain is for scoring breakdowns. Setting this significantly\n # increases query latency.\n if OPENSEARCH_EXPLAIN_ENABLED:\n final_keyword_search_query[\"explain\"] = True\n\n return final_keyword_search_query\n\n @staticmethod\n def get_semantic_search_query(\n query_embedding: list[float],\n num_hits: int,\n tenant_state: TenantState,\n index_filters: IndexFilters,\n include_hidden: bool,\n ) -> dict[str, Any]:\n \"\"\"Returns a final semantic search query.\n\n This query can be directly supplied to the OpenSearch client.\n\n TODO(andrei): There is some duplicated logic in this function with\n others in this file.\n\n Args:\n query_embedding: The vector embedding of the text to query for.\n num_hits: The final number of hits to return.\n tenant_state: Tenant state containing the tenant ID.\n index_filters: Filters for the semantic search query.\n include_hidden: Whether to include hidden documents.\n\n Returns:\n A dictionary representing the final semantic search query.\n \"\"\"\n if num_hits > DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW:\n raise ValueError(\n f\"Bug: num_hits ({num_hits}) is greater than the current maximum allowed \"\n f\"result window ({DEFAULT_OPENSEARCH_MAX_RESULT_WINDOW}).\"\n )\n\n semantic_search_filters = DocumentQuery._get_search_filters(\n tenant_state=tenant_state,\n include_hidden=include_hidden,\n # TODO(andrei): We've done no filtering for PUBLIC_DOC_PAT up to\n # now. This should not cause any issues but it can introduce\n # redundant filters in queries that may affect performance.\n access_control_list=index_filters.access_control_list,\n source_types=index_filters.source_type or [],\n tags=index_filters.tags or [],\n document_sets=index_filters.document_set or [],\n project_id_filter=index_filters.project_id_filter,\n persona_id_filter=index_filters.persona_id_filter,\n time_cutoff=index_filters.time_cutoff,\n min_chunk_index=None,\n max_chunk_index=None,\n attached_document_ids=index_filters.attached_document_ids,\n hierarchy_node_ids=index_filters.hierarchy_node_ids,\n )\n\n semantic_search_query = (\n DocumentQuery._get_content_vector_similarity_search_query(\n query_embedding,\n vector_candidates=num_hits,\n search_filters=semantic_search_filters,\n )\n )\n\n final_semantic_search_query: dict[str, Any] = {\n \"query\": semantic_search_query,\n \"size\": num_hits,\n \"timeout\": f\"{DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S}s\",\n # Exclude retrieving the vector fields in order to save on\n # retrieval cost as we don't need them upstream.\n \"_source\": {\n \"excludes\": [TITLE_VECTOR_FIELD_NAME, CONTENT_VECTOR_FIELD_NAME]\n },\n }\n\n if not OPENSEARCH_PROFILING_DISABLED:\n final_semantic_search_query[\"profile\"] = True\n\n # Explain is for scoring breakdowns. Setting this significantly\n # increases query latency.\n if OPENSEARCH_EXPLAIN_ENABLED:\n final_semantic_search_query[\"explain\"] = True\n\n return final_semantic_search_query\n\n @staticmethod\n def get_random_search_query(\n tenant_state: TenantState,\n index_filters: IndexFilters,\n num_to_retrieve: int,\n ) -> dict[str, Any]:\n \"\"\"Returns a final search query that gets document chunks randomly.\n\n Args:\n tenant_state: Tenant state containing the tenant ID.\n index_filters: Filters for the random search query.\n num_to_retrieve: Number of document chunks to retrieve.\n\n Returns:\n A dictionary representing the final random search query.\n \"\"\"\n search_filters = DocumentQuery._get_search_filters(\n tenant_state=tenant_state,\n include_hidden=False,\n access_control_list=index_filters.access_control_list,\n source_types=index_filters.source_type or [],\n tags=index_filters.tags or [],\n document_sets=index_filters.document_set or [],\n project_id_filter=index_filters.project_id_filter,\n persona_id_filter=index_filters.persona_id_filter,\n time_cutoff=index_filters.time_cutoff,\n min_chunk_index=None,\n max_chunk_index=None,\n attached_document_ids=index_filters.attached_document_ids,\n hierarchy_node_ids=index_filters.hierarchy_node_ids,\n )\n final_random_search_query = {\n \"query\": {\n \"function_score\": {\n \"query\": {\"bool\": {\"filter\": search_filters}},\n # See\n # https://docs.opensearch.org/latest/query-dsl/compound/function-score/#the-random-score-function\n \"random_score\": {\n # We'll use a different seed per invocation.\n \"seed\": random.randint(0, 1_000_000),\n # Some field which has a unique value per document\n # chunk.\n \"field\": \"_seq_no\",\n },\n # Replaces whatever score was computed in the query.\n \"boost_mode\": \"replace\",\n }\n },\n \"size\": num_to_retrieve,\n \"timeout\": f\"{DEFAULT_OPENSEARCH_QUERY_TIMEOUT_S}s\",\n # Exclude retrieving the vector fields in order to save on\n # retrieval cost as we don't need them upstream.\n \"_source\": {\n \"excludes\": [TITLE_VECTOR_FIELD_NAME, CONTENT_VECTOR_FIELD_NAME]\n },\n }\n if not OPENSEARCH_PROFILING_DISABLED:\n final_random_search_query[\"profile\"] = True\n\n return final_random_search_query\n\n @staticmethod\n def _get_hybrid_search_subqueries(\n query_text: str,\n query_vector: list[float],\n # The default number of neighbors to consider for knn vector similarity\n # search. This is higher than the number of results because the scoring\n # is hybrid. For a detailed breakdown, see where the default value is\n # set.\n vector_candidates: int = DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES,\n ) -> list[dict[str, Any]]:\n \"\"\"Returns subqueries for hybrid search.\n\n Each of these subqueries are the \"hybrid\" component of this search. We\n search on various things and combine results.\n\n The return of this function is not sufficient to be directly supplied to\n the OpenSearch client. See get_hybrid_search_query.\n\n Normalization is not performed here.\n The weights of each of these subqueries should be configured in a search\n pipeline.\n\n The exact subqueries executed depend on the\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION setting.\n\n NOTE: For OpenSearch, 5 is the maximum number of query clauses allowed\n in a single hybrid query. Source:\n https://docs.opensearch.org/latest/query-dsl/compound/hybrid/\n\n NOTE: Each query is independent during the search phase; there is no\n backfilling of scores for missing query components. What this means is\n that if a document was a good vector match but did not show up for\n keyword, it gets a score of 0 for the keyword component of the hybrid\n scoring. This is not as bad as just disregarding a score though as there\n is normalization applied after. So really it is \"increasing\" the missing\n score compared to if it was included and the range was renormalized.\n This does however mean that between docs that have high scores for say\n the vector field, the keyword scores between them are completely ignored\n unless they also showed up in the keyword query as a reasonably high\n match. TLDR, this is a bit of unique funky behavior but it seems ok.\n\n NOTE: Options considered and rejected:\n - minimum_should_match: Since it's hybrid search and users often provide\n semantic queries, there is often a lot of terms, and very low number\n of meaningful keywords (and a low ratio of keywords).\n - fuzziness AUTO: Typo tolerance (0/1/2 edit distance by term length).\n It's mostly for typos as the analyzer (\"english\" by default) already\n does some stemming and tokenization. In testing datasets, this makes\n recall slightly worse. It also is less performant so not really any\n reason to do it.\n\n Args:\n query_text: The text of the query to search for.\n query_vector: The vector embedding of the query to search for.\n num_candidates: The number of candidates to consider for vector\n similarity search.\n \"\"\"\n # Build sub-queries for hybrid search. Order must match normalization\n # pipeline weights.\n if (\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION\n is HybridSearchSubqueryConfiguration.TITLE_VECTOR_CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD\n ):\n return [\n DocumentQuery._get_title_vector_similarity_search_query(\n query_vector, vector_candidates\n ),\n DocumentQuery._get_content_vector_similarity_search_query(\n query_vector, vector_candidates\n ),\n DocumentQuery._get_title_content_combined_keyword_search_query(\n query_text\n ),\n ]\n elif (\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION\n is HybridSearchSubqueryConfiguration.CONTENT_VECTOR_TITLE_CONTENT_COMBINED_KEYWORD\n ):\n return [\n DocumentQuery._get_content_vector_similarity_search_query(\n query_vector, vector_candidates\n ),\n DocumentQuery._get_title_content_combined_keyword_search_query(\n query_text\n ),\n ]\n else:\n raise ValueError(\n f\"Bug: Unhandled hybrid search subquery configuration: {HYBRID_SEARCH_SUBQUERY_CONFIGURATION}\"\n )\n\n @staticmethod\n def _get_title_vector_similarity_search_query(\n query_vector: list[float],\n vector_candidates: int = DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES,\n ) -> dict[str, Any]:\n return {\n \"knn\": {\n TITLE_VECTOR_FIELD_NAME: {\n \"vector\": query_vector,\n \"k\": vector_candidates,\n }\n }\n }\n\n @staticmethod\n def _get_content_vector_similarity_search_query(\n query_vector: list[float],\n vector_candidates: int = DEFAULT_NUM_HYBRID_SUBQUERY_CANDIDATES,\n search_filters: list[dict[str, Any]] | None = None,\n ) -> dict[str, Any]:\n query = {\n \"knn\": {\n CONTENT_VECTOR_FIELD_NAME: {\n \"vector\": query_vector,\n \"k\": vector_candidates,\n }\n }\n }\n\n if search_filters is not None:\n query[\"knn\"][CONTENT_VECTOR_FIELD_NAME][\"filter\"] = {\n \"bool\": {\"filter\": search_filters}\n }\n\n return query\n\n @staticmethod\n def _get_title_content_combined_keyword_search_query(\n query_text: str,\n search_filters: list[dict[str, Any]] | None = None,\n ) -> dict[str, Any]:\n query = {\n \"bool\": {\n \"should\": [\n {\n \"match\": {\n TITLE_FIELD_NAME: {\n \"query\": query_text,\n \"operator\": \"or\",\n # The title fields are strongly discounted as\n # they are included in the content. This just\n # acts as a minor boost.\n \"boost\": 0.1,\n }\n }\n },\n {\n \"match_phrase\": {\n TITLE_FIELD_NAME: {\n \"query\": query_text,\n \"slop\": 1,\n \"boost\": 0.2,\n }\n }\n },\n {\n # Analyzes the query and returns results which match any\n # of the query's terms. More matches result in higher\n # scores.\n \"match\": {\n CONTENT_FIELD_NAME: {\n \"query\": query_text,\n \"operator\": \"or\",\n \"boost\": 1.0,\n }\n }\n },\n {\n # Matches an exact phrase in a specified order.\n \"match_phrase\": {\n CONTENT_FIELD_NAME: {\n \"query\": query_text,\n # The number of words permitted between words of\n # a query phrase and still result in a match.\n \"slop\": 1,\n \"boost\": 1.5,\n }\n }\n },\n ],\n # Ensures at least one match subquery from the query is present\n # in the document. This defaults to 1, unless a filter or must\n # clause is supplied, in which case it defaults to 0.\n \"minimum_should_match\": 1,\n }\n }\n\n if search_filters is not None:\n query[\"bool\"][\"filter\"] = search_filters\n\n return query\n\n @staticmethod\n def _get_search_filters(\n tenant_state: TenantState,\n include_hidden: bool,\n access_control_list: list[str] | None,\n source_types: list[DocumentSource],\n tags: list[Tag],\n document_sets: list[str],\n project_id_filter: int | None,\n persona_id_filter: int | None,\n time_cutoff: datetime | None,\n min_chunk_index: int | None,\n max_chunk_index: int | None,\n max_chunk_size: int | None = None,\n document_id: str | None = None,\n # Assistant knowledge filters\n attached_document_ids: list[str] | None = None,\n hierarchy_node_ids: list[int] | None = None,\n ) -> list[dict[str, Any]]:\n \"\"\"Returns filters to be passed into the \"filter\" key of a search query.\n\n The \"filter\" key applies a logical AND operator to its elements, so\n every subfilter must evaluate to true in order for the document to be\n retrieved. This function returns a list of such subfilters.\n See https://docs.opensearch.org/latest/query-dsl/compound/bool/.\n\n TODO(ENG-3874): The terms queries returned by this function can be made\n more performant for large cardinality sets by sorting the values by\n their UTF-8 byte order.\n\n TODO(ENG-3875): This function can take even better advantage of filter\n caching by grouping \"static\" filters together into one sub-clause.\n\n Args:\n tenant_state: Tenant state containing the tenant ID.\n include_hidden: Whether to include hidden documents.\n access_control_list: Access control list for the documents to\n retrieve. If None, there is no restriction on the documents that\n can be retrieved. If not None, only public documents can be\n retrieved, or non-public documents where at least one acl\n provided here is present in the document's acl list.\n source_types: If supplied, only documents of one of these source\n types will be retrieved.\n tags: If supplied, only documents with an entry in their metadata\n list corresponding to a tag will be retrieved.\n document_sets: If supplied, only documents with at least one\n document set ID from this list will be retrieved.\n project_id_filter: If not None, only documents with this project ID\n in user projects will be retrieved. Additive — only applied\n when a knowledge scope already exists.\n persona_id_filter: If not None, only documents whose personas array\n contains this persona ID will be retrieved. Primary — creates\n a knowledge scope on its own.\n time_cutoff: Time cutoff for the documents to retrieve. If not None,\n Documents which were last updated before this date will not be\n returned. For documents which do not have a value for their last\n updated time, we assume some default age of\n ASSUMED_DOCUMENT_AGE_DAYS for when the document was last\n updated.\n min_chunk_index: The minimum chunk index to retrieve, inclusive. If\n None, no minimum chunk index will be applied.\n max_chunk_index: The maximum chunk index to retrieve, inclusive. If\n None, no maximum chunk index will be applied.\n max_chunk_size: The type of chunk to retrieve, specified by the\n maximum number of tokens it can hold. If None, no filter will be\n applied for this. Defaults to None.\n NOTE: See DocumentChunk.max_chunk_size.\n document_id: The document ID to retrieve. If None, no filter will be\n applied for this. Defaults to None.\n attached_document_ids: Document IDs explicitly attached to the\n assistant. If provided along with hierarchy_node_ids, documents\n matching EITHER criteria will be retrieved (OR logic).\n hierarchy_node_ids: Hierarchy node IDs (folders/spaces) attached to\n the assistant. Matches chunks where ancestor_hierarchy_node_ids\n contains any of these values.\n\n Raises:\n ValueError: document_id and attached_document_ids were supplied\n together. This is not allowed because they operate on the same\n schema field, and it does not semantically make sense to use\n them together.\n ValueError: Too many of one of the collection arguments was\n supplied.\n\n Returns:\n A list of filters to be passed into the \"filter\" key of a search\n query.\n \"\"\"\n\n def _get_acl_visibility_filter(\n access_control_list: list[str],\n ) -> dict[str, dict[str, list[TermQuery[bool] | TermsQuery[str]] | int]]:\n \"\"\"Returns a filter for the access control list.\n\n Since this returns an isolated bool should clause, it can be cached\n in OpenSearch independently of other clauses in _get_search_filters.\n\n Args:\n access_control_list: The access control list to restrict\n documents to.\n\n Raises:\n ValueError: The number of access control list entries is greater\n than MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY.\n\n Returns:\n A filter for the access control list.\n \"\"\"\n # Logical OR operator on its elements.\n acl_visibility_filter: dict[str, dict[str, Any]] = {\n \"bool\": {\n \"should\": [{\"term\": {PUBLIC_FIELD_NAME: {\"value\": True}}}],\n \"minimum_should_match\": 1,\n }\n }\n if access_control_list:\n if len(access_control_list) > MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY:\n raise ValueError(\n f\"Too many access control list entries: {len(access_control_list)}. Max allowed: {MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY}.\"\n )\n # Use terms instead of a list of term within a should clause\n # because Lucene will optimize the filtering for large sets of\n # terms. Small sets of terms are not expected to perform any\n # differently than individual term clauses.\n acl_subclause: TermsQuery[str] = {\n \"terms\": {ACCESS_CONTROL_LIST_FIELD_NAME: list(access_control_list)}\n }\n acl_visibility_filter[\"bool\"][\"should\"].append(acl_subclause)\n return acl_visibility_filter\n\n def _get_source_type_filter(\n source_types: list[DocumentSource],\n ) -> TermsQuery[str]:\n \"\"\"Returns a filter for the source types.\n\n Since this returns an isolated terms clause, it can be cached in\n OpenSearch independently of other clauses in _get_search_filters.\n\n Args:\n source_types: The source types to restrict documents to.\n\n Raises:\n ValueError: The number of source types is greater than\n MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY.\n ValueError: An empty list was supplied.\n\n Returns:\n A filter for the source types.\n \"\"\"\n if not source_types:\n raise ValueError(\n \"source_types cannot be empty if trying to create a source type filter.\"\n )\n if len(source_types) > MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY:\n raise ValueError(\n f\"Too many source types: {len(source_types)}. Max allowed: {MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY}.\"\n )\n # Use terms instead of a list of term within a should clause because\n # Lucene will optimize the filtering for large sets of terms. Small\n # sets of terms are not expected to perform any differently than\n # individual term clauses.\n return {\n \"terms\": {\n SOURCE_TYPE_FIELD_NAME: [\n source_type.value for source_type in source_types\n ]\n }\n }\n\n def _get_tag_filter(tags: list[Tag]) -> TermsQuery[str]:\n \"\"\"Returns a filter for the tags.\n\n Since this returns an isolated terms clause, it can be cached in\n OpenSearch independently of other clauses in _get_search_filters.\n\n Args:\n tags: The tags to restrict documents to.\n\n Raises:\n ValueError: The number of tags is greater than\n MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY.\n ValueError: An empty list was supplied.\n\n Returns:\n A filter for the tags.\n \"\"\"\n if not tags:\n raise ValueError(\n \"tags cannot be empty if trying to create a tag filter.\"\n )\n if len(tags) > MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY:\n raise ValueError(\n f\"Too many tags: {len(tags)}. Max allowed: {MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY}.\"\n )\n # Kind of an abstraction leak, see\n # convert_metadata_dict_to_list_of_strings for why metadata list\n # entries are expected to look this way.\n tag_str_list = [\n f\"{tag.tag_key}{INDEX_SEPARATOR}{tag.tag_value}\" for tag in tags\n ]\n # Use terms instead of a list of term within a should clause because\n # Lucene will optimize the filtering for large sets of terms. Small\n # sets of terms are not expected to perform any differently than\n # individual term clauses.\n return {\"terms\": {METADATA_LIST_FIELD_NAME: tag_str_list}}\n\n def _get_document_set_filter(document_sets: list[str]) -> TermsQuery[str]:\n \"\"\"Returns a filter for the document sets.\n\n Since this returns an isolated terms clause, it can be cached in\n OpenSearch independently of other clauses in _get_search_filters.\n\n Args:\n document_sets: The document sets to restrict documents to.\n\n Raises:\n ValueError: The number of document sets is greater than\n MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY.\n ValueError: An empty list was supplied.\n\n Returns:\n A filter for the document sets.\n \"\"\"\n if not document_sets:\n raise ValueError(\n \"document_sets cannot be empty if trying to create a document set filter.\"\n )\n if len(document_sets) > MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY:\n raise ValueError(\n f\"Too many document sets: {len(document_sets)}. Max allowed: {MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY}.\"\n )\n # Use terms instead of a list of term within a should clause because\n # Lucene will optimize the filtering for large sets of terms. Small\n # sets of terms are not expected to perform any differently than\n # individual term clauses.\n return {\"terms\": {DOCUMENT_SETS_FIELD_NAME: list(document_sets)}}\n\n def _get_user_project_filter(project_id: int) -> TermQuery[int]:\n return {\"term\": {USER_PROJECTS_FIELD_NAME: {\"value\": project_id}}}\n\n def _get_persona_filter(persona_id: int) -> TermQuery[int]:\n return {\"term\": {PERSONAS_FIELD_NAME: {\"value\": persona_id}}}\n\n def _get_time_cutoff_filter(time_cutoff: datetime) -> dict[str, Any]:\n # Convert to UTC if not already so the cutoff is comparable to the\n # document data.\n time_cutoff = set_or_convert_timezone_to_utc(time_cutoff)\n # Logical OR operator on its elements.\n time_cutoff_filter: dict[str, Any] = {\n \"bool\": {\"should\": [], \"minimum_should_match\": 1}\n }\n time_cutoff_filter[\"bool\"][\"should\"].append(\n {\n \"range\": {\n LAST_UPDATED_FIELD_NAME: {\"gte\": int(time_cutoff.timestamp())}\n }\n }\n )\n if time_cutoff < datetime.now(timezone.utc) - timedelta(\n days=ASSUMED_DOCUMENT_AGE_DAYS\n ):\n # Since the time cutoff is older than ASSUMED_DOCUMENT_AGE_DAYS\n # ago, we include documents which have no\n # LAST_UPDATED_FIELD_NAME value.\n time_cutoff_filter[\"bool\"][\"should\"].append(\n {\n \"bool\": {\n \"must_not\": {\"exists\": {\"field\": LAST_UPDATED_FIELD_NAME}}\n }\n }\n )\n return time_cutoff_filter\n\n def _get_chunk_index_filter(\n min_chunk_index: int | None, max_chunk_index: int | None\n ) -> dict[str, Any]:\n range_clause: dict[str, Any] = {\"range\": {CHUNK_INDEX_FIELD_NAME: {}}}\n if min_chunk_index is not None:\n range_clause[\"range\"][CHUNK_INDEX_FIELD_NAME][\"gte\"] = min_chunk_index\n if max_chunk_index is not None:\n range_clause[\"range\"][CHUNK_INDEX_FIELD_NAME][\"lte\"] = max_chunk_index\n return range_clause\n\n def _get_attached_document_id_filter(\n doc_ids: list[str],\n ) -> TermsQuery[str]:\n \"\"\"\n Returns a filter for documents explicitly attached to an assistant.\n\n Since this returns an isolated terms clause, it can be cached in\n OpenSearch independently of other clauses in _get_search_filters.\n\n Args:\n doc_ids: The document IDs to restrict documents to.\n\n Raises:\n ValueError: The number of document IDs is greater than\n MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY.\n ValueError: An empty list was supplied.\n\n Returns:\n A filter for the document IDs.\n \"\"\"\n if not doc_ids:\n raise ValueError(\n \"doc_ids cannot be empty if trying to create a document ID filter.\"\n )\n if len(doc_ids) > MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY:\n raise ValueError(\n f\"Too many document IDs: {len(doc_ids)}. Max allowed: {MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY}.\"\n )\n # Use terms instead of a list of term within a should clause because\n # Lucene will optimize the filtering for large sets of terms. Small\n # sets of terms are not expected to perform any differently than\n # individual term clauses.\n return {\"terms\": {DOCUMENT_ID_FIELD_NAME: list(doc_ids)}}\n\n def _get_hierarchy_node_filter(\n node_ids: list[int],\n ) -> TermsQuery[int]:\n \"\"\"\n Returns a filter for chunks whose ancestors include any of the given\n hierarchy nodes.\n\n Since this returns an isolated terms clause, it can be cached in\n OpenSearch independently of other clauses in _get_search_filters.\n\n Args:\n node_ids: The hierarchy node IDs to restrict documents to.\n\n Raises:\n ValueError: The number of hierarchy node IDs is greater than\n MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY.\n ValueError: An empty list was supplied.\n\n Returns:\n A filter for the hierarchy node IDs.\n \"\"\"\n if not node_ids:\n raise ValueError(\n \"node_ids cannot be empty if trying to create a hierarchy node ID filter.\"\n )\n if len(node_ids) > MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY:\n raise ValueError(\n f\"Too many hierarchy node IDs: {len(node_ids)}. Max allowed: {MAX_NUM_TERMS_ALLOWED_IN_TERMS_QUERY}.\"\n )\n # Use terms instead of a list of term within a should clause because\n # Lucene will optimize the filtering for large sets of terms. Small\n # sets of terms are not expected to perform any differently than\n # individual term clauses.\n return {\"terms\": {ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME: list(node_ids)}}\n\n if document_id is not None and attached_document_ids is not None:\n raise ValueError(\n \"document_id and attached_document_ids cannot be used together.\"\n )\n\n filter_clauses: list[dict[str, Any]] = []\n\n if not include_hidden:\n filter_clauses.append({\"term\": {HIDDEN_FIELD_NAME: {\"value\": False}}})\n\n if access_control_list is not None:\n # If an access control list is provided, the caller can only\n # retrieve public documents, and non-public documents where at least\n # one acl provided here is present in the document's acl list. If\n # there is explicitly no list provided, we make no restrictions on\n # the documents that can be retrieved.\n filter_clauses.append(_get_acl_visibility_filter(access_control_list))\n\n if source_types:\n # If at least one source type is provided, the caller will only\n # retrieve documents whose source type is present in this input\n # list.\n filter_clauses.append(_get_source_type_filter(source_types))\n\n if tags:\n # If at least one tag is provided, the caller will only retrieve\n # documents where at least one tag provided here is present in the\n # document's metadata list.\n filter_clauses.append(_get_tag_filter(tags))\n\n # Knowledge scope: explicit knowledge attachments restrict what an\n # assistant can see. When none are set the assistant searches\n # everything.\n #\n # persona_id_filter is a primary trigger — a persona with user files IS\n # explicit knowledge, so it can start a knowledge scope on its own.\n #\n # project_id_filter is additive — it widens the scope to also cover\n # overflowing project files but never restricts on its own (a chat\n # inside a project should still search team knowledge).\n has_knowledge_scope = (\n attached_document_ids\n or hierarchy_node_ids\n or document_sets\n or persona_id_filter is not None\n )\n\n if has_knowledge_scope:\n # Since this returns an isolated bool should clause, it can be\n # cached in OpenSearch independently of other clauses in\n # _get_search_filters.\n knowledge_filter: dict[str, Any] = {\n \"bool\": {\"should\": [], \"minimum_should_match\": 1}\n }\n if attached_document_ids:\n knowledge_filter[\"bool\"][\"should\"].append(\n _get_attached_document_id_filter(attached_document_ids)\n )\n if hierarchy_node_ids:\n knowledge_filter[\"bool\"][\"should\"].append(\n _get_hierarchy_node_filter(hierarchy_node_ids)\n )\n if document_sets:\n knowledge_filter[\"bool\"][\"should\"].append(\n _get_document_set_filter(document_sets)\n )\n if persona_id_filter is not None:\n knowledge_filter[\"bool\"][\"should\"].append(\n _get_persona_filter(persona_id_filter)\n )\n if project_id_filter is not None:\n knowledge_filter[\"bool\"][\"should\"].append(\n _get_user_project_filter(project_id_filter)\n )\n filter_clauses.append(knowledge_filter)\n\n if time_cutoff is not None:\n # If a time cutoff is provided, the caller will only retrieve\n # documents where the document was last updated at or after the time\n # cutoff. For documents which do not have a value for\n # LAST_UPDATED_FIELD_NAME, we assume some default age for the\n # purposes of time cutoff.\n filter_clauses.append(_get_time_cutoff_filter(time_cutoff))\n\n if min_chunk_index is not None or max_chunk_index is not None:\n filter_clauses.append(\n _get_chunk_index_filter(min_chunk_index, max_chunk_index)\n )\n\n if document_id is not None:\n filter_clauses.append(\n {\"term\": {DOCUMENT_ID_FIELD_NAME: {\"value\": document_id}}}\n )\n\n if max_chunk_size is not None:\n filter_clauses.append(\n {\"term\": {MAX_CHUNK_SIZE_FIELD_NAME: {\"value\": max_chunk_size}}}\n )\n\n if tenant_state.multitenant:\n filter_clauses.append(\n {\"term\": {TENANT_ID_FIELD_NAME: {\"value\": tenant_state.tenant_id}}}\n )\n\n return filter_clauses\n\n @staticmethod\n def _get_match_highlights_configuration() -> dict[str, Any]:\n \"\"\"\n Gets configuration for returning match highlights for a hit.\n \"\"\"\n match_highlights_configuration: dict[str, Any] = {\n \"fields\": {\n CONTENT_FIELD_NAME: {\n # See https://docs.opensearch.org/latest/search-plugins/searching-data/highlight/#highlighter-types\n \"type\": \"unified\",\n # The length in chars of a match snippet. Somewhat\n # arbitrarily-chosen. The Vespa codepath limited total\n # highlights length to 400 chars. fragment_size *\n # number_of_fragments = 400 should be good enough.\n \"fragment_size\": 100,\n # The number of snippets to return per field per document\n # hit.\n \"number_of_fragments\": 4,\n # These tags wrap matched keywords and they match what Vespa\n # used to return. Use them to minimize changes to our code.\n \"pre_tags\": [\"\"],\n \"post_tags\": [\"\"],\n }\n }\n }\n\n return match_highlights_configuration\n" }, { "path": "backend/onyx/document_index/opensearch/string_filtering.py", "content": "import re\n\nMAX_DOCUMENT_ID_ENCODED_LENGTH: int = 512\n\n\nclass DocumentIDTooLongError(ValueError):\n \"\"\"Raised when a document ID is too long for OpenSearch after filtering.\"\"\"\n\n\ndef filter_and_validate_document_id(\n document_id: str, max_encoded_length: int = MAX_DOCUMENT_ID_ENCODED_LENGTH\n) -> str:\n \"\"\"\n Filters and validates a document ID such that it can be used as an ID in\n OpenSearch.\n\n OpenSearch imposes the following restrictions on IDs:\n - Must not be an empty string.\n - Must not exceed 512 bytes.\n - Must not contain any control characters (newline, etc.).\n - Must not contain URL-unsafe characters (#, ?, /, %, &, etc.).\n\n For extra resilience, this function simply removes all characters that are\n not alphanumeric or one of _.-~.\n\n Any query on document ID should use this function.\n\n Args:\n document_id: The document ID to filter and validate.\n max_encoded_length: The maximum length of the document ID after\n filtering in bytes. Compared with >= for extra resilience, so\n encoded values of this length will fail.\n\n Raises:\n DocumentIDTooLongError: If the document ID is too long after filtering.\n ValueError: If the document ID is empty after filtering.\n\n Returns:\n str: The filtered document ID.\n \"\"\"\n filtered_document_id = re.sub(r\"[^A-Za-z0-9_.\\-~]\", \"\", document_id)\n if not filtered_document_id:\n raise ValueError(f\"Document ID {document_id} is empty after filtering.\")\n if len(filtered_document_id.encode(\"utf-8\")) >= max_encoded_length:\n raise DocumentIDTooLongError(\n f\"Document ID {document_id} is too long after filtering.\"\n )\n return filtered_document_id\n" }, { "path": "backend/onyx/document_index/vespa/__init__.py", "content": "" }, { "path": "backend/onyx/document_index/vespa/app_config/schemas/danswer_chunk.sd.jinja", "content": "schema {{ schema_name }} {\n # source, type, target triplets for kg_relationships\n struct kg_relationship {\n field source type string {}\n field rel_type type string {}\n field target type string {}\n }\n\n document {{ schema_name }} {\n {% if multi_tenant %}\n field tenant_id type string {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n {% endif %}\n # Not to be confused with the UUID generated for this chunk which is called documentid by default\n field document_id type string {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n field chunk_id type int {\n indexing: summary | attribute\n }\n # Displayed in the UI as the main identifier for the doc\n field semantic_identifier type string {\n indexing: summary | attribute\n }\n # Must have an additional field for whether to skip title embeddings\n # This information cannot be extracted from either the title field nor title embedding\n field skip_title type bool {\n indexing: attribute\n }\n # May not always match the `semantic_identifier` e.g. for Slack docs the\n # `semantic_identifier` will be the channel name, but the `title` will be empty\n field title type string {\n indexing: summary | index | attribute\n index: enable-bm25\n }\n field content type string {\n indexing: summary | index\n index: enable-bm25\n }\n # duplication of `content` is far from ideal, but is needed for\n # non-gram based highlighting for now. If the capability to re-use a\n # single field to do both is added, `content_summary` should be removed\n field content_summary type string {\n indexing: summary | index\n summary: dynamic\n }\n # Title embedding (x1)\n field title_embedding type tensor<{{ embedding_precision }}>(x[{{ dim }}]) {\n indexing: attribute | index\n attribute {\n distance-metric: angular\n }\n }\n # Content embeddings (chunk + optional mini chunks embeddings)\n # \"t\" and \"x\" are arbitrary names, not special keywords\n field embeddings type tensor<{{ embedding_precision }}>(t{},x[{{ dim }}]) {\n indexing: attribute | index\n attribute {\n distance-metric: angular\n }\n }\n # Starting section of the doc, currently unused as it has been replaced by match highlighting\n field blurb type string {\n indexing: summary | attribute\n }\n field image_file_name type string {\n indexing: summary | attribute\n }\n # https://docs.vespa.ai/en/attributes.html potential enum store for speed, but probably not worth it\n field source_type type string {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n # Can also index links https://docs.vespa.ai/en/reference/schema-reference.html#attribute\n # URL type matching\n field source_links type string {\n indexing: summary | attribute\n }\n field section_continuation type bool {\n indexing: summary | attribute\n }\n # Technically this one should be int, but can't change without causing breaks to existing index\n field boost type float {\n indexing: summary | attribute\n }\n field hidden type bool {\n indexing: summary | attribute\n rank: filter\n }\n # Field to indicate whether a short chunk is a low content chunk\n field aggregated_chunk_boost_factor type float {\n indexing: attribute\n }\n\n # Separate array fields for knowledge graph data\n field kg_entities type array {\n indexing: summary | attribute\n attribute: fast-search\n }\n\n field kg_relationships type array {\n indexing: summary\n struct-field source {\n indexing: attribute\n attribute: fast-search\n }\n struct-field rel_type {\n indexing: attribute\n attribute: fast-search\n }\n struct-field target {\n indexing: attribute\n attribute: fast-search\n }\n }\n\n field kg_terms type array {\n indexing: summary | attribute\n attribute: fast-search\n }\n\n # Needs to have a separate Attribute list for efficient filtering\n field metadata_list type array {\n indexing: summary | attribute\n rank:filter\n attribute: fast-search\n }\n # If chunk is a large chunk, this will contain the ids of the smaller chunks\n field large_chunk_reference_ids type array {\n indexing: summary | attribute\n }\n field metadata type string {\n indexing: summary | attribute\n }\n field chunk_context type string {\n indexing: summary | attribute\n }\n field doc_summary type string {\n indexing: summary | attribute\n }\n field metadata_suffix type string {\n indexing: summary | attribute\n }\n field doc_updated_at type int {\n indexing: summary | attribute\n }\n field primary_owners type array {\n indexing: summary | attribute\n }\n field secondary_owners type array {\n indexing: summary | attribute\n }\n field access_control_list type weightedset {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n field document_sets type weightedset {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n field user_file type int {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n field user_folder type int {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n field user_project type array {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n field personas type array {\n indexing: summary | attribute\n rank: filter\n attribute: fast-search\n }\n }\n\n # If using different tokenization settings, the fieldset has to be removed, and the field must\n # be specified in the yql like:\n # + 'or ({grammar: \"weakAnd\", defaultIndex:\"title\"}userInput(@query)) '\n # + 'or ({grammar: \"weakAnd\", defaultIndex:\"content\"}userInput(@query)) '\n # Note: for BM-25, the ngram size (and whether ngrams are used) changes the range of the scores\n fieldset default {\n fields: content, title\n }\n\n rank-profile default_rank {\n inputs {\n query(decay_factor) double\n }\n\n function inline document_boost() {\n # 0.5 to 2x score: piecewise sigmoid function stretched out by factor of 3\n # meaning requires 3x the number of feedback votes to have default sigmoid effect\n expression: if(attribute(boost) < 0, 0.5 + (1 / (1 + exp(-attribute(boost) / 3))), 2 / (1 + exp(-attribute(boost) / 3)))\n }\n\n function inline document_age() {\n # Time in years (91.3 days ~= 3 Months ~= 1 fiscal quarter if no age found)\n expression: max(if(isNan(attribute(doc_updated_at)) == 1, 7890000, now() - attribute(doc_updated_at)) / 31536000, 0)\n }\n\n function inline aggregated_chunk_boost() {\n # Aggregated boost factor, currently only used for information content classification\n expression: if(isNan(attribute(aggregated_chunk_boost_factor)) == 1, 1.0, attribute(aggregated_chunk_boost_factor))\n }\n\n # Document score decays from 1 to 0.75 as age of last updated time increases\n function inline recency_bias() {\n expression: max(1 / (1 + query(decay_factor) * document_age), 0.75)\n }\n\n match-features: recency_bias\n }\n\n rank-profile hybrid_search_semantic_base_{{ dim }} inherits default, default_rank {\n inputs {\n query(query_embedding) tensor(x[{{ dim }}])\n }\n\n function title_vector_score() {\n expression {\n # If no good matching titles, then it should use the context embeddings rather than having some\n # irrelevant title have a vector score of 1. This way at least it will be the doc with the highest\n # matching content score getting the full score\n max(closeness(field, embeddings), closeness(field, title_embedding))\n }\n }\n\n # First phase must be vector to allow hits that have no keyword matches\n first-phase {\n expression: query(title_content_ratio) * closeness(field, title_embedding) + (1 - query(title_content_ratio)) * closeness(field, embeddings)\n }\n\n # Weighted average between Vector Search and BM-25\n global-phase {\n expression {\n (\n # Weighted Vector Similarity Score\n (\n query(alpha) * (\n (query(title_content_ratio) * normalize_linear(title_vector_score))\n +\n ((1 - query(title_content_ratio)) * normalize_linear(closeness(field, embeddings)))\n )\n )\n\n +\n\n # Weighted Keyword Similarity Score\n # Note: for the BM25 Title score, it requires decent stopword removal in the query\n # This needs to be the case so there aren't irrelevant titles being normalized to a score of 1\n (\n (1 - query(alpha)) * (\n (query(title_content_ratio) * normalize_linear(bm25(title)))\n +\n ((1 - query(title_content_ratio)) * normalize_linear(bm25(content)))\n )\n )\n )\n # Boost based on user feedback\n * document_boost\n # Decay factor based on time document was last updated\n * recency_bias\n # Boost based on aggregated boost calculation\n * aggregated_chunk_boost\n }\n # Target hits for hybrid retrieval should be at least this value.\n rerank-count: 1000\n }\n\n match-features {\n bm25(title)\n bm25(content)\n closeness(field, title_embedding)\n closeness(field, embeddings)\n document_boost\n recency_bias\n aggregated_chunk_boost\n closest(embeddings)\n }\n }\n\n\n rank-profile hybrid_search_keyword_base_{{ dim }} inherits default, default_rank {\n inputs {\n query(query_embedding) tensor(x[{{ dim }}])\n }\n\n function title_vector_score() {\n expression {\n # If no good matching titles, then it should use the context embeddings rather than having some\n # irrelevant title have a vector score of 1. This way at least it will be the doc with the highest\n # matching content score getting the full score\n max(closeness(field, embeddings), closeness(field, title_embedding))\n }\n }\n\n # First phase must be vector to allow hits that have no keyword matches\n first-phase {\n expression: query(title_content_ratio) * bm25(title) + (1 - query(title_content_ratio)) * bm25(content)\n }\n\n # Weighted average between Vector Search and BM-25\n global-phase {\n expression {\n (\n # Weighted Vector Similarity Score\n (\n query(alpha) * (\n (query(title_content_ratio) * normalize_linear(title_vector_score))\n +\n ((1 - query(title_content_ratio)) * normalize_linear(closeness(field, embeddings)))\n )\n )\n\n +\n\n # Weighted Keyword Similarity Score\n # Note: for the BM25 Title score, it requires decent stopword removal in the query\n # This needs to be the case so there aren't irrelevant titles being normalized to a score of 1\n (\n (1 - query(alpha)) * (\n (query(title_content_ratio) * normalize_linear(bm25(title)))\n +\n ((1 - query(title_content_ratio)) * normalize_linear(bm25(content)))\n )\n )\n )\n # Boost based on user feedback\n * document_boost\n # Decay factor based on time document was last updated\n * recency_bias\n # Boost based on aggregated boost calculation\n * aggregated_chunk_boost\n }\n # Target hits for hybrid retrieval should be at least this value.\n rerank-count: 1000\n }\n\n match-features {\n bm25(title)\n bm25(content)\n closeness(field, title_embedding)\n closeness(field, embeddings)\n document_boost\n recency_bias\n aggregated_chunk_boost\n closest(embeddings)\n }\n }\n\n # Used when searching from the admin UI for a specific doc to hide / boost\n # Very heavily prioritize title\n rank-profile admin_search inherits default, default_rank {\n first-phase {\n expression: bm25(content) + (5 * bm25(title))\n }\n }\n\n rank-profile random_ inherits default {\n first-phase {\n expression: random\n }\n }\n}\n" }, { "path": "backend/onyx/document_index/vespa/app_config/services.xml.jinja", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n 1\n \n \n {{ document_elements }}\n \n \n \n \n \n \n \n \n 0.85\n \n \n \n \n \n \n \n {{ num_search_threads }}\n \n \n \n \n \n \n 3\n 750\n 350\n 300\n \n \n" }, { "path": "backend/onyx/document_index/vespa/app_config/validation-overrides.xml.jinja", "content": "\n schema-removal\n indexing-change\n field-type-change\n\n" }, { "path": "backend/onyx/document_index/vespa/chunk_retrieval.py", "content": "import json\nimport string\nimport time\nfrom collections.abc import Callable\nfrom collections.abc import Mapping\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import cast\n\nimport httpx\nfrom retry import retry\n\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n FINISHED_VISITING_SLICE_CONTINUATION_TOKEN,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.transformer import (\n FIELDS_NEEDED_FOR_TRANSFORMATION,\n)\nfrom onyx.configs.app_configs import LOG_VESPA_TIMING_INFORMATION\nfrom onyx.configs.app_configs import VESPA_LANGUAGE_OVERRIDE\nfrom onyx.configs.app_configs import VESPA_MIGRATION_REQUEST_TIMEOUT_S\nfrom onyx.configs.app_configs import VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunkUncleaned\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\nfrom onyx.document_index.vespa.shared_utils.vespa_request_builders import (\n build_vespa_filters,\n)\nfrom onyx.document_index.vespa.shared_utils.vespa_request_builders import (\n build_vespa_id_based_retrieval_yql,\n)\nfrom onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST\nfrom onyx.document_index.vespa_constants import BLURB\nfrom onyx.document_index.vespa_constants import BOOST\nfrom onyx.document_index.vespa_constants import CHUNK_CONTEXT\nfrom onyx.document_index.vespa_constants import CHUNK_ID\nfrom onyx.document_index.vespa_constants import CONTENT\nfrom onyx.document_index.vespa_constants import CONTENT_SUMMARY\nfrom onyx.document_index.vespa_constants import DOC_SUMMARY\nfrom onyx.document_index.vespa_constants import DOC_UPDATED_AT\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.document_index.vespa_constants import HIDDEN\nfrom onyx.document_index.vespa_constants import IMAGE_FILE_NAME\nfrom onyx.document_index.vespa_constants import LARGE_CHUNK_REFERENCE_IDS\nfrom onyx.document_index.vespa_constants import MAX_ID_SEARCH_QUERY_SIZE\nfrom onyx.document_index.vespa_constants import MAX_OR_CONDITIONS\nfrom onyx.document_index.vespa_constants import METADATA\nfrom onyx.document_index.vespa_constants import METADATA_SUFFIX\nfrom onyx.document_index.vespa_constants import PRIMARY_OWNERS\nfrom onyx.document_index.vespa_constants import SEARCH_ENDPOINT\nfrom onyx.document_index.vespa_constants import SECONDARY_OWNERS\nfrom onyx.document_index.vespa_constants import SECTION_CONTINUATION\nfrom onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER\nfrom onyx.document_index.vespa_constants import SOURCE_LINKS\nfrom onyx.document_index.vespa_constants import SOURCE_TYPE\nfrom onyx.document_index.vespa_constants import TENANT_ID\nfrom onyx.document_index.vespa_constants import TITLE\nfrom onyx.document_index.vespa_constants import YQL_BASE\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef _process_dynamic_summary(\n dynamic_summary: str, max_summary_length: int = 400\n) -> list[str]:\n if not dynamic_summary:\n return []\n\n current_length = 0\n processed_summary: list[str] = []\n for summary_section in dynamic_summary.split(\"\"):\n # if we're past the desired max length, break at the last word\n if current_length + len(summary_section) >= max_summary_length:\n summary_section = summary_section[: max_summary_length - current_length]\n summary_section = summary_section.lstrip() # remove any leading whitespace\n\n # handle the case where the truncated section is either just a\n # single (partial) word or if it's empty\n first_space = summary_section.find(\" \")\n if first_space == -1:\n # add ``...`` to previous section\n if processed_summary:\n processed_summary[-1] += \"...\"\n break\n\n # handle the valid truncated section case\n summary_section = summary_section.rsplit(\" \", 1)[0]\n if summary_section[-1] in string.punctuation:\n summary_section = summary_section[:-1]\n summary_section += \"...\"\n processed_summary.append(summary_section)\n break\n\n processed_summary.append(summary_section)\n current_length += len(summary_section)\n\n return processed_summary\n\n\ndef _vespa_hit_to_inference_chunk(\n hit: dict[str, Any], null_score: bool = False\n) -> InferenceChunkUncleaned:\n fields = cast(dict[str, Any], hit[\"fields\"])\n\n # parse fields that are stored as strings, but are really json / datetime\n metadata = json.loads(fields[METADATA]) if METADATA in fields else {}\n updated_at = (\n datetime.fromtimestamp(fields[DOC_UPDATED_AT], tz=timezone.utc)\n if DOC_UPDATED_AT in fields\n else None\n )\n\n match_highlights = _process_dynamic_summary(\n # fallback to regular `content` if the `content_summary` field\n # isn't present\n dynamic_summary=hit[\"fields\"].get(CONTENT_SUMMARY, hit[\"fields\"][CONTENT]),\n )\n semantic_identifier = fields.get(SEMANTIC_IDENTIFIER, \"\")\n if not semantic_identifier:\n logger.error(\n f\"Chunk with blurb: {fields.get(BLURB, 'Unknown')[:50]}... has no Semantic Identifier\"\n )\n\n source_links = fields.get(SOURCE_LINKS, {})\n source_links_dict_unprocessed = (\n json.loads(source_links) if isinstance(source_links, str) else source_links\n )\n source_links_dict = {\n int(k): v\n for k, v in cast(dict[str, str], source_links_dict_unprocessed).items()\n }\n\n return InferenceChunkUncleaned(\n chunk_id=fields[CHUNK_ID],\n blurb=fields.get(BLURB, \"\"), # Unused\n content=fields[CONTENT], # Includes extra title prefix and metadata suffix;\n # also sometimes context for contextual rag\n source_links=source_links_dict or {0: \"\"},\n section_continuation=fields[SECTION_CONTINUATION],\n document_id=fields[DOCUMENT_ID],\n source_type=fields[SOURCE_TYPE],\n # still called `image_file_name` in Vespa for backwards compatibility\n image_file_id=fields.get(IMAGE_FILE_NAME),\n title=fields.get(TITLE),\n semantic_identifier=fields[SEMANTIC_IDENTIFIER],\n boost=fields.get(BOOST, 1),\n score=None if null_score else hit.get(\"relevance\", 0),\n hidden=fields.get(HIDDEN, False),\n primary_owners=fields.get(PRIMARY_OWNERS),\n secondary_owners=fields.get(SECONDARY_OWNERS),\n large_chunk_reference_ids=fields.get(LARGE_CHUNK_REFERENCE_IDS, []),\n metadata=metadata,\n metadata_suffix=fields.get(METADATA_SUFFIX),\n doc_summary=fields.get(DOC_SUMMARY, \"\"),\n chunk_context=fields.get(CHUNK_CONTEXT, \"\"),\n match_highlights=match_highlights,\n updated_at=updated_at,\n )\n\n\ndef get_chunks_via_visit_api(\n chunk_request: VespaChunkRequest,\n index_name: str,\n filters: IndexFilters,\n field_names: list[str] | None = None,\n get_large_chunks: bool = False,\n short_tensor_format: bool = False,\n) -> list[dict]:\n # Constructing the URL for the Visit API\n # NOTE: visit API uses the same URL as the document API, but with different params\n url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)\n\n # build the list of fields to retrieve\n field_set_list = (\n [f\"{field_name}\" for field_name in field_names] if field_names else []\n )\n acl_fieldset_entry = f\"{ACCESS_CONTROL_LIST}\"\n if (\n field_set_list\n and filters.access_control_list\n and acl_fieldset_entry not in field_set_list\n ):\n field_set_list.append(acl_fieldset_entry)\n\n if MULTI_TENANT:\n tenant_id_fieldset_entry = f\"{TENANT_ID}\"\n if field_set_list and tenant_id_fieldset_entry not in field_set_list:\n field_set_list.append(tenant_id_fieldset_entry)\n\n if field_set_list:\n field_set = f\"{index_name}:\" + \",\".join(field_set_list)\n else:\n field_set = None\n\n # build filters\n selection = f\"{index_name}.document_id=='{chunk_request.document_id}'\"\n\n if chunk_request.is_capped:\n selection += f\" and {index_name}.chunk_id>={chunk_request.min_chunk_ind or 0}\"\n selection += f\" and {index_name}.chunk_id<={chunk_request.max_chunk_ind}\"\n if not get_large_chunks:\n selection += f\" and {index_name}.large_chunk_reference_ids == null\"\n\n # enforcing tenant_id through a == condition\n if MULTI_TENANT:\n if filters.tenant_id:\n selection += f\" and {index_name}.tenant_id=='{filters.tenant_id}'\"\n else:\n raise ValueError(\"Tenant ID is required for multi-tenant\")\n\n # Setting up the selection criteria in the query parameters\n params = {\n # NOTE: Document Selector Language doesn't allow `contains`, so we can't check\n # for the ACL in the selection. Instead, we have to check as a postfilter\n \"selection\": selection,\n \"continuation\": None,\n \"wantedDocumentCount\": 1_000,\n \"fieldSet\": field_set,\n }\n # Vespa can supply tensors in various different formats. This explicitly\n # asks to retrieve tensor data in \"short-value\" format.\n if short_tensor_format:\n params[\"format.tensors\"] = \"short-value\"\n\n document_chunks: list[dict] = []\n while True:\n try:\n filtered_params = {k: v for k, v in params.items() if v is not None}\n with get_vespa_http_client() as http_client:\n response = http_client.get(url, params=filtered_params)\n response.raise_for_status()\n except httpx.HTTPError as e:\n error_base = \"Failed to query Vespa\"\n logger.error(\n f\"{error_base}:\\n\"\n f\"Request URL: {e.request.url}\\n\"\n f\"Request Headers: {e.request.headers}\\n\"\n f\"Request Payload: {params}\\n\"\n f\"Exception: {str(e)}\"\n )\n raise httpx.HTTPError(error_base) from e\n\n # Check if the response contains any documents\n response_data = response.json()\n\n if \"documents\" in response_data:\n for document in response_data[\"documents\"]:\n if filters.access_control_list:\n document_acl = document[\"fields\"].get(ACCESS_CONTROL_LIST)\n if not document_acl or not any(\n user_acl_entry in document_acl\n for user_acl_entry in filters.access_control_list\n ):\n continue\n\n if MULTI_TENANT:\n if not filters.tenant_id:\n raise ValueError(\"Tenant ID is required for multi-tenant\")\n document_tenant_id = document[\"fields\"].get(TENANT_ID)\n if document_tenant_id != filters.tenant_id:\n logger.error(\n f\"Skipping document {document['document_id']} because \"\n f\"it does not belong to tenant {filters.tenant_id}. \"\n \"This should never happen.\"\n )\n continue\n\n document_chunks.append(document)\n\n # Check for continuation token to handle pagination\n if \"continuation\" in response_data and response_data[\"continuation\"]:\n params[\"continuation\"] = response_data[\"continuation\"]\n else:\n break # Exit loop if no continuation token\n\n return document_chunks\n\n\ndef get_all_chunks_paginated(\n index_name: str,\n tenant_state: TenantState,\n continuation_token_map: dict[int, str | None],\n page_size: int,\n) -> tuple[list[dict], dict[int, str | None]]:\n \"\"\"Gets all chunks in Vespa matching the filters, paginated.\n\n Uses the Visit API with slicing. Each continuation token map entry is for a\n different slice. The number of entries determines the number of slices.\n\n Args:\n index_name: The name of the Vespa index to visit.\n tenant_state: The tenant state to filter by.\n continuation_token_map: Map of slice ID to a token returned by Vespa\n representing a page offset. None to start from the beginning of the\n slice.\n page_size: Best-effort batch size for the visit. Defaults to 1,000.\n\n Returns:\n Tuple of (list of chunk dicts, next continuation token or None). The\n continuation token is None when the visit is complete.\n \"\"\"\n\n def _get_all_chunks_paginated_for_slice(\n index_name: str,\n tenant_state: TenantState,\n slice_id: int,\n total_slices: int,\n continuation_token: str | None,\n page_size: int,\n ) -> tuple[list[dict], str | None]:\n if continuation_token == FINISHED_VISITING_SLICE_CONTINUATION_TOKEN:\n logger.debug(\n f\"Slice {slice_id} has finished visiting. Returning empty list and {FINISHED_VISITING_SLICE_CONTINUATION_TOKEN}.\"\n )\n return [], FINISHED_VISITING_SLICE_CONTINUATION_TOKEN\n\n url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)\n\n selection: str = f\"{index_name}.large_chunk_reference_ids == null\"\n if MULTI_TENANT:\n selection += f\" and {index_name}.tenant_id=='{tenant_state.tenant_id}'\"\n\n field_set = f\"{index_name}:\" + \",\".join(FIELDS_NEEDED_FOR_TRANSFORMATION)\n\n params: dict[str, str | int | None] = {\n \"selection\": selection,\n \"fieldSet\": field_set,\n \"wantedDocumentCount\": page_size,\n \"format.tensors\": \"short-value\",\n \"slices\": total_slices,\n \"sliceId\": slice_id,\n # When exceeded, Vespa should return gracefully with partial\n # results. Even if no hits are returned, Vespa should still return a\n # new continuation token representing a new spot in the linear\n # traversal.\n \"timeout\": VESPA_MIGRATION_SERVER_SIDE_REQUEST_TIMEOUT,\n }\n if continuation_token is not None:\n params[\"continuation\"] = continuation_token\n\n response: httpx.Response | None = None\n start_time = time.monotonic()\n try:\n with get_vespa_http_client(\n # When exceeded, an exception is raised in our code. No progress\n # is saved, and the task will retry this spot in the traversal\n # later.\n timeout=VESPA_MIGRATION_REQUEST_TIMEOUT_S\n ) as http_client:\n response = http_client.get(url, params=params)\n response.raise_for_status()\n except httpx.HTTPError as e:\n error_base = (\n f\"Failed to get chunks from Vespa slice {slice_id} with continuation token \"\n f\"{continuation_token} in {time.monotonic() - start_time:.3f} seconds.\"\n )\n logger.exception(\n f\"Request URL: {e.request.url}\\nRequest Headers: {e.request.headers}\\nRequest Payload: {params}\\n\"\n )\n error_message = (\n response.json().get(\"message\") if response else \"No response\"\n )\n logger.error(\"Error message from response: %s\", error_message)\n raise httpx.HTTPError(error_base) from e\n\n response_data = response.json()\n\n # NOTE: If we see a falsey value for \"continuation\" in the response we\n # assume we are done and return\n # FINISHED_VISITING_SLICE_CONTINUATION_TOKEN instead.\n next_continuation_token = (\n response_data.get(\"continuation\")\n or FINISHED_VISITING_SLICE_CONTINUATION_TOKEN\n )\n chunks = [chunk[\"fields\"] for chunk in response_data.get(\"documents\", [])]\n if next_continuation_token == FINISHED_VISITING_SLICE_CONTINUATION_TOKEN:\n logger.debug(\n f\"Slice {slice_id} has finished visiting. Returning {len(chunks)} chunks and {next_continuation_token}.\"\n )\n return chunks, next_continuation_token\n\n total_slices = len(continuation_token_map)\n if total_slices < 1:\n raise ValueError(\"continuation_token_map must have at least one entry.\")\n # We want to guarantee that these invocations are ordered by slice_id,\n # because we read in the same order below when parsing parallel_results.\n functions_with_args: list[tuple[Callable, tuple]] = [\n (\n _get_all_chunks_paginated_for_slice,\n (\n index_name,\n tenant_state,\n slice_id,\n total_slices,\n continuation_token,\n page_size,\n ),\n )\n for slice_id, continuation_token in sorted(continuation_token_map.items())\n ]\n\n parallel_results = run_functions_tuples_in_parallel(\n functions_with_args, allow_failures=True\n )\n if len(parallel_results) != total_slices:\n raise RuntimeError(\n f\"Expected {total_slices} parallel results, but got {len(parallel_results)}.\"\n )\n\n chunks: list[dict] = []\n next_continuation_token_map: dict[int, str | None] = {\n key: value for key, value in continuation_token_map.items()\n }\n for i, parallel_result in enumerate(parallel_results):\n if i not in next_continuation_token_map:\n raise RuntimeError(f\"Slice {i} is not in the continuation token map.\")\n if parallel_result is None:\n logger.error(\n f\"Failed to get chunks for slice {i} of {total_slices}. \"\n \"The continuation token for this slice will not be updated.\"\n )\n continue\n chunks.extend(parallel_result[0])\n next_continuation_token_map[i] = parallel_result[1]\n\n return chunks, next_continuation_token_map\n\n\n# TODO(rkuo): candidate for removal if not being used\n# @retry(tries=10, delay=1, backoff=2)\n# def get_all_vespa_ids_for_document_id(\n# document_id: str,\n# index_name: str,\n# filters: IndexFilters | None = None,\n# get_large_chunks: bool = False,\n# ) -> list[str]:\n# document_chunks = get_chunks_via_visit_api(\n# chunk_request=VespaChunkRequest(document_id=document_id),\n# index_name=index_name,\n# filters=filters or IndexFilters(access_control_list=None),\n# field_names=[DOCUMENT_ID],\n# get_large_chunks=get_large_chunks,\n# )\n# return [chunk[\"id\"].split(\"::\", 1)[-1] for chunk in document_chunks]\n\n\ndef parallel_visit_api_retrieval(\n index_name: str,\n chunk_requests: list[VespaChunkRequest],\n filters: IndexFilters,\n get_large_chunks: bool = False,\n) -> list[InferenceChunkUncleaned]:\n functions_with_args: list[tuple[Callable, tuple]] = [\n (\n get_chunks_via_visit_api,\n (chunk_request, index_name, filters, get_large_chunks),\n )\n for chunk_request in chunk_requests\n ]\n\n parallel_results = run_functions_tuples_in_parallel(\n functions_with_args, allow_failures=True\n )\n\n # Any failures to retrieve would give a None, drop the Nones and empty lists\n vespa_chunk_sets = [res for res in parallel_results if res]\n\n flattened_vespa_chunks = []\n for chunk_set in vespa_chunk_sets:\n flattened_vespa_chunks.extend(chunk_set)\n\n inference_chunks = [\n _vespa_hit_to_inference_chunk(chunk, null_score=True)\n for chunk in flattened_vespa_chunks\n ]\n\n return inference_chunks\n\n\n@retry(tries=3, delay=1, backoff=2)\ndef query_vespa(\n query_params: Mapping[str, str | int | float],\n) -> list[InferenceChunkUncleaned]:\n if \"query\" in query_params and not cast(str, query_params[\"query\"]).strip():\n raise ValueError(\"No/empty query received\")\n\n params = dict(\n **query_params,\n **(\n {\n \"presentation.timing\": True,\n }\n if LOG_VESPA_TIMING_INFORMATION\n else {}\n ),\n )\n\n if VESPA_LANGUAGE_OVERRIDE:\n params[\"language\"] = VESPA_LANGUAGE_OVERRIDE\n\n try:\n with get_vespa_http_client() as http_client:\n response = http_client.post(SEARCH_ENDPOINT, json=params)\n response.raise_for_status()\n except httpx.HTTPError as e:\n response_text = (\n e.response.text if isinstance(e, httpx.HTTPStatusError) else None\n )\n status_code = (\n e.response.status_code if isinstance(e, httpx.HTTPStatusError) else None\n )\n yql_value = params.get(\"yql\", \"\")\n yql_length = len(str(yql_value))\n\n # Log each detail on its own line so log collectors capture them\n # as separate entries rather than truncating a single multiline msg\n logger.error(\n f\"Failed to query Vespa | \"\n f\"status={status_code} | \"\n f\"yql_length={yql_length} | \"\n f\"exception={str(e)}\"\n )\n if response_text:\n logger.error(f\"Vespa error response: {response_text[:1000]}\")\n logger.error(f\"Vespa request URL: {e.request.url}\")\n\n # Re-raise with diagnostics so callers see what actually went wrong\n raise httpx.HTTPError(\n f\"Failed to query Vespa (status={status_code}, \" f\"yql_length={yql_length})\"\n ) from e\n\n response_json: dict[str, Any] = response.json()\n\n if LOG_VESPA_TIMING_INFORMATION:\n logger.debug(\"Vespa timing info: %s\", response_json.get(\"timing\"))\n hits = response_json[\"root\"].get(\"children\", [])\n\n if not hits:\n logger.warning(\n f\"No hits found for YQL Query: {query_params.get('yql', 'No YQL Query')}\"\n )\n logger.debug(f\"Vespa Response: {response.text}\")\n\n for hit in hits:\n if hit[\"fields\"].get(CONTENT) is None:\n identifier = hit[\"fields\"].get(\"documentid\") or hit[\"id\"]\n logger.error(\n f\"Vespa Index with Vespa ID {identifier} has no contents. \"\n f\"This is invalid because the vector is not meaningful and keywordsearch cannot \"\n f\"fetch this document\"\n )\n\n filtered_hits = [hit for hit in hits if hit[\"fields\"].get(CONTENT) is not None]\n\n inference_chunks = [_vespa_hit_to_inference_chunk(hit) for hit in filtered_hits]\n\n try:\n num_retrieved_inference_chunks = len(inference_chunks)\n num_retrieved_document_ids = len(\n set([chunk.document_id for chunk in inference_chunks])\n )\n logger.info(\n f\"Retrieved {num_retrieved_inference_chunks} inference chunks for {num_retrieved_document_ids} documents\"\n )\n except Exception as e:\n # Debug logging only, should not fail the retrieval\n logger.error(f\"Error logging retrieval statistics: {e}\")\n\n # Good Debugging Spot\n return inference_chunks\n\n\ndef _get_chunks_via_batch_search(\n index_name: str,\n chunk_requests: list[VespaChunkRequest],\n filters: IndexFilters,\n get_large_chunks: bool = False,\n) -> list[InferenceChunkUncleaned]:\n if not chunk_requests:\n return []\n\n filters_str = build_vespa_filters(filters=filters, include_hidden=True)\n\n yql = (\n YQL_BASE.format(index_name=index_name)\n + filters_str\n + build_vespa_id_based_retrieval_yql(chunk_requests[0])\n )\n chunk_requests.pop(0)\n\n for request in chunk_requests:\n yql += \" or \" + build_vespa_id_based_retrieval_yql(request)\n params: dict[str, str | int | float] = {\n \"yql\": yql,\n \"hits\": MAX_ID_SEARCH_QUERY_SIZE,\n }\n\n inference_chunks = query_vespa(params)\n if not get_large_chunks:\n inference_chunks = [\n chunk for chunk in inference_chunks if not chunk.large_chunk_reference_ids\n ]\n inference_chunks.sort(key=lambda chunk: chunk.chunk_id)\n return inference_chunks\n\n\ndef batch_search_api_retrieval(\n index_name: str,\n chunk_requests: list[VespaChunkRequest],\n filters: IndexFilters,\n get_large_chunks: bool = False,\n) -> list[InferenceChunkUncleaned]:\n retrieved_chunks: list[InferenceChunkUncleaned] = []\n capped_requests: list[VespaChunkRequest] = []\n uncapped_requests: list[VespaChunkRequest] = []\n chunk_count = 0\n for req_ind, request in enumerate(chunk_requests, start=1):\n # All requests without a chunk range are uncapped\n # Uncapped requests are retrieved using the Visit API\n range = request.range\n if range is None:\n uncapped_requests.append(request)\n continue\n\n if (\n chunk_count + range > MAX_ID_SEARCH_QUERY_SIZE\n or req_ind % MAX_OR_CONDITIONS == 0\n ):\n retrieved_chunks.extend(\n _get_chunks_via_batch_search(\n index_name=index_name,\n chunk_requests=capped_requests,\n filters=filters,\n get_large_chunks=get_large_chunks,\n )\n )\n capped_requests = []\n chunk_count = 0\n capped_requests.append(request)\n chunk_count += range\n\n if capped_requests:\n retrieved_chunks.extend(\n _get_chunks_via_batch_search(\n index_name=index_name,\n chunk_requests=capped_requests,\n filters=filters,\n get_large_chunks=get_large_chunks,\n )\n )\n\n if uncapped_requests:\n logger.debug(f\"Retrieving {len(uncapped_requests)} uncapped requests\")\n retrieved_chunks.extend(\n parallel_visit_api_retrieval(\n index_name, uncapped_requests, filters, get_large_chunks\n )\n )\n\n return retrieved_chunks\n" }, { "path": "backend/onyx/document_index/vespa/deletion.py", "content": "import concurrent.futures\nfrom uuid import UUID\n\nimport httpx\nfrom retry import retry\n\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.document_index.vespa_constants import NUM_THREADS\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nCONTENT_SUMMARY = \"content_summary\"\n\n\n@retry(tries=10, delay=1, backoff=2)\ndef _retryable_http_delete(http_client: httpx.Client, url: str) -> None:\n res = http_client.delete(url)\n res.raise_for_status()\n\n\ndef _delete_vespa_chunk(\n doc_chunk_id: UUID, index_name: str, http_client: httpx.Client\n) -> None:\n try:\n _retryable_http_delete(\n http_client,\n f\"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{doc_chunk_id}\",\n )\n except httpx.HTTPStatusError as e:\n logger.error(f\"Failed to delete chunk, details: {e.response.text}\")\n raise\n\n\ndef delete_vespa_chunks(\n doc_chunk_ids: list[UUID],\n index_name: str,\n http_client: httpx.Client,\n executor: concurrent.futures.ThreadPoolExecutor | None = None,\n) -> None:\n \"\"\"Deletes a list of chunks from a Vespa index in parallel.\n\n Args:\n doc_chunk_ids: List of chunk IDs to delete.\n index_name: Name of the index to delete from.\n http_client: HTTP client to use for the request.\n executor: Executor to use for the request.\n \"\"\"\n external_executor = True\n\n if not executor:\n external_executor = False\n executor = concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS)\n\n try:\n chunk_deletion_future = {\n executor.submit(\n _delete_vespa_chunk, doc_chunk_id, index_name, http_client\n ): doc_chunk_id\n for doc_chunk_id in doc_chunk_ids\n }\n for future in concurrent.futures.as_completed(chunk_deletion_future):\n # Will raise exception if the deletion raised an exception\n future.result()\n\n finally:\n if not external_executor:\n executor.shutdown(wait=True)\n" }, { "path": "backend/onyx/document_index/vespa/index.py", "content": "import concurrent.futures\nimport io\nimport logging\nimport os\nimport re\nimport time\nimport urllib\nimport zipfile\nfrom collections.abc import Iterable\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom typing import BinaryIO\nfrom typing import cast\nfrom typing import List\n\nimport httpx\nimport jinja2\nimport requests\nfrom pydantic import BaseModel\nfrom retry import retry\n\nfrom onyx.configs.app_configs import BLURB_SIZE\nfrom onyx.configs.chat_configs import NUM_RETURNED_HITS\nfrom onyx.configs.chat_configs import TITLE_CONTENT_RATIO\nfrom onyx.configs.chat_configs import VESPA_SEARCHER_THREADS\nfrom onyx.configs.constants import KV_REINDEX_KEY\nfrom onyx.configs.constants import RETURN_SEPARATOR\nfrom onyx.context.search.enums import QueryType\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceChunkUncleaned\nfrom onyx.context.search.models import QueryExpansionType\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.document_index.document_index_utils import get_uuid_from_chunk_info\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import (\n DocumentInsertionRecord as OldDocumentInsertionRecord,\n)\nfrom onyx.document_index.interfaces import EnrichedDocumentIndexingInfo\nfrom onyx.document_index.interfaces import IndexBatchParams\nfrom onyx.document_index.interfaces import MinimalDocumentIndexingInfo\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces import VespaDocumentFields\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\nfrom onyx.document_index.interfaces_new import DocumentSectionRequest\nfrom onyx.document_index.interfaces_new import IndexingMetadata\nfrom onyx.document_index.interfaces_new import MetadataUpdateRequest\nfrom onyx.document_index.vespa.chunk_retrieval import query_vespa\nfrom onyx.document_index.vespa.indexing_utils import BaseHTTPXClientContext\nfrom onyx.document_index.vespa.indexing_utils import check_for_final_chunk_existence\nfrom onyx.document_index.vespa.indexing_utils import GlobalHTTPXClientContext\nfrom onyx.document_index.vespa.indexing_utils import TemporaryHTTPXClientContext\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\nfrom onyx.document_index.vespa.shared_utils.vespa_request_builders import (\n build_vespa_filters,\n)\nfrom onyx.document_index.vespa.vespa_document_index import TenantState\nfrom onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex\nfrom onyx.document_index.vespa_constants import BATCH_SIZE\nfrom onyx.document_index.vespa_constants import CONTENT_SUMMARY\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.document_index.vespa_constants import NUM_THREADS\nfrom onyx.document_index.vespa_constants import VESPA_APPLICATION_ENDPOINT\nfrom onyx.document_index.vespa_constants import VESPA_TIMEOUT\nfrom onyx.document_index.vespa_constants import YQL_BASE\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.key_value_store.factory import get_shared_kv_store\nfrom onyx.kg.utils.formatting_utils import split_relationship_id\nfrom onyx.utils.batching import batch_generator\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom shared_configs.model_server_models import Embedding\n\nlogger = setup_logger()\n\n# Set the logging level to WARNING to ignore INFO and DEBUG logs\nhttpx_logger = logging.getLogger(\"httpx\")\nhttpx_logger.setLevel(logging.WARNING)\n\n\n@dataclass\nclass _VespaUpdateRequest:\n document_id: str\n url: str\n update_request: dict[str, dict]\n\n\nclass KGVespaChunkUpdateRequest(BaseModel):\n document_id: str\n chunk_id: int\n url: str\n update_request: dict[str, dict]\n\n\nclass KGUChunkUpdateRequest(BaseModel):\n \"\"\"\n Update KG fields for a document\n \"\"\"\n\n document_id: str\n chunk_id: int\n core_entity: str\n entities: set[str] | None = None\n relationships: set[str] | None = None\n terms: set[str] | None = None\n\n\nclass KGUDocumentUpdateRequest(BaseModel):\n \"\"\"\n Update KG fields for a document\n \"\"\"\n\n document_id: str\n entities: set[str]\n relationships: set[str]\n terms: set[str]\n\n\ndef generate_kg_update_request(\n kg_update_request: KGUChunkUpdateRequest,\n) -> dict[str, dict]:\n kg_update_dict: dict[str, dict] = {}\n\n if kg_update_request.entities is not None:\n kg_update_dict[\"kg_entities\"] = {\"assign\": list(kg_update_request.entities)}\n\n if kg_update_request.relationships is not None:\n kg_update_dict[\"kg_relationships\"] = {\"assign\": []}\n for relationship in kg_update_request.relationships:\n source, rel_type, target = split_relationship_id(relationship)\n kg_update_dict[\"kg_relationships\"][\"assign\"].append(\n {\n \"source\": source,\n \"rel_type\": rel_type,\n \"target\": target,\n }\n )\n\n return kg_update_dict\n\n\ndef in_memory_zip_from_file_bytes(file_contents: dict[str, bytes]) -> BinaryIO:\n zip_buffer = io.BytesIO()\n with zipfile.ZipFile(zip_buffer, \"w\", zipfile.ZIP_DEFLATED) as zipf:\n for filename, content in file_contents.items():\n zipf.writestr(filename, content)\n zip_buffer.seek(0)\n return zip_buffer\n\n\ndef _create_document_xml_lines(doc_names: list[str | None] | list[str]) -> str:\n doc_lines = [\n f''\n for doc_name in doc_names\n if doc_name\n ]\n return \"\\n\".join(doc_lines)\n\n\ndef add_ngrams_to_schema(schema_content: str) -> str:\n # Add the match blocks containing gram and gram-size to title and content fields\n schema_content = re.sub(\n r\"(field title type string \\{[^}]*indexing: summary \\| index \\| attribute)\",\n r\"\\1\\n match {\\n gram\\n gram-size: 3\\n }\",\n schema_content,\n )\n schema_content = re.sub(\n r\"(field content type string \\{[^}]*indexing: summary \\| index)\",\n r\"\\1\\n match {\\n gram\\n gram-size: 3\\n }\",\n schema_content,\n )\n return schema_content\n\n\ndef cleanup_chunks(chunks: list[InferenceChunkUncleaned]) -> list[InferenceChunk]:\n def _remove_title(chunk: InferenceChunkUncleaned) -> str:\n if not chunk.title or not chunk.content:\n return chunk.content\n\n if chunk.content.startswith(chunk.title):\n return chunk.content[len(chunk.title) :].lstrip()\n\n # BLURB SIZE is by token instead of char but each token is at least 1 char\n # If this prefix matches the content, it's assumed the title was prepended\n if chunk.content.startswith(chunk.title[:BLURB_SIZE]):\n return (\n chunk.content.split(RETURN_SEPARATOR, 1)[-1]\n if RETURN_SEPARATOR in chunk.content\n else chunk.content\n )\n\n return chunk.content\n\n def _remove_metadata_suffix(chunk: InferenceChunkUncleaned) -> str:\n if not chunk.metadata_suffix:\n return chunk.content\n return chunk.content.removesuffix(chunk.metadata_suffix).rstrip(\n RETURN_SEPARATOR\n )\n\n def _remove_contextual_rag(chunk: InferenceChunkUncleaned) -> str:\n # remove document summary\n if chunk.content.startswith(chunk.doc_summary):\n chunk.content = chunk.content[len(chunk.doc_summary) :].lstrip()\n # remove chunk context\n if chunk.content.endswith(chunk.chunk_context):\n chunk.content = chunk.content[\n : len(chunk.content) - len(chunk.chunk_context)\n ].rstrip()\n return chunk.content\n\n for chunk in chunks:\n chunk.content = _remove_title(chunk)\n chunk.content = _remove_metadata_suffix(chunk)\n chunk.content = _remove_contextual_rag(chunk)\n\n return [chunk.to_inference_chunk() for chunk in chunks]\n\n\nclass VespaIndex(DocumentIndex):\n VESPA_SCHEMA_JINJA_FILENAME = \"danswer_chunk.sd.jinja\"\n\n def __init__(\n self,\n index_name: str,\n secondary_index_name: str | None,\n large_chunks_enabled: bool,\n secondary_large_chunks_enabled: bool | None,\n multitenant: bool = False,\n httpx_client: httpx.Client | None = None,\n ) -> None:\n self.index_name = index_name\n self.secondary_index_name = secondary_index_name\n\n self.large_chunks_enabled = large_chunks_enabled\n self.secondary_large_chunks_enabled = secondary_large_chunks_enabled\n\n self.multitenant = multitenant\n\n # Temporary until we refactor the entirety of this class.\n self.httpx_client = httpx_client\n\n self.httpx_client_context: BaseHTTPXClientContext\n if httpx_client:\n self.httpx_client_context = GlobalHTTPXClientContext(httpx_client)\n else:\n self.httpx_client_context = TemporaryHTTPXClientContext(\n get_vespa_http_client\n )\n\n self.index_to_large_chunks_enabled: dict[str, bool] = {}\n self.index_to_large_chunks_enabled[index_name] = large_chunks_enabled\n if secondary_index_name and secondary_large_chunks_enabled:\n self.index_to_large_chunks_enabled[secondary_index_name] = (\n secondary_large_chunks_enabled\n )\n\n def ensure_indices_exist(\n self,\n primary_embedding_dim: int,\n primary_embedding_precision: EmbeddingPrecision,\n secondary_index_embedding_dim: int | None,\n secondary_index_embedding_precision: EmbeddingPrecision | None,\n ) -> None:\n if MULTI_TENANT:\n logger.info(\n \"Skipping Vespa index setup for multitenant (would wipe all indices)\"\n )\n return None\n\n jinja_env = jinja2.Environment()\n\n deploy_url = f\"{VESPA_APPLICATION_ENDPOINT}/tenant/default/prepareandactivate\"\n logger.notice(f\"Deploying Vespa application package to {deploy_url}\")\n\n vespa_schema_path = os.path.join(\n os.getcwd(), \"onyx\", \"document_index\", \"vespa\", \"app_config\"\n )\n schema_jinja_file = os.path.join(\n vespa_schema_path, \"schemas\", VespaIndex.VESPA_SCHEMA_JINJA_FILENAME\n )\n services_jinja_file = os.path.join(vespa_schema_path, \"services.xml.jinja\")\n overrides_jinja_file = os.path.join(\n vespa_schema_path, \"validation-overrides.xml.jinja\"\n )\n\n with open(services_jinja_file, \"r\") as services_f:\n schema_names = [self.index_name, self.secondary_index_name]\n doc_lines = _create_document_xml_lines(schema_names)\n\n services_template_str = services_f.read()\n services_template = jinja_env.from_string(services_template_str)\n services = services_template.render(\n document_elements=doc_lines,\n num_search_threads=str(VESPA_SEARCHER_THREADS),\n )\n\n kv_store = get_shared_kv_store()\n\n needs_reindexing = False\n try:\n needs_reindexing = cast(bool, kv_store.load(KV_REINDEX_KEY))\n except Exception:\n logger.debug(\"Could not load the reindexing flag. Using ngrams\")\n\n # Vespa requires an override to erase data including the indices we're no longer using\n # It also has a 30 day cap from current so we set it to 7 dynamically\n with open(overrides_jinja_file, \"r\") as overrides_f:\n overrides_template_str = overrides_f.read()\n overrides_template = jinja_env.from_string(overrides_template_str)\n\n now = datetime.now()\n date_in_7_days = now + timedelta(days=7)\n formatted_date = date_in_7_days.strftime(\"%Y-%m-%d\")\n overrides = overrides_template.render(\n until_date=formatted_date,\n )\n\n zip_dict = {\n \"services.xml\": services.encode(\"utf-8\"),\n \"validation-overrides.xml\": overrides.encode(\"utf-8\"),\n }\n\n with open(schema_jinja_file, \"r\") as schema_f:\n template_str = schema_f.read()\n\n template = jinja_env.from_string(template_str)\n schema = template.render(\n multi_tenant=MULTI_TENANT,\n schema_name=self.index_name,\n dim=primary_embedding_dim,\n embedding_precision=primary_embedding_precision.value,\n )\n\n schema = add_ngrams_to_schema(schema) if needs_reindexing else schema\n zip_dict[f\"schemas/{schema_names[0]}.sd\"] = schema.encode(\"utf-8\")\n\n if self.secondary_index_name:\n if secondary_index_embedding_dim is None:\n raise ValueError(\"Secondary index embedding dimension is required\")\n if secondary_index_embedding_precision is None:\n raise ValueError(\"Secondary index embedding precision is required\")\n\n upcoming_schema = template.render(\n multi_tenant=MULTI_TENANT,\n schema_name=self.secondary_index_name,\n dim=secondary_index_embedding_dim,\n embedding_precision=secondary_index_embedding_precision.value,\n )\n\n zip_dict[f\"schemas/{schema_names[1]}.sd\"] = upcoming_schema.encode(\"utf-8\")\n\n zip_file = in_memory_zip_from_file_bytes(zip_dict)\n\n headers = {\"Content-Type\": \"application/zip\"}\n response = requests.post(deploy_url, headers=headers, data=zip_file)\n if response.status_code != 200:\n logger.error(\n f\"Failed to prepare Vespa Onyx Index. Response: {response.text}\"\n )\n raise RuntimeError(\n f\"Failed to prepare Vespa Onyx Index. Response: {response.text}\"\n )\n\n @staticmethod\n def register_multitenant_indices(\n indices: list[str],\n embedding_dims: list[int],\n embedding_precisions: list[EmbeddingPrecision],\n ) -> None:\n if not MULTI_TENANT:\n raise ValueError(\"Multi-tenant is not enabled\")\n\n deploy_url = f\"{VESPA_APPLICATION_ENDPOINT}/tenant/default/prepareandactivate\"\n logger.info(f\"Deploying Vespa application package to {deploy_url}\")\n\n vespa_schema_path = os.path.join(\n os.getcwd(), \"onyx\", \"document_index\", \"vespa\", \"app_config\"\n )\n schema_jinja_file = os.path.join(\n vespa_schema_path, \"schemas\", VespaIndex.VESPA_SCHEMA_JINJA_FILENAME\n )\n services_jinja_file = os.path.join(vespa_schema_path, \"services.xml.jinja\")\n overrides_jinja_file = os.path.join(\n vespa_schema_path, \"validation-overrides.xml.jinja\"\n )\n\n jinja_env = jinja2.Environment()\n\n # Generate schema names from index settings\n with open(services_jinja_file, \"r\") as services_f:\n schema_names = [index_name for index_name in indices]\n doc_lines = _create_document_xml_lines(schema_names)\n\n services_template_str = services_f.read()\n services_template = jinja_env.from_string(services_template_str)\n services = services_template.render(\n document_elements=doc_lines,\n num_search_threads=str(VESPA_SEARCHER_THREADS),\n )\n\n kv_store = get_shared_kv_store()\n\n needs_reindexing = False\n try:\n needs_reindexing = cast(bool, kv_store.load(KV_REINDEX_KEY))\n except Exception:\n logger.debug(\"Could not load the reindexing flag. Using ngrams\")\n\n # Vespa requires an override to erase data including the indices we're no longer using\n # It also has a 30 day cap from current so we set it to 7 dynamically\n with open(overrides_jinja_file, \"r\") as overrides_f:\n overrides_template_str = overrides_f.read()\n overrides_template = jinja_env.from_string(overrides_template_str)\n\n now = datetime.now()\n date_in_7_days = now + timedelta(days=7)\n formatted_date = date_in_7_days.strftime(\"%Y-%m-%d\")\n overrides = overrides_template.render(\n until_date=formatted_date,\n )\n\n zip_dict = {\n \"services.xml\": services.encode(\"utf-8\"),\n \"validation-overrides.xml\": overrides.encode(\"utf-8\"),\n }\n\n with open(schema_jinja_file, \"r\") as schema_f:\n schema_template_str = schema_f.read()\n\n schema_template = jinja_env.from_string(schema_template_str)\n\n for i, index_name in enumerate(indices):\n embedding_dim = embedding_dims[i]\n embedding_precision = embedding_precisions[i]\n logger.info(\n f\"Creating index: {index_name} with embedding dimension: {embedding_dim}\"\n )\n\n schema = schema_template.render(\n multi_tenant=MULTI_TENANT,\n schema_name=index_name,\n dim=embedding_dim,\n embedding_precision=embedding_precision.value,\n )\n\n schema = add_ngrams_to_schema(schema) if needs_reindexing else schema\n zip_dict[f\"schemas/{index_name}.sd\"] = schema.encode(\"utf-8\")\n\n zip_file = in_memory_zip_from_file_bytes(zip_dict)\n\n headers = {\"Content-Type\": \"application/zip\"}\n response = requests.post(deploy_url, headers=headers, data=zip_file)\n\n if response.status_code != 200:\n raise RuntimeError(\n f\"Failed to prepare Vespa Onyx Indexes. Response: {response.text}\"\n )\n\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk],\n index_batch_params: IndexBatchParams,\n ) -> set[OldDocumentInsertionRecord]:\n \"\"\"\n NOTE: Do NOT consider the secondary index here. A separate indexing\n pipeline will be responsible for indexing to the secondary index. This\n design is not ideal and we should reconsider this when revamping index\n swapping.\n \"\"\"\n if len(index_batch_params.doc_id_to_previous_chunk_cnt) != len(\n index_batch_params.doc_id_to_new_chunk_cnt\n ):\n raise ValueError(\"Bug: Length of doc ID to chunk maps does not match.\")\n doc_id_to_chunk_cnt_diff = {\n doc_id: IndexingMetadata.ChunkCounts(\n old_chunk_cnt=index_batch_params.doc_id_to_previous_chunk_cnt[doc_id],\n new_chunk_cnt=index_batch_params.doc_id_to_new_chunk_cnt[doc_id],\n )\n for doc_id in index_batch_params.doc_id_to_previous_chunk_cnt.keys()\n }\n indexing_metadata = IndexingMetadata(\n doc_id_to_chunk_cnt_diff=doc_id_to_chunk_cnt_diff,\n )\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(),\n multitenant=MULTI_TENANT,\n )\n if tenant_state.multitenant != self.multitenant:\n raise ValueError(\n f\"Bug: Multitenant mismatch. Expected {tenant_state.multitenant}, got {self.multitenant}.\"\n )\n if (\n tenant_state.multitenant\n and tenant_state.tenant_id != index_batch_params.tenant_id\n ):\n raise ValueError(\n f\"Bug: Tenant ID mismatch. Expected {tenant_state.tenant_id}, got {index_batch_params.tenant_id}.\"\n )\n vespa_document_index = VespaDocumentIndex(\n index_name=self.index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=self.large_chunks_enabled,\n httpx_client=self.httpx_client,\n )\n # This conversion from list to set only to be converted again to a list\n # upstream is suboptimal and only temporary until we refactor the\n # entirety of this class.\n document_insertion_records = vespa_document_index.index(\n chunks, indexing_metadata\n )\n return set(\n [\n OldDocumentInsertionRecord(\n document_id=doc_insertion_record.document_id,\n already_existed=doc_insertion_record.already_existed,\n )\n for doc_insertion_record in document_insertion_records\n ]\n )\n\n @classmethod\n def _apply_updates_batched(\n cls,\n updates: list[_VespaUpdateRequest],\n httpx_client: httpx.Client,\n batch_size: int = BATCH_SIZE,\n ) -> None:\n \"\"\"Runs a batch of updates in parallel via the ThreadPoolExecutor.\"\"\"\n\n def _update_chunk(\n update: _VespaUpdateRequest, http_client: httpx.Client\n ) -> httpx.Response:\n logger.debug(\n f\"Updating with request to {update.url} with body {update.update_request}\"\n )\n return http_client.put(\n update.url,\n headers={\"Content-Type\": \"application/json\"},\n json=update.update_request,\n )\n\n # NOTE: using `httpx` here since `requests` doesn't support HTTP2. This is beneficient for\n # indexing / updates / deletes since we have to make a large volume of requests.\n\n with (\n concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor,\n httpx_client as http_client,\n ):\n for update_batch in batch_generator(updates, batch_size):\n future_to_document_id = {\n executor.submit(\n _update_chunk,\n update,\n http_client,\n ): update.document_id\n for update in update_batch\n }\n for future in concurrent.futures.as_completed(future_to_document_id):\n res = future.result()\n try:\n res.raise_for_status()\n except requests.HTTPError as e:\n failure_msg = f\"Failed to update document: {future_to_document_id[future]}\"\n raise requests.HTTPError(failure_msg) from e\n\n @classmethod\n def _apply_kg_chunk_updates_batched(\n cls,\n updates: list[KGVespaChunkUpdateRequest],\n httpx_client: httpx.Client,\n batch_size: int = BATCH_SIZE,\n ) -> None:\n \"\"\"Runs a batch of updates in parallel via the ThreadPoolExecutor.\"\"\"\n\n @retry(tries=3, delay=1, backoff=2, jitter=(0.0, 1.0))\n def _kg_update_chunk(\n update: KGVespaChunkUpdateRequest, http_client: httpx.Client\n ) -> httpx.Response:\n return http_client.put(\n update.url,\n headers={\"Content-Type\": \"application/json\"},\n json=update.update_request,\n )\n\n # NOTE: using `httpx` here since `requests` doesn't support HTTP2. This is beneficient for\n # indexing / updates / deletes since we have to make a large volume of requests.\n\n with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:\n for update_batch in batch_generator(updates, batch_size):\n future_to_document_id = {\n executor.submit(\n _kg_update_chunk,\n update,\n httpx_client,\n ): update.document_id\n for update in update_batch\n }\n for future in concurrent.futures.as_completed(future_to_document_id):\n res = future.result()\n try:\n res.raise_for_status()\n except requests.HTTPError as e:\n failure_msg = f\"Failed to update document {future_to_document_id[future]}\\nResponse: {res.text}\"\n raise requests.HTTPError(failure_msg) from e\n\n def kg_chunk_updates(\n self, kg_update_requests: list[KGUChunkUpdateRequest], tenant_id: str\n ) -> None:\n\n processed_updates_requests: list[KGVespaChunkUpdateRequest] = []\n update_start = time.monotonic()\n\n # Build the _VespaUpdateRequest objects\n\n for kg_update_request in kg_update_requests:\n kg_update_dict: dict[str, dict] = {\n \"fields\": generate_kg_update_request(kg_update_request)\n }\n if not kg_update_dict[\"fields\"]:\n logger.error(\"Update request received but nothing to update\")\n continue\n\n doc_chunk_id = get_uuid_from_chunk_info(\n document_id=kg_update_request.document_id,\n chunk_id=kg_update_request.chunk_id,\n tenant_id=tenant_id,\n large_chunk_id=None,\n )\n\n processed_updates_requests.append(\n KGVespaChunkUpdateRequest(\n document_id=kg_update_request.document_id,\n chunk_id=kg_update_request.chunk_id,\n url=f\"{DOCUMENT_ID_ENDPOINT.format(index_name=self.index_name)}/{doc_chunk_id}\",\n update_request=kg_update_dict,\n )\n )\n\n with self.httpx_client_context as httpx_client:\n self._apply_kg_chunk_updates_batched(\n processed_updates_requests, httpx_client\n )\n logger.debug(\n \"Updated %d vespa documents in %.2f seconds\",\n len(processed_updates_requests),\n time.monotonic() - update_start,\n )\n\n def update_single(\n self,\n doc_id: str,\n *,\n chunk_count: int | None,\n tenant_id: str,\n fields: VespaDocumentFields | None,\n user_fields: VespaDocumentUserFields | None,\n ) -> None:\n \"\"\"Note: if the document id does not exist, the update will be a no-op and the\n function will complete with no errors or exceptions.\n Handle other exceptions if you wish to implement retry behavior\n\n NOTE: Remember to handle the secondary index here. There is no separate\n pipeline for updating chunks in the secondary index. This design is not\n ideal and we should reconsider this when revamping index swapping.\n \"\"\"\n if fields is None and user_fields is None:\n logger.warning(\n f\"Tried to update document {doc_id} with no updated fields or user fields.\"\n )\n return\n\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(),\n multitenant=MULTI_TENANT,\n )\n if tenant_state.multitenant != self.multitenant:\n raise ValueError(\n f\"Bug: Multitenant mismatch. Expected {tenant_state.multitenant}, got {self.multitenant}.\"\n )\n if tenant_state.multitenant and tenant_state.tenant_id != tenant_id:\n raise ValueError(\n f\"Bug: Tenant ID mismatch. Expected {tenant_state.tenant_id}, got {tenant_id}.\"\n )\n\n project_ids: set[int] | None = None\n # NOTE: Empty user_projects is semantically different from None\n # user_projects.\n if user_fields is not None and user_fields.user_projects is not None:\n project_ids = set(user_fields.user_projects)\n persona_ids: set[int] | None = None\n # NOTE: Empty personas is semantically different from None personas.\n if user_fields is not None and user_fields.personas is not None:\n persona_ids = set(user_fields.personas)\n update_request = MetadataUpdateRequest(\n document_ids=[doc_id],\n doc_id_to_chunk_cnt={\n doc_id: chunk_count if chunk_count is not None else -1\n }, # NOTE: -1 represents an unknown chunk count.\n access=fields.access if fields is not None else None,\n document_sets=fields.document_sets if fields is not None else None,\n boost=fields.boost if fields is not None else None,\n hidden=fields.hidden if fields is not None else None,\n project_ids=project_ids,\n persona_ids=persona_ids,\n )\n\n indices = [self.index_name]\n if self.secondary_index_name:\n indices.append(self.secondary_index_name)\n\n for index_name in indices:\n vespa_document_index = VespaDocumentIndex(\n index_name=index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=self.index_to_large_chunks_enabled.get(\n index_name, False\n ),\n httpx_client=self.httpx_client,\n )\n vespa_document_index.update([update_request])\n\n def delete_single(\n self,\n doc_id: str,\n *,\n tenant_id: str,\n chunk_count: int | None,\n ) -> int:\n \"\"\"\n NOTE: Remember to handle the secondary index here. There is no separate\n pipeline for deleting chunks in the secondary index. This design is not\n ideal and we should reconsider this when revamping index swapping.\n \"\"\"\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(),\n multitenant=MULTI_TENANT,\n )\n if tenant_state.multitenant != self.multitenant:\n raise ValueError(\n f\"Bug: Multitenant mismatch. Expected {tenant_state.multitenant}, got {self.multitenant}.\"\n )\n if tenant_state.multitenant and tenant_state.tenant_id != tenant_id:\n raise ValueError(\n f\"Bug: Tenant ID mismatch. Expected {tenant_state.tenant_id}, got {tenant_id}.\"\n )\n indices = [self.index_name]\n if self.secondary_index_name:\n indices.append(self.secondary_index_name)\n\n total_chunks_deleted = 0\n for index_name in indices:\n vespa_document_index = VespaDocumentIndex(\n index_name=index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=self.index_to_large_chunks_enabled.get(\n index_name, False\n ),\n httpx_client=self.httpx_client,\n )\n total_chunks_deleted += vespa_document_index.delete(\n document_id=doc_id, chunk_count=chunk_count\n )\n\n return total_chunks_deleted\n\n def id_based_retrieval(\n self,\n chunk_requests: list[VespaChunkRequest],\n filters: IndexFilters,\n batch_retrieval: bool = False,\n get_large_chunks: bool = False, # noqa: ARG002\n ) -> list[InferenceChunk]:\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(),\n multitenant=MULTI_TENANT,\n )\n vespa_document_index = VespaDocumentIndex(\n index_name=self.index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=self.large_chunks_enabled,\n httpx_client=self.httpx_client,\n )\n generic_chunk_requests: list[DocumentSectionRequest] = []\n for chunk_request in chunk_requests:\n generic_chunk_requests.append(\n DocumentSectionRequest(\n document_id=chunk_request.document_id,\n min_chunk_ind=chunk_request.min_chunk_ind,\n max_chunk_ind=chunk_request.max_chunk_ind,\n )\n )\n return vespa_document_index.id_based_retrieval(\n chunk_requests=generic_chunk_requests,\n filters=filters,\n batch_retrieval=batch_retrieval,\n )\n\n @log_function_time(print_only=True, debug_only=True)\n def hybrid_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n final_keywords: list[str] | None,\n filters: IndexFilters,\n hybrid_alpha: float, # noqa: ARG002\n time_decay_multiplier: float, # noqa: ARG002\n num_to_retrieve: int,\n ranking_profile_type: QueryExpansionType = QueryExpansionType.SEMANTIC,\n title_content_ratio: float | None = TITLE_CONTENT_RATIO, # noqa: ARG002\n ) -> list[InferenceChunk]:\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(),\n multitenant=MULTI_TENANT,\n )\n vespa_document_index = VespaDocumentIndex(\n index_name=self.index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=self.large_chunks_enabled,\n httpx_client=self.httpx_client,\n )\n if not (\n ranking_profile_type == QueryExpansionType.KEYWORD\n or ranking_profile_type == QueryExpansionType.SEMANTIC\n ):\n raise ValueError(\n f\"Bug: Received invalid ranking profile type: {ranking_profile_type}\"\n )\n query_type = (\n QueryType.KEYWORD\n if ranking_profile_type == QueryExpansionType.KEYWORD\n else QueryType.SEMANTIC\n )\n return vespa_document_index.hybrid_retrieval(\n query,\n query_embedding,\n final_keywords,\n query_type,\n filters,\n num_to_retrieve,\n )\n\n def admin_retrieval(\n self,\n query: str,\n query_embedding: Embedding, # noqa: ARG002\n filters: IndexFilters,\n num_to_retrieve: int = NUM_RETURNED_HITS,\n ) -> list[InferenceChunk]:\n vespa_where_clauses = build_vespa_filters(filters, include_hidden=True)\n yql = (\n YQL_BASE.format(index_name=self.index_name)\n + vespa_where_clauses\n + '({grammar: \"weakAnd\"}userInput(@query) '\n # `({defaultIndex: \"content_summary\"}userInput(@query))` section is\n # needed for highlighting while the N-gram highlighting is broken /\n # not working as desired\n + f'or ({{defaultIndex: \"{CONTENT_SUMMARY}\"}}userInput(@query)))'\n )\n\n params: dict[str, str | int] = {\n \"yql\": yql,\n \"query\": query,\n \"hits\": num_to_retrieve,\n \"ranking.profile\": \"admin_search\",\n \"timeout\": VESPA_TIMEOUT,\n }\n\n return cleanup_chunks(query_vespa(params))\n\n # Retrieves chunk information for a document:\n # - Determines the last indexed chunk\n # - Identifies if the document uses the old or new chunk ID system\n # This data is crucial for Vespa document updates without relying on the visit API.\n @classmethod\n def enrich_basic_chunk_info(\n cls,\n index_name: str,\n http_client: httpx.Client,\n document_id: str,\n previous_chunk_count: int | None = None,\n new_chunk_count: int = 0,\n ) -> EnrichedDocumentIndexingInfo:\n last_indexed_chunk = previous_chunk_count\n\n # If the document has no `chunk_count` in the database, we know that it\n # has the old chunk ID system and we must check for the final chunk index\n is_old_version = False\n if last_indexed_chunk is None:\n is_old_version = True\n minimal_doc_info = MinimalDocumentIndexingInfo(\n doc_id=document_id, chunk_start_index=new_chunk_count\n )\n last_indexed_chunk = check_for_final_chunk_existence(\n minimal_doc_info=minimal_doc_info,\n start_index=new_chunk_count,\n index_name=index_name,\n http_client=http_client,\n )\n\n enriched_doc_info = EnrichedDocumentIndexingInfo(\n doc_id=document_id,\n chunk_start_index=new_chunk_count,\n chunk_end_index=last_indexed_chunk,\n old_version=is_old_version,\n )\n return enriched_doc_info\n\n @classmethod\n def delete_entries_by_tenant_id(\n cls,\n *,\n tenant_id: str,\n index_name: str,\n ) -> int:\n \"\"\"\n Deletes all entries in the specified index with the given tenant_id.\n\n Currently unused, but we anticipate this being useful. The entire flow does not\n use the httpx connection pool of an instance.\n\n Parameters:\n tenant_id (str): The tenant ID whose documents are to be deleted.\n index_name (str): The name of the index from which to delete documents.\n\n Returns:\n int: The number of documents deleted.\n \"\"\"\n logger.info(\n f\"Deleting entries with tenant_id: {tenant_id} from index: {index_name}\"\n )\n\n # Step 1: Retrieve all document IDs with the given tenant_id\n document_ids = cls._get_all_document_ids_by_tenant_id(tenant_id, index_name)\n\n if not document_ids:\n logger.info(\n f\"No documents found with tenant_id: {tenant_id} in index: {index_name}\"\n )\n return 0\n\n # Step 2: Delete documents in batches\n delete_requests = [\n _VespaDeleteRequest(document_id=doc_id, index_name=index_name)\n for doc_id in document_ids\n ]\n\n cls._apply_deletes_batched(delete_requests)\n return len(document_ids)\n\n @classmethod\n def _get_all_document_ids_by_tenant_id(\n cls, tenant_id: str, index_name: str\n ) -> List[str]:\n \"\"\"\n Retrieves all document IDs with the specified tenant_id, handling pagination.\n\n Internal helper function for delete_entries_by_tenant_id.\n\n Parameters:\n tenant_id (str): The tenant ID to search for.\n index_name (str): The name of the index to search in.\n\n Returns:\n List[str]: A list of document IDs matching the tenant_id.\n \"\"\"\n offset = 0\n limit = 1000 # Vespa's maximum hits per query\n document_ids = []\n\n logger.debug(\n f\"Starting document ID retrieval for tenant_id: {tenant_id} in index: {index_name}\"\n )\n\n while True:\n # Construct the query to fetch document IDs\n query_params = {\n \"yql\": f'select id from sources * where tenant_id contains \"{tenant_id}\";',\n \"offset\": str(offset),\n \"hits\": str(limit),\n \"timeout\": \"10s\",\n \"format\": \"json\",\n \"summary\": \"id\",\n }\n\n url = f\"{VESPA_APPLICATION_ENDPOINT}/search/\"\n\n logger.debug(\n f\"Querying for document IDs with tenant_id: {tenant_id}, offset: {offset}\"\n )\n\n with get_vespa_http_client() as http_client:\n response = http_client.get(url, params=query_params, timeout=None)\n response.raise_for_status()\n\n search_result = response.json()\n hits = search_result.get(\"root\", {}).get(\"children\", [])\n\n if not hits:\n break\n\n for hit in hits:\n doc_id = hit.get(\"id\")\n if doc_id:\n document_ids.append(doc_id)\n\n offset += limit # Move to the next page\n\n logger.debug(\n f\"Retrieved {len(document_ids)} document IDs for tenant_id: {tenant_id}\"\n )\n return document_ids\n\n @classmethod\n def _apply_deletes_batched(\n cls,\n delete_requests: List[\"_VespaDeleteRequest\"],\n batch_size: int = BATCH_SIZE,\n ) -> None:\n \"\"\"\n Deletes documents in batches using multiple threads.\n\n Internal helper function for delete_entries_by_tenant_id.\n\n This is a class method and does not use the httpx pool of the instance.\n This is OK because we don't use this method often.\n\n Parameters:\n delete_requests (List[_VespaDeleteRequest]): The list of delete requests.\n batch_size (int): The number of documents to delete in each batch.\n \"\"\"\n\n def _delete_document(\n delete_request: \"_VespaDeleteRequest\", http_client: httpx.Client\n ) -> None:\n logger.debug(f\"Deleting document with ID {delete_request.document_id}\")\n response = http_client.delete(\n delete_request.url,\n headers={\"Content-Type\": \"application/json\"},\n timeout=None,\n )\n response.raise_for_status()\n\n logger.debug(f\"Starting batch deletion for {len(delete_requests)} documents\")\n\n with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:\n with get_vespa_http_client() as http_client:\n for batch_start in range(0, len(delete_requests), batch_size):\n batch = delete_requests[batch_start : batch_start + batch_size]\n\n future_to_document_id = {\n executor.submit(\n _delete_document,\n delete_request,\n http_client,\n ): delete_request.document_id\n for delete_request in batch\n }\n\n for future in concurrent.futures.as_completed(\n future_to_document_id\n ):\n doc_id = future_to_document_id[future]\n try:\n future.result()\n logger.debug(f\"Successfully deleted document: {doc_id}\")\n except httpx.HTTPError as e:\n logger.error(f\"Failed to delete document {doc_id}: {e}\")\n # Optionally, implement retry logic or error handling here\n\n logger.info(\"Batch deletion completed\")\n\n def random_retrieval(\n self,\n filters: IndexFilters,\n num_to_retrieve: int = 10,\n ) -> list[InferenceChunk]:\n \"\"\"Retrieve random chunks matching the filters using Vespa's random ranking\n\n This method is currently used for random chunk retrieval in the context of\n assistant starter message creation (passed as sample context for usage by the assistant).\n \"\"\"\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(),\n multitenant=MULTI_TENANT,\n )\n vespa_document_index = VespaDocumentIndex(\n index_name=self.index_name,\n tenant_state=tenant_state,\n large_chunks_enabled=self.large_chunks_enabled,\n httpx_client=self.httpx_client,\n )\n return vespa_document_index.random_retrieval(\n filters=filters,\n num_to_retrieve=num_to_retrieve,\n )\n\n\nclass _VespaDeleteRequest:\n def __init__(self, document_id: str, index_name: str) -> None:\n self.document_id = document_id\n # Encode the document ID to ensure it's safe for use in the URL\n encoded_doc_id = urllib.parse.quote_plus(self.document_id)\n self.url = f\"{VESPA_APPLICATION_ENDPOINT}/document/v1/{index_name}/{index_name}/docid/{encoded_doc_id}\"\n" }, { "path": "backend/onyx/document_index/vespa/indexing_utils.py", "content": "import concurrent.futures\nimport json\nimport random\nimport time\nimport uuid\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom collections.abc import Callable\nfrom datetime import datetime\nfrom datetime import timezone\nfrom http import HTTPStatus\n\nimport httpx\nfrom retry import retry\n\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n get_experts_stores_representations,\n)\nfrom onyx.document_index.chunk_content_enrichment import (\n generate_enriched_content_for_chunk_text,\n)\nfrom onyx.document_index.document_index_utils import get_uuid_from_chunk\nfrom onyx.document_index.document_index_utils import get_uuid_from_chunk_info_old\nfrom onyx.document_index.interfaces import MinimalDocumentIndexingInfo\nfrom onyx.document_index.vespa.shared_utils.utils import (\n replace_invalid_doc_id_characters,\n)\nfrom onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST\nfrom onyx.document_index.vespa_constants import AGGREGATED_CHUNK_BOOST_FACTOR\nfrom onyx.document_index.vespa_constants import BLURB\nfrom onyx.document_index.vespa_constants import BOOST\nfrom onyx.document_index.vespa_constants import CHUNK_CONTEXT\nfrom onyx.document_index.vespa_constants import CHUNK_ID\nfrom onyx.document_index.vespa_constants import CONTENT\nfrom onyx.document_index.vespa_constants import CONTENT_SUMMARY\nfrom onyx.document_index.vespa_constants import DOC_SUMMARY\nfrom onyx.document_index.vespa_constants import DOC_UPDATED_AT\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.document_index.vespa_constants import DOCUMENT_SETS\nfrom onyx.document_index.vespa_constants import EMBEDDINGS\nfrom onyx.document_index.vespa_constants import FULL_CHUNK_EMBEDDING_KEY\nfrom onyx.document_index.vespa_constants import IMAGE_FILE_NAME\nfrom onyx.document_index.vespa_constants import LARGE_CHUNK_REFERENCE_IDS\nfrom onyx.document_index.vespa_constants import METADATA\nfrom onyx.document_index.vespa_constants import METADATA_LIST\nfrom onyx.document_index.vespa_constants import METADATA_SUFFIX\nfrom onyx.document_index.vespa_constants import NUM_THREADS\nfrom onyx.document_index.vespa_constants import PERSONAS\nfrom onyx.document_index.vespa_constants import PRIMARY_OWNERS\nfrom onyx.document_index.vespa_constants import SECONDARY_OWNERS\nfrom onyx.document_index.vespa_constants import SECTION_CONTINUATION\nfrom onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER\nfrom onyx.document_index.vespa_constants import SKIP_TITLE_EMBEDDING\nfrom onyx.document_index.vespa_constants import SOURCE_LINKS\nfrom onyx.document_index.vespa_constants import SOURCE_TYPE\nfrom onyx.document_index.vespa_constants import TENANT_ID\nfrom onyx.document_index.vespa_constants import TITLE\nfrom onyx.document_index.vespa_constants import TITLE_EMBEDDING\nfrom onyx.document_index.vespa_constants import USER_PROJECT\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.text_processing import remove_invalid_unicode_chars\n\n\nlogger = setup_logger()\n\n# Retry configuration constants\nINDEXING_MAX_RETRIES = 5\nINDEXING_BASE_DELAY = 1.0\nINDEXING_MAX_DELAY = 60.0\n\n\n@retry(tries=3, delay=1, backoff=2)\ndef _does_doc_chunk_exist(\n doc_chunk_id: uuid.UUID, index_name: str, http_client: httpx.Client\n) -> bool:\n doc_url = f\"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{doc_chunk_id}\"\n doc_fetch_response = http_client.get(doc_url)\n if doc_fetch_response.status_code == 404:\n return False\n\n if doc_fetch_response.status_code != 200:\n logger.debug(f\"Failed to check for document with URL {doc_url}\")\n raise RuntimeError(\n f\"Unexpected fetch document by ID value from Vespa: \"\n f\"error={doc_fetch_response.status_code} \"\n f\"index={index_name} \"\n f\"doc_chunk_id={doc_chunk_id}\"\n )\n return True\n\n\ndef _vespa_get_updated_at_attribute(t: datetime | None) -> int | None:\n if not t:\n return None\n\n if t.tzinfo != timezone.utc:\n raise ValueError(\"Connectors must provide document update time in UTC\")\n\n return int(t.timestamp())\n\n\ndef get_existing_documents_from_chunks(\n chunks: list[DocMetadataAwareIndexChunk],\n index_name: str,\n http_client: httpx.Client,\n executor: concurrent.futures.ThreadPoolExecutor | None = None,\n) -> set[str]:\n external_executor = True\n\n if not executor:\n external_executor = False\n executor = concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS)\n\n document_ids: set[str] = set()\n try:\n chunk_existence_future = {\n executor.submit(\n _does_doc_chunk_exist,\n get_uuid_from_chunk(chunk),\n index_name,\n http_client,\n ): chunk\n for chunk in chunks\n }\n for future in concurrent.futures.as_completed(chunk_existence_future):\n chunk = chunk_existence_future[future]\n chunk_already_existed = future.result()\n if chunk_already_existed:\n document_ids.add(chunk.source_document.id)\n\n finally:\n if not external_executor:\n executor.shutdown(wait=True)\n\n return document_ids\n\n\ndef _index_vespa_chunk(\n chunk: DocMetadataAwareIndexChunk,\n index_name: str,\n http_client: httpx.Client,\n multitenant: bool,\n) -> None:\n json_header = {\n \"Content-Type\": \"application/json\",\n }\n document = chunk.source_document\n\n # No minichunk documents in vespa, minichunk vectors are stored in the chunk itself\n\n vespa_chunk_id = str(get_uuid_from_chunk(chunk))\n\n embeddings = chunk.embeddings\n\n embeddings_name_vector_map = {FULL_CHUNK_EMBEDDING_KEY: embeddings.full_embedding}\n\n if embeddings.mini_chunk_embeddings:\n for ind, m_c_embed in enumerate(embeddings.mini_chunk_embeddings):\n embeddings_name_vector_map[f\"mini_chunk_{ind}\"] = m_c_embed\n\n title = document.get_title_for_document_index()\n\n metadata_json = document.metadata\n cleaned_metadata_json: dict[str, str | list[str]] = {}\n for key, value in metadata_json.items():\n cleaned_key = remove_invalid_unicode_chars(key)\n if isinstance(value, list):\n cleaned_metadata_json[cleaned_key] = [\n remove_invalid_unicode_chars(item) for item in value\n ]\n else:\n cleaned_metadata_json[cleaned_key] = remove_invalid_unicode_chars(value)\n\n metadata_list = document.get_metadata_str_attributes()\n if metadata_list:\n metadata_list = [\n remove_invalid_unicode_chars(metadata) for metadata in metadata_list\n ]\n\n vespa_document_fields = {\n DOCUMENT_ID: document.id,\n CHUNK_ID: chunk.chunk_id,\n BLURB: remove_invalid_unicode_chars(chunk.blurb),\n TITLE: remove_invalid_unicode_chars(title) if title else None,\n SKIP_TITLE_EMBEDDING: not title,\n # For the BM25 index, the keyword suffix is used, the vector is already generated with the more\n # natural language representation of the metadata section\n CONTENT: remove_invalid_unicode_chars(\n generate_enriched_content_for_chunk_text(chunk)\n ),\n # This duplication of `content` is needed for keyword highlighting\n # Note that it's not exactly the same as the actual content\n # which contains the title prefix and metadata suffix\n CONTENT_SUMMARY: remove_invalid_unicode_chars(chunk.content),\n SOURCE_TYPE: str(document.source.value),\n SOURCE_LINKS: json.dumps(chunk.source_links),\n SEMANTIC_IDENTIFIER: remove_invalid_unicode_chars(document.semantic_identifier),\n SECTION_CONTINUATION: chunk.section_continuation,\n LARGE_CHUNK_REFERENCE_IDS: chunk.large_chunk_reference_ids,\n METADATA: json.dumps(cleaned_metadata_json),\n # Save as a list for efficient extraction as an Attribute\n METADATA_LIST: metadata_list,\n METADATA_SUFFIX: remove_invalid_unicode_chars(chunk.metadata_suffix_keyword),\n CHUNK_CONTEXT: chunk.chunk_context,\n DOC_SUMMARY: chunk.doc_summary,\n EMBEDDINGS: embeddings_name_vector_map,\n TITLE_EMBEDDING: chunk.title_embedding,\n DOC_UPDATED_AT: _vespa_get_updated_at_attribute(document.doc_updated_at),\n PRIMARY_OWNERS: get_experts_stores_representations(document.primary_owners),\n SECONDARY_OWNERS: get_experts_stores_representations(document.secondary_owners),\n # the only `set` vespa has is `weightedset`, so we have to give each\n # element an arbitrary weight\n # rkuo: acl, docset and boost metadata are also updated through the metadata sync queue\n # which only calls VespaIndex.update\n ACCESS_CONTROL_LIST: {acl_entry: 1 for acl_entry in chunk.access.to_acl()},\n DOCUMENT_SETS: {document_set: 1 for document_set in chunk.document_sets},\n # still called `image_file_name` in Vespa for backwards compatibility\n IMAGE_FILE_NAME: chunk.image_file_id,\n USER_PROJECT: chunk.user_project if chunk.user_project is not None else [],\n PERSONAS: chunk.personas if chunk.personas is not None else [],\n BOOST: chunk.boost,\n AGGREGATED_CHUNK_BOOST_FACTOR: chunk.aggregated_chunk_boost_factor,\n }\n\n if multitenant:\n if chunk.tenant_id:\n vespa_document_fields[TENANT_ID] = chunk.tenant_id\n vespa_url = f\"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{vespa_chunk_id}\"\n logger.debug(f'Indexing to URL \"{vespa_url}\"')\n\n # Retry logic with exponential backoff for rate limiting\n for attempt in range(INDEXING_MAX_RETRIES):\n try:\n res = http_client.post(\n vespa_url, headers=json_header, json={\"fields\": vespa_document_fields}\n )\n res.raise_for_status()\n return # Success, exit the function\n except httpx.HTTPStatusError as e:\n # Handle 429 rate limiting specifically\n if e.response.status_code == HTTPStatus.TOO_MANY_REQUESTS:\n if attempt < INDEXING_MAX_RETRIES - 1:\n # Calculate exponential backoff with jitter\n delay = min(\n INDEXING_BASE_DELAY * (2**attempt), INDEXING_MAX_DELAY\n ) * random.uniform(0.5, 1.0)\n logger.warning(\n f\"Rate limited while indexing document '{document.id}' \"\n f\"(attempt {attempt + 1}/{INDEXING_MAX_RETRIES}). \"\n f\"Vespa response: '{e.response.text}'. \"\n f\"Backing off for {delay:.2f} seconds.\"\n )\n time.sleep(delay)\n continue\n else:\n raise RuntimeError(\n f\"Failed to index document '{document.id}' after {INDEXING_MAX_RETRIES} attempts due to rate limiting\"\n ) from e\n elif e.response.status_code == HTTPStatus.INSUFFICIENT_STORAGE:\n logger.error(\n f\"Failed to index document: '{document.id}'. Got response: '{e.response.text}'\"\n )\n logger.error(\n \"NOTE: HTTP Status 507 Insufficient Storage usually means \"\n \"you need to allocate more memory or disk space to the \"\n \"Vespa/index container.\"\n )\n raise\n else:\n # For other HTTP errors, check if retryable\n if e.response.status_code in (\n HTTPStatus.BAD_REQUEST,\n HTTPStatus.UNAUTHORIZED,\n HTTPStatus.FORBIDDEN,\n HTTPStatus.NOT_FOUND,\n ):\n # Non-retryable errors - fail immediately\n logger.error(\n f\"Non-retryable HTTP {e.response.status_code} error for document '{document.id}'\"\n )\n raise\n # Retry other errors with shorter backoff\n if attempt < INDEXING_MAX_RETRIES - 1:\n delay = INDEXING_BASE_DELAY * (1.5**attempt)\n logger.warning(\n f\"HTTP error {e.response.status_code} while indexing document '{document.id}' \"\n f\"(attempt {attempt + 1}/{INDEXING_MAX_RETRIES}). Retrying in {delay:.2f} seconds.\"\n )\n time.sleep(delay)\n continue\n else:\n logger.exception(\n f\"Failed to index document: '{document.id}'. Got response: '{e.response.text}'\"\n )\n raise\n except Exception as e:\n # For non-HTTP errors, use simple retry logic\n if attempt < INDEXING_MAX_RETRIES - 1:\n delay = INDEXING_BASE_DELAY * (1.5**attempt)\n logger.warning(\n f\"Error while indexing document '{document.id}' \"\n f\"(attempt {attempt + 1}/{INDEXING_MAX_RETRIES}): {str(e)}. \"\n f\"Retrying in {delay:.2f} seconds.\"\n )\n time.sleep(delay)\n continue\n else:\n logger.exception(f\"Failed to index document: '{document.id}'\")\n raise\n\n\ndef batch_index_vespa_chunks(\n chunks: list[DocMetadataAwareIndexChunk],\n index_name: str,\n http_client: httpx.Client,\n multitenant: bool,\n executor: concurrent.futures.ThreadPoolExecutor | None = None,\n) -> None:\n \"\"\"Indexes a list of chunks in a Vespa index in parallel.\n\n Args:\n chunks: List of chunks to index.\n index_name: Name of the index to index into.\n http_client: HTTP client to use for the request.\n multitenant: Whether the index is multitenant.\n executor: Executor to use for the request.\n \"\"\"\n external_executor = True\n\n if not executor:\n external_executor = False\n executor = concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS)\n\n try:\n chunk_index_future = {\n executor.submit(\n _index_vespa_chunk, chunk, index_name, http_client, multitenant\n ): chunk\n for chunk in chunks\n }\n for future in concurrent.futures.as_completed(chunk_index_future):\n # Will raise exception if any indexing raised an exception\n future.result()\n\n finally:\n if not external_executor:\n executor.shutdown(wait=True)\n\n\ndef clean_chunk_id_copy(\n chunk: DocMetadataAwareIndexChunk,\n) -> DocMetadataAwareIndexChunk:\n clean_chunk = chunk.model_copy(\n update={\n \"source_document\": chunk.source_document.model_copy(\n update={\n \"id\": replace_invalid_doc_id_characters(chunk.source_document.id)\n }\n )\n }\n )\n return clean_chunk\n\n\ndef check_for_final_chunk_existence(\n minimal_doc_info: MinimalDocumentIndexingInfo,\n start_index: int,\n index_name: str,\n http_client: httpx.Client,\n) -> int:\n index = start_index\n while True:\n doc_chunk_id = get_uuid_from_chunk_info_old(\n document_id=minimal_doc_info.doc_id,\n chunk_id=index,\n large_chunk_reference_ids=[],\n )\n if not _does_doc_chunk_exist(doc_chunk_id, index_name, http_client):\n return index\n index += 1\n\n\nclass BaseHTTPXClientContext(ABC):\n \"\"\"Abstract base class for an HTTPX client context manager.\"\"\"\n\n @abstractmethod\n def __enter__(self) -> httpx.Client:\n pass\n\n @abstractmethod\n def __exit__(self, exc_type, exc_value, traceback): # type: ignore\n pass\n\n\nclass GlobalHTTPXClientContext(BaseHTTPXClientContext):\n \"\"\"Context manager for a global HTTPX client that does not close it.\"\"\"\n\n def __init__(self, client: httpx.Client):\n self._client = client\n\n def __enter__(self) -> httpx.Client:\n return self._client # Reuse the global client\n\n def __exit__(self, exc_type, exc_value, traceback): # type: ignore\n pass # Do nothing; don't close the global client\n\n\nclass TemporaryHTTPXClientContext(BaseHTTPXClientContext):\n \"\"\"Context manager for a temporary HTTPX client that closes it after use.\"\"\"\n\n def __init__(self, client_factory: Callable[[], httpx.Client]):\n self._client_factory = client_factory\n self._client: httpx.Client | None = None # Client will be created in __enter__\n\n def __enter__(self) -> httpx.Client:\n self._client = self._client_factory() # Create a new client\n return self._client\n\n def __exit__(self, exc_type, exc_value, traceback): # type: ignore\n if self._client:\n self._client.close()\n" }, { "path": "backend/onyx/document_index/vespa/kg_interactions.py", "content": "from onyx.db.document import get_document_kg_entities_and_relationships\nfrom onyx.db.document import get_num_chunks_for_document\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.document_index.vespa.index import KGUChunkUpdateRequest\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef update_kg_chunks_vespa_info(\n kg_update_requests: list[KGUChunkUpdateRequest],\n index_name: str,\n tenant_id: str,\n) -> None:\n \"\"\" \"\"\"\n # Use the existing visit API infrastructure\n vespa_index = VespaIndex(\n index_name=index_name,\n secondary_index_name=None,\n large_chunks_enabled=False,\n secondary_large_chunks_enabled=False,\n multitenant=MULTI_TENANT,\n httpx_client=None,\n )\n\n vespa_index.kg_chunk_updates(\n kg_update_requests=kg_update_requests, tenant_id=tenant_id\n )\n\n\ndef get_kg_vespa_info_update_requests_for_document(\n document_id: str,\n) -> list[KGUChunkUpdateRequest]:\n \"\"\"Get the kg_info update requests for a document.\"\"\"\n # get all entities and relationships tied to the document\n with get_session_with_current_tenant() as db_session:\n entities, relationships = get_document_kg_entities_and_relationships(\n db_session, document_id\n )\n\n # create the kg vespa info\n kg_entities = {entity.id_name for entity in entities}\n kg_relationships = {relationship.id_name for relationship in relationships}\n\n # get chunks in the document\n with get_session_with_current_tenant() as db_session:\n num_chunks = get_num_chunks_for_document(db_session, document_id)\n\n # get vespa update requests\n return [\n KGUChunkUpdateRequest(\n document_id=document_id,\n chunk_id=chunk_id,\n core_entity=\"unused\",\n entities=kg_entities,\n relationships=kg_relationships or None,\n )\n for chunk_id in range(num_chunks)\n ]\n" }, { "path": "backend/onyx/document_index/vespa/shared_utils/utils.py", "content": "import time\nfrom typing import cast\n\nimport httpx\n\nfrom onyx.configs.app_configs import MANAGED_VESPA\nfrom onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH\nfrom onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH\nfrom onyx.configs.app_configs import VESPA_REQUEST_TIMEOUT\nfrom onyx.document_index.vespa_constants import VESPA_APP_CONTAINER_URL\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# NOTE: This does not seem to be used in reality despite the Vespa Docs pointing to this code\n# See here for reference: https://docs.vespa.ai/en/documents.html\n# https://github.com/vespa-engine/vespa/blob/master/vespajlib/src/main/java/com/yahoo/text/Text.java\n\n# Define allowed ASCII characters\nALLOWED_ASCII_CHARS: list[bool] = [False] * 0x80\nALLOWED_ASCII_CHARS[0x9] = True # tab\nALLOWED_ASCII_CHARS[0xA] = True # newline\nALLOWED_ASCII_CHARS[0xD] = True # carriage return\nfor i in range(0x20, 0x7F):\n ALLOWED_ASCII_CHARS[i] = True # printable ASCII chars\nALLOWED_ASCII_CHARS[0x7F] = True # del - discouraged, but allowed\n\n\ndef is_text_character(codepoint: int) -> bool:\n \"\"\"Returns whether the given codepoint is a valid text character.\"\"\"\n if codepoint < 0x80:\n return ALLOWED_ASCII_CHARS[codepoint]\n if codepoint < 0xD800:\n return True\n if codepoint <= 0xDFFF:\n return False\n if codepoint < 0xFDD0:\n return True\n if codepoint <= 0xFDEF:\n return False\n if codepoint >= 0x10FFFE:\n return False\n return (codepoint & 0xFFFF) < 0xFFFE\n\n\ndef replace_invalid_doc_id_characters(text: str) -> str:\n \"\"\"Replaces invalid document ID characters in text.\n NOTE: this must be called at the start of every vespa-related operation or else we\n risk discrepancies -> silent failures on deletion/update/insertion.\"\"\"\n # There may be a more complete set of replacements that need to be made but Vespa docs are unclear\n # and users only seem to be running into this error with single quotes\n return text.replace(\"'\", \"_\")\n\n\ndef get_vespa_http_client(\n no_timeout: bool = False, http2: bool = True, timeout: int | None = None\n) -> httpx.Client:\n \"\"\"\n Configures and returns an HTTP client for communicating with Vespa,\n including authentication if needed.\n \"\"\"\n return httpx.Client(\n cert=(\n cast(tuple[str, str], (VESPA_CLOUD_CERT_PATH, VESPA_CLOUD_KEY_PATH))\n if MANAGED_VESPA\n else None\n ),\n verify=False if not MANAGED_VESPA else True,\n timeout=None if no_timeout else (timeout or VESPA_REQUEST_TIMEOUT),\n http2=http2,\n )\n\n\ndef wait_for_vespa_with_timeout(wait_interval: int = 5, wait_limit: int = 60) -> bool:\n \"\"\"Waits for Vespa to become ready subject to a timeout.\n Returns True if Vespa is ready, False otherwise.\"\"\"\n\n time_start = time.monotonic()\n logger.info(\"Vespa: Readiness probe starting.\")\n while True:\n url = f\"{VESPA_APP_CONTAINER_URL}/state/v1/health\"\n try:\n client = get_vespa_http_client()\n response = client.get(url)\n response.raise_for_status()\n\n response_dict = response.json()\n if response_dict[\"status\"][\"code\"] == \"up\":\n logger.info(\"Vespa: Readiness probe succeeded. Continuing...\")\n return True\n except Exception as e:\n logger.warning(\n f\"Vespa: Readiness probe failed trying to connect to {url}. Exception: {e}\"\n )\n\n time_elapsed = time.monotonic() - time_start\n if time_elapsed > wait_limit:\n logger.info(\n f\"Vespa: Readiness probe did not succeed within the timeout ({wait_limit} seconds).\"\n )\n return False\n\n logger.info(\n f\"Vespa: Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={wait_limit:.1f}\"\n )\n\n time.sleep(wait_interval)\n" }, { "path": "backend/onyx/document_index/vespa/shared_utils/vespa_request_builders.py", "content": "from datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom onyx.configs.constants import INDEX_SEPARATOR\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST\nfrom onyx.document_index.vespa_constants import CHUNK_ID\nfrom onyx.document_index.vespa_constants import DOC_UPDATED_AT\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID\nfrom onyx.document_index.vespa_constants import DOCUMENT_SETS\nfrom onyx.document_index.vespa_constants import HIDDEN\nfrom onyx.document_index.vespa_constants import METADATA_LIST\nfrom onyx.document_index.vespa_constants import PERSONAS\nfrom onyx.document_index.vespa_constants import SOURCE_TYPE\nfrom onyx.document_index.vespa_constants import TENANT_ID\nfrom onyx.document_index.vespa_constants import USER_PROJECT\nfrom onyx.kg.utils.formatting_utils import split_relationship_id\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef build_tenant_id_filter(tenant_id: str) -> str:\n return f'({TENANT_ID} contains \"{tenant_id}\")'\n\n\ndef build_vespa_filters(\n filters: IndexFilters,\n *,\n include_hidden: bool = False,\n remove_trailing_and: bool = False, # Set to True when using as a complete Vespa query\n) -> str:\n def _build_or_filters(key: str, vals: list[str] | None) -> str:\n \"\"\"For string-based 'contains' filters, e.g. WSET fields or array fields.\n Returns a bare clause like '(key contains \"v1\" or key contains \"v2\")' or \"\".\"\"\"\n if not key or not vals:\n return \"\"\n eq_elems = [f'{key} contains \"{val}\"' for val in vals if val]\n if not eq_elems:\n return \"\"\n return f\"({' or '.join(eq_elems)})\"\n\n def _build_weighted_set_filter(key: str, vals: list[str] | None) -> str:\n \"\"\"Build a Vespa weightedSet filter for large value lists.\n\n Uses Vespa's native weightedSet() operator instead of OR-chained\n 'contains' clauses. This is critical for fields like\n access_control_list where a single user may have tens of thousands\n of ACL entries — OR clauses at that scale cause Vespa to reject\n the query with HTTP 400.\"\"\"\n if not key or not vals:\n return \"\"\n filtered = [val for val in vals if val]\n if not filtered:\n return \"\"\n items = \", \".join(f'\"{val}\":1' for val in filtered)\n return f\"weightedSet({key}, {{{items}}})\"\n\n def _build_int_or_filters(key: str, vals: list[int] | None) -> str:\n \"\"\"For an integer field filter.\n Returns a bare clause or \"\".\"\"\"\n if vals is None or not vals:\n return \"\"\n eq_elems = [f\"{key} = {val}\" for val in vals]\n return f\"({' or '.join(eq_elems)})\"\n\n def _build_kg_filter(\n kg_entities: list[str] | None,\n kg_relationships: list[str] | None,\n kg_terms: list[str] | None,\n ) -> str:\n if not kg_entities and not kg_relationships and not kg_terms:\n return \"\"\n\n combined_filter_parts = []\n\n def _build_kge(entity: str) -> str:\n GENERAL = \"::*\"\n if entity.endswith(GENERAL):\n return f'({{prefix: true}}\"{entity.split(GENERAL, 1)[0]}\")'\n else:\n return f'\"{entity}\"'\n\n if kg_entities:\n filter_parts = []\n for kg_entity in kg_entities:\n filter_parts.append(f\"(kg_entities contains {_build_kge(kg_entity)})\")\n combined_filter_parts.append(f\"({' or '.join(filter_parts)})\")\n\n # TODO: handle complex nested relationship logic (e.g., A participated, and B or C participated)\n if kg_relationships:\n filter_parts = []\n for kg_relationship in kg_relationships:\n source, rel_type, target = split_relationship_id(kg_relationship)\n filter_parts.append(\n \"(kg_relationships contains sameElement(\"\n f\"source contains {_build_kge(source)},\"\n f'rel_type contains \"{rel_type}\",'\n f\"target contains {_build_kge(target)}))\"\n )\n combined_filter_parts.append(f\"{' and '.join(filter_parts)}\")\n\n # TODO: remove kg terms entirely from prompts and codebase\n\n return f\"({' and '.join(combined_filter_parts)})\"\n\n def _build_kg_source_filters(\n kg_sources: list[str] | None,\n ) -> str:\n if not kg_sources:\n return \"\"\n\n source_phrases = [f'{DOCUMENT_ID} contains \"{source}\"' for source in kg_sources]\n return f\"({' or '.join(source_phrases)})\"\n\n def _build_kg_chunk_id_zero_only_filter(\n kg_chunk_id_zero_only: bool,\n ) -> str:\n if not kg_chunk_id_zero_only:\n return \"\"\n return \"(chunk_id = 0)\"\n\n def _build_time_filter(\n cutoff: datetime | None,\n untimed_doc_cutoff: timedelta = timedelta(days=92),\n ) -> str:\n if not cutoff:\n return \"\"\n include_untimed = datetime.now(timezone.utc) - untimed_doc_cutoff > cutoff\n cutoff_secs = int(cutoff.timestamp())\n\n if include_untimed:\n return f\"!({DOC_UPDATED_AT} < {cutoff_secs})\"\n return f\"({DOC_UPDATED_AT} >= {cutoff_secs})\"\n\n def _build_user_project_filter(\n project_id: int | None,\n ) -> str:\n if project_id is None:\n return \"\"\n try:\n pid = int(project_id)\n except Exception:\n return \"\"\n return f'({USER_PROJECT} contains \"{pid}\")'\n\n def _build_persona_filter(\n persona_id: int | None,\n ) -> str:\n if persona_id is None:\n return \"\"\n try:\n pid = int(persona_id)\n except Exception:\n logger.warning(f\"Invalid persona ID: {persona_id}\")\n return \"\"\n return f'({PERSONAS} contains \"{pid}\")'\n\n def _append(parts: list[str], clause: str) -> None:\n if clause:\n parts.append(clause)\n\n # Collect all top-level filter clauses, then join with \" and \" at the end.\n filter_parts: list[str] = []\n\n if not include_hidden:\n filter_parts.append(f\"!({HIDDEN}=true)\")\n\n # TODO: add error condition if MULTI_TENANT and no tenant_id filter is set\n if filters.tenant_id and MULTI_TENANT:\n filter_parts.append(build_tenant_id_filter(filters.tenant_id))\n\n # ACL filters — use weightedSet for efficient matching against the\n # access_control_list weightedset field. OR-chaining thousands\n # of 'contains' clauses causes Vespa to reject the query (HTTP 400)\n # for users with large numbers of external permission groups.\n if filters.access_control_list is not None:\n _append(\n filter_parts,\n _build_weighted_set_filter(\n ACCESS_CONTROL_LIST, filters.access_control_list\n ),\n )\n\n # Source type filters\n source_strs = (\n [s.value for s in filters.source_type] if filters.source_type else None\n )\n _append(filter_parts, _build_or_filters(SOURCE_TYPE, source_strs))\n\n # Tag filters\n tag_attributes = None\n if filters.tags:\n tag_attributes = [\n f\"{tag.tag_key}{INDEX_SEPARATOR}{tag.tag_value}\" for tag in filters.tags\n ]\n _append(filter_parts, _build_or_filters(METADATA_LIST, tag_attributes))\n\n # Knowledge scope: explicit knowledge attachments restrict what an\n # assistant can see. When none are set, the assistant can see\n # everything.\n #\n # persona_id_filter is a primary trigger — a persona with user files IS\n # explicit knowledge, so it can start a knowledge scope on its own.\n #\n # project_id_filter is additive — it widens the scope to also cover\n # overflowing project files but never restricts on its own (a chat\n # inside a project should still search team knowledge).\n knowledge_scope_parts: list[str] = []\n\n _append(\n knowledge_scope_parts, _build_or_filters(DOCUMENT_SETS, filters.document_set)\n )\n _append(knowledge_scope_parts, _build_persona_filter(filters.persona_id_filter))\n\n # project_id_filter only widens an existing scope.\n if knowledge_scope_parts:\n _append(\n knowledge_scope_parts,\n _build_user_project_filter(filters.project_id_filter),\n )\n\n if len(knowledge_scope_parts) > 1:\n filter_parts.append(\"(\" + \" or \".join(knowledge_scope_parts) + \")\")\n elif len(knowledge_scope_parts) == 1:\n filter_parts.append(knowledge_scope_parts[0])\n\n # Time filter\n _append(filter_parts, _build_time_filter(filters.time_cutoff))\n\n # # Knowledge Graph Filters\n # _append(filter_parts, _build_kg_filter(\n # kg_entities=filters.kg_entities,\n # kg_relationships=filters.kg_relationships,\n # kg_terms=filters.kg_terms,\n # ))\n\n # _append(filter_parts, _build_kg_source_filters(filters.kg_sources))\n\n # _append(filter_parts, _build_kg_chunk_id_zero_only_filter(\n # filters.kg_chunk_id_zero_only or False\n # ))\n\n filter_str = \" and \".join(filter_parts)\n\n if filter_str and not remove_trailing_and:\n filter_str += \" and \"\n\n return filter_str\n\n\ndef build_vespa_id_based_retrieval_yql(\n chunk_request: VespaChunkRequest,\n) -> str:\n id_based_retrieval_yql_section = (\n f'({DOCUMENT_ID} contains \"{chunk_request.document_id}\"'\n )\n\n if chunk_request.is_capped:\n id_based_retrieval_yql_section += (\n f\" and {CHUNK_ID} >= {chunk_request.min_chunk_ind or 0}\"\n )\n id_based_retrieval_yql_section += (\n f\" and {CHUNK_ID} <= {chunk_request.max_chunk_ind}\"\n )\n\n id_based_retrieval_yql_section += \")\"\n return id_based_retrieval_yql_section\n" }, { "path": "backend/onyx/document_index/vespa/vespa_document_index.py", "content": "import concurrent.futures\nimport logging\nimport random\nfrom collections.abc import Generator\nfrom collections.abc import Iterable\nfrom typing import Any\nfrom uuid import UUID\n\nimport httpx\nfrom pydantic import BaseModel\nfrom retry import retry\n\nfrom onyx.configs.app_configs import MAX_CHUNKS_PER_DOC_BATCH\nfrom onyx.configs.app_configs import RECENCY_BIAS_MULTIPLIER\nfrom onyx.configs.app_configs import RERANK_COUNT\nfrom onyx.configs.chat_configs import DOC_TIME_DECAY\nfrom onyx.configs.chat_configs import HYBRID_ALPHA\nfrom onyx.configs.chat_configs import TITLE_CONTENT_RATIO\nfrom onyx.context.search.enums import QueryType\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.document_index.chunk_content_enrichment import cleanup_content_for_chunks\nfrom onyx.document_index.document_index_utils import get_document_chunk_ids\nfrom onyx.document_index.document_index_utils import get_uuid_from_chunk_info\nfrom onyx.document_index.interfaces import EnrichedDocumentIndexingInfo\nfrom onyx.document_index.interfaces import MinimalDocumentIndexingInfo\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces_new import DocumentIndex\nfrom onyx.document_index.interfaces_new import DocumentInsertionRecord\nfrom onyx.document_index.interfaces_new import DocumentSectionRequest\nfrom onyx.document_index.interfaces_new import IndexingMetadata\nfrom onyx.document_index.interfaces_new import MetadataUpdateRequest\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.vespa.chunk_retrieval import batch_search_api_retrieval\nfrom onyx.document_index.vespa.chunk_retrieval import get_all_chunks_paginated\nfrom onyx.document_index.vespa.chunk_retrieval import get_chunks_via_visit_api\nfrom onyx.document_index.vespa.chunk_retrieval import (\n parallel_visit_api_retrieval,\n)\nfrom onyx.document_index.vespa.chunk_retrieval import query_vespa\nfrom onyx.document_index.vespa.deletion import delete_vespa_chunks\nfrom onyx.document_index.vespa.indexing_utils import BaseHTTPXClientContext\nfrom onyx.document_index.vespa.indexing_utils import batch_index_vespa_chunks\nfrom onyx.document_index.vespa.indexing_utils import check_for_final_chunk_existence\nfrom onyx.document_index.vespa.indexing_utils import clean_chunk_id_copy\nfrom onyx.document_index.vespa.indexing_utils import GlobalHTTPXClientContext\nfrom onyx.document_index.vespa.indexing_utils import TemporaryHTTPXClientContext\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\nfrom onyx.document_index.vespa.shared_utils.utils import (\n replace_invalid_doc_id_characters,\n)\nfrom onyx.document_index.vespa.shared_utils.vespa_request_builders import (\n build_vespa_filters,\n)\nfrom onyx.document_index.vespa_constants import BATCH_SIZE\nfrom onyx.document_index.vespa_constants import CHUNK_ID\nfrom onyx.document_index.vespa_constants import CONTENT_SUMMARY\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.document_index.vespa_constants import NUM_THREADS\nfrom onyx.document_index.vespa_constants import SEARCH_ENDPOINT\nfrom onyx.document_index.vespa_constants import VESPA_TIMEOUT\nfrom onyx.document_index.vespa_constants import YQL_BASE\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.tools.tool_implementations.search.constants import KEYWORD_QUERY_HYBRID_ALPHA\nfrom onyx.utils.batching import batch_generator\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.model_server_models import Embedding\n\n\nlogger = setup_logger(__name__)\n# Set the logging level to WARNING to ignore INFO and DEBUG logs from httpx. By\n# default it emits INFO-level logs for every request.\nhttpx_logger = logging.getLogger(\"httpx\")\nhttpx_logger.setLevel(logging.WARNING)\n\n\ndef _enrich_basic_chunk_info(\n index_name: str,\n http_client: httpx.Client,\n document_id: str,\n previous_chunk_count: int | None,\n new_chunk_count: int,\n) -> EnrichedDocumentIndexingInfo:\n \"\"\"Determines which chunks need to be deleted during document reindexing.\n\n When a document is reindexed, it may have fewer chunks than before. This\n function identifies the range of old chunks that need to be deleted by\n comparing the new chunk count with the previous chunk count.\n\n Example:\n If a document previously had 10 chunks (0-9) and now has 7 chunks (0-6),\n this function identifies that chunks 7-9 need to be deleted.\n\n Args:\n index_name: The Vespa index/schema name.\n http_client: HTTP client for making requests to Vespa.\n document_id: The Vespa-sanitized ID of the document being reindexed.\n previous_chunk_count: The total number of chunks the document had before\n reindexing. None for documents using the legacy chunk ID system.\n new_chunk_count: The total number of chunks the document has after\n reindexing. This becomes the starting index for deletion since\n chunks are 0-indexed.\n\n Returns:\n EnrichedDocumentIndexingInfo with chunk_start_index set to\n new_chunk_count (where deletion begins) and chunk_end_index set to\n previous_chunk_count (where deletion ends).\n \"\"\"\n # Technically last indexed chunk index +1.\n last_indexed_chunk = previous_chunk_count\n # If the document has no `chunk_count` in the database, we know that it\n # has the old chunk ID system and we must check for the final chunk index.\n is_old_version = False\n if last_indexed_chunk is None:\n is_old_version = True\n minimal_doc_info = MinimalDocumentIndexingInfo(\n doc_id=document_id, chunk_start_index=new_chunk_count\n )\n last_indexed_chunk = check_for_final_chunk_existence(\n minimal_doc_info=minimal_doc_info,\n start_index=new_chunk_count,\n index_name=index_name,\n http_client=http_client,\n )\n\n assert (\n last_indexed_chunk is not None and last_indexed_chunk >= 0\n ), f\"Bug: Last indexed chunk index is None or less than 0 for document: {document_id}.\"\n\n enriched_doc_info = EnrichedDocumentIndexingInfo(\n doc_id=document_id,\n chunk_start_index=new_chunk_count,\n chunk_end_index=last_indexed_chunk,\n old_version=is_old_version,\n )\n return enriched_doc_info\n\n\n@retry(\n tries=3,\n delay=1,\n backoff=2,\n exceptions=httpx.HTTPError,\n)\ndef _update_single_chunk(\n doc_chunk_id: UUID,\n index_name: str,\n doc_id: str,\n http_client: httpx.Client,\n update_request: MetadataUpdateRequest,\n) -> None:\n \"\"\"Updates a single document chunk in Vespa.\n\n TODO(andrei): Couldn't this be batched?\n\n Args:\n doc_chunk_id: The ID of the chunk to update.\n index_name: The index the chunk belongs to.\n doc_id: The ID of the document the chunk belongs to. Used only for\n logging.\n http_client: The HTTP client to use to make the request.\n update_request: Metadata update request object received in the bulk\n update method containing fields to update.\n \"\"\"\n\n class _Boost(BaseModel):\n model_config = {\"frozen\": True}\n assign: float\n\n class _DocumentSets(BaseModel):\n model_config = {\"frozen\": True}\n assign: dict[str, int]\n\n class _AccessControl(BaseModel):\n model_config = {\"frozen\": True}\n assign: dict[str, int]\n\n class _Hidden(BaseModel):\n model_config = {\"frozen\": True}\n assign: bool\n\n class _UserProjects(BaseModel):\n model_config = {\"frozen\": True}\n assign: list[int]\n\n class _Personas(BaseModel):\n model_config = {\"frozen\": True}\n assign: list[int]\n\n class _VespaPutFields(BaseModel):\n model_config = {\"frozen\": True}\n # The names of these fields are based the Vespa schema. Changes to the\n # schema require changes here. These names were originally found in\n # backend/onyx/document_index/vespa_constants.py.\n boost: _Boost | None = None\n document_sets: _DocumentSets | None = None\n access_control_list: _AccessControl | None = None\n hidden: _Hidden | None = None\n user_project: _UserProjects | None = None\n personas: _Personas | None = None\n\n class _VespaPutRequest(BaseModel):\n model_config = {\"frozen\": True}\n fields: _VespaPutFields\n\n boost_update: _Boost | None = (\n _Boost(assign=update_request.boost)\n if update_request.boost is not None\n else None\n )\n document_sets_update: _DocumentSets | None = (\n _DocumentSets(\n assign={document_set: 1 for document_set in update_request.document_sets}\n )\n if update_request.document_sets is not None\n else None\n )\n access_update: _AccessControl | None = (\n _AccessControl(\n assign={acl_entry: 1 for acl_entry in update_request.access.to_acl()}\n )\n if update_request.access is not None\n else None\n )\n hidden_update: _Hidden | None = (\n _Hidden(assign=update_request.hidden)\n if update_request.hidden is not None\n else None\n )\n user_projects_update: _UserProjects | None = (\n _UserProjects(assign=list(update_request.project_ids))\n if update_request.project_ids is not None\n else None\n )\n personas_update: _Personas | None = (\n _Personas(assign=list(update_request.persona_ids))\n if update_request.persona_ids is not None\n else None\n )\n\n vespa_put_fields = _VespaPutFields(\n boost=boost_update,\n document_sets=document_sets_update,\n access_control_list=access_update,\n hidden=hidden_update,\n user_project=user_projects_update,\n personas=personas_update,\n )\n\n vespa_put_request = _VespaPutRequest(\n fields=vespa_put_fields,\n )\n\n vespa_url = f\"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{doc_chunk_id}?create=true\"\n\n try:\n resp = http_client.put(\n vespa_url,\n headers={\"Content-Type\": \"application/json\"},\n json=vespa_put_request.model_dump(\n exclude_none=True\n ), # NOTE: Important to not produce null fields in the json.\n )\n resp.raise_for_status()\n except httpx.HTTPStatusError as e:\n logger.error(\n f\"Failed to update doc chunk {doc_chunk_id} (doc_id={doc_id}). \"\n f\"Code: {e.response.status_code}. Details: {e.response.text}\"\n )\n # Re-raise so the @retry decorator will catch and retry, unless the\n # status code is < 5xx, in which case wrap the exception in something\n # other than an HTTPError to skip retries.\n if e.response.status_code >= 500:\n raise\n raise RuntimeError(\n f\"Non-retryable error updating chunk {doc_chunk_id}: {e}\"\n ) from e\n\n\nclass VespaDocumentIndex(DocumentIndex):\n \"\"\"Vespa-specific implementation of the DocumentIndex interface.\n\n This class provides document indexing, retrieval, and management operations\n for a Vespa search engine instance. It handles the complete lifecycle of\n document chunks within a specific Vespa index/schema.\n \"\"\"\n\n def __init__(\n self,\n index_name: str,\n tenant_state: TenantState,\n large_chunks_enabled: bool,\n httpx_client: httpx.Client | None = None,\n ) -> None:\n self._index_name = index_name\n self._tenant_id = tenant_state.tenant_id\n self._large_chunks_enabled = large_chunks_enabled\n # NOTE: using `httpx` here since `requests` doesn't support HTTP2. This\n # is beneficial for indexing / updates / deletes since we have to make a\n # large volume of requests.\n self._httpx_client_context: BaseHTTPXClientContext\n if httpx_client:\n # Use the provided client. Because this client is presumed global,\n # it does not close after exiting a context manager.\n self._httpx_client_context = GlobalHTTPXClientContext(httpx_client)\n else:\n # We did not receive a client, so create one what will close after\n # exiting a context manager.\n self._httpx_client_context = TemporaryHTTPXClientContext(\n get_vespa_http_client\n )\n self._multitenant = tenant_state.multitenant\n\n def verify_and_create_index_if_necessary(\n self, embedding_dim: int, embedding_precision: EmbeddingPrecision\n ) -> None:\n raise NotImplementedError\n\n def index(\n self,\n chunks: Iterable[DocMetadataAwareIndexChunk],\n indexing_metadata: IndexingMetadata,\n ) -> list[DocumentInsertionRecord]:\n doc_id_to_chunk_cnt_diff = indexing_metadata.doc_id_to_chunk_cnt_diff\n doc_id_to_previous_chunk_cnt = {\n doc_id: chunk_cnt_diff.old_chunk_cnt\n for doc_id, chunk_cnt_diff in doc_id_to_chunk_cnt_diff.items()\n }\n doc_id_to_new_chunk_cnt = {\n doc_id: chunk_cnt_diff.new_chunk_cnt\n for doc_id, chunk_cnt_diff in doc_id_to_chunk_cnt_diff.items()\n }\n assert (\n len(doc_id_to_chunk_cnt_diff)\n == len(doc_id_to_previous_chunk_cnt)\n == len(doc_id_to_new_chunk_cnt)\n ), \"Bug: Doc ID to chunk maps have different lengths.\"\n\n # Vespa has restrictions on valid characters, yet document IDs come from\n # external w.r.t. this class. We need to sanitize them.\n #\n # Instead of materializing all cleaned chunks upfront, we stream them\n # through a generator that cleans IDs and builds the original-ID mapping\n # incrementally as chunks flow into Vespa.\n def _clean_and_track(\n chunks_iter: Iterable[DocMetadataAwareIndexChunk],\n id_map: dict[str, str],\n seen_ids: set[str],\n ) -> Generator[DocMetadataAwareIndexChunk, None, None]:\n \"\"\"Cleans chunk IDs and builds the original-ID mapping\n incrementally as chunks flow through, avoiding a separate\n materialization pass.\"\"\"\n for chunk in chunks_iter:\n original_id = chunk.source_document.id\n cleaned = clean_chunk_id_copy(chunk)\n cleaned_id = cleaned.source_document.id\n # Needed so the final DocumentInsertionRecord returned can have\n # the original document ID. cleaned_chunks might not contain IDs\n # exactly as callers supplied them.\n id_map[cleaned_id] = original_id\n seen_ids.add(cleaned_id)\n yield cleaned\n\n new_document_id_to_original_document_id: dict[str, str] = {}\n all_cleaned_doc_ids: set[str] = set()\n\n existing_docs: set[str] = set()\n\n with (\n concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor,\n self._httpx_client_context as http_client,\n ):\n # We require the start and end index for each document in order to\n # know precisely which chunks to delete. This information exists for\n # documents that have `chunk_count` in the database, but not for\n # `old_version` documents.\n enriched_doc_infos: list[EnrichedDocumentIndexingInfo] = [\n _enrich_basic_chunk_info(\n index_name=self._index_name,\n http_client=http_client,\n document_id=doc_id,\n previous_chunk_count=doc_id_to_previous_chunk_cnt[doc_id],\n new_chunk_count=doc_id_to_new_chunk_cnt[doc_id],\n )\n for doc_id in doc_id_to_chunk_cnt_diff.keys()\n # TODO(andrei), WARNING: Don't we need to sanitize these doc IDs?\n ]\n\n for enriched_doc_info in enriched_doc_infos:\n # If the document has previously indexed chunks, we know it\n # previously existed and this is a reindex.\n if enriched_doc_info.chunk_end_index:\n existing_docs.add(enriched_doc_info.doc_id)\n\n # Now, for each doc, we know exactly where to start and end our\n # deletion. So let's generate the chunk IDs for each chunk to\n # delete.\n # WARNING: This code seems to use\n # indexing_metadata.doc_id_to_chunk_cnt_diff as the source of truth\n # for which chunks to delete. This implies that the onus is on the\n # caller to ensure doc_id_to_chunk_cnt_diff only contains docs\n # relevant to the chunks argument to this method. This should not be\n # the contract of DocumentIndex; and this code is only a refactor\n # from old code. It would seem we should use all_cleaned_doc_ids as\n # the source of truth.\n chunks_to_delete = get_document_chunk_ids(\n enriched_document_info_list=enriched_doc_infos,\n tenant_id=self._tenant_id,\n large_chunks_enabled=self._large_chunks_enabled,\n )\n\n # Delete old Vespa documents.\n for doc_chunk_ids_batch in batch_generator(chunks_to_delete, BATCH_SIZE):\n delete_vespa_chunks(\n doc_chunk_ids=doc_chunk_ids_batch,\n index_name=self._index_name,\n http_client=http_client,\n executor=executor,\n )\n\n # Insert new Vespa documents, streaming through the cleaning\n # pipeline so chunks are never fully materialized.\n cleaned_chunks = _clean_and_track(\n chunks,\n new_document_id_to_original_document_id,\n all_cleaned_doc_ids,\n )\n for chunk_batch in batch_generator(\n cleaned_chunks, min(BATCH_SIZE, MAX_CHUNKS_PER_DOC_BATCH)\n ):\n batch_index_vespa_chunks(\n chunks=chunk_batch,\n index_name=self._index_name,\n http_client=http_client,\n multitenant=self._multitenant,\n executor=executor,\n )\n\n return [\n DocumentInsertionRecord(\n document_id=new_document_id_to_original_document_id[cleaned_doc_id],\n already_existed=cleaned_doc_id in existing_docs,\n )\n for cleaned_doc_id in all_cleaned_doc_ids\n ]\n\n def delete(self, document_id: str, chunk_count: int | None = None) -> int:\n total_chunks_deleted = 0\n\n sanitized_doc_id = replace_invalid_doc_id_characters(document_id)\n\n with (\n concurrent.futures.ThreadPoolExecutor(max_workers=NUM_THREADS) as executor,\n self._httpx_client_context as http_client,\n ):\n enriched_doc_info = _enrich_basic_chunk_info(\n index_name=self._index_name,\n http_client=http_client,\n document_id=sanitized_doc_id,\n previous_chunk_count=chunk_count,\n new_chunk_count=0,\n )\n chunks_to_delete = get_document_chunk_ids(\n enriched_document_info_list=[enriched_doc_info],\n tenant_id=self._tenant_id,\n large_chunks_enabled=self._large_chunks_enabled,\n )\n\n for doc_chunk_ids_batch in batch_generator(chunks_to_delete, BATCH_SIZE):\n total_chunks_deleted += len(doc_chunk_ids_batch)\n delete_vespa_chunks(\n doc_chunk_ids=doc_chunk_ids_batch,\n index_name=self._index_name,\n http_client=http_client,\n executor=executor,\n )\n\n return total_chunks_deleted\n\n def update(\n self,\n update_requests: list[MetadataUpdateRequest],\n ) -> None:\n # WARNING: This method can be called by vespa_metadata_sync_task, which\n # is kicked off by check_for_vespa_sync_task, notably before a document\n # has finished indexing. In this way, chunk_count below could be unknown\n # even for chunks not on the \"old\" chunk ID system; i.e. there could be\n # a race condition. Passing in None to _enrich_basic_chunk_info should\n # handle this, but a higher level TODO might be to not run update at all\n # on connectors that are still indexing, and therefore do not yet have a\n # chunk count because update_docs_chunk_count__no_commit has not been\n # run yet.\n with self._httpx_client_context as httpx_client:\n # Each invocation of this method can contain multiple update requests.\n for update_request in update_requests:\n # Each update request can correspond to multiple documents.\n for doc_id in update_request.document_ids:\n # NOTE: -1 represents an unknown chunk count.\n chunk_count = update_request.doc_id_to_chunk_cnt[doc_id]\n sanitized_doc_id = replace_invalid_doc_id_characters(doc_id)\n enriched_doc_info = _enrich_basic_chunk_info(\n index_name=self._index_name,\n http_client=httpx_client,\n document_id=sanitized_doc_id,\n previous_chunk_count=chunk_count if chunk_count >= 0 else None,\n new_chunk_count=0, # WARNING: This semantically makes no sense and is misusing this function.\n )\n\n doc_chunk_ids = get_document_chunk_ids(\n enriched_document_info_list=[enriched_doc_info],\n tenant_id=self._tenant_id,\n large_chunks_enabled=self._large_chunks_enabled,\n )\n\n for doc_chunk_id in doc_chunk_ids:\n _update_single_chunk(\n doc_chunk_id,\n self._index_name,\n # NOTE: Used only for logging, raw ID is ok here.\n doc_id,\n httpx_client,\n update_request,\n )\n\n logger.info(\n f\"Updated {len(doc_chunk_ids)} chunks for document {doc_id}.\"\n )\n\n def id_based_retrieval(\n self,\n chunk_requests: list[DocumentSectionRequest],\n filters: IndexFilters,\n batch_retrieval: bool = False,\n ) -> list[InferenceChunk]:\n sanitized_chunk_requests = [\n VespaChunkRequest(\n document_id=replace_invalid_doc_id_characters(\n chunk_request.document_id\n ),\n min_chunk_ind=chunk_request.min_chunk_ind,\n max_chunk_ind=chunk_request.max_chunk_ind,\n )\n for chunk_request in chunk_requests\n ]\n\n if batch_retrieval:\n return cleanup_content_for_chunks(\n batch_search_api_retrieval(\n index_name=self._index_name,\n chunk_requests=sanitized_chunk_requests,\n filters=filters,\n # No one was passing in this parameter in the legacy\n # interface, it always defaulted to False.\n get_large_chunks=False,\n )\n )\n return cleanup_content_for_chunks(\n parallel_visit_api_retrieval(\n index_name=self._index_name,\n chunk_requests=sanitized_chunk_requests,\n filters=filters,\n # No one was passing in this parameter in the legacy interface,\n # it always defaulted to False.\n get_large_chunks=False,\n )\n )\n\n def hybrid_retrieval(\n self,\n query: str,\n query_embedding: Embedding,\n final_keywords: list[str] | None,\n query_type: QueryType,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n vespa_where_clauses = build_vespa_filters(filters)\n # Avoid over-fetching a very large candidate set for global-phase reranking.\n # Keep enough headroom for quality while capping cost on larger indices.\n target_hits = min(max(4 * num_to_retrieve, 100), RERANK_COUNT)\n\n yql = (\n YQL_BASE.format(index_name=self._index_name)\n + vespa_where_clauses\n + f\"(({{targetHits: {target_hits}}}nearestNeighbor(embeddings, query_embedding)) \"\n + f\"or ({{targetHits: {target_hits}}}nearestNeighbor(title_embedding, query_embedding)) \"\n + 'or ({grammar: \"weakAnd\"}userInput(@query)) '\n + f'or ({{defaultIndex: \"{CONTENT_SUMMARY}\"}}userInput(@query)))'\n )\n\n final_query = \" \".join(final_keywords) if final_keywords else query\n\n ranking_profile = (\n f\"hybrid_search_{query_type.value}_base_{len(query_embedding)}\"\n )\n\n logger.info(f\"Selected ranking profile: {ranking_profile}\")\n\n logger.debug(f\"Query YQL: {yql}\")\n\n # In this interface we do not pass in hybrid alpha. Tracing the codepath\n # of the legacy Vespa interface, it so happens that KEYWORD always\n # corresponds to an alpha of 0.2 (from KEYWORD_QUERY_HYBRID_ALPHA), and\n # SEMANTIC to 0.5 (from HYBRID_ALPHA). HYBRID_ALPHA_KEYWORD was only\n # used in dead code so we do not use it here.\n hybrid_alpha = (\n KEYWORD_QUERY_HYBRID_ALPHA\n if query_type == QueryType.KEYWORD\n else HYBRID_ALPHA\n )\n\n params: dict[str, str | int | float] = {\n \"yql\": yql,\n \"query\": final_query,\n \"input.query(query_embedding)\": str(query_embedding),\n \"input.query(decay_factor)\": str(DOC_TIME_DECAY * RECENCY_BIAS_MULTIPLIER),\n \"input.query(alpha)\": hybrid_alpha,\n \"input.query(title_content_ratio)\": TITLE_CONTENT_RATIO,\n \"hits\": num_to_retrieve,\n \"ranking.profile\": ranking_profile,\n \"timeout\": VESPA_TIMEOUT,\n }\n\n return cleanup_content_for_chunks(query_vespa(params))\n\n def keyword_retrieval(\n self,\n query: str,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n raise NotImplementedError\n\n def semantic_retrieval(\n self,\n query_embedding: Embedding,\n filters: IndexFilters,\n num_to_retrieve: int,\n ) -> list[InferenceChunk]:\n raise NotImplementedError\n\n def random_retrieval(\n self,\n filters: IndexFilters,\n num_to_retrieve: int = 100,\n dirty: bool | None = None, # noqa: ARG002\n ) -> list[InferenceChunk]:\n vespa_where_clauses = build_vespa_filters(filters, remove_trailing_and=True)\n\n yql = YQL_BASE.format(index_name=self._index_name) + vespa_where_clauses\n\n random_seed = random.randint(0, 1_000_000)\n\n params: dict[str, str | int | float] = {\n \"yql\": yql,\n \"hits\": num_to_retrieve,\n \"timeout\": VESPA_TIMEOUT,\n \"ranking.profile\": \"random_\",\n \"ranking.properties.random.seed\": random_seed,\n }\n\n return cleanup_content_for_chunks(query_vespa(params))\n\n def get_raw_document_chunks(self, document_id: str) -> list[dict[str, Any]]:\n \"\"\"Gets all raw document chunks for a document as returned by Vespa.\n\n Used in the Vespa migration task.\n\n Args:\n document_id: The ID of the document to get chunks for.\n\n Returns:\n List of raw document chunks.\n \"\"\"\n # Vespa doc IDs are sanitized using replace_invalid_doc_id_characters.\n sanitized_document_id = replace_invalid_doc_id_characters(document_id)\n chunk_request = VespaChunkRequest(document_id=sanitized_document_id)\n raw_chunks = get_chunks_via_visit_api(\n chunk_request=chunk_request,\n index_name=self._index_name,\n filters=IndexFilters(access_control_list=None, tenant_id=self._tenant_id),\n get_large_chunks=False,\n short_tensor_format=True,\n )\n # Vespa returns other metadata around the actual document chunk. The raw\n # chunk we're interested in is in the \"fields\" field.\n raw_document_chunks = [chunk[\"fields\"] for chunk in raw_chunks]\n return raw_document_chunks\n\n def get_all_raw_document_chunks_paginated(\n self,\n continuation_token_map: dict[int, str | None],\n page_size: int,\n ) -> tuple[list[dict[str, Any]], dict[int, str | None]]:\n \"\"\"Gets all the chunks in Vespa, paginated.\n\n Used in the chunk-level Vespa-to-OpenSearch migration task.\n\n Args:\n continuation_token: Token returned by Vespa representing a page\n offset. None to start from the beginning. Defaults to None.\n page_size: Best-effort batch size for the visit.\n\n Returns:\n Tuple of (list of chunk dicts, next continuation token or None). The\n continuation token is None when the visit is complete.\n \"\"\"\n raw_chunks, next_continuation_token_map = get_all_chunks_paginated(\n index_name=self._index_name,\n tenant_state=TenantState(\n tenant_id=self._tenant_id, multitenant=MULTI_TENANT\n ),\n continuation_token_map=continuation_token_map,\n page_size=page_size,\n )\n return raw_chunks, next_continuation_token_map\n\n def index_raw_chunks(self, chunks: list[dict[str, Any]]) -> None:\n \"\"\"Indexes raw document chunks into Vespa.\n\n To only be used in tests. Not for production.\n \"\"\"\n json_header = {\n \"Content-Type\": \"application/json\",\n }\n with self._httpx_client_context as http_client:\n for chunk in chunks:\n chunk_id = str(\n get_uuid_from_chunk_info(\n document_id=chunk[DOCUMENT_ID],\n chunk_id=chunk[CHUNK_ID],\n tenant_id=self._tenant_id,\n )\n )\n vespa_url = f\"{DOCUMENT_ID_ENDPOINT.format(index_name=self._index_name)}/{chunk_id}\"\n response = http_client.post(\n vespa_url,\n headers=json_header,\n json={\"fields\": chunk},\n )\n response.raise_for_status()\n\n def get_chunk_count(self) -> int:\n \"\"\"Returns the exact number of document chunks in Vespa for this tenant.\n\n Uses the Vespa Search API with `limit 0` and `ranking.profile=unranked`\n to get an exact count without fetching any document data.\n\n Includes large chunks. There is no way to filter these out using the\n Search API.\n \"\"\"\n where_clause = (\n f'tenant_id contains \"{self._tenant_id}\"' if self._multitenant else \"true\"\n )\n yql = f\"select documentid from {self._index_name} where {where_clause} limit 0\"\n params: dict[str, str | int] = {\n \"yql\": yql,\n \"ranking.profile\": \"unranked\",\n \"timeout\": VESPA_TIMEOUT,\n }\n\n with get_vespa_http_client() as http_client:\n response = http_client.post(SEARCH_ENDPOINT, json=params)\n response.raise_for_status()\n response_data = response.json()\n return response_data[\"root\"][\"fields\"][\"totalCount\"]\n" }, { "path": "backend/onyx/document_index/vespa_constants.py", "content": "from onyx.configs.app_configs import VESPA_CLOUD_URL\nfrom onyx.configs.app_configs import VESPA_CONFIG_SERVER_HOST\nfrom onyx.configs.app_configs import VESPA_HOST\nfrom onyx.configs.app_configs import VESPA_PORT\nfrom onyx.configs.app_configs import VESPA_TENANT_PORT\nfrom onyx.configs.constants import SOURCE_TYPE\n\n# config server\n\n\nVESPA_CONFIG_SERVER_URL = (\n VESPA_CLOUD_URL or f\"http://{VESPA_CONFIG_SERVER_HOST}:{VESPA_TENANT_PORT}\"\n)\nVESPA_APPLICATION_ENDPOINT = f\"{VESPA_CONFIG_SERVER_URL}/application/v2\"\n\n# main search application\nVESPA_APP_CONTAINER_URL = VESPA_CLOUD_URL or f\"http://{VESPA_HOST}:{VESPA_PORT}\"\n\n\n# danswer_chunk below is defined in vespa/app_configs/schemas/danswer_chunk.sd.jinja\nDOCUMENT_ID_ENDPOINT = (\n f\"{VESPA_APP_CONTAINER_URL}/document/v1/default/{{index_name}}/docid\"\n)\n\n# the default document id endpoint is http://localhost:8080/document/v1/default/danswer_chunk/docid\n\nSEARCH_ENDPOINT = f\"{VESPA_APP_CONTAINER_URL}/search/\"\n\n# Since Vespa doesn't allow batching of inserts / updates, we use threads to\n# parallelize the operations.\nNUM_THREADS = 32\nMAX_ID_SEARCH_QUERY_SIZE = 400\n# Suspect that adding too many \"or\" conditions will cause Vespa to timeout and return\n# an empty list of hits (with no error status and coverage: 0 and degraded)\nMAX_OR_CONDITIONS = 10\n# up from 500ms for now, since we've seen quite a few timeouts\n# in the long term, we are looking to improve the performance of Vespa\n# so that we can bring this back to default\nVESPA_TIMEOUT = \"10s\"\n# The size of the batch to use for batched operations like inserts / updates.\n# The batch will likely be sent to a threadpool of size NUM_THREADS.\nBATCH_SIZE = 128\n\nTENANT_ID = \"tenant_id\"\nDOCUMENT_ID = \"document_id\"\nCHUNK_ID = \"chunk_id\"\nBLURB = \"blurb\"\nCONTENT = \"content\"\nSOURCE_LINKS = \"source_links\"\nSEMANTIC_IDENTIFIER = \"semantic_identifier\"\nTITLE = \"title\"\nSKIP_TITLE_EMBEDDING = \"skip_title\"\nSECTION_CONTINUATION = \"section_continuation\"\nEMBEDDINGS = \"embeddings\"\nTITLE_EMBEDDING = \"title_embedding\"\nACCESS_CONTROL_LIST = \"access_control_list\"\nDOCUMENT_SETS = \"document_sets\"\nUSER_FILE = \"user_file\"\nUSER_FOLDER = \"user_folder\"\nUSER_PROJECT = \"user_project\"\nPERSONAS = \"personas\"\nLARGE_CHUNK_REFERENCE_IDS = \"large_chunk_reference_ids\"\nMETADATA = \"metadata\"\nMETADATA_LIST = \"metadata_list\"\nMETADATA_SUFFIX = \"metadata_suffix\"\nDOC_SUMMARY = \"doc_summary\"\nCHUNK_CONTEXT = \"chunk_context\"\nBOOST = \"boost\"\nAGGREGATED_CHUNK_BOOST_FACTOR = \"aggregated_chunk_boost_factor\"\nDOC_UPDATED_AT = \"doc_updated_at\" # Indexed as seconds since epoch\nPRIMARY_OWNERS = \"primary_owners\"\nSECONDARY_OWNERS = \"secondary_owners\"\nRECENCY_BIAS = \"recency_bias\"\nHIDDEN = \"hidden\"\n# for legacy reasons, called `name` in Vespa despite it really being an ID\nIMAGE_FILE_NAME = \"image_file_name\"\n\n# Specific to Vespa, needed for highlighting matching keywords / section\nCONTENT_SUMMARY = \"content_summary\"\n\nFULL_CHUNK_EMBEDDING_KEY = \"full_chunk\"\n\n\nYQL_BASE = (\n f\"select \"\n f\"documentid, \"\n f\"{DOCUMENT_ID}, \"\n f\"{CHUNK_ID}, \"\n f\"{BLURB}, \"\n f\"{CONTENT}, \"\n f\"{SOURCE_TYPE}, \"\n f\"{SOURCE_LINKS}, \"\n f\"{SEMANTIC_IDENTIFIER}, \"\n f\"{TITLE}, \"\n f\"{SECTION_CONTINUATION}, \"\n f\"{IMAGE_FILE_NAME}, \"\n f\"{BOOST}, \"\n f\"{AGGREGATED_CHUNK_BOOST_FACTOR}, \"\n f\"{HIDDEN}, \"\n f\"{DOC_UPDATED_AT}, \"\n f\"{PRIMARY_OWNERS}, \"\n f\"{SECONDARY_OWNERS}, \"\n f\"{LARGE_CHUNK_REFERENCE_IDS}, \"\n f\"{METADATA}, \"\n f\"{METADATA_SUFFIX}, \"\n f\"{DOC_SUMMARY}, \"\n f\"{CHUNK_CONTEXT}, \"\n f\"{CONTENT_SUMMARY} \"\n f\"from {{index_name}} where \"\n)\n" }, { "path": "backend/onyx/error_handling/__init__.py", "content": "" }, { "path": "backend/onyx/error_handling/error_codes.py", "content": "\"\"\"\nStandardized error codes for the Onyx backend.\n\nUsage:\n from onyx.error_handling.error_codes import OnyxErrorCode\n from onyx.error_handling.exceptions import OnyxError\n\n raise OnyxError(OnyxErrorCode.UNAUTHENTICATED, \"Token expired\")\n\"\"\"\n\nfrom enum import Enum\n\n\nclass OnyxErrorCode(Enum):\n \"\"\"\n Each member is a tuple of (error_code_string, http_status_code).\n\n The error_code_string is a stable, machine-readable identifier that\n API consumers can match on. The http_status_code is the default HTTP\n status to return.\n \"\"\"\n\n # ------------------------------------------------------------------\n # Authentication (401)\n # ------------------------------------------------------------------\n UNAUTHENTICATED = (\"UNAUTHENTICATED\", 401)\n INVALID_TOKEN = (\"INVALID_TOKEN\", 401)\n TOKEN_EXPIRED = (\"TOKEN_EXPIRED\", 401)\n CSRF_FAILURE = (\"CSRF_FAILURE\", 403)\n\n # ------------------------------------------------------------------\n # Authorization (403)\n # ------------------------------------------------------------------\n UNAUTHORIZED = (\"UNAUTHORIZED\", 403)\n INSUFFICIENT_PERMISSIONS = (\"INSUFFICIENT_PERMISSIONS\", 403)\n ADMIN_ONLY = (\"ADMIN_ONLY\", 403)\n EE_REQUIRED = (\"EE_REQUIRED\", 403)\n SINGLE_TENANT_ONLY = (\"SINGLE_TENANT_ONLY\", 403)\n ENV_VAR_GATED = (\"ENV_VAR_GATED\", 403)\n\n # ------------------------------------------------------------------\n # Validation / Bad Request (400)\n # ------------------------------------------------------------------\n VALIDATION_ERROR = (\"VALIDATION_ERROR\", 400)\n INVALID_INPUT = (\"INVALID_INPUT\", 400)\n MISSING_REQUIRED_FIELD = (\"MISSING_REQUIRED_FIELD\", 400)\n QUERY_REJECTED = (\"QUERY_REJECTED\", 400)\n\n # ------------------------------------------------------------------\n # Not Found (404)\n # ------------------------------------------------------------------\n NOT_FOUND = (\"NOT_FOUND\", 404)\n CONNECTOR_NOT_FOUND = (\"CONNECTOR_NOT_FOUND\", 404)\n CREDENTIAL_NOT_FOUND = (\"CREDENTIAL_NOT_FOUND\", 404)\n PERSONA_NOT_FOUND = (\"PERSONA_NOT_FOUND\", 404)\n DOCUMENT_NOT_FOUND = (\"DOCUMENT_NOT_FOUND\", 404)\n SESSION_NOT_FOUND = (\"SESSION_NOT_FOUND\", 404)\n USER_NOT_FOUND = (\"USER_NOT_FOUND\", 404)\n\n # ------------------------------------------------------------------\n # Conflict (409)\n # ------------------------------------------------------------------\n CONFLICT = (\"CONFLICT\", 409)\n DUPLICATE_RESOURCE = (\"DUPLICATE_RESOURCE\", 409)\n\n # ------------------------------------------------------------------\n # Rate Limiting / Quotas (429 / 402)\n # ------------------------------------------------------------------\n RATE_LIMITED = (\"RATE_LIMITED\", 429)\n SEAT_LIMIT_EXCEEDED = (\"SEAT_LIMIT_EXCEEDED\", 402)\n\n # ------------------------------------------------------------------\n # Payload (413)\n # ------------------------------------------------------------------\n PAYLOAD_TOO_LARGE = (\"PAYLOAD_TOO_LARGE\", 413)\n\n # ------------------------------------------------------------------\n # Connector / Credential Errors (400-range)\n # ------------------------------------------------------------------\n CONNECTOR_VALIDATION_FAILED = (\"CONNECTOR_VALIDATION_FAILED\", 400)\n CREDENTIAL_INVALID = (\"CREDENTIAL_INVALID\", 400)\n CREDENTIAL_EXPIRED = (\"CREDENTIAL_EXPIRED\", 401)\n\n # ------------------------------------------------------------------\n # Server Errors (5xx)\n # ------------------------------------------------------------------\n INTERNAL_ERROR = (\"INTERNAL_ERROR\", 500)\n NOT_IMPLEMENTED = (\"NOT_IMPLEMENTED\", 501)\n SERVICE_UNAVAILABLE = (\"SERVICE_UNAVAILABLE\", 503)\n BAD_GATEWAY = (\"BAD_GATEWAY\", 502)\n LLM_PROVIDER_ERROR = (\"LLM_PROVIDER_ERROR\", 502)\n HOOK_EXECUTION_FAILED = (\"HOOK_EXECUTION_FAILED\", 502)\n GATEWAY_TIMEOUT = (\"GATEWAY_TIMEOUT\", 504)\n\n def __init__(self, code: str, status_code: int) -> None:\n self.code = code\n self.status_code = status_code\n\n def detail(self, message: str | None = None) -> dict[str, str]:\n \"\"\"Build a structured error detail dict.\n\n Returns a dict like:\n {\"error_code\": \"UNAUTHENTICATED\", \"detail\": \"Token expired\"}\n\n If no message is supplied, the error code itself is used as the detail.\n \"\"\"\n return {\n \"error_code\": self.code,\n \"detail\": message or self.code,\n }\n" }, { "path": "backend/onyx/error_handling/exceptions.py", "content": "\"\"\"OnyxError — the single exception type for all Onyx business errors.\n\nRaise ``OnyxError`` instead of ``HTTPException`` in business code. A global\nFastAPI exception handler (registered via ``register_onyx_exception_handlers``)\nconverts it into a JSON response with the standard\n``{\"error_code\": \"...\", \"detail\": \"...\"}`` shape.\n\nUsage::\n\n from onyx.error_handling.error_codes import OnyxErrorCode\n from onyx.error_handling.exceptions import OnyxError\n\n raise OnyxError(OnyxErrorCode.NOT_FOUND, \"Session not found\")\n\nFor upstream errors with a dynamic HTTP status (e.g. billing service),\nuse ``status_code_override``::\n\n raise OnyxError(\n OnyxErrorCode.BAD_GATEWAY,\n detail,\n status_code_override=upstream_status,\n )\n\"\"\"\n\nfrom fastapi import FastAPI\nfrom fastapi import Request\nfrom fastapi.responses import JSONResponse\n\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass OnyxError(Exception):\n \"\"\"Structured error that maps to a specific ``OnyxErrorCode``.\n\n Attributes:\n error_code: The ``OnyxErrorCode`` enum member.\n detail: Human-readable detail (defaults to the error code string).\n status_code: HTTP status — either overridden or from the error code.\n \"\"\"\n\n def __init__(\n self,\n error_code: OnyxErrorCode,\n detail: str | None = None,\n *,\n status_code_override: int | None = None,\n ) -> None:\n resolved_detail = detail or error_code.code\n super().__init__(resolved_detail)\n self.error_code = error_code\n self.detail = resolved_detail\n self._status_code_override = status_code_override\n\n @property\n def status_code(self) -> int:\n return self._status_code_override or self.error_code.status_code\n\n\ndef log_onyx_error(exc: OnyxError) -> None:\n detail = exc.detail\n status_code = exc.status_code\n if status_code >= 500:\n logger.error(f\"OnyxError {exc.error_code.code}: {detail}\")\n elif status_code >= 400:\n logger.warning(f\"OnyxError {exc.error_code.code}: {detail}\")\n\n\ndef onyx_error_to_json_response(exc: OnyxError) -> JSONResponse:\n return JSONResponse(\n status_code=exc.status_code,\n content=exc.error_code.detail(exc.detail),\n )\n\n\ndef register_onyx_exception_handlers(app: FastAPI) -> None:\n \"\"\"Register a global handler that converts ``OnyxError`` to JSON responses.\n\n Must be called *after* the app is created but *before* it starts serving.\n The handler logs at WARNING for 4xx and ERROR for 5xx.\n \"\"\"\n\n @app.exception_handler(OnyxError)\n async def _handle_onyx_error(\n request: Request, # noqa: ARG001\n exc: OnyxError,\n ) -> JSONResponse:\n log_onyx_error(exc)\n return onyx_error_to_json_response(exc)\n" }, { "path": "backend/onyx/evals/README.md", "content": "# Onyx Evaluations\n\nThis directory contains the evaluation framework for testing and measuring the performance of Onyx's chat and retrieval systems.\n\n## Overview\n\nThe evaluation system uses [Braintrust](https://www.braintrust.dev/) to run automated evaluations against test datasets. It measures the quality of responses generated by Onyx's chat system and can be used to track performance improvements over time.\n\n## Prerequisites\n\n**Important**: The model server must be running in order for evals to work properly. Make sure your model server is up and running before executing any evaluations.\n\n## Running Evaluations\n\nKick off a remote job\n```bash\nonyx/backend$ python -m dotenv -f .vscode/.env run -- python onyx/evals/eval_cli.py --remote --api-key --search-permissions-email --remote --remote-dataset-name Simple\n```\n\nYou can also run the CLI directly from the command line:\n\n```bash\nonyx$ python -m dotenv -f .vscode/.env run -- python backend/onyx/evals/eval_cli.py --local-dataset-path backend/onyx/evals/data/eval.json --search-permissions-email richard@onyx.app\n```\nSave the env var ONYX_EVAL_API_KEY in your .env file so you don't have to specify it every time for triggering remote runs.\nYou'll need to create an API key in the admin panel to run evals.\n\n\n### Production Environment\n\n### Local Development\n\nFor local development, use the `eval_cli.py` script. We recommend starting it from the VS Code launch configuration for the best debugging experience.\n\n#### Using VS Code Launch Configuration\n\n1. Open VS Code in the project root\n2. Go to the \"Run and Debug\" panel (Ctrl/Cmd + Shift + D)\n3. Select \"Eval CLI\" from the dropdown\n4. Click the play button or press F5\n\nThis will run the evaluation with the following default settings:\n- Uses the local data file at `evals/data/data.json`\n- Enables verbose output\n- Sets up proper environment variables and Python path\n\n#### CLI Options\n\n- `--local-data-path`: Path to local JSON file containing test data (defaults to `evals/data/data.json`)\n- `--remote-dataset-name`: Name of remote Braintrust dataset\n- `--braintrust-project`: Braintrust project name (overrides `BRAINTRUST_PROJECT` env var)\n- `--verbose`: Enable verbose output\n- `--no-send-logs`: Skip sending logs to Braintrust (useful for local testing)\n- `--local-only`: Run evals locally without Braintrust, output results to CLI only\n\n## Test Data\n\nThe evaluation system uses test data stored in `evals/data/data.json`. This file contains a list of test cases, each with:\n- `input`: The question or prompt to test\n\nExample test case:\n```json\n{\n \"input\": {\n \"message\": \"What is the capital of France?\"\n }\n}\n```\n\n### Per-Test Configuration\n\nConfigure tool forcing, assertions, and model settings per-test by adding optional fields to each test case.\n\n#### Tool Configuration\n\n- `force_tools`: List of tool type names to force for this specific test\n- `expected_tools`: List of tool type names expected to be called\n- `require_all_tools`: If true, all expected tools must be called (default: false)\n\n#### Model Configuration\n\n- `model`: Model version to use (e.g., \"gpt-4o\", \"claude-3-5-sonnet\")\n- `model_provider`: Model provider (e.g., \"openai\", \"anthropic\")\n- `temperature`: Temperature for the model (default: 0.0)\n\nExample with tool and model configuration:\n```json\n[\n {\n \"input\": {\n \"message\": \"Find information about Python programming\"\n },\n \"expected_tools\": [\"SearchTool\"],\n \"force_tools\": [\"SearchTool\"],\n \"model\": \"gpt-4o\"\n },\n {\n \"input\": {\n \"message\": \"Search the web for recent news about AI\"\n },\n \"expected_tools\": [\"WebSearchTool\"],\n \"model\": \"claude-3-5-sonnet\",\n \"model_provider\": \"anthropic\"\n },\n {\n \"input\": {\n \"message\": \"Calculate 2 + 2\"\n },\n \"expected_tools\": [\"PythonTool\"],\n \"temperature\": 0.5\n }\n]\n```\n\n### Multi-Turn Evaluations\n\nFor testing realistic multi-turn conversations where each turn may require different tools, use the `messages` array format instead of a single `message`:\n\n```json\n{\n \"input\": {\n \"messages\": [\n {\n \"message\": \"What's the latest news about OpenAI today?\",\n \"expected_tools\": [\"WebSearchTool\", \"OpenURLTool\"]\n },\n {\n \"message\": \"Now search our internal docs for our OpenAI integration guide\",\n \"expected_tools\": [\"SearchTool\"]\n },\n {\n \"message\": \"Thanks, that's helpful!\",\n \"expected_tools\": []\n }\n ]\n }\n}\n```\n\nEach message in the `messages` array can have its own configuration:\n- `message`: The user message text (required)\n- `expected_tools`: List of tool types expected to be called for this turn\n- `require_all_tools`: If true, all expected tools must be called (default: false)\n- `force_tools`: List of tool types to force for this turn\n- `model`: Model version override for this turn\n- `model_provider`: Model provider override for this turn\n- `temperature`: Temperature override for this turn\n\nMulti-turn evals run within a single chat session, so the model has full context of previous turns when responding.\n\n### Available Tool Types\n\nThe following built-in tool types can be used:\n- `SearchTool`: Internal document search\n- `WebSearchTool`: Internet/web search\n- `ImageGenerationTool`: Image generation\n- `PythonTool`: Python code execution\n- `OpenURLTool`: Open and read URLs\n\n### Braintrust Dashboard\n\nAfter running evaluations, you can view results in the Braintrust dashboard. The evaluation will report:\n- `tool_assertion`: Score of 1.0 if tool assertions passed (or no assertions configured), 0.0 if failed\n- Metadata including `tools_called`, `tools_called_count`, and assertion details\n" }, { "path": "backend/onyx/evals/eval.py", "content": "import time\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\n\nfrom sqlalchemy import Engine\nfrom sqlalchemy import event\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.orm import sessionmaker\nfrom sqlalchemy.orm.session import SessionTransaction\n\nfrom onyx.chat.chat_state import ChatStateContainer\nfrom onyx.chat.models import ChatFullResponse\nfrom onyx.chat.process_message import gather_stream_full\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.configs.constants import DEFAULT_PERSONA_ID\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.engine.sql_engine import get_sqlalchemy_engine\nfrom onyx.db.users import get_user_by_email\nfrom onyx.evals.models import ChatFullEvalResult\nfrom onyx.evals.models import EvalationAck\nfrom onyx.evals.models import EvalConfigurationOptions\nfrom onyx.evals.models import EvalMessage\nfrom onyx.evals.models import EvalProvider\nfrom onyx.evals.models import EvalTimings\nfrom onyx.evals.models import EvalToolResult\nfrom onyx.evals.models import MultiTurnEvalResult\nfrom onyx.evals.models import ToolAssertion\nfrom onyx.evals.provider import get_provider\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.server.query_and_chat.models import AUTO_PLACE_AFTER_LATEST_MESSAGE\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\n@contextmanager\ndef isolated_ephemeral_session_factory(\n engine: Engine,\n) -> Generator[Callable[[], Session], None, None]:\n \"\"\"\n Create a session factory that creates sessions that run in a transaction that gets rolled back.\n This is useful for running evals without any lasting db side effects.\n \"\"\"\n tenant_id = get_current_tenant_id()\n schema_translate_map = {None: tenant_id}\n conn = engine.connect().execution_options(schema_translate_map=schema_translate_map)\n outer_tx = conn.begin()\n Maker = sessionmaker(bind=conn, expire_on_commit=False, future=True)\n\n def make_session() -> Session:\n s = Maker()\n s.begin_nested()\n\n @event.listens_for(s, \"after_transaction_end\")\n def _restart_savepoint(\n session: Session, transaction: SessionTransaction\n ) -> None:\n if transaction.nested and not (\n transaction._parent is not None and transaction._parent.nested\n ):\n session.begin_nested()\n\n return s\n\n try:\n yield make_session\n finally:\n outer_tx.rollback()\n conn.close()\n\n\ndef _chat_full_response_to_eval_result(\n full: ChatFullResponse,\n stream_start_time: float,\n) -> ChatFullEvalResult:\n \"\"\"Map ChatFullResponse from gather_stream_full to eval result components.\"\"\"\n tools_called = [tc.tool_name for tc in full.tool_calls]\n tool_call_details: list[dict[str, Any]] = [\n {\"tool_name\": tc.tool_name, \"tool_arguments\": tc.tool_arguments}\n for tc in full.tool_calls\n ]\n stream_end_time = time.time()\n total_ms = (stream_end_time - stream_start_time) * 1000\n timings = EvalTimings(\n total_ms=total_ms,\n llm_first_token_ms=None,\n tool_execution_ms={},\n stream_processing_ms=total_ms,\n )\n return ChatFullEvalResult(\n answer=full.answer,\n tools_called=tools_called,\n tool_call_details=tool_call_details,\n citations=full.citation_info,\n timings=timings,\n )\n\n\ndef evaluate_tool_assertions(\n tools_called: list[str],\n assertions: ToolAssertion | None,\n) -> tuple[bool | None, str | None]:\n \"\"\"\n Evaluate tool assertions against the tools that were called.\n\n Args:\n tools_called: List of tool names that were called during evaluation\n assertions: Tool assertions to check, or None if no assertions\n\n Returns:\n Tuple of (passed, details) where:\n - passed: True if assertions passed, False if failed, None if no assertions\n - details: Human-readable explanation of the result\n \"\"\"\n if assertions is None:\n return None, None\n\n expected_tools = set(assertions.expected_tools)\n called_tools = set(tools_called)\n\n if assertions.require_all:\n # All expected tools must be called\n missing_tools = expected_tools - called_tools\n if missing_tools:\n return False, (\n f\"Missing expected tools: {sorted(missing_tools)}. Called tools: {sorted(called_tools)}\"\n )\n return True, (\n f\"All expected tools called: {sorted(expected_tools)}. Called tools: {sorted(called_tools)}\"\n )\n else:\n # At least one expected tool must be called\n matched_tools = expected_tools & called_tools\n if not matched_tools:\n return False, (\n f\"None of expected tools called. Expected one of: {sorted(expected_tools)}. Called tools: {sorted(called_tools)}\"\n )\n return True, (\n f\"Expected tool(s) called: {sorted(matched_tools)}. Called tools: {sorted(called_tools)}\"\n )\n\n\ndef _get_answer_with_tools(\n eval_input: dict[str, Any],\n configuration: EvalConfigurationOptions,\n) -> EvalToolResult:\n \"\"\"\n Get answer from the chat system with full tool call tracking.\n\n Args:\n eval_input: Dictionary containing:\n - 'message': The user message to send\n - 'force_tools' (optional): List of tool types to force for this input\n - 'expected_tools' (optional): List of tool types expected to be called\n - 'require_all_tools' (optional): If true, all expected tools must be called\n - 'model' (optional): Model version to use (e.g., \"gpt-4o\", \"claude-3-5-sonnet\")\n - 'model_provider' (optional): Model provider (e.g., \"openai\", \"anthropic\")\n - 'temperature' (optional): Temperature for the model\n configuration: Evaluation configuration options\n\n Returns:\n EvalToolResult containing the answer and tool call information\n \"\"\"\n engine = get_sqlalchemy_engine()\n with isolated_ephemeral_session_factory(engine) as SessionLocal:\n with SessionLocal() as db_session:\n full_configuration = configuration.get_configuration(db_session)\n\n # Handle per-input tool forcing (from data file)\n forced_tool_ids: list[int] = []\n input_force_tools = eval_input.get(\"force_tools\", [])\n if input_force_tools:\n from onyx.db.tools import get_builtin_tool\n from onyx.tools.built_in_tools import BUILT_IN_TOOL_MAP\n\n for tool_type in input_force_tools:\n if tool_type in BUILT_IN_TOOL_MAP:\n tool_id = get_builtin_tool(\n db_session, BUILT_IN_TOOL_MAP[tool_type]\n ).id\n if tool_id not in forced_tool_ids:\n forced_tool_ids.append(tool_id)\n\n # Build tool assertions from per-input config\n tool_assertions: ToolAssertion | None = None\n input_expected_tools = eval_input.get(\"expected_tools\", [])\n if input_expected_tools:\n tool_assertions = ToolAssertion(\n expected_tools=input_expected_tools,\n require_all=eval_input.get(\"require_all_tools\", False),\n )\n\n # Handle per-input model configuration\n llm_override = full_configuration.llm\n input_model = eval_input.get(\"model\")\n input_model_provider = eval_input.get(\"model_provider\")\n input_temperature = eval_input.get(\"temperature\")\n\n if input_model or input_model_provider or input_temperature is not None:\n # Create a new LLMOverride with per-input values, falling back to config\n llm_override = LLMOverride(\n model_provider=input_model_provider or llm_override.model_provider,\n model_version=input_model or llm_override.model_version,\n temperature=(\n input_temperature\n if input_temperature is not None\n else llm_override.temperature\n ),\n )\n\n user = get_user_by_email(configuration.search_permissions_email, db_session)\n if not user:\n raise ValueError(\n f\"User not found for email: {configuration.search_permissions_email}\"\n )\n\n forced_tool_id = forced_tool_ids[0] if forced_tool_ids else None\n request = SendMessageRequest(\n message=eval_input[\"message\"],\n llm_override=llm_override,\n allowed_tool_ids=full_configuration.allowed_tool_ids,\n forced_tool_id=forced_tool_id,\n chat_session_info=ChatSessionCreationRequest(\n persona_id=DEFAULT_PERSONA_ID,\n description=\"Eval session\",\n ),\n )\n\n stream_start_time = time.time()\n state_container = ChatStateContainer()\n packets = handle_stream_message_objects(\n new_msg_req=request,\n user=user,\n db_session=db_session,\n external_state_container=state_container,\n )\n full = gather_stream_full(packets, state_container)\n\n result = _chat_full_response_to_eval_result(full, stream_start_time)\n\n # Evaluate tool assertions\n assertion_passed, assertion_details = evaluate_tool_assertions(\n result.tools_called, tool_assertions\n )\n\n logger.info(\n f\"Eval completed. Tools called: {result.tools_called}.\\n\"\n f\"Assertion passed: {assertion_passed}. Details: {assertion_details}\\n\"\n )\n\n return EvalToolResult(\n answer=result.answer,\n tools_called=result.tools_called,\n tool_call_details=result.tool_call_details,\n citations=result.citations,\n assertion_passed=assertion_passed,\n assertion_details=assertion_details,\n timings=result.timings,\n )\n\n\ndef _get_multi_turn_answer_with_tools(\n eval_input: dict[str, Any],\n configuration: EvalConfigurationOptions,\n) -> MultiTurnEvalResult:\n \"\"\"\n Get answers from a multi-turn conversation with tool call tracking for each turn.\n\n Args:\n eval_input: Dictionary containing:\n - 'messages': List of message dicts, each with:\n - 'message': The user message text\n - 'expected_tools' (optional): List of expected tool types\n - 'require_all_tools' (optional): If true, all expected tools must be called\n - 'model' (optional): Model version override for this turn\n - 'model_provider' (optional): Provider override for this turn\n - 'temperature' (optional): Temperature override for this turn\n - 'force_tools' (optional): List of tool types to force\n configuration: Evaluation configuration options\n\n Returns:\n MultiTurnEvalResult containing per-turn results and aggregate metrics\n \"\"\"\n messages_data = eval_input.get(\"messages\", [])\n if not messages_data:\n raise ValueError(\"Multi-turn eval requires 'messages' array in input\")\n\n # Parse messages into EvalMessage objects\n messages: list[EvalMessage] = []\n for msg_data in messages_data:\n messages.append(\n EvalMessage(\n message=msg_data[\"message\"],\n expected_tools=msg_data.get(\"expected_tools\", []),\n require_all_tools=msg_data.get(\"require_all_tools\", False),\n model=msg_data.get(\"model\"),\n model_provider=msg_data.get(\"model_provider\"),\n temperature=msg_data.get(\"temperature\"),\n force_tools=msg_data.get(\"force_tools\", []),\n )\n )\n\n turn_results: list[EvalToolResult] = []\n\n engine = get_sqlalchemy_engine()\n with isolated_ephemeral_session_factory(engine) as SessionLocal:\n with SessionLocal() as db_session:\n full_configuration = configuration.get_configuration(db_session)\n\n user = get_user_by_email(configuration.search_permissions_email, db_session)\n if not user:\n raise ValueError(\n f\"User not found for email: {configuration.search_permissions_email}\"\n )\n # Cache user_id to avoid SQLAlchemy expiration issues\n user_id = user.id\n\n # Create a single chat session for all turns\n chat_session = create_chat_session(\n db_session=db_session,\n description=\"Multi-turn eval session\",\n user_id=user_id,\n persona_id=DEFAULT_PERSONA_ID,\n onyxbot_flow=True,\n )\n chat_session_id = chat_session.id\n\n # Process each turn sequentially\n for turn_idx, msg in enumerate(messages):\n logger.info(\n f\"Processing turn {turn_idx + 1}/{len(messages)}: {msg.message[:50]}...\"\n )\n\n # Handle per-turn tool forcing\n forced_tool_ids: list[int] = []\n if msg.force_tools:\n from onyx.db.tools import get_builtin_tool\n from onyx.tools.built_in_tools import BUILT_IN_TOOL_MAP\n\n for tool_type in msg.force_tools:\n if tool_type in BUILT_IN_TOOL_MAP:\n tool_id = get_builtin_tool(\n db_session, BUILT_IN_TOOL_MAP[tool_type]\n ).id\n if tool_id not in forced_tool_ids:\n forced_tool_ids.append(tool_id)\n\n # Build tool assertions for this turn\n tool_assertions: ToolAssertion | None = None\n if msg.expected_tools:\n tool_assertions = ToolAssertion(\n expected_tools=msg.expected_tools,\n require_all=msg.require_all_tools,\n )\n\n # Handle per-turn model configuration\n llm_override = full_configuration.llm\n if msg.model or msg.model_provider or msg.temperature is not None:\n llm_override = LLMOverride(\n model_provider=msg.model_provider\n or llm_override.model_provider,\n model_version=msg.model or llm_override.model_version,\n temperature=(\n msg.temperature\n if msg.temperature is not None\n else llm_override.temperature\n ),\n )\n\n # Create request for this turn using SendMessageRequest (same API as handle_stream_message_objects)\n # Use AUTO_PLACE_AFTER_LATEST_MESSAGE to chain messages\n forced_tool_id = forced_tool_ids[0] if forced_tool_ids else None\n request = SendMessageRequest(\n chat_session_id=chat_session_id,\n parent_message_id=AUTO_PLACE_AFTER_LATEST_MESSAGE,\n message=msg.message,\n llm_override=llm_override,\n allowed_tool_ids=full_configuration.allowed_tool_ids,\n forced_tool_id=forced_tool_id,\n )\n\n # Stream and gather results for this turn via handle_stream_message_objects + gather_stream_full\n stream_start_time = time.time()\n state_container = ChatStateContainer()\n packets = handle_stream_message_objects(\n new_msg_req=request,\n user=user,\n db_session=db_session,\n external_state_container=state_container,\n )\n full = gather_stream_full(packets, state_container)\n\n result = _chat_full_response_to_eval_result(full, stream_start_time)\n\n # Evaluate tool assertions for this turn\n assertion_passed, assertion_details = evaluate_tool_assertions(\n result.tools_called, tool_assertions\n )\n\n logger.info(\n f\"Turn {turn_idx + 1} completed. Tools called: {result.tools_called}.\\n\"\n f\"Assertion passed: {assertion_passed}. Details: {assertion_details}\\n\"\n )\n\n turn_results.append(\n EvalToolResult(\n answer=result.answer,\n tools_called=result.tools_called,\n tool_call_details=result.tool_call_details,\n citations=result.citations,\n assertion_passed=assertion_passed,\n assertion_details=assertion_details,\n timings=result.timings,\n )\n )\n\n # Calculate aggregate metrics\n pass_count = sum(1 for r in turn_results if r.assertion_passed is True)\n fail_count = sum(1 for r in turn_results if r.assertion_passed is False)\n # Consider \"all passed\" only if there are no failures\n # (turns with no assertions don't count as failures)\n all_passed = fail_count == 0\n\n return MultiTurnEvalResult(\n turn_results=turn_results,\n all_passed=all_passed,\n pass_count=pass_count,\n fail_count=fail_count,\n total_turns=len(turn_results),\n )\n\n\ndef run_eval(\n configuration: EvalConfigurationOptions,\n data: list[dict[str, Any]] | None = None,\n remote_dataset_name: str | None = None,\n provider: EvalProvider = get_provider(),\n) -> EvalationAck:\n if data is not None and remote_dataset_name is not None:\n raise ValueError(\"Cannot specify both data and remote_dataset_name\")\n\n if data is None and remote_dataset_name is None:\n raise ValueError(\"Must specify either data or remote_dataset_name\")\n\n return provider.eval(\n task=lambda eval_input: _get_answer_with_tools(eval_input, configuration),\n configuration=configuration,\n data=data,\n remote_dataset_name=remote_dataset_name,\n multi_turn_task=lambda eval_input: _get_multi_turn_answer_with_tools(\n eval_input, configuration\n ),\n )\n" }, { "path": "backend/onyx/evals/eval_cli.py", "content": "#!/usr/bin/env python3\n\"\"\"\nCLI for running evaluations with local configurations.\n\"\"\"\n\nimport argparse\nimport json\nimport logging\nimport os\nfrom typing import Any\n\nimport braintrust\nimport requests\n\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE\nfrom onyx.configs.constants import POSTGRES_WEB_APP_NAME\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.evals.eval import run_eval\nfrom onyx.evals.models import EvalationAck\nfrom onyx.evals.models import EvalConfigurationOptions\nfrom onyx.evals.provider import get_provider\nfrom onyx.tracing.setup import setup_tracing\n\n\ndef setup_session_factory() -> None:\n SqlEngine.set_app_name(POSTGRES_WEB_APP_NAME)\n SqlEngine.init_engine(\n pool_size=POSTGRES_API_SERVER_POOL_SIZE,\n max_overflow=POSTGRES_API_SERVER_POOL_OVERFLOW,\n )\n\n\ndef load_data_local(\n local_data_path: str,\n) -> list[dict[str, Any]]:\n if not os.path.isfile(local_data_path):\n raise ValueError(f\"Local data file does not exist: {local_data_path}\")\n with open(local_data_path, \"r\") as f:\n return json.load(f)\n\n\ndef configure_logging_for_evals(verbose: bool) -> None:\n \"\"\"Set logging level to WARNING to reduce noise during evals.\"\"\"\n if verbose:\n return\n\n # Set environment variable for any future logger creation\n os.environ[\"LOG_LEVEL\"] = \"WARNING\"\n\n # Force WARNING level for root logger and its handlers\n root = logging.getLogger()\n root.setLevel(logging.WARNING)\n for handler in root.handlers:\n handler.setLevel(logging.WARNING)\n\n # Force WARNING level for all existing loggers and their handlers\n for name in list(logging.Logger.manager.loggerDict.keys()):\n logger = logging.getLogger(name)\n logger.setLevel(logging.WARNING)\n for handler in logger.handlers:\n handler.setLevel(logging.WARNING)\n\n # Set a basic config to ensure new loggers also use WARNING\n logging.basicConfig(level=logging.WARNING, force=True)\n\n\ndef run_local(\n local_data_path: str | None,\n remote_dataset_name: str | None,\n search_permissions_email: str | None = None,\n no_send_logs: bool = False,\n local_only: bool = False,\n verbose: bool = False,\n) -> EvalationAck:\n \"\"\"\n Run evaluation with local configurations.\n\n Tool forcing and assertions are configured per-test in the data file using:\n - force_tools: List of tool type names to force\n - expected_tools: List of tool type names expected to be called\n - require_all_tools: If true, all expected tools must be called\n\n Args:\n local_data_path: Path to local JSON file\n remote_dataset_name: Name of remote Braintrust dataset\n search_permissions_email: Optional email address to impersonate for the evaluation\n no_send_logs: Whether to skip sending logs to Braintrust\n local_only: If True, use LocalEvalProvider (CLI output only, no Braintrust)\n\n Returns:\n EvalationAck: The evaluation result\n \"\"\"\n setup_session_factory()\n configure_logging_for_evals(\n verbose=verbose,\n )\n # Only setup tracing if not running in local-only mode\n if not local_only:\n setup_tracing()\n\n if search_permissions_email is None:\n raise ValueError(\"search_permissions_email is required for local evaluation\")\n\n configuration = EvalConfigurationOptions(\n search_permissions_email=search_permissions_email,\n dataset_name=remote_dataset_name or \"local\",\n no_send_logs=no_send_logs,\n )\n\n # Get the appropriate provider\n provider = get_provider(local_only=local_only)\n\n if remote_dataset_name:\n score = run_eval(\n configuration=configuration,\n remote_dataset_name=remote_dataset_name,\n provider=provider,\n )\n else:\n if local_data_path is None:\n raise ValueError(\n \"local_data_path or remote_dataset_name is required for local evaluation\"\n )\n data = load_data_local(local_data_path)\n score = run_eval(configuration=configuration, data=data, provider=provider)\n\n return score\n\n\ndef run_remote(\n base_url: str,\n api_key: str,\n remote_dataset_name: str,\n search_permissions_email: str,\n payload: dict[str, Any] | None = None,\n) -> dict[str, Any]:\n \"\"\"\n Trigger an eval pipeline execution on a remote server.\n\n Tool forcing and assertions are configured per-test in the dataset.\n\n Args:\n base_url: Base URL of the remote server (e.g., \"https://test.onyx.app\")\n api_key: API key for authentication\n remote_dataset_name: Name of remote Braintrust dataset\n search_permissions_email: Email address to use for the evaluation.\n payload: Optional payload to send with the request\n\n Returns:\n Response from the remote server\n\n Raises:\n requests.RequestException: If the request fails\n \"\"\"\n if payload is None:\n payload = {}\n\n payload[\"search_permissions_email\"] = search_permissions_email\n payload[\"dataset_name\"] = remote_dataset_name\n\n url = f\"{base_url}/api/evals/eval_run\"\n headers = {\n \"Authorization\": f\"Bearer {api_key}\",\n \"Content-Type\": \"application/json\",\n }\n response = requests.post(url, headers=headers, json=payload)\n\n response.raise_for_status()\n return response.json()\n\n\ndef main() -> None:\n \"\"\"Main CLI entry point.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Run evaluations with local configurations\"\n )\n\n parser.add_argument(\n \"--local-data-path\",\n type=str,\n help=\"Path to local JSON file containing test data\",\n )\n\n parser.add_argument(\n \"--remote-dataset-name\",\n type=str,\n help=\"Name of remote Braintrust dataset\",\n )\n\n parser.add_argument(\n \"--braintrust-project\",\n type=str,\n help=\"Braintrust project name\",\n default=\"Onyx\",\n )\n\n parser.add_argument(\"--verbose\", action=\"store_true\", help=\"Enable verbose output\")\n\n # Remote eval arguments\n parser.add_argument(\n \"--base-url\",\n type=str,\n default=\"https://test.onyx.app\",\n help=\"Base URL of the remote server (default: https://test.onyx.app)\",\n )\n\n parser.add_argument(\n \"--api-key\",\n type=str,\n help=\"API key for authentication with the remote server\",\n )\n\n parser.add_argument(\n \"--remote\",\n action=\"store_true\",\n help=\"Run evaluation on remote server instead of locally\",\n )\n\n parser.add_argument(\n \"--search-permissions-email\",\n type=str,\n help=\"Email address to impersonate for the evaluation\",\n )\n\n parser.add_argument(\n \"--no-send-logs\",\n action=\"store_true\",\n help=\"Do not send logs to the remote server\",\n default=False,\n )\n\n parser.add_argument(\n \"--local-only\",\n action=\"store_true\",\n help=\"Run evals locally without Braintrust, output results to CLI only\",\n default=False,\n )\n\n args = parser.parse_args()\n\n if args.local_data_path:\n print(f\"Loading data from local file: {args.local_data_path}\")\n elif args.remote_dataset_name:\n if args.local_only:\n raise ValueError(\n \"--local-only cannot be used with --remote-dataset-name. Use --local-data-path with a local JSON file instead.\"\n )\n print(f\"Loading data from remote dataset: {args.remote_dataset_name}\")\n dataset = braintrust.init_dataset(\n project=args.braintrust_project, name=args.remote_dataset_name\n )\n dataset_size = len(list(dataset.fetch()))\n print(f\"Dataset size: {dataset_size}\")\n if args.remote:\n if not args.api_key:\n print(\"Using API Key from ONYX_EVAL_API_KEY\")\n api_key: str = (\n args.api_key if args.api_key else os.environ.get(\"ONYX_EVAL_API_KEY\", \"\")\n )\n print(f\"Running evaluation on remote server: {args.base_url}\")\n\n if args.search_permissions_email:\n print(f\"Using search permissions email: {args.search_permissions_email}\")\n\n try:\n result = run_remote(\n args.base_url,\n api_key,\n args.remote_dataset_name,\n search_permissions_email=args.search_permissions_email,\n )\n print(f\"Remote evaluation triggered successfully: {result}\")\n except requests.RequestException as e:\n print(f\"Error triggering remote evaluation: {e}\")\n return\n else:\n if args.local_only:\n print(\"Running in local-only mode (no Braintrust)\")\n else:\n print(f\"Using Braintrust project: {args.braintrust_project}\")\n\n if args.search_permissions_email:\n print(f\"Using search permissions email: {args.search_permissions_email}\")\n\n run_local(\n local_data_path=args.local_data_path,\n remote_dataset_name=args.remote_dataset_name,\n search_permissions_email=args.search_permissions_email,\n no_send_logs=args.no_send_logs,\n local_only=args.local_only,\n verbose=args.verbose,\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/evals/models.py", "content": "from abc import ABC\nfrom abc import abstractmethod\nfrom collections.abc import Callable\nfrom typing import Any\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.tools import get_builtin_tool\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.tools.built_in_tools import BUILT_IN_TOOL_MAP\n\n\nclass ToolAssertion(BaseModel):\n \"\"\"Assertion about expected tool usage during evaluation.\"\"\"\n\n expected_tools: list[str] # Tool type names that should be called\n require_all: bool = False # If True, ALL expected tools must be called\n\n\nclass EvalTimings(BaseModel):\n \"\"\"Timing information for eval execution.\"\"\"\n\n total_ms: float # Total time for the eval\n llm_first_token_ms: float | None = None # Time to first token from LLM\n tool_execution_ms: dict[str, float] = Field(\n default_factory=dict\n ) # Per-tool timings\n stream_processing_ms: float | None = None # Time to process the stream\n\n\nclass ChatFullEvalResult(BaseModel):\n \"\"\"Raw eval components from ChatFullResponse (before tool assertions).\"\"\"\n\n answer: str\n tools_called: list[str]\n tool_call_details: list[dict[str, Any]]\n citations: list[CitationInfo]\n timings: EvalTimings\n\n\nclass EvalToolResult(BaseModel):\n \"\"\"Result of a single eval with tool call information.\"\"\"\n\n answer: str\n tools_called: list[str] # Names of tools that were called\n tool_call_details: list[dict[str, Any]] # Full tool call info\n citations: list[CitationInfo] # Citations used in the answer\n assertion_passed: bool | None = None # None if no assertion configured\n assertion_details: str | None = None # Explanation of pass/fail\n timings: EvalTimings | None = None # Timing information for the eval\n\n\nclass EvalMessage(BaseModel):\n \"\"\"Single message in a multi-turn evaluation conversation.\"\"\"\n\n message: str # The message text to send\n expected_tools: list[str] = Field(\n default_factory=list\n ) # Expected tools for this turn\n require_all_tools: bool = False # If True, ALL expected tools must be called\n # Per-message model configuration overrides\n model: str | None = None\n model_provider: str | None = None\n temperature: float | None = None\n force_tools: list[str] = Field(default_factory=list) # Tools to force for this turn\n\n\nclass MultiTurnEvalResult(BaseModel):\n \"\"\"Result of a multi-turn evaluation containing per-message results.\"\"\"\n\n turn_results: list[EvalToolResult] # Results for each turn/message\n all_passed: bool # True if all turn assertions passed\n pass_count: int # Number of turns that passed\n fail_count: int # Number of turns that failed\n total_turns: int # Total number of turns\n\n\nclass EvalConfiguration(BaseModel):\n llm: LLMOverride = Field(default_factory=LLMOverride)\n search_permissions_email: str\n allowed_tool_ids: list[int]\n\n\nclass EvalConfigurationOptions(BaseModel):\n builtin_tool_types: list[str] = list(BUILT_IN_TOOL_MAP.keys())\n llm: LLMOverride = LLMOverride(\n model_provider=None,\n model_version=\"gpt-4o\",\n temperature=0.0,\n )\n search_permissions_email: str\n dataset_name: str\n no_send_logs: bool = False\n # Optional override for Braintrust project (defaults to BRAINTRUST_PROJECT env var)\n braintrust_project: str | None = None\n # Optional experiment name for the eval run (shows in Braintrust UI)\n experiment_name: str | None = None\n\n def get_configuration(self, db_session: Session) -> EvalConfiguration:\n return EvalConfiguration(\n llm=self.llm,\n search_permissions_email=self.search_permissions_email,\n allowed_tool_ids=[\n get_builtin_tool(db_session, BUILT_IN_TOOL_MAP[tool]).id\n for tool in self.builtin_tool_types\n ],\n )\n\n\nclass EvalationAck(BaseModel):\n success: bool\n\n\nclass EvalProvider(ABC):\n @abstractmethod\n def eval(\n self,\n task: Callable[[dict[str, Any]], EvalToolResult],\n configuration: EvalConfigurationOptions,\n data: list[dict[str, Any]] | None = None,\n remote_dataset_name: str | None = None,\n multi_turn_task: \"Callable[[dict[str, Any]], MultiTurnEvalResult] | None\" = None,\n ) -> EvalationAck:\n pass\n" }, { "path": "backend/onyx/evals/one_off/create_braintrust_dataset.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to create a Braintrust dataset from the DR Master Question & Metric Sheet CSV.\n\nThis script:\n1. Parses the CSV file\n2. Filters records where \"Should we use it\" is TRUE and \"web-only\" is in categories\n3. Creates a Braintrust dataset with Question as input and research_type metadata\n\nUsage:\n python create_braintrust_dataset.py --dataset-name \"MyDataset\"\n python create_braintrust_dataset.py --dataset-name \"MyDataset\" --csv-path \"/path/to/csv\"\n\"\"\"\n\nimport argparse\nimport csv\nimport os\nimport sys\nfrom typing import Any\nfrom typing import Dict\nfrom typing import List\n\nfrom onyx.configs.app_configs import BRAINTRUST_API_KEY\n\ntry:\n from braintrust import init_dataset\nexcept ImportError:\n print(\n \"Error: braintrust package not found. Please install it with: pip install braintrust\"\n )\n sys.exit(1)\n\n\ndef column_letter_to_index(column_letter: str) -> int:\n \"\"\"Convert Google Sheets column letter (A, B, C, etc.) to 0-based index.\"\"\"\n result = 0\n for char in column_letter.upper():\n result = result * 26 + (ord(char) - ord(\"A\") + 1)\n return result - 1\n\n\ndef parse_csv_file(csv_path: str) -> List[Dict[str, Any]]:\n \"\"\"Parse the CSV file and extract relevant records.\"\"\"\n records = []\n\n with open(csv_path, \"r\", encoding=\"utf-8\") as file:\n # Skip the first few header rows and read the actual data\n lines = file.readlines()\n\n # Find the actual data start (skip header rows)\n data_start = 0\n for i, line in enumerate(lines):\n if \"Should we use it?\" in line:\n data_start = i + 1\n break\n\n # Parse the CSV data starting from the data_start line\n csv_reader = csv.reader(lines[data_start:])\n\n # Define Google Sheets column references for easy modification\n SHOULD_USE_COL = \"C\" # \"Should we use it?\"\n QUESTION_COL = \"H\" # \"Question\"\n EXPECTED_DEPTH_COL = \"J\" # \"Expected Depth\"\n CATEGORIES_COL = \"M\" # \"Categories\"\n OPENAI_DEEP_COL = \"AA\" # \"OpenAI Deep Answer\"\n OPENAI_THINKING_COL = \"O\" # \"OpenAI Thinking Answer\"\n\n for row_num, row in enumerate(csv_reader, start=data_start + 1):\n if len(row) < 15: # Ensure we have enough columns\n continue\n\n # Extract relevant fields using Google Sheets column references\n should_use = (\n row[column_letter_to_index(SHOULD_USE_COL)].strip().upper()\n if len(row) > column_letter_to_index(SHOULD_USE_COL)\n else \"\"\n )\n question = (\n row[column_letter_to_index(QUESTION_COL)].strip()\n if len(row) > column_letter_to_index(QUESTION_COL)\n else \"\"\n )\n expected_depth = (\n row[column_letter_to_index(EXPECTED_DEPTH_COL)].strip()\n if len(row) > column_letter_to_index(EXPECTED_DEPTH_COL)\n else \"\"\n )\n categories = (\n row[column_letter_to_index(CATEGORIES_COL)].strip()\n if len(row) > column_letter_to_index(CATEGORIES_COL)\n else \"\"\n )\n openai_deep_answer = (\n row[column_letter_to_index(OPENAI_DEEP_COL)].strip()\n if len(row) > column_letter_to_index(OPENAI_DEEP_COL)\n else \"\"\n )\n openai_thinking_answer = (\n row[column_letter_to_index(OPENAI_THINKING_COL)].strip()\n if len(row) > column_letter_to_index(OPENAI_THINKING_COL)\n else \"\"\n )\n\n # Filter records: should_use = TRUE and categories contains \"web-only\"\n if (\n should_use == \"TRUE\" and \"web-only\" in categories and question\n ): # Ensure question is not empty\n if expected_depth == \"Deep\":\n records.extend(\n [\n {\n \"question\": question\n + \". All info is contained in the quesiton. DO NOT ask any clarifying questions.\",\n \"research_type\": \"DEEP\",\n \"categories\": categories,\n \"expected_depth\": expected_depth,\n \"expected_answer\": openai_deep_answer,\n \"row_number\": row_num,\n }\n ]\n )\n else:\n records.extend(\n [\n {\n \"question\": question,\n \"research_type\": \"THOUGHTFUL\",\n \"categories\": categories,\n \"expected_depth\": expected_depth,\n \"expected_answer\": openai_thinking_answer,\n \"row_number\": row_num,\n }\n ]\n )\n\n return records\n\n\ndef create_braintrust_dataset(records: List[Dict[str, Any]], dataset_name: str) -> None:\n \"\"\"Create a Braintrust dataset with the filtered records.\"\"\"\n\n # Check if BRAINTRUST_API_KEY is set\n if BRAINTRUST_API_KEY == \"\":\n print(\"WARNING: BRAINTRUST_API_KEY environment variable is not set.\")\n print(\n \"The script will show what would be inserted but won't actually create the dataset.\"\n )\n print(\n \"To actually create the dataset, set your BRAINTRUST_API_KEY environment variable.\"\n )\n print()\n\n # Show what would be inserted\n print(\n f\"Would create Braintrust dataset '{dataset_name}' with {len(records)} records:\"\n )\n for i, record in enumerate(records, 1):\n print(f\"Record {i}/{len(records)}:\")\n print(f\" Question: {record['question'][:100]}...\")\n print(f\" Research Type: {record['research_type']}\")\n print(f\" Expected Answer: {record['expected_answer'][:100]}...\")\n print()\n return\n\n # Initialize the dataset\n dataset = init_dataset(\"Onyx\", dataset_name, api_key=BRAINTRUST_API_KEY)\n\n print(f\"Creating Braintrust dataset with {len(records)} records...\")\n\n # Insert records into the dataset\n for i, record in enumerate(records, 1):\n record_id = dataset.insert(\n {\"message\": record[\"question\"], \"research_type\": record[\"research_type\"]},\n expected=record[\"expected_answer\"],\n )\n print(f\"Inserted record {i}/{len(records)}: ID {record_id}\")\n print(f\" Question: {record['question'][:100]}...\")\n print(f\" Research Type: {record['research_type']}\")\n print(f\" Expected Answer: {record['expected_answer'][:100]}...\")\n print()\n\n # Flush to ensure all records are sent\n dataset.flush()\n print(f\"Successfully created dataset with {len(records)} records!\")\n\n\ndef main() -> None:\n \"\"\"Main function to run the script.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Create a Braintrust dataset from the DR Master Question & Metric Sheet CSV\"\n )\n parser.add_argument(\n \"--dataset-name\", required=True, help=\"Name of the Braintrust dataset to create\"\n )\n parser.add_argument(\n \"--csv-path\",\n default=\"/Users/richardguan/onyx/backend/onyx/evals/data/DR Master Question & Metric Sheet - Sheet1.csv\",\n help=\"Path to the CSV file (default: %(default)s)\",\n )\n\n args = parser.parse_args()\n\n csv_path = args.csv_path\n dataset_name = args.dataset_name\n\n if not os.path.exists(csv_path):\n print(f\"Error: CSV file not found at {csv_path}\")\n sys.exit(1)\n\n print(\"Parsing CSV file...\")\n records = parse_csv_file(csv_path)\n\n print(f\"Found {len(records)} records matching criteria:\")\n print(\"- Should we use it = TRUE\")\n print(\"- Categories contains 'web-only'\")\n print(\"- Question is not empty\")\n print()\n\n if not records:\n print(\"No records found matching the criteria!\")\n sys.exit(1)\n\n # Show summary of research types\n deep_count = sum(1 for r in records if r[\"research_type\"] == \"DEEP\")\n thoughtful_count = sum(1 for r in records if r[\"research_type\"] == \"THOUGHTFUL\")\n\n print(\"Research type breakdown:\")\n print(f\" DEEP: {deep_count}\")\n print(f\" THOUGHTFUL: {thoughtful_count}\")\n print()\n\n # Create the Braintrust dataset\n create_braintrust_dataset(records, dataset_name)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/evals/provider.py", "content": "from onyx.evals.models import EvalProvider\nfrom onyx.evals.providers.braintrust import BraintrustEvalProvider\nfrom onyx.evals.providers.local import LocalEvalProvider\n\n\ndef get_provider(local_only: bool = False) -> EvalProvider:\n \"\"\"\n Get the appropriate eval provider.\n\n Args:\n local_only: If True, use LocalEvalProvider (CLI output only, no Braintrust).\n If False, use BraintrustEvalProvider.\n\n Returns:\n The appropriate EvalProvider instance.\n \"\"\"\n if local_only:\n return LocalEvalProvider()\n return BraintrustEvalProvider()\n" }, { "path": "backend/onyx/evals/providers/braintrust.py", "content": "from collections.abc import Callable\nfrom typing import Any\nfrom typing import Union\n\nfrom braintrust import Eval\nfrom braintrust import EvalCase\nfrom braintrust import init_dataset\nfrom braintrust import Score\n\nfrom onyx.configs.app_configs import BRAINTRUST_MAX_CONCURRENCY\nfrom onyx.configs.app_configs import BRAINTRUST_PROJECT\nfrom onyx.evals.models import EvalationAck\nfrom onyx.evals.models import EvalConfigurationOptions\nfrom onyx.evals.models import EvalProvider\nfrom onyx.evals.models import EvalToolResult\nfrom onyx.evals.models import MultiTurnEvalResult\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Union type for both single and multi-turn results\nEvalResult = Union[EvalToolResult, MultiTurnEvalResult]\n\n\ndef tool_assertion_scorer(\n input: dict[str, Any], output: EvalResult, expected: EvalResult | None\n) -> Score:\n \"\"\"\n Scorer that checks if tool assertions passed.\n\n Handles both single-turn (EvalToolResult) and multi-turn (MultiTurnEvalResult) outputs.\n\n Args:\n input: The input data for the evaluation case.\n output: The actual output from the task.\n expected: The expected output (unused for this scorer).\n\n Returns:\n Score with value 1.0 if passed or no assertions, 0.0 if failed.\n \"\"\"\n # input and expected are unused but required by Braintrust scorer signature\n _ = input, expected\n\n # Handle multi-turn results\n if isinstance(output, MultiTurnEvalResult):\n # Calculate score based on pass rate\n if output.total_turns == 0:\n score = 1.0\n else:\n # Score is the ratio of passed assertions\n assertions_evaluated = output.pass_count + output.fail_count\n if assertions_evaluated == 0:\n score = 1.0 # No assertions configured\n else:\n score = output.pass_count / assertions_evaluated\n\n return Score(\n name=\"tool_assertion\",\n score=score,\n metadata={\n \"is_multi_turn\": True,\n \"total_turns\": output.total_turns,\n \"pass_count\": output.pass_count,\n \"fail_count\": output.fail_count,\n \"all_passed\": output.all_passed,\n \"turn_details\": [\n {\n \"tools_called\": r.tools_called,\n \"assertion_passed\": r.assertion_passed,\n \"assertion_details\": r.assertion_details,\n }\n for r in output.turn_results\n ],\n },\n )\n\n # Handle single-turn results (EvalToolResult)\n if output.assertion_passed is None:\n # No assertions configured - return passing score\n return Score(\n name=\"tool_assertion\",\n score=1.0,\n metadata={\n \"is_multi_turn\": False,\n \"tools_called\": output.tools_called,\n \"tools_called_count\": len(output.tools_called),\n \"assertion_configured\": False,\n },\n )\n\n return Score(\n name=\"tool_assertion\",\n score=1.0 if output.assertion_passed else 0.0,\n metadata={\n \"is_multi_turn\": False,\n \"tools_called\": output.tools_called,\n \"tools_called_count\": len(output.tools_called),\n \"assertion_passed\": output.assertion_passed,\n \"assertion_details\": output.assertion_details,\n \"tool_call_details\": output.tool_call_details,\n },\n )\n\n\nclass BraintrustEvalProvider(EvalProvider):\n def eval(\n self,\n task: Callable[[dict[str, Any]], EvalToolResult],\n configuration: EvalConfigurationOptions,\n data: list[dict[str, Any]] | None = None,\n remote_dataset_name: str | None = None,\n multi_turn_task: Callable[[dict[str, Any]], MultiTurnEvalResult] | None = None,\n ) -> EvalationAck:\n if data is not None and remote_dataset_name is not None:\n raise ValueError(\"Cannot specify both data and remote_dataset_name\")\n if data is None and remote_dataset_name is None:\n raise ValueError(\"Must specify either data or remote_dataset_name\")\n\n # Create a wrapper task that dispatches to the appropriate handler\n def dispatch_task(eval_input: dict[str, Any]) -> EvalResult:\n if \"messages\" in eval_input and multi_turn_task is not None:\n return multi_turn_task(eval_input)\n return task(eval_input)\n\n project_name = configuration.braintrust_project or BRAINTRUST_PROJECT\n experiment_name = configuration.experiment_name\n\n eval_data: Any = None\n if remote_dataset_name is not None:\n eval_data = init_dataset(project=project_name, name=remote_dataset_name)\n else:\n if data:\n eval_data = [\n EvalCase(\n input={\n **item.get(\"input\", {}),\n # Pass through per-test tool configuration (for single-turn)\n \"force_tools\": item.get(\"force_tools\", []),\n \"expected_tools\": item.get(\"expected_tools\", []),\n \"require_all_tools\": item.get(\"require_all_tools\", False),\n # Pass through per-test model configuration\n \"model\": item.get(\"model\"),\n \"model_provider\": item.get(\"model_provider\"),\n \"temperature\": item.get(\"temperature\"),\n },\n expected=item.get(\"expected\"),\n )\n for item in data\n ]\n\n metadata = configuration.model_dump()\n\n Eval( # type: ignore[misc]\n name=project_name,\n experiment_name=experiment_name,\n data=eval_data,\n task=dispatch_task,\n scores=[tool_assertion_scorer],\n metadata=metadata,\n max_concurrency=BRAINTRUST_MAX_CONCURRENCY,\n no_send_logs=configuration.no_send_logs,\n )\n return EvalationAck(success=True)\n" }, { "path": "backend/onyx/evals/providers/local.py", "content": "\"\"\"\nLocal eval provider that runs evaluations and outputs results to the CLI.\nNo external dependencies like Braintrust required.\n\"\"\"\n\nfrom collections.abc import Callable\nfrom typing import Any\n\nfrom onyx.evals.models import EvalationAck\nfrom onyx.evals.models import EvalConfigurationOptions\nfrom onyx.evals.models import EvalProvider\nfrom onyx.evals.models import EvalToolResult\nfrom onyx.evals.models import MultiTurnEvalResult\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# ANSI color codes\nGREEN = \"\\033[92m\"\nRED = \"\\033[91m\"\nYELLOW = \"\\033[93m\"\nBLUE = \"\\033[94m\"\nBOLD = \"\\033[1m\"\nRESET = \"\\033[0m\"\nDIM = \"\\033[2m\"\n\n\ndef _display_single_turn_result(\n result: EvalToolResult,\n passed_count: list[int],\n failed_count: list[int],\n no_assertion_count: list[int],\n) -> None:\n \"\"\"Display results for a single turn and update counters.\"\"\"\n # Display timing trace\n if result.timings:\n print(f\" {BOLD}Trace:{RESET}\")\n print(f\" Total: {result.timings.total_ms:.0f}ms\")\n if result.timings.llm_first_token_ms is not None:\n print(f\" First token: {result.timings.llm_first_token_ms:.0f}ms\")\n if result.timings.tool_execution_ms:\n for tool_name, duration_ms in result.timings.tool_execution_ms.items():\n print(f\" {tool_name}: {duration_ms:.0f}ms\")\n\n # Display tools called\n tools_str = \", \".join(result.tools_called) if result.tools_called else \"(none)\"\n print(f\" Tools called: {BLUE}{tools_str}{RESET}\")\n\n # Display assertion result\n if result.assertion_passed is None:\n print(f\" Assertion: {YELLOW}N/A{RESET} - No assertion configured\")\n no_assertion_count[0] += 1\n elif result.assertion_passed:\n print(f\" Assertion: {GREEN}PASS{RESET} - {result.assertion_details}\")\n passed_count[0] += 1\n else:\n print(f\" Assertion: {RED}FAIL{RESET} - {result.assertion_details}\")\n failed_count[0] += 1\n\n # Display truncated answer\n answer = result.answer\n truncated_answer = answer[:200] + \"...\" if len(answer) > 200 else answer\n truncated_answer = truncated_answer.replace(\"\\n\", \" \")\n print(f\" Answer: {truncated_answer}\")\n\n\nclass LocalEvalProvider(EvalProvider):\n \"\"\"\n Eval provider that runs evaluations locally and prints results to the CLI.\n Does not require Braintrust or any external service.\n \"\"\"\n\n def eval(\n self,\n task: Callable[[dict[str, Any]], EvalToolResult],\n configuration: EvalConfigurationOptions, # noqa: ARG002\n data: list[dict[str, Any]] | None = None,\n remote_dataset_name: str | None = None,\n multi_turn_task: Callable[[dict[str, Any]], MultiTurnEvalResult] | None = None,\n ) -> EvalationAck:\n if remote_dataset_name is not None:\n raise ValueError(\n \"LocalEvalProvider does not support remote datasets. Use --local-data-path with a local JSON file.\"\n )\n\n if data is None:\n raise ValueError(\"data is required for LocalEvalProvider\")\n\n total = len(data)\n # Use lists to allow mutation in helper function\n passed = [0]\n failed = [0]\n no_assertion = [0]\n\n print(f\"\\n{BOLD}Running {total} evaluation(s)...{RESET}\\n\")\n print(\"=\" * 60)\n\n for i, item in enumerate(data, 1):\n input_data = item.get(\"input\", {})\n\n # Check if this is a multi-turn eval (has 'messages' array)\n if \"messages\" in input_data:\n self._run_multi_turn_eval(\n i, total, item, multi_turn_task, passed, failed, no_assertion\n )\n else:\n self._run_single_turn_eval(\n i, total, item, task, passed, failed, no_assertion\n )\n\n # Summary\n print(\"\\n\" + \"=\" * 60)\n total_with_assertions = passed[0] + failed[0]\n if total_with_assertions > 0:\n pass_rate = (passed[0] / total_with_assertions) * 100\n print(\n f\"{BOLD}Summary:{RESET} {passed[0]}/{total_with_assertions} passed ({pass_rate:.1f}%)\"\n )\n else:\n print(f\"{BOLD}Summary:{RESET} No assertions configured\")\n\n print(f\" {GREEN}Passed:{RESET} {passed[0]}\")\n print(f\" {RED}Failed:{RESET} {failed[0]}\")\n if no_assertion[0] > 0:\n print(f\" {YELLOW}No assertion:{RESET} {no_assertion[0]}\")\n print(\"=\" * 60 + \"\\n\")\n\n # Return success if no failures\n return EvalationAck(success=(failed[0] == 0))\n\n def _run_single_turn_eval(\n self,\n i: int,\n total: int,\n item: dict[str, Any],\n task: Callable[[dict[str, Any]], EvalToolResult],\n passed: list[int],\n failed: list[int],\n no_assertion: list[int],\n ) -> None:\n \"\"\"Run a single-turn evaluation.\"\"\"\n # Build input with tool and model config\n eval_input = {\n **item.get(\"input\", {}),\n # Tool configuration\n \"force_tools\": item.get(\"force_tools\", []),\n \"expected_tools\": item.get(\"expected_tools\", []),\n \"require_all_tools\": item.get(\"require_all_tools\", False),\n # Model configuration\n \"model\": item.get(\"model\"),\n \"model_provider\": item.get(\"model_provider\"),\n \"temperature\": item.get(\"temperature\"),\n }\n\n message = eval_input.get(\"message\", \"(no message)\")\n truncated_message = message[:50] + \"...\" if len(message) > 50 else message\n\n # Show model if specified\n model_info = \"\"\n if item.get(\"model\"):\n model_info = f\" [{item.get('model')}]\"\n\n print(f'\\n{BOLD}[{i}/{total}]{RESET} \"{truncated_message}\"{model_info}')\n\n try:\n result = task(eval_input)\n _display_single_turn_result(result, passed, failed, no_assertion)\n except Exception as e:\n print(f\" {RED}ERROR:{RESET} {e}\")\n failed[0] += 1\n logger.exception(f\"Error running eval for input: {message}\")\n\n def _run_multi_turn_eval(\n self,\n i: int,\n total: int,\n item: dict[str, Any],\n multi_turn_task: Callable[[dict[str, Any]], MultiTurnEvalResult] | None,\n passed: list[int],\n failed: list[int],\n no_assertion: list[int],\n ) -> None:\n \"\"\"Run a multi-turn evaluation.\"\"\"\n if multi_turn_task is None:\n print(\n f\"\\n{BOLD}[{i}/{total}]{RESET} {RED}ERROR:{RESET} Multi-turn task not configured\"\n )\n failed[0] += 1\n return\n\n input_data = item.get(\"input\", {})\n messages = input_data.get(\"messages\", [])\n num_turns = len(messages)\n\n # Show first message as preview\n first_msg = (\n messages[0].get(\"message\", \"(no message)\") if messages else \"(no messages)\"\n )\n truncated_first = first_msg[:40] + \"...\" if len(first_msg) > 40 else first_msg\n\n print(f\"\\n{BOLD}[{i}/{total}] Multi-turn ({num_turns} turns){RESET}\")\n print(f' First: \"{truncated_first}\"')\n\n try:\n # Pass the full input with messages\n eval_input = {**input_data}\n result = multi_turn_task(eval_input)\n\n # Display each turn's result\n for turn_idx, turn_result in enumerate(result.turn_results):\n turn_msg = messages[turn_idx].get(\"message\", \"\")\n truncated_turn = (\n turn_msg[:40] + \"...\" if len(turn_msg) > 40 else turn_msg\n )\n print(f'\\n {DIM}Turn {turn_idx + 1}:{RESET} \"{truncated_turn}\"')\n _display_single_turn_result(turn_result, passed, failed, no_assertion)\n\n # Show multi-turn summary\n status = (\n f\"{GREEN}ALL PASSED{RESET}\"\n if result.all_passed\n else f\"{RED}SOME FAILED{RESET}\"\n )\n print(\n f\"\\n {BOLD}Multi-turn result:{RESET} {status} ({result.pass_count}/{result.total_turns} turns passed)\"\n )\n\n except Exception as e:\n print(f\" {RED}ERROR:{RESET} {e}\")\n failed[0] += 1\n logger.exception(f\"Error running multi-turn eval: {first_msg}\")\n" }, { "path": "backend/onyx/feature_flags/__init__.py", "content": "" }, { "path": "backend/onyx/feature_flags/factory.py", "content": "from onyx.configs.app_configs import DEV_MODE\nfrom onyx.feature_flags.interface import FeatureFlagProvider\nfrom onyx.feature_flags.interface import NoOpFeatureFlagProvider\nfrom onyx.utils.variable_functionality import (\n fetch_versioned_implementation_with_fallback,\n)\nfrom shared_configs.configs import MULTI_TENANT\n\n\ndef get_default_feature_flag_provider() -> FeatureFlagProvider:\n \"\"\"\n Get the default feature flag provider implementation.\n\n Returns the PostHog-based provider in Enterprise Edition when available,\n otherwise returns a no-op provider that always returns False.\n\n This function is designed for dependency injection - callers should\n use this factory rather than directly instantiating providers.\n\n Returns:\n FeatureFlagProvider: The configured feature flag provider instance\n \"\"\"\n if MULTI_TENANT or DEV_MODE:\n return fetch_versioned_implementation_with_fallback(\n module=\"onyx.feature_flags.factory\",\n attribute=\"get_posthog_feature_flag_provider\",\n fallback=lambda: NoOpFeatureFlagProvider(),\n )()\n return NoOpFeatureFlagProvider()\n" }, { "path": "backend/onyx/feature_flags/feature_flags_keys.py", "content": "\"\"\"\nFeature flag keys used throughout the application.\nCentralizes feature flag key definitions to avoid magic strings.\n\"\"\"\n" }, { "path": "backend/onyx/feature_flags/flags.py", "content": "" }, { "path": "backend/onyx/feature_flags/interface.py", "content": "import abc\nfrom typing import Any\nfrom uuid import UUID\n\nfrom onyx.db.models import User\nfrom shared_configs.configs import ENVIRONMENT\n\n\nclass FeatureFlagProvider(abc.ABC):\n \"\"\"\n Abstract base class for feature flag providers.\n\n Implementations should provide vendor-specific logic for checking\n whether a feature flag is enabled for a given user.\n \"\"\"\n\n @abc.abstractmethod\n def feature_enabled(\n self,\n flag_key: str,\n user_id: UUID,\n user_properties: dict[str, Any] | None = None,\n ) -> bool:\n \"\"\"\n Check if a feature flag is enabled for a user.\n\n Args:\n flag_key: The identifier for the feature flag to check\n user_id: The unique identifier for the user\n user_properties: Optional dictionary of user properties/attributes\n that may influence flag evaluation\n\n Returns:\n True if the feature is enabled for the user, False otherwise\n \"\"\"\n raise NotImplementedError\n\n def feature_enabled_for_user_tenant(\n self, flag_key: str, user: User, tenant_id: str\n ) -> bool:\n \"\"\"\n Check if a feature flag is enabled for a user.\n \"\"\"\n return self.feature_enabled(\n flag_key,\n # For anonymous/unauthenticated users, use a fixed UUID as fallback\n user.id if user else UUID(\"caa1e0cd-6ee6-4550-b1ec-8affaef4bf83\"),\n user_properties={\n \"tenant_id\": tenant_id,\n \"email\": user.email if user else \"anonymous@onyx.app\",\n },\n )\n\n\nclass NoOpFeatureFlagProvider(FeatureFlagProvider):\n \"\"\"\n No-operation feature flag provider that always returns False.\n\n Used as a fallback when no real feature flag provider is available\n (e.g., in MIT version without PostHog).\n \"\"\"\n\n def feature_enabled(\n self,\n flag_key: str, # noqa: ARG002\n user_id: UUID, # noqa: ARG002\n user_properties: dict[str, Any] | None = None, # noqa: ARG002\n ) -> bool:\n environment = ENVIRONMENT\n if environment == \"local\":\n return True\n return False\n" }, { "path": "backend/onyx/federated_connectors/__init__.py", "content": "" }, { "path": "backend/onyx/federated_connectors/factory.py", "content": "\"\"\"Factory for creating federated connector instances.\"\"\"\n\nimport importlib\nfrom typing import Any\nfrom typing import Type\n\nfrom onyx.configs.constants import FederatedConnectorSource\nfrom onyx.federated_connectors.interfaces import FederatedConnector\nfrom onyx.federated_connectors.registry import FEDERATED_CONNECTOR_CLASS_MAP\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass FederatedConnectorMissingException(Exception):\n pass\n\n\n# Cache for already imported federated connector classes\n_federated_connector_cache: dict[FederatedConnectorSource, Type[FederatedConnector]] = (\n {}\n)\n\n\ndef _load_federated_connector_class(\n source: FederatedConnectorSource,\n) -> Type[FederatedConnector]:\n \"\"\"Dynamically load and cache a federated connector class.\"\"\"\n if source in _federated_connector_cache:\n return _federated_connector_cache[source]\n\n if source not in FEDERATED_CONNECTOR_CLASS_MAP:\n raise FederatedConnectorMissingException(\n f\"Federated connector not found for source={source}\"\n )\n\n mapping = FEDERATED_CONNECTOR_CLASS_MAP[source]\n\n try:\n module = importlib.import_module(mapping.module_path)\n connector_class = getattr(module, mapping.class_name)\n _federated_connector_cache[source] = connector_class\n return connector_class\n except (ImportError, AttributeError) as e:\n raise FederatedConnectorMissingException(\n f\"Failed to import {mapping.class_name} from {mapping.module_path}: {e}\"\n )\n\n\ndef get_federated_connector(\n source: FederatedConnectorSource,\n credentials: dict[str, Any],\n) -> FederatedConnector:\n \"\"\"Get an instance of the appropriate federated connector.\"\"\"\n connector_cls = get_federated_connector_cls(source)\n return connector_cls(credentials)\n\n\ndef get_federated_connector_cls(\n source: FederatedConnectorSource,\n) -> Type[FederatedConnector]:\n \"\"\"Get the class of the appropriate federated connector.\"\"\"\n return _load_federated_connector_class(source)\n" }, { "path": "backend/onyx/federated_connectors/federated_retrieval.py", "content": "from collections import defaultdict\nfrom collections.abc import Callable\nfrom typing import Any\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FederatedConnectorSource\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.federated import (\n get_federated_connector_document_set_mappings_by_document_set_names,\n)\nfrom onyx.db.federated import list_federated_connector_oauth_tokens\nfrom onyx.db.models import FederatedConnector__DocumentSet\nfrom onyx.db.slack_bot import fetch_slack_bots\nfrom onyx.federated_connectors.factory import get_federated_connector\nfrom onyx.federated_connectors.interfaces import FederatedConnector\nfrom onyx.onyxbot.slack.models import SlackContext\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass FederatedRetrievalInfo(BaseModel):\n model_config = ConfigDict(arbitrary_types_allowed=True)\n\n retrieval_function: Callable[[ChunkIndexRequest], list[InferenceChunk]]\n source: FederatedConnectorSource\n\n\ndef get_federated_retrieval_functions(\n db_session: Session,\n user_id: UUID | None,\n source_types: list[DocumentSource] | None,\n document_set_names: list[str] | None,\n slack_context: SlackContext | None = None,\n) -> list[FederatedRetrievalInfo]:\n\n # Check for Slack bot context first (regardless of user_id)\n if slack_context:\n logger.debug(\"Slack context detected, checking for Slack bot setup...\")\n\n # Slack federated search requires a Slack federated connector to be linked\n # via document sets. If no document sets are provided, skip Slack federated search.\n if not document_set_names:\n logger.debug(\n \"Skipping Slack federated search: no document sets provided, \"\n \"Slack federated connector must be linked via document sets\"\n )\n return []\n\n # Check if any Slack federated connector is associated with the document sets\n # and extract its config (entities) for channel filtering\n slack_federated_connector_config: dict[str, Any] | None = None\n slack_federated_mappings = (\n get_federated_connector_document_set_mappings_by_document_set_names(\n db_session, document_set_names\n )\n )\n for mapping in slack_federated_mappings:\n if (\n mapping.federated_connector is not None\n and mapping.federated_connector.source\n == FederatedConnectorSource.FEDERATED_SLACK\n ):\n slack_federated_connector_config = (\n mapping.federated_connector.config or {}\n )\n logger.debug(\n f\"Found Slack federated connector config: {slack_federated_connector_config}\"\n )\n break\n\n if slack_federated_connector_config is None:\n logger.debug(\n f\"Skipping Slack federated search: document sets {document_set_names} \"\n \"are not associated with any Slack federated connector\"\n )\n # Return empty list - no Slack federated search for this context\n return []\n\n try:\n slack_bots = fetch_slack_bots(db_session)\n logger.debug(f\"Found {len(slack_bots)} Slack bots\")\n\n # First try to find a bot with user token\n tenant_slack_bot = next(\n (bot for bot in slack_bots if bot.enabled and bot.user_token), None\n )\n if tenant_slack_bot:\n logger.debug(f\"Selected bot with user_token: {tenant_slack_bot.name}\")\n else:\n # Fall back to any enabled bot without user token\n tenant_slack_bot = next(\n (bot for bot in slack_bots if bot.enabled), None\n )\n if tenant_slack_bot:\n logger.debug(\n f\"Selected bot without user_token: {tenant_slack_bot.name} (limited functionality)\"\n )\n else:\n logger.warning(\"No enabled Slack bots found\")\n\n if tenant_slack_bot:\n federated_retrieval_infos_slack = []\n\n # Use user_token if available, otherwise fall back to bot_token\n # Unwrap SensitiveValue for backend API calls\n access_token = (\n tenant_slack_bot.user_token.get_value(apply_mask=False)\n if tenant_slack_bot.user_token\n else (\n tenant_slack_bot.bot_token.get_value(apply_mask=False)\n if tenant_slack_bot.bot_token\n else \"\"\n )\n )\n if not tenant_slack_bot.user_token:\n logger.warning(\n f\"Using bot_token for Slack search (limited functionality): {tenant_slack_bot.name}\"\n )\n\n # For bot context, we don't need real OAuth credentials\n credentials = {\n \"client_id\": \"bot-context\", # Placeholder for bot context\n \"client_secret\": \"bot-context\", # Placeholder for bot context\n }\n\n # Create Slack federated connector\n connector = get_federated_connector(\n FederatedConnectorSource.FEDERATED_SLACK,\n credentials,\n )\n\n # Capture variables by value to avoid lambda closure issues\n # Unwrap SensitiveValue for backend API calls\n bot_token = (\n tenant_slack_bot.bot_token.get_value(apply_mask=False)\n if tenant_slack_bot.bot_token\n else \"\"\n )\n\n # Use connector config for channel filtering (guaranteed to exist at this point)\n connector_entities = slack_federated_connector_config\n logger.debug(\n f\"Using Slack federated connector entities for bot context: {connector_entities}\"\n )\n\n def create_slack_retrieval_function(\n conn: FederatedConnector,\n token: str,\n ctx: SlackContext,\n bot_tok: str,\n entities: dict[str, Any],\n ) -> Callable[[ChunkIndexRequest], list[InferenceChunk]]:\n def retrieval_fn(query: ChunkIndexRequest) -> list[InferenceChunk]:\n return conn.search(\n query,\n entities, # Use connector-level entities for channel filtering\n access_token=token,\n limit=None, # Let connector use its own max_messages_per_query config\n slack_event_context=ctx,\n bot_token=bot_tok,\n )\n\n return retrieval_fn\n\n federated_retrieval_infos_slack.append(\n FederatedRetrievalInfo(\n retrieval_function=create_slack_retrieval_function(\n connector,\n access_token,\n slack_context,\n bot_token,\n connector_entities,\n ),\n source=FederatedConnectorSource.FEDERATED_SLACK,\n )\n )\n logger.debug(\n f\"Added Slack federated search for bot, returning {len(federated_retrieval_infos_slack)} retrieval functions\"\n )\n return federated_retrieval_infos_slack\n\n except Exception as e:\n logger.warning(f\"Could not setup Slack bot federated search: {e}\")\n # Fall through to regular federated connector logic\n\n if user_id is None:\n # No user ID provided and no Slack context, return empty\n logger.warning(\n \"No user ID provided and no Slack context, returning empty retrieval functions\"\n )\n return []\n\n federated_connector__document_set_pairs = (\n (\n get_federated_connector_document_set_mappings_by_document_set_names(\n db_session, document_set_names\n )\n )\n if document_set_names\n else []\n )\n federated_connector_id_to_document_sets: dict[\n int, list[FederatedConnector__DocumentSet]\n ] = defaultdict(list)\n for pair in federated_connector__document_set_pairs:\n federated_connector_id_to_document_sets[pair.federated_connector_id].append(\n pair\n )\n\n # At this point, user_id is guaranteed to be not None since we're in the else branch\n assert user_id is not None\n\n # If no source types are specified, don't use any federated connectors\n if source_types is None:\n logger.debug(\"No source types specified, skipping all federated connectors\")\n return []\n\n federated_retrieval_infos: list[FederatedRetrievalInfo] = []\n federated_oauth_tokens = list_federated_connector_oauth_tokens(db_session, user_id)\n for oauth_token in federated_oauth_tokens:\n # Slack is handled separately inside SearchTool\n if (\n oauth_token.federated_connector.source\n == FederatedConnectorSource.FEDERATED_SLACK\n ):\n logger.debug(\n \"Skipping Slack federated connector in user OAuth path - handled by SearchTool\"\n )\n continue\n\n if (\n oauth_token.federated_connector.source.to_non_federated_source()\n not in source_types\n ):\n continue\n\n document_set_associations = federated_connector_id_to_document_sets[\n oauth_token.federated_connector_id\n ]\n\n # if document set names are specified by the user, skip federated connectors that are\n # not associated with any of the document sets\n if document_set_names and not document_set_associations:\n continue\n\n # Only use connector-level config (no junction table entities)\n entities = oauth_token.federated_connector.config or {}\n\n connector = get_federated_connector(\n oauth_token.federated_connector.source,\n oauth_token.federated_connector.credentials.get_value(apply_mask=False),\n )\n\n # Capture variables by value to avoid lambda closure issues\n access_token = oauth_token.token.get_value(apply_mask=False)\n\n def create_retrieval_function(\n conn: FederatedConnector,\n ent: dict[str, Any],\n token: str,\n ) -> Callable[[ChunkIndexRequest], list[InferenceChunk]]:\n return lambda query: conn.search(\n query,\n ent,\n access_token=token,\n limit=None, # Let connector use its own max_messages_per_query config\n )\n\n federated_retrieval_infos.append(\n FederatedRetrievalInfo(\n retrieval_function=create_retrieval_function(\n connector, entities, access_token\n ),\n source=oauth_token.federated_connector.source,\n )\n )\n return federated_retrieval_infos\n" }, { "path": "backend/onyx/federated_connectors/interfaces.py", "content": "from abc import ABC\nfrom abc import abstractmethod\nfrom typing import Any\nfrom typing import Dict\n\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.federated_connectors.models import CredentialField\nfrom onyx.federated_connectors.models import EntityField\nfrom onyx.federated_connectors.models import OAuthResult\nfrom onyx.onyxbot.slack.models import SlackContext\n\n\nclass FederatedConnector(ABC):\n \"\"\"Base interface that all federated connectors must implement.\"\"\"\n\n @abstractmethod\n def __init__(self, credentials: dict[str, Any]):\n \"\"\"\n Initialize the connector with credentials + validate their structure.\n\n Args:\n credentials: Dictionary of credentials to initialize the connector with\n \"\"\"\n self.credentials = credentials\n\n @abstractmethod\n def validate_entities(self, entities: Dict[str, Any]) -> bool:\n \"\"\"\n Validate that the provided entities match the expected structure.\n\n Args:\n entities: Dictionary of entities to validate\n\n Returns:\n True if entities are valid, False otherwise\n\n Note: This method is used for backward compatibility with document-set level entities.\n For connector-level config validation, use validate_config() instead.\n \"\"\"\n\n def validate_config(self, config: Dict[str, Any]) -> bool:\n \"\"\"\n Validate that the provided config matches the expected structure.\n\n This is an alias for validate_entities() to provide clearer semantics\n when validating connector-level configuration.\n\n Args:\n config: Dictionary of configuration to validate\n\n Returns:\n True if config is valid, False otherwise\n \"\"\"\n return self.validate_entities(config)\n\n @classmethod\n @abstractmethod\n def configuration_schema(cls) -> Dict[str, EntityField]:\n \"\"\"\n Return the specification of what configuration fields are available for this connector.\n\n Returns:\n Dictionary where keys are configuration field names and values are EntityField objects\n describing the expected structure and constraints.\n \"\"\"\n\n @classmethod\n @abstractmethod\n def credentials_schema(cls) -> Dict[str, CredentialField]:\n \"\"\"\n Return the specification of what credentials are required for this connector.\n\n Returns:\n Dictionary where keys are credential field names and values are CredentialField objects\n describing the expected structure, validation rules, and security properties.\n \"\"\"\n\n @abstractmethod\n def authorize(self, redirect_uri: str) -> str:\n \"\"\"\n Generate the OAuth authorization URL.\n\n Returns:\n The URL where users should be redirected to authorize the application\n \"\"\"\n\n @abstractmethod\n def callback(self, callback_data: Dict[str, Any], redirect_uri: str) -> OAuthResult:\n \"\"\"\n Handle the OAuth callback and exchange the authorization code for tokens.\n\n Args:\n callback_data: The data received from the OAuth callback (query params, etc.)\n redirect_uri: The OAuth redirect URI used in the authorization request\n\n Returns:\n Standardized OAuthResult containing tokens and metadata\n \"\"\"\n\n @abstractmethod\n def search(\n self,\n query: ChunkIndexRequest,\n entities: dict[str, Any],\n access_token: str,\n limit: int | None = None,\n # Slack-specific parameters\n slack_event_context: SlackContext | None = None,\n bot_token: str | None = None,\n ) -> list[InferenceChunk]:\n \"\"\"\n Perform a federated search using the provided query and entities.\n\n Args:\n query: The search query\n entities: Connector-level config (entity filtering configuration)\n access_token: The OAuth access token\n limit: Maximum number of results to return\n slack_event_context: Slack-specific context (only used by Slack bot)\n bot_token: Slack bot token (only used by Slack bot)\n\n Returns:\n Search results in a standardized format\n \"\"\"\n" }, { "path": "backend/onyx/federated_connectors/models.py", "content": "from datetime import datetime\nfrom typing import Any\nfrom typing import Dict\nfrom typing import Optional\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\n\nclass FieldSpec(BaseModel):\n \"\"\"Model for describing a field specification.\"\"\"\n\n type: str = Field(\n ..., description=\"The type of the field (e.g., 'str', 'bool', 'list[str]')\"\n )\n description: str = Field(\n ..., description=\"Description of what this field represents\"\n )\n required: bool = Field(default=False, description=\"Whether this field is required\")\n default: Optional[Any] = Field(\n default=None, description=\"Default value if not provided\"\n )\n example: Optional[Any] = Field(\n default=None, description=\"Example value for documentation\"\n )\n secret: bool = Field(\n default=False, description=\"Whether this field contains sensitive data\"\n )\n\n\nclass EntityField(FieldSpec):\n \"\"\"Model for describing an entity field in the entities specification.\"\"\"\n\n\nclass CredentialField(FieldSpec):\n \"\"\"Model for describing a credential field in the credentials specification.\"\"\"\n\n\nclass OAuthResult(BaseModel):\n \"\"\"Standardized OAuth result that all federated connectors should return from callback.\"\"\"\n\n access_token: Optional[str] = Field(\n default=None, description=\"The bot access token for bot operations\"\n )\n user_token: Optional[str] = Field(\n default=None,\n description=\"The user access token for user-scoped operations like federated search\",\n )\n token_type: Optional[str] = Field(\n default=None, description=\"Token type (usually 'bearer')\"\n )\n scope: Optional[str] = Field(default=None, description=\"Granted scopes\")\n expires_at: Optional[datetime] = Field(\n default=None, description=\"When the token expires\"\n )\n refresh_token: Optional[str] = Field(\n default=None, description=\"Refresh token if applicable\"\n )\n\n # Additional fields that might be useful\n team: Optional[Dict[str, Any]] = Field(\n default=None, description=\"Team/workspace information\"\n )\n user: Optional[Dict[str, Any]] = Field(default=None, description=\"User information\")\n raw_response: Optional[Dict[str, Any]] = Field(\n default=None, description=\"Raw response for debugging\"\n )\n\n # Pydantic V2 automatically serializes datetime to ISO format, so no custom encoder needed\n" }, { "path": "backend/onyx/federated_connectors/oauth_utils.py", "content": "\"\"\"Generic OAuth utilities for federated connectors API layer.\"\"\"\n\nimport base64\nimport json\nimport uuid\nfrom typing import Any\n\nfrom onyx.cache.factory import get_cache_backend\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nOAUTH_STATE_PREFIX = \"federated_oauth\"\nOAUTH_STATE_TTL = 300 # 5 minutes\n\n\nclass OAuthSession:\n \"\"\"Represents an OAuth session stored in the cache backend.\"\"\"\n\n def __init__(\n self,\n federated_connector_id: int,\n user_id: str,\n redirect_uri: str | None = None,\n additional_data: dict[str, Any] | None = None,\n ):\n self.federated_connector_id = federated_connector_id\n self.user_id = user_id\n self.redirect_uri = redirect_uri\n self.additional_data = additional_data or {}\n\n def to_dict(self) -> dict[str, Any]:\n return {\n \"federated_connector_id\": self.federated_connector_id,\n \"user_id\": self.user_id,\n \"redirect_uri\": self.redirect_uri,\n \"additional_data\": self.additional_data,\n }\n\n @classmethod\n def from_dict(cls, data: dict[str, Any]) -> \"OAuthSession\":\n return cls(\n federated_connector_id=data[\"federated_connector_id\"],\n user_id=data[\"user_id\"],\n redirect_uri=data.get(\"redirect_uri\"),\n additional_data=data.get(\"additional_data\", {}),\n )\n\n\ndef generate_oauth_state(\n federated_connector_id: int,\n user_id: str,\n redirect_uri: str | None = None,\n additional_data: dict[str, Any] | None = None,\n ttl: int = OAUTH_STATE_TTL,\n) -> str:\n \"\"\"\n Generate a secure state parameter and store session data in the cache backend.\n\n Args:\n federated_connector_id: ID of the federated connector\n user_id: ID of the user initiating OAuth\n redirect_uri: Optional redirect URI after OAuth completion\n additional_data: Any additional data to store with the session\n ttl: Time-to-live in seconds for the cache key\n\n Returns:\n Base64-encoded state parameter\n \"\"\"\n # Generate a random UUID for the state\n state_uuid = uuid.uuid4()\n state_b64 = base64.urlsafe_b64encode(state_uuid.bytes).decode(\"utf-8\").rstrip(\"=\")\n\n session = OAuthSession(\n federated_connector_id=federated_connector_id,\n user_id=user_id,\n redirect_uri=redirect_uri,\n additional_data=additional_data,\n )\n\n cache = get_cache_backend()\n cache_key = f\"{OAUTH_STATE_PREFIX}:{state_uuid}\"\n cache.set(cache_key, json.dumps(session.to_dict()), ex=ttl)\n\n logger.info(\n f\"Generated OAuth state for federated_connector_id={federated_connector_id}, user_id={user_id}, state={state_b64}\"\n )\n\n return state_b64\n\n\ndef verify_oauth_state(state: str) -> OAuthSession:\n \"\"\"\n Verify OAuth state parameter and retrieve session data.\n\n Args:\n state: Base64-encoded state parameter from OAuth callback\n\n Returns:\n OAuthSession if state is valid, None otherwise\n \"\"\"\n # Add padding if needed for base64 decoding\n padded_state = state + \"=\" * (-len(state) % 4)\n\n # Decode base64 to get UUID bytes\n state_bytes = base64.urlsafe_b64decode(padded_state)\n state_uuid = uuid.UUID(bytes=state_bytes)\n\n cache = get_cache_backend()\n cache_key = f\"{OAUTH_STATE_PREFIX}:{state_uuid}\"\n\n session_data = cache.get(cache_key)\n if not session_data:\n raise ValueError(f\"OAuth state not found: {state}\")\n\n cache.delete(cache_key)\n\n session_dict = json.loads(session_data)\n return OAuthSession.from_dict(session_dict)\n\n\ndef get_oauth_callback_uri() -> str:\n \"\"\"\n Generate the OAuth callback URI for a federated connector.\n\n Returns:\n The callback URI\n \"\"\"\n # Use the frontend callback page as the OAuth redirect URI\n # The frontend will then make an API call to process the callback\n return f\"{WEB_DOMAIN}/federated/oauth/callback\"\n\n\ndef add_state_to_oauth_url(base_oauth_url: str, state: str) -> str:\n \"\"\"\n Add state parameter to an OAuth URL.\n\n Args:\n base_oauth_url: The base OAuth URL from the connector\n state: The state parameter to add\n\n Returns:\n The OAuth URL with state parameter added\n \"\"\"\n # Check if URL already has query parameters\n separator = \"&\" if \"?\" in base_oauth_url else \"?\"\n return f\"{base_oauth_url}{separator}state={state}\"\n" }, { "path": "backend/onyx/federated_connectors/registry.py", "content": "\"\"\"Registry mapping for federated connector classes.\"\"\"\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import FederatedConnectorSource\n\n\nclass FederatedConnectorMapping(BaseModel):\n module_path: str\n class_name: str\n\n\n# Mapping of FederatedConnectorSource to connector details for lazy loading\nFEDERATED_CONNECTOR_CLASS_MAP = {\n FederatedConnectorSource.FEDERATED_SLACK: FederatedConnectorMapping(\n module_path=\"onyx.federated_connectors.slack.federated_connector\",\n class_name=\"SlackFederatedConnector\",\n ),\n}\n" }, { "path": "backend/onyx/federated_connectors/slack/__init__.py", "content": "# Slack federated connector module\n" }, { "path": "backend/onyx/federated_connectors/slack/federated_connector.py", "content": "from datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Any\nfrom urllib.parse import urlencode\n\nimport requests\nfrom pydantic import ValidationError\nfrom slack_sdk import WebClient\nfrom typing_extensions import override\n\nfrom onyx.context.search.federated.slack_search import slack_retrieval\nfrom onyx.context.search.models import ChunkIndexRequest\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.federated_connectors.interfaces import FederatedConnector\nfrom onyx.federated_connectors.models import CredentialField\nfrom onyx.federated_connectors.models import EntityField\nfrom onyx.federated_connectors.models import OAuthResult\nfrom onyx.federated_connectors.slack.models import SlackCredentials\nfrom onyx.federated_connectors.slack.models import SlackEntities\nfrom onyx.onyxbot.slack.models import SlackContext\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nSCOPES = [\n \"channels:read\",\n \"groups:read\",\n \"im:read\",\n \"mpim:read\",\n \"search:read\",\n \"channels:history\",\n \"groups:history\",\n \"im:history\",\n \"mpim:history\",\n \"users:read\",\n \"users.profile:read\",\n]\n\n\nclass SlackFederatedConnector(FederatedConnector):\n def __init__(self, credentials: dict[str, Any]):\n self.slack_credentials = SlackCredentials(**credentials)\n\n @override\n def validate_entities(self, entities: dict[str, Any]) -> bool:\n \"\"\"Check the entities and verify that they match the expected structure/all values are valid.\n\n For Slack federated search, we expect:\n - channels: list[str] (list of channel names or IDs)\n - include_dm: bool (whether to include direct messages)\n \"\"\"\n try:\n # Use Pydantic model for validation\n SlackEntities(**entities)\n return True\n except ValidationError as e:\n logger.warning(f\"Validation error for Slack entities: {e}\")\n return False\n except Exception as e:\n logger.error(f\"Error validating Slack entities: {e}\")\n return False\n\n @classmethod\n def entities_schema(cls) -> dict[str, EntityField]:\n \"\"\"Return the specifications of what entity configuration fields are available for Slack.\n\n This is the canonical schema definition for Slack entities.\n \"\"\"\n return {\n \"exclude_channels\": EntityField(\n type=\"list[str]\",\n description=\"Exclude the following channels from search. Glob patterns are supported.\",\n required=False,\n example=[\"secure-channel\", \"private-*\", \"customer*\"],\n ),\n \"search_all_channels\": EntityField(\n type=\"bool\",\n description=\"Search all accessible channels. If not set, must specify channels below.\",\n required=False,\n default=False,\n example=False,\n ),\n \"channels\": EntityField(\n type=\"list[str]\",\n description=\"Search the following channels\",\n required=False,\n example=[\"general\", \"eng*\", \"product-*\"],\n ),\n \"include_dm\": EntityField(\n type=\"bool\",\n description=\"Include user direct messages in search results\",\n required=False,\n default=False,\n example=False,\n ),\n \"include_group_dm\": EntityField(\n type=\"bool\",\n description=\"Include group direct messages (multi-person DMs) in search results\",\n required=False,\n default=False,\n example=False,\n ),\n \"include_private_channels\": EntityField(\n type=\"bool\",\n description=\"Include private channels in search results (user must have access)\",\n required=False,\n default=False,\n example=False,\n ),\n \"default_search_days\": EntityField(\n type=\"int\",\n description=\"Maximum number of days to search back. Increasing this value degrades answer quality.\",\n required=False,\n default=30,\n example=30,\n ),\n \"max_messages_per_query\": EntityField(\n type=\"int\",\n description=(\n \"Maximum number of messages to retrieve per search query. \"\n \"Higher values provide more context but may be slower.\"\n ),\n required=False,\n default=25,\n example=25,\n ),\n }\n\n @classmethod\n def configuration_schema(cls) -> dict[str, EntityField]:\n \"\"\"Wrapper for backwards compatibility - delegates to entities_schema().\"\"\"\n return cls.entities_schema()\n\n @classmethod\n @override\n def credentials_schema(cls) -> dict[str, CredentialField]:\n \"\"\"Return the specification of what credentials are required for Slack connector.\"\"\"\n return {\n \"client_id\": CredentialField(\n type=\"str\",\n description=\"Slack app client ID from your Slack app configuration\",\n required=True,\n example=\"1234567890.1234567890123\",\n secret=False,\n ),\n \"client_secret\": CredentialField(\n type=\"str\",\n description=\"Slack app client secret from your Slack app configuration\",\n required=True,\n example=\"1a2b3c4d5e6f7g8h9i0j1k2l3m4n5o6p\",\n secret=True,\n ),\n }\n\n @override\n def authorize(self, redirect_uri: str) -> str:\n \"\"\"Get back the OAuth URL for Slack authorization.\n\n Returns the URL where users should be redirected to authorize the application.\n Note: State parameter will be added by the API layer.\n \"\"\"\n # Build OAuth URL with proper parameters (no state - handled by API layer)\n params = {\n \"client_id\": self.slack_credentials.client_id,\n \"user_scope\": \" \".join(SCOPES),\n \"redirect_uri\": redirect_uri,\n }\n\n # Build query string\n oauth_url = f\"https://slack.com/oauth/v2/authorize?{urlencode(params)}\"\n\n logger.info(\"Generated Slack OAuth authorization URL\")\n return oauth_url\n\n @override\n def callback(self, callback_data: dict[str, Any], redirect_uri: str) -> OAuthResult:\n \"\"\"Handle the response from the OAuth flow and return it in a standard format.\n\n Args:\n callback_data: The data received from the OAuth callback (state already validated by API layer)\n\n Returns:\n Standardized OAuthResult\n \"\"\"\n # Extract authorization code from callback\n auth_code = callback_data.get(\"code\")\n error = callback_data.get(\"error\")\n\n if error:\n raise RuntimeError(f\"OAuth error received: {error}\")\n\n if not auth_code:\n raise ValueError(\"No authorization code received\")\n\n # Exchange authorization code for access token\n token_response = self._exchange_code_for_token(auth_code, redirect_uri)\n\n if not token_response.get(\"ok\"):\n raise RuntimeError(\n f\"Failed to exchange authorization code for token: {token_response.get('error')}\"\n )\n\n # Build team info\n team_info = None\n if \"team\" in token_response:\n team_info = {\n \"id\": token_response[\"team\"][\"id\"],\n \"name\": token_response[\"team\"][\"name\"],\n }\n\n # Build user info and extract OAuth tokens\n if \"authed_user\" not in token_response:\n raise RuntimeError(\"Missing authed_user in OAuth response from Slack\")\n\n authed_user = token_response[\"authed_user\"]\n user_info = {\n \"id\": authed_user[\"id\"],\n \"scope\": authed_user.get(\"scope\"),\n \"token_type\": authed_user.get(\"token_type\"),\n }\n\n # Extract OAuth tokens - bot token from root, user token from authed_user\n user_token = authed_user.get(\"access_token\") # User token\n refresh_token = authed_user.get(\"refresh_token\")\n token_type = authed_user.get(\"token_type\", \"bearer\")\n scope = authed_user.get(\"scope\")\n\n # Calculate expires_at from expires_in if present\n expires_at = None\n if \"expires_in\" in authed_user:\n expires_at = datetime.now(timezone.utc) + timedelta(\n seconds=authed_user[\"expires_in\"]\n )\n\n return OAuthResult(\n access_token=user_token, # Bot token for bot operations\n token_type=token_type,\n scope=scope,\n expires_at=expires_at,\n refresh_token=refresh_token,\n team=team_info,\n user=user_info,\n raw_response=token_response,\n )\n\n def _exchange_code_for_token(self, code: str, redirect_uri: str) -> dict[str, Any]:\n \"\"\"Exchange authorization code for access token.\n\n Args:\n code: Authorization code from OAuth callback\n\n Returns:\n Token response from Slack API\n \"\"\"\n response = requests.post(\n \"https://slack.com/api/oauth.v2.access\",\n data={\n \"client_id\": self.slack_credentials.client_id,\n \"client_secret\": self.slack_credentials.client_secret,\n \"code\": code,\n \"redirect_uri\": redirect_uri,\n },\n )\n response.raise_for_status()\n return response.json()\n\n @override\n def search(\n self,\n query: ChunkIndexRequest,\n entities: dict[str, Any],\n access_token: str,\n limit: int | None = None,\n slack_event_context: SlackContext | None = None,\n bot_token: str | None = None,\n ) -> list[InferenceChunk]:\n \"\"\"Perform a federated search on Slack.\n\n Args:\n query: The search query\n entities: Connector-level config (entity filtering configuration)\n access_token: The OAuth access token\n limit: Maximum number of results to return\n slack_event_context: Optional Slack context for slack bot\n bot_token: Optional bot token for slack bot\n\n Returns:\n Search results in SlackSearchResponse format\n \"\"\"\n logger.debug(f\"Slack federated search called with entities: {entities}\")\n\n # Get team_id from Slack API for caching and filtering\n team_id = None\n try:\n slack_client = WebClient(token=access_token)\n auth_response = slack_client.auth_test()\n auth_response.validate()\n\n # Cast response.data to dict for type checking\n auth_data: dict[str, Any] = auth_response.data # type: ignore\n team_id = auth_data.get(\"team_id\")\n logger.debug(f\"Slack team_id: {team_id}\")\n except Exception as e:\n logger.warning(f\"Could not fetch team_id from Slack API: {e}\")\n\n with get_session_with_current_tenant() as db_session:\n return slack_retrieval(\n query,\n access_token,\n db_session,\n entities=entities,\n limit=limit,\n slack_event_context=slack_event_context,\n bot_token=bot_token,\n team_id=team_id,\n )\n" }, { "path": "backend/onyx/federated_connectors/slack/models.py", "content": "from typing import Optional\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom pydantic import field_validator\nfrom pydantic import model_validator\n\n\nclass SlackEntities(BaseModel):\n \"\"\"Pydantic model for Slack federated search entities.\"\"\"\n\n # Channel filtering\n search_all_channels: bool = Field(\n default=True,\n description=\"Search all accessible channels. If not set, must specify channels below.\",\n )\n channels: Optional[list[str]] = Field(\n default=None,\n description=\"List of Slack channel names to search across.\",\n )\n exclude_channels: Optional[list[str]] = Field(\n default=None,\n description=\"List of channel names or patterns to exclude e.g. 'private-*, customer-*, secure-channel'.\",\n )\n\n # Direct message filtering\n include_dm: bool = Field(\n default=True,\n description=\"Include user direct messages in search results\",\n )\n include_group_dm: bool = Field(\n default=True,\n description=\"Include group direct messages (multi-person DMs) in search results\",\n )\n\n # Private channel filtering\n include_private_channels: bool = Field(\n default=True,\n description=\"Include private channels in search results (user must have access)\",\n )\n\n # Date range filtering\n default_search_days: int = Field(\n default=30,\n description=\"Maximum number of days to search back. Increasing this value degrades answer quality.\",\n )\n\n # Message count per slack request\n max_messages_per_query: int = Field(\n default=10,\n description=(\n \"Maximum number of messages to retrieve per search query. \"\n \"Higher values increase API calls and may trigger rate limits.\"\n ),\n )\n\n @field_validator(\"default_search_days\")\n @classmethod\n def validate_default_search_days(cls, v: int) -> int:\n \"\"\"Validate default_search_days is positive and reasonable\"\"\"\n if v < 1:\n raise ValueError(\"default_search_days must be at least 1\")\n if v > 365:\n raise ValueError(\"default_search_days cannot exceed 365 days\")\n return v\n\n @field_validator(\"max_messages_per_query\")\n @classmethod\n def validate_max_messages_per_query(cls, v: int) -> int:\n \"\"\"Validate max_messages_per_query is positive and reasonable\"\"\"\n if v < 1:\n raise ValueError(\"max_messages_per_query must be at least 1\")\n if v > 100:\n raise ValueError(\"max_messages_per_query cannot exceed 100\")\n return v\n\n @field_validator(\"channels\")\n @classmethod\n def validate_channels(cls, v: Optional[list[str]]) -> Optional[list[str]]:\n \"\"\"Validate each channel is a non-empty string\"\"\"\n if v is not None:\n if not isinstance(v, list):\n raise ValueError(\"channels must be a list\")\n for channel in v:\n if not isinstance(channel, str) or not channel.strip():\n raise ValueError(\"Each channel must be a non-empty string\")\n return v\n\n @field_validator(\"exclude_channels\")\n @classmethod\n def validate_exclude_patterns(cls, v: Optional[list[str]]) -> Optional[list[str]]:\n \"\"\"Validate each exclude pattern is a non-empty string\"\"\"\n if v is None:\n return v\n\n for pattern in v:\n if not isinstance(pattern, str) or not pattern.strip():\n raise ValueError(\"Each exclude pattern must be a non-empty string\")\n\n return v\n\n @model_validator(mode=\"after\")\n def validate_channel_config(self) -> \"SlackEntities\":\n \"\"\"Validate search_all_channels configuration\"\"\"\n # If search_all_channels is False, channels list must be provided\n if not self.search_all_channels:\n if self.channels is None or len(self.channels) == 0:\n raise ValueError(\n \"Must specify at least one channel when search_all_channels is False\"\n )\n\n return self\n\n\nclass SlackCredentials(BaseModel):\n \"\"\"Slack federated connector credentials.\"\"\"\n\n client_id: str = Field(..., description=\"Slack app client ID\")\n client_secret: str = Field(..., description=\"Slack app client secret\")\n\n @field_validator(\"client_id\")\n @classmethod\n def validate_client_id(cls, v: str) -> str:\n if not v or not v.strip():\n raise ValueError(\"Client ID cannot be empty\")\n return v.strip()\n\n @field_validator(\"client_secret\")\n @classmethod\n def validate_client_secret(cls, v: str) -> str:\n if not v or not v.strip():\n raise ValueError(\"Client secret cannot be empty\")\n return v.strip()\n\n\nclass SlackTeamInfo(BaseModel):\n \"\"\"Information about a Slack team/workspace.\"\"\"\n\n id: str = Field(..., description=\"Team ID\")\n name: str = Field(..., description=\"Team name\")\n domain: Optional[str] = Field(default=None, description=\"Team domain\")\n\n\nclass SlackUserInfo(BaseModel):\n \"\"\"Information about a Slack user.\"\"\"\n\n id: str = Field(..., description=\"User ID\")\n team_id: Optional[str] = Field(default=None, description=\"Team ID\")\n name: Optional[str] = Field(default=None, description=\"User name\")\n email: Optional[str] = Field(default=None, description=\"User email\")\n\n\nclass SlackSearchResult(BaseModel):\n \"\"\"Individual search result from Slack.\"\"\"\n\n channel: str = Field(..., description=\"Channel where the message was found\")\n timestamp: str = Field(..., description=\"Message timestamp\")\n user: Optional[str] = Field(default=None, description=\"User who sent the message\")\n text: str = Field(..., description=\"Message text\")\n permalink: Optional[str] = Field(\n default=None, description=\"Permalink to the message\"\n )\n score: Optional[float] = Field(default=None, description=\"Search relevance score\")\n\n # Additional context\n thread_ts: Optional[str] = Field(\n default=None, description=\"Thread timestamp if in a thread\"\n )\n reply_count: Optional[int] = Field(\n default=None, description=\"Number of replies if it's a thread\"\n )\n\n\nclass SlackSearchResponse(BaseModel):\n \"\"\"Response from Slack federated search.\"\"\"\n\n query: str = Field(..., description=\"The search query\")\n total_count: int = Field(..., description=\"Total number of results\")\n results: list[SlackSearchResult] = Field(..., description=\"Search results\")\n next_cursor: Optional[str] = Field(\n default=None, description=\"Cursor for pagination\"\n )\n\n # Metadata\n channels_searched: Optional[list[str]] = Field(\n default=None, description=\"Channels that were searched\"\n )\n search_time_ms: Optional[int] = Field(\n default=None, description=\"Time taken to search in milliseconds\"\n )\n" }, { "path": "backend/onyx/file_processing/__init__.py", "content": "" }, { "path": "backend/onyx/file_processing/enums.py", "content": "from enum import Enum\n\n\nclass HtmlBasedConnectorTransformLinksStrategy(str, Enum):\n # remove links entirely\n STRIP = \"strip\"\n # turn HTML links into markdown links\n MARKDOWN = \"markdown\"\n" }, { "path": "backend/onyx/file_processing/extract_file_text.py", "content": "import csv\nimport gc\nimport io\nimport json\nimport os\nimport re\nimport zipfile\nfrom collections.abc import Callable\nfrom collections.abc import Iterator\nfrom collections.abc import Sequence\nfrom email.parser import Parser as EmailParser\nfrom io import BytesIO\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import IO\nfrom typing import NamedTuple\nfrom typing import Optional\nfrom typing import TYPE_CHECKING\nfrom zipfile import BadZipFile\n\nimport chardet\nimport openpyxl\nfrom openpyxl.worksheet.worksheet import Worksheet\nfrom PIL import Image\n\nfrom onyx.configs.constants import ONYX_METADATA_FILENAME\nfrom onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled\nfrom onyx.file_processing.file_types import OnyxFileExtensions\nfrom onyx.file_processing.file_types import OnyxMimeTypes\nfrom onyx.file_processing.file_types import PRESENTATION_MIME_TYPE\nfrom onyx.file_processing.file_types import WORD_PROCESSING_MIME_TYPE\nfrom onyx.file_processing.html_utils import parse_html_page_basic\nfrom onyx.file_processing.unstructured import get_unstructured_api_key\nfrom onyx.file_processing.unstructured import unstructured_to_text\nfrom onyx.utils.logger import setup_logger\n\nif TYPE_CHECKING:\n from markitdown import MarkItDown\nlogger = setup_logger()\n\nTEXT_SECTION_SEPARATOR = \"\\n\\n\"\n\n_MARKITDOWN_CONVERTER: Optional[\"MarkItDown\"] = None\n\nKNOWN_OPENPYXL_BUGS = [\n \"Value must be either numerical or a string containing a wildcard\",\n \"File contains no valid workbook part\",\n \"Unable to read workbook: could not read stylesheet from None\",\n \"Colors must be aRGB hex values\",\n]\n\n\ndef get_markitdown_converter() -> \"MarkItDown\":\n global _MARKITDOWN_CONVERTER\n from markitdown import MarkItDown\n\n if _MARKITDOWN_CONVERTER is None:\n _MARKITDOWN_CONVERTER = MarkItDown(enable_plugins=False)\n return _MARKITDOWN_CONVERTER\n\n\ndef get_file_ext(file_path_or_name: str | Path) -> str:\n _, extension = os.path.splitext(file_path_or_name)\n return extension.lower()\n\n\ndef is_text_file(file: IO[bytes]) -> bool:\n \"\"\"\n checks if the first 1024 bytes only contain printable or whitespace characters\n if it does, then we say it's a plaintext file\n \"\"\"\n raw_data = file.read(1024)\n file.seek(0)\n text_chars = bytearray({7, 8, 9, 10, 12, 13, 27} | set(range(0x20, 0x100)) - {0x7F})\n return all(c in text_chars for c in raw_data)\n\n\ndef detect_encoding(file: IO[bytes]) -> str:\n raw_data = file.read(50000)\n file.seek(0)\n encoding = chardet.detect(raw_data)[\"encoding\"] or \"utf-8\"\n return encoding\n\n\ndef is_macos_resource_fork_file(file_name: str) -> bool:\n return os.path.basename(file_name).startswith(\"._\") and file_name.startswith(\n \"__MACOSX\"\n )\n\n\ndef to_bytesio(stream: IO[bytes]) -> BytesIO:\n if isinstance(stream, BytesIO):\n return stream\n data = stream.read() # consumes the stream!\n return BytesIO(data)\n\n\ndef load_files_from_zip(\n zip_file_io: IO,\n ignore_macos_resource_fork_files: bool = True,\n ignore_dirs: bool = True,\n) -> Iterator[tuple[zipfile.ZipInfo, IO[Any]]]:\n \"\"\"\n Iterates through files in a zip archive, yielding (ZipInfo, file handle) pairs.\n \"\"\"\n with zipfile.ZipFile(zip_file_io, \"r\") as zip_file:\n for file_info in zip_file.infolist():\n if ignore_dirs and file_info.is_dir():\n continue\n\n if (\n ignore_macos_resource_fork_files\n and is_macos_resource_fork_file(file_info.filename)\n ) or file_info.filename == ONYX_METADATA_FILENAME:\n continue\n\n with zip_file.open(file_info.filename, \"r\") as subfile:\n # Try to match by exact filename first\n yield file_info, subfile\n\n\ndef _extract_onyx_metadata(line: str) -> dict | None:\n \"\"\"\n Example: first line has:\n \n or\n #ONYX_METADATA={\"title\":\"...\"}\n \"\"\"\n html_comment_pattern = r\"\"\n hashtag_pattern = r\"#ONYX_METADATA=\\{(.*?)\\}\"\n\n html_comment_match = re.search(html_comment_pattern, line)\n hashtag_match = re.search(hashtag_pattern, line)\n\n if html_comment_match:\n json_str = html_comment_match.group(1)\n elif hashtag_match:\n json_str = hashtag_match.group(1)\n else:\n return None\n\n try:\n return json.loads(\"{\" + json_str + \"}\")\n except json.JSONDecodeError:\n return None\n\n\ndef read_text_file(\n file: IO,\n encoding: str = \"utf-8\",\n errors: str = \"replace\",\n ignore_onyx_metadata: bool = True,\n) -> tuple[str, dict]:\n \"\"\"\n For plain text files. Optionally extracts Onyx metadata from the first line.\n \"\"\"\n metadata = {}\n file_content_raw = \"\"\n for ind, line in enumerate(file):\n # decode\n try:\n line = line.decode(encoding) if isinstance(line, bytes) else line\n except UnicodeDecodeError:\n line = (\n line.decode(encoding, errors=errors)\n if isinstance(line, bytes)\n else line\n )\n\n # optionally parse metadata in the first line\n if ind == 0 and not ignore_onyx_metadata:\n potential_meta = _extract_onyx_metadata(line)\n if potential_meta is not None:\n metadata = potential_meta\n continue\n\n file_content_raw += line\n\n return file_content_raw, metadata\n\n\ndef pdf_to_text(file: IO[Any], pdf_pass: str | None = None) -> str:\n \"\"\"\n Extract text from a PDF. For embedded images, a more complex approach is needed.\n This is a minimal approach returning text only.\n \"\"\"\n text, _, _ = read_pdf_file(file, pdf_pass)\n return text\n\n\ndef read_pdf_file(\n file: IO[Any],\n pdf_pass: str | None = None,\n extract_images: bool = False,\n image_callback: Callable[[bytes, str], None] | None = None,\n) -> tuple[str, dict[str, Any], Sequence[tuple[bytes, str]]]:\n \"\"\"\n Returns the text, basic PDF metadata, and optionally extracted images.\n \"\"\"\n from pypdf import PdfReader\n from pypdf.errors import PdfStreamError\n\n metadata: dict[str, Any] = {}\n extracted_images: list[tuple[bytes, str]] = []\n try:\n pdf_reader = PdfReader(file)\n\n if pdf_reader.is_encrypted and pdf_pass is not None:\n decrypt_success = False\n try:\n decrypt_success = pdf_reader.decrypt(pdf_pass) != 0\n except Exception:\n logger.error(\"Unable to decrypt pdf\")\n\n if not decrypt_success:\n return \"\", metadata, []\n elif pdf_reader.is_encrypted:\n logger.warning(\"No Password for an encrypted PDF, returning empty text.\")\n return \"\", metadata, []\n\n # Basic PDF metadata\n if pdf_reader.metadata is not None:\n for key, value in pdf_reader.metadata.items():\n clean_key = key.lstrip(\"/\")\n if isinstance(value, str) and value.strip():\n metadata[clean_key] = value\n elif isinstance(value, list) and all(\n isinstance(item, str) for item in value\n ):\n metadata[clean_key] = \", \".join(value)\n\n text = TEXT_SECTION_SEPARATOR.join(\n page.extract_text() for page in pdf_reader.pages\n )\n\n if extract_images:\n for page_num, page in enumerate(pdf_reader.pages):\n for image_file_object in page.images:\n image = Image.open(io.BytesIO(image_file_object.data))\n img_byte_arr = io.BytesIO()\n image.save(img_byte_arr, format=image.format)\n img_bytes = img_byte_arr.getvalue()\n\n image_format = image.format.lower() if image.format else \"png\"\n image_name = f\"page_{page_num + 1}_image_{image_file_object.name}.{image_format}\"\n if image_callback is not None:\n # Stream image out immediately\n image_callback(img_bytes, image_name)\n else:\n extracted_images.append((img_bytes, image_name))\n\n return text, metadata, extracted_images\n\n except PdfStreamError:\n logger.exception(\"Invalid PDF file\")\n except Exception:\n logger.exception(\"Failed to read PDF\")\n\n return \"\", metadata, []\n\n\ndef extract_docx_images(docx_bytes: IO[Any]) -> Iterator[tuple[bytes, str]]:\n \"\"\"\n Given the bytes of a docx file, extract all the images.\n Returns a list of tuples (image_bytes, image_name).\n \"\"\"\n try:\n with zipfile.ZipFile(docx_bytes) as z:\n for name in z.namelist():\n if name.startswith(\"word/media/\"):\n yield (z.read(name), name.split(\"/\")[-1])\n except Exception:\n logger.exception(\"Failed to extract all docx images\")\n\n\ndef read_docx_file(\n file: IO[Any],\n file_name: str = \"\",\n extract_images: bool = False,\n image_callback: Callable[[bytes, str], None] | None = None,\n) -> tuple[str, Sequence[tuple[bytes, str]]]:\n \"\"\"\n Extract text from a docx.\n Return (text_content, list_of_images).\n\n The caller can choose to provide a callback to handle images with the intent\n of avoiding materializing the list of images in memory.\n The images list returned is empty in this case.\n \"\"\"\n md = get_markitdown_converter()\n from markitdown import (\n StreamInfo,\n FileConversionException,\n UnsupportedFormatException,\n )\n\n try:\n doc = md.convert(\n to_bytesio(file), stream_info=StreamInfo(mimetype=WORD_PROCESSING_MIME_TYPE)\n )\n except (\n BadZipFile,\n ValueError,\n FileConversionException,\n UnsupportedFormatException,\n ) as e:\n logger.warning(\n f\"Failed to extract docx {file_name or 'docx file'}: {e}. Attempting to read as text file.\"\n )\n\n # May be an invalid docx, but still a valid text file\n file.seek(0)\n encoding = detect_encoding(file)\n text_content_raw, _ = read_text_file(\n file, encoding=encoding, ignore_onyx_metadata=False\n )\n return text_content_raw or \"\", []\n\n file.seek(0)\n\n if extract_images:\n if image_callback is None:\n return doc.markdown, list(extract_docx_images(to_bytesio(file)))\n # If a callback is provided, iterate and stream images without accumulating\n try:\n for img_file_bytes, img_file_name in extract_docx_images(to_bytesio(file)):\n image_callback(img_file_bytes, img_file_name)\n except Exception:\n logger.exception(\"Failed to stream docx images\")\n return doc.markdown, []\n\n\ndef pptx_to_text(file: IO[Any], file_name: str = \"\") -> str:\n md = get_markitdown_converter()\n from markitdown import (\n StreamInfo,\n FileConversionException,\n UnsupportedFormatException,\n )\n\n stream_info = StreamInfo(\n mimetype=PRESENTATION_MIME_TYPE, filename=file_name or None, extension=\".pptx\"\n )\n try:\n presentation = md.convert(to_bytesio(file), stream_info=stream_info)\n except (\n BadZipFile,\n ValueError,\n FileConversionException,\n UnsupportedFormatException,\n ) as e:\n error_str = f\"Failed to extract text from {file_name or 'pptx file'}: {e}\"\n logger.warning(error_str)\n return \"\"\n return presentation.markdown\n\n\ndef _worksheet_to_matrix(\n worksheet: Worksheet,\n) -> list[list[str]]:\n \"\"\"\n Converts a singular worksheet to a matrix of values\n \"\"\"\n rows: list[list[str]] = []\n for worksheet_row in worksheet.iter_rows(min_row=1, values_only=True):\n row = [\"\" if cell is None else str(cell) for cell in worksheet_row]\n rows.append(row)\n\n return rows\n\n\ndef _clean_worksheet_matrix(matrix: list[list[str]]) -> list[list[str]]:\n \"\"\"\n Cleans a worksheet matrix by removing rows if there are N consecutive empty\n rows and removing cols if there are M consecutive empty columns\n \"\"\"\n MAX_EMPTY_ROWS = 2 # Runs longer than this are capped to max_empty; shorter runs are preserved as-is\n MAX_EMPTY_COLS = 2\n\n # Row cleanup\n matrix = _remove_empty_runs(matrix, max_empty=MAX_EMPTY_ROWS)\n\n if not matrix:\n return matrix\n\n # Column cleanup — determine which columns to keep without transposing.\n num_cols = len(matrix[0])\n keep_cols = _columns_to_keep(matrix, num_cols, max_empty=MAX_EMPTY_COLS)\n if len(keep_cols) < num_cols:\n matrix = [[row[c] for c in keep_cols] for row in matrix]\n\n return matrix\n\n\ndef _columns_to_keep(\n matrix: list[list[str]], num_cols: int, max_empty: int\n) -> list[int]:\n \"\"\"Return the indices of columns to keep after removing empty-column runs.\n\n Uses the same logic as ``_remove_empty_runs`` but operates on column\n indices so no transpose is needed.\n \"\"\"\n kept: list[int] = []\n empty_buffer: list[int] = []\n\n for col_idx in range(num_cols):\n col_is_empty = all(not row[col_idx] for row in matrix)\n if col_is_empty:\n empty_buffer.append(col_idx)\n else:\n kept.extend(empty_buffer[:max_empty])\n kept.append(col_idx)\n empty_buffer = []\n\n return kept\n\n\ndef _remove_empty_runs(\n rows: list[list[str]],\n max_empty: int,\n) -> list[list[str]]:\n \"\"\"Removes entire runs of empty rows when the run length exceeds max_empty.\n\n Leading empty runs are capped to max_empty, just like interior runs.\n Trailing empty rows are always dropped since there is no subsequent\n non-empty row to flush them.\n \"\"\"\n result: list[list[str]] = []\n empty_buffer: list[list[str]] = []\n\n for row in rows:\n # Check if empty\n if not any(row):\n if len(empty_buffer) < max_empty:\n empty_buffer.append(row)\n else:\n # Add upto max empty rows onto the result - that's what we allow\n result.extend(empty_buffer[:max_empty])\n # Add the new non-empty row\n result.append(row)\n empty_buffer = []\n\n return result\n\n\ndef xlsx_to_text(file: IO[Any], file_name: str = \"\") -> str:\n # TODO: switch back to this approach in a few months when markitdown\n # fixes their handling of excel files\n\n # md = get_markitdown_converter()\n # stream_info = StreamInfo(\n # mimetype=SPREADSHEET_MIME_TYPE, filename=file_name or None, extension=\".xlsx\"\n # )\n # try:\n # workbook = md.convert(to_bytesio(file), stream_info=stream_info)\n # except (\n # BadZipFile,\n # ValueError,\n # FileConversionException,\n # UnsupportedFormatException,\n # ) as e:\n # error_str = f\"Failed to extract text from {file_name or 'xlsx file'}: {e}\"\n # if file_name.startswith(\"~\"):\n # logger.debug(error_str + \" (this is expected for files with ~)\")\n # else:\n # logger.warning(error_str)\n # return \"\"\n # return workbook.markdown\n try:\n workbook = openpyxl.load_workbook(file, read_only=True)\n except BadZipFile as e:\n error_str = f\"Failed to extract text from {file_name or 'xlsx file'}: {e}\"\n if file_name.startswith(\"~\"):\n logger.debug(error_str + \" (this is expected for files with ~)\")\n else:\n logger.warning(error_str)\n return \"\"\n except Exception as e:\n if any(s in str(e) for s in KNOWN_OPENPYXL_BUGS):\n logger.error(\n f\"Failed to extract text from {file_name or 'xlsx file'}. This happens due to a bug in openpyxl. {e}\"\n )\n return \"\"\n raise\n\n text_content = []\n for sheet in workbook.worksheets:\n sheet_matrix = _clean_worksheet_matrix(_worksheet_to_matrix(sheet))\n buf = io.StringIO()\n writer = csv.writer(buf, lineterminator=\"\\n\")\n writer.writerows(sheet_matrix)\n text_content.append(buf.getvalue().rstrip(\"\\n\"))\n return TEXT_SECTION_SEPARATOR.join(text_content)\n\n\ndef eml_to_text(file: IO[Any]) -> str:\n encoding = detect_encoding(file)\n text_file = io.TextIOWrapper(file, encoding=encoding)\n parser = EmailParser()\n try:\n message = parser.parse(text_file)\n finally:\n try:\n # Keep underlying upload handle open for downstream consumers.\n raw_file = text_file.detach()\n except Exception as detach_error:\n logger.warning(\n f\"Failed to detach TextIOWrapper for EML upload, using original file: {detach_error}\"\n )\n raw_file = file\n try:\n raw_file.seek(0)\n except Exception:\n pass\n\n text_content = []\n for part in message.walk():\n if part.get_content_type().startswith(\"text/plain\"):\n payload = part.get_payload()\n if isinstance(payload, str):\n text_content.append(payload)\n elif isinstance(payload, list):\n text_content.extend(item for item in payload if isinstance(item, str))\n else:\n logger.warning(f\"Unexpected payload type: {type(payload)}\")\n return TEXT_SECTION_SEPARATOR.join(text_content)\n\n\ndef epub_to_text(file: IO[Any]) -> str:\n with zipfile.ZipFile(file) as epub:\n text_content = []\n for item in epub.infolist():\n if item.filename.endswith(\".xhtml\") or item.filename.endswith(\".html\"):\n with epub.open(item) as html_file:\n text_content.append(parse_html_page_basic(html_file))\n return TEXT_SECTION_SEPARATOR.join(text_content)\n\n\ndef file_io_to_text(file: IO[Any]) -> str:\n encoding = detect_encoding(file)\n file_content, _ = read_text_file(file, encoding=encoding)\n return file_content\n\n\ndef extract_file_text(\n file: IO[Any],\n file_name: str,\n break_on_unprocessable: bool = True,\n extension: str | None = None,\n) -> str:\n \"\"\"\n Legacy function that returns *only text*, ignoring embedded images.\n For backward-compatibility in code that only wants text.\n\n NOTE: Ignoring seems to be defined as returning an empty string for files it can't\n handle (such as images).\n \"\"\"\n extension_to_function: dict[str, Callable[[IO[Any]], str]] = {\n \".pdf\": pdf_to_text,\n \".docx\": lambda f: read_docx_file(f, file_name)[0], # no images\n \".pptx\": lambda f: pptx_to_text(f, file_name),\n \".xlsx\": lambda f: xlsx_to_text(f, file_name),\n \".eml\": eml_to_text,\n \".epub\": epub_to_text,\n \".html\": parse_html_page_basic,\n }\n\n try:\n if get_unstructured_api_key():\n try:\n return unstructured_to_text(file, file_name)\n except Exception as unstructured_error:\n logger.error(\n f\"Failed to process with Unstructured: {str(unstructured_error)}. Falling back to normal processing.\"\n )\n if extension is None:\n extension = get_file_ext(file_name)\n\n if extension in OnyxFileExtensions.TEXT_AND_DOCUMENT_EXTENSIONS:\n func = extension_to_function.get(extension, file_io_to_text)\n file.seek(0)\n return func(file)\n\n # If unknown extension, maybe it's a text file\n file.seek(0)\n if is_text_file(file):\n return file_io_to_text(file)\n\n raise ValueError(\"Unknown file extension or not recognized as text data\")\n\n except Exception as e:\n if break_on_unprocessable:\n raise RuntimeError(\n f\"Failed to process file {file_name or 'Unknown'}: {str(e)}\"\n ) from e\n logger.warning(f\"Failed to process file {file_name or 'Unknown'}: {str(e)}\")\n return \"\"\n\n\nclass ExtractionResult(NamedTuple):\n \"\"\"Structured result from text and image extraction from various file types.\"\"\"\n\n text_content: str\n embedded_images: Sequence[tuple[bytes, str]]\n metadata: dict[str, Any]\n\n\ndef extract_result_from_text_file(file: IO[Any]) -> ExtractionResult:\n encoding = detect_encoding(file)\n text_content_raw, file_metadata = read_text_file(\n file, encoding=encoding, ignore_onyx_metadata=False\n )\n return ExtractionResult(\n text_content=text_content_raw,\n embedded_images=[],\n metadata=file_metadata,\n )\n\n\ndef extract_text_and_images(\n file: IO[Any],\n file_name: str,\n pdf_pass: str | None = None,\n content_type: str | None = None,\n image_callback: Callable[[bytes, str], None] | None = None,\n) -> ExtractionResult:\n \"\"\"\n Primary new function for the updated connector.\n Returns structured extraction result with text content, embedded images, and metadata.\n\n Args:\n file: File-like object to extract content from.\n file_name: Name of the file (used to determine extension/type).\n pdf_pass: Optional password for encrypted PDFs.\n content_type: Optional MIME type override for the file.\n image_callback: Optional callback for streaming image extraction. When provided,\n embedded images are passed to this callback one at a time as (bytes, filename)\n instead of being accumulated in the returned ExtractionResult.embedded_images\n list. This is a memory optimization for large documents with many images -\n the caller can process/store each image immediately rather than holding all\n images in memory. When using a callback, ExtractionResult.embedded_images\n will be an empty list.\n\n Returns:\n ExtractionResult containing text_content, embedded_images (empty if callback used),\n and metadata extracted from the file.\n \"\"\"\n res = _extract_text_and_images(\n file, file_name, pdf_pass, content_type, image_callback\n )\n # Clean up any temporary objects and force garbage collection\n unreachable = gc.collect()\n logger.info(f\"Unreachable objects: {unreachable}\")\n\n return res\n\n\ndef _extract_text_and_images(\n file: IO[Any],\n file_name: str,\n pdf_pass: str | None = None,\n content_type: str | None = None,\n image_callback: Callable[[bytes, str], None] | None = None,\n) -> ExtractionResult:\n file.seek(0)\n\n if get_unstructured_api_key():\n try:\n text_content = unstructured_to_text(file, file_name)\n return ExtractionResult(\n text_content=text_content, embedded_images=[], metadata={}\n )\n except Exception as e:\n logger.error(\n f\"Failed to process with Unstructured: {str(e)}. Falling back to normal processing.\"\n )\n file.seek(0) # Reset file pointer just in case\n\n # When we upload a document via a connector or MyDocuments, we extract and store the content of files\n # with content types in UploadMimeTypes.DOCUMENT_MIME_TYPES as plain text files.\n # As a result, the file name extension may differ from the original content type.\n # We process files with a plain text content type first to handle this scenario.\n if content_type in OnyxMimeTypes.TEXT_MIME_TYPES:\n return extract_result_from_text_file(file)\n\n # Default processing\n try:\n extension = get_file_ext(file_name)\n # docx example for embedded images\n if extension == \".docx\":\n text_content, images = read_docx_file(\n file, file_name, extract_images=True, image_callback=image_callback\n )\n return ExtractionResult(\n text_content=text_content, embedded_images=images, metadata={}\n )\n\n # PDF example: we do not show complicated PDF image extraction here\n # so we simply extract text for now and skip images.\n if extension == \".pdf\":\n text_content, pdf_metadata, images = read_pdf_file(\n file,\n pdf_pass,\n extract_images=get_image_extraction_and_analysis_enabled(),\n image_callback=image_callback,\n )\n return ExtractionResult(\n text_content=text_content, embedded_images=images, metadata=pdf_metadata\n )\n\n # For PPTX, XLSX, EML, etc., we do not show embedded image logic here.\n # You can do something similar to docx if needed.\n if extension == \".pptx\":\n return ExtractionResult(\n text_content=pptx_to_text(file, file_name=file_name),\n embedded_images=[],\n metadata={},\n )\n\n if extension == \".xlsx\":\n return ExtractionResult(\n text_content=xlsx_to_text(file, file_name=file_name),\n embedded_images=[],\n metadata={},\n )\n\n if extension == \".eml\":\n return ExtractionResult(\n text_content=eml_to_text(file), embedded_images=[], metadata={}\n )\n\n if extension == \".epub\":\n return ExtractionResult(\n text_content=epub_to_text(file), embedded_images=[], metadata={}\n )\n\n if extension == \".html\":\n return ExtractionResult(\n text_content=parse_html_page_basic(file),\n embedded_images=[],\n metadata={},\n )\n\n # If we reach here and it's a recognized text extension\n if extension in OnyxFileExtensions.PLAIN_TEXT_EXTENSIONS:\n return extract_result_from_text_file(file)\n\n # If it's an image file or something else, we do not parse embedded images from them\n # just return empty text\n return ExtractionResult(text_content=\"\", embedded_images=[], metadata={})\n\n except Exception as e:\n logger.exception(f\"Failed to extract text/images from {file_name}: {e}\")\n return ExtractionResult(text_content=\"\", embedded_images=[], metadata={})\n\n\ndef docx_to_txt_filename(file_path: str) -> str:\n return file_path.rsplit(\".\", 1)[0] + \".txt\"\n" }, { "path": "backend/onyx/file_processing/file_types.py", "content": "PRESENTATION_MIME_TYPE = (\n \"application/vnd.openxmlformats-officedocument.presentationml.presentation\"\n)\n\nSPREADSHEET_MIME_TYPE = (\n \"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet\"\n)\nWORD_PROCESSING_MIME_TYPE = (\n \"application/vnd.openxmlformats-officedocument.wordprocessingml.document\"\n)\nPDF_MIME_TYPE = \"application/pdf\"\nPLAIN_TEXT_MIME_TYPE = \"text/plain\"\n\n\nclass OnyxMimeTypes:\n IMAGE_MIME_TYPES = {\"image/jpg\", \"image/jpeg\", \"image/png\", \"image/webp\"}\n CSV_MIME_TYPES = {\"text/csv\"}\n TABULAR_MIME_TYPES = CSV_MIME_TYPES | {SPREADSHEET_MIME_TYPE}\n TEXT_MIME_TYPES = {\n PLAIN_TEXT_MIME_TYPE,\n \"text/markdown\",\n \"text/x-markdown\",\n \"text/x-log\",\n \"text/x-config\",\n \"text/tab-separated-values\",\n \"application/json\",\n \"application/xml\",\n \"text/xml\",\n \"application/x-yaml\",\n \"application/yaml\",\n \"text/yaml\",\n \"text/x-yaml\",\n }\n DOCUMENT_MIME_TYPES = {\n PDF_MIME_TYPE,\n WORD_PROCESSING_MIME_TYPE,\n PRESENTATION_MIME_TYPE,\n \"message/rfc822\",\n \"application/epub+zip\",\n }\n\n ALLOWED_MIME_TYPES = IMAGE_MIME_TYPES.union(\n TEXT_MIME_TYPES, DOCUMENT_MIME_TYPES, TABULAR_MIME_TYPES\n )\n\n EXCLUDED_IMAGE_TYPES = {\n \"image/bmp\",\n \"image/tiff\",\n \"image/gif\",\n \"image/svg+xml\",\n \"image/avif\",\n }\n\n\nclass OnyxFileExtensions:\n TABULAR_EXTENSIONS = {\n \".csv\",\n \".tsv\",\n \".xlsx\",\n }\n PLAIN_TEXT_EXTENSIONS = {\n \".txt\",\n \".md\",\n \".mdx\",\n \".conf\",\n \".log\",\n \".json\",\n \".csv\",\n \".tsv\",\n \".xml\",\n \".yml\",\n \".yaml\",\n \".sql\",\n }\n DOCUMENT_EXTENSIONS = {\n \".pdf\",\n \".docx\",\n \".pptx\",\n \".xlsx\",\n \".eml\",\n \".epub\",\n \".html\",\n }\n IMAGE_EXTENSIONS = {\n \".png\",\n \".jpg\",\n \".jpeg\",\n \".webp\",\n }\n\n TEXT_AND_DOCUMENT_EXTENSIONS = PLAIN_TEXT_EXTENSIONS.union(DOCUMENT_EXTENSIONS)\n\n ALL_ALLOWED_EXTENSIONS = TEXT_AND_DOCUMENT_EXTENSIONS.union(IMAGE_EXTENSIONS)\n" }, { "path": "backend/onyx/file_processing/html_utils.py", "content": "import re\nfrom copy import copy\nfrom dataclasses import dataclass\nfrom io import BytesIO\nfrom typing import IO\n\nimport bs4\n\nfrom onyx.configs.app_configs import HTML_BASED_CONNECTOR_TRANSFORM_LINKS_STRATEGY\nfrom onyx.configs.app_configs import PARSE_WITH_TRAFILATURA\nfrom onyx.configs.app_configs import WEB_CONNECTOR_IGNORED_CLASSES\nfrom onyx.configs.app_configs import WEB_CONNECTOR_IGNORED_ELEMENTS\nfrom onyx.file_processing.enums import HtmlBasedConnectorTransformLinksStrategy\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nMINTLIFY_UNWANTED = [\"sticky\", \"hidden\"]\n\n\n@dataclass\nclass ParsedHTML:\n title: str | None\n cleaned_text: str\n\n\ndef strip_excessive_newlines_and_spaces(document: str) -> str:\n # collapse repeated spaces into one\n document = re.sub(r\" +\", \" \", document)\n # remove trailing spaces\n document = re.sub(r\" +[\\n\\r]\", \"\\n\", document)\n # remove repeated newlines\n document = re.sub(r\"[\\n\\r]+\", \"\\n\", document)\n return document.strip()\n\n\ndef strip_newlines(document: str) -> str:\n # HTML might contain newlines which are just whitespaces to a browser\n return re.sub(r\"[\\n\\r]+\", \" \", document)\n\n\ndef format_element_text(element_text: str, link_href: str | None) -> str:\n element_text_no_newlines = strip_newlines(element_text)\n\n if (\n not link_href\n or HTML_BASED_CONNECTOR_TRANSFORM_LINKS_STRATEGY\n == HtmlBasedConnectorTransformLinksStrategy.STRIP\n ):\n return element_text_no_newlines\n\n return f\"[{element_text_no_newlines}]({link_href})\"\n\n\ndef parse_html_with_trafilatura(html_content: str) -> str:\n \"\"\"Parse HTML content using trafilatura.\"\"\"\n import trafilatura # type: ignore\n from trafilatura.settings import use_config # type: ignore\n\n config = use_config()\n config.set(\"DEFAULT\", \"include_links\", \"True\")\n config.set(\"DEFAULT\", \"include_tables\", \"True\")\n config.set(\"DEFAULT\", \"include_images\", \"True\")\n config.set(\"DEFAULT\", \"include_formatting\", \"True\")\n\n extracted_text = trafilatura.extract(html_content, config=config)\n return strip_excessive_newlines_and_spaces(extracted_text) if extracted_text else \"\"\n\n\ndef format_document_soup(\n document: bs4.BeautifulSoup, table_cell_separator: str = \"\\t\"\n) -> str:\n \"\"\"Format html to a flat text document.\n\n The following goals:\n - Newlines from within the HTML are removed (as browser would ignore them as well).\n - Repeated newlines/spaces are removed (as browsers would ignore them).\n - Newlines only before and after headlines and paragraphs or when explicit (br or pre tag)\n - Table columns/rows are separated by newline\n - List elements are separated by newline and start with a hyphen\n \"\"\"\n text = \"\"\n list_element_start = False\n verbatim_output = 0\n in_table = False\n last_added_newline = False\n link_href: str | None = None\n\n for e in document.descendants:\n verbatim_output -= 1\n if isinstance(e, bs4.element.NavigableString):\n if isinstance(e, (bs4.element.Comment, bs4.element.Doctype)):\n continue\n element_text = e.text\n if in_table:\n # Tables are represented in natural language with rows separated by newlines\n # Can't have newlines then in the table elements\n element_text = element_text.replace(\"\\n\", \" \").strip()\n\n # Some tags are translated to spaces but in the logic underneath this section, we\n # translate them to newlines as a browser should render them such as with br\n # This logic here avoids a space after newline when it shouldn't be there.\n if last_added_newline and element_text.startswith(\" \"):\n element_text = element_text[1:]\n last_added_newline = False\n\n if element_text:\n content_to_add = (\n element_text\n if verbatim_output > 0\n else format_element_text(element_text, link_href)\n )\n\n # Don't join separate elements without any spacing\n if (text and not text[-1].isspace()) and (\n content_to_add and not content_to_add[0].isspace()\n ):\n text += \" \"\n\n text += content_to_add\n\n list_element_start = False\n elif isinstance(e, bs4.element.Tag):\n # table is standard HTML element\n if e.name == \"table\":\n in_table = True\n # tr is for rows\n elif e.name == \"tr\" and in_table:\n text += \"\\n\"\n # td for data cell, th for header\n elif e.name in [\"td\", \"th\"] and in_table:\n text += table_cell_separator\n elif e.name == \"/table\":\n in_table = False\n elif in_table:\n # don't handle other cases while in table\n pass\n elif e.name == \"a\":\n href_value = e.get(\"href\", None)\n # mostly for typing, having multiple hrefs is not valid HTML\n link_href = (\n href_value[0] if isinstance(href_value, list) else href_value\n )\n elif e.name == \"/a\":\n link_href = None\n elif e.name in [\"p\", \"div\"]:\n if not list_element_start:\n text += \"\\n\"\n elif e.name in [\"h1\", \"h2\", \"h3\", \"h4\"]:\n text += \"\\n\"\n list_element_start = False\n last_added_newline = True\n elif e.name == \"br\":\n text += \"\\n\"\n list_element_start = False\n last_added_newline = True\n elif e.name == \"li\":\n text += \"\\n- \"\n list_element_start = True\n elif e.name == \"pre\":\n if verbatim_output <= 0:\n verbatim_output = len(list(e.childGenerator()))\n return strip_excessive_newlines_and_spaces(text)\n\n\ndef parse_html_page_basic(text: str | BytesIO | IO[bytes]) -> str:\n soup = bs4.BeautifulSoup(text, \"lxml\")\n return format_document_soup(soup)\n\n\ndef web_html_cleanup(\n page_content: str | bs4.BeautifulSoup,\n mintlify_cleanup_enabled: bool = True,\n additional_element_types_to_discard: list[str] | None = None,\n) -> ParsedHTML:\n if isinstance(page_content, str):\n soup = bs4.BeautifulSoup(page_content, \"lxml\")\n else:\n soup = page_content\n\n title_tag = soup.find(\"title\")\n title = None\n if title_tag and title_tag.text:\n title = title_tag.text\n title_tag.extract()\n\n # Heuristics based cleaning of elements based on css classes\n unwanted_classes = copy(WEB_CONNECTOR_IGNORED_CLASSES)\n if mintlify_cleanup_enabled:\n unwanted_classes.extend(MINTLIFY_UNWANTED)\n for undesired_element in unwanted_classes:\n [\n tag.extract()\n for tag in soup.find_all(\n class_=lambda x: x and undesired_element in x.split()\n )\n ]\n\n for undesired_tag in WEB_CONNECTOR_IGNORED_ELEMENTS:\n [tag.extract() for tag in soup.find_all(undesired_tag)]\n\n if additional_element_types_to_discard:\n for undesired_tag in additional_element_types_to_discard:\n [tag.extract() for tag in soup.find_all(undesired_tag)]\n\n soup_string = str(soup)\n page_text = \"\"\n\n if PARSE_WITH_TRAFILATURA:\n try:\n page_text = parse_html_with_trafilatura(soup_string)\n if not page_text:\n raise ValueError(\"Empty content returned by trafilatura.\")\n except Exception as e:\n logger.info(f\"Trafilatura parsing failed: {e}. Falling back on bs4.\")\n page_text = format_document_soup(soup)\n else:\n page_text = format_document_soup(soup)\n\n # 200B is ZeroWidthSpace which we don't care for\n cleaned_text = page_text.replace(\"\\u200b\", \"\")\n\n return ParsedHTML(title=title, cleaned_text=cleaned_text)\n" }, { "path": "backend/onyx/file_processing/image_summarization.py", "content": "import base64\nfrom io import BytesIO\n\nfrom PIL import Image\n\nfrom onyx.configs.app_configs import IMAGE_SUMMARIZATION_SYSTEM_PROMPT\nfrom onyx.configs.app_configs import IMAGE_SUMMARIZATION_USER_PROMPT\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import ContentPart\nfrom onyx.llm.models import ImageContentPart\nfrom onyx.llm.models import ImageUrlDetail\nfrom onyx.llm.models import SystemMessage\nfrom onyx.llm.models import TextContentPart\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.b64 import get_image_type_from_bytes\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass UnsupportedImageFormatError(ValueError):\n \"\"\"Raised when an image uses a MIME type unsupported by the summarization flow.\"\"\"\n\n\ndef prepare_image_bytes(image_data: bytes) -> str:\n \"\"\"Prepare image bytes for summarization.\n Resizes image if it's larger than 20MB. Encodes image as a base64 string.\"\"\"\n image_data = _resize_image_if_needed(image_data)\n\n # encode image (base64)\n encoded_image = _encode_image_for_llm_prompt(image_data)\n\n return encoded_image\n\n\ndef summarize_image_pipeline(\n llm: LLM,\n image_data: bytes,\n query: str | None = None,\n system_prompt: str | None = None,\n) -> str:\n \"\"\"Pipeline to generate a summary of an image.\n Resizes images if it is bigger than 20MB. Encodes image as a base64 string.\n And finally uses the Default LLM to generate a textual summary of the image.\"\"\"\n # resize image if it's bigger than 20MB\n encoded_image = prepare_image_bytes(image_data)\n\n summary = _summarize_image(\n encoded_image,\n llm,\n query,\n system_prompt,\n )\n\n return summary\n\n\ndef summarize_image_with_error_handling(\n llm: LLM | None,\n image_data: bytes,\n context_name: str,\n system_prompt: str = IMAGE_SUMMARIZATION_SYSTEM_PROMPT,\n user_prompt_template: str = IMAGE_SUMMARIZATION_USER_PROMPT,\n) -> str | None:\n \"\"\"Wrapper function that handles error cases and configuration consistently.\n\n Args:\n llm: The LLM with vision capabilities to use for summarization\n image_data: The raw image bytes\n context_name: Name or title of the image for context\n system_prompt: System prompt to use for the LLM\n user_prompt_template: User prompt to use (without title)\n\n Returns:\n The image summary text, or None if summarization failed or is disabled\n \"\"\"\n if llm is None:\n return None\n\n # Prepend the image filename to the user prompt\n user_prompt = (\n f\"The image has the file name '{context_name}'.\\n{user_prompt_template}\"\n )\n try:\n return summarize_image_pipeline(llm, image_data, user_prompt, system_prompt)\n except UnsupportedImageFormatError:\n magic_hex = image_data[:8].hex() if image_data else \"empty\"\n logger.info(\n \"Skipping image summarization due to unsupported MIME type \"\n \"for %s (magic_bytes=%s, size=%d bytes)\",\n context_name,\n magic_hex,\n len(image_data),\n )\n return None\n\n\ndef _summarize_image(\n encoded_image: str,\n llm: LLM,\n query: str | None = None,\n system_prompt: str | None = None,\n) -> str:\n \"\"\"Use default LLM (if it is multimodal) to generate a summary of an image.\"\"\"\n\n messages: list[ChatCompletionMessage] = []\n\n if system_prompt:\n messages.append(SystemMessage(content=system_prompt))\n\n content: list[ContentPart] = []\n if query:\n content.append(TextContentPart(text=query))\n content.append(ImageContentPart(image_url=ImageUrlDetail(url=encoded_image)))\n\n messages.append(\n UserMessage(\n content=content,\n ),\n )\n\n try:\n # Call LLM with Braintrust tracing\n with llm_generation_span(\n llm=llm,\n flow=\"image_summarization\",\n input_messages=[{\"type\": \"image_summarization_request\"}],\n ) as span_generation:\n # Note: We don't include the actual image in the span input to avoid bloating traces\n response = llm.invoke(messages)\n record_llm_response(span_generation, response)\n summary = llm_response_to_string(response)\n\n return summary\n\n except Exception as e:\n # Extract structured details from LiteLLM exceptions when available,\n # rather than dumping the full messages payload (which contains base64\n # image data and produces enormous, unreadable error logs).\n str_e = str(e)\n if len(str_e) > 512:\n str_e = str_e[:512] + \"... (truncated)\"\n parts = [f\"Summarization failed: {type(e).__name__}: {str_e}\"]\n status_code = getattr(e, \"status_code\", None)\n llm_provider = getattr(e, \"llm_provider\", None)\n model = getattr(e, \"model\", None)\n if status_code is not None:\n parts.append(f\"status_code={status_code}\")\n if llm_provider is not None:\n parts.append(f\"llm_provider={llm_provider}\")\n if model is not None:\n parts.append(f\"model={model}\")\n raise ValueError(\" | \".join(parts)) from e\n\n\ndef _encode_image_for_llm_prompt(image_data: bytes) -> str:\n \"\"\"Prepare a data URL with the correct MIME type for the LLM message.\"\"\"\n try:\n mime_type = get_image_type_from_bytes(image_data)\n except ValueError as exc:\n raise UnsupportedImageFormatError(\n \"Unsupported image format for summarization\"\n ) from exc\n\n base64_encoded_data = base64.b64encode(image_data).decode(\"utf-8\")\n\n return f\"data:{mime_type};base64,{base64_encoded_data}\"\n\n\ndef _resize_image_if_needed(image_data: bytes, max_size_mb: int = 20) -> bytes:\n \"\"\"Resize image if it's larger than the specified max size in MB.\"\"\"\n max_size_bytes = max_size_mb * 1024 * 1024\n\n if len(image_data) > max_size_bytes:\n with Image.open(BytesIO(image_data)) as img:\n # Reduce dimensions for better size reduction\n img.thumbnail((1024, 1024), Image.Resampling.LANCZOS)\n output = BytesIO()\n\n # Save with lower quality for compression\n img.save(output, format=\"JPEG\", quality=85)\n resized_data = output.getvalue()\n\n return resized_data\n\n return image_data\n" }, { "path": "backend/onyx/file_processing/image_utils.py", "content": "from io import BytesIO\nfrom typing import Tuple\n\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.models import ImageSection\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef store_image_and_create_section(\n image_data: bytes,\n file_id: str,\n display_name: str,\n link: str | None = None,\n media_type: str = \"application/octet-stream\",\n file_origin: FileOrigin = FileOrigin.OTHER,\n) -> Tuple[ImageSection, str | None]:\n \"\"\"\n Stores an image in FileStore and creates an ImageSection object without summarization.\n\n Args:\n image_data: Raw image bytes\n file_id: Base identifier for the file\n display_name: Human-readable name for the image\n media_type: MIME type of the image\n file_origin: Origin of the file (e.g., CONFLUENCE, GOOGLE_DRIVE, etc.)\n\n Returns:\n Tuple containing:\n - ImageSection object with image reference\n - The file_id in FileStore or None if storage failed\n \"\"\"\n # Storage logic\n try:\n file_store = get_default_file_store()\n file_id = file_store.save_file(\n content=BytesIO(image_data),\n display_name=display_name,\n file_origin=file_origin,\n file_type=media_type,\n file_id=file_id,\n )\n except Exception as e:\n logger.error(f\"Failed to store image: {e}\")\n raise e\n\n # Create an ImageSection with empty text (will be filled by LLM later in the pipeline)\n return (\n ImageSection(image_file_id=file_id, link=link),\n file_id,\n )\n" }, { "path": "backend/onyx/file_processing/password_validation.py", "content": "from collections.abc import Callable\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom typing import IO\n\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nPASSWORD_PROTECTED_FILES = [\n \".pdf\",\n \".docx\",\n \".pptx\",\n \".xlsx\",\n]\n\n\n@contextmanager\ndef preserve_position(file: IO[Any]) -> Generator[IO[Any], None, None]:\n \"\"\"Preserves the file's cursor position\"\"\"\n pos = file.tell()\n try:\n file.seek(0)\n yield file\n finally:\n file.seek(pos)\n\n\ndef is_pdf_protected(file: IO[Any]) -> bool:\n from pypdf import PdfReader\n\n with preserve_position(file):\n reader = PdfReader(file)\n\n return bool(reader.is_encrypted)\n\n\ndef is_docx_protected(file: IO[Any]) -> bool:\n return is_office_file_protected(file)\n\n\ndef is_pptx_protected(file: IO[Any]) -> bool:\n return is_office_file_protected(file)\n\n\ndef is_xlsx_protected(file: IO[Any]) -> bool:\n return is_office_file_protected(file)\n\n\ndef is_office_file_protected(file: IO[Any]) -> bool:\n import msoffcrypto # type: ignore[import-untyped]\n\n with preserve_position(file):\n office = msoffcrypto.OfficeFile(file)\n\n return office.is_encrypted()\n\n\ndef is_file_password_protected(\n file: IO[Any],\n file_name: str,\n extension: str | None = None,\n) -> bool:\n extension_to_function: dict[str, Callable[[IO[Any]], bool]] = {\n \".pdf\": is_pdf_protected,\n \".docx\": is_docx_protected,\n \".pptx\": is_pptx_protected,\n \".xlsx\": is_xlsx_protected,\n }\n\n if not extension:\n extension = get_file_ext(file_name)\n\n if extension not in PASSWORD_PROTECTED_FILES:\n return False\n\n if extension not in extension_to_function:\n logger.warning(\n f\"Extension={extension} can be password protected, but no function found\"\n )\n return False\n\n func = extension_to_function[extension]\n\n return func(file)\n" }, { "path": "backend/onyx/file_processing/unstructured.py", "content": "from typing import Any\nfrom typing import cast\nfrom typing import IO\nfrom typing import TYPE_CHECKING\n\nfrom onyx.configs.constants import KV_UNSTRUCTURED_API_KEY\nfrom onyx.key_value_store.factory import get_kv_store\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.utils.logger import setup_logger\n\nif TYPE_CHECKING:\n from unstructured_client.models import operations\n\n\nlogger = setup_logger()\n\n\ndef get_unstructured_api_key() -> str | None:\n kv_store = get_kv_store()\n try:\n return cast(str, kv_store.load(KV_UNSTRUCTURED_API_KEY))\n except KvKeyNotFoundError:\n return None\n\n\ndef update_unstructured_api_key(api_key: str) -> None:\n kv_store = get_kv_store()\n kv_store.store(KV_UNSTRUCTURED_API_KEY, api_key)\n\n\ndef delete_unstructured_api_key() -> None:\n kv_store = get_kv_store()\n kv_store.delete(KV_UNSTRUCTURED_API_KEY)\n\n\ndef _sdk_partition_request(\n file: IO[Any], file_name: str, **kwargs: Any\n) -> \"operations.PartitionRequest\":\n from unstructured_client.models import operations\n from unstructured_client.models import shared\n\n file.seek(0, 0)\n try:\n request = operations.PartitionRequest(\n partition_parameters=shared.PartitionParameters(\n files=shared.Files(content=file.read(), file_name=file_name),\n **kwargs,\n ),\n )\n return request\n except Exception as e:\n logger.error(f\"Error creating partition request for file {file_name}: {str(e)}\")\n raise\n\n\ndef unstructured_to_text(file: IO[Any], file_name: str) -> str:\n from unstructured.staging.base import dict_to_elements\n from unstructured_client import UnstructuredClient\n\n logger.debug(f\"Starting to read file: {file_name}\")\n req = _sdk_partition_request(file, file_name, strategy=\"fast\")\n\n unstructured_client = UnstructuredClient(api_key_auth=get_unstructured_api_key())\n\n response = unstructured_client.general.partition(request=req)\n\n if response.status_code != 200:\n err = f\"Received unexpected status code {response.status_code} from Unstructured API.\"\n logger.error(err)\n raise ValueError(err)\n\n elements = dict_to_elements(response.elements or [])\n return \"\\n\\n\".join(str(el) for el in elements)\n" }, { "path": "backend/onyx/file_store/README.md", "content": "# Onyx File Store\n\nThe Onyx file store provides a unified interface for storing files and large binary objects in S3-compatible storage systems. It supports AWS S3, MinIO, Azure Blob Storage, Digital Ocean Spaces, and other S3-compatible services.\n\n## Architecture\n\nThe file store uses a single database table (`file_record`) to store file metadata while the actual file content is stored in external S3-compatible storage. This approach provides scalability, cost-effectiveness, and decouples file storage from the database.\n\n### Database Schema\n\nThe `file_record` table contains the following columns:\n\n- `file_id` (primary key): Unique identifier for the file\n- `display_name`: Human-readable name for the file\n- `file_origin`: Origin/source of the file (enum)\n- `file_type`: MIME type of the file\n- `file_metadata`: Additional metadata as JSON\n- `bucket_name`: External storage bucket/container name\n- `object_key`: External storage object key/path\n- `created_at`: Timestamp when the file was created\n- `updated_at`: Timestamp when the file was last updated\n\n## Storage Backend\n\n### S3-Compatible Storage\n\nStores files in external S3-compatible storage systems while keeping metadata in the database.\n\n**Pros:**\n- Scalable storage\n- Cost-effective for large files\n- CDN integration possible\n- Decoupled from database\n- Wide ecosystem support\n\n**Cons:**\n- Additional infrastructure required\n- Network dependency\n- Eventual consistency considerations\n\n## Configuration\n\nAll configuration is handled via environment variables. The system requires S3-compatible storage to be configured.\n\n### AWS S3\n\n```bash\nS3_FILE_STORE_BUCKET_NAME=your-bucket-name # Defaults to 'onyx-file-store-bucket'\nS3_FILE_STORE_PREFIX=onyx-files # Optional, defaults to 'onyx-files'\n\n# AWS credentials (use one of these methods):\n# 1. Environment variables\nS3_AWS_ACCESS_KEY_ID=your-access-key\nS3_AWS_SECRET_ACCESS_KEY=your-secret-key\nAWS_REGION_NAME=us-east-2 # Optional, defaults to 'us-east-2'\n\n# 2. IAM roles (recommended for EC2/ECS deployments)\n# No additional configuration needed if using IAM roles\n```\n\n### MinIO\n\n```bash\nS3_FILE_STORE_BUCKET_NAME=your-bucket-name\nS3_ENDPOINT_URL=http://localhost:9000 # MinIO endpoint\nS3_AWS_ACCESS_KEY_ID=minioadmin\nS3_AWS_SECRET_ACCESS_KEY=minioadmin\nAWS_REGION_NAME=us-east-1 # Any region name\nS3_VERIFY_SSL=false # Optional, defaults to false\n```\n\n### Digital Ocean Spaces\n\n```bash\nS3_FILE_STORE_BUCKET_NAME=your-space-name\nS3_ENDPOINT_URL=https://nyc3.digitaloceanspaces.com\nS3_AWS_ACCESS_KEY_ID=your-spaces-key\nS3_AWS_SECRET_ACCESS_KEY=your-spaces-secret\nAWS_REGION_NAME=nyc3\n```\n\n### Other S3-Compatible Services\n\nThe file store works with any S3-compatible service. Simply configure:\n- `S3_FILE_STORE_BUCKET_NAME`: Your bucket/container name\n- `S3_ENDPOINT_URL`: The service endpoint URL\n- `S3_AWS_ACCESS_KEY_ID` and `S3_AWS_SECRET_ACCESS_KEY`: Your credentials\n- `AWS_REGION_NAME`: The region (any valid region name)\n\n## Implementation\n\nThe system uses the `S3BackedFileStore` class that implements the abstract `FileStore` interface. The database uses generic column names (`bucket_name`, `object_key`) to maintain compatibility with different S3-compatible services.\n\n### File Store Interface\n\nThe `FileStore` abstract base class defines the following methods:\n\n- `initialize()`: Initialize the storage backend (create bucket if needed)\n- `has_file(file_id, file_origin, file_type)`: Check if a file exists\n- `save_file(content, display_name, file_origin, file_type, file_metadata, file_id)`: Save a file\n- `read_file(file_id, mode, use_tempfile)`: Read file content\n- `read_file_record(file_id)`: Get file metadata from database\n- `delete_file(file_id)`: Delete a file and its metadata\n- `get_file_with_mime_type(file_id)`: Get file with parsed MIME type\n\n## Usage Example\n\n```python\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.configs.constants import FileOrigin\n\n# Get the configured file store\nfile_store = get_default_file_store(db_session)\n\n# Initialize the storage backend (creates bucket if needed)\nfile_store.initialize()\n\n# Save a file\nwith open(\"example.pdf\", \"rb\") as f:\n file_id = file_store.save_file(\n content=f,\n display_name=\"Important Document.pdf\",\n file_origin=FileOrigin.OTHER,\n file_type=\"application/pdf\",\n file_metadata={\"department\": \"engineering\", \"version\": \"1.0\"}\n )\n\n# Check if a file exists\nexists = file_store.has_file(\n file_id=file_id,\n file_origin=FileOrigin.OTHER,\n file_type=\"application/pdf\"\n)\n\n# Read a file\nfile_content = file_store.read_file(file_id)\n\n# Read file with temporary file (for large files)\nfile_content = file_store.read_file(file_id, use_tempfile=True)\n\n# Get file metadata\nfile_record = file_store.read_file_record(file_id)\n\n# Get file with MIME type detection\nfile_with_mime = file_store.get_file_with_mime_type(file_id)\n\n# Delete a file\nfile_store.delete_file(file_id)\n```\n\n## Initialization\n\nWhen deploying the application, ensure that:\n\n1. The S3-compatible storage service is accessible\n2. Credentials are properly configured\n3. The bucket specified in `S3_FILE_STORE_BUCKET_NAME` exists or the service account has permissions to create it\n4. Call `file_store.initialize()` during application startup to ensure the bucket exists\n\nThe file store will automatically create the bucket if it doesn't exist and the credentials have sufficient permissions.\n " }, { "path": "backend/onyx/file_store/constants.py", "content": "MAX_IN_MEMORY_SIZE = 30 * 1024 * 1024 # 30MB\nSTANDARD_CHUNK_SIZE = 10 * 1024 * 1024 # 10MB chunks\n" }, { "path": "backend/onyx/file_store/document_batch_storage.py", "content": "import json\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom enum import Enum\nfrom io import StringIO\nfrom typing import List\nfrom typing import Optional\nfrom typing import TypeAlias\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.connectors.models import DocExtractionContext\nfrom onyx.connectors.models import DocIndexingContext\nfrom onyx.connectors.models import Document\nfrom onyx.file_store.file_store import FileStore\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass DocumentBatchStorageStateType(str, Enum):\n EXTRACTION = \"extraction\"\n INDEXING = \"indexing\"\n\n\nDocumentStorageState: TypeAlias = DocExtractionContext | DocIndexingContext\n\nSTATE_TYPE_TO_MODEL: dict[str, type[DocumentStorageState]] = {\n DocumentBatchStorageStateType.EXTRACTION.value: DocExtractionContext,\n DocumentBatchStorageStateType.INDEXING.value: DocIndexingContext,\n}\n\n\nclass BatchStoragePathInfo(BaseModel):\n cc_pair_id: int\n index_attempt_id: int\n batch_num: int\n\n\nclass DocumentBatchStorage(ABC):\n \"\"\"Abstract base class for document batch storage implementations.\"\"\"\n\n def __init__(self, cc_pair_id: int, index_attempt_id: int):\n self.cc_pair_id = cc_pair_id\n self.index_attempt_id = index_attempt_id\n self.base_path = f\"{self._per_cc_pair_base_path()}/{index_attempt_id}\"\n\n @abstractmethod\n def store_batch(self, batch_num: int, documents: List[Document]) -> None:\n \"\"\"Store a batch of documents.\"\"\"\n\n @abstractmethod\n def get_batch(self, batch_num: int) -> Optional[List[Document]]:\n \"\"\"Retrieve a batch of documents.\"\"\"\n\n @abstractmethod\n def delete_batch_by_name(self, batch_file_name: str) -> None:\n \"\"\"Delete a specific batch.\"\"\"\n\n @abstractmethod\n def delete_batch_by_num(self, batch_num: int) -> None:\n \"\"\"Delete a specific batch.\"\"\"\n\n @abstractmethod\n def cleanup_all_batches(self) -> None:\n \"\"\"Clean up all batches and state for this index attempt.\"\"\"\n\n @abstractmethod\n def get_all_batches_for_cc_pair(self) -> list[str]:\n \"\"\"Get all IDs of batches stored in the file store.\"\"\"\n\n @abstractmethod\n def update_old_batches_to_new_index_attempt(self, batch_names: list[str]) -> None:\n \"\"\"Update all batches to the new index attempt.\"\"\"\n \"\"\"\n This is used when we need to re-issue docprocessing tasks for a new index attempt.\n We need to update the batch file names to the new index attempt ID.\n \"\"\"\n\n @abstractmethod\n def extract_path_info(self, path: str) -> BatchStoragePathInfo | None:\n \"\"\"Extract path info from a path.\"\"\"\n\n def _serialize_documents(self, documents: list[Document]) -> str:\n \"\"\"Serialize documents to JSON string.\"\"\"\n # Use mode='json' to properly serialize datetime and other complex types\n return json.dumps([doc.model_dump(mode=\"json\") for doc in documents], indent=2)\n\n def _deserialize_documents(self, data: str) -> list[Document]:\n \"\"\"Deserialize documents from JSON string.\"\"\"\n doc_dicts = json.loads(data)\n return [\n Document.model_validate(self._normalize_doc_dict(doc_dict))\n for doc_dict in doc_dicts\n ]\n\n def _normalize_doc_dict(self, doc_dict: dict) -> dict:\n \"\"\"Normalize document dict to handle legacy data with non-string metadata values.\n\n Before the _convert_to_metadata_value fix, Salesforce connector stored raw\n types (bool, float, None) in metadata. This converts them to strings for\n backward compatibility.\n \"\"\"\n if \"metadata\" not in doc_dict:\n return doc_dict\n\n metadata = doc_dict[\"metadata\"]\n if not isinstance(metadata, dict):\n return doc_dict\n\n normalized_metadata: dict[str, str | list[str]] = {}\n converted_keys: list[str] = []\n for key, value in metadata.items():\n if isinstance(value, list):\n normalized_metadata[key] = [str(item) for item in value]\n elif isinstance(value, str):\n normalized_metadata[key] = value\n else:\n # Convert bool, int, float, None to string\n converted_keys.append(f\"{key}={type(value).__name__}\")\n normalized_metadata[key] = str(value)\n\n if converted_keys:\n doc_id = doc_dict.get(\"id\", \"unknown\")\n logger.warning(\n f\"Normalized legacy metadata for document {doc_id}: {converted_keys}\"\n )\n\n doc_dict[\"metadata\"] = normalized_metadata\n return doc_dict\n\n def _per_cc_pair_base_path(self) -> str:\n \"\"\"Get the base path for the cc pair.\"\"\"\n return f\"iab/{self.cc_pair_id}\"\n\n\nclass FileStoreDocumentBatchStorage(DocumentBatchStorage):\n \"\"\"FileStore-based implementation of document batch storage.\"\"\"\n\n def __init__(self, cc_pair_id: int, index_attempt_id: int, file_store: FileStore):\n super().__init__(cc_pair_id, index_attempt_id)\n self.file_store = file_store\n\n def _get_batch_file_name(self, batch_num: int) -> str:\n \"\"\"Generate file name for a document batch.\"\"\"\n return f\"{self.base_path}/{batch_num}.json\"\n\n def store_batch(self, batch_num: int, documents: list[Document]) -> None:\n \"\"\"Store a batch of documents using FileStore.\"\"\"\n file_name = self._get_batch_file_name(batch_num)\n try:\n data = self._serialize_documents(documents)\n content = StringIO(data)\n\n self.file_store.save_file(\n file_id=file_name,\n content=content,\n display_name=f\"Document Batch {batch_num}\",\n file_origin=FileOrigin.OTHER,\n file_type=\"application/json\",\n file_metadata={\n \"batch_num\": batch_num,\n \"document_count\": str(len(documents)),\n },\n )\n\n logger.debug(\n f\"Stored batch {batch_num} with {len(documents)} documents to FileStore as {file_name}\"\n )\n except Exception as e:\n logger.error(f\"Failed to store batch {batch_num}: {e}\")\n raise\n\n def get_batch(self, batch_num: int) -> list[Document] | None:\n \"\"\"Retrieve a batch of documents from FileStore.\"\"\"\n file_name = self._get_batch_file_name(batch_num)\n try:\n # Check if file exists\n if not self.file_store.has_file(\n file_id=file_name,\n file_origin=FileOrigin.OTHER,\n file_type=\"application/json\",\n ):\n logger.warning(\n f\"Batch {batch_num} not found in FileStore with name {file_name}\"\n )\n return None\n\n content_io = self.file_store.read_file(file_name)\n data = content_io.read().decode(\"utf-8\")\n\n documents = self._deserialize_documents(data)\n logger.debug(\n f\"Retrieved batch {batch_num} with {len(documents)} documents from FileStore\"\n )\n return documents\n except Exception as e:\n logger.error(f\"Failed to retrieve batch {batch_num}: {e}\")\n raise\n\n def delete_batch_by_name(self, batch_file_name: str) -> None:\n \"\"\"Delete a specific batch from FileStore.\"\"\"\n self.file_store.delete_file(batch_file_name)\n logger.debug(f\"Deleted batch {batch_file_name} from FileStore\")\n\n def delete_batch_by_num(self, batch_num: int) -> None:\n \"\"\"Delete a specific batch from FileStore.\"\"\"\n batch_file_name = self._get_batch_file_name(batch_num)\n self.delete_batch_by_name(batch_file_name)\n logger.debug(f\"Deleted batch num {batch_num} {batch_file_name} from FileStore\")\n\n def cleanup_all_batches(self) -> None:\n \"\"\"Clean up all batches for this index attempt.\"\"\"\n for batch_file_name in self.get_all_batches_for_cc_pair():\n self.delete_batch_by_name(batch_file_name)\n\n def get_all_batches_for_cc_pair(self) -> list[str]:\n \"\"\"Get all IDs of batches stored in the file store for the cc pair\n this batch store was initialized with.\n This includes any batches left over from a previous\n indexing attempt that need to be processed.\n \"\"\"\n return [\n file.file_id\n for file in self.file_store.list_files_by_prefix(\n self._per_cc_pair_base_path()\n )\n ]\n\n def update_old_batches_to_new_index_attempt(self, batch_names: list[str]) -> None:\n \"\"\"Update all batches to the new index attempt.\"\"\"\n for batch_file_name in batch_names:\n path_info = self.extract_path_info(batch_file_name)\n if path_info is None:\n logger.warning(\n f\"Could not extract path info from batch file: {batch_file_name}\"\n )\n continue\n new_batch_file_name = self._get_batch_file_name(path_info.batch_num)\n self.file_store.change_file_id(batch_file_name, new_batch_file_name)\n\n def extract_path_info(self, path: str) -> BatchStoragePathInfo | None:\n \"\"\"Extract path info from a path.\"\"\"\n path_spl = path.split(\"/\")\n # TODO: remove this in a few months, just for backwards compatibility\n if len(path_spl) == 3:\n path_spl = [\"iab\"] + path_spl\n try:\n _, cc_pair_id, index_attempt_id, batch_num = path_spl\n return BatchStoragePathInfo(\n cc_pair_id=int(cc_pair_id),\n index_attempt_id=int(index_attempt_id),\n batch_num=int(batch_num.split(\".\")[0]), # remove .json\n )\n except Exception as e:\n logger.error(f\"Failed to extract path info from {path}: {e}\")\n return None\n\n\ndef get_document_batch_storage(\n cc_pair_id: int, index_attempt_id: int\n) -> DocumentBatchStorage:\n \"\"\"Factory function to get the configured document batch storage implementation.\"\"\"\n # The get_default_file_store will now correctly use S3BackedFileStore\n # or other configured stores based on environment variables\n file_store = get_default_file_store()\n return FileStoreDocumentBatchStorage(cc_pair_id, index_attempt_id, file_store)\n" }, { "path": "backend/onyx/file_store/file_store.py", "content": "import hashlib\nimport tempfile\nimport uuid\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import cast\nfrom typing import IO\nfrom typing import NotRequired\nfrom typing import TypedDict\n\nimport boto3\nimport puremagic\nfrom botocore.config import Config\nfrom botocore.exceptions import ClientError\nfrom mypy_boto3_s3 import S3Client\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import AWS_REGION_NAME\nfrom onyx.configs.app_configs import S3_AWS_ACCESS_KEY_ID\nfrom onyx.configs.app_configs import S3_AWS_SECRET_ACCESS_KEY\nfrom onyx.configs.app_configs import S3_ENDPOINT_URL\nfrom onyx.configs.app_configs import S3_FILE_STORE_BUCKET_NAME\nfrom onyx.configs.app_configs import S3_FILE_STORE_PREFIX\nfrom onyx.configs.app_configs import S3_GENERATE_LOCAL_CHECKSUM\nfrom onyx.configs.app_configs import S3_VERIFY_SSL\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none\nfrom onyx.db.file_record import delete_filerecord_by_file_id\nfrom onyx.db.file_record import get_filerecord_by_file_id\nfrom onyx.db.file_record import get_filerecord_by_file_id_optional\nfrom onyx.db.file_record import get_filerecord_by_prefix\nfrom onyx.db.file_record import upsert_filerecord\nfrom onyx.db.models import FileRecord\nfrom onyx.db.models import FileRecord as FileStoreModel\nfrom onyx.file_store.s3_key_utils import generate_s3_key\nfrom onyx.utils.file import FileWithMimeType\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\nclass S3PutKwargs(TypedDict):\n ChecksumSHA256: NotRequired[str]\n\n\nclass FileStore(ABC):\n \"\"\"\n An abstraction for storing files and large binary objects.\n \"\"\"\n\n @abstractmethod\n def initialize(self) -> None:\n \"\"\"\n Should generally be called once before any other methods are called.\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def has_file(\n self,\n file_id: str,\n file_origin: FileOrigin,\n file_type: str,\n ) -> bool:\n \"\"\"\n Check if a file exists in the blob store\n\n Parameters:\n - file_id: Unique ID of the file to check for\n - file_origin: Origin of the file\n - file_type: Type of the file\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def save_file(\n self,\n content: IO,\n display_name: str | None,\n file_origin: FileOrigin,\n file_type: str,\n file_metadata: dict[str, Any] | None = None,\n file_id: str | None = None,\n ) -> str:\n \"\"\"\n Save a file to the blob store\n\n Parameters:\n - content: Contents of the file\n - display_name: Display name of the file to save\n - file_origin: Origin of the file\n - file_type: Type of the file\n - file_metadata: Additional metadata for the file\n - file_id: Unique ID of the file to save. If not provided, a random UUID will be generated.\n It is generally NOT recommended to provide this.\n\n Returns:\n The unique ID of the file that was saved.\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def read_file(\n self, file_id: str, mode: str | None = None, use_tempfile: bool = False\n ) -> IO[bytes]:\n \"\"\"\n Read the content of a given file by the ID\n\n Parameters:\n - file_id: Unique ID of file to read\n - mode: Mode to open the file (e.g. 'b' for binary)\n - use_tempfile: Whether to use a temporary file to store the contents\n in order to avoid loading the entire file into memory\n\n Returns:\n Contents of the file and metadata dict\n \"\"\"\n\n @abstractmethod\n def read_file_record(self, file_id: str) -> FileStoreModel:\n \"\"\"\n Read the file record by the ID\n \"\"\"\n\n @abstractmethod\n def get_file_size(\n self, file_id: str, db_session: Session | None = None\n ) -> int | None:\n \"\"\"\n Get the size of a file in bytes.\n Optionally provide a db_session for database access.\n \"\"\"\n\n @abstractmethod\n def delete_file(self, file_id: str) -> None:\n \"\"\"\n Delete a file by its ID.\n\n Parameters:\n - file_name: Name of file to delete\n \"\"\"\n\n @abstractmethod\n def get_file_with_mime_type(self, file_id: str) -> FileWithMimeType | None:\n \"\"\"\n Get the file + parse out the mime type.\n \"\"\"\n\n @abstractmethod\n def change_file_id(self, old_file_id: str, new_file_id: str) -> None:\n \"\"\"\n Change the file ID of an existing file.\n\n Parameters:\n - old_file_id: Current file ID\n - new_file_id: New file ID to assign\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def list_files_by_prefix(self, prefix: str) -> list[FileRecord]:\n \"\"\"\n List all file IDs that start with the given prefix.\n \"\"\"\n\n\nclass S3BackedFileStore(FileStore):\n \"\"\"Isn't necessarily S3, but is any S3-compatible storage (e.g. MinIO)\"\"\"\n\n def __init__(\n self,\n bucket_name: str,\n aws_access_key_id: str | None = None,\n aws_secret_access_key: str | None = None,\n aws_region_name: str | None = None,\n s3_endpoint_url: str | None = None,\n s3_prefix: str | None = None,\n s3_verify_ssl: bool = True,\n ) -> None:\n self._s3_client: S3Client | None = None\n self._bucket_name = bucket_name\n self._aws_access_key_id = aws_access_key_id\n self._aws_secret_access_key = aws_secret_access_key\n self._aws_region_name = aws_region_name or \"us-east-2\"\n self._s3_endpoint_url = s3_endpoint_url\n self._s3_prefix = s3_prefix or \"onyx-files\"\n self._s3_verify_ssl = s3_verify_ssl\n\n def _get_s3_client(self) -> S3Client:\n \"\"\"Initialize S3 client if not already done\"\"\"\n if self._s3_client is None:\n try:\n client_kwargs: dict[str, Any] = {\n \"service_name\": \"s3\",\n \"region_name\": self._aws_region_name,\n }\n\n # Add endpoint URL if specified (for MinIO, etc.)\n if self._s3_endpoint_url:\n client_kwargs[\"endpoint_url\"] = self._s3_endpoint_url\n client_kwargs[\"config\"] = Config(\n signature_version=\"s3v4\",\n s3={\"addressing_style\": \"path\"}, # Required for MinIO\n )\n # Disable SSL verification if requested (for local development)\n if not self._s3_verify_ssl:\n import urllib3\n\n urllib3.disable_warnings(\n urllib3.exceptions.InsecureRequestWarning\n )\n client_kwargs[\"verify\"] = False\n\n if self._aws_access_key_id and self._aws_secret_access_key:\n # Use explicit credentials\n client_kwargs.update(\n {\n \"aws_access_key_id\": self._aws_access_key_id,\n \"aws_secret_access_key\": self._aws_secret_access_key,\n }\n )\n self._s3_client = boto3.client(**client_kwargs)\n else:\n # Use IAM role or default credentials (not typically used with MinIO)\n self._s3_client = boto3.client(**client_kwargs)\n\n except Exception as e:\n logger.error(f\"Failed to initialize S3 client: {e}\")\n raise RuntimeError(f\"Failed to initialize S3 client: {e}\")\n\n return self._s3_client\n\n def _get_bucket_name(self) -> str:\n \"\"\"Get S3 bucket name from configuration\"\"\"\n if not self._bucket_name:\n raise RuntimeError(\"S3 bucket name is required for S3 file store\")\n return self._bucket_name\n\n def _get_s3_key(self, file_name: str) -> str:\n \"\"\"Generate S3 key from file name with tenant ID prefix\"\"\"\n tenant_id = get_current_tenant_id()\n\n s3_key = generate_s3_key(\n file_name=file_name,\n prefix=self._s3_prefix,\n tenant_id=tenant_id,\n max_key_length=1024,\n )\n\n # Log if truncation occurred (when the key is exactly at the limit)\n if len(s3_key) == 1024:\n logger.info(f\"File name was too long and was truncated: {file_name}\")\n\n return s3_key\n\n def initialize(self) -> None:\n \"\"\"Initialize the S3 file store by ensuring the bucket exists\"\"\"\n s3_client = self._get_s3_client()\n bucket_name = self._get_bucket_name()\n\n # Check if bucket exists\n try:\n s3_client.head_bucket(Bucket=bucket_name)\n logger.info(f\"S3 bucket '{bucket_name}' already exists\")\n except ClientError as e:\n error_code = e.response[\"Error\"][\"Code\"]\n if error_code == \"404\":\n # Bucket doesn't exist, create it\n logger.info(f\"Creating S3 bucket '{bucket_name}'\")\n\n # For AWS S3, we need to handle region-specific bucket creation\n region = (\n s3_client._client_config.region_name\n if hasattr(s3_client, \"_client_config\")\n else None\n )\n\n if region and region != \"us-east-1\":\n # For regions other than us-east-1, we need to specify LocationConstraint\n s3_client.create_bucket(\n Bucket=bucket_name,\n CreateBucketConfiguration={\"LocationConstraint\": region},\n )\n else:\n # For us-east-1 or MinIO/other S3-compatible services\n s3_client.create_bucket(Bucket=bucket_name)\n\n logger.info(f\"Successfully created S3 bucket '{bucket_name}'\")\n elif error_code == \"403\":\n # Bucket exists but we don't have permission to access it\n logger.warning(\n f\"S3 bucket '{bucket_name}' exists but access is forbidden\"\n )\n raise RuntimeError(\n f\"Access denied to S3 bucket '{bucket_name}'. Check credentials and permissions.\"\n )\n else:\n # Some other error occurred\n logger.error(f\"Failed to check S3 bucket '{bucket_name}': {e}\")\n raise RuntimeError(f\"Failed to check S3 bucket '{bucket_name}': {e}\")\n\n def has_file(\n self,\n file_id: str,\n file_origin: FileOrigin,\n file_type: str,\n db_session: Session | None = None,\n ) -> bool:\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n file_record = get_filerecord_by_file_id_optional(\n file_id=file_id, db_session=db_session\n )\n return (\n file_record is not None\n and file_record.file_origin == file_origin\n and file_record.file_type == file_type\n )\n\n def save_file(\n self,\n content: IO,\n display_name: str | None,\n file_origin: FileOrigin,\n file_type: str,\n file_metadata: dict[str, Any] | None = None,\n file_id: str | None = None,\n db_session: Session | None = None,\n ) -> str:\n if file_id is None:\n file_id = str(uuid.uuid4())\n\n s3_client = self._get_s3_client()\n bucket_name = self._get_bucket_name()\n s3_key = self._get_s3_key(file_id)\n\n hash256 = \"\"\n sha256_hash = hashlib.sha256()\n kwargs: S3PutKwargs = {}\n\n # FIX: Optimize checksum generation to avoid creating extra copies in memory\n # Read content from IO object\n if hasattr(content, \"read\"):\n file_content = content.read()\n if S3_GENERATE_LOCAL_CHECKSUM:\n # FIX: Don't convert to string first (creates unnecessary copy)\n # Work directly with bytes\n if isinstance(file_content, bytes):\n sha256_hash.update(file_content)\n else:\n sha256_hash.update(str(file_content).encode())\n hash256 = sha256_hash.hexdigest()\n kwargs[\"ChecksumSHA256\"] = hash256\n if hasattr(content, \"seek\"):\n content.seek(0) # Reset position for potential re-reads\n else:\n file_content = content\n\n # Upload to S3\n\n s3_client.put_object(\n Bucket=bucket_name,\n Key=s3_key,\n Body=file_content,\n ContentType=file_type,\n **kwargs,\n )\n\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n # Save metadata to database\n upsert_filerecord(\n file_id=file_id,\n display_name=display_name or file_id,\n file_origin=file_origin,\n file_type=file_type,\n bucket_name=bucket_name,\n object_key=s3_key,\n db_session=db_session,\n file_metadata=file_metadata,\n )\n db_session.commit()\n\n return file_id\n\n def read_file(\n self,\n file_id: str,\n mode: str | None = None, # noqa: ARG002\n use_tempfile: bool = False,\n db_session: Session | None = None,\n ) -> IO[bytes]:\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n file_record = get_filerecord_by_file_id(\n file_id=file_id, db_session=db_session\n )\n\n s3_client = self._get_s3_client()\n try:\n response = s3_client.get_object(\n Bucket=file_record.bucket_name, Key=file_record.object_key\n )\n except ClientError:\n logger.error(f\"Failed to read file {file_id} from S3\")\n raise\n\n # FIX: Stream file content instead of loading entire file into memory\n # This prevents OOM issues with large files (500MB+ PDFs, etc.)\n if use_tempfile:\n # Stream directly to temp file to avoid holding entire file in memory\n temp_file = tempfile.NamedTemporaryFile(mode=\"w+b\", delete=True)\n # Stream in 8MB chunks to reduce memory footprint\n for chunk in response[\"Body\"].iter_chunks(chunk_size=8 * 1024 * 1024):\n temp_file.write(chunk)\n temp_file.seek(0)\n return temp_file\n else:\n # For BytesIO, we still need to read into memory (legacy behavior)\n # but at least we're not creating duplicate copies\n file_content = response[\"Body\"].read()\n return BytesIO(file_content)\n\n def read_file_record(\n self, file_id: str, db_session: Session | None = None\n ) -> FileStoreModel:\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n file_record = get_filerecord_by_file_id(\n file_id=file_id, db_session=db_session\n )\n return file_record\n\n def get_file_size(\n self, file_id: str, db_session: Session | None = None\n ) -> int | None:\n \"\"\"\n Get the size of a file in bytes by querying S3 metadata.\n \"\"\"\n try:\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n file_record = get_filerecord_by_file_id(\n file_id=file_id, db_session=db_session\n )\n\n s3_client = self._get_s3_client()\n response = s3_client.head_object(\n Bucket=file_record.bucket_name, Key=file_record.object_key\n )\n return response.get(\"ContentLength\")\n except Exception as e:\n logger.warning(f\"Error getting file size for {file_id}: {e}\")\n return None\n\n def delete_file(self, file_id: str, db_session: Session | None = None) -> None:\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n try:\n file_record = get_filerecord_by_file_id(\n file_id=file_id, db_session=db_session\n )\n if not file_record.bucket_name:\n logger.error(\n f\"File record {file_id} with key {file_record.object_key} \"\n \"has no bucket name, cannot delete from filestore\"\n )\n delete_filerecord_by_file_id(file_id=file_id, db_session=db_session)\n db_session.commit()\n return\n\n # Delete from external storage\n s3_client = self._get_s3_client()\n try:\n s3_client.delete_object(\n Bucket=file_record.bucket_name, Key=file_record.object_key\n )\n except ClientError as e:\n # If the object doesn't exist in file store, treat it as success\n # since the end goal (object not existing) is achieved\n if e.response.get(\"Error\", {}).get(\"Code\") == \"NoSuchKey\":\n logger.warning(\n f\"delete_file: File {file_id} not found in file store (key: {file_record.object_key}), \"\n \"cleaning up database record.\"\n )\n else:\n raise\n\n # Delete metadata from database\n delete_filerecord_by_file_id(file_id=file_id, db_session=db_session)\n\n db_session.commit()\n\n except Exception:\n db_session.rollback()\n raise\n\n def change_file_id(\n self, old_file_id: str, new_file_id: str, db_session: Session | None = None\n ) -> None:\n with get_session_with_current_tenant_if_none(db_session) as db_session:\n try:\n # Get the existing file record\n old_file_record = get_filerecord_by_file_id(\n file_id=old_file_id, db_session=db_session\n )\n\n # Generate new S3 key for the new file ID\n new_s3_key = self._get_s3_key(new_file_id)\n\n # Copy S3 object to new key\n s3_client = self._get_s3_client()\n bucket_name = self._get_bucket_name()\n\n copy_source = (\n f\"{old_file_record.bucket_name}/{old_file_record.object_key}\"\n )\n\n s3_client.copy_object(\n CopySource=copy_source,\n Bucket=bucket_name,\n Key=new_s3_key,\n MetadataDirective=\"COPY\",\n )\n\n # Create new file record with new file_id\n # Cast file_metadata to the expected type\n file_metadata = cast(\n dict[Any, Any] | None, old_file_record.file_metadata\n )\n\n upsert_filerecord(\n file_id=new_file_id,\n display_name=old_file_record.display_name,\n file_origin=old_file_record.file_origin,\n file_type=old_file_record.file_type,\n bucket_name=bucket_name,\n object_key=new_s3_key,\n db_session=db_session,\n file_metadata=file_metadata,\n )\n\n # Delete old S3 object\n s3_client.delete_object(\n Bucket=old_file_record.bucket_name, Key=old_file_record.object_key\n )\n\n # Delete old file record\n delete_filerecord_by_file_id(file_id=old_file_id, db_session=db_session)\n\n db_session.commit()\n\n except Exception as e:\n db_session.rollback()\n logger.exception(\n f\"Failed to change file ID from {old_file_id} to {new_file_id}: {e}\"\n )\n raise\n\n def get_file_with_mime_type(self, file_id: str) -> FileWithMimeType | None:\n mime_type: str = \"application/octet-stream\"\n try:\n file_io = self.read_file(file_id, mode=\"b\")\n file_content = file_io.read()\n matches = puremagic.magic_string(file_content)\n if matches:\n mime_type = cast(str, matches[0].mime_type)\n return FileWithMimeType(data=file_content, mime_type=mime_type)\n except Exception:\n return None\n\n def list_files_by_prefix(self, prefix: str) -> list[FileRecord]:\n \"\"\"\n List all file IDs that start with the given prefix.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n file_records = get_filerecord_by_prefix(\n prefix=prefix, db_session=db_session\n )\n return file_records\n\n\ndef get_s3_file_store() -> S3BackedFileStore:\n \"\"\"\n Returns the S3 file store implementation.\n \"\"\"\n\n # Get bucket name - this is required\n bucket_name = S3_FILE_STORE_BUCKET_NAME\n if not bucket_name:\n raise RuntimeError(\n \"S3_FILE_STORE_BUCKET_NAME configuration is required for S3 file store\"\n )\n\n return S3BackedFileStore(\n bucket_name=bucket_name,\n aws_access_key_id=S3_AWS_ACCESS_KEY_ID,\n aws_secret_access_key=S3_AWS_SECRET_ACCESS_KEY,\n aws_region_name=AWS_REGION_NAME,\n s3_endpoint_url=S3_ENDPOINT_URL,\n s3_prefix=S3_FILE_STORE_PREFIX,\n s3_verify_ssl=S3_VERIFY_SSL,\n )\n\n\ndef get_default_file_store() -> FileStore:\n \"\"\"\n Returns the configured file store implementation based on FILE_STORE_BACKEND.\n\n When FILE_STORE_BACKEND=postgres (default):\n - Files are stored in PostgreSQL using Large Objects.\n - No external storage service (S3/MinIO) is required.\n\n When FILE_STORE_BACKEND=s3:\n - Supports AWS S3, MinIO, and other S3-compatible storage.\n - Configuration via environment variables:\n - S3_FILE_STORE_BUCKET_NAME, S3_ENDPOINT_URL, S3_AWS_ACCESS_KEY_ID, etc.\n \"\"\"\n from onyx.configs.app_configs import FILE_STORE_BACKEND\n from onyx.configs.constants import FileStoreType\n\n if FileStoreType(FILE_STORE_BACKEND) == FileStoreType.POSTGRES:\n from onyx.file_store.postgres_file_store import PostgresBackedFileStore\n\n return PostgresBackedFileStore()\n\n return get_s3_file_store()\n" }, { "path": "backend/onyx/file_store/models.py", "content": "import base64\nfrom enum import Enum\nfrom typing import NotRequired\nfrom typing_extensions import TypedDict # noreorder\n\nfrom pydantic import BaseModel\n\n\nclass ChatFileType(str, Enum):\n # Image types only contain the binary data\n IMAGE = \"image\"\n # Doc types are saved as both the binary, and the parsed text\n DOC = \"document\"\n # Plain text only contain the text\n PLAIN_TEXT = \"plain_text\"\n # Tabular data files (CSV, XLSX)\n TABULAR = \"tabular\"\n\n def is_text_file(self) -> bool:\n return self in (\n ChatFileType.PLAIN_TEXT,\n ChatFileType.DOC,\n ChatFileType.TABULAR,\n )\n\n def use_metadata_only(self) -> bool:\n \"\"\"File types where we can ignore the file content\n and only use the metadata.\"\"\"\n return self in (ChatFileType.TABULAR,)\n\n\nclass FileDescriptor(TypedDict):\n \"\"\"NOTE: is a `TypedDict` so it can be used as a type hint for a JSONB column\n in Postgres\"\"\"\n\n id: str\n type: ChatFileType\n name: NotRequired[str | None]\n user_file_id: NotRequired[str | None]\n\n\nclass InMemoryChatFile(BaseModel):\n file_id: str\n content: bytes\n file_type: ChatFileType\n filename: str | None = None\n\n def to_base64(self) -> str:\n if self.file_type == ChatFileType.IMAGE:\n return base64.b64encode(self.content).decode()\n else:\n raise RuntimeError(\n \"Should not be trying to convert a non-image file to base64\"\n )\n\n def to_file_descriptor(self) -> FileDescriptor:\n return {\n \"id\": str(self.file_id),\n \"type\": self.file_type,\n \"name\": self.filename,\n \"user_file_id\": str(self.file_id) if self.file_id else None,\n }\n" }, { "path": "backend/onyx/file_store/postgres_file_store.py", "content": "\"\"\"PostgreSQL-backed file store using Large Objects.\n\nStores file content directly in PostgreSQL via the Large Object facility,\neliminating the need for an external S3/MinIO service.\n\"\"\"\n\nimport tempfile\nimport uuid\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import cast\nfrom typing import IO\n\nimport puremagic\nfrom psycopg2.extensions import connection as Psycopg2Connection\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none\nfrom onyx.db.file_content import delete_file_content_by_file_id\nfrom onyx.db.file_content import get_file_content_by_file_id\nfrom onyx.db.file_content import get_file_content_by_file_id_optional\nfrom onyx.db.file_content import transfer_file_content_file_id\nfrom onyx.db.file_content import upsert_file_content\nfrom onyx.db.file_record import delete_filerecord_by_file_id\nfrom onyx.db.file_record import get_filerecord_by_file_id\nfrom onyx.db.file_record import get_filerecord_by_file_id_optional\nfrom onyx.db.file_record import get_filerecord_by_prefix\nfrom onyx.db.file_record import upsert_filerecord\nfrom onyx.db.models import FileRecord\nfrom onyx.db.models import FileRecord as FileStoreModel\nfrom onyx.file_store.file_store import FileStore\nfrom onyx.utils.file import FileWithMimeType\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nPOSTGRES_BUCKET_SENTINEL = \"postgres\"\nSTREAM_CHUNK_SIZE = 8 * 1024 * 1024 # 8 MB\n\n\ndef _get_raw_connection(db_session: Session) -> Psycopg2Connection:\n \"\"\"Extract the raw psycopg2 connection from a SQLAlchemy session.\"\"\"\n raw_conn = db_session.connection().connection.dbapi_connection\n if raw_conn is None:\n raise ValueError(\"Failed to get raw connection from session\")\n return cast(Psycopg2Connection, raw_conn)\n\n\ndef _create_large_object(raw_conn: Psycopg2Connection, data: bytes) -> int:\n \"\"\"Create a new Large Object, write data, and return the OID.\"\"\"\n lobj = raw_conn.lobject(0, \"wb\")\n lobj.write(data)\n oid: int = lobj.oid\n lobj.close()\n return oid\n\n\ndef _read_large_object(raw_conn: Psycopg2Connection, oid: int) -> bytes:\n \"\"\"Read all bytes from a Large Object.\"\"\"\n lobj = raw_conn.lobject(oid, \"rb\")\n data: bytes = lobj.read()\n lobj.close()\n return data\n\n\ndef _read_large_object_to_tempfile(raw_conn: Psycopg2Connection, oid: int) -> IO[bytes]:\n \"\"\"Stream a Large Object into a temporary file to avoid OOM on large files.\"\"\"\n lobj = raw_conn.lobject(oid, \"rb\")\n temp = tempfile.NamedTemporaryFile(mode=\"w+b\", delete=True)\n while True:\n chunk = lobj.read(STREAM_CHUNK_SIZE)\n if not chunk:\n break\n temp.write(chunk)\n lobj.close()\n temp.seek(0)\n return temp\n\n\ndef _delete_large_object(raw_conn: Any, oid: int) -> None:\n \"\"\"Unlink (delete) a Large Object by OID.\"\"\"\n lobj = raw_conn.lobject(oid, \"n\")\n lobj.unlink()\n\n\nclass PostgresBackedFileStore(FileStore):\n \"\"\"File store backed entirely by PostgreSQL.\n\n Metadata lives in `file_record`, content lives in PostgreSQL Large Objects\n with OID references tracked in `file_content`.\n \"\"\"\n\n def initialize(self) -> None:\n # Nothing to do — tables are created by Alembic migrations.\n pass\n\n def has_file(\n self,\n file_id: str,\n file_origin: FileOrigin,\n file_type: str,\n db_session: Session | None = None,\n ) -> bool:\n with get_session_with_current_tenant_if_none(db_session) as session:\n record = get_filerecord_by_file_id_optional(\n file_id=file_id, db_session=session\n )\n return (\n record is not None\n and record.file_origin == file_origin\n and record.file_type == file_type\n )\n\n def save_file(\n self,\n content: IO,\n display_name: str | None,\n file_origin: FileOrigin,\n file_type: str,\n file_metadata: dict[str, Any] | None = None,\n file_id: str | None = None,\n db_session: Session | None = None,\n ) -> str:\n if file_id is None:\n file_id = str(uuid.uuid4())\n\n file_bytes = self._read_content_bytes(content)\n created_lo = False\n\n with get_session_with_current_tenant_if_none(db_session) as session:\n raw_conn, oid = None, None\n try:\n raw_conn = _get_raw_connection(session)\n\n # Look up existing content so we can unlink the old\n # Large Object after a successful overwrite.\n existing = get_file_content_by_file_id_optional(\n file_id=file_id, db_session=session\n )\n old_oid = existing.lobj_oid if existing else None\n\n oid = _create_large_object(raw_conn, file_bytes)\n created_lo = True\n\n upsert_filerecord(\n file_id=file_id,\n display_name=display_name or file_id,\n file_origin=file_origin,\n file_type=file_type,\n bucket_name=POSTGRES_BUCKET_SENTINEL,\n object_key=str(oid),\n db_session=session,\n file_metadata=file_metadata,\n )\n upsert_file_content(\n file_id=file_id,\n lobj_oid=oid,\n file_size=len(file_bytes),\n db_session=session,\n )\n\n # Unlink the previous Large Object to avoid orphans\n if old_oid is not None and old_oid != oid:\n try:\n _delete_large_object(raw_conn, old_oid)\n except Exception:\n logger.warning(\n f\"Failed to unlink old large object {old_oid} for file {file_id}\"\n )\n\n session.commit()\n except Exception as e:\n session.rollback()\n try:\n if created_lo and raw_conn is not None and oid is not None:\n _delete_large_object(raw_conn, oid)\n except Exception:\n logger.exception(\n f\"Failed to delete large object {oid} for file {file_id}\"\n )\n raise e\n\n return file_id\n\n def read_file(\n self,\n file_id: str,\n mode: str | None = None, # noqa: ARG002\n use_tempfile: bool = False,\n db_session: Session | None = None,\n ) -> IO[bytes]:\n with get_session_with_current_tenant_if_none(db_session) as session:\n file_content = get_file_content_by_file_id(\n file_id=file_id, db_session=session\n )\n raw_conn = _get_raw_connection(session)\n\n if use_tempfile:\n return _read_large_object_to_tempfile(raw_conn, file_content.lobj_oid)\n\n data = _read_large_object(raw_conn, file_content.lobj_oid)\n return BytesIO(data)\n\n def read_file_record(\n self, file_id: str, db_session: Session | None = None\n ) -> FileStoreModel:\n with get_session_with_current_tenant_if_none(db_session) as session:\n return get_filerecord_by_file_id(file_id=file_id, db_session=session)\n\n def get_file_size(\n self, file_id: str, db_session: Session | None = None\n ) -> int | None:\n try:\n with get_session_with_current_tenant_if_none(db_session) as session:\n record = get_file_content_by_file_id(\n file_id=file_id, db_session=session\n )\n return record.file_size\n except Exception as e:\n logger.warning(f\"Error getting file size for {file_id}: {e}\")\n return None\n\n def delete_file(self, file_id: str, db_session: Session | None = None) -> None:\n with get_session_with_current_tenant_if_none(db_session) as session:\n try:\n file_content = get_file_content_by_file_id(\n file_id=file_id, db_session=session\n )\n raw_conn = _get_raw_connection(session)\n\n try:\n _delete_large_object(raw_conn, file_content.lobj_oid)\n except Exception:\n logger.warning(\n f\"Large object {file_content.lobj_oid} for file {file_id} not found, cleaning up records only.\"\n )\n\n delete_file_content_by_file_id(file_id=file_id, db_session=session)\n delete_filerecord_by_file_id(file_id=file_id, db_session=session)\n session.commit()\n except Exception:\n session.rollback()\n raise\n\n def get_file_with_mime_type(self, file_id: str) -> FileWithMimeType | None:\n mime_type = \"application/octet-stream\"\n try:\n file_io = self.read_file(file_id, mode=\"b\")\n except Exception:\n return None\n\n file_content = file_io.read()\n try:\n matches = puremagic.magic_string(file_content)\n if matches:\n mime_type = cast(str, matches[0].mime_type)\n except puremagic.PureError:\n pass\n\n return FileWithMimeType(data=file_content, mime_type=mime_type)\n\n def change_file_id(\n self, old_file_id: str, new_file_id: str, db_session: Session | None = None\n ) -> None:\n with get_session_with_current_tenant_if_none(db_session) as session:\n try:\n old_record = get_filerecord_by_file_id(\n file_id=old_file_id, db_session=session\n )\n file_metadata = cast(dict[Any, Any] | None, old_record.file_metadata)\n\n # 1. Create the new file_record so the FK target exists\n upsert_filerecord(\n file_id=new_file_id,\n display_name=old_record.display_name,\n file_origin=old_record.file_origin,\n file_type=old_record.file_type,\n bucket_name=POSTGRES_BUCKET_SENTINEL,\n object_key=old_record.object_key,\n db_session=session,\n file_metadata=file_metadata,\n )\n\n # 2. Move file_content in-place — the LO OID is never\n # shared between two rows.\n transfer_file_content_file_id(\n old_file_id=old_file_id,\n new_file_id=new_file_id,\n db_session=session,\n )\n\n # 3. Remove the now-orphaned old file_record\n delete_filerecord_by_file_id(file_id=old_file_id, db_session=session)\n\n session.commit()\n except Exception as e:\n session.rollback()\n logger.exception(\n f\"Failed to change file ID from {old_file_id} to {new_file_id}: {e}\"\n )\n raise\n\n def list_files_by_prefix(self, prefix: str) -> list[FileRecord]:\n with get_session_with_current_tenant() as session:\n return get_filerecord_by_prefix(prefix=prefix, db_session=session)\n\n @staticmethod\n def _read_content_bytes(content: IO) -> bytes:\n \"\"\"Normalize an IO object into raw bytes.\"\"\"\n if hasattr(content, \"read\"):\n raw = content.read()\n else:\n raw = content\n\n if isinstance(raw, str):\n return raw.encode(\"utf-8\")\n return raw\n" }, { "path": "backend/onyx/file_store/s3_key_utils.py", "content": "\"\"\"\nS3 key sanitization utilities for ensuring AWS S3 compatibility.\n\nThis module provides utilities for sanitizing file names to be compatible with\nAWS S3 object key naming guidelines while ensuring uniqueness when significant\nsanitization occurs.\n\nReference: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html\n\"\"\"\n\nimport hashlib\nimport re\nimport urllib.parse\nfrom re import Match\n\n# Constants for S3 key generation\nHASH_LENGTH = 64 # SHA256 hex digest length\nHASH_SEPARATOR_LENGTH = 1 # Length of underscore separator\nHASH_WITH_SEPARATOR_LENGTH = HASH_LENGTH + HASH_SEPARATOR_LENGTH\n\n\ndef _encode_special_char(match: Match[str]) -> str:\n \"\"\"Helper function to URL encode special characters.\"\"\"\n return urllib.parse.quote(match.group(0), safe=\"\")\n\n\ndef sanitize_s3_key_name(file_name: str) -> str:\n \"\"\"\n Sanitize file name to be S3-compatible according to AWS guidelines.\n\n This method:\n 1. Replaces problematic characters with safe alternatives\n 2. URL-encodes characters that might require special handling\n 3. Ensures the result is safe for S3 object keys\n 4. Adds uniqueness when significant sanitization occurs\n\n Args:\n file_name: The original file name to sanitize\n\n Returns:\n A sanitized file name that is S3-compatible\n\n Reference: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html\n \"\"\"\n if not file_name:\n return \"unnamed_file\"\n\n original_name = file_name\n\n # Characters to avoid completely (replace with underscore)\n # These are characters that AWS recommends avoiding\n avoid_chars = r'[\\\\{}^%`\\[\\]\"<>#|~/]'\n\n # Replace avoided characters with underscore\n sanitized = re.sub(avoid_chars, \"_\", file_name)\n # Characters that might require special handling but are allowed\n # We'll URL encode these to be safe\n special_chars = r\"[&$@=;:+,?\\s]\"\n\n sanitized = re.sub(special_chars, _encode_special_char, sanitized)\n\n # Handle non-ASCII characters by URL encoding them\n # This ensures Unicode characters are properly handled\n needs_unicode_encoding = False\n try:\n # Try to encode as ASCII to check if it contains non-ASCII chars\n sanitized.encode(\"ascii\")\n except UnicodeEncodeError:\n needs_unicode_encoding = True\n # Contains non-ASCII characters, URL encode the entire string\n # but preserve safe ASCII characters\n sanitized = urllib.parse.quote(\n sanitized,\n safe=\"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_.()!*\",\n )\n\n # Ensure we don't have consecutive periods at the start (relative path issue)\n sanitized = re.sub(r\"^\\.+\", \"\", sanitized)\n\n # Remove any trailing periods to avoid download issues\n sanitized = sanitized.rstrip(\".\")\n\n # Remove multiple separators\n sanitized = re.sub(r\"[-_]{2,}\", \"-\", sanitized)\n\n # If sanitization resulted in empty string, use a default\n if not sanitized:\n sanitized = \"sanitized_file\"\n\n # Check if significant sanitization occurred and add uniqueness if needed\n significant_changes = (\n # Check if we replaced many characters\n len(re.findall(avoid_chars, original_name)) > 3\n or\n # Check if we had to URL encode Unicode characters\n needs_unicode_encoding\n or\n # Check if the sanitized name is very different in length (expansion due to encoding)\n len(sanitized) > len(original_name) * 2\n or\n # Check if the original had many special characters\n len(re.findall(special_chars, original_name)) > 5\n )\n\n if significant_changes:\n # Add a short hash to ensure uniqueness while keeping some readability\n name_hash = hashlib.sha256(original_name.encode(\"utf-8\")).hexdigest()[:8]\n\n # Try to preserve file extension if it exists and is reasonable\n if \".\" in sanitized and len(sanitized.split(\".\")[-1]) <= 10:\n name_parts = sanitized.rsplit(\".\", 1)\n sanitized = f\"{name_parts[0]}_{name_hash}.{name_parts[1]}\"\n else:\n sanitized = f\"{sanitized}_{name_hash}\"\n\n return sanitized\n\n\ndef generate_s3_key(\n file_name: str, prefix: str, tenant_id: str, max_key_length: int = 1024\n) -> str:\n \"\"\"\n Generate a complete S3 key from file name with prefix and tenant ID.\n\n Args:\n file_name: The original file name\n prefix: S3 key prefix (e.g., 'onyx-files')\n tenant_id: Tenant identifier\n max_key_length: Maximum allowed S3 key length (default: 1024)\n\n Returns:\n A complete S3 key that fits within the length limit\n \"\"\"\n # Strip slashes from prefix and tenant_id to avoid double slashes\n prefix_clean = prefix.strip(\"/\")\n tenant_clean = tenant_id.strip(\"/\")\n\n # Sanitize the file name first\n sanitized_file_name = sanitize_s3_key_name(file_name)\n\n # Handle long file names that could exceed S3's key limit\n # S3 key format: {prefix}/{tenant_id}/{file_name}\n prefix_and_tenant_parts = [prefix_clean, tenant_clean]\n prefix_and_tenant = \"/\".join(prefix_and_tenant_parts) + \"/\"\n max_file_name_length = max_key_length - len(prefix_and_tenant)\n\n if len(sanitized_file_name) < max_file_name_length:\n return \"/\".join(prefix_and_tenant_parts + [sanitized_file_name])\n\n # For very long file names, use hash-based approach to ensure uniqueness\n # Use the original file name for the hash to maintain consistency\n file_hash = hashlib.sha256(file_name.encode(\"utf-8\")).hexdigest()\n\n # Calculate how much space we have for the readable part\n # Reserve space for hash (64 chars) + underscore separator (1 char)\n readable_part_max_length = max(0, max_file_name_length - HASH_WITH_SEPARATOR_LENGTH)\n\n if readable_part_max_length > 0:\n # Use first part of sanitized name + hash to maintain some readability\n readable_part = sanitized_file_name[:readable_part_max_length]\n truncated_name = f\"{readable_part}_{file_hash}\"\n else:\n # If no space for readable part, just use hash\n truncated_name = file_hash\n\n return \"/\".join(prefix_and_tenant_parts + [truncated_name])\n" }, { "path": "backend/onyx/file_store/utils.py", "content": "import base64\nfrom collections.abc import Callable\nfrom io import BytesIO\nfrom typing import cast\nfrom uuid import UUID\n\nimport requests\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.db.models import UserFile\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.file_store.models import ChatFileType\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.file_store.models import InMemoryChatFile\nfrom onyx.server.query_and_chat.chat_utils import mime_type_to_chat_file_type\nfrom onyx.utils.b64 import get_image_type\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom onyx.utils.timing import log_function_time\n\nlogger = setup_logger()\n\n\ndef plaintext_file_name_for_id(file_id: str) -> str:\n \"\"\"Generate a consistent file name for storing plaintext content of a file.\"\"\"\n return f\"plaintext_{file_id}\"\n\n\ndef store_plaintext(file_id: str, plaintext_content: str) -> bool:\n \"\"\"\n Store plaintext content for a file in the file store.\n\n Args:\n file_id: The ID of the file (user_file or artifact_file)\n plaintext_content: The plaintext content to store\n\n Returns:\n bool: True if storage was successful, False otherwise\n \"\"\"\n if not plaintext_content:\n return False\n\n plaintext_file_name = plaintext_file_name_for_id(file_id)\n try:\n file_store = get_default_file_store()\n file_content = BytesIO(plaintext_content.encode(\"utf-8\"))\n file_store.save_file(\n content=file_content,\n display_name=f\"Plaintext for {file_id}\",\n file_origin=FileOrigin.PLAINTEXT_CACHE,\n file_type=\"text/plain\",\n file_id=plaintext_file_name,\n )\n return True\n except Exception as e:\n logger.warning(f\"Failed to store plaintext for {file_id}: {e}\")\n return False\n\n\n# --- Convenience wrappers for callers that use user-file UUIDs ---\n\n\ndef user_file_id_to_plaintext_file_name(user_file_id: UUID) -> str:\n \"\"\"Generate a consistent file name for storing plaintext content of a user file.\"\"\"\n return plaintext_file_name_for_id(str(user_file_id))\n\n\ndef store_user_file_plaintext(user_file_id: UUID, plaintext_content: str) -> bool:\n \"\"\"Store plaintext content for a user file (delegates to :func:`store_plaintext`).\"\"\"\n return store_plaintext(str(user_file_id), plaintext_content)\n\n\ndef load_chat_file_by_id(file_id: str) -> InMemoryChatFile:\n \"\"\"Load a file directly from the file store using its file_record ID.\n\n This is the fallback path for chat-attached files that don't have a\n corresponding row in the ``user_file`` table.\"\"\"\n file_store = get_default_file_store()\n file_record = file_store.read_file_record(file_id)\n chat_file_type = mime_type_to_chat_file_type(file_record.file_type)\n\n file_io = file_store.read_file(file_id, mode=\"b\")\n return InMemoryChatFile(\n file_id=file_id,\n content=file_io.read(),\n file_type=chat_file_type,\n filename=file_record.display_name,\n )\n\n\ndef load_user_file(file_id: UUID, db_session: Session) -> InMemoryChatFile:\n status = \"not_loaded\"\n\n user_file = db_session.query(UserFile).filter(UserFile.id == file_id).first()\n if not user_file:\n raise ValueError(f\"User file with id {file_id} not found\")\n\n # Get the file record to determine the appropriate chat file type\n file_store = get_default_file_store()\n file_record = file_store.read_file_record(user_file.file_id)\n\n # Determine appropriate chat file type based on the original file's MIME type\n chat_file_type = mime_type_to_chat_file_type(file_record.file_type)\n\n # Try to load plaintext version first\n plaintext_file_name = user_file_id_to_plaintext_file_name(file_id)\n\n # check for plain text normalized version first, then use original file otherwise\n try:\n file_io = file_store.read_file(plaintext_file_name, mode=\"b\")\n # Metadata-only file types preserve their original type so\n # downstream injection paths can route them correctly.\n if chat_file_type.use_metadata_only():\n plaintext_chat_file_type = chat_file_type\n elif file_io is not None:\n # if we have plaintext for image (which happens when image\n # extraction is enabled), we use PLAIN_TEXT type\n plaintext_chat_file_type = ChatFileType.PLAIN_TEXT\n else:\n plaintext_chat_file_type = (\n ChatFileType.PLAIN_TEXT\n if chat_file_type != ChatFileType.IMAGE\n else chat_file_type\n )\n\n chat_file = InMemoryChatFile(\n file_id=str(user_file.file_id),\n content=file_io.read(),\n file_type=plaintext_chat_file_type,\n filename=user_file.name,\n )\n status = \"plaintext\"\n return chat_file\n except Exception as e:\n logger.warning(f\"Failed to load plaintext for user file {user_file.id}: {e}\")\n # Fall back to original file if plaintext not available\n file_io = file_store.read_file(user_file.file_id, mode=\"b\")\n\n chat_file = InMemoryChatFile(\n file_id=str(user_file.file_id),\n content=file_io.read(),\n file_type=chat_file_type,\n filename=user_file.name,\n )\n status = \"original\"\n return chat_file\n finally:\n logger.debug(\n f\"load_user_file finished: file_id={user_file.file_id} chat_file_type={chat_file_type} status={status}\"\n )\n\n\ndef load_in_memory_chat_files(\n user_file_ids: list[UUID],\n db_session: Session,\n) -> list[InMemoryChatFile]:\n \"\"\"\n Loads the actual content of user files specified by individual IDs and those\n within specified project IDs into memory.\n\n Args:\n user_file_ids: A list of specific UserFile IDs to load.\n db_session: The SQLAlchemy database session.\n\n Returns:\n A list of InMemoryChatFile objects, each containing the file content (as bytes),\n file ID, file type, and filename. Prioritizes loading plaintext versions if available.\n \"\"\"\n # Use parallel execution to load files concurrently\n return cast(\n list[InMemoryChatFile],\n run_functions_tuples_in_parallel(\n # 1. Load files specified by individual IDs\n [(load_user_file, (file_id, db_session)) for file_id in user_file_ids]\n ),\n )\n\n\ndef get_user_files(\n user_file_ids: list[UUID],\n db_session: Session,\n) -> list[UserFile]:\n \"\"\"\n Fetches UserFile database records based on provided file and project IDs.\n\n Args:\n user_file_ids: A list of specific UserFile IDs to fetch.\n db_session: The SQLAlchemy database session.\n\n Returns:\n A list containing UserFile SQLAlchemy model objects corresponding to the\n specified file IDs and all files within the specified project IDs.\n It does NOT return the actual file content.\n \"\"\"\n user_files: list[UserFile] = []\n\n # 1. Fetch UserFile records for specific file IDs\n for user_file_id in user_file_ids:\n # Query the database for a UserFile with the matching ID\n user_file = (\n db_session.query(UserFile).filter(UserFile.id == user_file_id).first()\n )\n # If found, add it to the list\n if user_file is not None:\n user_files.append(user_file)\n\n # 3. Return the combined list of UserFile database objects\n return user_files\n\n\ndef validate_user_files_ownership(\n user_file_ids: list[UUID],\n user_id: UUID | None,\n db_session: Session,\n) -> list[UserFile]:\n \"\"\"\n Fetches all UserFile database records for a given user.\n \"\"\"\n user_files = get_user_files(user_file_ids, db_session)\n current_user_files = []\n for user_file in user_files:\n # Note: if user_id is None, then all files should be None as well\n # (since auth must be disabled in this case)\n if user_file.user_id != user_id:\n raise ValueError(\n f\"User {user_id} does not have access to file {user_file.id}\"\n )\n current_user_files.append(user_file)\n\n return current_user_files\n\n\ndef save_file_from_url(url: str) -> str:\n response = requests.get(url)\n response.raise_for_status()\n\n file_io = BytesIO(response.content)\n file_store = get_default_file_store()\n file_id = file_store.save_file(\n content=file_io,\n display_name=\"GeneratedImage\",\n file_origin=FileOrigin.CHAT_IMAGE_GEN,\n file_type=\"image/png;base64\",\n )\n return file_id\n\n\ndef save_file_from_base64(base64_string: str) -> str:\n file_store = get_default_file_store()\n file_id = file_store.save_file(\n content=BytesIO(base64.b64decode(base64_string)),\n display_name=\"GeneratedImage\",\n file_origin=FileOrigin.CHAT_IMAGE_GEN,\n file_type=get_image_type(base64_string),\n )\n return file_id\n\n\ndef save_file(\n url: str | None = None,\n base64_data: str | None = None,\n) -> str:\n \"\"\"Save a file from either a URL or base64 encoded string.\n\n Args:\n url: URL to download file from\n base64_data: Base64 encoded file data\n\n Returns:\n The unique ID of the saved file\n\n Raises:\n ValueError: If neither url nor base64_data is provided, or if both are provided\n \"\"\"\n if url is not None and base64_data is not None:\n raise ValueError(\"Cannot specify both url and base64_data\")\n\n if url is not None:\n return save_file_from_url(url)\n elif base64_data is not None:\n return save_file_from_base64(base64_data)\n else:\n raise ValueError(\"Must specify either url or base64_data\")\n\n\ndef save_files(urls: list[str], base64_files: list[str]) -> list[str]:\n # NOTE: be explicit about typing so that if we change things, we get notified\n funcs: list[\n tuple[\n Callable[[str | None, str | None], str],\n tuple[str | None, str | None],\n ]\n ] = [(save_file, (url, None)) for url in urls] + [\n (save_file, (None, base64_file)) for base64_file in base64_files\n ]\n\n return run_functions_tuples_in_parallel(funcs)\n\n\n@log_function_time(print_only=True)\ndef verify_user_files(\n user_files: list[FileDescriptor],\n user_id: UUID | None,\n db_session: Session,\n project_id: int | None = None,\n) -> None:\n \"\"\"\n Verify that all provided file descriptors belong to the specified user.\n For project files (those without user_file_id), verifies access through project ownership.\n\n Args:\n user_files: List of file descriptors to verify\n user_id: The user ID to check ownership against\n db_session: The SQLAlchemy database session\n project_id: Optional project ID to verify project file access against\n\n Raises:\n ValueError: If any file does not belong to the user or is not found\n \"\"\"\n from onyx.db.models import Project__UserFile\n from onyx.db.projects import check_project_ownership\n\n # Extract user_file_ids and project file_ids from the file descriptors\n user_file_ids = []\n project_file_ids = []\n\n for file_descriptor in user_files:\n # Check if this file descriptor has a user_file_id\n if file_descriptor.get(\"user_file_id\"):\n try:\n user_file_ids.append(UUID(file_descriptor[\"user_file_id\"]))\n except (ValueError, TypeError):\n logger.warning(\n f\"Invalid user_file_id in file descriptor: {file_descriptor['user_file_id']}\"\n )\n continue\n else:\n # This is a project file - use the 'id' field which is the file_id\n if file_descriptor.get(\"id\"):\n project_file_ids.append(file_descriptor[\"id\"])\n\n # Verify user files (existing logic)\n if user_file_ids:\n validate_user_files_ownership(user_file_ids, user_id, db_session)\n\n # Verify project files\n if project_file_ids:\n if project_id is None:\n raise ValueError(\n \"Project files provided but no project_id specified for verification\"\n )\n\n # Verify user owns the project\n if not check_project_ownership(project_id, user_id, db_session):\n raise ValueError(\n f\"User {user_id} does not have access to project {project_id}\"\n )\n\n # Verify all project files belong to the specified project\n user_files_in_project = (\n db_session.query(UserFile)\n .join(Project__UserFile)\n .filter(\n Project__UserFile.project_id == project_id,\n UserFile.file_id.in_(project_file_ids),\n )\n .all()\n )\n\n # Check if all files were found in the project\n found_file_ids = {uf.file_id for uf in user_files_in_project}\n missing_files = set(project_file_ids) - found_file_ids\n\n if missing_files:\n raise ValueError(\n f\"Files {missing_files} are not associated with project {project_id}\"\n )\n\n\ndef build_frontend_file_url(file_id: str) -> str:\n return f\"/api/chat/file/{file_id}\"\n\n\ndef build_full_frontend_file_url(file_id: str) -> str:\n return f\"{WEB_DOMAIN}/api/chat/file/{file_id}\"\n" }, { "path": "backend/onyx/hooks/__init__.py", "content": "" }, { "path": "backend/onyx/hooks/api_dependencies.py", "content": "from onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom shared_configs.configs import MULTI_TENANT\n\n\ndef require_hook_enabled() -> None:\n \"\"\"FastAPI dependency that gates all hook management endpoints.\n\n Hooks are only available in single-tenant / self-hosted EE deployments.\n\n Use as: Depends(require_hook_enabled)\n \"\"\"\n if MULTI_TENANT:\n raise OnyxError(\n OnyxErrorCode.SINGLE_TENANT_ONLY,\n \"Hooks are not available in multi-tenant deployments\",\n )\n" }, { "path": "backend/onyx/hooks/executor.py", "content": "\"\"\"CE hook executor.\n\nHookSkipped and HookSoftFailed are real classes kept here because\nprocess_message.py (CE code) uses isinstance checks against them.\n\nexecute_hook is the public entry point. It dispatches to _execute_hook_impl\nvia fetch_versioned_implementation so that:\n - CE: onyx.hooks.executor._execute_hook_impl → no-op, returns HookSkipped()\n - EE: ee.onyx.hooks.executor._execute_hook_impl → real HTTP call\n\"\"\"\n\nfrom typing import Any\nfrom typing import TypeVar\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import HookPoint\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\n\nclass HookSkipped:\n \"\"\"No active hook configured for this hook point.\"\"\"\n\n\nclass HookSoftFailed:\n \"\"\"Hook was called but failed with SOFT fail strategy — continuing.\"\"\"\n\n\nT = TypeVar(\"T\", bound=BaseModel)\n\n\ndef _execute_hook_impl(\n *,\n db_session: Session, # noqa: ARG001\n hook_point: HookPoint, # noqa: ARG001\n payload: dict[str, Any], # noqa: ARG001\n response_type: type[T], # noqa: ARG001\n) -> T | HookSkipped | HookSoftFailed:\n \"\"\"CE no-op — hooks are not available without EE.\"\"\"\n return HookSkipped()\n\n\ndef execute_hook(\n *,\n db_session: Session,\n hook_point: HookPoint,\n payload: dict[str, Any],\n response_type: type[T],\n) -> T | HookSkipped | HookSoftFailed:\n \"\"\"Execute the hook for the given hook point.\n\n Dispatches to the versioned implementation so EE gets the real executor\n and CE gets the no-op stub, without any changes at the call site.\n \"\"\"\n impl = fetch_versioned_implementation(\"onyx.hooks.executor\", \"_execute_hook_impl\")\n return impl(\n db_session=db_session,\n hook_point=hook_point,\n payload=payload,\n response_type=response_type,\n )\n" }, { "path": "backend/onyx/hooks/models.py", "content": "from datetime import datetime\nfrom enum import Enum\nfrom typing import Annotated\nfrom typing import Any\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\nfrom pydantic import field_validator\nfrom pydantic import model_validator\nfrom pydantic import SecretStr\n\nfrom onyx.db.enums import HookFailStrategy\nfrom onyx.db.enums import HookPoint\n\nNonEmptySecretStr = Annotated[SecretStr, Field(min_length=1)]\n\n\n# ---------------------------------------------------------------------------\n# Request models\n# ---------------------------------------------------------------------------\n\n\nclass HookCreateRequest(BaseModel):\n name: str = Field(min_length=1)\n hook_point: HookPoint\n endpoint_url: str = Field(min_length=1)\n api_key: NonEmptySecretStr | None = None\n fail_strategy: HookFailStrategy | None = None # if None, uses HookPointSpec default\n timeout_seconds: float | None = Field(\n default=None, gt=0\n ) # if None, uses HookPointSpec default\n\n @field_validator(\"name\", \"endpoint_url\")\n @classmethod\n def no_whitespace_only(cls, v: str) -> str:\n if not v.strip():\n raise ValueError(\"cannot be whitespace-only.\")\n return v\n\n\nclass HookUpdateRequest(BaseModel):\n name: str | None = None\n endpoint_url: str | None = None\n api_key: NonEmptySecretStr | None = None\n fail_strategy: HookFailStrategy | None = None\n timeout_seconds: float | None = Field(default=None, gt=0)\n\n @model_validator(mode=\"after\")\n def require_at_least_one_field(self) -> \"HookUpdateRequest\":\n if not self.model_fields_set:\n raise ValueError(\"At least one field must be provided for an update.\")\n if \"name\" in self.model_fields_set and not (self.name or \"\").strip():\n raise ValueError(\"name cannot be cleared.\")\n if (\n \"endpoint_url\" in self.model_fields_set\n and not (self.endpoint_url or \"\").strip()\n ):\n raise ValueError(\"endpoint_url cannot be cleared.\")\n if \"fail_strategy\" in self.model_fields_set and self.fail_strategy is None:\n raise ValueError(\n \"fail_strategy cannot be null; omit the field to leave it unchanged.\"\n )\n if \"timeout_seconds\" in self.model_fields_set and self.timeout_seconds is None:\n raise ValueError(\n \"timeout_seconds cannot be null; omit the field to leave it unchanged.\"\n )\n return self\n\n\n# ---------------------------------------------------------------------------\n# Response models\n# ---------------------------------------------------------------------------\n\n\nclass HookPointMetaResponse(BaseModel):\n hook_point: HookPoint\n display_name: str\n description: str\n docs_url: str | None\n input_schema: dict[str, Any]\n output_schema: dict[str, Any]\n default_timeout_seconds: float\n default_fail_strategy: HookFailStrategy\n fail_hard_description: str\n\n\nclass HookResponse(BaseModel):\n id: int\n name: str\n hook_point: HookPoint\n # Nullable to match the DB column — endpoint_url is required on creation but\n # future hook point types may not use an external endpoint (e.g. built-in handlers).\n endpoint_url: str | None\n # Partially-masked API key (e.g. \"abcd••••••••wxyz\"), or None if no key is set.\n api_key_masked: str | None\n fail_strategy: HookFailStrategy\n timeout_seconds: float # always resolved — None from request is replaced with spec default before DB write\n is_active: bool\n is_reachable: bool | None\n creator_email: str | None\n created_at: datetime\n updated_at: datetime\n\n\nclass HookValidateStatus(str, Enum):\n passed = \"passed\" # server responded (any status except 401/403)\n auth_failed = \"auth_failed\" # server responded with 401 or 403\n timeout = (\n \"timeout\" # TCP connected, but read/write timed out (server exists but slow)\n )\n cannot_connect = \"cannot_connect\" # could not connect to the server\n\n\nclass HookValidateResponse(BaseModel):\n status: HookValidateStatus\n error_message: str | None = None\n\n\nclass HookExecutionRecord(BaseModel):\n error_message: str | None = None\n status_code: int | None = None\n duration_ms: int | None = None\n created_at: datetime\n" }, { "path": "backend/onyx/hooks/points/__init__.py", "content": "" }, { "path": "backend/onyx/hooks/points/base.py", "content": "from typing import Any\nfrom typing import ClassVar\n\nfrom pydantic import BaseModel\n\nfrom onyx.db.enums import HookFailStrategy\nfrom onyx.db.enums import HookPoint\n\n\n_REQUIRED_ATTRS = (\n \"hook_point\",\n \"display_name\",\n \"description\",\n \"default_timeout_seconds\",\n \"fail_hard_description\",\n \"default_fail_strategy\",\n \"payload_model\",\n \"response_model\",\n)\n\n\nclass HookPointSpec:\n \"\"\"Static metadata and contract for a pipeline hook point.\n\n Each concrete subclass represents exactly one hook point and is instantiated\n once at startup, registered in onyx.hooks.registry._REGISTRY. Prefer\n get_hook_point_spec() or get_all_specs() from the registry over direct\n instantiation.\n\n Each hook point is a concrete subclass of this class. Onyx engineers\n own these definitions — customers never touch this code.\n\n Subclasses must define all attributes as class-level constants.\n payload_model and response_model must be Pydantic BaseModel subclasses;\n input_schema and output_schema are derived from them automatically.\n \"\"\"\n\n hook_point: HookPoint\n display_name: str\n description: str\n default_timeout_seconds: float\n fail_hard_description: str\n default_fail_strategy: HookFailStrategy\n docs_url: str | None = None\n\n payload_model: ClassVar[type[BaseModel]]\n response_model: ClassVar[type[BaseModel]]\n\n # Computed once at class definition time from payload_model / response_model.\n input_schema: ClassVar[dict[str, Any]]\n output_schema: ClassVar[dict[str, Any]]\n\n def __init_subclass__(cls, **kwargs: object) -> None:\n \"\"\"Enforce that every subclass declares all required class attributes.\n\n Called automatically by Python whenever a class inherits from HookPointSpec.\n Raises TypeError at import time if any required attribute is missing or if\n payload_model / response_model are not Pydantic BaseModel subclasses.\n input_schema and output_schema are derived automatically from the models.\n \"\"\"\n super().__init_subclass__(**kwargs)\n missing = [attr for attr in _REQUIRED_ATTRS if not hasattr(cls, attr)]\n if missing:\n raise TypeError(f\"{cls.__name__} must define class attributes: {missing}\")\n for attr in (\"payload_model\", \"response_model\"):\n val = getattr(cls, attr, None)\n if val is None or not (\n isinstance(val, type) and issubclass(val, BaseModel)\n ):\n raise TypeError(\n f\"{cls.__name__}.{attr} must be a Pydantic BaseModel subclass, got {val!r}\"\n )\n cls.input_schema = cls.payload_model.model_json_schema()\n cls.output_schema = cls.response_model.model_json_schema()\n" }, { "path": "backend/onyx/hooks/points/document_ingestion.py", "content": "from pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.db.enums import HookFailStrategy\nfrom onyx.db.enums import HookPoint\nfrom onyx.hooks.points.base import HookPointSpec\n\n\nclass DocumentIngestionSection(BaseModel):\n \"\"\"Represents a single section of a document — either text or image, not both.\n\n Text section: set `text`, leave `image_file_id` null.\n Image section: set `image_file_id`, leave `text` null.\n \"\"\"\n\n text: str | None = Field(\n default=None,\n description=\"Text content of this section. Set for text sections, null for image sections.\",\n )\n link: str | None = Field(\n default=None,\n description=\"Optional URL associated with this section. Preserve the original link from the payload if you want it retained.\",\n )\n image_file_id: str | None = Field(\n default=None,\n description=(\n \"Opaque identifier for an image stored in the file store. \"\n \"The image content is not included — this field signals that the section is an image. \"\n \"Hooks can use its presence to reorder or drop image sections, but cannot read or modify the image itself.\"\n ),\n )\n\n\nclass DocumentIngestionOwner(BaseModel):\n display_name: str | None = Field(\n default=None,\n description=\"Human-readable name of the owner.\",\n )\n email: str | None = Field(\n default=None,\n description=\"Email address of the owner.\",\n )\n\n\nclass DocumentIngestionPayload(BaseModel):\n document_id: str = Field(\n description=\"Unique identifier for the document. Read-only — changes are ignored.\"\n )\n title: str | None = Field(description=\"Title of the document.\")\n semantic_identifier: str = Field(\n description=\"Human-readable identifier used for display (e.g. file name, page title).\"\n )\n source: str = Field(\n description=(\n \"Connector source type (e.g. confluence, slack, google_drive). \"\n \"Read-only — changes are ignored. \"\n \"Full list of values: https://github.com/onyx-dot-app/onyx/blob/main/backend/onyx/configs/constants.py#L195\"\n )\n )\n sections: list[DocumentIngestionSection] = Field(\n description=\"Sections of the document. Includes both text sections (text set, image_file_id null) and image sections (image_file_id set, text null).\"\n )\n metadata: dict[str, list[str]] = Field(\n description=\"Key-value metadata attached to the document. Values are always a list of strings.\"\n )\n doc_updated_at: str | None = Field(\n description=\"ISO 8601 UTC timestamp of the last update at the source, or null if unknown. Example: '2024-03-15T10:30:00+00:00'.\"\n )\n primary_owners: list[DocumentIngestionOwner] | None = Field(\n description=\"Primary owners of the document, or null if not available.\"\n )\n secondary_owners: list[DocumentIngestionOwner] | None = Field(\n description=\"Secondary owners of the document, or null if not available.\"\n )\n\n\nclass DocumentIngestionResponse(BaseModel):\n # Intentionally permissive — customer endpoints may return extra fields.\n sections: list[DocumentIngestionSection] | None = Field(\n description=\"The sections to index, in the desired order. Reorder, drop, or modify sections freely. Null or empty list drops the document.\"\n )\n rejection_reason: str | None = Field(\n default=None,\n description=\"Logged when sections is null or empty. Falls back to a generic message if omitted.\",\n )\n\n\nclass DocumentIngestionSpec(HookPointSpec):\n \"\"\"Hook point that runs on every document before it enters the indexing pipeline.\n\n Call site: immediately after Onyx's internal validation and before the\n indexing pipeline begins — no partial writes have occurred yet.\n\n If a Document Ingestion hook is configured, it takes precedence —\n Document Ingestion Light will not run. Configure only one per deployment.\n\n Supported use cases:\n - Document filtering: drop documents based on content or metadata\n - Content rewriting: redact PII or normalize text before indexing\n \"\"\"\n\n hook_point = HookPoint.DOCUMENT_INGESTION\n display_name = \"Document Ingestion\"\n description = (\n \"Runs on every document before it enters the indexing pipeline. \"\n \"Allows filtering, rewriting, or dropping documents.\"\n )\n default_timeout_seconds = 30.0\n fail_hard_description = \"The document will not be indexed.\"\n default_fail_strategy = HookFailStrategy.HARD\n docs_url = \"https://docs.onyx.app/admins/advanced_configs/hook_extensions#document-ingestion\"\n\n payload_model = DocumentIngestionPayload\n response_model = DocumentIngestionResponse\n" }, { "path": "backend/onyx/hooks/points/query_processing.py", "content": "from pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import Field\n\nfrom onyx.db.enums import HookFailStrategy\nfrom onyx.db.enums import HookPoint\nfrom onyx.hooks.points.base import HookPointSpec\n\n\nclass QueryProcessingPayload(BaseModel):\n model_config = ConfigDict(extra=\"forbid\")\n\n query: str = Field(description=\"The raw query string exactly as the user typed it.\")\n user_email: str | None = Field(\n description=\"Email of the user submitting the query, or null if unauthenticated.\"\n )\n chat_session_id: str = Field(\n description=\"UUID of the chat session, formatted as a hyphenated lowercase string (e.g. '550e8400-e29b-41d4-a716-446655440000'). Always present — the session is guaranteed to exist by the time this hook fires.\"\n )\n\n\nclass QueryProcessingResponse(BaseModel):\n # Intentionally permissive — customer endpoints may return extra fields.\n query: str | None = Field(\n default=None,\n description=(\n \"The query to use in the pipeline. \"\n \"Null, empty string, whitespace-only, or absent = reject the query.\"\n ),\n )\n rejection_message: str | None = Field(\n default=None,\n description=\"Message shown to the user when the query is rejected. Falls back to a generic message if not provided.\",\n )\n\n\nclass QueryProcessingSpec(HookPointSpec):\n \"\"\"Hook point that runs on every user query before it enters the pipeline.\n\n Call site: inside handle_stream_message_objects() in\n backend/onyx/chat/process_message.py, immediately after message_text is\n assigned from the request and before create_new_chat_message() saves it.\n\n This is the earliest possible point in the query pipeline:\n - Raw query — unmodified, exactly as the user typed it\n - No side effects yet — message has not been saved to DB\n - User identity is available for user-specific logic\n\n Supported use cases:\n - Query rejection: block queries based on content or user context\n - Query rewriting: normalize, expand, or modify the query\n - PII removal: scrub sensitive data before the LLM sees it\n - Access control: reject queries from certain users or groups\n - Query auditing: log or track queries based on business rules\n \"\"\"\n\n hook_point = HookPoint.QUERY_PROCESSING\n display_name = \"Query Processing\"\n description = (\n \"Runs on every user query before it enters the pipeline. \"\n \"Allows rewriting, filtering, or rejecting queries.\"\n )\n default_timeout_seconds = 5.0 # user is actively waiting — keep tight\n fail_hard_description = (\n \"The query will be blocked and the user will see an error message.\"\n )\n default_fail_strategy = HookFailStrategy.HARD\n docs_url = (\n \"https://docs.onyx.app/admins/advanced_configs/hook_extensions#query-processing\"\n )\n\n payload_model = QueryProcessingPayload\n response_model = QueryProcessingResponse\n" }, { "path": "backend/onyx/hooks/registry.py", "content": "from onyx.db.enums import HookPoint\nfrom onyx.hooks.points.base import HookPointSpec\nfrom onyx.hooks.points.document_ingestion import DocumentIngestionSpec\nfrom onyx.hooks.points.query_processing import QueryProcessingSpec\n\n# Internal: use `monkeypatch.setattr(registry_module, \"_REGISTRY\", {...})` to override in tests.\n_REGISTRY: dict[HookPoint, HookPointSpec] = {\n HookPoint.DOCUMENT_INGESTION: DocumentIngestionSpec(),\n HookPoint.QUERY_PROCESSING: QueryProcessingSpec(),\n}\n\n\ndef validate_registry() -> None:\n \"\"\"Assert that every HookPoint enum value has a registered spec.\n\n Call once at application startup (e.g. from the FastAPI lifespan hook).\n Raises RuntimeError if any hook point is missing a spec.\n \"\"\"\n missing = set(HookPoint) - set(_REGISTRY)\n if missing:\n raise RuntimeError(\n f\"Hook point(s) have no registered spec: {missing}. \"\n \"Add an entry to onyx.hooks.registry._REGISTRY.\"\n )\n\n\ndef get_hook_point_spec(hook_point: HookPoint) -> HookPointSpec:\n \"\"\"Returns the spec for a given hook point.\n\n Raises ValueError if the hook point has no registered spec — this is a\n programmer error; every HookPoint enum value must have a corresponding spec\n in _REGISTRY.\n \"\"\"\n try:\n return _REGISTRY[hook_point]\n except KeyError:\n raise ValueError(\n f\"No spec registered for hook point {hook_point!r}. \"\n \"Add an entry to onyx.hooks.registry._REGISTRY.\"\n )\n\n\ndef get_all_specs() -> list[HookPointSpec]:\n \"\"\"Returns the specs for all registered hook points.\"\"\"\n return list(_REGISTRY.values())\n" }, { "path": "backend/onyx/httpx/httpx_pool.py", "content": "import threading\nfrom typing import Any\n\nimport httpx\n\n\ndef make_default_kwargs() -> dict[str, Any]:\n return {\n \"http2\": True,\n \"limits\": httpx.Limits(),\n }\n\n\nclass HttpxPool:\n \"\"\"Class to manage a global httpx Client instance\"\"\"\n\n _clients: dict[str, httpx.Client] = {}\n _lock: threading.Lock = threading.Lock()\n\n # Default parameters for creation\n\n def __init__(self) -> None:\n pass\n\n @classmethod\n def _init_client(cls, **kwargs: Any) -> httpx.Client:\n \"\"\"Private helper method to create and return an httpx.Client.\"\"\"\n merged_kwargs = {**(make_default_kwargs()), **kwargs}\n return httpx.Client(**merged_kwargs)\n\n @classmethod\n def init_client(cls, name: str, **kwargs: Any) -> None:\n \"\"\"Allow the caller to init the client with extra params.\"\"\"\n with cls._lock:\n if name not in cls._clients:\n cls._clients[name] = cls._init_client(**kwargs)\n\n @classmethod\n def close_client(cls, name: str) -> None:\n \"\"\"Allow the caller to close the client.\"\"\"\n with cls._lock:\n client = cls._clients.pop(name, None)\n if client:\n client.close()\n\n @classmethod\n def close_all(cls) -> None:\n \"\"\"Close all registered clients.\"\"\"\n with cls._lock:\n for client in cls._clients.values():\n client.close()\n cls._clients.clear()\n\n @classmethod\n def get(cls, name: str) -> httpx.Client:\n \"\"\"Gets the httpx.Client. Will init to default settings if not init'd.\"\"\"\n with cls._lock:\n if name not in cls._clients:\n cls._clients[name] = cls._init_client()\n return cls._clients[name]\n" }, { "path": "backend/onyx/image_gen/__init__.py", "content": "" }, { "path": "backend/onyx/image_gen/exceptions.py", "content": "class ImageProviderError(Exception):\n pass\n\n\nclass ImageProviderCredentialsError(ImageProviderError):\n pass\n" }, { "path": "backend/onyx/image_gen/factory.py", "content": "from enum import Enum\n\nfrom onyx.image_gen.interfaces import ImageGenerationProvider\nfrom onyx.image_gen.interfaces import ImageGenerationProviderCredentials\nfrom onyx.image_gen.providers.azure_img_gen import AzureImageGenerationProvider\nfrom onyx.image_gen.providers.openai_img_gen import OpenAIImageGenerationProvider\nfrom onyx.image_gen.providers.vertex_img_gen import VertexImageGenerationProvider\n\n\nclass ImageGenerationProviderName(str, Enum):\n AZURE = \"azure\"\n OPENAI = \"openai\"\n VERTEX_AI = \"vertex_ai\"\n\n\nPROVIDERS: dict[ImageGenerationProviderName, type[ImageGenerationProvider]] = {\n ImageGenerationProviderName.AZURE: AzureImageGenerationProvider,\n ImageGenerationProviderName.OPENAI: OpenAIImageGenerationProvider,\n ImageGenerationProviderName.VERTEX_AI: VertexImageGenerationProvider,\n}\n\n\ndef get_image_generation_provider(\n provider: str,\n credentials: ImageGenerationProviderCredentials,\n) -> ImageGenerationProvider:\n provider_cls = _get_provider_cls(provider)\n return provider_cls.build_from_credentials(credentials)\n\n\ndef validate_credentials(\n provider: str,\n credentials: ImageGenerationProviderCredentials,\n) -> bool:\n provider_cls = _get_provider_cls(provider)\n return provider_cls.validate_credentials(credentials)\n\n\ndef _get_provider_cls(provider: str) -> type[ImageGenerationProvider]:\n try:\n provider_enum = ImageGenerationProviderName(provider)\n except ValueError:\n raise ValueError(f\"Invalid image generation provider: {provider}\")\n return PROVIDERS[provider_enum]\n" }, { "path": "backend/onyx/image_gen/interfaces.py", "content": "from __future__ import annotations\n\nimport abc\nfrom typing import Any\nfrom typing import TYPE_CHECKING\n\nfrom pydantic import BaseModel\n\nfrom onyx.image_gen.exceptions import ImageProviderCredentialsError\n\nif TYPE_CHECKING:\n from litellm.types.utils import ImageResponse as ImageGenerationResponse\n\n\nclass ImageGenerationProviderCredentials(BaseModel):\n api_key: str | None = None\n api_base: str | None = None\n api_version: str | None = None\n deployment_name: str | None = None\n custom_config: dict[str, str] | None = None\n\n\nclass ReferenceImage(BaseModel):\n data: bytes\n mime_type: str\n\n\nclass ImageGenerationProvider(abc.ABC):\n @property\n def supports_reference_images(self) -> bool:\n return False\n\n @property\n def max_reference_images(self) -> int:\n return 0\n\n @classmethod\n @abc.abstractmethod\n def validate_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> bool:\n \"\"\"Returns true if sufficient credentials are given to build this provider.\"\"\"\n raise NotImplementedError(\"validate_credentials not implemented\")\n\n @classmethod\n def build_from_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> ImageGenerationProvider:\n if not cls.validate_credentials(credentials):\n raise ImageProviderCredentialsError(\n f\"Invalid image generation credentials: {credentials}\"\n )\n return cls._build_from_credentials(credentials)\n\n @classmethod\n @abc.abstractmethod\n def _build_from_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> ImageGenerationProvider:\n \"\"\"\n Given credentials, builds an instance of the provider.\n Should NOT be called directly - use build_from_credentials instead.\n\n AssertionError if credentials are invalid.\n \"\"\"\n raise NotImplementedError(\"build_from_credentials not implemented\")\n\n @abc.abstractmethod\n def generate_image(\n self,\n prompt: str,\n model: str,\n size: str,\n n: int,\n quality: str | None = None,\n reference_images: list[ReferenceImage] | None = None,\n **kwargs: Any,\n ) -> ImageGenerationResponse:\n \"\"\"Generates an image based on a prompt.\"\"\"\n raise NotImplementedError(\"generate_image not implemented\")\n" }, { "path": "backend/onyx/image_gen/providers/azure_img_gen.py", "content": "from __future__ import annotations\n\nfrom typing import Any\nfrom typing import TYPE_CHECKING\n\nfrom onyx.image_gen.interfaces import ImageGenerationProvider\nfrom onyx.image_gen.interfaces import ImageGenerationProviderCredentials\nfrom onyx.image_gen.interfaces import ReferenceImage\n\nif TYPE_CHECKING:\n from onyx.image_gen.interfaces import ImageGenerationResponse\n\n\nclass AzureImageGenerationProvider(ImageGenerationProvider):\n _GPT_IMAGE_MODEL_PREFIX = \"gpt-image-\"\n _DALL_E_2_MODEL_NAME = \"dall-e-2\"\n\n def __init__(\n self,\n api_key: str,\n api_base: str,\n api_version: str,\n deployment_name: str | None = None,\n ):\n self._api_key = api_key\n self._api_base = api_base\n self._api_version = api_version\n self._deployment_name = deployment_name\n\n @classmethod\n def validate_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> bool:\n return all(\n [\n credentials.api_key,\n credentials.api_base,\n credentials.api_version,\n ]\n )\n\n @classmethod\n def _build_from_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> AzureImageGenerationProvider:\n assert credentials.api_key\n assert credentials.api_base\n assert credentials.api_version\n\n return cls(\n api_key=credentials.api_key,\n api_base=credentials.api_base,\n api_version=credentials.api_version,\n deployment_name=credentials.deployment_name,\n )\n\n @property\n def supports_reference_images(self) -> bool:\n return True\n\n @property\n def max_reference_images(self) -> int:\n # Azure GPT image models support up to 16 input images for edits.\n return 16\n\n def _normalize_model_name(self, model: str) -> str:\n return model.rsplit(\"/\", 1)[-1]\n\n def _model_supports_image_edits(self, model: str) -> bool:\n normalized_model = self._normalize_model_name(model)\n return (\n normalized_model.startswith(self._GPT_IMAGE_MODEL_PREFIX)\n or normalized_model == self._DALL_E_2_MODEL_NAME\n )\n\n def generate_image(\n self,\n prompt: str,\n model: str,\n size: str,\n n: int,\n quality: str | None = None,\n reference_images: list[ReferenceImage] | None = None,\n **kwargs: Any,\n ) -> ImageGenerationResponse:\n deployment = self._deployment_name or model\n model_name = f\"azure/{deployment}\"\n\n if reference_images:\n if not self._model_supports_image_edits(model):\n raise ValueError(\n f\"Model '{model}' does not support image edits with reference images.\"\n )\n\n normalized_model = self._normalize_model_name(model)\n if (\n normalized_model == self._DALL_E_2_MODEL_NAME\n and len(reference_images) > 1\n ):\n raise ValueError(\n \"Model 'dall-e-2' only supports a single reference image for edits.\"\n )\n\n from litellm import image_edit\n\n return image_edit(\n image=[image.data for image in reference_images],\n prompt=prompt,\n model=model_name,\n api_key=self._api_key,\n api_base=self._api_base,\n api_version=self._api_version,\n size=size,\n n=n,\n quality=quality,\n **kwargs,\n )\n\n from litellm import image_generation\n\n return image_generation(\n prompt=prompt,\n model=model_name,\n api_key=self._api_key,\n api_base=self._api_base,\n api_version=self._api_version,\n size=size,\n n=n,\n quality=quality,\n **kwargs,\n )\n" }, { "path": "backend/onyx/image_gen/providers/openai_img_gen.py", "content": "from __future__ import annotations\n\nfrom typing import Any\nfrom typing import TYPE_CHECKING\n\nfrom onyx.image_gen.interfaces import ImageGenerationProvider\nfrom onyx.image_gen.interfaces import ImageGenerationProviderCredentials\nfrom onyx.image_gen.interfaces import ReferenceImage\n\nif TYPE_CHECKING:\n from onyx.image_gen.interfaces import ImageGenerationResponse\n\n\nclass OpenAIImageGenerationProvider(ImageGenerationProvider):\n _GPT_IMAGE_MODEL_PREFIX = \"gpt-image-\"\n _DALL_E_2_MODEL_NAME = \"dall-e-2\"\n\n def __init__(\n self,\n api_key: str,\n api_base: str | None = None,\n ):\n self._api_key = api_key\n self._api_base = api_base\n\n @classmethod\n def validate_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> bool:\n return bool(credentials.api_key)\n\n @classmethod\n def _build_from_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> OpenAIImageGenerationProvider:\n assert credentials.api_key\n\n return cls(\n api_key=credentials.api_key,\n api_base=credentials.api_base,\n )\n\n @property\n def supports_reference_images(self) -> bool:\n return True\n\n @property\n def max_reference_images(self) -> int:\n # GPT image models support up to 16 input images for edits.\n return 16\n\n def _normalize_model_name(self, model: str) -> str:\n return model.rsplit(\"/\", 1)[-1]\n\n def _model_supports_image_edits(self, model: str) -> bool:\n normalized_model = self._normalize_model_name(model)\n return (\n normalized_model.startswith(self._GPT_IMAGE_MODEL_PREFIX)\n or normalized_model == self._DALL_E_2_MODEL_NAME\n )\n\n def generate_image(\n self,\n prompt: str,\n model: str,\n size: str,\n n: int,\n quality: str | None = None,\n reference_images: list[ReferenceImage] | None = None,\n **kwargs: Any,\n ) -> ImageGenerationResponse:\n if reference_images:\n if not self._model_supports_image_edits(model):\n raise ValueError(\n f\"Model '{model}' does not support image edits with reference images.\"\n )\n\n normalized_model = self._normalize_model_name(model)\n if (\n normalized_model == self._DALL_E_2_MODEL_NAME\n and len(reference_images) > 1\n ):\n raise ValueError(\n \"Model 'dall-e-2' only supports a single reference image for edits.\"\n )\n\n from litellm import image_edit\n\n return image_edit(\n image=[image.data for image in reference_images],\n prompt=prompt,\n model=model,\n api_key=self._api_key,\n api_base=self._api_base,\n size=size,\n n=n,\n quality=quality,\n **kwargs,\n )\n\n from litellm import image_generation\n\n return image_generation(\n prompt=prompt,\n model=model,\n api_key=self._api_key,\n api_base=self._api_base,\n size=size,\n n=n,\n quality=quality,\n **kwargs,\n )\n" }, { "path": "backend/onyx/image_gen/providers/vertex_img_gen.py", "content": "from __future__ import annotations\n\nimport base64\nimport json\nfrom datetime import datetime\nfrom typing import Any\nfrom typing import TYPE_CHECKING\n\nfrom pydantic import BaseModel\n\nfrom onyx.image_gen.exceptions import ImageProviderCredentialsError\nfrom onyx.image_gen.interfaces import ImageGenerationProvider\nfrom onyx.image_gen.interfaces import ImageGenerationProviderCredentials\nfrom onyx.image_gen.interfaces import ReferenceImage\n\nif TYPE_CHECKING:\n from onyx.image_gen.interfaces import ImageGenerationResponse\n\n\nclass VertexCredentials(BaseModel):\n vertex_credentials: str\n vertex_location: str\n project_id: str\n\n\nclass VertexImageGenerationProvider(ImageGenerationProvider):\n def __init__(\n self,\n vertex_credentials: VertexCredentials,\n ):\n self._vertex_credentials = vertex_credentials.vertex_credentials\n self._vertex_location = vertex_credentials.vertex_location\n self._vertex_project = vertex_credentials.project_id\n\n @classmethod\n def validate_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> bool:\n try:\n _parse_to_vertex_credentials(credentials)\n return True\n except ImageProviderCredentialsError:\n return False\n\n @classmethod\n def _build_from_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials,\n ) -> VertexImageGenerationProvider:\n vertex_credentials = _parse_to_vertex_credentials(credentials)\n\n return cls(\n vertex_credentials=vertex_credentials,\n )\n\n @property\n def supports_reference_images(self) -> bool:\n return True\n\n @property\n def max_reference_images(self) -> int:\n # Gemini image editing supports up to 14 input images.\n return 14\n\n def generate_image(\n self,\n prompt: str,\n model: str,\n size: str,\n n: int,\n quality: str | None = None,\n reference_images: list[ReferenceImage] | None = None,\n **kwargs: Any,\n ) -> ImageGenerationResponse:\n if reference_images:\n return self._generate_image_with_reference_images(\n prompt=prompt,\n model=model,\n size=size,\n n=n,\n reference_images=reference_images,\n )\n\n from litellm import image_generation\n\n return image_generation(\n prompt=prompt,\n model=model,\n size=size,\n n=n,\n quality=quality,\n vertex_location=self._vertex_location,\n vertex_credentials=self._vertex_credentials,\n vertex_project=self._vertex_project,\n **kwargs,\n )\n\n def _generate_image_with_reference_images(\n self,\n prompt: str,\n model: str,\n size: str,\n n: int,\n reference_images: list[ReferenceImage],\n ) -> ImageGenerationResponse:\n from google import genai\n from google.genai import types as genai_types\n from google.oauth2 import service_account\n from litellm.types.utils import ImageObject\n from litellm.types.utils import ImageResponse\n\n service_account_info = json.loads(self._vertex_credentials)\n credentials = service_account.Credentials.from_service_account_info(\n service_account_info,\n scopes=[\"https://www.googleapis.com/auth/cloud-platform\"],\n )\n\n client = genai.Client(\n vertexai=True,\n project=self._vertex_project,\n location=self._vertex_location,\n credentials=credentials,\n )\n\n parts: list[genai_types.Part] = [\n genai_types.Part.from_bytes(data=image.data, mime_type=image.mime_type)\n for image in reference_images\n ]\n parts.append(genai_types.Part.from_text(text=prompt))\n\n config = genai_types.GenerateContentConfig(\n response_modalities=[\"TEXT\", \"IMAGE\"],\n candidate_count=max(1, n),\n image_config=genai_types.ImageConfig(\n aspect_ratio=_map_size_to_aspect_ratio(size)\n ),\n )\n model_name = model.replace(\"vertex_ai/\", \"\")\n response = client.models.generate_content(\n model=model_name,\n contents=genai_types.Content(\n role=\"user\",\n parts=parts,\n ),\n config=config,\n )\n\n generated_data: list[ImageObject] = []\n for candidate in response.candidates or []:\n candidate_content = candidate.content\n if not candidate_content:\n continue\n\n for part in candidate_content.parts or []:\n inline_data = part.inline_data\n if not inline_data or inline_data.data is None:\n continue\n\n if isinstance(inline_data.data, bytes):\n b64_json = base64.b64encode(inline_data.data).decode(\"utf-8\")\n elif isinstance(inline_data.data, str):\n b64_json = inline_data.data\n else:\n continue\n\n generated_data.append(\n ImageObject(\n b64_json=b64_json,\n revised_prompt=prompt,\n )\n )\n\n if not generated_data:\n raise RuntimeError(\"No image data returned from Vertex AI.\")\n\n return ImageResponse(\n created=int(datetime.now().timestamp()),\n data=generated_data,\n )\n\n\ndef _map_size_to_aspect_ratio(size: str) -> str:\n return {\n \"1024x1024\": \"1:1\",\n \"1792x1024\": \"16:9\",\n \"1024x1792\": \"9:16\",\n \"1536x1024\": \"3:2\",\n \"1024x1536\": \"2:3\",\n }.get(size, \"1:1\")\n\n\ndef _parse_to_vertex_credentials(\n credentials: ImageGenerationProviderCredentials,\n) -> VertexCredentials:\n custom_config = credentials.custom_config\n\n if not custom_config:\n raise ImageProviderCredentialsError(\"Custom config is required\")\n\n vertex_credentials = custom_config.get(\"vertex_credentials\")\n vertex_location = custom_config.get(\"vertex_location\")\n\n if not vertex_credentials:\n raise ImageProviderCredentialsError(\"Vertex credentials are required\")\n\n if not vertex_location:\n raise ImageProviderCredentialsError(\"Vertex location is required\")\n\n vertex_json = json.loads(vertex_credentials)\n vertex_project = vertex_json.get(\"project_id\")\n\n if not vertex_project:\n raise ImageProviderCredentialsError(\"Project ID is required\")\n\n return VertexCredentials(\n vertex_credentials=vertex_credentials,\n vertex_location=vertex_location,\n project_id=vertex_project,\n )\n" }, { "path": "backend/onyx/indexing/__init__.py", "content": "" }, { "path": "backend/onyx/indexing/adapters/document_indexing_adapter.py", "content": "import contextlib\nfrom collections.abc import Generator\n\nfrom sqlalchemy.engine.util import TransactionalContext\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.access import get_access_for_documents\nfrom onyx.access.models import DocumentAccess\nfrom onyx.configs.constants import DEFAULT_BOOST\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import IndexAttemptMetadata\nfrom onyx.db.chunk import update_chunk_boost_components__no_commit\nfrom onyx.db.document import fetch_chunk_counts_for_documents\nfrom onyx.db.document import mark_document_as_indexed_for_cc_pair__no_commit\nfrom onyx.db.document import prepare_to_modify_documents\nfrom onyx.db.document import update_docs_chunk_count__no_commit\nfrom onyx.db.document import update_docs_last_modified__no_commit\nfrom onyx.db.document import update_docs_updated_at__no_commit\nfrom onyx.db.document_set import fetch_document_sets_for_documents\nfrom onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext\nfrom onyx.indexing.indexing_pipeline import index_doc_batch_prepare\nfrom onyx.indexing.models import ChunkEnrichmentContext\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.indexing.models import IndexChunk\nfrom onyx.indexing.models import UpdatableChunkData\nfrom onyx.redis.redis_hierarchy import get_ancestors_from_raw_id\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass DocumentIndexingBatchAdapter:\n \"\"\"Default adapter: handles DB prep, locking, metadata enrichment, and finalize.\n\n Keeps orchestration logic in the pipeline and side-effects in the adapter.\n \"\"\"\n\n def __init__(\n self,\n db_session: Session,\n connector_id: int,\n credential_id: int,\n tenant_id: str,\n index_attempt_metadata: IndexAttemptMetadata,\n ):\n self.db_session = db_session\n self.connector_id = connector_id\n self.credential_id = credential_id\n self.tenant_id = tenant_id\n self.index_attempt_metadata = index_attempt_metadata\n\n def prepare(\n self, documents: list[Document], ignore_time_skip: bool\n ) -> DocumentBatchPrepareContext | None:\n \"\"\"Upsert docs, map CC pairs, return context or mark as indexed if no-op.\"\"\"\n context = index_doc_batch_prepare(\n documents=documents,\n index_attempt_metadata=self.index_attempt_metadata,\n db_session=self.db_session,\n ignore_time_skip=ignore_time_skip,\n )\n\n if not context:\n # even though we didn't actually index anything, we should still\n # mark them as \"completed\" for the CC Pair in order to make the\n # counts match\n mark_document_as_indexed_for_cc_pair__no_commit(\n connector_id=self.index_attempt_metadata.connector_id,\n credential_id=self.index_attempt_metadata.credential_id,\n document_ids=[doc.id for doc in documents],\n db_session=self.db_session,\n )\n self.db_session.commit()\n\n return context\n\n @contextlib.contextmanager\n def lock_context(\n self, documents: list[Document]\n ) -> Generator[TransactionalContext, None, None]:\n \"\"\"Acquire transaction/row locks on docs for the critical section.\"\"\"\n with prepare_to_modify_documents(\n db_session=self.db_session, document_ids=[doc.id for doc in documents]\n ) as transaction:\n yield transaction\n\n def prepare_enrichment(\n self,\n context: DocumentBatchPrepareContext,\n tenant_id: str,\n chunks: list[DocAwareChunk],\n ) -> \"DocumentChunkEnricher\":\n \"\"\"Do all DB lookups once and return a per-chunk enricher.\"\"\"\n updatable_ids = [doc.id for doc in context.updatable_docs]\n\n doc_id_to_new_chunk_cnt: dict[str, int] = {\n doc_id: 0 for doc_id in updatable_ids\n }\n for chunk in chunks:\n if chunk.source_document.id in doc_id_to_new_chunk_cnt:\n doc_id_to_new_chunk_cnt[chunk.source_document.id] += 1\n\n no_access = DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n )\n\n return DocumentChunkEnricher(\n doc_id_to_access_info=get_access_for_documents(\n document_ids=updatable_ids, db_session=self.db_session\n ),\n doc_id_to_document_set={\n document_id: document_sets\n for document_id, document_sets in fetch_document_sets_for_documents(\n document_ids=updatable_ids, db_session=self.db_session\n )\n },\n doc_id_to_ancestor_ids=self._get_ancestor_ids_for_documents(\n context.updatable_docs, tenant_id\n ),\n id_to_boost_map=context.id_to_boost_map,\n doc_id_to_previous_chunk_cnt={\n document_id: chunk_count\n for document_id, chunk_count in fetch_chunk_counts_for_documents(\n document_ids=updatable_ids,\n db_session=self.db_session,\n )\n },\n doc_id_to_new_chunk_cnt=dict(doc_id_to_new_chunk_cnt),\n no_access=no_access,\n tenant_id=tenant_id,\n )\n\n def _get_ancestor_ids_for_documents(\n self,\n documents: list[Document],\n tenant_id: str,\n ) -> dict[str, list[int]]:\n \"\"\"\n Get ancestor hierarchy node IDs for a batch of documents.\n\n Uses Redis cache for fast lookups - no DB calls are made unless\n there's a cache miss. Documents provide parent_hierarchy_raw_node_id\n directly from the connector.\n\n Returns a mapping from document_id to list of ancestor node IDs.\n \"\"\"\n if not documents:\n return {}\n\n redis_client = get_redis_client(tenant_id=tenant_id)\n result: dict[str, list[int]] = {}\n\n for doc in documents:\n # Use parent_hierarchy_raw_node_id directly from the document\n # If None, get_ancestors_from_raw_id will return just the SOURCE node\n ancestors = get_ancestors_from_raw_id(\n redis_client=redis_client,\n source=doc.source,\n parent_hierarchy_raw_node_id=doc.parent_hierarchy_raw_node_id,\n db_session=self.db_session,\n )\n result[doc.id] = ancestors\n\n return result\n\n def post_index(\n self,\n context: DocumentBatchPrepareContext,\n updatable_chunk_data: list[UpdatableChunkData],\n filtered_documents: list[Document],\n enrichment: ChunkEnrichmentContext,\n ) -> None:\n \"\"\"Finalize DB updates, store plaintext, and mark docs as indexed.\"\"\"\n updatable_ids = [doc.id for doc in context.updatable_docs]\n last_modified_ids = []\n ids_to_new_updated_at = {}\n for doc in context.updatable_docs:\n last_modified_ids.append(doc.id)\n # doc_updated_at is the source's idea (on the other end of the connector)\n # of when the doc was last modified\n if doc.doc_updated_at is None:\n continue\n ids_to_new_updated_at[doc.id] = doc.doc_updated_at\n\n update_docs_updated_at__no_commit(\n ids_to_new_updated_at=ids_to_new_updated_at, db_session=self.db_session\n )\n\n update_docs_last_modified__no_commit(\n document_ids=last_modified_ids, db_session=self.db_session\n )\n\n update_docs_chunk_count__no_commit(\n document_ids=updatable_ids,\n doc_id_to_chunk_count=enrichment.doc_id_to_new_chunk_cnt,\n db_session=self.db_session,\n )\n\n # these documents can now be counted as part of the CC Pairs\n # document count, so we need to mark them as indexed\n # NOTE: even documents we skipped since they were already up\n # to date should be counted here in order to maintain parity\n # between CC Pair and index attempt counts\n mark_document_as_indexed_for_cc_pair__no_commit(\n connector_id=self.index_attempt_metadata.connector_id,\n credential_id=self.index_attempt_metadata.credential_id,\n document_ids=[doc.id for doc in filtered_documents],\n db_session=self.db_session,\n )\n\n # save the chunk boost components to postgres\n update_chunk_boost_components__no_commit(\n chunk_data=updatable_chunk_data, db_session=self.db_session\n )\n\n self.db_session.commit()\n\n\nclass DocumentChunkEnricher:\n \"\"\"Pre-computed metadata for per-chunk enrichment of connector documents.\"\"\"\n\n def __init__(\n self,\n doc_id_to_access_info: dict[str, DocumentAccess],\n doc_id_to_document_set: dict[str, list[str]],\n doc_id_to_ancestor_ids: dict[str, list[int]],\n id_to_boost_map: dict[str, int],\n doc_id_to_previous_chunk_cnt: dict[str, int],\n doc_id_to_new_chunk_cnt: dict[str, int],\n no_access: DocumentAccess,\n tenant_id: str,\n ) -> None:\n self._doc_id_to_access_info = doc_id_to_access_info\n self._doc_id_to_document_set = doc_id_to_document_set\n self._doc_id_to_ancestor_ids = doc_id_to_ancestor_ids\n self._id_to_boost_map = id_to_boost_map\n self._no_access = no_access\n self._tenant_id = tenant_id\n self.doc_id_to_previous_chunk_cnt = doc_id_to_previous_chunk_cnt\n self.doc_id_to_new_chunk_cnt = doc_id_to_new_chunk_cnt\n\n def enrich_chunk(\n self, chunk: IndexChunk, score: float\n ) -> DocMetadataAwareIndexChunk:\n return DocMetadataAwareIndexChunk.from_index_chunk(\n index_chunk=chunk,\n access=self._doc_id_to_access_info.get(\n chunk.source_document.id, self._no_access\n ),\n document_sets=set(\n self._doc_id_to_document_set.get(chunk.source_document.id, [])\n ),\n user_project=[],\n personas=[],\n boost=(\n self._id_to_boost_map[chunk.source_document.id]\n if chunk.source_document.id in self._id_to_boost_map\n else DEFAULT_BOOST\n ),\n tenant_id=self._tenant_id,\n aggregated_chunk_boost_factor=score,\n ancestor_hierarchy_node_ids=self._doc_id_to_ancestor_ids[\n chunk.source_document.id\n ],\n )\n" }, { "path": "backend/onyx/indexing/adapters/user_file_indexing_adapter.py", "content": "from __future__ import annotations\n\nimport contextlib\nimport datetime\nimport time\nfrom collections import defaultdict\nfrom collections.abc import Generator\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy.exc import OperationalError\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\nfrom sqlalchemy.orm.session import TransactionalContext\n\nfrom onyx.access.access import get_access_for_user_files\nfrom onyx.access.models import DocumentAccess\nfrom onyx.configs.constants import DEFAULT_BOOST\nfrom onyx.configs.constants import NotificationType\nfrom onyx.connectors.models import Document\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import Persona\nfrom onyx.db.models import UserFile\nfrom onyx.db.notification import create_notification\nfrom onyx.db.user_file import fetch_chunk_counts_for_user_files\nfrom onyx.db.user_file import fetch_persona_ids_for_user_files\nfrom onyx.db.user_file import fetch_user_project_ids_for_user_files\nfrom onyx.file_store.utils import store_user_file_plaintext\nfrom onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext\nfrom onyx.indexing.models import ChunkEnrichmentContext\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.indexing.models import IndexChunk\nfrom onyx.indexing.models import UpdatableChunkData\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.natural_language_processing.utils import count_tokens\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_NUM_LOCK_ATTEMPTS = 3\nretry_delay = 0.5\n\n\ndef _acquire_user_file_locks(db_session: Session, user_file_ids: list[str]) -> bool:\n \"\"\"Acquire locks for the specified user files.\"\"\"\n # Convert to UUIDs for the DB comparison\n user_file_uuid_list = [UUID(user_file_id) for user_file_id in user_file_ids]\n stmt = (\n select(UserFile.id)\n .where(UserFile.id.in_(user_file_uuid_list))\n .with_for_update(nowait=True)\n )\n # will raise exception if any of the documents are already locked\n documents = db_session.scalars(stmt).all()\n\n # make sure we found every document\n if len(documents) != len(set(user_file_ids)):\n logger.warning(\"Didn't find row for all specified user file IDs. Aborting.\")\n return False\n\n return True\n\n\nclass UserFileIndexingAdapter:\n def __init__(self, tenant_id: str, db_session: Session):\n self.tenant_id = tenant_id\n self.db_session = db_session\n\n def prepare(\n self,\n documents: list[Document],\n ignore_time_skip: bool, # noqa: ARG002\n ) -> DocumentBatchPrepareContext:\n return DocumentBatchPrepareContext(\n updatable_docs=documents,\n id_to_boost_map={}, # TODO(subash): add boost map\n )\n\n @contextlib.contextmanager\n def lock_context(\n self, documents: list[Document]\n ) -> Generator[TransactionalContext, None, None]:\n self.db_session.commit() # ensure that we're not in a transaction\n lock_acquired = False\n for i in range(_NUM_LOCK_ATTEMPTS):\n try:\n with self.db_session.begin() as transaction:\n lock_acquired = _acquire_user_file_locks(\n db_session=self.db_session,\n user_file_ids=[doc.id for doc in documents],\n )\n if lock_acquired:\n yield transaction\n break\n except OperationalError as e:\n logger.warning(\n f\"Failed to acquire locks for user files on attempt {i}, retrying. Error: {e}\"\n )\n\n time.sleep(retry_delay)\n\n if not lock_acquired:\n raise RuntimeError(\n f\"Failed to acquire locks after {_NUM_LOCK_ATTEMPTS} attempts for user files: {[doc.id for doc in documents]}\"\n )\n\n def prepare_enrichment(\n self,\n context: DocumentBatchPrepareContext,\n tenant_id: str,\n chunks: list[DocAwareChunk],\n ) -> UserFileChunkEnricher:\n \"\"\"Do all DB lookups and pre-compute file metadata from chunks.\"\"\"\n updatable_ids = [doc.id for doc in context.updatable_docs]\n\n doc_id_to_new_chunk_cnt: dict[str, int] = defaultdict(int)\n content_by_file: dict[str, list[str]] = defaultdict(list)\n for chunk in chunks:\n doc_id_to_new_chunk_cnt[chunk.source_document.id] += 1\n content_by_file[chunk.source_document.id].append(chunk.content)\n\n no_access = DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n )\n\n user_file_id_to_project_ids = fetch_user_project_ids_for_user_files(\n user_file_ids=updatable_ids,\n db_session=self.db_session,\n )\n user_file_id_to_persona_ids = fetch_persona_ids_for_user_files(\n user_file_ids=updatable_ids,\n db_session=self.db_session,\n )\n user_file_id_to_access: dict[str, DocumentAccess] = get_access_for_user_files(\n user_file_ids=updatable_ids,\n db_session=self.db_session,\n )\n user_file_id_to_previous_chunk_cnt: dict[str, int] = {\n user_file_id: chunk_count\n for user_file_id, chunk_count in fetch_chunk_counts_for_user_files(\n user_file_ids=updatable_ids,\n db_session=self.db_session,\n )\n }\n\n # Initialize tokenizer used for token count calculation\n try:\n llm = get_default_llm()\n llm_tokenizer = get_tokenizer(\n model_name=llm.config.model_name,\n provider_type=llm.config.model_provider,\n )\n except Exception as e:\n logger.error(f\"Error getting tokenizer: {e}\")\n llm_tokenizer = None\n\n user_file_id_to_raw_text: dict[str, str] = {}\n user_file_id_to_token_count: dict[str, int | None] = {}\n for user_file_id in updatable_ids:\n contents = content_by_file.get(user_file_id)\n if contents:\n combined_content = \" \".join(contents)\n user_file_id_to_raw_text[str(user_file_id)] = combined_content\n token_count: int = (\n count_tokens(combined_content, llm_tokenizer)\n if llm_tokenizer\n else 0\n )\n user_file_id_to_token_count[str(user_file_id)] = token_count\n else:\n user_file_id_to_raw_text[str(user_file_id)] = \"\"\n user_file_id_to_token_count[str(user_file_id)] = None\n\n return UserFileChunkEnricher(\n user_file_id_to_access=user_file_id_to_access,\n user_file_id_to_project_ids=user_file_id_to_project_ids,\n user_file_id_to_persona_ids=user_file_id_to_persona_ids,\n doc_id_to_previous_chunk_cnt=user_file_id_to_previous_chunk_cnt,\n doc_id_to_new_chunk_cnt=dict(doc_id_to_new_chunk_cnt),\n user_file_id_to_raw_text=user_file_id_to_raw_text,\n user_file_id_to_token_count=user_file_id_to_token_count,\n no_access=no_access,\n tenant_id=tenant_id,\n )\n\n def _notify_assistant_owners_if_files_ready(\n self, user_files: list[UserFile]\n ) -> None:\n \"\"\"\n Check if all files for associated assistants are processed and notify owners.\n Only sends notification when all files for an assistant are COMPLETED.\n \"\"\"\n for user_file in user_files:\n if user_file.status == UserFileStatus.COMPLETED:\n for assistant in user_file.assistants:\n # Skip assistants without owners\n if assistant.user_id is None:\n continue\n\n # Check if all OTHER files for this assistant are completed\n # (we already know current file is completed from the outer check)\n all_files_completed = all(\n f.status == UserFileStatus.COMPLETED\n for f in assistant.user_files\n if f.id != user_file.id\n )\n\n if all_files_completed:\n create_notification(\n user_id=assistant.user_id,\n notif_type=NotificationType.ASSISTANT_FILES_READY,\n db_session=self.db_session,\n title=\"Your files are ready!\",\n description=f\"All files for agent {assistant.name} have been processed and are now available.\",\n additional_data={\n \"persona_id\": assistant.id,\n \"link\": f\"/assistants/{assistant.id}\",\n },\n autocommit=False,\n )\n\n def post_index(\n self,\n context: DocumentBatchPrepareContext,\n updatable_chunk_data: list[UpdatableChunkData], # noqa: ARG002\n filtered_documents: list[Document], # noqa: ARG002\n enrichment: ChunkEnrichmentContext,\n ) -> None:\n assert isinstance(enrichment, UserFileChunkEnricher)\n user_file_ids = [doc.id for doc in context.updatable_docs]\n\n user_files = (\n self.db_session.query(UserFile)\n .options(selectinload(UserFile.assistants).selectinload(Persona.user_files))\n .filter(UserFile.id.in_(user_file_ids))\n .all()\n )\n for user_file in user_files:\n # don't update the status if the user file is being deleted\n if user_file.status != UserFileStatus.DELETING:\n user_file.status = UserFileStatus.COMPLETED\n user_file.last_project_sync_at = datetime.datetime.now(\n datetime.timezone.utc\n )\n user_file.chunk_count = enrichment.doc_id_to_new_chunk_cnt.get(\n str(user_file.id), 0\n )\n user_file.token_count = enrichment.user_file_id_to_token_count[\n str(user_file.id)\n ]\n\n # Notify assistant owners if all their files are now processed\n self._notify_assistant_owners_if_files_ready(user_files)\n\n self.db_session.commit()\n\n # Store the plaintext in the file store for faster retrieval\n # NOTE: this creates its own session to avoid committing the overall\n # transaction.\n for user_file_id, raw_text in enrichment.user_file_id_to_raw_text.items():\n store_user_file_plaintext(\n user_file_id=UUID(user_file_id),\n plaintext_content=raw_text,\n )\n\n\nclass UserFileChunkEnricher:\n \"\"\"Pre-computed metadata for per-chunk enrichment of user-uploaded files.\"\"\"\n\n def __init__(\n self,\n user_file_id_to_access: dict[str, DocumentAccess],\n user_file_id_to_project_ids: dict[str, list[int]],\n user_file_id_to_persona_ids: dict[str, list[int]],\n doc_id_to_previous_chunk_cnt: dict[str, int],\n doc_id_to_new_chunk_cnt: dict[str, int],\n user_file_id_to_raw_text: dict[str, str],\n user_file_id_to_token_count: dict[str, int | None],\n no_access: DocumentAccess,\n tenant_id: str,\n ) -> None:\n self._user_file_id_to_access = user_file_id_to_access\n self._user_file_id_to_project_ids = user_file_id_to_project_ids\n self._user_file_id_to_persona_ids = user_file_id_to_persona_ids\n self._no_access = no_access\n self._tenant_id = tenant_id\n self.doc_id_to_previous_chunk_cnt = doc_id_to_previous_chunk_cnt\n self.doc_id_to_new_chunk_cnt = doc_id_to_new_chunk_cnt\n self.user_file_id_to_raw_text = user_file_id_to_raw_text\n self.user_file_id_to_token_count = user_file_id_to_token_count\n\n def enrich_chunk(\n self, chunk: IndexChunk, score: float\n ) -> DocMetadataAwareIndexChunk:\n return DocMetadataAwareIndexChunk.from_index_chunk(\n index_chunk=chunk,\n access=self._user_file_id_to_access.get(\n chunk.source_document.id, self._no_access\n ),\n document_sets=set(),\n user_project=self._user_file_id_to_project_ids.get(\n chunk.source_document.id, []\n ),\n personas=self._user_file_id_to_persona_ids.get(\n chunk.source_document.id, []\n ),\n boost=DEFAULT_BOOST,\n tenant_id=self._tenant_id,\n aggregated_chunk_boost_factor=score,\n )\n" }, { "path": "backend/onyx/indexing/chunk_batch_store.py", "content": "import pickle\nimport shutil\nimport tempfile\nfrom collections.abc import Iterator\nfrom pathlib import Path\n\nfrom onyx.indexing.models import IndexChunk\n\n\nclass ChunkBatchStore:\n \"\"\"Manages serialization of embedded chunks to a temporary directory.\n\n Owns the temp directory lifetime and provides save/load/stream/scrub\n operations.\n\n Use as a context manager to ensure cleanup::\n\n with ChunkBatchStore() as store:\n store.save(chunks, batch_idx=0)\n for chunk in store.stream():\n ...\n \"\"\"\n\n _EXT = \".pkl\"\n\n def __init__(self) -> None:\n self._tmpdir: Path | None = None\n\n # -- context manager -----------------------------------------------------\n\n def __enter__(self) -> \"ChunkBatchStore\":\n self._tmpdir = Path(tempfile.mkdtemp(prefix=\"onyx_embeddings_\"))\n return self\n\n def __exit__(self, *_exc: object) -> None:\n if self._tmpdir is not None:\n shutil.rmtree(self._tmpdir, ignore_errors=True)\n self._tmpdir = None\n\n @property\n def _dir(self) -> Path:\n assert self._tmpdir is not None, \"ChunkBatchStore used outside context manager\"\n return self._tmpdir\n\n # -- storage primitives --------------------------------------------------\n\n def save(self, chunks: list[IndexChunk], batch_idx: int) -> None:\n \"\"\"Serialize a batch of embedded chunks to disk.\"\"\"\n with open(self._dir / f\"batch_{batch_idx}{self._EXT}\", \"wb\") as f:\n pickle.dump(chunks, f)\n\n def _load(self, batch_file: Path) -> list[IndexChunk]:\n \"\"\"Deserialize a batch of embedded chunks from a file.\"\"\"\n with open(batch_file, \"rb\") as f:\n return pickle.load(f)\n\n def _batch_files(self) -> list[Path]:\n \"\"\"Return batch files sorted by numeric index.\"\"\"\n return sorted(\n self._dir.glob(f\"batch_*{self._EXT}\"),\n key=lambda p: int(p.stem.removeprefix(\"batch_\")),\n )\n\n # -- higher-level operations ---------------------------------------------\n\n def stream(self) -> Iterator[IndexChunk]:\n \"\"\"Yield all chunks across all batch files.\n\n Each call returns a fresh generator, so the data can be iterated\n multiple times (e.g. once per document index).\n \"\"\"\n for batch_file in self._batch_files():\n yield from self._load(batch_file)\n\n def scrub_failed_docs(self, failed_doc_ids: set[str]) -> None:\n \"\"\"Remove chunks belonging to *failed_doc_ids* from all batch files.\n\n When a document fails embedding in batch N, earlier batches may\n already contain successfully embedded chunks for that document.\n This ensures the output is all-or-nothing per document.\n \"\"\"\n for batch_file in self._batch_files():\n batch_chunks = self._load(batch_file)\n cleaned = [\n c for c in batch_chunks if c.source_document.id not in failed_doc_ids\n ]\n if len(cleaned) != len(batch_chunks):\n with open(batch_file, \"wb\") as f:\n pickle.dump(cleaned, f)\n" }, { "path": "backend/onyx/indexing/chunker.py", "content": "from typing import cast\n\nfrom chonkie import SentenceChunker\n\nfrom onyx.configs.app_configs import AVERAGE_SUMMARY_EMBEDDINGS\nfrom onyx.configs.app_configs import BLURB_SIZE\nfrom onyx.configs.app_configs import LARGE_CHUNK_RATIO\nfrom onyx.configs.app_configs import MINI_CHUNK_SIZE\nfrom onyx.configs.app_configs import SKIP_METADATA_IN_CHUNK\nfrom onyx.configs.app_configs import USE_CHUNK_SUMMARY\nfrom onyx.configs.app_configs import USE_DOCUMENT_SUMMARY\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import RETURN_SEPARATOR\nfrom onyx.configs.constants import SECTION_SEPARATOR\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n get_metadata_keys_to_ignore,\n)\nfrom onyx.connectors.models import IndexingDocument\nfrom onyx.connectors.models import Section\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.llm.utils import MAX_CONTEXT_TOKENS\nfrom onyx.natural_language_processing.utils import BaseTokenizer\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.text_processing import clean_text\nfrom onyx.utils.text_processing import shared_precompare_cleanup\nfrom shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE\nfrom shared_configs.configs import STRICT_CHUNK_TOKEN_LIMIT\n\n# Not supporting overlaps, we need a clean combination of chunks and it is unclear if overlaps\n# actually help quality at all\nCHUNK_OVERLAP = 0\n# Fairly arbitrary numbers but the general concept is we don't want the title/metadata to\n# overwhelm the actual contents of the chunk\nMAX_METADATA_PERCENTAGE = 0.25\nCHUNK_MIN_CONTENT = 256\n\nlogger = setup_logger()\n\n\ndef _get_metadata_suffix_for_document_index(\n metadata: dict[str, str | list[str]], include_separator: bool = False\n) -> tuple[str, str]:\n \"\"\"\n Returns the metadata as a natural language string representation with all of the keys and values\n for the vector embedding and a string of all of the values for the keyword search.\n \"\"\"\n if not metadata:\n return \"\", \"\"\n\n metadata_str = \"Metadata:\\n\"\n metadata_values = []\n for key, value in metadata.items():\n if key in get_metadata_keys_to_ignore():\n continue\n\n value_str = \", \".join(value) if isinstance(value, list) else value\n\n if isinstance(value, list):\n metadata_values.extend(value)\n else:\n metadata_values.append(value)\n\n metadata_str += f\"\\t{key} - {value_str}\\n\"\n\n metadata_semantic = metadata_str.strip()\n metadata_keyword = \" \".join(metadata_values)\n\n if include_separator:\n return RETURN_SEPARATOR + metadata_semantic, RETURN_SEPARATOR + metadata_keyword\n return metadata_semantic, metadata_keyword\n\n\ndef _combine_chunks(chunks: list[DocAwareChunk], large_chunk_id: int) -> DocAwareChunk:\n \"\"\"\n Combines multiple DocAwareChunks into one large chunk (for \"multipass\" mode),\n appending the content and adjusting source_links accordingly.\n \"\"\"\n merged_chunk = DocAwareChunk(\n source_document=chunks[0].source_document,\n chunk_id=chunks[0].chunk_id,\n blurb=chunks[0].blurb,\n content=chunks[0].content,\n source_links=chunks[0].source_links or {},\n image_file_id=None,\n section_continuation=(chunks[0].chunk_id > 0),\n title_prefix=chunks[0].title_prefix,\n metadata_suffix_semantic=chunks[0].metadata_suffix_semantic,\n metadata_suffix_keyword=chunks[0].metadata_suffix_keyword,\n large_chunk_reference_ids=[chunk.chunk_id for chunk in chunks],\n mini_chunk_texts=None,\n large_chunk_id=large_chunk_id,\n chunk_context=\"\",\n doc_summary=\"\",\n contextual_rag_reserved_tokens=0,\n )\n\n offset = 0\n for i in range(1, len(chunks)):\n merged_chunk.content += SECTION_SEPARATOR + chunks[i].content\n\n offset += len(SECTION_SEPARATOR) + len(chunks[i - 1].content)\n for link_offset, link_text in (chunks[i].source_links or {}).items():\n if merged_chunk.source_links is None:\n merged_chunk.source_links = {}\n merged_chunk.source_links[link_offset + offset] = link_text\n\n return merged_chunk\n\n\ndef generate_large_chunks(chunks: list[DocAwareChunk]) -> list[DocAwareChunk]:\n \"\"\"\n Generates larger \"grouped\" chunks by combining sets of smaller chunks.\n \"\"\"\n large_chunks = []\n for idx, i in enumerate(range(0, len(chunks), LARGE_CHUNK_RATIO)):\n chunk_group = chunks[i : i + LARGE_CHUNK_RATIO]\n if len(chunk_group) > 1:\n large_chunk = _combine_chunks(chunk_group, idx)\n large_chunks.append(large_chunk)\n return large_chunks\n\n\nclass Chunker:\n \"\"\"\n Chunks documents into smaller chunks for indexing.\n \"\"\"\n\n def __init__(\n self,\n tokenizer: BaseTokenizer,\n enable_multipass: bool = False,\n enable_large_chunks: bool = False,\n enable_contextual_rag: bool = False,\n blurb_size: int = BLURB_SIZE,\n include_metadata: bool = not SKIP_METADATA_IN_CHUNK,\n chunk_token_limit: int = DOC_EMBEDDING_CONTEXT_SIZE,\n chunk_overlap: int = CHUNK_OVERLAP,\n mini_chunk_size: int = MINI_CHUNK_SIZE,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> None:\n self.include_metadata = include_metadata\n self.chunk_token_limit = chunk_token_limit\n self.enable_multipass = enable_multipass\n self.enable_large_chunks = enable_large_chunks\n self.enable_contextual_rag = enable_contextual_rag\n if enable_contextual_rag:\n assert (\n USE_CHUNK_SUMMARY or USE_DOCUMENT_SUMMARY\n ), \"Contextual RAG requires at least one of chunk summary and document summary enabled\"\n self.default_contextual_rag_reserved_tokens = MAX_CONTEXT_TOKENS * (\n int(USE_CHUNK_SUMMARY) + int(USE_DOCUMENT_SUMMARY)\n )\n self.tokenizer = tokenizer\n self.callback = callback\n\n self.max_context = 0\n self.prompt_tokens = 0\n\n # Create a token counter function that returns the count instead of the tokens\n def token_counter(text: str) -> int:\n return len(tokenizer.encode(text))\n\n self.blurb_splitter = SentenceChunker(\n tokenizer_or_token_counter=token_counter,\n chunk_size=blurb_size,\n chunk_overlap=0,\n return_type=\"texts\",\n )\n\n self.chunk_splitter = SentenceChunker(\n tokenizer_or_token_counter=token_counter,\n chunk_size=chunk_token_limit,\n chunk_overlap=chunk_overlap,\n return_type=\"texts\",\n )\n\n self.mini_chunk_splitter = (\n SentenceChunker(\n tokenizer_or_token_counter=token_counter,\n chunk_size=mini_chunk_size,\n chunk_overlap=0,\n return_type=\"texts\",\n )\n if enable_multipass\n else None\n )\n\n def _split_oversized_chunk(self, text: str, content_token_limit: int) -> list[str]:\n \"\"\"\n Splits the text into smaller chunks based on token count to ensure\n no chunk exceeds the content_token_limit.\n \"\"\"\n tokens = self.tokenizer.tokenize(text)\n chunks = []\n start = 0\n total_tokens = len(tokens)\n while start < total_tokens:\n end = min(start + content_token_limit, total_tokens)\n token_chunk = tokens[start:end]\n chunk_text = \" \".join(token_chunk)\n chunks.append(chunk_text)\n start = end\n return chunks\n\n def _extract_blurb(self, text: str) -> str:\n \"\"\"\n Extract a short blurb from the text (first chunk of size `blurb_size`).\n \"\"\"\n # chunker is in `text` mode\n texts = cast(list[str], self.blurb_splitter.chunk(text))\n if not texts:\n return \"\"\n return texts[0]\n\n def _get_mini_chunk_texts(self, chunk_text: str) -> list[str] | None:\n \"\"\"\n For \"multipass\" mode: additional sub-chunks (mini-chunks) for use in certain embeddings.\n \"\"\"\n if self.mini_chunk_splitter and chunk_text.strip():\n # chunker is in `text` mode\n return cast(list[str], self.mini_chunk_splitter.chunk(chunk_text))\n return None\n\n # ADDED: extra param image_url to store in the chunk\n def _create_chunk(\n self,\n document: IndexingDocument,\n chunks_list: list[DocAwareChunk],\n text: str,\n links: dict[int, str],\n is_continuation: bool = False,\n title_prefix: str = \"\",\n metadata_suffix_semantic: str = \"\",\n metadata_suffix_keyword: str = \"\",\n image_file_id: str | None = None,\n ) -> None:\n \"\"\"\n Helper to create a new DocAwareChunk, append it to chunks_list.\n \"\"\"\n new_chunk = DocAwareChunk(\n source_document=document,\n chunk_id=len(chunks_list),\n blurb=self._extract_blurb(text),\n content=text,\n source_links=links or {0: \"\"},\n image_file_id=image_file_id,\n section_continuation=is_continuation,\n title_prefix=title_prefix,\n metadata_suffix_semantic=metadata_suffix_semantic,\n metadata_suffix_keyword=metadata_suffix_keyword,\n mini_chunk_texts=self._get_mini_chunk_texts(text),\n large_chunk_id=None,\n doc_summary=\"\",\n chunk_context=\"\",\n contextual_rag_reserved_tokens=0, # set per-document in _handle_single_document\n )\n chunks_list.append(new_chunk)\n\n def _chunk_document_with_sections(\n self,\n document: IndexingDocument,\n sections: list[Section],\n title_prefix: str,\n metadata_suffix_semantic: str,\n metadata_suffix_keyword: str,\n content_token_limit: int,\n ) -> list[DocAwareChunk]:\n \"\"\"\n Loops through sections of the document, converting them into one or more chunks.\n Works with processed sections that are base Section objects.\n \"\"\"\n chunks: list[DocAwareChunk] = []\n link_offsets: dict[int, str] = {}\n chunk_text = \"\"\n\n for section_idx, section in enumerate(sections):\n # Get section text and other attributes\n section_text = clean_text(str(section.text or \"\"))\n section_link_text = section.link or \"\"\n image_url = section.image_file_id\n\n # If there is no useful content, skip\n if not section_text and (not document.title or section_idx > 0):\n logger.warning(\n f\"Skipping empty or irrelevant section in doc {document.semantic_identifier}, link={section_link_text}\"\n )\n continue\n\n # CASE 1: If this section has an image, force a separate chunk\n if image_url:\n # First, if we have any partially built text chunk, finalize it\n if chunk_text.strip():\n self._create_chunk(\n document,\n chunks,\n chunk_text,\n link_offsets,\n is_continuation=False,\n title_prefix=title_prefix,\n metadata_suffix_semantic=metadata_suffix_semantic,\n metadata_suffix_keyword=metadata_suffix_keyword,\n )\n chunk_text = \"\"\n link_offsets = {}\n\n # Create a chunk specifically for this image section\n # (Using the text summary that was generated during processing)\n self._create_chunk(\n document,\n chunks,\n section_text,\n links={0: section_link_text} if section_link_text else {},\n image_file_id=image_url,\n title_prefix=title_prefix,\n metadata_suffix_semantic=metadata_suffix_semantic,\n metadata_suffix_keyword=metadata_suffix_keyword,\n )\n # Continue to next section\n continue\n\n # CASE 2: Normal text section\n section_token_count = len(self.tokenizer.encode(section_text))\n\n # If the section is large on its own, split it separately\n if section_token_count > content_token_limit:\n if chunk_text.strip():\n self._create_chunk(\n document,\n chunks,\n chunk_text,\n link_offsets,\n False,\n title_prefix,\n metadata_suffix_semantic,\n metadata_suffix_keyword,\n )\n chunk_text = \"\"\n link_offsets = {}\n\n # chunker is in `text` mode\n split_texts = cast(list[str], self.chunk_splitter.chunk(section_text))\n for i, split_text in enumerate(split_texts):\n # If even the split_text is bigger than strict limit, further split\n if (\n STRICT_CHUNK_TOKEN_LIMIT\n and len(self.tokenizer.encode(split_text)) > content_token_limit\n ):\n smaller_chunks = self._split_oversized_chunk(\n split_text, content_token_limit\n )\n for j, small_chunk in enumerate(smaller_chunks):\n self._create_chunk(\n document,\n chunks,\n small_chunk,\n {0: section_link_text},\n is_continuation=(j != 0),\n title_prefix=title_prefix,\n metadata_suffix_semantic=metadata_suffix_semantic,\n metadata_suffix_keyword=metadata_suffix_keyword,\n )\n else:\n self._create_chunk(\n document,\n chunks,\n split_text,\n {0: section_link_text},\n is_continuation=(i != 0),\n title_prefix=title_prefix,\n metadata_suffix_semantic=metadata_suffix_semantic,\n metadata_suffix_keyword=metadata_suffix_keyword,\n )\n continue\n\n # If we can still fit this section into the current chunk, do so\n current_token_count = len(self.tokenizer.encode(chunk_text))\n current_offset = len(shared_precompare_cleanup(chunk_text))\n next_section_tokens = (\n len(self.tokenizer.encode(SECTION_SEPARATOR)) + section_token_count\n )\n\n if next_section_tokens + current_token_count <= content_token_limit:\n if chunk_text:\n chunk_text += SECTION_SEPARATOR\n chunk_text += section_text\n link_offsets[current_offset] = section_link_text\n else:\n # finalize the existing chunk\n self._create_chunk(\n document,\n chunks,\n chunk_text,\n link_offsets,\n False,\n title_prefix,\n metadata_suffix_semantic,\n metadata_suffix_keyword,\n )\n # start a new chunk\n link_offsets = {0: section_link_text}\n chunk_text = section_text\n\n # finalize any leftover text chunk\n if chunk_text.strip() or not chunks:\n self._create_chunk(\n document,\n chunks,\n chunk_text,\n link_offsets or {0: \"\"}, # safe default\n False,\n title_prefix,\n metadata_suffix_semantic,\n metadata_suffix_keyword,\n )\n return chunks\n\n def _handle_single_document(\n self, document: IndexingDocument\n ) -> list[DocAwareChunk]:\n # Specifically for reproducing an issue with gmail\n if document.source == DocumentSource.GMAIL:\n logger.debug(f\"Chunking {document.semantic_identifier}\")\n\n # Title prep\n title = self._extract_blurb(document.get_title_for_document_index() or \"\")\n title_prefix = title + RETURN_SEPARATOR if title else \"\"\n title_tokens = len(self.tokenizer.encode(title_prefix))\n\n # Metadata prep\n metadata_suffix_semantic = \"\"\n metadata_suffix_keyword = \"\"\n metadata_tokens = 0\n if self.include_metadata:\n (\n metadata_suffix_semantic,\n metadata_suffix_keyword,\n ) = _get_metadata_suffix_for_document_index(\n document.metadata, include_separator=True\n )\n metadata_tokens = len(self.tokenizer.encode(metadata_suffix_semantic))\n\n # If metadata is too large, skip it in the semantic content\n if metadata_tokens >= self.chunk_token_limit * MAX_METADATA_PERCENTAGE:\n metadata_suffix_semantic = \"\"\n metadata_tokens = 0\n\n single_chunk_fits = True\n doc_token_count = 0\n if self.enable_contextual_rag:\n doc_content = document.get_text_content()\n tokenized_doc = self.tokenizer.tokenize(doc_content)\n doc_token_count = len(tokenized_doc)\n\n # check if doc + title + metadata fits in a single chunk. If so, no need for contextual RAG\n single_chunk_fits = (\n doc_token_count + title_tokens + metadata_tokens\n <= self.chunk_token_limit\n )\n\n # expand the size of the context used for contextual rag based on whether chunk context and doc summary are used\n context_size = 0\n if (\n self.enable_contextual_rag\n and not single_chunk_fits\n and not AVERAGE_SUMMARY_EMBEDDINGS\n ):\n context_size += self.default_contextual_rag_reserved_tokens\n\n # Adjust content token limit to accommodate title + metadata\n content_token_limit = (\n self.chunk_token_limit - title_tokens - metadata_tokens - context_size\n )\n\n # first check: if there is not enough actual chunk content when including contextual rag,\n # then don't do contextual rag\n if content_token_limit <= CHUNK_MIN_CONTENT:\n context_size = 0 # Don't do contextual RAG\n # revert to previous content token limit\n content_token_limit = (\n self.chunk_token_limit - title_tokens - metadata_tokens\n )\n\n # If there is not enough context remaining then just index the chunk with no prefix/suffix\n if content_token_limit <= CHUNK_MIN_CONTENT:\n # Not enough space left, so revert to full chunk without the prefix\n content_token_limit = self.chunk_token_limit\n title_prefix = \"\"\n metadata_suffix_semantic = \"\"\n\n # Use processed_sections if available (IndexingDocument), otherwise use original sections\n sections_to_chunk = document.processed_sections\n\n normal_chunks = self._chunk_document_with_sections(\n document,\n sections_to_chunk,\n title_prefix,\n metadata_suffix_semantic,\n metadata_suffix_keyword,\n content_token_limit,\n )\n\n # Optional \"multipass\" large chunk creation\n if self.enable_multipass and self.enable_large_chunks:\n large_chunks = generate_large_chunks(normal_chunks)\n normal_chunks.extend(large_chunks)\n\n for chunk in normal_chunks:\n chunk.contextual_rag_reserved_tokens = context_size\n\n return normal_chunks\n\n def chunk(self, documents: list[IndexingDocument]) -> list[DocAwareChunk]:\n \"\"\"\n Takes in a list of documents and chunks them into smaller chunks for indexing\n while persisting the document metadata.\n\n Works with both standard Document objects and IndexingDocument objects with processed_sections.\n \"\"\"\n final_chunks: list[DocAwareChunk] = []\n for document in documents:\n if self.callback and self.callback.should_stop():\n raise RuntimeError(\"Chunker.chunk: Stop signal detected\")\n\n chunks = self._handle_single_document(document)\n final_chunks.extend(chunks)\n\n if self.callback:\n self.callback.progress(\"Chunker.chunk\", len(chunks))\n\n return final_chunks\n" }, { "path": "backend/onyx/indexing/content_classification.py", "content": "" }, { "path": "backend/onyx/indexing/embedder.py", "content": "import time\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom collections import defaultdict\n\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorStopSignal\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.db.models import SearchSettings\nfrom onyx.document_index.chunk_content_enrichment import (\n generate_enriched_content_for_chunk_embedding,\n)\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.indexing.models import ChunkEmbedding\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.indexing.models import IndexChunk\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.pydantic_util import shallow_model_dump\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.configs import INDEXING_MODEL_SERVER_HOST\nfrom shared_configs.configs import INDEXING_MODEL_SERVER_PORT\nfrom shared_configs.enums import EmbeddingProvider\nfrom shared_configs.enums import EmbedTextType\nfrom shared_configs.model_server_models import Embedding\n\n\nlogger = setup_logger()\n\n\nclass IndexingEmbedder(ABC):\n \"\"\"Converts chunks into chunks with embeddings. Note that one chunk may have\n multiple embeddings associated with it.\"\"\"\n\n def __init__(\n self,\n model_name: str,\n normalize: bool,\n query_prefix: str | None,\n passage_prefix: str | None,\n provider_type: EmbeddingProvider | None,\n api_key: str | None,\n api_url: str | None,\n api_version: str | None,\n deployment_name: str | None,\n reduced_dimension: int | None,\n callback: IndexingHeartbeatInterface | None,\n ):\n self.model_name = model_name\n self.normalize = normalize\n self.query_prefix = query_prefix\n self.passage_prefix = passage_prefix\n self.provider_type = provider_type\n self.api_key = api_key\n self.api_url = api_url\n self.api_version = api_version\n self.deployment_name = deployment_name\n\n self.embedding_model = EmbeddingModel(\n model_name=model_name,\n query_prefix=query_prefix,\n passage_prefix=passage_prefix,\n normalize=normalize,\n api_key=api_key,\n provider_type=provider_type,\n api_url=api_url,\n api_version=api_version,\n deployment_name=deployment_name,\n reduced_dimension=reduced_dimension,\n # The below are globally set, this flow always uses the indexing one\n server_host=INDEXING_MODEL_SERVER_HOST,\n server_port=INDEXING_MODEL_SERVER_PORT,\n retrim_content=True,\n callback=callback,\n )\n\n @abstractmethod\n def embed_chunks(\n self,\n chunks: list[DocAwareChunk],\n tenant_id: str | None = None,\n request_id: str | None = None,\n ) -> list[IndexChunk]:\n raise NotImplementedError\n\n\nclass DefaultIndexingEmbedder(IndexingEmbedder):\n def __init__(\n self,\n model_name: str,\n normalize: bool,\n query_prefix: str | None,\n passage_prefix: str | None,\n provider_type: EmbeddingProvider | None = None,\n api_key: str | None = None,\n api_url: str | None = None,\n api_version: str | None = None,\n deployment_name: str | None = None,\n reduced_dimension: int | None = None,\n callback: IndexingHeartbeatInterface | None = None,\n ):\n super().__init__(\n model_name,\n normalize,\n query_prefix,\n passage_prefix,\n provider_type,\n api_key,\n api_url,\n api_version,\n deployment_name,\n reduced_dimension,\n callback,\n )\n\n @log_function_time()\n def embed_chunks(\n self,\n chunks: list[DocAwareChunk],\n tenant_id: str | None = None,\n request_id: str | None = None,\n ) -> list[IndexChunk]:\n \"\"\"Adds embeddings to the chunks, the title and metadata suffixes are added to the chunk as well\n if they exist. If there is no space for it, it would have been thrown out at the chunking step.\n \"\"\"\n # All chunks at this point must have some non-empty content\n flat_chunk_texts: list[str] = []\n large_chunks_present = False\n for chunk in chunks:\n if chunk.large_chunk_reference_ids:\n large_chunks_present = True\n chunk_text = (\n generate_enriched_content_for_chunk_embedding(chunk)\n ) or chunk.source_document.get_title_for_document_index()\n\n if not chunk_text:\n # This should never happen, the document would have been dropped\n # before getting to this point\n raise ValueError(f\"Chunk has no content: {chunk.to_short_descriptor()}\")\n\n flat_chunk_texts.append(chunk_text)\n\n if chunk.mini_chunk_texts:\n if chunk.large_chunk_reference_ids:\n # A large chunk does not contain mini chunks, if it matches the large chunk\n # with a high score, then mini chunks would not be used anyway\n # otherwise it should match the normal chunk\n raise RuntimeError(\"Large chunk contains mini chunks\")\n flat_chunk_texts.extend(chunk.mini_chunk_texts)\n\n embeddings = self.embedding_model.encode(\n texts=flat_chunk_texts,\n text_type=EmbedTextType.PASSAGE,\n large_chunks_present=large_chunks_present,\n tenant_id=tenant_id,\n request_id=request_id,\n )\n\n chunk_titles = {\n chunk.source_document.get_title_for_document_index() for chunk in chunks\n }\n\n # Drop any None or empty strings\n # If there is no title or the title is empty, the title embedding field will be null\n # which is ok, it just won't contribute at all to the scoring.\n chunk_titles_list = [title for title in chunk_titles if title]\n\n # Cache the Title embeddings to only have to do it once\n title_embed_dict: dict[str, Embedding] = {}\n if chunk_titles_list:\n title_embeddings = self.embedding_model.encode(\n chunk_titles_list,\n text_type=EmbedTextType.PASSAGE,\n tenant_id=tenant_id,\n request_id=request_id,\n )\n title_embed_dict.update(\n {\n title: vector\n for title, vector in zip(chunk_titles_list, title_embeddings)\n }\n )\n\n # Mapping embeddings to chunks\n embedded_chunks: list[IndexChunk] = []\n embedding_ind_start = 0\n for chunk in chunks:\n num_embeddings = 1 + (\n len(chunk.mini_chunk_texts) if chunk.mini_chunk_texts else 0\n )\n chunk_embeddings = embeddings[\n embedding_ind_start : embedding_ind_start + num_embeddings\n ]\n\n title = chunk.source_document.get_title_for_document_index()\n\n title_embedding = None\n if title:\n if title in title_embed_dict:\n # Using cached value to avoid recalculating for every chunk\n title_embedding = title_embed_dict[title]\n else:\n logger.error(\n \"Title had to be embedded separately, this should not happen!\"\n )\n title_embedding = self.embedding_model.encode(\n [title],\n text_type=EmbedTextType.PASSAGE,\n tenant_id=tenant_id,\n request_id=request_id,\n )[0]\n title_embed_dict[title] = title_embedding\n\n new_embedded_chunk = IndexChunk.model_construct(\n **shallow_model_dump(chunk),\n embeddings=ChunkEmbedding(\n full_embedding=chunk_embeddings[0],\n mini_chunk_embeddings=chunk_embeddings[1:],\n ),\n title_embedding=title_embedding,\n )\n embedded_chunks.append(new_embedded_chunk)\n embedding_ind_start += num_embeddings\n\n return embedded_chunks\n\n @classmethod\n def from_db_search_settings(\n cls,\n search_settings: SearchSettings,\n callback: IndexingHeartbeatInterface | None = None,\n ) -> \"DefaultIndexingEmbedder\":\n return cls(\n model_name=search_settings.model_name,\n normalize=search_settings.normalize,\n query_prefix=search_settings.query_prefix,\n passage_prefix=search_settings.passage_prefix,\n provider_type=search_settings.provider_type,\n api_key=search_settings.api_key,\n api_url=search_settings.api_url,\n api_version=search_settings.api_version,\n deployment_name=search_settings.deployment_name,\n reduced_dimension=search_settings.reduced_dimension,\n callback=callback,\n )\n\n\ndef embed_chunks_with_failure_handling(\n chunks: list[DocAwareChunk],\n embedder: IndexingEmbedder,\n tenant_id: str | None = None,\n request_id: str | None = None,\n) -> tuple[list[IndexChunk], list[ConnectorFailure]]:\n \"\"\"Tries to embed all chunks in one large batch. If that batch fails for any reason,\n goes document by document to isolate the failure(s).\n \"\"\"\n\n # TODO(rkuo): this doesn't disambiguate calls to the model server on retries.\n # Improve this if needed.\n\n # First try to embed all chunks in one batch\n try:\n return (\n embedder.embed_chunks(\n chunks=chunks, tenant_id=tenant_id, request_id=request_id\n ),\n [],\n )\n except ConnectorStopSignal as e:\n logger.warning(\n \"Connector stop signal detected in embed_chunks_with_failure_handling\"\n )\n raise e\n except Exception:\n logger.exception(\"Failed to embed chunk batch. Trying individual docs.\")\n # wait a couple seconds to let any rate limits or temporary issues resolve\n time.sleep(2)\n\n # Try embedding each document's chunks individually\n chunks_by_doc: dict[str, list[DocAwareChunk]] = defaultdict(list)\n for chunk in chunks:\n chunks_by_doc[chunk.source_document.id].append(chunk)\n\n embedded_chunks: list[IndexChunk] = []\n failures: list[ConnectorFailure] = []\n\n for doc_id, chunks_for_doc in chunks_by_doc.items():\n try:\n doc_embedded_chunks = embedder.embed_chunks(\n chunks=chunks_for_doc, tenant_id=tenant_id, request_id=request_id\n )\n embedded_chunks.extend(doc_embedded_chunks)\n except Exception as e:\n logger.exception(f\"Failed to embed chunks for document '{doc_id}'\")\n failures.append(\n ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=doc_id,\n document_link=(\n chunks_for_doc[0].get_link() if chunks_for_doc else None\n ),\n ),\n failure_message=str(e),\n exception=e,\n )\n )\n\n return embedded_chunks, failures\n" }, { "path": "backend/onyx/indexing/indexing_heartbeat.py", "content": "from abc import ABC\nfrom abc import abstractmethod\n\n\nclass IndexingHeartbeatInterface(ABC):\n \"\"\"Defines a callback interface to be passed to\n to run_indexing_entrypoint.\"\"\"\n\n @abstractmethod\n def should_stop(self) -> bool:\n \"\"\"Signal to stop the looping function in flight.\"\"\"\n\n @abstractmethod\n def progress(self, tag: str, amount: int) -> None:\n \"\"\"Send progress updates to the caller.\n Amount can be a positive number to indicate progress or <= 0\n just to act as a keep-alive.\n \"\"\"\n" }, { "path": "backend/onyx/indexing/indexing_pipeline.py", "content": "from collections import defaultdict\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\nfrom contextlib import contextmanager\nfrom typing import Protocol\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DEFAULT_CONTEXTUAL_RAG_LLM_NAME\nfrom onyx.configs.app_configs import DEFAULT_CONTEXTUAL_RAG_LLM_PROVIDER\nfrom onyx.configs.app_configs import ENABLE_CONTEXTUAL_RAG\nfrom onyx.configs.app_configs import MAX_CHUNKS_PER_DOC_BATCH\nfrom onyx.configs.app_configs import MAX_DOCUMENT_CHARS\nfrom onyx.configs.app_configs import MAX_TOKENS_FOR_FULL_INCLUSION\nfrom onyx.configs.app_configs import USE_CHUNK_SUMMARY\nfrom onyx.configs.app_configs import USE_DOCUMENT_SUMMARY\nfrom onyx.configs.llm_configs import get_image_extraction_and_analysis_enabled\nfrom onyx.connectors.cross_connector_utils.miscellaneous_utils import (\n get_experts_stores_representations,\n)\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import ConnectorStopSignal\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import IndexAttemptMetadata\nfrom onyx.connectors.models import IndexingDocument\nfrom onyx.connectors.models import Section\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.document import get_documents_by_ids\nfrom onyx.db.document import upsert_document_by_connector_credential_pair\nfrom onyx.db.document import upsert_documents\nfrom onyx.db.enums import HookPoint\nfrom onyx.db.hierarchy import link_hierarchy_nodes_to_documents\nfrom onyx.db.models import Document as DBDocument\nfrom onyx.db.models import IndexModelStatus\nfrom onyx.db.search_settings import get_active_search_settings\nfrom onyx.db.tag import upsert_document_tags\nfrom onyx.document_index.document_index_utils import (\n get_multipass_config,\n)\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import DocumentInsertionRecord\nfrom onyx.document_index.interfaces import DocumentMetadata\nfrom onyx.document_index.interfaces import IndexBatchParams\nfrom onyx.file_processing.image_summarization import summarize_image_with_error_handling\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.hooks.executor import execute_hook\nfrom onyx.hooks.executor import HookSkipped\nfrom onyx.hooks.executor import HookSoftFailed\nfrom onyx.hooks.points.document_ingestion import DocumentIngestionOwner\nfrom onyx.hooks.points.document_ingestion import DocumentIngestionPayload\nfrom onyx.hooks.points.document_ingestion import DocumentIngestionResponse\nfrom onyx.hooks.points.document_ingestion import DocumentIngestionSection\nfrom onyx.indexing.chunk_batch_store import ChunkBatchStore\nfrom onyx.indexing.chunker import Chunker\nfrom onyx.indexing.embedder import embed_chunks_with_failure_handling\nfrom onyx.indexing.embedder import IndexingEmbedder\nfrom onyx.indexing.models import DocAwareChunk\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.indexing.models import IndexingBatchAdapter\nfrom onyx.indexing.models import UpdatableChunkData\nfrom onyx.indexing.vector_db_insertion import write_chunks_to_vector_db_with_backoff\nfrom onyx.llm.factory import get_default_llm_with_vision\nfrom onyx.llm.factory import get_llm_for_contextual_rag\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.multi_llm import LLMRateLimitError\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.llm.utils import MAX_CONTEXT_TOKENS\nfrom onyx.natural_language_processing.utils import BaseTokenizer\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.natural_language_processing.utils import tokenizer_trim_middle\nfrom onyx.prompts.contextual_retrieval import CONTEXTUAL_RAG_PROMPT1\nfrom onyx.prompts.contextual_retrieval import CONTEXTUAL_RAG_PROMPT2\nfrom onyx.prompts.contextual_retrieval import DOCUMENT_SUMMARY_PROMPT\nfrom onyx.utils.batching import batch_generator\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.postgres_sanitization import sanitize_documents_for_postgres\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom onyx.utils.timing import log_function_time\n\n\nlogger = setup_logger()\n\n\nclass DocumentBatchPrepareContext(BaseModel):\n updatable_docs: list[Document]\n id_to_boost_map: dict[str, int]\n indexable_docs: list[IndexingDocument] = []\n model_config = ConfigDict(arbitrary_types_allowed=True)\n\n\nclass IndexingPipelineResult(BaseModel):\n # number of documents that are completely new (e.g. did\n # not exist as a part of this OR any other connector)\n new_docs: int\n # NOTE: need total_docs, since the pipeline can skip some docs\n # (e.g. not even insert them into Postgres)\n total_docs: int\n # number of chunks that were inserted into Vespa\n total_chunks: int\n\n failures: list[ConnectorFailure]\n\n @classmethod\n def empty(cls, total_docs: int) -> \"IndexingPipelineResult\":\n return cls(\n new_docs=0,\n total_docs=total_docs,\n total_chunks=0,\n failures=[],\n )\n\n\nclass ChunkEmbeddingResult(BaseModel):\n successful_chunk_ids: list[tuple[int, str]] # (chunk_id, document_id)\n connector_failures: list[ConnectorFailure]\n\n\nclass IndexingPipelineProtocol(Protocol):\n def __call__(\n self,\n document_batch: list[Document],\n index_attempt_metadata: IndexAttemptMetadata,\n ) -> IndexingPipelineResult: ...\n\n\ndef _upsert_documents_in_db(\n documents: list[Document],\n index_attempt_metadata: IndexAttemptMetadata,\n db_session: Session,\n) -> None:\n # Metadata here refers to basic document info, not metadata about the actual content\n document_metadata_list: list[DocumentMetadata] = []\n for doc in documents:\n first_link = next(\n (section.link for section in doc.sections if section.link), \"\"\n )\n db_doc_metadata = DocumentMetadata(\n connector_id=index_attempt_metadata.connector_id,\n credential_id=index_attempt_metadata.credential_id,\n document_id=doc.id,\n semantic_identifier=doc.semantic_identifier,\n first_link=first_link,\n primary_owners=get_experts_stores_representations(doc.primary_owners),\n secondary_owners=get_experts_stores_representations(doc.secondary_owners),\n from_ingestion_api=doc.from_ingestion_api,\n external_access=doc.external_access,\n doc_metadata=doc.doc_metadata,\n # parent_hierarchy_node_id is resolved in docfetching using Redis cache\n parent_hierarchy_node_id=doc.parent_hierarchy_node_id,\n )\n document_metadata_list.append(db_doc_metadata)\n\n upsert_documents(db_session, document_metadata_list)\n\n # Insert document content metadata\n for doc in documents:\n upsert_document_tags(\n document_id=doc.id,\n source=doc.source,\n metadata=doc.metadata,\n db_session=db_session,\n )\n\n\ndef _get_failed_doc_ids(failures: list[ConnectorFailure]) -> set[str]:\n \"\"\"Extract document IDs from a list of connector failures.\"\"\"\n return {f.failed_document.document_id for f in failures if f.failed_document}\n\n\ndef _embed_chunks_to_store(\n chunks: list[DocAwareChunk],\n embedder: IndexingEmbedder,\n tenant_id: str,\n request_id: str | None,\n store: ChunkBatchStore,\n) -> ChunkEmbeddingResult:\n \"\"\"Embed chunks in batches, spilling each batch to *store*.\n\n If a document fails embedding in any batch, its chunks are excluded from\n all batches (including earlier ones already written) so that the output\n is all-or-nothing per document.\n \"\"\"\n successful_chunk_ids: list[tuple[int, str]] = []\n all_embedding_failures: list[ConnectorFailure] = []\n # Track failed doc IDs across all batches so that a failure in batch N\n # causes chunks for that doc to be skipped in batch N+1 and stripped\n # from earlier batches.\n all_failed_doc_ids: set[str] = set()\n\n for batch_idx, chunk_batch in enumerate(\n batch_generator(chunks, MAX_CHUNKS_PER_DOC_BATCH)\n ):\n # Skip chunks belonging to documents that failed in earlier batches.\n chunk_batch = [\n c for c in chunk_batch if c.source_document.id not in all_failed_doc_ids\n ]\n if not chunk_batch:\n continue\n\n logger.debug(f\"Embedding batch {batch_idx}: {len(chunk_batch)} chunks\")\n\n chunks_with_embeddings, embedding_failures = embed_chunks_with_failure_handling(\n chunks=chunk_batch,\n embedder=embedder,\n tenant_id=tenant_id,\n request_id=request_id,\n )\n all_embedding_failures.extend(embedding_failures)\n all_failed_doc_ids.update(_get_failed_doc_ids(embedding_failures))\n\n # Only keep successfully embedded chunks for non-failed docs.\n chunks_with_embeddings = [\n c\n for c in chunks_with_embeddings\n if c.source_document.id not in all_failed_doc_ids\n ]\n\n successful_chunk_ids.extend(\n (c.chunk_id, c.source_document.id) for c in chunks_with_embeddings\n )\n\n store.save(chunks_with_embeddings, batch_idx)\n del chunks_with_embeddings\n\n # Scrub earlier batches for docs that failed in later batches.\n if all_failed_doc_ids:\n store.scrub_failed_docs(all_failed_doc_ids)\n successful_chunk_ids = [\n (chunk_id, doc_id)\n for chunk_id, doc_id in successful_chunk_ids\n if doc_id not in all_failed_doc_ids\n ]\n\n return ChunkEmbeddingResult(\n successful_chunk_ids=successful_chunk_ids,\n connector_failures=all_embedding_failures,\n )\n\n\n@contextmanager\ndef embed_and_stream(\n chunks: list[DocAwareChunk],\n embedder: IndexingEmbedder,\n tenant_id: str,\n request_id: str | None,\n) -> Generator[tuple[ChunkEmbeddingResult, ChunkBatchStore], None, None]:\n \"\"\"Embed chunks to disk and yield a ``(result, store)`` pair.\n\n The store owns the temp directory — files are cleaned up when the context\n manager exits.\n\n Usage::\n\n with embed_and_stream(chunks, embedder, tenant_id, req_id) as (result, store):\n for chunk in store.stream():\n ...\n \"\"\"\n with ChunkBatchStore() as store:\n result = _embed_chunks_to_store(\n chunks=chunks,\n embedder=embedder,\n tenant_id=tenant_id,\n request_id=request_id,\n store=store,\n )\n yield result, store\n\n\ndef get_doc_ids_to_update(\n documents: list[Document], db_docs: list[DBDocument]\n) -> list[Document]:\n \"\"\"Figures out which documents actually need to be updated. If a document is already present\n and the `updated_at` hasn't changed, we shouldn't need to do anything with it.\n\n NB: Still need to associate the document in the DB if multiple connectors are\n indexing the same doc.\"\"\"\n id_update_time_map = {\n doc.id: doc.doc_updated_at for doc in db_docs if doc.doc_updated_at\n }\n\n updatable_docs: list[Document] = []\n for doc in documents:\n if (\n doc.id in id_update_time_map\n and doc.doc_updated_at\n and doc.doc_updated_at <= id_update_time_map[doc.id]\n ):\n continue\n updatable_docs.append(doc)\n\n return updatable_docs\n\n\ndef index_doc_batch_with_handler(\n *,\n chunker: Chunker,\n embedder: IndexingEmbedder,\n document_indices: list[DocumentIndex],\n document_batch: list[Document],\n request_id: str | None,\n tenant_id: str,\n db_session: Session,\n adapter: IndexingBatchAdapter,\n ignore_time_skip: bool = False,\n enable_contextual_rag: bool = False,\n llm: LLM | None = None,\n) -> IndexingPipelineResult:\n try:\n index_pipeline_result = index_doc_batch(\n chunker=chunker,\n embedder=embedder,\n document_indices=document_indices,\n document_batch=document_batch,\n request_id=request_id,\n tenant_id=tenant_id,\n db_session=db_session,\n adapter=adapter,\n ignore_time_skip=ignore_time_skip,\n enable_contextual_rag=enable_contextual_rag,\n llm=llm,\n )\n\n except ConnectorStopSignal as e:\n logger.warning(\"Connector stop signal detected in index_doc_batch_with_handler\")\n raise e\n except Exception as e:\n # don't log the batch directly, it's too much text\n document_ids = [doc.id for doc in document_batch]\n logger.exception(f\"Failed to index document batch: {document_ids}\")\n\n index_pipeline_result = IndexingPipelineResult(\n new_docs=0,\n total_docs=len(document_batch),\n total_chunks=0,\n failures=[\n ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=document.id,\n document_link=(\n document.sections[0].link if document.sections else None\n ),\n ),\n failure_message=str(e),\n exception=e,\n )\n for document in document_batch\n ],\n )\n\n return index_pipeline_result\n\n\ndef index_doc_batch_prepare(\n documents: list[Document],\n index_attempt_metadata: IndexAttemptMetadata,\n db_session: Session,\n ignore_time_skip: bool = False,\n) -> DocumentBatchPrepareContext | None:\n \"\"\"Sets up the documents in the relational DB (source of truth) for permissions, metadata, etc.\n This preceeds indexing it into the actual document index.\"\"\"\n documents = sanitize_documents_for_postgres(documents)\n\n # Create a trimmed list of docs that don't have a newer updated at\n # Shortcuts the time-consuming flow on connector index retries\n document_ids: list[str] = [document.id for document in documents]\n db_docs: list[DBDocument] = get_documents_by_ids(\n db_session=db_session,\n document_ids=document_ids,\n )\n\n updatable_docs = (\n get_doc_ids_to_update(documents=documents, db_docs=db_docs)\n if not ignore_time_skip\n else documents\n )\n if len(updatable_docs) != len(documents):\n updatable_doc_ids = [doc.id for doc in updatable_docs]\n skipped_doc_ids = [\n doc.id for doc in documents if doc.id not in updatable_doc_ids\n ]\n logger.info(\n f\"Skipping {len(skipped_doc_ids)} documents because they are up to date. Skipped doc IDs: {skipped_doc_ids}\"\n )\n\n # for all updatable docs, upsert into the DB\n # Does not include doc_updated_at which is also used to indicate a successful update\n if updatable_docs:\n _upsert_documents_in_db(\n documents=updatable_docs,\n index_attempt_metadata=index_attempt_metadata,\n db_session=db_session,\n )\n\n logger.info(\n f\"Upserted {len(updatable_docs)} changed docs out of {len(documents)} total docs into the DB\"\n )\n\n # for all docs, upsert the document to cc pair relationship\n upsert_document_by_connector_credential_pair(\n db_session,\n index_attempt_metadata.connector_id,\n index_attempt_metadata.credential_id,\n document_ids,\n )\n\n # Link hierarchy nodes to documents for sources where pages can be both\n # hierarchy nodes AND documents (e.g., Notion, Confluence).\n # This must happen after documents are upserted due to FK constraint.\n if documents:\n link_hierarchy_nodes_to_documents(\n db_session=db_session,\n document_ids=document_ids,\n source=documents[0].source,\n commit=False, # We'll commit with the rest of the transaction\n )\n\n # No docs to process because the batch is empty or every doc was already indexed\n if not updatable_docs:\n return None\n\n id_to_boost_map = {doc.id: doc.boost for doc in db_docs}\n return DocumentBatchPrepareContext(\n updatable_docs=updatable_docs, id_to_boost_map=id_to_boost_map\n )\n\n\ndef filter_documents(document_batch: list[Document]) -> list[Document]:\n documents: list[Document] = []\n total_chars_in_batch = 0\n skipped_too_long = []\n\n for document in document_batch:\n empty_contents = not any(\n isinstance(section, TextSection)\n and section.text is not None\n and section.text.strip()\n for section in document.sections\n )\n if (\n (not document.title or not document.title.strip())\n and not document.semantic_identifier.strip()\n and empty_contents\n ):\n # Skip documents that have neither title nor content\n # If the document doesn't have either, then there is no useful information in it\n # This is again verified later in the pipeline after chunking but at that point there should\n # already be no documents that are empty.\n logger.warning(\n f\"Skipping document with ID {document.id} as it has neither title nor content.\"\n )\n continue\n\n if document.title is not None and not document.title.strip() and empty_contents:\n # The title is explicitly empty (\"\" and not None) and the document is empty\n # so when building the chunk text representation, it will be empty and unuseable\n logger.warning(\n f\"Skipping document with ID {document.id} as the chunks will be empty.\"\n )\n continue\n\n section_chars = sum(\n (\n len(section.text)\n if isinstance(section, TextSection) and section.text is not None\n else 0\n )\n for section in document.sections\n )\n doc_total_chars = (\n len(document.title or document.semantic_identifier) + section_chars\n )\n\n if MAX_DOCUMENT_CHARS and doc_total_chars > MAX_DOCUMENT_CHARS:\n # Skip documents that are too long, later on there are more memory intensive steps done on the text\n # and the container will run out of memory and crash. Several other checks are included upstream but\n # those are at the connector level so a catchall is still needed.\n # Assumption here is that files that are that long, are generated files and not the type users\n # generally care for.\n logger.warning(\n f\"Skipping document with ID {document.id} as it is too long \"\n f\"({doc_total_chars:,} chars, max={MAX_DOCUMENT_CHARS:,})\"\n )\n skipped_too_long.append((document.id, doc_total_chars))\n continue\n\n total_chars_in_batch += doc_total_chars\n documents.append(document)\n\n # Log batch statistics for OOM debugging\n if documents:\n avg_chars = total_chars_in_batch / len(documents)\n # Get the source from the first document (all in batch should be same source)\n source = documents[0].source.value if documents[0].source else \"unknown\"\n logger.debug(\n f\"Document batch filter [{source}]: {len(documents)} docs kept, {len(skipped_too_long)} skipped (too long). \"\n f\"Total chars: {total_chars_in_batch:,}, Avg: {avg_chars:,.0f} chars/doc\"\n )\n if skipped_too_long:\n logger.warning(\n f\"Skipped oversized documents [{source}]: {skipped_too_long[:5]}\"\n ) # Log first 5\n\n return documents\n\n\ndef process_image_sections(documents: list[Document]) -> list[IndexingDocument]:\n \"\"\"\n Process all sections in documents by:\n 1. Converting both TextSection and ImageSection objects to base Section objects\n 2. Processing ImageSections to generate text summaries using a vision-capable LLM\n 3. Returning IndexingDocument objects with both original and processed sections\n\n Args:\n documents: List of documents with TextSection | ImageSection objects\n\n Returns:\n List of IndexingDocument objects with processed_sections as list[Section]\n \"\"\"\n # Check if image extraction and analysis is enabled before trying to get a vision LLM\n if not get_image_extraction_and_analysis_enabled():\n llm = None\n else:\n # Only get the vision LLM if image processing is enabled\n llm = get_default_llm_with_vision()\n\n if not llm:\n if get_image_extraction_and_analysis_enabled():\n logger.warning(\n \"Image analysis is enabled but no vision-capable LLM is \"\n \"available — images will not be summarized. Configure a \"\n \"vision model in the admin LLM settings.\"\n )\n # Even without LLM, we still convert to IndexingDocument with base Sections\n return [\n IndexingDocument(\n **document.model_dump(),\n processed_sections=[\n Section(\n text=section.text if isinstance(section, TextSection) else \"\",\n link=section.link,\n image_file_id=(\n section.image_file_id\n if isinstance(section, ImageSection)\n else None\n ),\n )\n for section in document.sections\n ],\n )\n for document in documents\n ]\n\n indexed_documents: list[IndexingDocument] = []\n\n for document in documents:\n processed_sections: list[Section] = []\n\n for section in document.sections:\n # For ImageSection, process and create base Section with both text and image_file_id\n if isinstance(section, ImageSection):\n # Default section with image path preserved - ensure text is always a string\n processed_section = Section(\n link=section.link,\n image_file_id=section.image_file_id,\n text=\"\", # Initialize with empty string\n )\n\n # Try to get image summary\n try:\n file_store = get_default_file_store()\n\n file_record = file_store.read_file_record(\n file_id=section.image_file_id\n )\n if not file_record:\n logger.warning(\n f\"Image file {section.image_file_id} not found in FileStore\"\n )\n\n processed_section.text = \"[Image could not be processed]\"\n else:\n # Get the image data\n image_data_io = file_store.read_file(\n file_id=section.image_file_id\n )\n image_data = image_data_io.read()\n summary = summarize_image_with_error_handling(\n llm=llm,\n image_data=image_data,\n context_name=file_record.display_name or \"Image\",\n )\n\n if summary:\n processed_section.text = summary\n else:\n processed_section.text = \"[Image could not be summarized]\"\n except Exception as e:\n logger.error(f\"Error processing image section: {e}\")\n processed_section.text = \"[Error processing image]\"\n\n processed_sections.append(processed_section)\n\n # For TextSection, create a base Section with text and link\n elif isinstance(section, TextSection):\n processed_section = Section(\n text=section.text or \"\", # Ensure text is always a string, not None\n link=section.link,\n image_file_id=None,\n )\n processed_sections.append(processed_section)\n\n # Create IndexingDocument with original sections and processed_sections\n indexed_document = IndexingDocument(\n **document.model_dump(), processed_sections=processed_sections\n )\n indexed_documents.append(indexed_document)\n\n return indexed_documents\n\n\ndef add_document_summaries(\n chunks_by_doc: list[DocAwareChunk],\n llm: LLM,\n tokenizer: BaseTokenizer,\n trunc_doc_tokens: int,\n) -> list[int] | None:\n \"\"\"\n Adds a document summary to a list of chunks from the same document.\n Returns the number of tokens in the document.\n \"\"\"\n\n doc_tokens = []\n # this is value is the same for each chunk in the document; 0 indicates\n # There is not enough space for contextual RAG (the chunk content\n # and possibly metadata took up too much space)\n if chunks_by_doc[0].contextual_rag_reserved_tokens == 0:\n return None\n\n doc_tokens = tokenizer.encode(chunks_by_doc[0].source_document.get_text_content())\n doc_content = tokenizer_trim_middle(doc_tokens, trunc_doc_tokens, tokenizer)\n\n # Apply prompt caching: cache the static prompt, document content is the suffix\n # Note: For document summarization, there's no cacheable prefix since the document changes\n # So we just pass the full prompt without caching\n summary_prompt = DOCUMENT_SUMMARY_PROMPT.format(document=doc_content)\n prompt_msg = UserMessage(content=summary_prompt)\n\n response = llm.invoke(prompt_msg, max_tokens=MAX_CONTEXT_TOKENS)\n doc_summary = llm_response_to_string(response)\n\n for chunk in chunks_by_doc:\n chunk.doc_summary = doc_summary\n\n return doc_tokens\n\n\ndef add_chunk_summaries(\n chunks_by_doc: list[DocAwareChunk],\n llm: LLM,\n tokenizer: BaseTokenizer,\n trunc_doc_chunk_tokens: int,\n doc_tokens: list[int] | None,\n) -> None:\n \"\"\"\n Adds chunk summaries to the chunks grouped by document id.\n Chunk summaries look at the chunk as well as the entire document (or a summary,\n if the document is too long) and describe how the chunk relates to the document.\n \"\"\"\n # all chunks within a document have the same contextual_rag_reserved_tokens\n if chunks_by_doc[0].contextual_rag_reserved_tokens == 0:\n return\n\n # use values computed in above doc summary section if available\n doc_tokens = doc_tokens or tokenizer.encode(\n chunks_by_doc[0].source_document.get_text_content()\n )\n doc_content = tokenizer_trim_middle(doc_tokens, trunc_doc_chunk_tokens, tokenizer)\n\n # only compute doc summary if needed\n doc_info = (\n doc_content\n if len(doc_tokens) <= MAX_TOKENS_FOR_FULL_INCLUSION\n else chunks_by_doc[0].doc_summary\n )\n if not doc_info:\n # This happens if the document is too long AND document summaries are turned off\n # In this case we compute a doc summary using the LLM\n fallback_prompt = UserMessage(\n content=DOCUMENT_SUMMARY_PROMPT.format(document=doc_content)\n )\n response = llm.invoke(fallback_prompt, max_tokens=MAX_CONTEXT_TOKENS)\n doc_info = llm_response_to_string(response)\n\n from onyx.llm.prompt_cache.processor import process_with_prompt_cache\n\n context_prompt1 = CONTEXTUAL_RAG_PROMPT1.format(document=doc_info)\n\n def assign_context(chunk: DocAwareChunk) -> None:\n context_prompt2 = CONTEXTUAL_RAG_PROMPT2.format(chunk=chunk.content)\n try:\n # Apply prompt caching: cache the document context (prompt1), chunk content is the suffix\n # For string inputs with continuation=True, the result will be a concatenated string\n processed_prompt, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=UserMessage(content=context_prompt1),\n suffix=UserMessage(content=context_prompt2),\n continuation=True, # Append chunk to the document context\n )\n\n response = llm.invoke(processed_prompt, max_tokens=MAX_CONTEXT_TOKENS)\n chunk.chunk_context = llm_response_to_string(response)\n\n except LLMRateLimitError as e:\n # Erroring during chunker is undesirable, so we log the error and continue\n # TODO: for v2, add robust retry logic\n logger.exception(f\"Rate limit adding chunk summary: {e}\", exc_info=e)\n chunk.chunk_context = \"\"\n except Exception as e:\n logger.exception(f\"Error adding chunk summary: {e}\", exc_info=e)\n chunk.chunk_context = \"\"\n\n run_functions_tuples_in_parallel(\n [(assign_context, (chunk,)) for chunk in chunks_by_doc]\n )\n\n\ndef add_contextual_summaries(\n chunks: list[DocAwareChunk],\n llm: LLM,\n tokenizer: BaseTokenizer,\n chunk_token_limit: int,\n) -> list[DocAwareChunk]:\n \"\"\"\n Adds Document summary and chunk-within-document context to the chunks\n based on which environment variables are set.\n \"\"\"\n doc2chunks = defaultdict(list)\n for chunk in chunks:\n doc2chunks[chunk.source_document.id].append(chunk)\n\n # The number of tokens allowed for the document when computing a document summary\n trunc_doc_summary_tokens = llm.config.max_input_tokens - len(\n tokenizer.encode(DOCUMENT_SUMMARY_PROMPT)\n )\n\n prompt_tokens = len(\n tokenizer.encode(CONTEXTUAL_RAG_PROMPT1 + CONTEXTUAL_RAG_PROMPT2)\n )\n # The number of tokens allowed for the document when computing a\n # \"chunk in context of document\" summary\n trunc_doc_chunk_tokens = (\n llm.config.max_input_tokens - prompt_tokens - chunk_token_limit\n )\n for chunks_by_doc in doc2chunks.values():\n doc_tokens = None\n if USE_DOCUMENT_SUMMARY:\n doc_tokens = add_document_summaries(\n chunks_by_doc, llm, tokenizer, trunc_doc_summary_tokens\n )\n\n if USE_CHUNK_SUMMARY:\n add_chunk_summaries(\n chunks_by_doc, llm, tokenizer, trunc_doc_chunk_tokens, doc_tokens\n )\n\n return chunks\n\n\ndef _verify_indexing_completeness(\n insertion_records: list[DocumentInsertionRecord],\n write_failures: list[ConnectorFailure],\n embedding_failed_doc_ids: set[str],\n updatable_ids: list[str],\n document_index_name: str,\n) -> None:\n \"\"\"Verify that every updatable document was either indexed or reported as failed.\"\"\"\n all_returned_doc_ids = (\n {r.document_id for r in insertion_records}\n | {f.failed_document.document_id for f in write_failures if f.failed_document}\n | embedding_failed_doc_ids\n )\n if all_returned_doc_ids != set(updatable_ids):\n raise RuntimeError(\n f\"Some documents were not successfully indexed. \"\n f\"Updatable IDs: {updatable_ids}, \"\n f\"Returned IDs: {all_returned_doc_ids}. \"\n f\"This should never happen. \"\n f\"This occured for document index {document_index_name}\"\n )\n\n\ndef _apply_document_ingestion_hook(\n documents: list[Document],\n db_session: Session,\n) -> list[Document]:\n \"\"\"Apply the Document Ingestion hook to each document in the batch.\n\n - HookSkipped / HookSoftFailed → document passes through unchanged.\n - Response with sections=None → document is dropped (logged).\n - Response with sections → document sections are replaced with the hook's output.\n \"\"\"\n\n def _build_payload(doc: Document) -> DocumentIngestionPayload:\n return DocumentIngestionPayload(\n document_id=doc.id or \"\",\n title=doc.title,\n semantic_identifier=doc.semantic_identifier,\n source=doc.source.value if doc.source is not None else \"\",\n sections=[\n DocumentIngestionSection(\n text=s.text if isinstance(s, TextSection) else None,\n link=s.link,\n image_file_id=(\n s.image_file_id if isinstance(s, ImageSection) else None\n ),\n )\n for s in doc.sections\n ],\n metadata={\n k: v if isinstance(v, list) else [v] for k, v in doc.metadata.items()\n },\n doc_updated_at=(\n doc.doc_updated_at.isoformat() if doc.doc_updated_at else None\n ),\n primary_owners=(\n [\n DocumentIngestionOwner(\n display_name=o.get_semantic_name() or None,\n email=o.email,\n )\n for o in doc.primary_owners\n ]\n if doc.primary_owners\n else None\n ),\n secondary_owners=(\n [\n DocumentIngestionOwner(\n display_name=o.get_semantic_name() or None,\n email=o.email,\n )\n for o in doc.secondary_owners\n ]\n if doc.secondary_owners\n else None\n ),\n )\n\n def _apply_result(\n doc: Document,\n hook_result: DocumentIngestionResponse | HookSkipped | HookSoftFailed,\n ) -> Document | None:\n \"\"\"Return the modified doc, original doc (skip/soft-fail), or None (drop).\"\"\"\n if isinstance(hook_result, (HookSkipped, HookSoftFailed)):\n return doc\n if not hook_result.sections:\n reason = hook_result.rejection_reason or \"Document rejected by hook\"\n logger.info(\n f\"Document ingestion hook dropped document doc_id={doc.id!r}: {reason}\"\n )\n return None\n new_sections: list[TextSection | ImageSection] = []\n for s in hook_result.sections:\n if s.image_file_id is not None:\n new_sections.append(\n ImageSection(image_file_id=s.image_file_id, link=s.link)\n )\n elif s.text is not None:\n new_sections.append(TextSection(text=s.text, link=s.link))\n else:\n logger.warning(\n f\"Document ingestion hook returned a section with neither text nor \"\n f\"image_file_id for doc_id={doc.id!r} — skipping section.\"\n )\n if not new_sections:\n logger.info(\n f\"Document ingestion hook produced no valid sections for doc_id={doc.id!r} — dropping document.\"\n )\n return None\n return doc.model_copy(update={\"sections\": new_sections})\n\n if not documents:\n return documents\n\n # Run the hook for the first document. If it returns HookSkipped the hook\n # is not configured — skip the remaining N-1 DB lookups.\n first_doc = documents[0]\n first_payload = _build_payload(first_doc).model_dump()\n first_hook_result = execute_hook(\n db_session=db_session,\n hook_point=HookPoint.DOCUMENT_INGESTION,\n payload=first_payload,\n response_type=DocumentIngestionResponse,\n )\n if isinstance(first_hook_result, HookSkipped):\n return documents\n\n result: list[Document] = []\n first_applied = _apply_result(first_doc, first_hook_result)\n if first_applied is not None:\n result.append(first_applied)\n\n for doc in documents[1:]:\n payload = _build_payload(doc).model_dump()\n hook_result = execute_hook(\n db_session=db_session,\n hook_point=HookPoint.DOCUMENT_INGESTION,\n payload=payload,\n response_type=DocumentIngestionResponse,\n )\n applied = _apply_result(doc, hook_result)\n if applied is not None:\n result.append(applied)\n\n return result\n\n\n@log_function_time(debug_only=True)\ndef index_doc_batch(\n *,\n document_batch: list[Document],\n chunker: Chunker,\n embedder: IndexingEmbedder,\n document_indices: list[DocumentIndex],\n request_id: str | None,\n tenant_id: str,\n db_session: Session,\n adapter: IndexingBatchAdapter,\n enable_contextual_rag: bool = False,\n llm: LLM | None = None,\n ignore_time_skip: bool = False,\n filter_fnc: Callable[[list[Document]], list[Document]] = filter_documents,\n) -> IndexingPipelineResult:\n \"\"\"End-to-end indexing for a pre-batched set of documents.\"\"\"\n \"\"\"Takes different pieces of the indexing pipeline and applies it to a batch of documents\n Note that the documents should already be batched at this point so that it does not inflate the\n memory requirements\n\n Returns a tuple where the first element is the number of new docs and the\n second element is the number of chunks.\"\"\"\n\n # Log connector info for debugging OOM issues\n connector_id = getattr(adapter, \"connector_id\", None)\n credential_id = getattr(adapter, \"credential_id\", None)\n logger.debug(\n f\"Starting index_doc_batch: connector_id={connector_id}, \"\n f\"credential_id={credential_id}, tenant_id={tenant_id}, \"\n f\"num_docs={len(document_batch)}\"\n )\n\n filtered_documents = filter_fnc(document_batch)\n filtered_documents = _apply_document_ingestion_hook(filtered_documents, db_session)\n context = adapter.prepare(filtered_documents, ignore_time_skip)\n if not context:\n return IndexingPipelineResult.empty(len(filtered_documents))\n\n # Convert documents to IndexingDocument objects with processed section\n # logger.debug(\"Processing image sections\")\n context.indexable_docs = process_image_sections(context.updatable_docs)\n\n doc_descriptors = [\n {\n \"doc_id\": doc.id,\n \"doc_length\": doc.get_total_char_length(),\n }\n for doc in context.indexable_docs\n ]\n logger.debug(f\"Starting indexing process for documents: {doc_descriptors}\")\n\n logger.debug(\"Starting chunking\")\n # NOTE: no special handling for failures here, since the chunker is not\n # a common source of failure for the indexing pipeline\n chunks: list[DocAwareChunk] = chunker.chunk(context.indexable_docs)\n llm_tokenizer: BaseTokenizer | None = None\n\n # contextual RAG\n if enable_contextual_rag:\n assert llm is not None, \"must provide an LLM for contextual RAG\"\n llm_tokenizer = get_tokenizer(\n model_name=llm.config.model_name,\n provider_type=llm.config.model_provider,\n )\n\n # Because the chunker's tokens are different from the LLM's tokens,\n # We add a fudge factor to ensure we truncate prompts to the LLM's token limit\n chunks = add_contextual_summaries(\n chunks=chunks,\n llm=llm,\n tokenizer=llm_tokenizer,\n chunk_token_limit=chunker.chunk_token_limit * 2,\n )\n\n logger.debug(\"Starting embedding\")\n with embed_and_stream(chunks, embedder, tenant_id, request_id) as (\n embedding_result,\n chunk_store,\n ):\n updatable_ids = [doc.id for doc in context.updatable_docs]\n updatable_chunk_data = [\n UpdatableChunkData(\n chunk_id=chunk_id,\n document_id=document_id,\n boost_score=1.0,\n )\n for chunk_id, document_id in embedding_result.successful_chunk_ids\n ]\n\n embedding_failed_doc_ids = _get_failed_doc_ids(\n embedding_result.connector_failures\n )\n\n # Filter to only successfully embedded chunks so\n # doc_id_to_new_chunk_cnt reflects what's actually written to Vespa.\n embedded_chunks = [\n c for c in chunks if c.source_document.id not in embedding_failed_doc_ids\n ]\n\n # Acquires a lock on the documents so that no other process can modify\n # them. Not needed until here, since this is when the actual race\n # condition with vector db can occur.\n with adapter.lock_context(context.updatable_docs):\n enricher = adapter.prepare_enrichment(\n context=context,\n tenant_id=tenant_id,\n chunks=embedded_chunks,\n )\n\n index_batch_params = IndexBatchParams(\n doc_id_to_previous_chunk_cnt=enricher.doc_id_to_previous_chunk_cnt,\n doc_id_to_new_chunk_cnt=enricher.doc_id_to_new_chunk_cnt,\n tenant_id=tenant_id,\n large_chunks_enabled=chunker.enable_large_chunks,\n )\n\n primary_doc_idx_insertion_records: list[DocumentInsertionRecord] | None = (\n None\n )\n primary_doc_idx_vector_db_write_failures: list[ConnectorFailure] | None = (\n None\n )\n\n for document_index in document_indices:\n\n def _enriched_stream() -> Iterator[DocMetadataAwareIndexChunk]:\n for chunk in chunk_store.stream():\n yield enricher.enrich_chunk(chunk, 1.0)\n\n insertion_records, write_failures = (\n write_chunks_to_vector_db_with_backoff(\n document_index=document_index,\n make_chunks=_enriched_stream,\n index_batch_params=index_batch_params,\n )\n )\n\n _verify_indexing_completeness(\n insertion_records=insertion_records,\n write_failures=write_failures,\n embedding_failed_doc_ids=embedding_failed_doc_ids,\n updatable_ids=updatable_ids,\n document_index_name=document_index.__class__.__name__,\n )\n # We treat the first document index we got as the primary one used\n # for reporting the state of indexing.\n if primary_doc_idx_insertion_records is None:\n primary_doc_idx_insertion_records = insertion_records\n if primary_doc_idx_vector_db_write_failures is None:\n primary_doc_idx_vector_db_write_failures = write_failures\n\n adapter.post_index(\n context=context,\n updatable_chunk_data=updatable_chunk_data,\n filtered_documents=filtered_documents,\n enrichment=enricher,\n )\n\n assert primary_doc_idx_insertion_records is not None\n assert primary_doc_idx_vector_db_write_failures is not None\n return IndexingPipelineResult(\n new_docs=sum(\n 1 for r in primary_doc_idx_insertion_records if not r.already_existed\n ),\n total_docs=len(filtered_documents),\n total_chunks=len(embedding_result.successful_chunk_ids),\n failures=primary_doc_idx_vector_db_write_failures\n + embedding_result.connector_failures,\n )\n\n\ndef run_indexing_pipeline(\n *,\n document_batch: list[Document],\n request_id: str | None,\n embedder: IndexingEmbedder,\n document_indices: list[DocumentIndex],\n db_session: Session,\n tenant_id: str,\n adapter: IndexingBatchAdapter,\n chunker: Chunker | None = None,\n ignore_time_skip: bool = False,\n) -> IndexingPipelineResult:\n \"\"\"Builds a pipeline which takes in a list (batch) of docs and indexes them.\"\"\"\n all_search_settings = get_active_search_settings(db_session)\n if (\n all_search_settings.secondary\n and all_search_settings.secondary.status == IndexModelStatus.FUTURE\n ):\n search_settings = all_search_settings.secondary\n else:\n search_settings = all_search_settings.primary\n\n multipass_config = get_multipass_config(search_settings)\n\n enable_contextual_rag = (\n search_settings.enable_contextual_rag or ENABLE_CONTEXTUAL_RAG\n )\n llm = None\n if enable_contextual_rag:\n llm = get_llm_for_contextual_rag(\n search_settings.contextual_rag_llm_name or DEFAULT_CONTEXTUAL_RAG_LLM_NAME,\n search_settings.contextual_rag_llm_provider\n or DEFAULT_CONTEXTUAL_RAG_LLM_PROVIDER,\n )\n\n chunker = chunker or Chunker(\n tokenizer=embedder.embedding_model.tokenizer,\n enable_multipass=multipass_config.multipass_indexing,\n enable_large_chunks=multipass_config.enable_large_chunks,\n enable_contextual_rag=enable_contextual_rag,\n # after every doc, update status in case there are a bunch of really long docs\n )\n\n return index_doc_batch_with_handler(\n chunker=chunker,\n embedder=embedder,\n document_indices=document_indices,\n document_batch=document_batch,\n request_id=request_id,\n tenant_id=tenant_id,\n db_session=db_session,\n adapter=adapter,\n enable_contextual_rag=enable_contextual_rag,\n llm=llm,\n ignore_time_skip=ignore_time_skip,\n )\n" }, { "path": "backend/onyx/indexing/models.py", "content": "import contextlib\nfrom collections.abc import Generator\nfrom typing import Optional\nfrom typing import Protocol\nfrom typing import TYPE_CHECKING\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.connectors.models import Document\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.db.enums import SwitchoverType\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.pydantic_util import shallow_model_dump\nfrom shared_configs.enums import EmbeddingProvider\nfrom shared_configs.model_server_models import Embedding\n\nif TYPE_CHECKING:\n from onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext\nfrom sqlalchemy.engine.util import TransactionalContext\n\nif TYPE_CHECKING:\n from onyx.db.models import SearchSettings\n\n\nlogger = setup_logger()\n\n\nclass ChunkEmbedding(BaseModel):\n full_embedding: Embedding\n mini_chunk_embeddings: list[Embedding]\n\n\nclass BaseChunk(BaseModel):\n chunk_id: int\n # The first sentence(s) of the first Section of the chunk\n blurb: str\n content: str\n # Holds the link and the offsets into the raw Chunk text\n source_links: dict[int, str] | None\n image_file_id: str | None\n # True if this Chunk's start is not at the start of a Section\n # TODO(andrei): This is deprecated as of the OpenSearch migration. Remove.\n # Do not use.\n section_continuation: bool\n\n\nclass DocAwareChunk(BaseChunk):\n # During indexing flow, we have access to a complete \"Document\"\n # During inference we only have access to the document id and do not reconstruct the Document\n source_document: Document\n\n # This could be an empty string if the title is too long and taking up too much of the chunk\n # This does not mean necessarily that the document does not have a title\n title_prefix: str\n\n # During indexing we also (optionally) build a metadata string from the metadata dict\n # This is also indexed so that we can strip it out after indexing, this way it supports\n # multiple iterations of metadata representation for backwards compatibility\n metadata_suffix_semantic: str\n metadata_suffix_keyword: str\n\n # This is the number of tokens reserved for contextual RAG\n # in the chunk. doc_summary and chunk_context conbined should\n # contain at most this many tokens.\n contextual_rag_reserved_tokens: int\n # This is the summary for the document generated for contextual RAG\n doc_summary: str\n # This is the context for this chunk generated for contextual RAG\n chunk_context: str\n\n mini_chunk_texts: list[str] | None\n\n large_chunk_id: int | None\n\n large_chunk_reference_ids: list[int] = Field(default_factory=list)\n\n def to_short_descriptor(self) -> str:\n \"\"\"Used when logging the identity of a chunk\"\"\"\n return f\"{self.source_document.to_short_descriptor()} Chunk ID: {self.chunk_id}\"\n\n def get_link(self) -> str | None:\n return (\n self.source_document.sections[0].link\n if self.source_document.sections\n else None\n )\n\n\nclass IndexChunk(DocAwareChunk):\n embeddings: ChunkEmbedding\n title_embedding: Embedding | None\n\n\n# TODO(rkuo): currently, this extra metadata sent during indexing is just for speed,\n# but full consistency happens on background sync\nclass DocMetadataAwareIndexChunk(IndexChunk):\n \"\"\"An `IndexChunk` that contains all necessary metadata to be indexed. This includes\n the following:\n\n access: holds all information about which users should have access to the\n source document for this chunk.\n document_sets: all document sets the source document for this chunk is a part\n of. This is used for filtering / personas.\n boost: influences the ranking of this chunk at query time. Positive -> ranked higher,\n negative -> ranked lower. Not included in aggregated boost calculation\n for legacy reasons.\n aggregated_chunk_boost_factor: represents the aggregated chunk-level boost (currently: information content)\n \"\"\"\n\n tenant_id: str\n access: \"DocumentAccess\"\n document_sets: set[str]\n user_project: list[int]\n personas: list[int]\n boost: int\n aggregated_chunk_boost_factor: float\n # Full ancestor path from root hierarchy node to document's parent.\n # Stored as an integer array in OpenSearch for hierarchy-based filtering.\n # Empty list means no hierarchy info (document excluded from hierarchy searches).\n ancestor_hierarchy_node_ids: list[int]\n\n @classmethod\n def from_index_chunk(\n cls,\n index_chunk: IndexChunk,\n access: \"DocumentAccess\",\n document_sets: set[str],\n user_project: list[int],\n personas: list[int],\n boost: int,\n aggregated_chunk_boost_factor: float,\n tenant_id: str,\n ancestor_hierarchy_node_ids: list[int] | None = None,\n ) -> \"DocMetadataAwareIndexChunk\":\n return cls.model_construct(\n **shallow_model_dump(index_chunk),\n access=access,\n document_sets=document_sets,\n user_project=user_project,\n personas=personas,\n boost=boost,\n aggregated_chunk_boost_factor=aggregated_chunk_boost_factor,\n tenant_id=tenant_id,\n ancestor_hierarchy_node_ids=ancestor_hierarchy_node_ids or [],\n )\n\n\nclass EmbeddingModelDetail(BaseModel):\n id: int | None = None\n model_name: str\n normalize: bool\n query_prefix: str | None\n passage_prefix: str | None\n api_url: str | None = None\n provider_type: EmbeddingProvider | None = None\n api_key: str | None = None\n\n # This disables the \"model_\" protected namespace for pydantic\n model_config = {\"protected_namespaces\": ()}\n\n @classmethod\n def from_db_model(\n cls,\n search_settings: \"SearchSettings\",\n ) -> \"EmbeddingModelDetail\":\n api_key = None\n if (\n search_settings.cloud_provider is not None\n and search_settings.cloud_provider.api_key is not None\n ):\n api_key = search_settings.cloud_provider.api_key.get_value(apply_mask=True)\n\n return cls(\n id=search_settings.id,\n model_name=search_settings.model_name,\n normalize=search_settings.normalize,\n query_prefix=search_settings.query_prefix,\n passage_prefix=search_settings.passage_prefix,\n provider_type=search_settings.provider_type,\n api_key=api_key,\n api_url=search_settings.api_url,\n )\n\n\n# Additional info needed for indexing time\nclass IndexingSetting(EmbeddingModelDetail):\n model_dim: int\n index_name: str | None\n multipass_indexing: bool\n embedding_precision: EmbeddingPrecision\n reduced_dimension: int | None = None\n\n switchover_type: SwitchoverType = SwitchoverType.REINDEX\n enable_contextual_rag: bool\n contextual_rag_llm_name: str | None = None\n contextual_rag_llm_provider: str | None = None\n\n # This disables the \"model_\" protected namespace for pydantic\n model_config = {\"protected_namespaces\": ()}\n\n @property\n def final_embedding_dim(self) -> int:\n if self.reduced_dimension:\n return self.reduced_dimension\n return self.model_dim\n\n @classmethod\n def from_db_model(cls, search_settings: \"SearchSettings\") -> \"IndexingSetting\":\n return cls(\n model_name=search_settings.model_name,\n model_dim=search_settings.model_dim,\n normalize=search_settings.normalize,\n query_prefix=search_settings.query_prefix,\n passage_prefix=search_settings.passage_prefix,\n provider_type=search_settings.provider_type,\n index_name=search_settings.index_name,\n multipass_indexing=search_settings.multipass_indexing,\n embedding_precision=search_settings.embedding_precision,\n reduced_dimension=search_settings.reduced_dimension,\n switchover_type=search_settings.switchover_type,\n enable_contextual_rag=search_settings.enable_contextual_rag,\n )\n\n\nclass MultipassConfig(BaseModel):\n multipass_indexing: bool\n enable_large_chunks: bool\n\n\nclass UpdatableChunkData(BaseModel):\n chunk_id: int\n document_id: str\n boost_score: float\n\n\nclass ChunkEnrichmentContext(Protocol):\n \"\"\"Returned by prepare_enrichment. Holds pre-computed metadata lookups\n and provides per-chunk enrichment.\"\"\"\n\n doc_id_to_previous_chunk_cnt: dict[str, int]\n doc_id_to_new_chunk_cnt: dict[str, int]\n\n def enrich_chunk(\n self, chunk: IndexChunk, score: float\n ) -> DocMetadataAwareIndexChunk: ...\n\n\nclass IndexingBatchAdapter(Protocol):\n def prepare(\n self, documents: list[Document], ignore_time_skip: bool\n ) -> Optional[\"DocumentBatchPrepareContext\"]: ...\n\n @contextlib.contextmanager\n def lock_context(\n self, documents: list[Document]\n ) -> Generator[TransactionalContext, None, None]:\n \"\"\"Provide a transaction/row-lock context for critical updates.\"\"\"\n\n def prepare_enrichment(\n self,\n context: \"DocumentBatchPrepareContext\",\n tenant_id: str,\n chunks: list[DocAwareChunk],\n ) -> ChunkEnrichmentContext:\n \"\"\"Prepare per-chunk enrichment data (access, document sets, boost, etc.).\n\n Precondition: ``chunks`` have already been through the embedding step\n (i.e. they are ``IndexChunk`` instances with populated embeddings,\n passed here as the base ``DocAwareChunk`` type).\n \"\"\"\n ...\n\n def post_index(\n self,\n context: \"DocumentBatchPrepareContext\",\n updatable_chunk_data: list[UpdatableChunkData],\n filtered_documents: list[Document],\n enrichment: ChunkEnrichmentContext,\n ) -> None: ...\n" }, { "path": "backend/onyx/indexing/vector_db_insertion.py", "content": "import time\nfrom collections.abc import Callable\nfrom collections.abc import Iterable\nfrom http import HTTPStatus\nfrom itertools import chain\nfrom itertools import groupby\n\nimport httpx\n\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import DocumentInsertionRecord\nfrom onyx.document_index.interfaces import IndexBatchParams\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _log_insufficient_storage_error(e: Exception) -> None:\n if isinstance(e, httpx.HTTPStatusError):\n if e.response.status_code == HTTPStatus.INSUFFICIENT_STORAGE:\n logger.error(\n \"NOTE: HTTP Status 507 Insufficient Storage indicates \"\n \"you need to allocate more memory or disk space to the \"\n \"Vespa/index container.\"\n )\n\n\ndef write_chunks_to_vector_db_with_backoff(\n document_index: DocumentIndex,\n make_chunks: Callable[[], Iterable[DocMetadataAwareIndexChunk]],\n index_batch_params: IndexBatchParams,\n) -> tuple[list[DocumentInsertionRecord], list[ConnectorFailure]]:\n \"\"\"Tries to insert all chunks in one large batch. If that batch fails for any reason,\n goes document by document to isolate the failure(s).\n\n IMPORTANT: must pass in whole documents at a time not individual chunks, since the\n vector DB interface assumes that all chunks for a single document are present. The\n chunks must also be in contiguous batches\n \"\"\"\n # first try to write the chunks to the vector db\n try:\n return (\n list(\n document_index.index(\n chunks=make_chunks(),\n index_batch_params=index_batch_params,\n )\n ),\n [],\n )\n except Exception as e:\n logger.exception(\n \"Failed to write chunk batch to vector db. Trying individual docs.\"\n )\n\n # give some specific logging on this common failure case.\n _log_insufficient_storage_error(e)\n\n # wait a couple seconds just to give the vector db a chance to recover\n time.sleep(2)\n\n insertion_records: list[DocumentInsertionRecord] = []\n failures: list[ConnectorFailure] = []\n\n def key(chunk: DocMetadataAwareIndexChunk) -> str:\n return chunk.source_document.id\n\n seen_doc_ids: set[str] = set()\n for doc_id, chunks_for_doc in groupby(make_chunks(), key=key):\n if doc_id in seen_doc_ids:\n raise RuntimeError(\n f\"Doc chunks are not arriving in order. Current doc_id={doc_id}, seen_doc_ids={list(seen_doc_ids)}\"\n )\n seen_doc_ids.add(doc_id)\n\n first_chunk = next(chunks_for_doc)\n chunks_for_doc = chain([first_chunk], chunks_for_doc)\n\n try:\n insertion_records.extend(\n document_index.index(\n chunks=chunks_for_doc,\n index_batch_params=index_batch_params,\n )\n )\n except Exception as e:\n logger.exception(\n f\"Failed to write document chunks for '{doc_id}' to vector db\"\n )\n\n # give some specific logging on this common failure case.\n _log_insufficient_storage_error(e)\n\n failures.append(\n ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=doc_id,\n document_link=first_chunk.get_link(),\n ),\n failure_message=str(e),\n exception=e,\n )\n )\n\n return insertion_records, failures\n" }, { "path": "backend/onyx/key_value_store/__init__.py", "content": "" }, { "path": "backend/onyx/key_value_store/factory.py", "content": "from onyx.key_value_store.interface import KeyValueStore\nfrom onyx.key_value_store.store import PgRedisKVStore\nfrom shared_configs.configs import DEFAULT_REDIS_PREFIX\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n\ndef get_kv_store() -> KeyValueStore:\n # In the Multi Tenant case, the tenant context is picked up automatically, it does not need to be passed in\n # It's read from the global thread level variable\n return PgRedisKVStore()\n\n\ndef get_shared_kv_store() -> KeyValueStore:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(DEFAULT_REDIS_PREFIX)\n try:\n return get_kv_store()\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n" }, { "path": "backend/onyx/key_value_store/interface.py", "content": "import abc\nfrom typing import cast\n\nfrom onyx.utils.special_types import JSON_ro\n\n\nclass KvKeyNotFoundError(Exception):\n pass\n\n\ndef unwrap_str(val: JSON_ro) -> str:\n \"\"\"Unwrap a string stored as {\"value\": str} in the encrypted KV store.\n Also handles legacy plain-string values cached in Redis.\"\"\"\n if isinstance(val, dict):\n try:\n return cast(str, val[\"value\"])\n except KeyError:\n raise ValueError(\n f\"Expected dict with 'value' key, got keys: {list(val.keys())}\"\n )\n return cast(str, val)\n\n\nclass KeyValueStore:\n # In the Multi Tenant case, the tenant context is picked up automatically, it does not need to be passed in\n # It's read from the global thread level variable\n @abc.abstractmethod\n def store(self, key: str, val: JSON_ro, encrypt: bool = False) -> None:\n raise NotImplementedError\n\n @abc.abstractmethod\n def load(self, key: str, refresh_cache: bool = False) -> JSON_ro:\n raise NotImplementedError\n\n @abc.abstractmethod\n def delete(self, key: str) -> None:\n raise NotImplementedError\n" }, { "path": "backend/onyx/key_value_store/store.py", "content": "import json\nfrom typing import cast\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import KVStore\nfrom onyx.key_value_store.interface import KeyValueStore\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.special_types import JSON_ro\n\n\nlogger = setup_logger()\n\n\nREDIS_KEY_PREFIX = \"onyx_kv_store:\"\nKV_REDIS_KEY_EXPIRATION = 60 * 60 * 24 # 1 Day\n\n\nclass PgRedisKVStore(KeyValueStore):\n def __init__(self, cache: CacheBackend | None = None) -> None:\n self._cache = cache\n\n def _get_cache(self) -> CacheBackend:\n if self._cache is None:\n from onyx.cache.factory import get_cache_backend\n\n self._cache = get_cache_backend()\n return self._cache\n\n def store(self, key: str, val: JSON_ro, encrypt: bool = False) -> None:\n # Not encrypted in Cache backend (typically Redis), but encrypted in Postgres\n try:\n self._get_cache().set(\n REDIS_KEY_PREFIX + key, json.dumps(val), ex=KV_REDIS_KEY_EXPIRATION\n )\n except Exception as e:\n # Fallback gracefully to Postgres if Cache backend fails\n logger.error(\n f\"Failed to set value in Cache backend for key '{key}': {str(e)}\"\n )\n\n encrypted_val = val if encrypt else None\n plain_val = val if not encrypt else None\n with get_session_with_current_tenant() as db_session:\n obj = db_session.query(KVStore).filter_by(key=key).first()\n if obj:\n obj.value = plain_val\n obj.encrypted_value = encrypted_val # type: ignore[assignment]\n else:\n obj = KVStore(key=key, value=plain_val, encrypted_value=encrypted_val)\n db_session.query(KVStore).filter_by(key=key).delete() # just in case\n db_session.add(obj)\n db_session.commit()\n\n def load(self, key: str, refresh_cache: bool = False) -> JSON_ro:\n if not refresh_cache:\n try:\n cached = self._get_cache().get(REDIS_KEY_PREFIX + key)\n if cached is not None:\n return json.loads(cached.decode(\"utf-8\"))\n except Exception as e:\n logger.error(\n f\"Failed to get value from cache for key '{key}': {str(e)}\"\n )\n\n with get_session_with_current_tenant() as db_session:\n obj = db_session.query(KVStore).filter_by(key=key).first()\n if not obj:\n raise KvKeyNotFoundError\n\n if obj.value is not None:\n value = obj.value\n elif obj.encrypted_value is not None:\n # Unwrap SensitiveValue - this is internal backend use\n value = obj.encrypted_value.get_value(apply_mask=False)\n else:\n value = None\n\n try:\n self._get_cache().set(\n REDIS_KEY_PREFIX + key,\n json.dumps(value),\n ex=KV_REDIS_KEY_EXPIRATION,\n )\n except Exception as e:\n logger.error(f\"Failed to set value in cache for key '{key}': {str(e)}\")\n\n return cast(JSON_ro, value)\n\n def delete(self, key: str) -> None:\n try:\n self._get_cache().delete(REDIS_KEY_PREFIX + key)\n except Exception as e:\n logger.error(f\"Failed to delete value from cache for key '{key}': {str(e)}\")\n\n with get_session_with_current_tenant() as db_session:\n result = db_session.query(KVStore).filter_by(key=key).delete()\n if result == 0:\n raise KvKeyNotFoundError\n db_session.commit()\n" }, { "path": "backend/onyx/kg/clustering/clustering.py", "content": "import time\nfrom collections.abc import Generator\nfrom typing import cast\n\nfrom rapidfuzz.fuzz import ratio\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy import func\nfrom sqlalchemy import text\n\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.kg_configs import KG_CLUSTERING_RETRIEVE_THRESHOLD\nfrom onyx.configs.kg_configs import KG_CLUSTERING_THRESHOLD\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.entities import KGEntity\nfrom onyx.db.entities import KGEntityExtractionStaging\nfrom onyx.db.entities import merge_entities\nfrom onyx.db.entities import transfer_entity\nfrom onyx.db.kg_config import get_kg_config_settings\nfrom onyx.db.kg_config import validate_kg_settings\nfrom onyx.db.models import Document\nfrom onyx.db.models import KGEntityType\nfrom onyx.db.models import KGRelationshipExtractionStaging\nfrom onyx.db.models import KGRelationshipTypeExtractionStaging\nfrom onyx.db.relationships import transfer_relationship\nfrom onyx.db.relationships import transfer_relationship_type\nfrom onyx.db.relationships import upsert_relationship\nfrom onyx.db.relationships import upsert_relationship_type\nfrom onyx.document_index.vespa.kg_interactions import (\n get_kg_vespa_info_update_requests_for_document,\n)\nfrom onyx.document_index.vespa.kg_interactions import update_kg_chunks_vespa_info\nfrom onyx.kg.models import KGGroundingType\nfrom onyx.kg.utils.formatting_utils import make_relationship_id\nfrom onyx.kg.utils.lock_utils import extend_lock\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\nlogger = setup_logger()\n\n\ndef _get_batch_untransferred_grounded_entities(\n batch_size: int,\n) -> Generator[list[KGEntityExtractionStaging], None, None]:\n while True:\n with get_session_with_current_tenant() as db_session:\n batch = (\n db_session.query(KGEntityExtractionStaging)\n .join(\n KGEntityType,\n KGEntityExtractionStaging.entity_type_id_name\n == KGEntityType.id_name,\n )\n .filter(\n KGEntityType.grounding == KGGroundingType.GROUNDED,\n KGEntityExtractionStaging.transferred_id_name.is_(None),\n )\n .limit(batch_size)\n .all()\n )\n if not batch:\n break\n yield batch\n\n\ndef _get_batch_untransferred_relationship_types(\n batch_size: int,\n) -> Generator[list[KGRelationshipTypeExtractionStaging], None, None]:\n while True:\n with get_session_with_current_tenant() as db_session:\n batch = (\n db_session.query(KGRelationshipTypeExtractionStaging)\n .filter(KGRelationshipTypeExtractionStaging.transferred.is_(False))\n .limit(batch_size)\n .all()\n )\n if not batch:\n break\n yield batch\n\n\ndef _get_batch_untransferred_relationships(\n batch_size: int,\n) -> Generator[list[KGRelationshipExtractionStaging], None, None]:\n while True:\n with get_session_with_current_tenant() as db_session:\n batch = (\n db_session.query(KGRelationshipExtractionStaging)\n .filter(KGRelationshipExtractionStaging.transferred.is_(False))\n .limit(batch_size)\n .all()\n )\n if not batch:\n break\n yield batch\n\n\ndef _get_batch_entities_with_parent(\n batch_size: int,\n) -> Generator[list[KGEntityExtractionStaging], None, None]:\n offset = 0\n\n while True:\n with get_session_with_current_tenant() as db_session:\n batch = (\n db_session.query(KGEntityExtractionStaging)\n .filter(KGEntityExtractionStaging.parent_key.isnot(None))\n .order_by(KGEntityExtractionStaging.id_name)\n .offset(offset)\n .limit(batch_size)\n .all()\n )\n if not batch:\n break\n # we can't filter out \"\"s earlier as it will mess up the pagination\n yield [entity for entity in batch if entity.parent_key != \"\"]\n offset += batch_size\n\n\ndef _get_batch_kg_processed_documents(\n batch_size: int,\n) -> Generator[list[Document], None, None]:\n offset = 0\n\n while True:\n with get_session_with_current_tenant() as db_session:\n batch = (\n db_session.query(Document)\n .join(\n KGEntityExtractionStaging,\n Document.id == KGEntityExtractionStaging.document_id,\n )\n .filter(\n KGEntityExtractionStaging.transferred_id_name.is_not(None),\n )\n .order_by(Document.id)\n .offset(offset)\n .limit(batch_size)\n .all()\n )\n if not batch:\n break\n yield batch\n offset += batch_size\n\n\ndef _cluster_one_grounded_entity(\n entity: KGEntityExtractionStaging,\n) -> tuple[KGEntity, bool]:\n \"\"\"\n Cluster a single grounded entity.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n # get entity name and filtering conditions\n if entity.document_id is not None:\n entity_name = cast(\n str,\n db_session.query(Document.semantic_id)\n .filter(Document.id == entity.document_id)\n .scalar(),\n ).lower()\n filtering = [KGEntity.document_id.is_(None)]\n else:\n entity_name = entity.name.lower()\n filtering = []\n\n # skip those with numbers so we don't cluster version1 and version2, etc.\n similar_entities: list[KGEntity] = []\n if not any(char.isdigit() for char in entity_name):\n # find similar entities, uses GIN index, very efficient\n db_session.execute(\n text(\n \"SET pg_trgm.similarity_threshold = \"\n + str(KG_CLUSTERING_RETRIEVE_THRESHOLD)\n )\n )\n similar_entities = (\n db_session.query(KGEntity)\n .filter(\n # find entities of the same type with a similar name\n *filtering,\n KGEntity.entity_type_id_name == entity.entity_type_id_name,\n getattr(func, POSTGRES_DEFAULT_SCHEMA).similarity_op(\n KGEntity.name, entity_name\n ),\n )\n .all()\n )\n\n # find best match\n best_score = -1.0\n best_entity = None\n for similar in similar_entities:\n # skip those with numbers so we don't cluster version1 and version2, etc.\n if any(char.isdigit() for char in similar.name):\n continue\n score = ratio(similar.name, entity_name)\n if score >= KG_CLUSTERING_THRESHOLD * 100 and score > best_score:\n best_score = score\n best_entity = similar\n\n # if there is a match, update the entity, otherwise create a new one\n with get_session_with_current_tenant() as db_session:\n if best_entity:\n logger.debug(f\"Merged {entity.name} with {best_entity.name}\")\n update_vespa = (\n best_entity.document_id is None and entity.document_id is not None\n )\n transferred_entity = merge_entities(\n db_session=db_session, parent=best_entity, child=entity\n )\n else:\n update_vespa = entity.document_id is not None\n transferred_entity = transfer_entity(db_session=db_session, entity=entity)\n\n db_session.commit()\n\n return transferred_entity, update_vespa\n\n\ndef _create_one_parent_child_relationship(entity: KGEntityExtractionStaging) -> None:\n \"\"\"\n Creates a relationship between the entity and its parent, if it exists.\n Then, updates the entity's parent to the next ancestor.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n # find the next ancestor\n parent = (\n db_session.query(KGEntity)\n .filter(KGEntity.entity_key == entity.parent_key)\n .first()\n )\n\n if parent is not None:\n # create parent child relationship and relationship type\n upsert_relationship_type(\n db_session=db_session,\n source_entity_type=parent.entity_type_id_name,\n relationship_type=\"has_subcomponent\",\n target_entity_type=entity.entity_type_id_name,\n )\n relationship_id_name = make_relationship_id(\n parent.id_name,\n \"has_subcomponent\",\n cast(str, entity.transferred_id_name),\n )\n upsert_relationship(\n db_session=db_session,\n relationship_id_name=relationship_id_name,\n source_document_id=entity.document_id,\n )\n\n next_ancestor = parent.parent_key or \"\"\n else:\n next_ancestor = \"\"\n\n # set the staging entity's parent to the next ancestor\n # if there is no parent or next ancestor, set to \"\" to differentiate from None\n # None will mess up the pagination in _get_batch_entities_with_parent\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.id_name == entity.id_name\n ).update({\"parent_key\": next_ancestor})\n db_session.commit()\n\n\ndef _transfer_one_relationship(\n relationship: KGRelationshipExtractionStaging,\n) -> None:\n with get_session_with_current_tenant() as db_session:\n # get the translations\n staging_entity_id_names = {\n relationship.source_node,\n relationship.target_node,\n }\n entity_translations: dict[str, str] = {\n entity.id_name: entity.transferred_id_name\n for entity in db_session.query(KGEntityExtractionStaging)\n .filter(KGEntityExtractionStaging.id_name.in_(staging_entity_id_names))\n .all()\n if entity.transferred_id_name is not None\n }\n if len(entity_translations) != len(staging_entity_id_names):\n logger.error(\n f\"Missing entity translations for {staging_entity_id_names - entity_translations.keys()}\"\n )\n return\n\n # transfer the relationship\n transfer_relationship(\n db_session=db_session,\n relationship=relationship,\n entity_translations=entity_translations,\n )\n db_session.commit()\n\n\ndef kg_clustering(\n tenant_id: str,\n index_name: str,\n lock: RedisLock,\n processing_chunk_batch_size: int = 16,\n) -> None:\n \"\"\"\n Here we will cluster the extractions based on their cluster frameworks.\n Initially, this will only focus on grounded entities with pre-determined\n relationships, so 'clustering' is actually not yet required.\n However, we may need to reconcile entities coming from different sources.\n\n The primary purpose of this function is to populate the actual KG tables\n from the temp_extraction tables.\n\n This will change with deep extraction, where grounded-sourceless entities\n can be extracted and then need to be clustered.\n \"\"\"\n logger.info(f\"Starting kg clustering for tenant {tenant_id}\")\n\n kg_config_settings = get_kg_config_settings()\n validate_kg_settings(kg_config_settings)\n\n last_lock_time = time.monotonic()\n\n # Cluster and transfer grounded entities sequentially\n start_time = time.monotonic()\n i_batch = 0\n for i_batch, untransferred_grounded_entities in enumerate(\n _get_batch_untransferred_grounded_entities(\n batch_size=processing_chunk_batch_size\n )\n ):\n for entity in untransferred_grounded_entities:\n _cluster_one_grounded_entity(entity)\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n # logger.debug(f\"Transferred entities batch {i}\")\n # NOTE: we assume every entity is transferred, as we currently only have grounded entities\n time_delta = time.monotonic() - start_time\n logger.info(\n f\"Finished transferring {i_batch + 1} entity batches in {time_delta:.2f}s\"\n )\n\n # Create parent-child relationships in parallel\n for _ in range(kg_config_settings.KG_MAX_PARENT_RECURSION_DEPTH):\n for root_entities in _get_batch_entities_with_parent(\n batch_size=processing_chunk_batch_size\n ):\n run_functions_tuples_in_parallel(\n [\n (_create_one_parent_child_relationship, (root_entity,))\n for root_entity in root_entities\n ]\n )\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n logger.info(\"Finished creating all parent-child relationships\")\n\n # Transfer the relationship types (no need to do in parallel as there's only a few)\n start_time = time.monotonic()\n i_batch = 0\n for i_batch, relationship_types in enumerate(\n _get_batch_untransferred_relationship_types(\n batch_size=processing_chunk_batch_size\n )\n ):\n with get_session_with_current_tenant() as db_session:\n for relationship_type in relationship_types:\n transfer_relationship_type(db_session, relationship_type)\n db_session.commit()\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n # logger.debug(f\"Transferred relationship types batch {i}\")\n time_delta = time.monotonic() - start_time\n logger.info(\n f\"Finished transferring {i_batch + 1} relationship type batches in {time_delta:.2f}s\"\n )\n\n # Transfer the relationships in parallel\n start_time = time.monotonic()\n i_batch = 0\n for i_batch, relationships in enumerate(\n _get_batch_untransferred_relationships(batch_size=processing_chunk_batch_size)\n ):\n run_functions_tuples_in_parallel(\n [\n (_transfer_one_relationship, (relationship,))\n for relationship in relationships\n ]\n )\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n # logger.debug(f\"Transferred relationships batch {i}\")\n time_delta = time.monotonic() - start_time\n logger.info(\n f\"Finished transferring {i_batch + 1} relationship batches in {time_delta:.2f}s\"\n )\n\n # Update vespa for each document\n start_time = time.monotonic()\n i_batch = 0\n for i_batch, documents in enumerate(\n _get_batch_kg_processed_documents(batch_size=processing_chunk_batch_size)\n ):\n batch_update_requests = run_functions_tuples_in_parallel(\n [\n (get_kg_vespa_info_update_requests_for_document, (document.id,))\n for document in documents\n ]\n )\n for update_requests, document in zip(batch_update_requests, documents):\n try:\n update_kg_chunks_vespa_info(update_requests, index_name, tenant_id)\n except Exception as e:\n logger.error(f\"Error updating vespa for document {document.id}: {e}\")\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n # logger.debug(f\"Updated vespa for documents batch {i}\")\n time_delta = time.monotonic() - start_time\n logger.info(\n f\"Finished updating {i_batch + 1} document batches in {time_delta:.2f}s\"\n )\n\n # Delete the transferred objects from the staging tables\n try:\n with get_session_with_current_tenant() as db_session:\n db_session.query(KGRelationshipExtractionStaging).filter(\n KGRelationshipExtractionStaging.transferred.is_(True)\n ).delete(synchronize_session=False)\n db_session.commit()\n except Exception as e:\n logger.error(f\"Error deleting relationships: {e}\")\n\n try:\n with get_session_with_current_tenant() as db_session:\n db_session.query(KGRelationshipTypeExtractionStaging).filter(\n KGRelationshipTypeExtractionStaging.transferred.is_(True)\n ).delete(synchronize_session=False)\n db_session.commit()\n except Exception as e:\n logger.error(f\"Error deleting relationship types: {e}\")\n\n try:\n with get_session_with_current_tenant() as db_session:\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.transferred_id_name.is_not(None)\n ).delete(synchronize_session=False)\n db_session.commit()\n except Exception as e:\n logger.error(f\"Error deleting entities: {e}\")\n logger.info(\"Finished deleting all transferred staging entries\")\n" }, { "path": "backend/onyx/kg/clustering/normalizations.py", "content": "import re\nfrom collections import defaultdict\nfrom typing import cast\n\nimport numpy as np\nfrom rapidfuzz.distance.DamerauLevenshtein import normalized_similarity\nfrom sqlalchemy import desc\nfrom sqlalchemy import Float\nfrom sqlalchemy import func\nfrom sqlalchemy import MetaData\nfrom sqlalchemy import select\nfrom sqlalchemy import String\nfrom sqlalchemy import Table\nfrom sqlalchemy.dialects.postgresql import ARRAY\n\nfrom onyx.configs.kg_configs import KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT\nfrom onyx.configs.kg_configs import KG_NORMALIZATION_RERANK_NGRAM_WEIGHTS\nfrom onyx.configs.kg_configs import KG_NORMALIZATION_RERANK_THRESHOLD\nfrom onyx.configs.kg_configs import KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import KGEntity\nfrom onyx.db.relationships import get_relationships_for_entity_type_pairs\nfrom onyx.kg.models import NormalizedEntities\nfrom onyx.kg.models import NormalizedRelationships\nfrom onyx.kg.utils.embeddings import encode_string_batch\nfrom onyx.kg.utils.formatting_utils import format_entity_id_for_models\nfrom onyx.kg.utils.formatting_utils import get_attributes\nfrom onyx.kg.utils.formatting_utils import get_entity_type\nfrom onyx.kg.utils.formatting_utils import make_entity_w_attributes\nfrom onyx.kg.utils.formatting_utils import make_relationship_id\nfrom onyx.kg.utils.formatting_utils import split_entity_id\nfrom onyx.kg.utils.formatting_utils import split_relationship_id\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\nlogger = setup_logger()\n\n\nalphanum_regex = re.compile(r\"[^a-z0-9]+\")\nrem_email_regex = re.compile(r\"(?<=\\S)@([a-z0-9-]+)\\.([a-z]{2,6})$\")\n\n\ndef _ngrams(sequence: str, n: int) -> list[tuple[str, ...]]:\n \"\"\"Generate n-grams from a sequence.\"\"\"\n return [tuple(sequence[i : i + n]) for i in range(len(sequence) - n + 1)]\n\n\ndef _clean_name(entity_name: str) -> str:\n \"\"\"\n Clean an entity string by removing non-alphanumeric characters and email addresses.\n If the name after cleaning is empty, return the original name in lowercase.\n \"\"\"\n cleaned_entity = entity_name.casefold()\n return (\n alphanum_regex.sub(\"\", rem_email_regex.sub(\"\", cleaned_entity))\n or cleaned_entity\n )\n\n\ndef _normalize_one_entity(\n entity: str,\n attributes: dict[str, str],\n allowed_docs_temp_view_name: str | None = None,\n) -> str | None:\n \"\"\"\n Matches a single entity to the best matching entity of the same type.\n \"\"\"\n entity_type, entity_name = split_entity_id(entity)\n if entity_name == \"*\":\n return entity\n\n cleaned_entity = _clean_name(entity_name)\n\n # narrow filter to subtype if requested\n type_filters = [KGEntity.entity_type_id_name == entity_type]\n if \"subtype\" in attributes:\n type_filters.append(\n KGEntity.attributes.op(\"@>\")({\"subtype\": attributes[\"subtype\"]})\n )\n\n # step 1: find entities containing the entity_name or something similar\n with get_session_with_current_tenant() as db_session:\n # get allowed documents\n metadata = MetaData()\n if allowed_docs_temp_view_name is None:\n raise ValueError(\"allowed_docs_temp_view_name is not available\")\n\n effective_schema_allowed_docs_temp_view_name = (\n allowed_docs_temp_view_name.split(\".\")[-1]\n )\n\n allowed_docs_temp_view = Table(\n effective_schema_allowed_docs_temp_view_name,\n metadata,\n autoload_with=db_session.get_bind(),\n )\n\n # generate trigrams of the queried entity Q\n query_trigrams = db_session.query(\n getattr(func, POSTGRES_DEFAULT_SCHEMA)\n .show_trgm(cleaned_entity)\n .cast(ARRAY(String(3)))\n .label(\"trigrams\")\n ).cte(\"query\")\n\n candidates = cast(\n list[tuple[str, str, float]],\n db_session.query(\n KGEntity.id_name,\n KGEntity.name,\n (\n # for each entity E, compute score = | Q ∩ E | / min(|Q|, |E|)\n func.cardinality(\n func.array(\n select(func.unnest(KGEntity.name_trigrams))\n .correlate(KGEntity)\n .intersect(\n select(\n func.unnest(query_trigrams.c.trigrams)\n ).correlate(query_trigrams)\n )\n .scalar_subquery()\n )\n ).cast(Float)\n / func.least(\n func.cardinality(query_trigrams.c.trigrams),\n func.cardinality(KGEntity.name_trigrams),\n )\n ).label(\"score\"),\n )\n .select_from(KGEntity, query_trigrams)\n .outerjoin(\n allowed_docs_temp_view,\n KGEntity.document_id == allowed_docs_temp_view.c.allowed_doc_id,\n )\n .filter(\n *type_filters,\n KGEntity.name_trigrams.overlap(query_trigrams.c.trigrams),\n # Add filter for allowed docs - either document_id is NULL or it's in allowed_docs\n (\n KGEntity.document_id.is_(None)\n | allowed_docs_temp_view.c.allowed_doc_id.isnot(None)\n ),\n )\n .order_by(desc(\"score\"))\n .limit(KG_NORMALIZATION_RETRIEVE_ENTITIES_LIMIT)\n .all(),\n )\n if not candidates:\n return None\n\n # step 2: do a weighted ngram analysis and damerau levenshtein distance to rerank\n n1, n2, n3 = (\n set(_ngrams(cleaned_entity, 1)),\n set(_ngrams(cleaned_entity, 2)),\n set(_ngrams(cleaned_entity, 3)),\n )\n for i, (candidate_id_name, candidate_name, _) in enumerate(candidates):\n cleaned_candidate = _clean_name(candidate_name)\n h_n1, h_n2, h_n3 = (\n set(_ngrams(cleaned_candidate, 1)),\n set(_ngrams(cleaned_candidate, 2)),\n set(_ngrams(cleaned_candidate, 3)),\n )\n\n # compute ngram overlap, renormalize scores if the names are too short for larger ngrams\n grams_used = min(2, len(cleaned_entity) - 1, len(cleaned_candidate) - 1)\n W_n1, W_n2, W_n3 = KG_NORMALIZATION_RERANK_NGRAM_WEIGHTS\n ngram_score = (\n # compute | Q ∩ E | / min(|Q|, |E|) for unigrams and bigrams (trigrams already computed)\n W_n1 * len(n1 & h_n1) / max(1, min(len(n1), len(h_n1)))\n + W_n2 * len(n2 & h_n2) / max(1, min(len(n2), len(h_n2)))\n + W_n3 * len(n3 & h_n3) / max(1, min(len(n3), len(h_n3)))\n ) / (W_n1, W_n1 + W_n2, 1.0)[grams_used]\n\n # compute damerau levenshtein distance to fuzzy match against typos\n W_leven = KG_NORMALIZATION_RERANK_LEVENSHTEIN_WEIGHT\n leven_score = normalized_similarity(cleaned_entity, cleaned_candidate)\n\n # combine scores\n score = (1.0 - W_leven) * ngram_score + W_leven * leven_score\n candidates[i] = (candidate_id_name, candidate_name, score)\n candidates = list(\n sorted(\n filter(lambda x: x[2] > KG_NORMALIZATION_RERANK_THRESHOLD, candidates),\n key=lambda x: x[2],\n reverse=True,\n )\n )\n if not candidates:\n return None\n\n return candidates[0][0]\n\n\ndef _get_existing_normalized_relationships(\n raw_relationships: list[str],\n) -> dict[str, dict[str, list[str]]]:\n \"\"\"\n Get existing normalized relationships from the database.\n \"\"\"\n\n relationship_type_map: dict[str, dict[str, list[str]]] = defaultdict(\n lambda: defaultdict(list)\n )\n relationship_pairs = list(\n {\n (\n get_entity_type(split_relationship_id(relationship)[0]),\n get_entity_type(split_relationship_id(relationship)[2]),\n )\n for relationship in raw_relationships\n }\n )\n\n with get_session_with_current_tenant() as db_session:\n relationships = get_relationships_for_entity_type_pairs(\n db_session, relationship_pairs\n )\n\n for relationship in relationships:\n relationship_type_map[relationship.source_entity_type_id_name][\n relationship.target_entity_type_id_name\n ].append(relationship.id_name)\n\n return relationship_type_map\n\n\ndef normalize_entities(\n raw_entities: list[str],\n raw_entities_w_attributes: list[str],\n allowed_docs_temp_view_name: str | None = None,\n) -> NormalizedEntities:\n \"\"\"\n Match each entity against a list of normalized entities using fuzzy matching.\n Returns the best matching normalized entity for each input entity.\n\n Args:\n raw_entities: list of entity strings to normalize, w/o attributes\n raw_entities_w_attributes: list of entity strings to normalize, w/ attributes\n\n Returns:\n list of normalized entity strings\n \"\"\"\n normalized_entities: list[str] = []\n normalized_entities_w_attributes: list[str] = []\n normalized_map: dict[str, str] = {}\n\n entity_attributes = [\n get_attributes(attr_entity) for attr_entity in raw_entities_w_attributes\n ]\n\n mapping: list[str | None] = run_functions_tuples_in_parallel(\n [\n (_normalize_one_entity, (entity, attributes, allowed_docs_temp_view_name))\n for entity, attributes in zip(raw_entities, entity_attributes)\n ]\n )\n for entity, attributes, normalized_entity in zip(\n raw_entities, entity_attributes, mapping\n ):\n if normalized_entity is not None:\n normalized_entities.append(normalized_entity)\n normalized_entities_w_attributes.append(\n make_entity_w_attributes(normalized_entity, attributes)\n )\n normalized_map[entity] = format_entity_id_for_models(normalized_entity)\n else:\n logger.warning(f\"No normalized entity found for {entity}\")\n normalized_map[entity] = format_entity_id_for_models(entity)\n\n return NormalizedEntities(\n entities=normalized_entities,\n entities_w_attributes=normalized_entities_w_attributes,\n entity_normalization_map=normalized_map,\n )\n\n\ndef normalize_relationships(\n raw_relationships: list[str], entity_normalization_map: dict[str, str]\n) -> NormalizedRelationships:\n \"\"\"\n Normalize relationships using entity mappings and relationship string matching.\n\n Args:\n relationships: list of relationships in format \"source__relation__target\"\n entity_normalization_map: Mapping of raw entities to normalized ones (or None)\n\n Returns:\n NormalizedRelationships containing normalized relationships and mapping\n \"\"\"\n # Placeholder for normalized relationship structure\n nor_relationships = _get_existing_normalized_relationships(raw_relationships)\n\n normalized_rels: list[str] = []\n normalization_map: dict[str, str] = {}\n\n for raw_rel in raw_relationships:\n # 1. Split and normalize entities\n try:\n source, rel_string, target = split_relationship_id(raw_rel)\n except ValueError:\n raise ValueError(f\"Invalid relationship format: {raw_rel}\")\n\n # Check if entities are in normalization map and not None\n norm_source = entity_normalization_map.get(source)\n norm_target = entity_normalization_map.get(target)\n\n if norm_source is None or norm_target is None:\n logger.warning(f\"No normalized entities found for {raw_rel}\")\n continue\n\n # 2. Find candidate normalized relationships\n candidate_rels = []\n norm_source_type = get_entity_type(format_entity_id_for_models(norm_source))\n norm_target_type = get_entity_type(format_entity_id_for_models(norm_target))\n if (\n norm_source_type in nor_relationships\n and norm_target_type in nor_relationships[norm_source_type]\n ):\n candidate_rels = [\n split_relationship_id(rel)[1]\n for rel in nor_relationships[norm_source_type][norm_target_type]\n ]\n\n if not candidate_rels:\n logger.warning(f\"No candidate relationships found for {raw_rel}\")\n continue\n\n # 3. Encode and find best match\n strings_to_encode = [rel_string] + candidate_rels\n vectors = encode_string_batch(strings_to_encode)\n\n # Get raw relation vector and candidate vectors\n raw_vector = vectors[0]\n candidate_vectors = vectors[1:]\n\n # Calculate dot products\n dot_products = np.dot(candidate_vectors, raw_vector)\n best_match_idx = np.argmax(dot_products)\n\n # Create normalized relationship\n norm_rel = make_relationship_id(\n norm_source, candidate_rels[best_match_idx], norm_target\n )\n normalized_rels.append(norm_rel)\n normalization_map[raw_rel] = norm_rel\n\n return NormalizedRelationships(\n relationships=normalized_rels, relationship_normalization_map=normalization_map\n )\n" }, { "path": "backend/onyx/kg/extractions/extraction_processing.py", "content": "import time\nfrom typing import Any\n\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.db.connector import get_kg_enabled_connectors\nfrom onyx.db.document import get_document_updated_at\nfrom onyx.db.document import get_skipped_kg_documents\nfrom onyx.db.document import get_unprocessed_kg_document_batch_for_connector\nfrom onyx.db.document import update_document_kg_info\nfrom onyx.db.document import update_document_kg_stage\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.entities import delete_from_kg_entities__no_commit\nfrom onyx.db.entities import upsert_staging_entity\nfrom onyx.db.entity_type import get_entity_types\nfrom onyx.db.kg_config import get_kg_config_settings\nfrom onyx.db.kg_config import validate_kg_settings\nfrom onyx.db.models import Document\nfrom onyx.db.models import KGStage\nfrom onyx.db.relationships import delete_from_kg_relationships__no_commit\nfrom onyx.db.relationships import upsert_staging_relationship\nfrom onyx.db.relationships import upsert_staging_relationship_type\nfrom onyx.kg.models import KGClassificationInstructions\nfrom onyx.kg.models import KGDocumentDeepExtractionResults\nfrom onyx.kg.models import KGEnhancedDocumentMetadata\nfrom onyx.kg.models import KGEntityTypeInstructions\nfrom onyx.kg.models import KGExtractionInstructions\nfrom onyx.kg.models import KGImpliedExtractionResults\nfrom onyx.kg.utils.extraction_utils import EntityTypeMetadataTracker\nfrom onyx.kg.utils.extraction_utils import (\n get_batch_documents_metadata,\n)\nfrom onyx.kg.utils.extraction_utils import kg_deep_extraction\nfrom onyx.kg.utils.extraction_utils import (\n kg_implied_extraction,\n)\nfrom onyx.kg.utils.formatting_utils import extract_relationship_type_id\nfrom onyx.kg.utils.formatting_utils import get_entity_type\nfrom onyx.kg.utils.formatting_utils import split_entity_id\nfrom onyx.kg.utils.formatting_utils import split_relationship_id\nfrom onyx.kg.utils.lock_utils import extend_lock\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\n\nlogger = setup_logger()\n\n\ndef _get_classification_extraction_instructions() -> (\n dict[str | None, dict[str, KGEntityTypeInstructions]]\n):\n \"\"\"\n Prepare the classification instructions for the given source.\n \"\"\"\n\n classification_instructions_dict: dict[\n str | None, dict[str, KGEntityTypeInstructions]\n ] = {}\n\n with get_session_with_current_tenant() as db_session:\n entity_types = get_entity_types(db_session, active=True)\n\n for entity_type in entity_types:\n grounded_source_name = entity_type.grounded_source_name\n\n if grounded_source_name not in classification_instructions_dict:\n classification_instructions_dict[grounded_source_name] = {}\n\n if grounded_source_name is None:\n continue\n\n attributes = entity_type.parsed_attributes\n classification_attributes = {\n option: info\n for option, info in attributes.classification_attributes.items()\n if info.extraction\n }\n classification_options = \", \".join(classification_attributes.keys())\n classification_enabled = (\n len(classification_options) > 0 and len(classification_attributes) > 0\n )\n\n classification_instructions_dict[grounded_source_name][entity_type.id_name] = (\n KGEntityTypeInstructions(\n metadata_attribute_conversion=attributes.metadata_attribute_conversion,\n classification_instructions=KGClassificationInstructions(\n classification_enabled=classification_enabled,\n classification_options=classification_options,\n classification_class_definitions=classification_attributes,\n ),\n extraction_instructions=KGExtractionInstructions(\n deep_extraction=entity_type.deep_extraction,\n active=entity_type.active,\n ),\n entity_filter_attributes=attributes.entity_filter_attributes,\n )\n )\n\n return classification_instructions_dict\n\n\ndef _get_batch_documents_enhanced_metadata(\n unprocessed_document_batch: list[Document],\n source_type_classification_extraction_instructions: dict[\n str, KGEntityTypeInstructions\n ],\n connector_source: str,\n) -> dict[str, KGEnhancedDocumentMetadata]:\n \"\"\"\n Get the entity types for the given unprocessed documents.\n \"\"\"\n\n kg_document_meta_data_dict: dict[str, KGEnhancedDocumentMetadata] = {\n document.id: KGEnhancedDocumentMetadata(\n entity_type=None,\n metadata_attribute_conversion=None,\n document_metadata=None,\n deep_extraction=False,\n classification_enabled=False,\n classification_instructions=None,\n skip=True,\n )\n for document in unprocessed_document_batch\n }\n\n batch_entity = None\n if len(source_type_classification_extraction_instructions) == 1:\n # if source only has one entity type, the document must be of that type\n batch_entity = list(source_type_classification_extraction_instructions.keys())[\n 0\n ]\n\n # the documents can be of multiple entity types. We need to identify the entity type for each document\n batch_metadata = get_batch_documents_metadata(\n [\n unprocessed_document.id\n for unprocessed_document in unprocessed_document_batch\n ],\n connector_source,\n )\n\n for metadata in batch_metadata:\n document_id = metadata.document_id\n doc_entity = None\n\n if not isinstance(document_id, str):\n continue\n\n chunk_metadata = metadata.source_metadata\n\n if batch_entity:\n doc_entity = batch_entity\n else:\n # TODO: make this a helper function\n if not chunk_metadata:\n continue\n\n for (\n potential_entity_type\n ) in source_type_classification_extraction_instructions.keys():\n potential_entity_type_attribute_filters = (\n source_type_classification_extraction_instructions[\n potential_entity_type\n ].entity_filter_attributes\n or {}\n )\n\n if not potential_entity_type_attribute_filters:\n continue\n\n if all(\n chunk_metadata.get(attribute)\n == potential_entity_type_attribute_filters.get(attribute)\n for attribute in potential_entity_type_attribute_filters\n ):\n doc_entity = potential_entity_type\n break\n\n if doc_entity is None:\n continue\n\n entity_instructions = source_type_classification_extraction_instructions[\n doc_entity\n ]\n\n kg_document_meta_data_dict[document_id] = KGEnhancedDocumentMetadata(\n entity_type=doc_entity,\n metadata_attribute_conversion=(\n source_type_classification_extraction_instructions[\n doc_entity\n ].metadata_attribute_conversion\n ),\n document_metadata=chunk_metadata,\n deep_extraction=entity_instructions.extraction_instructions.deep_extraction,\n classification_enabled=entity_instructions.classification_instructions.classification_enabled,\n classification_instructions=entity_instructions.classification_instructions,\n skip=False,\n )\n\n return kg_document_meta_data_dict\n\n\ndef kg_extraction(\n tenant_id: str,\n index_name: str,\n lock: RedisLock,\n processing_chunk_batch_size: int = 8,\n) -> None:\n \"\"\"\n This extraction will try to extract from all chunks that have not been kg-processed yet.\n\n Approach:\n - Get all connectors that are enabled for KG extraction\n - For each enabled connector:\n - Get unprocessed documents (using a generator)\n - For each batch of unprocessed documents:\n - Classify each document to select proper ones\n - Get and extract from chunks\n - Update chunks in Vespa\n - Update temporary KG extraction tables\n - Update document table to set kg_extracted = True\n \"\"\"\n\n logger.info(f\"Starting kg extraction for tenant {tenant_id}\")\n\n kg_config_settings = get_kg_config_settings()\n validate_kg_settings(kg_config_settings)\n\n # get connector ids that are enabled for KG extraction\n with get_session_with_current_tenant() as db_session:\n kg_enabled_connectors = get_kg_enabled_connectors(db_session)\n\n document_classification_extraction_instructions = (\n _get_classification_extraction_instructions()\n )\n\n # get entity type info\n with get_session_with_current_tenant() as db_session:\n all_entity_types = get_entity_types(db_session)\n active_entity_types = {\n entity_type.id_name\n for entity_type in get_entity_types(db_session, active=True)\n }\n\n # entity_type: (metadata: conversion property)\n entity_metadata_conversion_instructions = {\n entity_type.id_name: entity_type.parsed_attributes.metadata_attribute_conversion\n for entity_type in all_entity_types\n }\n\n # Track which metadata attributes are possible for each entity type\n metadata_tracker = EntityTypeMetadataTracker()\n metadata_tracker.import_typeinfo()\n\n last_lock_time = time.monotonic()\n\n # Iterate over connectors that are enabled for KG extraction\n for kg_enabled_connector in kg_enabled_connectors:\n connector_id = kg_enabled_connector.id\n connector_coverage_days = kg_enabled_connector.kg_coverage_days\n connector_source = kg_enabled_connector.source\n\n document_batch_counter = 0\n\n # iterate over un-kg-processed documents in connector\n while True:\n # get a batch of unprocessed documents\n with get_session_with_current_tenant() as db_session:\n unprocessed_document_batch = (\n get_unprocessed_kg_document_batch_for_connector(\n db_session,\n connector_id,\n kg_coverage_start=kg_config_settings.KG_COVERAGE_START_DATE,\n kg_max_coverage_days=connector_coverage_days\n or kg_config_settings.KG_MAX_COVERAGE_DAYS,\n batch_size=processing_chunk_batch_size,\n )\n )\n\n if len(unprocessed_document_batch) == 0:\n logger.info(\n f\"No unprocessed documents found for connector {connector_id}. Processed {document_batch_counter} batches.\"\n )\n break\n\n document_batch_counter += 1\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n logger.info(f\"Processing document batch {document_batch_counter}\")\n\n # Get the document attributes and entity types\n batch_metadata = _get_batch_documents_enhanced_metadata(\n unprocessed_document_batch,\n document_classification_extraction_instructions.get(\n connector_source, {}\n ),\n connector_source,\n )\n\n # mark docs in unprocessed_document_batch as EXTRACTING\n for unprocessed_document in unprocessed_document_batch:\n if batch_metadata[unprocessed_document.id].entity_type is None:\n # info for after the connector has been processed\n kg_stage = KGStage.SKIPPED\n logger.debug(\n f\"Document {unprocessed_document.id} is not of any entity type\"\n )\n elif batch_metadata[unprocessed_document.id].skip:\n # info for after the connector has been processed. But no message as there may be many\n # purposefully skipped documents\n kg_stage = KGStage.SKIPPED\n else:\n kg_stage = KGStage.EXTRACTING\n\n with get_session_with_current_tenant() as db_session:\n update_document_kg_stage(\n db_session,\n unprocessed_document.id,\n kg_stage,\n )\n\n if kg_stage == KGStage.EXTRACTING:\n delete_from_kg_relationships__no_commit(\n db_session, [unprocessed_document.id]\n )\n delete_from_kg_entities__no_commit(\n db_session, [unprocessed_document.id]\n )\n db_session.commit()\n\n # Iterate over batches of unprocessed documents\n # For each document:\n # - extract implied entities and relationships\n # - if deep extraction is enabled, extract entities and relationships with LLM\n # - if deep extraction and classification are enabled, classify document\n # - update postgres with\n # - extracted entities (with classification) and relationships\n # - kg_stage of the processed document\n\n documents_to_process = [x.id for x in unprocessed_document_batch]\n batch_implied_extraction: dict[str, KGImpliedExtractionResults] = {}\n batch_deep_extraction_args: list[\n tuple[str, KGEnhancedDocumentMetadata, KGImpliedExtractionResults]\n ] = []\n\n for unprocessed_document in unprocessed_document_batch:\n if (\n unprocessed_document.id not in documents_to_process\n or batch_metadata[unprocessed_document.id].entity_type is None\n or batch_metadata[unprocessed_document.id].skip\n ):\n with get_session_with_current_tenant() as db_session:\n update_document_kg_stage(\n db_session,\n unprocessed_document.id,\n KGStage.SKIPPED,\n )\n db_session.commit()\n continue\n\n # 1. perform (implicit) KG 'extractions' on the documents that should be processed\n # This is really about assigning document meta-data to KG entities/relationships or KG entity attributes\n # General approach:\n # - vendor emails to Employee-type entities + relationship to current primary grounded entity\n # - external account emails to Account-type entities + relationship to current primary grounded entity\n # - non-email owners to KG current entity's attributes, no relationships\n # We also collect email addresses of vendors and external accounts to inform chunk processing\n batch_implied_extraction[unprocessed_document.id] = (\n kg_implied_extraction(\n unprocessed_document,\n batch_metadata[unprocessed_document.id],\n active_entity_types,\n kg_config_settings,\n )\n )\n\n # 2. prepare inputs for deep extraction and classification\n if batch_metadata[unprocessed_document.id].deep_extraction:\n batch_deep_extraction_args.append(\n (\n unprocessed_document.id,\n batch_metadata[unprocessed_document.id],\n batch_implied_extraction[unprocessed_document.id],\n )\n )\n\n # 2. perform deep extraction and classification in parallel\n batch_deep_extraction_func_calls = [\n (\n kg_deep_extraction,\n (\n *arg,\n tenant_id,\n index_name,\n kg_config_settings,\n ),\n )\n for arg in batch_deep_extraction_args\n ]\n batch_deep_extractions: dict[str, KGDocumentDeepExtractionResults] = {\n document_id: result\n for document_id, result in zip(\n documents_to_process,\n run_functions_tuples_in_parallel(batch_deep_extraction_func_calls),\n )\n }\n\n # Collect entities and relationships to upsert\n batch_entities: list[tuple[str | None, str]] = []\n batch_relationships: list[tuple[str, str]] = []\n entity_classification: dict[str, str] = {}\n\n for document_id, implied_metadata in batch_implied_extraction.items():\n batch_entities += [\n (None, entity) for entity in implied_metadata.implied_entities\n ]\n batch_entities.append((document_id, implied_metadata.document_entity))\n batch_relationships += [\n (document_id, relationship)\n for relationship in implied_metadata.implied_relationships\n ]\n\n for document_id, deep_extraction_result in batch_deep_extractions.items():\n batch_entities += [\n (None, entity)\n for entity in deep_extraction_result.deep_extracted_entities\n ]\n for relationship in deep_extraction_result.deep_extracted_relationships:\n source_entity, _, target_entity = split_relationship_id(\n relationship\n )\n if (\n source_entity in active_entity_types\n and target_entity in active_entity_types\n ):\n batch_relationships += [(document_id, relationship)]\n\n classification_result = deep_extraction_result.classification_result\n if not classification_result:\n continue\n entity_classification[classification_result.document_entity] = (\n classification_result.classification_class\n )\n\n # Populate the KG database with the extracted entities, relationships, and terms\n for potential_document_id, entity in batch_entities:\n # verify the entity is valid\n parts = split_entity_id(entity)\n if len(parts) != 2:\n logger.error(\n f\"Invalid entity {entity} in aggregated_kg_extractions.entities\"\n )\n continue\n\n entity_type, entity_name = parts\n entity_type = entity_type.upper()\n entity_name = entity_name.capitalize()\n\n if entity_type not in active_entity_types:\n continue\n\n try:\n with get_session_with_current_tenant() as db_session:\n entity_attributes: dict[str, Any] = {}\n\n if potential_document_id:\n entity_attributes = (\n batch_metadata[potential_document_id].document_metadata\n or {}\n )\n\n # only keep selected attributes (and translate the attribute names)\n metadata_attributes = entity_metadata_conversion_instructions[\n entity_type\n ]\n keep_attributes = {\n metadata_attributes[attr_name].name: attr_val\n for attr_name, attr_val in entity_attributes.items()\n if (\n attr_name in metadata_attributes\n and metadata_attributes[attr_name].keep\n )\n }\n\n # add the classification result to the attributes\n if entity in entity_classification:\n keep_attributes[\"classification\"] = entity_classification[\n entity\n ]\n\n event_time = None\n if potential_document_id:\n event_time = get_document_updated_at(\n potential_document_id, db_session\n )\n\n upserted_entity = upsert_staging_entity(\n db_session=db_session,\n name=entity_name,\n entity_type=entity_type,\n document_id=potential_document_id,\n occurrences=1,\n attributes=keep_attributes,\n event_time=event_time,\n )\n metadata_tracker.track_metadata(\n entity_type, upserted_entity.attributes\n )\n\n db_session.commit()\n except Exception as e:\n logger.error(f\"Error adding entity {entity}. Error message: {e}\")\n\n for document_id, relationship in batch_relationships:\n relationship_split = split_relationship_id(relationship)\n\n if len(relationship_split) != 3:\n logger.error(\n f\"Invalid relationship {relationship} in aggregated_kg_extractions.relationships\"\n )\n continue\n\n source_entity, relationship_type, target_entity = relationship_split\n\n source_entity_type = get_entity_type(source_entity)\n target_entity_type = get_entity_type(target_entity)\n\n if (\n source_entity_type not in active_entity_types\n or target_entity_type not in active_entity_types\n ):\n continue\n\n relationship_type_id_name = extract_relationship_type_id(relationship)\n\n with get_session_with_current_tenant() as db_session:\n try:\n upsert_staging_relationship_type(\n db_session=db_session,\n source_entity_type=source_entity_type.upper(),\n relationship_type=relationship_type,\n target_entity_type=target_entity_type.upper(),\n definition=False,\n extraction_count=1,\n )\n db_session.commit()\n except Exception as e:\n logger.error(\n f\"Error adding relationship type {relationship_type_id_name} to the database: {e}\"\n )\n\n with get_session_with_current_tenant() as db_session:\n try:\n upsert_staging_relationship(\n db_session=db_session,\n relationship_id_name=relationship,\n source_document_id=document_id,\n occurrences=1,\n )\n db_session.commit()\n except Exception as e:\n logger.error(\n f\"Error adding relationship {relationship} to the database: {e}\"\n )\n\n # Populate the Documents table with the kg information for the documents\n\n for processed_document in documents_to_process:\n with get_session_with_current_tenant() as db_session:\n update_document_kg_info(\n db_session,\n processed_document,\n KGStage.EXTRACTED,\n )\n db_session.commit()\n\n # Update the the Skipped Docs back to Not Started\n with get_session_with_current_tenant() as db_session:\n skipped_documents = get_skipped_kg_documents(db_session)\n for document_id in skipped_documents:\n update_document_kg_stage(\n db_session,\n document_id,\n KGStage.NOT_STARTED,\n )\n db_session.commit()\n\n metadata_tracker.export_typeinfo()\n" }, { "path": "backend/onyx/kg/models.py", "content": "from datetime import datetime\nfrom enum import Enum\nfrom typing import Any\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.kg_configs import KG_DEFAULT_MAX_PARENT_RECURSION_DEPTH\n\n\n# Note: make sure to write a migration if adding a non-nullable field or removing a field\nclass KGConfigSettings(BaseModel):\n KG_EXPOSED: bool = False\n KG_ENABLED: bool = False\n KG_VENDOR: str | None = None\n KG_VENDOR_DOMAINS: list[str] = []\n KG_IGNORE_EMAIL_DOMAINS: list[str] = []\n KG_COVERAGE_START: str = datetime(1970, 1, 1).strftime(\"%Y-%m-%d\")\n KG_MAX_COVERAGE_DAYS: int = 10000\n KG_MAX_PARENT_RECURSION_DEPTH: int = KG_DEFAULT_MAX_PARENT_RECURSION_DEPTH\n KG_BETA_PERSONA_ID: int | None = None\n\n @property\n def KG_COVERAGE_START_DATE(self) -> datetime:\n return datetime.strptime(self.KG_COVERAGE_START, \"%Y-%m-%d\")\n\n\nclass KGGroundingType(str, Enum):\n UNGROUNDED = \"ungrounded\"\n GROUNDED = \"grounded\"\n\n\nclass KGAttributeTrackType(str, Enum):\n VALUE = \"value\"\n LIST = \"list\"\n\n\nclass KGAttributeTrackInfo(BaseModel):\n type: KGAttributeTrackType\n values: set[str] | None\n\n\nclass KGAttributeEntityOption(str, Enum):\n FROM_EMAIL = \"from_email\" # use email to determine type (ACCOUNT or EMPLOYEE)\n\n\nclass KGAttributeImplicationProperty(BaseModel):\n # type of implied entity to create\n # if str, will create an implied entity of that type\n # if KGAttributeEntityOption, will determine the type based on the option\n implied_entity_type: str | KGAttributeEntityOption\n # name of the implied relationship to create (from implied entity to this entity)\n implied_relationship_name: str\n\n\nclass KGAttributeProperty(BaseModel):\n # name of attribute to map metadata to\n name: str\n # whether to keep this attribute in the entity\n keep: bool\n # properties for creating implied entities and relations from this metadata\n implication_property: KGAttributeImplicationProperty | None = None\n\n\nclass KGEntityTypeClassificationInfo(BaseModel):\n extraction: bool\n description: str\n\n\nclass KGEntityTypeAttributes(BaseModel):\n # information on how to use the metadata to extract attributes, implied entities, and relations\n metadata_attribute_conversion: dict[str, KGAttributeProperty] = {}\n # a metadata key: value pair to match for to differentiate entities from the same source\n entity_filter_attributes: dict[str, Any] = {}\n # mapping of classification names to their corresponding classification info\n classification_attributes: dict[str, KGEntityTypeClassificationInfo] = {}\n\n # mapping of attribute names to their allowed values, populated during extraction\n attribute_values: dict[str, KGAttributeTrackInfo | None] = {}\n\n\nclass KGEntityTypeDefinition(BaseModel):\n description: str\n grounding: KGGroundingType\n grounded_source_name: DocumentSource | None\n active: bool = False\n attributes: KGEntityTypeAttributes = KGEntityTypeAttributes()\n entity_values: list[str] = []\n\n\nclass KGChunkFormat(BaseModel):\n connector_id: int | None = None\n document_id: str\n chunk_id: int\n title: str\n content: str\n primary_owners: list[str]\n secondary_owners: list[str]\n source_type: str\n metadata: dict[str, str | list[str]] | None = None\n\n\nclass KGPerson(BaseModel):\n name: str\n company: str\n employee: bool\n\n\nclass NormalizedEntities(BaseModel):\n entities: list[str]\n entities_w_attributes: list[str]\n entity_normalization_map: dict[str, str]\n\n\nclass NormalizedRelationships(BaseModel):\n relationships: list[str]\n relationship_normalization_map: dict[str, str]\n\n\nclass KGMetadataContent(BaseModel):\n document_id: str\n source_type: str\n source_metadata: dict[str, Any] | None = None\n\n\nclass KGClassificationInstructions(BaseModel):\n classification_enabled: bool\n classification_options: str\n classification_class_definitions: dict[str, KGEntityTypeClassificationInfo]\n\n\nclass KGExtractionInstructions(BaseModel):\n deep_extraction: bool\n active: bool\n\n\nclass KGEntityTypeInstructions(BaseModel):\n metadata_attribute_conversion: dict[str, KGAttributeProperty]\n classification_instructions: KGClassificationInstructions\n extraction_instructions: KGExtractionInstructions\n entity_filter_attributes: dict[str, Any] | None = None\n\n\nclass KGEnhancedDocumentMetadata(BaseModel):\n entity_type: str | None\n metadata_attribute_conversion: dict[str, KGAttributeProperty] | None\n document_metadata: dict[str, Any] | None\n deep_extraction: bool\n classification_enabled: bool\n classification_instructions: KGClassificationInstructions | None\n skip: bool\n\n\nclass KGConnectorData(BaseModel):\n id: int\n source: str\n kg_coverage_days: int | None\n\n\nclass KGStage(str, Enum):\n EXTRACTED = \"extracted\"\n NORMALIZED = \"normalized\"\n FAILED = \"failed\"\n SKIPPED = \"skipped\"\n NOT_STARTED = \"not_started\"\n EXTRACTING = \"extracting\"\n DO_NOT_EXTRACT = \"do_not_extract\"\n\n\nclass KGClassificationResult(BaseModel):\n document_entity: str\n classification_class: str\n\n\nclass KGImpliedExtractionResults(BaseModel):\n document_entity: str\n implied_entities: set[str]\n implied_relationships: set[str]\n company_participant_emails: set[str]\n account_participant_emails: set[str]\n\n\nclass KGDocumentDeepExtractionResults(BaseModel):\n classification_result: KGClassificationResult | None\n deep_extracted_entities: set[str]\n deep_extracted_relationships: set[str]\n\n\nclass KGException(Exception):\n pass\n" }, { "path": "backend/onyx/kg/resets/reset_index.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.db.document import reset_all_document_kg_stages\nfrom onyx.db.models import Connector\nfrom onyx.db.models import KGEntity\nfrom onyx.db.models import KGEntityExtractionStaging\nfrom onyx.db.models import KGEntityType\nfrom onyx.db.models import KGRelationship\nfrom onyx.db.models import KGRelationshipExtractionStaging\nfrom onyx.db.models import KGRelationshipType\nfrom onyx.db.models import KGRelationshipTypeExtractionStaging\n\n\ndef reset_full_kg_index__commit(db_session: Session) -> None:\n \"\"\"\n Resets the knowledge graph index.\n \"\"\"\n\n db_session.query(KGRelationship).delete()\n db_session.query(KGRelationshipType).delete()\n db_session.query(KGEntity).delete()\n db_session.query(KGRelationshipExtractionStaging).delete()\n db_session.query(KGEntityExtractionStaging).delete()\n db_session.query(KGRelationshipTypeExtractionStaging).delete()\n # Update all connectors to disable KG processing\n db_session.query(Connector).update({\"kg_processing_enabled\": False})\n\n # Only reset grounded entity types\n db_session.query(KGEntityType).filter(\n KGEntityType.grounded_source_name.isnot(None)\n ).update({\"active\": False})\n\n reset_all_document_kg_stages(db_session)\n\n db_session.commit()\n" }, { "path": "backend/onyx/kg/resets/reset_source.py", "content": "from redis.lock import Lock as RedisLock\nfrom sqlalchemy import or_\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import Connector\nfrom onyx.db.models import Document\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import KGEntity\nfrom onyx.db.models import KGEntityExtractionStaging\nfrom onyx.db.models import KGEntityType\nfrom onyx.db.models import KGRelationship\nfrom onyx.db.models import KGRelationshipExtractionStaging\nfrom onyx.db.models import KGRelationshipType\nfrom onyx.db.models import KGRelationshipTypeExtractionStaging\nfrom onyx.db.models import KGStage\nfrom onyx.kg.resets.reset_index import reset_full_kg_index__commit\nfrom onyx.kg.resets.reset_vespa import reset_vespa_kg_index\n\n\ndef reset_source_kg_index(\n source_name: str | None, tenant_id: str, index_name: str, lock: RedisLock\n) -> None:\n \"\"\"\n Resets the knowledge graph index and vespa for a source.\n \"\"\"\n # reset vespa for the source\n reset_vespa_kg_index(tenant_id, index_name, lock, source_name)\n\n with get_session_with_current_tenant() as db_session:\n if source_name is None:\n reset_full_kg_index__commit(db_session)\n return\n\n # get all the entity types for the given source\n entity_types = [\n et.id_name\n for et in db_session.query(KGEntityType)\n .filter(KGEntityType.grounded_source_name == source_name)\n .all()\n ]\n if not entity_types:\n raise ValueError(f\"There are no entity types for the source {source_name}\")\n\n # delete the entity type from the knowledge graph\n for entity_type in entity_types:\n db_session.query(KGRelationship).filter(\n or_(\n KGRelationship.source_node_type == entity_type,\n KGRelationship.target_node_type == entity_type,\n )\n ).delete()\n db_session.query(KGRelationshipType).filter(\n or_(\n KGRelationshipType.source_entity_type_id_name == entity_type,\n KGRelationshipType.target_entity_type_id_name == entity_type,\n )\n ).delete()\n db_session.query(KGEntity).filter(\n KGEntity.entity_type_id_name == entity_type\n ).delete()\n db_session.query(KGRelationshipExtractionStaging).filter(\n or_(\n KGRelationshipExtractionStaging.source_node_type == entity_type,\n KGRelationshipExtractionStaging.target_node_type == entity_type,\n )\n ).delete()\n db_session.query(KGEntityExtractionStaging).filter(\n KGEntityExtractionStaging.entity_type_id_name == entity_type\n ).delete()\n db_session.query(KGRelationshipTypeExtractionStaging).filter(\n or_(\n KGRelationshipTypeExtractionStaging.source_entity_type_id_name\n == entity_type,\n KGRelationshipTypeExtractionStaging.target_entity_type_id_name\n == entity_type,\n )\n ).delete()\n db_session.commit()\n\n with get_session_with_current_tenant() as db_session:\n # get all the documents for the given source\n kg_connectors = [\n connector.id\n for connector in db_session.query(Connector)\n .filter(Connector.source == DocumentSource(source_name))\n .all()\n ]\n document_ids = [\n cc_pair.id\n for cc_pair in db_session.query(DocumentByConnectorCredentialPair)\n .filter(DocumentByConnectorCredentialPair.connector_id.in_(kg_connectors))\n .all()\n ]\n\n # reset the kg stage for the documents\n db_session.query(Document).filter(Document.id.in_(document_ids)).update(\n {\"kg_stage\": KGStage.NOT_STARTED}\n )\n db_session.commit()\n" }, { "path": "backend/onyx/kg/resets/reset_vespa.py", "content": "import time\nfrom typing import Any\n\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.document import get_num_chunks_for_document\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import Connector\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import KGEntityType\nfrom onyx.document_index.document_index_utils import get_uuid_from_chunk_info\nfrom onyx.document_index.vespa.index import KGVespaChunkUpdateRequest\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\nfrom onyx.kg.utils.lock_utils import extend_lock\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\n\nlogger = setup_logger()\n\n\ndef _reset_vespa_for_doc(document_id: str, tenant_id: str, index_name: str) -> None:\n vespa_index = VespaIndex(\n index_name=index_name,\n secondary_index_name=None,\n large_chunks_enabled=False,\n secondary_large_chunks_enabled=False,\n multitenant=MULTI_TENANT,\n httpx_client=None,\n )\n\n reset_update_dict: dict[str, Any] = {\n \"fields\": {\n \"kg_entities\": {\"assign\": []},\n \"kg_relationships\": {\"assign\": []},\n \"kg_terms\": {\"assign\": []},\n }\n }\n\n with get_session_with_current_tenant() as db_session:\n num_chunks = get_num_chunks_for_document(db_session, document_id)\n\n vespa_requests: list[KGVespaChunkUpdateRequest] = []\n for chunk_num in range(num_chunks):\n doc_chunk_id = get_uuid_from_chunk_info(\n document_id=document_id,\n chunk_id=chunk_num,\n tenant_id=tenant_id,\n large_chunk_id=None,\n )\n vespa_requests.append(\n KGVespaChunkUpdateRequest(\n document_id=document_id,\n chunk_id=chunk_num,\n url=f\"{DOCUMENT_ID_ENDPOINT.format(index_name=vespa_index.index_name)}/{doc_chunk_id}\",\n update_request=reset_update_dict,\n )\n )\n\n with vespa_index.httpx_client_context as httpx_client:\n vespa_index._apply_kg_chunk_updates_batched(vespa_requests, httpx_client)\n\n\ndef reset_vespa_kg_index(\n tenant_id: str, index_name: str, lock: RedisLock, source_name: str | None = None\n) -> None:\n \"\"\"\n Reset the kg info in vespa for all documents of a given source name,\n or all documents from kg grounded sources if source_name is None.\n \"\"\"\n logger.info(\n f\"Resetting kg vespa index {index_name} for tenant {tenant_id}, source: {source_name if source_name else 'all'}\"\n )\n\n last_lock_time = time.monotonic()\n\n # Get all documents that need a vespa reset\n with get_session_with_current_tenant() as db_session:\n if source_name:\n # get all connectors of the given source name\n kg_connectors = [\n connector.id\n for connector in db_session.query(Connector)\n .filter(Connector.source == DocumentSource(source_name))\n .all()\n ]\n else:\n # get all connectors that have kg enabled\n kg_sources = [\n DocumentSource(et.grounded_source_name)\n for et in db_session.query(KGEntityType)\n .filter(\n KGEntityType.grounded_source_name.is_not(None),\n KGEntityType.active.is_(True),\n )\n .distinct()\n .all()\n ]\n kg_connectors = [\n connector.id\n for connector in db_session.query(Connector)\n .filter(Connector.source.in_(kg_sources))\n .all()\n ]\n\n # Get all the documents for the given connectors\n document_ids = [\n cc_pair.id\n for cc_pair in db_session.query(DocumentByConnectorCredentialPair)\n .filter(DocumentByConnectorCredentialPair.connector_id.in_(kg_connectors))\n .all()\n ]\n\n # Reset the kg fields\n for document_id in document_ids:\n _reset_vespa_for_doc(document_id, tenant_id, index_name)\n last_lock_time = extend_lock(\n lock, CELERY_GENERIC_BEAT_LOCK_TIMEOUT, last_lock_time\n )\n\n logger.info(\n f\"Finished resetting kg vespa index {index_name} for tenant {tenant_id}, source: {source_name if source_name else 'all'}\"\n )\n" }, { "path": "backend/onyx/kg/setup/kg_default_entity_definitions.py", "content": "from typing import cast\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.entity_type import KGEntityType\nfrom onyx.db.kg_config import get_kg_config_settings\nfrom onyx.db.kg_config import validate_kg_settings\nfrom onyx.kg.models import KGAttributeEntityOption\nfrom onyx.kg.models import KGAttributeImplicationProperty\nfrom onyx.kg.models import KGAttributeProperty\nfrom onyx.kg.models import KGEntityTypeAttributes\nfrom onyx.kg.models import KGEntityTypeClassificationInfo\nfrom onyx.kg.models import KGEntityTypeDefinition\nfrom onyx.kg.models import KGGroundingType\n\n\ndef get_default_entity_types(vendor_name: str) -> dict[str, KGEntityTypeDefinition]:\n return {\n \"LINEAR\": KGEntityTypeDefinition(\n description=\"A formal Linear ticket about a product issue or improvement request.\",\n attributes=KGEntityTypeAttributes(\n metadata_attribute_conversion={\n \"team\": KGAttributeProperty(name=\"team\", keep=True),\n \"state\": KGAttributeProperty(name=\"state\", keep=True),\n \"priority\": KGAttributeProperty(name=\"priority\", keep=True),\n \"estimate\": KGAttributeProperty(name=\"estimate\", keep=True),\n \"created_at\": KGAttributeProperty(name=\"created_at\", keep=True),\n \"started_at\": KGAttributeProperty(name=\"started_at\", keep=True),\n \"completed_at\": KGAttributeProperty(name=\"completed_at\", keep=True),\n \"due_date\": KGAttributeProperty(name=\"due_date\", keep=True),\n \"creator\": KGAttributeProperty(\n name=\"creator\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_creator_of\",\n ),\n ),\n \"assignee\": KGAttributeProperty(\n name=\"assignee\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_assignee_of\",\n ),\n ),\n },\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.LINEAR,\n ),\n \"JIRA\": KGEntityTypeDefinition(\n description=(\n \"A formal Jira ticket about a product issue or improvement request.\"\n ),\n attributes=KGEntityTypeAttributes(\n metadata_attribute_conversion={\n \"issuetype\": KGAttributeProperty(name=\"subtype\", keep=True),\n \"status\": KGAttributeProperty(name=\"status\", keep=True),\n \"priority\": KGAttributeProperty(name=\"priority\", keep=True),\n \"project_name\": KGAttributeProperty(name=\"project\", keep=True),\n \"created\": KGAttributeProperty(name=\"created_at\", keep=True),\n \"updated\": KGAttributeProperty(name=\"updated_at\", keep=True),\n \"resolution_date\": KGAttributeProperty(\n name=\"completed_at\", keep=True\n ),\n \"duedate\": KGAttributeProperty(name=\"due_date\", keep=True),\n \"reporter_email\": KGAttributeProperty(\n name=\"creator\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_creator_of\",\n ),\n ),\n \"assignee_email\": KGAttributeProperty(\n name=\"assignee\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_assignee_of\",\n ),\n ),\n # not using implication property as that only captures 1 depth\n \"key\": KGAttributeProperty(name=\"key\", keep=True),\n \"parent\": KGAttributeProperty(name=\"parent\", keep=True),\n },\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.JIRA,\n ),\n \"GITHUB_PR\": KGEntityTypeDefinition(\n description=\"A formal engineering request to merge proposed changes into the codebase.\",\n attributes=KGEntityTypeAttributes(\n metadata_attribute_conversion={\n \"repo\": KGAttributeProperty(name=\"repository\", keep=True),\n \"state\": KGAttributeProperty(name=\"state\", keep=True),\n \"num_commits\": KGAttributeProperty(name=\"num_commits\", keep=True),\n \"num_files_changed\": KGAttributeProperty(\n name=\"num_files_changed\", keep=True\n ),\n \"labels\": KGAttributeProperty(name=\"labels\", keep=True),\n \"merged\": KGAttributeProperty(name=\"merged\", keep=True),\n \"merged_at\": KGAttributeProperty(name=\"merged_at\", keep=True),\n \"closed_at\": KGAttributeProperty(name=\"closed_at\", keep=True),\n \"created_at\": KGAttributeProperty(name=\"created_at\", keep=True),\n \"updated_at\": KGAttributeProperty(name=\"updated_at\", keep=True),\n \"user\": KGAttributeProperty(\n name=\"creator\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_creator_of\",\n ),\n ),\n \"assignees\": KGAttributeProperty(\n name=\"assignees\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_assignee_of\",\n ),\n ),\n },\n entity_filter_attributes={\"object_type\": \"PullRequest\"},\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.GITHUB,\n ),\n \"GITHUB_ISSUE\": KGEntityTypeDefinition(\n description=\"A formal engineering ticket about an issue, idea, inquiry, or task.\",\n attributes=KGEntityTypeAttributes(\n metadata_attribute_conversion={\n \"repo\": KGAttributeProperty(name=\"repository\", keep=True),\n \"state\": KGAttributeProperty(name=\"state\", keep=True),\n \"labels\": KGAttributeProperty(name=\"labels\", keep=True),\n \"closed_at\": KGAttributeProperty(name=\"closed_at\", keep=True),\n \"created_at\": KGAttributeProperty(name=\"created_at\", keep=True),\n \"updated_at\": KGAttributeProperty(name=\"updated_at\", keep=True),\n \"user\": KGAttributeProperty(\n name=\"creator\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_creator_of\",\n ),\n ),\n \"assignees\": KGAttributeProperty(\n name=\"assignees\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=KGAttributeEntityOption.FROM_EMAIL,\n implied_relationship_name=\"is_assignee_of\",\n ),\n ),\n },\n entity_filter_attributes={\"object_type\": \"Issue\"},\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.GITHUB,\n ),\n \"FIREFLIES\": KGEntityTypeDefinition(\n description=(\n f\"A phone call transcript between us ({vendor_name}) and another account or individuals, or an internal meeting.\"\n ),\n attributes=KGEntityTypeAttributes(\n classification_attributes={\n \"customer\": KGEntityTypeClassificationInfo(\n extraction=True,\n description=\"a call with representatives of one or more customers prospects\",\n ),\n \"internal\": KGEntityTypeClassificationInfo(\n extraction=True,\n description=\"a call between employees of the vendor's company (a vendor-internal call)\",\n ),\n \"interview\": KGEntityTypeClassificationInfo(\n extraction=True,\n description=(\n \"a call with an individual who is interviewed or is discussing potential employment with the vendor\"\n ),\n ),\n \"other\": KGEntityTypeClassificationInfo(\n extraction=True,\n description=(\n \"a call with representatives of companies having a different reason for the call \"\n \"(investment, partnering, etc.)\"\n ),\n ),\n },\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.FIREFLIES,\n ),\n \"ACCOUNT\": KGEntityTypeDefinition(\n description=(\n \"A company that was, is, or potentially could be a customer of the vendor \"\n f\"('us, {vendor_name}'). Note that {vendor_name} can never be an ACCOUNT.\"\n ),\n attributes=KGEntityTypeAttributes(\n entity_filter_attributes={\"object_type\": \"Account\"},\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.SALESFORCE,\n ),\n \"OPPORTUNITY\": KGEntityTypeDefinition(\n description=\"A sales opportunity.\",\n attributes=KGEntityTypeAttributes(\n metadata_attribute_conversion={\n \"name\": KGAttributeProperty(name=\"name\", keep=True),\n \"stage_name\": KGAttributeProperty(name=\"stage\", keep=True),\n \"type\": KGAttributeProperty(name=\"type\", keep=True),\n \"amount\": KGAttributeProperty(name=\"amount\", keep=True),\n \"fiscal_year\": KGAttributeProperty(name=\"fiscal_year\", keep=True),\n \"fiscal_quarter\": KGAttributeProperty(\n name=\"fiscal_quarter\", keep=True\n ),\n \"is_closed\": KGAttributeProperty(name=\"is_closed\", keep=True),\n \"close_date\": KGAttributeProperty(name=\"close_date\", keep=True),\n \"probability\": KGAttributeProperty(\n name=\"close_probability\", keep=True\n ),\n \"created_date\": KGAttributeProperty(name=\"created_at\", keep=True),\n \"last_modified_date\": KGAttributeProperty(\n name=\"updated_at\", keep=True\n ),\n \"account\": KGAttributeProperty(\n name=\"account\",\n keep=False,\n implication_property=KGAttributeImplicationProperty(\n implied_entity_type=\"ACCOUNT\",\n implied_relationship_name=\"is_account_of\",\n ),\n ),\n },\n entity_filter_attributes={\"object_type\": \"Opportunity\"},\n ),\n grounding=KGGroundingType.GROUNDED,\n grounded_source_name=DocumentSource.SALESFORCE,\n ),\n \"VENDOR\": KGEntityTypeDefinition(\n description=f\"The Vendor {vendor_name}, 'us'\",\n grounding=KGGroundingType.GROUNDED,\n active=True,\n grounded_source_name=None,\n ),\n \"EMPLOYEE\": KGEntityTypeDefinition(\n description=(\n f\"A person who speaks on behalf of 'our' company (the VENDOR {vendor_name}), \"\n \"NOT of another account. Therefore, employees of other companies \"\n \"are NOT included here. If in doubt, do NOT extract.\"\n ),\n grounding=KGGroundingType.GROUNDED,\n active=False,\n grounded_source_name=None,\n ),\n }\n\n\ndef populate_missing_default_entity_types__commit(db_session: Session) -> None:\n \"\"\"\n Populates the database with the missing default entity types.\n \"\"\"\n kg_config_settings = get_kg_config_settings()\n validate_kg_settings(kg_config_settings)\n\n vendor_name = cast(str, kg_config_settings.KG_VENDOR)\n\n existing_entity_types = {et.id_name for et in db_session.query(KGEntityType).all()}\n\n default_entity_types = get_default_entity_types(vendor_name=vendor_name)\n for entity_type_id_name, entity_type_definition in default_entity_types.items():\n if entity_type_id_name in existing_entity_types:\n continue\n\n grounded_source_name = (\n entity_type_definition.grounded_source_name.value\n if entity_type_definition.grounded_source_name\n else None\n )\n kg_entity_type = KGEntityType(\n id_name=entity_type_id_name,\n description=entity_type_definition.description,\n attributes=entity_type_definition.attributes.model_dump(),\n grounding=entity_type_definition.grounding,\n grounded_source_name=grounded_source_name,\n active=entity_type_definition.active,\n )\n db_session.add(kg_entity_type)\n db_session.commit()\n" }, { "path": "backend/onyx/kg/utils/embeddings.py", "content": "from typing import List\n\nimport numpy as np\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.natural_language_processing.search_nlp_models import EmbedTextType\nfrom shared_configs.configs import MODEL_SERVER_HOST\nfrom shared_configs.configs import MODEL_SERVER_PORT\n\n\ndef encode_string_batch(strings: List[str]) -> np.ndarray:\n with get_session_with_current_tenant() as db_session:\n current_search_settings = get_current_search_settings(db_session)\n model = EmbeddingModel.from_db_model(\n search_settings=current_search_settings,\n server_host=MODEL_SERVER_HOST,\n server_port=MODEL_SERVER_PORT,\n )\n # Get embeddings while session is still open\n embedding = model.encode(strings, text_type=EmbedTextType.QUERY)\n return np.array(embedding)\n" }, { "path": "backend/onyx/kg/utils/extraction_utils.py", "content": "import json\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCallTypes\nfrom onyx.configs.kg_configs import KG_METADATA_TRACKING_THRESHOLD\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.entities import get_kg_entity_by_document\nfrom onyx.db.entity_type import get_entity_types\nfrom onyx.db.kg_config import KGConfigSettings\nfrom onyx.db.models import Document\nfrom onyx.db.models import KGEntityType\nfrom onyx.db.models import KGRelationshipType\nfrom onyx.db.tag import get_structured_tags_for_document\nfrom onyx.kg.models import KGAttributeEntityOption\nfrom onyx.kg.models import KGAttributeTrackInfo\nfrom onyx.kg.models import KGAttributeTrackType\nfrom onyx.kg.models import KGChunkFormat\nfrom onyx.kg.models import KGClassificationInstructions\nfrom onyx.kg.models import KGClassificationResult\nfrom onyx.kg.models import KGDocumentDeepExtractionResults\nfrom onyx.kg.models import KGEnhancedDocumentMetadata\nfrom onyx.kg.models import KGImpliedExtractionResults\nfrom onyx.kg.models import KGMetadataContent\nfrom onyx.kg.utils.formatting_utils import extract_email\nfrom onyx.kg.utils.formatting_utils import get_entity_type\nfrom onyx.kg.utils.formatting_utils import kg_email_processing\nfrom onyx.kg.utils.formatting_utils import make_entity_id\nfrom onyx.kg.utils.formatting_utils import make_relationship_id\nfrom onyx.kg.utils.formatting_utils import make_relationship_type_id\nfrom onyx.kg.vespa.vespa_interactions import get_document_vespa_contents\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.prompts.kg_prompts import CALL_CHUNK_PREPROCESSING_PROMPT\nfrom onyx.prompts.kg_prompts import CALL_DOCUMENT_CLASSIFICATION_PROMPT\nfrom onyx.prompts.kg_prompts import GENERAL_CHUNK_PREPROCESSING_PROMPT\nfrom onyx.prompts.kg_prompts import MASTER_EXTRACTION_PROMPT\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef get_entity_types_str(active: bool | None = None) -> str:\n \"\"\"\n Format the entity types into a string for the LLM.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n entity_types = get_entity_types(db_session, active)\n\n entity_types_list: list[str] = []\n for entity_type in entity_types:\n if entity_type.description:\n entity_description = \"\\n - Description: \" + entity_type.description\n else:\n entity_description = \"\"\n\n if entity_type.entity_values:\n allowed_values = \"\\n - Allowed Values: \" + \", \".join(\n entity_type.entity_values\n )\n else:\n allowed_values = \"\"\n\n attributes = entity_type.parsed_attributes\n\n entity_type_attribute_list: list[str] = []\n for attribute, values in attributes.attribute_values.items():\n entity_type_attribute_list.append(\n f\"{attribute}: {trackinfo_to_str(values)}\"\n )\n\n if attributes.classification_attributes:\n entity_type_attribute_list.append(\n # TODO: restructure classification attribute to be a dict of attribute name to classification info\n # e.g., {scope: {internal: prompt, external: prompt}, sentiment: {positive: prompt, negative: prompt}}\n \"classification: one of: \"\n + \", \".join(attributes.classification_attributes.keys())\n )\n if entity_type_attribute_list:\n entity_attributes = \"\\n - Attributes:\\n - \" + \"\\n - \".join(\n entity_type_attribute_list\n )\n else:\n entity_attributes = \"\"\n\n entity_types_list.append(\n entity_type.id_name\n + entity_description\n + allowed_values\n + entity_attributes\n )\n\n return \"\\n\".join(entity_types_list)\n\n\ndef get_relationship_types_str(active: bool | None = None) -> str:\n \"\"\"\n Format the relationship types into a string for the LLM.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n active_filters = []\n if active is not None:\n active_filters.append(KGRelationshipType.active == active)\n\n relationship_types = (\n db_session.query(KGRelationshipType).filter(*active_filters).all()\n )\n\n relationship_types_list = []\n for rel_type in relationship_types:\n # Format as \"source_type__relationship_type__target_type\"\n formatted_type = make_relationship_type_id(\n rel_type.source_entity_type_id_name,\n rel_type.type,\n rel_type.target_entity_type_id_name,\n )\n relationship_types_list.append(formatted_type)\n\n return \"\\n\".join(relationship_types_list)\n\n\ndef kg_process_owners(\n owner_emails: list[str],\n document_entity_id: str,\n relationship_type: str,\n kg_config_settings: KGConfigSettings,\n active_entity_types: set[str],\n) -> tuple[set[str], set[str], set[str], set[str]]:\n owner_entities: set[str] = set()\n owner_relationships: set[str] = set()\n company_participant_emails: set[str] = set()\n account_participant_emails: set[str] = set()\n\n for owner_email in owner_emails:\n if extract_email(owner_email) is None:\n continue\n\n process_results = kg_process_person(\n owner_email,\n document_entity_id,\n relationship_type,\n kg_config_settings,\n active_entity_types,\n )\n if process_results is None:\n continue\n\n (\n owner_entity,\n owner_relationship,\n company_participant_email,\n account_participant_email,\n ) = process_results\n\n owner_entities.add(owner_entity)\n owner_relationships.add(owner_relationship)\n if company_participant_email:\n company_participant_emails.add(company_participant_email)\n if account_participant_email:\n account_participant_emails.add(account_participant_email)\n\n return (\n owner_entities,\n owner_relationships,\n company_participant_emails,\n account_participant_emails,\n )\n\n\ndef kg_implied_extraction(\n document: Document,\n doc_metadata: KGEnhancedDocumentMetadata,\n active_entity_types: set[str],\n kg_config_settings: KGConfigSettings,\n) -> KGImpliedExtractionResults:\n \"\"\"\n Generate entities, relationships, and attributes for a document.\n \"\"\"\n\n # Get document entity and metadata stuff from the KGEnhancedDocumentMetadata\n document_entity_type = doc_metadata.entity_type\n document_metadata = doc_metadata.document_metadata or {}\n metadata_attribute_conversion = doc_metadata.metadata_attribute_conversion\n if document_entity_type is None or metadata_attribute_conversion is None:\n raise ValueError(\"Entity type and metadata attributes are required\")\n\n implied_entities: set[str] = set()\n implied_relationships: set[str] = set()\n\n # Quantity needed for call processing - participants from vendor\n company_participant_emails: set[str] = set()\n # Quantity needed for call processing - external participants\n account_participant_emails: set[str] = set()\n\n # Chunk treatment variables\n\n document_is_from_call = document_entity_type.lower() in (\n call_type.value.lower() for call_type in OnyxCallTypes\n )\n\n # Get core entity\n\n document_id = document.id\n primary_owners = document.primary_owners\n secondary_owners = document.secondary_owners\n\n with get_session_with_current_tenant() as db_session:\n document_entity = get_kg_entity_by_document(db_session, document_id)\n\n if document_entity:\n document_entity_id = document_entity.id_name\n else:\n document_entity_id = make_entity_id(document_entity_type, document_id)\n\n # Get implied entities and relationships from primary/secondary owners\n\n if document_is_from_call:\n (\n implied_entities,\n implied_relationships,\n company_participant_emails,\n account_participant_emails,\n ) = kg_process_owners(\n owner_emails=(primary_owners or []) + (secondary_owners or []),\n document_entity_id=document_entity_id,\n relationship_type=\"participates_in\",\n kg_config_settings=kg_config_settings,\n active_entity_types=active_entity_types,\n )\n else:\n (\n implied_entities,\n implied_relationships,\n company_participant_emails,\n account_participant_emails,\n ) = kg_process_owners(\n owner_emails=primary_owners or [],\n document_entity_id=document_entity_id,\n relationship_type=\"leads\",\n kg_config_settings=kg_config_settings,\n active_entity_types=active_entity_types,\n )\n\n (\n participant_entities,\n participant_relationships,\n company_emails,\n account_emails,\n ) = kg_process_owners(\n owner_emails=secondary_owners or [],\n document_entity_id=document_entity_id,\n relationship_type=\"participates_in\",\n kg_config_settings=kg_config_settings,\n active_entity_types=active_entity_types,\n )\n implied_entities.update(participant_entities)\n implied_relationships.update(participant_relationships)\n company_participant_emails.update(company_emails)\n account_participant_emails.update(account_emails)\n\n # Get implied entities and relationships from document metadata\n for metadata, value in document_metadata.items():\n # get implication property for this metadata\n if metadata not in metadata_attribute_conversion:\n continue\n if (\n implication_property := metadata_attribute_conversion[\n metadata\n ].implication_property\n ) is None:\n continue\n\n if not isinstance(value, str) and not isinstance(value, list):\n continue\n values: list[str] = [value] if isinstance(value, str) else value\n\n # create implied entities and relationships\n for item in values:\n if (\n implication_property.implied_entity_type\n == KGAttributeEntityOption.FROM_EMAIL\n ):\n # determine entity type from email\n email = extract_email(item)\n if email is None:\n continue\n process_results = kg_process_person(\n email=email,\n document_entity_id=document_entity_id,\n relationship_type=implication_property.implied_relationship_name,\n kg_config_settings=kg_config_settings,\n active_entity_types=active_entity_types,\n )\n if process_results is None:\n continue\n\n (implied_entity, implied_relationship, _, _) = process_results\n implied_entities.add(implied_entity)\n implied_relationships.add(implied_relationship)\n else:\n # use the given entity type\n entity_type = implication_property.implied_entity_type\n if entity_type not in active_entity_types:\n continue\n\n implied_entity = make_entity_id(entity_type, item)\n implied_entities.add(implied_entity)\n implied_relationships.add(\n make_relationship_id(\n implied_entity,\n implication_property.implied_relationship_name,\n document_entity_id,\n )\n )\n\n return KGImpliedExtractionResults(\n document_entity=document_entity_id,\n implied_entities=implied_entities,\n implied_relationships=implied_relationships,\n company_participant_emails=company_participant_emails,\n account_participant_emails=account_participant_emails,\n )\n\n\ndef kg_deep_extraction(\n document_id: str,\n metadata: KGEnhancedDocumentMetadata,\n implied_extraction: KGImpliedExtractionResults,\n tenant_id: str,\n index_name: str,\n kg_config_settings: KGConfigSettings,\n) -> KGDocumentDeepExtractionResults:\n \"\"\"\n Perform deep extraction and classification on the document.\n \"\"\"\n result = KGDocumentDeepExtractionResults(\n classification_result=None,\n deep_extracted_entities=set(),\n deep_extracted_relationships=set(),\n )\n\n entity_types_str = get_entity_types_str(active=True)\n relationship_types_str = get_relationship_types_str(active=True)\n\n for i, chunk_batch in enumerate(\n get_document_vespa_contents(document_id, index_name, tenant_id)\n ):\n # use first batch for classification\n if i == 0 and metadata.classification_enabled:\n if not metadata.classification_instructions:\n raise ValueError(\n \"Classification is enabled but no instructions are provided\"\n )\n result.classification_result = kg_classify_document(\n document_entity=implied_extraction.document_entity,\n chunk_batch=chunk_batch,\n implied_extraction=implied_extraction,\n classification_instructions=metadata.classification_instructions,\n kg_config_settings=kg_config_settings,\n )\n\n # deep extract from this chunk batch\n chunk_batch_results = kg_deep_extract_chunks(\n document_entity=implied_extraction.document_entity,\n chunk_batch=chunk_batch,\n implied_extraction=implied_extraction,\n kg_config_settings=kg_config_settings,\n entity_types_str=entity_types_str,\n relationship_types_str=relationship_types_str,\n )\n if chunk_batch_results is not None:\n result.deep_extracted_entities.update(\n chunk_batch_results.deep_extracted_entities\n )\n result.deep_extracted_relationships.update(\n chunk_batch_results.deep_extracted_relationships\n )\n\n return result\n\n\ndef kg_classify_document(\n document_entity: str,\n chunk_batch: list[KGChunkFormat],\n implied_extraction: KGImpliedExtractionResults,\n classification_instructions: KGClassificationInstructions,\n kg_config_settings: KGConfigSettings,\n) -> KGClassificationResult | None:\n # currently, classification is only done for calls\n # TODO: add support (or use same prompt and format) for non-call documents\n entity_type = get_entity_type(document_entity)\n if entity_type not in (call_type.value for call_type in OnyxCallTypes):\n return None\n\n # prepare prompt\n implied_extraction.document_entity\n company_participants = implied_extraction.company_participant_emails\n account_participants = implied_extraction.account_participant_emails\n content = (\n f\"Title: {chunk_batch[0].title}:\\nVendor Participants:\\n\"\n + \"\".join(f\" - {participant}\\n\" for participant in company_participants)\n + \"Other Participants:\\n\"\n + \"\".join(f\" - {participant}\\n\" for participant in account_participants)\n + \"Call Content:\\n\"\n + \"\\n\".join(chunk.content for chunk in chunk_batch)\n )\n category_list = {\n cls: definition.description\n for cls, definition in classification_instructions.classification_class_definitions.items()\n }\n prompt = CALL_DOCUMENT_CLASSIFICATION_PROMPT.format(\n beginning_of_call_content=content,\n category_list=category_list,\n category_options=classification_instructions.classification_options,\n vendor=kg_config_settings.KG_VENDOR,\n )\n\n # classify with LLM with Braintrust tracing\n llm = get_default_llm()\n try:\n prompt_msg = UserMessage(content=prompt)\n with llm_generation_span(\n llm=llm, flow=\"kg_document_classification\", input_messages=[prompt_msg]\n ) as span_generation:\n response = llm.invoke(prompt_msg)\n record_llm_response(span_generation, response)\n raw_classification_result = llm_response_to_string(response)\n\n classification_result = (\n raw_classification_result.replace(\"```json\", \"\").replace(\"```\", \"\").strip()\n )\n # no json parsing here because of reasoning output\n classification_class = classification_result.split(\"CATEGORY:\")[1].strip()\n\n if (\n classification_class\n in classification_instructions.classification_class_definitions\n ):\n return KGClassificationResult(\n document_entity=document_entity,\n classification_class=classification_class,\n )\n except Exception as e:\n logger.error(f\"Failed to classify document {document_entity}. Error: {str(e)}\")\n return None\n\n\ndef kg_deep_extract_chunks(\n document_entity: str,\n chunk_batch: list[KGChunkFormat],\n implied_extraction: KGImpliedExtractionResults,\n kg_config_settings: KGConfigSettings,\n entity_types_str: str,\n relationship_types_str: str,\n) -> KGDocumentDeepExtractionResults | None:\n # currently, calls are treated differently\n # TODO: either treat some other documents differently too, or ideally all the same way\n entity_type = get_entity_type(document_entity)\n is_call = entity_type in (call_type.value for call_type in OnyxCallTypes)\n\n content = \"\\n\".join(chunk.content for chunk in chunk_batch)\n\n # prepare prompt\n if is_call:\n company_participants_str = \"\".join(\n f\" - {participant}\\n\"\n for participant in implied_extraction.company_participant_emails\n )\n account_participants_str = \"\".join(\n f\" - {participant}\\n\"\n for participant in implied_extraction.account_participant_emails\n )\n llm_context = CALL_CHUNK_PREPROCESSING_PROMPT.format(\n participant_string=company_participants_str,\n account_participant_string=account_participants_str,\n vendor=kg_config_settings.KG_VENDOR,\n content=content,\n )\n else:\n llm_context = GENERAL_CHUNK_PREPROCESSING_PROMPT.format(\n vendor=kg_config_settings.KG_VENDOR,\n content=content,\n )\n prompt = MASTER_EXTRACTION_PROMPT.format(\n entity_types=entity_types_str,\n relationship_types=relationship_types_str,\n ).replace(\"---content---\", llm_context)\n\n # extract with LLM with Braintrust tracing\n llm = get_default_llm()\n try:\n prompt_msg = UserMessage(content=prompt)\n with llm_generation_span(\n llm=llm, flow=\"kg_deep_extraction\", input_messages=[prompt_msg]\n ) as span_generation:\n response = llm.invoke(prompt_msg)\n record_llm_response(span_generation, response)\n raw_extraction_result = llm_response_to_string(response)\n\n cleaned_response = (\n raw_extraction_result.replace(\"{{\", \"{\")\n .replace(\"}}\", \"}\")\n .replace(\"```json\\n\", \"\")\n .replace(\"\\n```\", \"\")\n .replace(\"\\n\", \"\")\n )\n first_bracket = cleaned_response.find(\"{\")\n last_bracket = cleaned_response.rfind(\"}\")\n cleaned_response = cleaned_response[first_bracket : last_bracket + 1]\n parsed_result = json.loads(cleaned_response)\n return KGDocumentDeepExtractionResults(\n classification_result=None,\n deep_extracted_entities=set(parsed_result.get(\"entities\", [])),\n deep_extracted_relationships={\n rel.replace(\" \", \"_\") for rel in parsed_result.get(\"relationships\", [])\n },\n )\n except Exception as e:\n failed_chunks = [chunk.chunk_id for chunk in chunk_batch]\n logger.error(\n f\"Failed to process chunks {failed_chunks} from document {document_entity}. Error: {str(e)}\"\n )\n return None\n\n\ndef kg_process_person(\n email: str,\n document_entity_id: str,\n relationship_type: str,\n kg_config_settings: KGConfigSettings,\n active_entity_types: set[str],\n) -> tuple[str, str, str, str] | None:\n \"\"\"\n Create an employee or account entity from an email address, and a relationship to\n the entity from the document that the email is from.\n\n Returns:\n tuple containing (person_entity, person_relationship, company_participant_email,\n and account_participant_email), or None if the created entity is not of an\n active entity type or is from an ignored email domain.\n \"\"\"\n kg_person = kg_email_processing(email, kg_config_settings)\n if any(\n domain.lower() in kg_person.company.lower()\n for domain in kg_config_settings.KG_IGNORE_EMAIL_DOMAINS\n ):\n return None\n\n person_entity = None\n if kg_person.employee and \"EMPLOYEE\" in active_entity_types:\n person_entity = make_entity_id(\"EMPLOYEE\", kg_person.name)\n elif not kg_person.employee and \"ACCOUNT\" in active_entity_types:\n person_entity = make_entity_id(\"ACCOUNT\", kg_person.company)\n\n if person_entity:\n is_account = person_entity.startswith(\"ACCOUNT\")\n participant_email = f\"{kg_person.name} -- ({kg_person.company})\"\n return (\n person_entity,\n make_relationship_id(person_entity, relationship_type, document_entity_id),\n participant_email if not is_account else \"\",\n participant_email if is_account else \"\",\n )\n\n return None\n\n\ndef get_batch_documents_metadata(\n document_ids: list[str], connector_source: str\n) -> list[KGMetadataContent]:\n \"\"\"\n Gets the metadata for a batch of documents.\n \"\"\"\n batch_metadata: list[KGMetadataContent] = []\n source_type = DocumentSource(connector_source).value\n\n with get_session_with_current_tenant() as db_session:\n for document_id in document_ids:\n # get document metadata\n metadata = get_structured_tags_for_document(document_id, db_session)\n\n batch_metadata.append(\n KGMetadataContent(\n document_id=document_id,\n source_type=source_type,\n source_metadata=metadata,\n )\n )\n return batch_metadata\n\n\ndef trackinfo_to_str(trackinfo: KGAttributeTrackInfo | None) -> str:\n \"\"\"Convert trackinfo to an LLM friendly string\"\"\"\n if trackinfo is None:\n return \"\"\n\n if trackinfo.type == KGAttributeTrackType.LIST:\n if trackinfo.values is None:\n return \"a list of any suitable values\"\n return \"a list with possible values: \" + \", \".join(trackinfo.values)\n elif trackinfo.type == KGAttributeTrackType.VALUE:\n if trackinfo.values is None:\n return \"any suitable value\"\n return \"one of: \" + \", \".join(trackinfo.values)\n\n\ndef trackinfo_to_dict(trackinfo: KGAttributeTrackInfo | None) -> dict | None:\n if trackinfo is None:\n return None\n return {\n \"type\": trackinfo.type,\n \"values\": (list(trackinfo.values) if trackinfo.values else None),\n }\n\n\nclass EntityTypeMetadataTracker:\n def __init__(self) -> None:\n \"\"\"\n Tracks the possible values the metadata attributes can take for each entity type.\n \"\"\"\n # entity type -> attribute -> trackinfo\n self.entity_attr_info: dict[str, dict[str, KGAttributeTrackInfo | None]] = {}\n self.entity_allowed_attrs: dict[str, set[str]] = {}\n\n def import_typeinfo(self) -> None:\n \"\"\"\n Loads the metadata tracking information from the database.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n entity_types = db_session.query(KGEntityType).all()\n\n for entity_type in entity_types:\n self.entity_attr_info[entity_type.id_name] = (\n entity_type.parsed_attributes.attribute_values\n )\n self.entity_allowed_attrs[entity_type.id_name] = {\n attr.name\n for attr in entity_type.parsed_attributes.metadata_attribute_conversion.values()\n }\n\n def export_typeinfo(self) -> None:\n \"\"\"\n Exports the metadata tracking information to the database.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n for entity_type_id_name, attribute_values in self.entity_attr_info.items():\n db_session.query(KGEntityType).filter(\n KGEntityType.id_name == entity_type_id_name\n ).update(\n {\n KGEntityType.attributes: KGEntityType.attributes.op(\"||\")(\n {\n \"attribute_values\": {\n attr: trackinfo_to_dict(info)\n for attr, info in attribute_values.items()\n }\n }\n )\n },\n synchronize_session=False,\n )\n db_session.commit()\n\n def track_metadata(\n self, entity_type: str, attributes: dict[str, str | list[str]]\n ) -> None:\n \"\"\"\n Tracks which values are possible for the given attributes.\n If the attribute value is a list, we track the values in the list rather than the list itself.\n If we see to many different values, we stop tracking the attribute.\n \"\"\"\n for attribute, value in attributes.items():\n # ignore types/metadata we are not tracking\n if entity_type not in self.entity_attr_info:\n continue\n if attribute not in self.entity_allowed_attrs[entity_type]:\n continue\n\n # determine if the attribute is a list or a value\n trackinfo = self.entity_attr_info[entity_type].get(attribute, None)\n if trackinfo is None:\n trackinfo = KGAttributeTrackInfo(\n type=(\n KGAttributeTrackType.VALUE\n if isinstance(value, str)\n else KGAttributeTrackType.LIST\n ),\n values=set(),\n )\n self.entity_attr_info[entity_type][attribute] = trackinfo\n\n # None means marked as don't track\n if trackinfo.values is None:\n continue\n\n # track the value\n if isinstance(value, str):\n trackinfo.values.add(value)\n else:\n trackinfo.type = KGAttributeTrackType.LIST\n trackinfo.values.update(value)\n\n # if we see to many different values, we stop tracking\n if len(trackinfo.values) > KG_METADATA_TRACKING_THRESHOLD:\n trackinfo.values = None\n" }, { "path": "backend/onyx/kg/utils/formatting_utils.py", "content": "import re\n\nfrom onyx.db.kg_config import KGConfigSettings\nfrom onyx.kg.models import KGPerson\n\n\ndef format_entity_id(entity_id_name: str) -> str:\n return make_entity_id(*split_entity_id(entity_id_name))\n\n\ndef make_entity_id(entity_type: str, entity_name: str) -> str:\n return f\"{entity_type.upper()}::{entity_name.lower()}\"\n\n\ndef split_entity_id(entity_id_name: str) -> list[str]:\n return entity_id_name.split(\"::\")\n\n\ndef get_entity_type(entity_id_name: str) -> str:\n return entity_id_name.split(\"::\", 1)[0].upper()\n\n\ndef format_entity_id_for_models(entity_id_name: str) -> str:\n entity_split = entity_id_name.split(\"::\")\n if len(entity_split) == 2:\n entity_type, entity_name = entity_split\n separator = \"::\"\n elif len(entity_split) > 2:\n raise ValueError(f\"Entity {entity_id_name} is not in the correct format\")\n else:\n entity_name = entity_id_name\n separator = entity_type = \"\"\n\n formatted_entity_type = entity_type.strip().upper()\n formatted_entity_name = entity_name.strip().replace('\"', \"\").replace(\"'\", \"\")\n\n return f\"{formatted_entity_type}{separator}{formatted_entity_name}\"\n\n\ndef get_attributes(entity_w_attributes: str) -> dict[str, str]:\n \"\"\"\n Extract attributes from an entity string.\n E.g., \"TYPE::Entity--[attr1: value1, attr2: value2]\" -> {\"attr1\": \"value1\", \"attr2\": \"value2\"}\n \"\"\"\n attr_split = entity_w_attributes.split(\"--\")\n if len(attr_split) != 2:\n raise ValueError(f\"Invalid entity with attributes: {entity_w_attributes}\")\n\n match = re.search(r\"\\[(.*)\\]\", attr_split[1])\n if not match:\n return {}\n\n attr_list_str = match.group(1)\n return {\n attr_split[0].strip(): attr_split[1].strip()\n for attr in attr_list_str.split(\",\")\n if len(attr_split := attr.split(\":\", 1)) == 2\n }\n\n\ndef make_entity_w_attributes(entity: str, attributes: dict[str, str]) -> str:\n return f\"{entity}--[{', '.join(f'{k}: {v}' for k, v in attributes.items())}]\"\n\n\ndef format_relationship_id(relationship_id_name: str) -> str:\n return make_relationship_id(*split_relationship_id(relationship_id_name))\n\n\ndef make_relationship_id(\n source_node: str, relationship_type: str, target_node: str\n) -> str:\n return f\"{format_entity_id(source_node)}__{relationship_type.lower()}__{format_entity_id(target_node)}\"\n\n\ndef split_relationship_id(relationship_id_name: str) -> list[str]:\n return relationship_id_name.split(\"__\")\n\n\ndef format_relationship_type_id(relationship_type_id_name: str) -> str:\n return make_relationship_type_id(\n *split_relationship_type_id(relationship_type_id_name)\n )\n\n\ndef make_relationship_type_id(\n source_node_type: str, relationship_type: str, target_node_type: str\n) -> str:\n return f\"{source_node_type.upper()}__{relationship_type.lower()}__{target_node_type.upper()}\"\n\n\ndef split_relationship_type_id(relationship_type_id_name: str) -> list[str]:\n return relationship_type_id_name.split(\"__\")\n\n\ndef extract_relationship_type_id(relationship_id_name: str) -> str:\n source_node, relationship_type, target_node = split_relationship_id(\n relationship_id_name\n )\n return make_relationship_type_id(\n get_entity_type(source_node), relationship_type, get_entity_type(target_node)\n )\n\n\ndef extract_email(email: str) -> str | None:\n \"\"\"\n Extract an email from an arbitrary string (if any).\n Only the first email is returned.\n \"\"\"\n match = re.search(r\"([A-Za-z0-9._+-]+@[A-Za-z0-9-]+(?:\\.[A-Za-z0-9-]+)+)\", email)\n return match.group(0) if match else None\n\n\ndef kg_email_processing(email: str, kg_config_settings: KGConfigSettings) -> KGPerson:\n \"\"\"\n Process the email.\n \"\"\"\n name, company_domain = email.split(\"@\")\n assert isinstance(company_domain, str)\n assert isinstance(kg_config_settings.KG_VENDOR_DOMAINS, list)\n assert isinstance(kg_config_settings.KG_VENDOR, str)\n\n employee = any(\n domain in company_domain for domain in kg_config_settings.KG_VENDOR_DOMAINS\n )\n if employee:\n company = kg_config_settings.KG_VENDOR\n else:\n # TODO: maybe store a list of domains for each account and use that to match\n # right now, gmail and other random domains are being converted into accounts\n company = company_domain.title()\n\n return KGPerson(name=name, company=company, employee=employee)\n" }, { "path": "backend/onyx/kg/utils/lock_utils.py", "content": "import time\n\nfrom redis.lock import Lock as RedisLock\n\n\ndef extend_lock(lock: RedisLock, timeout: int, last_lock_time: float) -> float:\n current_time = time.monotonic()\n if current_time - last_lock_time >= (timeout / 4):\n lock.reacquire()\n last_lock_time = current_time\n\n return last_lock_time\n" }, { "path": "backend/onyx/kg/vespa/vespa_interactions.py", "content": "import json\nfrom collections.abc import Generator\n\nfrom onyx.document_index.vespa.chunk_retrieval import get_chunks_via_visit_api\nfrom onyx.document_index.vespa.chunk_retrieval import VespaChunkRequest\nfrom onyx.document_index.vespa.index import IndexFilters\nfrom onyx.kg.models import KGChunkFormat\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef get_document_vespa_contents(\n document_id: str,\n index_name: str,\n tenant_id: str,\n batch_size: int = 8,\n) -> Generator[list[KGChunkFormat], None, None]:\n \"\"\"\n Retrieves chunks from Vespa for the given document IDs and converts them to KGChunks.\n\n Args:\n document_id (str): ID of the document to fetch chunks for\n index_name (str): Name of the Vespa index\n tenant_id (str): ID of the tenant\n batch_size (int): Number of chunks to fetch per batch\n\n Yields:\n list[KGChunk]: Batches of chunks ready for KG processing\n \"\"\"\n\n current_batch: list[KGChunkFormat] = []\n\n # get all chunks for the document\n # TODO: revisit the visit function\n chunks = get_chunks_via_visit_api(\n chunk_request=VespaChunkRequest(document_id=document_id),\n index_name=index_name,\n filters=IndexFilters(access_control_list=None, tenant_id=tenant_id),\n field_names=[\n \"document_id\",\n \"chunk_id\",\n \"title\",\n \"content\",\n \"metadata\",\n \"primary_owners\",\n \"secondary_owners\",\n \"source_type\",\n ],\n get_large_chunks=False,\n )\n\n # Convert Vespa chunks to KGChunks\n # kg_chunks: list[KGChunkFormat] = []\n\n for i, chunk in enumerate(chunks):\n fields = chunk[\"fields\"]\n if isinstance(fields.get(\"metadata\", {}), str):\n fields[\"metadata\"] = json.loads(fields[\"metadata\"])\n current_batch.append(\n KGChunkFormat(\n connector_id=None, # We may need to adjust this\n document_id=fields.get(\"document_id\"),\n chunk_id=fields.get(\"chunk_id\"),\n primary_owners=fields.get(\"primary_owners\", []),\n secondary_owners=fields.get(\"secondary_owners\", []),\n source_type=fields.get(\"source_type\", \"\"),\n title=fields.get(\"title\", \"\"),\n content=fields.get(\"content\", \"\"),\n metadata=fields.get(\"metadata\", {}),\n )\n )\n\n if len(current_batch) >= batch_size:\n yield current_batch\n current_batch = []\n\n # Yield any remaining chunks\n if current_batch:\n yield current_batch\n" }, { "path": "backend/onyx/llm/__init__.py", "content": "" }, { "path": "backend/onyx/llm/constants.py", "content": "\"\"\"\nLLM Constants\n\nCentralized constants for LLM providers, vendors, and display names.\n\"\"\"\n\nfrom enum import Enum\n\n\n# Provider names\nclass LlmProviderNames(str, Enum):\n \"\"\"\n Canonical string identifiers for LLM providers.\n \"\"\"\n\n OPENAI = \"openai\"\n ANTHROPIC = \"anthropic\"\n GOOGLE = \"google\"\n BEDROCK = \"bedrock\"\n BEDROCK_CONVERSE = \"bedrock_converse\"\n VERTEX_AI = \"vertex_ai\"\n OPENROUTER = \"openrouter\"\n AZURE = \"azure\"\n OLLAMA_CHAT = \"ollama_chat\"\n LM_STUDIO = \"lm_studio\"\n MISTRAL = \"mistral\"\n LITELLM_PROXY = \"litellm_proxy\"\n BIFROST = \"bifrost\"\n\n def __str__(self) -> str:\n \"\"\"Needed so things like:\n\n f\"{LlmProviderNames.OPENAI}/\" gives back \"openai/\" instead of \"LlmProviderNames.OPENAI/\"\n \"\"\"\n return self.value\n\n\nWELL_KNOWN_PROVIDER_NAMES = [\n LlmProviderNames.OPENAI,\n LlmProviderNames.ANTHROPIC,\n LlmProviderNames.VERTEX_AI,\n LlmProviderNames.BEDROCK,\n LlmProviderNames.OPENROUTER,\n LlmProviderNames.AZURE,\n LlmProviderNames.OLLAMA_CHAT,\n LlmProviderNames.LM_STUDIO,\n LlmProviderNames.LITELLM_PROXY,\n LlmProviderNames.BIFROST,\n]\n\n\n# Proper capitalization for known providers and vendors\nPROVIDER_DISPLAY_NAMES: dict[str, str] = {\n LlmProviderNames.OPENAI: \"OpenAI\",\n LlmProviderNames.ANTHROPIC: \"Anthropic\",\n LlmProviderNames.GOOGLE: \"Google\",\n LlmProviderNames.BEDROCK: \"Bedrock\",\n LlmProviderNames.BEDROCK_CONVERSE: \"Bedrock\",\n LlmProviderNames.VERTEX_AI: \"Vertex AI\",\n LlmProviderNames.OPENROUTER: \"OpenRouter\",\n LlmProviderNames.AZURE: \"Azure\",\n \"ollama\": \"Ollama\",\n LlmProviderNames.OLLAMA_CHAT: \"Ollama\",\n LlmProviderNames.LM_STUDIO: \"LM Studio\",\n LlmProviderNames.LITELLM_PROXY: \"LiteLLM Proxy\",\n LlmProviderNames.BIFROST: \"Bifrost\",\n \"groq\": \"Groq\",\n \"anyscale\": \"Anyscale\",\n \"deepseek\": \"DeepSeek\",\n \"xai\": \"xAI\",\n LlmProviderNames.MISTRAL: \"Mistral\",\n \"mistralai\": \"Mistral\", # Alias used by some providers\n \"cohere\": \"Cohere\",\n \"perplexity\": \"Perplexity\",\n \"amazon\": \"Amazon\",\n \"meta\": \"Meta\",\n \"meta-llama\": \"Meta\", # Alias used by some providers\n \"ai21\": \"AI21\",\n \"nvidia\": \"NVIDIA\",\n \"databricks\": \"Databricks\",\n \"alibaba\": \"Alibaba\",\n \"qwen\": \"Qwen\",\n \"microsoft\": \"Microsoft\",\n \"gemini\": \"Gemini\",\n \"stability\": \"Stability\",\n \"writer\": \"Writer\",\n}\n\n# Map vendors to their brand names (used for provider_display_name generation)\nVENDOR_BRAND_NAMES: dict[str, str] = {\n \"anthropic\": \"Claude\",\n \"openai\": \"GPT\",\n \"google\": \"Gemini\",\n \"amazon\": \"Nova\",\n \"meta\": \"Llama\",\n \"mistral\": \"Mistral\",\n \"cohere\": \"Command\",\n \"deepseek\": \"DeepSeek\",\n \"xai\": \"Grok\",\n \"perplexity\": \"Sonar\",\n \"ai21\": \"Jamba\",\n \"nvidia\": \"Nemotron\",\n \"qwen\": \"Qwen\",\n \"alibaba\": \"Qwen\",\n \"writer\": \"Palmyra\",\n}\n\n# Aggregator providers that host models from multiple vendors\nAGGREGATOR_PROVIDERS: set[str] = {\n LlmProviderNames.BEDROCK,\n LlmProviderNames.BEDROCK_CONVERSE,\n LlmProviderNames.OPENROUTER,\n LlmProviderNames.OLLAMA_CHAT,\n LlmProviderNames.LM_STUDIO,\n LlmProviderNames.VERTEX_AI,\n LlmProviderNames.AZURE,\n LlmProviderNames.LITELLM_PROXY,\n LlmProviderNames.BIFROST,\n}\n\n# Model family name mappings for display name generation\n# Used by Bedrock display name generator\nBEDROCK_MODEL_NAME_MAPPINGS: dict[str, str] = {\n \"claude\": \"Claude\",\n \"llama\": \"Llama\",\n \"mistral\": \"Mistral\",\n \"mixtral\": \"Mixtral\",\n \"titan\": \"Titan\",\n \"nova\": \"Nova\",\n \"jamba\": \"Jamba\",\n \"command\": \"Command\",\n \"deepseek\": \"DeepSeek\",\n}\n\n# Used by Ollama display name generator\nOLLAMA_MODEL_NAME_MAPPINGS: dict[str, str] = {\n \"llama\": \"Llama\",\n \"qwen\": \"Qwen\",\n \"mistral\": \"Mistral\",\n \"deepseek\": \"DeepSeek\",\n \"gemma\": \"Gemma\",\n \"phi\": \"Phi\",\n \"codellama\": \"Code Llama\",\n \"starcoder\": \"StarCoder\",\n \"wizardcoder\": \"WizardCoder\",\n \"vicuna\": \"Vicuna\",\n \"orca\": \"Orca\",\n \"dolphin\": \"Dolphin\",\n \"nous\": \"Nous\",\n \"neural\": \"Neural\",\n \"mixtral\": \"Mixtral\",\n \"falcon\": \"Falcon\",\n \"yi\": \"Yi\",\n \"command\": \"Command\",\n \"zephyr\": \"Zephyr\",\n \"openchat\": \"OpenChat\",\n \"solar\": \"Solar\",\n}\n\n# Bedrock model token limits (AWS doesn't expose this via API)\n# Note: Many Bedrock model IDs include context length suffix (e.g., \":200k\")\n# which is parsed first. This mapping is for models without suffixes.\n# Sources:\n# - LiteLLM model_prices_and_context_window.json\n# - AWS Bedrock documentation and announcement blogs\nBEDROCK_MODEL_TOKEN_LIMITS: dict[str, int] = {\n # Anthropic Claude models (new naming: claude-{tier}-{version})\n \"claude-opus-4\": 200000,\n \"claude-sonnet-4\": 200000,\n \"claude-haiku-4\": 200000,\n # Anthropic Claude models (old naming: claude-{version})\n \"claude-4\": 200000,\n \"claude-3-7\": 200000,\n \"claude-3-5\": 200000,\n \"claude-3\": 200000,\n \"claude-v2\": 100000,\n \"claude-instant\": 100000,\n # Amazon Nova models (from LiteLLM)\n \"nova-premier\": 1000000,\n \"nova-pro\": 300000,\n \"nova-lite\": 300000,\n \"nova-2-lite\": 1000000, # Nova 2 Lite has 1M context\n \"nova-2-sonic\": 128000,\n \"nova-micro\": 128000,\n # Amazon Titan models (from LiteLLM: all text models are 42K)\n \"titan-text-premier\": 42000,\n \"titan-text-express\": 42000,\n \"titan-text-lite\": 42000,\n \"titan-tg1\": 8000,\n # Meta Llama models (Llama 3 base = 8K, Llama 3.1+ = 128K)\n \"llama4\": 128000,\n \"llama3-3\": 128000,\n \"llama3-2\": 128000,\n \"llama3-1\": 128000,\n \"llama3-8b\": 8000,\n \"llama3-70b\": 8000,\n # Mistral models (Large 2+ = 128K, original Large/Small = 32K)\n \"mistral-large-3\": 128000,\n \"mistral-large-2407\": 128000, # Mistral Large 2\n \"mistral-large-2402\": 32000, # Original Mistral Large\n \"mistral-large\": 128000, # Default to newer version\n \"mistral-small\": 32000,\n \"mistral-7b\": 32000,\n \"mixtral-8x7b\": 32000,\n \"pixtral\": 128000,\n \"ministral\": 128000,\n \"magistral\": 128000,\n \"voxtral\": 32000,\n # Cohere models\n \"command-r-plus\": 128000,\n \"command-r\": 128000,\n # DeepSeek models\n \"deepseek\": 64000,\n # Google Gemma models\n \"gemma-3\": 128000,\n \"gemma-2\": 8000,\n \"gemma\": 8000,\n # Qwen models\n \"qwen3\": 128000,\n \"qwen2\": 128000,\n # NVIDIA models\n \"nemotron\": 128000,\n # Writer Palmyra models\n \"palmyra\": 128000,\n # Moonshot Kimi\n \"kimi\": 128000,\n # Minimax\n \"minimax\": 128000,\n # OpenAI (via Bedrock)\n \"gpt-oss\": 128000,\n # AI21 models (from LiteLLM: Jamba 1.5 = 256K, Jamba Instruct = 70K)\n \"jamba-1-5\": 256000,\n \"jamba-instruct\": 70000,\n \"jamba\": 256000, # Default to newer version\n}\n\n\n# Models that should keep their hyphenated format in display names\n# These are model families where the hyphen is part of the brand name\nHYPHENATED_MODEL_NAMES: set[str] = {\n \"gpt-oss\",\n}\n\n\n# General model prefix to vendor mapping (used as fallback when enrichment data is missing)\n# This covers common model families across all providers\nMODEL_PREFIX_TO_VENDOR: dict[str, str] = {\n # Google\n \"gemini\": \"google\",\n \"gemma\": \"google\",\n \"palm\": \"google\",\n # Anthropic\n \"claude\": \"anthropic\",\n # OpenAI\n \"gpt\": \"openai\",\n \"o1\": \"openai\",\n \"o3\": \"openai\",\n \"o4\": \"openai\",\n \"chatgpt\": \"openai\",\n # Meta\n \"llama\": \"meta\",\n \"codellama\": \"meta\",\n # Mistral\n \"mistral\": \"mistral\",\n \"mixtral\": \"mistral\",\n \"codestral\": \"mistral\",\n \"ministral\": \"mistral\",\n \"pixtral\": \"mistral\",\n \"magistral\": \"mistral\",\n # Cohere\n \"command\": \"cohere\",\n \"aya\": \"cohere\",\n # Amazon\n \"nova\": \"amazon\",\n \"titan\": \"amazon\",\n # AI21\n \"jamba\": \"ai21\",\n # DeepSeek\n \"deepseek\": \"deepseek\",\n # Alibaba/Qwen\n \"qwen\": \"alibaba\",\n \"qwq\": \"alibaba\",\n # Microsoft\n \"phi\": \"microsoft\",\n # NVIDIA\n \"nemotron\": \"nvidia\",\n # xAI\n \"grok\": \"xai\",\n}\n\n\n# Ollama model prefix to vendor mapping (for grouping models by vendor)\nOLLAMA_MODEL_TO_VENDOR: dict[str, str] = {\n \"llama\": \"Meta\",\n \"codellama\": \"Meta\",\n \"qwen\": \"Alibaba\",\n \"qwq\": \"Alibaba\",\n \"mistral\": \"Mistral\",\n \"ministral\": \"Mistral\",\n \"mixtral\": \"Mistral\",\n \"deepseek\": \"DeepSeek\",\n \"gemma\": \"Google\",\n \"phi\": \"Microsoft\",\n \"command\": \"Cohere\",\n \"aya\": \"Cohere\",\n \"falcon\": \"TII\",\n \"yi\": \"01.AI\",\n \"starcoder\": \"BigCode\",\n \"wizardcoder\": \"WizardLM\",\n \"vicuna\": \"LMSYS\",\n \"openchat\": \"OpenChat\",\n \"solar\": \"Upstage\",\n \"orca\": \"Microsoft\",\n \"dolphin\": \"Cognitive Computations\",\n \"nous\": \"Nous Research\",\n \"neural\": \"Intel\",\n \"zephyr\": \"HuggingFace\",\n \"granite\": \"IBM\",\n \"nemotron\": \"NVIDIA\",\n \"smollm\": \"HuggingFace\",\n}\n" }, { "path": "backend/onyx/llm/cost.py", "content": "\"\"\"LLM cost calculation utilities.\"\"\"\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef calculate_llm_cost_cents(\n model_name: str,\n prompt_tokens: int,\n completion_tokens: int,\n) -> float:\n \"\"\"\n Calculate the cost in cents for an LLM API call.\n\n Uses litellm's cost_per_token function to get current pricing.\n Returns 0 if the model is not found or on any error.\n \"\"\"\n try:\n import litellm\n\n # cost_per_token returns (prompt_cost, completion_cost) in USD\n prompt_cost_usd, completion_cost_usd = litellm.cost_per_token(\n model=model_name,\n prompt_tokens=prompt_tokens,\n completion_tokens=completion_tokens,\n )\n\n # Convert to cents (multiply by 100)\n total_cost_cents = (prompt_cost_usd + completion_cost_usd) * 100\n return total_cost_cents\n\n except Exception as e:\n # Log but don't fail - unknown models or errors shouldn't block usage\n logger.debug(\n f\"Could not calculate cost for model {model_name}: {e}. Assuming cost is 0.\"\n )\n return 0.0\n" }, { "path": "backend/onyx/llm/factory.py", "content": "from collections.abc import Callable\nfrom typing import Any\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.model_configs import GEN_AI_TEMPERATURE\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import LLMModelFlowType\nfrom onyx.db.llm import can_user_access_llm_provider\nfrom onyx.db.llm import fetch_default_llm_model\nfrom onyx.db.llm import fetch_default_vision_model\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import fetch_existing_models\nfrom onyx.db.llm import fetch_llm_provider_view\nfrom onyx.db.llm import fetch_user_group_ids\nfrom onyx.db.models import Persona\nfrom onyx.db.models import User\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.multi_llm import LitellmLLM\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.llm.utils import get_max_input_tokens_from_llm_provider\nfrom onyx.llm.utils import model_supports_image_input\nfrom onyx.llm.well_known_providers.constants import (\n PROVIDERS_WITH_SPECIAL_API_KEY_HANDLING,\n)\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.utils.headers import build_llm_extra_headers\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _build_provider_extra_headers(\n provider: str, custom_config: dict[str, str] | None\n) -> dict[str, str]:\n if provider in PROVIDERS_WITH_SPECIAL_API_KEY_HANDLING and custom_config:\n raw = custom_config.get(PROVIDERS_WITH_SPECIAL_API_KEY_HANDLING[provider])\n api_key = raw.strip() if raw else None\n if not api_key:\n return {}\n return {\n \"Authorization\": (\n api_key\n if api_key.lower().startswith(\"bearer \")\n else f\"Bearer {api_key}\"\n )\n }\n\n # Passing these will put Onyx on the OpenRouter leaderboard\n elif provider == LlmProviderNames.OPENROUTER:\n return {\n \"HTTP-Referer\": \"https://onyx.app\",\n \"X-Title\": \"Onyx\",\n }\n\n return {}\n\n\ndef _get_model_configured_max_input_tokens(\n llm_provider: LLMProviderView,\n model_name: str,\n) -> int | None:\n for model_configuration in llm_provider.model_configurations:\n if model_configuration.name == model_name:\n return model_configuration.max_input_tokens\n return None\n\n\ndef _build_model_kwargs(\n provider: str,\n configured_max_input_tokens: int | None,\n) -> dict[str, Any]:\n model_kwargs: dict[str, Any] = {}\n if (\n provider == LlmProviderNames.OLLAMA_CHAT\n and configured_max_input_tokens\n and configured_max_input_tokens > 0\n ):\n model_kwargs[\"num_ctx\"] = configured_max_input_tokens\n return model_kwargs\n\n\ndef get_llm_for_persona(\n persona: Persona | None,\n user: User,\n llm_override: LLMOverride | None = None,\n additional_headers: dict[str, str] | None = None,\n) -> LLM:\n if persona is None:\n logger.warning(\"No persona provided, using default LLM\")\n return get_default_llm()\n\n provider_name_override = llm_override.model_provider if llm_override else None\n model_version_override = llm_override.model_version if llm_override else None\n temperature_override = llm_override.temperature if llm_override else None\n\n provider_name = provider_name_override or persona.llm_model_provider_override\n if not provider_name:\n return get_default_llm(\n temperature=temperature_override or GEN_AI_TEMPERATURE,\n additional_headers=additional_headers,\n )\n\n with get_session_with_current_tenant() as db_session:\n provider_model = fetch_existing_llm_provider(provider_name, db_session)\n if not provider_model:\n raise ValueError(\"No LLM provider found\")\n\n # Fetch user group IDs for access control check\n user_group_ids = fetch_user_group_ids(db_session, user)\n\n if not can_user_access_llm_provider(\n provider_model, user_group_ids, persona, user.role == UserRole.ADMIN\n ):\n logger.warning(\n \"User %s with persona %s cannot access provider %s. Falling back to default provider.\",\n user.id,\n persona.id,\n provider_model.name,\n )\n return get_default_llm(\n temperature=temperature_override or GEN_AI_TEMPERATURE,\n additional_headers=additional_headers,\n )\n\n llm_provider = LLMProviderView.from_model(provider_model)\n\n model = model_version_override or persona.llm_model_version_override\n if not model:\n raise ValueError(\"No model name found\")\n\n return llm_from_provider(\n model_name=model,\n llm_provider=llm_provider,\n temperature=temperature_override,\n additional_headers=additional_headers,\n )\n\n\ndef get_default_llm_with_vision(\n timeout: int | None = None,\n temperature: float | None = None,\n additional_headers: dict[str, str] | None = None,\n) -> LLM | None:\n \"\"\"Get an LLM that supports image input, with the following priority:\n 1. Use the designated default vision provider if it exists and supports image input\n 2. Fall back to the first LLM provider that supports image input\n\n Returns None if no providers exist or if no provider supports images.\n \"\"\"\n\n def create_vision_llm(provider: LLMProviderView, model: str) -> LLM:\n \"\"\"Helper to create an LLM if the provider supports image input.\"\"\"\n return llm_from_provider(\n model_name=model,\n llm_provider=provider,\n timeout=timeout,\n temperature=temperature,\n additional_headers=additional_headers,\n )\n\n provider_map = {}\n with get_session_with_current_tenant() as db_session:\n # Try the default vision provider first\n default_model = fetch_default_vision_model(db_session)\n if default_model:\n if model_supports_image_input(\n default_model.name, default_model.llm_provider.provider\n ):\n logger.info(\n \"Using default vision model: %s (provider=%s)\",\n default_model.name,\n default_model.llm_provider.provider,\n )\n return create_vision_llm(\n LLMProviderView.from_model(default_model.llm_provider),\n default_model.name,\n )\n else:\n logger.warning(\n \"Default vision model %s (provider=%s) does not support \"\n \"image input — falling back to searching all providers\",\n default_model.name,\n default_model.llm_provider.provider,\n )\n\n # Fall back to searching all providers\n models = fetch_existing_models(\n db_session=db_session,\n flow_types=[LLMModelFlowType.VISION, LLMModelFlowType.CHAT],\n )\n\n if not models:\n logger.warning(\n \"No LLM models with VISION or CHAT flow type found — \"\n \"image summarization will be disabled\"\n )\n return None\n\n for model in models:\n if model.llm_provider_id not in provider_map:\n provider_map[model.llm_provider_id] = LLMProviderView.from_model(\n model.llm_provider\n )\n\n # Search for viable vision model followed by chat models\n # Sort models from VISION to CHAT priority\n sorted_models = sorted(\n models,\n key=lambda x: (\n LLMModelFlowType.VISION in x.llm_model_flow_types,\n LLMModelFlowType.CHAT in x.llm_model_flow_types,\n ),\n reverse=True,\n )\n\n for model in sorted_models:\n if model_supports_image_input(model.name, model.llm_provider.provider):\n logger.info(\n \"Using fallback vision model: %s (provider=%s)\",\n model.name,\n model.llm_provider.provider,\n )\n return create_vision_llm(\n provider_map[model.llm_provider_id],\n model.name,\n )\n\n checked_models = [\n f\"{m.name} (provider={m.llm_provider.provider})\" for m in sorted_models\n ]\n logger.warning(\n \"No vision-capable model found among %d candidates: %s — \"\n \"image summarization will be disabled\",\n len(sorted_models),\n \", \".join(checked_models),\n )\n return None\n\n\ndef llm_from_provider(\n model_name: str,\n llm_provider: LLMProviderView,\n timeout: int | None = None,\n temperature: float | None = None,\n additional_headers: dict[str, str] | None = None,\n) -> LLM:\n configured_max_input_tokens = _get_model_configured_max_input_tokens(\n llm_provider=llm_provider, model_name=model_name\n )\n model_kwargs = _build_model_kwargs(\n provider=llm_provider.provider,\n configured_max_input_tokens=configured_max_input_tokens,\n )\n max_input_tokens = (\n configured_max_input_tokens\n if configured_max_input_tokens\n else get_max_input_tokens_from_llm_provider(\n llm_provider=llm_provider, model_name=model_name\n )\n )\n return get_llm(\n provider=llm_provider.provider,\n model=model_name,\n deployment_name=llm_provider.deployment_name,\n api_key=llm_provider.api_key,\n api_base=llm_provider.api_base,\n api_version=llm_provider.api_version,\n custom_config=llm_provider.custom_config,\n timeout=timeout,\n temperature=temperature,\n additional_headers=additional_headers,\n max_input_tokens=max_input_tokens,\n model_kwargs=model_kwargs,\n )\n\n\ndef get_llm_for_contextual_rag(model_name: str, model_provider: str) -> LLM:\n with get_session_with_current_tenant() as db_session:\n llm_provider = fetch_llm_provider_view(db_session, model_provider)\n if not llm_provider:\n raise ValueError(\"No LLM provider with name {} found\".format(model_provider))\n return llm_from_provider(\n model_name=model_name,\n llm_provider=llm_provider,\n )\n\n\ndef get_default_llm(\n timeout: int | None = None,\n temperature: float | None = None,\n additional_headers: dict[str, str] | None = None,\n) -> LLM:\n with get_session_with_current_tenant() as db_session:\n model = fetch_default_llm_model(db_session)\n\n if not model:\n raise ValueError(\"No default LLM model found\")\n\n return llm_from_provider(\n model_name=model.name,\n llm_provider=LLMProviderView.from_model(model.llm_provider),\n timeout=timeout,\n temperature=temperature,\n additional_headers=additional_headers,\n )\n\n\ndef get_llm(\n provider: str,\n model: str,\n max_input_tokens: int,\n deployment_name: str | None,\n api_key: str | None = None,\n api_base: str | None = None,\n api_version: str | None = None,\n custom_config: dict[str, str] | None = None,\n temperature: float | None = None,\n timeout: int | None = None,\n additional_headers: dict[str, str] | None = None,\n model_kwargs: dict[str, Any] | None = None,\n) -> LLM:\n if temperature is None:\n temperature = GEN_AI_TEMPERATURE\n\n extra_headers = build_llm_extra_headers(additional_headers)\n\n # NOTE: this is needed since Ollama API key is optional\n # User may access Ollama cloud via locally hosted instance (logged in)\n # or just via the cloud API (not logged in, using API key)\n provider_extra_headers = _build_provider_extra_headers(provider, custom_config)\n if provider_extra_headers:\n extra_headers.update(provider_extra_headers)\n\n return LitellmLLM(\n model_provider=provider,\n model_name=model,\n deployment_name=deployment_name,\n api_key=api_key,\n api_base=api_base,\n api_version=api_version,\n timeout=timeout,\n temperature=temperature,\n custom_config=custom_config,\n extra_headers=extra_headers,\n model_kwargs=model_kwargs or {},\n max_input_tokens=max_input_tokens,\n )\n\n\ndef get_llm_tokenizer_encode_func(llm: LLM) -> Callable[[str], list[int]]:\n \"\"\"Get the tokenizer encode function for an LLM.\n\n Args:\n llm: The LLM instance to get the tokenizer for\n\n Returns:\n A callable that encodes a string into a list of token IDs\n \"\"\"\n llm_provider = llm.config.model_provider\n llm_model_name = llm.config.model_name\n\n llm_tokenizer = get_tokenizer(\n model_name=llm_model_name,\n provider_type=llm_provider,\n )\n return llm_tokenizer.encode\n\n\ndef get_llm_token_counter(llm: LLM) -> Callable[[str], int]:\n tokenizer_encode_func = get_llm_tokenizer_encode_func(llm)\n return lambda text: len(tokenizer_encode_func(text))\n" }, { "path": "backend/onyx/llm/interfaces.py", "content": "import abc\nfrom collections.abc import Iterator\n\nfrom braintrust import traced\nfrom pydantic import BaseModel\n\nfrom onyx.llm.model_response import ModelResponse\nfrom onyx.llm.model_response import ModelResponseStream\nfrom onyx.llm.models import LanguageModelInput\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import ToolChoiceOptions\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass LLMUserIdentity(BaseModel):\n user_id: str | None = None\n session_id: str | None = None\n\n\nclass LLMConfig(BaseModel):\n model_provider: str\n model_name: str\n temperature: float\n api_key: str | None = None\n api_base: str | None = None\n api_version: str | None = None\n deployment_name: str | None = None\n custom_config: dict[str, str] | None = None\n max_input_tokens: int\n # This disables the \"model_\" protected namespace for pydantic\n model_config = {\"protected_namespaces\": ()}\n\n\nclass LLM(abc.ABC):\n @property\n @abc.abstractmethod\n def config(self) -> LLMConfig:\n raise NotImplementedError\n\n @traced(name=\"invoke llm\", type=\"llm\")\n def invoke(\n self,\n prompt: LanguageModelInput,\n tools: list[dict] | None = None,\n tool_choice: ToolChoiceOptions | None = None,\n structured_response_format: dict | None = None,\n timeout_override: int | None = None,\n max_tokens: int | None = None,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n user_identity: LLMUserIdentity | None = None,\n ) -> \"ModelResponse\":\n raise NotImplementedError\n\n def stream(\n self,\n prompt: LanguageModelInput,\n tools: list[dict] | None = None,\n tool_choice: ToolChoiceOptions | None = None,\n structured_response_format: dict | None = None,\n timeout_override: int | None = None,\n max_tokens: int | None = None,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n user_identity: LLMUserIdentity | None = None,\n ) -> Iterator[ModelResponseStream]:\n raise NotImplementedError\n" }, { "path": "backend/onyx/llm/litellm_singleton/__init__.py", "content": "\"\"\"\nSingleton module for litellm configuration.\nThis ensures litellm is configured exactly once when first imported.\nAll other modules should import litellm from here instead of directly.\n\"\"\"\n\nimport litellm\n\nfrom .config import initialize_litellm\nfrom .monkey_patches import apply_monkey_patches\n\ninitialize_litellm()\napply_monkey_patches()\n\n# Export the configured litellm module and model\n__all__ = [\"litellm\"]\n" }, { "path": "backend/onyx/llm/litellm_singleton/config.py", "content": "import json\nfrom pathlib import Path\n\nimport litellm\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef configure_litellm_settings() -> None:\n # If a user configures a different model and it doesn't support all the same\n # parameters like frequency and presence, just ignore them\n litellm.drop_params = True\n litellm.telemetry = False\n litellm.modify_params = True\n litellm.add_function_to_prompt = False\n litellm.suppress_debug_info = True\n\n\n# TODO: We might not need to register ollama_chat in addition to ollama but let's just do it for good measure for now.\ndef register_ollama_models() -> None:\n litellm.register_model(\n model_cost={\n # GPT-OSS models\n \"ollama_chat/gpt-oss:120b-cloud\": {\"supports_function_calling\": True},\n \"ollama_chat/gpt-oss:120b\": {\"supports_function_calling\": True},\n \"ollama_chat/gpt-oss:20b-cloud\": {\"supports_function_calling\": True},\n \"ollama_chat/gpt-oss:20b\": {\"supports_function_calling\": True},\n \"ollama/gpt-oss:120b-cloud\": {\"supports_function_calling\": True},\n \"ollama/gpt-oss:120b\": {\"supports_function_calling\": True},\n \"ollama/gpt-oss:20b-cloud\": {\"supports_function_calling\": True},\n \"ollama/gpt-oss:20b\": {\"supports_function_calling\": True},\n # DeepSeek models\n \"ollama_chat/deepseek-r1:latest\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:1.5b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:7b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:8b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:14b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:32b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:70b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-r1:671b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-v3.1:latest\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-v3.1:671b\": {\"supports_function_calling\": True},\n \"ollama_chat/deepseek-v3.1:671b-cloud\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:latest\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:1.5b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:7b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:8b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:14b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:32b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:70b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-r1:671b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-v3.1:latest\": {\"supports_function_calling\": True},\n \"ollama/deepseek-v3.1:671b\": {\"supports_function_calling\": True},\n \"ollama/deepseek-v3.1:671b-cloud\": {\"supports_function_calling\": True},\n # Gemma3 models\n \"ollama_chat/gemma3:latest\": {\"supports_function_calling\": True},\n \"ollama_chat/gemma3:270m\": {\"supports_function_calling\": True},\n \"ollama_chat/gemma3:1b\": {\"supports_function_calling\": True},\n \"ollama_chat/gemma3:4b\": {\"supports_function_calling\": True},\n \"ollama_chat/gemma3:12b\": {\"supports_function_calling\": True},\n \"ollama_chat/gemma3:27b\": {\"supports_function_calling\": True},\n \"ollama/gemma3:latest\": {\"supports_function_calling\": True},\n \"ollama/gemma3:270m\": {\"supports_function_calling\": True},\n \"ollama/gemma3:1b\": {\"supports_function_calling\": True},\n \"ollama/gemma3:4b\": {\"supports_function_calling\": True},\n \"ollama/gemma3:12b\": {\"supports_function_calling\": True},\n \"ollama/gemma3:27b\": {\"supports_function_calling\": True},\n # Qwen models\n \"ollama_chat/qwen3-coder:latest\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-coder:30b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-coder:480b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-coder:480b-cloud\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:latest\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:2b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:4b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:8b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:30b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:32b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:235b\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:235b-cloud\": {\"supports_function_calling\": True},\n \"ollama_chat/qwen3-vl:235b-instruct-cloud\": {\n \"supports_function_calling\": True\n },\n \"ollama/qwen3-coder:latest\": {\"supports_function_calling\": True},\n \"ollama/qwen3-coder:30b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-coder:480b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-coder:480b-cloud\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:latest\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:2b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:4b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:8b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:30b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:32b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:235b\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:235b-cloud\": {\"supports_function_calling\": True},\n \"ollama/qwen3-vl:235b-instruct-cloud\": {\"supports_function_calling\": True},\n # Kimi\n \"ollama_chat/kimi-k2:1t\": {\"supports_function_calling\": True},\n \"ollama_chat/kimi-k2:1t-cloud\": {\"supports_function_calling\": True},\n \"ollama/kimi-k2:1t\": {\"supports_function_calling\": True},\n \"ollama/kimi-k2:1t-cloud\": {\"supports_function_calling\": True},\n # GLM\n \"ollama_chat/glm-4.6:cloud\": {\"supports_function_calling\": True},\n \"ollama_chat/glm-4.6\": {\"supports_function_calling\": True},\n \"ollama/glm-4.6\": {\"supports_function_calling\": True},\n \"ollama/glm-4.6-cloud\": {\"supports_function_calling\": True},\n }\n )\n\n\ndef load_model_metadata_enrichments() -> None:\n \"\"\"\n Load model metadata enrichments from JSON file and merge into litellm.model_cost.\n\n This adds model_vendor, display_name, and model_version fields\n to litellm's model_cost dict. These fields are used by the UI to display\n models grouped by vendor with human-friendly names.\n\n Once LiteLLM accepts our upstream PR to add these fields natively,\n this function and the JSON file can be removed.\n \"\"\"\n enrichments_path = Path(__file__).parent.parent / \"model_metadata_enrichments.json\"\n\n if not enrichments_path.exists():\n logger.warning(f\"Model metadata enrichments file not found: {enrichments_path}\")\n return\n\n try:\n with open(enrichments_path) as f:\n enrichments = json.load(f)\n\n # Merge enrichments into litellm.model_cost\n for model_key, metadata in enrichments.items():\n if model_key in litellm.model_cost:\n # Update existing entry with our metadata\n litellm.model_cost[model_key].update(metadata)\n else:\n # Model not in litellm.model_cost - add it with just our metadata\n litellm.model_cost[model_key] = metadata\n\n logger.info(f\"Loaded model metadata enrichments for {len(enrichments)} models\")\n\n # Clear the model name parser cache since enrichments are now loaded\n # This ensures any parsing done before enrichments were loaded gets refreshed\n try:\n from onyx.llm.model_name_parser import parse_litellm_model_name\n\n parse_litellm_model_name.cache_clear()\n except ImportError:\n pass # Parser not yet imported, no cache to clear\n except Exception as e:\n logger.error(f\"Failed to load model metadata enrichments: {e}\")\n\n\ndef initialize_litellm() -> None:\n configure_litellm_settings()\n register_ollama_models()\n load_model_metadata_enrichments()\n" }, { "path": "backend/onyx/llm/litellm_singleton/monkey_patches.py", "content": "\"\"\"\nLiteLLM Monkey Patches\n\nThis module addresses the following issues in LiteLLM:\n\nStatus checked against LiteLLM v1.81.6-nightly (2026-02-02):\n\n1. Ollama Streaming Reasoning Content (_patch_ollama_chunk_parser):\n - LiteLLM's chunk_parser doesn't properly handle reasoning content in streaming\n responses from Ollama\n - Processes native \"thinking\" field from Ollama responses\n - Also handles ... tags in content for models that use that format\n - Tracks reasoning state to properly separate thinking from regular content\n STATUS: STILL NEEDED - LiteLLM has a bug where it only yields thinking content on\n the first two chunks, then stops (lines 504-510). Our patch correctly yields\n ALL thinking chunks. The upstream logic sets finished_reasoning_content=True\n on the second chunk instead of when regular content starts.\n\n2. OpenAI Responses API Parallel Tool Calls (_patch_openai_responses_parallel_tool_calls):\n - LiteLLM's translate_responses_chunk_to_openai_stream hardcodes index=0 for all tool calls\n - This breaks parallel tool calls where multiple functions are called simultaneously\n - The OpenAI Responses API provides output_index in streaming events to track which\n tool call each event belongs to\n STATUS: STILL NEEDED - LiteLLM hardcodes index=0 in translate_responses_chunk_to_openai_stream\n for response.output_item.added (line 962), response.function_call_arguments.delta\n (line 989), and response.output_item.done (line 1033). Our patch uses output_index\n from the event to properly track parallel tool calls.\n\n3. OpenAI Responses API Non-Streaming (_patch_openai_responses_transform_response):\n - LiteLLM's transform_response doesn't properly concatenate multiple reasoning\n summary parts in non-streaming responses\n - Multiple ReasoningSummaryItem objects should be joined with newlines\n STATUS: STILL NEEDED - LiteLLM's _convert_response_output_to_choices (lines 366-370)\n only keeps the LAST summary item text, discarding earlier parts. Our patch\n concatenates all summary texts with double newlines.\n\n4. Azure Responses API Fake Streaming (_patch_azure_responses_should_fake_stream):\n - LiteLLM uses \"fake streaming\" (MockResponsesAPIStreamingIterator) for models\n not in its database, which buffers the entire response before yielding\n - This causes poor time-to-first-token for Azure custom model deployments\n - Azure's Responses API supports native streaming, so we force real streaming\n STATUS: STILL NEEDED - AzureOpenAIResponsesAPIConfig does NOT override should_fake_stream,\n so it inherits from OpenAIResponsesAPIConfig which returns True for models not\n in litellm.utils.supports_native_streaming(). Custom Azure deployments will\n still use fake streaming without this patch.\n\n# Note: 5 and 6 are to supress a warning and may fix usage info but is not strictly required for the app to run\n5. Responses API Usage Format Mismatch (_patch_responses_api_usage_format):\n - LiteLLM uses model_construct as a fallback in multiple places when\n ResponsesAPIResponse validation fails\n - This bypasses the usage validator, allowing chat completion format usage\n (completion_tokens, prompt_tokens) to be stored instead of Responses API format\n (input_tokens, output_tokens)\n - When model_dump() is later called, Pydantic emits a serialization warning\n STATUS: STILL NEEDED - Multiple files use model_construct which bypasses validation:\n openai/responses/transformation.py, chatgpt/responses/transformation.py,\n manus/responses/transformation.py, volcengine/responses/transformation.py,\n and handler.py. Our patch wraps ResponsesAPIResponse.model_construct itself\n to transform usage in all code paths.\n\n6. Logging Usage Transformation Warning (_patch_logging_assembled_streaming_response):\n - LiteLLM's _get_assembled_streaming_response in litellm_logging.py transforms\n ResponseAPIUsage to chat completion format and sets it as a dict on the\n ResponsesAPIResponse.usage field\n - This replaces the proper ResponseAPIUsage object with a dict, causing Pydantic\n to emit a serialization warning when model_dump() is called later\n STATUS: STILL NEEDED - litellm_core_utils/litellm_logging.py lines 3185-3199 set\n usage as a dict with chat completion format instead of keeping it as\n ResponseAPIUsage. Our patch creates a deep copy before modification.\n\n7. Responses API metadata=None TypeError (_patch_responses_metadata_none):\n - LiteLLM's @client decorator wrapper in utils.py uses kwargs.get(\"metadata\", {})\n to check for router calls, but when metadata is explicitly None (key exists with\n value None), the default {} is not used\n - This causes \"argument of type 'NoneType' is not iterable\" TypeError which swallows\n the real exception (e.g. AuthenticationError for wrong API key)\n - Surfaces as: APIConnectionError: OpenAIException - argument of type 'NoneType' is\n not iterable\n STATUS: STILL NEEDED - litellm/utils.py wrapper function (line 1721) does not guard\n against metadata being explicitly None. Triggered when Responses API bridge\n passes **litellm_params containing metadata=None.\n\"\"\"\n\nimport time\nimport uuid\nfrom typing import Any\nfrom typing import cast\nfrom typing import List\nfrom typing import Optional\n\nfrom litellm.completion_extras.litellm_responses_transformation.transformation import (\n LiteLLMResponsesTransformationHandler,\n)\nfrom litellm.completion_extras.litellm_responses_transformation.transformation import (\n OpenAiResponsesToChatCompletionStreamIterator,\n)\nfrom litellm.llms.ollama.chat.transformation import OllamaChatCompletionResponseIterator\nfrom litellm.llms.ollama.common_utils import OllamaError\nfrom litellm.types.utils import ChatCompletionUsageBlock\nfrom litellm.types.utils import ModelResponseStream\n\n\ndef _patch_ollama_chunk_parser() -> None:\n \"\"\"\n Patches OllamaChatCompletionResponseIterator.chunk_parser to properly handle\n reasoning content and content in streaming responses.\n \"\"\"\n if (\n getattr(OllamaChatCompletionResponseIterator.chunk_parser, \"__name__\", \"\")\n == \"_patched_chunk_parser\"\n ):\n return\n\n def _patched_chunk_parser(self: Any, chunk: dict) -> ModelResponseStream:\n try:\n \"\"\"\n Expected chunk format:\n {\n \"model\": \"llama3.1\",\n \"created_at\": \"2025-05-24T02:12:05.859654Z\",\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"\",\n \"tool_calls\": [{\n \"function\": {\n \"name\": \"get_latest_album_ratings\",\n \"arguments\": {\n \"artist_name\": \"Taylor Swift\"\n }\n }\n }]\n },\n \"done_reason\": \"stop\",\n \"done\": true,\n ...\n }\n Need to:\n - convert 'message' to 'delta'\n - return finish_reason when done is true\n - return usage when done is true\n \"\"\"\n from litellm.types.utils import Delta\n from litellm.types.utils import StreamingChoices\n\n # process tool calls - if complete function arg - add id to tool call\n tool_calls = chunk[\"message\"].get(\"tool_calls\")\n if tool_calls is not None:\n for tool_call in tool_calls:\n function_args = tool_call.get(\"function\").get(\"arguments\")\n if function_args is not None and len(function_args) > 0:\n is_function_call_complete = self._is_function_call_complete(\n function_args\n )\n if is_function_call_complete:\n tool_call[\"id\"] = str(uuid.uuid4())\n\n # PROCESS REASONING CONTENT\n reasoning_content: Optional[str] = None\n content: Optional[str] = None\n thinking_content = chunk[\"message\"].get(\"thinking\")\n if thinking_content: # Truthy check: skips None and empty string \"\"\n reasoning_content = thinking_content\n if self.started_reasoning_content is False:\n self.started_reasoning_content = True\n if chunk[\"message\"].get(\"content\") is not None:\n message_content = chunk[\"message\"].get(\"content\")\n # Track whether we are inside ... tagged content.\n in_think_tag_block = bool(getattr(self, \"_in_think_tag_block\", False))\n if \"\" in message_content:\n message_content = message_content.replace(\"\", \"\")\n self.started_reasoning_content = True\n self.finished_reasoning_content = False\n in_think_tag_block = True\n if \"\" in message_content and self.started_reasoning_content:\n message_content = message_content.replace(\"\", \"\")\n self.finished_reasoning_content = True\n in_think_tag_block = False\n\n # For native Ollama \"thinking\" streams, content without active\n # think tags indicates a transition into regular assistant output.\n if (\n self.started_reasoning_content\n and not self.finished_reasoning_content\n and not in_think_tag_block\n and not thinking_content\n ):\n self.finished_reasoning_content = True\n\n self._in_think_tag_block = in_think_tag_block\n\n # When Ollama returns both \"thinking\" and \"content\" in the same\n # chunk, preserve both instead of classifying content as reasoning.\n if thinking_content and not in_think_tag_block:\n content = message_content\n elif (\n self.started_reasoning_content\n and not self.finished_reasoning_content\n ):\n reasoning_content = message_content\n else:\n content = message_content\n\n delta = Delta(\n content=content,\n reasoning_content=reasoning_content,\n tool_calls=tool_calls,\n )\n if chunk[\"done\"] is True:\n finish_reason = chunk.get(\"done_reason\", \"stop\")\n choices = [\n StreamingChoices(\n delta=delta,\n finish_reason=finish_reason,\n )\n ]\n else:\n choices = [\n StreamingChoices(\n delta=delta,\n )\n ]\n\n usage = ChatCompletionUsageBlock(\n prompt_tokens=chunk.get(\"prompt_eval_count\", 0),\n completion_tokens=chunk.get(\"eval_count\", 0),\n total_tokens=chunk.get(\"prompt_eval_count\", 0)\n + chunk.get(\"eval_count\", 0),\n )\n\n return ModelResponseStream(\n id=str(uuid.uuid4()),\n object=\"chat.completion.chunk\",\n created=int(time.time()), # ollama created_at is in UTC\n usage=usage,\n model=chunk[\"model\"],\n choices=choices,\n )\n except KeyError as e:\n raise OllamaError(\n message=f\"KeyError: {e}, Got unexpected response from Ollama: {chunk}\",\n status_code=400,\n headers={\"Content-Type\": \"application/json\"},\n )\n except Exception as e:\n raise e\n\n OllamaChatCompletionResponseIterator.chunk_parser = _patched_chunk_parser # type: ignore[method-assign]\n\n\ndef _patch_openai_responses_parallel_tool_calls() -> None:\n \"\"\"\n Patches OpenAiResponsesToChatCompletionStreamIterator to properly handle:\n 1. Parallel tool calls by using output_index from streaming events\n 2. Reasoning summary sections by inserting newlines between different summary indices\n\n LiteLLM's implementation hardcodes index=0 for all tool calls, breaking parallel tool calls.\n The OpenAI Responses API provides output_index in each event to track which tool call\n the event belongs to.\n\n STATUS: STILL NEEDED - LiteLLM hardcodes index=0 in translate_responses_chunk_to_openai_stream\n for response.output_item.added (line 962), response.function_call_arguments.delta\n (line 989), and response.output_item.done (line 1033). Our patch uses output_index\n from the event to properly track parallel tool calls.\n \"\"\"\n if (\n getattr(\n OpenAiResponsesToChatCompletionStreamIterator.chunk_parser,\n \"__name__\",\n \"\",\n )\n == \"_patched_responses_chunk_parser\"\n ):\n return\n\n def _patched_responses_chunk_parser(\n self: Any, chunk: dict\n ) -> \"ModelResponseStream\":\n from pydantic import BaseModel\n\n from litellm.types.llms.openai import (\n ChatCompletionToolCallFunctionChunk,\n ResponsesAPIStreamEvents,\n )\n from litellm.types.utils import (\n ChatCompletionToolCallChunk,\n Delta,\n ModelResponseStream,\n StreamingChoices,\n )\n\n parsed_chunk = chunk\n if not parsed_chunk:\n raise ValueError(\"Chat provider: Empty parsed_chunk\")\n\n if isinstance(parsed_chunk, BaseModel):\n parsed_chunk = parsed_chunk.model_dump()\n if not isinstance(parsed_chunk, dict):\n raise ValueError(f\"Chat provider: Invalid chunk type {type(parsed_chunk)}\")\n\n event_type = parsed_chunk.get(\"type\")\n if isinstance(event_type, ResponsesAPIStreamEvents):\n event_type = event_type.value\n\n # Get the output_index for proper parallel tool call tracking\n output_index = parsed_chunk.get(\"output_index\", 0)\n\n if event_type == \"response.output_item.added\":\n output_item = parsed_chunk.get(\"item\", {})\n if output_item.get(\"type\") == \"function_call\":\n provider_specific_fields = output_item.get(\"provider_specific_fields\")\n if provider_specific_fields and not isinstance(\n provider_specific_fields, dict\n ):\n provider_specific_fields = (\n dict(provider_specific_fields)\n if hasattr(provider_specific_fields, \"__dict__\")\n else {}\n )\n\n function_chunk = ChatCompletionToolCallFunctionChunk(\n name=output_item.get(\"name\", None),\n arguments=parsed_chunk.get(\"arguments\", \"\"),\n )\n if provider_specific_fields:\n function_chunk[\"provider_specific_fields\"] = (\n provider_specific_fields\n )\n\n tool_call_chunk = ChatCompletionToolCallChunk(\n id=output_item.get(\"call_id\"),\n index=output_index, # Use output_index for parallel tool calls\n type=\"function\",\n function=function_chunk,\n )\n if provider_specific_fields:\n tool_call_chunk.provider_specific_fields = provider_specific_fields # type: ignore\n\n return ModelResponseStream(\n choices=[\n StreamingChoices(\n index=0,\n delta=Delta(tool_calls=[tool_call_chunk]),\n finish_reason=None,\n )\n ]\n )\n\n elif event_type == \"response.function_call_arguments.delta\":\n content_part: Optional[str] = parsed_chunk.get(\"delta\", None)\n if content_part:\n return ModelResponseStream(\n choices=[\n StreamingChoices(\n index=0,\n delta=Delta(\n tool_calls=[\n ChatCompletionToolCallChunk(\n id=None,\n index=output_index, # Use output_index for parallel tool calls\n type=\"function\",\n function=ChatCompletionToolCallFunctionChunk(\n name=None, arguments=content_part\n ),\n )\n ]\n ),\n finish_reason=None,\n )\n ]\n )\n else:\n raise ValueError(\n f\"Chat provider: Invalid function argument delta {parsed_chunk}\"\n )\n\n elif event_type == \"response.output_item.done\":\n output_item = parsed_chunk.get(\"item\", {})\n if output_item.get(\"type\") == \"function_call\":\n provider_specific_fields = output_item.get(\"provider_specific_fields\")\n if provider_specific_fields and not isinstance(\n provider_specific_fields, dict\n ):\n provider_specific_fields = (\n dict(provider_specific_fields)\n if hasattr(provider_specific_fields, \"__dict__\")\n else {}\n )\n\n function_chunk = ChatCompletionToolCallFunctionChunk(\n name=output_item.get(\"name\", None),\n arguments=\"\", # responses API sends everything again, we don't need it\n )\n if provider_specific_fields:\n function_chunk[\"provider_specific_fields\"] = (\n provider_specific_fields\n )\n\n tool_call_chunk = ChatCompletionToolCallChunk(\n id=output_item.get(\"call_id\"),\n index=output_index, # Use output_index for parallel tool calls\n type=\"function\",\n function=function_chunk,\n )\n if provider_specific_fields:\n tool_call_chunk.provider_specific_fields = provider_specific_fields # type: ignore\n\n return ModelResponseStream(\n choices=[\n StreamingChoices(\n index=0,\n delta=Delta(tool_calls=[tool_call_chunk]),\n finish_reason=\"tool_calls\",\n )\n ]\n )\n\n elif event_type == \"response.reasoning_summary_text.delta\":\n # Handle reasoning summary with newlines between sections\n content_part = parsed_chunk.get(\"delta\", None)\n if content_part:\n summary_index = parsed_chunk.get(\"summary_index\", 0)\n\n # Track the last summary index to insert newlines between parts\n last_summary_index = getattr(\n self, \"_last_reasoning_summary_index\", None\n )\n if (\n last_summary_index is not None\n and summary_index != last_summary_index\n ):\n # New summary part started, prepend newlines to separate them\n content_part = \"\\n\\n\" + content_part\n self._last_reasoning_summary_index = summary_index\n\n return ModelResponseStream(\n choices=[\n StreamingChoices(\n index=cast(int, summary_index),\n delta=Delta(reasoning_content=content_part),\n )\n ]\n )\n\n # For all other event types, use the original static method\n return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(\n parsed_chunk\n )\n\n _patched_responses_chunk_parser.__name__ = \"_patched_responses_chunk_parser\"\n OpenAiResponsesToChatCompletionStreamIterator.chunk_parser = _patched_responses_chunk_parser # type: ignore[method-assign]\n\n\ndef _patch_openai_responses_transform_response() -> None:\n \"\"\"\n Patches LiteLLMResponsesTransformationHandler.transform_response to properly\n concatenate multiple reasoning summary parts with newlines in non-streaming responses.\n \"\"\"\n # Store the original method\n original_transform_response = (\n LiteLLMResponsesTransformationHandler.transform_response\n )\n\n if (\n getattr(\n original_transform_response,\n \"__name__\",\n \"\",\n )\n == \"_patched_transform_response\"\n ):\n return\n\n def _patched_transform_response(\n self: Any,\n model: str,\n raw_response: Any,\n model_response: Any,\n logging_obj: Any,\n request_data: dict,\n messages: List[Any],\n optional_params: dict,\n litellm_params: dict,\n encoding: Any,\n api_key: Optional[str] = None,\n json_mode: Optional[bool] = None,\n ) -> Any:\n \"\"\"\n Patched transform_response that properly concatenates reasoning summary parts\n with newlines.\n \"\"\"\n from openai.types.responses.response import Response as ResponsesAPIResponse\n from openai.types.responses.response_reasoning_item import ResponseReasoningItem\n\n # Check if raw_response has reasoning items that need concatenation\n if isinstance(raw_response, ResponsesAPIResponse) and raw_response.output:\n for item in raw_response.output:\n if isinstance(item, ResponseReasoningItem) and item.summary:\n # Concatenate summary texts with double newlines\n summary_texts = []\n for summary_item in item.summary:\n text = getattr(summary_item, \"text\", \"\")\n if text:\n summary_texts.append(text)\n\n if len(summary_texts) > 1:\n # Modify the first summary item to contain all concatenated text\n combined_text = \"\\n\\n\".join(summary_texts)\n if hasattr(item.summary[0], \"text\"):\n # Create a modified copy of the response with concatenated text\n # Since OpenAI types are typically frozen, we need to work around this\n # by modifying the object after the fact or using the result\n pass # The fix is applied in the result processing below\n\n # Call the original method\n result = original_transform_response(\n self,\n model,\n raw_response,\n model_response,\n logging_obj,\n request_data,\n messages,\n optional_params,\n litellm_params,\n encoding,\n api_key,\n json_mode,\n )\n\n # Post-process: If there are multiple summary items, fix the reasoning_content\n if isinstance(raw_response, ResponsesAPIResponse) and raw_response.output:\n for item in raw_response.output:\n if isinstance(item, ResponseReasoningItem) and item.summary:\n if len(item.summary) > 1:\n # Concatenate all summary texts with double newlines\n summary_texts = []\n for summary_item in item.summary:\n text = getattr(summary_item, \"text\", \"\")\n if text:\n summary_texts.append(text)\n\n if summary_texts:\n combined_text = \"\\n\\n\".join(summary_texts)\n # Update the reasoning_content in the result choices\n if hasattr(result, \"choices\"):\n for choice in result.choices:\n if hasattr(choice, \"message\") and hasattr(\n choice.message, \"reasoning_content\"\n ):\n choice.message.reasoning_content = combined_text\n break # Only process the first reasoning item\n\n return result\n\n _patched_transform_response.__name__ = \"_patched_transform_response\"\n LiteLLMResponsesTransformationHandler.transform_response = _patched_transform_response # type: ignore[method-assign]\n\n\ndef _patch_azure_responses_should_fake_stream() -> None:\n \"\"\"\n Patches AzureOpenAIResponsesAPIConfig.should_fake_stream to always return False.\n\n By default, LiteLLM uses \"fake streaming\" (MockResponsesAPIStreamingIterator) for models\n not in its database. This causes Azure custom model deployments to buffer the entire\n response before yielding, resulting in poor time-to-first-token.\n\n Azure's Responses API supports native streaming, so we override this to always use\n real streaming (SyncResponsesAPIStreamingIterator).\n \"\"\"\n from litellm.llms.azure.responses.transformation import (\n AzureOpenAIResponsesAPIConfig,\n )\n\n if (\n getattr(AzureOpenAIResponsesAPIConfig.should_fake_stream, \"__name__\", \"\")\n == \"_patched_should_fake_stream\"\n ):\n return\n\n def _patched_should_fake_stream(\n self: Any, # noqa: ARG001\n model: Optional[str], # noqa: ARG001\n stream: Optional[bool], # noqa: ARG001\n custom_llm_provider: Optional[str] = None, # noqa: ARG001\n ) -> bool:\n # Azure Responses API supports native streaming - never fake it\n return False\n\n _patched_should_fake_stream.__name__ = \"_patched_should_fake_stream\"\n AzureOpenAIResponsesAPIConfig.should_fake_stream = _patched_should_fake_stream # type: ignore[method-assign]\n\n\ndef _patch_responses_api_usage_format() -> None:\n \"\"\"\n Patches ResponsesAPIResponse.model_construct to properly transform usage data\n from chat completion format to Responses API format.\n\n LiteLLM uses model_construct as a fallback in multiple places when ResponsesAPIResponse\n validation fails. This bypasses the usage validator, allowing usage data in chat\n completion format (completion_tokens, prompt_tokens) to be stored instead of Responses\n API format (input_tokens, output_tokens), causing Pydantic serialization warnings.\n\n This patch wraps model_construct to transform usage before construction, ensuring\n the correct type regardless of which code path calls model_construct.\n\n Affected locations in LiteLLM:\n - litellm/llms/openai/responses/transformation.py (lines 183, 563)\n - litellm/llms/chatgpt/responses/transformation.py (line 153)\n - litellm/llms/manus/responses/transformation.py (lines 243, 334)\n - litellm/llms/volcengine/responses/transformation.py (line 280)\n - litellm/completion_extras/litellm_responses_transformation/handler.py (line 51)\n \"\"\"\n from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse\n\n original_model_construct = ResponsesAPIResponse.model_construct\n\n if getattr(original_model_construct, \"_is_patched\", False):\n return\n\n @classmethod # type: ignore[misc]\n def _patched_model_construct(\n cls: Any,\n _fields_set: Optional[set[str]] = None,\n **values: Any,\n ) -> \"ResponsesAPIResponse\":\n \"\"\"\n Patched model_construct that ensures usage is a ResponseAPIUsage object.\n \"\"\"\n # Transform usage if present and not already the correct type\n if \"usage\" in values and values[\"usage\"] is not None:\n usage = values[\"usage\"]\n if not isinstance(usage, ResponseAPIUsage):\n if isinstance(usage, dict):\n values = dict(values) # Don't mutate original\n # Check if it's in chat completion format\n if \"prompt_tokens\" in usage or \"completion_tokens\" in usage:\n # Transform from chat completion format\n values[\"usage\"] = ResponseAPIUsage(\n input_tokens=usage.get(\"prompt_tokens\", 0),\n output_tokens=usage.get(\"completion_tokens\", 0),\n total_tokens=usage.get(\"total_tokens\", 0),\n )\n elif \"input_tokens\" in usage or \"output_tokens\" in usage:\n # Already in Responses API format, just convert to proper type\n values[\"usage\"] = ResponseAPIUsage(\n input_tokens=usage.get(\"input_tokens\", 0),\n output_tokens=usage.get(\"output_tokens\", 0),\n total_tokens=usage.get(\"total_tokens\", 0),\n )\n\n # Call original model_construct (need to call it as unbound method)\n return original_model_construct.__func__(cls, _fields_set, **values) # type: ignore[attr-defined]\n\n _patched_model_construct._is_patched = True # type: ignore[attr-defined]\n ResponsesAPIResponse.model_construct = _patched_model_construct # type: ignore[method-assign, assignment]\n\n\ndef _patch_logging_assembled_streaming_response() -> None:\n \"\"\"\n Patches LiteLLMLoggingObj._get_assembled_streaming_response to create a deep copy\n of the ResponsesAPIResponse before modifying its usage field.\n\n The original code transforms usage to chat completion format and sets it as a dict\n directly on the ResponsesAPIResponse.usage field. This mutates the original object,\n causing Pydantic serialization warnings when model_dump() is called later because\n the usage field contains a dict instead of the expected ResponseAPIUsage type.\n\n This patch creates a copy of the response before modification, preserving the\n original object with its proper ResponseAPIUsage type.\n \"\"\"\n from litellm import LiteLLMLoggingObj\n from litellm.responses.utils import ResponseAPILoggingUtils\n from litellm.types.llms.openai import (\n ResponseAPIUsage,\n ResponseCompletedEvent,\n ResponsesAPIResponse,\n )\n from litellm.types.utils import ModelResponse, TextCompletionResponse\n\n original_method = LiteLLMLoggingObj._get_assembled_streaming_response\n\n if getattr(original_method, \"_is_patched\", False):\n return\n\n def _patched_get_assembled_streaming_response(\n self: Any, # noqa: ARG001\n result: Any,\n start_time: Any, # noqa: ARG001\n end_time: Any, # noqa: ARG001\n is_async: bool, # noqa: ARG001\n streaming_chunks: List[Any], # noqa: ARG001\n ) -> Any:\n \"\"\"\n Patched version that creates a copy before modifying usage.\n\n The original LiteLLM code transforms usage to chat completion format and\n sets it directly as a dict, which causes Pydantic serialization warnings.\n This patch uses model_construct to rebuild the response with the transformed\n usage, ensuring proper typing.\n \"\"\"\n if isinstance(result, ModelResponse):\n return result\n elif isinstance(result, TextCompletionResponse):\n return result\n elif isinstance(result, ResponseCompletedEvent):\n # Get the original response data\n original_response = result.response\n response_data = original_response.model_dump()\n\n # Transform usage if present\n if isinstance(original_response.usage, ResponseAPIUsage):\n transformed_usage = (\n ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(\n original_response.usage\n )\n )\n # Put the transformed usage (in chat completion format) into response_data\n # Our patched model_construct will convert it back to ResponseAPIUsage\n response_data[\"usage\"] = (\n transformed_usage.model_dump()\n if hasattr(transformed_usage, \"model_dump\")\n else dict(transformed_usage)\n )\n\n # Rebuild using model_construct - our patch ensures usage is properly typed\n response_copy = ResponsesAPIResponse.model_construct(**response_data)\n\n # Copy hidden params\n if hasattr(original_response, \"_hidden_params\"):\n response_copy._hidden_params = dict(original_response._hidden_params)\n\n return response_copy\n else:\n return None\n\n _patched_get_assembled_streaming_response._is_patched = True # type: ignore[attr-defined]\n LiteLLMLoggingObj._get_assembled_streaming_response = _patched_get_assembled_streaming_response # type: ignore[method-assign]\n\n\ndef _patch_responses_metadata_none() -> None:\n \"\"\"\n Patches litellm.responses to normalize metadata=None to metadata={} in kwargs.\n\n LiteLLM's @client decorator wrapper in utils.py (line 1721) does:\n _is_litellm_router_call = \"model_group\" in kwargs.get(\"metadata\", {})\n When metadata is explicitly None in kwargs, kwargs.get(\"metadata\", {}) returns\n None (the key exists, so the default is not used), causing:\n TypeError: argument of type 'NoneType' is not iterable\n\n This swallows the real exception (e.g. AuthenticationError) and surfaces as:\n APIConnectionError: OpenAIException - argument of type 'NoneType' is not iterable\n\n This happens when the Responses API bridge calls litellm.responses() with\n **litellm_params which may contain metadata=None.\n\n STATUS: STILL NEEDED - litellm/utils.py wrapper function uses kwargs.get(\"metadata\", {})\n which does not guard against metadata being explicitly None. Same pattern exists\n on line 1407 for async path.\n \"\"\"\n import litellm as _litellm\n from functools import wraps\n\n original_responses = _litellm.responses\n\n if getattr(original_responses, \"_metadata_patched\", False):\n return\n\n @wraps(original_responses)\n def _patched_responses(*args: Any, **kwargs: Any) -> Any:\n if kwargs.get(\"metadata\") is None:\n kwargs[\"metadata\"] = {}\n return original_responses(*args, **kwargs)\n\n _patched_responses._metadata_patched = True # type: ignore[attr-defined]\n _litellm.responses = _patched_responses\n\n\ndef apply_monkey_patches() -> None:\n \"\"\"\n Apply all necessary monkey patches to LiteLLM for compatibility.\n\n This includes:\n - Patching OllamaChatCompletionResponseIterator.chunk_parser for streaming content\n - Patching translate_responses_chunk_to_openai_stream for parallel tool calls\n - Patching LiteLLMResponsesTransformationHandler.transform_response for non-streaming responses\n - Patching AzureOpenAIResponsesAPIConfig.should_fake_stream to enable native streaming\n - Patching ResponsesAPIResponse.model_construct to fix usage format in all code paths\n - Patching LiteLLMLoggingObj._get_assembled_streaming_response to avoid mutating original response\n - Patching litellm.responses to fix metadata=None causing TypeError in error handling\n \"\"\"\n _patch_ollama_chunk_parser()\n _patch_openai_responses_parallel_tool_calls()\n _patch_openai_responses_transform_response()\n _patch_azure_responses_should_fake_stream()\n _patch_responses_api_usage_format()\n _patch_logging_assembled_streaming_response()\n _patch_responses_metadata_none()\n" }, { "path": "backend/onyx/llm/model_metadata_enrichments.json", "content": "{\n \"ai21.j2-mid-v1\": {\n \"display_name\": \"J2 Mid\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"v1\"\n },\n \"ai21.j2-ultra-v1\": {\n \"display_name\": \"J2 Ultra\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"v1\"\n },\n \"ai21.jamba-1-5-large-v1:0\": {\n \"display_name\": \"Jamba 1.5 Large\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"v1:0\"\n },\n \"ai21.jamba-1-5-mini-v1:0\": {\n \"display_name\": \"Jamba 1.5 Mini\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"v1:0\"\n },\n \"ai21.jamba-instruct-v1:0\": {\n \"display_name\": \"Jamba Instruct\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"v1:0\"\n },\n \"amazon.nova-lite-v1:0\": {\n \"display_name\": \"Nova Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"amazon.nova-micro-v1:0\": {\n \"display_name\": \"Nova Micro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"amazon.nova-pro-v1:0\": {\n \"display_name\": \"Nova Pro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"amazon.titan-text-express-v1\": {\n \"display_name\": \"Titan Text Express\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1\"\n },\n \"amazon.titan-text-lite-v1\": {\n \"display_name\": \"Titan Text Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1\"\n },\n \"amazon.titan-text-premier-v1:0\": {\n \"display_name\": \"Titan Text Premier\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620-v1:0\"\n },\n \"anthropic.claude-3-5-sonnet-20241022-v2:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20241022-v2:0\"\n },\n \"anthropic.claude-3-sonnet-20240229-v1:0\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240229-v1:0\"\n },\n \"anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001-v1:0\"\n },\n \"anthropic.claude-haiku-4-5@20251001\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001\"\n },\n \"anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"anthropic.claude-opus-4-1-20250805-v1:0\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250805-v1:0\"\n },\n \"anthropic.claude-opus-4-20250514-v1:0\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514-v1:0\"\n },\n \"anthropic.claude-opus-4-5-20251101-v1:0\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251101-v1:0\"\n },\n \"anthropic.claude-sonnet-4-20250514-v1:0\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514-v1:0\"\n },\n \"anthropic.claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929-v1:0\"\n },\n \"anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"apac.amazon.nova-lite-v1:0\": {\n \"display_name\": \"Nova Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"apac.amazon.nova-micro-v1:0\": {\n \"display_name\": \"Nova Micro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"apac.amazon.nova-pro-v1:0\": {\n \"display_name\": \"Nova Pro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"apac.anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620-v1:0\"\n },\n \"apac.anthropic.claude-3-5-sonnet-20241022-v2:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20241022-v2:0\"\n },\n \"apac.anthropic.claude-3-sonnet-20240229-v1:0\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240229-v1:0\"\n },\n \"apac.anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001-v1:0\"\n },\n \"apac.anthropic.claude-sonnet-4-20250514-v1:0\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514-v1:0\"\n },\n \"au.anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001-v1:0\"\n },\n \"au.anthropic.claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929-v1:0\"\n },\n \"azure/claude-haiku-4-5\": {\n \"display_name\": \"Claude Haiku\",\n \"model_vendor\": \"anthropic\"\n },\n \"azure/claude-opus-4-1\": {\n \"display_name\": \"Claude Opus\",\n \"model_vendor\": \"anthropic\"\n },\n \"azure/claude-sonnet-4-5\": {\n \"display_name\": \"Claude Sonnet\",\n \"model_vendor\": \"anthropic\"\n },\n \"azure/codex-mini\": {\n \"display_name\": \"Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/command-r-plus\": {\n \"display_name\": \"Command R Plus\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"latest\"\n },\n \"azure/computer-use-preview\": {\n \"display_name\": \"Computer Use Preview\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"preview\"\n },\n \"azure/container\": {\n \"display_name\": \"Container\",\n \"model_vendor\": \"azure\",\n \"model_version\": \"latest\"\n },\n \"azure/eu/gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"azure/eu/gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"azure/eu/gpt-4o-mini-2024-07-18\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-07-18\"\n },\n \"azure/eu/gpt-4o-mini-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Mini Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/eu/gpt-4o-realtime-preview-2024-10-01\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-10-01\"\n },\n \"azure/eu/gpt-4o-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/eu/gpt-5-2025-08-07\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/eu/gpt-5-mini-2025-08-07\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/eu/gpt-5-nano-2025-08-07\": {\n \"display_name\": \"GPT 5 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/eu/gpt-5.1\": {\n \"display_name\": \"GPT-5.1\",\n \"model_vendor\": \"openai\"\n },\n \"azure/eu/gpt-5.1-chat\": {\n \"display_name\": \"GPT-5.1 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"azure/eu/gpt-5.1-codex\": {\n \"display_name\": \"GPT-5.1 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"azure/eu/gpt-5.1-codex-mini\": {\n \"display_name\": \"GPT-5.1 Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/eu/o1-2024-12-17\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/eu/o1-mini-2024-09-12\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-09-12\"\n },\n \"azure/eu/o1-preview-2024-09-12\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-09-12\"\n },\n \"azure/eu/o3-mini-2025-01-31\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-01-31\"\n },\n \"azure/global-standard/gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"azure/global-standard/gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"azure/global-standard/gpt-4o-mini\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/global/gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"azure/global/gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"azure/global/gpt-5.1\": {\n \"display_name\": \"GPT-5.1\",\n \"model_vendor\": \"openai\"\n },\n \"azure/global/gpt-5.1-chat\": {\n \"display_name\": \"GPT-5.1 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"azure/global/gpt-5.1-codex\": {\n \"display_name\": \"GPT-5.1 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"azure/global/gpt-5.1-codex-mini\": {\n \"display_name\": \"GPT-5.1 Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-3.5-turbo\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-3.5-turbo-0125\": {\n \"display_name\": \"GPT 3.5 Turbo 0125\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0125\"\n },\n \"azure/gpt-3.5-turbo-instruct-0914\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0914\"\n },\n \"azure/gpt-35-turbo\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-35-turbo-0125\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0125\"\n },\n \"azure/gpt-35-turbo-0301\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0301\"\n },\n \"azure/gpt-35-turbo-0613\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"azure/gpt-35-turbo-1106\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"1106\"\n },\n \"azure/gpt-35-turbo-16k\": {\n \"display_name\": \"GPT-3.5 Turbo 16K\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-35-turbo-16k-0613\": {\n \"display_name\": \"GPT-3.5 Turbo 16K\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"azure/gpt-35-turbo-instruct\": {\n \"display_name\": \"GPT-3.5 Turbo Instruct\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-35-turbo-instruct-0914\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0914\"\n },\n \"azure/gpt-4\": {\n \"display_name\": \"GPT-4\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4-0125-preview\": {\n \"display_name\": \"GPT 4 0125 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0125\"\n },\n \"azure/gpt-4-0613\": {\n \"display_name\": \"GPT 4 0613\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"azure/gpt-4-1106-preview\": {\n \"display_name\": \"GPT 4 1106 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"1106\"\n },\n \"azure/gpt-4-32k\": {\n \"display_name\": \"GPT-4 32K\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4-32k-0613\": {\n \"display_name\": \"GPT 4 32k 0613\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"azure/gpt-4-turbo\": {\n \"display_name\": \"GPT-4 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4-turbo-2024-04-09\": {\n \"display_name\": \"GPT-4 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-04-09\"\n },\n \"azure/gpt-4-turbo-vision-preview\": {\n \"display_name\": \"GPT-4 Turbo Vision Preview\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4.1\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4.1-2025-04-14\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"azure/gpt-4.1-mini\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4.1-mini-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"azure/gpt-4.1-nano\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4.1-nano-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"azure/gpt-4.5-preview\": {\n \"display_name\": \"GPT-4.5 Preview\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4o\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4o-2024-05-13\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-05-13\"\n },\n \"azure/gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"azure/gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"azure/gpt-4o-audio-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/gpt-4o-mini\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4o-mini-2024-07-18\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-07-18\"\n },\n \"azure/gpt-4o-mini-audio-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/gpt-4o-mini-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Mini Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/gpt-4o-mini-transcribe\": {\n \"display_name\": \"GPT-4o Mini Transcribe\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4o-mini-tts\": {\n \"display_name\": \"GPT-4o Mini TTS\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4o-realtime-preview-2024-10-01\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-10-01\"\n },\n \"azure/gpt-4o-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/gpt-4o-transcribe\": {\n \"display_name\": \"GPT-4o Transcribe\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-4o-transcribe-diarize\": {\n \"display_name\": \"GPT-4o Transcribe Diarize\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5-2025-08-07\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/gpt-5-chat\": {\n \"display_name\": \"GPT-5 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5-chat-latest\": {\n \"display_name\": \"GPT 5 Chat\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/gpt-5-codex\": {\n \"display_name\": \"GPT-5 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5-mini\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5-mini-2025-08-07\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/gpt-5-nano\": {\n \"display_name\": \"GPT-5 Nano\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5-nano-2025-08-07\": {\n \"display_name\": \"GPT 5 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/gpt-5-pro\": {\n \"display_name\": \"GPT-5 Pro\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5.1\": {\n \"display_name\": \"GPT-5.1\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5.1-2025-11-13\": {\n \"display_name\": \"GPT 5.1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-11-13\"\n },\n \"azure/gpt-5.1-chat\": {\n \"display_name\": \"GPT-5.1 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5.1-chat-2025-11-13\": {\n \"display_name\": \"GPT 5.1 Chat\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-11-13\"\n },\n \"azure/gpt-5.1-codex\": {\n \"display_name\": \"GPT-5.1 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5.1-codex-2025-11-13\": {\n \"display_name\": \"GPT-5.1 Codex\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-11-13\"\n },\n \"azure/gpt-5.1-codex-mini\": {\n \"display_name\": \"GPT-5.1 Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/gpt-5.1-codex-mini-2025-11-13\": {\n \"display_name\": \"GPT-5.1 Codex Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-11-13\"\n },\n \"azure/gpt-audio-2025-08-28\": {\n \"display_name\": \"GPT Audio\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-28\"\n },\n \"azure/gpt-audio-mini-2025-10-06\": {\n \"display_name\": \"GPT Audio Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-10-06\"\n },\n \"azure/gpt-realtime-2025-08-28\": {\n \"display_name\": \"GPT Realtime\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-28\"\n },\n \"azure/gpt-realtime-mini-2025-10-06\": {\n \"display_name\": \"GPT Realtime Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-10-06\"\n },\n \"azure/mistral-large-2402\": {\n \"display_name\": \"Mistral Large 24.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2402\"\n },\n \"azure/mistral-large-latest\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"azure/o1\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o1-2024-12-17\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/o1-mini\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o1-mini-2024-09-12\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-09-12\"\n },\n \"azure/o1-preview\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o1-preview-2024-09-12\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-09-12\"\n },\n \"azure/o3\": {\n \"display_name\": \"o3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o3-2025-04-16\": {\n \"display_name\": \"o3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-16\"\n },\n \"azure/o3-deep-research\": {\n \"display_name\": \"O3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o3-mini\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o3-mini-2025-01-31\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-01-31\"\n },\n \"azure/o3-pro\": {\n \"display_name\": \"O3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o3-pro-2025-06-10\": {\n \"display_name\": \"O3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-06-10\"\n },\n \"azure/o4-mini\": {\n \"display_name\": \"o4 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"azure/o4-mini-2025-04-16\": {\n \"display_name\": \"o4 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-16\"\n },\n \"azure/us/gpt-4.1-2025-04-14\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"azure/us/gpt-4.1-mini-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"azure/us/gpt-4.1-nano-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"azure/us/gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"azure/us/gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"azure/us/gpt-4o-mini-2024-07-18\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-07-18\"\n },\n \"azure/us/gpt-4o-mini-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Mini Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/us/gpt-4o-realtime-preview-2024-10-01\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-10-01\"\n },\n \"azure/us/gpt-4o-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/us/gpt-5-2025-08-07\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/us/gpt-5-mini-2025-08-07\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/us/gpt-5-nano-2025-08-07\": {\n \"display_name\": \"GPT 5 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"azure/us/gpt-5.1\": {\n \"display_name\": \"GPT-5.1\",\n \"model_vendor\": \"openai\"\n },\n \"azure/us/gpt-5.1-chat\": {\n \"display_name\": \"GPT-5.1 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"azure/us/gpt-5.1-codex\": {\n \"display_name\": \"GPT-5.1 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"azure/us/gpt-5.1-codex-mini\": {\n \"display_name\": \"GPT-5.1 Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"azure/us/o1-2024-12-17\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"azure/us/o1-mini-2024-09-12\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-09-12\"\n },\n \"azure/us/o1-preview-2024-09-12\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-09-12\"\n },\n \"azure/us/o3-2025-04-16\": {\n \"display_name\": \"o3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-16\"\n },\n \"azure/us/o3-mini-2025-01-31\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-01-31\"\n },\n \"azure/us/o4-mini-2025-04-16\": {\n \"display_name\": \"o4 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-16\"\n },\n \"azure_ai/Llama-3.2-11B-Vision-Instruct\": {\n \"display_name\": \"Llama 3.2 11B Vision Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Llama-3.2-90B-Vision-Instruct\": {\n \"display_name\": \"Llama 3.2 90B Vision Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Llama-3.3-70B-Instruct\": {\n \"display_name\": \"Llama 3.3 70B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8\": {\n \"display_name\": \"Llama 4 Maverick 17B 128E Instruct FP8\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Llama-4-Scout-17B-16E-Instruct\": {\n \"display_name\": \"Llama 4 Scout 17B 16E Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/MAI-DS-R1\": {\n \"display_name\": \"MAI-DS-R1\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Meta-Llama-3-70B-Instruct\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Meta-Llama-3.1-405B-Instruct\": {\n \"display_name\": \"Llama 3.1 405B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Meta-Llama-3.1-70B-Instruct\": {\n \"display_name\": \"Llama 3.1 70B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Meta-Llama-3.1-8B-Instruct\": {\n \"display_name\": \"Llama 3.1 8B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"azure_ai/Phi-3-medium-128k-instruct\": {\n \"display_name\": \"Phi 3 Medium 128k Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3-medium-4k-instruct\": {\n \"display_name\": \"Phi 3 Medium 4k Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3-mini-128k-instruct\": {\n \"display_name\": \"Phi 3 Mini 128k Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3-mini-4k-instruct\": {\n \"display_name\": \"Phi 3 Mini 4k Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3-small-128k-instruct\": {\n \"display_name\": \"Phi 3 Small 128k Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3-small-8k-instruct\": {\n \"display_name\": \"Phi 3 Small 8k Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3.5-MoE-instruct\": {\n \"display_name\": \"Phi 3.5 MOE Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3.5-mini-instruct\": {\n \"display_name\": \"Phi 3.5 Mini Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-3.5-vision-instruct\": {\n \"display_name\": \"Phi 3.5 Vision Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-4\": {\n \"display_name\": \"Phi 4\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-4-mini-instruct\": {\n \"display_name\": \"Phi 4 Mini Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-4-mini-reasoning\": {\n \"display_name\": \"Phi 4 Mini Reasoning\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-4-multimodal-instruct\": {\n \"display_name\": \"Phi 4 Multimodal Instruct\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/Phi-4-reasoning\": {\n \"display_name\": \"Phi 4 Reasoning\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/deepseek-r1\": {\n \"display_name\": \"DeepSeek R1\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/deepseek-v3\": {\n \"display_name\": \"DeepSeek V3\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"v3\"\n },\n \"azure_ai/deepseek-v3-0324\": {\n \"display_name\": \"DeepSeek v3 0324\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"0324\"\n },\n \"azure_ai/global/grok-3\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/global/grok-3-mini\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/grok-3\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/grok-3-mini\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/grok-4\": {\n \"display_name\": \"Grok 4\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/grok-4-fast-non-reasoning\": {\n \"display_name\": \"Grok 4 Fast Non Reasoning\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/grok-4-fast-reasoning\": {\n \"display_name\": \"Grok 4 Fast Reasoning\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/grok-code-fast-1\": {\n \"display_name\": \"Grok Code Fast 1\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/jais-30b-chat\": {\n \"display_name\": \"Jais 30B Chat\",\n \"model_vendor\": \"g42\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/jamba-instruct\": {\n \"display_name\": \"Jamba Instruct\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/ministral-3b\": {\n \"display_name\": \"Ministral 3B\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/mistral-large\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/mistral-large-2407\": {\n \"display_name\": \"Mistral Large 24.07\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2407\"\n },\n \"azure_ai/mistral-large-latest\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/mistral-medium-2505\": {\n \"display_name\": \"Mistral Medium 2505\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2505\"\n },\n \"azure_ai/mistral-nemo\": {\n \"display_name\": \"Mistral Nemo\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/mistral-small\": {\n \"display_name\": \"Mistral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"azure_ai/mistral-small-2503\": {\n \"display_name\": \"Mistral Small 2503\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2503\"\n },\n \"bedrock/*/1-month-commitment/cohere.command-light-text-v14\": {\n \"display_name\": \"Command Light Text\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v14\"\n },\n \"bedrock/*/1-month-commitment/cohere.command-text-v14\": {\n \"display_name\": \"Command Text\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v14\"\n },\n \"bedrock/*/6-month-commitment/cohere.command-light-text-v14\": {\n \"display_name\": \"Command Light Text\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v14\"\n },\n \"bedrock/*/6-month-commitment/cohere.command-text-v14\": {\n \"display_name\": \"Command Text\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v14\"\n },\n \"bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/ap-northeast-1/6-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/ap-northeast-1/6-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/ap-northeast-1/6-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/ap-northeast-1/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/ap-northeast-1/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/ap-northeast-1/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/ap-south-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/ca-central-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/ca-central-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/eu-central-1/1-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/eu-central-1/1-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/eu-central-1/1-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/eu-central-1/6-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/eu-central-1/6-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/eu-central-1/6-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/eu-central-1/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/eu-central-1/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/eu-central-1/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/eu-west-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/eu-west-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/eu-west-2/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/eu-west-2/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/eu-west-3/mistral.mistral-7b-instruct-v0:2\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:2\"\n },\n \"bedrock/eu-west-3/mistral.mistral-large-2402-v1:0\": {\n \"display_name\": \"Mistral Large 24.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2402-v1:0\"\n },\n \"bedrock/eu-west-3/mistral.mixtral-8x7b-instruct-v0:1\": {\n \"display_name\": \"Mixtral 8x7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:1\"\n },\n \"bedrock/invoke/anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620-v1:0\"\n },\n \"bedrock/sa-east-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/sa-east-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-east-1/1-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-east-1/1-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-east-1/1-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/us-east-1/6-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-east-1/6-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-east-1/6-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/us-east-1/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-east-1/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-east-1/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/us-east-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-east-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-east-1/mistral.mistral-7b-instruct-v0:2\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:2\"\n },\n \"bedrock/us-east-1/mistral.mistral-large-2402-v1:0\": {\n \"display_name\": \"Mistral Large 24.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2402-v1:0\"\n },\n \"bedrock/us-east-1/mistral.mixtral-8x7b-instruct-v0:1\": {\n \"display_name\": \"Mixtral 8x7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:1\"\n },\n \"bedrock/us-gov-east-1/amazon.nova-pro-v1:0\": {\n \"display_name\": \"Nova Pro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-east-1/amazon.titan-text-express-v1\": {\n \"display_name\": \"Titan Text Express\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-gov-east-1/amazon.titan-text-lite-v1\": {\n \"display_name\": \"Titan Text Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-gov-east-1/amazon.titan-text-premier-v1:0\": {\n \"display_name\": \"Titan Text Premier\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-east-1/anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620-v1:0\"\n },\n \"bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929-v1:0\"\n },\n \"bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-west-1/amazon.nova-pro-v1:0\": {\n \"display_name\": \"Nova Pro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-west-1/amazon.titan-text-express-v1\": {\n \"display_name\": \"Titan Text Express\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-gov-west-1/amazon.titan-text-lite-v1\": {\n \"display_name\": \"Titan Text Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-gov-west-1/amazon.titan-text-premier-v1:0\": {\n \"display_name\": \"Titan Text Premier\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-west-1/anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620-v1:0\"\n },\n \"bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929-v1:0\"\n },\n \"bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-west-1/meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-west-1/meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"bedrock/us-west-2/1-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-west-2/1-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-west-2/1-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/us-west-2/6-month-commitment/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-west-2/6-month-commitment/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-west-2/6-month-commitment/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/us-west-2/anthropic.claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-west-2/anthropic.claude-v1\": {\n \"display_name\": \"Claude\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"bedrock/us-west-2/anthropic.claude-v2:1\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v2:1\"\n },\n \"bedrock/us-west-2/mistral.mistral-7b-instruct-v0:2\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:2\"\n },\n \"bedrock/us-west-2/mistral.mistral-large-2402-v1:0\": {\n \"display_name\": \"Mistral Large 24.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2402-v1:0\"\n },\n \"bedrock/us-west-2/mistral.mixtral-8x7b-instruct-v0:1\": {\n \"display_name\": \"Mixtral 8x7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:1\"\n },\n \"chat-bison\": {\n \"display_name\": \"Chat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"chat-bison-32k\": {\n \"display_name\": \"Chat Bison 32k\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"chat-bison-32k@002\": {\n \"display_name\": \"Chat Bison 32k\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"chat-bison@001\": {\n \"display_name\": \"Chat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"chat-bison@002\": {\n \"display_name\": \"Chat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"chatgpt-4o-latest\": {\n \"display_name\": \"ChatGPT 4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"claude-3-5-sonnet-20240620\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620\"\n },\n \"claude-3-5-sonnet-20241022\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20241022\"\n },\n \"claude-3-5-sonnet-latest\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"latest\"\n },\n \"claude-4-opus-20250514\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"claude-4-sonnet-20250514\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"claude-haiku-4-5\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"claude-haiku-4-5-20251001\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001\"\n },\n \"claude-opus-4-1\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\"\n },\n \"claude-opus-4-1-20250805\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250805\"\n },\n \"claude-opus-4-1@20250805\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250805\"\n },\n \"claude-opus-4-20250514\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"claude-opus-4-5\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"claude-opus-4-6\": {\n \"display_name\": \"Claude Opus 4.6\",\n \"model_vendor\": \"anthropic\"\n },\n \"claude-opus-4-5-20251101\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251101\"\n },\n \"claude-sonnet-4-20250514\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"claude-sonnet-4-5\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"claude-sonnet-4-6\": {\n \"display_name\": \"Claude Sonnet 4.6\",\n \"model_vendor\": \"anthropic\"\n },\n \"claude-sonnet-4-5-20250929\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929\"\n },\n \"claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929-v1:0\"\n },\n \"codechat-bison\": {\n \"display_name\": \"Codechat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"codechat-bison-32k\": {\n \"display_name\": \"Codechat Bison 32k\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"codechat-bison-32k@002\": {\n \"display_name\": \"Codechat Bison 32k\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"codechat-bison@001\": {\n \"display_name\": \"Codechat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"codechat-bison@002\": {\n \"display_name\": \"Codechat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"codechat-bison@latest\": {\n \"display_name\": \"Codechat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"codex-mini-latest\": {\n \"display_name\": \"Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"cohere.command-light-text-v14\": {\n \"display_name\": \"Command Light Text\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v14\"\n },\n \"cohere.command-r-plus-v1:0\": {\n \"display_name\": \"Command R Plus\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v1:0\"\n },\n \"cohere.command-r-v1:0\": {\n \"display_name\": \"Command R\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v1:0\"\n },\n \"cohere.command-text-v14\": {\n \"display_name\": \"Command Text\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"v14\"\n },\n \"computer-use-preview\": {\n \"display_name\": \"Computer Use Preview\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"preview\"\n },\n \"deepseek.v3-v1:0\": {\n \"display_name\": \"DeepSeek V3\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"v1:0\"\n },\n \"deepseek/deepseek-chat\": {\n \"display_name\": \"DeepSeek Chat\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"deepseek/deepseek-coder\": {\n \"display_name\": \"DeepSeek Coder\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"deepseek/deepseek-r1\": {\n \"display_name\": \"DeepSeek R1\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"deepseek/deepseek-reasoner\": {\n \"display_name\": \"DeepSeek Reasoner\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"deepseek/deepseek-v3\": {\n \"display_name\": \"DeepSeek V3\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"v3\"\n },\n \"eu.amazon.nova-lite-v1:0\": {\n \"display_name\": \"Nova Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"eu.amazon.nova-micro-v1:0\": {\n \"display_name\": \"Nova Micro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"eu.amazon.nova-pro-v1:0\": {\n \"display_name\": \"Nova Pro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"eu.anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620-v1:0\"\n },\n \"eu.anthropic.claude-3-5-sonnet-20241022-v2:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20241022-v2:0\"\n },\n \"eu.anthropic.claude-3-sonnet-20240229-v1:0\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240229-v1:0\"\n },\n \"eu.anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001-v1:0\"\n },\n \"eu.anthropic.claude-opus-4-1-20250805-v1:0\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250805-v1:0\"\n },\n \"eu.anthropic.claude-opus-4-20250514-v1:0\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514-v1:0\"\n },\n \"eu.anthropic.claude-sonnet-4-20250514-v1:0\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514-v1:0\"\n },\n \"eu.anthropic.claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929-v1:0\"\n },\n \"eu.meta.llama3-2-1b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 1B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"eu.meta.llama3-2-3b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 3B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1:0\"\n },\n \"eu.mistral.pixtral-large-2502-v1:0\": {\n \"display_name\": \"Pixtral Large 25.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"2502-v1:0\"\n },\n \"eu.twelvelabs.pegasus-1-2-v1:0\": {\n \"display_name\": \"Pegasus 1.2\",\n \"model_vendor\": \"twelvelabs\",\n \"model_version\": \"1.2-v1:0\"\n },\n \"ft:gpt-3.5-turbo\": {\n \"display_name\": \"Ft:gpt 3.5 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"ft:gpt-3.5-turbo-0125\": {\n \"display_name\": \"GPT-3.5 Turbo (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0125\"\n },\n \"ft:gpt-3.5-turbo-0613\": {\n \"display_name\": \"GPT-3.5 Turbo (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"ft:gpt-3.5-turbo-1106\": {\n \"display_name\": \"GPT-3.5 Turbo (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"1106\"\n },\n \"ft:gpt-4-0613\": {\n \"display_name\": \"GPT-4 (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"ft:gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"ft:gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"ft:gpt-4o-mini-2024-07-18\": {\n \"display_name\": \"GPT-4o Mini (Fine-tuned)\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-07-18\"\n },\n \"gemini-1.0-pro\": {\n \"display_name\": \"Gemini 1.0 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini-1.0-pro-001\": {\n \"display_name\": \"Gemini 1.0 Pro 001\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini-1.0-pro-002\": {\n \"display_name\": \"Gemini 1.0 Pro 002\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"gemini-1.0-ultra\": {\n \"display_name\": \"Gemini 1.0 Ultra\",\n \"model_vendor\": \"google\"\n },\n \"gemini-1.0-ultra-001\": {\n \"display_name\": \"Gemini 1.0 Ultra 001\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini-1.5-flash\": {\n \"display_name\": \"Gemini 1.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini-1.5-flash-001\": {\n \"display_name\": \"Gemini 1.5 Flash 001\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini-1.5-flash-002\": {\n \"display_name\": \"Gemini 1.5 Flash 002\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"gemini-1.5-flash-exp-0827\": {\n \"display_name\": \"Gemini 1.5 Flash Exp 0827\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0827\"\n },\n \"gemini-1.5-flash-preview-0514\": {\n \"display_name\": \"Gemini 1.5 Flash Preview 0514\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0514\"\n },\n \"gemini-1.5-pro\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini-1.5-pro-001\": {\n \"display_name\": \"Gemini 1.5 Pro 001\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini-1.5-pro-002\": {\n \"display_name\": \"Gemini 1.5 Pro 002\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"gemini-1.5-pro-preview-0215\": {\n \"display_name\": \"Gemini 1.5 Pro Preview 0215\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0215\"\n },\n \"gemini-1.5-pro-preview-0409\": {\n \"display_name\": \"Gemini 1.5 Pro Preview 0409\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0409\"\n },\n \"gemini-1.5-pro-preview-0514\": {\n \"display_name\": \"Gemini 1.5 Pro Preview 0514\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0514\"\n },\n \"gemini-2.0-flash\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.0-flash-001\": {\n \"display_name\": \"Gemini 2.0 Flash 001\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini-2.0-flash-exp\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.0-flash-lite\": {\n \"display_name\": \"Gemini 2.0 Flash Lite\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.0-flash-lite-001\": {\n \"display_name\": \"Gemini 2.0 Flash Lite 001\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini-2.0-flash-live-preview-04-09\": {\n \"display_name\": \"Gemini 2.0 Flash Live Preview 04 09\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.0-flash-thinking-exp\": {\n \"display_name\": \"Gemini 2.0 Flash Thinking\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.0-flash-thinking-exp-01-21\": {\n \"display_name\": \"Gemini 2.0 Flash Thinking Exp 01 21\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.0-pro-exp-02-05\": {\n \"display_name\": \"Gemini 2.0 Pro Exp 02 05\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-flash\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-flash-lite\": {\n \"display_name\": \"Gemini 2.5 Flash Lite\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-flash-lite-preview-06-17\": {\n \"display_name\": \"Gemini 2.5 Flash Lite Preview 06 17\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-flash-lite-preview-09-2025\": {\n \"display_name\": \"Gemini 2.5 Flash Lite Preview 09 2025\",\n \"model_vendor\": \"google\",\n \"model_version\": \"2025\"\n },\n \"gemini-2.5-flash-preview-04-17\": {\n \"display_name\": \"Gemini 2.5 Flash Preview 04 17\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-flash-preview-05-20\": {\n \"display_name\": \"Gemini 2.5 Flash Preview 05 20\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-flash-preview-09-2025\": {\n \"display_name\": \"Gemini 2.5 Flash Preview 09 2025\",\n \"model_vendor\": \"google\",\n \"model_version\": \"2025\"\n },\n \"gemini-2.5-pro\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-pro-exp-03-25\": {\n \"display_name\": \"Gemini 2.5 Pro Exp 03 25\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-pro-preview-03-25\": {\n \"display_name\": \"Gemini 2.5 Pro Preview 03 25\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-pro-preview-05-06\": {\n \"display_name\": \"Gemini 2.5 Pro Preview 05 06\",\n \"model_vendor\": \"google\"\n },\n \"gemini-2.5-pro-preview-06-05\": {\n \"display_name\": \"Gemini 2.5 Pro Preview 06 05\",\n \"model_vendor\": \"google\"\n },\n \"gemini-3-pro-preview\": {\n \"display_name\": \"Gemini 3 Pro Preview\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini-3-flash-preview\": {\n \"display_name\": \"Gemini 3 Flash Preview\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini-flash-experimental\": {\n \"display_name\": \"Gemini Flash Experimental\",\n \"model_vendor\": \"google\",\n \"model_version\": \"experimental\"\n },\n \"gemini-pro\": {\n \"display_name\": \"Gemini Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini-pro-experimental\": {\n \"display_name\": \"Gemini Pro Experimental\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-1.5-flash\": {\n \"display_name\": \"Gemini 1.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-1.5-flash-001\": {\n \"display_name\": \"Gemini 1.5 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/gemini-1.5-flash-002\": {\n \"display_name\": \"Gemini 1.5 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"gemini/gemini-1.5-flash-8b\": {\n \"display_name\": \"Gemini 1.5 Flash 8B\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-1.5-flash-8b-exp-0827\": {\n \"display_name\": \"Gemini 1.5 Flash 8B\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0827\"\n },\n \"gemini/gemini-1.5-flash-8b-exp-0924\": {\n \"display_name\": \"Gemini 1.5 Flash 8B\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0924\"\n },\n \"gemini/gemini-1.5-flash-exp-0827\": {\n \"display_name\": \"Gemini 1.5 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0827\"\n },\n \"gemini/gemini-1.5-flash-latest\": {\n \"display_name\": \"Gemini 1.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-1.5-pro\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-1.5-pro-001\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/gemini-1.5-pro-002\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"gemini/gemini-1.5-pro-exp-0801\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0801\"\n },\n \"gemini/gemini-1.5-pro-exp-0827\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"0827\"\n },\n \"gemini/gemini-1.5-pro-latest\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash-001\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/gemini-2.0-flash-exp\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash-lite\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash-lite-preview-02-05\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash-live-001\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/gemini-2.0-flash-preview-image-generation\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash-thinking-exp\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-flash-thinking-exp-01-21\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.0-pro-exp-02-05\": {\n \"display_name\": \"Gemini 2.0\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-image\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-image-preview\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-lite\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-lite-preview-06-17\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-lite-preview-09-2025\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"2025\"\n },\n \"gemini/gemini-2.5-flash-preview-04-17\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-preview-05-20\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-flash-preview-09-2025\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"2025\"\n },\n \"gemini/gemini-2.5-flash-preview-tts\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-pro\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-pro-exp-03-25\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-pro-preview-03-25\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-pro-preview-05-06\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-pro-preview-06-05\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-2.5-pro-preview-tts\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-3-pro-image-preview\": {\n \"display_name\": \"Gemini 1.0 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini/gemini-3-pro-preview\": {\n \"display_name\": \"Gemini 1.0 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini/gemini-embedding-001\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/gemini-exp-1114\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"experimental\"\n },\n \"gemini/gemini-exp-1206\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"experimental\"\n },\n \"gemini/gemini-flash-latest\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"gemini/gemini-flash-lite-latest\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"gemini/gemini-gemma-2-27b-it\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"gemini/gemini-gemma-2-9b-it\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"gemini/gemini-live-2.5-flash-preview-native-audio-09-2025\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini/gemini-pro\": {\n \"display_name\": \"Gemini 1.0 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemini-pro-vision\": {\n \"display_name\": \"Gemini 1.0 Pro\",\n \"model_vendor\": \"google\"\n },\n \"gemini/gemma-3-27b-it\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"gemini/imagen-3.0-fast-generate-001\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/imagen-3.0-generate-001\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/imagen-3.0-generate-002\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"002\"\n },\n \"gemini/imagen-4.0-fast-generate-001\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/imagen-4.0-generate-001\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/imagen-4.0-ultra-generate-001\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/learnlm-1.5-pro-experimental\": {\n \"display_name\": \"Gemini 1.5 Pro\",\n \"model_vendor\": \"google\",\n \"model_version\": \"experimental\"\n },\n \"gemini/veo-2.0-generate-001\": {\n \"display_name\": \"Gemini 2.0\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"gemini/veo-3.0-fast-generate-preview\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini/veo-3.0-generate-preview\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini/veo-3.1-fast-generate-preview\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"gemini/veo-3.1-generate-preview\": {\n \"display_name\": \"Gemini\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"global.anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001\"\n },\n \"global.anthropic.claude-sonnet-4-20250514-v1:0\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"global.anthropic.claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929\"\n },\n \"gpt-3.5-turbo\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-3.5-turbo-0125\": {\n \"display_name\": \"GPT 3.5 Turbo 0125\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0125\"\n },\n \"gpt-3.5-turbo-0301\": {\n \"display_name\": \"GPT 3.5 Turbo 0301\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0301\"\n },\n \"gpt-3.5-turbo-0613\": {\n \"display_name\": \"GPT 3.5 Turbo 0613\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"gpt-3.5-turbo-1106\": {\n \"display_name\": \"GPT 3.5 Turbo 1106\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"1106\"\n },\n \"gpt-3.5-turbo-16k\": {\n \"display_name\": \"GPT-3.5 Turbo 16K\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-3.5-turbo-16k-0613\": {\n \"display_name\": \"GPT 3.5 Turbo 16k 0613\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"gpt-4\": {\n \"display_name\": \"GPT-4\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4-0125-preview\": {\n \"display_name\": \"GPT-4 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0125\"\n },\n \"gpt-4-0314\": {\n \"display_name\": \"GPT-4\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0314\"\n },\n \"gpt-4-0613\": {\n \"display_name\": \"GPT-4\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"gpt-4-1106-preview\": {\n \"display_name\": \"GPT-4 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"1106\"\n },\n \"gpt-4-1106-vision-preview\": {\n \"display_name\": \"GPT-4 Vision Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"1106\"\n },\n \"gpt-4-32k\": {\n \"display_name\": \"GPT-4 32K\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4-32k-0314\": {\n \"display_name\": \"GPT-4 32K\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0314\"\n },\n \"gpt-4-32k-0613\": {\n \"display_name\": \"GPT-4 32K\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"0613\"\n },\n \"gpt-4-turbo\": {\n \"display_name\": \"GPT-4 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4-turbo-2024-04-09\": {\n \"display_name\": \"GPT-4 Turbo\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-04-09\"\n },\n \"gpt-4-turbo-preview\": {\n \"display_name\": \"GPT-4 Turbo Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4-vision-preview\": {\n \"display_name\": \"GPT-4 Vision Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4.1\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4.1-2025-04-14\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"gpt-4.1-mini\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4.1-mini-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"gpt-4.1-nano\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4.1-nano-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"gpt-4.5-preview\": {\n \"display_name\": \"GPT-4.5 Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4.5-preview-2025-02-27\": {\n \"display_name\": \"GPT-4.5 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-02-27\"\n },\n \"gpt-4o\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-2024-05-13\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-05-13\"\n },\n \"gpt-4o-2024-08-06\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-08-06\"\n },\n \"gpt-4o-2024-11-20\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-11-20\"\n },\n \"gpt-4o-audio-preview\": {\n \"display_name\": \"GPT-4o Audio Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-audio-preview-2024-10-01\": {\n \"display_name\": \"GPT-4o Audio Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-10-01\"\n },\n \"gpt-4o-audio-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Audio Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"gpt-4o-audio-preview-2025-06-03\": {\n \"display_name\": \"GPT-4o Audio Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-06-03\"\n },\n \"gpt-4o-mini\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-mini-2024-07-18\": {\n \"display_name\": \"GPT-4o Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-07-18\"\n },\n \"gpt-4o-mini-audio-preview\": {\n \"display_name\": \"GPT-4o Mini Audio Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-mini-audio-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Mini Audio Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"gpt-4o-mini-realtime-preview\": {\n \"display_name\": \"GPT-4o Mini Realtime Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-mini-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Mini Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"gpt-4o-mini-search-preview\": {\n \"display_name\": \"GPT 4o Mini Search Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-mini-search-preview-2025-03-11\": {\n \"display_name\": \"GPT 4o Mini Search Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-03-11\"\n },\n \"gpt-4o-realtime-preview\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-realtime-preview-2024-10-01\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-10-01\"\n },\n \"gpt-4o-realtime-preview-2024-12-17\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"gpt-4o-realtime-preview-2025-06-03\": {\n \"display_name\": \"GPT-4o Realtime Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-06-03\"\n },\n \"gpt-4o-search-preview\": {\n \"display_name\": \"GPT 4o Search Preview\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-4o-search-preview-2025-03-11\": {\n \"display_name\": \"GPT 4o Search Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-03-11\"\n },\n \"gpt-5\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-2025-08-07\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"gpt-5-chat\": {\n \"display_name\": \"GPT 5 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-chat-latest\": {\n \"display_name\": \"GPT 5 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-codex\": {\n \"display_name\": \"GPT-5 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-mini\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-mini-2025-08-07\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"gpt-5-nano\": {\n \"display_name\": \"GPT 5 Nano\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-nano-2025-08-07\": {\n \"display_name\": \"GPT 5 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-07\"\n },\n \"gpt-5-pro\": {\n \"display_name\": \"GPT-5 Pro\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5-pro-2025-10-06\": {\n \"display_name\": \"GPT-5 Pro\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-10-06\"\n },\n \"gpt-5.4\": {\n \"display_name\": \"GPT-5.4\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.2-pro-2025-12-11\": {\n \"display_name\": \"GPT-5.2 Pro\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-12-11\"\n },\n \"gpt-5.2-pro\": {\n \"display_name\": \"GPT-5.2 Pro\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.2-chat-latest\": {\n \"display_name\": \"GPT 5.2 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.2-2025-12-11\": {\n \"display_name\": \"GPT 5.2\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-12-11\"\n },\n \"gpt-5.2\": {\n \"display_name\": \"GPT 5.2\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.1\": {\n \"display_name\": \"GPT 5.1\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.1-2025-11-13\": {\n \"display_name\": \"GPT 5.1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-11-13\"\n },\n \"gpt-5.1-chat-latest\": {\n \"display_name\": \"GPT 5.1 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.1-codex\": {\n \"display_name\": \"GPT-5.1 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-5.1-codex-mini\": {\n \"display_name\": \"GPT-5.1 Codex Mini\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-image-1-mini\": {\n \"display_name\": \"GPT Image 1 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"gpt-realtime\": {\n \"display_name\": \"GPT Realtime\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"gpt-realtime-2025-08-28\": {\n \"display_name\": \"GPT Realtime\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-08-28\"\n },\n \"gpt-realtime-mini\": {\n \"display_name\": \"GPT Realtime Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"jp.anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001\"\n },\n \"jp.anthropic.claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929\"\n },\n \"medlm-large\": {\n \"display_name\": \"MedLM Large\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"medlm-medium\": {\n \"display_name\": \"MedLM Medium\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"meta.llama2-13b-chat-v1\": {\n \"display_name\": \"Llama 2 13B Chat\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1\"\n },\n \"meta.llama2-70b-chat-v1\": {\n \"display_name\": \"Llama 2 70B Chat\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"v1\"\n },\n \"meta.llama3-1-405b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.1 405B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-1-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.1 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-1-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.1 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-2-11b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 11B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-2-1b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 1B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-2-3b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 3B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-2-90b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 90B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama3-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama4-maverick-17b-instruct-v1:0\": {\n \"display_name\": \"Llama 4 Maverick 17B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"meta.llama4-scout-17b-instruct-v1:0\": {\n \"display_name\": \"Llama 4 Scout 17B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"mistral.mistral-7b-instruct-v0:2\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0:2\"\n },\n \"mistral.mistral-large-2402-v1:0\": {\n \"display_name\": \"Mistral Large 24.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"1:0\"\n },\n \"mistral.mistral-large-2407-v1:0\": {\n \"display_name\": \"Mistral Large 24.07\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"1:0\"\n },\n \"mistral.mistral-small-2402-v1:0\": {\n \"display_name\": \"Mistral Small 24.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"1:0\"\n },\n \"mistral.mixtral-8x7b-instruct-v0:1\": {\n \"display_name\": \"Mixtral 8x7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"0:1\"\n },\n \"mistral/codestral-2405\": {\n \"display_name\": \"Codestral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/codestral-embed\": {\n \"display_name\": \"Codestral Embed\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/codestral-embed-2505\": {\n \"display_name\": \"Codestral Embed\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/codestral-latest\": {\n \"display_name\": \"Codestral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/codestral-mamba-latest\": {\n \"display_name\": \"Codestral Mamba\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/devstral-medium-2507\": {\n \"display_name\": \"Devstral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/devstral-small-2505\": {\n \"display_name\": \"Devstral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/devstral-small-2507\": {\n \"display_name\": \"Devstral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/magistral-medium-2506\": {\n \"display_name\": \"Magistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/magistral-medium-2509\": {\n \"display_name\": \"Magistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/magistral-medium-latest\": {\n \"display_name\": \"Magistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/magistral-small-2506\": {\n \"display_name\": \"Magistral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/magistral-small-latest\": {\n \"display_name\": \"Magistral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-embed\": {\n \"display_name\": \"Mistral Embed\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-large-2402\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-large-2407\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-large-2411\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-large-latest\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-medium\": {\n \"display_name\": \"Mistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-medium-2312\": {\n \"display_name\": \"Mistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-medium-2505\": {\n \"display_name\": \"Mistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-medium-latest\": {\n \"display_name\": \"Mistral Medium\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-ocr-2505-completion\": {\n \"display_name\": \"Mistral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-ocr-latest\": {\n \"display_name\": \"Mistral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-small\": {\n \"display_name\": \"Mistral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-small-latest\": {\n \"display_name\": \"Mistral Small\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/mistral-tiny\": {\n \"display_name\": \"Mistral Tiny\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/open-codestral-mamba\": {\n \"display_name\": \"Codestral Mamba\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/open-mistral-7b\": {\n \"display_name\": \"Open Mistral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/open-mistral-nemo\": {\n \"display_name\": \"Open Mistral Nemo\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/open-mistral-nemo-2407\": {\n \"display_name\": \"Open Mistral Nemo\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/open-mixtral-8x22b\": {\n \"display_name\": \"Open Mixtral 8x22B\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/open-mixtral-8x7b\": {\n \"display_name\": \"Open Mixtral 8x7B\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/pixtral-12b-2409\": {\n \"display_name\": \"Pixtral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/pixtral-large-2411\": {\n \"display_name\": \"Pixtral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"mistral/pixtral-large-latest\": {\n \"display_name\": \"Pixtral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"o1\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o1-2024-12-17\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-12-17\"\n },\n \"o1-mini\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o1-mini-2024-09-12\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o1-preview\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o1-preview-2024-09-12\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o1-pro\": {\n \"display_name\": \"o1 Pro\",\n \"model_vendor\": \"openai\"\n },\n \"o1-pro-2025-03-19\": {\n \"display_name\": \"o1 Pro\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-03-19\"\n },\n \"o3\": {\n \"display_name\": \"o3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o3-2025-04-16\": {\n \"display_name\": \"o3\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o3-deep-research\": {\n \"display_name\": \"o3 Deep Research\",\n \"model_vendor\": \"openai\"\n },\n \"o3-deep-research-2025-06-26\": {\n \"display_name\": \"o3 Deep Research\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-06-26\"\n },\n \"o3-mini\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o3-mini-2025-01-31\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o3-pro\": {\n \"display_name\": \"o3 Pro\",\n \"model_vendor\": \"openai\"\n },\n \"o3-pro-2025-06-10\": {\n \"display_name\": \"o3 Pro\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-06-10\"\n },\n \"o4-mini\": {\n \"display_name\": \"o4 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o4-mini-2025-04-16\": {\n \"display_name\": \"o4 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"o4-mini-deep-research\": {\n \"display_name\": \"o4 Mini Deep Research\",\n \"model_vendor\": \"openai\"\n },\n \"o4-mini-deep-research-2025-06-26\": {\n \"display_name\": \"o4 Mini Deep Research\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-06-26\"\n },\n \"ollama/codegeex4\": {\n \"display_name\": \"CodeGeeX4\",\n \"model_vendor\": \"zhipu\",\n \"model_version\": \"latest\"\n },\n \"ollama/codegemma\": {\n \"display_name\": \"Codegemma\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"ollama/codellama\": {\n \"display_name\": \"CodeLlama\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"latest\"\n },\n \"ollama/deepseek-coder-v2-base\": {\n \"display_name\": \"DeepSeek Coder v2 Base\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"ollama/deepseek-coder-v2-instruct\": {\n \"display_name\": \"DeepSeek Coder v2 Instruct\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"ollama/deepseek-coder-v2-lite-base\": {\n \"display_name\": \"DeepSeek Coder v2 Lite Base\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"ollama/deepseek-coder-v2-lite-instruct\": {\n \"display_name\": \"DeepSeek Coder v2 Lite Instruct\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"ollama/deepseek-v3.1:671b-cloud\": {\n \"display_name\": \"DeepSeek V3.1:671B Cloud\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"ollama/gpt-oss:120b-cloud\": {\n \"display_name\": \"GPT Open-Source 120B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"ollama/gpt-oss:20b-cloud\": {\n \"display_name\": \"GPT Open-Source 20B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"ollama/internlm2_5-20b-chat\": {\n \"display_name\": \"InternLM 2.5 20B Chat\",\n \"model_vendor\": \"shanghai-ai-lab\",\n \"model_version\": \"latest\"\n },\n \"ollama/llama2\": {\n \"display_name\": \"Llama 2\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama2-uncensored\": {\n \"display_name\": \"Llama 2 Uncensored\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama2:13b\": {\n \"display_name\": \"Llama 2:13B\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama2:70b\": {\n \"display_name\": \"Llama 2:70B\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama2:7b\": {\n \"display_name\": \"Llama 2:7B\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama3\": {\n \"display_name\": \"Llama 3\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama3.1\": {\n \"display_name\": \"Llama 3.1\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama3:70b\": {\n \"display_name\": \"Llama 3:70B\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/llama3:8b\": {\n \"display_name\": \"Llama 3:8B\",\n \"model_vendor\": \"meta\"\n },\n \"ollama/mistral\": {\n \"display_name\": \"Mistral\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"ollama/mistral-7B-Instruct-v0.1\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0.1\"\n },\n \"ollama/mistral-7B-Instruct-v0.2\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"v0.2\"\n },\n \"ollama/mistral-large-instruct-2407\": {\n \"display_name\": \"Mistral Large Instruct 24.07\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"ollama/mixtral-8x22B-Instruct-v0.1\": {\n \"display_name\": \"Mixtral 8x22B Instruct V0.1\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"ollama/mixtral-8x7B-Instruct-v0.1\": {\n \"display_name\": \"Mixtral 8x7B Instruct V0.1\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"ollama/orca-mini\": {\n \"display_name\": \"Orca Mini\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"ollama/qwen3-coder:480b-cloud\": {\n \"display_name\": \"Qwen3 Coder:480B Cloud\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"ollama/vicuna\": {\n \"display_name\": \"Vicuna\",\n \"model_vendor\": \"lmsys\",\n \"model_version\": \"latest\"\n },\n \"openai.gpt-oss-120b-1:0\": {\n \"display_name\": \"GPT Open-Source 120B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"v1:0\"\n },\n \"openai.gpt-oss-20b-1:0\": {\n \"display_name\": \"GPT Open-Source 20B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"v1:0\"\n },\n \"openai/container\": {\n \"display_name\": \"Container\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/agentica-org/deepcoder-14b-preview\": {\n \"display_name\": \"DeepCoder 14B Preview\",\n \"model_vendor\": \"agentica\"\n },\n \"openrouter/ai21/jamba-1-5-large\": {\n \"display_name\": \"Jamba 1.5 Large\",\n \"model_vendor\": \"ai21\"\n },\n \"openrouter/ai21/jamba-1-5-mini\": {\n \"display_name\": \"Jamba 1.5 Mini\",\n \"model_vendor\": \"ai21\"\n },\n \"openrouter/ai21/jamba-large-1.7\": {\n \"display_name\": \"Jamba Large 1.7\",\n \"model_vendor\": \"ai21\"\n },\n \"openrouter/aion-labs/aion-1.0\": {\n \"display_name\": \"AION 1.0\",\n \"model_vendor\": \"aion-labs\"\n },\n \"openrouter/alibaba/qwen-2.5-72b-instruct\": {\n \"display_name\": \"Qwen 2.5 72B Instruct\",\n \"model_vendor\": \"alibaba\"\n },\n \"openrouter/alibaba/qwen-2.5-coder-32b-instruct\": {\n \"display_name\": \"Qwen 2.5 Coder 32B\",\n \"model_vendor\": \"alibaba\"\n },\n \"openrouter/alibaba/tongyi-deepresearch-30b-a3b\": {\n \"display_name\": \"Tongyi DeepResearch 30B\",\n \"model_vendor\": \"alibaba\"\n },\n \"openrouter/alibaba/tongyi-deepresearch-30b-a3b:free\": {\n \"display_name\": \"Tongyi DeepResearch 30B (Free)\",\n \"model_vendor\": \"alibaba\"\n },\n \"openrouter/anthropic/claude-2\": {\n \"display_name\": \"Claude 2\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"latest\"\n },\n \"openrouter/anthropic/claude-3-sonnet\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\"\n },\n \"openrouter/anthropic/claude-3.5-sonnet\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"latest\"\n },\n \"openrouter/anthropic/claude-3.5-sonnet:beta\": {\n \"display_name\": \"Claude Sonnet 3.5:beta\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"latest\"\n },\n \"openrouter/anthropic/claude-haiku-4.5\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"latest\"\n },\n \"openrouter/anthropic/claude-instant-v1\": {\n \"display_name\": \"Claude Instant\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"v1\"\n },\n \"openrouter/anthropic/claude-opus-4\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\"\n },\n \"openrouter/anthropic/claude-opus-4.1\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\"\n },\n \"openrouter/anthropic/claude-opus-4.5\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"openrouter/anthropic/claude-sonnet-4\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\"\n },\n \"openrouter/anthropic/claude-sonnet-4.5\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"openrouter/baidu/ernie-4.5-300b-a47b\": {\n \"display_name\": \"ERNIE 4.5 300B\",\n \"model_vendor\": \"baidu\"\n },\n \"openrouter/baidu/ernie-4.5-vl-28b-a3b\": {\n \"display_name\": \"ERNIE 4.5 VL 28B\",\n \"model_vendor\": \"baidu\"\n },\n \"openrouter/bytedance/ui-tars-1.5-7b\": {\n \"display_name\": \"UI-TARS 1.5 7B\",\n \"model_vendor\": \"bytedance\",\n \"model_version\": \"latest\"\n },\n \"openrouter/cognitivecomputations/dolphin-mixtral-8x7b\": {\n \"display_name\": \"Dolphin Mixtral 8x7B\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/cohere/command-a\": {\n \"display_name\": \"Command A\",\n \"model_vendor\": \"cohere\"\n },\n \"openrouter/cohere/command-r\": {\n \"display_name\": \"Command R\",\n \"model_vendor\": \"cohere\"\n },\n \"openrouter/cohere/command-r-08-2024\": {\n \"display_name\": \"Command R\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"08-2024\"\n },\n \"openrouter/cohere/command-r-plus\": {\n \"display_name\": \"Command R Plus\",\n \"model_vendor\": \"cohere\"\n },\n \"openrouter/cohere/command-r-plus-08-2024\": {\n \"display_name\": \"Command R Plus\",\n \"model_vendor\": \"cohere\",\n \"model_version\": \"08-2024\"\n },\n \"openrouter/databricks/dbrx-instruct\": {\n \"display_name\": \"DBRX Instruct\",\n \"model_vendor\": \"databricks\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepcogito/cogito-v2-preview-deepseek-671b\": {\n \"display_name\": \"Cogito V2 Preview DeepSeek 671B\",\n \"model_vendor\": \"deepcogito\"\n },\n \"openrouter/deepcogito/cogito-v2-preview-llama-109b-moe\": {\n \"display_name\": \"Cogito V2 Preview Llama 109B MoE\",\n \"model_vendor\": \"deepcogito\"\n },\n \"openrouter/deepseek/deepseek-chat\": {\n \"display_name\": \"DeepSeek Chat\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepseek/deepseek-chat-v3-0324\": {\n \"display_name\": \"DeepSeek Chat v3 0324\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepseek/deepseek-chat-v3.1\": {\n \"display_name\": \"DeepSeek Chat V3.1\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepseek/deepseek-coder\": {\n \"display_name\": \"DeepSeek Coder\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepseek/deepseek-r1\": {\n \"display_name\": \"DeepSeek R1\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepseek/deepseek-r1-0528\": {\n \"display_name\": \"DeepSeek R1 0528\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"openrouter/deepseek/deepseek-v3.2-exp\": {\n \"display_name\": \"DeepSeek V3.2\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"experimental\"\n },\n \"openrouter/fireworks/firellava-13b\": {\n \"display_name\": \"FireLLaVA 13B\",\n \"model_vendor\": \"fireworks\",\n \"model_version\": \"latest\"\n },\n \"openrouter/google/gemini-2.0-flash-001\": {\n \"display_name\": \"Gemini 2.0 Flash\",\n \"model_vendor\": \"google\",\n \"model_version\": \"001\"\n },\n \"openrouter/google/gemini-2.5-flash\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemini-2.5-pro\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemini-3-pro-preview\": {\n \"display_name\": \"Gemini 3 Pro Preview\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"openrouter/google/gemini-pro-1.5\": {\n \"display_name\": \"Gemini Pro 1.5\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemini-pro-vision\": {\n \"display_name\": \"Gemini Pro Vision\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemma-2-27b-it\": {\n \"display_name\": \"Gemma 2 27B\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemma-2-9b-it\": {\n \"display_name\": \"Gemma 2 9B\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemma-2-9b-it:free\": {\n \"display_name\": \"Gemma 2 9B (Free)\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemma-3n-e4b-it\": {\n \"display_name\": \"Gemma 3N E4B\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/gemma-3n-e4b-it:free\": {\n \"display_name\": \"Gemma 3N E4B (Free)\",\n \"model_vendor\": \"google\"\n },\n \"openrouter/google/palm-2-chat-bison\": {\n \"display_name\": \"PaLM 2 Chat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"openrouter/google/palm-2-codechat-bison\": {\n \"display_name\": \"PaLM 2 Codechat Bison\",\n \"model_vendor\": \"google\",\n \"model_version\": \"latest\"\n },\n \"openrouter/gryphe/mythomax-l2-13b\": {\n \"display_name\": \"MythoMax L2 13B\",\n \"model_vendor\": \"gryphe\",\n \"model_version\": \"latest\"\n },\n \"openrouter/inclusionai/ring-1t\": {\n \"display_name\": \"Ring 1T\",\n \"model_vendor\": \"inclusionai\"\n },\n \"openrouter/jondurbin/airoboros-l2-70b-2.1\": {\n \"display_name\": \"Airoboros L2 70B\",\n \"model_vendor\": \"jondurbin\"\n },\n \"openrouter/mancer/weaver\": {\n \"display_name\": \"Weaver\",\n \"model_vendor\": \"mancer\",\n \"model_version\": \"latest\"\n },\n \"openrouter/meta-llama/codellama-34b-instruct\": {\n \"display_name\": \"CodeLlama 34B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/meta-llama/llama-2-13b-chat\": {\n \"display_name\": \"Llama 2 13B Chat\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/meta-llama/llama-2-70b-chat\": {\n \"display_name\": \"Llama 2 70B Chat\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/meta-llama/llama-3-70b-instruct\": {\n \"display_name\": \"Llama 3 70B Instruct\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/meta-llama/llama-3-70b-instruct:nitro\": {\n \"display_name\": \"Llama 3 70B Instruct:nitro\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/meta-llama/llama-3-8b-instruct:extended\": {\n \"display_name\": \"Llama 3 8B Instruct:extended\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/meta-llama/llama-3-8b-instruct:free\": {\n \"display_name\": \"Llama 3 8B Instruct:free\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/microsoft/wizardlm-2-8x22b:nitro\": {\n \"display_name\": \"WizardLM 2 8x22B\",\n \"model_vendor\": \"microsoft\",\n \"model_version\": \"latest\"\n },\n \"openrouter/minimax/minimax-m2\": {\n \"display_name\": \"MiniMax M2\",\n \"model_vendor\": \"minimax\",\n \"model_version\": \"latest\"\n },\n \"openrouter/mistralai/mistral-7b-instruct\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/mistralai/mistral-7b-instruct:free\": {\n \"display_name\": \"Mistral 7B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/mistralai/mistral-large\": {\n \"display_name\": \"Mistral Large\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/mistralai/mistral-small-3.1-24b-instruct\": {\n \"display_name\": \"Mistral Small 3.1 24B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/mistralai/mistral-small-3.2-24b-instruct\": {\n \"display_name\": \"Mistral Small 3.2 24B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/mistralai/mixtral-8x22b-instruct\": {\n \"display_name\": \"Mixtral 8x22B Instruct\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"openrouter/nousresearch/nous-hermes-llama2-13b\": {\n \"display_name\": \"Nous Hermes Llama 2 13B\",\n \"model_vendor\": \"meta\"\n },\n \"openrouter/openai/gpt-3.5-turbo\": {\n \"display_name\": \"GPT-3.5 Turbo\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-3.5-turbo-16k\": {\n \"display_name\": \"GPT-3.5 Turbo 16K\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4\": {\n \"display_name\": \"GPT-4\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4-vision-preview\": {\n \"display_name\": \"GPT-4 Vision Preview\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4.1\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4.1-2025-04-14\": {\n \"display_name\": \"GPT-4.1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"openrouter/openai/gpt-4.1-mini\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4.1-mini-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"openrouter/openai/gpt-4.1-nano\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4.1-nano-2025-04-14\": {\n \"display_name\": \"GPT-4.1 Nano\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2025-04-14\"\n },\n \"openrouter/openai/gpt-4o\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-4o-2024-05-13\": {\n \"display_name\": \"GPT-4o\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"2024-05-13\"\n },\n \"openrouter/openai/gpt-5\": {\n \"display_name\": \"GPT-5\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-5-chat\": {\n \"display_name\": \"GPT 5 Chat\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-5-codex\": {\n \"display_name\": \"GPT-5 Codex\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-5-mini\": {\n \"display_name\": \"GPT-5 Mini\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-5-nano\": {\n \"display_name\": \"GPT 5 Nano\",\n \"model_vendor\": \"openai\"\n },\n \"openrouter/openai/gpt-oss-120b\": {\n \"display_name\": \"GPT Open-Source 120B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/gpt-oss-20b\": {\n \"display_name\": \"GPT Open-Source 20B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o1\": {\n \"display_name\": \"o1\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o1-mini\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o1-mini-2024-09-12\": {\n \"display_name\": \"o1 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o1-preview\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o1-preview-2024-09-12\": {\n \"display_name\": \"o1 Preview\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o3-mini\": {\n \"display_name\": \"o3 Mini\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/openai/o3-mini-high\": {\n \"display_name\": \"O3 Mini High\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/pygmalionai/mythalion-13b\": {\n \"display_name\": \"Mythalion 13B\",\n \"model_vendor\": \"pygmalionai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/qwen/qwen-2.5-coder-32b-instruct\": {\n \"display_name\": \"Qwen 2.5 Coder 32B Instruct\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"openrouter/qwen/qwen-vl-plus\": {\n \"display_name\": \"Qwen Vl Plus\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"openrouter/qwen/qwen3-coder\": {\n \"display_name\": \"Qwen3 Coder\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"openrouter/switchpoint/router\": {\n \"display_name\": \"SwitchPoint Router\",\n \"model_vendor\": \"switchpoint\",\n \"model_version\": \"latest\"\n },\n \"openrouter/undi95/remm-slerp-l2-13b\": {\n \"display_name\": \"ReMM SLERP L2 13B\",\n \"model_vendor\": \"undi95\",\n \"model_version\": \"latest\"\n },\n \"openrouter/x-ai/grok-4\": {\n \"display_name\": \"Grok 4\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/x-ai/grok-4-fast:free\": {\n \"display_name\": \"Grok 4 Fast:free\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"openrouter/z-ai/glm-4.6\": {\n \"display_name\": \"GLM 4.6\",\n \"model_vendor\": \"zhipu\",\n \"model_version\": \"latest\"\n },\n \"openrouter/z-ai/glm-4.6:exacto\": {\n \"display_name\": \"GLM 4.6 Exacto\",\n \"model_vendor\": \"zhipu\",\n \"model_version\": \"latest\"\n },\n \"qwen.qwen3-235b-a22b-2507-v1:0\": {\n \"display_name\": \"Qwen.qwen3 235B A22b 2507\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"1:0\"\n },\n \"qwen.qwen3-32b-v1:0\": {\n \"display_name\": \"Qwen.qwen3 32B\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"1:0\"\n },\n \"qwen.qwen3-coder-30b-a3b-v1:0\": {\n \"display_name\": \"Qwen.qwen3 Coder 30B A3b\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"1:0\"\n },\n \"qwen.qwen3-coder-480b-a35b-v1:0\": {\n \"display_name\": \"Qwen.qwen3 Coder 480B A35b\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"1:0\"\n },\n \"twelvelabs.pegasus-1-2-v1:0\": {\n \"display_name\": \"Pegasus 1.2\",\n \"model_vendor\": \"twelvelabs\",\n \"model_version\": \"v1:0\"\n },\n \"us.amazon.nova-lite-v1:0\": {\n \"display_name\": \"Nova Lite\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"1:0\"\n },\n \"us.amazon.nova-micro-v1:0\": {\n \"display_name\": \"Nova Micro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"1:0\"\n },\n \"us.amazon.nova-premier-v1:0\": {\n \"display_name\": \"Nova Premier\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"v1:0\"\n },\n \"us.amazon.nova-pro-v1:0\": {\n \"display_name\": \"Nova Pro\",\n \"model_vendor\": \"amazon\",\n \"model_version\": \"1:0\"\n },\n \"us.anthropic.claude-3-5-sonnet-20240620-v1:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620\"\n },\n \"us.anthropic.claude-3-5-sonnet-20241022-v2:0\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20241022\"\n },\n \"us.anthropic.claude-3-sonnet-20240229-v1:0\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240229\"\n },\n \"us.anthropic.claude-haiku-4-5-20251001-v1:0\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001\"\n },\n \"us.anthropic.claude-opus-4-1-20250805-v1:0\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250805\"\n },\n \"us.anthropic.claude-opus-4-20250514-v1:0\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"us.anthropic.claude-opus-4-5-20251101-v1:0\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251101\"\n },\n \"us.anthropic.claude-sonnet-4-20250514-v1:0\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"us.anthropic.claude-sonnet-4-5-20250929-v1:0\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929\"\n },\n \"us.deepseek.r1-v1:0\": {\n \"display_name\": \"DeepSeek R1\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"v1:0\"\n },\n \"us.meta.llama3-1-405b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.1 405B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-1-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.1 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-1-8b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.1 8B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-2-11b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 11B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-2-1b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 1B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-2-3b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 3B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-2-90b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.2 90B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama3-3-70b-instruct-v1:0\": {\n \"display_name\": \"Llama 3.3 70B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama4-maverick-17b-instruct-v1:0\": {\n \"display_name\": \"Llama 4 Maverick 17B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.meta.llama4-scout-17b-instruct-v1:0\": {\n \"display_name\": \"Llama 4 Scout 17B Instruct\",\n \"model_vendor\": \"meta\",\n \"model_version\": \"1:0\"\n },\n \"us.mistral.pixtral-large-2502-v1:0\": {\n \"display_name\": \"Pixtral Large 25.02\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"1:0\"\n },\n \"us.twelvelabs.pegasus-1-2-v1:0\": {\n \"display_name\": \"Pegasus 1.2\",\n \"model_vendor\": \"twelvelabs\",\n \"model_version\": \"v1:0\"\n },\n \"vertex_ai/claude-3-5-sonnet\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-3-5-sonnet@20240620\": {\n \"display_name\": \"Claude Sonnet 3.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240620\"\n },\n \"vertex_ai/claude-3-sonnet\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-3-sonnet@20240229\": {\n \"display_name\": \"Claude Sonnet 3\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20240229\"\n },\n \"vertex_ai/claude-haiku-4-5\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-haiku-4-5@20251001\": {\n \"display_name\": \"Claude Haiku 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251001\"\n },\n \"vertex_ai/claude-opus-4\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-opus-4-1\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-opus-4-1@20250805\": {\n \"display_name\": \"Claude Opus 4.1\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250805\"\n },\n \"vertex_ai/claude-opus-4-5\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-opus-4-5@20251101\": {\n \"display_name\": \"Claude Opus 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20251101\"\n },\n \"vertex_ai/claude-opus-4@20250514\": {\n \"display_name\": \"Claude Opus 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"vertex_ai/claude-sonnet-4\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-sonnet-4-5\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\"\n },\n \"vertex_ai/claude-sonnet-4-5@20250929\": {\n \"display_name\": \"Claude Sonnet 4.5\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250929\"\n },\n \"vertex_ai/claude-sonnet-4@20250514\": {\n \"display_name\": \"Claude Sonnet 4\",\n \"model_vendor\": \"anthropic\",\n \"model_version\": \"20250514\"\n },\n \"vertex_ai/codestral-2\": {\n \"display_name\": \"Codestral 2\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/codestral-2501\": {\n \"display_name\": \"Codestral 25.01\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/codestral-2@001\": {\n \"display_name\": \"Codestral 2@001\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/codestral@2405\": {\n \"display_name\": \"Codestral@2405\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/codestral@latest\": {\n \"display_name\": \"Codestral@latest\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/deepseek-ai/deepseek-r1-0528-maas\": {\n \"display_name\": \"DeepSeek R1 0528 Maas\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/deepseek-ai/deepseek-v3.1-maas\": {\n \"display_name\": \"DeepSeek V3.1 Maas\",\n \"model_vendor\": \"deepseek\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/gemini-2.5-flash\": {\n \"display_name\": \"Gemini 2.5 Flash\",\n \"model_vendor\": \"google\"\n },\n \"vertex_ai/gemini-2.5-flash-lite\": {\n \"display_name\": \"Gemini 2.5 Flash Lite\",\n \"model_vendor\": \"google\"\n },\n \"vertex_ai/gemini-2.5-pro\": {\n \"display_name\": \"Gemini 2.5 Pro\",\n \"model_vendor\": \"google\"\n },\n \"vertex_ai/gemini-3-pro-preview\": {\n \"display_name\": \"Gemini 3 Pro Preview\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"vertex_ai/gemini-3-flash-preview\": {\n \"display_name\": \"Gemini 3 Flash Preview\",\n \"model_vendor\": \"google\",\n \"model_version\": \"preview\"\n },\n \"vertex_ai/jamba-1.5\": {\n \"display_name\": \"Jamba 1.5\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/jamba-1.5-large\": {\n \"display_name\": \"Jamba 1.5 Large\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/jamba-1.5-large@001\": {\n \"display_name\": \"Jamba 1.5 Large@001\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/jamba-1.5-mini\": {\n \"display_name\": \"Jamba 1.5 Mini\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/jamba-1.5-mini@001\": {\n \"display_name\": \"Jamba 1.5 Mini@001\",\n \"model_vendor\": \"ai21\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/meta/llama-3.1-405b-instruct-maas\": {\n \"display_name\": \"Llama 3.1 405B Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-3.1-70b-instruct-maas\": {\n \"display_name\": \"Llama 3.1 70B Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-3.1-8b-instruct-maas\": {\n \"display_name\": \"Llama 3.1 8B Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-3.2-90b-vision-instruct-maas\": {\n \"display_name\": \"Llama 3.2 90B Vision Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-4-maverick-17b-128e-instruct-maas\": {\n \"display_name\": \"Llama 4 Maverick 17B 128e Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-4-maverick-17b-16e-instruct-maas\": {\n \"display_name\": \"Llama 4 Maverick 17B 16e Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-4-scout-17b-128e-instruct-maas\": {\n \"display_name\": \"Llama 4 Scout 17B 128e Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama-4-scout-17b-16e-instruct-maas\": {\n \"display_name\": \"Llama 4 Scout 17B 16e Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama3-405b-instruct-maas\": {\n \"display_name\": \"Llama 3 405B Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama3-70b-instruct-maas\": {\n \"display_name\": \"Llama 3 70B Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/meta/llama3-8b-instruct-maas\": {\n \"display_name\": \"Llama 3 8B Instruct Maas\",\n \"model_vendor\": \"meta\"\n },\n \"vertex_ai/minimaxai/minimax-m2-maas\": {\n \"display_name\": \"MiniMax M2\",\n \"model_vendor\": \"minimax\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-large-2411\": {\n \"display_name\": \"Mistral Large 24.11\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-large@2407\": {\n \"display_name\": \"Mistral Large@24.07\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-large@2411-001\": {\n \"display_name\": \"Mistral Large@24.11 001\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"001\"\n },\n \"vertex_ai/mistral-large@latest\": {\n \"display_name\": \"Mistral Large@latest\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-medium-3\": {\n \"display_name\": \"Mistral Medium 3\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-medium-3@001\": {\n \"display_name\": \"Mistral Medium 3@001\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-nemo@2407\": {\n \"display_name\": \"Mistral Nemo@24.07\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-nemo@latest\": {\n \"display_name\": \"Mistral Nemo@latest\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-small-2503\": {\n \"display_name\": \"Mistral Small 2503\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistral-small-2503@001\": {\n \"display_name\": \"Mistral Small 2503@001\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistralai/codestral-2\": {\n \"display_name\": \"Codestral 2\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistralai/codestral-2@001\": {\n \"display_name\": \"Codestral 2@001\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistralai/mistral-medium-3\": {\n \"display_name\": \"Mistral Medium 3\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/mistralai/mistral-medium-3@001\": {\n \"display_name\": \"Mistral Medium 3@001\",\n \"model_vendor\": \"mistral\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/moonshotai/kimi-k2-thinking-maas\": {\n \"display_name\": \"Kimi K2 Thinking\",\n \"model_vendor\": \"moonshot\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/openai/gpt-oss-120b-maas\": {\n \"display_name\": \"GPT Open-Source 120B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/openai/gpt-oss-20b-maas\": {\n \"display_name\": \"GPT Open-Source 20B\",\n \"model_vendor\": \"openai\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas\": {\n \"display_name\": \"Qwen3 235B A22b Instruct 2507 Maas\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas\": {\n \"display_name\": \"Qwen3 Coder 480B A35b Instruct Maas\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas\": {\n \"display_name\": \"Qwen3 Next 80B A3b Instruct Maas\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas\": {\n \"display_name\": \"Qwen3 Next 80B A3b Thinking Maas\",\n \"model_vendor\": \"alibaba\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-2\": {\n \"display_name\": \"Grok 2\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-2-1212\": {\n \"display_name\": \"Grok 2\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-2-latest\": {\n \"display_name\": \"Grok 2\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-2-vision\": {\n \"display_name\": \"Grok 2 Vision\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-2-vision-1212\": {\n \"display_name\": \"Grok 2 Vision\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-2-vision-latest\": {\n \"display_name\": \"Grok 2 Vision\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-beta\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-fast-beta\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-fast-latest\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-latest\": {\n \"display_name\": \"Grok 3\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-mini\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-mini-beta\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-mini-fast\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-mini-fast-beta\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-mini-fast-latest\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-3-mini-latest\": {\n \"display_name\": \"Grok 3 Mini\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-0709\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-1-fast\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-1-fast-non-reasoning\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-1-fast-non-reasoning-latest\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-1-fast-reasoning\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-1-fast-reasoning-latest\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-fast-non-reasoning\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-fast-reasoning\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-4-latest\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-beta\": {\n \"display_name\": \"Grok Beta\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-code-fast\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-code-fast-1\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-code-fast-1-0825\": {\n \"display_name\": \"Grok\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n },\n \"xai/grok-vision-beta\": {\n \"display_name\": \"Grok Vision\",\n \"model_vendor\": \"xai\",\n \"model_version\": \"latest\"\n }\n}\n" }, { "path": "backend/onyx/llm/model_name_parser.py", "content": "\"\"\"\nLiteLLM Model Name Parser\n\nParses LiteLLM model strings and returns structured metadata for UI display.\nAll metadata comes from litellm's model_cost dictionary. Until this upstream patch to LiteLLM\nis merged (https://github.com/BerriAI/litellm/pull/17330), we use the model_metadata_enrichments.json\nto add these fields at server startup.\n\nEnrichment fields:\n- display_name: Human-friendly name (e.g., \"Claude 3.5 Sonnet\")\n- model_vendor: The company that made the model (anthropic, openai, meta, etc.)\n- model_version: Version string (e.g., \"20241022-v2:0\", \"v1:0\")\n\nThe parser only extracts provider and region from the model key - everything\nelse comes from enrichment.\n\"\"\"\n\nimport re\nfrom functools import lru_cache\n\nfrom pydantic import BaseModel\n\nfrom onyx.llm.constants import AGGREGATOR_PROVIDERS\nfrom onyx.llm.constants import HYPHENATED_MODEL_NAMES\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.constants import MODEL_PREFIX_TO_VENDOR\nfrom onyx.llm.constants import PROVIDER_DISPLAY_NAMES\nfrom onyx.llm.constants import VENDOR_BRAND_NAMES\n\n\nclass ParsedModelName(BaseModel):\n \"\"\"Structured representation of a parsed LiteLLM model name.\"\"\"\n\n raw_name: str # Original: \"bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0\"\n provider: str # \"bedrock\", \"azure\", \"openai\", etc. (the API route)\n vendor: str | None = None # From enrichment: \"anthropic\", \"openai\", \"meta\", etc.\n version: str | None = None # From enrichment: \"20241022-v2:0\", \"v1:0\", etc.\n region: str | None = None # Extracted: \"us\", \"eu\", or None\n display_name: str # From enrichment: \"Claude 3.5 Sonnet\"\n provider_display_name: str # Generated: \"Claude (Bedrock - Anthropic)\"\n\n\ndef _get_model_info(model_key: str) -> dict:\n \"\"\"Get model info from litellm.model_cost.\"\"\"\n from onyx.llm.litellm_singleton import litellm\n\n # Try exact key first\n info = litellm.model_cost.get(model_key)\n if info:\n return info\n\n # Try without provider prefix (e.g., \"bedrock/anthropic.claude-...\" -> \"anthropic.claude-...\")\n if \"/\" in model_key:\n return litellm.model_cost.get(model_key.split(\"/\", 1)[-1], {})\n\n return {}\n\n\ndef _extract_provider(model_key: str) -> str:\n \"\"\"Extract provider from model key prefix.\"\"\"\n from onyx.llm.litellm_singleton import litellm\n\n if \"/\" in model_key:\n return model_key.split(\"/\")[0]\n\n # No prefix - try to get from litellm.model_cost\n info = litellm.model_cost.get(model_key, {})\n litellm_provider = info.get(\"litellm_provider\", \"\")\n\n if litellm_provider:\n # Normalize vertex_ai variants\n if litellm_provider.startswith(LlmProviderNames.VERTEX_AI):\n return LlmProviderNames.VERTEX_AI\n return litellm_provider\n\n return \"unknown\"\n\n\ndef _extract_region(model_key: str) -> str | None:\n \"\"\"Extract region from model key (e.g., us., eu., apac. prefix).\"\"\"\n base = model_key.split(\"/\")[-1].lower()\n\n for prefix in [\"us.\", \"eu.\", \"apac.\", \"global.\", \"us-gov.\"]:\n if base.startswith(prefix):\n return prefix.rstrip(\".\")\n\n return None\n\n\ndef _format_name(name: str | None) -> str:\n \"\"\"Format provider or vendor name with proper capitalization.\"\"\"\n if not name:\n return \"Unknown\"\n return PROVIDER_DISPLAY_NAMES.get(name.lower(), name.replace(\"_\", \" \").title())\n\n\ndef _infer_vendor_from_model_name(model_name: str) -> str | None:\n \"\"\"\n Infer vendor from model name patterns when enrichment data is missing.\n\n Uses MODEL_PREFIX_TO_VENDOR mapping to match model name prefixes.\n Returns lowercase vendor name for consistency with enrichment data.\n\n Examples:\n \"gemini-3-flash-preview\" → \"google\"\n \"claude-3-5-sonnet\" → \"anthropic\"\n \"llama-3.1-70b\" → \"meta\"\n \"\"\"\n try:\n # Get the base model name (remove provider prefix if present)\n base_name = model_name.split(\"/\")[-1].lower()\n\n # Try to match against known prefixes (sorted by length to match longest first)\n for prefix in sorted(MODEL_PREFIX_TO_VENDOR.keys(), key=len, reverse=True):\n if base_name.startswith(prefix):\n return MODEL_PREFIX_TO_VENDOR[prefix]\n except Exception:\n pass\n\n return None\n\n\ndef _generate_display_name_from_model(model_name: str) -> str:\n \"\"\"\n Generate a human-friendly display name from a model identifier.\n\n Used as fallback when the model is not in enrichment data.\n Cleans up the raw model name by removing provider prefixes and\n formatting version numbers nicely.\n\n Examples:\n \"vertex_ai/gemini-3-flash-preview\" → \"Gemini 3 Flash Preview\"\n \"gemini-2.5-pro-exp-03-25\" → \"Gemini 2.5 Pro\"\n \"claude-3-5-sonnet-20241022\" → \"Claude 3.5 Sonnet\"\n \"gpt-oss:120b\" → \"GPT-OSS 120B\" (hyphenated exception)\n \"\"\"\n try:\n # Remove provider prefix if present\n base_name = model_name.split(\"/\")[-1]\n\n # Remove tag suffix (e.g., :14b, :latest) - handle separately\n size_suffix = \"\"\n if \":\" in base_name:\n base_name, tag = base_name.rsplit(\":\", 1)\n # Keep size tags like \"14b\", \"70b\", \"120b\"\n if re.match(r\"^\\d+[bBmM]$\", tag):\n size_suffix = f\" {tag.upper()}\"\n\n # Check if this is a hyphenated model that should keep its format\n base_name_lower = base_name.lower()\n for hyphenated in HYPHENATED_MODEL_NAMES:\n if base_name_lower.startswith(hyphenated):\n # Keep the hyphenated prefix, uppercase it\n return hyphenated.upper() + size_suffix\n\n # Remove common suffixes: date stamps, version numbers\n cleaned = base_name\n # Remove date stamps like -20241022, @20250219, -2024-08-06\n cleaned = re.sub(r\"[-@]\\d{4}-?\\d{2}-?\\d{2}\", \"\", cleaned)\n # Remove experimental/preview date suffixes like -exp-03-25\n cleaned = re.sub(r\"-exp-\\d{2}-\\d{2}\", \"\", cleaned)\n # Remove version suffixes like -v1, -v2\n cleaned = re.sub(r\"-v\\d+$\", \"\", cleaned)\n\n # Convert separators to spaces\n cleaned = cleaned.replace(\"-\", \" \").replace(\"_\", \" \")\n\n # Clean up version numbers: \"3 5\" → \"3.5\", \"2 5\" → \"2.5\"\n # But only for single digits that look like version numbers\n cleaned = re.sub(r\"(\\d) (\\d)(?!\\d)\", r\"\\1.\\2\", cleaned)\n\n # Title case each word, preserving version numbers\n words = cleaned.split()\n result_words = []\n for word in words:\n if word.isdigit() or re.match(r\"^\\d+\\.?\\d*$\", word):\n # Keep numbers as-is\n result_words.append(word)\n elif word.lower() in (\"pro\", \"lite\", \"mini\", \"flash\", \"preview\", \"ultra\"):\n # Common suffixes get title case\n result_words.append(word.title())\n else:\n # Title case other words\n result_words.append(word.title())\n\n return \" \".join(result_words) + size_suffix\n except Exception:\n return model_name\n\n\ndef _generate_provider_display_name(provider: str, vendor: str | None) -> str:\n \"\"\"\n Generate provider display name with model brand and vendor info.\n\n Examples:\n - Direct OpenAI: \"GPT (OpenAI)\"\n - Bedrock via Anthropic: \"Claude (Bedrock - Anthropic)\"\n - Vertex AI via Google: \"Gemini (Vertex AI - Google)\"\n \"\"\"\n provider_nice = _format_name(provider)\n vendor_nice = _format_name(vendor) if vendor else None\n brand = VENDOR_BRAND_NAMES.get(vendor.lower()) if vendor else None\n\n # For aggregator providers, show: Brand (Provider - Vendor)\n if provider.lower() in AGGREGATOR_PROVIDERS:\n if brand and vendor_nice:\n return f\"{brand} ({provider_nice} - {vendor_nice})\"\n elif vendor_nice:\n return f\"{provider_nice} - {vendor_nice}\"\n return provider_nice\n\n # For direct providers, show: Brand (Provider)\n if brand:\n return f\"{brand} ({provider_nice})\"\n\n return provider_nice\n\n\n@lru_cache(maxsize=1024)\ndef parse_litellm_model_name(raw_name: str) -> ParsedModelName:\n \"\"\"\n Parse a LiteLLM model string into structured data.\n\n Metadata comes from enrichment when available, with fallback logic\n for models not in the enrichment data.\n\n Args:\n raw_name: The LiteLLM model string\n\n Returns:\n ParsedModelName with all components from enrichment or fallback\n \"\"\"\n model_info = _get_model_info(raw_name)\n\n # Extract from key (not in enrichment)\n provider = _extract_provider(raw_name)\n region = _extract_region(raw_name)\n\n # Get from enrichment, with fallbacks for unenriched models\n vendor = model_info.get(\"model_vendor\") or _infer_vendor_from_model_name(raw_name)\n version = model_info.get(\"model_version\")\n display_name = model_info.get(\"display_name\") or _generate_display_name_from_model(\n raw_name\n )\n\n # Generate provider display name\n provider_display_name = _generate_provider_display_name(provider, vendor)\n\n return ParsedModelName(\n raw_name=raw_name,\n provider=provider,\n vendor=vendor,\n version=version,\n region=region,\n display_name=display_name,\n provider_display_name=provider_display_name,\n )\n" }, { "path": "backend/onyx/llm/model_response.py", "content": "from __future__ import annotations\n\nfrom typing import Any\nfrom typing import List\nfrom typing import TYPE_CHECKING\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\n\nclass FunctionCall(BaseModel):\n arguments: str | None = None\n name: str | None = None\n\n\nclass ChatCompletionMessageToolCall(BaseModel):\n id: str\n type: str = \"function\"\n function: FunctionCall\n\n\nclass ChatCompletionDeltaToolCall(BaseModel):\n id: str | None = None\n index: int = 0\n type: str = \"function\"\n function: FunctionCall | None = None\n\n\nclass Delta(BaseModel):\n content: str | None = None\n reasoning_content: str | None = None\n tool_calls: List[ChatCompletionDeltaToolCall] = Field(default_factory=list)\n\n\nclass StreamingChoice(BaseModel):\n finish_reason: str | None = None\n index: int = 0\n delta: Delta = Field(default_factory=Delta)\n\n\nclass Usage(BaseModel):\n completion_tokens: int\n prompt_tokens: int\n total_tokens: int\n cache_creation_input_tokens: int\n cache_read_input_tokens: int\n\n\nclass ModelResponseStream(BaseModel):\n id: str\n created: str\n choice: StreamingChoice\n usage: Usage | None = None\n\n\nif TYPE_CHECKING:\n from litellm.types.utils import ModelResponseStream as LiteLLMModelResponseStream\n\n\nclass Message(BaseModel):\n content: str | None = None\n role: str = \"assistant\"\n tool_calls: List[ChatCompletionMessageToolCall] | None = None\n reasoning_content: str | None = None\n\n\nclass Choice(BaseModel):\n finish_reason: str | None = None\n index: int = 0\n message: Message = Field(default_factory=Message)\n\n\nclass ModelResponse(BaseModel):\n id: str\n created: str\n choice: Choice\n usage: Usage | None = None\n\n\nif TYPE_CHECKING:\n from litellm.types.utils import (\n ModelResponse as LiteLLMModelResponse,\n ModelResponseStream as LiteLLMModelResponseStream,\n )\n\n\ndef _parse_function_call(\n function_payload: dict[str, Any] | None,\n) -> FunctionCall | None:\n \"\"\"Parse a function call payload into a FunctionCall object.\"\"\"\n if not function_payload or not isinstance(function_payload, dict):\n return None\n return FunctionCall(\n arguments=function_payload.get(\"arguments\"),\n name=function_payload.get(\"name\"),\n )\n\n\ndef _parse_delta_tool_calls(\n tool_calls: list[dict[str, Any]] | None,\n) -> list[ChatCompletionDeltaToolCall]:\n \"\"\"Parse tool calls for streaming responses (delta format).\"\"\"\n if not tool_calls:\n return []\n\n parsed_tool_calls: list[ChatCompletionDeltaToolCall] = []\n for tool_call in tool_calls:\n parsed_tool_calls.append(\n ChatCompletionDeltaToolCall(\n id=tool_call.get(\"id\"),\n index=tool_call.get(\"index\", 0),\n type=tool_call.get(\"type\", \"function\"),\n function=_parse_function_call(tool_call.get(\"function\")),\n )\n )\n return parsed_tool_calls\n\n\ndef _parse_message_tool_calls(\n tool_calls: list[dict[str, Any]] | None,\n) -> list[ChatCompletionMessageToolCall]:\n \"\"\"Parse tool calls for non-streaming responses (message format).\"\"\"\n if not tool_calls:\n return []\n\n parsed_tool_calls: list[ChatCompletionMessageToolCall] = []\n for tool_call in tool_calls:\n function_call = _parse_function_call(tool_call.get(\"function\"))\n if not function_call:\n continue\n\n parsed_tool_calls.append(\n ChatCompletionMessageToolCall(\n id=tool_call.get(\"id\", \"\"),\n type=tool_call.get(\"type\", \"function\"),\n function=function_call,\n )\n )\n return parsed_tool_calls\n\n\ndef _validate_and_extract_base_fields(\n response_data: dict[str, Any], error_prefix: str\n) -> tuple[str, str, dict[str, Any]]:\n \"\"\"\n Validate and extract common fields (id, created, first choice) from a LiteLLM response.\n\n Returns:\n Tuple of (id, created, choice_data)\n \"\"\"\n response_id = response_data.get(\"id\")\n created = response_data.get(\"created\")\n if response_id is None or created is None:\n raise ValueError(f\"{error_prefix} must include 'id' and 'created'.\")\n\n choices: list[dict[str, Any]] = response_data.get(\"choices\") or []\n if not choices:\n raise ValueError(f\"{error_prefix} must include at least one choice.\")\n\n return str(response_id), str(created), choices[0] or {}\n\n\ndef _usage_from_usage_data(usage_data: dict[str, Any]) -> Usage:\n # NOTE: sometimes the usage data dictionary has these keys and the values are None\n # hence the \"or 0\" instead of just using default values\n return Usage(\n completion_tokens=usage_data.get(\"completion_tokens\") or 0,\n prompt_tokens=usage_data.get(\"prompt_tokens\") or 0,\n total_tokens=usage_data.get(\"total_tokens\") or 0,\n cache_creation_input_tokens=usage_data.get(\"cache_creation_input_tokens\") or 0,\n cache_read_input_tokens=usage_data.get(\n \"cache_read_input_tokens\",\n (usage_data.get(\"prompt_tokens_details\") or {}).get(\"cached_tokens\"),\n )\n or 0,\n )\n\n\ndef from_litellm_model_response_stream(\n response: \"LiteLLMModelResponseStream\",\n) -> ModelResponseStream:\n \"\"\"\n Convert a LiteLLM ModelResponseStream into the simplified Onyx representation.\n \"\"\"\n response_data = response.model_dump()\n response_id, created, choice_data = _validate_and_extract_base_fields(\n response_data, \"LiteLLM response stream\"\n )\n\n delta_data: dict[str, Any] = choice_data.get(\"delta\") or {}\n parsed_delta = Delta(\n content=delta_data.get(\"content\"),\n reasoning_content=delta_data.get(\"reasoning_content\"),\n tool_calls=_parse_delta_tool_calls(delta_data.get(\"tool_calls\")),\n )\n\n streaming_choice = StreamingChoice(\n finish_reason=choice_data.get(\"finish_reason\"),\n index=choice_data.get(\"index\", 0),\n delta=parsed_delta,\n )\n\n usage_data = response_data.get(\"usage\")\n return ModelResponseStream(\n id=response_id,\n created=created,\n choice=streaming_choice,\n usage=(_usage_from_usage_data(usage_data) if usage_data else None),\n )\n\n\ndef from_litellm_model_response(\n response: \"LiteLLMModelResponse\",\n) -> ModelResponse:\n \"\"\"\n Convert a LiteLLM ModelResponse into the simplified Onyx representation.\n \"\"\"\n response_data = response.model_dump()\n response_id, created, choice_data = _validate_and_extract_base_fields(\n response_data, \"LiteLLM response\"\n )\n\n message_data: dict[str, Any] = choice_data.get(\"message\") or {}\n parsed_tool_calls = _parse_message_tool_calls(message_data.get(\"tool_calls\"))\n\n message = Message(\n content=message_data.get(\"content\"),\n role=message_data.get(\"role\", \"assistant\"),\n tool_calls=parsed_tool_calls if parsed_tool_calls else None,\n reasoning_content=message_data.get(\"reasoning_content\"),\n )\n\n choice = Choice(\n finish_reason=choice_data.get(\"finish_reason\"),\n index=choice_data.get(\"index\", 0),\n message=message,\n )\n\n usage_data = response_data.get(\"usage\")\n return ModelResponse(\n id=response_id,\n created=created,\n choice=choice,\n usage=(_usage_from_usage_data(usage_data) if usage_data else None),\n )\n" }, { "path": "backend/onyx/llm/models.py", "content": "from enum import Enum\nfrom typing import Literal\n\nfrom pydantic import BaseModel\n\n\nclass ToolChoiceOptions(str, Enum):\n REQUIRED = \"required\"\n AUTO = \"auto\"\n NONE = \"none\"\n\n\nclass ReasoningEffort(str, Enum):\n \"\"\"Reasoning effort levels for models that support extended thinking.\n\n Different providers map these values differently:\n - OpenAI: Uses \"low\", \"medium\", \"high\" directly for reasoning_effort. Recently added \"none\" for 5 series\n which is like \"minimal\"\n - Claude: Uses budget_tokens with different values for each level\n - Gemini: Uses \"none\", \"low\", \"medium\", \"high\" for thinking_budget (via litellm mapping)\n \"\"\"\n\n AUTO = \"auto\"\n OFF = \"off\"\n LOW = \"low\"\n MEDIUM = \"medium\"\n HIGH = \"high\"\n\n\n# OpenAI reasoning effort mapping\n# Note: OpenAI API does not support \"auto\" - valid values are: none, minimal, low, medium, high, xhigh\nOPENAI_REASONING_EFFORT: dict[ReasoningEffort, str] = {\n ReasoningEffort.AUTO: \"medium\", # Default to medium when auto is requested\n ReasoningEffort.OFF: \"none\",\n ReasoningEffort.LOW: \"low\",\n ReasoningEffort.MEDIUM: \"medium\",\n ReasoningEffort.HIGH: \"high\",\n}\n\n# Anthropic reasoning effort to budget tokens mapping\n# Loosely based on budgets from LiteLLM but this ensures it's not updated without our knowing from a version bump.\nANTHROPIC_REASONING_EFFORT_BUDGET: dict[ReasoningEffort, int] = {\n ReasoningEffort.AUTO: 2048,\n ReasoningEffort.LOW: 1024,\n ReasoningEffort.MEDIUM: 2048,\n ReasoningEffort.HIGH: 4096,\n}\n\n\n# Content part structures for multimodal messages\n# The classes in this mirror the OpenAI Chat Completions message types and work well with routers like LiteLLM\nclass TextContentPart(BaseModel):\n type: Literal[\"text\"] = \"text\"\n text: str\n # Some providers (e.g. Anthropic/Gemini) support prompt caching controls on content blocks.\n cache_control: dict | None = None\n\n\nclass ImageUrlDetail(BaseModel):\n url: str\n detail: Literal[\"auto\", \"low\", \"high\"] | None = None\n\n\nclass ImageContentPart(BaseModel):\n type: Literal[\"image_url\"] = \"image_url\"\n image_url: ImageUrlDetail\n\n\nContentPart = TextContentPart | ImageContentPart\n\n\n# Tool call structures\nclass FunctionCall(BaseModel):\n name: str\n arguments: str\n\n\nclass ToolCall(BaseModel):\n type: Literal[\"function\"] = \"function\"\n id: str\n function: FunctionCall\n\n\n# Message types\n\n\n# Base class for all cacheable messages\nclass CacheableMessage(BaseModel):\n # Some providers support prompt caching controls at the message level (passed through via LiteLLM).\n cache_control: dict | None = None\n\n\nclass SystemMessage(CacheableMessage):\n role: Literal[\"system\"] = \"system\"\n content: str\n\n\nclass UserMessage(CacheableMessage):\n role: Literal[\"user\"] = \"user\"\n content: str | list[ContentPart]\n\n\nclass AssistantMessage(CacheableMessage):\n role: Literal[\"assistant\"] = \"assistant\"\n content: str | None = None\n tool_calls: list[ToolCall] | None = None\n\n\nclass ToolMessage(CacheableMessage):\n role: Literal[\"tool\"] = \"tool\"\n content: str\n tool_call_id: str\n\n\n# Union type for all OpenAI Chat Completions messages\nChatCompletionMessage = SystemMessage | UserMessage | AssistantMessage | ToolMessage\n# Allows for passing in a string directly. This is provided for convenience and is wrapped as a UserMessage.\nLanguageModelInput = list[ChatCompletionMessage] | ChatCompletionMessage\n" }, { "path": "backend/onyx/llm/multi_llm.py", "content": "import os\nimport threading\nfrom collections.abc import Iterator\nfrom contextlib import contextmanager\nfrom contextlib import nullcontext\nfrom typing import Any\nfrom typing import cast\nfrom typing import TYPE_CHECKING\nfrom typing import Union\n\nfrom onyx.configs.app_configs import MOCK_LLM_RESPONSE\nfrom onyx.configs.chat_configs import LLM_SOCKET_READ_TIMEOUT\nfrom onyx.configs.model_configs import GEN_AI_TEMPERATURE\nfrom onyx.configs.model_configs import LITELLM_EXTRA_BODY\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.cost import calculate_llm_cost_cents\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMConfig\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.interfaces import ReasoningEffort\nfrom onyx.llm.interfaces import ToolChoiceOptions\nfrom onyx.llm.model_response import ModelResponse\nfrom onyx.llm.model_response import ModelResponseStream\nfrom onyx.llm.model_response import Usage\nfrom onyx.llm.models import ANTHROPIC_REASONING_EFFORT_BUDGET\nfrom onyx.llm.models import OPENAI_REASONING_EFFORT\nfrom onyx.llm.request_context import get_llm_mock_response\nfrom onyx.llm.utils import build_litellm_passthrough_kwargs\nfrom onyx.llm.utils import is_true_openai_model\nfrom onyx.llm.utils import model_is_reasoning_model\nfrom onyx.llm.well_known_providers.constants import AWS_ACCESS_KEY_ID_KWARG\nfrom onyx.llm.well_known_providers.constants import (\n AWS_ACCESS_KEY_ID_KWARG_ENV_VAR_FORMAT,\n)\nfrom onyx.llm.well_known_providers.constants import (\n AWS_BEARER_TOKEN_BEDROCK_KWARG_ENV_VAR_FORMAT,\n)\nfrom onyx.llm.well_known_providers.constants import AWS_REGION_NAME_KWARG\nfrom onyx.llm.well_known_providers.constants import AWS_REGION_NAME_KWARG_ENV_VAR_FORMAT\nfrom onyx.llm.well_known_providers.constants import AWS_SECRET_ACCESS_KEY_KWARG\nfrom onyx.llm.well_known_providers.constants import (\n AWS_SECRET_ACCESS_KEY_KWARG_ENV_VAR_FORMAT,\n)\nfrom onyx.llm.well_known_providers.constants import LM_STUDIO_API_KEY_CONFIG_KEY\nfrom onyx.llm.well_known_providers.constants import OLLAMA_API_KEY_CONFIG_KEY\nfrom onyx.llm.well_known_providers.constants import VERTEX_CREDENTIALS_FILE_KWARG\nfrom onyx.llm.well_known_providers.constants import (\n VERTEX_CREDENTIALS_FILE_KWARG_ENV_VAR_FORMAT,\n)\nfrom onyx.llm.well_known_providers.constants import VERTEX_LOCATION_KWARG\nfrom onyx.utils.encryption import mask_string\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_env_lock = threading.Lock()\n\nif TYPE_CHECKING:\n from litellm import CustomStreamWrapper\n from litellm import HTTPHandler\n\n\n_LLM_PROMPT_LONG_TERM_LOG_CATEGORY = \"llm_prompt\"\nLEGACY_MAX_TOKENS_KWARG = \"max_tokens\"\nSTANDARD_MAX_TOKENS_KWARG = \"max_completion_tokens\"\n_VERTEX_ANTHROPIC_MODELS_REJECTING_OUTPUT_CONFIG = (\n \"claude-opus-4-5\",\n \"claude-opus-4-6\",\n)\n\n\nclass LLMTimeoutError(Exception):\n \"\"\"\n Exception raised when an LLM call times out.\n \"\"\"\n\n\nclass LLMRateLimitError(Exception):\n \"\"\"\n Exception raised when an LLM call is rate limited.\n \"\"\"\n\n\ndef _prompt_to_dicts(prompt: LanguageModelInput) -> list[dict[str, Any]]:\n \"\"\"Convert Pydantic message models to dictionaries for LiteLLM.\n\n LiteLLM expects messages to be dictionaries (with .get() method),\n not Pydantic models. This function serializes the messages.\n \"\"\"\n if isinstance(prompt, list):\n return [msg.model_dump(exclude_none=True) for msg in prompt]\n return [prompt.model_dump(exclude_none=True)]\n\n\ndef _normalize_content(raw: Any) -> str:\n \"\"\"Normalize a message content field to a plain string.\n\n Content can be a string, None, or a list of content-block dicts\n (e.g. [{\"type\": \"text\", \"text\": \"...\"}]).\n \"\"\"\n if raw is None:\n return \"\"\n if isinstance(raw, str):\n return raw\n if isinstance(raw, list):\n return \"\\n\".join(\n block.get(\"text\", \"\") if isinstance(block, dict) else str(block)\n for block in raw\n )\n return str(raw)\n\n\ndef _strip_tool_content_from_messages(\n messages: list[dict[str, Any]],\n) -> list[dict[str, Any]]:\n \"\"\"Convert tool-related messages to plain text.\n\n Bedrock's Converse API requires toolConfig when messages contain\n toolUse/toolResult content blocks. When no tools are provided for the\n current request, we must convert any tool-related history into plain text\n to avoid the \"toolConfig field must be defined\" error.\n\n This is the same approach used by _OllamaHistoryMessageFormatter.\n \"\"\"\n result: list[dict[str, Any]] = []\n for msg in messages:\n role = msg.get(\"role\")\n tool_calls = msg.get(\"tool_calls\")\n\n if role == \"assistant\" and tool_calls:\n # Convert structured tool calls to text representation\n tool_call_lines = []\n for tc in tool_calls:\n func = tc.get(\"function\", {})\n name = func.get(\"name\", \"unknown\")\n args = func.get(\"arguments\", \"{}\")\n tc_id = tc.get(\"id\", \"\")\n tool_call_lines.append(\n f\"[Tool Call] name={name} id={tc_id} args={args}\"\n )\n\n existing_content = _normalize_content(msg.get(\"content\"))\n parts = (\n [existing_content] + tool_call_lines\n if existing_content\n else tool_call_lines\n )\n new_msg = {\n \"role\": \"assistant\",\n \"content\": \"\\n\".join(parts),\n }\n result.append(new_msg)\n\n elif role == \"tool\":\n # Convert tool response to user message with text content\n tool_call_id = msg.get(\"tool_call_id\", \"\")\n content = _normalize_content(msg.get(\"content\"))\n tool_result_text = f\"[Tool Result] id={tool_call_id}\\n{content}\"\n # Merge into previous user message if it is also a converted\n # tool result to avoid consecutive user messages (Bedrock requires\n # strict user/assistant alternation).\n if (\n result\n and result[-1][\"role\"] == \"user\"\n and \"[Tool Result]\" in result[-1].get(\"content\", \"\")\n ):\n result[-1][\"content\"] += \"\\n\\n\" + tool_result_text\n else:\n result.append({\"role\": \"user\", \"content\": tool_result_text})\n\n else:\n result.append(msg)\n\n return result\n\n\ndef _fix_tool_user_message_ordering(\n messages: list[dict[str, Any]],\n) -> list[dict[str, Any]]:\n \"\"\"Insert a synthetic assistant message between tool and user messages.\n\n Some models (e.g. Mistral on Azure) require strict message ordering where\n a user message cannot immediately follow a tool message. This function\n inserts a minimal assistant message to bridge the gap.\n \"\"\"\n if len(messages) < 2:\n return messages\n\n result: list[dict[str, Any]] = [messages[0]]\n for msg in messages[1:]:\n prev_role = result[-1].get(\"role\")\n curr_role = msg.get(\"role\")\n if prev_role == \"tool\" and curr_role == \"user\":\n result.append({\"role\": \"assistant\", \"content\": \"Noted. Continuing.\"})\n result.append(msg)\n return result\n\n\ndef _messages_contain_tool_content(messages: list[dict[str, Any]]) -> bool:\n \"\"\"Check if any messages contain tool-related content blocks.\"\"\"\n for msg in messages:\n if msg.get(\"role\") == \"tool\":\n return True\n if msg.get(\"role\") == \"assistant\" and msg.get(\"tool_calls\"):\n return True\n return False\n\n\ndef _prompt_contains_tool_call_history(prompt: LanguageModelInput) -> bool:\n \"\"\"Check if the prompt contains any assistant messages with tool_calls.\n\n When Anthropic's extended thinking is enabled, the API requires every\n assistant message to start with a thinking block before any tool_use\n blocks. Since we don't preserve thinking_blocks (they carry\n cryptographic signatures that can't be reconstructed), we must skip\n the thinking param whenever history contains prior tool-calling turns.\n \"\"\"\n from onyx.llm.models import AssistantMessage\n\n msgs = prompt if isinstance(prompt, list) else [prompt]\n return any(isinstance(msg, AssistantMessage) and msg.tool_calls for msg in msgs)\n\n\ndef _is_vertex_model_rejecting_output_config(model_name: str) -> bool:\n normalized_model_name = model_name.lower()\n return any(\n blocked_model in normalized_model_name\n for blocked_model in _VERTEX_ANTHROPIC_MODELS_REJECTING_OUTPUT_CONFIG\n )\n\n\nclass LitellmLLM(LLM):\n \"\"\"Uses Litellm library to allow easy configuration to use a multitude of LLMs\n See https://python.langchain.com/docs/integrations/chat/litellm\"\"\"\n\n def __init__(\n self,\n api_key: str | None,\n model_provider: str,\n model_name: str,\n max_input_tokens: int,\n timeout: int | None = None,\n api_base: str | None = None,\n api_version: str | None = None,\n deployment_name: str | None = None,\n custom_llm_provider: str | None = None,\n temperature: float | None = None,\n custom_config: dict[str, str] | None = None,\n extra_headers: dict[str, str] | None = None,\n extra_body: dict | None = LITELLM_EXTRA_BODY,\n model_kwargs: dict[str, Any] | None = None,\n ):\n # Timeout in seconds for each socket read operation (i.e., max time between\n # receiving data chunks/tokens). This is NOT a total request timeout - a\n # request can run indefinitely as long as data keeps arriving within this\n # window. If the LLM pauses for longer than this timeout between chunks,\n # a ReadTimeout is raised.\n self._timeout = timeout\n if timeout is None:\n self._timeout = LLM_SOCKET_READ_TIMEOUT\n\n self._temperature = GEN_AI_TEMPERATURE if temperature is None else temperature\n\n self._model_provider = model_provider\n self._model_version = model_name\n self._api_key = api_key\n self._deployment_name = deployment_name\n self._api_base = api_base\n self._api_version = api_version\n self._custom_llm_provider = custom_llm_provider\n self._max_input_tokens = max_input_tokens\n self._custom_config = custom_config\n\n # Create a dictionary for model-specific arguments if it's None\n model_kwargs = model_kwargs or {}\n\n if custom_config:\n for k, v in custom_config.items():\n if model_provider == LlmProviderNames.VERTEX_AI:\n if k == VERTEX_CREDENTIALS_FILE_KWARG:\n model_kwargs[k] = v\n elif k == VERTEX_CREDENTIALS_FILE_KWARG_ENV_VAR_FORMAT:\n model_kwargs[VERTEX_CREDENTIALS_FILE_KWARG] = v\n elif k == VERTEX_LOCATION_KWARG:\n model_kwargs[k] = v\n elif model_provider == LlmProviderNames.OLLAMA_CHAT:\n if k == OLLAMA_API_KEY_CONFIG_KEY:\n model_kwargs[\"api_key\"] = v\n elif model_provider == LlmProviderNames.LM_STUDIO:\n if k == LM_STUDIO_API_KEY_CONFIG_KEY:\n model_kwargs[\"api_key\"] = v\n elif model_provider == LlmProviderNames.BEDROCK:\n if k == AWS_REGION_NAME_KWARG:\n model_kwargs[k] = v\n elif k == AWS_REGION_NAME_KWARG_ENV_VAR_FORMAT:\n model_kwargs[AWS_REGION_NAME_KWARG] = v\n elif k == AWS_BEARER_TOKEN_BEDROCK_KWARG_ENV_VAR_FORMAT:\n model_kwargs[\"api_key\"] = v\n elif k == AWS_ACCESS_KEY_ID_KWARG:\n model_kwargs[k] = v\n elif k == AWS_ACCESS_KEY_ID_KWARG_ENV_VAR_FORMAT:\n model_kwargs[AWS_ACCESS_KEY_ID_KWARG] = v\n elif k == AWS_SECRET_ACCESS_KEY_KWARG:\n model_kwargs[k] = v\n elif k == AWS_SECRET_ACCESS_KEY_KWARG_ENV_VAR_FORMAT:\n model_kwargs[AWS_SECRET_ACCESS_KEY_KWARG] = v\n\n # LM Studio: LiteLLM defaults to \"fake-api-key\" when no key is provided,\n # which LM Studio rejects. Ensure we always pass an explicit key (or empty\n # string) to prevent LiteLLM from injecting its fake default.\n if model_provider == LlmProviderNames.LM_STUDIO:\n model_kwargs.setdefault(\"api_key\", \"\")\n\n # Users provide the server root (e.g. http://localhost:1234) but LiteLLM\n # needs /v1 for OpenAI-compatible calls.\n if self._api_base is not None:\n base = self._api_base.rstrip(\"/\")\n self._api_base = base if base.endswith(\"/v1\") else f\"{base}/v1\"\n model_kwargs[\"api_base\"] = self._api_base\n\n # Default vertex_location to \"global\" if not provided for Vertex AI\n # Latest gemini models are only available through the global region\n if (\n model_provider == LlmProviderNames.VERTEX_AI\n and VERTEX_LOCATION_KWARG not in model_kwargs\n ):\n model_kwargs[VERTEX_LOCATION_KWARG] = \"global\"\n\n # Bifrost: OpenAI-compatible proxy that expects model names in\n # provider/model format (e.g. \"anthropic/claude-sonnet-4-6\").\n # We route through LiteLLM's openai provider with the Bifrost base URL,\n # and ensure /v1 is appended.\n if model_provider == LlmProviderNames.BIFROST:\n self._custom_llm_provider = \"openai\"\n if self._api_base is not None:\n base = self._api_base.rstrip(\"/\")\n self._api_base = base if base.endswith(\"/v1\") else f\"{base}/v1\"\n model_kwargs[\"api_base\"] = self._api_base\n\n # This is needed for Ollama to do proper function calling\n if model_provider == LlmProviderNames.OLLAMA_CHAT and api_base is not None:\n model_kwargs[\"api_base\"] = api_base\n if extra_headers:\n model_kwargs.update({\"extra_headers\": extra_headers})\n if extra_body:\n model_kwargs.update({\"extra_body\": extra_body})\n\n self._model_kwargs = model_kwargs\n\n def _safe_model_config(self) -> dict:\n dump = self.config.model_dump()\n dump[\"api_key\"] = mask_string(dump.get(\"api_key\") or \"\")\n custom_config = dump.get(\"custom_config\")\n if isinstance(custom_config, dict):\n # Mask sensitive values in custom_config\n masked_config = {}\n for k, v in custom_config.items():\n masked_config[k] = mask_string(v) if v else v\n dump[\"custom_config\"] = masked_config\n return dump\n\n def _track_llm_cost(self, usage: Usage) -> None:\n \"\"\"\n Track LLM usage cost for Onyx-managed API keys.\n\n This is called after every LLM call completes (streaming or non-streaming).\n Cost is only tracked if:\n 1. Usage limits are enabled for this deployment\n 2. The API key is one of Onyx's managed default keys\n \"\"\"\n\n from onyx.server.usage_limits import is_usage_limits_enabled\n\n if not is_usage_limits_enabled():\n return\n\n from onyx.server.usage_limits import is_onyx_managed_api_key\n\n if not is_onyx_managed_api_key(self._api_key):\n return\n # Import here to avoid circular imports\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n from onyx.db.usage import increment_usage\n from onyx.db.usage import UsageType\n\n # Calculate cost in cents\n cost_cents = calculate_llm_cost_cents(\n model_name=self._model_version,\n prompt_tokens=usage.prompt_tokens,\n completion_tokens=usage.completion_tokens,\n )\n\n if cost_cents <= 0:\n return\n\n try:\n with get_session_with_current_tenant() as db_session:\n increment_usage(db_session, UsageType.LLM_COST, cost_cents)\n db_session.commit()\n except Exception as e:\n # Log but don't fail the LLM call if tracking fails\n logger.warning(f\"Failed to track LLM cost: {e}\")\n\n def _completion(\n self,\n prompt: LanguageModelInput,\n tools: list[dict] | None,\n tool_choice: ToolChoiceOptions | None,\n stream: bool,\n parallel_tool_calls: bool,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n structured_response_format: dict | None = None,\n timeout_override: int | None = None,\n max_tokens: int | None = None,\n user_identity: LLMUserIdentity | None = None,\n client: \"HTTPHandler | None\" = None,\n ) -> Union[\"ModelResponse\", \"CustomStreamWrapper\"]:\n # Lazy loading to avoid memory bloat for non-inference flows\n from onyx.llm.litellm_singleton import litellm\n from litellm.exceptions import Timeout, RateLimitError\n\n #########################\n # Flags that modify the final arguments\n #########################\n is_claude_model = \"claude\" in self.config.model_name.lower()\n is_reasoning = model_is_reasoning_model(\n self.config.model_name, self.config.model_provider\n )\n # All OpenAI models will use responses API for consistency\n # Responses API is needed to get reasoning packets from OpenAI models\n is_openai_model = is_true_openai_model(\n self.config.model_provider, self.config.model_name\n )\n is_ollama = self._model_provider == LlmProviderNames.OLLAMA_CHAT\n is_mistral = self._model_provider == LlmProviderNames.MISTRAL\n is_vertex_ai = self._model_provider == LlmProviderNames.VERTEX_AI\n # Some Vertex Anthropic models reject output_config.\n # Keep this guard until LiteLLM/Vertex accept the field for these models.\n is_vertex_model_rejecting_output_config = (\n is_vertex_ai\n and _is_vertex_model_rejecting_output_config(self.config.model_name)\n )\n\n #########################\n # Build arguments\n #########################\n # Optional kwargs - should only be passed to LiteLLM under certain conditions\n optional_kwargs: dict[str, Any] = {}\n\n # Model name\n is_bifrost = self._model_provider == LlmProviderNames.BIFROST\n model_provider = (\n f\"{self.config.model_provider}/responses\"\n if is_openai_model # Uses litellm's completions -> responses bridge\n else self.config.model_provider\n )\n if is_bifrost:\n # Bifrost expects model names in provider/model format\n # (e.g. \"anthropic/claude-sonnet-4-6\") sent directly to its\n # OpenAI-compatible endpoint. We use custom_llm_provider=\"openai\"\n # so LiteLLM doesn't try to route based on the provider prefix.\n model = self.config.deployment_name or self.config.model_name\n else:\n model = f\"{model_provider}/{self.config.deployment_name or self.config.model_name}\"\n\n # Tool choice\n if is_claude_model and tool_choice == ToolChoiceOptions.REQUIRED:\n # Claude models will not use reasoning if tool_choice is required\n # let it choose tools automatically so reasoning can still be used\n tool_choice = ToolChoiceOptions.AUTO\n\n # If no tools are provided, tool_choice should be None\n if not tools:\n tool_choice = None\n\n # Temperature\n temperature = 1 if is_reasoning else self._temperature\n\n if stream and not is_vertex_model_rejecting_output_config:\n optional_kwargs[\"stream_options\"] = {\"include_usage\": True}\n\n # Note, there is a reasoning_effort parameter in LiteLLM but it is completely jank and does not work for any\n # of the major providers. Not setting it sets it to OFF.\n if (\n is_reasoning\n # The default of this parameter not set is surprisingly not the equivalent of an Auto but is actually Off\n and reasoning_effort != ReasoningEffort.OFF\n and not is_vertex_model_rejecting_output_config\n ):\n if is_openai_model:\n # OpenAI API does not accept reasoning params for GPT 5 chat models\n # (neither reasoning nor reasoning_effort are accepted)\n # even though they are reasoning models (bug in OpenAI)\n if \"-chat\" not in model:\n optional_kwargs[\"reasoning\"] = {\n \"effort\": OPENAI_REASONING_EFFORT[reasoning_effort],\n \"summary\": \"auto\",\n }\n\n elif is_claude_model:\n budget_tokens: int | None = ANTHROPIC_REASONING_EFFORT_BUDGET.get(\n reasoning_effort\n )\n\n # Anthropic requires every assistant message with tool_use\n # blocks to start with a thinking block that carries a\n # cryptographic signature. We don't preserve those blocks\n # across turns, so skip thinking when the history already\n # contains tool-calling assistant messages. LiteLLM's\n # modify_params workaround doesn't cover all providers\n # (notably Bedrock).\n can_enable_thinking = (\n budget_tokens is not None\n and not _prompt_contains_tool_call_history(prompt)\n )\n\n if can_enable_thinking:\n assert budget_tokens is not None # mypy\n if max_tokens is not None:\n # Anthropic has a weird rule where max token has to be at least as much as budget tokens if set\n # and the minimum budget tokens is 1024\n # Will note that overwriting a developer set max tokens is not ideal but is the best we can do for now\n # It is better to allow the LLM to output more reasoning tokens even if it results in a fairly small tool\n # call as compared to reducing the budget for reasoning.\n max_tokens = max(budget_tokens + 1, max_tokens)\n optional_kwargs[\"thinking\"] = {\n \"type\": \"enabled\",\n \"budget_tokens\": budget_tokens,\n }\n\n # LiteLLM just does some mapping like this anyway but is incomplete for Anthropic\n optional_kwargs.pop(\"reasoning_effort\", None)\n\n else:\n # Hope for the best from LiteLLM\n if reasoning_effort in [\n ReasoningEffort.LOW,\n ReasoningEffort.MEDIUM,\n ReasoningEffort.HIGH,\n ]:\n optional_kwargs[\"reasoning_effort\"] = reasoning_effort.value\n else:\n optional_kwargs[\"reasoning_effort\"] = ReasoningEffort.MEDIUM.value\n\n if tools:\n # OpenAI will error if parallel_tool_calls is True and tools are not specified\n optional_kwargs[\"parallel_tool_calls\"] = parallel_tool_calls\n\n if structured_response_format:\n optional_kwargs[\"response_format\"] = structured_response_format\n\n if not (is_claude_model or is_ollama or is_mistral) or is_bifrost:\n # Litellm bug: tool_choice is dropped silently if not specified here for OpenAI\n # However, this param breaks Anthropic and Mistral models,\n # so it must be conditionally included unless the request is\n # routed through Bifrost's OpenAI-compatible endpoint.\n # Additionally, tool_choice is not supported by Ollama and causes warnings if included.\n # See also, https://github.com/ollama/ollama/issues/11171\n optional_kwargs[\"allowed_openai_params\"] = [\"tool_choice\"]\n\n # Passthrough kwargs\n passthrough_kwargs = build_litellm_passthrough_kwargs(\n model_kwargs=self._model_kwargs,\n user_identity=user_identity,\n )\n\n try:\n # NOTE: must pass in None instead of empty strings otherwise litellm\n # can have some issues with bedrock.\n # NOTE: Sometimes _model_kwargs may have an \"api_key\" kwarg\n # depending on what the caller passes in for custom_config. If it\n # does we allow it to clobber _api_key.\n if \"api_key\" not in passthrough_kwargs:\n passthrough_kwargs[\"api_key\"] = self._api_key or None\n\n # We only need to set environment variables if custom config is set\n env_ctx = (\n temporary_env_and_lock(self._custom_config)\n if self._custom_config\n else nullcontext()\n )\n with env_ctx:\n messages = _prompt_to_dicts(prompt)\n\n # Bedrock's Converse API requires toolConfig when messages\n # contain toolUse/toolResult content blocks. When no tools are\n # provided for this request but the history contains tool\n # content from previous turns, strip it to plain text.\n is_bedrock = self._model_provider in {\n LlmProviderNames.BEDROCK,\n LlmProviderNames.BEDROCK_CONVERSE,\n }\n if (\n is_bedrock\n and not tools\n and _messages_contain_tool_content(messages)\n ):\n messages = _strip_tool_content_from_messages(messages)\n\n # Some models (e.g. Mistral) reject a user message\n # immediately after a tool message. Insert a synthetic\n # assistant bridge message to satisfy the ordering\n # constraint. Check both the provider and the deployment/\n # model name to catch Mistral hosted on Azure.\n model_or_deployment = (\n self._deployment_name or self._model_version or \"\"\n ).lower()\n is_mistral_model = is_mistral or \"mistral\" in model_or_deployment\n if is_mistral_model:\n messages = _fix_tool_user_message_ordering(messages)\n\n # Only pass tool_choice when tools are present — some providers (e.g. Fireworks)\n # reject requests where tool_choice is explicitly null.\n if tools and tool_choice is not None:\n optional_kwargs[\"tool_choice\"] = tool_choice\n\n response = litellm.completion(\n mock_response=get_llm_mock_response() or MOCK_LLM_RESPONSE,\n model=model,\n base_url=self._api_base or None,\n api_version=self._api_version or None,\n custom_llm_provider=self._custom_llm_provider or None,\n messages=messages,\n tools=tools,\n stream=stream,\n temperature=temperature,\n timeout=timeout_override or self._timeout,\n max_tokens=max_tokens,\n client=client,\n **optional_kwargs,\n **passthrough_kwargs,\n )\n return response\n except Exception as e:\n # for break pointing\n if isinstance(e, Timeout):\n raise LLMTimeoutError(e)\n\n elif isinstance(e, RateLimitError):\n raise LLMRateLimitError(e)\n\n raise e\n\n @property\n def config(self) -> LLMConfig:\n return LLMConfig(\n model_provider=self._model_provider,\n model_name=self._model_version,\n temperature=self._temperature,\n api_key=self._api_key,\n api_base=self._api_base,\n api_version=self._api_version,\n deployment_name=self._deployment_name,\n custom_config=self._custom_config,\n max_input_tokens=self._max_input_tokens,\n )\n\n def invoke(\n self,\n prompt: LanguageModelInput,\n tools: list[dict] | None = None,\n tool_choice: ToolChoiceOptions | None = None,\n structured_response_format: dict | None = None,\n timeout_override: int | None = None,\n max_tokens: int | None = None,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n user_identity: LLMUserIdentity | None = None,\n ) -> ModelResponse:\n from litellm import HTTPHandler\n from litellm import ModelResponse as LiteLLMModelResponse\n\n from onyx.llm.model_response import from_litellm_model_response\n\n # HTTPHandler Threading & Connection Pool Notes:\n # =============================================\n # We create an isolated HTTPHandler ONLY for true OpenAI models (not OpenAI-compatible\n # providers like glm-4.7, DeepSeek, etc.). This distinction is critical:\n #\n # 1. WHY ONLY TRUE OPENAI MODELS:\n # - True OpenAI models use litellm's \"responses API\" path which expects HTTPHandler\n # - OpenAI-compatible providers (model_provider=\"openai\" with non-OpenAI models)\n # use the standard completion path which expects OpenAI SDK client objects\n # - Passing HTTPHandler to OpenAI-compatible providers causes:\n # AttributeError: 'HTTPHandler' object has no attribute 'api_key'\n # (because _get_openai_client() calls openai_client.api_key on line ~929)\n #\n # 2. WHY ISOLATED HTTPHandler FOR OPENAI:\n # - Prevents \"Bad file descriptor\" errors when multiple threads stream concurrently\n # - Shared connection pools can have stale connections or abandoned streams that\n # corrupt the pool state for other threads\n # - Each request gets its own fresh httpx.Client via HTTPHandler\n #\n # 3. WHY OTHER PROVIDERS DON'T NEED THIS:\n # - Other providers (Anthropic, Bedrock, etc.) use litellm.module_level_client\n # which handles concurrency appropriately\n # - httpx.Client itself IS thread-safe for concurrent requests\n # - The issue is specific to OpenAI's responses API path and connection reuse\n #\n # 4. PITFALL - is_true_openai_model() CHECK:\n # - Must use is_true_openai_model() NOT just check model_provider == \"openai\"\n # - Many OpenAI-compatible providers set model_provider=\"openai\" but are NOT true\n # OpenAI models (glm-4.7, DeepSeek, local proxies, etc.)\n # - is_true_openai_model() checks both provider AND model name patterns\n #\n # This note may not be entirely accurate as there is a lot of complexity in the LiteLLM codebase around this\n # and not every model path was traced thoroughly. It is also possible that in future versions of LiteLLM\n # they will realize that their OpenAI handling is not threadsafe. Hope they will just fix it.\n client = None\n if is_true_openai_model(self.config.model_provider, self.config.model_name):\n client = HTTPHandler(timeout=timeout_override or self._timeout)\n\n try:\n # When custom_config is set, env vars are temporarily injected\n # under a global lock. Using stream=True here means the lock is\n # only held during connection setup (not the full inference).\n # The chunks are then collected outside the lock and reassembled\n # into a single ModelResponse via stream_chunk_builder.\n from litellm import stream_chunk_builder\n from litellm import CustomStreamWrapper as LiteLLMCustomStreamWrapper\n\n stream_response = cast(\n LiteLLMCustomStreamWrapper,\n self._completion(\n prompt=prompt,\n tools=tools,\n tool_choice=tool_choice,\n stream=True,\n structured_response_format=structured_response_format,\n timeout_override=timeout_override,\n max_tokens=max_tokens,\n parallel_tool_calls=True,\n reasoning_effort=reasoning_effort,\n user_identity=user_identity,\n client=client,\n ),\n )\n chunks = list(stream_response)\n response = cast(\n LiteLLMModelResponse,\n stream_chunk_builder(chunks),\n )\n\n model_response = from_litellm_model_response(response)\n\n # Track LLM cost for Onyx-managed API keys\n if model_response.usage:\n self._track_llm_cost(model_response.usage)\n\n return model_response\n finally:\n if client is not None:\n client.close()\n\n def stream(\n self,\n prompt: LanguageModelInput,\n tools: list[dict] | None = None,\n tool_choice: ToolChoiceOptions | None = None,\n structured_response_format: dict | None = None,\n timeout_override: int | None = None,\n max_tokens: int | None = None,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n user_identity: LLMUserIdentity | None = None,\n ) -> Iterator[ModelResponseStream]:\n from litellm import CustomStreamWrapper as LiteLLMCustomStreamWrapper\n from litellm import HTTPHandler\n\n from onyx.llm.model_response import from_litellm_model_response_stream\n\n # HTTPHandler Threading & Connection Pool Notes:\n # =============================================\n # See invoke() method for full explanation. Key points for streaming:\n #\n # 1. SAME RESTRICTIONS APPLY:\n # - HTTPHandler ONLY for true OpenAI models (use is_true_openai_model())\n # - OpenAI-compatible providers will fail with AttributeError on api_key\n #\n # 2. STREAMING-SPECIFIC CONCERNS:\n # - \"Bad file descriptor\" errors are MORE common during streaming because:\n # a) Streams hold connections open longer, increasing conflict window\n # b) Multiple concurrent streams (e.g., deep research) share the pool\n # c) Abandoned/interrupted streams can leave connections in bad state\n #\n # 3. ABANDONED STREAM PITFALL:\n # - If callers abandon this generator without fully consuming it (e.g.,\n # early return, exception, or break), the finally block won't execute\n # until the generator is garbage collected\n # - This is acceptable because:\n # a) CPython's refcounting typically finalizes generators promptly\n # b) Each HTTPHandler has its own isolated connection pool\n # c) httpx has built-in connection timeouts as a fallback\n # - If abandoned streams become problematic, consider using contextlib\n # or explicit stream.close() at call sites\n #\n # 4. WHY NOT USE SHARED HTTPHandler:\n # - litellm's InMemoryCache (used for client caching) is NOT thread-safe\n # - Shared pools can have connections corrupted by other threads\n # - Per-request HTTPHandler eliminates cross-thread interference\n client = None\n if is_true_openai_model(self.config.model_provider, self.config.model_name):\n client = HTTPHandler(timeout=timeout_override or self._timeout)\n\n try:\n response = cast(\n LiteLLMCustomStreamWrapper,\n self._completion(\n prompt=prompt,\n tools=tools,\n tool_choice=tool_choice,\n stream=True,\n structured_response_format=structured_response_format,\n timeout_override=timeout_override,\n max_tokens=max_tokens,\n parallel_tool_calls=True,\n reasoning_effort=reasoning_effort,\n user_identity=user_identity,\n client=client,\n ),\n )\n\n for chunk in response:\n model_response = from_litellm_model_response_stream(chunk)\n\n # Track LLM cost when usage info is available (typically in the last chunk)\n if model_response.usage:\n self._track_llm_cost(model_response.usage)\n\n yield model_response\n finally:\n if client is not None:\n client.close()\n\n\n@contextmanager\ndef temporary_env_and_lock(env_variables: dict[str, str]) -> Iterator[None]:\n \"\"\"\n Temporarily sets the environment variables to the given values.\n Code path is locked while the environment variables are set.\n Then cleans up the environment and frees the lock.\n \"\"\"\n with _env_lock:\n logger.debug(\"Acquired lock in temporary_env_and_lock\")\n # Store original values (None if key didn't exist)\n original_values: dict[str, str | None] = {\n key: os.environ.get(key) for key in env_variables\n }\n try:\n os.environ.update(env_variables)\n yield\n finally:\n for key, original_value in original_values.items():\n if original_value is None:\n os.environ.pop(key, None) # Remove if it didn't exist before\n else:\n os.environ[key] = original_value # Restore original value\n\n logger.debug(\"Released lock in temporary_env_and_lock\")\n" }, { "path": "backend/onyx/llm/override_models.py", "content": "\"\"\"Overrides sent over the wire / stored in the DB\n\nNOTE: these models are used in many places, so have to be\nkepy in a separate file to avoid circular imports.\n\"\"\"\n\nfrom pydantic import BaseModel\n\n\nclass LLMOverride(BaseModel):\n \"\"\"Per-request LLM settings that override persona defaults.\n\n All fields are optional — only the fields that differ from the persona's\n configured LLM need to be supplied. Used both over the wire (API requests)\n and for multi-model comparison, where one override is supplied per model.\n\n Attributes:\n model_provider: LLM provider slug (e.g. ``\"openai\"``, ``\"anthropic\"``).\n When ``None``, the persona's default provider is used.\n model_version: Specific model version string (e.g. ``\"gpt-4o\"``).\n When ``None``, the persona's default model is used.\n temperature: Sampling temperature in ``[0, 2]``. When ``None``, the\n persona's default temperature is used.\n display_name: Human-readable label shown in the UI for this model,\n e.g. ``\"GPT-4 Turbo\"``. Optional; falls back to ``model_version``\n when not set.\n \"\"\"\n\n model_provider: str | None = None\n model_version: str | None = None\n temperature: float | None = None\n display_name: str | None = None\n\n # This disables the \"model_\" protected namespace for pydantic\n model_config = {\"protected_namespaces\": ()}\n\n\nclass PromptOverride(BaseModel):\n system_prompt: str | None = None\n task_prompt: str | None = None\n" }, { "path": "backend/onyx/llm/prompt_cache/README.md", "content": "# Prompt Caching Framework\n\nA comprehensive prompt-caching mechanism for enabling cost savings across multiple LLM providers by leveraging provider-side prompt token caching.\n\n## Overview\n\nThe prompt caching framework provides a unified interface for enabling prompt caching across different LLM providers. It supports both **implicit caching** (automatic provider-side caching) and **explicit caching** (with cache control parameters).\n\n## Features\n\n- **Provider Support**: OpenAI (implicit), Anthropic (explicit), Vertex AI (explicit)\n- **Flexible Input**: Supports both `str` and `Sequence[ChatCompletionMessage]` inputs\n- **Continuation Handling**: Smart merging of cacheable prefix and suffix messages\n- **Best-Effort**: Gracefully degrades if caching fails\n- **Tenant-Aware**: Automatic tenant isolation for multi-tenant deployments\n- **Configurable**: Enable/disable via environment variable\n\n## Quick Start\n\n### Basic Usage\n\n```python\nfrom onyx.llm.prompt_cache import process_with_prompt_cache\nfrom onyx.llm.models import SystemMessage, UserMessage\n\n# Assume you have an LLM instance with a config property\n# llm = get_your_llm_instance()\n\n# Define cacheable prefix (static context) using Pydantic message models\ncacheable_prefix = [\n SystemMessage(role=\"system\", content=\"You are a helpful assistant.\"),\n UserMessage(role=\"user\", content=\"Context: ...\") # Static context\n]\n\n# Define suffix (dynamic user input)\nsuffix = [UserMessage(role=\"user\", content=\"What is the weather?\")]\n\n# Process with caching - pass llm_config, not the llm instance\nprocessed_prompt, cache_metadata = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=False,\n)\n\n# Make LLM call with processed prompt\nresponse = llm.invoke(processed_prompt)\n```\n\n### Using String Inputs\n\n```python\n# Both prefix and suffix can be strings\ncacheable_prefix = \"You are a helpful assistant. Context: ...\"\nsuffix = \"What is the weather?\"\n\nprocessed_prompt, cache_metadata = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=False,\n)\n\nresponse = llm.invoke(processed_prompt)\n```\n\n### Continuation Flag\n\nWhen `continuation=True`, the suffix is appended to the last message of the cacheable prefix:\n\n```python\n# Without continuation (default)\n# Result: [system_msg, prefix_user_msg, suffix_user_msg]\n\n# With continuation=True\n# Result: [system_msg, prefix_user_msg + suffix_user_msg]\nprocessed_prompt, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=True, # Merge suffix into last prefix message\n)\n```\n\n**Note**: If `cacheable_prefix` is a string, it remains in its own content block even when `continuation=True`.\n\n## Provider-Specific Behavior\n\n### OpenAI\n- **Caching Type**: Implicit (automatic)\n- **Behavior**: No special parameters needed. Provider automatically caches prefixes >1024 tokens.\n- **Cache Lifetime**: Up to 1 hour\n- **Cost Savings**: 50% discount on cached tokens\n\n### Anthropic\n- **Caching Type**: Explicit (requires `cache_control` parameter)\n- **Behavior**: Automatically adds `cache_control={\"type\": \"ephemeral\"}` to the **last message** of the cacheable prefix\n- **Cache Lifetime**: 5 minutes (default)\n- **Limitations**: Supports up to 4 cache breakpoints\n\n### Vertex AI\n- **Caching Type**: Explicit (with `cache_control` parameter)\n- **Behavior**: Adds `cache_control={\"type\": \"ephemeral\"}` to **all content blocks** in cacheable messages. String content is converted to array format with the cache control attached.\n- **Cache Lifetime**: 5 minutes\n- **Future**: Full context caching with block number management (deferred to future PR)\n\n## Configuration\n\n### Environment Variables\n\n- `ENABLE_PROMPT_CACHING`: Enable/disable prompt caching (default: `true`)\n ```bash\n export ENABLE_PROMPT_CACHING=false # Disable caching\n ```\n\n## Architecture\n\n### Core Components\n\n1. **`processor.py`**: Main entry point (`process_with_prompt_cache`)\n2. **`cache_manager.py`**: Cache metadata storage and retrieval\n3. **`models.py`**: Pydantic models for cache metadata (`CacheMetadata`)\n4. **`providers/`**: Provider-specific adapters\n5. **`utils.py`**: Shared utility functions\n\n### Provider Adapters\n\nEach provider has its own adapter in `providers/`:\n\n| File | Class | Description |\n|------|-------|-------------|\n| `base.py` | `PromptCacheProvider` | Abstract base class for all providers |\n| `openai.py` | `OpenAIPromptCacheProvider` | Implicit caching (no transformation) |\n| `anthropic.py` | `AnthropicPromptCacheProvider` | Explicit caching with `cache_control` on last message |\n| `vertex.py` | `VertexAIPromptCacheProvider` | Explicit caching with `cache_control` on all content blocks |\n| `noop.py` | `NoOpPromptCacheProvider` | Fallback for unsupported providers |\n\nEach adapter implements:\n- `supports_caching()`: Whether caching is supported\n- `prepare_messages_for_caching()`: Transform messages for caching\n- `extract_cache_metadata()`: Extract metadata from responses\n- `get_cache_ttl_seconds()`: Cache TTL\n\n## Best Practices\n\n1. **Cache Static Content**: Use cacheable prefix for system prompts, static context, and instructions that don't change between requests.\n\n2. **Keep Dynamic Content in Suffix**: User queries, search results, and other dynamic content should be in the suffix.\n\n3. **Monitor Cache Effectiveness**: Check logs for cache hits/misses and adjust your caching strategy accordingly.\n\n4. **Provider Selection**: Different providers have different caching characteristics - choose based on your use case.\n\n## Error Handling\n\nThe framework is **best-effort** - if caching fails, it gracefully falls back to non-cached behavior:\n\n- Cache lookup failures: Logged and continue without caching\n- Provider adapter failures: Fall back to no-op adapter\n- Cache storage failures: Logged and continue (caching is best-effort)\n- Invalid cache metadata: Cleared and proceed without cache\n\n## Future Enhancements\n\n- **Explicit Caching for Vertex AI**: Full block number tracking and management\n- **Cache Analytics**: Detailed metrics on cache effectiveness and cost savings\n- **Advanced Strategies**: More sophisticated cache key generation and invalidation\n- **Distributed Caching**: Shared caches across instances\n\n## Examples\n\nSee `backend/tests/external_dependency_unit/llm/test_prompt_caching.py` for detailed integration test examples.\n" }, { "path": "backend/onyx/llm/prompt_cache/__init__.py", "content": "\"\"\"Prompt caching framework for LLM providers.\n\nThis module provides a framework for enabling prompt caching across different\nLLM providers. It supports both implicit caching (automatic provider-side caching)\nand explicit caching (with cache metadata management).\n\"\"\"\n\nfrom onyx.llm.prompt_cache.cache_manager import CacheManager\nfrom onyx.llm.prompt_cache.cache_manager import generate_cache_key_hash\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.llm.prompt_cache.processor import process_with_prompt_cache\nfrom onyx.llm.prompt_cache.providers.anthropic import AnthropicPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.factory import get_provider_adapter\nfrom onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.openai import OpenAIPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.vertex import VertexAIPromptCacheProvider\nfrom onyx.llm.prompt_cache.utils import combine_messages_with_continuation\nfrom onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform\n\n__all__ = [\n \"AnthropicPromptCacheProvider\",\n \"CacheManager\",\n \"CacheMetadata\",\n \"combine_messages_with_continuation\",\n \"generate_cache_key_hash\",\n \"get_provider_adapter\",\n \"NoOpPromptCacheProvider\",\n \"OpenAIPromptCacheProvider\",\n \"prepare_messages_with_cacheable_transform\",\n \"process_with_prompt_cache\",\n \"PromptCacheProvider\",\n \"VertexAIPromptCacheProvider\",\n]\n" }, { "path": "backend/onyx/llm/prompt_cache/cache_manager.py", "content": "\"\"\"Cache manager for storing and retrieving prompt cache metadata.\"\"\"\n\nimport hashlib\nimport json\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom onyx.configs.model_configs import PROMPT_CACHE_REDIS_TTL_MULTIPLIER\nfrom onyx.key_value_store.store import PgRedisKVStore\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nREDIS_KEY_PREFIX = \"prompt_cache:\"\n# Cache TTL multiplier - store caches slightly longer than provider TTL\n# This allows for some clock skew and ensures we don't lose cache metadata prematurely\n# Value is configurable via PROMPT_CACHE_REDIS_TTL_MULTIPLIER env var (default: 1.2)\nCACHE_TTL_MULTIPLIER = PROMPT_CACHE_REDIS_TTL_MULTIPLIER\n\n\nclass CacheManager:\n \"\"\"Manages storage and retrieval of prompt cache metadata.\"\"\"\n\n def __init__(self, kv_store: PgRedisKVStore | None = None) -> None:\n \"\"\"Initialize the cache manager.\n\n Args:\n kv_store: Optional key-value store. If None, creates a new PgRedisKVStore.\n \"\"\"\n self._kv_store = kv_store or PgRedisKVStore()\n\n def _build_cache_key(\n self,\n provider: str,\n model_name: str,\n cache_key_hash: str,\n tenant_id: str | None = None,\n ) -> str:\n \"\"\"Build a Redis/PostgreSQL key for cache metadata.\n\n Args:\n provider: LLM provider name (e.g., \"openai\", \"anthropic\")\n model_name: Model name\n cache_key_hash: Hash of the cacheable prefix content\n tenant_id: Tenant ID. If None, uses current tenant from context.\n\n Returns:\n Cache key string\n \"\"\"\n if tenant_id is None:\n tenant_id = get_current_tenant_id()\n return f\"{REDIS_KEY_PREFIX}{tenant_id}:{provider}:{model_name}:{cache_key_hash}\"\n\n def store_cache_metadata(\n self,\n metadata: CacheMetadata,\n ) -> None:\n \"\"\"Store cache metadata.\n\n Args:\n metadata: Cache metadata to store\n ttl_seconds: Optional TTL in seconds. If None, uses provider default.\n \"\"\"\n try:\n cache_key = self._build_cache_key(\n metadata.provider,\n metadata.model_name,\n metadata.cache_key,\n metadata.tenant_id,\n )\n\n # Update last_accessed timestamp\n metadata.last_accessed = datetime.now(timezone.utc)\n\n # Serialize metadata\n metadata_dict = metadata.model_dump(mode=\"json\")\n\n # Store in key-value store\n # Note: PgRedisKVStore doesn't support TTL directly, but Redis will\n # handle expiration. For PostgreSQL persistence, we rely on cleanup\n # based on last_accessed timestamp.\n self._kv_store.store(cache_key, metadata_dict, encrypt=False)\n\n logger.debug(\n f\"Stored cache metadata: provider={metadata.provider}, \"\n f\"model={metadata.model_name}, cache_key={metadata.cache_key[:16]}..., \"\n f\"tenant_id={metadata.tenant_id}\"\n )\n except Exception as e:\n # Best-effort: log and continue\n logger.warning(f\"Failed to store cache metadata: {str(e)}\")\n\n def retrieve_cache_metadata(\n self,\n provider: str,\n model_name: str,\n cache_key_hash: str,\n tenant_id: str | None = None,\n ) -> CacheMetadata | None:\n \"\"\"Retrieve cache metadata.\n\n Args:\n provider: LLM provider name\n model_name: Model name\n cache_key_hash: Hash of the cacheable prefix content\n tenant_id: Tenant ID. If None, uses current tenant from context.\n\n Returns:\n CacheMetadata if found, None otherwise\n \"\"\"\n try:\n cache_key = self._build_cache_key(\n provider, model_name, cache_key_hash, tenant_id\n )\n metadata_dict = self._kv_store.load(cache_key, refresh_cache=False)\n\n # Deserialize metadata\n metadata = CacheMetadata.model_validate(metadata_dict)\n\n # Update last_accessed timestamp\n metadata.last_accessed = datetime.now(timezone.utc)\n self.store_cache_metadata(metadata)\n\n logger.debug(\n f\"Retrieved cache metadata: provider={provider}, \"\n f\"model={model_name}, cache_key={cache_key_hash[:16]}..., \"\n f\"tenant_id={tenant_id}\"\n )\n return metadata\n except Exception as e:\n # Best-effort: log and continue\n logger.debug(f\"Cache metadata not found or error retrieving: {str(e)}\")\n return None\n\n def delete_cache_metadata(\n self,\n provider: str,\n model_name: str,\n cache_key_hash: str,\n tenant_id: str | None = None,\n ) -> None:\n \"\"\"Delete cache metadata.\n\n Args:\n provider: LLM provider name\n model_name: Model name\n cache_key_hash: Hash of the cacheable prefix content\n tenant_id: Tenant ID. If None, uses current tenant from context.\n \"\"\"\n try:\n cache_key = self._build_cache_key(\n provider, model_name, cache_key_hash, tenant_id\n )\n self._kv_store.delete(cache_key)\n logger.debug(\n f\"Deleted cache metadata for provider={provider}, model={model_name}, cache_key={cache_key_hash[:16]}...\"\n )\n except Exception as e:\n # Best-effort: log and continue\n logger.warning(f\"Failed to delete cache metadata: {str(e)}\")\n\n\ndef _make_json_serializable(obj: object) -> object:\n \"\"\"Recursively convert objects to JSON-serializable types.\n\n Handles Pydantic models, dicts, lists, and other common types.\n \"\"\"\n if hasattr(obj, \"model_dump\"):\n # Pydantic v2 model\n return obj.model_dump(mode=\"json\")\n elif hasattr(obj, \"dict\"):\n # Pydantic v1 model or similar\n return _make_json_serializable(obj.dict())\n elif isinstance(obj, dict):\n return {k: _make_json_serializable(v) for k, v in obj.items()}\n elif isinstance(obj, (list, tuple)):\n return [_make_json_serializable(item) for item in obj]\n elif isinstance(obj, (str, int, float, bool, type(None))):\n return obj\n else:\n # Fallback: convert to string representation\n return str(obj)\n\n\ndef generate_cache_key_hash(\n cacheable_prefix: LanguageModelInput,\n provider: str,\n model_name: str,\n tenant_id: str,\n) -> str:\n \"\"\"Generate a deterministic cache key hash from cacheable prefix.\n\n Args:\n cacheable_prefix: Single message or list of messages to hash\n provider: LLM provider name\n model_name: Model name\n tenant_id: Tenant ID\n\n Returns:\n SHA256 hash as hex string\n \"\"\"\n # Normalize to list for consistent hashing; _make_json_serializable handles Pydantic models\n messages = (\n cacheable_prefix if isinstance(cacheable_prefix, list) else [cacheable_prefix]\n )\n messages_dict = [_make_json_serializable(msg) for msg in messages]\n\n # Serialize messages in a deterministic way\n # Include only content, roles, and order - exclude timestamps or dynamic fields\n serialized = json.dumps(\n {\n \"messages\": messages_dict,\n \"provider\": provider,\n \"model\": model_name,\n \"tenant_id\": tenant_id,\n },\n sort_keys=True,\n separators=(\",\", \":\"),\n )\n return hashlib.sha256(serialized.encode(\"utf-8\")).hexdigest()\n" }, { "path": "backend/onyx/llm/prompt_cache/models.py", "content": "\"\"\"Interfaces and data structures for prompt caching.\"\"\"\n\nfrom datetime import datetime\n\nfrom pydantic import BaseModel\n\n\nclass CacheMetadata(BaseModel):\n \"\"\"Metadata for cached prompt prefixes.\"\"\"\n\n cache_key: str\n provider: str\n model_name: str\n tenant_id: str\n created_at: datetime\n last_accessed: datetime\n # Provider-specific metadata\n # TODO: Add explicit caching support in future PR\n # vertex_block_numbers: dict[str, str] | None = None # message_hash -> block_number\n # anthropic_cache_id: str | None = None\n" }, { "path": "backend/onyx/llm/prompt_cache/processor.py", "content": "\"\"\"Main processor for prompt caching.\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom onyx.configs.model_configs import ENABLE_PROMPT_CACHING\nfrom onyx.llm.interfaces import LLMConfig\nfrom onyx.llm.models import LanguageModelInput\nfrom onyx.llm.prompt_cache.cache_manager import generate_cache_key_hash\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.llm.prompt_cache.providers.factory import get_provider_adapter\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\n# TODO: test with a history containing images\ndef process_with_prompt_cache(\n llm_config: LLMConfig,\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool = False,\n) -> tuple[LanguageModelInput, CacheMetadata | None]:\n \"\"\"Process prompt with caching support.\n\n This function takes a cacheable prefix and suffix, processes them according to\n the LLM provider's caching capabilities, and returns the combined messages\n ready for LLM API calls along with optional cache metadata.\n\n Args:\n llm: The LLM instance (used to determine provider and model)\n cacheable_prefix: Optional cacheable prefix. If None, no caching is attempted.\n suffix: The non-cacheable suffix to append\n continuation: If True, suffix should be appended to the last message\n of cacheable_prefix rather than being separate messages\n\n Returns:\n Tuple of (processed_prompt, cache_metadata_to_store)\n - processed_prompt: Combined and transformed messages ready for LLM API call\n - cache_metadata_to_store: Optional cache metadata for post-processing\n (currently None for implicit caching, will be populated in future PR\n for explicit caching)\n \"\"\"\n # Check if prompt caching is enabled\n if not ENABLE_PROMPT_CACHING:\n logger.debug(\"Prompt caching is disabled via configuration\")\n # Fall back to no-op behavior\n from onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider\n\n noop_adapter = NoOpPromptCacheProvider()\n combined = noop_adapter.prepare_messages_for_caching(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n cache_metadata=None,\n )\n return combined, None\n\n # If no cacheable prefix, return suffix unchanged\n if cacheable_prefix is None:\n logger.debug(\"No cacheable prefix provided, skipping caching\")\n return suffix, None\n\n # Get provider adapter\n provider_adapter = get_provider_adapter(llm_config)\n\n # If provider doesn't support caching, combine and return unchanged\n if not provider_adapter.supports_caching():\n logger.debug(\n f\"Provider {llm_config.model_provider} does not support caching, combining messages without caching\"\n )\n # Use no-op adapter to combine messages\n from onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider\n\n noop_adapter = NoOpPromptCacheProvider()\n combined = noop_adapter.prepare_messages_for_caching(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n cache_metadata=None,\n )\n return combined, None\n\n # Generate cache key for cacheable prefix\n tenant_id = get_current_tenant_id()\n cache_key_hash = generate_cache_key_hash(\n cacheable_prefix=cacheable_prefix,\n provider=llm_config.model_provider,\n model_name=llm_config.model_name,\n tenant_id=tenant_id,\n )\n\n # For implicit caching: Skip cache lookup (providers handle caching automatically)\n # TODO (explicit caching - future PR): Look up cache metadata in CacheManager\n cache_metadata: CacheMetadata | None = None\n\n # Use provider adapter to prepare messages with caching\n try:\n processed_prompt = provider_adapter.prepare_messages_for_caching(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n cache_metadata=cache_metadata,\n )\n\n logger.debug(\n f\"Processed prompt with caching: provider={llm_config.model_provider}, \"\n f\"model={llm_config.model_name}, cache_key={cache_key_hash[:16]}..., \"\n f\"continuation={continuation}\"\n )\n\n # Create cache metadata for tracking (even for implicit caching)\n # This allows us to track cache usage and effectiveness\n cache_metadata = CacheMetadata(\n cache_key=cache_key_hash,\n provider=llm_config.model_provider,\n model_name=llm_config.model_name,\n tenant_id=tenant_id,\n created_at=datetime.now(timezone.utc),\n last_accessed=datetime.now(timezone.utc),\n )\n\n return processed_prompt, cache_metadata\n\n except Exception as e:\n # Best-effort: log error and fall back to no-op behavior\n logger.warning(\n f\"Error processing prompt with caching for provider={llm_config.model_provider}: {str(e)}. \"\n \"Falling back to non-cached behavior.\"\n )\n # Fall back to no-op adapter\n from onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider\n\n noop_adapter = NoOpPromptCacheProvider()\n combined = noop_adapter.prepare_messages_for_caching(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n cache_metadata=None,\n )\n return combined, None\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/__init__.py", "content": "\"\"\"Provider adapters for prompt caching.\"\"\"\n\nfrom onyx.llm.prompt_cache.providers.anthropic import AnthropicPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.factory import get_provider_adapter\nfrom onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.openai import OpenAIPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.vertex import VertexAIPromptCacheProvider\n\n__all__ = [\n \"AnthropicPromptCacheProvider\",\n \"get_provider_adapter\",\n \"NoOpPromptCacheProvider\",\n \"OpenAIPromptCacheProvider\",\n \"PromptCacheProvider\",\n \"VertexAIPromptCacheProvider\",\n]\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/anthropic.py", "content": "\"\"\"Anthropic provider adapter for prompt caching.\"\"\"\n\nfrom collections.abc import Sequence\n\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform\nfrom onyx.llm.prompt_cache.utils import revalidate_message_from_original\n\n\ndef _add_anthropic_cache_control(\n messages: Sequence[ChatCompletionMessage],\n) -> Sequence[ChatCompletionMessage]:\n \"\"\"Add cache_control parameter to messages for Anthropic caching.\n\n Args:\n messages: Messages to transform\n\n Returns:\n Messages with cache_control added\n \"\"\"\n last_message_dict = dict(messages[-1])\n last_message_dict[\"cache_control\"] = {\"type\": \"ephemeral\"}\n last_message = revalidate_message_from_original(\n original=messages[-1], mutated=last_message_dict\n )\n return list(messages[:-1]) + [last_message]\n\n\nclass AnthropicPromptCacheProvider(PromptCacheProvider):\n \"\"\"Anthropic adapter for prompt caching (explicit caching with cache_control).\n implicit caching = just need to ensure byte-equivalent prefixes, and the provider\n auto-detects and reuses them.\n explicit caching = the caller must do _something_ to enable provider-side caching.\n In this case, anthropic supports explicit caching via the cache_control parameter:\n https://platform.claude.com/docs/en/build-with-claude/prompt-caching\n \"\"\"\n\n def supports_caching(self) -> bool:\n \"\"\"Anthropic supports explicit prompt caching.\"\"\"\n return True\n\n def prepare_messages_for_caching(\n self,\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool,\n cache_metadata: CacheMetadata | None, # noqa: ARG002\n ) -> LanguageModelInput:\n \"\"\"Prepare messages for Anthropic caching.\n\n Anthropic requires cache_control parameter on cacheable messages.\n We add cache_control={\"type\": \"ephemeral\"} to all cacheable prefix messages.\n\n Args:\n cacheable_prefix: Optional cacheable prefix\n suffix: Non-cacheable suffix\n continuation: Whether to append suffix to last prefix message\n cache_metadata: Cache metadata (for future explicit caching support)\n\n Returns:\n Combined messages with cache_control on cacheable messages\n \"\"\"\n return prepare_messages_with_cacheable_transform(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n transform_cacheable=_add_anthropic_cache_control,\n )\n\n def extract_cache_metadata(\n self,\n response: dict, # noqa: ARG002\n cache_key: str, # noqa: ARG002\n ) -> CacheMetadata | None:\n \"\"\"Extract cache metadata from Anthropic response.\n\n Anthropic may return cache identifiers in the response.\n For now, we don't extract detailed metadata (future explicit caching support).\n\n Args:\n response: Anthropic API response dictionary\n cache_key: Cache key used for this request\n\n Returns:\n CacheMetadata if extractable, None otherwise\n \"\"\"\n # TODO: Extract cache identifiers from response when implementing explicit caching\n return None\n\n def get_cache_ttl_seconds(self) -> int:\n \"\"\"Get cache TTL for Anthropic (5 minutes default).\"\"\"\n return 300\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/base.py", "content": "\"\"\"Base interface for provider-specific prompt caching adapters.\"\"\"\n\nfrom abc import ABC\nfrom abc import abstractmethod\n\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.prompt_cache.models import CacheMetadata\n\n\nclass PromptCacheProvider(ABC):\n \"\"\"Abstract base class for provider-specific prompt caching logic.\"\"\"\n\n @abstractmethod\n def supports_caching(self) -> bool:\n \"\"\"Whether this provider supports prompt caching.\n\n Returns:\n True if caching is supported, False otherwise\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def prepare_messages_for_caching(\n self,\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool,\n cache_metadata: CacheMetadata | None,\n ) -> LanguageModelInput:\n \"\"\"Transform messages to enable caching.\n\n Args:\n cacheable_prefix: Optional cacheable prefix (can be str or Sequence[ChatCompletionMessage])\n suffix: Non-cacheable suffix (can be str or Sequence[ChatCompletionMessage])\n continuation: If True, suffix should be appended to the last message\n of cacheable_prefix rather than being separate messages.\n Note: When cacheable_prefix is a string, it should remain in its own\n content block even if continuation=True.\n cache_metadata: Optional cache metadata from previous requests\n\n Returns:\n Combined and transformed messages ready for LLM API call\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def extract_cache_metadata(\n self,\n response: dict, # Provider-specific response object\n cache_key: str,\n ) -> CacheMetadata | None:\n \"\"\"Extract cache metadata from API response.\n\n Args:\n response: Provider-specific response dictionary\n cache_key: Cache key used for this request\n\n Returns:\n CacheMetadata if extractable, None otherwise\n \"\"\"\n raise NotImplementedError\n\n @abstractmethod\n def get_cache_ttl_seconds(self) -> int:\n \"\"\"Get cache TTL in seconds for this provider.\n\n Returns:\n TTL in seconds\n \"\"\"\n raise NotImplementedError\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/factory.py", "content": "\"\"\"Factory for creating provider-specific prompt cache adapters.\"\"\"\n\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LLMConfig\nfrom onyx.llm.prompt_cache.providers.anthropic import AnthropicPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.noop import NoOpPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.openai import OpenAIPromptCacheProvider\nfrom onyx.llm.prompt_cache.providers.vertex import VertexAIPromptCacheProvider\n\nANTHROPIC_BEDROCK_TAG = \"anthropic.\"\n\n\ndef get_provider_adapter(llm_config: LLMConfig) -> PromptCacheProvider:\n \"\"\"Get the appropriate prompt cache provider adapter for a given provider.\n\n Args:\n provider: Provider name (e.g., \"openai\", \"anthropic\", \"vertex_ai\")\n\n Returns:\n PromptCacheProvider instance for the given provider\n \"\"\"\n if llm_config.model_provider == LlmProviderNames.OPENAI:\n return OpenAIPromptCacheProvider()\n elif llm_config.model_provider == LlmProviderNames.ANTHROPIC or (\n llm_config.model_provider == LlmProviderNames.BEDROCK\n and ANTHROPIC_BEDROCK_TAG in llm_config.model_name\n ):\n return AnthropicPromptCacheProvider()\n elif llm_config.model_provider == LlmProviderNames.VERTEX_AI:\n return VertexAIPromptCacheProvider()\n else:\n # Default to no-op for providers without caching support\n return NoOpPromptCacheProvider()\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/noop.py", "content": "\"\"\"No-op provider adapter for providers without caching support.\"\"\"\n\nfrom onyx.llm.models import LanguageModelInput\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform\n\n\nclass NoOpPromptCacheProvider(PromptCacheProvider):\n \"\"\"No-op adapter for providers that don't support prompt caching.\"\"\"\n\n def supports_caching(self) -> bool:\n \"\"\"No-op providers don't support caching.\"\"\"\n return False\n\n def prepare_messages_for_caching(\n self,\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool,\n cache_metadata: CacheMetadata | None, # noqa: ARG002\n ) -> LanguageModelInput:\n \"\"\"Return messages unchanged (no caching support).\n\n Args:\n cacheable_prefix: Optional cacheable prefix (can be str or Sequence[ChatCompletionMessage])\n suffix: Non-cacheable suffix (can be str or Sequence[ChatCompletionMessage])\n continuation: Whether to append suffix to last prefix message.\n Note: When cacheable_prefix is a string, it remains in its own content block.\n cache_metadata: Cache metadata (ignored)\n\n Returns:\n Combined messages (prefix + suffix)\n \"\"\"\n # No transformation needed for no-op provider\n return prepare_messages_with_cacheable_transform(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n transform_cacheable=None,\n )\n\n def extract_cache_metadata(\n self,\n response: dict, # noqa: ARG002\n cache_key: str, # noqa: ARG002\n ) -> CacheMetadata | None:\n \"\"\"No cache metadata to extract.\"\"\"\n return None\n\n def get_cache_ttl_seconds(self) -> int:\n \"\"\"Return default TTL (not used for no-op).\"\"\"\n return 0\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/openai.py", "content": "\"\"\"OpenAI provider adapter for prompt caching.\"\"\"\n\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform\n\n\nclass OpenAIPromptCacheProvider(PromptCacheProvider):\n \"\"\"OpenAI adapter for prompt caching (implicit caching).\"\"\"\n\n def supports_caching(self) -> bool:\n \"\"\"OpenAI supports automatic prompt caching.\"\"\"\n return True\n\n def prepare_messages_for_caching(\n self,\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool,\n cache_metadata: CacheMetadata | None, # noqa: ARG002\n ) -> LanguageModelInput:\n \"\"\"Prepare messages for OpenAI caching.\n\n OpenAI handles caching automatically, so we just normalize and combine\n the messages. The provider will automatically cache prefixes >1024 tokens.\n\n Args:\n cacheable_prefix: Optional cacheable prefix\n suffix: Non-cacheable suffix\n continuation: Whether to append suffix to last prefix message\n cache_metadata: Cache metadata (ignored for implicit caching)\n\n Returns:\n Combined messages ready for LLM API call\n \"\"\"\n # No transformation needed for OpenAI (implicit caching)\n return prepare_messages_with_cacheable_transform(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n transform_cacheable=None,\n )\n\n def extract_cache_metadata(\n self,\n response: dict, # noqa: ARG002\n cache_key: str, # noqa: ARG002\n ) -> CacheMetadata | None:\n \"\"\"Extract cache metadata from OpenAI response.\n\n OpenAI responses may include cached_tokens in the usage field.\n For implicit caching, we don't need to store much metadata.\n\n Args:\n response: OpenAI API response dictionary\n cache_key: Cache key used for this request\n\n Returns:\n CacheMetadata if extractable, None otherwise\n \"\"\"\n # For implicit caching, OpenAI handles everything automatically\n # We could extract cached_tokens from response.get(\"usage\", {}).get(\"cached_tokens\")\n # but for now, we don't need to store metadata for implicit caching\n return None\n\n def get_cache_ttl_seconds(self) -> int:\n \"\"\"Get cache TTL for OpenAI (1 hour max).\"\"\"\n return 3600\n" }, { "path": "backend/onyx/llm/prompt_cache/providers/vertex.py", "content": "\"\"\"Vertex AI provider adapter for prompt caching.\"\"\"\n\nfrom collections.abc import Sequence\n\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.prompt_cache.models import CacheMetadata\nfrom onyx.llm.prompt_cache.providers.base import PromptCacheProvider\nfrom onyx.llm.prompt_cache.utils import prepare_messages_with_cacheable_transform\nfrom onyx.llm.prompt_cache.utils import revalidate_message_from_original\n\n\nclass VertexAIPromptCacheProvider(PromptCacheProvider):\n \"\"\"Vertex AI adapter for prompt caching (implicit caching for this PR).\"\"\"\n\n def supports_caching(self) -> bool:\n \"\"\"Vertex AI supports prompt caching (implicit and explicit).\"\"\"\n return True\n\n def prepare_messages_for_caching(\n self,\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool,\n cache_metadata: CacheMetadata | None, # noqa: ARG002\n ) -> LanguageModelInput:\n \"\"\"Prepare messages for Vertex AI caching.\n\n For implicit caching we attach cache_control={\"type\": \"ephemeral\"} to every\n cacheable prefix message so Vertex/Gemini can reuse them automatically.\n Explicit context caching (with cache blocks) will be added in a future PR.\n\n Args:\n cacheable_prefix: Optional cacheable prefix\n suffix: Non-cacheable suffix\n continuation: Whether to append suffix to last prefix message\n cache_metadata: Cache metadata (for future explicit caching support)\n\n Returns:\n Combined messages ready for LLM API call\n \"\"\"\n # For implicit caching, no transformation needed (Vertex handles caching automatically)\n # TODO (explicit caching - future PR):\n # - Check cache_metadata for vertex_block_numbers\n # - Create transform function that replaces messages with cache_block_id if available\n # - Or adds cache_control parameter if not using cached blocks\n return prepare_messages_with_cacheable_transform(\n cacheable_prefix=cacheable_prefix,\n suffix=suffix,\n continuation=continuation,\n transform_cacheable=None, # TODO: support explicit caching\n )\n\n def extract_cache_metadata(\n self,\n response: dict, # noqa: ARG002\n cache_key: str, # noqa: ARG002\n ) -> CacheMetadata | None:\n \"\"\"Extract cache metadata from Vertex AI response.\n\n For this PR (implicit caching): Extract basic cache usage info if available.\n TODO (explicit caching - future PR): Extract block numbers from response\n and store in metadata.\n\n Args:\n response: Vertex AI API response dictionary\n cache_key: Cache key used for this request\n\n Returns:\n CacheMetadata if extractable, None otherwise\n \"\"\"\n # For implicit caching, Vertex handles everything automatically\n # TODO (explicit caching - future PR):\n # - Extract cache block numbers from response\n # - Store in cache_metadata.vertex_block_numbers\n return None\n\n def get_cache_ttl_seconds(self) -> int:\n \"\"\"Get cache TTL for Vertex AI (5 minutes).\"\"\"\n return 300\n\n\ndef _add_vertex_cache_control(\n messages: Sequence[ChatCompletionMessage],\n) -> Sequence[ChatCompletionMessage]:\n \"\"\"Add cache_control inside content blocks for Vertex AI/Gemini caching.\n\n Gemini requires cache_control to be on a content block within the content array,\n not at the message level. This function converts string content to the array format\n and adds cache_control to the last content block in each cacheable message.\n \"\"\"\n # NOTE: unfortunately we need a much more sophisticated mechnism to support\n # explict caching with vertex in the presence of tools and system messages\n # (since they're supposed to be stripped out when setting cache_control)\n # so we're deferring this to a future PR.\n updated: list[ChatCompletionMessage] = []\n for message in messages:\n mutated = dict(message)\n content = mutated.get(\"content\")\n\n if isinstance(content, str):\n # Convert string content to array format with cache_control\n mutated[\"content\"] = [\n {\n \"type\": \"text\",\n \"text\": content,\n \"cache_control\": {\"type\": \"ephemeral\"},\n }\n ]\n elif isinstance(content, list) and content:\n # Content is already an array - add cache_control to last block\n new_content = []\n for i, block in enumerate(content):\n if isinstance(block, dict):\n block_copy = dict(block)\n # Add cache_control to the last content block\n if i == len(content) - 1:\n block_copy[\"cache_control\"] = {\"type\": \"ephemeral\"}\n new_content.append(block_copy)\n else:\n new_content.append(block)\n mutated[\"content\"] = new_content\n\n updated.append(revalidate_message_from_original(message, mutated))\n return updated\n" }, { "path": "backend/onyx/llm/prompt_cache/utils.py", "content": "# pyright: reportMissingTypeStubs=false\n\"\"\"Utility functions for prompt caching.\"\"\"\n\nimport json\nfrom collections.abc import Callable\nfrom collections.abc import Sequence\nfrom typing import Any\n\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import LanguageModelInput\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef combine_messages_with_continuation(\n prefix_msgs: Sequence[ChatCompletionMessage],\n suffix_msgs: Sequence[ChatCompletionMessage],\n continuation: bool,\n) -> list[ChatCompletionMessage]:\n \"\"\"Combine prefix and suffix messages, handling continuation flag.\n\n Args:\n prefix_msgs: Normalized cacheable prefix messages\n suffix_msgs: Normalized suffix messages\n continuation: If True, append suffix content to the last message of prefix\n was_prefix_string: Deprecated, no longer used\n\n Returns:\n Combined messages\n \"\"\"\n if not continuation or not prefix_msgs:\n return list(prefix_msgs) + list(suffix_msgs)\n # Append suffix content to last message of prefix\n result = list(prefix_msgs)\n last_msg = dict(result[-1])\n suffix_first = dict(suffix_msgs[0]) if suffix_msgs else {}\n\n # Combine content\n if \"content\" in last_msg and \"content\" in suffix_first:\n if isinstance(last_msg[\"content\"], str) and isinstance(\n suffix_first[\"content\"], str\n ):\n last_msg[\"content\"] = last_msg[\"content\"] + suffix_first[\"content\"]\n else:\n # Handle list content (multimodal)\n prefix_content = (\n last_msg[\"content\"]\n if isinstance(last_msg[\"content\"], list)\n else [{\"type\": \"text\", \"text\": last_msg[\"content\"]}]\n )\n suffix_content = (\n suffix_first[\"content\"]\n if isinstance(suffix_first[\"content\"], list)\n else [{\"type\": \"text\", \"text\": suffix_first[\"content\"]}]\n )\n last_msg[\"content\"] = prefix_content + suffix_content\n\n result[-1] = revalidate_message_from_original(original=result[-1], mutated=last_msg)\n result.extend(suffix_msgs[1:])\n return result\n\n\ndef revalidate_message_from_original(\n original: ChatCompletionMessage,\n mutated: dict[str, Any],\n) -> ChatCompletionMessage:\n \"\"\"Rebuild a mutated message using the original BaseModel type.\n\n Some providers need to add cache metadata to messages. Re-run validation against\n the original message's Pydantic class so union discrimination (by role) stays\n intact.\n \"\"\"\n cls = original.__class__\n try:\n return cls.model_validate_json(json.dumps(mutated))\n except Exception:\n return cls.model_validate(mutated)\n\n\ndef prepare_messages_with_cacheable_transform(\n cacheable_prefix: LanguageModelInput | None,\n suffix: LanguageModelInput,\n continuation: bool,\n transform_cacheable: (\n Callable[[Sequence[ChatCompletionMessage]], Sequence[ChatCompletionMessage]]\n | None\n ) = None,\n) -> LanguageModelInput:\n \"\"\"Prepare messages for caching with optional transformation of cacheable prefix.\n\n This is a shared utility that handles the common flow:\n 1. Normalize inputs\n 2. Optionally transform cacheable messages\n 3. Combine with continuation handling\n\n Args:\n cacheable_prefix: Optional cacheable prefix\n suffix: Non-cacheable suffix\n continuation: Whether to append suffix to last prefix message\n transform_cacheable: Optional function to transform cacheable messages\n (e.g., add cache_control parameter). If None, messages are used as-is.\n\n Returns:\n Combined messages ready for LLM API call\n \"\"\"\n if cacheable_prefix is None:\n return suffix\n\n prefix_msgs = (\n cacheable_prefix if isinstance(cacheable_prefix, list) else [cacheable_prefix]\n )\n suffix_msgs = suffix if isinstance(suffix, list) else [suffix]\n\n # Apply transformation to cacheable messages if provided\n if transform_cacheable is not None:\n prefix_msgs = list(transform_cacheable(prefix_msgs))\n\n return combine_messages_with_continuation(\n prefix_msgs=prefix_msgs, suffix_msgs=suffix_msgs, continuation=continuation\n )\n" }, { "path": "backend/onyx/llm/request_context.py", "content": "import contextvars\n\n\n_LLM_MOCK_RESPONSE_CONTEXTVAR: contextvars.ContextVar[str | None] = (\n contextvars.ContextVar(\"llm_mock_response\", default=None)\n)\n\n\ndef get_llm_mock_response() -> str | None:\n return _LLM_MOCK_RESPONSE_CONTEXTVAR.get()\n\n\ndef set_llm_mock_response(mock_response: str | None) -> contextvars.Token[str | None]:\n return _LLM_MOCK_RESPONSE_CONTEXTVAR.set(mock_response)\n\n\ndef reset_llm_mock_response(token: contextvars.Token[str | None]) -> None:\n try:\n _LLM_MOCK_RESPONSE_CONTEXTVAR.reset(token)\n except ValueError:\n # Streaming requests can cross execution contexts.\n # Best effort clear to avoid crashing request teardown in integration mode.\n _LLM_MOCK_RESPONSE_CONTEXTVAR.set(None)\n" }, { "path": "backend/onyx/llm/utils.py", "content": "import copy\nimport re\nfrom collections.abc import Callable\nfrom functools import lru_cache\nfrom typing import Any\nfrom typing import cast\nfrom typing import TYPE_CHECKING\n\nfrom sqlalchemy import select\n\nfrom onyx.configs.app_configs import LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS\nfrom onyx.configs.app_configs import MAX_TOKENS_FOR_FULL_INCLUSION\nfrom onyx.configs.app_configs import SEND_USER_METADATA_TO_LLM_PROVIDER\nfrom onyx.configs.app_configs import USE_CHUNK_SUMMARY\nfrom onyx.configs.app_configs import USE_DOCUMENT_SUMMARY\nfrom onyx.configs.model_configs import GEN_AI_MAX_TOKENS\nfrom onyx.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS\nfrom onyx.configs.model_configs import GEN_AI_NUM_RESERVED_OUTPUT_TOKENS\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import LLMModelFlowType\nfrom onyx.db.models import LLMProvider\nfrom onyx.db.models import ModelConfiguration\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.model_response import ModelResponse\nfrom onyx.llm.models import UserMessage\nfrom onyx.prompts.contextual_retrieval import CONTEXTUAL_RAG_TOKEN_ESTIMATE\nfrom onyx.prompts.contextual_retrieval import DOCUMENT_SUMMARY_TOKEN_ESTIMATE\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE\n\n\nif TYPE_CHECKING:\n from onyx.server.manage.llm.models import LLMProviderView\n\n\nlogger = setup_logger()\n\nMAX_CONTEXT_TOKENS = 100\nONE_MILLION = 1_000_000\nCHUNKS_PER_DOC_ESTIMATE = 5\nMAX_LITELLM_USER_ID_LENGTH = 64\n_TWELVE_LABS_PEGASUS_MODEL_NAMES = [\n \"us.twelvelabs.pegasus-1-2-v1:0\",\n \"us.twelvelabs.pegasus-1-2-v1\",\n \"twelvelabs/us.twelvelabs.pegasus-1-2-v1:0\",\n \"twelvelabs/us.twelvelabs.pegasus-1-2-v1\",\n]\n_TWELVE_LABS_PEGASUS_OUTPUT_TOKENS = max(512, GEN_AI_MODEL_FALLBACK_MAX_TOKENS // 4)\nCUSTOM_LITELLM_MODEL_OVERRIDES: dict[str, dict[str, Any]] = {\n model_name: {\n \"max_input_tokens\": GEN_AI_MODEL_FALLBACK_MAX_TOKENS,\n \"max_output_tokens\": _TWELVE_LABS_PEGASUS_OUTPUT_TOKENS,\n \"max_tokens\": GEN_AI_MODEL_FALLBACK_MAX_TOKENS,\n \"supports_reasoning\": False,\n \"supports_vision\": False,\n }\n for model_name in _TWELVE_LABS_PEGASUS_MODEL_NAMES\n}\n\n\ndef truncate_litellm_user_id(user_id: str) -> str:\n \"\"\"Truncate the LiteLLM `user` field maximum length.\"\"\"\n if len(user_id) <= MAX_LITELLM_USER_ID_LENGTH:\n return user_id\n logger.warning(\n \"User's ID exceeds %d chars (len=%d); truncating for Litellm logging compatibility.\",\n MAX_LITELLM_USER_ID_LENGTH,\n len(user_id),\n )\n return user_id[:MAX_LITELLM_USER_ID_LENGTH]\n\n\ndef build_litellm_passthrough_kwargs(\n model_kwargs: dict[str, Any],\n user_identity: LLMUserIdentity | None,\n) -> dict[str, Any]:\n \"\"\"Build kwargs passed through directly to LiteLLM.\n\n Returns `model_kwargs` unchanged unless we need to add user/session metadata,\n in which case a copy is returned to avoid cross-request mutation.\n \"\"\"\n\n if not (SEND_USER_METADATA_TO_LLM_PROVIDER and user_identity):\n return model_kwargs\n\n passthrough_kwargs = copy.deepcopy(model_kwargs)\n\n if user_identity.user_id:\n passthrough_kwargs[\"user\"] = truncate_litellm_user_id(user_identity.user_id)\n\n if user_identity.session_id:\n existing_metadata = passthrough_kwargs.get(\"metadata\")\n metadata: dict[str, Any] | None\n if existing_metadata is None:\n metadata = {}\n elif isinstance(existing_metadata, dict):\n metadata = copy.deepcopy(existing_metadata)\n else:\n metadata = None\n\n if metadata is not None:\n metadata[\"session_id\"] = user_identity.session_id\n passthrough_kwargs[\"metadata\"] = metadata\n\n return passthrough_kwargs\n\n\ndef _unwrap_nested_exception(error: Exception) -> Exception:\n \"\"\"\n Traverse common exception wrappers to surface the underlying LiteLLM error.\n \"\"\"\n visited: set[int] = set()\n current = error\n for _ in range(100):\n visited.add(id(current))\n candidate: Exception | None = None\n cause = getattr(current, \"__cause__\", None)\n if isinstance(cause, Exception):\n candidate = cause\n elif (\n hasattr(current, \"args\")\n and len(getattr(current, \"args\")) == 1\n and isinstance(current.args[0], Exception)\n ):\n candidate = current.args[0]\n if candidate is None or id(candidate) in visited:\n break\n current = candidate\n return current\n\n\ndef litellm_exception_to_error_msg(\n e: Exception,\n llm: LLM,\n fallback_to_error_msg: bool = False,\n custom_error_msg_mappings: (\n dict[str, str] | None\n ) = LITELLM_CUSTOM_ERROR_MESSAGE_MAPPINGS,\n) -> tuple[str, str, bool]:\n \"\"\"Convert a LiteLLM exception to a user-friendly error message with classification.\n\n Returns:\n tuple: (error_message, error_code, is_retryable)\n - error_message: User-friendly error description\n - error_code: Categorized error code for frontend display\n - is_retryable: Whether the user should try again\n \"\"\"\n from litellm.exceptions import BadRequestError\n from litellm.exceptions import AuthenticationError\n from litellm.exceptions import PermissionDeniedError\n from litellm.exceptions import NotFoundError\n from litellm.exceptions import UnprocessableEntityError\n from litellm.exceptions import RateLimitError\n from litellm.exceptions import ContextWindowExceededError\n from litellm.exceptions import APIConnectionError\n from litellm.exceptions import APIError\n from litellm.exceptions import Timeout\n from litellm.exceptions import ContentPolicyViolationError\n from litellm.exceptions import BudgetExceededError\n from litellm.exceptions import ServiceUnavailableError\n\n core_exception = _unwrap_nested_exception(e)\n error_msg = str(core_exception)\n error_code = \"UNKNOWN_ERROR\"\n is_retryable = True\n\n if custom_error_msg_mappings:\n for error_msg_pattern, custom_error_msg in custom_error_msg_mappings.items():\n if error_msg_pattern in error_msg:\n return custom_error_msg, \"CUSTOM_ERROR\", True\n\n if isinstance(core_exception, BadRequestError):\n error_msg = \"Bad request: The server couldn't process your request. Please check your input.\"\n error_code = \"BAD_REQUEST\"\n is_retryable = True\n elif isinstance(core_exception, AuthenticationError):\n error_msg = \"Authentication failed: Please check your API key and credentials.\"\n error_code = \"AUTH_ERROR\"\n is_retryable = False\n elif isinstance(core_exception, PermissionDeniedError):\n error_msg = (\n \"Permission denied: You don't have the necessary permissions for this operation. \"\n \"Ensure you have access to this model.\"\n )\n error_code = \"PERMISSION_DENIED\"\n is_retryable = False\n elif isinstance(core_exception, NotFoundError):\n error_msg = \"Resource not found: The requested resource doesn't exist.\"\n error_code = \"NOT_FOUND\"\n is_retryable = False\n elif isinstance(core_exception, UnprocessableEntityError):\n error_msg = \"Unprocessable entity: The server couldn't process your request due to semantic errors.\"\n error_code = \"UNPROCESSABLE_ENTITY\"\n is_retryable = True\n elif isinstance(core_exception, RateLimitError):\n provider_name = (\n llm.config.model_provider\n if llm is not None and llm.config.model_provider\n else \"The LLM provider\"\n )\n upstream_detail: str | None = None\n message_attr = getattr(core_exception, \"message\", None)\n if message_attr:\n upstream_detail = str(message_attr)\n elif hasattr(core_exception, \"api_error\"):\n api_error = core_exception.api_error\n if isinstance(api_error, dict):\n upstream_detail = (\n api_error.get(\"message\")\n or api_error.get(\"detail\")\n or api_error.get(\"error\")\n )\n if not upstream_detail:\n upstream_detail = str(core_exception)\n upstream_detail = str(upstream_detail).strip()\n if \":\" in upstream_detail and upstream_detail.lower().startswith(\n \"ratelimiterror\"\n ):\n upstream_detail = upstream_detail.split(\":\", 1)[1].strip()\n upstream_detail_lower = upstream_detail.lower()\n if (\n \"insufficient_quota\" in upstream_detail_lower\n or \"exceeded your current quota\" in upstream_detail_lower\n ):\n error_msg = (\n f\"{provider_name} quota exceeded: {upstream_detail}\"\n if upstream_detail\n else f\"{provider_name} quota exceeded: Verify billing and quota for this API key.\"\n )\n error_code = \"BUDGET_EXCEEDED\"\n is_retryable = False\n else:\n error_msg = (\n f\"{provider_name} rate limit: {upstream_detail}\"\n if upstream_detail\n else f\"{provider_name} rate limit exceeded: Please slow down your requests and try again later.\"\n )\n error_code = \"RATE_LIMIT\"\n is_retryable = True\n elif isinstance(core_exception, ServiceUnavailableError):\n provider_name = (\n llm.config.model_provider\n if llm is not None and llm.config.model_provider\n else \"The LLM provider\"\n )\n # Check if this is specifically the Bedrock \"Too many connections\" error\n if \"Too many connections\" in error_msg or \"BedrockException\" in error_msg:\n error_msg = (\n f\"{provider_name} is experiencing high connection volume and cannot process your request right now. \"\n \"This typically happens when there are too many simultaneous requests to the AI model. \"\n \"Please wait a moment and try again. If this persists, contact your system administrator \"\n \"to review connection limits and retry configurations.\"\n )\n else:\n # Generic 503 Service Unavailable\n error_msg = f\"{provider_name} service error: {str(core_exception)}\"\n error_code = \"SERVICE_UNAVAILABLE\"\n is_retryable = True\n elif isinstance(core_exception, ContextWindowExceededError):\n error_msg = (\n \"Context window exceeded: Your input is too long for the model to process.\"\n )\n if llm is not None:\n try:\n max_context = get_max_input_tokens(\n model_name=llm.config.model_name,\n model_provider=llm.config.model_provider,\n )\n error_msg += f\" Your invoked model ({llm.config.model_name}) has a maximum context size of {max_context}.\"\n except Exception:\n logger.warning(\n \"Unable to get maximum input token for LiteLLM exception handling\"\n )\n error_code = \"CONTEXT_TOO_LONG\"\n is_retryable = False\n elif isinstance(core_exception, ContentPolicyViolationError):\n error_msg = \"Content policy violation: Your request violates the content policy. Please revise your input.\"\n error_code = \"CONTENT_POLICY\"\n is_retryable = False\n elif isinstance(core_exception, APIConnectionError):\n error_msg = \"API connection error: Failed to connect to the API. Please check your internet connection.\"\n error_code = \"CONNECTION_ERROR\"\n is_retryable = True\n elif isinstance(core_exception, BudgetExceededError):\n error_msg = (\n \"Budget exceeded: You've exceeded your allocated budget for API usage.\"\n )\n error_code = \"BUDGET_EXCEEDED\"\n is_retryable = False\n elif isinstance(core_exception, Timeout):\n error_msg = \"Request timed out: The operation took too long to complete. Please try again.\"\n error_code = \"CONNECTION_ERROR\"\n is_retryable = True\n elif isinstance(core_exception, APIError):\n error_msg = f\"API error: An error occurred while communicating with the API. Details: {str(core_exception)}\"\n error_code = \"API_ERROR\"\n is_retryable = True\n elif not fallback_to_error_msg:\n error_msg = \"An unexpected error occurred while processing your request. Please try again later.\"\n error_code = \"UNKNOWN_ERROR\"\n is_retryable = True\n\n return error_msg, error_code, is_retryable\n\n\ndef llm_response_to_string(message: ModelResponse) -> str:\n if not isinstance(message.choice.message.content, str):\n raise RuntimeError(\"LLM message not in expected format.\")\n\n return message.choice.message.content\n\n\ndef check_number_of_tokens(\n text: str, encode_fn: Callable[[str], list] | None = None\n) -> int:\n \"\"\"Gets the number of tokens in the provided text, using the provided encoding\n function. If none is provided, default to the tiktoken encoder used by GPT-3.5\n and GPT-4.\n \"\"\"\n import tiktoken\n\n if encode_fn is None:\n encode_fn = tiktoken.get_encoding(\"cl100k_base\").encode\n\n return len(encode_fn(text))\n\n\ndef test_llm(llm: LLM) -> str | None:\n # try for up to 2 timeouts (e.g. 10 seconds in total)\n error_msg = None\n for _ in range(2):\n try:\n llm.invoke(UserMessage(content=\"Do not respond\"), max_tokens=50)\n return None\n except Exception as e:\n error_msg = str(e)\n logger.warning(f\"Failed to call LLM with the following error: {error_msg}\")\n\n return error_msg\n\n\n@lru_cache(maxsize=1) # the copy.deepcopy is expensive, so we cache the result\ndef get_model_map() -> dict:\n import litellm\n\n DIVIDER = \"/\"\n\n original_map = cast(dict[str, dict], litellm.model_cost)\n starting_map = copy.deepcopy(original_map)\n for key in original_map:\n if DIVIDER in key:\n truncated_key = key.split(DIVIDER)[-1]\n # make sure not to overwrite an original key\n if truncated_key in original_map:\n continue\n\n # if there are multiple possible matches, choose the most \"detailed\"\n # one as a heuristic. \"detailed\" = the description of the model\n # has the most filled out fields.\n existing_truncated_value = starting_map.get(truncated_key)\n potential_truncated_value = original_map[key]\n if not existing_truncated_value or len(potential_truncated_value) > len(\n existing_truncated_value\n ):\n starting_map[truncated_key] = potential_truncated_value\n\n for model_name, model_metadata in CUSTOM_LITELLM_MODEL_OVERRIDES.items():\n if model_name in starting_map:\n continue\n starting_map[model_name] = copy.deepcopy(model_metadata)\n\n # NOTE: outside of the explicit CUSTOM_LITELLM_MODEL_OVERRIDES,\n # we avoid hard-coding additional models here. Ollama, for example,\n # allows the user to specify their desired max context window, and it's\n # unlikely to be standard across users even for the same model\n # (it heavily depends on their hardware). For those cases, we rely on\n # GEN_AI_MODEL_FALLBACK_MAX_TOKENS to cover this.\n # for model_name in [\n # \"llama3.2\",\n # \"llama3.2:1b\",\n # \"llama3.2:3b\",\n # \"llama3.2:11b\",\n # \"llama3.2:90b\",\n # ]:\n # starting_map[f\"ollama/{model_name}\"] = {\n # \"max_tokens\": 128000,\n # \"max_input_tokens\": 128000,\n # \"max_output_tokens\": 128000,\n # }\n\n return starting_map\n\n\ndef _strip_extra_provider_from_model_name(model_name: str) -> str:\n return model_name.split(\"/\")[1] if \"/\" in model_name else model_name\n\n\ndef _strip_colon_from_model_name(model_name: str) -> str:\n return \":\".join(model_name.split(\":\")[:-1]) if \":\" in model_name else model_name\n\n\ndef find_model_obj(model_map: dict, provider: str, model_name: str) -> dict | None:\n stripped_model_name = _strip_extra_provider_from_model_name(model_name)\n\n model_names = [\n model_name,\n _strip_extra_provider_from_model_name(model_name),\n # Remove leading extra provider. Usually for cases where user has a\n # customer model proxy which appends another prefix\n # remove :XXXX from the end, if present. Needed for ollama.\n _strip_colon_from_model_name(model_name),\n _strip_colon_from_model_name(stripped_model_name),\n ]\n\n # Filter out None values and deduplicate model names\n filtered_model_names = [name for name in model_names if name]\n\n # First try all model names with provider prefix\n for model_name in filtered_model_names:\n model_obj = model_map.get(f\"{provider}/{model_name}\")\n if model_obj:\n return model_obj\n\n # Then try all model names without provider prefix\n for model_name in filtered_model_names:\n model_obj = model_map.get(model_name)\n if model_obj:\n return model_obj\n\n return None\n\n\ndef get_llm_contextual_cost(\n llm: LLM,\n) -> float:\n \"\"\"\n Approximate the cost of using the given LLM for indexing with Contextual RAG.\n\n We use a precomputed estimate for the number of tokens in the contextualizing prompts,\n and we assume that every chunk is maximized in terms of content and context.\n We also assume that every document is maximized in terms of content, as currently if\n a document is longer than a certain length, its summary is used instead of the full content.\n\n We expect that the first assumption will overestimate more than the second one\n underestimates, so this should be a fairly conservative price estimate. Also,\n this does not account for the cost of documents that fit within a single chunk\n which do not get contextualized.\n \"\"\"\n\n import litellm\n\n # calculate input costs\n num_tokens = ONE_MILLION\n num_input_chunks = num_tokens // DOC_EMBEDDING_CONTEXT_SIZE\n\n # We assume that the documents are MAX_TOKENS_FOR_FULL_INCLUSION tokens long\n # on average.\n num_docs = num_tokens // MAX_TOKENS_FOR_FULL_INCLUSION\n\n num_input_tokens = 0\n num_output_tokens = 0\n\n if not USE_CHUNK_SUMMARY and not USE_DOCUMENT_SUMMARY:\n return 0\n\n if USE_CHUNK_SUMMARY:\n # Each per-chunk prompt includes:\n # - The prompt tokens\n # - the document tokens\n # - the chunk tokens\n\n # for each chunk, we prompt the LLM with the contextual RAG prompt\n # and the full document content (or the doc summary, so this is an overestimate)\n num_input_tokens += num_input_chunks * (\n CONTEXTUAL_RAG_TOKEN_ESTIMATE + MAX_TOKENS_FOR_FULL_INCLUSION\n )\n\n # in aggregate, each chunk content is used as a prompt input once\n # so the full input size is covered\n num_input_tokens += num_tokens\n\n # A single MAX_CONTEXT_TOKENS worth of output is generated per chunk\n num_output_tokens += num_input_chunks * MAX_CONTEXT_TOKENS\n\n # going over each doc once means all the tokens, plus the prompt tokens for\n # the summary prompt. This CAN happen even when USE_DOCUMENT_SUMMARY is false,\n # since doc summaries are used for longer documents when USE_CHUNK_SUMMARY is true.\n # So, we include this unconditionally to overestimate.\n num_input_tokens += num_tokens + num_docs * DOCUMENT_SUMMARY_TOKEN_ESTIMATE\n num_output_tokens += num_docs * MAX_CONTEXT_TOKENS\n\n try:\n usd_per_prompt, usd_per_completion = litellm.cost_per_token(\n model=llm.config.model_name,\n prompt_tokens=num_input_tokens,\n completion_tokens=num_output_tokens,\n )\n except Exception:\n logger.exception(\n \"An unexpected error occurred while calculating cost for model \"\n f\"{llm.config.model_name} (potentially due to malformed name). \"\n \"Assuming cost is 0.\"\n )\n return 0\n\n # Costs are in USD dollars per million tokens\n return usd_per_prompt + usd_per_completion\n\n\ndef llm_max_input_tokens(\n model_map: dict,\n model_name: str,\n model_provider: str,\n) -> int:\n \"\"\"Best effort attempt to get the max input tokens for the LLM.\"\"\"\n if GEN_AI_MAX_TOKENS:\n # This is an override, so always return this\n logger.info(f\"Using override GEN_AI_MAX_TOKENS: {GEN_AI_MAX_TOKENS}\")\n return GEN_AI_MAX_TOKENS\n\n model_obj = find_model_obj(\n model_map,\n model_provider,\n model_name,\n )\n if not model_obj:\n logger.warning(\n f\"Model '{model_name}' not found in LiteLLM. Falling back to {GEN_AI_MODEL_FALLBACK_MAX_TOKENS} tokens.\"\n )\n return GEN_AI_MODEL_FALLBACK_MAX_TOKENS\n\n if \"max_input_tokens\" in model_obj:\n return model_obj[\"max_input_tokens\"]\n\n if \"max_tokens\" in model_obj:\n return model_obj[\"max_tokens\"]\n\n logger.warning(\n f\"No max tokens found for '{model_name}'. Falling back to {GEN_AI_MODEL_FALLBACK_MAX_TOKENS} tokens.\"\n )\n return GEN_AI_MODEL_FALLBACK_MAX_TOKENS\n\n\ndef get_llm_max_output_tokens(\n model_map: dict,\n model_name: str,\n model_provider: str,\n) -> int:\n \"\"\"Best effort attempt to get the max output tokens for the LLM.\"\"\"\n default_output_tokens = int(GEN_AI_MODEL_FALLBACK_MAX_TOKENS)\n\n model_obj = model_map.get(f\"{model_provider}/{model_name}\")\n if not model_obj:\n model_obj = model_map.get(model_name)\n\n if not model_obj:\n logger.warning(\n f\"Model '{model_name}' not found in LiteLLM. Falling back to {default_output_tokens} output tokens.\"\n )\n return default_output_tokens\n\n if \"max_output_tokens\" in model_obj:\n return model_obj[\"max_output_tokens\"]\n\n # Fallback to a fraction of max_tokens if max_output_tokens is not specified\n if \"max_tokens\" in model_obj:\n return int(model_obj[\"max_tokens\"] * 0.1)\n\n logger.warning(\n f\"No max output tokens found for '{model_name}'. Falling back to {default_output_tokens} output tokens.\"\n )\n return default_output_tokens\n\n\ndef get_max_input_tokens(\n model_name: str,\n model_provider: str,\n output_tokens: int = GEN_AI_NUM_RESERVED_OUTPUT_TOKENS,\n) -> int:\n # NOTE: we previously used `litellm.get_max_tokens()`, but despite the name, this actually\n # returns the max OUTPUT tokens. Under the hood, this uses the `litellm.model_cost` dict,\n # and there is no other interface to get what we want. This should be okay though, since the\n # `model_cost` dict is a named public interface:\n # https://litellm.vercel.app/docs/completion/token_usage#7-model_cost\n # model_map is litellm.model_cost\n litellm_model_map = get_model_map()\n\n input_toks = (\n llm_max_input_tokens(\n model_name=model_name,\n model_provider=model_provider,\n model_map=litellm_model_map,\n )\n - output_tokens\n )\n\n if input_toks <= 0:\n return GEN_AI_MODEL_FALLBACK_MAX_TOKENS\n\n return input_toks\n\n\ndef get_max_input_tokens_from_llm_provider(\n llm_provider: \"LLMProviderView\",\n model_name: str,\n) -> int:\n \"\"\"Get max input tokens for a model, with fallback chain.\n\n Fallback order:\n 1. Use max_input_tokens from model_configuration (populated from source APIs\n like OpenRouter, Ollama, or our Bedrock mapping)\n 2. Look up in litellm.model_cost dictionary\n 3. Fall back to GEN_AI_MODEL_FALLBACK_MAX_TOKENS (32000)\n\n Most dynamic providers (OpenRouter, Ollama) provide context_length via their\n APIs. Bedrock doesn't expose this, so we parse from model ID suffix (:200k)\n or use BEDROCK_MODEL_TOKEN_LIMITS mapping. The 32000 fallback is only hit for\n unknown models not in any of these sources.\n \"\"\"\n max_input_tokens = None\n for model_configuration in llm_provider.model_configurations:\n if model_configuration.name == model_name:\n max_input_tokens = model_configuration.max_input_tokens\n return (\n max_input_tokens\n if max_input_tokens\n else get_max_input_tokens(\n model_provider=llm_provider.name,\n model_name=model_name,\n )\n )\n\n\ndef get_bedrock_token_limit(model_id: str) -> int:\n \"\"\"Look up token limit for a Bedrock model.\n\n AWS Bedrock API doesn't expose token limits directly. This function\n attempts to determine the limit from multiple sources.\n\n Lookup order:\n 1. Parse from model ID suffix (e.g., \":200k\" → 200000)\n 2. Check LiteLLM's model_cost dictionary\n 3. Fall back to our hardcoded BEDROCK_MODEL_TOKEN_LIMITS mapping\n 4. Default to 32000 if not found anywhere\n \"\"\"\n from onyx.llm.constants import BEDROCK_MODEL_TOKEN_LIMITS\n\n model_id_lower = model_id.lower()\n\n # 1. Try to parse context length from model ID suffix\n # Format: \"model-name:version:NNNk\" where NNN is the context length in thousands\n # Examples: \":200k\", \":128k\", \":1000k\", \":8k\", \":4k\"\n context_match = re.search(r\":(\\d+)k\\b\", model_id_lower)\n if context_match:\n return int(context_match.group(1)) * 1000\n\n # 2. Check LiteLLM's model_cost dictionary\n try:\n model_map = get_model_map()\n # Try with bedrock/ prefix first, then without\n for key in [f\"bedrock/{model_id}\", model_id]:\n if key in model_map:\n model_info = model_map[key]\n if \"max_input_tokens\" in model_info:\n return model_info[\"max_input_tokens\"]\n if \"max_tokens\" in model_info:\n return model_info[\"max_tokens\"]\n except Exception:\n pass # Fall through to mapping\n\n # 3. Try our hardcoded mapping (longest match first)\n for pattern, limit in sorted(\n BEDROCK_MODEL_TOKEN_LIMITS.items(), key=lambda x: -len(x[0])\n ):\n if pattern in model_id_lower:\n return limit\n\n # 4. Default fallback\n return GEN_AI_MODEL_FALLBACK_MAX_TOKENS\n\n\ndef model_supports_image_input(model_name: str, model_provider: str) -> bool:\n # First, try to read an explicit configuration from the model_configuration table\n try:\n with get_session_with_current_tenant() as db_session:\n model_config = db_session.scalar(\n select(ModelConfiguration)\n .join(\n LLMProvider,\n ModelConfiguration.llm_provider_id == LLMProvider.id,\n )\n .where(\n ModelConfiguration.name == model_name,\n LLMProvider.provider == model_provider,\n )\n )\n if (\n model_config\n and LLMModelFlowType.VISION in model_config.llm_model_flow_types\n ):\n return True\n except Exception as e:\n logger.warning(\n f\"Failed to query database for {model_provider} model {model_name} image support: {e}\"\n )\n\n # Fallback to looking up the model in the litellm model_cost dict\n return litellm_thinks_model_supports_image_input(model_name, model_provider)\n\n\ndef litellm_thinks_model_supports_image_input(\n model_name: str, model_provider: str\n) -> bool:\n \"\"\"Generally should call `model_supports_image_input` unless you already know that\n `model_supports_image_input` from the DB is not set OR you need to avoid the performance\n hit of querying the DB.\"\"\"\n try:\n model_obj = find_model_obj(get_model_map(), model_provider, model_name)\n if not model_obj:\n logger.warning(\n f\"No litellm entry found for {model_provider}/{model_name}, this model may or may not support image input.\"\n )\n return False\n # The or False here is because sometimes the dict contains the key but the value is None\n return model_obj.get(\"supports_vision\", False) or False\n except Exception:\n logger.exception(\n f\"Failed to get model object for {model_provider}/{model_name}\"\n )\n return False\n\n\ndef model_is_reasoning_model(model_name: str, model_provider: str) -> bool:\n import litellm\n\n model_map = get_model_map()\n try:\n model_obj = find_model_obj(\n model_map,\n model_provider,\n model_name,\n )\n if model_obj and \"supports_reasoning\" in model_obj:\n return model_obj[\"supports_reasoning\"]\n\n # Fallback: try using litellm.supports_reasoning() for newer models\n try:\n # logger.debug(\"Falling back to `litellm.supports_reasoning`\")\n full_model_name = (\n f\"{model_provider}/{model_name}\"\n if model_provider not in model_name\n else model_name\n )\n return litellm.supports_reasoning(model=full_model_name)\n except Exception:\n logger.exception(\n f\"Failed to check if {model_provider}/{model_name} supports reasoning\"\n )\n return False\n\n except Exception:\n logger.exception(\n f\"Failed to get model object for {model_provider}/{model_name}\"\n )\n return False\n\n\ndef is_true_openai_model(model_provider: str, model_name: str) -> bool:\n \"\"\"\n Determines if a model is a true OpenAI model or just using OpenAI-compatible API.\n\n LiteLLM uses the \"openai\" provider for any OpenAI-compatible server (e.g. vLLM, LiteLLM proxy),\n but this function checks if the model is actually from OpenAI's model registry.\n\n This function is used primarily to determine if we should use the responses API.\n OpenAI models from OpenAI and Azure should use responses.\n \"\"\"\n\n if model_provider not in {\n LlmProviderNames.OPENAI,\n LlmProviderNames.LITELLM_PROXY,\n LlmProviderNames.AZURE,\n }:\n return False\n\n model_map = get_model_map()\n\n def _check_if_model_name_is_openai_provider(model_name: str) -> bool:\n if model_name not in model_map:\n return False\n return model_map[model_name].get(\"litellm_provider\") == LlmProviderNames.OPENAI\n\n try:\n # Check if any model exists in litellm's registry with openai prefix\n # If it's registered as \"openai/model-name\", it's a real OpenAI model\n if f\"{LlmProviderNames.OPENAI}/{model_name}\" in model_map:\n return True\n\n if _check_if_model_name_is_openai_provider(model_name):\n return True\n\n if model_name.startswith(f\"{LlmProviderNames.AZURE}/\"):\n model_name_with_azure_removed = \"/\".join(model_name.split(\"/\")[1:])\n if _check_if_model_name_is_openai_provider(model_name_with_azure_removed):\n return True\n\n return False\n\n except Exception:\n logger.exception(\n f\"Failed to determine if {model_provider}/{model_name} is a true OpenAI model\"\n )\n return False\n\n\ndef model_needs_formatting_reenabled(model_name: str) -> bool:\n # See https://simonwillison.net/tags/markdown/ for context on why this is needed\n # for OpenAI reasoning models to have correct markdown generation\n\n # Models that need formatting re-enabled\n model_names = [\"gpt-5.1\", \"gpt-5\", \"o3\", \"o1\"]\n\n # Pattern matches if any of these model names appear with word boundaries\n # Word boundaries include: start/end of string, space, hyphen, or forward slash\n pattern = (\n r\"(?:^|[\\s\\-/])(\"\n + \"|\".join(re.escape(name) for name in model_names)\n + r\")(?:$|[\\s\\-/])\"\n )\n\n if re.search(pattern, model_name):\n return True\n\n return False\n" }, { "path": "backend/onyx/llm/well_known_providers/auto_update_models.py", "content": "\"\"\"Pydantic models for GitHub-hosted Auto LLM configuration.\"\"\"\n\nfrom datetime import datetime\nfrom typing import Any\n\nfrom pydantic import BaseModel\nfrom pydantic import field_validator\n\nfrom onyx.llm.well_known_providers.models import SimpleKnownModel\n\n\nclass LLMProviderRecommendation(BaseModel):\n \"\"\"Configuration for a single provider in the GitHub config.\n\n Schema matches the plan:\n - default_model: The default model config (can be string or object with name)\n - additional_visible_models: List of additional visible model configs\n \"\"\"\n\n default_model: SimpleKnownModel\n additional_visible_models: list[SimpleKnownModel] = []\n\n @field_validator(\"default_model\", mode=\"before\")\n @classmethod\n def normalize_default_model(cls, v: Any) -> dict[str, Any]:\n \"\"\"Allow default_model to be a string (model name) or object.\"\"\"\n if isinstance(v, str):\n return {\"name\": v}\n return v\n\n\nclass LLMRecommendations(BaseModel):\n \"\"\"Root configuration object fetched from GitHub.\"\"\"\n\n version: str\n updated_at: datetime\n providers: dict[str, LLMProviderRecommendation]\n\n def get_visible_models(self, provider_name: str) -> list[SimpleKnownModel]:\n \"\"\"Get the set of models that should be visible by default for a provider.\"\"\"\n if provider_name in self.providers:\n provider_config = self.providers[provider_name]\n return [provider_config.default_model] + list(\n provider_config.additional_visible_models\n )\n return []\n\n def get_default_model(self, provider_name: str) -> SimpleKnownModel | None:\n \"\"\"Get the default model for a provider.\"\"\"\n if provider_name in self.providers:\n provider_config = self.providers[provider_name]\n return provider_config.default_model\n return None\n" }, { "path": "backend/onyx/llm/well_known_providers/auto_update_service.py", "content": "\"\"\"Service for fetching and syncing LLM model configurations from GitHub.\n\nThis service manages Auto mode LLM providers, where models and configuration\nare managed centrally via a GitHub-hosted JSON file. In Auto mode:\n- Model list is controlled by GitHub config\n- Model visibility is controlled by GitHub config\n- Default model is controlled by GitHub config\n- Admin only needs to provide API credentials\n\"\"\"\n\nfrom datetime import datetime\n\nimport httpx\nfrom sqlalchemy.orm import Session\n\nfrom onyx.cache.factory import get_cache_backend\nfrom onyx.configs.app_configs import AUTO_LLM_CONFIG_URL\nfrom onyx.db.llm import fetch_auto_mode_providers\nfrom onyx.db.llm import sync_auto_mode_models\nfrom onyx.llm.well_known_providers.auto_update_models import LLMRecommendations\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_CACHE_KEY_LAST_UPDATED_AT = \"auto_llm_update:last_updated_at\"\n_CACHE_TTL_SECONDS = 60 * 60 * 24 # 24 hours\n\n\ndef _get_cached_last_updated_at() -> datetime | None:\n try:\n value = get_cache_backend().get(_CACHE_KEY_LAST_UPDATED_AT)\n if value is not None:\n return datetime.fromisoformat(value.decode(\"utf-8\"))\n except Exception as e:\n logger.warning(f\"Failed to get cached last_updated_at: {e}\")\n return None\n\n\ndef _set_cached_last_updated_at(updated_at: datetime) -> None:\n try:\n get_cache_backend().set(\n _CACHE_KEY_LAST_UPDATED_AT,\n updated_at.isoformat(),\n ex=_CACHE_TTL_SECONDS,\n )\n except Exception as e:\n logger.warning(f\"Failed to set cached last_updated_at: {e}\")\n\n\ndef fetch_llm_recommendations_from_github(\n timeout: float = 30.0,\n) -> LLMRecommendations | None:\n \"\"\"Fetch LLM configuration from GitHub.\n\n Returns:\n GitHubLLMConfig if successful, None on error.\n \"\"\"\n if not AUTO_LLM_CONFIG_URL:\n logger.debug(\"AUTO_LLM_CONFIG_URL not configured, skipping fetch\")\n return None\n\n try:\n with httpx.Client(timeout=timeout) as client:\n response = client.get(AUTO_LLM_CONFIG_URL)\n response.raise_for_status()\n\n data = response.json()\n return LLMRecommendations.model_validate(data)\n except httpx.HTTPError as e:\n logger.error(f\"Failed to fetch LLM config from GitHub: {e}\")\n return None\n except Exception as e:\n logger.error(f\"Error parsing LLM config: {e}\")\n return None\n\n\ndef sync_llm_models_from_github(\n db_session: Session,\n force: bool = False,\n) -> dict[str, int]:\n \"\"\"Sync models from GitHub config to database for all Auto mode providers.\n\n In Auto mode, EVERYTHING is controlled by GitHub config:\n - Model list\n - Model visibility (is_visible)\n - Default model\n - Fast default model\n\n Args:\n db_session: Database session\n config: GitHub LLM configuration\n force: If True, skip the updated_at check and force sync\n\n Returns:\n Dict of provider_name -> number of changes made.\n \"\"\"\n results: dict[str, int] = {}\n\n # Get all providers in Auto mode\n auto_providers = fetch_auto_mode_providers(db_session)\n if not auto_providers:\n logger.debug(\"No providers in Auto mode found\")\n return {}\n\n # Fetch config from GitHub\n config = fetch_llm_recommendations_from_github()\n if not config:\n logger.warning(\"Failed to fetch GitHub config\")\n return {}\n\n # Skip if we've already processed this version (unless forced)\n last_updated_at = _get_cached_last_updated_at()\n if not force and last_updated_at and config.updated_at <= last_updated_at:\n logger.debug(\"GitHub config unchanged, skipping sync\")\n _set_cached_last_updated_at(config.updated_at)\n return {}\n\n for provider in auto_providers:\n provider_type = provider.provider # e.g., \"openai\", \"anthropic\"\n\n if provider_type not in config.providers:\n logger.debug(\n f\"No config for provider type '{provider_type}' in GitHub config\"\n )\n continue\n\n # Sync models - this replaces the model list entirely for Auto mode\n changes = sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config,\n )\n\n if changes > 0:\n results[provider.name] = changes\n logger.info(\n f\"Applied {changes} model changes to provider '{provider.name}'\"\n )\n\n _set_cached_last_updated_at(config.updated_at)\n return results\n\n\ndef reset_cache() -> None:\n \"\"\"Reset the cache timestamp. Useful for testing.\"\"\"\n try:\n get_cache_backend().delete(_CACHE_KEY_LAST_UPDATED_AT)\n except Exception as e:\n logger.warning(f\"Failed to reset cache: {e}\")\n" }, { "path": "backend/onyx/llm/well_known_providers/constants.py", "content": "from onyx.llm.constants import LlmProviderNames\n\nOPENAI_PROVIDER_NAME = \"openai\"\n\nBEDROCK_PROVIDER_NAME = \"bedrock\"\n\n\nOLLAMA_PROVIDER_NAME = \"ollama_chat\"\nOLLAMA_API_KEY_CONFIG_KEY = \"OLLAMA_API_KEY\"\n\nLM_STUDIO_PROVIDER_NAME = \"lm_studio\"\nLM_STUDIO_API_KEY_CONFIG_KEY = \"LM_STUDIO_API_KEY\"\n\nLITELLM_PROXY_PROVIDER_NAME = \"litellm_proxy\"\n\nBIFROST_PROVIDER_NAME = \"bifrost\"\n\n# Providers that use optional Bearer auth from custom_config\nPROVIDERS_WITH_SPECIAL_API_KEY_HANDLING: dict[str, str] = {\n LlmProviderNames.OLLAMA_CHAT: OLLAMA_API_KEY_CONFIG_KEY,\n LlmProviderNames.LM_STUDIO: LM_STUDIO_API_KEY_CONFIG_KEY,\n}\n\n# OpenRouter\nOPENROUTER_PROVIDER_NAME = \"openrouter\"\n\nANTHROPIC_PROVIDER_NAME = \"anthropic\"\n\nAZURE_PROVIDER_NAME = \"azure\"\n\n\nVERTEXAI_PROVIDER_NAME = \"vertex_ai\"\nVERTEX_CREDENTIALS_FILE_KWARG = \"vertex_credentials\"\nVERTEX_CREDENTIALS_FILE_KWARG_ENV_VAR_FORMAT = \"CREDENTIALS_FILE\"\nVERTEX_LOCATION_KWARG = \"vertex_location\"\n\nAWS_REGION_NAME_KWARG = \"aws_region_name\"\nAWS_REGION_NAME_KWARG_ENV_VAR_FORMAT = \"AWS_REGION_NAME\"\nAWS_BEARER_TOKEN_BEDROCK_KWARG_ENV_VAR_FORMAT = \"AWS_BEARER_TOKEN_BEDROCK\"\nAWS_ACCESS_KEY_ID_KWARG = \"aws_access_key_id\"\nAWS_ACCESS_KEY_ID_KWARG_ENV_VAR_FORMAT = \"AWS_ACCESS_KEY_ID\"\nAWS_SECRET_ACCESS_KEY_KWARG = \"aws_secret_access_key\"\nAWS_SECRET_ACCESS_KEY_KWARG_ENV_VAR_FORMAT = \"AWS_SECRET_ACCESS_KEY\"\n" }, { "path": "backend/onyx/llm/well_known_providers/llm_provider_options.py", "content": "import json\nimport pathlib\nimport threading\nimport time\n\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.constants import PROVIDER_DISPLAY_NAMES\nfrom onyx.llm.constants import WELL_KNOWN_PROVIDER_NAMES\nfrom onyx.llm.utils import get_max_input_tokens\nfrom onyx.llm.utils import model_supports_image_input\nfrom onyx.llm.well_known_providers.auto_update_models import LLMRecommendations\nfrom onyx.llm.well_known_providers.auto_update_service import (\n fetch_llm_recommendations_from_github,\n)\nfrom onyx.llm.well_known_providers.constants import ANTHROPIC_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import AZURE_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import BEDROCK_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import BIFROST_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import LITELLM_PROXY_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import LM_STUDIO_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import OLLAMA_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import OPENAI_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import OPENROUTER_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.constants import VERTEXAI_PROVIDER_NAME\nfrom onyx.llm.well_known_providers.models import WellKnownLLMProviderDescriptor\nfrom onyx.server.manage.llm.models import ModelConfigurationView\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_RECOMMENDATIONS_CACHE_TTL_SECONDS = 300\n_recommendations_cache_lock = threading.Lock()\n_cached_recommendations: LLMRecommendations | None = None\n_cached_recommendations_time: float = 0.0\n\n\ndef _get_provider_to_models_map() -> dict[str, list[str]]:\n \"\"\"Lazy-load provider model mappings to avoid importing litellm at module level.\n\n Dynamic providers (Bedrock, Ollama, OpenRouter) return empty lists here\n because their models are fetched directly from the source API, which is\n more up-to-date than LiteLLM's static lists.\n \"\"\"\n return {\n OPENAI_PROVIDER_NAME: get_openai_model_names(),\n BEDROCK_PROVIDER_NAME: [], # Dynamic - fetched from AWS API\n ANTHROPIC_PROVIDER_NAME: get_anthropic_model_names(),\n VERTEXAI_PROVIDER_NAME: get_vertexai_model_names(),\n OLLAMA_PROVIDER_NAME: [], # Dynamic - fetched from Ollama API\n LM_STUDIO_PROVIDER_NAME: [], # Dynamic - fetched from LM Studio API\n OPENROUTER_PROVIDER_NAME: [], # Dynamic - fetched from OpenRouter API\n LITELLM_PROXY_PROVIDER_NAME: [], # Dynamic - fetched from LiteLLM proxy API\n BIFROST_PROVIDER_NAME: [], # Dynamic - fetched from Bifrost API\n }\n\n\ndef _load_bundled_recommendations() -> LLMRecommendations:\n json_path = pathlib.Path(__file__).parent / \"recommended-models.json\"\n with open(json_path, \"r\") as f:\n json_config = json.load(f)\n return LLMRecommendations.model_validate(json_config)\n\n\ndef get_recommendations() -> LLMRecommendations:\n \"\"\"Get the recommendations, with an in-memory cache to avoid\n hitting GitHub on every API request.\"\"\"\n global _cached_recommendations, _cached_recommendations_time\n\n now = time.monotonic()\n if (\n _cached_recommendations is not None\n and (now - _cached_recommendations_time) < _RECOMMENDATIONS_CACHE_TTL_SECONDS\n ):\n return _cached_recommendations\n\n with _recommendations_cache_lock:\n # Double-check after acquiring lock\n if (\n _cached_recommendations is not None\n and (time.monotonic() - _cached_recommendations_time)\n < _RECOMMENDATIONS_CACHE_TTL_SECONDS\n ):\n return _cached_recommendations\n\n recommendations_from_github = fetch_llm_recommendations_from_github()\n result = recommendations_from_github or _load_bundled_recommendations()\n\n _cached_recommendations = result\n _cached_recommendations_time = time.monotonic()\n return result\n\n\ndef is_obsolete_model(model_name: str, provider: str) -> bool:\n \"\"\"Check if a model is obsolete and should be filtered out.\n\n Filters models that are 2+ major versions behind or deprecated.\n This is the single source of truth for obsolete model detection.\n \"\"\"\n model_lower = model_name.lower()\n\n # OpenAI obsolete models\n if provider == LlmProviderNames.OPENAI:\n # GPT-3 models are obsolete\n if \"gpt-3\" in model_lower:\n return True\n # Legacy models\n deprecated = {\n \"text-davinci-003\",\n \"text-davinci-002\",\n \"text-curie-001\",\n \"text-babbage-001\",\n \"text-ada-001\",\n \"davinci\",\n \"curie\",\n \"babbage\",\n \"ada\",\n }\n if model_lower in deprecated:\n return True\n\n # Anthropic obsolete models\n if provider == LlmProviderNames.ANTHROPIC:\n if \"claude-2\" in model_lower or \"claude-instant\" in model_lower:\n return True\n\n # Vertex AI obsolete models\n if provider == LlmProviderNames.VERTEX_AI:\n if \"gemini-1.0\" in model_lower:\n return True\n if \"palm\" in model_lower or \"bison\" in model_lower:\n return True\n\n return False\n\n\ndef get_openai_model_names() -> list[str]:\n \"\"\"Get OpenAI model names dynamically from litellm.\"\"\"\n import re\n import litellm\n\n # TODO: remove these lists once we have a comprehensive model configuration page\n # The ideal flow should be: fetch all available models --> filter by type\n # --> allow user to modify filters and select models based on current context\n non_chat_model_terms = {\n \"embed\",\n \"audio\",\n \"tts\",\n \"whisper\",\n \"dall-e\",\n \"image\",\n \"moderation\",\n \"sora\",\n \"container\",\n }\n deprecated_model_terms = {\"babbage\", \"davinci\", \"gpt-3.5\", \"gpt-4-\"}\n excluded_terms = non_chat_model_terms | deprecated_model_terms\n\n # NOTE: We are explicitly excluding all \"timestamped\" models\n # because they are mostly just noise in the admin configuration panel\n # e.g. gpt-4o-2025-07-16, gpt-3.5-turbo-0613, etc.\n date_pattern = re.compile(r\"-\\d{4}\")\n\n def is_valid_model(model: str) -> bool:\n model_lower = model.lower()\n return not any(\n ex in model_lower for ex in excluded_terms\n ) and not date_pattern.search(model)\n\n return sorted(\n (\n model.removeprefix(\"openai/\")\n for model in litellm.open_ai_chat_completion_models\n if is_valid_model(model)\n ),\n reverse=True,\n )\n\n\ndef get_anthropic_model_names() -> list[str]:\n \"\"\"Get Anthropic model names dynamically from litellm.\"\"\"\n import litellm\n\n # Models to exclude from Anthropic's model list (deprecated or duplicates)\n _IGNORABLE_ANTHROPIC_MODELS = {\n \"claude-2\",\n \"claude-instant-1\",\n \"anthropic/claude-3-5-sonnet-20241022\",\n }\n\n return sorted(\n [\n model\n for model in litellm.anthropic_models\n if model not in _IGNORABLE_ANTHROPIC_MODELS\n and not is_obsolete_model(model, LlmProviderNames.ANTHROPIC)\n ],\n reverse=True,\n )\n\n\ndef get_vertexai_model_names() -> list[str]:\n \"\"\"Get Vertex AI model names dynamically from litellm model_cost.\"\"\"\n import litellm\n\n # Combine all vertex model sets\n vertex_models: set[str] = set()\n vertex_model_sets = [\n \"vertex_chat_models\",\n \"vertex_language_models\",\n \"vertex_anthropic_models\",\n \"vertex_llama3_models\",\n \"vertex_mistral_models\",\n \"vertex_ai_ai21_models\",\n \"vertex_deepseek_models\",\n ]\n for attr in vertex_model_sets:\n if hasattr(litellm, attr):\n vertex_models.update(getattr(litellm, attr))\n\n # Also extract from model_cost for any models not in the sets\n for key in litellm.model_cost.keys():\n if key.startswith(\"vertex_ai/\"):\n model_name = key.replace(\"vertex_ai/\", \"\")\n vertex_models.add(model_name)\n\n return sorted(\n [\n model\n for model in vertex_models\n if \"embed\" not in model.lower()\n and \"image\" not in model.lower()\n and \"video\" not in model.lower()\n and \"code\" not in model.lower()\n and \"veo\" not in model.lower() # video generation\n and \"live\" not in model.lower() # live/streaming models\n and \"tts\" not in model.lower() # text-to-speech\n and \"native-audio\" not in model.lower() # audio models\n and \"/\" not in model # filter out prefixed models like openai/gpt-oss\n and \"search_api\" not in model.lower() # not a model\n and \"-maas\" not in model.lower() # marketplace models\n and not is_obsolete_model(model, LlmProviderNames.VERTEX_AI)\n ],\n reverse=True,\n )\n\n\ndef model_configurations_for_provider(\n provider_name: str, llm_recommendations: LLMRecommendations\n) -> list[ModelConfigurationView]:\n recommended_visible_models = llm_recommendations.get_visible_models(provider_name)\n recommended_visible_models_names = [m.name for m in recommended_visible_models]\n\n # Preserve provider-defined ordering while de-duplicating.\n model_names: list[str] = []\n seen_model_names: set[str] = set()\n for model_name in (\n fetch_models_for_provider(provider_name) + recommended_visible_models_names\n ):\n if model_name in seen_model_names:\n continue\n seen_model_names.add(model_name)\n model_names.append(model_name)\n\n # Vertex model list can be large and mixed-vendor; alphabetical ordering\n # makes model discovery easier in admin selection UIs.\n if provider_name == VERTEXAI_PROVIDER_NAME:\n model_names = sorted(model_names, key=str.lower)\n\n return [\n ModelConfigurationView(\n name=model_name,\n is_visible=model_name in recommended_visible_models_names,\n max_input_tokens=get_max_input_tokens(model_name, provider_name),\n supports_image_input=model_supports_image_input(model_name, provider_name),\n )\n for model_name in model_names\n ]\n\n\ndef fetch_available_well_known_llms() -> list[WellKnownLLMProviderDescriptor]:\n llm_recommendations = get_recommendations()\n\n well_known_llms = []\n for provider_name in WELL_KNOWN_PROVIDER_NAMES:\n model_configurations = model_configurations_for_provider(\n provider_name, llm_recommendations\n )\n well_known_llms.append(\n WellKnownLLMProviderDescriptor(\n name=provider_name,\n known_models=model_configurations,\n recommended_default_model=llm_recommendations.get_default_model(\n provider_name\n ),\n )\n )\n return well_known_llms\n\n\ndef fetch_models_for_provider(provider_name: str) -> list[str]:\n return _get_provider_to_models_map().get(provider_name, [])\n\n\ndef fetch_model_names_for_provider_as_set(provider_name: str) -> set[str] | None:\n model_names = fetch_models_for_provider(provider_name)\n return set(model_names) if model_names else None\n\n\ndef fetch_visible_model_names_for_provider_as_set(\n provider_name: str,\n) -> set[str] | None:\n \"\"\"Get visible model names for a provider.\n\n Note: Since we no longer maintain separate visible model lists,\n this returns all models (same as fetch_model_names_for_provider_as_set).\n Kept for backwards compatibility with alembic migrations.\n \"\"\"\n return fetch_model_names_for_provider_as_set(provider_name)\n\n\ndef get_provider_display_name(provider_name: str) -> str:\n \"\"\"Get human-friendly display name for an Onyx-supported provider.\n\n First checks Onyx-specific display names, then falls back to\n PROVIDER_DISPLAY_NAMES from constants.\n \"\"\"\n # Display names for Onyx-supported LLM providers (used in admin UI provider selection).\n # These override PROVIDER_DISPLAY_NAMES for Onyx-specific branding.\n _ONYX_PROVIDER_DISPLAY_NAMES: dict[str, str] = {\n OPENAI_PROVIDER_NAME: \"ChatGPT (OpenAI)\",\n OLLAMA_PROVIDER_NAME: \"Ollama\",\n LM_STUDIO_PROVIDER_NAME: \"LM Studio\",\n ANTHROPIC_PROVIDER_NAME: \"Claude (Anthropic)\",\n AZURE_PROVIDER_NAME: \"Azure OpenAI\",\n BEDROCK_PROVIDER_NAME: \"Amazon Bedrock\",\n VERTEXAI_PROVIDER_NAME: \"Google Vertex AI\",\n OPENROUTER_PROVIDER_NAME: \"OpenRouter\",\n LITELLM_PROXY_PROVIDER_NAME: \"LiteLLM Proxy\",\n }\n\n if provider_name in _ONYX_PROVIDER_DISPLAY_NAMES:\n return _ONYX_PROVIDER_DISPLAY_NAMES[provider_name]\n return PROVIDER_DISPLAY_NAMES.get(\n provider_name.lower(), provider_name.replace(\"_\", \" \").title()\n )\n\n\ndef fetch_default_model_for_provider(provider_name: str) -> str | None:\n \"\"\"Fetch the default model for a provider.\n\n First checks the GitHub-hosted recommended-models.json config (via fetch_github_config),\n then falls back to hardcoded defaults if unavailable.\n \"\"\"\n llm_recommendations = get_recommendations()\n default_model = llm_recommendations.get_default_model(provider_name)\n return default_model.name if default_model else None\n" }, { "path": "backend/onyx/llm/well_known_providers/models.py", "content": "from enum import Enum\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.server.manage.llm.models import ModelConfigurationView\n\n\nclass CustomConfigKeyType(str, Enum):\n # used for configuration values that require manual input\n # i.e., textual API keys (e.g., \"abcd1234\")\n TEXT_INPUT = \"text_input\"\n\n # used for configuration values that require a file to be selected/drag-and-dropped\n # i.e., file based credentials (e.g., \"/path/to/credentials/file.json\")\n FILE_INPUT = \"file_input\"\n\n # used for configuration values that require a selection from predefined options\n SELECT = \"select\"\n\n\nclass SimpleKnownModel(BaseModel):\n name: str\n display_name: str | None = None\n\n\nclass WellKnownLLMProviderDescriptor(BaseModel):\n name: str\n\n # NOTE: the recommended visible models are encoded in the known_models list\n known_models: list[ModelConfigurationView] = Field(default_factory=list)\n recommended_default_model: SimpleKnownModel | None = None\n" }, { "path": "backend/onyx/llm/well_known_providers/recommended-models.json", "content": "{\n \"version\": \"1.1\",\n \"updated_at\": \"2026-03-05T00:00:00Z\",\n \"providers\": {\n \"openai\": {\n \"default_model\": { \"name\": \"gpt-5.4\" },\n \"additional_visible_models\": [\n { \"name\": \"gpt-5.4\" },\n { \"name\": \"gpt-5.2\" }\n ]\n },\n \"anthropic\": {\n \"default_model\": \"claude-opus-4-6\",\n \"additional_visible_models\": [\n {\n \"name\": \"claude-opus-4-6\",\n \"display_name\": \"Claude Opus 4.6\"\n },\n {\n \"name\": \"claude-sonnet-4-6\",\n \"display_name\": \"Claude Sonnet 4.6\"\n },\n {\n \"name\": \"claude-opus-4-5\",\n \"display_name\": \"Claude Opus 4.5\"\n },\n {\n \"name\": \"claude-sonnet-4-5\",\n \"display_name\": \"Claude Sonnet 4.5\"\n }\n ]\n },\n \"vertex_ai\": {\n \"default_model\": \"gemini-3-pro-preview\",\n \"additional_visible_models\": [\n {\n \"name\": \"gemini-3-pro-preview\",\n \"display_name\": \"Gemini 3 Pro\"\n },\n {\n \"name\": \"gemini-3-flash-preview\",\n \"display_name\": \"Gemini 3 Flash\"\n }\n ]\n },\n \"openrouter\": {\n \"default_model\": \"z-ai/glm-4.7\",\n \"additional_visible_models\": [\n {\n \"name\": \"z-ai/glm-4.7\",\n \"display_name\": \"GLM 4.7\"\n },\n {\n \"name\": \"deepseek/deepseek-v3.2\",\n \"display_name\": \"DeepSeek V3.2\"\n },\n {\n \"name\": \"qwen/qwen3-235b-a22b-2507\",\n \"display_name\": \"Qwen3 235B A22B Instruct 2507\"\n },\n {\n \"name\": \"moonshotai/kimi-k2-0905\",\n \"display_name\": \"Kimi K2 0905\"\n }\n ]\n }\n }\n}\n" }, { "path": "backend/onyx/main.py", "content": "import logging\nimport sys\nimport traceback\nimport warnings\nfrom collections.abc import AsyncGenerator\nfrom contextlib import asynccontextmanager\nfrom typing import Any\nfrom typing import cast\n\nimport sentry_sdk\nimport uvicorn\nfrom fastapi import APIRouter\nfrom fastapi import FastAPI\nfrom fastapi import HTTPException\nfrom fastapi import Request\nfrom fastapi import status\nfrom fastapi.exceptions import RequestValidationError\nfrom fastapi.middleware.cors import CORSMiddleware\nfrom fastapi.responses import JSONResponse\nfrom fastapi.routing import APIRoute\nfrom httpx_oauth.clients.google import GoogleOAuth2\nfrom httpx_oauth.clients.openid import BASE_SCOPES\nfrom httpx_oauth.clients.openid import OpenID\nfrom sentry_sdk.integrations.fastapi import FastApiIntegration\nfrom sentry_sdk.integrations.starlette import StarletteIntegration\nfrom starlette.types import Lifespan\n\nfrom onyx import __version__\nfrom onyx.auth.schemas import UserCreate\nfrom onyx.auth.schemas import UserRead\nfrom onyx.auth.schemas import UserUpdate\nfrom onyx.auth.users import auth_backend\nfrom onyx.auth.users import create_onyx_oauth_router\nfrom onyx.auth.users import fastapi_users\nfrom onyx.cache.interface import CacheBackendType\nfrom onyx.configs.app_configs import APP_API_PREFIX\nfrom onyx.configs.app_configs import APP_HOST\nfrom onyx.configs.app_configs import APP_PORT\nfrom onyx.configs.app_configs import AUTH_RATE_LIMITING_ENABLED\nfrom onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.app_configs import CACHE_BACKEND\nfrom onyx.configs.app_configs import DISABLE_VECTOR_DB\nfrom onyx.configs.app_configs import LOG_ENDPOINT_LATENCY\nfrom onyx.configs.app_configs import OAUTH_CLIENT_ID\nfrom onyx.configs.app_configs import OAUTH_CLIENT_SECRET\nfrom onyx.configs.app_configs import OAUTH_ENABLED\nfrom onyx.configs.app_configs import OIDC_PKCE_ENABLED\nfrom onyx.configs.app_configs import OIDC_SCOPE_OVERRIDE\nfrom onyx.configs.app_configs import OPENID_CONFIG_URL\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_READ_ONLY_POOL_OVERFLOW\nfrom onyx.configs.app_configs import POSTGRES_API_SERVER_READ_ONLY_POOL_SIZE\nfrom onyx.configs.app_configs import SYSTEM_RECURSION_LIMIT\nfrom onyx.configs.app_configs import USER_AUTH_SECRET\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import AuthType\nfrom onyx.configs.constants import POSTGRES_WEB_APP_NAME\nfrom onyx.db.engine.async_sql_engine import get_sqlalchemy_async_engine\nfrom onyx.db.engine.connection_warmup import warm_up_connections\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.error_handling.exceptions import register_onyx_exception_handlers\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.hooks.registry import validate_registry\nfrom onyx.server.api_key.api import router as api_key_router\nfrom onyx.server.auth_check import check_router_auth\nfrom onyx.server.documents.cc_pair import router as cc_pair_router\nfrom onyx.server.documents.connector import router as connector_router\nfrom onyx.server.documents.credential import router as credential_router\nfrom onyx.server.documents.document import router as document_router\nfrom onyx.server.documents.standard_oauth import router as standard_oauth_router\nfrom onyx.server.features.build.api.api import public_build_router\nfrom onyx.server.features.build.api.api import router as build_router\nfrom onyx.server.features.default_assistant.api import (\n router as default_assistant_router,\n)\nfrom onyx.server.features.document_set.api import router as document_set_router\nfrom onyx.server.features.hierarchy.api import router as hierarchy_router\nfrom onyx.server.features.input_prompt.api import (\n admin_router as admin_input_prompt_router,\n)\nfrom onyx.server.features.input_prompt.api import (\n basic_router as input_prompt_router,\n)\nfrom onyx.server.features.mcp.api import admin_router as mcp_admin_router\nfrom onyx.server.features.mcp.api import router as mcp_router\nfrom onyx.server.features.notifications.api import router as notification_router\nfrom onyx.server.features.oauth_config.api import (\n admin_router as admin_oauth_config_router,\n)\nfrom onyx.server.features.oauth_config.api import router as oauth_config_router\nfrom onyx.server.features.password.api import router as password_router\nfrom onyx.server.features.persona.api import admin_agents_router\nfrom onyx.server.features.persona.api import admin_router as admin_persona_router\nfrom onyx.server.features.persona.api import agents_router\nfrom onyx.server.features.persona.api import basic_router as persona_router\nfrom onyx.server.features.projects.api import router as projects_router\nfrom onyx.server.features.tool.api import admin_router as admin_tool_router\nfrom onyx.server.features.tool.api import router as tool_router\nfrom onyx.server.features.user_oauth_token.api import router as user_oauth_token_router\nfrom onyx.server.features.web_search.api import router as web_search_router\nfrom onyx.server.federated.api import router as federated_router\nfrom onyx.server.kg.api import admin_router as kg_admin_router\nfrom onyx.server.manage.administrative import router as admin_router\nfrom onyx.server.manage.code_interpreter.api import (\n admin_router as code_interpreter_admin_router,\n)\nfrom onyx.server.manage.discord_bot.api import router as discord_bot_router\nfrom onyx.server.manage.embedding.api import admin_router as embedding_admin_router\nfrom onyx.server.manage.embedding.api import basic_router as embedding_router\nfrom onyx.server.manage.get_state import router as state_router\nfrom onyx.server.manage.image_generation.api import (\n admin_router as image_generation_admin_router,\n)\nfrom onyx.server.manage.llm.api import admin_router as llm_admin_router\nfrom onyx.server.manage.llm.api import basic_router as llm_router\nfrom onyx.server.manage.opensearch_migration.api import (\n admin_router as opensearch_migration_admin_router,\n)\nfrom onyx.server.manage.search_settings import router as search_settings_router\nfrom onyx.server.manage.slack_bot import router as slack_bot_management_router\nfrom onyx.server.manage.users import router as user_router\nfrom onyx.server.manage.voice.api import admin_router as voice_admin_router\nfrom onyx.server.manage.voice.user_api import router as voice_router\nfrom onyx.server.manage.voice.websocket_api import router as voice_websocket_router\nfrom onyx.server.manage.web_search.api import (\n admin_router as web_search_admin_router,\n)\nfrom onyx.server.metrics.postgres_connection_pool import (\n setup_postgres_connection_pool_metrics,\n)\nfrom onyx.server.metrics.prometheus_setup import setup_prometheus_metrics\nfrom onyx.server.middleware.latency_logging import add_latency_logging_middleware\nfrom onyx.server.middleware.rate_limiting import close_auth_limiter\nfrom onyx.server.middleware.rate_limiting import get_auth_rate_limiters\nfrom onyx.server.middleware.rate_limiting import setup_auth_limiter\nfrom onyx.server.onyx_api.ingestion import router as onyx_api_router\nfrom onyx.server.pat.api import router as pat_router\nfrom onyx.server.query_and_chat.chat_backend import router as chat_router\nfrom onyx.server.query_and_chat.query_backend import (\n admin_router as admin_query_router,\n)\nfrom onyx.server.query_and_chat.query_backend import basic_router as query_router\nfrom onyx.server.saml import router as saml_router\nfrom onyx.server.settings.api import admin_router as settings_admin_router\nfrom onyx.server.settings.api import basic_router as settings_router\nfrom onyx.server.token_rate_limits.api import (\n router as token_rate_limit_settings_router,\n)\nfrom onyx.server.utils import BasicAuthenticationError\nfrom onyx.setup import setup_multitenant_onyx\nfrom onyx.setup import setup_onyx\nfrom onyx.tracing.setup import setup_tracing\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.logger import setup_uvicorn_logger\nfrom onyx.utils.middleware import add_endpoint_context_middleware\nfrom onyx.utils.middleware import add_onyx_request_id_middleware\nfrom onyx.utils.telemetry import get_or_generate_uuid\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\nfrom shared_configs.configs import CORS_ALLOWED_ORIGIN\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import SENTRY_DSN\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nwarnings.filterwarnings(\n \"ignore\", category=ResourceWarning, message=r\"Unclosed client session\"\n)\nwarnings.filterwarnings(\n \"ignore\", category=ResourceWarning, message=r\"Unclosed connector\"\n)\n\nlogger = setup_logger()\n\nfile_handlers = [\n h for h in logger.logger.handlers if isinstance(h, logging.FileHandler)\n]\n\nsetup_uvicorn_logger(shared_file_handlers=file_handlers)\n\n\ndef validation_exception_handler(request: Request, exc: Exception) -> JSONResponse:\n if not isinstance(exc, RequestValidationError):\n logger.error(\n f\"Unexpected exception type in validation_exception_handler - {type(exc)}\"\n )\n raise exc\n\n exc_str = f\"{exc}\".replace(\"\\n\", \" \").replace(\" \", \" \")\n logger.exception(f\"{request}: {exc_str}\")\n content = {\"status_code\": 422, \"message\": exc_str, \"data\": None}\n return JSONResponse(content=content, status_code=422)\n\n\ndef value_error_handler(_: Request, exc: Exception) -> JSONResponse:\n if not isinstance(exc, ValueError):\n logger.error(f\"Unexpected exception type in value_error_handler - {type(exc)}\")\n raise exc\n\n try:\n raise (exc)\n except Exception:\n # log stacktrace\n logger.exception(\"ValueError\")\n return JSONResponse(\n status_code=400,\n content={\"message\": str(exc)},\n )\n\n\ndef use_route_function_names_as_operation_ids(app: FastAPI) -> None:\n \"\"\"\n OpenAPI generation defaults to naming the operation with the\n function + route + HTTP method, which usually looks very redundant.\n\n This function changes the operation IDs to be just the function name.\n\n Should be called only after all routes have been added.\n \"\"\"\n for route in app.routes:\n if isinstance(route, APIRoute):\n route.operation_id = route.name\n\n\ndef include_router_with_global_prefix_prepended(\n application: FastAPI, router: APIRouter, **kwargs: Any\n) -> None:\n \"\"\"Adds the global prefix to all routes in the router.\"\"\"\n processed_global_prefix = f\"/{APP_API_PREFIX.strip('/')}\" if APP_API_PREFIX else \"\"\n\n passed_in_prefix = cast(str | None, kwargs.get(\"prefix\"))\n if passed_in_prefix:\n final_prefix = f\"{processed_global_prefix}/{passed_in_prefix.strip('/')}\"\n else:\n final_prefix = f\"{processed_global_prefix}\"\n final_kwargs: dict[str, Any] = {\n **kwargs,\n \"prefix\": final_prefix,\n }\n\n application.include_router(router, **final_kwargs)\n\n\ndef include_auth_router_with_prefix(\n application: FastAPI,\n router: APIRouter,\n prefix: str | None = None,\n tags: list[str] | None = None,\n) -> None:\n \"\"\"Wrapper function to include an 'auth' router with prefix + rate-limiting dependencies.\"\"\"\n final_tags = tags or [\"auth\"]\n include_router_with_global_prefix_prepended(\n application,\n router,\n prefix=prefix,\n tags=final_tags,\n dependencies=get_auth_rate_limiters(),\n )\n\n\ndef validate_cache_backend_settings() -> None:\n \"\"\"Validate that CACHE_BACKEND=postgres is only used with DISABLE_VECTOR_DB.\n\n The Postgres cache backend eliminates the Redis dependency, but only works\n when Celery is not running (which requires DISABLE_VECTOR_DB=true).\n \"\"\"\n if CACHE_BACKEND == CacheBackendType.POSTGRES and not DISABLE_VECTOR_DB:\n raise RuntimeError(\n \"CACHE_BACKEND=postgres requires DISABLE_VECTOR_DB=true. \"\n \"The Postgres cache backend is only supported in no-vector-DB \"\n \"deployments where Celery is replaced by the in-process task runner.\"\n )\n\n\ndef validate_no_vector_db_settings() -> None:\n \"\"\"Validate that DISABLE_VECTOR_DB is not combined with incompatible settings.\n\n Raises RuntimeError if DISABLE_VECTOR_DB is set alongside MULTI_TENANT or ENABLE_CRAFT,\n since these modes require infrastructure that is removed in no-vector-DB deployments.\n \"\"\"\n if not DISABLE_VECTOR_DB:\n return\n\n if MULTI_TENANT:\n raise RuntimeError(\n \"DISABLE_VECTOR_DB cannot be used with MULTI_TENANT. \"\n \"Multi-tenant deployments require the vector database for \"\n \"per-tenant document indexing and search. Run in single-tenant \"\n \"mode when disabling the vector database.\"\n )\n\n from onyx.server.features.build.configs import ENABLE_CRAFT\n\n if ENABLE_CRAFT:\n raise RuntimeError(\n \"DISABLE_VECTOR_DB cannot be used with ENABLE_CRAFT. \"\n \"Onyx Craft requires background workers for sandbox lifecycle \"\n \"management, which are removed in no-vector-DB deployments. \"\n \"Disable Craft (ENABLE_CRAFT=false) when disabling the vector database.\"\n )\n\n\n@asynccontextmanager\nasync def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # noqa: ARG001\n validate_no_vector_db_settings()\n validate_cache_backend_settings()\n validate_registry()\n\n # Set recursion limit\n if SYSTEM_RECURSION_LIMIT is not None:\n sys.setrecursionlimit(SYSTEM_RECURSION_LIMIT)\n logger.notice(f\"System recursion limit set to {SYSTEM_RECURSION_LIMIT}\")\n\n SqlEngine.set_app_name(POSTGRES_WEB_APP_NAME)\n\n SqlEngine.init_engine(\n pool_size=POSTGRES_API_SERVER_POOL_SIZE,\n max_overflow=POSTGRES_API_SERVER_POOL_OVERFLOW,\n )\n SqlEngine.get_engine()\n\n SqlEngine.init_readonly_engine(\n pool_size=POSTGRES_API_SERVER_READ_ONLY_POOL_SIZE,\n max_overflow=POSTGRES_API_SERVER_READ_ONLY_POOL_OVERFLOW,\n )\n\n # Register pool metrics now that engines are created.\n # HTTP instrumentation is set up earlier in get_application() since it\n # adds middleware (which Starlette forbids after the app has started).\n setup_postgres_connection_pool_metrics(\n engines={\n \"sync\": SqlEngine.get_engine(),\n \"async\": get_sqlalchemy_async_engine(),\n \"readonly\": SqlEngine.get_readonly_engine(),\n },\n )\n\n verify_auth = fetch_versioned_implementation(\n \"onyx.auth.users\", \"verify_auth_setting\"\n )\n\n # Will throw exception if an issue is found\n verify_auth()\n\n if OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET:\n logger.notice(\"Both OAuth Client ID and Secret are configured.\")\n\n # Initialize tracing if credentials are provided\n setup_tracing()\n\n # fill up Postgres connection pools\n await warm_up_connections()\n\n if not MULTI_TENANT:\n # We cache this at the beginning so there is no delay in the first telemetry\n CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA)\n get_or_generate_uuid()\n\n # If we are multi-tenant, we need to only set up initial public tables\n with get_session_with_current_tenant() as db_session:\n setup_onyx(db_session, POSTGRES_DEFAULT_SCHEMA)\n # set up the file store (e.g. create bucket if needed). On multi-tenant,\n # this is done via IaC\n get_default_file_store().initialize()\n else:\n setup_multitenant_onyx()\n\n if not MULTI_TENANT:\n # don't emit a metric for every pod rollover/restart\n optional_telemetry(\n record_type=RecordType.VERSION, data={\"version\": __version__}\n )\n\n if AUTH_RATE_LIMITING_ENABLED:\n await setup_auth_limiter()\n\n if DISABLE_VECTOR_DB:\n from onyx.background.periodic_poller import recover_stuck_user_files\n from onyx.background.periodic_poller import start_periodic_poller\n\n recover_stuck_user_files(POSTGRES_DEFAULT_SCHEMA)\n start_periodic_poller(POSTGRES_DEFAULT_SCHEMA)\n\n yield\n\n if DISABLE_VECTOR_DB:\n from onyx.background.periodic_poller import stop_periodic_poller\n\n stop_periodic_poller()\n\n SqlEngine.reset_engine()\n\n if AUTH_RATE_LIMITING_ENABLED:\n await close_auth_limiter()\n\n\ndef log_http_error(request: Request, exc: Exception) -> JSONResponse:\n status_code = getattr(exc, \"status_code\", 500)\n\n if isinstance(exc, BasicAuthenticationError):\n # For BasicAuthenticationError, just log a brief message without stack trace\n # (almost always spammy)\n logger.debug(f\"Authentication failed: {str(exc)}\")\n\n elif status_code == 404 and request.url.path == \"/metrics\":\n # Log 404 errors for the /metrics endpoint with debug level\n logger.debug(f\"404 error for /metrics endpoint: {str(exc)}\")\n\n elif status_code >= 400:\n error_msg = f\"{str(exc)}\\n\"\n error_msg += \"\".join(traceback.format_tb(exc.__traceback__))\n logger.error(error_msg)\n\n detail = exc.detail if isinstance(exc, HTTPException) else str(exc)\n return JSONResponse(\n status_code=status_code,\n content={\"detail\": detail},\n )\n\n\ndef get_application(lifespan_override: Lifespan | None = None) -> FastAPI:\n application = FastAPI(\n title=\"Onyx Backend\",\n version=__version__,\n description=\"Onyx API for AI-powered chat with search, document indexing, agents, actions, and more\",\n servers=[\n {\"url\": f\"{WEB_DOMAIN.rstrip('/')}/api\", \"description\": \"Onyx API Server\"}\n ],\n lifespan=lifespan_override or lifespan,\n )\n if SENTRY_DSN:\n sentry_sdk.init(\n dsn=SENTRY_DSN,\n integrations=[StarletteIntegration(), FastApiIntegration()],\n traces_sample_rate=0.1,\n release=__version__,\n )\n logger.info(\"Sentry initialized\")\n else:\n logger.debug(\"Sentry DSN not provided, skipping Sentry initialization\")\n\n application.add_exception_handler(status.HTTP_400_BAD_REQUEST, log_http_error)\n application.add_exception_handler(status.HTTP_401_UNAUTHORIZED, log_http_error)\n application.add_exception_handler(status.HTTP_403_FORBIDDEN, log_http_error)\n application.add_exception_handler(status.HTTP_404_NOT_FOUND, log_http_error)\n application.add_exception_handler(\n status.HTTP_500_INTERNAL_SERVER_ERROR, log_http_error\n )\n\n register_onyx_exception_handlers(application)\n\n include_router_with_global_prefix_prepended(application, password_router)\n include_router_with_global_prefix_prepended(application, chat_router)\n include_router_with_global_prefix_prepended(application, query_router)\n include_router_with_global_prefix_prepended(application, document_router)\n include_router_with_global_prefix_prepended(application, user_router)\n include_router_with_global_prefix_prepended(application, admin_query_router)\n include_router_with_global_prefix_prepended(application, admin_router)\n include_router_with_global_prefix_prepended(application, connector_router)\n include_router_with_global_prefix_prepended(application, credential_router)\n include_router_with_global_prefix_prepended(application, input_prompt_router)\n include_router_with_global_prefix_prepended(application, admin_input_prompt_router)\n include_router_with_global_prefix_prepended(application, cc_pair_router)\n include_router_with_global_prefix_prepended(application, projects_router)\n include_router_with_global_prefix_prepended(application, public_build_router)\n include_router_with_global_prefix_prepended(application, build_router)\n include_router_with_global_prefix_prepended(application, document_set_router)\n include_router_with_global_prefix_prepended(application, hierarchy_router)\n include_router_with_global_prefix_prepended(application, search_settings_router)\n include_router_with_global_prefix_prepended(\n application, slack_bot_management_router\n )\n include_router_with_global_prefix_prepended(application, discord_bot_router)\n include_router_with_global_prefix_prepended(application, persona_router)\n include_router_with_global_prefix_prepended(application, admin_persona_router)\n include_router_with_global_prefix_prepended(application, agents_router)\n include_router_with_global_prefix_prepended(application, admin_agents_router)\n include_router_with_global_prefix_prepended(application, default_assistant_router)\n include_router_with_global_prefix_prepended(application, notification_router)\n include_router_with_global_prefix_prepended(application, tool_router)\n include_router_with_global_prefix_prepended(application, admin_tool_router)\n include_router_with_global_prefix_prepended(application, oauth_config_router)\n include_router_with_global_prefix_prepended(application, admin_oauth_config_router)\n include_router_with_global_prefix_prepended(application, user_oauth_token_router)\n include_router_with_global_prefix_prepended(application, state_router)\n include_router_with_global_prefix_prepended(application, onyx_api_router)\n include_router_with_global_prefix_prepended(application, settings_router)\n include_router_with_global_prefix_prepended(application, settings_admin_router)\n include_router_with_global_prefix_prepended(application, llm_admin_router)\n include_router_with_global_prefix_prepended(application, kg_admin_router)\n include_router_with_global_prefix_prepended(application, llm_router)\n include_router_with_global_prefix_prepended(\n application, code_interpreter_admin_router\n )\n include_router_with_global_prefix_prepended(\n application, image_generation_admin_router\n )\n include_router_with_global_prefix_prepended(application, embedding_admin_router)\n include_router_with_global_prefix_prepended(application, embedding_router)\n include_router_with_global_prefix_prepended(application, web_search_router)\n include_router_with_global_prefix_prepended(application, web_search_admin_router)\n include_router_with_global_prefix_prepended(application, voice_admin_router)\n include_router_with_global_prefix_prepended(application, voice_router)\n include_router_with_global_prefix_prepended(application, voice_websocket_router)\n include_router_with_global_prefix_prepended(\n application, opensearch_migration_admin_router\n )\n include_router_with_global_prefix_prepended(\n application, token_rate_limit_settings_router\n )\n include_router_with_global_prefix_prepended(application, api_key_router)\n include_router_with_global_prefix_prepended(application, standard_oauth_router)\n include_router_with_global_prefix_prepended(application, federated_router)\n include_router_with_global_prefix_prepended(application, mcp_router)\n include_router_with_global_prefix_prepended(application, mcp_admin_router)\n\n include_router_with_global_prefix_prepended(application, pat_router)\n\n if AUTH_TYPE == AuthType.BASIC or AUTH_TYPE == AuthType.CLOUD:\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_auth_router(auth_backend),\n prefix=\"/auth\",\n )\n\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_register_router(UserRead, UserCreate),\n prefix=\"/auth\",\n )\n\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_reset_password_router(),\n prefix=\"/auth\",\n )\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_verify_router(UserRead),\n prefix=\"/auth\",\n )\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_users_router(UserRead, UserUpdate),\n prefix=\"/users\",\n )\n\n # Register Google OAuth when AUTH_TYPE is GOOGLE_OAUTH, or when\n # AUTH_TYPE is BASIC and OAuth credentials are configured\n if AUTH_TYPE == AuthType.GOOGLE_OAUTH or (\n AUTH_TYPE == AuthType.BASIC and OAUTH_ENABLED\n ):\n oauth_client = GoogleOAuth2(\n OAUTH_CLIENT_ID,\n OAUTH_CLIENT_SECRET,\n scopes=[\"openid\", \"email\", \"profile\"],\n )\n include_auth_router_with_prefix(\n application,\n create_onyx_oauth_router(\n oauth_client,\n auth_backend,\n USER_AUTH_SECRET,\n associate_by_email=True,\n is_verified_by_default=True,\n redirect_url=f\"{WEB_DOMAIN}/auth/oauth/callback\",\n ),\n prefix=\"/auth/oauth\",\n )\n\n # Need logout router for GOOGLE_OAUTH only (BASIC already has it from above)\n if AUTH_TYPE == AuthType.GOOGLE_OAUTH:\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_logout_router(auth_backend),\n prefix=\"/auth\",\n )\n\n if AUTH_TYPE == AuthType.OIDC:\n # Ensure we request offline_access for refresh tokens\n try:\n oidc_scopes = list(OIDC_SCOPE_OVERRIDE or BASE_SCOPES)\n if \"offline_access\" not in oidc_scopes:\n oidc_scopes.append(\"offline_access\")\n except Exception as e:\n logger.warning(f\"Error configuring OIDC scopes: {e}\")\n # Fall back to default scopes if there's an error\n oidc_scopes = BASE_SCOPES\n\n include_auth_router_with_prefix(\n application,\n create_onyx_oauth_router(\n OpenID(\n OAUTH_CLIENT_ID,\n OAUTH_CLIENT_SECRET,\n OPENID_CONFIG_URL,\n # Use the configured scopes\n base_scopes=oidc_scopes,\n ),\n auth_backend,\n USER_AUTH_SECRET,\n associate_by_email=True,\n is_verified_by_default=True,\n redirect_url=f\"{WEB_DOMAIN}/auth/oidc/callback\",\n enable_pkce=OIDC_PKCE_ENABLED,\n ),\n prefix=\"/auth/oidc\",\n )\n\n # need basic auth router for `logout` endpoint\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_auth_router(auth_backend),\n prefix=\"/auth\",\n )\n\n elif AUTH_TYPE == AuthType.SAML:\n include_auth_router_with_prefix(\n application,\n saml_router,\n )\n\n if (\n AUTH_TYPE == AuthType.CLOUD\n or AUTH_TYPE == AuthType.BASIC\n or AUTH_TYPE == AuthType.GOOGLE_OAUTH\n or AUTH_TYPE == AuthType.OIDC\n ):\n # Add refresh token endpoint for OAuth as well\n include_auth_router_with_prefix(\n application,\n fastapi_users.get_refresh_router(auth_backend),\n prefix=\"/auth\",\n )\n\n application.add_exception_handler(\n RequestValidationError, validation_exception_handler\n )\n\n application.add_exception_handler(ValueError, value_error_handler)\n\n application.add_middleware(\n CORSMiddleware,\n allow_origins=CORS_ALLOWED_ORIGIN, # Configurable via environment variable\n allow_credentials=True,\n allow_methods=[\"*\"],\n allow_headers=[\"*\"],\n )\n if LOG_ENDPOINT_LATENCY:\n add_latency_logging_middleware(application, logger)\n\n add_onyx_request_id_middleware(application, \"API\", logger)\n\n # Set endpoint context for per-endpoint DB pool attribution metrics.\n # Must be registered after all routes are added.\n add_endpoint_context_middleware(application)\n\n # HTTP request metrics (latency histograms, in-progress gauge, slow request\n # counter). Must be called here — before the app starts — because the\n # instrumentator adds middleware via app.add_middleware().\n setup_prometheus_metrics(application)\n\n # Ensure all routes have auth enabled or are explicitly marked as public\n check_router_auth(application)\n\n use_route_function_names_as_operation_ids(application)\n\n return application\n\n\n# NOTE: needs to be outside of the `if __name__ == \"__main__\"` block so that the\n# app is exportable\nset_is_ee_based_on_env_variable()\napp = fetch_versioned_implementation(module=\"onyx.main\", attribute=\"get_application\")\n\n\nif __name__ == \"__main__\":\n logger.notice(\n f\"Starting Onyx Backend version {__version__} on http://{APP_HOST}:{str(APP_PORT)}/\"\n )\n\n if global_version.is_ee_version():\n logger.notice(\"Running Enterprise Edition\")\n\n uvicorn.run(app, host=APP_HOST, port=APP_PORT)\n" }, { "path": "backend/onyx/mcp_server/README.md", "content": "# Onyx MCP Server\n\n## Overview\n\nThe Onyx MCP server allows LLMs to connect to your Onyx instance and access its knowledge base and search capabilities through the [Model Context Protocol (MCP)](https://modelcontextprotocol.io/).\n\nWith the Onyx MCP Server, you can search your knowledgebase,\ngive your LLMs web search, and upload and manage documents in Onyx.\n\nAll access controls are managed within the main Onyx application.\n\n### Authentication\n\nProvide an Onyx Personal Access Token or API Key in the `Authorization` header as a Bearer token.\nThe MCP server quickly validates and passes through the token on every request.\n\nDepending on usage, the MCP Server may support OAuth and stdio in the future.\n\n### Default Configuration\n- **Transport**: HTTP POST (MCP over HTTP)\n- **Port**: 8090 (shares domain with API server)\n- **Framework**: FastMCP with FastAPI wrapper\n- **Database**: None (all work delegates to the API server)\n\n### Architecture\n\nThe MCP server is built on [FastMCP](https://github.com/jlowin/fastmcp) and runs alongside the main Onyx API server:\n\n```\n┌─────────────────┐\n│ LLM Client │\n│ (Claude, etc) │\n└────────┬────────┘\n │ MCP over HTTP\n │ (POST with bearer)\n ▼\n┌─────────────────┐\n│ MCP Server │\n│ Port 8090 │\n│ ├─ Auth │\n│ ├─ Tools │\n│ └─ Resources │\n└────────┬────────┘\n │ Internal HTTP\n │ (authenticated)\n ▼\n┌─────────────────┐\n│ API Server │\n│ Port 8080 │\n│ ├─ /me (auth) │\n│ ├─ Search APIs │\n│ └─ ACL checks │\n└─────────────────┘\n```\n\n## Configuring MCP Clients\n\n### Claude Desktop\n\nAdd to your Claude Desktop configuration (`~/Library/Application Support/Claude/claude_desktop_config.json` on macOS):\n\n```json\n{\n \"mcpServers\": {\n \"onyx\": {\n \"url\": \"https://[YOUR_ONYX_DOMAIN]:8090/\",\n \"transport\": \"http\",\n \"headers\": {\n \"Authorization\": \"Bearer YOUR_ONYX_TOKEN_HERE\"\n }\n }\n }\n}\n```\n\n### Other MCP Clients\n\nMost MCP clients support HTTP transport with custom headers. Refer to your client's documentation for configuration details.\n\n## Capabilities\n\n### Tools\n\nThe server provides three tools for searching and retrieving information:\n\n1. `search_indexed_documents`\nSearch the user's private knowledge base indexed in Onyx. Returns ranked documents with content snippets, scores, and metadata.\n\n2. `search_web`\nSearch the public internet for current events and general knowledge. Returns web search results with titles, URLs, and snippets.\n\n3. `open_urls`\nRetrieve the complete text content from specific web URLs. Useful for fetching full page content after finding relevant URLs via `search_web`.\n\n### Resources\n\n1. `indexed_sources`\nLists all document sources currently indexed in the tenant (e.g., `\"confluence\"`, `\"github\"`). Use these values to filter results when calling `search_indexed_documents`.\n\n## Local Development\n\n### Running the MCP Server\n\nThe MCP Server automatically launches with the `Run All Onyx Services` task from the default launch.json.\n\nYou can also independently launch the Server via the vscode debugger.\n\n### Testing with MCP Inspector\n\nThe [MCP Inspector](https://github.com/modelcontextprotocol/inspector) is a debugging tool for MCP servers:\n\n```bash\nnpx @modelcontextprotocol/inspector http://localhost:8090/\n```\n\n**Setup in Inspector:**\n\n1. Ignore the OAuth configuration menus\n2. Open the **Authentication** tab\n3. Select **Bearer Token** authentication\n4. Paste your Onyx bearer token\n5. Click **Connect**\n\nOnce connected, you can:\n- Browse available tools\n- Test tool calls with different parameters\n- View request/response payloads\n- Debug authentication issues\n\n### Health Check\n\nVerify the server is running:\n\n```bash\ncurl http://localhost:8090/health\n```\n\nExpected response:\n```json\n{\n \"status\": \"healthy\",\n \"service\": \"mcp_server\"\n}\n```\n\n### Environment Variables\n\n**MCP Server Configuration:**\n- `MCP_SERVER_ENABLED`: Enable MCP server (set to \"true\" to enable, default: disabled)\n- `MCP_SERVER_PORT`: Port for MCP server (default: 8090)\n- `MCP_SERVER_CORS_ORIGINS`: Comma-separated CORS origins (optional)\n\n**API Server Connection:**\n- `API_SERVER_PROTOCOL`: Protocol for API server connection (default: \"http\")\n- `API_SERVER_HOST`: Hostname for API server connection (default: \"127.0.0.1\")\n- `API_SERVER_URL_OVERRIDE_FOR_HTTP_REQUESTS`: Optional override URL. If set, takes precedence over the protocol/host variables. Used for self-hosting the MCP server with Onyx Cloud as the backend." }, { "path": "backend/onyx/mcp_server/api.py", "content": "\"\"\"MCP server with FastAPI wrapper.\"\"\"\n\nfrom collections.abc import AsyncGenerator\nfrom contextlib import asynccontextmanager\n\nfrom fastapi import FastAPI\nfrom fastapi import Request\nfrom fastapi.middleware.cors import CORSMiddleware\nfrom fastapi.responses import JSONResponse\nfrom fastapi.responses import Response\nfrom fastmcp import FastMCP\nfrom starlette.datastructures import MutableHeaders\nfrom starlette.middleware.base import RequestResponseEndpoint\nfrom starlette.types import Receive\nfrom starlette.types import Scope\nfrom starlette.types import Send\n\nfrom onyx.configs.app_configs import MCP_SERVER_CORS_ORIGINS\nfrom onyx.mcp_server.auth import OnyxTokenVerifier\nfrom onyx.mcp_server.utils import shutdown_http_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nlogger.info(\"Creating Onyx MCP Server...\")\n\nmcp_server = FastMCP(\n name=\"Onyx MCP Server\",\n version=\"1.0.0\",\n auth=OnyxTokenVerifier(),\n)\n\n# Import tools and resources AFTER mcp_server is created to avoid circular imports\n# Components register themselves via decorators on the shared mcp_server instance\nfrom onyx.mcp_server.tools import search # noqa: E402, F401\nfrom onyx.mcp_server.resources import indexed_sources # noqa: E402, F401\n\nlogger.info(\"MCP server instance created\")\n\n\ndef create_mcp_fastapi_app() -> FastAPI:\n \"\"\"Create FastAPI app wrapping MCP server with auth and shared client lifecycle.\"\"\"\n mcp_asgi_app = mcp_server.http_app(path=\"/\")\n\n async def _ensure_streamable_accept_header(\n scope: Scope, receive: Receive, send: Send\n ) -> None:\n \"\"\"Ensure Accept header includes types required by FastMCP streamable HTTP.\"\"\"\n if scope.get(\"type\") == \"http\":\n headers = MutableHeaders(scope=scope)\n accept = headers.get(\"accept\", \"\")\n accept_lower = accept.lower()\n\n if (\n not accept\n or accept == \"*/*\"\n or \"application/json\" not in accept_lower\n or \"text/event-stream\" not in accept_lower\n ):\n headers[\"accept\"] = \"application/json, text/event-stream\"\n\n await mcp_asgi_app(scope, receive, send)\n\n @asynccontextmanager\n async def combined_lifespan(app: FastAPI) -> AsyncGenerator[None, None]:\n \"\"\"Initializes MCP session manager.\"\"\"\n logger.info(\"MCP server starting up\")\n\n try:\n async with mcp_asgi_app.lifespan(app):\n yield\n finally:\n logger.info(\"MCP server shutting down\")\n await shutdown_http_client()\n\n app = FastAPI(\n title=\"Onyx MCP Server\",\n description=\"HTTP POST transport with bearer auth delegated to API /me\",\n version=\"1.0.0\",\n lifespan=combined_lifespan,\n )\n\n # Public health check endpoint (bypasses MCP auth)\n @app.middleware(\"http\")\n async def health_check(\n request: Request, call_next: RequestResponseEndpoint\n ) -> Response:\n if request.url.path.rstrip(\"/\") == \"/health\":\n return JSONResponse({\"status\": \"healthy\", \"service\": \"mcp_server\"})\n return await call_next(request)\n\n # Authentication is handled by FastMCP's OnyxTokenVerifier (see auth.py)\n\n if MCP_SERVER_CORS_ORIGINS:\n logger.info(f\"CORS origins: {MCP_SERVER_CORS_ORIGINS}\")\n app.add_middleware(\n CORSMiddleware,\n allow_origins=MCP_SERVER_CORS_ORIGINS,\n allow_credentials=True,\n allow_methods=[\"*\"],\n allow_headers=[\"*\"],\n )\n\n app.mount(\"/\", _ensure_streamable_accept_header)\n\n return app\n\n\nmcp_app = create_mcp_fastapi_app()\n" }, { "path": "backend/onyx/mcp_server/auth.py", "content": "\"\"\"Authentication helpers for the Onyx MCP server.\"\"\"\n\nfrom typing import Optional\n\nfrom fastmcp.server.auth.auth import AccessToken\nfrom fastmcp.server.auth.auth import TokenVerifier\n\nfrom onyx.mcp_server.utils import get_http_client\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import build_api_server_url_for_http_requests\n\nlogger = setup_logger()\n\n\nclass OnyxTokenVerifier(TokenVerifier):\n \"\"\"Validates bearer tokens by delegating to the API server.\"\"\"\n\n async def verify_token(self, token: str) -> Optional[AccessToken]:\n \"\"\"Call API /me to verify the token, return minimal AccessToken on success.\"\"\"\n try:\n response = await get_http_client().get(\n f\"{build_api_server_url_for_http_requests(respect_env_override_if_set=True)}/me\",\n headers={\"Authorization\": f\"Bearer {token}\"},\n )\n except Exception as exc:\n logger.error(\n \"MCP server failed to reach API /me for authentication: %s\",\n exc,\n exc_info=True,\n )\n return None\n\n if response.status_code != 200:\n logger.warning(\n \"API server rejected MCP auth token with status %s\",\n response.status_code,\n )\n return None\n\n return AccessToken(\n token=token,\n client_id=\"mcp\",\n scopes=[\"mcp:use\"],\n expires_at=None,\n resource=None,\n claims={},\n )\n" }, { "path": "backend/onyx/mcp_server/mcp.json.template", "content": "{\n \"mcpServers\": {\n \"Onyx\": {\n \"url\": \"https://cloud.onyx.app/mcp\",\n \"headers\": {\n \"Authorization\": \"Bearer [YOUR PAT OR API KEY HERE]\"\n }\n }\n }\n }" }, { "path": "backend/onyx/mcp_server/resources/__init__.py", "content": "\"\"\"Resource registrations for the Onyx MCP server.\"\"\"\n\n# Import resource modules so decorators execute when the package loads.\nfrom onyx.mcp_server.resources import indexed_sources # noqa: F401\n" }, { "path": "backend/onyx/mcp_server/resources/indexed_sources.py", "content": "\"\"\"Resources that expose metadata for the Onyx MCP server.\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any\n\nfrom onyx.mcp_server.api import mcp_server\nfrom onyx.mcp_server.utils import get_indexed_sources\nfrom onyx.mcp_server.utils import require_access_token\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n@mcp_server.resource(\n \"resource://indexed_sources\",\n name=\"indexed_sources\",\n description=(\n \"Enumerate the user's document sources that are currently indexed in Onyx.\"\n \"This can be used to discover filters for the `search_indexed_documents` tool.\"\n ),\n mime_type=\"application/json\",\n)\nasync def indexed_sources_resource() -> dict[str, Any]:\n \"\"\"Return the list of indexed source types for search filtering.\"\"\"\n\n access_token = require_access_token()\n\n sources = await get_indexed_sources(access_token)\n\n logger.info(\n \"Onyx MCP Server: indexed_sources resource returning %s entries\",\n len(sources),\n )\n\n return {\n \"indexed_sources\": sorted(sources),\n }\n" }, { "path": "backend/onyx/mcp_server/tools/__init__.py", "content": "\"\"\"Tool registrations for the Onyx MCP server.\"\"\"\n\n# Import tool modules so decorators execute when the package is imported.\nfrom onyx.mcp_server.tools import search # noqa: F401\n" }, { "path": "backend/onyx/mcp_server/tools/search.py", "content": "\"\"\"Search tools for MCP server - document and web search.\"\"\"\n\nfrom datetime import datetime\nfrom typing import Any\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.mcp_server.api import mcp_server\nfrom onyx.mcp_server.utils import get_http_client\nfrom onyx.mcp_server.utils import get_indexed_sources\nfrom onyx.mcp_server.utils import require_access_token\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import build_api_server_url_for_http_requests\nfrom onyx.utils.variable_functionality import global_version\n\nlogger = setup_logger()\n\n\n@mcp_server.tool()\nasync def search_indexed_documents(\n query: str,\n source_types: list[str] | None = None,\n time_cutoff: str | None = None,\n limit: int = 10,\n) -> dict[str, Any]:\n \"\"\"\n Search the user's knowledge base indexed in Onyx.\n Use this tool for information that is not public knowledge and specific to the user,\n their team, their work, or their organization/company.\n\n Note: In CE mode, this tool uses the chat endpoint internally which invokes an LLM\n on every call, consuming tokens and adding latency.\n Additionally, CE callers receive a truncated snippet (blurb) instead of a full document chunk,\n but this should still be sufficient for most use cases. CE mode functionality should be swapped\n when a dedicated CE search endpoint is implemented.\n\n In EE mode, the dedicated search endpoint is used instead.\n\n To find a list of available sources, use the `indexed_sources` resource.\n Returns chunks of text as search results with snippets, scores, and metadata.\n\n Example usage:\n ```\n {\n \"query\": \"What is the latest status of PROJ-1234 and what is the next development item?\",\n \"source_types\": [\"jira\", \"google_drive\", \"github\"],\n \"time_cutoff\": \"2025-11-24T00:00:00Z\",\n \"limit\": 10,\n }\n ```\n \"\"\"\n logger.info(\n f\"Onyx MCP Server: document search: query='{query}', sources={source_types}, limit={limit}\"\n )\n\n # Parse time_cutoff string to datetime if provided\n time_cutoff_dt: datetime | None = None\n if time_cutoff:\n try:\n time_cutoff_dt = datetime.fromisoformat(time_cutoff.replace(\"Z\", \"+00:00\"))\n except ValueError as e:\n logger.warning(\n f\"Onyx MCP Server: Invalid time_cutoff format '{time_cutoff}': {e}. Continuing without time filter.\"\n )\n # Continue with no time_cutoff instead of returning an error\n time_cutoff_dt = None\n\n # Initialize source_type_enums early to avoid UnboundLocalError\n source_type_enums: list[DocumentSource] | None = None\n\n # Get authenticated user from FastMCP's access token\n access_token = require_access_token()\n\n try:\n sources = await get_indexed_sources(access_token)\n except Exception as e:\n # Error fetching sources (network error, API failure, etc.)\n logger.error(\n \"Onyx MCP Server: Error checking indexed sources: %s\",\n e,\n exc_info=True,\n )\n return {\n \"documents\": [],\n \"total_results\": 0,\n \"query\": query,\n \"error\": (f\"Failed to check indexed sources: {str(e)}. \"),\n }\n\n if not sources:\n logger.info(\"Onyx MCP Server: No indexed sources available for tenant\")\n return {\n \"documents\": [],\n \"total_results\": 0,\n \"query\": query,\n \"message\": (\n \"No document sources are indexed yet. Add connectors or upload data \"\n \"through Onyx before calling onyx_search_documents.\"\n ),\n }\n\n # Convert source_types strings to DocumentSource enums if provided\n # Invalid values will be handled by the API server\n if source_types is not None:\n source_type_enums = []\n for src in source_types:\n try:\n source_type_enums.append(DocumentSource(src.lower()))\n except ValueError:\n logger.warning(\n f\"Onyx MCP Server: Invalid source type '{src}' - will be ignored by server\"\n )\n\n # Build filters dict only with non-None values\n filters: dict[str, Any] | None = None\n if source_type_enums or time_cutoff_dt:\n filters = {}\n if source_type_enums:\n filters[\"source_type\"] = [src.value for src in source_type_enums]\n if time_cutoff_dt:\n filters[\"time_cutoff\"] = time_cutoff_dt.isoformat()\n\n is_ee = global_version.is_ee_version()\n base_url = build_api_server_url_for_http_requests(respect_env_override_if_set=True)\n auth_headers = {\"Authorization\": f\"Bearer {access_token.token}\"}\n\n search_request: dict[str, Any]\n if is_ee:\n # EE: use the dedicated search endpoint (no LLM invocation)\n search_request = {\n \"search_query\": query,\n \"filters\": filters,\n \"num_docs_fed_to_llm_selection\": limit,\n \"run_query_expansion\": False,\n \"include_content\": True,\n \"stream\": False,\n }\n endpoint = f\"{base_url}/search/send-search-message\"\n error_key = \"error\"\n docs_key = \"search_docs\"\n content_field = \"content\"\n else:\n # CE: fall back to the chat endpoint (invokes LLM, consumes tokens)\n search_request = {\n \"message\": query,\n \"stream\": False,\n \"chat_session_info\": {},\n }\n if filters:\n search_request[\"internal_search_filters\"] = filters\n endpoint = f\"{base_url}/chat/send-chat-message\"\n error_key = \"error_msg\"\n docs_key = \"top_documents\"\n content_field = \"blurb\"\n\n try:\n response = await get_http_client().post(\n endpoint,\n json=search_request,\n headers=auth_headers,\n )\n response.raise_for_status()\n result = response.json()\n\n # Check for error in response\n if result.get(error_key):\n return {\n \"documents\": [],\n \"total_results\": 0,\n \"query\": query,\n \"error\": result.get(error_key),\n }\n\n documents = [\n {\n \"semantic_identifier\": doc.get(\"semantic_identifier\"),\n \"content\": doc.get(content_field),\n \"source_type\": doc.get(\"source_type\"),\n \"link\": doc.get(\"link\"),\n \"score\": doc.get(\"score\"),\n }\n for doc in result.get(docs_key, [])\n ]\n\n # NOTE: search depth is controlled by the backend persona defaults, not `limit`.\n # `limit` only caps the returned list; fewer results may be returned if the\n # backend retrieves fewer documents than requested.\n documents = documents[:limit]\n\n logger.info(\n f\"Onyx MCP Server: Internal search returned {len(documents)} results\"\n )\n return {\n \"documents\": documents,\n \"total_results\": len(documents),\n \"query\": query,\n }\n except Exception as e:\n logger.error(f\"Onyx MCP Server: Document search error: {e}\", exc_info=True)\n return {\n \"error\": f\"Document search failed: {str(e)}\",\n \"documents\": [],\n \"query\": query,\n }\n\n\n@mcp_server.tool()\nasync def search_web(\n query: str,\n limit: int = 5,\n) -> dict[str, Any]:\n \"\"\"\n Search the public internet for general knowledge, current events, and publicly available information.\n Use this tool for information that is publicly available on the web,\n such as news, documentation, general facts, or when the user's private knowledge base doesn't contain relevant information.\n\n Returns web search results with titles, URLs, and snippets (NOT full content). Use `open_urls` to fetch full page content.\n\n Example usage:\n ```\n {\n \"query\": \"React 19 migration guide to use react compiler\",\n \"limit\": 5\n }\n ```\n \"\"\"\n logger.info(f\"Onyx MCP Server: Web search: query='{query}', limit={limit}\")\n\n access_token = require_access_token()\n\n try:\n request_payload = {\"queries\": [query], \"max_results\": limit}\n response = await get_http_client().post(\n f\"{build_api_server_url_for_http_requests(respect_env_override_if_set=True)}/web-search/search-lite\",\n json=request_payload,\n headers={\"Authorization\": f\"Bearer {access_token.token}\"},\n )\n response.raise_for_status()\n response_payload = response.json()\n results = response_payload.get(\"results\", [])\n return {\n \"results\": results,\n \"query\": query,\n }\n except Exception as e:\n logger.error(f\"Onyx MCP Server: Web search error: {e}\", exc_info=True)\n return {\n \"error\": f\"Web search failed: {str(e)}\",\n \"results\": [],\n \"query\": query,\n }\n\n\n@mcp_server.tool()\nasync def open_urls(\n urls: list[str],\n) -> dict[str, Any]:\n \"\"\"\n Retrieve the complete text content from specific web URLs.\n Use this tool when you need to access full content from known URLs,\n such as documentation pages or articles returned by the `search_web` tool.\n\n Useful for following up on web search results when snippets do not provide enough information.\n\n Returns the full text content of each URL along with metadata like title and content type.\n\n Example usage:\n ```\n {\n \"urls\": [\"https://react.dev/versions\", \"https://react.dev/learn/react-compiler\",\"https://react.dev/learn/react-compiler/introduction\"]\n }\n ```\n \"\"\"\n logger.info(f\"Onyx MCP Server: Open URL: fetching {len(urls)} URLs\")\n\n access_token = require_access_token()\n\n try:\n response = await get_http_client().post(\n f\"{build_api_server_url_for_http_requests(respect_env_override_if_set=True)}/web-search/open-urls\",\n json={\"urls\": urls},\n headers={\"Authorization\": f\"Bearer {access_token.token}\"},\n )\n response.raise_for_status()\n response_payload = response.json()\n results = response_payload.get(\"results\", [])\n return {\n \"results\": results,\n }\n except Exception as e:\n logger.error(f\"Onyx MCP Server: URL fetch error: {e}\", exc_info=True)\n return {\n \"error\": f\"URL fetch failed: {str(e)}\",\n \"results\": [],\n }\n" }, { "path": "backend/onyx/mcp_server/utils.py", "content": "\"\"\"Utility helpers for the Onyx MCP server.\"\"\"\n\nfrom __future__ import annotations\n\nimport httpx\nfrom fastmcp.server.auth.auth import AccessToken\nfrom fastmcp.server.dependencies import get_access_token\n\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import build_api_server_url_for_http_requests\n\nlogger = setup_logger()\n\n# Shared HTTP client reused across requests\n_http_client: httpx.AsyncClient | None = None\n\n\ndef require_access_token() -> AccessToken:\n \"\"\"\n Get and validate the access token from the current request.\n\n Raises:\n ValueError: If no access token is present in the request.\n\n Returns:\n AccessToken: The validated access token.\n \"\"\"\n access_token = get_access_token()\n if not access_token:\n raise ValueError(\n \"MCP Server requires an Onyx access token to authenticate your request\"\n )\n return access_token\n\n\ndef get_http_client() -> httpx.AsyncClient:\n \"\"\"Return a shared async HTTP client.\"\"\"\n global _http_client\n if _http_client is None:\n _http_client = httpx.AsyncClient(timeout=60.0)\n return _http_client\n\n\nasync def shutdown_http_client() -> None:\n \"\"\"Close the shared HTTP client when the server shuts down.\"\"\"\n global _http_client\n if _http_client is not None:\n await _http_client.aclose()\n _http_client = None\n\n\nasync def get_indexed_sources(\n access_token: AccessToken,\n) -> list[str]:\n \"\"\"\n Fetch indexed document sources for the current user/tenant.\n\n Returns:\n List of indexed source strings. Empty list if no sources are indexed.\n \"\"\"\n headers = {\"Authorization\": f\"Bearer {access_token.token}\"}\n try:\n response = await get_http_client().get(\n f\"{build_api_server_url_for_http_requests(respect_env_override_if_set=True)}/manage/indexed-sources\",\n headers=headers,\n )\n response.raise_for_status()\n payload = response.json()\n sources = payload.get(\"sources\", [])\n if not isinstance(sources, list):\n raise ValueError(\"Unexpected response shape for indexed sources\")\n return [str(source) for source in sources]\n except (httpx.HTTPStatusError, httpx.RequestError, ValueError):\n # Re-raise known exception types (httpx errors and validation errors)\n logger.error(\n \"Onyx MCP Server: Failed to fetch indexed sources\",\n exc_info=True,\n )\n raise\n except Exception as exc:\n # Wrap unexpected exceptions\n logger.error(\n \"Onyx MCP Server: Unexpected error fetching indexed sources\",\n exc_info=True,\n )\n raise RuntimeError(f\"Failed to fetch indexed sources: {exc}\") from exc\n" }, { "path": "backend/onyx/mcp_server_main.py", "content": "\"\"\"Entry point for MCP server - HTTP POST transport with API key auth.\"\"\"\n\nimport uvicorn\n\nfrom onyx.configs.app_configs import MCP_SERVER_ENABLED\nfrom onyx.configs.app_configs import MCP_SERVER_HOST\nfrom onyx.configs.app_configs import MCP_SERVER_PORT\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef main() -> None:\n \"\"\"Run the MCP server.\"\"\"\n if not MCP_SERVER_ENABLED:\n logger.info(\"MCP server is disabled (MCP_SERVER_ENABLED=false)\")\n return\n\n logger.info(f\"Starting MCP server on {MCP_SERVER_HOST}:{MCP_SERVER_PORT}\")\n\n from onyx.mcp_server.api import mcp_app\n\n uvicorn.run(\n mcp_app,\n host=MCP_SERVER_HOST,\n port=MCP_SERVER_PORT,\n log_config=None,\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/natural_language_processing/__init__.py", "content": "" }, { "path": "backend/onyx/natural_language_processing/constants.py", "content": "\"\"\"\nConstants for natural language processing, including embedding and reranking models.\n\nThis file contains constants moved from model_server to support the gradual migration\nof API-based calls to bypass the model server.\n\"\"\"\n\nfrom shared_configs.enums import EmbeddingProvider\nfrom shared_configs.enums import EmbedTextType\n\n\n# Default model names for different providers\nDEFAULT_OPENAI_MODEL = \"text-embedding-3-small\"\nDEFAULT_COHERE_MODEL = \"embed-english-light-v3.0\"\nDEFAULT_VOYAGE_MODEL = \"voyage-large-2-instruct\"\nDEFAULT_VERTEX_MODEL = \"text-embedding-005\"\n\n\nclass EmbeddingModelTextType:\n \"\"\"Mapping of Onyx text types to provider-specific text types.\"\"\"\n\n PROVIDER_TEXT_TYPE_MAP = {\n EmbeddingProvider.COHERE: {\n EmbedTextType.QUERY: \"search_query\",\n EmbedTextType.PASSAGE: \"search_document\",\n },\n EmbeddingProvider.VOYAGE: {\n EmbedTextType.QUERY: \"query\",\n EmbedTextType.PASSAGE: \"document\",\n },\n EmbeddingProvider.GOOGLE: {\n EmbedTextType.QUERY: \"RETRIEVAL_QUERY\",\n EmbedTextType.PASSAGE: \"RETRIEVAL_DOCUMENT\",\n },\n }\n\n @staticmethod\n def get_type(provider: EmbeddingProvider, text_type: EmbedTextType) -> str:\n \"\"\"Get provider-specific text type string.\"\"\"\n return EmbeddingModelTextType.PROVIDER_TEXT_TYPE_MAP[provider][text_type]\n" }, { "path": "backend/onyx/natural_language_processing/english_stopwords.py", "content": "import re\n\nENGLISH_STOPWORDS = [\n \"a\",\n \"about\",\n \"above\",\n \"after\",\n \"again\",\n \"against\",\n \"ain\",\n \"all\",\n \"am\",\n \"an\",\n \"and\",\n \"any\",\n \"are\",\n \"aren\",\n \"aren't\",\n \"as\",\n \"at\",\n \"be\",\n \"because\",\n \"been\",\n \"before\",\n \"being\",\n \"below\",\n \"between\",\n \"both\",\n \"but\",\n \"by\",\n \"can\",\n \"couldn\",\n \"couldn't\",\n \"d\",\n \"did\",\n \"didn\",\n \"didn't\",\n \"do\",\n \"does\",\n \"doesn\",\n \"doesn't\",\n \"doing\",\n \"don\",\n \"don't\",\n \"down\",\n \"during\",\n \"each\",\n \"few\",\n \"for\",\n \"from\",\n \"further\",\n \"had\",\n \"hadn\",\n \"hadn't\",\n \"has\",\n \"hasn\",\n \"hasn't\",\n \"have\",\n \"haven\",\n \"haven't\",\n \"having\",\n \"he\",\n \"he'd\",\n \"he'll\",\n \"he's\",\n \"her\",\n \"here\",\n \"hers\",\n \"herself\",\n \"him\",\n \"himself\",\n \"his\",\n \"how\",\n \"i\",\n \"i'd\",\n \"i'll\",\n \"i'm\",\n \"i've\",\n \"if\",\n \"in\",\n \"into\",\n \"is\",\n \"isn\",\n \"isn't\",\n \"it\",\n \"it'd\",\n \"it'll\",\n \"it's\",\n \"its\",\n \"itself\",\n \"just\",\n \"ll\",\n \"m\",\n \"ma\",\n \"me\",\n \"mightn\",\n \"mightn't\",\n \"more\",\n \"most\",\n \"mustn\",\n \"mustn't\",\n \"my\",\n \"myself\",\n \"needn\",\n \"needn't\",\n \"no\",\n \"nor\",\n \"not\",\n \"now\",\n \"o\",\n \"of\",\n \"off\",\n \"on\",\n \"once\",\n \"only\",\n \"or\",\n \"other\",\n \"our\",\n \"ours\",\n \"ourselves\",\n \"out\",\n \"over\",\n \"own\",\n \"re\",\n \"s\",\n \"same\",\n \"shan\",\n \"shan't\",\n \"she\",\n \"she'd\",\n \"she'll\",\n \"she's\",\n \"should\",\n \"should've\",\n \"shouldn\",\n \"shouldn't\",\n \"so\",\n \"some\",\n \"such\",\n \"t\",\n \"than\",\n \"that\",\n \"that'll\",\n \"the\",\n \"their\",\n \"theirs\",\n \"them\",\n \"themselves\",\n \"then\",\n \"there\",\n \"these\",\n \"they\",\n \"they'd\",\n \"they'll\",\n \"they're\",\n \"they've\",\n \"this\",\n \"those\",\n \"through\",\n \"to\",\n \"too\",\n \"under\",\n \"until\",\n \"up\",\n \"ve\",\n \"very\",\n \"was\",\n \"wasn\",\n \"wasn't\",\n \"we\",\n \"we'd\",\n \"we'll\",\n \"we're\",\n \"we've\",\n \"were\",\n \"weren\",\n \"weren't\",\n \"what\",\n \"when\",\n \"where\",\n \"which\",\n \"while\",\n \"who\",\n \"whom\",\n \"why\",\n \"will\",\n \"with\",\n \"won\",\n \"won't\",\n \"wouldn\",\n \"wouldn't\",\n \"y\",\n \"you\",\n \"you'd\",\n \"you'll\",\n \"you're\",\n \"you've\",\n \"your\",\n \"yours\",\n \"yourself\",\n \"yourselves\",\n]\n\nENGLISH_STOPWORDS_SET = frozenset(ENGLISH_STOPWORDS)\n\n\ndef strip_stopwords(text: str) -> list[str]:\n \"\"\"Remove English stopwords from text.\n\n Matching is case-insensitive and ignores leading/trailing punctuation\n on each word. Internal punctuation (like apostrophes in contractions)\n is preserved for matching, so \"you're\" matches the stopword \"you're\"\n but \"youre\" would not.\n \"\"\"\n words = text.split()\n result = []\n\n for word in words:\n # Strip leading/trailing punctuation to get the core word for comparison\n # This preserves internal punctuation like apostrophes\n core = re.sub(r\"^[^\\w']+|[^\\w']+$\", \"\", word)\n if core.lower() not in ENGLISH_STOPWORDS_SET:\n result.append(word)\n\n return result\n" }, { "path": "backend/onyx/natural_language_processing/exceptions.py", "content": "class ModelServerRateLimitError(Exception):\n \"\"\"\n Exception raised for rate limiting errors from the model server.\n \"\"\"\n\n\nclass CohereBillingLimitError(Exception):\n \"\"\"\n Raised when Cohere rejects requests because the billing cap is reached.\n \"\"\"\n" }, { "path": "backend/onyx/natural_language_processing/search_nlp_models.py", "content": "import asyncio\nimport json\nimport os\nimport threading\nimport time\nfrom collections.abc import Callable\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import ThreadPoolExecutor\nfrom functools import partial\nfrom functools import wraps\nfrom types import TracebackType\nfrom typing import Any\nfrom typing import cast\n\nimport aioboto3 # type: ignore\nimport httpx\nimport requests\nimport voyageai # type: ignore[import-untyped]\nfrom cohere import AsyncClient as CohereAsyncClient\nfrom cohere.core.api_error import ApiError\nfrom google.oauth2 import service_account\nfrom httpx import HTTPError\nfrom requests import JSONDecodeError\nfrom requests import RequestException\nfrom requests import Response\nfrom retry import retry\n\nfrom onyx.configs.app_configs import INDEXING_EMBEDDING_MODEL_NUM_THREADS\nfrom onyx.configs.app_configs import LARGE_CHUNK_RATIO\nfrom onyx.configs.model_configs import BATCH_SIZE_ENCODE_CHUNKS\nfrom onyx.configs.model_configs import (\n BATCH_SIZE_ENCODE_CHUNKS_FOR_API_EMBEDDING_SERVICES,\n)\nfrom onyx.connectors.models import ConnectorStopSignal\nfrom onyx.db.models import SearchSettings\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.natural_language_processing.constants import DEFAULT_COHERE_MODEL\nfrom onyx.natural_language_processing.constants import DEFAULT_OPENAI_MODEL\nfrom onyx.natural_language_processing.constants import DEFAULT_VERTEX_MODEL\nfrom onyx.natural_language_processing.constants import DEFAULT_VOYAGE_MODEL\nfrom onyx.natural_language_processing.constants import EmbeddingModelTextType\nfrom onyx.natural_language_processing.exceptions import CohereBillingLimitError\nfrom onyx.natural_language_processing.exceptions import ModelServerRateLimitError\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.natural_language_processing.utils import tokenizer_trim_content\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.search_nlp_models_utils import pass_aws_key\nfrom onyx.utils.text_processing import remove_invalid_unicode_chars\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.configs import API_BASED_EMBEDDING_TIMEOUT\nfrom shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE\nfrom shared_configs.configs import INDEXING_ONLY\nfrom shared_configs.configs import MODEL_SERVER_HOST\nfrom shared_configs.configs import MODEL_SERVER_PORT\nfrom shared_configs.configs import OPENAI_EMBEDDING_TIMEOUT\nfrom shared_configs.configs import SKIP_WARM_UP\nfrom shared_configs.configs import VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE\nfrom shared_configs.enums import EmbeddingProvider\nfrom shared_configs.enums import EmbedTextType\nfrom shared_configs.enums import RerankerProvider\nfrom shared_configs.model_server_models import Embedding\nfrom shared_configs.model_server_models import EmbedRequest\nfrom shared_configs.model_server_models import EmbedResponse\nfrom shared_configs.model_server_models import IntentRequest\nfrom shared_configs.model_server_models import IntentResponse\nfrom shared_configs.model_server_models import RerankRequest\nfrom shared_configs.model_server_models import RerankResponse\nfrom shared_configs.utils import batch_list\n\nlogger = setup_logger()\n\n# If we are not only indexing, dont want retry very long\n_RETRY_DELAY = 10 if INDEXING_ONLY else 0.1\n_RETRY_TRIES = 10 if INDEXING_ONLY else 2\n\n# OpenAI only allows 2048 embeddings to be computed at once\n_OPENAI_MAX_INPUT_LEN = 2048\n# Cohere allows up to 96 embeddings in a single embedding calling\n_COHERE_MAX_INPUT_LEN = 96\n\n# Authentication error string constants\n_AUTH_ERROR_401 = \"401\"\n_AUTH_ERROR_UNAUTHORIZED = \"unauthorized\"\n_AUTH_ERROR_INVALID_API_KEY = \"invalid api key\"\n_AUTH_ERROR_PERMISSION = \"permission\"\n\n# Thread-local storage for event loops\n# This prevents creating thousands of event loops during batch processing,\n# which was causing severe memory leaks with API-based embedding providers\n_thread_local = threading.local()\n\n\ndef _get_or_create_event_loop() -> asyncio.AbstractEventLoop:\n \"\"\"Get or create a thread-local event loop for API embedding calls.\n\n This prevents creating a new event loop for every batch during embedding,\n which was causing memory leaks. Instead, each thread reuses the same loop.\n\n Returns:\n asyncio.AbstractEventLoop: The thread-local event loop\n \"\"\"\n if (\n not hasattr(_thread_local, \"loop\")\n or _thread_local.loop is None\n or _thread_local.loop.is_closed()\n ):\n _thread_local.loop = asyncio.new_event_loop()\n asyncio.set_event_loop(_thread_local.loop)\n return _thread_local.loop\n\n\ndef cleanup_embedding_thread_locals() -> None:\n \"\"\"Clean up thread-local event loops to prevent memory leaks.\n\n This should be called after each task completes to ensure that\n event loops and their associated resources are properly released.\n Thread-local storage persists across Celery tasks when using the\n thread pool, so explicit cleanup is necessary.\n\n NOTE: This must be called from the SAME thread that created the event loop.\n For ThreadPoolExecutor-based embedding, this cleanup happens automatically\n via the _cleanup_thread_local wrapper.\n \"\"\"\n if hasattr(_thread_local, \"loop\") and _thread_local.loop is not None:\n loop = _thread_local.loop\n if not loop.is_closed():\n # Cancel all pending tasks in the event loop\n try:\n # Ensure loop is set as current event loop before accessing tasks\n asyncio.set_event_loop(loop)\n pending = asyncio.all_tasks(loop)\n if pending:\n logger.debug(\n f\"Cleaning up event loop with {len(pending)} pending tasks in thread {threading.current_thread().name}\"\n )\n for task in pending:\n task.cancel()\n # Run the loop briefly to allow cancelled tasks to complete\n loop.run_until_complete(\n asyncio.gather(*pending, return_exceptions=True)\n )\n except Exception as e:\n # If gathering tasks fails, just close the loop\n logger.debug(f\"Error gathering tasks during cleanup: {e}\")\n\n # Close the event loop\n loop.close()\n logger.debug(\n f\"Closed event loop in thread {threading.current_thread().name}\"\n )\n\n # Clear the thread-local reference\n _thread_local.loop = None\n\n\ndef _cleanup_thread_local(func: Callable) -> Callable:\n \"\"\"Decorator to ensure thread-local cleanup after function execution.\n\n This wraps functions that run in ThreadPoolExecutor threads to ensure\n that thread-local event loops are cleaned up after each execution,\n preventing memory leaks from persistent thread-local storage.\n \"\"\"\n\n @wraps(func)\n def wrapper(*args: Any, **kwargs: Any) -> Any:\n try:\n return func(*args, **kwargs)\n finally:\n # Clean up thread-local event loop after this thread's work is done\n cleanup_embedding_thread_locals()\n\n return wrapper\n\n\nWARM_UP_STRINGS = [\n \"Onyx is amazing!\",\n \"Check out our easy deployment guide at\",\n \"https://docs.onyx.app/deployment/getting_started/quickstart\",\n]\n\n\ndef clean_model_name(model_str: str) -> str:\n return model_str.replace(\"/\", \"_\").replace(\"-\", \"_\").replace(\".\", \"_\")\n\n\ndef build_model_server_url(\n model_server_host: str,\n model_server_port: int,\n) -> str:\n model_server_url = f\"{model_server_host}:{model_server_port}\"\n\n # use protocol if provided\n if \"http\" in model_server_url:\n return model_server_url\n\n # otherwise default to http\n return f\"http://{model_server_url}\"\n\n\ndef is_authentication_error(error: Exception) -> bool:\n \"\"\"Check if an exception is related to authentication issues.\n\n Args:\n error: The exception to check\n\n Returns:\n bool: True if the error appears to be authentication-related\n \"\"\"\n error_str = str(error).lower()\n return (\n _AUTH_ERROR_401 in error_str\n or _AUTH_ERROR_UNAUTHORIZED in error_str\n or _AUTH_ERROR_INVALID_API_KEY in error_str\n or _AUTH_ERROR_PERMISSION in error_str\n )\n\n\ndef format_embedding_error(\n error: Exception,\n service_name: str,\n model: str | None,\n provider: EmbeddingProvider,\n sanitized_api_key: str | None = None,\n status_code: int | None = None,\n) -> str:\n \"\"\"\n Format a standardized error string for embedding errors.\n \"\"\"\n detail = f\"Status {status_code}\" if status_code else f\"{type(error)}\"\n\n return (\n f\"{'HTTP error' if status_code else 'Exception'} embedding text with {service_name} - {detail}: \"\n f\"Model: {model} \"\n f\"Provider: {provider} \"\n f\"API Key: {sanitized_api_key} \"\n f\"Exception: {error}\"\n )\n\n\n# Custom exception for authentication errors\nclass AuthenticationError(Exception):\n \"\"\"Raised when authentication fails with a provider.\"\"\"\n\n def __init__(self, provider: str, message: str = \"API key is invalid or expired\"):\n self.provider = provider\n self.message = message\n super().__init__(f\"{provider} authentication failed: {message}\")\n\n\nclass CloudEmbedding:\n def __init__(\n self,\n api_key: str,\n provider: EmbeddingProvider,\n api_url: str | None = None,\n api_version: str | None = None,\n timeout: int = API_BASED_EMBEDDING_TIMEOUT,\n ) -> None:\n self.provider = provider\n self.api_key = api_key\n self.api_url = api_url\n self.api_version = api_version\n self.timeout = timeout\n self.http_client = httpx.AsyncClient(timeout=timeout)\n self._closed = False\n self.sanitized_api_key = api_key[:4] + \"********\" + api_key[-4:]\n\n async def _embed_openai(\n self, texts: list[str], model: str | None, reduced_dimension: int | None\n ) -> list[Embedding]:\n if not model:\n model = DEFAULT_OPENAI_MODEL\n\n import openai\n\n # Use the OpenAI specific timeout for this one\n client = openai.AsyncOpenAI(\n api_key=self.api_key, timeout=OPENAI_EMBEDDING_TIMEOUT\n )\n\n final_embeddings: list[Embedding] = []\n\n for text_batch in batch_list(texts, _OPENAI_MAX_INPUT_LEN):\n response = await client.embeddings.create(\n input=text_batch,\n model=model,\n dimensions=reduced_dimension or openai.omit,\n )\n final_embeddings.extend(\n [embedding.embedding for embedding in response.data]\n )\n return final_embeddings\n\n async def _embed_cohere(\n self, texts: list[str], model: str | None, embedding_type: str\n ) -> list[Embedding]:\n if not model:\n model = DEFAULT_COHERE_MODEL\n\n client = CohereAsyncClient(api_key=self.api_key)\n\n final_embeddings: list[Embedding] = []\n for text_batch in batch_list(texts, _COHERE_MAX_INPUT_LEN):\n # Does not use the same tokenizer as the Onyx API server but it's approximately the same\n # empirically it's only off by a very few tokens so it's not a big deal\n response = await client.embed(\n texts=text_batch,\n model=model,\n input_type=embedding_type,\n truncate=\"END\",\n )\n final_embeddings.extend(cast(list[Embedding], response.embeddings))\n return final_embeddings\n\n async def _embed_voyage(\n self, texts: list[str], model: str | None, embedding_type: str\n ) -> list[Embedding]:\n if not model:\n model = DEFAULT_VOYAGE_MODEL\n\n client = voyageai.AsyncClient(\n api_key=self.api_key, timeout=API_BASED_EMBEDDING_TIMEOUT\n )\n\n response = await client.embed(\n texts=texts,\n model=model,\n input_type=embedding_type,\n truncation=True,\n )\n return response.embeddings\n\n async def _embed_azure(\n self, texts: list[str], model: str | None\n ) -> list[Embedding]:\n from litellm import aembedding\n\n response = await aembedding(\n model=model,\n input=texts,\n timeout=API_BASED_EMBEDDING_TIMEOUT,\n api_key=self.api_key,\n api_base=self.api_url,\n api_version=self.api_version,\n )\n embeddings = [embedding[\"embedding\"] for embedding in response.data]\n return embeddings\n\n async def _embed_vertex(\n self,\n texts: list[str],\n model: str | None,\n embedding_type: str,\n reduced_dimension: int | None,\n ) -> list[Embedding]:\n from google import genai\n from google.genai import types as genai_types\n\n if not model:\n model = DEFAULT_VERTEX_MODEL\n\n service_account_info = json.loads(self.api_key)\n credentials = service_account.Credentials.from_service_account_info(\n service_account_info,\n scopes=[\"https://www.googleapis.com/auth/cloud-platform\"],\n )\n project_id = service_account_info[\"project_id\"]\n location = (\n service_account_info.get(\"location\")\n or os.environ.get(\"GOOGLE_CLOUD_LOCATION\")\n or \"us-central1\"\n )\n\n client = genai.Client(\n vertexai=True,\n project=project_id,\n location=location,\n credentials=credentials,\n )\n\n embed_config = genai_types.EmbedContentConfig(\n task_type=embedding_type,\n output_dimensionality=reduced_dimension,\n auto_truncate=True,\n )\n\n async def _embed_batch(batch_texts: list[str]) -> list[Embedding]:\n content_requests: list[Any] = [\n genai_types.Content(parts=[genai_types.Part(text=text)])\n for text in batch_texts\n ]\n response = await client.aio.models.embed_content(\n model=model,\n contents=content_requests,\n config=embed_config,\n )\n\n if not response.embeddings:\n raise RuntimeError(\"Received empty embeddings from Google GenAI.\")\n\n embeddings: list[Embedding] = []\n for idx, embedding in enumerate(response.embeddings):\n if embedding.values is None:\n raise RuntimeError(\n f\"Missing embedding values for input at index {idx}.\"\n )\n embeddings.append(embedding.values)\n return embeddings\n\n # Process VertexAI batches sequentially to avoid additional intra-task fanout.\n # The higher-level thread pool already provides concurrency; running these\n # requests in parallel here was causing excessive memory usage.\n batches = [\n texts[i : i + VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE]\n for i in range(0, len(texts), VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE)\n ]\n all_embeddings: list[Embedding] = []\n\n logger.debug(\n f\"VertexAI embedding: processing {len(texts)} texts in {len(batches)} batches \"\n f\"(batch_size={VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE})\"\n )\n\n try:\n for batch_idx, batch in enumerate(batches):\n batch_embeddings = await _embed_batch(batch)\n all_embeddings.extend(batch_embeddings)\n\n # Log progress for large batches to track memory usage patterns\n if batch_idx % 10 == 0 and batch_idx > 0:\n logger.debug(\n f\"VertexAI embedding progress: batch {batch_idx}/{len(batches)}, total_embeddings={len(all_embeddings)}\"\n )\n\n logger.debug(\n f\"VertexAI embedding completed: {len(all_embeddings)} embeddings generated\"\n )\n return all_embeddings\n finally:\n # Ensure client is closed with a timeout to prevent hanging on stuck sessions\n try:\n await asyncio.wait_for(client.aio.aclose(), timeout=5.0)\n except asyncio.TimeoutError:\n logger.warning(\"Google GenAI client aclose() timed out after 5s\")\n except Exception as e:\n logger.warning(f\"Error closing Google GenAI client: {e}\")\n\n async def _embed_litellm_proxy(\n self, texts: list[str], model_name: str | None\n ) -> list[Embedding]:\n if not model_name:\n raise ValueError(\"Model name is required for LiteLLM proxy embedding.\")\n\n if not self.api_url:\n raise ValueError(\"API URL is required for LiteLLM proxy embedding.\")\n\n headers = (\n {} if not self.api_key else {\"Authorization\": f\"Bearer {self.api_key}\"}\n )\n\n response = await self.http_client.post(\n self.api_url,\n json={\n \"model\": model_name,\n \"input\": texts,\n },\n headers=headers,\n )\n response.raise_for_status()\n result = response.json()\n return [embedding[\"embedding\"] for embedding in result[\"data\"]]\n\n @retry(tries=_RETRY_TRIES, delay=_RETRY_DELAY)\n async def embed(\n self,\n *,\n texts: list[str],\n text_type: EmbedTextType,\n model_name: str | None = None,\n deployment_name: str | None = None,\n reduced_dimension: int | None = None,\n ) -> list[Embedding]:\n import openai\n\n try:\n if self.provider == EmbeddingProvider.OPENAI:\n return await self._embed_openai(texts, model_name, reduced_dimension)\n elif self.provider == EmbeddingProvider.AZURE:\n return await self._embed_azure(texts, f\"azure/{deployment_name}\")\n elif self.provider == EmbeddingProvider.LITELLM:\n return await self._embed_litellm_proxy(texts, model_name)\n\n embedding_type = EmbeddingModelTextType.get_type(self.provider, text_type)\n if self.provider == EmbeddingProvider.COHERE:\n return await self._embed_cohere(texts, model_name, embedding_type)\n elif self.provider == EmbeddingProvider.VOYAGE:\n return await self._embed_voyage(texts, model_name, embedding_type)\n elif self.provider == EmbeddingProvider.GOOGLE:\n return await self._embed_vertex(\n texts, model_name, embedding_type, reduced_dimension\n )\n else:\n raise ValueError(f\"Unsupported provider: {self.provider}\")\n except openai.AuthenticationError:\n raise AuthenticationError(provider=\"OpenAI\")\n except httpx.HTTPStatusError as e:\n if e.response.status_code == 401:\n raise AuthenticationError(provider=str(self.provider))\n\n error_string = format_embedding_error(\n e,\n str(self.provider),\n model_name or deployment_name,\n self.provider,\n sanitized_api_key=self.sanitized_api_key,\n status_code=e.response.status_code,\n )\n logger.error(error_string)\n logger.debug(f\"Exception texts: {texts}\")\n\n raise RuntimeError(error_string)\n except Exception as e:\n if is_authentication_error(e):\n raise AuthenticationError(provider=str(self.provider))\n\n error_string = format_embedding_error(\n e,\n str(self.provider),\n model_name or deployment_name,\n self.provider,\n sanitized_api_key=self.sanitized_api_key,\n )\n logger.error(error_string)\n logger.debug(f\"Exception texts: {texts}\")\n\n raise RuntimeError(error_string)\n\n @staticmethod\n def create(\n api_key: str,\n provider: EmbeddingProvider,\n api_url: str | None = None,\n api_version: str | None = None,\n ) -> \"CloudEmbedding\":\n logger.debug(f\"Creating Embedding instance for provider: {provider}\")\n return CloudEmbedding(api_key, provider, api_url, api_version)\n\n async def aclose(self) -> None:\n \"\"\"Explicitly close the client.\"\"\"\n if not self._closed:\n await self.http_client.aclose()\n self._closed = True\n\n async def __aenter__(self) -> \"CloudEmbedding\":\n return self\n\n async def __aexit__(\n self,\n exc_type: type[BaseException] | None,\n exc_val: BaseException | None,\n exc_tb: TracebackType | None,\n ) -> None:\n await self.aclose()\n\n def __del__(self) -> None:\n \"\"\"Finalizer to warn about unclosed clients.\"\"\"\n if not self._closed:\n logger.warning(\n \"CloudEmbedding was not properly closed. Use 'async with' or call aclose()\"\n )\n\n\n# API-based reranking functions (moved from model server)\nasync def cohere_rerank_api(\n query: str, docs: list[str], model_name: str, api_key: str\n) -> list[float]:\n cohere_client = CohereAsyncClient(api_key=api_key)\n try:\n response = await cohere_client.rerank(\n query=query, documents=docs, model=model_name\n )\n except ApiError as err:\n if err.status_code == 402:\n logger.warning(\n \"Cohere rerank request rejected due to billing cap. Falling back to retrieval ordering until billing resets.\"\n )\n raise CohereBillingLimitError(\n \"Cohere billing limit reached for reranking\"\n ) from err\n raise\n results = response.results\n sorted_results = sorted(results, key=lambda item: item.index)\n return [result.relevance_score for result in sorted_results]\n\n\nasync def cohere_rerank_aws(\n query: str,\n docs: list[str],\n model_name: str,\n region_name: str,\n aws_access_key_id: str,\n aws_secret_access_key: str,\n) -> list[float]:\n session = aioboto3.Session(\n aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key\n )\n async with session.client(\n \"bedrock-runtime\", region_name=region_name\n ) as bedrock_client:\n body = json.dumps(\n {\n \"query\": query,\n \"documents\": docs,\n \"api_version\": 2,\n }\n )\n # Invoke the Bedrock model asynchronously\n response = await bedrock_client.invoke_model(\n modelId=model_name,\n accept=\"application/json\",\n contentType=\"application/json\",\n body=body,\n )\n\n # Read the response asynchronously\n response_body = json.loads(await response[\"body\"].read())\n\n # Extract and sort the results\n results = response_body.get(\"results\", [])\n sorted_results = sorted(results, key=lambda item: item[\"index\"])\n\n return [result[\"relevance_score\"] for result in sorted_results]\n\n\nasync def litellm_rerank(\n query: str, docs: list[str], api_url: str, model_name: str, api_key: str | None\n) -> list[float]:\n headers = {} if not api_key else {\"Authorization\": f\"Bearer {api_key}\"}\n async with httpx.AsyncClient() as client:\n response = await client.post(\n api_url,\n json={\n \"model\": model_name,\n \"query\": query,\n \"documents\": docs,\n },\n headers=headers,\n )\n response.raise_for_status()\n result = response.json()\n return [\n item[\"relevance_score\"]\n for item in sorted(result[\"results\"], key=lambda x: x[\"index\"])\n ]\n\n\nclass EmbeddingModel:\n def __init__(\n self,\n server_host: str, # Changes depending on indexing or inference\n server_port: int,\n model_name: str | None,\n normalize: bool,\n query_prefix: str | None,\n passage_prefix: str | None,\n api_key: str | None,\n api_url: str | None,\n provider_type: EmbeddingProvider | None,\n retrim_content: bool = False,\n callback: IndexingHeartbeatInterface | None = None,\n api_version: str | None = None,\n deployment_name: str | None = None,\n reduced_dimension: int | None = None,\n ) -> None:\n self.api_key = api_key\n self.provider_type = provider_type\n self.query_prefix = query_prefix\n self.passage_prefix = passage_prefix\n self.normalize = normalize\n self.model_name = model_name\n self.retrim_content = retrim_content\n self.api_url = api_url\n self.api_version = api_version\n self.deployment_name = deployment_name\n self.reduced_dimension = reduced_dimension\n self.tokenizer = get_tokenizer(\n model_name=model_name, provider_type=provider_type\n )\n self.callback = callback\n\n # Only build model server endpoint for local models\n if self.provider_type is None:\n model_server_url = build_model_server_url(server_host, server_port)\n self.embed_server_endpoint: str | None = (\n f\"{model_server_url}/encoder/bi-encoder-embed\"\n )\n else:\n # API providers don't need model server endpoint\n self.embed_server_endpoint = None\n\n async def _make_direct_api_call(\n self,\n embed_request: EmbedRequest,\n tenant_id: str | None = None, # noqa: ARG002\n request_id: str | None = None, # noqa: ARG002\n ) -> EmbedResponse:\n \"\"\"Make direct API call to cloud provider, bypassing model server.\"\"\"\n if self.provider_type is None:\n raise ValueError(\"Provider type is required for direct API calls\")\n\n if self.api_key is None:\n logger.error(\"API key not provided for cloud model\")\n raise RuntimeError(\"API key not provided for cloud model\")\n\n # Check for prefix usage with cloud models\n if embed_request.manual_query_prefix or embed_request.manual_passage_prefix:\n logger.warning(\"Prefix provided for cloud model, which is not supported\")\n raise ValueError(\n \"Prefix string is not valid for cloud models. Cloud models take an explicit text type instead.\"\n )\n\n if not all(embed_request.texts):\n logger.error(\"Empty strings provided for embedding\")\n raise ValueError(\"Empty strings are not allowed for embedding.\")\n\n if not embed_request.texts:\n logger.error(\"No texts provided for embedding\")\n raise ValueError(\"No texts provided for embedding.\")\n\n start_time = time.monotonic()\n total_chars = sum(len(text) for text in embed_request.texts)\n\n logger.info(\n f\"Embedding {len(embed_request.texts)} texts with {total_chars} total characters with provider: {self.provider_type}\"\n )\n\n async with CloudEmbedding(\n api_key=self.api_key,\n provider=self.provider_type,\n api_url=self.api_url,\n api_version=self.api_version,\n ) as cloud_model:\n embeddings = await cloud_model.embed(\n texts=embed_request.texts,\n model_name=embed_request.model_name,\n deployment_name=embed_request.deployment_name,\n text_type=embed_request.text_type,\n reduced_dimension=embed_request.reduced_dimension,\n )\n\n if any(embedding is None for embedding in embeddings):\n error_message = \"Embeddings contain None values\\n\"\n error_message += \"Corresponding texts:\\n\"\n error_message += \"\\n\".join(embed_request.texts)\n logger.error(error_message)\n raise ValueError(error_message)\n\n elapsed = time.monotonic() - start_time\n logger.info(\n f\"event=embedding_provider \"\n f\"texts={len(embed_request.texts)} \"\n f\"chars={total_chars} \"\n f\"provider={self.provider_type} \"\n f\"elapsed={elapsed:.2f}\"\n )\n\n return EmbedResponse(embeddings=embeddings)\n\n def _make_model_server_request(\n self,\n embed_request: EmbedRequest,\n tenant_id: str | None = None,\n request_id: str | None = None,\n ) -> EmbedResponse:\n if self.embed_server_endpoint is None:\n raise ValueError(\"Model server endpoint is not configured for local models\")\n\n # Store the endpoint in a local variable to help mypy understand it's not None\n endpoint = self.embed_server_endpoint\n\n def _make_request() -> Response:\n headers = {}\n if tenant_id:\n headers[\"X-Onyx-Tenant-ID\"] = tenant_id\n\n if request_id:\n headers[\"X-Onyx-Request-ID\"] = request_id\n\n response = requests.post(\n endpoint,\n headers=headers,\n json=embed_request.model_dump(),\n )\n # signify that this is a rate limit error\n if response.status_code == 429:\n raise ModelServerRateLimitError(response.text)\n\n response.raise_for_status()\n return response\n\n final_make_request_func = _make_request\n\n # if the text type is a passage, add some default\n # retries + handling for rate limiting\n if embed_request.text_type == EmbedTextType.PASSAGE:\n final_make_request_func = retry(\n tries=3,\n delay=5,\n exceptions=(RequestException, ValueError, JSONDecodeError),\n )(final_make_request_func)\n # use 10 second delay as per Azure suggestion\n final_make_request_func = retry(\n tries=10, delay=10, exceptions=ModelServerRateLimitError\n )(final_make_request_func)\n\n response: Response | None = None\n\n try:\n response = final_make_request_func()\n return EmbedResponse(**response.json())\n except requests.HTTPError as e:\n if not response:\n raise HTTPError(\"HTTP error occurred - response is None.\") from e\n\n try:\n error_detail = response.json().get(\"detail\", str(e))\n except Exception:\n error_detail = response.text\n raise HTTPError(f\"HTTP error occurred: {error_detail}\") from e\n except requests.RequestException as e:\n raise HTTPError(f\"Request failed: {str(e)}\") from e\n\n def _batch_encode_texts(\n self,\n texts: list[str],\n text_type: EmbedTextType,\n batch_size: int,\n max_seq_length: int,\n num_threads: int = INDEXING_EMBEDDING_MODEL_NUM_THREADS,\n tenant_id: str | None = None,\n request_id: str | None = None,\n ) -> list[Embedding]:\n text_batches = batch_list(texts, batch_size)\n\n logger.debug(f\"Encoding {len(texts)} texts in {len(text_batches)} batches\")\n\n embeddings: list[Embedding] = []\n\n @_cleanup_thread_local\n def process_batch(\n batch_idx: int,\n batch_len: int,\n text_batch: list[str],\n tenant_id: str | None = None,\n request_id: str | None = None,\n ) -> tuple[int, list[Embedding]]:\n if self.callback:\n if self.callback.should_stop():\n raise ConnectorStopSignal(\n \"_batch_encode_texts detected stop signal\"\n )\n\n embed_request = EmbedRequest(\n model_name=self.model_name,\n texts=text_batch,\n api_version=self.api_version,\n deployment_name=self.deployment_name,\n max_context_length=max_seq_length,\n normalize_embeddings=self.normalize,\n api_key=self.api_key,\n provider_type=self.provider_type,\n text_type=text_type,\n manual_query_prefix=self.query_prefix,\n manual_passage_prefix=self.passage_prefix,\n api_url=self.api_url,\n reduced_dimension=self.reduced_dimension,\n )\n\n start_time = time.monotonic()\n\n # Route between direct API calls and model server calls\n if self.provider_type is not None:\n # For API providers, make direct API call\n # Use thread-local event loop to prevent memory leaks from creating\n # thousands of event loops during batch processing\n loop = _get_or_create_event_loop()\n response = loop.run_until_complete(\n self._make_direct_api_call(\n embed_request, tenant_id=tenant_id, request_id=request_id\n )\n )\n else:\n # For local models, use model server\n response = self._make_model_server_request(\n embed_request, tenant_id=tenant_id, request_id=request_id\n )\n\n end_time = time.monotonic()\n\n processing_time = end_time - start_time\n logger.debug(\n f\"EmbeddingModel.process_batch: Batch {batch_idx}/{batch_len} processing time: {processing_time:.2f} seconds\"\n )\n\n return batch_idx, response.embeddings\n\n # only multi thread if:\n # 1. num_threads is greater than 1\n # 2. we are using an API-based embedding model (provider_type is not None)\n # 3. there are more than 1 batch (no point in threading if only 1)\n if num_threads >= 1 and self.provider_type and len(text_batches) > 1:\n with ThreadPoolExecutor(max_workers=num_threads) as executor:\n future_to_batch = {\n executor.submit(\n partial(\n process_batch,\n idx,\n len(text_batches),\n batch,\n tenant_id=tenant_id,\n request_id=request_id,\n )\n ): idx\n for idx, batch in enumerate(text_batches, start=1)\n }\n\n # Collect results in order\n batch_results: list[tuple[int, list[Embedding]]] = []\n for future in as_completed(future_to_batch):\n try:\n result = future.result()\n batch_results.append(result)\n except Exception as e:\n logger.exception(\"Embedding model failed to process batch\")\n raise e\n\n # Sort by batch index and extend embeddings\n batch_results.sort(key=lambda x: x[0])\n for _, batch_embeddings in batch_results:\n embeddings.extend(batch_embeddings)\n else:\n # Original sequential processing\n for idx, text_batch in enumerate(text_batches, start=1):\n _, batch_embeddings = process_batch(\n idx,\n len(text_batches),\n text_batch,\n tenant_id=tenant_id,\n request_id=request_id,\n )\n embeddings.extend(batch_embeddings)\n\n return embeddings\n\n @log_function_time(print_only=True, debug_only=True)\n def encode(\n self,\n texts: list[str],\n text_type: EmbedTextType,\n large_chunks_present: bool = False,\n local_embedding_batch_size: int = BATCH_SIZE_ENCODE_CHUNKS,\n api_embedding_batch_size: int = BATCH_SIZE_ENCODE_CHUNKS_FOR_API_EMBEDDING_SERVICES,\n max_seq_length: int = DOC_EMBEDDING_CONTEXT_SIZE,\n tenant_id: str | None = None,\n request_id: str | None = None,\n ) -> list[Embedding]:\n if not texts or not all(texts):\n raise ValueError(f\"Empty or missing text for embedding: {texts}\")\n\n if large_chunks_present:\n max_seq_length *= LARGE_CHUNK_RATIO\n\n if self.retrim_content:\n # This is applied during indexing as a catchall for overly long titles (or other uncapped fields)\n # Note that this uses just the default tokenizer which may also lead to very minor miscountings\n # However this slight miscounting is very unlikely to have any material impact.\n texts = [\n tokenizer_trim_content(\n content=text,\n desired_length=max_seq_length,\n tokenizer=self.tokenizer,\n )\n for text in texts\n ]\n\n # Remove invalid Unicode characters (e.g., unpaired surrogates from malformed documents)\n # that would cause UTF-8 encoding errors when sent to embedding providers\n texts = [remove_invalid_unicode_chars(text) or \"<>\" for text in texts]\n\n batch_size = (\n api_embedding_batch_size\n if self.provider_type\n else local_embedding_batch_size\n )\n\n return self._batch_encode_texts(\n texts=texts,\n text_type=text_type,\n batch_size=batch_size,\n max_seq_length=max_seq_length,\n tenant_id=tenant_id,\n request_id=request_id,\n )\n\n @classmethod\n def from_db_model(\n cls,\n search_settings: SearchSettings,\n server_host: str, # Changes depending on indexing or inference\n server_port: int,\n retrim_content: bool = False,\n ) -> \"EmbeddingModel\":\n return cls(\n server_host=server_host,\n server_port=server_port,\n model_name=search_settings.model_name,\n normalize=search_settings.normalize,\n query_prefix=search_settings.query_prefix,\n passage_prefix=search_settings.passage_prefix,\n api_key=search_settings.api_key,\n provider_type=search_settings.provider_type,\n api_url=search_settings.api_url,\n retrim_content=retrim_content,\n api_version=search_settings.api_version,\n deployment_name=search_settings.deployment_name,\n reduced_dimension=search_settings.reduced_dimension,\n )\n\n\nclass RerankingModel:\n def __init__(\n self,\n model_name: str,\n provider_type: RerankerProvider | None,\n api_key: str | None,\n api_url: str | None,\n model_server_host: str = MODEL_SERVER_HOST,\n model_server_port: int = MODEL_SERVER_PORT,\n ) -> None:\n self.model_name = model_name\n self.provider_type = provider_type\n self.api_key = api_key\n self.api_url = api_url\n\n # Only build model server endpoint for local models\n if self.provider_type is None:\n model_server_url = build_model_server_url(\n model_server_host, model_server_port\n )\n self.rerank_server_endpoint: str | None = (\n model_server_url + \"/encoder/cross-encoder-scores\"\n )\n else:\n # API providers don't need model server endpoint\n self.rerank_server_endpoint = None\n\n async def _make_direct_rerank_call(\n self, query: str, passages: list[str]\n ) -> list[float]:\n \"\"\"Make direct API call to cloud provider, bypassing model server.\"\"\"\n if self.provider_type is None:\n raise ValueError(\"Provider type is required for direct API calls\")\n\n if self.api_key is None:\n raise ValueError(\"API key is required for cloud provider\")\n\n if self.provider_type == RerankerProvider.COHERE:\n return await cohere_rerank_api(\n query, passages, self.model_name, self.api_key\n )\n elif self.provider_type == RerankerProvider.BEDROCK:\n aws_access_key_id, aws_secret_access_key, aws_region = pass_aws_key(\n self.api_key\n )\n return await cohere_rerank_aws(\n query,\n passages,\n self.model_name,\n aws_region,\n aws_access_key_id,\n aws_secret_access_key,\n )\n elif self.provider_type == RerankerProvider.LITELLM:\n if self.api_url is None:\n raise ValueError(\"API URL is required for LiteLLM reranking.\")\n return await litellm_rerank(\n query, passages, self.api_url, self.model_name, self.api_key\n )\n else:\n raise ValueError(f\"Unsupported reranking provider: {self.provider_type}\")\n\n def predict(self, query: str, passages: list[str]) -> list[float]:\n # Route between direct API calls and model server calls\n if self.provider_type is not None:\n # For API providers, make direct API call\n loop = asyncio.new_event_loop()\n try:\n asyncio.set_event_loop(loop)\n return loop.run_until_complete(\n self._make_direct_rerank_call(query, passages)\n )\n finally:\n loop.close()\n else:\n # For local models, use model server\n if self.rerank_server_endpoint is None:\n raise ValueError(\n \"Rerank server endpoint is not configured for local models\"\n )\n\n rerank_request = RerankRequest(\n query=query,\n documents=passages,\n model_name=self.model_name,\n provider_type=self.provider_type,\n api_key=self.api_key,\n api_url=self.api_url,\n )\n\n response = requests.post(\n self.rerank_server_endpoint, json=rerank_request.model_dump()\n )\n response.raise_for_status()\n\n return RerankResponse(**response.json()).scores\n\n\nclass QueryAnalysisModel:\n def __init__(\n self,\n model_server_host: str = MODEL_SERVER_HOST,\n model_server_port: int = MODEL_SERVER_PORT,\n # Lean heavily towards not throwing out keywords\n keyword_percent_threshold: float = 0.1,\n # Lean towards semantic which is the default\n semantic_percent_threshold: float = 0.4,\n ) -> None:\n model_server_url = build_model_server_url(model_server_host, model_server_port)\n self.intent_server_endpoint = model_server_url + \"/custom/query-analysis\"\n self.keyword_percent_threshold = keyword_percent_threshold\n self.semantic_percent_threshold = semantic_percent_threshold\n\n def predict(\n self,\n query: str,\n ) -> tuple[bool, list[str]]:\n intent_request = IntentRequest(\n query=query,\n keyword_percent_threshold=self.keyword_percent_threshold,\n semantic_percent_threshold=self.semantic_percent_threshold,\n )\n\n response = requests.post(\n self.intent_server_endpoint, json=intent_request.model_dump()\n )\n response.raise_for_status()\n\n response_model = IntentResponse(**response.json())\n\n return response_model.is_keyword, response_model.keywords\n\n\ndef warm_up_retry(\n func: Callable[..., Any],\n tries: int = 20,\n delay: int = 5,\n *args: Any, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n) -> Callable[..., Any]:\n @wraps(func)\n def wrapper(*args: Any, **kwargs: Any) -> Any:\n exceptions = []\n for attempt in range(tries):\n try:\n return func(*args, **kwargs)\n except Exception as e:\n exceptions.append(e)\n logger.info(\n f\"Attempt {attempt + 1}/{tries} failed; retrying in {delay} seconds...\"\n )\n time.sleep(delay)\n raise Exception(f\"All retries failed: {exceptions}\")\n\n return wrapper\n\n\ndef warm_up_bi_encoder(\n embedding_model: EmbeddingModel,\n non_blocking: bool = False,\n) -> None:\n if SKIP_WARM_UP:\n return\n\n warm_up_str = \" \".join(WARM_UP_STRINGS)\n\n logger.debug(f\"Warming up encoder model: {embedding_model.model_name}\")\n get_tokenizer(\n model_name=embedding_model.model_name,\n provider_type=embedding_model.provider_type,\n ).encode(warm_up_str)\n\n def _warm_up() -> None:\n try:\n embedding_model.encode(texts=[warm_up_str], text_type=EmbedTextType.QUERY)\n logger.debug(\n f\"Warm-up complete for encoder model: {embedding_model.model_name}\"\n )\n except Exception as e:\n logger.warning(\n f\"Warm-up request failed for encoder model {embedding_model.model_name}: {e}\"\n )\n\n if non_blocking:\n threading.Thread(target=_warm_up, daemon=True).start()\n logger.debug(\n f\"Started non-blocking warm-up for encoder model: {embedding_model.model_name}\"\n )\n else:\n retry_encode = warm_up_retry(embedding_model.encode)\n retry_encode(texts=[warm_up_str], text_type=EmbedTextType.QUERY)\n\n\n# No longer used\ndef warm_up_cross_encoder(\n rerank_model_name: str,\n non_blocking: bool = False,\n) -> None:\n if SKIP_WARM_UP:\n return\n\n logger.debug(f\"Warming up reranking model: {rerank_model_name}\")\n\n reranking_model = RerankingModel(\n model_name=rerank_model_name,\n provider_type=None,\n api_url=None,\n api_key=None,\n )\n\n def _warm_up() -> None:\n try:\n reranking_model.predict(WARM_UP_STRINGS[0], WARM_UP_STRINGS[1:])\n logger.debug(f\"Warm-up complete for reranking model: {rerank_model_name}\")\n except Exception as e:\n logger.warning(\n f\"Warm-up request failed for reranking model {rerank_model_name}: {e}\"\n )\n\n if non_blocking:\n threading.Thread(target=_warm_up, daemon=True).start()\n logger.debug(\n f\"Started non-blocking warm-up for reranking model: {rerank_model_name}\"\n )\n else:\n retry_rerank = warm_up_retry(reranking_model.predict)\n retry_rerank(WARM_UP_STRINGS[0], WARM_UP_STRINGS[1:])\n" }, { "path": "backend/onyx/natural_language_processing/utils.py", "content": "import os\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom copy import copy\n\nfrom tokenizers import Encoding # type: ignore[import-untyped]\nfrom tokenizers import Tokenizer\n\nfrom onyx.configs.model_configs import DOCUMENT_ENCODER_MODEL\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import DOC_EMBEDDING_CONTEXT_SIZE\nfrom shared_configs.enums import EmbeddingProvider\n\nTRIM_SEP_PAT = \"\\n... {n} tokens removed...\\n\"\n\nlogger = setup_logger()\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\nos.environ[\"HF_HUB_DISABLE_TELEMETRY\"] = \"1\"\n\n\nclass BaseTokenizer(ABC):\n @abstractmethod\n def encode(self, string: str) -> list[int]:\n pass\n\n @abstractmethod\n def tokenize(self, string: str) -> list[str]:\n pass\n\n @abstractmethod\n def decode(self, tokens: list[int]) -> str:\n pass\n\n\nclass TiktokenTokenizer(BaseTokenizer):\n _instances: dict[str, \"TiktokenTokenizer\"] = {}\n\n def __new__(cls, model_name: str) -> \"TiktokenTokenizer\":\n if model_name not in cls._instances:\n cls._instances[model_name] = super(TiktokenTokenizer, cls).__new__(cls)\n return cls._instances[model_name]\n\n def __init__(self, model_name: str):\n if not hasattr(self, \"encoder\"):\n import tiktoken\n\n self.encoder = tiktoken.encoding_for_model(model_name)\n\n def encode(self, string: str) -> list[int]:\n # this ignores special tokens that the model is trained on, see encode_ordinary for details\n return self.encoder.encode_ordinary(string)\n\n def tokenize(self, string: str) -> list[str]:\n encoded = self.encode(string)\n decoded = [self.encoder.decode([token]) for token in encoded]\n\n if len(decoded) != len(encoded):\n logger.warning(\n f\"OpenAI tokenized length {len(decoded)} does not match encoded length {len(encoded)} for string: {string}\"\n )\n\n return decoded\n\n def decode(self, tokens: list[int]) -> str:\n return self.encoder.decode(tokens)\n\n\nclass HuggingFaceTokenizer(BaseTokenizer):\n def __init__(self, model_name: str):\n self.encoder: Tokenizer = Tokenizer.from_pretrained(model_name)\n\n def _safer_encode(self, string: str) -> Encoding:\n \"\"\"\n Encode a string using the HuggingFaceTokenizer, but if it fails,\n encode the string as ASCII and decode it back to a string. This helps\n in cases where the string has weird characters like \\udeb4.\n \"\"\"\n try:\n return self.encoder.encode(string, add_special_tokens=False)\n except Exception:\n return self.encoder.encode(\n string.encode(\"ascii\", \"ignore\").decode(), add_special_tokens=False\n )\n\n def encode(self, string: str) -> list[int]:\n # this returns no special tokens\n return self._safer_encode(string).ids\n\n def tokenize(self, string: str) -> list[str]:\n return self._safer_encode(string).tokens\n\n def decode(self, tokens: list[int]) -> str:\n return self.encoder.decode(tokens)\n\n\n_TOKENIZER_CACHE: dict[tuple[EmbeddingProvider | None, str | None], BaseTokenizer] = {}\n\n\ndef _check_tokenizer_cache(\n model_provider: EmbeddingProvider | None, model_name: str | None\n) -> BaseTokenizer:\n global _TOKENIZER_CACHE\n id_tuple = (model_provider, model_name)\n\n if id_tuple not in _TOKENIZER_CACHE:\n tokenizer = None\n\n if model_name:\n tokenizer = _try_initialize_tokenizer(model_name, model_provider)\n\n if not tokenizer:\n logger.info(\n f\"Falling back to default embedding model tokenizer: {DOCUMENT_ENCODER_MODEL}\"\n )\n tokenizer = _get_default_tokenizer()\n\n _TOKENIZER_CACHE[id_tuple] = tokenizer\n\n return _TOKENIZER_CACHE[id_tuple]\n\n\ndef _try_initialize_tokenizer(\n model_name: str, model_provider: EmbeddingProvider | None\n) -> BaseTokenizer | None:\n tokenizer: BaseTokenizer | None = None\n\n if model_provider is not None:\n # Try using TiktokenTokenizer first if model_provider exists\n try:\n tokenizer = TiktokenTokenizer(model_name)\n logger.info(f\"Initialized TiktokenTokenizer for: {model_name}\")\n return tokenizer\n except Exception as tiktoken_error:\n logger.debug(\n f\"TiktokenTokenizer not available for model {model_name}: {tiktoken_error}\"\n )\n else:\n # If no provider specified, try HuggingFaceTokenizer\n try:\n tokenizer = HuggingFaceTokenizer(model_name)\n logger.info(f\"Initialized HuggingFaceTokenizer for: {model_name}\")\n return tokenizer\n except Exception as hf_error:\n logger.warning(\n f\"Failed to initialize HuggingFaceTokenizer for {model_name}: {hf_error}\"\n )\n\n # If both initializations fail, return None\n return None\n\n\n_DEFAULT_TOKENIZER: BaseTokenizer | None = None\n\n\ndef _get_default_tokenizer() -> BaseTokenizer:\n \"\"\"Lazy-load the default tokenizer to avoid loading it at module import time.\"\"\"\n global _DEFAULT_TOKENIZER\n if _DEFAULT_TOKENIZER is None:\n _DEFAULT_TOKENIZER = HuggingFaceTokenizer(DOCUMENT_ENCODER_MODEL)\n return _DEFAULT_TOKENIZER\n\n\ndef get_tokenizer(\n model_name: str | None, provider_type: EmbeddingProvider | str | None\n) -> BaseTokenizer:\n if isinstance(provider_type, str):\n try:\n provider_type = EmbeddingProvider(provider_type)\n except ValueError:\n logger.debug(\n f\"Invalid provider_type '{provider_type}'. Falling back to default tokenizer.\"\n )\n return _get_default_tokenizer()\n return _check_tokenizer_cache(provider_type, model_name)\n\n\n# Max characters per encode() call.\n_ENCODE_CHUNK_SIZE = 500_000\n\n\ndef count_tokens(\n text: str,\n tokenizer: BaseTokenizer,\n token_limit: int | None = None,\n) -> int:\n \"\"\"Count tokens, chunking the input to avoid tiktoken stack overflow.\n\n If token_limit is provided and the text is large enough to require\n multiple chunks (> 500k chars), stops early once the count exceeds it.\n When early-exiting, the returned value exceeds token_limit but may be\n less than the true full token count.\n \"\"\"\n if len(text) <= _ENCODE_CHUNK_SIZE:\n return len(tokenizer.encode(text))\n total = 0\n for start in range(0, len(text), _ENCODE_CHUNK_SIZE):\n total += len(tokenizer.encode(text[start : start + _ENCODE_CHUNK_SIZE]))\n if token_limit is not None and total > token_limit:\n return total # Already over — skip remaining chunks\n return total\n\n\ndef tokenizer_trim_content(\n content: str, desired_length: int, tokenizer: BaseTokenizer\n) -> str:\n tokens = tokenizer.encode(content)\n if len(tokens) <= desired_length:\n return content\n\n return tokenizer.decode(tokens[:desired_length])\n\n\ndef tokenizer_trim_middle(\n tokens: list[int], desired_length: int, tokenizer: BaseTokenizer\n) -> str:\n if len(tokens) <= desired_length:\n return tokenizer.decode(tokens)\n sep_str = TRIM_SEP_PAT.format(n=len(tokens) - desired_length)\n sep_tokens = tokenizer.encode(sep_str)\n slice_size = (desired_length - len(sep_tokens)) // 2\n assert slice_size > 0, \"Slice size is not positive, desired length is too short\"\n return (\n tokenizer.decode(tokens[:slice_size])\n + sep_str\n + tokenizer.decode(tokens[-slice_size:])\n )\n\n\ndef tokenizer_trim_chunks(\n chunks: list[InferenceChunk],\n tokenizer: BaseTokenizer,\n max_chunk_toks: int = DOC_EMBEDDING_CONTEXT_SIZE,\n) -> list[InferenceChunk]:\n new_chunks = copy(chunks)\n for ind, chunk in enumerate(new_chunks):\n new_content = tokenizer_trim_content(chunk.content, max_chunk_toks, tokenizer)\n if len(new_content) != len(chunk.content):\n new_chunk = copy(chunk)\n new_chunk.content = new_content\n new_chunks[ind] = new_chunk\n return new_chunks\n" }, { "path": "backend/onyx/onyxbot/discord/DISCORD_MULTITENANT_README.md", "content": "# Discord Bot Multitenant Architecture\n\nThis document analyzes how the Discord cache manager and API client coordinate to handle multitenant API keys from a single Discord client.\n\n## Overview\n\nThe Discord bot uses a **single-client, multi-tenant** architecture where one `OnyxDiscordClient` instance serves multiple tenants (organizations) simultaneously. Tenant isolation is achieved through:\n\n- **Cache Manager**: Maps Discord guilds to tenants and stores per-tenant API keys\n- **API Client**: Stateless HTTP client that accepts dynamic API keys per request\n\n```\n┌─────────────────────────────────────────────────────────────────────┐\n│ OnyxDiscordClient │\n│ │\n│ ┌─────────────────────────┐ ┌─────────────────────────────┐ │\n│ │ DiscordCacheManager │ │ OnyxAPIClient │ │\n│ │ │ │ │ │\n│ │ guild_id → tenant_id │───▶│ send_chat_message( │ │\n│ │ tenant_id → api_key │ │ message, │ │\n│ │ │ │ api_key=, │ │\n│ └─────────────────────────┘ │ persona_id=... │ │\n│ │ ) │ │\n│ └─────────────────────────────┘ │\n└─────────────────────────────────────────────────────────────────────┘\n```\n\n---\n\n## Component Details\n\n### 1. Cache Manager (`backend/onyx/onyxbot/discord/cache.py`)\n\nThe `DiscordCacheManager` maintains two critical in-memory mappings:\n\n```python\nclass DiscordCacheManager:\n _guild_tenants: dict[int, str] # guild_id → tenant_id\n _api_keys: dict[str, str] # tenant_id → api_key\n _lock: asyncio.Lock # Concurrency control\n```\n\n#### Key Responsibilities\n\n| Function | Purpose |\n|----------|---------|\n| `get_tenant(guild_id)` | O(1) lookup: guild → tenant |\n| `get_api_key(tenant_id)` | O(1) lookup: tenant → API key |\n| `refresh_all()` | Full cache rebuild from database |\n| `refresh_guild()` | Incremental update for single guild |\n\n#### API Key Provisioning Strategy\n\nAPI keys are **lazily provisioned** - only created when first needed:\n\n```python\nasync def _load_tenant_data(self, tenant_id: str) -> tuple[list[int], str | None]:\n needs_key = tenant_id not in self._api_keys\n\n with get_session_with_tenant(tenant_id) as db:\n # Load guild configs\n configs = get_discord_bot_configs(db)\n guild_ids = [c.guild_id for c in configs if c.enabled]\n\n # Only provision API key if not already cached\n api_key = None\n if needs_key:\n api_key = get_or_create_discord_service_api_key(db, tenant_id)\n\n return guild_ids, api_key\n```\n\nThis optimization avoids repeated database calls for API key generation.\n\n#### Concurrency Control\n\nAll write operations acquire an async lock to prevent race conditions:\n\n```python\nasync def refresh_all(self) -> None:\n async with self._lock:\n # Safe to modify _guild_tenants and _api_keys\n for tenant_id in get_all_tenant_ids():\n guild_ids, api_key = await self._load_tenant_data(tenant_id)\n # Update mappings...\n```\n\nRead operations (`get_tenant`, `get_api_key`) are lock-free since Python dict lookups are atomic.\n\n---\n\n### 2. API Client (`backend/onyx/onyxbot/discord/api_client.py`)\n\nThe `OnyxAPIClient` is a **stateless async HTTP client** that communicates with Onyx API pods.\n\n#### Key Design: Per-Request API Key Injection\n\n```python\nclass OnyxAPIClient:\n async def send_chat_message(\n self,\n message: str,\n api_key: str, # Injected per-request\n persona_id: int | None,\n ...\n ) -> ChatFullResponse:\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {api_key}\", # Tenant-specific auth\n }\n # Make request...\n```\n\nThe client accepts `api_key` as a parameter to each method, enabling **dynamic tenant selection at request time**. This design allows a single client instance to serve multiple tenants:\n\n```python\n# Same client, different tenants\nawait api_client.send_chat_message(msg, api_key=key_for_tenant_1, ...)\nawait api_client.send_chat_message(msg, api_key=key_for_tenant_2, ...)\n```\n\n---\n\n## Coordination Flow\n\n### Message Processing Pipeline\n\nWhen a Discord message arrives, the client coordinates cache and API client:\n\n```python\nasync def on_message(self, message: Message) -> None:\n guild_id = message.guild.id\n\n # Step 1: Cache lookup - guild → tenant\n tenant_id = self.cache.get_tenant(guild_id)\n if not tenant_id:\n return # Guild not registered\n\n # Step 2: Cache lookup - tenant → API key\n api_key = self.cache.get_api_key(tenant_id)\n if not api_key:\n logger.warning(f\"No API key for tenant {tenant_id}\")\n return\n\n # Step 3: API call with tenant-specific credentials\n await process_chat_message(\n message=message,\n api_key=api_key, # Tenant-specific\n persona_id=persona_id, # Tenant-specific\n api_client=self.api_client,\n )\n```\n\n### Startup Sequence\n\n```python\nasync def setup_hook(self) -> None:\n # 1. Initialize API client (create aiohttp session)\n await self.api_client.initialize()\n\n # 2. Populate cache with all tenants\n await self.cache.refresh_all()\n\n # 3. Start background refresh task\n self._cache_refresh_task = self.loop.create_task(\n self._periodic_cache_refresh() # Every 60 seconds\n )\n```\n\n### Shutdown Sequence\n\n```python\nasync def close(self) -> None:\n # 1. Cancel background refresh\n if self._cache_refresh_task:\n self._cache_refresh_task.cancel()\n\n # 2. Close Discord connection\n await super().close()\n\n # 3. Close API client session\n await self.api_client.close()\n\n # 4. Clear cache\n self.cache.clear()\n```\n\n---\n\n## Tenant Isolation Mechanisms\n\n### 1. Per-Tenant API Keys\n\nEach tenant has a dedicated service API key:\n\n```python\n# backend/onyx/db/discord_bot.py\ndef get_or_create_discord_service_api_key(db_session: Session, tenant_id: str) -> str:\n existing = get_discord_service_api_key(db_session)\n if existing:\n return regenerate_key(existing)\n\n # Create LIMITED role key (chat-only permissions)\n return insert_api_key(\n db_session=db_session,\n api_key_args=APIKeyArgs(\n name=DISCORD_SERVICE_API_KEY_NAME,\n role=UserRole.LIMITED, # Minimal permissions\n ),\n user_id=None, # Service account (system-owned)\n ).api_key\n```\n\n### 2. Database Context Variables\n\nThe cache uses context variables for proper tenant-scoped DB sessions:\n\n```python\ncontext_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\ntry:\n with get_session_with_tenant(tenant_id) as db:\n # All DB operations scoped to this tenant\n ...\nfinally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(context_token)\n```\n\n### 3. Enterprise Gating Support\n\nGated tenants are filtered during cache refresh:\n\n```python\ngated_tenants = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.product_gating\",\n \"get_gated_tenants\",\n set(),\n)()\n\nfor tenant_id in get_all_tenant_ids():\n if tenant_id in gated_tenants:\n continue # Skip gated tenants\n```\n\n---\n\n## Cache Refresh Strategy\n\n| Trigger | Method | Scope |\n|---------|--------|-------|\n| Startup | `refresh_all()` | All tenants |\n| Periodic (60s) | `refresh_all()` | All tenants |\n| Guild registration | `refresh_guild()` | Single tenant |\n\n### Error Handling\n\n- **Tenant-level errors**: Logged and skipped (doesn't stop other tenants)\n- **Missing API key**: Bot silently ignores messages from that guild\n- **Network errors**: Logged, cache continues with stale data until next refresh\n\n---\n\n## Key Design Insights\n\n1. **Single Client, Multiple Tenants**: One `OnyxAPIClient` and one `DiscordCacheManager` instance serves all tenants via dynamic API key injection.\n\n2. **Cache-First Architecture**: Guild lookups are O(1) in-memory; API keys are cached after first provisioning to avoid repeated DB calls.\n\n3. **Graceful Degradation**: If an API key is missing or stale, the bot simply doesn't respond (no crash or error propagation).\n\n4. **Thread Safety Without Blocking**: `asyncio.Lock` prevents race conditions while maintaining async concurrency for reads.\n\n5. **Lazy Provisioning**: API keys are only created when first needed, then cached for performance.\n\n6. **Stateless API Client**: The HTTP client holds no tenant state - all tenant context is injected per-request via the `api_key` parameter.\n\n---\n\n## File References\n\n| Component | Path |\n|-----------|------|\n| Cache Manager | `backend/onyx/onyxbot/discord/cache.py` |\n| API Client | `backend/onyx/onyxbot/discord/api_client.py` |\n| Discord Client | `backend/onyx/onyxbot/discord/client.py` |\n| API Key DB Operations | `backend/onyx/db/discord_bot.py` |\n| Cache Manager Tests | `backend/tests/unit/onyx/onyxbot/discord/test_cache_manager.py` |\n| API Client Tests | `backend/tests/unit/onyx/onyxbot/discord/test_api_client.py` |" }, { "path": "backend/onyx/onyxbot/discord/api_client.py", "content": "\"\"\"Async HTTP client for communicating with Onyx API pods.\"\"\"\n\nimport aiohttp\n\nfrom onyx.chat.models import ChatFullResponse\nfrom onyx.onyxbot.discord.constants import API_REQUEST_TIMEOUT\nfrom onyx.onyxbot.discord.exceptions import APIConnectionError\nfrom onyx.onyxbot.discord.exceptions import APIResponseError\nfrom onyx.onyxbot.discord.exceptions import APITimeoutError\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.models import MessageOrigin\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import build_api_server_url_for_http_requests\n\nlogger = setup_logger()\n\n\nclass OnyxAPIClient:\n \"\"\"Async HTTP client for sending chat requests to Onyx API pods.\n\n This client manages an aiohttp session for making non-blocking HTTP\n requests to the Onyx API server. It handles authentication with per-tenant\n API keys and multi-tenant routing.\n\n Usage:\n client = OnyxAPIClient()\n await client.initialize()\n try:\n response = await client.send_chat_message(\n message=\"What is our deployment process?\",\n tenant_id=\"tenant_123\",\n api_key=\"dn_xxx...\",\n persona_id=1,\n )\n print(response.answer)\n finally:\n await client.close()\n \"\"\"\n\n def __init__(\n self,\n timeout: int = API_REQUEST_TIMEOUT,\n ) -> None:\n \"\"\"Initialize the API client.\n\n Args:\n timeout: Request timeout in seconds.\n \"\"\"\n # Helm chart uses API_SERVER_URL_OVERRIDE_FOR_HTTP_REQUESTS to set the base URL\n # TODO: Ideally, this override is only used when someone is launching an Onyx service independently\n self._base_url = build_api_server_url_for_http_requests(\n respect_env_override_if_set=True\n ).rstrip(\"/\")\n self._timeout = timeout\n self._session: aiohttp.ClientSession | None = None\n\n async def initialize(self) -> None:\n \"\"\"Create the aiohttp session.\n\n Must be called before making any requests. The session is created\n with a total timeout and connection timeout.\n \"\"\"\n if self._session is not None:\n logger.warning(\"API client session already initialized\")\n return\n\n timeout = aiohttp.ClientTimeout(\n total=self._timeout,\n connect=30, # 30 seconds to establish connection\n )\n self._session = aiohttp.ClientSession(timeout=timeout)\n logger.info(f\"API client initialized with base URL: {self._base_url}\")\n\n async def close(self) -> None:\n \"\"\"Close the aiohttp session.\n\n Should be called when shutting down the bot to properly release\n resources.\n \"\"\"\n if self._session is not None:\n await self._session.close()\n self._session = None\n logger.info(\"API client session closed\")\n\n @property\n def is_initialized(self) -> bool:\n \"\"\"Check if the session is initialized.\"\"\"\n return self._session is not None\n\n async def send_chat_message(\n self,\n message: str,\n api_key: str,\n persona_id: int | None = None,\n ) -> ChatFullResponse:\n \"\"\"Send a chat message to the Onyx API server and get a response.\n\n This method sends a non-streaming chat request to the API server. The response\n contains the complete answer with any citations and metadata.\n\n Args:\n message: The user's message to process.\n api_key: The API key for authentication.\n persona_id: Optional persona ID to use for the response.\n\n Returns:\n ChatFullResponse containing the answer, citations, and metadata.\n\n Raises:\n APIConnectionError: If unable to connect to the API.\n APITimeoutError: If the request times out.\n APIResponseError: If the API returns an error response.\n \"\"\"\n if self._session is None:\n raise APIConnectionError(\n \"API client not initialized. Call initialize() first.\"\n )\n\n url = f\"{self._base_url}/chat/send-chat-message\"\n\n # Build request payload\n request = SendMessageRequest(\n message=message,\n stream=False,\n origin=MessageOrigin.DISCORDBOT,\n chat_session_info=ChatSessionCreationRequest(\n persona_id=persona_id if persona_id is not None else 0,\n ),\n )\n\n # Build headers\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {api_key}\",\n }\n\n try:\n async with self._session.post(\n url,\n json=request.model_dump(mode=\"json\"),\n headers=headers,\n ) as response:\n if response.status == 401:\n raise APIResponseError(\n \"Authentication failed - invalid API key\",\n status_code=401,\n )\n elif response.status == 403:\n raise APIResponseError(\n \"Access denied - insufficient permissions\",\n status_code=403,\n )\n elif response.status == 404:\n raise APIResponseError(\n \"API endpoint not found\",\n status_code=404,\n )\n elif response.status >= 500:\n error_text = await response.text()\n raise APIResponseError(\n f\"Server error: {error_text}\",\n status_code=response.status,\n )\n elif response.status >= 400:\n error_text = await response.text()\n raise APIResponseError(\n f\"Request error: {error_text}\",\n status_code=response.status,\n )\n\n # Parse successful response\n data = await response.json()\n response_obj = ChatFullResponse.model_validate(data)\n\n if response_obj.error_msg:\n logger.warning(f\"Chat API returned error: {response_obj.error_msg}\")\n\n return response_obj\n\n except aiohttp.ClientConnectorError as e:\n logger.error(f\"Failed to connect to API: {e}\")\n raise APIConnectionError(\n f\"Failed to connect to API at {self._base_url}: {e}\"\n ) from e\n\n except TimeoutError as e:\n logger.error(f\"API request timed out after {self._timeout}s\")\n raise APITimeoutError(\n f\"Request timed out after {self._timeout} seconds\"\n ) from e\n\n except aiohttp.ClientError as e:\n logger.error(f\"HTTP client error: {e}\")\n raise APIConnectionError(f\"HTTP client error: {e}\") from e\n\n async def health_check(self) -> bool:\n \"\"\"Check if the API server is healthy.\n\n Returns:\n True if the API server is reachable and healthy, False otherwise.\n \"\"\"\n if self._session is None:\n logger.warning(\"API client not initialized. Call initialize() first.\")\n return False\n\n try:\n url = f\"{self._base_url}/health\"\n async with self._session.get(\n url, timeout=aiohttp.ClientTimeout(total=10)\n ) as response:\n return response.status == 200\n except Exception as e:\n logger.warning(f\"API server health check failed: {e}\")\n return False\n" }, { "path": "backend/onyx/onyxbot/discord/cache.py", "content": "\"\"\"Multi-tenant cache for Discord bot guild-tenant mappings and API keys.\"\"\"\n\nimport asyncio\n\nfrom onyx.db.discord_bot import get_guild_configs\nfrom onyx.db.discord_bot import get_or_create_discord_service_api_key\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.onyxbot.discord.exceptions import CacheError\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\n\nclass DiscordCacheManager:\n \"\"\"Caches guild->tenant mappings and tenant->API key mappings.\n\n Refreshed on startup, periodically (every 60s), and when guilds register.\n \"\"\"\n\n def __init__(self) -> None:\n self._guild_tenants: dict[int, str] = {} # guild_id -> tenant_id\n self._api_keys: dict[str, str] = {} # tenant_id -> api_key\n self._lock = asyncio.Lock()\n self._initialized = False\n\n @property\n def is_initialized(self) -> bool:\n return self._initialized\n\n async def refresh_all(self) -> None:\n \"\"\"Full cache refresh from all tenants.\"\"\"\n async with self._lock:\n logger.info(\"Starting Discord cache refresh\")\n\n new_guild_tenants: dict[int, str] = {}\n new_api_keys: dict[str, str] = {}\n\n try:\n gated = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.product_gating\",\n \"get_gated_tenants\",\n set(),\n )()\n\n tenant_ids = await asyncio.to_thread(get_all_tenant_ids)\n for tenant_id in tenant_ids:\n if tenant_id in gated:\n continue\n\n context_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n guild_ids, api_key = await self._load_tenant_data(tenant_id)\n if not guild_ids:\n logger.debug(f\"No guilds found for tenant {tenant_id}\")\n continue\n\n if not api_key:\n logger.warning(\n \"Discord service API key missing for tenant that has registered guilds. \"\n f\"{tenant_id} will not be handled in this refresh cycle.\"\n )\n continue\n\n for guild_id in guild_ids:\n new_guild_tenants[guild_id] = tenant_id\n\n new_api_keys[tenant_id] = api_key\n except Exception as e:\n logger.warning(f\"Failed to refresh tenant {tenant_id}: {e}\")\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(context_token)\n\n self._guild_tenants = new_guild_tenants\n self._api_keys = new_api_keys\n self._initialized = True\n\n logger.info(\n f\"Cache refresh complete: {len(new_guild_tenants)} guilds, {len(new_api_keys)} tenants\"\n )\n\n except Exception as e:\n logger.error(f\"Cache refresh failed: {e}\")\n raise CacheError(f\"Failed to refresh cache: {e}\") from e\n\n async def refresh_guild(self, guild_id: int, tenant_id: str) -> None:\n \"\"\"Add a single guild to cache after registration.\"\"\"\n async with self._lock:\n logger.info(f\"Refreshing cache for guild {guild_id} (tenant: {tenant_id})\")\n\n guild_ids, api_key = await self._load_tenant_data(tenant_id)\n\n if guild_id in guild_ids:\n self._guild_tenants[guild_id] = tenant_id\n if api_key:\n self._api_keys[tenant_id] = api_key\n logger.info(f\"Cache updated for guild {guild_id}\")\n else:\n logger.warning(f\"Guild {guild_id} not found or disabled\")\n\n async def _load_tenant_data(self, tenant_id: str) -> tuple[list[int], str | None]:\n \"\"\"Load guild IDs and provision API key if needed.\n\n Returns:\n (active_guild_ids, api_key) - api_key is the cached key if available,\n otherwise a newly created key. Returns None if no guilds found.\n \"\"\"\n cached_key = self._api_keys.get(tenant_id)\n\n def _sync() -> tuple[list[int], str | None]:\n with get_session_with_tenant(tenant_id=tenant_id) as db:\n configs = get_guild_configs(db)\n guild_ids = [\n config.guild_id\n for config in configs\n if config.enabled and config.guild_id is not None\n ]\n\n if not guild_ids:\n return [], None\n\n if not cached_key:\n new_key = get_or_create_discord_service_api_key(db, tenant_id)\n db.commit()\n return guild_ids, new_key\n\n return guild_ids, cached_key\n\n return await asyncio.to_thread(_sync)\n\n def get_tenant(self, guild_id: int) -> str | None:\n \"\"\"Get tenant ID for a guild.\"\"\"\n return self._guild_tenants.get(guild_id)\n\n def get_api_key(self, tenant_id: str) -> str | None:\n \"\"\"Get API key for a tenant.\"\"\"\n return self._api_keys.get(tenant_id)\n\n def remove_guild(self, guild_id: int) -> None:\n \"\"\"Remove a guild from cache.\"\"\"\n self._guild_tenants.pop(guild_id, None)\n\n def get_all_guild_ids(self) -> list[int]:\n \"\"\"Get all cached guild IDs.\"\"\"\n return list(self._guild_tenants.keys())\n\n def clear(self) -> None:\n \"\"\"Clear all caches.\"\"\"\n self._guild_tenants.clear()\n self._api_keys.clear()\n self._initialized = False\n" }, { "path": "backend/onyx/onyxbot/discord/client.py", "content": "\"\"\"Discord bot client with integrated message handling.\"\"\"\n\nimport asyncio\nimport time\n\nimport discord\nfrom discord.ext import commands\n\nfrom onyx.configs.app_configs import DISCORD_BOT_INVOKE_CHAR\nfrom onyx.onyxbot.discord.api_client import OnyxAPIClient\nfrom onyx.onyxbot.discord.cache import DiscordCacheManager\nfrom onyx.onyxbot.discord.constants import CACHE_REFRESH_INTERVAL\nfrom onyx.onyxbot.discord.handle_commands import handle_dm\nfrom onyx.onyxbot.discord.handle_commands import handle_registration_command\nfrom onyx.onyxbot.discord.handle_commands import handle_sync_channels_command\nfrom onyx.onyxbot.discord.handle_message import process_chat_message\nfrom onyx.onyxbot.discord.handle_message import should_respond\nfrom onyx.onyxbot.discord.utils import get_bot_token\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass OnyxDiscordClient(commands.Bot):\n \"\"\"Discord bot client with integrated cache, API client, and message handling.\n\n This client handles:\n - Guild registration via !register command\n - Message processing with persona-based responses\n - Thread context for conversation continuity\n - Multi-tenant support via cached API keys\n \"\"\"\n\n def __init__(self, command_prefix: str = DISCORD_BOT_INVOKE_CHAR) -> None:\n intents = discord.Intents.default()\n intents.message_content = True\n intents.members = True\n\n super().__init__(command_prefix=command_prefix, intents=intents)\n\n self.ready = False\n self.cache = DiscordCacheManager()\n self.api_client = OnyxAPIClient()\n self._cache_refresh_task: asyncio.Task | None = None\n\n # -------------------------------------------------------------------------\n # Lifecycle Methods\n # -------------------------------------------------------------------------\n\n async def setup_hook(self) -> None:\n \"\"\"Called before on_ready. Initialize components.\"\"\"\n logger.info(\"Initializing Discord bot components...\")\n\n # Initialize API client\n await self.api_client.initialize()\n\n # Initial cache load\n await self.cache.refresh_all()\n\n # Start periodic cache refresh\n self._cache_refresh_task = self.loop.create_task(self._periodic_cache_refresh())\n\n logger.info(\"Discord bot components initialized\")\n\n async def _periodic_cache_refresh(self) -> None:\n \"\"\"Background task to refresh cache periodically.\"\"\"\n while not self.is_closed():\n await asyncio.sleep(CACHE_REFRESH_INTERVAL)\n try:\n await self.cache.refresh_all()\n except Exception as e:\n logger.error(f\"Cache refresh failed: {e}\")\n\n async def on_ready(self) -> None:\n \"\"\"Bot connected and ready.\"\"\"\n if self.ready:\n return\n\n if not self.user:\n raise RuntimeError(\"Critical error: Discord Bot user not found\")\n\n logger.info(f\"Discord Bot connected as {self.user} (ID: {self.user.id})\")\n logger.info(f\"Connected to {len(self.guilds)} guild(s)\")\n logger.info(f\"Cached {len(self.cache.get_all_guild_ids())} registered guild(s)\")\n\n self.ready = True\n\n async def close(self) -> None:\n \"\"\"Graceful shutdown.\"\"\"\n logger.info(\"Shutting down Discord bot...\")\n\n # Cancel cache refresh task\n if self._cache_refresh_task:\n self._cache_refresh_task.cancel()\n try:\n await self._cache_refresh_task\n except asyncio.CancelledError:\n pass\n\n # Close Discord connection first - stops new commands from triggering cache ops\n if not self.is_closed():\n await super().close()\n\n # Close API client\n await self.api_client.close()\n\n # Clear cache (safe now - no concurrent operations possible)\n self.cache.clear()\n\n self.ready = False\n logger.info(\"Discord bot shutdown complete\")\n\n # -------------------------------------------------------------------------\n # Message Handling\n # -------------------------------------------------------------------------\n\n async def on_message(self, message: discord.Message) -> None:\n \"\"\"Main message handler.\"\"\"\n # mypy\n if not self.user:\n raise RuntimeError(\"Critical error: Discord Bot user not found\")\n\n try:\n # Ignore bot messages\n if message.author.bot:\n return\n\n # Ignore thread starter messages (empty reference nodes that don't contain content)\n if message.type == discord.MessageType.thread_starter_message:\n return\n\n # Handle DMs\n if isinstance(message.channel, discord.DMChannel):\n await handle_dm(message)\n return\n\n # Must have a guild\n if not message.guild or not message.guild.id:\n return\n\n guild_id = message.guild.id\n\n # Check for registration command first\n if await handle_registration_command(message, self.cache):\n return\n\n # Look up guild in cache\n tenant_id = self.cache.get_tenant(guild_id)\n\n # Check for sync-channels command (requires registered guild)\n if await handle_sync_channels_command(message, tenant_id, self):\n return\n\n if not tenant_id:\n # Guild not registered, ignore\n return\n\n # Get API key\n api_key = self.cache.get_api_key(tenant_id)\n if not api_key:\n logger.warning(f\"No API key cached for tenant {tenant_id}\")\n return\n\n # Check if bot should respond\n should_respond_context = await should_respond(message, tenant_id, self.user)\n\n if not should_respond_context.should_respond:\n return\n\n logger.debug(\n f\"Processing message: '{message.content[:50]}' in \"\n f\"#{getattr(message.channel, 'name', 'unknown')} ({message.guild.name}), \"\n f\"persona_id={should_respond_context.persona_id}\"\n )\n\n # Process the message\n await process_chat_message(\n message=message,\n api_key=api_key,\n persona_id=should_respond_context.persona_id,\n thread_only_mode=should_respond_context.thread_only_mode,\n api_client=self.api_client,\n bot_user=self.user,\n )\n\n except Exception as e:\n logger.exception(f\"Error processing message: {e}\")\n\n\n# -----------------------------------------------------------------------------\n# Entry Point\n# -----------------------------------------------------------------------------\n\n\ndef main() -> None:\n \"\"\"Main entry point for Discord bot.\"\"\"\n from onyx.db.engine.sql_engine import SqlEngine\n from onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\n\n logger.info(\"Starting Onyx Discord Bot...\")\n\n # Initialize the database engine (required before any DB operations)\n SqlEngine.init_engine(pool_size=20, max_overflow=5)\n\n # Initialize EE features based on environment\n set_is_ee_based_on_env_variable()\n\n counter = 0\n while True:\n token = get_bot_token()\n if not token:\n if counter % 180 == 0:\n logger.info(\n \"Discord bot is dormant. Waiting for token configuration...\"\n )\n counter += 1\n time.sleep(5)\n continue\n counter = 0\n bot = OnyxDiscordClient()\n\n try:\n # bot.run() handles SIGINT/SIGTERM and calls close() automatically\n bot.run(token)\n\n except Exception:\n logger.exception(\"Fatal error in Discord bot\")\n raise\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/onyxbot/discord/constants.py", "content": "\"\"\"Discord bot constants.\"\"\"\n\n# API settings\nAPI_REQUEST_TIMEOUT: int = 3 * 60 # 3 minutes\n\n# Cache settings\nCACHE_REFRESH_INTERVAL: int = 60 # 1 minute\n\n# Message settings\nMAX_MESSAGE_LENGTH: int = 2000 # Discord's character limit\nMAX_CONTEXT_MESSAGES: int = 10 # Max messages to include in conversation context\n# Note: Discord.py's add_reaction() requires unicode emoji, not :name: format\nTHINKING_EMOJI: str = \"🤔\" # U+1F914 - Thinking Face\nSUCCESS_EMOJI: str = \"✅\" # U+2705 - White Heavy Check Mark\nERROR_EMOJI: str = \"❌\" # U+274C - Cross Mark\n\n# Command prefix\nREGISTER_COMMAND: str = \"register\"\nSYNC_CHANNELS_COMMAND: str = \"sync-channels\"\n" }, { "path": "backend/onyx/onyxbot/discord/exceptions.py", "content": "\"\"\"Custom exception classes for Discord bot.\"\"\"\n\n\nclass DiscordBotError(Exception):\n \"\"\"Base exception for Discord bot errors.\"\"\"\n\n\nclass RegistrationError(DiscordBotError):\n \"\"\"Error during guild registration.\"\"\"\n\n\nclass SyncChannelsError(DiscordBotError):\n \"\"\"Error during channel sync.\"\"\"\n\n\nclass APIError(DiscordBotError):\n \"\"\"Base API error.\"\"\"\n\n\nclass CacheError(DiscordBotError):\n \"\"\"Error during cache operations.\"\"\"\n\n\nclass APIConnectionError(APIError):\n \"\"\"Failed to connect to API.\"\"\"\n\n\nclass APITimeoutError(APIError):\n \"\"\"Request timed out.\"\"\"\n\n\nclass APIResponseError(APIError):\n \"\"\"API returned an error response.\"\"\"\n\n def __init__(self, message: str, status_code: int | None = None):\n super().__init__(message)\n self.status_code = status_code\n" }, { "path": "backend/onyx/onyxbot/discord/handle_commands.py", "content": "\"\"\"Discord bot command handlers for registration and channel sync.\"\"\"\n\nimport asyncio\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport discord\n\nfrom onyx.configs.app_configs import DISCORD_BOT_INVOKE_CHAR\nfrom onyx.configs.constants import ONYX_DISCORD_URL\nfrom onyx.db.discord_bot import bulk_create_channel_configs\nfrom onyx.db.discord_bot import get_guild_config_by_discord_id\nfrom onyx.db.discord_bot import get_guild_config_by_internal_id\nfrom onyx.db.discord_bot import get_guild_config_by_registration_key\nfrom onyx.db.discord_bot import sync_channel_configs\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.utils import DiscordChannelView\nfrom onyx.onyxbot.discord.cache import DiscordCacheManager\nfrom onyx.onyxbot.discord.constants import REGISTER_COMMAND\nfrom onyx.onyxbot.discord.constants import SYNC_CHANNELS_COMMAND\nfrom onyx.onyxbot.discord.exceptions import RegistrationError\nfrom onyx.onyxbot.discord.exceptions import SyncChannelsError\nfrom onyx.server.manage.discord_bot.utils import parse_discord_registration_key\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\n\nasync def handle_dm(message: discord.Message) -> None:\n \"\"\"Handle direct messages.\"\"\"\n dm_response = (\n \"**I can't respond to DMs** :sweat:\\n\\n\"\n f\"Please chat with me in a server channel, or join the official \"\n f\"[Onyx Discord]({ONYX_DISCORD_URL}) for help!\"\n )\n await message.channel.send(dm_response)\n\n\n# -------------------------------------------------------------------------\n# Helper functions for error handling\n# -------------------------------------------------------------------------\n\n\nasync def _try_dm_author(message: discord.Message, content: str) -> bool:\n \"\"\"Attempt to DM the message author. Returns True if successful.\"\"\"\n logger.debug(f\"Responding in Discord DM with {content}\")\n try:\n await message.author.send(content)\n return True\n except (discord.Forbidden, discord.HTTPException) as e:\n # User has DMs disabled or other error\n logger.warning(f\"Failed to DM author {message.author.id}: {e}\")\n except Exception as e:\n logger.exception(f\"Unexpected error DMing author {message.author.id}: {e}\")\n return False\n\n\nasync def _try_delete_message(message: discord.Message) -> bool:\n \"\"\"Attempt to delete a message. Returns True if successful.\"\"\"\n logger.debug(f\"Deleting potentially sensitive message {message.id}\")\n try:\n await message.delete()\n return True\n except (discord.Forbidden, discord.HTTPException) as e:\n # Bot lacks permission or other error\n logger.warning(f\"Failed to delete message {message.id}: {e}\")\n except Exception as e:\n logger.exception(f\"Unexpected error deleting message {message.id}: {e}\")\n return False\n\n\nasync def _try_react_x(message: discord.Message) -> bool:\n \"\"\"Attempt to react to a message with ❌. Returns True if successful.\"\"\"\n try:\n await message.add_reaction(\"❌\")\n return True\n except (discord.Forbidden, discord.HTTPException) as e:\n # Bot lacks permission or other error\n logger.warning(f\"Failed to react to message {message.id}: {e}\")\n except Exception as e:\n logger.exception(f\"Unexpected error reacting to message {message.id}: {e}\")\n return False\n\n\n# -------------------------------------------------------------------------\n# Registration\n# -------------------------------------------------------------------------\n\n\nasync def handle_registration_command(\n message: discord.Message,\n cache: DiscordCacheManager,\n) -> bool:\n \"\"\"Handle !register command. Returns True if command was handled.\"\"\"\n content = message.content.strip()\n\n # Check for !register command\n if not content.startswith(f\"{DISCORD_BOT_INVOKE_CHAR}{REGISTER_COMMAND}\"):\n return False\n\n # Must be in a server\n if not message.guild:\n await _try_dm_author(\n message, \"This command can only be used in a server channel.\"\n )\n return True\n\n guild_name = message.guild.name\n logger.info(f\"Registration command received: {guild_name}\")\n\n try:\n # Parse the registration key\n parts = content.split(maxsplit=1)\n if len(parts) < 2:\n raise RegistrationError(\n \"Invalid registration key format. Please check the key and try again.\"\n )\n\n registration_key = parts[1].strip()\n\n if not message.author or not isinstance(message.author, discord.Member):\n raise RegistrationError(\n \"You need to be a server administrator to register the bot.\"\n )\n\n # Check permissions - require admin or manage_guild\n if not message.author.guild_permissions.administrator:\n if not message.author.guild_permissions.manage_guild:\n raise RegistrationError(\n \"You need **Administrator** or **Manage Server** permissions to register this bot.\"\n )\n\n await _register_guild(message, registration_key, cache)\n logger.info(f\"Registration successful: {guild_name}\")\n await message.reply(\n \":white_check_mark: **Successfully registered!**\\n\\n\"\n \"This server is now connected to Onyx. \"\n \"I'll respond to messages based on your server and channel settings set in Onyx.\"\n )\n except RegistrationError as e:\n logger.debug(f\"Registration failed: {guild_name}, error={e}\")\n await _try_dm_author(message, f\":x: **Registration failed.**\\n\\n{e}\")\n await _try_delete_message(message)\n except Exception:\n logger.exception(f\"Registration failed unexpectedly: {guild_name}\")\n await _try_dm_author(\n message,\n \":x: **Registration failed.**\\n\\nAn unexpected error occurred. Please try again later.\",\n )\n await _try_delete_message(message)\n\n return True\n\n\nasync def _register_guild(\n message: discord.Message,\n registration_key: str,\n cache: DiscordCacheManager,\n) -> None:\n \"\"\"Register a guild with a registration key.\"\"\"\n if not message.guild:\n # mypy, even though we already know that message.guild is not None\n raise RegistrationError(\"This command can only be used in a server.\")\n\n logger.info(f\"Guild '{message.guild.name}' attempting to register Discord bot\")\n registration_key = registration_key.strip()\n\n # Parse tenant_id from registration key\n parsed = parse_discord_registration_key(registration_key)\n if parsed is None:\n raise RegistrationError(\n \"Invalid registration key format. Please check the key and try again.\"\n )\n\n tenant_id = parsed\n\n logger.info(f\"Parsed tenant_id {tenant_id} from registration key\")\n\n # Check if this guild is already registered to any tenant\n guild_id = message.guild.id\n existing_tenant = cache.get_tenant(guild_id)\n if existing_tenant is not None:\n logger.warning(\n f\"Guild {guild_id} is already registered to tenant {existing_tenant}\"\n )\n raise RegistrationError(\n \"This server is already registered.\\n\\nOnyxBot can only connect one Discord server to one Onyx workspace.\"\n )\n\n context_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n guild = message.guild\n guild_name = guild.name\n\n # Collect all text channels from the guild\n channels = get_text_channels(guild)\n logger.info(f\"Found {len(channels)} text channels in guild '{guild_name}'\")\n\n # Validate and update in database\n def _sync_register() -> int:\n with get_session_with_tenant(tenant_id=tenant_id) as db:\n # Find the guild config by registration key\n config = get_guild_config_by_registration_key(db, registration_key)\n if not config:\n raise RegistrationError(\n \"Registration key not found.\\n\\n\"\n \"The key may have expired or been deleted. \"\n \"Please generate a new one from the Onyx admin panel.\"\n )\n\n # Check if already used\n if config.guild_id is not None:\n raise RegistrationError(\n \"This registration key has already been used.\\n\\n\"\n \"Each key can only be used once. \"\n \"Please generate a new key from the Onyx admin panel.\"\n )\n\n # Update the guild config\n config.guild_id = guild_id\n config.guild_name = guild_name\n config.registered_at = datetime.now(timezone.utc)\n\n # Create channel configs for all text channels\n bulk_create_channel_configs(db, config.id, channels)\n\n db.commit()\n return config.id\n\n await asyncio.to_thread(_sync_register)\n\n # Refresh cache for this guild\n await cache.refresh_guild(guild_id, tenant_id)\n\n logger.info(\n f\"Guild '{guild_name}' registered with {len(channels)} channel configs\"\n )\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(context_token)\n\n\ndef get_text_channels(guild: discord.Guild) -> list[DiscordChannelView]:\n \"\"\"Get all text channels from a guild as DiscordChannelView objects.\"\"\"\n channels: list[DiscordChannelView] = []\n for channel in guild.channels:\n # Include text channels and forum channels (where threads can be created)\n if isinstance(channel, (discord.TextChannel, discord.ForumChannel)):\n # Check if channel is private (not visible to @everyone)\n everyone_perms = channel.permissions_for(guild.default_role)\n is_private = not everyone_perms.view_channel\n\n logger.debug(\n f\"Found channel: #{channel.name}, type={channel.type.name}, is_private={is_private}\"\n )\n\n channels.append(\n DiscordChannelView(\n channel_id=channel.id,\n channel_name=channel.name,\n channel_type=channel.type.name, # \"text\" or \"forum\"\n is_private=is_private,\n )\n )\n\n logger.debug(f\"Retrieved {len(channels)} channels from guild '{guild.name}'\")\n return channels\n\n\n# -------------------------------------------------------------------------\n# Sync Channels\n# -------------------------------------------------------------------------\n\n\nasync def handle_sync_channels_command(\n message: discord.Message,\n tenant_id: str | None,\n bot: discord.Client,\n) -> bool:\n \"\"\"Handle !sync-channels command. Returns True if command was handled.\"\"\"\n content = message.content.strip()\n\n # Check for !sync-channels command\n if not content.startswith(f\"{DISCORD_BOT_INVOKE_CHAR}{SYNC_CHANNELS_COMMAND}\"):\n return False\n\n # Must be in a server\n if not message.guild:\n await _try_dm_author(\n message, \"This command can only be used in a server channel.\"\n )\n return True\n\n guild_name = message.guild.name\n logger.info(f\"Sync-channels command received: {guild_name}\")\n\n try:\n # Must be registered\n if not tenant_id:\n raise SyncChannelsError(\n \"This server is not registered. Please register it first.\"\n )\n\n # Check permissions - require admin or manage_guild\n if not message.author or not isinstance(message.author, discord.Member):\n raise SyncChannelsError(\n \"You need to be a server administrator to sync channels.\"\n )\n\n if not message.author.guild_permissions.administrator:\n if not message.author.guild_permissions.manage_guild:\n raise SyncChannelsError(\n \"You need **Administrator** or **Manage Server** permissions to sync channels.\"\n )\n\n # Get guild config ID\n def _get_guild_config_id() -> int | None:\n with get_session_with_tenant(tenant_id=tenant_id) as db:\n if not message.guild:\n raise SyncChannelsError(\n \"Server not found. This shouldn't happen. Please contact Onyx support.\"\n )\n config = get_guild_config_by_discord_id(db, message.guild.id)\n return config.id if config else None\n\n guild_config_id = await asyncio.to_thread(_get_guild_config_id)\n\n if not guild_config_id:\n raise SyncChannelsError(\n \"Server config not found. This shouldn't happen. Please contact Onyx support.\"\n )\n\n # Perform the sync\n added, removed, updated = await sync_guild_channels(\n guild_config_id, tenant_id, bot\n )\n logger.info(\n f\"Sync-channels successful: {guild_name}, added={added}, removed={removed}, updated={updated}\"\n )\n await message.reply(\n f\":white_check_mark: **Channel sync complete!**\\n\\n\"\n f\"* **{added}** new channel(s) added\\n\"\n f\"* **{removed}** deleted channel(s) removed\\n\"\n f\"* **{updated}** channel name(s) updated\\n\\n\"\n \"New channels are disabled by default. Enable them in the Onyx admin panel.\"\n )\n except SyncChannelsError as e:\n logger.debug(f\"Sync-channels failed: {guild_name}, error={e}\")\n await _try_dm_author(message, f\":x: **Channel sync failed.**\\n\\n{e}\")\n await _try_react_x(message)\n except Exception:\n logger.exception(f\"Sync-channels failed unexpectedly: {guild_name}\")\n await _try_dm_author(\n message,\n \":x: **Channel sync failed.**\\n\\nAn unexpected error occurred. Please try again later.\",\n )\n await _try_react_x(message)\n\n return True\n\n\nasync def sync_guild_channels(\n guild_config_id: int,\n tenant_id: str,\n bot: discord.Client,\n) -> tuple[int, int, int]:\n \"\"\"Sync channel configs with current Discord channels for a guild.\n\n Fetches current channels from Discord and syncs with database:\n - Creates configs for new channels (disabled by default)\n - Removes configs for deleted channels\n - Updates names for existing channels if changed\n\n Args:\n guild_config_id: Internal ID of the guild config\n tenant_id: Tenant ID for database access\n bot: Discord bot client\n\n Returns:\n (added_count, removed_count, updated_count)\n\n Raises:\n ValueError: If guild config not found or guild not registered\n \"\"\"\n context_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)\n try:\n # Get guild_id from config\n def _get_guild_id() -> int | None:\n with get_session_with_tenant(tenant_id=tenant_id) as db:\n config = get_guild_config_by_internal_id(db, guild_config_id)\n if not config:\n return None\n return config.guild_id\n\n guild_id = await asyncio.to_thread(_get_guild_id)\n\n if guild_id is None:\n raise ValueError(\n f\"Guild config {guild_config_id} not found or not registered\"\n )\n\n # Get the guild from Discord\n guild = bot.get_guild(guild_id)\n if not guild:\n raise ValueError(f\"Guild {guild_id} not found in Discord cache\")\n\n # Get current channels from Discord\n channels = get_text_channels(guild)\n logger.info(f\"Syncing {len(channels)} channels for guild '{guild.name}'\")\n\n # Sync with database\n def _sync() -> tuple[int, int, int]:\n with get_session_with_tenant(tenant_id=tenant_id) as db:\n added, removed, updated = sync_channel_configs(\n db, guild_config_id, channels\n )\n db.commit()\n return added, removed, updated\n\n added, removed, updated = await asyncio.to_thread(_sync)\n\n logger.info(\n f\"Channel sync complete for guild '{guild.name}': added={added}, removed={removed}, updated={updated}\"\n )\n\n return added, removed, updated\n\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(context_token)\n" }, { "path": "backend/onyx/onyxbot/discord/handle_message.py", "content": "\"\"\"Discord bot message handling and response logic.\"\"\"\n\nimport asyncio\n\nimport discord\nfrom pydantic import BaseModel\n\nfrom onyx.chat.models import ChatFullResponse\nfrom onyx.db.discord_bot import get_channel_config_by_discord_ids\nfrom onyx.db.discord_bot import get_guild_config_by_discord_id\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.models import DiscordChannelConfig\nfrom onyx.db.models import DiscordGuildConfig\nfrom onyx.onyxbot.discord.api_client import OnyxAPIClient\nfrom onyx.onyxbot.discord.constants import MAX_CONTEXT_MESSAGES\nfrom onyx.onyxbot.discord.constants import MAX_MESSAGE_LENGTH\nfrom onyx.onyxbot.discord.constants import THINKING_EMOJI\nfrom onyx.onyxbot.discord.exceptions import APIError\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# Message types with actual content (excludes system notifications like \"user joined\")\nCONTENT_MESSAGE_TYPES = (\n discord.MessageType.default,\n discord.MessageType.reply,\n discord.MessageType.thread_starter_message,\n)\n\n\nclass ShouldRespondContext(BaseModel):\n \"\"\"Context for whether the bot should respond to a message.\"\"\"\n\n should_respond: bool\n persona_id: int | None\n thread_only_mode: bool\n\n\n# -------------------------------------------------------------------------\n# Response Logic\n# -------------------------------------------------------------------------\n\n\nasync def should_respond(\n message: discord.Message,\n tenant_id: str,\n bot_user: discord.ClientUser,\n) -> ShouldRespondContext:\n \"\"\"Determine if bot should respond and which persona to use.\"\"\"\n if not message.guild:\n logger.warning(\"Received a message that isn't in a server.\")\n return ShouldRespondContext(\n should_respond=False, persona_id=None, thread_only_mode=False\n )\n\n guild_id = message.guild.id\n channel_id = message.channel.id\n bot_mentioned = bot_user in message.mentions\n\n def _get_configs() -> tuple[DiscordGuildConfig | None, DiscordChannelConfig | None]:\n with get_session_with_tenant(tenant_id=tenant_id) as db:\n guild_config = get_guild_config_by_discord_id(db, guild_id)\n if not guild_config or not guild_config.enabled:\n return None, None\n\n # For threads, use parent channel ID\n actual_channel_id = channel_id\n if isinstance(message.channel, discord.Thread) and message.channel.parent:\n actual_channel_id = message.channel.parent.id\n\n channel_config = get_channel_config_by_discord_ids(\n db, guild_id, actual_channel_id\n )\n return guild_config, channel_config\n\n guild_config, channel_config = await asyncio.to_thread(_get_configs)\n\n if not guild_config or not channel_config or not channel_config.enabled:\n return ShouldRespondContext(\n should_respond=False, persona_id=None, thread_only_mode=False\n )\n\n # Determine persona (channel override or guild default)\n persona_id = channel_config.persona_override_id or guild_config.default_persona_id\n\n # Check mention requirement (with exceptions for implicit invocation)\n if channel_config.require_bot_invocation and not bot_mentioned:\n if not await check_implicit_invocation(message, bot_user):\n return ShouldRespondContext(\n should_respond=False, persona_id=None, thread_only_mode=False\n )\n\n return ShouldRespondContext(\n should_respond=True,\n persona_id=persona_id,\n thread_only_mode=channel_config.thread_only_mode,\n )\n\n\nasync def check_implicit_invocation(\n message: discord.Message,\n bot_user: discord.ClientUser,\n) -> bool:\n \"\"\"Check if the bot should respond without explicit mention.\n\n Returns True if:\n 1. User is replying to a bot message\n 2. User is in a thread owned by the bot\n 3. User is in a thread created from a bot message\n \"\"\"\n # Check if replying to a bot message\n if message.reference and message.reference.message_id:\n try:\n referenced_msg = await message.channel.fetch_message(\n message.reference.message_id\n )\n if referenced_msg.author.id == bot_user.id:\n logger.debug(\n f\"Implicit invocation via reply: '{message.content[:50]}...'\"\n )\n return True\n except (discord.NotFound, discord.HTTPException):\n pass\n\n # Check thread-related conditions\n if isinstance(message.channel, discord.Thread):\n thread = message.channel\n\n # Bot owns the thread\n if thread.owner_id == bot_user.id:\n logger.debug(\n f\"Implicit invocation via bot-owned thread: '{message.content[:50]}...' in #{thread.name}\"\n )\n return True\n\n # Thread was created from a bot message\n if thread.parent and not isinstance(thread.parent, discord.ForumChannel):\n try:\n starter = await thread.parent.fetch_message(thread.id)\n if starter.author.id == bot_user.id:\n logger.debug(\n f\"Implicit invocation via bot-started thread: '{message.content[:50]}...' in #{thread.name}\"\n )\n return True\n except (discord.NotFound, discord.HTTPException):\n pass\n\n return False\n\n\n# -------------------------------------------------------------------------\n# Message Processing\n# -------------------------------------------------------------------------\n\n\nasync def process_chat_message(\n message: discord.Message,\n api_key: str,\n persona_id: int | None,\n thread_only_mode: bool,\n api_client: OnyxAPIClient,\n bot_user: discord.ClientUser,\n) -> None:\n \"\"\"Process a message and send response.\"\"\"\n try:\n await message.add_reaction(THINKING_EMOJI)\n except discord.DiscordException:\n logger.warning(\n f\"Failed to add thinking reaction to message: '{message.content[:50]}...'\"\n )\n\n try:\n # Build conversation context\n context = await _build_conversation_context(message, bot_user)\n\n # Prepare full message content\n parts = []\n if context:\n parts.append(context)\n if isinstance(message.channel, discord.Thread):\n if isinstance(message.channel.parent, discord.ForumChannel):\n parts.append(f\"Forum post title: {message.channel.name}\")\n parts.append(\n f\"Current message from @{message.author.display_name}: {format_message_content(message)}\"\n )\n\n # Send to API\n response = await api_client.send_chat_message(\n message=\"\\n\\n\".join(parts),\n api_key=api_key,\n persona_id=persona_id,\n )\n\n # Format response with citations\n answer = response.answer or \"I couldn't generate a response.\"\n answer = _append_citations(answer, response)\n\n await send_response(message, answer, thread_only_mode)\n\n try:\n await message.remove_reaction(THINKING_EMOJI, bot_user)\n except discord.DiscordException:\n pass\n\n except APIError as e:\n logger.error(f\"API error processing message: {e}\")\n await send_error_response(message, bot_user)\n except Exception as e:\n logger.exception(f\"Error processing chat message: {e}\")\n await send_error_response(message, bot_user)\n\n\nasync def _build_conversation_context(\n message: discord.Message,\n bot_user: discord.ClientUser,\n) -> str | None:\n \"\"\"Build conversation context from thread history or reply chain.\"\"\"\n if isinstance(message.channel, discord.Thread):\n return await _build_thread_context(message, bot_user)\n elif message.reference:\n return await _build_reply_chain_context(message, bot_user)\n return None\n\n\ndef _append_citations(answer: str, response: ChatFullResponse) -> str:\n \"\"\"Append citation sources to the answer if present.\"\"\"\n if not response.citation_info or not response.top_documents:\n return answer\n\n cited_docs: list[tuple[int, str, str | None]] = []\n for citation in response.citation_info:\n doc = next(\n (\n d\n for d in response.top_documents\n if d.document_id == citation.document_id\n ),\n None,\n )\n if doc:\n cited_docs.append(\n (\n citation.citation_number,\n doc.semantic_identifier or \"Source\",\n doc.link,\n )\n )\n\n if not cited_docs:\n return answer\n\n cited_docs.sort(key=lambda x: x[0])\n citations = \"\\n\\n**Sources:**\\n\"\n for num, name, link in cited_docs[:5]:\n if link:\n citations += f\"{num}. [{name}](<{link}>)\\n\"\n else:\n citations += f\"{num}. {name}\\n\"\n\n return answer + citations\n\n\n# -------------------------------------------------------------------------\n# Context Building\n# -------------------------------------------------------------------------\n\n\nasync def _build_reply_chain_context(\n message: discord.Message,\n bot_user: discord.ClientUser,\n) -> str | None:\n \"\"\"Build context by following the reply chain backwards.\"\"\"\n if not message.reference or not message.reference.message_id:\n return None\n\n try:\n messages: list[discord.Message] = []\n current = message\n\n # Follow reply chain backwards up to MAX_CONTEXT_MESSAGES\n while (\n current.reference\n and current.reference.message_id\n and len(messages) < MAX_CONTEXT_MESSAGES\n ):\n try:\n parent = await message.channel.fetch_message(\n current.reference.message_id\n )\n messages.append(parent)\n current = parent\n except (discord.NotFound, discord.HTTPException):\n break\n\n if not messages:\n return None\n\n messages.reverse() # Chronological order\n\n logger.debug(\n f\"Built reply chain context: {len(messages)} messages in #{getattr(message.channel, 'name', 'unknown')}\"\n )\n\n return _format_messages_as_context(messages, bot_user)\n\n except Exception as e:\n logger.warning(f\"Failed to build reply chain context: {e}\")\n return None\n\n\nasync def _build_thread_context(\n message: discord.Message,\n bot_user: discord.ClientUser,\n) -> str | None:\n \"\"\"Build context from thread message history.\"\"\"\n if not isinstance(message.channel, discord.Thread):\n return None\n\n try:\n thread = message.channel\n messages: list[discord.Message] = []\n\n # Fetch recent messages (excluding current)\n async for msg in thread.history(limit=MAX_CONTEXT_MESSAGES, oldest_first=False):\n if msg.id != message.id:\n messages.append(msg)\n\n # Include thread starter message and its reply chain if not already present\n if thread.parent and not isinstance(thread.parent, discord.ForumChannel):\n try:\n starter = await thread.parent.fetch_message(thread.id)\n if starter.id != message.id and not any(\n m.id == starter.id for m in messages\n ):\n messages.append(starter)\n\n # Trace back through the starter's reply chain for more context\n current = starter\n while (\n current.reference\n and current.reference.message_id\n and len(messages) < MAX_CONTEXT_MESSAGES\n ):\n try:\n parent = await thread.parent.fetch_message(\n current.reference.message_id\n )\n if not any(m.id == parent.id for m in messages):\n messages.append(parent)\n current = parent\n except (discord.NotFound, discord.HTTPException):\n break\n except (discord.NotFound, discord.HTTPException):\n pass\n\n if not messages:\n return None\n\n messages.sort(key=lambda m: m.id) # Chronological order\n logger.debug(\n f\"Built thread context: {len(messages)} messages in #{thread.name}\"\n )\n\n return _format_messages_as_context(messages, bot_user)\n\n except Exception as e:\n logger.warning(f\"Failed to build thread context: {e}\")\n return None\n\n\ndef _format_messages_as_context(\n messages: list[discord.Message],\n bot_user: discord.ClientUser,\n) -> str | None:\n \"\"\"Format a list of messages into a conversation context string.\"\"\"\n formatted = []\n for msg in messages:\n if msg.type not in CONTENT_MESSAGE_TYPES:\n continue\n\n sender = (\n \"OnyxBot\" if msg.author.id == bot_user.id else f\"@{msg.author.display_name}\"\n )\n formatted.append(f\"{sender}: {format_message_content(msg)}\")\n\n if not formatted:\n return None\n\n return (\n \"You are a Discord bot named OnyxBot.\\n\"\n 'Always assume that [user] is the same as the \"Current message\" author.'\n \"Conversation history:\\n\"\n \"---\\n\" + \"\\n\".join(formatted) + \"\\n---\"\n )\n\n\n# -------------------------------------------------------------------------\n# Message Formatting\n# -------------------------------------------------------------------------\n\n\ndef format_message_content(message: discord.Message) -> str:\n \"\"\"Format message content with readable mentions.\"\"\"\n content = message.content\n\n for user in message.mentions:\n content = content.replace(f\"<@{user.id}>\", f\"@{user.display_name}\")\n content = content.replace(f\"<@!{user.id}>\", f\"@{user.display_name}\")\n\n for role in message.role_mentions:\n content = content.replace(f\"<@&{role.id}>\", f\"@{role.name}\")\n\n for channel in message.channel_mentions:\n content = content.replace(f\"<#{channel.id}>\", f\"#{channel.name}\")\n\n return content\n\n\n# -------------------------------------------------------------------------\n# Response Sending\n# -------------------------------------------------------------------------\n\n\nasync def send_response(\n message: discord.Message,\n content: str,\n thread_only_mode: bool,\n) -> None:\n \"\"\"Send response based on thread_only_mode setting.\"\"\"\n chunks = _split_message(content)\n\n if isinstance(message.channel, discord.Thread):\n for chunk in chunks:\n await message.channel.send(chunk)\n elif thread_only_mode:\n thread_name = f\"OnyxBot <> {message.author.display_name}\"[:100]\n thread = await message.create_thread(name=thread_name)\n for chunk in chunks:\n await thread.send(chunk)\n else:\n for i, chunk in enumerate(chunks):\n if i == 0:\n await message.reply(chunk)\n else:\n await message.channel.send(chunk)\n\n\ndef _split_message(content: str) -> list[str]:\n \"\"\"Split content into chunks that fit Discord's message limit.\"\"\"\n chunks = []\n while content:\n if len(content) <= MAX_MESSAGE_LENGTH:\n chunks.append(content)\n break\n\n # Find a good split point\n split_at = MAX_MESSAGE_LENGTH\n for sep in [\"\\n\\n\", \"\\n\", \". \", \" \"]:\n idx = content.rfind(sep, 0, MAX_MESSAGE_LENGTH)\n if idx > MAX_MESSAGE_LENGTH // 2:\n split_at = idx + len(sep)\n break\n\n chunks.append(content[:split_at])\n content = content[split_at:]\n\n return chunks\n\n\nasync def send_error_response(\n message: discord.Message,\n bot_user: discord.ClientUser,\n) -> None:\n \"\"\"Send error response and clean up reaction.\"\"\"\n try:\n await message.remove_reaction(THINKING_EMOJI, bot_user)\n except discord.DiscordException:\n pass\n\n error_msg = \"Sorry, I encountered an error processing your message. You may want to contact Onyx for support :sweat_smile:\"\n\n try:\n if isinstance(message.channel, discord.Thread):\n await message.channel.send(error_msg)\n else:\n thread = await message.create_thread(\n name=f\"Response to {message.author.display_name}\"[:100]\n )\n await thread.send(error_msg)\n except discord.DiscordException:\n pass\n" }, { "path": "backend/onyx/onyxbot/discord/utils.py", "content": "from onyx.configs.app_configs import AUTH_TYPE\nfrom onyx.configs.app_configs import DISCORD_BOT_TOKEN\nfrom onyx.configs.constants import AuthType\nfrom onyx.db.discord_bot import get_discord_bot_config\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.sensitive import SensitiveValue\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\nlogger = setup_logger()\n\n\ndef get_bot_token() -> str | None:\n \"\"\"Get Discord bot token from env var or database.\n\n Priority:\n 1. DISCORD_BOT_TOKEN env var (always takes precedence)\n 2. For self-hosted: DiscordBotConfig in database (default tenant)\n 3. For Cloud: should always have env var set\n\n Returns:\n Bot token string, or None if not configured.\n \"\"\"\n # Environment variable takes precedence\n if DISCORD_BOT_TOKEN:\n return DISCORD_BOT_TOKEN\n\n # Cloud should always have env var; if not, return None\n if AUTH_TYPE == AuthType.CLOUD:\n logger.warning(\"Cloud deployment missing DISCORD_BOT_TOKEN env var\")\n return None\n\n # Self-hosted: check database for bot config\n try:\n with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db:\n config = get_discord_bot_config(db)\n except Exception as e:\n logger.error(f\"Failed to get bot token from database: {e}\")\n return None\n if config and config.bot_token:\n if isinstance(config.bot_token, SensitiveValue):\n return config.bot_token.get_value(apply_mask=False)\n return config.bot_token\n return None\n" }, { "path": "backend/onyx/onyxbot/slack/blocks.py", "content": "from datetime import datetime\nfrom typing import cast\n\nimport pytz\nimport timeago # type: ignore\nfrom slack_sdk.models.blocks import ActionsBlock\nfrom slack_sdk.models.blocks import Block\nfrom slack_sdk.models.blocks import ButtonElement\nfrom slack_sdk.models.blocks import ContextBlock\nfrom slack_sdk.models.blocks import DividerBlock\nfrom slack_sdk.models.blocks import HeaderBlock\nfrom slack_sdk.models.blocks import Option\nfrom slack_sdk.models.blocks import RadioButtonsElement\nfrom slack_sdk.models.blocks import SectionBlock\nfrom slack_sdk.models.blocks.basic_components import MarkdownTextObject\nfrom slack_sdk.models.blocks.block_elements import ImageElement\n\nfrom onyx.chat.models import ChatBasicResponse\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import SearchFeedbackType\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_NUM_DOCS_TO_DISPLAY\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.chat import get_chat_session_by_message_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import ChannelConfig\nfrom onyx.onyxbot.slack.constants import CONTINUE_IN_WEB_UI_ACTION_ID\nfrom onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_RESOLVED_ACTION_ID\nfrom onyx.onyxbot.slack.constants import IMMEDIATE_RESOLVED_BUTTON_ACTION_ID\nfrom onyx.onyxbot.slack.constants import KEEP_TO_YOURSELF_ACTION_ID\nfrom onyx.onyxbot.slack.constants import LIKE_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import SHOW_EVERYONE_ACTION_ID\nfrom onyx.onyxbot.slack.formatting import format_slack_message\nfrom onyx.onyxbot.slack.icons import source_to_github_img_link\nfrom onyx.onyxbot.slack.models import ActionValuesEphemeralMessage\nfrom onyx.onyxbot.slack.models import ActionValuesEphemeralMessageChannelConfig\nfrom onyx.onyxbot.slack.models import ActionValuesEphemeralMessageMessageInfo\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.onyxbot.slack.utils import build_continue_in_web_ui_id\nfrom onyx.onyxbot.slack.utils import build_feedback_id\nfrom onyx.onyxbot.slack.utils import build_publish_ephemeral_message_id\nfrom onyx.onyxbot.slack.utils import remove_slack_text_interactions\nfrom onyx.onyxbot.slack.utils import translate_vespa_highlight_to_slack\nfrom onyx.utils.text_processing import decode_escapes\n\n_MAX_BLURB_LEN = 45\n\n\ndef _format_doc_updated_at(updated_at: datetime | None) -> str | None:\n \"\"\"Convert document timestamps to a human friendly relative string.\"\"\"\n if updated_at is None:\n return None\n\n if updated_at.tzinfo is None or updated_at.tzinfo.utcoffset(updated_at) is None:\n aware_updated_at = updated_at.replace(tzinfo=pytz.utc)\n else:\n aware_updated_at = updated_at.astimezone(pytz.utc)\n\n return timeago.format(aware_updated_at, datetime.now(pytz.utc))\n\n\ndef get_feedback_reminder_blocks(thread_link: str, include_followup: bool) -> Block:\n text = (\n f\"Please provide feedback on <{thread_link}|this answer>. \"\n \"This is essential to help us to improve the quality of the answers. \"\n \"Please rate it by clicking the `Helpful` or `Not helpful` button. \"\n )\n if include_followup:\n text += \"\\n\\nIf you need more help, click the `I need more help from a human!` button. \"\n\n text += \"\\n\\nThanks!\"\n\n return SectionBlock(text=text)\n\n\ndef _split_text(text: str, limit: int = 3000) -> list[str]:\n if len(text) <= limit:\n return [text]\n\n chunks = []\n while text:\n if len(text) <= limit:\n chunks.append(text)\n break\n\n # Find the nearest space before the limit to avoid splitting a word\n split_at = text.rfind(\" \", 0, limit)\n if split_at == -1: # No spaces found, force split\n split_at = limit\n\n chunk = text[:split_at]\n chunks.append(chunk)\n text = text[split_at:].lstrip() # Remove leading spaces from the next chunk\n\n return chunks\n\n\ndef _clean_markdown_link_text(text: str) -> str:\n # Remove any newlines within the text\n return format_slack_message(text).replace(\"\\n\", \" \").strip()\n\n\ndef _build_qa_feedback_block(\n message_id: int, feedback_reminder_id: str | None = None\n) -> Block:\n return ActionsBlock(\n block_id=build_feedback_id(message_id),\n elements=[\n ButtonElement(\n action_id=LIKE_BLOCK_ACTION_ID,\n text=\"👍 Helpful\",\n value=feedback_reminder_id,\n ),\n ButtonElement(\n action_id=DISLIKE_BLOCK_ACTION_ID,\n text=\"👎 Not helpful\",\n value=feedback_reminder_id,\n ),\n ],\n )\n\n\ndef _build_ephemeral_publication_block(\n channel_id: str, # noqa: ARG001\n chat_message_id: int,\n message_info: SlackMessageInfo,\n original_question_ts: str,\n channel_conf: ChannelConfig,\n feedback_reminder_id: str | None = None,\n) -> Block:\n # check whether the message is in a thread\n if (\n message_info is not None\n and message_info.msg_to_respond is not None\n and message_info.thread_to_respond is not None\n and (message_info.msg_to_respond == message_info.thread_to_respond)\n ):\n respond_ts = None\n else:\n respond_ts = original_question_ts\n\n action_values_ephemeral_message_channel_config = (\n ActionValuesEphemeralMessageChannelConfig(\n channel_name=channel_conf.get(\"channel_name\"),\n respond_tag_only=channel_conf.get(\"respond_tag_only\"),\n respond_to_bots=channel_conf.get(\"respond_to_bots\"),\n is_ephemeral=channel_conf.get(\"is_ephemeral\", False),\n respond_member_group_list=channel_conf.get(\"respond_member_group_list\"),\n answer_filters=channel_conf.get(\"answer_filters\"),\n follow_up_tags=channel_conf.get(\"follow_up_tags\"),\n show_continue_in_web_ui=channel_conf.get(\"show_continue_in_web_ui\", False),\n )\n )\n\n action_values_ephemeral_message_message_info = (\n ActionValuesEphemeralMessageMessageInfo(\n bypass_filters=message_info.bypass_filters,\n channel_to_respond=message_info.channel_to_respond,\n msg_to_respond=message_info.msg_to_respond,\n email=message_info.email,\n sender_id=message_info.sender_id,\n thread_messages=[],\n is_slash_command=message_info.is_slash_command,\n is_bot_dm=message_info.is_bot_dm,\n thread_to_respond=respond_ts,\n )\n )\n\n action_values_ephemeral_message = ActionValuesEphemeralMessage(\n original_question_ts=original_question_ts,\n feedback_reminder_id=feedback_reminder_id,\n chat_message_id=chat_message_id,\n message_info=action_values_ephemeral_message_message_info,\n channel_conf=action_values_ephemeral_message_channel_config,\n )\n\n return ActionsBlock(\n block_id=build_publish_ephemeral_message_id(original_question_ts),\n elements=[\n ButtonElement(\n action_id=SHOW_EVERYONE_ACTION_ID,\n text=\"📢 Share with Everyone\",\n value=action_values_ephemeral_message.model_dump_json(),\n ),\n ButtonElement(\n action_id=KEEP_TO_YOURSELF_ACTION_ID,\n text=\"🤫 Keep to Yourself\",\n value=action_values_ephemeral_message.model_dump_json(),\n ),\n ],\n )\n\n\ndef get_document_feedback_blocks() -> Block:\n return SectionBlock(\n text=(\n \"- 'Up-Boost' if this document is a good source of information and should be \"\n \"shown more often.\\n\"\n \"- 'Down-boost' if this document is a poor source of information and should be \"\n \"shown less often.\\n\"\n \"- 'Hide' if this document is deprecated and should never be shown anymore.\"\n ),\n accessory=RadioButtonsElement(\n options=[\n Option(\n text=\":thumbsup: Up-Boost\",\n value=SearchFeedbackType.ENDORSE.value,\n ),\n Option(\n text=\":thumbsdown: Down-Boost\",\n value=SearchFeedbackType.REJECT.value,\n ),\n Option(\n text=\":x: Hide\",\n value=SearchFeedbackType.HIDE.value,\n ),\n ]\n ),\n )\n\n\ndef _build_doc_feedback_block(\n message_id: int,\n document_id: str,\n document_rank: int,\n) -> ButtonElement:\n feedback_id = build_feedback_id(message_id, document_id, document_rank)\n return ButtonElement(\n action_id=FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID,\n value=feedback_id,\n text=\"Give Feedback\",\n )\n\n\ndef get_restate_blocks(\n msg: str,\n is_slash_command: bool,\n) -> list[Block]:\n # Only the slash command needs this context because the user doesn't see their own input\n if not is_slash_command:\n return []\n\n return [\n HeaderBlock(text=\"Responding to the Query\"),\n SectionBlock(text=f\"```{msg}```\"),\n ]\n\n\ndef _build_documents_blocks(\n documents: list[SearchDoc],\n message_id: int | None,\n num_docs_to_display: int = ONYX_BOT_NUM_DOCS_TO_DISPLAY,\n) -> list[Block]:\n header_text = \"Reference Documents\"\n seen_docs_identifiers = set()\n section_blocks: list[Block] = [HeaderBlock(text=header_text)]\n included_docs = 0\n for rank, d in enumerate(documents):\n if d.document_id in seen_docs_identifiers:\n continue\n seen_docs_identifiers.add(d.document_id)\n\n # Strip newlines from the semantic identifier for Slackbot formatting\n doc_sem_id = d.semantic_identifier.replace(\"\\n\", \" \")\n if d.source_type == DocumentSource.SLACK.value:\n doc_sem_id = \"#\" + doc_sem_id\n\n used_chars = len(doc_sem_id) + 3\n match_str = translate_vespa_highlight_to_slack(d.match_highlights, used_chars)\n\n included_docs += 1\n\n header_line = f\"{doc_sem_id}\\n\"\n if d.link:\n header_line = f\"<{d.link}|{doc_sem_id}>\\n\"\n\n updated_at_line = \"\"\n updated_at_str = _format_doc_updated_at(d.updated_at)\n if updated_at_str:\n updated_at_line = f\"_Updated {updated_at_str}_\\n\"\n\n body_text = f\">{remove_slack_text_interactions(match_str)}\"\n\n block_text = header_line + updated_at_line + body_text\n\n feedback: ButtonElement | dict = {}\n if message_id is not None:\n feedback = _build_doc_feedback_block(\n message_id=message_id,\n document_id=d.document_id,\n document_rank=rank,\n )\n\n section_blocks.append(\n SectionBlock(text=block_text, accessory=feedback),\n )\n\n section_blocks.append(DividerBlock())\n\n if included_docs >= num_docs_to_display:\n break\n\n return section_blocks\n\n\ndef _build_sources_blocks(\n cited_documents: list[tuple[int, SearchDoc]],\n num_docs_to_display: int = ONYX_BOT_NUM_DOCS_TO_DISPLAY,\n) -> list[Block]:\n if not cited_documents:\n return [\n SectionBlock(\n text=\"*Warning*: no sources were cited for this answer, so it may be unreliable 😔\"\n )\n ]\n\n seen_docs_identifiers = set()\n section_blocks: list[Block] = [SectionBlock(text=\"*Sources:*\")]\n included_docs = 0\n for citation_num, d in cited_documents:\n if d.document_id in seen_docs_identifiers:\n continue\n seen_docs_identifiers.add(d.document_id)\n\n doc_sem_id = d.semantic_identifier\n if d.source_type == DocumentSource.SLACK.value:\n # for legacy reasons, before the switch to how Slack semantic identifiers are constructed\n if \"#\" not in doc_sem_id:\n doc_sem_id = \"#\" + doc_sem_id\n\n # this is needed to try and prevent the line from overflowing\n # if it does overflow, the image gets placed above the title and it\n # looks bad\n doc_sem_id = (\n doc_sem_id[:_MAX_BLURB_LEN] + \"...\"\n if len(doc_sem_id) > _MAX_BLURB_LEN\n else doc_sem_id\n )\n\n owner_str = f\"By {d.primary_owners[0]}\" if d.primary_owners else None\n days_ago_str = _format_doc_updated_at(d.updated_at)\n final_metadata_str = \" | \".join(\n ([owner_str] if owner_str else [])\n + ([days_ago_str] if days_ago_str else [])\n )\n\n document_title = _clean_markdown_link_text(doc_sem_id)\n img_link = source_to_github_img_link(d.source_type)\n\n section_blocks.append(\n ContextBlock(\n elements=(\n [\n ImageElement(\n image_url=img_link,\n alt_text=f\"{d.source_type.value} logo\",\n )\n ]\n if img_link\n else []\n )\n + [\n (\n MarkdownTextObject(text=f\"{document_title}\")\n if d.link == \"\"\n else MarkdownTextObject(\n text=f\"*<{d.link}|[{citation_num}] {document_title}>*\\n{final_metadata_str}\"\n )\n ),\n ]\n )\n )\n\n if included_docs >= num_docs_to_display:\n break\n\n return section_blocks\n\n\ndef _priority_ordered_documents_blocks(\n answer: ChatBasicResponse,\n) -> list[Block]:\n top_docs = answer.top_documents if answer.top_documents else None\n if not top_docs:\n return []\n\n document_blocks = _build_documents_blocks(\n documents=top_docs,\n message_id=answer.message_id,\n )\n if document_blocks:\n document_blocks = [DividerBlock()] + document_blocks\n return document_blocks\n\n\ndef _build_citations_blocks(\n answer: ChatBasicResponse,\n) -> list[Block]:\n top_docs = answer.top_documents\n citations = answer.citation_info or []\n cited_docs: list[tuple[int, SearchDoc]] = []\n for citation_info in citations:\n matching_doc = next(\n (d for d in top_docs if d.document_id == citation_info.document_id),\n None,\n )\n if matching_doc:\n cited_docs.append((citation_info.citation_number, matching_doc))\n\n cited_docs.sort()\n citations_block = _build_sources_blocks(cited_documents=cited_docs)\n return citations_block\n\n\ndef _build_main_response_blocks(\n answer: ChatBasicResponse,\n) -> list[Block]:\n # TODO: add back in later when auto-filtering is implemented\n # if (\n # retrieval_info.applied_time_cutoff\n # or retrieval_info.recency_bias_multiplier > 1\n # or retrieval_info.applied_source_filters\n # ):\n # filter_text = \"Filters: \"\n # if retrieval_info.applied_source_filters:\n # sources_str = \", \".join(\n # [s.value for s in retrieval_info.applied_source_filters]\n # )\n # filter_text += f\"`Sources in [{sources_str}]`\"\n # if (\n # retrieval_info.applied_time_cutoff\n # or retrieval_info.recency_bias_multiplier > 1\n # ):\n # filter_text += \" and \"\n # if retrieval_info.applied_time_cutoff is not None:\n # time_str = retrieval_info.applied_time_cutoff.strftime(\"%b %d, %Y\")\n # filter_text += f\"`Docs Updated >= {time_str}` \"\n # if retrieval_info.recency_bias_multiplier > 1:\n # if retrieval_info.applied_time_cutoff is not None:\n # filter_text += \"+ \"\n # filter_text += \"`Prioritize Recently Updated Docs`\"\n\n # filter_block = SectionBlock(text=f\"_{filter_text}_\")\n\n # replaces markdown links with slack format links\n formatted_answer = format_slack_message(answer.answer)\n answer_processed = decode_escapes(remove_slack_text_interactions(formatted_answer))\n answer_blocks = [SectionBlock(text=text) for text in _split_text(answer_processed)]\n\n return cast(list[Block], answer_blocks)\n\n\ndef _build_continue_in_web_ui_block(\n message_id: int | None,\n) -> Block:\n if message_id is None:\n raise ValueError(\"No message id provided to build continue in web ui block\")\n with get_session_with_current_tenant() as db_session:\n chat_session = get_chat_session_by_message_id(\n db_session=db_session,\n message_id=message_id,\n )\n return ActionsBlock(\n block_id=build_continue_in_web_ui_id(message_id),\n elements=[\n ButtonElement(\n action_id=CONTINUE_IN_WEB_UI_ACTION_ID,\n text=\"Continue Chat in Onyx!\",\n style=\"primary\",\n url=f\"{WEB_DOMAIN}/chat?slackChatId={chat_session.id}\",\n ),\n ],\n )\n\n\ndef _build_follow_up_block(message_id: int | None) -> ActionsBlock:\n return ActionsBlock(\n block_id=build_feedback_id(message_id) if message_id is not None else None,\n elements=[\n ButtonElement(\n action_id=IMMEDIATE_RESOLVED_BUTTON_ACTION_ID,\n style=\"primary\",\n text=\"I'm all set!\",\n ),\n ButtonElement(\n action_id=FOLLOWUP_BUTTON_ACTION_ID,\n style=\"danger\",\n text=\"I need more help from a human!\",\n ),\n ],\n )\n\n\ndef build_follow_up_resolved_blocks(\n tag_ids: list[str], group_ids: list[str]\n) -> list[Block]:\n tag_str = \" \".join([f\"<@{tag}>\" for tag in tag_ids])\n if tag_str:\n tag_str += \" \"\n\n group_str = \" \".join([f\"\" for group_id in group_ids])\n if group_str:\n group_str += \" \"\n\n text = (\n tag_str\n + group_str\n + \"Someone has requested more help.\\n\\n:point_down:Please mark this resolved after answering!\"\n )\n text_block = SectionBlock(text=text)\n button_block = ActionsBlock(\n elements=[\n ButtonElement(\n action_id=FOLLOWUP_BUTTON_RESOLVED_ACTION_ID,\n style=\"primary\",\n text=\"Mark Resolved\",\n )\n ]\n )\n return [text_block, button_block]\n\n\ndef build_slack_response_blocks(\n answer: ChatBasicResponse,\n message_info: SlackMessageInfo,\n channel_conf: ChannelConfig | None,\n feedback_reminder_id: str | None,\n skip_ai_feedback: bool = False,\n offer_ephemeral_publication: bool = False,\n skip_restated_question: bool = False,\n) -> list[Block]:\n \"\"\"\n This function is a top level function that builds all the blocks for the Slack response.\n It also handles combining all the blocks together.\n \"\"\"\n # If called with the OnyxBot slash command, the question is lost so we have to reshow it\n if not skip_restated_question:\n restate_question_block = get_restate_blocks(\n message_info.thread_messages[-1].message, message_info.is_slash_command\n )\n else:\n restate_question_block = []\n\n answer_blocks = _build_main_response_blocks(answer)\n\n web_follow_up_block = []\n if channel_conf and channel_conf.get(\"show_continue_in_web_ui\"):\n web_follow_up_block.append(\n _build_continue_in_web_ui_block(\n message_id=answer.message_id,\n )\n )\n\n follow_up_block = []\n if (\n channel_conf\n and channel_conf.get(\"follow_up_tags\") is not None\n and not channel_conf.get(\"is_ephemeral\", False)\n ):\n follow_up_block.append(_build_follow_up_block(message_id=answer.message_id))\n\n publish_ephemeral_message_block = []\n\n if (\n offer_ephemeral_publication\n and answer.message_id is not None\n and message_info.msg_to_respond is not None\n and channel_conf is not None\n ):\n publish_ephemeral_message_block.append(\n _build_ephemeral_publication_block(\n channel_id=message_info.channel_to_respond,\n chat_message_id=answer.message_id,\n original_question_ts=message_info.msg_to_respond,\n message_info=message_info,\n channel_conf=channel_conf,\n feedback_reminder_id=feedback_reminder_id,\n )\n )\n\n ai_feedback_block: list[Block] = []\n\n if answer.message_id is not None and not skip_ai_feedback:\n ai_feedback_block.append(\n _build_qa_feedback_block(\n message_id=answer.message_id,\n feedback_reminder_id=feedback_reminder_id,\n )\n )\n\n citations_blocks = []\n if answer.citation_info:\n citations_blocks = _build_citations_blocks(answer)\n\n citations_divider = [DividerBlock()] if citations_blocks else []\n buttons_divider = [DividerBlock()] if web_follow_up_block or follow_up_block else []\n\n all_blocks = (\n restate_question_block\n + answer_blocks\n + publish_ephemeral_message_block\n + ai_feedback_block\n + citations_divider\n + citations_blocks\n + buttons_divider\n + web_follow_up_block\n + follow_up_block\n )\n\n return all_blocks\n" }, { "path": "backend/onyx/onyxbot/slack/config.py", "content": "import os\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.db.slack_channel_config import (\n fetch_slack_channel_config_for_channel_or_default,\n)\nfrom onyx.db.slack_channel_config import fetch_slack_channel_configs\n\nVALID_SLACK_FILTERS = [\n \"answerable_prefilter\",\n \"well_answered_postfilter\",\n \"questionmark_prefilter\",\n]\n\n\ndef get_slack_channel_config_for_bot_and_channel(\n db_session: Session,\n slack_bot_id: int,\n channel_name: str | None,\n) -> SlackChannelConfig:\n slack_bot_config = fetch_slack_channel_config_for_channel_or_default(\n db_session=db_session, slack_bot_id=slack_bot_id, channel_name=channel_name\n )\n if not slack_bot_config:\n raise ValueError(\n \"No default configuration has been set for this Slack bot. This should not be possible.\"\n )\n\n return slack_bot_config\n\n\ndef validate_channel_name(\n db_session: Session,\n current_slack_bot_id: int,\n channel_name: str,\n current_slack_channel_config_id: int | None,\n) -> str:\n \"\"\"Make sure that this channel_name does not exist in other Slack channel configs.\n Returns a cleaned up channel name (e.g. '#' removed if present)\"\"\"\n slack_bot_configs = fetch_slack_channel_configs(\n db_session=db_session,\n slack_bot_id=current_slack_bot_id,\n )\n cleaned_channel_name = channel_name.lstrip(\"#\").lower()\n for slack_channel_config in slack_bot_configs:\n if slack_channel_config.id == current_slack_channel_config_id:\n continue\n\n if cleaned_channel_name == slack_channel_config.channel_config[\"channel_name\"]:\n raise ValueError(\n f\"Channel name '{channel_name}' already exists in \"\n \"another Slack channel config with in Slack Bot with name: \"\n f\"{slack_channel_config.slack_bot.name}\"\n )\n\n return cleaned_channel_name\n\n\n# Scaling configurations for multi-tenant Slack channel handling\nTENANT_LOCK_EXPIRATION = 1800 # How long a pod can hold exclusive access to a tenant before other pods can acquire it\nTENANT_HEARTBEAT_INTERVAL = (\n 15 # How often pods send heartbeats to indicate they are still processing a tenant\n)\nTENANT_HEARTBEAT_EXPIRATION = (\n 60 # How long before a tenant's heartbeat expires, allowing other pods to take over\n)\nTENANT_ACQUISITION_INTERVAL = 60 # How often pods attempt to acquire unprocessed tenants and checks for new tokens\n\nMAX_TENANTS_PER_POD = int(os.getenv(\"MAX_TENANTS_PER_POD\", 50))\n" }, { "path": "backend/onyx/onyxbot/slack/constants.py", "content": "import re\nfrom enum import Enum\n\n# Matches Slack channel references like <#C097NBWMY8Y> or <#C097NBWMY8Y|channel-name>\nSLACK_CHANNEL_REF_PATTERN = re.compile(r\"<#([A-Z0-9]+)(?:\\|([^>]+))?>\")\n\nLIKE_BLOCK_ACTION_ID = \"feedback-like\"\nDISLIKE_BLOCK_ACTION_ID = \"feedback-dislike\"\nSHOW_EVERYONE_ACTION_ID = \"show-everyone\"\nKEEP_TO_YOURSELF_ACTION_ID = \"keep-to-yourself\"\nCONTINUE_IN_WEB_UI_ACTION_ID = \"continue-in-web-ui\"\nFEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID = \"feedback-doc-button\"\nIMMEDIATE_RESOLVED_BUTTON_ACTION_ID = \"immediate-resolved-button\"\nFOLLOWUP_BUTTON_ACTION_ID = \"followup-button\"\nFOLLOWUP_BUTTON_RESOLVED_ACTION_ID = \"followup-resolved-button\"\nVIEW_DOC_FEEDBACK_ID = \"view-doc-feedback\"\nGENERATE_ANSWER_BUTTON_ACTION_ID = \"generate-answer-button\"\n\n\nclass FeedbackVisibility(str, Enum):\n PRIVATE = \"private\"\n ANONYMOUS = \"anonymous\"\n PUBLIC = \"public\"\n" }, { "path": "backend/onyx/onyxbot/slack/formatting.py", "content": "import re\nfrom collections.abc import Callable\nfrom typing import Any\n\nfrom mistune import create_markdown\nfrom mistune import HTMLRenderer\n\n# Tags that should be replaced with a newline (line-break and block-level elements)\n_HTML_NEWLINE_TAG_PATTERN = re.compile(\n r\"|\",\n re.IGNORECASE,\n)\n\n# Strips HTML tags but excludes autolinks like and \n_HTML_TAG_PATTERN = re.compile(\n r\"<(?!https?://|mailto:)/?[a-zA-Z][^>]*>\",\n)\n\n# Matches fenced code blocks (``` ... ```) so we can skip sanitization inside them\n_FENCED_CODE_BLOCK_PATTERN = re.compile(r\"```[\\s\\S]*?```\")\n\n# Matches the start of any markdown link: [text]( or [[n]](\n# The inner group handles nested brackets for citation links like [[1]](.\n_MARKDOWN_LINK_PATTERN = re.compile(r\"\\[(?:[^\\[\\]]|\\[[^\\]]*\\])*\\]\\(\")\n\n# Matches Slack-style links that LLMs sometimes output directly.\n# Mistune doesn't recognise this syntax, so text() would escape the angle\n# brackets and Slack would render them as literal text instead of links.\n_SLACK_LINK_PATTERN = re.compile(r\"<(https?://[^|>]+)\\|([^>]+)>\")\n\n\ndef _sanitize_html(text: str) -> str:\n \"\"\"Strip HTML tags from a text fragment.\n\n Block-level closing tags and
    are converted to newlines.\n All other HTML tags are removed. Autolinks () are preserved.\n \"\"\"\n text = _HTML_NEWLINE_TAG_PATTERN.sub(\"\\n\", text)\n text = _HTML_TAG_PATTERN.sub(\"\", text)\n return text\n\n\ndef _transform_outside_code_blocks(\n message: str, transform: Callable[[str], str]\n) -> str:\n \"\"\"Apply *transform* only to text outside fenced code blocks.\"\"\"\n parts = _FENCED_CODE_BLOCK_PATTERN.split(message)\n code_blocks = _FENCED_CODE_BLOCK_PATTERN.findall(message)\n\n result: list[str] = []\n for i, part in enumerate(parts):\n result.append(transform(part))\n if i < len(code_blocks):\n result.append(code_blocks[i])\n\n return \"\".join(result)\n\n\ndef _extract_link_destination(message: str, start_idx: int) -> tuple[str, int | None]:\n \"\"\"Extract markdown link destination, allowing nested parentheses in the URL.\"\"\"\n depth = 0\n i = start_idx\n\n while i < len(message):\n curr = message[i]\n if curr == \"\\\\\":\n i += 2\n continue\n\n if curr == \"(\":\n depth += 1\n elif curr == \")\":\n if depth == 0:\n return message[start_idx:i], i\n depth -= 1\n i += 1\n\n return message[start_idx:], None\n\n\ndef _normalize_link_destinations(message: str) -> str:\n \"\"\"Wrap markdown link URLs in angle brackets so the parser handles special chars safely.\n\n Markdown link syntax [text](url) breaks when the URL contains unescaped\n parentheses, spaces, or other special characters. Wrapping the URL in angle\n brackets — [text]() — tells the parser to treat everything inside as\n a literal URL. This applies to all links, not just citations.\n \"\"\"\n if \"](\" not in message:\n return message\n\n normalized_parts: list[str] = []\n cursor = 0\n\n while match := _MARKDOWN_LINK_PATTERN.search(message, cursor):\n normalized_parts.append(message[cursor : match.end()])\n destination_start = match.end()\n destination, end_idx = _extract_link_destination(message, destination_start)\n if end_idx is None:\n normalized_parts.append(message[destination_start:])\n return \"\".join(normalized_parts)\n\n already_wrapped = destination.startswith(\"<\") and destination.endswith(\">\")\n if destination and not already_wrapped:\n destination = f\"<{destination}>\"\n\n normalized_parts.append(destination)\n normalized_parts.append(\")\")\n cursor = end_idx + 1\n\n normalized_parts.append(message[cursor:])\n return \"\".join(normalized_parts)\n\n\ndef _convert_slack_links_to_markdown(message: str) -> str:\n \"\"\"Convert Slack-style links to standard markdown [text](url).\n\n LLMs sometimes emit Slack mrkdwn link syntax directly. Mistune doesn't\n recognise it, so the angle brackets would be escaped by text() and Slack\n would render the link as literal text instead of a clickable link.\n \"\"\"\n return _transform_outside_code_blocks(\n message, lambda text: _SLACK_LINK_PATTERN.sub(r\"[\\2](\\1)\", text)\n )\n\n\ndef format_slack_message(message: str | None) -> str:\n if message is None:\n return \"\"\n message = _transform_outside_code_blocks(message, _sanitize_html)\n message = _convert_slack_links_to_markdown(message)\n normalized_message = _normalize_link_destinations(message)\n md = create_markdown(renderer=SlackRenderer(), plugins=[\"strikethrough\", \"table\"])\n result = md(normalized_message)\n # With HTMLRenderer, result is always str (not AST list)\n assert isinstance(result, str)\n return result.rstrip(\"\\n\")\n\n\nclass SlackRenderer(HTMLRenderer):\n \"\"\"Renders markdown as Slack mrkdwn format instead of HTML.\n\n Overrides all HTMLRenderer methods that produce HTML tags to ensure\n no raw HTML ever appears in Slack messages.\n \"\"\"\n\n SPECIALS: dict[str, str] = {\"&\": \"&\", \"<\": \"<\", \">\": \">\"}\n\n def __init__(self) -> None:\n super().__init__()\n self._table_headers: list[str] = []\n self._current_row_cells: list[str] = []\n\n def escape_special(self, text: str) -> str:\n for special, replacement in self.SPECIALS.items():\n text = text.replace(special, replacement)\n return text\n\n def heading(self, text: str, level: int, **attrs: Any) -> str: # noqa: ARG002\n return f\"*{text}*\\n\\n\"\n\n def emphasis(self, text: str) -> str:\n return f\"_{text}_\"\n\n def strong(self, text: str) -> str:\n return f\"*{text}*\"\n\n def strikethrough(self, text: str) -> str:\n return f\"~{text}~\"\n\n def list(self, text: str, ordered: bool, **attrs: Any) -> str: # noqa: ARG002\n lines = text.split(\"\\n\")\n count = 0\n for i, line in enumerate(lines):\n if line.startswith(\"li: \"):\n count += 1\n prefix = f\"{count}. \" if ordered else \"• \"\n lines[i] = f\"{prefix}{line[4:]}\"\n return \"\\n\".join(lines) + \"\\n\"\n\n def list_item(self, text: str) -> str:\n return f\"li: {text}\\n\"\n\n def link(self, text: str, url: str, title: str | None = None) -> str:\n escaped_url = self.escape_special(url)\n if text:\n return f\"<{escaped_url}|{text}>\"\n if title:\n return f\"<{escaped_url}|{title}>\"\n return f\"<{escaped_url}>\"\n\n def image(self, text: str, url: str, title: str | None = None) -> str:\n escaped_url = self.escape_special(url)\n display_text = title or text\n return f\"<{escaped_url}|{display_text}>\" if display_text else f\"<{escaped_url}>\"\n\n def codespan(self, text: str) -> str:\n return f\"`{text}`\"\n\n def block_code(self, code: str, info: str | None = None) -> str: # noqa: ARG002\n return f\"```\\n{code.rstrip(chr(10))}\\n```\\n\\n\"\n\n def linebreak(self) -> str:\n return \"\\n\"\n\n def thematic_break(self) -> str:\n return \"---\\n\\n\"\n\n def block_quote(self, text: str) -> str:\n lines = text.strip().split(\"\\n\")\n quoted = \"\\n\".join(f\">{line}\" for line in lines)\n return quoted + \"\\n\\n\"\n\n def block_html(self, html: str) -> str:\n return _sanitize_html(html) + \"\\n\\n\"\n\n def block_error(self, text: str) -> str:\n return f\"```\\n{text}\\n```\\n\\n\"\n\n def text(self, text: str) -> str:\n # Only escape the three entities Slack recognizes: & < >\n # HTMLRenderer.text() also escapes \" to " which Slack renders\n # as literal " text since Slack doesn't recognize that entity.\n return self.escape_special(text)\n\n # -- Table rendering (converts markdown tables to vertical cards) --\n\n def table_cell(\n self,\n text: str,\n align: str | None = None, # noqa: ARG002\n head: bool = False, # noqa: ARG002\n ) -> str:\n if head:\n self._table_headers.append(text.strip())\n else:\n self._current_row_cells.append(text.strip())\n return \"\"\n\n def table_head(self, text: str) -> str: # noqa: ARG002\n self._current_row_cells = []\n return \"\"\n\n def table_row(self, text: str) -> str: # noqa: ARG002\n cells = self._current_row_cells\n self._current_row_cells = []\n # First column becomes the bold title, remaining columns are bulleted fields\n lines: list[str] = []\n if cells:\n title = cells[0]\n if title:\n # Avoid double-wrapping if cell already contains bold markup\n if title.startswith(\"*\") and title.endswith(\"*\") and len(title) > 1:\n lines.append(title)\n else:\n lines.append(f\"*{title}*\")\n for i, cell in enumerate(cells[1:], start=1):\n if i < len(self._table_headers):\n lines.append(f\" • {self._table_headers[i]}: {cell}\")\n else:\n lines.append(f\" • {cell}\")\n return \"\\n\".join(lines) + \"\\n\\n\"\n\n def table_body(self, text: str) -> str:\n return text\n\n def table(self, text: str) -> str:\n self._table_headers = []\n self._current_row_cells = []\n return text + \"\\n\"\n\n def paragraph(self, text: str) -> str:\n return f\"{text}\\n\\n\"\n" }, { "path": "backend/onyx/onyxbot/slack/handlers/__init__.py", "content": "" }, { "path": "backend/onyx/onyxbot/slack/handlers/handle_buttons.py", "content": "import json\nfrom typing import Any\nfrom typing import cast\n\nfrom slack_sdk import WebClient\nfrom slack_sdk.models.blocks import SectionBlock\nfrom slack_sdk.models.views import View\nfrom slack_sdk.socket_mode.request import SocketModeRequest\nfrom slack_sdk.webhook import WebhookClient\n\nfrom onyx.chat.models import ChatBasicResponse\nfrom onyx.chat.process_message import remove_answer_citations\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import SearchFeedbackType\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_FOLLOWUP_EMOJI\nfrom onyx.connectors.slack.utils import expert_info_from_slack_id\nfrom onyx.context.search.models import SavedSearchDoc\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.chat import get_chat_message\nfrom onyx.db.chat import translate_db_message_to_chat_message_detail\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.feedback import create_chat_message_feedback\nfrom onyx.db.feedback import create_doc_retrieval_feedback\nfrom onyx.db.users import get_user_by_email\nfrom onyx.onyxbot.slack.blocks import build_follow_up_resolved_blocks\nfrom onyx.onyxbot.slack.blocks import build_slack_response_blocks\nfrom onyx.onyxbot.slack.blocks import get_document_feedback_blocks\nfrom onyx.onyxbot.slack.config import get_slack_channel_config_for_bot_and_channel\nfrom onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FeedbackVisibility\nfrom onyx.onyxbot.slack.constants import KEEP_TO_YOURSELF_ACTION_ID\nfrom onyx.onyxbot.slack.constants import LIKE_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import SHOW_EVERYONE_ACTION_ID\nfrom onyx.onyxbot.slack.constants import VIEW_DOC_FEEDBACK_ID\nfrom onyx.onyxbot.slack.handlers.handle_message import (\n remove_scheduled_feedback_reminder,\n)\nfrom onyx.onyxbot.slack.handlers.handle_regular_answer import (\n handle_regular_answer,\n)\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.onyxbot.slack.utils import build_feedback_id\nfrom onyx.onyxbot.slack.utils import decompose_action_id\nfrom onyx.onyxbot.slack.utils import fetch_group_ids_from_names\nfrom onyx.onyxbot.slack.utils import fetch_slack_user_ids_from_emails\nfrom onyx.onyxbot.slack.utils import get_channel_name_from_id\nfrom onyx.onyxbot.slack.utils import get_feedback_visibility\nfrom onyx.onyxbot.slack.utils import read_slack_thread\nfrom onyx.onyxbot.slack.utils import respond_in_thread_or_channel\nfrom onyx.onyxbot.slack.utils import TenantSocketModeClient\nfrom onyx.onyxbot.slack.utils import update_emote_react\nfrom onyx.server.query_and_chat.models import ChatMessageDetail\nfrom onyx.server.query_and_chat.streaming_models import CitationInfo\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\ndef _convert_document_ids_to_citation_info(\n citation_dict: dict[int, str], top_documents: list[SavedSearchDoc]\n) -> list[CitationInfo]:\n citation_list_with_document_id = []\n # Build a set of valid document_ids from top_documents for validation\n valid_document_ids = {doc.document_id for doc in top_documents}\n\n for citation_num, document_id in citation_dict.items():\n if document_id is not None and document_id in valid_document_ids:\n citation_list_with_document_id.append(\n CitationInfo(\n citation_number=citation_num,\n document_id=document_id,\n )\n )\n return citation_list_with_document_id\n\n\ndef _build_citation_list(chat_message_detail: ChatMessageDetail) -> list[CitationInfo]:\n citation_dict = chat_message_detail.citations\n if citation_dict is None:\n return []\n else:\n top_documents = (\n chat_message_detail.context_docs if chat_message_detail.context_docs else []\n )\n citation_list = _convert_document_ids_to_citation_info(\n citation_dict, top_documents\n )\n return citation_list\n\n\ndef handle_doc_feedback_button(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n) -> None:\n if not (actions := req.payload.get(\"actions\")):\n logger.error(\"Missing actions. Unable to build the source feedback view\")\n return\n\n # Extracts the feedback_id coming from the 'source feedback' button\n # and generates a new one for the View, to keep track of the doc info\n query_event_id, doc_id, doc_rank = decompose_action_id(actions[0].get(\"value\"))\n external_id = build_feedback_id(query_event_id, doc_id, doc_rank)\n\n channel_id = req.payload[\"container\"][\"channel_id\"]\n thread_ts = req.payload[\"container\"].get(\"thread_ts\", None)\n\n data = View(\n type=\"modal\",\n callback_id=VIEW_DOC_FEEDBACK_ID,\n external_id=external_id,\n # We use the private metadata to keep track of the channel id and thread ts\n private_metadata=f\"{channel_id}_{thread_ts}\",\n title=\"Give Feedback\",\n blocks=[get_document_feedback_blocks()],\n submit=\"send\",\n close=\"cancel\",\n )\n\n client.web_client.views_open(\n trigger_id=req.payload[\"trigger_id\"], view=data.to_dict()\n )\n\n\ndef handle_generate_answer_button(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n) -> None:\n channel_id = req.payload[\"channel\"][\"id\"]\n channel_name = req.payload[\"channel\"][\"name\"]\n message_ts = req.payload[\"message\"][\"ts\"]\n thread_ts = req.payload[\"container\"].get(\"thread_ts\", None)\n user_id = req.payload[\"user\"][\"id\"]\n expert_info = expert_info_from_slack_id(user_id, client.web_client, user_cache={})\n email = expert_info.email if expert_info else None\n\n if not thread_ts:\n raise ValueError(\"Missing thread_ts in the payload\")\n\n thread_messages = read_slack_thread(\n tenant_id=client._tenant_id,\n channel=channel_id,\n thread=thread_ts,\n client=client.web_client,\n )\n # remove all assistant messages till we get to the last user message\n # we want the new answer to be generated off of the last \"question\" in\n # the thread\n for i in range(len(thread_messages) - 1, -1, -1):\n if thread_messages[i].role == MessageType.USER:\n break\n if thread_messages[i].role == MessageType.ASSISTANT:\n thread_messages.pop(i)\n\n # tell the user that we're working on it\n # Send an ephemeral message to the user that we're generating the answer\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=channel_id,\n receiver_ids=[user_id],\n text=\"I'm working on generating a full answer for you. This may take a moment...\",\n thread_ts=thread_ts,\n )\n\n with get_session_with_current_tenant() as db_session:\n slack_channel_config = get_slack_channel_config_for_bot_and_channel(\n db_session=db_session,\n slack_bot_id=client.slack_bot_id,\n channel_name=channel_name,\n )\n\n handle_regular_answer(\n message_info=SlackMessageInfo(\n thread_messages=thread_messages,\n channel_to_respond=channel_id,\n msg_to_respond=cast(str, message_ts or thread_ts),\n thread_to_respond=cast(str, thread_ts or message_ts),\n sender_id=user_id or None,\n email=email or None,\n bypass_filters=True,\n is_slash_command=False,\n is_bot_dm=False,\n ),\n slack_channel_config=slack_channel_config,\n receiver_ids=None,\n client=client.web_client,\n channel=channel_id,\n logger=logger,\n feedback_reminder_id=None,\n )\n\n\ndef handle_publish_ephemeral_message_button(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n action_id: str,\n) -> None:\n \"\"\"\n This function handles the Share with Everyone/Keep for Yourself buttons\n for ephemeral messages.\n \"\"\"\n channel_id = req.payload[\"channel\"][\"id\"]\n ephemeral_message_ts = req.payload[\"container\"][\"message_ts\"]\n\n slack_sender_id = req.payload[\"user\"][\"id\"]\n response_url = req.payload[\"response_url\"]\n webhook = WebhookClient(url=response_url)\n\n # The additional data required that was added to buttons.\n # Specifically, this contains the message_info, channel_conf information\n # and some additional attributes.\n value_dict = json.loads(req.payload[\"actions\"][0][\"value\"])\n\n original_question_ts = value_dict.get(\"original_question_ts\")\n if not original_question_ts:\n raise ValueError(\"Missing original_question_ts in the payload\")\n if not ephemeral_message_ts:\n raise ValueError(\"Missing ephemeral_message_ts in the payload\")\n\n feedback_reminder_id = value_dict.get(\"feedback_reminder_id\")\n\n slack_message_info = SlackMessageInfo(**value_dict[\"message_info\"])\n channel_conf = value_dict.get(\"channel_conf\")\n\n user_email = value_dict.get(\"message_info\", {}).get(\"email\")\n\n chat_message_id = value_dict.get(\"chat_message_id\")\n\n # Obtain onyx_user and chat_message information\n if not chat_message_id:\n raise ValueError(\"Missing chat_message_id in the payload\")\n\n with get_session_with_current_tenant() as db_session:\n onyx_user = get_user_by_email(user_email, db_session)\n if not onyx_user:\n raise ValueError(\"Cannot determine onyx_user_id from email in payload\")\n try:\n chat_message = get_chat_message(chat_message_id, onyx_user.id, db_session)\n except ValueError:\n chat_message = get_chat_message(\n chat_message_id, None, db_session\n ) # is this good idea?\n except Exception as e:\n logger.error(f\"Failed to get chat message: {e}\")\n raise e\n\n chat_message_detail = translate_db_message_to_chat_message_detail(chat_message)\n\n # construct the proper citation format and then the answer in the suitable format\n # we need to construct the blocks.\n citation_list = _build_citation_list(chat_message_detail)\n\n if chat_message_detail.context_docs:\n top_documents: list[SearchDoc] = [\n SearchDoc.from_saved_search_doc(doc)\n for doc in chat_message_detail.context_docs\n ]\n else:\n top_documents = []\n\n onyx_bot_answer = ChatBasicResponse(\n answer=chat_message_detail.message,\n answer_citationless=remove_answer_citations(chat_message_detail.message),\n top_documents=top_documents,\n message_id=chat_message_id,\n error_msg=None,\n citation_info=citation_list,\n )\n\n # Note: we need to use the webhook and the respond_url to update/delete ephemeral messages\n if action_id == SHOW_EVERYONE_ACTION_ID:\n # Convert to non-ephemeral message in thread\n try:\n webhook.send(\n response_type=\"ephemeral\",\n text=\"\",\n blocks=[],\n replace_original=True,\n delete_original=True,\n )\n except Exception as e:\n logger.error(f\"Failed to send webhook: {e}\")\n\n # remove handling of empheremal block and add AI feedback.\n all_blocks = build_slack_response_blocks(\n answer=onyx_bot_answer,\n message_info=slack_message_info,\n channel_conf=channel_conf,\n feedback_reminder_id=feedback_reminder_id,\n skip_ai_feedback=False,\n offer_ephemeral_publication=False,\n skip_restated_question=True,\n )\n try:\n # Post in thread as non-ephemeral message\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=channel_id,\n receiver_ids=None, # If respond_member_group_list is set, send to them. TODO: check!\n text=\"Hello! Onyx has some results for you!\",\n blocks=all_blocks,\n thread_ts=original_question_ts,\n # don't unfurl, since otherwise we will have 5+ previews which makes the message very long\n unfurl=False,\n send_as_ephemeral=False,\n )\n except Exception as e:\n logger.error(f\"Failed to publish ephemeral message: {e}\")\n raise e\n\n elif action_id == KEEP_TO_YOURSELF_ACTION_ID:\n # Keep as ephemeral message in channel or thread, but remove the publish button and add feedback button\n\n changed_blocks = build_slack_response_blocks(\n answer=onyx_bot_answer,\n message_info=slack_message_info,\n channel_conf=channel_conf,\n feedback_reminder_id=feedback_reminder_id,\n skip_ai_feedback=False,\n offer_ephemeral_publication=False,\n skip_restated_question=True,\n )\n\n try:\n if slack_message_info.thread_to_respond is not None:\n # There seems to be a bug in slack where an update within the thread\n # actually leads to the update to be posted in the channel. Therefore,\n # for now we delete the original ephemeral message and post a new one\n # if the ephemeral message is in a thread.\n webhook.send(\n response_type=\"ephemeral\",\n text=\"\",\n blocks=[],\n replace_original=True,\n delete_original=True,\n )\n\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=channel_id,\n receiver_ids=[slack_sender_id],\n text=\"Your personal response, sent as an ephemeral message.\",\n blocks=changed_blocks,\n thread_ts=original_question_ts,\n # don't unfurl, since otherwise we will have 5+ previews which makes the message very long\n unfurl=False,\n send_as_ephemeral=True,\n )\n else:\n # This works fine if the ephemeral message is in the channel\n webhook.send(\n response_type=\"ephemeral\",\n text=\"Your personal response, sent as an ephemeral message.\",\n blocks=changed_blocks,\n replace_original=True,\n delete_original=False,\n )\n except Exception as e:\n logger.error(f\"Failed to send webhook: {e}\")\n\n\ndef handle_slack_feedback(\n feedback_id: str,\n feedback_type: str,\n feedback_msg_reminder: str,\n client: WebClient,\n user_id_to_post_confirmation: str,\n channel_id_to_post_confirmation: str,\n thread_ts_to_post_confirmation: str,\n) -> None:\n message_id, doc_id, doc_rank = decompose_action_id(feedback_id)\n\n # Get Onyx user from Slack ID\n expert_info = expert_info_from_slack_id(\n user_id_to_post_confirmation, client, user_cache={}\n )\n email = expert_info.email if expert_info else None\n\n with get_session_with_current_tenant() as db_session:\n onyx_user = get_user_by_email(email, db_session) if email else None\n if feedback_type in [LIKE_BLOCK_ACTION_ID, DISLIKE_BLOCK_ACTION_ID]:\n create_chat_message_feedback(\n is_positive=feedback_type == LIKE_BLOCK_ACTION_ID,\n feedback_text=\"\",\n chat_message_id=message_id,\n user_id=onyx_user.id if onyx_user else None,\n db_session=db_session,\n )\n remove_scheduled_feedback_reminder(\n client=client,\n channel=user_id_to_post_confirmation,\n msg_id=feedback_msg_reminder,\n )\n elif feedback_type in [\n SearchFeedbackType.ENDORSE.value,\n SearchFeedbackType.REJECT.value,\n SearchFeedbackType.HIDE.value,\n ]:\n if doc_id is None or doc_rank is None:\n raise ValueError(\"Missing information for Document Feedback\")\n\n if feedback_type == SearchFeedbackType.ENDORSE.value:\n feedback = SearchFeedbackType.ENDORSE\n elif feedback_type == SearchFeedbackType.REJECT.value:\n feedback = SearchFeedbackType.REJECT\n else:\n feedback = SearchFeedbackType.HIDE\n\n create_doc_retrieval_feedback(\n message_id=message_id,\n document_id=doc_id,\n document_rank=doc_rank,\n db_session=db_session,\n clicked=False, # Not tracking this for Slack\n feedback=feedback,\n )\n else:\n logger.error(f\"Feedback type '{feedback_type}' not supported\")\n\n if get_feedback_visibility() == FeedbackVisibility.PRIVATE or feedback_type not in [\n LIKE_BLOCK_ACTION_ID,\n DISLIKE_BLOCK_ACTION_ID,\n ]:\n client.chat_postEphemeral(\n channel=channel_id_to_post_confirmation,\n user=user_id_to_post_confirmation,\n thread_ts=thread_ts_to_post_confirmation,\n text=\"Thanks for your feedback!\",\n )\n else:\n feedback_response_txt = (\n \"liked\" if feedback_type == LIKE_BLOCK_ACTION_ID else \"disliked\"\n )\n\n if get_feedback_visibility() == FeedbackVisibility.ANONYMOUS:\n msg = f\"A user has {feedback_response_txt} the AI Answer\"\n else:\n msg = f\"<@{user_id_to_post_confirmation}> has {feedback_response_txt} the AI Answer\"\n\n respond_in_thread_or_channel(\n client=client,\n channel=channel_id_to_post_confirmation,\n text=msg,\n thread_ts=thread_ts_to_post_confirmation,\n unfurl=False,\n )\n\n\ndef handle_followup_button(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n) -> None:\n action_id = None\n if actions := req.payload.get(\"actions\"):\n action = cast(dict[str, Any], actions[0])\n action_id = cast(str, action.get(\"block_id\"))\n\n channel_id = req.payload[\"container\"][\"channel_id\"]\n thread_ts = req.payload[\"container\"].get(\"thread_ts\", None)\n\n update_emote_react(\n emoji=ONYX_BOT_FOLLOWUP_EMOJI,\n channel=channel_id,\n message_ts=thread_ts,\n remove=False,\n client=client.web_client,\n )\n\n tag_ids: list[str] = []\n group_ids: list[str] = []\n with get_session_with_current_tenant() as db_session:\n channel_name, is_dm = get_channel_name_from_id(\n client=client.web_client, channel_id=channel_id\n )\n slack_channel_config = get_slack_channel_config_for_bot_and_channel(\n db_session=db_session,\n slack_bot_id=client.slack_bot_id,\n channel_name=channel_name,\n )\n if slack_channel_config:\n tag_names = slack_channel_config.channel_config.get(\"follow_up_tags\")\n remaining = None\n if tag_names:\n tag_ids, remaining = fetch_slack_user_ids_from_emails(\n tag_names, client.web_client\n )\n if remaining:\n group_ids, _ = fetch_group_ids_from_names(remaining, client.web_client)\n\n blocks = build_follow_up_resolved_blocks(tag_ids=tag_ids, group_ids=group_ids)\n\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=channel_id,\n text=\"Received your request for more help\",\n blocks=blocks,\n thread_ts=thread_ts,\n unfurl=False,\n )\n\n if action_id is not None:\n message_id, _, _ = decompose_action_id(action_id)\n\n create_chat_message_feedback(\n is_positive=None,\n feedback_text=\"\",\n chat_message_id=message_id,\n user_id=None, # no \"user\" for Slack bot for now\n db_session=db_session,\n required_followup=True,\n )\n\n\ndef get_clicker_name(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n) -> str:\n clicker_name = req.payload.get(\"user\", {}).get(\"name\", \"Someone\")\n clicker_real_name = None\n try:\n clicker = client.web_client.users_info(user=req.payload[\"user\"][\"id\"])\n clicker_real_name = (\n cast(dict, clicker.data).get(\"user\", {}).get(\"profile\", {}).get(\"real_name\")\n )\n except Exception:\n # Likely a scope issue\n pass\n\n if clicker_real_name:\n clicker_name = clicker_real_name\n\n return clicker_name\n\n\ndef handle_followup_resolved_button(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n immediate: bool = False,\n) -> None:\n channel_id = req.payload[\"container\"][\"channel_id\"]\n message_ts = req.payload[\"container\"][\"message_ts\"]\n thread_ts = req.payload[\"container\"].get(\"thread_ts\", None)\n\n clicker_name = get_clicker_name(req, client)\n\n update_emote_react(\n emoji=ONYX_BOT_FOLLOWUP_EMOJI,\n channel=channel_id,\n message_ts=thread_ts,\n remove=True,\n client=client.web_client,\n )\n\n # Delete the message with the option to mark resolved\n if not immediate:\n response = client.web_client.chat_delete(\n channel=channel_id,\n ts=message_ts,\n )\n\n if not response.get(\"ok\"):\n logger.error(\"Unable to delete message for resolved\")\n\n if immediate:\n msg_text = f\"{clicker_name} has marked this question as resolved!\"\n else:\n msg_text = (\n f\"{clicker_name} has marked this question as resolved! \"\n f'\\n\\n You can always click the \"I need more help button\" to let the team '\n f\"know that your problem still needs attention.\"\n )\n\n resolved_block = SectionBlock(text=msg_text)\n\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=channel_id,\n text=\"Your request for help as been addressed!\",\n blocks=[resolved_block],\n thread_ts=thread_ts,\n unfurl=False,\n )\n" }, { "path": "backend/onyx/onyxbot/slack/handlers/handle_message.py", "content": "import datetime\n\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\n\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_FEEDBACK_REMINDER\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_REACT_EMOJI\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.db.user_preferences import activate_user\nfrom onyx.db.users import add_slack_user_if_not_exists\nfrom onyx.db.users import get_user_by_email\nfrom onyx.onyxbot.slack.blocks import get_feedback_reminder_blocks\nfrom onyx.onyxbot.slack.handlers.handle_regular_answer import (\n handle_regular_answer,\n)\nfrom onyx.onyxbot.slack.handlers.handle_standard_answers import (\n handle_standard_answers,\n)\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.onyxbot.slack.utils import fetch_slack_user_ids_from_emails\nfrom onyx.onyxbot.slack.utils import fetch_user_ids_from_groups\nfrom onyx.onyxbot.slack.utils import respond_in_thread_or_channel\nfrom onyx.onyxbot.slack.utils import slack_usage_report\nfrom onyx.onyxbot.slack.utils import update_emote_react\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom shared_configs.configs import SLACK_CHANNEL_ID\n\nlogger_base = setup_logger()\n\n\ndef send_msg_ack_to_user(details: SlackMessageInfo, client: WebClient) -> None:\n if details.is_slash_command and details.sender_id:\n respond_in_thread_or_channel(\n client=client,\n channel=details.channel_to_respond,\n thread_ts=details.msg_to_respond,\n receiver_ids=[details.sender_id],\n text=\"Hi, we're evaluating your query :face_with_monocle:\",\n )\n return\n\n update_emote_react(\n emoji=ONYX_BOT_REACT_EMOJI,\n channel=details.channel_to_respond,\n message_ts=details.msg_to_respond,\n remove=False,\n client=client,\n )\n\n\ndef schedule_feedback_reminder(\n details: SlackMessageInfo, include_followup: bool, client: WebClient\n) -> str | None:\n logger = setup_logger(extra={SLACK_CHANNEL_ID: details.channel_to_respond})\n\n if not ONYX_BOT_FEEDBACK_REMINDER:\n logger.info(\"Scheduled feedback reminder disabled...\")\n return None\n\n try:\n permalink = client.chat_getPermalink(\n channel=details.channel_to_respond,\n message_ts=details.msg_to_respond, # type:ignore\n )\n except SlackApiError as e:\n logger.error(f\"Unable to generate the feedback reminder permalink: {e}\")\n return None\n\n now = datetime.datetime.now()\n future = now + datetime.timedelta(minutes=ONYX_BOT_FEEDBACK_REMINDER)\n\n try:\n response = client.chat_scheduleMessage(\n channel=details.sender_id, # type:ignore\n post_at=int(future.timestamp()),\n blocks=[\n get_feedback_reminder_blocks(\n thread_link=permalink.data[\"permalink\"], # type:ignore\n include_followup=include_followup,\n )\n ],\n text=\"\",\n )\n logger.info(\"Scheduled feedback reminder configured\")\n return response.data[\"scheduled_message_id\"] # type:ignore\n except SlackApiError as e:\n logger.error(f\"Unable to generate the feedback reminder message: {e}\")\n return None\n\n\ndef remove_scheduled_feedback_reminder(\n client: WebClient, channel: str | None, msg_id: str\n) -> None:\n logger = setup_logger(extra={SLACK_CHANNEL_ID: channel})\n\n try:\n client.chat_deleteScheduledMessage(\n channel=channel, # type:ignore\n scheduled_message_id=msg_id,\n )\n logger.info(\"Scheduled feedback reminder deleted\")\n except SlackApiError as e:\n if e.response[\"error\"] == \"invalid_scheduled_message_id\":\n logger.info(\n \"Unable to delete the scheduled message. It must have already been posted\"\n )\n\n\ndef handle_message(\n message_info: SlackMessageInfo,\n slack_channel_config: SlackChannelConfig,\n client: WebClient,\n feedback_reminder_id: str | None,\n) -> bool:\n \"\"\"Potentially respond to the user message depending on filters and if an answer was generated\n\n Returns True if need to respond with an additional message to the user(s) after this\n function is finished. True indicates an unexpected failure that needs to be communicated\n Query thrown out by filters due to config does not count as a failure that should be notified\n Onyx failing to answer/retrieve docs does count and should be notified\n \"\"\"\n channel = message_info.channel_to_respond\n\n logger = setup_logger(extra={SLACK_CHANNEL_ID: channel})\n\n messages = message_info.thread_messages\n sender_id = message_info.sender_id\n bypass_filters = message_info.bypass_filters\n is_slash_command = message_info.is_slash_command\n is_bot_dm = message_info.is_bot_dm\n\n action = \"slack_message\"\n if is_slash_command:\n action = \"slack_slash_message\"\n elif bypass_filters:\n action = \"slack_tag_message\"\n elif is_bot_dm:\n action = \"slack_dm_message\"\n slack_usage_report(action=action, sender_id=sender_id, client=client)\n\n document_set_names: list[str] | None = None\n persona = slack_channel_config.persona if slack_channel_config else None\n if persona:\n document_set_names = [\n document_set.name for document_set in persona.document_sets\n ]\n\n respond_tag_only = False\n respond_member_group_list = None\n\n channel_conf = None\n if slack_channel_config and slack_channel_config.channel_config:\n channel_conf = slack_channel_config.channel_config\n if not bypass_filters and \"answer_filters\" in channel_conf:\n if (\n \"questionmark_prefilter\" in channel_conf[\"answer_filters\"]\n and \"?\" not in messages[-1].message\n ):\n logger.info(\n \"Skipping message since it does not contain a question mark\"\n )\n return False\n\n logger.info(\n \"Found slack bot config for channel. Restricting bot to use document \"\n f\"sets: {document_set_names}, \"\n f\"validity checks enabled: {channel_conf.get('answer_filters', 'NA')}\"\n )\n\n respond_tag_only = channel_conf.get(\"respond_tag_only\") or False\n respond_member_group_list = channel_conf.get(\"respond_member_group_list\", None)\n\n # Only default config can be disabled.\n # If channel config is disabled, bot should not respond to this message (including DMs)\n if slack_channel_config.channel_config.get(\"disabled\"):\n logger.info(\"Skipping message: OnyxBot is disabled for this channel\")\n return False\n\n # If bot should only respond to tags and is not tagged nor in a DM, skip message\n if respond_tag_only and not bypass_filters and not is_bot_dm:\n logger.info(\"Skipping message: OnyxBot only responds to tags in this channel\")\n return False\n\n # List of user id to send message to, if None, send to everyone in channel\n send_to: list[str] | None = None\n missing_users: list[str] | None = None\n if respond_member_group_list:\n send_to, missing_ids = fetch_slack_user_ids_from_emails(\n respond_member_group_list, client\n )\n\n user_ids, missing_users = fetch_user_ids_from_groups(missing_ids, client)\n send_to = list(set(send_to + user_ids)) if send_to else user_ids\n\n if missing_users:\n logger.warning(f\"Failed to find these users/groups: {missing_users}\")\n\n # If configured to respond to team members only, then cannot be used with a /OnyxBot command\n # which would just respond to the sender\n if send_to and is_slash_command:\n if sender_id:\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n receiver_ids=[sender_id],\n text=\"The OnyxBot slash command is not enabled for this channel\",\n thread_ts=None,\n )\n\n try:\n send_msg_ack_to_user(message_info, client)\n except SlackApiError as e:\n logger.error(f\"Was not able to react to user message due to: {e}\")\n\n with get_session_with_current_tenant() as db_session:\n if message_info.email:\n existing_user = get_user_by_email(message_info.email, db_session)\n if existing_user is None:\n # New user — check seat availability before creating\n check_seat_fn = fetch_ee_implementation_or_noop(\n \"onyx.db.license\",\n \"check_seat_availability\",\n None,\n )\n # noop returns None when called; real function returns SeatAvailabilityResult\n seat_result = check_seat_fn(db_session=db_session)\n if seat_result is not None and not seat_result.available:\n logger.info(\n f\"Blocked new Slack user {message_info.email}: {seat_result.error_message}\"\n )\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n thread_ts=message_info.msg_to_respond,\n text=(\n \"We weren't able to respond because your organization \"\n \"has reached its user seat limit. Since this is your \"\n \"first time interacting with the bot, a new account \"\n \"could not be created for you. Please contact your \"\n \"Onyx administrator to add more seats.\"\n ),\n )\n return False\n\n elif (\n not existing_user.is_active\n and existing_user.account_type == AccountType.BOT\n ):\n check_seat_fn = fetch_ee_implementation_or_noop(\n \"onyx.db.license\",\n \"check_seat_availability\",\n None,\n )\n seat_result = check_seat_fn(db_session=db_session)\n if seat_result is not None and not seat_result.available:\n logger.info(\n f\"Blocked inactive Slack user {message_info.email}: {seat_result.error_message}\"\n )\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n thread_ts=message_info.msg_to_respond,\n text=(\n \"We weren't able to respond because your organization \"\n \"has reached its user seat limit. Your account is \"\n \"currently deactivated and cannot be reactivated \"\n \"until more seats are available. Please contact \"\n \"your Onyx administrator.\"\n ),\n )\n return False\n\n activate_user(existing_user, db_session)\n invalidate_license_cache_fn = fetch_ee_implementation_or_noop(\n \"onyx.db.license\",\n \"invalidate_license_cache\",\n None,\n )\n invalidate_license_cache_fn()\n logger.info(f\"Reactivated inactive Slack user {message_info.email}\")\n\n add_slack_user_if_not_exists(db_session, message_info.email)\n\n # first check if we need to respond with a standard answer\n # standard answers should be published in a thread\n used_standard_answer = handle_standard_answers(\n message_info=message_info,\n receiver_ids=send_to,\n slack_channel_config=slack_channel_config,\n logger=logger,\n client=client,\n db_session=db_session,\n )\n if used_standard_answer:\n return False\n\n # if no standard answer applies, try a regular answer\n issue_with_regular_answer = handle_regular_answer(\n message_info=message_info,\n slack_channel_config=slack_channel_config,\n receiver_ids=send_to,\n client=client,\n channel=channel,\n logger=logger,\n feedback_reminder_id=feedback_reminder_id,\n )\n return issue_with_regular_answer\n" }, { "path": "backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py", "content": "import functools\nfrom collections.abc import Callable\nfrom typing import Any\nfrom typing import Optional\nfrom typing import TypeVar\n\nfrom retry import retry\nfrom slack_sdk import WebClient\n\nfrom onyx.auth.users import get_anonymous_user\nfrom onyx.chat.models import ChatBasicResponse\nfrom onyx.chat.process_message import gather_stream\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.configs.constants import DEFAULT_PERSONA_ID\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_DISABLE_DOCS_ONLY_ANSWER\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_DISPLAY_ERROR_MSGS\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_NUM_RETRIES\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_REACT_EMOJI\nfrom onyx.context.search.models import BaseFilters\nfrom onyx.context.search.models import Tag\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.db.models import User\nfrom onyx.db.persona import get_persona_by_id\nfrom onyx.db.users import get_user_by_email\nfrom onyx.onyxbot.slack.blocks import build_slack_response_blocks\nfrom onyx.onyxbot.slack.constants import SLACK_CHANNEL_REF_PATTERN\nfrom onyx.onyxbot.slack.handlers.utils import send_team_member_message\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.onyxbot.slack.models import ThreadMessage\nfrom onyx.onyxbot.slack.utils import get_channel_from_id\nfrom onyx.onyxbot.slack.utils import get_channel_name_from_id\nfrom onyx.onyxbot.slack.utils import respond_in_thread_or_channel\nfrom onyx.onyxbot.slack.utils import SlackRateLimiter\nfrom onyx.onyxbot.slack.utils import update_emote_react\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.models import MessageOrigin\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.utils.logger import OnyxLoggingAdapter\n\nsrl = SlackRateLimiter()\n\nRT = TypeVar(\"RT\") # return type\n\n\ndef resolve_channel_references(\n message: str,\n client: WebClient,\n logger: OnyxLoggingAdapter,\n) -> tuple[str, list[Tag]]:\n \"\"\"Parse Slack channel references from a message, resolve IDs to names,\n replace the raw markup with readable #channel-name, and return channel tags\n for search filtering.\"\"\"\n tags: list[Tag] = []\n channel_matches = SLACK_CHANNEL_REF_PATTERN.findall(message)\n seen_channel_ids: set[str] = set()\n\n for channel_id, channel_name_from_markup in channel_matches:\n if channel_id in seen_channel_ids:\n continue\n seen_channel_ids.add(channel_id)\n\n channel_name = channel_name_from_markup or None\n\n if not channel_name:\n try:\n channel_info = get_channel_from_id(client=client, channel_id=channel_id)\n channel_name = channel_info.get(\"name\") or None\n except Exception:\n logger.warning(f\"Failed to resolve channel name for ID: {channel_id}\")\n\n if not channel_name:\n continue\n\n # Replace raw Slack markup with readable channel name\n if channel_name_from_markup:\n message = message.replace(\n f\"<#{channel_id}|{channel_name_from_markup}>\",\n f\"#{channel_name}\",\n )\n else:\n message = message.replace(\n f\"<#{channel_id}>\",\n f\"#{channel_name}\",\n )\n tags.append(Tag(tag_key=\"Channel\", tag_value=channel_name))\n\n return message, tags\n\n\ndef rate_limits(\n client: WebClient, channel: str, thread_ts: Optional[str]\n) -> Callable[[Callable[..., RT]], Callable[..., RT]]:\n def decorator(func: Callable[..., RT]) -> Callable[..., RT]:\n @functools.wraps(func)\n def wrapper(*args: Any, **kwargs: Any) -> RT:\n if not srl.is_available():\n func_randid, position = srl.init_waiter()\n srl.notify(client, channel, position, thread_ts)\n while not srl.is_available():\n srl.waiter(func_randid)\n srl.acquire_slot()\n return func(*args, **kwargs)\n\n return wrapper\n\n return decorator\n\n\ndef build_slack_context_str(\n messages: list[ThreadMessage], channel_name: str | None\n) -> str | None:\n if not messages:\n return None\n\n if channel_name:\n slack_context_str = f\"The following is a thread in Slack in channel {channel_name}:\\n====================\\n\"\n else:\n slack_context_str = (\n \"The following is a thread from Slack:\\n====================\\n\"\n )\n\n message_strs: list[str] = []\n for message in messages:\n if message.role == MessageType.USER:\n message_text = f\"{message.sender or 'Unknown User'}:\\n{message.message}\"\n elif message.role == MessageType.ASSISTANT:\n message_text = f\"AI:\\n{message.message}\"\n else:\n message_text = f\"{message.role.value.upper()}:\\n{message.message}\"\n message_strs.append(message_text)\n\n return slack_context_str + \"\\n\\n\".join(message_strs)\n\n\ndef handle_regular_answer(\n message_info: SlackMessageInfo,\n slack_channel_config: SlackChannelConfig,\n receiver_ids: list[str] | None,\n client: WebClient,\n channel: str,\n logger: OnyxLoggingAdapter,\n feedback_reminder_id: str | None,\n num_retries: int = ONYX_BOT_NUM_RETRIES,\n should_respond_with_error_msgs: bool = ONYX_BOT_DISPLAY_ERROR_MSGS,\n disable_docs_only_answer: bool = ONYX_BOT_DISABLE_DOCS_ONLY_ANSWER,\n) -> bool:\n channel_conf = slack_channel_config.channel_config\n\n messages = message_info.thread_messages\n\n message_ts_to_respond_to = message_info.msg_to_respond\n is_slash_command = message_info.is_slash_command\n\n # Capture whether response mode for channel is ephemeral. Even if the channel is set\n # to respond with an ephemeral message, we still send as non-ephemeral if\n # the message is a dm with the Onyx bot.\n send_as_ephemeral = (\n slack_channel_config.channel_config.get(\"is_ephemeral\", False)\n or message_info.is_slash_command\n ) and not message_info.is_bot_dm\n\n # If the channel is configured to respond with an ephemeral message,\n # or the message is a dm to the Onyx bot, we should use the proper onyx user from the email.\n # This will make documents privately accessible to the user available to Onyx Bot answers.\n # Otherwise - if not ephemeral or DM to Onyx Bot - we use anonymous user to restrict\n # to public docs.\n\n if message_info.email:\n with get_session_with_current_tenant() as db_session:\n found_user = get_user_by_email(message_info.email, db_session)\n user = found_user if found_user else get_anonymous_user()\n else:\n user = get_anonymous_user()\n\n target_thread_ts = (\n None\n if send_as_ephemeral and len(message_info.thread_messages) < 2\n else message_ts_to_respond_to\n )\n target_receiver_ids = (\n [message_info.sender_id]\n if message_info.sender_id and send_as_ephemeral\n else receiver_ids\n )\n\n document_set_names: list[str] | None = None\n # If no persona is specified, use the default search based persona\n # This way slack flow always has a persona\n persona = slack_channel_config.persona\n if not persona:\n logger.warning(\"No persona found for channel config, using default persona\")\n with get_session_with_current_tenant() as db_session:\n persona = get_persona_by_id(DEFAULT_PERSONA_ID, user, db_session)\n document_set_names = [\n document_set.name for document_set in persona.document_sets\n ]\n else:\n logger.info(f\"Using persona {persona.name} for channel config\")\n document_set_names = [\n document_set.name for document_set in persona.document_sets\n ]\n\n user_message = messages[-1]\n history_messages = messages[:-1]\n\n # Resolve any <#CHANNEL_ID> references in the user message to readable\n # channel names and extract channel tags for search filtering\n resolved_message, channel_tags = resolve_channel_references(\n message=user_message.message,\n client=client,\n logger=logger,\n )\n\n user_message = ThreadMessage(\n message=resolved_message,\n sender=user_message.sender,\n role=user_message.role,\n )\n\n channel_name, _ = get_channel_name_from_id(\n client=client,\n channel_id=channel,\n )\n\n # NOTE: only the message history will contain the person asking. This is likely\n # fine since the most common use case for this info is when referring to a user\n # who previously posted in the thread.\n slack_context_str = build_slack_context_str(history_messages, channel_name)\n\n if not message_ts_to_respond_to and not is_slash_command:\n # if the message is not \"/onyx\" command, then it should have a message ts to respond to\n raise RuntimeError(\n \"No message timestamp to respond to in `handle_message`. This should never happen.\"\n )\n\n @retry(\n tries=num_retries,\n delay=0.25,\n backoff=2,\n )\n @rate_limits(client=client, channel=channel, thread_ts=message_ts_to_respond_to)\n def _get_slack_answer(\n new_message_request: SendMessageRequest,\n slack_context_str: str | None,\n onyx_user: User,\n ) -> ChatBasicResponse:\n with get_session_with_current_tenant() as db_session:\n packets = handle_stream_message_objects(\n new_msg_req=new_message_request,\n user=onyx_user,\n db_session=db_session,\n bypass_acl=False,\n additional_context=slack_context_str,\n slack_context=message_info.slack_context,\n )\n answer = gather_stream(packets)\n\n if answer.error_msg:\n raise RuntimeError(answer.error_msg)\n\n return answer\n\n try:\n # By leaving time_cutoff and favor_recent as None, and setting enable_auto_detect_filters\n # it allows the slack flow to extract out filters from the user query\n filters = BaseFilters(\n source_type=None,\n document_set=document_set_names,\n time_cutoff=None,\n tags=channel_tags if channel_tags else None,\n )\n\n new_message_request = SendMessageRequest(\n message=user_message.message,\n allowed_tool_ids=None,\n forced_tool_id=None,\n file_descriptors=[],\n internal_search_filters=filters,\n deep_research=False,\n origin=MessageOrigin.SLACKBOT,\n chat_session_info=ChatSessionCreationRequest(\n persona_id=persona.id,\n ),\n )\n\n # if it's a DM or ephemeral message, answer based on private documents.\n # otherwise, answer based on public documents ONLY as to not leak information.\n can_search_over_private_docs = message_info.is_bot_dm or send_as_ephemeral\n answer = _get_slack_answer(\n new_message_request=new_message_request,\n onyx_user=user if can_search_over_private_docs else get_anonymous_user(),\n slack_context_str=slack_context_str,\n )\n\n # If a channel filter was applied but no results were found, override\n # the LLM response to avoid hallucinated answers about unindexed channels\n if channel_tags and not answer.citation_info and not answer.top_documents:\n channel_names = \", \".join(f\"#{tag.tag_value}\" for tag in channel_tags)\n answer.answer = (\n f\"No indexed data found for {channel_names}. \"\n \"This channel may not be indexed, or there may be no messages \"\n \"matching your query within it.\"\n )\n\n except Exception as e:\n logger.exception(\n f\"Unable to process message - did not successfully answer in {num_retries} attempts\"\n )\n # Optionally, respond in thread with the error message, Used primarily\n # for debugging purposes\n if should_respond_with_error_msgs:\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n receiver_ids=target_receiver_ids,\n text=f\"Encountered exception when trying to answer: \\n\\n```{e}```\",\n thread_ts=target_thread_ts,\n send_as_ephemeral=send_as_ephemeral,\n )\n\n # In case of failures, don't keep the reaction there permanently\n update_emote_react(\n emoji=ONYX_BOT_REACT_EMOJI,\n channel=message_info.channel_to_respond,\n message_ts=message_info.msg_to_respond,\n remove=True,\n client=client,\n )\n\n return True\n\n # Got an answer at this point, can remove reaction and give results\n if not is_slash_command: # Slash commands don't have reactions\n update_emote_react(\n emoji=ONYX_BOT_REACT_EMOJI,\n channel=message_info.channel_to_respond,\n message_ts=message_info.msg_to_respond,\n remove=True,\n client=client,\n )\n\n if not answer.answer and disable_docs_only_answer:\n logger.notice(\n \"Unable to find answer - not responding since the `ONYX_BOT_DISABLE_DOCS_ONLY_ANSWER` env variable is set\"\n )\n return True\n\n only_respond_if_citations = (\n channel_conf\n and \"well_answered_postfilter\" in channel_conf.get(\"answer_filters\", [])\n )\n\n if (\n only_respond_if_citations\n and not answer.citation_info\n and not message_info.bypass_filters\n and not channel_tags\n ):\n logger.error(\n f\"Unable to find citations to answer: '{answer.answer}' - not answering!\"\n )\n # Optionally, respond in thread with the error message\n # Used primarily for debugging purposes\n if should_respond_with_error_msgs:\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n receiver_ids=target_receiver_ids,\n text=\"Found no citations or quotes when trying to answer.\",\n thread_ts=target_thread_ts,\n send_as_ephemeral=send_as_ephemeral,\n )\n return True\n\n if (\n send_as_ephemeral\n and target_receiver_ids is not None\n and len(target_receiver_ids) == 1\n ):\n offer_ephemeral_publication = True\n skip_ai_feedback = True\n else:\n offer_ephemeral_publication = False\n skip_ai_feedback = False\n\n all_blocks = build_slack_response_blocks(\n message_info=message_info,\n answer=answer,\n channel_conf=channel_conf,\n feedback_reminder_id=feedback_reminder_id,\n offer_ephemeral_publication=offer_ephemeral_publication,\n skip_ai_feedback=skip_ai_feedback,\n )\n\n # NOTE(rkuo): Slack has a maximum block list size of 50.\n # we should modify build_slack_response_blocks to respect the max\n # but enforcing the hard limit here is the last resort.\n all_blocks = all_blocks[:50]\n\n try:\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n receiver_ids=target_receiver_ids,\n text=\"Hello! Onyx has some results for you!\",\n blocks=all_blocks,\n thread_ts=target_thread_ts,\n # don't unfurl, since otherwise we will have 5+ previews which makes the message very long\n unfurl=False,\n send_as_ephemeral=send_as_ephemeral,\n )\n\n # For DM (ephemeral message), we need to create a thread via a normal message so the user can see\n # the ephemeral message. This also will give the user a notification which ephemeral message does not.\n # if there is no message_ts_to_respond_to, and we have made it this far, then this is a /onyx message\n # so we shouldn't send_team_member_message\n if (\n target_receiver_ids\n and message_ts_to_respond_to is not None\n and not send_as_ephemeral\n and target_thread_ts is not None\n ):\n send_team_member_message(\n client=client,\n channel=channel,\n thread_ts=target_thread_ts,\n receiver_ids=target_receiver_ids,\n send_as_ephemeral=send_as_ephemeral,\n )\n\n return False\n\n except Exception:\n logger.exception(\n f\"Unable to process message - could not respond in slack in {num_retries} attempts\"\n )\n return True\n" }, { "path": "backend/onyx/onyxbot/slack/handlers/handle_standard_answers.py", "content": "from slack_sdk import WebClient\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.utils.logger import OnyxLoggingAdapter\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\nlogger = setup_logger()\n\n\ndef handle_standard_answers(\n message_info: SlackMessageInfo,\n receiver_ids: list[str] | None,\n slack_channel_config: SlackChannelConfig,\n logger: OnyxLoggingAdapter,\n client: WebClient,\n db_session: Session,\n) -> bool:\n \"\"\"Returns whether one or more Standard Answer message blocks were\n emitted by the Slack bot\"\"\"\n versioned_handle_standard_answers = fetch_versioned_implementation(\n \"onyx.onyxbot.slack.handlers.handle_standard_answers\",\n \"_handle_standard_answers\",\n )\n return versioned_handle_standard_answers(\n message_info=message_info,\n receiver_ids=receiver_ids,\n slack_channel_config=slack_channel_config,\n logger=logger,\n client=client,\n db_session=db_session,\n )\n\n\ndef _handle_standard_answers(\n message_info: SlackMessageInfo, # noqa: ARG001\n receiver_ids: list[str] | None, # noqa: ARG001\n slack_channel_config: SlackChannelConfig, # noqa: ARG001\n logger: OnyxLoggingAdapter, # noqa: ARG001\n client: WebClient, # noqa: ARG001\n db_session: Session, # noqa: ARG001\n) -> bool:\n \"\"\"\n Standard Answers are a paid Enterprise Edition feature. This is the fallback\n function handling the case where EE features are not enabled.\n\n Always returns false i.e. since EE features are not enabled, we NEVER create any\n Slack message blocks.\n \"\"\"\n return False\n" }, { "path": "backend/onyx/onyxbot/slack/handlers/utils.py", "content": "from slack_sdk import WebClient\n\nfrom onyx.onyxbot.slack.utils import respond_in_thread_or_channel\n\n\ndef send_team_member_message(\n client: WebClient,\n channel: str,\n thread_ts: str,\n receiver_ids: list[str] | None = None, # noqa: ARG001\n send_as_ephemeral: bool = False,\n) -> None:\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n text=(\n \"👋 Hi, we've just gathered and forwarded the relevant \"\n + \"information to the team. They'll get back to you shortly!\"\n ),\n thread_ts=thread_ts,\n receiver_ids=None,\n send_as_ephemeral=send_as_ephemeral,\n )\n" }, { "path": "backend/onyx/onyxbot/slack/icons.py", "content": "from onyx.configs.constants import DocumentSource\n\n\ndef source_to_github_img_link(source: DocumentSource) -> str | None:\n # TODO: store these images somewhere better\n if source == DocumentSource.WEB.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/Web.png\"\n if source == DocumentSource.FILE.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/File.png\"\n if source == DocumentSource.GOOGLE_SITES.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/GoogleSites.png\"\n if source == DocumentSource.SLACK.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Slack.png\"\n if source == DocumentSource.GMAIL.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Gmail.png\"\n if source == DocumentSource.GOOGLE_DRIVE.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/GoogleDrive.png\"\n if source == DocumentSource.GITHUB.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Github.png\"\n if source == DocumentSource.GITLAB.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Gitlab.png\"\n if source == DocumentSource.CONFLUENCE.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/Confluence.png\"\n if source == DocumentSource.JIRA.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/Jira.png\"\n if source == DocumentSource.NOTION.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Notion.png\"\n if source == DocumentSource.ZENDESK.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/Zendesk.png\"\n if source == DocumentSource.GONG.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Gong.png\"\n if source == DocumentSource.LINEAR.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Linear.png\"\n if source == DocumentSource.PRODUCTBOARD.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Productboard.webp\"\n if source == DocumentSource.SLAB.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/SlabLogo.png\"\n if source == DocumentSource.ZULIP.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Zulip.png\"\n if source == DocumentSource.GURU.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/Guru.png\"\n if source == DocumentSource.HUBSPOT.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/HubSpot.png\"\n if source == DocumentSource.DOCUMENT360.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Document360.png\"\n if source == DocumentSource.BOOKSTACK.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Bookstack.png\"\n if source == DocumentSource.OUTLINE.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Outline.png\"\n if source == DocumentSource.LOOPIO.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Loopio.png\"\n if source == DocumentSource.SHAREPOINT.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/web/public/Sharepoint.png\"\n if source == DocumentSource.REQUESTTRACKER.value:\n # just use file icon for now\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/File.png\"\n if source == DocumentSource.INGESTION_API.value:\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/File.png\"\n\n return \"https://raw.githubusercontent.com/onyx-dot-app/onyx/main/backend/slackbot_images/File.png\"\n" }, { "path": "backend/onyx/onyxbot/slack/listener.py", "content": "import os\nimport signal\nimport sys\nimport threading\nimport time\nfrom collections.abc import Callable\nfrom contextvars import Token\nfrom threading import Event\nfrom types import FrameType\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\n\nimport psycopg2.errors\nfrom prometheus_client import Gauge\nfrom prometheus_client import start_http_server\nfrom redis.lock import Lock\nfrom redis.lock import Lock as RedisLock\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\nfrom slack_sdk.http_retry import ConnectionErrorRetryHandler\nfrom slack_sdk.http_retry import RateLimitErrorRetryHandler\nfrom slack_sdk.http_retry import RetryHandler\nfrom slack_sdk.socket_mode.request import SocketModeRequest\nfrom slack_sdk.socket_mode.response import SocketModeResponse\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.configs.app_configs import POD_NAME\nfrom onyx.configs.app_configs import POD_NAMESPACE\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import OnyxRedisLocks\nfrom onyx.configs.onyxbot_configs import NOTIFY_SLACKBOT_NO_ANSWER\nfrom onyx.connectors.slack.utils import expert_info_from_slack_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.db.models import SlackBot\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.db.slack_bot import fetch_slack_bot\nfrom onyx.db.slack_bot import fetch_slack_bots\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder\nfrom onyx.onyxbot.slack.config import get_slack_channel_config_for_bot_and_channel\nfrom onyx.onyxbot.slack.config import MAX_TENANTS_PER_POD\nfrom onyx.onyxbot.slack.config import TENANT_ACQUISITION_INTERVAL\nfrom onyx.onyxbot.slack.config import TENANT_HEARTBEAT_EXPIRATION\nfrom onyx.onyxbot.slack.config import TENANT_HEARTBEAT_INTERVAL\nfrom onyx.onyxbot.slack.config import TENANT_LOCK_EXPIRATION\nfrom onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_ACTION_ID\nfrom onyx.onyxbot.slack.constants import FOLLOWUP_BUTTON_RESOLVED_ACTION_ID\nfrom onyx.onyxbot.slack.constants import GENERATE_ANSWER_BUTTON_ACTION_ID\nfrom onyx.onyxbot.slack.constants import IMMEDIATE_RESOLVED_BUTTON_ACTION_ID\nfrom onyx.onyxbot.slack.constants import KEEP_TO_YOURSELF_ACTION_ID\nfrom onyx.onyxbot.slack.constants import LIKE_BLOCK_ACTION_ID\nfrom onyx.onyxbot.slack.constants import SHOW_EVERYONE_ACTION_ID\nfrom onyx.onyxbot.slack.constants import VIEW_DOC_FEEDBACK_ID\nfrom onyx.onyxbot.slack.handlers.handle_buttons import handle_doc_feedback_button\nfrom onyx.onyxbot.slack.handlers.handle_buttons import handle_followup_button\nfrom onyx.onyxbot.slack.handlers.handle_buttons import (\n handle_followup_resolved_button,\n)\nfrom onyx.onyxbot.slack.handlers.handle_buttons import (\n handle_generate_answer_button,\n)\nfrom onyx.onyxbot.slack.handlers.handle_buttons import (\n handle_publish_ephemeral_message_button,\n)\nfrom onyx.onyxbot.slack.handlers.handle_buttons import handle_slack_feedback\nfrom onyx.onyxbot.slack.handlers.handle_message import handle_message\nfrom onyx.onyxbot.slack.handlers.handle_message import (\n remove_scheduled_feedback_reminder,\n)\nfrom onyx.onyxbot.slack.handlers.handle_message import schedule_feedback_reminder\nfrom onyx.onyxbot.slack.models import SlackContext\nfrom onyx.onyxbot.slack.models import SlackMessageInfo\nfrom onyx.onyxbot.slack.models import ThreadMessage\nfrom onyx.onyxbot.slack.utils import check_message_limit\nfrom onyx.onyxbot.slack.utils import decompose_action_id\nfrom onyx.onyxbot.slack.utils import get_channel_name_from_id\nfrom onyx.onyxbot.slack.utils import get_channel_type_from_id\nfrom onyx.onyxbot.slack.utils import get_onyx_bot_auth_ids\nfrom onyx.onyxbot.slack.utils import read_slack_thread\nfrom onyx.onyxbot.slack.utils import remove_onyx_bot_tag\nfrom onyx.onyxbot.slack.utils import respond_in_thread_or_channel\nfrom onyx.onyxbot.slack.utils import TenantSocketModeClient\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.manage.models import SlackBotTokens\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom onyx.utils.variable_functionality import set_is_ee_based_on_env_variable\nfrom shared_configs.configs import DISALLOWED_SLACK_BOT_TENANT_LIST\nfrom shared_configs.configs import MODEL_SERVER_HOST\nfrom shared_configs.configs import MODEL_SERVER_PORT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.configs import SLACK_CHANNEL_ID\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n# Prometheus metric for HPA\nactive_tenants_gauge = Gauge(\n \"active_tenants\",\n \"Number of active tenants handled by this pod\",\n [\"namespace\", \"pod\"],\n)\n\n# In rare cases, some users have been experiencing a massive amount of trivial messages coming through\n# to the Slack Bot with trivial messages. Adding this to avoid exploding LLM costs while we track down\n# the cause.\n_SLACK_GREETINGS_TO_IGNORE = {\n \"Welcome back!\",\n \"It's going to be a great day.\",\n \"Salutations!\",\n \"Greetings!\",\n \"Feeling great!\",\n \"Hi there\",\n \":wave:\",\n}\n\n# This is always (currently) the user id of Slack's official slackbot\n_OFFICIAL_SLACKBOT_USER_ID = \"USLACKBOT\"\n\n# Fields to exclude from Slack payload logging\n# Intention is to not log slack message content\n_EXCLUDED_SLACK_PAYLOAD_FIELDS = {\"text\", \"blocks\"}\n\n\nclass SlackbotHandler:\n def __init__(self) -> None:\n logger.info(\"Initializing SlackbotHandler\")\n self.tenant_ids: set[str] = set()\n # The keys for these dictionaries are tuples of (tenant_id, slack_bot_id)\n self.socket_clients: Dict[tuple[str, int], TenantSocketModeClient] = {}\n self.slack_bot_tokens: Dict[tuple[str, int], SlackBotTokens] = {}\n\n # Store Redis lock objects here so we can release them properly\n self.redis_locks: Dict[str, Lock] = {}\n\n self.running = True\n self.pod_id = os.environ.get(\"HOSTNAME\", \"unknown_pod\")\n self._shutdown_event = Event()\n\n self._lock = threading.Lock()\n\n logger.info(f\"Pod ID: {self.pod_id}\")\n\n # Set up signal handlers for graceful shutdown\n signal.signal(signal.SIGTERM, self.shutdown)\n signal.signal(signal.SIGINT, self.shutdown)\n logger.info(\"Signal handlers registered\")\n\n # Start the Prometheus metrics server\n logger.info(\"Starting Prometheus metrics server\")\n start_http_server(8000)\n logger.info(\"Prometheus metrics server started\")\n\n # Start background threads\n logger.info(\"Starting background threads\")\n self.acquire_thread = threading.Thread(\n target=self.acquire_tenants_loop, daemon=True\n )\n self.heartbeat_thread = threading.Thread(\n target=self.heartbeat_loop, daemon=True\n )\n\n self.acquire_thread.start()\n self.heartbeat_thread.start()\n\n logger.info(\"Background threads started\")\n\n def acquire_tenants_loop(self) -> None:\n while not self._shutdown_event.is_set():\n try:\n self.acquire_tenants()\n\n # After we finish acquiring and managing Slack bots,\n # set the gauge to the number of active tenants (those with Slack bots).\n active_tenants_gauge.labels(namespace=POD_NAMESPACE, pod=POD_NAME).set(\n len(self.tenant_ids)\n )\n logger.debug(\n f\"Current active tenants with Slack bots: {len(self.tenant_ids)}\"\n )\n except Exception as e:\n logger.exception(f\"Error in Slack acquisition: {e}\")\n self._shutdown_event.wait(timeout=TENANT_ACQUISITION_INTERVAL)\n\n def heartbeat_loop(self) -> None:\n \"\"\"This heartbeats into redis.\n\n NOTE(rkuo): this is not thread-safe with acquire_tenants_loop and will\n occasionally exception. Fix it!\n \"\"\"\n while not self._shutdown_event.is_set():\n try:\n with self._lock:\n tenant_ids = self.tenant_ids.copy()\n\n SlackbotHandler.send_heartbeats(self.pod_id, tenant_ids)\n logger.debug(f\"Sent heartbeats for {len(tenant_ids)} active tenants\")\n except Exception as e:\n logger.exception(f\"Error in heartbeat loop: {e}\")\n self._shutdown_event.wait(timeout=TENANT_HEARTBEAT_INTERVAL)\n\n def _manage_clients_per_tenant(\n self, db_session: Session, tenant_id: str, bot: SlackBot\n ) -> None:\n \"\"\"\n - If the tokens are missing or empty, close the socket client and remove them.\n - If the tokens have changed, close the existing socket client and reconnect.\n - If the tokens are new, warm up the model and start a new socket client.\n \"\"\"\n tenant_bot_pair = (tenant_id, bot.id)\n\n # If the tokens are missing or empty, close the socket client and remove them.\n if not bot.bot_token or not bot.app_token:\n logger.debug(\n f\"No Slack bot tokens found for tenant={tenant_id}, bot {bot.id}\"\n )\n if tenant_bot_pair in self.socket_clients:\n self.socket_clients[tenant_bot_pair].close()\n del self.socket_clients[tenant_bot_pair]\n del self.slack_bot_tokens[tenant_bot_pair]\n return\n\n slack_bot_tokens = SlackBotTokens(\n bot_token=bot.bot_token.get_value(apply_mask=False),\n app_token=bot.app_token.get_value(apply_mask=False),\n )\n\n tokens_exist = tenant_bot_pair in self.slack_bot_tokens\n tokens_changed = (\n tokens_exist and slack_bot_tokens != self.slack_bot_tokens[tenant_bot_pair]\n )\n if not tokens_exist or tokens_changed:\n if tokens_exist:\n logger.info(\n f\"Slack Bot tokens changed for tenant={tenant_id}, bot {bot.id}; reconnecting\"\n )\n else:\n # Warm up the model if needed\n search_settings = get_current_search_settings(db_session)\n embedding_model = EmbeddingModel.from_db_model(\n search_settings=search_settings,\n server_host=MODEL_SERVER_HOST,\n server_port=MODEL_SERVER_PORT,\n )\n warm_up_bi_encoder(embedding_model=embedding_model)\n\n self.slack_bot_tokens[tenant_bot_pair] = slack_bot_tokens\n\n # Close any existing connection first\n if tenant_bot_pair in self.socket_clients:\n self.socket_clients[tenant_bot_pair].close()\n\n socket_client = self.start_socket_client(\n bot.id, tenant_id, slack_bot_tokens\n )\n if socket_client:\n # Ensure tenant is tracked as active\n self.socket_clients[tenant_id, bot.id] = socket_client\n\n logger.info(\n f\"Started SocketModeClient: {tenant_id=} {socket_client.bot_name=} {bot.id=}\"\n )\n\n self.tenant_ids.add(tenant_id)\n\n def acquire_tenants(self) -> None:\n \"\"\"\n - Attempt to acquire a Redis lock for each tenant.\n - If acquired, check if that tenant actually has Slack bots.\n - If yes, store them in self.tenant_ids and manage the socket connections.\n - If a tenant in self.tenant_ids no longer has Slack bots, remove it (and release the lock in this scope).\n \"\"\"\n\n token: Token[str | None]\n\n # tenants that are disabled (e.g. their trial is over and haven't subscribed)\n # for non-cloud, this will return an empty set\n gated_tenants = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.product_gating\",\n \"get_gated_tenants\",\n set(),\n )()\n all_active_tenants = [\n tenant_id\n for tenant_id in get_all_tenant_ids()\n if tenant_id not in gated_tenants\n ]\n\n # 1) Try to acquire locks for new tenants\n for tenant_id in all_active_tenants:\n if (\n DISALLOWED_SLACK_BOT_TENANT_LIST is not None\n and tenant_id in DISALLOWED_SLACK_BOT_TENANT_LIST\n ):\n logger.debug(f\"Tenant {tenant_id} is disallowed; skipping.\")\n continue\n\n # Already acquired in a previous loop iteration?\n if tenant_id in self.tenant_ids:\n continue\n\n # Respect max tenant limit per pod\n if len(self.tenant_ids) >= MAX_TENANTS_PER_POD:\n logger.info(\n f\"Max tenants per pod reached, not acquiring more: {MAX_TENANTS_PER_POD=}\"\n )\n break\n\n redis_client = get_redis_client(tenant_id=tenant_id)\n # Acquire a Redis lock (non-blocking)\n # thread_local=False because the shutdown event is handled\n # on an arbitrary thread\n rlock: RedisLock = redis_client.lock(\n OnyxRedisLocks.SLACK_BOT_LOCK,\n timeout=TENANT_LOCK_EXPIRATION,\n thread_local=False,\n )\n lock_acquired = rlock.acquire(blocking=False)\n\n if not lock_acquired and not DEV_MODE:\n logger.debug(\n f\"Another pod holds the lock for tenant {tenant_id}, skipping.\"\n )\n continue\n\n if lock_acquired:\n logger.debug(f\"Acquired lock for tenant {tenant_id}.\")\n self.redis_locks[tenant_id] = rlock\n else:\n # DEV_MODE will skip the lock acquisition guard\n logger.debug(\n f\"Running in DEV_MODE. Not enforcing lock for {tenant_id}.\"\n )\n\n # Now check if this tenant actually has Slack bots\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(\n tenant_id or POSTGRES_DEFAULT_SCHEMA\n )\n try:\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n bots: list[SlackBot] = []\n try:\n bots = list(fetch_slack_bots(db_session=db_session))\n except KvKeyNotFoundError:\n # No Slackbot tokens, pass\n pass\n except psycopg2.errors.UndefinedTable:\n logger.error(\n \"Undefined table error in fetch_slack_bots. Tenant schema may need fixing.\"\n )\n except Exception as e:\n logger.exception(\n f\"Error fetching Slack bots for tenant {tenant_id}: {e}\"\n )\n\n if bots:\n # Mark as active tenant\n self.tenant_ids.add(tenant_id)\n for bot in bots:\n self._manage_clients_per_tenant(\n db_session=db_session,\n tenant_id=tenant_id,\n bot=bot,\n )\n else:\n # If no Slack bots, release lock immediately (unless in DEV_MODE)\n if lock_acquired and not DEV_MODE:\n rlock.release()\n del self.redis_locks[tenant_id]\n logger.debug(\n f\"No Slack bots for tenant {tenant_id}; lock released (if held).\"\n )\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n # 2) Make sure tenants we're handling still have Slack bots\n # and haven't been suspended (gated)\n for tenant_id in list(self.tenant_ids):\n if tenant_id in gated_tenants:\n logger.info(\n f\"Tenant {tenant_id} is now gated (suspended). Disconnecting.\"\n )\n self._remove_tenant(tenant_id)\n if tenant_id in self.redis_locks and not DEV_MODE:\n try:\n self.redis_locks[tenant_id].release()\n del self.redis_locks[tenant_id]\n except Exception as e:\n logger.error(\n f\"Error releasing lock for gated tenant {tenant_id}: {e}\"\n )\n continue\n\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(\n tenant_id or POSTGRES_DEFAULT_SCHEMA\n )\n redis_client = get_redis_client(tenant_id=tenant_id)\n\n try:\n with get_session_with_current_tenant() as db_session:\n # Attempt to fetch Slack bots\n try:\n bots = list(fetch_slack_bots(db_session=db_session))\n except KvKeyNotFoundError:\n # No Slackbot tokens, pass (and remove below)\n bots = []\n except Exception as e:\n logger.exception(f\"Error handling tenant {tenant_id}: {e}\")\n bots = []\n\n if not bots:\n logger.info(\n f\"Tenant {tenant_id} no longer has Slack bots. Removing.\"\n )\n self._remove_tenant(tenant_id)\n\n # NOTE: We release the lock here (in the same scope it was acquired)\n if tenant_id in self.redis_locks and not DEV_MODE:\n try:\n self.redis_locks[tenant_id].release()\n del self.redis_locks[tenant_id]\n logger.info(f\"Released lock for tenant {tenant_id}\")\n except Exception as e:\n logger.error(\n f\"Error releasing lock for tenant {tenant_id}: {e}\"\n )\n else:\n # Manage or reconnect Slack bot sockets\n for bot in bots:\n self._manage_clients_per_tenant(\n db_session=db_session,\n tenant_id=tenant_id,\n bot=bot,\n )\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n def _remove_tenant(self, tenant_id: str) -> None:\n \"\"\"\n Helper to remove a tenant from `self.tenant_ids` and close any socket clients.\n (Lock release now happens in `acquire_tenants()`, not here.)\n \"\"\"\n socket_client_list = list(self.socket_clients.items())\n # Close all socket clients for this tenant\n for (t_id, slack_bot_id), client in socket_client_list:\n if t_id == tenant_id:\n client.close()\n del self.socket_clients[(t_id, slack_bot_id)]\n del self.slack_bot_tokens[(t_id, slack_bot_id)]\n logger.info(\n f\"Stopped SocketModeClient for tenant: {t_id}, app: {slack_bot_id}\"\n )\n\n # Remove from active set\n if tenant_id in self.tenant_ids:\n self.tenant_ids.remove(tenant_id)\n\n @staticmethod\n def send_heartbeats(pod_id: str, tenant_ids: set[str]) -> None:\n current_time = int(time.time())\n logger.debug(f\"Sending heartbeats for {len(tenant_ids)} active tenants\")\n for tenant_id in tenant_ids:\n redis_client = get_redis_client(tenant_id=tenant_id)\n heartbeat_key = f\"{OnyxRedisLocks.SLACK_BOT_HEARTBEAT_PREFIX}:{pod_id}\"\n redis_client.set(\n heartbeat_key, current_time, ex=TENANT_HEARTBEAT_EXPIRATION\n )\n\n @staticmethod\n def start_socket_client(\n slack_bot_id: int, tenant_id: str, slack_bot_tokens: SlackBotTokens\n ) -> TenantSocketModeClient | None:\n \"\"\"Returns the socket client if this succeeds\"\"\"\n socket_client: TenantSocketModeClient = _get_socket_client(\n slack_bot_tokens, tenant_id, slack_bot_id\n )\n\n try:\n bot_info = socket_client.web_client.auth_test()\n\n if bot_info[\"ok\"]:\n bot_user_id = bot_info[\"user_id\"]\n user_info = socket_client.web_client.users_info(user=bot_user_id)\n if user_info[\"ok\"]:\n bot_name = (\n user_info[\"user\"][\"real_name\"] or user_info[\"user\"][\"name\"]\n )\n socket_client.bot_name = bot_name\n # logger.info(\n # f\"Started socket client for Slackbot with name '{bot_name}' (tenant: {tenant_id}, app: {slack_bot_id})\"\n # )\n except SlackApiError as e:\n # Only error out if we get a not_authed error\n if \"not_authed\" in str(e):\n # for some reason we want to add the tenant to the list when this happens?\n logger.error(\n f\"Authentication error - Invalid or expired credentials: {tenant_id=} {slack_bot_id=}. Error: {e}\"\n )\n return None\n\n # Log other Slack API errors but continue\n logger.error(\n f\"Slack API error fetching bot info: {e} for tenant: {tenant_id}, app: {slack_bot_id}\"\n )\n except Exception as e:\n # Log other exceptions but continue\n logger.error(\n f\"Error fetching bot info: {e} for tenant: {tenant_id}, app: {slack_bot_id}\"\n )\n\n # Append the event handler\n process_slack_event = create_process_slack_event()\n socket_client.socket_mode_request_listeners.append(process_slack_event) # type: ignore\n\n # Establish a WebSocket connection to the Socket Mode servers\n # logger.debug(\n # f\"Connecting socket client for tenant: {tenant_id}, app: {slack_bot_id}\"\n # )\n socket_client.connect()\n # logger.info(\n # f\"Started SocketModeClient for tenant: {tenant_id}, app: {slack_bot_id}\"\n # )\n\n return socket_client\n\n @staticmethod\n def stop_socket_clients(\n pod_id: str, socket_clients: Dict[tuple[str, int], TenantSocketModeClient]\n ) -> None:\n socket_client_list = list(socket_clients.items())\n length = len(socket_client_list)\n\n x = 0\n for (tenant_id, slack_bot_id), client in socket_client_list:\n x += 1\n client.close()\n logger.info(\n f\"Stopped SocketModeClient {x}/{length}: {pod_id=} {tenant_id=} {slack_bot_id=}\"\n )\n\n def shutdown(\n self,\n signum: int | None, # noqa: ARG002\n frame: FrameType | None, # noqa: ARG002\n ) -> None:\n if not self.running:\n return\n\n logger.info(\"Shutting down gracefully\")\n self.running = False\n self._shutdown_event.set() # set the shutdown event\n\n # wait for threads to detect the event and exit\n self.acquire_thread.join(timeout=60.0)\n self.heartbeat_thread.join(timeout=60.0)\n\n # Stop all socket clients\n logger.info(f\"Stopping {len(self.socket_clients)} socket clients\")\n SlackbotHandler.stop_socket_clients(self.pod_id, self.socket_clients)\n\n # Release locks for all tenants we currently hold\n logger.info(f\"Releasing locks for {len(self.tenant_ids)} tenants\")\n for tenant_id in list(self.tenant_ids):\n if tenant_id in self.redis_locks:\n try:\n self.redis_locks[tenant_id].release()\n logger.info(f\"Released lock for tenant {tenant_id}\")\n except Exception as e:\n logger.error(f\"Error releasing lock for tenant {tenant_id}: {e}\")\n finally:\n del self.redis_locks[tenant_id]\n\n # Wait for background threads to finish (with a timeout)\n logger.info(\"Waiting for background threads to finish...\")\n self.acquire_thread.join(timeout=5)\n self.heartbeat_thread.join(timeout=5)\n\n logger.info(\"Shutdown complete\")\n sys.exit(0)\n\n\ndef sanitize_slack_payload(payload: dict) -> dict:\n \"\"\"Remove message content from Slack payload for logging\"\"\"\n sanitized = {\n k: v for k, v in payload.items() if k not in _EXCLUDED_SLACK_PAYLOAD_FIELDS\n }\n if \"event\" in sanitized and isinstance(sanitized[\"event\"], dict):\n sanitized[\"event\"] = {\n k: v\n for k, v in sanitized[\"event\"].items()\n if k not in _EXCLUDED_SLACK_PAYLOAD_FIELDS\n }\n return sanitized\n\n\ndef prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -> bool:\n \"\"\"True to keep going, False to ignore this Slack request\"\"\"\n\n # skip cases where the bot is disabled in the web UI\n tenant_id = get_current_tenant_id()\n\n bot_token_user_id, bot_token_bot_id = get_onyx_bot_auth_ids(\n tenant_id, client.web_client\n )\n logger.info(f\"prefilter_requests: {bot_token_user_id=} {bot_token_bot_id=}\")\n\n with get_session_with_current_tenant() as db_session:\n slack_bot = fetch_slack_bot(\n db_session=db_session, slack_bot_id=client.slack_bot_id\n )\n if not slack_bot:\n logger.error(\n f\"Slack bot with ID '{client.slack_bot_id}' not found. Skipping request.\"\n )\n return False\n\n if not slack_bot.enabled:\n logger.info(\n f\"Slack bot with ID '{client.slack_bot_id}' is disabled. Skipping request.\"\n )\n return False\n\n if req.type == \"events_api\":\n # Verify channel is valid\n event = cast(dict[str, Any], req.payload.get(\"event\", {}))\n msg = cast(str | None, event.get(\"text\"))\n channel = cast(str | None, event.get(\"channel\"))\n channel_specific_logger = setup_logger(extra={SLACK_CHANNEL_ID: channel})\n\n # This should never happen, but we can't continue without a channel since\n # we can't send a response without it\n if not channel:\n channel_specific_logger.warning(\"Found message without channel - skipping\")\n return False\n\n if not msg:\n channel_specific_logger.warning(\n \"Cannot respond to empty message - skipping\"\n )\n return False\n\n if (\n req.payload.setdefault(\"event\", {}).get(\"user\", \"\")\n == _OFFICIAL_SLACKBOT_USER_ID\n ):\n channel_specific_logger.info(\n \"Ignoring messages from Slack's official Slackbot\"\n )\n return False\n\n if (\n msg in _SLACK_GREETINGS_TO_IGNORE\n or remove_onyx_bot_tag(tenant_id, msg, client=client.web_client)\n in _SLACK_GREETINGS_TO_IGNORE\n ):\n channel_specific_logger.error(\n f\"Ignoring weird Slack greeting message: '{msg}'\"\n )\n channel_specific_logger.error(\n f\"Weird Slack greeting message payload: '{req.payload}'\"\n )\n return False\n\n # Ensure that the message is a new message of expected type\n event_type = event.get(\"type\")\n event.get(\"channel_type\")\n\n if event_type not in [\"app_mention\", \"message\"]:\n return False\n\n bot_token_user_id, bot_token_bot_id = get_onyx_bot_auth_ids(\n tenant_id, client.web_client\n )\n if event_type == \"message\":\n is_onyx_bot_msg = False\n is_tagged = False\n\n event_user = event.get(\"user\", \"\")\n event_bot_id = event.get(\"bot_id\", \"\")\n\n is_dm = event.get(\"channel_type\") == \"im\"\n if bot_token_user_id and f\"<@{bot_token_user_id}>\" in msg:\n is_tagged = True\n\n if bot_token_user_id and bot_token_user_id in event_user:\n is_onyx_bot_msg = True\n\n if bot_token_bot_id and bot_token_bot_id in event_bot_id:\n is_onyx_bot_msg = True\n\n # OnyxBot should never respond to itself\n if is_onyx_bot_msg:\n logger.info(\"Ignoring message from OnyxBot (self-message)\")\n return False\n\n # DMs with the bot don't pick up the @OnyxBot so we have to keep the\n # caught events_api\n if is_tagged and not is_dm:\n # Let the tag flow handle this case, don't reply twice\n return False\n\n # Check if this is a bot message (either via bot_profile or bot_message subtype)\n is_bot_message = bool(\n event.get(\"bot_profile\") or event.get(\"subtype\") == \"bot_message\"\n )\n if is_bot_message:\n channel_name, _ = get_channel_name_from_id(\n client=client.web_client, channel_id=channel\n )\n with get_session_with_current_tenant() as db_session:\n slack_channel_config = get_slack_channel_config_for_bot_and_channel(\n db_session=db_session,\n slack_bot_id=client.slack_bot_id,\n channel_name=channel_name,\n )\n\n # If OnyxBot is not specifically tagged and the channel is not set to respond to bots, ignore the message\n if (not bot_token_user_id or bot_token_user_id not in msg) and (\n not slack_channel_config\n or not slack_channel_config.channel_config.get(\"respond_to_bots\")\n ):\n channel_specific_logger.info(\n \"Ignoring message from bot since respond_to_bots is disabled\"\n )\n return False\n\n # Ignore things like channel_join, channel_leave, etc.\n # NOTE: \"file_share\" is just a message with a file attachment, so we\n # should not ignore it\n message_subtype = event.get(\"subtype\")\n if message_subtype not in [None, \"file_share\", \"bot_message\"]:\n channel_specific_logger.info(\n f\"Ignoring message with subtype '{message_subtype}' since it is a special message type\"\n )\n return False\n\n message_ts = event.get(\"ts\")\n thread_ts = event.get(\"thread_ts\")\n # Pick the root of the thread (if a thread exists)\n # Can respond in thread if it's an \"im\" directly to Onyx or @OnyxBot is tagged\n if (\n thread_ts\n and message_ts != thread_ts\n and event_type != \"app_mention\"\n and event.get(\"channel_type\") != \"im\"\n ):\n channel_specific_logger.debug(\n \"Skipping message since it is not the root of a thread\"\n )\n return False\n\n msg = cast(str, event.get(\"text\", \"\"))\n if not msg:\n channel_specific_logger.error(\"Unable to process empty message\")\n return False\n\n if req.type == \"slash_commands\":\n # Verify that there's an associated channel\n channel = req.payload.get(\"channel_id\")\n channel_specific_logger = setup_logger(extra={SLACK_CHANNEL_ID: channel})\n\n if not channel:\n channel_specific_logger.error(\n \"Received OnyxBot command without channel - skipping\"\n )\n return False\n\n sender = req.payload.get(\"user_id\")\n if not sender:\n channel_specific_logger.error(\n \"Cannot respond to OnyxBot command without sender to respond to.\"\n )\n return False\n\n if not check_message_limit():\n return False\n\n # Don't log Slack message content\n logger.debug(\n f\"Handling Slack request: {client.bot_name=} '{sanitize_slack_payload(req.payload)=}'\"\n )\n return True\n\n\ndef process_feedback(req: SocketModeRequest, client: TenantSocketModeClient) -> None:\n if actions := req.payload.get(\"actions\"):\n action = cast(dict[str, Any], actions[0])\n feedback_type = cast(str, action.get(\"action_id\"))\n feedback_msg_reminder = cast(str, action.get(\"value\"))\n feedback_id = cast(str, action.get(\"block_id\"))\n channel_id = cast(str, req.payload[\"container\"][\"channel_id\"])\n thread_ts = cast(\n str,\n req.payload[\"container\"].get(\"thread_ts\")\n or req.payload[\"container\"].get(\"message_ts\"),\n )\n else:\n logger.error(\"Unable to process feedback. Action not found\")\n return\n\n user_id = cast(str, req.payload[\"user\"][\"id\"])\n\n handle_slack_feedback(\n feedback_id=feedback_id,\n feedback_type=feedback_type,\n feedback_msg_reminder=feedback_msg_reminder,\n client=client.web_client,\n user_id_to_post_confirmation=user_id,\n channel_id_to_post_confirmation=channel_id,\n thread_ts_to_post_confirmation=thread_ts,\n )\n\n query_event_id, _, _ = decompose_action_id(feedback_id)\n logger.info(f\"Successfully handled QA feedback for event: {query_event_id}\")\n\n\ndef build_request_details(\n req: SocketModeRequest, client: TenantSocketModeClient\n) -> SlackMessageInfo:\n tagged: bool = False\n\n tenant_id = get_current_tenant_id()\n if req.type == \"events_api\":\n event = cast(dict[str, Any], req.payload[\"event\"])\n msg = cast(str, event[\"text\"])\n channel = cast(str, event[\"channel\"])\n\n # Check for both app_mention events and messages containing bot tag\n bot_token_user_id, _ = get_onyx_bot_auth_ids(tenant_id, client.web_client)\n message_ts = event.get(\"ts\")\n thread_ts = event.get(\"thread_ts\")\n sender_id = event.get(\"user\") or None\n expert_info = expert_info_from_slack_id(\n sender_id, client.web_client, user_cache={}\n )\n email = expert_info.email if expert_info else None\n\n msg = remove_onyx_bot_tag(tenant_id, msg, client=client.web_client)\n\n logger.info(f\"Received Slack message: {msg}\")\n\n event_type = event.get(\"type\")\n if event_type == \"app_mention\":\n tagged = True\n\n if event_type == \"message\":\n if bot_token_user_id:\n if f\"<@{bot_token_user_id}>\" in msg:\n tagged = True\n\n if tagged:\n logger.debug(\"User tagged OnyxBot\")\n\n # Build Slack context for federated search\n # Get proper channel type from Slack API instead of relying on event.channel_type\n channel_type = get_channel_type_from_id(client.web_client, channel)\n\n slack_context = SlackContext(\n channel_type=channel_type,\n channel_id=channel,\n user_id=sender_id or \"unknown\",\n message_ts=message_ts,\n )\n logger.info(\n f\"build_request_details: Capturing Slack context: \"\n f\"channel_type={channel_type} channel_id={channel} message_ts={message_ts}\"\n )\n\n if thread_ts != message_ts and thread_ts is not None:\n thread_messages: list[ThreadMessage] = read_slack_thread(\n tenant_id=tenant_id,\n channel=channel,\n thread=thread_ts,\n client=client.web_client,\n )\n else:\n sender_display_name = None\n if expert_info:\n sender_display_name = expert_info.display_name\n if sender_display_name is None:\n sender_display_name = (\n f\"{expert_info.first_name} {expert_info.last_name}\"\n if expert_info.last_name\n else expert_info.first_name\n )\n if sender_display_name is None:\n sender_display_name = expert_info.email\n thread_messages = [\n ThreadMessage(\n message=msg, sender=sender_display_name, role=MessageType.USER\n )\n ]\n\n return SlackMessageInfo(\n thread_messages=thread_messages,\n channel_to_respond=channel,\n msg_to_respond=cast(str, message_ts or thread_ts),\n thread_to_respond=cast(str, thread_ts or message_ts),\n sender_id=sender_id,\n email=email,\n bypass_filters=tagged,\n is_slash_command=False,\n is_bot_dm=event.get(\"channel_type\") == \"im\",\n slack_context=slack_context, # Add Slack context for federated search\n )\n\n elif req.type == \"slash_commands\":\n channel = req.payload[\"channel_id\"]\n channel_name = req.payload[\"channel_name\"]\n msg = req.payload[\"text\"]\n sender = req.payload[\"user_id\"]\n expert_info = expert_info_from_slack_id(\n sender, client.web_client, user_cache={}\n )\n email = expert_info.email if expert_info else None\n\n # Get proper channel type for slash commands too\n channel_type = get_channel_type_from_id(client.web_client, channel)\n\n slack_context = SlackContext(\n channel_type=channel_type,\n channel_id=channel,\n user_id=sender,\n message_ts=None, # Slash commands don't have a message timestamp\n )\n logger.info(\n f\"build_request_details: Capturing Slack context for slash command: channel_type={channel_type} channel_id={channel}\"\n )\n\n single_msg = ThreadMessage(message=msg, sender=None, role=MessageType.USER)\n\n return SlackMessageInfo(\n thread_messages=[single_msg],\n channel_to_respond=channel,\n msg_to_respond=None,\n thread_to_respond=None,\n sender_id=sender,\n email=email,\n bypass_filters=True,\n is_slash_command=True,\n is_bot_dm=channel_name == \"directmessage\",\n slack_context=slack_context, # Add Slack context for federated search\n )\n\n raise RuntimeError(\"Programming fault, this should never happen.\")\n\n\ndef apologize_for_fail(\n details: SlackMessageInfo,\n client: TenantSocketModeClient,\n) -> None:\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=details.channel_to_respond,\n thread_ts=details.msg_to_respond,\n text=\"Sorry, we weren't able to find anything relevant :cold_sweat:\",\n )\n\n\ndef process_message(\n req: SocketModeRequest,\n client: TenantSocketModeClient,\n notify_no_answer: bool = NOTIFY_SLACKBOT_NO_ANSWER,\n) -> None:\n tenant_id = get_current_tenant_id()\n if req.type == \"events_api\":\n event = cast(dict[str, Any], req.payload[\"event\"])\n event_type = event.get(\"type\")\n logger.info(\n f\"process_message start: {tenant_id=} {req.type=} {req.envelope_id=} {event_type=}\"\n )\n else:\n logger.info(\n f\"process_message start: {tenant_id=} {req.type=} {req.envelope_id=}\"\n )\n\n # Throw out requests that can't or shouldn't be handled\n if not prefilter_requests(req, client):\n logger.info(\n f\"process_message prefiltered: {tenant_id=} {req.type=} {req.envelope_id=}\"\n )\n return\n\n details = build_request_details(req, client)\n channel = details.channel_to_respond\n channel_name, is_dm = get_channel_name_from_id(\n client=client.web_client, channel_id=channel\n )\n\n with get_session_with_current_tenant() as db_session:\n slack_channel_config = get_slack_channel_config_for_bot_and_channel(\n db_session=db_session,\n slack_bot_id=client.slack_bot_id,\n channel_name=channel_name,\n )\n\n follow_up = bool(\n slack_channel_config.channel_config\n and slack_channel_config.channel_config.get(\"follow_up_tags\") is not None\n )\n\n feedback_reminder_id = schedule_feedback_reminder(\n details=details, client=client.web_client, include_followup=follow_up\n )\n\n failed = handle_message(\n message_info=details,\n slack_channel_config=slack_channel_config,\n client=client.web_client,\n feedback_reminder_id=feedback_reminder_id,\n )\n\n if failed:\n if feedback_reminder_id:\n remove_scheduled_feedback_reminder(\n client=client.web_client,\n channel=details.sender_id,\n msg_id=feedback_reminder_id,\n )\n # Skipping answering due to pre-filtering is not considered a failure\n if notify_no_answer:\n apologize_for_fail(details, client)\n\n logger.info(\n f\"process_message finished: success={not failed} {tenant_id=} {req.type=} {req.envelope_id=}\"\n )\n\n\ndef acknowledge_message(req: SocketModeRequest, client: TenantSocketModeClient) -> None:\n response = SocketModeResponse(envelope_id=req.envelope_id)\n client.send_socket_mode_response(response)\n\n\ndef action_routing(req: SocketModeRequest, client: TenantSocketModeClient) -> None:\n if actions := req.payload.get(\"actions\"):\n action = cast(dict[str, Any], actions[0])\n\n if action[\"action_id\"] in [DISLIKE_BLOCK_ACTION_ID, LIKE_BLOCK_ACTION_ID]:\n # AI Answer feedback\n return process_feedback(req, client)\n elif action[\"action_id\"] in [\n SHOW_EVERYONE_ACTION_ID,\n KEEP_TO_YOURSELF_ACTION_ID,\n ]:\n # Publish ephemeral message or keep hidden in main channel\n return handle_publish_ephemeral_message_button(\n req, client, action[\"action_id\"]\n )\n elif action[\"action_id\"] == FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID:\n # Activation of the \"source feedback\" button\n return handle_doc_feedback_button(req, client)\n elif action[\"action_id\"] == FOLLOWUP_BUTTON_ACTION_ID:\n return handle_followup_button(req, client)\n elif action[\"action_id\"] == IMMEDIATE_RESOLVED_BUTTON_ACTION_ID:\n return handle_followup_resolved_button(req, client, immediate=True)\n elif action[\"action_id\"] == FOLLOWUP_BUTTON_RESOLVED_ACTION_ID:\n return handle_followup_resolved_button(req, client, immediate=False)\n elif action[\"action_id\"] == GENERATE_ANSWER_BUTTON_ACTION_ID:\n return handle_generate_answer_button(req, client)\n\n\ndef view_routing(req: SocketModeRequest, client: TenantSocketModeClient) -> None:\n if view := req.payload.get(\"view\"):\n if view[\"callback_id\"] == VIEW_DOC_FEEDBACK_ID:\n return process_feedback(req, client)\n\n\ndef _extract_channel_from_request(req: SocketModeRequest) -> str | None:\n \"\"\"Best-effort channel extraction from any Slack request type.\"\"\"\n if req.type == \"events_api\":\n return cast(dict[str, Any], req.payload.get(\"event\", {})).get(\"channel\")\n elif req.type == \"slash_commands\":\n return req.payload.get(\"channel_id\")\n elif req.type == \"interactive\":\n container = req.payload.get(\"container\", {})\n return container.get(\"channel_id\") or req.payload.get(\"channel\", {}).get(\"id\")\n return None\n\n\ndef _check_tenant_gated(client: TenantSocketModeClient, req: SocketModeRequest) -> bool:\n \"\"\"Check if the current tenant is gated (suspended or license expired).\n\n Multi-tenant: checks the gated tenants Redis set (populated by control plane).\n Self-hosted: checks the cached license metadata for expiry.\n\n Returns True if blocked.\n \"\"\"\n from onyx.server.settings.models import ApplicationStatus\n\n # Multi-tenant path: control plane marks gated tenants in Redis\n is_gated: bool = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.product_gating\",\n \"is_tenant_gated\",\n False,\n )(get_current_tenant_id())\n\n # Self-hosted path: check license metadata cache\n if not is_gated:\n get_cached_metadata = fetch_ee_implementation_or_noop(\n \"onyx.db.license\",\n \"get_cached_license_metadata\",\n None,\n )\n metadata = get_cached_metadata()\n if metadata is not None:\n if metadata.status == ApplicationStatus.GATED_ACCESS:\n is_gated = True\n\n if not is_gated:\n return False\n\n # Only notify once per user action:\n # - Skip bot messages (avoids feedback loop from our own response)\n # - Skip app_mention events (Slack fires both app_mention AND message\n # for @mentions; we respond on the message event only)\n event = req.payload.get(\"event\", {}) if req.type == \"events_api\" else {}\n is_bot_event = bool(\n event.get(\"bot_id\")\n or event.get(\"bot_profile\")\n or event.get(\"subtype\") == \"bot_message\"\n )\n is_duplicate_mention = event.get(\"type\") == \"app_mention\"\n if not is_bot_event and not is_duplicate_mention:\n channel = _extract_channel_from_request(req)\n thread_ts = event.get(\"thread_ts\") or event.get(\"ts\")\n if channel:\n respond_in_thread_or_channel(\n client=client.web_client,\n channel=channel,\n thread_ts=thread_ts,\n text=(\n \"Your organization's subscription has expired. Please contact your Onyx administrator to restore access.\"\n ),\n )\n logger.info(f\"Blocked Slack request for gated tenant {get_current_tenant_id()}\")\n return True\n\n\ndef create_process_slack_event() -> (\n Callable[[TenantSocketModeClient, SocketModeRequest], None]\n):\n def process_slack_event(\n client: TenantSocketModeClient, req: SocketModeRequest\n ) -> None:\n # Always respond right away, if Slack doesn't receive these frequently enough\n # it will assume the Bot is DEAD!!! :(\n acknowledge_message(req, client)\n\n if _check_tenant_gated(client, req):\n return\n\n try:\n if req.type == \"interactive\":\n if req.payload.get(\"type\") == \"block_actions\":\n return action_routing(req, client)\n elif req.payload.get(\"type\") == \"view_submission\":\n return view_routing(req, client)\n elif req.type == \"events_api\" or req.type == \"slash_commands\":\n return process_message(req, client)\n except Exception:\n logger.exception(\"Failed to process slack event\")\n\n return process_slack_event\n\n\ndef _get_socket_client(\n slack_bot_tokens: SlackBotTokens, tenant_id: str, slack_bot_id: int\n) -> TenantSocketModeClient:\n # For more info on how to set this up, checkout the docs:\n # https://docs.onyx.app/admins/getting_started/slack_bot_setup\n\n # use the retry handlers built into the slack sdk\n connection_error_retry_handler = ConnectionErrorRetryHandler()\n rate_limit_error_retry_handler = RateLimitErrorRetryHandler(max_retry_count=7)\n slack_retry_handlers: list[RetryHandler] = [\n connection_error_retry_handler,\n rate_limit_error_retry_handler,\n ]\n\n return TenantSocketModeClient(\n # This app-level token will be used only for establishing a connection\n app_token=slack_bot_tokens.app_token,\n web_client=WebClient(\n token=slack_bot_tokens.bot_token, retry_handlers=slack_retry_handlers\n ),\n tenant_id=tenant_id,\n slack_bot_id=slack_bot_id,\n )\n\n\nif __name__ == \"__main__\":\n # Initialize the SqlEngine\n SqlEngine.init_engine(pool_size=20, max_overflow=5)\n\n # Initialize the tenant handler which will manage tenant connections\n logger.info(\"Starting SlackbotHandler\")\n tenant_handler = SlackbotHandler()\n\n set_is_ee_based_on_env_variable()\n\n try:\n # Keep the main thread alive\n while tenant_handler.running:\n time.sleep(1)\n\n except Exception:\n logger.exception(\"Fatal error in main thread\")\n tenant_handler.shutdown(None, None)\n" }, { "path": "backend/onyx/onyxbot/slack/models.py", "content": "from enum import Enum\nfrom typing import Literal\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import MessageType\n\n\nclass ChannelType(str, Enum):\n \"\"\"Slack channel types.\"\"\"\n\n IM = \"im\" # Direct message\n MPIM = \"mpim\" # Multi-person direct message\n PRIVATE_CHANNEL = \"private_channel\" # Private channel\n PUBLIC_CHANNEL = \"public_channel\" # Public channel\n UNKNOWN = \"unknown\" # Unknown channel type\n\n\nclass SlackContext(BaseModel):\n \"\"\"Context information for Slack bot interactions.\"\"\"\n\n channel_type: ChannelType\n channel_id: str\n user_id: str\n message_ts: str | None = None # Used as request ID for log correlation\n\n\nclass ThreadMessage(BaseModel):\n message: str\n sender: str | None = None\n role: MessageType = MessageType.USER\n\n\nclass SlackMessageInfo(BaseModel):\n thread_messages: list[ThreadMessage]\n channel_to_respond: str\n msg_to_respond: str | None\n thread_to_respond: str | None\n sender_id: str | None\n email: str | None\n bypass_filters: bool # User has tagged @OnyxBot\n is_slash_command: bool # User is using /OnyxBot\n is_bot_dm: bool # User is direct messaging to OnyxBot\n slack_context: SlackContext | None = None\n\n\n# Models used to encode the relevant data for the ephemeral message actions\nclass ActionValuesEphemeralMessageMessageInfo(BaseModel):\n bypass_filters: bool | None\n channel_to_respond: str | None\n msg_to_respond: str | None\n email: str | None\n sender_id: str | None\n thread_messages: list[ThreadMessage] | None\n is_slash_command: bool | None\n is_bot_dm: bool | None\n thread_to_respond: str | None\n\n\nclass ActionValuesEphemeralMessageChannelConfig(BaseModel):\n channel_name: str | None\n respond_tag_only: bool | None\n respond_to_bots: bool | None\n is_ephemeral: bool\n respond_member_group_list: list[str] | None\n answer_filters: (\n list[Literal[\"well_answered_postfilter\", \"questionmark_prefilter\"]] | None\n )\n follow_up_tags: list[str] | None\n show_continue_in_web_ui: bool\n\n\nclass ActionValuesEphemeralMessage(BaseModel):\n original_question_ts: str | None\n feedback_reminder_id: str | None\n chat_message_id: int\n message_info: ActionValuesEphemeralMessageMessageInfo\n channel_conf: ActionValuesEphemeralMessageChannelConfig\n" }, { "path": "backend/onyx/onyxbot/slack/utils.py", "content": "import logging\nimport random\nimport re\nimport string\nimport threading\nimport time\nimport uuid\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom typing import cast\n\nfrom retry import retry\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\nfrom slack_sdk.models.blocks import Block\nfrom slack_sdk.models.blocks import SectionBlock\nfrom slack_sdk.models.metadata import Metadata\nfrom slack_sdk.socket_mode import SocketModeClient\n\nfrom onyx.configs.app_configs import DISABLE_TELEMETRY\nfrom onyx.configs.constants import ID_SEPARATOR\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_FEEDBACK_VISIBILITY\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_MAX_QPM\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_MAX_WAIT_TIME\nfrom onyx.configs.onyxbot_configs import ONYX_BOT_NUM_RETRIES\nfrom onyx.configs.onyxbot_configs import (\n ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD,\n)\nfrom onyx.configs.onyxbot_configs import (\n ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS,\n)\nfrom onyx.connectors.slack.utils import SlackTextCleaner\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.users import get_user_by_email\nfrom onyx.onyxbot.slack.constants import FeedbackVisibility\nfrom onyx.onyxbot.slack.models import ChannelType\nfrom onyx.onyxbot.slack.models import ThreadMessage\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import optional_telemetry\nfrom onyx.utils.telemetry import RecordType\nfrom onyx.utils.text_processing import replace_whitespaces_w_space\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\nslack_token_user_ids: dict[str, str | None] = {}\nslack_token_bot_ids: dict[str, str | None] = {}\nslack_token_lock = threading.Lock()\n\n_ONYX_BOT_MESSAGE_COUNT: int = 0\n_ONYX_BOT_COUNT_START_TIME: float = time.time()\n\n\ndef get_onyx_bot_auth_ids(\n tenant_id: str, web_client: WebClient\n) -> tuple[str | None, str | None]:\n \"\"\"Returns a tuple of user_id and bot_id.\"\"\"\n\n user_id: str | None\n bot_id: str | None\n\n global slack_token_user_ids\n global slack_token_bot_ids\n\n with slack_token_lock:\n user_id = slack_token_user_ids.get(tenant_id)\n bot_id = slack_token_bot_ids.get(tenant_id)\n\n if user_id is None or bot_id is None:\n response = web_client.auth_test()\n user_id = response.get(\"user_id\")\n bot_id = response.get(\"bot_id\")\n with slack_token_lock:\n slack_token_user_ids[tenant_id] = user_id\n slack_token_bot_ids[tenant_id] = bot_id\n\n return user_id, bot_id\n\n\ndef get_channel_type_from_id(web_client: WebClient, channel_id: str) -> ChannelType:\n \"\"\"\n Get the channel type from a channel ID using Slack API.\n Returns: ChannelType enum value\n \"\"\"\n try:\n channel_info = web_client.conversations_info(channel=channel_id)\n if channel_info.get(\"ok\") and channel_info.get(\"channel\"):\n channel: dict[str, Any] = channel_info.get(\"channel\", {})\n\n if channel.get(\"is_im\"):\n return ChannelType.IM # Direct message\n elif channel.get(\"is_mpim\"):\n return ChannelType.MPIM # Multi-person direct message\n elif channel.get(\"is_private\"):\n return ChannelType.PRIVATE_CHANNEL # Private channel\n elif channel.get(\"is_channel\"):\n return ChannelType.PUBLIC_CHANNEL # Public channel\n else:\n logger.warning(\n f\"Could not determine channel type for {channel_id}, defaulting to unknown\"\n )\n return ChannelType.UNKNOWN\n else:\n logger.warning(f\"Invalid channel info response for {channel_id}\")\n return ChannelType.UNKNOWN\n except Exception as e:\n logger.warning(\n f\"Error getting channel info for {channel_id}, defaulting to unknown: {e}\"\n )\n return ChannelType.UNKNOWN\n\n\ndef check_message_limit() -> bool:\n \"\"\"\n This isnt a perfect solution.\n High traffic at the end of one period and start of another could cause\n the limit to be exceeded.\n \"\"\"\n if ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD <= 0:\n return True\n global _ONYX_BOT_MESSAGE_COUNT\n global _ONYX_BOT_COUNT_START_TIME\n time_since_start = time.time() - _ONYX_BOT_COUNT_START_TIME\n if time_since_start > ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS:\n _ONYX_BOT_MESSAGE_COUNT = 0\n _ONYX_BOT_COUNT_START_TIME = time.time()\n if (_ONYX_BOT_MESSAGE_COUNT + 1) > ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD:\n logger.error(\n f\"OnyxBot has reached the message limit {ONYX_BOT_RESPONSE_LIMIT_PER_TIME_PERIOD}\"\n f\" for the time period {ONYX_BOT_RESPONSE_LIMIT_TIME_PERIOD_SECONDS} seconds.\"\n \" These limits are configurable in backend/onyx/configs/onyxbot_configs.py\"\n )\n return False\n _ONYX_BOT_MESSAGE_COUNT += 1\n return True\n\n\ndef update_emote_react(\n emoji: str,\n channel: str,\n message_ts: str | None,\n remove: bool,\n client: WebClient,\n) -> None:\n if not message_ts:\n action = \"remove\" if remove else \"add\"\n logger.error(f\"update_emote_react - no message specified: {channel=} {action=}\")\n return\n\n if remove:\n try:\n client.reactions_remove(\n name=emoji,\n channel=channel,\n timestamp=message_ts,\n )\n except SlackApiError as e:\n logger.error(f\"Failed to remove Reaction due to: {e}\")\n\n return\n\n try:\n client.reactions_add(\n name=emoji,\n channel=channel,\n timestamp=message_ts,\n )\n except SlackApiError as e:\n logger.error(f\"Was not able to react to user message due to: {e}\")\n\n return\n\n\ndef remove_onyx_bot_tag(tenant_id: str, message_str: str, client: WebClient) -> str:\n bot_token_user_id, _ = get_onyx_bot_auth_ids(tenant_id, web_client=client)\n return re.sub(rf\"<@{bot_token_user_id}>\\s*\", \"\", message_str)\n\n\ndef _check_for_url_in_block(block: Block) -> bool:\n \"\"\"\n Check if the block has a key that contains \"url\" in it\n \"\"\"\n block_dict = block.to_dict()\n\n def check_dict_for_url(d: dict) -> bool:\n for key, value in d.items():\n if \"url\" in key.lower():\n return True\n if isinstance(value, dict):\n if check_dict_for_url(value):\n return True\n elif isinstance(value, list):\n for item in value:\n if isinstance(item, dict) and check_dict_for_url(item):\n return True\n return False\n\n return check_dict_for_url(block_dict)\n\n\ndef _build_error_block(error_message: str) -> Block:\n \"\"\"\n Build an error block to display in slack so that the user can see\n the error without completely breaking\n \"\"\"\n display_text = (\n \"There was an error displaying all of the Onyx answers.\"\n f\" Please let an admin or an onyx developer know. Error: {error_message}\"\n )\n return SectionBlock(text=display_text)\n\n\n@retry(\n tries=ONYX_BOT_NUM_RETRIES,\n delay=0.25,\n backoff=2,\n logger=cast(logging.Logger, logger),\n)\ndef respond_in_thread_or_channel(\n client: WebClient,\n channel: str,\n thread_ts: str | None,\n text: str | None = None,\n blocks: list[Block] | None = None,\n receiver_ids: list[str] | None = None,\n metadata: Metadata | None = None,\n unfurl: bool = True,\n send_as_ephemeral: bool | None = True, # noqa: ARG001\n) -> list[str]:\n if not text and not blocks:\n raise ValueError(\"One of `text` or `blocks` must be provided\")\n\n message_ids: list[str] = []\n if not receiver_ids:\n try:\n response = client.chat_postMessage(\n channel=channel,\n text=text,\n blocks=blocks,\n thread_ts=thread_ts,\n metadata=metadata,\n unfurl_links=unfurl,\n unfurl_media=unfurl,\n )\n except Exception as e:\n blocks_str = str(blocks)[:1024] # truncate block logging\n logger.warning(f\"Failed to post message: {e} \\n blocks: {blocks_str}\")\n logger.warning(\"Trying again without blocks that have urls\")\n\n if not blocks:\n raise e\n\n blocks_without_urls = [\n block for block in blocks if not _check_for_url_in_block(block)\n ]\n blocks_without_urls.append(_build_error_block(str(e)))\n\n # Try again wtihout blocks containing url\n response = client.chat_postMessage(\n channel=channel,\n text=text,\n blocks=blocks_without_urls,\n thread_ts=thread_ts,\n metadata=metadata,\n unfurl_links=unfurl,\n unfurl_media=unfurl,\n )\n\n message_ids.append(response[\"message_ts\"])\n else:\n for receiver in receiver_ids:\n try:\n response = client.chat_postEphemeral(\n channel=channel,\n user=receiver,\n text=text,\n blocks=blocks,\n thread_ts=thread_ts,\n metadata=metadata,\n unfurl_links=unfurl,\n unfurl_media=unfurl,\n )\n except Exception as e:\n blocks_str = str(blocks)[:1024] # truncate block logging\n logger.warning(f\"Failed to post message: {e} \\n blocks: {blocks_str}\")\n logger.warning(\"Trying again without blocks that have urls\")\n\n if not blocks:\n raise e\n\n blocks_without_urls = [\n block for block in blocks if not _check_for_url_in_block(block)\n ]\n blocks_without_urls.append(_build_error_block(str(e)))\n\n # Try again wtihout blocks containing url\n response = client.chat_postEphemeral(\n channel=channel,\n user=receiver,\n text=text,\n blocks=blocks_without_urls,\n thread_ts=thread_ts,\n metadata=metadata,\n unfurl_links=unfurl,\n unfurl_media=unfurl,\n )\n\n message_ids.append(response[\"message_ts\"])\n\n return message_ids\n\n\ndef build_feedback_id(\n message_id: int,\n document_id: str | None = None,\n document_rank: int | None = None,\n) -> str:\n unique_prefix = \"\".join(random.choice(string.ascii_letters) for _ in range(10))\n if document_id is not None:\n if not document_id or document_rank is None:\n raise ValueError(\"Invalid document, missing information\")\n if ID_SEPARATOR in document_id:\n raise ValueError(\n \"Separator pattern should not already exist in document id\"\n )\n feedback_id = ID_SEPARATOR.join(\n [str(message_id), document_id, str(document_rank)]\n )\n else:\n feedback_id = str(message_id)\n\n return unique_prefix + ID_SEPARATOR + feedback_id\n\n\ndef build_publish_ephemeral_message_id(\n original_question_ts: str,\n) -> str:\n return \"publish_ephemeral_message__\" + original_question_ts\n\n\ndef build_continue_in_web_ui_id(\n message_id: int,\n) -> str:\n unique_prefix = str(uuid.uuid4())[:10]\n return unique_prefix + ID_SEPARATOR + str(message_id)\n\n\ndef decompose_action_id(feedback_id: str) -> tuple[int, str | None, int | None]:\n \"\"\"Decompose into query_id, document_id, document_rank, see above function\"\"\"\n try:\n components = feedback_id.split(ID_SEPARATOR)\n if len(components) != 2 and len(components) != 4:\n raise ValueError(\"Feedback ID does not contain right number of elements\")\n\n if len(components) == 2:\n return int(components[-1]), None, None\n\n return int(components[1]), components[2], int(components[3])\n\n except Exception as e:\n logger.error(e)\n raise ValueError(\"Received invalid Feedback Identifier\")\n\n\ndef get_view_values(state_values: dict[str, Any]) -> dict[str, str]:\n \"\"\"Extract view values\n\n Args:\n state_values (dict): The Slack view-submission values\n\n Returns:\n dict: keys/values of the view state content\n \"\"\"\n view_values = {}\n for _, view_data in state_values.items():\n for k, v in view_data.items():\n if (\n \"selected_option\" in v\n and isinstance(v[\"selected_option\"], dict)\n and \"value\" in v[\"selected_option\"]\n ):\n view_values[k] = v[\"selected_option\"][\"value\"]\n elif \"selected_options\" in v and isinstance(v[\"selected_options\"], list):\n view_values[k] = [\n x[\"value\"] for x in v[\"selected_options\"] if \"value\" in x\n ]\n elif \"selected_date\" in v:\n view_values[k] = v[\"selected_date\"]\n elif \"value\" in v:\n view_values[k] = v[\"value\"]\n return view_values\n\n\ndef translate_vespa_highlight_to_slack(match_strs: list[str], used_chars: int) -> str:\n def _replace_highlight(s: str) -> str:\n s = re.sub(r\"(?<=[^\\s])(.*?)\", r\"\\1\", s)\n s = s.replace(\"\", \"*\").replace(\"\", \"*\")\n return s\n\n final_matches = [\n replace_whitespaces_w_space(_replace_highlight(match_str)).strip()\n for match_str in match_strs\n if match_str\n ]\n combined = \"... \".join(final_matches)\n\n # Slack introduces \"Show More\" after 300 on desktop which is ugly\n # But don't trim the message if there is still a highlight after 300 chars\n remaining = 300 - used_chars\n if len(combined) > remaining and \"*\" not in combined[remaining:]:\n combined = combined[: remaining - 3] + \"...\"\n\n return combined\n\n\ndef remove_slack_text_interactions(slack_str: str) -> str:\n slack_str = SlackTextCleaner.replace_tags_basic(slack_str)\n slack_str = SlackTextCleaner.replace_channels_basic(slack_str)\n slack_str = SlackTextCleaner.replace_special_mentions(slack_str)\n slack_str = SlackTextCleaner.replace_special_catchall(slack_str)\n slack_str = SlackTextCleaner.add_zero_width_whitespace_after_tag(slack_str)\n return slack_str\n\n\ndef get_channel_from_id(client: WebClient, channel_id: str) -> dict[str, Any]:\n response = client.conversations_info(channel=channel_id)\n response.validate()\n return response[\"channel\"]\n\n\ndef get_channel_name_from_id(\n client: WebClient, channel_id: str\n) -> tuple[str | None, bool]:\n try:\n channel_info = get_channel_from_id(client, channel_id)\n name = channel_info.get(\"name\")\n is_dm = any([channel_info.get(\"is_im\"), channel_info.get(\"is_mpim\")])\n return name, is_dm\n except SlackApiError as e:\n logger.exception(f\"Couldn't fetch channel name from id: {channel_id}\")\n raise e\n\n\ndef fetch_slack_user_ids_from_emails(\n user_emails: list[str], client: WebClient\n) -> tuple[list[str], list[str]]:\n user_ids: list[str] = []\n failed_to_find: list[str] = []\n for email in user_emails:\n try:\n user = client.users_lookupByEmail(email=email)\n user_ids.append(user.data[\"user\"][\"id\"]) # type: ignore\n except Exception:\n logger.error(f\"Was not able to find slack user by email: {email}\")\n failed_to_find.append(email)\n\n return user_ids, failed_to_find\n\n\ndef fetch_user_ids_from_groups(\n given_names: list[str], client: WebClient\n) -> tuple[list[str], list[str]]:\n user_ids: list[str] = []\n failed_to_find: list[str] = []\n try:\n response = client.usergroups_list()\n if not isinstance(response.data, dict):\n logger.error(\"Error fetching user groups\")\n return user_ids, given_names\n\n all_group_data = response.data.get(\"usergroups\", [])\n name_id_map = {d[\"name\"]: d[\"id\"] for d in all_group_data}\n handle_id_map = {d[\"handle\"]: d[\"id\"] for d in all_group_data}\n for given_name in given_names:\n group_id = name_id_map.get(given_name) or handle_id_map.get(\n given_name.lstrip(\"@\")\n )\n if not group_id:\n failed_to_find.append(given_name)\n continue\n try:\n response = client.usergroups_users_list(usergroup=group_id)\n if isinstance(response.data, dict):\n user_ids.extend(response.data.get(\"users\", []))\n else:\n failed_to_find.append(given_name)\n except Exception as e:\n logger.error(f\"Error fetching user group ids: {str(e)}\")\n failed_to_find.append(given_name)\n except Exception as e:\n logger.error(f\"Error fetching user groups: {str(e)}\")\n failed_to_find = given_names\n\n return user_ids, failed_to_find\n\n\ndef fetch_group_ids_from_names(\n given_names: list[str], client: WebClient\n) -> tuple[list[str], list[str]]:\n group_data: list[str] = []\n failed_to_find: list[str] = []\n\n try:\n response = client.usergroups_list()\n if not isinstance(response.data, dict):\n logger.error(\"Error fetching user groups\")\n return group_data, given_names\n\n all_group_data = response.data.get(\"usergroups\", [])\n\n name_id_map = {d[\"name\"]: d[\"id\"] for d in all_group_data}\n handle_id_map = {d[\"handle\"]: d[\"id\"] for d in all_group_data}\n\n for given_name in given_names:\n id = handle_id_map.get(given_name.lstrip(\"@\"))\n id = id or name_id_map.get(given_name)\n if id:\n group_data.append(id)\n else:\n failed_to_find.append(given_name)\n except Exception as e:\n failed_to_find = given_names\n logger.error(f\"Error fetching user groups: {str(e)}\")\n\n return group_data, failed_to_find\n\n\ndef fetch_user_semantic_id_from_id(\n user_id: str | None, client: WebClient\n) -> str | None:\n if not user_id:\n return None\n\n response = client.users_info(user=user_id)\n if not response[\"ok\"]:\n return None\n\n user: dict = cast(dict[Any, dict], response.data).get(\"user\", {})\n\n return (\n user.get(\"real_name\")\n or user.get(\"name\")\n or user.get(\"profile\", {}).get(\"email\")\n )\n\n\ndef read_slack_thread(\n tenant_id: str, channel: str, thread: str, client: WebClient\n) -> list[ThreadMessage]:\n thread_messages: list[ThreadMessage] = []\n response = client.conversations_replies(channel=channel, ts=thread)\n replies = cast(dict, response.data).get(\"messages\", [])\n for reply in replies:\n if \"user\" in reply and \"bot_id\" not in reply:\n message = reply[\"text\"]\n user_sem_id = (\n fetch_user_semantic_id_from_id(reply.get(\"user\"), client)\n or \"Unknown User\"\n )\n message_type = MessageType.USER\n else:\n blocks: Any\n is_onyx_bot_response = False\n\n reply_user = reply.get(\"user\")\n reply_bot_id = reply.get(\"bot_id\")\n\n self_slack_bot_user_id, self_slack_bot_bot_id = get_onyx_bot_auth_ids(\n tenant_id, client\n )\n if reply_user is not None and reply_user == self_slack_bot_user_id:\n is_onyx_bot_response = True\n\n if reply_bot_id is not None and reply_bot_id == self_slack_bot_bot_id:\n is_onyx_bot_response = True\n\n if is_onyx_bot_response:\n # OnyxBot response\n message_type = MessageType.ASSISTANT\n user_sem_id = \"Assistant\"\n\n # OnyxBot responses have both text and blocks\n # The useful content is in the blocks, specifically the first block unless there are\n # auto-detected filters\n blocks = reply.get(\"blocks\")\n if not blocks:\n logger.warning(f\"OnyxBot response has no blocks: {reply}\")\n continue\n\n message = blocks[0].get(\"text\", {}).get(\"text\")\n\n # If auto-detected filters are on, use the second block for the actual answer\n # The first block is the auto-detected filters\n if message is not None and message.startswith(\"_Filters\"):\n if len(blocks) < 2:\n logger.warning(f\"Only filter blocks found: {reply}\")\n continue\n # This is the OnyxBot answer format, if there is a change to how we respond,\n # this will need to be updated to get the correct \"answer\" portion\n message = reply[\"blocks\"][1].get(\"text\", {}).get(\"text\")\n else:\n # Other bots are not counted as the LLM response which only comes from Onyx\n message_type = MessageType.USER\n bot_user_name = fetch_user_semantic_id_from_id(\n reply.get(\"user\"), client\n )\n user_sem_id = bot_user_name or \"Unknown\" + \" Bot\"\n\n # For other bots, just use the text as we have no way of knowing that the\n # useful portion is\n message = reply.get(\"text\")\n if not message:\n message = blocks[0].get(\"text\", {}).get(\"text\")\n\n if not message:\n logger.warning(\"Skipping Slack thread message, no text found\")\n continue\n\n message = remove_onyx_bot_tag(tenant_id, message, client=client)\n thread_messages.append(\n ThreadMessage(message=message, sender=user_sem_id, role=message_type)\n )\n\n return thread_messages\n\n\ndef slack_usage_report(action: str, sender_id: str | None, client: WebClient) -> None:\n if DISABLE_TELEMETRY:\n return\n\n onyx_user = None\n sender_email = None\n try:\n sender_email = client.users_info(user=sender_id).data[\"user\"][\"profile\"][\"email\"] # type: ignore\n except Exception:\n logger.warning(\"Unable to find sender email\")\n\n if sender_email is not None:\n with get_session_with_current_tenant() as db_session:\n onyx_user = get_user_by_email(email=sender_email, db_session=db_session)\n\n optional_telemetry(\n record_type=RecordType.USAGE,\n data={\"action\": action},\n user_id=str(onyx_user.id) if onyx_user else \"Non-Onyx-Or-No-Auth-User\",\n )\n\n\nclass SlackRateLimiter:\n def __init__(self) -> None:\n self.max_qpm: int | None = ONYX_BOT_MAX_QPM\n self.max_wait_time = ONYX_BOT_MAX_WAIT_TIME\n self.active_question = 0\n self.last_reset_time = time.time()\n self.waiting_questions: list[int] = []\n\n def refill(self) -> None:\n # If elapsed time is greater than the period, reset the active question count\n if (time.time() - self.last_reset_time) > 60:\n self.active_question = 0\n self.last_reset_time = time.time()\n\n def notify(\n self, client: WebClient, channel: str, position: int, thread_ts: str | None\n ) -> None:\n respond_in_thread_or_channel(\n client=client,\n channel=channel,\n receiver_ids=None,\n text=f\"Your question has been queued. You are in position {position}.\\nPlease wait a moment :hourglass_flowing_sand:\",\n thread_ts=thread_ts,\n )\n\n def is_available(self) -> bool:\n if self.max_qpm is None:\n return True\n\n self.refill()\n return self.active_question < self.max_qpm\n\n def acquire_slot(self) -> None:\n self.active_question += 1\n\n def init_waiter(self) -> tuple[int, int]:\n func_randid = random.getrandbits(128)\n self.waiting_questions.append(func_randid)\n position = self.waiting_questions.index(func_randid) + 1\n\n return func_randid, position\n\n def waiter(self, func_randid: int) -> None:\n if self.max_qpm is None:\n return\n\n wait_time = 0\n while (\n self.active_question >= self.max_qpm\n or self.waiting_questions[0] != func_randid\n ):\n if wait_time > self.max_wait_time:\n raise TimeoutError\n time.sleep(2)\n wait_time += 2\n self.refill()\n\n del self.waiting_questions[0]\n\n\ndef get_feedback_visibility() -> FeedbackVisibility:\n try:\n return FeedbackVisibility(ONYX_BOT_FEEDBACK_VISIBILITY.lower())\n except ValueError:\n return FeedbackVisibility.PRIVATE\n\n\nclass TenantSocketModeClient(SocketModeClient):\n def __init__(self, tenant_id: str, slack_bot_id: int, *args: Any, **kwargs: Any):\n super().__init__(*args, **kwargs)\n self._tenant_id = tenant_id\n self.slack_bot_id = slack_bot_id\n self.bot_name: str = \"Unnamed\"\n\n @contextmanager\n def _set_tenant_context(self) -> Generator[None, None, None]:\n token = None\n try:\n if self._tenant_id:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(self._tenant_id)\n yield\n finally:\n if token:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n def enqueue_message(self, message: str) -> None:\n with self._set_tenant_context():\n super().enqueue_message(message)\n\n def process_message(self) -> None:\n with self._set_tenant_context():\n super().process_message()\n\n def run_message_listeners(self, message: dict, raw_message: str) -> None:\n with self._set_tenant_context():\n super().run_message_listeners(message, raw_message)\n" }, { "path": "backend/onyx/prompts/__init__.py", "content": "" }, { "path": "backend/onyx/prompts/basic_memory.py", "content": "# ruff: noqa: E501, W605 start\n\n# Note that the user_basic_information is only included if we have at least 1 of the following: user_name, user_email, user_role\n# This is included because sometimes we need to know the user's name or basic info to best generate the memory.\nFULL_MEMORY_UPDATE_PROMPT = \"\"\"\nYou are a memory update agent that helps the user add or update memories. You are given a list of existing memories and a new memory to add. \\\nJust as context, you are also given the last few user messages from the conversation which generated the new memory. You must determine if the memory is brand new or if it is related to an existing memory. \\\nIf the new memory is an update to an existing memory or contradicts an existing memory, it should be treated as an update and you should reference the existing memory by memory_id (see below). \\\nThe memory should omit the user's name and direct reference to the user - for example, a memory like \"Yuhong prefers dark mode.\" should be modified to \"Prefers dark mode.\" (if the user's name is Yuhong).\n\n# Truncated chat history\n{chat_history}{user_basic_information}\n\n# User's existing memories\n{existing_memories}\n\n# New memory the user wants to insert\n{new_memory}\n\n# Response Style\nYou MUST respond in a json which follows the following format and keys:\n```json\n{{\n \"operation\": \"add or update\",\n \"memory_id\": \"if the operation is update, the id of the memory to update, otherwise null\",\n \"memory_text\": \"the text of the memory to add or update\"\n}}\n```\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n\nMEMORY_USER_BASIC_INFORMATION_PROMPT = \"\"\"\n\n# User Basic Information\nUser name: {user_name}\nUser email: {user_email}\nUser role: {user_role}\n\"\"\"\n" }, { "path": "backend/onyx/prompts/chat_prompts.py", "content": "# ruff: noqa: E501, W605 start\n\nfrom onyx.prompts.constants import REMINDER_TAG_NO_HEADER\n\n\nDATETIME_REPLACEMENT_PAT = \"{{CURRENT_DATETIME}}\"\nCITATION_GUIDANCE_REPLACEMENT_PAT = \"{{CITATION_GUIDANCE}}\"\nREMINDER_TAG_REPLACEMENT_PAT = \"{{REMINDER_TAG_DESCRIPTION}}\"\n\n\n# Note this uses a string pattern replacement so the user can also include it in their custom prompts. Keeps the replacement logic simple\n# This is editable by the user in the admin UI.\n# The first line is intended to help guide the general feel/behavior of the system.\nDEFAULT_SYSTEM_PROMPT = f\"\"\"\nYou are an expert assistant who is truthful, nuanced, insightful, and efficient. \\\nYour goal is to deeply understand the user's intent, think step-by-step through complex problems, provide clear and accurate answers, and proactively anticipate helpful follow-up information. \\\nWhenever there is any ambiguity around the user's query (or more information would be helpful), you use available tools (if any) to get more context.\n\nThe current date is {DATETIME_REPLACEMENT_PAT}.{CITATION_GUIDANCE_REPLACEMENT_PAT}\n\n# Response Style\nYou use different text styles, bolding, emojis (sparingly), block quotes, and other formatting to make your responses more readable and engaging.\nYou use proper Markdown and LaTeX to format your responses for math, scientific, and chemical formulas, symbols, etc.: '$$\\\\n[expression]\\\\n$$' for standalone cases and '\\\\( [expression] \\\\)' when inline.\nFor code you prefer to use Markdown and specify the language.\nYou can use horizontal rules (---) to separate sections of your responses.\nYou can use Markdown tables to format your responses for data, lists, and other structured information.\n\n{REMINDER_TAG_REPLACEMENT_PAT}\n\"\"\".lstrip()\n\n\nCOMPANY_NAME_BLOCK = \"\"\"\nThe user is at an organization called `{company_name}`.\n\"\"\"\n\nCOMPANY_DESCRIPTION_BLOCK = \"\"\"\nOrganization description: {company_description}\n\"\"\"\n\n# This is added to the system prompt prior to the tools section and is applied only if search tools have been run\nREQUIRE_CITATION_GUIDANCE = \"\"\"\n\nCRITICAL: If referencing knowledge from searches, cite relevant statements INLINE using the format [1], [2], [3], etc. to reference the \"document\" field. \\\nDO NOT provide any links following the citations. Cite inline as opposed to leaving all citations until the very end of the response.\n\"\"\"\n\n\n# Reminder message if any search tool has been run anytime in the chat turn\nCITATION_REMINDER = \"\"\"\nRemember to provide inline citations in the format [1], [2], [3], etc. based on the \"document\" field of the documents.\n\"\"\".strip()\n\nLAST_CYCLE_CITATION_REMINDER = \"\"\"\nYou are on your last cycle and no longer have any tool calls available. You must answer the query now to the best of your ability.\n\"\"\".strip()\n\n\n# Reminder message that replaces the usual reminder if web_search was the last tool call\nOPEN_URL_REMINDER = \"\"\"\nRemember that after using web_search, you are encouraged to open some pages to get more context unless the query is completely answered by the snippets.\nOpen the pages that look the most promising and high quality by calling the open_url tool with an array of URLs. Open as many as you want.\n\nIf you do have enough to answer, remember to provide INLINE citations using the \"document\" field in the format [1], [2], [3], etc.\n\"\"\".strip()\n\n\nIMAGE_GEN_REMINDER = \"\"\"\nVery briefly describe the image(s) generated. Do not include any links or attachments.\n\"\"\".strip()\n\n\nFILE_REMINDER = \"\"\"\nYour code execution generated file(s) with download links.\nIf you reference or share these files, use the exact markdown format [filename](file_link) with the file_link from the execution result.\n\"\"\".strip()\n\n\n# Specifically for OpenAI models, this prefix needs to be in place for the model to output markdown and correct styling\nCODE_BLOCK_MARKDOWN = \"Formatting re-enabled. \"\n\n# This is just for Slack context today\nADDITIONAL_CONTEXT_PROMPT = \"\"\"\nHere is some additional context which may be relevant to the user query:\n\n{additional_context}\n\"\"\".strip()\n\n\nTOOL_CALL_RESPONSE_CROSS_MESSAGE = \"\"\"\nThis tool call completed but the results are no longer accessible.\n\"\"\".strip()\n\n# This is used to add the current date and time to the prompt in the case where the Agent should be aware of the current\n# date and time but the replacement pattern is not present in the prompt.\nADDITIONAL_INFO = \"\\n\\nAdditional Information:\\n\\t- {datetime_info}.\"\n\n\nCHAT_NAMING_SYSTEM_PROMPT = f\"\"\"\nGiven the conversation history, provide a SHORT name for the conversation. Focus the name on the important keywords to convey the topic of the conversation. \\\nMake sure the name is in the same language as the user's first message.\n\n{REMINDER_TAG_NO_HEADER}\n\nIMPORTANT: DO NOT OUTPUT ANYTHING ASIDE FROM THE NAME. MAKE IT AS CONCISE AS POSSIBLE. NEVER USE MORE THAN 5 WORDS, LESS IS FINE.\n\"\"\".strip()\n\n\nCHAT_NAMING_REMINDER = \"\"\"\nProvide a short name for the conversation. Refer to other messages in the conversation (not including this one) to determine the language of the name.\n\nIMPORTANT: DO NOT OUTPUT ANYTHING ASIDE FROM THE NAME. MAKE IT AS CONCISE AS POSSIBLE. NEVER USE MORE THAN 5 WORDS, LESS IS FINE.\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/chat_tools.py", "content": "# These prompts are to support tool calling. Currently not used in the main flow or via any configs\n# The current generation of LLM is too unreliable for this task.\n# Onyx retrieval call as a tool option\nDANSWER_TOOL_NAME = \"Current Search\"\nDANSWER_TOOL_DESCRIPTION = \"A search tool that can find information on any topic including up to date and proprietary knowledge.\"\n\n\n# Tool calling format inspired from LangChain\nTOOL_TEMPLATE = \"\"\"\nTOOLS\n------\nYou can use tools to look up information that may be helpful in answering the user's \\\noriginal question. The available tools are:\n\n{tool_overviews}\n\nRESPONSE FORMAT INSTRUCTIONS\n----------------------------\nWhen responding to me, please output a response in one of two formats:\n\n**Option 1:**\nUse this if you want to use a tool. Markdown code snippet formatted in the following schema:\n\n```json\n{{\n \"action\": string, \\\\ The action to take. {tool_names}\n \"action_input\": string \\\\ The input to the action\n}}\n```\n\n**Option #2:**\nUse this if you want to respond directly to the user. Markdown code snippet formatted in the following schema:\n\n```json\n{{\n \"action\": \"Final Answer\",\n \"action_input\": string \\\\ You should put what you want to return to use here\n}}\n```\n\"\"\"\n\n# For the case where the user has not configured any tools to call, but still using the tool-flow\n# expected format\nTOOL_LESS_PROMPT = \"\"\"\nRespond with a markdown code snippet in the following schema:\n\n```json\n{{\n \"action\": \"Final Answer\",\n \"action_input\": string \\\\ You should put what you want to return to use here\n}}\n```\n\"\"\"\n\n\n# Second part of the prompt to include the user query\nUSER_INPUT = \"\"\"\nUSER'S INPUT\n--------------------\nHere is the user's input \\\n(remember to respond with a markdown code snippet of a json blob with a single action, and NOTHING else):\n\n{user_input}\n\"\"\"\n\n\n# After the tool call, this is the following message to get a final answer\n# Tools are not chained currently, the system must provide an answer after calling a tool\nTOOL_FOLLOWUP = \"\"\"\nTOOL RESPONSE:\n---------------------\n{tool_output}\n\nUSER'S INPUT\n--------------------\nOkay, so what is the response to my last comment? If using information obtained from the tools you must \\\nmention it explicitly without mentioning the tool names - I have forgotten all TOOL RESPONSES!\nIf the tool response is not useful, ignore it completely.\n{optional_reminder}{hint}\nIMPORTANT! You MUST respond with a markdown code snippet of a json blob with a single action, and NOTHING else.\n\"\"\"\n\n\n# If no tools were used, but retrieval is enabled, then follow up with this message to get the final answer\nTOOL_LESS_FOLLOWUP = \"\"\"\nRefer to the following documents when responding to my final query. Ignore any documents that are not relevant.\n\nCONTEXT DOCUMENTS:\n---------------------\n{context_str}\n\nFINAL QUERY:\n--------------------\n{user_query}\n\n{hint_text}\n\"\"\"\n" }, { "path": "backend/onyx/prompts/compression_prompts.py", "content": "# Prompts for chat history compression via summarization.\n\n# ruff: noqa: E501, W605 start\n# Cutoff marker helps the LLM focus on summarizing only messages before this point.\n# This improves \"needle in haystack\" accuracy by explicitly marking where to stop with an exact pattern which is also placed in locations easily attended to by the LLM (last user message and system prompt).\nCONTEXT_CUTOFF_START_MARKER = \"\"\nCONTEXT_CUTOFF_END_MARKER = \"\"\n\nSUMMARIZATION_CUTOFF_MARKER = f\"{CONTEXT_CUTOFF_START_MARKER} Stop summarizing the rest of the conversation past this point. {CONTEXT_CUTOFF_END_MARKER}\"\n\nSUMMARIZATION_PROMPT = f\"\"\"\nYou are a summarization system. Your task is to produce a detailed and accurate summary of a chat conversation up to a specified cutoff message. The cutoff will be marked by the string {CONTEXT_CUTOFF_START_MARKER}. \\\nIMPORTANT: Do not explicitly mention anything about the cutoff in your response. Do not situate the summary with respect to the cutoff. The context cutoff is only a system injected marker.\n\n# Guidelines\n- Only consider messages that occur at or before the cutoff point. Use the messages after it purely as context without including any of it in the summary.\n- Preserve factual correctness and intent; do not infer or speculate.\n- The summary should be information dense and detailed.\n- The summary should be in paragraph format and long enough to capture all of the most prominent details.\n\n# Focus on\n- Key topics discussed.\n- Decisions made, tools used, and conclusions reached.\n- Open questions or unresolved items.\n- Important constraints, preferences, or assumptions stated.\n- Omit small talk, repetition, and stylistic filler unless it affects meaning.\n\"\"\".strip()\n\nPROGRESSIVE_SUMMARY_SYSTEM_PROMPT_BLOCK = \"\"\"\n\n# Existing summary\nThere is a previous summary of the conversation. Build on top of this when constructing the new overall summary of the conversation:\n{previous_summary}\n\"\"\".rstrip()\n\nUSER_REMINDER = f\"Help summarize the conversation up to the cutoff point (do not mention anything related to the cutoff directly in your response). It should be a long form summary of the conversation up to the cutoff point as marked by {CONTEXT_CUTOFF_START_MARKER}. Be thorough.\"\n\nPROGRESSIVE_USER_REMINDER = f\"Update the existing summary by incorporating the new messages up to the cutoff point as marked by {CONTEXT_CUTOFF_START_MARKER} (do not mention anything related to the cutoff directly in your response). Be thorough and maintain the long form summary format.\"\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/constants.py", "content": "# ruff: noqa: E501, W605 start\nCODE_BLOCK_PAT = \"```\\n{}\\n```\"\nTRIPLE_BACKTICK = \"```\"\nSYSTEM_REMINDER_TAG_OPEN = \"\"\nSYSTEM_REMINDER_TAG_CLOSE = \"\"\n\n# Tags format inspired by Anthropic and OpenCode\nREMINDER_TAG_NO_HEADER = f\"\"\"\nUser messages may include {SYSTEM_REMINDER_TAG_OPEN} and {SYSTEM_REMINDER_TAG_CLOSE} tags. These {SYSTEM_REMINDER_TAG_OPEN} tags contain useful information and reminders. \\\nThey are automatically added by the system and are not actual user inputs. Behave in accordance to these instructions if relevant, and continue normally if they are not.\n\"\"\".strip()\n\nREMINDER_TAG_DESCRIPTION = f\"\"\"\n# System Reminders\n{REMINDER_TAG_NO_HEADER}\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/contextual_retrieval.py", "content": "# NOTE: the prompt separation is partially done for efficiency; previously I tried\n# to do it all in one prompt with sequential format() calls but this will cause a backend\n# error when the document contains any {} as python will expect the {} to be filled by\n# format() arguments\n\n# ruff: noqa: E501, W605 start\nCONTEXTUAL_RAG_PROMPT1 = \"\"\"\n{document}\n\nHere is the chunk we want to situate within the whole document\"\"\"\n\nCONTEXTUAL_RAG_PROMPT2 = \"\"\"\n{chunk}\n\nPlease give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.\n\"\"\".rstrip()\n\nCONTEXTUAL_RAG_TOKEN_ESTIMATE = 64 # 19 + 45\n\nDOCUMENT_SUMMARY_PROMPT = \"\"\"\n{document}\n\nPlease give a short succinct summary of the entire document. Answer only with the succinct summary and nothing else.\n\"\"\".rstrip()\n\nDOCUMENT_SUMMARY_TOKEN_ESTIMATE = 50\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/deep_research/__init__.py", "content": "" }, { "path": "backend/onyx/prompts/deep_research/dr_tool_prompts.py", "content": "GENERATE_PLAN_TOOL_NAME = \"generate_plan\"\n\n\nGENERATE_REPORT_TOOL_NAME = \"generate_report\"\n\n\nRESEARCH_AGENT_TOOL_NAME = \"research_agent\"\n\n\n# This is to ensure that even the non-reasoning models can have an ok time with this more complex flow.\nTHINK_TOOL_NAME = \"think_tool\"\n\n\n# ruff: noqa: E501, W605 start\n\n# Hard for the open_url tool to be called for a ton of search results all at once so limit to 3\nWEB_SEARCH_TOOL_DESCRIPTION = \"\"\"\n\n## web_search\nUse the `web_search` tool to get search results from the web. You should use this tool to get context for your research. These should be optimized for search engines like Google. \\\nUse concise and specific queries and avoid merging multiple queries into one. You can call web_search with multiple queries at once (3 max) but generally only do this when there is a clear opportunity for parallel searching. \\\nIf you use multiple queries, ensure that the queries are related in topic but not similar such that the results would be redundant.\n\"\"\"\n\n# This one is mostly similar to the one for the main flow but there won't be any user specified URLs to open.\nOPEN_URLS_TOOL_DESCRIPTION = f\"\"\"\n\n## open_urls\nUse the `open_urls` tool to read the content of one or more URLs. Use this tool to access the contents of the most promising web pages from your searches. \\\nYou can open many URLs at once by passing multiple URLs in the array if multiple pages seem promising. Prioritize the most promising pages and reputable sources. \\\nYou should almost always use open_urls after a web_search call and sometimes after reasoning with the {THINK_TOOL_NAME} tool.\n\"\"\"\n\nOPEN_URLS_TOOL_DESCRIPTION_REASONING = \"\"\"\n\n## open_urls\nUse the `open_urls` tool to read the content of one or more URLs. Use this tool to access the contents of the most promising web pages from your searches. \\\nYou can open many URLs at once by passing multiple URLs in the array if multiple pages seem promising. Prioritize the most promising pages and reputable sources. \\\nYou should almost always use open_urls after a web_search call.\n\"\"\"\n\n# NOTE: Internal search tool uses the same description as the default flow, not duplicating here.\n\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/deep_research/orchestration_layer.py", "content": "from onyx.prompts.deep_research.dr_tool_prompts import GENERATE_PLAN_TOOL_NAME\nfrom onyx.prompts.deep_research.dr_tool_prompts import GENERATE_REPORT_TOOL_NAME\nfrom onyx.prompts.deep_research.dr_tool_prompts import RESEARCH_AGENT_TOOL_NAME\nfrom onyx.prompts.deep_research.dr_tool_prompts import THINK_TOOL_NAME\n\n\n# ruff: noqa: E501, W605 start\nCLARIFICATION_PROMPT = f\"\"\"\nYou are a clarification agent that runs prior to deep research. Assess whether you need to ask clarifying questions, or if the user has already provided enough information for you to start research. \\\nCRITICAL - Never directly answer the user's query, you must only ask clarifying questions or call the `{GENERATE_PLAN_TOOL_NAME}` tool.\n\nIf the user query is already very detailed or lengthy (more than 3 sentences), do not ask for clarification and instead call the `{GENERATE_PLAN_TOOL_NAME}` tool.\n\nFor context, the date is {{current_datetime}}.\n\nBe conversational and friendly, prefer saying \"could you\" rather than \"I need\" etc.\n\nIf you need to ask questions, follow these guidelines:\n- Be concise and do not ask more than 5 questions.\n- If there are ambiguous terms or questions, ask the user to clarify.\n- Your questions should be a numbered list for clarity.\n- Respond in the same language as the user's query.\n- Make sure to gather all the information needed to carry out the research task in a concise, well-structured manner.{{internal_search_clarification_guidance}}\n- Wrap up with a quick sentence on what the clarification will help with, it's ok to reference the user query closely here.\n\"\"\".strip()\n\n\nINTERNAL_SEARCH_CLARIFICATION_GUIDANCE = \"\"\"\n- The deep research system is connected with organization internal document search and web search capabilities. In cases where it is unclear which source is more appropriate, ask the user to clarify.\n\"\"\"\n\n# Here there is a bit of combating model behavior which during alignment may be overly tuned to be cautious about access to data and feasibility.\n# Sometimes the model will just apologize and claim the task is not possible, hence the long section following CRITICAL.\nRESEARCH_PLAN_PROMPT = \"\"\"\nYou are a research planner agent that generates the high level approach for deep research on a user query. Analyze the query carefully and break it down into main concepts and areas of exploration. \\\nStick closely to the user query and stay on topic but be curious and avoid duplicate or overlapping exploration directions. \\\nBe sure to take into account the time sensitive aspects of the research topic and make sure to emphasize up to date information where appropriate. \\\nFocus on providing thorough research of the user's query over being helpful.\n\nCRITICAL - You MUST only output the research plan for the deep research flow and nothing else, you are not responding to the user. \\\nDo not worry about the feasibility of the plan or access to data or tools, a different deep research flow will handle that.\n\nFor context, the date is {current_datetime}.\n\nThe research plan should be formatted as a numbered list of steps and have 6 or less individual steps.\n\nEach step should be a standalone exploration question or topic that can be researched independently but may build on previous steps. The plan should be in the same language as the user's query.\n\nOutput only the numbered list of steps with no additional prefix or suffix.\n\"\"\".strip()\n\n\n# Specifically for some models, it really struggles to not just answer the user when there are questions about internal knowledge.\n# A reminder (specifically the fact that it's also a User type message) helps to prevent this.\nRESEARCH_PLAN_REMINDER = \"\"\"\nRemember to only output the research plan and nothing else. Do not worry about the feasibility of the plan or data access.\n\nYour response must only be a numbered list of steps with no additional prefix or suffix.\n\"\"\".strip()\n\n\nORCHESTRATOR_PROMPT = f\"\"\"\nYou are an orchestrator agent for deep research. Your job is to conduct research by calling the {RESEARCH_AGENT_TOOL_NAME} tool with high level research tasks. \\\nThis delegates the lower level research work to the {RESEARCH_AGENT_TOOL_NAME} which will provide back the results of the research.\n\nFor context, the date is {{current_datetime}}.\n\nBefore calling {GENERATE_REPORT_TOOL_NAME}, reason to double check that all aspects of the user's query have been well researched and that all key topics around the plan have been researched. \\\nThere are cases where new discoveries from research may lead to a deviation from the original research plan.\nIn these cases, ensure that the new directions are thoroughly investigated prior to calling {GENERATE_REPORT_TOOL_NAME}.\n\nNEVER output normal response tokens, you must only call tools.\n\n# Tools\nYou have currently used {{current_cycle_count}} of {{max_cycles}} max research cycles. You do not need to use all cycles.\n\n## {RESEARCH_AGENT_TOOL_NAME}\nThe research task provided to the {RESEARCH_AGENT_TOOL_NAME} should be reasonably high level with a clear direction for investigation. \\\nIt should not be a single short query, rather it should be 1 (or 2 if necessary) descriptive sentences that outline the direction of the investigation. \\\nThe research task should be in the same language as the overall research plan.\n\nCRITICAL - the {RESEARCH_AGENT_TOOL_NAME} only receives the task and has no additional context about the user's query, research plan, other research agents, or message history. \\\nYou absolutely must provide all of the context needed to complete the task in the argument to the {RESEARCH_AGENT_TOOL_NAME}.{{internal_search_research_task_guidance}}\n\nYou should call the {RESEARCH_AGENT_TOOL_NAME} MANY times before completing with the {GENERATE_REPORT_TOOL_NAME} tool.\n\nYou are encouraged to call the {RESEARCH_AGENT_TOOL_NAME} in parallel if the research tasks are not dependent on each other, which is typically the case. NEVER call more than 3 {RESEARCH_AGENT_TOOL_NAME} calls in parallel.\n\n## {GENERATE_REPORT_TOOL_NAME}\nYou should call the {GENERATE_REPORT_TOOL_NAME} tool if any of the following conditions are met:\n- You have researched all of the relevant topics of the research plan.\n- You have shifted away from the original research plan and believe that you are done.\n- You have all of the information needed to thoroughly answer all aspects of the user's query.\n- The last research cycle yielded minimal new information and future cycles are unlikely to yield more information.\n\n## {THINK_TOOL_NAME}\nCRITICAL - use the {THINK_TOOL_NAME} to reason between every call to the {RESEARCH_AGENT_TOOL_NAME} and before calling {GENERATE_REPORT_TOOL_NAME}. You should treat this as chain-of-thought reasoning to think deeply on what to do next. \\\nBe curious, identify knowledge gaps and consider new potential directions of research. Use paragraph format, do not use bullet points or lists.\n\nNEVER use the {THINK_TOOL_NAME} in parallel with other {RESEARCH_AGENT_TOOL_NAME} or {GENERATE_REPORT_TOOL_NAME}.\n\nBefore calling {GENERATE_REPORT_TOOL_NAME}, double check that all aspects of the user's query have been researched and that all key topics around the plan have been researched (unless you have gone in a different direction).\n\n# Research Plan\n{{research_plan}}\n\"\"\".strip()\n\n\nINTERNAL_SEARCH_RESEARCH_TASK_GUIDANCE = \"\"\"\n If necessary, clarify if the research agent should focus mostly on organization internal searches, web searches, or a combination of both. If the task doesn't require a clear priority, don't add sourcing guidance.\n\"\"\".strip(\n \"\\n\"\n)\n\n\nUSER_ORCHESTRATOR_PROMPT = \"\"\"\nRemember to refer to the system prompt and follow how to use the tools. Call the {THINK_TOOL_NAME} between every call to the {RESEARCH_AGENT_TOOL_NAME} and before calling {GENERATE_REPORT_TOOL_NAME}. Never run more than 3 {RESEARCH_AGENT_TOOL_NAME} calls in parallel.\n\nDon't mention this reminder or underlying details about the system.\n\"\"\".strip()\n\n\nFINAL_REPORT_PROMPT = \"\"\"\nYou are the final answer generator for a deep research task. Your job is to produce a thorough, balanced, and comprehensive answer on the research question provided by the user. \\\nYou have access to high-quality, diverse sources collected by secondary research agents as well as their analysis of the sources.\n\nIMPORTANT - You get straight to the point, never providing a title and avoiding lengthy introductions/preambles.\n\nFor context, the date is {current_datetime}.\n\nUsers have explicitly selected the deep research mode and will expect a long and detailed answer. It is ok and encouraged that your response is several pages long. \\\nStructure your response logically into relevant sections. You may find it helpful to reference the research plan to help structure your response but do not limit yourself to what is contained in the plan.\n\nYou use different text styles and formatting to make the response easier to read. You may use markdown rarely when necessary to make the response more digestible.\n\nProvide inline citations in the format [1], [2], [3], etc. based on the citations included by the research agents.\n\"\"\".strip()\n\n\nUSER_FINAL_REPORT_QUERY = f\"\"\"\nThe original research plan is included below (use it as a helpful reference but do not limit yourself to this):\n```\n{{research_plan}}\n```\n\nBased on all of the context provided in the research history, provide a comprehensive, well structured, and insightful answer to the user's previous query. \\\nCRITICAL: be extremely thorough in your response and address all relevant aspects of the query.\n\nIgnore the format styles of the intermediate {RESEARCH_AGENT_TOOL_NAME} reports, those are not end user facing and different from your task.\n\nProvide inline citations in the format [1], [2], [3], etc. based on the citations included by the research agents. The citations should be just a number in a bracket, nothing additional.\n\"\"\".strip()\n\n\n# Reasoning Model Variants of the prompts\nORCHESTRATOR_PROMPT_REASONING = f\"\"\"\nYou are an orchestrator agent for deep research. Your job is to conduct research by calling the {RESEARCH_AGENT_TOOL_NAME} tool with high level research tasks. \\\nThis delegates the lower level research work to the {RESEARCH_AGENT_TOOL_NAME} which will provide back the results of the research.\n\nFor context, the date is {{current_datetime}}.\n\nBefore calling {GENERATE_REPORT_TOOL_NAME}, reason to double check that all aspects of the user's query have been well researched and that all key topics around the plan have been researched.\nThere are cases where new discoveries from research may lead to a deviation from the original research plan. In these cases, ensure that the new directions are thoroughly investigated prior to calling {GENERATE_REPORT_TOOL_NAME}.\n\nBetween calls, think deeply on what to do next. Be curious, identify knowledge gaps and consider new potential directions of research. Use paragraph format for your reasoning, do not use bullet points or lists.\n\nNEVER output normal response tokens, you must only call tools.\n\n# Tools\nYou have currently used {{current_cycle_count}} of {{max_cycles}} max research cycles. You do not need to use all cycles.\n\n## {RESEARCH_AGENT_TOOL_NAME}\nThe research task provided to the {RESEARCH_AGENT_TOOL_NAME} should be reasonably high level with a clear direction for investigation. \\\nIt should not be a single short query, rather it should be 1 (or 2 if necessary) descriptive sentences that outline the direction of the investigation. \\\nThe research task should be in the same language as the overall research plan.\n\nCRITICAL - the {RESEARCH_AGENT_TOOL_NAME} only receives the task and has no additional context about the user's query, research plan, or message history. \\\nYou absolutely must provide all of the context needed to complete the task in the argument to the {RESEARCH_AGENT_TOOL_NAME}.{{internal_search_research_task_guidance}}\n\nYou should call the {RESEARCH_AGENT_TOOL_NAME} MANY times before completing with the {GENERATE_REPORT_TOOL_NAME} tool.\n\nYou are encouraged to call the {RESEARCH_AGENT_TOOL_NAME} in parallel if the research tasks are not dependent on each other, which is typically the case. NEVER call more than 3 {RESEARCH_AGENT_TOOL_NAME} calls in parallel.\n\n## {GENERATE_REPORT_TOOL_NAME}\nYou should call the {GENERATE_REPORT_TOOL_NAME} tool if any of the following conditions are met:\n- You have researched all of the relevant topics of the research plan.\n- You have shifted away from the original research plan and believe that you are done.\n- You have all of the information needed to thoroughly answer all aspects of the user's query.\n- The last research cycle yielded minimal new information and future cycles are unlikely to yield more information.\n\n# Research Plan\n{{research_plan}}\n\"\"\".strip()\n\n\nUSER_ORCHESTRATOR_PROMPT_REASONING = \"\"\"\nRemember to refer to the system prompt and follow how to use the tools. \\\nYou are encouraged to call the {RESEARCH_AGENT_TOOL_NAME} in parallel when the research tasks are not dependent on each other, but never call more than 3 {RESEARCH_AGENT_TOOL_NAME} calls in parallel.\n\nDon't mention this reminder or underlying details about the system.\n\"\"\".strip()\n\n\n# Only for the first cycle, we encourage the model to research more, since it is unlikely that it has already addressed all parts of the plan at this point.\nFIRST_CYCLE_REMINDER_TOKENS = 100\nFIRST_CYCLE_REMINDER = \"\"\"\nMake sure all parts of the user question and the plan have been thoroughly explored before calling generate_report. If new interesting angles have been revealed from the research, you may deviate from the plan to research new directions.\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/deep_research/research_agent.py", "content": "from onyx.prompts.deep_research.dr_tool_prompts import GENERATE_REPORT_TOOL_NAME\nfrom onyx.prompts.deep_research.dr_tool_prompts import THINK_TOOL_NAME\n\n\nMAX_RESEARCH_CYCLES = 8\n\n# ruff: noqa: E501, W605 start\nRESEARCH_AGENT_PROMPT = f\"\"\"\nYou are a highly capable, thoughtful, and precise research agent that conducts research on a specific topic. Prefer being thorough in research over being helpful. Be curious but stay strictly on topic. \\\nYou iteratively call the tools available to you including {{available_tools}} until you have completed your research at which point you call the {GENERATE_REPORT_TOOL_NAME} tool.\n\nNEVER output normal response tokens, you must only call tools.\n\nFor context, the date is {{current_datetime}}.\n\n# Tools\nYou have a limited number of cycles to complete your research and you do not have to use all cycles. You are on cycle {{current_cycle_count}} of {MAX_RESEARCH_CYCLES}.\\\n{{optional_internal_search_tool_description}}\\\n{{optional_web_search_tool_description}}\\\n{{optional_open_url_tool_description}}\n## {THINK_TOOL_NAME}\nCRITICAL - use the think tool after every set of searches and reads (so search, read some pages, then think and repeat). \\\nYou MUST use the {THINK_TOOL_NAME} before calling the web_search tool for all calls to web_search except for the first call. \\\nUse the {THINK_TOOL_NAME} before calling the {GENERATE_REPORT_TOOL_NAME} tool.\n\nAfter a set of searches + reads, use the {THINK_TOOL_NAME} to analyze the results and plan the next steps.\n- Reflect on the key information found with relation to the task.\n- Reason thoroughly about what could be missing, the knowledge gaps, and what queries might address them, \\\nor why there is enough information to answer the research task comprehensively.\n\n## {GENERATE_REPORT_TOOL_NAME}\nOnce you have completed your research, call the `{GENERATE_REPORT_TOOL_NAME}` tool. \\\nYou should only call this tool after you have fully researched the topic. \\\nConsider other potential areas of research and weigh that against the materials already gathered before calling this tool.\n\"\"\".strip()\n\n\nRESEARCH_REPORT_PROMPT = \"\"\"\nYou are a highly capable and precise research sub-agent that has conducted research on a specific topic. \\\nYour job is now to organize the findings to return a comprehensive report that preserves all relevant statements and information that has been gathered in the existing messages. \\\nThe report will be seen by another agent instead of a user so keep it free of formatting or commentary and instead focus on the facts only. \\\nDo not give it a title, do not break it down into sections, and do not provide any of your own conclusions/analysis.\n\nYou may see a list of tool calls in the history but you do not have access to tools anymore. You should only use the information in the history to create the report.\n\nCRITICAL - This report should be as long as necessary to return ALL of the information that the researcher has gathered. It should be several pages long so as to capture as much detail as possible from the research. \\\nIt cannot be stressed enough that this report must be EXTREMELY THOROUGH and COMPREHENSIVE. Only this report is going to be returned, so it's CRUCIAL that you don't lose any details from the raw messages.\n\nRemove any obviously irrelevant or duplicative information.\n\nIf a statement seems not trustworthy or is contradictory to other statements, it is important to flag it.\n\nWrite the report in the same language as the provided task.\n\nCite all sources INLINE using the format [1], [2], [3], etc. based on the `document` field of the source. \\\nCite inline as opposed to leaving all citations until the very end of the response.\n\"\"\"\n\n\nUSER_REPORT_QUERY = \"\"\"\nPlease write me a comprehensive report on the research topic given the context above. As a reminder, the original topic was:\n{research_topic}\n\nRemember to include AS MUCH INFORMATION AS POSSIBLE and as faithful to the original sources as possible. \\\nKeep it free of formatting and focus on the facts only. Be sure to include all context for each fact to avoid misinterpretation or misattribution. \\\nRespond in the same language as the topic provided above.\n\nCite every fact INLINE using the format [1], [2], [3], etc. based on the `document` field of the source.\n\nCRITICAL - BE EXTREMELY THOROUGH AND COMPREHENSIVE, YOUR RESPONSE SHOULD BE SEVERAL PAGES LONG.\n\"\"\"\n\n\n# Reasoning Model Variants of the prompts\nRESEARCH_AGENT_PROMPT_REASONING = f\"\"\"\nYou are a highly capable, thoughtful, and precise research agent that conducts research on a specific topic. Prefer being thorough in research over being helpful. Be curious but stay strictly on topic. \\\nYou iteratively call the tools available to you including {{available_tools}} until you have completed your research at which point you call the {GENERATE_REPORT_TOOL_NAME} tool. Between calls, think about the results of the previous tool call and plan the next steps. \\\nReason thoroughly about what could be missing, identify knowledge gaps, and what queries might address them. Or consider why there is enough information to answer the research task comprehensively.\n\nOnce you have completed your research, call the `{GENERATE_REPORT_TOOL_NAME}` tool.\n\nNEVER output normal response tokens, you must only call tools.\n\nFor context, the date is {{current_datetime}}.\n\n# Tools\nYou have a limited number of cycles to complete your research and you do not have to use all cycles. You are on cycle {{current_cycle_count}} of {MAX_RESEARCH_CYCLES}.\\\n{{optional_internal_search_tool_description}}\\\n{{optional_web_search_tool_description}}\\\n{{optional_open_url_tool_description}}\n## {GENERATE_REPORT_TOOL_NAME}\nOnce you have completed your research, call the `{GENERATE_REPORT_TOOL_NAME}` tool. You should only call this tool after you have fully researched the topic.\n\"\"\".strip()\n\n\nOPEN_URL_REMINDER_RESEARCH_AGENT = \"\"\"\nRemember that after using web_search, you are encouraged to open some pages to get more context unless the query is completely answered by the snippets.\nOpen the pages that look the most promising and high quality by calling the open_url tool with an array of URLs.\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/federated_search.py", "content": "from onyx.configs.app_configs import MAX_SLACK_QUERY_EXPANSIONS\n\nSLACK_QUERY_EXPANSION_PROMPT = f\"\"\"\nRewrite the user's query into at most {MAX_SLACK_QUERY_EXPANSIONS} keyword-only queries for Slack's keyword search.\n\nSlack search behavior:\n- Pure keyword AND search (no semantics)\n- More words = fewer matches, so keep queries concise (1-3 words)\n\nALWAYS include:\n- Person names (e.g., \"Sarah Chen\", \"Mike Johnson\") - people search for messages from/about specific people\n- Project/product names, technical terms, proper nouns\n- Actual content words: \"performance\", \"bug\", \"deployment\", \"API\", \"error\"\n\nDO NOT include:\n- Meta-words: \"topics\", \"conversations\", \"discussed\", \"summary\", \"messages\"\n- Temporal: \"today\", \"yesterday\", \"week\", \"month\", \"recent\", \"last\"\n- Channel names: \"general\", \"eng-general\", \"random\"\n\nExamples:\n\nQuery: \"what are the big topics in eng-general this week?\"\nOutput:\n\nQuery: \"messages with Sarah about the deployment\"\nOutput:\nSarah deployment\nSarah\ndeployment\n\nQuery: \"what did Mike say about the budget?\"\nOutput:\nMike budget\nMike\nbudget\n\nQuery: \"performance issues in eng-general\"\nOutput:\nperformance issues\nperformance\nissues\n\nQuery: \"what did we discuss about the API migration?\"\nOutput:\nAPI migration\nAPI\nmigration\n\nNow process this query:\n\n{{query}}\n\nOutput (keywords only, one per line, NO explanations or commentary):\n\"\"\"\n\nSLACK_DATE_EXTRACTION_PROMPT = \"\"\"\nExtract the date range from the user's query and return it in a structured format.\n\nCurrent date context:\n- Today: {today}\n- Current time: {current_time}\n\nGuidelines:\n1. Return a JSON object with \"days_back\" (integer) indicating how many days back to search\n2. If no date/time is mentioned, return {{\"days_back\": null}}\n3. Interpret relative dates accurately:\n - \"today\" or \"today's\" = 0 days back\n - \"yesterday\" = 1 day back\n - \"last week\" = 7 days back\n - \"last month\" = 30 days back\n - \"last X days\" = X days back\n - \"past X days\" = X days back\n - \"this week\" = 7 days back\n - \"this month\" = 30 days back\n4. For creative expressions, interpret intent:\n - \"recent\" = 7 days back\n - \"recently\" = 7 days back\n - \"lately\" = 14 days back\n5. Always be conservative - if uncertain, use a longer time range\n\nUser query: {query}\n\nReturn ONLY a valid JSON object in this format: {{\"days_back\": }}\nNothing else.\n\"\"\"\n" }, { "path": "backend/onyx/prompts/filter_extration.py", "content": "# The following prompts are used for extracting filters to apply along with the query in the\n# document index. For example, a filter for dates or a filter by source type such as GitHub\n# or Slack\nSOURCES_KEY = \"sources\"\n\n# Smaller followup prompts in time_filter.py\nTIME_FILTER_PROMPT = \"\"\"\nYou are a tool to identify time filters to apply to a user query for a downstream search \\\napplication. The downstream application is able to use a recency bias or apply a hard cutoff to \\\nremove all documents before the cutoff. Identify the correct filters to apply for the user query.\n\nThe current day and time is {current_day_time_str}.\n\nAlways answer with ONLY a json which contains the keys \"filter_type\", \"filter_value\", \\\n\"value_multiple\" and \"date\".\n\nThe valid values for \"filter_type\" are \"hard cutoff\", \"favors recent\", or \"not time sensitive\".\nThe valid values for \"filter_value\" are \"day\", \"week\", \"month\", \"quarter\", \"half\", or \"year\".\nThe valid values for \"value_multiple\" is any number.\nThe valid values for \"date\" is a date in format MM/DD/YYYY, ALWAYS follow this format.\n\"\"\".strip()\n\n\n# Smaller followup prompts in source_filter.py\n# Known issue: LLMs like GPT-3.5 try to generalize. If the valid sources contains \"web\" but not\n# \"confluence\" and the user asks for confluence related things, the LLM will select \"web\" since\n# confluence is accessed as a website. This cannot be fixed without also reducing the capability\n# to match things like repository->github, website->web, etc.\n# This is generally not a big issue though as if the company has confluence, hopefully they add\n# a connector for it or the user is aware that confluence has not been added.\nSOURCE_FILTER_PROMPT = f\"\"\"\nGiven a user query, extract relevant source filters for use in a downstream search tool.\nRespond with a json containing the source filters or null if no specific sources are referenced.\nONLY extract sources when the user is explicitly limiting the scope of where information is \\\ncoming from.\nThe user may provide invalid source filters, ignore those.\n\nThe valid sources are:\n{{valid_sources}}\n{{web_source_warning}}\n{{file_source_warning}}\n\n\nALWAYS answer with ONLY a json with the key \"{SOURCES_KEY}\". \\\nThe value for \"{SOURCES_KEY}\" must be null or a list of valid sources.\n\nSample Response:\n{{sample_response}}\n\"\"\".strip()\n\nWEB_SOURCE_WARNING = \"\"\"\nNote: The \"web\" source only applies to when the user specifies \"website\" in the query. \\\nIt does not apply to tools such as Confluence, GitHub, etc. that have a website.\n\"\"\".strip()\n\nFILE_SOURCE_WARNING = \"\"\"\nNote: The \"file\" source only applies to when the user refers to uploaded files in the query.\n\"\"\".strip()\n\n\n# Use the following for easy viewing of prompts\nif __name__ == \"__main__\":\n print(TIME_FILTER_PROMPT)\n print(\"------------------\")\n print(SOURCE_FILTER_PROMPT)\n" }, { "path": "backend/onyx/prompts/image_analysis.py", "content": "# Used for creating embeddings of images for vector search\nDEFAULT_IMAGE_SUMMARIZATION_SYSTEM_PROMPT = \"\"\"\nYou are an assistant for summarizing images for retrieval.\nSummarize the content of the following image and be as precise as possible.\nThe summary will be embedded and used to retrieve the original image.\nTherefore, write a concise summary of the image that is optimized for retrieval.\n\"\"\"\n\n# Prompt for generating image descriptions with filename context\nDEFAULT_IMAGE_SUMMARIZATION_USER_PROMPT = \"\"\"\nDescribe precisely and concisely what the image shows.\n\"\"\"\n\n\n# Used for analyzing images in response to user queries at search time\nDEFAULT_IMAGE_ANALYSIS_SYSTEM_PROMPT = (\n \"You are an AI assistant specialized in describing images.\\n\"\n \"You will receive a user question plus an image URL. Provide a concise textual answer.\\n\"\n \"Focus on aspects of the image that are relevant to the user's question.\\n\"\n \"Be specific and detailed about visual elements that directly address the query.\\n\"\n)\n" }, { "path": "backend/onyx/prompts/kg_prompts.py", "content": "# Standards\nSEPARATOR_LINE = \"-------\"\nSEPARATOR_LINE_LONG = \"---------------\"\nNO_EXTRACTION = \"No extraction of knowledge graph objects was feasible.\"\nYES = \"yes\"\nNO = \"no\"\n\n# Framing/Support/Template Prompts\nENTITY_TYPE_SETTING_PROMPT = f\"\"\"\n{SEPARATOR_LINE}\n{{entity_types}}\n{SEPARATOR_LINE}\n\"\"\".strip()\n\nRELATIONSHIP_TYPE_SETTING_PROMPT = f\"\"\"\nHere are the types of relationships:\n{SEPARATOR_LINE}\n{{relationship_types}}\n{SEPARATOR_LINE}\n\"\"\".strip()\n\nEXTRACTION_FORMATTING_PROMPT = r\"\"\"\n{{\"entities\": [::' (please use that capitalization). If allowed options \\\nare provided above, you can only extract those types of entities! Again, there should be an 'Other' \\\noption. Pick this if none of the others apply.>],\n\"relationships\": [::__\\\n__::'>],\n\"terms\": ['>]\n}}\n\"\"\".strip()\n\nQUERY_ENTITY_EXTRACTION_FORMATTING_PROMPT = r\"\"\"\n{{\"entities\": [::' (please use that capitalization)>. Each entity \\\nalso should be followed by a list of comma-separated attribute filters for the entity, if referred to in the \\\nquestion for that entity. CRITICAL: you can only use attributes that are mentioned above for the \\\nentity type in question. Example: 'ACCOUNT::* -- [account_type: customer, status: active]' if the question is \\\n'list all customer accounts', and ACCOUNT was an entity type with these attribute key/values allowed.] \\\n\"time_filter\": \n}}\n\"\"\".strip()\n\nQUERY_RELATIONSHIP_EXTRACTION_FORMATTING_PROMPT = r\"\"\"\n{{\"relationships\": [::__\\\n__::'>]\n}}\n\"\"\".strip()\n\nEXAMPLE_1 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::Nike\", \"CONCERN::*\"],\n \"relationships\": [\"ACCOUNT::Nike__had__CONCERN::*\"], \"terms\": []}}\n\"\"\".strip()\n\nEXAMPLE_2 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::Nike\", \"CONCERN::performance\"],\n \"relationships\": [\"ACCOUNT::*__had_issues__CONCERN::performance\"], \"terms\": [\"performance issue\"]}}\n\"\"\".strip()\n\nEXAMPLE_3 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::Nike\", \"CONCERN::performance\", \"CONCERN::user_experience\"],\n \"relationships\": [\"ACCOUNT::Nike__had__CONCERN::performance\",\n \"ACCOUNT::Nike__solved__CONCERN::user_experience\"],\n \"terms\": [\"performance\", \"user experience\"]}}\n\"\"\".strip()\n\nEXAMPLE_4 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::Nike\", \"FEATURE::dashboard\", \"CONCERN::performance\"],\n \"relationships\": [\"ACCOUNT::Nike__had__CONCERN::performance\",\n \"ACCOUNT::Nike__had_issues__FEATURE::dashboard\",\n \"ACCOUNT::NIKE__gets_value_from__FEATURE::dashboard\"],\n \"terms\": [\"value\", \"performance\"]}}\n\"\"\".strip()\n\nRELATIONSHIP_EXAMPLE_1 = r\"\"\"\n'Which issues did Nike report?' and the extracted entities were found to be:\n\n \"ACCOUNT::Nike\", \"CONCERN::*\"\n\nthen a valid relationship extraction could be:\n\n{{\"relationships\": [\"ACCOUNT::Nike__had__CONCERN::*\"]}}\n\"\"\".strip()\n\nRELATIONSHIP_EXAMPLE_2 = r\"\"\"\n'Did Nike say anything about performance issues?' and the extracted entities were found to be:\n\n\"ACCOUNT::Nike\", \"CONCERN::performance\"\n\nthen a much more suitable relationship extraction could be:\n{{\"relationships\": [\"ACCOUNT::*__had_issues__CONCERN::performance\"]}}\n\"\"\".strip()\n\nRELATIONSHIP_EXAMPLE_3 = r\"\"\"\n'Did Nike report some performance issues with our solution? And were they happy that the user experience issue got solved?', \\\nand the extracted entities were found to be:\n\n\"ACCOUNT::Nike\", \"CONCERN::performance\", \"CONCERN::user_experience\"\n\nthen a valid relationship extraction could be:\n\n{{\"relationships\": [\"ACCOUNT::Nike__had__CONCERN::performance\",\n \"ACCOUNT::Nike__solved__CONCERN::user_experience\"]}}\n\"\"\".strip()\n\nRELATIONSHIP_EXAMPLE_4 = r\"\"\"\n'Nike reported some performance issues with our dashboard solution, but do they think it delivers great value nevertheless?' \\\nand the extracted entities were found to be:\n\n\"ACCOUNT::Nike\", \"FEATURE::dashboard\", \"CONCERN::performance\"\n\nthen a valid relationship extraction could be:\nExample 4:\n\n{{\"relationships\": [\"ACCOUNT::Nike__had__CONCERN::performance\",\n \"ACCOUNT::Nike__had_issues__FEATURE::dashboard\",\n \"ACCOUNT::NIKE__gets_value_from__FEATURE::dashboard\"]}}\n\nExplanation:\n - Nike did report performance concerns\n - Nike had problems with the dashboard, which is a feature\n - We are interested in the value relationship between Nike and the dashboard feature\n\n\"\"\".strip()\n\nRELATIONSHIP_EXAMPLE_5 = r\"\"\"\n'In which emails did Nike discuss their issues with the dashboard?' \\\nand the extracted entities were found to be:\n\n\"ACCOUNT::Nike\", \"FEATURE::dashboard\", \"EMAIL::*\"\n\nthen a valid relationship extraction could be:\n\n{{\"relationships\": [\"ACCOUNT::Nike__had__CONCERN::*\",\n \"ACCOUNT::Nike__had_issues__FEATURE::dashboard\",\n \"ACCOUNT::NIKE__in__EMAIL::*\",\n \"EMAIL::*__discusses__FEATURE::dashboard\",\n \"EMAIL::*Nike__had__CONCERN::* \"]}}\nExplanation:\n - Nike did report unspecified concerns\n - Nike had problems with the dashboard, which is a feature\n - We are interested in emails that Nike exchanged with us\n\"\"\".strip()\n\nRELATIONSHIP_EXAMPLE_6 = r\"\"\"\n'List the last 5 emails that Lisa exchanged with Nike:' \\\nand the extracted entities were found to be:\n\n\"ACCOUNT::Nike\", \"EMAIL::*\", \"EMPLOYEE::Lisa\"\n\nthen a valid relationship extraction could be:\n\n{{\"relationships\": [\"ACCOUNT::Nike__had__CONCERN::*\",\n \"ACCOUNT::Nike__had_issues__FEATURE::dashboard\",\n \"ACCOUNT::NIKE__in__EMAIL::*\"]}}\nExplanation:\n - Nike did report unspecified concerns\n - Nike had problems with the dashboard, which is a feature\n - We are interested in emails that Nike exchanged with us\n\"\"\".strip()\n\n\nENTITY_EXAMPLE_1 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::Nike--[]\", \"CONCERN::*--[]\"]}}\n\"\"\".strip()\n\nENTITY_EXAMPLE_2 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::Nike--[]\", \"CONCERN::performance--[]\"]}}\n\"\"\".strip()\n\nENTITY_EXAMPLE_3 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::*--[]\", \"CONCERN::performance--[]\", \"CONCERN::user_experience--[]\"]}}\n\"\"\".strip()\n\nENTITY_EXAMPLE_4 = r\"\"\"\n{{\"entities\": [\"ACCOUNT::*--[]\", \"CONCERN::performance--[degree: severe]\"]}}\n\"\"\".strip()\n\nMASTER_EXTRACTION_PROMPT = f\"\"\"\nYou are an expert in the area of knowledge extraction in order to construct a knowledge graph. You are given a text \\\nand asked to extract entities, relationships, and terms from it that you can reliably identify.\n\nHere are the entity types that are available for extraction. Some of them may have a description, others \\\nshould be obvious. Also, for a given entity allowed options may be provided. If allowed options are provided, \\\nyou can only extract those types of entities! If no allowed options are provided, take your best guess.\n\nYou can ONLY extract entities of these types and relationships between objects of these types:\n{SEPARATOR_LINE}\n{ENTITY_TYPE_SETTING_PROMPT}\n{SEPARATOR_LINE}\nPlease format your answer in this format:\n{SEPARATOR_LINE}\n{EXTRACTION_FORMATTING_PROMPT}\n{SEPARATOR_LINE}\n\nThe list above here is the exclusive, only list of entities you can choose from!\n\nHere are some important additional instructions. (For the purpose of illustration, assume that ]\n \"ACCOUNT\", \"CONCERN\", and \"FEATURE\" are all in the list of entity types above, and shown actual \\\nentities fall into allowed options. Note that this \\\nis just assumed for these examples, but you MUST use only the entities above for the actual extraction!)\n\n- You can either extract specific entities if a specific entity is referred to, or you can refer to the entity type.\n* if the entity type is referred to in general, you would use '*' as the entity name in the extraction.\nAs an example, if the text would say:\n 'Nike reported that they had issues'\nthen a valid extraction could be:\nExample 1:\n{EXAMPLE_1}\n\n* If on the other hand the text would say:\n'Nike reported that they had performance issues'\nthen a much more suitable extraction could be:\nExample 2:\n{EXAMPLE_2}\n\n- You can extract multiple relationships between the same two entity types.\nAs an example, if the text would say:\n'Nike reported some performance issues with our solution, but they are very happy that the user experience issue got solved.'\nthen a valid extraction could be:\nExample 3:\n{EXAMPLE_3}\n\n- You can extract multiple relationships between the same two actual entities if you think that \\\nthere are multiple relationships between them based on the text.\nAs an example, if the text would say:\n'Nike reported some performance issues with our dashboard solution, but they think it delivers great value.'\nthen a valid extraction could be:\nExample 4:\n{EXAMPLE_4}\n\nNote that effectively a three-way relationship (Nike - performance issues - dashboard) extracted as two individual \\\nrelationships.\n\n- Again,\n - you should only extract entities belonging to the entity types above - but do extract all that you \\\ncan reliably identify in the text\n - use refer to 'all' entities in an entity type listed above by using '*' as the entity name\n - only extract important relationships that signify something non-trivial, expressing things like \\\nneeds, wants, likes, dislikes, plans, interests, lack of interests, problems the account is having, etc.\n - you MUST only use the initial list of entities provided! Ignore the entities in the examples unless \\\nthey are also part of the initial list of entities! This is essential!\n - only extract relationships between the entities extracted first!\n\n\n{SEPARATOR_LINE}\n\nHere is the text you are asked to extract knowledge from, if needed with additional information about any participants:\n{SEPARATOR_LINE}\n---content---\n{SEPARATOR_LINE}\n\"\"\".strip()\n\n\nQUERY_ENTITY_EXTRACTION_PROMPT = f\"\"\"\nYou are an expert in the area of knowledge extraction and using knowledge graphs. You are given a question \\\nand asked to extract entities (with attributes if applicable) that you can reliably identify, which will then\nbe matched with a known entity in the knowledge graph. You are also asked to extract time constraints information \\\nfrom the QUESTION. Some time constraints will be captured by entity attributes if \\\nthe entity type has a fitting attribute (example: 'created_at' could be a candidate for that), other times\nwe will extract an explicit time filter if no attribute fits. (Note regarding 'last', 'first', etc.: DO NOT \\\nimply the need for a time filter just because the question asks for something that is not the current date. \\\nThey will relate to ordering that we will handle separately later).\n\nIn case useful, today is ---today_date--- and the user asking is ---user_name---, which may or may not be relevant.\nHere are the entity types that are available for extraction. Some of them may have \\\na description, others should be obvious. Also, notice that some may have attributes associated with them, which will \\\nbe important later.\nYou can ONLY extract entities of these types:\n{SEPARATOR_LINE}\n{ENTITY_TYPE_SETTING_PROMPT}\n{SEPARATOR_LINE}\n\nThe list above here is the exclusive, only list of entities you can choose from!\n\nAlso, note that there are fixed relationship types between these entities. Please consider those \\\nas well so to make sure that you are not missing implicit entities! Implicit entities are often \\\nin verbs ('emailed to', 'talked to', ...). Also, they may be used to connect entities that are \\\nclearly in the question.\n\n{SEPARATOR_LINE}\n{RELATIONSHIP_TYPE_SETTING_PROMPT}\n{SEPARATOR_LINE}\n\nHere are some important additional instructions. (For the purpose of illustration, assume that \\\n \"ACCOUNT\", \"CONCERN\", \"EMAIL\", and \"FEATURE\" are all in the list of entity types above, and the \\\nattribute options for \"CONCERN\" include 'degree' with possible values that include 'severe'. Note that this \\\nis just assumed for these examples, but you MUST use only the entities above for the actual extraction!)\n\n- You can either extract specific entities if a specific entity is referred to, or you can refer to the entity type.\n* if the entity type is referred to in general, you would use '*' as the entity name in the extraction.\nAs an example, if the question would say:\n 'Which issues did Nike report?'\nthen a valid entity and term extraction could be:\nExample 1:\n{ENTITY_EXAMPLE_1}\n\n* If on the other hand the question would say:\n'Did Nike say anything about performance issues?'\nthen a much more suitable entity and term extraction could be:\nExample 2:\n{ENTITY_EXAMPLE_2}\n\n* Then, if the question is:\n'Who reported performance issues?'\nthen a suitable entity and term extraction could be:\nExample 3:\n{ENTITY_EXAMPLE_3}\n\n* Then, if we inquire about an entity with a specific attribute :\n'Who reported severe performance issues?'\nthen a suitable entity and term extraction could be:\nExample 3:\n{ENTITY_EXAMPLE_4}\n\n- Again,\n - you should only extract entities belonging to the entity types above - but do extract all that you \\\ncan reliably identify in the text\n - if you refer to all/any/an unspecified entity of an entity type listed above, use '*' as the entity name\n - similarly, if a specific entity type is referred to in general, you should use '*' as the entity name\n - you MUST only use the initial list of entities provided! Ignore the entities in the examples unless \\\nthey are also part of the initial list of entities! This is essential!\n - don't forget to provide answers also to the event filtering and whether documents need to be inspected!\n - 'who' often refers to individuals or accounts.\n - see whether any of the entities are supposed to be narrowed down by an attribute value. The precise attribute \\\nand the value would need to be taken from the specification, as the question may use different words and the \\\nactual attribute may be implied.\n - don't just look at the entities that are mentioned in the question but also those that the question \\\nmay be about.\n - be very careful that you only extract attributes that are listed above for the entity type in question! Do \\\nnot make up attributes even if they are implied! Particularly if there is a relationship type that would \\\nactually represent that information, you MUST not extract the information as an attribute. We \\\nwill extract the relationship type later.\n - For the values of attributes, look at the possible values above! For example 'open' may refer to \\\n'backlog', 'todo', 'in progress', etc. In cases like that construct a ';'-separated list of values that you think may fit \\\nwhat is implied in the question (in the exanple: 'open; backlog; todo; in progress').\n\nAlso, if you think the name or the title of an entity is given but name or title are not mentioned \\\nexplicitly as an attribute, then you should indeed extract the name/title as the entity name.\n\n{SEPARATOR_LINE}\n\nHere is the question you are asked to extract desired entities and time filters from:\n{SEPARATOR_LINE}\n---content---\n{SEPARATOR_LINE}\n\nPlease format your answer in this format:\n{SEPARATOR_LINE}\n{QUERY_ENTITY_EXTRACTION_FORMATTING_PROMPT}\n{SEPARATOR_LINE}\n\n\"\"\".strip()\n\n\nQUERY_RELATIONSHIP_EXTRACTION_PROMPT = f\"\"\"\nYou are an expert in the area of knowledge extraction and using knowledge graphs. You are given a question \\\nand previously you were asked to identify known entities in the question. Now you are asked to extract \\\nthe relationships between the entities you have identified earlier.\n\nFirst off as background, here are the entity types that are known to the system:\n{SEPARATOR_LINE}\n---entity_types---\n{SEPARATOR_LINE}\n\n\nHere are the entities you have identified earlier:\n{SEPARATOR_LINE}\n---identified_entities---\n{SEPARATOR_LINE}\n\nNote that the notation for the entities is ::.\n\nHere are the options for the relationship types(!) between the entities you have identified earlier \\\nas well as relationship types between the identified entities and other entities \\\nnot explicitly mentioned:\n{SEPARATOR_LINE}\n---relationship_type_options---\n{SEPARATOR_LINE}\n\nThese types are, if any were identified, formatted as \\\n____, and they \\\nlimit the allowed relationships that you can extract. You would then though use the actual full entities as in:\n\n::____::.\n\nNote: should be a word or two that captures the nature \\\nof the relationship. Common relationships may be: 'likes', 'dislikes', 'uses', 'is interested in', 'mentions', \\\n'addresses', 'participates in', etc., but look at the text to find the most appropriate relationship. \\\nUse spaces here for word separation.\n\nPlease format your answer in this format:\n{SEPARATOR_LINE}\n{QUERY_RELATIONSHIP_EXTRACTION_FORMATTING_PROMPT}\n{SEPARATOR_LINE}\n\nThe list above here is the exclusive, only list of entities and relationship types you can choose from!\n\nHere are some important additional instructions. (For the purpose of illustration, assume that ]\n \"ACCOUNT\", \"CONCERN\", and \"FEATURE\" are all in the list of entity types above. Note that this \\\nis just assumed for these examples, but you MUST use only the entities above for the actual extraction!)\n\n- You can either extract specific entities if a specific entity is referred to, or you can refer to the entity type.\n* if the entity type is referred to in general, you would use '*' as the entity name in the extraction.\n\nAs an example, if the question would say:\n\n{RELATIONSHIP_EXAMPLE_1}\n\n* If on the other hand the question would say:\n\n{RELATIONSHIP_EXAMPLE_2}\n\n- You can extract multiple relationships between the same two entity types.\nFor example 3, if the question would say:\n\n{RELATIONSHIP_EXAMPLE_3}\n\n- You can extract multiple relationships between the same two actual entities if you think that \\\nthere are multiple relationships between them based on the question.\nAs an example, if the question would say:\n\n{RELATIONSHIP_EXAMPLE_4}\n\nNote that effectively a three-way relationship (Nike - performance issues - dashboard) extracted as two individual \\\nrelationships.\n\n- Again,\n - you can only extract relationships between the entities extracted earlier\n - you can only extract the relationships that match the listed relationship types\n - if in doubt and there are multiple relationships between the same two entities, you can extract \\\nall of those that may fit with the question.\n - be really thinking through the question which type of relationships should be extracted and which should not.\n\nOther important notes:\n - For questions that really try to explore in general what a certain entity was involved in like 'what did Paul Smith do \\\nin the last 3 months?', and Paul Smith has been extracted i.e. as an entity of type 'EMPLOYEE', then you need to extract \\\nall of the possible relationships an empoyee Paul Smith could have.\n - You are not forced to use all or any of the relationship types listed above. Really look at the question to \\\n determine which relationships are explicitly or implicitly referred to in the question.\n\n{SEPARATOR_LINE}\n\nHere is the question you are asked to extract desired entities, relationships, and terms from:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\"\"\".strip()\n\n\nGENERAL_CHUNK_PREPROCESSING_PROMPT = \"\"\"\nThis is a part of a document that you need to extract information (entities, relationships) from.\n\nNote: when you extract relationships, please make sure that:\n - if you see a relationship for one of our employees, you should extract the relationship both for the employee AND \\\n VENDOR::{vendor}.\n - if you see a relationship for one of the representatives of other accounts, you should extract the relationship \\\nonly for the account ACCOUNT::!\n\n--\nAnd here is the content:\n{content}\n\"\"\".strip()\n\n\n### Source-specific prompts\n\nCALL_CHUNK_PREPROCESSING_PROMPT = \"\"\"\nThis is a call between employees of the VENDOR's company and representatives of one or more ACCOUNTs (usually one). \\\nWhen you extract information based on the instructions, please make sure that you properly attribute the information \\\nto the correct employee and account. \\\n\nHere are the participants (name component of email) from us ({vendor}):\n{participant_string}\n\nHere are the participants (name component of email) from the other account(s):\n{account_participant_string}\n\nIn the text it should be easy to associate a name with the email, and then with the account ('us' vs 'them'). If in doubt, \\\nlook at the context and try to identify whether the statement comes from the other account. If you are not sure, ignore.\n\nNote: when you extract relationships, please make sure that:\n - if you see a relationship for one of our employees, you should extract the relationship both for the employee AND \\\n VENDOR::{vendor}.\n - if you see a relationship for one of the representatives of other accounts, you should extract the relationship \\\nonly for the account ACCOUNT::!\n\n--\nAnd here is the content:\n{content}\n\"\"\".strip()\n\n\nCALL_DOCUMENT_CLASSIFICATION_PROMPT = \"\"\"\nThis is the beginning of a call between employees of the VENDOR's company ({vendor}) and other participants.\n\nYour task is to classify the call into one of the following categories:\n{category_options}\n\nPlease also consider the participants when you perform your classification task - they can be important indicators \\\nfor the category.\n\nPlease format your answer as a string in the format:\n\nREASONING: - CATEGORY: \n\n--\nAnd here is the beginning of the call, including title and participants:\n\n{beginning_of_call_content}\n\"\"\".strip()\n\n\nSTRATEGY_GENERATION_PROMPT = f\"\"\"\nNow you need to decide what type of strategy to use to answer a given question, how ultimately \\\nthe answer should be formatted to match the user's expectation, and what an appropriate question \\\nto/about 'one object or one set of objects' may be, should the answer logically benefit from a divide \\\nand conquer strategy, or it naturally relates to one or few individual objects. Also, you are \\\nsupposed to determine whether a divide and conquer strategy would be appropriate.\n\n\nHere are the entity types that are available in the knowledge graph:\n{SEPARATOR_LINE}\n---possible_entities---\n{SEPARATOR_LINE}\n\nHere are the relationship types that are available in the knowledge graph:\n{SEPARATOR_LINE}\n---possible_relationships---\n{SEPARATOR_LINE}\n\nHere is the question whose answer is ultimately sought:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nAnd here are the entities and relationships that have been extracted earlier from this question:\n{SEPARATOR_LINE}\n---entities---\n---relationships---\n{SEPARATOR_LINE}\n\nHere are more instructions:\n\na) Regarding the strategy, there are three aspects to it:\n\na1) \"Search Type\":\nShould the question be answered as a SEARCH ('filtered search'), or as a SQL ('SQL query search')?\n\nThe options are:\n1. SEARCH: A filtered search simply uses the entities and relationships that you extracted earlier and \\\napplies them as filters to search the underlying documents, which are properly indexed. Examples are \\\n'what did Nike say about the Analyzer product?', or 'what did I say in my calls with Nike about pricing?'. So this \\\nis used really when there is *no implicit or explicit constraint or requirements* on underlying source documents \\\noutside of filters, and there is no ordering, no limiting their number, etc. So use this for a question that \\\ntries to get information *across* documents which may be filtered by their related relationships and entities, but without \\\nother constraints.\n\n2. SQL: Choose this option if the question either requires counting of entities (e.g. 'how many calls...'), or \\\nif the query refers to specific entities that first need to be identified and then analyzed/searched/listed. \\\nExamples here are 'what did I say about pricing in my call with Nike last week?' (the specific call needs to \\\nbe identified first and then analyzed), \\\n'what are the next steps of our two largest opportunities?', or 'summarize my 3 most recent customer calls'. So \\\nthis is used if there *are implicit constraints* on the underlying source documents beyond filtering, including \\\nordering, limiting, etc. Use this also if the answer expects to analyze each source independently as part \\\nof the overall answer.\n\nNote:\n - here, you should look at the extracted entities and relationships and judge whether using them as filters \\\n(using an *and*) would be appropriate to identify the range of underlying sources, or whether more \\\ncalculations would be needed to find the underlying sources ('last 2...', etc.) .\n - It is also *critical* to look at the attributes of the entities! You only can use the given attributes (and their\n values, if given) as where conditions etc in a SQL statement. So if you think you would 'want\n to' have a where condition but there is not appropriate attribute, then you should not use the SQL strategy\n but the SEARCH strategy. (A Search can always look through data and see what is the best fit, SQL needs to\n be more specific.). On the other hand, if the question maps well to the entities and attributes, then\n SQL may be a good choice.\n - Likely, if there are questions 'about something', then this only is used in a SQL statement or a filter \\\n if it shows up as an entity or relationship in the extracted entities and relationships. Otherwise, it will \\\n be part of the analysis/search. not the document identification.\n - again, note that we can only FILTER (SEARCH) or COMPUTE (SQL) using the extracted entities (and their attributes)\n and relationships. \\\n So do not think that if there is another term in the question, it should be included in the SQL statement. \\\n It cannot.\n\n\na2) \"Search Strategy\":\nIf a SQL search is chosen, i.e., documents have to be identified first, there are two approaches:\n1. SIMPLE: You think you can answer the question using a database that is aware of the entities, relationships \\\nabove, and is generally suitable if it is enough to either list or count entities, return dates, etc. Usually, \\\n'SIMPLE' is chosen for questions of the form 'how many...' (always), or 'list the...' (often), 'when was...', \\\n'what did (someone) work on...'etc. Often it is also used in cases like 'what did John work on since April?'. Here, \\\nthe user would expect to just see the list. So chose 'SIMPLE' here unless there are REALLY CLEAR \\\nfollow-up instructions for each item (like 'summarize...' , 'analyze...', 'what are the main points of...'.) If \\\nit is a 'what did...'-type question, choose 'SIMPLE'!\n\n2. DEEP: You think you really should ALSO leverage the actual text of sources to answer the question, which sits \\\nin a vector database. Examples are 'what is discussed in...', 'summarize', 'what is the discussion about...',\\\n'how does... relate to...', 'are there any mentions of... in..', 'what are the main points in...', \\\n'what are the next steps...', etc. Those are usually questions 'about' \\\nthe entities retrieved from the knowledge graph, or questions about the underlying sources.\n\nYour task is to decide which of the two strategies to use.\n\na3) \"Relationship Detection\":\nYou need to evaluate whether the question involves any relationships between entities (of the same type) \\\nor between entities and relationships. Respond with 'RELATIONSHIPS' or 'NO_RELATIONSHIPS'.\n\nb) Regarding the format of the answer: there are also two types of formats available to you:\n\n1. LIST: The user would expect an answer as a bullet point list of objects, likely with text associated with each \\\nbullet point (or sub-bullet). This will be clearer once the data is available.\n2. TEXT: The user would expect the questions to be answered in text form.\n\nYour task is to decide which of the two formats to use.\n\n\nc) Regarding the broken down question for one object:\n\nAlways generate a broken_down_question if the question pertains ultimately to a specific objects, even if it seems to be \\\na singular object.\n\n- If the question is of type 'how many...', or similar, then imagine that the individual objects have been \\\nfound and you want to ask each object something that illustrates why/in what what that object relates to the \\\nquestion. (question: 'How many cars are fast?' -> broken_down_question: 'How fast is this car?')\n\n- Assume the answer would either i) best be generated by first analyzing one object at a time, then aggregating \\\nthe results, or ii) directly relates to one or few objects found through matching suitable criteria.\n\n- The key is to drop any filtering/criteria matching as the objects are already filtered by the criteria. Also, do not \\\ntry to verify here whether the object in question actually satisfies a filter criteria, but rather see \\\nwhat it says/does etc. In other words, use this to identify more details about the object, as it relates \\\nto the original question.\n(Example: question: 'What did our oil & gas customers say about the new product?' -> broken_down_question: \\\n'What did this customer say about the new product?',\nor:\nquestion: 'What was in the email from Frank?' -> broken_down_question: 'What is in this email?')\n\n\nd) Regarding the divide and conquer strategy:\n\nYou are supposed to decide whether a divide and conquer strategy would be appropriate. That means, do you think \\\nthat in order to answer the question, it would be good to first analyze one object at a time, and then aggregate the \\\nresults? Or should the information rather be analyzed as a whole? This would be 'yes' or 'no'.\n\nPlease answer in json format in this form:\n\n{{\n \"search_type\": ,\n \"search_strategy\": ,\n \"relationship_detection\": ,\n \"format\": ,\n \"broken_down_question\": ,\n \"divide_and_conquer\": \n}}\n\nDo not include any other text or explanations.\n\"\"\"\n\nSOURCE_DETECTION_PROMPT = f\"\"\"\nYou are an expert in generating, understanding and analyzing SQL statements.\n\nYou are given an original SQL statement that returns a list of entities from a table or \\\nan aggregation of entities from a table. Your task will be to \\\nidentify the source documents that are relevant to what the SQL statement is returning.\n\nThe task is actually quite simple. There are two tables involved - relationship_table and entity_table. \\\nrelationship_table was used to generate the original SQL statement. Again, returning entities \\\nor aggregations of entities. The second table, entity_table contains the entities and \\\nthe corresponding source_documents. All you need to do is to appropriately join the \\\nentity_table table on the entities that would be retrieved from the original SQL statement, \\\nand then return the source_documents from the entity_table table.\n\nFor your orientation, the relationship_table table has this structure:\n - Table name: relationship_table\n - Columns:\n - relationship (str): The name of the RELATIONSHIP, combining the nature of the relationship and the names of the entities. \\\nIt is of the form \\\n::____:: \\\n[example: ACCOUNT::Nike__has__CONCERN::performance]. Note that this is NOT UNIQUE!\n - source_entity (str): the id of the source ENTITY/NODE in the relationship [example: ACCOUNT::Nike]\n - source_entity_attributes (json): the attributes of the source entity/node [example: {{\"account_type\": \"customer\"}}]\n - target_entity (str): the id of the target ENTITY/NODE in the relationship [example: CONCERN::performance]\n - target_entity_attributes (json): the attributes of the target entity/node [example: {{\"degree\": \"severe\"}}]\n - source_entity_type (str): the type of the source entity/node [example: ACCOUNT]. Only the entity types provided \\\n below are valid.\n - target_entity_type (str): the type of the target entity/node [example: CONCERN]. Only the entity types provided \\\n below are valid.\n - relationship_type (str): the type of the relationship, formatted as \\\n____. So the explicit entity_names have \\\nbeen removed. [example: ACCOUNT__has__CONCERN]\n - source_date (str): the 'event' date of the source document [example: 2021-01-01]\n\nThe second table, entity_table, has this structure:\n - Table name: entity_table\n - Columns:\n - entity (str): The name of the ENTITY, which is unique in this table. source_entity and target_entity \\\nin the relationship_table table are the same as entity in this table.\n - source_document (str): the id of the document that contains the entity.\n\nAgain, ultimately, your task is to join the entity_table table on the entities that would be retrieved from the \\\noriginal SQL statement, and then return the source_documents from the entity_table table.\n\nThe way to do that is to create a common table expression for the original SQL statement and join the \\\nentity_table table suitably on the entities.\n\nHere is the *original* SQL statement:\n{SEPARATOR_LINE}\n---original_sql_statement---\n{SEPARATOR_LINE}\n\nPlease structure your answer using , ,, start and end tags as in:\n\n[think very briefly through the problem step by step, not more than 2-3 sentences] \\\n[the new SQL statement that returns the source documents involved in the original SQL statement]\n\"\"\".strip()\n\nENTITY_SOURCE_DETECTION_PROMPT = f\"\"\"\nYou are an expert in generating, understanding and analyzing SQL statements.\n\nYou are given a SQL statement that returned an aggregation of entities in a table. \\\nYour task will be to identify the source documents for the entities involved in \\\nthe answer. For example, should the original SQL statement be \\\n'SELECT COUNT(entity) FROM entity_table where entity_type = \"ACCOUNT\"' \\\nthen you should return the source documents that contain the entities of type 'ACCOUNT'.\n\nThe table has this structure:\n - Table name: entity_table\n - Columns:\n - entity (str): The name of the ENTITY, combining the nature of the entity and the id of the entity. \\\nIt is of the form :: [example: ACCOUNT::625482894].\n - entity_type (str): the type of the entity [example: ACCOUNT].\n - entity_attributes (json): the attributes of the entity [example: {{\"priority\": \"high\", \"status\": \"active\"}}]\n - source_document (str): the id of the document that contains the entity. Note that the combination of \\\nid_name and source_document IS UNIQUE!\n - source_date (timestamp): the 'event' date of the source document [example: 2025-04-25 21:43:31.054741+00]\n\nSpecifically, the table contains the 'source_document' column, which is the id of the source document that \\\ncontains the core information about the entity. Make sure that you do not return more documents, i.e. if there \\\nis a limit on source documents in the original SQL statement, the new SQL statement needs to have \\\nthe same limit.\n\nCRITICAL NOTES:\n - Only return source documents and nothing else!\n\nYour task is then to create a new SQL statement that returns the source documents that are relevant to what the \\\noriginal SQL statement is returning. So the source document of every row used in the original SQL statement should \\\nbe included in the result of the new SQL statement, and then you should apply a 'distinct'.\n\nHere is the *original* SQL statement:\n{SEPARATOR_LINE}\n---original_sql_statement---\n{SEPARATOR_LINE}\n\nPlease structure your answer using , ,, start and end tags as in:\n\n[think very briefly through the problem step by step, not more than 2-3 sentences] \\\n[the new SQL statement that returns the source documents involved in the original SQL statement]\n\"\"\".strip()\n\n\nENTITY_TABLE_DESCRIPTION = f\"\"\"\\\n - Table name: entity_table\n - Columns:\n - entity (str): The name of the ENTITY, combining the nature of the entity and the id of the entity. \\\nIt is of the form :: [example: ACCOUNT::625482894].\n - entity_type (str): the type of the entity [example: ACCOUNT].\n - entity_attributes (json): the attributes of the entity [example: {{\"priority\": \"high\", \"status\": \"active\"}}]\n - source_document (str): the id of the document that contains the entity. Note that the combination of \\\nid_name and source_document IS UNIQUE!\n - source_date (timestamp): the 'event' date of the source document [example: 2025-04-25 21:43:31.054741+00]\n\n{SEPARATOR_LINE}\n\nImportantly, here are the entity (node) types that you can use, with a short description of what they mean. You may need to \\\nidentify the proper entity type through its description. Also notice the allowed attributes for each entity type and \\\ntheir values, if provided. Of particular importance is the 'subtype' attribute, if provided, as this is how \\\nthe entity type may also often be referred to.\n{SEPARATOR_LINE}\n---entity_types---\n{SEPARATOR_LINE}\n\"\"\"\n\nRELATIONSHIP_TABLE_DESCRIPTION = f\"\"\"\\\n - Table name: relationship_table\n - Columns:\n - relationship (str): The name of the RELATIONSHIP, combining the nature of the relationship and the names of the entities. \\\nIt is of the form \\\n::____:: \\\n[example: ACCOUNT::Nike__has__CONCERN::performance]. Note that this is NOT UNIQUE!\n - source_entity (str): the id of the source ENTITY/NODE in the relationship [example: ACCOUNT::Nike]\n - source_entity_attributes (json): the attributes of the source entity/node [example: {{\"account_type\": \"customer\"}}]\n - target_entity (str): the id of the target ENTITY/NODE in the relationship [example: CONCERN::performance]\n - target_entity_attributes (json): the attributes of the target entity/node [example: {{\"degree\": \"severe\"}}]\n - source_entity_type (str): the type of the source entity/node [example: ACCOUNT]. Only the entity types provided \\\n below are valid.\n - target_entity_type (str): the type of the target entity/node [example: CONCERN]. Only the entity types provided \\\n below are valid.\n - relationship_type (str): the type of the relationship, formatted as \\\n____. So the explicit entity_names have \\\nbeen removed. [example: ACCOUNT__has__CONCERN]\n - source_document (str): the id of the document that contains the relationship. Note that the combination of \\\nid_name and source_document IS UNIQUE!\n - source_date (timestamp): the 'event' date of the source document [example: 2025-04-25 21:43:31.054741+00]\n\n{SEPARATOR_LINE}\n\nImportantly, here are the entity (node) types that you can use, with a short description of what they mean. You may need to \\\nidentify the proper entity type through its description. Also notice the allowed attributes for each entity type and \\\ntheir values, if provided. Of particular importance is the 'subtype' attribute, if provided, as this is how \\\nthe entity type may also often be referred to.\n{SEPARATOR_LINE}\n---entity_types---\n{SEPARATOR_LINE}\n\nHere are the relationship types that are in the table, denoted as ____.\nIn the table, the actual relationships are not quite of this form, but each is followed by '::' \\\nin the relationship id as shown above.\n{SEPARATOR_LINE}\n---relationship_types---\n{SEPARATOR_LINE}\n\"\"\"\n\n\nSIMPLE_SQL_PROMPT = f\"\"\"\nYou are an expert in generating a SQL statement that only uses ONE TABLE that captures RELATIONSHIPS \\\nbetween TWO ENTITIES. The table has the following structure:\n\n{SEPARATOR_LINE}\n{RELATIONSHIP_TABLE_DESCRIPTION}\n\nHere is the question you are supposed to translate into a SQL statement:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nTo help you, we already have identified the entities and relationships that the SQL statement likely *should* use (but note the \\\nexception below!). The entities also contain the list of attributes and attribute values that should specify the entity. \\\nThe format is ::--[:, \\\n:, ...].\n{SEPARATOR_LINE}\nIdentified entities with attributes in query:\n\n---query_entities_with_attributes---\n\nThese are the entities that should be used in the SQL statement. However, \\\nnote that these are the entities (with potential attributes) that were *matches* of Knowledge Graph identified with the \\\nentities originally identified in the original question. A such, they may have id names that may not mean much by themselves, \\\neg ACCOUNT::a74f332. Here is the mapping of entities originally identified (whose role in the query should be obvious) with \\\nthe entities that were matched to them in the Knowledge Graph:\n\n---entity_explanation_string---\n\n--\n\nHere are relationships that were identified as explicitly or implicitly referred to in the question:\n\n---query_relationships---\n\n(Again, if applicable, the entities contained in the relationships are the same as the entities in the \\\nquery_entities_with_attributes, and those are the correct ones to use in the SQL statement.)\n\n{SEPARATOR_LINE}\n\nCRITICAL SPECIAL CASE:\n - if an identified entity is of the form ::*, or an identified relationship contains an \\\nentity of this form, this refers to *any* entity of that type. Correspondingly, the SQL query should use the *entity type*, \\\nand possibly the relationship type, but not the entity with the * itself. \\\nExample: if you see 'ACCOUNT::*', that means any account matches. So if you are supposed to count the 'ACCOUNT::*', \\\nyou should count the entities of entity_type 'ACCOUNT'.\n\n\nIMPORTANT NOTES:\n- The id_name of each relationship has the format \\\n____.\n- The relationship id_names are NOT UNIQUE, only the combinations of relationship id_name and source_document_id are unique. \\\nThat is because each relationship is extracted from a document. So make sure you use the proper 'distinct's!\n- If the SQL contains a 'SELECT DISTINCT' clause and an ORDER BY clause, then you MUST include the columns from the ORDER BY \\\nclause ALSO IN THE SELECT DISTINCT CLAUSE! This is very important! (This is a postgres db., so this is a MUST!). \\\nYou MUST NOT have a column in the ORDER BY clause that is not ALSO in the SELECT DISTINCT clause!\n- If you join the relationship table on itself using the source_node or target_node, you need to make sure that you also \\\njoin on the source_document_id.\n- The id_name of each node/entity has the format ::, where 'entity_type_id_name' \\\nand 'name' are columns and \\\n the values and can be used for filtering.\n- The table can be joined on itself on source nodes and/or target nodes if needed.\n- the SQL statement MUST ultimately only return NODES/ENTITIES (not relationships!), or aggregations of \\\nentities/nodes(count, avg, max, min, etc.). \\\nAgain, DO NOT compose a SQL statement that returns id_name of relationships.\n- You CAN ONLY return ENTITIES or COUNTS (or other aggregations) of ENTITIES, or you can return \\\nsource_date (but only if the question asks for event dates or times). DO NOT return \\\nsource documents or counts of source documents, or relationships or counts of relationships! \\\nThose can only appear in where clauses, ordering etc., but they cannot be returned or ultimately \\\ncounted here! source_date and date operations can appear in select statements, particularly if \\\nthere is time ordering or grouping involved.\n- ENTITIES can be target_entity or source_entity. Think about the allowed relationships and the \\\nquestion to decide which one you want!\n- It is ok to generate nested SQL as long as it is correct postgres syntax!\n- Attributes are stored in the attributes json field. As this is postgres, querying for those must be done as \\\n\"attributes ->> '' = ''\".\n- The SELECT clause MUST only contain entities or aggregations/counts of entities, or, in cases the \\\nquestion was about dates or times, then it can also include source_date. But source_document MUST NEVER appear \\\nin the SELECT clause!\n- Again, NEVER count or retrieve source documents in SELECT CLAUSE, whether it is in combination with \\\nentities, with a distinct, etc. NO source_document in SELECT CLAUSE! So NEVER produce a \\\n'SELECT COUNT(source_entity, source_document)...'\n- Please think about whether you are interested in source entities or target entities! For that purpose, \\\nconsider the allowed relationship types to make sure you select or count the correct one!\n- Again, ALWAYS make sure that EACH COLUMN in an ORDER-BY clause IS ALSO IN THE SELECT CLAUSE! Remind yourself \\\nof that in the reasoning.\n- Be careful with dates! Often a date will refer to the source data, which is the date when \\\nan underlying piece of information was updated. However, if the attributes of an entity contain \\\ntime information as well (like 'started_at', 'completed_at', etc.), then you should really look at \\\nthe wording to see whether you should use a date in the attributes or the event date.\n- Dates are ALWAYS in string format of the form YYYY-MM-DD, for source date as well as for date-like the attributes! \\\nSo please use that format, particularly if you use data comparisons (>, <, ...)\n- Again, NO 'relationship' or 'source_document' in the SELECT CLAUSE, be it as direct columns are in aggregations!\n- Careful with SORT! Really think in which order you want to sort if you have multiple columns you \\\nwant 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 \\\nvariable as the first column to sort by.\n- When doing a SORT on an attribute value of an entity, you MUST also apply a WHERE clause to filter \\\nfor entities that have the attribute value set. For example, if you want to sort the target entity \\\nby the attribute 'created_date', you must also have a WHERE clause that checks whether the target \\\nentity attribute contains 'created_date'. This is vital for proper ordering with null values.\n- Usually, you will want to retrieve or count entities, maybe with attributes. But you almost always want to \\\nhave entities involved in the SELECT clause.\n- Questions like 'What did Paul work on last week?' should generally be handled by finding all entities \\\nthat reasonably relate to 'work entities' that are i) related to Paul, and ii) that were created or \\\nupdated (by him) last week. So this would likely be a UNION of multiple queries.\n- If you do joins consider the possibility that the second entity does not exist for all examples. \\\nTherefore joins should generally be LEFT joins (or RIGHT joins) as appropriate. Think about which \\\nentities you are interested in, and which ones provides attributes.\nAnother important note:\n - For questions that really try to explore what a certain entity was involved in like 'what did Paul Smith do \\\nin the last 3 months?', and Paul Smith has been extracted ie as an entity of type 'EMPLOYEE', you will \\\nwant to consider all entities that Paul Smith may be related to that satisfy any potential other conditions.\n- Joins should always be made on entities, not source documents!\n- Try to be as efficient as possible.\n\nAPPROACH:\nPlease think through this step by step. Make sure that you include all columns in the ORDER BY clause \\\nalso in the SELECT DISTINCT clause, \\\nif applicable! And again, joins should generally be LEFT JOINS!\n\nAlso, in case it is important, today is ---today_date--- and the user/employee asking is ---user_name---.\n\nPlease structure your answer using , , , start and end tags as in:\n\n[think through the logic but do so extremely briefly! Not more than 3-4 sentences.]\n[the SQL statement that you generate to satisfy the task]\n\"\"\".strip()\n\n# TODO: remove following before merging after enough testing\nSIMPLE_SQL_CORRECTION_PROMPT = f\"\"\"\nYou are an expert in reviewing and fixing SQL statements.\n\nHere is a draft SQL statement that you should consider as generally capturing the information intended. \\\nHowever, it may or may not be syntactically 100% for our postgresql database.\n\nGuidance:\n - Think about whether attributes should be numbers or strings. You may need to convert them.\n - If we use SELECT DISTINCT we need to have the ORDER BY columns in the \\\nSELECT statement as well! And it needs to be in the EXACT FORM! So if a \\\nconversion took place, make sure to include the conversion in the SELECT and the ORDER BY clause!\n - never should 'source_document' be in the SELECT clause! Remove if present!\n - if there are joins, they must be on entities, never source documents\n - if there are joins, consider the possibility that the second entity does not exist for all examples.\\\n Therefore consider using LEFT joins (or RIGHT joins) as appropriate.\n\nDraft SQL:\n{SEPARATOR_LINE}\n---draft_sql---\n{SEPARATOR_LINE}\n\nPlease structure your answer using , , , start and end tags as in:\n\n[think briefly through the problem step by step]\n[the corrected (or original one, if correct) SQL statement]\n\"\"\".strip()\n\nSIMPLE_ENTITY_SQL_PROMPT = f\"\"\"\nYou are an expert in generating a SQL statement that only uses ONE TABLE that captures ENTITIES \\\nand their attributes and other data. The table has the following structure:\n\n{SEPARATOR_LINE}\n{ENTITY_TABLE_DESCRIPTION}\n\nHere is the question you are supposed to translate into a SQL statement:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nTo help you, we already have identified the entities that the SQL statement likely *should* use (but note the \\\nexception below!). The entities as written below also contain the list of attributes and attribute values \\\nthat should specify the entity. \\\nThe format is ::--[:, \\\n:, ...].\n{SEPARATOR_LINE}\nIdentified entities with attributes in query:\n\n---query_entities_with_attributes---\n\nThese are the entities that should be used in the SQL statement. However, \\\nnote that these are the entities (with potential attributes) that were *matches* of Knowledge Graph identified with the \\\nentities originally identified in the original question. As such, they may have id names that may not mean much by themselves, \\\neg ACCOUNT::a74f332. Here is the mapping of entities originally identified (whose role in the query should be obvious) with \\\nthe entities that were matched to them in the Knowledge Graph:\n\n---entity_explanation_string---\n\n--\n\n\n{SEPARATOR_LINE}\n\nCRITICAL SPECIAL CASE:\n - if an identified entity is of the form ::*, or an identified relationship contains an \\\nentity of this form, this refers to *any* entity of that type. Correspondingly, the SQL query should use the *entity type*, \\\nbut not the entity with the * itself. \\\nExample: if you see 'ACCOUNT::*', that means any account matches. So if you are supposed to count the 'ACCOUNT::*', \\\nyou should count the entities of entity_type 'ACCOUNT'.\n\n\nIMPORTANT NOTES:\n- The entities are unique in the table.\n- If the SQL contains a 'SELECT DISTINCT' clause and an ORDER BY clause, then you MUST include the columns from the ORDER BY \\\nclause ALSO IN THE SELECT DISTINCT CLAUSE! This is very important! (This is a postgres db., so this is a MUST!). \\\nYou MUST NOT have a column in the ORDER BY clause that is not ALSO in the SELECT DISTINCT clause!\n- The table cannot be joined on itself.\n- You CAN ONLY return ENTITIES or COUNTS (or other aggregations) of ENTITIES, or you can return \\\nsource_date (but only if the question asks for event dates or times, and then the \\\ncorresponding entity must also be returned).\n- Generally, the query can only return ENTITIES or aggregations of ENTITIES:\n - if individual entities are returned, then you MUST also return the source_document. \\\nIf the source date was requested, you can return that too.\n - if aggregations of entities are returned, then you can only aggregate the entities.\n- Attributes are stored in the attributes json field. As this is postgres, querying for those must be done as \\\n\"attributes ->> '' = ''\".\n- Again, ALWAYS make sure that EACH COLUMN in an ORDER-BY clause IS ALSO IN THE SELECT CLAUSE! Remind yourself \\\nof that in the reasoning.\n- Be careful with dates! Often a date will refer to the source data, which is the date when \\\nan underlying piece of information was updated. However, if the attributes of an entity may contain \\\ntime information as well (like 'started_at', 'completed_at', etc.), then you should really look at \\\nthe wording to see whether you should use a date in the attributes or the event date.\n- Dates are ALWAYS in string format of the form YYYY-MM-DD, for source date as well as for date-like the attributes! \\\nSo please use that format, particularly if you use data comparisons (>, <, ...)\n- Careful with SORT! Really think in which order you want to sort if you have multiple columns you \\\nwant 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 \\\nvariable as the first column to sort by.\n- When doing a SORT on an attribute value of an entity, you MUST also apply a WHERE clause to filter \\\nfor entities that have the attribute value set. For example, if you want to sort the target entity \\\nby the attribute 'created_date', you must also have a WHERE clause that checks whether the target \\\nentity attribute contains 'created_date'. This is vital for proper ordering with null values.\n- Usually, you will want to retrieve or count entities, maybe with attributes. But you almost always want to \\\nhave entities involved in the SELECT clause.\n- You MUST ONLY rely on the entity attributes provided! This is essential! Do not assume \\\nother attributes exist...they don't! Note that there will often be a search using the results \\\nof this query. So if there is information in the question that does not fit the provided attributes, \\\nyou should not use it here but rely on the later search!\n- Try to be as efficient as possible.\n\nAPPROACH:\nPlease think through this step by step. Make sure that you include all columns in the ORDER BY clause \\\nalso in the SELECT DISTINCT clause, \\\nif applicable!\n\nAlso, in case it is important, today is ---today_date--- and the user/employee asking is ---user_name---.\n\nPlease structure your answer using , , , start and end tags as in:\n\n[think through the logic but do so extremely briefly! Not more than 3-4 sentences.]\n[the SQL statement that you generate to satisfy the task]\n\"\"\".strip()\n\nSIMPLE_SQL_ERROR_FIX_PROMPT = f\"\"\"\nYou are an expert at fixing SQL statements. You will be provided with a SQL statement that aims to address \\\na question, but it contains an error. Your task is to fix the SQL statement, based on the error message.\n\nHere is the description of the table that the SQL statement is supposed to use:\n---table_description---\n\nHere is the question you are supposed to translate into a SQL statement:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nHere is the SQL statement that you should fix:\n{SEPARATOR_LINE}\n---sql_statement---\n{SEPARATOR_LINE}\n\nHere is the error message that was returned:\n{SEPARATOR_LINE}\n---error_message---\n{SEPARATOR_LINE}\n\nNote that in the case the error states the sql statement did not return any results, it is possible that the \\\nsql statement is correct, but the question is not addressable with the information in the knowledge graph. \\\nIf you are absolutely certain that is the case, you may return the original sql statement.\n\nHere are a couple common errors that you may encounter:\n- source_document is in the SELECT clause -> remove it\n- columns used in ORDER BY must also appear in the SELECT DISTINCT clause\n- consider carefully the type of the columns you are using, especially for attributes. You may need to cast them\n- dates are ALWAYS in string format of the form YYYY-MM-DD, for source date as well as for date-like the attributes! \\\nSo please use that format, particularly if you use data comparisons (>, <, ...)\n- attributes are stored in the attributes json field. As this is postgres, querying for those must be done as \\\n\"attributes ->> '' = ''\" (or \"attributes ? ''\" to check for existence).\n- if you are using joins and the sql returned no joins, make sure you are using the appropriate join type (LEFT, RIGHT, etc.) \\\nit is possible that the second entity does not exist for all examples.\n- (ignore if using entity_table) if using the relationship_table and the sql returned no results, make sure you are \\\nselecting the correct column! Use the available relationship types to determine whether to use the source or target entity.\n\nAPPROACH:\nPlease think through this step by step. Please also bear in mind that the sql statement is written in postgres syntax.\n\nAlso, in case it is important, today is ---today_date--- and the user/employee asking is ---user_name---.\n\nPlease structure your answer using , , , start and end tags as in:\n\n[think through the logic but do so extremely briefly! Not more than 3-4 sentences.]\n[the SQL statement that you generate to satisfy the task]\n\"\"\"\n\n\nSEARCH_FILTER_CONSTRUCTION_PROMPT = f\"\"\"\nYou need to prepare a search across text segments that contain the information necessary to \\\nanswer a question. The text segments have tags that can be used to filter for the relevant segments. \\\nKey are suitable entities and relationships of a knowledge graph, as well as underlying source documents.\n\nYour overall task is to find the filters and structures that are needed to filtering a database to \\\nproperly address a user question.\n\nYou will be given:\n - the user question\n - a description of all of the potential entity types involved\n - a list of 'global' entities and relationships that should be filtered by, given the question\n - the structure of a schema that was used to derive additional entity filters\n - a SQL statement that was generated to derive those filters\n - the results that were generated using the SQL statement. This can have multiple rows, \\\nand those will be the 'local' filters (which will later mean that each retrieved result will \\\nneed to match at least one of the conditions that you will generate).\n - the results of another query that asked for the underlying source documents that resulted \\\nin the answers of the SQL statement\n\n\nHere is the information:\n\n1) The overall user question\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\n2) Here is a description of all of the entity types:\n{SEPARATOR_LINE}\n---entity_type_descriptions---\n{SEPARATOR_LINE}\n\n3) Here are the lists of entity and relationship filters that were derived from the question:\n{SEPARATOR_LINE}\nEntity filters:\n\n---entity_filters---\n\n--\n\nRelationship filters:\n\n---relationship_filters---\n\n{SEPARATOR_LINE}\n\n4) Here are the columns of a table in a database that has a lot of knowledge about the \\\ndata:\n{SEPARATOR_LINE}\n - relationship (str): The name of the RELATIONSHIP, combining the nature of the relationship and the names of the entities. \\\nIt is of the form \\\n::____:: \\\n[example: ACCOUNT::Nike__has__CONCERN::performance]. Note that this is NOT UNIQUE!\n - source_entity (str): the id of the source ENTITY/NODE in the relationship [example: ACCOUNT::Nike]\n - source_entity_attributes (json): the attributes of the source entity/node [example: {{\"account_type\": \"customer\"}}]\n - target_entity (str): the id of the target ENTITY/NODE in the relationship [example: CONCERN::performance]\n - target_entity_attributes (json): the attributes of the target entity/node [example: {{\"degree\": \"severe\"}}]\n - source_entity_type (str): the type of the source entity/node [example: ACCOUNT]. Only the entity types provided \\\n below are valid.\n - target_entity_type (str): the type of the target entity/node [example: CONCERN]. Only the entity types provided \\\n below are valid.\n - relationship_type (str): the type of the relationship, formatted as \\\n____. So the explicit entity_names have \\\nbeen removed. [example: ACCOUNT__has__CONCERN]\n - source_document (str): the id of the document that contains the relationship. Note that the combination of \\\nid_name and source_document IS UNIQUE!\n - source_date (str): the 'event' date of the source document [example: 2021-01-01]\n\n{SEPARATOR_LINE}\n\n5) Here is a query that was generated for that table to provide additional filters:\n{SEPARATOR_LINE}\n---sql_query---\n{SEPARATOR_LINE}\n\n6) Here are the results of that SQL query. (Consider the schema description and the \\\nstructure of the entities to interpret the results)\n{SEPARATOR_LINE}\n---sql_results---\n{SEPARATOR_LINE}\n\n7) Here are the results of the other query that provided the underlying source documents \\\nusing the schema:\n{SEPARATOR_LINE}\n---source_document_results---\n{SEPARATOR_LINE}\n\nHere is the detailed set of tasks that you should perform, including the proper output format for you:\n\nPlease reply as a json dictionary in this form:\n\n{{\n \"global_entity_filters\": ,\n \"global_relationship_filters\": ,\n \"local_entity_filters\": ,\n \"source_document_filters\": ,\n \"structure\": \n}}\n\nAgain - DO NOT FORGET - here is the user question that motivates this whole task:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nYour json dictionary answer:\n\"\"\".strip()\n\nOUTPUT_FORMAT_NO_EXAMPLES_PROMPT = f\"\"\"\nYou need to format an answer to a research question. \\\nYou will see what the desired output is, the original question, and the unformatted answer to the research question. \\\nYour purpose is to generate the answer respecting the desired format.\n\nNotes:\n - Note that you are a language model and that answers may or may not be perfect. To communicate \\\nthis to the user, consider phrases like 'I found [10 accounts]...', or 'Here are a number of [goals] that \\\nI found...]\n- Please DO NOT mention the explicit output format in your answer. Just use it to inform the formatting.\n\nHere is the unformatted answer to the research question:\n{SEPARATOR_LINE}\n---introductory_answer---\n{SEPARATOR_LINE}\n\nHere is the original question:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nAnd finally, here is the desired output format:\n{SEPARATOR_LINE}\n---output_format---\n{SEPARATOR_LINE}\n\nPlease start generating the answer, without any explanation. There should be no real modifications to \\\nthe text, after all, all you need to do here is formatting. \\\n\nYour Answer:\n\"\"\".strip()\n\n\nOUTPUT_FORMAT_PROMPT = f\"\"\"\nYou need to format the answers to a research question that was generated using one or more objects. \\\nAn overall introductory answer may be provided to you, as well as the research results for each individual object. \\\nYou will also be provided with the original question as background, and the desired format. \\\n\nYour purpose is to generate a consolidated and FORMATTED answer that starts of with the introductory \\\nanswer, and then formats the research results for each individual object in the desired format. \\\nDo not add any other text please!\n\nNotes:\n - Note that you are a language model and that answers may or may not be perfect. To communicate \\\nthis to the user, consider phrases like 'I found [10 accounts]...', or 'Here are a number of [goals] that \\\nI found...]\n- Please DO NOT mention the explicit output format in your answer. Just use it to inform the formatting.\n- DO NOT add any content to the introductory answer!\n\n\nHere is the original question for your background:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nHere is the desired output format:\n{SEPARATOR_LINE}\n---output_format---\n{SEPARATOR_LINE}\n\nHere is the introductory answer:\n{SEPARATOR_LINE}\n---introductory_answer---\n{SEPARATOR_LINE}\n\nHere are the research results that you should - respecting the target format- return in a formatted way:\n{SEPARATOR_LINE}\n---research_results---\n{SEPARATOR_LINE}\n\nPlease start generating the answer, without any explanation. After all, all you need to do here is formatting. \\\n\n\nYour Answer:\n\"\"\".strip()\n\nOUTPUT_FORMAT_NO_OVERALL_ANSWER_PROMPT = f\"\"\"\nYou need to format the return of research on multiple objects. The research results will be given \\\nto you as a string. You will also see what the desired output is, as well as the original question. \\\nYour purpose is to generate the answer respecting the desired format.\n\nNotes:\n - Note that you are a language model and that answers may or may not be perfect. To communicate \\\nthis to the user, consider phrases like 'I found [10 accounts]...', or 'Here are a number of [goals] that \\\nI found...]\n- Please DO NOT mention the explicit output format in your answer. Just use it to inform the formatting.\n - Often, you are also provided with a list of explicit examples. If - AND ONLY IF - the list is not \\\nempty, then these should be listed at the end with the text:\n'...\nHere are some examples of what I found:\n\n...'\n - Again if the list of examples is an empty string then skip this section! Do not use the \\\nresults data for this purpose instead! (They will already be handled in the answer.)\n- Even if the desired output format is 'text', make sure that you keep the individual research results \\\nseparated by bullet points, and mention the object name first, followed by a new line. The object name \\\nis at the beginning of the research result, and should be in the format ::.\n\n\nHere is the original question:\n{SEPARATOR_LINE}\n---question---\n{SEPARATOR_LINE}\n\nAnd finally, here is the desired output format:\n{SEPARATOR_LINE}\n---output_format---\n{SEPARATOR_LINE}\n\nHere are the research results that you should properly format:\n{SEPARATOR_LINE}\n---research_results---\n{SEPARATOR_LINE}\n\nPlease start generating the answer, without any explanation. After all, all you need to do here is formatting. \\\n\n\nYour Answer:\n\"\"\".strip()\n\nKG_OBJECT_SOURCE_RESEARCH_PROMPT = f\"\"\"\nYou are an expert in extracting relevant structured information from a list of documents that \\\nshould relate to one object. You are presented with a list of documents that have been determined to be \\\nrelevant to the task of interest. Your goal is to extract the information asked around these topics:\nYou should look at the documents - in no particular order! - and extract the information that relates \\\nto a question:\n{SEPARATOR_LINE}\n{{question}}\n{SEPARATOR_LINE}\n\nHere are the documents you are supposed to search through:\n--\n{{document_text}}\n{SEPARATOR_LINE}\nNote: in this case, please do NOT cite your sources. This is very important!\n\nPlease now generate the answer to the question given the documents:\n\"\"\".strip()\n\nKG_SEARCH_PROMPT = f\"\"\"\nYou are an expert in extracting relevant structured information from a list of documents that \\\nshould relate to one object. You are presented with a list of documents that have been determined to be \\\nrelevant to the task of interest. Your goal is to extract the information asked around these topics:\nYou should look at the documents and extract the information that relates \\\nto a question:\n{SEPARATOR_LINE}\n{{question}}\n{SEPARATOR_LINE}\n\nHere are the documents you are supposed to search through:\n--\n{{document_text}}\n{SEPARATOR_LINE}\nNote: in this case, please DO cite your sources. This is very important! \\\nUse the format []. Ie, use [1], [2], and NOT [1,2] if \\\nthere are two documents to cite, etc. \\\n\n\nPlease now generate the answer to the question given the documents:\n\"\"\".strip()\n\n# KG Beta Assistant System Prompt\nKG_BETA_ASSISTANT_SYSTEM_PROMPT = \"\"\"\"You are a knowledge graph assistant that helps users explore and \\\nunderstand relationships between entities.\"\"\"\n\nKG_BETA_ASSISTANT_TASK_PROMPT = \"\"\"\"Help users explore and understand the knowledge graph by answering \\\nquestions about entities and their relationships.\"\"\"\n\n\n# Just in case, for best practice, send a system message with key rules.\n# (The db user permissions executing the SQL will avoid issues anyway,\n# but it does not hurt to to put multiple checks in place.)\nSQL_INSTRUCTIONS_RELATIONSHIP_PROMPT = \"\"\"\nYou are an expert at generating SQL queries to answer questions about a knowledge graph.\n\nYou will be given a lot of instructions later, but here rules that MUST BE FOLLOWED:\n - the SQL generated MUST only use the table one table named 'relationship_table'. \\\nThis table is not a table that can be defined or overwritten by the user and the resulting SQL \\\nstatement, it MUST be seen as an existing table in the database.\n - self-joins of the 'relationship_table' are allowed, as well as common table expressions \\\n that reference only the 'relationship_table'.\n - no other table or view can in any way or shape be \\\ninvolved in the generated SQL.\n - no other database operations can be generated except for those that query the 'relationship_table'. \\\n(WHERE, GROUP BY, etc. are certainly allowed, but no other database table can be used in the generated SQL.)\n\"\"\"\n\nSQL_INSTRUCTIONS_ENTITY_PROMPT = \"\"\"\nYou are an expert at generating SQL queries to answer questions about a knowledge graph.\n\nYou will be given a lot of instructions later, but here rules that MUST BE FOLLOWED:\n - the SQL generated MUST only use the table one table named 'entity_table'. \\\nThis table is not a table that can be defined or overwritten by the user and the resulting SQL \\\nstatement, it MUST be seen as an existing table in the database.\n - common table expressions that reference only the 'entity_table' are allowed.\n - no other table or view of a potential underlying schema can in any way or shape be \\\ninvolved in the generated SQL.\n - no other database operations can be generated except for those that query the 'entity_table'. \\\n(WHERE, GROUP BY, etc. are certainly allowed, but no other database table can be used in the generated SQL.)\n\"\"\"\n" }, { "path": "backend/onyx/prompts/prompt_template.py", "content": "import re\n\nfrom onyx.prompts.prompt_utils import replace_current_datetime_tag\n\n\nclass PromptTemplate:\n \"\"\"\n A class for building prompt templates with placeholders.\n Useful when building templates with json schemas, as {} will not work with f-strings.\n Unlike string.replace, this class will raise an error if the fields are missing.\n \"\"\"\n\n DEFAULT_PATTERN = r\"---([a-zA-Z0-9_]+)---\"\n\n def __init__(self, template: str, pattern: str = DEFAULT_PATTERN):\n self._pattern_str = pattern\n self._pattern = re.compile(pattern)\n self._template = template\n self._fields: set[str] = set(self._pattern.findall(template))\n\n def build(self, **kwargs: str) -> str:\n \"\"\"\n Build the prompt template with the given fields.\n Will raise an error if the fields are missing.\n Will ignore fields that are not in the template.\n \"\"\"\n missing = self._fields - set(kwargs.keys())\n if missing:\n raise ValueError(f\"Missing required fields: {missing}.\")\n built = self._replace_fields(kwargs)\n return self._postprocess(built)\n\n def partial_build(self, **kwargs: str) -> \"PromptTemplate\":\n \"\"\"\n Returns another PromptTemplate with the given fields replaced.\n Will ignore fields that are not in the template.\n \"\"\"\n new_template = self._replace_fields(kwargs)\n return PromptTemplate(new_template, self._pattern_str)\n\n def _replace_fields(self, field_vals: dict[str, str]) -> str:\n def repl(match: re.Match) -> str:\n key = match.group(1)\n return field_vals.get(key, match.group(0))\n\n return self._pattern.sub(repl, self._template)\n\n def _postprocess(self, text: str) -> str:\n \"\"\"Apply global replacements such as [[CURRENT_DATETIME]].\"\"\"\n if not text:\n return text\n # Ensure [[CURRENT_DATETIME]] matches shared prompt formatting\n return replace_current_datetime_tag(\n text,\n full_sentence=True,\n include_day_of_week=True,\n )\n" }, { "path": "backend/onyx/prompts/prompt_utils.py", "content": "from datetime import datetime\nfrom typing import cast\n\nfrom langchain_core.messages import BaseMessage\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.prompts.chat_prompts import ADDITIONAL_INFO\nfrom onyx.prompts.chat_prompts import CITATION_GUIDANCE_REPLACEMENT_PAT\nfrom onyx.prompts.chat_prompts import COMPANY_DESCRIPTION_BLOCK\nfrom onyx.prompts.chat_prompts import COMPANY_NAME_BLOCK\nfrom onyx.prompts.chat_prompts import DATETIME_REPLACEMENT_PAT\nfrom onyx.prompts.chat_prompts import REMINDER_TAG_REPLACEMENT_PAT\nfrom onyx.prompts.chat_prompts import REQUIRE_CITATION_GUIDANCE\nfrom onyx.prompts.constants import CODE_BLOCK_PAT\nfrom onyx.prompts.constants import REMINDER_TAG_DESCRIPTION\nfrom onyx.server.settings.store import load_settings\nfrom onyx.utils.logger import setup_logger\n\n\nlogger = setup_logger()\n\n\n_BASIC_TIME_STR = \"The current date is {datetime_info}.\"\n\n\ndef get_current_llm_day_time(\n include_day_of_week: bool = True,\n full_sentence: bool = True,\n include_hour_min: bool = False,\n) -> str:\n current_datetime = datetime.now()\n # Format looks like: \"October 16, 2023 14:30\" if include_hour_min, otherwise \"October 16, 2023\"\n formatted_datetime = (\n current_datetime.strftime(\"%B %d, %Y %H:%M\")\n if include_hour_min\n else current_datetime.strftime(\"%B %d, %Y\")\n )\n day_of_week = current_datetime.strftime(\"%A\")\n if full_sentence:\n return f\"The current day and time is {day_of_week} {formatted_datetime}\"\n if include_day_of_week:\n return f\"{day_of_week} {formatted_datetime}\"\n return f\"{formatted_datetime}\"\n\n\ndef replace_current_datetime_tag(\n prompt_str: str,\n *,\n full_sentence: bool = False,\n include_day_of_week: bool = True,\n) -> str:\n datetime_str = get_current_llm_day_time(\n full_sentence=full_sentence,\n include_day_of_week=include_day_of_week,\n )\n\n if DATETIME_REPLACEMENT_PAT in prompt_str:\n prompt_str = prompt_str.replace(DATETIME_REPLACEMENT_PAT, datetime_str)\n\n return prompt_str\n\n\ndef replace_citation_guidance_tag(\n prompt_str: str,\n *,\n should_cite_documents: bool = False,\n include_all_guidance: bool = False,\n) -> tuple[str, bool]:\n \"\"\"\n Replace {{CITATION_GUIDANCE}} placeholder with citation guidance if needed.\n\n Returns:\n tuple[str, bool]: (prompt_with_replacement, should_append_fallback)\n - prompt_with_replacement: The prompt with placeholder replaced (or unchanged if not present)\n - should_append_fallback: True if citation guidance should be appended\n (placeholder is not present and citations are needed)\n \"\"\"\n placeholder_was_present = CITATION_GUIDANCE_REPLACEMENT_PAT in prompt_str\n\n if not placeholder_was_present:\n # Placeholder not present - caller should append if citations are needed\n should_append = (\n should_cite_documents or include_all_guidance\n ) and REQUIRE_CITATION_GUIDANCE not in prompt_str\n return prompt_str, should_append\n\n citation_guidance = (\n REQUIRE_CITATION_GUIDANCE\n if should_cite_documents or include_all_guidance\n else \"\"\n )\n\n prompt_str = prompt_str.replace(\n CITATION_GUIDANCE_REPLACEMENT_PAT,\n citation_guidance,\n )\n\n return prompt_str, False\n\n\ndef replace_reminder_tag(prompt_str: str) -> str:\n \"\"\"Replace {{REMINDER_TAG_DESCRIPTION}} with the reminder tag content.\"\"\"\n if REMINDER_TAG_REPLACEMENT_PAT in prompt_str:\n prompt_str = prompt_str.replace(\n REMINDER_TAG_REPLACEMENT_PAT, REMINDER_TAG_DESCRIPTION\n )\n\n return prompt_str\n\n\ndef handle_onyx_date_awareness(\n prompt_str: str,\n # We always replace the pattern {{CURRENT_DATETIME}} if it shows up\n # but if it doesn't show up and the prompt is datetime aware, add it to the prompt at the end.\n datetime_aware: bool = False,\n) -> str:\n \"\"\"\n If there is a {{CURRENT_DATETIME}} tag, replace it with the current date and time no matter what.\n If the prompt is datetime aware, and there are no datetime tags, add it to the prompt.\n Do nothing otherwise.\n This can later be expanded to support other tags.\n \"\"\"\n\n prompt_with_datetime = replace_current_datetime_tag(\n prompt_str,\n full_sentence=False,\n include_day_of_week=True,\n )\n if prompt_with_datetime != prompt_str:\n return prompt_with_datetime\n\n if datetime_aware:\n return prompt_str + ADDITIONAL_INFO.format(\n datetime_info=_BASIC_TIME_STR.format(\n datetime_info=get_current_llm_day_time()\n )\n )\n\n return prompt_str\n\n\ndef get_company_context() -> str | None:\n prompt_str = None\n try:\n workspace_settings = load_settings()\n company_name = workspace_settings.company_name\n company_description = workspace_settings.company_description\n\n if not company_name and not company_description:\n return None\n\n prompt_str = \"\"\n if company_name:\n prompt_str += COMPANY_NAME_BLOCK.format(company_name=company_name)\n if company_description:\n prompt_str += COMPANY_DESCRIPTION_BLOCK.format(\n company_description=company_description\n )\n return prompt_str\n except Exception as e:\n logger.error(f\"Error handling company awareness: {e}\")\n return None\n\n\n# Maps connector enum string to a more natural language representation for the LLM\n# If not on the list, uses the original but slightly cleaned up, see below\nCONNECTOR_NAME_MAP = {\n \"web\": \"Website\",\n \"requesttracker\": \"Request Tracker\",\n \"github\": \"GitHub\",\n \"file\": \"File Upload\",\n}\n\n\ndef clean_up_source(source_str: str) -> str:\n if source_str in CONNECTOR_NAME_MAP:\n return CONNECTOR_NAME_MAP[source_str]\n return source_str.replace(\"_\", \" \").title()\n\n\ndef build_doc_context_str(\n semantic_identifier: str,\n source_type: DocumentSource,\n content: str,\n metadata_dict: dict[str, str | list[str]],\n updated_at: datetime | None,\n ind: int,\n include_metadata: bool = True,\n) -> str:\n context_str = \"\"\n if include_metadata:\n context_str += f\"DOCUMENT {ind}: {semantic_identifier}\\n\"\n context_str += f\"Source: {clean_up_source(source_type)}\\n\"\n\n for k, v in metadata_dict.items():\n if isinstance(v, list):\n v_str = \", \".join(v)\n context_str += f\"{k.capitalize()}: {v_str}\\n\"\n else:\n context_str += f\"{k.capitalize()}: {v}\\n\"\n\n if updated_at:\n update_str = updated_at.strftime(\"%B %d, %Y %H:%M\")\n context_str += f\"Updated: {update_str}\\n\"\n context_str += f\"{CODE_BLOCK_PAT.format(content.strip())}\\n\\n\\n\"\n return context_str\n\n\n_PER_MESSAGE_TOKEN_BUFFER = 7\n\n\ndef find_last_index(lst: list[int], max_prompt_tokens: int) -> int:\n \"\"\"From the back, find the index of the last element to include\n before the list exceeds the maximum\"\"\"\n running_sum = 0\n\n if not lst:\n logger.warning(\"Empty message history passed to find_last_index\")\n return 0\n\n last_ind = 0\n for i in range(len(lst) - 1, -1, -1):\n running_sum += lst[i] + _PER_MESSAGE_TOKEN_BUFFER\n if running_sum > max_prompt_tokens:\n last_ind = i + 1\n break\n\n if last_ind >= len(lst):\n logger.error(\n f\"Last message alone is too large! max_prompt_tokens: {max_prompt_tokens}, message_token_counts: {lst}\"\n )\n raise ValueError(\"Last message alone is too large!\")\n\n return last_ind\n\n\ndef drop_messages_history_overflow(\n messages_with_token_cnts: list[tuple[BaseMessage, int]],\n max_allowed_tokens: int,\n) -> list[BaseMessage]:\n \"\"\"As message history grows, messages need to be dropped starting from the furthest in the past.\n The System message should be kept if at all possible and the latest user input which is inserted in the\n prompt template must be included\"\"\"\n\n final_messages: list[BaseMessage] = []\n messages, token_counts = cast(\n tuple[list[BaseMessage], list[int]], zip(*messages_with_token_cnts)\n )\n system_msg = (\n final_messages[0]\n if final_messages and final_messages[0].type == \"system\"\n else None\n )\n\n history_msgs = messages[:-1]\n final_msg = messages[-1]\n if final_msg.type != \"human\":\n if final_msg.type != \"tool\":\n raise ValueError(\"Last message must be user input OR a tool result\")\n else:\n final_msgs = messages[-3:]\n history_msgs = messages[:-3]\n else:\n final_msgs = [final_msg]\n\n # Start dropping from the history if necessary\n ind_prev_msg_start = find_last_index(\n token_counts, max_prompt_tokens=max_allowed_tokens\n )\n\n if system_msg and ind_prev_msg_start <= len(history_msgs):\n final_messages.append(system_msg)\n\n final_messages.extend(history_msgs[ind_prev_msg_start:])\n final_messages.extend(final_msgs)\n\n return final_messages\n" }, { "path": "backend/onyx/prompts/search_prompts.py", "content": "# How it works and rationale:\n# First - this works best empirically across multiple LLMs, some of this is back-explaining reasons based on results.\n#\n# The system prompt is kept simple and as similar to typical system prompts as possible to stay within training distribution.\n# The history is passed through as a list of messages, this should allow the LLM to more easily understand what is going on.\n# The special tokens and separators let the LLM more easily disregard no longer relevant past messages.\n# The last message is dynamically created and has a detailed description of the actual task.\n# This is based on the assumption that users give much more varied requests in their prompts and LLMs are well adjusted to this.\n# The proximity of the instructions and the lack of any breaks should also let the LLM follow the task more clearly.\n#\n# For document verification, the history is not included as the queries should ideally be standalone enough.\n# To keep it simple, it is just a single simple prompt.\n\n\nSEMANTIC_QUERY_REPHRASE_SYSTEM_PROMPT = \"\"\"\nYou are an assistant that reformulates the last user message into a standalone, self-contained query suitable for \\\nsemantic search. Your goal is to output a single natural language query that captures the full meaning of the user's \\\nmost recent message. It should be fully semantic and natural language unless the user query is already a keyword query. \\\nWhen relevant, you bring in context from the history or knowledge about the user.\n\nThe current date is {current_date}.\n\"\"\"\n\nSEMANTIC_QUERY_REPHRASE_USER_PROMPT = \"\"\"\nGiven the chat history above (if any) and the final user query (provided below), provide a standalone query that is as\nrepresentative of the user query as possible. In most cases, it should be exactly the same as the last user query. \\\nIt should be fully semantic and natural language unless the user query is already a keyword query. \\\nFocus on the last user message, in most cases the history and extra context should be ignored.\n\nFor a query like \"What are the use cases for product X\", your output should remain \"What are the use cases for product X\". \\\nIt should remain semantic, and as close to the original query as possible. There is nothing additional needed \\\nfrom the history or that should be removed / replaced from the query.\n\nFor modifications, you can:\n1. Insert relevant context from the chat history. For example:\n\"How do I set it up?\" -> \"How do I set up software Y?\" (assuming the conversation was about software Y)\n\n2. Remove asks or requests not related to the searching. For example:\n\"Can you summarize the calls with example company\" -> \"calls with example company\"\n\"Can you find me the document that goes over all of the software to set up on an engineer's first day?\" -> \\\n\"all of the software to set up on an engineer's first day\"\n\n3. Fill in relevant information about the user. For example:\n\"What document did I write last week?\" -> \"What document did John Doe write last week?\" (assuming the user is John Doe)\n{additional_context}\n=========================\nCRITICAL: ONLY provide the standalone query and nothing else.\n\nFinal user query:\n{user_query}\n\"\"\".strip()\n\n\nKEYWORD_REPHRASE_SYSTEM_PROMPT = \"\"\"\nYou are an assistant that reformulates the last user message into a set of standalone keyword queries suitable for a keyword \\\nsearch engine. Your goal is to output keyword queries that optimize finding relevant documents to answer the user query. \\\nWhen relevant, you bring in context from the history or knowledge about the user.\n\nThe current date is {current_date}.\n\"\"\"\n\n\nKEYWORD_REPHRASE_USER_PROMPT = \"\"\"\nGiven the chat history above (if any) and the final user query (provided below), provide a set of keyword only queries that can\nhelp find relevant documents. Provide a single query per line (where each query consists of one or more keywords). \\\nThe queries must be purely keywords and not contain any natural language. \\\nEach query should have as few keywords as necessary to represent the user's search intent.\n\nGuidelines:\n- Do not provide more than 3 queries.\n- Do not replace or expand niche, proprietary, or obscure terms\n- Focus on the last user message, in most cases the history and any extra context should be ignored.\n{additional_context}\n=========================\nCRITICAL: ONLY provide the keyword queries, one set of keywords per line and nothing else.\n\nFinal user query:\n{user_query}\n\"\"\".strip()\n\n\nREPHRASE_CONTEXT_PROMPT = \"\"\"\nIn most cases the following additional context is not needed. If relevant, here is some information about the user:\n{user_info}\n\nHere are some memories about the user:\n{memories}\n\"\"\"\n\n\n# This prompt is intended to be fairly lenient since there are additional filters downstream.\n# There are now multiple places for misleading docs to get dropped so each one can be a bit more lax.\n# As models get better, it's likely better to include more context than not, some questionably\n# useful stuff may be helpful downstream.\n# Adding the ! option to allow better models to handle questions where all of the documents are\n# necessary to make a good determination.\n# If a document is by far the best and is a very obvious inclusion, add a ! after the section_id to indicate that it should \\\n# be included in full. Example output: [8, 2!, 5].\nDOCUMENT_SELECTION_PROMPT = \"\"\"\nSelect the most relevant document sections for the user's query (maximum {max_sections}).{extra_instructions}\n\n# Document Sections\n```\n{formatted_doc_sections}\n```\n\n# User Query\n```\n{user_query}\n```\n\n# Selection Criteria\n- Choose sections most relevant to answering the query, if at all in doubt, include the section.\n- Even if only a tiny part of the section is relevant, include it.\n- It is ok to select multiple sections from the same document.\n- Consider indirect connections and supporting context to be valuable.\n- If the section is not directly helpful but the document seems relevant, there is an opportunity \\\nlater to expand the section and read more from the document so include the section.\n\n# Output Format\nReturn ONLY section_ids as a comma-separated list, ordered by relevance:\n[most_relevant_section_id, second_most_relevant_section_id, ...]\n\nSection IDs:\n\"\"\".strip()\n\nTRY_TO_FILL_TO_MAX_INSTRUCTIONS = \"\"\"\nTry to fill the list to the maximum number of sections if possible without including non-relevant or misleading sections.\n\"\"\"\n\n\n# Some models are trained heavily to reason in the actual output so we allow some flexibility in the prompt.\n# Downstream of the model, we will attempt to parse the output to extract the number.\n# This inference will not have a system prompt as it's a single message task more like the traditional ones.\n# LLMs should do better with just this type of next word prediction.\n# Opted to not include metadata here as the doc was already selected by the previous step that has it.\n# Also hopefully it leans not throwing out documents as there are not many bad ones that make it to this stage.\n# If anything, it's mostly because of something misleading, otherwise this step should be treated as 95% expansion/filtering.\nDOCUMENT_CONTEXT_SELECTION_PROMPT = \"\"\"\nAnalyze the relevance of document sections to a search query and classify according to the categories \\\ndescribed at the end of the prompt.\n\n# Document Title / Metadata\n```\n{document_title}\n```\n\n# Section Above:\n```\n{section_above}\n```\n\n# Main Section:\n```\n{main_section}\n```\n\n# Section Below:\n```\n{section_below}\n```\n\n# User Query:\n```\n{user_query}\n```\n\n# Classification Categories:\n**0 - NOT_RELEVANT**\n- Main section and surrounding sections do not help answer the query or provide meaningful, relevant information.\n- Appears on topic but refers to a different context or subject (could lead to potential confusion or misdirection). \\\nIt is important to avoid conflating different contexts and subjects - if the document is related to the query but not about \\\nthe correct subject. Example: \"How much did we quote ACME for project X\", \"ACME paid us $100,000 for project Y\".\n\n**1 - MAIN_SECTION_ONLY**\n- Main section contains useful information relevant to the query.\n- Adjacent sections do not provide additional directly relevant information.\n\n**2 - INCLUDE_ADJACENT_SECTIONS**\n- The main section AND adjacent sections are all useful for answering the user query.\n- The surrounding sections provide relevant information that does not exist in the main section.\n- Even if only 1 of the adjacent sections is useful or there is a small piece in either that is useful.\n- Additional unseen sections are unlikely to contain valuable related information.\n\n**3 - INCLUDE_FULL_DOCUMENT**\n- Additional unseen sections are likely to contain valuable related information to the query.\n\n## Additional Decision Notes\n- If only a small piece of the document is useful - use classification 1 or 2, do not use 0.\n- If the document is on topic and provides additional context that might be useful in \\\ncombination with other documents - use classification 1, 2 or 3, do not use 0.\n\nCRITICAL: ONLY output the NUMBER of the situation most applicable to the query and sections provided (0, 1, 2, or 3).\n\nSituation Number:\n\"\"\".strip()\n" }, { "path": "backend/onyx/prompts/tool_prompts.py", "content": "# ruff: noqa: E501, W605 start\n# If there are any tools, this section is included, the sections below are for the available tools\nTOOL_SECTION_HEADER = \"\\n# Tools\\n\\n\"\n\n\n# This section is included if there are search type tools, currently internal_search and web_search\nTOOL_DESCRIPTION_SEARCH_GUIDANCE = \"\"\"\nFor questions that can be answered from existing knowledge, answer the user directly without using any tools. \\\nIf you suspect your knowledge is outdated or for topics where things are rapidly changing, use search tools to get more context. \\\nFor statements that may be describing or referring to a document, run a search for the document. \\\nIn ambiguous cases, favor searching to get more context.\n\nWhen using any search type tool, do not make any assumptions and stay as faithful to the user's query as possible. \\\nBetween internal and web search (if both are available), think about if the user's query is likely better answered by team internal sources or online web pages. \\\nWhen searching for information, if the initial results cannot fully answer the user's query, try again with different tools or arguments. \\\nDo not repeat the same or very similar queries if it already has been run in the chat history.\n\nIf it is unclear which tool to use, consider using multiple in parallel to be efficient with time.\n\"\"\".lstrip()\n\n\nINTERNAL_SEARCH_GUIDANCE = \"\"\"\n## internal_search\nUse the `internal_search` tool to search connected applications for information. Some examples of when to use `internal_search` include:\n- Internal information: any time where there may be some information stored in internal applications that could help better answer the query.\n- Niche/Specific information: information that is likely not found in public sources, things specific to a project or product, team, process, etc.\n- Keyword Queries: queries that are heavily keyword based are often internal document search queries.\n- Ambiguity: questions about something that is not widely known or understood.\nNever provide more than 3 queries at once to `internal_search`.\n\"\"\".lstrip()\n\n\nWEB_SEARCH_GUIDANCE = \"\"\"\n## web_search\nUse the `web_search` tool to access up-to-date information from the web. Some examples of when to use `web_search` include:\n- Freshness: when the answer might be enhanced by up-to-date information on a topic. Very important for topics that are changing or evolving.\n- Accuracy: if the cost of outdated/inaccurate information is high.\n- Niche Information: when detailed info is not widely known or understood (but is likely found on the internet).{site_colon_disabled}\n\"\"\".lstrip()\n\nWEB_SEARCH_SITE_DISABLED_GUIDANCE = \"\"\"\nDo not use the \"site:\" operator in your web search queries.\n\"\"\".lstrip()\n\n\nOPEN_URLS_GUIDANCE = \"\"\"\n## open_url\nUse the `open_url` tool to read the content of one or more URLs. Use this tool to access the contents of the most promising web pages from your web searches or user specified URLs. \\\nYou can open many URLs at once by passing multiple URLs in the array if multiple pages seem promising. Prioritize the most promising pages and reputable sources. \\\nDo not open URLs that are image files like .png, .jpg, etc.\nYou should almost always use open_url after a web_search call. Use this tool when a user asks about a specific provided URL.\n\"\"\".lstrip()\n\nPYTHON_TOOL_GUIDANCE = \"\"\"\n## python\nUse the `python` tool to execute Python code in an isolated sandbox. The tool will respond with the output of the execution or time out after 60.0 seconds.\nAny files uploaded to the chat will be automatically be available in the execution environment's current directory. \\\nThe current directory in the file system can be used to save and persist user files. Files written to the current directory will be returned with a `file_link`. \\\nUse this to give the user a way to download the file OR to display generated images.\nInternet access for this session is disabled. Do not make external web requests or API calls as they will fail.\nUse `openpyxl` to read and write Excel files. You have access to libraries like numpy, pandas, scipy, matplotlib, and PIL.\nIMPORTANT: each call to this tool is independent. Variables from previous calls will NOT be available in the current call.\n\"\"\".lstrip()\n\nGENERATE_IMAGE_GUIDANCE = \"\"\"\n## generate_image\nNEVER use generate_image unless the user specifically requests an image.\nFor edits/variations of a previously generated image, pass `reference_image_file_ids` with\nthe `file_id` values returned by earlier `generate_image` tool results.\n\"\"\".lstrip()\n\nMEMORY_GUIDANCE = \"\"\"\n## add_memory\nUse the `add_memory` tool for facts shared by the user that should be remembered for future conversations. \\\nOnly add memories that are specific, likely to remain true, and likely to be useful later. \\\nFocus on enduring preferences, long-term goals, stable constraints, and explicit \"remember this\" type requests.\n\"\"\".lstrip()\n\nTOOL_CALL_FAILURE_PROMPT = \"\"\"\nLLM attempted to call a tool but failed. Most likely the tool name or arguments were misspelled.\n\"\"\".strip()\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/prompts/user_info.py", "content": "# ruff: noqa: E501, W605 start\nUSER_INFORMATION_HEADER = \"\\n# User Information\\n\\n\"\n\nBASIC_INFORMATION_PROMPT = \"\"\"\n## Basic Information\nUser name: {user_name}\nUser email: {user_email}{user_role}\n\"\"\".lstrip()\n\n# This line only shows up if the user has configured their role.\nUSER_ROLE_PROMPT = \"\"\"\nUser role: {user_role}\n\"\"\".lstrip()\n\n# Team information should be a paragraph style description of the user's team.\nTEAM_INFORMATION_PROMPT = \"\"\"\n## Team Information\n{team_information}\n\"\"\".lstrip()\n\n# User preferences should be a paragraph style description of the user's preferences.\nUSER_PREFERENCES_PROMPT = \"\"\"\n## User Preferences\n{user_preferences}\n\"\"\".lstrip()\n\n# User memories should look something like:\n# - Memory 1\n# - Memory 2\n# - Memory 3\nUSER_MEMORIES_PROMPT = \"\"\"\n## User Memories\n{user_memories}\n\"\"\".lstrip()\n\n# ruff: noqa: E501, W605 end\n" }, { "path": "backend/onyx/redis/iam_auth.py", "content": "\"\"\"\nRedis IAM Authentication Module\nThis module provides Redis IAM authentication functionality for AWS ElastiCache.\nUnlike RDS IAM auth, Redis IAM auth relies on IAM roles and policies rather than\ngenerating authentication tokens.\nKey functions:\n- configure_redis_iam_auth: Configure Redis connection parameters for IAM auth\n- create_redis_ssl_context_if_iam: Create SSL context for secure connections\n\"\"\"\n\nimport ssl\nfrom typing import Any\n\n\ndef configure_redis_iam_auth(connection_kwargs: dict[str, Any]) -> None:\n \"\"\"\n Configure Redis connection parameters for IAM authentication.\n Modifies the connection_kwargs dict in-place to:\n 1. Remove password (not needed with IAM)\n 2. Enable SSL with system CA certificates\n 3. Set proper SSL context for secure connections\n \"\"\"\n # Remove password as it's not needed with IAM authentication\n if \"password\" in connection_kwargs:\n del connection_kwargs[\"password\"]\n\n # Ensure SSL is enabled for IAM authentication\n connection_kwargs[\"ssl\"] = True\n connection_kwargs[\"ssl_context\"] = create_redis_ssl_context_if_iam()\n\n\ndef create_redis_ssl_context_if_iam() -> ssl.SSLContext:\n \"\"\"Create an SSL context for Redis IAM authentication using system CA certificates.\"\"\"\n # Use system CA certificates by default - no need for additional CA files\n ssl_context = ssl.create_default_context()\n ssl_context.check_hostname = True\n ssl_context.verify_mode = ssl.CERT_REQUIRED\n return ssl_context\n" }, { "path": "backend/onyx/redis/redis_connector.py", "content": "import redis\n\nfrom onyx.redis.redis_connector_delete import RedisConnectorDelete\nfrom onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync\nfrom onyx.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync\nfrom onyx.redis.redis_connector_prune import RedisConnectorPrune\nfrom onyx.redis.redis_connector_stop import RedisConnectorStop\nfrom onyx.redis.redis_pool import get_redis_client\n\n\n# TODO: reduce dependence on redis\nclass RedisConnector:\n \"\"\"Composes several classes to simplify interacting with a connector and its\n associated background tasks / associated redis interactions.\"\"\"\n\n def __init__(self, tenant_id: str, cc_pair_id: int) -> None:\n \"\"\"id: a connector credential pair id\"\"\"\n\n self.tenant_id: str = tenant_id\n self.cc_pair_id: int = cc_pair_id\n self.redis: redis.Redis = get_redis_client(tenant_id=tenant_id)\n\n self.stop = RedisConnectorStop(tenant_id, cc_pair_id, self.redis)\n self.prune = RedisConnectorPrune(tenant_id, cc_pair_id, self.redis)\n self.delete = RedisConnectorDelete(tenant_id, cc_pair_id, self.redis)\n self.permissions = RedisConnectorPermissionSync(\n tenant_id, cc_pair_id, self.redis\n )\n self.external_group_sync = RedisConnectorExternalGroupSync(\n tenant_id, cc_pair_id, self.redis\n )\n\n @staticmethod\n def get_id_from_fence_key(key: str) -> str | None:\n \"\"\"\n Extracts the object ID from a fence key in the format `PREFIX_fence_X`.\n\n Args:\n key (str): The fence key string.\n\n Returns:\n Optional[int]: The extracted ID if the key is in the correct format, otherwise None.\n \"\"\"\n parts = key.split(\"_\")\n if len(parts) != 3:\n return None\n\n object_id = parts[2]\n return object_id\n\n @staticmethod\n def get_id_from_task_id(task_id: str) -> str | None:\n \"\"\"\n Extracts the object ID from a task ID string.\n\n This method assumes the task ID is formatted as `prefix_objectid_suffix`, where:\n - `prefix` is an arbitrary string (e.g., the name of the task or entity),\n - `objectid` is the ID you want to extract,\n - `suffix` is another arbitrary string (e.g., a UUID).\n\n Example:\n If the input `task_id` is `documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc`,\n this method will return the string `\"1\"`.\n\n Args:\n task_id (str): The task ID string from which to extract the object ID.\n\n Returns:\n str | None: The extracted object ID if the task ID is in the correct format, otherwise None.\n \"\"\"\n # example: task_id=documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc\n parts = task_id.split(\"_\")\n if len(parts) != 3:\n return None\n\n object_id = parts[1]\n return object_id\n\n def db_lock_key(self, search_settings_id: int) -> str:\n \"\"\"\n Key for the db lock for an indexing attempt.\n Prevents multiple modifications to the current indexing attempt row\n from multiple docfetching/docprocessing tasks.\n \"\"\"\n return f\"da_lock:indexing:db_{self.cc_pair_id}/{search_settings_id}\"\n" }, { "path": "backend/onyx/redis/redis_connector_delete.py", "content": "import time\nfrom datetime import datetime\nfrom typing import cast\nfrom uuid import uuid4\n\nimport redis\nfrom celery import Celery\nfrom pydantic import BaseModel\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DB_YIELD_PER_DEFAULT\nfrom onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.document import construct_document_id_select_for_connector_credential_pair\n\n\nclass RedisConnectorDeletePayload(BaseModel):\n num_tasks: int | None\n submitted: datetime\n\n\nclass RedisConnectorDelete:\n \"\"\"Manages interactions with redis for deletion tasks. Should only be accessed\n through RedisConnector.\"\"\"\n\n PREFIX = \"connectordeletion\"\n FENCE_PREFIX = f\"{PREFIX}_fence\" # \"connectordeletion_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n TASKSET_PREFIX = f\"{PREFIX}_taskset\" # \"connectordeletion_taskset\"\n TASKSET_TTL = FENCE_TTL\n\n # used to signal the overall workflow is still active\n # it's impossible to get the exact state of the system at a single point in time\n # so we need a signal with a TTL to bridge gaps in our checks\n ACTIVE_PREFIX = PREFIX + \"_active\"\n ACTIVE_TTL = 3600\n\n def __init__(self, tenant_id: str, id: int, redis: redis.Redis) -> None:\n self.tenant_id: str = tenant_id\n self.id = id\n self.redis = redis\n\n self.fence_key: str = f\"{self.FENCE_PREFIX}_{id}\"\n self.taskset_key = f\"{self.TASKSET_PREFIX}_{id}\"\n\n self.active_key = f\"{self.ACTIVE_PREFIX}_{id}\"\n\n def taskset_clear(self) -> None:\n self.redis.delete(self.taskset_key)\n\n def get_remaining(self) -> int:\n # todo: move into fence\n remaining = cast(int, self.redis.scard(self.taskset_key))\n return remaining\n\n @property\n def fenced(self) -> bool:\n return bool(self.redis.exists(self.fence_key))\n\n @property\n def payload(self) -> RedisConnectorDeletePayload | None:\n # read related data and evaluate/print task progress\n fence_bytes = cast(bytes, self.redis.get(self.fence_key))\n if fence_bytes is None:\n return None\n\n fence_str = fence_bytes.decode(\"utf-8\")\n payload = RedisConnectorDeletePayload.model_validate_json(cast(str, fence_str))\n\n return payload\n\n def set_fence(self, payload: RedisConnectorDeletePayload | None) -> None:\n if not payload:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, payload.model_dump_json(), ex=self.FENCE_TTL)\n self.redis.sadd(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n\n def set_active(self) -> None:\n \"\"\"This sets a signal to keep the permissioning flow from getting cleaned up within\n the expiration time.\n\n The slack in timing is needed to avoid race conditions where simply checking\n the celery queue and task status could result in race conditions.\"\"\"\n self.redis.set(self.active_key, 0, ex=self.ACTIVE_TTL)\n\n def active(self) -> bool:\n return bool(self.redis.exists(self.active_key))\n\n def _generate_task_id(self) -> str:\n # celery's default task id format is \"dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # we prefix the task id so it's easier to keep track of who created the task\n # aka \"connectordeletion_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n\n return f\"{self.PREFIX}_{self.id}_{uuid4()}\"\n\n def generate_tasks(\n self,\n celery_app: Celery,\n db_session: Session,\n lock: RedisLock,\n ) -> int | None:\n \"\"\"Returns None if the cc_pair doesn't exist.\n Otherwise, returns an int with the number of generated tasks.\"\"\"\n last_lock_time = time.monotonic()\n\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=int(self.id),\n )\n if not cc_pair:\n return None\n\n num_tasks_sent = 0\n\n stmt = construct_document_id_select_for_connector_credential_pair(\n cc_pair.connector_id, cc_pair.credential_id\n )\n for doc_id in db_session.scalars(stmt).yield_per(DB_YIELD_PER_DEFAULT):\n doc_id = cast(str, doc_id)\n current_time = time.monotonic()\n if current_time - last_lock_time >= (\n CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4\n ):\n lock.reacquire()\n last_lock_time = current_time\n\n custom_task_id = self._generate_task_id()\n\n # add to the tracking taskset in redis BEFORE creating the celery task.\n # note that for the moment we are using a single taskset key, not differentiated by cc_pair id\n self.redis.sadd(self.taskset_key, custom_task_id)\n self.redis.expire(self.taskset_key, self.TASKSET_TTL)\n\n # Priority on sync's triggered by new indexing should be medium\n celery_app.send_task(\n OnyxCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_TASK,\n kwargs=dict(\n document_id=doc_id,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n tenant_id=self.tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_DELETION,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n ignore_result=True,\n )\n\n num_tasks_sent += 1\n\n return num_tasks_sent\n\n def reset(self) -> None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.active_key)\n self.redis.delete(self.taskset_key)\n self.redis.delete(self.fence_key)\n\n @staticmethod\n def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:\n taskset_key = f\"{RedisConnectorDelete.TASKSET_PREFIX}_{id}\"\n r.srem(taskset_key, task_id)\n return\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n \"\"\"Deletes all redis values for all connectors\"\"\"\n for key in r.scan_iter(RedisConnectorDelete.ACTIVE_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorDelete.TASKSET_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorDelete.FENCE_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_connector_doc_perm_sync.py", "content": "import time\nfrom datetime import datetime\nfrom logging import Logger\nfrom typing import Any\nfrom typing import cast\nfrom typing import NamedTuple\n\nimport redis\nfrom pydantic import BaseModel\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.access.models import ElementExternalAccess\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\n\n\nclass PermissionSyncResult(NamedTuple):\n \"\"\"Result of a permission sync operation.\n\n Attributes:\n num_updated: Number of documents successfully updated\n num_errors: Number of documents that failed to update\n \"\"\"\n\n num_updated: int\n num_errors: int\n\n\nclass RedisConnectorPermissionSyncPayload(BaseModel):\n id: str\n submitted: datetime\n started: datetime | None\n celery_task_id: str | None\n\n\nclass RedisConnectorPermissionSync:\n \"\"\"Manages interactions with redis for doc permission sync tasks. Should only be accessed\n through RedisConnector.\"\"\"\n\n PREFIX = \"connectordocpermissionsync\"\n\n FENCE_PREFIX = f\"{PREFIX}_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n\n # phase 1 - geneartor task and progress signals\n GENERATORTASK_PREFIX = f\"{PREFIX}+generator\" # connectorpermissions+generator\n GENERATOR_PROGRESS_PREFIX = (\n PREFIX + \"_generator_progress\"\n ) # connectorpermissions_generator_progress\n GENERATOR_COMPLETE_PREFIX = (\n PREFIX + \"_generator_complete\"\n ) # connectorpermissions_generator_complete\n\n TASKSET_PREFIX = f\"{PREFIX}_taskset\" # connectorpermissions_taskset\n SUBTASK_PREFIX = f\"{PREFIX}+sub\" # connectorpermissions+sub\n\n # used to signal the overall workflow is still active\n # it's impossible to get the exact state of the system at a single point in time\n # so we need a signal with a TTL to bridge gaps in our checks\n ACTIVE_PREFIX = PREFIX + \"_active\"\n ACTIVE_TTL = CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT * 2\n\n def __init__(self, tenant_id: str, id: int, redis: redis.Redis) -> None:\n self.tenant_id: str = tenant_id\n self.id = id\n self.redis = redis\n\n self.fence_key: str = f\"{self.FENCE_PREFIX}_{id}\"\n self.generator_task_key = f\"{self.GENERATORTASK_PREFIX}_{id}\"\n self.generator_progress_key = f\"{self.GENERATOR_PROGRESS_PREFIX}_{id}\"\n self.generator_complete_key = f\"{self.GENERATOR_COMPLETE_PREFIX}_{id}\"\n\n self.taskset_key = f\"{self.TASKSET_PREFIX}_{id}\"\n\n self.subtask_prefix: str = f\"{self.SUBTASK_PREFIX}_{id}\"\n self.active_key = f\"{self.ACTIVE_PREFIX}_{id}\"\n\n def taskset_clear(self) -> None:\n self.redis.delete(self.taskset_key)\n\n def generator_clear(self) -> None:\n self.redis.delete(self.generator_progress_key)\n self.redis.delete(self.generator_complete_key)\n\n def get_remaining(self) -> int:\n remaining = cast(int, self.redis.scard(self.taskset_key))\n return remaining\n\n def get_active_task_count(self) -> int:\n \"\"\"Count of active permission sync tasks\"\"\"\n count = 0\n for _ in self.redis.sscan_iter(\n OnyxRedisConstants.ACTIVE_FENCES,\n RedisConnectorPermissionSync.FENCE_PREFIX + \"*\",\n count=SCAN_ITER_COUNT_DEFAULT,\n ):\n count += 1\n return count\n\n @property\n def fenced(self) -> bool:\n return bool(self.redis.exists(self.fence_key))\n\n @property\n def payload(self) -> RedisConnectorPermissionSyncPayload | None:\n # read related data and evaluate/print task progress\n fence_bytes = cast(Any, self.redis.get(self.fence_key))\n if fence_bytes is None:\n return None\n\n fence_str = fence_bytes.decode(\"utf-8\")\n payload = RedisConnectorPermissionSyncPayload.model_validate_json(\n cast(str, fence_str)\n )\n\n return payload\n\n def set_fence(\n self,\n payload: RedisConnectorPermissionSyncPayload | None,\n ) -> None:\n if not payload:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, payload.model_dump_json(), ex=self.FENCE_TTL)\n self.redis.sadd(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n\n def set_active(self) -> None:\n \"\"\"This sets a signal to keep the permissioning flow from getting cleaned up within\n the expiration time.\n\n The slack in timing is needed to avoid race conditions where simply checking\n the celery queue and task status could result in race conditions.\"\"\"\n self.redis.set(self.active_key, 0, ex=self.ACTIVE_TTL)\n\n def active(self) -> bool:\n return bool(self.redis.exists(self.active_key))\n\n @property\n def generator_complete(self) -> int | None:\n \"\"\"the fence payload is an int representing the starting number of\n permission sync tasks to be processed ... just after the generator completes.\"\"\"\n fence_bytes = self.redis.get(self.generator_complete_key)\n if fence_bytes is None:\n return None\n\n if fence_bytes == b\"None\":\n return None\n\n fence_int = int(cast(bytes, fence_bytes).decode())\n return fence_int\n\n @generator_complete.setter\n def generator_complete(self, payload: int | None) -> None:\n \"\"\"Set the payload to an int to set the fence, otherwise if None it will\n be deleted\"\"\"\n if payload is None:\n self.redis.delete(self.generator_complete_key)\n return\n\n self.redis.set(self.generator_complete_key, payload, ex=self.FENCE_TTL)\n\n def update_db(\n self,\n lock: RedisLock | None,\n new_permissions: list[ElementExternalAccess],\n source_string: str,\n connector_id: int,\n credential_id: int,\n task_logger: Logger | None = None,\n ) -> PermissionSyncResult:\n \"\"\"Update permissions for documents and hierarchy nodes.\n\n Returns:\n PermissionSyncResult containing counts of successful updates and errors\n \"\"\"\n last_lock_time = time.monotonic()\n\n element_update_permissions_fn = fetch_versioned_implementation(\n \"onyx.background.celery.tasks.doc_permission_syncing.tasks\",\n \"element_update_permissions\",\n )\n\n num_permissions = 0\n num_errors = 0\n # Create a task for each permission sync\n for permissions in new_permissions:\n current_time = time.monotonic()\n if lock and current_time - last_lock_time >= (\n CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4\n ):\n lock.reacquire()\n last_lock_time = current_time\n\n if (\n permissions.external_access.num_entries\n > permissions.external_access.MAX_NUM_ENTRIES\n ):\n if task_logger:\n num_users = len(permissions.external_access.external_user_emails)\n num_groups = len(\n permissions.external_access.external_user_group_ids\n )\n element_id = (\n permissions.doc_id\n if isinstance(permissions, DocExternalAccess)\n else permissions.raw_node_id\n )\n task_logger.warning(\n f\"Permissions length exceeded, skipping...: \"\n f\"{element_id} \"\n f\"{num_users=} {num_groups=} \"\n f\"{permissions.external_access.MAX_NUM_ENTRIES=}\"\n )\n continue\n\n # NOTE(rkuo): this used to fire a task instead of directly writing to the DB,\n # but the permissions can be excessively large if sent over the wire.\n # On the other hand, the downside of doing db updates here is that we can\n # block and fail if we can't make the calls to the DB ... but that's probably\n # a rare enough case to be acceptable.\n\n # This can internally exception due to db issues but still continue\n # Catch exceptions per-element to avoid breaking the entire sync\n try:\n element_update_permissions_fn(\n self.tenant_id,\n permissions,\n source_string,\n connector_id,\n credential_id,\n )\n\n num_permissions += 1\n except Exception:\n num_errors += 1\n if task_logger:\n element_id = (\n permissions.doc_id\n if isinstance(permissions, DocExternalAccess)\n else permissions.raw_node_id\n )\n task_logger.exception(\n f\"Failed to update permissions for element {element_id}\"\n )\n # Continue processing other elements\n\n return PermissionSyncResult(num_updated=num_permissions, num_errors=num_errors)\n\n def reset(self) -> None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.active_key)\n self.redis.delete(self.generator_progress_key)\n self.redis.delete(self.generator_complete_key)\n self.redis.delete(self.taskset_key)\n self.redis.delete(self.fence_key)\n\n @staticmethod\n def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:\n taskset_key = f\"{RedisConnectorPermissionSync.TASKSET_PREFIX}_{id}\"\n r.srem(taskset_key, task_id)\n return\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n \"\"\"Deletes all redis values for all connectors\"\"\"\n for key in r.scan_iter(RedisConnectorPermissionSync.ACTIVE_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorPermissionSync.TASKSET_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(\n RedisConnectorPermissionSync.GENERATOR_COMPLETE_PREFIX + \"*\"\n ):\n r.delete(key)\n\n for key in r.scan_iter(\n RedisConnectorPermissionSync.GENERATOR_PROGRESS_PREFIX + \"*\"\n ):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorPermissionSync.FENCE_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_connector_ext_group_sync.py", "content": "from datetime import datetime\nfrom typing import cast\n\nimport redis\nfrom celery import Celery\nfrom pydantic import BaseModel\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT\n\n\nclass RedisConnectorExternalGroupSyncPayload(BaseModel):\n id: str\n submitted: datetime\n started: datetime | None\n celery_task_id: str | None\n\n\nclass RedisConnectorExternalGroupSync:\n \"\"\"Manages interactions with redis for external group syncing tasks. Should only be accessed\n through RedisConnector.\"\"\"\n\n PREFIX = \"connectorexternalgroupsync\"\n\n FENCE_PREFIX = f\"{PREFIX}_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n\n # phase 1 - geneartor task and progress signals\n GENERATORTASK_PREFIX = f\"{PREFIX}+generator\" # connectorexternalgroupsync+generator\n GENERATOR_PROGRESS_PREFIX = (\n PREFIX + \"_generator_progress\"\n ) # connectorexternalgroupsync_generator_progress\n GENERATOR_COMPLETE_PREFIX = (\n PREFIX + \"_generator_complete\"\n ) # connectorexternalgroupsync_generator_complete\n\n TASKSET_PREFIX = f\"{PREFIX}_taskset\" # connectorexternalgroupsync_taskset\n SUBTASK_PREFIX = f\"{PREFIX}+sub\" # connectorexternalgroupsync+sub\n\n # used to signal the overall workflow is still active\n # it's impossible to get the exact state of the system at a single point in time\n # so we need a signal with a TTL to bridge gaps in our checks\n ACTIVE_PREFIX = PREFIX + \"_active\"\n ACTIVE_TTL = 3600\n\n def __init__(self, tenant_id: str, id: int, redis: redis.Redis) -> None:\n self.tenant_id: str = tenant_id\n self.id = id\n self.redis = redis\n\n self.fence_key: str = f\"{self.FENCE_PREFIX}_{id}\"\n self.generator_task_key = f\"{self.GENERATORTASK_PREFIX}_{id}\"\n self.generator_progress_key = f\"{self.GENERATOR_PROGRESS_PREFIX}_{id}\"\n self.generator_complete_key = f\"{self.GENERATOR_COMPLETE_PREFIX}_{id}\"\n\n self.taskset_key = f\"{self.TASKSET_PREFIX}_{id}\"\n\n self.subtask_prefix: str = f\"{self.SUBTASK_PREFIX}_{id}\"\n self.active_key = f\"{self.ACTIVE_PREFIX}_{id}\"\n\n def taskset_clear(self) -> None:\n self.redis.delete(self.taskset_key)\n\n def generator_clear(self) -> None:\n self.redis.delete(self.generator_progress_key)\n self.redis.delete(self.generator_complete_key)\n\n def get_remaining(self) -> int:\n # todo: move into fence\n remaining = cast(int, self.redis.scard(self.taskset_key))\n return remaining\n\n def get_active_task_count(self) -> int:\n \"\"\"Count of active external group syncing tasks\"\"\"\n count = 0\n for _ in self.redis.sscan_iter(\n OnyxRedisConstants.ACTIVE_FENCES,\n RedisConnectorExternalGroupSync.FENCE_PREFIX + \"*\",\n count=SCAN_ITER_COUNT_DEFAULT,\n ):\n count += 1\n return count\n\n @property\n def fenced(self) -> bool:\n return bool(self.redis.exists(self.fence_key))\n\n @property\n def payload(self) -> RedisConnectorExternalGroupSyncPayload | None:\n # read related data and evaluate/print task progress\n fence_raw = self.redis.get(self.fence_key)\n if fence_raw is None:\n return None\n\n fence_bytes = cast(bytes, fence_raw)\n fence_str = fence_bytes.decode(\"utf-8\")\n payload = RedisConnectorExternalGroupSyncPayload.model_validate_json(\n cast(str, fence_str)\n )\n\n return payload\n\n def set_fence(\n self,\n payload: RedisConnectorExternalGroupSyncPayload | None,\n ) -> None:\n if not payload:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, payload.model_dump_json(), ex=self.FENCE_TTL)\n self.redis.sadd(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n\n def set_active(self) -> None:\n \"\"\"This sets a signal to keep the permissioning flow from getting cleaned up within\n the expiration time.\n\n The slack in timing is needed to avoid race conditions where simply checking\n the celery queue and task status could result in race conditions.\"\"\"\n self.redis.set(self.active_key, 0, ex=self.ACTIVE_TTL)\n\n def active(self) -> bool:\n return bool(self.redis.exists(self.active_key))\n\n @property\n def generator_complete(self) -> int | None:\n \"\"\"the fence payload is an int representing the starting number of\n external group syncing tasks to be processed ... just after the generator completes.\n \"\"\"\n fence_bytes = self.redis.get(self.generator_complete_key)\n if fence_bytes is None:\n return None\n\n if fence_bytes == b\"None\":\n return None\n\n fence_int = int(cast(bytes, fence_bytes).decode())\n return fence_int\n\n @generator_complete.setter\n def generator_complete(self, payload: int | None) -> None:\n \"\"\"Set the payload to an int to set the fence, otherwise if None it will\n be deleted\"\"\"\n if payload is None:\n self.redis.delete(self.generator_complete_key)\n return\n\n self.redis.set(self.generator_complete_key, payload, ex=self.FENCE_TTL)\n\n def generate_tasks(\n self,\n celery_app: Celery,\n db_session: Session,\n lock: RedisLock | None,\n ) -> int | None:\n pass\n\n def reset(self) -> None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.active_key)\n self.redis.delete(self.generator_progress_key)\n self.redis.delete(self.generator_complete_key)\n self.redis.delete(self.taskset_key)\n self.redis.delete(self.fence_key)\n\n @staticmethod\n def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:\n taskset_key = f\"{RedisConnectorExternalGroupSync.TASKSET_PREFIX}_{id}\"\n r.srem(taskset_key, task_id)\n return\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n \"\"\"Deletes all redis values for all connectors\"\"\"\n for key in r.scan_iter(RedisConnectorExternalGroupSync.ACTIVE_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorExternalGroupSync.TASKSET_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(\n RedisConnectorExternalGroupSync.GENERATOR_COMPLETE_PREFIX + \"*\"\n ):\n r.delete(key)\n\n for key in r.scan_iter(\n RedisConnectorExternalGroupSync.GENERATOR_PROGRESS_PREFIX + \"*\"\n ):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorExternalGroupSync.FENCE_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_connector_index.py", "content": "from datetime import datetime\n\nfrom pydantic import BaseModel\n\n\nclass RedisConnectorIndexPayload(BaseModel):\n index_attempt_id: int | None\n started: datetime | None\n submitted: datetime\n celery_task_id: str | None\n" }, { "path": "backend/onyx/redis/redis_connector_prune.py", "content": "import time\nfrom datetime import datetime\nfrom typing import cast\nfrom uuid import uuid4\n\nimport redis\nfrom celery import Celery\nfrom pydantic import BaseModel\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import CELERY_PRUNING_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT\n\n\nclass RedisConnectorPrunePayload(BaseModel):\n id: str\n submitted: datetime\n started: datetime | None\n celery_task_id: str | None\n\n\nclass RedisConnectorPrune:\n \"\"\"Manages interactions with redis for pruning tasks. Should only be accessed\n through RedisConnector.\"\"\"\n\n PREFIX = \"connectorpruning\"\n\n FENCE_PREFIX = f\"{PREFIX}_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n\n # phase 1 - geneartor task and progress signals\n GENERATORTASK_PREFIX = f\"{PREFIX}+generator\" # connectorpruning+generator\n GENERATOR_PROGRESS_PREFIX = (\n PREFIX + \"_generator_progress\"\n ) # connectorpruning_generator_progress\n GENERATOR_COMPLETE_PREFIX = (\n PREFIX + \"_generator_complete\"\n ) # connectorpruning_generator_complete\n\n TASKSET_PREFIX = f\"{PREFIX}_taskset\" # connectorpruning_taskset\n TASKSET_TTL = FENCE_TTL\n SUBTASK_PREFIX = f\"{PREFIX}+sub\" # connectorpruning+sub\n\n # used to signal the overall workflow is still active\n # it's impossible to get the exact state of the system at a single point in time\n # so we need a signal with a TTL to bridge gaps in our checks\n ACTIVE_PREFIX = PREFIX + \"_active\"\n ACTIVE_TTL = CELERY_PRUNING_LOCK_TIMEOUT * 2\n\n def __init__(self, tenant_id: str, id: int, redis: redis.Redis) -> None:\n self.tenant_id: str = tenant_id\n self.id = id\n self.redis = redis\n\n self.fence_key: str = f\"{self.FENCE_PREFIX}_{id}\"\n self.generator_task_key = f\"{self.GENERATORTASK_PREFIX}_{id}\"\n self.generator_progress_key = f\"{self.GENERATOR_PROGRESS_PREFIX}_{id}\"\n self.generator_complete_key = f\"{self.GENERATOR_COMPLETE_PREFIX}_{id}\"\n\n self.taskset_key = f\"{self.TASKSET_PREFIX}_{id}\"\n\n self.subtask_prefix: str = f\"{self.SUBTASK_PREFIX}_{id}\"\n self.active_key = f\"{self.ACTIVE_PREFIX}_{id}\"\n\n def taskset_clear(self) -> None:\n self.redis.delete(self.taskset_key)\n\n def generator_clear(self) -> None:\n self.redis.delete(self.generator_progress_key)\n self.redis.delete(self.generator_complete_key)\n\n def get_remaining(self) -> int:\n # todo: move into fence\n remaining = cast(int, self.redis.scard(self.taskset_key))\n return remaining\n\n def get_active_task_count(self) -> int:\n \"\"\"Count of active pruning tasks\"\"\"\n count = 0\n for _ in self.redis.sscan_iter(\n OnyxRedisConstants.ACTIVE_FENCES,\n RedisConnectorPrune.FENCE_PREFIX + \"*\",\n count=SCAN_ITER_COUNT_DEFAULT,\n ):\n count += 1\n return count\n\n @property\n def fenced(self) -> bool:\n return bool(self.redis.exists(self.fence_key))\n\n @property\n def payload(self) -> RedisConnectorPrunePayload | None:\n # read related data and evaluate/print task progress\n fence_bytes = cast(bytes, self.redis.get(self.fence_key))\n if fence_bytes is None:\n return None\n\n fence_str = fence_bytes.decode(\"utf-8\")\n payload = RedisConnectorPrunePayload.model_validate_json(cast(str, fence_str))\n\n return payload\n\n def set_fence(\n self,\n payload: RedisConnectorPrunePayload | None,\n ) -> None:\n if not payload:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, payload.model_dump_json(), ex=self.FENCE_TTL)\n self.redis.sadd(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n\n def set_active(self) -> None:\n \"\"\"This sets a signal to keep the permissioning flow from getting cleaned up within\n the expiration time.\n\n The slack in timing is needed to avoid race conditions where simply checking\n the celery queue and task status could result in race conditions.\"\"\"\n self.redis.set(self.active_key, 0, ex=self.ACTIVE_TTL)\n\n def active(self) -> bool:\n return bool(self.redis.exists(self.active_key))\n\n @property\n def generator_complete(self) -> int | None:\n \"\"\"the fence payload is an int representing the starting number of\n pruning tasks to be processed ... just after the generator completes.\"\"\"\n fence_bytes = self.redis.get(self.generator_complete_key)\n if fence_bytes is None:\n return None\n\n fence_int = int(cast(bytes, fence_bytes))\n return fence_int\n\n @generator_complete.setter\n def generator_complete(self, payload: int | None) -> None:\n \"\"\"Set the payload to an int to set the fence, otherwise if None it will\n be deleted\"\"\"\n if payload is None:\n self.redis.delete(self.generator_complete_key)\n return\n\n self.redis.set(self.generator_complete_key, payload, ex=self.FENCE_TTL)\n\n def generate_tasks(\n self,\n documents_to_prune: set[str],\n celery_app: Celery,\n db_session: Session,\n lock: RedisLock | None,\n ) -> int | None:\n last_lock_time = time.monotonic()\n\n async_results = []\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=int(self.id),\n )\n if not cc_pair:\n return None\n\n for doc_id in documents_to_prune:\n current_time = time.monotonic()\n if lock and current_time - last_lock_time >= (\n CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4\n ):\n lock.reacquire()\n last_lock_time = current_time\n\n # celery's default task id format is \"dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # the actual redis key is \"celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # we prefix the task id so it's easier to keep track of who created the task\n # aka \"documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n custom_task_id = f\"{self.subtask_prefix}_{uuid4()}\"\n\n # add to the tracking taskset in redis BEFORE creating the celery task.\n self.redis.sadd(self.taskset_key, custom_task_id)\n self.redis.expire(self.taskset_key, self.TASKSET_TTL)\n\n # Priority on sync's triggered by new indexing should be medium\n result = celery_app.send_task(\n OnyxCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_TASK,\n kwargs=dict(\n document_id=doc_id,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n tenant_id=self.tenant_id,\n ),\n queue=OnyxCeleryQueues.CONNECTOR_DELETION,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n ignore_result=True,\n )\n\n async_results.append(result)\n\n return len(async_results)\n\n def reset(self) -> None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.active_key)\n self.redis.delete(self.generator_progress_key)\n self.redis.delete(self.generator_complete_key)\n self.redis.delete(self.taskset_key)\n self.redis.delete(self.fence_key)\n\n @staticmethod\n def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:\n taskset_key = f\"{RedisConnectorPrune.TASKSET_PREFIX}_{id}\"\n r.srem(taskset_key, task_id)\n return\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n \"\"\"Deletes all redis values for all connectors\"\"\"\n for key in r.scan_iter(RedisConnectorPrune.ACTIVE_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorPrune.TASKSET_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorPrune.GENERATOR_COMPLETE_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorPrune.GENERATOR_PROGRESS_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorPrune.FENCE_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_connector_stop.py", "content": "import redis\n\n\nclass RedisConnectorStop:\n \"\"\"Manages interactions with redis for stop signaling. Should only be accessed\n through RedisConnector.\"\"\"\n\n PREFIX = \"connectorstop\"\n FENCE_PREFIX = f\"{PREFIX}_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n\n # if this timeout is exceeded, the caller may decide to take more\n # drastic measures\n TIMEOUT_PREFIX = f\"{PREFIX}_timeout\"\n TIMEOUT_TTL = 300\n\n def __init__(self, tenant_id: str, id: int, redis: redis.Redis) -> None:\n self.tenant_id: str = tenant_id\n self.id: int = id\n self.redis = redis\n\n self.fence_key: str = f\"{self.FENCE_PREFIX}_{id}\"\n self.timeout_key: str = f\"{self.TIMEOUT_PREFIX}_{id}\"\n\n @property\n def fenced(self) -> bool:\n return bool(self.redis.exists(self.fence_key))\n\n def set_fence(self, value: bool) -> None:\n if not value:\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, 0, ex=self.FENCE_TTL)\n\n @property\n def timed_out(self) -> bool:\n return not bool(self.redis.exists(self.timeout_key))\n\n def set_timeout(self) -> None:\n \"\"\"After calling this, call timed_out to determine if the timeout has been\n exceeded.\"\"\"\n self.redis.set(f\"{self.timeout_key}\", 0, ex=self.TIMEOUT_TTL)\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n for key in r.scan_iter(RedisConnectorStop.FENCE_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisConnectorStop.TIMEOUT_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_connector_utils.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import TaskStatus\nfrom onyx.db.models import TaskQueueState\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.server.documents.models import DeletionAttemptSnapshot\n\n\ndef _get_deletion_status(\n connector_id: int,\n credential_id: int,\n db_session: Session,\n tenant_id: str,\n) -> TaskQueueState | None:\n \"\"\"We no longer store TaskQueueState in the DB for a deletion attempt.\n This function populates TaskQueueState by just checking redis.\n \"\"\"\n cc_pair = get_connector_credential_pair(\n connector_id=connector_id, credential_id=credential_id, db_session=db_session\n )\n if not cc_pair:\n return None\n\n redis_connector = RedisConnector(tenant_id, cc_pair.id)\n if redis_connector.delete.fenced:\n return TaskQueueState(\n task_id=\"\",\n task_name=redis_connector.delete.fence_key,\n status=TaskStatus.STARTED,\n )\n\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n return TaskQueueState(\n task_id=\"\",\n task_name=redis_connector.delete.fence_key,\n status=TaskStatus.PENDING,\n )\n\n return None\n\n\ndef get_deletion_attempt_snapshot(\n connector_id: int,\n credential_id: int,\n db_session: Session,\n tenant_id: str,\n) -> DeletionAttemptSnapshot | None:\n deletion_task = _get_deletion_status(\n connector_id, credential_id, db_session, tenant_id\n )\n if not deletion_task:\n return None\n\n return DeletionAttemptSnapshot(\n connector_id=connector_id,\n credential_id=credential_id,\n status=deletion_task.status,\n )\n" }, { "path": "backend/onyx/redis/redis_document_set.py", "content": "import time\nfrom typing import cast\nfrom uuid import uuid4\n\nimport redis\nfrom celery import Celery\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DB_YIELD_PER_DEFAULT\nfrom onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.db.document_set import construct_document_id_select_by_docset\nfrom onyx.redis.redis_object_helper import RedisObjectHelper\n\n\nclass RedisDocumentSet(RedisObjectHelper):\n PREFIX = \"documentset\"\n FENCE_PREFIX = PREFIX + \"_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n TASKSET_PREFIX = PREFIX + \"_taskset\"\n TASKSET_TTL = FENCE_TTL\n\n def __init__(self, tenant_id: str, id: int) -> None:\n super().__init__(tenant_id, str(id))\n\n @property\n def fenced(self) -> bool:\n return bool(self.redis.exists(self.fence_key))\n\n def set_fence(self, payload: int | None) -> None:\n if payload is None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, payload, ex=self.FENCE_TTL)\n self.redis.sadd(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n\n @property\n def payload(self) -> int | None:\n bytes = self.redis.get(self.fence_key)\n if bytes is None:\n return None\n\n progress = int(cast(int, bytes))\n return progress\n\n def generate_tasks(\n self,\n max_tasks: int, # noqa: ARG002\n celery_app: Celery,\n db_session: Session,\n redis_client: Redis,\n lock: RedisLock,\n tenant_id: str,\n ) -> tuple[int, int] | None:\n \"\"\"Max tasks is ignored for now until we can build the logic to mark the\n document set up to date over multiple batches.\n \"\"\"\n last_lock_time = time.monotonic()\n\n num_tasks_sent = 0\n\n stmt = construct_document_id_select_by_docset(int(self._id), current_only=False)\n for doc_id in db_session.scalars(stmt).yield_per(DB_YIELD_PER_DEFAULT):\n doc_id = cast(str, doc_id)\n current_time = time.monotonic()\n if current_time - last_lock_time >= (\n CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4\n ):\n lock.reacquire()\n last_lock_time = current_time\n\n # celery's default task id format is \"dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # the key for the result is \"celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # we prefix the task id so it's easier to keep track of who created the task\n # aka \"documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n custom_task_id = f\"{self.task_id_prefix}_{uuid4()}\"\n\n # add to the set BEFORE creating the task.\n redis_client.sadd(self.taskset_key, custom_task_id)\n redis_client.expire(self.taskset_key, self.TASKSET_TTL)\n\n celery_app.send_task(\n OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,\n kwargs=dict(document_id=doc_id, tenant_id=tenant_id),\n queue=OnyxCeleryQueues.VESPA_METADATA_SYNC,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n )\n\n num_tasks_sent += 1\n\n return num_tasks_sent, num_tasks_sent\n\n def reset(self) -> None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.taskset_key)\n self.redis.delete(self.fence_key)\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n for key in r.scan_iter(RedisDocumentSet.TASKSET_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_hierarchy.py", "content": "\"\"\"Redis cache operations for hierarchy node ancestor resolution.\n\nThis module provides a Redis-based cache for hierarchy node parent relationships,\nenabling fast ancestor path resolution without repeated database queries.\n\nThe cache stores node_id -> parent_id mappings for all hierarchy nodes of a given\nsource type. When resolving ancestors for a document, we walk up the tree using\nRedis lookups instead of database queries.\n\nCache Strategy:\n- Nodes are cached per source type with a 6-hour TTL\n- During docfetching, nodes are added to cache as they're upserted to Postgres\n- If the cache is stale (TTL expired during long-running job), one worker does\n a full refresh from DB while others wait\n- If a node is still not found after refresh, we log an error and fall back to\n using only the SOURCE-type node as the ancestor\n\"\"\"\n\nfrom typing import cast\nfrom typing import TYPE_CHECKING\n\nfrom pydantic import BaseModel\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.db.hierarchy import ensure_source_node_exists as db_ensure_source_node_exists\nfrom onyx.db.hierarchy import get_all_hierarchy_nodes_for_source\nfrom onyx.utils.logger import setup_logger\n\nif TYPE_CHECKING:\n from onyx.db.models import HierarchyNode as DBHierarchyNode\n\nlogger = setup_logger()\n\n# Cache TTL: 6 hours in seconds\nHIERARCHY_CACHE_TTL_SECONDS = 6 * 60 * 60\n\n# Lock timeout for cache refresh: 5 minutes\nHIERARCHY_CACHE_LOCK_TIMEOUT_SECONDS = 5 * 60\n\n# Lock acquisition timeout: 60 seconds\nHIERARCHY_CACHE_LOCK_ACQUIRE_TIMEOUT_SECONDS = 60\n\nMAX_DEPTH = 1000\n\n\nclass HierarchyNodeCacheEntry(BaseModel):\n \"\"\"Represents a hierarchy node for caching purposes.\"\"\"\n\n node_id: int\n parent_id: int | None\n node_type: HierarchyNodeType\n raw_node_id: str\n\n @classmethod\n def from_db_model(cls, node: \"DBHierarchyNode\") -> \"HierarchyNodeCacheEntry\":\n \"\"\"Create a cache entry from a SQLAlchemy HierarchyNode model.\"\"\"\n return cls(\n node_id=node.id,\n parent_id=node.parent_id,\n node_type=node.node_type,\n raw_node_id=node.raw_node_id,\n )\n\n\ndef _cache_key(source: DocumentSource) -> str:\n \"\"\"Get the Redis hash key for hierarchy node cache of a given source.\n\n This hash stores: node_id -> \"parent_id:node_type\"\n \"\"\"\n return f\"hierarchy_cache:{source.value}\"\n\n\ndef _raw_id_cache_key(source: DocumentSource) -> str:\n \"\"\"Get the Redis hash key for raw_node_id -> node_id mapping.\n\n This hash stores: raw_node_id -> node_id\n \"\"\"\n return f\"hierarchy_cache_rawid:{source.value}\"\n\n\ndef _source_node_key(source: DocumentSource) -> str:\n \"\"\"Get the Redis key for the SOURCE-type node ID of a given source.\n\n This is a simple string key storing the database ID of the SOURCE node.\n \"\"\"\n return f\"hierarchy_source_node:{source.value}\"\n\n\ndef _loading_lock_key(source: DocumentSource) -> str:\n \"\"\"Get the Redis lock key for cache loading of a given source.\"\"\"\n return f\"hierarchy_cache_loading:{source.value}\"\n\n\ndef _construct_parent_value(parent_id: int | None, node_type: HierarchyNodeType) -> str:\n \"\"\"Construct the cached value string from parent_id and node_type.\n\n Format: \"parent_id:node_type\" where parent_id is empty string if None.\n \"\"\"\n parent_str = str(parent_id) if parent_id is not None else \"\"\n return f\"{parent_str}:{node_type.value}\"\n\n\ndef _unpack_parent_value(value: str) -> tuple[int | None, HierarchyNodeType | None]:\n \"\"\"Unpack a cached value string back into (parent_id, node_type).\n\n Returns None for invalid values.\n \"\"\"\n parts = value.split(\":\", 1)\n parent_str = parts[0]\n node_type_str = parts[1] if len(parts) > 1 else \"\"\n parent_id = int(parent_str) if parent_str else None\n\n node_type = HierarchyNodeType(node_type_str) if node_type_str else None\n\n return parent_id, node_type\n\n\ndef cache_hierarchy_node(\n redis_client: Redis,\n source: DocumentSource,\n entry: HierarchyNodeCacheEntry,\n) -> None:\n \"\"\"\n Add or update a single hierarchy node in the Redis cache.\n\n Called during docfetching when nodes are upserted to Postgres.\n Stores the parent chain mapping, raw_id -> node_id mapping, and\n SOURCE node ID (if this is a SOURCE-type node).\n\n Args:\n redis_client: Redis client with tenant prefixing\n source: The document source (e.g., CONFLUENCE, GOOGLE_DRIVE)\n entry: The hierarchy node cache entry\n \"\"\"\n cache_key = _cache_key(source)\n raw_id_key = _raw_id_cache_key(source)\n\n # Store parent chain: node_id -> \"parent_id:node_type\"\n value = _construct_parent_value(entry.parent_id, entry.node_type)\n redis_client.hset(cache_key, str(entry.node_id), value)\n\n # Store raw_id -> node_id mapping\n redis_client.hset(raw_id_key, entry.raw_node_id, str(entry.node_id))\n\n # If this is the SOURCE node, store its ID in the dedicated key\n if entry.node_type == HierarchyNodeType.SOURCE:\n source_node_key = _source_node_key(source)\n redis_client.set(source_node_key, str(entry.node_id))\n redis_client.expire(source_node_key, HIERARCHY_CACHE_TTL_SECONDS)\n\n # Refresh TTL on every write (ensures cache stays alive during long indexing)\n redis_client.expire(cache_key, HIERARCHY_CACHE_TTL_SECONDS)\n redis_client.expire(raw_id_key, HIERARCHY_CACHE_TTL_SECONDS)\n\n\ndef cache_hierarchy_nodes_batch(\n redis_client: Redis,\n source: DocumentSource,\n entries: list[HierarchyNodeCacheEntry],\n) -> None:\n \"\"\"\n Add or update multiple hierarchy nodes in the Redis cache.\n\n Args:\n redis_client: Redis client with tenant prefixing\n source: The document source\n entries: List of HierarchyNodeCacheEntry objects\n \"\"\"\n if not entries:\n return\n\n cache_key = _cache_key(source)\n raw_id_key = _raw_id_cache_key(source)\n source_node_key = _source_node_key(source)\n\n # Build mappings for batch insert\n parent_mapping: dict[str, str] = {}\n raw_id_mapping: dict[str, str] = {}\n source_node_id: int | None = None\n\n for entry in entries:\n parent_mapping[str(entry.node_id)] = _construct_parent_value(\n entry.parent_id, entry.node_type\n )\n raw_id_mapping[entry.raw_node_id] = str(entry.node_id)\n\n # Track the SOURCE node if we encounter it\n if entry.node_type == HierarchyNodeType.SOURCE:\n source_node_id = entry.node_id\n\n # Use hset with mapping for batch insert\n redis_client.hset(cache_key, mapping=parent_mapping)\n redis_client.hset(raw_id_key, mapping=raw_id_mapping)\n\n # Cache the SOURCE node ID if found\n if source_node_id is not None:\n redis_client.set(source_node_key, str(source_node_id))\n redis_client.expire(source_node_key, HIERARCHY_CACHE_TTL_SECONDS)\n\n redis_client.expire(cache_key, HIERARCHY_CACHE_TTL_SECONDS)\n redis_client.expire(raw_id_key, HIERARCHY_CACHE_TTL_SECONDS)\n\n\ndef evict_hierarchy_nodes_from_cache(\n redis_client: Redis,\n source: DocumentSource,\n raw_node_ids: list[str],\n) -> None:\n \"\"\"Remove specific hierarchy nodes from the Redis cache.\n\n Deletes entries from both the parent-chain hash and the raw_id→node_id hash.\n \"\"\"\n if not raw_node_ids:\n return\n\n cache_key = _cache_key(source)\n raw_id_key = _raw_id_cache_key(source)\n\n # Look up node_ids so we can remove them from the parent-chain hash\n raw_values = cast(list[str | None], redis_client.hmget(raw_id_key, raw_node_ids))\n node_id_strs = [v for v in raw_values if v is not None]\n\n if node_id_strs:\n redis_client.hdel(cache_key, *node_id_strs)\n redis_client.hdel(raw_id_key, *raw_node_ids)\n\n\ndef get_node_id_from_raw_id(\n redis_client: Redis,\n source: DocumentSource,\n raw_node_id: str,\n) -> tuple[int | None, bool]:\n \"\"\"\n Get the database node_id for a raw_node_id from the cache.\n\n Returns:\n Tuple of (node_id or None, found_in_cache)\n - If found_in_cache is False, the raw_id doesn't exist in cache\n - If found_in_cache is True, node_id is the database ID\n \"\"\"\n raw_id_key = _raw_id_cache_key(source)\n value = redis_client.hget(raw_id_key, raw_node_id)\n\n if value is None:\n return None, False\n\n # Decode bytes if needed\n value_str: str\n if isinstance(value, bytes):\n value_str = value.decode(\"utf-8\")\n else:\n value_str = str(value)\n\n return int(value_str), True\n\n\ndef get_parent_id_from_cache(\n redis_client: Redis,\n source: DocumentSource,\n node_id: int,\n) -> tuple[int | None, bool]:\n \"\"\"\n Get the parent_id for a node from the cache.\n\n Returns:\n Tuple of (parent_id or None, found_in_cache)\n - If found_in_cache is False, the node doesn't exist in cache\n - If found_in_cache is True, parent_id is the actual parent (or None for root)\n \"\"\"\n cache_key = _cache_key(source)\n value = redis_client.hget(cache_key, str(node_id))\n\n if value is None:\n return None, False\n\n # Decode bytes if needed\n value_str: str\n if isinstance(value, bytes):\n value_str = value.decode(\"utf-8\")\n else:\n value_str = str(value)\n\n parent_id, _ = _unpack_parent_value(value_str)\n return parent_id, True\n\n\ndef is_cache_populated(redis_client: Redis, source: DocumentSource) -> bool:\n \"\"\"Check if the cache has any entries for this source.\"\"\"\n cache_key = _cache_key(source)\n # redis.exists returns int (number of keys that exist)\n exists_result: int = redis_client.exists(cache_key) # type: ignore[assignment]\n return exists_result > 0\n\n\ndef refresh_hierarchy_cache_from_db(\n redis_client: Redis,\n db_session: Session,\n source: DocumentSource,\n) -> None:\n \"\"\"\n Refresh the entire hierarchy cache for a source from the database.\n\n This function acquires a distributed lock to ensure only one worker\n performs the refresh. Other workers will wait for the refresh to complete.\n\n Args:\n redis_client: Redis client with tenant prefixing\n db_session: SQLAlchemy session for database access\n source: The document source to refresh\n \"\"\"\n\n lock_key = _loading_lock_key(source)\n\n # Try to acquire lock - if we can't get it, someone else is refreshing\n lock: RedisLock = redis_client.lock(\n lock_key,\n timeout=HIERARCHY_CACHE_LOCK_TIMEOUT_SECONDS,\n blocking=True,\n blocking_timeout=HIERARCHY_CACHE_LOCK_ACQUIRE_TIMEOUT_SECONDS,\n )\n\n acquired = lock.acquire(blocking=True)\n if not acquired:\n logger.warning(\n f\"Could not acquire lock for hierarchy cache refresh for source {source.value} - another worker may be refreshing\"\n )\n return\n\n try:\n # Always refresh from DB when called - new nodes may have been added\n # since the cache was last populated. The lock ensures only one worker\n # does the refresh at a time.\n logger.info(f\"Refreshing hierarchy cache for source {source.value} from DB\")\n\n # Load all nodes for this source from DB\n nodes = get_all_hierarchy_nodes_for_source(db_session, source)\n\n if not nodes:\n logger.warning(f\"No hierarchy nodes found in DB for source {source.value}\")\n return\n\n # Batch insert into cache\n cache_entries = [HierarchyNodeCacheEntry.from_db_model(node) for node in nodes]\n cache_hierarchy_nodes_batch(redis_client, source, cache_entries)\n\n logger.info(\n f\"Refreshed hierarchy cache for {source.value} with {len(nodes)} nodes\"\n )\n\n finally:\n try:\n lock.release()\n except Exception as e:\n logger.warning(f\"Error releasing hierarchy cache lock: {e}\")\n\n\ndef _walk_ancestor_chain(\n redis_client: Redis,\n source: DocumentSource,\n start_node_id: int,\n db_session: Session,\n) -> list[int]:\n \"\"\"\n Walk up the hierarchy tree from a node, collecting all ancestor IDs.\n\n Internal helper used by both get_ancestors_from_node_id and\n get_ancestors_from_raw_id.\n \"\"\"\n ancestors: list[int] = []\n current_id: int | None = start_node_id\n visited: set[int] = set()\n\n while current_id is not None and len(ancestors) < MAX_DEPTH:\n if current_id in visited:\n logger.error(\n f\"Cycle detected in hierarchy for source {source.value} at node {current_id}. Ancestors so far: {ancestors}\"\n )\n break\n\n visited.add(current_id)\n ancestors.append(current_id)\n\n parent_id, found = get_parent_id_from_cache(redis_client, source, current_id)\n\n if not found:\n logger.debug(\n f\"Cache miss for hierarchy node {current_id} of source {source.value}, attempting refresh\"\n )\n refresh_hierarchy_cache_from_db(redis_client, db_session, source)\n parent_id, found = get_parent_id_from_cache(\n redis_client, source, current_id\n )\n\n if not found:\n logger.error(\n f\"Hierarchy node {current_id} not found in cache for source {source.value} even after refresh.\"\n )\n break\n\n current_id = parent_id\n\n if len(ancestors) >= MAX_DEPTH:\n logger.error(\n f\"Hit max depth {MAX_DEPTH} traversing hierarchy for source \"\n f\"{source.value}. Possible infinite loop or very deep hierarchy.\"\n )\n\n return ancestors\n\n\ndef get_ancestors_from_raw_id(\n redis_client: Redis,\n source: DocumentSource,\n parent_hierarchy_raw_node_id: str | None,\n db_session: Session,\n) -> list[int]:\n \"\"\"\n Get all ancestor hierarchy node IDs from a raw_node_id.\n\n This is the main entry point for getting ancestors from a document's\n parent_hierarchy_raw_node_id. It resolves the raw_id to a database ID\n via Redis cache, then walks up the tree.\n\n No DB calls are made unless the cache is stale.\n\n Args:\n redis_client: Redis client with tenant prefixing\n source: The document source\n parent_hierarchy_raw_node_id: The document's parent raw node ID (from connector)\n db_session: DB session for cache refresh if needed\n\n Returns:\n List of ancestor hierarchy node IDs from parent to root (inclusive).\n Returns list with just SOURCE node ID if parent is None or not found.\n \"\"\"\n # If no parent specified, return just the SOURCE node\n if parent_hierarchy_raw_node_id is None:\n source_node_id = get_source_node_id_from_cache(redis_client, db_session, source)\n return [source_node_id] if source_node_id else []\n\n # Resolve raw_id to node_id via Redis\n node_id, found = get_node_id_from_raw_id(\n redis_client, source, parent_hierarchy_raw_node_id\n )\n\n if not found:\n # Cache miss - try refresh\n logger.debug(\n f\"Cache miss for raw_node_id '{parent_hierarchy_raw_node_id}' of source {source.value}, attempting refresh\"\n )\n refresh_hierarchy_cache_from_db(redis_client, db_session, source)\n node_id, found = get_node_id_from_raw_id(\n redis_client, source, parent_hierarchy_raw_node_id\n )\n\n if not found or node_id is None:\n logger.error(\n f\"Raw node ID '{parent_hierarchy_raw_node_id}' not found in cache \"\n f\"for source {source.value}. Falling back to SOURCE node only.\"\n )\n source_node_id = get_source_node_id_from_cache(redis_client, db_session, source)\n return [source_node_id] if source_node_id else []\n\n # Walk up the ancestor chain\n return _walk_ancestor_chain(redis_client, source, node_id, db_session)\n\n\ndef get_source_node_id_from_cache(\n redis_client: Redis,\n db_session: Session,\n source: DocumentSource,\n) -> int | None:\n \"\"\"\n Get the SOURCE-type node ID for a given source from cache.\n\n If not in cache and db_session is provided, refreshes from DB.\n\n Returns:\n The ID of the SOURCE node, or None if not found.\n \"\"\"\n source_node_key = _source_node_key(source)\n\n # Try to get from dedicated SOURCE node key\n value = redis_client.get(source_node_key)\n if value is not None:\n if isinstance(value, bytes):\n value = value.decode(\"utf-8\")\n if not isinstance(value, str):\n raise ValueError(f\"SOURCE node value is not a string: {value}\")\n return int(value)\n\n # Not in cache - try refresh from DB\n refresh_hierarchy_cache_from_db(redis_client, db_session, source)\n\n # Try again after refresh\n value = redis_client.get(source_node_key)\n if value is not None:\n if isinstance(value, bytes):\n value = value.decode(\"utf-8\")\n if not isinstance(value, str):\n raise ValueError(f\"SOURCE node value is not a string: {value}\")\n return int(value)\n\n logger.error(f\"SOURCE node not found for source {source.value}\")\n return None\n\n\ndef clear_hierarchy_cache(redis_client: Redis, source: DocumentSource) -> None:\n \"\"\"Clear the hierarchy cache for a source (useful for testing).\"\"\"\n cache_key = _cache_key(source)\n raw_id_key = _raw_id_cache_key(source)\n source_node_key = _source_node_key(source)\n redis_client.delete(cache_key)\n redis_client.delete(raw_id_key)\n redis_client.delete(source_node_key)\n\n\ndef ensure_source_node_exists(\n redis_client: Redis,\n db_session: Session,\n source: DocumentSource,\n) -> int:\n \"\"\"\n Ensure that a SOURCE-type hierarchy node exists for the given source and cache it.\n\n This is the primary entry point for ensuring hierarchy infrastructure is set up\n for a source before processing documents. It should be called early in the\n indexing pipeline (e.g., at the start of docfetching or hierarchy fetching).\n\n The function:\n 1. Checks Redis cache for existing SOURCE node ID\n 2. If not cached, ensures the SOURCE node exists in the database\n 3. Caches the SOURCE node in Redis for fast subsequent lookups\n\n This is idempotent and safe to call multiple times concurrently.\n\n Args:\n redis_client: Redis client with tenant prefixing\n db_session: SQLAlchemy session for database operations\n source: The document source type (e.g., GOOGLE_DRIVE, CONFLUENCE)\n\n Returns:\n The database ID of the SOURCE-type hierarchy node\n \"\"\"\n # First check if we already have it cached\n source_node_key = _source_node_key(source)\n cached_value = redis_client.get(source_node_key)\n\n if cached_value is not None:\n value_str: str\n if isinstance(cached_value, bytes):\n value_str = cached_value.decode(\"utf-8\")\n else:\n value_str = str(cached_value)\n return int(value_str)\n\n # Not cached - ensure it exists in DB and cache it\n source_node = db_ensure_source_node_exists(db_session, source, commit=True)\n\n # Cache the SOURCE node\n cache_entry = HierarchyNodeCacheEntry.from_db_model(source_node)\n cache_hierarchy_node(redis_client, source, cache_entry)\n\n logger.info(\n f\"Ensured SOURCE node exists and cached for {source.value}: id={source_node.id}\"\n )\n\n return source_node.id\n" }, { "path": "backend/onyx/redis/redis_object_helper.py", "content": "from abc import ABC\nfrom abc import abstractmethod\n\nfrom celery import Celery\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.redis.redis_pool import get_redis_client\n\n\nclass RedisObjectHelper(ABC):\n PREFIX = \"base\"\n FENCE_PREFIX = PREFIX + \"_fence\"\n TASKSET_PREFIX = PREFIX + \"_taskset\"\n\n def __init__(self, tenant_id: str, id: str):\n self._tenant_id: str = tenant_id\n self._id: str = id\n self.redis = get_redis_client(tenant_id=tenant_id)\n\n @property\n def task_id_prefix(self) -> str:\n return f\"{self.PREFIX}_{self._id}\"\n\n @property\n def fence_key(self) -> str:\n # example: documentset_fence_1\n return f\"{self.FENCE_PREFIX}_{self._id}\"\n\n @property\n def taskset_key(self) -> str:\n # example: documentset_taskset_1\n return f\"{self.TASKSET_PREFIX}_{self._id}\"\n\n @staticmethod\n def get_id_from_fence_key(key: str) -> str | None:\n \"\"\"\n Extracts the object ID from a fence key in the format `PREFIX_fence_X`.\n\n Args:\n key (str): The fence key string.\n\n Returns:\n Optional[int]: The extracted ID if the key is in the correct format, otherwise None.\n \"\"\"\n parts = key.split(\"_\")\n if len(parts) != 3:\n return None\n\n object_id = parts[2]\n return object_id\n\n @staticmethod\n def get_id_from_task_id(task_id: str) -> str | None:\n \"\"\"\n Extracts the object ID from a task ID string.\n\n This method assumes the task ID is formatted as `prefix_objectid_suffix`, where:\n - `prefix` is an arbitrary string (e.g., the name of the task or entity),\n - `objectid` is the ID you want to extract,\n - `suffix` is another arbitrary string (e.g., a UUID).\n\n Example:\n If the input `task_id` is `documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc`,\n this method will return the string `\"1\"`.\n\n Args:\n task_id (str): The task ID string from which to extract the object ID.\n\n Returns:\n str | None: The extracted object ID if the task ID is in the correct format, otherwise None.\n \"\"\"\n # example: task_id=documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc\n parts = task_id.split(\"_\")\n if len(parts) != 3:\n return None\n\n object_id = parts[1]\n return object_id\n\n @abstractmethod\n def generate_tasks(\n self,\n max_tasks: int,\n celery_app: Celery,\n db_session: Session,\n redis_client: Redis,\n lock: RedisLock,\n tenant_id: str,\n ) -> tuple[int, int] | None:\n \"\"\"First element should be the number of actual tasks generated, second should\n be the number of docs that were candidates to be synced for the cc pair.\n\n The need for this is when we are syncing stale docs referenced by multiple\n connectors. In a single pass across multiple cc pairs, we only want a task\n for be created for a particular document id the first time we see it.\n The rest can be skipped.\"\"\"\n" }, { "path": "backend/onyx/redis/redis_pool.py", "content": "import asyncio\nimport functools\nimport json\nimport ssl\nimport threading\nfrom collections.abc import Callable\nfrom typing import Any\nfrom typing import cast\nfrom typing import Optional\n\nimport redis\nfrom fastapi import Request\nfrom redis import asyncio as aioredis\nfrom redis.client import Redis\nfrom redis.lock import Lock as RedisLock\n\nfrom onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX\nfrom onyx.configs.app_configs import REDIS_DB_NUMBER\nfrom onyx.configs.app_configs import REDIS_HEALTH_CHECK_INTERVAL\nfrom onyx.configs.app_configs import REDIS_HOST\nfrom onyx.configs.app_configs import REDIS_PASSWORD\nfrom onyx.configs.app_configs import REDIS_POOL_MAX_CONNECTIONS\nfrom onyx.configs.app_configs import REDIS_PORT\nfrom onyx.configs.app_configs import REDIS_REPLICA_HOST\nfrom onyx.configs.app_configs import REDIS_SSL\nfrom onyx.configs.app_configs import REDIS_SSL_CA_CERTS\nfrom onyx.configs.app_configs import REDIS_SSL_CERT_REQS\nfrom onyx.configs.app_configs import USE_REDIS_IAM_AUTH\nfrom onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME\nfrom onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS\nfrom onyx.redis.iam_auth import configure_redis_iam_auth\nfrom onyx.redis.iam_auth import create_redis_ssl_context_if_iam\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import DEFAULT_REDIS_PREFIX\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nSCAN_ITER_COUNT_DEFAULT = 4096\n\n\nclass TenantRedis(redis.Redis):\n def __init__(self, tenant_id: str, *args: Any, **kwargs: Any) -> None:\n super().__init__(*args, **kwargs)\n self.tenant_id: str = tenant_id\n\n def _prefixed(self, key: str | bytes | memoryview) -> str | bytes | memoryview:\n prefix: str = f\"{self.tenant_id}:\"\n if isinstance(key, str):\n if key.startswith(prefix):\n return key\n else:\n return prefix + key\n elif isinstance(key, bytes):\n prefix_bytes = prefix.encode()\n if key.startswith(prefix_bytes):\n return key\n else:\n return prefix_bytes + key\n elif isinstance(key, memoryview):\n key_bytes = key.tobytes()\n prefix_bytes = prefix.encode()\n if key_bytes.startswith(prefix_bytes):\n return key\n else:\n return memoryview(prefix_bytes + key_bytes)\n else:\n raise TypeError(f\"Unsupported key type: {type(key)}\")\n\n def _prefix_method(self, method: Callable) -> Callable:\n @functools.wraps(method)\n def wrapper(*args: Any, **kwargs: Any) -> Any:\n if \"name\" in kwargs:\n kwargs[\"name\"] = self._prefixed(kwargs[\"name\"])\n elif len(args) > 0:\n args = (self._prefixed(args[0]),) + args[1:]\n return method(*args, **kwargs)\n\n return wrapper\n\n def _prefix_scan_iter(self, method: Callable) -> Callable:\n @functools.wraps(method)\n def wrapper(*args: Any, **kwargs: Any) -> Any:\n # Prefix the match pattern if provided\n if \"match\" in kwargs:\n kwargs[\"match\"] = self._prefixed(kwargs[\"match\"])\n elif len(args) > 0:\n args = (self._prefixed(args[0]),) + args[1:]\n\n # Get the iterator\n iterator = method(*args, **kwargs)\n\n # Remove prefix from returned keys\n prefix = f\"{self.tenant_id}:\".encode()\n prefix_len = len(prefix)\n\n for key in iterator:\n if isinstance(key, bytes) and key.startswith(prefix):\n yield key[prefix_len:]\n else:\n yield key\n\n return wrapper\n\n def __getattribute__(self, item: str) -> Any:\n original_attr = super().__getattribute__(item)\n methods_to_wrap = [\n \"lock\",\n \"unlock\",\n \"get\",\n \"set\",\n \"setex\",\n \"delete\",\n \"exists\",\n \"incrby\",\n \"hset\",\n \"hget\",\n \"getset\",\n \"owned\",\n \"reacquire\",\n \"create_lock\",\n \"startswith\",\n \"smembers\",\n \"sismember\",\n \"sadd\",\n \"srem\",\n \"scard\",\n \"hexists\",\n \"hset\",\n \"hdel\",\n \"ttl\",\n \"pttl\",\n ] # Regular methods that need simple prefixing\n\n if item == \"scan_iter\" or item == \"sscan_iter\":\n return self._prefix_scan_iter(original_attr)\n elif item in methods_to_wrap and callable(original_attr):\n return self._prefix_method(original_attr)\n return original_attr\n\n\nclass RedisPool:\n _instance: Optional[\"RedisPool\"] = None\n _lock: threading.Lock = threading.Lock()\n _pool: redis.BlockingConnectionPool\n _replica_pool: redis.BlockingConnectionPool\n\n def __new__(cls) -> \"RedisPool\":\n if not cls._instance:\n with cls._lock:\n if not cls._instance:\n cls._instance = super(RedisPool, cls).__new__(cls)\n cls._instance._init_pools()\n return cls._instance\n\n def _init_pools(self) -> None:\n self._pool = RedisPool.create_pool(ssl=REDIS_SSL)\n self._replica_pool = RedisPool.create_pool(\n host=REDIS_REPLICA_HOST, ssl=REDIS_SSL\n )\n\n def get_client(self, tenant_id: str) -> Redis:\n return TenantRedis(tenant_id, connection_pool=self._pool)\n\n def get_replica_client(self, tenant_id: str) -> Redis:\n return TenantRedis(tenant_id, connection_pool=self._replica_pool)\n\n def get_raw_client(self) -> Redis:\n \"\"\"\n Returns a Redis client with direct access to the primary connection pool,\n without tenant prefixing.\n \"\"\"\n return redis.Redis(connection_pool=self._pool)\n\n def get_raw_replica_client(self) -> Redis:\n \"\"\"\n Returns a Redis client with direct access to the replica connection pool,\n without tenant prefixing.\n \"\"\"\n return redis.Redis(connection_pool=self._replica_pool)\n\n @staticmethod\n def create_pool(\n host: str = REDIS_HOST,\n port: int = REDIS_PORT,\n db: int = REDIS_DB_NUMBER,\n password: str = REDIS_PASSWORD,\n max_connections: int = REDIS_POOL_MAX_CONNECTIONS,\n ssl_ca_certs: str | None = REDIS_SSL_CA_CERTS,\n ssl_cert_reqs: str = REDIS_SSL_CERT_REQS,\n ssl: bool = False,\n ) -> redis.BlockingConnectionPool:\n \"\"\"\n Create a Redis connection pool with appropriate SSL configuration.\n SSL Configuration Priority:\n 1. IAM Authentication (USE_REDIS_IAM_AUTH=true): Uses system CA certificates\n 2. Regular SSL (REDIS_SSL=true): Uses custom SSL configuration\n 3. No SSL: Standard connection without encryption\n Note: IAM authentication automatically enables SSL and takes precedence\n over regular SSL configuration to ensure proper security.\n\n We use BlockingConnectionPool because it will block and wait for a connection\n rather than error if max_connections is reached. This is far more deterministic\n behavior and aligned with how we want to use Redis.\"\"\"\n\n # Using ConnectionPool is not well documented.\n # Useful examples: https://github.com/redis/redis-py/issues/780\n\n # Handle IAM authentication\n if USE_REDIS_IAM_AUTH:\n # For IAM authentication, we don't use password\n # and ensure SSL is enabled with proper context\n ssl_context = create_redis_ssl_context_if_iam()\n return redis.BlockingConnectionPool(\n host=host,\n port=port,\n db=db,\n password=None, # No password with IAM auth\n max_connections=max_connections,\n timeout=None,\n health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,\n socket_keepalive=True,\n socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,\n connection_class=redis.SSLConnection,\n ssl_context=ssl_context, # Use IAM auth SSL context\n )\n\n if ssl:\n return redis.BlockingConnectionPool(\n host=host,\n port=port,\n db=db,\n password=password,\n max_connections=max_connections,\n timeout=None,\n health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,\n socket_keepalive=True,\n socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,\n connection_class=redis.SSLConnection,\n ssl_ca_certs=ssl_ca_certs,\n ssl_cert_reqs=ssl_cert_reqs,\n )\n\n return redis.BlockingConnectionPool(\n host=host,\n port=port,\n db=db,\n password=password,\n max_connections=max_connections,\n timeout=None,\n health_check_interval=REDIS_HEALTH_CHECK_INTERVAL,\n socket_keepalive=True,\n socket_keepalive_options=REDIS_SOCKET_KEEPALIVE_OPTIONS,\n )\n\n\nredis_pool = RedisPool()\n\n\n# # Usage example\n# redis_pool = RedisPool()\n# redis_client = redis_pool.get_client()\n\n# # Example of setting and getting a value\n# redis_client.set('key', 'value')\n# value = redis_client.get('key')\n# print(value.decode()) # Output: 'value'\n\n\ndef get_redis_client(\n *,\n # This argument will be deprecated in the future\n tenant_id: str | None = None,\n) -> Redis:\n \"\"\"\n Returns a Redis client with tenant-specific key prefixing.\n\n This ensures proper data isolation between tenants by automatically\n prefixing all Redis keys with the tenant ID.\n\n Use this when working with tenant-specific data that should be\n isolated from other tenants.\n \"\"\"\n if tenant_id is None:\n tenant_id = get_current_tenant_id()\n\n return redis_pool.get_client(tenant_id)\n\n\ndef get_redis_replica_client(\n *,\n # this argument will be deprecated in the future\n tenant_id: str | None = None,\n) -> Redis:\n \"\"\"\n Returns a Redis replica client with tenant-specific key prefixing.\n\n Similar to get_redis_client(), but connects to a read replica when available.\n This ensures proper data isolation between tenants by automatically\n prefixing all Redis keys with the tenant ID.\n\n Use this for read-heavy operations on tenant-specific data.\n \"\"\"\n if tenant_id is None:\n tenant_id = get_current_tenant_id()\n\n return redis_pool.get_replica_client(tenant_id)\n\n\ndef get_shared_redis_client() -> Redis:\n \"\"\"\n Returns a Redis client with a shared namespace prefix.\n\n Unlike tenant-specific clients, this uses a common prefix for all keys,\n creating a shared namespace accessible across all tenants.\n\n Use this for data that should be shared across the application and\n isn't specific to any individual tenant.\n \"\"\"\n return redis_pool.get_client(DEFAULT_REDIS_PREFIX)\n\n\ndef get_shared_redis_replica_client() -> Redis:\n \"\"\"\n Returns a Redis replica client with a shared namespace prefix.\n\n Similar to get_shared_redis_client(), but connects to a read replica when available.\n Uses a common prefix for all keys, creating a shared namespace.\n\n Use this for read-heavy operations on data that should be shared\n across the application.\n \"\"\"\n return redis_pool.get_replica_client(DEFAULT_REDIS_PREFIX)\n\n\ndef get_raw_redis_client() -> Redis:\n \"\"\"\n Returns a Redis client that doesn't apply tenant prefixing to keys.\n\n Use this only when you need to access Redis directly without tenant isolation\n or any key prefixing. Typically needed for integrating with external systems\n or libraries that have inflexible key requirements.\n\n Warning: Be careful with this client as it bypasses tenant isolation.\n \"\"\"\n return redis_pool.get_raw_client()\n\n\ndef get_raw_redis_replica_client() -> Redis:\n \"\"\"\n Returns a Redis replica client that doesn't apply tenant prefixing to keys.\n\n Similar to get_raw_redis_client(), but connects to a read replica when available.\n Use this for read-heavy operations that need direct Redis access without\n tenant isolation or key prefixing.\n\n Warning: Be careful with this client as it bypasses tenant isolation.\n \"\"\"\n return redis_pool.get_raw_replica_client()\n\n\nSSL_CERT_REQS_MAP = {\n \"none\": ssl.CERT_NONE,\n \"optional\": ssl.CERT_OPTIONAL,\n \"required\": ssl.CERT_REQUIRED,\n}\n\n\n_async_redis_connection: aioredis.Redis | None = None\n_async_lock = asyncio.Lock()\n\n\nasync def get_async_redis_connection() -> aioredis.Redis:\n \"\"\"\n Provides a shared async Redis connection, using the same configs (host, port, SSL, etc.).\n Ensures that the connection is created only once (lazily) and reused for all future calls.\n \"\"\"\n global _async_redis_connection\n\n # If we haven't yet created an async Redis connection, we need to create one\n if _async_redis_connection is None:\n # Acquire the lock to ensure that only one coroutine attempts to create the connection\n async with _async_lock:\n # Double-check inside the lock to avoid race conditions\n if _async_redis_connection is None:\n # Load env vars or your config variables\n\n connection_kwargs: dict[str, Any] = {\n \"host\": REDIS_HOST,\n \"port\": REDIS_PORT,\n \"db\": REDIS_DB_NUMBER,\n \"password\": REDIS_PASSWORD,\n \"max_connections\": REDIS_POOL_MAX_CONNECTIONS,\n \"health_check_interval\": REDIS_HEALTH_CHECK_INTERVAL,\n \"socket_keepalive\": True,\n \"socket_keepalive_options\": REDIS_SOCKET_KEEPALIVE_OPTIONS,\n }\n\n if USE_REDIS_IAM_AUTH:\n configure_redis_iam_auth(connection_kwargs)\n elif REDIS_SSL:\n ssl_context = ssl.create_default_context()\n\n if REDIS_SSL_CA_CERTS:\n ssl_context.load_verify_locations(REDIS_SSL_CA_CERTS)\n ssl_context.check_hostname = False\n\n # Map your string to the proper ssl.CERT_* constant\n ssl_context.verify_mode = SSL_CERT_REQS_MAP.get(\n REDIS_SSL_CERT_REQS, ssl.CERT_NONE\n )\n\n connection_kwargs[\"ssl\"] = ssl_context\n\n # Create a new Redis connection (or connection pool) with SSL configuration\n _async_redis_connection = aioredis.Redis(**connection_kwargs)\n\n # Return the established connection (or pool) for all future operations\n return _async_redis_connection\n\n\nasync def retrieve_auth_token_data(token: str) -> dict | None:\n \"\"\"Validate auth token against Redis and return token data.\n\n Args:\n token: The raw authentication token string.\n\n Returns:\n Token data dict if valid, None if invalid/expired.\n \"\"\"\n try:\n redis = await get_async_redis_connection()\n redis_key = REDIS_AUTH_KEY_PREFIX + token\n token_data_str = await redis.get(redis_key)\n\n if not token_data_str:\n logger.debug(f\"Token key {redis_key} not found or expired in Redis\")\n return None\n\n return json.loads(token_data_str)\n except json.JSONDecodeError:\n logger.error(\"Error decoding token data from Redis\")\n return None\n except Exception as e:\n logger.error(f\"Unexpected error in retrieve_auth_token_data: {str(e)}\")\n raise ValueError(f\"Unexpected error in retrieve_auth_token_data: {str(e)}\")\n\n\nasync def retrieve_auth_token_data_from_redis(request: Request) -> dict | None:\n \"\"\"Validate auth token from request cookie. Wrapper for backwards compatibility.\"\"\"\n token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)\n if not token:\n logger.debug(\"No auth token cookie found\")\n return None\n return await retrieve_auth_token_data(token)\n\n\n# WebSocket token prefix (separate from regular auth tokens)\nREDIS_WS_TOKEN_PREFIX = \"ws_token:\"\n# WebSocket tokens expire after 60 seconds\nWS_TOKEN_TTL_SECONDS = 60\n# Rate limit: max tokens per user per window\nWS_TOKEN_RATE_LIMIT_MAX = 10\nWS_TOKEN_RATE_LIMIT_WINDOW_SECONDS = 60\nREDIS_WS_TOKEN_RATE_LIMIT_PREFIX = \"ws_token_rate:\"\n\n\nclass WsTokenRateLimitExceeded(Exception):\n \"\"\"Raised when a user exceeds the WS token generation rate limit.\"\"\"\n\n\nasync def store_ws_token(token: str, user_id: str) -> None:\n \"\"\"Store a short-lived WebSocket authentication token in Redis.\n\n Args:\n token: The generated WS token.\n user_id: The user ID to associate with this token.\n\n Raises:\n WsTokenRateLimitExceeded: If the user has exceeded the rate limit.\n \"\"\"\n redis = await get_async_redis_connection()\n\n # Atomically increment and check rate limit to avoid TOCTOU races\n rate_limit_key = REDIS_WS_TOKEN_RATE_LIMIT_PREFIX + user_id\n pipe = redis.pipeline()\n pipe.incr(rate_limit_key)\n pipe.expire(rate_limit_key, WS_TOKEN_RATE_LIMIT_WINDOW_SECONDS)\n results = await pipe.execute()\n new_count = results[0]\n\n if new_count > WS_TOKEN_RATE_LIMIT_MAX:\n # Over limit — decrement back since we won't use this slot\n await redis.decr(rate_limit_key)\n logger.warning(f\"WS token rate limit exceeded for user {user_id}\")\n raise WsTokenRateLimitExceeded(\n f\"Rate limit exceeded. Maximum {WS_TOKEN_RATE_LIMIT_MAX} tokens per minute.\"\n )\n\n # Store the actual token\n redis_key = REDIS_WS_TOKEN_PREFIX + token\n token_data = json.dumps({\"sub\": user_id})\n await redis.set(redis_key, token_data, ex=WS_TOKEN_TTL_SECONDS)\n\n\nasync def retrieve_ws_token_data(token: str) -> dict | None:\n \"\"\"Validate a WebSocket token and return the token data.\n\n This uses GETDEL for atomic get-and-delete to prevent race conditions\n where the same token could be used twice.\n\n Args:\n token: The WS token to validate.\n\n Returns:\n Token data dict with 'sub' (user ID) if valid, None if invalid/expired.\n \"\"\"\n try:\n redis = await get_async_redis_connection()\n redis_key = REDIS_WS_TOKEN_PREFIX + token\n\n # Atomic get-and-delete to prevent race conditions (Redis 6.2+)\n token_data_str = await redis.getdel(redis_key)\n\n if not token_data_str:\n return None\n\n return json.loads(token_data_str)\n except json.JSONDecodeError:\n logger.error(\"Error decoding WS token data from Redis\")\n return None\n except Exception as e:\n logger.error(f\"Unexpected error in retrieve_ws_token_data: {str(e)}\")\n return None\n\n\ndef redis_lock_dump(lock: RedisLock, r: Redis) -> None:\n # diagnostic logging for lock errors\n name = lock.name\n ttl = r.ttl(name)\n locked = lock.locked()\n owned = lock.owned()\n local_token: str | None = lock.local.token\n\n remote_token_raw = r.get(lock.name)\n if remote_token_raw:\n remote_token_bytes = cast(bytes, remote_token_raw)\n remote_token = remote_token_bytes.decode(\"utf-8\")\n else:\n remote_token = None\n\n logger.warning(\n f\"RedisLock diagnostic: \"\n f\"name={name} \"\n f\"locked={locked} \"\n f\"owned={owned} \"\n f\"local_token={local_token} \"\n f\"remote_token={remote_token} \"\n f\"ttl={ttl}\"\n )\n" }, { "path": "backend/onyx/redis/redis_usergroup.py", "content": "import time\nfrom typing import cast\nfrom uuid import uuid4\n\nimport redis\nfrom celery import Celery\nfrom redis import Redis\nfrom redis.lock import Lock as RedisLock\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DB_YIELD_PER_DEFAULT\nfrom onyx.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import OnyxRedisConstants\nfrom onyx.redis.redis_object_helper import RedisObjectHelper\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\n\nclass RedisUserGroup(RedisObjectHelper):\n PREFIX = \"usergroup\"\n FENCE_PREFIX = PREFIX + \"_fence\"\n FENCE_TTL = 7 * 24 * 60 * 60 # 7 days - defensive TTL to prevent memory leaks\n TASKSET_PREFIX = PREFIX + \"_taskset\"\n TASKSET_TTL = FENCE_TTL\n\n def __init__(self, tenant_id: str, id: int) -> None:\n super().__init__(tenant_id, str(id))\n\n @property\n def fenced(self) -> bool:\n if self.redis.exists(self.fence_key):\n return True\n\n return False\n\n def set_fence(self, payload: int | None) -> None:\n if payload is None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.fence_key)\n return\n\n self.redis.set(self.fence_key, payload, ex=self.FENCE_TTL)\n self.redis.sadd(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n\n @property\n def payload(self) -> int | None:\n bytes = self.redis.get(self.fence_key)\n if bytes is None:\n return None\n\n progress = int(cast(int, bytes))\n return progress\n\n def generate_tasks(\n self,\n max_tasks: int, # noqa: ARG002\n celery_app: Celery,\n db_session: Session,\n redis_client: Redis,\n lock: RedisLock,\n tenant_id: str,\n ) -> tuple[int, int] | None:\n \"\"\"Max tasks is ignored for now until we can build the logic to mark the\n user group up to date over multiple batches.\n \"\"\"\n last_lock_time = time.monotonic()\n num_tasks_sent = 0\n\n if not global_version.is_ee_version():\n return 0, 0\n\n try:\n construct_document_id_select_by_usergroup = fetch_versioned_implementation(\n \"onyx.db.user_group\",\n \"construct_document_id_select_by_usergroup\",\n )\n except ModuleNotFoundError:\n return 0, 0\n\n stmt = construct_document_id_select_by_usergroup(int(self._id))\n for doc_id in db_session.scalars(stmt).yield_per(DB_YIELD_PER_DEFAULT):\n doc_id = cast(str, doc_id)\n current_time = time.monotonic()\n if current_time - last_lock_time >= (\n CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4\n ):\n lock.reacquire()\n last_lock_time = current_time\n\n # celery's default task id format is \"dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # the key for the result is \"celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n # we prefix the task id so it's easier to keep track of who created the task\n # aka \"documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac\"\n custom_task_id = f\"{self.task_id_prefix}_{uuid4()}\"\n\n # add to the set BEFORE creating the task.\n redis_client.sadd(self.taskset_key, custom_task_id)\n redis_client.expire(self.taskset_key, self.TASKSET_TTL)\n\n celery_app.send_task(\n OnyxCeleryTask.VESPA_METADATA_SYNC_TASK,\n kwargs=dict(document_id=doc_id, tenant_id=tenant_id),\n queue=OnyxCeleryQueues.VESPA_METADATA_SYNC,\n task_id=custom_task_id,\n priority=OnyxCeleryPriority.MEDIUM,\n )\n\n num_tasks_sent += 1\n\n return num_tasks_sent, num_tasks_sent\n\n def reset(self) -> None:\n self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key)\n self.redis.delete(self.taskset_key)\n self.redis.delete(self.fence_key)\n\n @staticmethod\n def reset_all(r: redis.Redis) -> None:\n for key in r.scan_iter(RedisUserGroup.TASKSET_PREFIX + \"*\"):\n r.delete(key)\n\n for key in r.scan_iter(RedisUserGroup.FENCE_PREFIX + \"*\"):\n r.delete(key)\n" }, { "path": "backend/onyx/redis/redis_utils.py", "content": "from onyx.redis.redis_connector_delete import RedisConnectorDelete\nfrom onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync\nfrom onyx.redis.redis_connector_prune import RedisConnectorPrune\nfrom onyx.redis.redis_document_set import RedisDocumentSet\nfrom onyx.redis.redis_usergroup import RedisUserGroup\n\n\ndef is_fence(key_bytes: bytes) -> bool:\n key_str = key_bytes.decode(\"utf-8\")\n if key_str.startswith(RedisDocumentSet.FENCE_PREFIX):\n return True\n if key_str.startswith(RedisUserGroup.FENCE_PREFIX):\n return True\n if key_str.startswith(RedisConnectorDelete.FENCE_PREFIX):\n return True\n if key_str.startswith(RedisConnectorPrune.FENCE_PREFIX):\n return True\n if key_str.startswith(RedisConnectorPermissionSync.FENCE_PREFIX):\n return True\n\n return False\n" }, { "path": "backend/onyx/secondary_llm_flows/__init__.py", "content": "" }, { "path": "backend/onyx/secondary_llm_flows/chat_session_naming.py", "content": "from onyx.chat.llm_step import translate_history_to_llm_format\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.configs.constants import MessageType\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.utils import llm_response_to_string\nfrom onyx.prompts.chat_prompts import CHAT_NAMING_REMINDER\nfrom onyx.prompts.chat_prompts import CHAT_NAMING_SYSTEM_PROMPT\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef generate_chat_session_name(\n chat_history: list[ChatMessageSimple],\n llm: LLM,\n) -> str:\n system_prompt = ChatMessageSimple(\n message=CHAT_NAMING_SYSTEM_PROMPT,\n token_count=100,\n message_type=MessageType.SYSTEM,\n )\n\n reminder_prompt = ChatMessageSimple(\n message=CHAT_NAMING_REMINDER,\n token_count=100,\n message_type=MessageType.USER_REMINDER,\n )\n\n complete_message_history = [system_prompt] + chat_history + [reminder_prompt]\n\n llm_facing_history = translate_history_to_llm_format(\n complete_message_history, llm.config\n )\n\n # Call LLM with Braintrust tracing\n with llm_generation_span(\n llm=llm, flow=\"chat_session_naming\", input_messages=llm_facing_history\n ) as span_generation:\n response = llm.invoke(llm_facing_history, reasoning_effort=ReasoningEffort.OFF)\n record_llm_response(span_generation, response)\n new_name_raw = llm_response_to_string(response)\n\n return new_name_raw.strip().strip('\"')\n" }, { "path": "backend/onyx/secondary_llm_flows/document_filter.py", "content": "import json\nimport re\n\nfrom onyx.context.search.models import ContextExpansionType\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import InferenceSection\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import UserMessage\nfrom onyx.prompts.search_prompts import DOCUMENT_CONTEXT_SELECTION_PROMPT\nfrom onyx.prompts.search_prompts import DOCUMENT_SELECTION_PROMPT\nfrom onyx.prompts.search_prompts import TRY_TO_FILL_TO_MAX_INSTRUCTIONS\nfrom onyx.tools.tool_implementations.search.constants import (\n MAX_CHUNKS_FOR_RELEVANCE,\n)\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef select_chunks_for_relevance(\n section: InferenceSection,\n max_chunks: int = MAX_CHUNKS_FOR_RELEVANCE,\n) -> list[InferenceChunk]:\n \"\"\"Select a subset of chunks from a section based on center chunk position.\n\n Logic:\n - Always include the center chunk\n - If there are chunks directly next to it by index, grab the preceding and following\n - Otherwise grab 2 in the direction that does exist (2 before or 2 after)\n - If there are not enough in either direction, just grab what's available\n - If there are no other chunks, just use the central chunk\n\n Args:\n section: InferenceSection with center_chunk and chunks\n max_chunks: Maximum number of chunks to select (default: MAX_CHUNKS_FOR_RELEVANCE)\n\n Returns:\n List of selected InferenceChunks ordered by position\n \"\"\"\n if max_chunks <= 0:\n return []\n\n center_chunk = section.center_chunk\n all_chunks = section.chunks\n\n # Find the index of the center chunk in the chunks list\n try:\n center_index = next(\n i\n for i, chunk in enumerate(all_chunks)\n if chunk.chunk_id == center_chunk.chunk_id\n )\n except StopIteration:\n # If center chunk not found in chunks list, just return center chunk\n return [center_chunk]\n\n if max_chunks == 1:\n return [center_chunk]\n\n # Calculate how many chunks to take before and after\n chunks_needed = max_chunks - 1 # minus 1 for center chunk\n\n # Determine available chunks before and after center\n chunks_before_available = center_index\n chunks_after_available = len(all_chunks) - center_index - 1\n\n # Start with balanced distribution (1 before, 1 after for max_chunks=3)\n chunks_before = min(chunks_needed // 2, chunks_before_available)\n chunks_after = min(chunks_needed // 2, chunks_after_available)\n\n # Allocate remaining chunks to whichever direction has availability\n remaining = chunks_needed - chunks_before - chunks_after\n if remaining > 0:\n # Try to add more chunks before center if available\n if chunks_before_available > chunks_before:\n additional_before = min(remaining, chunks_before_available - chunks_before)\n chunks_before += additional_before\n remaining -= additional_before\n # Try to add more chunks after center if available\n if remaining > 0 and chunks_after_available > chunks_after:\n additional_after = min(remaining, chunks_after_available - chunks_after)\n chunks_after += additional_after\n\n # Select the chunks\n start_index = center_index - chunks_before\n end_index = center_index + chunks_after + 1 # +1 to include center and chunks after\n\n return all_chunks[start_index:end_index]\n\n\ndef classify_section_relevance(\n document_title: str,\n section_text: str,\n user_query: str,\n llm: LLM,\n section_above_text: str | None,\n section_below_text: str | None,\n) -> ContextExpansionType:\n \"\"\"Use LLM to classify section relevance and determine context expansion type.\n\n Args:\n section_text: The text content of the section to classify\n user_query: The user's search query\n llm: LLM instance to use for classification\n section_above_text: Text content from chunks above the section\n section_below_text: Text content from chunks below the section\n\n Returns:\n ContextExpansionType indicating how the section should be expanded\n \"\"\"\n # Build the prompt\n prompt_text = DOCUMENT_CONTEXT_SELECTION_PROMPT.format(\n document_title=document_title,\n main_section=section_text,\n section_above=section_above_text if section_above_text else \"N/A\",\n section_below=section_below_text if section_below_text else \"N/A\",\n user_query=user_query,\n )\n\n # Default to MAIN_SECTION_ONLY\n default_classification = ContextExpansionType.MAIN_SECTION_ONLY\n\n # Call LLM for classification with Braintrust tracing\n try:\n prompt_msg = UserMessage(content=prompt_text)\n with llm_generation_span(\n llm=llm, flow=\"classify_section_relevance\", input_messages=[prompt_msg]\n ) as span_generation:\n response = llm.invoke(\n prompt=prompt_msg,\n reasoning_effort=ReasoningEffort.OFF,\n )\n record_llm_response(span_generation, response)\n llm_response = response.choice.message.content\n\n if not llm_response:\n logger.warning(\n \"LLM returned empty response for context selection, defaulting to MAIN_SECTION_ONLY\"\n )\n classification = default_classification\n else:\n # Parse the response to extract the situation number (0-3)\n numbers = re.findall(r\"\\b[0-3]\\b\", llm_response)\n if numbers:\n situation = int(numbers[-1])\n # Map situation number to ContextExpansionType\n situation_to_type = {\n 0: ContextExpansionType.NOT_RELEVANT,\n 1: ContextExpansionType.MAIN_SECTION_ONLY,\n 2: ContextExpansionType.INCLUDE_ADJACENT_SECTIONS,\n 3: ContextExpansionType.FULL_DOCUMENT,\n }\n classification = situation_to_type.get(\n situation, default_classification\n )\n else:\n logger.warning(\n f\"Could not parse situation number from LLM response: {llm_response}\"\n )\n classification = default_classification\n\n except Exception as e:\n logger.error(f\"Error calling LLM for context selection: {e}\")\n classification = default_classification\n\n # To save some effort down the line, if there is nothing surrounding, don't allow a classification of adjacent or whole doc\n if (\n not section_above_text\n and not section_below_text\n and classification != ContextExpansionType.NOT_RELEVANT\n ):\n classification = ContextExpansionType.MAIN_SECTION_ONLY\n\n return classification\n\n\ndef select_sections_for_expansion(\n sections: list[InferenceSection],\n user_query: str,\n llm: LLM,\n max_sections: int = 10,\n max_chunks_per_section: int | None = MAX_CHUNKS_FOR_RELEVANCE,\n try_to_fill_to_max: bool = False,\n) -> tuple[list[InferenceSection], list[str] | None]:\n \"\"\"Use LLM to select the most relevant document sections for expansion.\n\n Args:\n sections: List of InferenceSection objects to select from\n user_query: The user's search query\n llm: LLM instance to use for selection\n max_sections: Maximum number of sections to select (default: 10)\n max_chunks_per_section: Maximum chunks to consider per section (default: MAX_CHUNKS_FOR_RELEVANCE)\n\n Returns:\n A tuple of:\n - Filtered list of InferenceSection objects selected by the LLM\n - List of document IDs for sections marked with \"!\" by the LLM, or None if none.\n Note: The \"!\" marker support exists in parsing but is not currently used because\n the prompt does not instruct the LLM to use it.\n \"\"\"\n if not sections:\n return [], None\n\n # Create a mapping of section ID to section\n section_map: dict[str, InferenceSection] = {}\n sections_dict: list[dict[str, str | int | list[str]]] = []\n\n for idx, section in enumerate(sections):\n # Create a unique ID for each section\n section_id = f\"{idx}\"\n section_map[section_id] = section\n\n # Format the section for the LLM\n chunk = section.center_chunk\n\n # Combine primary and secondary owners for authors\n authors = None\n if chunk.primary_owners or chunk.secondary_owners:\n authors = []\n if chunk.primary_owners:\n authors.extend(chunk.primary_owners)\n if chunk.secondary_owners:\n authors.extend(chunk.secondary_owners)\n\n # Format updated_at as ISO string if available\n updated_at_str = None\n if chunk.updated_at:\n updated_at_str = chunk.updated_at.isoformat()\n\n # Convert metadata to JSON string\n metadata_str = json.dumps(chunk.metadata)\n\n # Select only the most relevant chunks from the section to avoid flooding\n # the LLM with too much content from documents with many matching sections\n if max_chunks_per_section is not None:\n selected_chunks = select_chunks_for_relevance(\n section, max_chunks_per_section\n )\n selected_content = \" \".join(chunk.content for chunk in selected_chunks)\n else:\n selected_content = section.combined_content\n\n section_dict: dict[str, str | int | list[str]] = {\n \"section_id\": idx,\n \"title\": chunk.semantic_identifier,\n }\n\n # Only include updated_at if not None\n if updated_at_str is not None:\n section_dict[\"updated_at\"] = updated_at_str\n\n # Only include authors if not None\n if authors is not None:\n section_dict[\"authors\"] = authors\n\n section_dict[\"source_type\"] = str(chunk.source_type)\n section_dict[\"metadata\"] = metadata_str\n section_dict[\"content\"] = selected_content\n\n sections_dict.append(section_dict)\n\n # Build the prompt\n extra_instructions = TRY_TO_FILL_TO_MAX_INSTRUCTIONS if try_to_fill_to_max else \"\"\n prompt_text = UserMessage(\n content=DOCUMENT_SELECTION_PROMPT.format(\n max_sections=max_sections,\n extra_instructions=extra_instructions,\n formatted_doc_sections=json.dumps(sections_dict, indent=2),\n user_query=user_query,\n )\n )\n\n # Call LLM for selection with Braintrust tracing\n try:\n with llm_generation_span(\n llm=llm, flow=\"select_sections_for_expansion\", input_messages=[prompt_text]\n ) as span_generation:\n response = llm.invoke(\n prompt=[prompt_text], reasoning_effort=ReasoningEffort.OFF\n )\n record_llm_response(span_generation, response)\n llm_response = response.choice.message.content\n\n if not llm_response:\n logger.warning(\n \"LLM returned empty response for document selection, returning first max_sections\"\n )\n return sections[:max_sections], None\n\n # Parse the response to extract section IDs\n # Look for patterns like [1, 2, 3] or [1,2,3] with flexible whitespace/newlines\n # Also handle unbracketed comma-separated lists like \"1, 2, 3\"\n # Track which sections have \"!\" marker (e.g., \"1, 2!, 3\" or \"[1, 2!, 3]\")\n section_ids = []\n sections_with_exclamation = set() # Track section IDs that have \"!\" marker\n\n # First try to find a bracketed list\n bracket_pattern = r\"\\[([^\\]]+)\\]\"\n bracket_match = re.search(bracket_pattern, llm_response)\n\n if bracket_match:\n # Extract the content between brackets\n list_content = bracket_match.group(1)\n # Split by comma, preserving the parts\n parts = [part.strip() for part in list_content.split(\",\")]\n for part in parts:\n # Check if this part has an exclamation mark\n has_exclamation = \"!\" in part\n # Extract the number (digits only)\n numbers = re.findall(r\"\\d+\", part)\n if numbers:\n section_id = numbers[0]\n section_ids.append(section_id)\n if has_exclamation:\n sections_with_exclamation.add(section_id)\n else:\n # Try to find an unbracketed comma-separated list\n # Look for patterns like \"1, 2, 3\" or \"1, 2!, 3\"\n # This regex finds sequences of digits optionally followed by \"!\" and separated by commas\n comma_list_pattern = r\"\\b\\d+!?\\b(?:\\s*,\\s*\\b\\d+!?\\b)*\"\n comma_match = re.search(comma_list_pattern, llm_response)\n\n if comma_match:\n # Extract the matched comma-separated list\n list_content = comma_match.group(0)\n parts = [part.strip() for part in list_content.split(\",\")]\n for part in parts:\n # Check if this part has an exclamation mark\n has_exclamation = \"!\" in part\n # Extract the number (digits only)\n numbers = re.findall(r\"\\d+\", part)\n if numbers:\n section_id = numbers[0]\n section_ids.append(section_id)\n if has_exclamation:\n sections_with_exclamation.add(section_id)\n else:\n # Fallback: try to extract all numbers from the response\n # Also check for \"!\" after numbers\n number_pattern = r\"\\b(\\d+)(!)?\\b\"\n matches = re.finditer(number_pattern, llm_response)\n for match in matches:\n section_id = match.group(1)\n has_exclamation = match.group(2) == \"!\"\n section_ids.append(section_id)\n if has_exclamation:\n sections_with_exclamation.add(section_id)\n\n if not section_ids:\n logger.warning(\n f\"Could not parse section IDs from LLM response: {llm_response}\"\n )\n return sections[:max_sections], None\n\n # Filter sections based on LLM selection\n # Skip out-of-range IDs and don't count them toward max_sections\n selected_sections = []\n document_ids_with_exclamation = [] # Collect document_ids for sections with \"!\"\n num_sections = len(sections)\n\n for section_id_str in section_ids:\n # Convert to int\n try:\n section_id_int = int(section_id_str)\n except ValueError:\n logger.warning(f\"Could not convert section ID to int: {section_id_str}\")\n continue\n\n # Check if in valid range\n if section_id_int < 0 or section_id_int >= num_sections:\n logger.warning(\n f\"Section ID {section_id_int} is out of range [0, {num_sections - 1}], skipping\"\n )\n continue\n\n # Convert back to string for section_map lookup\n section_id = str(section_id_int)\n if section_id in section_map:\n section = section_map[section_id]\n selected_sections.append(section)\n\n # If this section has an exclamation mark, collect its document_id\n if section_id_str in sections_with_exclamation:\n document_id = section.center_chunk.document_id\n if document_id not in document_ids_with_exclamation:\n document_ids_with_exclamation.append(document_id)\n\n # Stop if we've reached max_sections valid selections\n if len(selected_sections) >= max_sections:\n break\n\n if not selected_sections:\n logger.warning(\n \"No valid sections selected from LLM response, returning first max_sections\"\n )\n return sections[:max_sections], None\n\n # Collect all selected document IDs\n selected_document_ids = [\n section.center_chunk.document_id for section in selected_sections\n ]\n\n logger.debug(\n f\"LLM selected {len(selected_sections)} valid sections from {len(sections)} total candidates. \"\n f\"Selected document IDs: {selected_document_ids}. \"\n f\"Document IDs with exclamation: {document_ids_with_exclamation if document_ids_with_exclamation else []}\"\n )\n\n # Return document_ids if any sections had exclamation marks, otherwise None\n return selected_sections, (\n document_ids_with_exclamation if document_ids_with_exclamation else None\n )\n\n except Exception as e:\n logger.error(f\"Error calling LLM for document selection: {e}\")\n return sections[:max_sections], None\n" }, { "path": "backend/onyx/secondary_llm_flows/memory_update.py", "content": "from onyx.configs.constants import MessageType\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import UserMessage\nfrom onyx.prompts.basic_memory import FULL_MEMORY_UPDATE_PROMPT\nfrom onyx.tools.models import ChatMinimalTextMessage\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.text_processing import parse_llm_json_response\n\nlogger = setup_logger()\n\n# Maximum number of user messages to include\nMAX_USER_MESSAGES = 3\nMAX_CHARS_PER_MESSAGE = 500\n\n\ndef _format_chat_history(chat_history: list[ChatMinimalTextMessage]) -> str:\n user_messages = [\n msg for msg in chat_history if msg.message_type == MessageType.USER\n ]\n\n if not user_messages:\n return \"No chat history available.\"\n\n # Take the last N user messages\n recent_user_messages = user_messages[-MAX_USER_MESSAGES:]\n\n formatted_parts = []\n for i, msg in enumerate(recent_user_messages, start=1):\n if len(msg.message) > MAX_CHARS_PER_MESSAGE:\n truncated_message = msg.message[:MAX_CHARS_PER_MESSAGE] + \"[...truncated]\"\n else:\n truncated_message = msg.message\n formatted_parts.append(f\"\\nUser message:\\n{truncated_message}\\n\")\n\n return \"\".join(formatted_parts).strip()\n\n\ndef _format_existing_memories(existing_memories: list[str]) -> str:\n \"\"\"Format existing memories as a numbered list (1-indexed for readability).\"\"\"\n if not existing_memories:\n return \"No existing memories.\"\n\n formatted_lines = []\n for i, memory in enumerate(existing_memories, start=1):\n formatted_lines.append(f\"{i}. {memory}\")\n\n return \"\\n\".join(formatted_lines)\n\n\ndef _format_user_basic_information(\n user_name: str | None,\n user_email: str | None,\n user_role: str | None,\n) -> str:\n \"\"\"Format user basic information, only including fields that have values.\"\"\"\n lines = []\n if user_name:\n lines.append(f\"User name: {user_name}\")\n if user_email:\n lines.append(f\"User email: {user_email}\")\n if user_role:\n lines.append(f\"User role: {user_role}\")\n\n if not lines:\n return \"\"\n\n return \"\\n\\n# User Basic Information\\n\" + \"\\n\".join(lines)\n\n\ndef process_memory_update(\n new_memory: str,\n existing_memories: list[str],\n chat_history: list[ChatMinimalTextMessage],\n llm: LLM,\n user_name: str | None = None,\n user_email: str | None = None,\n user_role: str | None = None,\n) -> tuple[str, int | None]:\n \"\"\"\n Determine if a memory should be added or updated.\n\n Uses the LLM to analyze the new memory against existing memories and\n determine whether to add it as new or update an existing memory.\n\n Args:\n new_memory: The new memory text from the memory tool\n existing_memories: List of existing memory strings\n chat_history: Recent chat history for context\n llm: LLM instance to use for the decision\n user_name: Optional user name for context\n user_email: Optional user email for context\n user_role: Optional user role for context\n\n Returns:\n Tuple of (memory_text, index_to_replace)\n - memory_text: The final memory text to store\n - index_to_replace: Index in existing_memories to replace, or None if adding new\n \"\"\"\n # Format inputs for the prompt\n formatted_chat_history = _format_chat_history(chat_history)\n formatted_memories = _format_existing_memories(existing_memories)\n formatted_user_info = _format_user_basic_information(\n user_name, user_email, user_role\n )\n\n # Build the prompt\n prompt = FULL_MEMORY_UPDATE_PROMPT.format(\n chat_history=formatted_chat_history,\n user_basic_information=formatted_user_info,\n existing_memories=formatted_memories,\n new_memory=new_memory,\n )\n\n # Call LLM with Braintrust tracing\n try:\n prompt_msg = UserMessage(content=prompt)\n with llm_generation_span(\n llm=llm, flow=\"memory_update\", input_messages=[prompt_msg]\n ) as span_generation:\n response = llm.invoke(\n prompt=prompt_msg, reasoning_effort=ReasoningEffort.OFF\n )\n record_llm_response(span_generation, response)\n content = response.choice.message.content\n except Exception as e:\n logger.warning(f\"LLM invocation failed for memory update: {e}\")\n return (new_memory, None)\n\n # Handle empty response\n if not content:\n logger.warning(\n \"LLM returned empty response for memory update, defaulting to add\"\n )\n return (new_memory, None)\n\n # Parse JSON response\n parsed_response = parse_llm_json_response(content)\n\n if not parsed_response:\n logger.warning(\n f\"Failed to parse JSON from LLM response: {content[:200]}..., defaulting to add\"\n )\n return (new_memory, None)\n\n # Extract fields from response\n operation = parsed_response.get(\"operation\", \"add\").lower()\n memory_id = parsed_response.get(\"memory_id\")\n memory_text = parsed_response.get(\"memory_text\", new_memory)\n\n # Ensure memory_text is valid\n if not memory_text or not isinstance(memory_text, str):\n memory_text = new_memory\n\n # Handle add operation\n if operation == \"add\":\n logger.debug(\"Memory update operation: add\")\n return (memory_text, None)\n\n # Handle update operation\n if operation == \"update\":\n # Validate memory_id\n if memory_id is None:\n logger.warning(\"Update operation specified but no memory_id provided\")\n return (memory_text, None)\n\n # Convert memory_id to integer if it's a string\n try:\n memory_id_int = int(memory_id)\n except (ValueError, TypeError):\n logger.warning(f\"Invalid memory_id format: {memory_id}\")\n return (memory_text, None)\n\n # Convert from 1-indexed (LLM response) to 0-indexed (internal)\n index_to_replace = memory_id_int - 1\n\n # Validate index is in range\n if index_to_replace < 0 or index_to_replace >= len(existing_memories):\n logger.warning(\n f\"memory_id {memory_id_int} out of range (1-{len(existing_memories)}), defaulting to add\"\n )\n return (memory_text, None)\n\n logger.debug(f\"Memory update operation: update at index {index_to_replace}\")\n return (memory_text, index_to_replace)\n\n # Unknown operation, default to add\n logger.warning(f\"Unknown operation '{operation}', defaulting to add\")\n return (memory_text, None)\n" }, { "path": "backend/onyx/secondary_llm_flows/query_expansion.py", "content": "from onyx.configs.constants import MessageType\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.models import AssistantMessage\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import ReasoningEffort\nfrom onyx.llm.models import SystemMessage\nfrom onyx.llm.models import UserMessage\nfrom onyx.prompts.prompt_utils import get_current_llm_day_time\nfrom onyx.prompts.search_prompts import KEYWORD_REPHRASE_SYSTEM_PROMPT\nfrom onyx.prompts.search_prompts import KEYWORD_REPHRASE_USER_PROMPT\nfrom onyx.prompts.search_prompts import REPHRASE_CONTEXT_PROMPT\nfrom onyx.prompts.search_prompts import SEMANTIC_QUERY_REPHRASE_SYSTEM_PROMPT\nfrom onyx.prompts.search_prompts import SEMANTIC_QUERY_REPHRASE_USER_PROMPT\nfrom onyx.tools.models import ChatMinimalTextMessage\nfrom onyx.tracing.llm_utils import llm_generation_span\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _build_additional_context(\n user_info: str | None = None,\n memories: list[str] | None = None,\n) -> str:\n \"\"\"Build additional context section for query rephrasing/expansion.\n\n Returns empty string if both user_info and memories are None/empty.\n Otherwise returns formatted context with \"N/A\" for missing fields.\n \"\"\"\n has_user_info = user_info and user_info.strip()\n has_memories = memories and any(m.strip() for m in memories)\n\n if not has_user_info and not has_memories:\n return \"\"\n\n formatted_user_info = user_info if has_user_info else \"N/A\"\n formatted_memories = (\n \"\\n\".join(f\"- {memory}\" for memory in memories)\n if has_memories and memories\n else \"N/A\"\n )\n\n return REPHRASE_CONTEXT_PROMPT.format(\n user_info=formatted_user_info,\n memories=formatted_memories,\n )\n\n\ndef _build_message_history(\n history: list[ChatMinimalTextMessage],\n) -> list[ChatCompletionMessage]:\n \"\"\"Convert ChatMinimalTextMessage list to ChatCompletionMessage list.\"\"\"\n messages: list[ChatCompletionMessage] = []\n\n for msg in history:\n if msg.message_type == MessageType.USER:\n user_msg = UserMessage(content=msg.message)\n messages.append(user_msg)\n elif msg.message_type == MessageType.ASSISTANT:\n assistant_msg = AssistantMessage(content=msg.message)\n messages.append(assistant_msg)\n\n return messages\n\n\ndef semantic_query_rephrase(\n history: list[ChatMinimalTextMessage],\n llm: LLM,\n user_info: str | None = None,\n memories: list[str] | None = None,\n) -> str:\n \"\"\"Rephrase a query into a standalone query using chat history context.\n\n Converts the user's query into a self-contained search query that incorporates\n relevant context from the chat history and optional user information/memories.\n\n Args:\n history: Chat message history. Must contain at least one user message.\n llm: Language model to use for rephrasing\n user_info: Optional user information for personalization\n memories: Optional user memories for personalization\n\n Returns:\n Rephrased standalone query string\n\n Raises:\n ValueError: If history is empty or contains no user messages\n RuntimeError: If LLM fails to generate a rephrased query\n \"\"\"\n if not history:\n raise ValueError(\"History cannot be empty for query rephrasing\")\n\n # Find the last user message in the history\n last_user_message_idx = None\n for i in range(len(history) - 1, -1, -1):\n if history[i].message_type == MessageType.USER:\n last_user_message_idx = i\n break\n\n if last_user_message_idx is None:\n raise ValueError(\"History must contain at least one user message\")\n\n # Extract the last user query\n user_query = history[last_user_message_idx].message\n\n # Build additional context section\n additional_context = _build_additional_context(user_info, memories)\n\n current_datetime_str = get_current_llm_day_time(\n include_day_of_week=True, full_sentence=False\n )\n\n # Build system message with current date\n system_msg = SystemMessage(\n content=SEMANTIC_QUERY_REPHRASE_SYSTEM_PROMPT.format(\n current_date=current_datetime_str\n )\n )\n\n # Convert chat history to message format (excluding the last user message and everything after it)\n messages: list[ChatCompletionMessage] = [system_msg]\n messages.extend(_build_message_history(history[:last_user_message_idx]))\n\n # Add the last message as the user prompt with instructions\n final_user_msg = UserMessage(\n content=SEMANTIC_QUERY_REPHRASE_USER_PROMPT.format(\n additional_context=additional_context, user_query=user_query\n )\n )\n messages.append(final_user_msg)\n\n # Call LLM and return result with Braintrust tracing\n with llm_generation_span(\n llm=llm, flow=\"semantic_query_rephrase\", input_messages=messages\n ) as span_generation:\n response = llm.invoke(prompt=messages, reasoning_effort=ReasoningEffort.OFF)\n record_llm_response(span_generation, response)\n final_query = response.choice.message.content\n\n if not final_query:\n # It's ok if some other queries fail, this one is likely the best one\n # It also can't fail in parsing so we should be able to guarantee a valid query here.\n raise RuntimeError(\"LLM failed to generate a rephrased query\")\n\n return final_query\n\n\ndef keyword_query_expansion(\n history: list[ChatMinimalTextMessage],\n llm: LLM,\n user_info: str | None = None,\n memories: list[str] | None = None,\n) -> list[str] | None:\n \"\"\"Expand a query into multiple keyword-only queries using chat history context.\n\n Converts the user's query into a set of keyword-based search queries (max 3)\n that incorporate relevant context from the chat history and optional user\n information/memories. Returns a list of keyword queries.\n\n Args:\n history: Chat message history. Must contain at least one user message.\n llm: Language model to use for keyword expansion\n user_info: Optional user information for personalization\n memories: Optional user memories for personalization\n\n Returns:\n List of keyword-only query strings (max 3), or empty list if generation fails\n\n Raises:\n ValueError: If history is empty or contains no user messages\n \"\"\"\n if not history:\n raise ValueError(\"History cannot be empty for keyword query expansion\")\n\n # Find the last user message in the history\n last_user_message_idx = None\n for i in range(len(history) - 1, -1, -1):\n if history[i].message_type == MessageType.USER:\n last_user_message_idx = i\n break\n\n if last_user_message_idx is None:\n raise ValueError(\"History must contain at least one user message\")\n\n # Extract the last user query\n user_query = history[last_user_message_idx].message\n\n # Build additional context section\n additional_context = _build_additional_context(user_info, memories)\n\n current_datetime_str = get_current_llm_day_time(\n include_day_of_week=True, full_sentence=False\n )\n\n # Build system message with current date\n system_msg = SystemMessage(\n content=KEYWORD_REPHRASE_SYSTEM_PROMPT.format(current_date=current_datetime_str)\n )\n\n # Convert chat history to message format (excluding the last user message and everything after it)\n messages: list[ChatCompletionMessage] = [system_msg]\n messages.extend(_build_message_history(history[:last_user_message_idx]))\n\n # Add the last message as the user prompt with instructions\n final_user_msg = UserMessage(\n content=KEYWORD_REPHRASE_USER_PROMPT.format(\n additional_context=additional_context, user_query=user_query\n )\n )\n messages.append(final_user_msg)\n\n # Call LLM and return result with Braintrust tracing\n with llm_generation_span(\n llm=llm, flow=\"keyword_query_expansion\", input_messages=messages\n ) as span_generation:\n response = llm.invoke(prompt=messages, reasoning_effort=ReasoningEffort.OFF)\n record_llm_response(span_generation, response)\n content = response.choice.message.content\n\n # Parse the response - each line is a separate keyword query\n if not content:\n return []\n\n queries = [line.strip() for line in content.strip().split(\"\\n\") if line.strip()]\n return queries\n" }, { "path": "backend/onyx/secondary_llm_flows/source_filter.py", "content": "from sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.llm.interfaces import LLM\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef strings_to_document_sources(source_strs: list[str]) -> list[DocumentSource]:\n sources = []\n for s in source_strs:\n try:\n sources.append(DocumentSource(s))\n except ValueError:\n logger.warning(f\"Failed to translate {s} to a DocumentSource\")\n return sources\n\n\ndef extract_source_filter(\n query: str, llm: LLM, db_session: Session\n) -> list[DocumentSource] | None:\n # Can reference onyx/prompts/filter_extration.py for previous implementation prompts\n raise NotImplementedError(\"This function should not be getting called right now\")\n" }, { "path": "backend/onyx/secondary_llm_flows/time_filter.py", "content": "from datetime import datetime\nfrom datetime import timezone\n\nfrom dateutil.parser import parse\n\nfrom onyx.llm.interfaces import LLM\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef best_match_time(time_str: str) -> datetime | None:\n preferred_formats = [\"%m/%d/%Y\", \"%m-%d-%Y\"]\n\n for fmt in preferred_formats:\n try:\n # As we don't know if the user is interacting with the API server from\n # the same timezone as the API server, just assume the queries are UTC time\n # the few hours offset (if any) shouldn't make any significant difference\n dt = datetime.strptime(time_str, fmt)\n return dt.replace(tzinfo=timezone.utc)\n except ValueError:\n continue\n\n # If the above formats don't match, try using dateutil's parser\n try:\n dt = parse(time_str)\n return (\n dt.astimezone(timezone.utc)\n if dt.tzinfo\n else dt.replace(tzinfo=timezone.utc)\n )\n except ValueError:\n return None\n\n\ndef extract_time_filter(query: str, llm: LLM) -> tuple[datetime | None, bool]:\n \"\"\"Returns a datetime if a hard time filter should be applied for the given query\n Additionally returns a bool, True if more recently updated Documents should be\n heavily favored\"\"\"\n raise NotImplementedError(\"This function should not be getting called right now\")\n\n\n# def _get_time_filter_messages(query: str) -> list[dict[str, str]]:\n# messages = [\n# {\n# \"role\": \"system\",\n# \"content\": TIME_FILTER_PROMPT.format(\n# current_day_time_str=get_current_llm_day_time()\n# ),\n# },\n# {\n# \"role\": \"user\",\n# \"content\": \"What documents in Confluence were written in the last two quarters\",\n# },\n# {\n# \"role\": \"assistant\",\n# \"content\": json.dumps(\n# {\n# \"filter_type\": \"hard cutoff\",\n# \"filter_value\": \"quarter\",\n# \"value_multiple\": 2,\n# }\n# ),\n# },\n# {\"role\": \"user\", \"content\": \"What's the latest on project Corgies?\"},\n# {\n# \"role\": \"assistant\",\n# \"content\": json.dumps({\"filter_type\": \"favor recent\"}),\n# },\n# {\n# \"role\": \"user\",\n# \"content\": \"Which customer asked about security features in February of 2022?\",\n# },\n# {\n# \"role\": \"assistant\",\n# \"content\": json.dumps(\n# {\"filter_type\": \"hard cutoff\", \"date\": \"02/01/2022\"}\n# ),\n# },\n# {\"role\": \"user\", \"content\": query},\n# ]\n# return messages\n\n# def _extract_time_filter_from_llm_out(\n# model_out: str,\n# ) -> tuple[datetime | None, bool]:\n# \"\"\"Returns a datetime for a hard cutoff and a bool for if the\"\"\"\n# try:\n# model_json = json.loads(model_out, strict=False)\n# except json.JSONDecodeError:\n# return None, False\n\n# # If filter type is not present, just assume something has gone wrong\n# # Potentially model has identified a date and just returned that but\n# # better to be conservative and not identify the wrong filter.\n# if \"filter_type\" not in model_json:\n# return None, False\n\n# if \"hard\" in model_json[\"filter_type\"] or \"recent\" in model_json[\"filter_type\"]:\n# favor_recent = \"recent\" in model_json[\"filter_type\"]\n\n# if \"date\" in model_json:\n# extracted_time = best_match_time(model_json[\"date\"])\n# if extracted_time is not None:\n# # LLM struggles to understand the concept of not sensitive within a time range\n# # So if a time is extracted, just go with that alone\n# return extracted_time, False\n\n# time_diff = None\n# multiplier = 1.0\n\n# if \"value_multiple\" in model_json:\n# try:\n# multiplier = float(model_json[\"value_multiple\"])\n# except ValueError:\n# pass\n\n# if \"filter_value\" in model_json:\n# filter_value = model_json[\"filter_value\"]\n# if \"day\" in filter_value:\n# time_diff = timedelta(days=multiplier)\n# elif \"week\" in filter_value:\n# time_diff = timedelta(weeks=multiplier)\n# elif \"month\" in filter_value:\n# # Have to just use the average here, too complicated to calculate exact day\n# # based on current day etc.\n# time_diff = timedelta(days=multiplier * 30.437)\n# elif \"quarter\" in filter_value:\n# time_diff = timedelta(days=multiplier * 91.25)\n# elif \"year\" in filter_value:\n# time_diff = timedelta(days=multiplier * 365)\n\n# if time_diff is not None:\n# current = datetime.now(timezone.utc)\n# # LLM struggles to understand the concept of not sensitive within a time range\n# # So if a time is extracted, just go with that alone\n# return current - time_diff, False\n\n# # If we failed to extract a hard filter, just pass back the value of favor recent\n# return None, favor_recent\n\n# return None, False\n\n# messages = _get_time_filter_messages(query)\n# filled_llm_prompt = dict_based_prompt_to_langchain_prompt(messages)\n# model_output = message_to_string(llm.invoke_langchain(filled_llm_prompt))\n# logger.debug(model_output)\n\n# return _extract_time_filter_from_llm_out(model_output)\n" }, { "path": "backend/onyx/seeding/__init__.py", "content": "" }, { "path": "backend/onyx/server/__init__.py", "content": "" }, { "path": "backend/onyx/server/api_key/api.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_admin_user\nfrom onyx.db.api_key import ApiKeyDescriptor\nfrom onyx.db.api_key import fetch_api_keys\nfrom onyx.db.api_key import insert_api_key\nfrom onyx.db.api_key import regenerate_api_key\nfrom onyx.db.api_key import remove_api_key\nfrom onyx.db.api_key import update_api_key\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.server.api_key.models import APIKeyArgs\n\n\nrouter = APIRouter(prefix=\"/admin/api-key\")\n\n\n@router.get(\"\")\ndef list_api_keys(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[ApiKeyDescriptor]:\n return fetch_api_keys(db_session)\n\n\n@router.post(\"\")\ndef create_api_key(\n api_key_args: APIKeyArgs,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> ApiKeyDescriptor:\n return insert_api_key(db_session, api_key_args, user.id)\n\n\n@router.post(\"/{api_key_id}/regenerate\")\ndef regenerate_existing_api_key(\n api_key_id: int,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> ApiKeyDescriptor:\n return regenerate_api_key(db_session, api_key_id)\n\n\n@router.patch(\"/{api_key_id}\")\ndef update_existing_api_key(\n api_key_id: int,\n api_key_args: APIKeyArgs,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> ApiKeyDescriptor:\n return update_api_key(db_session, api_key_id, api_key_args)\n\n\n@router.delete(\"/{api_key_id}\")\ndef delete_api_key(\n api_key_id: int,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> None:\n remove_api_key(db_session, api_key_id)\n" }, { "path": "backend/onyx/server/api_key/models.py", "content": "from pydantic import BaseModel\n\nfrom onyx.auth.schemas import UserRole\n\n\nclass APIKeyArgs(BaseModel):\n name: str | None = None\n role: UserRole = UserRole.BASIC\n" }, { "path": "backend/onyx/server/api_key_usage.py", "content": "\"\"\"API key and PAT usage tracking for cloud usage limits.\"\"\"\n\nfrom fastapi import Depends\nfrom fastapi import Request\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.api_key import get_hashed_api_key_from_request\nfrom onyx.auth.pat import get_hashed_pat_from_request\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.usage import increment_usage\nfrom onyx.db.usage import UsageType\nfrom onyx.server.usage_limits import check_usage_and_raise\nfrom onyx.server.usage_limits import is_usage_limits_enabled\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\ndef check_api_key_usage(\n request: Request,\n db_session: Session = Depends(get_session),\n) -> None:\n \"\"\"\n FastAPI dependency that checks and tracks API key/PAT usage limits.\n\n This should be added as a dependency to endpoints that accept API key\n or PAT authentication and should be usage-limited.\n \"\"\"\n if not is_usage_limits_enabled():\n return\n\n # Check if request is authenticated via API key or PAT\n is_api_key_request = get_hashed_api_key_from_request(request) is not None\n is_pat_request = get_hashed_pat_from_request(request) is not None\n\n if not is_api_key_request and not is_pat_request:\n return\n\n tenant_id = get_current_tenant_id()\n\n # Check usage limit\n check_usage_and_raise(\n db_session=db_session,\n usage_type=UsageType.API_CALLS,\n tenant_id=tenant_id,\n pending_amount=1,\n )\n\n # Increment usage counter\n increment_usage(\n db_session=db_session,\n usage_type=UsageType.API_CALLS,\n amount=1,\n )\n db_session.commit()\n" }, { "path": "backend/onyx/server/auth_check.py", "content": "from typing import cast\n\nfrom fastapi import FastAPI\nfrom fastapi.dependencies.models import Dependant\nfrom starlette.routing import BaseRoute\n\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_chat_accessible_user\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.auth.users import current_limited_user\nfrom onyx.auth.users import current_user\nfrom onyx.auth.users import current_user_from_websocket\nfrom onyx.auth.users import current_user_with_expired_token\nfrom onyx.configs.app_configs import APP_API_PREFIX\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\n\nPUBLIC_ENDPOINT_SPECS = [\n # built-in documentation functions\n (\"/openapi.json\", {\"GET\", \"HEAD\"}),\n (\"/docs\", {\"GET\", \"HEAD\"}),\n (\"/docs/oauth2-redirect\", {\"GET\", \"HEAD\"}),\n (\"/redoc\", {\"GET\", \"HEAD\"}),\n # should always be callable, will just return 401 if not authenticated\n (\"/me\", {\"GET\"}),\n # just returns 200 to validate that the server is up\n (\"/health\", {\"GET\"}),\n # just returns auth type, needs to be accessible before the user is logged\n # in to determine what flow to give the user\n (\"/auth/type\", {\"GET\"}),\n # just gets the version of Onyx (e.g. 0.3.11)\n (\"/version\", {\"GET\"}),\n # Gets stable and beta versions for Onyx docker images\n (\"/versions\", {\"GET\"}),\n # stuff related to basic auth\n (\"/auth/refresh\", {\"POST\"}),\n (\"/auth/register\", {\"POST\"}),\n (\"/auth/login\", {\"POST\"}),\n (\"/auth/logout\", {\"POST\"}),\n (\"/auth/forgot-password\", {\"POST\"}),\n (\"/auth/reset-password\", {\"POST\"}),\n (\"/auth/request-verify-token\", {\"POST\"}),\n (\"/auth/verify\", {\"POST\"}),\n (\"/users/me\", {\"GET\"}),\n (\"/users/me\", {\"PATCH\"}),\n (\"/users/{id}\", {\"GET\"}),\n (\"/users/{id}\", {\"PATCH\"}),\n (\"/users/{id}\", {\"DELETE\"}),\n # oauth\n (\"/auth/oauth/authorize\", {\"GET\"}),\n (\"/auth/oauth/callback\", {\"GET\"}),\n # oidc\n (\"/auth/oidc/authorize\", {\"GET\"}),\n (\"/auth/oidc/callback\", {\"GET\"}),\n # saml\n (\"/auth/saml/authorize\", {\"GET\"}),\n (\"/auth/saml/callback\", {\"POST\"}),\n (\"/auth/saml/callback\", {\"GET\"}),\n (\"/auth/saml/logout\", {\"POST\"}),\n # anonymous user on cloud\n (\"/tenants/anonymous-user\", {\"POST\"}),\n (\"/metrics\", {\"GET\"}), # added by prometheus_fastapi_instrumentator\n # craft webapp proxy — access enforced per-session via sharing_scope in handler\n (\"/build/sessions/{session_id}/webapp\", {\"GET\"}),\n (\"/build/sessions/{session_id}/webapp/{path:path}\", {\"GET\"}),\n]\n\n\ndef is_route_in_spec_list(\n route: BaseRoute, public_endpoint_specs: list[tuple[str, set[str]]]\n) -> bool:\n if not hasattr(route, \"path\") or not hasattr(route, \"methods\"):\n return False\n\n # try adding the prefix AND not adding the prefix, since some endpoints\n # are not prefixed (e.g. /openapi.json)\n if (route.path, route.methods) in public_endpoint_specs:\n return True\n\n processed_global_prefix = f\"/{APP_API_PREFIX.strip('/')}\" if APP_API_PREFIX else \"\"\n if not processed_global_prefix:\n return False\n\n for endpoint_spec in public_endpoint_specs:\n base_path, methods = endpoint_spec\n prefixed_path = f\"{processed_global_prefix}/{base_path.strip('/')}\"\n\n if prefixed_path == route.path and route.methods == methods:\n return True\n\n return False\n\n\ndef check_router_auth(\n application: FastAPI,\n public_endpoint_specs: list[tuple[str, set[str]]] = PUBLIC_ENDPOINT_SPECS,\n) -> None:\n \"\"\"Ensures that all endpoints on the passed in application either\n (1) have auth enabled OR\n (2) are explicitly marked as a public endpoint\n \"\"\"\n\n control_plane_dep = fetch_ee_implementation_or_noop(\n \"onyx.server.tenants.access\", \"control_plane_dep\"\n )\n current_cloud_superuser = fetch_ee_implementation_or_noop(\n \"onyx.auth.users\", \"current_cloud_superuser\"\n )\n verify_scim_token = fetch_ee_implementation_or_noop(\n \"onyx.server.scim.auth\", \"verify_scim_token\"\n )\n\n for route in application.routes:\n # explicitly marked as public\n if is_route_in_spec_list(route, public_endpoint_specs):\n continue\n\n # check for auth\n found_auth = False\n route_dependant_obj = cast(\n Dependant | None, route.dependant if hasattr(route, \"dependant\") else None\n )\n if route_dependant_obj:\n for dependency in route_dependant_obj.dependencies:\n depends_fn = dependency.cache_key[0]\n if (\n depends_fn == current_limited_user\n or depends_fn == current_user\n or depends_fn == current_admin_user\n or depends_fn == current_curator_or_admin_user\n or depends_fn == current_user_with_expired_token\n or depends_fn == current_chat_accessible_user\n or depends_fn == current_user_from_websocket\n or depends_fn == control_plane_dep\n or depends_fn == current_cloud_superuser\n or depends_fn == verify_scim_token\n ):\n found_auth = True\n break\n\n if not found_auth:\n # uncomment to print out all route(s) that are missing auth\n # print(f\"(\\\"{route.path}\\\", {set(route.methods)}),\")\n\n raise RuntimeError(\n f\"Did not find user dependency in private route - {route}\"\n )\n" }, { "path": "backend/onyx/server/documents/__init__.py", "content": "" }, { "path": "backend/onyx/server/documents/cc_pair.py", "content": "from datetime import datetime\nfrom http import HTTPStatus\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi.responses import JSONResponse\nfrom sqlalchemy import select\nfrom sqlalchemy.exc import IntegrityError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.auth.users import current_user\nfrom onyx.background.celery.tasks.pruning.tasks import (\n try_creating_prune_generator_task,\n)\nfrom onyx.background.celery.versioned_apps.client import app as client_app\nfrom onyx.background.indexing.models import IndexAttemptErrorPydantic\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.connectors.exceptions import ValidationError\nfrom onyx.connectors.factory import validate_ccpair_for_user\nfrom onyx.db.connector import delete_connector\nfrom onyx.db.connector_credential_pair import add_credential_to_connector\nfrom onyx.db.connector_credential_pair import (\n get_connector_credential_pair_from_id_for_user,\n)\nfrom onyx.db.connector_credential_pair import remove_credential_from_connector\nfrom onyx.db.connector_credential_pair import (\n update_connector_credential_pair_from_id,\n)\nfrom onyx.db.connector_credential_pair import verify_user_has_access_to_cc_pair\nfrom onyx.db.document import get_document_counts_for_cc_pairs\nfrom onyx.db.document import get_documents_for_cc_pair\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.index_attempt import count_index_attempt_errors_for_cc_pair\nfrom onyx.db.index_attempt import count_index_attempts_for_cc_pair\nfrom onyx.db.index_attempt import get_index_attempt_errors_for_cc_pair\nfrom onyx.db.index_attempt import get_latest_index_attempt_for_cc_pair_id\nfrom onyx.db.index_attempt import (\n get_latest_successful_index_attempt_for_cc_pair_id,\n)\nfrom onyx.db.index_attempt import get_paginated_index_attempts_for_cc_pair_id\nfrom onyx.db.indexing_coordination import IndexingCoordination\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import User\nfrom onyx.db.permission_sync_attempt import (\n get_latest_doc_permission_sync_attempt_for_cc_pair,\n)\nfrom onyx.db.permission_sync_attempt import (\n get_recent_doc_permission_sync_attempts_for_cc_pair,\n)\nfrom onyx.redis.redis_connector import RedisConnector\nfrom onyx.redis.redis_connector_utils import get_deletion_attempt_snapshot\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import CCPairFullInfo\nfrom onyx.server.documents.models import CCPropertyUpdateRequest\nfrom onyx.server.documents.models import CCStatusUpdateRequest\nfrom onyx.server.documents.models import ConnectorCredentialPairIdentifier\nfrom onyx.server.documents.models import ConnectorCredentialPairMetadata\nfrom onyx.server.documents.models import DocumentSyncStatus\nfrom onyx.server.documents.models import IndexAttemptSnapshot\nfrom onyx.server.documents.models import PaginatedReturn\nfrom onyx.server.documents.models import PermissionSyncAttemptSnapshot\nfrom onyx.server.models import StatusResponse\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\nrouter = APIRouter(prefix=\"/manage\")\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/index-attempts\", tags=PUBLIC_API_TAGS)\ndef get_cc_pair_index_attempts(\n cc_pair_id: int,\n page_num: int = Query(0, ge=0),\n page_size: int = Query(10, ge=1, le=1000),\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> PaginatedReturn[IndexAttemptSnapshot]:\n if user:\n user_has_access = verify_user_has_access_to_cc_pair(\n cc_pair_id, db_session, user, get_editable=False\n )\n if not user_has_access:\n raise HTTPException(\n status_code=400, detail=\"CC Pair not found for current user permissions\"\n )\n\n total_count = count_index_attempts_for_cc_pair(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n index_attempts = get_paginated_index_attempts_for_cc_pair_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n page=page_num,\n page_size=page_size,\n )\n return PaginatedReturn(\n items=[\n IndexAttemptSnapshot.from_index_attempt_db_model(index_attempt)\n for index_attempt in index_attempts\n ],\n total_items=total_count,\n )\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/permission-sync-attempts\")\ndef get_cc_pair_permission_sync_attempts(\n cc_pair_id: int,\n page_num: int = Query(0, ge=0),\n page_size: int = Query(10, ge=1, le=1000),\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> PaginatedReturn[PermissionSyncAttemptSnapshot]:\n if user:\n user_has_access = verify_user_has_access_to_cc_pair(\n cc_pair_id, db_session, user, get_editable=False\n )\n if not user_has_access:\n raise HTTPException(\n status_code=400, detail=\"CC Pair not found for current user permissions\"\n )\n\n # Get all permission sync attempts for this cc pair\n all_attempts = get_recent_doc_permission_sync_attempts_for_cc_pair(\n cc_pair_id=cc_pair_id,\n limit=1000,\n db_session=db_session,\n )\n\n start_idx = page_num * page_size\n end_idx = start_idx + page_size\n paginated_attempts = all_attempts[start_idx:end_idx]\n items = [\n PermissionSyncAttemptSnapshot.from_permission_sync_attempt_db_model(attempt)\n for attempt in paginated_attempts\n ]\n\n return PaginatedReturn(\n items=items,\n total_items=len(all_attempts),\n )\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}\", tags=PUBLIC_API_TAGS)\ndef get_cc_pair_full_info(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> CCPairFullInfo:\n tenant_id = get_current_tenant_id()\n\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id, db_session, user, get_editable=False\n )\n if not cc_pair:\n raise HTTPException(\n status_code=404, detail=\"CC Pair not found for current user permissions\"\n )\n editable_cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id, db_session, user, get_editable=True\n )\n is_editable_for_current_user = editable_cc_pair is not None\n\n document_count_info_list = list(\n get_document_counts_for_cc_pairs(\n db_session=db_session,\n cc_pairs=[\n ConnectorCredentialPairIdentifier(\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n )\n ],\n )\n )\n documents_indexed = (\n document_count_info_list[0][-1] if document_count_info_list else 0\n )\n\n latest_attempt = get_latest_index_attempt_for_cc_pair_id(\n db_session=db_session,\n connector_credential_pair_id=cc_pair_id,\n secondary_index=False,\n only_finished=False,\n )\n\n latest_successful_attempt = get_latest_successful_index_attempt_for_cc_pair_id(\n db_session=db_session,\n connector_credential_pair_id=cc_pair_id,\n )\n\n # Get latest permission sync attempt for status\n latest_permission_sync_attempt = None\n if cc_pair.access_type == AccessType.SYNC:\n latest_permission_sync_attempt = (\n get_latest_doc_permission_sync_attempt_for_cc_pair(\n db_session=db_session,\n connector_credential_pair_id=cc_pair_id,\n )\n )\n\n return CCPairFullInfo.from_models(\n cc_pair_model=cc_pair,\n number_of_index_attempts=count_index_attempts_for_cc_pair(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n ),\n last_index_attempt=latest_attempt,\n last_successful_index_time=(\n latest_successful_attempt.time_started\n if latest_successful_attempt\n else None\n ),\n latest_deletion_attempt=get_deletion_attempt_snapshot(\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n db_session=db_session,\n tenant_id=tenant_id,\n ),\n num_docs_indexed=documents_indexed,\n is_editable_for_current_user=is_editable_for_current_user,\n indexing=bool(\n latest_attempt and latest_attempt.status == IndexingStatus.IN_PROGRESS\n ),\n last_permission_sync_attempt_status=(\n latest_permission_sync_attempt.status\n if latest_permission_sync_attempt\n else None\n ),\n permission_syncing=bool(\n latest_permission_sync_attempt\n and latest_permission_sync_attempt.status\n == PermissionSyncStatus.IN_PROGRESS\n ),\n last_permission_sync_attempt_finished=(\n latest_permission_sync_attempt.time_finished\n if latest_permission_sync_attempt\n else None\n ),\n last_permission_sync_attempt_error_message=(\n latest_permission_sync_attempt.error_message\n if latest_permission_sync_attempt\n else None\n ),\n )\n\n\n@router.put(\"/admin/cc-pair/{cc_pair_id}/status\", tags=PUBLIC_API_TAGS)\ndef update_cc_pair_status(\n cc_pair_id: int,\n status_update_request: CCStatusUpdateRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> JSONResponse:\n \"\"\"This method returns nearly immediately. It simply sets some signals and\n optimistically assumes any running background processes will clean themselves up.\n This is done to improve the perceived end user experience.\n\n Returns HTTPStatus.OK if everything finished.\n \"\"\"\n tenant_id = get_current_tenant_id()\n\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=True,\n )\n\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"Connection not found for current user's permissions\",\n )\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n if status_update_request.status == ConnectorCredentialPairStatus.PAUSED:\n redis_connector.stop.set_fence(True)\n\n # Request cancellation for any active indexing attempts for this cc_pair\n active_attempts = (\n db_session.execute(\n select(IndexAttempt).where(\n IndexAttempt.connector_credential_pair_id == cc_pair_id,\n IndexAttempt.status.in_(\n [IndexingStatus.NOT_STARTED, IndexingStatus.IN_PROGRESS]\n ),\n )\n )\n .scalars()\n .all()\n )\n\n for attempt in active_attempts:\n try:\n IndexingCoordination.request_cancellation(db_session, attempt.id)\n # Revoke the task to prevent it from running\n if attempt.celery_task_id:\n client_app.control.revoke(attempt.celery_task_id)\n logger.info(\n f\"Requested cancellation for active indexing attempt {attempt.id} \"\n f\"due to connector pause: cc_pair={cc_pair_id}\"\n )\n except Exception:\n logger.exception(\n f\"Failed to request cancellation for indexing attempt {attempt.id}\"\n )\n\n else:\n redis_connector.stop.set_fence(False)\n\n update_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n status=status_update_request.status,\n )\n\n db_session.commit()\n\n # this speeds up the start of indexing by firing the check immediately\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_INDEXING,\n kwargs=dict(tenant_id=tenant_id),\n priority=OnyxCeleryPriority.HIGH,\n )\n\n return JSONResponse(\n status_code=HTTPStatus.OK, content={\"message\": str(HTTPStatus.OK)}\n )\n\n\n@router.put(\"/admin/cc-pair/{cc_pair_id}/name\")\ndef update_cc_pair_name(\n cc_pair_id: int,\n new_name: str,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[int]:\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=True,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400, detail=\"CC Pair not found for current user's permissions\"\n )\n\n try:\n cc_pair.name = new_name\n db_session.commit()\n return StatusResponse(\n success=True, message=\"Name updated successfully\", data=cc_pair_id\n )\n except IntegrityError:\n db_session.rollback()\n raise HTTPException(status_code=400, detail=\"Name must be unique\")\n\n\n@router.put(\"/admin/cc-pair/{cc_pair_id}/property\")\ndef update_cc_pair_property(\n cc_pair_id: int,\n update_request: CCPropertyUpdateRequest, # in seconds\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[int]:\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=True,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400, detail=\"CC Pair not found for current user's permissions\"\n )\n\n # Can we centralize logic for updating connector properties\n # so that we don't need to manually validate everywhere?\n if update_request.name == \"refresh_frequency\":\n cc_pair.connector.refresh_freq = int(update_request.value)\n cc_pair.connector.validate_refresh_freq()\n db_session.commit()\n\n msg = \"Refresh frequency updated successfully\"\n elif update_request.name == \"pruning_frequency\":\n cc_pair.connector.prune_freq = int(update_request.value)\n cc_pair.connector.validate_prune_freq()\n db_session.commit()\n\n msg = \"Pruning frequency updated successfully\"\n else:\n raise HTTPException(\n status_code=400, detail=f\"Property name {update_request.name} is not valid.\"\n )\n\n return StatusResponse(success=True, message=msg, data=cc_pair_id)\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/last_pruned\")\ndef get_cc_pair_last_pruned(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> datetime | None:\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"cc_pair not found for current user's permissions\",\n )\n\n return cc_pair.last_pruned\n\n\n@router.post(\"/admin/cc-pair/{cc_pair_id}/prune\", tags=PUBLIC_API_TAGS)\ndef prune_cc_pair(\n cc_pair_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[list[int]]:\n \"\"\"Triggers pruning on a particular cc_pair immediately\"\"\"\n tenant_id = get_current_tenant_id()\n\n cc_pair = get_connector_credential_pair_from_id_for_user(\n cc_pair_id=cc_pair_id,\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n if not cc_pair:\n raise HTTPException(\n status_code=400,\n detail=\"Connection not found for current user's permissions\",\n )\n\n r = get_redis_client()\n\n redis_connector = RedisConnector(tenant_id, cc_pair_id)\n if redis_connector.prune.fenced:\n raise HTTPException(\n status_code=HTTPStatus.CONFLICT,\n detail=\"Pruning task already in progress.\",\n )\n\n logger.info(\n f\"Pruning cc_pair: cc_pair={cc_pair_id} \"\n f\"connector={cc_pair.connector_id} \"\n f\"credential={cc_pair.credential_id} \"\n f\"{cc_pair.connector.name} connector.\"\n )\n payload_id = try_creating_prune_generator_task(\n client_app, cc_pair, db_session, r, tenant_id\n )\n if not payload_id:\n raise HTTPException(\n status_code=HTTPStatus.INTERNAL_SERVER_ERROR,\n detail=\"Pruning task creation failed.\",\n )\n\n logger.info(f\"Pruning queued: cc_pair={cc_pair.id} id={payload_id}\")\n\n return StatusResponse(\n success=True,\n message=\"Successfully created the pruning task.\",\n )\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/get-docs-sync-status\")\ndef get_docs_sync_status(\n cc_pair_id: int,\n _: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[DocumentSyncStatus]:\n all_docs_for_cc_pair = get_documents_for_cc_pair(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n )\n return [DocumentSyncStatus.from_model(doc) for doc in all_docs_for_cc_pair]\n\n\n@router.get(\"/admin/cc-pair/{cc_pair_id}/errors\", tags=PUBLIC_API_TAGS)\ndef get_cc_pair_indexing_errors(\n cc_pair_id: int,\n include_resolved: bool = Query(False),\n page_num: int = Query(0, ge=0),\n page_size: int = Query(10, ge=1, le=100),\n _: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> PaginatedReturn[IndexAttemptErrorPydantic]:\n \"\"\"Gives back all errors for a given CC Pair. Allows pagination based on page and page_size params.\n\n Args:\n cc_pair_id: ID of the connector-credential pair to get errors for\n include_resolved: Whether to include resolved errors in the results\n page_num: Page number for pagination, starting at 0\n page_size: Number of errors to return per page\n _: Current user, must be curator or admin\n db_session: Database session\n\n Returns:\n Paginated list of indexing errors for the CC pair.\n \"\"\"\n total_count = count_index_attempt_errors_for_cc_pair(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n unresolved_only=not include_resolved,\n )\n\n index_attempt_errors = get_index_attempt_errors_for_cc_pair(\n db_session=db_session,\n cc_pair_id=cc_pair_id,\n unresolved_only=not include_resolved,\n page=page_num,\n page_size=page_size,\n )\n return PaginatedReturn(\n items=[IndexAttemptErrorPydantic.from_model(e) for e in index_attempt_errors],\n total_items=total_count,\n )\n\n\n@router.put(\n \"/connector/{connector_id}/credential/{credential_id}\", tags=PUBLIC_API_TAGS\n)\ndef associate_credential_to_connector(\n connector_id: int,\n credential_id: int,\n metadata: ConnectorCredentialPairMetadata,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n tenant_id: str = Depends(get_current_tenant_id),\n) -> StatusResponse[int]:\n \"\"\"NOTE(rkuo): internally discussed and the consensus is this endpoint\n and create_connector_with_mock_credential should be combined.\n\n The intent of this endpoint is to handle connectors that actually need credentials.\n \"\"\"\n\n fetch_ee_implementation_or_noop(\n \"onyx.db.user_group\", \"validate_object_creation_for_user\", None\n )(\n db_session=db_session,\n user=user,\n target_group_ids=metadata.groups,\n object_is_public=metadata.access_type == AccessType.PUBLIC,\n object_is_perm_sync=metadata.access_type == AccessType.SYNC,\n object_is_new=True,\n )\n\n try:\n validate_ccpair_for_user(\n connector_id, credential_id, metadata.access_type, db_session\n )\n\n response = add_credential_to_connector(\n db_session=db_session,\n user=user,\n connector_id=connector_id,\n credential_id=credential_id,\n cc_pair_name=metadata.name,\n access_type=metadata.access_type,\n auto_sync_options=metadata.auto_sync_options,\n groups=metadata.groups,\n processing_mode=metadata.processing_mode,\n )\n\n # trigger indexing immediately\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_INDEXING,\n priority=OnyxCeleryPriority.HIGH,\n kwargs={\"tenant_id\": tenant_id},\n )\n\n logger.info(\n f\"associate_credential_to_connector - running check_for_indexing: cc_pair={response.data}\"\n )\n\n return response\n except ValidationError as e:\n # If validation fails, delete the connector and commit the changes\n # Ensures we don't leave invalid connectors in the database\n # NOTE: consensus is that it makes sense to unify connector and ccpair creation flows\n # which would rid us of needing to handle cases like these\n delete_connector(db_session, connector_id)\n db_session.commit()\n\n raise HTTPException(\n status_code=400, detail=\"Connector validation error: \" + str(e)\n )\n except IntegrityError as e:\n logger.error(f\"IntegrityError: {e}\")\n delete_connector(db_session, connector_id)\n db_session.commit()\n\n raise HTTPException(status_code=400, detail=\"Name must be unique\")\n\n except Exception as e:\n logger.exception(f\"Unexpected error: {e}\")\n\n raise HTTPException(status_code=500, detail=\"Unexpected error\")\n\n\n@router.delete(\n \"/connector/{connector_id}/credential/{credential_id}\", tags=PUBLIC_API_TAGS\n)\ndef dissociate_credential_from_connector(\n connector_id: int,\n credential_id: int,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[int]:\n return remove_credential_from_connector(\n connector_id, credential_id, user, db_session\n )\n" }, { "path": "backend/onyx/server/documents/connector.py", "content": "import json\nimport math\nimport mimetypes\nimport os\nimport zipfile\nfrom datetime import datetime\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import cast\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import File\nfrom fastapi import Form\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi import Request\nfrom fastapi import Response\nfrom fastapi import UploadFile\nfrom google.oauth2.credentials import Credentials\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.email_utils import send_email\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_chat_accessible_user\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.auth.users import current_user\nfrom onyx.background.celery.tasks.pruning.tasks import (\n try_creating_prune_generator_task,\n)\nfrom onyx.background.celery.versioned_apps.client import app as client_app\nfrom onyx.configs.app_configs import EMAIL_CONFIGURED\nfrom onyx.configs.app_configs import ENABLED_CONNECTOR_TYPES\nfrom onyx.configs.app_configs import MOCK_CONNECTOR_FILE_PATH\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.configs.constants import MilestoneRecordType\nfrom onyx.configs.constants import ONYX_METADATA_FILENAME\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.factory import validate_ccpair_for_user\nfrom onyx.connectors.google_utils.google_auth import (\n get_google_oauth_creds,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n build_service_account_creds,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n delete_google_app_cred,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n delete_service_account_key,\n)\nfrom onyx.connectors.google_utils.google_kv import get_auth_url\nfrom onyx.connectors.google_utils.google_kv import (\n get_google_app_cred,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n get_service_account_key,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n update_credential_access_tokens,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n upsert_google_app_cred,\n)\nfrom onyx.connectors.google_utils.google_kv import (\n upsert_service_account_key,\n)\nfrom onyx.connectors.google_utils.google_kv import verify_csrf\nfrom onyx.connectors.google_utils.shared_constants import DB_CREDENTIALS_DICT_TOKEN_KEY\nfrom onyx.connectors.google_utils.shared_constants import (\n GoogleOAuthAuthenticationMethod,\n)\nfrom onyx.db.connector import create_connector\nfrom onyx.db.connector import delete_connector\nfrom onyx.db.connector import fetch_connector_by_id\nfrom onyx.db.connector import fetch_connectors\nfrom onyx.db.connector import fetch_unique_document_sources\nfrom onyx.db.connector import get_connector_credential_ids\nfrom onyx.db.connector import mark_ccpair_with_indexing_trigger\nfrom onyx.db.connector import update_connector\nfrom onyx.db.connector_credential_pair import add_credential_to_connector\nfrom onyx.db.connector_credential_pair import (\n fetch_connector_credential_pair_for_connector,\n)\nfrom onyx.db.connector_credential_pair import get_cc_pair_groups_for_ids\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs_for_user\nfrom onyx.db.connector_credential_pair import (\n get_connector_credential_pairs_for_user_parallel,\n)\nfrom onyx.db.connector_credential_pair import verify_user_has_access_to_cc_pair\nfrom onyx.db.credentials import cleanup_gmail_credentials\nfrom onyx.db.credentials import cleanup_google_drive_credentials\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.credentials import delete_service_account_credentials\nfrom onyx.db.credentials import fetch_credential_by_id_for_user\nfrom onyx.db.deletion_attempt import check_deletion_attempt_is_allowed\nfrom onyx.db.document import get_document_counts_for_all_cc_pairs\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingMode\nfrom onyx.db.enums import ProcessingMode\nfrom onyx.db.federated import fetch_all_federated_connectors_parallel\nfrom onyx.db.index_attempt import get_index_attempts_for_cc_pair\nfrom onyx.db.index_attempt import get_latest_index_attempts_by_status\nfrom onyx.db.index_attempt import get_latest_index_attempts_parallel\nfrom onyx.db.index_attempt import (\n get_latest_successful_index_attempts_parallel,\n)\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import FederatedConnector\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexingStatus\nfrom onyx.db.models import User\nfrom onyx.db.models import UserRole\nfrom onyx.file_processing.file_types import PLAIN_TEXT_MIME_TYPE\nfrom onyx.file_processing.file_types import WORD_PROCESSING_MIME_TYPE\nfrom onyx.file_store.file_store import FileStore\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import AuthStatus\nfrom onyx.server.documents.models import AuthUrl\nfrom onyx.server.documents.models import ConnectorBase\nfrom onyx.server.documents.models import ConnectorCredentialPairIdentifier\nfrom onyx.server.documents.models import ConnectorFileInfo\nfrom onyx.server.documents.models import ConnectorFilesResponse\nfrom onyx.server.documents.models import ConnectorIndexingStatusLite\nfrom onyx.server.documents.models import ConnectorIndexingStatusLiteResponse\nfrom onyx.server.documents.models import ConnectorRequestSubmission\nfrom onyx.server.documents.models import ConnectorSnapshot\nfrom onyx.server.documents.models import ConnectorStatus\nfrom onyx.server.documents.models import ConnectorUpdateRequest\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.server.documents.models import CredentialSnapshot\nfrom onyx.server.documents.models import DocsCountOperator\nfrom onyx.server.documents.models import FailedConnectorIndexingStatus\nfrom onyx.server.documents.models import FileUploadResponse\nfrom onyx.server.documents.models import GDriveCallback\nfrom onyx.server.documents.models import GmailCallback\nfrom onyx.server.documents.models import GoogleAppCredentials\nfrom onyx.server.documents.models import GoogleServiceAccountCredentialRequest\nfrom onyx.server.documents.models import GoogleServiceAccountKey\nfrom onyx.server.documents.models import IndexedSourcesResponse\nfrom onyx.server.documents.models import IndexingStatusRequest\nfrom onyx.server.documents.models import ObjectCreationIdResponse\nfrom onyx.server.documents.models import RunConnectorRequest\nfrom onyx.server.documents.models import SourceSummary\nfrom onyx.server.federated.models import FederatedConnectorStatus\nfrom onyx.server.models import StatusResponse\nfrom onyx.server.utils_vector_db import require_vector_db\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.telemetry import mt_cloud_telemetry\nfrom onyx.utils.threadpool_concurrency import CallableProtocol\nfrom onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n_GMAIL_CREDENTIAL_ID_COOKIE_NAME = \"gmail_credential_id\"\n_GOOGLE_DRIVE_CREDENTIAL_ID_COOKIE_NAME = \"google_drive_credential_id\"\n_INDEXING_STATUS_PAGE_SIZE = 10\n\nSEEN_ZIP_DETAIL = \"Only one zip file is allowed per file connector, \\\nuse the ingestion APIs for multiple files\"\n\nrouter = APIRouter(prefix=\"/manage\", dependencies=[Depends(require_vector_db)])\n\n\n\"\"\"Admin only API endpoints\"\"\"\n\n\n@router.get(\"/admin/connector/gmail/app-credential\")\ndef check_google_app_gmail_credentials_exist(\n _: User = Depends(current_curator_or_admin_user),\n) -> dict[str, str]:\n try:\n return {\"client_id\": get_google_app_cred(DocumentSource.GMAIL).web.client_id}\n except KvKeyNotFoundError:\n raise HTTPException(status_code=404, detail=\"Google App Credentials not found\")\n\n\n@router.put(\"/admin/connector/gmail/app-credential\")\ndef upsert_google_app_gmail_credentials(\n app_credentials: GoogleAppCredentials, _: User = Depends(current_admin_user)\n) -> StatusResponse:\n try:\n upsert_google_app_cred(app_credentials, DocumentSource.GMAIL)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully saved Google App Credentials\"\n )\n\n\n@router.delete(\"/admin/connector/gmail/app-credential\")\ndef delete_google_app_gmail_credentials(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n try:\n delete_google_app_cred(DocumentSource.GMAIL)\n cleanup_gmail_credentials(db_session=db_session)\n except KvKeyNotFoundError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully deleted Google App Credentials\"\n )\n\n\n@router.get(\"/admin/connector/google-drive/app-credential\")\ndef check_google_app_credentials_exist(\n _: User = Depends(current_curator_or_admin_user),\n) -> dict[str, str]:\n try:\n return {\n \"client_id\": get_google_app_cred(DocumentSource.GOOGLE_DRIVE).web.client_id\n }\n except KvKeyNotFoundError:\n raise HTTPException(status_code=404, detail=\"Google App Credentials not found\")\n\n\n@router.put(\"/admin/connector/google-drive/app-credential\")\ndef upsert_google_app_credentials(\n app_credentials: GoogleAppCredentials, _: User = Depends(current_admin_user)\n) -> StatusResponse:\n try:\n upsert_google_app_cred(app_credentials, DocumentSource.GOOGLE_DRIVE)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully saved Google App Credentials\"\n )\n\n\n@router.delete(\"/admin/connector/google-drive/app-credential\")\ndef delete_google_app_credentials(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n try:\n delete_google_app_cred(DocumentSource.GOOGLE_DRIVE)\n cleanup_google_drive_credentials(db_session=db_session)\n except KvKeyNotFoundError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully deleted Google App Credentials\"\n )\n\n\n@router.get(\"/admin/connector/gmail/service-account-key\")\ndef check_google_service_gmail_account_key_exist(\n _: User = Depends(current_curator_or_admin_user),\n) -> dict[str, str]:\n try:\n return {\n \"service_account_email\": get_service_account_key(\n DocumentSource.GMAIL\n ).client_email\n }\n except KvKeyNotFoundError:\n raise HTTPException(\n status_code=404, detail=\"Google Service Account Key not found\"\n )\n\n\n@router.put(\"/admin/connector/gmail/service-account-key\")\ndef upsert_google_service_gmail_account_key(\n service_account_key: GoogleServiceAccountKey, _: User = Depends(current_admin_user)\n) -> StatusResponse:\n try:\n upsert_service_account_key(service_account_key, DocumentSource.GMAIL)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully saved Google Service Account Key\"\n )\n\n\n@router.delete(\"/admin/connector/gmail/service-account-key\")\ndef delete_google_service_gmail_account_key(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n try:\n delete_service_account_key(DocumentSource.GMAIL)\n cleanup_gmail_credentials(db_session=db_session)\n except KvKeyNotFoundError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully deleted Google Service Account Key\"\n )\n\n\n@router.get(\"/admin/connector/google-drive/service-account-key\")\ndef check_google_service_account_key_exist(\n _: User = Depends(current_curator_or_admin_user),\n) -> dict[str, str]:\n try:\n return {\n \"service_account_email\": get_service_account_key(\n DocumentSource.GOOGLE_DRIVE\n ).client_email\n }\n except KvKeyNotFoundError:\n raise HTTPException(\n status_code=404, detail=\"Google Service Account Key not found\"\n )\n\n\n@router.put(\"/admin/connector/google-drive/service-account-key\")\ndef upsert_google_service_account_key(\n service_account_key: GoogleServiceAccountKey, _: User = Depends(current_admin_user)\n) -> StatusResponse:\n try:\n upsert_service_account_key(service_account_key, DocumentSource.GOOGLE_DRIVE)\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully saved Google Service Account Key\"\n )\n\n\n@router.delete(\"/admin/connector/google-drive/service-account-key\")\ndef delete_google_service_account_key(\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n try:\n delete_service_account_key(DocumentSource.GOOGLE_DRIVE)\n cleanup_google_drive_credentials(db_session=db_session)\n except KvKeyNotFoundError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n return StatusResponse(\n success=True, message=\"Successfully deleted Google Service Account Key\"\n )\n\n\n@router.put(\"/admin/connector/google-drive/service-account-credential\")\ndef upsert_service_account_credential(\n service_account_credential_request: GoogleServiceAccountCredentialRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> ObjectCreationIdResponse:\n \"\"\"Special API which allows the creation of a credential for a service account.\n Combines the input with the saved service account key to create an entry in the\n `Credential` table.\"\"\"\n try:\n credential_base = build_service_account_creds(\n DocumentSource.GOOGLE_DRIVE,\n primary_admin_email=service_account_credential_request.google_primary_admin,\n name=\"Service Account (uploaded)\",\n )\n except KvKeyNotFoundError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n # first delete all existing service account credentials\n delete_service_account_credentials(user, db_session, DocumentSource.GOOGLE_DRIVE)\n # `user=None` since this credential is not a personal credential\n credential = create_credential(\n credential_data=credential_base, user=user, db_session=db_session\n )\n return ObjectCreationIdResponse(id=credential.id)\n\n\n@router.put(\"/admin/connector/gmail/service-account-credential\")\ndef upsert_gmail_service_account_credential(\n service_account_credential_request: GoogleServiceAccountCredentialRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> ObjectCreationIdResponse:\n \"\"\"Special API which allows the creation of a credential for a service account.\n Combines the input with the saved service account key to create an entry in the\n `Credential` table.\"\"\"\n try:\n credential_base = build_service_account_creds(\n DocumentSource.GMAIL,\n primary_admin_email=service_account_credential_request.google_primary_admin,\n )\n except KvKeyNotFoundError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n # first delete all existing service account credentials\n delete_service_account_credentials(user, db_session, DocumentSource.GMAIL)\n # `user=None` since this credential is not a personal credential\n credential = create_credential(\n credential_data=credential_base, user=user, db_session=db_session\n )\n return ObjectCreationIdResponse(id=credential.id)\n\n\n@router.get(\"/admin/connector/google-drive/check-auth/{credential_id}\")\ndef check_drive_tokens(\n credential_id: int,\n user: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> AuthStatus:\n db_credentials = fetch_credential_by_id_for_user(credential_id, user, db_session)\n if not db_credentials or not db_credentials.credential_json:\n return AuthStatus(authenticated=False)\n\n credential_json = db_credentials.credential_json.get_value(apply_mask=False)\n if DB_CREDENTIALS_DICT_TOKEN_KEY not in credential_json:\n return AuthStatus(authenticated=False)\n token_json_str = str(credential_json[DB_CREDENTIALS_DICT_TOKEN_KEY])\n google_drive_creds = get_google_oauth_creds(\n token_json_str=token_json_str,\n source=DocumentSource.GOOGLE_DRIVE,\n )\n if google_drive_creds is None:\n return AuthStatus(authenticated=False)\n return AuthStatus(authenticated=True)\n\n\ndef save_zip_metadata_to_file_store(\n zf: zipfile.ZipFile, file_store: FileStore\n) -> str | None:\n \"\"\"\n Extract .onyx_metadata.json from zip and save to file store.\n Returns the file_id or None if no metadata file exists.\n \"\"\"\n try:\n metadata_file_info = zf.getinfo(ONYX_METADATA_FILENAME)\n with zf.open(metadata_file_info, \"r\") as metadata_file:\n metadata_bytes = metadata_file.read()\n\n # Validate that it's valid JSON before saving\n try:\n json.loads(metadata_bytes)\n except json.JSONDecodeError as e:\n logger.warning(f\"Unable to load {ONYX_METADATA_FILENAME}: {e}\")\n raise HTTPException(\n status_code=400,\n detail=f\"Unable to load {ONYX_METADATA_FILENAME}: {e}\",\n )\n\n # Save to file store\n file_id = file_store.save_file(\n content=BytesIO(metadata_bytes),\n display_name=ONYX_METADATA_FILENAME,\n file_origin=FileOrigin.CONNECTOR_METADATA,\n file_type=\"application/json\",\n )\n return file_id\n except KeyError:\n logger.info(f\"No {ONYX_METADATA_FILENAME} file\")\n return None\n\n\ndef is_zip_file(file: UploadFile) -> bool:\n \"\"\"\n Check if the file is a zip file by content type or filename.\n \"\"\"\n return bool(\n (\n file.content_type\n and file.content_type.startswith(\n (\n \"application/zip\",\n \"application/x-zip-compressed\", # May be this in Windows\n \"application/x-zip\",\n \"multipart/x-zip\",\n )\n )\n )\n or (file.filename and file.filename.lower().endswith(\".zip\"))\n )\n\n\ndef upload_files(\n files: list[UploadFile],\n file_origin: FileOrigin = FileOrigin.CONNECTOR,\n unzip: bool = True,\n) -> FileUploadResponse:\n\n # Skip directories and known macOS metadata entries\n def should_process_file(file_path: str) -> bool:\n normalized_path = os.path.normpath(file_path)\n return not any(part.startswith(\".\") for part in normalized_path.split(os.sep))\n\n deduped_file_paths = []\n deduped_file_names = []\n zip_metadata_file_id: str | None = None\n try:\n file_store = get_default_file_store()\n seen_zip = False\n for file in files:\n if not file.filename:\n logger.warning(\"File has no filename, skipping\")\n continue\n\n if is_zip_file(file):\n if seen_zip:\n raise HTTPException(status_code=400, detail=SEEN_ZIP_DETAIL)\n seen_zip = True\n\n # Validate the zip by opening it (catches corrupt/non-zip files)\n with zipfile.ZipFile(file.file, \"r\") as zf:\n if unzip:\n zip_metadata_file_id = save_zip_metadata_to_file_store(\n zf, file_store\n )\n for file_info in zf.namelist():\n if zf.getinfo(file_info).is_dir():\n continue\n\n if not should_process_file(file_info):\n continue\n\n sub_file_bytes = zf.read(file_info)\n\n mime_type, __ = mimetypes.guess_type(file_info)\n if mime_type is None:\n mime_type = \"application/octet-stream\"\n\n file_id = file_store.save_file(\n content=BytesIO(sub_file_bytes),\n display_name=os.path.basename(file_info),\n file_origin=file_origin,\n file_type=mime_type,\n )\n deduped_file_paths.append(file_id)\n deduped_file_names.append(os.path.basename(file_info))\n continue\n\n # Store the zip as-is (unzip=False)\n file.file.seek(0)\n file_id = file_store.save_file(\n content=file.file,\n display_name=file.filename,\n file_origin=file_origin,\n file_type=file.content_type or \"application/zip\",\n )\n deduped_file_paths.append(file_id)\n deduped_file_names.append(file.filename)\n continue\n\n # Since we can't render docx files in the UI,\n # we store them in the file store as plain text\n if file.content_type == WORD_PROCESSING_MIME_TYPE:\n # Lazy load to avoid importing markitdown when not needed\n from onyx.file_processing.extract_file_text import read_docx_file\n\n text, _ = read_docx_file(file.file, file.filename)\n file_id = file_store.save_file(\n content=BytesIO(text.encode(\"utf-8\")),\n display_name=file.filename,\n file_origin=file_origin,\n file_type=PLAIN_TEXT_MIME_TYPE,\n )\n\n else:\n file_id = file_store.save_file(\n content=file.file,\n display_name=file.filename,\n file_origin=file_origin,\n file_type=file.content_type or \"text/plain\",\n )\n deduped_file_paths.append(file_id)\n deduped_file_names.append(file.filename)\n\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n return FileUploadResponse(\n file_paths=deduped_file_paths,\n file_names=deduped_file_names,\n zip_metadata_file_id=zip_metadata_file_id,\n )\n\n\ndef _normalize_file_names_for_backwards_compatibility(\n file_locations: list[str], file_names: list[str]\n) -> list[str]:\n \"\"\"\n Ensures file_names list is the same length as file_locations for backwards compatibility.\n In legacy data, file_names might not exist or be shorter than file_locations.\n If file_names is shorter, pads it with corresponding file_locations values.\n \"\"\"\n return file_names + file_locations[len(file_names) :]\n\n\ndef _fetch_and_check_file_connector_cc_pair_permissions(\n connector_id: int,\n user: User,\n db_session: Session,\n require_editable: bool,\n) -> ConnectorCredentialPair:\n cc_pair = fetch_connector_credential_pair_for_connector(db_session, connector_id)\n if cc_pair is None:\n raise HTTPException(\n status_code=404,\n detail=\"No Connector-Credential Pair found for this connector\",\n )\n\n has_requested_access = verify_user_has_access_to_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n user=user,\n get_editable=require_editable,\n )\n if has_requested_access:\n return cc_pair\n\n # Special case: global curators should be able to manage files\n # for public file connectors even when they are not the creator.\n if (\n require_editable\n and user.role == UserRole.GLOBAL_CURATOR\n and cc_pair.access_type == AccessType.PUBLIC\n ):\n return cc_pair\n\n raise HTTPException(\n status_code=403,\n detail=\"Access denied. User cannot manage files for this connector.\",\n )\n\n\n@router.post(\"/admin/connector/file/upload\", tags=PUBLIC_API_TAGS)\ndef upload_files_api(\n files: list[UploadFile],\n unzip: bool = True,\n _: User = Depends(current_curator_or_admin_user),\n) -> FileUploadResponse:\n return upload_files(files, FileOrigin.OTHER, unzip=unzip)\n\n\n@router.get(\"/admin/connector/{connector_id}/files\", tags=PUBLIC_API_TAGS)\ndef list_connector_files(\n connector_id: int,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> ConnectorFilesResponse:\n \"\"\"List all files in a file connector.\"\"\"\n connector = fetch_connector_by_id(connector_id, db_session)\n if connector is None:\n raise HTTPException(status_code=404, detail=\"Connector not found\")\n\n if connector.source != DocumentSource.FILE:\n raise HTTPException(\n status_code=400, detail=\"This endpoint only works with file connectors\"\n )\n\n _ = _fetch_and_check_file_connector_cc_pair_permissions(\n connector_id=connector_id,\n user=user,\n db_session=db_session,\n require_editable=False,\n )\n\n file_locations = connector.connector_specific_config.get(\"file_locations\", [])\n file_names = connector.connector_specific_config.get(\"file_names\", [])\n\n # Normalize file_names for backwards compatibility with legacy data\n file_names = _normalize_file_names_for_backwards_compatibility(\n file_locations, file_names\n )\n\n file_store = get_default_file_store()\n files = []\n\n for file_id, file_name in zip(file_locations, file_names):\n try:\n file_record = file_store.read_file_record(file_id)\n file_size = None\n upload_date = None\n if file_record:\n file_size = file_store.get_file_size(file_id)\n upload_date = (\n file_record.created_at.isoformat()\n if file_record.created_at\n else None\n )\n files.append(\n ConnectorFileInfo(\n file_id=file_id,\n file_name=file_name,\n file_size=file_size,\n upload_date=upload_date,\n )\n )\n except Exception as e:\n logger.warning(f\"Error reading file record for {file_id}: {e}\")\n # Include file with basic info even if record fetch fails\n files.append(\n ConnectorFileInfo(\n file_id=file_id,\n file_name=file_name,\n )\n )\n\n return ConnectorFilesResponse(files=files)\n\n\n@router.post(\"/admin/connector/{connector_id}/files/update\", tags=PUBLIC_API_TAGS)\ndef update_connector_files(\n connector_id: int,\n files: list[UploadFile] | None = File(None),\n file_ids_to_remove: str = Form(\"[]\"),\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> FileUploadResponse:\n \"\"\"\n Update files in a connector by adding new files and/or removing existing ones.\n This is an atomic operation that validates, updates the connector config, and triggers indexing.\n \"\"\"\n files = files or []\n connector = fetch_connector_by_id(connector_id, db_session)\n if connector is None:\n raise HTTPException(status_code=404, detail=\"Connector not found\")\n\n if connector.source != DocumentSource.FILE:\n raise HTTPException(\n status_code=400, detail=\"This endpoint only works with file connectors\"\n )\n\n # Get the connector-credential pair for indexing/pruning triggers\n # and validate user permissions for file management.\n cc_pair = _fetch_and_check_file_connector_cc_pair_permissions(\n connector_id=connector_id,\n user=user,\n db_session=db_session,\n require_editable=True,\n )\n\n # Parse file IDs to remove\n try:\n file_ids_list = json.loads(file_ids_to_remove)\n except json.JSONDecodeError:\n raise HTTPException(status_code=400, detail=\"Invalid file_ids_to_remove format\")\n\n if not isinstance(file_ids_list, list):\n raise HTTPException(\n status_code=400,\n detail=\"file_ids_to_remove must be a JSON-encoded list\",\n )\n\n # Get current connector config\n current_config = connector.connector_specific_config\n current_file_locations = current_config.get(\"file_locations\", [])\n current_file_names = current_config.get(\"file_names\", [])\n current_zip_metadata_file_id = current_config.get(\"zip_metadata_file_id\")\n\n # Load existing metadata from file store if available\n file_store = get_default_file_store()\n current_zip_metadata: dict[str, Any] = {}\n if current_zip_metadata_file_id:\n try:\n metadata_io = file_store.read_file(\n file_id=current_zip_metadata_file_id, mode=\"b\"\n )\n metadata_bytes = metadata_io.read()\n loaded_metadata = json.loads(metadata_bytes)\n if isinstance(loaded_metadata, list):\n current_zip_metadata = {d[\"filename\"]: d for d in loaded_metadata}\n else:\n current_zip_metadata = loaded_metadata\n except Exception as e:\n logger.warning(f\"Failed to load existing metadata file: {e}\")\n raise HTTPException(\n status_code=500,\n detail=\"Failed to load existing connector metadata file\",\n )\n\n # Upload new files if any\n new_file_paths = []\n new_file_names_list = []\n new_zip_metadata_file_id: str | None = None\n new_zip_metadata: dict[str, Any] = {}\n\n if files and len(files) > 0:\n upload_response = upload_files(files, FileOrigin.CONNECTOR)\n new_file_paths = upload_response.file_paths\n new_file_names_list = upload_response.file_names\n new_zip_metadata_file_id = upload_response.zip_metadata_file_id\n\n # Load new metadata from file store if available\n if new_zip_metadata_file_id:\n try:\n metadata_io = file_store.read_file(\n file_id=new_zip_metadata_file_id, mode=\"b\"\n )\n metadata_bytes = metadata_io.read()\n loaded_metadata = json.loads(metadata_bytes)\n if isinstance(loaded_metadata, list):\n new_zip_metadata = {d[\"filename\"]: d for d in loaded_metadata}\n else:\n new_zip_metadata = loaded_metadata\n except Exception as e:\n logger.warning(f\"Failed to load new metadata file: {e}\")\n\n # Remove specified files\n files_to_remove_set = set(file_ids_list)\n\n # Normalize file_names for backwards compatibility with legacy data\n current_file_names = _normalize_file_names_for_backwards_compatibility(\n current_file_locations, current_file_names\n )\n\n remaining_file_locations = []\n remaining_file_names = []\n removed_file_names = set()\n\n for file_id, file_name in zip(current_file_locations, current_file_names):\n if file_id not in files_to_remove_set:\n remaining_file_locations.append(file_id)\n remaining_file_names.append(file_name)\n else:\n removed_file_names.add(file_name)\n\n # Combine remaining files with new files\n final_file_locations = remaining_file_locations + new_file_paths\n final_file_names = remaining_file_names + new_file_names_list\n\n # Validate that at least one file remains\n if not final_file_locations:\n raise HTTPException(\n status_code=400,\n detail=\"Cannot remove all files from connector. At least one file must remain.\",\n )\n\n # Merge and filter metadata (remove metadata for deleted files)\n final_zip_metadata = {\n key: value\n for key, value in current_zip_metadata.items()\n if key not in removed_file_names\n }\n final_zip_metadata.update(new_zip_metadata)\n\n # Save merged metadata to file store if we have any metadata\n final_zip_metadata_file_id: str | None = None\n if final_zip_metadata:\n final_zip_metadata_file_id = file_store.save_file(\n content=BytesIO(json.dumps(final_zip_metadata).encode(\"utf-8\")),\n display_name=ONYX_METADATA_FILENAME,\n file_origin=FileOrigin.CONNECTOR_METADATA,\n file_type=\"application/json\",\n )\n\n # Update connector config\n updated_config = {\n **current_config,\n \"file_locations\": final_file_locations,\n \"file_names\": final_file_names,\n \"zip_metadata_file_id\": final_zip_metadata_file_id,\n }\n # Remove old zip_metadata dict if present (backwards compatibility cleanup)\n updated_config.pop(\"zip_metadata\", None)\n\n connector_base = ConnectorBase(\n name=connector.name,\n source=connector.source,\n input_type=connector.input_type,\n connector_specific_config=updated_config,\n refresh_freq=connector.refresh_freq,\n prune_freq=connector.prune_freq,\n indexing_start=connector.indexing_start,\n )\n\n updated_connector = update_connector(connector_id, connector_base, db_session)\n if updated_connector is None:\n raise HTTPException(\n status_code=500, detail=\"Failed to update connector configuration\"\n )\n\n # Trigger re-indexing for new files and pruning for removed files\n try:\n tenant_id = get_current_tenant_id()\n\n # If files were added, mark for UPDATE indexing (only new docs)\n if new_file_paths:\n mark_ccpair_with_indexing_trigger(\n cc_pair.id, IndexingMode.UPDATE, db_session\n )\n\n # Send task to check for indexing immediately\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_INDEXING,\n kwargs={\"tenant_id\": tenant_id},\n priority=OnyxCeleryPriority.HIGH,\n )\n logger.info(\n f\"Marked cc_pair {cc_pair.id} for UPDATE indexing (new files) for connector {connector_id}\"\n )\n\n # If files were removed, trigger pruning immediately\n if file_ids_list:\n r = get_redis_client()\n payload_id = try_creating_prune_generator_task(\n client_app, cc_pair, db_session, r, tenant_id\n )\n if payload_id:\n logger.info(\n f\"Triggered pruning for cc_pair {cc_pair.id} (removed files) for connector \"\n f\"{connector_id}, payload_id={payload_id}\"\n )\n else:\n logger.warning(\n f\"Failed to trigger pruning for cc_pair {cc_pair.id} (removed files) for connector {connector_id}\"\n )\n except Exception as e:\n logger.error(f\"Failed to trigger re-indexing after file update: {e}\")\n\n return FileUploadResponse(\n file_paths=final_file_locations,\n file_names=final_file_names,\n zip_metadata_file_id=final_zip_metadata_file_id,\n )\n\n\n@router.get(\"/admin/connector\", tags=PUBLIC_API_TAGS)\ndef get_connectors_by_credential(\n _: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n credential: int | None = None,\n) -> list[ConnectorSnapshot]:\n \"\"\"Get a list of connectors. Allow filtering by a specific credential id.\"\"\"\n\n connectors = fetch_connectors(db_session)\n\n filtered_connectors = []\n for connector in connectors:\n if connector.source == DocumentSource.INGESTION_API:\n # don't include INGESTION_API, as it's a system level\n # connector not manageable by the user\n continue\n\n if credential is not None:\n found = False\n for cc_pair in connector.credentials:\n if credential == cc_pair.credential_id:\n found = True\n break\n\n if not found:\n continue\n\n filtered_connectors.append(ConnectorSnapshot.from_connector_db_model(connector))\n\n return filtered_connectors\n\n\n# Retrieves most recent failure cases for connectors that are currently failing\n@router.get(\"/admin/connector/failed-indexing-status\", tags=PUBLIC_API_TAGS)\ndef get_currently_failed_indexing_status(\n secondary_index: bool = False,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n get_editable: bool = Query(\n False, description=\"If true, return editable document sets\"\n ),\n) -> list[FailedConnectorIndexingStatus]:\n # Get the latest failed indexing attempts\n latest_failed_indexing_attempts = get_latest_index_attempts_by_status(\n secondary_index=secondary_index,\n db_session=db_session,\n status=IndexingStatus.FAILED,\n )\n\n # Get the latest successful indexing attempts\n latest_successful_indexing_attempts = get_latest_index_attempts_by_status(\n secondary_index=secondary_index,\n db_session=db_session,\n status=IndexingStatus.SUCCESS,\n )\n\n # Get all connector credential pairs\n cc_pairs = get_connector_credential_pairs_for_user(\n db_session=db_session,\n user=user,\n get_editable=get_editable,\n )\n\n # Filter out failed attempts that have a more recent successful attempt\n filtered_failed_attempts = [\n failed_attempt\n for failed_attempt in latest_failed_indexing_attempts\n if not any(\n success_attempt.connector_credential_pair_id\n == failed_attempt.connector_credential_pair_id\n and success_attempt.time_updated > failed_attempt.time_updated\n for success_attempt in latest_successful_indexing_attempts\n )\n ]\n\n # Filter cc_pairs to include only those with failed attempts\n cc_pairs = [\n cc_pair\n for cc_pair in cc_pairs\n if any(\n attempt.connector_credential_pair == cc_pair\n for attempt in filtered_failed_attempts\n )\n ]\n\n # Create a mapping of cc_pair_id to its latest failed index attempt\n cc_pair_to_latest_index_attempt = {\n attempt.connector_credential_pair_id: attempt\n for attempt in filtered_failed_attempts\n }\n\n indexing_statuses = []\n\n for cc_pair in cc_pairs:\n # Skip DefaultCCPair\n if cc_pair.name == \"DefaultCCPair\":\n continue\n\n latest_index_attempt = cc_pair_to_latest_index_attempt.get(cc_pair.id)\n\n indexing_statuses.append(\n FailedConnectorIndexingStatus(\n cc_pair_id=cc_pair.id,\n name=cc_pair.name,\n error_msg=(\n latest_index_attempt.error_msg if latest_index_attempt else None\n ),\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n is_deletable=check_deletion_attempt_is_allowed(\n connector_credential_pair=cc_pair,\n db_session=db_session,\n allow_scheduled=True,\n )\n is None,\n )\n )\n\n return indexing_statuses\n\n\n@router.get(\"/admin/connector/status\", tags=PUBLIC_API_TAGS)\ndef get_connector_status(\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[ConnectorStatus]:\n # This method is only used document set and group creation/editing\n # Therefore, it is okay to get non-editable, but public cc_pairs\n cc_pairs = get_connector_credential_pairs_for_user(\n db_session=db_session,\n user=user,\n eager_load_connector=True,\n eager_load_credential=True,\n eager_load_user=True,\n get_editable=False,\n )\n\n group_cc_pair_relationships = get_cc_pair_groups_for_ids(\n db_session=db_session,\n cc_pair_ids=[cc_pair.id for cc_pair in cc_pairs],\n )\n group_cc_pair_relationships_dict: dict[int, list[int]] = {}\n for relationship in group_cc_pair_relationships:\n group_cc_pair_relationships_dict.setdefault(relationship.cc_pair_id, []).append(\n relationship.user_group_id\n )\n\n # Pre-compute credential_ids per connector to avoid N+1 lazy loads\n connector_to_credential_ids: dict[int, list[int]] = {}\n for cc_pair in cc_pairs:\n connector_to_credential_ids.setdefault(cc_pair.connector_id, []).append(\n cc_pair.credential_id\n )\n\n return [\n ConnectorStatus(\n cc_pair_id=cc_pair.id,\n name=cc_pair.name,\n connector=ConnectorSnapshot.from_connector_db_model(\n cc_pair.connector,\n credential_ids=connector_to_credential_ids.get(\n cc_pair.connector_id, []\n ),\n ),\n credential=CredentialSnapshot.from_credential_db_model(cc_pair.credential),\n access_type=cc_pair.access_type,\n groups=group_cc_pair_relationships_dict.get(cc_pair.id, []),\n )\n for cc_pair in cc_pairs\n if cc_pair.name != \"DefaultCCPair\" and cc_pair.connector and cc_pair.credential\n ]\n\n\n@router.post(\"/admin/connector/indexing-status\", tags=PUBLIC_API_TAGS)\ndef get_connector_indexing_status(\n request: IndexingStatusRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[ConnectorIndexingStatusLiteResponse]:\n tenant_id = get_current_tenant_id()\n\n # NOTE: If the connector is deleting behind the scenes,\n # accessing cc_pairs can be inconsistent and members like\n # connector or credential may be None.\n # Additional checks are done to make sure the connector and credential still exist.\n # TODO: make this one query ... possibly eager load or wrap in a read transaction\n # to avoid the complexity of trying to error check throughout the function\n\n # see https://stackoverflow.com/questions/75758327/\n # sqlalchemy-method-connection-for-bind-is-already-in-progress\n # for why we can't pass in the current db_session to these functions\n\n if MOCK_CONNECTOR_FILE_PATH:\n import json\n\n with open(MOCK_CONNECTOR_FILE_PATH, \"r\") as f:\n raw_data = json.load(f)\n connector_indexing_statuses = [\n ConnectorIndexingStatusLite(**status) for status in raw_data\n ]\n return [\n ConnectorIndexingStatusLiteResponse(\n source=DocumentSource.FILE,\n summary=SourceSummary(\n total_connectors=100,\n active_connectors=100,\n public_connectors=100,\n total_docs_indexed=100000,\n ),\n current_page=1,\n total_pages=1,\n indexing_statuses=connector_indexing_statuses,\n )\n ]\n\n parallel_functions: list[tuple[CallableProtocol, tuple[Any, ...]]] = [\n # Get editable connector/credential pairs\n (\n lambda: get_connector_credential_pairs_for_user_parallel(\n user, True, None, True, True, False, True, request.source\n ),\n (),\n ),\n # Get federated connectors\n (fetch_all_federated_connectors_parallel, ()),\n # Get most recent index attempts\n (\n lambda: get_latest_index_attempts_parallel(\n request.secondary_index, True, False\n ),\n (),\n ),\n # Get most recent finished index attempts\n (\n lambda: get_latest_index_attempts_parallel(\n request.secondary_index, True, True\n ),\n (),\n ),\n # Get most recent successful index attempts\n (\n lambda: get_latest_successful_index_attempts_parallel(\n request.secondary_index,\n ),\n (),\n ),\n ]\n\n if user and user.role == UserRole.ADMIN:\n (\n editable_cc_pairs,\n federated_connectors,\n latest_index_attempts,\n latest_finished_index_attempts,\n latest_successful_index_attempts,\n ) = run_functions_tuples_in_parallel(parallel_functions)\n non_editable_cc_pairs = []\n else:\n parallel_functions.append(\n (\n lambda: get_connector_credential_pairs_for_user_parallel(\n user, False, None, True, True, False, True, request.source\n ),\n (),\n ),\n )\n\n (\n editable_cc_pairs,\n federated_connectors,\n latest_index_attempts,\n latest_finished_index_attempts,\n latest_successful_index_attempts,\n non_editable_cc_pairs,\n ) = run_functions_tuples_in_parallel(parallel_functions)\n\n # Cast results to proper types\n non_editable_cc_pairs = cast(list[ConnectorCredentialPair], non_editable_cc_pairs)\n editable_cc_pairs = cast(list[ConnectorCredentialPair], editable_cc_pairs)\n federated_connectors = cast(list[FederatedConnector], federated_connectors)\n latest_index_attempts = cast(list[IndexAttempt], latest_index_attempts)\n latest_finished_index_attempts = cast(\n list[IndexAttempt], latest_finished_index_attempts\n )\n latest_successful_index_attempts = cast(\n list[IndexAttempt], latest_successful_index_attempts\n )\n\n document_count_info = get_document_counts_for_all_cc_pairs(db_session)\n\n # Create lookup dictionaries for efficient access\n cc_pair_to_document_cnt: dict[tuple[int, int], int] = {\n (connector_id, credential_id): cnt\n for connector_id, credential_id, cnt in document_count_info\n }\n\n def _attempt_lookup(\n attempts: list[IndexAttempt],\n ) -> dict[int, IndexAttempt]:\n return {attempt.connector_credential_pair_id: attempt for attempt in attempts}\n\n cc_pair_to_latest_index_attempt = _attempt_lookup(latest_index_attempts)\n cc_pair_to_latest_finished_index_attempt = _attempt_lookup(\n latest_finished_index_attempts\n )\n cc_pair_to_latest_successful_index_attempt = _attempt_lookup(\n latest_successful_index_attempts\n )\n\n def build_connector_indexing_status(\n cc_pair: ConnectorCredentialPair,\n is_editable: bool,\n ) -> ConnectorIndexingStatusLite | None:\n if cc_pair.name == \"DefaultCCPair\":\n return None\n\n latest_attempt = cc_pair_to_latest_index_attempt.get(cc_pair.id)\n latest_finished_attempt = cc_pair_to_latest_finished_index_attempt.get(\n cc_pair.id\n )\n latest_successful_attempt = cc_pair_to_latest_successful_index_attempt.get(\n cc_pair.id\n )\n doc_count = cc_pair_to_document_cnt.get(\n (cc_pair.connector_id, cc_pair.credential_id), 0\n )\n\n return _get_connector_indexing_status_lite(\n cc_pair,\n latest_attempt,\n latest_finished_attempt,\n (\n latest_successful_attempt.time_started\n if latest_successful_attempt\n else None\n ),\n is_editable,\n doc_count,\n )\n\n # Process editable cc_pairs\n editable_statuses: list[ConnectorIndexingStatusLite] = []\n for cc_pair in editable_cc_pairs:\n status = build_connector_indexing_status(cc_pair, True)\n if status:\n editable_statuses.append(status)\n\n # Process non-editable cc_pairs\n non_editable_statuses: list[ConnectorIndexingStatusLite] = []\n for cc_pair in non_editable_cc_pairs:\n status = build_connector_indexing_status(cc_pair, False)\n if status:\n non_editable_statuses.append(status)\n\n # Process federated connectors\n federated_statuses: list[FederatedConnectorStatus] = []\n for federated_connector in federated_connectors:\n federated_status = FederatedConnectorStatus(\n id=federated_connector.id,\n source=federated_connector.source,\n name=f\"{federated_connector.source.replace('_', ' ').title()}\",\n )\n\n federated_statuses.append(federated_status)\n\n source_to_summary: dict[DocumentSource, SourceSummary] = {}\n\n # Apply filters only if any are provided\n has_filters = bool(\n request.access_type_filters\n or request.last_status_filters\n or (\n request.docs_count_operator is not None\n and request.docs_count_value is not None\n )\n or request.name_filter\n )\n\n if has_filters:\n editable_statuses = _apply_connector_status_filters(\n editable_statuses,\n request.access_type_filters,\n request.last_status_filters,\n request.docs_count_operator,\n request.docs_count_value,\n request.name_filter,\n )\n non_editable_statuses = _apply_connector_status_filters(\n non_editable_statuses,\n request.access_type_filters,\n request.last_status_filters,\n request.docs_count_operator,\n request.docs_count_value,\n request.name_filter,\n )\n federated_statuses = _apply_federated_connector_status_filters(\n federated_statuses,\n request.name_filter,\n )\n\n # Calculate source summary\n for connector_status in (\n editable_statuses + non_editable_statuses + federated_statuses\n ):\n if isinstance(connector_status, FederatedConnectorStatus):\n source = connector_status.source.to_non_federated_source()\n else:\n source = connector_status.source\n\n # Skip if source is None (federated connectors without mapping)\n if source is None:\n continue\n\n if source not in source_to_summary:\n source_to_summary[source] = SourceSummary(\n total_connectors=0,\n active_connectors=0,\n public_connectors=0,\n total_docs_indexed=0,\n )\n source_to_summary[source].total_connectors += 1\n if isinstance(connector_status, ConnectorIndexingStatusLite):\n if connector_status.cc_pair_status == ConnectorCredentialPairStatus.ACTIVE:\n source_to_summary[source].active_connectors += 1\n if connector_status.access_type == AccessType.PUBLIC:\n source_to_summary[source].public_connectors += 1\n source_to_summary[\n source\n ].total_docs_indexed += connector_status.docs_indexed\n\n # Track admin page visit for analytics\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id),\n event=MilestoneRecordType.VISITED_ADMIN_PAGE,\n )\n\n # Group statuses by source for pagination\n source_to_all_statuses: dict[\n DocumentSource, list[ConnectorIndexingStatusLite | FederatedConnectorStatus]\n ] = {}\n # Group by source\n for connector_status in (\n editable_statuses + non_editable_statuses + federated_statuses\n ):\n if isinstance(connector_status, FederatedConnectorStatus):\n source = connector_status.source.to_non_federated_source()\n else:\n source = connector_status.source\n\n # Skip if source is None (federated connectors without mapping)\n if source is None:\n continue\n\n if source not in source_to_all_statuses:\n source_to_all_statuses[source] = []\n source_to_all_statuses[source].append(connector_status)\n\n # Create paginated response objects by source\n response_list: list[ConnectorIndexingStatusLiteResponse] = []\n\n source_list = list(source_to_all_statuses.keys())\n source_list.sort()\n\n for source in source_list:\n statuses = source_to_all_statuses[source]\n # Get current page for this source (default to page 1, 1-indexed)\n current_page = request.source_to_page.get(source, 1)\n\n # Calculate start and end indices for pagination (convert to 0-indexed)\n start_idx = (current_page - 1) * _INDEXING_STATUS_PAGE_SIZE\n end_idx = start_idx + _INDEXING_STATUS_PAGE_SIZE\n\n if request.get_all_connectors:\n page_statuses = statuses\n else:\n # Get the page slice for this source\n page_statuses = statuses[start_idx:end_idx]\n\n # Create response object for this source\n if page_statuses: # Only include sources that have data on this page\n response_list.append(\n ConnectorIndexingStatusLiteResponse(\n source=source,\n summary=source_to_summary[source],\n current_page=current_page,\n total_pages=math.ceil(len(statuses) / _INDEXING_STATUS_PAGE_SIZE),\n indexing_statuses=page_statuses,\n )\n )\n\n return response_list\n\n\ndef _get_connector_indexing_status_lite(\n cc_pair: ConnectorCredentialPair,\n latest_index_attempt: IndexAttempt | None,\n latest_finished_index_attempt: IndexAttempt | None,\n last_successful_index_time: datetime | None,\n is_editable: bool,\n document_cnt: int,\n) -> ConnectorIndexingStatusLite | None:\n # TODO remove this to enable ingestion API\n if cc_pair.name == \"DefaultCCPair\":\n return None\n\n connector = cc_pair.connector\n credential = cc_pair.credential\n if not connector or not credential:\n # This may happen if background deletion is happening\n return None\n\n in_progress = bool(\n latest_index_attempt\n and latest_index_attempt.status == IndexingStatus.IN_PROGRESS\n )\n\n return ConnectorIndexingStatusLite(\n cc_pair_id=cc_pair.id,\n name=cc_pair.name,\n source=cc_pair.connector.source,\n access_type=cc_pair.access_type,\n cc_pair_status=cc_pair.status,\n is_editable=is_editable,\n in_progress=in_progress,\n in_repeated_error_state=cc_pair.in_repeated_error_state,\n last_finished_status=(\n latest_finished_index_attempt.status\n if latest_finished_index_attempt\n else None\n ),\n last_status=latest_index_attempt.status if latest_index_attempt else None,\n last_success=last_successful_index_time,\n docs_indexed=document_cnt,\n latest_index_attempt_docs_indexed=(\n latest_index_attempt.total_docs_indexed if latest_index_attempt else None\n ),\n )\n\n\ndef _apply_connector_status_filters(\n statuses: list[ConnectorIndexingStatusLite],\n access_type_filters: list[AccessType],\n last_status_filters: list[IndexingStatus],\n docs_count_operator: DocsCountOperator | None,\n docs_count_value: int | None,\n name_filter: str | None,\n) -> list[ConnectorIndexingStatusLite]:\n \"\"\"Apply filters to a list of ConnectorIndexingStatusLite objects\"\"\"\n filtered_statuses: list[ConnectorIndexingStatusLite] = []\n\n for status in statuses:\n # Filter by access type\n if access_type_filters and status.access_type not in access_type_filters:\n continue\n\n # Filter by last status\n if last_status_filters and status.last_status not in last_status_filters:\n continue\n\n # Filter by document count\n if docs_count_operator and docs_count_value is not None:\n if docs_count_operator == DocsCountOperator.GREATER_THAN and not (\n status.docs_indexed > docs_count_value\n ):\n continue\n elif docs_count_operator == DocsCountOperator.LESS_THAN and not (\n status.docs_indexed < docs_count_value\n ):\n continue\n elif (\n docs_count_operator == DocsCountOperator.EQUAL_TO\n and status.docs_indexed != docs_count_value\n ):\n continue\n\n # Filter by name\n if status.name:\n if name_filter and name_filter.lower() not in status.name.lower():\n continue\n else:\n if name_filter:\n continue\n\n filtered_statuses.append(status)\n\n return filtered_statuses\n\n\ndef _apply_federated_connector_status_filters(\n statuses: list[FederatedConnectorStatus],\n name_filter: str | None,\n) -> list[FederatedConnectorStatus]:\n filtered_statuses: list[FederatedConnectorStatus] = []\n\n for status in statuses:\n if name_filter and name_filter.lower() not in status.name.lower():\n continue\n\n filtered_statuses.append(status)\n\n return filtered_statuses\n\n\ndef _validate_connector_allowed(source: DocumentSource) -> None:\n valid_connectors = [\n x for x in ENABLED_CONNECTOR_TYPES.replace(\"_\", \"\").split(\",\") if x\n ]\n if not valid_connectors:\n return\n for connector_type in valid_connectors:\n if source.value.lower().replace(\"_\", \"\") == connector_type:\n return\n\n raise ValueError(\n \"This connector type has been disabled by your system admin. Please contact them to get it enabled if you wish to use it.\"\n )\n\n\n@router.post(\"/admin/connector\", tags=PUBLIC_API_TAGS)\ndef create_connector_from_model(\n connector_data: ConnectorUpdateRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> ObjectCreationIdResponse:\n tenant_id = get_current_tenant_id()\n\n try:\n _validate_connector_allowed(connector_data.source)\n\n fetch_ee_implementation_or_noop(\n \"onyx.db.user_group\", \"validate_object_creation_for_user\", None\n )(\n db_session=db_session,\n user=user,\n target_group_ids=connector_data.groups,\n object_is_public=connector_data.access_type == AccessType.PUBLIC,\n object_is_perm_sync=connector_data.access_type == AccessType.SYNC,\n object_is_new=True,\n )\n connector_base = connector_data.to_connector_base()\n connector_response = create_connector(\n db_session=db_session,\n connector_data=connector_base,\n )\n\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id),\n event=MilestoneRecordType.CREATED_CONNECTOR,\n )\n\n return connector_response\n except ValueError as e:\n logger.error(f\"Error creating connector: {e}\")\n raise HTTPException(status_code=400, detail=str(e))\n\n\n@router.post(\"/admin/connector-with-mock-credential\")\ndef create_connector_with_mock_credential(\n connector_data: ConnectorUpdateRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n tenant_id = get_current_tenant_id()\n\n fetch_ee_implementation_or_noop(\n \"onyx.db.user_group\", \"validate_object_creation_for_user\", None\n )(\n db_session=db_session,\n user=user,\n target_group_ids=connector_data.groups,\n object_is_public=connector_data.access_type == AccessType.PUBLIC,\n object_is_perm_sync=connector_data.access_type == AccessType.SYNC,\n )\n try:\n _validate_connector_allowed(connector_data.source)\n connector_response = create_connector(\n db_session=db_session,\n connector_data=connector_data,\n )\n\n mock_credential = CredentialBase(\n credential_json={},\n admin_public=True,\n source=connector_data.source,\n )\n credential = create_credential(\n credential_data=mock_credential,\n user=user,\n db_session=db_session,\n )\n\n # Store the created connector and credential IDs\n connector_id = cast(int, connector_response.id)\n credential_id = credential.id\n\n validate_ccpair_for_user(\n connector_id=connector_id,\n credential_id=credential_id,\n access_type=connector_data.access_type,\n db_session=db_session,\n )\n response = add_credential_to_connector(\n db_session=db_session,\n user=user,\n connector_id=connector_id,\n credential_id=credential_id,\n access_type=connector_data.access_type,\n cc_pair_name=connector_data.name,\n groups=connector_data.groups,\n )\n\n # trigger indexing immediately\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_INDEXING,\n priority=OnyxCeleryPriority.HIGH,\n kwargs={\"tenant_id\": tenant_id},\n )\n\n logger.info(\n f\"create_connector_with_mock_credential - running check_for_indexing: cc_pair={response.data}\"\n )\n\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id),\n event=MilestoneRecordType.CREATED_CONNECTOR,\n )\n return response\n\n except ConnectorValidationError as e:\n raise HTTPException(\n status_code=400, detail=\"Connector validation error: \" + str(e)\n )\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n\n@router.patch(\"/admin/connector/{connector_id}\", tags=PUBLIC_API_TAGS)\ndef update_connector_from_model(\n connector_id: int,\n connector_data: ConnectorUpdateRequest,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> ConnectorSnapshot | StatusResponse[int]:\n cc_pair = fetch_connector_credential_pair_for_connector(db_session, connector_id)\n try:\n _validate_connector_allowed(connector_data.source)\n fetch_ee_implementation_or_noop(\n \"onyx.db.user_group\", \"validate_object_creation_for_user\", None\n )(\n db_session=db_session,\n user=user,\n target_group_ids=connector_data.groups,\n object_is_public=connector_data.access_type == AccessType.PUBLIC,\n object_is_perm_sync=connector_data.access_type == AccessType.SYNC,\n object_is_owned_by_user=cc_pair and user and cc_pair.creator_id == user.id,\n )\n connector_base = connector_data.to_connector_base()\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n updated_connector = update_connector(connector_id, connector_base, db_session)\n if updated_connector is None:\n raise HTTPException(\n status_code=404, detail=f\"Connector {connector_id} does not exist\"\n )\n\n return ConnectorSnapshot(\n id=updated_connector.id,\n name=updated_connector.name,\n source=updated_connector.source,\n input_type=updated_connector.input_type,\n connector_specific_config=updated_connector.connector_specific_config,\n refresh_freq=updated_connector.refresh_freq,\n prune_freq=updated_connector.prune_freq,\n credential_ids=[\n association.credential.id for association in updated_connector.credentials\n ],\n indexing_start=updated_connector.indexing_start,\n time_created=updated_connector.time_created,\n time_updated=updated_connector.time_updated,\n )\n\n\n@router.delete(\n \"/admin/connector/{connector_id}\",\n response_model=StatusResponse[int],\n tags=PUBLIC_API_TAGS,\n)\ndef delete_connector_by_id(\n connector_id: int,\n _: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[int]:\n try:\n with db_session.begin():\n return delete_connector(\n db_session=db_session,\n connector_id=connector_id,\n )\n except AssertionError:\n raise HTTPException(status_code=400, detail=\"Connector is not deletable\")\n\n\n@router.post(\"/admin/connector/run-once\", tags=PUBLIC_API_TAGS)\ndef connector_run_once(\n run_info: RunConnectorRequest,\n _: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse[int]:\n \"\"\"Used to trigger indexing on a set of cc_pairs associated with a\n single connector.\"\"\"\n tenant_id = get_current_tenant_id()\n\n connector_id = run_info.connector_id\n specified_credential_ids = run_info.credential_ids\n\n try:\n possible_credential_ids = get_connector_credential_ids(\n run_info.connector_id, db_session\n )\n except ValueError:\n raise HTTPException(\n status_code=404,\n detail=f\"Connector by id {connector_id} does not exist.\",\n )\n\n if not specified_credential_ids:\n credential_ids = possible_credential_ids\n else:\n if set(specified_credential_ids).issubset(set(possible_credential_ids)):\n credential_ids = specified_credential_ids\n else:\n raise HTTPException(\n status_code=400,\n detail=\"Not all specified credentials are associated with connector\",\n )\n\n if not credential_ids:\n raise HTTPException(\n status_code=400,\n detail=\"Connector has no valid credentials, cannot create index attempts.\",\n )\n try:\n num_triggers = trigger_indexing_for_cc_pair(\n credential_ids,\n connector_id,\n run_info.from_beginning,\n tenant_id,\n db_session,\n )\n except ValueError as e:\n raise HTTPException(status_code=400, detail=str(e))\n\n logger.info(\"connector_run_once - running check_for_indexing\")\n\n msg = f\"Marked {num_triggers} index attempts with indexing triggers.\"\n return StatusResponse(\n success=True,\n message=msg,\n data=num_triggers,\n )\n\n\n\"\"\"Endpoints for basic users\"\"\"\n\n\n@router.get(\"/connector/gmail/authorize/{credential_id}\")\ndef gmail_auth(\n response: Response, credential_id: str, _: User = Depends(current_user)\n) -> AuthUrl:\n # set a cookie that we can read in the callback (used for `verify_csrf`)\n response.set_cookie(\n key=_GMAIL_CREDENTIAL_ID_COOKIE_NAME,\n value=credential_id,\n httponly=True,\n max_age=600,\n )\n return AuthUrl(auth_url=get_auth_url(int(credential_id), DocumentSource.GMAIL))\n\n\n@router.get(\"/connector/google-drive/authorize/{credential_id}\")\ndef google_drive_auth(\n response: Response, credential_id: str, _: User = Depends(current_user)\n) -> AuthUrl:\n # set a cookie that we can read in the callback (used for `verify_csrf`)\n response.set_cookie(\n key=_GOOGLE_DRIVE_CREDENTIAL_ID_COOKIE_NAME,\n value=credential_id,\n httponly=True,\n max_age=600,\n )\n return AuthUrl(\n auth_url=get_auth_url(int(credential_id), DocumentSource.GOOGLE_DRIVE)\n )\n\n\n@router.get(\"/connector/gmail/callback\")\ndef gmail_callback(\n request: Request,\n callback: GmailCallback = Depends(),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n credential_id_cookie = request.cookies.get(_GMAIL_CREDENTIAL_ID_COOKIE_NAME)\n if credential_id_cookie is None or not credential_id_cookie.isdigit():\n raise HTTPException(\n status_code=401, detail=\"Request did not pass CSRF verification.\"\n )\n credential_id = int(credential_id_cookie)\n verify_csrf(credential_id, callback.state)\n credentials: Credentials | None = update_credential_access_tokens(\n callback.code,\n credential_id,\n user,\n db_session,\n DocumentSource.GMAIL,\n GoogleOAuthAuthenticationMethod.UPLOADED,\n )\n if credentials is None:\n raise HTTPException(\n status_code=500, detail=\"Unable to fetch Gmail access tokens\"\n )\n\n return StatusResponse(success=True, message=\"Updated Gmail access tokens\")\n\n\n@router.get(\"/connector/google-drive/callback\")\ndef google_drive_callback(\n request: Request,\n callback: GDriveCallback = Depends(),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n credential_id_cookie = request.cookies.get(_GOOGLE_DRIVE_CREDENTIAL_ID_COOKIE_NAME)\n if credential_id_cookie is None or not credential_id_cookie.isdigit():\n raise HTTPException(\n status_code=401, detail=\"Request did not pass CSRF verification.\"\n )\n credential_id = int(credential_id_cookie)\n verify_csrf(credential_id, callback.state)\n\n credentials: Credentials | None = update_credential_access_tokens(\n callback.code,\n credential_id,\n user,\n db_session,\n DocumentSource.GOOGLE_DRIVE,\n GoogleOAuthAuthenticationMethod.UPLOADED,\n )\n if credentials is None:\n raise HTTPException(\n status_code=500, detail=\"Unable to fetch Google Drive access tokens\"\n )\n\n return StatusResponse(success=True, message=\"Updated Google Drive access tokens\")\n\n\n@router.get(\"/connector\", tags=PUBLIC_API_TAGS)\ndef get_connectors(\n _: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> list[ConnectorSnapshot]:\n connectors = fetch_connectors(db_session)\n return [\n ConnectorSnapshot.from_connector_db_model(connector)\n for connector in connectors\n # don't include INGESTION_API, as it's not a \"real\"\n # connector like those created by the user\n if connector.source != DocumentSource.INGESTION_API\n ]\n\n\n@router.get(\"/indexed-sources\", tags=PUBLIC_API_TAGS)\ndef get_indexed_sources(\n _: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> IndexedSourcesResponse:\n sources = sorted(\n fetch_unique_document_sources(db_session), key=lambda source: source.value\n )\n return IndexedSourcesResponse(sources=sources)\n\n\n@router.get(\"/connector/{connector_id}\", tags=PUBLIC_API_TAGS)\ndef get_connector_by_id(\n connector_id: int,\n _: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> ConnectorSnapshot | StatusResponse[int]:\n connector = fetch_connector_by_id(connector_id, db_session)\n if connector is None:\n raise HTTPException(\n status_code=404, detail=f\"Connector {connector_id} does not exist\"\n )\n\n return ConnectorSnapshot(\n id=connector.id,\n name=connector.name,\n source=connector.source,\n indexing_start=connector.indexing_start,\n input_type=connector.input_type,\n connector_specific_config=connector.connector_specific_config,\n refresh_freq=connector.refresh_freq,\n prune_freq=connector.prune_freq,\n credential_ids=[\n association.credential.id for association in connector.credentials\n ],\n time_created=connector.time_created,\n time_updated=connector.time_updated,\n )\n\n\n@router.post(\"/connector-request\")\ndef submit_connector_request(\n request_data: ConnectorRequestSubmission,\n user: User = Depends(current_user),\n) -> StatusResponse:\n \"\"\"\n Submit a connector request for Cloud deployments.\n Tracks via PostHog telemetry and sends email to hello@onyx.app.\n \"\"\"\n tenant_id = get_current_tenant_id()\n connector_name = request_data.connector_name.strip()\n\n if not connector_name:\n raise HTTPException(status_code=400, detail=\"Connector name cannot be empty\")\n\n user_email = user.email\n\n # Track connector request via PostHog telemetry (Cloud only)\n from shared_configs.configs import MULTI_TENANT\n\n if MULTI_TENANT:\n mt_cloud_telemetry(\n tenant_id=tenant_id,\n distinct_id=str(user.id),\n event=MilestoneRecordType.REQUESTED_CONNECTOR,\n properties={\n \"connector_name\": connector_name,\n \"user_email\": user.email,\n },\n )\n\n # Send email notification (if email is configured)\n if EMAIL_CONFIGURED:\n try:\n subject = \"Onyx Craft Connector Request\"\n email_body_text = f\"\"\"A new connector request has been submitted:\n\nConnector Name: {connector_name}\nUser Email: {user_email or \"Not provided (anonymous user)\"}\nTenant ID: {tenant_id}\n\"\"\"\n email_body_html = f\"\"\"\n\n

    A new connector request has been submitted:

    \n
      \n
    • Connector Name: {connector_name}
      • \n
    • User Email: {user_email or \"Not provided (anonymous user)\"}
      • \n
    • Tenant ID: {tenant_id}
      • \n
    \n\n\"\"\"\n\n send_email(\n user_email=\"hello@onyx.app\",\n subject=subject,\n html_body=email_body_html,\n text_body=email_body_text,\n )\n logger.info(\n f\"Connector request email sent to hello@onyx.app for connector: {connector_name}\"\n )\n except Exception as e:\n # Log error but don't fail the request if email fails\n logger.error(\n f\"Failed to send connector request email for {connector_name}: {e}\"\n )\n\n logger.info(\n f\"Connector request submitted: {connector_name} by user {user_email or 'anonymous'} (tenant: {tenant_id})\"\n )\n\n return StatusResponse(\n success=True,\n message=\"Connector request submitted successfully. We'll prioritize popular requests!\",\n )\n\n\nclass BasicCCPairInfo(BaseModel):\n has_successful_run: bool\n source: DocumentSource\n status: ConnectorCredentialPairStatus\n\n\n@router.get(\"/connector-status\", tags=PUBLIC_API_TAGS)\ndef get_basic_connector_indexing_status(\n user: User = Depends(current_chat_accessible_user),\n db_session: Session = Depends(get_session),\n) -> list[BasicCCPairInfo]:\n cc_pairs = get_connector_credential_pairs_for_user(\n db_session=db_session,\n eager_load_connector=True,\n get_editable=False,\n user=user,\n )\n\n # NOTE: This endpoint excludes Craft connectors\n return [\n BasicCCPairInfo(\n has_successful_run=cc_pair.last_successful_index_time is not None,\n source=cc_pair.connector.source,\n status=cc_pair.status,\n )\n for cc_pair in cc_pairs\n if cc_pair.connector.source != DocumentSource.INGESTION_API\n and cc_pair.processing_mode == ProcessingMode.REGULAR\n ]\n\n\ndef trigger_indexing_for_cc_pair(\n specified_credential_ids: list[int],\n connector_id: int,\n from_beginning: bool,\n tenant_id: str,\n db_session: Session,\n) -> int:\n try:\n possible_credential_ids = get_connector_credential_ids(connector_id, db_session)\n except ValueError as e:\n raise ValueError(f\"Connector by id {connector_id} does not exist: {str(e)}\")\n\n if not specified_credential_ids:\n credential_ids = possible_credential_ids\n else:\n if set(specified_credential_ids).issubset(set(possible_credential_ids)):\n credential_ids = specified_credential_ids\n else:\n raise ValueError(\n \"Not all specified credentials are associated with connector\"\n )\n\n if not credential_ids:\n raise ValueError(\n \"Connector has no valid credentials, cannot create index attempts.\"\n )\n\n # Prevents index attempts for cc pairs that already have an index attempt currently running\n skipped_credentials = [\n credential_id\n for credential_id in credential_ids\n if get_index_attempts_for_cc_pair(\n cc_pair_identifier=ConnectorCredentialPairIdentifier(\n connector_id=connector_id,\n credential_id=credential_id,\n ),\n only_current=True,\n db_session=db_session,\n disinclude_finished=True,\n )\n ]\n\n connector_credential_pairs = [\n get_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n )\n for credential_id in credential_ids\n if credential_id not in skipped_credentials\n ]\n\n num_triggers = 0\n for cc_pair in connector_credential_pairs:\n if cc_pair is not None:\n indexing_mode = IndexingMode.UPDATE\n if from_beginning:\n indexing_mode = IndexingMode.REINDEX\n\n mark_ccpair_with_indexing_trigger(cc_pair.id, indexing_mode, db_session)\n num_triggers += 1\n\n logger.info(\n f\"connector_run_once - marking cc_pair with indexing trigger: \"\n f\"connector={connector_id} \"\n f\"cc_pair={cc_pair.id} \"\n f\"indexing_trigger={indexing_mode}\"\n )\n\n priority = OnyxCeleryPriority.HIGH\n\n # run the beat task to pick up the triggers immediately\n logger.info(f\"Sending indexing check task with priority {priority}\")\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_INDEXING,\n priority=priority,\n kwargs={\"tenant_id\": tenant_id},\n )\n\n return num_triggers\n" }, { "path": "backend/onyx/server/documents/credential.py", "content": "import json\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import File\nfrom fastapi import Form\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi import UploadFile\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_admin_user\nfrom onyx.auth.users import current_curator_or_admin_user\nfrom onyx.auth.users import current_user\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.connectors.factory import validate_ccpair_for_user\nfrom onyx.db.credentials import alter_credential\nfrom onyx.db.credentials import cleanup_gmail_credentials\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.credentials import CREDENTIAL_PERMISSIONS_TO_IGNORE\nfrom onyx.db.credentials import delete_credential\nfrom onyx.db.credentials import delete_credential_for_user\nfrom onyx.db.credentials import fetch_credential_by_id_for_user\nfrom onyx.db.credentials import fetch_credentials_by_source_for_user\nfrom onyx.db.credentials import fetch_credentials_for_user\nfrom onyx.db.credentials import swap_credentials_connector\nfrom onyx.db.credentials import update_credential\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import DocumentSource\nfrom onyx.db.models import User\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.server.documents.models import CredentialDataUpdateRequest\nfrom onyx.server.documents.models import CredentialSnapshot\nfrom onyx.server.documents.models import CredentialSwapRequest\nfrom onyx.server.documents.models import ObjectCreationIdResponse\nfrom onyx.server.documents.private_key_types import FILE_TYPE_TO_FILE_PROCESSOR\nfrom onyx.server.documents.private_key_types import PrivateKeyFileTypes\nfrom onyx.server.documents.private_key_types import ProcessPrivateKeyFileProtocol\nfrom onyx.server.models import StatusResponse\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\nlogger = setup_logger()\n\n\nrouter = APIRouter(prefix=\"/manage\", tags=PUBLIC_API_TAGS)\n\n\ndef _ignore_credential_permissions(source: DocumentSource) -> bool:\n return source in CREDENTIAL_PERMISSIONS_TO_IGNORE\n\n\n\"\"\"Admin-only endpoints\"\"\"\n\n\n@router.get(\"/admin/credential\")\ndef list_credentials_admin(\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> list[CredentialSnapshot]:\n \"\"\"Lists all public credentials\"\"\"\n credentials = fetch_credentials_for_user(\n db_session=db_session,\n user=user,\n get_editable=False,\n )\n return [\n CredentialSnapshot.from_credential_db_model(credential)\n for credential in credentials\n ]\n\n\n@router.get(\"/admin/similar-credentials/{source_type}\")\ndef get_cc_source_full_info(\n source_type: DocumentSource,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n get_editable: bool = Query(\n False, description=\"If true, return editable credentials\"\n ),\n) -> list[CredentialSnapshot]:\n credentials = fetch_credentials_by_source_for_user(\n db_session=db_session,\n user=user,\n document_source=source_type,\n get_editable=get_editable,\n )\n\n return [\n CredentialSnapshot.from_credential_db_model(credential)\n for credential in credentials\n ]\n\n\n@router.delete(\"/admin/credential/{credential_id}\")\ndef delete_credential_by_id_admin(\n credential_id: int,\n _: User = Depends(current_admin_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n \"\"\"Same as the user endpoint, but can delete any credential (not just the user's own)\"\"\"\n delete_credential(db_session=db_session, credential_id=credential_id)\n return StatusResponse(\n success=True, message=\"Credential deleted successfully\", data=credential_id\n )\n\n\n@router.put(\"/admin/credential/swap\")\ndef swap_credentials_for_connector(\n credential_swap_req: CredentialSwapRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n validate_ccpair_for_user(\n credential_swap_req.connector_id,\n credential_swap_req.new_credential_id,\n credential_swap_req.access_type,\n db_session,\n )\n\n connector_credential_pair = swap_credentials_connector(\n new_credential_id=credential_swap_req.new_credential_id,\n connector_id=credential_swap_req.connector_id,\n db_session=db_session,\n user=user,\n )\n\n return StatusResponse(\n success=True,\n message=\"Credential swapped successfully\",\n data=connector_credential_pair.id,\n )\n\n\n@router.post(\"/credential\")\ndef create_credential_from_model(\n credential_info: CredentialBase,\n user: User = Depends(current_curator_or_admin_user),\n db_session: Session = Depends(get_session),\n) -> ObjectCreationIdResponse:\n if not _ignore_credential_permissions(credential_info.source):\n fetch_ee_implementation_or_noop(\n \"onyx.db.user_group\", \"validate_object_creation_for_user\", None\n )(\n db_session=db_session,\n user=user,\n target_group_ids=credential_info.groups,\n object_is_public=credential_info.curator_public,\n )\n\n # Temporary fix for empty Google App credentials\n if credential_info.source == DocumentSource.GMAIL:\n cleanup_gmail_credentials(db_session=db_session)\n\n credential = create_credential(credential_info, user, db_session)\n return ObjectCreationIdResponse(\n id=credential.id,\n credential=CredentialSnapshot.from_credential_db_model(credential),\n )\n\n\n@router.post(\"/credential/private-key\")\ndef create_credential_with_private_key(\n credential_json: str = Form(...),\n admin_public: bool = Form(False),\n curator_public: bool = Form(False),\n groups: list[int] = Form([]),\n name: str | None = Form(None),\n source: str = Form(...),\n user: User = Depends(current_curator_or_admin_user),\n uploaded_file: UploadFile = File(...),\n field_key: str = Form(...),\n type_definition_key: str = Form(...),\n db_session: Session = Depends(get_session),\n) -> ObjectCreationIdResponse:\n try:\n credential_data = json.loads(credential_json)\n except json.JSONDecodeError as e:\n raise HTTPException(\n status_code=400,\n detail=f\"Invalid JSON in credential_json: {str(e)}\",\n )\n\n private_key_processor: ProcessPrivateKeyFileProtocol | None = (\n FILE_TYPE_TO_FILE_PROCESSOR.get(PrivateKeyFileTypes(type_definition_key))\n )\n if private_key_processor is None:\n raise HTTPException(\n status_code=400,\n detail=\"Invalid type definition key for private key file\",\n )\n private_key_content: str = private_key_processor(uploaded_file)\n\n credential_data[field_key] = private_key_content\n\n credential_info = CredentialBase(\n credential_json=credential_data,\n admin_public=admin_public,\n curator_public=curator_public,\n groups=groups,\n name=name,\n source=DocumentSource(source),\n )\n\n if not _ignore_credential_permissions(DocumentSource(source)):\n fetch_ee_implementation_or_noop(\n \"onyx.db.user_group\", \"validate_object_creation_for_user\", None\n )(\n db_session=db_session,\n user=user,\n target_group_ids=groups,\n object_is_public=curator_public,\n )\n\n # Temporary fix for empty Google App credentials\n if DocumentSource(source) == DocumentSource.GMAIL:\n cleanup_gmail_credentials(db_session=db_session)\n\n credential = create_credential(credential_info, user, db_session)\n return ObjectCreationIdResponse(\n id=credential.id,\n credential=CredentialSnapshot.from_credential_db_model(credential),\n )\n\n\n\"\"\"Endpoints for all\"\"\"\n\n\n@router.get(\"/credential\")\ndef list_credentials(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> list[CredentialSnapshot]:\n credentials = fetch_credentials_for_user(db_session=db_session, user=user)\n return [\n CredentialSnapshot.from_credential_db_model(credential)\n for credential in credentials\n ]\n\n\n@router.get(\"/credential/{credential_id}\")\ndef get_credential_by_id(\n credential_id: int,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> CredentialSnapshot | StatusResponse[int]:\n credential = fetch_credential_by_id_for_user(\n credential_id,\n user,\n db_session,\n get_editable=False,\n )\n if credential is None:\n raise HTTPException(\n status_code=401,\n detail=f\"Credential {credential_id} does not exist or does not belong to user\",\n )\n\n return CredentialSnapshot.from_credential_db_model(credential)\n\n\n@router.put(\"/admin/credential/{credential_id}\")\ndef update_credential_data(\n credential_id: int,\n credential_update: CredentialDataUpdateRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> CredentialBase:\n credential = alter_credential(\n credential_id,\n credential_update.name,\n credential_update.credential_json,\n user,\n db_session,\n )\n\n if credential is None:\n raise HTTPException(\n status_code=401,\n detail=f\"Credential {credential_id} does not exist or does not belong to user\",\n )\n\n return CredentialSnapshot.from_credential_db_model(credential)\n\n\n@router.put(\"/admin/credential/private-key/{credential_id}\")\ndef update_credential_private_key(\n credential_id: int,\n name: str = Form(...),\n credential_json: str = Form(...),\n uploaded_file: UploadFile = File(...),\n field_key: str = Form(...),\n type_definition_key: str = Form(...),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> CredentialBase:\n try:\n credential_data = json.loads(credential_json)\n except json.JSONDecodeError as e:\n raise HTTPException(\n status_code=400,\n detail=f\"Invalid JSON in credential_json: {str(e)}\",\n )\n\n private_key_processor: ProcessPrivateKeyFileProtocol | None = (\n FILE_TYPE_TO_FILE_PROCESSOR.get(PrivateKeyFileTypes(type_definition_key))\n )\n if private_key_processor is None:\n raise HTTPException(\n status_code=400,\n detail=\"Invalid type definition key for private key file\",\n )\n private_key_content: str = private_key_processor(uploaded_file)\n credential_data[field_key] = private_key_content\n\n credential = alter_credential(\n credential_id,\n name,\n credential_data,\n user,\n db_session,\n )\n\n if credential is None:\n raise HTTPException(\n status_code=401,\n detail=f\"Credential {credential_id} does not exist or does not belong to user\",\n )\n\n return CredentialSnapshot.from_credential_db_model(credential)\n\n\n@router.patch(\"/credential/{credential_id}\")\ndef update_credential_from_model(\n credential_id: int,\n credential_data: CredentialBase,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> CredentialSnapshot | StatusResponse[int]:\n updated_credential = update_credential(\n credential_id, credential_data, user, db_session\n )\n if updated_credential is None:\n raise HTTPException(\n status_code=401,\n detail=f\"Credential {credential_id} does not exist or does not belong to user\",\n )\n\n # Get credential_json value - use masking for API responses\n credential_json_value = (\n updated_credential.credential_json.get_value(apply_mask=True)\n if updated_credential.credential_json\n else {}\n )\n\n return CredentialSnapshot(\n source=updated_credential.source,\n id=updated_credential.id,\n credential_json=credential_json_value,\n user_id=updated_credential.user_id,\n name=updated_credential.name,\n admin_public=updated_credential.admin_public,\n time_created=updated_credential.time_created,\n time_updated=updated_credential.time_updated,\n curator_public=updated_credential.curator_public,\n )\n\n\n@router.delete(\"/credential/{credential_id}\")\ndef delete_credential_by_id(\n credential_id: int,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n delete_credential_for_user(\n credential_id,\n user,\n db_session,\n )\n\n return StatusResponse(\n success=True, message=\"Credential deleted successfully\", data=credential_id\n )\n\n\n@router.delete(\"/credential/force/{credential_id}\")\ndef force_delete_credential_by_id(\n credential_id: int,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> StatusResponse:\n delete_credential_for_user(credential_id, user, db_session, True)\n\n return StatusResponse(\n success=True, message=\"Credential deleted successfully\", data=credential_id\n )\n" }, { "path": "backend/onyx/server/documents/document.py", "content": "from fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_user\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.context.search.preprocessing.access_filters import (\n build_access_filters_for_user,\n)\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.document_index.factory import get_default_document_index\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.natural_language_processing.utils import get_tokenizer\nfrom onyx.prompts.prompt_utils import build_doc_context_str\nfrom onyx.server.documents.models import ChunkInfo\nfrom onyx.server.documents.models import DocumentInfo\nfrom onyx.server.utils_vector_db import require_vector_db\n\n\nrouter = APIRouter(prefix=\"/document\")\n\n\n# Have to use a query parameter as FastAPI is interpreting the URL type document_ids\n# as a different path\n@router.get(\"/document-size-info\", dependencies=[Depends(require_vector_db)])\ndef get_document_info(\n document_id: str = Query(...),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> DocumentInfo:\n search_settings = get_current_search_settings(db_session)\n # This flow is for search so we do not get all indices.\n document_index = get_default_document_index(search_settings, None, db_session)\n\n user_acl_filters = build_access_filters_for_user(user, db_session)\n inference_chunks = document_index.id_based_retrieval(\n chunk_requests=[VespaChunkRequest(document_id=document_id)],\n filters=IndexFilters(access_control_list=user_acl_filters),\n )\n\n if not inference_chunks:\n raise HTTPException(status_code=404, detail=\"Document not found\")\n\n contents = [chunk.content for chunk in inference_chunks]\n\n combined_contents = \"\\n\".join(contents)\n\n # get actual document context used for LLM\n first_chunk = inference_chunks[0]\n tokenizer_encode = get_tokenizer(\n provider_type=search_settings.provider_type,\n model_name=search_settings.model_name,\n ).encode\n full_context_str = build_doc_context_str(\n semantic_identifier=first_chunk.semantic_identifier,\n source_type=first_chunk.source_type,\n content=combined_contents,\n metadata_dict=first_chunk.metadata,\n updated_at=first_chunk.updated_at,\n ind=0,\n )\n\n return DocumentInfo(\n num_chunks=len(inference_chunks),\n num_tokens=len(tokenizer_encode(full_context_str)),\n )\n\n\n@router.get(\"/chunk-info\", dependencies=[Depends(require_vector_db)])\ndef get_chunk_info(\n document_id: str = Query(...),\n chunk_id: int = Query(...),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> ChunkInfo:\n search_settings = get_current_search_settings(db_session)\n # This flow is for search so we do not get all indices.\n document_index = get_default_document_index(search_settings, None, db_session)\n\n user_acl_filters = build_access_filters_for_user(user, db_session)\n chunk_request = VespaChunkRequest(\n document_id=document_id,\n min_chunk_ind=chunk_id,\n max_chunk_ind=chunk_id,\n )\n\n inference_chunks = document_index.id_based_retrieval(\n chunk_requests=[chunk_request],\n filters=IndexFilters(access_control_list=user_acl_filters),\n batch_retrieval=True,\n )\n\n if not inference_chunks:\n raise HTTPException(status_code=404, detail=\"Chunk not found\")\n\n chunk_content = inference_chunks[0].content\n\n tokenizer_encode = get_tokenizer(\n provider_type=search_settings.provider_type,\n model_name=search_settings.model_name,\n ).encode\n\n return ChunkInfo(\n content=chunk_content, num_tokens=len(tokenizer_encode(chunk_content))\n )\n" }, { "path": "backend/onyx/server/documents/document_utils.py", "content": "from cryptography.hazmat.primitives.serialization import pkcs12\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef _is_password_related_error(error: Exception) -> bool:\n \"\"\"\n Check if the exception indicates a password-related issue rather than a format issue.\n \"\"\"\n error_msg = str(error).lower()\n password_keywords = [\"mac\", \"integrity\", \"password\", \"authentication\", \"verify\"]\n return any(keyword in error_msg for keyword in password_keywords)\n\n\ndef validate_pkcs12_content(file_bytes: bytes) -> bool:\n \"\"\"\n Validate that the file content is actually a PKCS#12 file.\n This performs basic format validation without requiring passwords.\n \"\"\"\n try:\n # Basic file size check\n if len(file_bytes) < 10:\n logger.debug(\"File too small to be a valid PKCS#12 file\")\n return False\n\n # Check for PKCS#12 magic bytes/ASN.1 structure\n # PKCS#12 files start with ASN.1 SEQUENCE tag (0x30)\n if file_bytes[0] != 0x30:\n logger.debug(\"File does not start with ASN.1 SEQUENCE tag\")\n return False\n\n # Try to parse the outer ASN.1 structure without password validation\n # This checks if the file has the basic PKCS#12 structure\n try:\n # Attempt to load just to validate the basic format\n # We expect this to fail due to password, but it should fail with a specific error\n pkcs12.load_key_and_certificates(file_bytes, password=None)\n return True\n except ValueError as e:\n # Check if the error is related to password (expected) vs format issues\n if _is_password_related_error(e):\n # These errors indicate the file format is correct but password is wrong/missing\n logger.debug(\n f\"PKCS#12 format appears valid, password-related error: {e}\"\n )\n return True\n else:\n # Other ValueError likely indicates format issues\n logger.debug(f\"PKCS#12 format validation failed: {e}\")\n return False\n except Exception as e:\n # Try with empty password as fallback\n try:\n pkcs12.load_key_and_certificates(file_bytes, password=b\"\")\n return True\n except ValueError as e2:\n if _is_password_related_error(e2):\n logger.debug(\n f\"PKCS#12 format appears valid with empty password attempt: {e2}\"\n )\n return True\n else:\n logger.debug(\n f\"PKCS#12 validation failed on both attempts: {e}, {e2}\"\n )\n return False\n except Exception:\n logger.debug(f\"PKCS#12 validation failed: {e}\")\n return False\n\n except Exception as e:\n logger.debug(f\"Unexpected error during PKCS#12 validation: {e}\")\n return False\n" }, { "path": "backend/onyx/server/documents/models.py", "content": "from collections.abc import Sequence\nfrom datetime import datetime\nfrom datetime import timezone\nfrom datetime import UTC\nfrom enum import Enum\nfrom typing import Any\nfrom typing import Generic\nfrom typing import TypeVar\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\nfrom pydantic import Field\n\nfrom onyx.configs.app_configs import MASK_CREDENTIAL_PREFIX\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.enums import ProcessingMode\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import DocPermissionSyncAttempt\nfrom onyx.db.models import Document as DbDocument\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexingStatus\nfrom onyx.db.models import TaskStatus\nfrom onyx.server.federated.models import FederatedConnectorStatus\nfrom onyx.utils.variable_functionality import fetch_ee_implementation_or_noop\n\n\nclass DocumentSyncStatus(BaseModel):\n doc_id: str\n last_synced: datetime | None\n last_modified: datetime | None\n\n @classmethod\n def from_model(cls, doc: DbDocument) -> \"DocumentSyncStatus\":\n return DocumentSyncStatus(\n doc_id=doc.id,\n last_synced=doc.last_synced,\n last_modified=doc.last_modified,\n )\n\n\nclass DocumentInfo(BaseModel):\n num_chunks: int\n num_tokens: int\n\n\nclass ChunkInfo(BaseModel):\n content: str\n num_tokens: int\n\n\nclass IndexedSourcesResponse(BaseModel):\n model_config = ConfigDict(use_enum_values=True)\n sources: list[DocumentSource]\n\n\nclass DeletionAttemptSnapshot(BaseModel):\n connector_id: int\n credential_id: int\n status: TaskStatus\n\n\nclass ConnectorBase(BaseModel):\n name: str\n source: DocumentSource\n input_type: InputType\n connector_specific_config: dict[str, Any]\n # In seconds, None for one time index with no refresh\n refresh_freq: int | None = None\n prune_freq: int | None = None\n indexing_start: datetime | None = None\n\n\nclass ConnectorUpdateRequest(ConnectorBase):\n access_type: AccessType\n groups: list[int] = Field(default_factory=list)\n\n def to_connector_base(self) -> ConnectorBase:\n return ConnectorBase(**self.model_dump(exclude={\"access_type\", \"groups\"}))\n\n\nclass ConnectorSnapshot(ConnectorBase):\n id: int\n credential_ids: list[int]\n time_created: datetime\n time_updated: datetime\n source: DocumentSource\n\n @classmethod\n def from_connector_db_model(\n cls, connector: Connector, credential_ids: list[int] | None = None\n ) -> \"ConnectorSnapshot\":\n return ConnectorSnapshot(\n id=connector.id,\n name=connector.name,\n source=connector.source,\n input_type=connector.input_type,\n connector_specific_config=connector.connector_specific_config,\n refresh_freq=connector.refresh_freq,\n prune_freq=connector.prune_freq,\n credential_ids=(\n credential_ids\n or [association.credential.id for association in connector.credentials]\n ),\n indexing_start=connector.indexing_start,\n time_created=connector.time_created,\n time_updated=connector.time_updated,\n )\n\n\nclass CredentialSwapRequest(BaseModel):\n new_credential_id: int\n connector_id: int\n access_type: AccessType\n\n\nclass CredentialDataUpdateRequest(BaseModel):\n name: str\n credential_json: dict[str, Any]\n\n\nclass CredentialBase(BaseModel):\n credential_json: dict[str, Any]\n # if `true`, then all Admins will have access to the credential\n admin_public: bool\n source: DocumentSource\n name: str | None = None\n curator_public: bool = False\n groups: list[int] = Field(default_factory=list)\n\n\nclass CredentialSnapshot(CredentialBase):\n id: int\n user_id: UUID | None\n user_email: str | None = None\n time_created: datetime\n time_updated: datetime\n\n @classmethod\n def from_credential_db_model(cls, credential: Credential) -> \"CredentialSnapshot\":\n # Get the credential_json value with appropriate masking\n if credential.credential_json is None:\n credential_json_value: dict[str, Any] = {}\n elif MASK_CREDENTIAL_PREFIX:\n credential_json_value = credential.credential_json.get_value(\n apply_mask=True\n )\n else:\n credential_json_value = credential.credential_json.get_value(\n apply_mask=False\n )\n\n return CredentialSnapshot(\n id=credential.id,\n credential_json=credential_json_value,\n user_id=credential.user_id,\n user_email=credential.user.email if credential.user else None,\n admin_public=credential.admin_public,\n time_created=credential.time_created,\n time_updated=credential.time_updated,\n source=credential.source or DocumentSource.NOT_APPLICABLE,\n name=credential.name,\n curator_public=credential.curator_public,\n )\n\n\nclass IndexAttemptSnapshot(BaseModel):\n id: int\n status: IndexingStatus | None\n from_beginning: bool\n new_docs_indexed: int # only includes completely new docs\n total_docs_indexed: int # includes docs that are updated\n docs_removed_from_index: int\n error_msg: str | None\n error_count: int\n full_exception_trace: str | None\n time_started: str | None\n time_updated: str\n poll_range_start: datetime | None = None\n poll_range_end: datetime | None = None\n\n @classmethod\n def from_index_attempt_db_model(\n cls, index_attempt: IndexAttempt\n ) -> \"IndexAttemptSnapshot\":\n return IndexAttemptSnapshot(\n id=index_attempt.id,\n status=index_attempt.status,\n from_beginning=index_attempt.from_beginning,\n new_docs_indexed=index_attempt.new_docs_indexed or 0,\n total_docs_indexed=index_attempt.total_docs_indexed or 0,\n docs_removed_from_index=index_attempt.docs_removed_from_index or 0,\n error_msg=index_attempt.error_msg,\n error_count=len(index_attempt.error_rows),\n full_exception_trace=index_attempt.full_exception_trace,\n time_started=(\n index_attempt.time_started.isoformat()\n if index_attempt.time_started\n else None\n ),\n time_updated=index_attempt.time_updated.isoformat(),\n poll_range_start=index_attempt.poll_range_start,\n poll_range_end=index_attempt.poll_range_end,\n )\n\n\n# These are the types currently supported by the pagination hook\n# More api endpoints can be refactored and be added here for use with the pagination hook\nPaginatedType = TypeVar(\"PaginatedType\", bound=BaseModel)\n\n\nclass PermissionSyncAttemptSnapshot(BaseModel):\n id: int\n status: PermissionSyncStatus\n error_message: str | None\n total_docs_synced: int\n docs_with_permission_errors: int\n time_created: str\n time_started: str | None\n time_finished: str | None\n\n @classmethod\n def from_permission_sync_attempt_db_model(\n cls, attempt: DocPermissionSyncAttempt\n ) -> \"PermissionSyncAttemptSnapshot\":\n return PermissionSyncAttemptSnapshot(\n id=attempt.id,\n status=attempt.status,\n error_message=attempt.error_message,\n total_docs_synced=attempt.total_docs_synced or 0,\n docs_with_permission_errors=attempt.docs_with_permission_errors or 0,\n time_created=attempt.time_created.isoformat(),\n time_started=(\n attempt.time_started.isoformat() if attempt.time_started else None\n ),\n time_finished=(\n attempt.time_finished.isoformat() if attempt.time_finished else None\n ),\n )\n\n\nclass PaginatedReturn(BaseModel, Generic[PaginatedType]):\n items: list[PaginatedType]\n total_items: int\n\n\nclass CCPairFullInfo(BaseModel):\n id: int\n name: str\n status: ConnectorCredentialPairStatus\n in_repeated_error_state: bool\n num_docs_indexed: int\n connector: ConnectorSnapshot\n credential: CredentialSnapshot\n number_of_index_attempts: int\n last_index_attempt_status: IndexingStatus | None\n latest_deletion_attempt: DeletionAttemptSnapshot | None\n access_type: AccessType\n is_editable_for_current_user: bool\n deletion_failure_message: str | None\n indexing: bool\n creator: UUID | None\n creator_email: str | None\n\n # information on syncing/indexing\n last_indexed: datetime | None\n last_pruned: datetime | None\n # accounts for both doc sync and group sync\n last_full_permission_sync: datetime | None\n overall_indexing_speed: float | None\n latest_checkpoint_description: str | None\n\n # permission sync attempt status\n last_permission_sync_attempt_status: PermissionSyncStatus | None\n permission_syncing: bool\n last_permission_sync_attempt_finished: datetime | None\n last_permission_sync_attempt_error_message: str | None\n\n @classmethod\n def _get_last_full_permission_sync(\n cls, cc_pair_model: ConnectorCredentialPair\n ) -> datetime | None:\n check_if_source_requires_external_group_sync = fetch_ee_implementation_or_noop(\n \"onyx.external_permissions.sync_params\",\n \"source_requires_external_group_sync\",\n noop_return_value=False,\n )\n check_if_source_requires_doc_sync = fetch_ee_implementation_or_noop(\n \"onyx.external_permissions.sync_params\",\n \"source_requires_doc_sync\",\n noop_return_value=False,\n )\n\n needs_group_sync = check_if_source_requires_external_group_sync(\n cc_pair_model.connector.source\n )\n needs_doc_sync = check_if_source_requires_doc_sync(\n cc_pair_model.connector.source\n )\n\n last_group_sync = (\n cc_pair_model.last_time_external_group_sync\n if needs_group_sync\n else datetime.now(UTC)\n )\n last_doc_sync = (\n cc_pair_model.last_time_perm_sync if needs_doc_sync else datetime.now(UTC)\n )\n\n # if either is still None at this point, it means sync is necessary but\n # has never completed.\n if last_group_sync is None or last_doc_sync is None:\n return None\n\n return min(last_group_sync, last_doc_sync)\n\n @classmethod\n def from_models(\n cls,\n cc_pair_model: ConnectorCredentialPair,\n latest_deletion_attempt: DeletionAttemptSnapshot | None,\n number_of_index_attempts: int,\n last_index_attempt: IndexAttempt | None,\n num_docs_indexed: int, # not ideal, but this must be computed separately\n is_editable_for_current_user: bool,\n indexing: bool,\n last_successful_index_time: datetime | None = None,\n last_permission_sync_attempt_status: PermissionSyncStatus | None = None,\n permission_syncing: bool = False,\n last_permission_sync_attempt_finished: datetime | None = None,\n last_permission_sync_attempt_error_message: str | None = None,\n ) -> \"CCPairFullInfo\":\n # figure out if we need to artificially deflate the number of docs indexed.\n # This is required since the total number of docs indexed by a CC Pair is\n # updated before the new docs for an indexing attempt. If we don't do this,\n # there is a mismatch between these two numbers which may confuse users.\n last_indexing_status = last_index_attempt.status if last_index_attempt else None\n if (\n # only need to do this if the last indexing attempt is still in progress\n last_indexing_status == IndexingStatus.IN_PROGRESS\n and number_of_index_attempts == 1\n and last_index_attempt\n and last_index_attempt.new_docs_indexed\n ):\n num_docs_indexed = (\n last_index_attempt.new_docs_indexed if last_index_attempt else 0\n )\n\n overall_indexing_speed = num_docs_indexed / (\n (\n datetime.now(tz=timezone.utc) - cc_pair_model.connector.time_created\n ).total_seconds()\n / 60\n )\n\n return cls(\n id=cc_pair_model.id,\n name=cc_pair_model.name,\n status=cc_pair_model.status,\n in_repeated_error_state=cc_pair_model.in_repeated_error_state,\n num_docs_indexed=num_docs_indexed,\n connector=ConnectorSnapshot.from_connector_db_model(\n cc_pair_model.connector,\n credential_ids=[cc_pair_model.credential_id],\n ),\n credential=CredentialSnapshot.from_credential_db_model(\n cc_pair_model.credential\n ),\n number_of_index_attempts=number_of_index_attempts,\n last_index_attempt_status=last_indexing_status,\n latest_deletion_attempt=latest_deletion_attempt,\n access_type=cc_pair_model.access_type,\n is_editable_for_current_user=is_editable_for_current_user,\n deletion_failure_message=cc_pair_model.deletion_failure_message,\n indexing=indexing,\n creator=cc_pair_model.creator_id,\n creator_email=(\n cc_pair_model.creator.email if cc_pair_model.creator else None\n ),\n last_indexed=last_successful_index_time,\n last_pruned=cc_pair_model.last_pruned,\n last_full_permission_sync=cls._get_last_full_permission_sync(cc_pair_model),\n overall_indexing_speed=overall_indexing_speed,\n latest_checkpoint_description=None,\n last_permission_sync_attempt_status=last_permission_sync_attempt_status,\n permission_syncing=permission_syncing,\n last_permission_sync_attempt_finished=last_permission_sync_attempt_finished,\n last_permission_sync_attempt_error_message=last_permission_sync_attempt_error_message,\n )\n\n\nclass CeleryTaskStatus(BaseModel):\n id: str\n name: str\n status: TaskStatus\n start_time: datetime | None\n register_time: datetime | None\n\n\nclass FailedConnectorIndexingStatus(BaseModel):\n \"\"\"Simplified version of ConnectorIndexingStatus for failed indexing attempts\"\"\"\n\n cc_pair_id: int\n name: str\n error_msg: str | None\n is_deletable: bool\n connector_id: int\n credential_id: int\n\n\nclass ConnectorStatus(BaseModel):\n \"\"\"\n Represents the status of a connector,\n including indexing status elated information\n \"\"\"\n\n cc_pair_id: int\n name: str\n connector: ConnectorSnapshot\n credential: CredentialSnapshot\n access_type: AccessType\n groups: list[int]\n\n\nclass ConnectorIndexingStatus(ConnectorStatus):\n \"\"\"Represents the full indexing status of a connector\"\"\"\n\n cc_pair_status: ConnectorCredentialPairStatus\n # this is separate from the `status` above, since a connector can be `INITIAL_INDEXING`, `ACTIVE`,\n # or `PAUSED` and still be in a repeated error state.\n in_repeated_error_state: bool\n owner: str\n last_finished_status: IndexingStatus | None\n last_status: IndexingStatus | None\n last_success: datetime | None\n latest_index_attempt: IndexAttemptSnapshot | None\n docs_indexed: int\n in_progress: bool\n\n\nclass DocsCountOperator(str, Enum):\n GREATER_THAN = \">\"\n LESS_THAN = \"<\"\n EQUAL_TO = \"=\"\n\n\nclass ConnectorIndexingStatusLite(BaseModel):\n cc_pair_id: int\n name: str\n source: DocumentSource\n access_type: AccessType\n cc_pair_status: ConnectorCredentialPairStatus\n in_progress: bool\n in_repeated_error_state: bool\n last_finished_status: IndexingStatus | None\n last_status: IndexingStatus | None\n last_success: datetime | None\n is_editable: bool\n docs_indexed: int\n latest_index_attempt_docs_indexed: int | None\n\n\nclass SourceSummary(BaseModel):\n total_connectors: int\n active_connectors: int\n public_connectors: int\n total_docs_indexed: int\n\n\nclass ConnectorIndexingStatusLiteResponse(BaseModel):\n source: DocumentSource\n summary: SourceSummary\n current_page: int\n total_pages: int\n indexing_statuses: Sequence[ConnectorIndexingStatusLite | FederatedConnectorStatus]\n\n\nclass ConnectorCredentialPairIdentifier(BaseModel):\n connector_id: int\n credential_id: int\n\n\nclass ConnectorCredentialPairMetadata(BaseModel):\n name: str\n access_type: AccessType\n auto_sync_options: dict[str, Any] | None = None\n groups: list[int] = Field(default_factory=list)\n processing_mode: ProcessingMode = ProcessingMode.REGULAR\n\n\nclass CCStatusUpdateRequest(BaseModel):\n status: ConnectorCredentialPairStatus\n\n\nclass ConnectorCredentialPairDescriptor(BaseModel):\n id: int\n name: str\n connector: ConnectorSnapshot\n credential: CredentialSnapshot\n access_type: AccessType\n\n\nclass CCPairSummary(BaseModel):\n \"\"\"Simplified connector-credential pair information with just essential data\"\"\"\n\n id: int\n name: str\n source: DocumentSource\n access_type: AccessType\n\n @classmethod\n def from_cc_pair_descriptor(\n cls, descriptor: ConnectorCredentialPairDescriptor\n ) -> \"CCPairSummary\":\n return cls(\n id=descriptor.id,\n name=descriptor.name,\n source=descriptor.connector.source,\n access_type=descriptor.access_type,\n )\n\n\nclass RunConnectorRequest(BaseModel):\n connector_id: int\n credential_ids: list[int] | None = None\n from_beginning: bool = False\n\n\nclass ConnectorRequestSubmission(BaseModel):\n connector_name: str\n\n\nclass CCPropertyUpdateRequest(BaseModel):\n name: str\n value: str\n\n\n\"\"\"Connectors Models\"\"\"\n\n\nclass GoogleAppWebCredentials(BaseModel):\n client_id: str\n project_id: str\n auth_uri: str\n token_uri: str\n auth_provider_x509_cert_url: str\n client_secret: str\n redirect_uris: list[str]\n javascript_origins: list[str]\n\n\nclass GoogleAppCredentials(BaseModel):\n web: GoogleAppWebCredentials\n\n\nclass GoogleServiceAccountKey(BaseModel):\n type: str\n project_id: str\n private_key_id: str\n private_key: str\n client_email: str\n client_id: str\n auth_uri: str\n token_uri: str\n auth_provider_x509_cert_url: str\n client_x509_cert_url: str\n universe_domain: str\n\n\nclass GoogleServiceAccountCredentialRequest(BaseModel):\n google_primary_admin: str | None = None # email of user to impersonate\n\n\nclass FileUploadResponse(BaseModel):\n file_paths: list[str]\n file_names: list[str]\n zip_metadata_file_id: str | None # File ID pointing to metadata in file store\n\n\nclass ConnectorFileInfo(BaseModel):\n file_id: str\n file_name: str\n file_size: int | None = None\n upload_date: str | None = None\n\n\nclass ConnectorFilesResponse(BaseModel):\n files: list[ConnectorFileInfo]\n\n\nclass ObjectCreationIdResponse(BaseModel):\n id: int\n credential: CredentialSnapshot | None = None\n\n\nclass AuthStatus(BaseModel):\n authenticated: bool\n\n\nclass AuthUrl(BaseModel):\n auth_url: str\n\n\nclass GmailCallback(BaseModel):\n state: str\n code: str\n\n\nclass GDriveCallback(BaseModel):\n state: str\n code: str\n\n\nclass IndexingStatusRequest(BaseModel):\n secondary_index: bool = False\n source: DocumentSource | None = None\n access_type_filters: list[AccessType] = Field(default_factory=list)\n last_status_filters: list[IndexingStatus] = Field(default_factory=list)\n docs_count_operator: DocsCountOperator | None = None\n docs_count_value: int | None = None\n name_filter: str | None = None\n source_to_page: dict[DocumentSource, int] = Field(default_factory=dict)\n get_all_connectors: bool = False\n" }, { "path": "backend/onyx/server/documents/private_key_types.py", "content": "import base64\nfrom enum import Enum\nfrom typing import Protocol\n\nfrom fastapi import HTTPException\nfrom fastapi import UploadFile\n\nfrom onyx.server.documents.document_utils import validate_pkcs12_content\n\n\nclass ProcessPrivateKeyFileProtocol(Protocol):\n def __call__(self, file: UploadFile) -> str:\n \"\"\"\n Accepts a file-like object, validates the file (e.g., checks extension and content),\n and returns its contents as a base64-encoded string if valid.\n Raises an exception if validation fails.\n \"\"\"\n ...\n\n\nclass PrivateKeyFileTypes(Enum):\n SHAREPOINT_PFX_FILE = \"sharepoint_pfx_file\"\n\n\ndef process_sharepoint_private_key_file(file: UploadFile) -> str:\n \"\"\"\n Process and validate a private key file upload.\n\n Validates both the file extension and file content to ensure it's a valid PKCS#12 file.\n Content validation prevents attacks that rely on file extension spoofing.\n \"\"\"\n # First check file extension (basic filter)\n if not (file.filename and file.filename.lower().endswith(\".pfx\")):\n raise HTTPException(\n status_code=400, detail=\"Invalid file type. Only .pfx files are supported.\"\n )\n\n # Read file content for validation and processing\n private_key_bytes = file.file.read()\n\n # Validate file content to prevent extension spoofing attacks\n if not validate_pkcs12_content(private_key_bytes):\n raise HTTPException(\n status_code=400,\n detail=\"Invalid file content. The uploaded file does not appear to be a valid PKCS#12 (.pfx) file.\",\n )\n\n # Convert to base64 if validation passes\n pfx_64 = base64.b64encode(private_key_bytes).decode(\"ascii\")\n return pfx_64\n\n\nFILE_TYPE_TO_FILE_PROCESSOR: dict[\n PrivateKeyFileTypes, ProcessPrivateKeyFileProtocol\n] = {\n PrivateKeyFileTypes.SHAREPOINT_PFX_FILE: process_sharepoint_private_key_file,\n}\n" }, { "path": "backend/onyx/server/documents/standard_oauth.py", "content": "import json\nimport uuid\nfrom typing import Annotated\nfrom typing import cast\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi import Request\nfrom pydantic import BaseModel\nfrom pydantic import ValidationError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_user\nfrom onyx.configs.app_configs import WEB_DOMAIN\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import OAuthConnector\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import User\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.utils.logger import setup_logger\nfrom onyx.utils.subclasses import find_all_subclasses_in_package\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/connector/oauth\")\n\n_OAUTH_STATE_KEY_FMT = \"oauth_state:{state}\"\n_OAUTH_STATE_EXPIRATION_SECONDS = 10 * 60 # 10 minutes\n_DESIRED_RETURN_URL_KEY = \"desired_return_url\"\n_ADDITIONAL_KWARGS_KEY = \"additional_kwargs\"\n\n# Cache for OAuth connectors, populated at module load time\n_OAUTH_CONNECTORS: dict[DocumentSource, type[OAuthConnector]] = {}\n\n\ndef _discover_oauth_connectors() -> dict[DocumentSource, type[OAuthConnector]]:\n \"\"\"Walk through the connectors package to find all OAuthConnector implementations\"\"\"\n global _OAUTH_CONNECTORS\n if _OAUTH_CONNECTORS: # Return cached connectors if already discovered\n return _OAUTH_CONNECTORS\n\n # Import submodules using package-based discovery to avoid sys.path mutations\n oauth_connectors = find_all_subclasses_in_package(\n cast(type[OAuthConnector], OAuthConnector), \"onyx.connectors\"\n )\n\n _OAUTH_CONNECTORS = {cls.oauth_id(): cls for cls in oauth_connectors}\n return _OAUTH_CONNECTORS\n\n\n# Discover OAuth connectors at module load time\n_discover_oauth_connectors()\n\n\ndef _get_additional_kwargs(\n request: Request, connector_cls: type[OAuthConnector], args_to_ignore: list[str]\n) -> dict[str, str]:\n # get additional kwargs from request\n # e.g. anything except for desired_return_url\n additional_kwargs_dict = {\n k: v for k, v in request.query_params.items() if k not in args_to_ignore\n }\n try:\n # validate\n connector_cls.AdditionalOauthKwargs(**additional_kwargs_dict)\n except ValidationError:\n raise HTTPException(\n status_code=400,\n detail=(\n f\"Invalid additional kwargs. Got {additional_kwargs_dict}, expected \"\n f\"{connector_cls.AdditionalOauthKwargs.model_json_schema()}\"\n ),\n )\n\n return additional_kwargs_dict\n\n\nclass AuthorizeResponse(BaseModel):\n redirect_url: str\n\n\n@router.get(\"/authorize/{source}\")\ndef oauth_authorize(\n request: Request,\n source: DocumentSource,\n desired_return_url: Annotated[str | None, Query()] = None,\n _: User = Depends(current_user),\n) -> AuthorizeResponse:\n \"\"\"Initiates the OAuth flow by redirecting to the provider's auth page\"\"\"\n\n tenant_id = get_current_tenant_id()\n oauth_connectors = _discover_oauth_connectors()\n\n if source not in oauth_connectors:\n raise HTTPException(status_code=400, detail=f\"Unknown OAuth source: {source}\")\n\n connector_cls = oauth_connectors[source]\n base_url = WEB_DOMAIN\n\n # get additional kwargs from request\n # e.g. anything except for desired_return_url\n additional_kwargs = _get_additional_kwargs(\n request, connector_cls, [\"desired_return_url\"]\n )\n\n # store state in redis\n if not desired_return_url:\n desired_return_url = f\"{base_url}/admin/connectors/{source}?step=0\"\n redis_client = get_redis_client(tenant_id=tenant_id)\n state = str(uuid.uuid4())\n redis_client.set(\n _OAUTH_STATE_KEY_FMT.format(state=state),\n json.dumps(\n {\n _DESIRED_RETURN_URL_KEY: desired_return_url,\n _ADDITIONAL_KWARGS_KEY: additional_kwargs,\n }\n ),\n ex=_OAUTH_STATE_EXPIRATION_SECONDS,\n )\n\n return AuthorizeResponse(\n redirect_url=connector_cls.oauth_authorization_url(\n base_url, state, additional_kwargs\n )\n )\n\n\nclass CallbackResponse(BaseModel):\n redirect_url: str\n\n\n@router.get(\"/callback/{source}\")\ndef oauth_callback(\n source: DocumentSource,\n code: Annotated[str, Query()],\n state: Annotated[str, Query()],\n db_session: Session = Depends(get_session),\n user: User = Depends(current_user),\n) -> CallbackResponse:\n \"\"\"Handles the OAuth callback and exchanges the code for tokens\"\"\"\n oauth_connectors = _discover_oauth_connectors()\n\n if source not in oauth_connectors:\n raise HTTPException(status_code=400, detail=f\"Unknown OAuth source: {source}\")\n\n connector_cls = oauth_connectors[source]\n\n # get state from redis\n redis_client = get_redis_client()\n oauth_state_bytes = cast(\n bytes, redis_client.get(_OAUTH_STATE_KEY_FMT.format(state=state))\n )\n if not oauth_state_bytes:\n raise HTTPException(status_code=400, detail=\"Invalid OAuth state\")\n oauth_state = json.loads(oauth_state_bytes.decode(\"utf-8\"))\n\n desired_return_url = cast(str, oauth_state[_DESIRED_RETURN_URL_KEY])\n additional_kwargs = cast(dict[str, str], oauth_state[_ADDITIONAL_KWARGS_KEY])\n\n base_url = WEB_DOMAIN\n token_info = connector_cls.oauth_code_to_token(base_url, code, additional_kwargs)\n\n # Create a new credential with the token info\n credential_data = CredentialBase(\n credential_json=token_info,\n admin_public=True, # Or based on some logic/parameter\n source=source,\n name=f\"{source.title()} OAuth Credential\",\n )\n\n credential = create_credential(\n credential_data=credential_data,\n user=user,\n db_session=db_session,\n )\n\n # TODO: use a library for url handling\n sep = \"&\" if \"?\" in desired_return_url else \"?\"\n return CallbackResponse(\n redirect_url=f\"{desired_return_url}{sep}credentialId={credential.id}\"\n )\n\n\nclass OAuthAdditionalKwargDescription(BaseModel):\n name: str\n display_name: str\n description: str\n\n\nclass OAuthDetails(BaseModel):\n oauth_enabled: bool\n additional_kwargs: list[OAuthAdditionalKwargDescription]\n\n\n@router.get(\"/details/{source}\")\ndef oauth_details(\n source: DocumentSource,\n _: User = Depends(current_user),\n) -> OAuthDetails:\n oauth_connectors = _discover_oauth_connectors()\n\n if source not in oauth_connectors:\n return OAuthDetails(\n oauth_enabled=False,\n additional_kwargs=[],\n )\n\n connector_cls = oauth_connectors[source]\n\n additional_kwarg_descriptions = []\n for key, value in connector_cls.AdditionalOauthKwargs.model_json_schema()[\n \"properties\"\n ].items():\n additional_kwarg_descriptions.append(\n OAuthAdditionalKwargDescription(\n name=key,\n display_name=value.get(\"title\", key),\n description=value.get(\"description\", \"\"),\n )\n )\n\n return OAuthDetails(\n oauth_enabled=True,\n additional_kwargs=additional_kwarg_descriptions,\n )\n" }, { "path": "backend/onyx/server/evals/__init__.py", "content": "" }, { "path": "backend/onyx/server/evals/models.py", "content": "from pydantic import BaseModel\n\n\nclass EvalRunAck(BaseModel):\n \"\"\"Response model for evaluation runs\"\"\"\n\n success: bool\n" }, { "path": "backend/onyx/server/features/__init__.py", "content": "" }, { "path": "backend/onyx/server/features/build/.gitignore", "content": "sandbox/kubernetes/docker/templates/venv/**\nsandbox/kubernetes/docker/demo_data/**\n" }, { "path": "backend/onyx/server/features/build/AGENTS.template.md", "content": "# AGENTS.md\n\nYou are an AI agent powering **Onyx Craft**. You create interactive web applications, dashboards, and documents from company knowledge. You run in a secure sandbox with access to the user's knowledge sources. The knowledge sources you have are organization context like meeting notes, emails, slack messages, and other organizational data that you must use to answer your question.\n\n{{USER_CONTEXT}}\n\n## Configuration\n\n- **LLM**: {{LLM_PROVIDER_NAME}} / {{LLM_MODEL_NAME}}\n- **Next.js**: Running on port {{NEXTJS_PORT}} (already started — do NOT run `npm run dev`)\n {{DISABLED_TOOLS_SECTION}}\n\n## Environment\n\nEphemeral VM with Python 3.11 and Node v22. Virtual environment at `.venv/` includes numpy, pandas, matplotlib, scipy.\n\nInstall packages: `pip install ` or `npm install ` (from `outputs/web`).\n\n{{ORG_INFO_SECTION}}\n\n## Skills\n\n{{AVAILABLE_SKILLS_SECTION}}\n\nRead the relevant SKILL.md before starting work that the skill covers.\n\n## Recommended Task Approach Methodology\n\nWhen presented with a task, you typically:\n\n1. Analyze the request to understand what's being asked\n2. Break down complex problems into manageable steps and sub-questions\n3. Use appropriate tools and methods to address each step\n4. Provide clear communication throughout the process\n5. Deliver results in a helpful and organized manner\n\nFollow this two-step pattern for most tasks:\n\n### Step 1: Information Retrieval\n\n1. **Search** knowledge sources using `find`, `grep`, or direct file reads. Start your search at the root of the `files/` directory\nto get a general grasp of what subdirectories to further explore, especially when looking for a person. their name may be a proper noun\nor strictly lowercase.\n2. **Extract** relevant data from JSON documents\n3. **Summarize** key findings before proceeding\n\n**Tip**: Use `find`, `grep`, or `glob` to search files directly rather than navigating directories one at a time.\n\n### Step 2: Output Generation\n\n1. **Choose format**: Web app for interactive/visual, Markdown for reports, or direct response for quick answers\n2. **Build** the output using retrieved information\n3. **Verify** the output renders correctly and includes accurate data\n\n## Behavior Guidelines\n\n- **Accuracy**: Do not make any assumptions about the user. Any conclusions you reach must be supported by the provided data.\n\n- **Completeness**: For any tasks requiring data from the knowledge sources, you should make sure to look at ALL sources that may be relevant to the user's questions and use that in your final response. Make sure you check Google Drive if applicable\n - **Explicitly state** which sources were checked and which had no relevant data\n - **Search ALL knowledge sources** for the person's name/email, not just the obvious ones when answering questions about a person's activites.\n\n- **Task Management**: For any non-trivial task involving multiple steps, you should organize your work and track progress. This helps users understand what you're doing and ensures nothing is missed.\n\n- **Verification**: For important work, include a verification step to double-check your output. This could involve testing functionality, reviewing for accuracy, or validating against requirements.\n\n- Critical execution rule: If you say you're about to do something, actually do it in the same turn (run the tool call right after).\n\n- Check off completed TODOs before reporting progress.\n\n- Your main goal is to follow the USER's instructions at each message\n\n- Don't mention tool names to the user; describe actions naturally.\n\n## Knowledge Sources\n\nThe `files/` directory contains JSON documents from various knowledge sources. Here's what's available:\n\n{{KNOWLEDGE_SOURCES_SECTION}}\n\n### Document Format\n\nFiles are JSON with: `title`, `source`, `metadata`, `sections[{text, link}]`.\n\n**Important**: The `files/` directory is read-only. Do NOT attempt to write to it.\n\n## Outputs\n\nAll outputs go in the `outputs/` directory.\n\n| Format | Use For |\n| ------------ | ---------------------------------------- |\n| **Web App** | Interactive dashboards, data exploration |\n| **Markdown** | Reports, analyses, documentation |\n| **Response** | Quick answers, lookups |\n\nYou can also generate other output formats if you think they more directly answer the user's question\n\n### Web Apps\n\nUse `outputs/web` with Next.js 16.1.1, React v19, Tailwind, Recharts, shadcn/ui.\n\n\n\n### Markdown\n\nSave to `outputs/markdown/*.md`. Use clear headings and tables.\n\n## Questions to Ask\n\n- Did you check all relevant sources that could be useful in addressing the user's question?\n- Did you generate the correct output format that the user requested?\n- Did you answer the user's question thoroughly?\n" }, { "path": "backend/onyx/server/features/build/__init__.py", "content": "# Build feature module\n" }, { "path": "backend/onyx/server/features/build/api/api.py", "content": "import re\nfrom collections.abc import Iterator\nfrom pathlib import Path\nfrom uuid import UUID\n\nimport httpx\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi import Request\nfrom fastapi import Response\nfrom fastapi.responses import RedirectResponse\nfrom fastapi.responses import StreamingResponse\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_user\nfrom onyx.auth.users import optional_user\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs_for_user\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import ProcessingMode\nfrom onyx.db.enums import SharingScope\nfrom onyx.db.index_attempt import get_latest_index_attempt_for_cc_pair_id\nfrom onyx.db.models import BuildSession\nfrom onyx.db.models import User\nfrom onyx.server.features.build.api.messages_api import router as messages_router\nfrom onyx.server.features.build.api.models import BuildConnectorInfo\nfrom onyx.server.features.build.api.models import BuildConnectorListResponse\nfrom onyx.server.features.build.api.models import BuildConnectorStatus\nfrom onyx.server.features.build.api.models import RateLimitResponse\nfrom onyx.server.features.build.api.rate_limit import get_user_rate_limit_status\nfrom onyx.server.features.build.api.sessions_api import router as sessions_router\nfrom onyx.server.features.build.api.user_library import router as user_library_router\nfrom onyx.server.features.build.db.sandbox import get_sandbox_by_user_id\nfrom onyx.server.features.build.sandbox import get_sandbox_manager\nfrom onyx.server.features.build.session.manager import SessionManager\nfrom onyx.server.features.build.utils import is_onyx_craft_enabled\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n_TEMPLATES_DIR = Path(__file__).parent / \"templates\"\n_WEBAPP_HMR_FIXER_TEMPLATE = (_TEMPLATES_DIR / \"webapp_hmr_fixer.js\").read_text()\n\n\ndef require_onyx_craft_enabled(user: User = Depends(current_user)) -> User:\n \"\"\"\n Dependency that checks if Onyx Craft is enabled for the user.\n Raises HTTP 403 if Onyx Craft is disabled via feature flag.\n \"\"\"\n if not is_onyx_craft_enabled(user):\n raise HTTPException(\n status_code=403,\n detail=\"Onyx Craft is not available\",\n )\n return user\n\n\nrouter = APIRouter(prefix=\"/build\", dependencies=[Depends(require_onyx_craft_enabled)])\n\n# Include sub-routers for sessions, messages, and user library\nrouter.include_router(sessions_router, tags=[\"build\"])\nrouter.include_router(messages_router, tags=[\"build\"])\nrouter.include_router(user_library_router, tags=[\"build\"])\n\n\n# -----------------------------------------------------------------------------\n# Rate Limiting\n# -----------------------------------------------------------------------------\n\n\n@router.get(\"/limit\", response_model=RateLimitResponse)\ndef get_rate_limit(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> RateLimitResponse:\n \"\"\"Get rate limit information for the current user.\"\"\"\n return get_user_rate_limit_status(user, db_session)\n\n\n# -----------------------------------------------------------------------------\n# Build Connectors\n# -----------------------------------------------------------------------------\n\n\n@router.get(\"/connectors\", response_model=BuildConnectorListResponse)\ndef get_build_connectors(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> BuildConnectorListResponse:\n \"\"\"Get all connectors for the build admin panel.\n\n Returns connector-credential pairs with simplified status information.\n On the build configure page, all users (including admins) only see connectors\n they own/created. Users can create new connectors if they don't have one of a type.\n \"\"\"\n # Fetch both FILE_SYSTEM (standard connectors) and RAW_BINARY (User Library) connectors\n file_system_cc_pairs = get_connector_credential_pairs_for_user(\n db_session=db_session,\n user=user,\n get_editable=False,\n eager_load_connector=True,\n eager_load_credential=True,\n processing_mode=ProcessingMode.FILE_SYSTEM,\n )\n raw_binary_cc_pairs = get_connector_credential_pairs_for_user(\n db_session=db_session,\n user=user,\n get_editable=False,\n eager_load_connector=True,\n eager_load_credential=True,\n processing_mode=ProcessingMode.RAW_BINARY,\n )\n cc_pairs = file_system_cc_pairs + raw_binary_cc_pairs\n\n # Filter to only show connectors created by the current user\n # All users (including admins) must create their own connectors on the build configure page\n if user:\n cc_pairs = [cc_pair for cc_pair in cc_pairs if cc_pair.creator_id == user.id]\n\n connectors: list[BuildConnectorInfo] = []\n for cc_pair in cc_pairs:\n # Skip ingestion API connectors and default pairs\n if cc_pair.connector.source == DocumentSource.INGESTION_API:\n continue\n if cc_pair.name == \"DefaultCCPair\":\n continue\n\n # Determine status\n error_message: str | None = None\n has_ever_succeeded = cc_pair.last_successful_index_time is not None\n\n if cc_pair.status == ConnectorCredentialPairStatus.DELETING:\n status = BuildConnectorStatus.DELETING\n elif cc_pair.status == ConnectorCredentialPairStatus.INVALID:\n # If connector has succeeded before but credentials are now invalid,\n # show as connected_with_errors so user can still disable demo data\n if has_ever_succeeded:\n status = BuildConnectorStatus.CONNECTED_WITH_ERRORS\n error_message = \"Connector credentials are invalid\"\n else:\n status = BuildConnectorStatus.ERROR\n error_message = \"Connector credentials are invalid\"\n else:\n # Check latest index attempt for errors\n latest_attempt = get_latest_index_attempt_for_cc_pair_id(\n db_session=db_session,\n connector_credential_pair_id=cc_pair.id,\n secondary_index=False,\n only_finished=True,\n )\n\n if latest_attempt and latest_attempt.status == IndexingStatus.FAILED:\n # If connector has succeeded before but latest attempt failed,\n # show as connected_with_errors\n if has_ever_succeeded:\n status = BuildConnectorStatus.CONNECTED_WITH_ERRORS\n else:\n status = BuildConnectorStatus.ERROR\n error_message = latest_attempt.error_msg\n elif (\n latest_attempt\n and latest_attempt.status == IndexingStatus.COMPLETED_WITH_ERRORS\n ):\n # Completed with errors - if it has succeeded before, show as connected_with_errors\n if has_ever_succeeded:\n status = BuildConnectorStatus.CONNECTED_WITH_ERRORS\n else:\n status = BuildConnectorStatus.ERROR\n error_message = \"Indexing completed with errors\"\n elif cc_pair.status == ConnectorCredentialPairStatus.PAUSED:\n status = BuildConnectorStatus.CONNECTED\n elif cc_pair.last_successful_index_time is None:\n # Never successfully indexed - check if currently indexing\n # First check cc_pair status for scheduled/initial indexing\n if cc_pair.status in (\n ConnectorCredentialPairStatus.SCHEDULED,\n ConnectorCredentialPairStatus.INITIAL_INDEXING,\n ):\n status = BuildConnectorStatus.INDEXING\n else:\n in_progress_attempt = get_latest_index_attempt_for_cc_pair_id(\n db_session=db_session,\n connector_credential_pair_id=cc_pair.id,\n secondary_index=False,\n only_finished=False,\n )\n if (\n in_progress_attempt\n and in_progress_attempt.status == IndexingStatus.IN_PROGRESS\n ):\n status = BuildConnectorStatus.INDEXING\n elif (\n in_progress_attempt\n and in_progress_attempt.status == IndexingStatus.NOT_STARTED\n ):\n status = BuildConnectorStatus.INDEXING\n else:\n # Has a finished attempt but never succeeded - likely error\n status = BuildConnectorStatus.ERROR\n error_message = (\n latest_attempt.error_msg\n if latest_attempt\n else \"Initial indexing failed\"\n )\n else:\n status = BuildConnectorStatus.CONNECTED\n\n connectors.append(\n BuildConnectorInfo(\n cc_pair_id=cc_pair.id,\n connector_id=cc_pair.connector.id,\n credential_id=cc_pair.credential.id,\n source=cc_pair.connector.source.value,\n name=cc_pair.name or cc_pair.connector.name or \"Unnamed\",\n status=status,\n docs_indexed=0, # Would need to query for this\n last_indexed=cc_pair.last_successful_index_time,\n error_message=error_message,\n )\n )\n\n return BuildConnectorListResponse(connectors=connectors)\n\n\n# Headers to skip when proxying.\n# Hop-by-hop headers must not be forwarded, and set-cookie is stripped to\n# prevent LLM-generated apps from setting cookies on the parent Onyx domain.\nEXCLUDED_HEADERS = {\n \"content-encoding\",\n \"content-length\",\n \"transfer-encoding\",\n \"connection\",\n \"set-cookie\",\n}\n\n\ndef _stream_response(response: httpx.Response) -> Iterator[bytes]:\n \"\"\"Stream the response content in chunks.\"\"\"\n for chunk in response.iter_bytes(chunk_size=8192):\n yield chunk\n\n\ndef _inject_hmr_fixer(content: bytes, session_id: str) -> bytes:\n \"\"\"Inject a script that stubs root-scoped Next HMR websocket connections.\"\"\"\n base = f\"/api/build/sessions/{session_id}/webapp\"\n script = f\"\"\n text = content.decode(\"utf-8\")\n text = re.sub(\n r\"(]*>)\",\n lambda m: m.group(0) + script,\n text,\n count=1,\n flags=re.IGNORECASE,\n )\n return text.encode(\"utf-8\")\n\n\ndef _rewrite_asset_paths(content: bytes, session_id: str) -> bytes:\n \"\"\"Rewrite Next.js asset paths to go through the proxy.\"\"\"\n webapp_base_path = f\"/api/build/sessions/{session_id}/webapp\"\n escaped_webapp_base_path = webapp_base_path.replace(\"/\", r\"\\/\")\n hmr_paths = (\"/_next/webpack-hmr\", \"/_next/hmr\")\n\n text = content.decode(\"utf-8\")\n # Anchor on delimiter so already-prefixed URLs (from assetPrefix) aren't double-rewritten.\n for delim in ('\"', \"'\", \"(\"):\n text = text.replace(f\"{delim}/_next/\", f\"{delim}{webapp_base_path}/_next/\")\n text = re.sub(\n rf\"{re.escape(delim)}https?://[^/\\\"')]+/_next/\",\n f\"{delim}{webapp_base_path}/_next/\",\n text,\n )\n text = re.sub(\n rf\"{re.escape(delim)}wss?://[^/\\\"')]+/_next/\",\n f\"{delim}{webapp_base_path}/_next/\",\n text,\n )\n text = text.replace(r\"\\/_next\\/\", rf\"{escaped_webapp_base_path}\\/_next\\/\")\n text = re.sub(\n r\"https?:\\\\\\/\\\\\\/[^\\\"']+?\\\\\\/_next\\\\\\/\",\n rf\"{escaped_webapp_base_path}\\/_next\\/\",\n text,\n )\n text = re.sub(\n r\"wss?:\\\\\\/\\\\\\/[^\\\"']+?\\\\\\/_next\\\\\\/\",\n rf\"{escaped_webapp_base_path}\\/_next\\/\",\n text,\n )\n for hmr_path in hmr_paths:\n escaped_hmr_path = hmr_path.replace(\"/\", r\"\\/\")\n text = text.replace(\n f\"{webapp_base_path}{hmr_path}\",\n hmr_path,\n )\n text = text.replace(\n f\"{escaped_webapp_base_path}{escaped_hmr_path}\",\n escaped_hmr_path,\n )\n text = re.sub(\n r'\"(/(?:[a-zA-Z0-9_-]+/)*[a-zA-Z0-9_-]+\\.json)\"',\n f'\"{webapp_base_path}\\\\1\"',\n text,\n )\n text = re.sub(\n r\"'(/(?:[a-zA-Z0-9_-]+/)*[a-zA-Z0-9_-]+\\.json)'\",\n f\"'{webapp_base_path}\\\\1'\",\n text,\n )\n text = text.replace('\"/favicon.ico', f'\"{webapp_base_path}/favicon.ico')\n return text.encode(\"utf-8\")\n\n\ndef _rewrite_proxy_response_headers(\n headers: dict[str, str], session_id: str\n) -> dict[str, str]:\n \"\"\"Rewrite response headers that can leak root-scoped asset URLs.\"\"\"\n link = headers.get(\"link\")\n if link:\n webapp_base_path = f\"/api/build/sessions/{session_id}/webapp\"\n rewritten_link = re.sub(\n r\"]+/_next/\",\n f\"<{webapp_base_path}/_next/\",\n link,\n )\n rewritten_link = rewritten_link.replace(\n \" str:\n \"\"\"Get the internal URL for a session's Next.js server.\n\n Uses the sandbox manager to get the correct URL for both local and\n Kubernetes environments.\n\n Args:\n session_id: The build session ID\n db_session: Database session\n\n Returns:\n Internal URL to proxy requests to\n\n Raises:\n HTTPException: If session not found, port not allocated, or sandbox not found\n \"\"\"\n\n session = db_session.get(BuildSession, session_id)\n if not session:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n if session.nextjs_port is None:\n raise HTTPException(status_code=503, detail=\"Session port not allocated\")\n if session.user_id is None:\n raise HTTPException(status_code=404, detail=\"User not found\")\n\n sandbox = get_sandbox_by_user_id(db_session, session.user_id)\n if sandbox is None:\n raise HTTPException(status_code=404, detail=\"Sandbox not found\")\n\n sandbox_manager = get_sandbox_manager()\n return sandbox_manager.get_webapp_url(sandbox.id, session.nextjs_port)\n\n\ndef _proxy_request(\n path: str, request: Request, session_id: UUID, db_session: Session\n) -> StreamingResponse | Response:\n \"\"\"Proxy a request to the sandbox's Next.js server.\"\"\"\n base_url = _get_sandbox_url(session_id, db_session)\n\n # Build the target URL\n target_url = f\"{base_url}/{path.lstrip('/')}\"\n\n # Include query params if present\n if request.query_params:\n target_url = f\"{target_url}?{request.query_params}\"\n\n logger.debug(f\"Proxying request to: {target_url}\")\n\n try:\n # Make the request to the target URL\n with httpx.Client(timeout=30.0, follow_redirects=True) as client:\n response = client.get(\n target_url,\n headers={\n key: value\n for key, value in request.headers.items()\n if key.lower() not in (\"host\", \"content-length\")\n },\n )\n\n # Build response headers, excluding hop-by-hop headers\n response_headers = {\n key: value\n for key, value in response.headers.items()\n if key.lower() not in EXCLUDED_HEADERS\n }\n response_headers = _rewrite_proxy_response_headers(\n response_headers, str(session_id)\n )\n\n content_type = response.headers.get(\"content-type\", \"\")\n\n # For HTML/CSS/JS responses, rewrite asset paths\n if any(ct in content_type for ct in REWRITABLE_CONTENT_TYPES):\n content = _rewrite_asset_paths(response.content, str(session_id))\n if \"text/html\" in content_type:\n content = _inject_hmr_fixer(content, str(session_id))\n return Response(\n content=content,\n status_code=response.status_code,\n headers=response_headers,\n media_type=content_type,\n )\n\n return StreamingResponse(\n content=_stream_response(response),\n status_code=response.status_code,\n headers=response_headers,\n media_type=content_type or None,\n )\n\n except httpx.TimeoutException:\n logger.error(f\"Timeout while proxying request to {target_url}\")\n raise HTTPException(status_code=504, detail=\"Gateway timeout\")\n except httpx.RequestError as e:\n logger.error(f\"Error proxying request to {target_url}: {e}\")\n raise HTTPException(status_code=502, detail=\"Bad gateway\")\n\n\ndef _check_webapp_access(\n session_id: UUID, user: User | None, db_session: Session\n) -> BuildSession:\n \"\"\"Check if user can access a session's webapp.\n\n - public_global: accessible by anyone (no auth required)\n - public_org: accessible by any authenticated user\n - private: only accessible by the session owner\n \"\"\"\n session = db_session.get(BuildSession, session_id)\n if not session:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n if session.sharing_scope == SharingScope.PUBLIC_GLOBAL:\n return session\n if user is None:\n raise HTTPException(status_code=401, detail=\"Authentication required\")\n if session.sharing_scope == SharingScope.PRIVATE and session.user_id != user.id:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n return session\n\n\n_OFFLINE_HTML_PATH = _TEMPLATES_DIR / \"webapp_offline.html\"\n\n\ndef _offline_html_response() -> Response:\n \"\"\"Return a branded Craft HTML page when the sandbox is not reachable.\n\n Design mirrors the default Craft web template (outputs/web/app/page.tsx):\n terminal window aesthetic with Minecraft-themed typing animation.\n\n \"\"\"\n html = _OFFLINE_HTML_PATH.read_text()\n return Response(content=html, status_code=503, media_type=\"text/html\")\n\n\n# Public router for webapp proxy — no authentication required\n# (access controlled per-session via sharing_scope)\npublic_build_router = APIRouter(prefix=\"/build\")\n\n\n@public_build_router.get(\"/sessions/{session_id}/webapp\", response_model=None)\n@public_build_router.get(\n \"/sessions/{session_id}/webapp/{path:path}\", response_model=None\n)\ndef get_webapp(\n session_id: UUID,\n request: Request,\n path: str = \"\",\n user: User | None = Depends(optional_user),\n db_session: Session = Depends(get_session),\n) -> StreamingResponse | Response:\n \"\"\"Proxy the webapp for a specific session (root and subpaths).\n\n Accessible without authentication when sharing_scope is public_global.\n Returns a friendly offline page when the sandbox is not running.\n \"\"\"\n try:\n _check_webapp_access(session_id, user, db_session)\n except HTTPException as e:\n if e.status_code == 401:\n return RedirectResponse(url=\"/auth/login\", status_code=302)\n raise\n try:\n return _proxy_request(path, request, session_id, db_session)\n except HTTPException as e:\n if e.status_code in (502, 503, 504):\n return _offline_html_response()\n raise\n\n\n# =============================================================================\n# Sandbox Management Endpoints\n# =============================================================================\n\n\n@router.post(\"/sandbox/reset\", response_model=None)\ndef reset_sandbox(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"Reset the user's sandbox by terminating it and cleaning up all sessions.\n\n This endpoint terminates the user's shared sandbox container/pod and\n cleans up all session workspaces. Useful for \"start fresh\" functionality.\n\n After calling this endpoint, the next session creation will provision a\n new sandbox.\n \"\"\"\n session_manager = SessionManager(db_session)\n\n try:\n success = session_manager.terminate_user_sandbox(user.id)\n if not success:\n raise HTTPException(\n status_code=404,\n detail=\"No sandbox found for user\",\n )\n db_session.commit()\n except HTTPException:\n raise\n except Exception as e:\n db_session.rollback()\n logger.error(f\"Failed to reset sandbox for user {user.id}: {e}\")\n raise HTTPException(\n status_code=500,\n detail=f\"Failed to reset sandbox: {e}\",\n )\n\n return Response(status_code=204)\n" }, { "path": "backend/onyx/server/features/build/api/messages_api.py", "content": "\"\"\"API endpoints for Build Mode message management.\"\"\"\n\nfrom collections.abc import Generator\nfrom uuid import UUID\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import HTTPException\nfrom fastapi.responses import StreamingResponse\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_user\nfrom onyx.configs.constants import PUBLIC_API_TAGS\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import User\nfrom onyx.server.features.build.api.models import MessageListResponse\nfrom onyx.server.features.build.api.models import MessageRequest\nfrom onyx.server.features.build.api.models import MessageResponse\nfrom onyx.server.features.build.db.sandbox import get_sandbox_by_user_id\nfrom onyx.server.features.build.db.sandbox import update_sandbox_heartbeat\nfrom onyx.server.features.build.session.manager import RateLimitError\nfrom onyx.server.features.build.session.manager import SessionManager\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nrouter = APIRouter()\n\n\ndef check_build_rate_limits(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> None:\n \"\"\"\n Dependency to check build mode rate limits before processing the request.\n\n Raises HTTPException(429) if rate limit is exceeded.\n Follows the same pattern as chat's check_token_rate_limits.\n \"\"\"\n session_manager = SessionManager(db_session)\n\n try:\n session_manager.check_rate_limit(user)\n except RateLimitError as e:\n raise HTTPException(\n status_code=429,\n detail=str(e),\n )\n\n\n@router.get(\"/sessions/{session_id}/messages\", tags=PUBLIC_API_TAGS)\ndef list_messages(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> MessageListResponse:\n \"\"\"Get all messages for a build session.\"\"\"\n session_manager = SessionManager(db_session)\n\n messages = session_manager.list_messages(session_id, user.id)\n\n if messages is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n return MessageListResponse(\n messages=[MessageResponse.from_model(msg) for msg in messages]\n )\n\n\n@router.post(\"/sessions/{session_id}/send-message\", tags=PUBLIC_API_TAGS)\ndef send_message(\n session_id: UUID,\n request: MessageRequest,\n user: User = Depends(current_user),\n _rate_limit_check: None = Depends(check_build_rate_limits),\n) -> StreamingResponse:\n \"\"\"\n Send a message to the CLI agent and stream the response.\n\n Enforces rate limiting before executing the agent (via dependency).\n Returns a Server-Sent Events (SSE) stream with the agent's response.\n\n Follows the same pattern as /chat/send-chat-message for consistency.\n \"\"\"\n\n def stream_generator() -> Generator[str, None, None]:\n \"\"\"Stream generator that manages its own database session.\n\n This is necessary because StreamingResponse consumes the generator\n AFTER the endpoint returns, at which point FastAPI's dependency-injected\n db_session has already been closed. By creating a new session inside\n the generator, we ensure the session remains open for the entire\n streaming duration.\n \"\"\"\n # Capture user info needed for streaming (user object may not be available\n # after the endpoint returns due to dependency cleanup)\n user_id = user.id\n message_content = request.content\n\n with get_session_with_current_tenant() as db_session:\n # Update sandbox heartbeat - this is the only place we track activity\n # for determining when a sandbox should be put to sleep\n sandbox = get_sandbox_by_user_id(db_session, user.id)\n if sandbox and sandbox.status.is_active():\n update_sandbox_heartbeat(db_session, sandbox.id)\n\n session_manager = SessionManager(db_session)\n yield from session_manager.send_message(\n session_id, user_id, message_content\n )\n\n # Stream the CLI agent's response\n return StreamingResponse(\n stream_generator(),\n media_type=\"text/event-stream\",\n headers={\n \"Cache-Control\": \"no-cache\",\n \"Connection\": \"keep-alive\",\n \"X-Accel-Buffering\": \"no\", # Disable nginx buffering\n },\n )\n" }, { "path": "backend/onyx/server/features/build/api/models.py", "content": "from datetime import datetime\nfrom enum import Enum\nfrom typing import Any\nfrom typing import TYPE_CHECKING\nfrom typing import Union\n\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.enums import ArtifactType\nfrom onyx.db.enums import BuildSessionStatus\nfrom onyx.db.enums import SandboxStatus\nfrom onyx.db.enums import SharingScope\nfrom onyx.server.features.build.sandbox.models import (\n FilesystemEntry as FileSystemEntry,\n)\n\nif TYPE_CHECKING:\n from onyx.db.models import Sandbox\n from onyx.db.models import BuildSession\n\n\n# ===== Session Models =====\nclass SessionCreateRequest(BaseModel):\n \"\"\"Request to create a new build session.\"\"\"\n\n name: str | None = None # Optional session name\n demo_data_enabled: bool = True # Whether to enable demo org_info data in sandbox\n user_work_area: str | None = None # User's work area (e.g., \"engineering\")\n user_level: str | None = None # User's level (e.g., \"ic\", \"manager\")\n # LLM selection from user's cookie\n llm_provider_type: str | None = None # Provider type (e.g., \"anthropic\", \"openai\")\n llm_model_name: str | None = None # Model name (e.g., \"claude-opus-4-5\")\n\n\nclass SessionUpdateRequest(BaseModel):\n \"\"\"Request to update a build session.\n\n If name is None, the session name will be auto-generated using LLM.\n \"\"\"\n\n name: str | None = None\n\n\nclass SessionNameGenerateResponse(BaseModel):\n \"\"\"Response containing a generated session name.\"\"\"\n\n name: str\n\n\nclass SandboxResponse(BaseModel):\n \"\"\"Sandbox metadata in session response.\"\"\"\n\n id: str\n status: SandboxStatus\n container_id: str | None\n created_at: datetime\n last_heartbeat: datetime | None\n\n @classmethod\n def from_model(cls, sandbox: Any) -> \"SandboxResponse\":\n \"\"\"Convert Sandbox ORM model to response.\"\"\"\n return cls(\n id=str(sandbox.id),\n status=sandbox.status,\n container_id=sandbox.container_id,\n created_at=sandbox.created_at,\n last_heartbeat=sandbox.last_heartbeat,\n )\n\n\nclass ArtifactResponse(BaseModel):\n \"\"\"Artifact metadata in session response.\"\"\"\n\n id: str\n session_id: str\n type: ArtifactType\n name: str\n path: str\n preview_url: str | None\n created_at: datetime\n updated_at: datetime\n\n @classmethod\n def from_model(cls, artifact: Any) -> \"ArtifactResponse\":\n \"\"\"Convert Artifact ORM model to response.\"\"\"\n return cls(\n id=str(artifact.id),\n session_id=str(artifact.session_id),\n type=artifact.type,\n name=artifact.name,\n path=artifact.path,\n preview_url=getattr(artifact, \"preview_url\", None),\n created_at=artifact.created_at,\n updated_at=artifact.updated_at,\n )\n\n\nclass SessionResponse(BaseModel):\n \"\"\"Response containing session details.\"\"\"\n\n id: str\n user_id: str | None\n name: str | None\n status: BuildSessionStatus\n created_at: datetime\n last_activity_at: datetime\n nextjs_port: int | None\n sandbox: SandboxResponse | None\n artifacts: list[ArtifactResponse]\n sharing_scope: SharingScope\n\n @classmethod\n def from_model(\n cls, session: \"BuildSession\", sandbox: Union[\"Sandbox\", None] = None\n ) -> \"SessionResponse\":\n \"\"\"Convert BuildSession ORM model to response.\n\n Args:\n session: BuildSession ORM model\n sandbox: Optional Sandbox ORM model. Since sandboxes are now user-owned\n (not session-owned), the sandbox must be passed separately.\n \"\"\"\n return cls(\n id=str(session.id),\n user_id=str(session.user_id) if session.user_id else None,\n name=session.name,\n status=session.status,\n created_at=session.created_at,\n last_activity_at=session.last_activity_at,\n nextjs_port=session.nextjs_port,\n sandbox=(SandboxResponse.from_model(sandbox) if sandbox else None),\n artifacts=[ArtifactResponse.from_model(a) for a in session.artifacts],\n sharing_scope=session.sharing_scope,\n )\n\n\nclass DetailedSessionResponse(SessionResponse):\n \"\"\"Extended session response with sandbox state details.\n\n Used for single-session endpoints where we compute expensive fields\n like session_loaded_in_sandbox.\n \"\"\"\n\n session_loaded_in_sandbox: bool\n\n @classmethod\n def from_session_response(\n cls,\n base: SessionResponse,\n session_loaded_in_sandbox: bool,\n ) -> \"DetailedSessionResponse\":\n return cls(\n **base.model_dump(),\n session_loaded_in_sandbox=session_loaded_in_sandbox,\n )\n\n\nclass SessionListResponse(BaseModel):\n \"\"\"Response containing list of sessions.\"\"\"\n\n sessions: list[SessionResponse]\n\n\nclass SetSessionSharingRequest(BaseModel):\n \"\"\"Request to set the sharing scope of a session.\"\"\"\n\n sharing_scope: SharingScope\n\n\nclass SetSessionSharingResponse(BaseModel):\n \"\"\"Response after setting session sharing scope.\"\"\"\n\n session_id: str\n sharing_scope: SharingScope\n\n\n# ===== Message Models =====\nclass MessageRequest(BaseModel):\n \"\"\"Request to send a message to the CLI agent.\"\"\"\n\n content: str\n\n\nclass MessageResponse(BaseModel):\n \"\"\"Response containing message details.\n\n All message data is stored in message_metadata as JSON (the raw ACP packet).\n The turn_index groups all assistant responses under the user prompt they respond to.\n\n Packet types in message_metadata:\n - user_message: {type: \"user_message\", content: {...}}\n - agent_message: {type: \"agent_message\", content: {...}}\n - agent_thought: {type: \"agent_thought\", content: {...}}\n - tool_call_progress: {type: \"tool_call_progress\", status: \"completed\", ...}\n - agent_plan_update: {type: \"agent_plan_update\", entries: [...]}\n \"\"\"\n\n id: str\n session_id: str\n turn_index: int\n type: MessageType\n message_metadata: dict[str, Any]\n created_at: datetime\n\n @classmethod\n def from_model(cls, message: Any) -> \"MessageResponse\":\n \"\"\"Convert BuildMessage ORM model to response.\"\"\"\n return cls(\n id=str(message.id),\n session_id=str(message.session_id),\n turn_index=message.turn_index,\n type=message.type,\n message_metadata=message.message_metadata,\n created_at=message.created_at,\n )\n\n\nclass MessageListResponse(BaseModel):\n \"\"\"Response containing list of messages.\"\"\"\n\n messages: list[MessageResponse]\n\n\n# ===== Legacy Models (for compatibility with other code) =====\nclass CreateSessionRequest(BaseModel):\n task: str\n available_sources: list[str] | None = None\n\n\nclass CreateSessionResponse(BaseModel):\n session_id: str\n\n\nclass ExecuteRequest(BaseModel):\n task: str\n context: str | None = None\n\n\nclass ArtifactInfo(BaseModel):\n artifact_type: str # \"webapp\", \"file\", \"markdown\", \"image\"\n path: str\n filename: str\n mime_type: str | None = None\n\n\nclass SessionStatus(BaseModel):\n session_id: str\n status: str # \"idle\", \"running\", \"completed\", \"failed\"\n webapp_url: str | None = None\n\n\nclass DirectoryListing(BaseModel):\n path: str # Current directory path\n entries: list[FileSystemEntry] # Contents\n\n\nclass WebappInfo(BaseModel):\n has_webapp: bool # Whether a webapp exists in outputs/web\n webapp_url: str | None # URL to access the webapp (e.g., http://localhost:3015)\n status: str # Sandbox status (running, terminated, etc.)\n ready: bool # Whether the NextJS dev server is actually responding\n sharing_scope: SharingScope\n\n\n# ===== File Upload Models =====\nclass UploadResponse(BaseModel):\n \"\"\"Response after successful file upload.\"\"\"\n\n filename: str # Sanitized filename\n path: str # Relative path in sandbox (e.g., \"attachments/doc.pdf\")\n size_bytes: int # File size in bytes\n\n\n# ===== Rate Limit Models =====\nclass RateLimitResponse(BaseModel):\n \"\"\"Rate limit information.\"\"\"\n\n is_limited: bool\n limit_type: str # \"weekly\" or \"total\"\n messages_used: int\n limit: int\n reset_timestamp: str | None = None\n\n\n# ===== Pre-Provisioned Session Check Models =====\nclass PreProvisionedCheckResponse(BaseModel):\n \"\"\"Response for checking if a pre-provisioned session is still valid (empty).\"\"\"\n\n valid: bool # True if session exists and has no messages\n session_id: str | None = None # Session ID if valid, None otherwise\n\n\n# ===== Build Connector Models =====\nclass BuildConnectorStatus(str, Enum):\n \"\"\"Status of a build connector.\"\"\"\n\n NOT_CONNECTED = \"not_connected\"\n CONNECTED = \"connected\"\n CONNECTED_WITH_ERRORS = \"connected_with_errors\"\n INDEXING = \"indexing\"\n ERROR = \"error\"\n DELETING = \"deleting\"\n\n\nclass BuildConnectorInfo(BaseModel):\n \"\"\"Simplified connector info for build admin panel.\"\"\"\n\n cc_pair_id: int\n connector_id: int\n credential_id: int\n source: str\n name: str\n status: BuildConnectorStatus\n docs_indexed: int\n last_indexed: datetime | None\n error_message: str | None = None\n\n\nclass BuildConnectorListResponse(BaseModel):\n \"\"\"List of build connectors.\"\"\"\n\n connectors: list[BuildConnectorInfo]\n\n\n# ===== Suggestion Bubble Models =====\nclass SuggestionTheme(str, Enum):\n \"\"\"Theme/category of a follow-up suggestion.\"\"\"\n\n ADD = \"add\"\n QUESTION = \"question\"\n\n\nclass SuggestionBubble(BaseModel):\n \"\"\"A single follow-up suggestion bubble.\"\"\"\n\n theme: SuggestionTheme\n text: str\n\n\nclass GenerateSuggestionsRequest(BaseModel):\n \"\"\"Request to generate follow-up suggestions.\"\"\"\n\n user_message: str # First user message\n assistant_message: str # First assistant text response (accumulated)\n\n\nclass GenerateSuggestionsResponse(BaseModel):\n \"\"\"Response containing generated suggestions.\"\"\"\n\n suggestions: list[SuggestionBubble]\n\n\nclass PptxPreviewResponse(BaseModel):\n \"\"\"Response with PPTX slide preview metadata.\"\"\"\n\n slide_count: int\n slide_paths: list[str] # Relative paths to slide JPEGs within session workspace\n cached: bool # Whether result was served from cache\n" }, { "path": "backend/onyx/server/features/build/api/packet_logger.py", "content": "\"\"\"Comprehensive packet and ACP event logger for build mode debugging.\n\nLogs all packets, JSON-RPC messages, and ACP events during build mode streaming.\nProvides detailed tracing for the entire agent loop and communication flow.\n\nLog output locations (in priority order):\n1. /var/log/onyx/packets.log (for Docker - mounted to host via docker-compose volumes)\n2. backend/log/packets.log (for local dev without Docker)\n3. backend/onyx/server/features/build/packets.log (fallback)\n\nEnable logging by setting LOG_LEVEL=DEBUG or BUILD_PACKET_LOGGING=true.\n\nFeatures:\n- Rotating log with max 5000 lines (configurable via BUILD_PACKET_LOG_MAX_LINES)\n- Automatically trims oldest entries when limit is exceeded\n- Visual separators between message streams for easy reading\n\"\"\"\n\nimport json\nimport logging\nimport os\nimport threading\nimport time\nfrom pathlib import Path\nfrom typing import Any\nfrom uuid import UUID\n\n# Default max lines to keep in the log file (acts like a deque)\nDEFAULT_MAX_LOG_LINES = 5000\n\n\nclass PacketLogger:\n \"\"\"Comprehensive logger for ACP/OpenCode communication and packet streaming.\n\n Logs:\n - All JSON-RPC requests sent to the agent\n - All JSON-RPC responses/notifications received from the agent\n - All ACP events emitted during streaming\n - Session and sandbox lifecycle events\n - Timing information for debugging performance\n\n The log file is kept to a maximum number of lines (default 5000) to prevent\n unbounded growth. When the limit is exceeded, the oldest lines are trimmed.\n \"\"\"\n\n _instance: \"PacketLogger | None\" = None\n _initialized: bool\n\n def __new__(cls) -> \"PacketLogger\":\n if cls._instance is None:\n cls._instance = super().__new__(cls)\n cls._instance._initialized = False\n return cls._instance\n\n def __init__(self) -> None:\n if self._initialized:\n return\n\n self._initialized = True\n # Enable via LOG_LEVEL=DEBUG or BUILD_PACKET_LOGGING=true\n log_level = os.getenv(\"LOG_LEVEL\", \"\").upper()\n packet_logging = os.getenv(\"BUILD_PACKET_LOGGING\", \"\").lower()\n self._enabled = log_level == \"DEBUG\" or packet_logging in (\"true\", \"1\", \"yes\")\n self._logger: logging.Logger | None = None\n self._log_file_path: Path | None = None\n self._session_start_times: dict[str, float] = {}\n\n # Max lines to keep in log file\n try:\n self._max_lines = int(\n os.getenv(\"BUILD_PACKET_LOG_MAX_LINES\", str(DEFAULT_MAX_LOG_LINES))\n )\n except ValueError:\n self._max_lines = DEFAULT_MAX_LOG_LINES\n\n # Lock for thread-safe file operations\n self._file_lock = threading.Lock()\n\n # Track approximate line count to avoid reading file too often\n self._approx_line_count = 0\n self._lines_since_last_trim = 0\n # Trim every N lines written to avoid constant file reads\n self._trim_interval = 500\n\n if self._enabled:\n self._setup_logger()\n\n def _get_log_file_path(self) -> Path:\n \"\"\"Determine the best log file path based on environment.\n\n Priority:\n 1. /var/log/onyx/packets.log - Docker environment (mounted to host)\n 2. backend/log/packets.log - Local dev (same dir as other logs)\n 3. backend/onyx/server/features/build/packets.log - Fallback\n \"\"\"\n # Option 1: Docker environment - use /var/log/onyx which is mounted\n docker_log_dir = Path(\"/var/log/onyx\")\n if docker_log_dir.exists() and docker_log_dir.is_dir():\n return docker_log_dir / \"packets.log\"\n\n # Option 2: Local dev - use backend/log directory (same as other debug logs)\n # Navigate from this file to backend/log\n backend_dir = Path(__file__).parents[4] # up to backend/\n local_log_dir = backend_dir / \"log\"\n if local_log_dir.exists() and local_log_dir.is_dir():\n return local_log_dir / \"packets.log\"\n\n # Option 3: Fallback to build directory\n build_dir = Path(__file__).parents[1]\n return build_dir / \"packets.log\"\n\n def _setup_logger(self) -> None:\n \"\"\"Set up the file handler for packet logging.\"\"\"\n self._log_file_path = self._get_log_file_path()\n\n # Ensure parent directory exists\n self._log_file_path.parent.mkdir(parents=True, exist_ok=True)\n\n self._logger = logging.getLogger(\"build.packets\")\n self._logger.setLevel(logging.DEBUG)\n self._logger.propagate = False\n\n self._logger.handlers.clear()\n\n # Use append mode\n handler = logging.FileHandler(self._log_file_path, mode=\"a\", encoding=\"utf-8\")\n handler.setLevel(logging.DEBUG)\n # Include timestamp in each log entry\n handler.setFormatter(\n logging.Formatter(\n \"%(asctime)s.%(msecs)03d | %(message)s\", \"%Y-%m-%d %H:%M:%S\"\n )\n )\n\n self._logger.addHandler(handler)\n\n # Initialize line count from existing file\n self._init_line_count()\n\n def _init_line_count(self) -> None:\n \"\"\"Initialize the approximate line count from the existing log file.\"\"\"\n if not self._log_file_path or not self._log_file_path.exists():\n self._approx_line_count = 0\n return\n\n try:\n with open(self._log_file_path, \"r\", encoding=\"utf-8\", errors=\"ignore\") as f:\n self._approx_line_count = sum(1 for _ in f)\n except Exception:\n self._approx_line_count = 0\n\n def _maybe_trim_log(self) -> None:\n \"\"\"Trim the log file if it exceeds the max line limit.\n\n This is called periodically (every _trim_interval lines) to avoid\n reading the file on every write.\n \"\"\"\n self._lines_since_last_trim += 1\n\n if self._lines_since_last_trim < self._trim_interval:\n return\n\n self._lines_since_last_trim = 0\n self._trim_log_file()\n\n def _trim_log_file(self) -> None:\n \"\"\"Trim the log file to keep only the last max_lines.\"\"\"\n if not self._log_file_path or not self._log_file_path.exists():\n return\n\n with self._file_lock:\n try:\n # Read all lines\n with open(\n self._log_file_path, \"r\", encoding=\"utf-8\", errors=\"ignore\"\n ) as f:\n lines = f.readlines()\n\n current_count = len(lines)\n self._approx_line_count = current_count\n\n # If under limit, nothing to do\n if current_count <= self._max_lines:\n return\n\n # Keep only the last max_lines\n lines_to_keep = lines[-self._max_lines :]\n\n # Close the logger's file handler temporarily\n if self._logger:\n for handler in self._logger.handlers:\n handler.close()\n\n # Rewrite the file with trimmed content\n with open(self._log_file_path, \"w\", encoding=\"utf-8\") as f:\n f.writelines(lines_to_keep)\n\n # Reopen the handler\n if self._logger:\n self._logger.handlers.clear()\n handler = logging.FileHandler(\n self._log_file_path, mode=\"a\", encoding=\"utf-8\"\n )\n handler.setLevel(logging.DEBUG)\n handler.setFormatter(\n logging.Formatter(\n \"%(asctime)s.%(msecs)03d | %(message)s\", \"%Y-%m-%d %H:%M:%S\"\n )\n )\n self._logger.addHandler(handler)\n\n self._approx_line_count = len(lines_to_keep)\n\n except Exception:\n pass # Silently ignore errors during trim\n\n def clear_log_file(self) -> None:\n \"\"\"Clear the log file contents.\n\n Note: With the rotating log approach, this is optional. The log will\n automatically trim itself. But this can still be useful to start fresh.\n \"\"\"\n if not self._enabled or not self._log_file_path:\n return\n\n with self._file_lock:\n try:\n # Close the logger's file handler temporarily\n if self._logger:\n for handler in self._logger.handlers:\n handler.close()\n\n # Truncate the file\n with open(self._log_file_path, \"w\", encoding=\"utf-8\") as f:\n f.write(\"\") # Empty the file\n\n # Reopen the handler\n if self._logger:\n self._logger.handlers.clear()\n handler = logging.FileHandler(\n self._log_file_path, mode=\"a\", encoding=\"utf-8\"\n )\n handler.setLevel(logging.DEBUG)\n handler.setFormatter(\n logging.Formatter(\n \"%(asctime)s.%(msecs)03d | %(message)s\", \"%Y-%m-%d %H:%M:%S\"\n )\n )\n self._logger.addHandler(handler)\n\n self._approx_line_count = 0\n self._lines_since_last_trim = 0\n\n except Exception:\n pass # Silently ignore errors\n\n @property\n def is_enabled(self) -> bool:\n \"\"\"Check if logging is enabled.\"\"\"\n return self._enabled and self._logger is not None\n\n def _format_uuid(self, value: Any) -> str:\n \"\"\"Format UUID for logging (shortened for readability).\"\"\"\n if isinstance(value, UUID):\n return str(value)[:8]\n if isinstance(value, str) and len(value) >= 8:\n return value[:8]\n return str(value)\n\n def _write_log(self, message: str) -> None:\n \"\"\"Internal method to write a log message and trigger trim check.\n\n Args:\n message: The formatted log message\n \"\"\"\n if not self._logger:\n return\n\n self._logger.debug(message)\n self._maybe_trim_log()\n\n def log(self, packet_type: str, payload: dict[str, Any] | None = None) -> None:\n \"\"\"Log a packet as JSON.\n\n Args:\n packet_type: The type of packet\n payload: The packet payload\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n output = json.dumps(payload, indent=2, default=str) if payload else \"{}\"\n self._write_log(f\"[PACKET] {packet_type}\\n{output}\")\n except Exception:\n self._write_log(f\"[PACKET] {packet_type}\\n{payload}\")\n\n def log_raw(self, label: str, data: Any) -> None:\n \"\"\"Log raw data with a label.\n\n Args:\n label: A label for this log entry\n data: Any data to log\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n if isinstance(data, (dict, list)):\n output = json.dumps(data, indent=2, default=str)\n else:\n output = str(data)\n self._write_log(f\"[RAW] {label}\\n{output}\")\n except Exception:\n self._write_log(f\"[RAW] {label}\\n{data}\")\n\n # =========================================================================\n # JSON-RPC Communication Logging\n # =========================================================================\n\n def log_jsonrpc_request(\n self,\n method: str,\n request_id: int | None,\n params: dict[str, Any] | None = None,\n context: str = \"\",\n ) -> None:\n \"\"\"Log a JSON-RPC request being sent to the agent.\n\n Args:\n method: The JSON-RPC method name\n request_id: The request ID (None for notifications)\n params: The request parameters\n context: Additional context (e.g., \"local\", \"k8s\")\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n req_type = \"REQUEST\" if request_id is not None else \"NOTIFICATION\"\n ctx_prefix = f\"[{context}] \" if context else \"\"\n params_str = json.dumps(params, indent=2, default=str) if params else \"{}\"\n id_str = f\" id={request_id}\" if request_id is not None else \"\"\n self._write_log(\n f\"{ctx_prefix}[JSONRPC-OUT] {req_type} {method}{id_str}\\n{params_str}\"\n )\n except Exception as e:\n self._write_log(f\"[JSONRPC-OUT] {method} (logging error: {e})\")\n\n def log_jsonrpc_response(\n self,\n request_id: int | None,\n result: dict[str, Any] | None = None,\n error: dict[str, Any] | None = None,\n context: str = \"\",\n ) -> None:\n \"\"\"Log a JSON-RPC response received from the agent.\n\n Args:\n request_id: The request ID this is responding to\n result: The result payload (if success)\n error: The error payload (if error)\n context: Additional context (e.g., \"local\", \"k8s\")\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n ctx_prefix = f\"[{context}] \" if context else \"\"\n id_str = f\" id={request_id}\" if request_id is not None else \"\"\n if error:\n error_str = json.dumps(error, indent=2, default=str)\n self._write_log(\n f\"{ctx_prefix}[JSONRPC-IN] RESPONSE{id_str} ERROR\\n{error_str}\"\n )\n else:\n result_str = (\n json.dumps(result, indent=2, default=str) if result else \"{}\"\n )\n self._write_log(\n f\"{ctx_prefix}[JSONRPC-IN] RESPONSE{id_str}\\n{result_str}\"\n )\n except Exception as e:\n self._write_log(f\"[JSONRPC-IN] RESPONSE (logging error: {e})\")\n\n def log_jsonrpc_notification(\n self,\n method: str,\n params: dict[str, Any] | None = None,\n context: str = \"\",\n ) -> None:\n \"\"\"Log a JSON-RPC notification received from the agent.\n\n Args:\n method: The notification method name\n params: The notification parameters\n context: Additional context (e.g., \"local\", \"k8s\")\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n ctx_prefix = f\"[{context}] \" if context else \"\"\n params_str = json.dumps(params, indent=2, default=str) if params else \"{}\"\n self._write_log(\n f\"{ctx_prefix}[JSONRPC-IN] NOTIFICATION {method}\\n{params_str}\"\n )\n except Exception as e:\n self._write_log(f\"[JSONRPC-IN] NOTIFICATION {method} (logging error: {e})\")\n\n def log_jsonrpc_raw_message(\n self,\n direction: str,\n message: dict[str, Any] | str,\n context: str = \"\",\n ) -> None:\n \"\"\"Log a raw JSON-RPC message (for debugging parsing issues).\n\n Args:\n direction: \"IN\" or \"OUT\"\n message: The raw message (dict or string)\n context: Additional context\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n ctx_prefix = f\"[{context}] \" if context else \"\"\n if isinstance(message, dict):\n msg_str = json.dumps(message, indent=2, default=str)\n else:\n msg_str = str(message)\n self._write_log(f\"{ctx_prefix}[JSONRPC-RAW-{direction}]\\n{msg_str}\")\n except Exception as e:\n self._write_log(f\"[JSONRPC-RAW-{direction}] (logging error: {e})\")\n\n # =========================================================================\n # ACP Event Logging\n # =========================================================================\n\n def log_acp_event(\n self,\n event_type: str,\n event_data: dict[str, Any],\n sandbox_id: UUID | str | None = None,\n session_id: UUID | str | None = None,\n ) -> None:\n \"\"\"Log an ACP event being emitted.\n\n Args:\n event_type: The ACP event type (e.g., \"agent_message_chunk\")\n event_data: The full event data\n sandbox_id: The sandbox ID (optional, for context)\n session_id: The session ID (optional, for context)\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n ctx_parts = []\n if sandbox_id:\n ctx_parts.append(f\"sandbox={self._format_uuid(sandbox_id)}\")\n if session_id:\n ctx_parts.append(f\"session={self._format_uuid(session_id)}\")\n ctx = f\" ({', '.join(ctx_parts)})\" if ctx_parts else \"\"\n\n # For message chunks, show truncated content for readability\n display_data = event_data.copy()\n if event_type in (\"agent_message_chunk\", \"agent_thought_chunk\"):\n content = display_data.get(\"content\", {})\n if isinstance(content, dict) and \"text\" in content:\n text = content.get(\"text\", \"\")\n if len(text) > 200:\n display_data[\"content\"] = {\n **content,\n \"text\": text[:200] + f\"... ({len(text)} chars total)\",\n }\n\n event_str = json.dumps(display_data, indent=2, default=str)\n self._write_log(f\"[ACP-EVENT] {event_type}{ctx}\\n{event_str}\")\n except Exception as e:\n self._write_log(f\"[ACP-EVENT] {event_type} (logging error: {e})\")\n\n def log_acp_event_yielded(\n self,\n event_type: str,\n event_obj: Any,\n sandbox_id: UUID | str | None = None,\n session_id: UUID | str | None = None,\n ) -> None:\n \"\"\"Log an ACP event object being yielded from the generator.\n\n Args:\n event_type: The ACP event type\n event_obj: The Pydantic event object\n sandbox_id: The sandbox ID (optional)\n session_id: The session ID (optional)\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n if hasattr(event_obj, \"model_dump\"):\n event_data = event_obj.model_dump(mode=\"json\", by_alias=True)\n else:\n event_data = {\"raw\": str(event_obj)}\n self.log_acp_event(event_type, event_data, sandbox_id, session_id)\n except Exception as e:\n self._write_log(f\"[ACP-EVENT] {event_type} (logging error: {e})\")\n\n # =========================================================================\n # Session and Sandbox Lifecycle Logging\n # =========================================================================\n\n def log_session_start(\n self,\n session_id: UUID | str,\n sandbox_id: UUID | str,\n message_preview: str = \"\",\n ) -> None:\n \"\"\"Log the start of a message streaming session.\n\n Args:\n session_id: The session ID\n sandbox_id: The sandbox ID\n message_preview: First 100 chars of the user message\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n session_key = str(session_id)\n self._session_start_times[session_key] = time.time()\n\n preview = (\n message_preview[:100] + \"...\"\n if len(message_preview) > 100\n else message_preview\n )\n self._write_log(\n f\"[SESSION-START] session={self._format_uuid(session_id)} \"\n f\"sandbox={self._format_uuid(sandbox_id)}\\n\"\n f\" message: {preview}\"\n )\n\n def log_session_end(\n self,\n session_id: UUID | str,\n success: bool = True,\n error: str | None = None,\n events_count: int = 0,\n ) -> None:\n \"\"\"Log the end of a message streaming session.\n\n Args:\n session_id: The session ID\n success: Whether the session completed successfully\n error: Error message if failed\n events_count: Number of events emitted\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n session_key = str(session_id)\n start_time = self._session_start_times.pop(session_key, None)\n duration_ms = (time.time() - start_time) * 1000 if start_time else 0\n\n status = \"SUCCESS\" if success else \"FAILED\"\n error_str = f\"\\n error: {error}\" if error else \"\"\n self._write_log(\n f\"[SESSION-END] session={self._format_uuid(session_id)} \"\n f\"status={status} duration={duration_ms:.0f}ms events={events_count}\"\n f\"{error_str}\"\n )\n\n def log_acp_client_start(\n self,\n sandbox_id: UUID | str,\n session_id: UUID | str,\n cwd: str,\n context: str = \"\",\n ) -> None:\n \"\"\"Log ACP client initialization.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n cwd: Working directory\n context: \"local\" or \"k8s\"\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n ctx_prefix = f\"[{context}] \" if context else \"\"\n self._write_log(\n f\"{ctx_prefix}[ACP-CLIENT-START] \"\n f\"sandbox={self._format_uuid(sandbox_id)} \"\n f\"session={self._format_uuid(session_id)}\\n\"\n f\" cwd: {cwd}\"\n )\n\n def log_acp_client_stop(\n self,\n sandbox_id: UUID | str,\n session_id: UUID | str,\n context: str = \"\",\n ) -> None:\n \"\"\"Log ACP client shutdown.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n context: \"local\" or \"k8s\"\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n ctx_prefix = f\"[{context}] \" if context else \"\"\n self._write_log(\n f\"{ctx_prefix}[ACP-CLIENT-STOP] sandbox={self._format_uuid(sandbox_id)} session={self._format_uuid(session_id)}\"\n )\n\n # =========================================================================\n # Streaming State Logging\n # =========================================================================\n\n def log_streaming_state_update(\n self,\n session_id: UUID | str,\n state_type: str,\n details: dict[str, Any] | None = None,\n ) -> None:\n \"\"\"Log streaming state changes.\n\n Args:\n session_id: The session ID\n state_type: Type of state change (e.g., \"chunk_accumulated\", \"saved_to_db\")\n details: Additional details\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n try:\n details_str = \"\"\n if details:\n details_str = \"\\n\" + json.dumps(details, indent=2, default=str)\n self._write_log(\n f\"[STREAMING-STATE] session={self._format_uuid(session_id)} type={state_type}{details_str}\"\n )\n except Exception as e:\n self._write_log(f\"[STREAMING-STATE] {state_type} (logging error: {e})\")\n\n def log_sse_emit(\n self,\n event_type: str,\n session_id: UUID | str | None = None,\n ) -> None:\n \"\"\"Log SSE event being emitted to frontend.\n\n Args:\n event_type: The event type being emitted\n session_id: The session ID\n \"\"\"\n if not self._enabled or not self._logger:\n return\n\n session_str = f\" session={self._format_uuid(session_id)}\" if session_id else \"\"\n self._write_log(f\"[SSE-EMIT] {event_type}{session_str}\")\n\n\n# Singleton instance\n_packet_logger: PacketLogger | None = None\n\n\ndef get_packet_logger() -> PacketLogger:\n \"\"\"Get the singleton packet logger instance.\"\"\"\n global _packet_logger\n if _packet_logger is None:\n _packet_logger = PacketLogger()\n return _packet_logger\n\n\ndef log_separator(label: str = \"\") -> None:\n \"\"\"Log a visual separator for readability in the log file.\n\n Args:\n label: Optional label for the separator\n \"\"\"\n logger = get_packet_logger()\n if not logger.is_enabled or not logger._logger:\n return\n\n separator = \"=\" * 80\n if label:\n logger._write_log(f\"\\n{separator}\\n{label}\\n{separator}\")\n else:\n logger._write_log(f\"\\n{separator}\")\n" }, { "path": "backend/onyx/server/features/build/api/packets.py", "content": "\"\"\"Build Mode packet types for streaming agent responses.\n\nThis module defines CUSTOM Onyx packet types that extend ACP (Agent Client Protocol).\nACP events are passed through directly from the agent - this module only contains\nOnyx-specific extensions like artifacts and file operations.\n\nAll packets use SSE (Server-Sent Events) format with `event: message` and include\na `type` field to distinguish packet types.\n\nACP events (passed through directly from acp.schema):\n- agent_message_chunk: Text/image content from agent\n- agent_thought_chunk: Agent's internal reasoning\n- tool_call_start: Tool invocation started\n- tool_call_progress: Tool execution progress/result\n- agent_plan_update: Agent's execution plan\n- current_mode_update: Agent mode change\n- prompt_response: Agent finished processing\n- error: An error occurred\n\nCustom Onyx packets (defined here):\n- error: Onyx-specific errors (e.g., session not found)\n\nBased on:\n- Agent Client Protocol (ACP): https://agentclientprotocol.com\n\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import Literal\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\n\n# =============================================================================\n# Base Packet Type\n# =============================================================================\n\n\nclass BasePacket(BaseModel):\n \"\"\"Base packet with common fields for all custom Onyx packet types.\"\"\"\n\n type: str\n timestamp: str = Field(\n default_factory=lambda: datetime.now(tz=timezone.utc).isoformat()\n )\n\n\n# =============================================================================\n# Custom Onyx Packets\n# =============================================================================\n\n\nclass ErrorPacket(BasePacket):\n \"\"\"An Onyx-specific error occurred (e.g., session not found, sandbox not running).\"\"\"\n\n type: Literal[\"error\"] = \"error\"\n message: str\n code: int | None = None\n details: dict[str, Any] | None = None\n\n\n# =============================================================================\n# Union Type for Custom Onyx Packets\n# =============================================================================\n\nBuildPacket = ErrorPacket\n" }, { "path": "backend/onyx/server/features/build/api/rate_limit.py", "content": "\"\"\"Rate limiting logic for Build Mode.\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom typing import Literal\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import User\nfrom onyx.feature_flags.factory import get_default_feature_flag_provider\nfrom onyx.server.features.build.api.models import RateLimitResponse\nfrom onyx.server.features.build.api.subscription_check import is_user_subscribed\nfrom onyx.server.features.build.configs import CRAFT_PAID_USER_RATE_LIMIT\nfrom onyx.server.features.build.db.rate_limit import count_user_messages_in_window\nfrom onyx.server.features.build.db.rate_limit import count_user_messages_total\nfrom onyx.server.features.build.db.rate_limit import get_oldest_message_timestamp\nfrom onyx.server.features.build.utils import CRAFT_HAS_USAGE_LIMITS\nfrom shared_configs.configs import MULTI_TENANT\n\n# Default limit for free/non-subscribed users (not configurable)\nFREE_USER_RATE_LIMIT = 5\n\n\ndef _should_skip_rate_limiting(user: User) -> bool:\n \"\"\"\n Check if rate limiting should be skipped for this user.\n\n Currently grants unlimited usage to dev tenant users (tenant_dev).\n Controlled via PostHog feature flag.\n\n Returns:\n True to skip rate limiting (unlimited), False to apply normal limits\n \"\"\"\n # NOTE: We can modify the posthog flag to return more detail about a limit\n # i.e. can set variable limits per user and tenant via PostHog instead of env vars\n # to avoid re-deploying on every limit change\n\n feature_flag_provider = get_default_feature_flag_provider()\n # Flag returns True for users who SHOULD be rate limited\n # We negate to get: True = skip rate limiting\n has_rate_limit = feature_flag_provider.feature_enabled(\n CRAFT_HAS_USAGE_LIMITS,\n user.id,\n )\n return not has_rate_limit\n\n\ndef get_user_rate_limit_status(\n user: User,\n db_session: Session,\n) -> RateLimitResponse:\n \"\"\"\n Get the rate limit status for a user.\n\n Rate limits:\n - Cloud (MULTI_TENANT=true):\n - Subscribed users: CRAFT_PAID_USER_RATE_LIMIT messages per week\n (configurable, default 25)\n - Non-subscribed users: 5 messages (lifetime total)\n - Per-user overrides via PostHog feature flag\n - Self-hosted (MULTI_TENANT=false):\n - Unlimited (no rate limiting)\n\n Args:\n user: The authenticated user\n db_session: Database session\n\n Returns:\n RateLimitResponse with current limit status\n \"\"\"\n # Self-hosted deployments have no rate limits\n if not MULTI_TENANT:\n return RateLimitResponse(\n is_limited=False,\n limit_type=\"weekly\",\n messages_used=0,\n limit=0, # 0 indicates unlimited\n reset_timestamp=None,\n )\n\n # Check if user should skip rate limiting (e.g., dev tenant users)\n if _should_skip_rate_limiting(user):\n return RateLimitResponse(\n is_limited=False,\n limit_type=\"weekly\",\n messages_used=-1,\n limit=0, # 0 indicates unlimited\n reset_timestamp=None,\n )\n\n # Determine subscription status\n is_subscribed = is_user_subscribed(user, db_session)\n\n # Get limit based on subscription status\n limit = CRAFT_PAID_USER_RATE_LIMIT if is_subscribed else FREE_USER_RATE_LIMIT\n\n # Limit type: weekly for subscribed users, total for free\n limit_type: Literal[\"weekly\", \"total\"] = \"weekly\" if is_subscribed else \"total\"\n\n # Count messages\n if limit_type == \"weekly\":\n # Subscribed: rolling 7-day window\n cutoff_time = datetime.now(tz=timezone.utc) - timedelta(days=7)\n messages_used = count_user_messages_in_window(user.id, cutoff_time, db_session)\n\n # Calculate reset timestamp (when oldest message ages out)\n # Only show reset time if user is at or over the limit\n if messages_used >= limit:\n oldest_msg = get_oldest_message_timestamp(user.id, cutoff_time, db_session)\n if oldest_msg:\n reset_time = oldest_msg + timedelta(days=7)\n reset_timestamp = reset_time.isoformat()\n else:\n reset_timestamp = None\n else:\n reset_timestamp = None\n else:\n # Non-subscribed: lifetime total\n messages_used = count_user_messages_total(user.id, db_session)\n reset_timestamp = None\n\n return RateLimitResponse(\n is_limited=messages_used >= limit,\n limit_type=limit_type,\n messages_used=messages_used,\n limit=limit,\n reset_timestamp=reset_timestamp,\n )\n" }, { "path": "backend/onyx/server/features/build/api/sessions_api.py", "content": "\"\"\"API endpoints for Build Mode session management.\"\"\"\n\nfrom uuid import UUID\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import File\nfrom fastapi import HTTPException\nfrom fastapi import Response\nfrom fastapi import UploadFile\nfrom sqlalchemy import exists\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_user\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import BuildSessionStatus\nfrom onyx.db.enums import SandboxStatus\nfrom onyx.db.models import BuildMessage\nfrom onyx.db.models import User\nfrom onyx.redis.redis_pool import get_redis_client\nfrom onyx.server.features.build.api.models import ArtifactResponse\nfrom onyx.server.features.build.api.models import DetailedSessionResponse\nfrom onyx.server.features.build.api.models import DirectoryListing\nfrom onyx.server.features.build.api.models import GenerateSuggestionsRequest\nfrom onyx.server.features.build.api.models import GenerateSuggestionsResponse\nfrom onyx.server.features.build.api.models import PptxPreviewResponse\nfrom onyx.server.features.build.api.models import PreProvisionedCheckResponse\nfrom onyx.server.features.build.api.models import SessionCreateRequest\nfrom onyx.server.features.build.api.models import SessionListResponse\nfrom onyx.server.features.build.api.models import SessionNameGenerateResponse\nfrom onyx.server.features.build.api.models import SessionResponse\nfrom onyx.server.features.build.api.models import SessionUpdateRequest\nfrom onyx.server.features.build.api.models import SetSessionSharingRequest\nfrom onyx.server.features.build.api.models import SetSessionSharingResponse\nfrom onyx.server.features.build.api.models import SuggestionBubble\nfrom onyx.server.features.build.api.models import SuggestionTheme\nfrom onyx.server.features.build.api.models import UploadResponse\nfrom onyx.server.features.build.api.models import WebappInfo\nfrom onyx.server.features.build.configs import SANDBOX_BACKEND\nfrom onyx.server.features.build.configs import SandboxBackend\nfrom onyx.server.features.build.db.build_session import allocate_nextjs_port\nfrom onyx.server.features.build.db.build_session import get_build_session\nfrom onyx.server.features.build.db.build_session import set_build_session_sharing_scope\nfrom onyx.server.features.build.db.sandbox import get_latest_snapshot_for_session\nfrom onyx.server.features.build.db.sandbox import get_sandbox_by_user_id\nfrom onyx.server.features.build.db.sandbox import update_sandbox_heartbeat\nfrom onyx.server.features.build.db.sandbox import update_sandbox_status__no_commit\nfrom onyx.server.features.build.sandbox import get_sandbox_manager\nfrom onyx.server.features.build.session.manager import SessionManager\nfrom onyx.server.features.build.session.manager import UploadLimitExceededError\nfrom onyx.server.features.build.utils import sanitize_filename\nfrom onyx.server.features.build.utils import validate_file\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/sessions\")\n\n\n# =============================================================================\n# Session Management Endpoints\n# =============================================================================\n\n\n@router.get(\"\", response_model=SessionListResponse)\ndef list_sessions(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> SessionListResponse:\n \"\"\"List all build sessions for the current user.\"\"\"\n session_manager = SessionManager(db_session)\n\n sessions = session_manager.list_sessions(user.id)\n\n # Get the user's sandbox (shared across all sessions)\n sandbox = get_sandbox_by_user_id(db_session, user.id)\n\n return SessionListResponse(\n sessions=[SessionResponse.from_model(session, sandbox) for session in sessions]\n )\n\n\n# Lock timeout for session creation (should be longer than max provision time)\nSESSION_CREATE_LOCK_TIMEOUT_SECONDS = 300\n\n\n@router.post(\"\", response_model=DetailedSessionResponse)\ndef create_session(\n request: SessionCreateRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> DetailedSessionResponse:\n \"\"\"\n Create or get an existing empty build session.\n\n Creates a sandbox with the necessary file structure and returns a session ID.\n Uses SessionManager for session and sandbox provisioning.\n\n This endpoint is atomic - if sandbox provisioning fails, no database\n records are created (transaction is rolled back).\n\n Uses Redis lock to prevent race conditions when multiple requests try to\n create/provision a session for the same user concurrently.\n \"\"\"\n tenant_id = get_current_tenant_id()\n redis_client = get_redis_client(tenant_id=tenant_id)\n\n # Lock on user_id to prevent concurrent session creation for the same user\n # This prevents race conditions where two requests both see sandbox as SLEEPING\n # and both try to provision, with one deleting the other's work\n lock_key = f\"session_create:{user.id}\"\n lock = redis_client.lock(lock_key, timeout=SESSION_CREATE_LOCK_TIMEOUT_SECONDS)\n\n # blocking=True means wait if another create is in progress\n acquired = lock.acquire(\n blocking=True, blocking_timeout=SESSION_CREATE_LOCK_TIMEOUT_SECONDS\n )\n if not acquired:\n raise HTTPException(\n status_code=503,\n detail=\"Session creation timed out waiting for lock\",\n )\n\n try:\n session_manager = SessionManager(db_session)\n build_session = session_manager.get_or_create_empty_session(\n user.id,\n user_work_area=(\n request.user_work_area if request.demo_data_enabled else None\n ),\n user_level=request.user_level if request.demo_data_enabled else None,\n llm_provider_type=request.llm_provider_type,\n llm_model_name=request.llm_model_name,\n demo_data_enabled=request.demo_data_enabled,\n )\n db_session.commit()\n\n sandbox = get_sandbox_by_user_id(db_session, user.id)\n base_response = SessionResponse.from_model(build_session, sandbox)\n return DetailedSessionResponse.from_session_response(\n base_response, session_loaded_in_sandbox=True\n )\n except ValueError as e:\n logger.exception(\"Session creation failed\")\n db_session.rollback()\n raise HTTPException(status_code=429, detail=str(e))\n except Exception as e:\n db_session.rollback()\n logger.error(f\"Session creation failed: {e}\")\n raise HTTPException(status_code=500, detail=f\"Session creation failed: {e}\")\n finally:\n if lock.owned():\n lock.release()\n\n\n@router.get(\"/{session_id}\", response_model=DetailedSessionResponse)\ndef get_session_details(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> DetailedSessionResponse:\n \"\"\"\n Get details of a specific build session.\n\n Returns session_loaded_in_sandbox to indicate if the session workspace\n exists in the running sandbox.\n \"\"\"\n session_manager = SessionManager(db_session)\n\n session = session_manager.get_session(session_id, user.id)\n\n if session is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n # Get the user's sandbox to include in response\n sandbox = get_sandbox_by_user_id(db_session, user.id)\n\n # Check if session workspace exists in the sandbox\n session_loaded = False\n if sandbox and sandbox.status == SandboxStatus.RUNNING:\n sandbox_manager = get_sandbox_manager()\n session_loaded = sandbox_manager.session_workspace_exists(\n sandbox.id, session_id\n )\n\n base_response = SessionResponse.from_model(session, sandbox)\n return DetailedSessionResponse.from_session_response(\n base_response, session_loaded_in_sandbox=session_loaded\n )\n\n\n@router.get(\n \"/{session_id}/pre-provisioned-check\", response_model=PreProvisionedCheckResponse\n)\ndef check_pre_provisioned_session(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> PreProvisionedCheckResponse:\n \"\"\"\n Check if a pre-provisioned session is still valid (empty).\n\n Used by the frontend to poll and detect when another tab has used\n the session. A session is considered valid if it has no messages yet.\n\n Returns:\n - valid=True, session_id= if the session is still empty\n - valid=False, session_id=None if the session has messages or doesn't exist\n \"\"\"\n session = get_build_session(session_id, user.id, db_session)\n\n if session is None:\n return PreProvisionedCheckResponse(valid=False, session_id=None)\n\n # Check if session is still empty (no messages = pre-provisioned)\n has_messages = db_session.query(\n exists().where(BuildMessage.session_id == session_id)\n ).scalar()\n\n if not has_messages:\n return PreProvisionedCheckResponse(valid=True, session_id=str(session_id))\n\n # Session has messages - it's no longer a valid pre-provisioned session\n return PreProvisionedCheckResponse(valid=False, session_id=None)\n\n\n@router.post(\"/{session_id}/generate-name\", response_model=SessionNameGenerateResponse)\ndef generate_session_name(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> SessionNameGenerateResponse:\n \"\"\"Generate a session name using LLM based on the first user message.\"\"\"\n session_manager = SessionManager(db_session)\n\n generated_name = session_manager.generate_session_name(session_id, user.id)\n\n if generated_name is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n return SessionNameGenerateResponse(name=generated_name)\n\n\n@router.post(\n \"/{session_id}/generate-suggestions\", response_model=GenerateSuggestionsResponse\n)\ndef generate_suggestions(\n session_id: UUID,\n request: GenerateSuggestionsRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> GenerateSuggestionsResponse:\n \"\"\"Generate follow-up suggestions based on the first exchange in a session.\"\"\"\n session_manager = SessionManager(db_session)\n\n # Verify session exists and belongs to user\n session = session_manager.get_session(session_id, user.id)\n if session is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n # Generate suggestions\n suggestions_data = session_manager.generate_followup_suggestions(\n user_message=request.user_message,\n assistant_message=request.assistant_message,\n )\n\n # Convert to response model\n suggestions = [\n SuggestionBubble(\n theme=SuggestionTheme(item[\"theme\"]),\n text=item[\"text\"],\n )\n for item in suggestions_data\n ]\n\n return GenerateSuggestionsResponse(suggestions=suggestions)\n\n\n@router.put(\"/{session_id}/name\", response_model=SessionResponse)\ndef update_session_name(\n session_id: UUID,\n request: SessionUpdateRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> SessionResponse:\n \"\"\"Update the name of a build session.\"\"\"\n session_manager = SessionManager(db_session)\n\n session = session_manager.update_session_name(session_id, user.id, request.name)\n\n if session is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n # Get the user's sandbox to include in response\n sandbox = get_sandbox_by_user_id(db_session, user.id)\n return SessionResponse.from_model(session, sandbox)\n\n\n@router.patch(\"/{session_id}/public\")\ndef set_session_public(\n session_id: UUID,\n request: SetSessionSharingRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> SetSessionSharingResponse:\n \"\"\"Set the sharing scope of a build session's webapp.\"\"\"\n updated = set_build_session_sharing_scope(\n session_id, user.id, request.sharing_scope, db_session\n )\n if not updated:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n return SetSessionSharingResponse(\n session_id=str(session_id),\n sharing_scope=updated.sharing_scope,\n )\n\n\n@router.delete(\"/{session_id}\", response_model=None)\ndef delete_session(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"Delete a build session and all associated data.\n\n This endpoint is atomic - if sandbox termination fails, the session\n is NOT deleted (transaction is rolled back).\n \"\"\"\n session_manager = SessionManager(db_session)\n\n try:\n success = session_manager.delete_session(session_id, user.id)\n if not success:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n db_session.commit()\n except HTTPException:\n # Re-raise HTTP exceptions (like 404) without rollback\n raise\n except Exception as e:\n # Sandbox termination failed - rollback to preserve session\n db_session.rollback()\n logger.error(f\"Failed to delete session {session_id}: {e}\")\n raise HTTPException(\n status_code=500,\n detail=f\"Failed to delete session: {e}\",\n )\n\n return Response(status_code=204)\n\n\n# Lock timeout should be longer than max restore time (5 minutes)\nRESTORE_LOCK_TIMEOUT_SECONDS = 300\n\n\n@router.post(\"/{session_id}/restore\", response_model=DetailedSessionResponse)\ndef restore_session(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> DetailedSessionResponse:\n \"\"\"Restore sandbox and load session snapshot. Blocks until complete.\n\n Uses Redis lock to ensure only one restore runs per sandbox at a time.\n If another restore is in progress, waits for it to complete.\n\n Handles two cases:\n 1. Sandbox is SLEEPING: Re-provision pod, then load session snapshot\n 2. Sandbox is RUNNING but session not loaded: Just load session snapshot\n\n Returns immediately if session workspace already exists in pod.\n Always returns session_loaded_in_sandbox=True on success.\n \"\"\"\n session = get_build_session(session_id, user.id, db_session)\n if not session:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n sandbox = get_sandbox_by_user_id(db_session, user.id)\n if not sandbox:\n raise HTTPException(status_code=404, detail=\"Sandbox not found\")\n\n # If sandbox is already running, check if session workspace exists\n sandbox_manager = get_sandbox_manager()\n tenant_id = get_current_tenant_id()\n\n # Need to do some work - acquire Redis lock\n redis_client = get_redis_client(tenant_id=tenant_id)\n lock_key = f\"sandbox_restore:{sandbox.id}\"\n lock = redis_client.lock(lock_key, timeout=RESTORE_LOCK_TIMEOUT_SECONDS)\n\n # Non-blocking: if another restore is already running, return 409 immediately\n # instead of making the user wait. The frontend will retry.\n acquired = lock.acquire(blocking=False)\n if not acquired:\n raise HTTPException(\n status_code=409,\n detail=\"Restore already in progress\",\n )\n\n try:\n # Re-fetch sandbox status (may have changed while waiting for lock)\n db_session.refresh(sandbox)\n\n # Also re-check if session workspace exists (another request may have\n # restored it while we were waiting)\n if sandbox.status == SandboxStatus.RUNNING:\n is_healthy = sandbox_manager.health_check(sandbox.id, timeout=10.0)\n if is_healthy and sandbox_manager.session_workspace_exists(\n sandbox.id, session_id\n ):\n session.status = BuildSessionStatus.ACTIVE\n update_sandbox_heartbeat(db_session, sandbox.id)\n base_response = SessionResponse.from_model(session, sandbox)\n return DetailedSessionResponse.from_session_response(\n base_response, session_loaded_in_sandbox=True\n )\n\n if not is_healthy:\n logger.warning(\n f\"Sandbox {sandbox.id} marked as RUNNING but pod is unhealthy/missing. Entering recovery mode.\"\n )\n # Terminate to clean up any lingering K8s resources\n sandbox_manager.terminate(sandbox.id)\n\n update_sandbox_status__no_commit(\n db_session, sandbox.id, SandboxStatus.TERMINATED\n )\n db_session.commit()\n db_session.refresh(sandbox)\n # Fall through to TERMINATED handling below\n\n session_manager = SessionManager(db_session)\n llm_config = session_manager._get_llm_config(None, None)\n\n if sandbox.status in (SandboxStatus.SLEEPING, SandboxStatus.TERMINATED):\n # Mark as PROVISIONING before the long-running provision() call\n # so other requests know work is in progress\n update_sandbox_status__no_commit(\n db_session, sandbox.id, SandboxStatus.PROVISIONING\n )\n db_session.commit()\n\n sandbox_manager.provision(\n sandbox_id=sandbox.id,\n user_id=user.id,\n tenant_id=tenant_id,\n llm_config=llm_config,\n )\n\n # Mark as RUNNING after successful provision\n update_sandbox_status__no_commit(\n db_session, sandbox.id, SandboxStatus.RUNNING\n )\n db_session.commit()\n\n # 2. Check if session workspace needs to be loaded\n if sandbox.status == SandboxStatus.RUNNING:\n workspace_exists = sandbox_manager.session_workspace_exists(\n sandbox.id, session_id\n )\n\n if not workspace_exists:\n # Allocate port if not already set (needed for both snapshot restore and fresh setup)\n if not session.nextjs_port:\n session.nextjs_port = allocate_nextjs_port(db_session)\n # Commit port allocation before long-running operations\n db_session.commit()\n\n # Only Kubernetes backend supports snapshot restoration\n snapshot = None\n if SANDBOX_BACKEND == SandboxBackend.KUBERNETES:\n snapshot = get_latest_snapshot_for_session(db_session, session_id)\n\n if snapshot:\n try:\n sandbox_manager.restore_snapshot(\n sandbox_id=sandbox.id,\n session_id=session_id,\n snapshot_storage_path=snapshot.storage_path,\n tenant_id=tenant_id,\n nextjs_port=session.nextjs_port,\n llm_config=llm_config,\n use_demo_data=session.demo_data_enabled,\n )\n session.status = BuildSessionStatus.ACTIVE\n db_session.commit()\n except Exception as e:\n logger.error(\n f\"Snapshot restore failed for session {session_id}: {e}\"\n )\n session.nextjs_port = None\n db_session.commit()\n raise\n else:\n # No snapshot - set up fresh workspace\n sandbox_manager.setup_session_workspace(\n sandbox_id=sandbox.id,\n session_id=session_id,\n llm_config=llm_config,\n nextjs_port=session.nextjs_port,\n )\n session.status = BuildSessionStatus.ACTIVE\n db_session.commit()\n else:\n logger.warning(\n f\"Sandbox {sandbox.id} status is {sandbox.status} after re-provision, expected RUNNING\"\n )\n\n except Exception as e:\n logger.error(f\"Failed to restore session {session_id}: {e}\", exc_info=True)\n raise HTTPException(\n status_code=500,\n detail=f\"Failed to restore session: {e}\",\n )\n finally:\n if lock.owned():\n lock.release()\n\n # Update heartbeat to mark sandbox as active after successful restore\n update_sandbox_heartbeat(db_session, sandbox.id)\n\n base_response = SessionResponse.from_model(session, sandbox)\n return DetailedSessionResponse.from_session_response(\n base_response, session_loaded_in_sandbox=True\n )\n\n\n# =============================================================================\n# Artifact Endpoints\n# =============================================================================\n\n\n@router.get(\n \"/{session_id}/artifacts\",\n response_model=list[ArtifactResponse],\n)\ndef list_artifacts(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> list[dict]:\n \"\"\"List artifacts generated in the session.\"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n artifacts = session_manager.list_artifacts(session_id, user_id)\n if artifacts is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n return artifacts\n\n\n@router.get(\"/{session_id}/files\", response_model=DirectoryListing)\ndef list_directory(\n session_id: UUID,\n path: str = \"\",\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> DirectoryListing:\n \"\"\"\n List files and directories in the sandbox.\n\n Args:\n session_id: The session ID\n path: Relative path from sandbox root (empty string for root)\n\n Returns:\n DirectoryListing with sorted entries (directories first, then files)\n \"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n try:\n listing = session_manager.list_directory(session_id, user_id, path)\n except ValueError as e:\n error_message = str(e)\n if \"path traversal\" in error_message.lower():\n raise HTTPException(status_code=403, detail=\"Access denied\")\n elif \"not found\" in error_message.lower():\n raise HTTPException(status_code=404, detail=\"Directory not found\")\n elif \"not a directory\" in error_message.lower():\n raise HTTPException(status_code=400, detail=\"Path is not a directory\")\n raise HTTPException(status_code=400, detail=error_message)\n\n if listing is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n return listing\n\n\n@router.get(\"/{session_id}/artifacts/{path:path}\")\ndef download_artifact(\n session_id: UUID,\n path: str,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"Download a specific artifact file.\"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n try:\n result = session_manager.download_artifact(session_id, user_id, path)\n except ValueError as e:\n error_message = str(e)\n if (\n \"path traversal\" in error_message.lower()\n or \"access denied\" in error_message.lower()\n ):\n raise HTTPException(status_code=403, detail=\"Access denied\")\n elif \"directory\" in error_message.lower():\n raise HTTPException(status_code=400, detail=\"Cannot download directory\")\n raise HTTPException(status_code=400, detail=error_message)\n\n if result is None:\n raise HTTPException(status_code=404, detail=\"Artifact not found\")\n\n content, mime_type, filename = result\n\n # Handle Unicode filenames in Content-Disposition header\n # HTTP headers require Latin-1 encoding, so we use RFC 5987 for Unicode\n try:\n # Try Latin-1 encoding first (ASCII-compatible filenames)\n filename.encode(\"latin-1\")\n content_disposition = f'attachment; filename=\"{filename}\"'\n except UnicodeEncodeError:\n # Use RFC 5987 encoding for Unicode filenames\n from urllib.parse import quote\n\n encoded_filename = quote(filename, safe=\"\")\n content_disposition = f\"attachment; filename*=UTF-8''{encoded_filename}\"\n\n return Response(\n content=content,\n media_type=mime_type,\n headers={\n \"Content-Disposition\": content_disposition,\n },\n )\n\n\n@router.get(\"/{session_id}/export-docx/{path:path}\")\ndef export_docx(\n session_id: UUID,\n path: str,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"Export a markdown file as DOCX.\"\"\"\n session_manager = SessionManager(db_session)\n\n try:\n result = session_manager.export_docx(session_id, user.id, path)\n except ValueError as e:\n error_message = str(e)\n if (\n \"path traversal\" in error_message.lower()\n or \"access denied\" in error_message.lower()\n ):\n raise HTTPException(status_code=403, detail=\"Access denied\")\n raise HTTPException(status_code=400, detail=error_message)\n\n if result is None:\n raise HTTPException(status_code=404, detail=\"File not found\")\n\n docx_bytes, filename = result\n\n try:\n filename.encode(\"latin-1\")\n content_disposition = f'attachment; filename=\"{filename}\"'\n except UnicodeEncodeError:\n from urllib.parse import quote\n\n encoded_filename = quote(filename, safe=\"\")\n content_disposition = f\"attachment; filename*=UTF-8''{encoded_filename}\"\n\n return Response(\n content=docx_bytes,\n media_type=\"application/vnd.openxmlformats-officedocument.wordprocessingml.document\",\n headers={\"Content-Disposition\": content_disposition},\n )\n\n\n@router.get(\"/{session_id}/pptx-preview/{path:path}\")\ndef get_pptx_preview(\n session_id: UUID,\n path: str,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> PptxPreviewResponse:\n \"\"\"Generate slide image previews for a PPTX file.\"\"\"\n session_manager = SessionManager(db_session)\n\n try:\n result = session_manager.get_pptx_preview(session_id, user.id, path)\n except ValueError as e:\n error_message = str(e)\n if (\n \"path traversal\" in error_message.lower()\n or \"access denied\" in error_message.lower()\n ):\n raise HTTPException(status_code=403, detail=\"Access denied\")\n raise HTTPException(status_code=400, detail=error_message)\n\n if result is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n return PptxPreviewResponse(**result)\n\n\n@router.get(\"/{session_id}/webapp-info\", response_model=WebappInfo)\ndef get_webapp_info(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> WebappInfo:\n \"\"\"\n Get webapp information for a session.\n\n Returns whether a webapp exists, its URL, and the sandbox status.\n \"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n webapp_info = session_manager.get_webapp_info(session_id, user_id)\n\n if webapp_info is None:\n raise HTTPException(status_code=404, detail=\"Session not found\")\n\n return WebappInfo(**webapp_info)\n\n\n@router.get(\"/{session_id}/webapp-download\")\ndef download_webapp(\n session_id: UUID,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"\n Download the webapp directory as a zip file.\n\n Returns the entire outputs/web directory as a zip archive.\n \"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n result = session_manager.download_webapp_zip(session_id, user_id)\n\n if result is None:\n raise HTTPException(status_code=404, detail=\"Webapp not found\")\n\n zip_bytes, filename = result\n\n return Response(\n content=zip_bytes,\n media_type=\"application/zip\",\n headers={\n \"Content-Disposition\": f'attachment; filename=\"{filename}\"',\n },\n )\n\n\n@router.get(\"/{session_id}/download-directory/{path:path}\")\ndef download_directory(\n session_id: UUID,\n path: str,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"\n Download a directory as a zip file.\n\n Returns the specified directory as a zip archive.\n \"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n try:\n result = session_manager.download_directory(session_id, user_id, path)\n except ValueError as e:\n error_message = str(e)\n if \"path traversal\" in error_message.lower():\n raise HTTPException(status_code=403, detail=\"Access denied\")\n raise HTTPException(status_code=400, detail=error_message)\n\n if result is None:\n raise HTTPException(status_code=404, detail=\"Directory not found\")\n\n zip_bytes, filename = result\n\n return Response(\n content=zip_bytes,\n media_type=\"application/zip\",\n headers={\n \"Content-Disposition\": f'attachment; filename=\"{filename}\"',\n },\n )\n\n\n@router.post(\"/{session_id}/upload\", response_model=UploadResponse)\ndef upload_file_endpoint(\n session_id: UUID,\n file: UploadFile = File(...),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> UploadResponse:\n \"\"\"Upload a file to the session's sandbox.\n\n The file will be placed in the sandbox's attachments directory.\n \"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n if not file.filename:\n raise HTTPException(status_code=400, detail=\"File has no filename\")\n\n # Read file content (use sync file interface)\n content = file.file.read()\n\n # Validate file (extension, mime type, size)\n is_valid, error = validate_file(file.filename, file.content_type, len(content))\n if not is_valid:\n raise HTTPException(status_code=400, detail=error)\n\n # Sanitize filename\n safe_filename = sanitize_filename(file.filename)\n\n try:\n relative_path, _ = session_manager.upload_file(\n session_id=session_id,\n user_id=user_id,\n filename=safe_filename,\n content=content,\n )\n except UploadLimitExceededError as e:\n # Return 429 for limit exceeded errors\n raise HTTPException(status_code=429, detail=str(e))\n except ValueError as e:\n error_message = str(e)\n if \"not found\" in error_message.lower():\n raise HTTPException(status_code=404, detail=error_message)\n raise HTTPException(status_code=400, detail=error_message)\n\n return UploadResponse(\n filename=safe_filename,\n path=relative_path,\n size_bytes=len(content),\n )\n\n\n@router.delete(\"/{session_id}/files/{path:path}\", response_model=None)\ndef delete_file_endpoint(\n session_id: UUID,\n path: str,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> Response:\n \"\"\"Delete a file from the session's sandbox.\n\n Args:\n session_id: The session ID\n path: Relative path to the file (e.g., \"attachments/doc.pdf\")\n \"\"\"\n user_id: UUID = user.id\n session_manager = SessionManager(db_session)\n\n try:\n deleted = session_manager.delete_file(session_id, user_id, path)\n except ValueError as e:\n error_message = str(e)\n if \"path traversal\" in error_message.lower():\n raise HTTPException(status_code=403, detail=\"Access denied\")\n elif \"not found\" in error_message.lower():\n raise HTTPException(status_code=404, detail=error_message)\n elif \"directory\" in error_message.lower():\n raise HTTPException(status_code=400, detail=\"Cannot delete directory\")\n raise HTTPException(status_code=400, detail=error_message)\n\n if not deleted:\n raise HTTPException(status_code=404, detail=\"File not found\")\n\n return Response(status_code=204)\n" }, { "path": "backend/onyx/server/features/build/api/subscription_check.py", "content": "\"\"\"Subscription detection for Build Mode rate limiting.\"\"\"\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.app_configs import DEV_MODE\nfrom onyx.db.models import User\nfrom onyx.server.usage_limits import is_tenant_on_trial_fn\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\n\ndef is_user_subscribed(user: User, db_session: Session) -> bool: # noqa: ARG001\n \"\"\"\n Check if a user has an active subscription.\n\n For cloud (MULTI_TENANT=true):\n - Checks Stripe billing via control plane\n - Returns True if tenant is NOT on trial (subscribed = NOT on trial)\n\n For self-hosted (MULTI_TENANT=false):\n - Checks license metadata\n - Returns True if license status is ACTIVE\n\n Args:\n user: The user object (None for unauthenticated users)\n db_session: Database session\n\n Returns:\n True if user has active subscription, False otherwise\n \"\"\"\n if DEV_MODE:\n return True\n\n if user is None:\n return False\n\n if MULTI_TENANT:\n # Cloud: check Stripe billing via control plane\n tenant_id = get_current_tenant_id()\n try:\n on_trial = is_tenant_on_trial_fn(tenant_id)\n # Subscribed = NOT on trial\n return not on_trial\n except Exception as e:\n logger.warning(f\"Subscription check failed for tenant {tenant_id}: {e}\")\n # Default to non-subscribed (safer/more restrictive)\n return False\n\n return True\n" }, { "path": "backend/onyx/server/features/build/api/templates/webapp_hmr_fixer.js", "content": "(function () {\n var WEBAPP_BASE = \"__WEBAPP_BASE__\";\n var PROXIED_NEXT_PREFIX = WEBAPP_BASE + \"/_next/\";\n var PROXIED_HMR_PREFIX = WEBAPP_BASE + \"/_next/webpack-hmr\";\n var PROXIED_ALT_HMR_PREFIX = WEBAPP_BASE + \"/_next/hmr\";\n\n function isHmrWebSocketUrl(url) {\n if (!url) return false;\n try {\n var parsedUrl = new URL(String(url), window.location.href);\n return (\n parsedUrl.pathname.indexOf(\"/_next/webpack-hmr\") === 0 ||\n parsedUrl.pathname.indexOf(\"/_next/hmr\") === 0 ||\n parsedUrl.pathname.indexOf(PROXIED_HMR_PREFIX) === 0 ||\n parsedUrl.pathname.indexOf(PROXIED_ALT_HMR_PREFIX) === 0\n );\n } catch (e) {}\n if (typeof url === \"string\") {\n return (\n url.indexOf(\"/_next/webpack-hmr\") === 0 ||\n url.indexOf(\"/_next/hmr\") === 0 ||\n url.indexOf(PROXIED_HMR_PREFIX) === 0 ||\n url.indexOf(PROXIED_ALT_HMR_PREFIX) === 0\n );\n }\n return false;\n }\n\n function rewriteNextAssetUrl(url) {\n if (!url) return url;\n try {\n var parsedUrl = new URL(String(url), window.location.href);\n if (parsedUrl.pathname.indexOf(PROXIED_NEXT_PREFIX) === 0) {\n return parsedUrl.pathname + parsedUrl.search + parsedUrl.hash;\n }\n if (parsedUrl.pathname.indexOf(\"/_next/\") === 0) {\n return (\n WEBAPP_BASE + parsedUrl.pathname + parsedUrl.search + parsedUrl.hash\n );\n }\n } catch (e) {}\n if (typeof url === \"string\") {\n if (url.indexOf(PROXIED_NEXT_PREFIX) === 0) {\n return url;\n }\n if (url.indexOf(\"/_next/\") === 0) {\n return WEBAPP_BASE + url;\n }\n }\n return url;\n }\n\n function createEvent(eventType) {\n return typeof Event === \"function\"\n ? new Event(eventType)\n : { type: eventType };\n }\n\n function MockHmrWebSocket(url) {\n this.url = String(url);\n this.readyState = 1;\n this.bufferedAmount = 0;\n this.extensions = \"\";\n this.protocol = \"\";\n this.binaryType = \"blob\";\n this.onopen = null;\n this.onmessage = null;\n this.onerror = null;\n this.onclose = null;\n this._l = {};\n var socket = this;\n setTimeout(function () {\n socket._d(\"open\", createEvent(\"open\"));\n }, 0);\n }\n\n MockHmrWebSocket.CONNECTING = 0;\n MockHmrWebSocket.OPEN = 1;\n MockHmrWebSocket.CLOSING = 2;\n MockHmrWebSocket.CLOSED = 3;\n\n MockHmrWebSocket.prototype.addEventListener = function (eventType, callback) {\n (this._l[eventType] || (this._l[eventType] = [])).push(callback);\n };\n\n MockHmrWebSocket.prototype.removeEventListener = function (\n eventType,\n callback,\n ) {\n var listeners = this._l[eventType] || [];\n this._l[eventType] = listeners.filter(function (listener) {\n return listener !== callback;\n });\n };\n\n MockHmrWebSocket.prototype._d = function (eventType, eventValue) {\n var listeners = this._l[eventType] || [];\n for (var i = 0; i < listeners.length; i++) {\n listeners[i].call(this, eventValue);\n }\n var handler = this[\"on\" + eventType];\n if (typeof handler === \"function\") {\n handler.call(this, eventValue);\n }\n };\n\n MockHmrWebSocket.prototype.send = function () {};\n\n MockHmrWebSocket.prototype.close = function (code, reason) {\n if (this.readyState >= 2) return;\n this.readyState = 3;\n var closeEvent = createEvent(\"close\");\n closeEvent.code = code === undefined ? 1000 : code;\n closeEvent.reason = reason || \"\";\n closeEvent.wasClean = true;\n this._d(\"close\", closeEvent);\n };\n\n if (window.WebSocket) {\n var OriginalWebSocket = window.WebSocket;\n window.WebSocket = function (url, protocols) {\n if (isHmrWebSocketUrl(url)) {\n return new MockHmrWebSocket(rewriteNextAssetUrl(url));\n }\n return protocols === undefined\n ? new OriginalWebSocket(url)\n : new OriginalWebSocket(url, protocols);\n };\n window.WebSocket.prototype = OriginalWebSocket.prototype;\n Object.setPrototypeOf(window.WebSocket, OriginalWebSocket);\n [\"CONNECTING\", \"OPEN\", \"CLOSING\", \"CLOSED\"].forEach(function (stateKey) {\n window.WebSocket[stateKey] = OriginalWebSocket[stateKey];\n });\n }\n})();\n" }, { "path": "backend/onyx/server/features/build/api/templates/webapp_offline.html", "content": "\n\n \n \n \n \n Craft — Starting up\n \n \n \n
    \n
    \n
    \n
    \n
    \n crafting_table\n
    \n
    \n />\n Sandbox is asleep...\n
    \n
    \n

    \n Ask the owner to open their Craft session to wake it up.\n

    \n \n\n" }, { "path": "backend/onyx/server/features/build/api/user_library.py", "content": "\"\"\"API endpoints for User Library file management in Craft.\n\nThis module provides endpoints for uploading and managing raw binary files\n(xlsx, pptx, docx, csv, etc.) that are stored directly in S3 for sandbox access.\n\nFiles are stored at:\n s3://{bucket}/{tenant_id}/knowledge/{user_id}/user_library/{path}\n\nAnd synced to sandbox at:\n /workspace/files/user_library/{path}\n\nKnown Issues / TODOs:\n - Memory: Upload endpoints read entire file content into memory (up to 500MB).\n Should be refactored to stream uploads directly to S3 via multipart upload\n for better memory efficiency under concurrent load.\n - Transaction safety: Multi-file uploads are not atomic. If the endpoint fails\n mid-batch (e.g., file 3 of 5 exceeds storage quota), files 1-2 are already\n persisted to S3 and DB. A partial upload is not catastrophic but the response\n implies atomicity that doesn't exist.\n\"\"\"\n\nimport hashlib\nimport mimetypes\nimport re\nimport zipfile\nfrom datetime import datetime\nfrom datetime import timezone\nfrom io import BytesIO\nfrom typing import Any\n\nfrom fastapi import APIRouter\nfrom fastapi import Depends\nfrom fastapi import File\nfrom fastapi import Form\nfrom fastapi import HTTPException\nfrom fastapi import Query\nfrom fastapi import UploadFile\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.users import current_user\nfrom onyx.background.celery.versioned_apps.client import app as celery_app\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.connector_credential_pair import update_connector_credential_pair\nfrom onyx.db.document import upsert_document_by_connector_credential_pair\nfrom onyx.db.document import upsert_documents\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import User\nfrom onyx.document_index.interfaces import DocumentMetadata\nfrom onyx.server.features.build.configs import USER_LIBRARY_MAX_FILE_SIZE_BYTES\nfrom onyx.server.features.build.configs import USER_LIBRARY_MAX_FILES_PER_UPLOAD\nfrom onyx.server.features.build.configs import USER_LIBRARY_MAX_TOTAL_SIZE_BYTES\nfrom onyx.server.features.build.configs import USER_LIBRARY_SOURCE_DIR\nfrom onyx.server.features.build.db.user_library import get_or_create_craft_connector\nfrom onyx.server.features.build.db.user_library import get_user_storage_bytes\nfrom onyx.server.features.build.indexing.persistent_document_writer import (\n get_persistent_document_writer,\n)\nfrom onyx.server.features.build.indexing.persistent_document_writer import (\n PersistentDocumentWriter,\n)\nfrom onyx.server.features.build.indexing.persistent_document_writer import (\n S3PersistentDocumentWriter,\n)\nfrom onyx.server.features.build.utils import sanitize_filename as api_sanitize_filename\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import get_current_tenant_id\n\nlogger = setup_logger()\n\nrouter = APIRouter(prefix=\"/user-library\")\n\n\n# =============================================================================\n# Pydantic Models\n# =============================================================================\n\n\nclass LibraryEntryResponse(BaseModel):\n \"\"\"Response for a single library entry (file or directory).\"\"\"\n\n id: str # document_id\n name: str\n path: str\n is_directory: bool\n file_size: int | None\n mime_type: str | None\n sync_enabled: bool\n created_at: datetime\n children: list[\"LibraryEntryResponse\"] | None = None\n\n\nclass CreateDirectoryRequest(BaseModel):\n \"\"\"Request to create a virtual directory.\"\"\"\n\n name: str\n parent_path: str = \"/\"\n\n\nclass UploadResponse(BaseModel):\n \"\"\"Response after successful file upload.\"\"\"\n\n entries: list[LibraryEntryResponse]\n total_uploaded: int\n total_size_bytes: int\n\n\nclass ToggleSyncResponse(BaseModel):\n \"\"\"Response after toggling file sync.\"\"\"\n\n success: bool\n sync_enabled: bool\n\n\nclass DeleteFileResponse(BaseModel):\n \"\"\"Response after deleting a file.\"\"\"\n\n success: bool\n deleted: str\n\n\n# =============================================================================\n# Helper Functions\n# =============================================================================\n\n\ndef _sanitize_path(path: str) -> str:\n \"\"\"Sanitize a file path, removing traversal attempts and normalizing.\n\n Removes '..' and '.' segments and ensures the path starts with '/'.\n Only allows alphanumeric characters, hyphens, underscores, dots, spaces,\n and forward slashes. All other characters are stripped.\n \"\"\"\n parts = path.split(\"/\")\n sanitized_parts: list[str] = []\n for p in parts:\n if not p or p == \"..\" or p == \".\":\n continue\n # Strip any character not in the whitelist\n cleaned = re.sub(r\"[^a-zA-Z0-9\\-_. ]\", \"\", p)\n if cleaned:\n sanitized_parts.append(cleaned)\n return \"/\" + \"/\".join(sanitized_parts)\n\n\ndef _build_document_id(user_id: str, path: str) -> str:\n \"\"\"Build a document ID for a craft file.\n\n Deterministic: re-uploading the same file to the same path will produce the\n same document ID, allowing upsert to overwrite the previous record.\n\n Uses a hash of the path to avoid collisions from separator replacement\n (e.g., \"/a/b_c\" vs \"/a_b/c\" would collide with naive slash-to-underscore).\n \"\"\"\n path_hash = hashlib.sha256(path.encode()).hexdigest()[:16]\n return f\"CRAFT_FILE__{user_id}__{path_hash}\"\n\n\ndef _trigger_sandbox_sync(\n user_id: str, tenant_id: str, source: str | None = None\n) -> None:\n \"\"\"Trigger sandbox file sync task.\n\n Args:\n user_id: The user ID whose sandbox should be synced\n tenant_id: The tenant ID for S3 path construction\n source: Optional source type (e.g., \"user_library\"). If specified,\n only syncs that source's directory with --delete flag.\n \"\"\"\n celery_app.send_task(\n OnyxCeleryTask.SANDBOX_FILE_SYNC,\n kwargs={\"user_id\": user_id, \"tenant_id\": tenant_id, \"source\": source},\n queue=OnyxCeleryQueues.SANDBOX,\n )\n\n\ndef _validate_zip_contents(\n zip_file: zipfile.ZipFile,\n existing_usage: int,\n) -> None:\n \"\"\"Validate zip file contents before extraction.\n\n Checks file count limit and total decompressed size against storage quota.\n Raises HTTPException on validation failure.\n \"\"\"\n if len(zip_file.namelist()) > USER_LIBRARY_MAX_FILES_PER_UPLOAD:\n raise HTTPException(\n status_code=400,\n detail=f\"Zip contains too many files. Maximum is {USER_LIBRARY_MAX_FILES_PER_UPLOAD}.\",\n )\n\n # Zip bomb protection: check total decompressed size before extracting\n declared_total = sum(\n info.file_size for info in zip_file.infolist() if not info.is_dir()\n )\n if existing_usage + declared_total > USER_LIBRARY_MAX_TOTAL_SIZE_BYTES:\n raise HTTPException(\n status_code=400,\n detail=(\n f\"Zip decompressed size ({declared_total // (1024 * 1024)}MB) would exceed storage limit.\"\n ),\n )\n\n\ndef _verify_ownership_and_get_document(\n document_id: str,\n user: User,\n db_session: Session,\n) -> Any:\n \"\"\"Verify the user owns the document and return it.\n\n Raises HTTPException on authorization failure or if document not found.\n \"\"\"\n from onyx.db.document import get_document\n\n user_prefix = f\"CRAFT_FILE__{user.id}__\"\n if not document_id.startswith(user_prefix):\n raise HTTPException(\n status_code=403, detail=\"Not authorized to modify this file\"\n )\n\n doc = get_document(document_id, db_session)\n if doc is None:\n raise HTTPException(status_code=404, detail=\"File not found\")\n\n return doc\n\n\ndef _store_and_track_file(\n *,\n writer: \"PersistentDocumentWriter | S3PersistentDocumentWriter\",\n file_path: str,\n content: bytes,\n content_type: str | None,\n user_id: str,\n connector_id: int,\n credential_id: int,\n db_session: Session,\n) -> tuple[str, str]:\n \"\"\"Write a file to storage and upsert its document record.\n\n Returns:\n Tuple of (document_id, storage_key)\n \"\"\"\n storage_key = writer.write_raw_file(\n path=file_path,\n content=content,\n content_type=content_type,\n )\n\n doc_id = _build_document_id(user_id, file_path)\n doc_metadata = DocumentMetadata(\n connector_id=connector_id,\n credential_id=credential_id,\n document_id=doc_id,\n semantic_identifier=f\"{USER_LIBRARY_SOURCE_DIR}{file_path}\",\n first_link=storage_key,\n doc_metadata={\n \"storage_key\": storage_key,\n \"file_path\": file_path,\n \"file_size\": len(content),\n \"mime_type\": content_type,\n \"is_directory\": False,\n },\n )\n upsert_documents(db_session, [doc_metadata])\n upsert_document_by_connector_credential_pair(\n db_session, connector_id, credential_id, [doc_id]\n )\n\n return doc_id, storage_key\n\n\n# =============================================================================\n# API Endpoints\n# =============================================================================\n\n\n@router.get(\"/tree\")\ndef get_library_tree(\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> list[LibraryEntryResponse]:\n \"\"\"Get user's uploaded files as a tree structure.\n\n Returns all CRAFT_FILE documents for the user, organized hierarchically.\n \"\"\"\n from onyx.db.document import get_documents_by_source\n\n # Get CRAFT_FILE documents for this user (filtered at SQL level)\n user_docs = get_documents_by_source(\n db_session=db_session,\n source=DocumentSource.CRAFT_FILE,\n creator_id=user.id,\n )\n\n # Build tree structure\n entries: list[LibraryEntryResponse] = []\n now = datetime.now(timezone.utc)\n for doc in user_docs:\n doc_metadata = doc.doc_metadata or {}\n entries.append(\n LibraryEntryResponse(\n id=doc.id,\n name=doc.semantic_id.split(\"/\")[-1] if doc.semantic_id else \"unknown\",\n path=doc.semantic_id or \"\",\n is_directory=doc_metadata.get(\"is_directory\", False),\n file_size=doc_metadata.get(\"file_size\"),\n mime_type=doc_metadata.get(\"mime_type\"),\n sync_enabled=not doc_metadata.get(\"sync_disabled\", False),\n created_at=doc.last_modified or now,\n )\n )\n\n return entries\n\n\n@router.post(\"/upload\")\nasync def upload_files(\n files: list[UploadFile] = File(...),\n path: str = Form(\"/\"),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> UploadResponse:\n \"\"\"Upload files directly to S3 and track in PostgreSQL.\n\n Files are stored as raw binary (no text extraction) for access by\n the sandbox agent using Python libraries like openpyxl, python-pptx, etc.\n \"\"\"\n tenant_id = get_current_tenant_id()\n if tenant_id is None:\n raise HTTPException(status_code=500, detail=\"Tenant ID not found\")\n\n # Validate file count\n if len(files) > USER_LIBRARY_MAX_FILES_PER_UPLOAD:\n raise HTTPException(\n status_code=400,\n detail=f\"Too many files. Maximum is {USER_LIBRARY_MAX_FILES_PER_UPLOAD} per upload.\",\n )\n\n # Check cumulative storage usage\n existing_usage = get_user_storage_bytes(db_session, user.id)\n\n # Get or create connector\n connector_id, credential_id = get_or_create_craft_connector(db_session, user)\n\n # Get the persistent document writer\n writer = get_persistent_document_writer(\n user_id=str(user.id),\n tenant_id=tenant_id,\n )\n\n uploaded_entries: list[LibraryEntryResponse] = []\n total_size = 0\n now = datetime.now(timezone.utc)\n\n # Sanitize the base path\n base_path = _sanitize_path(path)\n\n for file in files:\n # TODO: Stream directly to S3 via multipart upload instead of reading\n # entire file into memory. With 500MB max file size, this can OOM under\n # concurrent uploads.\n content = await file.read()\n file_size = len(content)\n\n # Validate individual file size\n if file_size > USER_LIBRARY_MAX_FILE_SIZE_BYTES:\n raise HTTPException(\n status_code=400,\n detail=f\"File '{file.filename}' exceeds maximum size of {USER_LIBRARY_MAX_FILE_SIZE_BYTES // (1024 * 1024)}MB\",\n )\n\n # Validate cumulative storage (existing + this upload batch)\n total_size += file_size\n if existing_usage + total_size > USER_LIBRARY_MAX_TOTAL_SIZE_BYTES:\n raise HTTPException(\n status_code=400,\n detail=f\"Total storage would exceed maximum of {USER_LIBRARY_MAX_TOTAL_SIZE_BYTES // (1024 * 1024 * 1024)}GB\",\n )\n\n # Sanitize filename\n safe_filename = api_sanitize_filename(file.filename or \"unnamed\")\n file_path = f\"{base_path}/{safe_filename}\".replace(\"//\", \"/\")\n\n doc_id, _ = _store_and_track_file(\n writer=writer,\n file_path=file_path,\n content=content,\n content_type=file.content_type,\n user_id=str(user.id),\n connector_id=connector_id,\n credential_id=credential_id,\n db_session=db_session,\n )\n\n uploaded_entries.append(\n LibraryEntryResponse(\n id=doc_id,\n name=safe_filename,\n path=file_path,\n is_directory=False,\n file_size=file_size,\n mime_type=file.content_type,\n sync_enabled=True,\n created_at=now,\n )\n )\n\n # Mark connector as having succeeded (sets last_successful_index_time)\n # This allows the demo data toggle to be disabled\n update_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n status=ConnectorCredentialPairStatus.ACTIVE,\n net_docs=len(uploaded_entries),\n run_dt=now,\n )\n\n # Trigger sandbox sync for user_library source only\n _trigger_sandbox_sync(str(user.id), tenant_id, source=USER_LIBRARY_SOURCE_DIR)\n\n logger.info(\n f\"Uploaded {len(uploaded_entries)} files ({total_size} bytes) for user {user.id}\"\n )\n\n return UploadResponse(\n entries=uploaded_entries,\n total_uploaded=len(uploaded_entries),\n total_size_bytes=total_size,\n )\n\n\n@router.post(\"/upload-zip\")\nasync def upload_zip(\n file: UploadFile = File(...),\n path: str = Form(\"/\"),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> UploadResponse:\n \"\"\"Upload and extract a zip file, storing each extracted file to S3.\n\n Preserves the directory structure from the zip file.\n \"\"\"\n tenant_id = get_current_tenant_id()\n if tenant_id is None:\n raise HTTPException(status_code=500, detail=\"Tenant ID not found\")\n\n # Read zip content\n content = await file.read()\n if len(content) > USER_LIBRARY_MAX_TOTAL_SIZE_BYTES:\n raise HTTPException(\n status_code=400,\n detail=f\"Zip file exceeds maximum size of {USER_LIBRARY_MAX_TOTAL_SIZE_BYTES // (1024 * 1024 * 1024)}GB\",\n )\n\n # Check cumulative storage usage\n existing_usage = get_user_storage_bytes(db_session, user.id)\n\n # Get or create connector\n connector_id, credential_id = get_or_create_craft_connector(db_session, user)\n\n # Get the persistent document writer\n writer = get_persistent_document_writer(\n user_id=str(user.id),\n tenant_id=tenant_id,\n )\n\n uploaded_entries: list[LibraryEntryResponse] = []\n total_size = 0\n\n # Extract zip contents into a subfolder named after the zip file\n zip_name = api_sanitize_filename(file.filename or \"upload\")\n if zip_name.lower().endswith(\".zip\"):\n zip_name = zip_name[:-4]\n folder_path = f\"{_sanitize_path(path)}/{zip_name}\".replace(\"//\", \"/\")\n base_path = folder_path\n\n now = datetime.now(timezone.utc)\n\n # Track all directory paths we need to create records for\n directory_paths: set[str] = set()\n\n try:\n with zipfile.ZipFile(BytesIO(content), \"r\") as zip_file:\n _validate_zip_contents(zip_file, existing_usage)\n\n for zip_info in zip_file.infolist():\n # Skip hidden files and __MACOSX\n if (\n zip_info.filename.startswith(\"__MACOSX\")\n or \"/.\" in zip_info.filename\n ):\n continue\n\n # Skip directories - we'll create records from file paths below\n if zip_info.is_dir():\n continue\n\n # Read file content\n file_content = zip_file.read(zip_info.filename)\n file_size = len(file_content)\n\n # Validate individual file size\n if file_size > USER_LIBRARY_MAX_FILE_SIZE_BYTES:\n logger.warning(f\"Skipping '{zip_info.filename}' - exceeds max size\")\n continue\n\n total_size += file_size\n\n # Validate cumulative storage\n if existing_usage + total_size > USER_LIBRARY_MAX_TOTAL_SIZE_BYTES:\n raise HTTPException(\n status_code=400,\n detail=f\"Total storage would exceed maximum of {USER_LIBRARY_MAX_TOTAL_SIZE_BYTES // (1024 * 1024 * 1024)}GB\",\n )\n\n # Build path preserving zip structure\n sanitized_zip_path = _sanitize_path(zip_info.filename)\n file_path = f\"{base_path}{sanitized_zip_path}\".replace(\"//\", \"/\")\n file_name = file_path.split(\"/\")[-1]\n\n # Collect all intermediate directories for this file\n parts = file_path.split(\"/\")\n for i in range(\n 2, len(parts)\n ): # start at 2 to skip empty + first segment\n directory_paths.add(\"/\".join(parts[:i]))\n\n # Guess content type\n content_type, _ = mimetypes.guess_type(file_name)\n\n doc_id, _ = _store_and_track_file(\n writer=writer,\n file_path=file_path,\n content=file_content,\n content_type=content_type,\n user_id=str(user.id),\n connector_id=connector_id,\n credential_id=credential_id,\n db_session=db_session,\n )\n\n uploaded_entries.append(\n LibraryEntryResponse(\n id=doc_id,\n name=file_name,\n path=file_path,\n is_directory=False,\n file_size=file_size,\n mime_type=content_type,\n sync_enabled=True,\n created_at=now,\n )\n )\n\n except zipfile.BadZipFile:\n raise HTTPException(status_code=400, detail=\"Invalid zip file\")\n\n # Create directory document records so they appear in the tree view\n if directory_paths:\n dir_doc_ids: list[str] = []\n for dir_path in sorted(directory_paths):\n dir_doc_id = _build_document_id(str(user.id), dir_path)\n dir_doc_ids.append(dir_doc_id)\n dir_metadata = DocumentMetadata(\n connector_id=connector_id,\n credential_id=credential_id,\n document_id=dir_doc_id,\n semantic_identifier=f\"{USER_LIBRARY_SOURCE_DIR}{dir_path}\",\n first_link=\"\",\n doc_metadata={\"is_directory\": True},\n )\n upsert_documents(db_session, [dir_metadata])\n upsert_document_by_connector_credential_pair(\n db_session, connector_id, credential_id, dir_doc_ids\n )\n\n # Mark connector as having succeeded (sets last_successful_index_time)\n # This allows the demo data toggle to be disabled\n update_connector_credential_pair(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential_id,\n status=ConnectorCredentialPairStatus.ACTIVE,\n net_docs=len(uploaded_entries),\n run_dt=now,\n )\n\n # Trigger sandbox sync for user_library source only\n _trigger_sandbox_sync(str(user.id), tenant_id, source=USER_LIBRARY_SOURCE_DIR)\n\n logger.info(\n f\"Extracted {len(uploaded_entries)} files ({total_size} bytes) from zip for user {user.id}\"\n )\n\n return UploadResponse(\n entries=uploaded_entries,\n total_uploaded=len(uploaded_entries),\n total_size_bytes=total_size,\n )\n\n\n@router.post(\"/directories\")\ndef create_directory(\n request: CreateDirectoryRequest,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> LibraryEntryResponse:\n \"\"\"Create a virtual directory.\n\n Directories are tracked as documents with is_directory=True.\n No S3 object is created (S3 doesn't have real directories).\n \"\"\"\n # Get or create connector\n connector_id, credential_id = get_or_create_craft_connector(db_session, user)\n\n # Build path\n parent_path = _sanitize_path(request.parent_path)\n safe_name = api_sanitize_filename(request.name)\n dir_path = f\"{parent_path}/{safe_name}\".replace(\"//\", \"/\")\n\n # Track in document table\n doc_id = _build_document_id(str(user.id), dir_path)\n doc_metadata = DocumentMetadata(\n connector_id=connector_id,\n credential_id=credential_id,\n document_id=doc_id,\n semantic_identifier=f\"{USER_LIBRARY_SOURCE_DIR}{dir_path}\",\n first_link=\"\",\n doc_metadata={\n \"is_directory\": True,\n },\n )\n upsert_documents(db_session, [doc_metadata])\n upsert_document_by_connector_credential_pair(\n db_session, connector_id, credential_id, [doc_id]\n )\n db_session.commit()\n\n return LibraryEntryResponse(\n id=doc_id,\n name=safe_name,\n path=dir_path,\n is_directory=True,\n file_size=None,\n mime_type=None,\n sync_enabled=True,\n created_at=datetime.now(timezone.utc),\n )\n\n\n@router.patch(\"/files/{document_id}/toggle\")\ndef toggle_file_sync(\n document_id: str,\n enabled: bool = Query(...),\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> ToggleSyncResponse:\n \"\"\"Enable/disable syncing a file to sandboxes.\n\n When sync is disabled, the file's metadata is updated with sync_disabled=True.\n The sandbox sync task will exclude these files when syncing to the sandbox.\n\n If the item is a directory, all children are also toggled.\n \"\"\"\n from onyx.db.document import get_documents_by_source\n from onyx.db.document import update_document_metadata__no_commit\n\n tenant_id = get_current_tenant_id()\n if tenant_id is None:\n raise HTTPException(status_code=500, detail=\"Tenant ID not found\")\n\n doc = _verify_ownership_and_get_document(document_id, user, db_session)\n\n # Update metadata for this document\n new_metadata = dict(doc.doc_metadata or {})\n new_metadata[\"sync_disabled\"] = not enabled\n update_document_metadata__no_commit(db_session, document_id, new_metadata)\n\n # If this is a directory, also toggle all children\n doc_metadata = doc.doc_metadata or {}\n if doc_metadata.get(\"is_directory\"):\n folder_path = doc.semantic_id\n if folder_path:\n all_docs = get_documents_by_source(\n db_session=db_session,\n source=DocumentSource.CRAFT_FILE,\n creator_id=user.id,\n )\n for child_doc in all_docs:\n if child_doc.semantic_id and child_doc.semantic_id.startswith(\n folder_path + \"/\"\n ):\n child_metadata = dict(child_doc.doc_metadata or {})\n child_metadata[\"sync_disabled\"] = not enabled\n update_document_metadata__no_commit(\n db_session, child_doc.id, child_metadata\n )\n\n db_session.commit()\n\n return ToggleSyncResponse(success=True, sync_enabled=enabled)\n\n\n@router.delete(\"/files/{document_id}\")\ndef delete_file(\n document_id: str,\n user: User = Depends(current_user),\n db_session: Session = Depends(get_session),\n) -> DeleteFileResponse:\n \"\"\"Delete a file from both S3 and the document table.\"\"\"\n from onyx.db.document import delete_document_by_id__no_commit\n\n tenant_id = get_current_tenant_id()\n if tenant_id is None:\n raise HTTPException(status_code=500, detail=\"Tenant ID not found\")\n\n doc = _verify_ownership_and_get_document(document_id, user, db_session)\n\n # Delete from storage if it's a file (not directory)\n doc_metadata = doc.doc_metadata or {}\n if not doc_metadata.get(\"is_directory\"):\n file_path = doc_metadata.get(\"file_path\")\n if file_path:\n writer = get_persistent_document_writer(\n user_id=str(user.id),\n tenant_id=tenant_id,\n )\n try:\n if isinstance(writer, S3PersistentDocumentWriter):\n writer.delete_raw_file_by_path(file_path)\n else:\n writer.delete_raw_file(file_path)\n except Exception as e:\n logger.warning(f\"Failed to delete file at path {file_path}: {e}\")\n else:\n # Fallback for documents created before file_path was stored\n storage_key = doc_metadata.get(\"storage_key\") or doc_metadata.get(\"s3_key\")\n if storage_key:\n writer = get_persistent_document_writer(\n user_id=str(user.id),\n tenant_id=tenant_id,\n )\n try:\n if isinstance(writer, S3PersistentDocumentWriter):\n writer.delete_raw_file(storage_key)\n else:\n logger.warning(\n f\"Cannot delete file in local mode without file_path: {document_id}\"\n )\n except Exception as e:\n logger.warning(\n f\"Failed to delete storage object {storage_key}: {e}\"\n )\n\n # Delete from document table\n delete_document_by_id__no_commit(db_session, document_id)\n db_session.commit()\n\n # Trigger sync to apply changes\n _trigger_sandbox_sync(str(user.id), tenant_id, source=USER_LIBRARY_SOURCE_DIR)\n\n return DeleteFileResponse(success=True, deleted=document_id)\n" }, { "path": "backend/onyx/server/features/build/configs.py", "content": "import os\nfrom enum import Enum\nfrom pathlib import Path\n\n\nclass SandboxBackend(str, Enum):\n \"\"\"Backend mode for sandbox operations.\n\n LOCAL: Development mode - no snapshots, no automatic cleanup\n KUBERNETES: Production mode - full snapshots and cleanup\n \"\"\"\n\n LOCAL = \"local\"\n KUBERNETES = \"kubernetes\"\n\n\n# Sandbox backend mode (controls snapshot and cleanup behavior)\n# \"local\" = no snapshots, no cleanup (for development)\n# \"kubernetes\" = full snapshots and cleanup (for production)\nSANDBOX_BACKEND = SandboxBackend(os.environ.get(\"SANDBOX_BACKEND\", \"local\"))\n\n# Base directory path for persistent document storage (local filesystem)\n# Example: /var/onyx/file-system or /app/file-system\nPERSISTENT_DOCUMENT_STORAGE_PATH = os.environ.get(\n \"PERSISTENT_DOCUMENT_STORAGE_PATH\", \"/app/file-system\"\n)\n\n# Demo Data Path\n# Local: Source tree path (relative to this file)\n# Kubernetes: Baked into container image at /workspace/demo_data\n_THIS_FILE = Path(__file__)\nDEMO_DATA_PATH = str(\n _THIS_FILE.parent / \"sandbox\" / \"kubernetes\" / \"docker\" / \"demo_data\"\n)\n\n# Sandbox filesystem paths\nSANDBOX_BASE_PATH = os.environ.get(\"SANDBOX_BASE_PATH\", \"/tmp/onyx-sandboxes\")\nOUTPUTS_TEMPLATE_PATH = os.environ.get(\"OUTPUTS_TEMPLATE_PATH\", \"/templates/outputs\")\nVENV_TEMPLATE_PATH = os.environ.get(\"VENV_TEMPLATE_PATH\", \"/templates/venv\")\n\n# Sandbox agent configuration\nSANDBOX_AGENT_COMMAND = os.environ.get(\"SANDBOX_AGENT_COMMAND\", \"opencode\").split()\n\n# OpenCode disabled tools (comma-separated list)\n# Available tools: bash, edit, write, read, grep, glob, list, lsp, patch,\n# skill, todowrite, todoread, webfetch, question\n# Example: \"question,webfetch\" to disable user questions and web fetching\n_disabled_tools_str = os.environ.get(\"OPENCODE_DISABLED_TOOLS\", \"question\")\nOPENCODE_DISABLED_TOOLS: list[str] = [\n t.strip() for t in _disabled_tools_str.split(\",\") if t.strip()\n]\n\n# Sandbox lifecycle configuration\nSANDBOX_IDLE_TIMEOUT_SECONDS = int(\n os.environ.get(\"SANDBOX_IDLE_TIMEOUT_SECONDS\", \"3600\")\n)\nSANDBOX_MAX_CONCURRENT_PER_ORG = int(\n os.environ.get(\"SANDBOX_MAX_CONCURRENT_PER_ORG\", \"10\")\n)\n\n# Sandbox snapshot storage\nSANDBOX_SNAPSHOTS_BUCKET = os.environ.get(\n \"SANDBOX_SNAPSHOTS_BUCKET\", \"sandbox-snapshots\"\n)\n\n# Next.js preview server port range\nSANDBOX_NEXTJS_PORT_START = int(os.environ.get(\"SANDBOX_NEXTJS_PORT_START\", \"3010\"))\nSANDBOX_NEXTJS_PORT_END = int(os.environ.get(\"SANDBOX_NEXTJS_PORT_END\", \"3100\"))\n\n# File upload configuration\nMAX_UPLOAD_FILE_SIZE_MB = int(os.environ.get(\"BUILD_MAX_UPLOAD_FILE_SIZE_MB\", \"50\"))\nMAX_UPLOAD_FILE_SIZE_BYTES = MAX_UPLOAD_FILE_SIZE_MB * 1024 * 1024\nMAX_UPLOAD_FILES_PER_SESSION = int(\n os.environ.get(\"BUILD_MAX_UPLOAD_FILES_PER_SESSION\", \"20\")\n)\nMAX_TOTAL_UPLOAD_SIZE_MB = int(os.environ.get(\"BUILD_MAX_TOTAL_UPLOAD_SIZE_MB\", \"200\"))\nMAX_TOTAL_UPLOAD_SIZE_BYTES = MAX_TOTAL_UPLOAD_SIZE_MB * 1024 * 1024\nATTACHMENTS_DIRECTORY = \"attachments\"\n\n# ============================================================================\n# Kubernetes Sandbox Configuration\n# Only used when SANDBOX_BACKEND = \"kubernetes\"\n# ============================================================================\n\n# Namespace where sandbox pods are created\nSANDBOX_NAMESPACE = os.environ.get(\"SANDBOX_NAMESPACE\", \"onyx-sandboxes\")\n\n# Container image for sandbox pods\n# Should include Next.js template, opencode CLI, and demo_data zip\nSANDBOX_CONTAINER_IMAGE = os.environ.get(\n \"SANDBOX_CONTAINER_IMAGE\", \"onyxdotapp/sandbox:v0.1.5\"\n)\n\n# S3 bucket for sandbox file storage (snapshots, knowledge files, uploads)\n# Path structure: s3://{bucket}/{tenant_id}/snapshots/{session_id}/{snapshot_id}.tar.gz\n# s3://{bucket}/{tenant_id}/knowledge/{user_id}/\n# s3://{bucket}/{tenant_id}/uploads/{session_id}/\nSANDBOX_S3_BUCKET = os.environ.get(\"SANDBOX_S3_BUCKET\", \"onyx-sandbox-files\")\n\n# Service account for sandbox pods (NO IRSA - no AWS API access)\nSANDBOX_SERVICE_ACCOUNT_NAME = os.environ.get(\n \"SANDBOX_SERVICE_ACCOUNT_NAME\", \"sandbox-runner\"\n)\n\n# Service account for init container (has IRSA for S3 access)\nSANDBOX_FILE_SYNC_SERVICE_ACCOUNT = os.environ.get(\n \"SANDBOX_FILE_SYNC_SERVICE_ACCOUNT\", \"sandbox-file-sync\"\n)\n\nENABLE_CRAFT = os.environ.get(\"ENABLE_CRAFT\", \"false\").lower() == \"true\"\n\n# ============================================================================\n# SSE Streaming Configuration\n# ============================================================================\n\n# SSE keepalive interval in seconds - send keepalive comment if no events\nSSE_KEEPALIVE_INTERVAL = float(os.environ.get(\"SSE_KEEPALIVE_INTERVAL\", \"15.0\"))\n\n# ============================================================================\n# ACP (Agent Communication Protocol) Configuration\n# ============================================================================\n\n# Timeout for ACP message processing in seconds\n# This is the maximum time to wait for a complete response from the agent\nACP_MESSAGE_TIMEOUT = float(os.environ.get(\"ACP_MESSAGE_TIMEOUT\", \"900.0\"))\n\n# ============================================================================\n# Rate Limiting Configuration\n# ============================================================================\n\n# Base rate limit for paid/subscribed users (messages per week)\n# Free users always get 5 messages total (not configurable)\n# Per-user overrides are managed via PostHog feature flag \"craft-has-usage-limits\"\nCRAFT_PAID_USER_RATE_LIMIT = int(os.environ.get(\"CRAFT_PAID_USER_RATE_LIMIT\", \"25\"))\n\n# ============================================================================\n# User Library Configuration\n# For user-uploaded raw files (xlsx, pptx, docx, etc.) in Craft\n# ============================================================================\n\n# Maximum size per file in MB (default 500MB)\nUSER_LIBRARY_MAX_FILE_SIZE_MB = int(\n os.environ.get(\"USER_LIBRARY_MAX_FILE_SIZE_MB\", \"500\")\n)\nUSER_LIBRARY_MAX_FILE_SIZE_BYTES = USER_LIBRARY_MAX_FILE_SIZE_MB * 1024 * 1024\n\n# Maximum total storage per user in GB (default 10GB)\nUSER_LIBRARY_MAX_TOTAL_SIZE_GB = int(\n os.environ.get(\"USER_LIBRARY_MAX_TOTAL_SIZE_GB\", \"10\")\n)\nUSER_LIBRARY_MAX_TOTAL_SIZE_BYTES = USER_LIBRARY_MAX_TOTAL_SIZE_GB * 1024 * 1024 * 1024\n\n# Maximum files per single upload request (default 100)\nUSER_LIBRARY_MAX_FILES_PER_UPLOAD = int(\n os.environ.get(\"USER_LIBRARY_MAX_FILES_PER_UPLOAD\", \"100\")\n)\n\n# String constants for User Library entities\nUSER_LIBRARY_CONNECTOR_NAME = \"User Library\"\nUSER_LIBRARY_CREDENTIAL_NAME = \"User Library Credential\"\nUSER_LIBRARY_SOURCE_DIR = \"user_library\"\n" }, { "path": "backend/onyx/server/features/build/db/__init__.py", "content": "# Database operations for the build feature\n" }, { "path": "backend/onyx/server/features/build/db/build_session.py", "content": "\"\"\"Database operations for Build Mode sessions.\"\"\"\n\nfrom datetime import datetime\nfrom typing import Any\nfrom uuid import UUID\n\nfrom sqlalchemy import desc\nfrom sqlalchemy import exists\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import selectinload\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.enums import BuildSessionStatus\nfrom onyx.db.enums import SandboxStatus\nfrom onyx.db.enums import SharingScope\nfrom onyx.db.models import Artifact\nfrom onyx.db.models import BuildMessage\nfrom onyx.db.models import BuildSession\nfrom onyx.db.models import LLMProvider as LLMProviderModel\nfrom onyx.db.models import Sandbox\nfrom onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_END\nfrom onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_START\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef create_build_session__no_commit(\n user_id: UUID,\n db_session: Session,\n name: str | None = None,\n demo_data_enabled: bool = True,\n) -> BuildSession:\n \"\"\"Create a new build session for the given user.\n\n NOTE: This function uses flush() instead of commit(). The caller is\n responsible for committing the transaction when ready.\n\n Args:\n user_id: The user ID\n db_session: Database session\n name: Optional session name\n demo_data_enabled: Whether this session uses demo data (default True)\n \"\"\"\n session = BuildSession(\n user_id=user_id,\n name=name,\n status=BuildSessionStatus.ACTIVE,\n demo_data_enabled=demo_data_enabled,\n )\n db_session.add(session)\n db_session.flush()\n\n logger.info(\n f\"Created build session {session.id} for user {user_id} (demo_data={demo_data_enabled})\"\n )\n return session\n\n\ndef get_build_session(\n session_id: UUID,\n user_id: UUID,\n db_session: Session,\n) -> BuildSession | None:\n \"\"\"Get a build session by ID, ensuring it belongs to the user.\"\"\"\n return (\n db_session.query(BuildSession)\n .filter(\n BuildSession.id == session_id,\n BuildSession.user_id == user_id,\n )\n .one_or_none()\n )\n\n\ndef get_user_build_sessions(\n user_id: UUID,\n db_session: Session,\n limit: int = 100,\n) -> list[BuildSession]:\n \"\"\"Get all build sessions for a user that have at least one message.\n\n Excludes empty (pre-provisioned) sessions from the listing.\n \"\"\"\n # Subquery to check if session has any messages\n has_messages = exists().where(BuildMessage.session_id == BuildSession.id)\n\n return (\n db_session.query(BuildSession)\n .filter(\n BuildSession.user_id == user_id,\n has_messages, # Only sessions with messages\n )\n .order_by(desc(BuildSession.created_at))\n .limit(limit)\n .all()\n )\n\n\ndef get_empty_session_for_user(\n user_id: UUID,\n db_session: Session,\n demo_data_enabled: bool | None = None,\n) -> BuildSession | None:\n \"\"\"Get an empty (pre-provisioned) session for the user if one exists.\n\n Returns a session with no messages, or None if all sessions have messages.\n\n Args:\n user_id: The user ID\n db_session: Database session\n demo_data_enabled: Match sessions with this demo_data setting.\n If None, matches any session regardless of setting.\n \"\"\"\n # Subquery to check if session has any messages\n has_messages = exists().where(BuildMessage.session_id == BuildSession.id)\n\n query = db_session.query(BuildSession).filter(\n BuildSession.user_id == user_id,\n ~has_messages, # Sessions with no messages only\n )\n\n if demo_data_enabled is not None:\n query = query.filter(BuildSession.demo_data_enabled == demo_data_enabled)\n\n return query.first()\n\n\ndef update_session_activity(\n session_id: UUID,\n db_session: Session,\n) -> None:\n \"\"\"Update the last activity timestamp for a session.\"\"\"\n session = (\n db_session.query(BuildSession)\n .filter(BuildSession.id == session_id)\n .one_or_none()\n )\n if session:\n session.last_activity_at = datetime.utcnow()\n db_session.commit()\n\n\ndef update_session_status(\n session_id: UUID,\n status: BuildSessionStatus,\n db_session: Session,\n) -> None:\n \"\"\"Update the status of a build session.\"\"\"\n session = (\n db_session.query(BuildSession)\n .filter(BuildSession.id == session_id)\n .one_or_none()\n )\n if session:\n session.status = status\n db_session.commit()\n logger.info(f\"Updated build session {session_id} status to {status}\")\n\n\ndef set_build_session_sharing_scope(\n session_id: UUID,\n user_id: UUID,\n sharing_scope: SharingScope,\n db_session: Session,\n) -> BuildSession | None:\n \"\"\"Set the sharing scope of a build session.\n\n Only the session owner can change this setting.\n Returns the updated session, or None if not found/unauthorized.\n \"\"\"\n session = get_build_session(session_id, user_id, db_session)\n if not session:\n return None\n session.sharing_scope = sharing_scope\n db_session.commit()\n logger.info(f\"Set build session {session_id} sharing_scope={sharing_scope}\")\n return session\n\n\ndef delete_build_session__no_commit(\n session_id: UUID,\n user_id: UUID,\n db_session: Session,\n) -> bool:\n \"\"\"Delete a build session and all related data.\n\n NOTE: This function uses flush() instead of commit(). The caller is\n responsible for committing the transaction when ready.\n \"\"\"\n session = get_build_session(session_id, user_id, db_session)\n if not session:\n return False\n\n db_session.delete(session)\n db_session.flush()\n logger.info(f\"Deleted build session {session_id}\")\n return True\n\n\n# Sandbox operations\n# NOTE: Most sandbox operations have moved to sandbox.py\n# These remain here for convenience in session-related workflows\n\n\ndef update_sandbox_status(\n sandbox_id: UUID,\n status: SandboxStatus,\n db_session: Session,\n container_id: str | None = None,\n) -> None:\n \"\"\"Update the status of a sandbox.\"\"\"\n sandbox = db_session.query(Sandbox).filter(Sandbox.id == sandbox_id).one_or_none()\n if sandbox:\n sandbox.status = status\n if container_id is not None:\n sandbox.container_id = container_id\n sandbox.last_heartbeat = datetime.utcnow()\n db_session.commit()\n logger.info(f\"Updated sandbox {sandbox_id} status to {status}\")\n\n\ndef update_sandbox_heartbeat(\n sandbox_id: UUID,\n db_session: Session,\n) -> None:\n \"\"\"Update the heartbeat timestamp for a sandbox.\"\"\"\n sandbox = db_session.query(Sandbox).filter(Sandbox.id == sandbox_id).one_or_none()\n if sandbox:\n sandbox.last_heartbeat = datetime.utcnow()\n db_session.commit()\n\n\n# Artifact operations\ndef create_artifact(\n session_id: UUID,\n artifact_type: str,\n path: str,\n name: str,\n db_session: Session,\n) -> Artifact:\n \"\"\"Create a new artifact record.\"\"\"\n artifact = Artifact(\n session_id=session_id,\n type=artifact_type,\n path=path,\n name=name,\n )\n db_session.add(artifact)\n db_session.commit()\n db_session.refresh(artifact)\n\n logger.info(f\"Created artifact {artifact.id} for session {session_id}\")\n return artifact\n\n\ndef get_session_artifacts(\n session_id: UUID,\n db_session: Session,\n) -> list[Artifact]:\n \"\"\"Get all artifacts for a session.\"\"\"\n return (\n db_session.query(Artifact)\n .filter(Artifact.session_id == session_id)\n .order_by(desc(Artifact.created_at))\n .all()\n )\n\n\ndef update_artifact(\n artifact_id: UUID,\n db_session: Session,\n path: str | None = None,\n name: str | None = None,\n) -> None:\n \"\"\"Update artifact metadata.\"\"\"\n artifact = (\n db_session.query(Artifact).filter(Artifact.id == artifact_id).one_or_none()\n )\n if artifact:\n if path is not None:\n artifact.path = path\n if name is not None:\n artifact.name = name\n artifact.updated_at = datetime.utcnow()\n db_session.commit()\n logger.info(f\"Updated artifact {artifact_id}\")\n\n\n# Message operations\ndef create_message(\n session_id: UUID,\n message_type: MessageType,\n turn_index: int,\n message_metadata: dict[str, Any],\n db_session: Session,\n) -> BuildMessage:\n \"\"\"Create a new message in a build session.\n\n All message data is stored in message_metadata as JSON.\n\n Args:\n session_id: Session UUID\n message_type: Type of message (USER, ASSISTANT, SYSTEM)\n turn_index: 0-indexed user message number this message belongs to\n message_metadata: Required structured data (the raw ACP packet JSON)\n db_session: Database session\n \"\"\"\n message = BuildMessage(\n session_id=session_id,\n turn_index=turn_index,\n type=message_type,\n message_metadata=message_metadata,\n )\n db_session.add(message)\n db_session.commit()\n db_session.refresh(message)\n\n logger.info(\n f\"Created {message_type.value} message {message.id} for session {session_id} \"\n f\"turn={turn_index} type={message_metadata.get('type')}\"\n )\n return message\n\n\ndef update_message(\n message_id: UUID,\n message_metadata: dict[str, Any],\n db_session: Session,\n) -> BuildMessage | None:\n \"\"\"Update an existing message's metadata.\n\n Used for upserting agent_plan_update messages.\n\n Args:\n message_id: The message UUID to update\n message_metadata: New metadata to set\n db_session: Database session\n\n Returns:\n Updated BuildMessage or None if not found\n \"\"\"\n message = (\n db_session.query(BuildMessage).filter(BuildMessage.id == message_id).first()\n )\n if message is None:\n return None\n\n message.message_metadata = message_metadata\n db_session.commit()\n db_session.refresh(message)\n\n logger.info(\n f\"Updated message {message_id} metadata type={message_metadata.get('type')}\"\n )\n return message\n\n\ndef upsert_agent_plan(\n session_id: UUID,\n turn_index: int,\n plan_metadata: dict[str, Any],\n db_session: Session,\n existing_plan_id: UUID | None = None,\n) -> BuildMessage:\n \"\"\"Upsert an agent plan - update if exists, create if not.\n\n Each session/turn should only have one agent_plan_update message.\n This function updates the existing plan message or creates a new one.\n\n Args:\n session_id: Session UUID\n turn_index: Current turn index\n plan_metadata: The agent_plan_update packet data\n db_session: Database session\n existing_plan_id: ID of existing plan message to update (if known)\n\n Returns:\n The created or updated BuildMessage\n \"\"\"\n if existing_plan_id:\n # Fast path: we know the plan ID\n updated = update_message(existing_plan_id, plan_metadata, db_session)\n if updated:\n return updated\n\n # Check if a plan already exists for this session/turn\n existing_plan = (\n db_session.query(BuildMessage)\n .filter(\n BuildMessage.session_id == session_id,\n BuildMessage.turn_index == turn_index,\n BuildMessage.message_metadata[\"type\"].astext == \"agent_plan_update\",\n )\n .first()\n )\n\n if existing_plan:\n existing_plan.message_metadata = plan_metadata\n db_session.commit()\n db_session.refresh(existing_plan)\n logger.info(\n f\"Updated agent_plan_update message {existing_plan.id} for session {session_id}\"\n )\n return existing_plan\n\n # Create new plan message\n return create_message(\n session_id=session_id,\n message_type=MessageType.ASSISTANT,\n turn_index=turn_index,\n message_metadata=plan_metadata,\n db_session=db_session,\n )\n\n\ndef get_session_messages(\n session_id: UUID,\n db_session: Session,\n) -> list[BuildMessage]:\n \"\"\"Get all messages for a session, ordered by turn index and creation time.\"\"\"\n return (\n db_session.query(BuildMessage)\n .filter(BuildMessage.session_id == session_id)\n .order_by(BuildMessage.turn_index, BuildMessage.created_at)\n .all()\n )\n\n\ndef _is_port_available(port: int) -> bool:\n \"\"\"Check if a port is available by attempting to bind to it.\n\n Checks both IPv4 and IPv6 wildcard addresses to properly detect\n if anything is listening on the port, regardless of address family.\n \"\"\"\n import socket\n\n logger.debug(f\"Checking if port {port} is available\")\n\n # Check IPv4 wildcard (0.0.0.0) - this will detect any IPv4 listener\n try:\n with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:\n sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)\n sock.bind((\"0.0.0.0\", port))\n logger.debug(f\"Port {port} IPv4 wildcard bind successful\")\n except OSError as e:\n logger.debug(f\"Port {port} IPv4 wildcard not available: {e}\")\n return False\n\n # Check IPv6 wildcard (::) - this will detect any IPv6 listener\n try:\n with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock:\n sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)\n # IPV6_V6ONLY must be False to allow dual-stack behavior\n sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)\n sock.bind((\"::\", port))\n logger.debug(f\"Port {port} IPv6 wildcard bind successful\")\n except OSError as e:\n logger.debug(f\"Port {port} IPv6 wildcard not available: {e}\")\n return False\n\n logger.debug(f\"Port {port} is available\")\n return True\n\n\ndef allocate_nextjs_port(db_session: Session) -> int:\n \"\"\"Allocate an available port for a new session.\n\n Finds the first available port in the configured range by checking\n both database allocations and system-level port availability.\n\n Args:\n db_session: Database session for querying allocated ports\n\n Returns:\n An available port number\n\n Raises:\n RuntimeError: If no ports are available in the configured range\n \"\"\"\n from onyx.db.models import BuildSession\n\n # Get all currently allocated ports from active sessions\n allocated_ports = set(\n db_session.query(BuildSession.nextjs_port)\n .filter(BuildSession.nextjs_port.isnot(None))\n .all()\n )\n allocated_ports = {port[0] for port in allocated_ports if port[0] is not None}\n\n # Find first port that's not in DB and not currently bound\n for port in range(SANDBOX_NEXTJS_PORT_START, SANDBOX_NEXTJS_PORT_END):\n if port not in allocated_ports and _is_port_available(port):\n return port\n\n raise RuntimeError(\n f\"No available ports in range [{SANDBOX_NEXTJS_PORT_START}, {SANDBOX_NEXTJS_PORT_END})\"\n )\n\n\ndef mark_user_sessions_idle__no_commit(db_session: Session, user_id: UUID) -> int:\n \"\"\"Mark all ACTIVE sessions for a user as IDLE.\n\n Called when a sandbox goes to sleep so the frontend knows these sessions\n need restoration before they can be used again.\n\n Args:\n db_session: Database session\n user_id: The user whose sessions should be marked idle\n\n Returns:\n Number of sessions updated\n \"\"\"\n result = (\n db_session.query(BuildSession)\n .filter(\n BuildSession.user_id == user_id,\n BuildSession.status == BuildSessionStatus.ACTIVE,\n )\n .update({BuildSession.status: BuildSessionStatus.IDLE})\n )\n db_session.flush()\n logger.info(f\"Marked {result} sessions as IDLE for user {user_id}\")\n return result\n\n\ndef clear_nextjs_ports_for_user(db_session: Session, user_id: UUID) -> int:\n \"\"\"Clear nextjs_port for all sessions belonging to a user.\n\n Called when sandbox goes to sleep to release port allocations.\n\n Args:\n db_session: Database session\n user_id: The user whose sessions should have ports cleared\n\n Returns:\n Number of sessions updated\n \"\"\"\n result = (\n db_session.query(BuildSession)\n .filter(\n BuildSession.user_id == user_id,\n BuildSession.nextjs_port.isnot(None),\n )\n .update({BuildSession.nextjs_port: None})\n )\n db_session.flush()\n logger.info(f\"Cleared {result} nextjs_port allocations for user {user_id}\")\n return result\n\n\ndef fetch_llm_provider_by_type_for_build_mode(\n db_session: Session, provider_type: str\n) -> LLMProviderView | None:\n \"\"\"Fetch an LLM provider by its provider type (e.g., \"anthropic\", \"openai\").\n\n Resolution priority:\n 1. First try to find a provider named \"build-mode-{type}\" (e.g., \"build-mode-anthropic\")\n 2. If not found, fall back to any provider that matches the type\n\n Args:\n db_session: Database session\n provider_type: The provider type (e.g., \"anthropic\", \"openai\", \"openrouter\")\n\n Returns:\n LLMProviderView if found, None otherwise\n \"\"\"\n from onyx.db.llm import fetch_existing_llm_provider\n\n # First try to find a \"build-mode-{type}\" provider\n build_mode_name = f\"build-mode-{provider_type}\"\n provider_model = fetch_existing_llm_provider(\n name=build_mode_name, db_session=db_session\n )\n\n # If not found, fall back to any provider that matches the type\n if not provider_model:\n provider_model = db_session.scalar(\n select(LLMProviderModel)\n .where(LLMProviderModel.provider == provider_type)\n .options(\n selectinload(LLMProviderModel.model_configurations),\n selectinload(LLMProviderModel.groups),\n selectinload(LLMProviderModel.personas),\n )\n )\n\n if not provider_model:\n return None\n return LLMProviderView.from_model(provider_model)\n" }, { "path": "backend/onyx/server/features/build/db/rate_limit.py", "content": "\"\"\"Database queries for Build Mode rate limiting.\"\"\"\n\nfrom datetime import datetime\nfrom uuid import UUID\n\nfrom sqlalchemy import func\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.models import BuildMessage\nfrom onyx.db.models import BuildSession\n\n\ndef count_user_messages_in_window(\n user_id: UUID,\n cutoff_time: datetime,\n db_session: Session,\n) -> int:\n \"\"\"\n Count USER messages for a user since cutoff_time.\n\n Args:\n user_id: The user's UUID\n cutoff_time: Only count messages created at or after this time\n db_session: Database session\n\n Returns:\n Number of USER messages in the time window\n \"\"\"\n return (\n db_session.query(func.count(BuildMessage.id))\n .join(BuildSession, BuildMessage.session_id == BuildSession.id)\n .filter(\n BuildSession.user_id == user_id,\n BuildMessage.type == MessageType.USER,\n BuildMessage.created_at >= cutoff_time,\n )\n .scalar()\n or 0\n )\n\n\ndef count_user_messages_total(user_id: UUID, db_session: Session) -> int:\n \"\"\"\n Count all USER messages for a user (lifetime total).\n\n Args:\n user_id: The user's UUID\n db_session: Database session\n\n Returns:\n Total number of USER messages\n \"\"\"\n return (\n db_session.query(func.count(BuildMessage.id))\n .join(BuildSession, BuildMessage.session_id == BuildSession.id)\n .filter(\n BuildSession.user_id == user_id,\n BuildMessage.type == MessageType.USER,\n )\n .scalar()\n or 0\n )\n\n\ndef get_oldest_message_timestamp(\n user_id: UUID,\n cutoff_time: datetime,\n db_session: Session,\n) -> datetime | None:\n \"\"\"\n Get the timestamp of the oldest USER message in the time window.\n\n Used to calculate when the rate limit will reset (when the oldest\n message ages out of the rolling window).\n\n Args:\n user_id: The user's UUID\n cutoff_time: Only consider messages created at or after this time\n db_session: Database session\n\n Returns:\n Timestamp of oldest message in window, or None if no messages\n \"\"\"\n return (\n db_session.query(BuildMessage.created_at)\n .join(BuildSession, BuildMessage.session_id == BuildSession.id)\n .filter(\n BuildSession.user_id == user_id,\n BuildMessage.type == MessageType.USER,\n BuildMessage.created_at >= cutoff_time,\n )\n .order_by(BuildMessage.created_at.asc())\n .limit(1)\n .scalar()\n )\n" }, { "path": "backend/onyx/server/features/build/db/sandbox.py", "content": "\"\"\"Database operations for CLI agent sandbox management.\"\"\"\n\nimport datetime\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import func\nfrom sqlalchemy import or_\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import SandboxStatus\nfrom onyx.db.models import Sandbox\nfrom onyx.db.models import Snapshot\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef create_sandbox__no_commit(\n db_session: Session,\n user_id: UUID,\n) -> Sandbox:\n \"\"\"Create a new sandbox record for a user.\n\n Sets last_heartbeat to now so that:\n 1. The sandbox has a proper idle timeout baseline from creation\n 2. Long-running provisioning doesn't cause the sandbox to appear \"old\"\n when it transitions to RUNNING\n\n NOTE: This function uses flush() instead of commit(). The caller is\n responsible for committing the transaction when ready.\n \"\"\"\n sandbox = Sandbox(\n user_id=user_id,\n status=SandboxStatus.PROVISIONING,\n last_heartbeat=datetime.datetime.now(datetime.timezone.utc),\n )\n db_session.add(sandbox)\n db_session.flush()\n return sandbox\n\n\ndef get_sandbox_by_user_id(db_session: Session, user_id: UUID) -> Sandbox | None:\n \"\"\"Get sandbox by user ID (primary lookup method).\"\"\"\n stmt = select(Sandbox).where(Sandbox.user_id == user_id)\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef get_sandbox_by_session_id(db_session: Session, session_id: UUID) -> Sandbox | None:\n \"\"\"Get sandbox by session ID (compatibility function).\n\n This function provides backwards compatibility during the transition to\n user-owned sandboxes. It looks up the session's user_id, then finds the\n user's sandbox.\n\n NOTE: This will be removed in a future phase when all callers are updated\n to use get_sandbox_by_user_id() directly.\n \"\"\"\n from onyx.db.models import BuildSession\n\n stmt = select(BuildSession.user_id).where(BuildSession.id == session_id)\n result = db_session.execute(stmt).scalar_one_or_none()\n if result is None:\n return None\n\n return get_sandbox_by_user_id(db_session, result)\n\n\ndef get_sandbox_by_id(db_session: Session, sandbox_id: UUID) -> Sandbox | None:\n \"\"\"Get sandbox by its ID.\"\"\"\n stmt = select(Sandbox).where(Sandbox.id == sandbox_id)\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef update_sandbox_status__no_commit(\n db_session: Session,\n sandbox_id: UUID,\n status: SandboxStatus,\n) -> Sandbox:\n \"\"\"Update sandbox status.\n\n When transitioning to RUNNING, also sets last_heartbeat to now. This ensures\n newly provisioned sandboxes have a proper idle timeout baseline (rather than\n being immediately considered idle due to NULL heartbeat).\n\n NOTE: This function uses flush() instead of commit(). The caller is\n responsible for committing the transaction when ready.\n \"\"\"\n sandbox = get_sandbox_by_id(db_session, sandbox_id)\n if not sandbox:\n raise ValueError(f\"Sandbox {sandbox_id} not found\")\n\n sandbox.status = status\n\n # Set heartbeat when sandbox becomes active to establish idle timeout baseline\n if status == SandboxStatus.RUNNING:\n sandbox.last_heartbeat = datetime.datetime.now(datetime.timezone.utc)\n\n db_session.flush()\n return sandbox\n\n\ndef update_sandbox_heartbeat(db_session: Session, sandbox_id: UUID) -> Sandbox:\n \"\"\"Update sandbox last_heartbeat to now.\"\"\"\n sandbox = get_sandbox_by_id(db_session, sandbox_id)\n if not sandbox:\n raise ValueError(f\"Sandbox {sandbox_id} not found\")\n\n sandbox.last_heartbeat = datetime.datetime.now(datetime.timezone.utc)\n db_session.commit()\n return sandbox\n\n\ndef get_idle_sandboxes(\n db_session: Session, idle_threshold_seconds: int\n) -> list[Sandbox]:\n \"\"\"Get sandboxes that have been idle longer than threshold.\n\n Also includes sandboxes with NULL heartbeat, but only if they were created\n before the threshold (to avoid sweeping up brand-new sandboxes that may have\n NULL heartbeat due to edge cases like older rows or manual inserts).\n \"\"\"\n threshold_time = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(\n seconds=idle_threshold_seconds\n )\n\n stmt = select(Sandbox).where(\n Sandbox.status == SandboxStatus.RUNNING,\n or_(\n Sandbox.last_heartbeat < threshold_time,\n and_(\n Sandbox.last_heartbeat.is_(None),\n Sandbox.created_at < threshold_time,\n ),\n ),\n )\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef get_running_sandbox_count_by_tenant(\n db_session: Session,\n tenant_id: str, # noqa: ARG001\n) -> int:\n \"\"\"Get count of running sandboxes for a tenant (for limit enforcement).\n\n Note: tenant_id parameter is kept for API compatibility but is not used\n since Sandbox model no longer has tenant_id. This function returns\n the count of all running sandboxes.\n \"\"\"\n stmt = select(func.count(Sandbox.id)).where(Sandbox.status == SandboxStatus.RUNNING)\n result = db_session.execute(stmt).scalar()\n return result or 0\n\n\ndef create_snapshot__no_commit(\n db_session: Session,\n session_id: UUID,\n storage_path: str,\n size_bytes: int,\n) -> Snapshot:\n \"\"\"Create a snapshot record for a session.\n\n NOTE: Uses flush() instead of commit(). The caller (cleanup task) is\n responsible for committing after all snapshots + status updates are done,\n so the entire operation is atomic.\n \"\"\"\n snapshot = Snapshot(\n session_id=session_id,\n storage_path=storage_path,\n size_bytes=size_bytes,\n )\n db_session.add(snapshot)\n db_session.flush()\n return snapshot\n\n\ndef get_latest_snapshot_for_session(\n db_session: Session, session_id: UUID\n) -> Snapshot | None:\n \"\"\"Get most recent snapshot for a session.\"\"\"\n stmt = (\n select(Snapshot)\n .where(Snapshot.session_id == session_id)\n .order_by(Snapshot.created_at.desc())\n .limit(1)\n )\n return db_session.execute(stmt).scalar_one_or_none()\n\n\ndef get_snapshots_for_session(db_session: Session, session_id: UUID) -> list[Snapshot]:\n \"\"\"Get all snapshots for a session, ordered by creation time descending.\"\"\"\n stmt = (\n select(Snapshot)\n .where(Snapshot.session_id == session_id)\n .order_by(Snapshot.created_at.desc())\n )\n return list(db_session.execute(stmt).scalars().all())\n\n\ndef delete_old_snapshots(\n db_session: Session,\n tenant_id: str, # noqa: ARG001\n retention_days: int,\n) -> int:\n \"\"\"Delete snapshots older than retention period, return count deleted.\n\n Note: tenant_id parameter is kept for API compatibility but is not used\n since Snapshot model no longer has tenant_id. This function deletes\n all snapshots older than the retention period.\n \"\"\"\n cutoff_time = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(\n days=retention_days\n )\n\n stmt = select(Snapshot).where(\n Snapshot.created_at < cutoff_time,\n )\n old_snapshots = db_session.execute(stmt).scalars().all()\n\n count = 0\n for snapshot in old_snapshots:\n db_session.delete(snapshot)\n count += 1\n\n if count > 0:\n db_session.commit()\n\n return count\n\n\ndef delete_snapshot(db_session: Session, snapshot_id: UUID) -> bool:\n \"\"\"Delete a specific snapshot by ID. Returns True if deleted, False if not found.\"\"\"\n stmt = select(Snapshot).where(Snapshot.id == snapshot_id)\n snapshot = db_session.execute(stmt).scalar_one_or_none()\n\n if not snapshot:\n return False\n\n db_session.delete(snapshot)\n db_session.commit()\n return True\n" }, { "path": "backend/onyx/server/features/build/db/user_library.py", "content": "\"\"\"Database operations for User Library (CRAFT_FILE connector).\n\nHandles storage quota queries and connector/credential setup for the\nUser Library feature in Craft.\n\"\"\"\n\nfrom uuid import UUID\n\nfrom sqlalchemy import and_\nfrom sqlalchemy import cast\nfrom sqlalchemy import func\nfrom sqlalchemy import Integer\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.connector import create_connector\nfrom onyx.db.connector import fetch_connectors\nfrom onyx.db.connector_credential_pair import add_credential_to_connector\nfrom onyx.db.connector_credential_pair import (\n get_connector_credential_pairs_for_user,\n)\nfrom onyx.db.credentials import create_credential\nfrom onyx.db.credentials import fetch_credentials_for_user\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ProcessingMode\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Document as DbDocument\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import User\nfrom onyx.server.documents.models import ConnectorBase\nfrom onyx.server.documents.models import CredentialBase\nfrom onyx.server.features.build.configs import USER_LIBRARY_CONNECTOR_NAME\nfrom onyx.server.features.build.configs import USER_LIBRARY_CREDENTIAL_NAME\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef get_user_storage_bytes(db_session: Session, user_id: UUID) -> int:\n \"\"\"Get total storage usage for a user's library files.\n\n Uses SQL aggregation to sum file_size from doc_metadata JSONB for all\n CRAFT_FILE documents owned by this user, avoiding loading all documents\n into Python memory.\n \"\"\"\n stmt = (\n select(\n func.coalesce(\n func.sum(\n cast(\n DbDocument.doc_metadata[\"file_size\"].as_string(),\n Integer,\n )\n ),\n 0,\n )\n )\n .join(\n DocumentByConnectorCredentialPair,\n DbDocument.id == DocumentByConnectorCredentialPair.id,\n )\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .join(\n Connector,\n ConnectorCredentialPair.connector_id == Connector.id,\n )\n .where(Connector.source == DocumentSource.CRAFT_FILE)\n .where(ConnectorCredentialPair.creator_id == user_id)\n .where(DbDocument.doc_metadata[\"is_directory\"].as_boolean().is_not(True))\n )\n result = db_session.execute(stmt).scalar()\n return int(result or 0)\n\n\ndef get_or_create_craft_connector(db_session: Session, user: User) -> tuple[int, int]:\n \"\"\"Get or create the CRAFT_FILE connector for a user.\n\n Returns:\n Tuple of (connector_id, credential_id)\n\n Note: We need to create a credential even though CRAFT_FILE doesn't require\n authentication. This is because Onyx's connector-credential pair system\n requires a credential for all connectors. The credential is empty ({}).\n\n This function handles recovery from partial creation failures by detecting\n orphaned connectors (connectors without cc_pairs) and completing their setup.\n \"\"\"\n # Check if user already has a complete CRAFT_FILE cc_pair\n cc_pairs = get_connector_credential_pairs_for_user(\n db_session=db_session,\n user=user,\n get_editable=False,\n eager_load_connector=True,\n eager_load_credential=True,\n processing_mode=ProcessingMode.RAW_BINARY,\n )\n\n for cc_pair in cc_pairs:\n if (\n cc_pair.connector.source == DocumentSource.CRAFT_FILE\n and cc_pair.creator_id == user.id\n ):\n return cc_pair.connector.id, cc_pair.credential.id\n\n # No cc_pair for this user — find or create the shared CRAFT_FILE connector\n existing_connectors = fetch_connectors(\n db_session, sources=[DocumentSource.CRAFT_FILE]\n )\n connector_id: int | None = None\n for conn in existing_connectors:\n if conn.name == USER_LIBRARY_CONNECTOR_NAME:\n connector_id = conn.id\n break\n\n if connector_id is None:\n connector_data = ConnectorBase(\n name=USER_LIBRARY_CONNECTOR_NAME,\n source=DocumentSource.CRAFT_FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={\"disabled_paths\": []},\n refresh_freq=None,\n prune_freq=None,\n )\n connector_response = create_connector(\n db_session=db_session,\n connector_data=connector_data,\n )\n connector_id = connector_response.id\n\n # Try to reuse an existing User Library credential for this user\n existing_credentials = fetch_credentials_for_user(\n db_session=db_session,\n user=user,\n )\n credential = None\n for cred in existing_credentials:\n if (\n cred.source == DocumentSource.CRAFT_FILE\n and cred.name == USER_LIBRARY_CREDENTIAL_NAME\n ):\n credential = cred\n break\n\n if credential is None:\n credential_data = CredentialBase(\n credential_json={},\n admin_public=False,\n source=DocumentSource.CRAFT_FILE,\n name=USER_LIBRARY_CREDENTIAL_NAME,\n )\n credential = create_credential(\n credential_data=credential_data,\n user=user,\n db_session=db_session,\n )\n\n # Link them with RAW_BINARY processing mode\n add_credential_to_connector(\n db_session=db_session,\n connector_id=connector_id,\n credential_id=credential.id,\n user=user,\n cc_pair_name=USER_LIBRARY_CONNECTOR_NAME,\n access_type=AccessType.PRIVATE,\n groups=None,\n processing_mode=ProcessingMode.RAW_BINARY,\n )\n\n db_session.commit()\n return connector_id, credential.id\n" }, { "path": "backend/onyx/server/features/build/indexing/persistent_document_writer.py", "content": "\"\"\"\nPersistent Document Writer for writing indexed documents to local filesystem or S3 with\nhierarchical directory structure that mirrors the source organization.\n\nLocal mode (SandboxBackend.LOCAL):\n Writes to local filesystem at {PERSISTENT_DOCUMENT_STORAGE_PATH}/{tenant_id}/knowledge/{user_id}/...\n\nKubernetes mode (SandboxBackend.KUBERNETES):\n Writes to S3 at s3://{SANDBOX_S3_BUCKET}/{tenant_id}/knowledge/{user_id}/...\n This is the same location that kubernetes_sandbox_manager.py reads from when\n provisioning sandboxes.\n\nBoth modes use consistent tenant/user-segregated paths for multi-tenant isolation.\n\"\"\"\n\nimport hashlib\nimport json\nimport unicodedata\nfrom pathlib import Path\nfrom typing import Any\n\nfrom botocore.exceptions import ClientError\nfrom mypy_boto3_s3.client import S3Client\n\nfrom onyx.connectors.models import Document\nfrom onyx.server.features.build.configs import PERSISTENT_DOCUMENT_STORAGE_PATH\nfrom onyx.server.features.build.configs import SANDBOX_BACKEND\nfrom onyx.server.features.build.configs import SANDBOX_S3_BUCKET\nfrom onyx.server.features.build.configs import SandboxBackend\nfrom onyx.server.features.build.s3.s3_client import build_s3_client\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# =============================================================================\n# Shared Utilities for Path Building\n# =============================================================================\n\n\ndef sanitize_path_component(component: str, replace_slash: bool = True) -> str:\n \"\"\"Sanitize a path component for file system / S3 key safety.\n\n Args:\n component: The path component to sanitize\n replace_slash: If True, replaces forward slashes (needed for local filesystem).\n Set to False for S3 where `/` is a valid delimiter.\n\n Returns:\n Sanitized path component safe for use in file paths or S3 keys\n \"\"\"\n # First, normalize Unicode to decomposed form and remove combining characters\n # This handles cases like accented characters, while also filtering format chars\n normalized = unicodedata.normalize(\"NFKD\", component)\n\n # Filter out Unicode format/control characters (categories Cf, Cc)\n # This removes invisible chars like U+2060 (WORD JOINER), zero-width spaces, etc.\n sanitized = \"\".join(\n c for c in normalized if unicodedata.category(c) not in (\"Cf\", \"Cc\")\n )\n\n # Replace spaces with underscores\n sanitized = sanitized.replace(\" \", \"_\")\n # Replace problematic characters\n if replace_slash:\n sanitized = sanitized.replace(\"/\", \"_\")\n sanitized = sanitized.replace(\"\\\\\", \"_\").replace(\":\", \"_\")\n sanitized = sanitized.replace(\"<\", \"_\").replace(\">\", \"_\").replace(\"|\", \"_\")\n sanitized = sanitized.replace('\"', \"_\").replace(\"?\", \"_\").replace(\"*\", \"_\")\n return sanitized.strip() or \"unnamed\"\n\n\ndef sanitize_filename(name: str, replace_slash: bool = True) -> str:\n \"\"\"Sanitize name for use as filename.\n\n Args:\n name: The filename to sanitize\n replace_slash: Passed through to sanitize_path_component\n\n Returns:\n Sanitized filename, truncated with hash suffix if too long\n \"\"\"\n sanitized = sanitize_path_component(name, replace_slash=replace_slash)\n if len(sanitized) > 200:\n # Keep first 150 chars + hash suffix for uniqueness\n hash_suffix = hashlib.sha256(name.encode()).hexdigest()[:16]\n return f\"{sanitized[:150]}_{hash_suffix}\"\n return sanitized\n\n\ndef normalize_leading_slash(path: str) -> str:\n \"\"\"Ensure a path starts with exactly one leading slash.\"\"\"\n return \"/\" + path.lstrip(\"/\")\n\n\ndef get_base_filename(doc: Document, replace_slash: bool = True) -> str:\n \"\"\"Get base filename from document, preferring semantic identifier.\n\n Args:\n doc: The document to get filename for\n replace_slash: Passed through to sanitize_filename\n\n Returns:\n Sanitized base filename (without extension)\n \"\"\"\n name = doc.semantic_identifier or doc.title or doc.id\n return sanitize_filename(name, replace_slash=replace_slash)\n\n\ndef build_document_subpath(doc: Document, replace_slash: bool = True) -> list[str]:\n \"\"\"Build the source/hierarchy path components from a document.\n\n Returns path components like: [source, hierarchy_part1, hierarchy_part2, ...]\n\n This is the common part of the path that comes after user/tenant segregation.\n\n Args:\n doc: The document to build path for\n replace_slash: Passed through to sanitize_path_component\n\n Returns:\n List of sanitized path components\n \"\"\"\n parts: list[str] = []\n\n # Source type (e.g., \"google_drive\", \"confluence\")\n parts.append(doc.source.value)\n\n # Get hierarchy from doc_metadata\n hierarchy: dict[str, Any] = (\n doc.doc_metadata.get(\"hierarchy\", {}) if doc.doc_metadata else {}\n )\n source_path: list[str] = hierarchy.get(\"source_path\", [])\n\n if source_path:\n parts.extend(\n [\n sanitize_path_component(p, replace_slash=replace_slash)\n for p in source_path\n ]\n )\n\n return parts\n\n\ndef resolve_duplicate_filename(\n doc: Document,\n base_filename: str,\n has_duplicates: bool,\n replace_slash: bool = True,\n) -> str:\n \"\"\"Resolve filename, appending ID suffix if there are duplicates.\n\n Args:\n doc: The document (for ID extraction)\n base_filename: The base filename without extension\n has_duplicates: Whether there are other docs with the same base filename\n replace_slash: Passed through to sanitize_path_component\n\n Returns:\n Final filename with .json extension\n \"\"\"\n if has_duplicates:\n id_suffix = sanitize_path_component(doc.id, replace_slash=replace_slash)\n if len(id_suffix) > 50:\n id_suffix = hashlib.sha256(doc.id.encode()).hexdigest()[:16]\n return f\"{base_filename}_{id_suffix}.json\"\n return f\"{base_filename}.json\"\n\n\ndef serialize_document(doc: Document) -> dict[str, Any]:\n \"\"\"Serialize a document to a dictionary for JSON storage.\n\n Args:\n doc: The document to serialize\n\n Returns:\n Dictionary representation of the document\n \"\"\"\n return {\n \"id\": doc.id,\n \"semantic_identifier\": doc.semantic_identifier,\n \"title\": doc.title,\n \"source\": doc.source.value,\n \"doc_updated_at\": (\n doc.doc_updated_at.isoformat() if doc.doc_updated_at else None\n ),\n \"metadata\": doc.metadata,\n \"doc_metadata\": doc.doc_metadata,\n \"sections\": [\n {\"text\": s.text if hasattr(s, \"text\") else None, \"link\": s.link}\n for s in doc.sections\n ],\n \"primary_owners\": [o.model_dump() for o in (doc.primary_owners or [])],\n \"secondary_owners\": [o.model_dump() for o in (doc.secondary_owners or [])],\n }\n\n\n# =============================================================================\n# Classes\n# =============================================================================\n\n\nclass PersistentDocumentWriter:\n \"\"\"Writes indexed documents to local filesystem with hierarchical structure.\n\n Documents are stored in tenant/user-segregated paths:\n {base_path}/{tenant_id}/knowledge/{user_id}/{source}/{hierarchy}/document.json\n\n This enables per-tenant and per-user isolation for sandbox access control.\n \"\"\"\n\n def __init__(\n self,\n base_path: str,\n tenant_id: str,\n user_id: str,\n ):\n self.base_path = Path(base_path)\n self.tenant_id = tenant_id\n self.user_id = user_id\n\n def write_documents(self, documents: list[Document]) -> list[str]:\n \"\"\"Write documents to local filesystem, returns written file paths.\"\"\"\n written_paths: list[str] = []\n\n # Build a map of base filenames to detect duplicates\n # Key: (directory_path, base_filename) -> list of docs with that name\n filename_map: dict[tuple[Path, str], list[Document]] = {}\n\n for doc in documents:\n dir_path = self._build_directory_path(doc)\n base_filename = get_base_filename(doc, replace_slash=True)\n key = (dir_path, base_filename)\n if key not in filename_map:\n filename_map[key] = []\n filename_map[key].append(doc)\n\n # Now write documents, appending ID if there are duplicates\n for (dir_path, base_filename), docs in filename_map.items():\n has_duplicates = len(docs) > 1\n for doc in docs:\n filename = resolve_duplicate_filename(\n doc, base_filename, has_duplicates, replace_slash=True\n )\n path = dir_path / filename\n self._write_document(doc, path)\n written_paths.append(str(path))\n\n return written_paths\n\n def _build_directory_path(self, doc: Document) -> Path:\n \"\"\"Build directory path from document metadata.\n\n Documents are stored under tenant/user-segregated paths:\n {base_path}/{tenant_id}/knowledge/{user_id}/{source}/{hierarchy}/\n\n This enables per-tenant and per-user isolation for sandbox access control.\n \"\"\"\n # Tenant and user segregation prefix (matches S3 path structure)\n parts = [self.tenant_id, \"knowledge\", self.user_id]\n # Add source and hierarchy from document\n parts.extend(build_document_subpath(doc, replace_slash=True))\n\n return self.base_path / \"/\".join(parts)\n\n def _write_document(self, doc: Document, path: Path) -> None:\n \"\"\"Serialize and write document to filesystem.\"\"\"\n content = serialize_document(doc)\n\n # Create parent directories if they don't exist\n path.parent.mkdir(parents=True, exist_ok=True)\n\n # Write the JSON file\n with open(path, \"w\", encoding=\"utf-8\") as f:\n json.dump(content, f, indent=2, default=str)\n\n logger.debug(f\"Wrote document to {path}\")\n\n def write_raw_file(\n self,\n path: str,\n content: bytes,\n content_type: str | None = None, # noqa: ARG002\n ) -> str:\n \"\"\"Write a raw binary file to local filesystem (for User Library).\n\n Unlike write_documents which serializes Document objects to JSON, this method\n writes raw binary content directly. Used for user-uploaded files like xlsx, pptx.\n\n Args:\n path: Relative path within user's library (e.g., \"/project-data/financials.xlsx\")\n content: Raw binary content to write\n content_type: MIME type of the file (stored as metadata, unused locally)\n\n Returns:\n Full filesystem path where file was written\n \"\"\"\n # Build full path: {base_path}/{tenant}/knowledge/{user}/user_library/{path}\n normalized_path = normalize_leading_slash(path)\n full_path = (\n self.base_path\n / self.tenant_id\n / \"knowledge\"\n / self.user_id\n / \"user_library\"\n / normalized_path.lstrip(\"/\")\n )\n\n # Create parent directories if they don't exist\n full_path.parent.mkdir(parents=True, exist_ok=True)\n\n # Write the raw binary content\n with open(full_path, \"wb\") as f:\n f.write(content)\n\n logger.debug(f\"Wrote raw file to {full_path}\")\n return str(full_path)\n\n def delete_raw_file(self, path: str) -> None:\n \"\"\"Delete a raw file from local filesystem.\n\n Args:\n path: Relative path within user's library (e.g., \"/project-data/financials.xlsx\")\n \"\"\"\n # Build full path\n normalized_path = normalize_leading_slash(path)\n full_path = (\n self.base_path\n / self.tenant_id\n / \"knowledge\"\n / self.user_id\n / \"user_library\"\n / normalized_path.lstrip(\"/\")\n )\n\n if full_path.exists():\n full_path.unlink()\n logger.debug(f\"Deleted raw file at {full_path}\")\n else:\n logger.warning(f\"File not found for deletion: {full_path}\")\n\n\nclass S3PersistentDocumentWriter:\n \"\"\"Writes indexed documents to S3 with hierarchical structure.\n\n Documents are stored in tenant/user-segregated paths:\n s3://{bucket}/{tenant_id}/knowledge/{user_id}/{source}/{hierarchy}/document.json\n\n This matches the location that KubernetesSandboxManager reads from when\n provisioning sandboxes (via the sidecar container's s5cmd sync command).\n \"\"\"\n\n def __init__(self, tenant_id: str, user_id: str):\n \"\"\"Initialize S3PersistentDocumentWriter.\n\n Args:\n tenant_id: Tenant identifier for multi-tenant isolation\n user_id: User ID for user-segregated storage paths\n \"\"\"\n self.tenant_id = tenant_id\n self.user_id = user_id\n self.bucket = SANDBOX_S3_BUCKET\n self._s3_client: S3Client | None = None\n\n def _get_s3_client(self) -> S3Client:\n \"\"\"Lazily initialize S3 client.\n\n Uses the craft-specific boto3 client which only supports IAM roles (IRSA).\n \"\"\"\n if self._s3_client is None:\n self._s3_client = build_s3_client()\n return self._s3_client\n\n def write_documents(self, documents: list[Document]) -> list[str]:\n \"\"\"Write documents to S3, returns written S3 keys.\n\n Args:\n documents: List of documents to write\n\n Returns:\n List of S3 keys that were written\n \"\"\"\n written_keys: list[str] = []\n\n # Build a map of base keys to detect duplicates\n # Key: (directory_prefix, base_filename) -> list of docs with that name\n key_map: dict[tuple[str, str], list[Document]] = {}\n\n for doc in documents:\n dir_prefix = self._build_directory_path(doc)\n base_filename = get_base_filename(doc, replace_slash=False)\n key = (dir_prefix, base_filename)\n if key not in key_map:\n key_map[key] = []\n key_map[key].append(doc)\n\n # Now write documents, appending ID if there are duplicates\n s3_client = self._get_s3_client()\n\n for (dir_prefix, base_filename), docs in key_map.items():\n has_duplicates = len(docs) > 1\n for doc in docs:\n filename = resolve_duplicate_filename(\n doc, base_filename, has_duplicates, replace_slash=False\n )\n s3_key = f\"{dir_prefix}/{filename}\"\n self._write_document(s3_client, doc, s3_key)\n written_keys.append(s3_key)\n\n return written_keys\n\n def _build_directory_path(self, doc: Document) -> str:\n \"\"\"Build S3 key prefix from document metadata.\n\n Documents are stored under tenant/user-segregated paths:\n {tenant_id}/knowledge/{user_id}/{source}/{hierarchy}/\n\n This matches the path that KubernetesSandboxManager syncs from:\n s5cmd sync \"s3://{bucket}/{tenant_id}/knowledge/{user_id}/*\" /workspace/files/\n \"\"\"\n # Tenant and user segregation (matches K8s sandbox init container path)\n parts = [self.tenant_id, \"knowledge\", self.user_id]\n # Add source and hierarchy from document\n parts.extend(build_document_subpath(doc, replace_slash=False))\n\n return \"/\".join(parts)\n\n def _write_document(self, s3_client: S3Client, doc: Document, s3_key: str) -> None:\n \"\"\"Serialize and write document to S3.\"\"\"\n content = serialize_document(doc)\n json_content = json.dumps(content, indent=2, default=str)\n\n try:\n s3_client.put_object(\n Bucket=self.bucket,\n Key=s3_key,\n Body=json_content.encode(\"utf-8\"),\n ContentType=\"application/json\",\n )\n logger.debug(f\"Wrote document to s3://{self.bucket}/{s3_key}\")\n except ClientError as e:\n logger.error(f\"Failed to write to S3: {e}\")\n raise\n\n def write_raw_file(\n self,\n path: str,\n content: bytes,\n content_type: str | None = None,\n ) -> str:\n \"\"\"Write a raw binary file to S3 (for User Library).\n\n Unlike write_documents which serializes Document objects to JSON, this method\n writes raw binary content directly. Used for user-uploaded files like xlsx, pptx.\n\n Args:\n path: Relative path within user's library (e.g., \"/project-data/financials.xlsx\")\n content: Raw binary content to write\n content_type: MIME type of the file\n\n Returns:\n S3 key where file was written\n \"\"\"\n # Build S3 key: {tenant}/knowledge/{user}/user_library/{path}\n normalized_path = path.lstrip(\"/\")\n s3_key = (\n f\"{self.tenant_id}/knowledge/{self.user_id}/user_library/{normalized_path}\"\n )\n\n s3_client = self._get_s3_client()\n\n try:\n s3_client.put_object(\n Bucket=self.bucket,\n Key=s3_key,\n Body=content,\n ContentType=content_type or \"application/octet-stream\",\n )\n logger.debug(f\"Wrote raw file to s3://{self.bucket}/{s3_key}\")\n return s3_key\n except ClientError as e:\n logger.error(f\"Failed to write raw file to S3: {e}\")\n raise\n\n def delete_raw_file(self, s3_key: str) -> None:\n \"\"\"Delete a raw file from S3.\n\n Args:\n s3_key: Full S3 key of the file to delete\n \"\"\"\n s3_client = self._get_s3_client()\n\n try:\n s3_client.delete_object(Bucket=self.bucket, Key=s3_key)\n logger.debug(f\"Deleted raw file at s3://{self.bucket}/{s3_key}\")\n except ClientError as e:\n logger.error(f\"Failed to delete raw file from S3: {e}\")\n raise\n\n def delete_raw_file_by_path(self, path: str) -> None:\n \"\"\"Delete a raw file from S3 by its relative path.\n\n Args:\n path: Relative path within user's library (e.g., \"/project-data/financials.xlsx\")\n \"\"\"\n normalized_path = path.lstrip(\"/\")\n s3_key = (\n f\"{self.tenant_id}/knowledge/{self.user_id}/user_library/{normalized_path}\"\n )\n self.delete_raw_file(s3_key)\n\n\ndef get_persistent_document_writer(\n user_id: str,\n tenant_id: str,\n) -> PersistentDocumentWriter | S3PersistentDocumentWriter:\n \"\"\"Factory function to create a PersistentDocumentWriter with default configuration.\n\n Args:\n user_id: User ID for user-segregated storage paths.\n tenant_id: Tenant ID for multi-tenant isolation.\n\n Both local and S3 modes use consistent tenant/user-segregated paths:\n - Local: {base_path}/{tenant_id}/knowledge/{user_id}/...\n - S3: s3://{bucket}/{tenant_id}/knowledge/{user_id}/...\n\n Returns:\n PersistentDocumentWriter for local mode, S3PersistentDocumentWriter for K8s mode\n \"\"\"\n if SANDBOX_BACKEND == SandboxBackend.LOCAL:\n return PersistentDocumentWriter(\n base_path=PERSISTENT_DOCUMENT_STORAGE_PATH,\n tenant_id=tenant_id,\n user_id=user_id,\n )\n elif SANDBOX_BACKEND == SandboxBackend.KUBERNETES:\n return S3PersistentDocumentWriter(\n tenant_id=tenant_id,\n user_id=user_id,\n )\n else:\n raise ValueError(f\"Unknown sandbox backend: {SANDBOX_BACKEND}\")\n" }, { "path": "backend/onyx/server/features/build/s3/s3_client.py", "content": "import boto3\nfrom mypy_boto3_s3.client import S3Client\n\nfrom onyx.configs.app_configs import AWS_REGION_NAME\n\n\ndef build_s3_client() -> S3Client:\n \"\"\"Build an S3 client using IAM roles (IRSA)\"\"\"\n return boto3.client(\"s3\", region_name=AWS_REGION_NAME)\n" }, { "path": "backend/onyx/server/features/build/sandbox/README.md", "content": "# Onyx Sandbox System\n\nThis directory contains the implementation of Onyx's sandbox system for running OpenCode agents in isolated environments.\n\n## Overview\n\nThe sandbox system provides isolated execution environments where OpenCode agents can build web applications, run code, and interact with knowledge files. Each sandbox includes:\n\n- **Next.js development environment** - Lightweight Next.js scaffold with shadcn/ui and Recharts for building UIs\n- **Python virtual environment** - Pre-installed packages for data processing\n- **OpenCode agent** - AI coding agent with access to tools and MCP servers\n- **Knowledge files** - Access to indexed documents and user uploads\n\n## Architecture\n\n### Deployment Modes\n\n1. **Local Mode** (`SANDBOX_BACKEND=local`)\n - Sandboxes run as directories on the local filesystem\n - No automatic cleanup or snapshots\n - Suitable for development and testing\n\n2. **Kubernetes Mode** (`SANDBOX_BACKEND=kubernetes`)\n - Sandboxes run as Kubernetes pods\n - Automatic snapshots to S3\n - Auto-cleanup of idle sandboxes\n - Production-ready with resource isolation\n\n### Directory Structure\n\n```\n/workspace/ # Sandbox root (in container)\n├── outputs/ # Working directory\n│ ├── web/ # Lightweight Next.js app (shadcn/ui, Recharts)\n│ ├── slides/ # Generated presentations\n│ ├── markdown/ # Generated documents\n│ └── graphs/ # Generated visualizations\n├── .venv/ # Python virtual environment\n├── files/ # Symlink to knowledge files\n├── attachments/ # User uploads\n├── AGENTS.md # Agent instructions\n└── .opencode/\n └── skills/ # Agent skills\n```\n\n## Setup\n\n### Running via Docker/Kubernetes (Zero Setup!) 🎉\n\n**No setup required!** Just build and deploy:\n\n```bash\n# Build backend image (includes both templates)\ncd backend\ndocker build -f Dockerfile.sandbox-templates -t onyxdotapp/backend:latest .\n\n# Build sandbox container (lightweight runner)\ncd onyx/server/features/build/sandbox/kubernetes/docker\ndocker build -t onyxdotapp/sandbox:latest .\n\n# Deploy with docker-compose or kubectl - sandboxes work immediately!\n```\n\n**How it works:**\n\n- **Backend image**: Contains both templates at build time:\n - Web template at `/templates/outputs/web` (lightweight Next.js scaffold, ~2MB)\n - Python venv template at `/templates/venv` (pre-installed packages, ~50MB)\n- **Init container** (Kubernetes only): Syncs knowledge files from S3\n- **Sandbox startup**: Runs `npm install` (for fresh dependency locks) + `next dev`\n\n### Running Backend Directly (Without Docker)\n\n**Only needed if you're running the Onyx backend outside of Docker.** Most developers use Docker and can skip this section.\n\nIf you're running the backend Python process directly on your machine, you need templates at `/templates/`:\n\n#### Web Template\n\nThe web template is a lightweight Next.js app (Next.js 16, React 19, shadcn/ui, Recharts) checked into the codebase at `backend/onyx/server/features/build/templates/outputs/web/`.\n\nFor local development, create a symlink to this template:\n\n```bash\nsudo mkdir -p /templates/outputs\nsudo ln -s $(pwd)/backend/onyx/server/features/build/templates/outputs/web /templates/outputs/web\n```\n\n#### Python Venv Template\n\nIf you don't have a venv template, create it:\n\n```bash\n# Use the utility script\ncd backend\npython -m onyx.server.features.build.sandbox.util.build_venv_template\n\n# Or manually\npython3 -m venv /templates/venv\n/templates/venv/bin/pip install -r backend/onyx/server/features/build/sandbox/kubernetes/docker/initial-requirements.txt\n```\n\n#### System Dependencies (for PPTX skill)\n\nThe PPTX skill requires LibreOffice and Poppler for PDF conversion and thumbnail generation:\n\n**macOS:**\n\n```bash\nbrew install poppler\nbrew install --cask libreoffice\n```\n\nEnsure `soffice` is on your PATH:\n\n```bash\nexport PATH=\"/Applications/LibreOffice.app/Contents/MacOS:$PATH\"\n```\n\n**Linux (Debian/Ubuntu):**\n\n```bash\nsudo apt-get install libreoffice-impress poppler-utils\n```\n\n**That's it!** When sandboxes are created:\n\n1. Web template is copied from `/templates/outputs/web`\n2. Python venv is copied from `/templates/venv`\n3. `npm install` runs automatically to install fresh Next.js dependencies\n\n## OpenCode Configuration\n\nEach sandbox includes an OpenCode agent configured with:\n\n- **LLM Provider**: Anthropic, OpenAI, Google, Bedrock, or Azure\n- **Extended thinking**: High reasoning effort / thinking budgets for complex tasks\n- **Tool permissions**: File operations, bash commands, web access\n- **Disabled tools**: Configurable via `OPENCODE_DISABLED_TOOLS` env var\n\nConfiguration is generated dynamically in `templates/opencode_config.py`.\n\n## Key Components\n\n### Managers\n\n- **`base.py`** - Abstract base class defining the sandbox interface\n- **`local/manager.py`** - Filesystem-based sandbox manager for local development\n- **`kubernetes/manager.py`** - Kubernetes-based sandbox manager for production\n\n### Managers (Shared)\n\n- **`manager/directory_manager.py`** - Creates sandbox directory structure and copies templates\n- **`manager/snapshot_manager.py`** - Handles snapshot creation and restoration\n\n### Utilities\n\n- **`util/opencode_config.py`** - Generates OpenCode configuration with MCP support\n- **`util/agent_instructions.py`** - Generates agent instructions (AGENTS.md)\n- **`util/build_venv_template.py`** - Utility to build Python venv template for local development\n\n### Templates\n\n- **`../templates/outputs/web/`** - Lightweight Next.js scaffold (shadcn/ui, Recharts) versioned with the backend code\n\n### Kubernetes Specific\n\n- **`kubernetes/docker/Dockerfile`** - Sandbox container image (runs Next.js + OpenCode)\n- **`kubernetes/docker/entrypoint.sh`** - Container startup script\n\n## Environment Variables\n\n### Core Settings\n\n```bash\n# Sandbox backend mode\nSANDBOX_BACKEND=local|kubernetes # Default: local\n\n# Template paths (local mode)\nOUTPUTS_TEMPLATE_PATH=/templates/outputs # Default: /templates/outputs\nVENV_TEMPLATE_PATH=/templates/venv # Default: /templates/venv\n\n# Sandbox base path (local mode)\nSANDBOX_BASE_PATH=/tmp/onyx-sandboxes # Default: /tmp/onyx-sandboxes\n\n# OpenCode configuration\nOPENCODE_DISABLED_TOOLS=question # Comma-separated list, default: question\n```\n\n### Kubernetes Settings\n\n```bash\n# Kubernetes namespace\nSANDBOX_NAMESPACE=onyx-sandboxes # Default: onyx-sandboxes\n\n# Container image\nSANDBOX_CONTAINER_IMAGE=onyxdotapp/sandbox:latest\n\n# S3 bucket for snapshots and files\nSANDBOX_S3_BUCKET=onyx-sandbox-files # Default: onyx-sandbox-files\n\n# Service accounts\nSANDBOX_SERVICE_ACCOUNT_NAME=sandbox-runner # No AWS access\nSANDBOX_FILE_SYNC_SERVICE_ACCOUNT=sandbox-file-sync # Has S3 access via IRSA\n```\n\n### Lifecycle Settings\n\n```bash\n# Idle timeout before cleanup (seconds)\nSANDBOX_IDLE_TIMEOUT_SECONDS=900 # Default: 900 (15 minutes)\n\n# Max concurrent sandboxes per organization\nSANDBOX_MAX_CONCURRENT_PER_ORG=10 # Default: 10\n\n# Next.js port range (local mode)\nSANDBOX_NEXTJS_PORT_START=3010 # Default: 3010\nSANDBOX_NEXTJS_PORT_END=3100 # Default: 3100\n```\n\n## Testing\n\n### Integration Tests\n\n```bash\n# Test local sandbox provisioning\nuv run pytest backend/tests/integration/sandbox/test_local_sandbox.py\n\n# Test Kubernetes sandbox provisioning (requires k8s cluster)\nuv run pytest backend/tests/integration/sandbox/test_kubernetes_sandbox.py\n```\n\n### Manual Testing\n\n```bash\n# Start a local sandbox session\ncurl -X POST http://localhost:3000/api/build/session \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"user_id\": \"user-123\",\n \"file_system_path\": \"/path/to/files\"\n }'\n\n# Send a message to the agent\ncurl -X POST http://localhost:3000/api/build/session/{session_id}/message \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"message\": \"Create a simple web page\"\n }'\n```\n\n## Troubleshooting\n\n### Sandbox Stuck in PROVISIONING (Kubernetes)\n\n**Symptoms**: Sandbox status never changes from `PROVISIONING`\n\n**Solutions**:\n\n- Check pod logs: `kubectl logs -n onyx-sandboxes sandbox-{sandbox-id}`\n- Check init container: `kubectl logs -n onyx-sandboxes sandbox-{sandbox-id} -c file-sync`\n- Verify init container completed: `kubectl describe pod -n onyx-sandboxes sandbox-{sandbox-id}`\n- Check S3 bucket access: Ensure init container service account has IRSA configured\n\n### Next.js Server Won't Start\n\n**Symptoms**: Sandbox provisioned but web preview doesn't load\n\n**Solutions**:\n\n- **Local mode**: Check if port is already in use\n- **Docker/K8s**: Check container logs: `kubectl logs -n onyx-sandboxes sandbox-{sandbox-id}`\n- Verify npm install succeeded (check entrypoint.sh logs)\n- Check that web template was copied: `kubectl exec -n onyx-sandboxes sandbox-{sandbox-id} -- ls /workspace/outputs/web`\n\n### Templates Not Found (Local Mode)\n\n**Symptoms**: `RuntimeError: Sandbox templates are missing`\n\n**Solution**: Set up templates as described in the \"Local Development\" section above:\n\n```bash\n# Symlink web template\nsudo ln -s $(pwd)/backend/onyx/server/features/build/templates/outputs/web /templates/outputs/web\n\n# Create Python venv\npython3 -m venv /templates/venv\n/templates/venv/bin/pip install -r backend/onyx/server/features/build/sandbox/kubernetes/docker/initial-requirements.txt\n```\n\n### Permission Denied\n\n**Symptoms**: `Permission denied` error accessing `/templates/`\n\n**Solution**: Either use sudo when creating symlinks, or use custom paths:\n\n```bash\nexport OUTPUTS_TEMPLATE_PATH=$HOME/.onyx/templates/outputs\nexport VENV_TEMPLATE_PATH=$HOME/.onyx/templates/venv\n\n# Then symlink to your home directory\nmkdir -p $HOME/.onyx/templates/outputs\nln -s $(pwd)/backend/onyx/server/features/build/templates/outputs/web $HOME/.onyx/templates/outputs/web\n```\n\n## Security Considerations\n\n### Sandbox Isolation\n\n- **Kubernetes pods** run with restricted security context (non-root, no privilege escalation)\n- **Init containers** have S3 access for file sync, but main sandbox container does NOT\n- **Network policies** can restrict sandbox egress traffic\n- **Resource limits** prevent resource exhaustion\n\n### Credentials Management\n\n- LLM API keys are passed as environment variables (not stored in sandbox)\n- User file access is read-only via symlinks\n- Snapshots are isolated per tenant in S3\n\n## Development\n\n### Adding New MCP Servers\n\n1. Add MCP configuration to `templates/opencode_config.py`:\n\n ```python\n config[\"mcp\"] = {\n \"my-mcp\": {\n \"type\": \"local\",\n \"command\": [\"npx\", \"@my/mcp@latest\"],\n \"enabled\": True,\n }\n }\n ```\n\n2. Install required npm packages in web template (if needed)\n\n3. Rebuild Docker image and templates\n\n### Modifying Agent Instructions\n\nEdit `AGENTS.template.md` in the build directory. This is populated with dynamic content by `templates/agent_instructions.py`.\n\n### Adding New Tools/Permissions\n\nUpdate `templates/opencode_config.py` to add/remove tool permissions in the `permission` section.\n\n## Template Details\n\n### Web Template\n\nThe lightweight Next.js template (`backend/onyx/server/features/build/templates/outputs/web/`) includes:\n\n- **Framework**: Next.js 16.1.4 with React 19.2.3\n- **UI Library**: shadcn/ui components with Radix UI primitives\n- **Styling**: Tailwind CSS v4 with custom theming support\n- **Charts**: Recharts for data visualization\n- **Size**: ~2MB (excluding node_modules, which are installed fresh per sandbox)\n\nThis template provides a modern development environment without the complexity of the full Onyx application, allowing agents to build custom UIs quickly.\n\n### Python Venv Template\n\nThe Python venv (`/templates/venv/`) includes packages from `initial-requirements.txt`:\n\n- Data processing: pandas, numpy, polars\n- HTTP clients: requests, httpx\n- Utilities: python-dotenv, pydantic\n\n## References\n\n- [OpenCode Documentation](https://docs.opencode.ai)\n- [Next.js Documentation](https://nextjs.org/docs)\n- [shadcn/ui Components](https://ui.shadcn.com)\n" }, { "path": "backend/onyx/server/features/build/sandbox/__init__.py", "content": "\"\"\"\nSandbox module for CLI agent filesystem-based isolation.\n\nThis module provides lightweight sandbox management for CLI-based AI agent sessions.\nEach sandbox is a directory on the local filesystem or a Kubernetes pod.\n\nUsage:\n from onyx.server.features.build.sandbox import get_sandbox_manager\n\n # Get the appropriate sandbox manager based on SANDBOX_BACKEND config\n sandbox_manager = get_sandbox_manager()\n\n # Use the sandbox manager\n sandbox_info = sandbox_manager.provision(...)\n\nModule structure:\n - base.py: SandboxManager ABC and get_sandbox_manager() factory\n - models.py: Shared Pydantic models\n - local/: Local filesystem-based implementation for development\n - kubernetes/: Kubernetes pod-based implementation for production\n - internal/: Shared internal utilities (snapshot manager)\n\"\"\"\n\nfrom onyx.server.features.build.sandbox.base import get_sandbox_manager\nfrom onyx.server.features.build.sandbox.base import SandboxManager\nfrom onyx.server.features.build.sandbox.local.local_sandbox_manager import (\n LocalSandboxManager,\n)\nfrom onyx.server.features.build.sandbox.models import FilesystemEntry\nfrom onyx.server.features.build.sandbox.models import SandboxInfo\nfrom onyx.server.features.build.sandbox.models import SnapshotInfo\n\n__all__ = [\n # Factory function (preferred)\n \"get_sandbox_manager\",\n # Interface\n \"SandboxManager\",\n # Implementations\n \"LocalSandboxManager\",\n # Models\n \"SandboxInfo\",\n \"SnapshotInfo\",\n \"FilesystemEntry\",\n]\n" }, { "path": "backend/onyx/server/features/build/sandbox/base.py", "content": "\"\"\"Abstract base class and factory for sandbox operations.\n\nSandboxManager is the abstract interface for sandbox lifecycle management.\nUse get_sandbox_manager() to get the appropriate implementation based on SANDBOX_BACKEND.\n\nIMPORTANT: SandboxManager implementations must NOT interface with the database directly.\nAll database operations should be handled by the caller (SessionManager, Celery tasks, etc.).\n\nArchitecture Note (User-Shared Sandbox Model):\n- One sandbox (container/pod) is shared across all of a user's sessions\n- provision() creates the user's sandbox with shared files/ directory\n- setup_session_workspace() creates per-session workspace within the sandbox\n- cleanup_session_workspace() removes session workspace on session delete\n- terminate() destroys the entire sandbox (all sessions)\n\"\"\"\n\nimport threading\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom collections.abc import Generator\nfrom typing import Any\nfrom uuid import UUID\n\nfrom onyx.server.features.build.configs import SANDBOX_BACKEND\nfrom onyx.server.features.build.configs import SandboxBackend\nfrom onyx.server.features.build.sandbox.models import FilesystemEntry\nfrom onyx.server.features.build.sandbox.models import LLMProviderConfig\nfrom onyx.server.features.build.sandbox.models import SandboxInfo\nfrom onyx.server.features.build.sandbox.models import SnapshotResult\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n# ACPEvent is a union type defined in both local and kubernetes modules\n# Using Any here to avoid circular imports - the actual type checking\n# happens in the implementation modules\nACPEvent = Any\n\n\nclass SandboxManager(ABC):\n \"\"\"Abstract interface for sandbox operations.\n\n Defines the contract for sandbox lifecycle management including:\n - Provisioning and termination (user-level)\n - Session workspace setup and cleanup (session-level)\n - Snapshot creation (session-level)\n - Health checks\n - Agent communication (session-level)\n - Filesystem operations (session-level)\n\n Directory Structure:\n $SANDBOX_ROOT/\n ├── files/ # SHARED - symlink to user's persistent documents\n └── sessions/\n ├── $session_id_1/ # Per-session workspace\n │ ├── outputs/ # Agent output for this session\n │ │ └── web/ # Next.js app\n │ ├── venv/ # Python virtual environment\n │ ├── skills/ # Opencode skills\n │ ├── AGENTS.md # Agent instructions\n │ ├── opencode.json # LLM config\n │ └── attachments/\n └── $session_id_2/\n └── ...\n\n IMPORTANT: Implementations must NOT interface with the database directly.\n All database operations should be handled by the caller.\n\n Use get_sandbox_manager() to get the appropriate implementation.\n \"\"\"\n\n @abstractmethod\n def provision(\n self,\n sandbox_id: UUID,\n user_id: UUID,\n tenant_id: str,\n llm_config: LLMProviderConfig,\n ) -> SandboxInfo:\n \"\"\"Provision a new sandbox for a user.\n\n Creates the sandbox container/directory with:\n - sessions/ directory for per-session workspaces\n\n NOTE: This does NOT set up session-specific workspaces.\n Call setup_session_workspace() after provisioning to create a session workspace.\n\n Args:\n sandbox_id: Unique identifier for the sandbox\n user_id: User identifier who owns this sandbox\n tenant_id: Tenant identifier for multi-tenant isolation\n llm_config: LLM provider configuration (for default config)\n\n Returns:\n SandboxInfo with the provisioned sandbox details\n\n Raises:\n RuntimeError: If provisioning fails\n \"\"\"\n ...\n\n @abstractmethod\n def terminate(self, sandbox_id: UUID) -> None:\n \"\"\"Terminate a sandbox and clean up all resources.\n\n Destroys the entire sandbox including all session workspaces.\n Use cleanup_session_workspace() to remove individual sessions.\n\n Args:\n sandbox_id: The sandbox ID to terminate\n \"\"\"\n ...\n\n @abstractmethod\n def setup_session_workspace(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n llm_config: LLMProviderConfig,\n nextjs_port: int,\n file_system_path: str | None = None,\n snapshot_path: str | None = None,\n user_name: str | None = None,\n user_role: str | None = None,\n user_work_area: str | None = None,\n user_level: str | None = None,\n use_demo_data: bool = False,\n excluded_user_library_paths: list[str] | None = None,\n ) -> None:\n \"\"\"Set up a session workspace within an existing sandbox.\n\n Creates the per-session directory structure:\n - sessions/$session_id/outputs/ (from snapshot or template)\n - sessions/$session_id/venv/\n - sessions/$session_id/skills/\n - sessions/$session_id/files/ (symlink to demo data or user files)\n - sessions/$session_id/AGENTS.md\n - sessions/$session_id/opencode.json\n - sessions/$session_id/attachments/\n - sessions/$session_id/org_info/ (if demo data enabled)\n\n Args:\n sandbox_id: The sandbox ID (must be provisioned)\n session_id: The session ID for this workspace\n llm_config: LLM provider configuration for opencode.json\n file_system_path: Path to user's knowledge/source files\n snapshot_path: Optional storage path to restore outputs from\n user_name: User's name for personalization in AGENTS.md\n user_role: User's role/title for personalization in AGENTS.md\n user_work_area: User's work area for demo persona (e.g., \"engineering\")\n user_level: User's level for demo persona (e.g., \"ic\", \"manager\")\n use_demo_data: If True, symlink files/ to demo data; else to user files\n excluded_user_library_paths: List of paths within user_library to exclude\n from the sandbox (e.g., [\"/data/file.xlsx\"]). Only applies when\n use_demo_data=False. Files at these paths won't be accessible.\n\n Raises:\n RuntimeError: If workspace setup fails\n \"\"\"\n ...\n\n @abstractmethod\n def cleanup_session_workspace(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n nextjs_port: int | None = None,\n ) -> None:\n \"\"\"Clean up a session workspace (on session delete).\n\n 1. Stop the Next.js dev server if running on nextjs_port\n 2. Remove the session directory: sessions/$session_id/\n\n Does NOT terminate the sandbox - other sessions may still be using it.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID to clean up\n nextjs_port: Optional port where Next.js server is running\n \"\"\"\n ...\n\n @abstractmethod\n def create_snapshot(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n tenant_id: str,\n ) -> SnapshotResult | None:\n \"\"\"Create a snapshot of a session's outputs and attachments directories.\n\n Captures session-specific user data:\n - sessions/$session_id/outputs/ (generated artifacts, web apps)\n - sessions/$session_id/attachments/ (user uploaded files)\n\n Does NOT include: venv, skills, AGENTS.md, opencode.json, files symlink\n (these are regenerated during restore)\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID to snapshot\n tenant_id: Tenant identifier for storage path\n\n Returns:\n SnapshotResult with storage path and size, or None if:\n - Snapshots are disabled for this backend\n - No outputs directory exists (nothing to snapshot)\n\n Raises:\n RuntimeError: If snapshot creation fails\n \"\"\"\n ...\n\n @abstractmethod\n def restore_snapshot(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n snapshot_storage_path: str,\n tenant_id: str,\n nextjs_port: int,\n llm_config: LLMProviderConfig,\n use_demo_data: bool = False,\n ) -> None:\n \"\"\"Restore a session workspace from a snapshot.\n\n For Kubernetes: Downloads and extracts the snapshot, regenerates config files.\n For Local: No-op since workspaces persist on disk (no snapshots).\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID to restore\n snapshot_storage_path: Path to the snapshot in storage\n tenant_id: Tenant identifier for storage access\n nextjs_port: Port number for the NextJS dev server\n llm_config: LLM provider configuration for opencode.json\n use_demo_data: If True, symlink files/ to demo data\n\n Raises:\n RuntimeError: If snapshot restoration fails\n \"\"\"\n ...\n\n @abstractmethod\n def session_workspace_exists(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n ) -> bool:\n \"\"\"Check if a session's workspace directory exists in the sandbox.\n\n Used to determine if we need to restore from snapshot.\n Checks for sessions/$session_id/outputs/ directory.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID to check\n\n Returns:\n True if the session workspace exists, False otherwise\n \"\"\"\n ...\n\n @abstractmethod\n def health_check(self, sandbox_id: UUID, timeout: float = 60.0) -> bool:\n \"\"\"Check if the sandbox is healthy.\n\n Args:\n sandbox_id: The sandbox ID to check\n\n Returns:\n True if sandbox is healthy, False otherwise\n \"\"\"\n ...\n\n @abstractmethod\n def send_message(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n message: str,\n ) -> Generator[ACPEvent, None, None]:\n \"\"\"Send a message to the CLI agent and stream typed ACP events.\n\n The agent runs in the session-specific workspace:\n sessions/$session_id/\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID (determines workspace directory)\n message: The message content to send\n\n Yields:\n Typed ACP schema event objects\n\n Raises:\n RuntimeError: If agent communication fails\n \"\"\"\n ...\n\n @abstractmethod\n def list_directory(\n self, sandbox_id: UUID, session_id: UUID, path: str\n ) -> list[FilesystemEntry]:\n \"\"\"List contents of a directory in the session's outputs directory.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n path: Relative path within sessions/$session_id/outputs/\n\n Returns:\n List of FilesystemEntry objects sorted by directory first, then name\n\n Raises:\n ValueError: If path traversal attempted or path is not a directory\n \"\"\"\n ...\n\n @abstractmethod\n def read_file(self, sandbox_id: UUID, session_id: UUID, path: str) -> bytes:\n \"\"\"Read a file from the session's workspace.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n path: Relative path within sessions/$session_id/\n\n Returns:\n File contents as bytes\n\n Raises:\n ValueError: If path traversal attempted or path is not a file\n \"\"\"\n ...\n\n @abstractmethod\n def upload_file(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n filename: str,\n content: bytes,\n ) -> str:\n \"\"\"Upload a file to the session's attachments directory.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n filename: Sanitized filename\n content: File content as bytes\n\n Returns:\n Relative path where file was saved (e.g., \"attachments/doc.pdf\")\n\n Raises:\n RuntimeError: If upload fails\n \"\"\"\n ...\n\n @abstractmethod\n def delete_file(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n path: str,\n ) -> bool:\n \"\"\"Delete a file from the session's workspace.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n path: Relative path to the file (e.g., \"attachments/doc.pdf\")\n\n Returns:\n True if file was deleted, False if not found\n\n Raises:\n ValueError: If path traversal attempted\n \"\"\"\n ...\n\n @abstractmethod\n def get_upload_stats(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n ) -> tuple[int, int]:\n \"\"\"Get current file count and total size for a session's attachments.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n\n Returns:\n Tuple of (file_count, total_size_bytes)\n \"\"\"\n ...\n\n @abstractmethod\n def get_webapp_url(self, sandbox_id: UUID, port: int) -> str:\n \"\"\"Get the webapp URL for a session's Next.js server.\n\n Returns the appropriate URL based on the backend:\n - Local: Returns localhost URL with port\n - Kubernetes: Returns internal cluster service URL\n\n Args:\n sandbox_id: The sandbox ID\n port: The session's allocated Next.js port\n\n Returns:\n URL to access the webapp\n \"\"\"\n ...\n\n @abstractmethod\n def generate_pptx_preview(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n pptx_path: str,\n cache_dir: str,\n ) -> tuple[list[str], bool]:\n \"\"\"Convert PPTX to slide JPEG images for preview, with caching.\n\n Checks if cache_dir already has slides. If the PPTX is newer than the\n cached images (or no cache exists), runs soffice -> pdftoppm pipeline.\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n pptx_path: Relative path to the PPTX file within the session workspace\n cache_dir: Relative path for the cache directory\n (e.g., \"outputs/.pptx-preview/abc123\")\n\n Returns:\n Tuple of (slide_paths, cached) where slide_paths is a list of\n relative paths to slide JPEG images (within session workspace)\n and cached indicates whether the result was served from cache.\n\n Raises:\n ValueError: If file not found or conversion fails\n \"\"\"\n ...\n\n @abstractmethod\n def sync_files(\n self,\n sandbox_id: UUID,\n user_id: UUID,\n tenant_id: str,\n source: str | None = None,\n ) -> bool:\n \"\"\"Sync files from S3 to the sandbox's /workspace/files directory.\n\n For Kubernetes backend: Executes `s5cmd sync` in the file-sync sidecar container.\n For Local backend: No-op since files are directly accessible via symlink.\n\n This is idempotent - only downloads changed files. File visibility in\n sessions is controlled via filtered symlinks in setup_session_workspace(),\n not at the sync level.\n\n Args:\n sandbox_id: The sandbox UUID\n user_id: The user ID (for S3 path construction)\n tenant_id: The tenant ID (for S3 path construction)\n source: Optional source type (e.g., \"gmail\", \"google_drive\").\n If None, syncs all sources. If specified, only syncs\n that source's directory.\n\n Returns:\n True if sync was successful, False otherwise.\n \"\"\"\n ...\n\n def ensure_nextjs_running(\n self,\n sandbox_id: UUID,\n session_id: UUID,\n nextjs_port: int,\n ) -> None:\n \"\"\"Ensure the Next.js server is running for a session.\n\n Default is a no-op — only meaningful for local backends that manage\n process lifecycles directly (e.g., LocalSandboxManager).\n\n Args:\n sandbox_id: The sandbox ID\n session_id: The session ID\n nextjs_port: The port the Next.js server should be listening on\n \"\"\"\n\n\n# Singleton instance cache for the factory\n_sandbox_manager_instance: SandboxManager | None = None\n_sandbox_manager_lock = threading.Lock()\n\n\ndef get_sandbox_manager() -> SandboxManager:\n \"\"\"Get the appropriate SandboxManager implementation based on SANDBOX_BACKEND.\n\n Returns:\n SandboxManager instance:\n - LocalSandboxManager for local backend (development)\n - KubernetesSandboxManager for kubernetes backend (production)\n \"\"\"\n global _sandbox_manager_instance\n\n if _sandbox_manager_instance is None:\n with _sandbox_manager_lock:\n if _sandbox_manager_instance is None:\n if SANDBOX_BACKEND == SandboxBackend.LOCAL:\n from onyx.server.features.build.sandbox.local.local_sandbox_manager import (\n LocalSandboxManager,\n )\n\n _sandbox_manager_instance = LocalSandboxManager()\n elif SANDBOX_BACKEND == SandboxBackend.KUBERNETES:\n from onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import (\n KubernetesSandboxManager,\n )\n\n _sandbox_manager_instance = KubernetesSandboxManager()\n logger.info(\"Using KubernetesSandboxManager for sandbox operations\")\n else:\n raise ValueError(f\"Unknown sandbox backend: {SANDBOX_BACKEND}\")\n\n return _sandbox_manager_instance\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/__init__.py", "content": "\"\"\"Kubernetes-based sandbox implementation.\n\nThis module provides the KubernetesSandboxManager for production deployments\nthat run sandboxes as isolated Kubernetes pods.\n\nInternal implementation details (acp_http_client) are in the internal/\nsubdirectory and should not be used directly.\n\"\"\"\n\nfrom onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import (\n KubernetesSandboxManager,\n)\n\n__all__ = [\n \"KubernetesSandboxManager\",\n]\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/Dockerfile", "content": "# Sandbox Container Image\n#\n# User-shared sandbox model:\n# - One pod per user, shared across all user's sessions\n# - Session workspaces created via kubectl exec (setup_session_workspace)\n# - OpenCode agent runs via kubectl exec when needed\n#\n# Directory structure (created by init container + session setup):\n# /workspace/\n# ├── demo_data/ # Demo data (baked into image, for demo sessions)\n# ├── files/ # User's knowledge files (synced from S3)\n# ├── skills/ # Agent skills (baked into image, copied per-session)\n# ├── templates/ # Output templates (baked into image)\n# └── sessions/ # Per-session workspaces (created via exec)\n# └── $session_id/\n# ├── files/ # Symlink to /workspace/demo_data or /workspace/files\n# ├── outputs/\n# ├── AGENTS.md\n# └── opencode.json\n\nFROM node:20-slim\n\n# Install system dependencies\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n python3 \\\n python3-pip \\\n python3-venv \\\n curl \\\n git \\\n procps \\\n unzip \\\n \\\n libreoffice-core \\\n libreoffice-common \\\n libreoffice-impress \\\n libreoffice-draw \\\n poppler-utils \\\n gcc \\\n libc6-dev \\\n fontconfig \\\n fonts-dejavu-core \\\n fonts-liberation \\\n && rm -rf /var/lib/apt/lists/*\n\n# Create non-root user (matches pod securityContext)\n# Handle existing user/group with UID/GID 1000 in base image\nRUN EXISTING_USER=$(id -nu 1000 2>/dev/null || echo \"\"); \\\n EXISTING_GROUP=$(getent group 1000 | cut -d: -f1 2>/dev/null || echo \"\"); \\\n if [ -n \"$EXISTING_GROUP\" ] && [ \"$EXISTING_GROUP\" != \"sandbox\" ]; then \\\n groupmod -n sandbox $EXISTING_GROUP; \\\n elif [ -z \"$EXISTING_GROUP\" ]; then \\\n groupadd -g 1000 sandbox; \\\n fi; \\\n if [ -n \"$EXISTING_USER\" ] && [ \"$EXISTING_USER\" != \"sandbox\" ]; then \\\n usermod -l sandbox -g sandbox $EXISTING_USER; \\\n usermod -d /home/sandbox -m sandbox; \\\n usermod -s /bin/bash sandbox; \\\n elif [ -z \"$EXISTING_USER\" ]; then \\\n useradd -u 1000 -g sandbox -m -s /bin/bash sandbox; \\\n fi\n\n# Create workspace directories\nRUN mkdir -p workspace/sessions /workspace/files /workspace/templates /workspace/demo_data && \\\n chown -R sandbox:sandbox /workspace\n\n# Copy outputs template (web app scaffold, without node_modules)\nCOPY --exclude=.next --exclude=node_modules templates/outputs /workspace/templates/outputs\nRUN chown -R sandbox:sandbox /workspace/templates\n\n# Copy and extract demo data from zip file\n# Zip contains demo_data/ as root folder\nCOPY demo_data.zip /tmp/demo_data.zip\nRUN unzip -q /tmp/demo_data.zip -d /workspace && \\\n rm /tmp/demo_data.zip && \\\n chown -R sandbox:sandbox /workspace/demo_data\n\n# Copy and install Python requirements into a venv\nCOPY initial-requirements.txt /tmp/initial-requirements.txt\nRUN python3 -m venv /workspace/.venv && \\\n /workspace/.venv/bin/pip install --upgrade pip && \\\n /workspace/.venv/bin/pip install -r /tmp/initial-requirements.txt && \\\n rm /tmp/initial-requirements.txt && \\\n chown -R sandbox:sandbox /workspace/.venv\n\n# Add venv to PATH so python/pip use it by default\nENV PATH=\"/workspace/.venv/bin:${PATH}\"\n\n# Install pptxgenjs globally for creating presentations from scratch\nRUN npm install -g pptxgenjs\n\n# Install opencode CLI as sandbox user so it goes to their home directory\nUSER sandbox\nRUN curl -fsSL https://opencode.ai/install | bash\nUSER root\n\n# Add opencode to PATH (installs to ~/.opencode/bin)\nENV PATH=\"/home/sandbox/.opencode/bin:${PATH}\"\n\n# Copy agent skills (symlinked into each session's .opencode/skills/ at setup time)\nCOPY --exclude=__pycache__ skills/ /workspace/skills/\n\n# Set ownership\nRUN chown -R sandbox:sandbox /workspace\n\n# Copy scripts\nCOPY generate_agents_md.py /usr/local/bin/generate_agents_md.py\nRUN chmod +x /usr/local/bin/generate_agents_md.py\n\n# Switch to non-root user\nUSER sandbox\nWORKDIR /workspace\n\n# Expose ports\n# - 3000: Next.js dev server (started per-session if needed)\n# - 8081: OpenCode ACP HTTP server (started via exec)\nEXPOSE 3000 8081\n\n# Keep container alive - all work done via kubectl exec\nCMD [\"sleep\", \"infinity\"]\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/README.md", "content": "# Sandbox Container Image\n\nThis directory contains the Dockerfile and resources for building the Onyx Craft sandbox container image.\n\n## Directory Structure\n\n```\ndocker/\n├── Dockerfile # Main container image definition\n├── demo_data.zip # Demo data (extracted to /workspace/demo_data)\n├── skills/ # Agent skills (image-generation, pptx, etc.)\n├── templates/\n│ └── outputs/ # Web app scaffold template (Next.js)\n├── initial-requirements.txt # Python packages pre-installed in sandbox\n├── generate_agents_md.py # Script to generate AGENTS.md for sessions\n└── README.md # This file\n```\n\n## Building the Image\n\nThe sandbox image must be built for **amd64** architecture since our Kubernetes cluster runs on x86_64 nodes.\n\n### Build for amd64 only (fastest)\n\n```bash\ncd backend/onyx/server/features/build/sandbox/kubernetes/docker\ndocker build --platform linux/amd64 -t onyxdotapp/sandbox:v0.1.x .\ndocker push onyxdotapp/sandbox:v0.1.x\n```\n\n### Build multi-arch (recommended for flexibility)\n\n```bash\ndocker buildx build --platform linux/amd64,linux/arm64 \\\n -t onyxdotapp/sandbox:v0.1.x \\\n --push .\n```\n\n### Update the `latest` tag\n\nAfter pushing a versioned tag, update `latest`:\n\n```bash\ndocker tag onyxdotapp/sandbox:v0.1.x onyxdotapp/sandbox:latest\ndocker push onyxdotapp/sandbox:latest\n```\n\nOr with buildx:\n\n```bash\ndocker buildx build --platform linux/amd64,linux/arm64 \\\n -t onyxdotapp/sandbox:v0.1.x \\\n -t onyxdotapp/sandbox:latest \\\n --push .\n```\n\n## Deploying a New Version\n\n1. **Build and push** the new image (see above)\n\n2. **Update the ConfigMap** in `cloud-deployment-yamls/danswer/configmap/env-configmap.yaml`:\n ```yaml\n SANDBOX_CONTAINER_IMAGE: \"onyxdotapp/sandbox:v0.1.x\"\n ```\n\n3. **Apply the ConfigMap**:\n ```bash\n kubectl apply -f configmap/env-configmap.yaml\n ```\n\n4. **Restart the API server** to pick up the new config:\n ```bash\n kubectl rollout restart deployment/api-server -n danswer\n ```\n\n5. **Delete existing sandbox pods** (they will be recreated with the new image):\n ```bash\n kubectl delete pods -n onyx-sandboxes -l app.kubernetes.io/component=sandbox\n ```\n\n## What's Baked Into the Image\n\n- **Base**: `node:20-slim` (Debian-based)\n- **Demo data**: `/workspace/demo_data/` - sample files for demo sessions\n- **Skills**: `/workspace/skills/` - agent skills (image-generation, pptx, etc.)\n- **Templates**: `/workspace/templates/outputs/` - Next.js web app scaffold\n- **Python venv**: `/workspace/.venv/` with packages from `initial-requirements.txt`\n- **OpenCode CLI**: Installed in `/home/sandbox/.opencode/bin/`\n\n## Runtime Directory Structure\n\nWhen a session is created, the following structure is set up in the pod:\n\n```\n/workspace/\n├── demo_data/ # Baked into image\n├── files/ # Mounted volume, synced from S3\n├── skills/ # Baked into image (agent skills)\n├── templates/ # Baked into image\n└── sessions/\n └── $session_id/\n ├── .opencode/\n │ └── skills/ # Symlink to /workspace/skills\n ├── files/ # Symlink to /workspace/demo_data or /workspace/files\n ├── outputs/ # Copied from templates, contains web app\n ├── attachments/ # User-uploaded files\n ├── org_info/ # Demo persona info (if demo mode)\n ├── AGENTS.md # Instructions for the AI agent\n └── opencode.json # OpenCode configuration\n```\n\n## Troubleshooting\n\n### Verify image exists on Docker Hub\n\n```bash\ncurl -s \"https://hub.docker.com/v2/repositories/onyxdotapp/sandbox/tags\" | jq '.results[].name'\n```\n\n### Check what image a pod is using\n\n```bash\nkubectl get pod -n onyx-sandboxes -o jsonpath='{.spec.containers[?(@.name==\"sandbox\")].image}'\n```\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/generate_agents_md.py", "content": "#!/usr/bin/env python3\n\"\"\"Generate AGENTS.md by scanning the files directory and populating the template.\n\nThis script runs during session setup, AFTER files have been synced from S3\nand the files symlink has been created. It reads an existing AGENTS.md (which\ncontains the {{KNOWLEDGE_SOURCES_SECTION}} placeholder), replaces the\nplaceholder by scanning the knowledge source directory, and writes it back.\n\nUsage:\n python3 generate_agents_md.py \n\nArguments:\n agents_md_path: Path to the AGENTS.md file to update in place\n files_path: Path to the files directory to scan for knowledge sources\n\"\"\"\n\nimport sys\nfrom pathlib import Path\n\n# Type alias for connector info entries\nConnectorInfoEntry = dict[str, str | int]\n\n# Connector information for generating knowledge sources section\n# Keys are normalized (lowercase, underscores) directory names\n# Each entry has: summary (with optional {subdirs}), file_pattern, scan_depth\n# NOTE: This is duplicated from agent_instructions.py to avoid circular imports\nCONNECTOR_INFO: dict[str, ConnectorInfoEntry] = {\n \"google_drive\": {\n \"summary\": \"Documents and files from Google Drive. This may contain information about a user and work they have done.\",\n \"file_pattern\": \"`FILE_NAME.json`\",\n \"scan_depth\": 0,\n },\n \"gmail\": {\n \"summary\": \"Email conversations and threads\",\n \"file_pattern\": \"`FILE_NAME.json`\",\n \"scan_depth\": 0,\n },\n \"linear\": {\n \"summary\": \"Engineering tickets from teams: {subdirs}\",\n \"file_pattern\": \"`[TEAM]/[TICKET_ID]_TICKET_TITLE.json`\",\n \"scan_depth\": 2,\n },\n \"slack\": {\n \"summary\": \"Team messages from channels: {subdirs}\",\n \"file_pattern\": \"`[CHANNEL]/[AUTHOR]_in_[CHANNEL]__[MSG].json`\",\n \"scan_depth\": 1,\n },\n \"github\": {\n \"summary\": \"Pull requests and code from: {subdirs}\",\n \"file_pattern\": \"`[ORG]/[REPO]/pull_requests/[PR_NUMBER]__[PR_TITLE].json`\",\n \"scan_depth\": 2,\n },\n \"fireflies\": {\n \"summary\": \"Meeting transcripts from: {subdirs}\",\n \"file_pattern\": \"`[YYYY-MM]/CALL_TITLE.json`\",\n \"scan_depth\": 1,\n },\n \"hubspot\": {\n \"summary\": \"CRM data including: {subdirs}\",\n \"file_pattern\": \"`[TYPE]/[RECORD_NAME].json`\",\n \"scan_depth\": 1,\n },\n \"notion\": {\n \"summary\": \"Documentation and notes: {subdirs}\",\n \"file_pattern\": \"`PAGE_TITLE.json`\",\n \"scan_depth\": 1,\n },\n \"user_library\": {\n \"summary\": \"User-uploaded files (spreadsheets, documents, presentations, etc.)\",\n \"file_pattern\": \"Any file format\",\n \"scan_depth\": 1,\n },\n}\nDEFAULT_SCAN_DEPTH = 1\n\n\ndef _normalize_connector_name(name: str) -> str:\n \"\"\"Normalize a connector directory name for lookup.\"\"\"\n return name.lower().replace(\" \", \"_\").replace(\"-\", \"_\")\n\n\ndef _scan_directory_to_depth(\n directory: Path, current_depth: int, max_depth: int, indent: str = \" \"\n) -> list[str]:\n \"\"\"Recursively scan directory up to max_depth levels.\"\"\"\n if current_depth >= max_depth:\n return []\n\n lines: list[str] = []\n try:\n subdirs = sorted(\n d for d in directory.iterdir() if d.is_dir() and not d.name.startswith(\".\")\n )\n\n for subdir in subdirs[:10]: # Limit to 10 per level\n lines.append(f\"{indent}- {subdir.name}/\")\n\n # Recurse if we haven't hit max depth\n if current_depth + 1 < max_depth:\n nested = _scan_directory_to_depth(\n subdir, current_depth + 1, max_depth, indent + \" \"\n )\n lines.extend(nested)\n\n if len(subdirs) > 10:\n lines.append(f\"{indent}- ... and {len(subdirs) - 10} more\")\n except Exception:\n pass\n\n return lines\n\n\ndef build_knowledge_sources_section(files_path: Path) -> str:\n \"\"\"Build combined knowledge sources section with summary, structure, and file patterns.\n\n This creates a single section per connector that includes:\n - What kind of data it contains (with actual subdirectory names)\n - The directory structure\n - The file naming pattern\n\n Args:\n files_path: Path to the files directory\n\n Returns:\n Formatted knowledge sources section\n \"\"\"\n if not files_path.exists():\n return \"No knowledge sources available.\"\n\n sections: list[str] = []\n try:\n for item in sorted(files_path.iterdir()):\n if not item.is_dir() or item.name.startswith(\".\"):\n continue\n\n normalized = _normalize_connector_name(item.name)\n info = CONNECTOR_INFO.get(normalized, {})\n\n # Get subdirectory names\n subdirs: list[str] = []\n try:\n subdirs = sorted(\n d.name\n for d in item.iterdir()\n if d.is_dir() and not d.name.startswith(\".\")\n )[:5]\n except Exception:\n pass\n\n # Build summary with subdirs\n summary_template = str(info.get(\"summary\", f\"Data from {item.name}\"))\n if \"{subdirs}\" in summary_template and subdirs:\n subdir_str = \", \".join(subdirs)\n if len(subdirs) == 5:\n subdir_str += \", ...\"\n summary = summary_template.format(subdirs=subdir_str)\n elif \"{subdirs}\" in summary_template:\n summary = summary_template.replace(\": {subdirs}\", \"\").replace(\n \" {subdirs}\", \"\"\n )\n else:\n summary = summary_template\n\n # Build connector section\n file_pattern = str(info.get(\"file_pattern\", \"\"))\n scan_depth = int(info.get(\"scan_depth\", DEFAULT_SCAN_DEPTH))\n\n lines = [f\"### {item.name}/\"]\n lines.append(f\"{summary}.\\n\")\n # Add directory structure if depth > 0\n if scan_depth > 0:\n lines.append(\"Directory structure:\\n\")\n nested = _scan_directory_to_depth(item, 0, scan_depth, \"\")\n if nested:\n lines.append(\"\")\n lines.extend(nested)\n\n lines.append(f\"\\nFile format: {file_pattern}\")\n\n sections.append(\"\\n\".join(lines))\n except Exception as e:\n print(\n f\"Warning: Error building knowledge sources section: {e}\", file=sys.stderr\n )\n return \"Error scanning knowledge sources.\"\n\n if not sections:\n return \"No knowledge sources available.\"\n\n return \"\\n\\n\".join(sections)\n\n\ndef main() -> None:\n \"\"\"Main entry point for container startup script.\n\n Reads an existing AGENTS.md, replaces the {{KNOWLEDGE_SOURCES_SECTION}}\n placeholder by scanning the files directory, and writes it back.\n\n Usage:\n python3 generate_agents_md.py \n \"\"\"\n if len(sys.argv) != 3:\n print(\n f\"Usage: {sys.argv[0]} \",\n file=sys.stderr,\n )\n sys.exit(1)\n\n agents_md_path = Path(sys.argv[1])\n files_path = Path(sys.argv[2])\n\n if not agents_md_path.exists():\n print(f\"Error: {agents_md_path} not found\", file=sys.stderr)\n sys.exit(1)\n\n template = agents_md_path.read_text()\n\n # Resolve symlinks (handles both direct symlinks and dirs containing symlinks)\n resolved_files_path = files_path.resolve()\n\n knowledge_sources_section = build_knowledge_sources_section(resolved_files_path)\n\n # Replace placeholder and write back\n content = template.replace(\n \"{{KNOWLEDGE_SOURCES_SECTION}}\", knowledge_sources_section\n )\n agents_md_path.write_text(content)\n print(f\"Populated knowledge sources in {agents_md_path}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/initial-requirements.txt", "content": "defusedxml>=0.7.1\ngoogle-genai>=1.0.0\nlxml>=5.0.0\nmarkitdown>=0.1.2\nmatplotlib==3.9.1\nmatplotlib-inline>=0.1.7\nmatplotlib-venn>=1.1.2\nnumpy==1.26.4\nopencv-python>=4.11.0.86\nopenpyxl>=3.1.5\npandas==2.2.2\npdfplumber>=0.11.7\nPillow>=10.0.0\npydantic>=2.11.9\npython-pptx>=1.0.2\nscikit-image>=0.25.2\nscikit-learn>=1.7.2\nscipy>=1.16.2\nseaborn>=0.13.2\nxgboost>=3.0.5" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/run-test.sh", "content": "#!/bin/bash\n# Run Kubernetes sandbox integration tests\n#\n# This script:\n# 1. Builds the onyx-backend Docker image\n# 2. Loads it into the kind cluster\n# 3. Deletes/recreates the test pod\n# 4. Waits for the pod to be ready\n# 5. Runs the pytest command inside the pod\n#\n# Usage:\n# ./run-test.sh [test_name]\n#\n# Examples:\n# ./run-test.sh # Run all tests\n# ./run-test.sh test_kubernetes_sandbox_provision # Run specific test\n\nset -e\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\nPROJECT_ROOT=\"$(cd \"$SCRIPT_DIR/../../../../../../../..\" && pwd)\"\nNAMESPACE=\"onyx-sandboxes\"\nPOD_NAME=\"sandbox-test\"\nIMAGE_NAME=\"onyxdotapp/onyx-backend:latest\"\nTEST_FILE=\"onyx/server/features/build/sandbox/kubernetes/test_kubernetes_sandbox.py\"\nENV_FILE=\"$PROJECT_ROOT/.vscode/.env\"\n\nORIGINAL_TEST_FILE=\"$PROJECT_ROOT/backend/tests/external_dependency_unit/craft/test_kubernetes_sandbox.py\"\ncp \"$ORIGINAL_TEST_FILE\" \"$PROJECT_ROOT/backend/$TEST_FILE\"\n\n# Optional: specific test to run\nTEST_NAME=\"${1:-}\"\n\n# Build env var arguments from .vscode/.env file for passing to the container\nENV_VARS=()\nif [ -f \"$ENV_FILE\" ]; then\n echo \"=== Loading environment variables from .vscode/.env ===\"\n while IFS= read -r line || [ -n \"$line\" ]; do\n # Skip empty lines and comments\n [[ -z \"$line\" || \"$line\" =~ ^[[:space:]]*# ]] && continue\n # Skip lines without =\n [[ \"$line\" != *\"=\"* ]] && continue\n # Add to env vars array\n ENV_VARS+=(\"$line\")\n done < \"$ENV_FILE\"\n echo \"Loaded ${#ENV_VARS[@]} environment variables\"\nelse\n echo \"Warning: .vscode/.env not found, running without additional env vars\"\nfi\n\necho \"=== Building onyx-backend Docker image ===\"\ncd \"$PROJECT_ROOT/backend\"\ndocker build -t \"$IMAGE_NAME\" -f Dockerfile .\n\nrm \"$PROJECT_ROOT/backend/$TEST_FILE\"\n\necho \"=== Loading image into kind cluster ===\"\nkind load docker-image \"$IMAGE_NAME\" --name onyx 2>/dev/null || \\\n kind load docker-image \"$IMAGE_NAME\" 2>/dev/null || \\\n echo \"Warning: Could not load into kind. If using minikube, run: minikube image load $IMAGE_NAME\"\n\necho \"=== Deleting existing test pod (if any) ===\"\nkubectl delete pod \"$POD_NAME\" -n \"$NAMESPACE\" --ignore-not-found=true\n\necho \"=== Creating test pod ===\"\nkubectl apply -f \"$SCRIPT_DIR/test-job.yaml\"\n\necho \"=== Waiting for pod to be ready ===\"\nkubectl wait --for=condition=Ready pod/\"$POD_NAME\" -n \"$NAMESPACE\" --timeout=120s\n\necho \"=== Running tests ===\"\nif [ -n \"$TEST_NAME\" ]; then\n kubectl exec -it \"$POD_NAME\" -n \"$NAMESPACE\" -- \\\n env \"${ENV_VARS[@]}\" pytest \"$TEST_FILE::$TEST_NAME\" -v -s\nelse\n kubectl exec -it \"$POD_NAME\" -n \"$NAMESPACE\" -- \\\n env \"${ENV_VARS[@]}\" pytest \"$TEST_FILE\" -v -s\nfi\n\necho \"=== Tests complete ===\"\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/image-generation/SKILL.md", "content": "---\nname: image-generation\ndescription: Generate images using nano banana.\n---\n\n# Image Generation Skill\n\nGenerate images using Nano Banana (Google Gemini Image API). Supports text-to-image and image-to-image generation with configurable options.\n\n## Setup\n\n### Dependencies\n\n```bash\npip install google-genai Pillow\n```\n\n### Environment Variable\n\nSet your API key:\n\n```bash\nexport GEMINI_API_KEY=\"your_api_key_here\"\n```\n\n## Usage\n\n### Basic Text-to-Image\n\n```bash\npython scripts/generate.py --prompt \"A futuristic city at sunset with neon lights\" --output city.png\n```\n\n### With Aspect Ratio\n\n```bash\npython scripts/generate.py \\\n --prompt \"Mountain landscape with a lake\" \\\n --output landscape.png \\\n --aspect-ratio 16:9\n```\n\n### Image-to-Image Mode\n\nUse a reference image to guide generation:\n\n```bash\npython scripts/generate.py \\\n --prompt \"Make it look like a watercolor painting\" \\\n --input-image original.png \\\n --output watercolor.png\n```\n\n### Generate Multiple Images\n\n```bash\npython scripts/generate.py \\\n --prompt \"Abstract colorful art\" \\\n --output art.png \\\n --num-images 3\n```\n\n## Arguments\n\n| Argument | Short | Required | Default | Description |\n|----------|-------|----------|---------|-------------|\n| `--prompt` | `-p` | Yes | — | Text prompt describing the desired image |\n| `--output` | `-o` | No | `output.png` | Output path for the generated image |\n| `--model` | `-m` | No | `gemini-2.0-flash-preview-image-generation` | Model to use for generation |\n| `--input-image` | `-i` | No | — | Reference image for image-to-image mode |\n| `--aspect-ratio` | `-a` | No | — | Aspect ratio: `1:1`, `16:9`, `9:16`, `4:3`, `3:4` |\n| `--num-images` | `-n` | No | `1` | Number of images to generate |\n\n## Available Models\n\n- `gemini-2.0-flash-preview-image-generation` - Fast, optimized for speed and lower latency\n- `imagen-3.0-generate-002` - High quality image generation\n\n## Programmatic Usage\n\nImport the function directly in Python:\n\n```python\nfrom scripts.generate import generate_image\n\npaths = generate_image(\n prompt=\"A serene mountain lake under moonlight\",\n output_path=\"./outputs/lake.png\",\n aspect_ratio=\"16:9\",\n num_images=2,\n)\n```\n\n## Tips\n\n- **Detailed prompts work better**: Instead of \"a cat\", try \"a fluffy orange tabby cat sitting on a windowsill, soft morning light, photorealistic\"\n- **Specify style**: Include style keywords like \"digital art\", \"oil painting\", \"photorealistic\", \"anime style\"\n- **Use aspect ratios**: Match the aspect ratio to your intended use (16:9 for landscapes, 9:16 for portraits/mobile)\n- **Image-to-image**: Great for style transfer, variations, or guided modifications of existing images\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/image-generation/scripts/generate.py", "content": "#!/usr/bin/env python3\n\"\"\"\nImage generation script using Nano Banana (Google Gemini Image API).\n\nSupports text-to-image and image-to-image generation with configurable options.\n\"\"\"\n\nimport argparse\nimport base64\nimport os\nimport sys\nfrom io import BytesIO\nfrom pathlib import Path\n\nfrom PIL import Image\n\n\ndef load_image_as_base64(image_path: str) -> tuple[str, str]:\n \"\"\"Load an image file and return base64 data and mime type.\"\"\"\n path = Path(image_path)\n if not path.exists():\n raise FileNotFoundError(f\"Image not found: {image_path}\")\n\n # Determine mime type from extension\n ext = path.suffix.lower()\n mime_types = {\n \".png\": \"image/png\",\n \".jpg\": \"image/jpeg\",\n \".jpeg\": \"image/jpeg\",\n \".gif\": \"image/gif\",\n \".webp\": \"image/webp\",\n }\n mime_type = mime_types.get(ext, \"image/png\")\n\n with open(image_path, \"rb\") as f:\n data = base64.b64encode(f.read()).decode(\"utf-8\")\n\n return data, mime_type\n\n\ndef generate_image(\n prompt: str,\n output_path: str,\n model: str = \"gemini-3-pro-image-preview\",\n input_image: str | None = None,\n aspect_ratio: str | None = None, # noqa: ARG001\n num_images: int = 1,\n) -> list[str]:\n \"\"\"\n Generate image(s) using Google Gemini / Nano Banana API.\n\n Args:\n prompt: Text description for image generation.\n output_path: Path to save the generated image(s).\n model: Model ID to use for generation.\n input_image: Optional path to reference image for image-to-image mode.\n aspect_ratio: Aspect ratio (e.g., \"1:1\", \"16:9\", \"9:16\", \"4:3\", \"3:4\").\n num_images: Number of images to generate.\n\n Returns:\n List of paths to saved images.\n \"\"\"\n api_key = os.environ.get(\"GEMINI_API_KEY\") or os.environ.get(\"GENAI_API_KEY\")\n if not api_key:\n raise ValueError(\n \"API key not found. Set GEMINI_API_KEY or GENAI_API_KEY environment variable.\"\n )\n\n # lazy importing since very heavy libs\n from google import genai\n from google.genai import types\n\n client = genai.Client(api_key=api_key)\n\n # Build content parts\n parts: list[types.Part] = []\n\n # Add reference image if provided (image-to-image mode)\n if input_image:\n img_data, mime_type = load_image_as_base64(input_image)\n parts.append(\n types.Part.from_bytes(\n data=base64.b64decode(img_data),\n mime_type=mime_type,\n )\n )\n\n # Add text prompt\n parts.append(types.Part.from_text(text=prompt))\n\n # Build generation config\n generate_config = types.GenerateContentConfig(\n response_modalities=[\"TEXT\", \"IMAGE\"],\n )\n\n saved_paths: list[str] = []\n output_dir = Path(output_path).parent\n output_dir.mkdir(parents=True, exist_ok=True)\n\n base_name = Path(output_path).stem\n extension = Path(output_path).suffix or \".png\"\n\n for i in range(num_images):\n response = client.models.generate_content(\n model=model,\n contents=types.Content(parts=parts),\n config=generate_config,\n )\n\n # Validate response\n if not response.candidates:\n raise ValueError(\"No candidates returned from the API\")\n\n candidate = response.candidates[0]\n if not candidate.content or not candidate.content.parts:\n raise ValueError(\"No content parts returned from the API\")\n\n # Process response parts\n image_count = 0\n for part in candidate.content.parts:\n if part.inline_data is not None and part.inline_data.data is not None:\n # Extract and save the image\n image_data = part.inline_data.data\n image = Image.open(BytesIO(image_data))\n\n # Generate output filename\n if num_images == 1 and image_count == 0:\n save_path = output_path\n else:\n save_path = str(\n output_dir / f\"{base_name}_{i + 1}_{image_count + 1}{extension}\"\n )\n\n image.save(save_path)\n saved_paths.append(save_path)\n print(f\"Saved: {save_path}\")\n image_count += 1\n elif part.text:\n # Print any text response from the model\n print(f\"Model response: {part.text}\")\n\n return saved_paths\n\n\ndef main() -> None:\n \"\"\"Main entry point for CLI usage.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Generate images using Nano Banana (Google Gemini Image API).\",\n formatter_class=argparse.RawDescriptionHelpFormatter,\n epilog=\"\"\"\nExamples:\n # Basic text-to-image generation\n python generate.py --prompt \"A futuristic city at sunset\" --output city.png\n\n # Generate with specific aspect ratio\n python generate.py --prompt \"Mountain landscape\" --output landscape.png --aspect-ratio 16:9\n\n # Image-to-image mode (use reference image)\n python generate.py --prompt \"Make it more colorful\" --input-image ref.png --output colorful.png\n\n # Generate multiple images\n python generate.py --prompt \"Abstract art\" --output art.png --num-images 3\n\"\"\",\n )\n\n parser.add_argument(\n \"--prompt\",\n \"-p\",\n type=str,\n required=True,\n help=\"Text prompt describing the desired image.\",\n )\n parser.add_argument(\n \"--output\",\n \"-o\",\n type=str,\n default=\"output.png\",\n help=\"Output path for the generated image (default: output.png).\",\n )\n parser.add_argument(\n \"--model\",\n \"-m\",\n type=str,\n default=\"gemini-3-pro-image-preview\",\n help=\"Model to use (default: gemini-3-pro-image-preview).\",\n )\n parser.add_argument(\n \"--input-image\",\n \"-i\",\n type=str,\n help=\"Path to reference image for image-to-image generation.\",\n )\n parser.add_argument(\n \"--aspect-ratio\",\n \"-a\",\n type=str,\n choices=[\"1:1\", \"16:9\", \"9:16\", \"4:3\", \"3:4\"],\n help=\"Aspect ratio for the generated image.\",\n )\n parser.add_argument(\n \"--num-images\",\n \"-n\",\n type=int,\n default=1,\n help=\"Number of images to generate (default: 1).\",\n )\n\n args = parser.parse_args()\n\n try:\n saved_paths = generate_image(\n prompt=args.prompt,\n output_path=args.output,\n model=args.model,\n input_image=args.input_image,\n aspect_ratio=args.aspect_ratio,\n num_images=args.num_images,\n )\n\n print(f\"\\nSuccessfully generated {len(saved_paths)} image(s):\")\n for path in saved_paths:\n print(f\" - {path}\")\n\n except Exception as e:\n print(f\"Error: {e}\", file=sys.stderr)\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/SKILL.md", "content": "---\nname: pptx\ndescription: \"Use this skill any time a .pptx file is involved in any way — as input, output, or both. This includes: creating slide decks, pitch decks, or presentations; reading, parsing, or extracting text from any .pptx file (even if the extracted content will be used elsewhere, like in an email or summary); editing, modifying, or updating existing presentations; combining or splitting slide files; working with templates, layouts, speaker notes, or comments. Trigger whenever the user mentions \\\"deck,\\\" \\\"slides,\\\" \\\"presentation,\\\" or references a .pptx filename, regardless of what they plan to do with the content afterward. If a .pptx file needs to be opened, created, or touched, use this skill.\"\nlicense: Proprietary. LICENSE.txt has complete terms\n---\n\n# PPTX Skill\n\n> **Path convention**: All commands run from the **session workspace** (your working directory). Never `cd` into the skill directory. Prefix all skill scripts with `.opencode/skills/pptx/`. All generated files (unpacked dirs, output presentations, thumbnails, PDFs, images) go in `outputs/`.\n\n## Quick Reference\n\n| Task | Guide |\n|------|-------|\n| Read/analyze content | `python -m markitdown presentation.pptx` |\n| Edit or create from template | Read [editing.md](editing.md) |\n| Create from scratch | Read [pptxgenjs.md](pptxgenjs.md) |\n\n---\n\n## Reading Content\n\n```bash\n# Text extraction\npython -m markitdown presentation.pptx\n\n# Visual overview\npython .opencode/skills/pptx/scripts/thumbnail.py presentation.pptx\n\n# Raw XML\npython .opencode/skills/pptx/scripts/office/unpack.py presentation.pptx outputs/unpacked/\n```\n\n---\n\n## Editing Workflow\n\n**Read [editing.md](editing.md) for full details.**\n\n1. Analyze template with `thumbnail.py`\n2. Unpack → manipulate slides → edit content → clean → pack\n\n---\n\n## Creating from Scratch\n\n**Read [pptxgenjs.md](pptxgenjs.md) for full details.**\n\nUse when no template or reference presentation is available.\n\n---\n\n## Design Ideas\n\n**Don't create boring slides.** Plain bullets on a white background won't impress anyone. Consider ideas from this list for each slide.\n\n### Before Starting\n\n- **Pick a bold, content-informed color palette**: The palette should feel designed for THIS topic. If swapping your colors into a completely different presentation would still \"work,\" you haven't made specific enough choices.\n- **Dominance over equality**: One color should dominate (60-70% visual weight), with 1-2 supporting tones and one sharp accent. Never give all colors equal weight.\n- **Dark/light contrast**: Dark backgrounds for title + conclusion slides, light for content (\"sandwich\" structure). Or commit to dark throughout for a premium feel.\n- **Commit to a visual motif**: Pick ONE distinctive element and repeat it — rounded image frames, icons in colored circles, thick single-side borders. Carry it across every slide.\n\n### Color Palettes\n\nChoose colors that match your topic — don't default to generic blue. Use these palettes as inspiration:\n\n| Theme | Primary | Secondary | Accent |\n|-------|---------|-----------|--------|\n| **Midnight Executive** | `1E2761` (navy) | `CADCFC` (ice blue) | `FFFFFF` (white) |\n| **Forest & Moss** | `2C5F2D` (forest) | `97BC62` (moss) | `F5F5F5` (cream) |\n| **Coral Energy** | `F96167` (coral) | `F9E795` (gold) | `2F3C7E` (navy) |\n| **Warm Terracotta** | `B85042` (terracotta) | `E7E8D1` (sand) | `A7BEAE` (sage) |\n| **Ocean Gradient** | `065A82` (deep blue) | `1C7293` (teal) | `21295C` (midnight) |\n| **Charcoal Minimal** | `36454F` (charcoal) | `F2F2F2` (off-white) | `212121` (black) |\n| **Teal Trust** | `028090` (teal) | `00A896` (seafoam) | `02C39A` (mint) |\n| **Berry & Cream** | `6D2E46` (berry) | `A26769` (dusty rose) | `ECE2D0` (cream) |\n| **Sage Calm** | `84B59F` (sage) | `69A297` (eucalyptus) | `50808E` (slate) |\n| **Cherry Bold** | `990011` (cherry) | `FCF6F5` (off-white) | `2F3C7E` (navy) |\n\n### For Each Slide\n\n**Every slide needs a visual element** — image, chart, icon, or shape. Text-only slides are forgettable.\n\n**Layout options:**\n- Two-column (text left, illustration on right)\n- Icon + text rows (icon in colored circle, bold header, description below)\n- 2x2 or 2x3 grid (image on one side, grid of content blocks on other)\n- Half-bleed image (full left or right side) with content overlay\n\n**Data display:**\n- Large stat callouts (big numbers 60-72pt with small labels below)\n- Comparison columns (before/after, pros/cons, side-by-side options)\n- Timeline or process flow (numbered steps, arrows)\n\n**Visual polish:**\n- Icons in small colored circles next to section headers\n- Italic accent text for key stats or taglines\n\n### Typography\n\n**Choose an interesting font pairing** — don't default to Arial. Pick a header font with personality and pair it with a clean body font.\n\n| Header Font | Body Font |\n|-------------|-----------|\n| Georgia | Calibri |\n| Arial Black | Arial |\n| Calibri | Calibri Light |\n| Cambria | Calibri |\n| Trebuchet MS | Calibri |\n| Impact | Arial |\n| Palatino | Garamond |\n| Consolas | Calibri |\n\n| Element | Size |\n|---------|------|\n| Slide title | 36-44pt bold |\n| Section header | 20-24pt bold |\n| Body text | 14-16pt |\n| Captions | 10-12pt muted |\n\n### Spacing\n\n- 0.5\" minimum margins\n- 0.3-0.5\" between content blocks\n- Leave breathing room—don't fill every inch\n\n### Avoid (Common Mistakes)\n\n- **Don't repeat the same layout** — vary columns, cards, and callouts across slides\n- **Don't center body text** — left-align paragraphs and lists; center only titles\n- **Don't skimp on size contrast** — titles need 36pt+ to stand out from 14-16pt body\n- **Don't default to blue** — pick colors that reflect the specific topic\n- **Don't mix spacing randomly** — choose 0.3\" or 0.5\" gaps and use consistently\n- **Don't style one slide and leave the rest plain** — commit fully or keep it simple throughout\n- **Don't create text-only slides** — add images, icons, charts, or visual elements; avoid plain title + bullets\n- **Don't forget text box padding** — when aligning lines or shapes with text edges, set `margin: 0` on the text box or offset the shape to account for padding\n- **Don't use low-contrast elements** — icons AND text need strong contrast against the background; avoid light text on light backgrounds or dark text on dark backgrounds\n- **NEVER use accent lines under titles** — these are a hallmark of AI-generated slides; use whitespace or background color instead\n\n---\n\n## QA (Required)\n\n**Assume there are problems. Your job is to find them.**\n\nYour first render is almost never correct. Approach QA as a bug hunt, not a confirmation step. If you found zero issues on first inspection, you weren't looking hard enough.\n\n### Content QA\n\n```bash\npython -m markitdown output.pptx\n```\n\nCheck for missing content, typos, wrong order.\n\n**When using templates, check for leftover placeholder text:**\n\n```bash\npython -m markitdown output.pptx | grep -iE \"xxxx|lorem|ipsum|this.*(page|slide).*layout\"\n```\n\nIf grep returns results, fix them before declaring success.\n\n### Visual QA\n\n**⚠️ USE SUBAGENTS** — even for 2-3 slides. You've been staring at the code and will see what you expect, not what's there. Subagents have fresh eyes.\n\nConvert slides to images (see [Converting to Images](#converting-to-images)), then use this prompt:\n\n```\nVisually inspect these slides. Assume there are issues — find them.\n\nLook for:\n- Overlapping elements (text through shapes, lines through words, stacked elements)\n- Text overflow or cut off at edges/box boundaries\n- Decorative lines positioned for single-line text but title wrapped to two lines\n- Source citations or footers colliding with content above\n- Elements too close (< 0.3\" gaps) or cards/sections nearly touching\n- Uneven gaps (large empty area in one place, cramped in another)\n- Insufficient margin from slide edges (< 0.5\")\n- Columns or similar elements not aligned consistently\n- Low-contrast text (e.g., light gray text on cream-colored background)\n- Low-contrast icons (e.g., dark icons on dark backgrounds without a contrasting circle)\n- Text boxes too narrow causing excessive wrapping\n- Leftover placeholder content\n\nFor each slide, list issues or areas of concern, even if minor.\n\nRead and analyze these images:\n1. /path/to/slide-01.jpg (Expected: [brief description])\n2. /path/to/slide-02.jpg (Expected: [brief description])\n\nReport ALL issues found, including minor ones.\n```\n\n### Verification Loop\n\n1. Generate slides → Convert to images → Inspect\n2. **List issues found** (if none found, look again more critically)\n3. Fix issues\n4. **Re-verify affected slides** — one fix often creates another problem\n5. Repeat until a full pass reveals no new issues\n\n**Do not declare success until you've completed at least one fix-and-verify cycle.**\n\n---\n\n## Converting to Images\n\nConvert presentations to individual slide images for visual inspection:\n\n```bash\npython .opencode/skills/pptx/scripts/office/soffice.py --headless --convert-to pdf outputs/output.pptx\npdftoppm -jpeg -r 150 outputs/output.pdf outputs/slide\n```\n\nThis creates `slide-01.jpg`, `slide-02.jpg`, etc.\n\nTo re-render specific slides after fixes:\n\n```bash\npdftoppm -jpeg -r 150 -f N -l N outputs/output.pdf outputs/slide-fixed\n```\n\n---\n\n## Dependencies\n\n- `pip install \"markitdown[pptx]\"` - text extraction\n- `pip install Pillow` - thumbnail grids\n- `npm install -g pptxgenjs` - creating from scratch\n- LibreOffice (`soffice`) - PDF conversion (auto-configured for sandboxed environments via `.opencode/skills/pptx/scripts/office/soffice.py`)\n- Poppler (`pdftoppm`) - PDF to images\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/editing.md", "content": "# Editing Presentations\n\n> **Path convention**: All commands run from the **session workspace**. Never `cd` into the skill directory. Prefix all skill scripts with `.opencode/skills/pptx/`. All generated files go in `outputs/`.\n\n## Template-Based Workflow\n\nWhen using an existing presentation as a template:\n\n1. **Analyze existing slides**:\n ```bash\n python .opencode/skills/pptx/scripts/thumbnail.py template.pptx outputs/thumbnails\n python -m markitdown template.pptx\n ```\n Review `outputs/thumbnails.jpg` to see layouts, and markitdown output to see placeholder text.\n\n2. **Plan slide mapping**: For each content section, choose a template slide.\n\n ⚠️ **USE VARIED LAYOUTS** — monotonous presentations are a common failure mode. Don't default to basic title + bullet slides. Actively seek out:\n - Multi-column layouts (2-column, 3-column)\n - Image + text combinations\n - Full-bleed images with text overlay\n - Quote or callout slides\n - Section dividers\n - Stat/number callouts\n - Icon grids or icon + text rows\n\n **Avoid:** Repeating the same text-heavy layout for every slide.\n\n Match content type to layout style (e.g., key points → bullet slide, team info → multi-column, testimonials → quote slide).\n\n3. **Unpack**: `python .opencode/skills/pptx/scripts/office/unpack.py template.pptx outputs/unpacked/`\n\n4. **Build presentation** (do this yourself, not with subagents):\n - Delete unwanted slides (remove from ``)\n - Duplicate slides you want to reuse (`add_slide.py`)\n - Reorder slides in ``\n - **Complete all structural changes before step 5**\n\n5. **Edit content**: Update text in each `slide{N}.xml`.\n **Use subagents here if available** — slides are separate XML files, so subagents can edit in parallel.\n\n6. **Clean**: `python .opencode/skills/pptx/scripts/clean.py outputs/unpacked/`\n\n7. **Pack**: `python .opencode/skills/pptx/scripts/office/pack.py outputs/unpacked/ outputs/output.pptx --original template.pptx`\n\n---\n\n## Scripts\n\n| Script | Purpose |\n|--------|---------|\n| `unpack.py` | Extract and pretty-print PPTX |\n| `add_slide.py` | Duplicate slide or create from layout |\n| `clean.py` | Remove orphaned files |\n| `pack.py` | Repack with validation |\n| `thumbnail.py` | Create visual grid of slides |\n\n### unpack.py\n\n```bash\npython .opencode/skills/pptx/scripts/office/unpack.py input.pptx outputs/unpacked/\n```\n\nExtracts PPTX, pretty-prints XML, escapes smart quotes.\n\n### add_slide.py\n\n```bash\npython .opencode/skills/pptx/scripts/add_slide.py outputs/unpacked/ slide2.xml # Duplicate slide\npython .opencode/skills/pptx/scripts/add_slide.py outputs/unpacked/ slideLayout2.xml # From layout\n```\n\nPrints `` to add to `` at desired position.\n\n### clean.py\n\n```bash\npython .opencode/skills/pptx/scripts/clean.py outputs/unpacked/\n```\n\nRemoves slides not in ``, unreferenced media, orphaned rels.\n\n### pack.py\n\n```bash\npython .opencode/skills/pptx/scripts/office/pack.py outputs/unpacked/ outputs/output.pptx --original input.pptx\n```\n\nValidates, repairs, condenses XML, re-encodes smart quotes.\n\n### thumbnail.py\n\n```bash\npython .opencode/skills/pptx/scripts/thumbnail.py input.pptx outputs/thumbnails [--cols N]\n```\n\nCreates `outputs/thumbnails.jpg` with slide filenames as labels. Default 3 columns, max 12 per grid.\n\n**Use for template analysis only** (choosing layouts). For visual QA, use `soffice` + `pdftoppm` to create full-resolution individual slide images—see SKILL.md.\n\n---\n\n## Slide Operations\n\nSlide order is in `outputs/unpacked/ppt/presentation.xml` → ``.\n\n**Reorder**: Rearrange `` elements.\n\n**Delete**: Remove ``, then run `clean.py`.\n\n**See available layouts**: `ls outputs/unpacked/ppt/slideLayouts/`\n\n**Add**: Use `add_slide.py`. Never manually copy slide files—the script handles notes references, Content_Types.xml, and relationship IDs that manual copying misses.\n\n---\n\n## Editing Content\n\n**Subagents:** If available, use them here (after completing step 4). Each slide is a separate XML file, so subagents can edit in parallel. In your prompt to subagents, include:\n- The slide file path(s) to edit\n- **\"Use the Edit tool for all changes\"**\n- The formatting rules and common pitfalls below\n\nFor each slide:\n1. Read the slide's XML\n2. Identify ALL placeholder content—text, images, charts, icons, captions\n3. Replace each placeholder with final content\n\n**Use the Edit tool, not sed or Python scripts.** The Edit tool forces specificity about what to replace and where, yielding better reliability.\n\n### Formatting Rules\n\n- **Bold all headers, subheadings, and inline labels**: Use `b=\"1\"` on ``. This includes:\n - Slide titles\n - Section headers within a slide\n - Inline labels like (e.g.: \"Status:\", \"Description:\") at the start of a line\n- **Never use unicode bullets (•)**: Use proper list formatting with `` or ``\n- **Bullet consistency**: Let bullets inherit from the layout. Only specify `` or ``.\n\n---\n\n## Common Pitfalls\n\n### Template Adaptation\n\nWhen source content has fewer items than the template:\n- **Remove excess elements entirely** (images, shapes, text boxes), don't just clear text\n- Check for orphaned visuals after clearing text content\n- Run visual QA to catch mismatched counts\n\nWhen replacing text with different length content:\n- **Shorter replacements**: Usually safe\n- **Longer replacements**: May overflow or wrap unexpectedly\n- Test with visual QA after text changes\n- Consider truncating or splitting content to fit the template's design constraints\n\n**Template slots ≠ Source items**: If template has 4 team members but source has 3 users, delete the 4th member's entire group (image + text boxes), not just the text.\n\n### Multi-Item Content\n\nIf source has multiple items (numbered lists, multiple sections), create separate `` elements for each — **never concatenate into one string**.\n\n**❌ WRONG** — all items in one paragraph:\n```xml\n\n Step 1: Do the first thing. Step 2: Do the second thing.\n\n```\n\n**✅ CORRECT** — separate paragraphs with bold headers:\n```xml\n\n \n Step 1\n\n\n \n Do the first thing.\n\n\n \n Step 2\n\n\n```\n\nCopy `` from the original paragraph to preserve line spacing. Use `b=\"1\"` on headers.\n\n### Smart Quotes\n\nHandled automatically by unpack/pack. But the Edit tool converts smart quotes to ASCII.\n\n**When adding new text with quotes, use XML entities:**\n\n```xml\nthe “Agreement”\n```\n\n| Character | Name | Unicode | XML Entity |\n|-----------|------|---------|------------|\n| `“` | Left double quote | U+201C | `“` |\n| `”` | Right double quote | U+201D | `”` |\n| `‘` | Left single quote | U+2018 | `‘` |\n| `’` | Right single quote | U+2019 | `’` |\n\n### Other\n\n- **Whitespace**: Use `xml:space=\"preserve\"` on `` with leading/trailing spaces\n- **XML parsing**: Use `defusedxml.minidom`, not `xml.etree.ElementTree` (corrupts namespaces)\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/pptxgenjs.md", "content": "# PptxGenJS Tutorial\n\n## Setup & Basic Structure\n\n```javascript\nconst pptxgen = require(\"pptxgenjs\");\n\nlet pres = new pptxgen();\npres.layout = 'LAYOUT_16x9'; // or 'LAYOUT_16x10', 'LAYOUT_4x3', 'LAYOUT_WIDE'\npres.author = 'Your Name';\npres.title = 'Presentation Title';\n\nlet slide = pres.addSlide();\nslide.addText(\"Hello World!\", { x: 0.5, y: 0.5, fontSize: 36, color: \"363636\" });\n\npres.writeFile({ fileName: \"Presentation.pptx\" });\n```\n\n## Layout Dimensions\n\nSlide dimensions (coordinates in inches):\n- `LAYOUT_16x9`: 10\" × 5.625\" (default)\n- `LAYOUT_16x10`: 10\" × 6.25\"\n- `LAYOUT_4x3`: 10\" × 7.5\"\n- `LAYOUT_WIDE`: 13.3\" × 7.5\"\n\n---\n\n## Text & Formatting\n\n```javascript\n// Basic text\nslide.addText(\"Simple Text\", {\n x: 1, y: 1, w: 8, h: 2, fontSize: 24, fontFace: \"Arial\",\n color: \"363636\", bold: true, align: \"center\", valign: \"middle\"\n});\n\n// Character spacing (use charSpacing, not letterSpacing which is silently ignored)\nslide.addText(\"SPACED TEXT\", { x: 1, y: 1, w: 8, h: 1, charSpacing: 6 });\n\n// Rich text arrays\nslide.addText([\n { text: \"Bold \", options: { bold: true } },\n { text: \"Italic \", options: { italic: true } }\n], { x: 1, y: 3, w: 8, h: 1 });\n\n// Multi-line text (requires breakLine: true)\nslide.addText([\n { text: \"Line 1\", options: { breakLine: true } },\n { text: \"Line 2\", options: { breakLine: true } },\n { text: \"Line 3\" } // Last item doesn't need breakLine\n], { x: 0.5, y: 0.5, w: 8, h: 2 });\n\n// Text box margin (internal padding)\nslide.addText(\"Title\", {\n x: 0.5, y: 0.3, w: 9, h: 0.6,\n margin: 0 // Use 0 when aligning text with other elements like shapes or icons\n});\n```\n\n**Tip:** Text boxes have internal margin by default. Set `margin: 0` when you need text to align precisely with shapes, lines, or icons at the same x-position.\n\n---\n\n## Lists & Bullets\n\n```javascript\n// ✅ CORRECT: Multiple bullets\nslide.addText([\n { text: \"First item\", options: { bullet: true, breakLine: true } },\n { text: \"Second item\", options: { bullet: true, breakLine: true } },\n { text: \"Third item\", options: { bullet: true } }\n], { x: 0.5, y: 0.5, w: 8, h: 3 });\n\n// ❌ WRONG: Never use unicode bullets\nslide.addText(\"• First item\", { ... }); // Creates double bullets\n\n// Sub-items and numbered lists\n{ text: \"Sub-item\", options: { bullet: true, indentLevel: 1 } }\n{ text: \"First\", options: { bullet: { type: \"number\" }, breakLine: true } }\n```\n\n---\n\n## Shapes\n\n```javascript\nslide.addShape(pres.shapes.RECTANGLE, {\n x: 0.5, y: 0.8, w: 1.5, h: 3.0,\n fill: { color: \"FF0000\" }, line: { color: \"000000\", width: 2 }\n});\n\nslide.addShape(pres.shapes.OVAL, { x: 4, y: 1, w: 2, h: 2, fill: { color: \"0000FF\" } });\n\nslide.addShape(pres.shapes.LINE, {\n x: 1, y: 3, w: 5, h: 0, line: { color: \"FF0000\", width: 3, dashType: \"dash\" }\n});\n\n// With transparency\nslide.addShape(pres.shapes.RECTANGLE, {\n x: 1, y: 1, w: 3, h: 2,\n fill: { color: \"0088CC\", transparency: 50 }\n});\n\n// Rounded rectangle (rectRadius only works with ROUNDED_RECTANGLE, not RECTANGLE)\n// ⚠️ Don't pair with rectangular accent overlays — they won't cover rounded corners. Use RECTANGLE instead.\nslide.addShape(pres.shapes.ROUNDED_RECTANGLE, {\n x: 1, y: 1, w: 3, h: 2,\n fill: { color: \"FFFFFF\" }, rectRadius: 0.1\n});\n\n// With shadow\nslide.addShape(pres.shapes.RECTANGLE, {\n x: 1, y: 1, w: 3, h: 2,\n fill: { color: \"FFFFFF\" },\n shadow: { type: \"outer\", color: \"000000\", blur: 6, offset: 2, angle: 135, opacity: 0.15 }\n});\n```\n\nShadow options:\n\n| Property | Type | Range | Notes |\n|----------|------|-------|-------|\n| `type` | string | `\"outer\"`, `\"inner\"` | |\n| `color` | string | 6-char hex (e.g. `\"000000\"`) | No `#` prefix, no 8-char hex — see Common Pitfalls |\n| `blur` | number | 0-100 pt | |\n| `offset` | number | 0-200 pt | **Must be non-negative** — negative values corrupt the file |\n| `angle` | number | 0-359 degrees | Direction the shadow falls (135 = bottom-right, 270 = upward) |\n| `opacity` | number | 0.0-1.0 | Use this for transparency, never encode in color string |\n\nTo cast a shadow upward (e.g. on a footer bar), use `angle: 270` with a positive offset — do **not** use a negative offset.\n\n**Note**: Gradient fills are not natively supported. Use a gradient image as a background instead.\n\n---\n\n## Images\n\n### Image Sources\n\n```javascript\n// From file path\nslide.addImage({ path: \"images/chart.png\", x: 1, y: 1, w: 5, h: 3 });\n\n// From URL\nslide.addImage({ path: \"https://example.com/image.jpg\", x: 1, y: 1, w: 5, h: 3 });\n\n// From base64 (faster, no file I/O)\nslide.addImage({ data: \"image/png;base64,iVBORw0KGgo...\", x: 1, y: 1, w: 5, h: 3 });\n```\n\n### Image Options\n\n```javascript\nslide.addImage({\n path: \"image.png\",\n x: 1, y: 1, w: 5, h: 3,\n rotate: 45, // 0-359 degrees\n rounding: true, // Circular crop\n transparency: 50, // 0-100\n flipH: true, // Horizontal flip\n flipV: false, // Vertical flip\n altText: \"Description\", // Accessibility\n hyperlink: { url: \"https://example.com\" }\n});\n```\n\n### Image Sizing Modes\n\n```javascript\n// Contain - fit inside, preserve ratio\n{ sizing: { type: 'contain', w: 4, h: 3 } }\n\n// Cover - fill area, preserve ratio (may crop)\n{ sizing: { type: 'cover', w: 4, h: 3 } }\n\n// Crop - cut specific portion\n{ sizing: { type: 'crop', x: 0.5, y: 0.5, w: 2, h: 2 } }\n```\n\n### Calculate Dimensions (preserve aspect ratio)\n\n```javascript\nconst origWidth = 1978, origHeight = 923, maxHeight = 3.0;\nconst calcWidth = maxHeight * (origWidth / origHeight);\nconst centerX = (10 - calcWidth) / 2;\n\nslide.addImage({ path: \"image.png\", x: centerX, y: 1.2, w: calcWidth, h: maxHeight });\n```\n\n### Supported Formats\n\n- **Standard**: PNG, JPG, GIF (animated GIFs work in Microsoft 365)\n- **SVG**: Works in modern PowerPoint/Microsoft 365\n\n---\n\n## Icons\n\nUse react-icons to generate SVG icons, then rasterize to PNG for universal compatibility.\n\n### Setup\n\n```javascript\nconst React = require(\"react\");\nconst ReactDOMServer = require(\"react-dom/server\");\nconst sharp = require(\"sharp\");\nconst { FaCheckCircle, FaChartLine } = require(\"react-icons/fa\");\n\nfunction renderIconSvg(IconComponent, color = \"#000000\", size = 256) {\n return ReactDOMServer.renderToStaticMarkup(\n React.createElement(IconComponent, { color, size: String(size) })\n );\n}\n\nasync function iconToBase64Png(IconComponent, color, size = 256) {\n const svg = renderIconSvg(IconComponent, color, size);\n const pngBuffer = await sharp(Buffer.from(svg)).png().toBuffer();\n return \"image/png;base64,\" + pngBuffer.toString(\"base64\");\n}\n```\n\n### Add Icon to Slide\n\n```javascript\nconst iconData = await iconToBase64Png(FaCheckCircle, \"#4472C4\", 256);\n\nslide.addImage({\n data: iconData,\n x: 1, y: 1, w: 0.5, h: 0.5 // Size in inches\n});\n```\n\n**Note**: Use size 256 or higher for crisp icons. The size parameter controls the rasterization resolution, not the display size on the slide (which is set by `w` and `h` in inches).\n\n### Icon Libraries\n\nInstall: `npm install -g react-icons react react-dom sharp`\n\nPopular icon sets in react-icons:\n- `react-icons/fa` - Font Awesome\n- `react-icons/md` - Material Design\n- `react-icons/hi` - Heroicons\n- `react-icons/bi` - Bootstrap Icons\n\n---\n\n## Slide Backgrounds\n\n```javascript\n// Solid color\nslide.background = { color: \"F1F1F1\" };\n\n// Color with transparency\nslide.background = { color: \"FF3399\", transparency: 50 };\n\n// Image from URL\nslide.background = { path: \"https://example.com/bg.jpg\" };\n\n// Image from base64\nslide.background = { data: \"image/png;base64,iVBORw0KGgo...\" };\n```\n\n---\n\n## Tables\n\n```javascript\nslide.addTable([\n [\"Header 1\", \"Header 2\"],\n [\"Cell 1\", \"Cell 2\"]\n], {\n x: 1, y: 1, w: 8, h: 2,\n border: { pt: 1, color: \"999999\" }, fill: { color: \"F1F1F1\" }\n});\n\n// Advanced with merged cells\nlet tableData = [\n [{ text: \"Header\", options: { fill: { color: \"6699CC\" }, color: \"FFFFFF\", bold: true } }, \"Cell\"],\n [{ text: \"Merged\", options: { colspan: 2 } }]\n];\nslide.addTable(tableData, { x: 1, y: 3.5, w: 8, colW: [4, 4] });\n```\n\n---\n\n## Charts\n\n```javascript\n// Bar chart\nslide.addChart(pres.charts.BAR, [{\n name: \"Sales\", labels: [\"Q1\", \"Q2\", \"Q3\", \"Q4\"], values: [4500, 5500, 6200, 7100]\n}], {\n x: 0.5, y: 0.6, w: 6, h: 3, barDir: 'col',\n showTitle: true, title: 'Quarterly Sales'\n});\n\n// Line chart\nslide.addChart(pres.charts.LINE, [{\n name: \"Temp\", labels: [\"Jan\", \"Feb\", \"Mar\"], values: [32, 35, 42]\n}], { x: 0.5, y: 4, w: 6, h: 3, lineSize: 3, lineSmooth: true });\n\n// Pie chart\nslide.addChart(pres.charts.PIE, [{\n name: \"Share\", labels: [\"A\", \"B\", \"Other\"], values: [35, 45, 20]\n}], { x: 7, y: 1, w: 5, h: 4, showPercent: true });\n```\n\n### Better-Looking Charts\n\nDefault charts look dated. Apply these options for a modern, clean appearance:\n\n```javascript\nslide.addChart(pres.charts.BAR, chartData, {\n x: 0.5, y: 1, w: 9, h: 4, barDir: \"col\",\n\n // Custom colors (match your presentation palette)\n chartColors: [\"0D9488\", \"14B8A6\", \"5EEAD4\"],\n\n // Clean background\n chartArea: { fill: { color: \"FFFFFF\" }, roundedCorners: true },\n\n // Muted axis labels\n catAxisLabelColor: \"64748B\",\n valAxisLabelColor: \"64748B\",\n\n // Subtle grid (value axis only)\n valGridLine: { color: \"E2E8F0\", size: 0.5 },\n catGridLine: { style: \"none\" },\n\n // Data labels on bars\n showValue: true,\n dataLabelPosition: \"outEnd\",\n dataLabelColor: \"1E293B\",\n\n // Hide legend for single series\n showLegend: false,\n});\n```\n\n**Key styling options:**\n- `chartColors: [...]` - hex colors for series/segments\n- `chartArea: { fill, border, roundedCorners }` - chart background\n- `catGridLine/valGridLine: { color, style, size }` - grid lines (`style: \"none\"` to hide)\n- `lineSmooth: true` - curved lines (line charts)\n- `legendPos: \"r\"` - legend position: \"b\", \"t\", \"l\", \"r\", \"tr\"\n\n---\n\n## Slide Masters\n\n```javascript\npres.defineSlideMaster({\n title: 'TITLE_SLIDE', background: { color: '283A5E' },\n objects: [{\n placeholder: { options: { name: 'title', type: 'title', x: 1, y: 2, w: 8, h: 2 } }\n }]\n});\n\nlet titleSlide = pres.addSlide({ masterName: \"TITLE_SLIDE\" });\ntitleSlide.addText(\"My Title\", { placeholder: \"title\" });\n```\n\n---\n\n## Common Pitfalls\n\n⚠️ These issues cause file corruption, visual bugs, or broken output. Avoid them.\n\n1. **NEVER use \"#\" with hex colors** - causes file corruption\n ```javascript\n color: \"FF0000\" // ✅ CORRECT\n color: \"#FF0000\" // ❌ WRONG\n ```\n\n2. **NEVER encode opacity in hex color strings** - 8-char colors (e.g., `\"00000020\"`) corrupt the file. Use the `opacity` property instead.\n ```javascript\n shadow: { type: \"outer\", blur: 6, offset: 2, color: \"00000020\" } // ❌ CORRUPTS FILE\n shadow: { type: \"outer\", blur: 6, offset: 2, color: \"000000\", opacity: 0.12 } // ✅ CORRECT\n ```\n\n3. **Use `bullet: true`** - NEVER unicode symbols like \"•\" (creates double bullets)\n\n4. **Use `breakLine: true`** between array items or text runs together\n\n5. **Avoid `lineSpacing` with bullets** - causes excessive gaps; use `paraSpaceAfter` instead\n\n6. **Each presentation needs fresh instance** - don't reuse `pptxgen()` objects\n\n7. **NEVER reuse option objects across calls** - PptxGenJS mutates objects in-place (e.g. converting shadow values to EMU). Sharing one object between multiple calls corrupts the second shape.\n ```javascript\n const shadow = { type: \"outer\", blur: 6, offset: 2, color: \"000000\", opacity: 0.15 };\n slide.addShape(pres.shapes.RECTANGLE, { shadow, ... }); // ❌ second call gets already-converted values\n slide.addShape(pres.shapes.RECTANGLE, { shadow, ... });\n\n const makeShadow = () => ({ type: \"outer\", blur: 6, offset: 2, color: \"000000\", opacity: 0.15 });\n slide.addShape(pres.shapes.RECTANGLE, { shadow: makeShadow(), ... }); // ✅ fresh object each time\n slide.addShape(pres.shapes.RECTANGLE, { shadow: makeShadow(), ... });\n ```\n\n8. **Don't use `ROUNDED_RECTANGLE` with accent borders** - rectangular overlay bars won't cover rounded corners. Use `RECTANGLE` instead.\n ```javascript\n // ❌ WRONG: Accent bar doesn't cover rounded corners\n slide.addShape(pres.shapes.ROUNDED_RECTANGLE, { x: 1, y: 1, w: 3, h: 1.5, fill: { color: \"FFFFFF\" } });\n slide.addShape(pres.shapes.RECTANGLE, { x: 1, y: 1, w: 0.08, h: 1.5, fill: { color: \"0891B2\" } });\n\n // ✅ CORRECT: Use RECTANGLE for clean alignment\n slide.addShape(pres.shapes.RECTANGLE, { x: 1, y: 1, w: 3, h: 1.5, fill: { color: \"FFFFFF\" } });\n slide.addShape(pres.shapes.RECTANGLE, { x: 1, y: 1, w: 0.08, h: 1.5, fill: { color: \"0891B2\" } });\n ```\n\n---\n\n## Quick Reference\n\n- **Shapes**: RECTANGLE, OVAL, LINE, ROUNDED_RECTANGLE\n- **Charts**: BAR, LINE, PIE, DOUGHNUT, SCATTER, BUBBLE, RADAR\n- **Layouts**: LAYOUT_16x9 (10\"×5.625\"), LAYOUT_16x10, LAYOUT_4x3, LAYOUT_WIDE\n- **Alignment**: \"left\", \"center\", \"right\"\n- **Chart data labels**: \"outEnd\", \"inEnd\", \"center\"\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/__init__.py", "content": "" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/add_slide.py", "content": "\"\"\"Add a new slide to an unpacked PPTX directory.\n\nUsage: python add_slide.py \n\nThe source can be:\n - A slide file (e.g., slide2.xml) - duplicates the slide\n - A layout file (e.g., slideLayout2.xml) - creates from layout\n\nExamples:\n python add_slide.py unpacked/ slide2.xml\n # Duplicates slide2, creates slide5.xml\n\n python add_slide.py unpacked/ slideLayout2.xml\n # Creates slide5.xml from slideLayout2.xml\n\nTo see available layouts: ls unpacked/ppt/slideLayouts/\n\nPrints the element to add to presentation.xml.\n\"\"\"\n\nimport re\nimport shutil\nimport sys\nfrom pathlib import Path\n\n\ndef get_next_slide_number(slides_dir: Path) -> int:\n existing = [\n int(m.group(1))\n for f in slides_dir.glob(\"slide*.xml\")\n if (m := re.match(r\"slide(\\d+)\\.xml\", f.name))\n ]\n return max(existing) + 1 if existing else 1\n\n\ndef create_slide_from_layout(unpacked_dir: Path, layout_file: str) -> None:\n slides_dir = unpacked_dir / \"ppt\" / \"slides\"\n rels_dir = slides_dir / \"_rels\"\n layouts_dir = unpacked_dir / \"ppt\" / \"slideLayouts\"\n\n layout_path = layouts_dir / layout_file\n if not layout_path.exists():\n print(f\"Error: {layout_path} not found\", file=sys.stderr)\n sys.exit(1)\n\n next_num = get_next_slide_number(slides_dir)\n dest = f\"slide{next_num}.xml\"\n dest_slide = slides_dir / dest\n dest_rels = rels_dir / f\"{dest}.rels\"\n\n slide_xml = \"\"\"\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\"\"\"\n dest_slide.write_text(slide_xml, encoding=\"utf-8\")\n\n rels_dir.mkdir(exist_ok=True)\n rels_xml = f\"\"\"\n\n \n\"\"\"\n dest_rels.write_text(rels_xml, encoding=\"utf-8\")\n\n _add_to_content_types(unpacked_dir, dest)\n\n rid = _add_to_presentation_rels(unpacked_dir, dest)\n\n next_slide_id = _get_next_slide_id(unpacked_dir)\n\n print(f\"Created {dest} from {layout_file}\")\n print(\n f'Add to presentation.xml : '\n )\n\n\ndef duplicate_slide(unpacked_dir: Path, source: str) -> None:\n slides_dir = unpacked_dir / \"ppt\" / \"slides\"\n rels_dir = slides_dir / \"_rels\"\n\n source_slide = slides_dir / source\n\n if not source_slide.exists():\n print(f\"Error: {source_slide} not found\", file=sys.stderr)\n sys.exit(1)\n\n next_num = get_next_slide_number(slides_dir)\n dest = f\"slide{next_num}.xml\"\n dest_slide = slides_dir / dest\n\n source_rels = rels_dir / f\"{source}.rels\"\n dest_rels = rels_dir / f\"{dest}.rels\"\n\n shutil.copy2(source_slide, dest_slide)\n\n if source_rels.exists():\n shutil.copy2(source_rels, dest_rels)\n\n rels_content = dest_rels.read_text(encoding=\"utf-8\")\n rels_content = re.sub(\n r'\\s*]*Type=\"[^\"]*notesSlide\"[^>]*/>\\s*',\n \"\\n\",\n rels_content,\n )\n dest_rels.write_text(rels_content, encoding=\"utf-8\")\n\n _add_to_content_types(unpacked_dir, dest)\n\n rid = _add_to_presentation_rels(unpacked_dir, dest)\n\n next_slide_id = _get_next_slide_id(unpacked_dir)\n\n print(f\"Created {dest} from {source}\")\n print(\n f'Add to presentation.xml : '\n )\n\n\ndef _add_to_content_types(unpacked_dir: Path, dest: str) -> None:\n content_types_path = unpacked_dir / \"[Content_Types].xml\"\n content_types = content_types_path.read_text(encoding=\"utf-8\")\n\n content_type = (\n \"application/vnd.openxmlformats-officedocument.presentationml.slide+xml\"\n )\n new_override = (\n f''\n )\n\n if f\"/ppt/slides/{dest}\" not in content_types:\n content_types = content_types.replace(\"\", f\" {new_override}\\n\")\n content_types_path.write_text(content_types, encoding=\"utf-8\")\n\n\ndef _add_to_presentation_rels(unpacked_dir: Path, dest: str) -> str:\n pres_rels_path = unpacked_dir / \"ppt\" / \"_rels\" / \"presentation.xml.rels\"\n pres_rels = pres_rels_path.read_text(encoding=\"utf-8\")\n\n rids = [int(m) for m in re.findall(r'Id=\"rId(\\d+)\"', pres_rels)]\n next_rid = max(rids) + 1 if rids else 1\n rid = f\"rId{next_rid}\"\n\n slide_type = (\n \"http://schemas.openxmlformats.org/officeDocument/2006/relationships/slide\"\n )\n new_rel = f''\n\n if f\"slides/{dest}\" not in pres_rels:\n pres_rels = pres_rels.replace(\n \"\", f\" {new_rel}\\n\"\n )\n pres_rels_path.write_text(pres_rels, encoding=\"utf-8\")\n\n return rid\n\n\ndef _get_next_slide_id(unpacked_dir: Path) -> int:\n pres_path = unpacked_dir / \"ppt\" / \"presentation.xml\"\n pres_content = pres_path.read_text(encoding=\"utf-8\")\n slide_ids = [int(m) for m in re.findall(r']*id=\"(\\d+)\"', pres_content)]\n return max(slide_ids) + 1 if slide_ids else 256\n\n\ndef parse_source(source: str) -> tuple[str, str | None]:\n if source.startswith(\"slideLayout\") and source.endswith(\".xml\"):\n return (\"layout\", source)\n\n return (\"slide\", None)\n\n\nif __name__ == \"__main__\":\n if len(sys.argv) != 3:\n print(\"Usage: python add_slide.py \", file=sys.stderr)\n print(\"\", file=sys.stderr)\n print(\"Source can be:\", file=sys.stderr)\n print(\" slide2.xml - duplicate an existing slide\", file=sys.stderr)\n print(\" slideLayout2.xml - create from a layout template\", file=sys.stderr)\n print(\"\", file=sys.stderr)\n print(\n \"To see available layouts: ls /ppt/slideLayouts/\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n unpacked_dir = Path(sys.argv[1])\n source = sys.argv[2]\n\n if not unpacked_dir.exists():\n print(f\"Error: {unpacked_dir} not found\", file=sys.stderr)\n sys.exit(1)\n\n source_type, layout_file = parse_source(source)\n\n if source_type == \"layout\" and layout_file is not None:\n create_slide_from_layout(unpacked_dir, layout_file)\n else:\n duplicate_slide(unpacked_dir, source)\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/clean.py", "content": "\"\"\"Remove unreferenced files from an unpacked PPTX directory.\n\nUsage: python clean.py \n\nExample:\n python clean.py unpacked/\n\nThis script removes:\n- Orphaned slides (not in sldIdLst) and their relationships\n- [trash] directory (unreferenced files)\n- Orphaned .rels files for deleted resources\n- Unreferenced media, embeddings, charts, diagrams, drawings, ink files\n- Unreferenced theme files\n- Unreferenced notes slides\n- Content-Type overrides for deleted files\n\"\"\"\n\nimport re\nimport sys\nfrom pathlib import Path\n\nimport defusedxml.minidom\n\n\ndef get_slides_in_sldidlst(unpacked_dir: Path) -> set[str]:\n pres_path = unpacked_dir / \"ppt\" / \"presentation.xml\"\n pres_rels_path = unpacked_dir / \"ppt\" / \"_rels\" / \"presentation.xml.rels\"\n\n if not pres_path.exists() or not pres_rels_path.exists():\n return set()\n\n rels_dom = defusedxml.minidom.parse(str(pres_rels_path))\n rid_to_slide = {}\n for rel in rels_dom.getElementsByTagName(\"Relationship\"):\n rid = rel.getAttribute(\"Id\")\n target = rel.getAttribute(\"Target\")\n rel_type = rel.getAttribute(\"Type\")\n if \"slide\" in rel_type and target.startswith(\"slides/\"):\n rid_to_slide[rid] = target.replace(\"slides/\", \"\")\n\n pres_content = pres_path.read_text(encoding=\"utf-8\")\n referenced_rids = set(re.findall(r']*r:id=\"([^\"]+)\"', pres_content))\n\n return {rid_to_slide[rid] for rid in referenced_rids if rid in rid_to_slide}\n\n\ndef remove_orphaned_slides(unpacked_dir: Path) -> list[str]:\n slides_dir = unpacked_dir / \"ppt\" / \"slides\"\n slides_rels_dir = slides_dir / \"_rels\"\n pres_rels_path = unpacked_dir / \"ppt\" / \"_rels\" / \"presentation.xml.rels\"\n\n if not slides_dir.exists():\n return []\n\n referenced_slides = get_slides_in_sldidlst(unpacked_dir)\n removed = []\n\n for slide_file in slides_dir.glob(\"slide*.xml\"):\n if slide_file.name not in referenced_slides:\n rel_path = slide_file.relative_to(unpacked_dir)\n slide_file.unlink()\n removed.append(str(rel_path))\n\n rels_file = slides_rels_dir / f\"{slide_file.name}.rels\"\n if rels_file.exists():\n rels_file.unlink()\n removed.append(str(rels_file.relative_to(unpacked_dir)))\n\n if removed and pres_rels_path.exists():\n rels_dom = defusedxml.minidom.parse(str(pres_rels_path))\n changed = False\n\n for rel in list(rels_dom.getElementsByTagName(\"Relationship\")):\n target = rel.getAttribute(\"Target\")\n if target.startswith(\"slides/\"):\n slide_name = target.replace(\"slides/\", \"\")\n if slide_name not in referenced_slides:\n if rel.parentNode:\n rel.parentNode.removeChild(rel)\n changed = True\n\n if changed:\n with open(pres_rels_path, \"wb\") as f:\n f.write(rels_dom.toxml(encoding=\"utf-8\"))\n\n return removed\n\n\ndef remove_trash_directory(unpacked_dir: Path) -> list[str]:\n trash_dir = unpacked_dir / \"[trash]\"\n removed = []\n\n if trash_dir.exists() and trash_dir.is_dir():\n for file_path in trash_dir.iterdir():\n if file_path.is_file():\n rel_path = file_path.relative_to(unpacked_dir)\n removed.append(str(rel_path))\n file_path.unlink()\n trash_dir.rmdir()\n\n return removed\n\n\ndef get_slide_referenced_files(unpacked_dir: Path) -> set:\n referenced = set()\n slides_rels_dir = unpacked_dir / \"ppt\" / \"slides\" / \"_rels\"\n\n if not slides_rels_dir.exists():\n return referenced\n\n for rels_file in slides_rels_dir.glob(\"*.rels\"):\n dom = defusedxml.minidom.parse(str(rels_file))\n for rel in dom.getElementsByTagName(\"Relationship\"):\n target = rel.getAttribute(\"Target\")\n if not target:\n continue\n target_path = (rels_file.parent.parent / target).resolve()\n try:\n referenced.add(target_path.relative_to(unpacked_dir.resolve()))\n except ValueError:\n pass\n\n return referenced\n\n\ndef remove_orphaned_rels_files(unpacked_dir: Path) -> list[str]:\n resource_dirs = [\"charts\", \"diagrams\", \"drawings\"]\n removed = []\n slide_referenced = get_slide_referenced_files(unpacked_dir)\n\n for dir_name in resource_dirs:\n rels_dir = unpacked_dir / \"ppt\" / dir_name / \"_rels\"\n if not rels_dir.exists():\n continue\n\n for rels_file in rels_dir.glob(\"*.rels\"):\n resource_file = rels_dir.parent / rels_file.name.replace(\".rels\", \"\")\n try:\n resource_rel_path = resource_file.resolve().relative_to(\n unpacked_dir.resolve()\n )\n except ValueError:\n continue\n\n if not resource_file.exists() or resource_rel_path not in slide_referenced:\n rels_file.unlink()\n rel_path = rels_file.relative_to(unpacked_dir)\n removed.append(str(rel_path))\n\n return removed\n\n\ndef get_referenced_files(unpacked_dir: Path) -> set:\n referenced = set()\n\n for rels_file in unpacked_dir.rglob(\"*.rels\"):\n dom = defusedxml.minidom.parse(str(rels_file))\n for rel in dom.getElementsByTagName(\"Relationship\"):\n target = rel.getAttribute(\"Target\")\n if not target:\n continue\n target_path = (rels_file.parent.parent / target).resolve()\n try:\n referenced.add(target_path.relative_to(unpacked_dir.resolve()))\n except ValueError:\n pass\n\n return referenced\n\n\ndef remove_orphaned_files(unpacked_dir: Path, referenced: set) -> list[str]:\n resource_dirs = [\n \"media\",\n \"embeddings\",\n \"charts\",\n \"diagrams\",\n \"tags\",\n \"drawings\",\n \"ink\",\n ]\n removed = []\n\n for dir_name in resource_dirs:\n dir_path = unpacked_dir / \"ppt\" / dir_name\n if not dir_path.exists():\n continue\n\n for file_path in dir_path.glob(\"*\"):\n if not file_path.is_file():\n continue\n rel_path = file_path.relative_to(unpacked_dir)\n if rel_path not in referenced:\n file_path.unlink()\n removed.append(str(rel_path))\n\n theme_dir = unpacked_dir / \"ppt\" / \"theme\"\n if theme_dir.exists():\n for file_path in theme_dir.glob(\"theme*.xml\"):\n rel_path = file_path.relative_to(unpacked_dir)\n if rel_path not in referenced:\n file_path.unlink()\n removed.append(str(rel_path))\n theme_rels = theme_dir / \"_rels\" / f\"{file_path.name}.rels\"\n if theme_rels.exists():\n theme_rels.unlink()\n removed.append(str(theme_rels.relative_to(unpacked_dir)))\n\n notes_dir = unpacked_dir / \"ppt\" / \"notesSlides\"\n if notes_dir.exists():\n for file_path in notes_dir.glob(\"*.xml\"):\n if not file_path.is_file():\n continue\n rel_path = file_path.relative_to(unpacked_dir)\n if rel_path not in referenced:\n file_path.unlink()\n removed.append(str(rel_path))\n\n notes_rels_dir = notes_dir / \"_rels\"\n if notes_rels_dir.exists():\n for file_path in notes_rels_dir.glob(\"*.rels\"):\n notes_file = notes_dir / file_path.name.replace(\".rels\", \"\")\n if not notes_file.exists():\n file_path.unlink()\n removed.append(str(file_path.relative_to(unpacked_dir)))\n\n return removed\n\n\ndef update_content_types(unpacked_dir: Path, removed_files: list[str]) -> None:\n ct_path = unpacked_dir / \"[Content_Types].xml\"\n if not ct_path.exists():\n return\n\n dom = defusedxml.minidom.parse(str(ct_path))\n changed = False\n\n for override in list(dom.getElementsByTagName(\"Override\")):\n part_name = override.getAttribute(\"PartName\").lstrip(\"/\")\n if part_name in removed_files:\n if override.parentNode:\n override.parentNode.removeChild(override)\n changed = True\n\n if changed:\n with open(ct_path, \"wb\") as f:\n f.write(dom.toxml(encoding=\"utf-8\"))\n\n\ndef clean_unused_files(unpacked_dir: Path) -> list[str]:\n all_removed = []\n\n slides_removed = remove_orphaned_slides(unpacked_dir)\n all_removed.extend(slides_removed)\n\n trash_removed = remove_trash_directory(unpacked_dir)\n all_removed.extend(trash_removed)\n\n while True:\n removed_rels = remove_orphaned_rels_files(unpacked_dir)\n referenced = get_referenced_files(unpacked_dir)\n removed_files = remove_orphaned_files(unpacked_dir, referenced)\n\n total_removed = removed_rels + removed_files\n if not total_removed:\n break\n\n all_removed.extend(total_removed)\n\n if all_removed:\n update_content_types(unpacked_dir, all_removed)\n\n return all_removed\n\n\nif __name__ == \"__main__\":\n if len(sys.argv) != 2:\n print(\"Usage: python clean.py \", file=sys.stderr)\n print(\"Example: python clean.py unpacked/\", file=sys.stderr)\n sys.exit(1)\n\n unpacked_dir = Path(sys.argv[1])\n\n if not unpacked_dir.exists():\n print(f\"Error: {unpacked_dir} not found\", file=sys.stderr)\n sys.exit(1)\n\n removed = clean_unused_files(unpacked_dir)\n\n if removed:\n print(f\"Removed {len(removed)} unreferenced files:\")\n for f in removed:\n print(f\" {f}\")\n else:\n print(\"No unreferenced files found\")\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/helpers/__init__.py", "content": "" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/helpers/merge_runs.py", "content": "\"\"\"Merge adjacent runs with identical formatting in DOCX.\n\nMerges adjacent elements that have identical properties.\nWorks on runs in paragraphs and inside tracked changes (, ).\n\nAlso:\n- Removes rsid attributes from runs (revision metadata that doesn't affect rendering)\n- Removes proofErr elements (spell/grammar markers that block merging)\n\"\"\"\n\nfrom pathlib import Path\n\nimport defusedxml.minidom\n\n\ndef merge_runs(input_dir: str) -> tuple[int, str]:\n doc_xml = Path(input_dir) / \"word\" / \"document.xml\"\n\n if not doc_xml.exists():\n return 0, f\"Error: {doc_xml} not found\"\n\n try:\n dom = defusedxml.minidom.parseString(doc_xml.read_text(encoding=\"utf-8\"))\n root = dom.documentElement\n\n _remove_elements(root, \"proofErr\")\n _strip_run_rsid_attrs(root)\n\n containers = {run.parentNode for run in _find_elements(root, \"r\")}\n\n merge_count = 0\n for container in containers:\n merge_count += _merge_runs_in(container)\n\n doc_xml.write_bytes(dom.toxml(encoding=\"UTF-8\"))\n return merge_count, f\"Merged {merge_count} runs\"\n\n except Exception as e:\n return 0, f\"Error: {e}\"\n\n\ndef _find_elements(root, tag: str) -> list:\n results = []\n\n def traverse(node):\n if node.nodeType == node.ELEMENT_NODE:\n name = node.localName or node.tagName\n if name == tag or name.endswith(f\":{tag}\"):\n results.append(node)\n for child in node.childNodes:\n traverse(child)\n\n traverse(root)\n return results\n\n\ndef _get_child(parent, tag: str):\n for child in parent.childNodes:\n if child.nodeType == child.ELEMENT_NODE:\n name = child.localName or child.tagName\n if name == tag or name.endswith(f\":{tag}\"):\n return child\n return None\n\n\ndef _get_children(parent, tag: str) -> list:\n results = []\n for child in parent.childNodes:\n if child.nodeType == child.ELEMENT_NODE:\n name = child.localName or child.tagName\n if name == tag or name.endswith(f\":{tag}\"):\n results.append(child)\n return results\n\n\ndef _is_adjacent(elem1, elem2) -> bool:\n node = elem1.nextSibling\n while node:\n if node == elem2:\n return True\n if node.nodeType == node.ELEMENT_NODE:\n return False\n if node.nodeType == node.TEXT_NODE and node.data.strip():\n return False\n node = node.nextSibling\n return False\n\n\ndef _remove_elements(root, tag: str):\n for elem in _find_elements(root, tag):\n if elem.parentNode:\n elem.parentNode.removeChild(elem)\n\n\ndef _strip_run_rsid_attrs(root):\n for run in _find_elements(root, \"r\"):\n for attr in list(run.attributes.values()):\n if \"rsid\" in attr.name.lower():\n run.removeAttribute(attr.name)\n\n\ndef _merge_runs_in(container) -> int:\n merge_count = 0\n run = _first_child_run(container)\n\n while run:\n while True:\n next_elem = _next_element_sibling(run)\n if next_elem and _is_run(next_elem) and _can_merge(run, next_elem):\n _merge_run_content(run, next_elem)\n container.removeChild(next_elem)\n merge_count += 1\n else:\n break\n\n _consolidate_text(run)\n run = _next_sibling_run(run)\n\n return merge_count\n\n\ndef _first_child_run(container):\n for child in container.childNodes:\n if child.nodeType == child.ELEMENT_NODE and _is_run(child):\n return child\n return None\n\n\ndef _next_element_sibling(node):\n sibling = node.nextSibling\n while sibling:\n if sibling.nodeType == sibling.ELEMENT_NODE:\n return sibling\n sibling = sibling.nextSibling\n return None\n\n\ndef _next_sibling_run(node):\n sibling = node.nextSibling\n while sibling:\n if sibling.nodeType == sibling.ELEMENT_NODE:\n if _is_run(sibling):\n return sibling\n sibling = sibling.nextSibling\n return None\n\n\ndef _is_run(node) -> bool:\n name = node.localName or node.tagName\n return name == \"r\" or name.endswith(\":r\")\n\n\ndef _can_merge(run1, run2) -> bool:\n rpr1 = _get_child(run1, \"rPr\")\n rpr2 = _get_child(run2, \"rPr\")\n\n if (rpr1 is None) != (rpr2 is None):\n return False\n if rpr1 is None:\n return True\n return rpr1.toxml() == rpr2.toxml()\n\n\ndef _merge_run_content(target, source):\n for child in list(source.childNodes):\n if child.nodeType == child.ELEMENT_NODE:\n name = child.localName or child.tagName\n if name != \"rPr\" and not name.endswith(\":rPr\"):\n target.appendChild(child)\n\n\ndef _consolidate_text(run):\n t_elements = _get_children(run, \"t\")\n\n for i in range(len(t_elements) - 1, 0, -1):\n curr, prev = t_elements[i], t_elements[i - 1]\n\n if _is_adjacent(prev, curr):\n prev_text = prev.firstChild.data if prev.firstChild else \"\"\n curr_text = curr.firstChild.data if curr.firstChild else \"\"\n merged = prev_text + curr_text\n\n if prev.firstChild:\n prev.firstChild.data = merged\n else:\n prev.appendChild(run.ownerDocument.createTextNode(merged))\n\n if merged.startswith(\" \") or merged.endswith(\" \"):\n prev.setAttribute(\"xml:space\", \"preserve\")\n elif prev.hasAttribute(\"xml:space\"):\n prev.removeAttribute(\"xml:space\")\n\n run.removeChild(curr)\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/helpers/simplify_redlines.py", "content": "\"\"\"Simplify tracked changes by merging adjacent w:ins or w:del elements.\n\nMerges adjacent elements from the same author into a single element.\nSame for elements. This makes heavily-redlined documents easier to\nwork with by reducing the number of tracked change wrappers.\n\nRules:\n- Only merges w:ins with w:ins, w:del with w:del (same element type)\n- Only merges if same author (ignores timestamp differences)\n- Only merges if truly adjacent (only whitespace between them)\n\"\"\"\n\nimport xml.etree.ElementTree as ET\nimport zipfile\nfrom pathlib import Path\n\nimport defusedxml.minidom\n\nWORD_NS = \"http://schemas.openxmlformats.org/wordprocessingml/2006/main\"\n\n\ndef simplify_redlines(input_dir: str) -> tuple[int, str]:\n doc_xml = Path(input_dir) / \"word\" / \"document.xml\"\n\n if not doc_xml.exists():\n return 0, f\"Error: {doc_xml} not found\"\n\n try:\n dom = defusedxml.minidom.parseString(doc_xml.read_text(encoding=\"utf-8\"))\n root = dom.documentElement\n\n merge_count = 0\n\n containers = _find_elements(root, \"p\") + _find_elements(root, \"tc\")\n\n for container in containers:\n merge_count += _merge_tracked_changes_in(container, \"ins\")\n merge_count += _merge_tracked_changes_in(container, \"del\")\n\n doc_xml.write_bytes(dom.toxml(encoding=\"UTF-8\"))\n return merge_count, f\"Simplified {merge_count} tracked changes\"\n\n except Exception as e:\n return 0, f\"Error: {e}\"\n\n\ndef _merge_tracked_changes_in(container, tag: str) -> int:\n merge_count = 0\n\n tracked = [\n child\n for child in container.childNodes\n if child.nodeType == child.ELEMENT_NODE and _is_element(child, tag)\n ]\n\n if len(tracked) < 2:\n return 0\n\n i = 0\n while i < len(tracked) - 1:\n curr = tracked[i]\n next_elem = tracked[i + 1]\n\n if _can_merge_tracked(curr, next_elem):\n _merge_tracked_content(curr, next_elem)\n container.removeChild(next_elem)\n tracked.pop(i + 1)\n merge_count += 1\n else:\n i += 1\n\n return merge_count\n\n\ndef _is_element(node, tag: str) -> bool:\n name = node.localName or node.tagName\n return name == tag or name.endswith(f\":{tag}\")\n\n\ndef _get_author(elem) -> str:\n author = elem.getAttribute(\"w:author\")\n if not author:\n for attr in elem.attributes.values():\n if attr.localName == \"author\" or attr.name.endswith(\":author\"):\n return attr.value\n return author\n\n\ndef _can_merge_tracked(elem1, elem2) -> bool:\n if _get_author(elem1) != _get_author(elem2):\n return False\n\n node = elem1.nextSibling\n while node and node != elem2:\n if node.nodeType == node.ELEMENT_NODE:\n return False\n if node.nodeType == node.TEXT_NODE and node.data.strip():\n return False\n node = node.nextSibling\n\n return True\n\n\ndef _merge_tracked_content(target, source):\n while source.firstChild:\n child = source.firstChild\n source.removeChild(child)\n target.appendChild(child)\n\n\ndef _find_elements(root, tag: str) -> list:\n results = []\n\n def traverse(node):\n if node.nodeType == node.ELEMENT_NODE:\n name = node.localName or node.tagName\n if name == tag or name.endswith(f\":{tag}\"):\n results.append(node)\n for child in node.childNodes:\n traverse(child)\n\n traverse(root)\n return results\n\n\ndef get_tracked_change_authors(doc_xml_path: Path) -> dict[str, int]:\n if not doc_xml_path.exists():\n return {}\n\n try:\n tree = ET.parse(doc_xml_path)\n root = tree.getroot()\n except ET.ParseError:\n return {}\n\n namespaces = {\"w\": WORD_NS}\n author_attr = f\"{{{WORD_NS}}}author\"\n\n authors: dict[str, int] = {}\n for tag in [\"ins\", \"del\"]:\n for elem in root.findall(f\".//w:{tag}\", namespaces):\n author = elem.get(author_attr)\n if author:\n authors[author] = authors.get(author, 0) + 1\n\n return authors\n\n\ndef _get_authors_from_docx(docx_path: Path) -> dict[str, int]:\n try:\n with zipfile.ZipFile(docx_path, \"r\") as zf:\n if \"word/document.xml\" not in zf.namelist():\n return {}\n with zf.open(\"word/document.xml\") as f:\n tree = ET.parse(f)\n root = tree.getroot()\n\n namespaces = {\"w\": WORD_NS}\n author_attr = f\"{{{WORD_NS}}}author\"\n\n authors: dict[str, int] = {}\n for tag in [\"ins\", \"del\"]:\n for elem in root.findall(f\".//w:{tag}\", namespaces):\n author = elem.get(author_attr)\n if author:\n authors[author] = authors.get(author, 0) + 1\n return authors\n except (zipfile.BadZipFile, ET.ParseError):\n return {}\n\n\ndef infer_author(\n modified_dir: Path, original_docx: Path, default: str = \"Claude\"\n) -> str:\n modified_xml = modified_dir / \"word\" / \"document.xml\"\n modified_authors = get_tracked_change_authors(modified_xml)\n\n if not modified_authors:\n return default\n\n original_authors = _get_authors_from_docx(original_docx)\n\n new_changes: dict[str, int] = {}\n for author, count in modified_authors.items():\n original_count = original_authors.get(author, 0)\n diff = count - original_count\n if diff > 0:\n new_changes[author] = diff\n\n if not new_changes:\n return default\n\n if len(new_changes) == 1:\n return next(iter(new_changes))\n\n raise ValueError(\n f\"Multiple authors added new changes: {new_changes}. Cannot infer which author to validate.\"\n )\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/pack.py", "content": "\"\"\"Pack a directory into a DOCX, PPTX, or XLSX file.\n\nValidates with auto-repair, condenses XML formatting, and creates the Office file.\n\nUsage:\n python pack.py [--original ] [--validate true|false]\n\nExamples:\n python pack.py unpacked/ output.docx --original input.docx\n python pack.py unpacked/ output.pptx --validate false\n\"\"\"\n\nimport argparse\nimport shutil\nimport sys\nimport tempfile\nimport zipfile\nfrom pathlib import Path\n\nimport defusedxml.minidom\nfrom validators import DOCXSchemaValidator\nfrom validators import PPTXSchemaValidator\nfrom validators import RedliningValidator\n\n\ndef pack(\n input_directory: str,\n output_file: str,\n original_file: str | None = None,\n validate: bool = True,\n infer_author_func=None,\n) -> tuple[None, str]:\n input_dir = Path(input_directory)\n output_path = Path(output_file)\n suffix = output_path.suffix.lower()\n\n if not input_dir.is_dir():\n return None, f\"Error: {input_dir} is not a directory\"\n\n if suffix not in {\".docx\", \".pptx\", \".xlsx\"}:\n return None, f\"Error: {output_file} must be a .docx, .pptx, or .xlsx file\"\n\n if validate and original_file:\n original_path = Path(original_file)\n if original_path.exists():\n success, output = _run_validation(\n input_dir, original_path, suffix, infer_author_func\n )\n if output:\n print(output)\n if not success:\n return None, f\"Error: Validation failed for {input_dir}\"\n\n with tempfile.TemporaryDirectory() as temp_dir:\n temp_content_dir = Path(temp_dir) / \"content\"\n shutil.copytree(input_dir, temp_content_dir)\n\n for pattern in [\"*.xml\", \"*.rels\"]:\n for xml_file in temp_content_dir.rglob(pattern):\n _condense_xml(xml_file)\n\n output_path.parent.mkdir(parents=True, exist_ok=True)\n with zipfile.ZipFile(output_path, \"w\", zipfile.ZIP_DEFLATED) as zf:\n for f in temp_content_dir.rglob(\"*\"):\n if f.is_file():\n zf.write(f, f.relative_to(temp_content_dir))\n\n return None, f\"Successfully packed {input_dir} to {output_file}\"\n\n\ndef _run_validation(\n unpacked_dir: Path,\n original_file: Path,\n suffix: str,\n infer_author_func=None,\n) -> tuple[bool, str | None]:\n output_lines = []\n validators = []\n\n if suffix == \".docx\":\n author = \"Claude\"\n if infer_author_func:\n try:\n author = infer_author_func(unpacked_dir, original_file)\n except ValueError as e:\n print(f\"Warning: {e} Using default author 'Claude'.\", file=sys.stderr)\n\n validators = [\n DOCXSchemaValidator(unpacked_dir, original_file),\n RedliningValidator(unpacked_dir, original_file, author=author),\n ]\n elif suffix == \".pptx\":\n validators = [PPTXSchemaValidator(unpacked_dir, original_file)]\n\n if not validators:\n return True, None\n\n total_repairs = sum(v.repair() for v in validators)\n if total_repairs:\n output_lines.append(f\"Auto-repaired {total_repairs} issue(s)\")\n\n success = all(v.validate() for v in validators)\n\n if success:\n output_lines.append(\"All validations PASSED!\")\n\n return success, \"\\n\".join(output_lines) if output_lines else None\n\n\ndef _condense_xml(xml_file: Path) -> None:\n try:\n with open(xml_file, encoding=\"utf-8\") as f:\n dom = defusedxml.minidom.parse(f)\n\n for element in dom.getElementsByTagName(\"*\"):\n if element.tagName.endswith(\":t\"):\n continue\n\n for child in list(element.childNodes):\n if (\n child.nodeType == child.TEXT_NODE\n and child.nodeValue\n and child.nodeValue.strip() == \"\"\n ) or child.nodeType == child.COMMENT_NODE:\n element.removeChild(child)\n\n xml_file.write_bytes(dom.toxml(encoding=\"UTF-8\"))\n except Exception as e:\n print(f\"ERROR: Failed to parse {xml_file.name}: {e}\", file=sys.stderr)\n raise\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(\n description=\"Pack a directory into a DOCX, PPTX, or XLSX file\"\n )\n parser.add_argument(\"input_directory\", help=\"Unpacked Office document directory\")\n parser.add_argument(\"output_file\", help=\"Output Office file (.docx/.pptx/.xlsx)\")\n parser.add_argument(\n \"--original\",\n help=\"Original file for validation comparison\",\n )\n parser.add_argument(\n \"--validate\",\n type=lambda x: x.lower() == \"true\",\n default=True,\n metavar=\"true|false\",\n help=\"Run validation with auto-repair (default: true)\",\n )\n args = parser.parse_args()\n\n _, message = pack(\n args.input_directory,\n args.output_file,\n original_file=args.original,\n validate=args.validate,\n )\n print(message)\n\n if \"Error\" in message:\n sys.exit(1)\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-chart.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-chartDrawing.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-diagram.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-lockedCanvas.xsd", "content": "\n\n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-main.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-picture.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-spreadsheetDrawing.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/dml-wordprocessingDrawing.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/pml.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-additionalCharacteristics.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-bibliography.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-commonSimpleTypes.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-customXmlDataProperties.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-customXmlSchemaProperties.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-documentPropertiesCustom.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-documentPropertiesExtended.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-documentPropertiesVariantTypes.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-math.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/shared-relationshipReference.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/sml.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/vml-main.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/vml-officeDrawing.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/vml-presentationDrawing.xsd", "content": "\n\n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/vml-spreadsheetDrawing.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/vml-wordprocessingDrawing.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/wml.xsd", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ISO-IEC29500-4_2016/xml.xsd", "content": "\n\n\n \n \n See http://www.w3.org/XML/1998/namespace.html and\n http://www.w3.org/TR/REC-xml for information about this namespace.\n\n This schema document describes the XML namespace, in a form\n suitable for import by other schema documents. \n\n Note that local names in this namespace are intended to be defined\n only by the World Wide Web Consortium or its subgroups. The\n following names are currently defined in this namespace and should\n not be used with conflicting semantics by any Working Group,\n specification, or document instance:\n\n base (as an attribute name): denotes an attribute whose value\n provides a URI to be used as the base for interpreting any\n relative URIs in the scope of the element on which it\n appears; its value is inherited. This name is reserved\n by virtue of its definition in the XML Base specification.\n\n lang (as an attribute name): denotes an attribute whose value\n is a language code for the natural language of the content of\n any element; its value is inherited. This name is reserved\n by virtue of its definition in the XML specification.\n \n space (as an attribute name): denotes an attribute whose\n value is a keyword indicating what whitespace processing\n discipline is intended for the content of the element; its\n value is inherited. This name is reserved by virtue of its\n definition in the XML specification.\n\n Father (in any context at all): denotes Jon Bosak, the chair of \n the original XML Working Group. This name is reserved by \n the following decision of the W3C XML Plenary and \n XML Coordination groups:\n\n In appreciation for his vision, leadership and dedication\n the W3C XML Plenary on this 10th day of February, 2000\n reserves for Jon Bosak in perpetuity the XML name\n xml:Father\n \n \n\n \n This schema defines attributes and an attribute group\n suitable for use by\n schemas wishing to allow xml:base, xml:lang or xml:space attributes\n on elements they define.\n\n To enable this, such a schema must import this schema\n for the XML namespace, e.g. as follows:\n <schema . . .>\n . . .\n <import namespace=\"http://www.w3.org/XML/1998/namespace\"\n schemaLocation=\"http://www.w3.org/2001/03/xml.xsd\"/>\n\n Subsequently, qualified reference to any of the attributes\n or the group defined below will have the desired effect, e.g.\n\n <type . . .>\n . . .\n <attributeGroup ref=\"xml:specialAttrs\"/>\n \n will define a type which will schema-validate an instance\n element with any of those attributes\n \n\n \n In keeping with the XML Schema WG's standard versioning\n policy, this schema document will persist at\n http://www.w3.org/2001/03/xml.xsd.\n At the date of issue it can also be found at\n http://www.w3.org/2001/xml.xsd.\n The schema document at that URI may however change in the future,\n in order to remain compatible with the latest version of XML Schema\n itself. In other words, if the XML Schema namespace changes, the version\n of this document at\n http://www.w3.org/2001/xml.xsd will change\n accordingly; the version at\n http://www.w3.org/2001/03/xml.xsd will not change.\n \n \n\n \n \n In due course, we should install the relevant ISO 2- and 3-letter\n codes as the enumerated possible values . . .\n \n \n\n \n \n \n \n \n \n \n \n\n \n \n See http://www.w3.org/TR/xmlbase/ for\n information about this attribute.\n \n \n\n \n \n \n \n \n\n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ecma/fouth-edition/opc-contentTypes.xsd", "content": "\n\n\n \n \n \n\n \n \n \n \n \n \n\n \n \n \n \n\n \n \n \n \n\n \n \n \n \n \n\n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ecma/fouth-edition/opc-coreProperties.xsd", "content": "\n\n\n \n \n \n\n \n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n \n \n \n \n \n \n\n \n \n \n \n \n \n \n\n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ecma/fouth-edition/opc-digSig.xsd", "content": "\n\n\n \n \n \n\n \n \n \n \n \n \n\n \n \n \n \n \n \n \n\n \n \n \n \n \n \n \n\n \n \n \n \n \n\n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/ecma/fouth-edition/opc-relationships.xsd", "content": "\n\n\n \n \n\n \n \n \n \n \n\n \n \n \n \n \n \n \n \n \n \n\n \n \n \n \n \n \n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/mce/mc.xsd", "content": "\n\n\n \n\n \n \n\n \n\t\n\t\n\n\n\t\n\t\t\n\t\t\t\n\t\t\t\t\n\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\t\n\t\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\n\t\t\t\t\n\t\t\t\t\n\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\t\n\t\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\t\n\t\t\t\t\t\n\t\t\t\t\n\t\t\t\n\t\t\t\n\t\t\t\n\t\t\t\n\t\t\t\n\t\t\n\t\n\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-2010.xsd", "content": " \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-2012.xsd", "content": " \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-2018.xsd", "content": " \n \n \n \n \n \n \n \n \n \n \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-cex-2018.xsd", "content": " \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-cid-2016.xsd", "content": " \n \n \n \n \n \n \n \n \n \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-sdtdatahash-2020.xsd", "content": " \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/schemas/microsoft/wml-symex-2015.xsd", "content": " \n \n \n \n \n \n \n \n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/soffice.py", "content": "\"\"\"\nHelper for running LibreOffice (soffice) in environments where AF_UNIX\nsockets may be blocked (e.g., sandboxed VMs). Detects the restriction\nat runtime and applies an LD_PRELOAD shim if needed.\n\nUsage:\n from office.soffice import run_soffice, get_soffice_env\n\n # Option 1 – run soffice directly\n result = run_soffice([\"--headless\", \"--convert-to\", \"pdf\", \"input.docx\"])\n\n # Option 2 – get env dict for your own subprocess calls\n env = get_soffice_env()\n subprocess.run([\"soffice\", ...], env=env)\n\"\"\"\n\nimport os\nimport socket\nimport subprocess\nimport tempfile\nfrom pathlib import Path\n\n\ndef get_soffice_env() -> dict:\n env = os.environ.copy()\n env[\"SAL_USE_VCLPLUGIN\"] = \"svp\"\n\n if _needs_shim():\n shim = _ensure_shim()\n env[\"LD_PRELOAD\"] = str(shim)\n\n return env\n\n\ndef run_soffice(args: list[str], **kwargs) -> subprocess.CompletedProcess:\n env = get_soffice_env()\n return subprocess.run([\"soffice\"] + args, env=env, **kwargs)\n\n\n_SHIM_SO = Path(tempfile.gettempdir()) / \"lo_socket_shim.so\"\n\n\ndef _needs_shim() -> bool:\n try:\n s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)\n s.close()\n return False\n except OSError:\n return True\n\n\ndef _ensure_shim() -> Path:\n if _SHIM_SO.exists():\n return _SHIM_SO\n\n src = Path(tempfile.gettempdir()) / \"lo_socket_shim.c\"\n src.write_text(_SHIM_SOURCE)\n subprocess.run(\n [\"gcc\", \"-shared\", \"-fPIC\", \"-o\", str(_SHIM_SO), str(src), \"-ldl\"],\n check=True,\n capture_output=True,\n )\n src.unlink()\n return _SHIM_SO\n\n\n_SHIM_SOURCE = r\"\"\"\n#define _GNU_SOURCE\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nstatic int (*real_socket)(int, int, int);\nstatic int (*real_socketpair)(int, int, int, int[2]);\nstatic int (*real_listen)(int, int);\nstatic int (*real_accept)(int, struct sockaddr *, socklen_t *);\nstatic int (*real_close)(int);\nstatic int (*real_read)(int, void *, size_t);\n\n/* Per-FD bookkeeping (FDs >= 1024 are passed through unshimmed). */\nstatic int is_shimmed[1024];\nstatic int peer_of[1024];\nstatic int wake_r[1024]; /* accept() blocks reading this */\nstatic int wake_w[1024]; /* close() writes to this */\nstatic int listener_fd = -1; /* FD that received listen() */\n\n__attribute__((constructor))\nstatic void init(void) {\n real_socket = dlsym(RTLD_NEXT, \"socket\");\n real_socketpair = dlsym(RTLD_NEXT, \"socketpair\");\n real_listen = dlsym(RTLD_NEXT, \"listen\");\n real_accept = dlsym(RTLD_NEXT, \"accept\");\n real_close = dlsym(RTLD_NEXT, \"close\");\n real_read = dlsym(RTLD_NEXT, \"read\");\n for (int i = 0; i < 1024; i++) {\n peer_of[i] = -1;\n wake_r[i] = -1;\n wake_w[i] = -1;\n }\n}\n\n/* ---- socket ---------------------------------------------------------- */\nint socket(int domain, int type, int protocol) {\n if (domain == AF_UNIX) {\n int fd = real_socket(domain, type, protocol);\n if (fd >= 0) return fd;\n /* socket(AF_UNIX) blocked – fall back to socketpair(). */\n int sv[2];\n if (real_socketpair(domain, type, protocol, sv) == 0) {\n if (sv[0] >= 0 && sv[0] < 1024) {\n is_shimmed[sv[0]] = 1;\n peer_of[sv[0]] = sv[1];\n int wp[2];\n if (pipe(wp) == 0) {\n wake_r[sv[0]] = wp[0];\n wake_w[sv[0]] = wp[1];\n }\n }\n return sv[0];\n }\n errno = EPERM;\n return -1;\n }\n return real_socket(domain, type, protocol);\n}\n\n/* ---- listen ---------------------------------------------------------- */\nint listen(int sockfd, int backlog) {\n if (sockfd >= 0 && sockfd < 1024 && is_shimmed[sockfd]) {\n listener_fd = sockfd;\n return 0;\n }\n return real_listen(sockfd, backlog);\n}\n\n/* ---- accept ---------------------------------------------------------- */\nint accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {\n if (sockfd >= 0 && sockfd < 1024 && is_shimmed[sockfd]) {\n /* Block until close() writes to the wake pipe. */\n if (wake_r[sockfd] >= 0) {\n char buf;\n real_read(wake_r[sockfd], &buf, 1);\n }\n errno = ECONNABORTED;\n return -1;\n }\n return real_accept(sockfd, addr, addrlen);\n}\n\n/* ---- close ----------------------------------------------------------- */\nint close(int fd) {\n if (fd >= 0 && fd < 1024 && is_shimmed[fd]) {\n int was_listener = (fd == listener_fd);\n is_shimmed[fd] = 0;\n\n if (wake_w[fd] >= 0) { /* unblock accept() */\n char c = 0;\n write(wake_w[fd], &c, 1);\n real_close(wake_w[fd]);\n wake_w[fd] = -1;\n }\n if (wake_r[fd] >= 0) { real_close(wake_r[fd]); wake_r[fd] = -1; }\n if (peer_of[fd] >= 0) { real_close(peer_of[fd]); peer_of[fd] = -1; }\n\n if (was_listener)\n _exit(0); /* conversion done – exit */\n }\n return real_close(fd);\n}\n\"\"\"\n\n\nif __name__ == \"__main__\":\n import sys\n\n result = run_soffice(sys.argv[1:])\n sys.exit(result.returncode)\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/unpack.py", "content": "\"\"\"Unpack Office files (DOCX, PPTX, XLSX) for editing.\n\nExtracts the ZIP archive, pretty-prints XML files, and optionally:\n- Merges adjacent runs with identical formatting (DOCX only)\n- Simplifies adjacent tracked changes from same author (DOCX only)\n\nUsage:\n python unpack.py [options]\n\nExamples:\n python unpack.py document.docx unpacked/\n python unpack.py presentation.pptx unpacked/\n python unpack.py document.docx unpacked/ --merge-runs false\n\"\"\"\n\nimport argparse\nimport sys\nimport zipfile\nfrom pathlib import Path\n\nimport defusedxml.minidom\nfrom helpers.merge_runs import merge_runs as do_merge_runs\nfrom helpers.simplify_redlines import simplify_redlines as do_simplify_redlines\n\nSMART_QUOTE_REPLACEMENTS = {\n \"\\u201c\": \"“\",\n \"\\u201d\": \"”\",\n \"\\u2018\": \"‘\",\n \"\\u2019\": \"’\",\n}\n\n\ndef unpack(\n input_file: str,\n output_directory: str,\n merge_runs: bool = True,\n simplify_redlines: bool = True,\n) -> tuple[None, str]:\n input_path = Path(input_file)\n output_path = Path(output_directory)\n suffix = input_path.suffix.lower()\n\n if not input_path.exists():\n return None, f\"Error: {input_file} does not exist\"\n\n if suffix not in {\".docx\", \".pptx\", \".xlsx\"}:\n return None, f\"Error: {input_file} must be a .docx, .pptx, or .xlsx file\"\n\n try:\n output_path.mkdir(parents=True, exist_ok=True)\n\n with zipfile.ZipFile(input_path, \"r\") as zf:\n zf.extractall(output_path)\n\n xml_files = list(output_path.rglob(\"*.xml\")) + list(output_path.rglob(\"*.rels\"))\n for xml_file in xml_files:\n _pretty_print_xml(xml_file)\n\n message = f\"Unpacked {input_file} ({len(xml_files)} XML files)\"\n\n if suffix == \".docx\":\n if simplify_redlines:\n simplify_count, _ = do_simplify_redlines(str(output_path))\n message += f\", simplified {simplify_count} tracked changes\"\n\n if merge_runs:\n merge_count, _ = do_merge_runs(str(output_path))\n message += f\", merged {merge_count} runs\"\n\n for xml_file in xml_files:\n _escape_smart_quotes(xml_file)\n\n return None, message\n\n except zipfile.BadZipFile:\n return None, f\"Error: {input_file} is not a valid Office file\"\n except Exception as e:\n return None, f\"Error unpacking: {e}\"\n\n\ndef _pretty_print_xml(xml_file: Path) -> None:\n try:\n content = xml_file.read_text(encoding=\"utf-8\")\n dom = defusedxml.minidom.parseString(content)\n xml_file.write_bytes(dom.toprettyxml(indent=\" \", encoding=\"utf-8\"))\n except Exception:\n pass\n\n\ndef _escape_smart_quotes(xml_file: Path) -> None:\n try:\n content = xml_file.read_text(encoding=\"utf-8\")\n for char, entity in SMART_QUOTE_REPLACEMENTS.items():\n content = content.replace(char, entity)\n xml_file.write_text(content, encoding=\"utf-8\")\n except Exception:\n pass\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(\n description=\"Unpack an Office file (DOCX, PPTX, XLSX) for editing\"\n )\n parser.add_argument(\"input_file\", help=\"Office file to unpack\")\n parser.add_argument(\"output_directory\", help=\"Output directory\")\n parser.add_argument(\n \"--merge-runs\",\n type=lambda x: x.lower() == \"true\",\n default=True,\n metavar=\"true|false\",\n help=\"Merge adjacent runs with identical formatting (DOCX only, default: true)\",\n )\n parser.add_argument(\n \"--simplify-redlines\",\n type=lambda x: x.lower() == \"true\",\n default=True,\n metavar=\"true|false\",\n help=\"Merge adjacent tracked changes from same author (DOCX only, default: true)\",\n )\n args = parser.parse_args()\n\n _, message = unpack(\n args.input_file,\n args.output_directory,\n merge_runs=args.merge_runs,\n simplify_redlines=args.simplify_redlines,\n )\n print(message)\n\n if \"Error\" in message:\n sys.exit(1)\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/validate.py", "content": "\"\"\"\nCommand line tool to validate Office document XML files against XSD schemas and tracked changes.\n\nUsage:\n python validate.py [--original ] [--auto-repair] [--author NAME]\n\nThe first argument can be either:\n- An unpacked directory containing the Office document XML files\n- A packed Office file (.docx/.pptx/.xlsx) which will be unpacked to a temp directory\n\nAuto-repair fixes:\n- paraId/durableId values that exceed OOXML limits\n- Missing xml:space=\"preserve\" on w:t elements with whitespace\n\"\"\"\n\nimport argparse\nimport sys\nimport tempfile\nimport zipfile\nfrom pathlib import Path\n\nfrom validators import DOCXSchemaValidator\nfrom validators import PPTXSchemaValidator\nfrom validators import RedliningValidator\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Validate Office document XML files\")\n parser.add_argument(\n \"path\",\n help=\"Path to unpacked directory or packed Office file (.docx/.pptx/.xlsx)\",\n )\n parser.add_argument(\n \"--original\",\n required=False,\n default=None,\n help=(\n \"Path to original file (.docx/.pptx/.xlsx). If omitted, all XSD errors \"\n \"are reported and redlining validation is skipped.\"\n ),\n )\n parser.add_argument(\n \"-v\",\n \"--verbose\",\n action=\"store_true\",\n help=\"Enable verbose output\",\n )\n parser.add_argument(\n \"--auto-repair\",\n action=\"store_true\",\n help=\"Automatically repair common issues (hex IDs, whitespace preservation)\",\n )\n parser.add_argument(\n \"--author\",\n default=\"Claude\",\n help=\"Author name for redlining validation (default: Claude)\",\n )\n args = parser.parse_args()\n\n path = Path(args.path)\n assert path.exists(), f\"Error: {path} does not exist\"\n\n original_file = None\n if args.original:\n original_file = Path(args.original)\n assert original_file.is_file(), f\"Error: {original_file} is not a file\"\n assert original_file.suffix.lower() in [\n \".docx\",\n \".pptx\",\n \".xlsx\",\n ], f\"Error: {original_file} must be a .docx, .pptx, or .xlsx file\"\n\n file_extension = (original_file or path).suffix.lower()\n assert file_extension in [\n \".docx\",\n \".pptx\",\n \".xlsx\",\n ], f\"Error: Cannot determine file type from {path}. Use --original or provide a .docx/.pptx/.xlsx file.\"\n\n if path.is_file() and path.suffix.lower() in [\".docx\", \".pptx\", \".xlsx\"]:\n temp_dir = tempfile.mkdtemp()\n with zipfile.ZipFile(path, \"r\") as zf:\n zf.extractall(temp_dir)\n unpacked_dir = Path(temp_dir)\n else:\n assert path.is_dir(), f\"Error: {path} is not a directory or Office file\"\n unpacked_dir = path\n\n match file_extension:\n case \".docx\":\n validators = [\n DOCXSchemaValidator(unpacked_dir, original_file, verbose=args.verbose),\n ]\n if original_file:\n validators.append(\n RedliningValidator(\n unpacked_dir,\n original_file,\n verbose=args.verbose,\n author=args.author,\n )\n )\n case \".pptx\":\n validators = [\n PPTXSchemaValidator(unpacked_dir, original_file, verbose=args.verbose),\n ]\n case _:\n print(f\"Error: Validation not supported for file type {file_extension}\")\n sys.exit(1)\n\n if args.auto_repair:\n total_repairs = sum(v.repair() for v in validators)\n if total_repairs:\n print(f\"Auto-repaired {total_repairs} issue(s)\")\n\n success = all(v.validate() for v in validators)\n\n if success:\n print(\"All validations PASSED!\")\n\n sys.exit(0 if success else 1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/validators/__init__.py", "content": "\"\"\"\nValidation modules for Word document processing.\n\"\"\"\n\nfrom .base import BaseSchemaValidator\nfrom .docx import DOCXSchemaValidator\nfrom .pptx import PPTXSchemaValidator\nfrom .redlining import RedliningValidator\n\n__all__ = [\n \"BaseSchemaValidator\",\n \"DOCXSchemaValidator\",\n \"PPTXSchemaValidator\",\n \"RedliningValidator\",\n]\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/validators/base.py", "content": "\"\"\"\nBase validator with common validation logic for document files.\n\"\"\"\n\nimport re\nfrom pathlib import Path\n\nimport defusedxml.minidom\nimport lxml.etree\n\n\nclass BaseSchemaValidator:\n IGNORED_VALIDATION_ERRORS = [\n \"hyphenationZone\",\n \"purl.org/dc/terms\",\n ]\n\n UNIQUE_ID_REQUIREMENTS = {\n \"comment\": (\"id\", \"file\"),\n \"commentrangestart\": (\"id\", \"file\"),\n \"commentrangeend\": (\"id\", \"file\"),\n \"bookmarkstart\": (\"id\", \"file\"),\n \"bookmarkend\": (\"id\", \"file\"),\n \"sldid\": (\"id\", \"file\"),\n \"sldmasterid\": (\"id\", \"global\"),\n \"sldlayoutid\": (\"id\", \"global\"),\n \"cm\": (\"authorid\", \"file\"),\n \"sheet\": (\"sheetid\", \"file\"),\n \"definedname\": (\"id\", \"file\"),\n \"cxnsp\": (\"id\", \"file\"),\n \"sp\": (\"id\", \"file\"),\n \"pic\": (\"id\", \"file\"),\n \"grpsp\": (\"id\", \"file\"),\n }\n\n EXCLUDED_ID_CONTAINERS = {\n \"sectionlst\",\n }\n\n ELEMENT_RELATIONSHIP_TYPES = {}\n\n SCHEMA_MAPPINGS = {\n \"word\": \"ISO-IEC29500-4_2016/wml.xsd\",\n \"ppt\": \"ISO-IEC29500-4_2016/pml.xsd\",\n \"xl\": \"ISO-IEC29500-4_2016/sml.xsd\",\n \"[Content_Types].xml\": \"ecma/fouth-edition/opc-contentTypes.xsd\",\n \"app.xml\": \"ISO-IEC29500-4_2016/shared-documentPropertiesExtended.xsd\",\n \"core.xml\": \"ecma/fouth-edition/opc-coreProperties.xsd\",\n \"custom.xml\": \"ISO-IEC29500-4_2016/shared-documentPropertiesCustom.xsd\",\n \".rels\": \"ecma/fouth-edition/opc-relationships.xsd\",\n \"people.xml\": \"microsoft/wml-2012.xsd\",\n \"commentsIds.xml\": \"microsoft/wml-cid-2016.xsd\",\n \"commentsExtensible.xml\": \"microsoft/wml-cex-2018.xsd\",\n \"commentsExtended.xml\": \"microsoft/wml-2012.xsd\",\n \"chart\": \"ISO-IEC29500-4_2016/dml-chart.xsd\",\n \"theme\": \"ISO-IEC29500-4_2016/dml-main.xsd\",\n \"drawing\": \"ISO-IEC29500-4_2016/dml-main.xsd\",\n }\n\n MC_NAMESPACE = \"http://schemas.openxmlformats.org/markup-compatibility/2006\"\n XML_NAMESPACE = \"http://www.w3.org/XML/1998/namespace\"\n\n PACKAGE_RELATIONSHIPS_NAMESPACE = (\n \"http://schemas.openxmlformats.org/package/2006/relationships\"\n )\n OFFICE_RELATIONSHIPS_NAMESPACE = (\n \"http://schemas.openxmlformats.org/officeDocument/2006/relationships\"\n )\n CONTENT_TYPES_NAMESPACE = (\n \"http://schemas.openxmlformats.org/package/2006/content-types\"\n )\n\n MAIN_CONTENT_FOLDERS = {\"word\", \"ppt\", \"xl\"}\n\n OOXML_NAMESPACES = {\n \"http://schemas.openxmlformats.org/officeDocument/2006/math\",\n \"http://schemas.openxmlformats.org/officeDocument/2006/relationships\",\n \"http://schemas.openxmlformats.org/schemaLibrary/2006/main\",\n \"http://schemas.openxmlformats.org/drawingml/2006/main\",\n \"http://schemas.openxmlformats.org/drawingml/2006/chart\",\n \"http://schemas.openxmlformats.org/drawingml/2006/chartDrawing\",\n \"http://schemas.openxmlformats.org/drawingml/2006/diagram\",\n \"http://schemas.openxmlformats.org/drawingml/2006/picture\",\n \"http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing\",\n \"http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing\",\n \"http://schemas.openxmlformats.org/wordprocessingml/2006/main\",\n \"http://schemas.openxmlformats.org/presentationml/2006/main\",\n \"http://schemas.openxmlformats.org/spreadsheetml/2006/main\",\n \"http://schemas.openxmlformats.org/officeDocument/2006/sharedTypes\",\n \"http://www.w3.org/XML/1998/namespace\",\n }\n\n def __init__(self, unpacked_dir, original_file=None, verbose=False):\n self.unpacked_dir = Path(unpacked_dir).resolve()\n self.original_file = Path(original_file) if original_file else None\n self.verbose = verbose\n\n self.schemas_dir = Path(__file__).parent.parent / \"schemas\"\n\n patterns = [\"*.xml\", \"*.rels\"]\n self.xml_files = [\n f for pattern in patterns for f in self.unpacked_dir.rglob(pattern)\n ]\n\n if not self.xml_files:\n print(f\"Warning: No XML files found in {self.unpacked_dir}\")\n\n def validate(self):\n raise NotImplementedError(\"Subclasses must implement the validate method\")\n\n def repair(self) -> int:\n return self.repair_whitespace_preservation()\n\n def repair_whitespace_preservation(self) -> int:\n repairs = 0\n\n for xml_file in self.xml_files:\n try:\n content = xml_file.read_text(encoding=\"utf-8\")\n dom = defusedxml.minidom.parseString(content)\n modified = False\n\n for elem in dom.getElementsByTagName(\"*\"):\n if elem.tagName.endswith(\":t\") and elem.firstChild:\n text = elem.firstChild.nodeValue\n if text and (\n text.startswith((\" \", \"\\t\")) or text.endswith((\" \", \"\\t\"))\n ):\n if elem.getAttribute(\"xml:space\") != \"preserve\":\n elem.setAttribute(\"xml:space\", \"preserve\")\n text_preview = (\n repr(text[:30]) + \"...\"\n if len(text) > 30\n else repr(text)\n )\n print(\n f\" Repaired: {xml_file.name}: Added xml:space='preserve' to {elem.tagName}: {text_preview}\"\n )\n repairs += 1\n modified = True\n\n if modified:\n xml_file.write_bytes(dom.toxml(encoding=\"UTF-8\"))\n\n except Exception:\n pass\n\n return repairs\n\n def validate_xml(self):\n errors = []\n\n for xml_file in self.xml_files:\n try:\n lxml.etree.parse(str(xml_file))\n except lxml.etree.XMLSyntaxError as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Line {e.lineno}: {e.msg}\"\n )\n except Exception as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Unexpected error: {str(e)}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} XML violations:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - All XML files are well-formed\")\n return True\n\n def validate_namespaces(self):\n errors = []\n\n for xml_file in self.xml_files:\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n declared = set(root.nsmap.keys()) - {None}\n\n for attr_val in [\n v for k, v in root.attrib.items() if k.endswith(\"Ignorable\")\n ]:\n undeclared = set(attr_val.split()) - declared\n errors.extend(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Namespace '{ns}' in Ignorable but not declared\"\n for ns in undeclared\n )\n except lxml.etree.XMLSyntaxError:\n continue\n\n if errors:\n print(f\"FAILED - {len(errors)} namespace issues:\")\n for error in errors:\n print(error)\n return False\n if self.verbose:\n print(\"PASSED - All namespace prefixes properly declared\")\n return True\n\n def validate_unique_ids(self):\n errors = []\n global_ids = {}\n\n for xml_file in self.xml_files:\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n file_ids = {}\n\n mc_elements = root.xpath(\n \".//mc:AlternateContent\", namespaces={\"mc\": self.MC_NAMESPACE}\n )\n for elem in mc_elements:\n elem.getparent().remove(elem)\n\n for elem in root.iter():\n tag = (\n elem.tag.split(\"}\")[-1].lower()\n if \"}\" in elem.tag\n else elem.tag.lower()\n )\n\n if tag in self.UNIQUE_ID_REQUIREMENTS:\n in_excluded_container = any(\n ancestor.tag.split(\"}\")[-1].lower()\n in self.EXCLUDED_ID_CONTAINERS\n for ancestor in elem.iterancestors()\n )\n if in_excluded_container:\n continue\n\n attr_name, scope = self.UNIQUE_ID_REQUIREMENTS[tag]\n\n id_value = None\n for attr, value in elem.attrib.items():\n attr_local = (\n attr.split(\"}\")[-1].lower()\n if \"}\" in attr\n else attr.lower()\n )\n if attr_local == attr_name:\n id_value = value\n break\n\n if id_value is not None:\n if scope == \"global\":\n if id_value in global_ids:\n prev_file, prev_line, prev_tag = global_ids[\n id_value\n ]\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {elem.sourceline}: Global ID '{id_value}' in <{tag}> \"\n f\"already used in {prev_file} at line {prev_line} in <{prev_tag}>\"\n )\n else:\n global_ids[id_value] = (\n xml_file.relative_to(self.unpacked_dir),\n elem.sourceline,\n tag,\n )\n elif scope == \"file\":\n key = (tag, attr_name)\n if key not in file_ids:\n file_ids[key] = {}\n\n if id_value in file_ids[key]:\n prev_line = file_ids[key][id_value]\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {elem.sourceline}: Duplicate {attr_name}='{id_value}' in <{tag}> \"\n f\"(first occurrence at line {prev_line})\"\n )\n else:\n file_ids[key][id_value] = elem.sourceline\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} ID uniqueness violations:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - All required IDs are unique\")\n return True\n\n def validate_file_references(self):\n errors = []\n\n rels_files = list(self.unpacked_dir.rglob(\"*.rels\"))\n\n if not rels_files:\n if self.verbose:\n print(\"PASSED - No .rels files found\")\n return True\n\n all_files = []\n for file_path in self.unpacked_dir.rglob(\"*\"):\n if (\n file_path.is_file()\n and file_path.name != \"[Content_Types].xml\"\n and not file_path.name.endswith(\".rels\")\n ):\n all_files.append(file_path.resolve())\n\n all_referenced_files = set()\n\n if self.verbose:\n print(\n f\"Found {len(rels_files)} .rels files and {len(all_files)} target files\"\n )\n\n for rels_file in rels_files:\n try:\n rels_root = lxml.etree.parse(str(rels_file)).getroot()\n\n rels_dir = rels_file.parent\n\n referenced_files = set()\n broken_refs = []\n\n for rel in rels_root.findall(\n \".//ns:Relationship\",\n namespaces={\"ns\": self.PACKAGE_RELATIONSHIPS_NAMESPACE},\n ):\n target = rel.get(\"Target\")\n if target and not target.startswith((\"http\", \"mailto:\")):\n if target.startswith(\"/\"):\n target_path = self.unpacked_dir / target.lstrip(\"/\")\n elif rels_file.name == \".rels\":\n target_path = self.unpacked_dir / target\n else:\n base_dir = rels_dir.parent\n target_path = base_dir / target\n\n try:\n target_path = target_path.resolve()\n if target_path.exists() and target_path.is_file():\n referenced_files.add(target_path)\n all_referenced_files.add(target_path)\n else:\n broken_refs.append((target, rel.sourceline))\n except (OSError, ValueError):\n broken_refs.append((target, rel.sourceline))\n\n if broken_refs:\n rel_path = rels_file.relative_to(self.unpacked_dir)\n for broken_ref, line_num in broken_refs:\n errors.append(\n f\" {rel_path}: Line {line_num}: Broken reference to {broken_ref}\"\n )\n\n except Exception as e:\n rel_path = rels_file.relative_to(self.unpacked_dir)\n errors.append(f\" Error parsing {rel_path}: {e}\")\n\n unreferenced_files = set(all_files) - all_referenced_files\n\n if unreferenced_files:\n for unref_file in sorted(unreferenced_files):\n unref_rel_path = unref_file.relative_to(self.unpacked_dir)\n errors.append(f\" Unreferenced file: {unref_rel_path}\")\n\n if errors:\n print(f\"FAILED - Found {len(errors)} relationship validation errors:\")\n for error in errors:\n print(error)\n print(\n \"CRITICAL: These errors will cause the document to appear corrupt. \"\n + \"Broken references MUST be fixed, \"\n + \"and unreferenced files MUST be referenced or removed.\"\n )\n return False\n else:\n if self.verbose:\n print(\n \"PASSED - All references are valid and all files are properly referenced\"\n )\n return True\n\n def validate_all_relationship_ids(self):\n import lxml.etree\n\n errors = []\n\n for xml_file in self.xml_files:\n if xml_file.suffix == \".rels\":\n continue\n\n rels_dir = xml_file.parent / \"_rels\"\n rels_file = rels_dir / f\"{xml_file.name}.rels\"\n\n if not rels_file.exists():\n continue\n\n try:\n rels_root = lxml.etree.parse(str(rels_file)).getroot()\n rid_to_type = {}\n\n for rel in rels_root.findall(\n f\".//{{{self.PACKAGE_RELATIONSHIPS_NAMESPACE}}}Relationship\"\n ):\n rid = rel.get(\"Id\")\n rel_type = rel.get(\"Type\", \"\")\n if rid:\n if rid in rid_to_type:\n rels_rel_path = rels_file.relative_to(self.unpacked_dir)\n errors.append(\n f\" {rels_rel_path}: Line {rel.sourceline}: \"\n f\"Duplicate relationship ID '{rid}' (IDs must be unique)\"\n )\n type_name = (\n rel_type.split(\"/\")[-1] if \"/\" in rel_type else rel_type\n )\n rid_to_type[rid] = type_name\n\n xml_root = lxml.etree.parse(str(xml_file)).getroot()\n\n r_ns = self.OFFICE_RELATIONSHIPS_NAMESPACE\n rid_attrs_to_check = [\"id\", \"embed\", \"link\"]\n for elem in xml_root.iter():\n for attr_name in rid_attrs_to_check:\n rid_attr = elem.get(f\"{{{r_ns}}}{attr_name}\")\n if not rid_attr:\n continue\n xml_rel_path = xml_file.relative_to(self.unpacked_dir)\n elem_name = (\n elem.tag.split(\"}\")[-1] if \"}\" in elem.tag else elem.tag\n )\n\n if rid_attr not in rid_to_type:\n errors.append(\n f\" {xml_rel_path}: Line {elem.sourceline}: \"\n f\"<{elem_name}> r:{attr_name} references non-existent relationship '{rid_attr}' \"\n f\"(valid IDs: {', '.join(sorted(rid_to_type.keys())[:5])}{'...' if len(rid_to_type) > 5 else ''})\"\n )\n elif attr_name == \"id\" and self.ELEMENT_RELATIONSHIP_TYPES:\n expected_type = self._get_expected_relationship_type(\n elem_name\n )\n if expected_type:\n actual_type = rid_to_type[rid_attr]\n if expected_type not in actual_type.lower():\n errors.append(\n f\" {xml_rel_path}: Line {elem.sourceline}: \"\n f\"<{elem_name}> references '{rid_attr}' which points to '{actual_type}' \"\n f\"but should point to a '{expected_type}' relationship\"\n )\n\n except Exception as e:\n xml_rel_path = xml_file.relative_to(self.unpacked_dir)\n errors.append(f\" Error processing {xml_rel_path}: {e}\")\n\n if errors:\n print(f\"FAILED - Found {len(errors)} relationship ID reference errors:\")\n for error in errors:\n print(error)\n print(\"\\nThese ID mismatches will cause the document to appear corrupt!\")\n return False\n else:\n if self.verbose:\n print(\"PASSED - All relationship ID references are valid\")\n return True\n\n def _get_expected_relationship_type(self, element_name):\n elem_lower = element_name.lower()\n\n if elem_lower in self.ELEMENT_RELATIONSHIP_TYPES:\n return self.ELEMENT_RELATIONSHIP_TYPES[elem_lower]\n\n if elem_lower.endswith(\"id\") and len(elem_lower) > 2:\n prefix = elem_lower[:-2]\n if prefix.endswith(\"master\"):\n return prefix.lower()\n elif prefix.endswith(\"layout\"):\n return prefix.lower()\n else:\n if prefix == \"sld\":\n return \"slide\"\n return prefix.lower()\n\n if elem_lower.endswith(\"reference\") and len(elem_lower) > 9:\n prefix = elem_lower[:-9]\n return prefix.lower()\n\n return None\n\n def validate_content_types(self):\n errors = []\n\n content_types_file = self.unpacked_dir / \"[Content_Types].xml\"\n if not content_types_file.exists():\n print(\"FAILED - [Content_Types].xml file not found\")\n return False\n\n try:\n root = lxml.etree.parse(str(content_types_file)).getroot()\n declared_parts = set()\n declared_extensions = set()\n\n for override in root.findall(\n f\".//{{{self.CONTENT_TYPES_NAMESPACE}}}Override\"\n ):\n part_name = override.get(\"PartName\")\n if part_name is not None:\n declared_parts.add(part_name.lstrip(\"/\"))\n\n for default in root.findall(\n f\".//{{{self.CONTENT_TYPES_NAMESPACE}}}Default\"\n ):\n extension = default.get(\"Extension\")\n if extension is not None:\n declared_extensions.add(extension.lower())\n\n declarable_roots = {\n \"sld\",\n \"sldLayout\",\n \"sldMaster\",\n \"presentation\",\n \"document\",\n \"workbook\",\n \"worksheet\",\n \"theme\",\n }\n\n media_extensions = {\n \"png\": \"image/png\",\n \"jpg\": \"image/jpeg\",\n \"jpeg\": \"image/jpeg\",\n \"gif\": \"image/gif\",\n \"bmp\": \"image/bmp\",\n \"tiff\": \"image/tiff\",\n \"wmf\": \"image/x-wmf\",\n \"emf\": \"image/x-emf\",\n }\n\n all_files = list(self.unpacked_dir.rglob(\"*\"))\n all_files = [f for f in all_files if f.is_file()]\n\n for xml_file in self.xml_files:\n path_str = str(xml_file.relative_to(self.unpacked_dir)).replace(\n \"\\\\\", \"/\"\n )\n\n if any(\n skip in path_str\n for skip in [\".rels\", \"[Content_Types]\", \"docProps/\", \"_rels/\"]\n ):\n continue\n\n try:\n root_tag = lxml.etree.parse(str(xml_file)).getroot().tag\n root_name = root_tag.split(\"}\")[-1] if \"}\" in root_tag else root_tag\n\n if root_name in declarable_roots and path_str not in declared_parts:\n errors.append(\n f\" {path_str}: File with <{root_name}> root not declared in [Content_Types].xml\"\n )\n\n except Exception:\n continue\n\n for file_path in all_files:\n if file_path.suffix.lower() in {\".xml\", \".rels\"}:\n continue\n if file_path.name == \"[Content_Types].xml\":\n continue\n if \"_rels\" in file_path.parts or \"docProps\" in file_path.parts:\n continue\n\n extension = file_path.suffix.lstrip(\".\").lower()\n if extension and extension not in declared_extensions:\n if extension in media_extensions:\n relative_path = file_path.relative_to(self.unpacked_dir)\n msg = (\n f\" {relative_path}: File with extension '{extension}' \"\n f\"not declared in [Content_Types].xml - should add: \"\n f''\n )\n errors.append(msg)\n\n except Exception as e:\n errors.append(f\" Error parsing [Content_Types].xml: {e}\")\n\n if errors:\n print(f\"FAILED - Found {len(errors)} content type declaration errors:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\n \"PASSED - All content files are properly declared in [Content_Types].xml\"\n )\n return True\n\n def validate_file_against_xsd(self, xml_file, verbose=False):\n xml_file = Path(xml_file).resolve()\n unpacked_dir = self.unpacked_dir.resolve()\n\n is_valid, current_errors = self._validate_single_file_xsd(\n xml_file, unpacked_dir\n )\n\n if is_valid is None:\n return None, set()\n elif is_valid:\n return True, set()\n\n original_errors = self._get_original_file_errors(xml_file)\n\n assert current_errors is not None\n new_errors = current_errors - original_errors\n\n new_errors = {\n e\n for e in new_errors\n if not any(pattern in e for pattern in self.IGNORED_VALIDATION_ERRORS)\n }\n\n if new_errors:\n if verbose:\n relative_path = xml_file.relative_to(unpacked_dir)\n print(f\"FAILED - {relative_path}: {len(new_errors)} new error(s)\")\n for error in list(new_errors)[:3]:\n truncated = error[:250] + \"...\" if len(error) > 250 else error\n print(f\" - {truncated}\")\n return False, new_errors\n else:\n if verbose:\n print(\n f\"PASSED - No new errors (original had {len(current_errors)} errors)\"\n )\n return True, set()\n\n def validate_against_xsd(self):\n new_errors = []\n original_error_count = 0\n valid_count = 0\n skipped_count = 0\n\n for xml_file in self.xml_files:\n relative_path = str(xml_file.relative_to(self.unpacked_dir))\n is_valid, new_file_errors = self.validate_file_against_xsd(\n xml_file, verbose=False\n )\n\n if is_valid is None:\n skipped_count += 1\n continue\n elif is_valid and not new_file_errors:\n valid_count += 1\n continue\n elif is_valid:\n original_error_count += 1\n valid_count += 1\n continue\n\n new_errors.append(f\" {relative_path}: {len(new_file_errors)} new error(s)\")\n for error in list(new_file_errors)[:3]:\n new_errors.append(\n f\" - {error[:250]}...\" if len(error) > 250 else f\" - {error}\"\n )\n\n if self.verbose:\n print(f\"Validated {len(self.xml_files)} files:\")\n print(f\" - Valid: {valid_count}\")\n print(f\" - Skipped (no schema): {skipped_count}\")\n if original_error_count:\n print(f\" - With original errors (ignored): {original_error_count}\")\n print(\n f\" - With NEW errors: {len(new_errors) > 0 and len([e for e in new_errors if not e.startswith(' ')]) or 0}\"\n )\n\n if new_errors:\n print(\"\\nFAILED - Found NEW validation errors:\")\n for error in new_errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"\\nPASSED - No new XSD validation errors introduced\")\n return True\n\n def _get_schema_path(self, xml_file):\n if xml_file.name in self.SCHEMA_MAPPINGS:\n return self.schemas_dir / self.SCHEMA_MAPPINGS[xml_file.name]\n\n if xml_file.suffix == \".rels\":\n return self.schemas_dir / self.SCHEMA_MAPPINGS[\".rels\"]\n\n if \"charts/\" in str(xml_file) and xml_file.name.startswith(\"chart\"):\n return self.schemas_dir / self.SCHEMA_MAPPINGS[\"chart\"]\n\n if \"theme/\" in str(xml_file) and xml_file.name.startswith(\"theme\"):\n return self.schemas_dir / self.SCHEMA_MAPPINGS[\"theme\"]\n\n if xml_file.parent.name in self.MAIN_CONTENT_FOLDERS:\n return self.schemas_dir / self.SCHEMA_MAPPINGS[xml_file.parent.name]\n\n return None\n\n def _clean_ignorable_namespaces(self, xml_doc):\n xml_string = lxml.etree.tostring(xml_doc, encoding=\"unicode\")\n xml_copy = lxml.etree.fromstring(xml_string)\n\n for elem in xml_copy.iter():\n attrs_to_remove = []\n\n for attr in elem.attrib:\n if \"{\" in attr:\n ns = attr.split(\"}\")[0][1:]\n if ns not in self.OOXML_NAMESPACES:\n attrs_to_remove.append(attr)\n\n for attr in attrs_to_remove:\n del elem.attrib[attr]\n\n self._remove_ignorable_elements(xml_copy)\n\n return lxml.etree.ElementTree(xml_copy)\n\n def _remove_ignorable_elements(self, root):\n elements_to_remove = []\n\n for elem in list(root):\n if not hasattr(elem, \"tag\") or callable(elem.tag):\n continue\n\n tag_str = str(elem.tag)\n if tag_str.startswith(\"{\"):\n ns = tag_str.split(\"}\")[0][1:]\n if ns not in self.OOXML_NAMESPACES:\n elements_to_remove.append(elem)\n continue\n\n self._remove_ignorable_elements(elem)\n\n for elem in elements_to_remove:\n root.remove(elem)\n\n def _preprocess_for_mc_ignorable(self, xml_doc):\n root = xml_doc.getroot()\n\n if f\"{{{self.MC_NAMESPACE}}}Ignorable\" in root.attrib:\n del root.attrib[f\"{{{self.MC_NAMESPACE}}}Ignorable\"]\n\n return xml_doc\n\n def _validate_single_file_xsd(self, xml_file, base_path):\n schema_path = self._get_schema_path(xml_file)\n if not schema_path:\n return None, None\n\n try:\n with open(schema_path, \"rb\") as xsd_file:\n parser = lxml.etree.XMLParser()\n xsd_doc = lxml.etree.parse(\n xsd_file, parser=parser, base_url=str(schema_path)\n )\n schema = lxml.etree.XMLSchema(xsd_doc)\n\n with open(xml_file, \"r\") as f:\n xml_doc = lxml.etree.parse(f)\n\n xml_doc, _ = self._remove_template_tags_from_text_nodes(xml_doc)\n xml_doc = self._preprocess_for_mc_ignorable(xml_doc)\n\n relative_path = xml_file.relative_to(base_path)\n if (\n relative_path.parts\n and relative_path.parts[0] in self.MAIN_CONTENT_FOLDERS\n ):\n xml_doc = self._clean_ignorable_namespaces(xml_doc)\n\n if schema.validate(xml_doc):\n return True, set()\n else:\n errors = set()\n for error in schema.error_log:\n errors.add(error.message)\n return False, errors\n\n except Exception as e:\n return False, {str(e)}\n\n def _get_original_file_errors(self, xml_file):\n if self.original_file is None:\n return set()\n\n import tempfile\n import zipfile\n\n xml_file = Path(xml_file).resolve()\n unpacked_dir = self.unpacked_dir.resolve()\n relative_path = xml_file.relative_to(unpacked_dir)\n\n with tempfile.TemporaryDirectory() as temp_dir:\n temp_path = Path(temp_dir)\n\n with zipfile.ZipFile(self.original_file, \"r\") as zip_ref:\n zip_ref.extractall(temp_path)\n\n original_xml_file = temp_path / relative_path\n\n if not original_xml_file.exists():\n return set()\n\n is_valid, errors = self._validate_single_file_xsd(\n original_xml_file, temp_path\n )\n return errors if errors else set()\n\n def _remove_template_tags_from_text_nodes(self, xml_doc):\n warnings = []\n template_pattern = re.compile(r\"\\{\\{[^}]*\\}\\}\")\n\n xml_string = lxml.etree.tostring(xml_doc, encoding=\"unicode\")\n xml_copy = lxml.etree.fromstring(xml_string)\n\n def process_text_content(text, content_type):\n if not text:\n return text\n matches = list(template_pattern.finditer(text))\n if matches:\n for match in matches:\n warnings.append(\n f\"Found template tag in {content_type}: {match.group()}\"\n )\n return template_pattern.sub(\"\", text)\n return text\n\n for elem in xml_copy.iter():\n if not hasattr(elem, \"tag\") or callable(elem.tag):\n continue\n tag_str = str(elem.tag)\n if tag_str.endswith(\"}t\") or tag_str == \"t\":\n continue\n\n elem.text = process_text_content(elem.text, \"text content\")\n elem.tail = process_text_content(elem.tail, \"tail content\")\n\n return lxml.etree.ElementTree(xml_copy), warnings\n\n\nif __name__ == \"__main__\":\n raise RuntimeError(\"This module should not be run directly.\")\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/validators/docx.py", "content": "\"\"\"\nValidator for Word document XML files against XSD schemas.\n\"\"\"\n\nimport random\nimport re\nimport tempfile\nimport zipfile\n\nimport defusedxml.minidom\nimport lxml.etree\n\nfrom .base import BaseSchemaValidator\n\n\nclass DOCXSchemaValidator(BaseSchemaValidator):\n WORD_2006_NAMESPACE = \"http://schemas.openxmlformats.org/wordprocessingml/2006/main\"\n W14_NAMESPACE = \"http://schemas.microsoft.com/office/word/2010/wordml\"\n W16CID_NAMESPACE = \"http://schemas.microsoft.com/office/word/2016/wordml/cid\"\n\n ELEMENT_RELATIONSHIP_TYPES = {}\n\n def validate(self):\n if not self.validate_xml():\n return False\n\n all_valid = True\n if not self.validate_namespaces():\n all_valid = False\n\n if not self.validate_unique_ids():\n all_valid = False\n\n if not self.validate_file_references():\n all_valid = False\n\n if not self.validate_content_types():\n all_valid = False\n\n if not self.validate_against_xsd():\n all_valid = False\n\n if not self.validate_whitespace_preservation():\n all_valid = False\n\n if not self.validate_deletions():\n all_valid = False\n\n if not self.validate_insertions():\n all_valid = False\n\n if not self.validate_all_relationship_ids():\n all_valid = False\n\n if not self.validate_id_constraints():\n all_valid = False\n\n if not self.validate_comment_markers():\n all_valid = False\n\n self.compare_paragraph_counts()\n\n return all_valid\n\n def validate_whitespace_preservation(self):\n errors = []\n\n for xml_file in self.xml_files:\n if xml_file.name != \"document.xml\":\n continue\n\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n\n for elem in root.iter(f\"{{{self.WORD_2006_NAMESPACE}}}t\"):\n if elem.text:\n text = elem.text\n if re.search(r\"^[ \\t\\n\\r]\", text) or re.search(\n r\"[ \\t\\n\\r]$\", text\n ):\n xml_space_attr = f\"{{{self.XML_NAMESPACE}}}space\"\n if (\n xml_space_attr not in elem.attrib\n or elem.attrib[xml_space_attr] != \"preserve\"\n ):\n text_preview = (\n repr(text)[:50] + \"...\"\n if len(repr(text)) > 50\n else repr(text)\n )\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {elem.sourceline}: w:t element with whitespace \"\n f\"missing xml:space='preserve': {text_preview}\"\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} whitespace preservation violations:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - All whitespace is properly preserved\")\n return True\n\n def validate_deletions(self):\n errors = []\n\n for xml_file in self.xml_files:\n if xml_file.name != \"document.xml\":\n continue\n\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n namespaces = {\"w\": self.WORD_2006_NAMESPACE}\n\n for t_elem in root.xpath(\".//w:del//w:t\", namespaces=namespaces):\n if t_elem.text:\n text_preview = (\n repr(t_elem.text)[:50] + \"...\"\n if len(repr(t_elem.text)) > 50\n else repr(t_elem.text)\n )\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {t_elem.sourceline}: found within : {text_preview}\"\n )\n\n for instr_elem in root.xpath(\n \".//w:del//w:instrText\", namespaces=namespaces\n ):\n text_preview = (\n repr(instr_elem.text or \"\")[:50] + \"...\"\n if len(repr(instr_elem.text or \"\")) > 50\n else repr(instr_elem.text or \"\")\n )\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {instr_elem.sourceline}: found within (use ): {text_preview}\"\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} deletion validation violations:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - No w:t elements found within w:del elements\")\n return True\n\n def count_paragraphs_in_unpacked(self):\n count = 0\n\n for xml_file in self.xml_files:\n if xml_file.name != \"document.xml\":\n continue\n\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n paragraphs = root.findall(f\".//{{{self.WORD_2006_NAMESPACE}}}p\")\n count = len(paragraphs)\n except Exception as e:\n print(f\"Error counting paragraphs in unpacked document: {e}\")\n\n return count\n\n def count_paragraphs_in_original(self):\n original = self.original_file\n if original is None:\n return 0\n\n count = 0\n\n try:\n with tempfile.TemporaryDirectory() as temp_dir:\n with zipfile.ZipFile(original, \"r\") as zip_ref:\n zip_ref.extractall(temp_dir)\n\n doc_xml_path = temp_dir + \"/word/document.xml\"\n root = lxml.etree.parse(doc_xml_path).getroot()\n\n paragraphs = root.findall(f\".//{{{self.WORD_2006_NAMESPACE}}}p\")\n count = len(paragraphs)\n\n except Exception as e:\n print(f\"Error counting paragraphs in original document: {e}\")\n\n return count\n\n def validate_insertions(self):\n errors = []\n\n for xml_file in self.xml_files:\n if xml_file.name != \"document.xml\":\n continue\n\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n namespaces = {\"w\": self.WORD_2006_NAMESPACE}\n\n invalid_elements = root.xpath(\n \".//w:ins//w:delText[not(ancestor::w:del)]\", namespaces=namespaces\n )\n\n for elem in invalid_elements:\n text_preview = (\n repr(elem.text or \"\")[:50] + \"...\"\n if len(repr(elem.text or \"\")) > 50\n else repr(elem.text or \"\")\n )\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {elem.sourceline}: within : {text_preview}\"\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} insertion validation violations:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - No w:delText elements within w:ins elements\")\n return True\n\n def compare_paragraph_counts(self):\n original_count = self.count_paragraphs_in_original()\n new_count = self.count_paragraphs_in_unpacked()\n\n diff = new_count - original_count\n diff_str = f\"+{diff}\" if diff > 0 else str(diff)\n print(f\"\\nParagraphs: {original_count} → {new_count} ({diff_str})\")\n\n def _parse_id_value(self, val: str, base: int = 16) -> int:\n return int(val, base)\n\n def validate_id_constraints(self):\n errors = []\n para_id_attr = f\"{{{self.W14_NAMESPACE}}}paraId\"\n durable_id_attr = f\"{{{self.W16CID_NAMESPACE}}}durableId\"\n\n for xml_file in self.xml_files:\n try:\n for elem in lxml.etree.parse(str(xml_file)).iter():\n if val := elem.get(para_id_attr):\n if self._parse_id_value(val, base=16) >= 0x80000000:\n errors.append(\n f\" {xml_file.name}:{elem.sourceline}: paraId={val} >= 0x80000000\"\n )\n\n if val := elem.get(durable_id_attr):\n if xml_file.name == \"numbering.xml\":\n try:\n if self._parse_id_value(val, base=10) >= 0x7FFFFFFF:\n errors.append(\n f\" {xml_file.name}:{elem.sourceline}: durableId={val} >= 0x7FFFFFFF\"\n )\n except ValueError:\n errors.append(\n f\" {xml_file.name}:{elem.sourceline}: durableId={val} must be decimal in numbering.xml\"\n )\n else:\n if self._parse_id_value(val, base=16) >= 0x7FFFFFFF:\n errors.append(\n f\" {xml_file.name}:{elem.sourceline}: durableId={val} >= 0x7FFFFFFF\"\n )\n except Exception:\n pass\n\n if errors:\n print(f\"FAILED - {len(errors)} ID constraint violations:\")\n for e in errors:\n print(e)\n elif self.verbose:\n print(\"PASSED - All paraId/durableId values within constraints\")\n return not errors\n\n def validate_comment_markers(self):\n errors = []\n\n document_xml = None\n comments_xml = None\n for xml_file in self.xml_files:\n if xml_file.name == \"document.xml\" and \"word\" in str(xml_file):\n document_xml = xml_file\n elif xml_file.name == \"comments.xml\":\n comments_xml = xml_file\n\n if not document_xml:\n if self.verbose:\n print(\"PASSED - No document.xml found (skipping comment validation)\")\n return True\n\n try:\n doc_root = lxml.etree.parse(str(document_xml)).getroot()\n namespaces = {\"w\": self.WORD_2006_NAMESPACE}\n\n range_starts = {\n elem.get(f\"{{{self.WORD_2006_NAMESPACE}}}id\")\n for elem in doc_root.xpath(\n \".//w:commentRangeStart\", namespaces=namespaces\n )\n }\n range_ends = {\n elem.get(f\"{{{self.WORD_2006_NAMESPACE}}}id\")\n for elem in doc_root.xpath(\n \".//w:commentRangeEnd\", namespaces=namespaces\n )\n }\n references = {\n elem.get(f\"{{{self.WORD_2006_NAMESPACE}}}id\")\n for elem in doc_root.xpath(\n \".//w:commentReference\", namespaces=namespaces\n )\n }\n\n orphaned_ends = range_ends - range_starts\n for comment_id in sorted(\n orphaned_ends, key=lambda x: int(x) if x and x.isdigit() else 0\n ):\n errors.append(\n f' document.xml: commentRangeEnd id=\"{comment_id}\" has no matching commentRangeStart'\n )\n\n orphaned_starts = range_starts - range_ends\n for comment_id in sorted(\n orphaned_starts, key=lambda x: int(x) if x and x.isdigit() else 0\n ):\n errors.append(\n f' document.xml: commentRangeStart id=\"{comment_id}\" has no matching commentRangeEnd'\n )\n\n comment_ids = set()\n if comments_xml and comments_xml.exists():\n comments_root = lxml.etree.parse(str(comments_xml)).getroot()\n comment_ids = {\n elem.get(f\"{{{self.WORD_2006_NAMESPACE}}}id\")\n for elem in comments_root.xpath(\n \".//w:comment\", namespaces=namespaces\n )\n }\n\n marker_ids = range_starts | range_ends | references\n invalid_refs = marker_ids - comment_ids\n for comment_id in sorted(\n invalid_refs, key=lambda x: int(x) if x and x.isdigit() else 0\n ):\n if comment_id:\n errors.append(\n f' document.xml: marker id=\"{comment_id}\" references non-existent comment'\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(f\" Error parsing XML: {e}\")\n\n if errors:\n print(f\"FAILED - {len(errors)} comment marker violations:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - All comment markers properly paired\")\n return True\n\n def repair(self) -> int:\n repairs = super().repair()\n repairs += self.repair_durableId()\n return repairs\n\n def repair_durableId(self) -> int:\n repairs = 0\n\n for xml_file in self.xml_files:\n try:\n content = xml_file.read_text(encoding=\"utf-8\")\n dom = defusedxml.minidom.parseString(content)\n modified = False\n\n for elem in dom.getElementsByTagName(\"*\"):\n if not elem.hasAttribute(\"w16cid:durableId\"):\n continue\n\n durable_id = elem.getAttribute(\"w16cid:durableId\")\n needs_repair = False\n\n if xml_file.name == \"numbering.xml\":\n try:\n needs_repair = (\n self._parse_id_value(durable_id, base=10) >= 0x7FFFFFFF\n )\n except ValueError:\n needs_repair = True\n else:\n try:\n needs_repair = (\n self._parse_id_value(durable_id, base=16) >= 0x7FFFFFFF\n )\n except ValueError:\n needs_repair = True\n\n if needs_repair:\n value = random.randint(1, 0x7FFFFFFE)\n if xml_file.name == \"numbering.xml\":\n new_id = str(value)\n else:\n new_id = f\"{value:08X}\"\n\n elem.setAttribute(\"w16cid:durableId\", new_id)\n print(\n f\" Repaired: {xml_file.name}: durableId {durable_id} → {new_id}\"\n )\n repairs += 1\n modified = True\n\n if modified:\n xml_file.write_bytes(dom.toxml(encoding=\"UTF-8\"))\n\n except Exception:\n pass\n\n return repairs\n\n\nif __name__ == \"__main__\":\n raise RuntimeError(\"This module should not be run directly.\")\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/validators/pptx.py", "content": "\"\"\"\nValidator for PowerPoint presentation XML files against XSD schemas.\n\"\"\"\n\nimport re\n\nfrom .base import BaseSchemaValidator\n\n\nclass PPTXSchemaValidator(BaseSchemaValidator):\n PRESENTATIONML_NAMESPACE = (\n \"http://schemas.openxmlformats.org/presentationml/2006/main\"\n )\n\n ELEMENT_RELATIONSHIP_TYPES = {\n \"sldid\": \"slide\",\n \"sldmasterid\": \"slidemaster\",\n \"notesmasterid\": \"notesmaster\",\n \"sldlayoutid\": \"slidelayout\",\n \"themeid\": \"theme\",\n \"tablestyleid\": \"tablestyles\",\n }\n\n def validate(self):\n if not self.validate_xml():\n return False\n\n all_valid = True\n if not self.validate_namespaces():\n all_valid = False\n\n if not self.validate_unique_ids():\n all_valid = False\n\n if not self.validate_uuid_ids():\n all_valid = False\n\n if not self.validate_file_references():\n all_valid = False\n\n if not self.validate_slide_layout_ids():\n all_valid = False\n\n if not self.validate_content_types():\n all_valid = False\n\n if not self.validate_against_xsd():\n all_valid = False\n\n if not self.validate_notes_slide_references():\n all_valid = False\n\n if not self.validate_all_relationship_ids():\n all_valid = False\n\n if not self.validate_no_duplicate_slide_layouts():\n all_valid = False\n\n return all_valid\n\n def validate_uuid_ids(self):\n import lxml.etree\n\n errors = []\n uuid_pattern = re.compile(\n r\"^[\\{\\(]?[0-9A-Fa-f]{8}-?[0-9A-Fa-f]{4}-?[0-9A-Fa-f]{4}-?[0-9A-Fa-f]{4}-?[0-9A-Fa-f]{12}[\\}\\)]?$\"\n )\n\n for xml_file in self.xml_files:\n try:\n root = lxml.etree.parse(str(xml_file)).getroot()\n\n for elem in root.iter():\n for attr, value in elem.attrib.items():\n attr_name = attr.split(\"}\")[-1].lower()\n if attr_name == \"id\" or attr_name.endswith(\"id\"):\n if self._looks_like_uuid(value):\n if not uuid_pattern.match(value):\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: \"\n f\"Line {elem.sourceline}: ID '{value}' appears to be \"\n \"a UUID but contains invalid hex characters\"\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {xml_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} UUID ID validation errors:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - All UUID-like IDs contain valid hex values\")\n return True\n\n def _looks_like_uuid(self, value):\n clean_value = value.strip(\"{}()\").replace(\"-\", \"\")\n return len(clean_value) == 32 and all(c.isalnum() for c in clean_value)\n\n def validate_slide_layout_ids(self):\n import lxml.etree\n\n errors = []\n\n slide_masters = list(self.unpacked_dir.glob(\"ppt/slideMasters/*.xml\"))\n\n if not slide_masters:\n if self.verbose:\n print(\"PASSED - No slide masters found\")\n return True\n\n for slide_master in slide_masters:\n try:\n root = lxml.etree.parse(str(slide_master)).getroot()\n\n rels_file = slide_master.parent / \"_rels\" / f\"{slide_master.name}.rels\"\n\n if not rels_file.exists():\n errors.append(\n f\" {slide_master.relative_to(self.unpacked_dir)}: \"\n f\"Missing relationships file: {rels_file.relative_to(self.unpacked_dir)}\"\n )\n continue\n\n rels_root = lxml.etree.parse(str(rels_file)).getroot()\n\n valid_layout_rids = set()\n for rel in rels_root.findall(\n f\".//{{{self.PACKAGE_RELATIONSHIPS_NAMESPACE}}}Relationship\"\n ):\n rel_type = rel.get(\"Type\", \"\")\n if \"slideLayout\" in rel_type:\n valid_layout_rids.add(rel.get(\"Id\"))\n\n for sld_layout_id in root.findall(\n f\".//{{{self.PRESENTATIONML_NAMESPACE}}}sldLayoutId\"\n ):\n r_id = sld_layout_id.get(\n f\"{{{self.OFFICE_RELATIONSHIPS_NAMESPACE}}}id\"\n )\n layout_id = sld_layout_id.get(\"id\")\n\n if r_id and r_id not in valid_layout_rids:\n errors.append(\n f\" {slide_master.relative_to(self.unpacked_dir)}: \"\n f\"Line {sld_layout_id.sourceline}: sldLayoutId with id='{layout_id}' \"\n f\"references r:id='{r_id}' which is not found in slide layout relationships\"\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {slide_master.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(f\"FAILED - Found {len(errors)} slide layout ID validation errors:\")\n for error in errors:\n print(error)\n print(\n \"Remove invalid references or add missing slide layouts to the relationships file.\"\n )\n return False\n else:\n if self.verbose:\n print(\"PASSED - All slide layout IDs reference valid slide layouts\")\n return True\n\n def validate_no_duplicate_slide_layouts(self):\n import lxml.etree\n\n errors = []\n slide_rels_files = list(self.unpacked_dir.glob(\"ppt/slides/_rels/*.xml.rels\"))\n\n for rels_file in slide_rels_files:\n try:\n root = lxml.etree.parse(str(rels_file)).getroot()\n\n layout_rels = [\n rel\n for rel in root.findall(\n f\".//{{{self.PACKAGE_RELATIONSHIPS_NAMESPACE}}}Relationship\"\n )\n if \"slideLayout\" in rel.get(\"Type\", \"\")\n ]\n\n if len(layout_rels) > 1:\n errors.append(\n f\" {rels_file.relative_to(self.unpacked_dir)}: has {len(layout_rels)} slideLayout references\"\n )\n\n except Exception as e:\n errors.append(\n f\" {rels_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n if errors:\n print(\"FAILED - Found slides with duplicate slideLayout references:\")\n for error in errors:\n print(error)\n return False\n else:\n if self.verbose:\n print(\"PASSED - All slides have exactly one slideLayout reference\")\n return True\n\n def validate_notes_slide_references(self):\n import lxml.etree\n\n errors = []\n notes_slide_references = {}\n\n slide_rels_files = list(self.unpacked_dir.glob(\"ppt/slides/_rels/*.xml.rels\"))\n\n if not slide_rels_files:\n if self.verbose:\n print(\"PASSED - No slide relationship files found\")\n return True\n\n for rels_file in slide_rels_files:\n try:\n root = lxml.etree.parse(str(rels_file)).getroot()\n\n for rel in root.findall(\n f\".//{{{self.PACKAGE_RELATIONSHIPS_NAMESPACE}}}Relationship\"\n ):\n rel_type = rel.get(\"Type\", \"\")\n if \"notesSlide\" in rel_type:\n target = rel.get(\"Target\", \"\")\n if target:\n normalized_target = target.replace(\"../\", \"\")\n\n slide_name = rels_file.stem.replace(\".xml\", \"\")\n\n if normalized_target not in notes_slide_references:\n notes_slide_references[normalized_target] = []\n notes_slide_references[normalized_target].append(\n (slide_name, rels_file)\n )\n\n except (lxml.etree.XMLSyntaxError, Exception) as e:\n errors.append(\n f\" {rels_file.relative_to(self.unpacked_dir)}: Error: {e}\"\n )\n\n for target, references in notes_slide_references.items():\n if len(references) > 1:\n slide_names = [ref[0] for ref in references]\n errors.append(\n f\" Notes slide '{target}' is referenced by multiple slides: {', '.join(slide_names)}\"\n )\n for slide_name, rels_file in references:\n errors.append(f\" - {rels_file.relative_to(self.unpacked_dir)}\")\n\n if errors:\n print(\n f\"FAILED - Found {len([e for e in errors if not e.startswith(' ')])} notes slide reference validation errors:\"\n )\n for error in errors:\n print(error)\n print(\"Each slide may optionally have its own slide file.\")\n return False\n else:\n if self.verbose:\n print(\"PASSED - All notes slide references are unique\")\n return True\n\n\nif __name__ == \"__main__\":\n raise RuntimeError(\"This module should not be run directly.\")\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/office/validators/redlining.py", "content": "\"\"\"\nValidator for tracked changes in Word documents.\n\"\"\"\n\nimport subprocess\nimport tempfile\nimport zipfile\nfrom pathlib import Path\n\n\nclass RedliningValidator:\n def __init__(self, unpacked_dir, original_docx, verbose=False, author=\"Claude\"):\n self.unpacked_dir = Path(unpacked_dir)\n self.original_docx = Path(original_docx)\n self.verbose = verbose\n self.author = author\n self.namespaces = {\n \"w\": \"http://schemas.openxmlformats.org/wordprocessingml/2006/main\"\n }\n\n def repair(self) -> int:\n return 0\n\n def validate(self):\n modified_file = self.unpacked_dir / \"word\" / \"document.xml\"\n if not modified_file.exists():\n print(f\"FAILED - Modified document.xml not found at {modified_file}\")\n return False\n\n try:\n import xml.etree.ElementTree as ET\n\n tree = ET.parse(modified_file)\n root = tree.getroot()\n\n del_elements = root.findall(\".//w:del\", self.namespaces)\n ins_elements = root.findall(\".//w:ins\", self.namespaces)\n\n author_del_elements = [\n elem\n for elem in del_elements\n if elem.get(f\"{{{self.namespaces['w']}}}author\") == self.author\n ]\n author_ins_elements = [\n elem\n for elem in ins_elements\n if elem.get(f\"{{{self.namespaces['w']}}}author\") == self.author\n ]\n\n if not author_del_elements and not author_ins_elements:\n if self.verbose:\n print(f\"PASSED - No tracked changes by {self.author} found.\")\n return True\n\n except Exception:\n pass\n\n with tempfile.TemporaryDirectory() as temp_dir:\n temp_path = Path(temp_dir)\n\n try:\n with zipfile.ZipFile(self.original_docx, \"r\") as zip_ref:\n zip_ref.extractall(temp_path)\n except Exception as e:\n print(f\"FAILED - Error unpacking original docx: {e}\")\n return False\n\n original_file = temp_path / \"word\" / \"document.xml\"\n if not original_file.exists():\n print(\n f\"FAILED - Original document.xml not found in {self.original_docx}\"\n )\n return False\n\n try:\n import xml.etree.ElementTree as ET\n\n modified_tree = ET.parse(modified_file)\n modified_root = modified_tree.getroot()\n original_tree = ET.parse(original_file)\n original_root = original_tree.getroot()\n except ET.ParseError as e:\n print(f\"FAILED - Error parsing XML files: {e}\")\n return False\n\n self._remove_author_tracked_changes(original_root)\n self._remove_author_tracked_changes(modified_root)\n\n modified_text = self._extract_text_content(modified_root)\n original_text = self._extract_text_content(original_root)\n\n if modified_text != original_text:\n error_message = self._generate_detailed_diff(\n original_text, modified_text\n )\n print(error_message)\n return False\n\n if self.verbose:\n print(f\"PASSED - All changes by {self.author} are properly tracked\")\n return True\n\n def _generate_detailed_diff(self, original_text, modified_text):\n error_parts = [\n f\"FAILED - Document text doesn't match after removing {self.author}'s tracked changes\",\n \"\",\n \"Likely causes:\",\n \" 1. Modified text inside another author's or tags\",\n \" 2. Made edits without proper tracked changes\",\n \" 3. Didn't nest inside when deleting another's insertion\",\n \"\",\n \"For pre-redlined documents, use correct patterns:\",\n \" - To reject another's INSERTION: Nest inside their \",\n \" - To restore another's DELETION: Add new AFTER their \",\n \"\",\n ]\n\n git_diff = self._get_git_word_diff(original_text, modified_text)\n if git_diff:\n error_parts.extend([\"Differences:\", \"============\", git_diff])\n else:\n error_parts.append(\"Unable to generate word diff (git not available)\")\n\n return \"\\n\".join(error_parts)\n\n def _get_git_word_diff(self, original_text, modified_text):\n try:\n with tempfile.TemporaryDirectory() as temp_dir:\n temp_path = Path(temp_dir)\n\n original_file = temp_path / \"original.txt\"\n modified_file = temp_path / \"modified.txt\"\n\n original_file.write_text(original_text, encoding=\"utf-8\")\n modified_file.write_text(modified_text, encoding=\"utf-8\")\n\n result = subprocess.run(\n [\n \"git\",\n \"diff\",\n \"--word-diff=plain\",\n \"--word-diff-regex=.\",\n \"-U0\",\n \"--no-index\",\n str(original_file),\n str(modified_file),\n ],\n capture_output=True,\n text=True,\n )\n\n if result.stdout.strip():\n lines = result.stdout.split(\"\\n\")\n content_lines = []\n in_content = False\n for line in lines:\n if line.startswith(\"@@\"):\n in_content = True\n continue\n if in_content and line.strip():\n content_lines.append(line)\n\n if content_lines:\n return \"\\n\".join(content_lines)\n\n result = subprocess.run(\n [\n \"git\",\n \"diff\",\n \"--word-diff=plain\",\n \"-U0\",\n \"--no-index\",\n str(original_file),\n str(modified_file),\n ],\n capture_output=True,\n text=True,\n )\n\n if result.stdout.strip():\n lines = result.stdout.split(\"\\n\")\n content_lines = []\n in_content = False\n for line in lines:\n if line.startswith(\"@@\"):\n in_content = True\n continue\n if in_content and line.strip():\n content_lines.append(line)\n return \"\\n\".join(content_lines)\n\n except (subprocess.CalledProcessError, FileNotFoundError, Exception):\n pass\n\n return None\n\n def _remove_author_tracked_changes(self, root):\n ins_tag = f\"{{{self.namespaces['w']}}}ins\"\n del_tag = f\"{{{self.namespaces['w']}}}del\"\n author_attr = f\"{{{self.namespaces['w']}}}author\"\n\n for parent in root.iter():\n to_remove = []\n for child in parent:\n if child.tag == ins_tag and child.get(author_attr) == self.author:\n to_remove.append(child)\n for elem in to_remove:\n parent.remove(elem)\n\n deltext_tag = f\"{{{self.namespaces['w']}}}delText\"\n t_tag = f\"{{{self.namespaces['w']}}}t\"\n\n for parent in root.iter():\n to_process = []\n for child in parent:\n if child.tag == del_tag and child.get(author_attr) == self.author:\n to_process.append((child, list(parent).index(child)))\n\n for del_elem, del_index in reversed(to_process):\n for elem in del_elem.iter():\n if elem.tag == deltext_tag:\n elem.tag = t_tag\n\n for child in reversed(list(del_elem)):\n parent.insert(del_index, child)\n parent.remove(del_elem)\n\n def _extract_text_content(self, root):\n p_tag = f\"{{{self.namespaces['w']}}}p\"\n t_tag = f\"{{{self.namespaces['w']}}}t\"\n\n paragraphs = []\n for p_elem in root.findall(f\".//{p_tag}\"):\n text_parts = []\n for t_elem in p_elem.findall(f\".//{t_tag}\"):\n if t_elem.text:\n text_parts.append(t_elem.text)\n paragraph_text = \"\".join(text_parts)\n if paragraph_text:\n paragraphs.append(paragraph_text)\n\n return \"\\n\".join(paragraphs)\n\n\nif __name__ == \"__main__\":\n raise RuntimeError(\"This module should not be run directly.\")\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/preview.py", "content": "\"\"\"Generate slide preview images from a PowerPoint file.\n\nConverts PPTX -> PDF -> JPEG slides with caching. If cached slides\nalready exist and are up-to-date, returns them without reconverting.\n\nOutput protocol (stdout):\n Line 1: status — one of CACHED, GENERATED, ERROR_NOT_FOUND, ERROR_NO_PDF\n Lines 2+: sorted absolute paths to slide-*.jpg files\n\nUsage:\n python preview.py /path/to/file.pptx /path/to/cache_dir\n\"\"\"\n\nimport os\nimport subprocess\nimport sys\nfrom pathlib import Path\n\n# Allow importing office.soffice from the scripts directory\nsys.path.insert(0, str(Path(__file__).resolve().parent))\n\nfrom office.soffice import run_soffice\n\nCONVERSION_DPI = 150\n\n\ndef _find_slides(directory: Path) -> list[str]:\n \"\"\"Find slide-*.jpg files in directory, sorted by page number.\"\"\"\n slides = list(directory.glob(\"slide-*.jpg\"))\n slides.sort(key=lambda p: int(p.stem.split(\"-\")[-1]))\n return [str(s) for s in slides]\n\n\ndef main() -> None:\n if len(sys.argv) != 3:\n print(f\"Usage: {sys.argv[0]} \", file=sys.stderr)\n sys.exit(1)\n\n pptx_path = Path(sys.argv[1])\n cache_dir = Path(sys.argv[2])\n\n if not pptx_path.is_file():\n print(\"ERROR_NOT_FOUND\")\n return\n\n # Check cache: if slides exist and are at least as new as the PPTX, reuse them\n cached_slides = _find_slides(cache_dir)\n if cached_slides:\n pptx_mtime = os.path.getmtime(pptx_path)\n oldest_slide_mtime = min(os.path.getmtime(s) for s in cached_slides)\n if oldest_slide_mtime >= pptx_mtime:\n print(\"CACHED\")\n for slide in cached_slides:\n print(slide)\n return\n # Stale cache — remove old slides\n for slide in cached_slides:\n os.remove(slide)\n\n cache_dir.mkdir(parents=True, exist_ok=True)\n\n # Convert PPTX -> PDF via LibreOffice\n result = run_soffice(\n [\n \"--headless\",\n \"--convert-to\",\n \"pdf\",\n \"--outdir\",\n str(cache_dir),\n str(pptx_path),\n ],\n capture_output=True,\n text=True,\n )\n if result.returncode != 0:\n print(\"CONVERSION_ERROR\", file=sys.stderr)\n sys.exit(1)\n\n # Find the generated PDF\n pdfs = sorted(cache_dir.glob(\"*.pdf\"))\n if not pdfs:\n print(\"ERROR_NO_PDF\")\n return\n\n pdf_file = pdfs[0]\n\n # Convert PDF -> JPEG slides\n result = subprocess.run(\n [\n \"pdftoppm\",\n \"-jpeg\",\n \"-r\",\n str(CONVERSION_DPI),\n str(pdf_file),\n str(cache_dir / \"slide\"),\n ],\n capture_output=True,\n text=True,\n )\n if result.returncode != 0:\n print(\"CONVERSION_ERROR\", file=sys.stderr)\n sys.exit(1)\n\n # Clean up PDF\n pdf_file.unlink(missing_ok=True)\n\n slides = _find_slides(cache_dir)\n print(\"GENERATED\")\n for slide in slides:\n print(slide)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx/scripts/thumbnail.py", "content": "\"\"\"Create thumbnail grids from PowerPoint presentation slides.\n\nCreates a grid layout of slide thumbnails for quick visual analysis.\nLabels each thumbnail with its XML filename (e.g., slide1.xml).\nHidden slides are shown with a placeholder pattern.\n\nUsage:\n python thumbnail.py input.pptx [output_prefix] [--cols N]\n\nExamples:\n python thumbnail.py presentation.pptx\n # Creates: thumbnails.jpg\n\n python thumbnail.py template.pptx grid --cols 4\n # Creates: grid.jpg (or grid-1.jpg, grid-2.jpg for large decks)\n\"\"\"\n\nimport argparse\nimport subprocess\nimport sys\nimport tempfile\nimport zipfile\nfrom pathlib import Path\n\nimport defusedxml.minidom\nfrom office.soffice import get_soffice_env\nfrom PIL import Image\nfrom PIL import ImageDraw\nfrom PIL import ImageFont\n\nTHUMBNAIL_WIDTH = 300\nCONVERSION_DPI = 100\nMAX_COLS = 6\nDEFAULT_COLS = 3\nJPEG_QUALITY = 95\nGRID_PADDING = 20\nBORDER_WIDTH = 2\nFONT_SIZE_RATIO = 0.10\nLABEL_PADDING_RATIO = 0.4\n\n\ndef main():\n parser = argparse.ArgumentParser(\n description=\"Create thumbnail grids from PowerPoint slides.\"\n )\n parser.add_argument(\"input\", help=\"Input PowerPoint file (.pptx)\")\n parser.add_argument(\n \"output_prefix\",\n nargs=\"?\",\n default=\"thumbnails\",\n help=\"Output prefix for image files (default: thumbnails)\",\n )\n parser.add_argument(\n \"--cols\",\n type=int,\n default=DEFAULT_COLS,\n help=f\"Number of columns (default: {DEFAULT_COLS}, max: {MAX_COLS})\",\n )\n\n args = parser.parse_args()\n\n cols = min(args.cols, MAX_COLS)\n if args.cols > MAX_COLS:\n print(f\"Warning: Columns limited to {MAX_COLS}\")\n\n input_path = Path(args.input)\n if not input_path.exists() or input_path.suffix.lower() != \".pptx\":\n print(f\"Error: Invalid PowerPoint file: {args.input}\", file=sys.stderr)\n sys.exit(1)\n\n output_path = Path(f\"{args.output_prefix}.jpg\")\n\n try:\n slide_info = get_slide_info(input_path)\n\n with tempfile.TemporaryDirectory() as temp_dir:\n temp_path = Path(temp_dir)\n visible_images = convert_to_images(input_path, temp_path)\n\n if not visible_images and not any(s[\"hidden\"] for s in slide_info):\n print(\"Error: No slides found\", file=sys.stderr)\n sys.exit(1)\n\n slides = build_slide_list(slide_info, visible_images, temp_path)\n\n grid_files = create_grids(slides, cols, THUMBNAIL_WIDTH, output_path)\n\n print(f\"Created {len(grid_files)} grid(s):\")\n for grid_file in grid_files:\n print(f\" {grid_file}\")\n\n except Exception as e:\n print(f\"Error: {e}\", file=sys.stderr)\n sys.exit(1)\n\n\ndef get_slide_info(pptx_path: Path) -> list[dict]:\n with zipfile.ZipFile(pptx_path, \"r\") as zf:\n rels_content = zf.read(\"ppt/_rels/presentation.xml.rels\").decode(\"utf-8\")\n rels_dom = defusedxml.minidom.parseString(rels_content)\n\n rid_to_slide = {}\n for rel in rels_dom.getElementsByTagName(\"Relationship\"):\n rid = rel.getAttribute(\"Id\")\n target = rel.getAttribute(\"Target\")\n rel_type = rel.getAttribute(\"Type\")\n if \"slide\" in rel_type and target.startswith(\"slides/\"):\n rid_to_slide[rid] = target.replace(\"slides/\", \"\")\n\n pres_content = zf.read(\"ppt/presentation.xml\").decode(\"utf-8\")\n pres_dom = defusedxml.minidom.parseString(pres_content)\n\n slides = []\n for sld_id in pres_dom.getElementsByTagName(\"p:sldId\"):\n rid = sld_id.getAttribute(\"r:id\")\n if rid in rid_to_slide:\n hidden = sld_id.getAttribute(\"show\") == \"0\"\n slides.append({\"name\": rid_to_slide[rid], \"hidden\": hidden})\n\n return slides\n\n\ndef build_slide_list(\n slide_info: list[dict],\n visible_images: list[Path],\n temp_dir: Path,\n) -> list[tuple[Path, str]]:\n if visible_images:\n with Image.open(visible_images[0]) as img:\n placeholder_size = img.size\n else:\n placeholder_size = (1920, 1080)\n\n slides = []\n visible_idx = 0\n\n for info in slide_info:\n if info[\"hidden\"]:\n placeholder_path = temp_dir / f\"hidden-{info['name']}.jpg\"\n placeholder_img = create_hidden_placeholder(placeholder_size)\n placeholder_img.save(placeholder_path, \"JPEG\")\n slides.append((placeholder_path, f\"{info['name']} (hidden)\"))\n else:\n if visible_idx < len(visible_images):\n slides.append((visible_images[visible_idx], info[\"name\"]))\n visible_idx += 1\n\n return slides\n\n\ndef create_hidden_placeholder(size: tuple[int, int]) -> Image.Image:\n img = Image.new(\"RGB\", size, color=\"#F0F0F0\")\n draw = ImageDraw.Draw(img)\n line_width = max(5, min(size) // 100)\n draw.line([(0, 0), size], fill=\"#CCCCCC\", width=line_width)\n draw.line([(size[0], 0), (0, size[1])], fill=\"#CCCCCC\", width=line_width)\n return img\n\n\ndef convert_to_images(pptx_path: Path, temp_dir: Path) -> list[Path]:\n pdf_path = temp_dir / f\"{pptx_path.stem}.pdf\"\n\n result = subprocess.run(\n [\n \"soffice\",\n \"--headless\",\n \"--convert-to\",\n \"pdf\",\n \"--outdir\",\n str(temp_dir),\n str(pptx_path),\n ],\n capture_output=True,\n text=True,\n env=get_soffice_env(),\n )\n if result.returncode != 0 or not pdf_path.exists():\n raise RuntimeError(\"PDF conversion failed\")\n\n result = subprocess.run(\n [\n \"pdftoppm\",\n \"-jpeg\",\n \"-r\",\n str(CONVERSION_DPI),\n str(pdf_path),\n str(temp_dir / \"slide\"),\n ],\n capture_output=True,\n text=True,\n )\n if result.returncode != 0:\n raise RuntimeError(\"Image conversion failed\")\n\n return sorted(temp_dir.glob(\"slide-*.jpg\"))\n\n\ndef create_grids(\n slides: list[tuple[Path, str]],\n cols: int,\n width: int,\n output_path: Path,\n) -> list[str]:\n max_per_grid = cols * (cols + 1)\n grid_files = []\n\n for chunk_idx, start_idx in enumerate(range(0, len(slides), max_per_grid)):\n end_idx = min(start_idx + max_per_grid, len(slides))\n chunk_slides = slides[start_idx:end_idx]\n\n grid = create_grid(chunk_slides, cols, width)\n\n if len(slides) <= max_per_grid:\n grid_filename = output_path\n else:\n stem = output_path.stem\n suffix = output_path.suffix\n grid_filename = output_path.parent / f\"{stem}-{chunk_idx + 1}{suffix}\"\n\n grid_filename.parent.mkdir(parents=True, exist_ok=True)\n grid.save(str(grid_filename), quality=JPEG_QUALITY)\n grid_files.append(str(grid_filename))\n\n return grid_files\n\n\ndef create_grid(\n slides: list[tuple[Path, str]],\n cols: int,\n width: int,\n) -> Image.Image:\n font_size = int(width * FONT_SIZE_RATIO)\n label_padding = int(font_size * LABEL_PADDING_RATIO)\n\n with Image.open(slides[0][0]) as img:\n aspect = img.height / img.width\n height = int(width * aspect)\n\n rows = (len(slides) + cols - 1) // cols\n grid_w = cols * width + (cols + 1) * GRID_PADDING\n grid_h = rows * (height + font_size + label_padding * 2) + (rows + 1) * GRID_PADDING\n\n grid = Image.new(\"RGB\", (grid_w, grid_h), \"white\")\n draw = ImageDraw.Draw(grid)\n\n try:\n font = ImageFont.load_default(size=font_size)\n except Exception:\n font = ImageFont.load_default()\n\n for i, (img_path, slide_name) in enumerate(slides):\n row, col = i // cols, i % cols\n x = col * width + (col + 1) * GRID_PADDING\n y_base = (\n row * (height + font_size + label_padding * 2) + (row + 1) * GRID_PADDING\n )\n\n label = slide_name\n bbox = draw.textbbox((0, 0), label, font=font)\n text_w = bbox[2] - bbox[0]\n draw.text(\n (x + (width - text_w) // 2, y_base + label_padding),\n label,\n fill=\"black\",\n font=font,\n )\n\n y_thumbnail = y_base + label_padding + font_size + label_padding\n\n with Image.open(img_path) as img:\n img.thumbnail((width, height), Image.Resampling.LANCZOS)\n w, h = img.size\n tx = x + (width - w) // 2\n ty = y_thumbnail + (height - h) // 2\n grid.paste(img, (tx, ty))\n\n if BORDER_WIDTH > 0:\n draw.rectangle(\n [\n (tx - BORDER_WIDTH, ty - BORDER_WIDTH),\n (tx + w + BORDER_WIDTH - 1, ty + h + BORDER_WIDTH - 1),\n ],\n outline=\"gray\",\n width=BORDER_WIDTH,\n )\n\n return grid\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/.gitignore", "content": "# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.\n\n# dependencies\n/node_modules\n/.pnp\n.pnp.*\n.yarn/*\n!.yarn/patches\n!.yarn/plugins\n!.yarn/releases\n!.yarn/versions\n\n# testing\n/coverage\n\n# next.js\n/.next/\n/out/\n\n# production\n/build\n\n# misc\n.DS_Store\n*.pem\n\n# debug\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\n.pnpm-debug.log*\n\n# env files (can opt-in for committing if needed)\n.env*\n\n# vercel\n.vercel\n\n# typescript\n*.tsbuildinfo\nnext-env.d.ts\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/AGENTS.md", "content": "# AGENTS.md\n\nThis file provides guidance to AI agents when working on the web application within this directory.\n\n## Important Notes\n\n- **The development server is already running** at a dynamically allocated port. Do NOT run `npm run dev` yourself.\n- **We do NOT use a `src` directory** - all code lives directly in the root folders (`app/`, `components/`, `lib/`, etc.)\n- If the app needs pre-computation (data processing, API calls, etc.), create a bash or python script called `prepare.sh`/`prepare.py` at the root of this directory\n- **CRITICAL: Create small, modular components** - Do NOT write everything in `page.tsx`. Break your UI into small, reusable components in the `components/` directory. Each component should have a single responsibility and be in its own file.\n\n## Data Preparation Scripts\n\n**CRITICAL: Always re-run data scripts after modifying them.**\n\nIf a `prepare.sh` or `prepare.py` script exists at the root of this directory, it is responsible for generating/loading data that the frontend consumes. \n\n### When to Run the Script\n\nYou MUST run the data preparation script:\n1. **After creating** the script for the first time\n2. **After modifying** the script logic (new data sources, changed processing, etc.)\n3. **After updating** any data files the script reads from\n4. **Before testing** the frontend if you're unsure if data is fresh\n\n### How to Run\n\n```bash\n# For bash scripts\nbash prepare.sh\n\n# For python scripts\npython prepare.py\n```\n\n### Common Mistake\n\n❌ **Updating the script but forgetting to run it** - This leaves stale data in place and the frontend won't reflect your changes. Always run the script immediately after modifying it.\n\n## Commands\n\n```bash\nnpm run dev # Start development server (DO NOT RUN - already running)\nnpm run lint # Run ESLint\n```\n\n## Architecture\n\nThis is a **Next.js 16.1.1** application using the **App Router** with **React 19** and **TypeScript**. It serves as a component showcase/template built on shadcn/ui.\n\n### File Organization Philosophy\n\n**Prioritize small, incremental file writes.** Break your application into many small components rather than monolithic page files.\n\n#### Component Organization\n\n```\ncomponents/\n├── dashboard/ # Feature-specific components\n│ ├── stats-card.tsx\n│ ├── activity-feed.tsx\n│ └── recent-items.tsx\n├── charts/ # Chart components\n│ ├── line-chart.tsx\n│ ├── bar-chart.tsx\n│ └── pie-chart.tsx\n├── data/ # Data display components\n│ ├── data-table.tsx\n│ ├── filter-bar.tsx\n│ └── sort-controls.tsx\n└── layout/ # Layout components\n ├── header.tsx\n ├── sidebar.tsx\n └── footer.tsx\n```\n\n#### Page Structure\n\nPages (`app/page.tsx`) should be **thin orchestration layers** that compose components:\n\n```typescript\n// ✅ GOOD - page.tsx is just composition\nimport { StatsCard } from \"@/components/dashboard/stats-card\";\nimport { ActivityFeed } from \"@/components/dashboard/activity-feed\";\nimport { RecentItems } from \"@/components/dashboard/recent-items\";\n\nexport default function DashboardPage() {\n return (\n
    \n

    Dashboard

    \n
    \n \n \n \n
    \n
    \n \n \n
    \n
    \n );\n}\n\n// ❌ BAD - Everything in page.tsx (500+ lines of mixed logic)\nexport default function DashboardPage() {\n // ... 500 lines of component logic, state, handlers, JSX ...\n}\n```\n\n#### Component Granularity\n\nCreate a new component file when:\n- A UI section has distinct functionality (e.g., `user-profile-card.tsx`)\n- Logic exceeds ~50-100 lines\n- A pattern is reused 2+ times\n- Testing/maintenance would benefit from isolation\n\n**Example: Dashboard Feature**\n\nInstead of writing everything in `app/page.tsx`:\n\n```typescript\n// components/dashboard/stats-card.tsx\nexport function StatsCard({ title, value, trend }: StatsCardProps) {\n return (\n \n \n {title}\n \n \n
    {value}
    \n {trend &&

    {trend}

    }\n     \n     \n );\n}\n\n// components/dashboard/activity-feed.tsx\nexport function ActivityFeed() {\n // Activity feed logic here\n}\n\n// components/dashboard/recent-items.tsx\nexport function RecentItems() {\n // Recent items logic here\n}\n```\n\n#### Benefits of Small Components\n\n1. **Incremental Development**: Write one component at a time, test, iterate\n2. **Better Diffs**: Smaller files = clearer git diffs and easier reviews\n3. **Reusability**: Components can be imported across pages\n4. **Maintainability**: Easier to locate and fix issues\n5. **Hot Reload Efficiency**: Changes to small files reload faster\n6. **Parallel Development**: Multiple features can be worked on independently\n\n### Tech Stack\n\n- **Framework**: Next.js 16.1.1 with App Router\n- **React**: React 19\n- **Language**: TypeScript\n- **Styling**: Tailwind CSS v4 with CSS variables in OKLCH color space\n- **Charts**: recharts for data visualization\n- **UI Components**: shadcn/ui (53 components) built on Radix UI primitives\n- **Variants**: class-variance-authority (CVA) for component variants\n- **Class Merging**: `cn()` utility in `lib/utils.ts` (clsx + tailwind-merge)\n- **Theme**: Dark mode enforced (via `dark` class on ``)\n\n### Key Directories\n\n- `app/` - Next.js App Router pages and layouts\n- `components/ui/` - shadcn/ui component library (Button, Card, Dialog, etc.)\n- `components/` - App-specific components\n- `hooks/` - Custom React hooks (e.g., `use-mobile.ts`)\n- `lib/` - Utilities (`cn()` function)\n\n### Component Patterns\n\n- **Compound Components**: Components like `DropdownMenu`, `Dialog`, `Select` export multiple sub-components (Trigger, Content, Item)\n- **Variants via CVA**: Use `variants` prop for size/style variations (e.g., `buttonVariants`)\n- **Radix UI Primitives**: UI components wrap Radix for accessibility\n\n### Path Aliases\n\nAll imports use `@/` alias (e.g., `@/components/ui/button`, `@/lib/utils`)\n\n### shadcn/ui Configuration\n\nLocated in `components.json`:\n\n- Style: `radix-nova`\n- RSC enabled\n- Icons: lucide-react\n\n### Theme Variables\n\nGlobal CSS variables defined in `app/globals.css` control colors, radius, and spacing. **Dark mode is enforced site-wide** via the `dark` class on the `` element in `app/layout.tsx`. All styling should assume dark mode is active.\n\n### Dark Mode Priority\n\n- **Dark mode is the default and only theme** - do not design for light mode\n- The `dark` class is permanently set on `` in `layout.tsx`\n- Use dark-appropriate colors: `bg-background`, `text-foreground`, etc.\n- Ensure sufficient contrast for dark backgrounds\n- Test all components in dark mode only\n\n## Styling Guidelines\n\n### CRITICAL: Use Only shadcn/ui Components\n\n**MINIMIZE freestyling and creating custom components.** This application uses a complete, professionally designed component library (shadcn/ui). You MUST use the existing components from `components/ui/` for most UI needs.\n\n#### Available shadcn/ui Components\n\nAll components are in `components/ui/`. Import using `@/components/ui/component-name`.\n\n**Layout & Structure:**\n\n- `Card` (`card.tsx`) - Content containers with CardHeader, CardTitle, CardDescription, CardContent, CardFooter\n- `Separator` (`separator.tsx`) - Horizontal/vertical dividers\n- `Tabs` (`tabs.tsx`) - Tabbed interfaces with Tabs, TabsList, TabsTrigger, TabsContent\n- `ScrollArea` (`scroll-area.tsx`) - Styled scrollable regions\n- `Resizable` (`resizable.tsx`) - Resizable panel layouts\n- `Drawer` (`drawer.tsx`) - Bottom/side drawer overlays\n- `Sidebar` (`sidebar.tsx`) - Application sidebar layout\n- `AspectRatio` (`aspect-ratio.tsx`) - Maintain aspect ratios\n\n**Forms & Inputs:**\n\n- `Button` (`button.tsx`) - Primary, secondary, destructive, outline, ghost, link variants\n- `ButtonGroup` (`button-group.tsx`) - Group of related buttons\n- `Input` (`input.tsx`) - Text inputs with various states\n- `InputGroup` (`input-group.tsx`) - Input with addons/icons\n- `Textarea` (`textarea.tsx`) - Multi-line text input\n- `Checkbox` (`checkbox.tsx`) - Checkboxes with indeterminate state\n- `RadioGroup` (`radio-group.tsx`) - Radio button groups\n- `Switch` (`switch.tsx`) - Toggle switches\n- `Select` (`select.tsx`) - Dropdown select menus\n- `NativeSelect` (`native-select.tsx`) - Native HTML select\n- `Combobox` (`combobox.tsx`) - Autocomplete select with search\n- `Command` (`command.tsx`) - Command palette/search interface\n- `Field` (`field.tsx`) - Form field wrapper with label and error\n- `Label` (`label.tsx`) - Form labels with proper accessibility\n- `Slider` (`slider.tsx`) - Range sliders\n- `Calendar` (`calendar.tsx`) - Date picker calendar\n- `Toggle` (`toggle.tsx`) - Toggle button\n- `ToggleGroup` (`toggle-group.tsx`) - Group of toggle buttons\n\n**Navigation:**\n\n- `NavigationMenu` (`navigation-menu.tsx`) - Complex navigation menus\n- `Menubar` (`menubar.tsx`) - Application menu bar\n- `Breadcrumb` (`breadcrumb.tsx`) - Breadcrumb navigation\n- `Pagination` (`pagination.tsx`) - Page navigation controls\n\n**Feedback & Overlays:**\n\n- `Dialog` (`dialog.tsx`) - Modal dialogs\n- `AlertDialog` (`alert-dialog.tsx`) - Confirmation dialogs\n- `Sheet` (`sheet.tsx`) - Side sheets/panels\n- `Popover` (`popover.tsx`) - Floating popovers\n- `HoverCard` (`hover-card.tsx`) - Hover-triggered cards\n- `Tooltip` (`tooltip.tsx`) - Tooltips on hover\n- `Sonner` (`sonner.tsx`) - Toast notifications\n- `Alert` (`alert.tsx`) - Static alert messages\n- `Progress` (`progress.tsx`) - Progress bars\n- `Skeleton` (`skeleton.tsx`) - Loading skeletons\n- `Spinner` (`spinner.tsx`) - Loading spinners\n- `Empty` (`empty.tsx`) - Empty state placeholder\n\n**Menus & Dropdowns:**\n\n- `DropdownMenu` (`dropdown-menu.tsx`) - Dropdown menus with submenus\n- `ContextMenu` (`context-menu.tsx`) - Right-click context menus\n\n**Data Display:**\n\n- `Table` (`table.tsx`) - Data tables with Table, TableHeader, TableBody, TableRow, TableCell, etc.\n- `Badge` (`badge.tsx`) - Status badges and tags\n- `Avatar` (`avatar.tsx`) - User avatars with fallbacks\n- `Accordion` (`accordion.tsx`) - Collapsible content sections\n- `Collapsible` (`collapsible.tsx`) - Simple collapse/expand\n- `Carousel` (`carousel.tsx`) - Image/content carousels\n- `Item` (`item.tsx`) - List item component\n- `Kbd` (`kbd.tsx`) - Keyboard shortcut display\n\n**Data Visualization:**\n\n- `Chart` (`chart.tsx`) - Chart wrapper with ChartContainer, ChartTooltip, ChartTooltipContent, ChartLegend, ChartLegendContent\n\n### Component Usage Principles\n\n#### 1. **Never Create Custom Components**\n\n```typescript\n// ❌ WRONG - Do not create freestyle components\nfunction CustomCard({ title, children }) {\n return (\n
    \n

    {title}

    \n {children}\n
    \n );\n}\n\n// ✅ CORRECT - Use shadcn Card\nimport { Card, CardHeader, CardTitle, CardContent } from \"@/components/ui/card\";\n\nfunction MyComponent() {\n return (\n \n \n Title\n \n Content here\n \n );\n}\n```\n\n#### 2. **Use Component Variants, Don't Style Directly**\n\n```typescript\n// ❌ WRONG - Applying custom Tailwind classes\n\n\n// ✅ CORRECT - Use Button variants\nimport { Button } from \"@/components/ui/button\";\n\n\n\n\n\n\n\n```\n\n#### 3. **Compose Compound Components**\n\nMany shadcn components export multiple sub-components. Use them as designed:\n\n```typescript\n// ✅ Dropdown Menu Composition\nimport {\n DropdownMenu,\n DropdownMenuTrigger,\n DropdownMenuContent,\n DropdownMenuItem,\n DropdownMenuSeparator,\n DropdownMenuLabel,\n} from \"@/components/ui/dropdown-menu\";\n\n\n \n \n \n \n Actions\n \n Edit\n Delete\n \n\n```\n\n#### 4. **Use Layout Components for Structure**\n\n```typescript\n// ✅ Use Card for content sections\nimport { Card, CardHeader, CardTitle, CardDescription, CardContent, CardFooter } from \"@/components/ui/card\";\n\n\n \n Dashboard\n Overview of your data\n \n \n {/* Your content */}\n \n \n \n \n\n```\n\n### Styling Rules\n\n#### 1. **Spacing & Layout**\n\nUse Tailwind's utility classes for spacing, but stick to the design system:\n\n- Gap: `gap-2`, `gap-4`, `gap-6`, `gap-8`\n- Padding: `p-2`, `p-4`, `p-6`, `p-8`\n- Margins: Prefer `gap` and `space-y-*` over margins\n\n#### 2. **Colors**\n\nAll colors come from CSS variables in `app/globals.css`. Use semantic color classes:\n\n- `bg-background`, `bg-foreground`\n- `bg-card`, `text-card-foreground`\n- `bg-primary`, `text-primary-foreground`\n- `bg-secondary`, `text-secondary-foreground`\n- `bg-muted`, `text-muted-foreground`\n- `bg-accent`, `text-accent-foreground`\n- `bg-destructive`, `text-destructive-foreground`\n- `border-border`, `border-input`\n- `ring-ring`\n\n**DO NOT use arbitrary color values** like `bg-blue-500` or `text-red-600`.\n\n#### **CRITICAL: Color Contrast Pairing Rules**\n\n**Always pair background colors with their matching foreground colors.** The color system uses paired variables where each background has a corresponding text color designed for proper contrast.\n\n| Background Class | Text Class to Use | Description |\n|-----------------|-------------------|-------------|\n| `bg-background` | `text-foreground` | Main page background |\n| `bg-card` | `text-card-foreground` | Card containers |\n| `bg-primary` | `text-primary-foreground` | Primary buttons/accents |\n| `bg-secondary` | `text-secondary-foreground` | Secondary elements |\n| `bg-muted` | `text-muted-foreground` | Muted/subtle areas |\n| `bg-accent` | `text-accent-foreground` | Accent highlights |\n| `bg-destructive` | `text-destructive-foreground` | Error/delete actions |\n\n**Examples:**\n\n```typescript\n// ✅ CORRECT - Matching background and foreground pairs\n
    Content
    \n\n
    Subtle text
    \n\n// ❌ WRONG - Mismatched colors causing contrast issues\n
    Invisible text!
    \n
    May have poor contrast
    \n\n```\n\n**Key Rules:**\n\n1. **Never use the same color for background and text** (e.g., `bg-foreground text-foreground`)\n2. **Always use the `-foreground` variant for text** when using a colored background\n3. **For text on `bg-background`**, use `text-foreground` (primary) or `text-muted-foreground` (secondary)\n4. **Test visually** - if text is hard to read, you have a contrast problem\n\n#### 3. **Typography**\n\nUse Tailwind text utilities (no separate Typography component):\n\n- Headings: `text-xl font-semibold`, `text-2xl font-bold`, etc.\n- Body: `text-sm`, `text-base`\n- Secondary text: `text-muted-foreground`\n- Use semantic HTML: `

    `, `

    `, `

    `, etc.\n- **Always wrap text** - Use `max-w-prose` or `max-w-xl` for readable line lengths\n- **Prevent overflow** - Use `break-words` or `truncate` for long text that might overflow containers\n\n#### 4. **Responsive Design**\n\nUse Tailwind's responsive prefixes:\n\n```typescript\n

    \n {/* Responsive grid */}\n
    \n```\n\n#### 5. **Icons**\n\nUse Lucide React icons (already configured):\n\n```typescript\nimport { Check, X, ChevronDown, User } from \"lucide-react\";\n\n\n```\n\n### Data Visualization\n\nFor charts and data visualization, use the **shadcn/ui Chart components** (`@/components/ui/chart`) which wrap recharts with consistent theming. Charts should be **elegant, informative, and digestible at a glance**.\n\n#### Chart Design Principles\n\n1. **Clarity over complexity** - A chart should communicate ONE key insight immediately\n2. **Minimal visual noise** - Remove anything that doesn't add information\n3. **Consistent styling** - Use `ChartConfig` for colors, not arbitrary values\n4. **Responsive** - Always use `ChartContainer` (includes ResponsiveContainer)\n5. **Accessible** - Use `ChartTooltip` with `ChartTooltipContent` for proper styling\n\n#### Chart Type Selection\n\n| Data Type | Recommended Chart | Use Case |\n|-----------|-------------------|----------|\n| Trend over time | `LineChart` or `AreaChart` | Stock prices, user growth, metrics over days/months |\n| Comparing categories | `BarChart` | Revenue by product, users by region |\n| Part of whole | `PieChart` or `RadialBarChart` | Market share, budget allocation |\n| Distribution | `BarChart` (horizontal) | Survey responses, rating distribution |\n| Correlation | `ScatterChart` | Price vs. quality, age vs. income |\n\n#### shadcn/ui Chart Components\n\nAlways import from the shadcn chart component:\n\n```typescript\nimport {\n ChartContainer,\n ChartTooltip,\n ChartTooltipContent,\n ChartLegend,\n ChartLegendContent,\n type ChartConfig,\n} from \"@/components/ui/chart\";\nimport { LineChart, Line, XAxis, YAxis, CartesianGrid } from \"recharts\";\n```\n\n#### ChartConfig - Define Colors and Labels\n\nThe `ChartConfig` object defines colors and labels for your data series. This ensures consistent theming:\n\n```typescript\nconst chartConfig = {\n revenue: {\n label: \"Revenue\",\n color: \"var(--chart-1)\",\n },\n expenses: {\n label: \"Expenses\", \n color: \"var(--chart-2)\",\n },\n} satisfies ChartConfig;\n```\n\n#### Basic Line Chart Template\n\n```typescript\nimport {\n ChartContainer,\n ChartTooltip,\n ChartTooltipContent,\n type ChartConfig,\n} from \"@/components/ui/chart\";\nimport { LineChart, Line, XAxis, YAxis, CartesianGrid } from \"recharts\";\n\nconst chartConfig = {\n value: {\n label: \"Value\",\n color: \"var(--chart-1)\",\n },\n} satisfies ChartConfig;\n\n\n \n \n \n \n } />\n \n \n\n```\n\n#### Bar Chart with Multiple Series\n\n```typescript\nconst chartConfig = {\n revenue: {\n label: \"Revenue\",\n color: \"var(--chart-1)\",\n },\n expenses: {\n label: \"Expenses\",\n color: \"var(--chart-2)\",\n },\n} satisfies ChartConfig;\n\n\n \n \n \n \n } />\n } />\n \n \n \n\n```\n\n#### Pie/Donut Chart\n\n```typescript\nconst chartConfig = {\n desktop: { label: \"Desktop\", color: \"var(--chart-1)\" },\n mobile: { label: \"Mobile\", color: \"var(--chart-2)\" },\n tablet: { label: \"Tablet\", color: \"var(--chart-3)\" },\n} satisfies ChartConfig;\n\n\n \n } />\n \n } />\n \n\n```\n\n#### Chart Styling Rules\n\n**Colors (use CSS variables from globals.css):**\n- `var(--chart-1)` through `var(--chart-5)` - Primary chart colors\n- `var(--primary)` - For single-series emphasis\n- `var(--muted)` - For de-emphasized data\n\n**Color References in Charts:**\n- In `ChartConfig`: Use `color: \"var(--chart-1)\"`\n- In chart elements: Use `fill=\"var(--color-keyname)\"` or `stroke=\"var(--color-keyname)\"`\n- The `keyname` matches the key in your `ChartConfig`\n\n**Visual Cleanup:**\n- Set `tickLine={false}` and `axisLine={false}` on axes for cleaner look\n- Use `vertical={false}` on `CartesianGrid` for horizontal-only grid lines\n- Use `dot={false}` on line charts unless individual points matter\n- Add `radius={4}` to bars for rounded corners\n- Limit to 3-5 data series maximum per chart\n\n**Avoid:**\n- ❌ 3D effects\n- ❌ More than 5-6 colors in one chart\n- ❌ Legends with more than 5 items (simplify the data instead)\n- ❌ Dual Y-axes (confusing - use two separate charts)\n- ❌ Pie charts with more than 5-6 slices\n- ❌ Custom tooltip styling - use `ChartTooltipContent`\n\n#### Fallback to Raw Recharts\n\nIf shadcn/ui Chart components don't support a specific chart type (e.g., ScatterChart, ComposedChart, RadarChart), you can use recharts directly:\n\n```typescript\nimport { ScatterChart, Scatter, XAxis, YAxis, CartesianGrid, Tooltip, ResponsiveContainer } from \"recharts\";\n\n\n \n \n \n \n \n \n \n\n```\n\n**When using raw recharts:**\n- Still use CSS variables for colors (`var(--chart-1)`, etc.)\n- Match styling to shadcn conventions (tickLine={false}, axisLine={false})\n- Style tooltips to match the design system\n\n#### Data Accuracy Checklist\n\nBefore displaying a chart, verify:\n- [ ] `ChartConfig` keys match your data's `dataKey` values\n- [ ] Data values are correctly mapped to the right axes\n- [ ] Axis labels match the data units (%, $, count, etc.)\n- [ ] Time series data is sorted chronologically\n- [ ] No missing data points that would break the visualization\n- [ ] `ChartTooltip` with `ChartTooltipContent` is included\n- [ ] Chart title/context makes the insight clear\n\n### Common Patterns\n\n#### Loading States\n\n```typescript\nimport { Skeleton } from \"@/components/ui/skeleton\";\n\n{isLoading ? (\n \n) : (\n \n)}\n```\n\n#### Empty States\n\n```typescript\nimport { Empty, EmptyHeader, EmptyTitle, EmptyDescription, EmptyMedia } from \"@/components/ui/empty\";\nimport { Inbox } from \"lucide-react\";\n\n\n \n \n \n \n No data available\n \n There's nothing to display yet. Add some items to get started.\n \n \n\n```\n\n#### Interactive Lists\n\n```typescript\nimport { ScrollArea } from \"@/components/ui/scroll-area\";\nimport { ItemGroup, Item, ItemContent, ItemTitle, ItemDescription, ItemMedia } from \"@/components/ui/item\";\nimport { FileText } from \"lucide-react\";\n\n\n \n {items.map((item) => (\n \n \n \n \n \n {item.name}\n {item.description}\n \n \n ))}\n \n\n```\n\n#### Form Fields\n\n```typescript\nimport { Field, FieldLabel, FieldDescription, FieldError, FieldGroup } from \"@/components/ui/field\";\nimport { Input } from \"@/components/ui/input\";\nimport { Button } from \"@/components/ui/button\";\n\n\n \n Email\n \n We'll never share your email.\n \n \n Password\n \n Password must be at least 8 characters.\n \n \n\n```\n\n### What NOT To Do\n\n❌ **Don't create custom styled divs when a component exists**\n❌ **Don't use arbitrary Tailwind colors** (use CSS variables)\n❌ **Don't import UI libraries** like Material-UI, Ant Design, etc.\n❌ **Don't use inline styles** except for dynamic values\n❌ **Don't create custom form inputs** (use Field, Input, Select, etc. from components/ui)\n❌ **Don't add new dependencies** without checking if shadcn covers it\n❌ **Don't write everything in page.tsx** - break into separate component files\n❌ **Don't design for light mode** - this site is dark mode only\n❌ **Don't use `dark:` variants** - dark mode is always active, use base classes\n\n### Development Workflow\n\n1. **Plan the component structure** - Identify logical UI sections before writing code\n2. **Create components incrementally** - Write one small component file at a time\n3. **Test each component** - Verify it works before moving to the next\n4. **Compose in page.tsx** - Import and arrange your components in the page\n5. **Iterate** - Refine individual components without touching others\n\n### Summary\n\nThis application has a **complete, production-ready component library**. Your job is to:\n1. **Compose** shadcn/ui components (from `components/ui/`)\n2. **Create small, focused component files** (in `components/`)\n3. **Keep pages thin** - pages should orchestrate components, not contain implementation\n\nThink of yourself as assembling LEGO blocks—all the UI pieces you need already exist in `components/ui/`, and you create small, organized structures by composing them into feature-specific components.\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/app/globals.css", "content": "@import \"tailwindcss\";\n@import \"tw-animate-css\";\n@import \"shadcn/tailwind.css\";\n\n@custom-variant dark (&:is(.dark *));\n\n@theme inline {\n --color-background: var(--background);\n --color-foreground: var(--foreground);\n --font-sans: var(--font-sans);\n --font-mono: var(--font-geist-mono);\n --color-sidebar-ring: var(--sidebar-ring);\n --color-sidebar-border: var(--sidebar-border);\n --color-sidebar-accent-foreground: var(--sidebar-accent-foreground);\n --color-sidebar-accent: var(--sidebar-accent);\n --color-sidebar-primary-foreground: var(--sidebar-primary-foreground);\n --color-sidebar-primary: var(--sidebar-primary);\n --color-sidebar-foreground: var(--sidebar-foreground);\n --color-sidebar: var(--sidebar);\n --color-chart-5: var(--chart-5);\n --color-chart-4: var(--chart-4);\n --color-chart-3: var(--chart-3);\n --color-chart-2: var(--chart-2);\n --color-chart-1: var(--chart-1);\n --color-ring: var(--ring);\n --color-input: var(--input);\n --color-border: var(--border);\n --color-destructive: var(--destructive);\n --color-accent-foreground: var(--accent-foreground);\n --color-accent: var(--accent);\n --color-muted-foreground: var(--muted-foreground);\n --color-muted: var(--muted);\n --color-secondary-foreground: var(--secondary-foreground);\n --color-secondary: var(--secondary);\n --color-primary-foreground: var(--primary-foreground);\n --color-primary: var(--primary);\n --color-popover-foreground: var(--popover-foreground);\n --color-popover: var(--popover);\n --color-card-foreground: var(--card-foreground);\n --color-card: var(--card);\n --radius-sm: calc(var(--radius) - 4px);\n --radius-md: calc(var(--radius) - 2px);\n --radius-lg: var(--radius);\n --radius-xl: calc(var(--radius) + 4px);\n --radius-2xl: calc(var(--radius) + 8px);\n --radius-3xl: calc(var(--radius) + 12px);\n --radius-4xl: calc(var(--radius) + 16px);\n}\n\n:root {\n --background: oklch(1 0 0);\n --foreground: oklch(0.145 0 0);\n --card: oklch(1 0 0);\n --card-foreground: oklch(0.145 0 0);\n --popover: oklch(1 0 0);\n --popover-foreground: oklch(0.145 0 0);\n --primary: oklch(0.67 0.16 58);\n --primary-foreground: oklch(0.99 0.02 95);\n --secondary: oklch(0.967 0.001 286.375);\n --secondary-foreground: oklch(0.21 0.006 285.885);\n --muted: oklch(0.97 0 0);\n --muted-foreground: oklch(0.556 0 0);\n --accent: oklch(0.97 0 0);\n --accent-foreground: oklch(0.205 0 0);\n --destructive: oklch(0.58 0.22 27);\n --border: oklch(0.922 0 0);\n --input: oklch(0.922 0 0);\n --ring: oklch(0.708 0 0);\n --chart-1: oklch(0.88 0.15 92);\n --chart-2: oklch(0.77 0.16 70);\n --chart-3: oklch(0.67 0.16 58);\n --chart-4: oklch(0.56 0.15 49);\n --chart-5: oklch(0.47 0.12 46);\n --radius: 0.625rem;\n --sidebar: oklch(0.985 0 0);\n --sidebar-foreground: oklch(0.145 0 0);\n --sidebar-primary: oklch(0.67 0.16 58);\n --sidebar-primary-foreground: oklch(0.99 0.02 95);\n --sidebar-accent: oklch(0.97 0 0);\n --sidebar-accent-foreground: oklch(0.205 0 0);\n --sidebar-border: oklch(0.922 0 0);\n --sidebar-ring: oklch(0.708 0 0);\n}\n\n.dark {\n --background: oklch(0.145 0 0);\n --foreground: oklch(0.985 0 0);\n --card: oklch(0.205 0 0);\n --card-foreground: oklch(0.985 0 0);\n --popover: oklch(0.205 0 0);\n --popover-foreground: oklch(0.985 0 0);\n --primary: oklch(0.77 0.16 70);\n --primary-foreground: oklch(0.28 0.07 46);\n --secondary: oklch(0.274 0.006 286.033);\n --secondary-foreground: oklch(0.985 0 0);\n --muted: oklch(0.269 0 0);\n --muted-foreground: oklch(0.708 0 0);\n --accent: oklch(0.371 0 0);\n --accent-foreground: oklch(0.985 0 0);\n --destructive: oklch(0.704 0.191 22.216);\n --border: oklch(1 0 0 / 10%);\n --input: oklch(1 0 0 / 15%);\n --ring: oklch(0.556 0 0);\n /* Chart colors optimized for dark backgrounds - brighter and more vibrant */\n --chart-1: oklch(0.82 0.18 140);\n --chart-2: oklch(0.75 0.2 200);\n --chart-3: oklch(0.7 0.22 280);\n --chart-4: oklch(0.78 0.18 50);\n --chart-5: oklch(0.72 0.2 330);\n --sidebar: oklch(0.205 0 0);\n --sidebar-foreground: oklch(0.985 0 0);\n --sidebar-primary: oklch(0.77 0.16 70);\n --sidebar-primary-foreground: oklch(0.28 0.07 46);\n --sidebar-accent: oklch(0.269 0 0);\n --sidebar-accent-foreground: oklch(0.985 0 0);\n --sidebar-border: oklch(1 0 0 / 10%);\n --sidebar-ring: oklch(0.556 0 0);\n}\n\n@layer base {\n * {\n @apply border-border outline-ring/50;\n }\n body {\n @apply bg-background text-foreground;\n }\n}\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/app/layout.tsx", "content": "import type { Metadata } from \"next\";\nimport { Geist, Geist_Mono, Inter } from \"next/font/google\";\nimport \"./globals.css\";\n\nconst inter = Inter({ subsets: [\"latin\"], variable: \"--font-sans\" });\n\nconst geistSans = Geist({\n variable: \"--font-geist-sans\",\n subsets: [\"latin\"],\n});\n\nconst geistMono = Geist_Mono({\n variable: \"--font-geist-mono\",\n subsets: [\"latin\"],\n});\n\nexport const metadata: Metadata = {\n title: \"Onyx Craft\",\n description: \"Crafting your next great idea.\",\n};\n\nexport default function RootLayout({\n children,\n}: Readonly<{\n children: React.ReactNode;\n}>) {\n return (\n \n \n {children}\n \n \n );\n}\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/app/page.tsx", "content": "\"use client\";\n\nimport { useState, useEffect, useRef } from \"react\";\n\nconst messages = [\n \"Punching wood...\",\n \"Gathering resources...\",\n \"Placing blocks...\",\n \"Crafting your workspace...\",\n \"Mining for dependencies...\",\n \"Smelting the code...\",\n \"Enchanting with magic...\",\n \"World generation complete...\",\n \"/gamemode 1\",\n];\n\nconst MESSAGE_COUNT = messages.length;\nconst TYPE_DELAY = 40;\nconst LINE_PAUSE = 800;\nconst RESET_DELAY = 2000;\n\nexport default function CraftingLoader() {\n const [display, setDisplay] = useState({\n lines: [] as string[],\n currentText: \"\",\n });\n\n const lineIndexRef = useRef(0);\n const charIndexRef = useRef(0);\n const lastUpdateRef = useRef(0);\n const timeoutRef = useRef(undefined);\n const rafRef = useRef(undefined);\n\n useEffect(() => {\n let isActive = true;\n\n const update = (now: number) => {\n if (!isActive) return;\n\n const lineIdx = lineIndexRef.current;\n const charIdx = charIndexRef.current;\n\n if (lineIdx >= MESSAGE_COUNT) {\n timeoutRef.current = setTimeout(() => {\n if (!isActive) return;\n lineIndexRef.current = 0;\n charIndexRef.current = 0;\n setDisplay({ lines: [], currentText: \"\" });\n lastUpdateRef.current = performance.now();\n rafRef.current = requestAnimationFrame(update);\n }, RESET_DELAY);\n return;\n }\n\n const msg = messages[lineIdx];\n if (!msg) return;\n\n const elapsed = now - lastUpdateRef.current;\n\n if (charIdx < msg.length) {\n if (elapsed >= TYPE_DELAY) {\n charIndexRef.current = charIdx + 1;\n setDisplay((prev) => ({\n lines: prev.lines,\n currentText: msg.substring(0, charIdx + 1),\n }));\n lastUpdateRef.current = now;\n }\n } else if (elapsed >= LINE_PAUSE) {\n setDisplay((prev) => ({\n lines: [...prev.lines, msg],\n currentText: \"\",\n }));\n lineIndexRef.current = lineIdx + 1;\n charIndexRef.current = 0;\n lastUpdateRef.current = now;\n }\n\n rafRef.current = requestAnimationFrame(update);\n };\n\n lastUpdateRef.current = performance.now();\n rafRef.current = requestAnimationFrame(update);\n\n return () => {\n isActive = false;\n if (rafRef.current !== undefined) cancelAnimationFrame(rafRef.current);\n if (timeoutRef.current !== undefined) clearTimeout(timeoutRef.current);\n };\n }, []);\n\n const { lines, currentText } = display;\n const hasCurrentText = currentText.length > 0;\n\n return (\n
    \n
    \n
    \n
    \n
    \n
    \n \n crafting_table\n \n
    \n\n
    \n {lines.map((line, i) => (\n
    \n />\n {line}\n
    \n ))}\n {hasCurrentText && (\n
    \n />\n {currentText}\n \n
    \n )}\n {!hasCurrentText && (\n
    \n />\n \n
    \n )}\n
    \n
    \n\n

    \n Crafting your next great idea...\n

    \n
    \n );\n}\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/app/site.webmanifest", "content": "{\"name\":\"\",\"short_name\":\"\",\"icons\":[{\"src\":\"/android-chrome-192x192.png\",\"sizes\":\"192x192\",\"type\":\"image/png\"},{\"src\":\"/android-chrome-512x512.png\",\"sizes\":\"512x512\",\"type\":\"image/png\"}],\"theme_color\":\"#ffffff\",\"background_color\":\"#ffffff\",\"display\":\"standalone\"}" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/component-example.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\n\nimport { Example, ExampleWrapper } from \"@/components/example\";\nimport {\n AlertDialog,\n AlertDialogAction,\n AlertDialogCancel,\n AlertDialogContent,\n AlertDialogDescription,\n AlertDialogFooter,\n AlertDialogHeader,\n AlertDialogMedia,\n AlertDialogTitle,\n AlertDialogTrigger,\n} from \"@/components/ui/alert-dialog\";\nimport { Badge } from \"@/components/ui/badge\";\nimport { Button } from \"@/components/ui/button\";\nimport {\n Card,\n CardAction,\n CardContent,\n CardDescription,\n CardFooter,\n CardHeader,\n CardTitle,\n} from \"@/components/ui/card\";\nimport {\n Combobox,\n ComboboxContent,\n ComboboxEmpty,\n ComboboxInput,\n ComboboxItem,\n ComboboxList,\n} from \"@/components/ui/combobox\";\nimport {\n DropdownMenu,\n DropdownMenuCheckboxItem,\n DropdownMenuContent,\n DropdownMenuGroup,\n DropdownMenuItem,\n DropdownMenuLabel,\n DropdownMenuPortal,\n DropdownMenuRadioGroup,\n DropdownMenuRadioItem,\n DropdownMenuSeparator,\n DropdownMenuShortcut,\n DropdownMenuSub,\n DropdownMenuSubContent,\n DropdownMenuSubTrigger,\n DropdownMenuTrigger,\n} from \"@/components/ui/dropdown-menu\";\nimport { Field, FieldGroup, FieldLabel } from \"@/components/ui/field\";\nimport { Input } from \"@/components/ui/input\";\nimport {\n Select,\n SelectContent,\n SelectGroup,\n SelectItem,\n SelectTrigger,\n SelectValue,\n} from \"@/components/ui/select\";\nimport { Textarea } from \"@/components/ui/textarea\";\nimport {\n PlusIcon,\n BluetoothIcon,\n MoreVerticalIcon,\n FileIcon,\n FolderIcon,\n FolderOpenIcon,\n FileCodeIcon,\n MoreHorizontalIcon,\n FolderSearchIcon,\n SaveIcon,\n DownloadIcon,\n EyeIcon,\n LayoutIcon,\n PaletteIcon,\n SunIcon,\n MoonIcon,\n MonitorIcon,\n UserIcon,\n CreditCardIcon,\n SettingsIcon,\n KeyboardIcon,\n LanguagesIcon,\n BellIcon,\n MailIcon,\n ShieldIcon,\n HelpCircleIcon,\n FileTextIcon,\n LogOutIcon,\n} from \"lucide-react\";\n\nexport function ComponentExample() {\n return (\n \n \n \n \n );\n}\n\nfunction CardExample() {\n return (\n \n \n
    \n \n \n Observability Plus is replacing Monitoring\n \n Switch to the improved way to explore your data, with natural\n language. Monitoring will no longer be available on the Pro plan in\n November, 2025\n \n \n \n \n \n \n \n \n \n \n \n \n Allow accessory to connect?\n \n Do you want to allow the USB accessory to connect to this\n device?\n \n \n \n Don't allow\n Allow\n \n \n \n \n Warning\n \n \n \n \n );\n}\n\nconst frameworks = [\n \"Next.js\",\n \"SvelteKit\",\n \"Nuxt.js\",\n \"Remix\",\n \"Astro\",\n] as const;\n\nfunction FormExample() {\n const [notifications, setNotifications] = React.useState({\n email: true,\n sms: false,\n push: true,\n });\n const [theme, setTheme] = React.useState(\"light\");\n\n return (\n \n \n \n User Information\n Please fill in your details below\n \n \n \n \n \n \n \n File\n \n \n New File\n ⌘N\n \n \n \n New Folder\n ⇧⌘N\n \n \n \n \n Open Recent\n \n \n \n \n Recent Projects\n \n \n Project Alpha\n \n \n \n Project Beta\n \n \n \n \n More Projects\n \n \n \n \n \n Project Gamma\n \n \n \n Project Delta\n \n \n \n \n \n \n \n \n \n Browse...\n \n \n \n \n \n \n \n \n Save\n ⌘S\n \n \n \n Export\n ⇧⌘E\n \n \n \n \n View\n \n setNotifications({\n ...notifications,\n email: checked === true,\n })\n }\n >\n \n Show Sidebar\n \n \n setNotifications({\n ...notifications,\n sms: checked === true,\n })\n }\n >\n \n Show Status Bar\n \n \n \n \n Theme\n \n \n \n \n Appearance\n \n \n \n Light\n \n \n \n Dark\n \n \n \n System\n \n \n \n \n \n \n \n \n \n Account\n \n \n Profile\n ⇧⌘P\n \n \n \n Billing\n \n \n \n \n Settings\n \n \n \n \n Preferences\n \n \n Keyboard Shortcuts\n \n \n \n Language\n \n \n \n \n Notifications\n \n \n \n \n \n Notification Types\n \n \n setNotifications({\n ...notifications,\n push: checked === true,\n })\n }\n >\n \n Push Notifications\n \n \n setNotifications({\n ...notifications,\n email: checked === true,\n })\n }\n >\n \n Email Notifications\n \n \n \n \n \n \n \n \n \n \n Privacy & Security\n \n \n \n \n \n \n \n \n \n \n Help & Support\n \n \n \n Documentation\n \n \n \n \n \n \n Sign Out\n ⇧⌘Q\n \n \n \n \n \n \n \n
    \n \n
    \n \n Name\n \n \n \n Role\n \n \n
    \n \n \n Framework\n \n \n \n \n No frameworks found.\n \n {(item) => (\n \n {item}\n \n )}\n \n \n \n \n \n Comments\n \n \n \n \n \n \n     \n
    \n     \n     \n     \n );\n}\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/example.tsx", "content": "import { cn } from \"@/lib/utils\";\n\nfunction ExampleWrapper({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n
    \n \n
    \n );\n}\n\nfunction Example({\n title,\n children,\n className,\n containerClassName,\n ...props\n}: React.ComponentProps<\"div\"> & {\n title?: string;\n containerClassName?: string;\n}) {\n return (\n \n {title && (\n
    \n {title}\n
    \n )}\n \n {children}\n
    \n
    \n );\n}\n\nexport { ExampleWrapper, Example };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/accordion.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport { Accordion as AccordionPrimitive } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\nimport { ChevronDownIcon, ChevronUpIcon } from \"lucide-react\";\n\nfunction Accordion({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AccordionItem({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AccordionTrigger({\n className,\n children,\n ...props\n}: React.ComponentProps) {\n return (\n \n \n {children}\n \n \n );\n}\n\nfunction AccordionContent({\n className,\n children,\n ...props\n}: React.ComponentProps) {\n return (\n \n \n {children}\n
    \n \n );\n}\n\nexport { Accordion, AccordionItem, AccordionTrigger, AccordionContent };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/alert-dialog.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport { AlertDialog as AlertDialogPrimitive } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\nimport { Button } from \"@/components/ui/button\";\n\nfunction AlertDialog({\n ...props\n}: React.ComponentProps) {\n return ;\n}\n\nfunction AlertDialogTrigger({\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AlertDialogPortal({\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AlertDialogOverlay({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AlertDialogContent({\n className,\n size = \"default\",\n ...props\n}: React.ComponentProps & {\n size?: \"default\" | \"sm\";\n}) {\n return (\n \n \n \n \n );\n}\n\nfunction AlertDialogHeader({\n className,\n ...props\n}: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction AlertDialogFooter({\n className,\n ...props\n}: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction AlertDialogMedia({\n className,\n ...props\n}: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction AlertDialogTitle({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AlertDialogDescription({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AlertDialogAction({\n className,\n variant = \"default\",\n size = \"default\",\n ...props\n}: React.ComponentProps &\n Pick, \"variant\" | \"size\">) {\n return (\n \n );\n}\n\nfunction AlertDialogCancel({\n className,\n variant = \"outline\",\n size = \"default\",\n ...props\n}: React.ComponentProps &\n Pick, \"variant\" | \"size\">) {\n return (\n \n );\n}\n\nexport {\n AlertDialog,\n AlertDialogAction,\n AlertDialogCancel,\n AlertDialogContent,\n AlertDialogDescription,\n AlertDialogFooter,\n AlertDialogHeader,\n AlertDialogMedia,\n AlertDialogOverlay,\n AlertDialogPortal,\n AlertDialogTitle,\n AlertDialogTrigger,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/alert.tsx", "content": "import * as React from \"react\";\nimport { cva, type VariantProps } from \"class-variance-authority\";\n\nimport { cn } from \"@/lib/utils\";\n\nconst alertVariants = cva(\n \"grid gap-0.5 rounded-lg border px-2.5 py-2 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert\",\n {\n variants: {\n variant: {\n default: \"bg-card text-card-foreground\",\n destructive:\n \"text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current\",\n },\n },\n defaultVariants: {\n variant: \"default\",\n },\n },\n);\n\nfunction Alert({\n className,\n variant,\n ...props\n}: React.ComponentProps<\"div\"> & VariantProps) {\n return (\n \n );\n}\n\nfunction AlertTitle({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3\",\n className,\n )}\n {...props}\n />\n );\n}\n\nfunction AlertDescription({\n className,\n ...props\n}: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction AlertAction({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nexport { Alert, AlertTitle, AlertDescription, AlertAction };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/aspect-ratio.tsx", "content": "\"use client\";\n\nimport { AspectRatio as AspectRatioPrimitive } from \"radix-ui\";\n\nfunction AspectRatio({\n ...props\n}: React.ComponentProps) {\n return ;\n}\n\nexport { AspectRatio };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/avatar.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport { Avatar as AvatarPrimitive } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\n\nfunction Avatar({\n className,\n size = \"default\",\n ...props\n}: React.ComponentProps & {\n size?: \"default\" | \"sm\" | \"lg\";\n}) {\n return (\n \n );\n}\n\nfunction AvatarImage({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AvatarFallback({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction AvatarBadge({ className, ...props }: React.ComponentProps<\"span\">) {\n return (\n svg]:hidden\",\n \"group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2\",\n \"group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2\",\n className,\n )}\n {...props}\n />\n );\n}\n\nfunction AvatarGroup({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction AvatarGroupCount({\n className,\n ...props\n}: React.ComponentProps<\"div\">) {\n return (\n svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2\",\n className,\n )}\n {...props}\n />\n );\n}\n\nexport {\n Avatar,\n AvatarImage,\n AvatarFallback,\n AvatarGroup,\n AvatarGroupCount,\n AvatarBadge,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/badge.tsx", "content": "import * as React from \"react\";\nimport { cva, type VariantProps } from \"class-variance-authority\";\nimport { Slot } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\n\nconst badgeVariants = cva(\n \"h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge\",\n {\n variants: {\n variant: {\n default: \"bg-primary text-primary-foreground [a]:hover:bg-primary/80\",\n secondary:\n \"bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80\",\n destructive:\n \"bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20\",\n outline:\n \"border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground\",\n ghost:\n \"hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50\",\n link: \"text-primary underline-offset-4 hover:underline\",\n },\n },\n defaultVariants: {\n variant: \"default\",\n },\n },\n);\n\nfunction Badge({\n className,\n variant = \"default\",\n asChild = false,\n ...props\n}: React.ComponentProps<\"span\"> &\n VariantProps & { asChild?: boolean }) {\n const Comp = asChild ? Slot.Root : \"span\";\n\n return (\n \n );\n}\n\nexport { Badge, badgeVariants };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/breadcrumb.tsx", "content": "import * as React from \"react\";\nimport { Slot } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\nimport { ChevronRightIcon, MoreHorizontalIcon } from \"lucide-react\";\n\nfunction Breadcrumb({ className, ...props }: React.ComponentProps<\"nav\">) {\n return (\n \n );\n}\n\nfunction BreadcrumbList({ className, ...props }: React.ComponentProps<\"ol\">) {\n return (\n \n );\n}\n\nfunction BreadcrumbItem({ className, ...props }: React.ComponentProps<\"li\">) {\n return (\n \n );\n}\n\nfunction BreadcrumbLink({\n asChild,\n className,\n ...props\n}: React.ComponentProps<\"a\"> & {\n asChild?: boolean;\n}) {\n const Comp = asChild ? Slot.Root : \"a\";\n\n return (\n \n );\n}\n\nfunction BreadcrumbPage({ className, ...props }: React.ComponentProps<\"span\">) {\n return (\n \n );\n}\n\nfunction BreadcrumbSeparator({\n children,\n className,\n ...props\n}: React.ComponentProps<\"li\">) {\n return (\n svg]:size-3.5\", className)}\n {...props}\n >\n {children ?? }\n \n );\n}\n\nfunction BreadcrumbEllipsis({\n className,\n ...props\n}: React.ComponentProps<\"span\">) {\n return (\n svg]:size-4 flex items-center justify-center\",\n className,\n )}\n {...props}\n >\n \n More\n \n );\n}\n\nexport {\n Breadcrumb,\n BreadcrumbList,\n BreadcrumbItem,\n BreadcrumbLink,\n BreadcrumbPage,\n BreadcrumbSeparator,\n BreadcrumbEllipsis,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/button-group.tsx", "content": "import { cva, type VariantProps } from \"class-variance-authority\";\nimport { Slot } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\nimport { Separator } from \"@/components/ui/separator\";\n\nconst buttonGroupVariants = cva(\n \"has-[>[data-slot=button-group]]:gap-2 has-[select[aria-hidden=true]:last-child]:[&>[data-slot=select-trigger]:last-of-type]:rounded-r-lg flex w-fit items-stretch [&>*]:focus-visible:z-10 [&>*]:focus-visible:relative [&>[data-slot=select-trigger]:not([class*='w-'])]:w-fit [&>input]:flex-1\",\n {\n variants: {\n orientation: {\n horizontal:\n \"[&>[data-slot]:not(:has(~[data-slot]))]:rounded-r-lg! [&>*:not(:first-child)]:rounded-l-none [&>*:not(:first-child)]:border-l-0 [&>*:not(:last-child)]:rounded-r-none\",\n vertical:\n \"[&>[data-slot]:not(:has(~[data-slot]))]:rounded-b-lg! flex-col [&>*:not(:first-child)]:rounded-t-none [&>*:not(:first-child)]:border-t-0 [&>*:not(:last-child)]:rounded-b-none\",\n },\n },\n defaultVariants: {\n orientation: \"horizontal\",\n },\n },\n);\n\nfunction ButtonGroup({\n className,\n orientation,\n ...props\n}: React.ComponentProps<\"div\"> & VariantProps) {\n return (\n \n );\n}\n\nfunction ButtonGroupText({\n className,\n asChild = false,\n ...props\n}: React.ComponentProps<\"div\"> & {\n asChild?: boolean;\n}) {\n const Comp = asChild ? Slot.Root : \"div\";\n\n return (\n \n );\n}\n\nfunction ButtonGroupSeparator({\n className,\n orientation = \"vertical\",\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nexport {\n ButtonGroup,\n ButtonGroupSeparator,\n ButtonGroupText,\n buttonGroupVariants,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/button.tsx", "content": "import * as React from \"react\";\nimport { cva, type VariantProps } from \"class-variance-authority\";\nimport { Slot } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\n\nconst buttonVariants = cva(\n \"focus-visible:border-ring focus-visible:ring-ring/50 aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive dark:aria-invalid:border-destructive/50 rounded-lg border border-transparent bg-clip-padding text-sm font-medium focus-visible:ring-[3px] aria-invalid:ring-[3px] [&_svg:not([class*='size-'])]:size-4 inline-flex items-center justify-center whitespace-nowrap transition-all disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none shrink-0 [&_svg]:shrink-0 outline-none group/button select-none\",\n {\n variants: {\n variant: {\n default: \"bg-primary text-primary-foreground [a]:hover:bg-primary/80\",\n outline:\n \"border-border bg-background hover:bg-muted hover:text-foreground dark:bg-input/30 dark:border-input dark:hover:bg-input/50 aria-expanded:bg-muted aria-expanded:text-foreground\",\n secondary:\n \"bg-secondary text-secondary-foreground hover:bg-secondary/80 aria-expanded:bg-secondary aria-expanded:text-secondary-foreground\",\n ghost:\n \"hover:bg-muted hover:text-foreground dark:hover:bg-muted/50 aria-expanded:bg-muted aria-expanded:text-foreground\",\n destructive:\n \"bg-destructive/10 hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 dark:bg-destructive/20 text-destructive focus-visible:border-destructive/40 dark:hover:bg-destructive/30\",\n link: \"text-primary underline-offset-4 hover:underline\",\n },\n size: {\n default:\n \"h-8 gap-1.5 px-2.5 has-data-[icon=inline-end]:pr-2 has-data-[icon=inline-start]:pl-2\",\n xs: \"h-6 gap-1 rounded-[min(var(--radius-md),10px)] px-2 text-xs in-data-[slot=button-group]:rounded-lg has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&_svg:not([class*='size-'])]:size-3\",\n sm: \"h-7 gap-1 rounded-[min(var(--radius-md),12px)] px-2.5 text-[0.8rem] in-data-[slot=button-group]:rounded-lg has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&_svg:not([class*='size-'])]:size-3.5\",\n lg: \"h-9 gap-1.5 px-2.5 has-data-[icon=inline-end]:pr-3 has-data-[icon=inline-start]:pl-3\",\n icon: \"size-8\",\n \"icon-xs\":\n \"size-6 rounded-[min(var(--radius-md),10px)] in-data-[slot=button-group]:rounded-lg [&_svg:not([class*='size-'])]:size-3\",\n \"icon-sm\":\n \"size-7 rounded-[min(var(--radius-md),12px)] in-data-[slot=button-group]:rounded-lg\",\n \"icon-lg\": \"size-9\",\n },\n },\n defaultVariants: {\n variant: \"default\",\n size: \"default\",\n },\n },\n);\n\nfunction Button({\n className,\n variant = \"default\",\n size = \"default\",\n asChild = false,\n ...props\n}: React.ComponentProps<\"button\"> &\n VariantProps & {\n asChild?: boolean;\n }) {\n const Comp = asChild ? Slot.Root : \"button\";\n\n return (\n \n );\n}\n\nexport { Button, buttonVariants };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/calendar.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport {\n DayPicker,\n getDefaultClassNames,\n type DayButton,\n} from \"react-day-picker\";\n\nimport { cn } from \"@/lib/utils\";\nimport { Button, buttonVariants } from \"@/components/ui/button\";\nimport {\n ChevronLeftIcon,\n ChevronRightIcon,\n ChevronDownIcon,\n} from \"lucide-react\";\n\nfunction Calendar({\n className,\n classNames,\n showOutsideDays = true,\n captionLayout = \"label\",\n buttonVariant = \"ghost\",\n formatters,\n components,\n ...props\n}: React.ComponentProps & {\n buttonVariant?: React.ComponentProps[\"variant\"];\n}) {\n const defaultClassNames = getDefaultClassNames();\n\n return (\n svg]:rotate-180`,\n String.raw`rtl:**:[.rdp-button\\_previous>svg]:rotate-180`,\n className,\n )}\n captionLayout={captionLayout}\n formatters={{\n formatMonthDropdown: (date) =>\n date.toLocaleString(\"default\", { month: \"short\" }),\n ...formatters,\n }}\n classNames={{\n root: cn(\"w-fit\", defaultClassNames.root),\n months: cn(\n \"flex gap-4 flex-col md:flex-row relative\",\n defaultClassNames.months,\n ),\n month: cn(\"flex flex-col w-full gap-4\", defaultClassNames.month),\n nav: cn(\n \"flex items-center gap-1 w-full absolute top-0 inset-x-0 justify-between\",\n defaultClassNames.nav,\n ),\n button_previous: cn(\n buttonVariants({ variant: buttonVariant }),\n \"size-(--cell-size) aria-disabled:opacity-50 p-0 select-none\",\n defaultClassNames.button_previous,\n ),\n button_next: cn(\n buttonVariants({ variant: buttonVariant }),\n \"size-(--cell-size) aria-disabled:opacity-50 p-0 select-none\",\n defaultClassNames.button_next,\n ),\n month_caption: cn(\n \"flex items-center justify-center h-(--cell-size) w-full px-(--cell-size)\",\n defaultClassNames.month_caption,\n ),\n dropdowns: cn(\n \"w-full flex items-center text-sm font-medium justify-center h-(--cell-size) gap-1.5\",\n defaultClassNames.dropdowns,\n ),\n dropdown_root: cn(\n \"relative cn-calendar-dropdown-root rounded-(--cell-radius)\",\n defaultClassNames.dropdown_root,\n ),\n dropdown: cn(\n \"absolute bg-popover inset-0 opacity-0\",\n defaultClassNames.dropdown,\n ),\n caption_label: cn(\n \"select-none font-medium\",\n captionLayout === \"label\"\n ? \"text-sm\"\n : \"cn-calendar-caption-label rounded-(--cell-radius) flex items-center gap-1 text-sm [&>svg]:text-muted-foreground [&>svg]:size-3.5\",\n defaultClassNames.caption_label,\n ),\n table: \"w-full border-collapse\",\n weekdays: cn(\"flex\", defaultClassNames.weekdays),\n weekday: cn(\n \"text-muted-foreground rounded-(--cell-radius) flex-1 font-normal text-[0.8rem] select-none\",\n defaultClassNames.weekday,\n ),\n week: cn(\"flex w-full mt-2\", defaultClassNames.week),\n week_number_header: cn(\n \"select-none w-(--cell-size)\",\n defaultClassNames.week_number_header,\n ),\n week_number: cn(\n \"text-[0.8rem] select-none text-muted-foreground\",\n defaultClassNames.week_number,\n ),\n day: cn(\n \"relative w-full rounded-(--cell-radius) h-full p-0 text-center [&:last-child[data-selected=true]_button]:rounded-r-(--cell-radius) group/day aspect-square select-none\",\n props.showWeekNumber\n ? \"[&:nth-child(2)[data-selected=true]_button]:rounded-l-(--cell-radius)\"\n : \"[&:first-child[data-selected=true]_button]:rounded-l-(--cell-radius)\",\n defaultClassNames.day,\n ),\n range_start: cn(\n \"rounded-l-(--cell-radius) bg-muted relative after:bg-muted after:absolute after:inset-y-0 after:w-4 after:right-0 -z-0 isolate\",\n defaultClassNames.range_start,\n ),\n range_middle: cn(\"rounded-none\", defaultClassNames.range_middle),\n range_end: cn(\n \"rounded-r-(--cell-radius) bg-muted relative after:bg-muted-200 after:absolute after:inset-y-0 after:w-4 after:left-0 -z-0 isolate\",\n defaultClassNames.range_end,\n ),\n today: cn(\n \"bg-muted text-foreground rounded-(--cell-radius) data-[selected=true]:rounded-none\",\n defaultClassNames.today,\n ),\n outside: cn(\n \"text-muted-foreground aria-selected:text-muted-foreground\",\n defaultClassNames.outside,\n ),\n disabled: cn(\n \"text-muted-foreground opacity-50\",\n defaultClassNames.disabled,\n ),\n hidden: cn(\"invisible\", defaultClassNames.hidden),\n ...classNames,\n }}\n components={{\n Root: ({ className, rootRef, ...props }) => {\n return (\n \n );\n },\n Chevron: ({ className, orientation, ...props }) => {\n if (orientation === \"left\") {\n return (\n \n );\n }\n\n if (orientation === \"right\") {\n return (\n \n );\n }\n\n return (\n \n );\n },\n DayButton: CalendarDayButton,\n WeekNumber: ({ children, ...props }) => {\n return (\n \n
    \n {children}\n
    \n \n );\n },\n ...components,\n }}\n {...props}\n />\n );\n}\n\nfunction CalendarDayButton({\n className,\n day,\n modifiers,\n ...props\n}: React.ComponentProps) {\n const defaultClassNames = getDefaultClassNames();\n\n const ref = React.useRef(null);\n React.useEffect(() => {\n if (modifiers.focused) ref.current?.focus();\n }, [modifiers.focused]);\n\n return (\n span]:text-xs [&>span]:opacity-70\",\n defaultClassNames.day,\n className,\n )}\n {...props}\n />\n );\n}\n\nexport { Calendar, CalendarDayButton };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/card.tsx", "content": "import * as React from \"react\";\n\nimport { cn } from \"@/lib/utils\";\n\nfunction Card({\n className,\n size = \"default\",\n ...props\n}: React.ComponentProps<\"div\"> & { size?: \"default\" | \"sm\" }) {\n return (\n img:first-child]:pt-0 data-[size=sm]:gap-3 data-[size=sm]:py-3 data-[size=sm]:has-data-[slot=card-footer]:pb-0 *:[img:first-child]:rounded-t-xl *:[img:last-child]:rounded-b-xl group/card flex flex-col\",\n className,\n )}\n {...props}\n />\n );\n}\n\nfunction CardHeader({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction CardTitle({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction CardDescription({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction CardAction({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction CardContent({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nfunction CardFooter({ className, ...props }: React.ComponentProps<\"div\">) {\n return (\n \n );\n}\n\nexport {\n Card,\n CardHeader,\n CardFooter,\n CardTitle,\n CardAction,\n CardDescription,\n CardContent,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/carousel.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport useEmblaCarousel, {\n type UseEmblaCarouselType,\n} from \"embla-carousel-react\";\n\nimport { cn } from \"@/lib/utils\";\nimport { Button } from \"@/components/ui/button\";\nimport { ChevronLeftIcon, ChevronRightIcon } from \"lucide-react\";\n\ntype CarouselApi = UseEmblaCarouselType[1];\ntype UseCarouselParameters = Parameters;\ntype CarouselOptions = UseCarouselParameters[0];\ntype CarouselPlugin = UseCarouselParameters[1];\n\ntype CarouselProps = {\n opts?: CarouselOptions;\n plugins?: CarouselPlugin;\n orientation?: \"horizontal\" | \"vertical\";\n setApi?: (api: CarouselApi) => void;\n};\n\ntype CarouselContextProps = {\n carouselRef: ReturnType[0];\n api: ReturnType[1];\n scrollPrev: () => void;\n scrollNext: () => void;\n canScrollPrev: boolean;\n canScrollNext: boolean;\n} & CarouselProps;\n\nconst CarouselContext = React.createContext(null);\n\nfunction useCarousel() {\n const context = React.useContext(CarouselContext);\n\n if (!context) {\n throw new Error(\"useCarousel must be used within a \");\n }\n\n return context;\n}\n\nfunction Carousel({\n orientation = \"horizontal\",\n opts,\n setApi,\n plugins,\n className,\n children,\n ...props\n}: React.ComponentProps<\"div\"> & CarouselProps) {\n const [carouselRef, api] = useEmblaCarousel(\n {\n ...opts,\n axis: orientation === \"horizontal\" ? \"x\" : \"y\",\n },\n plugins,\n );\n const [canScrollPrev, setCanScrollPrev] = React.useState(false);\n const [canScrollNext, setCanScrollNext] = React.useState(false);\n\n const onSelect = React.useCallback((api: CarouselApi) => {\n if (!api) return;\n setCanScrollPrev(api.canScrollPrev());\n setCanScrollNext(api.canScrollNext());\n }, []);\n\n const scrollPrev = React.useCallback(() => {\n api?.scrollPrev();\n }, [api]);\n\n const scrollNext = React.useCallback(() => {\n api?.scrollNext();\n }, [api]);\n\n const handleKeyDown = React.useCallback(\n (event: React.KeyboardEvent) => {\n if (event.key === \"ArrowLeft\") {\n event.preventDefault();\n scrollPrev();\n } else if (event.key === \"ArrowRight\") {\n event.preventDefault();\n scrollNext();\n }\n },\n [scrollPrev, scrollNext],\n );\n\n React.useEffect(() => {\n if (!api || !setApi) return;\n setApi(api);\n }, [api, setApi]);\n\n React.useEffect(() => {\n if (!api) return;\n onSelect(api);\n api.on(\"reInit\", onSelect);\n api.on(\"select\", onSelect);\n\n return () => {\n api?.off(\"select\", onSelect);\n };\n }, [api, onSelect]);\n\n return (\n \n \n {children}\n
    \n \n );\n}\n\nfunction CarouselContent({ className, ...props }: React.ComponentProps<\"div\">) {\n const { carouselRef, orientation } = useCarousel();\n\n return (\n \n \n \n );\n}\n\nfunction CarouselItem({ className, ...props }: React.ComponentProps<\"div\">) {\n const { orientation } = useCarousel();\n\n return (\n \n );\n}\n\nfunction CarouselPrevious({\n className,\n variant = \"outline\",\n size = \"icon-sm\",\n ...props\n}: React.ComponentProps) {\n const { orientation, scrollPrev, canScrollPrev } = useCarousel();\n\n return (\n \n \n Previous slide\n \n );\n}\n\nfunction CarouselNext({\n className,\n variant = \"outline\",\n size = \"icon-sm\",\n ...props\n}: React.ComponentProps) {\n const { orientation, scrollNext, canScrollNext } = useCarousel();\n\n return (\n \n \n Next slide\n \n );\n}\n\nexport {\n type CarouselApi,\n Carousel,\n CarouselContent,\n CarouselItem,\n CarouselPrevious,\n CarouselNext,\n useCarousel,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/chart.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport * as RechartsPrimitive from \"recharts\";\n\nimport { cn } from \"@/lib/utils\";\n\n// Format: { THEME_NAME: CSS_SELECTOR }\nconst THEMES = { light: \"\", dark: \".dark\" } as const;\n\nexport type ChartConfig = {\n [k in string]: {\n label?: React.ReactNode;\n icon?: React.ComponentType;\n } & (\n | { color?: string; theme?: never }\n | { color?: never; theme: Record }\n );\n};\n\ntype ChartContextProps = {\n config: ChartConfig;\n};\n\nconst ChartContext = React.createContext(null);\n\nfunction useChart() {\n const context = React.useContext(ChartContext);\n\n if (!context) {\n throw new Error(\"useChart must be used within a \");\n }\n\n return context;\n}\n\nfunction ChartContainer({\n id,\n className,\n children,\n config,\n ...props\n}: React.ComponentProps<\"div\"> & {\n config: ChartConfig;\n children: React.ComponentProps<\n typeof RechartsPrimitive.ResponsiveContainer\n >[\"children\"];\n}) {\n const uniqueId = React.useId();\n const chartId = `chart-${id || uniqueId.replace(/:/g, \"\")}`;\n\n return (\n \n \n \n \n {children}\n \n \n \n );\n}\n\nconst ChartStyle = ({ id, config }: { id: string; config: ChartConfig }) => {\n const colorConfig = Object.entries(config).filter(\n ([, config]) => config.theme || config.color,\n );\n\n if (!colorConfig.length) {\n return null;\n }\n\n return (\n `\n${prefix} [data-chart=${id}] {\n${colorConfig\n .map(([key, itemConfig]) => {\n const color =\n itemConfig.theme?.[theme as keyof typeof itemConfig.theme] ||\n itemConfig.color;\n return color ? ` --color-${key}: ${color};` : null;\n })\n .join(\"\\n\")}\n}\n`,\n )\n .join(\"\\n\"),\n }}\n />\n );\n};\n\nconst ChartTooltip = RechartsPrimitive.Tooltip;\n\nfunction ChartTooltipContent({\n active,\n payload,\n className,\n indicator = \"dot\",\n hideLabel = false,\n hideIndicator = false,\n label,\n labelFormatter,\n labelClassName,\n formatter,\n color,\n nameKey,\n labelKey,\n}: React.ComponentProps &\n React.ComponentProps<\"div\"> & {\n hideLabel?: boolean;\n hideIndicator?: boolean;\n indicator?: \"line\" | \"dot\" | \"dashed\";\n nameKey?: string;\n labelKey?: string;\n }) {\n const { config } = useChart();\n\n const tooltipLabel = React.useMemo(() => {\n if (hideLabel || !payload?.length) {\n return null;\n }\n\n const [item] = payload;\n const key = `${labelKey || item?.dataKey || item?.name || \"value\"}`;\n const itemConfig = getPayloadConfigFromPayload(config, item, key);\n const value =\n !labelKey && typeof label === \"string\"\n ? config[label as keyof typeof config]?.label || label\n : itemConfig?.label;\n\n if (labelFormatter) {\n return (\n
    \n {labelFormatter(value, payload)}\n
    \n );\n }\n\n if (!value) {\n return null;\n }\n\n return
    {value}
    ;\n }, [\n label,\n labelFormatter,\n payload,\n hideLabel,\n labelClassName,\n config,\n labelKey,\n ]);\n\n if (!active || !payload?.length) {\n return null;\n }\n\n const nestLabel = payload.length === 1 && indicator !== \"dot\";\n\n return (\n \n {!nestLabel ? tooltipLabel : null}\n
    \n {payload\n .filter((item) => item.type !== \"none\")\n .map((item, index) => {\n const key = `${nameKey || item.name || item.dataKey || \"value\"}`;\n const itemConfig = getPayloadConfigFromPayload(config, item, key);\n const indicatorColor = color || item.payload.fill || item.color;\n\n return (\n svg]:text-muted-foreground flex w-full flex-wrap items-stretch gap-2 [&>svg]:h-2.5 [&>svg]:w-2.5\",\n indicator === \"dot\" && \"items-center\",\n )}\n >\n {formatter && item?.value !== undefined && item.name ? (\n formatter(item.value, item.name, item, index, item.payload)\n ) : (\n <>\n {itemConfig?.icon ? (\n \n ) : (\n !hideIndicator && (\n \n )\n )}\n \n
    \n {nestLabel ? tooltipLabel : null}\n \n {itemConfig?.label || item.name}\n \n
    \n {item.value && (\n \n {item.value.toLocaleString()}\n \n )}\n
    \n \n )}\n \n );\n })}\n \n \n );\n}\n\nconst ChartLegend = RechartsPrimitive.Legend;\n\nfunction ChartLegendContent({\n className,\n hideIcon = false,\n payload,\n verticalAlign = \"bottom\",\n nameKey,\n}: React.ComponentProps<\"div\"> &\n Pick & {\n hideIcon?: boolean;\n nameKey?: string;\n }) {\n const { config } = useChart();\n\n if (!payload?.length) {\n return null;\n }\n\n return (\n \n {payload\n .filter((item) => item.type !== \"none\")\n .map((item) => {\n const key = `${nameKey || item.dataKey || \"value\"}`;\n const itemConfig = getPayloadConfigFromPayload(config, item, key);\n\n return (\n svg]:text-muted-foreground flex items-center gap-1.5 [&>svg]:h-3 [&>svg]:w-3\",\n )}\n >\n {itemConfig?.icon && !hideIcon ? (\n \n ) : (\n \n )}\n {itemConfig?.label}\n \n );\n })}\n \n );\n}\n\nfunction getPayloadConfigFromPayload(\n config: ChartConfig,\n payload: unknown,\n key: string,\n) {\n if (typeof payload !== \"object\" || payload === null) {\n return undefined;\n }\n\n const payloadPayload =\n \"payload\" in payload &&\n typeof payload.payload === \"object\" &&\n payload.payload !== null\n ? payload.payload\n : undefined;\n\n let configLabelKey: string = key;\n\n if (\n key in payload &&\n typeof payload[key as keyof typeof payload] === \"string\"\n ) {\n configLabelKey = payload[key as keyof typeof payload] as string;\n } else if (\n payloadPayload &&\n key in payloadPayload &&\n typeof payloadPayload[key as keyof typeof payloadPayload] === \"string\"\n ) {\n configLabelKey = payloadPayload[\n key as keyof typeof payloadPayload\n ] as string;\n }\n\n return configLabelKey in config\n ? config[configLabelKey]\n : config[key as keyof typeof config];\n}\n\nexport {\n ChartContainer,\n ChartTooltip,\n ChartTooltipContent,\n ChartLegend,\n ChartLegendContent,\n ChartStyle,\n};\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/checkbox.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport { Checkbox as CheckboxPrimitive } from \"radix-ui\";\n\nimport { cn } from \"@/lib/utils\";\nimport { CheckIcon } from \"lucide-react\";\n\nfunction Checkbox({\n className,\n ...props\n}: React.ComponentProps) {\n return (\n \n svg]:size-3.5 grid place-content-center text-current transition-none\"\n >\n \n \n \n );\n}\n\nexport { Checkbox };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/collapsible.tsx", "content": "\"use client\";\n\nimport { Collapsible as CollapsiblePrimitive } from \"radix-ui\";\n\nfunction Collapsible({\n ...props\n}: React.ComponentProps) {\n return ;\n}\n\nfunction CollapsibleTrigger({\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nfunction CollapsibleContent({\n ...props\n}: React.ComponentProps) {\n return (\n \n );\n}\n\nexport { Collapsible, CollapsibleTrigger, CollapsibleContent };\n" }, { "path": "backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/outputs/web/components/ui/combobox.tsx", "content": "\"use client\";\n\nimport * as React from \"react\";\nimport { Combobox as ComboboxPrimitive } from \"@base-ui/react\";\n\nimport { cn } from \"@/lib/utils\";\nimport { Button } from \"@/components/ui/button\";\nimport {\n InputGroup,\n InputGroupAddon,\n InputGroupButton,\n InputGroupInput,\n} from \"@/components/ui/input-group\";\nimport { ChevronDownIcon, XIcon, CheckIcon } from \"lucide-react\";\n\nconst Combobox = ComboboxPrimitive.Root;\n\nfunction ComboboxValue({ ...props }: ComboboxPrimitive.Value.Props) {\n return ;\n}\n\nfunction ComboboxTrigger({\n className,\n children,\n ...props\n}: ComboboxPrimitive.Trigger.Props) {\n return (\n \n {children}\n \n \n );\n}\n\nfunction ComboboxClear({ className, ...props }: ComboboxPrimitive.Clear.Props) {\n return (\n }\n className={cn(className)}\n {...props}\n >\n \n \n );\n}\n\nfunction ComboboxInput({\n className,\n children,\n disabled = false,\n showTrigger = true,\n showClear = false,\n ...props\n}: ComboboxPrimitive.Input.Props & {\n showTrigger?: boolean;\n showClear?: boolean;\n}) {\n return (\n \n }\n {...props}\n />\n \n {showTrigger && (\n \n \n \n )}\n {showClear && }\n \n {children}\n \n );\n}\n\nfunction ComboboxContent({\n className,\n side = \"bottom\",\n sideOffset = 6,\n align = \"start\",\n alignOffset = 0,\n anchor,\n ...props\n}: ComboboxPrimitive.Popup.Props &\n Pick<\n ComboboxPrimitive.Positioner.Props,\n \"side\" | \"align\" | \"sideOffset\" | \"alignOffset\" | \"anchor\"\n >) {\n return (\n \n \n \n \n \n );\n}\n\nfunction ComboboxList({ className, ...props }: ComboboxPrimitive.List.Props) {\n return (\n \n );\n}\n\nfunction ComboboxItem({\n className,\n children,\n ...props\n}: ComboboxPrimitive.Item.Props) {\n return (\n \n {children}\n \n }\n >\n \n \n \n );\n}\n\nfunction ComboboxGroup({ className, ...props }: ComboboxPrimitive.Group.Props) {\n return (\n \n );\n}\n\nfunction ComboboxLabel({\n className,\n ...props\n}: ComboboxPrimitive.GroupLabel.Props) {\n return (\n \n );\n}\n\nfunction ComboboxCollection({ ...props }: ComboboxPrimitive.Collection.Props) {\n return (\n \n );\n}\n\nfunction ComboboxEmpty({ className, ...props }: ComboboxPrimitive.Empty.Props) {\n return (\n     <\\/a>/r /tmp/foss_notice.txt' README.md\nsed -i 's/utm_source=onyx_repo/utm_source=foss_repo/g' README.md\n\ngit add README.md\ngit commit -m \"README\"\n\necho \"=== Creating blob callback script ===\"\ncat > /tmp/license_replacer.py << 'PYEOF'\n#!/usr/bin/env python3\nimport sys\n\n# Read MIT license from file\nwith open('/tmp/mit_license.txt', 'rb') as f:\n MIT_LICENSE = f.read()\n\nimport git_filter_repo as fr\n\nreplaced_count = 0\n\ndef replace_license_blob_content(blob, metadata):\n \"\"\"Replace LICENSE blob content with MIT license based on content detection\"\"\"\n global replaced_count\n\n # Check if this blob looks like a license file\n # We'll replace any blob that contains the old Apache/custom license text\n if blob.data and len(blob.data) > 100:\n # Check for license-like content\n # Unfortunately, we don't have access to the path, so we can't just check that the path\n # is `LICENSE`.\n data_lower = blob.data.lower()\n if (\n b'portions of this software are licensed as follows' in data_lower and\n b'all third party components incorporated into the' in data_lower\n ):\n # Additional check: make sure it's actually a license file, not source code\n # License files typically don't have common code patterns\n if b'def ' not in blob.data and b'class ' not in blob.data and b'import ' not in blob.data[:200]:\n blob.data = MIT_LICENSE\n replaced_count += 1\n\nargs = fr.FilteringOptions.parse_args(['--force'], error_on_empty=False)\nfilter_obj = fr.RepoFilter(args, blob_callback=replace_license_blob_content)\nfilter_obj.run()\n\nprint(f\"Replaced {replaced_count} LICENSE blob(s)\", file=sys.stderr)\nPYEOF\n\necho \"=== Replacing LICENSE file in all commits ===\"\nchmod +x /tmp/license_replacer.py\n/tmp/license_replacer.py\n\necho \"=== Done building FOSS repo ===\"\n" }, { "path": "backend/scripts/onyx_openapi_schema.py", "content": "# export openapi schema without having to start the actual web server\n\n# helpful tips: https://github.com/fastapi/fastapi/issues/1173\n\nimport argparse\nimport json\nimport os\nimport subprocess\nimport sys\n\nfrom fastapi import FastAPI\nfrom fastapi.openapi.utils import get_openapi\n\nfrom onyx.main import app as app_fn\n\nOPENAPI_VERSION = \"3.1.0\"\n\n\ndef go(filename: str, tagged_for_docs: str | None = None) -> None:\n \"\"\"Generate OpenAPI schema.\n\n By default outputs tag-stripped schema (for client generation).\n If tagged_for_docs is provided, also outputs the original tagged version for docs.\n \"\"\"\n app: FastAPI = app_fn()\n app.openapi_version = OPENAPI_VERSION\n schema = get_openapi(\n title=app.title,\n version=app.version,\n openapi_version=app.openapi_version,\n description=app.description,\n routes=app.routes,\n )\n\n # Output tagged version for docs if requested\n if tagged_for_docs:\n with open(tagged_for_docs, \"w\") as f:\n json.dump(schema, f)\n print(f\"Wrote tagged OpenAPI schema to {tagged_for_docs}\")\n\n # Output stripped version (default) for client generation\n stripped = strip_tags_from_schema(schema)\n with open(filename, \"w\") as f:\n json.dump(stripped, f)\n print(f\"Wrote OpenAPI schema to {filename}.\")\n\n\ndef strip_tags_from_schema(schema: dict) -> dict:\n \"\"\"Strip tags from OpenAPI schema so openapi-generator puts all endpoints in DefaultApi.\"\"\"\n import copy\n\n schema = copy.deepcopy(schema)\n\n # Remove tags from all operations\n if \"paths\" in schema:\n for path_item in schema[\"paths\"].values():\n for operation in path_item.values():\n if isinstance(operation, dict) and \"tags\" in operation:\n del operation[\"tags\"]\n\n # Remove top-level tags definition\n if \"tags\" in schema:\n del schema[\"tags\"]\n\n return schema\n\n\ndef generate_client(openapi_json_path: str, strip_tags: bool = True) -> None:\n \"\"\"Generate Python client from OpenAPI schema using openapi-generator.\"\"\"\n import tempfile\n\n output_dir = os.path.join(os.path.dirname(openapi_json_path), \"onyx_openapi_client\")\n\n # Optionally strip tags so all endpoints go under DefaultApi\n schema_path = openapi_json_path\n if strip_tags:\n with open(openapi_json_path) as f:\n schema = json.load(f)\n stripped = strip_tags_from_schema(schema)\n fd, schema_path = tempfile.mkstemp(suffix=\".json\")\n with os.fdopen(fd, \"w\") as f:\n json.dump(stripped, f)\n print(f\"Stripped tags from schema, using temp file: {schema_path}\")\n\n cmd = [\n \"openapi-generator\",\n \"generate\",\n \"-i\",\n schema_path,\n \"-g\",\n \"python\",\n \"-o\",\n output_dir,\n \"--package-name\",\n \"onyx_openapi_client\",\n \"--skip-validate-spec\",\n \"--openapi-normalizer\",\n \"SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true\",\n ]\n\n print(\"Running openapi-generator...\")\n try:\n result = subprocess.run(cmd)\n if result.returncode == 0:\n print(f\"Generated Python client at {output_dir}\")\n else:\n print(\n \"Failed to generate Python client. See backend/tests/integration/README.md for setup instructions.\",\n file=sys.stderr,\n )\n finally:\n # Clean up temp file if we created one\n if strip_tags and schema_path != openapi_json_path:\n os.unlink(schema_path)\n\n\ndef main() -> None:\n parser = argparse.ArgumentParser(\n description=\"Export OpenAPI schema for Onyx API (does not require starting API server)\"\n )\n parser.add_argument(\n \"--filename\", \"-f\", help=\"Filename to write to\", default=\"openapi.json\"\n )\n parser.add_argument(\n \"--generate-python-client\",\n action=\"store_true\",\n help=\"Generate Python client schemas (needed for integration tests)\",\n )\n parser.add_argument(\n \"--tagged-for-docs\",\n help=\"Also output a tagged version for API docs (specify output path)\",\n )\n\n args = parser.parse_args()\n go(args.filename, tagged_for_docs=args.tagged_for_docs)\n\n if args.generate_python_client:\n # Schema is already stripped by go(), no need to strip again\n generate_client(args.filename, strip_tags=False)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/orphan_doc_cleanup_script.py", "content": "import concurrent.futures\nimport os\nimport sys\n\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\n\nfrom onyx.document_index.document_index_utils import get_multipass_config\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\n# makes it so `PYTHONPATH=.` is not required when running this script\nparent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\nsys.path.append(parent_dir)\n\nfrom onyx.context.search.models import IndexFilters # noqa: E402\nfrom onyx.document_index.interfaces import VespaChunkRequest # noqa: E402\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant # noqa: E402\nfrom onyx.db.document import delete_documents_complete__no_commit # noqa: E402\nfrom onyx.db.tag import delete_orphan_tags__no_commit # noqa: E402\nfrom onyx.db.search_settings import get_current_search_settings # noqa: E402\nfrom onyx.document_index.vespa.index import VespaIndex # noqa: E402\nfrom onyx.db.document import get_document # noqa: E402\n\nBATCH_SIZE = 100\n\n\ndef _get_orphaned_document_ids(db_session: Session, limit: int) -> list[str]:\n \"\"\"Get document IDs that don't have any entries in document_by_connector_credential_pair\"\"\"\n query = text(\n \"\"\"\n SELECT d.id\n FROM document d\n LEFT JOIN document_by_connector_credential_pair dbcc ON d.id = dbcc.id\n WHERE dbcc.id IS NULL\n LIMIT :limit\n \"\"\"\n )\n orphaned_ids = [doc_id[0] for doc_id in db_session.execute(query, {\"limit\": limit})]\n print(f\"Found {len(orphaned_ids)} orphaned documents in this batch\")\n return orphaned_ids\n\n\ndef main() -> None:\n with get_session_with_current_tenant() as db_session:\n total_processed = 0\n while True:\n # Get orphaned document IDs in batches\n orphaned_ids = _get_orphaned_document_ids(db_session, BATCH_SIZE)\n if not orphaned_ids:\n if total_processed == 0:\n print(\"No orphaned documents found\")\n else:\n print(\n f\"Finished processing all batches. Total documents processed: {total_processed}\"\n )\n return\n\n # Setup Vespa index\n search_settings = get_current_search_settings(db_session)\n multipass_config = get_multipass_config(search_settings)\n index_name = search_settings.index_name\n vespa_index = VespaIndex(\n index_name=index_name,\n secondary_index_name=None,\n large_chunks_enabled=multipass_config.enable_large_chunks,\n secondary_large_chunks_enabled=None,\n )\n\n # Delete chunks from Vespa first\n print(\"Deleting orphaned document chunks from Vespa\")\n successfully_vespa_deleted_doc_ids: list[str] = []\n # Process documents in parallel using ThreadPoolExecutor\n with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:\n\n def process_doc(doc_id: str) -> str | None:\n document = get_document(doc_id, db_session)\n if not document:\n return None\n # Check if document exists in Vespa first\n try:\n chunks = vespa_index.id_based_retrieval(\n chunk_requests=[\n VespaChunkRequest(document_id=doc_id, max_chunk_ind=2)\n ],\n filters=IndexFilters(access_control_list=None),\n batch_retrieval=True,\n )\n if not chunks:\n print(f\"Document {doc_id} not found in Vespa\")\n return doc_id\n except Exception as e:\n print(\n f\"Error checking if document {doc_id} exists in Vespa: {e}\"\n )\n return None\n\n try:\n print(f\"Deleting document {doc_id} in Vespa\")\n chunks_deleted = vespa_index.delete_single(\n doc_id,\n tenant_id=POSTGRES_DEFAULT_SCHEMA,\n chunk_count=document.chunk_count,\n )\n if chunks_deleted > 0:\n print(\n f\"Deleted {chunks_deleted} chunks for document {doc_id}\"\n )\n return doc_id\n except Exception as e:\n print(\n f\"Error deleting document {doc_id} in Vespa and will not delete from Postgres: {e}\"\n )\n return None\n\n # Submit all tasks and gather results\n futures = [\n executor.submit(process_doc, doc_id) for doc_id in orphaned_ids\n ]\n for future in concurrent.futures.as_completed(futures):\n doc_id = future.result()\n if doc_id:\n successfully_vespa_deleted_doc_ids.append(doc_id)\n\n # Delete documents from Postgres\n print(\"Deleting orphaned documents from Postgres\")\n try:\n delete_documents_complete__no_commit(\n db_session, successfully_vespa_deleted_doc_ids\n )\n delete_orphan_tags__no_commit(db_session)\n db_session.commit()\n except Exception as e:\n print(f\"Error deleting documents from Postgres: {e}\")\n break\n\n total_processed += len(successfully_vespa_deleted_doc_ids)\n print(\n f\"Successfully cleaned up {len(successfully_vespa_deleted_doc_ids)} orphaned documents in this batch\"\n )\n print(f\"Total documents processed so far: {total_processed}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/query_time_check/seed_dummy_docs.py", "content": "\"\"\"\nlaunch:\n- api server\n- postgres\n- vespa\n- model server (this is only needed so the api server can startup, no embedding is done)\n\nRun this script to seed the database with dummy documents.\nThen run test_query_times.py to test query times.\n\"\"\"\n\nimport random\nfrom datetime import datetime\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.document_index.document_index_utils import get_multipass_config\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.indexing.indexing_pipeline import IndexBatchParams\nfrom onyx.indexing.models import ChunkEmbedding\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom onyx.indexing.models import IndexChunk\nfrom onyx.utils.timing import log_function_time\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\nfrom shared_configs.model_server_models import Embedding\n\nTOTAL_DOC_SETS = 8\nTOTAL_ACL_ENTRIES_PER_CATEGORY = 80\n\n\ndef generate_random_embedding(dim: int) -> Embedding:\n return [random.uniform(-1, 1) for _ in range(dim)]\n\n\ndef generate_random_identifier() -> str:\n return f\"dummy_doc_{random.randint(1, 1000)}\"\n\n\ndef generate_dummy_chunk(\n doc_id: str,\n chunk_id: int,\n embedding_dim: int,\n number_of_acl_entries: int,\n number_of_document_sets: int,\n) -> DocMetadataAwareIndexChunk:\n document = Document(\n id=doc_id,\n source=DocumentSource.GOOGLE_DRIVE,\n sections=[],\n metadata={},\n semantic_identifier=generate_random_identifier(),\n )\n\n chunk = IndexChunk(\n chunk_id=chunk_id,\n blurb=f\"Blurb for chunk {chunk_id} of document {doc_id}.\",\n content=f\"Content for chunk {chunk_id} of document {doc_id}. This is dummy text for testing purposes.\",\n source_links={},\n section_continuation=False,\n source_document=document,\n title_prefix=f\"Title prefix for doc {doc_id}\",\n metadata_suffix_semantic=\"\",\n metadata_suffix_keyword=\"\",\n doc_summary=\"\",\n chunk_context=\"\",\n mini_chunk_texts=None,\n contextual_rag_reserved_tokens=0,\n embeddings=ChunkEmbedding(\n full_embedding=generate_random_embedding(embedding_dim),\n mini_chunk_embeddings=[],\n ),\n title_embedding=generate_random_embedding(embedding_dim),\n large_chunk_id=None,\n large_chunk_reference_ids=[],\n image_file_id=None,\n )\n\n document_set_names = []\n for i in range(number_of_document_sets):\n document_set_names.append(f\"Document Set {i}\")\n\n user_emails: list[str | None] = []\n user_groups: list[str] = []\n external_user_emails: list[str] = []\n external_user_group_ids: list[str] = []\n for i in range(number_of_acl_entries):\n user_emails.append(f\"user_{i}@example.com\")\n user_groups.append(f\"group_{i}\")\n external_user_emails.append(f\"external_user_{i}@example.com\")\n external_user_group_ids.append(f\"external_group_{i}\")\n\n return DocMetadataAwareIndexChunk.from_index_chunk(\n index_chunk=chunk,\n user_project=[],\n personas=[],\n access=DocumentAccess.build(\n user_emails=user_emails,\n user_groups=user_groups,\n external_user_emails=external_user_emails,\n external_user_group_ids=external_user_group_ids,\n is_public=random.choice([True, False]),\n ),\n document_sets={document_set for document_set in document_set_names},\n boost=random.randint(-1, 1),\n aggregated_chunk_boost_factor=random.random(),\n tenant_id=POSTGRES_DEFAULT_SCHEMA,\n )\n\n\n@log_function_time()\ndef do_insertion(\n vespa_index: VespaIndex, all_chunks: list[DocMetadataAwareIndexChunk]\n) -> None:\n insertion_records = vespa_index.index(\n chunks=all_chunks,\n index_batch_params=IndexBatchParams(\n doc_id_to_previous_chunk_cnt={},\n doc_id_to_new_chunk_cnt={},\n tenant_id=POSTGRES_DEFAULT_SCHEMA,\n large_chunks_enabled=False,\n ),\n )\n print(f\"Indexed {len(insertion_records)} documents.\")\n print(\n f\"New documents: {sum(1 for record in insertion_records if not record.already_existed)}\"\n )\n print(\n f\"Existing documents updated: {sum(1 for record in insertion_records if record.already_existed)}\"\n )\n\n\n@log_function_time()\ndef seed_dummy_docs(\n number_of_document_sets: int,\n number_of_acl_entries: int,\n num_docs: int = 1000,\n chunks_per_doc: int = 5,\n batch_size: int = 100,\n) -> None:\n with get_session_with_current_tenant() as db_session:\n search_settings = get_current_search_settings(db_session)\n multipass_config = get_multipass_config(search_settings)\n index_name = search_settings.index_name\n embedding_dim = search_settings.final_embedding_dim\n\n vespa_index = VespaIndex(\n index_name=index_name,\n secondary_index_name=None,\n large_chunks_enabled=multipass_config.enable_large_chunks,\n secondary_large_chunks_enabled=None,\n )\n print(index_name)\n\n all_chunks = []\n chunk_count = 0\n for doc_num in range(num_docs):\n doc_id = f\"dummy_doc_{doc_num}_{datetime.now().isoformat()}\"\n for chunk_num in range(chunks_per_doc):\n chunk = generate_dummy_chunk(\n doc_id=doc_id,\n chunk_id=chunk_num,\n embedding_dim=embedding_dim,\n number_of_acl_entries=number_of_acl_entries,\n number_of_document_sets=number_of_document_sets,\n )\n all_chunks.append(chunk)\n chunk_count += 1\n\n if len(all_chunks) >= chunks_per_doc * batch_size:\n do_insertion(vespa_index, all_chunks)\n print(\n f\"Indexed {chunk_count} chunks out of {num_docs * chunks_per_doc}.\"\n )\n print(\n f\"percentage: {chunk_count / (num_docs * chunks_per_doc) * 100:.2f}% \\n\"\n )\n all_chunks = []\n\n if all_chunks:\n do_insertion(vespa_index, all_chunks)\n\n\nif __name__ == \"__main__\":\n seed_dummy_docs(\n number_of_document_sets=TOTAL_DOC_SETS,\n number_of_acl_entries=TOTAL_ACL_ENTRIES_PER_CATEGORY,\n num_docs=100000,\n chunks_per_doc=5,\n batch_size=1000,\n )\n" }, { "path": "backend/scripts/query_time_check/test_query_times.py", "content": "# \"\"\"\n# RUN THIS AFTER SEED_DUMMY_DOCS.PY\n# \"\"\"\n\n# import random\n# import time\n\n# from onyx.agents.agent_search.shared_graph_utils.models import QueryExpansionType\n# from onyx.configs.constants import DocumentSource\n# from onyx.configs.model_configs import DOC_EMBEDDING_DIM\n# from onyx.context.search.models import IndexFilters\n# from onyx.db.engine.sql_engine import get_session_with_current_tenant\n# from onyx.db.search_settings import get_current_search_settings\n# from onyx.document_index.document_index_utils import get_multipass_config\n# from onyx.document_index.vespa.index import VespaIndex\n# from scripts.query_time_check.seed_dummy_docs import TOTAL_ACL_ENTRIES_PER_CATEGORY\n# from scripts.query_time_check.seed_dummy_docs import TOTAL_DOC_SETS\n# from shared_configs.model_server_models import Embedding\n\n# # make sure these are smaller than TOTAL_ACL_ENTRIES_PER_CATEGORY and TOTAL_DOC_SETS, respectively\n# NUMBER_OF_ACL_ENTRIES_PER_QUERY = 6\n# NUMBER_OF_DOC_SETS_PER_QUERY = 2\n\n\n# def get_slowest_99th_percentile(results: list[float]) -> float:\n# return sorted(results)[int(0.99 * len(results))]\n\n\n# # Generate random filters\n# def _random_filters() -> IndexFilters:\n# \"\"\"\n# Generate random filters for the query containing:\n# - NUMBER_OF_ACL_ENTRIES_PER_QUERY user emails\n# - NUMBER_OF_ACL_ENTRIES_PER_QUERY groups\n# - NUMBER_OF_ACL_ENTRIES_PER_QUERY external groups\n# - NUMBER_OF_DOC_SETS_PER_QUERY document sets\n# \"\"\"\n# access_control_list = [\n# f\"user_email:user_{random.randint(0, TOTAL_ACL_ENTRIES_PER_CATEGORY - 1)}@example.com\",\n# ]\n# acl_indices = random.sample(\n# range(TOTAL_ACL_ENTRIES_PER_CATEGORY), NUMBER_OF_ACL_ENTRIES_PER_QUERY\n# )\n# for i in acl_indices:\n# access_control_list.append(f\"group:group_{acl_indices[i]}\")\n# access_control_list.append(f\"external_group:external_group_{acl_indices[i]}\")\n\n# doc_sets = []\n# doc_set_indices = random.sample(\n# range(TOTAL_DOC_SETS), NUMBER_OF_ACL_ENTRIES_PER_QUERY\n# )\n# for i in doc_set_indices:\n# doc_sets.append(f\"document_set:Document Set {doc_set_indices[i]}\")\n\n# return IndexFilters(\n# source_type=[DocumentSource.GOOGLE_DRIVE],\n# document_set=doc_sets,\n# tags=[],\n# access_control_list=access_control_list,\n# )\n\n\n# def test_hybrid_retrieval_times(\n# number_of_queries: int,\n# ) -> None:\n# with get_session_with_current_tenant() as db_session:\n# search_settings = get_current_search_settings(db_session)\n# multipass_config = get_multipass_config(search_settings)\n# index_name = search_settings.index_name\n\n# vespa_index = VespaIndex(\n# index_name=index_name,\n# secondary_index_name=None,\n# large_chunks_enabled=multipass_config.enable_large_chunks,\n# secondary_large_chunks_enabled=None,\n# )\n\n# # Generate random queries\n# queries = [f\"Random Query {i}\" for i in range(number_of_queries)]\n\n# # Generate random embeddings\n# embeddings = [\n# Embedding([random.random() for _ in range(DOC_EMBEDDING_DIM)])\n# for _ in range(number_of_queries)\n# ]\n\n# total_time = 0.0\n# results = []\n# for i in range(number_of_queries):\n# start_time = time.time()\n\n# vespa_index.hybrid_retrieval(\n# query=queries[i],\n# query_embedding=embeddings[i],\n# final_keywords=None,\n# filters=_random_filters(),\n# hybrid_alpha=0.5,\n# time_decay_multiplier=1.0,\n# num_to_retrieve=50,\n# ranking_profile_type=QueryExpansionType.SEMANTIC,\n# offset=0,\n# title_content_ratio=0.5,\n# )\n\n# end_time = time.time()\n# query_time = end_time - start_time\n# total_time += query_time\n# results.append(query_time)\n\n# print(f\"Query {i+1}: {query_time:.4f} seconds\")\n\n# avg_time = total_time / number_of_queries\n# fast_time = min(results)\n# slow_time = max(results)\n# ninety_ninth_percentile = get_slowest_99th_percentile(results)\n# # Write results to a file\n# _OUTPUT_PATH = \"query_times_results_large_more.txt\"\n# with open(_OUTPUT_PATH, \"w\") as f:\n# f.write(f\"Average query time: {avg_time:.4f} seconds\\n\")\n# f.write(f\"Fastest query: {fast_time:.4f} seconds\\n\")\n# f.write(f\"Slowest query: {slow_time:.4f} seconds\\n\")\n# f.write(f\"99th percentile: {ninety_ninth_percentile:.4f} seconds\\n\")\n# print(f\"Results written to {_OUTPUT_PATH}\")\n\n# print(f\"\\nAverage query time: {avg_time:.4f} seconds\")\n# print(f\"Fastest query: {fast_time:.4f} seconds\")\n# print(f\"Slowest query: {max(results):.4f} seconds\")\n# print(f\"99th percentile: {get_slowest_99th_percentile(results):.4f} seconds\")\n\n\n# if __name__ == \"__main__\":\n# test_hybrid_retrieval_times(number_of_queries=1000)\n" }, { "path": "backend/scripts/reencrypt_secrets.py", "content": "\"\"\"Re-encrypt secrets under the current ENCRYPTION_KEY_SECRET.\n\nDecrypts all encrypted columns using the old key (or raw decode if the old key\nis empty), then re-encrypts them with the current ENCRYPTION_KEY_SECRET.\n\nUsage (docker):\n docker exec -it onyx-api_server-1 \\\n python -m scripts.reencrypt_secrets --old-key \"previous-key\"\n\nUsage (kubernetes):\n kubectl exec -it -- \\\n python -m scripts.reencrypt_secrets --old-key \"previous-key\"\n\nOmit --old-key (or pass \"\") if secrets were not previously encrypted.\n\nFor multi-tenant deployments, pass --tenant-id to target a specific tenant,\nor --all-tenants to iterate every tenant.\n\"\"\"\n\nimport argparse\nimport os\nimport sys\n\nparent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\nsys.path.append(parent_dir)\n\nfrom onyx.db.rotate_encryption_key import rotate_encryption_key # noqa: E402\nfrom onyx.db.engine.sql_engine import get_session_with_tenant # noqa: E402\nfrom onyx.db.engine.sql_engine import SqlEngine # noqa: E402\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids # noqa: E402\nfrom onyx.utils.variable_functionality import global_version # noqa: E402\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA # noqa: E402\n\n\ndef _run_for_tenant(tenant_id: str, old_key: str | None, dry_run: bool = False) -> None:\n print(f\"Re-encrypting secrets for tenant: {tenant_id}\")\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n results = rotate_encryption_key(db_session, old_key=old_key, dry_run=dry_run)\n\n if results:\n for col, count in results.items():\n print(\n f\" {col}: {count} row(s) {'would be ' if dry_run else ''}re-encrypted\"\n )\n else:\n print(\"No rows needed re-encryption.\")\n\n\ndef main() -> None:\n parser = argparse.ArgumentParser(\n description=\"Re-encrypt secrets under the current encryption key.\"\n )\n parser.add_argument(\n \"--old-key\",\n default=None,\n help=\"Previous encryption key. Omit or pass empty string if not applicable.\",\n )\n parser.add_argument(\n \"--dry-run\",\n action=\"store_true\",\n help=\"Show what would be re-encrypted without making changes.\",\n )\n\n tenant_group = parser.add_mutually_exclusive_group()\n tenant_group.add_argument(\n \"--tenant-id\",\n default=None,\n help=\"Target a specific tenant schema.\",\n )\n tenant_group.add_argument(\n \"--all-tenants\",\n action=\"store_true\",\n help=\"Iterate all tenants.\",\n )\n\n args = parser.parse_args()\n\n old_key = args.old_key if args.old_key else None\n\n global_version.set_ee()\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n if args.dry_run:\n print(\"DRY RUN — no changes will be made\")\n\n if args.all_tenants:\n tenant_ids = get_all_tenant_ids()\n print(f\"Found {len(tenant_ids)} tenant(s)\")\n failed_tenants: list[str] = []\n for tid in tenant_ids:\n try:\n _run_for_tenant(tid, old_key, dry_run=args.dry_run)\n except Exception as e:\n print(f\" ERROR for tenant {tid}: {e}\")\n failed_tenants.append(tid)\n if failed_tenants:\n print(f\"FAILED tenants ({len(failed_tenants)}): {failed_tenants}\")\n sys.exit(1)\n else:\n tenant_id = args.tenant_id or POSTGRES_DEFAULT_SCHEMA\n _run_for_tenant(tenant_id, old_key, dry_run=args.dry_run)\n\n print(\"Done.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/reset_indexes.py", "content": "# This file is purely for development use, not included in any builds\nimport os\nimport sys\nfrom time import sleep\n\nimport requests\nfrom requests.exceptions import RequestException\n\n# makes it so `PYTHONPATH=.` is not required when running this script\nparent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\nsys.path.append(parent_dir)\n\nfrom onyx.configs.app_configs import DOCUMENT_INDEX_NAME # noqa: E402\nfrom onyx.document_index.vespa.index import DOCUMENT_ID_ENDPOINT # noqa: E402\nfrom onyx.utils.logger import setup_logger # noqa: E402\n\nlogger = setup_logger()\n\n\ndef wipe_vespa_index() -> bool:\n \"\"\"\n Wipes the Vespa index by deleting all documents.\n \"\"\"\n continuation = None\n should_continue = True\n RETRIES = 3\n\n while should_continue:\n params = {\"selection\": \"true\", \"cluster\": DOCUMENT_INDEX_NAME}\n if continuation:\n params[\"continuation\"] = continuation\n\n for attempt in range(RETRIES):\n try:\n response = requests.delete(DOCUMENT_ID_ENDPOINT, params=params)\n response.raise_for_status()\n\n response_json = response.json()\n logger.info(f\"Response: {response_json}\")\n\n continuation = response_json.get(\"continuation\")\n should_continue = bool(continuation)\n break # Exit the retry loop if the request is successful\n\n except RequestException:\n logger.exception(\"Request failed\")\n sleep(2**attempt) # Exponential backoff\n else:\n logger.error(f\"Max retries ({RETRIES}) exceeded. Exiting.\")\n return False\n\n return True\n\n\ndef main() -> int:\n \"\"\"\n Main function to execute the script.\n \"\"\"\n try:\n succeeded = wipe_vespa_index()\n except Exception:\n logger.exception(\"wipe_vespa_index exceptioned.\")\n return 1\n\n if not succeeded:\n logger.info(\"Vespa index wipe failed.\")\n return 0\n\n logger.info(\"Vespa index wiped successfully.\")\n return 1\n\n\nif __name__ == \"__main__\":\n sys.exit(main())\n" }, { "path": "backend/scripts/reset_postgres.py", "content": "import os\nimport sys\n\nimport psycopg2\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_sqlalchemy_engine\n\n# makes it so `PYTHONPATH=.` is not required when running this script\nparent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\nsys.path.append(parent_dir)\n\nfrom onyx.configs.app_configs import POSTGRES_DB # noqa: E402\nfrom onyx.configs.app_configs import POSTGRES_HOST # noqa: E402\nfrom onyx.configs.app_configs import POSTGRES_PASSWORD # noqa: E402\nfrom onyx.configs.app_configs import POSTGRES_PORT # noqa: E402\nfrom onyx.configs.app_configs import POSTGRES_USER # noqa: E402\nfrom onyx.db.credentials import create_initial_public_credential # noqa: E402\n\n\ndef wipe_all_rows(database: str) -> None:\n conn = psycopg2.connect(\n dbname=database,\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n )\n cur = conn.cursor()\n\n # Disable triggers to prevent foreign key constraints from being checked\n cur.execute(\"SET session_replication_role = 'replica';\")\n\n # Fetch all table names in the current database\n cur.execute(\n \"\"\"\n SELECT tablename\n FROM pg_tables\n WHERE schemaname = 'public'\n \"\"\"\n )\n\n tables = cur.fetchall()\n\n for table in tables:\n table_name = table[0]\n\n # Don't touch migration history\n if table_name == \"alembic_version\":\n continue\n\n print(f\"Deleting all rows from {table_name}...\")\n cur.execute(f'DELETE FROM \"{table_name}\"')\n\n # Re-enable triggers\n cur.execute(\"SET session_replication_role = 'origin';\")\n\n conn.commit()\n cur.close()\n conn.close()\n print(\"Finished wiping all rows.\")\n\n\nif __name__ == \"__main__\":\n print(\"Cleaning up all Onyx tables\")\n wipe_all_rows(POSTGRES_DB)\n with Session(get_sqlalchemy_engine(), expire_on_commit=False) as db_session:\n create_initial_public_credential(db_session)\n print(\"To keep data consistent, it's best to wipe the document index as well.\")\n print(\n \"To be safe, it's best to restart the Onyx services (API Server and Background Tasks\"\n )\n" }, { "path": "backend/scripts/restart_containers.sh", "content": "#!/bin/bash\nset -e\n\nSCRIPT_DIR=\"$( cd \"$( dirname \"${BASH_SOURCE[0]}\" )\" &> /dev/null && pwd )\"\nCOMPOSE_FILE=\"$SCRIPT_DIR/../../deployment/docker_compose/docker-compose.yml\"\nCOMPOSE_DEV_FILE=\"$SCRIPT_DIR/../../deployment/docker_compose/docker-compose.dev.yml\"\n\nstop_and_remove_containers() {\n docker stop onyx_postgres onyx_vespa onyx_redis onyx_minio onyx_code_interpreter 2>/dev/null || true\n docker rm onyx_postgres onyx_vespa onyx_redis onyx_minio onyx_code_interpreter 2>/dev/null || true\n docker compose -f \"$COMPOSE_FILE\" -f \"$COMPOSE_DEV_FILE\" --profile opensearch-enabled stop opensearch 2>/dev/null || true\n docker compose -f \"$COMPOSE_FILE\" -f \"$COMPOSE_DEV_FILE\" --profile opensearch-enabled rm -f opensearch 2>/dev/null || true\n}\n\ncleanup() {\n echo \"Error occurred. Cleaning up...\"\n stop_and_remove_containers\n}\n\n# Trap errors and output a message, then cleanup\ntrap 'echo \"Error occurred on line $LINENO. Exiting script.\" >&2; cleanup' ERR\n\n# Usage of the script with optional volume arguments\n# ./restart_containers.sh [vespa_volume] [postgres_volume] [redis_volume]\n# [minio_volume] [--keep-opensearch-data]\n\nKEEP_OPENSEARCH_DATA=false\nPOSITIONAL_ARGS=()\nfor arg in \"$@\"; do\n if [[ \"$arg\" == \"--keep-opensearch-data\" ]]; then\n KEEP_OPENSEARCH_DATA=true\n else\n POSITIONAL_ARGS+=(\"$arg\")\n fi\ndone\n\nVESPA_VOLUME=${POSITIONAL_ARGS[0]:-\"\"}\nPOSTGRES_VOLUME=${POSITIONAL_ARGS[1]:-\"\"}\nREDIS_VOLUME=${POSITIONAL_ARGS[2]:-\"\"}\nMINIO_VOLUME=${POSITIONAL_ARGS[3]:-\"\"}\n\n# Stop and remove the existing containers\necho \"Stopping and removing existing containers...\"\nstop_and_remove_containers\n\n# Start the PostgreSQL container with optional volume\necho \"Starting PostgreSQL container...\"\nif [[ -n \"$POSTGRES_VOLUME\" ]]; then\n docker run -p 5432:5432 --name onyx_postgres -e POSTGRES_PASSWORD=password -d -v $POSTGRES_VOLUME:/var/lib/postgresql/data postgres -c max_connections=250\nelse\n docker run -p 5432:5432 --name onyx_postgres -e POSTGRES_PASSWORD=password -d postgres -c max_connections=250\nfi\n\n# Start the Vespa container with optional volume\necho \"Starting Vespa container...\"\nif [[ -n \"$VESPA_VOLUME\" ]]; then\n docker run --detach --name onyx_vespa --hostname vespa-container --publish 8081:8081 --publish 19071:19071 -v $VESPA_VOLUME:/opt/vespa/var vespaengine/vespa:8\nelse\n docker run --detach --name onyx_vespa --hostname vespa-container --publish 8081:8081 --publish 19071:19071 vespaengine/vespa:8\nfi\n\n# If OPENSEARCH_ADMIN_PASSWORD is not already set, try loading it from\n# .vscode/.env so existing dev setups that stored it there aren't silently\n# broken.\nVSCODE_ENV=\"$SCRIPT_DIR/../../.vscode/.env\"\nif [[ -z \"${OPENSEARCH_ADMIN_PASSWORD:-}\" && -f \"$VSCODE_ENV\" ]]; then\n set -a\n # shellcheck source=/dev/null\n source \"$VSCODE_ENV\"\n set +a\nfi\n\n# Start the OpenSearch container using the same service from docker-compose that\n# our users use, setting OPENSEARCH_INITIAL_ADMIN_PASSWORD from the env's\n# OPENSEARCH_ADMIN_PASSWORD if it exists, else defaulting to StrongPassword123!.\n# Pass --keep-opensearch-data to preserve the opensearch-data volume across\n# restarts, else the volume is deleted so the container starts fresh.\nif [[ \"$KEEP_OPENSEARCH_DATA\" == \"false\" ]]; then\n echo \"Deleting opensearch-data volume...\"\n docker volume rm onyx_opensearch-data 2>/dev/null || true\nfi\necho \"Starting OpenSearch container...\"\ndocker compose -f \"$COMPOSE_FILE\" -f \"$COMPOSE_DEV_FILE\" --profile opensearch-enabled up --force-recreate -d opensearch\n\n# Start the Redis container with optional volume\necho \"Starting Redis container...\"\nif [[ -n \"$REDIS_VOLUME\" ]]; then\n docker run --detach --name onyx_redis --publish 6379:6379 -v $REDIS_VOLUME:/data redis\nelse\n docker run --detach --name onyx_redis --publish 6379:6379 redis\nfi\n\n# Start the MinIO container with optional volume\necho \"Starting MinIO container...\"\nif [[ -n \"$MINIO_VOLUME\" ]]; then\n docker run --detach --name onyx_minio --publish 9004:9000 --publish 9005:9001 -e MINIO_ROOT_USER=minioadmin -e MINIO_ROOT_PASSWORD=minioadmin -v $MINIO_VOLUME:/data minio/minio server /data --console-address \":9001\"\nelse\n docker run --detach --name onyx_minio --publish 9004:9000 --publish 9005:9001 -e MINIO_ROOT_USER=minioadmin -e MINIO_ROOT_PASSWORD=minioadmin minio/minio server /data --console-address \":9001\"\nfi\n\n# Start the Code Interpreter container\necho \"Starting Code Interpreter container...\"\ndocker run --detach --name onyx_code_interpreter --publish 8000:8000 --user root -v /var/run/docker.sock:/var/run/docker.sock onyxdotapp/code-interpreter:latest bash ./entrypoint.sh code-interpreter-api\n\n# Ensure alembic runs in the correct directory (backend/)\nPARENT_DIR=\"$(dirname \"$SCRIPT_DIR\")\"\ncd \"$PARENT_DIR\"\n\n# Give Postgres a second to start\nsleep 1\n\n# Alembic should be configured in the virtualenv for this repo\nif [[ -f \"../.venv/bin/activate\" ]]; then\n source ../.venv/bin/activate\nelse\n echo \"Warning: Python virtual environment not found at .venv/bin/activate; alembic may not work.\"\nfi\n\n# Run Alembic upgrade\necho \"Running Alembic migration...\"\nalembic upgrade head\n\n# Run the following instead of the above if using MT cloud\n# alembic -n schema_private upgrade head\n\necho \"Containers restarted and migration completed.\"\n" }, { "path": "backend/scripts/resume_paused_connectors.py", "content": "import argparse\n\nimport requests\n\nAPI_SERVER_URL = \"http://localhost:3000\"\nAPI_KEY = \"onyx-api-key\" # API key here, if auth is enabled\n\n\ndef resume_paused_connectors(\n api_server_url: str,\n api_key: str | None,\n specific_connector_sources: list[str] | None = None,\n) -> None:\n headers = {\"Content-Type\": \"application/json\"}\n if api_key:\n headers[\"Authorization\"] = f\"Bearer {api_key}\"\n\n # Get all paused connectors\n response = requests.post(\n f\"{api_server_url}/api/manage/admin/connector/indexing-status\",\n headers=headers,\n json={\"get_all_connectors\": True},\n )\n response.raise_for_status()\n\n indexing_status_response = response.json()\n\n # Iterate over all connectors and resume paused ones\n for connectors_by_source in indexing_status_response:\n if (\n specific_connector_sources\n and connectors_by_source[\"source\"] not in specific_connector_sources\n ):\n print(f\"Skipping connector source: {connectors_by_source['source']}\")\n continue\n connectors = connectors_by_source[\"indexing_statuses\"]\n for connector in connectors:\n if connector.get(\"cc_pair_status\"):\n if connector[\"cc_pair_status\"] == \"PAUSED\":\n print(f\"Resuming connector: {connector['name']}\")\n response = requests.put(\n f\"{api_server_url}/api/manage/admin/cc-pair/{connector['cc_pair_id']}/status\",\n json={\"status\": \"ACTIVE\"},\n headers=headers,\n )\n response.raise_for_status()\n print(f\"Resumed connector: {connector['name']}\")\n else:\n print(f\"Connector {connector['name']} is not paused\")\n else:\n print(f\"Connector {connector['name']} is a Federated Connector\")\n\n\ndef main() -> None:\n parser = argparse.ArgumentParser(description=\"Resume paused connectors\")\n parser.add_argument(\n \"--api_server_url\",\n type=str,\n default=API_SERVER_URL,\n help=\"The URL of the API server to use. If not provided, will use the default.\",\n )\n parser.add_argument(\n \"--api_key\",\n type=str,\n default=None,\n help=\"The API key to use for authentication. If not provided, no authentication will be used.\",\n )\n parser.add_argument(\n \"--connector_sources\",\n type=str.lower,\n nargs=\"+\",\n help=\"The sources of the connectors to resume. If not provided, will resume all paused connectors.\",\n )\n args = parser.parse_args()\n\n resume_paused_connectors(args.api_server_url, args.api_key, args.connector_sources)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/run_industryrag_bench_questions.py", "content": "from __future__ import annotations\n\nimport argparse\nimport asyncio\nimport json\nimport logging\nimport sys\nimport time\nfrom dataclasses import asdict\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import TypedDict\nfrom typing import TypeGuard\n\nimport aiohttp\n\n\nlogging.basicConfig(\n level=logging.INFO,\n format=\"%(asctime)s %(levelname)s %(message)s\",\n)\nlogger = logging.getLogger(__name__)\n\n\nDEFAULT_API_BASE = \"http://localhost:3000\"\nINTERNAL_SEARCH_TOOL_NAME = \"internal_search\"\nINTERNAL_SEARCH_IN_CODE_TOOL_ID = \"SearchTool\"\nMAX_REQUEST_ATTEMPTS = 5\nRETRIABLE_STATUS_CODES = {429, 500, 502, 503, 504}\nQUESTION_TIMEOUT_SECONDS = 300\nQUESTION_RETRY_PAUSE_SECONDS = 30\nMAX_QUESTION_ATTEMPTS = 3\n\n\n@dataclass(frozen=True)\nclass QuestionRecord:\n question_id: str\n question: str\n\n\n@dataclass(frozen=True)\nclass AnswerRecord:\n question_id: str\n answer: str\n document_ids: list[str]\n\n\n@dataclass(frozen=True)\nclass FailedQuestionRecord:\n question_id: str\n error: str\n\n\nclass Citation(TypedDict, total=False):\n citation_number: int\n document_id: str\n\n\ndef parse_args() -> argparse.Namespace:\n parser = argparse.ArgumentParser(\n description=(\n \"Submit questions to Onyx chat with internal search forced and write \"\n \"answers to a JSONL file.\"\n )\n )\n parser.add_argument(\n \"--questions-file\",\n type=Path,\n required=True,\n help=\"Path to the input questions JSONL file.\",\n )\n parser.add_argument(\n \"--output-file\",\n type=Path,\n required=True,\n help=\"Path to the output answers JSONL file.\",\n )\n parser.add_argument(\n \"--api-key\",\n type=str,\n required=True,\n help=\"API key used to authenticate against Onyx.\",\n )\n parser.add_argument(\n \"--api-base\",\n type=str,\n default=DEFAULT_API_BASE,\n help=(\n \"Frontend base URL for Onyx. If `/api` is omitted, it will be added \"\n f\"automatically. Default: {DEFAULT_API_BASE}\"\n ),\n )\n parser.add_argument(\n \"--parallelism\",\n type=int,\n default=1,\n help=\"Number of questions to process in parallel. Default: 1.\",\n )\n parser.add_argument(\n \"--max-questions\",\n type=int,\n default=None,\n help=\"Optional cap on how many questions to process. Defaults to all.\",\n )\n return parser.parse_args()\n\n\ndef normalize_api_base(api_base: str) -> str:\n normalized = api_base.rstrip(\"/\")\n if normalized.endswith(\"/api\"):\n return normalized\n return f\"{normalized}/api\"\n\n\ndef load_completed_question_ids(output_file: Path) -> set[str]:\n if not output_file.exists():\n return set()\n\n completed_ids: set[str] = set()\n with output_file.open(\"r\", encoding=\"utf-8\") as file:\n for line in file:\n stripped = line.strip()\n if not stripped:\n continue\n try:\n record = json.loads(stripped)\n except json.JSONDecodeError:\n continue\n question_id = record.get(\"question_id\")\n if isinstance(question_id, str) and question_id:\n completed_ids.add(question_id)\n\n return completed_ids\n\n\ndef load_questions(questions_file: Path) -> list[QuestionRecord]:\n if not questions_file.exists():\n raise FileNotFoundError(f\"Questions file not found: {questions_file}\")\n\n questions: list[QuestionRecord] = []\n with questions_file.open(\"r\", encoding=\"utf-8\") as file:\n for line_number, line in enumerate(file, start=1):\n stripped_line = line.strip()\n if not stripped_line:\n continue\n\n try:\n payload = json.loads(stripped_line)\n except json.JSONDecodeError as exc:\n raise ValueError(\n f\"Invalid JSON on line {line_number} of {questions_file}\"\n ) from exc\n\n question_id = payload.get(\"question_id\")\n question = payload.get(\"question\")\n\n if not isinstance(question_id, str) or not question_id:\n raise ValueError(\n f\"Line {line_number} is missing a non-empty `question_id`.\"\n )\n if not isinstance(question, str) or not question:\n raise ValueError(\n f\"Line {line_number} is missing a non-empty `question`.\"\n )\n\n questions.append(QuestionRecord(question_id=question_id, question=question))\n\n return questions\n\n\nasync def read_json_response(\n response: aiohttp.ClientResponse,\n) -> dict[str, Any] | list[dict[str, Any]]:\n response_text = await response.text()\n if response.status >= 400:\n raise RuntimeError(\n f\"Request to {response.url} failed with {response.status}: {response_text}\"\n )\n\n try:\n payload = json.loads(response_text)\n except json.JSONDecodeError as exc:\n raise RuntimeError(\n f\"Request to {response.url} returned non-JSON content: {response_text}\"\n ) from exc\n\n if not isinstance(payload, (dict, list)):\n raise RuntimeError(\n f\"Unexpected response payload type from {response.url}: {type(payload)}\"\n )\n\n return payload\n\n\nasync def request_json_with_retries(\n session: aiohttp.ClientSession,\n method: str,\n url: str,\n headers: dict[str, str],\n json_payload: dict[str, Any] | None = None,\n) -> dict[str, Any] | list[dict[str, Any]]:\n backoff_seconds = 1.0\n\n for attempt in range(1, MAX_REQUEST_ATTEMPTS + 1):\n try:\n async with session.request(\n method=method,\n url=url,\n headers=headers,\n json=json_payload,\n ) as response:\n if (\n response.status in RETRIABLE_STATUS_CODES\n and attempt < MAX_REQUEST_ATTEMPTS\n ):\n response_text = await response.text()\n logger.warning(\n \"Retryable response from %s on attempt %s/%s: %s %s\",\n url,\n attempt,\n MAX_REQUEST_ATTEMPTS,\n response.status,\n response_text,\n )\n await asyncio.sleep(backoff_seconds)\n backoff_seconds *= 2\n continue\n\n return await read_json_response(response)\n except (aiohttp.ClientError, asyncio.TimeoutError) as exc:\n if attempt == MAX_REQUEST_ATTEMPTS:\n raise RuntimeError(\n f\"Request to {url} failed after {MAX_REQUEST_ATTEMPTS} attempts.\"\n ) from exc\n\n logger.warning(\n \"Request to %s failed on attempt %s/%s: %s\",\n url,\n attempt,\n MAX_REQUEST_ATTEMPTS,\n exc,\n )\n await asyncio.sleep(backoff_seconds)\n backoff_seconds *= 2\n\n raise RuntimeError(f\"Request to {url} failed unexpectedly.\")\n\n\ndef extract_document_ids(citation_info: object) -> list[str]:\n if not isinstance(citation_info, list):\n return []\n\n sorted_citations = sorted(\n (citation for citation in citation_info if _is_valid_citation(citation)),\n key=_citation_sort_key,\n )\n\n document_ids: list[str] = []\n seen_document_ids: set[str] = set()\n for citation in sorted_citations:\n document_id = citation[\"document_id\"]\n if document_id not in seen_document_ids:\n seen_document_ids.add(document_id)\n document_ids.append(document_id)\n\n return document_ids\n\n\ndef _is_valid_citation(citation: object) -> TypeGuard[Citation]:\n return (\n isinstance(citation, dict)\n and isinstance(citation.get(\"document_id\"), str)\n and bool(citation[\"document_id\"])\n )\n\n\ndef _citation_sort_key(citation: Citation) -> int:\n citation_number = citation.get(\"citation_number\")\n if isinstance(citation_number, int):\n return citation_number\n return sys.maxsize\n\n\nasync def fetch_internal_search_tool_id(\n session: aiohttp.ClientSession,\n api_base: str,\n headers: dict[str, str],\n) -> int:\n payload = await request_json_with_retries(\n session=session,\n method=\"GET\",\n url=f\"{api_base}/tool\",\n headers=headers,\n )\n\n if not isinstance(payload, list):\n raise RuntimeError(\"Expected `/tool` to return a list.\")\n\n for tool in payload:\n if not isinstance(tool, dict):\n continue\n\n if tool.get(\"in_code_tool_id\") == INTERNAL_SEARCH_IN_CODE_TOOL_ID:\n tool_id = tool.get(\"id\")\n if isinstance(tool_id, int):\n return tool_id\n\n for tool in payload:\n if not isinstance(tool, dict):\n continue\n\n if tool.get(\"name\") == INTERNAL_SEARCH_TOOL_NAME:\n tool_id = tool.get(\"id\")\n if isinstance(tool_id, int):\n return tool_id\n\n raise RuntimeError(\n \"Could not find the internal search tool in `/tool`. \"\n \"Make sure SearchTool is available for this environment.\"\n )\n\n\nasync def submit_question(\n session: aiohttp.ClientSession,\n api_base: str,\n headers: dict[str, str],\n internal_search_tool_id: int,\n question_record: QuestionRecord,\n) -> AnswerRecord:\n payload = {\n \"message\": question_record.question,\n \"chat_session_info\": {\"persona_id\": 0},\n \"parent_message_id\": None,\n \"file_descriptors\": [],\n \"allowed_tool_ids\": [internal_search_tool_id],\n \"forced_tool_id\": internal_search_tool_id,\n \"stream\": False,\n }\n\n response_payload = await request_json_with_retries(\n session=session,\n method=\"POST\",\n url=f\"{api_base}/chat/send-chat-message\",\n headers=headers,\n json_payload=payload,\n )\n\n if not isinstance(response_payload, dict):\n raise RuntimeError(\n \"Expected `/chat/send-chat-message` to return an object when `stream=false`.\"\n )\n\n answer = response_payload.get(\"answer_citationless\")\n if not isinstance(answer, str):\n answer = response_payload.get(\"answer\")\n\n if not isinstance(answer, str):\n raise RuntimeError(\n f\"Response for question {question_record.question_id} is missing `answer`.\"\n )\n\n return AnswerRecord(\n question_id=question_record.question_id,\n answer=answer,\n document_ids=extract_document_ids(response_payload.get(\"citation_info\")),\n )\n\n\nasync def generate_answers(\n questions: list[QuestionRecord],\n output_file: Path,\n api_base: str,\n api_key: str,\n parallelism: int,\n skipped: int,\n) -> None:\n if parallelism < 1:\n raise ValueError(\"`--parallelism` must be at least 1.\")\n\n headers = {\n \"Authorization\": f\"Bearer {api_key}\",\n \"Content-Type\": \"application/json\",\n }\n\n timeout = aiohttp.ClientTimeout(\n total=None,\n connect=30,\n sock_connect=30,\n sock_read=600,\n )\n connector = aiohttp.TCPConnector(limit=parallelism)\n\n output_file.parent.mkdir(parents=True, exist_ok=True)\n with output_file.open(\"a\", encoding=\"utf-8\") as file:\n async with aiohttp.ClientSession(\n timeout=timeout, connector=connector\n ) as session:\n internal_search_tool_id = await fetch_internal_search_tool_id(\n session=session,\n api_base=api_base,\n headers=headers,\n )\n logger.info(\"Using internal search tool id %s\", internal_search_tool_id)\n\n semaphore = asyncio.Semaphore(parallelism)\n progress_lock = asyncio.Lock()\n write_lock = asyncio.Lock()\n completed = 0\n successful = 0\n stuck_count = 0\n failed_questions: list[FailedQuestionRecord] = []\n remaining_count = len(questions)\n overall_total = remaining_count + skipped\n question_durations: list[float] = []\n run_start_time = time.monotonic()\n\n def print_progress() -> None:\n avg_time = (\n sum(question_durations) / len(question_durations)\n if question_durations\n else 0.0\n )\n elapsed = time.monotonic() - run_start_time\n eta = avg_time * (remaining_count - completed) / max(parallelism, 1)\n\n done = skipped + completed\n bar_width = 30\n filled = (\n int(bar_width * done / overall_total)\n if overall_total\n else bar_width\n )\n bar = \"█\" * filled + \"░\" * (bar_width - filled)\n pct = (done / overall_total * 100) if overall_total else 100.0\n\n parts = (\n f\"\\r{bar} {pct:5.1f}% \"\n f\"[{done}/{overall_total}] \"\n f\"avg {avg_time:.1f}s/q \"\n f\"elapsed {elapsed:.0f}s \"\n f\"ETA {eta:.0f}s \"\n f\"(ok:{successful} fail:{len(failed_questions)}\"\n )\n if stuck_count:\n parts += f\" stuck:{stuck_count}\"\n if skipped:\n parts += f\" skip:{skipped}\"\n parts += \")\"\n\n sys.stderr.write(parts)\n sys.stderr.flush()\n\n print_progress()\n\n async def process_question(question_record: QuestionRecord) -> None:\n nonlocal completed\n nonlocal successful\n nonlocal stuck_count\n\n last_error: Exception | None = None\n for attempt in range(1, MAX_QUESTION_ATTEMPTS + 1):\n q_start = time.monotonic()\n try:\n async with semaphore:\n result = await asyncio.wait_for(\n submit_question(\n session=session,\n api_base=api_base,\n headers=headers,\n internal_search_tool_id=internal_search_tool_id,\n question_record=question_record,\n ),\n timeout=QUESTION_TIMEOUT_SECONDS,\n )\n except asyncio.TimeoutError:\n async with progress_lock:\n stuck_count += 1\n logger.warning(\n \"Question %s timed out after %ss (attempt %s/%s, \"\n \"total stuck: %s) — retrying in %ss\",\n question_record.question_id,\n QUESTION_TIMEOUT_SECONDS,\n attempt,\n MAX_QUESTION_ATTEMPTS,\n stuck_count,\n QUESTION_RETRY_PAUSE_SECONDS,\n )\n print_progress()\n last_error = TimeoutError(\n f\"Timed out after {QUESTION_TIMEOUT_SECONDS}s \"\n f\"on attempt {attempt}/{MAX_QUESTION_ATTEMPTS}\"\n )\n await asyncio.sleep(QUESTION_RETRY_PAUSE_SECONDS)\n continue\n except Exception as exc:\n duration = time.monotonic() - q_start\n async with progress_lock:\n completed += 1\n question_durations.append(duration)\n failed_questions.append(\n FailedQuestionRecord(\n question_id=question_record.question_id,\n error=str(exc),\n )\n )\n logger.exception(\n \"Failed question %s (%s/%s)\",\n question_record.question_id,\n completed,\n remaining_count,\n )\n print_progress()\n return\n\n duration = time.monotonic() - q_start\n\n async with write_lock:\n file.write(json.dumps(asdict(result), ensure_ascii=False))\n file.write(\"\\n\")\n file.flush()\n\n async with progress_lock:\n completed += 1\n successful += 1\n question_durations.append(duration)\n print_progress()\n return\n\n # All attempts exhausted due to timeouts\n async with progress_lock:\n completed += 1\n failed_questions.append(\n FailedQuestionRecord(\n question_id=question_record.question_id,\n error=str(last_error),\n )\n )\n logger.error(\n \"Question %s failed after %s timeout attempts (%s/%s)\",\n question_record.question_id,\n MAX_QUESTION_ATTEMPTS,\n completed,\n remaining_count,\n )\n print_progress()\n\n await asyncio.gather(\n *(process_question(question_record) for question_record in questions)\n )\n\n # Final newline after progress bar\n sys.stderr.write(\"\\n\")\n sys.stderr.flush()\n\n total_elapsed = time.monotonic() - run_start_time\n avg_time = (\n sum(question_durations) / len(question_durations)\n if question_durations\n else 0.0\n )\n stuck_suffix = f\", {stuck_count} stuck timeouts\" if stuck_count else \"\"\n resume_suffix = (\n f\" — {skipped} previously completed, \"\n f\"{skipped + successful}/{overall_total} overall\"\n if skipped\n else \"\"\n )\n logger.info(\n \"Done: %s/%s successful in %.1fs (avg %.1fs/question%s)%s\",\n successful,\n remaining_count,\n total_elapsed,\n avg_time,\n stuck_suffix,\n resume_suffix,\n )\n\n if failed_questions:\n logger.warning(\n \"%s questions failed:\",\n len(failed_questions),\n )\n for failed_question in failed_questions:\n logger.warning(\n \"Failed question %s: %s\",\n failed_question.question_id,\n failed_question.error,\n )\n\n\ndef main() -> None:\n args = parse_args()\n questions = load_questions(args.questions_file)\n api_base = normalize_api_base(args.api_base)\n\n if args.max_questions is not None:\n if args.max_questions < 1:\n raise ValueError(\"`--max-questions` must be at least 1 when provided.\")\n questions = questions[: args.max_questions]\n\n completed_ids = load_completed_question_ids(args.output_file)\n logger.info(\n \"Found %s already-answered question IDs in %s\",\n len(completed_ids),\n args.output_file,\n )\n total_before_filter = len(questions)\n questions = [q for q in questions if q.question_id not in completed_ids]\n skipped = total_before_filter - len(questions)\n\n if skipped:\n logger.info(\n \"Resuming: %s/%s already answered, %s remaining\",\n skipped,\n total_before_filter,\n len(questions),\n )\n else:\n logger.info(\"Loaded %s questions from %s\", len(questions), args.questions_file)\n\n if not questions:\n logger.info(\"All questions already answered. Nothing to do.\")\n return\n\n logger.info(\"Writing answers to %s\", args.output_file)\n\n asyncio.run(\n generate_answers(\n questions=questions,\n output_file=args.output_file,\n api_base=api_base,\n api_key=args.api_key,\n parallelism=args.parallelism,\n skipped=skipped,\n )\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/save_load_state.py", "content": "# This file is purely for development use, not included in any builds\n# Remember to first to send over the schema information (run API Server)\nimport argparse\nimport json\nimport os\nimport subprocess\n\nimport requests\n\nfrom alembic import command\nfrom alembic.config import Config\nfrom onyx.configs.app_configs import POSTGRES_DB\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PASSWORD\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.document_index.vespa.index import DOCUMENT_ID_ENDPOINT\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\ndef save_postgres(filename: str, container_name: str) -> None:\n logger.notice(\"Attempting to take Postgres snapshot\")\n cmd = f\"docker exec {container_name} pg_dump -U {POSTGRES_USER} -h {POSTGRES_HOST} -p {POSTGRES_PORT} -W -F t {POSTGRES_DB}\"\n with open(filename, \"w\") as file:\n subprocess.run(\n cmd,\n shell=True,\n check=True,\n stdout=file,\n text=True,\n input=f\"{POSTGRES_PASSWORD}\\n\",\n )\n\n\ndef load_postgres(filename: str, container_name: str) -> None:\n logger.notice(\"Attempting to load Postgres snapshot\")\n try:\n alembic_cfg = Config(\"alembic.ini\")\n command.upgrade(alembic_cfg, \"head\")\n except Exception as e:\n logger.error(f\"Alembic upgrade failed: {e}\")\n\n host_file_path = os.path.abspath(filename)\n\n copy_cmd = f\"docker cp {host_file_path} {container_name}:/tmp/\"\n subprocess.run(copy_cmd, shell=True, check=True)\n\n container_file_path = f\"/tmp/{os.path.basename(filename)}\"\n\n restore_cmd = (\n f\"docker exec {container_name} pg_restore --clean -U {POSTGRES_USER} \"\n f\"-h localhost -p {POSTGRES_PORT} -d {POSTGRES_DB} -1 -F t {container_file_path}\"\n )\n subprocess.run(restore_cmd, shell=True, check=True)\n\n\ndef save_vespa(filename: str) -> None:\n logger.notice(\"Attempting to take Vespa snapshot\")\n continuation = \"\"\n params = {}\n doc_jsons: list[dict] = []\n while continuation is not None:\n if continuation:\n params = {\"continuation\": continuation}\n response = requests.get(DOCUMENT_ID_ENDPOINT, params=params)\n response.raise_for_status()\n found = response.json()\n continuation = found.get(\"continuation\")\n docs = found[\"documents\"]\n for doc in docs:\n doc_json = {\"update\": doc[\"id\"], \"create\": True, \"fields\": doc[\"fields\"]}\n doc_jsons.append(doc_json)\n\n with open(filename, \"w\") as jsonl_file:\n for doc in doc_jsons:\n json_str = json.dumps(doc)\n jsonl_file.write(json_str + \"\\n\")\n\n\ndef load_vespa(filename: str) -> None:\n headers = {\"Content-Type\": \"application/json\"}\n with open(filename, \"r\") as f:\n for line in f:\n new_doc = json.loads(line.strip())\n doc_id = new_doc[\"update\"].split(\"::\")[-1]\n response = requests.post(\n DOCUMENT_ID_ENDPOINT + \"/\" + doc_id,\n headers=headers,\n json=new_doc,\n )\n response.raise_for_status()\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(description=\"Onyx checkpoint saving and loading.\")\n parser.add_argument(\n \"--save\", action=\"store_true\", help=\"Save Onyx state to directory.\"\n )\n parser.add_argument(\n \"--load\", action=\"store_true\", help=\"Load Onyx state from save directory.\"\n )\n parser.add_argument(\n \"--postgres_container_name\",\n type=str,\n default=\"onyx-relational_db-1\",\n help=\"Name of the postgres container to dump\",\n )\n parser.add_argument(\n \"--checkpoint_dir\",\n type=str,\n default=os.path.join(\"..\", \"onyx_checkpoint\"),\n help=\"A directory to store temporary files to.\",\n )\n\n args = parser.parse_args()\n checkpoint_dir = args.checkpoint_dir\n postgres_container = args.postgres_container_name\n\n if not os.path.exists(checkpoint_dir):\n os.makedirs(checkpoint_dir)\n\n if not args.save and not args.load:\n raise ValueError(\"Must specify --save or --load\")\n\n if args.load:\n load_postgres(\n os.path.join(checkpoint_dir, \"postgres_snapshot.tar\"), postgres_container\n )\n load_vespa(os.path.join(checkpoint_dir, \"vespa_snapshot.jsonl\"))\n else:\n save_postgres(\n os.path.join(checkpoint_dir, \"postgres_snapshot.tar\"), postgres_container\n )\n save_vespa(os.path.join(checkpoint_dir, \"vespa_snapshot.jsonl\"))\n" }, { "path": "backend/scripts/setup_craft_templates.sh", "content": "#!/bin/sh\n# Setup Onyx Craft templates\n# This script is called on container startup to ensure Craft templates are ready\n# Set ENABLE_CRAFT=false to skip setup\n\n# Check if Craft is disabled\nif [ \"$ENABLE_CRAFT\" = \"false\" ] || [ \"$ENABLE_CRAFT\" = \"False\" ]; then\n echo \"Onyx Craft is disabled (ENABLE_CRAFT=false), skipping template setup\"\n exit 0\nfi\n\nset -e\n\n# Verify opencode CLI is available (installed in Dockerfile)\nif ! command -v opencode >/dev/null 2>&1; then\n echo \"ERROR: opencode CLI is not available but ENABLE_CRAFT is enabled.\" >&2\n echo \"opencode is required for Craft agent functionality. Ensure you are using Dockerfile\" >&2\n echo \"which includes the opencode CLI, or set ENABLE_CRAFT=false to disable Craft.\" >&2\n exit 1\nfi\n\nCRAFT_BASE=\"/app/onyx/server/features/build/sandbox/kubernetes/docker\"\nDEMO_DATA_ZIP=\"${CRAFT_BASE}/demo_data.zip\"\nDEMO_DATA_DIR=\"${CRAFT_BASE}/demo_data\"\n# Use environment variables if set, otherwise use defaults\nOUTPUTS_TEMPLATE_PATH=\"${OUTPUTS_TEMPLATE_PATH:-${CRAFT_BASE}/templates/outputs}\"\nVENV_TEMPLATE_PATH=\"${VENV_TEMPLATE_PATH:-${CRAFT_BASE}/templates/venv}\"\nWEB_TEMPLATE_PATH=\"${WEB_TEMPLATE_PATH:-${OUTPUTS_TEMPLATE_PATH}/web}\"\nREQUIREMENTS_PATH=\"${CRAFT_BASE}/initial-requirements.txt\"\n\necho \"Setting up Onyx Craft templates...\"\n\n# 1. Unzip demo_data.zip if demo_data directory doesn't exist\nif [ ! -d \"$DEMO_DATA_DIR\" ] && [ -f \"$DEMO_DATA_ZIP\" ]; then\n echo \" Extracting demo data...\"\n cd \"$CRAFT_BASE\" && unzip -q demo_data.zip || { echo \"ERROR: Failed to extract demo data\" >&2; exit 1; }\n echo \" Demo data extracted\"\nfi\n\n# 2. Create Python venv template if it doesn't exist\nif [ ! -d \"$VENV_TEMPLATE_PATH\" ] && [ -f \"$REQUIREMENTS_PATH\" ]; then\n echo \" Creating Python venv template (this may take 30-60 seconds)...\"\n python -m venv \"$VENV_TEMPLATE_PATH\"\n \"$VENV_TEMPLATE_PATH/bin/pip\" install --upgrade pip -q\n \"$VENV_TEMPLATE_PATH/bin/pip\" install -q -r \"$REQUIREMENTS_PATH\"\n echo \" Python venv template created\"\nfi\n\n# 3. Run npm install in web template\nif [ -d \"$WEB_TEMPLATE_PATH\" ]; then\n if ! command -v npm >/dev/null 2>&1; then\n echo \"ERROR: npm is not available but ENABLE_CRAFT is enabled.\" >&2\n echo \"npm is required for Craft web features. Ensure you are using Dockerfile\" >&2\n echo \"which includes Node.js, or set ENABLE_CRAFT=false to disable Craft.\" >&2\n exit 1\n fi\n # Always remove and reinstall to ensure correct architecture binaries\n if [ -d \"${WEB_TEMPLATE_PATH}/node_modules\" ]; then\n echo \" Removing existing node_modules...\"\n rm -rf \"${WEB_TEMPLATE_PATH}/node_modules\"\n fi\n echo \" Installing npm packages (this may take 1-2 minutes)...\"\n cd \"$WEB_TEMPLATE_PATH\" && npm install 2>&1 || { echo \"ERROR: npm install failed\" >&2; exit 1; }\n echo \" Web template dependencies installed\"\nfi\n\necho \"Craft template setup complete\"\n" }, { "path": "backend/scripts/sources_selection_analysis.py", "content": "import argparse\nimport json\nimport os\nimport sys\nimport time\nfrom datetime import datetime\nfrom os import listdir\nfrom os.path import isfile\nfrom os.path import join\nfrom typing import Optional\n\nimport requests\n\nfrom onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME\n\nparent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\nsys.path.append(parent_dir)\n\nfrom onyx.configs.app_configs import DOCUMENT_INDEX_NAME # noqa: E402\nfrom onyx.configs.constants import SOURCE_TYPE # noqa: E402\n\nANALYSIS_FOLDER = f\"{parent_dir}/scripts/.analysisfiles/\"\n\n\ndef color_output(\n text: str,\n model: Optional[str] = None,\n text_color: str = \"white\",\n bg_color: str = \"black\",\n text_style: str = \"normal\",\n text_prefix: str = \"\",\n) -> None:\n \"\"\"Color and print a text\n\n Args:\n text (str): The text to display\n model (str, optional): A pre-defined output model. Defaults to None.\n text_color (str, optional): Define the text color. Defaults to \"white\".\n bg_color (str, optional): Define the background color. Defaults to \"black\".\n text_style (str, optional): Define the text style. Defaults to \"normal\".\n text_prefix (str, optional): Set a text prefix. Defaults to \"\".\n \"\"\"\n if model:\n if model == \"alert\":\n text_color = \"black\"\n bg_color = \"red\"\n text_style = \"bold\"\n elif model == \"critical\":\n text_prefix = \"CRITICAL: \"\n text_color = \"white\"\n bg_color = \"red\"\n text_style = \"bold\"\n elif model == \"note\":\n text_color = \"yellow\"\n bg_color = \"transparent\"\n text_style = \"normal\"\n elif model == \"info\":\n text_prefix = \"INFO: \"\n text_color = \"black\"\n bg_color = \"yellow\"\n text_style = \"bold\"\n elif model == \"info2\":\n text_prefix = \"INFO: \"\n text_color = \"black\"\n bg_color = \"white\"\n text_style = \"bold\"\n elif model == \"valid\":\n text_prefix = \"INFO: \"\n text_color = \"white\"\n bg_color = \"green\"\n text_style = \"bold\"\n elif model == \"debug\":\n text_prefix = \"DEBUG: \"\n text_color = \"blue\"\n bg_color = \"transparent\"\n text_style = \"bold\"\n\n text_colors = {\n \"black\": 30,\n \"red\": 31,\n \"green\": 32,\n \"yellow\": 33,\n \"blue\": 34,\n \"purple\": 35,\n \"cian\": 36,\n \"white\": 37,\n }\n bg_colors = {\n \"black\": 40,\n \"red\": 41,\n \"green\": 42,\n \"yellow\": 43,\n \"blue\": 44,\n \"purple\": 45,\n \"cian\": 46,\n \"white\": 47,\n \"transparent\": 49,\n }\n text_styles = {\n \"normal\": 0,\n \"bold\": 1,\n \"light\": 2,\n \"italicized\": 3,\n \"underlined\": 4,\n \"blink\": 5,\n }\n print(\n f\"\\033[{text_styles[text_style]};{text_colors[text_color]};{bg_colors[bg_color]}m {text_prefix} {text} \\033[0;0m\"\n )\n\n\nclass CompareAnalysis:\n def __init__(\n self, query: str, previous_content: dict, new_content: dict, threshold: float\n ) -> None:\n \"\"\"Make the comparison between 2 analysis for a specific query\n\n Args:\n query (str): The analysed query\n previous_content (dict): The previous analysis content for the selected query\n new_content (dict): The new analysis content for the selected query\n threshold (float): The minimum difference (percentage) between scores to raise an anomaly\n \"\"\"\n self._query = query\n self._previous_content = previous_content\n self._new_content = new_content\n self._threshold = threshold\n\n def _identify_diff(self, content_key: str) -> list[dict]:\n \"\"\"Try to identify differences between the two analysis based\n on the selected analysis key.\n\n Args:\n content_key (str): The analysis item's key to compare the versions.\n Examples: score / document_id\n\n Returns:\n list[dict]: List of dict representing the information regarding the difference\n Format: {\n \"previous_rank\": XX,\n \"new_rank\": XX,\n \"document_id\": XXXX,\n \"previous_score\": XX,\n \"new_score\": XX,\n \"score_change_pct\": XX\n }\n \"\"\"\n changes = []\n\n previous_content = {\n k: v[content_key] for k, v in self._previous_content.items()\n }\n new_content = {k: v[content_key] for k, v in self._new_content.items()}\n\n if previous_content != new_content:\n for pos, data in previous_content.items():\n if data != new_content[pos]:\n try:\n score_change_pct = round(\n (\n abs(\n self._new_content[pos][\"score\"]\n - self._previous_content[pos][\"score\"]\n )\n / self._new_content[pos][\"score\"]\n )\n * 100.0,\n 2,\n )\n except ZeroDivisionError:\n score_change_pct = 0\n\n changes.append(\n {\n \"previous_rank\": pos,\n \"new_rank\": (\n pos\n if content_key == \"score\"\n else {\n \"x\": k for k, v in new_content.items() if v == data\n }.get(\"x\", \"not_ranked\")\n ),\n \"document_id\": self._previous_content[pos][\"document_id\"],\n \"previous_score\": self._previous_content[pos][\"score\"],\n \"new_score\": self._new_content[pos][\"score\"],\n \"score_change_pct\": score_change_pct,\n }\n )\n return changes\n\n def check_config_changes(\n self, previous_doc_rank: int | str, new_doc_rank: int\n ) -> None:\n \"\"\"Try to identify possible reasons why a change has been detected by\n checking the latest document update date or the boost value.\n\n Args:\n previous_doc_rank (int): The document rank for the previous analysis\n new_doc_rank (int): The document rank for the new analysis\n \"\"\"\n if isinstance(new_doc_rank, str) and new_doc_rank == \"not_ranked\":\n color_output(\n (\n \"NOTE: The document is missing in the 'current' analysis file. \"\n \"Unable to identify more details about the reason for the change.\"\n ),\n model=\"note\",\n )\n return None\n\n if (\n self._previous_content[previous_doc_rank][\"boost\"]\n != self._new_content[new_doc_rank][\"boost\"]\n ):\n color_output(\n \"NOTE: The 'boost' value has been changed which (maybe) explains the change.\",\n model=\"note\",\n )\n color_output(\n (\n f\"Previously it was '{self._previous_content[previous_doc_rank]['boost']}' \"\n f\"and now is set to '{self._new_content[new_doc_rank]['boost']}'\"\n ),\n model=\"note\",\n )\n if (\n self._previous_content[previous_doc_rank][\"updated_at\"]\n != self._new_content[new_doc_rank][\"updated_at\"]\n ):\n color_output(\"NOTE: The document seems to have been updated.\", model=\"note\")\n color_output(\n (\n f\"Previously the updated date was '{self._previous_content[previous_doc_rank]['updated_at']}' \"\n f\"and now is '{self._new_content[new_doc_rank]['updated_at']}'\"\n ),\n model=\"note\",\n )\n\n def check_documents_score(self) -> bool:\n \"\"\"Check if the scores have changed between analysis.\n\n Returns:\n bool: True if at least one change has been detected. False otherwise.\n \"\"\"\n color_output(\"Checking documents Score....\", model=\"info\")\n color_output(\n f\"Differences under '{self._threshold}%' are ignored (based on the '--threshold' argument)\",\n model=\"info\",\n )\n\n if diff := [\n x\n for x in self._identify_diff(\"score\")\n if x[\"score_change_pct\"] > self._threshold\n ]:\n color_output(\"<<<<< Changes detected >>>>>\", model=\"alert\")\n for change in diff:\n color_output(\"-\" * 100)\n color_output(\n (\n f\"The document '{change['document_id']}' (rank: {change['previous_rank']}) \"\n f\"score has a changed of {change['score_change_pct']}%\"\n )\n )\n color_output(f\"previous score: {change['previous_score']}\")\n color_output(f\"current score: {change['new_score']}\")\n self.check_config_changes(change[\"previous_rank\"], change[\"new_rank\"])\n\n color_output(\"<<<<< End of changes >>>>>\", model=\"alert\")\n color_output(f\"Number of changes detected {len(diff)}\", model=\"info\")\n else:\n color_output(\"No change detected\", model=\"valid\")\n color_output(\"Documents Score check completed.\", model=\"info\")\n\n return False if diff else True\n\n def check_documents_order(self) -> bool:\n \"\"\"Check if the selected documents are the same and in the same order.\n\n Returns:\n bool: True if at least one change has been detected. False otherwise.\n \"\"\"\n color_output(\"Checking documents Order....\", model=\"info\")\n\n if diff := self._identify_diff(\"document_id\"):\n color_output(\"<<<<< Changes detected >>>>>\", model=\"alert\")\n for change in diff:\n color_output(\"-\" * 100)\n color_output(\n (\n f\"The document '{change['document_id']}' was at a rank \"\n f\"'{change['previous_rank']}' but now is at rank '{change['new_rank']}'\"\n )\n )\n color_output(f\"previous score: {change['previous_score']}\")\n color_output(f\"current score: {change['new_score']}\")\n self.check_config_changes(change[\"previous_rank\"], change[\"new_rank\"])\n color_output(\"<<<<< End of changes >>>>>\", model=\"alert\")\n color_output(f\"Number of changes detected {len(diff)}\", model=\"info\")\n\n else:\n color_output(\"No change detected\", model=\"valid\")\n color_output(\"Documents order check completed.\", model=\"info\")\n\n return False if diff else True\n\n def __call__(self) -> None:\n \"\"\"Manage the analysis process\"\"\"\n if not self.check_documents_order():\n color_output(\n \"Skipping other checks as the documents order has changed\", model=\"info\"\n )\n return None\n\n self.check_documents_score()\n\n\nclass SelectionAnalysis:\n def __init__(\n self,\n exectype: str,\n analysisfiles: list = [],\n queries: list = [],\n threshold: float = 0.0,\n web_port: int = 3000,\n auth_cookie: str = \"\",\n wait: int = 10,\n ) -> None:\n \"\"\"\n\n Args:\n exectype (str): The execution mode (new or compare)\n analysisfiles (list, optional): List of analysis files to compare or if only one, to use as the base. Defaults to [].\n Requiered only by the 'compare' mode\n queries (list, optional): The queries to analysed. Defaults to [].\n Required only by the 'new' mode\n threshold (float, optional): The minimum difference (percentage) between scores to raise an anomaly\n web_port (int, optional): The port of the UI. Defaults to 3000 (local exec port)\n auth_cookie (str, optional): The Auth cookie value (fastapiusersauth). Defaults to None.\n wait (int, optional): The waiting time (in seconds) to respect between queries.\n It is helpful to avoid hitting the Generative AI rate limiting.\n \"\"\"\n self._exectype = exectype\n self._analysisfiles = analysisfiles\n self._queries = queries\n self._threshold = threshold\n self._web_port = web_port\n self._auth_cookie = auth_cookie\n self._wait = wait\n\n def _wait_between_queries(self, query: str) -> None:\n \"\"\"If there are remaining queries, waits for the defined time.\n\n Args:\n query (str): The latest executed query\n \"\"\"\n if query != self._queries[-1]:\n color_output(f\"Next query in {self._wait} seconds\", model=\"debug\")\n time.sleep(self._wait)\n\n def prepare(self) -> bool:\n \"\"\"Create the requirements to execute this script\n\n Returns:\n bool: True if all the requirements are setup. False otherwise\n \"\"\"\n try:\n os.makedirs(ANALYSIS_FOLDER, exist_ok=True)\n return True\n except Exception as e:\n color_output(f\"Unable to setup the requirements: {e}\", model=\"critical\")\n return False\n\n def do_request(self, query: str) -> dict:\n \"\"\"Request the Onyx API\n\n Args:\n query (str): A query\n\n Returns:\n dict: The Onyx API response content\n \"\"\"\n cookies = (\n {FASTAPI_USERS_AUTH_COOKIE_NAME: self._auth_cookie}\n if self._auth_cookie\n else {}\n )\n\n endpoint = f\"http://127.0.0.1:{self._web_port}/api/direct-qa\"\n query_json = {\n \"query\": query,\n \"collection\": DOCUMENT_INDEX_NAME,\n \"filters\": {SOURCE_TYPE: None},\n \"enable_auto_detect_filters\": True,\n \"search_type\": \"hybrid\",\n \"offset\": 0,\n \"favor_recent\": True,\n }\n try:\n response = requests.post(endpoint, json=query_json, cookies=cookies)\n if response.status_code != 200:\n color_output(\n (\n f\"something goes wrong while requesting the Onyx API for the query '{query}': {response.text}\"\n ),\n model=\"critical\",\n )\n sys.exit(1)\n except Exception as e:\n color_output(\n f\"Unable to request the Onyx API for the query '{query}': {e}\",\n model=\"critical\",\n )\n sys.exit(1)\n\n return json.loads(response.content)\n\n def get_analysis_files(self) -> list[str]:\n \"\"\"Returns the list of existing analysis files.\n\n Returns:\n list[str]: List of filename\n \"\"\"\n return [f for f in listdir(ANALYSIS_FOLDER) if isfile(join(ANALYSIS_FOLDER, f))]\n\n def get_analysis_file_content(self, filename: str) -> list[dict]:\n \"\"\"Returns the content of an analysis file\n\n Args:\n filename (str): The analysis filename\n\n Returns:\n list[dict]: Content of the selected file\n \"\"\"\n with open(f\"{ANALYSIS_FOLDER}{filename}\", \"r\") as f:\n return json.load(f)\n\n def extract_content(self, contents: dict) -> dict:\n \"\"\"Extract the content returns by the Onyx API\n\n Args:\n contents (dict): The onyx response content\n\n Returns:\n dict: Data regarding the selected sources document\n \"\"\"\n return {\n pos: doc\n for pos, doc in enumerate(\n sorted(\n contents[\"top_ranked_docs\"], key=lambda d: d[\"score\"], reverse=True\n )[:5]\n )\n }\n\n def save_analysisfile(self, content: list[dict]) -> Optional[str]:\n \"\"\"Save the extracted content\n\n Args:\n content (list[dict]): The content to save\n\n Returns:\n str: The filname\n \"\"\"\n filename = datetime.now().strftime(\"%Y_%m_%d-%I_%M_%S\")\n analysis_file = f\"{ANALYSIS_FOLDER}{filename}.json\"\n\n try:\n with open(analysis_file, \"w\") as f:\n json.dump(content, f, indent=4)\n except Exception as e:\n color_output(f\"Unable to create the analysis file: {e}\", model=\"critical\")\n return None\n\n color_output(f\"Analysis file created: {analysis_file}\", model=\"debug\")\n return analysis_file\n\n def new(self) -> Optional[str]:\n \"\"\"Manage the process to create a new analysis file\n based on the submitted queries\n\n Returns:\n str: The new filename with the analysis content\n \"\"\"\n if not self._queries:\n color_output(\"Missing queries\", model=\"critical\")\n sys.exit(1)\n\n color_output(\"Generating a new analysis file...\", model=\"debug\")\n analysisfile = []\n\n for query in self._queries:\n color_output(f\"Gathering data of the query: '{query}'\", model=\"info2\")\n contents = self.do_request(query)\n\n analysisfile.append(\n {\"query\": query, \"selected_documents\": self.extract_content(contents)}\n )\n color_output(\"Data gathered\", model=\"info2\")\n self._wait_between_queries(query)\n\n return self.save_analysisfile(analysisfile)\n\n def compare(\n self,\n previous_analysisfile_content: list[dict],\n new_analysisfile_content: list[dict],\n ) -> None:\n \"\"\"Manage the process to compare two analysis\n\n Args:\n previous_analysisfile_content (list): Previous content analysis\n new_analysisfile_content (list): New content analysis\n \"\"\"\n for query in self._queries:\n # Extract data regarding the selected source documents\n prev_querie_content = [\n x for x in previous_analysisfile_content if x[\"query\"] == query\n ][0][\"selected_documents\"]\n new_querie_content = [\n x for x in new_analysisfile_content if x[\"query\"] == query\n ][0][\"selected_documents\"]\n\n color_output(f\"Analysing the query: '{query}'\", model=\"info2\")\n CompareAnalysis(\n query, prev_querie_content, new_querie_content, self._threshold\n )()\n color_output(f\"Analyse completed for the query: '{query}'\", model=\"info2\")\n self._wait_between_queries(query)\n\n color_output(\"All the defined queries have been evaluated.\", model=\"info2\")\n\n def validate_analysisfiles(self) -> bool:\n \"\"\"Validate that the selected analysis files exist\n\n Returns:\n bool: True if all of them exist. False otherwise\n \"\"\"\n existing_analysisfiles = self.get_analysis_files()\n\n if missing_analysisfiles := [\n x for x in self._analysisfiles if x not in existing_analysisfiles\n ]:\n color_output(\n f\"Missing analysis file(s) '{', '.join(missing_analysisfiles)}' - NOT FOUND\",\n model=\"critical\",\n )\n analysisfiles = \"\\n \".join(existing_analysisfiles)\n color_output(\"Available analysis files:\", model=\"info2\")\n color_output(analysisfiles)\n return False\n\n return True\n\n def __call__(self) -> None:\n if not self.prepare():\n sys.exit(1)\n\n if self._exectype == \"new\":\n self.new()\n\n elif self._exectype == \"compare\":\n self._analysisfiles = [\n x.replace(\".json\", \"\") + \".json\" for x in self._analysisfiles\n ]\n\n if not self.validate_analysisfiles():\n sys.exit(1)\n\n color_output(\n \"Extracting queries from the existing analysis file...\", model=\"debug\"\n )\n previous_analysisfile_content = self.get_analysis_file_content(\n self._analysisfiles[0]\n )\n\n # Extract the queries\n self._queries = sorted([x[\"query\"] for x in previous_analysisfile_content])\n color_output(\n f\"Extracted queries: {', '.join(self._queries)}\", model=\"debug\"\n )\n\n if len(self._analysisfiles) == 1:\n if new_file := self.new():\n new_analysisfile_content = self.get_analysis_file_content(\n new_file.split(\"/\")[-1:][0]\n )\n return self.compare(\n previous_analysisfile_content, new_analysisfile_content\n )\n else:\n color_output(\n \"Unable to generate a new analysis file\", model=\"critical\"\n )\n sys.exit(1)\n else:\n color_output(\n (\n f\"For the rest of this execution, the analysis file '{self._analysisfiles[0]}' \"\n f\"is identified as 'previous' and '{self._analysisfiles[1]}' as 'current'\"\n ),\n model=\"info2\",\n )\n new_analysisfile_content = self.get_analysis_file_content(\n self._analysisfiles[1]\n )\n new_queries = sorted([x[\"query\"] for x in new_analysisfile_content])\n if new_queries != self._queries:\n color_output(\n \"Unable to compare analysis files as the queries are differents\",\n model=\"critical\",\n )\n sys.exit(1)\n self.compare(previous_analysisfile_content, new_analysisfile_content)\n\n\ndef validate_cmd_args(args: argparse.Namespace) -> bool:\n \"\"\"Validate the CMD arguments\n\n Args:\n args (argparse.Namespace): The argparse data input\n\n Returns:\n bool: True if the CMD arguments are valid. False otherwise\n \"\"\"\n if not args.execution:\n color_output(\n \"Missing argument. The execution mode ('--execution') must be defined ('new' or 'compare')\",\n model=\"critical\",\n )\n return False\n if args.execution == \"new\" and not args.q__queries:\n color_output(\n \"Missing argument. When the execution type is set to 'new' the '--queries' argument must be defined\",\n model=\"critical\",\n )\n return False\n elif args.execution == \"compare\":\n if not args.files:\n color_output(\n \"Missing argument. When the execution type is set to 'compare' the '--files' argument must be defined\",\n model=\"critical\",\n )\n return False\n elif len(args.files) > 2:\n color_output(\n \"Too many arguments. The '--files' argument cannot be repeated more than 2 times.\",\n model=\"critical\",\n )\n return False\n return True\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"-a\",\n \"--auth\",\n type=str,\n default=None,\n help=(\n \"Currently, to get this script working when the Onyx Auth is \"\n \"enabled, you must extract from the UI your cookie 'fastapiusersauth' \"\n \"and then set it using this argument\"\n ),\n )\n parser.add_argument(\n \"-e\",\n \"--execution\",\n type=str,\n choices=[\"new\", \"compare\"],\n default=None,\n help=(\n \"The execution type. Must be 'new' to generate a new analysis file \"\n \"or 'compare' to compare a previous execution with a new one based on the same queries\"\n ),\n )\n parser.add_argument(\n \"-f\",\n \"--files\",\n action=\"extend\",\n default=[],\n nargs=1,\n help=(\n \"Analysis file(s) to use for the comparison. Required if the execution arg is set \"\n \"to 'compare'. NOTE: By repeating this argument, you can make a comparison between \"\n \"two specific executions. If not repeated, a new execution will be performed and \"\n \"compared with the selected one.\"\n ),\n )\n parser.add_argument(\n \"-p\",\n \"--port\",\n type=int,\n default=3000,\n help=(\n \"The Onyx Web (not the API) port. We use the UI to forward the requests to the API. \"\n \"It should be '3000' for local dev and '80' if Onyx runs using docker compose.\"\n ),\n )\n parser.add_argument(\n \"-q--queries\",\n type=str,\n action=\"extend\",\n default=[],\n nargs=1,\n help=(\n \"The query to evaluate. Required if the execution arg is set to 'new'. \"\n \"NOTE: This argument can be repeated multiple times\"\n ),\n )\n parser.add_argument(\n \"-t\",\n \"--threshold\",\n type=float,\n default=0.0,\n help=\"The minimum score change (percentage) to detect an issue.\",\n )\n parser.add_argument(\n \"-w\",\n \"--wait\",\n type=int,\n default=10,\n help=(\n \"The waiting time (in seconds) to respect between queries. \"\n \"It is helpful to avoid hitting the Generative AI rate limiting.\"\n ),\n )\n\n args = parser.parse_args()\n if not validate_cmd_args(args):\n sys.exit(1)\n\n SelectionAnalysis(\n args.execution,\n args.files,\n args.q__queries,\n args.threshold,\n args.port,\n args.auth,\n args.wait,\n )()\n" }, { "path": "backend/scripts/supervisord_entrypoint.sh", "content": "#!/bin/sh\n# Entrypoint script for supervisord\n\n# Launch supervisord with environment variables available\nexec /usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf\n" }, { "path": "backend/scripts/tenant_cleanup/QUICK_START_NO_BASTION.md", "content": "# Quick Start: Tenant Cleanup Without Bastion\n\n## TL;DR - The Commands You Need\n\n```bash\n# Navigate to backend directory\ncd onyx/backend\n\n# Step 1: Generate CSV of tenants to clean (5-10 min)\nPYTHONPATH=. python scripts/tenant_cleanup/no_bastion_analyze_tenants.py\n\n# Step 2: Mark connectors for deletion (1-2 min)\nPYTHONPATH=. python scripts/tenant_cleanup/no_bastion_mark_connectors.py \\\n --csv gated_tenants_no_query_3mo_*.csv \\\n --force \\\n --concurrency 16\n\n# ⏰ WAIT 6+ hours for background deletion to complete\n\n# Step 3: Final cleanup (1-2 min)\nPYTHONPATH=. python scripts/tenant_cleanup/no_bastion_cleanup_tenants.py \\\n --csv gated_tenants_no_query_3mo_*.csv \\\n --force\n```\n\n## What Changed?\n\nInstead of the original scripts that require bastion access:\n- `analyze_current_tenants.py` → `no_bastion_analyze_tenants.py`\n- `mark_connectors_for_deletion.py` → `no_bastion_mark_connectors.py`\n- `cleanup_tenants.py` → `no_bastion_cleanup_tenants.py`\n\n**No environment variables needed!** All queries run directly from pods.\n\n## What You Need\n\n✅ `kubectl` access to your cluster\n✅ Running `celery-worker-user-file-processing` pods\n✅ Permission to exec into pods\n\n❌ No bastion host required\n❌ No SSH keys required\n❌ No environment variables required\n\n## Test Your Setup\n\n```bash\n# Check if you can find worker pods\nkubectl get po | grep celery-worker-user-file-processing | grep Running\n\n# If you see pods, you're ready to go!\n```\n\n## Important Notes\n\n1. **Step 2 triggers background deletion** - the actual document deletion happens asynchronously via Celery workers\n2. **You MUST wait** between Step 2 and Step 3 for deletion to complete (can take 6+ hours)\n3. **Monitor deletion progress** with: `kubectl logs -f `\n4. **All scripts verify tenant status** - they'll refuse to process active (non-GATED_ACCESS) tenants\n\n## Files Generated\n\n- `gated_tenants_no_query_3mo_YYYYMMDD_HHMMSS.csv` - List of tenants to clean\n- `cleaned_tenants.csv` - Successfully cleaned tenants with timestamps\n\n## Safety First\n\nThe scripts include multiple safety checks:\n- ✅ Verifies tenant status before any operation\n- ✅ Checks documents are deleted before dropping schemas\n- ✅ Prompts for confirmation on dangerous operations (unless `--force`)\n- ✅ Records all successful operations in real-time\n\n## Need More Details?\n\nSee [NO_BASTION_README.md](./NO_BASTION_README.md) for:\n- Detailed explanations of each step\n- Troubleshooting guide\n- How it works under the hood\n- Performance characteristics\n" }, { "path": "backend/scripts/tenant_cleanup/README.md", "content": "## How to Tenant Cleanup\n\nThree main steps.\n\n### Build a list of tenants to cleanup\n\nUse the `analyze_current_tenants.py` script:\n\n```\nPYTHONPATH=. \\\nCONTROL_PLANE_RDS_HOST= \\\nCONTROL_PLANE_RDS_PASSWORD= \\\nBASTION_HOST= \\\nPEM_FILE_LOCATION= \\\npython scripts/tenant_cleanup/analyze_current_tenants.py\n```\n\nThis will create a `.csv` called something like `gated_tenants_no_query_3mo_20251012_161102.csv` in the `backend` dir.\n\n\n### Delete all documents within these tenants\n\nUse the `mark_connectors_for_deletion.py` script:\n\n```\nPYTHONPATH=. \\\nCONTROL_PLANE_RDS_HOST= \\\nCONTROL_PLANE_RDS_PASSWORD= \\\nBASTION_HOST= \\\nPEM_FILE_LOCATION= \\\npython scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo_.csv --force\n```\n\nReplace `gated_tenants_no_query_3mo_.csv` with the CSV name from step (1).\n\nThis will update the data plane database to 1/ cancel all index attempts 2/ mark all connectors as up for deletion.\nWe now need to wait for the deletion to run.\n\nIt's done this way to re-use as much of the existing code + take advantage of existing infra for parallelized, long running jobs. These \ndeletion jobs can take a LONG time (>6hrs), so having it performed syncronously by a script is not really tenable.\n\n\n### Cleanup the tenants\n\nUse the `cleanup_tenants.py` script:\n\n```\nPYTHONPATH=. \\\nCONTROL_PLANE_RDS_HOST= \\\nCONTROL_PLANE_RDS_PASSWORD= \\\nBASTION_HOST= \\\nPEM_FILE_LOCATION= \\\npython scripts/tenant_cleanup/cleanup_tenants.py --csv gated_tenants_no_query_3mo_.csv --force\n```\n\nThis will drop the tenant schema from the data plane DB, cleanup the `user_tenant_mapping` table, and \nclean up any control plane DB tables associated with each tenant.\n\nNOTE: if the previous step has not completed, tenants with documents will throw an exception.\n" }, { "path": "backend/scripts/tenant_cleanup/analyze_current_tenants.py", "content": "#!/usr/bin/env python3\n\"\"\"\nFull tenant analysis script that:\n1. Finds a heavy worker pod\n2. Runs the tenant data collection script on the pod\n3. Analyzes the collected data\n\"\"\"\n\nimport argparse\nimport csv\nimport json\nimport os\nimport subprocess\nimport sys\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom pathlib import Path\nfrom typing import Any\n\nfrom scripts.tenant_cleanup.cleanup_utils import find_worker_pod\n\n\ndef collect_tenant_data(pod_name: str) -> list[dict[str, Any]]:\n \"\"\"Run the understand_tenants script on the pod and return the data.\"\"\"\n print(f\"\\nCollecting tenant data from pod {pod_name}...\")\n\n # Get the path to the understand_tenants script\n script_dir = Path(__file__).parent\n understand_tenants_script = script_dir / \"on_pod_scripts\" / \"understand_tenants.py\"\n\n if not understand_tenants_script.exists():\n raise FileNotFoundError(\n f\"understand_tenants.py not found at {understand_tenants_script}\"\n )\n\n # Copy script to pod\n print(\"Copying script to pod...\")\n subprocess.run(\n [\n \"kubectl\",\n \"cp\",\n str(understand_tenants_script),\n f\"{pod_name}:/tmp/understand_tenants.py\",\n ],\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n print(\"Executing script on pod (this may take a while)...\")\n result = subprocess.run(\n [\"kubectl\", \"exec\", pod_name, \"--\", \"python\", \"/tmp/understand_tenants.py\"],\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(result.stderr, file=sys.stderr)\n\n # Parse JSON from stdout\n try:\n tenant_data = json.loads(result.stdout)\n print(f\"Successfully collected data for {len(tenant_data)} tenants\")\n return tenant_data\n except json.JSONDecodeError as e:\n print(f\"Failed to parse JSON output: {e}\", file=sys.stderr)\n print(f\"stdout: {result.stdout[:500]}\", file=sys.stderr)\n raise\n\n\ndef collect_control_plane_data() -> list[dict[str, Any]]:\n \"\"\"Collect control plane data from the control plane database.\"\"\"\n print(\"\\nCollecting control plane data...\")\n\n rds_host = os.environ.get(\"CONTROL_PLANE_RDS_HOST\")\n if not rds_host:\n raise ValueError(\"CONTROL_PLANE_RDS_HOST is not set\")\n\n rds_password = os.environ.get(\"CONTROL_PLANE_RDS_PASSWORD\")\n if not rds_password:\n raise ValueError(\"CONTROL_PLANE_RDS_PASSWORD is not set\")\n\n db_url = f\"postgresql://postgres:{rds_password}@{rds_host}:5432/control\"\n\n bastion_host = os.environ.get(\"BASTION_HOST\")\n if not bastion_host:\n raise ValueError(\"BASTION_HOST is not set\")\n\n pem_file_location = os.environ.get(\"PEM_FILE_LOCATION\")\n if not pem_file_location:\n raise ValueError(\"PEM_FILE_LOCATION is not set\")\n\n full_cmd = (\n f\"ssh -i {pem_file_location} ec2-user@{bastion_host} \"\n f\"\\\"psql {db_url} -c '\\\\copy (SELECT * FROM tenant) \"\n f\"to '/tmp/control_plane_data.csv' with (format csv);'\\\"\"\n )\n\n result = subprocess.run(\n full_cmd,\n shell=True,\n check=True,\n capture_output=True,\n text=True,\n )\n\n # Copy the CSV file from the bastion to local machine\n copy_cmd = f\"scp -i {pem_file_location} ec2-user@{bastion_host}:/tmp/control_plane_data.csv .\"\n\n copy_result = subprocess.run(\n copy_cmd, shell=True, check=True, capture_output=True, text=True\n )\n\n if copy_result.stderr:\n print(f\"Copy warnings: {copy_result.stderr}\", file=sys.stderr)\n raise RuntimeError(\n \"Failed to copy control plane data from bastion to local machine\"\n )\n\n print(\"Control plane data copied to local machine as control_plane_data.csv\")\n\n print(result.stdout)\n\n # Read the CSV file and convert to list of dictionaries\n control_plane_data = []\n with open(\"control_plane_data.csv\", \"r\", newline=\"\", encoding=\"utf-8\") as csvfile:\n reader = csv.DictReader(csvfile)\n for row in reader:\n control_plane_data.append(row)\n\n return control_plane_data\n\n\ndef analyze_tenants(\n tenants: list[dict[str, Any]], control_plane_data: list[dict[str, Any]]\n) -> list[dict[str, Any]]:\n \"\"\"Analyze tenant activity data and return gated tenants with no query in last 3 months.\"\"\"\n\n print(f\"\\n{'=' * 80}\")\n print(f\"TENANT ANALYSIS REPORT - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\")\n print(f\"{'=' * 80}\")\n print(f\"Total tenants analyzed: {len(tenants)}\\n\")\n\n # Create a lookup dict for control plane data by tenant_id\n control_plane_lookup = {}\n for row in control_plane_data:\n # CSV has no header, columns are: tenant_id, stripe_customer_id, created_at,\n # stripe_subscription_quantity, contact_email, registration_origin, tenant_status\n if len(row) >= 7:\n tenant_id = list(row.values())[0] # First column is tenant_id\n tenant_status = list(row.values())[6] # 7th column is tenant_status\n control_plane_lookup[tenant_id] = tenant_status\n\n # Calculate cutoff dates\n one_month_cutoff = datetime.now(timezone.utc) - timedelta(days=30)\n three_month_cutoff = datetime.now(timezone.utc) - timedelta(days=90)\n\n # Categorize tenants into 4 groups\n gated_no_query_3_months = [] # GATED_ACCESS + no query in last 3 months\n gated_query_1_3_months = [] # GATED_ACCESS + query between 1-3 months\n gated_query_1_month = [] # GATED_ACCESS + query in last 1 month\n everyone_else = [] # All other tenants\n\n for tenant in tenants:\n tenant_id = tenant.get(\"tenant_id\")\n last_query_time = tenant.get(\"last_query_time\")\n tenant_status = control_plane_lookup.get(tenant_id, \"UNKNOWN\")\n\n is_gated = tenant_status == \"GATED_ACCESS\"\n\n # Parse last query time\n if last_query_time:\n query_time = datetime.fromisoformat(last_query_time.replace(\"Z\", \"+00:00\"))\n else:\n query_time = None\n\n # Categorize\n if is_gated:\n if query_time is None or query_time <= three_month_cutoff:\n gated_no_query_3_months.append(tenant)\n elif query_time <= one_month_cutoff:\n gated_query_1_3_months.append(tenant)\n else: # query_time > one_month_cutoff\n gated_query_1_month.append(tenant)\n else:\n everyone_else.append(tenant)\n\n # Calculate document counts for each group\n gated_no_query_docs = sum(\n t.get(\"num_documents\", 0) for t in gated_no_query_3_months\n )\n gated_1_3_month_docs = sum(\n t.get(\"num_documents\", 0) for t in gated_query_1_3_months\n )\n gated_1_month_docs = sum(t.get(\"num_documents\", 0) for t in gated_query_1_month)\n everyone_else_docs = sum(t.get(\"num_documents\", 0) for t in everyone_else)\n\n print(\"=\" * 80)\n print(\"TENANT CATEGORIZATION BY GATED ACCESS STATUS AND ACTIVITY\")\n print(\"=\" * 80)\n\n print(\"\\n1. GATED_ACCESS + No query in last 3 months:\")\n print(f\" Count: {len(gated_no_query_3_months):,}\")\n print(f\" Total documents: {gated_no_query_docs:,}\")\n print(\n f\" Avg documents per tenant: {gated_no_query_docs / len(gated_no_query_3_months) if gated_no_query_3_months else 0:.2f}\"\n )\n\n print(\"\\n2. GATED_ACCESS + Query between 1-3 months ago:\")\n print(f\" Count: {len(gated_query_1_3_months):,}\")\n print(f\" Total documents: {gated_1_3_month_docs:,}\")\n print(\n f\" Avg documents per tenant: {gated_1_3_month_docs / len(gated_query_1_3_months) if gated_query_1_3_months else 0:.2f}\"\n )\n\n print(\"\\n3. GATED_ACCESS + Query in last 1 month:\")\n print(f\" Count: {len(gated_query_1_month):,}\")\n print(f\" Total documents: {gated_1_month_docs:,}\")\n print(\n f\" Avg documents per tenant: {gated_1_month_docs / len(gated_query_1_month) if gated_query_1_month else 0:.2f}\"\n )\n\n print(\"\\n4. Everyone else (non-GATED_ACCESS):\")\n print(f\" Count: {len(everyone_else):,}\")\n print(f\" Total documents: {everyone_else_docs:,}\")\n print(\n f\" Avg documents per tenant: {everyone_else_docs / len(everyone_else) if everyone_else else 0:.2f}\"\n )\n\n total_docs = (\n gated_no_query_docs\n + gated_1_3_month_docs\n + gated_1_month_docs\n + everyone_else_docs\n )\n print(f\"\\nTotal documents across all tenants: {total_docs:,}\")\n\n # Top 100 tenants by document count\n print(\"\\n\" + \"=\" * 80)\n print(\"TOP 100 TENANTS BY DOCUMENT COUNT\")\n print(\"=\" * 80)\n\n # Sort all tenants by document count\n sorted_tenants = sorted(\n tenants, key=lambda t: t.get(\"num_documents\", 0), reverse=True\n )\n\n top_100 = sorted_tenants[:100]\n\n print(\n f\"\\n{'Rank':<6} {'Tenant ID':<45} {'Documents':>12} {'Users':>8} {'Last Query':<12} {'Group'}\"\n )\n print(\"-\" * 130)\n\n for idx, tenant in enumerate(top_100, 1):\n tenant_id = tenant.get(\"tenant_id\", \"Unknown\")\n num_docs = tenant.get(\"num_documents\", 0)\n num_users = tenant.get(\"num_users\", 0)\n last_query = tenant.get(\"last_query_time\", \"Never\")\n tenant_status = control_plane_lookup.get(tenant_id, \"UNKNOWN\")\n\n # Format the last query time\n if last_query and last_query != \"Never\":\n try:\n query_dt = datetime.fromisoformat(last_query.replace(\"Z\", \"+00:00\"))\n last_query_str = query_dt.strftime(\"%Y-%m-%d\")\n except Exception:\n last_query_str = last_query[:10] if len(last_query) > 10 else last_query\n else:\n last_query_str = \"Never\"\n\n # Determine group\n if tenant_status == \"GATED_ACCESS\":\n if last_query and last_query != \"Never\":\n query_time = datetime.fromisoformat(last_query.replace(\"Z\", \"+00:00\"))\n if query_time <= three_month_cutoff:\n group = \"Gated - No query (3mo)\"\n elif query_time <= one_month_cutoff:\n group = \"Gated - Query (1-3mo)\"\n else:\n group = \"Gated - Query (1mo)\"\n else:\n group = \"Gated - No query (3mo)\"\n else:\n group = f\"Other ({tenant_status})\"\n\n print(\n f\"{idx:<6} {tenant_id:<45} {num_docs:>12,} {num_users:>8} {last_query_str:<12} {group}\"\n )\n\n # Summary stats for top 100\n top_100_docs = sum(t.get(\"num_documents\", 0) for t in top_100)\n\n print(\"\\n\" + \"-\" * 110)\n print(f\"Top 100 total documents: {top_100_docs:,}\")\n print(\n f\"Percentage of all documents: {(top_100_docs / total_docs * 100) if total_docs > 0 else 0:.2f}%\"\n )\n\n # Additional insights\n print(\"\\n\" + \"=\" * 80)\n print(\"ADDITIONAL INSIGHTS\")\n print(\"=\" * 80)\n\n # Tenants with no documents\n no_docs = [t for t in tenants if t.get(\"num_documents\", 0) == 0]\n print(\n f\"\\nTenants with 0 documents: {len(no_docs):,} ({len(no_docs) / len(tenants) * 100:.2f}%)\"\n )\n\n # Tenants with no users\n no_users = [t for t in tenants if t.get(\"num_users\", 0) == 0]\n print(\n f\"Tenants with 0 users: {len(no_users):,} ({len(no_users) / len(tenants) * 100:.2f}%)\"\n )\n\n # Document distribution quartiles\n doc_counts = sorted([t.get(\"num_documents\", 0) for t in tenants])\n if doc_counts:\n print(\"\\nDocument count distribution:\")\n print(f\" Median: {doc_counts[len(doc_counts) // 2]:,}\")\n print(f\" 75th percentile: {doc_counts[int(len(doc_counts) * 0.75)]:,}\")\n print(f\" 90th percentile: {doc_counts[int(len(doc_counts) * 0.90)]:,}\")\n print(f\" 95th percentile: {doc_counts[int(len(doc_counts) * 0.95)]:,}\")\n print(f\" 99th percentile: {doc_counts[int(len(doc_counts) * 0.99)]:,}\")\n print(f\" Max: {doc_counts[-1]:,}\")\n\n return gated_no_query_3_months\n\n\ndef find_recent_tenant_data() -> tuple[list[dict[str, Any]] | None, str | None]:\n \"\"\"Find the most recent tenant data file if it's less than 7 days old.\"\"\"\n current_dir = Path.cwd()\n tenant_data_files = list(current_dir.glob(\"tenant_data_*.json\"))\n\n if not tenant_data_files:\n return None, None\n\n # Sort by modification time, most recent first\n tenant_data_files.sort(key=lambda p: p.stat().st_mtime, reverse=True)\n most_recent = tenant_data_files[0]\n\n # Check if file is less than 7 days old\n file_age = datetime.now().timestamp() - most_recent.stat().st_mtime\n seven_days_in_seconds = 7 * 24 * 60 * 60\n\n if file_age < seven_days_in_seconds:\n file_age_days = file_age / (24 * 60 * 60)\n print(\n f\"\\n✓ Found recent tenant data: {most_recent.name} (age: {file_age_days:.1f} days)\"\n )\n\n with open(most_recent, \"r\") as f:\n tenant_data = json.load(f)\n\n return tenant_data, str(most_recent)\n\n return None, None\n\n\ndef main() -> None:\n # Parse command-line arguments\n parser = argparse.ArgumentParser(\n description=\"Analyze tenant data and identify gated tenants with no recent queries\"\n )\n parser.add_argument(\n \"--skip-cache\",\n action=\"store_true\",\n help=\"Skip cached tenant data and collect fresh data from pod\",\n )\n args = parser.parse_args()\n\n try:\n # Step 0: Collect control plane data\n control_plane_data = collect_control_plane_data()\n\n # Step 1: Check for recent tenant data (< 7 days old) unless --skip-cache is set\n tenant_data = None\n cached_file = None\n\n if not args.skip_cache:\n tenant_data, cached_file = find_recent_tenant_data()\n\n if tenant_data:\n print(f\"Using cached tenant data from: {cached_file}\")\n print(f\"Total tenants in cache: {len(tenant_data)}\")\n else:\n if args.skip_cache:\n print(\"\\n⚠ Skipping cache (--skip-cache flag set)\")\n\n # Step 2a: Find the heavy worker pod\n pod_name = find_worker_pod()\n\n # Step 2b: Collect tenant data\n tenant_data = collect_tenant_data(pod_name)\n\n # Step 2c: Save raw data to file with timestamp\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n output_file = f\"tenant_data_{timestamp}.json\"\n with open(output_file, \"w\") as f:\n json.dump(tenant_data, f, indent=2, default=str)\n print(f\"\\n✓ Raw data saved to: {output_file}\")\n\n # Step 3: Analyze the data and get gated tenants without recent queries\n gated_no_query_3_months = analyze_tenants(tenant_data, control_plane_data)\n\n # Step 4: Export to CSV (sorted by num_documents descending)\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n csv_file = f\"gated_tenants_no_query_3mo_{timestamp}.csv\"\n\n # Sort by num_documents in descending order\n sorted_tenants = sorted(\n gated_no_query_3_months,\n key=lambda t: t.get(\"num_documents\", 0),\n reverse=True,\n )\n\n with open(csv_file, \"w\", newline=\"\", encoding=\"utf-8\") as csvfile:\n fieldnames = [\n \"tenant_id\",\n \"num_documents\",\n \"num_users\",\n \"last_query_time\",\n \"days_since_last_query\",\n ]\n writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n writer.writeheader()\n\n now = datetime.now(timezone.utc)\n for tenant in sorted_tenants:\n # Calculate days since last query\n last_query_time = tenant.get(\"last_query_time\")\n if last_query_time:\n try:\n query_dt = datetime.fromisoformat(\n last_query_time.replace(\"Z\", \"+00:00\")\n )\n days_since = str((now - query_dt).days)\n except Exception:\n days_since = \"N/A\"\n else:\n days_since = \"Never\"\n\n writer.writerow(\n {\n \"tenant_id\": tenant.get(\"tenant_id\", \"\"),\n \"num_documents\": tenant.get(\"num_documents\", 0),\n \"num_users\": tenant.get(\"num_users\", 0),\n \"last_query_time\": last_query_time or \"Never\",\n \"days_since_last_query\": days_since,\n }\n )\n\n print(f\"\\n✓ CSV exported to: {csv_file}\")\n print(\n f\" Total gated tenants with no query in last 3 months: {len(gated_no_query_3_months)}\"\n )\n\n except subprocess.CalledProcessError as e:\n print(f\"Error running command: {e}\", file=sys.stderr)\n if e.stderr:\n print(f\"stderr: {e.stderr}\", file=sys.stderr)\n sys.exit(1)\n except Exception as e:\n print(f\"Error: {e}\", file=sys.stderr)\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/check_no_bastion_setup.py", "content": "#!/usr/bin/env python3\n\"\"\"\nVerification script to check if your environment is ready for no-bastion tenant cleanup.\n\nUsage:\n python scripts/tenant_cleanup/check_no_bastion_setup.py\n\"\"\"\n\nimport subprocess\nimport sys\n\n\ndef print_header(text: str) -> None:\n \"\"\"Print a formatted header.\"\"\"\n print(f\"\\n{'=' * 80}\")\n print(f\" {text}\")\n print(f\"{'=' * 80}\\n\")\n\n\ndef check_kubectl_access() -> bool:\n \"\"\"Check if kubectl is installed and can access the cluster.\"\"\"\n print(\"Checking kubectl access...\")\n\n try:\n result = subprocess.run(\n [\"kubectl\", \"version\", \"--client\", \"--short\"],\n capture_output=True,\n text=True,\n timeout=5,\n )\n\n if result.returncode == 0:\n print(f\"✅ kubectl is installed: {result.stdout.strip()}\")\n\n # Try to access cluster\n result = subprocess.run(\n [\"kubectl\", \"get\", \"ns\"],\n capture_output=True,\n text=True,\n timeout=10,\n )\n\n if result.returncode == 0:\n print(\"✅ kubectl can access the cluster\")\n return True\n else:\n print(\"❌ kubectl cannot access the cluster\")\n print(f\" Error: {result.stderr}\")\n return False\n else:\n print(\"❌ kubectl is not installed or not in PATH\")\n return False\n\n except FileNotFoundError:\n print(\"❌ kubectl is not installed\")\n return False\n except subprocess.TimeoutExpired:\n print(\"❌ kubectl command timed out\")\n return False\n except Exception as e:\n print(f\"❌ Error checking kubectl: {e}\")\n return False\n\n\ndef check_worker_pods() -> tuple[bool, list[str]]:\n \"\"\"Check if worker pods are running.\"\"\"\n print(\"\\nChecking for worker pods...\")\n\n try:\n result = subprocess.run(\n [\"kubectl\", \"get\", \"po\"],\n capture_output=True,\n text=True,\n timeout=10,\n check=True,\n )\n\n lines = result.stdout.strip().split(\"\\n\")\n worker_pods = []\n\n for line in lines[1:]: # Skip header\n if \"celery-worker-user-file-processing\" in line and \"Running\" in line:\n pod_name = line.split()[0]\n worker_pods.append(pod_name)\n\n if worker_pods:\n print(f\"✅ Found {len(worker_pods)} running worker pod(s):\")\n for pod in worker_pods[:3]: # Show first 3\n print(f\" - {pod}\")\n if len(worker_pods) > 3:\n print(f\" ... and {len(worker_pods) - 3} more\")\n return True, worker_pods\n else:\n print(\"❌ No running celery-worker-user-file-processing pods found\")\n print(\" Available pods:\")\n for line in lines[1:6]: # Show first 5 pods\n print(f\" {line}\")\n return False, []\n\n except subprocess.CalledProcessError as e:\n print(f\"❌ Error getting pods: {e}\")\n return False, []\n except Exception as e:\n print(f\"❌ Error checking worker pods: {e}\")\n return False, []\n\n\ndef check_pod_exec_permission(pod_name: str) -> bool:\n \"\"\"Check if we can exec into a pod.\"\"\"\n print(\"\\nChecking pod exec permissions...\")\n\n try:\n result = subprocess.run(\n [\"kubectl\", \"exec\", pod_name, \"--\", \"echo\", \"test\"],\n capture_output=True,\n text=True,\n timeout=10,\n )\n\n if result.returncode == 0 and \"test\" in result.stdout:\n print(f\"✅ Can exec into pod: {pod_name}\")\n return True\n else:\n print(f\"❌ Cannot exec into pod: {pod_name}\")\n print(f\" Error: {result.stderr}\")\n return False\n\n except subprocess.TimeoutExpired:\n print(f\"❌ Exec command timed out for pod: {pod_name}\")\n return False\n except Exception as e:\n print(f\"❌ Error checking exec permission: {e}\")\n return False\n\n\ndef check_pod_db_access(pod_name: str) -> dict:\n \"\"\"Check if pod has database environment variables.\"\"\"\n print(\"\\nChecking database access from pod...\")\n\n checks = {\n \"control_plane\": False,\n \"data_plane\": False,\n }\n\n try:\n # Check for control plane DB env vars\n result = subprocess.run(\n [\"kubectl\", \"exec\", pod_name, \"--\", \"env\"],\n capture_output=True,\n text=True,\n timeout=10,\n )\n\n if result.returncode == 0:\n env_output = result.stdout\n\n # Check control plane access\n if any(\n var in env_output\n for var in [\n \"POSTGRES_CONTROL_URI\",\n \"POSTGRES_CONTROL_HOST\",\n ]\n ):\n print(\"✅ Pod has control plane database environment variables\")\n checks[\"control_plane\"] = True\n else:\n print(\n \"⚠️ Pod may not have control plane database environment variables\"\n )\n print(\" (This might be okay if they're dynamically loaded)\")\n\n # Check data plane access\n if any(\n var in env_output\n for var in [\"POSTGRES_URI\", \"POSTGRES_HOST\", \"DATABASE_URL\"]\n ):\n print(\"✅ Pod has data plane database environment variables\")\n checks[\"data_plane\"] = True\n else:\n print(\"❌ Pod does not have data plane database environment variables\")\n\n return checks\n\n except Exception as e:\n print(f\"❌ Error checking database access: {e}\")\n return checks\n\n\ndef check_required_scripts() -> bool:\n \"\"\"Check if the required on_pod_scripts exist.\"\"\"\n print(\"\\nChecking for required scripts...\")\n\n from pathlib import Path\n\n script_dir = Path(__file__).parent\n required_scripts = [\n \"on_pod_scripts/understand_tenants.py\",\n \"on_pod_scripts/execute_connector_deletion.py\",\n \"on_pod_scripts/check_documents_deleted.py\",\n \"on_pod_scripts/cleanup_tenant_schema.py\",\n \"on_pod_scripts/get_tenant_index_name.py\",\n \"on_pod_scripts/get_tenant_users.py\",\n ]\n\n all_exist = True\n for script in required_scripts:\n script_path = script_dir / script\n if script_path.exists():\n print(f\"✅ {script}\")\n else:\n print(f\"❌ {script} - NOT FOUND\")\n all_exist = False\n\n return all_exist\n\n\ndef main() -> None:\n print_header(\"No-Bastion Tenant Cleanup - Setup Verification\")\n\n all_checks_passed = True\n\n # 1. Check kubectl access\n if not check_kubectl_access():\n all_checks_passed = False\n\n # 2. Check for worker pods\n has_pods, worker_pods = check_worker_pods()\n if not has_pods:\n all_checks_passed = False\n print(\"\\n⚠️ Cannot proceed without running worker pods\")\n print_header(\"SETUP VERIFICATION FAILED\")\n sys.exit(1)\n\n # Use first worker pod for remaining checks\n test_pod = worker_pods[0]\n\n # 3. Check exec permissions\n if not check_pod_exec_permission(test_pod):\n all_checks_passed = False\n\n # 4. Check database access\n db_checks = check_pod_db_access(test_pod)\n if not db_checks[\"data_plane\"]:\n all_checks_passed = False\n\n # 5. Check required scripts\n if not check_required_scripts():\n all_checks_passed = False\n\n # Summary\n print_header(\"VERIFICATION SUMMARY\")\n\n if all_checks_passed and db_checks[\"control_plane\"]:\n print(\"✅ ALL CHECKS PASSED!\")\n print(\"\\nYou're ready to run tenant cleanup without bastion access.\")\n print(\"\\nNext steps:\")\n print(\"1. Read QUICK_START_NO_BASTION.md for commands\")\n print(\n \"2. Run: PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_analyze_tenants.py\"\n )\n sys.exit(0)\n elif all_checks_passed:\n print(\"⚠️ MOSTLY READY (with warnings)\")\n print(\"\\nYou can proceed, but control plane access may need verification.\")\n print(\"Try running Step 1 and see if it works.\")\n print(\"\\nNext steps:\")\n print(\n \"1. Run: PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_analyze_tenants.py\"\n )\n print(\"2. If it fails with DB errors, check pod environment variables\")\n sys.exit(0)\n else:\n print(\"❌ SETUP VERIFICATION FAILED\")\n print(\"\\nPlease fix the issues above before proceeding.\")\n print(\"\\nCommon fixes:\")\n print(\"- Install kubectl: https://kubernetes.io/docs/tasks/tools/\")\n print(\"- Configure cluster access: kubectl config use-context \")\n print(\"- Check pod status: kubectl get po\")\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/cleanup_tenants.py", "content": "#!/usr/bin/env python3\n\"\"\"\nTenant cleanup script that:\n1. Deletes all documents from Vespa\n2. Drops the data plane PostgreSQL schema\n3. Clean up control plane (tenants, subscription table)\n\nUsage:\n python backend/scripts/cleanup_tenant.py [--force]\n python backend/scripts/cleanup_tenant.py --csv [--force]\n\nArguments:\n tenant_id The tenant ID to clean up (required if not using --csv)\n --csv PATH Path to CSV file containing tenant IDs to clean up\n --force Skip all confirmation prompts (optional)\n\nExamples:\n python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789\n python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789 --force\n python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv\n python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv --force\n\"\"\"\n\nimport csv\nimport json\nimport signal\nimport subprocess\nimport sys\nfrom datetime import datetime\nfrom pathlib import Path\n\nfrom scripts.tenant_cleanup.cleanup_utils import confirm_step\nfrom scripts.tenant_cleanup.cleanup_utils import execute_control_plane_query\nfrom scripts.tenant_cleanup.cleanup_utils import find_worker_pod\nfrom scripts.tenant_cleanup.cleanup_utils import get_tenant_status\nfrom scripts.tenant_cleanup.cleanup_utils import read_tenant_ids_from_csv\nfrom scripts.tenant_cleanup.cleanup_utils import TenantNotFoundInControlPlaneError\n\n\ndef signal_handler(signum: int, frame: object) -> None: # noqa: ARG001\n \"\"\"Handle termination signals by killing active subprocess.\"\"\"\n sys.exit(1)\n\n\ndef get_tenant_index_name(pod_name: str, tenant_id: str) -> str:\n \"\"\"Get the default index name for the given tenant by running script on pod.\"\"\"\n print(f\"Getting default index name for tenant: {tenant_id}\")\n\n # Get the path to the script\n script_dir = Path(__file__).parent\n index_name_script = script_dir / \"on_pod_scripts\" / \"get_tenant_index_name.py\"\n\n if not index_name_script.exists():\n raise FileNotFoundError(\n f\"get_tenant_index_name.py not found at {index_name_script}\"\n )\n\n try:\n # Copy script to pod\n print(\" Copying script to pod...\")\n subprocess.run(\n [\n \"kubectl\",\n \"cp\",\n str(index_name_script),\n f\"{pod_name}:/tmp/get_tenant_index_name.py\",\n ],\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n print(\" Executing script on pod...\")\n result = subprocess.run(\n [\n \"kubectl\",\n \"exec\",\n pod_name,\n \"--\",\n \"python\",\n \"/tmp/get_tenant_index_name.py\",\n tenant_id,\n ],\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n\n if status == \"success\":\n index_name = result_data.get(\"index_name\")\n print(f\"✓ Found index name: {index_name}\")\n return index_name\n else:\n message = result_data.get(\"message\", \"Unknown error\")\n raise RuntimeError(f\"Failed to get index name: {message}\")\n\n except subprocess.CalledProcessError as e:\n print(\n f\"✗ Failed to get index name for tenant {tenant_id}: {e}\", file=sys.stderr\n )\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n print(\n f\"✗ Failed to get index name for tenant {tenant_id}: {e}\", file=sys.stderr\n )\n raise\n\n\ndef get_tenant_users(pod_name: str, tenant_id: str) -> list[str]:\n \"\"\"Get list of user emails from the tenant's data plane schema.\n\n Args:\n pod_name: The Kubernetes pod name to execute on\n tenant_id: The tenant ID to query\n\n Returns:\n List of user email addresses, or empty list if query fails\n \"\"\"\n print(f\"Fetching user emails for tenant: {tenant_id}\")\n\n # Get the path to the script\n script_dir = Path(__file__).parent\n get_users_script = script_dir / \"on_pod_scripts\" / \"get_tenant_users.py\"\n\n if not get_users_script.exists():\n raise FileNotFoundError(f\"get_tenant_users.py not found at {get_users_script}\")\n\n try:\n # Copy script to pod\n print(\" Copying script to pod...\")\n subprocess.run(\n [\n \"kubectl\",\n \"cp\",\n str(get_users_script),\n f\"{pod_name}:/tmp/get_tenant_users.py\",\n ],\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n print(\" Executing script on pod...\")\n result = subprocess.run(\n [\n \"kubectl\",\n \"exec\",\n pod_name,\n \"--\",\n \"python\",\n \"/tmp/get_tenant_users.py\",\n tenant_id,\n ],\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n\n if status == \"success\":\n users = result_data.get(\"users\", [])\n if users:\n print(f\"✓ Found {len(users)} user(s):\")\n for email in users:\n print(f\" - {email}\")\n else:\n print(\" No users found in tenant\")\n return users\n else:\n message = result_data.get(\"message\", \"Unknown error\")\n print(f\"⚠ Could not fetch users: {message}\")\n return []\n\n except subprocess.CalledProcessError as e:\n print(f\"⚠ Failed to get users for tenant {tenant_id}: {e}\")\n if e.stderr:\n print(f\" Error details: {e.stderr}\")\n return []\n except Exception as e:\n print(f\"⚠ Failed to get users for tenant {tenant_id}: {e}\")\n return []\n\n\ndef check_documents_deleted(pod_name: str, tenant_id: str) -> None:\n \"\"\"Check if all documents and connector credential pairs have been deleted.\n\n Raises RuntimeError if any ConnectorCredentialPairs or Documents remain.\n \"\"\"\n print(f\"Checking for remaining documents in tenant: {tenant_id}\")\n\n # Get the path to the script\n script_dir = Path(__file__).parent\n check_script = script_dir / \"on_pod_scripts\" / \"check_documents_deleted.py\"\n\n if not check_script.exists():\n raise FileNotFoundError(\n f\"check_documents_deleted.py not found at {check_script}\"\n )\n\n try:\n # Copy script to pod\n print(\" Copying script to pod...\")\n subprocess.run(\n [\n \"kubectl\",\n \"cp\",\n str(check_script),\n f\"{pod_name}:/tmp/check_documents_deleted.py\",\n ],\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n print(\" Executing check on pod...\")\n result = subprocess.run(\n [\n \"kubectl\",\n \"exec\",\n pod_name,\n \"--\",\n \"python\",\n \"/tmp/check_documents_deleted.py\",\n tenant_id,\n ],\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n\n if status == \"success\":\n message = result_data.get(\"message\")\n print(f\"✓ {message}\")\n elif status == \"not_found\":\n message = result_data.get(\"message\", \"Schema not found\")\n print(f\"⚠ {message}\")\n else:\n message = result_data.get(\"message\", \"Unknown error\")\n cc_count = result_data.get(\"connector_credential_pair_count\", 0)\n doc_count = result_data.get(\"document_count\", 0)\n error_details = f\"{message}\"\n if cc_count > 0 or doc_count > 0:\n error_details += f\"\\n ConnectorCredentialPairs: {cc_count}\\n Documents: {doc_count}\"\n raise RuntimeError(error_details)\n\n except subprocess.CalledProcessError as e:\n print(\n f\"✗ Failed to check documents for tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n print(\n f\"✗ Failed to check documents for tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n raise\n\n\ndef drop_data_plane_schema(pod_name: str, tenant_id: str) -> None:\n \"\"\"Drop the PostgreSQL schema for the given tenant by running script on pod.\"\"\"\n print(f\"Dropping data plane schema for tenant: {tenant_id}\")\n\n # Get the path to the cleanup script\n script_dir = Path(__file__).parent\n schema_cleanup_script = script_dir / \"on_pod_scripts\" / \"cleanup_tenant_schema.py\"\n\n if not schema_cleanup_script.exists():\n raise FileNotFoundError(\n f\"cleanup_tenant_schema.py not found at {schema_cleanup_script}\"\n )\n\n try:\n # Copy script to pod\n print(\" Copying script to pod...\")\n subprocess.run(\n [\n \"kubectl\",\n \"cp\",\n str(schema_cleanup_script),\n f\"{pod_name}:/tmp/cleanup_tenant_schema.py\",\n ],\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n print(\" Executing schema cleanup on pod...\")\n result = subprocess.run(\n [\n \"kubectl\",\n \"exec\",\n pod_name,\n \"--\",\n \"python\",\n \"/tmp/cleanup_tenant_schema.py\",\n tenant_id,\n ],\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n message = result_data.get(\"message\")\n\n if status == \"success\":\n print(f\"✓ {message}\")\n elif status == \"not_found\":\n print(f\"⚠ {message}\")\n else:\n print(f\"✗ {message}\", file=sys.stderr)\n raise RuntimeError(message)\n\n except subprocess.CalledProcessError as e:\n print(f\"✗ Failed to drop schema for tenant {tenant_id}: {e}\", file=sys.stderr)\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n print(f\"✗ Failed to drop schema for tenant {tenant_id}: {e}\", file=sys.stderr)\n raise\n\n\ndef cleanup_control_plane(tenant_id: str, force: bool = False) -> None:\n \"\"\"\n Clean up control plane data (tenants table, subscription table, etc.)\n\n Deletes from tables in this order:\n 1. tenant_notification (foreign key to tenant)\n 2. tenant_config (foreign key to tenant)\n 3. subscription (foreign key to tenant)\n 4. tenant (primary table)\n \"\"\"\n print(f\"Cleaning up control plane data for tenant: {tenant_id}\")\n\n # Delete in order respecting foreign key constraints\n delete_queries = [\n (\n \"tenant_notification\",\n \"DELETE FROM tenant_notification WHERE tenant_id = '{tenant_id}';\",\n ),\n (\"tenant_config\", \"DELETE FROM tenant_config WHERE tenant_id = '{tenant_id}';\"),\n (\"subscription\", \"DELETE FROM subscription WHERE tenant_id = '{tenant_id}';\"),\n (\"tenant\", \"DELETE FROM tenant WHERE tenant_id = '{tenant_id}';\"),\n ]\n\n try:\n for table_name, query in delete_queries:\n formatted_query = query.format(tenant_id=tenant_id)\n print(f\" Deleting from {table_name}...\")\n\n if not confirm_step(f\"Delete from {table_name}?\", force):\n print(f\" Skipping deletion from {table_name}\")\n continue\n\n result = execute_control_plane_query(formatted_query)\n\n if result.stdout:\n # Extract row count from output (e.g., \"DELETE 5\")\n print(f\" {result.stdout.strip()}\")\n\n print(f\"✓ Successfully cleaned up control plane data for tenant: {tenant_id}\")\n\n except subprocess.CalledProcessError as e:\n print(\n f\"✗ Failed to clean up control plane for tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n\n\ndef cleanup_tenant(tenant_id: str, pod_name: str, force: bool = False) -> bool:\n \"\"\"\n Main cleanup function that orchestrates all cleanup steps.\n\n Args:\n tenant_id: The tenant ID to clean up\n pod_name: The Kubernetes pod name to execute operations on\n force: If True, skip all confirmation prompts\n\n Returns:\n True if cleanup was performed, False if skipped\n \"\"\"\n print(f\"Starting cleanup for tenant: {tenant_id}\")\n\n # Track if tenant was not found in control plane (for force mode)\n tenant_not_found_in_control_plane = False\n\n # Check tenant status first\n print(f\"\\n{'=' * 80}\")\n try:\n tenant_status = get_tenant_status(tenant_id)\n\n # If tenant is not GATED_ACCESS, require explicit confirmation even in force mode\n if tenant_status and tenant_status != \"GATED_ACCESS\":\n print(\n f\"\\n⚠️ WARNING: Tenant status is '{tenant_status}', not 'GATED_ACCESS'!\"\n )\n print(\n \"This tenant may be active and should not be deleted without careful review.\"\n )\n print(f\"{'=' * 80}\\n\")\n\n if force:\n print(f\"Skipping cleanup for tenant {tenant_id} in force mode\")\n return False\n\n # Always ask for confirmation if not gated, even in force mode\n response = input(\n \"Are you ABSOLUTELY SURE you want to proceed? Type 'yes' to confirm: \"\n )\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - tenant is not GATED_ACCESS\")\n return False\n elif tenant_status == \"GATED_ACCESS\":\n print(\"✓ Tenant status is GATED_ACCESS - safe to proceed with cleanup\")\n elif tenant_status is None:\n print(\"⚠️ WARNING: Could not determine tenant status!\")\n\n if force:\n print(f\"Skipping cleanup for tenant {tenant_id} in force mode\")\n return False\n\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - could not verify tenant status\")\n return False\n except TenantNotFoundInControlPlaneError as e:\n # Tenant/table not found in control plane\n error_str = str(e)\n print(f\"⚠️ WARNING: Tenant not found in control plane: {error_str}\")\n tenant_not_found_in_control_plane = True\n\n if force:\n print(\n \"[FORCE MODE] Tenant not found in control plane - continuing with dataplane cleanup only\"\n )\n else:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - tenant not found in control plane\")\n return False\n except Exception as e:\n # Other errors (not \"not found\")\n error_str = str(e)\n print(f\"⚠️ WARNING: Failed to check tenant status: {error_str}\")\n\n if force:\n print(f\"Skipping cleanup for tenant {tenant_id} in force mode\")\n return False\n\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - could not verify tenant status\")\n return False\n print(f\"{'=' * 80}\\n\")\n\n # Fetch tenant users for informational purposes (non-blocking)\n # Skip in force mode as it's only informational\n if not force:\n print(f\"\\n{'=' * 80}\")\n try:\n get_tenant_users(pod_name, tenant_id)\n except Exception as e:\n print(f\"⚠ Could not fetch tenant users: {e}\")\n print(f\"{'=' * 80}\\n\")\n\n # Step 1: Make sure all documents are deleted\n print(f\"\\n{'=' * 80}\")\n print(\"Step 1/3: Checking for remaining ConnectorCredentialPairs and Documents\")\n print(f\"{'=' * 80}\")\n try:\n check_documents_deleted(pod_name, tenant_id)\n except Exception as e:\n print(f\"✗ Document check failed: {e}\", file=sys.stderr)\n print(\n \"\\nPlease ensure all ConnectorCredentialPairs and Documents are deleted before running cleanup.\"\n )\n print(\n \"You may need to mark connectors for deletion and wait for cleanup to complete.\"\n )\n return False\n print(f\"{'=' * 80}\\n\")\n\n # Step 2: Drop data plane schema\n if confirm_step(\n f\"Step 2/3: Drop data plane schema '{tenant_id}' (CASCADE - will delete all tables, functions, etc.)\",\n force,\n ):\n try:\n drop_data_plane_schema(pod_name, tenant_id)\n except Exception as e:\n print(f\"✗ Failed at schema cleanup step: {e}\", file=sys.stderr)\n if not force:\n response = input(\"Continue with control plane cleanup? (y/n): \")\n if response.lower() != \"y\":\n print(\"Cleanup aborted by user\")\n return False\n else:\n print(\"[FORCE MODE] Continuing despite schema cleanup failure\")\n else:\n print(\"Step 2 skipped by user\")\n\n # Step 3: Clean up control plane (skip if tenant not found in control plane with --force)\n if tenant_not_found_in_control_plane:\n print(f\"\\n{'=' * 80}\")\n print(\n \"Step 3/3: Skipping control plane cleanup (tenant not found in control plane)\"\n )\n print(f\"{'=' * 80}\\n\")\n elif confirm_step(\n \"Step 3/3: Delete control plane records (tenant_notification, tenant_config, subscription, tenant)\",\n force,\n ):\n try:\n cleanup_control_plane(tenant_id, force)\n except Exception as e:\n print(f\"✗ Failed at control plane cleanup step: {e}\", file=sys.stderr)\n if not force:\n print(\"Control plane cleanup failed\")\n else:\n print(\"[FORCE MODE] Control plane cleanup failed but continuing\")\n else:\n print(\"Step 3 skipped by user\")\n return False\n\n print(f\"\\n{'=' * 80}\")\n print(f\"✓ Cleanup completed for tenant: {tenant_id}\")\n print(f\"{'=' * 80}\")\n return True\n\n\ndef main() -> None:\n # Register signal handlers for graceful shutdown\n signal.signal(signal.SIGINT, signal_handler)\n signal.signal(signal.SIGTERM, signal_handler)\n\n if len(sys.argv) < 2:\n print(\"Usage: python backend/scripts/cleanup_tenant.py [--force]\")\n print(\n \" python backend/scripts/cleanup_tenant.py --csv [--force]\"\n )\n print(\"\\nArguments:\")\n print(\n \" tenant_id The tenant ID to clean up (required if not using --csv)\"\n )\n print(\" --csv PATH Path to CSV file containing tenant IDs to clean up\")\n print(\" --force Skip all confirmation prompts (optional)\")\n print(\"\\nExamples:\")\n print(\" python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789\")\n print(\n \" python backend/scripts/cleanup_tenant.py tenant_abc123-def456-789 --force\"\n )\n print(\n \" python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv\"\n )\n print(\n \" python backend/scripts/cleanup_tenant.py --csv gated_tenants_no_query_3mo.csv --force\"\n )\n sys.exit(1)\n\n # Parse arguments\n force = \"--force\" in sys.argv\n tenant_ids = []\n\n # Check for CSV mode\n if \"--csv\" in sys.argv:\n try:\n csv_index = sys.argv.index(\"--csv\")\n if csv_index + 1 >= len(sys.argv):\n print(\"Error: --csv flag requires a file path\", file=sys.stderr)\n sys.exit(1)\n\n csv_path = sys.argv[csv_index + 1]\n tenant_ids = read_tenant_ids_from_csv(csv_path)\n\n if not tenant_ids:\n print(\"Error: No tenant IDs found in CSV file\", file=sys.stderr)\n sys.exit(1)\n\n print(f\"Found {len(tenant_ids)} tenant(s) in CSV file: {csv_path}\")\n\n except Exception as e:\n print(f\"Error reading CSV file: {e}\", file=sys.stderr)\n sys.exit(1)\n else:\n # Single tenant mode\n tenant_ids = [sys.argv[1]]\n\n # Initial confirmation (unless --force is used)\n if not force:\n print(f\"\\n{'=' * 80}\")\n print(\"TENANT CLEANUP - CONFIRMATION REQUIRED\")\n print(f\"{'=' * 80}\")\n if len(tenant_ids) == 1:\n print(f\"Tenant ID: {tenant_ids[0]}\")\n else:\n print(f\"Number of tenants: {len(tenant_ids)}\")\n print(f\"Tenant IDs: {', '.join(tenant_ids[:5])}\")\n if len(tenant_ids) > 5:\n print(f\" ... and {len(tenant_ids) - 5} more\")\n\n print(\"Index Name: Will be fetched automatically when deleting Vespa documents\")\n print(\n f\"Mode: {'FORCE (no confirmations)' if force else 'Interactive (will ask for confirmation at each step)'}\"\n )\n print(\"\\nThis will:\")\n print(\" 1. Delete ALL Vespa documents for this tenant\")\n print(\" 2. Drop the data plane PostgreSQL schema (CASCADE)\")\n print(\" 3. Clean up control plane data:\")\n print(\" - Delete from tenant_notification table\")\n print(\" - Delete from tenant_config table\")\n print(\" - Delete from subscription table\")\n print(\" - Delete from tenant table\")\n print(f\"\\n{'=' * 80}\")\n print(\"WARNING: This operation is IRREVERSIBLE!\")\n print(f\"{'=' * 80}\\n\")\n\n response = input(\"Are you sure you want to proceed? Type 'yes' to confirm: \")\n\n if response.lower() != \"yes\":\n print(\"Cleanup aborted by user\")\n sys.exit(0)\n else:\n if len(tenant_ids) == 1:\n print(\n f\"⚠ FORCE MODE: Running cleanup for {tenant_ids[0]} without confirmations\"\n )\n else:\n print(\n f\"⚠ FORCE MODE: Running cleanup for {len(tenant_ids)} tenants without confirmations\"\n )\n\n # Find heavy worker pod once for all tenants\n try:\n pod_name = find_worker_pod()\n print(f\"✓ Found worker pod: {pod_name}\\n\")\n except Exception as e:\n print(f\"✗ Failed to find heavy worker pod: {e}\", file=sys.stderr)\n print(\"Cannot proceed with cleanup\")\n sys.exit(1)\n\n # Run cleanup for each tenant\n failed_tenants = []\n successful_tenants = []\n skipped_tenants = []\n\n # Open CSV file for writing successful cleanups in real-time\n csv_output_path = \"cleaned_tenants.csv\"\n with open(csv_output_path, \"w\", newline=\"\") as csv_file:\n csv_writer = csv.writer(csv_file)\n csv_writer.writerow([\"tenant_id\", \"cleaned_at\"])\n csv_file.flush() # Ensure header is written immediately\n\n print(f\"Writing successful cleanups to: {csv_output_path}\\n\")\n\n for idx, tenant_id in enumerate(tenant_ids, 1):\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(f\"Processing tenant {idx}/{len(tenant_ids)}: {tenant_id}\")\n print(f\"{'=' * 80}\")\n\n try:\n was_cleaned = cleanup_tenant(tenant_id, pod_name, force)\n\n if was_cleaned:\n # Only record if actually cleaned up (not skipped)\n successful_tenants.append(tenant_id)\n\n # Write to CSV immediately after successful cleanup\n timestamp = datetime.utcnow().isoformat()\n csv_writer.writerow([tenant_id, timestamp])\n csv_file.flush() # Ensure real-time write\n print(f\"✓ Recorded cleanup in {csv_output_path}\")\n else:\n skipped_tenants.append(tenant_id)\n print(f\"⚠ Tenant {tenant_id} was skipped (not recorded in CSV)\")\n\n except Exception as e:\n print(f\"✗ Cleanup failed for tenant {tenant_id}: {e}\", file=sys.stderr)\n failed_tenants.append((tenant_id, str(e)))\n\n # If not in force mode and there are more tenants, ask if we should continue\n if not force and idx < len(tenant_ids):\n response = input(\n f\"\\nContinue with remaining {len(tenant_ids) - idx} tenant(s)? (y/n): \"\n )\n if response.lower() != \"y\":\n print(\"Cleanup aborted by user\")\n break\n\n # Print summary\n if len(tenant_ids) == 1:\n if successful_tenants:\n print(f\"\\n✓ Successfully cleaned tenant written to: {csv_output_path}\")\n elif skipped_tenants:\n print(\"\\n⚠ Tenant was skipped\")\n elif len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(\"CLEANUP SUMMARY\")\n print(f\"{'=' * 80}\")\n print(f\"Total tenants: {len(tenant_ids)}\")\n print(f\"Successful: {len(successful_tenants)}\")\n print(f\"Skipped: {len(skipped_tenants)}\")\n print(f\"Failed: {len(failed_tenants)}\")\n print(f\"\\nSuccessfully cleaned tenants written to: {csv_output_path}\")\n\n if skipped_tenants:\n print(f\"\\nSkipped tenants ({len(skipped_tenants)}):\")\n for tenant_id in skipped_tenants:\n print(f\" - {tenant_id}\")\n\n if failed_tenants:\n print(f\"\\nFailed tenants ({len(failed_tenants)}):\")\n for tenant_id, error in failed_tenants:\n print(f\" - {tenant_id}: {error}\")\n\n print(f\"{'=' * 80}\")\n\n if failed_tenants:\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/cleanup_utils.py", "content": "import csv\nimport os\nimport random\nimport subprocess\nimport sys\nfrom dataclasses import dataclass\nfrom pathlib import Path\n\n\nclass TenantNotFoundInControlPlaneError(Exception):\n \"\"\"Exception raised when tenant/table is not found in control plane.\"\"\"\n\n\n@dataclass\nclass ControlPlaneConfig:\n \"\"\"Configuration for connecting to the control plane database.\"\"\"\n\n db_url: str\n bastion_host: str\n pem_file_location: str\n\n\ndef find_worker_pod() -> str:\n \"\"\"Find a user file processing worker pod using kubectl.\"\"\"\n print(\"Finding user file processing worker pod...\")\n\n result = subprocess.run(\n [\"kubectl\", \"get\", \"po\"], capture_output=True, text=True, check=True\n )\n\n # Parse output and find user file processing worker pod\n lines = result.stdout.strip().split(\"\\n\")\n lines = lines[1:] # Skip header\n random.shuffle(lines)\n for line in lines:\n if \"celery-worker-user-file-processing\" in line and \"Running\" in line:\n pod_name = line.split()[0]\n print(f\"Found pod: {pod_name}\")\n return pod_name\n\n raise RuntimeError(\"No running user file processing worker pod found\")\n\n\ndef confirm_step(message: str, force: bool = False) -> bool:\n \"\"\"Ask for confirmation before executing a step.\n\n Args:\n message: The confirmation message to display\n force: If True, skip confirmation and return True\n\n Returns:\n True if user confirms or force is True, False otherwise\n \"\"\"\n if force:\n print(f\"[FORCE MODE] Skipping confirmation: {message}\")\n return True\n\n print(f\"\\n{message}\")\n response = input(\"Proceed? (y/n): \")\n return response.lower() == \"y\"\n\n\ndef get_control_plane_config() -> ControlPlaneConfig:\n \"\"\"Get control plane database configuration from environment variables.\n\n Returns:\n ControlPlaneConfig with db_url, bastion_host, and pem_file_location\n\n Raises:\n ValueError: If any required environment variable is not set\n \"\"\"\n rds_host = os.environ.get(\"CONTROL_PLANE_RDS_HOST\")\n if not rds_host:\n raise ValueError(\"CONTROL_PLANE_RDS_HOST is not set\")\n\n rds_password = os.environ.get(\"CONTROL_PLANE_RDS_PASSWORD\")\n if not rds_password:\n raise ValueError(\"CONTROL_PLANE_RDS_PASSWORD is not set\")\n\n bastion_host = os.environ.get(\"BASTION_HOST\")\n if not bastion_host:\n raise ValueError(\"BASTION_HOST is not set\")\n\n pem_file_location = os.environ.get(\"PEM_FILE_LOCATION\")\n if not pem_file_location:\n raise ValueError(\"PEM_FILE_LOCATION is not set\")\n\n db_url = f\"postgresql://postgres:{rds_password}@{rds_host}:5432/control\"\n\n return ControlPlaneConfig(\n db_url=db_url,\n bastion_host=bastion_host,\n pem_file_location=pem_file_location,\n )\n\n\ndef execute_control_plane_query(\n query: str, tuple_only: bool = False\n) -> subprocess.CompletedProcess:\n \"\"\"Execute a SQL query against the control plane database via SSH.\n\n Args:\n query: The SQL query to execute\n tuple_only: If True, use psql's tuple-only mode (-t flag) for cleaner output\n\n Returns:\n subprocess.CompletedProcess with the result\n\n Raises:\n subprocess.CalledProcessError: If the command fails\n \"\"\"\n config = get_control_plane_config()\n db_url = config.db_url\n bastion_host = config.bastion_host\n pem_file_location = config.pem_file_location\n\n # Build psql flags\n psql_flags = \"-t\" if tuple_only else \"\"\n\n # Build the SSH command with proper escaping\n full_cmd = f'ssh -i {pem_file_location} ec2-user@{bastion_host} \"psql {db_url} {psql_flags} -c \\\\\"{query}\\\\\"\"'\n\n result = subprocess.run(\n full_cmd,\n shell=True,\n check=True,\n capture_output=True,\n text=True,\n )\n\n return result\n\n\ndef get_tenant_status(tenant_id: str) -> str | None:\n \"\"\"\n Get tenant status from control plane database.\n\n Returns:\n Tenant status string (e.g., 'GATED_ACCESS', 'ACTIVE') or None if not found\n\n Raises:\n TenantNotFoundInControlPlaneError: If the tenant table/relation does not exist\n \"\"\"\n print(f\"Fetching tenant status for tenant: {tenant_id}\")\n\n query = f\"SELECT application_status FROM tenant WHERE tenant_id = '{tenant_id}';\"\n\n try:\n result = execute_control_plane_query(query, tuple_only=True)\n\n # Parse the output - psql returns the value with whitespace\n status = result.stdout.strip()\n\n if status:\n print(f\"✓ Tenant status: {status}\")\n return status\n else:\n print(\"⚠ Tenant not found in control plane\")\n raise TenantNotFoundInControlPlaneError(\n f\"Tenant {tenant_id} not found in control plane database\"\n )\n except TenantNotFoundInControlPlaneError:\n # Re-raise without wrapping\n raise\n except subprocess.CalledProcessError as e:\n error_msg = e.stderr if e.stderr else str(e)\n print(\n f\"✗ Failed to get tenant status for {tenant_id}: {error_msg}\",\n file=sys.stderr,\n )\n return None\n\n\ndef read_tenant_ids_from_csv(csv_path: str) -> list[str]:\n \"\"\"Read tenant IDs from CSV file.\n\n Args:\n csv_path: Path to CSV file\n\n Returns:\n List of tenant IDs\n \"\"\"\n if not Path(csv_path).exists():\n raise FileNotFoundError(f\"CSV file not found: {csv_path}\")\n\n tenant_ids = []\n with open(csv_path, \"r\", newline=\"\", encoding=\"utf-8\") as csvfile:\n reader = csv.DictReader(csvfile)\n\n # Check if tenant_id column exists\n if not reader.fieldnames or \"tenant_id\" not in reader.fieldnames:\n raise ValueError(\n f\"CSV file must have a 'tenant_id' column. Found columns: {reader.fieldnames}\"\n )\n\n for row in reader:\n tenant_id = row.get(\"tenant_id\", \"\").strip()\n if tenant_id:\n tenant_ids.append(tenant_id)\n\n return tenant_ids\n" }, { "path": "backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py", "content": "#!/usr/bin/env python3\n\"\"\"\nMark connectors for deletion script that:\n1. Finds all connectors for the specified tenant(s)\n2. Cancels any scheduled indexing attempts\n3. Marks each connector credential pair as DELETING\n4. Triggers the cleanup task\n\nUsage:\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py [--force] [--concurrency N]\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv [--force] [--concurrency N]\n\nArguments:\n tenant_id The tenant ID to process (required if not using --csv)\n --csv PATH Path to CSV file containing tenant IDs to process\n --force Skip all confirmation prompts (optional)\n --concurrency N Process N tenants concurrently (default: 1)\n\nExamples:\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789 --force\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv --force\n python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py \\\n --csv gated_tenants_no_query_3mo.csv --force --concurrency 16\n\"\"\"\n\nimport subprocess\nimport sys\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import ThreadPoolExecutor\nfrom pathlib import Path\nfrom threading import Lock\nfrom typing import Any\n\nfrom scripts.tenant_cleanup.cleanup_utils import confirm_step\nfrom scripts.tenant_cleanup.cleanup_utils import find_worker_pod\nfrom scripts.tenant_cleanup.cleanup_utils import get_tenant_status\nfrom scripts.tenant_cleanup.cleanup_utils import read_tenant_ids_from_csv\n\n# Global lock for thread-safe printing\n_print_lock: Lock = Lock()\n\n\ndef safe_print(*args: Any, **kwargs: Any) -> None:\n \"\"\"Thread-safe print function.\"\"\"\n with _print_lock:\n print(*args, **kwargs)\n\n\ndef run_connector_deletion(pod_name: str, tenant_id: str) -> None:\n \"\"\"Mark all connector credential pairs for deletion.\n\n Args:\n pod_name: The Kubernetes pod name to execute on\n tenant_id: The tenant ID\n \"\"\"\n safe_print(\" Marking all connector credential pairs for deletion...\")\n\n # Get the path to the script\n script_dir = Path(__file__).parent\n mark_deletion_script = (\n script_dir / \"on_pod_scripts\" / \"execute_connector_deletion.py\"\n )\n\n if not mark_deletion_script.exists():\n raise FileNotFoundError(\n f\"execute_connector_deletion.py not found at {mark_deletion_script}\"\n )\n\n try:\n # Copy script to pod\n subprocess.run(\n [\n \"kubectl\",\n \"cp\",\n str(mark_deletion_script),\n f\"{pod_name}:/tmp/execute_connector_deletion.py\",\n ],\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n result = subprocess.run(\n [\n \"kubectl\",\n \"exec\",\n pod_name,\n \"--\",\n \"python\",\n \"/tmp/execute_connector_deletion.py\",\n tenant_id,\n \"--all\",\n ],\n )\n\n if result.returncode != 0:\n raise RuntimeError(result.stderr)\n\n except subprocess.CalledProcessError as e:\n safe_print(\n f\" ✗ Failed to mark all connector credential pairs for deletion: {e}\",\n file=sys.stderr,\n )\n if e.stderr:\n safe_print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n safe_print(\n f\" ✗ Failed to mark all connector credential pairs for deletion: {e}\",\n file=sys.stderr,\n )\n raise\n\n\ndef mark_tenant_connectors_for_deletion(\n tenant_id: str, pod_name: str, force: bool = False\n) -> None:\n \"\"\"\n Main function to mark all connectors for a tenant for deletion.\n\n Args:\n tenant_id: The tenant ID to process\n pod_name: The Kubernetes pod name to execute on\n force: If True, skip all confirmation prompts\n \"\"\"\n safe_print(f\"Processing connectors for tenant: {tenant_id}\")\n\n # Check tenant status first\n safe_print(f\"\\n{'=' * 80}\")\n try:\n tenant_status = get_tenant_status(tenant_id)\n\n # If tenant is not GATED_ACCESS, require explicit confirmation even in force mode\n if tenant_status and tenant_status != \"GATED_ACCESS\":\n safe_print(\n f\"\\n⚠️ WARNING: Tenant status is '{tenant_status}', not 'GATED_ACCESS'!\"\n )\n safe_print(\n \"This tenant may be active and should not have connectors deleted without careful review.\"\n )\n safe_print(f\"{'=' * 80}\\n\")\n\n # Always ask for confirmation if not gated, even in force mode\n # Note: In parallel mode with force, this will still block\n if not force:\n response = input(\n \"Are you ABSOLUTELY SURE you want to proceed? Type 'yes' to confirm: \"\n )\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - tenant is not GATED_ACCESS\")\n raise RuntimeError(f\"Tenant {tenant_id} is not GATED_ACCESS\")\n else:\n raise RuntimeError(f\"Tenant {tenant_id} is not GATED_ACCESS\")\n elif tenant_status == \"GATED_ACCESS\":\n safe_print(\"✓ Tenant status is GATED_ACCESS - safe to proceed\")\n elif tenant_status is None:\n safe_print(\"⚠️ WARNING: Could not determine tenant status!\")\n if not force:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - could not verify tenant status\")\n raise RuntimeError(\n f\"Could not verify tenant status for {tenant_id}\"\n )\n else:\n raise RuntimeError(f\"Could not verify tenant status for {tenant_id}\")\n except Exception as e:\n safe_print(f\"⚠️ WARNING: Failed to check tenant status: {e}\")\n if not force:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - could not verify tenant status\")\n raise\n else:\n raise RuntimeError(f\"Failed to check tenant status for {tenant_id}\")\n safe_print(f\"{'=' * 80}\\n\")\n\n # Confirm before proceeding (only in non-force mode)\n if not confirm_step(\n f\"Mark all connector credential pairs for deletion for tenant {tenant_id}?\",\n force,\n ):\n safe_print(\"Operation cancelled by user\")\n raise ValueError(\"Operation cancelled by user\")\n\n run_connector_deletion(pod_name, tenant_id)\n\n # Print summary\n safe_print(\n f\"✓ Marked all connector credential pairs for deletion for tenant {tenant_id}\"\n )\n\n\ndef main() -> None:\n if len(sys.argv) < 2:\n print(\n \"Usage: python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py [--force] [--concurrency N]\"\n )\n print(\n \" python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv [--force]\"\n \" [--concurrency N]\"\n )\n print(\"\\nArguments:\")\n print(\n \" tenant_id The tenant ID to process (required if not using --csv)\"\n )\n print(\" --csv PATH Path to CSV file containing tenant IDs to process\")\n print(\" --force Skip all confirmation prompts (optional)\")\n print(\" --concurrency N Process N tenants concurrently (default: 1)\")\n print(\"\\nExamples:\")\n print(\n \" python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789\"\n )\n print(\n \" python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py tenant_abc123-def456-789 --force\"\n )\n print(\n \" python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv\"\n )\n print(\n \" python backend/scripts/tenant_cleanup/mark_connectors_for_deletion.py --csv gated_tenants_no_query_3mo.csv \"\n \"--force --concurrency 16\"\n )\n sys.exit(1)\n\n # Parse arguments\n force = \"--force\" in sys.argv\n tenant_ids: list[str] = []\n\n # Parse concurrency\n concurrency: int = 1\n if \"--concurrency\" in sys.argv:\n try:\n concurrency_index = sys.argv.index(\"--concurrency\")\n if concurrency_index + 1 >= len(sys.argv):\n print(\"Error: --concurrency flag requires a number\", file=sys.stderr)\n sys.exit(1)\n concurrency = int(sys.argv[concurrency_index + 1])\n if concurrency < 1:\n print(\"Error: concurrency must be at least 1\", file=sys.stderr)\n sys.exit(1)\n except ValueError:\n print(\"Error: --concurrency value must be an integer\", file=sys.stderr)\n sys.exit(1)\n\n # Validate: concurrency > 1 requires --force\n if concurrency > 1 and not force:\n print(\n \"Error: --concurrency > 1 requires --force flag (interactive mode not supported with parallel processing)\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n # Check for CSV mode\n if \"--csv\" in sys.argv:\n try:\n csv_index: int = sys.argv.index(\"--csv\")\n if csv_index + 1 >= len(sys.argv):\n print(\"Error: --csv flag requires a file path\", file=sys.stderr)\n sys.exit(1)\n\n csv_path: str = sys.argv[csv_index + 1]\n tenant_ids = read_tenant_ids_from_csv(csv_path)\n\n if not tenant_ids:\n print(\"Error: No tenant IDs found in CSV file\", file=sys.stderr)\n sys.exit(1)\n\n print(f\"Found {len(tenant_ids)} tenant(s) in CSV file: {csv_path}\")\n\n except Exception as e:\n print(f\"Error reading CSV file: {e}\", file=sys.stderr)\n sys.exit(1)\n else:\n # Single tenant mode\n tenant_ids = [sys.argv[1]]\n\n # Find heavy worker pod once before processing\n try:\n print(\"Finding worker pod...\")\n pod_name: str = find_worker_pod()\n print(f\"✓ Using worker pod: {pod_name}\")\n except Exception as e:\n print(f\"✗ Failed to find heavy worker pod: {e}\", file=sys.stderr)\n print(\"Cannot proceed with marking connectors for deletion\")\n sys.exit(1)\n\n # Initial confirmation (unless --force is used)\n if not force:\n print(f\"\\n{'=' * 80}\")\n print(\"MARK CONNECTORS FOR DELETION - CONFIRMATION REQUIRED\")\n print(f\"{'=' * 80}\")\n if len(tenant_ids) == 1:\n print(f\"Tenant ID: {tenant_ids[0]}\")\n else:\n print(f\"Number of tenants: {len(tenant_ids)}\")\n print(f\"Tenant IDs: {', '.join(tenant_ids[:5])}\")\n if len(tenant_ids) > 5:\n print(f\" ... and {len(tenant_ids) - 5} more\")\n\n print(\n f\"Mode: {'FORCE (no confirmations)' if force else 'Interactive (will ask for confirmation at each step)'}\"\n )\n print(f\"Concurrency: {concurrency} tenant(s) at a time\")\n print(\"\\nThis will:\")\n print(\" 1. Fetch all connector credential pairs for each tenant\")\n print(\" 2. Cancel any scheduled indexing attempts for each connector\")\n print(\" 3. Mark each connector credential pair status as DELETING\")\n print(\" 4. Trigger the connector deletion task\")\n print(f\"\\n{'=' * 80}\")\n print(\"WARNING: This will mark connectors for deletion!\")\n print(\"The actual deletion will be performed by the background celery worker.\")\n print(f\"{'=' * 80}\\n\")\n\n response = input(\"Are you sure you want to proceed? Type 'yes' to confirm: \")\n\n if response.lower() != \"yes\":\n print(\"Operation aborted by user\")\n sys.exit(0)\n else:\n if len(tenant_ids) == 1:\n print(\n f\"⚠ FORCE MODE: Marking connectors for deletion for {tenant_ids[0]} without confirmations\"\n )\n else:\n print(\n f\"⚠ FORCE MODE: Marking connectors for deletion for {len(tenant_ids)} tenants \"\n f\"(concurrency: {concurrency}) without confirmations\"\n )\n\n # Process tenants (in parallel if concurrency > 1)\n failed_tenants: list[tuple[str, str]] = []\n successful_tenants: list[str] = []\n\n if concurrency == 1:\n # Sequential processing\n for idx, tenant_id in enumerate(tenant_ids, 1):\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(f\"Processing tenant {idx}/{len(tenant_ids)}: {tenant_id}\")\n print(f\"{'=' * 80}\")\n\n try:\n mark_tenant_connectors_for_deletion(tenant_id, pod_name, force)\n successful_tenants.append(tenant_id)\n except Exception as e:\n print(\n f\"✗ Failed to process tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n failed_tenants.append((tenant_id, str(e)))\n\n # If not in force mode and there are more tenants, ask if we should continue\n if not force and idx < len(tenant_ids):\n response = input(\n f\"\\nContinue with remaining {len(tenant_ids) - idx} tenant(s)? (y/n): \"\n )\n if response.lower() != \"y\":\n print(\"Operation aborted by user\")\n break\n else:\n # Parallel processing\n print(\n f\"\\nProcessing {len(tenant_ids)} tenant(s) with concurrency={concurrency}\"\n )\n\n def process_tenant(tenant_id: str) -> tuple[str, bool, str | None]:\n \"\"\"Process a single tenant. Returns (tenant_id, success, error_message).\"\"\"\n try:\n mark_tenant_connectors_for_deletion(tenant_id, pod_name, force)\n return (tenant_id, True, None)\n except Exception as e:\n return (tenant_id, False, str(e))\n\n with ThreadPoolExecutor(max_workers=concurrency) as executor:\n # Submit all tasks\n future_to_tenant = {\n executor.submit(process_tenant, tenant_id): tenant_id\n for tenant_id in tenant_ids\n }\n\n # Process results as they complete\n completed: int = 0\n for future in as_completed(future_to_tenant):\n completed += 1\n tenant_id, success, error = future.result()\n\n if success:\n successful_tenants.append(tenant_id)\n safe_print(\n f\"[{completed}/{len(tenant_ids)}] ✓ Successfully processed {tenant_id}\"\n )\n else:\n failed_tenants.append((tenant_id, error or \"Unknown error\"))\n safe_print(\n f\"[{completed}/{len(tenant_ids)}] ✗ Failed to process {tenant_id}: {error}\",\n file=sys.stderr,\n )\n\n # Print summary if multiple tenants\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(\"OPERATION SUMMARY\")\n print(f\"{'=' * 80}\")\n print(f\"Total tenants: {len(tenant_ids)}\")\n print(f\"Successful: {len(successful_tenants)}\")\n print(f\"Failed: {len(failed_tenants)}\")\n\n if failed_tenants:\n print(\"\\nFailed tenants:\")\n for tenant_id, error in failed_tenants:\n print(f\" - {tenant_id}: {error}\")\n\n print(f\"{'=' * 80}\")\n\n if failed_tenants:\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/no_bastion_analyze_tenants.py", "content": "#!/usr/bin/env python3\n\"\"\"\nTenant analysis script that works WITHOUT bastion access.\nControl plane and data plane are in SEPARATE clusters.\n\nUsage:\n PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_analyze_tenants.py \\\n [--skip-cache] \\\n [--data-plane-context ] \\\n [--control-plane-context ]\n\"\"\"\n\nimport argparse\nimport csv\nimport json\nimport subprocess\nimport sys\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom pathlib import Path\nfrom typing import Any\n\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import find_background_pod\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import find_worker_pod\n\n\ndef collect_tenant_data(\n pod_name: str, context: str | None = None\n) -> list[dict[str, Any]]:\n \"\"\"Run the understand_tenants script on the data plane pod.\"\"\"\n print(f\"\\nCollecting tenant data from data plane pod {pod_name}...\")\n\n # Get the path to the understand_tenants script\n script_dir = Path(__file__).parent\n understand_tenants_script = script_dir / \"on_pod_scripts\" / \"understand_tenants.py\"\n\n if not understand_tenants_script.exists():\n raise FileNotFoundError(\n f\"understand_tenants.py not found at {understand_tenants_script}\"\n )\n\n # Copy script to pod\n print(\"Copying script to pod...\")\n cmd_cp = [\n \"kubectl\",\n \"cp\",\n str(understand_tenants_script),\n f\"{pod_name}:/tmp/understand_tenants.py\",\n ]\n if context:\n cmd_cp.extend([\"--context\", context])\n\n subprocess.run(cmd_cp, check=True, capture_output=True)\n\n # Execute script on pod\n print(\"Executing script on pod (this may take a while)...\")\n cmd_exec = [\"kubectl\", \"exec\", pod_name]\n if context:\n cmd_exec.extend([\"--context\", context])\n cmd_exec.extend([\"--\", \"python\", \"/tmp/understand_tenants.py\"])\n\n result = subprocess.run(cmd_exec, capture_output=True, text=True, check=True)\n\n # Show progress messages from stderr\n if result.stderr:\n print(result.stderr, file=sys.stderr)\n\n # Parse JSON from stdout\n try:\n tenant_data = json.loads(result.stdout)\n print(f\"Successfully collected data for {len(tenant_data)} tenants\")\n return tenant_data\n except json.JSONDecodeError as e:\n print(f\"Failed to parse JSON output: {e}\", file=sys.stderr)\n print(f\"stdout: {result.stdout[:500]}\", file=sys.stderr)\n raise\n\n\ndef collect_control_plane_data_from_pod(\n pod_name: str, context: str | None = None\n) -> list[dict[str, Any]]:\n \"\"\"Collect control plane data by running a query on a control plane pod.\"\"\"\n print(f\"\\nCollecting control plane data from pod {pod_name}...\")\n\n # Create a script to query the control plane database\n query_script = \"\"\"\nimport json\nimport os\nfrom sqlalchemy import create_engine, text\n\n# Try to get database URL from various environment patterns\ncontrol_db_url = None\n\n# Pattern 1: POSTGRES_CONTROL_* variables\nif os.environ.get(\"POSTGRES_CONTROL_HOST\"):\n host = os.environ.get(\"POSTGRES_CONTROL_HOST\")\n port = os.environ.get(\"POSTGRES_CONTROL_PORT\", \"5432\")\n db = os.environ.get(\"POSTGRES_CONTROL_DB\", \"control\")\n user = os.environ.get(\"POSTGRES_CONTROL_USER\", \"postgres\")\n password = os.environ.get(\"POSTGRES_CONTROL_PASSWORD\", \"\")\n if password:\n control_db_url = f\"postgresql://{user}:{password}@{host}:{port}/{db}\"\n\n# Pattern 2: Standard POSTGRES_* variables (in control plane cluster)\nif not control_db_url and os.environ.get(\"POSTGRES_HOST\"):\n host = os.environ.get(\"POSTGRES_HOST\")\n port = os.environ.get(\"POSTGRES_PORT\", \"5432\")\n db = os.environ.get(\"POSTGRES_DB\", \"danswer\")\n user = os.environ.get(\"POSTGRES_USER\", \"postgres\")\n password = os.environ.get(\"POSTGRES_PASSWORD\", \"\")\n if password:\n control_db_url = f\"postgresql://{user}:{password}@{host}:{port}/{db}\"\n\nif not control_db_url:\n raise ValueError(\"Cannot determine control plane database connection\")\n\nengine = create_engine(control_db_url)\n\nwith engine.connect() as conn:\n result = conn.execute(\n text(\n \"SELECT tenant_id, stripe_customer_id, created_at, active_seats, \"\n \"creator_email, referral_source, application_status FROM tenant\"\n )\n )\n rows = [dict(row._mapping) for row in result]\n print(json.dumps(rows, default=str))\n\"\"\"\n\n # Write the script to a temp file\n script_path = \"/tmp/query_control_plane.py\"\n\n print(\" Creating control plane query script on pod...\")\n cmd_write = [\"kubectl\", \"exec\", pod_name]\n if context:\n cmd_write.extend([\"--context\", context])\n cmd_write.extend(\n [\"--\", \"bash\", \"-c\", f\"cat > {script_path} << 'EOF'\\n{query_script}\\nEOF\"]\n )\n\n subprocess.run(cmd_write, check=True, capture_output=True)\n\n # Execute the script on the pod\n print(\" Executing control plane query on pod...\")\n cmd_exec = [\"kubectl\", \"exec\", pod_name]\n if context:\n cmd_exec.extend([\"--context\", context])\n cmd_exec.extend([\"--\", \"python\", script_path])\n\n result = subprocess.run(cmd_exec, capture_output=True, text=True, check=True)\n\n # Parse JSON output\n try:\n control_plane_data = json.loads(result.stdout)\n print(\n f\"✓ Successfully collected {len(control_plane_data)} tenant records from control plane\"\n )\n return control_plane_data\n except json.JSONDecodeError as e:\n print(f\"Failed to parse JSON output: {e}\", file=sys.stderr)\n print(f\"stdout: {result.stdout[:500]}\", file=sys.stderr)\n raise\n\n\ndef analyze_tenants(\n tenants: list[dict[str, Any]], control_plane_data: list[dict[str, Any]]\n) -> list[dict[str, Any]]:\n \"\"\"Analyze tenant activity data and return gated tenants with no query in last 3 months.\"\"\"\n\n print(f\"\\n{'=' * 80}\")\n print(f\"TENANT ANALYSIS REPORT - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\")\n print(f\"{'=' * 80}\")\n print(f\"Total tenants analyzed: {len(tenants)}\\n\")\n\n # Create a lookup dict for control plane data by tenant_id\n control_plane_lookup = {}\n for row in control_plane_data:\n tenant_id = row.get(\"tenant_id\")\n tenant_status = row.get(\"application_status\")\n if tenant_id:\n control_plane_lookup[tenant_id] = tenant_status\n\n # Calculate cutoff dates\n one_month_cutoff = datetime.now(timezone.utc) - timedelta(days=30)\n three_month_cutoff = datetime.now(timezone.utc) - timedelta(days=90)\n\n # Categorize tenants into 4 groups\n gated_no_query_3_months = [] # GATED_ACCESS + no query in last 3 months\n gated_query_1_3_months = [] # GATED_ACCESS + query between 1-3 months\n gated_query_1_month = [] # GATED_ACCESS + query in last 1 month\n everyone_else = [] # All other tenants\n\n for tenant in tenants:\n tenant_id = tenant.get(\"tenant_id\")\n last_query_time = tenant.get(\"last_query_time\")\n tenant_status = control_plane_lookup.get(tenant_id, \"UNKNOWN\")\n\n is_gated = tenant_status == \"GATED_ACCESS\"\n\n # Parse last query time\n if last_query_time:\n query_time = datetime.fromisoformat(last_query_time.replace(\"Z\", \"+00:00\"))\n else:\n query_time = None\n\n # Categorize\n if is_gated:\n if query_time is None or query_time <= three_month_cutoff:\n gated_no_query_3_months.append(tenant)\n elif query_time <= one_month_cutoff:\n gated_query_1_3_months.append(tenant)\n else: # query_time > one_month_cutoff\n gated_query_1_month.append(tenant)\n else:\n everyone_else.append(tenant)\n\n # Calculate document counts for each group\n gated_no_query_docs = sum(\n t.get(\"num_documents\", 0) for t in gated_no_query_3_months\n )\n gated_1_3_month_docs = sum(\n t.get(\"num_documents\", 0) for t in gated_query_1_3_months\n )\n gated_1_month_docs = sum(t.get(\"num_documents\", 0) for t in gated_query_1_month)\n everyone_else_docs = sum(t.get(\"num_documents\", 0) for t in everyone_else)\n\n print(\"=\" * 80)\n print(\"TENANT CATEGORIZATION BY GATED ACCESS STATUS AND ACTIVITY\")\n print(\"=\" * 80)\n\n print(\"\\n1. GATED_ACCESS + No query in last 3 months:\")\n print(f\" Count: {len(gated_no_query_3_months):,}\")\n print(f\" Total documents: {gated_no_query_docs:,}\")\n print(\n f\" Avg documents per tenant: {gated_no_query_docs / len(gated_no_query_3_months) if gated_no_query_3_months else 0:.2f}\"\n )\n\n print(\"\\n2. GATED_ACCESS + Query between 1-3 months ago:\")\n print(f\" Count: {len(gated_query_1_3_months):,}\")\n print(f\" Total documents: {gated_1_3_month_docs:,}\")\n print(\n f\" Avg documents per tenant: {gated_1_3_month_docs / len(gated_query_1_3_months) if gated_query_1_3_months else 0:.2f}\"\n )\n\n print(\"\\n3. GATED_ACCESS + Query in last 1 month:\")\n print(f\" Count: {len(gated_query_1_month):,}\")\n print(f\" Total documents: {gated_1_month_docs:,}\")\n print(\n f\" Avg documents per tenant: {gated_1_month_docs / len(gated_query_1_month) if gated_query_1_month else 0:.2f}\"\n )\n\n print(\"\\n4. Everyone else (non-GATED_ACCESS):\")\n print(f\" Count: {len(everyone_else):,}\")\n print(f\" Total documents: {everyone_else_docs:,}\")\n print(\n f\" Avg documents per tenant: {everyone_else_docs / len(everyone_else) if everyone_else else 0:.2f}\"\n )\n\n total_docs = (\n gated_no_query_docs\n + gated_1_3_month_docs\n + gated_1_month_docs\n + everyone_else_docs\n )\n print(f\"\\nTotal documents across all tenants: {total_docs:,}\")\n\n # Top 100 tenants by document count\n print(\"\\n\" + \"=\" * 80)\n print(\"TOP 100 TENANTS BY DOCUMENT COUNT\")\n print(\"=\" * 80)\n\n # Sort all tenants by document count\n sorted_tenants = sorted(\n tenants, key=lambda t: t.get(\"num_documents\", 0), reverse=True\n )\n\n top_100 = sorted_tenants[:100]\n\n print(\n f\"\\n{'Rank':<6} {'Tenant ID':<45} {'Documents':>12} {'Users':>8} {'Last Query':<12} {'Group'}\"\n )\n print(\"-\" * 130)\n\n for idx, tenant in enumerate(top_100, 1):\n tenant_id = tenant.get(\"tenant_id\", \"Unknown\")\n num_docs = tenant.get(\"num_documents\", 0)\n num_users = tenant.get(\"num_users\", 0)\n last_query = tenant.get(\"last_query_time\", \"Never\")\n tenant_status = control_plane_lookup.get(tenant_id, \"UNKNOWN\")\n\n # Format the last query time\n if last_query and last_query != \"Never\":\n try:\n query_dt = datetime.fromisoformat(last_query.replace(\"Z\", \"+00:00\"))\n last_query_str = query_dt.strftime(\"%Y-%m-%d\")\n except Exception:\n last_query_str = last_query[:10] if len(last_query) > 10 else last_query\n else:\n last_query_str = \"Never\"\n\n # Determine group\n if tenant_status == \"GATED_ACCESS\":\n if last_query and last_query != \"Never\":\n query_time = datetime.fromisoformat(last_query.replace(\"Z\", \"+00:00\"))\n if query_time <= three_month_cutoff:\n group = \"Gated - No query (3mo)\"\n elif query_time <= one_month_cutoff:\n group = \"Gated - Query (1-3mo)\"\n else:\n group = \"Gated - Query (1mo)\"\n else:\n group = \"Gated - No query (3mo)\"\n else:\n group = f\"Other ({tenant_status})\"\n\n print(\n f\"{idx:<6} {tenant_id:<45} {num_docs:>12,} {num_users:>8} {last_query_str:<12} {group}\"\n )\n\n # Summary stats for top 100\n top_100_docs = sum(t.get(\"num_documents\", 0) for t in top_100)\n\n print(\"\\n\" + \"-\" * 110)\n print(f\"Top 100 total documents: {top_100_docs:,}\")\n print(\n f\"Percentage of all documents: {(top_100_docs / total_docs * 100) if total_docs > 0 else 0:.2f}%\"\n )\n\n # Additional insights\n print(\"\\n\" + \"=\" * 80)\n print(\"ADDITIONAL INSIGHTS\")\n print(\"=\" * 80)\n\n # Tenants with no documents\n no_docs = [t for t in tenants if t.get(\"num_documents\", 0) == 0]\n print(\n f\"\\nTenants with 0 documents: {len(no_docs):,} ({len(no_docs) / len(tenants) * 100:.2f}%)\"\n )\n\n # Tenants with no users\n no_users = [t for t in tenants if t.get(\"num_users\", 0) == 0]\n print(\n f\"Tenants with 0 users: {len(no_users):,} ({len(no_users) / len(tenants) * 100:.2f}%)\"\n )\n\n # Document distribution quartiles\n doc_counts = sorted([t.get(\"num_documents\", 0) for t in tenants])\n if doc_counts:\n print(\"\\nDocument count distribution:\")\n print(f\" Median: {doc_counts[len(doc_counts) // 2]:,}\")\n print(f\" 75th percentile: {doc_counts[int(len(doc_counts) * 0.75)]:,}\")\n print(f\" 90th percentile: {doc_counts[int(len(doc_counts) * 0.90)]:,}\")\n print(f\" 95th percentile: {doc_counts[int(len(doc_counts) * 0.95)]:,}\")\n print(f\" 99th percentile: {doc_counts[int(len(doc_counts) * 0.99)]:,}\")\n print(f\" Max: {doc_counts[-1]:,}\")\n\n return gated_no_query_3_months\n\n\ndef find_recent_tenant_data() -> tuple[list[dict[str, Any]] | None, str | None]:\n \"\"\"Find the most recent tenant data file if it's less than 7 days old.\"\"\"\n current_dir = Path.cwd()\n tenant_data_files = list(current_dir.glob(\"tenant_data_*.json\"))\n\n if not tenant_data_files:\n return None, None\n\n # Sort by modification time, most recent first\n tenant_data_files.sort(key=lambda p: p.stat().st_mtime, reverse=True)\n most_recent = tenant_data_files[0]\n\n # Check if file is less than 7 days old\n file_age = datetime.now().timestamp() - most_recent.stat().st_mtime\n seven_days_in_seconds = 7 * 24 * 60 * 60\n\n if file_age < seven_days_in_seconds:\n file_age_days = file_age / (24 * 60 * 60)\n print(\n f\"\\n✓ Found recent tenant data: {most_recent.name} (age: {file_age_days:.1f} days)\"\n )\n\n with open(most_recent, \"r\") as f:\n tenant_data = json.load(f)\n\n return tenant_data, str(most_recent)\n\n return None, None\n\n\ndef main() -> None:\n # Parse command-line arguments\n parser = argparse.ArgumentParser(\n description=\"Analyze tenant data WITHOUT bastion access - control plane and data plane are separate clusters\"\n )\n parser.add_argument(\n \"--skip-cache\",\n action=\"store_true\",\n help=\"Skip cached tenant data and collect fresh data from pod\",\n )\n parser.add_argument(\n \"--data-plane-context\",\n type=str,\n help=\"Kubectl context for data plane cluster (optional)\",\n )\n parser.add_argument(\n \"--control-plane-context\",\n type=str,\n help=\"Kubectl context for control plane cluster (optional)\",\n )\n args = parser.parse_args()\n\n try:\n # Step 1: Check for recent tenant data (< 7 days old) unless --skip-cache is set\n tenant_data = None\n cached_file = None\n\n if not args.skip_cache:\n tenant_data, cached_file = find_recent_tenant_data()\n\n if tenant_data:\n print(f\"Using cached tenant data from: {cached_file}\")\n print(f\"Total tenants in cache: {len(tenant_data)}\")\n else:\n if args.skip_cache:\n print(\"\\n⚠ Skipping cache (--skip-cache flag set)\")\n\n # Find data plane worker pod\n print(\"\\n\" + \"=\" * 80)\n print(\"CONNECTING TO DATA PLANE CLUSTER\")\n print(\"=\" * 80)\n data_plane_pod = find_worker_pod(args.data_plane_context)\n\n # Collect tenant data from data plane\n tenant_data = collect_tenant_data(data_plane_pod, args.data_plane_context)\n\n # Save raw data to file with timestamp\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n output_file = f\"tenant_data_{timestamp}.json\"\n with open(output_file, \"w\") as f:\n json.dump(tenant_data, f, indent=2, default=str)\n print(f\"\\n✓ Raw data saved to: {output_file}\")\n\n # Step 2: Collect control plane data from control plane cluster\n print(\"\\n\" + \"=\" * 80)\n print(\"CONNECTING TO CONTROL PLANE CLUSTER\")\n print(\"=\" * 80)\n control_plane_pod = find_background_pod(args.control_plane_context)\n control_plane_data = collect_control_plane_data_from_pod(\n control_plane_pod, args.control_plane_context\n )\n\n # Step 3: Analyze the data and get gated tenants without recent queries\n gated_no_query_3_months = analyze_tenants(tenant_data, control_plane_data)\n\n # Step 4: Export to CSV (sorted by num_documents descending)\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n csv_file = f\"gated_tenants_no_query_3mo_{timestamp}.csv\"\n\n # Sort by num_documents in descending order\n sorted_tenants = sorted(\n gated_no_query_3_months,\n key=lambda t: t.get(\"num_documents\", 0),\n reverse=True,\n )\n\n with open(csv_file, \"w\", newline=\"\", encoding=\"utf-8\") as csvfile:\n fieldnames = [\n \"tenant_id\",\n \"num_documents\",\n \"num_users\",\n \"last_query_time\",\n \"days_since_last_query\",\n ]\n writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n writer.writeheader()\n\n now = datetime.now(timezone.utc)\n for tenant in sorted_tenants:\n # Calculate days since last query\n last_query_time = tenant.get(\"last_query_time\")\n if last_query_time:\n try:\n query_dt = datetime.fromisoformat(\n last_query_time.replace(\"Z\", \"+00:00\")\n )\n days_since = str((now - query_dt).days)\n except Exception:\n days_since = \"N/A\"\n else:\n days_since = \"Never\"\n\n writer.writerow(\n {\n \"tenant_id\": tenant.get(\"tenant_id\", \"\"),\n \"num_documents\": tenant.get(\"num_documents\", 0),\n \"num_users\": tenant.get(\"num_users\", 0),\n \"last_query_time\": last_query_time or \"Never\",\n \"days_since_last_query\": days_since,\n }\n )\n\n print(f\"\\n✓ CSV exported to: {csv_file}\")\n print(\n f\" Total gated tenants with no query in last 3 months: {len(gated_no_query_3_months)}\"\n )\n\n except subprocess.CalledProcessError as e:\n print(f\"Error running command: {e}\", file=sys.stderr)\n if e.stderr:\n print(f\"stderr: {e.stderr}\", file=sys.stderr)\n sys.exit(1)\n except Exception as e:\n print(f\"Error: {e}\", file=sys.stderr)\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/no_bastion_cleanup_tenants.py", "content": "#!/usr/bin/env python3\n\"\"\"\nTenant cleanup script that works WITHOUT bastion access.\nAll queries run directly from pods.\nSupports two-cluster architecture (data plane and control plane in separate clusters).\n\nUsage:\n PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_cleanup_tenants.py \\\n --data-plane-context --control-plane-context [--force]\n\n PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_cleanup_tenants.py --csv \\\n --data-plane-context --control-plane-context [--force]\n\"\"\"\n\nimport csv\nimport json\nimport signal\nimport subprocess\nimport sys\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import ThreadPoolExecutor\nfrom datetime import datetime\nfrom pathlib import Path\nfrom threading import Lock\n\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import confirm_step\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import execute_control_plane_delete\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import find_background_pod\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import find_worker_pod\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import get_tenant_status\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import read_tenant_ids_from_csv\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import (\n TenantNotFoundInControlPlaneError,\n)\n\n\n# Global lock for thread-safe operations\n_print_lock: Lock = Lock()\n_csv_lock: Lock = Lock()\n\n\ndef signal_handler(signum: int, frame: object) -> None: # noqa: ARG001\n \"\"\"Handle termination signals by killing active subprocess.\"\"\"\n sys.exit(1)\n\n\ndef setup_scripts_on_pod(pod_name: str, context: str) -> None:\n \"\"\"Copy all required scripts to the pod once at the beginning.\n\n Args:\n pod_name: Pod to copy scripts to\n context: kubectl context for the cluster\n \"\"\"\n print(\"Setting up scripts on pod (one-time operation)...\")\n\n script_dir = Path(__file__).parent\n scripts_to_copy = [\n (\n \"on_pod_scripts/check_documents_deleted.py\",\n \"/tmp/check_documents_deleted.py\",\n ),\n (\"on_pod_scripts/cleanup_tenant_schema.py\", \"/tmp/cleanup_tenant_schema.py\"),\n (\"on_pod_scripts/get_tenant_users.py\", \"/tmp/get_tenant_users.py\"),\n (\"on_pod_scripts/get_tenant_index_name.py\", \"/tmp/get_tenant_index_name.py\"),\n ]\n\n for local_path, remote_path in scripts_to_copy:\n local_file = script_dir / local_path\n if not local_file.exists():\n raise FileNotFoundError(f\"Script not found: {local_file}\")\n\n cmd_cp = [\"kubectl\", \"cp\", \"--context\", context]\n cmd_cp.extend([str(local_file), f\"{pod_name}:{remote_path}\"])\n\n subprocess.run(cmd_cp, check=True, capture_output=True)\n\n print(\"✓ All scripts copied to pod\")\n\n\ndef get_tenant_index_name(pod_name: str, tenant_id: str, context: str) -> str:\n \"\"\"Get the default index name for the given tenant by running script on pod.\n\n Args:\n pod_name: Data plane pod to execute on\n tenant_id: Tenant ID to process\n context: kubectl context for data plane cluster\n \"\"\"\n print(f\"Getting default index name for tenant: {tenant_id}\")\n\n # Get the path to the script\n script_dir = Path(__file__).parent\n index_name_script = script_dir / \"on_pod_scripts\" / \"get_tenant_index_name.py\"\n\n if not index_name_script.exists():\n raise FileNotFoundError(\n f\"get_tenant_index_name.py not found at {index_name_script}\"\n )\n\n try:\n # Copy script to pod\n print(\" Copying script to pod...\")\n cmd_cp = [\"kubectl\", \"cp\", \"--context\", context]\n cmd_cp.extend(\n [\n str(index_name_script),\n f\"{pod_name}:/tmp/get_tenant_index_name.py\",\n ]\n )\n\n subprocess.run(\n cmd_cp,\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n print(\" Executing script on pod...\")\n cmd_exec = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_exec.extend(\n [\n \"--\",\n \"python\",\n \"/tmp/get_tenant_index_name.py\",\n tenant_id,\n ]\n )\n\n result = subprocess.run(\n cmd_exec,\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n\n if status == \"success\":\n index_name = result_data.get(\"index_name\")\n print(f\"✓ Found index name: {index_name}\")\n return index_name\n else:\n message = result_data.get(\"message\", \"Unknown error\")\n raise RuntimeError(f\"Failed to get index name: {message}\")\n\n except subprocess.CalledProcessError as e:\n print(\n f\"✗ Failed to get index name for tenant {tenant_id}: {e}\", file=sys.stderr\n )\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n print(\n f\"✗ Failed to get index name for tenant {tenant_id}: {e}\", file=sys.stderr\n )\n raise\n\n\ndef get_tenant_users(pod_name: str, tenant_id: str, context: str) -> list[str]:\n \"\"\"Get list of user emails from the tenant's data plane schema.\n\n Args:\n pod_name: Data plane pod to execute on\n tenant_id: Tenant ID to process\n context: kubectl context for data plane cluster\n \"\"\"\n # Script is already on pod from setup_scripts_on_pod()\n try:\n # Execute script on pod\n cmd_exec = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_exec.extend(\n [\n \"--\",\n \"python\",\n \"/tmp/get_tenant_users.py\",\n tenant_id,\n ]\n )\n\n result = subprocess.run(\n cmd_exec,\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n\n if status == \"success\":\n users = result_data.get(\"users\", [])\n if users:\n print(f\"✓ Found {len(users)} user(s):\")\n for email in users:\n print(f\" - {email}\")\n else:\n print(\" No users found in tenant\")\n return users\n else:\n message = result_data.get(\"message\", \"Unknown error\")\n print(f\"⚠ Could not fetch users: {message}\")\n return []\n\n except subprocess.CalledProcessError as e:\n print(f\"⚠ Failed to get users for tenant {tenant_id}: {e}\")\n if e.stderr:\n print(f\" Error details: {e.stderr}\")\n return []\n except Exception as e:\n print(f\"⚠ Failed to get users for tenant {tenant_id}: {e}\")\n return []\n\n\ndef check_documents_deleted(pod_name: str, tenant_id: str, context: str) -> None:\n \"\"\"Check if all documents and connector credential pairs have been deleted.\n\n Args:\n pod_name: Data plane pod to execute on\n tenant_id: Tenant ID to process\n context: kubectl context for data plane cluster\n \"\"\"\n # Script is already on pod from setup_scripts_on_pod()\n try:\n # Execute script on pod\n cmd_exec = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_exec.extend(\n [\n \"--\",\n \"python\",\n \"/tmp/check_documents_deleted.py\",\n tenant_id,\n ]\n )\n\n result = subprocess.run(\n cmd_exec,\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n\n if status == \"success\":\n message = result_data.get(\"message\")\n print(f\"✓ {message}\")\n elif status == \"not_found\":\n message = result_data.get(\"message\", \"Schema not found\")\n print(f\"⚠ {message}\")\n else:\n message = result_data.get(\"message\", \"Unknown error\")\n cc_count = result_data.get(\"connector_credential_pair_count\", 0)\n doc_count = result_data.get(\"document_count\", 0)\n error_details = f\"{message}\"\n if cc_count > 0 or doc_count > 0:\n error_details += f\"\\n ConnectorCredentialPairs: {cc_count}\\n Documents: {doc_count}\"\n raise RuntimeError(error_details)\n\n except subprocess.CalledProcessError as e:\n print(\n f\"✗ Failed to check documents for tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n print(\n f\"✗ Failed to check documents for tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n raise\n\n\ndef drop_data_plane_schema(pod_name: str, tenant_id: str, context: str) -> None:\n \"\"\"Drop the PostgreSQL schema for the given tenant by running script on pod.\n\n Args:\n pod_name: Data plane pod to execute on\n tenant_id: Tenant ID to process\n context: kubectl context for data plane cluster\n \"\"\"\n # Script is already on pod from setup_scripts_on_pod()\n try:\n # Execute script on pod\n cmd_exec = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_exec.extend(\n [\n \"--\",\n \"python\",\n \"/tmp/cleanup_tenant_schema.py\",\n tenant_id,\n ]\n )\n\n result = subprocess.run(\n cmd_exec,\n capture_output=True,\n text=True,\n check=True,\n )\n\n # Show progress messages from stderr\n if result.stderr:\n print(f\" {result.stderr}\", end=\"\")\n\n # Parse JSON result from stdout\n result_data = json.loads(result.stdout)\n status = result_data.get(\"status\")\n message = result_data.get(\"message\")\n\n if status == \"success\":\n print(f\"✓ {message}\")\n elif status == \"not_found\":\n print(f\"⚠ {message}\")\n else:\n print(f\"✗ {message}\", file=sys.stderr)\n raise RuntimeError(message)\n\n except subprocess.CalledProcessError as e:\n print(f\"✗ Failed to drop schema for tenant {tenant_id}: {e}\", file=sys.stderr)\n if e.stderr:\n print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n print(f\"✗ Failed to drop schema for tenant {tenant_id}: {e}\", file=sys.stderr)\n raise\n\n\ndef cleanup_control_plane(\n pod_name: str, tenant_id: str, context: str, force: bool = False\n) -> None:\n \"\"\"Clean up control plane data via pod queries.\n\n Args:\n pod_name: Control plane pod to execute on\n tenant_id: Tenant ID to process\n context: kubectl context for control plane cluster\n force: Skip confirmations if True\n \"\"\"\n print(f\"Cleaning up control plane data for tenant: {tenant_id}\")\n\n # Delete in order respecting foreign key constraints\n delete_queries = [\n (\n \"tenant_notification\",\n f\"DELETE FROM tenant_notification WHERE tenant_id = '{tenant_id}'\",\n ),\n (\"tenant_config\", f\"DELETE FROM tenant_config WHERE tenant_id = '{tenant_id}'\"),\n (\"subscription\", f\"DELETE FROM subscription WHERE tenant_id = '{tenant_id}'\"),\n (\"tenant\", f\"DELETE FROM tenant WHERE tenant_id = '{tenant_id}'\"),\n ]\n\n try:\n for table_name, query in delete_queries:\n print(f\" Deleting from {table_name}...\")\n\n if not confirm_step(f\"Delete from {table_name}?\", force):\n print(f\" Skipping deletion from {table_name}\")\n continue\n\n execute_control_plane_delete(pod_name, query, context)\n\n print(f\"✓ Successfully cleaned up control plane data for tenant: {tenant_id}\")\n\n except Exception as e:\n print(\n f\"✗ Failed to clean up control plane for tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n raise\n\n\ndef cleanup_tenant(\n tenant_id: str,\n data_plane_pod: str,\n control_plane_pod: str,\n data_plane_context: str,\n control_plane_context: str,\n force: bool = False,\n) -> bool:\n \"\"\"Main cleanup function that orchestrates all cleanup steps.\n\n Args:\n tenant_id: Tenant ID to process\n data_plane_pod: Data plane pod for schema operations\n control_plane_pod: Control plane pod for tenant record operations\n data_plane_context: kubectl context for data plane cluster\n control_plane_context: kubectl context for control plane cluster\n force: Skip confirmations if True\n \"\"\"\n print(f\"Starting cleanup for tenant: {tenant_id}\")\n\n # Track if tenant was not found in control plane (for force mode)\n tenant_not_found_in_control_plane = False\n\n # Check tenant status first (from control plane)\n print(f\"\\n{'=' * 80}\")\n try:\n tenant_status = get_tenant_status(\n control_plane_pod, tenant_id, control_plane_context\n )\n\n # If tenant is not GATED_ACCESS, require explicit confirmation even in force mode\n if tenant_status and tenant_status != \"GATED_ACCESS\":\n print(\n f\"\\n⚠️ WARNING: Tenant status is '{tenant_status}', not 'GATED_ACCESS'!\"\n )\n print(\n \"This tenant may be active and should not be deleted without careful review.\"\n )\n print(f\"{'=' * 80}\\n\")\n\n if force:\n print(f\"Skipping cleanup for tenant {tenant_id} in force mode\")\n return False\n\n # Always ask for confirmation if not gated\n response = input(\n \"Are you ABSOLUTELY SURE you want to proceed? Type 'yes' to confirm: \"\n )\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - tenant is not GATED_ACCESS\")\n return False\n elif tenant_status == \"GATED_ACCESS\":\n print(\"✓ Tenant status is GATED_ACCESS - safe to proceed with cleanup\")\n elif tenant_status is None:\n print(\"⚠️ WARNING: Could not determine tenant status!\")\n\n if force:\n print(f\"Skipping cleanup for tenant {tenant_id} in force mode\")\n return False\n\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - could not verify tenant status\")\n return False\n except TenantNotFoundInControlPlaneError as e:\n # Tenant/table not found in control plane\n error_str = str(e)\n print(f\"⚠️ WARNING: Tenant not found in control plane: {error_str}\")\n tenant_not_found_in_control_plane = True\n\n if force:\n print(\n \"[FORCE MODE] Tenant not found in control plane - continuing with dataplane cleanup only\"\n )\n else:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - tenant not found in control plane\")\n return False\n except Exception as e:\n # Other errors (not \"not found\")\n error_str = str(e)\n print(f\"⚠️ WARNING: Failed to check tenant status: {error_str}\")\n\n if force:\n print(f\"Skipping cleanup for tenant {tenant_id} in force mode\")\n return False\n\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n print(\"Cleanup aborted - could not verify tenant status\")\n return False\n print(f\"{'=' * 80}\\n\")\n\n # Fetch tenant users for informational purposes (non-blocking) from data plane\n if not force:\n print(f\"\\n{'=' * 80}\")\n try:\n get_tenant_users(data_plane_pod, tenant_id, data_plane_context)\n except Exception as e:\n print(f\"⚠ Could not fetch tenant users: {e}\")\n print(f\"{'=' * 80}\\n\")\n\n # Step 1: Make sure all documents are deleted (data plane)\n print(f\"\\n{'=' * 80}\")\n print(\"Step 1/3: Checking for remaining ConnectorCredentialPairs and Documents\")\n print(f\"{'=' * 80}\")\n try:\n check_documents_deleted(data_plane_pod, tenant_id, data_plane_context)\n except Exception as e:\n print(f\"✗ Document check failed: {e}\", file=sys.stderr)\n print(\n \"\\nPlease ensure all ConnectorCredentialPairs and Documents are deleted before running cleanup.\"\n )\n print(\n \"You may need to mark connectors for deletion and wait for cleanup to complete.\"\n )\n return False\n print(f\"{'=' * 80}\\n\")\n\n # Step 2: Drop data plane schema\n if confirm_step(\n f\"Step 2/3: Drop data plane schema '{tenant_id}' (CASCADE - will delete all tables, functions, etc.)\",\n force,\n ):\n try:\n drop_data_plane_schema(data_plane_pod, tenant_id, data_plane_context)\n except Exception as e:\n print(f\"✗ Failed at schema cleanup step: {e}\", file=sys.stderr)\n if not force:\n response = input(\"Continue with control plane cleanup? (y/n): \")\n if response.lower() != \"y\":\n print(\"Cleanup aborted by user\")\n return False\n else:\n print(\"[FORCE MODE] Continuing despite schema cleanup failure\")\n else:\n print(\"Step 2 skipped by user\")\n\n # Step 3: Clean up control plane (skip if tenant not found in control plane with --force)\n if tenant_not_found_in_control_plane:\n print(f\"\\n{'=' * 80}\")\n print(\n \"Step 3/3: Skipping control plane cleanup (tenant not found in control plane)\"\n )\n print(f\"{'=' * 80}\\n\")\n elif confirm_step(\n \"Step 3/3: Delete control plane records (tenant_notification, tenant_config, subscription, tenant)\",\n force,\n ):\n try:\n cleanup_control_plane(\n control_plane_pod, tenant_id, control_plane_context, force\n )\n except Exception as e:\n print(f\"✗ Failed at control plane cleanup step: {e}\", file=sys.stderr)\n if not force:\n print(\"Control plane cleanup failed\")\n else:\n print(\"[FORCE MODE] Control plane cleanup failed but continuing\")\n else:\n print(\"Step 3 skipped by user\")\n return False\n\n print(f\"\\n{'=' * 80}\")\n print(f\"✓ Cleanup completed for tenant: {tenant_id}\")\n print(f\"{'=' * 80}\")\n return True\n\n\ndef main() -> None:\n # Register signal handlers for graceful shutdown\n signal.signal(signal.SIGINT, signal_handler)\n signal.signal(signal.SIGTERM, signal_handler)\n\n if len(sys.argv) < 2:\n print(\n \"Usage: PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_cleanup_tenants.py \\\\\"\n )\n print(\n \" --data-plane-context --control-plane-context [--force]\"\n )\n print(\n \" PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_cleanup_tenants.py --csv \\\\\"\n )\n print(\n \" --data-plane-context --control-plane-context [--force]\"\n )\n print(\"\\nThis version runs ALL operations from pods (no bastion required)\")\n print(\"\\nArguments:\")\n print(\n \" tenant_id The tenant ID to clean up (required if not using --csv)\"\n )\n print(\n \" --csv PATH Path to CSV file containing tenant IDs to clean up\"\n )\n print(\" --force Skip all confirmation prompts (optional)\")\n print(\n \" --concurrency N Process N tenants concurrently (default: 1)\"\n )\n print(\n \" --data-plane-context CTX Kubectl context for data plane cluster (required)\"\n )\n print(\n \" --control-plane-context CTX Kubectl context for control plane cluster (required)\"\n )\n sys.exit(1)\n\n # Parse arguments\n force = \"--force\" in sys.argv\n tenant_ids = []\n\n # Parse concurrency\n concurrency: int = 1\n if \"--concurrency\" in sys.argv:\n try:\n concurrency_index = sys.argv.index(\"--concurrency\")\n if concurrency_index + 1 >= len(sys.argv):\n print(\"Error: --concurrency flag requires a number\", file=sys.stderr)\n sys.exit(1)\n concurrency = int(sys.argv[concurrency_index + 1])\n if concurrency < 1:\n print(\"Error: concurrency must be at least 1\", file=sys.stderr)\n sys.exit(1)\n except ValueError:\n print(\"Error: --concurrency value must be an integer\", file=sys.stderr)\n sys.exit(1)\n\n # Validate: concurrency > 1 requires --force\n if concurrency > 1 and not force:\n print(\n \"Error: --concurrency > 1 requires --force flag (interactive mode not supported with parallel processing)\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n # Parse contexts (required)\n data_plane_context: str | None = None\n control_plane_context: str | None = None\n\n if \"--data-plane-context\" in sys.argv:\n try:\n idx = sys.argv.index(\"--data-plane-context\")\n if idx + 1 >= len(sys.argv):\n print(\n \"Error: --data-plane-context requires a context name\",\n file=sys.stderr,\n )\n sys.exit(1)\n data_plane_context = sys.argv[idx + 1]\n except ValueError:\n pass\n\n if \"--control-plane-context\" in sys.argv:\n try:\n idx = sys.argv.index(\"--control-plane-context\")\n if idx + 1 >= len(sys.argv):\n print(\n \"Error: --control-plane-context requires a context name\",\n file=sys.stderr,\n )\n sys.exit(1)\n control_plane_context = sys.argv[idx + 1]\n except ValueError:\n pass\n\n # Validate required contexts\n if not data_plane_context:\n print(\n \"Error: --data-plane-context is required\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n if not control_plane_context:\n print(\n \"Error: --control-plane-context is required\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n # Check for CSV mode\n if \"--csv\" in sys.argv:\n try:\n csv_index = sys.argv.index(\"--csv\")\n if csv_index + 1 >= len(sys.argv):\n print(\"Error: --csv flag requires a file path\", file=sys.stderr)\n sys.exit(1)\n\n csv_path = sys.argv[csv_index + 1]\n tenant_ids = read_tenant_ids_from_csv(csv_path)\n\n if not tenant_ids:\n print(\"Error: No tenant IDs found in CSV file\", file=sys.stderr)\n sys.exit(1)\n\n print(f\"Found {len(tenant_ids)} tenant(s) in CSV file: {csv_path}\")\n\n except Exception as e:\n print(f\"Error reading CSV file: {e}\", file=sys.stderr)\n sys.exit(1)\n else:\n # Single tenant mode\n tenant_ids = [sys.argv[1]]\n\n # Initial confirmation (unless --force is used)\n if not force:\n print(f\"\\n{'=' * 80}\")\n print(\"TENANT CLEANUP - NO BASTION VERSION\")\n print(f\"{'=' * 80}\")\n if len(tenant_ids) == 1:\n print(f\"Tenant ID: {tenant_ids[0]}\")\n else:\n print(f\"Number of tenants: {len(tenant_ids)}\")\n print(f\"Tenant IDs: {', '.join(tenant_ids[:5])}\")\n if len(tenant_ids) > 5:\n print(f\" ... and {len(tenant_ids) - 5} more\")\n\n print(\"\\nThis will:\")\n print(\" 1. Check for remaining documents and connector credential pairs\")\n print(\" 2. Drop the data plane PostgreSQL schema (CASCADE)\")\n print(\" 3. Clean up control plane data (all via pod queries)\")\n print(f\"\\n{'=' * 80}\")\n print(\"WARNING: This operation is IRREVERSIBLE!\")\n print(f\"{'=' * 80}\\n\")\n\n response = input(\"Are you sure you want to proceed? Type 'yes' to confirm: \")\n\n if response.lower() != \"yes\":\n print(\"Cleanup aborted by user\")\n sys.exit(0)\n else:\n print(\n f\"⚠ FORCE MODE: Running cleanup for {len(tenant_ids)} tenant(s) without confirmations\"\n )\n\n # Find pods in both clusters before processing\n try:\n print(\"Finding data plane worker pod...\")\n data_plane_pod = find_worker_pod(data_plane_context)\n print(f\"✓ Using data plane worker pod: {data_plane_pod}\")\n\n print(\"Finding control plane pod...\")\n control_plane_pod = find_background_pod(control_plane_context)\n print(f\"✓ Using control plane pod: {control_plane_pod}\\n\")\n\n # Copy all scripts to data plane pod once\n setup_scripts_on_pod(data_plane_pod, data_plane_context)\n print()\n except Exception as e:\n print(f\"✗ Failed to find required pods or setup scripts: {e}\", file=sys.stderr)\n print(\"Cannot proceed with cleanup\")\n sys.exit(1)\n\n # Run cleanup for each tenant\n failed_tenants = []\n successful_tenants = []\n skipped_tenants = []\n\n # Open CSV file for writing successful cleanups in real-time\n csv_output_path = \"cleaned_tenants.csv\"\n with open(csv_output_path, \"w\", newline=\"\") as csv_file:\n csv_writer = csv.writer(csv_file)\n csv_writer.writerow([\"tenant_id\", \"cleaned_at\"])\n csv_file.flush()\n\n print(f\"Writing successful cleanups to: {csv_output_path}\\n\")\n\n if concurrency == 1:\n # Sequential processing\n for idx, tenant_id in enumerate(tenant_ids, 1):\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(f\"Processing tenant {idx}/{len(tenant_ids)}: {tenant_id}\")\n print(f\"{'=' * 80}\")\n\n try:\n was_cleaned = cleanup_tenant(\n tenant_id,\n data_plane_pod,\n control_plane_pod,\n data_plane_context,\n control_plane_context,\n force,\n )\n\n if was_cleaned:\n successful_tenants.append(tenant_id)\n\n # Write to CSV immediately after successful cleanup\n timestamp = datetime.utcnow().isoformat()\n csv_writer.writerow([tenant_id, timestamp])\n csv_file.flush()\n print(f\"✓ Recorded cleanup in {csv_output_path}\")\n else:\n skipped_tenants.append(tenant_id)\n print(f\"⚠ Tenant {tenant_id} was skipped (not recorded in CSV)\")\n\n except Exception as e:\n print(\n f\"✗ Cleanup failed for tenant {tenant_id}: {e}\", file=sys.stderr\n )\n failed_tenants.append((tenant_id, str(e)))\n\n # If not in force mode and there are more tenants, ask if we should continue\n if not force and idx < len(tenant_ids):\n response = input(\n f\"\\nContinue with remaining {len(tenant_ids) - idx} tenant(s)? (y/n): \"\n )\n if response.lower() != \"y\":\n print(\"Cleanup aborted by user\")\n break\n else:\n # Parallel processing\n print(\n f\"Processing {len(tenant_ids)} tenant(s) with concurrency={concurrency}\\n\"\n )\n\n def process_tenant(tenant_id: str) -> tuple[str, bool, str | None]:\n \"\"\"Process a single tenant. Returns (tenant_id, was_cleaned, error_message).\"\"\"\n try:\n was_cleaned = cleanup_tenant(\n tenant_id,\n data_plane_pod,\n control_plane_pod,\n data_plane_context,\n control_plane_context,\n force,\n )\n return (tenant_id, was_cleaned, None)\n except Exception as e:\n return (tenant_id, False, str(e))\n\n with ThreadPoolExecutor(max_workers=concurrency) as executor:\n # Submit all tasks\n future_to_tenant = {\n executor.submit(process_tenant, tenant_id): tenant_id\n for tenant_id in tenant_ids\n }\n\n # Process results as they complete\n completed = 0\n for future in as_completed(future_to_tenant):\n completed += 1\n tenant_id, was_cleaned, error = future.result()\n\n if error:\n with _print_lock:\n print(\n f\"[{completed}/{len(tenant_ids)}] ✗ Failed: {tenant_id}: {error}\",\n file=sys.stderr,\n )\n failed_tenants.append((tenant_id, error))\n elif was_cleaned:\n with _csv_lock:\n timestamp = datetime.utcnow().isoformat()\n csv_writer.writerow([tenant_id, timestamp])\n csv_file.flush()\n successful_tenants.append(tenant_id)\n with _print_lock:\n print(\n f\"[{completed}/{len(tenant_ids)}] ✓ Cleaned: {tenant_id}\"\n )\n else:\n skipped_tenants.append(tenant_id)\n with _print_lock:\n print(\n f\"[{completed}/{len(tenant_ids)}] ⊘ Skipped: {tenant_id}\"\n )\n\n # Print summary\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(\"CLEANUP SUMMARY\")\n print(f\"{'=' * 80}\")\n print(f\"Total tenants: {len(tenant_ids)}\")\n print(f\"Successful: {len(successful_tenants)}\")\n print(f\"Skipped: {len(skipped_tenants)}\")\n print(f\"Failed: {len(failed_tenants)}\")\n print(f\"\\nSuccessfully cleaned tenants written to: {csv_output_path}\")\n\n if skipped_tenants:\n print(f\"\\nSkipped tenants ({len(skipped_tenants)}):\")\n for tenant_id in skipped_tenants:\n print(f\" - {tenant_id}\")\n\n if failed_tenants:\n print(f\"\\nFailed tenants ({len(failed_tenants)}):\")\n for tenant_id, error in failed_tenants:\n print(f\" - {tenant_id}: {error}\")\n\n print(f\"{'=' * 80}\")\n\n if failed_tenants:\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/no_bastion_cleanup_utils.py", "content": "\"\"\"\nCleanup utilities that work WITHOUT bastion access.\nControl plane and data plane are in SEPARATE clusters.\n\"\"\"\n\nimport csv\nimport json\nimport subprocess\nimport sys\nfrom pathlib import Path\n\n\nclass TenantNotFoundInControlPlaneError(Exception):\n \"\"\"Exception raised when tenant/table is not found in control plane.\"\"\"\n\n\ndef find_worker_pod(context: str) -> str:\n \"\"\"Find a user file processing worker pod using kubectl.\n\n Args:\n context: kubectl context to use\n \"\"\"\n print(f\"Finding user file processing worker pod in context {context}...\")\n\n cmd = [\"kubectl\", \"get\", \"po\", \"--context\", context]\n\n result = subprocess.run(cmd, capture_output=True, text=True, check=True)\n\n # Parse output and find user file processing worker pod\n lines = result.stdout.strip().split(\"\\n\")\n lines = lines[1:] # Skip header\n\n import random\n\n random.shuffle(lines)\n\n for line in lines:\n if \"celery-worker-user-file-processing\" in line and \"Running\" in line:\n pod_name = line.split()[0]\n print(f\"Found pod: {pod_name}\")\n return pod_name\n\n raise RuntimeError(\"No running user file processing worker pod found\")\n\n\ndef find_background_pod(context: str) -> str:\n \"\"\"Find a pod for control plane operations.\n\n Args:\n context: kubectl context to use\n \"\"\"\n print(f\"Finding control plane pod in context {context}...\")\n\n cmd = [\"kubectl\", \"get\", \"po\", \"--context\", context]\n\n result = subprocess.run(cmd, capture_output=True, text=True, check=True)\n\n # Parse output and find suitable pod\n lines = result.stdout.strip().split(\"\\n\")\n lines = lines[1:] # Skip header\n\n import random\n\n random.shuffle(lines)\n\n # Try to find control plane pods\n for line in lines:\n if (\n any(\n name in line\n for name in [\n \"background-processing-deployment\",\n \"subscription-deployment\",\n \"tenants-deployment\",\n ]\n )\n and \"Running\" in line\n ):\n pod_name = line.split()[0]\n print(f\"Found pod: {pod_name}\")\n return pod_name\n\n raise RuntimeError(\"No suitable background pod found for control plane operations\")\n\n\ndef confirm_step(message: str, force: bool = False) -> bool:\n \"\"\"Ask for confirmation before executing a step.\n\n Args:\n message: The confirmation message to display\n force: If True, skip confirmation and return True\n\n Returns:\n True if user confirms or force is True, False otherwise\n \"\"\"\n if force:\n print(f\"[FORCE MODE] Skipping confirmation: {message}\")\n return True\n\n print(f\"\\n{message}\")\n response = input(\"Proceed? (y/n): \")\n return response.lower() == \"y\"\n\n\ndef execute_control_plane_query_from_pod(\n pod_name: str, query: str, context: str\n) -> dict:\n \"\"\"Execute a SQL query against control plane database from within a pod.\n\n Args:\n pod_name: The Kubernetes pod name to execute from\n query: The SQL query to execute\n context: kubectl context for control plane cluster\n\n Returns:\n Dict with 'success' bool, 'stdout' str, and optional 'error' str\n \"\"\"\n # Create a Python script to run the query\n # This script tries multiple environment variable patterns\n\n # NOTE: whuang 01/08/2026: POSTGRES_CONTROL_* don't exist. This uses pattern 2 currently.\n\n query_script = f'''\nimport os\nfrom sqlalchemy import create_engine, text\n\n# Try to get control plane database URL from various environment patterns\ncontrol_db_url = None\n\n# Pattern 1: POSTGRES_CONTROL_* variables\nif os.environ.get(\"POSTGRES_CONTROL_HOST\"):\n host = os.environ.get(\"POSTGRES_CONTROL_HOST\")\n port = os.environ.get(\"POSTGRES_CONTROL_PORT\", \"5432\")\n db = os.environ.get(\"POSTGRES_CONTROL_DB\", \"control\")\n user = os.environ.get(\"POSTGRES_CONTROL_USER\", \"postgres\")\n password = os.environ.get(\"POSTGRES_CONTROL_PASSWORD\", \"\")\n if password:\n control_db_url = f\"postgresql://{{user}}:{{password}}@{{host}}:{{port}}/{{db}}\"\n\n# Pattern 2: Standard POSTGRES_* variables (might point to control plane in this cluster)\nif not control_db_url and os.environ.get(\"POSTGRES_HOST\"):\n host = os.environ.get(\"POSTGRES_HOST\")\n port = os.environ.get(\"POSTGRES_PORT\", \"5432\")\n db = os.environ.get(\"POSTGRES_DB\", \"danswer\")\n user = os.environ.get(\"POSTGRES_USER\", \"postgres\")\n password = os.environ.get(\"POSTGRES_PASSWORD\", \"\")\n if password:\n control_db_url = f\"postgresql://{{user}}:{{password}}@{{host}}:{{port}}/{{db}}\"\n\n# Pattern 3: Direct URI\nif not control_db_url:\n control_db_url = os.environ.get(\"DATABASE_URL\") or os.environ.get(\"POSTGRES_URI\")\n\nif not control_db_url:\n raise ValueError(\"Cannot determine control plane database connection. No suitable environment variables found.\")\n\nengine = create_engine(control_db_url)\n\nwith engine.connect() as conn:\n result = conn.execute(text(\"\"\"{query}\"\"\"))\n\n # Check if this is a SELECT query\n if result.returns_rows:\n rows = [dict(row._mapping) for row in result]\n import json\n print(json.dumps(rows, default=str))\n else:\n # For INSERT/UPDATE/DELETE, print rowcount\n print(f\"{{result.rowcount}} rows affected\")\n\n conn.commit()\n'''\n\n # Write the script to a temp file on the pod\n script_path = \"/tmp/control_plane_query.py\"\n\n try:\n cmd_write = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_write.extend(\n [\n \"--\",\n \"bash\",\n \"-c\",\n f\"cat > {script_path} << 'EOFQUERY'\\n{query_script}\\nEOFQUERY\",\n ]\n )\n\n subprocess.run(\n cmd_write,\n check=True,\n capture_output=True,\n )\n\n # Execute the script\n cmd_exec = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_exec.extend([\"--\", \"python\", script_path])\n\n result = subprocess.run(\n cmd_exec,\n capture_output=True,\n text=True,\n check=True,\n )\n\n return {\n \"success\": True,\n \"stdout\": result.stdout.strip(),\n \"stderr\": result.stderr.strip() if result.stderr else \"\",\n }\n\n except subprocess.CalledProcessError as e:\n return {\n \"success\": False,\n \"stdout\": e.stdout if e.stdout else \"\",\n \"error\": e.stderr if e.stderr else str(e),\n }\n\n\ndef get_tenant_status(pod_name: str, tenant_id: str, context: str) -> str | None:\n \"\"\"\n Get tenant status from control plane database via pod.\n\n Args:\n pod_name: The pod to execute the query from\n tenant_id: The tenant ID to look up\n context: kubectl context for control plane cluster\n\n Returns:\n Tenant status string (e.g., 'GATED_ACCESS', 'ACTIVE') or None if not found\n\n Raises:\n TenantNotFoundInControlPlaneError: If the tenant record is not found in the table\n \"\"\"\n print(f\"Fetching tenant status for tenant: {tenant_id}\")\n\n query = f\"SELECT application_status FROM tenant WHERE tenant_id = '{tenant_id}'\"\n\n result = execute_control_plane_query_from_pod(pod_name, query, context)\n\n if not result[\"success\"]:\n error_msg = result.get(\"error\", \"Unknown error\")\n print(\n f\"✗ Failed to get tenant status for {tenant_id}: {error_msg}\",\n file=sys.stderr,\n )\n return None\n\n try:\n # Parse JSON output\n rows = json.loads(result[\"stdout\"])\n\n if rows and len(rows) > 0:\n status = rows[0].get(\"application_status\")\n if status:\n print(f\"✓ Tenant status: {status}\")\n return status\n\n # Tenant record not found in control plane table\n print(\"⚠ Tenant not found in control plane\")\n raise TenantNotFoundInControlPlaneError(\n f\"Tenant {tenant_id} not found in control plane database\"\n )\n\n except TenantNotFoundInControlPlaneError:\n # Re-raise without wrapping\n raise\n except (json.JSONDecodeError, KeyError, IndexError) as e:\n print(f\"✗ Failed to parse tenant status: {e}\", file=sys.stderr)\n return None\n\n\ndef execute_control_plane_delete(pod_name: str, query: str, context: str) -> bool:\n \"\"\"Execute a DELETE query against control plane database from pod.\n\n Args:\n pod_name: The pod to execute the query from\n query: The DELETE query to execute\n context: kubectl context for control plane cluster\n\n Returns:\n True if successful, False otherwise\n \"\"\"\n result = execute_control_plane_query_from_pod(pod_name, query, context)\n\n if result[\"success\"]:\n print(f\" {result['stdout']}\")\n return True\n else:\n print(f\" Error: {result.get('error', 'Unknown error')}\", file=sys.stderr)\n return False\n\n\ndef read_tenant_ids_from_csv(csv_path: str) -> list[str]:\n \"\"\"Read tenant IDs from CSV file.\n\n Args:\n csv_path: Path to CSV file\n\n Returns:\n List of tenant IDs\n \"\"\"\n if not Path(csv_path).exists():\n raise FileNotFoundError(f\"CSV file not found: {csv_path}\")\n\n tenant_ids = []\n with open(csv_path, \"r\", newline=\"\", encoding=\"utf-8\") as csvfile:\n reader = csv.DictReader(csvfile)\n\n # Check if tenant_id column exists\n if not reader.fieldnames or \"tenant_id\" not in reader.fieldnames:\n raise ValueError(\n f\"CSV file must have a 'tenant_id' column. Found columns: {reader.fieldnames}\"\n )\n\n for row in reader:\n tenant_id = row.get(\"tenant_id\", \"\").strip()\n if tenant_id:\n tenant_ids.append(tenant_id)\n\n return tenant_ids\n" }, { "path": "backend/scripts/tenant_cleanup/no_bastion_mark_connectors.py", "content": "#!/usr/bin/env python3\n\"\"\"\nMark connectors for deletion script that works WITHOUT bastion access.\nAll queries run directly from pods.\nSupports two-cluster architecture (data plane and control plane in separate clusters).\n\nUsage:\n PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_mark_connectors.py \\\n --data-plane-context --control-plane-context [--force]\n\n PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_mark_connectors.py --csv \\\n --data-plane-context --control-plane-context [--force] [--concurrency N]\n\"\"\"\n\nimport subprocess\nimport sys\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import ThreadPoolExecutor\nfrom pathlib import Path\nfrom threading import Lock\nfrom typing import Any\n\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import confirm_step\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import find_background_pod\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import find_worker_pod\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import get_tenant_status\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import read_tenant_ids_from_csv\nfrom scripts.tenant_cleanup.no_bastion_cleanup_utils import (\n TenantNotFoundInControlPlaneError,\n)\n\n# Global lock for thread-safe printing\n_print_lock: Lock = Lock()\n\n\ndef safe_print(*args: Any, **kwargs: Any) -> None:\n \"\"\"Thread-safe print function.\"\"\"\n with _print_lock:\n print(*args, **kwargs)\n\n\ndef run_connector_deletion(pod_name: str, tenant_id: str, context: str) -> None:\n \"\"\"Mark all connector credential pairs for deletion.\n\n Args:\n pod_name: Data plane pod to execute deletion on\n tenant_id: Tenant ID to process\n context: kubectl context for data plane cluster\n \"\"\"\n safe_print(\" Marking all connector credential pairs for deletion...\")\n\n # Get the path to the script\n script_dir = Path(__file__).parent\n mark_deletion_script = (\n script_dir / \"on_pod_scripts\" / \"execute_connector_deletion.py\"\n )\n\n if not mark_deletion_script.exists():\n raise FileNotFoundError(\n f\"execute_connector_deletion.py not found at {mark_deletion_script}\"\n )\n\n try:\n # Copy script to pod\n cmd_cp = [\"kubectl\", \"cp\", \"--context\", context]\n cmd_cp.extend(\n [\n str(mark_deletion_script),\n f\"{pod_name}:/tmp/execute_connector_deletion.py\",\n ]\n )\n\n subprocess.run(\n cmd_cp,\n check=True,\n capture_output=True,\n )\n\n # Execute script on pod\n cmd_exec = [\"kubectl\", \"exec\", \"--context\", context, pod_name]\n cmd_exec.extend(\n [\n \"--\",\n \"python\",\n \"/tmp/execute_connector_deletion.py\",\n tenant_id,\n \"--all\",\n ]\n )\n\n result = subprocess.run(cmd_exec)\n\n if result.returncode != 0:\n raise RuntimeError(result.stderr)\n\n except subprocess.CalledProcessError as e:\n safe_print(\n f\" ✗ Failed to mark all connector credential pairs for deletion: {e}\",\n file=sys.stderr,\n )\n if e.stderr:\n safe_print(f\" Error details: {e.stderr}\", file=sys.stderr)\n raise\n except Exception as e:\n safe_print(\n f\" ✗ Failed to mark all connector credential pairs for deletion: {e}\",\n file=sys.stderr,\n )\n raise\n\n\ndef mark_tenant_connectors_for_deletion(\n tenant_id: str,\n data_plane_pod: str,\n control_plane_pod: str,\n data_plane_context: str,\n control_plane_context: str,\n force: bool = False,\n) -> None:\n \"\"\"Main function to mark all connectors for a tenant for deletion.\n\n Args:\n tenant_id: Tenant ID to process\n data_plane_pod: Data plane pod for connector operations\n control_plane_pod: Control plane pod for status checks\n data_plane_context: kubectl context for data plane cluster\n control_plane_context: kubectl context for control plane cluster\n force: Skip confirmations if True\n \"\"\"\n safe_print(f\"Processing connectors for tenant: {tenant_id}\")\n\n # Check tenant status first (from control plane)\n safe_print(f\"\\n{'=' * 80}\")\n try:\n tenant_status = get_tenant_status(\n control_plane_pod, tenant_id, control_plane_context\n )\n\n # If tenant is not GATED_ACCESS, require explicit confirmation even in force mode\n if tenant_status and tenant_status != \"GATED_ACCESS\":\n safe_print(\n f\"\\n⚠️ WARNING: Tenant status is '{tenant_status}', not 'GATED_ACCESS'!\"\n )\n safe_print(\n \"This tenant may be active and should not have connectors deleted without careful review.\"\n )\n safe_print(f\"{'=' * 80}\\n\")\n\n # Always ask for confirmation if not gated, even in force mode\n if not force:\n response = input(\n \"Are you ABSOLUTELY SURE you want to proceed? Type 'yes' to confirm: \"\n )\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - tenant is not GATED_ACCESS\")\n raise RuntimeError(f\"Tenant {tenant_id} is not GATED_ACCESS\")\n else:\n raise RuntimeError(f\"Tenant {tenant_id} is not GATED_ACCESS\")\n elif tenant_status == \"GATED_ACCESS\":\n safe_print(\"✓ Tenant status is GATED_ACCESS - safe to proceed\")\n elif tenant_status is None:\n safe_print(\"⚠️ WARNING: Could not determine tenant status!\")\n if not force:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - could not verify tenant status\")\n raise RuntimeError(\n f\"Could not verify tenant status for {tenant_id}\"\n )\n else:\n raise RuntimeError(f\"Could not verify tenant status for {tenant_id}\")\n except TenantNotFoundInControlPlaneError as e:\n # Tenant/table not found in control plane\n error_str = str(e)\n safe_print(f\"⚠️ WARNING: Tenant not found in control plane: {error_str}\")\n\n if force:\n safe_print(\n \"[FORCE MODE] Tenant not found in control plane - continuing with connector deletion anyway\"\n )\n else:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - tenant not found in control plane\")\n raise RuntimeError(f\"Tenant {tenant_id} not found in control plane\")\n except RuntimeError:\n # Re-raise RuntimeError (from status checks above) without wrapping\n raise\n except Exception as e:\n safe_print(f\"⚠️ WARNING: Failed to check tenant status: {e}\")\n if not force:\n response = input(\"Continue anyway? Type 'yes' to confirm: \")\n if response.lower() != \"yes\":\n safe_print(\"Operation aborted - could not verify tenant status\")\n raise\n else:\n raise RuntimeError(f\"Failed to check tenant status for {tenant_id}\")\n safe_print(f\"{'=' * 80}\\n\")\n\n # Confirm before proceeding (only in non-force mode)\n if not confirm_step(\n f\"Mark all connector credential pairs for deletion for tenant {tenant_id}?\",\n force,\n ):\n safe_print(\"Operation cancelled by user\")\n raise ValueError(\"Operation cancelled by user\")\n\n run_connector_deletion(data_plane_pod, tenant_id, data_plane_context)\n\n # Print summary\n safe_print(\n f\"✓ Marked all connector credential pairs for deletion for tenant {tenant_id}\"\n )\n\n\ndef main() -> None:\n if len(sys.argv) < 2:\n print(\n \"Usage: PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_mark_connectors.py \\\\\"\n )\n print(\n \" --data-plane-context --control-plane-context [--force]\"\n )\n print(\n \" PYTHONPATH=. python scripts/tenant_cleanup/no_bastion_mark_connectors.py --csv \\\\\"\n )\n print(\n \" --data-plane-context --control-plane-context [--force] [--concurrency N]\"\n )\n print(\"\\nThis version runs ALL operations from pods (no bastion required)\")\n print(\"\\nArguments:\")\n print(\n \" tenant_id The tenant ID to process (required if not using --csv)\"\n )\n print(\n \" --csv PATH Path to CSV file containing tenant IDs to process\"\n )\n print(\" --force Skip all confirmation prompts (optional)\")\n print(\n \" --concurrency N Process N tenants concurrently (default: 1)\"\n )\n print(\n \" --data-plane-context CTX Kubectl context for data plane cluster (required)\"\n )\n print(\n \" --control-plane-context CTX Kubectl context for control plane cluster (required)\"\n )\n sys.exit(1)\n\n # Parse arguments\n force = \"--force\" in sys.argv\n tenant_ids: list[str] = []\n\n # Parse contexts (required)\n data_plane_context: str | None = None\n control_plane_context: str | None = None\n\n if \"--data-plane-context\" in sys.argv:\n try:\n idx = sys.argv.index(\"--data-plane-context\")\n if idx + 1 >= len(sys.argv):\n print(\n \"Error: --data-plane-context requires a context name\",\n file=sys.stderr,\n )\n sys.exit(1)\n data_plane_context = sys.argv[idx + 1]\n except ValueError:\n pass\n\n if \"--control-plane-context\" in sys.argv:\n try:\n idx = sys.argv.index(\"--control-plane-context\")\n if idx + 1 >= len(sys.argv):\n print(\n \"Error: --control-plane-context requires a context name\",\n file=sys.stderr,\n )\n sys.exit(1)\n control_plane_context = sys.argv[idx + 1]\n except ValueError:\n pass\n\n # Validate required contexts\n if not data_plane_context:\n print(\n \"Error: --data-plane-context is required\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n if not control_plane_context:\n print(\n \"Error: --control-plane-context is required\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n # Parse concurrency\n concurrency: int = 1\n if \"--concurrency\" in sys.argv:\n try:\n concurrency_index = sys.argv.index(\"--concurrency\")\n if concurrency_index + 1 >= len(sys.argv):\n print(\"Error: --concurrency flag requires a number\", file=sys.stderr)\n sys.exit(1)\n concurrency = int(sys.argv[concurrency_index + 1])\n if concurrency < 1:\n print(\"Error: concurrency must be at least 1\", file=sys.stderr)\n sys.exit(1)\n except ValueError:\n print(\"Error: --concurrency value must be an integer\", file=sys.stderr)\n sys.exit(1)\n\n # Validate: concurrency > 1 requires --force\n if concurrency > 1 and not force:\n print(\n \"Error: --concurrency > 1 requires --force flag (interactive mode not supported with parallel processing)\",\n file=sys.stderr,\n )\n sys.exit(1)\n\n # Check for CSV mode\n if \"--csv\" in sys.argv:\n try:\n csv_index: int = sys.argv.index(\"--csv\")\n if csv_index + 1 >= len(sys.argv):\n print(\"Error: --csv flag requires a file path\", file=sys.stderr)\n sys.exit(1)\n\n csv_path: str = sys.argv[csv_index + 1]\n tenant_ids = read_tenant_ids_from_csv(csv_path)\n\n if not tenant_ids:\n print(\"Error: No tenant IDs found in CSV file\", file=sys.stderr)\n sys.exit(1)\n\n print(f\"Found {len(tenant_ids)} tenant(s) in CSV file: {csv_path}\")\n\n except Exception as e:\n print(f\"Error reading CSV file: {e}\", file=sys.stderr)\n sys.exit(1)\n else:\n # Single tenant mode\n tenant_ids = [sys.argv[1]]\n\n # Find pods in both clusters before processing\n try:\n print(\"Finding data plane worker pod...\")\n data_plane_pod: str = find_worker_pod(data_plane_context)\n print(f\"✓ Using data plane worker pod: {data_plane_pod}\")\n\n print(\"Finding control plane pod...\")\n control_plane_pod: str = find_background_pod(control_plane_context)\n print(f\"✓ Using control plane pod: {control_plane_pod}\")\n except Exception as e:\n print(f\"✗ Failed to find required pods: {e}\", file=sys.stderr)\n print(\"Cannot proceed with marking connectors for deletion\")\n sys.exit(1)\n\n # Initial confirmation (unless --force is used)\n if not force:\n print(f\"\\n{'=' * 80}\")\n print(\"MARK CONNECTORS FOR DELETION - NO BASTION VERSION\")\n print(f\"{'=' * 80}\")\n if len(tenant_ids) == 1:\n print(f\"Tenant ID: {tenant_ids[0]}\")\n else:\n print(f\"Number of tenants: {len(tenant_ids)}\")\n print(f\"Tenant IDs: {', '.join(tenant_ids[:5])}\")\n if len(tenant_ids) > 5:\n print(f\" ... and {len(tenant_ids) - 5} more\")\n\n print(\n f\"Mode: {'FORCE (no confirmations)' if force else 'Interactive (will ask for confirmation at each step)'}\"\n )\n print(f\"Concurrency: {concurrency} tenant(s) at a time\")\n print(\"\\nThis will:\")\n print(\" 1. Fetch all connector credential pairs for each tenant\")\n print(\" 2. Cancel any scheduled indexing attempts for each connector\")\n print(\" 3. Mark each connector credential pair status as DELETING\")\n print(\" 4. Trigger the connector deletion task\")\n print(f\"\\n{'=' * 80}\")\n print(\"WARNING: This will mark connectors for deletion!\")\n print(\"The actual deletion will be performed by the background celery worker.\")\n print(f\"{'=' * 80}\\n\")\n\n response = input(\"Are you sure you want to proceed? Type 'yes' to confirm: \")\n\n if response.lower() != \"yes\":\n print(\"Operation aborted by user\")\n sys.exit(0)\n else:\n if len(tenant_ids) == 1:\n print(\n f\"⚠ FORCE MODE: Marking connectors for deletion for {tenant_ids[0]} without confirmations\"\n )\n else:\n print(\n f\"⚠ FORCE MODE: Marking connectors for deletion for {len(tenant_ids)} tenants \"\n f\"(concurrency: {concurrency}) without confirmations\"\n )\n\n # Process tenants (in parallel if concurrency > 1)\n failed_tenants: list[tuple[str, str]] = []\n successful_tenants: list[str] = []\n\n if concurrency == 1:\n # Sequential processing\n for idx, tenant_id in enumerate(tenant_ids, 1):\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(f\"Processing tenant {idx}/{len(tenant_ids)}: {tenant_id}\")\n print(f\"{'=' * 80}\")\n\n try:\n mark_tenant_connectors_for_deletion(\n tenant_id,\n data_plane_pod,\n control_plane_pod,\n data_plane_context,\n control_plane_context,\n force,\n )\n successful_tenants.append(tenant_id)\n except Exception as e:\n print(\n f\"✗ Failed to process tenant {tenant_id}: {e}\",\n file=sys.stderr,\n )\n failed_tenants.append((tenant_id, str(e)))\n\n # If not in force mode and there are more tenants, ask if we should continue\n if not force and idx < len(tenant_ids):\n response = input(\n f\"\\nContinue with remaining {len(tenant_ids) - idx} tenant(s)? (y/n): \"\n )\n if response.lower() != \"y\":\n print(\"Operation aborted by user\")\n break\n else:\n # Parallel processing\n print(\n f\"\\nProcessing {len(tenant_ids)} tenant(s) with concurrency={concurrency}\"\n )\n\n def process_tenant(tenant_id: str) -> tuple[str, bool, str | None]:\n \"\"\"Process a single tenant. Returns (tenant_id, success, error_message).\"\"\"\n try:\n mark_tenant_connectors_for_deletion(\n tenant_id,\n data_plane_pod,\n control_plane_pod,\n data_plane_context,\n control_plane_context,\n force,\n )\n return (tenant_id, True, None)\n except Exception as e:\n return (tenant_id, False, str(e))\n\n with ThreadPoolExecutor(max_workers=concurrency) as executor:\n # Submit all tasks\n future_to_tenant = {\n executor.submit(process_tenant, tenant_id): tenant_id\n for tenant_id in tenant_ids\n }\n\n # Process results as they complete\n completed: int = 0\n for future in as_completed(future_to_tenant):\n completed += 1\n tenant_id, success, error = future.result()\n\n if success:\n successful_tenants.append(tenant_id)\n safe_print(\n f\"[{completed}/{len(tenant_ids)}] ✓ Successfully processed {tenant_id}\"\n )\n else:\n failed_tenants.append((tenant_id, error or \"Unknown error\"))\n safe_print(\n f\"[{completed}/{len(tenant_ids)}] ✗ Failed to process {tenant_id}: {error}\",\n file=sys.stderr,\n )\n\n # Print summary if multiple tenants\n if len(tenant_ids) > 1:\n print(f\"\\n{'=' * 80}\")\n print(\"OPERATION SUMMARY\")\n print(f\"{'=' * 80}\")\n print(f\"Total tenants: {len(tenant_ids)}\")\n print(f\"Successful: {len(successful_tenants)}\")\n print(f\"Failed: {len(failed_tenants)}\")\n\n if failed_tenants:\n print(\"\\nFailed tenants:\")\n for tenant_id, error in failed_tenants:\n print(f\" - {tenant_id}: {error}\")\n\n print(f\"{'=' * 80}\")\n\n if failed_tenants:\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/check_documents_deleted.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to check for remaining ConnectorCredentialPairs and Documents in a tenant's schema.\nMust be run on a pod with access to the data plane PostgreSQL database.\n\nUsage:\n python check_documents_deleted.py \n\nOutput:\n JSON object with status, message, and counts of remaining records\n\"\"\"\n\nimport json\nimport sys\n\nfrom sqlalchemy import func\nfrom sqlalchemy import select\n\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Document\n\n\ndef check_documents_deleted(tenant_id: str) -> dict:\n \"\"\"\n Check for remaining ConnectorCredentialPairs and Documents in tenant schema.\n\n Args:\n tenant_id: The tenant ID to query\n\n Returns:\n Dictionary with status and counts of remaining records\n \"\"\"\n try:\n print(\n f\"Checking for remaining documents in tenant: {tenant_id}\",\n file=sys.stderr,\n )\n\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # Count ConnectorCredentialPairs\n cc_count = db_session.scalar(\n select(func.count()).select_from(ConnectorCredentialPair)\n )\n\n # Count Documents\n doc_count = db_session.scalar(select(func.count()).select_from(Document))\n\n # Handle None values from scalar (should not happen but mypy needs it)\n cc_count = cc_count or 0\n doc_count = doc_count or 0\n\n # If any records remain beyond acceptable thresholds, return error status\n is_deletable = cc_count == 0 or doc_count <= 5\n if not is_deletable:\n return {\n \"status\": \"error\",\n \"message\": (\n f\"Found {cc_count} ConnectorCredentialPair(s) and {doc_count} Document(s) \"\n \"still remaining. Must have 0 ConnectorCredentialPairs and no more than \"\n \"5 Documents before cleanup.\"\n ),\n \"connector_credential_pair_count\": cc_count,\n \"document_count\": doc_count,\n }\n\n # All clear\n return {\n \"status\": \"success\",\n \"message\": \"No ConnectorCredentialPairs or Documents found - safe to proceed\",\n \"connector_credential_pair_count\": 0,\n \"document_count\": 0,\n }\n\n except Exception as e:\n error_msg = str(e)\n print(f\"Error checking documents: {error_msg}\", file=sys.stderr)\n # Check if it's a schema not found error\n if \"does not exist\" in error_msg:\n return {\n \"status\": \"not_found\",\n \"message\": f\"Schema '{tenant_id}' does not exist\",\n }\n return {\"status\": \"error\", \"message\": f\"Error checking documents: {error_msg}\"}\n\n\ndef main() -> None:\n if len(sys.argv) != 2:\n print(\n json.dumps(\n {\n \"status\": \"error\",\n \"message\": \"Usage: python check_documents_deleted.py \",\n }\n )\n )\n sys.exit(1)\n\n tenant_id = sys.argv[1]\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n result = check_documents_deleted(tenant_id)\n print(json.dumps(result))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/cleanup_tenant_schema.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to drop a tenant's PostgreSQL schema.\nDesigned to be run on a heavy worker pod.\n\nUsage:\n python cleanup_tenant_schema.py \n\"\"\"\n\nimport json\nimport sys\n\nfrom sqlalchemy import text\n\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.engine.sql_engine import SqlEngine\n\n\ndef drop_data_plane_schema(tenant_id: str) -> dict[str, str]:\n \"\"\"Drop the PostgreSQL schema for the given tenant.\"\"\"\n print(f\"Dropping data plane schema for tenant: {tenant_id}\", file=sys.stderr)\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n try:\n with get_session_with_shared_schema() as session:\n # First, verify the schema exists\n check_schema_query = text(\n \"\"\"\n SELECT nspname\n FROM pg_namespace\n WHERE nspname = :schema_name\n \"\"\"\n )\n\n result = session.execute(\n check_schema_query, {\"schema_name\": tenant_id}\n ).fetchone()\n\n if not result:\n print(f\"Schema {tenant_id} does not exist\", file=sys.stderr)\n return {\n \"status\": \"not_found\",\n \"message\": f\"Schema {tenant_id} does not exist\",\n }\n\n # Drop the schema with CASCADE to remove all objects within it\n drop_schema_query = text(f'DROP SCHEMA IF EXISTS \"{tenant_id}\" CASCADE')\n session.execute(drop_schema_query)\n session.commit()\n\n print(f\"Successfully dropped schema: {tenant_id}\", file=sys.stderr)\n\n # Delete the tenant mapping from user_tenant_mapping table\n delete_mapping_query = text(\n \"\"\"\n DELETE FROM user_tenant_mapping\n WHERE tenant_id = :tenant_id\n \"\"\"\n )\n session.execute(delete_mapping_query, {\"tenant_id\": tenant_id})\n session.commit()\n\n print(\n f\"Successfully deleted tenant mapping for: {tenant_id}\", file=sys.stderr\n )\n return {\n \"status\": \"success\",\n \"message\": f\"Successfully dropped schema: {tenant_id}\",\n }\n\n except Exception as e:\n print(f\"Failed to drop schema for tenant {tenant_id}: {e}\", file=sys.stderr)\n return {\"status\": \"error\", \"message\": str(e)}\n\n\nif __name__ == \"__main__\":\n if len(sys.argv) < 2:\n print(\"Usage: python cleanup_tenant_schema.py \", file=sys.stderr)\n sys.exit(1)\n\n tenant_id = sys.argv[1]\n\n result = drop_data_plane_schema(tenant_id)\n\n # Output result as JSON to stdout for easy parsing\n print(json.dumps(result))\n\n # Exit with error code if failed\n if result[\"status\"] == \"error\":\n sys.exit(1)\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/execute_connector_deletion.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to mark connector credential pairs for deletion.\nRuns on a Kubernetes pod with access to the data plane database.\n\nUsage:\n # Mark a specific connector for deletion\n python mark_connector_for_deletion.py \n\n # Mark all connectors for deletion\n python mark_connector_for_deletion.py --all\n\nOutput:\n JSON to stdout with structure:\n {\n \"status\": \"success\" | \"error\",\n \"message\": str,\n \"deleted_count\": int (when using --all),\n \"timing\": {\n \"total_seconds\": float,\n \"per_connector\": [...]\n }\n }\n\"\"\"\n\nimport json\nimport sys\nimport time\nfrom typing import Any\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.versioned_apps.client import app as client_app\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.connector_credential_pair import get_connector_credential_pairs\nfrom onyx.db.connector_credential_pair import update_connector_credential_pair_from_id\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair\n\n\ndef mark_connector_for_deletion(\n tenant_id: str, cc_pair_id: int, db_session: Session | None = None\n) -> dict[str, Any]:\n \"\"\"Mark a connector credential pair for deletion.\n\n Args:\n tenant_id: The tenant ID\n cc_pair_id: The connector credential pair ID\n db_session: Optional database session (if None, creates a new one)\n\n Returns:\n Dict with status, message, and timing\n \"\"\"\n timing: dict[str, float] = {}\n start_time: float = time.time()\n\n try:\n print(\n f\"Marking connector credential pair {cc_pair_id} for deletion\",\n file=sys.stderr,\n )\n\n def _mark_deletion(db_sess: Session) -> dict[str, Any]:\n # Get the connector credential pair\n fetch_start: float = time.time()\n cc_pair = get_connector_credential_pair_from_id(\n db_session=db_sess,\n cc_pair_id=cc_pair_id,\n )\n timing[\"fetch_cc_pair_seconds\"] = time.time() - fetch_start\n\n if not cc_pair:\n return {\n \"status\": \"error\",\n \"message\": f\"Connector credential pair {cc_pair_id} not found\",\n \"timing\": timing,\n }\n\n # Cancel any scheduled indexing attempts\n print(\n f\"Canceling indexing attempts for CC pair {cc_pair_id}\",\n file=sys.stderr,\n )\n cancel_start: float = time.time()\n cancel_indexing_attempts_for_ccpair(\n cc_pair_id=cc_pair.id,\n db_session=db_sess,\n include_secondary_index=True,\n )\n timing[\"cancel_indexing_seconds\"] = time.time() - cancel_start\n\n # Mark as deleting\n print(\n f\"Updating CC pair {cc_pair_id} status to DELETING\",\n file=sys.stderr,\n )\n update_start: float = time.time()\n update_connector_credential_pair_from_id(\n db_session=db_sess,\n cc_pair_id=cc_pair.id,\n status=ConnectorCredentialPairStatus.DELETING,\n )\n timing[\"update_status_seconds\"] = time.time() - update_start\n\n commit_start: float = time.time()\n db_sess.commit()\n timing[\"commit_seconds\"] = time.time() - commit_start\n\n return {\n \"status\": \"success\",\n \"message\": f\"Marked connector credential pair {cc_pair_id} for deletion\",\n \"timing\": timing,\n }\n\n result: dict[str, Any]\n if db_session:\n result = _mark_deletion(db_session)\n else:\n with get_session_with_tenant(tenant_id=tenant_id) as db_sess:\n result = _mark_deletion(db_sess)\n\n # Trigger the deletion check task\n print(\n \"Triggering connector deletion check task\",\n file=sys.stderr,\n )\n task_start: float = time.time()\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,\n priority=OnyxCeleryPriority.HIGH,\n kwargs={\"tenant_id\": tenant_id},\n )\n timing[\"send_task_seconds\"] = time.time() - task_start\n timing[\"total_seconds\"] = time.time() - start_time\n\n result[\"timing\"] = timing\n return result\n\n except Exception as e:\n print(\n f\"Error marking connector for deletion: {e}\",\n file=sys.stderr,\n )\n timing[\"total_seconds\"] = time.time() - start_time\n return {\n \"status\": \"error\",\n \"message\": str(e),\n \"timing\": timing,\n }\n\n\ndef mark_all_connectors_for_deletion(tenant_id: str) -> dict[str, Any]:\n \"\"\"Mark all connector credential pairs for a tenant for deletion.\n\n Args:\n tenant_id: The tenant ID\n\n Returns:\n Dict with status, message, deleted_count, and timing\n \"\"\"\n overall_start: float = time.time()\n per_connector_timing: list[dict[str, Any]] = []\n\n try:\n print(\n f\"Marking all connector credential pairs for tenant {tenant_id} for deletion\",\n file=sys.stderr,\n )\n\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # Get all connector credential pairs\n fetch_all_start: float = time.time()\n cc_pairs = get_connector_credential_pairs(db_session=db_session)\n fetch_all_time: float = time.time() - fetch_all_start\n\n print(\n f\"Found {len(cc_pairs)} connector credential pairs to delete\",\n file=sys.stderr,\n )\n\n if not cc_pairs:\n return {\n \"status\": \"success\",\n \"message\": \"No connector credential pairs found for tenant\",\n \"deleted_count\": 0,\n \"timing\": {\n \"fetch_all_seconds\": fetch_all_time,\n \"total_seconds\": time.time() - overall_start,\n },\n }\n\n deleted_count: int = 0\n errors: list[str] = []\n\n for cc_pair in cc_pairs:\n connector_start: float = time.time()\n print(\n f\"Processing CC pair {cc_pair.id} ({deleted_count + 1}/{len(cc_pairs)})\",\n file=sys.stderr,\n )\n\n # Cancel any scheduled indexing attempts\n cancel_start: float = time.time()\n cancel_indexing_attempts_for_ccpair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n include_secondary_index=True,\n )\n cancel_time: float = time.time() - cancel_start\n\n # Mark as deleting\n update_start: float = time.time()\n try:\n update_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair.id,\n status=ConnectorCredentialPairStatus.DELETING,\n )\n deleted_count += 1\n except Exception as e:\n errors.append(f\"CC pair {cc_pair.id}: {str(e)}\")\n print(\n f\"Error updating CC pair {cc_pair.id}: {e}\",\n file=sys.stderr,\n )\n\n update_time: float = time.time() - update_start\n connector_total_time: float = time.time() - connector_start\n\n per_connector_timing.append(\n {\n \"cc_pair_id\": cc_pair.id,\n \"cancel_indexing_seconds\": cancel_time,\n \"update_status_seconds\": update_time,\n \"total_seconds\": connector_total_time,\n }\n )\n\n # Commit all changes\n commit_start: float = time.time()\n db_session.commit()\n commit_time: float = time.time() - commit_start\n\n # Trigger the deletion check task\n print(\n \"Triggering connector deletion check task\",\n file=sys.stderr,\n )\n task_start: float = time.time()\n client_app.send_task(\n OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,\n priority=OnyxCeleryPriority.HIGH,\n kwargs={\"tenant_id\": tenant_id},\n )\n task_time: float = time.time() - task_start\n\n total_time: float = time.time() - overall_start\n\n result: dict[str, Any] = {\n \"status\": \"success\",\n \"message\": f\"Marked {deleted_count} connector credential pairs for deletion\",\n \"deleted_count\": deleted_count,\n \"timing\": {\n \"fetch_all_seconds\": fetch_all_time,\n \"commit_seconds\": commit_time,\n \"send_task_seconds\": task_time,\n \"total_seconds\": total_time,\n \"per_connector\": per_connector_timing,\n },\n }\n\n if errors:\n result[\"errors\"] = errors\n\n return result\n\n except Exception as e:\n print(\n f\"Error marking all connectors for deletion: {e}\",\n file=sys.stderr,\n )\n return {\n \"status\": \"error\",\n \"message\": str(e),\n \"timing\": {\n \"total_seconds\": time.time() - overall_start,\n \"per_connector\": per_connector_timing,\n },\n }\n\n\ndef main() -> None:\n if len(sys.argv) < 2 or len(sys.argv) > 3:\n print(\n json.dumps(\n {\n \"status\": \"error\",\n \"message\": \"Usage: python mark_connector_for_deletion.py [|--all]\",\n }\n )\n )\n sys.exit(1)\n\n tenant_id: str = sys.argv[1]\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n result: dict[str, Any]\n # Check if we should mark all connectors or just one\n if len(sys.argv) == 3:\n second_arg: str = sys.argv[2]\n if second_arg == \"--all\":\n result = mark_all_connectors_for_deletion(tenant_id)\n else:\n try:\n cc_pair_id: int = int(second_arg)\n result = mark_connector_for_deletion(tenant_id, cc_pair_id)\n except ValueError:\n print(\n json.dumps(\n {\n \"status\": \"error\",\n \"message\": \"cc_pair_id must be an integer or use --all\",\n }\n )\n )\n sys.exit(1)\n else:\n # If only tenant_id is provided, show error\n print(\n json.dumps(\n {\n \"status\": \"error\",\n \"message\": \"Usage: python mark_connector_for_deletion.py [|--all]\",\n }\n )\n )\n sys.exit(1)\n\n print(json.dumps(result, indent=2))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_connectors.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to fetch connector credential pairs for a tenant.\nRuns on a Kubernetes pod with access to the data plane database.\n\nUsage:\n python get_tenant_connectors.py \n\nOutput:\n JSON to stdout with structure:\n {\n \"status\": \"success\" | \"error\",\n \"connectors\": [\n {\n \"id\": int,\n \"connector_id\": int,\n \"credential_id\": int,\n \"name\": str,\n \"status\": str\n },\n ...\n ] (if success),\n \"message\": str (if error)\n }\n\"\"\"\n\nimport json\nimport sys\n\nfrom sqlalchemy import select\n\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.models import ConnectorCredentialPair\n\n\ndef get_tenant_connectors(tenant_id: str) -> dict:\n \"\"\"Get all connector credential pairs for a tenant.\n\n Args:\n tenant_id: The tenant ID to query\n\n Returns:\n Dict with status and list of connectors or error message\n \"\"\"\n try:\n print(\n f\"Fetching connector credential pairs for tenant: {tenant_id}\",\n file=sys.stderr,\n )\n\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # Get all connector credential pairs\n stmt = select(ConnectorCredentialPair)\n cc_pairs = db_session.execute(stmt).scalars().all()\n\n connectors = [\n {\n \"id\": cc.id,\n \"connector_id\": cc.connector_id,\n \"credential_id\": cc.credential_id,\n \"name\": cc.name,\n \"status\": cc.status.value,\n }\n for cc in cc_pairs\n ]\n\n print(\n f\"Found {len(connectors)} connector credential pair(s)\",\n file=sys.stderr,\n )\n\n return {\n \"status\": \"success\",\n \"connectors\": connectors,\n }\n\n except Exception as e:\n print(f\"Error fetching connectors: {e}\", file=sys.stderr)\n return {\n \"status\": \"error\",\n \"message\": str(e),\n }\n\n\ndef main() -> None:\n if len(sys.argv) != 2:\n print(\n json.dumps(\n {\n \"status\": \"error\",\n \"message\": \"Usage: python get_tenant_connectors.py \",\n }\n )\n )\n sys.exit(1)\n\n tenant_id = sys.argv[1]\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n result = get_tenant_connectors(tenant_id)\n print(json.dumps(result))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_index_name.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to get the default index name for a tenant.\nDesigned to be run on a heavy worker pod.\n\nUsage:\n python get_tenant_index_name.py \n\"\"\"\n\nimport json\nimport sys\n\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.search_settings import get_current_search_settings\n\n\ndef get_tenant_index_name(tenant_id: str) -> dict[str, str]:\n \"\"\"Get the default index name for the given tenant.\"\"\"\n print(f\"Getting default index name for tenant: {tenant_id}\", file=sys.stderr)\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n try:\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n search_settings = get_current_search_settings(db_session)\n index_name = search_settings.index_name\n print(f\"Found index name: {index_name}\", file=sys.stderr)\n return {\"status\": \"success\", \"index_name\": index_name}\n\n except Exception as e:\n print(f\"Failed to get index name for tenant {tenant_id}: {e}\", file=sys.stderr)\n return {\"status\": \"error\", \"message\": str(e)}\n\n\nif __name__ == \"__main__\":\n if len(sys.argv) < 2:\n print(\"Usage: python get_tenant_index_name.py \", file=sys.stderr)\n sys.exit(1)\n\n tenant_id = sys.argv[1]\n\n result = get_tenant_index_name(tenant_id)\n\n # Output result as JSON to stdout for easy parsing\n print(json.dumps(result))\n\n # Exit with error code if failed\n if result[\"status\"] == \"error\":\n sys.exit(1)\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/get_tenant_users.py", "content": "#!/usr/bin/env python3\n\"\"\"\nScript to fetch user emails from a tenant's data plane schema.\nMust be run on a pod with access to the data plane PostgreSQL database.\n\nUsage:\n python get_tenant_users.py \n\nOutput:\n JSON object with status and users list\n\"\"\"\n\nimport json\nimport sys\n\nfrom sqlalchemy import select\n\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.models import User\n\n\ndef get_tenant_users(tenant_id: str) -> dict:\n \"\"\"\n Fetch user emails from the tenant's data plane schema.\n\n Args:\n tenant_id: The tenant ID to query\n\n Returns:\n Dictionary with status and users list\n \"\"\"\n try:\n print(f\"Querying users for tenant: {tenant_id}\", file=sys.stderr)\n\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # Query users from the tenant schema\n # Select only the email column\n user_email_column = User.__table__.c.email\n stmt = select(user_email_column).order_by(user_email_column)\n result = db_session.execute(stmt)\n users = [row[0] for row in result]\n\n return {\"status\": \"success\", \"users\": users}\n\n except Exception as e:\n error_msg = str(e)\n print(f\"Error fetching users: {error_msg}\", file=sys.stderr)\n # Check if it's a schema not found error\n if \"does not exist\" in error_msg:\n return {\n \"status\": \"not_found\",\n \"message\": f\"Schema '{tenant_id}' does not exist\",\n \"users\": [],\n }\n return {\"status\": \"error\", \"message\": error_msg, \"users\": []}\n\n\ndef main() -> None:\n if len(sys.argv) != 2:\n print(\n json.dumps(\n {\n \"status\": \"error\",\n \"message\": \"Usage: python get_tenant_users.py \",\n }\n )\n )\n sys.exit(1)\n\n tenant_id = sys.argv[1]\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n result = get_tenant_users(tenant_id)\n print(json.dumps(result))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/tenant_cleanup/on_pod_scripts/understand_tenants.py", "content": "import json\nimport sys\nfrom typing import Any\n\nfrom sqlalchemy import text\nfrom sqlalchemy.exc import ProgrammingError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom onyx.db.engine.sql_engine import SqlEngine\n\n\ndef get_tenant_activity_summary(session: Session) -> list[dict[str, Any]]:\n \"\"\"Return a list of dicts, one per tenant, with last query info, doc count, and user count.\"\"\"\n\n # Step 1: fetch all tenant schemas\n tenant_schemas = [\n row[0]\n for row in session.execute(\n text(\n \"\"\"\n SELECT nspname\n FROM pg_namespace\n WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'public')\n AND nspname NOT LIKE 'pg_toast%%'\n AND nspname NOT LIKE 'pg_temp%%'\n ORDER BY nspname\n \"\"\"\n )\n )\n ]\n\n print(f\"Found {len(tenant_schemas)} tenant schemas\", file=sys.stderr)\n\n summaries = []\n\n # Step 2: loop through each tenant schema\n for idx, schema in enumerate(tenant_schemas):\n if idx % 100 == 0:\n print(f\"Processing tenant {idx}/{len(tenant_schemas)}\", file=sys.stderr)\n\n try:\n # Use a single query to get all data at once\n query = text(\n f\"\"\"\n SELECT\n :tenant_id AS tenant_id,\n (\n SELECT time_sent\n FROM \"{schema}\".chat_message\n WHERE message_type = 'USER'\n ORDER BY time_sent DESC\n LIMIT 1\n ) AS last_query_time,\n (\n SELECT message\n FROM \"{schema}\".chat_message\n WHERE message_type = 'USER'\n ORDER BY time_sent DESC\n LIMIT 1\n ) AS last_query_text,\n (SELECT COUNT(*) FROM \"{schema}\".document) AS num_documents,\n (SELECT COUNT(*) FROM \"{schema}\".user) AS num_users\n \"\"\"\n )\n\n result = session.execute(query, {\"tenant_id\": schema}).mappings().first()\n\n if result:\n summaries.append(dict(result))\n\n except ProgrammingError as e:\n # schema may be missing a table\n print(f\"Error processing schema {schema}: {e}\", file=sys.stderr)\n session.rollback()\n continue\n except Exception as e:\n print(f\"Unexpected error processing schema {schema}: {e}\", file=sys.stderr)\n session.rollback()\n continue\n\n return summaries\n\n\ndef main() -> None:\n\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n\n with get_session_with_shared_schema() as session:\n summaries = get_tenant_activity_summary(session)\n\n print(json.dumps(summaries, indent=2, default=str))\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/test-openapi-key.py", "content": "VALID_MODEL_LIST = [\n \"gpt-4o-mini\",\n \"gpt-4o\",\n \"gpt-4-1106-preview\",\n \"gpt-4-vision-preview\",\n \"gpt-4\",\n \"gpt-4-0314\",\n \"gpt-4-0613\",\n \"gpt-4-32k\",\n \"gpt-4-32k-0314\",\n \"gpt-4-32k-0613\",\n \"gpt-3.5-turbo-0125\",\n \"gpt-3.5-turbo-1106\",\n \"gpt-3.5-turbo\",\n \"gpt-3.5-turbo-16k\",\n \"gpt-3.5-turbo-0301\",\n \"gpt-3.5-turbo-0613\",\n \"gpt-3.5-turbo-16k-0613\",\n]\n\n\nif __name__ == \"__main__\":\n from openai import OpenAI\n\n model_version = None\n while model_version not in VALID_MODEL_LIST:\n model_version = input(\"Please provide an OpenAI model version to test: \")\n if model_version not in VALID_MODEL_LIST:\n print(f\"Model must be from valid list: {', '.join(VALID_MODEL_LIST)}\")\n assert model_version\n\n api_key = input(\"Please provide an OpenAI API Key to test: \")\n client = OpenAI(\n api_key=api_key,\n )\n\n prompt = \"The boy went to the \"\n print(f\"Asking OpenAI to finish the sentence using {model_version}\")\n print(prompt)\n try:\n messages = [\n {\"role\": \"system\", \"content\": \"Finish the sentence\"},\n {\"role\": \"user\", \"content\": prompt},\n ]\n response = client.chat.completions.create(\n model=model_version,\n messages=messages, # type:ignore\n max_tokens=5,\n temperature=2,\n )\n print(response.choices[0].message.content)\n print(\"Success! Feel free to use this API key for Onyx.\")\n except Exception:\n print(\n \"Failed, provided API key is invalid for Onyx, please address the error from OpenAI.\"\n )\n raise\n" }, { "path": "backend/scripts/transform_openapi_for_docs.py", "content": "\"\"\"\nTransform OpenAPI schema for public documentation.\n\nFilters endpoints tagged with \"public\", converts auth to Bearer token,\nand removes internal parameters (tenant_id, db_session).\n\nUsage:\n python scripts/transform_openapi_for_docs.py -i generated/openapi.json -o openapi_docs.json\n\"\"\"\n\nimport argparse\nimport copy\nimport json\nfrom typing import Any\n\nPUBLIC_TAG = \"public\"\nDOCS_SERVER_URL = \"https://cloud.onyx.app/api\"\nINTERNAL_PARAMETERS = {\"tenant_id\", \"db_session\"}\n\n\ndef collect_schema_refs(obj: Any, refs: set[str]) -> None:\n \"\"\"Recursively collect all $ref references from an object.\"\"\"\n if isinstance(obj, dict):\n if \"$ref\" in obj:\n ref = obj[\"$ref\"]\n if ref.startswith(\"#/components/schemas/\"):\n refs.add(ref.split(\"/\")[-1])\n for value in obj.values():\n collect_schema_refs(value, refs)\n elif isinstance(obj, list):\n for item in obj:\n collect_schema_refs(item, refs)\n\n\ndef get_all_referenced_schemas(\n schemas: dict[str, Any], initial_refs: set[str]\n) -> set[str]:\n \"\"\"Get all schemas referenced by initial_refs, including nested references.\"\"\"\n all_refs = set(initial_refs)\n to_process = list(initial_refs)\n\n while to_process:\n schema_name = to_process.pop()\n if schema_name not in schemas:\n continue\n\n new_refs: set[str] = set()\n collect_schema_refs(schemas[schema_name], new_refs)\n\n for ref in new_refs:\n if ref not in all_refs:\n all_refs.add(ref)\n to_process.append(ref)\n\n return all_refs\n\n\ndef remove_internal_properties_from_schema(schema: dict[str, Any]) -> None:\n \"\"\"Recursively remove internal properties from a schema.\"\"\"\n if not isinstance(schema, dict):\n return\n\n if \"properties\" in schema and isinstance(schema[\"properties\"], dict):\n for prop_name in list(schema[\"properties\"].keys()):\n if prop_name in INTERNAL_PARAMETERS:\n del schema[\"properties\"][prop_name]\n\n if \"required\" in schema and isinstance(schema[\"required\"], list):\n schema[\"required\"] = [\n r for r in schema[\"required\"] if r not in INTERNAL_PARAMETERS\n ]\n if not schema[\"required\"]:\n del schema[\"required\"]\n\n for key in [\"allOf\", \"oneOf\", \"anyOf\"]:\n if key in schema and isinstance(schema[key], list):\n for item in schema[key]:\n remove_internal_properties_from_schema(item)\n\n if \"items\" in schema:\n remove_internal_properties_from_schema(schema[\"items\"])\n\n if \"additionalProperties\" in schema and isinstance(\n schema[\"additionalProperties\"], dict\n ):\n remove_internal_properties_from_schema(schema[\"additionalProperties\"])\n\n\ndef remove_internal_parameters(spec: dict[str, Any]) -> None:\n \"\"\"Remove internal parameters from all endpoints and schemas.\"\"\"\n for path_data in spec.get(\"paths\", {}).values():\n for method_data in path_data.values():\n if isinstance(method_data, dict) and \"parameters\" in method_data:\n method_data[\"parameters\"] = [\n p\n for p in method_data[\"parameters\"]\n if not (\n isinstance(p, dict) and p.get(\"name\") in INTERNAL_PARAMETERS\n )\n ]\n if not method_data[\"parameters\"]:\n del method_data[\"parameters\"]\n\n for schema in spec.get(\"components\", {}).get(\"schemas\", {}).values():\n remove_internal_properties_from_schema(schema)\n\n\ndef transform_openapi(input_spec: dict[str, Any]) -> dict[str, Any]:\n \"\"\"Transform the OpenAPI spec for public documentation.\"\"\"\n output_spec: dict[str, Any] = {\n \"openapi\": input_spec.get(\"openapi\", \"3.1.0\"),\n \"info\": {\n \"title\": \"Onyx API\",\n \"description\": \"Onyx API for AI-powered enterprise search and chat\",\n \"version\": input_spec.get(\"info\", {}).get(\"version\", \"1.0.0\"),\n },\n \"servers\": [{\"url\": DOCS_SERVER_URL}],\n \"paths\": {},\n \"components\": {\n \"schemas\": {},\n \"securitySchemes\": {\n \"BearerAuth\": {\n \"type\": \"http\",\n \"scheme\": \"bearer\",\n \"description\": \"Authorization header with Bearer token\",\n }\n },\n },\n }\n\n input_paths = input_spec.get(\"paths\", {})\n initial_refs: set[str] = set()\n\n for path, path_data in input_paths.items():\n for method, method_data in path_data.items():\n if not isinstance(method_data, dict):\n continue\n\n if PUBLIC_TAG in method_data.get(\"tags\", []):\n if path not in output_spec[\"paths\"]:\n output_spec[\"paths\"][path] = {}\n\n endpoint = copy.deepcopy(method_data)\n if \"security\" in endpoint:\n endpoint[\"security\"] = [{\"BearerAuth\": []}]\n output_spec[\"paths\"][path][method] = endpoint\n collect_schema_refs(method_data, initial_refs)\n\n input_schemas = input_spec.get(\"components\", {}).get(\"schemas\", {})\n all_refs = get_all_referenced_schemas(input_schemas, initial_refs)\n\n for schema_name in all_refs:\n if schema_name in input_schemas:\n output_spec[\"components\"][\"schemas\"][schema_name] = copy.deepcopy(\n input_schemas[schema_name]\n )\n\n remove_internal_parameters(output_spec)\n\n return output_spec\n\n\ndef main() -> None:\n parser = argparse.ArgumentParser(\n description=\"Transform OpenAPI schema for public documentation\"\n )\n parser.add_argument(\n \"--input\", \"-i\", default=\"openapi.json\", help=\"Input OpenAPI JSON file\"\n )\n parser.add_argument(\n \"--output\", \"-o\", default=\"openapi_docs.json\", help=\"Output OpenAPI JSON file\"\n )\n args = parser.parse_args()\n\n with open(args.input) as f:\n input_spec = json.load(f)\n\n output_spec = transform_openapi(input_spec)\n\n with open(args.output, \"w\") as f:\n json.dump(output_spec, f, indent=2)\n\n endpoint_count = sum(len(m) for m in output_spec[\"paths\"].values())\n schema_count = len(output_spec[\"components\"][\"schemas\"])\n print(f\"Wrote {args.output}: {endpoint_count} endpoints, {schema_count} schemas\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/scripts/upload_files_as_connectors.py", "content": "\"\"\"\nScript to upload files from a directory as individual file connectors in Onyx.\nEach file gets its own connector named after the file.\n\nUsage:\n python upload_files_as_connectors.py --data-dir /path/to/files --api-key YOUR_KEY\n python upload_files_as_connectors.py --data-dir /path/to/files --api-key YOUR_KEY --api-base http://onyxserver:3000\n python upload_files_as_connectors.py --data-dir /path/to/files --api-key YOUR_KEY --file-glob '*.zip'\n\nRequires:\n pip install requests\n\"\"\"\n\nimport argparse\nimport fnmatch\nimport os\nimport sys\nimport threading\nimport time\n\nimport requests\n\nREQUEST_TIMEOUT = 900 # 15 minutes\n\n\ndef _elapsed_printer(label: str, stop_event: threading.Event) -> None:\n \"\"\"Print a live elapsed-time counter until stop_event is set.\"\"\"\n start = time.monotonic()\n while not stop_event.wait(timeout=1):\n elapsed = int(time.monotonic() - start)\n m, s = divmod(elapsed, 60)\n print(f\"\\r {label} ... {m:02d}:{s:02d}\", end=\"\", flush=True)\n elapsed = int(time.monotonic() - start)\n m, s = divmod(elapsed, 60)\n print(f\"\\r {label} ... {m:02d}:{s:02d} done\")\n\n\ndef _timed_request(label: str, fn: object) -> requests.Response:\n \"\"\"Run a request function while displaying a live elapsed timer.\"\"\"\n stop = threading.Event()\n t = threading.Thread(target=_elapsed_printer, args=(label, stop), daemon=True)\n t.start()\n try:\n resp = fn() # type: ignore[operator]\n finally:\n stop.set()\n t.join()\n return resp\n\n\ndef upload_file(\n session: requests.Session, base_url: str, file_path: str\n) -> dict | None:\n \"\"\"Upload a single file and return the response with file_paths and file_names.\"\"\"\n with open(file_path, \"rb\") as f:\n resp = _timed_request(\n \"Uploading\",\n lambda: session.post(\n f\"{base_url}/api/manage/admin/connector/file/upload\",\n files={\"files\": (os.path.basename(file_path), f)},\n timeout=REQUEST_TIMEOUT,\n ),\n )\n if not resp.ok:\n print(f\" ERROR uploading: {resp.text}\")\n return None\n return resp.json()\n\n\ndef create_connector(\n session: requests.Session,\n base_url: str,\n name: str,\n file_paths: list[str],\n file_names: list[str],\n zip_metadata_file_id: str | None,\n) -> int | None:\n \"\"\"Create a file connector and return its ID.\"\"\"\n resp = _timed_request(\n \"Creating connector\",\n lambda: session.post(\n f\"{base_url}/api/manage/admin/connector\",\n json={\n \"name\": name,\n \"source\": \"file\",\n \"input_type\": \"load_state\",\n \"connector_specific_config\": {\n \"file_locations\": file_paths,\n \"file_names\": file_names,\n \"zip_metadata_file_id\": zip_metadata_file_id,\n },\n \"refresh_freq\": None,\n \"prune_freq\": None,\n \"indexing_start\": None,\n \"access_type\": \"public\",\n \"groups\": [],\n },\n timeout=REQUEST_TIMEOUT,\n ),\n )\n if not resp.ok:\n print(f\" ERROR creating connector: {resp.text}\")\n return None\n return resp.json()[\"id\"]\n\n\ndef create_credential(\n session: requests.Session, base_url: str, name: str\n) -> int | None:\n \"\"\"Create a dummy credential for the file connector.\"\"\"\n resp = session.post(\n f\"{base_url}/api/manage/credential\",\n json={\n \"credential_json\": {},\n \"admin_public\": True,\n \"source\": \"file\",\n \"curator_public\": True,\n \"groups\": [],\n \"name\": name,\n },\n timeout=REQUEST_TIMEOUT,\n )\n if not resp.ok:\n print(f\" ERROR creating credential: {resp.text}\")\n return None\n return resp.json()[\"id\"]\n\n\ndef link_credential(\n session: requests.Session,\n base_url: str,\n connector_id: int,\n credential_id: int,\n name: str,\n) -> bool:\n \"\"\"Link the connector to the credential (create CC pair).\"\"\"\n resp = session.put(\n f\"{base_url}/api/manage/connector/{connector_id}/credential/{credential_id}\",\n json={\n \"name\": name,\n \"access_type\": \"public\",\n \"groups\": [],\n \"auto_sync_options\": None,\n \"processing_mode\": \"REGULAR\",\n },\n timeout=REQUEST_TIMEOUT,\n )\n if not resp.ok:\n print(f\" ERROR linking credential: {resp.text}\")\n return False\n return True\n\n\ndef run_connector(\n session: requests.Session,\n base_url: str,\n connector_id: int,\n credential_id: int,\n) -> bool:\n \"\"\"Trigger the connector to start indexing.\"\"\"\n resp = session.post(\n f\"{base_url}/api/manage/admin/connector/run-once\",\n json={\n \"connector_id\": connector_id,\n \"credentialIds\": [credential_id],\n \"from_beginning\": False,\n },\n timeout=REQUEST_TIMEOUT,\n )\n if not resp.ok:\n print(f\" ERROR running connector: {resp.text}\")\n return False\n return True\n\n\ndef process_file(session: requests.Session, base_url: str, file_path: str) -> bool:\n \"\"\"Process a single file through the full connector creation flow.\"\"\"\n file_name = os.path.basename(file_path)\n connector_name = file_name\n print(f\"Processing: {file_name}\")\n\n # Step 1: Upload\n upload_resp = upload_file(session, base_url, file_path)\n if not upload_resp:\n return False\n\n # Step 2: Create connector\n connector_id = create_connector(\n session,\n base_url,\n name=f\"FileConnector-{connector_name}\",\n file_paths=upload_resp[\"file_paths\"],\n file_names=upload_resp[\"file_names\"],\n zip_metadata_file_id=upload_resp.get(\"zip_metadata_file_id\"),\n )\n if connector_id is None:\n return False\n\n # Step 3: Create credential\n credential_id = create_credential(session, base_url, name=connector_name)\n if credential_id is None:\n return False\n\n # Step 4: Link connector to credential\n if not link_credential(\n session, base_url, connector_id, credential_id, connector_name\n ):\n return False\n\n # Step 5: Trigger indexing\n if not run_connector(session, base_url, connector_id, credential_id):\n return False\n\n print(f\" OK (connector_id={connector_id})\")\n return True\n\n\ndef get_authenticated_session(api_key: str) -> requests.Session:\n \"\"\"Create a session authenticated with an API key.\"\"\"\n session = requests.Session()\n session.headers.update({\"Authorization\": f\"Bearer {api_key}\"})\n return session\n\n\ndef main() -> None:\n parser = argparse.ArgumentParser(\n description=\"Upload files as individual Onyx file connectors.\"\n )\n parser.add_argument(\n \"--data-dir\",\n required=True,\n help=\"Directory containing files to upload.\",\n )\n parser.add_argument(\n \"--api-base\",\n default=\"http://localhost:3000\",\n help=\"Base URL for the Onyx API (default: http://localhost:3000).\",\n )\n parser.add_argument(\n \"--api-key\",\n required=True,\n help=\"API key for authentication.\",\n )\n parser.add_argument(\n \"--file-glob\",\n default=None,\n help=\"Glob pattern to filter files (e.g. '*.json', '*.zip').\",\n )\n args = parser.parse_args()\n\n data_dir = args.data_dir\n base_url = args.api_base.rstrip(\"/\")\n api_key = args.api_key\n file_glob = args.file_glob\n\n if not os.path.isdir(data_dir):\n print(f\"Error: {data_dir} is not a directory\")\n sys.exit(1)\n\n script_path = os.path.realpath(__file__)\n files = sorted(\n os.path.join(data_dir, f)\n for f in os.listdir(data_dir)\n if os.path.isfile(os.path.join(data_dir, f))\n and os.path.realpath(os.path.join(data_dir, f)) != script_path\n and (file_glob is None or fnmatch.fnmatch(f, file_glob))\n )\n\n if not files:\n print(f\"No files found in {data_dir}\")\n sys.exit(1)\n\n print(f\"Found {len(files)} file(s) in {data_dir}\\n\")\n\n session = get_authenticated_session(api_key)\n\n success = 0\n failed = 0\n for file_path in files:\n if process_file(session, base_url, file_path):\n success += 1\n else:\n failed += 1\n # Small delay to avoid overwhelming the server\n time.sleep(0.5)\n\n print(f\"\\nDone: {success} succeeded, {failed} failed out of {len(files)} files.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "backend/shared_configs/__init__.py", "content": "" }, { "path": "backend/shared_configs/configs.py", "content": "import os\nfrom typing import Any\nfrom typing import List\nfrom urllib.parse import urlparse\n\n# Used for logging\nSLACK_CHANNEL_ID = \"channel_id\"\n\n# Skip model warmup at startup\n# Default to True (skip warmup) if not set, otherwise respect the value\nSKIP_WARM_UP = os.environ.get(\"SKIP_WARM_UP\", \"true\").lower() == \"true\"\n\n# Check if model server is disabled\nDISABLE_MODEL_SERVER = os.environ.get(\"DISABLE_MODEL_SERVER\", \"\").lower() == \"true\"\n\n# If model server is disabled, use \"disabled\" as host to trigger proper handling\nif DISABLE_MODEL_SERVER:\n MODEL_SERVER_HOST = \"disabled\"\n MODEL_SERVER_ALLOWED_HOST = \"disabled\"\n INDEXING_MODEL_SERVER_HOST = \"disabled\"\nelse:\n MODEL_SERVER_HOST = os.environ.get(\"MODEL_SERVER_HOST\") or \"localhost\"\n MODEL_SERVER_ALLOWED_HOST = os.environ.get(\"MODEL_SERVER_HOST\") or \"0.0.0.0\"\n INDEXING_MODEL_SERVER_HOST = (\n os.environ.get(\"INDEXING_MODEL_SERVER_HOST\") or MODEL_SERVER_HOST\n )\n\nMODEL_SERVER_PORT = int(os.environ.get(\"MODEL_SERVER_PORT\") or \"9000\")\n# Model server for indexing should use a separate one to not allow indexing to introduce delay\n# for inference\nINDEXING_MODEL_SERVER_PORT = int(\n os.environ.get(\"INDEXING_MODEL_SERVER_PORT\") or MODEL_SERVER_PORT\n)\n\n# Onyx custom Deep Learning Models\nCONNECTOR_CLASSIFIER_MODEL_REPO = \"Danswer/filter-extraction-model\"\nCONNECTOR_CLASSIFIER_MODEL_TAG = \"1.0.0\"\nINTENT_MODEL_VERSION = \"onyx-dot-app/hybrid-intent-token-classifier\"\n# INTENT_MODEL_TAG = \"v1.0.3\"\nINTENT_MODEL_TAG: str | None = None\n# Bi-Encoder, other details\nDOC_EMBEDDING_CONTEXT_SIZE = 512\n\n# Used to distinguish alternative indices\nALT_INDEX_SUFFIX = \"__danswer_alt_index\"\n\n# Used for loading defaults for automatic deployments and dev flows\n# For local, use: mixedbread-ai/mxbai-rerank-xsmall-v1\nDEFAULT_CROSS_ENCODER_MODEL_NAME = (\n os.environ.get(\"DEFAULT_CROSS_ENCODER_MODEL_NAME\") or None\n)\nDEFAULT_CROSS_ENCODER_API_KEY = os.environ.get(\"DEFAULT_CROSS_ENCODER_API_KEY\") or None\nDEFAULT_CROSS_ENCODER_PROVIDER_TYPE = (\n os.environ.get(\"DEFAULT_CROSS_ENCODER_PROVIDER_TYPE\") or None\n)\nDISABLE_RERANK_FOR_STREAMING = (\n os.environ.get(\"DISABLE_RERANK_FOR_STREAMING\", \"\").lower() == \"true\"\n)\n\n# This controls the minimum number of pytorch \"threads\" to allocate to the embedding\n# model. If torch finds more threads on its own, this value is not used.\nMIN_THREADS_ML_MODELS = int(os.environ.get(\"MIN_THREADS_ML_MODELS\") or 1)\n\n# Model server that has indexing only set will throw exception if used for reranking\n# or intent classification\nINDEXING_ONLY = os.environ.get(\"INDEXING_ONLY\", \"\").lower() == \"true\"\n\n# The process needs to have this for the log file to write to\n# otherwise, it will not create additional log files\n# This should just be the filename base without extension or path.\nLOG_FILE_NAME = os.environ.get(\"LOG_FILE_NAME\") or \"onyx\"\n\n# Enable generating persistent log files for local dev environments\nDEV_LOGGING_ENABLED = os.environ.get(\"DEV_LOGGING_ENABLED\", \"\").lower() == \"true\"\n# notset, debug, info, notice, warning, error, or critical\nLOG_LEVEL = os.environ.get(\"LOG_LEVEL\") or \"info\"\n\n# Timeout for API-based embedding models\n# NOTE: does not apply for Google VertexAI, since the python client doesn't\n# allow us to specify a custom timeout\nAPI_BASED_EMBEDDING_TIMEOUT = int(os.environ.get(\"API_BASED_EMBEDDING_TIMEOUT\", \"600\"))\n\n# Local batch size for VertexAI embedding models currently calibrated for item size of 512 tokens\n# NOTE: increasing this value may lead to API errors due to token limit exhaustion per call.\nVERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE = int(\n os.environ.get(\"VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE\", \"50\")\n)\n\n# Only used for OpenAI\nOPENAI_EMBEDDING_TIMEOUT = int(\n os.environ.get(\"OPENAI_EMBEDDING_TIMEOUT\", API_BASED_EMBEDDING_TIMEOUT)\n)\n\n# Whether or not to strictly enforce token limit for chunking.\nSTRICT_CHUNK_TOKEN_LIMIT = (\n os.environ.get(\"STRICT_CHUNK_TOKEN_LIMIT\", \"\").lower() == \"true\"\n)\n\n# Set up Sentry integration (for error logging)\nSENTRY_DSN = os.environ.get(\"SENTRY_DSN\")\n\n\n# Fields which should only be set on new search setting\nPRESERVED_SEARCH_FIELDS = [\n \"id\",\n \"provider_type\",\n \"api_key\",\n \"model_name\",\n \"api_url\",\n \"index_name\",\n \"multipass_indexing\",\n \"enable_contextual_rag\",\n \"model_dim\",\n \"normalize\",\n \"passage_prefix\",\n \"query_prefix\",\n]\n\n\ndef validate_cors_origin(origin: str) -> None:\n parsed = urlparse(origin)\n if parsed.scheme not in [\"http\", \"https\"] or not parsed.netloc:\n raise ValueError(f\"Invalid CORS origin: '{origin}'\")\n\n\n# Examples of valid values for the environment variable:\n# - \"\" (allow all origins)\n# - \"http://example.com\" (single origin)\n# - \"http://example.com,https://example.org\" (multiple origins)\n# - \"*\" (allow all origins)\nCORS_ALLOWED_ORIGIN_ENV = os.environ.get(\"CORS_ALLOWED_ORIGIN\", \"\")\n\n# Explicitly declare the type of CORS_ALLOWED_ORIGIN\nCORS_ALLOWED_ORIGIN: List[str]\n\nif CORS_ALLOWED_ORIGIN_ENV:\n # Split the environment variable into a list of origins\n CORS_ALLOWED_ORIGIN = [\n origin.strip()\n for origin in CORS_ALLOWED_ORIGIN_ENV.split(\",\")\n if origin.strip()\n ]\n # Validate each origin in the list\n for origin in CORS_ALLOWED_ORIGIN:\n validate_cors_origin(origin)\nelse:\n # If the environment variable is empty, allow all origins\n CORS_ALLOWED_ORIGIN = [\"*\"]\n\n\n# Multi-tenancy configuration\nMULTI_TENANT = os.environ.get(\"MULTI_TENANT\", \"\").lower() == \"true\"\n\n# Outside this file, should almost always use `POSTGRES_DEFAULT_SCHEMA` unless you\n# have a very good reason\nPOSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE = \"public\"\nPOSTGRES_DEFAULT_SCHEMA = (\n os.environ.get(\"POSTGRES_DEFAULT_SCHEMA\") or POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE\n)\nDEFAULT_REDIS_PREFIX = os.environ.get(\"DEFAULT_REDIS_PREFIX\") or \"default\"\n\n\nasync def async_return_default_schema(\n *args: Any, **kwargs: Any # noqa: ARG001\n) -> str: # noqa: ARG001\n return POSTGRES_DEFAULT_SCHEMA\n\n\n# Prefix used for all tenant ids\nTENANT_ID_PREFIX = \"tenant_\"\n\nDISALLOWED_SLACK_BOT_TENANT_IDS = os.environ.get(\"DISALLOWED_SLACK_BOT_TENANT_IDS\")\nDISALLOWED_SLACK_BOT_TENANT_LIST = (\n [\n tenant.strip()\n for tenant in DISALLOWED_SLACK_BOT_TENANT_IDS.split(\",\")\n if tenant.strip()\n ]\n if DISALLOWED_SLACK_BOT_TENANT_IDS\n else None\n)\n\nIGNORED_SYNCING_TENANT_IDS = os.environ.get(\"IGNORED_SYNCING_TENANT_IDS\")\nIGNORED_SYNCING_TENANT_LIST = (\n [\n tenant.strip()\n for tenant in IGNORED_SYNCING_TENANT_IDS.split(\",\")\n if tenant.strip()\n ]\n if IGNORED_SYNCING_TENANT_IDS\n else None\n)\n\nENVIRONMENT = os.environ.get(\"ENVIRONMENT\") or \"not_explicitly_set\"\n\n\n#####\n# Usage Limits Configuration (meant for cloud, off by default for self-hosted)\n#####\n# Whether usage limits are enforced (defaults to MULTI_TENANT value)\n_USAGE_LIMITS_ENABLED_RAW = os.environ.get(\"USAGE_LIMITS_ENABLED\")\nif _USAGE_LIMITS_ENABLED_RAW is not None:\n USAGE_LIMITS_ENABLED = _USAGE_LIMITS_ENABLED_RAW.lower() == \"true\"\nelse:\n # Default: enabled on cloud (MULTI_TENANT), disabled for self-hosted\n USAGE_LIMITS_ENABLED = MULTI_TENANT\n\n# Usage limit window in seconds (default: 1 week = 604800 seconds)\nUSAGE_LIMIT_WINDOW_SECONDS = int(os.environ.get(\"USAGE_LIMIT_WINDOW_SECONDS\", \"604800\"))\n\n# Per-week LLM usage cost limits in cents (e.g., 1000 = $10.00)\n# Trial users get lower limits than paid users\nUSAGE_LIMIT_LLM_COST_CENTS_TRIAL = int(\n os.environ.get(\"USAGE_LIMIT_LLM_COST_CENTS_TRIAL\", \"3200\") # $32.00 default\n)\nUSAGE_LIMIT_LLM_COST_CENTS_PAID = int(\n os.environ.get(\"USAGE_LIMIT_LLM_COST_CENTS_PAID\", \"6400\") # $64.00 default\n)\n\n# Per-week chunks indexed limits\nUSAGE_LIMIT_CHUNKS_INDEXED_TRIAL = int(\n os.environ.get(\"USAGE_LIMIT_CHUNKS_INDEXED_TRIAL\", 400_000)\n)\nUSAGE_LIMIT_CHUNKS_INDEXED_PAID = int(\n os.environ.get(\"USAGE_LIMIT_CHUNKS_INDEXED_PAID\", 4_000_000)\n)\n\n# Per-week API calls using API keys or Personal Access Tokens\nUSAGE_LIMIT_API_CALLS_TRIAL = int(os.environ.get(\"USAGE_LIMIT_API_CALLS_TRIAL\", \"0\"))\nUSAGE_LIMIT_API_CALLS_PAID = int(os.environ.get(\"USAGE_LIMIT_API_CALLS_PAID\", \"40000\"))\n\n# Per-week non-streaming API calls (more expensive, so lower limits)\nUSAGE_LIMIT_NON_STREAMING_CALLS_TRIAL = int(\n os.environ.get(\"USAGE_LIMIT_NON_STREAMING_CALLS_TRIAL\", \"0\")\n)\nUSAGE_LIMIT_NON_STREAMING_CALLS_PAID = int(\n os.environ.get(\"USAGE_LIMIT_NON_STREAMING_CALLS_PAID\", \"160\")\n)\n" }, { "path": "backend/shared_configs/contextvars.py", "content": "import contextvars\n\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\n\n# Context variable for the current tenant id\nCURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = (\n contextvars.ContextVar(\n \"current_tenant_id\", default=None if MULTI_TENANT else POSTGRES_DEFAULT_SCHEMA\n )\n)\n\n# set by every route in the API server\nINDEXING_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = (\n contextvars.ContextVar(\"indexing_request_id\", default=None)\n)\n\n# set by every route in the API server\nONYX_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = contextvars.ContextVar(\n \"onyx_request_id\", default=None\n)\n\n# Used to store cc pair id and index attempt id in multithreaded environments\nINDEX_ATTEMPT_INFO_CONTEXTVAR: contextvars.ContextVar[tuple[int, int] | None] = (\n contextvars.ContextVar(\"index_attempt_info\", default=None)\n)\n\n# Set by endpoint context middleware — used for per-endpoint DB pool attribution\nCURRENT_ENDPOINT_CONTEXTVAR: contextvars.ContextVar[str | None] = (\n contextvars.ContextVar(\"current_endpoint\", default=None)\n)\n\n\n\"\"\"Utils related to contextvars\"\"\"\n\n\ndef get_current_tenant_id() -> str:\n tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()\n if tenant_id is None:\n import traceback\n\n if not MULTI_TENANT:\n return POSTGRES_DEFAULT_SCHEMA\n\n stack_trace = traceback.format_stack()\n error_message = (\n \"Tenant ID is not set. This should never happen.\\nStack trace:\\n\"\n + \"\".join(stack_trace)\n )\n raise RuntimeError(error_message)\n return tenant_id\n" }, { "path": "backend/shared_configs/enums.py", "content": "from enum import Enum\n\n\nclass EmbeddingProvider(str, Enum):\n OPENAI = \"openai\"\n COHERE = \"cohere\"\n VOYAGE = \"voyage\"\n GOOGLE = \"google\"\n LITELLM = \"litellm\"\n AZURE = \"azure\"\n\n\nclass RerankerProvider(str, Enum):\n COHERE = \"cohere\"\n LITELLM = \"litellm\"\n BEDROCK = \"bedrock\"\n\n\nclass EmbedTextType(str, Enum):\n QUERY = \"query\"\n PASSAGE = \"passage\"\n\n\nclass WebSearchProviderType(str, Enum):\n GOOGLE_PSE = \"google_pse\"\n SERPER = \"serper\"\n EXA = \"exa\"\n SEARXNG = \"searxng\"\n BRAVE = \"brave\"\n\n\nclass WebContentProviderType(str, Enum):\n ONYX_WEB_CRAWLER = \"onyx_web_crawler\"\n FIRECRAWL = \"firecrawl\"\n EXA = \"exa\"\n" }, { "path": "backend/shared_configs/model_server_models.py", "content": "from pydantic import BaseModel\n\nfrom shared_configs.enums import EmbeddingProvider\nfrom shared_configs.enums import EmbedTextType\nfrom shared_configs.enums import RerankerProvider\n\n\nEmbedding = list[float]\n\n\nclass EmbedRequest(BaseModel):\n texts: list[str]\n # Can be none for cloud embedding model requests, error handling logic exists for other cases\n model_name: str | None = None\n deployment_name: str | None = None\n max_context_length: int\n normalize_embeddings: bool\n api_key: str | None = None\n provider_type: EmbeddingProvider | None = None\n text_type: EmbedTextType\n manual_query_prefix: str | None = None\n manual_passage_prefix: str | None = None\n api_url: str | None = None\n api_version: str | None = None\n\n # allows for the truncation of the vector to a lower dimension\n # to reduce memory usage. Currently only supported for OpenAI models.\n # will be ignored for other providers.\n reduced_dimension: int | None = None\n\n # This disables the \"model_\" protected namespace for pydantic\n model_config = {\"protected_namespaces\": ()}\n\n\nclass EmbedResponse(BaseModel):\n embeddings: list[Embedding]\n\n\nclass RerankRequest(BaseModel):\n query: str\n documents: list[str]\n model_name: str\n provider_type: RerankerProvider | None = None\n api_key: str | None = None\n api_url: str | None = None\n\n # This disables the \"model_\" protected namespace for pydantic\n model_config = {\"protected_namespaces\": ()}\n\n\nclass RerankResponse(BaseModel):\n scores: list[float]\n\n\nclass IntentRequest(BaseModel):\n query: str\n # Sequence classification threshold\n semantic_percent_threshold: float\n # Token classification threshold\n keyword_percent_threshold: float\n\n\nclass IntentResponse(BaseModel):\n is_keyword: bool\n keywords: list[str]\n" }, { "path": "backend/shared_configs/utils.py", "content": "from typing import TypeVar\n\n\nT = TypeVar(\"T\")\n\n\ndef batch_list(\n lst: list[T],\n batch_size: int,\n) -> list[list[T]]:\n return [lst[i : i + batch_size] for i in range(0, len(lst), batch_size)]\n" }, { "path": "backend/slackbot_images/README.md", "content": "This folder contains images needed by the Onyx Slack Bot. When possible, we use the images\nwithin `web/public`, but sometimes those images do not work for the Slack Bot.\n" }, { "path": "backend/supervisord.conf", "content": "[supervisord]\nnodaemon=true\nuser=root\nlogfile=/var/log/supervisord.log\nenvironment=PYTHONPATH=\"/app\"\n\n# region enable supervisorctl usage\n[supervisorctl]\nserverurl=unix:///tmp/supervisor.sock\n\n[unix_http_server]\nfile=/tmp/supervisor.sock\nchmod=0700\n\n[rpcinterface:supervisor]\nsupervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface\n# endregion enable supervisorctl usage\n\n# Background jobs that must be run async due to long time to completion\n# NOTE: due to an issue with Celery + SQLAlchemy\n# (https://github.com/celery/celery/issues/7007#issuecomment-1740139367)\n# we must use the threads pool instead of the default prefork pool for now\n# in order to avoid intermittent errors like:\n# `billiard.exceptions.WorkerLostError: Worker exited prematurely: signal 11 (SIGSEGV)`.\n#\n# This means workers will not be able take advantage of multiple CPU cores\n# on a system, but this should be okay for now since all our celery tasks are\n# relatively compute-light (e.g. they tend to just make a bunch of requests to\n# Vespa / Postgres)\n[program:celery_worker_primary]\ncommand=celery -A onyx.background.celery.versioned_apps.primary worker\n --loglevel=INFO\n --hostname=primary@%%n\n -Q celery\nstdout_logfile=/var/log/celery_worker_primary.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n# NOTE: only allowing configuration here and not in the other celery workers,\n# since this is often the bottleneck for \"sync\" jobs (e.g. document set syncing,\n# user group syncing, deletion, etc.)\n[program:celery_worker_light]\ncommand=celery -A onyx.background.celery.versioned_apps.light worker\n --loglevel=INFO\n --hostname=light@%%n\n -Q vespa_metadata_sync,connector_deletion,doc_permissions_upsert,checkpoint_cleanup,index_attempt_cleanup,opensearch_migration\nstdout_logfile=/var/log/celery_worker_light.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n[program:celery_worker_heavy]\ncommand=celery -A onyx.background.celery.versioned_apps.heavy worker\n --loglevel=INFO\n --hostname=heavy@%%n\n -Q connector_pruning,connector_doc_permissions_sync,connector_external_group_sync,csv_generation,sandbox\nstdout_logfile=/var/log/celery_worker_heavy.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n[program:celery_worker_docprocessing]\ncommand=celery -A onyx.background.celery.versioned_apps.docprocessing worker\n --loglevel=INFO\n --hostname=docprocessing@%%n\n -Q docprocessing\nstdout_logfile=/var/log/celery_worker_docprocessing.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n[program:celery_worker_user_file_processing]\ncommand=celery -A onyx.background.celery.versioned_apps.user_file_processing worker\n --loglevel=INFO\n --hostname=user_file_processing@%%n\n -Q user_file_processing,user_file_project_sync,user_file_delete\nstdout_logfile=/var/log/celery_worker_user_file_processing.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n[program:celery_worker_docfetching]\ncommand=celery -A onyx.background.celery.versioned_apps.docfetching worker\n --loglevel=INFO\n --hostname=docfetching@%%n\n -Q connector_doc_fetching\nstdout_logfile=/var/log/celery_worker_docfetching.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n[program:celery_worker_monitoring]\ncommand=celery -A onyx.background.celery.versioned_apps.monitoring worker\n --loglevel=INFO\n --hostname=monitoring@%%n\n -Q monitoring\nstdout_logfile=/var/log/celery_worker_monitoring.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartsecs=10\nstopasgroup=true\n\n\n# Job scheduler for periodic tasks\n[program:celery_beat]\ncommand=celery -A onyx.background.celery.versioned_apps.beat beat\nstdout_logfile=/var/log/celery_beat.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nstartsecs=10\nstopasgroup=true\n\n# watchdog to detect and restart the beat in case of inactivity\n# supervisord only restarts the process if it's dead\n# make sure this key matches ONYX_CELERY_BEAT_HEARTBEAT_KEY\n[program:supervisord_watchdog_celery_beat]\ncommand=python onyx/utils/supervisord_watchdog.py\n --conf /etc/supervisor/conf.d/supervisord.conf\n --key \"onyx:celery:beat:heartbeat\"\n --program celery_beat\nstdout_logfile=/var/log/supervisord_watchdog_celery_beat.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nstartsecs=10\nstopasgroup=true\n\n# Listens for Slack messages and responds with answers\n# for all channels that the OnyxBot has been added to.\n# If not setup, this will just fail 5 times and then stop.\n# More details on setup here: https://docs.onyx.app/admins/getting_started/slack_bot_setup\n[program:slack_bot]\ncommand=python onyx/onyxbot/slack/listener.py\nstdout_logfile=/var/log/slack_bot.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartretries=5\nstartsecs=60\n\n# Listens for Discord messages and responds with answers\n# for all guilds/channels that the OnyxBot has been added to.\n# If not configured, will continue to probe every 3 minutes for a Discord bot token.\n[program:discord_bot]\ncommand=python onyx/onyxbot/discord/client.py\nstdout_logfile=/var/log/discord_bot.log\nstdout_logfile_maxbytes=16MB\nredirect_stderr=true\nautorestart=true\nstartretries=5\nstartsecs=60\n\n# Pushes all logs from the above programs to stdout\n# No log rotation here, since it's stdout it's handled by the Docker container logging\n[program:log-redirect-handler]\ncommand=tail -qF\n /var/log/celery_beat.log\n /var/log/celery_worker_primary.log\n /var/log/celery_worker_light.log\n /var/log/celery_worker_heavy.log\n /var/log/celery_worker_docprocessing.log\n /var/log/celery_worker_monitoring.log\n /var/log/celery_worker_user_file_processing.log\n /var/log/celery_worker_docfetching.log\n /var/log/slack_bot.log\n /var/log/discord_bot.log\n /var/log/supervisord_watchdog_celery_beat.log\n /var/log/mcp_server.log\n /var/log/mcp_server.err.log\nstdout_logfile=/dev/stdout\nstdout_logfile_maxbytes = 0 # must be set to 0 when stdout_logfile=/dev/stdout\nautorestart=true\n" }, { "path": "backend/tests/README.md", "content": "# Backend Tests\n\n## Test Types\n\nThere are four test categories, ordered by increasing scope:\n\n### Unit Tests (`tests/unit/`)\n\nNo external services. Mock all I/O with `unittest.mock`. Use for complex, isolated\nlogic (e.g. citation processing, encryption).\n\n```bash\npytest -xv backend/tests/unit\n```\n\n### External Dependency Unit Tests (`tests/external_dependency_unit/`)\n\nExternal services (Postgres, Redis, Vespa, OpenAI, etc.) are running, but Onyx\napplication containers are not. Tests call functions directly and can mock selectively.\n\nUse when you need a real database or real API calls but want control over setup.\n\n```bash\npython -m dotenv -f .vscode/.env run -- pytest backend/tests/external_dependency_unit\n```\n\n### Integration Tests (`tests/integration/`)\n\nFull Onyx deployment running. No mocking. Prefer this over other test types when possible.\n\n```bash\npython -m dotenv -f .vscode/.env run -- pytest backend/tests/integration\n```\n\n### Playwright / E2E Tests (`web/tests/e2e/`)\n\nFull stack including web server. Use for frontend-backend coordination.\n\n```bash\nnpx playwright test \n```\n\n## Shared Fixtures\n\nShared fixtures live in `backend/tests/conftest.py`. Test subdirectories can define\ntheir own `conftest.py` for directory-scoped fixtures.\n\n## Running Tests Repeatedly (`pytest-repeat`)\n\nUse `pytest-repeat` to catch flaky tests by running them multiple times:\n\n```bash\n# Run a specific test 50 times\npytest --count=50 backend/tests/unit/path/to/test.py::test_name\n\n# Stop on first failure with -x\npytest --count=50 -x backend/tests/unit/path/to/test.py::test_name\n\n# Repeat an entire test file\npytest --count=10 backend/tests/unit/path/to/test_file.py\n```\n\n## Best Practices\n\n### Use `enable_ee` fixture instead of inlining\n\nEnables EE mode for a test, with proper teardown and cache clearing.\n\n```python\n# Whole file (in a test module, NOT in conftest.py)\npytestmark = pytest.mark.usefixtures(\"enable_ee\")\n\n# Whole directory — add an autouse wrapper to the directory's conftest.py\n@pytest.fixture(autouse=True)\ndef _enable_ee_for_directory(enable_ee: None) -> None: \n \"\"\"Wraps the shared enable_ee fixture with autouse for this directory.\"\"\"\n\n# Single test\ndef test_something(enable_ee: None) -> None: ...\n```\n\n**Note:** `pytestmark` in a `conftest.py` does NOT apply markers to tests in that\ndirectory — it only affects tests defined in the conftest itself (which is none).\nUse the autouse fixture wrapper pattern shown above instead.\n\nDo NOT inline `global_version.set_ee()` — always use the fixture.\n" }, { "path": "backend/tests/__init__.py", "content": "" }, { "path": "backend/tests/api/test_api.py", "content": "import os\nfrom collections.abc import Generator\nfrom typing import Any\n\nimport pytest\nfrom fastapi import FastAPI\nfrom fastapi.testclient import TestClient\n\nfrom onyx.configs.constants import DEV_VERSION_PATTERN\nfrom onyx.configs.constants import STABLE_VERSION_PATTERN\nfrom onyx.main import fetch_versioned_implementation\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n@pytest.fixture(scope=\"function\")\ndef client() -> Generator[TestClient, Any, None]:\n # Set environment variables\n os.environ[\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\"] = \"True\"\n\n # Initialize TestClient with the FastAPI app\n app: FastAPI = fetch_versioned_implementation(\n module=\"onyx.main\", attribute=\"get_application\"\n )()\n client = TestClient(app)\n yield client\n\n\n@pytest.mark.skip(\n reason=\"enable when we have a testing environment with preloaded data\"\n)\ndef test_handle_simplified_chat_message(client: TestClient) -> None:\n req: dict[str, Any] = {}\n\n req[\"persona_id\"] = 0\n req[\"description\"] = \"pytest\"\n response = client.post(\"/chat/create-chat-session\", json=req)\n chat_session_id = response.json()[\"chat_session_id\"]\n\n req = {}\n req[\"chat_session_id\"] = chat_session_id\n req[\"message\"] = \"hello\"\n\n response = client.post(\"/chat/send-message-simple-api\", json=req)\n assert response.status_code == 200\n\n\n@pytest.mark.skip(\n reason=\"enable when we have a testing environment with preloaded data\"\n)\ndef test_handle_send_message_simple_with_history(client: TestClient) -> None:\n req: dict[str, Any] = {}\n messages = []\n messages.append({\"message\": \"What sorts of questions can you answer for me?\"})\n # messages.append({\"message\":\n # \"I'd be happy to assist you with a wide range of questions related to Ramp's expense management platform. \"\n # \"I can help with topics such as:\\n\\n\"\n # \"1. Setting up and managing your Ramp account\\n\"\n # \"2. Using Ramp cards and making purchases\\n\"\n # \"3. Submitting and reviewing expenses\\n\"\n # \"4. Understanding Ramp's features and benefits\\n\"\n # \"5. Navigating the Ramp dashboard and mobile app\\n\"\n # \"6. Managing team spending and budgets\\n\"\n # \"7. Integrating Ramp with accounting software\\n\"\n # \"8. Troubleshooting common issues\\n\\n\"\n # \"Feel free to ask any specific questions you have about using Ramp, \"\n # \"and I'll do my best to provide clear and helpful answers. \"\n # \"Is there a particular area you'd like to know more about?\",\n # \"role\": \"assistant\"})\n # req[\"prompt_id\"] = 9\n # req[\"persona_id\"] = 6\n\n # Yoda\n req[\"persona_id\"] = 1\n messages.append(\n {\n \"message\": \"Answer questions for you, I can. \"\n \"About many topics, knowledge I have. \"\n \"But specific to documents provided, limited my responses are. \"\n \"Ask you may about:\\n\\n\"\n \"- User interviews and building trust with participants\\n\"\n \"- Designing effective surveys and survey questions \\n\"\n \"- Product analysis approaches\\n\"\n \"- Recruiting participants for research\\n\"\n \"- Discussion guides for user interviews\\n\"\n \"- Types of survey questions\\n\\n\"\n \"More there may be, but focus on these areas, the given context does. \"\n \"Specific questions you have, ask you should. Guide you I will, as best I can.\",\n \"role\": \"assistant\",\n }\n )\n # messages.append({\"message\": \"Where can I pilot a survey?\"})\n\n # messages.append({\"message\": \"How many data points should I collect to validate my solution?\"})\n messages.append({\"message\": \"What is solution validation research used for?\"})\n\n req[\"messages\"] = messages\n\n response = client.post(\"/chat/send-message-simple-with-history\", json=req)\n assert response.status_code == 200\n\n resp_json = response.json()\n\n # persona must have LLM relevance enabled for this to pass\n assert len(resp_json[\"llm_selected_doc_indices\"]) > 0\n\n\ndef test_versions_endpoint(client: TestClient) -> None:\n \"\"\"Test that /api/versions endpoint returns valid stable, dev, and migration configurations\"\"\"\n response = client.get(\"/versions\")\n assert response.status_code == 200\n\n data = response.json()\n\n # Verify the top-level structure\n assert \"stable\" in data\n assert \"dev\" in data\n assert \"migration\" in data\n\n # Verify stable configuration\n stable = data[\"stable\"]\n assert \"onyx\" in stable\n assert \"relational_db\" in stable\n assert \"index\" in stable\n assert \"nginx\" in stable\n\n # Verify stable version follows correct pattern (v1.2.3)\n # If this fails, revise latest Github release for typo or incorrect version name\n assert STABLE_VERSION_PATTERN.match(\n stable[\"onyx\"]\n ), f\"Stable version {stable['onyx']} doesn't match pattern v(number).(number).(number)\"\n\n # Verify dev configuration\n dev = data[\"dev\"]\n assert \"onyx\" in dev\n assert \"relational_db\" in dev\n assert \"index\" in dev\n assert \"nginx\" in dev\n\n # Verify dev version follows correct pattern (v1.2.3-beta.4)\n assert DEV_VERSION_PATTERN.match(\n dev[\"onyx\"]\n ), f\"Dev version {dev['onyx']} doesn't match pattern v(number).(number).(number)-beta.(number)\"\n\n # Verify migration configuration\n migration = data[\"migration\"]\n assert \"onyx\" in migration\n assert \"relational_db\" in migration\n assert \"index\" in migration\n assert \"nginx\" in migration\n\n # Verify migration has expected values\n assert migration[\"onyx\"] == \"airgapped-intfloat-nomic-migration\"\n assert migration[\"relational_db\"] == \"postgres:15.2-alpine\"\n assert migration[\"index\"] == \"vespaengine/vespa:8.277.17\"\n assert migration[\"nginx\"] == \"nginx:1.25.5-alpine\"\n\n # Verify versions are different between stable and dev\n assert stable[\"onyx\"] != dev[\"onyx\"], \"Stable and dev versions should be different\"\n\n # Additional validation: ensure all required fields are strings\n for config_name, config in [\n (\"stable\", stable),\n (\"dev\", dev),\n (\"migration\", migration),\n ]:\n for field_name, field_value in config.items():\n assert isinstance(\n field_value, str\n ), f\"{config_name}.{field_name} should be a string, got {type(field_value)}\"\n assert (\n field_value.strip() != \"\"\n ), f\"{config_name}.{field_name} should not be empty\"\n" }, { "path": "backend/tests/conftest.py", "content": "\"\"\"Root conftest — shared fixtures available to all test directories.\"\"\"\n\nfrom collections.abc import Generator\n\nimport pytest\n\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\n\n@pytest.fixture()\ndef enable_ee() -> Generator[None, None, None]:\n \"\"\"Temporarily enable EE mode for a single test.\n\n Restores the previous EE state and clears the versioned-implementation\n cache on teardown so state doesn't leak between tests.\n \"\"\"\n was_ee = global_version.is_ee_version()\n global_version.set_ee()\n fetch_versioned_implementation.cache_clear()\n yield\n if not was_ee:\n global_version.unset_ee()\n fetch_versioned_implementation.cache_clear()\n" }, { "path": "backend/tests/daily/conftest.py", "content": "import os\n\n# Set environment variables BEFORE any other imports to ensure they're picked up\n# by module-level code that reads env vars at import time\n# TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\nos.environ[\"LICENSE_ENFORCEMENT_ENABLED\"] = \"false\"\n\nfrom collections.abc import AsyncGenerator\nfrom collections.abc import Generator\nfrom contextlib import asynccontextmanager\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\nfrom dotenv import load_dotenv\nfrom fastapi import FastAPI\nfrom fastapi.testclient import TestClient\n\nfrom onyx.auth.users import current_admin_user\nfrom onyx.db.engine.sql_engine import get_session\nfrom onyx.db.models import UserRole\nfrom onyx.main import get_application\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\nload_dotenv()\n\n\n@asynccontextmanager\nasync def test_lifespan(\n app: FastAPI, # noqa: ARG001\n) -> AsyncGenerator[None, None]: # noqa: ARG001\n \"\"\"No-op lifespan for tests that don't need database or other services.\"\"\"\n yield\n\n\ndef mock_get_session() -> Generator[MagicMock, None, None]:\n \"\"\"Mock database session for tests that don't actually need DB access.\"\"\"\n yield MagicMock()\n\n\ndef mock_current_admin_user() -> MagicMock:\n \"\"\"Mock admin user for endpoints protected by current_admin_user.\"\"\"\n mock_admin = MagicMock()\n mock_admin.role = UserRole.ADMIN\n return mock_admin\n\n\n@pytest.fixture(scope=\"function\")\ndef client() -> Generator[TestClient, None, None]:\n # Initialize TestClient with the FastAPI app using a no-op test lifespan.\n # Patch out prometheus metrics setup to avoid \"Duplicated timeseries in\n # CollectorRegistry\" errors when multiple tests each create a new app\n # (prometheus registers metrics globally and rejects duplicate names).\n with patch(\"onyx.main.setup_prometheus_metrics\"):\n app: FastAPI = get_application(lifespan_override=test_lifespan)\n\n # Override the database session dependency with a mock\n # (these tests don't actually need DB access)\n app.dependency_overrides[get_session] = mock_get_session\n app.dependency_overrides[current_admin_user] = mock_current_admin_user\n\n # Use TestClient as a context manager to properly trigger lifespan\n with TestClient(app) as client:\n yield client\n\n # Clean up dependency overrides\n app.dependency_overrides.clear()\n" }, { "path": "backend/tests/daily/connectors/airtable/test_airtable_basic.py", "content": "import os\nfrom typing import cast\nfrom unittest.mock import MagicMock\n\nimport pytest\nfrom pydantic import BaseModel\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.airtable.airtable_connector import AirtableConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\n\nBASE_VIEW_ID = \"viwVUEJjWPd8XYjh8\"\n\n\nclass AirtableConfig(BaseModel):\n base_id: str\n table_identifier: str\n access_token: str\n\n\n@pytest.fixture(params=[True, False])\ndef airtable_config(request: pytest.FixtureRequest) -> AirtableConfig:\n table_identifier = (\n os.environ[\"AIRTABLE_TEST_TABLE_NAME\"]\n if request.param\n else os.environ[\"AIRTABLE_TEST_TABLE_ID\"]\n )\n return AirtableConfig(\n base_id=os.environ[\"AIRTABLE_TEST_BASE_ID\"],\n table_identifier=table_identifier,\n access_token=os.environ[\"AIRTABLE_ACCESS_TOKEN\"],\n )\n\n\ndef create_test_document(\n id: str,\n title: str,\n description: str,\n priority: str,\n status: str,\n # Link to another record is skipped for now\n # category: str,\n ticket_id: str,\n created_time: str,\n status_last_changed: str,\n submitted_by: str,\n assignee: str,\n days_since_status_change: int | None,\n attachments: list[tuple[str, str]] | None = None,\n all_fields_as_metadata: bool = False,\n share_id: str | None = None,\n view_id: str | None = None,\n) -> Document:\n base_id = os.environ.get(\"AIRTABLE_TEST_BASE_ID\")\n table_id = os.environ.get(\"AIRTABLE_TEST_TABLE_ID\")\n missing_vars = []\n if not base_id:\n missing_vars.append(\"AIRTABLE_TEST_BASE_ID\")\n if not table_id:\n missing_vars.append(\"AIRTABLE_TEST_TABLE_ID\")\n\n if missing_vars:\n raise RuntimeError(\n f\"Required environment variables not set: {', '.join(missing_vars)}. \"\n \"These variables are required to run Airtable connector tests.\"\n )\n link_base = f\"https://airtable.com/{base_id}\"\n if share_id:\n link_base = f\"{link_base}/{share_id}\"\n link_base = f\"{link_base}/{table_id}\"\n if view_id:\n link_base = f\"{link_base}/{view_id}\"\n\n sections = []\n\n if not all_fields_as_metadata:\n sections.extend(\n [\n TextSection(\n text=f\"Title:\\n------------------------\\n{title}\\n------------------------\",\n link=f\"{link_base}/{id}\",\n ),\n TextSection(\n text=f\"Description:\\n------------------------\\n{description}\\n------------------------\",\n link=f\"{link_base}/{id}\",\n ),\n ]\n )\n\n if attachments:\n for attachment_text, attachment_link in attachments:\n sections.append(\n TextSection(\n text=f\"Attachment:\\n------------------------\\n{attachment_text}\\n------------------------\",\n link=attachment_link,\n ),\n )\n\n metadata: dict[str, str | list[str]] = {\n # \"Category\": category,\n \"Assignee\": assignee,\n \"Submitted by\": submitted_by,\n \"Priority\": priority,\n \"Status\": status,\n \"Created time\": created_time,\n \"ID\": ticket_id,\n \"Status last changed\": status_last_changed,\n **(\n {\"Days since status change\": str(days_since_status_change)}\n if days_since_status_change is not None\n else {}\n ),\n }\n\n if all_fields_as_metadata:\n metadata.update(\n {\n \"Title\": title,\n \"Description\": description,\n }\n )\n\n return Document(\n id=f\"airtable__{id}\",\n sections=cast(list[TextSection | ImageSection], sections),\n source=DocumentSource.AIRTABLE,\n semantic_identifier=f\"{os.environ.get('AIRTABLE_TEST_TABLE_NAME', '')}: {title}\",\n metadata=metadata,\n doc_updated_at=None,\n primary_owners=None,\n secondary_owners=None,\n title=None,\n from_ingestion_api=False,\n additional_info=None,\n )\n\n\ndef compare_documents(\n actual_docs: list[Document], expected_docs: list[Document]\n) -> None:\n \"\"\"Utility function to compare actual and expected documents, ignoring order.\"\"\"\n actual_docs_dict = {doc.id: doc for doc in actual_docs}\n expected_docs_dict = {doc.id: doc for doc in expected_docs}\n\n assert actual_docs_dict.keys() == expected_docs_dict.keys(), \"Document ID mismatch\"\n\n for doc_id in actual_docs_dict:\n actual = actual_docs_dict[doc_id]\n expected = expected_docs_dict[doc_id]\n\n assert (\n actual.source == expected.source\n ), f\"Source mismatch for document {doc_id}\"\n assert (\n actual.semantic_identifier == expected.semantic_identifier\n ), f\"Semantic identifier mismatch for document {doc_id}\"\n assert (\n actual.metadata == expected.metadata\n ), f\"Metadata mismatch for document {doc_id}\"\n assert (\n actual.doc_updated_at == expected.doc_updated_at\n ), f\"Updated at mismatch for document {doc_id}\"\n assert (\n actual.primary_owners == expected.primary_owners\n ), f\"Primary owners mismatch for document {doc_id}\"\n assert (\n actual.secondary_owners == expected.secondary_owners\n ), f\"Secondary owners mismatch for document {doc_id}\"\n assert actual.title == expected.title, f\"Title mismatch for document {doc_id}\"\n assert (\n actual.from_ingestion_api == expected.from_ingestion_api\n ), f\"Ingestion API flag mismatch for document {doc_id}\"\n assert (\n actual.additional_info == expected.additional_info\n ), f\"Additional info mismatch for document {doc_id}\"\n\n # Compare sections\n assert len(actual.sections) == len(\n expected.sections\n ), f\"Number of sections mismatch for document {doc_id}\"\n for i, (actual_section, expected_section) in enumerate(\n zip(actual.sections, expected.sections)\n ):\n assert (\n actual_section.text == expected_section.text\n ), f\"Section {i} text mismatch for document {doc_id}\"\n assert (\n actual_section.link == expected_section.link\n ), f\"Section {i} link mismatch for document {doc_id}\"\n\n\ndef test_airtable_connector_basic(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n airtable_config: AirtableConfig,\n) -> None:\n \"\"\"Test behavior when all non-attachment fields are treated as metadata.\"\"\"\n connector = AirtableConnector(\n base_id=airtable_config.base_id,\n table_name_or_id=airtable_config.table_identifier,\n treat_all_non_attachment_fields_as_metadata=False,\n )\n connector.load_credentials(\n {\n \"airtable_access_token\": airtable_config.access_token,\n }\n )\n doc_batch_generator = connector.load_from_state()\n doc_batch = [\n doc for doc in next(doc_batch_generator) if not isinstance(doc, HierarchyNode)\n ]\n with pytest.raises(StopIteration):\n next(doc_batch_generator)\n\n assert len(doc_batch) == 2\n\n expected_docs = [\n create_test_document(\n id=\"rec8BnxDLyWeegOuO\",\n title=\"Slow Internet\",\n description=\"The internet connection is very slow.\",\n priority=\"Medium\",\n status=\"In Progress\",\n ticket_id=\"2\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n all_fields_as_metadata=False,\n view_id=BASE_VIEW_ID,\n ),\n create_test_document(\n id=\"reccSlIA4pZEFxPBg\",\n title=\"Printer Issue\",\n description=\"The office printer is not working.\",\n priority=\"High\",\n status=\"Open\",\n ticket_id=\"1\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n attachments=[\n (\n \"Test.pdf:\\ntesting!!!\",\n \"https://airtable.com/appCXJqDFS4gea8tn/tblRxFQsTlBBZdRY1/viwVUEJjWPd8XYjh8/reccSlIA4pZEFxPBg/fld1u21zkJACIvAEF/attlj2UBWNEDZngCc?blocks=hide\",\n )\n ],\n all_fields_as_metadata=False,\n view_id=BASE_VIEW_ID,\n ),\n ]\n\n # Compare documents using the utility function\n compare_documents(doc_batch, expected_docs)\n\n\ndef test_airtable_connector_url(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n airtable_config: AirtableConfig,\n) -> None:\n \"\"\"Test that passing an Airtable URL produces the same results as base_id + table_id.\"\"\"\n if not airtable_config.table_identifier.startswith(\"tbl\"):\n pytest.skip(\"URL test requires table ID, not table name\")\n\n url = f\"https://airtable.com/{airtable_config.base_id}/{airtable_config.table_identifier}/{BASE_VIEW_ID}\"\n connector = AirtableConnector(\n airtable_url=url,\n treat_all_non_attachment_fields_as_metadata=False,\n )\n connector.load_credentials({\"airtable_access_token\": airtable_config.access_token})\n\n doc_batch_generator = connector.load_from_state()\n doc_batch = [\n doc for doc in next(doc_batch_generator) if not isinstance(doc, HierarchyNode)\n ]\n with pytest.raises(StopIteration):\n next(doc_batch_generator)\n\n assert len(doc_batch) == 2\n\n expected_docs = [\n create_test_document(\n id=\"rec8BnxDLyWeegOuO\",\n title=\"Slow Internet\",\n description=\"The internet connection is very slow.\",\n priority=\"Medium\",\n status=\"In Progress\",\n ticket_id=\"2\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n all_fields_as_metadata=False,\n view_id=BASE_VIEW_ID,\n ),\n create_test_document(\n id=\"reccSlIA4pZEFxPBg\",\n title=\"Printer Issue\",\n description=\"The office printer is not working.\",\n priority=\"High\",\n status=\"Open\",\n ticket_id=\"1\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n attachments=[\n (\n \"Test.pdf:\\ntesting!!!\",\n f\"https://airtable.com/{airtable_config.base_id}/{airtable_config.table_identifier}/{BASE_VIEW_ID}/reccSlIA4pZEFxPBg/fld1u21zkJACIvAEF/attlj2UBWNEDZngCc?blocks=hide\",\n )\n ],\n all_fields_as_metadata=False,\n view_id=BASE_VIEW_ID,\n ),\n ]\n\n compare_documents(doc_batch, expected_docs)\n\n\ndef test_airtable_connector_index_all(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n airtable_config: AirtableConfig,\n) -> None:\n \"\"\"Test index_all mode discovers all bases/tables and returns documents.\n\n The test token has access to one base (\"Onyx\") with three tables:\n - Tickets: 3 records, 2 with content (1 empty record is skipped)\n - Support Categories: 4 records, all with Category Name field\n - Table 3: 3 records, 1 with content (2 empty records are skipped)\n Total expected: 7 documents\n \"\"\"\n connector = AirtableConnector()\n connector.load_credentials({\"airtable_access_token\": airtable_config.access_token})\n\n all_docs: list[Document] = []\n for batch in connector.load_from_state():\n for item in batch:\n if isinstance(item, Document):\n all_docs.append(item)\n\n # 2 from Tickets + 4 from Support Categories + 1 from Table 3 = 7\n assert len(all_docs) == 7\n\n docs_by_id = {d.id: d for d in all_docs}\n\n # Verify all expected document IDs are present\n expected_ids = {\n # Tickets\n \"airtable__rec8BnxDLyWeegOuO\",\n \"airtable__reccSlIA4pZEFxPBg\",\n # Support Categories\n \"airtable__rec5SgUDcHXcBc8kS\",\n \"airtable__recD3DQHc0BQkDaqX\",\n \"airtable__recPHdnWu1Q9ZxyTg\",\n \"airtable__recWbIElUDz9HjgMd\",\n # Table 3\n \"airtable__recNalBz02QU1LhbM\",\n }\n assert docs_by_id.keys() == expected_ids\n\n # In index_all mode, semantic identifiers include \"Base Name > Table Name: Primary Field\"\n assert (\n docs_by_id[\"airtable__rec8BnxDLyWeegOuO\"].semantic_identifier\n == \"Onyx > Tickets: Slow Internet\"\n )\n assert (\n docs_by_id[\"airtable__rec5SgUDcHXcBc8kS\"].semantic_identifier\n == \"Onyx > Support Categories: Software Development\"\n )\n assert (\n docs_by_id[\"airtable__recNalBz02QU1LhbM\"].semantic_identifier\n == \"Onyx > Table 3: A\"\n )\n\n # Verify hierarchy metadata on a Tickets doc\n tickets_doc = docs_by_id[\"airtable__rec8BnxDLyWeegOuO\"]\n assert tickets_doc.doc_metadata is not None\n hierarchy = tickets_doc.doc_metadata[\"hierarchy\"]\n assert hierarchy[\"source_path\"] == [\"Onyx\", \"Tickets\"]\n assert hierarchy[\"base_id\"] == airtable_config.base_id\n assert hierarchy[\"base_name\"] == \"Onyx\"\n assert hierarchy[\"table_name\"] == \"Tickets\"\n\n # Verify hierarchy on a Support Categories doc\n cat_doc = docs_by_id[\"airtable__rec5SgUDcHXcBc8kS\"]\n assert cat_doc.doc_metadata is not None\n assert cat_doc.doc_metadata[\"hierarchy\"][\"source_path\"] == [\n \"Onyx\",\n \"Support Categories\",\n ]\n\n\ndef test_airtable_connector_all_metadata(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n airtable_config: AirtableConfig,\n) -> None:\n connector = AirtableConnector(\n base_id=airtable_config.base_id,\n table_name_or_id=airtable_config.table_identifier,\n treat_all_non_attachment_fields_as_metadata=True,\n )\n connector.load_credentials(\n {\n \"airtable_access_token\": airtable_config.access_token,\n }\n )\n doc_batch_generator = connector.load_from_state()\n doc_batch = [\n doc for doc in next(doc_batch_generator) if not isinstance(doc, HierarchyNode)\n ]\n with pytest.raises(StopIteration):\n next(doc_batch_generator)\n\n # NOTE: one of the rows has no attachments -> no content -> no document\n assert len(doc_batch) == 1\n\n expected_docs = [\n create_test_document(\n id=\"reccSlIA4pZEFxPBg\",\n title=\"Printer Issue\",\n description=\"The office printer is not working.\",\n priority=\"High\",\n status=\"Open\",\n # Link to another record is skipped for now\n # category=\"Software Development\",\n ticket_id=\"1\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n attachments=[\n (\n \"Test.pdf:\\ntesting!!!\",\n # hard code link for now\n \"https://airtable.com/appCXJqDFS4gea8tn/tblRxFQsTlBBZdRY1/viwVUEJjWPd8XYjh8/reccSlIA4pZEFxPBg/fld1u21zkJACIvAEF/attlj2UBWNEDZngCc?blocks=hide\",\n )\n ],\n all_fields_as_metadata=True,\n view_id=BASE_VIEW_ID,\n ),\n ]\n\n # Compare documents using the utility function\n compare_documents(doc_batch, expected_docs)\n\n\ndef test_airtable_connector_with_share_and_view(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n airtable_config: AirtableConfig,\n) -> None:\n \"\"\"Test behavior when using share_id and view_id for URL generation.\"\"\"\n SHARE_ID = \"shrkfjEzDmLaDtK83\"\n\n connector = AirtableConnector(\n base_id=airtable_config.base_id,\n table_name_or_id=airtable_config.table_identifier,\n treat_all_non_attachment_fields_as_metadata=False,\n share_id=SHARE_ID,\n view_id=BASE_VIEW_ID,\n )\n connector.load_credentials(\n {\n \"airtable_access_token\": airtable_config.access_token,\n }\n )\n doc_batch_generator = connector.load_from_state()\n doc_batch = [\n doc for doc in next(doc_batch_generator) if not isinstance(doc, HierarchyNode)\n ]\n with pytest.raises(StopIteration):\n next(doc_batch_generator)\n\n assert len(doc_batch) == 2\n\n expected_docs = [\n create_test_document(\n id=\"rec8BnxDLyWeegOuO\",\n title=\"Slow Internet\",\n description=\"The internet connection is very slow.\",\n priority=\"Medium\",\n status=\"In Progress\",\n ticket_id=\"2\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n all_fields_as_metadata=False,\n share_id=SHARE_ID,\n view_id=BASE_VIEW_ID,\n ),\n create_test_document(\n id=\"reccSlIA4pZEFxPBg\",\n title=\"Printer Issue\",\n description=\"The office printer is not working.\",\n priority=\"High\",\n status=\"Open\",\n ticket_id=\"1\",\n created_time=\"2024-12-24T21:02:49.000Z\",\n status_last_changed=\"2024-12-24T21:02:49.000Z\",\n days_since_status_change=0,\n assignee=\"Chris Weaver (chris@onyx.app)\",\n submitted_by=\"Chris Weaver (chris@onyx.app)\",\n attachments=[\n (\n \"Test.pdf:\\ntesting!!!\",\n (\n f\"https://airtable.com/{airtable_config.base_id}/{SHARE_ID}/\"\n f\"{os.environ['AIRTABLE_TEST_TABLE_ID']}/{BASE_VIEW_ID}/reccSlIA4pZEFxPBg/\"\n \"fld1u21zkJACIvAEF/attlj2UBWNEDZngCc?blocks=hide\"\n ),\n )\n ],\n all_fields_as_metadata=False,\n share_id=SHARE_ID,\n view_id=BASE_VIEW_ID,\n ),\n ]\n\n # Compare documents using the utility function\n compare_documents(doc_batch, expected_docs)\n" }, { "path": "backend/tests/daily/connectors/bitbucket/conftest.py", "content": "from tests.load_env_vars import load_env_vars\n\n\n# Load environment variables at the module level\nload_env_vars()\n" }, { "path": "backend/tests/daily/connectors/bitbucket/test_bitbucket_checkpointed.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.bitbucket.connector import BitbucketConnector\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\n@pytest.fixture\ndef bitbucket_connector_for_checkpoint() -> BitbucketConnector:\n \"\"\"Daily fixture for Bitbucket checkpointed indexing.\n\n Env vars:\n - BITBUCKET_EMAIL: Bitbucket account email\n - BITBUCKET_API_TOKEN: Bitbucket app password/token\n - BITBUCKET_WORKSPACE: workspace id\n - BITBUCKET_REPOSITORIES: comma-separated slugs\n - BITBUCKET_PROJECTS: optional comma-separated project keys\n \"\"\"\n workspace = os.environ[\"BITBUCKET_WORKSPACE\"]\n repositories = os.environ.get(\"BITBUCKET_REPOSITORIES\")\n projects = os.environ.get(\"BITBUCKET_PROJECTS\")\n\n connector = BitbucketConnector(\n workspace=workspace,\n repositories=repositories,\n projects=projects,\n batch_size=10,\n )\n\n email = os.environ.get(\"BITBUCKET_EMAIL\")\n token = os.environ.get(\"BITBUCKET_API_TOKEN\")\n if not email or not token:\n pytest.skip(\"BITBUCKET_EMAIL or BITBUCKET_API_TOKEN not set in environment\")\n\n connector.load_credentials({\"bitbucket_email\": email, \"bitbucket_api_token\": token})\n return connector\n\n\ndef test_bitbucket_checkpointed_load(\n bitbucket_connector_for_checkpoint: BitbucketConnector,\n) -> None:\n # Use a broad window; results may be empty depending on repository state\n start = 1755004439 # Tue Aug 12 2025 13:13:59 UTC\n end = time.time()\n\n docs = load_all_from_connector(\n connector=bitbucket_connector_for_checkpoint,\n start=start,\n end=end,\n ).documents\n\n assert isinstance(docs, list)\n\n for doc in docs:\n assert doc.source == DocumentSource.BITBUCKET\n assert doc.metadata is not None\n assert doc.metadata.get(\"object_type\") == \"PullRequest\"\n assert \"id\" in doc.metadata\n assert \"state\" in doc.metadata\n assert \"title\" in doc.metadata\n assert \"updated_on\" in doc.metadata\n\n # Basic section checks\n assert len(doc.sections) >= 1\n section = doc.sections[0]\n assert isinstance(section.link, str)\n assert isinstance(section.text, str)\n" }, { "path": "backend/tests/daily/connectors/bitbucket/test_bitbucket_slim_connector.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.bitbucket.connector import BitbucketConnector\nfrom onyx.connectors.models import HierarchyNode\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\n@pytest.fixture\ndef bitbucket_connector_for_slim() -> BitbucketConnector:\n workspace = os.environ[\"BITBUCKET_WORKSPACE\"]\n repositories = os.environ.get(\"BITBUCKET_REPOSITORIES\")\n projects = os.environ.get(\"BITBUCKET_PROJECTS\")\n\n connector = BitbucketConnector(\n workspace=workspace,\n repositories=repositories,\n projects=projects,\n batch_size=10,\n )\n\n email = os.environ.get(\"BITBUCKET_EMAIL\")\n token = os.environ.get(\"BITBUCKET_API_TOKEN\")\n if not email or not token:\n pytest.skip(\"BITBUCKET_EMAIL or BITBUCKET_API_TOKEN not set in environment\")\n\n connector.load_credentials({\"bitbucket_email\": email, \"bitbucket_api_token\": token})\n return connector\n\n\ndef test_bitbucket_full_ids_subset_of_slim_ids(\n bitbucket_connector_for_slim: BitbucketConnector,\n) -> None:\n # Get all full doc IDs from load_from_state\n docs = load_all_from_connector(\n connector=bitbucket_connector_for_slim,\n start=0,\n end=time.time(),\n ).documents\n all_full_doc_ids: set[str] = set([doc.id for doc in docs])\n\n # Get all doc IDs from the slim connector\n all_slim_doc_ids: set[str] = set()\n for (\n slim_doc_batch\n ) in bitbucket_connector_for_slim.retrieve_all_slim_docs_perm_sync():\n all_slim_doc_ids.update(\n [doc.id for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # The set of full doc IDs should always be a subset of slim doc IDs\n assert all_full_doc_ids.issubset(all_slim_doc_ids)\n # Make sure we actually got some documents\n assert len(all_slim_doc_ids) > 0\n\n # Basic sanity checks if any docs exist\n if all_slim_doc_ids:\n example_id = next(iter(all_slim_doc_ids))\n assert example_id.startswith(f\"{DocumentSource.BITBUCKET.value}:\")\n" }, { "path": "backend/tests/daily/connectors/blob/test_blob_connector.py", "content": "import os\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom urllib.parse import parse_qs\nfrom urllib.parse import unquote\nfrom urllib.parse import urlparse\n\nimport pytest\n\nfrom onyx.configs.constants import BlobType\nfrom onyx.connectors.blob.connector import BlobStorageConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.file_processing.extract_file_text import get_file_ext\nfrom onyx.file_processing.file_types import OnyxFileExtensions\n\n\n@pytest.fixture\ndef blob_connector(request: pytest.FixtureRequest) -> BlobStorageConnector:\n \"\"\"Fixture requires (BlobType, bucket_name) and optional init kwargs.\n\n Param format: (BlobType, bucket_name, {optional init kwargs})\n - The 3rd element is optional and, if provided, must be a dict.\n - Extra kwargs are passed to BlobStorageConnector.__init__.\n\n Example:\n @pytest.mark.parametrize(\n \"blob_connector\",\n [(BlobType.S3, \"my-bucket\"), (BlobType.S3, \"my-bucket\", {\"prefix\": \"foo/\"})],\n indirect=True,\n )\n \"\"\"\n try:\n bucket_type, bucket_name, *rest = request.param\n except Exception as e:\n raise AssertionError(\n \"blob_connector requires (BlobType, bucket_name, [init_kwargs])\"\n ) from e\n\n init_kwargs = rest[0] if rest else {}\n if rest and not isinstance(init_kwargs, dict):\n raise AssertionError(\"init_kwargs must be a dict if provided\")\n\n if not isinstance(bucket_type, BlobType):\n bucket_type = BlobType(bucket_type)\n\n connector = BlobStorageConnector(\n bucket_type=bucket_type, bucket_name=bucket_name, **init_kwargs\n )\n\n if bucket_type == BlobType.S3:\n creds = {\n \"aws_access_key_id\": os.environ[\"AWS_ACCESS_KEY_ID_DAILY_CONNECTOR_TESTS\"],\n \"aws_secret_access_key\": os.environ[\n \"AWS_SECRET_ACCESS_KEY_DAILY_CONNECTOR_TESTS\"\n ],\n }\n elif bucket_type == BlobType.R2:\n creds = {\n \"account_id\": os.environ[\"R2_ACCOUNT_ID_DAILY_CONNECTOR_TESTS\"],\n \"r2_access_key_id\": os.environ[\"R2_ACCESS_KEY_ID_DAILY_CONNECTOR_TESTS\"],\n \"r2_secret_access_key\": os.environ[\n \"R2_SECRET_ACCESS_KEY_DAILY_CONNECTOR_TESTS\"\n ],\n }\n elif bucket_type == BlobType.GOOGLE_CLOUD_STORAGE:\n creds = {\n \"access_key_id\": os.environ[\"GCS_ACCESS_KEY_ID_DAILY_CONNECTOR_TESTS\"],\n \"secret_access_key\": os.environ[\n \"GCS_SECRET_ACCESS_KEY_DAILY_CONNECTOR_TESTS\"\n ],\n }\n else:\n # Until we figure out the Oracle log in, this fixture only supports S3, R2, and GCS.\n raise AssertionError(f\"Unsupported bucket type: {bucket_type}\")\n\n connector.load_credentials(creds)\n return connector\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\n@pytest.mark.parametrize(\n \"blob_connector\", [(BlobType.S3, \"onyx-connector-tests\")], indirect=True\n)\ndef test_blob_s3_connector(\n mock_get_api_key: MagicMock, # noqa: ARG001\n blob_connector: BlobStorageConnector,\n) -> None:\n \"\"\"\n Plain and document file types should be fully indexed.\n\n Multimedia and unknown file types will be indexed be skipped unless `set_allow_images`\n is called with `True`.\n\n This is intentional in order to allow searching by just the title even if we can't\n index the file content.\n \"\"\"\n all_docs: list[Document] = []\n document_batches = blob_connector.load_from_state()\n for doc_batch in document_batches:\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs.append(doc)\n\n assert len(all_docs) == 15\n\n for doc in all_docs:\n section = doc.sections[0]\n assert isinstance(section, TextSection)\n\n file_extension = get_file_ext(doc.semantic_identifier)\n if file_extension in OnyxFileExtensions.TEXT_AND_DOCUMENT_EXTENSIONS:\n assert len(section.text) > 0\n continue\n\n # unknown extension\n assert len(section.text) == 0\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\n@pytest.mark.parametrize(\n \"blob_connector\", [(BlobType.S3, \"s3-role-connector-test\")], indirect=True\n)\ndef test_blob_s3_cross_region_and_citation_link(\n mock_get_api_key: MagicMock, # noqa: ARG001\n blob_connector: BlobStorageConnector,\n) -> None:\n \"\"\"Buckets in a different region should be accessible and links should reflect the correct region.\n\n Validates that using the same credentials we can access a bucket in a\n different AWS region and that the generated object URL includes the bucket's\n region and is a valid S3 dashboard URL.\n \"\"\"\n\n assert blob_connector.bucket_region == \"ap-south-1\"\n\n # Load documents and validate the single object + its link\n all_docs: list[Document] = []\n for doc_batch in blob_connector.load_from_state():\n all_docs.extend(\n [doc for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # The test bucket contains exactly one object named \"Chapter 6.pdf\"\n assert len(all_docs) == 1\n doc = all_docs[0]\n assert doc.semantic_identifier == \"Chapter 6.pdf\"\n\n # Validate link\n assert len(doc.sections) >= 1\n link = doc.sections[0].link\n assert link is not None and isinstance(link, str) and len(link) > 0\n\n parsed = urlparse(link)\n # Expect the link to be the AWS S3 console object URL\n assert parsed.netloc == \"s3.console.aws.amazon.com\"\n assert parsed.path == \"/s3/object/s3-role-connector-test\"\n\n # Query should include region and prefix\n query = parse_qs(parsed.query)\n assert query.get(\"region\") == [\"ap-south-1\"]\n assert \"prefix\" in query and len(query[\"prefix\"]) == 1\n prefix_val = query[\"prefix\"][0]\n # The prefix (object key) should decode to the filename\n decoded_prefix = unquote(prefix_val)\n assert decoded_prefix == \"Chapter 6.pdf\" or decoded_prefix.endswith(\n \"/Chapter 6.pdf\"\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\n@pytest.mark.parametrize(\n \"blob_connector\", [(BlobType.R2, \"asia-pacific-bucket\")], indirect=True\n)\ndef test_blob_r2_connector(\n mock_get_api_key: MagicMock, # noqa: ARG001\n blob_connector: BlobStorageConnector,\n) -> None:\n \"\"\"Validate basic R2 connector creation and document loading\"\"\"\n\n all_docs: list[Document] = []\n for doc_batch in blob_connector.load_from_state():\n all_docs.extend(\n [doc for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n assert len(all_docs) >= 1\n doc = all_docs[0]\n assert len(doc.sections) >= 1\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\n@pytest.mark.parametrize(\n \"blob_connector\",\n [(BlobType.R2, \"onyx-daily-connector-tests\", {\"european_residency\": True})],\n indirect=True,\n)\ndef test_blob_r2_eu_residency_connector(\n mock_get_api_key: MagicMock, # noqa: ARG001\n blob_connector: BlobStorageConnector,\n) -> None:\n \"\"\"Validate R2 connector with European residency setting\"\"\"\n\n all_docs: list[Document] = []\n for doc_batch in blob_connector.load_from_state():\n all_docs.extend(\n [doc for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n assert len(all_docs) >= 1\n doc = all_docs[0]\n assert len(doc.sections) >= 1\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\n@pytest.mark.parametrize(\n \"blob_connector\", [(BlobType.GOOGLE_CLOUD_STORAGE, \"onyx-test-1\")], indirect=True\n)\ndef test_blob_gcs_connector(\n mock_get_api_key: MagicMock, # noqa: ARG001\n blob_connector: BlobStorageConnector,\n) -> None:\n all_docs: list[Document] = []\n for doc_batch in blob_connector.load_from_state():\n all_docs.extend(\n [doc for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # At least one object from the test bucket\n assert len(all_docs) >= 1\n doc = all_docs[0]\n assert len(doc.sections) >= 1\n" }, { "path": "backend/tests/daily/connectors/coda/README.md", "content": "# Coda Connector Test Suite\n\n## Overview\n\nThe `test_coda_connector.py` file contains comprehensive integration tests for the Coda connector. These tests validate that the connector properly:\n- Authenticates with the Coda API\n- Retrieves documents, pages, and tables\n- Generates properly structured Onyx `Document` objects\n- Handles batch processing correctly\n- Supports workspace scoping\n- Polls for recent updates\n- Handles error cases gracefully\n\n## Prerequisites\n\n1. **Coda API Access**: You need a valid Coda account with at least one workspace containing documents, pages, or tables\n2. **Coda Bearer Token**: Generate an API token from your Coda account settings\n3. **Python Environment**: Backend dependencies installed (see backend/requirements)\n4. **Test Data**: Ideally, your Coda workspace should have:\n - At least one document\n - At least one page within a document\n - At least one table within a document\n\n## Environment Variables\n\nThe test suite requires the following environment variables:\n\n### Required\n- **`CODA_BEARER_TOKEN`**: Your Coda API bearer token\n - Get this from: Coda Account Settings → API Settings → Generate API Token\n - Without this, tests will be skipped\n\n### Optional\n- **`CODA_BASE_URL`**: The Coda API base URL\n - Default: `https://coda.io/apis/v1`\n - Only override if using a different API endpoint\n\n- **`CODA_WORKSPACE_ID`**: A specific workspace ID to test workspace scoping\n - If not provided, workspace-scoped tests will be skipped\n - Find this by inspecting the Coda API response or your workspace URL\n\n## Running the Tests\n\n### Method 1: Run All Tests in the File\n\nFrom the `backend/` directory:\n\n```bash\n# Set environment variables and run all tests\nexport CODA_BEARER_TOKEN=\"your_token_here\"\npytest -v -s tests/daily/connectors/coda/test_coda_connector.py\n```\n\n### Method 2: Run a Specific Test Class\n\n```bash\n# Run only validation tests\nexport CODA_BEARER_TOKEN=\"your_token_here\"\npytest -v -s tests/daily/connectors/coda/test_coda_connector.py::TestCodaConnectorValidation\n\n# Run only load_from_state tests\npytest -v -s tests/daily/connectors/coda/test_coda_connector.py::TestLoadFromState\n```\n\n### Method 3: Run a Single Test\n\n```bash\n# Run a specific test function\nexport CODA_BEARER_TOKEN=\"your_token_here\"\npytest -v -s tests/daily/connectors/coda/test_coda_connector.py::TestLoadFromState::test_document_count_matches_expected\n```\n\n### Method 4: Using an Environment File\n\nCreate a `.env` file in `backend/tests/daily/connectors/coda/`:\n\n```bash\n# .env\nCODA_BEARER_TOKEN=your_token_here\nCODA_WORKSPACE_ID=your_workspace_id # Optional\n```\n\nThen run with dotenv:\n\n```bash\ncd backend\npython -m dotenv -f tests/daily/connectors/coda/.env run -- pytest -v -s tests/daily/connectors/coda/test_coda_connector.py\n```\n\n### Method 5: Direct Execution\n\nThe test file can be run directly:\n\n```bash\nexport CODA_BEARER_TOKEN=\"your_token_here\"\ncd backend/tests/daily/connectors/coda\npython test_coda_connector.py\n```\n\n## Test Structure\n\n### Test Classes\n\n1. **`TestCodaConnectorValidation`**\n - Validates connector settings and credentials\n - Tests authentication success and failure cases\n - Tests workspace-scoped connector validation\n\n2. **`TestLoadFromState`**\n - Tests full document retrieval via `load_from_state()`\n - Validates batch sizes, document counts, and structure\n - Checks document fields, metadata, and content\n - Verifies both page and table document generation\n - Tests the `index_page_content` configuration flag\n\n3. **`TestPollSource`**\n - Tests incremental updates via `poll_source()`\n - Validates time-range filtering\n - Checks that only updated documents are returned\n\n4. **`TestWorkspaceScoping`**\n - Tests the `workspace_id` filtering functionality\n - Validates that scoped connectors only retrieve documents from the specified workspace\n\n5. **`TestErrorHandling`**\n - Tests graceful handling of edge cases\n - Validates behavior with inaccessible content or empty tables\n\n## Common Test Patterns\n\n### Fixtures\n\nThe test suite uses pytest fixtures for setup:\n\n- **`coda_credentials`**: Loads and validates credentials from environment variables\n- **`connector`**: Creates a standard CodaConnector instance\n- **`workspace_scoped_connector`**: Creates a workspace-scoped connector (if `CODA_WORKSPACE_ID` is set)\n- **`reference_data`**: Fetches ground truth data from the Coda API for validation\n\n### Skipped Tests\n\nTests are automatically skipped when:\n- `CODA_BEARER_TOKEN` is not set\n- `CODA_WORKSPACE_ID` is not set (for workspace-scoped tests)\n- No documents, pages, or tables are found in the workspace\n\n## Troubleshooting\n\n### Tests are Skipped\n\n**Issue**: Tests show as \"SKIPPED\" instead of running\n\n**Solutions**:\n- Ensure `CODA_BEARER_TOKEN` is set and valid\n- Verify your Coda workspace has at least one document with pages or tables\n- For workspace tests, ensure `CODA_WORKSPACE_ID` is set\n\n### Authentication Errors\n\n**Issue**: Tests fail with authentication errors\n\n**Solutions**:\n- Verify your bearer token is valid and hasn't expired\n- Check that the token has appropriate API permissions\n- Ensure you're not hitting API rate limits\n\n### Document Count Mismatches\n\n**Issue**: Tests fail with \"Expected X documents but got Y\"\n\n**Possible Causes**:\n- API rate limiting causing partial data retrieval\n- Network issues during test execution\n- Changes to workspace data during test execution\n- Permission issues preventing access to some documents\n\n### Empty Content Errors\n\n**Issue**: Tests fail due to empty document content\n\n**Possible Causes**:\n- Pages without accessible content (permission issues)\n- Empty tables or pages in your workspace\n- The `index_page_content` flag set incorrectly\n\n## Test Execution Tips\n\n1. **Run tests during low-traffic times**: API rate limits may affect test reliability\n2. **Use a dedicated test workspace**: Avoid running tests on production workspaces with changing data\n3. **Check test output verbosity**: Use `-v` for verbose test names, `-s` to see print statements\n4. **Isolate failing tests**: Run specific test classes or functions to debug issues\n5. **Review fixture output**: The `reference_data` fixture prints warnings about API access issues\n\n## CI/CD Integration\n\nWhen integrating these tests into CI/CD pipelines:\n\n```yaml\n# Example GitHub Actions configuration\n- name: Run Coda Connector Tests\n env:\n CODA_BEARER_TOKEN: ${{ secrets.CODA_BEARER_TOKEN }}\n CODA_WORKSPACE_ID: ${{ secrets.CODA_WORKSPACE_ID }}\n run: |\n cd backend\n pytest -v tests/daily/connectors/coda/test_coda_connector.py\n```\n\nStore credentials as encrypted secrets in your CI/CD platform.\n\n## Expected Test Duration\n\n- Full test suite: ~30-60 seconds (depending on workspace size and API latency)\n- Individual test classes: ~5-15 seconds\n- Validation tests: <5 seconds\n\n## Additional Resources\n\n- [Coda API Documentation](https://coda.io/developers/apis/v1)\n- [Onyx Connector Documentation](../../../../onyx/connectors/README.md)\n- [pytest Documentation](https://docs.pytest.org/)\n" }, { "path": "backend/tests/daily/connectors/coda/test_coda_connector.py", "content": "import os\nimport time\nfrom collections.abc import Generator\nfrom typing import Any\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.coda.connector import CodaConnector\nfrom onyx.connectors.exceptions import CredentialInvalidError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\n\n\ndef connector_doc_generator(\n connector: CodaConnector,\n) -> Generator[Document, None, None]:\n for batch in connector.load_from_state():\n for doc in batch:\n if isinstance(doc, HierarchyNode):\n continue\n yield doc\n\n\n@pytest.fixture\ndef coda_credentials() -> dict[str, str]:\n \"\"\"Fixture to get and validate Coda credentials.\"\"\"\n bearer_token = os.environ.get(\"CODA_BEARER_TOKEN\")\n\n if not bearer_token:\n pytest.skip(\"CODA_BEARER_TOKEN not set\")\n\n return {\n \"coda_bearer_token\": bearer_token,\n }\n\n\n@pytest.fixture\ndef connector(coda_credentials: dict[str, str]) -> CodaConnector:\n \"\"\"Fixture to create and authenticate connector.\"\"\"\n conn = CodaConnector(batch_size=5, index_page_content=True)\n conn.load_credentials(coda_credentials)\n return conn\n\n\n@pytest.fixture\ndef workspace_scoped_connector(coda_credentials: dict[str, str]) -> CodaConnector:\n \"\"\"Fixture to create connector scoped to a specific workspace (if CODA_WORKSPACE_ID is set).\"\"\"\n workspace_id = os.environ.get(\"CODA_WORKSPACE_ID\")\n if not workspace_id:\n pytest.skip(\"CODA_WORKSPACE_ID not set - skipping workspace-scoped tests\")\n\n conn = CodaConnector(\n batch_size=5, index_page_content=True, workspace_id=workspace_id\n )\n conn.load_credentials(coda_credentials)\n return conn\n\n\n@pytest.fixture\ndef reference_data(connector: CodaConnector) -> dict[str, Any]:\n \"\"\"Fixture to fetch reference data from API for validation.\"\"\"\n all_docs = connector._list_all_docs()\n\n if not all_docs:\n pytest.skip(\"No docs found in Coda workspace\")\n\n expected_page_count = 0\n expected_table_count = 0\n pages_by_doc = {}\n tables_by_doc = {}\n\n for doc in all_docs:\n doc_id = doc.id\n\n try:\n pages = connector._list_pages_in_doc(doc_id)\n pages_by_doc[doc_id] = pages\n expected_page_count += len(pages)\n except Exception as e:\n print(f\"Warning: Could not fetch pages for doc {doc_id}: {e}\")\n pages_by_doc[doc_id] = []\n\n try:\n tables = connector._list_tables(doc_id)\n tables_by_doc[doc_id] = tables\n expected_table_count += len(tables)\n except Exception as e:\n print(f\"Warning: Could not fetch tables for doc {doc_id}: {e}\")\n tables_by_doc[doc_id] = []\n\n total_expected_documents = expected_page_count + expected_table_count\n\n if total_expected_documents == 0:\n pytest.skip(\"No pages or tables found in Coda workspace\")\n\n return {\n \"docs\": all_docs,\n \"total_pages\": expected_page_count,\n \"total_tables\": expected_table_count,\n \"total_documents\": total_expected_documents,\n \"pages_by_doc\": pages_by_doc,\n \"tables_by_doc\": tables_by_doc,\n }\n\n\nclass TestCodaConnectorValidation:\n \"\"\"Test suite for connector validation and credential handling.\"\"\"\n\n def test_validate_connector_settings_success(\n self, connector: CodaConnector\n ) -> None:\n \"\"\"Test that validate_connector_settings succeeds with valid credentials.\"\"\"\n # Should not raise any exceptions\n connector.validate_connector_settings()\n\n def test_validate_workspace_scoped_connector(\n self, workspace_scoped_connector: CodaConnector\n ) -> None:\n \"\"\"Test that workspace-scoped connector validates successfully.\"\"\"\n workspace_scoped_connector.validate_connector_settings()\n\n def test_load_credentials_invalid_token(self) -> None:\n \"\"\"Test that invalid credentials are rejected.\"\"\"\n conn = CodaConnector()\n\n with pytest.raises(CredentialInvalidError):\n conn.load_credentials(\n {\n \"coda_bearer_token\": \"invalid_token_12345\",\n }\n )\n\n\nclass TestLoadFromState:\n \"\"\"Test suite for load_from_state functionality.\"\"\"\n\n def test_returns_generator(self, connector: CodaConnector) -> None:\n \"\"\"Test that load_from_state returns a generator.\"\"\"\n gen = connector.load_from_state()\n assert isinstance(gen, Generator), \"load_from_state should return a Generator\"\n\n def test_batch_sizes_respect_config(\n self,\n connector: CodaConnector,\n reference_data: dict[str, Any], # noqa: ARG002\n ) -> None:\n \"\"\"Test that batches respect the configured batch_size.\"\"\"\n batch_size = connector.batch_size\n gen = connector.load_from_state()\n\n batch_sizes = []\n for batch in gen:\n batch_sizes.append(len(batch))\n assert (\n len(batch) <= batch_size\n ), f\"Batch size {len(batch)} exceeds configured {batch_size}\"\n\n for i, size in enumerate(batch_sizes[:-1]):\n assert (\n size == batch_size\n ), f\"Non-final batch {i} has size {size}, expected {batch_size}\"\n\n # Last batch may be smaller or equal\n if batch_sizes:\n assert batch_sizes[-1] <= batch_size\n\n def test_document_count_matches_expected(\n self, connector: CodaConnector, reference_data: dict[str, Any]\n ) -> None:\n \"\"\"Test that total documents match expected pages + tables.\"\"\"\n gen = connector.load_from_state()\n\n total_documents = sum(len(batch) for batch in gen)\n expected_count = reference_data[\"total_documents\"]\n\n assert total_documents == expected_count, (\n f\"Expected {expected_count} documents \"\n f\"({reference_data['total_pages']} pages + \"\n f\"{reference_data['total_tables']} tables) \"\n f\"but got {total_documents}\"\n )\n\n def test_document_required_fields(\n self,\n connector: CodaConnector,\n reference_data: dict[str, Any], # noqa: ARG002\n ) -> None:\n \"\"\"Test that all documents have required fields with valid values.\"\"\"\n gen = connector.load_from_state()\n\n for batch in gen:\n for doc in batch:\n assert isinstance(doc, Document)\n\n assert doc.id is not None, \"Document ID should not be None\"\n assert doc.id.startswith(\n \"coda-\"\n ), \"Document ID should start with 'coda-'\"\n assert (\n doc.source == DocumentSource.CODA\n ), \"Document source should be CODA\"\n assert (\n doc.semantic_identifier is not None\n ), \"Semantic identifier should not be None\"\n assert (\n doc.doc_updated_at is not None\n ), \"doc_updated_at should not be None\"\n\n assert (\n len(doc.sections) > 0\n ), \"Document should have at least one section\"\n for section in doc.sections:\n assert section.text is not None, \"Section text should not be None\"\n assert len(section.text) > 0, \"Section text should not be empty\"\n assert section.link is not None, \"Section link should not be None\"\n assert section.link.startswith(\n \"https://\"\n ), \"Section link should be a valid URL\"\n\n assert \"doc_id\" in doc.metadata, \"Metadata should contain doc_id\"\n assert (\n \"browser_link\" in doc.metadata\n ), \"Metadata should contain browser_link\"\n\n def test_document_types(\n self, connector: CodaConnector, reference_data: dict[str, Any]\n ) -> None:\n \"\"\"Test that both page and table documents are generated correctly.\"\"\"\n page_docs = []\n table_docs = []\n\n for doc in connector_doc_generator(connector):\n if \"coda-page-\" in doc.id:\n page_docs.append(doc)\n assert \"content_type\" in doc.metadata\n elif \"coda-table-\" in doc.id:\n table_docs.append(doc)\n assert \"row_count\" in doc.metadata\n\n # Verify we found both types (if both exist in the workspace)\n if reference_data[\"total_pages\"] > 0:\n assert len(page_docs) > 0, \"Should have found page documents\"\n\n if reference_data[\"total_tables\"] > 0:\n assert len(table_docs) > 0, \"Should have found table documents\"\n\n # Verify counts match\n assert (\n len(page_docs) == reference_data[\"total_pages\"]\n ), f\"Expected {reference_data['total_pages']} page documents, got {len(page_docs)}\"\n assert (\n len(table_docs) == reference_data[\"total_tables\"]\n ), f\"Expected {reference_data['total_tables']} table documents, got {len(table_docs)}\"\n\n def test_no_duplicate_documents(\n self,\n connector: CodaConnector,\n reference_data: dict[str, Any], # noqa: ARG002\n ) -> None:\n \"\"\"Test that no documents are yielded twice.\"\"\"\n document_ids = []\n for doc in connector_doc_generator(connector):\n document_ids.append(doc.id)\n\n unique_ids = set(document_ids)\n assert len(document_ids) == len(\n unique_ids\n ), f\"Found {len(document_ids) - len(unique_ids)} duplicate documents\"\n\n def test_all_docs_processed(\n self, connector: CodaConnector, reference_data: dict[str, Any]\n ) -> None:\n \"\"\"Test that content from all docs are included.\"\"\"\n processed_doc_ids = set()\n for doc in connector_doc_generator(connector):\n doc_id = doc.metadata.get(\"doc_id\")\n processed_doc_ids.add(doc_id)\n\n expected_doc_ids = {doc.id for doc in reference_data[\"docs\"]}\n\n expected_doc_ids_with_content = {\n doc_id\n for doc_id in expected_doc_ids\n if len(reference_data[\"pages_by_doc\"].get(doc_id, [])) > 0\n or len(reference_data[\"tables_by_doc\"].get(doc_id, [])) > 0\n }\n\n assert (\n processed_doc_ids == expected_doc_ids_with_content\n ), f\"Not all docs with content were processed. Expected {expected_doc_ids_with_content}, got {processed_doc_ids}\"\n\n def test_document_content_not_empty(\n self,\n connector: CodaConnector,\n reference_data: dict[str, Any], # noqa: ARG002\n ) -> None:\n \"\"\"Test that all documents have meaningful content.\"\"\"\n for doc in connector_doc_generator(connector):\n assert doc.semantic_identifier, \"Semantic identifier should not be empty\"\n assert (\n len(doc.semantic_identifier) > 0\n ), \"Semantic identifier should have content\"\n\n total_text_length = sum(len(section.text or \"\") for section in doc.sections)\n assert total_text_length > 0, f\"Document {doc.id} has no content\"\n\n def test_page_content_indexing(self, coda_credentials: dict[str, str]) -> None:\n \"\"\"Test that index_page_content flag works correctly.\"\"\"\n # page indexing disabled\n conn_no_content = CodaConnector(batch_size=5, index_page_content=False)\n conn_no_content.load_credentials(coda_credentials)\n\n # page indexing enabled\n conn_with_content = CodaConnector(batch_size=5, index_page_content=True)\n conn_with_content.load_credentials(coda_credentials)\n\n docs_no_content = []\n for batch in conn_no_content.load_from_state():\n for doc in batch:\n if isinstance(doc, HierarchyNode):\n continue\n if \"coda-page-\" in doc.id:\n docs_no_content.append(doc)\n break\n if docs_no_content:\n break\n\n docs_with_content = []\n for batch in conn_with_content.load_from_state():\n for doc in batch:\n if isinstance(doc, HierarchyNode):\n continue\n if \"coda-page-\" in doc.id:\n docs_with_content.append(doc)\n break\n if docs_with_content:\n break\n\n if docs_no_content and docs_with_content:\n no_content_length = sum(\n len(s.text or \"\") for s in docs_no_content[0].sections\n )\n with_content_length = sum(\n len(s.text or \"\") for s in docs_with_content[0].sections\n )\n\n assert (\n with_content_length >= no_content_length\n ), \"Content-indexed page should have at least as much text as non-indexed\"\n\n\nclass TestPollSource:\n \"\"\"Test suite for poll_source functionality.\"\"\"\n\n def test_poll_source_returns_generator(self, connector: CodaConnector) -> None:\n \"\"\"Test that poll_source returns a generator.\"\"\"\n current_time = time.time()\n start_time = current_time - 86400 # 24 hours\n\n gen = connector.poll_source(start_time, current_time)\n assert isinstance(gen, Generator), \"poll_source should return a Generator\"\n\n def test_poll_source_recent_updates(self, connector: CodaConnector) -> None:\n \"\"\"Test polling for recently updated documents.\"\"\"\n current_time = time.time()\n start_time = current_time - (86400 * 30)\n\n gen = connector.poll_source(start_time, current_time)\n\n documents = []\n for batch in gen:\n documents.extend(batch)\n\n # All returned documents should be updated within the time range\n for doc in documents:\n if isinstance(doc, HierarchyNode):\n continue\n assert doc.doc_updated_at is not None, \"doc_updated_at should not be None\"\n doc_timestamp = doc.doc_updated_at.timestamp()\n assert (\n start_time < doc_timestamp <= current_time\n ), f\"Document {doc.id} updated at {doc_timestamp} is outside range [{start_time}, {current_time}]\"\n\n def test_poll_source_no_updates_in_range(self, connector: CodaConnector) -> None:\n \"\"\"Test polling with a time range that has no updates.\"\"\"\n end_time = time.time() - (86400 * 365) # 1 year ago\n start_time = end_time - 86400 # 1 day before that\n\n gen = connector.poll_source(start_time, end_time)\n\n documents = []\n for batch in gen:\n documents.extend(batch)\n\n # Should return no documents (unless workspace is very old)\n print(f\"Found {len(documents)} documents updated over a year ago\")\n assert len(documents) == 0\n\n def test_poll_source_batch_sizes(self, connector: CodaConnector) -> None:\n \"\"\"Test that poll_source respects batch sizes.\"\"\"\n current_time = time.time()\n start_time = current_time - (86400 * 30)\n\n batch_size = connector.batch_size\n gen = connector.poll_source(start_time, current_time)\n\n for batch in gen:\n assert (\n len(batch) <= batch_size\n ), f\"Batch size {len(batch)} exceeds configured {batch_size}\"\n\n\nclass TestWorkspaceScoping:\n \"\"\"Test suite for workspace_id scoping functionality.\"\"\"\n\n def test_workspace_scoped_loads_subset(\n self,\n connector: CodaConnector,\n workspace_scoped_connector: CodaConnector,\n reference_data: dict[str, Any], # noqa: ARG002\n ) -> None:\n \"\"\"Test that workspace-scoped connector loads a subset of documents.\"\"\"\n all_docs = []\n for batch in connector.load_from_state():\n all_docs.extend(batch)\n\n scoped_docs = []\n for batch in workspace_scoped_connector.load_from_state():\n scoped_docs.extend(batch)\n\n # Scoped should be <= all docs\n assert len(scoped_docs) <= len(\n all_docs\n ), \"Workspace-scoped connector should return same or fewer documents\"\n\n workspace_id = workspace_scoped_connector.workspace_id\n for doc in scoped_docs:\n if isinstance(doc, HierarchyNode):\n continue\n doc_id = doc.metadata.get(\"doc_id\")\n assert isinstance(doc_id, str), \"doc_id should be a string\"\n coda_doc = workspace_scoped_connector._get_doc(doc_id)\n assert (\n coda_doc.workspace_id == workspace_id\n ), f\"Document {doc_id} has workspace {coda_doc.workspace_id}, expected {workspace_id}\"\n\n\nclass TestErrorHandling:\n \"\"\"Test suite for error handling and edge cases.\"\"\"\n\n def test_handles_missing_page_content_gracefully(\n self, connector: CodaConnector\n ) -> None:\n \"\"\"Test that connector handles pages without accessible content.\"\"\"\n gen = connector.load_from_state()\n\n documents = []\n for batch in gen:\n documents.extend(batch)\n\n assert (\n len(documents) > 0\n ), \"Should yield documents even if some content is inaccessible\"\n\n def test_handles_empty_tables_gracefully(self, connector: CodaConnector) -> None:\n \"\"\"Test that connector handles tables with no rows.\"\"\"\n for doc in connector_doc_generator(connector):\n if \"coda-table-\" in doc.id:\n assert len(doc.sections) > 0, \"Empty table should still have a section\"\n if doc.metadata.get(\"row_count\") == \"0\":\n assert (\n len(doc.sections) == 1\n ), \"Empty table should have exactly one section\"\n\n\nif __name__ == \"__main__\":\n pytest.main([__file__, \"-v\", \"-s\"])\n" }, { "path": "backend/tests/daily/connectors/confluence/models.py", "content": "from pydantic import BaseModel\n\nfrom ee.onyx.db.external_perm import ExternalUserGroup\n\n\nclass ExternalUserGroupSet(BaseModel):\n \"\"\"A version of ExternalUserGroup that uses a set for user_emails to avoid order-dependent comparisons.\"\"\"\n\n id: str\n user_emails: set[str]\n gives_anyone_access: bool\n\n @classmethod\n def from_model(\n cls, external_user_group: ExternalUserGroup\n ) -> \"ExternalUserGroupSet\":\n \"\"\"Convert from ExternalUserGroup to ExternalUserGroupSet.\"\"\"\n return cls(\n id=external_user_group.id,\n user_emails=set(external_user_group.user_emails),\n gives_anyone_access=external_user_group.gives_anyone_access,\n )\n" }, { "path": "backend/tests/daily/connectors/confluence/test_confluence_basic.py", "content": "import os\nimport time\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.connector import ConfluenceConnector\nfrom onyx.connectors.confluence.utils import AttachmentProcessingResult\nfrom onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\nfrom onyx.connectors.models import Document\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\ndef _make_connector(\n space: str, access_token: str, scoped_token: bool = False\n) -> ConfluenceConnector:\n connector = ConfluenceConnector(\n wiki_base=os.environ[\"CONFLUENCE_TEST_SPACE_URL\"],\n space=space,\n is_cloud=os.environ.get(\"CONFLUENCE_IS_CLOUD\", \"true\").lower() == \"true\",\n page_id=os.environ.get(\"CONFLUENCE_TEST_PAGE_ID\", \"\"),\n scoped_token=scoped_token,\n )\n\n credentials_provider = OnyxStaticCredentialsProvider(\n None,\n DocumentSource.CONFLUENCE,\n {\n \"confluence_username\": os.environ[\"CONFLUENCE_USER_NAME\"],\n \"confluence_access_token\": access_token,\n },\n )\n connector.set_credentials_provider(credentials_provider)\n return connector\n\n\n@pytest.fixture\ndef confluence_connector(space: str) -> ConfluenceConnector:\n return _make_connector(space, os.environ[\"CONFLUENCE_ACCESS_TOKEN\"].strip())\n\n\n@pytest.fixture\ndef confluence_connector_scoped(space: str) -> ConfluenceConnector:\n return _make_connector(\n space, os.environ[\"CONFLUENCE_ACCESS_TOKEN_SCOPED\"].strip(), scoped_token=True\n )\n\n\n@pytest.mark.parametrize(\"space\", [os.getenv(\"CONFLUENCE_TEST_SPACE\") or \"DailyConne\"])\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_confluence_connector_basic(\n mock_get_api_key: MagicMock, # noqa: ARG001\n confluence_connector: ConfluenceConnector,\n) -> None:\n _test_confluence_connector_basic(confluence_connector)\n\n\n@pytest.mark.parametrize(\"space\", [os.getenv(\"CONFLUENCE_TEST_SPACE\") or \"DailyConne\"])\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_confluence_connector_basic_scoped(\n mock_get_api_key: MagicMock, # noqa: ARG001\n confluence_connector_scoped: ConfluenceConnector,\n) -> None:\n _test_confluence_connector_basic(\n confluence_connector_scoped, expect_attachments=True\n )\n\n\ndef _test_confluence_connector_basic(\n confluence_connector: ConfluenceConnector, expect_attachments: bool = True\n) -> None:\n confluence_connector.set_allow_images(False)\n result = load_all_from_connector(confluence_connector, 0, time.time())\n doc_batch = result.documents\n hierarchy_nodes = result.hierarchy_nodes\n\n assert len(doc_batch) == (3 if expect_attachments else 2)\n\n # Hierarchy structure:\n # - Space \"DailyConne\" (root)\n # - Page \"DailyConnectorTestSpace Home\" (has attachments, so becomes hierarchy node)\n # - Attachment \"small-file.txt\"\n # - Page \"Page Within A Page\" (no children/attachments, not a hierarchy node)\n expected_hierarchy_count = 2 if expect_attachments else 1\n assert len(hierarchy_nodes) == expected_hierarchy_count, (\n f\"Expected {expected_hierarchy_count} hierarchy nodes but got {len(hierarchy_nodes)}. \"\n f\"Nodes: {[(n.raw_node_id, n.node_type, n.display_name) for n in hierarchy_nodes]}\"\n )\n\n # Verify hierarchy node structure\n space_node = next(\n (n for n in hierarchy_nodes if n.node_type.value == \"space\"), None\n )\n assert space_node is not None, \"Space hierarchy node not found\"\n assert space_node.raw_node_id == \"DailyConne\"\n assert space_node.display_name == \"DailyConnectorTestSpace\"\n assert space_node.raw_parent_id is None # Space is root\n\n if expect_attachments:\n home_page_node = next(\n (n for n in hierarchy_nodes if n.node_type.value == \"page\"), None\n )\n assert home_page_node is not None, \"Home page hierarchy node not found\"\n assert home_page_node.display_name == \"DailyConnectorTestSpace Home\"\n assert home_page_node.raw_parent_id == \"DailyConne\" # Parent is the space\n\n page_within_a_page_doc: Document | None = None\n page_doc: Document | None = None\n small_file_doc: Document | None = None\n\n for doc in doc_batch:\n if doc.semantic_identifier == \"DailyConnectorTestSpace Home\":\n page_doc = doc\n elif doc.semantic_identifier == \"Page Within A Page\":\n page_within_a_page_doc = doc\n elif doc.semantic_identifier == \"small-file.txt\":\n small_file_doc = doc\n else:\n print(f\"Unexpected doc: {doc.semantic_identifier}\")\n\n assert page_within_a_page_doc is not None\n assert page_within_a_page_doc.semantic_identifier == \"Page Within A Page\"\n assert page_within_a_page_doc.primary_owners\n assert page_within_a_page_doc.primary_owners[0].email == \"hagen@danswer.ai\"\n assert (\n page_within_a_page_doc.id\n == \"https://danswerai.atlassian.net/wiki/spaces/DailyConne/pages/200769540/Page+Within+A+Page\"\n )\n assert len(page_within_a_page_doc.sections) == 1\n\n page_within_a_page_section = page_within_a_page_doc.sections[0]\n page_within_a_page_text = \"@Chris Weaver loves cherry pie\"\n assert page_within_a_page_section.text == page_within_a_page_text\n assert (\n page_within_a_page_section.link\n == \"https://danswerai.atlassian.net/wiki/spaces/DailyConne/pages/200769540/Page+Within+A+Page\"\n )\n\n assert page_doc is not None\n assert page_doc.semantic_identifier == \"DailyConnectorTestSpace Home\"\n assert (\n page_doc.id == \"https://danswerai.atlassian.net/wiki/spaces/DailyConne/overview\"\n )\n assert page_doc.metadata[\"labels\"] == [\"testlabel\"]\n assert page_doc.primary_owners\n assert page_doc.primary_owners[0].email == \"hagen@danswer.ai\"\n assert (\n len(page_doc.sections) == 1\n ) # just page text, attachment text is separate doc\n\n page_section = page_doc.sections[0]\n assert (\n page_section.text\n == \"test123 \"\n + page_within_a_page_text\n + \"\\nsmall-file.txt\\nbig-file.txt\"\n )\n assert (\n page_section.link\n == \"https://danswerai.atlassian.net/wiki/spaces/DailyConne/overview\"\n )\n\n if expect_attachments:\n assert small_file_doc is not None\n text_attachment_section = small_file_doc.sections[0]\n assert text_attachment_section.text == \"small\"\n assert text_attachment_section.link\n assert text_attachment_section.link.split(\"?\")[0].endswith(\"small-file.txt\")\n\n\n@pytest.mark.parametrize(\"space\", [\"MI\"])\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_confluence_connector_skip_images(\n mock_get_api_key: MagicMock, # noqa: ARG001\n confluence_connector: ConfluenceConnector,\n) -> None:\n confluence_connector.set_allow_images(False)\n result = load_all_from_connector(confluence_connector, 0, time.time())\n doc_batch = result.documents\n hierarchy_nodes = result.hierarchy_nodes\n\n assert len(doc_batch) == 8\n assert sum(len(doc.sections) for doc in doc_batch) == 8\n\n # Hierarchy structure for MI space (when images are skipped):\n # - Space \"MI\" (Many Images)\n # - Page \"Many Images\" (home page, has children)\n # - Page \"Image formats\" (has children - the image pages)\n # Note: Image pages themselves don't become hierarchy nodes since images are skipped\n assert len(hierarchy_nodes) == 3, (\n f\"Expected 3 hierarchy nodes but got {len(hierarchy_nodes)}. \"\n f\"Nodes: {[(n.raw_node_id, n.node_type, n.display_name) for n in hierarchy_nodes]}\"\n )\n\n\ndef mock_process_image_attachment(\n *args: Any, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n) -> AttachmentProcessingResult:\n \"\"\"We need this mock to bypass DB access happening in the connector. Which shouldn't\n be done as a rule to begin with, but life is not perfect. Fix it later\"\"\"\n\n return AttachmentProcessingResult(\n text=\"Hi_text\",\n file_name=\"Hi_filename\",\n error=None,\n )\n\n\n@pytest.mark.parametrize(\"space\", [\"MI\"])\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\n@patch(\n \"onyx.connectors.confluence.utils._process_image_attachment\",\n side_effect=mock_process_image_attachment,\n)\ndef test_confluence_connector_allow_images(\n mock_get_api_key: MagicMock, # noqa: ARG001\n mock_process_image_attachment: MagicMock, # noqa: ARG001\n confluence_connector: ConfluenceConnector,\n) -> None:\n confluence_connector.set_allow_images(True)\n\n result = load_all_from_connector(confluence_connector, 0, time.time())\n doc_batch = result.documents\n hierarchy_nodes = result.hierarchy_nodes\n\n assert len(doc_batch) == 12\n assert sum(len(doc.sections) for doc in doc_batch) == 12\n\n # Hierarchy structure for MI space (when images are allowed):\n # - Space \"MI\" (Many Images)\n # - Page \"Many Images\" (home page)\n # - Page \"Image formats\" (has children)\n # - Page \"Dunder Mifflin Org Chart\" (has image attachments)\n # - Page \"List of Joey's Favorite Objects\" (has image attachments)\n # - Page \"Content\" (has image attachments)\n # Pages with image attachments become hierarchy nodes because attachments reference them\n assert len(hierarchy_nodes) == 6, (\n f\"Expected 6 hierarchy nodes but got {len(hierarchy_nodes)}. \"\n f\"Nodes: {[(n.raw_node_id, n.node_type, n.display_name) for n in hierarchy_nodes]}\"\n )\n" }, { "path": "backend/tests/daily/connectors/confluence/test_confluence_permissions_basic.py", "content": "import os\nimport time\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.confluence.connector import ConfluenceConnector\nfrom onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import SortOrder\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\n@pytest.fixture\ndef confluence_connector() -> ConfluenceConnector:\n connector = ConfluenceConnector(\n wiki_base=\"https://danswerai.atlassian.net\",\n is_cloud=True,\n )\n\n credentials_provider = OnyxStaticCredentialsProvider(\n None,\n DocumentSource.CONFLUENCE,\n {\n \"confluence_username\": os.environ[\"CONFLUENCE_USER_NAME\"],\n \"confluence_access_token\": os.environ[\"CONFLUENCE_ACCESS_TOKEN\"],\n },\n )\n connector.set_credentials_provider(credentials_provider)\n return connector\n\n\n# This should never fail because even if the docs in the cloud change,\n# the full doc ids retrieved should always be a subset of the slim doc ids\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_confluence_connector_permissions(\n mock_get_api_key: MagicMock, # noqa: ARG001\n confluence_connector: ConfluenceConnector,\n enable_ee: None, # noqa: ARG001\n) -> None:\n # Get all doc IDs from the full connector\n all_full_doc_ids = set()\n result = load_all_from_connector(confluence_connector, 0, time.time())\n doc_batch = result.documents\n hierarchy_nodes = result.hierarchy_nodes\n all_full_doc_ids.update([doc.id for doc in doc_batch])\n\n # Verify hierarchy nodes are returned and have valid structure\n # Note: The exact count depends on the current state of the Confluence instance\n assert len(hierarchy_nodes) > 0, \"Expected at least some hierarchy nodes\"\n\n # Verify all space nodes have no parent and all page nodes have a parent\n for node in hierarchy_nodes:\n if node.node_type.value == \"space\":\n assert (\n node.raw_parent_id is None\n ), f\"Space node {node.raw_node_id} should have no parent\"\n elif node.node_type.value == \"page\":\n assert (\n node.raw_parent_id is not None\n ), f\"Page node {node.raw_node_id} should have a parent\"\n\n # Get all doc IDs from the slim connector\n all_slim_doc_ids = set()\n for slim_doc_batch in confluence_connector.retrieve_all_slim_docs_perm_sync():\n all_slim_doc_ids.update(\n [doc.id for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # Find IDs that are in full but not in slim\n difference = all_full_doc_ids - all_slim_doc_ids\n\n # The set of full doc IDs should be always be a subset of the slim doc IDs\n assert all_full_doc_ids.issubset(\n all_slim_doc_ids\n ), f\"Full doc IDs are not a subset of slim doc IDs. Found {len(difference)} IDs in full docs but not in slim docs.\"\n\n\n@patch(\"ee.onyx.external_permissions.confluence.doc_sync.OnyxDBCredentialsProvider\")\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_confluence_connector_restriction_handling(\n mock_get_api_key: MagicMock, # noqa: ARG001\n mock_db_provider_class: MagicMock,\n enable_ee: None, # noqa: ARG001\n) -> None:\n # Test space key\n test_space_key = \"DailyPermS\"\n\n # Configure the mock provider instance that will be returned\n mock_provider_instance = MagicMock()\n mock_provider_instance.get_credentials.return_value = {\n \"confluence_username\": os.environ[\"CONFLUENCE_USER_NAME\"],\n \"confluence_access_token\": os.environ[\"CONFLUENCE_ACCESS_TOKEN\"],\n }\n # this prevents redis calls inside of OnyxConfluence\n mock_provider_instance.is_dynamic.return_value = False\n # Make the class return our configured instance when called\n mock_db_provider_class.return_value = mock_provider_instance\n\n # Mock the cc_pair to pass to the function\n mock_cc_pair = MagicMock(spec=ConnectorCredentialPair)\n # Mock the nested connector attribute and its config\n mock_cc_pair.connector = MagicMock()\n mock_cc_pair.connector.connector_specific_config = {\n \"wiki_base\": \"https://danswerai.atlassian.net\",\n \"is_cloud\": True,\n \"space\": test_space_key,\n }\n # Set a mock credential ID\n mock_cc_pair.credential_id = 1\n\n # Call the confluence_doc_sync function directly with the mock cc_pair\n def mock_fetch_all_docs_fn(\n sort_order: SortOrder | None = None, # noqa: ARG001\n ) -> list[DocumentRow]:\n return []\n\n def mock_fetch_all_docs_ids_fn() -> list[str]:\n return []\n\n doc_access_generator = confluence_doc_sync(\n mock_cc_pair, mock_fetch_all_docs_fn, mock_fetch_all_docs_ids_fn, None\n )\n doc_access_list = list(doc_access_generator)\n assert len(doc_access_list) == 7\n assert all(\n not doc_access.external_access.is_public for doc_access in doc_access_list\n )\n\n # if no restriction is applied, the groups should give access, so no need\n # for more emails outside of the owner\n non_restricted_emails = {\"chris@onyx.app\"}\n non_restricted_user_groups = {\n \"confluence-admins-danswerai\",\n \"org-admins\",\n \"atlassian-addons-admin\",\n \"confluence-users-danswerai\",\n }\n\n # if restriction is applied, only should be visible to shared users / groups\n restricted_emails = {\"chris@onyx.app\", \"hagen@danswer.ai\", \"oauth@onyx.app\"}\n restricted_user_groups = {\"confluence-admins-danswerai\"}\n\n extra_restricted_emails = {\"chris@onyx.app\", \"oauth@onyx.app\"}\n extra_restricted_user_groups: set[str] = set()\n\n # note that this is only allowed since yuhong@onyx.app is a member of the\n # confluence-admins-danswerai group\n special_restricted_emails = {\"chris@onyx.app\", \"yuhong@onyx.app\", \"oauth@onyx.app\"}\n special_restricted_user_groups: set[str] = set()\n\n # Check Root+Page+2 is public\n root_page_2 = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.endswith(\"Root+Page+2\")\n )\n assert root_page_2.external_access.external_user_emails == non_restricted_emails\n assert (\n root_page_2.external_access.external_user_group_ids\n == non_restricted_user_groups\n )\n\n # Check Overview page is public\n overview_page = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.lower().endswith(\"overview\")\n )\n assert (\n overview_page.external_access.external_user_emails == non_restricted_emails\n ), \"Overview page emails do not match expected values\"\n assert (\n overview_page.external_access.external_user_group_ids\n == non_restricted_user_groups\n ), \"Overview page groups do not match expected values\"\n\n # check root page is restricted\n root_page = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.endswith(\"Root+Page\")\n )\n assert (\n root_page.external_access.external_user_emails == restricted_emails\n ), \"Root page emails do not match expected values\"\n assert (\n root_page.external_access.external_user_group_ids == restricted_user_groups\n ), \"Root page groups do not match expected values\"\n\n # check child page has restriction propagated\n child_page = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.endswith(\"Child+Page\")\n )\n assert (\n child_page.external_access.external_user_emails == restricted_emails\n ), \"Child page emails do not match expected values\"\n assert (\n child_page.external_access.external_user_group_ids == restricted_user_groups\n ), \"Child page groups do not match expected values\"\n\n # check doubly nested child page has restriction propagated\n child_page_2 = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.endswith(\"Child+Page+2\")\n )\n assert (\n child_page_2.external_access.external_user_emails == restricted_emails\n ), \"Child page 2 emails do not match expected values\"\n assert (\n child_page_2.external_access.external_user_group_ids == restricted_user_groups\n ), \"Child page 2 groups do not match expected values\"\n\n # check child page w/ specific restrictions have those applied\n child_page_3 = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.endswith(\"Child+Page+3\")\n )\n assert (\n child_page_3.external_access.external_user_emails == extra_restricted_emails\n ), \"Child page 3 emails do not match expected values\"\n assert (\n child_page_3.external_access.external_user_group_ids\n == extra_restricted_user_groups\n ), \"Child page 3 groups do not match expected values\"\n\n # check child page w/ specific restrictions have those applied\n child_page_4 = next(\n d\n for d in doc_access_list\n if isinstance(d, DocExternalAccess) and d.doc_id.endswith(\"Child+Page+4\")\n )\n assert (\n child_page_4.external_access.external_user_emails == special_restricted_emails\n ), \"Child page 4 emails do not match expected values\"\n assert (\n child_page_4.external_access.external_user_group_ids\n == special_restricted_user_groups\n ), \"Child page 4 groups do not match expected values\"\n" }, { "path": "backend/tests/daily/connectors/confluence/test_confluence_user_email_overrides.py", "content": "import types\nfrom unittest.mock import patch\n\nfrom onyx.connectors.confluence.onyx_confluence import ConfluenceUser\nfrom onyx.connectors.confluence.onyx_confluence import OnyxConfluence\nfrom onyx.connectors.interfaces import CredentialsProviderInterface\n\n\nclass MockCredentialsProvider(CredentialsProviderInterface):\n def get_tenant_id(self) -> str:\n return \"test_tenant\"\n\n def get_provider_key(self) -> str:\n return \"test_provider\"\n\n def is_dynamic(self) -> bool:\n return False\n\n def get_credentials(self) -> dict[str, str]:\n return {\"confluence_access_token\": \"test_token\"}\n\n def set_credentials(self, credentials: dict[str, str]) -> None:\n pass\n\n def __enter__(self) -> \"MockCredentialsProvider\":\n return self\n\n def __exit__(\n self,\n exc_type: type[BaseException] | None,\n exc_val: BaseException | None,\n exc_tb: types.TracebackType | None,\n ) -> None:\n pass\n\n\ndef test_paginated_cql_user_retrieval_with_overrides() -> None:\n \"\"\"\n Tests that paginated_cql_user_retrieval yields users from the overrides\n when provided and is_cloud is False.\n \"\"\"\n mock_provider = MockCredentialsProvider()\n overrides = [\n {\n \"user_id\": \"override_user_1\",\n \"username\": \"override1\",\n \"display_name\": \"Override User One\",\n \"email\": \"override1@example.com\",\n \"type\": \"override\",\n },\n {\n \"user_id\": \"override_user_2\",\n \"username\": \"override2\",\n \"display_name\": \"Override User Two\",\n \"email\": \"override2@example.com\",\n \"type\": \"override\",\n },\n ]\n expected_users = [ConfluenceUser(**user_data) for user_data in overrides]\n\n confluence_client = OnyxConfluence(\n is_cloud=False, # Overrides are primarily for Server/DC\n url=\"http://dummy-confluence.com\",\n credentials_provider=mock_provider,\n confluence_user_profiles_override=overrides,\n )\n\n retrieved_users = list(confluence_client.paginated_cql_user_retrieval())\n\n assert len(retrieved_users) == len(expected_users)\n # Sort lists by user_id for order-independent comparison\n retrieved_users.sort(key=lambda u: u.user_id)\n expected_users.sort(key=lambda u: u.user_id)\n assert retrieved_users == expected_users\n\n\ndef test_paginated_cql_user_retrieval_no_overrides_server() -> None:\n \"\"\"\n Tests that paginated_cql_user_retrieval attempts to call the actual\n API pagination when no overrides are provided for Server/DC.\n \"\"\"\n mock_provider = MockCredentialsProvider()\n confluence_client = OnyxConfluence(\n is_cloud=False,\n url=\"http://dummy-confluence.com\",\n credentials_provider=mock_provider,\n confluence_user_profiles_override=None,\n )\n\n # Mock the internal pagination method to check if it's called\n with patch.object(confluence_client, \"_paginate_url\") as mock_paginate:\n mock_paginate.return_value = iter([]) # Return an empty iterator\n\n list(confluence_client.paginated_cql_user_retrieval())\n\n mock_paginate.assert_called_once_with(\"rest/api/user/list\", None)\n\n\ndef test_paginated_cql_user_retrieval_no_overrides_cloud() -> None:\n \"\"\"\n Tests that paginated_cql_user_retrieval attempts to call the actual\n API pagination when no overrides are provided for Cloud.\n \"\"\"\n mock_provider = MockCredentialsProvider()\n confluence_client = OnyxConfluence(\n is_cloud=True,\n url=\"http://dummy-confluence.com\", # URL doesn't matter much here due to mocking\n credentials_provider=mock_provider,\n confluence_user_profiles_override=None,\n )\n\n # Mock the internal pagination method to check if it's called\n with patch.object(confluence_client, \"_paginate_url\") as mock_paginate:\n mock_paginate.return_value = iter([]) # Return an empty iterator\n\n list(confluence_client.paginated_cql_user_retrieval())\n\n # Check that the cloud-specific user search URL is called\n mock_paginate.assert_called_once_with(\n \"rest/api/search/user?cql=type=user\",\n None,\n force_offset_pagination=True,\n )\n" }, { "path": "backend/tests/daily/connectors/conftest.py", "content": "from collections.abc import Generator\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\n\n@pytest.fixture\ndef mock_get_unstructured_api_key() -> Generator[MagicMock, None, None]:\n with patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n ) as mock:\n yield mock\n" }, { "path": "backend/tests/daily/connectors/discord/test_discord_connector.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.connectors.discord.connector import DiscordConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentSource\nfrom onyx.connectors.models import HierarchyNode\n\n\n@pytest.fixture\ndef discord_connector() -> DiscordConnector:\n connector = DiscordConnector()\n connector.load_credentials(\n {\"discord_bot_token\": os.environ[\"DISCORD_CONNECTOR_BOT_TOKEN\"]}\n )\n return connector\n\n\ndef test_discord_connector_basic(discord_connector: DiscordConnector) -> None:\n # If there are no Discord messages in the last 7 days, something has gone horribly wrong\n end_time = time.time()\n start_time = end_time - (7 * 24 * 60 * 60)\n doc_batch_generator = discord_connector.poll_source(start_time, end_time)\n\n doc_batch = next(doc_batch_generator)\n\n docs: list[Document] = []\n for doc in doc_batch:\n if not isinstance(doc, HierarchyNode):\n docs.append(doc)\n\n assert len(docs) > 0, \"No documents were retrieved from the connector\"\n\n # Check basic document structure\n doc = docs[0]\n assert doc.source == DocumentSource.DISCORD\n assert doc.id is not None\n assert doc.semantic_identifier is not None\n assert len(doc.sections) > 0\n assert doc.sections[0].text is not None\n assert doc.sections[0].link is not None\n" }, { "path": "backend/tests/daily/connectors/file/test_file_connector.py", "content": "import io\nfrom datetime import datetime\nfrom datetime import timezone\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\n\nfrom onyx.connectors.file.connector import LocalFileConnector\nfrom onyx.connectors.models import HierarchyNode\n\n\n@pytest.fixture\ndef mock_db_session() -> MagicMock:\n return MagicMock()\n\n\n@pytest.fixture\ndef mock_file_store() -> MagicMock:\n store = MagicMock()\n return store\n\n\n@pytest.fixture\ndef mock_filestore_record() -> MagicMock:\n record = MagicMock()\n record.file_id = uuid4()\n record.display_name = \"test.txt\"\n return record\n\n\n@patch(\"onyx.connectors.file.connector.get_default_file_store\")\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\", return_value=None\n)\ndef test_single_text_file_with_metadata(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_get_session: MagicMock,\n mock_db_session: MagicMock,\n mock_file_store: MagicMock,\n mock_filestore_record: MagicMock,\n) -> None:\n file_content = io.BytesIO(\n b'#ONYX_METADATA={\"link\": \"https://onyx.app\", \"file_display_name\":\"my display name\", \"tag_of_your_choice\": \"test-tag\", \\\n \"primary_owners\": [\"wenxi@onyx.app\"], \"secondary_owners\": [\"founders@onyx.app\"], \\\n \"doc_updated_at\": \"2001-01-01T00:00:00Z\"}\\n'\n b\"Test answer is 12345\"\n )\n mock_get_filestore = MagicMock()\n mock_get_filestore.return_value = mock_file_store\n mock_file_store.read_file_record.return_value = mock_filestore_record\n mock_get_session.return_value.__enter__.return_value = mock_db_session\n mock_file_store.read_file.return_value = file_content\n\n with patch(\n \"onyx.connectors.file.connector.get_default_file_store\",\n return_value=mock_file_store,\n ):\n connector = LocalFileConnector(\n file_locations=[\"test.txt\"], file_names=[\"test.txt\"], zip_metadata={}\n )\n batches = list(connector.load_from_state())\n\n assert len(batches) == 1\n docs = batches[0]\n assert len(docs) == 1\n doc = docs[0]\n assert not isinstance(doc, HierarchyNode)\n\n assert doc.sections[0].text == \"Test answer is 12345\"\n assert doc.sections[0].link == \"https://onyx.app\"\n assert doc.semantic_identifier == \"my display name\"\n assert doc.primary_owners[0].display_name == \"wenxi@onyx.app\" # type: ignore\n assert doc.secondary_owners[0].display_name == \"founders@onyx.app\" # type: ignore\n assert doc.doc_updated_at == datetime(2001, 1, 1, 0, 0, 0, tzinfo=timezone.utc)\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\", return_value=None\n)\ndef test_two_text_files_with_zip_metadata(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_db_session: MagicMock, # noqa: ARG001\n mock_file_store: MagicMock,\n) -> None:\n file1_content = io.BytesIO(b\"File 1 content\")\n file2_content = io.BytesIO(b\"File 2 content\")\n mock_get_filestore = MagicMock()\n mock_get_filestore.return_value = mock_file_store\n mock_file_store.read_file_record.side_effect = [\n MagicMock(file_id=str(uuid4()), display_name=\"file1.txt\"),\n MagicMock(file_id=str(uuid4()), display_name=\"file2.txt\"),\n ]\n mock_file_store.read_file.side_effect = [file1_content, file2_content]\n zip_metadata = {\n \"file1.txt\": {\n \"filename\": \"file1.txt\",\n \"file_display_name\": \"display 1\",\n \"link\": \"https://onyx.app/1\",\n \"primary_owners\": [\"alice@onyx.app\"],\n \"secondary_owners\": [\"bob@onyx.app\"],\n \"doc_updated_at\": \"2022-02-02T00:00:00Z\",\n },\n \"file2.txt\": {\n \"filename\": \"file2.txt\",\n \"file_display_name\": \"display 2\",\n \"link\": \"https://onyx.app/2\",\n \"primary_owners\": [\"carol@onyx.app\"],\n \"secondary_owners\": [\"dave@onyx.app\"],\n \"doc_updated_at\": \"2023-03-03T00:00:00Z\",\n },\n }\n\n with patch(\n \"onyx.connectors.file.connector.get_default_file_store\",\n return_value=mock_file_store,\n ):\n connector = LocalFileConnector(\n file_locations=[\"file1.txt\", \"file2.txt\"],\n file_names=[\"file1.txt\", \"file2.txt\"],\n zip_metadata=zip_metadata,\n )\n batches = list(connector.load_from_state())\n\n assert len(batches) == 1\n docs = batches[0]\n assert len(docs) == 2\n doc1, doc2 = docs\n assert not isinstance(doc1, HierarchyNode)\n assert not isinstance(doc2, HierarchyNode)\n\n assert doc1.sections[0].text == \"File 1 content\"\n assert doc1.sections[0].link == \"https://onyx.app/1\"\n assert doc1.semantic_identifier == \"display 1\"\n assert doc1.primary_owners[0].display_name == \"alice@onyx.app\" # type: ignore\n assert doc1.secondary_owners[0].display_name == \"bob@onyx.app\" # type: ignore\n assert doc1.doc_updated_at == datetime(2022, 2, 2, 0, 0, 0, tzinfo=timezone.utc)\n assert doc2.sections[0].text == \"File 2 content\"\n assert doc2.sections[0].link == \"https://onyx.app/2\"\n assert doc2.semantic_identifier == \"display 2\"\n assert doc2.primary_owners[0].display_name == \"carol@onyx.app\" # type: ignore\n assert doc2.secondary_owners[0].display_name == \"dave@onyx.app\" # type: ignore\n assert doc2.doc_updated_at == datetime(2023, 3, 3, 0, 0, 0, tzinfo=timezone.utc)\n" }, { "path": "backend/tests/daily/connectors/fireflies/test_fireflies_connector.py", "content": "import json\nimport os\nimport time\nfrom pathlib import Path\nfrom typing import Any\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.fireflies.connector import FirefliesConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\n\n\ndef load_test_data(file_name: str = \"test_fireflies_data.json\") -> dict[str, Any]:\n current_dir = Path(__file__).parent\n with open(current_dir / file_name, \"r\") as f:\n return json.load(f)\n\n\n@pytest.fixture\ndef fireflies_connector() -> FirefliesConnector:\n connector = FirefliesConnector()\n connector.load_credentials(\n {\"fireflies_api_key\": os.environ[\"FIREFLIES_API_KEY\"]},\n )\n return connector\n\n\n@pytest.mark.xfail(\n reason=\"We don't have the key that is stored in GitHub Secrets and the returned data is different than expected\",\n)\ndef test_fireflies_connector_basic(fireflies_connector: FirefliesConnector) -> None:\n test_data = load_test_data()\n\n connector_return_data: list[Document | HierarchyNode] = next(\n fireflies_connector.poll_source(0, time.time())\n )\n target_doc: Document | HierarchyNode = connector_return_data[0]\n if isinstance(target_doc, HierarchyNode):\n raise ValueError(\"Hierarchy node returned from connector\")\n\n assert target_doc is not None, \"No documents were retrieved from the connector\"\n assert (\n target_doc.primary_owners is not None\n ), \"No primary owners were retrieved from the connector\"\n\n assert target_doc.id == test_data[\"id\"]\n assert target_doc.semantic_identifier == test_data[\"semantic_identifier\"]\n assert target_doc.primary_owners[0].email == test_data[\"primary_owners\"]\n assert target_doc.secondary_owners == test_data[\"secondary_owners\"]\n assert str(target_doc.doc_updated_at) == test_data[\"doc_updated_at\"]\n\n assert (\n target_doc.source == DocumentSource.FIREFLIES\n ), \"Document source is not fireflies\"\n assert target_doc.metadata == test_data[\"metadata\"]\n\n # Check that the test data and the connector data contain the same section data\n assert {section.text for section in target_doc.sections} == {\n section[\"text\"] for section in test_data[\"sections\"]\n }\n assert {section.link for section in target_doc.sections} == {\n section[\"link\"] for section in test_data[\"sections\"]\n }\n" }, { "path": "backend/tests/daily/connectors/fireflies/test_fireflies_data.json", "content": "{\n \"id\": \"FIREFLIES_VcBdZpuV82rImQCA\",\n \"semantic_identifier\": \"Lead Generation Efforts\",\n \"primary_owners\": \"admin@onyx-test.com\",\n \"secondary_owners\": [],\n \"doc_updated_at\": \"2025-01-10 19:10:00+00:00\",\n \"metadata\": {\n \"meeting_date\": \"2025-01-10 19:10:00+00:00\",\n \"duration_min\": \"10\"\n },\n \"sections\": [\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=153.1\",\n \"text\": \"test_user_1 1: Hey, David, thanks for taking the time today.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=158.14\",\n \"text\": \"Test Admin Admin: Of course Sarah, It's nice to see you. Whenever you're ready.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=165.1\",\n \"text\": \"test_user_1 1: All right then, David, let's jump right in. How are the lead generation efforts for the new product launch looking?\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=171.084\",\n \"text\": \"Test Admin Admin: So far we've seen a good initial response, but we're facing a slight challenge with qualifying leads. The sales team is getting inquiries. Some aren't quite aligned with our ideal customer profile.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=191.86\",\n \"text\": \"test_user_1 1: That makes sense. Do you think we need to adjust our marketing messaging to better target the right audience?\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=202.26\",\n \"text\": \"Test Admin Admin: Absolutely. Maybe we could emphasize the key features that are most relevant to our target market in the marketing materials. What are your thoughts on refining the lead capture to gather more specific information?\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=225.99\",\n \"text\": \"test_user_1 1: I think that's a great idea. We could add additional qualifying questions to ensure we're capturing leads with the right needs.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=238.56\",\n \"text\": \"Test Admin Admin: On another note, how are the social media campaigns performing? Are we seeing good engagement with the new product launch post?\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=257.2\",\n \"text\": \"test_user_1 1: The engagement is positive, but we could potentially increase increase reach further with targeted ad campaigns and key platforms.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=268.91\",\n \"text\": \"Test Admin Admin: Agreed. Let's discuss a strategy to develop targeted ads that focus on the pain points our ideal customers are facing and how our product solves them.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=270.27\",\n \"text\": \"test_user_1 1: We can collaborate on creating specific ad copy that highlights these benefits.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=289.06\",\n \"text\": \"Test Admin Admin: All right, so to summarize, let's prioritize refining the lead capture form, develop targeted social media ads, and make sure our marketing method clearly aligns with our ideal customer profile.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=303.38\",\n \"text\": \"test_user_1 1: Yep. And let's schedule a follow up meeting in a week, review progress and discuss any adjustments.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=310.9\",\n \"text\": \"Test Admin Admin: Sounds good. I'll send you address updated lead form by the end of the day. Thanks, Sarah.\"\n },\n {\n \"link\": \"https://app.fireflies.ai/view/VcBdZpuV82rImQCA?t=319.19\",\n \"text\": \"test_user_1 1: Thank you David.\"\n }\n ]\n}\n" }, { "path": "backend/tests/daily/connectors/gitbook/test_gitbook_connector.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.gitbook.connector import GitbookConnector\nfrom onyx.connectors.models import HierarchyNode\n\n\n@pytest.fixture\ndef gitbook_connector() -> GitbookConnector:\n connector = GitbookConnector(\n space_id=os.environ[\"GITBOOK_SPACE_ID\"],\n )\n connector.load_credentials(\n {\n \"gitbook_api_key\": os.environ[\"GITBOOK_API_KEY\"],\n }\n )\n return connector\n\n\nNUM_PAGES = 3\n\n\ndef test_gitbook_connector_basic(gitbook_connector: GitbookConnector) -> None:\n doc_batch_generator = gitbook_connector.load_from_state()\n\n # Get first batch of documents\n doc_batch = next(doc_batch_generator)\n assert len(doc_batch) == NUM_PAGES\n\n # Verify first document structure\n main_doc = doc_batch[0]\n assert not isinstance(main_doc, HierarchyNode)\n\n # Basic document properties\n assert main_doc.id.startswith(\"gitbook-\")\n assert main_doc.semantic_identifier == \"Acme Corp Internal Handbook\"\n assert main_doc.source == DocumentSource.GITBOOK\n\n # Metadata checks\n assert \"path\" in main_doc.metadata\n assert \"type\" in main_doc.metadata\n assert \"kind\" in main_doc.metadata\n\n # Section checks\n assert len(main_doc.sections) == 1\n section = main_doc.sections[0]\n\n # Content specific checks\n content = section.text\n assert content is not None, \"Section text should not be None\"\n\n # Check for specific content elements\n assert \"* Fruit Shopping List:\" in content\n assert \"> test quote it doesn't mean anything\" in content\n\n # Check headings\n assert \"# Heading 1\" in content\n assert \"## Heading 2\" in content\n assert \"### Heading 3\" in content\n\n # Check task list\n assert \"- [ ] Uncompleted Task\" in content\n assert \"- [x] Completed Task\" in content\n\n # Check table content\n assert \"| ethereum | 10 | 3000 |\" in content\n assert \"| bitcoin | 2 | 98000 |\" in content\n\n # Check paragraph content\n assert \"New York City comprises 5 boroughs\" in content\n assert \"Empire State Building\" in content\n\n # Check code block (just verify presence of some unique code elements)\n assert \"function fizzBuzz(n)\" in content\n assert 'res.push(\"FizzBuzz\")' in content\n\n assert section.link # Should have a URL\n\n nested1 = doc_batch[1]\n assert not isinstance(nested1, HierarchyNode)\n assert nested1.id.startswith(\"gitbook-\")\n assert nested1.semantic_identifier == \"Nested1\"\n assert len(nested1.sections) == 1\n # extra newlines at the end, remove them to make test easier\n assert nested1.sections[0].text is not None\n assert nested1.sections[0].text.strip() == \"nested1\"\n assert nested1.source == DocumentSource.GITBOOK\n\n nested2 = doc_batch[2]\n assert not isinstance(nested2, HierarchyNode)\n assert nested2.id.startswith(\"gitbook-\")\n assert nested2.semantic_identifier == \"Nested2\"\n assert len(nested2.sections) == 1\n assert nested2.sections[0].text is not None\n assert nested2.sections[0].text.strip() == \"nested2\"\n assert nested2.source == DocumentSource.GITBOOK\n\n # Time-based polling test\n current_time = time.time()\n poll_docs = gitbook_connector.poll_source(0, current_time)\n poll_batch = next(poll_docs)\n assert len(poll_batch) == NUM_PAGES\n" }, { "path": "backend/tests/daily/connectors/github/test_github_basic.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.github.connector import GithubConnector\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\n@pytest.fixture\ndef github_connector() -> GithubConnector:\n connector = GithubConnector(\n repo_owner=\"onyx-dot-app\",\n repositories=\"documentation\",\n include_prs=True,\n include_issues=True,\n )\n connector.load_credentials(\n {\n \"github_access_token\": os.environ[\"ACCESS_TOKEN_GITHUB\"],\n }\n )\n return connector\n\n\ndef test_github_connector_basic(github_connector: GithubConnector) -> None:\n docs = load_all_from_connector(\n connector=github_connector,\n start=0,\n end=time.time(),\n ).documents\n assert len(docs) > 1 # We expect at least one PR and one Issue to exist\n\n # Test the first document's structure\n pr_doc = docs[0]\n issue_doc = docs[-1]\n\n # Verify basic document properties\n assert pr_doc.source == DocumentSource.GITHUB\n assert pr_doc.secondary_owners is None\n assert pr_doc.from_ingestion_api is False\n assert pr_doc.additional_info is None\n\n # Verify GitHub-specific properties\n assert \"github.com\" in pr_doc.id # Should be a GitHub URL\n\n # Verify PR-specific properties\n assert pr_doc.metadata is not None\n assert pr_doc.metadata.get(\"object_type\") == \"PullRequest\"\n assert \"id\" in pr_doc.metadata\n assert \"merged\" in pr_doc.metadata\n assert \"state\" in pr_doc.metadata\n assert \"user\" in pr_doc.metadata\n assert \"assignees\" in pr_doc.metadata\n assert pr_doc.metadata.get(\"repo\") == \"onyx-dot-app/documentation\"\n assert \"num_commits\" in pr_doc.metadata\n assert \"num_files_changed\" in pr_doc.metadata\n assert \"labels\" in pr_doc.metadata\n assert \"created_at\" in pr_doc.metadata\n\n # Verify Issue-specific properties\n assert issue_doc.metadata is not None\n assert issue_doc.metadata.get(\"object_type\") == \"Issue\"\n assert \"id\" in issue_doc.metadata\n assert \"state\" in issue_doc.metadata\n assert \"user\" in issue_doc.metadata\n assert \"assignees\" in issue_doc.metadata\n assert issue_doc.metadata.get(\"repo\") == \"onyx-dot-app/documentation\"\n assert \"labels\" in issue_doc.metadata\n assert \"created_at\" in issue_doc.metadata\n\n # Verify sections\n assert len(pr_doc.sections) == 1\n section = pr_doc.sections[0]\n assert section.link == pr_doc.id # Section link should match document ID\n assert isinstance(section.text, str) # Should have some text content\n" }, { "path": "backend/tests/daily/connectors/gitlab/test_gitlab_basic.py", "content": "import itertools\nimport os\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.gitlab.connector import GitlabConnector\nfrom onyx.connectors.models import HierarchyNode\n\n\n@pytest.fixture\ndef gitlab_connector() -> GitlabConnector:\n connector = GitlabConnector(\n project_owner=\"onyx2895818\",\n project_name=\"onyx\",\n include_mrs=True,\n include_issues=True,\n include_code_files=True, # Include code files in the test\n )\n # Ensure GITLAB_ACCESS_TOKEN and optionally GITLAB_URL are set in the environment\n gitlab_url = os.environ.get(\"GITLAB_URL\", \"https://gitlab.com\")\n gitlab_token = os.environ.get(\"GITLAB_ACCESS_TOKEN\")\n\n if not gitlab_token:\n pytest.skip(\"GITLAB_ACCESS_TOKEN environment variable not set.\")\n\n connector.load_credentials(\n {\n \"gitlab_access_token\": gitlab_token,\n \"gitlab_url\": gitlab_url,\n }\n )\n return connector\n\n\ndef test_gitlab_connector_basic(gitlab_connector: GitlabConnector) -> None:\n doc_batches = gitlab_connector.load_from_state()\n docs = list(itertools.chain(*doc_batches))\n # Assert right number of docs - Adjust if necessary based on test repo state\n assert len(docs) == 79\n\n # Find one of each type to validate\n validated_mr = False\n validated_issue = False\n validated_code_file = False\n gitlab_base_url = os.environ.get(\"GITLAB_URL\", \"https://gitlab.com\").split(\"//\")[-1]\n project_path = f\"{gitlab_connector.project_owner}/{gitlab_connector.project_name}\"\n\n # --- Specific Document Details to Validate ---\n target_mr_id = f\"https://{gitlab_base_url}/{project_path}/-/merge_requests/1\"\n target_issue_id = f\"https://{gitlab_base_url}/{project_path}/-/work_items/2\"\n target_code_file_semantic_id = \"README.md\"\n # ---\n\n for doc in docs:\n if isinstance(doc, HierarchyNode):\n continue\n # Verify basic document properties (common to all types)\n assert doc.source == DocumentSource.GITLAB\n assert doc.secondary_owners is None\n assert doc.from_ingestion_api is False\n assert doc.additional_info is None\n assert isinstance(doc.id, str)\n assert doc.metadata is not None\n assert \"type\" in doc.metadata\n doc_type = doc.metadata[\"type\"]\n\n # Verify sections (common structure)\n assert len(doc.sections) >= 1\n section = doc.sections[0]\n assert isinstance(section.link, str)\n assert gitlab_base_url in section.link\n assert isinstance(section.text, str)\n\n # --- Type-specific and Content Validation ---\n if doc.id == target_mr_id and doc_type == \"MergeRequest\":\n assert doc.metadata[\"state\"] == \"opened\"\n assert doc.semantic_identifier == \"Add awesome feature\"\n assert section.text == \"This MR implements the awesome feature\"\n assert doc.primary_owners is not None\n assert len(doc.primary_owners) == 1\n assert (\n doc.primary_owners[0].display_name == \"Test\"\n ) # Adjust if author changes\n assert doc.id == section.link\n validated_mr = True\n elif doc.id == target_issue_id and doc_type == \"ISSUE\":\n assert doc.metadata[\"state\"] == \"opened\"\n assert doc.semantic_identifier == \"Investigate performance issue\"\n assert (\n section.text\n == \"Investigate and resolve the performance degradation on endpoint X\"\n )\n assert doc.primary_owners is not None\n assert len(doc.primary_owners) == 1\n assert (\n doc.primary_owners[0].display_name == \"Test\"\n ) # Adjust if author changes\n assert doc.id == section.link\n validated_issue = True\n elif (\n doc.semantic_identifier == target_code_file_semantic_id\n and doc_type == \"CodeFile\"\n ):\n # ID is a git hash (e.g., 'd177...'), Link is the blob URL\n assert doc.id != section.link\n assert section.link.endswith(\"/README.md\")\n assert \"# onyx\" in section.text # Check for a known part of the content\n # Code files might not have primary owners assigned this way\n # assert len(doc.primary_owners) == 0\n validated_code_file = True\n\n # Generic validation for *any* document of the type if specific one not found yet\n elif doc_type == \"MergeRequest\" and not validated_mr:\n assert \"state\" in doc.metadata\n assert gitlab_base_url in doc.id # MR ID should be a URL\n assert doc.id == section.link # Link and ID are the same URL\n elif doc_type == \"ISSUE\" and not validated_issue:\n assert \"state\" in doc.metadata\n assert gitlab_base_url in doc.id # Issue ID should be a URL\n assert doc.id == section.link # Link and ID are the same URL\n elif doc_type == \"CodeFile\" and not validated_code_file:\n assert doc.id != section.link # ID is GID/hash, link is blob URL\n\n # Early exit optimization (optional)\n # if validated_mr and validated_issue and validated_code_file:\n # break\n\n # Assert that we found and validated the specific documents\n assert (\n validated_mr\n ), f\"Failed to find and validate the specific MergeRequest ({target_mr_id}).\"\n assert (\n validated_issue\n ), f\"Failed to find and validate the specific Issue ({target_issue_id}).\"\n assert (\n validated_code_file\n ), f\"Failed to find and validate the specific CodeFile ({target_code_file_semantic_id}).\"\n" }, { "path": "backend/tests/daily/connectors/gmail/conftest.py", "content": "import json\nimport os\nfrom collections.abc import Callable\n\nimport pytest\n\nfrom onyx.connectors.gmail.connector import GmailConnector\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_AUTHENTICATION_METHOD,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_TOKEN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GoogleOAuthAuthenticationMethod,\n)\nfrom tests.load_env_vars import load_env_vars\n\n\n# Load environment variables at the module level\nload_env_vars()\n\n\ndef parse_credentials(env_str: str) -> dict:\n \"\"\"\n Parse a double-escaped JSON string from environment variables into a Python dictionary.\n\n Args:\n env_str (str): The double-escaped JSON string from environment variables\n\n Returns:\n dict: Parsed OAuth credentials\n \"\"\"\n # first try normally\n try:\n return json.loads(env_str)\n except Exception:\n # First, try remove extra escaping backslashes\n unescaped = env_str.replace('\\\\\"', '\"')\n\n # remove leading / trailing quotes\n unescaped = unescaped.strip('\"')\n\n # Now parse the JSON\n return json.loads(unescaped)\n\n\n@pytest.fixture\ndef google_gmail_oauth_connector_factory() -> Callable[..., GmailConnector]:\n def _connector_factory(\n primary_admin_email: str = \"admin@onyx-test.com\",\n ) -> GmailConnector:\n print(\"Creating GmailConnector with OAuth credentials\")\n connector = GmailConnector()\n\n json_string = os.environ[\"GOOGLE_GMAIL_OAUTH_CREDENTIALS_JSON_STR\"]\n refried_json_string = json.dumps(parse_credentials(json_string))\n\n credentials_json = {\n DB_CREDENTIALS_DICT_TOKEN_KEY: refried_json_string,\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: primary_admin_email,\n DB_CREDENTIALS_AUTHENTICATION_METHOD: GoogleOAuthAuthenticationMethod.UPLOADED.value,\n }\n connector.load_credentials(credentials_json)\n return connector\n\n return _connector_factory\n\n\n@pytest.fixture\ndef google_gmail_service_acct_connector_factory() -> Callable[..., GmailConnector]:\n def _connector_factory(\n primary_admin_email: str = \"admin@onyx-test.com\",\n ) -> GmailConnector:\n print(\"Creating GmailConnector with service account credentials\")\n connector = GmailConnector()\n\n json_string = os.environ[\"GOOGLE_GMAIL_SERVICE_ACCOUNT_JSON_STR\"]\n refried_json_string = json.dumps(parse_credentials(json_string))\n\n # Load Service Account Credentials\n connector.load_credentials(\n {\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: refried_json_string,\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: primary_admin_email,\n DB_CREDENTIALS_AUTHENTICATION_METHOD: GoogleOAuthAuthenticationMethod.UPLOADED.value,\n }\n )\n return connector\n\n return _connector_factory\n" }, { "path": "backend/tests/daily/connectors/gmail/test_gmail_connector.py", "content": "from collections.abc import Callable\nfrom typing import Any\nfrom typing import cast\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nfrom onyx.connectors.gmail.connector import GmailConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom tests.unit.onyx.connectors.utils import load_everything_from_checkpoint_connector\n\n\n_THREAD_1_START_TIME = 1730568700\n_THREAD_1_END_TIME = 1730569000\n\n\"\"\"\nThis thread was 4 emails long:\n admin@onyx-test.com -> test-group-1@onyx-test.com (conaining test_user_1 and test_user_2)\n test_user_1@onyx-test.com -> admin@onyx-test.com\n admin@onyx-test.com -> test_user_2@onyx-test.com + BCC: test_user_3@onyx-test.com\n test_user_3@onyx-test.com -> admin@onyx-test.com\n\"\"\"\n_THREAD_1_BY_ID: dict[str, dict[str, Any]] = {\n \"192edefb315737c3\": {\n \"email\": \"admin@onyx-test.com\",\n \"sections_count\": 4,\n \"primary_owners\": set(\n [\n \"admin@onyx-test.com\",\n \"test_user_1@onyx-test.com\",\n \"test_user_3@onyx-test.com\",\n ]\n ),\n \"secondary_owners\": set(\n [\n \"test-group-1@onyx-test.com\",\n \"admin@onyx-test.com\",\n \"test_user_2@onyx-test.com\",\n \"test_user_3@onyx-test.com\",\n ]\n ),\n },\n \"192edf020d2f5def\": {\n \"email\": \"test_user_1@onyx-test.com\",\n \"sections_count\": 2,\n \"primary_owners\": set([\"admin@onyx-test.com\", \"test_user_1@onyx-test.com\"]),\n \"secondary_owners\": set([\"test-group-1@onyx-test.com\", \"admin@onyx-test.com\"]),\n },\n \"192edf020ae90aab\": {\n \"email\": \"test_user_2@onyx-test.com\",\n \"sections_count\": 2,\n \"primary_owners\": set([\"admin@onyx-test.com\"]),\n \"secondary_owners\": set(\n [\"test-group-1@onyx-test.com\", \"test_user_2@onyx-test.com\"]\n ),\n },\n \"192edf18316015fa\": {\n \"email\": \"test_user_3@onyx-test.com\",\n \"sections_count\": 2,\n \"primary_owners\": set([\"admin@onyx-test.com\", \"test_user_3@onyx-test.com\"]),\n \"secondary_owners\": set(\n [\n \"admin@onyx-test.com\",\n \"test_user_2@onyx-test.com\",\n \"test_user_3@onyx-test.com\",\n ]\n ),\n },\n}\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_slim_docs_retrieval(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_gmail_service_acct_connector_factory: Callable[..., GmailConnector],\n) -> None:\n print(\"\\n\\nRunning test_slim_docs_retrieval\")\n connector = google_gmail_service_acct_connector_factory()\n retrieved_slim_docs: list[SlimDocument] = []\n for doc_batch in connector.retrieve_all_slim_docs_perm_sync(\n _THREAD_1_START_TIME, _THREAD_1_END_TIME\n ):\n retrieved_slim_docs.extend(\n [doc for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n assert len(retrieved_slim_docs) == 4\n\n for doc in retrieved_slim_docs:\n assert doc.external_access is not None\n assert len(doc.external_access.external_user_emails) == 1\n user_email = next(iter(doc.external_access.external_user_emails))\n assert _THREAD_1_BY_ID[doc.id][\"email\"] == user_email\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_docs_retrieval(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_gmail_service_acct_connector_factory: Callable[..., GmailConnector],\n) -> None:\n print(\"\\n\\nRunning test_docs_retrieval\")\n connector = google_gmail_service_acct_connector_factory()\n retrieved_docs: list[Document] = []\n for doc_batch in load_everything_from_checkpoint_connector(\n connector, _THREAD_1_START_TIME, _THREAD_1_END_TIME\n ):\n assert all(isinstance(item, Document) for item in doc_batch.items)\n retrieved_docs.extend(cast(list[Document], doc_batch.items))\n\n assert len(retrieved_docs) == 4\n\n for doc in retrieved_docs:\n id = doc.id\n retrieved_primary_owner_emails: set[str | None] = set()\n retrieved_secondary_owner_emails: set[str | None] = set()\n if doc.primary_owners:\n retrieved_primary_owner_emails = set(\n [owner.email for owner in doc.primary_owners]\n )\n if doc.secondary_owners:\n retrieved_secondary_owner_emails = set(\n [owner.email for owner in doc.secondary_owners]\n )\n assert _THREAD_1_BY_ID[id][\"sections_count\"] == len(doc.sections)\n assert _THREAD_1_BY_ID[id][\"primary_owners\"] == retrieved_primary_owner_emails\n assert (\n _THREAD_1_BY_ID[id][\"secondary_owners\"] == retrieved_secondary_owner_emails\n )\n" }, { "path": "backend/tests/daily/connectors/gong/test_gong.py", "content": "import os\nimport time\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onyx.connectors.gong.connector import GongConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\n\n\n@pytest.fixture\ndef gong_connector() -> GongConnector:\n connector = GongConnector()\n\n connector.load_credentials(\n {\n \"gong_access_key\": os.environ[\"GONG_ACCESS_KEY\"],\n \"gong_access_key_secret\": os.environ[\"GONG_ACCESS_KEY_SECRET\"],\n }\n )\n\n return connector\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_gong_basic(\n mock_get_api_key: MagicMock, # noqa: ARG001\n gong_connector: GongConnector,\n) -> None:\n doc_batch_generator = gong_connector.poll_source(0, time.time())\n\n doc_batch = next(doc_batch_generator)\n with pytest.raises(StopIteration):\n next(doc_batch_generator)\n\n assert len(doc_batch) == 2\n\n docs: list[Document] = []\n for doc in doc_batch:\n if not isinstance(doc, HierarchyNode):\n docs.append(doc)\n\n assert docs[0].semantic_identifier == \"test with chris\"\n assert docs[1].semantic_identifier == \"Testing Gong\"\n" }, { "path": "backend/tests/daily/connectors/google_drive/conftest.py", "content": "import json\nimport os\nimport resource\nfrom collections.abc import Callable\n\nimport pytest\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_AUTHENTICATION_METHOD,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_TOKEN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n GoogleOAuthAuthenticationMethod,\n)\nfrom tests.load_env_vars import load_env_vars\n\n\n# Load environment variables at the module level\nload_env_vars()\n\n\n_USER_TO_OAUTH_CREDENTIALS_MAP = {\n \"admin@onyx-test.com\": \"GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR\",\n \"test_user_1@onyx-test.com\": \"GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR_TEST_USER_1\",\n}\n\n_USER_TO_SERVICE_ACCOUNT_CREDENTIALS_MAP = {\n \"admin@onyx-test.com\": \"GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR\",\n}\n\n\ndef parse_credentials(env_str: str) -> dict:\n \"\"\"\n Parse a double-escaped JSON string from environment variables into a Python dictionary.\n\n Args:\n env_str (str): The double-escaped JSON string from environment variables\n\n Returns:\n dict: Parsed OAuth credentials\n \"\"\"\n # first try normally\n try:\n return json.loads(env_str)\n except Exception:\n # First, try remove extra escaping backslashes\n unescaped = env_str.replace('\\\\\"', '\"')\n\n # remove leading / trailing quotes\n unescaped = unescaped.strip('\"')\n\n # Now parse the JSON\n return json.loads(unescaped)\n\n\ndef get_credentials_from_env(email: str, oauth: bool) -> dict:\n if oauth:\n raw_credential_string = os.environ[_USER_TO_OAUTH_CREDENTIALS_MAP[email]]\n else:\n raw_credential_string = os.environ[\n _USER_TO_SERVICE_ACCOUNT_CREDENTIALS_MAP[email]\n ]\n\n refried_credential_string = json.dumps(parse_credentials(raw_credential_string))\n\n cred_key = (\n DB_CREDENTIALS_DICT_TOKEN_KEY\n if oauth\n else DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY\n )\n return {\n cred_key: refried_credential_string,\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: email,\n DB_CREDENTIALS_AUTHENTICATION_METHOD: GoogleOAuthAuthenticationMethod.UPLOADED.value,\n }\n\n\n@pytest.fixture\ndef google_drive_oauth_uploaded_connector_factory() -> (\n Callable[..., GoogleDriveConnector]\n):\n def _connector_factory(\n primary_admin_email: str,\n include_shared_drives: bool,\n shared_drive_urls: str | None,\n include_my_drives: bool,\n my_drive_emails: str | None,\n shared_folder_urls: str | None,\n include_files_shared_with_me: bool,\n ) -> GoogleDriveConnector:\n print(\"Creating GoogleDriveConnector with OAuth credentials\")\n connector = GoogleDriveConnector(\n include_shared_drives=include_shared_drives,\n shared_drive_urls=shared_drive_urls,\n include_my_drives=include_my_drives,\n include_files_shared_with_me=include_files_shared_with_me,\n my_drive_emails=my_drive_emails,\n shared_folder_urls=shared_folder_urls,\n )\n\n credentials_json = get_credentials_from_env(primary_admin_email, oauth=True)\n connector.load_credentials(credentials_json)\n return connector\n\n return _connector_factory\n\n\n@pytest.fixture\ndef google_drive_service_acct_connector_factory() -> (\n Callable[..., GoogleDriveConnector]\n):\n def _connector_factory(\n primary_admin_email: str,\n include_shared_drives: bool,\n shared_drive_urls: str | None,\n include_my_drives: bool,\n my_drive_emails: str | None,\n shared_folder_urls: str | None,\n include_files_shared_with_me: bool,\n specific_user_emails: str | None = None,\n ) -> GoogleDriveConnector:\n print(\"Creating GoogleDriveConnector with service account credentials\")\n connector = GoogleDriveConnector(\n include_shared_drives=include_shared_drives,\n shared_drive_urls=shared_drive_urls,\n include_my_drives=include_my_drives,\n my_drive_emails=my_drive_emails,\n shared_folder_urls=shared_folder_urls,\n include_files_shared_with_me=include_files_shared_with_me,\n specific_user_emails=specific_user_emails,\n )\n\n # Load Service Account Credentials\n credentials_json = get_credentials_from_env(\n email=primary_admin_email, oauth=False\n )\n connector.load_credentials(credentials_json)\n return connector\n\n return _connector_factory\n\n\n@pytest.fixture(scope=\"session\", autouse=True)\ndef set_resource_limits() -> None:\n # the google sdk is aggressive about using up file descriptors and\n # macos is stingy ... these tests will fail randomly unless the descriptor limit is raised\n RLIMIT_MINIMUM = 2048\n soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)\n desired_soft = min(RLIMIT_MINIMUM, hard) # Pick your target here\n\n print(f\"Open file limit: soft={soft} hard={hard} soft_required={RLIMIT_MINIMUM}\")\n\n if soft < desired_soft:\n print(f\"Raising open file limit: {soft} -> {desired_soft}\")\n resource.setrlimit(resource.RLIMIT_NOFILE, (desired_soft, hard))\n\n soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)\n print(f\"New open file limit: soft={soft} hard={hard}\")\n return\n" }, { "path": "backend/tests/daily/connectors/google_drive/consts_and_utils.py", "content": "import time\nfrom collections.abc import Sequence\nfrom dataclasses import dataclass\nfrom dataclasses import field\nfrom dataclasses import replace\nfrom urllib.parse import urlparse\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.enums import HierarchyNodeType\nfrom tests.daily.connectors.utils import ConnectorOutput\nfrom tests.daily.connectors.utils import load_all_from_connector\n\nALL_FILES = list(range(0, 60))\nSHARED_DRIVE_FILES = list(range(20, 25))\n\n\nADMIN_FILE_IDS = list(range(0, 5))\nADMIN_FOLDER_3_FILE_IDS = list(range(65, 70)) # This folder is shared with test_user_1\nTEST_USER_1_FILE_IDS = list(range(5, 10))\nTEST_USER_2_FILE_IDS = list(range(10, 15))\nTEST_USER_3_FILE_IDS = list(range(15, 20))\nSHARED_DRIVE_1_FILE_IDS = list(range(20, 25))\nFOLDER_1_FILE_IDS = list(range(25, 30))\nFOLDER_1_1_FILE_IDS = list(range(30, 35))\nFOLDER_1_2_FILE_IDS = list(range(35, 40)) # This folder is public\nSHARED_DRIVE_2_FILE_IDS = list(range(40, 45))\nFOLDER_2_FILE_IDS = list(range(45, 50))\nFOLDER_2_1_FILE_IDS = list(range(50, 55))\nFOLDER_2_2_FILE_IDS = list(range(55, 60))\nSECTIONS_FILE_IDS = [61]\nFOLDER_3_FILE_IDS = list(range(62, 65))\n\nDONWLOAD_REVOKED_FILE_ID = 21\n\nPUBLIC_FOLDER_RANGE = FOLDER_1_2_FILE_IDS\nPUBLIC_FILE_IDS = list(range(55, 57))\nPUBLIC_RANGE = PUBLIC_FOLDER_RANGE + PUBLIC_FILE_IDS\n\nSHARED_DRIVE_1_URL = \"https://drive.google.com/drive/folders/0AC_OJ4BkMd4kUk9PVA\"\n# Group 1 is given access to this folder\nFOLDER_1_URL = (\n \"https://drive.google.com/drive/folders/1d3I7U3vUZMDziF1OQqYRkB8Jp2s_GWUn\"\n)\nFOLDER_1_1_URL = (\n \"https://drive.google.com/drive/folders/1aR33-zwzl_mnRAwH55GgtWTE-4A4yWWI\"\n)\nFOLDER_1_2_URL = (\n \"https://drive.google.com/drive/folders/1IO0X55VhvLXf4mdxzHxuKf4wxrDBB6jq\"\n)\nSHARED_DRIVE_2_URL = \"https://drive.google.com/drive/folders/0ABKspIh7P4f4Uk9PVA\"\nFOLDER_2_URL = (\n \"https://drive.google.com/drive/folders/1lNpCJ1teu8Se0louwL0oOHK9nEalskof\"\n)\nFOLDER_2_1_URL = (\n \"https://drive.google.com/drive/folders/1XeDOMWwxTDiVr9Ig2gKum3Zq_Wivv6zY\"\n)\nFOLDER_2_2_URL = (\n \"https://drive.google.com/drive/folders/1RKlsexA8h7NHvBAWRbU27MJotic7KXe3\"\n)\nFOLDER_3_URL = (\n \"https://drive.google.com/drive/folders/1LHibIEXfpUmqZ-XjBea44SocA91Nkveu\"\n)\nSECTIONS_FOLDER_URL = (\n \"https://drive.google.com/drive/u/5/folders/1loe6XJ-pJxu9YYPv7cF3Hmz296VNzA33\"\n)\n\n\ndef extract_folder_id_from_url(url: str) -> str:\n \"\"\"Extract the folder ID from a Google Drive URL.\"\"\"\n parsed = urlparse(url)\n # URL format: /drive/folders/{id} or /drive/u/{num}/folders/{id}\n parts = parsed.path.split(\"/\")\n # Find 'folders' and take the next segment\n for i, part in enumerate(parts):\n if part == \"folders\" and i + 1 < len(parts):\n return parts[i + 1]\n raise ValueError(f\"Could not extract folder ID from URL: {url}\")\n\n\n# Folder IDs extracted from URLs\nSHARED_DRIVE_1_ID = extract_folder_id_from_url(SHARED_DRIVE_1_URL)\nSHARED_DRIVE_2_ID = extract_folder_id_from_url(SHARED_DRIVE_2_URL)\nFOLDER_1_ID = extract_folder_id_from_url(FOLDER_1_URL)\nFOLDER_1_1_ID = extract_folder_id_from_url(FOLDER_1_1_URL)\nFOLDER_1_2_ID = extract_folder_id_from_url(FOLDER_1_2_URL)\nFOLDER_2_ID = extract_folder_id_from_url(FOLDER_2_URL)\nFOLDER_2_1_ID = extract_folder_id_from_url(FOLDER_2_1_URL)\nFOLDER_2_2_ID = extract_folder_id_from_url(FOLDER_2_2_URL)\nFOLDER_3_ID = extract_folder_id_from_url(FOLDER_3_URL)\nSECTIONS_FOLDER_ID = extract_folder_id_from_url(SECTIONS_FOLDER_URL)\nRESTRICTED_ACCESS_FOLDER_ID = \"1HK4wZ16ucz8QGywlcS87Y629W7i7KdeN\"\n\n\n# ============================================================================\n# FOLDER HIERARCHY DEFINITION\n# ============================================================================\n# This defines the expected folder hierarchy for our test Google Drive setup.\n#\n# Folder Hierarchy:\n# shared_drive_1 (0AC_OJ4BkMd4kUk9PVA)\n# ├── restricted_access_folder (1HK4wZ16ucz8QGywlcS87Y629W7i7KdeN)\n# └── folder_1 (1d3I7U3vUZMDziF1OQqYRkB8Jp2s_GWUn)\n# ├── folder_1_1 (1aR33-zwzl_mnRAwH55GgtWTE-4A4yWWI)\n# └── folder_1_2 (1IO0X55VhvLXf4mdxzHxuKf4wxrDBB6jq)\n#\n# shared_drive_2 (0ABKspIh7P4f4Uk9PVA)\n# ├── sections_folder (1loe6XJ-pJxu9YYPv7cF3Hmz296VNzA33)\n# └── folder_2 (1lNpCJ1teu8Se0louwL0oOHK9nEalskof)\n# ├── folder_2_1 (1XeDOMWwxTDiVr9Ig2gKum3Zq_Wivv6zY)\n# └── folder_2_2 (1RKlsexA8h7NHvBAWRbU27MJotic7KXe3)\n# ============================================================================\n\n\n@dataclass\nclass ExpectedHierarchyNode:\n \"\"\"Expected hierarchy node for test verification.\"\"\"\n\n raw_node_id: str\n display_name: str\n node_type: HierarchyNodeType\n # None means parent is the source root (shared drive or my drive)\n raw_parent_id: str | None = None\n children: list[\"ExpectedHierarchyNode\"] = field(default_factory=list)\n\n\n# Expected hierarchy for shared_drive_1\nEXPECTED_SHARED_DRIVE_1_HIERARCHY = ExpectedHierarchyNode(\n raw_node_id=SHARED_DRIVE_1_ID,\n display_name=\"Shared Drive 1\",\n node_type=HierarchyNodeType.SHARED_DRIVE,\n raw_parent_id=None,\n children=[\n ExpectedHierarchyNode(\n raw_node_id=RESTRICTED_ACCESS_FOLDER_ID,\n display_name=\"restricted_access\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=SHARED_DRIVE_1_ID,\n ),\n ExpectedHierarchyNode(\n raw_node_id=FOLDER_1_ID,\n display_name=\"folder 1\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=SHARED_DRIVE_1_ID,\n children=[\n ExpectedHierarchyNode(\n raw_node_id=FOLDER_1_1_ID,\n display_name=\"folder 1-1\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=FOLDER_1_ID,\n ),\n ExpectedHierarchyNode(\n raw_node_id=FOLDER_1_2_ID,\n display_name=\"folder 1-2\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=FOLDER_1_ID,\n ),\n ],\n ),\n ],\n)\n\n# Expected hierarchy for shared_drive_2\nEXPECTED_SHARED_DRIVE_2_HIERARCHY = ExpectedHierarchyNode(\n raw_node_id=SHARED_DRIVE_2_ID,\n display_name=\"Shared Drive 2\",\n node_type=HierarchyNodeType.SHARED_DRIVE,\n raw_parent_id=None,\n children=[\n ExpectedHierarchyNode(\n raw_node_id=SECTIONS_FOLDER_ID,\n display_name=\"sections\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=SHARED_DRIVE_2_ID,\n ),\n ExpectedHierarchyNode(\n raw_node_id=FOLDER_2_ID,\n display_name=\"folder 2\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=SHARED_DRIVE_2_ID,\n children=[\n ExpectedHierarchyNode(\n raw_node_id=FOLDER_2_1_ID,\n display_name=\"folder 2-1\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=FOLDER_2_ID,\n ),\n ExpectedHierarchyNode(\n raw_node_id=FOLDER_2_2_ID,\n display_name=\"folder 2-2\",\n node_type=HierarchyNodeType.FOLDER,\n raw_parent_id=FOLDER_2_ID,\n ),\n ],\n ),\n ],\n)\n\n\ndef flatten_hierarchy(\n expected: ExpectedHierarchyNode,\n) -> dict[str, ExpectedHierarchyNode]:\n \"\"\"Flatten an expected hierarchy tree into a dict keyed by raw_node_id.\"\"\"\n result = {expected.raw_node_id: expected}\n for child in expected.children:\n result.update(flatten_hierarchy(child))\n return result\n\n\ndef _node(\n raw_node_id: str,\n display_name: str,\n node_type: HierarchyNodeType,\n raw_parent_id: str | None = None,\n) -> ExpectedHierarchyNode:\n return ExpectedHierarchyNode(\n raw_node_id=raw_node_id,\n display_name=display_name,\n node_type=node_type,\n raw_parent_id=raw_parent_id,\n )\n\n\n# Flattened maps for easy lookup\nEXPECTED_SHARED_DRIVE_1_NODES = flatten_hierarchy(EXPECTED_SHARED_DRIVE_1_HIERARCHY)\nEXPECTED_SHARED_DRIVE_2_NODES = flatten_hierarchy(EXPECTED_SHARED_DRIVE_2_HIERARCHY)\n\nEXTERNAL_SHARED_FOLDER_URL = (\n \"https://drive.google.com/drive/folders/1sWC7Oi0aQGgifLiMnhTjvkhRWVeDa-XS\"\n)\nEXTERNAL_SHARED_FOLDER_ID = \"1sWC7Oi0aQGgifLiMnhTjvkhRWVeDa-XS\"\nEXTERNAL_SHARED_DOCS_IN_FOLDER = [\n \"https://docs.google.com/document/d/1Sywmv1-H6ENk2GcgieKou3kQHR_0te1mhIUcq8XlcdY\"\n]\nEXTERNAL_SHARED_DOC_SINGLETON = (\n \"https://docs.google.com/document/d/11kmisDfdvNcw5LYZbkdPVjTOdj-Uc5ma6Jep68xzeeA\"\n)\n\nSHARED_DRIVE_3_URL = \"https://drive.google.com/drive/folders/0AJYm2K_I_vtNUk9PVA\"\n\nRESTRICTED_ACCESS_FOLDER_URL = (\n \"https://drive.google.com/drive/folders/1HK4wZ16ucz8QGywlcS87Y629W7i7KdeN\"\n)\n\n# ============================================================================\n# PERMISSION SYNC TEST DRIVES\n# ============================================================================\n# These are separate shared drives used specifically for testing permission sync.\n# Each drive has different access levels:\n#\n# PERM_SYNC_DRIVE_ADMIN_ONLY: Only shared with admin\n# PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A: Shared with admin and test_user_1\n# PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B: Shared with admin and test_user_1\n# ============================================================================\n\nPERM_SYNC_DRIVE_ADMIN_ONLY_URL = (\n \"https://drive.google.com/drive/folders/0ACOrCU1EMD1hUk9PVA\"\n)\nPERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_URL = (\n \"https://drive.google.com/drive/folders/0ABec4pV29sMuUk9PVA\"\n)\nPERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_URL = (\n \"https://drive.google.com/drive/folders/0ANpbToRgjHD4Uk9PVA\"\n)\n\nPERM_SYNC_DRIVE_ADMIN_ONLY_ID = \"0ACOrCU1EMD1hUk9PVA\"\nPERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID = \"0ABec4pV29sMuUk9PVA\"\nPERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID = \"0ANpbToRgjHD4Uk9PVA\"\n\n# ============================================================================\n# ADDITIONAL DRIVES/FOLDERS ACCESSIBLE TO TEST_USER_1\n# ============================================================================\n# These are additional shared drives and folders that test_user_1 has access to.\n# They are returned as hierarchy nodes when running the connector as test_user_1.\n# ============================================================================\n\n# Additional shared drives accessible to test_user_1\nTEST_USER_1_MY_DRIVE_ID = \"0AFpeuWG1VyABUk9PVA\" # My Drive indicator for test_user_1\nTEST_USER_1_MY_DRIVE_FOLDER_ID = (\n \"1tF10nDFND-GE_IT0f6PjEn2Du6m2k-DE\" # Child folder (partial sharing)\n)\n\nTEST_USER_1_DRIVE_B_ID = (\n \"0AFskk4zfZm86Uk9PVA\" # My_super_special_shared_drive_suuuper_private\n)\nTEST_USER_1_DRIVE_B_FOLDER_ID = (\n \"1oIj7nigzvP5xI2F8BmibUA8R_J3AbBA-\" # Child folder (silliness)\n)\n\n# Other drives test_user_1 has access to\nTEST_USER_1_EXTRA_DRIVE_1_ID = \"0AL67XRMq9reYUk9PVA\" # Okay_fine_admin_I_will_share\nTEST_USER_1_EXTRA_DRIVE_2_ID = \"0ACeKoHrGKxCbUk9PVA\" # reee test\nTEST_USER_1_EXTRA_FOLDER_ID = (\n \"1i2Q1TNvUfZkH-A7RGyAqRuEI-3mHANku\" # read only no download test\n)\n\n# Additional shared drives in the organization that appear when running include_all tests\nADMIN_MY_DRIVE_ID = \"0ABTZwt798K7MUk9PVA\" # Admin's My Drive\nTEST_USER_2_MY_DRIVE = \"0ADjBZv2nEvJNUk9PVA\" # Test user 2's My Drive\nTEST_USER_3_MY_DRIVE_ID = \"0AKl0e4Wr5NW7Uk9PVA\" # Test user 3's My Drive\nPILL_FOLDER_ID = \"1FWzfA369tx9VT8scJ3LCOPBBuTBgt0OH\" # contains file with date pills\n\nPADDING_DRIVE_URLS = [\n \"0AOorXE6AfJRAUk9PVA\",\n \"0ANn2MSqGi74JUk9PVA\",\n \"0ANI_NFCPzaRwUk9PVA\",\n \"0ABu8fYjvA21dUk9PVA\",\n]\n\nADMIN_EMAIL = \"admin@onyx-test.com\"\nTEST_USER_1_EMAIL = \"test_user_1@onyx-test.com\"\nTEST_USER_2_EMAIL = \"test_user_2@onyx-test.com\"\nTEST_USER_3_EMAIL = \"test_user_3@onyx-test.com\"\n\n# Expected permissions for perm sync drives\n# Maps drive ID -> set of user emails with access\nPERM_SYNC_DRIVE_ACCESS_MAPPING: dict[str, set[str]] = {\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID: {ADMIN_EMAIL},\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID: {ADMIN_EMAIL, TEST_USER_1_EMAIL},\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID: {ADMIN_EMAIL, TEST_USER_1_EMAIL},\n}\n\n# ============================================================================\n# NON-SHARED-DRIVE HIERARCHY NODES\n# ============================================================================\n# These cover My Drive roots, perm sync drives, extra shared drives,\n# and standalone folders that appear in various tests.\n# Display names must match what the Google Drive API actually returns.\n# ============================================================================\n\nEXPECTED_FOLDER_3 = _node(\n FOLDER_3_ID, \"Folder 3\", HierarchyNodeType.FOLDER, ADMIN_MY_DRIVE_ID\n)\n\nEXPECTED_ADMIN_MY_DRIVE = _node(ADMIN_MY_DRIVE_ID, \"My Drive\", HierarchyNodeType.FOLDER)\nEXPECTED_TEST_USER_1_MY_DRIVE = _node(\n TEST_USER_1_MY_DRIVE_ID, \"My Drive\", HierarchyNodeType.FOLDER\n)\nEXPECTED_TEST_USER_1_MY_DRIVE_FOLDER = _node(\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n \"partial_sharing\",\n HierarchyNodeType.FOLDER,\n TEST_USER_1_MY_DRIVE_ID,\n)\nEXPECTED_TEST_USER_2_MY_DRIVE = _node(\n TEST_USER_2_MY_DRIVE, \"My Drive\", HierarchyNodeType.FOLDER\n)\nEXPECTED_TEST_USER_3_MY_DRIVE = _node(\n TEST_USER_3_MY_DRIVE_ID, \"My Drive\", HierarchyNodeType.FOLDER\n)\n\nEXPECTED_PERM_SYNC_DRIVE_ADMIN_ONLY = _node(\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n \"perm_sync_drive_0dc9d8b5-e243-4c2f-8678-2235958f7d7c\",\n HierarchyNodeType.SHARED_DRIVE,\n)\nEXPECTED_PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A = _node(\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n \"perm_sync_drive_785db121-0823-4ebe-8689-ad7f52405e32\",\n HierarchyNodeType.SHARED_DRIVE,\n)\nEXPECTED_PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B = _node(\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n \"perm_sync_drive_d8dc3649-3f65-4392-b87f-4b20e0389673\",\n HierarchyNodeType.SHARED_DRIVE,\n)\n\nEXPECTED_TEST_USER_1_DRIVE_B = _node(\n TEST_USER_1_DRIVE_B_ID,\n \"My_super_special_shared_drive_suuuper_private\",\n HierarchyNodeType.SHARED_DRIVE,\n)\nEXPECTED_TEST_USER_1_DRIVE_B_FOLDER = _node(\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n \"silliness\",\n HierarchyNodeType.FOLDER,\n TEST_USER_1_DRIVE_B_ID,\n)\nEXPECTED_TEST_USER_1_EXTRA_DRIVE_1 = _node(\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n \"Okay_Admin_fine_I_will_share\",\n HierarchyNodeType.SHARED_DRIVE,\n)\nEXPECTED_TEST_USER_1_EXTRA_DRIVE_2 = _node(\n TEST_USER_1_EXTRA_DRIVE_2_ID, \"reee test\", HierarchyNodeType.SHARED_DRIVE\n)\nEXPECTED_TEST_USER_1_EXTRA_FOLDER = _node(\n TEST_USER_1_EXTRA_FOLDER_ID,\n \"read only no download test\",\n HierarchyNodeType.FOLDER,\n)\n\nEXPECTED_PILL_FOLDER = _node(\n PILL_FOLDER_ID, \"pill_folder\", HierarchyNodeType.FOLDER, ADMIN_MY_DRIVE_ID\n)\nEXPECTED_EXTERNAL_SHARED_FOLDER = _node(\n EXTERNAL_SHARED_FOLDER_ID, \"Onyx-test\", HierarchyNodeType.FOLDER\n)\n\n# Comprehensive mapping of ALL known hierarchy nodes.\n# Every retrieved node is checked against this for display_name and node_type.\nALL_EXPECTED_HIERARCHY_NODES: dict[str, ExpectedHierarchyNode] = {\n **EXPECTED_SHARED_DRIVE_1_NODES,\n **EXPECTED_SHARED_DRIVE_2_NODES,\n FOLDER_3_ID: EXPECTED_FOLDER_3,\n ADMIN_MY_DRIVE_ID: EXPECTED_ADMIN_MY_DRIVE,\n TEST_USER_1_MY_DRIVE_ID: EXPECTED_TEST_USER_1_MY_DRIVE,\n TEST_USER_1_MY_DRIVE_FOLDER_ID: EXPECTED_TEST_USER_1_MY_DRIVE_FOLDER,\n TEST_USER_2_MY_DRIVE: EXPECTED_TEST_USER_2_MY_DRIVE,\n TEST_USER_3_MY_DRIVE_ID: EXPECTED_TEST_USER_3_MY_DRIVE,\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID: EXPECTED_PERM_SYNC_DRIVE_ADMIN_ONLY,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID: EXPECTED_PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID: EXPECTED_PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B,\n TEST_USER_1_DRIVE_B_ID: EXPECTED_TEST_USER_1_DRIVE_B,\n TEST_USER_1_DRIVE_B_FOLDER_ID: EXPECTED_TEST_USER_1_DRIVE_B_FOLDER,\n TEST_USER_1_EXTRA_DRIVE_1_ID: EXPECTED_TEST_USER_1_EXTRA_DRIVE_1,\n TEST_USER_1_EXTRA_DRIVE_2_ID: EXPECTED_TEST_USER_1_EXTRA_DRIVE_2,\n TEST_USER_1_EXTRA_FOLDER_ID: EXPECTED_TEST_USER_1_EXTRA_FOLDER,\n PILL_FOLDER_ID: EXPECTED_PILL_FOLDER,\n EXTERNAL_SHARED_FOLDER_ID: EXPECTED_EXTERNAL_SHARED_FOLDER,\n}\n\n# Dictionary for access permissions\n# All users have access to their own My Drive as well as public files\nACCESS_MAPPING: dict[str, list[int]] = {\n # Admin has access to everything in shared\n ADMIN_EMAIL: (\n ADMIN_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n + SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + SHARED_DRIVE_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + SECTIONS_FILE_IDS\n ),\n TEST_USER_1_EMAIL: (\n TEST_USER_1_FILE_IDS\n # This user has access to drive 1\n + SHARED_DRIVE_1_FILE_IDS\n # This user has redundant access to folder 1 because of group access\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n # This user has been given shared access to folder 3 in Admin's My Drive\n + ADMIN_FOLDER_3_FILE_IDS\n # This user has been given shared access to files 0 and 1 in Admin's My Drive\n + list(range(0, 2))\n ),\n TEST_USER_2_EMAIL: (\n TEST_USER_2_FILE_IDS\n # Group 1 includes this user, giving access to folder 1\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n # This folder is public\n + FOLDER_1_2_FILE_IDS\n # Folder 2-1 is shared with this user\n + FOLDER_2_1_FILE_IDS\n # This user has been given shared access to files 45 and 46 in folder 2\n + list(range(45, 47))\n ),\n # This user can only see his own files and public files\n TEST_USER_3_EMAIL: TEST_USER_3_FILE_IDS,\n}\n\nSPECIAL_FILE_ID_TO_CONTENT_MAP: dict[int, str] = {\n 61: (\n \"Title\\n\"\n \"This is a Google Doc with sections - \"\n \"Section 1\\n\"\n \"Section 1 content - \"\n \"Sub-Section 1-1\\n\"\n \"Sub-Section 1-1 content - \"\n \"Sub-Section 1-2\\n\"\n \"Sub-Section 1-2 content - \"\n \"Section 2\\n\"\n \"Section 2 content\"\n ),\n}\n\nMISC_SHARED_DRIVE_FNAMES = [\n \"asdfasdfsfad\",\n \"perm_sync_doc_0ABec4pV29sMuUk9PVA_a5ea8ec4-0440-4926-a43d-3aeef1c10bdd\",\n \"perm_sync_doc_0ACOrCU1EMD1hUk9PVA_651821cb-8140-42fe-a876-1a92012375c9\",\n \"perm_sync_doc_0ACOrCU1EMD1hUk9PVA_ab63b976-effb-49af-84e7-423d17a17dd7\",\n \"super secret thing that test user 1 can't see\",\n \"perm_sync_doc_0ABec4pV29sMuUk9PVA_419f2ef0-9815-4c69-8435-98b163c9c156\",\n \"Untitled documentfsdfsdfsdf\",\n \"bingle_bongle.txt\",\n \"bb4.txt\",\n \"bb3.txt\",\n \"bb2.txt\",\n]\n\nfile_name_template = \"file_{}.txt\"\nfile_text_template = \"This is file {}\"\n\n# This is done to prevent different tests from interfering with each other\n# So each test type should have its own valid prefix\n_VALID_PREFIX = \"file_\"\n\n\ndef filter_invalid_prefixes(names: set[str]) -> set[str]:\n return {name for name in names if name.startswith(_VALID_PREFIX)}\n\n\ndef print_discrepancies(\n expected: set[str],\n retrieved: set[str],\n) -> None:\n if expected != retrieved:\n expected_list = sorted(expected)\n retrieved_list = sorted(retrieved)\n print(expected_list)\n print(retrieved_list)\n print(\"Extra:\")\n print(sorted(retrieved - expected))\n print(\"Missing:\")\n print(sorted(expected - retrieved))\n\n\ndef _get_expected_file_content(file_id: int) -> str:\n if file_id in SPECIAL_FILE_ID_TO_CONTENT_MAP:\n return SPECIAL_FILE_ID_TO_CONTENT_MAP[file_id]\n\n return file_text_template.format(file_id)\n\n\ndef id_to_name(file_id: int) -> str:\n return file_name_template.format(file_id)\n\n\ndef assert_expected_docs_in_retrieved_docs(\n retrieved_docs: list[Document],\n expected_file_ids: Sequence[int],\n) -> None:\n \"\"\"NOTE: as far as i can tell this does NOT assert for an exact match.\n it only checks to see if that the expected file id's are IN the retrieved doc list\n \"\"\"\n\n expected_file_names = {id_to_name(file_id) for file_id in expected_file_ids}\n expected_file_texts = {\n _get_expected_file_content(file_id) for file_id in expected_file_ids\n }\n\n retrieved_docs.sort(key=lambda x: x.semantic_identifier)\n\n for doc in retrieved_docs:\n print(f\"retrieved doc: doc.semantic_identifier={doc.semantic_identifier}\")\n\n # Filter out invalid prefixes to prevent different tests from interfering with each other\n valid_retrieved_docs = [\n doc\n for doc in retrieved_docs\n if doc.semantic_identifier.startswith(_VALID_PREFIX)\n ]\n valid_retrieved_file_names = set(\n [doc.semantic_identifier for doc in valid_retrieved_docs]\n )\n valid_retrieved_texts = set(\n [\n \" - \".join(\n [\n section.text\n for section in doc.sections\n if isinstance(section, TextSection) and section.text is not None\n ]\n )\n for doc in valid_retrieved_docs\n ]\n )\n\n # Check file names\n print_discrepancies(\n expected=expected_file_names,\n retrieved=valid_retrieved_file_names,\n )\n assert expected_file_names == valid_retrieved_file_names\n\n # Check file texts\n print_discrepancies(\n expected=expected_file_texts,\n retrieved=valid_retrieved_texts,\n )\n assert expected_file_texts == valid_retrieved_texts\n\n\ndef load_connector_outputs(\n connector: GoogleDriveConnector,\n include_permissions: bool = False,\n) -> ConnectorOutput:\n \"\"\"Load all documents, failures, and hierarchy nodes from the connector.\"\"\"\n return load_all_from_connector(\n connector,\n 0,\n time.time(),\n include_permissions=include_permissions,\n )\n\n\ndef assert_hierarchy_nodes_match_expected(\n retrieved_nodes: list[HierarchyNode],\n expected_nodes: dict[str, ExpectedHierarchyNode],\n ignorable_node_ids: set[str] | None = None,\n) -> None:\n \"\"\"\n Assert that retrieved hierarchy nodes match expected structure.\n\n Checks node IDs, display names, node types, and parent relationships\n for EVERY retrieved node (global checks).\n\n Args:\n retrieved_nodes: List of HierarchyNode objects from the connector\n expected_nodes: Dict mapping raw_node_id -> ExpectedHierarchyNode with\n expected display_name, node_type, and raw_parent_id\n ignorable_node_ids: Optional set of node IDs that can be missing or extra\n without failing. Useful for non-deterministically returned nodes.\n \"\"\"\n expected_node_ids = set(expected_nodes.keys())\n retrieved_node_ids = {node.raw_node_id for node in retrieved_nodes}\n ignorable = ignorable_node_ids or set()\n\n missing = expected_node_ids - retrieved_node_ids - ignorable\n extra = retrieved_node_ids - expected_node_ids - ignorable\n\n if missing or extra:\n print(\"Expected hierarchy node IDs:\")\n print(sorted(expected_node_ids))\n print(\"Retrieved hierarchy node IDs:\")\n print(sorted(retrieved_node_ids))\n print(\"Extra (retrieved but not expected):\")\n print(sorted(retrieved_node_ids - expected_node_ids))\n print(\"Missing (expected but not retrieved):\")\n print(sorted(expected_node_ids - retrieved_node_ids))\n if ignorable:\n print(\"Ignorable node IDs:\")\n print(sorted(ignorable))\n\n assert (\n not missing and not extra\n ), f\"Hierarchy node mismatch. Missing: {missing}, Extra: {extra}\"\n\n for node in retrieved_nodes:\n if node.raw_node_id in ignorable and node.raw_node_id not in expected_nodes:\n continue\n\n assert (\n node.raw_node_id in expected_nodes\n ), f\"Node {node.raw_node_id} ({node.display_name}) not found in expected_nodes\"\n expected = expected_nodes[node.raw_node_id]\n\n assert (\n node.display_name == expected.display_name\n ), f\"Display name mismatch for node {node.raw_node_id}: expected '{expected.display_name}', got '{node.display_name}'\"\n assert (\n node.node_type == expected.node_type\n ), f\"Node type mismatch for node {node.raw_node_id}: expected '{expected.node_type}', got '{node.node_type}'\"\n if expected.raw_parent_id is not None:\n assert node.raw_parent_id == expected.raw_parent_id, (\n f\"Parent mismatch for node {node.raw_node_id} ({node.display_name}): \"\n f\"expected parent={expected.raw_parent_id}, got parent={node.raw_parent_id}\"\n )\n\n\ndef _pick(\n *node_ids: str,\n) -> dict[str, ExpectedHierarchyNode]:\n \"\"\"Pick nodes from ALL_EXPECTED_HIERARCHY_NODES by their IDs.\"\"\"\n return {nid: ALL_EXPECTED_HIERARCHY_NODES[nid] for nid in node_ids}\n\n\ndef _clear_parents(\n nodes: dict[str, ExpectedHierarchyNode],\n *node_ids: str,\n) -> dict[str, ExpectedHierarchyNode]:\n \"\"\"Return a shallow copy of nodes with the specified nodes' parents set to None.\n Useful for OAuth tests where the user can't resolve certain parents\n (e.g. a folder in another user's My Drive).\"\"\"\n result = dict(nodes)\n for nid in node_ids:\n result[nid] = replace(result[nid], raw_parent_id=None)\n return result\n\n\ndef get_expected_hierarchy_for_shared_drives(\n include_drive_1: bool = True,\n include_drive_2: bool = True,\n include_restricted_folder: bool = True,\n) -> dict[str, ExpectedHierarchyNode]:\n \"\"\"Get expected hierarchy nodes for shared drives.\"\"\"\n result: dict[str, ExpectedHierarchyNode] = {}\n\n if include_drive_1:\n result.update(EXPECTED_SHARED_DRIVE_1_NODES)\n if not include_restricted_folder:\n result.pop(RESTRICTED_ACCESS_FOLDER_ID, None)\n\n if include_drive_2:\n result.update(EXPECTED_SHARED_DRIVE_2_NODES)\n\n return result\n\n\ndef get_expected_hierarchy_for_folder_1() -> dict[str, ExpectedHierarchyNode]:\n \"\"\"Get expected hierarchy for folder_1 and its children only.\"\"\"\n return _pick(FOLDER_1_ID, FOLDER_1_1_ID, FOLDER_1_2_ID)\n\n\ndef get_expected_hierarchy_for_folder_2() -> dict[str, ExpectedHierarchyNode]:\n \"\"\"Get expected hierarchy for folder_2 and its children only.\"\"\"\n return _pick(FOLDER_2_ID, FOLDER_2_1_ID, FOLDER_2_2_ID)\n\n\ndef get_expected_hierarchy_for_test_user_1() -> dict[str, ExpectedHierarchyNode]:\n \"\"\"\n Get expected hierarchy for test_user_1's full access (OAuth).\n\n test_user_1 has access to:\n - shared_drive_1 and its contents (folder_1, folder_1_1, folder_1_2)\n - folder_3 (shared from admin's My Drive)\n - PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A and PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B\n - Additional drives/folders the user has access to\n\n NOTE: Folder 3 lives in the admin's My Drive. When running as an OAuth\n connector for test_user_1, the Google Drive API won't return the parent\n for Folder 3 because the user can't access the admin's My Drive root.\n \"\"\"\n result = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=False,\n include_restricted_folder=False,\n )\n result.update(\n _pick(\n FOLDER_3_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n TEST_USER_1_MY_DRIVE_ID,\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n TEST_USER_1_DRIVE_B_ID,\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n )\n )\n return _clear_parents(result, FOLDER_3_ID)\n\n\ndef get_expected_hierarchy_for_test_user_1_shared_drives_only() -> (\n dict[str, ExpectedHierarchyNode]\n):\n \"\"\"Expected hierarchy nodes when test_user_1 runs with include_shared_drives=True only.\"\"\"\n result = get_expected_hierarchy_for_test_user_1()\n for nid in (\n TEST_USER_1_MY_DRIVE_ID,\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n FOLDER_3_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n ):\n result.pop(nid, None)\n return result\n\n\ndef get_expected_hierarchy_for_test_user_1_shared_with_me_only() -> (\n dict[str, ExpectedHierarchyNode]\n):\n \"\"\"Expected hierarchy nodes when test_user_1 runs with include_files_shared_with_me=True only.\"\"\"\n return _clear_parents(\n _pick(FOLDER_3_ID, TEST_USER_1_EXTRA_FOLDER_ID),\n FOLDER_3_ID,\n )\n\n\ndef get_expected_hierarchy_for_test_user_1_my_drive_only() -> (\n dict[str, ExpectedHierarchyNode]\n):\n \"\"\"Expected hierarchy nodes when test_user_1 runs with include_my_drives=True only.\"\"\"\n return _pick(TEST_USER_1_MY_DRIVE_ID, TEST_USER_1_MY_DRIVE_FOLDER_ID)\n" }, { "path": "backend/tests/daily/connectors/google_drive/drive_id_mapping.json", "content": "{\n \"12\": \"https://drive.google.com/file/d/1u7nynrG4WuFZeuZs8yyhqJF_lbo-op-m\",\n \"10\": \"https://drive.google.com/file/d/1LFcVuXuXIdNJ7hkL0C40eYn_cQtryUVQ\",\n \"13\": \"https://drive.google.com/file/d/1muQMyYAJe0_F-HiDFIfFMt-4qsgMlREM\",\n \"11\": \"https://drive.google.com/file/d/1oHNtlsdJJtk7dE10NgH83Kn5_f2L-Su1\",\n \"14\": \"https://drive.google.com/file/d/1sAw-DrsqpnqLF5A8P59BZwIpt9-LrlaL\",\n \"18\": \"https://drive.google.com/file/d/1qqKH3esasdqV6ryEhdoSQezDPlKj11At\",\n \"17\": \"https://drive.google.com/file/d/1z08VsrCUTozpc5Quzb7mEDUwNkXU3foT\",\n \"15\": \"https://drive.google.com/file/d/1QQ6ZGyYP49IJNeGKNmqZISyVLzTOtK4v\",\n \"19\": \"https://drive.google.com/file/d/172as_pb7E15bXUd63mIIBRotk_tT7h56\",\n \"16\": \"https://drive.google.com/file/d/1552S6HEjJ81q8JXr46BtixQiVq9xlW_I\",\n \"5\": \"https://drive.google.com/file/d/1sv9epxLcNlgM6C-oPDeD_heFw7AIZMgp\",\n \"7\": \"https://drive.google.com/file/d/1S_S0LpQW90EUPPPjJX4jfu5p9gOQjiQF\",\n \"9\": \"https://drive.google.com/file/d/1wH2dBrWzmiGJ88ySHWu6srb7Jsj7qYbA\",\n \"8\": \"https://drive.google.com/file/d/14URUm6RKSZziH1lUtT6gs-xnCTWkXpSn\",\n \"6\": \"https://drive.google.com/file/d/1LBKBuTMRSss-kVw8ut3rMk51wSbTM95j\",\n \"3\": \"https://drive.google.com/file/d/1nNazkPrkuRXHFOl8gdA68pU2g8cy-h6n\",\n \"2\": \"https://drive.google.com/file/d/1miG_QpqXe2QIMApcrlNzaB6fsXW5WMFX\",\n \"4\": \"https://drive.google.com/file/d/1o-i8can6ciL1XXzy2pVUPHZEXEjBJi6C\",\n \"0\": \"https://drive.google.com/file/d/1d3Y59Sns8I0FIW9CtOAjVVLE2MEe_3nP\",\n \"1\": \"https://drive.google.com/file/d/1ipSqxJajs_NkfSKFxgltIMNc0ffdt-NX\",\n \"68\": \"https://drive.google.com/file/d/1rCBZsbhQ-ULWGztiKB0JYhFth9EChiSZ\",\n \"66\": \"https://drive.google.com/file/d/1WVAlbWcu9-Braa0aG6w3cShrY5dbIYcY\",\n \"67\": \"https://drive.google.com/file/d/1p44poOCdNLnVYMxTL9b3h-BXsOQ2RDgM\",\n \"69\": \"https://drive.google.com/file/d/1HFYsaqC14aE-EaobQdwkw0FOlAYMYqkV\",\n \"65\": \"https://drive.google.com/file/d/1RyE07CpTIDYMO3b-atwjWH6ZHFDjyoCl\",\n \"32\": \"https://drive.google.com/file/d/17egJ5W-0bvS2akLBqvxylTIViN0d9nG7\",\n \"28\": \"https://drive.google.com/file/d/1HNqSM2XGqgHnyNYT5wp8hyski18HMcfO\",\n \"37\": \"https://drive.google.com/file/d/16Tdu3gveWkFL0VBUzYSzKxFO4ffv-8h7\",\n \"30\": \"https://drive.google.com/file/d/1uj69jGyYnNOXXqKmLNIp-4KKrVC1qaPy\",\n \"25\": \"https://drive.google.com/file/d/1bw6NFlR4ZxOV6reQK1Oqeq_UaYFVpNV6\",\n \"33\": \"https://drive.google.com/file/d/1FkmXBkt__lOFXg_uhxLI0QIuxWbIGySL\",\n \"20\": \"https://drive.google.com/file/d/1r77uBVOHkuiDQFa9iz9FU8QbfjImOAjF\",\n \"24\": \"https://drive.google.com/file/d/1kwLrdhTgCdjNrOcSwRI14K3gXnS48xne\",\n \"39\": \"https://drive.google.com/file/d/1V3av9F47t44Nf3jcO12U6OIsjsX-B7L1\",\n \"29\": \"https://drive.google.com/file/d/172dCAUNaaoZX0RHqEi7Ev12eV930LtTa\",\n \"31\": \"https://drive.google.com/file/d/17zzfgMSWBVebWGnpSHKd6g1LFN4vn-YP\",\n \"38\": \"https://drive.google.com/file/d/1xOQvIBlBJ2swTGp78WkCZJUQ-d1F8pVu\",\n \"23\": \"https://drive.google.com/file/d/1X89y_CoTWWjh3BWq0ZgeGydCvg3gMZeJ\",\n \"34\": \"https://drive.google.com/file/d/1VNDhcbA_-Ckjp084hKyl9bwP4E3l9K_2\",\n \"47\": \"https://drive.google.com/file/d/1O8E7haA8WcJIma0iKcvebd4_dlC5Zr7S\",\n \"52\": \"https://drive.google.com/file/d/1o-ateliXHj4TyugOxb9zYYXwrkhFl4FX\",\n \"27\": \"https://drive.google.com/file/d/1aZ1CwNVWJt_OtIBVO-9zv1UUqXTDlM1F\",\n \"26\": \"https://drive.google.com/file/d/1qegrc27hYeECs0KexnEuuG0WQm-8Y9oZ\",\n \"59\": \"https://drive.google.com/file/d/1L9oWKHMTjQreGW_k8rNy7kBQ7c0FuXFm\",\n \"35\": \"https://drive.google.com/file/d/1NewjF092B9KKDBs-dpnZ9dzVl2GAs2LW\",\n \"49\": \"https://drive.google.com/file/d/1TsUrBlr2nxJtH122nKQ_GzdMc0DFFERB\",\n \"41\": \"https://drive.google.com/file/d/1gc2Vo3HZF-Bm_WhZ0zyFedWNfVL2BEol\",\n \"22\": \"https://drive.google.com/file/d/1iPfQeganYriuqHO2e5npUPeuX5VIbhG3\",\n \"36\": \"https://drive.google.com/file/d/1KyNoHRTfGMNR15dCRpcVW74l2z-wVm0V\",\n \"44\": \"https://drive.google.com/file/d/1PDuxwmrD20s54FHQIhXn3ucdFmXSX5kS\",\n \"21\": \"https://drive.google.com/file/d/1ZwO5cCfBJgGpZTIpoi8p2js8zuHT_qxe\",\n \"53\": \"https://drive.google.com/file/d/140NZAuAOoiqrNVqWmF4TPNv6njd_guwE\",\n \"50\": \"https://drive.google.com/file/d/1MBmy7nQi7pMwwIPZHJjB_iuQeO07QWsN\",\n \"54\": \"https://drive.google.com/file/d/1TtIJ-ULYWyv0yUvUVdfTPuBNlBt_j1Yd\",\n \"57\": \"https://drive.google.com/file/d/19V5d3NcR029AhGiRibk2nlTmFNCVGBgO\",\n \"43\": \"https://drive.google.com/file/d/1kLChcxIWZS_kHLEHThLcm7ekcgwYP0jF\",\n \"42\": \"https://drive.google.com/file/d/1HKW3C1B5vFYUuXmFieMKYAfq4CwtnEZ_\",\n \"48\": \"https://drive.google.com/file/d/1EJGd47XpWZDXJKWU0CGp84Hm7K47GNVt\",\n \"40\": \"https://drive.google.com/file/d/1Fr4dVKdOvth_O-Td8PTwgNGzZz8ridAl\",\n \"58\": \"https://drive.google.com/file/d/1lUFpiwE7ISzLbowHvCtEUj4sfG4w0Gst\",\n \"51\": \"https://drive.google.com/file/d/1V6fOoKgA8QSTJYWPP5GVHz8WFAQIRLNB\",\n \"45\": \"https://drive.google.com/file/d/1hSrPOwyxFEth4GWWN1e4BjBftmnKa8px\",\n \"46\": \"https://drive.google.com/file/d/1jCynzDt1r0EISpwcrFuk3RlKWHM9u7Mj\",\n \"55\": \"https://drive.google.com/file/d/1Db01f4I_Xn8Bs9piQgZU59ZWAeC2MaQm\",\n \"56\": \"https://drive.google.com/file/d/1NxVfwIxm6FVVR1XnxQNMWWbQEVX66cQm\",\n \"61\": \"https://docs.google.com/document/d/1eAaZJAqjXMZ2VvG_r04EGtn6EGcYycofdNUkDHEA8vY\"\n}" }, { "path": "backend/tests/daily/connectors/google_drive/test_admin_oauth.py", "content": "from collections.abc import Callable\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom tests.daily.connectors.google_drive.consts_and_utils import _pick\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FOLDER_3_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_MY_DRIVE_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_expected_docs_in_retrieved_docs,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_hierarchy_nodes_match_expected,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_shared_drives,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import PILL_FOLDER_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n RESTRICTED_ACCESS_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FOLDER_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_FOLDER_ID,\n)\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_all(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_all\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=True,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n my_drive_emails=None,\n shared_drive_urls=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n ADMIN_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n + SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + SHARED_DRIVE_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + SECTIONS_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.update(\n _pick(\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n ADMIN_MY_DRIVE_ID,\n PILL_FOLDER_ID,\n RESTRICTED_ACCESS_FOLDER_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n FOLDER_3_ID,\n )\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_shared_drives_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_shared_drives_only\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n my_drive_emails=None,\n shared_drive_urls=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + SHARED_DRIVE_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + SECTIONS_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.update(\n _pick(\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n RESTRICTED_ACCESS_FOLDER_ID,\n )\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_my_drives_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_my_drives_only\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=True,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n my_drive_emails=None,\n shared_drive_urls=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = ADMIN_FILE_IDS + ADMIN_FOLDER_3_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = _pick(\n FOLDER_3_ID,\n ADMIN_MY_DRIVE_ID,\n PILL_FOLDER_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_drive_one_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_drive_one_only\")\n drive_urls = [SHARED_DRIVE_1_URL]\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n my_drive_emails=None,\n shared_drive_urls=\",\".join([str(url) for url in drive_urls]),\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=False,\n include_restricted_folder=False,\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_folder_and_shared_drive(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_folder_and_shared_drive\")\n drive_urls = [SHARED_DRIVE_1_URL]\n folder_urls = [FOLDER_2_URL]\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=\",\".join([str(url) for url in folder_urls]),\n my_drive_emails=None,\n shared_drive_urls=\",\".join([str(url) for url in drive_urls]),\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.pop(SECTIONS_FOLDER_ID, None)\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_folders_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_folders_only\")\n folder_urls = [\n FOLDER_1_2_URL,\n FOLDER_2_1_URL,\n FOLDER_2_2_URL,\n FOLDER_3_URL,\n ]\n shared_drive_urls = [\n FOLDER_1_1_URL,\n ]\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=\",\".join([str(url) for url in folder_urls]),\n my_drive_emails=None,\n shared_drive_urls=\",\".join([str(url) for url in shared_drive_urls]),\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.pop(SECTIONS_FOLDER_ID, None)\n expected_nodes.update(_pick(ADMIN_MY_DRIVE_ID, FOLDER_3_ID))\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_personal_folders_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_personal_folders_only\")\n folder_urls = [\n FOLDER_3_URL,\n ]\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=\",\".join([str(url) for url in folder_urls]),\n my_drive_emails=None,\n shared_drive_urls=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = ADMIN_FOLDER_3_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = _pick(FOLDER_3_ID, ADMIN_MY_DRIVE_ID)\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n )\n" }, { "path": "backend/tests/daily/connectors/google_drive/test_drive_perm_sync.py", "content": "import copy\nimport json\nimport os\nfrom collections import defaultdict\nfrom collections.abc import Callable\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nfrom ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync\nfrom ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import SortOrder\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom tests.daily.connectors.google_drive.consts_and_utils import _pick\nfrom tests.daily.connectors.google_drive.consts_and_utils import ACCESS_MAPPING\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_MY_DRIVE_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_hierarchy_nodes_match_expected,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n EXTERNAL_SHARED_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n FOLDER_3_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_shared_drives,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ACCESS_MAPPING,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PILL_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import PUBLIC_RANGE\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n RESTRICTED_ACCESS_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_DRIVE_B_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_MY_DRIVE_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_2_MY_DRIVE,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_3_MY_DRIVE_ID,\n)\n\n\ndef _build_connector(\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> GoogleDriveConnector:\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=True,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n # don't need this anymore, it's been called in the factory\n connector.load_credentials = MagicMock() # type: ignore\n return connector\n\n\ndef test_gdrive_perm_sync_with_real_data(\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n enable_ee: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test gdrive_doc_sync and gdrive_group_sync with real data from the test drive.\n\n This test uses the real connector to make actual API calls to Google Drive\n and verifies the permission structure returned.\n \"\"\"\n # Create a mock cc_pair that will use our real connector\n mock_cc_pair = MagicMock(spec=ConnectorCredentialPair)\n mock_cc_pair.connector = MagicMock()\n mock_cc_pair.connector.connector_specific_config = {}\n mock_cc_pair.credential_id = 1\n # Import and use the mock helper\n from onyx.utils.sensitive import make_mock_sensitive_value\n\n mock_cc_pair.credential.credential_json = make_mock_sensitive_value({})\n mock_cc_pair.last_time_perm_sync = None\n mock_cc_pair.last_time_external_group_sync = None\n\n # Create a mock heartbeat\n mock_heartbeat = MagicMock(spec=IndexingHeartbeatInterface)\n mock_heartbeat.should_stop.return_value = False\n\n # Load drive_id_mapping.json\n with open(\n os.path.join(os.path.dirname(__file__), \"drive_id_mapping.json\"), \"r\"\n ) as f:\n drive_id_mapping = json.load(f)\n\n # Invert the mapping to get URL -> ID\n url_to_id_mapping = {url: int(id) for id, url in drive_id_mapping.items()}\n\n # Use the connector directly without mocking Google Drive API calls\n with patch(\n \"ee.onyx.external_permissions.google_drive.doc_sync.GoogleDriveConnector\",\n return_value=_build_connector(google_drive_service_acct_connector_factory),\n ):\n # Call the function under test\n def mock_fetch_all_docs_fn(\n sort_order: SortOrder | None = None, # noqa: ARG001\n ) -> list[DocumentRow]:\n return []\n\n def mock_fetch_all_docs_ids_fn() -> list[str]:\n return []\n\n doc_access_generator = gdrive_doc_sync(\n mock_cc_pair,\n mock_fetch_all_docs_fn,\n mock_fetch_all_docs_ids_fn,\n mock_heartbeat,\n )\n doc_access_list = list(doc_access_generator)\n\n # Verify we got some results\n assert len(doc_access_list) > 0\n print(f\"Found {len(doc_access_list)} documents with permissions\")\n\n # create new connector\n with patch(\n \"ee.onyx.external_permissions.google_drive.group_sync.GoogleDriveConnector\",\n return_value=_build_connector(google_drive_service_acct_connector_factory),\n ):\n external_user_group_generator = gdrive_group_sync(\"test_tenant\", mock_cc_pair)\n external_user_groups = list(external_user_group_generator)\n\n # map group ids to emails\n group_id_to_email_mapping: dict[str, set[str]] = defaultdict(set)\n groups_with_anyone_access: set[str] = set()\n for group in external_user_groups:\n for email in group.user_emails:\n group_id_to_email_mapping[group.id].add(email)\n\n if group.gives_anyone_access:\n groups_with_anyone_access.add(group.id)\n\n # Map documents to their permissions (flattening groups)\n doc_to_email_mapping: dict[str, set[str]] = {}\n doc_to_raw_result_mapping: dict[str, set[str]] = {}\n public_doc_ids: set[str] = set()\n\n for doc_access in doc_access_list:\n if not isinstance(doc_access, DocExternalAccess):\n continue\n doc_id = doc_access.doc_id\n # make sure they are new sets to avoid mutating the original\n doc_to_email_mapping[doc_id] = copy.deepcopy(\n doc_access.external_access.external_user_emails\n )\n doc_to_raw_result_mapping[doc_id] = copy.deepcopy(\n doc_access.external_access.external_user_emails\n )\n\n for group_id in doc_access.external_access.external_user_group_ids:\n doc_to_email_mapping[doc_id].update(group_id_to_email_mapping[group_id])\n doc_to_raw_result_mapping[doc_id].add(group_id)\n\n if doc_access.external_access.is_public:\n public_doc_ids.add(doc_id)\n\n if any(\n group_id in groups_with_anyone_access\n for group_id in doc_access.external_access.external_user_group_ids\n ):\n public_doc_ids.add(doc_id)\n\n # Check permissions based on drive_id_mapping.json and ACCESS_MAPPING\n # For each document URL that exists in our mapping\n checked_files = 0\n for doc_id, emails_with_access in doc_to_email_mapping.items():\n # Skip URLs that aren't in our mapping, we don't want new stuff to interfere\n # with the test.\n if doc_id not in url_to_id_mapping:\n continue\n\n file_numeric_id = url_to_id_mapping.get(doc_id)\n if file_numeric_id is None:\n raise ValueError(f\"File {doc_id} not found in drive_id_mapping.json\")\n\n checked_files += 1\n\n # Check which users should have access to this file according to ACCESS_MAPPING\n expected_users = set()\n for user_email, file_ids in ACCESS_MAPPING.items():\n if file_numeric_id in file_ids:\n expected_users.add(user_email)\n\n # Verify the permissions match\n if file_numeric_id in PUBLIC_RANGE:\n assert (\n doc_id in public_doc_ids\n ), f\"File {doc_id} (ID: {file_numeric_id}) should be public but is not in the public_doc_ids set\"\n else:\n assert expected_users == emails_with_access, (\n f\"File {doc_id} (ID: {file_numeric_id}) should be accessible to users {expected_users} \"\n f\"but is accessible to {emails_with_access}. Raw result: {doc_to_raw_result_mapping[doc_id]} \"\n )\n\n # Verify that we checked every file in ACCESS_MAPPING\n all_expected_files = set()\n for file_ids in ACCESS_MAPPING.values():\n all_expected_files.update(file_ids)\n\n checked_file_ids = {\n url_to_id_mapping[doc_id]\n for doc_id in doc_to_email_mapping\n if doc_id in url_to_id_mapping\n }\n\n assert all_expected_files == checked_file_ids, (\n f\"Not all expected files were checked. \"\n f\"Missing files: {all_expected_files - checked_file_ids}, \"\n f\"Extra files checked: {checked_file_ids - all_expected_files}\"\n )\n\n print(f\"Checked permissions for {checked_files} files from drive_id_mapping.json\")\n\n # Verify hierarchy nodes are returned with correct structure\n # Use include_permissions=True to populate external_access on hierarchy nodes\n hierarchy_connector = _build_connector(google_drive_service_acct_connector_factory)\n output = load_connector_outputs(hierarchy_connector, include_permissions=True)\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.update(\n _pick(\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n TEST_USER_1_MY_DRIVE_ID,\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n TEST_USER_1_DRIVE_B_ID,\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n ADMIN_MY_DRIVE_ID,\n TEST_USER_2_MY_DRIVE,\n TEST_USER_3_MY_DRIVE_ID,\n PILL_FOLDER_ID,\n RESTRICTED_ACCESS_FOLDER_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n EXTERNAL_SHARED_FOLDER_ID,\n FOLDER_3_ID,\n )\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n # Verify the perm sync drives are included in the hierarchy\n # These drives should have external_access set on their hierarchy nodes\n perm_sync_drive_nodes = [\n node\n for node in output.hierarchy_nodes\n if node.raw_node_id\n in {\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n }\n ]\n\n # Verify permissions on perm sync drive hierarchy nodes\n for node in perm_sync_drive_nodes:\n assert (\n node.external_access is not None\n ), f\"Hierarchy node {node.raw_node_id} has no external access\"\n expected_emails = PERM_SYNC_DRIVE_ACCESS_MAPPING.get(node.raw_node_id, set())\n actual_emails = node.external_access.external_user_emails\n assert actual_emails == expected_emails, (\n f\"Permission mismatch for perm sync drive {node.raw_node_id} ({node.display_name}): \"\n f\"expected {expected_emails}, got {actual_emails}\"\n )\n\n print(f\"Verified {len(output.hierarchy_nodes)} hierarchy nodes\")\n" }, { "path": "backend/tests/daily/connectors/google_drive/test_link_visibility_filter.py", "content": "from collections.abc import Iterable\nfrom typing import Any\nfrom unittest.mock import patch\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.google_drive.file_retrieval import has_link_only_permission\nfrom onyx.connectors.google_drive.models import DriveRetrievalStage\nfrom onyx.connectors.google_drive.models import RetrievedDriveFile\n\n\ndef _stub_run_functions(\n func_with_args: Iterable[tuple],\n max_workers: int = 8, # noqa: ARG001\n) -> list[Any]:\n return [func(*args) for func, args in func_with_args]\n\n\ndef _build_retrieved_file(\n permissions: list[dict[str, Any]],\n) -> RetrievedDriveFile:\n return RetrievedDriveFile(\n completion_stage=DriveRetrievalStage.OAUTH_FILES,\n drive_file={\n \"id\": \"file-id\",\n \"name\": \"Test File\",\n \"permissions\": permissions,\n },\n user_email=\"user@example.com\",\n )\n\n\ndef _prepare_connector(exclude: bool) -> GoogleDriveConnector:\n connector = GoogleDriveConnector(\n include_shared_drives=True,\n exclude_domain_link_only=exclude,\n )\n connector._creds = object() # type: ignore[assignment]\n connector._primary_admin_email = \"admin@example.com\"\n return connector\n\n\ndef test_has_link_only_permission_detects_domain_link() -> None:\n file = {\n \"permissions\": [\n {\"type\": \"domain\", \"allowFileDiscovery\": False},\n {\"type\": \"user\", \"emailAddress\": \"user@example.com\"},\n ]\n }\n assert has_link_only_permission(file) is True\n\n\ndef test_has_link_only_permission_detects_anyone_link() -> None:\n file = {\n \"permissions\": [\n {\"type\": \"anyone\", \"allowFileDiscovery\": False},\n ]\n }\n assert has_link_only_permission(file) is True\n\n\ndef test_has_link_only_permission_ignores_other_permissions() -> None:\n file = {\n \"permissions\": [\n {\"type\": \"domain\", \"allowFileDiscovery\": True},\n {\"type\": \"user\", \"emailAddress\": \"user@example.com\"},\n ]\n }\n assert has_link_only_permission(file) is False\n\n\ndef test_connector_skips_link_only_files_when_enabled() -> None:\n connector = _prepare_connector(exclude=True)\n retrieved_file = _build_retrieved_file(\n [{\"type\": \"domain\", \"allowFileDiscovery\": False}]\n )\n\n with (\n patch(\n \"onyx.connectors.google_drive.connector.run_functions_tuples_in_parallel\",\n side_effect=_stub_run_functions,\n ),\n patch(\n \"onyx.connectors.google_drive.connector.convert_drive_item_to_document\"\n ) as convert_mock,\n patch(\n \"onyx.connectors.google_drive.connector.GoogleDriveConnector._get_new_ancestors_for_files\"\n ) as get_new_ancestors_mock,\n ):\n convert_mock.return_value = \"doc\"\n checkpoint = connector.build_dummy_checkpoint()\n results = list(\n connector._convert_retrieved_files_to_documents(\n drive_files_iter=iter([retrieved_file]),\n checkpoint=checkpoint,\n include_permissions=False,\n )\n )\n\n assert results == []\n convert_mock.assert_not_called()\n get_new_ancestors_mock.assert_called_once()\n\n\ndef test_connector_processes_files_when_option_disabled() -> None:\n connector = _prepare_connector(exclude=False)\n retrieved_file = _build_retrieved_file(\n [{\"type\": \"domain\", \"allowFileDiscovery\": False}]\n )\n\n with (\n patch(\n \"onyx.connectors.google_drive.connector.run_functions_tuples_in_parallel\",\n side_effect=_stub_run_functions,\n ),\n patch(\n \"onyx.connectors.google_drive.connector.convert_drive_item_to_document\"\n ) as convert_mock,\n patch(\n \"onyx.connectors.google_drive.connector.GoogleDriveConnector._get_new_ancestors_for_files\"\n ) as get_new_ancestors_mock,\n ):\n convert_mock.return_value = \"doc\"\n checkpoint = connector.build_dummy_checkpoint()\n results = list(\n connector._convert_retrieved_files_to_documents(\n drive_files_iter=iter([retrieved_file]),\n checkpoint=checkpoint,\n include_permissions=False,\n )\n )\n\n assert len(results) == 1\n convert_mock.assert_called_once()\n get_new_ancestors_mock.assert_called_once()\n" }, { "path": "backend/tests/daily/connectors/google_drive/test_map_test_ids.py", "content": "#!/usr/bin/env python\n\nimport json\nimport os\n\nimport pytest\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom tests.daily.connectors.google_drive.conftest import get_credentials_from_env\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import file_name_template\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_FILE_IDS\n\n\ndef generate_test_id_to_drive_id_mapping() -> dict[int, str]:\n \"\"\"\n Generate a mapping from test file IDs to actual Google Drive file IDs.\n\n This is useful for writing tests that need to verify specific files\n are accessible to specific users.\n\n Returns:\n dict: Mapping from test file ID (int) to Google Drive file ID (str)\n \"\"\"\n # Set up the connector with real credentials\n connector = GoogleDriveConnector(\n include_shared_drives=True,\n include_my_drives=True,\n include_files_shared_with_me=False,\n )\n\n # Load credentials\n connector.load_credentials(get_credentials_from_env(email=ADMIN_EMAIL, oauth=False))\n\n # Get all documents from the connector\n docs = load_connector_outputs(connector).documents\n\n # Create a mapping from test file ID to actual Drive file ID\n test_id_to_drive_id = {}\n\n # Process all documents retrieved from Drive\n for doc in docs:\n # Check if this document's name matches our test file naming pattern (file_X.txt)\n if not doc.semantic_identifier.startswith(\n file_name_template.format(\"\").split(\"_\")[0]\n ):\n continue\n\n try:\n # Extract the test file ID from the filename (file_X.txt -> X)\n file_id_str = doc.semantic_identifier.split(\"_\")[1].split(\".\")[0]\n test_file_id = int(file_id_str)\n\n # Store the mapping from test ID to actual Drive ID\n # Extract Drive ID from document URL\n test_id_to_drive_id[test_file_id] = doc.id\n except (ValueError, IndexError):\n # Skip files that don't follow our naming convention\n continue\n\n # Print the mapping for all defined test file ID ranges\n all_test_ranges = {\n \"ADMIN_FILE_IDS\": ADMIN_FILE_IDS,\n \"TEST_USER_1_FILE_IDS\": TEST_USER_1_FILE_IDS,\n \"TEST_USER_2_FILE_IDS\": TEST_USER_2_FILE_IDS,\n \"TEST_USER_3_FILE_IDS\": TEST_USER_3_FILE_IDS,\n \"SHARED_DRIVE_1_FILE_IDS\": SHARED_DRIVE_1_FILE_IDS,\n \"SHARED_DRIVE_2_FILE_IDS\": SHARED_DRIVE_2_FILE_IDS,\n \"FOLDER_1_FILE_IDS\": FOLDER_1_FILE_IDS,\n \"FOLDER_1_1_FILE_IDS\": FOLDER_1_1_FILE_IDS,\n \"FOLDER_1_2_FILE_IDS\": FOLDER_1_2_FILE_IDS,\n \"FOLDER_2_FILE_IDS\": FOLDER_2_FILE_IDS,\n \"FOLDER_2_1_FILE_IDS\": FOLDER_2_1_FILE_IDS,\n \"FOLDER_2_2_FILE_IDS\": FOLDER_2_2_FILE_IDS,\n \"FOLDER_3_FILE_IDS\": FOLDER_3_FILE_IDS,\n }\n\n # Print the mapping for each test range\n for range_name, file_ids in all_test_ranges.items():\n print(f\"\\n{range_name}:\")\n for test_id in file_ids:\n drive_id = test_id_to_drive_id.get(test_id, \"NOT_FOUND\")\n print(f\" {test_id} -> {drive_id}\")\n\n return test_id_to_drive_id\n\n\n@pytest.mark.skipif(\n not os.getenv(\"RUN_MANUAL_TESTS\"),\n reason=\"This test maps test IDs to actual Google Drive IDs. Set RUN_MANUAL_TESTS=1 to run.\",\n)\ndef test_generate_drive_id_mapping() -> None:\n \"\"\"Test to generate mapping from test IDs to actual Google Drive IDs.\n\n This test is skipped by default as it requires real Google Drive credentials\n and is primarily used to generate mappings for other tests.\n\n Run with:\n\n RUN_MANUAL_TESTS=true pytest -xvs tests/daily/connectors/google_drive/test_map_test_ids.py::test_generate_drive_id_mapping\n \"\"\"\n mapping = generate_test_id_to_drive_id_mapping()\n assert mapping, \"Failed to generate any test ID to drive ID mappings\"\n\n # Write the mapping to a JSON file\n output_dir = os.path.dirname(os.path.abspath(__file__))\n mapping_file = os.path.join(output_dir, \"drive_id_mapping.json\")\n\n # Convert int keys to strings for JSON compatibility\n json_mapping = {str(k): v for k, v in mapping.items()}\n\n # Write the mapping to a JSON file\n with open(mapping_file, \"w\") as f:\n json.dump(json_mapping, f, indent=2)\n\n print(f\"\\nMapping written to: {mapping_file}\")\n raise RuntimeError(\"Mapping written to file, test complete\")\n" }, { "path": "backend/tests/daily/connectors/google_drive/test_sections.py", "content": "from collections.abc import Callable\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs\nfrom tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FOLDER_URL\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_google_drive_sections(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n oauth_connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=SECTIONS_FOLDER_URL,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n service_acct_connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=SECTIONS_FOLDER_URL,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n for connector in [oauth_connector, service_acct_connector]:\n output = load_connector_outputs(connector)\n retrieved_docs = output.documents\n\n # Verify we got the 1 doc with sections\n assert len(retrieved_docs) == 1\n\n # Verify each section has the expected structure\n doc = retrieved_docs[0]\n assert len(doc.sections) == 5\n\n header_section = doc.sections[0]\n assert header_section.text == \"Title\\nThis is a Google Doc with sections\"\n assert header_section.link is not None\n assert header_section.link.endswith(\n \"?tab=t.0#heading=h.hfjc17k6qwzt\"\n ) or header_section.link.endswith(\"?tab=t.0#heading=h.hfjc17k6qwzt\")\n\n section_1 = doc.sections[1]\n assert section_1.text == \"Section 1\\nSection 1 content\"\n assert section_1.link is not None\n assert section_1.link.endswith(\"?tab=t.0#heading=h.8slfx752a3g5\")\n\n section_2 = doc.sections[2]\n assert section_2.text == \"Sub-Section 1-1\\nSub-Section 1-1 content\"\n assert section_2.link is not None\n assert section_2.link.endswith(\"?tab=t.0#heading=h.4kj3ayade1bp\")\n\n section_3 = doc.sections[3]\n assert section_3.text == \"Sub-Section 1-2\\nSub-Section 1-2 content\"\n assert section_3.link is not None\n assert section_3.link.endswith(\"?tab=t.0#heading=h.pm6wrpzgk69l\")\n\n section_4 = doc.sections[4]\n assert section_4.text == \"Section 2\\nSection 2 content\"\n assert section_4.link is not None\n assert section_4.link.endswith(\"?tab=t.0#heading=h.2m0s9youe2k9\")\n" }, { "path": "backend/tests/daily/connectors/google_drive/test_service_acct.py", "content": "from collections.abc import Callable\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom urllib.parse import urlparse\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom tests.daily.connectors.google_drive.consts_and_utils import _pick\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FOLDER_3_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_MY_DRIVE_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_expected_docs_in_retrieved_docs,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_hierarchy_nodes_match_expected,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n EXTERNAL_SHARED_DOC_SINGLETON,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n EXTERNAL_SHARED_DOCS_IN_FOLDER,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n EXTERNAL_SHARED_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n EXTERNAL_SHARED_FOLDER_URL,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_2_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_2_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_shared_drives,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import id_to_name\nfrom tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n MISC_SHARED_DRIVE_FNAMES,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n PILL_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n RESTRICTED_ACCESS_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n RESTRICTED_ACCESS_FOLDER_URL,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import SECTIONS_FOLDER_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_DRIVE_B_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_EXTRA_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_1_MY_DRIVE_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_2_MY_DRIVE,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_3_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n TEST_USER_3_MY_DRIVE_ID,\n)\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_all(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_all\")\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=True,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n # Should get everything\n expected_file_ids = (\n ADMIN_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n + TEST_USER_1_FILE_IDS\n + TEST_USER_2_FILE_IDS\n + TEST_USER_3_FILE_IDS\n + SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + SHARED_DRIVE_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + SECTIONS_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.update(\n _pick(\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n TEST_USER_1_MY_DRIVE_ID,\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n TEST_USER_1_DRIVE_B_ID,\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n ADMIN_MY_DRIVE_ID,\n TEST_USER_2_MY_DRIVE,\n TEST_USER_3_MY_DRIVE_ID,\n PILL_FOLDER_ID,\n RESTRICTED_ACCESS_FOLDER_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n EXTERNAL_SHARED_FOLDER_ID,\n FOLDER_3_ID,\n )\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_shared_drives_only_with_size_threshold(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_shared_drives_only_with_size_threshold\")\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n\n # this threshold will skip one file\n connector.size_threshold = 16384\n\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + SHARED_DRIVE_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + SECTIONS_FILE_IDS\n )\n\n expected_file_names = {id_to_name(file_id) for file_id in expected_file_ids}\n expected_file_names.update(MISC_SHARED_DRIVE_FNAMES)\n retrieved_file_names = {doc.semantic_identifier for doc in output.documents}\n for name in expected_file_names - retrieved_file_names:\n print(f\"expected but did not retrieve: {name}\")\n for name in retrieved_file_names - expected_file_names:\n print(f\"retrieved but did not expect: {name}\")\n\n # 2 extra files from shared drive owned by non-admin and not shared with admin\n # TODO: added a file in a \"restricted\" folder, which the connector sometimes succeeds at finding\n # and adding. Specifically, our shared drive retrieval logic currently assumes that\n # \"having access to a shared drive\" means that the connector has access to all files in the shared drive.\n # therefore when a user successfully retrieves a shared drive, we mark it as \"done\". If that user's\n # access is restricted for a folder in the shared drive, the connector will not retrieve that folder.\n # If instead someone with FULL access to the shared drive retrieves it, the connector will retrieve\n # the folder and all its files. There is currently no consistency to the order of assignment of users\n # to shared drives, so this is a heisenbug. When we guarantee that restricted folders are retrieved,\n # we can change this to 52\n assert len(output.documents) == 50 or len(output.documents) == 51\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_shared_drives_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_shared_drives_only\")\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=True,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n\n output = load_connector_outputs(connector)\n\n # Should only get shared drives\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + SHARED_DRIVE_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + SECTIONS_FILE_IDS\n )\n\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n # 2 extra files from shared drive owned by non-admin and not shared with admin\n # another one flaky for unknown reasons\n # TODO: switch to 54 when restricted access issue is resolved\n assert len(output.documents) == 51 or len(output.documents) == 52\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.update(\n _pick(\n PERM_SYNC_DRIVE_ADMIN_ONLY_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_A_ID,\n PERM_SYNC_DRIVE_ADMIN_AND_USER_1_B_ID,\n TEST_USER_1_DRIVE_B_ID,\n TEST_USER_1_DRIVE_B_FOLDER_ID,\n TEST_USER_1_EXTRA_DRIVE_1_ID,\n TEST_USER_1_EXTRA_DRIVE_2_ID,\n RESTRICTED_ACCESS_FOLDER_ID,\n )\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_include_my_drives_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_include_my_drives_only\")\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=True,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n # Should only get everyone's My Drives\n expected_file_ids = (\n ADMIN_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n + TEST_USER_1_FILE_IDS\n + TEST_USER_2_FILE_IDS\n + TEST_USER_3_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = _pick(\n FOLDER_3_ID,\n ADMIN_MY_DRIVE_ID,\n TEST_USER_1_MY_DRIVE_ID,\n TEST_USER_1_MY_DRIVE_FOLDER_ID,\n TEST_USER_2_MY_DRIVE,\n TEST_USER_3_MY_DRIVE_ID,\n PILL_FOLDER_ID,\n TEST_USER_1_EXTRA_FOLDER_ID,\n EXTERNAL_SHARED_FOLDER_ID,\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_drive_one_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_drive_one_only\")\n urls = [SHARED_DRIVE_1_URL]\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=\",\".join([str(url) for url in urls]),\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n # We ignore shared_drive_urls if include_shared_drives is False\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=False,\n include_restricted_folder=False,\n )\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_folder_and_shared_drive(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_folder_and_shared_drive\")\n drive_urls = [SHARED_DRIVE_1_URL]\n folder_urls = [FOLDER_2_URL]\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_drive_urls=\",\".join([str(url) for url in drive_urls]),\n shared_folder_urls=\",\".join([str(url) for url in folder_urls]),\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n # Should get everything except for the top level files in drive 2\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + FOLDER_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.pop(SECTIONS_FOLDER_ID, None)\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n ignorable_node_ids={RESTRICTED_ACCESS_FOLDER_ID},\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_folders_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_folders_only\")\n folder_urls = [\n FOLDER_1_2_URL,\n FOLDER_2_1_URL,\n FOLDER_2_2_URL,\n FOLDER_3_URL,\n ]\n # This should get converted to a drive request and spit out a warning in the logs\n shared_drive_urls = [\n FOLDER_1_1_URL,\n ]\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_drive_urls=\",\".join([str(url) for url in shared_drive_urls]),\n shared_folder_urls=\",\".join([str(url) for url in folder_urls]),\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + FOLDER_2_1_FILE_IDS\n + FOLDER_2_2_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n expected_nodes = get_expected_hierarchy_for_shared_drives(\n include_drive_1=True,\n include_drive_2=True,\n include_restricted_folder=False,\n )\n expected_nodes.pop(SECTIONS_FOLDER_ID, None)\n expected_nodes.update(_pick(ADMIN_MY_DRIVE_ID, FOLDER_3_ID))\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=expected_nodes,\n )\n\n\ndef test_shared_folder_owned_by_external_user(\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_shared_folder_owned_by_external_user\")\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_drive_urls=None,\n shared_folder_urls=EXTERNAL_SHARED_FOLDER_URL,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_docs = EXTERNAL_SHARED_DOCS_IN_FOLDER\n\n assert len(output.documents) == len(expected_docs) # 1 for now\n assert expected_docs[0] in output.documents[0].id\n\n\ndef test_shared_with_me(\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_shared_with_me\")\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=True,\n include_files_shared_with_me=True,\n shared_drive_urls=None,\n shared_folder_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n print(output.documents)\n\n expected_file_ids = (\n ADMIN_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n + TEST_USER_1_FILE_IDS\n + TEST_USER_2_FILE_IDS\n + TEST_USER_3_FILE_IDS\n )\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n retrieved_ids = {urlparse(doc.id).path.split(\"/\")[-1] for doc in output.documents}\n for id in retrieved_ids:\n print(id)\n\n assert EXTERNAL_SHARED_DOC_SINGLETON.split(\"/\")[-1] in retrieved_ids\n assert EXTERNAL_SHARED_DOCS_IN_FOLDER[0].split(\"/\")[-1] in retrieved_ids\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_specific_emails(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_specific_emails\")\n my_drive_emails = [\n TEST_USER_1_EMAIL,\n TEST_USER_3_EMAIL,\n ]\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=\",\".join([str(email) for email in my_drive_emails]),\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = TEST_USER_1_FILE_IDS + TEST_USER_3_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef get_specific_folders_in_my_drive(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning get_specific_folders_in_my_drive\")\n folder_urls = [\n FOLDER_3_URL,\n ]\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=\",\".join([str(url) for url in folder_urls]),\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = ADMIN_FOLDER_3_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_specific_user_emails_restricted_folder(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_specific_user_emails_restricted_folder\")\n\n # Test with admin email - should get 1 doc\n admin_connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=RESTRICTED_ACCESS_FOLDER_URL,\n shared_drive_urls=None,\n my_drive_emails=None,\n specific_user_emails=ADMIN_EMAIL,\n )\n admin_output = load_connector_outputs(admin_connector)\n assert len(admin_output.documents) == 1\n\n # Test with test users - should get 0 docs\n test_users = [TEST_USER_1_EMAIL, TEST_USER_2_EMAIL, TEST_USER_3_EMAIL]\n test_connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=False,\n include_files_shared_with_me=False,\n shared_folder_urls=RESTRICTED_ACCESS_FOLDER_URL,\n shared_drive_urls=None,\n my_drive_emails=None,\n specific_user_emails=\",\".join(test_users),\n )\n test_output = load_connector_outputs(test_connector)\n assert len(test_output.documents) == 0\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_specific_user_email_shared_with_me(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_service_acct_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_specific_user_email_shared_with_me\")\n\n # Test with admin email - should get 1 doc\n connector = google_drive_service_acct_connector_factory(\n primary_admin_email=ADMIN_EMAIL,\n include_shared_drives=False,\n include_my_drives=True,\n include_files_shared_with_me=False, # This is what is set in the UI unfortunately\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n specific_user_emails=TEST_USER_1_EMAIL,\n )\n output = load_connector_outputs(connector)\n expected = [id_to_name(file_id) for file_id in TEST_USER_1_FILE_IDS]\n expected += [\"private_file\", \"shared_file\"] # in My Drive\n expected += [\"read only users can't download\"] # Shared with me\n\n expected += [id_to_name(file_id) for file_id in [0, 1] + ADMIN_FOLDER_3_FILE_IDS]\n\n # these are in shared drives\n # expected += ['perm_sync_doc_0ACOrCU1EMD1hUk9PVA_ab63b976-effb-49af-84e7-423d17a17dd7']\n # expected += ['file_22.txt'] # Shared drive\n\n doc_titles = set(doc.semantic_identifier for doc in output.documents)\n assert doc_titles == set(expected)\n" }, { "path": "backend/tests/daily/connectors/google_drive/test_user_1_oauth.py", "content": "from collections.abc import Callable\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nfrom onyx.connectors.google_drive.connector import GoogleDriveConnector\nfrom onyx.connectors.models import Document\nfrom tests.daily.connectors.google_drive.consts_and_utils import _clear_parents\nfrom tests.daily.connectors.google_drive.consts_and_utils import _pick\nfrom tests.daily.connectors.google_drive.consts_and_utils import ADMIN_FOLDER_3_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_expected_docs_in_retrieved_docs,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n assert_hierarchy_nodes_match_expected,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n DONWLOAD_REVOKED_FILE_ID,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_1_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_2_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_1_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import FOLDER_3_URL\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_test_user_1,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_test_user_1_my_drive_only,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_test_user_1_shared_drives_only,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import (\n get_expected_hierarchy_for_test_user_1_shared_with_me_only,\n)\nfrom tests.daily.connectors.google_drive.consts_and_utils import load_connector_outputs\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_FILE_IDS\nfrom tests.daily.connectors.google_drive.consts_and_utils import SHARED_DRIVE_1_ID\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_EMAIL\nfrom tests.daily.connectors.google_drive.consts_and_utils import TEST_USER_1_FILE_IDS\nfrom tests.daily.connectors.utils import ConnectorOutput\n\n\ndef _check_for_error(\n output: ConnectorOutput,\n expected_file_ids: list[int],\n) -> list[Document]:\n retrieved_docs = output.documents\n retrieved_failures = output.failures\n assert len(retrieved_failures) <= 1\n\n if len(retrieved_failures) == 1:\n fail_msg = retrieved_failures[0].failure_message\n assert \"HttpError 403\" in fail_msg\n assert f\"file_{DONWLOAD_REVOKED_FILE_ID}.txt\" in fail_msg\n\n expected_file_ids.remove(DONWLOAD_REVOKED_FILE_ID)\n return retrieved_docs\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_all(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_all\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=TEST_USER_1_EMAIL,\n include_files_shared_with_me=True,\n include_shared_drives=True,\n include_my_drives=True,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n TEST_USER_1_FILE_IDS\n + SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n + ADMIN_FOLDER_3_FILE_IDS\n + list(range(0, 2))\n )\n\n retrieved_docs = _check_for_error(output, expected_file_ids)\n\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=retrieved_docs,\n expected_file_ids=expected_file_ids,\n )\n\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=get_expected_hierarchy_for_test_user_1(),\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_shared_drives_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_shared_drives_only\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=TEST_USER_1_EMAIL,\n include_files_shared_with_me=False,\n include_shared_drives=True,\n include_my_drives=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = (\n SHARED_DRIVE_1_FILE_IDS\n + FOLDER_1_FILE_IDS\n + FOLDER_1_1_FILE_IDS\n + FOLDER_1_2_FILE_IDS\n )\n\n retrieved_docs = _check_for_error(output, expected_file_ids)\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=retrieved_docs,\n expected_file_ids=expected_file_ids,\n )\n\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=get_expected_hierarchy_for_test_user_1_shared_drives_only(),\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_shared_with_me_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_shared_with_me_only\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=TEST_USER_1_EMAIL,\n include_files_shared_with_me=True,\n include_shared_drives=False,\n include_my_drives=False,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = ADMIN_FOLDER_3_FILE_IDS + list(range(0, 2))\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=get_expected_hierarchy_for_test_user_1_shared_with_me_only(),\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_my_drive_only(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_my_drive_only\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=TEST_USER_1_EMAIL,\n include_files_shared_with_me=False,\n include_shared_drives=False,\n include_my_drives=True,\n shared_folder_urls=None,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = TEST_USER_1_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=get_expected_hierarchy_for_test_user_1_my_drive_only(),\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_shared_my_drive_folder(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_shared_my_drive_folder\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=TEST_USER_1_EMAIL,\n include_files_shared_with_me=False,\n include_shared_drives=False,\n include_my_drives=True,\n shared_folder_urls=FOLDER_3_URL,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = ADMIN_FOLDER_3_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=_clear_parents(_pick(FOLDER_3_ID), FOLDER_3_ID),\n )\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_shared_drive_folder(\n mock_get_api_key: MagicMock, # noqa: ARG001\n google_drive_oauth_uploaded_connector_factory: Callable[..., GoogleDriveConnector],\n) -> None:\n print(\"\\n\\nRunning test_shared_drive_folder\")\n connector = google_drive_oauth_uploaded_connector_factory(\n primary_admin_email=TEST_USER_1_EMAIL,\n include_files_shared_with_me=False,\n include_shared_drives=False,\n include_my_drives=True,\n shared_folder_urls=FOLDER_1_URL,\n shared_drive_urls=None,\n my_drive_emails=None,\n )\n output = load_connector_outputs(connector)\n\n expected_file_ids = FOLDER_1_FILE_IDS + FOLDER_1_1_FILE_IDS + FOLDER_1_2_FILE_IDS\n assert_expected_docs_in_retrieved_docs(\n retrieved_docs=output.documents,\n expected_file_ids=expected_file_ids,\n )\n\n assert_hierarchy_nodes_match_expected(\n retrieved_nodes=output.hierarchy_nodes,\n expected_nodes=_pick(\n SHARED_DRIVE_1_ID, FOLDER_1_ID, FOLDER_1_1_ID, FOLDER_1_2_ID\n ),\n )\n" }, { "path": "backend/tests/daily/connectors/highspot/test_highspot_connector.py", "content": "import json\nimport os\nimport time\nfrom datetime import datetime\nfrom pathlib import Path\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.highspot.connector import HighspotConnector\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\n\n\ndef load_test_data(file_name: str = \"test_highspot_data.json\") -> dict:\n \"\"\"Load test data from JSON file.\"\"\"\n current_dir = Path(__file__).parent\n with open(current_dir / file_name, \"r\") as f:\n return json.load(f)\n\n\n@pytest.fixture\ndef highspot_connector() -> HighspotConnector:\n \"\"\"Create a Highspot connector with credentials from environment variables.\"\"\"\n # Check if required environment variables are set\n if not os.environ.get(\"HIGHSPOT_KEY\") or not os.environ.get(\"HIGHSPOT_SECRET\"):\n pytest.fail(\"HIGHSPOT_KEY or HIGHSPOT_SECRET environment variables not set\")\n\n connector = HighspotConnector(\n spot_names=[\"Test content\"], # Use specific spot name instead of empty list\n batch_size=10, # Smaller batch size for testing\n )\n connector.load_credentials(\n {\n \"highspot_key\": os.environ[\"HIGHSPOT_KEY\"],\n \"highspot_secret\": os.environ[\"HIGHSPOT_SECRET\"],\n \"highspot_url\": os.environ.get(\n \"HIGHSPOT_URL\", \"https://api-su2.highspot.com/v1.0/\"\n ),\n }\n )\n return connector\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_highspot_connector_basic(\n mock_get_api_key: MagicMock, # noqa: ARG001\n highspot_connector: HighspotConnector,\n) -> None:\n \"\"\"Test basic functionality of the Highspot connector.\"\"\"\n all_docs: list[Document] = []\n test_data = load_test_data()\n target_test_doc_id = test_data.get(\"target_doc_id\")\n target_test_doc: Document | None = None\n\n # Test loading documents\n for doc_batch in highspot_connector.poll_source(0, time.time()):\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs.append(doc)\n if doc.id == f\"HIGHSPOT_{target_test_doc_id}\":\n target_test_doc = doc\n\n # Verify documents were loaded\n assert len(all_docs) > 0\n\n # If we have a specific test document ID, validate it\n if target_test_doc_id and target_test_doc is not None:\n assert target_test_doc.semantic_identifier == test_data.get(\n \"semantic_identifier\"\n )\n assert target_test_doc.source == DocumentSource.HIGHSPOT\n assert target_test_doc.metadata is not None\n\n assert len(target_test_doc.sections) == 1\n section = target_test_doc.sections[0]\n assert section.link is not None\n # Only check if content exists, as exact content might change\n assert section.text is not None\n assert len(section.text) > 0\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_highspot_connector_slim(\n mock_get_api_key: MagicMock, # noqa: ARG001\n highspot_connector: HighspotConnector,\n) -> None:\n \"\"\"Test slim document retrieval.\"\"\"\n # Get all doc IDs from the full connector\n all_full_doc_ids = set()\n for doc_batch in highspot_connector.load_from_state():\n all_full_doc_ids.update(\n [doc.id for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # Get all doc IDs from the slim connector\n all_slim_doc_ids = set()\n for slim_doc_batch in highspot_connector.retrieve_all_slim_docs_perm_sync():\n all_slim_doc_ids.update(\n [doc.id for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # The set of full doc IDs should be a subset of the slim doc IDs\n assert all_full_doc_ids.issubset(all_slim_doc_ids)\n # Make sure we actually got some documents\n assert len(all_slim_doc_ids) > 0\n\n\n\"\"\"This test might fail because of how Highspot handles changes to the document's\n\"updated at\" property. It is marked as expected to fail until we can confirm the behavior.\"\"\"\n\n\n@pytest.mark.xfail(reason=\"Highspot is not returning updated documents as expected.\")\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_highspot_connector_poll_source(\n mock_get_api_key: MagicMock, # noqa: ARG001\n highspot_connector: HighspotConnector,\n) -> None:\n \"\"\"Test poll_source functionality with date range filtering.\"\"\"\n # Define date range: April 3, 2025 to April 4, 2025\n start_date = datetime(2025, 4, 3, 0, 0, 0)\n end_date = datetime(2025, 4, 4, 23, 59, 59)\n\n # Convert to seconds since Unix epoch\n start_time = int(time.mktime(start_date.timetuple()))\n end_time = int(time.mktime(end_date.timetuple()))\n\n # Load test data for assertions\n test_data = load_test_data()\n poll_source_data = test_data.get(\"poll_source\", {})\n target_doc_id = poll_source_data.get(\"target_doc_id\")\n\n # Call poll_source with date range\n all_docs: list[Document] = []\n target_doc: Document | None = None\n\n for doc_batch in highspot_connector.poll_source(start_time, end_time):\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs.append(doc)\n if doc.id == f\"HIGHSPOT_{target_doc_id}\":\n target_doc = doc\n\n # Verify documents were loaded\n assert len(all_docs) > 0\n\n # Verify the specific test document was found and has correct properties\n assert target_doc is not None\n assert target_doc.semantic_identifier == poll_source_data.get(\"semantic_identifier\")\n assert target_doc.source == DocumentSource.HIGHSPOT\n assert target_doc.metadata is not None\n\n # Verify sections\n assert len(target_doc.sections) == 1\n section = target_doc.sections[0]\n assert section.link == poll_source_data.get(\"link\")\n assert section.text is not None\n assert len(section.text) > 0\n\n\ndef test_highspot_connector_validate_credentials(\n highspot_connector: HighspotConnector,\n) -> None:\n \"\"\"Test credential validation.\"\"\"\n assert highspot_connector.validate_credentials() is True\n" }, { "path": "backend/tests/daily/connectors/highspot/test_highspot_data.json", "content": "{\n \"target_doc_id\": \"67cd8eb35d3ee0487de2e704\",\n \"semantic_identifier\": \"Highspot in Action _ Salesforce Integration\",\n \"link\": \"https://www.highspot.com/items/67cd8eb35d3ee0487de2e704\",\n \"poll_source\": {\n \"target_doc_id\":\"67efb452c3f40bcca2b48ca5\",\n \"semantic_identifier\":\"Introduction to Intelligent Agents\",\n \"link\":\"https://www.highspot.com/items/67efb452c3f40bcca2b48ca5\"\n }\n}\n" }, { "path": "backend/tests/daily/connectors/hubspot/test_hubspot_connector.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onyx.connectors.hubspot.connector import AVAILABLE_OBJECT_TYPES\nfrom onyx.connectors.hubspot.connector import HubSpotConnector\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\n\n\nclass TestHubSpotConnector:\n \"\"\"Test HubSpot connector functionality using real API calls.\"\"\"\n\n @pytest.fixture\n def connector(self) -> HubSpotConnector:\n \"\"\"Create a HubSpot connector instance.\"\"\"\n return HubSpotConnector(batch_size=10)\n\n @pytest.fixture\n def credentials(self) -> dict[str, Any]:\n \"\"\"Provide test credentials.\"\"\"\n return {\"hubspot_access_token\": os.environ[\"HUBSPOT_ACCESS_TOKEN\"]}\n\n def test_credentials_properties_raise_exception_when_none(self) -> None:\n \"\"\"Test that access_token and portal_id properties raise exceptions when not set.\"\"\"\n connector = HubSpotConnector()\n\n # access_token should raise exception when not set\n with pytest.raises(ConnectorMissingCredentialError) as exc_info:\n _ = connector.access_token\n assert \"HubSpot access token not set\" in str(exc_info.value)\n\n # portal_id should raise exception when not set\n with pytest.raises(ConnectorMissingCredentialError) as exc_info:\n _ = connector.portal_id\n assert \"HubSpot portal ID not set\" in str(exc_info.value)\n\n def test_load_credentials(\n self, connector: HubSpotConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that credentials are loaded correctly.\"\"\"\n result = connector.load_credentials(credentials)\n\n assert result is None # Should return None on success\n assert connector.access_token == credentials[\"hubspot_access_token\"]\n assert connector.portal_id is not None\n assert isinstance(connector.portal_id, str)\n\n def test_load_from_state_basic_functionality(\n self, connector: HubSpotConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test basic load_from_state functionality.\"\"\"\n connector.load_credentials(credentials)\n\n # Get first batch of documents\n document_batches = connector.load_from_state()\n first_batch = next(document_batches, None)\n\n # Should have at least some documents\n assert first_batch is not None\n assert isinstance(first_batch, list)\n assert len(first_batch) > 0\n\n # Check document structure\n doc = first_batch[0]\n assert isinstance(doc, Document)\n assert doc.id.startswith(\"hubspot_\")\n assert doc.source.value == \"hubspot\"\n assert doc.semantic_identifier is not None\n assert doc.doc_updated_at is not None\n assert isinstance(doc.metadata, dict)\n assert \"object_type\" in doc.metadata\n assert doc.metadata[\"object_type\"] in [\"ticket\", \"company\", \"deal\", \"contact\"]\n\n # Check sections\n assert len(doc.sections) > 0\n assert doc.sections[0].text is not None\n assert doc.sections[0].link is not None\n\n def test_document_metadata_structure(\n self, connector: HubSpotConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that document metadata contains expected fields.\"\"\"\n connector.load_credentials(credentials)\n\n document_batches = connector.load_from_state()\n all_docs: list[Document] = []\n\n # Collect a few batches to test different object types\n batch_count = 0\n for batch in document_batches:\n all_docs.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n batch_count += 1\n if (\n batch_count >= 3 or len(all_docs) >= 20\n ): # Limit to avoid too many API calls\n break\n\n # Group documents by object type\n docs_by_type: dict[str, list[Document]] = {}\n for doc in all_docs:\n obj_type_value = doc.metadata[\"object_type\"]\n # Handle the case where metadata value could be a list\n obj_type = (\n obj_type_value if isinstance(obj_type_value, str) else obj_type_value[0]\n )\n if obj_type not in docs_by_type:\n docs_by_type[obj_type] = []\n docs_by_type[obj_type].append(doc)\n\n # Test each object type has expected metadata\n for obj_type, docs in docs_by_type.items():\n doc = docs[0] # Test first document of each type\n\n if obj_type == \"ticket\":\n assert \"ticket_id\" in doc.metadata\n assert doc.id.startswith(\"hubspot_ticket_\")\n elif obj_type == \"company\":\n assert \"company_id\" in doc.metadata\n assert doc.id.startswith(\"hubspot_company_\")\n\n elif obj_type == \"deal\":\n assert \"deal_id\" in doc.metadata\n assert doc.id.startswith(\"hubspot_deal_\")\n\n elif obj_type == \"contact\":\n assert \"contact_id\" in doc.metadata\n assert doc.id.startswith(\"hubspot_contact_\")\n\n # Check for associated object IDs in metadata (if they exist)\n potential_association_keys = [\n \"associated_contact_ids\",\n \"associated_company_ids\",\n \"associated_deal_ids\",\n \"associated_ticket_ids\",\n \"associated_note_ids\",\n ]\n\n for key in potential_association_keys:\n if key in doc.metadata:\n assert isinstance(doc.metadata[key], list)\n assert len(doc.metadata[key]) > 0\n assert all(isinstance(id_val, str) for id_val in doc.metadata[key])\n\n def test_associated_objects_as_sections(\n self, connector: HubSpotConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that associated objects are included as sections.\"\"\"\n connector.load_credentials(credentials)\n\n document_batches = connector.load_from_state()\n\n # Find a document with multiple sections (indicating associated objects)\n found_multi_section_doc = False\n batch_count = 0\n\n for batch in document_batches:\n for doc in batch:\n if isinstance(doc, HierarchyNode):\n continue\n if len(doc.sections) > 1:\n found_multi_section_doc = True\n\n # First section should be the main object\n main_section = doc.sections[0]\n assert main_section.text is not None\n assert main_section.link is not None\n\n # Additional sections should be associated objects\n for section in doc.sections[1:]:\n assert section.text is not None\n assert section.link is not None\n # Should contain object type information\n assert any(\n obj_type in section.text.lower()\n for obj_type in [\n \"contact:\",\n \"company:\",\n \"deal:\",\n \"ticket:\",\n \"note:\",\n ]\n )\n\n break\n\n if found_multi_section_doc:\n break\n\n batch_count += 1\n if batch_count >= 5: # Limit API calls\n break\n\n # Note: This test might not always pass if there are no associated objects\n # in the test HubSpot instance, but it validates the structure when they exist\n if found_multi_section_doc:\n print(\"✓ Found document with associated objects as sections\")\n else:\n print(\"⚠ No documents with associated objects found in test data\")\n\n def test_poll_source_functionality(\n self, connector: HubSpotConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test poll_source with time filtering.\"\"\"\n connector.load_credentials(credentials)\n\n # Test with a recent time range (last 30 days)\n end_time = datetime.now(timezone.utc)\n start_time = datetime.now(timezone.utc).replace(day=1) # Start of current month\n\n start_timestamp = int(start_time.timestamp())\n end_timestamp = int(end_time.timestamp())\n\n document_batches = connector.poll_source(start_timestamp, end_timestamp)\n\n # Should be able to get at least one batch\n first_batch = next(document_batches, None)\n\n if first_batch is not None:\n assert isinstance(first_batch, list)\n assert len(first_batch) > 0\n\n # Check that documents have proper timestamps\n for doc in first_batch:\n if isinstance(doc, HierarchyNode):\n continue\n assert doc.doc_updated_at is not None\n # Note: We don't strictly enforce the time range here since\n # the test data might not have recent updates\n else:\n print(\"⚠ No documents found in the specified time range\")\n\n def test_all_object_types_processed(\n self, connector: HubSpotConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Integration test to verify all object types are processed correctly.\"\"\"\n connector.load_credentials(credentials)\n\n document_batches = connector.load_from_state()\n all_docs: list[Document] = []\n object_types_found = set()\n\n # Collect several batches to ensure we see all object types\n batch_count = 0\n for batch in document_batches:\n all_docs.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n for doc in batch:\n if isinstance(doc, HierarchyNode):\n continue\n object_types_found.add(doc.metadata[\"object_type\"])\n\n batch_count += 1\n # Stop after we've seen all expected types or after reasonable number of batches\n if len(object_types_found) >= 4 or batch_count >= 10:\n break\n\n print(f\"Found {len(all_docs)} total documents\")\n print(f\"Object types found: {sorted(object_types_found)}\")\n\n # Should have at least some documents\n assert len(all_docs) > 0\n\n # Verify we can process multiple object types\n # Note: We don't require all 4 types since the test instance might not have all types\n assert len(object_types_found) >= 1\n\n # Verify document structure for each type found\n for obj_type in object_types_found:\n type_docs = [\n doc for doc in all_docs if doc.metadata[\"object_type\"] == obj_type\n ]\n assert len(type_docs) > 0\n\n # Check first document of this type\n doc = type_docs[0]\n assert doc.id.startswith(f\"hubspot_{obj_type}_\")\n assert doc.semantic_identifier is not None\n assert len(doc.sections) > 0\n assert doc.sections[0].text is not None\n assert doc.sections[0].link is not None\n\n # Check object-specific metadata\n if obj_type == \"company\":\n assert \"company_id\" in doc.metadata\n elif obj_type == \"deal\":\n assert \"deal_id\" in doc.metadata\n elif obj_type == \"contact\":\n assert \"contact_id\" in doc.metadata\n elif obj_type == \"ticket\":\n assert \"ticket_id\" in doc.metadata\n\n def test_init_default_object_types(self) -> None:\n \"\"\"Test that connector initializes with all object types by default.\"\"\"\n connector = HubSpotConnector()\n assert connector.object_types == AVAILABLE_OBJECT_TYPES\n assert \"tickets\" in connector.object_types\n assert \"companies\" in connector.object_types\n assert \"deals\" in connector.object_types\n assert \"contacts\" in connector.object_types\n\n def test_init_custom_object_types(self) -> None:\n \"\"\"Test that connector can be initialized with custom object types.\"\"\"\n custom_types = [\"tickets\", \"companies\"]\n connector = HubSpotConnector(object_types=custom_types)\n expected_set = {\"tickets\", \"companies\"}\n assert connector.object_types == expected_set\n assert \"tickets\" in connector.object_types\n assert \"companies\" in connector.object_types\n assert \"deals\" not in connector.object_types\n assert \"contacts\" not in connector.object_types\n\n def test_init_custom_object_types_from_list(self) -> None:\n \"\"\"Test that connector can be initialized with custom object types from a list (frontend format).\"\"\"\n custom_types_list = [\"tickets\", \"companies\"]\n connector = HubSpotConnector(object_types=custom_types_list)\n expected_set = {\"tickets\", \"companies\"}\n assert connector.object_types == expected_set\n assert \"tickets\" in connector.object_types\n assert \"companies\" in connector.object_types\n assert \"deals\" not in connector.object_types\n assert \"contacts\" not in connector.object_types\n\n def test_init_single_object_type(self) -> None:\n \"\"\"Test that connector can be initialized with a single object type.\"\"\"\n single_type = [\"deals\"]\n connector = HubSpotConnector(object_types=single_type)\n expected_set = {\"deals\"}\n assert connector.object_types == expected_set\n assert len(connector.object_types) == 1\n assert \"deals\" in connector.object_types\n\n def test_init_invalid_object_types(self) -> None:\n \"\"\"Test that connector raises error for invalid object types.\"\"\"\n invalid_types = [\"tickets\", \"invalid_type\", \"another_invalid\"]\n\n with pytest.raises(ValueError) as exc_info:\n HubSpotConnector(object_types=invalid_types)\n\n error_message = str(exc_info.value)\n assert \"Invalid object types\" in error_message\n assert \"invalid_type\" in error_message\n assert \"another_invalid\" in error_message\n assert \"Available types\" in error_message\n\n def test_init_empty_object_types(self) -> None:\n \"\"\"Test that connector can be initialized with empty object types set.\"\"\"\n empty_types: list[str] = []\n connector = HubSpotConnector(object_types=empty_types)\n expected_set: set[str] = set()\n assert connector.object_types == expected_set\n assert len(connector.object_types) == 0\n\n def test_selective_object_fetching_tickets_only(\n self, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that only tickets are fetched when configured.\"\"\"\n connector = HubSpotConnector(object_types=[\"tickets\"], batch_size=5)\n connector.load_credentials(credentials)\n\n document_batches = connector.load_from_state()\n all_docs: list[Document] = []\n\n # Collect a few batches\n batch_count = 0\n for batch in document_batches:\n all_docs.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n batch_count += 1\n if batch_count >= 3 or len(all_docs) >= 10:\n break\n\n # Should have documents\n if all_docs:\n # All documents should be tickets\n for doc in all_docs:\n assert doc.metadata[\"object_type\"] == \"ticket\"\n assert doc.id.startswith(\"hubspot_ticket_\")\n\n print(f\"✓ Successfully fetched {len(all_docs)} ticket documents only\")\n else:\n print(\"⚠ No ticket documents found in test data\")\n\n def test_selective_object_fetching_companies_and_deals(\n self, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that only companies and deals are fetched when configured.\"\"\"\n connector = HubSpotConnector(object_types=[\"companies\", \"deals\"], batch_size=5)\n connector.load_credentials(credentials)\n\n document_batches = connector.load_from_state()\n all_docs: list[Document] = []\n object_types_found = set()\n\n # Collect a few batches\n batch_count = 0\n for batch in document_batches:\n all_docs.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n for doc in batch:\n if isinstance(doc, HierarchyNode):\n continue\n object_types_found.add(doc.metadata[\"object_type\"])\n batch_count += 1\n if batch_count >= 3 or len(all_docs) >= 10:\n break\n\n if all_docs:\n # Should only have companies and deals\n assert object_types_found.issubset({\"company\", \"deal\"})\n assert \"ticket\" not in object_types_found\n assert \"contact\" not in object_types_found\n\n # Verify document structure\n for doc in all_docs:\n obj_type = doc.metadata[\"object_type\"]\n assert obj_type in [\"company\", \"deal\"]\n if obj_type == \"company\":\n assert doc.id.startswith(\"hubspot_company_\")\n elif obj_type == \"deal\":\n assert doc.id.startswith(\"hubspot_deal_\")\n\n print(\n f\"✓ Successfully fetched {len(all_docs)} documents of types: {object_types_found}\"\n )\n else:\n print(\"⚠ No company/deal documents found in test data\")\n\n def test_empty_object_types_fetches_nothing(\n self, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that no documents are fetched when object_types is empty.\"\"\"\n connector = HubSpotConnector(object_types=[], batch_size=5)\n connector.load_credentials(credentials)\n\n document_batches = connector.load_from_state()\n all_docs: list[Document] = []\n\n # Try to collect batches\n batch_count = 0\n for batch in document_batches:\n all_docs.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n batch_count += 1\n if batch_count >= 2: # Don't wait too long\n break\n\n # Should have no documents\n assert len(all_docs) == 0\n print(\"✓ No documents fetched with empty object_types as expected\")\n\n def test_poll_source_respects_object_types(\n self, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Test that poll_source respects the object_types configuration.\"\"\"\n connector = HubSpotConnector(object_types=[\"contacts\"], batch_size=5)\n connector.load_credentials(credentials)\n\n # Test with a recent time range\n end_time = datetime.now(timezone.utc)\n start_time = datetime.now(timezone.utc).replace(day=1)\n\n start_timestamp = int(start_time.timestamp())\n end_timestamp = int(end_time.timestamp())\n\n document_batches = connector.poll_source(start_timestamp, end_timestamp)\n all_docs: list[Document] = []\n\n # Collect a few batches\n batch_count = 0\n for batch in document_batches:\n all_docs.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n batch_count += 1\n if batch_count >= 2 or len(all_docs) >= 5:\n break\n\n if all_docs:\n # All documents should be contacts\n for doc in all_docs:\n assert doc.metadata[\"object_type\"] == \"contact\"\n assert doc.id.startswith(\"hubspot_contact_\")\n\n print(\n f\"✓ Poll source successfully fetched {len(all_docs)} contact documents only\"\n )\n else:\n print(\"⚠ No contact documents found in specified time range\")\n\n def test_object_types_immutability(self) -> None:\n \"\"\"Test that object_types set cannot be modified externally.\"\"\"\n original_types = [\"tickets\", \"companies\"]\n connector = HubSpotConnector(object_types=original_types)\n\n # Modifying the original list should not affect the connector\n original_types.append(\"deals\")\n assert \"deals\" not in connector.object_types\n assert connector.object_types == {\"tickets\", \"companies\"}\n\n # Trying to modify the connector's object_types should not affect the original\n connector_types = connector.object_types\n connector_types.add(\"contacts\")\n # The connector should still have the original types since we made a copy\n # Note: This test verifies our implementation makes a copy in __init__\n\n def test_url_generation(self) -> None:\n \"\"\"Test that URLs are generated correctly for different object types.\"\"\"\n connector = HubSpotConnector()\n connector.portal_id = \"12345\" # Mock portal ID\n\n # Test URL generation for each object type\n ticket_url = connector._get_object_url(\"tickets\", \"67890\")\n expected_ticket_url = \"https://app.hubspot.com/contacts/12345/record/0-5/67890\"\n assert ticket_url == expected_ticket_url\n\n company_url = connector._get_object_url(\"companies\", \"11111\")\n expected_company_url = \"https://app.hubspot.com/contacts/12345/record/0-2/11111\"\n assert company_url == expected_company_url\n\n deal_url = connector._get_object_url(\"deals\", \"22222\")\n expected_deal_url = \"https://app.hubspot.com/contacts/12345/record/0-3/22222\"\n assert deal_url == expected_deal_url\n\n contact_url = connector._get_object_url(\"contacts\", \"33333\")\n expected_contact_url = \"https://app.hubspot.com/contacts/12345/record/0-1/33333\"\n assert contact_url == expected_contact_url\n\n note_url = connector._get_object_url(\"notes\", \"44444\")\n expected_note_url = \"https://app.hubspot.com/contacts/12345/objects/0-4/44444\"\n assert note_url == expected_note_url\n\n def test_ticket_with_none_content(self) -> None:\n \"\"\"Test that tickets with None content are handled gracefully.\"\"\"\n connector = HubSpotConnector(object_types=[\"tickets\"], batch_size=10)\n connector._access_token = \"mock_token\"\n connector._portal_id = \"mock_portal_id\"\n\n # Create a mock ticket with None content\n mock_ticket = MagicMock()\n mock_ticket.id = \"12345\"\n mock_ticket.properties = {\n \"subject\": \"Test Ticket\",\n \"content\": None, # This is the key test case\n \"hs_ticket_priority\": \"HIGH\",\n }\n mock_ticket.updated_at = datetime.now(timezone.utc)\n\n # Mock the HubSpot API client\n mock_api_client = MagicMock()\n\n # Mock the API calls and associated object methods\n with (\n patch(\"onyx.connectors.hubspot.connector.HubSpot\") as MockHubSpot,\n patch.object(connector, \"_paginated_results\") as mock_paginated,\n patch.object(connector, \"_get_associated_objects\", return_value=[]),\n patch.object(connector, \"_get_associated_notes\", return_value=[]),\n ):\n MockHubSpot.return_value = mock_api_client\n mock_paginated.return_value = iter([mock_ticket])\n\n # This should not raise a validation error\n document_batches = connector._process_tickets()\n first_batch = next(document_batches, None)\n\n # Verify the document was created successfully\n assert first_batch is not None\n assert len(first_batch) == 1\n\n doc = first_batch[0]\n assert not isinstance(doc, HierarchyNode)\n assert doc.id == \"hubspot_ticket_12345\"\n assert doc.semantic_identifier == \"Test Ticket\"\n\n # Verify the first section has an empty string, not None\n assert len(doc.sections) > 0\n assert doc.sections[0].text == \"\" # Should be empty string, not None\n assert doc.sections[0].link is not None\n" }, { "path": "backend/tests/daily/connectors/imap/models.py", "content": "from pydantic import BaseModel\n\nfrom onyx.connectors.models import Document\nfrom tests.daily.connectors.utils import to_text_sections\n\n\nclass EmailDoc(BaseModel):\n subject: str\n recipients: set[str]\n body: str\n\n @classmethod\n def from_doc(cls, document: Document) -> \"EmailDoc\":\n # Acceptable to perform assertions since this class is only used in tests.\n assert document.title\n assert document.external_access\n\n body = \" \".join(to_text_sections(sections=iter(document.sections)))\n\n return cls(\n subject=document.title,\n recipients=document.external_access.external_user_emails,\n body=body,\n )\n" }, { "path": "backend/tests/daily/connectors/imap/test_imap_connector.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\nfrom onyx.connectors.imap.connector import ImapConnector\nfrom tests.daily.connectors.imap.models import EmailDoc\nfrom tests.daily.connectors.utils import (\n load_all_from_connector,\n)\n\n\n@pytest.fixture\ndef imap_connector() -> ImapConnector:\n host = os.environ.get(\"IMAP_HOST\")\n mailboxes_str = os.environ.get(\"IMAP_MAILBOXES\")\n username = os.environ.get(\"IMAP_USERNAME\")\n password = os.environ.get(\"IMAP_PASSWORD\")\n\n assert host\n mailboxes = (\n [mailbox.strip() for mailbox in mailboxes_str.split(\",\") if mailbox]\n if mailboxes_str\n else []\n )\n\n imap_connector = ImapConnector(\n host=host,\n mailboxes=mailboxes,\n )\n imap_connector.set_credentials_provider(\n OnyxStaticCredentialsProvider(\n tenant_id=None,\n connector_name=DocumentSource.IMAP,\n credential_json={\n \"imap_username\": username,\n \"imap_password\": password,\n },\n )\n )\n\n return imap_connector\n\n\n@pytest.mark.parametrize(\n \"expected_email_docs\",\n [\n [\n EmailDoc(\n subject=\"Testing\",\n recipients=set([\"admin@onyx-test.com\", \"raunak@onyx.app\"]),\n body=\"Hello, testing.\",\n ),\n EmailDoc(\n subject=\"Hello world\",\n recipients=set([\"admin@onyx-test.com\", \"r@rabh.io\", \"raunak@onyx.app\"]),\n body='Hello world, this is an email that contains multiple \"To\" recipients.',\n ),\n ]\n ],\n)\ndef test_imap_connector(\n imap_connector: ImapConnector,\n expected_email_docs: list[EmailDoc],\n) -> None:\n actual_email_docs = [\n EmailDoc.from_doc(document=document)\n for document in load_all_from_connector(\n connector=imap_connector,\n start=0,\n end=time.time(),\n include_permissions=True,\n ).documents\n ]\n\n assert actual_email_docs == expected_email_docs\n" }, { "path": "backend/tests/daily/connectors/jira/test_jira_basic.py", "content": "import os\nimport time\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.jira.connector import JiraConnector\nfrom onyx.connectors.models import Document\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\ndef _make_connector(scoped_token: bool = False) -> JiraConnector:\n connector = JiraConnector(\n jira_base_url=\"https://danswerai.atlassian.net\",\n project_key=\"AS\",\n comment_email_blacklist=[],\n scoped_token=scoped_token,\n )\n connector.load_credentials(\n {\n \"jira_user_email\": os.environ[\"JIRA_USER_EMAIL\"],\n \"jira_api_token\": (\n os.environ[\"JIRA_API_TOKEN_SCOPED\"]\n if scoped_token\n else os.environ[\"JIRA_API_TOKEN\"]\n ),\n }\n )\n return connector\n\n\n@pytest.fixture\ndef jira_connector() -> JiraConnector:\n return _make_connector()\n\n\n@pytest.fixture\ndef jira_connector_scoped() -> JiraConnector:\n return _make_connector(scoped_token=True)\n\n\n@pytest.fixture\ndef jira_connector_with_jql() -> JiraConnector:\n connector = JiraConnector(\n jira_base_url=\"https://danswerai.atlassian.net\",\n jql_query=\"project = 'AS' AND issuetype = Story\",\n comment_email_blacklist=[],\n )\n connector.load_credentials(\n {\n \"jira_user_email\": os.environ[\"JIRA_USER_EMAIL\"],\n \"jira_api_token\": os.environ[\"JIRA_API_TOKEN\"],\n }\n )\n connector.validate_connector_settings()\n\n return connector\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_jira_connector_basic(\n reset: None, # noqa: ARG001\n jira_connector: JiraConnector,\n) -> None:\n _test_jira_connector_basic(jira_connector)\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_jira_connector_basic_scoped(\n reset: None, # noqa: ARG001\n jira_connector_scoped: JiraConnector,\n) -> None:\n _test_jira_connector_basic(jira_connector_scoped)\n\n\ndef _test_jira_connector_basic(jira_connector: JiraConnector) -> None:\n docs = load_all_from_connector(\n connector=jira_connector,\n start=0,\n end=time.time(),\n ).documents\n assert len(docs) == 2\n\n # Find story and epic\n story: Document | None = None\n epic: Document | None = None\n for doc in docs:\n if doc.metadata[\"issuetype\"] == \"Story\":\n story = doc\n elif doc.metadata[\"issuetype\"] == \"Epic\":\n epic = doc\n\n assert story is not None\n assert epic is not None\n\n # Check task\n assert story.id == \"https://danswerai.atlassian.net/browse/AS-3\"\n assert story.semantic_identifier == \"AS-3: Magic Answers\"\n assert story.source == DocumentSource.JIRA\n assert story.metadata == {\n \"priority\": \"Medium\",\n \"status\": \"Done\",\n \"resolution\": \"Done\",\n \"resolution_date\": \"2025-05-29T15:33:31.031-0700\",\n \"reporter\": \"Chris Weaver\",\n \"assignee\": \"Chris Weaver\",\n \"issuetype\": \"Story\",\n \"created\": \"2025-04-16T16:44:06.716-0700\",\n \"reporter_email\": \"chris@onyx.app\",\n \"assignee_email\": \"chris@onyx.app\",\n \"project_name\": \"DailyConnectorTestProject\",\n \"project\": \"AS\",\n \"parent\": \"AS-4\",\n \"key\": \"AS-3\",\n \"updated\": \"2025-06-17T12:13:00.070-0700\",\n }\n assert story.secondary_owners is None\n assert story.title == \"AS-3 Magic Answers\"\n assert story.from_ingestion_api is False\n assert story.additional_info is None\n\n assert len(story.sections) == 1\n section = story.sections[0]\n assert (\n section.text\n == \"This is a critical request for super-human answer quality in Onyx! We need magic!\\n\"\n )\n assert section.link == \"https://danswerai.atlassian.net/browse/AS-3\"\n\n # Check epic\n assert epic.id == \"https://danswerai.atlassian.net/browse/AS-4\"\n assert epic.semantic_identifier == \"AS-4: EPIC\"\n assert epic.source == DocumentSource.JIRA\n assert epic.metadata == {\n \"priority\": \"Medium\",\n \"status\": \"Backlog\",\n \"reporter\": \"Founder Onyx\",\n \"assignee\": \"Chris Weaver\",\n \"issuetype\": \"Epic\",\n \"created\": \"2025-04-16T16:55:53.068-0700\",\n \"reporter_email\": \"founders@onyx.app\",\n \"assignee_email\": \"chris@onyx.app\",\n \"project_name\": \"DailyConnectorTestProject\",\n \"project\": \"AS\",\n \"key\": \"AS-4\",\n \"updated\": \"2025-05-29T14:43:05.312-0700\",\n }\n assert epic.secondary_owners is None\n assert epic.title == \"AS-4 EPIC\"\n assert epic.from_ingestion_api is False\n assert epic.additional_info is None\n\n assert len(epic.sections) == 1\n section = epic.sections[0]\n assert section.text == \"example_text\\n\"\n assert section.link == \"https://danswerai.atlassian.net/browse/AS-4\"\n\n\n@patch(\n \"onyx.file_processing.extract_file_text.get_unstructured_api_key\",\n return_value=None,\n)\ndef test_jira_connector_with_jql(\n reset: None, # noqa: ARG001\n jira_connector_with_jql: JiraConnector,\n) -> None:\n \"\"\"Test that JQL query functionality works correctly.\n\n This test verifies that when a JQL query is provided, only issues matching the query are returned.\n The JQL query used is \"project = \\'AS\\' AND issuetype = Story\", which should only return Story-type issues.\n \"\"\"\n docs = load_all_from_connector(\n connector=jira_connector_with_jql,\n start=0,\n end=time.time(),\n ).documents\n\n # Should only return Story-type issues\n assert len(docs) == 1\n\n # All documents should be Story-type\n for doc in docs:\n assert doc.metadata[\"issuetype\"] == \"Story\"\n\n # Verify it's the expected Story\n story = docs[0]\n assert story.id == \"https://danswerai.atlassian.net/browse/AS-3\"\n assert story.semantic_identifier == \"AS-3: Magic Answers\"\n assert story.metadata[\"issuetype\"] == \"Story\"\n" }, { "path": "backend/tests/daily/connectors/notion/test_notion_connector.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.notion.connector import NotionConnector\n\n\ndef compare_hierarchy_nodes(\n yielded_nodes: list[HierarchyNode],\n expected_nodes: list[HierarchyNode],\n) -> None:\n \"\"\"Compare yielded HierarchyNodes against expected ground truth.\n\n Compares nodes by their essential fields (raw_node_id, raw_parent_id, display_name, link).\n Order does not matter.\n \"\"\"\n if not expected_nodes:\n # Empty ground truth - skip comparison for now\n return\n\n yielded_set = {\n (n.raw_node_id, n.raw_parent_id, n.display_name, n.link) for n in yielded_nodes\n }\n expected_set = {\n (n.raw_node_id, n.raw_parent_id, n.display_name, n.link) for n in expected_nodes\n }\n\n missing = expected_set - yielded_set\n extra = yielded_set - expected_set\n\n assert not missing, f\"Missing expected HierarchyNodes: {missing}\"\n assert not extra, f\"Unexpected HierarchyNodes: {extra}\"\n\n\n@pytest.fixture\ndef notion_connector() -> NotionConnector:\n \"\"\"Create a NotionConnector with credentials from environment variables\"\"\"\n connector = NotionConnector()\n connector.load_credentials(\n {\n \"notion_integration_token\": os.environ[\"NOTION_INTEGRATION_TOKEN\"],\n }\n )\n return connector\n\n\ndef test_notion_connector_basic(notion_connector: NotionConnector) -> None:\n \"\"\"Test the NotionConnector with a real Notion page.\n\n Uses a Notion workspace under the onyx-test.com domain.\n \"\"\"\n doc_batch_generator = notion_connector.poll_source(0, time.time())\n\n # Collect all documents and hierarchy nodes from all batches\n documents: list[Document] = []\n hierarchy_nodes: list[HierarchyNode] = []\n for doc_batch in doc_batch_generator:\n for item in doc_batch:\n if isinstance(item, HierarchyNode):\n hierarchy_nodes.append(item)\n else:\n documents.append(item)\n\n # Verify document count\n assert (\n len(documents) == 5\n ), \"Expected exactly 5 documents (root, two children, table entry, and table entry child)\"\n\n # Verify HierarchyNodes against ground truth (empty for now)\n expected_hierarchy_nodes: list[HierarchyNode] = []\n compare_hierarchy_nodes(hierarchy_nodes, expected_hierarchy_nodes)\n\n # Find root and child documents by semantic identifier\n root_doc = None\n child1_doc = None\n child2_doc = None\n table_entry_doc = None\n table_entry_child_doc = None\n for doc in documents:\n if doc.semantic_identifier == \"Root\":\n root_doc = doc\n elif doc.semantic_identifier == \"Child1\":\n child1_doc = doc\n elif doc.semantic_identifier == \"Child2\":\n child2_doc = doc\n elif doc.semantic_identifier == \"table-entry01\":\n table_entry_doc = doc\n elif doc.semantic_identifier == \"Child-table-entry01\":\n table_entry_child_doc = doc\n\n assert root_doc is not None, \"Root document not found\"\n assert child1_doc is not None, \"Child1 document not found\"\n assert child2_doc is not None, \"Child2 document not found\"\n assert table_entry_doc is not None, \"Table entry document not found\"\n assert table_entry_child_doc is not None, \"Table entry child document not found\"\n\n # Verify root document structure\n assert root_doc.id is not None\n assert root_doc.source == DocumentSource.NOTION\n\n # Section checks for root\n assert len(root_doc.sections) == 1\n root_section = root_doc.sections[0]\n\n # Content specific checks for root\n assert root_section.text == \"\\nroot\"\n assert root_section.link is not None\n assert root_section.link.startswith(\"https://www.notion.so/\")\n\n # Verify child1 document structure\n assert child1_doc.id is not None\n assert child1_doc.source == DocumentSource.NOTION\n\n # Section checks for child1\n assert len(child1_doc.sections) == 1\n child1_section = child1_doc.sections[0]\n\n # Content specific checks for child1\n assert child1_section.text == \"\\nchild1\"\n assert child1_section.link is not None\n assert child1_section.link.startswith(\"https://www.notion.so/\")\n\n # Verify child2 document structure (includes database)\n assert child2_doc.id is not None\n assert child2_doc.source == DocumentSource.NOTION\n\n # Section checks for child2\n assert len(child2_doc.sections) == 2 # One for content, one for database\n child2_section = child2_doc.sections[0]\n child2_db_section = child2_doc.sections[1]\n\n # Content specific checks for child2\n assert child2_section.text == \"\\nchild2\"\n assert child2_section.link is not None\n assert child2_section.link.startswith(\"https://www.notion.so/\")\n\n # Database section checks for child2\n assert child2_db_section.text is not None\n assert child2_db_section.text.strip() != \"\" # Should contain some database content\n assert child2_db_section.link is not None\n assert child2_db_section.link.startswith(\"https://www.notion.so/\")\n\n # Verify table entry document structure\n assert table_entry_doc.id is not None\n assert table_entry_doc.source == DocumentSource.NOTION\n\n # Section checks for table entry\n assert len(table_entry_doc.sections) == 1\n table_entry_section = table_entry_doc.sections[0]\n\n # Content specific checks for table entry\n assert table_entry_section.text == \"\\ntable-entry01\"\n assert table_entry_section.link is not None\n assert table_entry_section.link.startswith(\"https://www.notion.so/\")\n\n # Verify table entry child document structure\n assert table_entry_child_doc.id is not None\n assert table_entry_child_doc.source == DocumentSource.NOTION\n\n # Section checks for table entry child\n assert len(table_entry_child_doc.sections) == 1\n table_entry_child_section = table_entry_child_doc.sections[0]\n\n # Content specific checks for table entry child\n assert table_entry_child_section.text == \"\\nchild-table-entry01\"\n assert table_entry_child_section.link is not None\n assert table_entry_child_section.link.startswith(\"https://www.notion.so/\")\n" }, { "path": "backend/tests/daily/connectors/outline/test_outline_connector.py", "content": "import os\nimport time\nfrom typing import Any\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.exceptions import ConnectorValidationError\nfrom onyx.connectors.exceptions import CredentialExpiredError\nfrom onyx.connectors.models import ConnectorMissingCredentialError\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.outline.connector import OutlineConnector\n\n\nclass TestOutlineConnector:\n \"\"\"Comprehensive test suite for the OutlineConnector.\"\"\"\n\n @pytest.fixture\n def connector(self) -> OutlineConnector:\n \"\"\"Create an Outline connector instance.\"\"\"\n return OutlineConnector(batch_size=10)\n\n @pytest.fixture\n def credentials(self) -> dict[str, Any]:\n \"\"\"Provide test credentials from environment variables.\"\"\"\n outline_base_url = os.environ.get(\"OUTLINE_BASE_URL\")\n outline_api_token = os.environ.get(\"OUTLINE_API_TOKEN\")\n\n if not outline_base_url or not outline_api_token:\n pytest.skip(\n \"OUTLINE_BASE_URL and OUTLINE_API_TOKEN environment variables must be set\"\n )\n\n return {\n \"outline_api_token\": outline_api_token,\n \"outline_base_url\": outline_base_url,\n }\n\n def test_credentials_missing_raises_exception(self) -> None:\n \"\"\"Should raise if credentials are missing.\"\"\"\n connector = OutlineConnector()\n\n with pytest.raises(ConnectorMissingCredentialError) as exc_info:\n list(connector.load_from_state())\n assert \"Outline\" in str(exc_info.value)\n\n def test_load_credentials(\n self, connector: OutlineConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Credentials should load correctly.\"\"\"\n result = connector.load_credentials(credentials)\n\n assert result is None\n assert connector.outline_client is not None\n assert connector.outline_client.api_token == credentials[\"outline_api_token\"]\n assert connector.outline_client.base_url == credentials[\n \"outline_base_url\"\n ].rstrip(\"/\")\n\n def test_outline_connector_basic(\n self, connector: OutlineConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Validate that connector fetches and structures documents properly.\"\"\"\n connector.load_credentials(credentials)\n\n documents: list[Document] = []\n for batch in connector.load_from_state():\n documents.extend(\n [doc for doc in batch if not isinstance(doc, HierarchyNode)]\n )\n\n assert len(documents) > 0, \"Expected at least one document/collection\"\n\n collections = [d for d in documents if d.metadata.get(\"type\") == \"collection\"]\n docs = [d for d in documents if d.metadata.get(\"type\") == \"document\"]\n\n assert len(collections) > 0, \"Should find at least one collection\"\n\n collection = collections[0]\n assert collection.id.startswith(\"outline_collection__\")\n assert collection.source == DocumentSource.OUTLINE\n assert collection.title is not None\n assert len(collection.sections) == 1\n assert collection.sections[0].text is not None\n assert collection.metadata[\"type\"] == \"collection\"\n\n if docs:\n document = docs[0]\n assert document.id.startswith(\"outline_document__\")\n assert document.source == DocumentSource.OUTLINE\n assert document.title is not None\n assert len(document.sections) == 1\n assert document.sections[0].text is not None\n assert document.metadata[\"type\"] == \"document\"\n\n section_link = document.sections[0].link\n assert section_link is not None\n assert \"/doc/\" in section_link\n\n def test_outline_connector_time_filtering(\n self, connector: OutlineConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Validate poll_source with time range filtering.\"\"\"\n connector.load_credentials(credentials)\n\n end_time = time.time()\n start_time = end_time - 30 * 24 * 60 * 60\n\n docs: list[Document] = []\n for batch in connector.poll_source(start_time, end_time):\n docs.extend([doc for doc in batch if not isinstance(doc, HierarchyNode)])\n\n for doc in docs:\n assert isinstance(doc, Document)\n assert doc.source == DocumentSource.OUTLINE\n if doc.doc_updated_at:\n assert start_time <= doc.doc_updated_at.timestamp() <= end_time\n\n def test_outline_connector_load_from_state(\n self, connector: OutlineConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"load_from_state should fetch documents.\"\"\"\n connector.load_credentials(credentials)\n\n gen = connector.load_from_state()\n batch = next(gen)\n assert isinstance(batch, list)\n\n for doc in batch:\n assert isinstance(doc, Document)\n assert doc.source == DocumentSource.OUTLINE\n\n def test_outline_connector_batch_processing(\n self, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Connector should respect batch size.\"\"\"\n small_batch_connector = OutlineConnector(batch_size=2)\n small_batch_connector.load_credentials(credentials)\n\n for batch in small_batch_connector.poll_source(0, time.time()):\n assert len(batch) <= 2\n break\n\n def test_outline_connector_document_types(\n self, connector: OutlineConnector, credentials: dict[str, Any]\n ) -> None:\n \"\"\"Validate metadata for collections and documents.\"\"\"\n connector.load_credentials(credentials)\n\n docs: list[Document] = []\n for batch in connector.poll_source(0, time.time()):\n docs.extend([doc for doc in batch if not isinstance(doc, HierarchyNode)])\n\n if docs:\n doc_types = {d.metadata[\"type\"] for d in docs}\n assert doc_types.issubset({\"document\", \"collection\"})\n\n for doc in docs:\n if doc.metadata[\"type\"] == \"document\":\n assert any(\n (s.text.strip() if s.text else None) for s in doc.sections\n )\n elif doc.metadata[\"type\"] == \"collection\":\n assert len(doc.sections) >= 1\n\n def test_outline_connector_invalid_credentials(self) -> None:\n \"\"\"Should raise with invalid/missing credentials.\"\"\"\n connector = OutlineConnector()\n\n # Missing everything\n with pytest.raises(ConnectorMissingCredentialError):\n connector.load_credentials({})\n\n # Missing base URL\n with pytest.raises(ConnectorMissingCredentialError):\n connector.load_credentials({\"outline_api_token\": \"token\"})\n\n # Missing token\n with pytest.raises(ConnectorMissingCredentialError):\n connector.load_credentials({\"outline_base_url\": \"https://example.com\"})\n\n # Invalid credentials will be caught during validation, not credential loading\n connector.load_credentials(\n {\n \"outline_base_url\": \"https://invalid.invalid\",\n \"outline_api_token\": \"invalid\",\n }\n )\n # Validation should catch invalid credentials\n with pytest.raises((CredentialExpiredError, ConnectorValidationError)):\n connector.validate_connector_settings()\n\n def test_outline_connector_invalid_url(self) -> None:\n \"\"\"Invalid URL should raise validation error during validation.\"\"\"\n connector = OutlineConnector()\n\n # Load credentials with invalid URL\n connector.load_credentials(\n {\n \"outline_base_url\": \"https://not-a-valid-url.invalid\",\n \"outline_api_token\": \"token\",\n }\n )\n\n # Validation should catch invalid URL\n with pytest.raises(ConnectorValidationError):\n connector.validate_connector_settings()\n" }, { "path": "backend/tests/daily/connectors/salesforce/test_salesforce_connector.py", "content": "import json\nimport os\nfrom datetime import datetime\nfrom datetime import timezone\nfrom pathlib import Path\nfrom typing import Any\nfrom typing import cast\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.salesforce.connector import SalesforceConnector\nfrom onyx.connectors.salesforce.utils import ACCOUNT_OBJECT_TYPE\n\n\ndef extract_key_value_pairs_to_set(\n list_of_unparsed_key_value_strings: list[str],\n) -> set[str]:\n set_of_key_value_pairs = set()\n for string_key_value_pairs in list_of_unparsed_key_value_strings:\n list_of_parsed_key_values = string_key_value_pairs.split(\"\\n\")\n for key_value_pair in list_of_parsed_key_values:\n set_of_key_value_pairs.add(key_value_pair.strip())\n return set_of_key_value_pairs\n\n\ndef _load_reference_data(\n file_name: str = \"test_salesforce_data.json\",\n) -> dict[str, str | list[str] | dict[str, Any] | list[dict[str, Any]]]:\n current_dir = Path(__file__).parent\n with open(current_dir / file_name, \"r\") as f:\n return json.load(f)\n\n\n@pytest.fixture\ndef salesforce_connector() -> SalesforceConnector:\n connector = SalesforceConnector(\n requested_objects=[ACCOUNT_OBJECT_TYPE, \"Contact\", \"Opportunity\"],\n )\n\n username = os.environ[\"SF_USERNAME\"]\n password = os.environ[\"SF_PASSWORD\"]\n security_token = os.environ[\"SF_SECURITY_TOKEN\"]\n\n connector.load_credentials(\n {\n \"sf_username\": username,\n \"sf_password\": password,\n \"sf_security_token\": security_token,\n }\n )\n return connector\n\n\n# TODO: make the credentials not expire\n@pytest.mark.skip(\n reason=(\n \"Credentials change over time, so this test will fail if run when the credentials expire.\"\n )\n)\ndef test_salesforce_connector_basic(salesforce_connector: SalesforceConnector) -> None:\n test_data = _load_reference_data()\n target_test_doc: Document | None = None\n all_docs: list[Document] = []\n for doc_batch in salesforce_connector.load_from_state():\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs.append(doc)\n if doc.id == test_data[\"id\"]:\n target_test_doc = doc\n break\n\n # The number of docs here seems to change actively so do a very loose check\n # as of 2025-03-28 it was around 32472\n assert len(all_docs) > 32000\n assert len(all_docs) < 40000\n\n assert target_test_doc is not None\n\n # Set of received links\n received_links: set[str] = set()\n # List of received text fields, which contain key-value pairs seperated by newlines\n received_text: list[str] = []\n\n # Iterate over the sections of the target test doc to extract the links and text\n for section in target_test_doc.sections:\n assert section.link\n assert section.text\n received_links.add(section.link)\n received_text.append(section.text)\n\n # Check that the received links match the expected links from the test data json\n expected_links = set(test_data[\"expected_links\"])\n assert received_links == expected_links\n\n # Check that the received key-value pairs from the text fields match the expected key-value pairs from the test data json\n expected_text = test_data[\"expected_text\"]\n if not isinstance(expected_text, list):\n raise ValueError(\"Expected text is not a list\")\n\n unparsed_expected_key_value_pairs: list[str] = cast(list[str], expected_text)\n received_key_value_pairs = extract_key_value_pairs_to_set(received_text)\n expected_key_value_pairs = extract_key_value_pairs_to_set(\n unparsed_expected_key_value_pairs\n )\n assert received_key_value_pairs == expected_key_value_pairs\n\n # Check that the rest of the fields match the expected fields from the test data json\n assert target_test_doc.source == DocumentSource.SALESFORCE\n assert target_test_doc.semantic_identifier == test_data[\"semantic_identifier\"]\n assert target_test_doc.metadata == test_data[\"metadata\"]\n\n assert target_test_doc.primary_owners is not None\n primary_owner = target_test_doc.primary_owners[0]\n expected_primary_owner = test_data[\"primary_owners\"]\n assert isinstance(expected_primary_owner, dict)\n assert primary_owner.email == expected_primary_owner[\"email\"]\n assert primary_owner.first_name == expected_primary_owner[\"first_name\"]\n assert primary_owner.last_name == expected_primary_owner[\"last_name\"]\n\n secondary_owners = (\n [owner.model_dump() for owner in target_test_doc.secondary_owners]\n if target_test_doc.secondary_owners\n else None\n )\n assert secondary_owners == test_data[\"secondary_owners\"]\n assert target_test_doc.title == test_data[\"title\"]\n\n\n@pytest.mark.skip(\n reason=(\n \"All Salesforce tests need to be re-thought + made less flakey. \"\n \"We need to handle credential resets + the rate limits (move to a smaller dataset)\"\n )\n)\ndef test_salesforce_connector_poll_source(\n salesforce_connector: SalesforceConnector,\n) -> None:\n\n intermediate_time = datetime(\n 2024, 6, 3, 0, 0, 0, tzinfo=timezone.utc\n ) # roughly 92 docs\n\n # intermediate_time = datetime(2024, 7, 1, 0, 0, 0, tzinfo=timezone.utc) # roughly 1100 to 1200 docs\n\n all_docs_1: list[Document] = []\n for doc_batch in salesforce_connector.poll_source(0, intermediate_time.timestamp()):\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs_1.append(doc)\n\n len_1 = len(all_docs_1)\n\n # NOTE: this is the correct document count.\n # If you were to inspect the underlying db, however, the partial download results in\n # an incomplete set of object relationships. This is expected.\n\n assert len_1 > 85 and len_1 < 100\n print(f\"all_docs_1 length: {len(all_docs_1)}\")\n\n # assert len_1 > 1100 and len_1 < 1200\n # print(f\"all_docs_1 length: {len(all_docs_1)}\")\n\n # leave this out for the moment because it's slow to process 30k docs\n # all_docs_2: list[Document] = []\n # for doc_batch in salesforce_connector.poll_source(\n # intermediate_time.timestamp(), time.time()\n # ):\n # for doc in doc_batch:\n # all_docs_2.append(doc)\n\n # len_2 = len(all_docs_2)\n # assert len_2 > 31000\n\n # print(f\"all_docs_2 length: {len(all_docs_2)}\")\n\n\n# TODO: make the credentials not expire\n@pytest.mark.skip(\n reason=(\n \"Credentials change over time, so this test will fail if run when the credentials expire.\"\n )\n)\ndef test_salesforce_connector_slim(salesforce_connector: SalesforceConnector) -> None:\n # Get all doc IDs from the full connector\n all_full_doc_ids = set()\n for doc_batch in salesforce_connector.load_from_state():\n all_full_doc_ids.update(\n [doc.id for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # Get all doc IDs from the slim connector\n all_slim_doc_ids = set()\n for slim_doc_batch in salesforce_connector.retrieve_all_slim_docs_perm_sync():\n all_slim_doc_ids.update(\n [doc.id for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # The set of full doc IDs should be always be a subset of the slim doc IDs\n assert all_full_doc_ids.issubset(all_slim_doc_ids)\n" }, { "path": "backend/tests/daily/connectors/salesforce/test_salesforce_data.json", "content": "{\n \"id\": \"SALESFORCE_001bm00000eu6n5AAA\",\n \"expected_links\": [\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESpEeAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESqd3AAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESoKiAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvDSAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrmHAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrl2AAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvejAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000EStlvAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESpPfAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrP9AAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvlMAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESt3JAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESoBkAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000EStw2AAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrkMAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESojKAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuLEAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESoSIAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESu2YAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvgSAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESurnAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrnqAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESoB5AAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuJuAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrfyAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/001bm00000eu6n5AAA\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESpUHAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESsgGAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESr7UAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESu1BAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESpqzAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESplZAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvJ3AAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESurKAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000EStSiAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuJFAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESu8xAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESqfzAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESqsrAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000EStoZAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESsIUAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESsAGAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESv8GAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrOKAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESoUmAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESudKAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuJ8AAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvf2AAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESw3qAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESugRAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESr18AAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESqV1AAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuLVAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESpjoAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESqULAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuCAAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrfpAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESp5YAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrMNAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000EStaUAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESt5LAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrtcAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESomaAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrtIAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESoToAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuWLAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESrWvAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESsJEAA1\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESsxwAAD\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvUgAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESvWjAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000EStBuAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESpZiAAL\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuhYAAT\",\n \"https://danswer-dev-ed.develop.my.salesforce.com/003bm00000ESuWAAA1\"\n ],\n \"expected_text\": [\n \"IsDeleted: false\\nBillingCity: Shaykh al \\u00e1\\u00b8\\u00a8ad\\u00c4\\u00abd\\nName: Voonder\\nCleanStatus: Pending\\nBillingStreet: 12 Cambridge Parkway\",\n \"Email: eslayqzs@icio.us\\nIsDeleted: false\\nLastName: Slay\\nIsEmailBounced: false\\nFirstName: Ebeneser\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ptweedgdh@umich.edu\\nIsDeleted: false\\nLastName: Tweed\\nIsEmailBounced: false\\nFirstName: Paulita\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ehurnellnlx@facebook.com\\nIsDeleted: false\\nLastName: Hurnell\\nIsEmailBounced: false\\nFirstName: Eliot\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ccarik4q4@google.it\\nIsDeleted: false\\nLastName: Carik\\nIsEmailBounced: false\\nFirstName: Chadwick\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: cvannozziina6@moonfruit.com\\nIsDeleted: false\\nLastName: Vannozzii\\nIsEmailBounced: false\\nFirstName: Christophorus\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: mikringill2kz@hugedomains.com\\nIsDeleted: false\\nLastName: Ikringill\\nIsEmailBounced: false\\nFirstName: Meghann\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: bgrinvalray@fda.gov\\nIsDeleted: false\\nLastName: Grinval\\nIsEmailBounced: false\\nFirstName: Berti\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: aollanderhr7@cam.ac.uk\\nIsDeleted: false\\nLastName: Ollander\\nIsEmailBounced: false\\nFirstName: Annemarie\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: rwhitesideq38@gravatar.com\\nIsDeleted: false\\nLastName: Whiteside\\nIsEmailBounced: false\\nFirstName: Rolando\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: vkrafthmz@techcrunch.com\\nIsDeleted: false\\nLastName: Kraft\\nIsEmailBounced: false\\nFirstName: Vidovik\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: jhillaut@4shared.com\\nIsDeleted: false\\nLastName: Hill\\nIsEmailBounced: false\\nFirstName: Janel\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: lralstonycs@discovery.com\\nIsDeleted: false\\nLastName: Ralston\\nIsEmailBounced: false\\nFirstName: Lorrayne\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: blyttlewba@networkadvertising.org\\nIsDeleted: false\\nLastName: Lyttle\\nIsEmailBounced: false\\nFirstName: Ban\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: pplummernvf@technorati.com\\nIsDeleted: false\\nLastName: Plummer\\nIsEmailBounced: false\\nFirstName: Pete\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: babrahamoffxpb@theatlantic.com\\nIsDeleted: false\\nLastName: Abrahamoff\\nIsEmailBounced: false\\nFirstName: Brander\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ahargieym0@homestead.com\\nIsDeleted: false\\nLastName: Hargie\\nIsEmailBounced: false\\nFirstName: Aili\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: hstotthp2@yelp.com\\nIsDeleted: false\\nLastName: Stott\\nIsEmailBounced: false\\nFirstName: Hartley\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: jganniclifftuvj@blinklist.com\\nIsDeleted: false\\nLastName: Ganniclifft\\nIsEmailBounced: false\\nFirstName: Jamima\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ldodelly8q@ed.gov\\nIsDeleted: false\\nLastName: Dodell\\nIsEmailBounced: false\\nFirstName: Lynde\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: rmilner3cp@smh.com.au\\nIsDeleted: false\\nLastName: Milner\\nIsEmailBounced: false\\nFirstName: Ralph\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: gghiriardellic19@state.tx.us\\nIsDeleted: false\\nLastName: Ghiriardelli\\nIsEmailBounced: false\\nFirstName: Garv\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: rhubatschfpu@nature.com\\nIsDeleted: false\\nLastName: Hubatsch\\nIsEmailBounced: false\\nFirstName: Rose\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: mtrenholme1ws@quantcast.com\\nIsDeleted: false\\nLastName: Trenholme\\nIsEmailBounced: false\\nFirstName: Mariejeanne\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: jmussettpbd@over-blog.com\\nIsDeleted: false\\nLastName: Mussett\\nIsEmailBounced: false\\nFirstName: Juliann\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: bgoroni145@illinois.edu\\nIsDeleted: false\\nLastName: Goroni\\nIsEmailBounced: false\\nFirstName: Bernarr\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: afalls3ph@theguardian.com\\nIsDeleted: false\\nLastName: Falls\\nIsEmailBounced: false\\nFirstName: Angelia\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: lswettjoi@go.com\\nIsDeleted: false\\nLastName: Swett\\nIsEmailBounced: false\\nFirstName: Levon\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: emullinsz38@dailymotion.com\\nIsDeleted: false\\nLastName: Mullins\\nIsEmailBounced: false\\nFirstName: Elsa\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ibernettehco@ebay.co.uk\\nIsDeleted: false\\nLastName: Bernette\\nIsEmailBounced: false\\nFirstName: Ingrid\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: trisleybtt@simplemachines.org\\nIsDeleted: false\\nLastName: Risley\\nIsEmailBounced: false\\nFirstName: Toma\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: rgypsonqx1@goodreads.com\\nIsDeleted: false\\nLastName: Gypson\\nIsEmailBounced: false\\nFirstName: Reed\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: cposvneri28@jiathis.com\\nIsDeleted: false\\nLastName: Posvner\\nIsEmailBounced: false\\nFirstName: Culley\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: awilmut2rz@geocities.jp\\nIsDeleted: false\\nLastName: Wilmut\\nIsEmailBounced: false\\nFirstName: Andy\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: aluckwellra5@exblog.jp\\nIsDeleted: false\\nLastName: Luckwell\\nIsEmailBounced: false\\nFirstName: Andreana\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: irollings26j@timesonline.co.uk\\nIsDeleted: false\\nLastName: Rollings\\nIsEmailBounced: false\\nFirstName: Ibrahim\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: gspireqpd@g.co\\nIsDeleted: false\\nLastName: Spire\\nIsEmailBounced: false\\nFirstName: Gaelan\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: sbezleyk2y@acquirethisname.com\\nIsDeleted: false\\nLastName: Bezley\\nIsEmailBounced: false\\nFirstName: Sindee\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: icollerrr@flickr.com\\nIsDeleted: false\\nLastName: Coller\\nIsEmailBounced: false\\nFirstName: Inesita\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: kfolliott1bo@nature.com\\nIsDeleted: false\\nLastName: Folliott\\nIsEmailBounced: false\\nFirstName: Kennan\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: kroofjfo@gnu.org\\nIsDeleted: false\\nLastName: Roof\\nIsEmailBounced: false\\nFirstName: Karlik\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: lcovotti8s4@rediff.com\\nIsDeleted: false\\nLastName: Covotti\\nIsEmailBounced: false\\nFirstName: Lucho\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: gpatriskson1rs@census.gov\\nIsDeleted: false\\nLastName: Patriskson\\nIsEmailBounced: false\\nFirstName: Gardener\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: spidgleyqvw@usgs.gov\\nIsDeleted: false\\nLastName: Pidgley\\nIsEmailBounced: false\\nFirstName: Simona\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: cbecarrak0i@over-blog.com\\nIsDeleted: false\\nLastName: Becarra\\nIsEmailBounced: false\\nFirstName: Cally\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: aparkman9td@bbc.co.uk\\nIsDeleted: false\\nLastName: Parkman\\nIsEmailBounced: false\\nFirstName: Agneta\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: bboddingtonhn@quantcast.com\\nIsDeleted: false\\nLastName: Boddington\\nIsEmailBounced: false\\nFirstName: Betta\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: dcasementx0p@cafepress.com\\nIsDeleted: false\\nLastName: Casement\\nIsEmailBounced: false\\nFirstName: Dannie\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: hzornbhe@latimes.com\\nIsDeleted: false\\nLastName: Zorn\\nIsEmailBounced: false\\nFirstName: Haleigh\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: cfifieldbjb@blogspot.com\\nIsDeleted: false\\nLastName: Fifield\\nIsEmailBounced: false\\nFirstName: Christalle\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ddewerson4t3@skype.com\\nIsDeleted: false\\nLastName: Dewerson\\nIsEmailBounced: false\\nFirstName: Dyann\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: khullock52p@sohu.com\\nIsDeleted: false\\nLastName: Hullock\\nIsEmailBounced: false\\nFirstName: Kellina\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: tfremantle32n@bandcamp.com\\nIsDeleted: false\\nLastName: Fremantle\\nIsEmailBounced: false\\nFirstName: Turner\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: sbernardtylp@nps.gov\\nIsDeleted: false\\nLastName: Bernardt\\nIsEmailBounced: false\\nFirstName: Selina\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: smcgettigan8kk@slideshare.net\\nIsDeleted: false\\nLastName: McGettigan\\nIsEmailBounced: false\\nFirstName: Sada\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: wdelafontvgn@businesswire.com\\nIsDeleted: false\\nLastName: Delafont\\nIsEmailBounced: false\\nFirstName: West\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: lbelsher9ne@indiatimes.com\\nIsDeleted: false\\nLastName: Belsher\\nIsEmailBounced: false\\nFirstName: Lou\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: cgoody27y@blogtalkradio.com\\nIsDeleted: false\\nLastName: Goody\\nIsEmailBounced: false\\nFirstName: Colene\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: cstodejzz@ucoz.ru\\nIsDeleted: false\\nLastName: Stode\\nIsEmailBounced: false\\nFirstName: Curcio\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: abromidgejb@china.com.cn\\nIsDeleted: false\\nLastName: Bromidge\\nIsEmailBounced: false\\nFirstName: Ariela\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ldelgardilloqvp@xrea.com\\nIsDeleted: false\\nLastName: Delgardillo\\nIsEmailBounced: false\\nFirstName: Lauralee\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: dcroal9t4@businessinsider.com\\nIsDeleted: false\\nLastName: Croal\\nIsEmailBounced: false\\nFirstName: Devlin\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: dclarageqzb@wordpress.com\\nIsDeleted: false\\nLastName: Clarage\\nIsEmailBounced: false\\nFirstName: Dre\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: dthirlwall3jf@taobao.com\\nIsDeleted: false\\nLastName: Thirlwall\\nIsEmailBounced: false\\nFirstName: Dareen\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: tkeddie2lj@wiley.com\\nIsDeleted: false\\nLastName: Keddie\\nIsEmailBounced: false\\nFirstName: Tandi\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: jrimingtoni3i@istockphoto.com\\nIsDeleted: false\\nLastName: Rimington\\nIsEmailBounced: false\\nFirstName: Judy\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: gtroynet@slashdot.org\\nIsDeleted: false\\nLastName: Troy\\nIsEmailBounced: false\\nFirstName: Gail\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: ebunneyh0n@meetup.com\\nIsDeleted: false\\nLastName: Bunney\\nIsEmailBounced: false\\nFirstName: Efren\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: yhaken8p3@slate.com\\nIsDeleted: false\\nLastName: Haken\\nIsEmailBounced: false\\nFirstName: Yard\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: nolliffeq6q@biblegateway.com\\nIsDeleted: false\\nLastName: Olliffe\\nIsEmailBounced: false\\nFirstName: Nani\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: bgalia9jz@odnoklassniki.ru\\nIsDeleted: false\\nLastName: Galia\\nIsEmailBounced: false\\nFirstName: Berrie\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: djedrzej3v1@google.com\\nIsDeleted: false\\nLastName: Jedrzej\\nIsEmailBounced: false\\nFirstName: Deanne\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: mcamiesh1t@fc2.com\\nIsDeleted: false\\nLastName: Camies\\nIsEmailBounced: false\\nFirstName: Mikaela\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: csunshineqni@state.tx.us\\nIsDeleted: false\\nLastName: Sunshine\\nIsEmailBounced: false\\nFirstName: Curtis\\nIsPriorityRecord: false\\nCleanStatus: Pending\",\n \"Email: fiannellib46@marriott.com\\nIsDeleted: false\\nLastName: Iannelli\\nIsEmailBounced: false\\nFirstName: Felicio\\nIsPriorityRecord: false\\nCleanStatus: Pending\"\n ],\n \"semantic_identifier\": \"Voonder\",\n \"metadata\": {\"object_type\": \"Account\"},\n \"primary_owners\": {\"email\": \"hagen@danswer.ai\", \"first_name\": \"Hagen\", \"last_name\": \"oneill\"},\n \"secondary_owners\": null,\n \"title\": null\n}\n" }, { "path": "backend/tests/daily/connectors/sharepoint/test_sharepoint_connector.py", "content": "import os\nimport time\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom datetime import timezone\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.sharepoint.connector import SharepointAuthMethod\nfrom onyx.connectors.sharepoint.connector import SharepointConnector\nfrom onyx.db.enums import HierarchyNodeType\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n# NOTE: Sharepoint site for tests is \"sharepoint-tests\"\n\n\n@dataclass\nclass ExpectedDocument:\n semantic_identifier: str\n content: str\n folder_path: str | None = None\n library: str = \"Shared Documents\" # Default to main library\n expected_link_substrings: list[str] | None = None\n\n\nEXPECTED_DOCUMENTS = [\n ExpectedDocument(\n semantic_identifier=\"test1.docx\",\n content=\"test1\",\n folder_path=\"test\",\n expected_link_substrings=[\"_layouts/15/Doc.aspx\", \"file=test1.docx\"],\n ),\n ExpectedDocument(\n semantic_identifier=\"test2.docx\",\n content=\"test2\",\n folder_path=\"test/nested with spaces\",\n expected_link_substrings=[\"_layouts/15/Doc.aspx\", \"file=test2.docx\"],\n ),\n ExpectedDocument(\n semantic_identifier=\"should-not-index-on-specific-folder.docx\",\n content=\"should-not-index-on-specific-folder\",\n folder_path=None, # root folder\n expected_link_substrings=[\n \"_layouts/15/Doc.aspx\",\n \"file=should-not-index-on-specific-folder.docx\",\n ],\n ),\n ExpectedDocument(\n semantic_identifier=\"other.docx\",\n content=\"other\",\n folder_path=None,\n library=\"Other Library\",\n expected_link_substrings=[\"_layouts/15/Doc.aspx\", \"file=other.docx\"],\n ),\n]\n\nEXPECTED_PAGES = [\n ExpectedDocument(\n semantic_identifier=\"CollabHome\",\n content=(\n \"# Home\\n\\nDisplay recent news.\\n\\n## News\\n\\nShow recent activities from your site\\n\\n\"\n \"## Site activity\\n\\n## Quick links\\n\\nLearn about a team site\\n\\nLearn how to add a page\\n\\n\"\n \"Add links to important documents and pages.\\n\\n## Quick links\\n\\nDocuments\\n\\n\"\n \"Add a document library\\n\\n## Document library\"\n ),\n folder_path=None,\n expected_link_substrings=[\"SitePages/CollabHome.aspx\"],\n ),\n ExpectedDocument(\n semantic_identifier=\"Home\",\n content=\"# Home\",\n folder_path=None,\n expected_link_substrings=[\"SitePages/Home.aspx\"],\n ),\n]\n\n\ndef verify_document_metadata(doc: Document) -> None:\n \"\"\"Verify common metadata that should be present on all documents.\"\"\"\n assert isinstance(doc.doc_updated_at, datetime)\n assert doc.doc_updated_at.tzinfo == timezone.utc\n assert doc.source == DocumentSource.SHAREPOINT\n assert doc.primary_owners is not None\n assert len(doc.primary_owners) == 1\n owner = doc.primary_owners[0]\n assert owner.display_name is not None\n assert owner.email is not None\n\n\ndef verify_document_content(doc: Document, expected: ExpectedDocument) -> None:\n \"\"\"Verify a document matches its expected content.\"\"\"\n assert doc.semantic_identifier == expected.semantic_identifier\n assert len(doc.sections) == 1\n assert doc.sections[0].text is not None\n assert expected.content == doc.sections[0].text\n\n if expected.expected_link_substrings is not None:\n actual_link = doc.sections[0].link\n assert actual_link is not None, (\n f\"Expected section link containing {expected.expected_link_substrings} \"\n f\"for '{expected.semantic_identifier}', but link was None\"\n )\n for substr in expected.expected_link_substrings:\n assert substr in actual_link, (\n f\"Section link for '{expected.semantic_identifier}' \"\n f\"missing expected substring '{substr}', \"\n f\"actual link: '{actual_link}'\"\n )\n\n verify_document_metadata(doc)\n\n\ndef find_document(documents: list[Document], semantic_identifier: str) -> Document:\n \"\"\"Find a document by its semantic identifier.\"\"\"\n matching_docs = [\n d for d in documents if d.semantic_identifier == semantic_identifier\n ]\n assert (\n len(matching_docs) == 1\n ), f\"Expected exactly one document with identifier {semantic_identifier}\"\n return matching_docs[0]\n\n\n@pytest.fixture\ndef mock_store_image() -> MagicMock:\n \"\"\"Mock store_image_and_create_section to return a predefined ImageSection.\"\"\"\n mock = MagicMock()\n mock.return_value = (\n ImageSection(image_file_id=\"mocked-file-id\", link=\"https://example.com/image\"),\n \"mocked-file-id\",\n )\n return mock\n\n\n@pytest.fixture\ndef sharepoint_credentials() -> dict[str, str]:\n return {\n \"sp_client_id\": os.environ[\"SHAREPOINT_CLIENT_ID\"],\n \"sp_client_secret\": os.environ[\"SHAREPOINT_CLIENT_SECRET\"],\n \"sp_directory_id\": os.environ[\"SHAREPOINT_CLIENT_DIRECTORY_ID\"],\n }\n\n\ndef test_sharepoint_connector_all_sites__docs_only(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n # Initialize connector with no sites\n connector = SharepointConnector(\n include_site_pages=False, include_site_documents=True\n )\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Not asserting expected sites because that can change in test tenant at any time\n # Finding any docs is good enough to verify that the connector is working\n document_batches = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n )\n assert document_batches, \"Should find documents from all sites\"\n\n\ndef test_sharepoint_connector_all_sites__pages_only(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n # Initialize connector with no docs\n connector = SharepointConnector(\n include_site_pages=True, include_site_documents=False\n )\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Not asserting expected sites because that can change in test tenant at any time\n # Finding any docs is good enough to verify that the connector is working\n document_batches = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n )\n assert document_batches, \"Should find site pages from all sites\"\n\n\ndef test_sharepoint_connector_specific_folder(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n # Initialize connector with the test site URL and specific folder\n connector = SharepointConnector(\n sites=[os.environ[\"SHAREPOINT_SITE\"] + \"/Shared Documents/test\"],\n include_site_pages=False,\n include_site_documents=True,\n )\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Get all documents\n found_documents: list[Document] = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n ).documents\n\n # Should only find documents in the test folder\n test_folder_docs = [\n doc\n for doc in EXPECTED_DOCUMENTS\n if doc.folder_path and doc.folder_path.startswith(\"test\")\n ]\n assert len(found_documents) == len(\n test_folder_docs\n ), \"Should only find documents in test folder\"\n\n # Verify each expected document\n for expected in test_folder_docs:\n doc = find_document(found_documents, expected.semantic_identifier)\n verify_document_content(doc, expected)\n\n\ndef test_sharepoint_connector_root_folder__docs_only(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n # Initialize connector with the base site URL\n connector = SharepointConnector(\n sites=[os.environ[\"SHAREPOINT_SITE\"]],\n include_site_pages=False,\n include_site_documents=True,\n )\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Get all documents\n found_documents: list[Document] = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n ).documents\n\n assert len(found_documents) == len(\n EXPECTED_DOCUMENTS\n ), \"Should find all documents in main library\"\n\n # Verify each expected document\n for expected in EXPECTED_DOCUMENTS:\n doc = find_document(found_documents, expected.semantic_identifier)\n verify_document_content(doc, expected)\n\n\ndef test_sharepoint_connector_other_library(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n # Initialize connector with the other library\n connector = SharepointConnector(\n sites=[\n os.environ[\"SHAREPOINT_SITE\"] + \"/Other Library\",\n ],\n include_site_pages=False,\n include_site_documents=True,\n )\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Get all documents\n found_documents: list[Document] = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n ).documents\n expected_documents: list[ExpectedDocument] = [\n doc for doc in EXPECTED_DOCUMENTS if doc.library == \"Other Library\"\n ]\n\n # Should find all documents in `Other Library`\n assert len(found_documents) == len(\n expected_documents\n ), \"Should find all documents in `Other Library`\"\n\n # Verify each expected document\n for expected in expected_documents:\n doc = find_document(found_documents, expected.semantic_identifier)\n verify_document_content(doc, expected)\n\n\ndef test_sharepoint_connector_poll(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n # Initialize connector with the base site URL\n connector = SharepointConnector(sites=[os.environ[\"SHAREPOINT_SITE\"]])\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Set time window to only capture test1.docx (modified at 2025-01-28 20:51:42+00:00)\n start = datetime(\n 2025, 1, 28, 20, 51, 30, tzinfo=timezone.utc\n ) # 12 seconds before\n end = datetime(2025, 1, 28, 20, 51, 50, tzinfo=timezone.utc) # 8 seconds after\n\n # Get documents within the time window\n found_documents: list[Document] = load_all_from_connector(\n connector=connector,\n start=start.timestamp(),\n end=end.timestamp(),\n ).documents\n\n # Should only find test1.docx\n assert (\n len(found_documents) == 1\n ), \"Should only find one document in the time window\"\n doc = found_documents[0]\n assert doc.semantic_identifier == \"test1.docx\"\n verify_document_content(\n doc,\n next(\n d for d in EXPECTED_DOCUMENTS if d.semantic_identifier == \"test1.docx\"\n ),\n )\n\n\ndef test_sharepoint_connector_pages(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n connector = SharepointConnector(\n sites=[os.environ[\"SHAREPOINT_SITE\"]],\n include_site_pages=True,\n include_site_documents=False,\n )\n\n connector.load_credentials(sharepoint_credentials)\n\n found_documents = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n ).documents\n\n assert len(found_documents) == len(\n EXPECTED_PAGES\n ), \"Should find all pages in test site\"\n\n for expected in EXPECTED_PAGES:\n doc = find_document(found_documents, expected.semantic_identifier)\n verify_document_content(doc, expected)\n\n\ndef verify_hierarchy_nodes(\n hierarchy_nodes: list[HierarchyNode],\n documents: list[Document],\n expected_site_url: str,\n) -> None:\n \"\"\"Verify hierarchy nodes have correct structure and relationships.\"\"\"\n # Build a set of all raw_node_ids for parent validation\n all_node_ids = {node.raw_node_id for node in hierarchy_nodes}\n\n # Track nodes by type\n site_nodes = [n for n in hierarchy_nodes if n.node_type == HierarchyNodeType.SITE]\n drive_nodes = [n for n in hierarchy_nodes if n.node_type == HierarchyNodeType.DRIVE]\n folder_nodes = [\n n for n in hierarchy_nodes if n.node_type == HierarchyNodeType.FOLDER\n ]\n\n # Verify we have at least one site node\n assert len(site_nodes) >= 1, \"Should have at least one SITE hierarchy node\"\n assert len(drive_nodes) >= 1, \"Should have at least one DRIVE hierarchy node\"\n assert len(folder_nodes) >= 1, \"Should have at least one FOLDER hierarchy node\"\n\n # Verify expected site is in hierarchy\n site_node_ids = {n.raw_node_id for n in site_nodes}\n assert (\n expected_site_url in site_node_ids\n ), f\"Expected site {expected_site_url} not found in hierarchy nodes. Found sites: {site_node_ids}\"\n\n # Verify no duplicate raw_node_ids\n assert len(all_node_ids) == len(\n hierarchy_nodes\n ), \"Should not have duplicate hierarchy nodes\"\n\n # Verify all hierarchy nodes have required fields\n for node in hierarchy_nodes:\n assert node.raw_node_id, \"All nodes should have raw_node_id\"\n assert node.display_name, \"All nodes should have display_name\"\n assert node.link, \"All nodes should have link\"\n assert node.node_type in [\n HierarchyNodeType.SITE,\n HierarchyNodeType.DRIVE,\n HierarchyNodeType.FOLDER,\n ], f\"Unexpected node type: {node.node_type}\"\n\n # Verify parent relationships\n for node in hierarchy_nodes:\n if node.node_type == HierarchyNodeType.SITE:\n # Sites should have no parent (direct child of SOURCE)\n assert node.raw_parent_id is None, \"SITE nodes should have no parent\"\n elif node.node_type == HierarchyNodeType.DRIVE:\n # Drives should have a site as parent\n assert node.raw_parent_id is not None, \"DRIVE nodes should have a parent\"\n assert (\n node.raw_parent_id in site_node_ids\n ), f\"DRIVE parent {node.raw_parent_id} should be a SITE node\"\n elif node.node_type == HierarchyNodeType.FOLDER:\n # Folders should have either a drive or another folder as parent\n assert node.raw_parent_id is not None, \"FOLDER nodes should have a parent\"\n assert (\n node.raw_parent_id in all_node_ids\n ), f\"FOLDER parent {node.raw_parent_id} should exist in hierarchy\"\n\n # Verify documents have parent_hierarchy_raw_node_id set\n for doc in documents:\n if doc.parent_hierarchy_raw_node_id:\n assert (\n doc.parent_hierarchy_raw_node_id in all_node_ids\n ), f\"Document {doc.semantic_identifier} parent {doc.parent_hierarchy_raw_node_id} should exist in hierarchy\"\n\n\ndef test_sharepoint_connector_hierarchy_nodes(\n mock_get_unstructured_api_key: MagicMock, # noqa: ARG001\n mock_store_image: MagicMock,\n sharepoint_credentials: dict[str, str],\n) -> None:\n \"\"\"Test that the SharePoint connector yields proper hierarchy nodes.\"\"\"\n with patch(\n \"onyx.connectors.sharepoint.connector.store_image_and_create_section\",\n mock_store_image,\n ):\n site_url = os.environ[\"SHAREPOINT_SITE\"]\n\n # Initialize connector with the test site\n connector = SharepointConnector(\n sites=[site_url],\n include_site_pages=True,\n include_site_documents=True,\n )\n\n # Load credentials\n connector.load_credentials(sharepoint_credentials)\n\n # Get all documents and hierarchy nodes\n result = load_all_from_connector(\n connector=connector,\n start=0,\n end=time.time(),\n )\n\n found_documents = result.documents\n hierarchy_nodes = result.hierarchy_nodes\n\n # Should have hierarchy nodes\n assert len(hierarchy_nodes) > 0, \"Should have hierarchy nodes\"\n\n # Verify hierarchy structure\n verify_hierarchy_nodes(hierarchy_nodes, found_documents, site_url)\n\n # Verify we have the expected node types\n node_types = {n.node_type for n in hierarchy_nodes}\n assert HierarchyNodeType.SITE in node_types, \"Should have SITE nodes\"\n assert HierarchyNodeType.DRIVE in node_types, \"Should have DRIVE nodes\"\n\n # Should have folder nodes if documents are in folders\n docs_in_folders = [d for d in EXPECTED_DOCUMENTS if d.folder_path]\n if docs_in_folders:\n assert (\n HierarchyNodeType.FOLDER in node_types\n ), \"Should have FOLDER nodes since documents are in folders\"\n\n # Verify all documents have parent_hierarchy_raw_node_id set\n for doc in found_documents:\n assert (\n doc.parent_hierarchy_raw_node_id is not None\n ), f\"Document {doc.semantic_identifier} should have parent_hierarchy_raw_node_id set\"\n\n\n@pytest.fixture\ndef sharepoint_cert_credentials() -> dict[str, str]:\n return {\n \"authentication_method\": SharepointAuthMethod.CERTIFICATE.value,\n \"sp_client_id\": os.environ[\"PERM_SYNC_SHAREPOINT_CLIENT_ID\"],\n \"sp_private_key\": os.environ[\"PERM_SYNC_SHAREPOINT_PRIVATE_KEY\"],\n \"sp_certificate_password\": os.environ[\n \"PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD\"\n ],\n \"sp_directory_id\": os.environ[\"PERM_SYNC_SHAREPOINT_DIRECTORY_ID\"],\n }\n\n\ndef test_resolve_tenant_domain_from_site_urls(\n sharepoint_cert_credentials: dict[str, str],\n) -> None:\n \"\"\"Verify that certificate auth resolves the tenant domain from site URLs\n without calling the /organization endpoint.\"\"\"\n site_url = os.environ[\"SHAREPOINT_SITE\"]\n connector = SharepointConnector(sites=[site_url])\n connector.load_credentials(sharepoint_cert_credentials)\n\n assert connector.sp_tenant_domain is not None\n assert len(connector.sp_tenant_domain) > 0\n # The tenant domain should match the first label of the site URL hostname\n from urllib.parse import urlsplit\n\n expected = urlsplit(site_url).hostname.split(\".\")[0] # type: ignore\n assert connector.sp_tenant_domain == expected\n\n\ndef test_resolve_tenant_domain_from_root_site(\n sharepoint_cert_credentials: dict[str, str],\n) -> None:\n \"\"\"Verify that certificate auth resolves the tenant domain via the root\n site endpoint when no site URLs are configured.\"\"\"\n connector = SharepointConnector(sites=[])\n connector.load_credentials(sharepoint_cert_credentials)\n\n assert connector.sp_tenant_domain is not None\n assert len(connector.sp_tenant_domain) > 0\n" }, { "path": "backend/tests/daily/connectors/slab/test_slab_connector.py", "content": "import json\nimport os\nimport time\nfrom pathlib import Path\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.slab.connector import SlabConnector\n\n\ndef load_test_data(file_name: str = \"test_slab_data.json\") -> dict[str, str]:\n current_dir = Path(__file__).parent\n with open(current_dir / file_name, \"r\") as f:\n return json.load(f)\n\n\n@pytest.fixture\ndef slab_connector() -> SlabConnector:\n connector = SlabConnector(\n base_url=\"https://onyx-test.slab.com/\",\n )\n connector.load_credentials(\n {\n \"slab_bot_token\": os.environ[\"SLAB_BOT_TOKEN\"],\n }\n )\n return connector\n\n\n@pytest.mark.xfail(\n reason=(\n \"Need a test account with a slab subscription to run this test.Trial only lasts 14 days.\"\n )\n)\ndef test_slab_connector_basic(slab_connector: SlabConnector) -> None:\n all_docs: list[Document] = []\n target_test_doc_id = \"jcp6cohu\"\n target_test_doc: Document | None = None\n for doc_batch in slab_connector.poll_source(0, time.time()):\n for doc in doc_batch:\n if not isinstance(doc, Document):\n continue\n all_docs.append(doc)\n if doc.id == target_test_doc_id:\n target_test_doc = doc\n\n assert len(all_docs) == 6\n assert target_test_doc is not None\n\n desired_test_data = load_test_data()\n assert (\n target_test_doc.semantic_identifier == desired_test_data[\"semantic_identifier\"]\n )\n assert target_test_doc.source == DocumentSource.SLAB\n assert target_test_doc.metadata == {}\n assert target_test_doc.primary_owners is None\n assert target_test_doc.secondary_owners is None\n assert target_test_doc.title is None\n assert target_test_doc.from_ingestion_api is False\n assert target_test_doc.additional_info is None\n\n assert len(target_test_doc.sections) == 1\n section = target_test_doc.sections[0]\n # Need to replace the weird apostrophe with a normal one\n assert section.text is not None\n assert section.text.replace(\"\\u2019\", \"'\") == desired_test_data[\"section_text\"]\n assert section.link == desired_test_data[\"link\"]\n\n\n@pytest.mark.xfail(\n reason=(\n \"Need a test account with a slab subscription to run this test.Trial only lasts 14 days.\"\n )\n)\ndef test_slab_connector_slim(slab_connector: SlabConnector) -> None:\n # Get all doc IDs from the full connector\n all_full_doc_ids = set()\n for doc_batch in slab_connector.load_from_state():\n all_full_doc_ids.update(\n [doc.id for doc in doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # Get all doc IDs from the slim connector\n all_slim_doc_ids = set()\n for slim_doc_batch in slab_connector.retrieve_all_slim_docs_perm_sync():\n all_slim_doc_ids.update(\n [doc.id for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # The set of full doc IDs should be always be a subset of the slim doc IDs\n assert all_full_doc_ids.issubset(all_slim_doc_ids)\n" }, { "path": "backend/tests/daily/connectors/slab/test_slab_data.json", "content": "{\n \"section_text\": \"Learn about Posts\\nWelcome\\nThis is a post, where you can edit, share, and collaborate in real time with your team. We'd love to show you how it works!\\nReading and editing\\nClick the mode button to toggle between read and edit modes. You can only make changes to a post when editing.\\nOrganize your posts\\nWhen in edit mode, you can add topics to a post, which will keep it organized for the right 👀 to see.\\nSmart mentions\\nMentions are references to users, posts, topics and third party tools that show details on hover. Paste in a link for automatic conversion.\\nLook back in time\\nYou are ready to begin writing. You can always bring back this tour in the help menu.\\nGreat job!\\nYou are ready to begin writing. You can always bring back this tour in the help menu.\\n\\n\",\n \"link\": \"https://onyx-test.slab.com/posts/learn-about-posts-jcp6cohu\",\n \"semantic_identifier\": \"Learn about Posts\"\n} " }, { "path": "backend/tests/daily/connectors/slack/conftest.py", "content": "import os\nfrom collections.abc import Generator\nfrom unittest.mock import MagicMock\n\nimport pytest\nfrom pytest import FixtureRequest\nfrom slack_sdk import WebClient\n\nfrom onyx.connectors.credentials_provider import OnyxStaticCredentialsProvider\nfrom onyx.connectors.slack.connector import SlackConnector\nfrom shared_configs.contextvars import get_current_tenant_id\n\n\n@pytest.fixture\ndef mock_slack_client() -> MagicMock:\n mock = MagicMock(spec=WebClient)\n return mock\n\n\n@pytest.fixture\ndef slack_connector(\n request: FixtureRequest,\n mock_slack_client: MagicMock,\n slack_credentials_provider: OnyxStaticCredentialsProvider,\n) -> Generator[SlackConnector]:\n channel: str | None = request.param if hasattr(request, \"param\") else None\n connector = SlackConnector(\n channels=[channel] if channel else None,\n channel_regex_enabled=False,\n use_redis=False,\n )\n connector.client = mock_slack_client\n connector.set_credentials_provider(credentials_provider=slack_credentials_provider)\n yield connector\n\n\n@pytest.fixture\ndef slack_credentials_provider() -> OnyxStaticCredentialsProvider:\n CI_ENV_VAR = \"SLACK_BOT_TOKEN\"\n LOCAL_ENV_VAR = \"ONYX_BOT_SLACK_BOT_TOKEN\"\n\n slack_bot_token = os.environ.get(CI_ENV_VAR, os.environ.get(LOCAL_ENV_VAR))\n if not slack_bot_token:\n raise RuntimeError(\n f\"No slack credentials found; either set the {CI_ENV_VAR} env-var or the {LOCAL_ENV_VAR} env-var\"\n )\n\n return OnyxStaticCredentialsProvider(\n tenant_id=get_current_tenant_id(),\n connector_name=\"slack\",\n credential_json={\n \"slack_bot_token\": slack_bot_token,\n },\n )\n" }, { "path": "backend/tests/daily/connectors/slack/test_slack_connector.py", "content": "import time\n\nimport pytest\n\nfrom onyx.connectors.slack.connector import SlackConnector\nfrom onyx.db.enums import HierarchyNodeType\nfrom tests.daily.connectors.utils import load_all_from_connector\nfrom tests.daily.connectors.utils import to_sections\nfrom tests.daily.connectors.utils import to_text_sections\n\n\ndef test_validate_slack_connector_settings(\n slack_connector: SlackConnector,\n) -> None:\n slack_connector.validate_connector_settings()\n\n\n@pytest.mark.parametrize(\n \"slack_connector,expected_messages,expected_channel_name\",\n [\n [\"general\", set(), \"general\"],\n [\"#general\", set(), \"general\"],\n [\n \"daily-connector-test-channel\",\n set(\n [\n \"Hello, world!\",\n \"\",\n \"Reply!\",\n \"Testing again...\",\n ]\n ),\n \"daily-connector-test-channel\",\n ],\n [\n \"#daily-connector-test-channel\",\n set(\n [\n \"Hello, world!\",\n \"\",\n \"Reply!\",\n \"Testing again...\",\n ]\n ),\n \"daily-connector-test-channel\",\n ],\n ],\n indirect=[\"slack_connector\"],\n)\ndef test_indexing_channels_with_message_count(\n slack_connector: SlackConnector,\n expected_messages: set[str],\n expected_channel_name: str,\n) -> None:\n if not slack_connector.client:\n raise RuntimeError(\"Web client must be defined\")\n\n result = load_all_from_connector(\n connector=slack_connector,\n start=0.0,\n end=time.time(),\n )\n docs = result.documents\n hierarchy_nodes = result.hierarchy_nodes\n\n # Verify messages\n actual_messages = set(to_text_sections(to_sections(docs)))\n assert expected_messages == actual_messages\n\n # Verify hierarchy nodes exist\n assert len(hierarchy_nodes) > 0, \"Expected at least one hierarchy node (channel)\"\n\n # Verify all hierarchy nodes are channels with correct structure\n for node in hierarchy_nodes:\n assert node.node_type == HierarchyNodeType.CHANNEL\n assert node.raw_parent_id is None # Direct child of SOURCE\n assert node.raw_node_id # Channel ID must be present\n assert node.display_name.startswith(\"#\") # e.g. \"#general\"\n\n # Verify the expected channel appears in the hierarchy nodes\n channel_display_names = {node.display_name for node in hierarchy_nodes}\n assert (\n f\"#{expected_channel_name}\" in channel_display_names\n ), f\"Expected channel '#{expected_channel_name}' not found in hierarchy nodes. Found: {channel_display_names}\"\n\n # Verify documents reference their parent channel\n channel_ids = {node.raw_node_id for node in hierarchy_nodes}\n for doc in docs:\n assert (\n doc.parent_hierarchy_raw_node_id is not None\n ), f\"Document '{doc.id}' has no parent_hierarchy_raw_node_id\"\n assert doc.parent_hierarchy_raw_node_id in channel_ids, (\n f\"Document '{doc.id}' has parent_hierarchy_raw_node_id=\"\n f\"'{doc.parent_hierarchy_raw_node_id}' which is not in \"\n f\"hierarchy nodes: {channel_ids}\"\n )\n\n\n@pytest.mark.parametrize(\n \"slack_connector\",\n [\n # w/o hashtag\n \"doesnt-exist\",\n # w/ hashtag\n \"#doesnt-exist\",\n ],\n indirect=True,\n)\ndef test_indexing_channels_that_dont_exist(\n slack_connector: SlackConnector,\n) -> None:\n if not slack_connector.client:\n raise RuntimeError(\"Web client must be defined\")\n\n with pytest.raises(\n ValueError,\n match=r\"Channel '.*' not found in workspace.*\",\n ):\n load_all_from_connector(\n connector=slack_connector,\n start=0.0,\n end=time.time(),\n ).documents\n" }, { "path": "backend/tests/daily/connectors/slack/test_slack_perm_sync.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.connectors.slack.connector import SlackConnector\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\nPUBLIC_CHANNEL_NAME = \"#daily-connector-test-channel\"\nPRIVATE_CHANNEL_NAME = \"#private-channel\"\nPRIVATE_CHANNEL_USERS = [\n \"admin@onyx-test.com\",\n \"test_user_1@onyx-test.com\",\n # user 2 added via a group\n \"test_user_2@onyx-test.com\",\n]\n\n# Predates any test workspace messages, so the result set should match\n# the \"no start time\" case while exercising the oldest= parameter.\nOLDEST_TS_2016 = datetime(2016, 1, 1, tzinfo=timezone.utc).timestamp()\n\npytestmark = pytest.mark.usefixtures(\"enable_ee\")\n\n\n@pytest.mark.parametrize(\n \"slack_connector\",\n [\n PUBLIC_CHANNEL_NAME,\n ],\n indirect=True,\n)\ndef test_load_from_checkpoint_access__public_channel(\n slack_connector: SlackConnector,\n) -> None:\n \"\"\"Test that load_from_checkpoint returns correct access information for documents.\"\"\"\n if not slack_connector.client:\n raise RuntimeError(\"Web client must be defined\")\n\n docs = load_all_from_connector(\n connector=slack_connector,\n start=0.0,\n end=time.time(),\n include_permissions=True,\n ).documents\n\n # We should have at least some documents\n assert len(docs) > 0, \"Expected to find at least one document\"\n\n for doc in docs:\n assert (\n doc.external_access is not None\n ), f\"Document {doc.id} should have external_access when using perm sync\"\n assert (\n doc.external_access.is_public is True\n ), f\"Document {doc.id} should have public access when using perm sync\"\n assert (\n doc.external_access.external_user_emails == set()\n ), f\"Document {doc.id} should have no external user emails when using perm sync\"\n assert (\n doc.external_access.external_user_group_ids == set()\n ), f\"Document {doc.id} should have no external user group ids when using perm sync\"\n\n\n@pytest.mark.parametrize(\n \"slack_connector\",\n [\n PRIVATE_CHANNEL_NAME,\n ],\n indirect=True,\n)\ndef test_load_from_checkpoint_access__private_channel(\n slack_connector: SlackConnector,\n) -> None:\n \"\"\"Test that load_from_checkpoint returns correct access information for documents.\"\"\"\n if not slack_connector.client:\n raise RuntimeError(\"Web client must be defined\")\n\n docs = load_all_from_connector(\n connector=slack_connector,\n start=0.0,\n end=time.time(),\n include_permissions=True,\n ).documents\n\n # We should have at least some documents\n assert len(docs) > 0, \"Expected to find at least one document\"\n\n for doc in docs:\n assert (\n doc.external_access is not None\n ), f\"Document {doc.id} should have external_access when using perm sync\"\n assert (\n doc.external_access.is_public is False\n ), f\"Document {doc.id} should have private access when using perm sync\"\n assert doc.external_access.external_user_emails == set(\n PRIVATE_CHANNEL_USERS\n ), f\"Document {doc.id} should have private channel users when using perm sync\"\n assert (\n doc.external_access.external_user_group_ids == set()\n ), f\"Document {doc.id} should have no external user group ids when using perm sync\"\n\n\n@pytest.mark.parametrize(\n \"slack_connector\",\n [\n PUBLIC_CHANNEL_NAME,\n ],\n indirect=True,\n)\n@pytest.mark.parametrize(\"start_ts\", [None, OLDEST_TS_2016])\ndef test_slim_documents_access__public_channel(\n slack_connector: SlackConnector,\n start_ts: float | None,\n) -> None:\n \"\"\"Test that retrieve_all_slim_docs_perm_sync returns correct access information for slim documents.\"\"\"\n if not slack_connector.client:\n raise RuntimeError(\"Web client must be defined\")\n\n slim_docs_generator = slack_connector.retrieve_all_slim_docs_perm_sync(\n start=start_ts,\n end=time.time(),\n )\n\n # Collect all slim documents from the generator\n all_slim_docs: list[SlimDocument] = []\n for slim_doc_batch in slim_docs_generator:\n all_slim_docs.extend(\n [doc for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # We should have at least some slim documents\n assert len(all_slim_docs) > 0, \"Expected to find at least one slim document\"\n\n for slim_doc in all_slim_docs:\n assert slim_doc.external_access is not None\n assert slim_doc.external_access.is_public is True\n assert slim_doc.external_access.external_user_emails == set()\n assert slim_doc.external_access.external_user_group_ids == set()\n\n\n@pytest.mark.parametrize(\n \"slack_connector\",\n [\n PRIVATE_CHANNEL_NAME,\n ],\n indirect=True,\n)\ndef test_slim_documents_access__private_channel(\n slack_connector: SlackConnector,\n) -> None:\n \"\"\"Test that retrieve_all_slim_docs_perm_sync returns correct access information for slim documents.\"\"\"\n if not slack_connector.client:\n raise RuntimeError(\"Web client must be defined\")\n\n slim_docs_generator = slack_connector.retrieve_all_slim_docs_perm_sync(\n start=None,\n end=time.time(),\n )\n\n # Collect all slim documents from the generator\n all_slim_docs: list[SlimDocument] = []\n for slim_doc_batch in slim_docs_generator:\n all_slim_docs.extend(\n [doc for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # We should have at least some slim documents\n assert len(all_slim_docs) > 0, \"Expected to find at least one slim document\"\n\n for slim_doc in all_slim_docs:\n assert slim_doc.external_access is not None\n assert slim_doc.external_access.is_public is False\n assert slim_doc.external_access.external_user_emails == set(\n PRIVATE_CHANNEL_USERS\n )\n assert slim_doc.external_access.external_user_group_ids == set()\n" }, { "path": "backend/tests/daily/connectors/teams/models.py", "content": "from pydantic import BaseModel\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.models import Document\n\n\nclass TeamsThread(BaseModel):\n thread: str\n external_access: ExternalAccess\n\n @classmethod\n def from_doc(cls, document: Document) -> \"TeamsThread\":\n assert (\n document.external_access\n ), f\"ExternalAccess should always be available, instead got {document=}\"\n\n return cls(\n thread=document.get_text_content(),\n external_access=document.external_access,\n )\n" }, { "path": "backend/tests/daily/connectors/teams/test_teams_connector.py", "content": "import os\nimport time\n\nimport pytest\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.teams.connector import TeamsConnector\nfrom tests.daily.connectors.teams.models import TeamsThread\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\nTEAMS_THREAD = [\n # Posted in \"Public Channel\"\n TeamsThread(\n thread=\"This is the first message in Onyx-Testing ...This is a reply!This is a second reply.Third.4th.5\",\n external_access=ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n ),\n ),\n TeamsThread(\n thread=\"Testing body.\",\n external_access=ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n ),\n ),\n TeamsThread(\n thread=\"Hello, world! Nice to meet you all.\",\n external_access=ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n ),\n ),\n # Posted in \"Private Channel (Raunak is excluded)\"\n TeamsThread(\n thread=\"This is a test post. Raunak should not be able to see this!\",\n external_access=ExternalAccess(\n external_user_emails=set([\"test@danswerai.onmicrosoft.com\"]),\n external_user_group_ids=set(),\n is_public=False,\n ),\n ),\n # Posted in \"Private Channel (Raunak is a member)\"\n TeamsThread(\n thread=\"This is a test post in a private channel that Raunak does have access to! Hello, Raunak!\"\n \"Hello, world! I am just a member in this chat, but not an owner.\",\n external_access=ExternalAccess(\n external_user_emails=set(\n [\"test@danswerai.onmicrosoft.com\", \"raunak@onyx.app\"]\n ),\n external_user_group_ids=set(),\n is_public=False,\n ),\n ),\n # Posted in \"Private Channel (Raunak owns)\"\n TeamsThread(\n thread=\"This is a test post in a private channel that Raunak is an owner of! Whoa!\"\n \"Hello, world! I am an owner of this chat. The power!\",\n external_access=ExternalAccess(\n external_user_emails=set(\n [\"test@danswerai.onmicrosoft.com\", \"raunak@onyx.app\"]\n ),\n external_user_group_ids=set(),\n is_public=False,\n ),\n ),\n]\n\n\n@pytest.fixture\ndef teams_credentials() -> dict[str, str]:\n app_id = os.environ[\"TEAMS_APPLICATION_ID\"]\n dir_id = os.environ[\"TEAMS_DIRECTORY_ID\"]\n secret = os.environ[\"TEAMS_SECRET\"]\n\n return {\n \"teams_client_id\": app_id,\n \"teams_directory_id\": dir_id,\n \"teams_client_secret\": secret,\n }\n\n\n@pytest.fixture\ndef teams_connector(\n teams_credentials: dict[str, str],\n) -> TeamsConnector:\n teams_connector = TeamsConnector(teams=[\"Onyx-Testing\"])\n teams_connector.load_credentials(teams_credentials)\n return teams_connector\n\n\ndef _build_map(threads: list[TeamsThread]) -> dict[str, TeamsThread]:\n map: dict[str, TeamsThread] = {}\n\n for thread in threads:\n assert thread.thread not in map, f\"Duplicate thread found in map; {thread=}\"\n map[thread.thread] = thread\n\n return map\n\n\ndef _assert_is_valid_external_access(\n external_access: ExternalAccess,\n) -> None:\n assert (\n not external_access.external_user_group_ids\n ), f\"{external_access.external_user_group_ids=} should be empty for MS Teams\"\n\n if external_access.is_public:\n assert (\n not external_access.external_user_emails\n ), f\"{external_access.external_user_emails=} should be empty for public channels\"\n else:\n assert (\n external_access.external_user_emails\n ), f\"{external_access.external_user_emails=} should contains at least one user for private channels\"\n\n\n@pytest.mark.parametrize(\n \"expected_teams_threads\",\n [TEAMS_THREAD],\n)\ndef test_loading_all_docs_from_teams_connector(\n teams_connector: TeamsConnector,\n expected_teams_threads: list[TeamsThread],\n) -> None:\n docs = list(\n load_all_from_connector(\n connector=teams_connector,\n start=0.0,\n end=time.time(),\n ).documents\n )\n actual_teams_threads = [TeamsThread.from_doc(doc) for doc in docs]\n actual_teams_threads_map = _build_map(threads=actual_teams_threads)\n expected_teams_threads_map = _build_map(threads=expected_teams_threads)\n\n # Assert that each thread document matches what we expect.\n assert actual_teams_threads_map == expected_teams_threads_map\n\n # Assert that all the `ExternalAccess` instances are well-formed.\n for thread in actual_teams_threads:\n _assert_is_valid_external_access(external_access=thread.external_access)\n\n\ndef test_slim_docs_retrieval_from_teams_connector(\n teams_connector: TeamsConnector,\n) -> None:\n slim_docs = [\n slim_doc\n for slim_doc_batch in teams_connector.retrieve_all_slim_docs_perm_sync()\n for slim_doc in slim_doc_batch\n ]\n\n for slim_doc in slim_docs:\n if isinstance(slim_doc, HierarchyNode):\n continue\n assert (\n slim_doc.external_access\n ), f\"ExternalAccess should always be available, instead got {slim_doc=}\"\n _assert_is_valid_external_access(external_access=slim_doc.external_access)\n\n\ndef test_load_from_checkpoint_with_perm_sync(\n teams_connector: TeamsConnector,\n enable_ee: None, # noqa: ARG001\n) -> None:\n \"\"\"Test that load_from_checkpoint_with_perm_sync returns documents with external_access.\n\n This verifies the CheckpointedConnectorWithPermSync interface is properly implemented.\n \"\"\"\n docs = load_all_from_connector(\n connector=teams_connector,\n start=0.0,\n end=time.time(),\n include_permissions=True, # Uses load_from_checkpoint_with_perm_sync\n ).documents\n\n # We should have at least some documents\n assert len(docs) > 0, \"Expected to find at least one document\"\n\n for doc in docs:\n assert (\n doc.external_access is not None\n ), f\"Document {doc.id} should have external_access when using perm sync\"\n _assert_is_valid_external_access(external_access=doc.external_access)\n" }, { "path": "backend/tests/daily/connectors/utils.py", "content": "from collections.abc import Iterator\nfrom typing import TypeVar\n\nfrom pydantic import BaseModel\n\nfrom onyx.connectors.connector_runner import CheckpointOutputWrapper\nfrom onyx.connectors.interfaces import CheckpointedConnector\nfrom onyx.connectors.interfaces import CheckpointedConnectorWithPermSync\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.models import ConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.models import ImageSection\nfrom onyx.connectors.models import TextSection\n\n_ITERATION_LIMIT = 100_000\n\nCT = TypeVar(\"CT\", bound=ConnectorCheckpoint)\n\n\nclass ConnectorOutput(BaseModel):\n \"\"\"Structured output from loading a connector.\"\"\"\n\n documents: list[Document]\n failures: list[ConnectorFailure]\n hierarchy_nodes: list[HierarchyNode]\n\n model_config = {\"arbitrary_types_allowed\": True}\n\n\ndef load_all_from_connector(\n connector: CheckpointedConnector[CT],\n start: SecondsSinceUnixEpoch,\n end: SecondsSinceUnixEpoch,\n include_permissions: bool = False,\n raise_on_failures: bool = True,\n) -> ConnectorOutput:\n \"\"\"\n Load all documents, hierarchy nodes, and failures from a connector.\n\n Returns a ConnectorOutput with documents, failures, and hierarchy_nodes separated.\n\n Also validates that parent hierarchy nodes are always yielded before their children:\n - For documents: parent must have been yielded before the document\n - For hierarchy nodes: after each batch, validates that all parents in the batch\n have been seen (either in the current batch or a previous batch)\n \"\"\"\n num_iterations = 0\n\n if include_permissions and not isinstance(\n connector, CheckpointedConnectorWithPermSync\n ):\n raise ValueError(\"Connector does not support permission syncing\")\n\n checkpoint = connector.build_dummy_checkpoint()\n documents: list[Document] = []\n failures: list[ConnectorFailure] = []\n hierarchy_nodes: list[HierarchyNode] = []\n\n # Track all seen hierarchy node raw_ids for parent validation\n seen_hierarchy_raw_ids: set[str] = set()\n\n while checkpoint.has_more:\n load_from_checkpoint_generator = (\n connector.load_from_checkpoint_with_perm_sync\n if include_permissions\n and isinstance(connector, CheckpointedConnectorWithPermSync)\n else connector.load_from_checkpoint\n )\n doc_batch_generator = CheckpointOutputWrapper[CT]()(\n load_from_checkpoint_generator(start, end, checkpoint)\n )\n\n # Collect hierarchy nodes from this batch (for end-of-batch validation)\n batch_hierarchy_nodes: list[HierarchyNode] = []\n\n for document, hierarchy_node, failure, next_checkpoint in doc_batch_generator:\n if hierarchy_node is not None:\n hierarchy_nodes.append(hierarchy_node)\n batch_hierarchy_nodes.append(hierarchy_node)\n # Add to seen set immediately so subsequent documents can reference it\n seen_hierarchy_raw_ids.add(hierarchy_node.raw_node_id)\n\n if failure is not None:\n failures.append(failure)\n\n if document is not None and isinstance(document, Document):\n documents.append(document)\n # Validate: document's parent must have been yielded before this document\n if document.parent_hierarchy_raw_node_id is not None:\n if (\n document.parent_hierarchy_raw_node_id\n not in seen_hierarchy_raw_ids\n ):\n raise AssertionError(\n f\"Document '{document.id}' \"\n f\"(semantic_identifier='{document.semantic_identifier}') \"\n f\"has parent_hierarchy_raw_node_id=\"\n f\"'{document.parent_hierarchy_raw_node_id}' \"\n f\"which was not yielded before this document. \"\n f\"Seen hierarchy IDs: {seen_hierarchy_raw_ids}\"\n )\n\n if next_checkpoint is not None:\n checkpoint = next_checkpoint\n\n # End-of-batch validation for hierarchy nodes:\n # Each node's parent must be in the current batch or a previous batch\n batch_hierarchy_raw_ids = {node.raw_node_id for node in batch_hierarchy_nodes}\n for node in batch_hierarchy_nodes:\n if node.raw_parent_id is None:\n continue # Root nodes have no parent\n\n parent_in_current_batch = node.raw_parent_id in batch_hierarchy_raw_ids\n parent_in_previous_batch = node.raw_parent_id in seen_hierarchy_raw_ids\n\n if not parent_in_current_batch and not parent_in_previous_batch:\n raise AssertionError(\n f\"HierarchyNode '{node.raw_node_id}' \"\n f\"(display_name='{node.display_name}') \"\n f\"has raw_parent_id='{node.raw_parent_id}' which was not yielded \"\n f\"in the current batch or any previous batch. \"\n f\"Seen hierarchy IDs: {seen_hierarchy_raw_ids}, \"\n f\"Current batch IDs: {batch_hierarchy_raw_ids}\"\n )\n\n num_iterations += 1\n if num_iterations > _ITERATION_LIMIT:\n raise RuntimeError(\"Too many iterations. Infinite loop?\")\n\n if raise_on_failures and failures:\n raise RuntimeError(f\"Failed to load documents: {failures}\")\n\n return ConnectorOutput(\n documents=documents,\n failures=failures,\n hierarchy_nodes=hierarchy_nodes,\n )\n\n\ndef to_sections(\n documents: list[Document],\n) -> Iterator[TextSection | ImageSection]:\n for doc in documents:\n for section in doc.sections:\n yield section\n\n\ndef to_text_sections(sections: Iterator[TextSection | ImageSection]) -> Iterator[str]:\n for section in sections:\n if isinstance(section, TextSection):\n yield section.text\n" }, { "path": "backend/tests/daily/connectors/web/test_web_connector.py", "content": "from concurrent.futures import ThreadPoolExecutor\n\nimport pytest\n\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.web.connector import WEB_CONNECTOR_VALID_SETTINGS\nfrom onyx.connectors.web.connector import WebConnector\n\nEXPECTED_QUOTE = (\n \"If you can't explain it to a six year old, you don't understand it yourself.\"\n)\n\n\n# NOTE(rkuo): we will probably need to adjust this test to point at our own test site\n# to avoid depending on a third party site\n@pytest.fixture\ndef quotes_to_scroll_web_connector(request: pytest.FixtureRequest) -> WebConnector:\n scroll_before_scraping = request.param\n connector = WebConnector(\n base_url=\"https://quotes.toscrape.com/scroll\",\n web_connector_type=WEB_CONNECTOR_VALID_SETTINGS.SINGLE.value,\n scroll_before_scraping=scroll_before_scraping,\n )\n return connector\n\n\n@pytest.mark.parametrize(\"quotes_to_scroll_web_connector\", [True], indirect=True)\ndef test_web_connector_scroll(quotes_to_scroll_web_connector: WebConnector) -> None:\n all_docs: list[Document] = []\n document_batches = quotes_to_scroll_web_connector.load_from_state()\n for doc_batch in document_batches:\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs.append(doc)\n\n assert len(all_docs) == 1\n doc = all_docs[0]\n assert doc.sections[0].text is not None\n assert EXPECTED_QUOTE in doc.sections[0].text\n\n\n@pytest.mark.parametrize(\"quotes_to_scroll_web_connector\", [False], indirect=True)\ndef test_web_connector_no_scroll(quotes_to_scroll_web_connector: WebConnector) -> None:\n all_docs: list[Document] = []\n document_batches = quotes_to_scroll_web_connector.load_from_state()\n for doc_batch in document_batches:\n for doc in doc_batch:\n if isinstance(doc, HierarchyNode):\n continue\n all_docs.append(doc)\n\n assert len(all_docs) == 1\n doc = all_docs[0]\n assert doc.sections[0].text is not None\n assert EXPECTED_QUOTE not in doc.sections[0].text\n\n\nMERCURY_EXPECTED_QUOTE = \"How can we help?\"\n\n\n@pytest.mark.xfail(\n reason=(\n \"flaky. maybe we can improve how we avoid triggering bot protection ormaybe this is just how it has to be.\"\n ),\n)\ndef test_web_connector_bot_protection() -> None:\n connector = WebConnector(\n base_url=\"https://support.mercury.com/hc\",\n web_connector_type=WEB_CONNECTOR_VALID_SETTINGS.SINGLE.value,\n )\n document_batches = list(connector.load_from_state())\n assert len(document_batches) == 1\n doc_batch = document_batches[0]\n assert len(doc_batch) == 1\n doc = doc_batch[0]\n assert not isinstance(doc, HierarchyNode)\n assert doc.sections[0].text is not None\n assert MERCURY_EXPECTED_QUOTE in doc.sections[0].text\n\n\ndef test_web_connector_recursive_www_redirect() -> None:\n # Check that https://onyx.app can be recursed if re-directed to www.onyx.app\n # Run in thread pool to avoid conflict with pytest-asyncio's event loop\n def _run_connector() -> list[Document]:\n connector = WebConnector(\n base_url=\"https://onyx.app\",\n web_connector_type=WEB_CONNECTOR_VALID_SETTINGS.RECURSIVE.value,\n )\n return [\n doc\n for batch in connector.load_from_state()\n for doc in batch\n if not isinstance(doc, HierarchyNode)\n ]\n\n with ThreadPoolExecutor(max_workers=1) as executor:\n future = executor.submit(_run_connector)\n documents = future.result()\n\n assert len(documents) > 1\n" }, { "path": "backend/tests/daily/connectors/zendesk/test_zendesk_connector.py", "content": "import json\nimport os\nimport time\nfrom pathlib import Path\nfrom typing import cast\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import HierarchyNode\nfrom onyx.connectors.zendesk.connector import ZendeskConnector\nfrom tests.daily.connectors.utils import load_all_from_connector\n\n\ndef load_test_data(file_name: str = \"test_zendesk_data.json\") -> dict[str, dict]:\n current_dir = Path(__file__).parent\n with open(current_dir / file_name, \"r\") as f:\n return json.load(f)\n\n\n@pytest.fixture\ndef zendesk_article_connector() -> ZendeskConnector:\n connector = ZendeskConnector(content_type=\"articles\")\n connector.load_credentials(get_credentials())\n return connector\n\n\n@pytest.fixture\ndef zendesk_ticket_connector() -> ZendeskConnector:\n connector = ZendeskConnector(content_type=\"tickets\")\n connector.load_credentials(get_credentials())\n return connector\n\n\ndef get_credentials() -> dict[str, str]:\n return {\n \"zendesk_subdomain\": os.environ[\"ZENDESK_SUBDOMAIN\"],\n \"zendesk_email\": os.environ[\"ZENDESK_EMAIL\"],\n \"zendesk_token\": os.environ[\"ZENDESK_TOKEN\"],\n }\n\n\n@pytest.mark.xfail(\n reason=(\n \"Cannot get Zendesk developer account to ensure zendesk account does not expire after 2 weeks\"\n )\n)\n@pytest.mark.parametrize(\n \"connector_fixture\", [\"zendesk_article_connector\", \"zendesk_ticket_connector\"]\n)\ndef test_zendesk_connector_basic(\n request: pytest.FixtureRequest, connector_fixture: str\n) -> None:\n connector = cast(ZendeskConnector, request.getfixturevalue(connector_fixture))\n test_data = load_test_data()\n all_docs: list[Document] = []\n target_test_doc_id: str\n if connector.content_type == \"articles\":\n target_test_doc_id = f\"article:{test_data['article']['id']}\"\n else:\n target_test_doc_id = f\"zendesk_ticket_{test_data['ticket']['id']}\"\n\n target_doc: Document | None = None\n\n for doc in load_all_from_connector(connector, 0, time.time()).documents:\n all_docs.append(doc)\n if doc.id == target_test_doc_id:\n target_doc = doc\n print(f\"target_doc {target_doc}\")\n\n assert len(all_docs) > 0, \"No documents were retrieved from the connector\"\n assert (\n target_doc is not None\n ), \"Target document was not found in the retrieved documents\"\n assert target_doc.source == DocumentSource.ZENDESK, \"Document source is not ZENDESK\"\n\n if connector.content_type == \"articles\":\n test_article = test_data[\"article\"]\n assert target_doc.semantic_identifier == test_article[\"semantic_identifier\"]\n assert target_doc.sections[0].link == test_article[\"sections\"][0][\"link\"]\n assert target_doc.source == test_article[\"source\"]\n assert target_doc.primary_owners is not None\n assert len(target_doc.primary_owners) == 1\n assert (\n target_doc.primary_owners[0].display_name\n == test_article[\"primary_owners\"][0][\"display_name\"]\n )\n assert (\n target_doc.primary_owners[0].email\n == test_article[\"primary_owners\"][0][\"email\"]\n )\n else:\n test_ticket = test_data[\"ticket\"]\n assert target_doc.semantic_identifier == test_ticket[\"semantic_identifier\"]\n assert target_doc.sections[0].link == test_ticket[\"sections\"][0][\"link\"]\n assert target_doc.source == test_ticket[\"source\"]\n assert target_doc.metadata[\"status\"] == test_ticket[\"metadata\"][\"status\"]\n assert target_doc.metadata[\"priority\"] == test_ticket[\"metadata\"][\"priority\"]\n assert target_doc.metadata[\"tags\"] == test_ticket[\"metadata\"][\"tags\"]\n assert (\n target_doc.metadata[\"ticket_type\"] == test_ticket[\"metadata\"][\"ticket_type\"]\n )\n\n\n@pytest.mark.xfail(\n reason=(\n \"Cannot get Zendesk developer account to ensure zendesk account does not expire after 2 weeks\"\n )\n)\ndef test_zendesk_connector_slim(zendesk_article_connector: ZendeskConnector) -> None:\n # Get full doc IDs\n all_full_doc_ids = set()\n for doc in load_all_from_connector(\n zendesk_article_connector, 0, time.time()\n ).documents:\n all_full_doc_ids.add(doc.id)\n\n # Get slim doc IDs\n all_slim_doc_ids = set()\n for slim_doc_batch in zendesk_article_connector.retrieve_all_slim_docs_perm_sync():\n all_slim_doc_ids.update(\n [doc.id for doc in slim_doc_batch if not isinstance(doc, HierarchyNode)]\n )\n\n # Full docs should be subset of slim docs\n assert all_full_doc_ids.issubset(\n all_slim_doc_ids\n ), f\"Full doc IDs {all_full_doc_ids} not subset of slim doc IDs {all_slim_doc_ids}\"\n" }, { "path": "backend/tests/daily/connectors/zendesk/test_zendesk_data.json", "content": "{\n \"article\": {\n \"id\": \"32502691728155\",\n \"sections\": [\n {\n \"link\": \"https://d3v-onyx.zendesk.com/hc/en-us/articles/32502691728155-How-can-agents-leverage-knowledge-to-help-customers\"\n }\n ],\n \"source\": \"zendesk\",\n \"semantic_identifier\": \"How can agents leverage knowledge to help customers?\",\n \"primary_owners\": [\n {\n \"display_name\": \"Dan Swer\",\n \"email\": \"admin@onyx-test.com\"\n }\n ]\n },\n \"ticket\": {\n \"id\": \"1\",\n \"sections\": [\n {\n \"link\": \"https://d3v-onyx.zendesk.com/agent/tickets/1\"\n }\n ],\n \"source\": \"zendesk\",\n \"semantic_identifier\": \"Ticket #1: SAMPLE TICKET: Meet the ticket\",\n \"metadata\": {\n \"status\": \"open\",\n \"priority\": \"normal\",\n \"tags\": [\"sample\", \"support\", \"zendesk\"],\n \"ticket_type\": \"incident\"\n }\n }\n}" }, { "path": "backend/tests/daily/embedding/test_embeddings.py", "content": "import os\n\nimport pytest\nfrom tenacity import retry\nfrom tenacity import retry_if_exception_type\nfrom tenacity import stop_after_attempt\nfrom tenacity import wait_exponential\n\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom shared_configs.enums import EmbedTextType\nfrom shared_configs.model_server_models import EmbeddingProvider\n\nVALID_SAMPLE = [\"hi\", \"hello my name is bob\", \"woah there!!!. 😃\"]\nVALID_LONG_SAMPLE = [\"hi \" * 999]\n# openai limit is 2048, cohere is supposed to be 96 but in practice that doesn't\n# seem to be true\nTOO_LONG_SAMPLE = [\"a\"] * 2500\n\n\ndef _run_embeddings(\n texts: list[str], embedding_model: EmbeddingModel, expected_dim: int\n) -> None:\n for text_type in [EmbedTextType.QUERY, EmbedTextType.PASSAGE]:\n embeddings = embedding_model.encode(texts, text_type)\n assert len(embeddings) == len(texts)\n assert len(embeddings[0]) == expected_dim\n\n\n@pytest.fixture\ndef openai_embedding_model() -> EmbeddingModel:\n return EmbeddingModel(\n server_host=\"localhost\",\n server_port=9000,\n model_name=\"text-embedding-3-small\",\n normalize=True,\n query_prefix=None,\n passage_prefix=None,\n api_key=os.environ[\"OPENAI_API_KEY\"],\n provider_type=EmbeddingProvider.OPENAI,\n api_url=None,\n )\n\n\ndef test_openai_embedding(openai_embedding_model: EmbeddingModel) -> None:\n _run_embeddings(VALID_SAMPLE, openai_embedding_model, 1536)\n _run_embeddings(TOO_LONG_SAMPLE, openai_embedding_model, 1536)\n\n\n@pytest.fixture\ndef cohere_embedding_model() -> EmbeddingModel:\n return EmbeddingModel(\n server_host=\"localhost\",\n server_port=9000,\n model_name=\"embed-english-light-v3.0\",\n normalize=True,\n query_prefix=None,\n passage_prefix=None,\n api_key=os.environ[\"COHERE_API_KEY\"],\n provider_type=EmbeddingProvider.COHERE,\n api_url=None,\n )\n\n\ndef test_cohere_embedding(cohere_embedding_model: EmbeddingModel) -> None:\n _run_embeddings(VALID_SAMPLE, cohere_embedding_model, 384)\n _run_embeddings(TOO_LONG_SAMPLE, cohere_embedding_model, 384)\n\n\n@pytest.fixture\ndef litellm_embedding_model() -> EmbeddingModel:\n return EmbeddingModel(\n server_host=\"localhost\",\n server_port=9000,\n model_name=\"text-embedding-3-small\",\n normalize=True,\n query_prefix=None,\n passage_prefix=None,\n api_key=os.environ[\"LITELLM_API_KEY\"],\n provider_type=EmbeddingProvider.LITELLM,\n api_url=os.environ[\"LITELLM_API_URL\"],\n )\n\n\n@pytest.mark.skip(reason=\"re-enable when we can get the correct litellm key and url\")\ndef test_litellm_embedding(litellm_embedding_model: EmbeddingModel) -> None:\n _run_embeddings(VALID_SAMPLE, litellm_embedding_model, 1536)\n _run_embeddings(TOO_LONG_SAMPLE, litellm_embedding_model, 1536)\n\n\n@pytest.fixture\ndef local_nomic_embedding_model() -> EmbeddingModel:\n return EmbeddingModel(\n server_host=\"localhost\",\n server_port=9000,\n model_name=\"nomic-ai/nomic-embed-text-v1\",\n normalize=True,\n query_prefix=\"search_query: \",\n passage_prefix=\"search_document: \",\n api_key=None,\n provider_type=None,\n api_url=None,\n )\n\n\ndef test_local_nomic_embedding(local_nomic_embedding_model: EmbeddingModel) -> None:\n _run_embeddings(VALID_SAMPLE, local_nomic_embedding_model, 768)\n _run_embeddings(TOO_LONG_SAMPLE, local_nomic_embedding_model, 768)\n\n\n@pytest.fixture\ndef azure_embedding_model() -> EmbeddingModel:\n return EmbeddingModel(\n server_host=\"localhost\",\n server_port=9000,\n model_name=\"text-embedding-3-small\",\n normalize=True,\n query_prefix=None,\n passage_prefix=None,\n api_key=os.environ[\"AZURE_API_KEY\"],\n provider_type=EmbeddingProvider.AZURE,\n api_url=os.environ[\"AZURE_API_URL\"],\n )\n\n\n# Azure has strict rate limits on their embedding API, so we retry with exponential\n# backoff to handle transient RateLimitError responses\n@retry(\n retry=retry_if_exception_type(RuntimeError),\n stop=stop_after_attempt(5),\n wait=wait_exponential(multiplier=1, min=1, max=10),\n reraise=True,\n)\ndef test_azure_embedding(azure_embedding_model: EmbeddingModel) -> None:\n _run_embeddings(VALID_SAMPLE, azure_embedding_model, 1536)\n _run_embeddings(TOO_LONG_SAMPLE, azure_embedding_model, 1536)\n\n\n# NOTE (chris): this test doesn't work, and I do not know why\n# def test_azure_embedding_model_rate_limit(azure_embedding_model: EmbeddingModel):\n# \"\"\"NOTE: this test relies on a very low rate limit for the Azure API +\n# this test only being run once in a 1 minute window\"\"\"\n# # VALID_LONG_SAMPLE is 999 tokens, so the second call should run into rate\n# # limits assuming the limit is 1000 tokens per minute\n# result = azure_embedding_model.encode(VALID_LONG_SAMPLE, EmbedTextType.QUERY)\n# assert len(result) == 1\n# assert len(result[0]) == 1536\n\n# # this should fail\n# with pytest.raises(ModelServerRateLimitError):\n# azure_embedding_model.encode(VALID_LONG_SAMPLE, EmbedTextType.QUERY)\n# azure_embedding_model.encode(VALID_LONG_SAMPLE, EmbedTextType.QUERY)\n# azure_embedding_model.encode(VALID_LONG_SAMPLE, EmbedTextType.QUERY)\n\n# # this should succeed, since passage requests retry up to 10 times\n# start = time.time()\n# result = azure_embedding_model.encode(VALID_LONG_SAMPLE, EmbedTextType.PASSAGE)\n# assert len(result) == 1\n# assert len(result[0]) == 1536\n# assert time.time() - start > 30 # make sure we waited, even though we hit rate limits\n" }, { "path": "backend/tests/daily/llm/test_bedrock.py", "content": "import os\nfrom typing import Any\n\nimport pytest\nfrom fastapi.testclient import TestClient\n\nfrom onyx.llm.constants import LlmProviderNames\n\n\n_DEFAULT_BEDROCK_MODEL = \"anthropic.claude-3-5-sonnet-20241022-v2:0\"\n\n\n@pytest.mark.xfail(\n reason=\"Credentials not yet available due to compliance work needed\",\n)\ndef test_bedrock_llm_configuration(client: TestClient) -> None:\n # Prepare the test request payload\n test_request: dict[str, Any] = {\n \"provider\": LlmProviderNames.BEDROCK,\n \"model\": _DEFAULT_BEDROCK_MODEL,\n \"api_key\": None,\n \"api_base\": None,\n \"api_version\": None,\n \"custom_config\": {\n \"AWS_REGION_NAME\": os.environ.get(\"AWS_REGION_NAME\", \"us-east-1\"),\n \"AWS_ACCESS_KEY_ID\": os.environ.get(\"AWS_ACCESS_KEY_ID\"),\n \"AWS_SECRET_ACCESS_KEY\": os.environ.get(\"AWS_SECRET_ACCESS_KEY\"),\n },\n \"model_configurations\": [{\"name\": _DEFAULT_BEDROCK_MODEL, \"is_visible\": True}],\n \"api_key_changed\": True,\n \"custom_config_changed\": True,\n }\n\n # Send the test request\n response = client.post(\"/admin/llm/test\", json=test_request)\n\n # Assert the response\n assert (\n response.status_code == 200\n ), f\"Expected status code 200, but got {response.status_code}. Response: {response.text}\"\n\n\ndef test_bedrock_llm_configuration_invalid_key(client: TestClient) -> None:\n # Prepare the test request payload with invalid credentials\n test_request: dict[str, Any] = {\n \"provider\": LlmProviderNames.BEDROCK,\n \"model\": _DEFAULT_BEDROCK_MODEL,\n \"api_key\": None,\n \"api_base\": None,\n \"api_version\": None,\n \"custom_config\": {\n \"AWS_REGION_NAME\": \"us-east-1\",\n \"AWS_ACCESS_KEY_ID\": \"invalid_access_key_id\",\n \"AWS_SECRET_ACCESS_KEY\": \"invalid_secret_access_key\",\n },\n \"model_configurations\": [{\"name\": _DEFAULT_BEDROCK_MODEL, \"is_visible\": True}],\n \"api_key_changed\": True,\n \"custom_config_changed\": True,\n }\n\n # Send the test request\n response = client.post(\"/admin/llm/test\", json=test_request)\n\n # Assert the response\n assert (\n response.status_code == 400\n ), f\"Expected status code 400, but got {response.status_code}. Response: {response.text}\"\n assert (\n \"Invalid credentials\" in response.text\n or \"Invalid Authentication\" in response.text\n ), f\"Expected error message about invalid credentials, but got: {response.text}\"\n" }, { "path": "backend/tests/external_dependency_unit/answer/conftest.py", "content": "import os\nfrom collections.abc import Iterator\nfrom collections.abc import Mapping\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\n\n\n# Counter for generating unique file IDs in mock file store\n_mock_file_id_counter = 0\n\n\ndef ensure_default_llm_provider(db_session: Session) -> None:\n \"\"\"Ensure a default LLM provider exists for tests that exercise chat flows.\"\"\"\n\n try:\n llm_provider_request = LLMProviderUpsertRequest(\n name=\"test-provider\",\n provider=LlmProviderNames.OPENAI,\n api_key=os.environ.get(\"OPENAI_API_KEY\", \"test\"),\n is_public=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\",\n is_visible=True,\n )\n ],\n groups=[],\n )\n provider = upsert_llm_provider(\n llm_provider_upsert_request=llm_provider_request,\n db_session=db_session,\n )\n update_default_provider(provider.id, \"gpt-4o-mini\", db_session)\n except Exception as exc: # pragma: no cover - only hits on duplicate setup issues\n # Rollback to clear the pending transaction state\n db_session.rollback()\n print(f\"Note: Could not create LLM provider: {exc}\")\n\n\n@pytest.fixture\ndef mock_nlp_embeddings_post() -> Iterator[None]:\n \"\"\"Patch model-server embedding HTTP calls used by NLP components.\"\"\"\n\n def _mock_post(\n url: str,\n json: Mapping[str, Any] | None = None,\n headers: Mapping[str, str] | None = None, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n ) -> MagicMock:\n resp = MagicMock()\n if \"encoder/bi-encoder-embed\" in url:\n num_texts = len(json.get(\"texts\", [])) if json else 1\n resp.status_code = 200\n resp.json.return_value = {\"embeddings\": [[0.0] * 768] * num_texts}\n resp.raise_for_status = MagicMock()\n return resp\n resp.status_code = 200\n resp.json.return_value = {}\n resp.raise_for_status = MagicMock()\n return resp\n\n with patch(\n \"onyx.natural_language_processing.search_nlp_models.requests.post\",\n side_effect=_mock_post,\n ):\n yield\n\n\n@pytest.fixture\ndef mock_gpu_status() -> Iterator[None]:\n \"\"\"Avoid hitting model server for GPU status checks.\"\"\"\n with patch(\n \"onyx.utils.gpu_utils._get_gpu_status_from_model_server\", return_value=False\n ):\n yield\n\n\n@pytest.fixture\ndef mock_vespa_query() -> Iterator[None]:\n \"\"\"Stub Vespa query to a safe empty response to avoid CI flakiness.\"\"\"\n with patch(\"onyx.document_index.vespa.index.query_vespa\", return_value=[]):\n yield\n\n\n@pytest.fixture\ndef mock_file_store() -> Iterator[None]:\n \"\"\"Mock the file store to avoid S3/storage dependencies in tests.\"\"\"\n global _mock_file_id_counter\n\n def _mock_save_file(*args: Any, **kwargs: Any) -> str: # noqa: ARG001\n global _mock_file_id_counter\n _mock_file_id_counter += 1\n # Return a predictable file ID for tests\n return \"123\"\n\n mock_store = MagicMock()\n mock_store.save_file.side_effect = _mock_save_file\n mock_store.initialize.return_value = None\n\n with patch(\n \"onyx.file_store.utils.get_default_file_store\",\n return_value=mock_store,\n ):\n yield\n\n\n@pytest.fixture\ndef mock_external_deps(\n mock_nlp_embeddings_post: None, # noqa: ARG001\n mock_gpu_status: None, # noqa: ARG001\n mock_vespa_query: None, # noqa: ARG001\n mock_file_store: None, # noqa: ARG001\n) -> Iterator[None]:\n \"\"\"Convenience fixture to enable all common external dependency mocks.\"\"\"\n yield\n" }, { "path": "backend/tests/external_dependency_unit/answer/stream_test_assertions.py", "content": "from __future__ import annotations\n\nfrom typing import cast\n\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.chat.models import CreateChatSessionID\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import ImageGenerationFinal\nfrom onyx.server.query_and_chat.streaming_models import OpenUrlDocuments\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta\n\n\ndef assert_answer_stream_part_correct(\n received: AnswerStreamPart, expected: AnswerStreamPart\n) -> None:\n assert isinstance(received, type(expected))\n\n if isinstance(received, Packet):\n r_packet = cast(Packet, received)\n e_packet = cast(Packet, expected)\n\n assert r_packet.placement == e_packet.placement\n\n if isinstance(r_packet.obj, SearchToolDocumentsDelta):\n assert isinstance(e_packet.obj, SearchToolDocumentsDelta)\n assert is_search_tool_document_delta_equal(r_packet.obj, e_packet.obj)\n return\n elif isinstance(r_packet.obj, OpenUrlDocuments):\n assert isinstance(e_packet.obj, OpenUrlDocuments)\n assert is_open_url_documents_equal(r_packet.obj, e_packet.obj)\n return\n elif isinstance(r_packet.obj, AgentResponseStart):\n assert isinstance(e_packet.obj, AgentResponseStart)\n assert is_agent_response_start_equal(r_packet.obj, e_packet.obj)\n return\n elif isinstance(r_packet.obj, ImageGenerationFinal):\n assert isinstance(e_packet.obj, ImageGenerationFinal)\n assert is_image_generation_final_equal(r_packet.obj, e_packet.obj)\n return\n\n assert r_packet.obj == e_packet.obj\n elif isinstance(received, MessageResponseIDInfo):\n # We're not going to make assumptions about what the user id / assistant id should be\n # So just return\n return\n elif isinstance(received, CreateChatSessionID):\n # Don't worry about same session ids\n return\n else:\n raise NotImplementedError(\"Not implemented\")\n\n\ndef _are_search_docs_equal(\n received: list[SearchDoc],\n expected: list[SearchDoc],\n) -> bool:\n \"\"\"\n What we care about:\n - All documents are present (order does not)\n - Expected document_id, link, blurb, source_type and hidden\n \"\"\"\n if len(received) != len(expected):\n return False\n\n received.sort(key=lambda x: x.document_id)\n expected.sort(key=lambda x: x.document_id)\n\n for received_document, expected_document in zip(received, expected):\n if received_document.document_id != expected_document.document_id:\n return False\n if received_document.link != expected_document.link:\n return False\n if received_document.blurb != expected_document.blurb:\n return False\n if received_document.source_type != expected_document.source_type:\n return False\n if received_document.hidden != expected_document.hidden:\n return False\n return True\n\n\ndef is_search_tool_document_delta_equal(\n received: SearchToolDocumentsDelta,\n expected: SearchToolDocumentsDelta,\n) -> bool:\n \"\"\"\n What we care about:\n - All documents are present (order does not)\n - Expected document_id, link, blurb, source_type and hidden\n \"\"\"\n received_documents = received.documents\n expected_documents = expected.documents\n\n return _are_search_docs_equal(received_documents, expected_documents)\n\n\ndef is_open_url_documents_equal(\n received: OpenUrlDocuments,\n expected: OpenUrlDocuments,\n) -> bool:\n \"\"\"\n What we care about:\n - All documents are present (order does not)\n - Expected document_id, link, blurb, source_type and hidden\n \"\"\"\n received_documents = received.documents\n expected_documents = expected.documents\n\n return _are_search_docs_equal(received_documents, expected_documents)\n\n\ndef is_agent_response_start_equal(\n received: AgentResponseStart,\n expected: AgentResponseStart,\n) -> bool:\n \"\"\"\n What we care about:\n - All documents are present (order does not)\n - Expected document_id, link, blurb, source_type and hidden\n \"\"\"\n received_documents = received.final_documents\n expected_documents = expected.final_documents\n\n if received_documents is None and expected_documents is None:\n return True\n if not received_documents or not expected_documents:\n return False\n\n return _are_search_docs_equal(received_documents, expected_documents)\n\n\ndef is_image_generation_final_equal(\n received: ImageGenerationFinal,\n expected: ImageGenerationFinal,\n) -> bool:\n \"\"\"\n What we care about:\n - Number of images are the same\n - On each image, url and file_id are aligned such that url=/api/chat/file/{file_id}\n - Revised prompt is expected\n - Shape is expected\n \"\"\"\n if len(received.images) != len(expected.images):\n return False\n\n for received_image, expected_image in zip(received.images, expected.images):\n if received_image.url != f\"/api/chat/file/{received_image.file_id}\":\n return False\n if received_image.revised_prompt != expected_image.revised_prompt:\n return False\n if received_image.shape != expected_image.shape:\n return False\n return True\n" }, { "path": "backend/tests/external_dependency_unit/answer/stream_test_builder.py", "content": "from __future__ import annotations\n\nfrom collections.abc import Iterator\n\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import OverallStop\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import ReasoningDone\nfrom onyx.server.query_and_chat.streaming_models import ReasoningStart\nfrom tests.external_dependency_unit.answer.stream_test_assertions import (\n assert_answer_stream_part_correct,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import (\n create_packet_with_agent_response_delta,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import (\n create_packet_with_reasoning_delta,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import create_placement\nfrom tests.external_dependency_unit.mock_llm import LLMResponse\nfrom tests.external_dependency_unit.mock_llm import MockLLMController\n\n\nclass StreamTestBuilder:\n def __init__(self, llm_controller: MockLLMController) -> None:\n self._llm_controller = llm_controller\n\n # List of (expected_packet, forward_count) tuples\n self._expected_packets_queue: list[tuple[Packet, int]] = []\n\n def add_response(self, response: LLMResponse) -> StreamTestBuilder:\n self._llm_controller.add_response(response)\n\n return self\n\n def add_responses_together(self, *responses: LLMResponse) -> StreamTestBuilder:\n \"\"\"Add multiple responses that should be emitted together in the same tick.\"\"\"\n self._llm_controller.add_responses_together(*responses)\n\n return self\n\n def expect(\n self, expected_pkt: Packet, forward: int | bool = True\n ) -> StreamTestBuilder:\n \"\"\"\n Add an expected packet to the queue.\n\n Args:\n expected_pkt: The packet to expect\n forward: Number of tokens to forward before expecting this packet.\n True = 1 token, False = 0 tokens, int = that many tokens.\n \"\"\"\n forward_count = 1 if forward is True else (0 if forward is False else forward)\n self._expected_packets_queue.append((expected_pkt, forward_count))\n\n return self\n\n def expect_packets(\n self, packets: list[Packet], forward: int | bool = True\n ) -> StreamTestBuilder:\n \"\"\"\n Add multiple expected packets to the queue.\n\n Args:\n packets: List of packets to expect\n forward: Number of tokens to forward before expecting EACH packet.\n True = 1 token per packet, False = 0 tokens, int = that many tokens per packet.\n \"\"\"\n forward_count = 1 if forward is True else (0 if forward is False else forward)\n for pkt in packets:\n self._expected_packets_queue.append((pkt, forward_count))\n\n return self\n\n def expect_reasoning(\n self,\n reasoning_tokens: list[str],\n turn_index: int,\n ) -> StreamTestBuilder:\n return (\n self.expect(\n Packet(\n placement=create_placement(turn_index),\n obj=ReasoningStart(),\n )\n )\n .expect_packets(\n [\n create_packet_with_reasoning_delta(token, turn_index)\n for token in reasoning_tokens\n ]\n )\n .expect(\n Packet(\n placement=create_placement(turn_index),\n obj=ReasoningDone(),\n )\n )\n )\n\n def expect_agent_response(\n self,\n answer_tokens: list[str],\n turn_index: int,\n final_documents: list[SearchDoc] | None = None,\n ) -> StreamTestBuilder:\n return (\n self.expect(\n Packet(\n placement=create_placement(turn_index),\n obj=AgentResponseStart(\n final_documents=final_documents,\n ),\n )\n )\n .expect_packets(\n [\n create_packet_with_agent_response_delta(token, turn_index)\n for token in answer_tokens\n ]\n )\n .expect(\n Packet(\n placement=create_placement(turn_index),\n obj=OverallStop(),\n )\n )\n )\n\n def run_and_validate(self, stream: Iterator[AnswerStreamPart]) -> None:\n while self._expected_packets_queue:\n expected_pkt, forward_count = self._expected_packets_queue.pop(0)\n if forward_count > 0:\n self._llm_controller.forward(forward_count)\n received_pkt = next(stream)\n\n assert_answer_stream_part_correct(received_pkt, expected_pkt)\n" }, { "path": "backend/tests/external_dependency_unit/answer/stream_test_utils.py", "content": "from __future__ import annotations\n\nfrom collections.abc import Iterator\nfrom uuid import UUID\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.chat_utils import create_chat_session_from_request\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.models import ChatSession\nfrom onyx.db.models import User\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import ReasoningDelta\nfrom tests.external_dependency_unit.mock_content_provider import MockWebContent\nfrom tests.external_dependency_unit.mock_search_provider import MockWebSearchResult\n\n\ndef create_placement(\n turn_index: int,\n tab_index: int = 0,\n sub_turn_index: int | None = None,\n model_index: int | None = 0,\n) -> Placement:\n return Placement(\n turn_index=turn_index,\n tab_index=tab_index,\n sub_turn_index=sub_turn_index,\n model_index=model_index,\n )\n\n\ndef submit_query(\n query: str,\n chat_session_id: UUID | None,\n db_session: Session,\n user: User,\n llm_override: LLMOverride | None = None,\n) -> Iterator[AnswerStreamPart]:\n request = SendMessageRequest(\n message=query,\n chat_session_id=chat_session_id,\n stream=True,\n chat_session_info=(\n ChatSessionCreationRequest() if chat_session_id is None else None\n ),\n llm_override=llm_override,\n )\n\n return handle_stream_message_objects(\n new_msg_req=request,\n user=user,\n db_session=db_session,\n )\n\n\ndef create_chat_session(\n db_session: Session,\n user: User,\n) -> ChatSession:\n return create_chat_session_from_request(\n chat_session_request=ChatSessionCreationRequest(),\n user_id=user.id,\n db_session=db_session,\n )\n\n\ndef create_packet_with_agent_response_delta(token: str, turn_index: int) -> Packet:\n return Packet(\n placement=create_placement(turn_index),\n obj=AgentResponseDelta(\n content=token,\n ),\n )\n\n\ndef create_packet_with_reasoning_delta(token: str, turn_index: int) -> Packet:\n return Packet(\n placement=create_placement(turn_index),\n obj=ReasoningDelta(\n reasoning=token,\n ),\n )\n\n\ndef create_web_search_doc(\n semantic_identifier: str,\n link: str,\n blurb: str,\n) -> SearchDoc:\n return SearchDoc(\n document_id=f\"WEB_SEARCH_DOC_{link}\",\n chunk_ind=0,\n semantic_identifier=semantic_identifier,\n link=link,\n blurb=blurb,\n source_type=DocumentSource.WEB,\n boost=1,\n hidden=False,\n metadata={},\n match_highlights=[],\n )\n\n\ndef mock_web_search_result_to_search_doc(result: MockWebSearchResult) -> SearchDoc:\n return create_web_search_doc(\n semantic_identifier=result.title,\n link=result.link,\n blurb=result.snippet,\n )\n\n\ndef mock_web_content_to_search_doc(content: MockWebContent) -> SearchDoc:\n return create_web_search_doc(\n semantic_identifier=content.title,\n link=content.url,\n blurb=content.title,\n )\n\n\ndef tokenise(text: str) -> list[str]:\n return [(token + \" \") for token in text.split(\" \")]\n" }, { "path": "backend/tests/external_dependency_unit/answer/test_answer_without_openai.py", "content": "from __future__ import annotations\n\nimport os\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.chat.models import StreamingError\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.enums import LLMModelFlowType\nfrom onyx.db.llm import fetch_existing_llm_providers\nfrom onyx.db.llm import remove_llm_provider\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef test_answer_with_only_anthropic_provider(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n \"\"\"Ensure chat still streams answers when only an Anthropic provider is configured.\"\"\"\n\n anthropic_api_key = os.environ.get(\"ANTHROPIC_API_KEY\")\n assert anthropic_api_key, \"ANTHROPIC_API_KEY environment variable must be set\"\n\n # Drop any existing providers so that only Anthropic is available.\n for provider in fetch_existing_llm_providers(db_session, [LLMModelFlowType.CHAT]):\n remove_llm_provider(db_session, provider.id)\n\n anthropic_model = \"claude-haiku-4-5-20251001\"\n provider_name = f\"anthropic-test-{uuid4().hex}\"\n\n anthropic_provider = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.ANTHROPIC,\n api_key=anthropic_api_key,\n is_public=True,\n groups=[],\n model_configurations=[\n ModelConfigurationUpsertRequest(name=anthropic_model, is_visible=True)\n ],\n api_key_changed=True,\n ),\n db_session=db_session,\n )\n\n try:\n update_default_provider(anthropic_provider.id, anthropic_model, db_session)\n\n test_user = create_test_user(db_session, email_prefix=\"anthropic_only\")\n chat_session = create_chat_session(\n db_session=db_session,\n description=\"Anthropic only chat\",\n user_id=test_user.id,\n persona_id=0,\n )\n\n chat_request = SendMessageRequest(\n message=\"hello\",\n chat_session_id=chat_session.id,\n )\n\n response_stream: list[AnswerStreamPart] = []\n for packet in handle_stream_message_objects(\n new_msg_req=chat_request,\n user=test_user,\n db_session=db_session,\n ):\n response_stream.append(packet)\n\n assert response_stream, \"Should receive streamed packets\"\n assert not any(\n isinstance(packet, StreamingError) for packet in response_stream\n ), \"No streaming errors expected with Anthropic provider\"\n\n has_message_id = any(\n isinstance(packet, MessageResponseIDInfo) for packet in response_stream\n )\n assert has_message_id, \"Should include reserved assistant message ID\"\n\n has_message_start = any(\n isinstance(packet, Packet) and isinstance(packet.obj, AgentResponseStart)\n for packet in response_stream\n )\n assert has_message_start, \"Stream should have a MessageStart packet\"\n\n has_message_delta = any(\n isinstance(packet, Packet) and isinstance(packet.obj, AgentResponseDelta)\n for packet in response_stream\n )\n assert has_message_delta, \"Stream should have a MessageDelta packet\"\n\n finally:\n remove_llm_provider(db_session, anthropic_provider.id)\n" }, { "path": "backend/tests/external_dependency_unit/answer/test_current_datetime_replacement.py", "content": "import re\nfrom datetime import datetime\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.chat.models import StreamingError\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.models import User\nfrom onyx.db.persona import get_persona_by_id\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef test_stream_chat_current_date_response(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n \"\"\"Smoke test that asking for current date yields a streamed response.\n\n This exercises the full chat path using the default persona, ensuring\n the system prompt makes it to the LLM and a response is returned.\n \"\"\"\n # Ensure LLM provider exists\n ensure_default_llm_provider(db_session)\n\n # Create user, persona, session\n test_user: User = create_test_user(db_session, email_prefix=\"test_current_date\")\n default_persona = get_persona_by_id(\n persona_id=0, user=test_user, db_session=db_session, is_for_edit=False\n )\n chat_session = create_chat_session(\n db_session=db_session,\n description=\"Test current date question\",\n user_id=test_user.id if test_user else None,\n persona_id=default_persona.id,\n )\n\n chat_request = SendMessageRequest(\n message=\"Please respond only with the current date in the format 'Weekday Month DD, YYYY'.\",\n chat_session_id=chat_session.id,\n )\n\n gen = handle_stream_message_objects(\n new_msg_req=chat_request,\n user=test_user,\n db_session=db_session,\n )\n\n raw: list[AnswerStreamPart] = []\n content = \"\"\n had_error = False\n\n for pkt in gen:\n raw.append(pkt)\n if hasattr(pkt, \"obj\") and isinstance(pkt.obj, AgentResponseDelta):\n if pkt.obj.content:\n content += pkt.obj.content\n if hasattr(pkt, \"obj\") and isinstance(pkt.obj, StreamingError):\n had_error = True\n break\n\n assert not had_error, \"Should not error when answering current date\"\n assert any(\n isinstance(p, MessageResponseIDInfo) for p in raw\n ), \"Should yield a message ID\"\n assert len(content) > 0, \"Should stream some assistant content\"\n\n # Validate the response contains a properly formatted current date string\n match = re.search(r\"[A-Za-z]+ [A-Za-z]+ \\d{1,2}, \\d{4}\", content)\n assert match, f\"Expected a date in content, got: {content[:200]}...\"\n\n timestamp_str = match.group(0)\n timestamp_dt = datetime.strptime(timestamp_str, \"%A %B %d, %Y\")\n now = datetime.now()\n\n assert timestamp_dt.strftime(\"%A\") == now.strftime(\n \"%A\"\n ), f\"Expected weekday {now.strftime('%A')}, got {timestamp_dt.strftime('%A')}\"\n assert timestamp_dt.strftime(\"%B\") == now.strftime(\n \"%B\"\n ), f\"Expected month {now.strftime('%B')}, got {timestamp_dt.strftime('%B')}\"\n assert timestamp_dt.day == now.day and timestamp_dt.year == now.year, (\n f\"Expected day {now.strftime('%d')} and year {now.strftime('%Y')}, \"\n f\"got {timestamp_dt.strftime('%d')} {timestamp_dt.strftime('%Y')}\"\n )\n" }, { "path": "backend/tests/external_dependency_unit/answer/test_stream_chat_message.py", "content": "from __future__ import annotations\n\nimport json\nfrom uuid import UUID\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.models import CreateChatSessionID\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseStart\nfrom onyx.server.query_and_chat.streaming_models import GeneratedImage\nfrom onyx.server.query_and_chat.streaming_models import ImageGenerationFinal\nfrom onyx.server.query_and_chat.streaming_models import ImageGenerationToolHeartbeat\nfrom onyx.server.query_and_chat.streaming_models import ImageGenerationToolStart\nfrom onyx.server.query_and_chat.streaming_models import OpenUrlDocuments\nfrom onyx.server.query_and_chat.streaming_models import OpenUrlStart\nfrom onyx.server.query_and_chat.streaming_models import OpenUrlUrls\nfrom onyx.server.query_and_chat.streaming_models import OverallStop\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import ReasoningDone\nfrom onyx.server.query_and_chat.streaming_models import ReasoningStart\nfrom onyx.server.query_and_chat.streaming_models import SearchToolDocumentsDelta\nfrom onyx.server.query_and_chat.streaming_models import SearchToolQueriesDelta\nfrom onyx.server.query_and_chat.streaming_models import SearchToolStart\nfrom onyx.server.query_and_chat.streaming_models import SectionEnd\nfrom onyx.server.query_and_chat.streaming_models import TopLevelBranching\nfrom tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider\nfrom tests.external_dependency_unit.answer.stream_test_assertions import (\n assert_answer_stream_part_correct,\n)\nfrom tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder\nfrom tests.external_dependency_unit.answer.stream_test_utils import create_chat_session\nfrom tests.external_dependency_unit.answer.stream_test_utils import (\n create_packet_with_agent_response_delta,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import (\n create_packet_with_reasoning_delta,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import create_placement\nfrom tests.external_dependency_unit.answer.stream_test_utils import (\n mock_web_content_to_search_doc,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import (\n mock_web_search_result_to_search_doc,\n)\nfrom tests.external_dependency_unit.answer.stream_test_utils import submit_query\nfrom tests.external_dependency_unit.answer.stream_test_utils import tokenise\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.mock_content_provider import MockWebContent\nfrom tests.external_dependency_unit.mock_content_provider import (\n use_mock_content_provider,\n)\nfrom tests.external_dependency_unit.mock_image_provider import (\n use_mock_image_generation_provider,\n)\nfrom tests.external_dependency_unit.mock_llm import LLMAnswerResponse\nfrom tests.external_dependency_unit.mock_llm import LLMReasoningResponse\nfrom tests.external_dependency_unit.mock_llm import LLMToolCallResponse\nfrom tests.external_dependency_unit.mock_llm import use_mock_llm\nfrom tests.external_dependency_unit.mock_search_pipeline import MockInternalSearchResult\nfrom tests.external_dependency_unit.mock_search_pipeline import use_mock_search_pipeline\nfrom tests.external_dependency_unit.mock_search_provider import MockWebSearchResult\nfrom tests.external_dependency_unit.mock_search_provider import use_mock_web_provider\n\n\ndef test_stream_chat_with_answer(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n \"\"\"Test that the stream chat with answer endpoint returns a valid answer.\"\"\"\n ensure_default_llm_provider(db_session)\n test_user = create_test_user(\n db_session, email_prefix=\"test_stream_chat_with_answer\"\n )\n\n query = \"What is the capital of France?\"\n answer = \"The capital of France is Paris.\"\n\n answer_tokens = tokenise(answer)\n\n with use_mock_llm() as mock_llm:\n handler = StreamTestBuilder(llm_controller=mock_llm)\n\n handler.add_response(LLMAnswerResponse(answer_tokens=answer_tokens))\n\n chat_session = create_chat_session(db_session=db_session, user=test_user)\n\n answer_stream = submit_query(\n query=query,\n chat_session_id=chat_session.id,\n db_session=db_session,\n user=test_user,\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=1,\n reserved_assistant_message_id=1,\n ),\n )\n\n handler.expect_agent_response(\n answer_tokens=answer_tokens,\n turn_index=0,\n ).run_and_validate(stream=answer_stream)\n\n with pytest.raises(StopIteration):\n next(answer_stream)\n\n\ndef test_stream_chat_with_answer_create_chat(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n ensure_default_llm_provider(db_session)\n test_user = create_test_user(\n db_session, email_prefix=\"test_stream_chat_with_answer_create_chat\"\n )\n\n query = \"Hi there friends\"\n answer = \"Hello friend\"\n\n tokens = [answer]\n\n with use_mock_llm() as mock_llm:\n handler = StreamTestBuilder(llm_controller=mock_llm)\n\n handler.add_response(LLMAnswerResponse(answer_tokens=tokens))\n\n answer_stream = submit_query(\n query=query,\n chat_session_id=None,\n db_session=db_session,\n user=test_user,\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=CreateChatSessionID(\n chat_session_id=UUID(\"123e4567-e89b-12d3-a456-426614174000\")\n ),\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=1,\n reserved_assistant_message_id=2,\n ),\n )\n\n handler.expect_agent_response(\n answer_tokens=tokens,\n turn_index=0,\n ).run_and_validate(stream=answer_stream)\n\n with pytest.raises(StopIteration):\n next(answer_stream)\n\n\ndef test_stream_chat_with_search_and_openurl_tools(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n ensure_default_llm_provider(db_session)\n test_user = create_test_user(\n db_session, email_prefix=\"test_stream_chat_with_search_tool\"\n )\n\n QUERY = \"What is the weather in Sydney?\"\n\n REASONING_RESPONSE_1 = (\n \"I need to perform a web search to get current weather details. \"\n \"I can use the search tool to do this.\"\n )\n\n WEB_QUERY_1 = \"weather in sydney\"\n WEB_QUERY_2 = \"current weather in sydney\"\n\n RESULTS1 = [\n MockWebSearchResult(\n title=\"Official Weather\",\n link=\"www.weather.com.au\",\n snippet=\"The current weather in Sydney is 20 degrees Celsius.\",\n ),\n MockWebSearchResult(\n title=\"Weather CHannel\",\n link=\"www.wc.com.au\",\n snippet=\"Morning is 10 degree Celsius, afternoon is 25 degrees Celsius.\",\n ),\n ]\n\n RESULTS2 = [\n MockWebSearchResult(\n title=\"Weather Now!\",\n link=\"www.weathernow.com.au\",\n snippet=\"The weather right now is sunny with a temperature of 22 degrees Celsius.\",\n )\n ]\n\n REASONING_RESPONSE_2 = \"I like weathernow and the official weather site\"\n\n QUERY_URLS_1 = [\"www.weathernow.com.au\", \"www.weather.com.au\"]\n\n CONTENT1 = [\n MockWebContent(\n title=\"Weather Now!\",\n url=\"www.weathernow.com.au\",\n content=\"The weather right now is sunny with a temperature of 22 degrees Celsius.\",\n ),\n MockWebContent(\n title=\"Weather Official\",\n url=\"www.weather.com.au\",\n content=\"The current weather in Sydney is 20 degrees Celsius.\",\n ),\n ]\n\n REASONING_RESPONSE_3 = (\n \"I now know everything that I need to know. \" \"I can now answer the question.\"\n )\n\n ANSWER_RESPONSE_1 = (\n \"The weather in Sydney is sunny with a temperature of 22 degrees celsius.\"\n )\n\n with (\n use_mock_llm() as mock_llm,\n use_mock_web_provider(db_session) as mock_web,\n use_mock_content_provider() as mock_content,\n ):\n handler = StreamTestBuilder(\n llm_controller=mock_llm,\n )\n\n chat_session = create_chat_session(db_session=db_session, user=test_user)\n\n answer_stream = submit_query(\n query=QUERY,\n chat_session_id=chat_session.id,\n db_session=db_session,\n user=test_user,\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=1,\n reserved_assistant_message_id=1,\n ),\n )\n\n # LLM Stream Response 1\n mock_web.add_results(WEB_QUERY_1, RESULTS1)\n mock_web.add_results(WEB_QUERY_2, RESULTS2)\n\n handler.add_response(\n LLMReasoningResponse(reasoning_tokens=tokenise(REASONING_RESPONSE_1))\n ).add_response(\n LLMToolCallResponse(\n tool_name=\"web_search\",\n tool_call_id=\"123\",\n tool_call_argument_tokens=[\n json.dumps({\"queries\": [WEB_QUERY_1, WEB_QUERY_2]})\n ],\n )\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=ReasoningStart(),\n )\n ).expect_packets(\n [\n create_packet_with_reasoning_delta(token, 0)\n for token in tokenise(REASONING_RESPONSE_1)\n ]\n ).expect(\n Packet(placement=create_placement(0), obj=ReasoningDone())\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=SearchToolStart(\n is_internet_search=True,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=SearchToolQueriesDelta(\n queries=[WEB_QUERY_1, WEB_QUERY_2],\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=SearchToolDocumentsDelta(\n documents=[\n mock_web_search_result_to_search_doc(result)\n for result in RESULTS1\n ]\n + [\n mock_web_search_result_to_search_doc(result)\n for result in RESULTS2\n ]\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=SectionEnd(),\n )\n ).run_and_validate(\n stream=answer_stream\n )\n\n # LLM Stream Response 2\n for content in CONTENT1:\n mock_content.add_content(content)\n\n handler.add_response(\n LLMReasoningResponse(reasoning_tokens=tokenise(REASONING_RESPONSE_2))\n ).add_response(\n LLMToolCallResponse(\n tool_name=\"open_url\",\n tool_call_id=\"123\",\n tool_call_argument_tokens=[json.dumps({\"urls\": QUERY_URLS_1})],\n )\n ).expect(\n Packet(\n placement=create_placement(2),\n obj=ReasoningStart(),\n )\n ).expect_packets(\n [\n create_packet_with_reasoning_delta(token, 2)\n for token in tokenise(REASONING_RESPONSE_2)\n ]\n ).expect(\n Packet(\n placement=create_placement(2),\n obj=ReasoningDone(),\n )\n ).expect(\n Packet(\n placement=create_placement(3),\n obj=OpenUrlStart(),\n )\n ).expect(\n Packet(\n placement=create_placement(3),\n obj=OpenUrlUrls(urls=[content.url for content in CONTENT1]),\n )\n ).expect(\n Packet(\n placement=create_placement(3),\n obj=OpenUrlDocuments(\n documents=[\n mock_web_content_to_search_doc(content) for content in CONTENT1\n ]\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(3),\n obj=SectionEnd(),\n )\n ).run_and_validate(\n stream=answer_stream\n )\n\n # LLM Stream Response 3\n handler.add_response(\n LLMReasoningResponse(reasoning_tokens=tokenise(REASONING_RESPONSE_3))\n ).add_response(\n LLMAnswerResponse(answer_tokens=tokenise(ANSWER_RESPONSE_1))\n ).expect(\n Packet(\n placement=create_placement(4),\n obj=ReasoningStart(),\n )\n ).expect_packets(\n [\n create_packet_with_reasoning_delta(token, 4)\n for token in tokenise(REASONING_RESPONSE_3)\n ]\n ).expect(\n Packet(\n placement=create_placement(4),\n obj=ReasoningDone(),\n )\n ).expect_agent_response(\n answer_tokens=tokenise(ANSWER_RESPONSE_1),\n turn_index=5,\n final_documents=[\n mock_web_search_result_to_search_doc(result) for result in RESULTS1\n ]\n + [mock_web_search_result_to_search_doc(result) for result in RESULTS2]\n + [mock_web_content_to_search_doc(content) for content in CONTENT1],\n ).run_and_validate(\n stream=answer_stream\n )\n\n with pytest.raises(StopIteration):\n next(answer_stream)\n\n\ndef test_image_generation_tool_no_reasoning(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n ensure_default_llm_provider(db_session)\n test_user = create_test_user(db_session, email_prefix=\"test_image_generation_tool\")\n\n QUERY = \"Create me an image of a dog on a rocketship\"\n\n IMAGE_DATA = (\n \"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfF\"\n \"cSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==\"\n )\n # Heartbeat interval is 5 seconds. A delay of 8 seconds ensures exactly 2 heartbeats:\n IMAGE_DELAY = 8.0\n\n ANSWER_RESPONSE = \"Here is a dog on a rocketship\"\n\n with (\n use_mock_llm() as mock_llm,\n use_mock_image_generation_provider() as mock_image_gen,\n ):\n handler = StreamTestBuilder(\n llm_controller=mock_llm,\n )\n\n chat_session = create_chat_session(db_session=db_session, user=test_user)\n\n answer_stream = submit_query(\n query=QUERY,\n chat_session_id=chat_session.id,\n db_session=db_session,\n user=test_user,\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=1,\n reserved_assistant_message_id=1,\n ),\n )\n\n # LLM Stream Response 1\n mock_image_gen.add_image(IMAGE_DATA, IMAGE_DELAY)\n mock_llm.set_max_timeout(\n IMAGE_DELAY + 5.0\n ) # Give enough buffer for image generation\n\n # The LLMToolCallResponse has 2 tokens (1 for tool name/id + 1 for arguments).\n # We need to forward all 2 tokens before the tool starts executing and emitting packets.\n # The tool then emits: start, heartbeats (during image generation), final, and section end.\n handler.add_response(\n LLMToolCallResponse(\n tool_name=\"generate_image\",\n tool_call_id=\"123\",\n tool_call_argument_tokens=[json.dumps({\"prompt\": QUERY})],\n )\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=ImageGenerationToolStart(),\n ),\n forward=2, # Forward both tool call tokens before expecting first packet\n ).expect_packets(\n [\n Packet(\n placement=create_placement(0),\n obj=ImageGenerationToolHeartbeat(),\n )\n ]\n * 2,\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=ImageGenerationFinal(\n images=[\n GeneratedImage(\n file_id=\"123\",\n url=\"/api/chat/file/123\",\n revised_prompt=QUERY,\n shape=\"square\",\n )\n ]\n ),\n ),\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=SectionEnd(),\n ),\n forward=False,\n ).run_and_validate(\n stream=answer_stream\n )\n\n # LLM Stream Response 2 - the answer comes after the tool call, so turn_index=1\n handler.add_response(\n LLMAnswerResponse(\n answer_tokens=tokenise(ANSWER_RESPONSE),\n )\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=AgentResponseStart(final_documents=None),\n )\n ).expect_packets(\n [\n create_packet_with_agent_response_delta(token, 1)\n for token in tokenise(ANSWER_RESPONSE)\n ]\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=OverallStop(),\n )\n ).run_and_validate(\n stream=answer_stream\n )\n\n with pytest.raises(StopIteration):\n next(answer_stream)\n\n\ndef test_parallel_internal_and_web_search_tool_calls(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n \"\"\"\n User asks a question\n LLM does some thinking\n LLM runs parallel tool calls for internal & web search\n\n -> Interal Search Branch performs seach + read ~10 documents\n -> Web Search: Searches the web for information\n\n LLM reads web documents\n LLM does thinking across all results\n LLM reads one more website\n LLM does more thinking\n LLM generates answer\n \"\"\"\n ensure_default_llm_provider(db_session)\n test_user = create_test_user(\n db_session, email_prefix=\"test_parallel_internal_and_web_search_tool_calls\"\n )\n\n AVALIABLE_CONNECTORS = [\n DocumentSource.GOOGLE_DRIVE,\n DocumentSource.CONFLUENCE,\n DocumentSource.LINEAR,\n DocumentSource.FIREFLIES,\n ]\n\n QUERY = \"How will forecasts against 2026 global GDP growth affect our Q2 strategy?\"\n\n THINKING_RESPONSE_1 = (\n \"I need to build more context around the user's query to answer it. \"\n \"I should look at GDP growth projections for 2026. \"\n \"I should also look at what the Q2 strategy is and what projects are included. \"\n \"I should perform both web and internal searches in parallel to get information efficiently.\"\n )\n\n WEB_QUERIES_1 = [\n \"2026 global GDP growth projections\",\n \"GDP growth 2026\",\n \"GDP forecast 2026\",\n ]\n\n WEB_RESULTS_1 = {\n WEB_QUERIES_1[0]: [\n MockWebSearchResult(\n title=\"World Economic Outlook Update, January 2026\",\n link=\"https://www.imf.org/weo/issues/2026/01/19/world-economic-outlook-update-january-2026\",\n snippet=\"Global growth is projected at 3.3 percent for 2026 and 3.2 percent for 2027...\",\n ),\n MockWebSearchResult(\n title=\"IMF sees steady global growth in 2026 as AI boom offsets ...\",\n link=\"https://www.reuters.com/article/us-world-economy-imf-idUSKBN2JU23E\",\n snippet=\"IMF forecasts 2026 global GDP growth at 3.3% even with stronger 2025 performance\",\n ),\n MockWebSearchResult(\n title=\"The Global Economy Is Forecast to Post...\",\n link=\"https://www.goldmansachs.com/insights/articles/123\",\n snippet=\"Global GDP is projected by Goldman Sachs Research to increase 2.8% in 2026\",\n ),\n ],\n WEB_QUERIES_1[1]: [\n MockWebSearchResult(\n title=\"US third-quarter economic growth revised slightly higher\",\n link=\"https://www.reuters.com/word/us-third-quarter-eco\",\n snippet=\"Gross domestic product increased at an upwardly revised 4.4% annualized rate, the ...\",\n ),\n MockWebSearchResult(\n title=\"US GDP Growth Is Projected to Outperform Economist ...\",\n link=\"https://www.goldmansachs.com/insights/articles/321\",\n snippet=\"US GDP is forecast to expand 2.5% in 2026 (fourth quarter, yoy), versus\",\n ),\n MockWebSearchResult(\n title=\"Gross Domestic Product\",\n link=\"https://www.bea.gov/data/gdp/gross-domestic-product\",\n snippet=\"Real gross domestic product (GDP) increased at an annual rate of 4.4 percent in the third quarter\",\n ),\n ],\n WEB_QUERIES_1[2]: [\n MockWebSearchResult(\n title=\"World Economic Outlook Update, January 2026\",\n link=\"https://www.imf.org/web/issues/2026/01/19/world-economic-outlook-update-january-2026\",\n snippet=\"Global growth is projected at 3.3 percent for 2026 and 3.2 percent for 2027...\",\n ),\n MockWebSearchResult(\n title=\"US GDP Growth Is Projected to Outperform Economist ...\",\n link=\"https://www.goldmansachs.com/insights/articles/321\",\n snippet=\"US GDP is forecast to expand 2.5% in 2026 (fourth quarter, yoy), versus\",\n ),\n MockWebSearchResult(\n title=\"Our economic outlook for the United States - Vanguard\",\n link=\"https://corporate.vanguard.com/content/corp/vemo\",\n snippet=\"We expect strong capital investment to remain a principal strength in the year ahead\",\n ),\n ],\n }\n\n INTERNAL_QUERIES_1 = [\"Q2 strategy 2026\", \"GDP growth 2026 projects\", \"Q2 projects\"]\n\n INTERNAL_RESULTS_1 = {\n INTERNAL_QUERIES_1[0]: [\n MockInternalSearchResult(\n document_id=\"123456789\",\n source_type=DocumentSource.GOOGLE_DRIVE,\n semantic_identifier=\"Q2 strategy 2026\",\n chunk_ind=11,\n ),\n MockInternalSearchResult(\n document_id=\"732190732173\",\n source_type=DocumentSource.FIREFLIES,\n semantic_identifier=\"What we think is going to happen in Q2\",\n chunk_ind=5,\n ),\n MockInternalSearchResult(\n document_id=\"12389123219\",\n source_type=DocumentSource.CONFLUENCE,\n semantic_identifier=\"Strategy roadmap for Q2 2026\",\n chunk_ind=7,\n ),\n ],\n INTERNAL_QUERIES_1[1]: [\n MockInternalSearchResult(\n document_id=\"123123\",\n source_type=DocumentSource.LINEAR,\n semantic_identifier=\"GDP growth 2026 projects\",\n chunk_ind=13,\n )\n ],\n INTERNAL_QUERIES_1[2]: [\n MockInternalSearchResult(\n document_id=\"98823643243\",\n source_type=DocumentSource.GOOGLE_DRIVE,\n semantic_identifier=\"Full list of Q2 projects\",\n chunk_ind=1,\n )\n ],\n }\n\n OPEN_URL_URLS_1 = [\n WEB_RESULTS_1[WEB_QUERIES_1[0]][0].link,\n WEB_RESULTS_1[WEB_QUERIES_1[0]][2].link,\n WEB_RESULTS_1[WEB_QUERIES_1[2]][0].link,\n ]\n\n OPEN_URL_DOCUMENTS_1 = [\n MockWebContent(\n title=WEB_RESULTS_1[WEB_QUERIES_1[0]][0].title,\n url=WEB_RESULTS_1[WEB_QUERIES_1[0]][0].link,\n content=\"Global growth is projected at 3.3 percent for 2026 and 3.2 percent for 2027...\",\n ),\n MockWebContent(\n title=WEB_RESULTS_1[WEB_QUERIES_1[0]][2].title,\n url=WEB_RESULTS_1[WEB_QUERIES_1[0]][2].link,\n content=\"Global growth is projected at 3.3 percent for 2026 and 3.2 percent for 2027...\",\n ),\n MockWebContent(\n title=WEB_RESULTS_1[WEB_QUERIES_1[2]][0].title,\n url=WEB_RESULTS_1[WEB_QUERIES_1[2]][0].link,\n content=\"Global growth is projected at 3.3 percent for 2026 and 3.2 percent for 2027...\",\n ),\n ]\n\n THINKING_RESPONSE_2 = (\n \"I now have a clear picture of the 2026 global GDP projections and the Q2 strategy. \"\n \"I would like to now about the outperform expections though...\"\n )\n\n OPEN_URL_URLS_2 = [WEB_RESULTS_1[WEB_QUERIES_1[1]][1].link]\n OPEN_URL_DOCUMENTS_2 = [\n MockWebContent(\n title=WEB_RESULTS_1[WEB_QUERIES_1[1]][1].title,\n url=WEB_RESULTS_1[WEB_QUERIES_1[1]][1].link,\n content=\"US GDP is forecast to expand 2.5% in 2026 (fourth quarter, yoy), versus\",\n )\n ]\n\n REASONING_RESPONSE_3 = (\n \"I now have all the information I need to answer the user's question.\"\n )\n\n ANSWER_RESPONSE = (\n \"We will have to change around some of our projects to accomodate the outperform expections. \"\n \"We should focus on aggresive expansion projects and prioritize them over cost-cutting initiatives.\"\n )\n\n expected_web_docs = []\n seen_web_results = set()\n for web_results in WEB_RESULTS_1.values():\n for web_result in web_results:\n key = (web_result.title, web_result.link)\n if key in seen_web_results:\n continue\n seen_web_results.add(key)\n expected_web_docs.append(mock_web_search_result_to_search_doc(web_result))\n\n expected_internal_docs = []\n seen_internal_results = set()\n for internal_results in INTERNAL_RESULTS_1.values():\n for internal_result in internal_results:\n key = (internal_result.semantic_identifier, internal_result.document_id)\n if key in seen_internal_results:\n continue\n seen_internal_results.add(key)\n expected_internal_docs.append(internal_result.to_search_doc())\n\n with (\n use_mock_llm() as mock_llm,\n use_mock_search_pipeline(\n connectors=AVALIABLE_CONNECTORS\n ) as mock_search_pipeline,\n use_mock_web_provider(db_session) as mock_web,\n use_mock_content_provider() as mock_content,\n ):\n for query, web_results in WEB_RESULTS_1.items():\n mock_web.add_results(query, web_results)\n\n for query, internal_results in INTERNAL_RESULTS_1.items():\n mock_search_pipeline.add_search_results(query, internal_results)\n\n handler = StreamTestBuilder(\n llm_controller=mock_llm,\n )\n\n chat_session = create_chat_session(db_session=db_session, user=test_user)\n\n answer_stream = submit_query(\n query=QUERY,\n chat_session_id=chat_session.id,\n db_session=db_session,\n user=test_user,\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=1,\n reserved_assistant_message_id=1,\n ),\n )\n\n # LLM Stream Response 1\n handler.add_response(\n LLMReasoningResponse(\n reasoning_tokens=tokenise(THINKING_RESPONSE_1),\n )\n ).add_responses_together(\n LLMToolCallResponse(\n tool_name=\"internal_search\",\n tool_call_id=\"123\",\n tool_call_argument_tokens=[json.dumps({\"queries\": INTERNAL_QUERIES_1})],\n ),\n LLMToolCallResponse(\n tool_name=\"web_search\",\n tool_call_id=\"321\",\n tool_call_argument_tokens=[json.dumps({\"queries\": WEB_QUERIES_1})],\n ),\n ).expect_reasoning(\n reasoning_tokens=tokenise(THINKING_RESPONSE_1),\n turn_index=0,\n ).expect(\n Packet(\n placement=create_placement(1),\n obj=TopLevelBranching(\n num_parallel_branches=2,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 0),\n obj=SearchToolStart(\n is_internet_search=False,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 1),\n obj=SearchToolStart(\n is_internet_search=True,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 0),\n obj=SearchToolQueriesDelta(\n queries=INTERNAL_QUERIES_1 + [QUERY],\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 0),\n obj=SearchToolDocumentsDelta(\n documents=expected_internal_docs,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 0),\n obj=SectionEnd(),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 1),\n obj=SearchToolQueriesDelta(\n queries=WEB_QUERIES_1,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 1),\n obj=SearchToolDocumentsDelta(\n documents=expected_web_docs,\n ),\n )\n ).expect(\n Packet(\n placement=create_placement(1, 1),\n obj=SectionEnd(),\n )\n ).run_and_validate(\n stream=answer_stream\n )\n\n # LLM Stream Response 2\n for content in OPEN_URL_DOCUMENTS_1:\n mock_content.add_content(content)\n\n handler.add_response(\n LLMToolCallResponse(\n tool_name=\"open_url\",\n tool_call_id=\"456\",\n tool_call_argument_tokens=[json.dumps({\"urls\": OPEN_URL_URLS_1})],\n )\n ).expect(\n Packet(\n placement=create_placement(2, 0),\n obj=OpenUrlStart(),\n ),\n forward=2, # Need both header + argument tokens for the tool call\n ).expect(\n Packet(\n placement=create_placement(2, 0),\n obj=OpenUrlUrls(urls=OPEN_URL_URLS_1),\n ),\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(2, 0),\n obj=OpenUrlDocuments(\n documents=[\n mock_web_content_to_search_doc(content)\n for content in OPEN_URL_DOCUMENTS_1\n ]\n ),\n ),\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(2, 0),\n obj=SectionEnd(),\n ),\n forward=False,\n ).run_and_validate(\n stream=answer_stream\n )\n\n # LLM Stream Response 3\n for content in OPEN_URL_DOCUMENTS_2:\n mock_content.add_content(content)\n\n handler.add_response(\n LLMReasoningResponse(\n reasoning_tokens=tokenise(THINKING_RESPONSE_2),\n )\n ).add_response(\n LLMToolCallResponse(\n tool_name=\"open_url\",\n tool_call_id=\"789\",\n tool_call_argument_tokens=[json.dumps({\"urls\": OPEN_URL_URLS_2})],\n )\n ).expect_reasoning(\n reasoning_tokens=tokenise(THINKING_RESPONSE_2),\n turn_index=3,\n ).expect(\n Packet(\n placement=create_placement(4),\n obj=OpenUrlStart(),\n )\n ).expect(\n Packet(placement=create_placement(4), obj=OpenUrlUrls(urls=OPEN_URL_URLS_2))\n ).expect(\n Packet(\n placement=create_placement(4),\n obj=OpenUrlDocuments(\n documents=[\n mock_web_content_to_search_doc(content)\n for content in OPEN_URL_DOCUMENTS_2\n ]\n ),\n ),\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(4),\n obj=SectionEnd(),\n )\n ).run_and_validate(\n stream=answer_stream\n )\n\n # LLM Stream Response 4\n handler.add_response(\n LLMReasoningResponse(\n reasoning_tokens=tokenise(REASONING_RESPONSE_3),\n )\n ).add_response(\n LLMAnswerResponse(\n answer_tokens=tokenise(ANSWER_RESPONSE),\n )\n ).expect_reasoning(\n reasoning_tokens=tokenise(REASONING_RESPONSE_3),\n turn_index=5,\n ).expect_agent_response(\n answer_tokens=tokenise(ANSWER_RESPONSE),\n turn_index=6,\n final_documents=expected_internal_docs\n + expected_web_docs\n + [\n mock_web_content_to_search_doc(content)\n for content in OPEN_URL_DOCUMENTS_1\n ]\n + [\n mock_web_content_to_search_doc(content)\n for content in OPEN_URL_DOCUMENTS_2\n ],\n ).run_and_validate(\n stream=answer_stream\n )\n\n # End stream\n with pytest.raises(StopIteration):\n next(answer_stream)\n" }, { "path": "backend/tests/external_dependency_unit/answer/test_stream_chat_message_objects.py", "content": "import uuid\nfrom typing import Any\nfrom unittest.mock import MagicMock\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.models import AnswerStreamPart\nfrom onyx.chat.models import StreamingError\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.models import User\nfrom onyx.db.persona import upsert_persona\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.streaming_models import AgentResponseDelta\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\n@pytest.mark.skip(reason=\"Temporarily disabled\")\ndef test_stream_chat_message_objects_without_web_search(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n mock_external_deps: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test that when web search is requested but the persona has no web search tool,\n the system handles it gracefully and returns a message explaining that web\n search is not available.\n \"\"\"\n\n # Mock the model server HTTP calls for embeddings\n def mock_post(\n url: str,\n json: dict[str, Any] | None = None,\n headers: dict[str, str] | None = None, # noqa: ARG001\n **kwargs: Any, # noqa: ARG001\n ) -> MagicMock:\n \"\"\"Mock requests.post for model server embedding calls\"\"\"\n mock_response = MagicMock()\n\n # Check if this is a call to the embedding endpoint\n if \"encoder/bi-encoder-embed\" in url:\n # Return a mock embedding response\n # The embedding dimension doesn't matter for this test,\n # just needs to be a valid response structure\n num_texts = len(json.get(\"texts\", [])) if json else 1\n mock_response.status_code = 200\n mock_response.json.return_value = {\n \"embeddings\": [[0.1] * 768]\n * num_texts # 768 is a common embedding dimension\n }\n mock_response.raise_for_status = MagicMock()\n return mock_response\n\n # For other URLs, return a generic success response\n mock_response.status_code = 200\n mock_response.json.return_value = {}\n mock_response.raise_for_status = MagicMock()\n return mock_response\n\n # First, ensure we have an LLM provider set up\n ensure_default_llm_provider(db_session)\n\n # Create a test user\n test_user: User = create_test_user(db_session, email_prefix=\"test_web_search\")\n\n # Create a test persona explicitly WITHOUT any tools (including web search)\n # This ensures the test doesn't rely on the state of the default persona\n test_persona = upsert_persona(\n user=None, # System persona\n name=f\"Test Persona {uuid.uuid4()}\",\n description=\"Test persona with no tools for web search test\",\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n system_prompt=None,\n task_prompt=None,\n datetime_aware=None,\n is_public=True,\n db_session=db_session,\n tool_ids=[], # Explicitly no tools\n document_set_ids=None,\n is_listed=True,\n )\n\n # Create a chat session with our test persona\n chat_session = create_chat_session(\n db_session=db_session,\n description=\"Test web search without tool\",\n user_id=test_user.id if test_user else None,\n persona_id=test_persona.id,\n )\n # Create the chat message request with a query that attempts to force web search\n chat_request = SendMessageRequest(\n message=\"run a web search for 'Onyx'\",\n chat_session_id=chat_session.id,\n )\n # Call handle_stream_message_objects\n response_generator = handle_stream_message_objects(\n new_msg_req=chat_request,\n user=test_user,\n db_session=db_session,\n )\n # Collect all packets from the response\n raw_answer_stream: list[AnswerStreamPart] = []\n message_content = \"\"\n error_occurred = False\n\n for packet in response_generator:\n raw_answer_stream.append(packet)\n if isinstance(packet, Packet):\n if isinstance(packet.obj, AgentResponseDelta):\n # Direct MessageDelta (if not wrapped)\n if packet.obj.content:\n message_content += packet.obj.content\n elif isinstance(packet.obj, StreamingError):\n error_occurred = True\n break\n\n assert not error_occurred, \"Should not have received a streaming error\"\n\n # Verify that we got a response\n assert len(raw_answer_stream) > 0, \"Should have received at least some packets\"\n\n # Check if we got MessageResponseIDInfo packet (indicating message was created)\n has_message_id = any(\n isinstance(packet, MessageResponseIDInfo) for packet in raw_answer_stream\n )\n assert has_message_id, \"Should have received a message ID packet\"\n\n assert len(message_content) > 0, \"Should have received some message content\"\n\n\ndef test_nothing() -> None:\n assert True, \"This test is just to ensure the test suite is running\"\n" }, { "path": "backend/tests/external_dependency_unit/background/test_periodic_task_claim.py", "content": "\"\"\"External dependency unit tests for periodic task claiming.\n\nTests ``_try_claim_task`` and ``_try_run_periodic_task`` against real\nPostgreSQL, verifying happy-path behavior and concurrent-access safety.\n\nThe claim mechanism uses a transaction-scoped advisory lock + a KVStore\ntimestamp for cross-instance dedup. The DB session is released before\nthe task runs, so long-running tasks don't hold connections.\n\"\"\"\n\nimport time\nfrom collections.abc import Generator\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import ThreadPoolExecutor\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\nfrom unittest.mock import MagicMock\nfrom uuid import uuid4\n\nimport pytest\n\nfrom onyx.background.periodic_poller import _PeriodicTaskDef\nfrom onyx.background.periodic_poller import _try_claim_task\nfrom onyx.background.periodic_poller import _try_run_periodic_task\nfrom onyx.background.periodic_poller import PERIODIC_TASK_KV_PREFIX\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.models import KVStore\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n_TEST_LOCK_BASE = 90_000\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef _init_engine() -> None:\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n\ndef _make_task(\n *,\n name: str | None = None,\n interval: float = 3600,\n lock_id: int | None = None,\n run_fn: MagicMock | None = None,\n) -> _PeriodicTaskDef:\n return _PeriodicTaskDef(\n name=name if name is not None else f\"test-{uuid4().hex[:8]}\",\n interval_seconds=interval,\n lock_id=lock_id if lock_id is not None else _TEST_LOCK_BASE,\n run_fn=run_fn if run_fn is not None else MagicMock(),\n )\n\n\n@pytest.fixture(autouse=True)\ndef _cleanup_kv(\n tenant_context: None, # noqa: ARG001\n) -> Generator[None, None, None]:\n yield\n with get_session_with_current_tenant() as db_session:\n db_session.query(KVStore).filter(\n KVStore.key.like(f\"{PERIODIC_TASK_KV_PREFIX}test-%\")\n ).delete(synchronize_session=False)\n db_session.commit()\n\n\n# ------------------------------------------------------------------\n# Happy-path: _try_claim_task\n# ------------------------------------------------------------------\n\n\nclass TestClaimHappyPath:\n def test_first_claim_succeeds(self) -> None:\n assert _try_claim_task(_make_task()) is True\n\n def test_first_claim_creates_kv_row(self) -> None:\n task = _make_task()\n _try_claim_task(task)\n\n with get_session_with_current_tenant() as db_session:\n row = (\n db_session.query(KVStore)\n .filter_by(key=PERIODIC_TASK_KV_PREFIX + task.name)\n .first()\n )\n assert row is not None\n assert row.value is not None\n\n def test_second_claim_within_interval_fails(self) -> None:\n task = _make_task(interval=3600)\n assert _try_claim_task(task) is True\n assert _try_claim_task(task) is False\n\n def test_claim_after_interval_succeeds(self) -> None:\n task = _make_task(interval=1)\n assert _try_claim_task(task) is True\n\n kv_key = PERIODIC_TASK_KV_PREFIX + task.name\n with get_session_with_current_tenant() as db_session:\n row = db_session.query(KVStore).filter_by(key=kv_key).first()\n assert row is not None\n row.value = (datetime.now(timezone.utc) - timedelta(seconds=10)).isoformat()\n db_session.commit()\n\n assert _try_claim_task(task) is True\n\n\n# ------------------------------------------------------------------\n# Happy-path: _try_run_periodic_task\n# ------------------------------------------------------------------\n\n\nclass TestRunHappyPath:\n def test_runs_task_and_updates_last_run_at(self) -> None:\n mock_fn = MagicMock()\n task = _make_task(run_fn=mock_fn)\n\n _try_run_periodic_task(task)\n\n mock_fn.assert_called_once()\n assert task.last_run_at > 0\n\n def test_skips_when_in_memory_interval_not_elapsed(self) -> None:\n mock_fn = MagicMock()\n task = _make_task(run_fn=mock_fn, interval=3600)\n task.last_run_at = time.monotonic()\n\n _try_run_periodic_task(task)\n\n mock_fn.assert_not_called()\n\n def test_skips_when_db_claim_blocked(self) -> None:\n name = f\"test-{uuid4().hex[:8]}\"\n lock_id = _TEST_LOCK_BASE + 10\n\n _try_claim_task(_make_task(name=name, lock_id=lock_id, interval=3600))\n\n mock_fn = MagicMock()\n task = _make_task(name=name, lock_id=lock_id, interval=3600, run_fn=mock_fn)\n _try_run_periodic_task(task)\n\n mock_fn.assert_not_called()\n\n def test_task_exception_does_not_propagate(self) -> None:\n task = _make_task(run_fn=MagicMock(side_effect=RuntimeError(\"boom\")))\n _try_run_periodic_task(task)\n\n def test_claim_committed_before_task_runs(self) -> None:\n \"\"\"The KV claim must be visible in the DB when run_fn executes.\"\"\"\n task_name = f\"test-order-{uuid4().hex[:8]}\"\n kv_key = PERIODIC_TASK_KV_PREFIX + task_name\n claim_visible: list[bool] = []\n\n def check_claim() -> None:\n with get_session_with_current_tenant() as db_session:\n row = db_session.query(KVStore).filter_by(key=kv_key).first()\n claim_visible.append(row is not None and row.value is not None)\n\n task = _PeriodicTaskDef(\n name=task_name,\n interval_seconds=3600,\n lock_id=_TEST_LOCK_BASE + 11,\n run_fn=check_claim,\n )\n\n _try_run_periodic_task(task)\n\n assert claim_visible == [True]\n\n\n# ------------------------------------------------------------------\n# Concurrency: only one claimer should win\n# ------------------------------------------------------------------\n\n\nclass TestClaimConcurrency:\n def test_concurrent_claims_single_winner(self) -> None:\n \"\"\"Many threads claim the same task — exactly one should succeed.\"\"\"\n num_threads = 20\n task_name = f\"test-race-{uuid4().hex[:8]}\"\n lock_id = _TEST_LOCK_BASE + 20\n\n def claim() -> bool:\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n return _try_claim_task(\n _PeriodicTaskDef(\n name=task_name,\n interval_seconds=3600,\n lock_id=lock_id,\n run_fn=lambda: None,\n )\n )\n\n results: list[bool] = []\n with ThreadPoolExecutor(max_workers=num_threads) as executor:\n futures = [executor.submit(claim) for _ in range(num_threads)]\n for future in as_completed(futures):\n results.append(future.result())\n\n winners = sum(1 for r in results if r)\n assert winners == 1, f\"Expected 1 winner, got {winners}\"\n\n def test_concurrent_run_single_execution(self) -> None:\n \"\"\"Many threads run the same task — run_fn fires exactly once.\"\"\"\n num_threads = 20\n task_name = f\"test-run-race-{uuid4().hex[:8]}\"\n lock_id = _TEST_LOCK_BASE + 21\n counter = MagicMock()\n\n def run() -> None:\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n _try_run_periodic_task(\n _PeriodicTaskDef(\n name=task_name,\n interval_seconds=3600,\n lock_id=lock_id,\n run_fn=counter,\n )\n )\n\n with ThreadPoolExecutor(max_workers=num_threads) as executor:\n futures = [executor.submit(run) for _ in range(num_threads)]\n for future in as_completed(futures):\n future.result()\n\n assert (\n counter.call_count == 1\n ), f\"Expected run_fn called once, got {counter.call_count}\"\n\n def test_no_errors_under_contention(self) -> None:\n \"\"\"All threads complete without exceptions under high contention.\"\"\"\n num_threads = 30\n task_name = f\"test-err-{uuid4().hex[:8]}\"\n lock_id = _TEST_LOCK_BASE + 22\n errors: list[Exception] = []\n\n def claim() -> bool:\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n return _try_claim_task(\n _PeriodicTaskDef(\n name=task_name,\n interval_seconds=3600,\n lock_id=lock_id,\n run_fn=lambda: None,\n )\n )\n\n with ThreadPoolExecutor(max_workers=num_threads) as executor:\n futures = [executor.submit(claim) for _ in range(num_threads)]\n for future in as_completed(futures):\n try:\n future.result()\n except Exception as e:\n errors.append(e)\n\n assert errors == [], f\"Got {len(errors)} errors: {errors}\"\n" }, { "path": "backend/tests/external_dependency_unit/background/test_startup_recovery.py", "content": "\"\"\"External dependency unit tests for startup recovery (Step 10g).\n\nSeeds ``UserFile`` records in stuck states (PROCESSING, DELETING,\nneeds_project_sync) then calls ``recover_stuck_user_files`` and verifies\nthe drain loops pick them up via ``FOR UPDATE SKIP LOCKED``.\n\nUses real PostgreSQL (via ``db_session`` / ``tenant_context`` fixtures).\nThe per-file ``*_impl`` functions are mocked so no real file store or\nconnector is needed — we only verify that recovery finds and dispatches\nthe correct files.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import UUID\nfrom uuid import uuid4\n\nimport pytest\nimport sqlalchemy as sa\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.periodic_poller import recover_stuck_user_files\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import UserFile\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n_IMPL_MODULE = \"onyx.background.celery.tasks.user_file_processing.tasks\"\n\n\ndef _create_user_file(\n db_session: Session,\n user_id: object,\n *,\n status: UserFileStatus = UserFileStatus.PROCESSING,\n needs_project_sync: bool = False,\n needs_persona_sync: bool = False,\n) -> UserFile:\n uf = UserFile(\n id=uuid4(),\n user_id=user_id,\n file_id=f\"test_file_{uuid4().hex[:8]}\",\n name=f\"test_{uuid4().hex[:8]}.txt\",\n file_type=\"text/plain\",\n status=status,\n needs_project_sync=needs_project_sync,\n needs_persona_sync=needs_persona_sync,\n )\n db_session.add(uf)\n db_session.commit()\n db_session.refresh(uf)\n return uf\n\n\ndef _fake_delete_impl(\n user_file_id: str,\n tenant_id: str, # noqa: ARG001\n redis_locking: bool, # noqa: ARG001\n) -> None:\n \"\"\"Mock side-effect: delete the row so the drain loop terminates.\"\"\"\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n with get_session_with_current_tenant() as session:\n session.execute(sa.delete(UserFile).where(UserFile.id == UUID(user_file_id)))\n session.commit()\n\n\ndef _fake_sync_impl(\n user_file_id: str,\n tenant_id: str, # noqa: ARG001\n redis_locking: bool, # noqa: ARG001\n) -> None:\n \"\"\"Mock side-effect: clear sync flags so the drain loop terminates.\"\"\"\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n with get_session_with_current_tenant() as session:\n session.execute(\n sa.update(UserFile)\n .where(UserFile.id == UUID(user_file_id))\n .values(needs_project_sync=False, needs_persona_sync=False)\n )\n session.commit()\n\n\n@pytest.fixture()\ndef _cleanup_user_files(db_session: Session) -> Generator[list[UserFile], None, None]:\n \"\"\"Track created UserFile rows and delete them after each test.\"\"\"\n created: list[UserFile] = []\n yield created\n for uf in created:\n existing = db_session.get(UserFile, uf.id)\n if existing:\n db_session.delete(existing)\n db_session.commit()\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestRecoverProcessingFiles:\n \"\"\"Files in PROCESSING status are re-processed via the processing drain loop.\"\"\"\n\n def test_processing_files_recovered(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"recovery_proc\")\n uf = _create_user_file(db_session, user.id, status=UserFileStatus.PROCESSING)\n _cleanup_user_files.append(uf)\n\n mock_impl = MagicMock()\n with patch(f\"{_IMPL_MODULE}.process_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert (\n str(uf.id) in called_ids\n ), f\"Expected file {uf.id} to be recovered but got: {called_ids}\"\n\n def test_completed_files_not_recovered(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"recovery_comp\")\n uf = _create_user_file(db_session, user.id, status=UserFileStatus.COMPLETED)\n _cleanup_user_files.append(uf)\n\n mock_impl = MagicMock()\n with patch(f\"{_IMPL_MODULE}.process_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert (\n str(uf.id) not in called_ids\n ), f\"COMPLETED file {uf.id} should not have been recovered\"\n\n\nclass TestRecoverDeletingFiles:\n \"\"\"Files in DELETING status are recovered via the delete drain loop.\"\"\"\n\n def test_deleting_files_recovered(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"recovery_del\")\n uf = _create_user_file(db_session, user.id, status=UserFileStatus.DELETING)\n # Row is deleted by _fake_delete_impl, so no cleanup needed.\n\n mock_impl = MagicMock(side_effect=_fake_delete_impl)\n with patch(f\"{_IMPL_MODULE}.delete_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert (\n str(uf.id) in called_ids\n ), f\"Expected file {uf.id} to be recovered for deletion but got: {called_ids}\"\n\n\nclass TestRecoverSyncFiles:\n \"\"\"Files needing project/persona sync are recovered via the sync drain loop.\"\"\"\n\n def test_needs_project_sync_recovered(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"recovery_sync\")\n uf = _create_user_file(\n db_session,\n user.id,\n status=UserFileStatus.COMPLETED,\n needs_project_sync=True,\n )\n _cleanup_user_files.append(uf)\n\n mock_impl = MagicMock(side_effect=_fake_sync_impl)\n with patch(f\"{_IMPL_MODULE}.project_sync_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert (\n str(uf.id) in called_ids\n ), f\"Expected file {uf.id} to be recovered for sync but got: {called_ids}\"\n\n def test_needs_persona_sync_recovered(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"recovery_psync\")\n uf = _create_user_file(\n db_session,\n user.id,\n status=UserFileStatus.COMPLETED,\n needs_persona_sync=True,\n )\n _cleanup_user_files.append(uf)\n\n mock_impl = MagicMock(side_effect=_fake_sync_impl)\n with patch(f\"{_IMPL_MODULE}.project_sync_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert (\n str(uf.id) in called_ids\n ), f\"Expected file {uf.id} to be recovered for persona sync but got: {called_ids}\"\n\n\nclass TestRecoveryMultipleFiles:\n \"\"\"Recovery processes all stuck files in one pass, not just the first.\"\"\"\n\n def test_multiple_processing_files(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"recovery_multi\")\n files = []\n for _ in range(3):\n uf = _create_user_file(\n db_session, user.id, status=UserFileStatus.PROCESSING\n )\n _cleanup_user_files.append(uf)\n files.append(uf)\n\n mock_impl = MagicMock()\n with patch(f\"{_IMPL_MODULE}.process_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = {call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list}\n expected_ids = {str(uf.id) for uf in files}\n assert expected_ids.issubset(\n called_ids\n ), f\"Expected all {len(files)} files to be recovered. Missing: {expected_ids - called_ids}\"\n\n\nclass TestTransientFailures:\n \"\"\"Drain loops skip failed files, process the rest, and terminate.\"\"\"\n\n def test_processing_failure_skips_and_continues(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"fail_proc\")\n uf_fail = _create_user_file(\n db_session, user.id, status=UserFileStatus.PROCESSING\n )\n uf_ok = _create_user_file(db_session, user.id, status=UserFileStatus.PROCESSING)\n _cleanup_user_files.extend([uf_fail, uf_ok])\n\n fail_id = str(uf_fail.id)\n\n def side_effect(\n *,\n user_file_id: str,\n tenant_id: str, # noqa: ARG001\n redis_locking: bool, # noqa: ARG001\n ) -> None:\n if user_file_id == fail_id:\n raise RuntimeError(\"transient failure\")\n\n mock_impl = MagicMock(side_effect=side_effect)\n with patch(f\"{_IMPL_MODULE}.process_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert fail_id in called_ids, \"Failed file should have been attempted\"\n assert str(uf_ok.id) in called_ids, \"Healthy file should have been processed\"\n assert called_ids.count(fail_id) == 1, \"Failed file retried — infinite loop\"\n assert called_ids.count(str(uf_ok.id)) == 1\n\n def test_delete_failure_skips_and_continues(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"fail_del\")\n uf_fail = _create_user_file(db_session, user.id, status=UserFileStatus.DELETING)\n uf_ok = _create_user_file(db_session, user.id, status=UserFileStatus.DELETING)\n _cleanup_user_files.append(uf_fail)\n\n fail_id = str(uf_fail.id)\n\n def side_effect(\n *, user_file_id: str, tenant_id: str, redis_locking: bool\n ) -> None:\n if user_file_id == fail_id:\n raise RuntimeError(\"transient failure\")\n _fake_delete_impl(user_file_id, tenant_id, redis_locking)\n\n mock_impl = MagicMock(side_effect=side_effect)\n with patch(f\"{_IMPL_MODULE}.delete_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert fail_id in called_ids, \"Failed file should have been attempted\"\n assert str(uf_ok.id) in called_ids, \"Healthy file should have been deleted\"\n assert called_ids.count(fail_id) == 1, \"Failed file retried — infinite loop\"\n assert called_ids.count(str(uf_ok.id)) == 1\n\n def test_sync_failure_skips_and_continues(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n _cleanup_user_files: list[UserFile],\n ) -> None:\n user = create_test_user(db_session, \"fail_sync\")\n uf_fail = _create_user_file(\n db_session,\n user.id,\n status=UserFileStatus.COMPLETED,\n needs_project_sync=True,\n )\n uf_ok = _create_user_file(\n db_session,\n user.id,\n status=UserFileStatus.COMPLETED,\n needs_persona_sync=True,\n )\n _cleanup_user_files.extend([uf_fail, uf_ok])\n\n fail_id = str(uf_fail.id)\n\n def side_effect(\n *, user_file_id: str, tenant_id: str, redis_locking: bool\n ) -> None:\n if user_file_id == fail_id:\n raise RuntimeError(\"transient failure\")\n _fake_sync_impl(user_file_id, tenant_id, redis_locking)\n\n mock_impl = MagicMock(side_effect=side_effect)\n with patch(f\"{_IMPL_MODULE}.project_sync_user_file_impl\", mock_impl):\n recover_stuck_user_files(TEST_TENANT_ID)\n\n called_ids = [call.kwargs[\"user_file_id\"] for call in mock_impl.call_args_list]\n assert fail_id in called_ids, \"Failed file should have been attempted\"\n assert str(uf_ok.id) in called_ids, \"Healthy file should have been synced\"\n assert called_ids.count(fail_id) == 1, \"Failed file retried — infinite loop\"\n assert called_ids.count(str(uf_ok.id)) == 1\n" }, { "path": "backend/tests/external_dependency_unit/cache/conftest.py", "content": "\"\"\"Fixtures for cache backend tests.\n\nRequires a running PostgreSQL instance (and Redis for parity tests).\nRun with::\n\n python -m dotenv -f .vscode/.env run -- pytest tests/external_dependency_unit/cache/\n\"\"\"\n\nfrom collections.abc import Generator\n\nimport pytest\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.cache.postgres_backend import PostgresCacheBackend\nfrom onyx.cache.redis_backend import RedisCacheBackend\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\n@pytest.fixture(scope=\"session\", autouse=True)\ndef _init_db() -> Generator[None, None, None]:\n \"\"\"Initialize DB engine. Assumes Postgres has migrations applied (e.g. via docker compose).\"\"\"\n SqlEngine.init_engine(pool_size=5, max_overflow=2)\n yield\n\n\n@pytest.fixture(autouse=True)\ndef _tenant_context() -> Generator[None, None, None]:\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n try:\n yield\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\n@pytest.fixture\ndef pg_cache() -> PostgresCacheBackend:\n return PostgresCacheBackend(TEST_TENANT_ID)\n\n\n@pytest.fixture\ndef redis_cache() -> RedisCacheBackend:\n from onyx.redis.redis_pool import redis_pool\n\n return RedisCacheBackend(redis_pool.get_client(TEST_TENANT_ID))\n\n\n@pytest.fixture(params=[\"postgres\", \"redis\"], ids=[\"postgres\", \"redis\"])\ndef cache(\n request: pytest.FixtureRequest,\n pg_cache: PostgresCacheBackend,\n redis_cache: RedisCacheBackend,\n) -> CacheBackend:\n if request.param == \"postgres\":\n return pg_cache\n return redis_cache\n" }, { "path": "backend/tests/external_dependency_unit/cache/test_cache_backend_parity.py", "content": "\"\"\"Parameterized tests that run the same CacheBackend operations against\nboth Redis and PostgreSQL, asserting identical return values.\n\nEach test runs twice (once per backend) via the ``cache`` fixture defined\nin conftest.py.\n\"\"\"\n\nimport time\nfrom uuid import uuid4\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.cache.interface import TTL_KEY_NOT_FOUND\nfrom onyx.cache.interface import TTL_NO_EXPIRY\n\n\ndef _key() -> str:\n return f\"parity_{uuid4().hex[:12]}\"\n\n\nclass TestKVParity:\n def test_get_missing(self, cache: CacheBackend) -> None:\n assert cache.get(_key()) is None\n\n def test_get_set(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, b\"value\")\n assert cache.get(k) == b\"value\"\n\n def test_overwrite(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, b\"a\")\n cache.set(k, b\"b\")\n assert cache.get(k) == b\"b\"\n\n def test_set_string(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, \"hello\")\n assert cache.get(k) == b\"hello\"\n\n def test_set_int(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, 42)\n assert cache.get(k) == b\"42\"\n\n def test_delete(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, b\"x\")\n cache.delete(k)\n assert cache.get(k) is None\n\n def test_exists(self, cache: CacheBackend) -> None:\n k = _key()\n assert not cache.exists(k)\n cache.set(k, b\"x\")\n assert cache.exists(k)\n\n\nclass TestTTLParity:\n def test_ttl_missing(self, cache: CacheBackend) -> None:\n assert cache.ttl(_key()) == TTL_KEY_NOT_FOUND\n\n def test_ttl_no_expiry(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, b\"x\")\n assert cache.ttl(k) == TTL_NO_EXPIRY\n\n def test_ttl_remaining(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, b\"x\", ex=10)\n remaining = cache.ttl(k)\n assert 8 <= remaining <= 10\n\n def test_set_with_ttl_expires(self, cache: CacheBackend) -> None:\n k = _key()\n cache.set(k, b\"x\", ex=1)\n assert cache.get(k) == b\"x\"\n time.sleep(1.5)\n assert cache.get(k) is None\n\n\nclass TestLockParity:\n def test_acquire_release(self, cache: CacheBackend) -> None:\n lock = cache.lock(f\"parity_lock_{uuid4().hex[:8]}\")\n assert lock.acquire(blocking=False)\n assert lock.owned()\n lock.release()\n assert not lock.owned()\n\n\nclass TestListParity:\n def test_rpush_blpop(self, cache: CacheBackend) -> None:\n k = f\"parity_list_{uuid4().hex[:8]}\"\n cache.rpush(k, b\"item\")\n result = cache.blpop([k], timeout=1)\n assert result is not None\n assert result[1] == b\"item\"\n\n def test_blpop_timeout(self, cache: CacheBackend) -> None:\n result = cache.blpop([f\"parity_empty_{uuid4().hex[:8]}\"], timeout=1)\n assert result is None\n" }, { "path": "backend/tests/external_dependency_unit/cache/test_kv_store_cache_layer.py", "content": "\"\"\"Tests for PgRedisKVStore's cache layer integration with CacheBackend.\n\nVerifies that the KV store correctly uses the CacheBackend for caching\nin front of PostgreSQL: cache hits, cache misses falling through to PG,\ncache population after PG reads, cache invalidation on delete, and\ngraceful degradation when the cache backend raises.\n\nRequires running PostgreSQL.\n\"\"\"\n\nimport json\nfrom collections.abc import Generator\nfrom unittest.mock import MagicMock\n\nimport pytest\nfrom sqlalchemy import delete\n\nfrom onyx.cache.interface import CacheBackend\nfrom onyx.cache.postgres_backend import PostgresCacheBackend\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.models import CacheStore\nfrom onyx.db.models import KVStore\nfrom onyx.key_value_store.interface import KvKeyNotFoundError\nfrom onyx.key_value_store.store import PgRedisKVStore\nfrom onyx.key_value_store.store import REDIS_KEY_PREFIX\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\n@pytest.fixture(autouse=True)\ndef _clean_kv() -> Generator[None, None, None]:\n yield\n with get_session_with_tenant(tenant_id=TEST_TENANT_ID) as session:\n session.execute(delete(KVStore))\n session.execute(delete(CacheStore))\n session.commit()\n\n\n@pytest.fixture\ndef kv_store(pg_cache: PostgresCacheBackend) -> PgRedisKVStore:\n return PgRedisKVStore(cache=pg_cache)\n\n\nclass TestStoreAndLoad:\n def test_store_populates_cache_and_pg(\n self, kv_store: PgRedisKVStore, pg_cache: PostgresCacheBackend\n ) -> None:\n kv_store.store(\"k1\", {\"hello\": \"world\"})\n\n cached = pg_cache.get(REDIS_KEY_PREFIX + \"k1\")\n assert cached is not None\n assert json.loads(cached) == {\"hello\": \"world\"}\n\n loaded = kv_store.load(\"k1\")\n assert loaded == {\"hello\": \"world\"}\n\n def test_load_returns_cached_value_without_pg_hit(\n self, pg_cache: PostgresCacheBackend\n ) -> None:\n \"\"\"If the cache already has the value, PG should not be queried.\"\"\"\n pg_cache.set(REDIS_KEY_PREFIX + \"cached_only\", json.dumps({\"from\": \"cache\"}))\n kv = PgRedisKVStore(cache=pg_cache)\n assert kv.load(\"cached_only\") == {\"from\": \"cache\"}\n\n def test_load_falls_through_to_pg_on_cache_miss(\n self, kv_store: PgRedisKVStore, pg_cache: PostgresCacheBackend\n ) -> None:\n kv_store.store(\"k2\", [1, 2, 3])\n\n pg_cache.delete(REDIS_KEY_PREFIX + \"k2\")\n assert pg_cache.get(REDIS_KEY_PREFIX + \"k2\") is None\n\n loaded = kv_store.load(\"k2\")\n assert loaded == [1, 2, 3]\n\n repopulated = pg_cache.get(REDIS_KEY_PREFIX + \"k2\")\n assert repopulated is not None\n assert json.loads(repopulated) == [1, 2, 3]\n\n def test_load_with_refresh_cache_skips_cache(\n self, kv_store: PgRedisKVStore, pg_cache: PostgresCacheBackend\n ) -> None:\n kv_store.store(\"k3\", \"original\")\n\n pg_cache.set(REDIS_KEY_PREFIX + \"k3\", json.dumps(\"stale\"))\n\n loaded = kv_store.load(\"k3\", refresh_cache=True)\n assert loaded == \"original\"\n\n\nclass TestDelete:\n def test_delete_removes_from_cache_and_pg(\n self, kv_store: PgRedisKVStore, pg_cache: PostgresCacheBackend\n ) -> None:\n kv_store.store(\"del_me\", \"bye\")\n kv_store.delete(\"del_me\")\n\n assert pg_cache.get(REDIS_KEY_PREFIX + \"del_me\") is None\n\n with pytest.raises(KvKeyNotFoundError):\n kv_store.load(\"del_me\")\n\n def test_delete_missing_key_raises(self, kv_store: PgRedisKVStore) -> None:\n with pytest.raises(KvKeyNotFoundError):\n kv_store.delete(\"nonexistent\")\n\n\nclass TestCacheFailureGracefulDegradation:\n def test_store_succeeds_when_cache_set_raises(self) -> None:\n failing_cache = MagicMock(spec=CacheBackend)\n failing_cache.set.side_effect = ConnectionError(\"cache down\")\n\n kv = PgRedisKVStore(cache=failing_cache)\n kv.store(\"resilient\", {\"data\": True})\n\n working_cache = MagicMock(spec=CacheBackend)\n working_cache.get.return_value = None\n kv_reader = PgRedisKVStore(cache=working_cache)\n loaded = kv_reader.load(\"resilient\")\n assert loaded == {\"data\": True}\n\n def test_load_falls_through_when_cache_get_raises(self) -> None:\n failing_cache = MagicMock(spec=CacheBackend)\n failing_cache.get.side_effect = ConnectionError(\"cache down\")\n failing_cache.set.side_effect = ConnectionError(\"cache down\")\n\n kv = PgRedisKVStore(cache=failing_cache)\n kv.store(\"survive\", 42)\n loaded = kv.load(\"survive\")\n assert loaded == 42\n" }, { "path": "backend/tests/external_dependency_unit/cache/test_postgres_cache_backend.py", "content": "\"\"\"Tests for PostgresCacheBackend against real PostgreSQL.\n\nCovers every method on the backend: KV CRUD, TTL behaviour, advisory\nlocks (acquire / release / contention), list operations (rpush / blpop),\nand the periodic cleanup function.\n\"\"\"\n\nimport time\nfrom uuid import uuid4\n\nfrom sqlalchemy import select\n\nfrom onyx.cache.interface import TTL_KEY_NOT_FOUND\nfrom onyx.cache.interface import TTL_NO_EXPIRY\nfrom onyx.cache.postgres_backend import cleanup_expired_cache_entries\nfrom onyx.cache.postgres_backend import PostgresCacheBackend\nfrom onyx.db.models import CacheStore\n\n\ndef _key() -> str:\n return f\"test_{uuid4().hex[:12]}\"\n\n\n# ------------------------------------------------------------------\n# Basic KV\n# ------------------------------------------------------------------\n\n\nclass TestKV:\n def test_get_set(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"hello\")\n assert pg_cache.get(k) == b\"hello\"\n\n def test_get_missing(self, pg_cache: PostgresCacheBackend) -> None:\n assert pg_cache.get(_key()) is None\n\n def test_set_overwrite(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"first\")\n pg_cache.set(k, b\"second\")\n assert pg_cache.get(k) == b\"second\"\n\n def test_set_string_value(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, \"string_val\")\n assert pg_cache.get(k) == b\"string_val\"\n\n def test_set_int_value(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, 42)\n assert pg_cache.get(k) == b\"42\"\n\n def test_delete(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"to_delete\")\n pg_cache.delete(k)\n assert pg_cache.get(k) is None\n\n def test_delete_missing_is_noop(self, pg_cache: PostgresCacheBackend) -> None:\n pg_cache.delete(_key())\n\n def test_exists(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n assert not pg_cache.exists(k)\n pg_cache.set(k, b\"x\")\n assert pg_cache.exists(k)\n\n\n# ------------------------------------------------------------------\n# TTL\n# ------------------------------------------------------------------\n\n\nclass TestTTL:\n def test_set_with_ttl_expires(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"ephemeral\", ex=1)\n assert pg_cache.get(k) == b\"ephemeral\"\n time.sleep(1.5)\n assert pg_cache.get(k) is None\n\n def test_ttl_no_expiry(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"forever\")\n assert pg_cache.ttl(k) == TTL_NO_EXPIRY\n\n def test_ttl_missing_key(self, pg_cache: PostgresCacheBackend) -> None:\n assert pg_cache.ttl(_key()) == TTL_KEY_NOT_FOUND\n\n def test_ttl_remaining(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"x\", ex=10)\n remaining = pg_cache.ttl(k)\n assert 8 <= remaining <= 10\n\n def test_ttl_expired_key(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"x\", ex=1)\n time.sleep(1.5)\n assert pg_cache.ttl(k) == TTL_KEY_NOT_FOUND\n\n def test_expire_adds_ttl(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"x\")\n assert pg_cache.ttl(k) == TTL_NO_EXPIRY\n pg_cache.expire(k, 10)\n assert 8 <= pg_cache.ttl(k) <= 10\n\n def test_exists_respects_ttl(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"x\", ex=1)\n assert pg_cache.exists(k)\n time.sleep(1.5)\n assert not pg_cache.exists(k)\n\n\n# ------------------------------------------------------------------\n# Locks\n# ------------------------------------------------------------------\n\n\nclass TestLock:\n def test_acquire_release(self, pg_cache: PostgresCacheBackend) -> None:\n lock = pg_cache.lock(f\"lock_{uuid4().hex[:8]}\")\n assert lock.acquire(blocking=False)\n assert lock.owned()\n lock.release()\n assert not lock.owned()\n\n def test_contention(self, pg_cache: PostgresCacheBackend) -> None:\n name = f\"contention_{uuid4().hex[:8]}\"\n lock1 = pg_cache.lock(name)\n lock2 = pg_cache.lock(name)\n\n assert lock1.acquire(blocking=False)\n assert not lock2.acquire(blocking=False)\n\n lock1.release()\n assert lock2.acquire(blocking=False)\n lock2.release()\n\n def test_context_manager(self, pg_cache: PostgresCacheBackend) -> None:\n with pg_cache.lock(f\"ctx_{uuid4().hex[:8]}\") as lock:\n assert lock.owned()\n assert not lock.owned()\n\n def test_blocking_timeout(self, pg_cache: PostgresCacheBackend) -> None:\n name = f\"timeout_{uuid4().hex[:8]}\"\n holder = pg_cache.lock(name)\n holder.acquire(blocking=False)\n\n waiter = pg_cache.lock(name, timeout=0.3)\n start = time.monotonic()\n assert not waiter.acquire(blocking=True, blocking_timeout=0.3)\n elapsed = time.monotonic() - start\n assert elapsed >= 0.25\n\n holder.release()\n\n\n# ------------------------------------------------------------------\n# List (rpush / blpop)\n# ------------------------------------------------------------------\n\n\nclass TestList:\n def test_rpush_blpop(self, pg_cache: PostgresCacheBackend) -> None:\n k = f\"list_{uuid4().hex[:8]}\"\n pg_cache.rpush(k, b\"item1\")\n result = pg_cache.blpop([k], timeout=1)\n assert result is not None\n assert result == (k.encode(), b\"item1\")\n\n def test_blpop_timeout(self, pg_cache: PostgresCacheBackend) -> None:\n result = pg_cache.blpop([f\"empty_{uuid4().hex[:8]}\"], timeout=1)\n assert result is None\n\n def test_fifo_order(self, pg_cache: PostgresCacheBackend) -> None:\n k = f\"fifo_{uuid4().hex[:8]}\"\n pg_cache.rpush(k, b\"first\")\n time.sleep(0.01)\n pg_cache.rpush(k, b\"second\")\n\n r1 = pg_cache.blpop([k], timeout=1)\n r2 = pg_cache.blpop([k], timeout=1)\n assert r1 is not None and r1[1] == b\"first\"\n assert r2 is not None and r2[1] == b\"second\"\n\n def test_multiple_keys(self, pg_cache: PostgresCacheBackend) -> None:\n k1 = f\"mk1_{uuid4().hex[:8]}\"\n k2 = f\"mk2_{uuid4().hex[:8]}\"\n pg_cache.rpush(k2, b\"from_k2\")\n\n result = pg_cache.blpop([k1, k2], timeout=1)\n assert result is not None\n assert result == (k2.encode(), b\"from_k2\")\n\n\n# ------------------------------------------------------------------\n# Cleanup\n# ------------------------------------------------------------------\n\n\nclass TestCleanup:\n def test_removes_expired_rows(self, pg_cache: PostgresCacheBackend) -> None:\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n\n k = _key()\n pg_cache.set(k, b\"stale\", ex=1)\n time.sleep(1.5)\n cleanup_expired_cache_entries()\n\n stmt = select(CacheStore.key).where(CacheStore.key == k)\n with get_session_with_current_tenant() as session:\n row = session.execute(stmt).first()\n assert row is None, \"expired row should be physically deleted\"\n\n def test_preserves_unexpired_rows(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"fresh\", ex=300)\n cleanup_expired_cache_entries()\n assert pg_cache.get(k) == b\"fresh\"\n\n def test_preserves_no_ttl_rows(self, pg_cache: PostgresCacheBackend) -> None:\n k = _key()\n pg_cache.set(k, b\"permanent\")\n cleanup_expired_cache_entries()\n assert pg_cache.get(k) == b\"permanent\"\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_docfetching_priority.py", "content": "\"\"\"\nExternal dependency unit tests for document processing job priority.\n\nTests that first-time indexing connectors (no last_successful_index_time)\nget higher priority than re-indexing jobs from connectors that have\npreviously completed indexing.\n\nUses real Redis for locking and real database objects for CC pairs and search settings.\n\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.tasks.docfetching.task_creation_utils import (\n try_creating_docfetching_task,\n)\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import SearchSettings\nfrom onyx.redis.redis_pool import get_redis_client\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\ndef _create_test_connector(db_session: Session, name: str) -> Connector:\n \"\"\"Create a test connector with all required fields.\"\"\"\n connector = Connector(\n name=name,\n source=DocumentSource.FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={},\n refresh_freq=3600,\n )\n db_session.add(connector)\n db_session.commit()\n db_session.refresh(connector)\n return connector\n\n\ndef _create_test_credential(db_session: Session) -> Credential:\n \"\"\"Create a test credential with all required fields.\"\"\"\n credential = Credential(\n name=f\"test_credential_{uuid4().hex[:8]}\",\n source=DocumentSource.FILE,\n credential_json={},\n admin_public=True,\n )\n db_session.add(credential)\n db_session.commit()\n db_session.refresh(credential)\n return credential\n\n\ndef _create_test_cc_pair(\n db_session: Session,\n connector: Connector,\n credential: Credential,\n status: ConnectorCredentialPairStatus,\n name: str,\n last_successful_index_time: datetime | None = None,\n) -> ConnectorCredentialPair:\n \"\"\"Create a connector credential pair with the specified status.\"\"\"\n cc_pair = ConnectorCredentialPair(\n name=name,\n connector_id=connector.id,\n credential_id=credential.id,\n status=status,\n access_type=AccessType.PUBLIC,\n last_successful_index_time=last_successful_index_time,\n )\n db_session.add(cc_pair)\n db_session.commit()\n db_session.refresh(cc_pair)\n return cc_pair\n\n\ndef _create_test_search_settings(\n db_session: Session, index_name: str\n) -> SearchSettings:\n \"\"\"Create test search settings with all required fields.\"\"\"\n search_settings = SearchSettings(\n model_name=\"test-model\",\n model_dim=768,\n normalize=True,\n query_prefix=\"\",\n passage_prefix=\"\",\n status=IndexModelStatus.PRESENT,\n index_name=index_name,\n embedding_precision=EmbeddingPrecision.FLOAT,\n )\n db_session.add(search_settings)\n db_session.commit()\n db_session.refresh(search_settings)\n return search_settings\n\n\nclass TestDocfetchingTaskPriorityWithRealObjects:\n \"\"\"\n Tests for document fetching task priority based on last_successful_index_time.\n\n Uses real Redis for locking and real database objects for CC pairs\n and search settings.\n \"\"\"\n\n @pytest.mark.parametrize(\n \"has_successful_index,expected_priority\",\n [\n # First-time indexing (no last_successful_index_time) should get HIGH priority\n (False, OnyxCeleryPriority.HIGH),\n # Re-indexing (has last_successful_index_time) should get MEDIUM priority\n (True, OnyxCeleryPriority.MEDIUM),\n ],\n )\n @patch(\n \"onyx.background.celery.tasks.docfetching.task_creation_utils.IndexingCoordination.try_create_index_attempt\"\n )\n def test_priority_based_on_last_successful_index_time(\n self,\n mock_try_create_index_attempt: MagicMock,\n db_session: Session,\n has_successful_index: bool,\n expected_priority: OnyxCeleryPriority,\n ) -> None:\n \"\"\"\n Test that first-time indexing connectors get higher priority than re-indexing.\n\n Priority is determined by last_successful_index_time:\n - None (never indexed): HIGH priority\n - Has timestamp (previously indexed): MEDIUM priority\n\n Uses real Redis for locking and real database objects.\n \"\"\"\n # Create unique names to avoid conflicts between test runs\n unique_suffix = uuid4().hex[:8]\n\n # Determine last_successful_index_time based on the test case\n last_successful_index_time = (\n datetime.now(timezone.utc) if has_successful_index else None\n )\n\n # Create real database objects\n connector = _create_test_connector(\n db_session, f\"test_connector_{has_successful_index}_{unique_suffix}\"\n )\n credential = _create_test_credential(db_session)\n cc_pair = _create_test_cc_pair(\n db_session,\n connector,\n credential,\n ConnectorCredentialPairStatus.ACTIVE,\n name=f\"test_cc_pair_{has_successful_index}_{unique_suffix}\",\n last_successful_index_time=last_successful_index_time,\n )\n search_settings = _create_test_search_settings(\n db_session, f\"test_index_{unique_suffix}\"\n )\n\n # Mock the index attempt creation to return a valid ID\n mock_try_create_index_attempt.return_value = 12345\n\n # Mock celery app to capture task submission\n mock_celery_app = MagicMock()\n mock_celery_app.send_task.return_value = MagicMock()\n\n # Use real Redis client\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n\n # Call the function with real objects\n result = try_creating_docfetching_task(\n celery_app=mock_celery_app,\n cc_pair=cc_pair,\n search_settings=search_settings,\n reindex=False,\n db_session=db_session,\n r=redis_client,\n tenant_id=TEST_TENANT_ID,\n )\n\n # Verify task was created\n assert result == 12345\n\n # Verify send_task was called with the expected priority\n mock_celery_app.send_task.assert_called_once()\n call_kwargs = mock_celery_app.send_task.call_args\n actual_priority = call_kwargs.kwargs[\"priority\"]\n assert (\n actual_priority == expected_priority\n ), f\"Expected priority {expected_priority} for has_successful_index={has_successful_index}, but got {actual_priority}\"\n\n @patch(\n \"onyx.background.celery.tasks.docfetching.task_creation_utils.IndexingCoordination.try_create_index_attempt\"\n )\n def test_no_task_created_when_deleting(\n self,\n mock_try_create_index_attempt: MagicMock,\n db_session: Session,\n ) -> None:\n \"\"\"Test that no task is created when connector is in DELETING status.\"\"\"\n unique_suffix = uuid4().hex[:8]\n\n connector = _create_test_connector(\n db_session, f\"test_connector_deleting_{unique_suffix}\"\n )\n credential = _create_test_credential(db_session)\n cc_pair = _create_test_cc_pair(\n db_session,\n connector,\n credential,\n ConnectorCredentialPairStatus.DELETING,\n name=f\"test_cc_pair_deleting_{unique_suffix}\",\n )\n search_settings = _create_test_search_settings(\n db_session, f\"test_index_deleting_{unique_suffix}\"\n )\n\n mock_celery_app = MagicMock()\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n\n result = try_creating_docfetching_task(\n celery_app=mock_celery_app,\n cc_pair=cc_pair,\n search_settings=search_settings,\n reindex=False,\n db_session=db_session,\n r=redis_client,\n tenant_id=TEST_TENANT_ID,\n )\n\n # Verify no task was created\n assert result is None\n mock_celery_app.send_task.assert_not_called()\n mock_try_create_index_attempt.assert_not_called()\n\n @patch(\n \"onyx.background.celery.tasks.docfetching.task_creation_utils.IndexingCoordination.try_create_index_attempt\"\n )\n def test_redis_lock_prevents_concurrent_task_creation(\n self,\n mock_try_create_index_attempt: MagicMock,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test that the Redis lock prevents concurrent task creation attempts.\n\n This test uses real Redis to verify the locking mechanism works correctly.\n When the lock is already held, the function should return None without\n attempting to create a task.\n \"\"\"\n unique_suffix = uuid4().hex[:8]\n\n connector = _create_test_connector(\n db_session, f\"test_connector_lock_{unique_suffix}\"\n )\n credential = _create_test_credential(db_session)\n cc_pair = _create_test_cc_pair(\n db_session,\n connector,\n credential,\n ConnectorCredentialPairStatus.INITIAL_INDEXING,\n name=f\"test_cc_pair_lock_{unique_suffix}\",\n )\n search_settings = _create_test_search_settings(\n db_session, f\"test_index_lock_{unique_suffix}\"\n )\n\n mock_try_create_index_attempt.return_value = 12345\n mock_celery_app = MagicMock()\n mock_celery_app.send_task.return_value = MagicMock()\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n\n # Acquire the lock before calling the function\n from onyx.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX\n\n lock = redis_client.lock(\n DANSWER_REDIS_FUNCTION_LOCK_PREFIX + \"try_creating_indexing_task\",\n timeout=30,\n )\n\n try:\n acquired = lock.acquire(blocking=False)\n assert acquired, \"Failed to acquire lock for test\"\n\n # Now try to create a task - should fail because lock is held\n result = try_creating_docfetching_task(\n celery_app=mock_celery_app,\n cc_pair=cc_pair,\n search_settings=search_settings,\n reindex=False,\n db_session=db_session,\n r=redis_client,\n tenant_id=TEST_TENANT_ID,\n )\n\n # Should return None because lock couldn't be acquired\n assert result is None\n mock_celery_app.send_task.assert_not_called()\n\n finally:\n # Always release the lock\n if lock.owned():\n lock.release()\n\n @patch(\n \"onyx.background.celery.tasks.docfetching.task_creation_utils.IndexingCoordination.try_create_index_attempt\"\n )\n def test_lock_released_after_successful_task_creation(\n self,\n mock_try_create_index_attempt: MagicMock,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test that the Redis lock is released after successful task creation.\n\n This verifies that subsequent calls can acquire the lock and create tasks.\n \"\"\"\n unique_suffix = uuid4().hex[:8]\n\n connector = _create_test_connector(\n db_session, f\"test_connector_release_{unique_suffix}\"\n )\n credential = _create_test_credential(db_session)\n cc_pair = _create_test_cc_pair(\n db_session,\n connector,\n credential,\n ConnectorCredentialPairStatus.INITIAL_INDEXING,\n name=f\"test_cc_pair_release_{unique_suffix}\",\n )\n search_settings = _create_test_search_settings(\n db_session, f\"test_index_release_{unique_suffix}\"\n )\n\n mock_try_create_index_attempt.return_value = 12345\n mock_celery_app = MagicMock()\n mock_celery_app.send_task.return_value = MagicMock()\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n\n # First call should succeed\n result1 = try_creating_docfetching_task(\n celery_app=mock_celery_app,\n cc_pair=cc_pair,\n search_settings=search_settings,\n reindex=False,\n db_session=db_session,\n r=redis_client,\n tenant_id=TEST_TENANT_ID,\n )\n assert result1 == 12345\n\n # Reset mocks for second call\n mock_celery_app.reset_mock()\n mock_try_create_index_attempt.reset_mock()\n mock_try_create_index_attempt.return_value = 67890\n\n # Second call should also succeed (lock was released)\n result2 = try_creating_docfetching_task(\n celery_app=mock_celery_app,\n cc_pair=cc_pair,\n search_settings=search_settings,\n reindex=False,\n db_session=db_session,\n r=redis_client,\n tenant_id=TEST_TENANT_ID,\n )\n assert result2 == 67890\n\n # Both calls should have submitted tasks\n mock_celery_app.send_task.assert_called_once()\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_docprocessing_priority.py", "content": "\"\"\"\nExternal dependency unit tests for docprocessing task priority.\n\nTests that docprocessing tasks spawned by connector_document_extraction\nget the correct priority based on last_successful_index_time.\n\nUses real database objects for CC pairs, search settings, and index attempts.\n\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.indexing.run_docfetching import connector_document_extraction\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.configs.constants import OnyxCeleryPriority\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import SearchSettings\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\ndef _create_test_connector(db_session: Session, name: str) -> Connector:\n \"\"\"Create a test connector with all required fields.\"\"\"\n connector = Connector(\n name=name,\n source=DocumentSource.FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={},\n refresh_freq=3600,\n )\n db_session.add(connector)\n db_session.commit()\n db_session.refresh(connector)\n return connector\n\n\ndef _create_test_credential(db_session: Session) -> Credential:\n \"\"\"Create a test credential with all required fields.\"\"\"\n credential = Credential(\n name=f\"test_credential_{uuid4().hex[:8]}\",\n source=DocumentSource.FILE,\n credential_json={},\n admin_public=True,\n )\n db_session.add(credential)\n db_session.commit()\n db_session.refresh(credential)\n return credential\n\n\ndef _create_test_cc_pair(\n db_session: Session,\n connector: Connector,\n credential: Credential,\n status: ConnectorCredentialPairStatus,\n name: str,\n last_successful_index_time: datetime | None = None,\n) -> ConnectorCredentialPair:\n \"\"\"Create a connector credential pair with the specified status.\"\"\"\n cc_pair = ConnectorCredentialPair(\n name=name,\n connector_id=connector.id,\n credential_id=credential.id,\n status=status,\n access_type=AccessType.PUBLIC,\n last_successful_index_time=last_successful_index_time,\n )\n db_session.add(cc_pair)\n db_session.commit()\n db_session.refresh(cc_pair)\n return cc_pair\n\n\ndef _create_test_search_settings(\n db_session: Session, index_name: str\n) -> SearchSettings:\n \"\"\"Create test search settings with all required fields.\"\"\"\n search_settings = SearchSettings(\n model_name=\"test-model\",\n model_dim=768,\n normalize=True,\n query_prefix=\"\",\n passage_prefix=\"\",\n status=IndexModelStatus.PRESENT,\n index_name=index_name,\n embedding_precision=EmbeddingPrecision.FLOAT,\n )\n db_session.add(search_settings)\n db_session.commit()\n db_session.refresh(search_settings)\n return search_settings\n\n\ndef _create_test_index_attempt(\n db_session: Session,\n cc_pair: ConnectorCredentialPair,\n search_settings: SearchSettings,\n from_beginning: bool = False,\n) -> IndexAttempt:\n \"\"\"Create a test index attempt with the specified cc_pair and search_settings.\"\"\"\n index_attempt = IndexAttempt(\n connector_credential_pair_id=cc_pair.id,\n search_settings_id=search_settings.id,\n from_beginning=from_beginning,\n status=IndexingStatus.IN_PROGRESS,\n celery_task_id=f\"test_celery_task_{uuid4().hex[:8]}\",\n )\n db_session.add(index_attempt)\n db_session.commit()\n db_session.refresh(index_attempt)\n return index_attempt\n\n\nclass TestDocprocessingPriorityInDocumentExtraction:\n \"\"\"\n Tests for docprocessing task priority within connector_document_extraction.\n\n Verifies that the priority passed to docprocessing tasks is determined\n by last_successful_index_time on the cc_pair.\n \"\"\"\n\n @pytest.mark.parametrize(\n \"has_successful_index,expected_priority\",\n [\n # First-time indexing (no last_successful_index_time) should get HIGH priority\n (False, OnyxCeleryPriority.HIGH),\n # Re-indexing (has last_successful_index_time) should get MEDIUM priority\n (True, OnyxCeleryPriority.MEDIUM),\n ],\n )\n @patch(\"onyx.background.indexing.run_docfetching.get_document_batch_storage\")\n @patch(\"onyx.background.indexing.run_docfetching.MemoryTracer\")\n @patch(\"onyx.background.indexing.run_docfetching._get_connector_runner\")\n @patch(\n \"onyx.background.indexing.run_docfetching.strip_null_characters\",\n side_effect=lambda batch: batch,\n )\n @patch(\n \"onyx.background.indexing.run_docfetching.get_recent_completed_attempts_for_cc_pair\"\n )\n @patch(\n \"onyx.background.indexing.run_docfetching.get_last_successful_attempt_poll_range_end\"\n )\n @patch(\"onyx.background.indexing.run_docfetching.save_checkpoint\")\n @patch(\"onyx.background.indexing.run_docfetching.get_latest_valid_checkpoint\")\n @patch(\"onyx.background.indexing.run_docfetching.get_redis_client\")\n @patch(\"onyx.background.indexing.run_docfetching.ensure_source_node_exists\")\n @patch(\"onyx.background.indexing.run_docfetching.get_source_node_id_from_cache\")\n @patch(\"onyx.background.indexing.run_docfetching.get_node_id_from_raw_id\")\n @patch(\"onyx.background.indexing.run_docfetching.cache_hierarchy_nodes_batch\")\n def test_docprocessing_priority_based_on_last_successful_index_time(\n self,\n mock_cache_hierarchy_nodes_batch: MagicMock, # noqa: ARG002\n mock_get_node_id_from_raw_id: MagicMock,\n mock_get_source_node_id_from_cache: MagicMock,\n mock_ensure_source_node_exists: MagicMock,\n mock_get_redis_client: MagicMock,\n mock_get_latest_valid_checkpoint: MagicMock,\n mock_save_checkpoint: MagicMock, # noqa: ARG002\n mock_get_last_successful_attempt_poll_range_end: MagicMock,\n mock_get_recent_completed_attempts: MagicMock,\n mock_strip_null_characters: MagicMock, # noqa: ARG002\n mock_get_connector_runner: MagicMock,\n mock_memory_tracer_class: MagicMock,\n mock_get_batch_storage: MagicMock,\n db_session: Session,\n has_successful_index: bool,\n expected_priority: OnyxCeleryPriority,\n ) -> None:\n \"\"\"\n Test that docprocessing tasks get the correct priority based on\n last_successful_index_time.\n\n Priority is determined by last_successful_index_time:\n - None (never indexed): HIGH priority\n - Has timestamp (previously indexed): MEDIUM priority\n\n Uses real database objects for CC pairs and search settings.\n \"\"\"\n unique_suffix = uuid4().hex[:8]\n\n # Determine last_successful_index_time based on the test case\n last_successful_index_time = (\n datetime.now(timezone.utc) if has_successful_index else None\n )\n\n # Create real database objects\n connector = _create_test_connector(\n db_session, f\"test_connector_docproc_{has_successful_index}_{unique_suffix}\"\n )\n credential = _create_test_credential(db_session)\n cc_pair = _create_test_cc_pair(\n db_session,\n connector,\n credential,\n ConnectorCredentialPairStatus.ACTIVE,\n name=f\"test_cc_pair_docproc_{has_successful_index}_{unique_suffix}\",\n last_successful_index_time=last_successful_index_time,\n )\n search_settings = _create_test_search_settings(\n db_session, f\"test_index_docproc_{unique_suffix}\"\n )\n index_attempt = _create_test_index_attempt(\n db_session, cc_pair, search_settings, from_beginning=False\n )\n\n # Setup mocks\n mock_batch_storage = MagicMock()\n mock_get_batch_storage.return_value = mock_batch_storage\n\n mock_memory_tracer = MagicMock()\n mock_memory_tracer_class.return_value = mock_memory_tracer\n\n # Mock Redis-related functions (not the focus of this test)\n # Configure mock Redis client to return None for common operations\n # as a safety net in case any patches don't work as expected\n mock_redis_client = MagicMock()\n mock_redis_client.get.return_value = None\n mock_redis_client.hget.return_value = None\n mock_redis_client.hset.return_value = None\n mock_redis_client.exists.return_value = 0\n mock_redis_client.expire.return_value = True\n mock_get_redis_client.return_value = mock_redis_client\n\n # Mock hierarchy/cache functions\n mock_ensure_source_node_exists.return_value = 1 # Return a valid node ID\n mock_get_source_node_id_from_cache.return_value = (\n 1 # Return a valid source node ID\n )\n mock_get_node_id_from_raw_id.return_value = (None, False) # (node_id, found)\n # cache_hierarchy_nodes_batch doesn't need a return value (returns None)\n\n # Create checkpoint mocks - initial checkpoint has_more=True, final has_more=False\n mock_initial_checkpoint = MagicMock(has_more=True)\n mock_final_checkpoint = MagicMock(has_more=False)\n\n # get_latest_valid_checkpoint returns (checkpoint, resuming_from_checkpoint)\n mock_get_latest_valid_checkpoint.return_value = (mock_initial_checkpoint, False)\n\n # Create a mock connector runner that yields one document batch\n mock_connector = MagicMock()\n mock_connector_runner = MagicMock()\n mock_connector_runner.connector = mock_connector\n # The connector runner yields (document_batch, hierarchy_nodes, failure, next_checkpoint)\n # We provide one batch of documents to trigger a send_task call\n mock_doc = MagicMock()\n mock_doc.to_short_descriptor.return_value = \"test_doc\"\n mock_doc.sections = []\n # Set to None to avoid Redis operations trying to resolve hierarchy\n mock_doc.parent_hierarchy_raw_node_id = None\n mock_doc.parent_hierarchy_node_id = None\n mock_connector_runner.run.return_value = iter(\n [([mock_doc], None, None, mock_final_checkpoint)]\n )\n mock_get_connector_runner.return_value = mock_connector_runner\n\n mock_get_recent_completed_attempts.return_value = iter([])\n mock_get_last_successful_attempt_poll_range_end.return_value = 0\n\n # Mock celery app to capture task submission\n mock_celery_app = MagicMock()\n mock_celery_app.send_task.return_value = MagicMock()\n\n # Call the function\n connector_document_extraction(\n app=mock_celery_app,\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n search_settings_id=search_settings.id,\n tenant_id=TEST_TENANT_ID,\n callback=None,\n )\n\n # Verify send_task was called with the expected priority for docprocessing\n assert mock_celery_app.send_task.called, \"send_task should have been called\"\n call_kwargs = mock_celery_app.send_task.call_args\n actual_priority = call_kwargs.kwargs[\"priority\"]\n assert (\n actual_priority == expected_priority\n ), f\"Expected priority {expected_priority} for has_successful_index={has_successful_index}, but got {actual_priority}\"\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_persona_file_sync.py", "content": "\"\"\"\nExternal dependency unit tests for persona file sync.\n\nValidates that:\n\n1. The check_for_user_file_project_sync beat task picks up UserFiles with\n needs_persona_sync=True (not just needs_project_sync).\n\n2. The process_single_user_file_project_sync worker task reads persona\n associations from the DB, passes persona_ids to the document index via\n VespaDocumentUserFields, and clears needs_persona_sync afterwards.\n\n3. upsert_persona correctly marks affected UserFiles with\n needs_persona_sync=True when file associations change.\n\nUses real Redis and PostgreSQL. Document index (Vespa) calls are mocked\nsince we only need to verify the arguments passed to update_single.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom unittest.mock import PropertyMock\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n check_for_user_file_project_sync,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n process_single_user_file_project_sync,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n user_file_project_sync_lock_key,\n)\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__UserFile\nfrom onyx.db.models import User\nfrom onyx.db.models import UserFile\nfrom onyx.db.persona import upsert_persona\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\nfrom onyx.redis.redis_pool import get_redis_client\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _create_completed_user_file(\n db_session: Session,\n user: User,\n needs_persona_sync: bool = False,\n needs_project_sync: bool = False,\n) -> UserFile:\n \"\"\"Insert a UserFile in COMPLETED status.\"\"\"\n uf = UserFile(\n id=uuid4(),\n user_id=user.id,\n file_id=f\"test_file_{uuid4().hex[:8]}\",\n name=f\"test_{uuid4().hex[:8]}.txt\",\n file_type=\"text/plain\",\n status=UserFileStatus.COMPLETED,\n needs_persona_sync=needs_persona_sync,\n needs_project_sync=needs_project_sync,\n chunk_count=5,\n )\n db_session.add(uf)\n db_session.commit()\n db_session.refresh(uf)\n return uf\n\n\ndef _create_test_persona(\n db_session: Session,\n user: User,\n user_files: list[UserFile] | None = None,\n) -> Persona:\n \"\"\"Create a minimal Persona via direct model insert.\"\"\"\n persona = Persona(\n name=f\"Test Persona {uuid4().hex[:8]}\",\n description=\"Test persona\",\n system_prompt=\"You are a test assistant\",\n task_prompt=\"Answer the question\",\n tools=[],\n document_sets=[],\n users=[user],\n groups=[],\n is_listed=True,\n is_public=True,\n display_priority=None,\n starter_messages=None,\n deleted=False,\n user_files=user_files or [],\n user_id=user.id,\n )\n db_session.add(persona)\n db_session.commit()\n db_session.refresh(persona)\n return persona\n\n\ndef _link_file_to_persona(\n db_session: Session, persona: Persona, user_file: UserFile\n) -> None:\n \"\"\"Create the join table row between a persona and a user file.\"\"\"\n link = Persona__UserFile(persona_id=persona.id, user_file_id=user_file.id)\n db_session.add(link)\n db_session.commit()\n\n\n_PATCH_QUEUE_DEPTH = \"onyx.background.celery.tasks.user_file_processing.tasks.get_user_file_project_sync_queue_depth\"\n\n\n@contextmanager\ndef _patch_task_app(task: Any, mock_app: MagicMock) -> Generator[None, None, None]:\n \"\"\"Patch the ``app`` property on a bound Celery task.\"\"\"\n task_instance = task.run.__self__\n with (\n patch.object(\n type(task_instance),\n \"app\",\n new_callable=PropertyMock,\n return_value=mock_app,\n ),\n patch(_PATCH_QUEUE_DEPTH, return_value=0),\n patch(\n \"onyx.background.celery.tasks.user_file_processing.tasks.celery_get_broker_client\",\n return_value=MagicMock(),\n ),\n ):\n yield\n\n\n# ---------------------------------------------------------------------------\n# Test: check_for_user_file_project_sync picks up persona sync\n# ---------------------------------------------------------------------------\n\n\nclass TestCheckSweepIncludesPersonaSync:\n \"\"\"The beat task must pick up files needing persona sync, not just project sync.\"\"\"\n\n def test_persona_sync_flag_enqueues_task(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file with needs_persona_sync=True (and COMPLETED) gets enqueued.\"\"\"\n user = create_test_user(db_session, \"persona_sweep\")\n uf = _create_completed_user_file(db_session, user, needs_persona_sync=True)\n\n mock_app = MagicMock()\n\n with _patch_task_app(check_for_user_file_project_sync, mock_app):\n check_for_user_file_project_sync.run(tenant_id=TEST_TENANT_ID)\n\n enqueued_ids = {\n call.kwargs[\"kwargs\"][\"user_file_id\"]\n for call in mock_app.send_task.call_args_list\n }\n assert str(uf.id) in enqueued_ids\n\n def test_neither_flag_does_not_enqueue(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file with both flags False is not enqueued.\"\"\"\n user = create_test_user(db_session, \"no_sync\")\n uf = _create_completed_user_file(db_session, user)\n\n mock_app = MagicMock()\n\n with _patch_task_app(check_for_user_file_project_sync, mock_app):\n check_for_user_file_project_sync.run(tenant_id=TEST_TENANT_ID)\n\n enqueued_ids = {\n call.kwargs[\"kwargs\"][\"user_file_id\"]\n for call in mock_app.send_task.call_args_list\n }\n assert str(uf.id) not in enqueued_ids\n\n def test_both_flags_enqueues_once(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file with BOTH flags True is enqueued exactly once.\"\"\"\n user = create_test_user(db_session, \"both_flags\")\n uf = _create_completed_user_file(\n db_session, user, needs_persona_sync=True, needs_project_sync=True\n )\n\n mock_app = MagicMock()\n\n with _patch_task_app(check_for_user_file_project_sync, mock_app):\n check_for_user_file_project_sync.run(tenant_id=TEST_TENANT_ID)\n\n matching_calls = [\n call\n for call in mock_app.send_task.call_args_list\n if call.kwargs[\"kwargs\"][\"user_file_id\"] == str(uf.id)\n ]\n assert len(matching_calls) == 1\n\n\n# ---------------------------------------------------------------------------\n# Test: process_single_user_file_project_sync passes persona_ids to index\n# ---------------------------------------------------------------------------\n\n_PATCH_GET_SETTINGS = (\n \"onyx.background.celery.tasks.user_file_processing.tasks.get_active_search_settings\"\n)\n_PATCH_GET_INDICES = (\n \"onyx.background.celery.tasks.user_file_processing.tasks.get_all_document_indices\"\n)\n_PATCH_HTTPX_INIT = (\n \"onyx.background.celery.tasks.user_file_processing.tasks.httpx_init_vespa_pool\"\n)\n_PATCH_DISABLE_VDB = (\n \"onyx.background.celery.tasks.user_file_processing.tasks.DISABLE_VECTOR_DB\"\n)\n\n\nclass TestSyncTaskWritesPersonaIds:\n \"\"\"The sync task reads persona associations and sends them to the index.\"\"\"\n\n def test_passes_persona_ids_to_update_single(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"After linking a file to a persona, sync sends the persona ID.\"\"\"\n user = create_test_user(db_session, \"sync_persona\")\n uf = _create_completed_user_file(db_session, user, needs_persona_sync=True)\n persona = _create_test_persona(db_session, user)\n _link_file_to_persona(db_session, persona, uf)\n\n mock_doc_index = MagicMock()\n mock_search_settings = MagicMock()\n mock_search_settings.primary = MagicMock()\n mock_search_settings.secondary = None\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n lock_key = user_file_project_sync_lock_key(str(uf.id))\n redis_client.delete(lock_key)\n\n with (\n patch(_PATCH_DISABLE_VDB, False),\n patch(_PATCH_HTTPX_INIT),\n patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings),\n patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]),\n ):\n process_single_user_file_project_sync.run(\n user_file_id=str(uf.id), tenant_id=TEST_TENANT_ID\n )\n\n mock_doc_index.update_single.assert_called_once()\n call_args = mock_doc_index.update_single.call_args\n user_fields: VespaDocumentUserFields = call_args.kwargs[\"user_fields\"]\n assert user_fields.personas is not None\n assert persona.id in user_fields.personas\n assert call_args.args[0] == str(uf.id)\n\n def test_clears_persona_sync_flag(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"After a successful sync the needs_persona_sync flag is cleared.\"\"\"\n user = create_test_user(db_session, \"sync_clear\")\n uf = _create_completed_user_file(db_session, user, needs_persona_sync=True)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n lock_key = user_file_project_sync_lock_key(str(uf.id))\n redis_client.delete(lock_key)\n\n with patch(_PATCH_DISABLE_VDB, True):\n process_single_user_file_project_sync.run(\n user_file_id=str(uf.id), tenant_id=TEST_TENANT_ID\n )\n\n db_session.refresh(uf)\n assert uf.needs_persona_sync is False\n\n def test_passes_both_project_and_persona_ids(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file linked to both a project and a persona gets both IDs.\"\"\"\n from onyx.db.models import Project__UserFile\n from onyx.db.models import UserProject\n\n user = create_test_user(db_session, \"sync_both\")\n uf = _create_completed_user_file(\n db_session, user, needs_persona_sync=True, needs_project_sync=True\n )\n persona = _create_test_persona(db_session, user)\n _link_file_to_persona(db_session, persona, uf)\n\n project = UserProject(user_id=user.id, name=\"test-project\", instructions=\"\")\n db_session.add(project)\n db_session.commit()\n db_session.refresh(project)\n\n link = Project__UserFile(project_id=project.id, user_file_id=uf.id)\n db_session.add(link)\n db_session.commit()\n\n mock_doc_index = MagicMock()\n mock_search_settings = MagicMock()\n mock_search_settings.primary = MagicMock()\n mock_search_settings.secondary = None\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n lock_key = user_file_project_sync_lock_key(str(uf.id))\n redis_client.delete(lock_key)\n\n with (\n patch(_PATCH_DISABLE_VDB, False),\n patch(_PATCH_HTTPX_INIT),\n patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings),\n patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]),\n ):\n process_single_user_file_project_sync.run(\n user_file_id=str(uf.id), tenant_id=TEST_TENANT_ID\n )\n\n call_kwargs = mock_doc_index.update_single.call_args.kwargs\n user_fields: VespaDocumentUserFields = call_kwargs[\"user_fields\"]\n assert user_fields.personas is not None\n assert user_fields.user_projects is not None\n assert persona.id in user_fields.personas\n assert project.id in user_fields.user_projects\n\n # Both flags should be cleared\n db_session.refresh(uf)\n assert uf.needs_persona_sync is False\n assert uf.needs_project_sync is False\n\n def test_deleted_persona_excluded_from_ids(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A soft-deleted persona should NOT appear in the persona_ids sent to Vespa.\"\"\"\n user = create_test_user(db_session, \"sync_deleted\")\n uf = _create_completed_user_file(db_session, user, needs_persona_sync=True)\n persona = _create_test_persona(db_session, user)\n _link_file_to_persona(db_session, persona, uf)\n\n persona.deleted = True\n db_session.commit()\n\n mock_doc_index = MagicMock()\n mock_search_settings = MagicMock()\n mock_search_settings.primary = MagicMock()\n mock_search_settings.secondary = None\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n lock_key = user_file_project_sync_lock_key(str(uf.id))\n redis_client.delete(lock_key)\n\n with (\n patch(_PATCH_DISABLE_VDB, False),\n patch(_PATCH_HTTPX_INIT),\n patch(_PATCH_GET_SETTINGS, return_value=mock_search_settings),\n patch(_PATCH_GET_INDICES, return_value=[mock_doc_index]),\n ):\n process_single_user_file_project_sync.run(\n user_file_id=str(uf.id), tenant_id=TEST_TENANT_ID\n )\n\n call_kwargs = mock_doc_index.update_single.call_args.kwargs\n user_fields: VespaDocumentUserFields = call_kwargs[\"user_fields\"]\n assert user_fields.personas is not None\n assert persona.id not in user_fields.personas\n\n\n# ---------------------------------------------------------------------------\n# Test: upsert_persona marks files for persona sync\n# ---------------------------------------------------------------------------\n\n\nclass TestUpsertPersonaMarksSyncFlag:\n \"\"\"upsert_persona must set needs_persona_sync on affected UserFiles.\"\"\"\n\n def test_creating_persona_with_files_marks_sync(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n user = create_test_user(db_session, \"upsert_create\")\n uf = _create_completed_user_file(db_session, user)\n assert uf.needs_persona_sync is False\n\n upsert_persona(\n user=user,\n name=f\"persona-{uuid4().hex[:8]}\",\n description=\"test\",\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n system_prompt=\"test\",\n task_prompt=\"test\",\n datetime_aware=None,\n is_public=True,\n db_session=db_session,\n user_file_ids=[uf.id],\n )\n\n db_session.refresh(uf)\n assert uf.needs_persona_sync is True\n\n def test_updating_persona_files_marks_both_old_and_new(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"When file associations change, both the removed and added files are flagged.\"\"\"\n user = create_test_user(db_session, \"upsert_update\")\n uf_old = _create_completed_user_file(db_session, user)\n uf_new = _create_completed_user_file(db_session, user)\n\n persona = upsert_persona(\n user=user,\n name=f\"persona-{uuid4().hex[:8]}\",\n description=\"test\",\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n system_prompt=\"test\",\n task_prompt=\"test\",\n datetime_aware=None,\n is_public=True,\n db_session=db_session,\n user_file_ids=[uf_old.id],\n )\n\n # Clear the flag from creation so we can observe the update\n uf_old.needs_persona_sync = False\n db_session.commit()\n\n # Now update the persona to swap files\n upsert_persona(\n user=user,\n name=persona.name,\n description=persona.description,\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n system_prompt=persona.system_prompt,\n task_prompt=persona.task_prompt,\n datetime_aware=None,\n is_public=persona.is_public,\n db_session=db_session,\n persona_id=persona.id,\n user_file_ids=[uf_new.id],\n )\n\n db_session.refresh(uf_old)\n db_session.refresh(uf_new)\n assert uf_old.needs_persona_sync is True, \"Removed file should be flagged\"\n assert uf_new.needs_persona_sync is True, \"Added file should be flagged\"\n\n def test_removing_all_files_marks_old_files(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Removing all files from a persona flags the previously associated files.\"\"\"\n user = create_test_user(db_session, \"upsert_remove\")\n uf = _create_completed_user_file(db_session, user)\n\n persona = upsert_persona(\n user=user,\n name=f\"persona-{uuid4().hex[:8]}\",\n description=\"test\",\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n system_prompt=\"test\",\n task_prompt=\"test\",\n datetime_aware=None,\n is_public=True,\n db_session=db_session,\n user_file_ids=[uf.id],\n )\n\n uf.needs_persona_sync = False\n db_session.commit()\n\n upsert_persona(\n user=user,\n name=persona.name,\n description=persona.description,\n llm_model_provider_override=None,\n llm_model_version_override=None,\n starter_messages=None,\n system_prompt=persona.system_prompt,\n task_prompt=persona.task_prompt,\n datetime_aware=None,\n is_public=persona.is_public,\n db_session=db_session,\n persona_id=persona.id,\n user_file_ids=[],\n )\n\n db_session.refresh(uf)\n assert uf.needs_persona_sync is True\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_pruning_hierarchy_nodes.py", "content": "\"\"\"\nExternal dependency unit tests for pruning hierarchy node extraction and DB persistence.\n\nVerifies that:\n1. extract_ids_from_runnable_connector correctly separates hierarchy nodes from doc IDs\n2. Extracted hierarchy nodes are correctly upserted to Postgres via upsert_hierarchy_nodes_batch\n3. Upserting is idempotent (running twice doesn't duplicate nodes)\n4. Document-to-hierarchy-node linkage is updated during pruning\n5. link_hierarchy_nodes_to_documents links nodes that are also documents\n6. HierarchyNodeByConnectorCredentialPair join table population and pruning\n7. Orphaned hierarchy node deletion and re-parenting\n\nUses a mock SlimConnectorWithPermSync that yields known hierarchy nodes and slim documents,\ncombined with a real PostgreSQL database for verifying persistence.\n\"\"\"\n\nfrom collections.abc import Iterator\nfrom typing import Any\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.access.models import ExternalAccess\nfrom onyx.background.celery.celery_utils import extract_ids_from_runnable_connector\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.interfaces import GenerateSlimDocumentOutput\nfrom onyx.connectors.interfaces import SecondsSinceUnixEpoch\nfrom onyx.connectors.interfaces import SlimConnectorWithPermSync\nfrom onyx.connectors.models import HierarchyNode as PydanticHierarchyNode\nfrom onyx.connectors.models import InputType\nfrom onyx.connectors.models import SlimDocument\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.db.hierarchy import delete_orphaned_hierarchy_nodes\nfrom onyx.db.hierarchy import ensure_source_node_exists\nfrom onyx.db.hierarchy import get_all_hierarchy_nodes_for_source\nfrom onyx.db.hierarchy import get_hierarchy_node_by_raw_id\nfrom onyx.db.hierarchy import link_hierarchy_nodes_to_documents\nfrom onyx.db.hierarchy import remove_stale_hierarchy_node_cc_pair_entries\nfrom onyx.db.hierarchy import reparent_orphaned_hierarchy_nodes\nfrom onyx.db.hierarchy import update_document_parent_hierarchy_nodes\nfrom onyx.db.hierarchy import upsert_hierarchy_node_cc_pair_entries\nfrom onyx.db.hierarchy import upsert_hierarchy_nodes_batch\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import Document as DbDocument\nfrom onyx.db.models import HierarchyNode as DBHierarchyNode\nfrom onyx.db.models import HierarchyNodeByConnectorCredentialPair\nfrom onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface\nfrom onyx.kg.models import KGStage\n\n# ---------------------------------------------------------------------------\n# Constants\n# ---------------------------------------------------------------------------\n\nTEST_SOURCE = DocumentSource.SLACK\n\nCHANNEL_A_ID = \"C_GENERAL\"\nCHANNEL_A_NAME = \"#general\"\nCHANNEL_B_ID = \"C_RANDOM\"\nCHANNEL_B_NAME = \"#random\"\nCHANNEL_C_ID = \"C_ENGINEERING\"\nCHANNEL_C_NAME = \"#engineering\"\n\nSLIM_DOC_IDS = [\"msg-001\", \"msg-002\", \"msg-003\"]\n\n\n# ---------------------------------------------------------------------------\n# Mock connector\n# ---------------------------------------------------------------------------\n\n\ndef _make_hierarchy_nodes() -> list[PydanticHierarchyNode]:\n \"\"\"Build a known set of hierarchy nodes resembling Slack channels.\"\"\"\n return [\n PydanticHierarchyNode(\n raw_node_id=CHANNEL_A_ID,\n raw_parent_id=None,\n display_name=CHANNEL_A_NAME,\n link=\"https://slack.example.com/channels/general\",\n node_type=HierarchyNodeType.CHANNEL,\n external_access=ExternalAccess(\n external_user_emails={\"alice@example.com\", \"bob@example.com\"},\n external_user_group_ids=set(),\n is_public=False,\n ),\n ),\n PydanticHierarchyNode(\n raw_node_id=CHANNEL_B_ID,\n raw_parent_id=None,\n display_name=CHANNEL_B_NAME,\n link=\"https://slack.example.com/channels/random\",\n node_type=HierarchyNodeType.CHANNEL,\n ),\n PydanticHierarchyNode(\n raw_node_id=CHANNEL_C_ID,\n raw_parent_id=None,\n display_name=CHANNEL_C_NAME,\n link=\"https://slack.example.com/channels/engineering\",\n node_type=HierarchyNodeType.CHANNEL,\n external_access=ExternalAccess(\n external_user_emails=set(),\n external_user_group_ids={\"eng-team\"},\n is_public=True,\n ),\n ),\n ]\n\n\nDOC_PARENT_MAP = {\n \"msg-001\": CHANNEL_A_ID,\n \"msg-002\": CHANNEL_A_ID,\n \"msg-003\": CHANNEL_B_ID,\n}\n\n\ndef _make_slim_docs() -> list[SlimDocument | PydanticHierarchyNode]:\n return [\n SlimDocument(id=doc_id, parent_hierarchy_raw_node_id=DOC_PARENT_MAP.get(doc_id))\n for doc_id in SLIM_DOC_IDS\n ]\n\n\nclass MockSlimConnectorWithPermSync(SlimConnectorWithPermSync):\n \"\"\"Yields a batch containing interleaved hierarchy nodes and slim docs.\"\"\"\n\n def load_credentials(\n self,\n credentials: dict[str, Any], # noqa: ARG002\n ) -> dict[str, Any] | None: # noqa: ARG002\n return None\n\n def retrieve_all_slim_docs_perm_sync(\n self,\n start: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n end: SecondsSinceUnixEpoch | None = None, # noqa: ARG002\n callback: IndexingHeartbeatInterface | None = None, # noqa: ARG002\n ) -> GenerateSlimDocumentOutput:\n return self._generate()\n\n def _generate(self) -> Iterator[list[SlimDocument | PydanticHierarchyNode]]:\n # First batch: hierarchy nodes + first slim doc\n batch_1: list[SlimDocument | PydanticHierarchyNode] = [\n *_make_hierarchy_nodes(),\n _make_slim_docs()[0],\n ]\n yield batch_1\n\n # Second batch: remaining slim docs only (no hierarchy nodes)\n yield _make_slim_docs()[1:]\n\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _create_cc_pair(\n db_session: Session,\n source: DocumentSource = TEST_SOURCE,\n) -> ConnectorCredentialPair:\n \"\"\"Create a real Connector + Credential + ConnectorCredentialPair for testing.\"\"\"\n connector = Connector(\n name=f\"Test {source.value} Connector\",\n source=source,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={},\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n source=source,\n credential_json={},\n admin_public=True,\n )\n db_session.add(credential)\n db_session.flush()\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=f\"Test {source.value} CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.PUBLIC,\n )\n db_session.add(cc_pair)\n db_session.commit()\n db_session.refresh(cc_pair)\n return cc_pair\n\n\ndef _cleanup_test_data(db_session: Session) -> None:\n \"\"\"Remove all test hierarchy nodes and documents to isolate tests.\"\"\"\n for doc_id in SLIM_DOC_IDS:\n db_session.query(DbDocument).filter(DbDocument.id == doc_id).delete()\n\n test_connector_ids_q = db_session.query(Connector.id).filter(\n Connector.source == TEST_SOURCE,\n Connector.name.like(\"Test %\"),\n )\n\n db_session.query(HierarchyNodeByConnectorCredentialPair).filter(\n HierarchyNodeByConnectorCredentialPair.connector_id.in_(test_connector_ids_q)\n ).delete(synchronize_session=\"fetch\")\n db_session.query(DBHierarchyNode).filter(\n DBHierarchyNode.source == TEST_SOURCE\n ).delete()\n db_session.flush()\n\n # Collect credential IDs before deleting cc_pairs (bulk query.delete()\n # bypasses ORM-level cascade, so credentials won't be auto-removed).\n credential_ids = [\n row[0]\n for row in db_session.query(ConnectorCredentialPair.credential_id)\n .filter(ConnectorCredentialPair.connector_id.in_(test_connector_ids_q))\n .all()\n ]\n\n db_session.query(ConnectorCredentialPair).filter(\n ConnectorCredentialPair.connector_id.in_(test_connector_ids_q)\n ).delete(synchronize_session=\"fetch\")\n db_session.query(Connector).filter(\n Connector.source == TEST_SOURCE,\n Connector.name.like(\"Test %\"),\n ).delete(synchronize_session=\"fetch\")\n if credential_ids:\n db_session.query(Credential).filter(Credential.id.in_(credential_ids)).delete(\n synchronize_session=\"fetch\"\n )\n db_session.commit()\n\n\ndef _create_test_documents(db_session: Session) -> list[DbDocument]:\n \"\"\"Insert minimal Document rows for our test doc IDs.\"\"\"\n docs = []\n for doc_id in SLIM_DOC_IDS:\n doc = DbDocument(\n id=doc_id,\n semantic_id=doc_id,\n kg_stage=KGStage.NOT_STARTED,\n )\n db_session.add(doc)\n docs.append(doc)\n db_session.commit()\n return docs\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\ndef test_pruning_extracts_hierarchy_nodes(\n db_session: Session, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"extract_ids_from_runnable_connector must separate hierarchy node IDs and\n document IDs into the correct buckets of the SlimConnectorExtractionResult.\"\"\"\n connector = MockSlimConnectorWithPermSync()\n\n result = extract_ids_from_runnable_connector(connector, callback=None)\n\n # raw_id_to_parent should contain ONLY document IDs, not hierarchy node IDs\n assert result.raw_id_to_parent.keys() == set(SLIM_DOC_IDS)\n\n # Hierarchy nodes should be the 3 channels\n assert len(result.hierarchy_nodes) == 3\n extracted_raw_ids = {n.raw_node_id for n in result.hierarchy_nodes}\n assert extracted_raw_ids == {CHANNEL_A_ID, CHANNEL_B_ID, CHANNEL_C_ID}\n\n\ndef test_pruning_upserts_hierarchy_nodes_to_db(db_session: Session) -> None:\n \"\"\"Full flow: extract hierarchy nodes from mock connector, upsert to Postgres,\n then verify the DB state (node count, parent relationships, permissions).\"\"\"\n _cleanup_test_data(db_session)\n\n # Step 1: ensure the SOURCE node exists (mirrors what the pruning task does)\n source_node = ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n\n # Step 2: extract from mock connector\n connector = MockSlimConnectorWithPermSync()\n result = extract_ids_from_runnable_connector(connector, callback=None)\n assert len(result.hierarchy_nodes) == 3\n\n # Step 3: upsert hierarchy nodes (public connector = False)\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=result.hierarchy_nodes,\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n assert len(upserted) == 3\n\n # Step 4: verify DB state\n all_nodes = get_all_hierarchy_nodes_for_source(db_session, TEST_SOURCE)\n # 3 channel nodes + 1 SOURCE node\n assert len(all_nodes) == 4\n\n # Verify each channel node\n channel_a = get_hierarchy_node_by_raw_id(db_session, CHANNEL_A_ID, TEST_SOURCE)\n assert channel_a is not None\n assert channel_a.display_name == CHANNEL_A_NAME\n assert channel_a.node_type == HierarchyNodeType.CHANNEL\n assert channel_a.link == \"https://slack.example.com/channels/general\"\n # Parent should be the SOURCE node (raw_parent_id was None)\n assert channel_a.parent_id == source_node.id\n # Permission fields for channel A (private, has user emails)\n assert channel_a.is_public is False\n assert channel_a.external_user_emails is not None\n assert set(channel_a.external_user_emails) == {\n \"alice@example.com\",\n \"bob@example.com\",\n }\n\n channel_b = get_hierarchy_node_by_raw_id(db_session, CHANNEL_B_ID, TEST_SOURCE)\n assert channel_b is not None\n assert channel_b.display_name == CHANNEL_B_NAME\n assert channel_b.parent_id == source_node.id\n # Channel B has no external_access -> defaults to not public, no emails/groups\n assert channel_b.is_public is False\n assert channel_b.external_user_emails is None\n assert channel_b.external_user_group_ids is None\n\n channel_c = get_hierarchy_node_by_raw_id(db_session, CHANNEL_C_ID, TEST_SOURCE)\n assert channel_c is not None\n assert channel_c.display_name == CHANNEL_C_NAME\n assert channel_c.parent_id == source_node.id\n # Channel C is public and has a group\n assert channel_c.is_public is True\n assert channel_c.external_user_group_ids is not None\n assert set(channel_c.external_user_group_ids) == {\"eng-team\"}\n\n\ndef test_pruning_upserts_hierarchy_nodes_public_connector(\n db_session: Session,\n) -> None:\n \"\"\"When the connector's access type is PUBLIC, all hierarchy nodes must be\n marked is_public=True regardless of their external_access settings.\"\"\"\n _cleanup_test_data(db_session)\n\n ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n\n connector = MockSlimConnectorWithPermSync()\n result = extract_ids_from_runnable_connector(connector, callback=None)\n\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=result.hierarchy_nodes,\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=True,\n )\n assert len(upserted) == 3\n\n # Every node should be public\n for node in upserted:\n assert node.is_public is True\n # Public connector forces emails/groups to None\n assert node.external_user_emails is None\n assert node.external_user_group_ids is None\n\n\ndef test_pruning_hierarchy_node_upsert_idempotency(db_session: Session) -> None:\n \"\"\"Upserting the same hierarchy nodes twice must not create duplicates.\n The second call should update existing rows in place.\"\"\"\n _cleanup_test_data(db_session)\n\n ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n\n nodes = _make_hierarchy_nodes()\n\n # First upsert\n first_result = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=nodes,\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n first_ids = {n.id for n in first_result}\n all_after_first = get_all_hierarchy_nodes_for_source(db_session, TEST_SOURCE)\n count_after_first = len(all_after_first)\n\n # Second upsert with the same nodes\n second_result = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=nodes,\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n second_ids = {n.id for n in second_result}\n all_after_second = get_all_hierarchy_nodes_for_source(db_session, TEST_SOURCE)\n count_after_second = len(all_after_second)\n\n # No new rows should have been created\n assert count_after_first == count_after_second\n # Same DB primary keys should have been returned\n assert first_ids == second_ids\n\n\ndef test_pruning_hierarchy_node_upsert_updates_fields(db_session: Session) -> None:\n \"\"\"Upserting a hierarchy node with changed fields should update the existing row.\"\"\"\n _cleanup_test_data(db_session)\n\n ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n\n original_node = PydanticHierarchyNode(\n raw_node_id=CHANNEL_A_ID,\n raw_parent_id=None,\n display_name=CHANNEL_A_NAME,\n link=\"https://slack.example.com/channels/general\",\n node_type=HierarchyNodeType.CHANNEL,\n )\n upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=[original_node],\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n\n # Now upsert again with updated display_name and permissions\n updated_node = PydanticHierarchyNode(\n raw_node_id=CHANNEL_A_ID,\n raw_parent_id=None,\n display_name=\"#general-renamed\",\n link=\"https://slack.example.com/channels/general-renamed\",\n node_type=HierarchyNodeType.CHANNEL,\n external_access=ExternalAccess(\n external_user_emails={\"new_user@example.com\"},\n external_user_group_ids=set(),\n is_public=True,\n ),\n )\n upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=[updated_node],\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n\n db_node = get_hierarchy_node_by_raw_id(db_session, CHANNEL_A_ID, TEST_SOURCE)\n assert db_node is not None\n assert db_node.display_name == \"#general-renamed\"\n assert db_node.link == \"https://slack.example.com/channels/general-renamed\"\n assert db_node.is_public is True\n assert db_node.external_user_emails is not None\n assert set(db_node.external_user_emails) == {\"new_user@example.com\"}\n\n\n# ---------------------------------------------------------------------------\n# Document-to-hierarchy-node linkage tests\n# ---------------------------------------------------------------------------\n\n\ndef test_extraction_preserves_parent_hierarchy_raw_node_id(\n db_session: Session, # noqa: ARG001\n) -> None:\n \"\"\"extract_ids_from_runnable_connector should carry the\n parent_hierarchy_raw_node_id from SlimDocument into the raw_id_to_parent dict.\"\"\"\n connector = MockSlimConnectorWithPermSync()\n result = extract_ids_from_runnable_connector(connector, callback=None)\n\n for doc_id, expected_parent in DOC_PARENT_MAP.items():\n assert (\n result.raw_id_to_parent[doc_id] == expected_parent\n ), f\"raw_id_to_parent[{doc_id}] should be {expected_parent}\"\n\n # Hierarchy node IDs should NOT be in raw_id_to_parent\n for channel_id in [CHANNEL_A_ID, CHANNEL_B_ID, CHANNEL_C_ID]:\n assert channel_id not in result.raw_id_to_parent\n\n\ndef test_update_document_parent_hierarchy_nodes(db_session: Session) -> None:\n \"\"\"update_document_parent_hierarchy_nodes should set\n Document.parent_hierarchy_node_id for each document in the mapping.\"\"\"\n _cleanup_test_data(db_session)\n\n source_node = ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=_make_hierarchy_nodes(),\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n node_id_by_raw = {n.raw_node_id: n.id for n in upserted}\n\n # Create documents with no parent set\n docs = _create_test_documents(db_session)\n for doc in docs:\n assert doc.parent_hierarchy_node_id is None\n\n # Build resolved map (same logic as _resolve_and_update_document_parents)\n resolved: dict[str, int | None] = {}\n for doc_id, raw_parent in DOC_PARENT_MAP.items():\n resolved[doc_id] = node_id_by_raw.get(raw_parent, source_node.id)\n\n updated = update_document_parent_hierarchy_nodes(\n db_session=db_session,\n doc_parent_map=resolved,\n commit=True,\n )\n assert updated == len(SLIM_DOC_IDS)\n\n # Verify each document now points to the correct hierarchy node\n db_session.expire_all()\n for doc_id, raw_parent in DOC_PARENT_MAP.items():\n tmp_doc = db_session.get(DbDocument, doc_id)\n assert tmp_doc is not None\n doc = tmp_doc\n expected_node_id = node_id_by_raw[raw_parent]\n assert (\n doc.parent_hierarchy_node_id == expected_node_id\n ), f\"Document {doc_id} should point to node for {raw_parent}\"\n\n\ndef test_update_document_parent_is_idempotent(db_session: Session) -> None:\n \"\"\"Running update_document_parent_hierarchy_nodes a second time with the\n same mapping should update zero rows.\"\"\"\n _cleanup_test_data(db_session)\n\n ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=_make_hierarchy_nodes(),\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n node_id_by_raw = {n.raw_node_id: n.id for n in upserted}\n _create_test_documents(db_session)\n\n resolved: dict[str, int | None] = {\n doc_id: node_id_by_raw[raw_parent]\n for doc_id, raw_parent in DOC_PARENT_MAP.items()\n }\n\n first_updated = update_document_parent_hierarchy_nodes(\n db_session=db_session,\n doc_parent_map=resolved,\n commit=True,\n )\n assert first_updated == len(SLIM_DOC_IDS)\n\n second_updated = update_document_parent_hierarchy_nodes(\n db_session=db_session,\n doc_parent_map=resolved,\n commit=True,\n )\n assert second_updated == 0\n\n\ndef test_link_hierarchy_nodes_to_documents_for_confluence(\n db_session: Session,\n) -> None:\n \"\"\"For sources in SOURCES_WITH_HIERARCHY_NODE_DOCUMENTS (e.g. Confluence),\n link_hierarchy_nodes_to_documents should set HierarchyNode.document_id\n when a hierarchy node's raw_node_id matches a document ID.\"\"\"\n _cleanup_test_data(db_session)\n confluence_source = DocumentSource.CONFLUENCE\n\n # Clean up any existing Confluence hierarchy nodes\n db_session.query(DBHierarchyNode).filter(\n DBHierarchyNode.source == confluence_source\n ).delete()\n db_session.commit()\n\n ensure_source_node_exists(db_session, confluence_source, commit=True)\n\n # Create a hierarchy node whose raw_node_id matches a document ID\n page_node_id = \"confluence-page-123\"\n nodes = [\n PydanticHierarchyNode(\n raw_node_id=page_node_id,\n raw_parent_id=None,\n display_name=\"Test Page\",\n link=\"https://wiki.example.com/page/123\",\n node_type=HierarchyNodeType.PAGE,\n ),\n ]\n upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=nodes,\n source=confluence_source,\n commit=True,\n is_connector_public=False,\n )\n\n # Verify the node exists but has no document_id yet\n db_node = get_hierarchy_node_by_raw_id(db_session, page_node_id, confluence_source)\n assert db_node is not None\n assert db_node.document_id is None\n\n # Create a document with the same ID as the hierarchy node\n doc = DbDocument(\n id=page_node_id,\n semantic_id=\"Test Page\",\n kg_stage=KGStage.NOT_STARTED,\n )\n db_session.add(doc)\n db_session.commit()\n\n # Link nodes to documents\n linked = link_hierarchy_nodes_to_documents(\n db_session=db_session,\n document_ids=[page_node_id],\n source=confluence_source,\n commit=True,\n )\n assert linked == 1\n\n # Verify the hierarchy node now has document_id set\n db_session.expire_all()\n db_node = get_hierarchy_node_by_raw_id(db_session, page_node_id, confluence_source)\n assert db_node is not None\n assert db_node.document_id == page_node_id\n\n # Cleanup\n db_session.query(DbDocument).filter(DbDocument.id == page_node_id).delete()\n db_session.query(DBHierarchyNode).filter(\n DBHierarchyNode.source == confluence_source\n ).delete()\n db_session.commit()\n\n\ndef test_link_hierarchy_nodes_skips_non_hierarchy_sources(\n db_session: Session,\n) -> None:\n \"\"\"link_hierarchy_nodes_to_documents should return 0 for sources that\n don't support hierarchy-node-as-document (e.g. Slack, Google Drive).\"\"\"\n linked = link_hierarchy_nodes_to_documents(\n db_session=db_session,\n document_ids=SLIM_DOC_IDS,\n source=TEST_SOURCE, # Slack — not in SOURCES_WITH_HIERARCHY_NODE_DOCUMENTS\n commit=False,\n )\n assert linked == 0\n\n\n# ---------------------------------------------------------------------------\n# Join table + pruning tests\n# ---------------------------------------------------------------------------\n\n\ndef test_upsert_hierarchy_node_cc_pair_entries(db_session: Session) -> None:\n \"\"\"upsert_hierarchy_node_cc_pair_entries should insert rows and be idempotent.\"\"\"\n _cleanup_test_data(db_session)\n ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n cc_pair = _create_cc_pair(db_session)\n\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=_make_hierarchy_nodes(),\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n node_ids = [n.id for n in upserted]\n\n # First call — should insert rows\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=node_ids,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n commit=True,\n )\n\n rows = (\n db_session.query(HierarchyNodeByConnectorCredentialPair)\n .filter(\n HierarchyNodeByConnectorCredentialPair.connector_id == cc_pair.connector_id,\n HierarchyNodeByConnectorCredentialPair.credential_id\n == cc_pair.credential_id,\n )\n .all()\n )\n assert len(rows) == 3\n\n # Second call — idempotent, same count\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=node_ids,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n commit=True,\n )\n rows_after = (\n db_session.query(HierarchyNodeByConnectorCredentialPair)\n .filter(\n HierarchyNodeByConnectorCredentialPair.connector_id == cc_pair.connector_id,\n HierarchyNodeByConnectorCredentialPair.credential_id\n == cc_pair.credential_id,\n )\n .all()\n )\n assert len(rows_after) == 3\n\n\ndef test_remove_stale_entries_and_delete_orphans(db_session: Session) -> None:\n \"\"\"After removing stale join-table entries, orphaned hierarchy nodes should\n be deleted and the SOURCE node should survive.\"\"\"\n _cleanup_test_data(db_session)\n source_node = ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n cc_pair = _create_cc_pair(db_session)\n\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=_make_hierarchy_nodes(),\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n all_ids = [n.id for n in upserted]\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=all_ids,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n commit=True,\n )\n\n # Now simulate a pruning run where only channel A survived\n channel_a = get_hierarchy_node_by_raw_id(db_session, CHANNEL_A_ID, TEST_SOURCE)\n assert channel_a is not None\n live_ids = {channel_a.id}\n\n stale_removed = remove_stale_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n live_hierarchy_node_ids=live_ids,\n commit=True,\n )\n assert stale_removed == 2\n\n # Delete orphaned nodes\n deleted_raw_ids = delete_orphaned_hierarchy_nodes(\n db_session=db_session,\n source=TEST_SOURCE,\n commit=True,\n )\n assert set(deleted_raw_ids) == {CHANNEL_B_ID, CHANNEL_C_ID}\n\n # Verify only channel A + SOURCE remain\n remaining = get_all_hierarchy_nodes_for_source(db_session, TEST_SOURCE)\n remaining_raw = {n.raw_node_id for n in remaining}\n assert remaining_raw == {CHANNEL_A_ID, source_node.raw_node_id}\n\n\ndef test_multi_cc_pair_prevents_premature_deletion(db_session: Session) -> None:\n \"\"\"A hierarchy node shared by two cc_pairs should NOT be deleted when only\n one cc_pair removes its association.\"\"\"\n _cleanup_test_data(db_session)\n ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n cc_pair_1 = _create_cc_pair(db_session)\n cc_pair_2 = _create_cc_pair(db_session)\n\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=_make_hierarchy_nodes(),\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n all_ids = [n.id for n in upserted]\n\n # cc_pair 1 owns all 3\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=all_ids,\n connector_id=cc_pair_1.connector_id,\n credential_id=cc_pair_1.credential_id,\n commit=True,\n )\n # cc_pair 2 also owns all 3\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=all_ids,\n connector_id=cc_pair_2.connector_id,\n credential_id=cc_pair_2.credential_id,\n commit=True,\n )\n\n # cc_pair 1 prunes — keeps none\n remove_stale_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n connector_id=cc_pair_1.connector_id,\n credential_id=cc_pair_1.credential_id,\n live_hierarchy_node_ids=set(),\n commit=True,\n )\n\n # Orphan deletion should find nothing because cc_pair 2 still references them\n deleted = delete_orphaned_hierarchy_nodes(\n db_session=db_session,\n source=TEST_SOURCE,\n commit=True,\n )\n assert deleted == []\n\n # All 3 nodes + SOURCE should still exist\n remaining = get_all_hierarchy_nodes_for_source(db_session, TEST_SOURCE)\n assert len(remaining) == 4\n\n\ndef test_reparent_orphaned_children(db_session: Session) -> None:\n \"\"\"After deleting a parent hierarchy node, its children should be\n re-parented to the SOURCE node.\"\"\"\n _cleanup_test_data(db_session)\n source_node = ensure_source_node_exists(db_session, TEST_SOURCE, commit=True)\n cc_pair = _create_cc_pair(db_session)\n\n # Create a parent node and a child node\n parent_node = PydanticHierarchyNode(\n raw_node_id=\"PARENT\",\n raw_parent_id=None,\n display_name=\"Parent\",\n node_type=HierarchyNodeType.CHANNEL,\n )\n child_node = PydanticHierarchyNode(\n raw_node_id=\"CHILD\",\n raw_parent_id=\"PARENT\",\n display_name=\"Child\",\n node_type=HierarchyNodeType.CHANNEL,\n )\n upserted = upsert_hierarchy_nodes_batch(\n db_session=db_session,\n nodes=[parent_node, child_node],\n source=TEST_SOURCE,\n commit=True,\n is_connector_public=False,\n )\n assert len(upserted) == 2\n\n parent_db = get_hierarchy_node_by_raw_id(db_session, \"PARENT\", TEST_SOURCE)\n child_db = get_hierarchy_node_by_raw_id(db_session, \"CHILD\", TEST_SOURCE)\n assert parent_db is not None and child_db is not None\n assert child_db.parent_id == parent_db.id\n\n # Associate only the child with a cc_pair (parent is orphaned)\n upsert_hierarchy_node_cc_pair_entries(\n db_session=db_session,\n hierarchy_node_ids=[child_db.id],\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n commit=True,\n )\n\n # Delete orphaned nodes (parent has no cc_pair entry)\n deleted = delete_orphaned_hierarchy_nodes(\n db_session=db_session,\n source=TEST_SOURCE,\n commit=True,\n )\n assert \"PARENT\" in deleted\n\n # Child should now have parent_id=NULL (SET NULL cascade)\n db_session.expire_all()\n child_db = get_hierarchy_node_by_raw_id(db_session, \"CHILD\", TEST_SOURCE)\n assert child_db is not None\n assert child_db.parent_id is None\n\n # Re-parent orphans to SOURCE\n reparented = reparent_orphaned_hierarchy_nodes(\n db_session=db_session,\n source=TEST_SOURCE,\n commit=True,\n )\n assert len(reparented) == 1\n\n db_session.expire_all()\n child_db = get_hierarchy_node_by_raw_id(db_session, \"CHILD\", TEST_SOURCE)\n assert child_db is not None\n assert child_db.parent_id == source_node.id\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_user_file_delete_queue.py", "content": "\"\"\"\nExternal dependency unit tests for user file delete queue protections.\n\nVerifies that the three mechanisms added to check_for_user_file_delete work\ncorrectly:\n\n1. Queue depth backpressure – when the broker queue exceeds\n USER_FILE_DELETE_MAX_QUEUE_DEPTH, no new tasks are enqueued.\n\n2. Per-file Redis guard key – if the guard key for a file already exists in\n Redis, that file is skipped even though it is still in DELETING status.\n\n3. Task expiry – every send_task call carries expires=\n CELERY_USER_FILE_DELETE_TASK_EXPIRES so that stale queued tasks are\n discarded by workers automatically.\n\nAlso verifies that delete_user_file_impl clears the guard key the moment\nit is picked up by a worker.\n\nUses real Redis (DB 0 via get_redis_client) and real PostgreSQL for UserFile\nrows. The Celery app is provided as a MagicMock injected via a PropertyMock\non the task class so no real broker is needed.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom unittest.mock import PropertyMock\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n _user_file_delete_lock_key,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n _user_file_delete_queued_key,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n check_for_user_file_delete,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n process_single_user_file_delete,\n)\nfrom onyx.configs.constants import CELERY_USER_FILE_DELETE_TASK_EXPIRES\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import USER_FILE_DELETE_MAX_QUEUE_DEPTH\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import UserFile\nfrom onyx.redis.redis_pool import get_redis_client\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n_PATCH_QUEUE_LEN = (\n \"onyx.background.celery.tasks.user_file_processing.tasks.celery_get_queue_length\"\n)\n\n\ndef _create_deleting_user_file(db_session: Session, user_id: object) -> UserFile:\n \"\"\"Insert a UserFile in DELETING status and return it.\"\"\"\n uf = UserFile(\n id=uuid4(),\n user_id=user_id,\n file_id=f\"test_file_{uuid4().hex[:8]}\",\n name=f\"test_{uuid4().hex[:8]}.txt\",\n file_type=\"text/plain\",\n status=UserFileStatus.DELETING,\n )\n db_session.add(uf)\n db_session.commit()\n db_session.refresh(uf)\n return uf\n\n\n@contextmanager\ndef _patch_task_app(task: Any, mock_app: MagicMock) -> Generator[None, None, None]:\n \"\"\"Patch the ``app`` property on *task*'s class so that ``self.app``\n inside the task function returns *mock_app*.\n\n With ``bind=True``, ``task.run`` is a bound method whose ``__self__`` is\n the actual task instance. We patch ``app`` on that instance's class\n (a unique Celery-generated Task subclass) so the mock is scoped to this\n task only.\n\n Also patches ``celery_get_broker_client`` so the mock app doesn't need\n a real broker URL.\n \"\"\"\n task_instance = task.run.__self__\n with (\n patch.object(\n type(task_instance),\n \"app\",\n new_callable=PropertyMock,\n return_value=mock_app,\n ),\n patch(\n \"onyx.background.celery.tasks.user_file_processing.tasks.celery_get_broker_client\",\n return_value=MagicMock(),\n ),\n ):\n yield\n\n\n# ---------------------------------------------------------------------------\n# Test classes\n# ---------------------------------------------------------------------------\n\n\nclass TestDeleteQueueDepthBackpressure:\n \"\"\"Protection 1: skip all enqueuing when the broker queue is too deep.\"\"\"\n\n def test_no_tasks_enqueued_when_queue_over_limit(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"When the queue depth exceeds the limit the beat cycle is skipped.\"\"\"\n user = create_test_user(db_session, \"del_bp_user\")\n _create_deleting_user_file(db_session, user.id)\n\n mock_app = MagicMock()\n\n with (\n _patch_task_app(check_for_user_file_delete, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=USER_FILE_DELETE_MAX_QUEUE_DEPTH + 1),\n ):\n check_for_user_file_delete.run(tenant_id=TEST_TENANT_ID)\n\n mock_app.send_task.assert_not_called()\n\n\nclass TestDeletePerFileGuardKey:\n \"\"\"Protection 2: per-file Redis guard key prevents duplicate enqueue.\"\"\"\n\n def test_guarded_file_not_re_enqueued(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file whose guard key is already set in Redis is skipped.\"\"\"\n user = create_test_user(db_session, \"del_guard_user\")\n uf = _create_deleting_user_file(db_session, user.id)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_delete_queued_key(uf.id)\n redis_client.setex(guard_key, CELERY_USER_FILE_DELETE_TASK_EXPIRES, 1)\n\n mock_app = MagicMock()\n\n try:\n with (\n _patch_task_app(check_for_user_file_delete, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=0),\n ):\n check_for_user_file_delete.run(tenant_id=TEST_TENANT_ID)\n\n # send_task must not have been called with this specific file's ID\n for call in mock_app.send_task.call_args_list:\n kwargs = call.kwargs.get(\"kwargs\", {})\n assert kwargs.get(\"user_file_id\") != str(\n uf.id\n ), f\"File {uf.id} should have been skipped because its guard key exists\"\n finally:\n redis_client.delete(guard_key)\n\n def test_guard_key_exists_in_redis_after_enqueue(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"After a file is enqueued its guard key is present in Redis with a TTL.\"\"\"\n user = create_test_user(db_session, \"del_guard_set_user\")\n uf = _create_deleting_user_file(db_session, user.id)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_delete_queued_key(uf.id)\n redis_client.delete(guard_key) # clean slate\n\n mock_app = MagicMock()\n\n try:\n with (\n _patch_task_app(check_for_user_file_delete, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=0),\n ):\n check_for_user_file_delete.run(tenant_id=TEST_TENANT_ID)\n\n assert redis_client.exists(\n guard_key\n ), \"Guard key should be set in Redis after enqueue\"\n ttl = int(redis_client.ttl(guard_key)) # type: ignore[arg-type]\n assert (\n 0 < ttl <= CELERY_USER_FILE_DELETE_TASK_EXPIRES\n ), f\"Guard key TTL {ttl}s is outside the expected range (0, {CELERY_USER_FILE_DELETE_TASK_EXPIRES}]\"\n finally:\n redis_client.delete(guard_key)\n\n\nclass TestDeleteTaskExpiry:\n \"\"\"Protection 3: every send_task call includes an expires value.\"\"\"\n\n def test_send_task_called_with_expires(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"send_task is called with the correct queue, task name, and expires.\"\"\"\n user = create_test_user(db_session, \"del_expires_user\")\n uf = _create_deleting_user_file(db_session, user.id)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_delete_queued_key(uf.id)\n redis_client.delete(guard_key)\n\n mock_app = MagicMock()\n\n try:\n with (\n _patch_task_app(check_for_user_file_delete, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=0),\n ):\n check_for_user_file_delete.run(tenant_id=TEST_TENANT_ID)\n\n # At least one task should have been submitted (for our file)\n assert (\n mock_app.send_task.call_count >= 1\n ), \"Expected at least one task to be submitted\"\n\n # Every submitted task must carry expires\n for call in mock_app.send_task.call_args_list:\n assert call.args[0] == OnyxCeleryTask.DELETE_SINGLE_USER_FILE\n assert call.kwargs.get(\"queue\") == OnyxCeleryQueues.USER_FILE_DELETE\n assert (\n call.kwargs.get(\"expires\") == CELERY_USER_FILE_DELETE_TASK_EXPIRES\n ), \"Task must be submitted with the correct expires value to prevent stale task accumulation\"\n finally:\n redis_client.delete(guard_key)\n\n\nclass TestDeleteWorkerClearsGuardKey:\n \"\"\"process_single_user_file_delete removes the guard key when it picks up a task.\"\"\"\n\n def test_guard_key_deleted_on_pickup(\n self,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"The guard key is deleted before the worker does any real work.\n\n We simulate an already-locked file so delete_user_file_impl returns\n early – but crucially, after the guard key deletion.\n \"\"\"\n user_file_id = str(uuid4())\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_delete_queued_key(user_file_id)\n\n # Simulate the guard key set when the beat enqueued the task\n redis_client.setex(guard_key, CELERY_USER_FILE_DELETE_TASK_EXPIRES, 1)\n assert redis_client.exists(guard_key), \"Guard key must exist before pickup\"\n\n # Hold the per-file delete lock so the worker exits early without\n # touching the database or file store.\n lock_key = _user_file_delete_lock_key(user_file_id)\n delete_lock = redis_client.lock(lock_key, timeout=10)\n acquired = delete_lock.acquire(blocking=False)\n assert acquired, \"Should be able to acquire the delete lock for this test\"\n\n try:\n process_single_user_file_delete.run(\n user_file_id=user_file_id,\n tenant_id=TEST_TENANT_ID,\n )\n finally:\n if delete_lock.owned():\n delete_lock.release()\n\n assert not redis_client.exists(\n guard_key\n ), \"Guard key should be deleted when the worker picks up the task\"\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_user_file_indexing_adapter.py", "content": "\"\"\"\nExternal dependency unit tests for UserFileIndexingAdapter metadata writing.\n\nValidates that prepare_enrichment produces DocMetadataAwareIndexChunk\nobjects with both `user_project` and `personas` fields populated correctly\nbased on actual DB associations.\n\nUses real PostgreSQL for UserFile/Persona/UserProject rows.\nMocks the LLM tokenizer and file store since they are not relevant here.\n\"\"\"\n\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import TextSection\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__UserFile\nfrom onyx.db.models import Project__UserFile\nfrom onyx.db.models import User\nfrom onyx.db.models import UserFile\nfrom onyx.db.models import UserProject\nfrom onyx.indexing.adapters.user_file_indexing_adapter import UserFileIndexingAdapter\nfrom onyx.indexing.indexing_pipeline import DocumentBatchPrepareContext\nfrom onyx.indexing.models import ChunkEmbedding\nfrom onyx.indexing.models import IndexChunk\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _create_user_file(db_session: Session, user: User) -> UserFile:\n uf = UserFile(\n id=uuid4(),\n user_id=user.id,\n file_id=f\"test_file_{uuid4().hex[:8]}\",\n name=f\"test_{uuid4().hex[:8]}.txt\",\n file_type=\"text/plain\",\n status=UserFileStatus.COMPLETED,\n chunk_count=1,\n )\n db_session.add(uf)\n db_session.commit()\n db_session.refresh(uf)\n return uf\n\n\ndef _create_persona(db_session: Session, user: User) -> Persona:\n persona = Persona(\n name=f\"Test Persona {uuid4().hex[:8]}\",\n description=\"Test persona\",\n system_prompt=\"test\",\n task_prompt=\"test\",\n tools=[],\n document_sets=[],\n users=[user],\n groups=[],\n is_listed=True,\n is_public=True,\n display_priority=None,\n starter_messages=None,\n deleted=False,\n user_id=user.id,\n )\n db_session.add(persona)\n db_session.commit()\n db_session.refresh(persona)\n return persona\n\n\ndef _create_project(db_session: Session, user: User) -> UserProject:\n project = UserProject(\n user_id=user.id,\n name=f\"project-{uuid4().hex[:8]}\",\n instructions=\"\",\n )\n db_session.add(project)\n db_session.commit()\n db_session.refresh(project)\n return project\n\n\ndef _make_index_chunk(user_file: UserFile) -> IndexChunk:\n \"\"\"Build a minimal IndexChunk whose source document ID matches the UserFile.\"\"\"\n doc = Document(\n id=str(user_file.id),\n source=DocumentSource.USER_FILE,\n semantic_identifier=user_file.name,\n sections=[TextSection(text=\"test chunk content\", link=None)],\n metadata={},\n )\n return IndexChunk(\n source_document=doc,\n chunk_id=0,\n blurb=\"test chunk\",\n content=\"test chunk content\",\n source_links={0: \"\"},\n image_file_id=None,\n section_continuation=False,\n title_prefix=\"\",\n metadata_suffix_semantic=\"\",\n metadata_suffix_keyword=\"\",\n contextual_rag_reserved_tokens=0,\n doc_summary=\"\",\n chunk_context=\"\",\n mini_chunk_texts=None,\n large_chunk_id=None,\n embeddings=ChunkEmbedding(\n full_embedding=[0.0] * 768,\n mini_chunk_embeddings=[],\n ),\n title_embedding=None,\n )\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestAdapterWritesBothMetadataFields:\n \"\"\"prepare_enrichment must populate user_project AND personas.\"\"\"\n\n @patch(\n \"onyx.indexing.adapters.user_file_indexing_adapter.get_default_llm\",\n side_effect=Exception(\"no LLM in test\"),\n )\n def test_file_linked_to_persona_gets_persona_id(\n self,\n _mock_llm: MagicMock,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n user = create_test_user(db_session, \"adapter_persona\")\n uf = _create_user_file(db_session, user)\n persona = _create_persona(db_session, user)\n\n db_session.add(Persona__UserFile(persona_id=persona.id, user_file_id=uf.id))\n db_session.commit()\n\n adapter = UserFileIndexingAdapter(\n tenant_id=TEST_TENANT_ID, db_session=db_session\n )\n chunk = _make_index_chunk(uf)\n doc = chunk.source_document\n context = DocumentBatchPrepareContext(updatable_docs=[doc], id_to_boost_map={})\n\n enricher = adapter.prepare_enrichment(\n context=context,\n tenant_id=TEST_TENANT_ID,\n chunks=[chunk],\n )\n aware_chunk = enricher.enrich_chunk(chunk, 1.0)\n\n assert persona.id in aware_chunk.personas\n assert aware_chunk.user_project == []\n\n @patch(\n \"onyx.indexing.adapters.user_file_indexing_adapter.get_default_llm\",\n side_effect=Exception(\"no LLM in test\"),\n )\n def test_file_linked_to_project_gets_project_id(\n self,\n _mock_llm: MagicMock,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n user = create_test_user(db_session, \"adapter_project\")\n uf = _create_user_file(db_session, user)\n project = _create_project(db_session, user)\n\n db_session.add(Project__UserFile(project_id=project.id, user_file_id=uf.id))\n db_session.commit()\n\n adapter = UserFileIndexingAdapter(\n tenant_id=TEST_TENANT_ID, db_session=db_session\n )\n chunk = _make_index_chunk(uf)\n context = DocumentBatchPrepareContext(\n updatable_docs=[chunk.source_document], id_to_boost_map={}\n )\n\n enricher = adapter.prepare_enrichment(\n context=context,\n tenant_id=TEST_TENANT_ID,\n chunks=[chunk],\n )\n aware_chunk = enricher.enrich_chunk(chunk, 1.0)\n\n assert project.id in aware_chunk.user_project\n assert aware_chunk.personas == []\n\n @patch(\n \"onyx.indexing.adapters.user_file_indexing_adapter.get_default_llm\",\n side_effect=Exception(\"no LLM in test\"),\n )\n def test_file_linked_to_both_gets_both_ids(\n self,\n _mock_llm: MagicMock,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n user = create_test_user(db_session, \"adapter_both\")\n uf = _create_user_file(db_session, user)\n persona = _create_persona(db_session, user)\n project = _create_project(db_session, user)\n\n db_session.add(Persona__UserFile(persona_id=persona.id, user_file_id=uf.id))\n db_session.add(Project__UserFile(project_id=project.id, user_file_id=uf.id))\n db_session.commit()\n\n adapter = UserFileIndexingAdapter(\n tenant_id=TEST_TENANT_ID, db_session=db_session\n )\n chunk = _make_index_chunk(uf)\n context = DocumentBatchPrepareContext(\n updatable_docs=[chunk.source_document], id_to_boost_map={}\n )\n\n enricher = adapter.prepare_enrichment(\n context=context,\n tenant_id=TEST_TENANT_ID,\n chunks=[chunk],\n )\n aware_chunk = enricher.enrich_chunk(chunk, 1.0)\n\n assert persona.id in aware_chunk.personas\n assert project.id in aware_chunk.user_project\n\n @patch(\n \"onyx.indexing.adapters.user_file_indexing_adapter.get_default_llm\",\n side_effect=Exception(\"no LLM in test\"),\n )\n def test_file_with_no_associations_gets_empty_lists(\n self,\n _mock_llm: MagicMock,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n user = create_test_user(db_session, \"adapter_empty\")\n uf = _create_user_file(db_session, user)\n\n adapter = UserFileIndexingAdapter(\n tenant_id=TEST_TENANT_ID, db_session=db_session\n )\n chunk = _make_index_chunk(uf)\n context = DocumentBatchPrepareContext(\n updatable_docs=[chunk.source_document], id_to_boost_map={}\n )\n\n enricher = adapter.prepare_enrichment(\n context=context,\n tenant_id=TEST_TENANT_ID,\n chunks=[chunk],\n )\n aware_chunk = enricher.enrich_chunk(chunk, 1.0)\n\n assert aware_chunk.personas == []\n assert aware_chunk.user_project == []\n\n @patch(\n \"onyx.indexing.adapters.user_file_indexing_adapter.get_default_llm\",\n side_effect=Exception(\"no LLM in test\"),\n )\n def test_multiple_personas_all_appear(\n self,\n _mock_llm: MagicMock,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file linked to multiple personas should have all their IDs.\"\"\"\n user = create_test_user(db_session, \"adapter_multi\")\n uf = _create_user_file(db_session, user)\n persona_a = _create_persona(db_session, user)\n persona_b = _create_persona(db_session, user)\n\n db_session.add(Persona__UserFile(persona_id=persona_a.id, user_file_id=uf.id))\n db_session.add(Persona__UserFile(persona_id=persona_b.id, user_file_id=uf.id))\n db_session.commit()\n\n adapter = UserFileIndexingAdapter(\n tenant_id=TEST_TENANT_ID, db_session=db_session\n )\n chunk = _make_index_chunk(uf)\n context = DocumentBatchPrepareContext(\n updatable_docs=[chunk.source_document], id_to_boost_map={}\n )\n\n enricher = adapter.prepare_enrichment(\n context=context,\n tenant_id=TEST_TENANT_ID,\n chunks=[chunk],\n )\n aware_chunk = enricher.enrich_chunk(chunk, 1.0)\n\n assert set(aware_chunk.personas) == {persona_a.id, persona_b.id}\n" }, { "path": "backend/tests/external_dependency_unit/celery/test_user_file_processing_queue.py", "content": "\"\"\"\nExternal dependency unit tests for user file processing queue protections.\n\nVerifies that the three mechanisms added to check_user_file_processing work\ncorrectly:\n\n1. Queue depth backpressure – when the broker queue exceeds\n USER_FILE_PROCESSING_MAX_QUEUE_DEPTH, no new tasks are enqueued.\n\n2. Per-file Redis guard key – if the guard key for a file already exists in\n Redis, that file is skipped even though it is still in PROCESSING status.\n\n3. Task expiry – every send_task call carries expires=\n CELERY_USER_FILE_PROCESSING_TASK_EXPIRES so that stale queued tasks are\n discarded by workers automatically.\n\nAlso verifies that process_single_user_file clears the guard key the moment\nit is picked up by a worker.\n\nUses real Redis (DB 0 via get_redis_client) and real PostgreSQL for UserFile\nrows. The Celery app is provided as a MagicMock injected via a PropertyMock\non the task class so no real broker is needed.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom unittest.mock import PropertyMock\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n _user_file_lock_key,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n _user_file_queued_key,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n check_user_file_processing,\n)\nfrom onyx.background.celery.tasks.user_file_processing.tasks import (\n process_single_user_file,\n)\nfrom onyx.configs.constants import CELERY_USER_FILE_PROCESSING_TASK_EXPIRES\nfrom onyx.configs.constants import OnyxCeleryQueues\nfrom onyx.configs.constants import OnyxCeleryTask\nfrom onyx.configs.constants import USER_FILE_PROCESSING_MAX_QUEUE_DEPTH\nfrom onyx.db.enums import UserFileStatus\nfrom onyx.db.models import UserFile\nfrom onyx.redis.redis_pool import get_redis_client\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n_PATCH_QUEUE_LEN = (\n \"onyx.background.celery.tasks.user_file_processing.tasks.celery_get_queue_length\"\n)\n\n\ndef _create_processing_user_file(db_session: Session, user_id: object) -> UserFile:\n \"\"\"Insert a UserFile in PROCESSING status and return it.\"\"\"\n uf = UserFile(\n id=uuid4(),\n user_id=user_id,\n file_id=f\"test_file_{uuid4().hex[:8]}\",\n name=f\"test_{uuid4().hex[:8]}.txt\",\n file_type=\"text/plain\",\n status=UserFileStatus.PROCESSING,\n )\n db_session.add(uf)\n db_session.commit()\n db_session.refresh(uf)\n return uf\n\n\n@contextmanager\ndef _patch_task_app(task: Any, mock_app: MagicMock) -> Generator[None, None, None]:\n \"\"\"Patch the ``app`` property on *task*'s class so that ``self.app``\n inside the task function returns *mock_app*.\n\n With ``bind=True``, ``task.run`` is a bound method whose ``__self__`` is\n the actual task instance. We patch ``app`` on that instance's class\n (a unique Celery-generated Task subclass) so the mock is scoped to this\n task only.\n \"\"\"\n task_instance = task.run.__self__\n with (\n patch.object(\n type(task_instance),\n \"app\",\n new_callable=PropertyMock,\n return_value=mock_app,\n ),\n patch(\n \"onyx.background.celery.tasks.user_file_processing.tasks.celery_get_broker_client\",\n return_value=MagicMock(),\n ),\n ):\n yield\n\n\n# ---------------------------------------------------------------------------\n# Test classes\n# ---------------------------------------------------------------------------\n\n\nclass TestQueueDepthBackpressure:\n \"\"\"Protection 1: skip all enqueuing when the broker queue is too deep.\"\"\"\n\n def test_no_tasks_enqueued_when_queue_over_limit(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"When the queue depth exceeds the limit the beat cycle is skipped.\"\"\"\n user = create_test_user(db_session, \"bp_user\")\n _create_processing_user_file(db_session, user.id)\n\n mock_app = MagicMock()\n\n with (\n _patch_task_app(check_user_file_processing, mock_app),\n patch(\n _PATCH_QUEUE_LEN, return_value=USER_FILE_PROCESSING_MAX_QUEUE_DEPTH + 1\n ),\n ):\n check_user_file_processing.run(tenant_id=TEST_TENANT_ID)\n\n mock_app.send_task.assert_not_called()\n\n\nclass TestPerFileGuardKey:\n \"\"\"Protection 2: per-file Redis guard key prevents duplicate enqueue.\"\"\"\n\n def test_guarded_file_not_re_enqueued(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A file whose guard key is already set in Redis is skipped.\"\"\"\n user = create_test_user(db_session, \"guard_user\")\n uf = _create_processing_user_file(db_session, user.id)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_queued_key(uf.id)\n redis_client.setex(guard_key, CELERY_USER_FILE_PROCESSING_TASK_EXPIRES, 1)\n\n mock_app = MagicMock()\n\n try:\n with (\n _patch_task_app(check_user_file_processing, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=0),\n ):\n check_user_file_processing.run(tenant_id=TEST_TENANT_ID)\n\n # send_task must not have been called with this specific file's ID\n for call in mock_app.send_task.call_args_list:\n kwargs = call.kwargs.get(\"kwargs\", {})\n assert kwargs.get(\"user_file_id\") != str(\n uf.id\n ), f\"File {uf.id} should have been skipped because its guard key exists\"\n finally:\n redis_client.delete(guard_key)\n\n def test_guard_key_exists_in_redis_after_enqueue(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"After a file is enqueued its guard key is present in Redis with a TTL.\"\"\"\n user = create_test_user(db_session, \"guard_set_user\")\n uf = _create_processing_user_file(db_session, user.id)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_queued_key(uf.id)\n redis_client.delete(guard_key) # clean slate\n\n mock_app = MagicMock()\n\n try:\n with (\n _patch_task_app(check_user_file_processing, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=0),\n ):\n check_user_file_processing.run(tenant_id=TEST_TENANT_ID)\n\n assert redis_client.exists(\n guard_key\n ), \"Guard key should be set in Redis after enqueue\"\n ttl = int(redis_client.ttl(guard_key)) # type: ignore[arg-type]\n assert (\n 0 < ttl <= CELERY_USER_FILE_PROCESSING_TASK_EXPIRES\n ), f\"Guard key TTL {ttl}s is outside the expected range (0, {CELERY_USER_FILE_PROCESSING_TASK_EXPIRES}]\"\n finally:\n redis_client.delete(guard_key)\n\n\nclass TestTaskExpiry:\n \"\"\"Protection 3: every send_task call includes an expires value.\"\"\"\n\n def test_send_task_called_with_expires(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"send_task is called with the correct queue, task name, and expires.\"\"\"\n user = create_test_user(db_session, \"expires_user\")\n uf = _create_processing_user_file(db_session, user.id)\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_queued_key(uf.id)\n redis_client.delete(guard_key)\n\n mock_app = MagicMock()\n\n try:\n with (\n _patch_task_app(check_user_file_processing, mock_app),\n patch(_PATCH_QUEUE_LEN, return_value=0),\n ):\n check_user_file_processing.run(tenant_id=TEST_TENANT_ID)\n\n # At least one task should have been submitted (for our file)\n assert (\n mock_app.send_task.call_count >= 1\n ), \"Expected at least one task to be submitted\"\n\n # Every submitted task must carry expires\n for call in mock_app.send_task.call_args_list:\n assert call.args[0] == OnyxCeleryTask.PROCESS_SINGLE_USER_FILE\n assert call.kwargs.get(\"queue\") == OnyxCeleryQueues.USER_FILE_PROCESSING\n assert (\n call.kwargs.get(\"expires\")\n == CELERY_USER_FILE_PROCESSING_TASK_EXPIRES\n ), \"Task must be submitted with the correct expires value to prevent stale task accumulation\"\n finally:\n redis_client.delete(guard_key)\n\n\nclass TestWorkerClearsGuardKey:\n \"\"\"process_single_user_file removes the guard key when it picks up a task.\"\"\"\n\n def test_guard_key_deleted_on_pickup(\n self,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"The guard key is deleted before the worker does any real work.\n\n We simulate an already-locked file so process_single_user_file returns\n early – but crucially, after the guard key deletion.\n \"\"\"\n user_file_id = str(uuid4())\n\n redis_client = get_redis_client(tenant_id=TEST_TENANT_ID)\n guard_key = _user_file_queued_key(user_file_id)\n\n # Simulate the guard key set when the beat enqueued the task\n redis_client.setex(guard_key, CELERY_USER_FILE_PROCESSING_TASK_EXPIRES, 1)\n assert redis_client.exists(guard_key), \"Guard key must exist before pickup\"\n\n # Hold the per-file processing lock so the worker exits early without\n # touching the database or file store.\n lock_key = _user_file_lock_key(user_file_id)\n processing_lock = redis_client.lock(lock_key, timeout=10)\n acquired = processing_lock.acquire(blocking=False)\n assert acquired, \"Should be able to acquire the processing lock for this test\"\n\n try:\n process_single_user_file.run(\n user_file_id=user_file_id,\n tenant_id=TEST_TENANT_ID,\n )\n finally:\n if processing_lock.owned():\n processing_lock.release()\n\n assert not redis_client.exists(\n guard_key\n ), \"Guard key should be deleted when the worker picks up the task\"\n" }, { "path": "backend/tests/external_dependency_unit/chat/test_user_reminder_message_type.py", "content": "\"\"\"\nTests for the USER_REMINDER message type handling in translate_history_to_llm_format.\n\nThese tests verify that:\n1. USER_REMINDER messages are wrapped with tags\n2. The wrapped messages are converted to UserMessage type for the LLM\n3. The tags are properly applied around the message content\n4. CODE_BLOCK_MARKDOWN is prepended to system messages for models that need it\n\"\"\"\n\nimport pytest\n\nfrom onyx.chat.llm_step import translate_history_to_llm_format\nfrom onyx.chat.models import ChatMessageSimple\nfrom onyx.configs.constants import MessageType\nfrom onyx.llm.interfaces import LLMConfig\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import SystemMessage\nfrom onyx.llm.models import UserMessage\nfrom onyx.prompts.chat_prompts import CODE_BLOCK_MARKDOWN\nfrom onyx.prompts.constants import SYSTEM_REMINDER_TAG_CLOSE\nfrom onyx.prompts.constants import SYSTEM_REMINDER_TAG_OPEN\n\n\ndef _ensure_list(\n result: list[ChatCompletionMessage] | ChatCompletionMessage,\n) -> list[ChatCompletionMessage]:\n \"\"\"Convert LanguageModelInput to a list for easier testing.\"\"\"\n if isinstance(result, list):\n return result\n return [result]\n\n\n@pytest.fixture\ndef mock_llm_config() -> LLMConfig:\n \"\"\"Create a minimal LLMConfig for testing.\"\"\"\n return LLMConfig(\n model_provider=\"openai\",\n model_name=\"gpt-4o-mini\",\n temperature=0.7,\n api_key=\"test-key\",\n api_base=None,\n api_version=None,\n max_input_tokens=128000,\n )\n\n\nclass TestUserReminderMessageType:\n \"\"\"Tests for USER_REMINDER message handling in translate_history_to_llm_format.\"\"\"\n\n def test_user_reminder_wrapped_with_tags(self, mock_llm_config: LLMConfig) -> None:\n \"\"\"Test that USER_REMINDER messages are wrapped with system-reminder tags.\"\"\"\n reminder_text = \"Remember to cite your sources.\"\n history = [\n ChatMessageSimple(\n message=reminder_text,\n token_count=10,\n message_type=MessageType.USER_REMINDER,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, mock_llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, UserMessage)\n assert msg.role == \"user\"\n # Verify the content starts and ends with the proper tags\n assert isinstance(msg.content, str)\n assert msg.content.startswith(SYSTEM_REMINDER_TAG_OPEN)\n assert msg.content.endswith(SYSTEM_REMINDER_TAG_CLOSE)\n # Verify the original message is inside the tags\n assert reminder_text in msg.content\n\n def test_user_reminder_tag_format(self, mock_llm_config: LLMConfig) -> None:\n \"\"\"Test the exact format of the system-reminder tag wrapping.\"\"\"\n reminder_text = \"This is a test reminder.\"\n history = [\n ChatMessageSimple(\n message=reminder_text,\n token_count=10,\n message_type=MessageType.USER_REMINDER,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, mock_llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, UserMessage)\n expected_content = (\n f\"{SYSTEM_REMINDER_TAG_OPEN}\\n{reminder_text}\\n{SYSTEM_REMINDER_TAG_CLOSE}\"\n )\n assert msg.content == expected_content\n\n def test_user_reminder_converted_to_user_message(\n self, mock_llm_config: LLMConfig\n ) -> None:\n \"\"\"Test that USER_REMINDER is converted to UserMessage (not a different type).\"\"\"\n history = [\n ChatMessageSimple(\n message=\"Test reminder\",\n token_count=5,\n message_type=MessageType.USER_REMINDER,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, mock_llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n # Should be a UserMessage since LLM APIs don't have a native reminder type\n assert isinstance(result[0], UserMessage)\n assert result[0].role == \"user\"\n\n def test_user_reminder_in_mixed_history(self, mock_llm_config: LLMConfig) -> None:\n \"\"\"Test USER_REMINDER handling when mixed with other message types.\"\"\"\n history = [\n ChatMessageSimple(\n message=\"You are a helpful assistant.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n ),\n ChatMessageSimple(\n message=\"Hello!\",\n token_count=5,\n message_type=MessageType.USER,\n ),\n ChatMessageSimple(\n message=\"Hi there! How can I help?\",\n token_count=10,\n message_type=MessageType.ASSISTANT,\n ),\n ChatMessageSimple(\n message=\"Remember to be concise.\",\n token_count=8,\n message_type=MessageType.USER_REMINDER,\n ),\n ]\n\n raw_result = translate_history_to_llm_format(history, mock_llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 4\n # Check the reminder message (last one)\n reminder_msg = result[3]\n assert isinstance(reminder_msg, UserMessage)\n assert isinstance(reminder_msg.content, str)\n assert reminder_msg.content.startswith(SYSTEM_REMINDER_TAG_OPEN)\n assert reminder_msg.content.endswith(SYSTEM_REMINDER_TAG_CLOSE)\n assert \"Remember to be concise.\" in reminder_msg.content\n\n # Check that regular USER message is NOT wrapped\n user_msg = result[1]\n assert isinstance(user_msg, UserMessage)\n assert user_msg.content == \"Hello!\" # No tags\n\n def test_regular_user_message_not_wrapped(self, mock_llm_config: LLMConfig) -> None:\n \"\"\"Test that regular USER messages are NOT wrapped with system-reminder tags.\"\"\"\n history = [\n ChatMessageSimple(\n message=\"This is a normal user message.\",\n token_count=10,\n message_type=MessageType.USER,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, mock_llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, UserMessage)\n # Regular user message should NOT have the tags\n assert isinstance(msg.content, str)\n assert SYSTEM_REMINDER_TAG_OPEN not in msg.content\n assert SYSTEM_REMINDER_TAG_CLOSE not in msg.content\n assert msg.content == \"This is a normal user message.\"\n\n\ndef _create_llm_config(model_name: str) -> LLMConfig:\n \"\"\"Create a LLMConfig with the specified model name.\"\"\"\n return LLMConfig(\n model_provider=\"openai\",\n model_name=model_name,\n temperature=0.7,\n api_key=\"test-key\",\n api_base=None,\n api_version=None,\n max_input_tokens=128000,\n )\n\n\nclass TestCodeBlockMarkdownFormatting:\n \"\"\"Tests for CODE_BLOCK_MARKDOWN prefix handling in translate_history_to_llm_format.\n\n OpenAI reasoning models (o1, o3, gpt-5) need a \"Formatting re-enabled. \" prefix\n in their system messages for correct markdown generation.\n \"\"\"\n\n def test_o1_model_prepends_markdown_to_string(self) -> None:\n \"\"\"Test that o1 model prepends CODE_BLOCK_MARKDOWN to string system message.\"\"\"\n llm_config = _create_llm_config(\"o1\")\n history = [\n ChatMessageSimple(\n message=\"You are a helpful assistant.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, SystemMessage)\n assert isinstance(msg.content, str)\n assert msg.content == CODE_BLOCK_MARKDOWN + \"You are a helpful assistant.\"\n\n def test_o3_model_prepends_markdown(self) -> None:\n \"\"\"Test that o3 model prepends CODE_BLOCK_MARKDOWN to system message.\"\"\"\n llm_config = _create_llm_config(\"o3-mini\")\n history = [\n ChatMessageSimple(\n message=\"System prompt here.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, SystemMessage)\n assert isinstance(msg.content, str)\n assert msg.content.startswith(CODE_BLOCK_MARKDOWN)\n\n def test_gpt5_model_prepends_markdown(self) -> None:\n \"\"\"Test that gpt-5 model prepends CODE_BLOCK_MARKDOWN to system message.\"\"\"\n llm_config = _create_llm_config(\"gpt-5\")\n history = [\n ChatMessageSimple(\n message=\"System prompt here.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, SystemMessage)\n assert isinstance(msg.content, str)\n assert msg.content.startswith(CODE_BLOCK_MARKDOWN)\n\n def test_gpt4o_does_not_prepend(self) -> None:\n \"\"\"Test that gpt-4o model does NOT prepend CODE_BLOCK_MARKDOWN.\"\"\"\n llm_config = _create_llm_config(\"gpt-4o\")\n history = [\n ChatMessageSimple(\n message=\"You are a helpful assistant.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, SystemMessage)\n assert isinstance(msg.content, str)\n # Should NOT have the prefix\n assert msg.content == \"You are a helpful assistant.\"\n assert not msg.content.startswith(CODE_BLOCK_MARKDOWN)\n\n def test_no_system_message_no_crash(self) -> None:\n \"\"\"Test that history without system message doesn't crash.\"\"\"\n llm_config = _create_llm_config(\"o1\")\n history = [\n ChatMessageSimple(\n message=\"Hello!\",\n token_count=5,\n message_type=MessageType.USER,\n )\n ]\n\n raw_result = translate_history_to_llm_format(history, llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 1\n msg = result[0]\n assert isinstance(msg, UserMessage)\n assert msg.content == \"Hello!\"\n\n def test_only_first_system_message_modified(self) -> None:\n \"\"\"Test that only the first system message gets the prefix.\"\"\"\n llm_config = _create_llm_config(\"o1\")\n history = [\n ChatMessageSimple(\n message=\"First system prompt.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n ),\n ChatMessageSimple(\n message=\"Hello!\",\n token_count=5,\n message_type=MessageType.USER,\n ),\n ChatMessageSimple(\n message=\"Second system prompt.\",\n token_count=10,\n message_type=MessageType.SYSTEM,\n ),\n ]\n\n raw_result = translate_history_to_llm_format(history, llm_config)\n result = _ensure_list(raw_result)\n\n assert len(result) == 3\n # First system message should have prefix\n first_sys = result[0]\n assert isinstance(first_sys, SystemMessage)\n assert isinstance(first_sys.content, str)\n assert first_sys.content.startswith(CODE_BLOCK_MARKDOWN)\n # Second system message should NOT have prefix (only first one is modified)\n second_sys = result[2]\n assert isinstance(second_sys, SystemMessage)\n assert isinstance(second_sys.content, str)\n assert not second_sys.content.startswith(CODE_BLOCK_MARKDOWN)\n" }, { "path": "backend/tests/external_dependency_unit/conftest.py", "content": "from collections.abc import Generator\nfrom uuid import uuid4\n\nimport pytest\nfrom fastapi_users.password import PasswordHelper\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import User\nfrom onyx.db.models import UserRole\nfrom onyx.file_store.file_store import get_default_file_store\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\nfrom tests.external_dependency_unit.full_setup import (\n ensure_full_deployment_setup,\n)\n\n\n@pytest.fixture(scope=\"function\")\ndef db_session() -> Generator[Session, None, None]:\n \"\"\"Create a database session for testing using the actual PostgreSQL database\"\"\"\n # Make sure that the db engine is initialized before any tests are run\n SqlEngine.init_engine(\n pool_size=10,\n max_overflow=5,\n )\n with get_session_with_current_tenant() as session:\n yield session\n\n\n@pytest.fixture(scope=\"session\")\ndef full_deployment_setup() -> Generator[None, None, None]:\n \"\"\"Optional fixture to perform full deployment-like setup on demand.\n\n Import and call tests.external_dependency_unit.startup.full_setup.ensure_full_deployment_setup\n to initialize Postgres defaults, Vespa indices, and seed initial docs.\n \"\"\"\n ensure_full_deployment_setup()\n yield\n\n\n@pytest.fixture(scope=\"function\")\ndef tenant_context() -> Generator[None, None, None]:\n \"\"\"Set up tenant context for testing\"\"\"\n # Set the tenant context for the test\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n try:\n yield\n finally:\n # Reset the tenant context after the test\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\ndef create_test_user(\n db_session: Session,\n email_prefix: str,\n role: UserRole = UserRole.BASIC,\n account_type: AccountType = AccountType.STANDARD,\n) -> User:\n \"\"\"Helper to create a test user with a unique email\"\"\"\n # Use UUID to ensure unique email addresses\n unique_email = f\"{email_prefix}_{uuid4().hex[:8]}@example.com\"\n\n password_helper = PasswordHelper()\n password = password_helper.generate()\n hashed_password = password_helper.hash(password)\n\n user = User(\n id=uuid4(),\n email=unique_email,\n hashed_password=hashed_password,\n is_active=True,\n is_superuser=False,\n is_verified=True,\n role=role,\n account_type=account_type,\n )\n db_session.add(user)\n db_session.commit()\n db_session.refresh(user)\n return user\n\n\n@pytest.fixture(scope=\"module\")\ndef initialize_file_store() -> Generator[None, None, None]:\n \"\"\"Initialize the file store for testing.\n\n Scoped to module level since file store initialization is idempotent\n and doesn't need to be reset between tests.\n \"\"\"\n get_default_file_store().initialize()\n yield\n" }, { "path": "backend/tests/external_dependency_unit/connectors/confluence/conftest.py", "content": "import os\nfrom typing import Any\n\nimport pytest\n\n\n@pytest.fixture\ndef confluence_connector_config() -> dict[str, Any]:\n url_base = os.environ.get(\"CONFLUENCE_TEST_SPACE_URL\")\n space_key = os.environ.get(\"CONFLUENCE_SPACE_KEY\")\n page_id = os.environ.get(\"CONFLUENCE_PAGE_ID\")\n is_cloud = os.environ.get(\"CONFLUENCE_IS_CLOUD\", \"true\").lower() == \"true\"\n\n assert url_base, \"CONFLUENCE_URL environment variable is required\"\n\n return {\n \"wiki_base\": url_base,\n \"is_cloud\": is_cloud,\n \"space\": space_key or \"\",\n \"page_id\": page_id or \"\",\n }\n\n\n@pytest.fixture\ndef confluence_credential_json() -> dict[str, Any]:\n username = os.environ.get(\"CONFLUENCE_USER_NAME\")\n access_token = os.environ.get(\"CONFLUENCE_ACCESS_TOKEN\")\n\n assert username, \"CONFLUENCE_USERNAME environment variable is required\"\n assert access_token, \"CONFLUENCE_ACCESS_TOKEN environment variable is required\"\n\n return {\n \"confluence_username\": username,\n \"confluence_access_token\": access_token,\n }\n" }, { "path": "backend/tests/external_dependency_unit/connectors/confluence/test_confluence_group_sync.py", "content": "from typing import Any\n\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom tests.daily.connectors.confluence.models import ExternalUserGroupSet\n\n\n# In order to get these tests to run, use the credentials from Bitwarden.\n# Search up \"ENV vars for local and Github tests\", and find the Confluence relevant key-value pairs.\n\n_EXPECTED_CONFLUENCE_GROUPS = [\n ExternalUserGroupSet(\n id=\"confluence-admins-danswerai\",\n user_emails={\"chris@onyx.app\", \"yuhong@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"org-admins\",\n user_emails={\n \"founders@onyx.app\",\n \"chris@onyx.app\",\n \"yuhong@onyx.app\",\n \"oauth@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"confluence-users-danswerai\",\n user_emails={\n \"chris@onyx.app\",\n \"hagen@danswer.ai\",\n \"founders@onyx.app\",\n \"pablo@onyx.app\",\n \"yuhong@onyx.app\",\n \"oauth@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-users-danswerai\",\n user_emails={\n \"hagen@danswer.ai\",\n \"founders@onyx.app\",\n \"pablo@onyx.app\",\n \"chris@onyx.app\",\n \"oauth@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-admins-danswerai\",\n user_emails={\"hagen@danswer.ai\", \"founders@onyx.app\", \"pablo@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"confluence-user-access-admins-danswerai\",\n user_emails={\"hagen@danswer.ai\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-user-access-admins-danswerai\",\n user_emails={\"hagen@danswer.ai\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"Yuhong Only No Chris Allowed\",\n user_emails={\"yuhong@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"All_Confluence_Users_Found_By_Onyx\",\n user_emails={\n \"chris@onyx.app\",\n \"founders@onyx.app\",\n \"hagen@danswer.ai\",\n \"pablo@onyx.app\",\n \"yuhong@onyx.app\",\n \"oauth@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"bitbucket-users-onyxai\",\n user_emails={\"founders@onyx.app\", \"oauth@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"bitbucket-admins-onyxai\",\n user_emails={\"founders@onyx.app\", \"oauth@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-servicemanagement-users-danswerai\",\n user_emails={\"oauth@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"no yuhong allowed\",\n user_emails={\"hagen@danswer.ai\", \"pablo@onyx.app\", \"chris@onyx.app\"},\n gives_anyone_access=False,\n ),\n]\n\n\ndef test_confluence_group_sync(\n db_session: Session,\n confluence_connector_config: dict[str, Any],\n confluence_credential_json: dict[str, Any],\n) -> None:\n connector = Connector(\n name=\"Test Connector\",\n source=DocumentSource.CONFLUENCE,\n input_type=InputType.POLL,\n connector_specific_config=confluence_connector_config,\n refresh_freq=None,\n prune_freq=None,\n indexing_start=None,\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n source=DocumentSource.CONFLUENCE,\n credential_json=confluence_credential_json,\n )\n db_session.add(credential)\n db_session.flush()\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.SYNC,\n auto_sync_options=None,\n )\n db_session.add(cc_pair)\n db_session.commit()\n db_session.refresh(cc_pair)\n\n tenant_id = get_current_tenant_id()\n group_sync_iter = confluence_group_sync(\n tenant_id=tenant_id,\n cc_pair=cc_pair,\n )\n\n expected_groups = {group.id: group for group in _EXPECTED_CONFLUENCE_GROUPS}\n actual_groups = {\n group.id: ExternalUserGroupSet.from_model(external_user_group=group)\n for group in group_sync_iter\n }\n assert expected_groups == actual_groups\n" }, { "path": "backend/tests/external_dependency_unit/connectors/google_drive/test_google_drive_group_sync.py", "content": "from collections.abc import Generator\nfrom unittest.mock import Mock\nfrom unittest.mock import patch\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.background.celery.tasks.external_group_syncing.tasks import (\n _perform_external_group_sync,\n)\nfrom ee.onyx.db.external_perm import ExternalUserGroup\nfrom onyx.access.utils import build_ext_group_name_for_onyx\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import AccountType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import PublicExternalUserGroup\nfrom onyx.db.models import User\nfrom onyx.db.models import User__ExternalUserGroupId\nfrom onyx.db.models import UserRole\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\ndef _create_ext_perm_user(db_session: Session, name: str) -> User:\n \"\"\"Create an external-permission user for group sync tests.\"\"\"\n return create_test_user(\n db_session,\n name,\n role=UserRole.EXT_PERM_USER,\n account_type=AccountType.EXT_PERM_USER,\n )\n\n\ndef _create_test_connector_credential_pair(\n db_session: Session, source: DocumentSource = DocumentSource.GOOGLE_DRIVE\n) -> ConnectorCredentialPair:\n \"\"\"Helper to create a test connector credential pair\"\"\"\n # For Google Drive, we need to include required config parameters\n connector_config = {}\n if source == DocumentSource.GOOGLE_DRIVE:\n connector_config = {\n \"include_shared_drives\": True, # At least one of these is required\n }\n\n connector = Connector(\n name=\"Test Connector\",\n source=source,\n input_type=InputType.POLL,\n connector_specific_config=connector_config,\n refresh_freq=None,\n prune_freq=None,\n indexing_start=None,\n )\n db_session.add(connector)\n db_session.flush() # To get the connector ID\n\n credential = Credential(\n source=source,\n credential_json={},\n user_id=None,\n )\n db_session.add(credential)\n db_session.flush() # To get the credential ID\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.SYNC,\n auto_sync_options=None,\n )\n db_session.add(cc_pair)\n db_session.commit()\n db_session.refresh(cc_pair)\n return cc_pair\n\n\ndef _get_user_external_groups(\n db_session: Session, cc_pair_id: int, include_stale: bool = False\n) -> list[User__ExternalUserGroupId]:\n \"\"\"Helper to get user external groups from database\"\"\"\n query = select(User__ExternalUserGroupId).where(\n User__ExternalUserGroupId.cc_pair_id == cc_pair_id\n )\n if not include_stale:\n query = query.where(User__ExternalUserGroupId.stale.is_(False))\n\n return list(db_session.scalars(query).all())\n\n\ndef _get_public_external_groups(\n db_session: Session, cc_pair_id: int, include_stale: bool = False\n) -> list[PublicExternalUserGroup]:\n \"\"\"Helper to get public external groups from database\"\"\"\n query = select(PublicExternalUserGroup).where(\n PublicExternalUserGroup.cc_pair_id == cc_pair_id\n )\n if not include_stale:\n query = query.where(PublicExternalUserGroup.stale.is_(False))\n\n return list(db_session.scalars(query).all())\n\n\nclass TestPerformExternalGroupSync:\n def test_initial_group_sync(self, db_session: Session) -> None:\n \"\"\"Test syncing external groups for the first time (initial sync)\"\"\"\n # Create test data\n user1 = _create_ext_perm_user(db_session, \"user1\")\n user2 = _create_ext_perm_user(db_session, \"user2\")\n user3 = _create_ext_perm_user(db_session, \"user3\")\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Mock external groups data as a generator that yields the expected groups\n mock_groups = [\n ExternalUserGroup(id=\"group1\", user_emails=[user1.email, user2.email]),\n ExternalUserGroup(id=\"group2\", user_emails=[user2.email, user3.email]),\n ExternalUserGroup(\n id=\"public_group\", user_emails=[user1.email], gives_anyone_access=True\n ),\n ]\n\n def mock_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n for group in mock_groups:\n yield group\n\n # Verify no groups exist initially\n assert len(_get_user_external_groups(db_session, cc_pair.id)) == 0\n assert len(_get_public_external_groups(db_session, cc_pair.id)) == 0\n\n with patch(\n \"ee.onyx.background.celery.tasks.external_group_syncing.tasks.get_source_perm_sync_config\"\n ) as mock_config:\n # Mock sync config\n mock_group_config = Mock()\n mock_group_config.group_sync_func = mock_group_sync_func\n\n mock_sync_config = Mock()\n mock_sync_config.group_sync_config = mock_group_config\n\n mock_config.return_value = mock_sync_config\n\n # Run the sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify user groups were created\n user_groups = _get_user_external_groups(db_session, cc_pair.id)\n assert (\n len(user_groups) == 5\n ) # user1+2 in group1, user2+3 in group2, user1 in public_group\n\n # Verify group names are properly prefixed\n expected_group1_id = build_ext_group_name_for_onyx(\n \"group1\", DocumentSource.GOOGLE_DRIVE\n )\n expected_group2_id = build_ext_group_name_for_onyx(\n \"group2\", DocumentSource.GOOGLE_DRIVE\n )\n expected_public_group_id = build_ext_group_name_for_onyx(\n \"public_group\", DocumentSource.GOOGLE_DRIVE\n )\n\n group_ids = {ug.external_user_group_id for ug in user_groups}\n assert expected_group1_id in group_ids\n assert expected_group2_id in group_ids\n assert expected_public_group_id in group_ids\n\n # Verify public group was created\n public_groups = _get_public_external_groups(db_session, cc_pair.id)\n assert len(public_groups) == 1\n assert public_groups[0].external_user_group_id == expected_public_group_id\n assert public_groups[0].stale is False\n\n # Verify all groups are not stale\n for ug in user_groups:\n assert ug.stale is False\n\n def test_update_existing_groups(self, db_session: Session) -> None:\n \"\"\"Test updating existing groups (adding/removing users)\"\"\"\n # Create test data\n user1 = _create_ext_perm_user(db_session, \"user1\")\n user2 = _create_ext_perm_user(db_session, \"user2\")\n user3 = _create_ext_perm_user(db_session, \"user3\")\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Initial sync with original groups\n def initial_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n yield ExternalUserGroup(id=\"group1\", user_emails=[user1.email, user2.email])\n yield ExternalUserGroup(id=\"group2\", user_emails=[user2.email])\n\n # For now, verify test setup is working\n assert len(_get_user_external_groups(db_session, cc_pair.id)) == 0\n\n with patch(\n \"ee.onyx.background.celery.tasks.external_group_syncing.tasks.get_source_perm_sync_config\"\n ) as mock_config:\n # Mock sync config\n mock_group_config = Mock()\n mock_group_config.group_sync_func = initial_group_sync_func\n\n mock_sync_config = Mock()\n mock_sync_config.group_sync_config = mock_group_config\n\n mock_config.return_value = mock_sync_config\n\n # Run initial sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify initial state\n initial_user_groups = _get_user_external_groups(db_session, cc_pair.id)\n assert (\n len(initial_user_groups) == 3\n ) # user1+user2 in group1, user2 in group2\n\n # Updated sync with modified groups\n def updated_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n # group1 now has user1 and user3 (user2 removed, user3 added)\n yield ExternalUserGroup(\n id=\"group1\", user_emails=[user1.email, user3.email]\n )\n # group2 now has all three users (user1 and user3 added)\n yield ExternalUserGroup(\n id=\"group2\", user_emails=[user1.email, user2.email, user3.email]\n )\n\n # Update the mock function\n mock_group_config.group_sync_func = updated_group_sync_func\n\n # Run updated sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify updated state\n updated_user_groups = _get_user_external_groups(db_session, cc_pair.id)\n assert (\n len(updated_user_groups) == 5\n ) # user1+user3 in group1, user1+user2+user3 in group2\n\n # Verify specific user-group mappings\n expected_group1_id = build_ext_group_name_for_onyx(\n \"group1\", DocumentSource.GOOGLE_DRIVE\n )\n expected_group2_id = build_ext_group_name_for_onyx(\n \"group2\", DocumentSource.GOOGLE_DRIVE\n )\n\n group1_users = {\n ug.user_id\n for ug in updated_user_groups\n if ug.external_user_group_id == expected_group1_id\n }\n group2_users = {\n ug.user_id\n for ug in updated_user_groups\n if ug.external_user_group_id == expected_group2_id\n }\n\n assert user1.id in group1_users and user3.id in group1_users\n assert user2.id not in group1_users # user2 was removed from group1\n assert (\n user1.id in group2_users\n and user2.id in group2_users\n and user3.id in group2_users\n )\n\n # Verify no stale groups remain\n for ug in updated_user_groups:\n assert ug.stale is False\n\n def test_remove_groups(self, db_session: Session) -> None:\n \"\"\"Test removing groups (groups that no longer exist in external system)\"\"\"\n # Create test data\n user1 = _create_ext_perm_user(db_session, \"user1\")\n user2 = _create_ext_perm_user(db_session, \"user2\")\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Initial sync with multiple groups\n def initial_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n yield ExternalUserGroup(id=\"group1\", user_emails=[user1.email, user2.email])\n yield ExternalUserGroup(id=\"group2\", user_emails=[user1.email])\n yield ExternalUserGroup(\n id=\"public_group\", user_emails=[user1.email], gives_anyone_access=True\n )\n\n assert len(_get_user_external_groups(db_session, cc_pair.id)) == 0\n assert len(_get_public_external_groups(db_session, cc_pair.id)) == 0\n\n with patch(\n \"ee.onyx.background.celery.tasks.external_group_syncing.tasks.get_source_perm_sync_config\"\n ) as mock_config:\n # Mock sync config\n mock_group_config = Mock()\n mock_group_config.group_sync_func = initial_group_sync_func\n\n mock_sync_config = Mock()\n mock_sync_config.group_sync_config = mock_group_config\n\n mock_config.return_value = mock_sync_config\n\n # Run initial sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify initial state\n initial_user_groups = _get_user_external_groups(db_session, cc_pair.id)\n initial_public_groups = _get_public_external_groups(db_session, cc_pair.id)\n assert (\n len(initial_user_groups) == 4\n ) # 2 in group1, 1 in group2, 1 in public_group\n assert len(initial_public_groups) == 1\n\n # Updated sync with only one group remaining\n def updated_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n # Only group1 remains, group2 and public_group are removed\n yield ExternalUserGroup(\n id=\"group1\", user_emails=[user1.email, user2.email]\n )\n\n # Update the mock function\n mock_group_config.group_sync_func = updated_group_sync_func\n\n # Run updated sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify updated state\n updated_user_groups = _get_user_external_groups(db_session, cc_pair.id)\n updated_public_groups = _get_public_external_groups(db_session, cc_pair.id)\n\n assert len(updated_user_groups) == 2 # Only group1 mappings remain\n assert len(updated_public_groups) == 0 # Public group was removed\n\n # Verify only group1 exists\n expected_group1_id = build_ext_group_name_for_onyx(\n \"group1\", DocumentSource.GOOGLE_DRIVE\n )\n group_ids = {ug.external_user_group_id for ug in updated_user_groups}\n assert group_ids == {expected_group1_id}\n\n # Verify stale groups were actually deleted from database\n all_user_groups_including_stale = _get_user_external_groups(\n db_session, cc_pair.id, include_stale=True\n )\n all_public_groups_including_stale = _get_public_external_groups(\n db_session, cc_pair.id, include_stale=True\n )\n\n assert len(all_user_groups_including_stale) == 2 # Only group1 mappings\n assert len(all_public_groups_including_stale) == 0 # Public group deleted\n\n def test_empty_group_sync(self, db_session: Session) -> None:\n \"\"\"Test syncing when no groups are returned (all groups removed)\"\"\"\n # Create test data\n user1 = _create_ext_perm_user(db_session, \"user1\")\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Initial sync with groups\n def initial_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n yield ExternalUserGroup(id=\"group1\", user_emails=[user1.email])\n\n with patch(\n \"ee.onyx.background.celery.tasks.external_group_syncing.tasks.get_source_perm_sync_config\"\n ) as mock_config:\n # Mock sync config\n mock_group_config = Mock()\n mock_group_config.group_sync_func = initial_group_sync_func\n\n mock_sync_config = Mock()\n mock_sync_config.group_sync_config = mock_group_config\n\n mock_config.return_value = mock_sync_config\n\n # Run initial sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify initial state\n initial_user_groups = _get_user_external_groups(db_session, cc_pair.id)\n assert len(initial_user_groups) == 1\n\n # Updated sync with no groups\n def empty_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n # No groups yielded\n return\n yield # This line is never reached but satisfies the generator type\n\n # Update the mock function\n mock_group_config.group_sync_func = empty_group_sync_func\n\n # Run updated sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify all groups were removed\n updated_user_groups = _get_user_external_groups(db_session, cc_pair.id)\n updated_public_groups = _get_public_external_groups(db_session, cc_pair.id)\n\n assert len(updated_user_groups) == 0\n assert len(updated_public_groups) == 0\n\n def test_batch_processing(self, db_session: Session) -> None:\n \"\"\"Test that large numbers of groups are processed in batches\"\"\"\n # Create many test users\n users = []\n for i in range(150): # More than the batch size of 100\n users.append(_create_ext_perm_user(db_session, f\"user{i}\"))\n\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Create a large group with many users\n def large_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n yield ExternalUserGroup(\n id=\"large_group\", user_emails=[user.email for user in users]\n )\n\n with patch(\n \"ee.onyx.background.celery.tasks.external_group_syncing.tasks.get_source_perm_sync_config\"\n ) as mock_config:\n # Mock sync config\n mock_group_config = Mock()\n mock_group_config.group_sync_func = large_group_sync_func\n\n mock_sync_config = Mock()\n mock_sync_config.group_sync_config = mock_group_config\n\n mock_config.return_value = mock_sync_config\n\n # Run the sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify all users were added to the group\n user_groups = _get_user_external_groups(db_session, cc_pair.id)\n assert len(user_groups) == 150\n\n # Verify all groups are not stale\n for ug in user_groups:\n assert ug.stale is False\n\n def test_mixed_regular_and_public_groups(self, db_session: Session) -> None:\n \"\"\"Test syncing a mix of regular and public groups\"\"\"\n # Create test data\n user1 = _create_ext_perm_user(db_session, \"user1\")\n user2 = _create_ext_perm_user(db_session, \"user2\")\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n def mixed_group_sync_func(\n tenant_id: str, # noqa: ARG001\n cc_pair: ConnectorCredentialPair, # noqa: ARG001\n ) -> Generator[ExternalUserGroup, None, None]:\n yield ExternalUserGroup(\n id=\"regular_group\", user_emails=[user1.email, user2.email]\n )\n yield ExternalUserGroup(\n id=\"public_group1\", user_emails=[user1.email], gives_anyone_access=True\n )\n yield ExternalUserGroup(\n id=\"public_group2\",\n user_emails=[], # Empty user list for public group\n gives_anyone_access=True,\n )\n\n with patch(\n \"ee.onyx.background.celery.tasks.external_group_syncing.tasks.get_source_perm_sync_config\"\n ) as mock_config:\n # Mock sync config\n mock_group_config = Mock()\n mock_group_config.group_sync_func = mixed_group_sync_func\n\n mock_sync_config = Mock()\n mock_sync_config.group_sync_config = mock_group_config\n\n mock_config.return_value = mock_sync_config\n\n # Run the sync\n _perform_external_group_sync(cc_pair.id, TEST_TENANT_ID)\n\n # Verify user groups\n user_groups = _get_user_external_groups(db_session, cc_pair.id)\n expected_regular_group_id = build_ext_group_name_for_onyx(\n \"regular_group\", DocumentSource.GOOGLE_DRIVE\n )\n expected_public_group1_id = build_ext_group_name_for_onyx(\n \"public_group1\", DocumentSource.GOOGLE_DRIVE\n )\n\n # Should have 2 users in regular_group + 1 user in public_group1 = 3 total\n assert len(user_groups) == 3\n\n regular_group_users = [\n ug\n for ug in user_groups\n if ug.external_user_group_id == expected_regular_group_id\n ]\n public_group1_users = [\n ug\n for ug in user_groups\n if ug.external_user_group_id == expected_public_group1_id\n ]\n\n assert len(regular_group_users) == 2\n assert len(public_group1_users) == 1\n\n # Verify public groups\n public_groups = _get_public_external_groups(db_session, cc_pair.id)\n assert len(public_groups) == 2 # public_group1 and public_group2\n\n public_group_ids = {pg.external_user_group_id for pg in public_groups}\n expected_public_group2_id = build_ext_group_name_for_onyx(\n \"public_group2\", DocumentSource.GOOGLE_DRIVE\n )\n assert expected_public_group1_id in public_group_ids\n assert expected_public_group2_id in public_group_ids\n" }, { "path": "backend/tests/external_dependency_unit/connectors/jira/conftest.py", "content": "import os\nfrom typing import Any\n\nimport pytest\n\n\n@pytest.fixture\ndef jira_connector_config() -> dict[str, Any]:\n jira_base_url = os.environ.get(\"JIRA_BASE_URL\", \"https://danswerai.atlassian.net\")\n\n return {\n \"jira_base_url\": jira_base_url,\n \"project_key\": \"\", # Empty to sync all projects\n \"scoped_token\": False,\n }\n\n\n@pytest.fixture\ndef jira_credential_json() -> dict[str, Any]:\n user_email = os.environ.get(\"JIRA_ADMIN_USER_EMAIL\", \"chris@onyx.app\")\n api_token = os.environ.get(\"JIRA_ADMIN_API_TOKEN\")\n\n assert user_email, \"JIRA_ADMIN_USER_EMAIL environment variable is required\"\n assert api_token, \"JIRA_ADMIN_API_TOKEN environment variable is required\"\n\n return {\n \"jira_user_email\": user_email,\n \"jira_api_token\": api_token,\n }\n" }, { "path": "backend/tests/external_dependency_unit/connectors/jira/test_jira_doc_sync.py", "content": "from typing import Any\n\nimport pytest\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync\nfrom onyx.access.models import DocExternalAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.utils import DocumentRow\nfrom onyx.db.utils import SortOrder\n\n\n# In order to get these tests to run, use the credentials from Bitwarden.\n# Search up \"ENV vars for local and Github tests\", and find the Jira relevant key-value pairs.\n# Required env vars: JIRA_USER_EMAIL, JIRA_API_TOKEN\n\npytestmark = pytest.mark.usefixtures(\"enable_ee\")\n\n\nclass DocExternalAccessSet(BaseModel):\n \"\"\"A version of DocExternalAccess that uses sets for comparison.\"\"\"\n\n doc_id: str\n external_user_emails: set[str]\n external_user_group_ids: set[str]\n is_public: bool\n\n @classmethod\n def from_doc_external_access(\n cls, doc_external_access: DocExternalAccess\n ) -> \"DocExternalAccessSet\":\n return cls(\n doc_id=doc_external_access.doc_id,\n external_user_emails=doc_external_access.external_access.external_user_emails,\n external_user_group_ids=doc_external_access.external_access.external_user_group_ids,\n is_public=doc_external_access.external_access.is_public,\n )\n\n\ndef test_jira_doc_sync(\n db_session: Session,\n jira_connector_config: dict[str, Any],\n jira_credential_json: dict[str, Any],\n) -> None:\n \"\"\"Test that Jira doc sync returns documents with correct permissions.\n\n This test uses the AS project which has applicationRole permission,\n meaning all documents should be marked as public.\n \"\"\"\n try:\n # Use AS project specifically for this test\n connector_config = {\n **jira_connector_config,\n \"project_key\": \"AS\", # DailyConnectorTestProject\n }\n\n connector = Connector(\n name=\"Test Jira Doc Sync Connector\",\n source=DocumentSource.JIRA,\n input_type=InputType.POLL,\n connector_specific_config=connector_config,\n refresh_freq=None,\n prune_freq=None,\n indexing_start=None,\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n source=DocumentSource.JIRA,\n credential_json=jira_credential_json,\n )\n db_session.add(credential)\n db_session.flush()\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test Jira Doc Sync CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.SYNC,\n auto_sync_options=None,\n )\n db_session.add(cc_pair)\n db_session.flush()\n db_session.refresh(cc_pair)\n\n # Mock functions - we don't have existing docs in the test DB\n def fetch_all_existing_docs_fn(\n sort_order: SortOrder | None = None, # noqa: ARG001\n ) -> list[DocumentRow]:\n return []\n\n def fetch_all_existing_docs_ids_fn() -> list[str]:\n return []\n\n doc_sync_iter = jira_doc_sync(\n cc_pair=cc_pair,\n fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,\n fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,\n )\n\n # Expected documents from the danswerai.atlassian.net Jira instance\n # The AS project has applicationRole permission, so all docs should be public\n _EXPECTED_JIRA_DOCS = [\n DocExternalAccessSet(\n doc_id=\"https://danswerai.atlassian.net/browse/AS-3\",\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n ),\n DocExternalAccessSet(\n doc_id=\"https://danswerai.atlassian.net/browse/AS-4\",\n external_user_emails=set(),\n external_user_group_ids=set(),\n is_public=True,\n ),\n ]\n\n expected_docs = {doc.doc_id: doc for doc in _EXPECTED_JIRA_DOCS}\n actual_docs = {\n doc.doc_id: DocExternalAccessSet.from_doc_external_access(doc)\n for doc in doc_sync_iter\n if isinstance(doc, DocExternalAccess)\n }\n assert (\n expected_docs == actual_docs\n ), f\"Expected docs: {expected_docs}\\nActual docs: {actual_docs}\"\n finally:\n db_session.rollback()\n\n\ndef test_jira_doc_sync_with_specific_permissions(\n db_session: Session,\n jira_connector_config: dict[str, Any],\n jira_credential_json: dict[str, Any],\n) -> None:\n \"\"\"Test that Jira doc sync returns documents with specific permissions.\n\n This test uses a project that has specific user permissions to verify\n that specific users are correctly extracted.\n \"\"\"\n try:\n # Use SUP project which has specific user permissions\n connector_config = {\n **jira_connector_config,\n \"project_key\": \"SUP\",\n }\n\n connector = Connector(\n name=\"Test Jira Doc Sync with Groups Connector\",\n source=DocumentSource.JIRA,\n input_type=InputType.POLL,\n connector_specific_config=connector_config,\n refresh_freq=None,\n prune_freq=None,\n indexing_start=None,\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n source=DocumentSource.JIRA,\n credential_json=jira_credential_json,\n )\n db_session.add(credential)\n db_session.flush()\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test Jira Doc Sync with Groups CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.SYNC,\n auto_sync_options=None,\n )\n db_session.add(cc_pair)\n db_session.flush()\n db_session.refresh(cc_pair)\n\n # Mock functions\n def fetch_all_existing_docs_fn(\n sort_order: SortOrder | None = None, # noqa: ARG001\n ) -> list[DocumentRow]:\n return []\n\n def fetch_all_existing_docs_ids_fn() -> list[str]:\n return []\n\n doc_sync_iter = jira_doc_sync(\n cc_pair=cc_pair,\n fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,\n fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,\n )\n\n docs = list(doc_sync_iter)\n\n # SUP project should have user-specific permissions (not public)\n assert len(docs) > 0, \"Expected at least one document from SUP project\"\n\n _EXPECTED_USER_EMAILS = set(\n [\"yuhong@onyx.app\", \"chris@onyx.app\", \"founders@onyx.app\"]\n )\n _EXPECTED_USER_GROUP_IDS = set([\"jira-users-danswerai\"])\n\n for doc in docs:\n if not isinstance(doc, DocExternalAccess):\n continue\n assert doc.doc_id.startswith(\"https://danswerai.atlassian.net/browse/SUP-\")\n # SUP project has specific users assigned, not applicationRole\n assert (\n not doc.external_access.is_public\n ), f\"Document {doc.doc_id} should not be public\"\n # Should have user emails\n assert doc.external_access.external_user_emails == _EXPECTED_USER_EMAILS\n assert (\n doc.external_access.external_user_group_ids == _EXPECTED_USER_GROUP_IDS\n )\n finally:\n db_session.rollback()\n" }, { "path": "backend/tests/external_dependency_unit/connectors/jira/test_jira_group_sync.py", "content": "from typing import Any\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.external_permissions.jira.group_sync import jira_group_sync\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom tests.daily.connectors.confluence.models import ExternalUserGroupSet\n\n\n# In order to get these tests to run, use the credentials from Bitwarden.\n# Search up \"ENV vars for local and Github tests\", and find the Jira relevant key-value pairs.\n# Required env vars: JIRA_USER_EMAIL, JIRA_API_TOKEN\n\npytestmark = pytest.mark.usefixtures(\"enable_ee\")\n\n# Expected groups from the danswerai.atlassian.net Jira instance\n# Note: These groups are shared with Confluence since they're both Atlassian products\n# App accounts (bots, integrations) are filtered out\n_EXPECTED_JIRA_GROUPS = [\n ExternalUserGroupSet(\n id=\"Yuhong Only No Chris Allowed\",\n user_emails={\"yuhong@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"confluence-admins-danswerai\",\n user_emails={\"chris@onyx.app\", \"yuhong@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"confluence-user-access-admins-danswerai\",\n user_emails={\"hagen@danswer.ai\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"confluence-users-danswerai\",\n user_emails={\n \"chris@onyx.app\",\n \"founders@onyx.app\",\n \"hagen@danswer.ai\",\n \"pablo@onyx.app\",\n \"yuhong@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-admins-danswerai\",\n user_emails={\"founders@onyx.app\", \"hagen@danswer.ai\", \"pablo@onyx.app\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-user-access-admins-danswerai\",\n user_emails={\"hagen@danswer.ai\"},\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"jira-users-danswerai\",\n user_emails={\n \"chris@onyx.app\",\n \"founders@onyx.app\",\n \"hagen@danswer.ai\",\n \"pablo@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"org-admins\",\n user_emails={\n \"chris@onyx.app\",\n \"founders@onyx.app\",\n \"yuhong@onyx.app\",\n },\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"bitbucket-admins-onyxai\",\n user_emails={\"founders@onyx.app\"}, # no Oauth, we skip \"app\" account in jira\n gives_anyone_access=False,\n ),\n ExternalUserGroupSet(\n id=\"bitbucket-users-onyxai\",\n user_emails={\"founders@onyx.app\"}, # no Oauth, we skip \"app\" account in jira\n gives_anyone_access=False,\n ),\n]\n\n\ndef test_jira_group_sync(\n db_session: Session,\n jira_connector_config: dict[str, Any],\n jira_credential_json: dict[str, Any],\n) -> None:\n try:\n connector = Connector(\n name=\"Test Jira Connector\",\n source=DocumentSource.JIRA,\n input_type=InputType.POLL,\n connector_specific_config=jira_connector_config,\n refresh_freq=None,\n prune_freq=None,\n indexing_start=None,\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n source=DocumentSource.JIRA,\n credential_json=jira_credential_json,\n )\n db_session.add(credential)\n db_session.flush()\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test Jira CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.SYNC,\n auto_sync_options=None,\n )\n db_session.add(cc_pair)\n db_session.flush()\n db_session.refresh(cc_pair)\n\n tenant_id = get_current_tenant_id()\n group_sync_iter = jira_group_sync(\n tenant_id=tenant_id,\n cc_pair=cc_pair,\n )\n\n expected_groups = {group.id: group for group in _EXPECTED_JIRA_GROUPS}\n actual_groups = {\n group.id: ExternalUserGroupSet.from_model(external_user_group=group)\n for group in group_sync_iter\n }\n assert expected_groups == actual_groups\n finally:\n db_session.rollback()\n" }, { "path": "backend/tests/external_dependency_unit/constants.py", "content": "TEST_TENANT_ID: str = \"public\"\n" }, { "path": "backend/tests/external_dependency_unit/craft/conftest.py", "content": "\"\"\"Fixtures for build mode tests.\"\"\"\n\nfrom collections.abc import Generator\nfrom uuid import uuid4\n\nimport pytest\nfrom fastapi_users.password import PasswordHelper\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.enums import AccountType\nfrom onyx.db.enums import BuildSessionStatus\nfrom onyx.db.models import BuildSession\nfrom onyx.db.models import User\nfrom onyx.db.models import UserRole\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\n@pytest.fixture(scope=\"function\")\ndef db_session() -> Generator[Session, None, None]:\n \"\"\"Create a database session for testing using the actual PostgreSQL database.\"\"\"\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n with get_session_with_current_tenant() as session:\n yield session\n\n\n@pytest.fixture(scope=\"function\")\ndef tenant_context() -> Generator[None, None, None]:\n \"\"\"Set up tenant context for testing.\"\"\"\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n try:\n yield\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\n@pytest.fixture(scope=\"function\")\ndef test_user(db_session: Session, tenant_context: None) -> User: # noqa: ARG001\n \"\"\"Create a test user for build session tests.\"\"\"\n unique_email = f\"build_test_{uuid4().hex[:8]}@example.com\"\n\n password_helper = PasswordHelper()\n password = password_helper.generate()\n hashed_password = password_helper.hash(password)\n\n user = User(\n id=uuid4(),\n email=unique_email,\n hashed_password=hashed_password,\n is_active=True,\n is_superuser=False,\n is_verified=True,\n role=UserRole.EXT_PERM_USER,\n account_type=AccountType.EXT_PERM_USER,\n )\n db_session.add(user)\n db_session.commit()\n db_session.refresh(user)\n return user\n\n\n@pytest.fixture(scope=\"function\")\ndef build_session(\n db_session: Session,\n test_user: User,\n tenant_context: None, # noqa: ARG001\n) -> BuildSession:\n \"\"\"Create a test build session.\"\"\"\n session = BuildSession(\n id=uuid4(),\n user_id=test_user.id,\n name=\"Test Build Session\",\n status=BuildSessionStatus.ACTIVE,\n )\n db_session.add(session)\n db_session.commit()\n db_session.refresh(session)\n return session\n" }, { "path": "backend/tests/external_dependency_unit/craft/test_build_packet_storage.py", "content": "\"\"\"\nTest suite for build mode packet storage.\n\nTests the new packet storage behavior:\n- All data stored in message_metadata as JSON (no content column)\n- turn_index tracks which user message each assistant message belongs to\n- Tool calls: Only save when status=\"completed\"\n- Message/thought chunks: Accumulated and saved as synthetic packets\n- Agent plan updates: Upserted (only latest kept per turn)\n\"\"\"\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import MessageType\nfrom onyx.db.models import BuildSession\nfrom onyx.server.features.build.db.build_session import create_message\nfrom onyx.server.features.build.db.build_session import get_session_messages\nfrom onyx.server.features.build.db.build_session import upsert_agent_plan\nfrom onyx.server.features.build.session.manager import BuildStreamingState\n\n\nclass TestBuildMessageStorage:\n \"\"\"Tests for build message storage in the database.\"\"\"\n\n def test_create_message_with_metadata(\n self,\n db_session: Session,\n build_session: BuildSession,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test creating a message with JSON metadata and turn_index.\"\"\"\n user_message_metadata = {\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Hello, world!\"},\n }\n\n message = create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=0,\n message_metadata=user_message_metadata,\n db_session=db_session,\n )\n\n assert message.id is not None\n assert message.session_id == build_session.id\n assert message.type == MessageType.USER\n assert message.turn_index == 0\n assert message.message_metadata == user_message_metadata\n assert message.message_metadata[\"type\"] == \"user_message\"\n assert message.message_metadata[\"content\"][\"text\"] == \"Hello, world!\"\n\n def test_create_multiple_messages_with_turn_index(\n self,\n db_session: Session,\n build_session: BuildSession,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test creating multiple messages with correct turn_index values.\"\"\"\n # First user message (turn 0)\n create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=0,\n message_metadata={\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"First question\"},\n },\n db_session=db_session,\n )\n\n # Assistant response (turn 0)\n create_message(\n session_id=build_session.id,\n message_type=MessageType.ASSISTANT,\n turn_index=0,\n message_metadata={\n \"type\": \"agent_message\",\n \"content\": {\"type\": \"text\", \"text\": \"First answer\"},\n },\n db_session=db_session,\n )\n\n # Second user message (turn 1)\n create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=1,\n message_metadata={\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Second question\"},\n },\n db_session=db_session,\n )\n\n # Assistant response (turn 1)\n create_message(\n session_id=build_session.id,\n message_type=MessageType.ASSISTANT,\n turn_index=1,\n message_metadata={\n \"type\": \"agent_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Second answer\"},\n },\n db_session=db_session,\n )\n\n # Verify messages\n messages = get_session_messages(build_session.id, db_session)\n assert len(messages) == 4\n\n # Check turn indices\n turn_0_messages = [m for m in messages if m.turn_index == 0]\n turn_1_messages = [m for m in messages if m.turn_index == 1]\n\n assert len(turn_0_messages) == 2\n assert len(turn_1_messages) == 2\n\n def test_tool_call_completed_storage(\n self,\n db_session: Session,\n build_session: BuildSession,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test storing only completed tool calls.\"\"\"\n # Create a user message first\n create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=0,\n message_metadata={\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Run a tool\"},\n },\n db_session=db_session,\n )\n\n # Create a completed tool call\n tool_call_packet = {\n \"type\": \"tool_call_progress\",\n \"toolCallId\": \"tool-123\",\n \"status\": \"completed\",\n \"kind\": \"bash\",\n \"title\": \"Running command\",\n \"rawOutput\": \"Command completed successfully\",\n \"timestamp\": \"2025-01-01T00:00:00Z\",\n }\n\n message = create_message(\n session_id=build_session.id,\n message_type=MessageType.ASSISTANT,\n turn_index=0,\n message_metadata=tool_call_packet,\n db_session=db_session,\n )\n\n assert message.message_metadata[\"type\"] == \"tool_call_progress\"\n assert message.message_metadata[\"status\"] == \"completed\"\n assert message.message_metadata[\"toolCallId\"] == \"tool-123\"\n\n def test_upsert_agent_plan(\n self,\n db_session: Session,\n build_session: BuildSession,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test upserting agent plan - only latest should be kept.\"\"\"\n # Create a user message first\n create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=0,\n message_metadata={\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Create a plan\"},\n },\n db_session=db_session,\n )\n\n # First plan\n plan1 = {\n \"type\": \"agent_plan_update\",\n \"entries\": [\n {\"id\": \"1\", \"status\": \"pending\", \"content\": \"Step 1\"},\n ],\n \"timestamp\": \"2025-01-01T00:00:00Z\",\n }\n\n plan_msg1 = upsert_agent_plan(\n session_id=build_session.id,\n turn_index=0,\n plan_metadata=plan1,\n db_session=db_session,\n )\n\n assert plan_msg1.message_metadata[\"entries\"][0][\"status\"] == \"pending\"\n\n # Update plan with new status\n plan2 = {\n \"type\": \"agent_plan_update\",\n \"entries\": [\n {\"id\": \"1\", \"status\": \"completed\", \"content\": \"Step 1\"},\n {\"id\": \"2\", \"status\": \"in_progress\", \"content\": \"Step 2\"},\n ],\n \"timestamp\": \"2025-01-01T00:01:00Z\",\n }\n\n plan_msg2 = upsert_agent_plan(\n session_id=build_session.id,\n turn_index=0,\n plan_metadata=plan2,\n db_session=db_session,\n existing_plan_id=plan_msg1.id,\n )\n\n # Should be the same message, updated\n assert plan_msg2.id == plan_msg1.id\n assert len(plan_msg2.message_metadata[\"entries\"]) == 2\n assert plan_msg2.message_metadata[\"entries\"][0][\"status\"] == \"completed\"\n\n # Verify only one plan message exists for this turn\n messages = get_session_messages(build_session.id, db_session)\n plan_messages = [\n m for m in messages if m.message_metadata.get(\"type\") == \"agent_plan_update\"\n ]\n assert len(plan_messages) == 1\n\n def test_upsert_agent_plan_without_existing_id(\n self,\n db_session: Session,\n build_session: BuildSession,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test upserting agent plan when we don't know the existing ID.\"\"\"\n # Create a user message first\n create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=0,\n message_metadata={\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Create a plan\"},\n },\n db_session=db_session,\n )\n\n # First plan - no existing ID\n plan1 = {\n \"type\": \"agent_plan_update\",\n \"entries\": [{\"id\": \"1\", \"status\": \"pending\", \"content\": \"Step 1\"}],\n }\n\n plan_msg1 = upsert_agent_plan(\n session_id=build_session.id,\n turn_index=0,\n plan_metadata=plan1,\n db_session=db_session,\n )\n\n # Second plan - still no existing ID, should find and update\n plan2 = {\n \"type\": \"agent_plan_update\",\n \"entries\": [{\"id\": \"1\", \"status\": \"completed\", \"content\": \"Step 1\"}],\n }\n\n plan_msg2 = upsert_agent_plan(\n session_id=build_session.id,\n turn_index=0,\n plan_metadata=plan2,\n db_session=db_session,\n )\n\n # Should be the same message\n assert plan_msg2.id == plan_msg1.id\n\n def test_streaming_flow_db_calls(\n self,\n db_session: Session,\n build_session: BuildSession,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test that streaming flow creates correct number of DB messages.\n\n Simulates:\n 1. Agent message chunks -> 1 message\n 2. Tool call -> 1 message\n 3. Agent message chunks -> 1 message\n\n This verifies that we save parts of the turn as they finish, rather than\n buffering everything into one giant message or losing granularity.\n \"\"\"\n # 0. Initial user message\n create_message(\n session_id=build_session.id,\n message_type=MessageType.USER,\n turn_index=0,\n message_metadata={\n \"type\": \"user_message\",\n \"content\": {\"type\": \"text\", \"text\": \"Do something\"},\n },\n db_session=db_session,\n )\n\n state = BuildStreamingState(turn_index=0)\n\n # 1. Stream agent message chunks\n state.add_message_chunk(\"Thinking\")\n state.add_message_chunk(\" about it...\")\n\n # Simulate switch to tool call (e.g. ToolCallStart event) -> finalize message\n # In SessionManager, this happens via state.should_finalize_chunks()\n if state.should_finalize_chunks(\"tool_call_start\"):\n msg_packet = state.finalize_message_chunks()\n if msg_packet:\n create_message(\n session_id=build_session.id,\n message_type=MessageType.ASSISTANT,\n turn_index=0,\n message_metadata=msg_packet,\n db_session=db_session,\n )\n state.clear_last_chunk_type()\n\n # 2. Handle completed tool call (immediate save)\n tool_packet = {\n \"type\": \"tool_call_progress\",\n \"toolCallId\": \"call_1\",\n \"status\": \"completed\",\n \"timestamp\": \"2025-01-01T00:00:00Z\",\n }\n create_message(\n session_id=build_session.id,\n message_type=MessageType.ASSISTANT,\n turn_index=0,\n message_metadata=tool_packet,\n db_session=db_session,\n )\n\n # 3. Stream more agent message chunks\n state.add_message_chunk(\"Done\")\n state.add_message_chunk(\" with tool.\")\n\n # End of stream -> finalize\n msg_packet = state.finalize_message_chunks()\n if msg_packet:\n create_message(\n session_id=build_session.id,\n message_type=MessageType.ASSISTANT,\n turn_index=0,\n message_metadata=msg_packet,\n db_session=db_session,\n )\n\n # Verify DB state\n messages = get_session_messages(build_session.id, db_session)\n # 1 user + 3 assistant = 4 total\n assert len(messages) == 4\n\n # Verify types/order\n assert messages[0].type == MessageType.USER\n\n assert messages[1].type == MessageType.ASSISTANT\n assert messages[1].message_metadata[\"content\"][\"text\"] == \"Thinking about it...\"\n\n assert messages[2].type == MessageType.ASSISTANT\n assert messages[2].message_metadata[\"type\"] == \"tool_call_progress\"\n\n assert messages[3].type == MessageType.ASSISTANT\n assert messages[3].message_metadata[\"content\"][\"text\"] == \"Done with tool.\"\n\n\nclass TestBuildStreamingState:\n \"\"\"Tests for BuildStreamingState class.\"\"\"\n\n def test_message_chunk_accumulation(self) -> None:\n \"\"\"Test accumulating message chunks.\"\"\"\n state = BuildStreamingState(turn_index=0)\n\n state.add_message_chunk(\"Hello, \")\n state.add_message_chunk(\"world!\")\n\n packet = state.finalize_message_chunks()\n\n assert packet is not None\n assert packet[\"type\"] == \"agent_message\"\n assert packet[\"content\"][\"text\"] == \"Hello, world!\"\n\n # After finalize, chunks should be cleared\n assert len(state.message_chunks) == 0\n\n def test_thought_chunk_accumulation(self) -> None:\n \"\"\"Test accumulating thought chunks.\"\"\"\n state = BuildStreamingState(turn_index=0)\n\n state.add_thought_chunk(\"Thinking about \")\n state.add_thought_chunk(\"the problem...\")\n\n packet = state.finalize_thought_chunks()\n\n assert packet is not None\n assert packet[\"type\"] == \"agent_thought\"\n assert packet[\"content\"][\"text\"] == \"Thinking about the problem...\"\n\n def test_should_finalize_chunks_on_type_change(self) -> None:\n \"\"\"Test detection of when to finalize chunks.\"\"\"\n state = BuildStreamingState(turn_index=0)\n\n # Add message chunk\n state.add_message_chunk(\"Hello\")\n\n # Should finalize when receiving non-message packet\n assert state.should_finalize_chunks(\"tool_call_start\") is True\n assert state.should_finalize_chunks(\"agent_plan_update\") is True\n assert state.should_finalize_chunks(\"agent_thought_chunk\") is True\n\n # Should NOT finalize for same type\n assert state.should_finalize_chunks(\"agent_message_chunk\") is False\n\n def test_finalize_returns_none_when_empty(self) -> None:\n \"\"\"Test that finalize returns None when no chunks accumulated.\"\"\"\n state = BuildStreamingState(turn_index=0)\n\n assert state.finalize_message_chunks() is None\n assert state.finalize_thought_chunks() is None\n" }, { "path": "backend/tests/external_dependency_unit/craft/test_file_upload.py", "content": "\"\"\"Tests for file upload functionality in build sessions.\n\nTests the file upload and delete operations for pre-provisioned sessions,\nincluding limit enforcement and SandboxManager delegation.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Generator\nfrom typing import TYPE_CHECKING\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import BuildSessionStatus\nfrom onyx.db.enums import SandboxStatus\nfrom onyx.db.models import BuildSession\nfrom onyx.db.models import Sandbox\nfrom onyx.db.models import User\nfrom onyx.server.features.build.configs import ATTACHMENTS_DIRECTORY\nfrom onyx.server.features.build.configs import MAX_TOTAL_UPLOAD_SIZE_BYTES\nfrom onyx.server.features.build.configs import MAX_UPLOAD_FILES_PER_SESSION\nfrom onyx.server.features.build.session.manager import UploadLimitExceededError\n\nif TYPE_CHECKING:\n from onyx.server.features.build.session.manager import SessionManager\n\n\n@pytest.fixture(scope=\"function\")\ndef sandbox(\n db_session: Session,\n test_user: User,\n tenant_context: None, # noqa: ARG001\n) -> Sandbox:\n \"\"\"Create a test sandbox for the user (sandboxes are per-user, not per-session).\"\"\"\n sandbox = Sandbox(\n id=uuid4(),\n user_id=test_user.id,\n status=SandboxStatus.RUNNING,\n )\n db_session.add(sandbox)\n db_session.commit()\n db_session.refresh(sandbox)\n return sandbox\n\n\n@pytest.fixture(scope=\"function\")\ndef build_session_with_user(\n db_session: Session,\n test_user: User,\n sandbox: Sandbox, # noqa: ARG001\n tenant_context: None, # noqa: ARG001\n) -> BuildSession:\n \"\"\"Create a test build session for a user who has a sandbox.\"\"\"\n session = BuildSession(\n id=uuid4(),\n user_id=test_user.id,\n name=\"Test Build Session\",\n status=BuildSessionStatus.ACTIVE,\n )\n db_session.add(session)\n db_session.commit()\n db_session.refresh(session)\n return session\n\n\n@pytest.fixture(scope=\"function\")\ndef mock_sandbox_manager() -> MagicMock:\n \"\"\"Create a mock sandbox manager.\"\"\"\n return MagicMock()\n\n\n@pytest.fixture(scope=\"function\")\ndef session_manager_with_mock(\n db_session: Session, mock_sandbox_manager: MagicMock\n) -> Generator[\"SessionManager\", None, None]:\n \"\"\"Create a SessionManager with mocked sandbox manager.\"\"\"\n # Import here to avoid module-level initialization issues\n with patch(\n \"onyx.server.features.build.session.manager.get_sandbox_manager\",\n return_value=mock_sandbox_manager,\n ):\n from onyx.server.features.build.session.manager import SessionManager\n\n manager = SessionManager(db_session)\n yield manager\n\n\nclass TestFileUpload:\n \"\"\"Tests for file upload functionality.\"\"\"\n\n def test_upload_file_delegates_to_sandbox_manager(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox,\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that uploading a file delegates to the sandbox manager.\"\"\"\n # Configure mocks\n mock_sandbox_manager.get_upload_stats.return_value = (0, 0)\n mock_sandbox_manager.upload_file.return_value = (\n f\"{ATTACHMENTS_DIRECTORY}/test.txt\"\n )\n\n # Upload a file\n content = b\"Hello, World!\"\n relative_path, size = session_manager_with_mock.upload_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n filename=\"test.txt\",\n content=content,\n )\n\n # Verify the sandbox manager was called correctly\n mock_sandbox_manager.upload_file.assert_called_once_with(\n sandbox_id=sandbox.id,\n session_id=build_session_with_user.id,\n filename=\"test.txt\",\n content=content,\n )\n assert relative_path == f\"{ATTACHMENTS_DIRECTORY}/test.txt\"\n assert size == len(content)\n\n def test_upload_file_returns_correct_path(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that upload returns the correct relative path.\"\"\"\n mock_sandbox_manager.get_upload_stats.return_value = (0, 0)\n mock_sandbox_manager.upload_file.return_value = (\n f\"{ATTACHMENTS_DIRECTORY}/document.pdf\"\n )\n\n relative_path, size = session_manager_with_mock.upload_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n filename=\"document.pdf\",\n content=b\"PDF content\",\n )\n\n assert relative_path == f\"{ATTACHMENTS_DIRECTORY}/document.pdf\"\n assert size == 11 # len(\"PDF content\")\n\n def test_upload_file_session_not_found(\n self,\n test_user: User,\n sandbox: Sandbox, # noqa: ARG002\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that uploading to a non-existent session raises ValueError.\"\"\"\n with pytest.raises(ValueError, match=\"Session not found\"):\n session_manager_with_mock.upload_file(\n session_id=uuid4(), # Non-existent session\n user_id=test_user.id,\n filename=\"test.txt\",\n content=b\"content\",\n )\n\n\nclass TestFileUploadLimits:\n \"\"\"Tests for file upload limit enforcement.\"\"\"\n\n def test_upload_file_count_limit_enforced(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that exceeding the file count limit raises an error.\"\"\"\n # Mock get_upload_stats to return max files already uploaded\n mock_sandbox_manager.get_upload_stats.return_value = (\n MAX_UPLOAD_FILES_PER_SESSION,\n 1000,\n )\n\n # Try to upload one more file\n with pytest.raises(UploadLimitExceededError, match=\"Maximum number of files\"):\n session_manager_with_mock.upload_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n filename=\"one_too_many.txt\",\n content=b\"content\",\n )\n\n # Verify upload_file was NOT called (limit check happens before)\n mock_sandbox_manager.upload_file.assert_not_called()\n\n def test_upload_total_size_limit_enforced(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that exceeding the total size limit raises an error.\"\"\"\n # Mock get_upload_stats to return almost at the limit\n existing_size = MAX_TOTAL_UPLOAD_SIZE_BYTES - 100 # 100 bytes under limit\n mock_sandbox_manager.get_upload_stats.return_value = (1, existing_size)\n\n # Try to upload a file that would exceed the limit\n with pytest.raises(UploadLimitExceededError, match=\"Total upload size limit\"):\n session_manager_with_mock.upload_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n filename=\"over_limit.txt\",\n content=b\"x\" * 200, # 200 bytes, would exceed by 100\n )\n\n # Verify upload_file was NOT called (limit check happens before)\n mock_sandbox_manager.upload_file.assert_not_called()\n\n def test_upload_succeeds_when_under_limits(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that upload succeeds when under limits.\"\"\"\n # Mock get_upload_stats to return well under limits\n mock_sandbox_manager.get_upload_stats.return_value = (5, 1000)\n mock_sandbox_manager.upload_file.return_value = (\n f\"{ATTACHMENTS_DIRECTORY}/test.txt\"\n )\n\n relative_path, size = session_manager_with_mock.upload_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n filename=\"test.txt\",\n content=b\"content\",\n )\n\n # Verify upload_file was called\n mock_sandbox_manager.upload_file.assert_called_once()\n assert relative_path == f\"{ATTACHMENTS_DIRECTORY}/test.txt\"\n\n\nclass TestFileDelete:\n \"\"\"Tests for file delete functionality.\"\"\"\n\n def test_delete_file_delegates_to_sandbox_manager(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox,\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that delete file delegates to the sandbox manager.\"\"\"\n mock_sandbox_manager.delete_file.return_value = True\n\n result = session_manager_with_mock.delete_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n path=f\"{ATTACHMENTS_DIRECTORY}/test.txt\",\n )\n\n assert result is True\n mock_sandbox_manager.delete_file.assert_called_once_with(\n sandbox_id=sandbox.id,\n session_id=build_session_with_user.id,\n path=f\"{ATTACHMENTS_DIRECTORY}/test.txt\",\n )\n\n def test_delete_file_returns_false_when_not_found(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that delete returns False when file doesn't exist.\"\"\"\n mock_sandbox_manager.delete_file.return_value = False\n\n result = session_manager_with_mock.delete_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n path=f\"{ATTACHMENTS_DIRECTORY}/nonexistent.txt\",\n )\n\n assert result is False\n\n def test_delete_file_session_not_found(\n self,\n test_user: User,\n sandbox: Sandbox, # noqa: ARG002\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that deleting from a non-existent session raises ValueError.\"\"\"\n with pytest.raises(ValueError, match=\"Session not found\"):\n session_manager_with_mock.delete_file(\n session_id=uuid4(), # Non-existent session\n user_id=test_user.id,\n path=f\"{ATTACHMENTS_DIRECTORY}/test.txt\",\n )\n\n\nclass TestPathSanitization:\n \"\"\"Tests for path sanitization in delete operations.\"\"\"\n\n def test_delete_file_rejects_path_traversal(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that paths with .. are rejected.\"\"\"\n # Configure mock to raise ValueError (simulating sandbox manager behavior)\n mock_sandbox_manager.delete_file.side_effect = ValueError(\n \"Invalid path: potential path traversal detected\"\n )\n\n with pytest.raises(ValueError, match=\"path traversal\"):\n session_manager_with_mock.delete_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n path=\"../../../etc/passwd\",\n )\n\n def test_delete_file_rejects_url_encoded_traversal(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that URL-encoded paths are rejected.\"\"\"\n mock_sandbox_manager.delete_file.side_effect = ValueError(\n \"Invalid path: potential path traversal detected\"\n )\n\n with pytest.raises(ValueError, match=\"path traversal\"):\n session_manager_with_mock.delete_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n path=\"attachments/%2e%2e/secret.txt\",\n )\n\n def test_delete_file_rejects_shell_metacharacters(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that shell metacharacters are rejected.\"\"\"\n mock_sandbox_manager.delete_file.side_effect = ValueError(\n \"Invalid path: contains disallowed characters\"\n )\n\n dangerous_paths = [\n \"attachments/file;rm -rf /\",\n \"attachments/file|cat /etc/passwd\",\n \"attachments/file`whoami`\",\n \"attachments/file$(id)\",\n \"attachments/file'test\",\n ]\n\n for dangerous_path in dangerous_paths:\n with pytest.raises(ValueError, match=\"disallowed characters\"):\n session_manager_with_mock.delete_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n path=dangerous_path,\n )\n # Reset mock for next iteration\n mock_sandbox_manager.delete_file.reset_mock()\n\n def test_delete_file_rejects_null_bytes(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that null bytes in paths are rejected.\"\"\"\n mock_sandbox_manager.delete_file.side_effect = ValueError(\n \"Invalid path: potential path traversal detected\"\n )\n\n with pytest.raises(ValueError, match=\"path traversal\"):\n session_manager_with_mock.delete_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n path=\"attachments/file.txt\\x00.jpg\",\n )\n\n\nclass TestFilenameCollision:\n \"\"\"Tests for filename collision handling.\"\"\"\n\n def test_upload_returns_collision_handled_path(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox, # noqa: ARG002\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that sandbox manager can return a renamed path for collisions.\"\"\"\n # Simulate sandbox manager handling collision by returning renamed path\n mock_sandbox_manager.get_upload_stats.return_value = (1, 100) # 1 existing file\n mock_sandbox_manager.upload_file.return_value = (\n f\"{ATTACHMENTS_DIRECTORY}/document_1.pdf\"\n )\n\n relative_path, size = session_manager_with_mock.upload_file(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n filename=\"document.pdf\",\n content=b\"PDF content\",\n )\n\n # Verify the collision-handled path is returned\n assert relative_path == f\"{ATTACHMENTS_DIRECTORY}/document_1.pdf\"\n\n\nclass TestGetUploadStats:\n \"\"\"Tests for get_upload_stats functionality.\"\"\"\n\n def test_get_upload_stats_delegates_to_sandbox_manager(\n self,\n test_user: User,\n build_session_with_user: BuildSession,\n sandbox: Sandbox,\n mock_sandbox_manager: MagicMock,\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that get_upload_stats delegates to the sandbox manager.\"\"\"\n mock_sandbox_manager.get_upload_stats.return_value = (3, 1500)\n\n file_count, total_size = session_manager_with_mock.get_upload_stats(\n session_id=build_session_with_user.id,\n user_id=test_user.id,\n )\n\n # Verify the sandbox manager was called correctly\n mock_sandbox_manager.get_upload_stats.assert_called_once_with(\n sandbox_id=sandbox.id,\n session_id=build_session_with_user.id,\n )\n assert file_count == 3\n assert total_size == 1500\n\n def test_get_upload_stats_session_not_found(\n self,\n test_user: User,\n sandbox: Sandbox, # noqa: ARG002\n session_manager_with_mock: \"SessionManager\",\n ) -> None:\n \"\"\"Test that getting stats for non-existent session raises ValueError.\"\"\"\n with pytest.raises(ValueError, match=\"Session not found\"):\n session_manager_with_mock.get_upload_stats(\n session_id=uuid4(), # Non-existent session\n user_id=test_user.id,\n )\n" }, { "path": "backend/tests/external_dependency_unit/craft/test_kubernetes_sandbox.py", "content": "\"\"\"Integration test for KubernetesSandboxManager.provision().\n\nThis test requires:\n- A running Kubernetes cluster (kind, minikube, or real cluster)\n- The SANDBOX_BACKEND=kubernetes environment variable\n- The sandbox namespace to exist (default: onyx-sandboxes)\n- Service accounts for sandbox (sandbox-runner, sandbox-file-sync)\n\nRun with:\n SANDBOX_BACKEND=kubernetes python -m dotenv -f .vscode/.env run -- \\\n pytest backend/tests/integration/tests/build/test_kubernetes_sandbox_provision.py -v\n\"\"\"\n\nimport time\nfrom uuid import UUID\nfrom uuid import uuid4\n\nimport pytest\nfrom kubernetes import client # type: ignore[import-untyped]\nfrom kubernetes import config\nfrom kubernetes.client.rest import ApiException # type: ignore[import-untyped]\nfrom kubernetes.stream import stream as k8s_stream # type: ignore[import-untyped]\n\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.enums import SandboxStatus\nfrom onyx.server.features.build.configs import SANDBOX_BACKEND\nfrom onyx.server.features.build.configs import SANDBOX_NAMESPACE\nfrom onyx.server.features.build.configs import SANDBOX_NEXTJS_PORT_START\nfrom onyx.server.features.build.configs import SandboxBackend\nfrom onyx.server.features.build.sandbox.base import ACPEvent\nfrom onyx.server.features.build.sandbox.kubernetes.kubernetes_sandbox_manager import (\n KubernetesSandboxManager,\n)\nfrom onyx.server.features.build.sandbox.models import LLMProviderConfig\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\nlogger = setup_logger()\n\n# Test constants\nTEST_TENANT_ID = \"test-tenant\"\nTEST_USER_ID = UUID(\"ee0dd46a-23dc-4128-abab-6712b3f4464c\")\n\n\ndef _is_kubernetes_available() -> None:\n \"\"\"Check if Kubernetes is available and configured.\"\"\"\n try:\n config.load_incluster_config()\n except config.ConfigException:\n config.load_kube_config()\n\n v1 = client.CoreV1Api()\n # List pods in sandbox namespace instead of namespaces (avoids cluster-scope permissions)\n v1.list_namespaced_pod(SANDBOX_NAMESPACE, limit=1)\n\n\ndef _get_kubernetes_client() -> client.CoreV1Api:\n \"\"\"Get a configured Kubernetes CoreV1Api client.\"\"\"\n try:\n config.load_incluster_config()\n except config.ConfigException:\n config.load_kube_config()\n return client.CoreV1Api()\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_kubernetes_sandbox_provision() -> None:\n \"\"\"Test that provision() creates a sandbox pod and DB record successfully.\n\n This is a happy path test that:\n 1. Creates a BuildSession in the database\n 2. Calls provision() to create a Kubernetes pod\n 3. Verifies the sandbox is created with RUNNING status\n 4. Cleans up by terminating the sandbox\n \"\"\"\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context (required for multi-tenant operations)\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n # Get the manager instance\n manager = KubernetesSandboxManager()\n\n sandbox_id = uuid4()\n\n # Create a test LLM config (values don't matter for this test)\n llm_config = LLMProviderConfig(\n provider=\"openai\",\n model_name=\"gpt-4\",\n api_key=\"test-key\",\n api_base=None,\n )\n\n try:\n # Call provision\n sandbox_info = manager.provision(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n llm_config=llm_config,\n )\n\n # Verify the return value\n assert sandbox_info.sandbox_id == sandbox_id\n assert sandbox_info.status == SandboxStatus.RUNNING\n assert sandbox_info.directory_path.startswith(\"k8s://\")\n\n # Verify Kubernetes resources exist\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n service_name = pod_name\n\n # Verify pod exists and is running\n pod = k8s_client.read_namespaced_pod(\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n )\n assert pod is not None\n assert pod.status.phase == \"Running\"\n\n # Verify service exists\n service = k8s_client.read_namespaced_service(\n name=service_name,\n namespace=SANDBOX_NAMESPACE,\n )\n assert service is not None\n assert service.spec.type == \"ClusterIP\"\n\n # Verify /workspace/templates/outputs directory exists and contains expected files\n exec_command = [\"/bin/sh\", \"-c\", \"ls -la /workspace/templates/outputs\"]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None\n print(f\"DEBUG: Contents of /workspace/templates/outputs:\\n{resp}\")\n assert (\n \"web\" in resp\n ), f\"/workspace/templates/outputs should contain web directory. Actual contents:\\n{resp}\"\n\n # Verify /workspace/templates/outputs/web/AGENTS.md file exists\n exec_command = [\n \"/bin/sh\",\n \"-c\",\n \"cat /workspace/templates/outputs/web/AGENTS.md\",\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None\n assert (\n len(resp) > 0\n ), \"/workspace/templates/outputs/web/AGENTS.md file should not be empty\"\n # Verify it contains expected content\n assert (\n \"Agent\" in resp or \"Instructions\" in resp or \"#\" in resp\n ), \"/workspace/templates/outputs/web/AGENTS.md should contain agent instructions\"\n\n # Verify /workspace/files directory exists and contains expected files\n exec_command = [\"/bin/sh\", \"-c\", \"find /workspace/files -type f | wc -l\"]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None\n file_count = int(resp.strip())\n assert (\n file_count == 1099\n ), f\"/workspace/files should contain 1099 files, but found {file_count}\"\n\n # start session\n session_id = uuid4()\n manager.setup_session_workspace(\n sandbox_id=sandbox_id,\n session_id=session_id,\n llm_config=llm_config,\n nextjs_port=SANDBOX_NEXTJS_PORT_START,\n file_system_path=None,\n snapshot_path=None,\n user_name=\"Test User\",\n user_role=\"Test Role\",\n )\n\n # Verify AGENTS.md file exists for the session\n exec_command = [\n \"/bin/sh\",\n \"-c\",\n f\"cat /workspace/sessions/{session_id}/AGENTS.md\",\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None\n assert len(resp) > 0, \"AGENTS.md file should not be empty\"\n # Verify it contains expected content (from template or default)\n assert \"Agent\" in resp or \"Instructions\" in resp or \"#\" in resp\n assert \"Test User\" in resp\n assert \"Test Role\" in resp\n\n # Verify opencode.json file exists for the session\n exec_command = [\n \"/bin/sh\",\n \"-c\",\n f\"cat /workspace/sessions/{session_id}/opencode.json\",\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None\n assert len(resp) > 0, \"opencode.json file should not be empty\"\n\n # verify that the outputs directory is copied over\n exec_command = [\n \"/bin/sh\",\n \"-c\",\n f\"ls -la /workspace/sessions/{session_id}/outputs\",\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None\n assert len(resp) > 0, \"outputs directory should not be empty\"\n assert \"web\" in resp, \"outputs directory should contain web directory\"\n\n finally:\n # Clean up: terminate the sandbox (no longer needs db_session)\n if sandbox_id:\n manager.terminate(sandbox_id)\n\n # Verify Kubernetes resources are cleaned up\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n\n # Give K8s a moment to delete resources\n time.sleep(2)\n\n # Verify pod is deleted (or being deleted)\n try:\n pod = k8s_client.read_namespaced_pod(\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n )\n # Pod might still exist but be terminating\n assert pod.metadata.deletion_timestamp is not None\n except ApiException as e:\n # 404 means pod was successfully deleted\n assert e.status == 404\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_kubernetes_sandbox_send_message() -> None:\n \"\"\"Test that send_message() communicates with the sandbox agent successfully.\n\n This test:\n 1. Creates a sandbox pod\n 2. Sends a simple message via send_message()\n 3. Verifies we receive ACP events back (agent responses)\n 4. Cleans up by terminating the sandbox\n \"\"\"\n from acp.schema import AgentMessageChunk\n from acp.schema import Error\n from acp.schema import PromptResponse\n\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context (required for multi-tenant operations)\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n # Get the manager instance\n manager = KubernetesSandboxManager()\n\n sandbox_id = uuid4()\n session_id = uuid4()\n\n # Create a test LLM config (values don't matter for this test)\n llm_config = LLMProviderConfig(\n provider=\"openai\",\n model_name=\"gpt-4\",\n api_key=\"test-key\",\n api_base=None,\n )\n\n try:\n # Provision the sandbox\n sandbox_info = manager.provision(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n llm_config=llm_config,\n )\n\n assert sandbox_info.status == SandboxStatus.RUNNING\n\n # Verify health check passes before sending message\n is_healthy = False\n for _ in range(10):\n is_healthy = manager.health_check(sandbox_id)\n if is_healthy:\n break\n time.sleep(10)\n\n assert is_healthy, \"Sandbox agent should be healthy before sending messages\"\n print(\"DEBUG: Sandbox agent is healthy\")\n\n manager.setup_session_workspace(\n sandbox_id, session_id, llm_config, nextjs_port=SANDBOX_NEXTJS_PORT_START\n )\n\n # Send a simple message\n events: list[ACPEvent] = []\n for event in manager.send_message(sandbox_id, session_id, \"What is 2 + 2?\"):\n events.append(event)\n\n # Verify we received events\n assert len(events) > 0, \"Should receive at least one event from send_message\"\n\n for event in events:\n print(f\"Recieved event: {event}\")\n\n # Check for errors\n errors = [e for e in events if isinstance(e, Error)]\n assert len(errors) == 0, f\"Should not receive errors: {errors}\"\n\n # Verify we received some agent message content or a final response\n message_chunks = [e for e in events if isinstance(e, AgentMessageChunk)]\n prompt_responses = [e for e in events if isinstance(e, PromptResponse)]\n\n assert (\n len(message_chunks) > 0 or len(prompt_responses) > 0\n ), \"Should receive either AgentMessageChunk or PromptResponse events\"\n\n # If we got a PromptResponse, verify it completed successfully\n if prompt_responses:\n final_response = prompt_responses[-1]\n assert (\n final_response.stop_reason is not None\n ), \"PromptResponse should have a stop_reason\"\n\n finally:\n # Clean up: terminate the sandbox\n if sandbox_id:\n manager.terminate(sandbox_id)\n\n # Verify Kubernetes resources are cleaned up\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n\n # Give K8s a moment to delete resources\n time.sleep(2)\n\n # Verify pod is deleted (or being deleted)\n try:\n pod = k8s_client.read_namespaced_pod(\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n )\n # Pod might still exist but be terminating\n assert pod.metadata.deletion_timestamp is not None\n except ApiException as e:\n # 404 means pod was successfully deleted\n assert e.status == 404\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_kubernetes_sandbox_webapp_passthrough() -> None:\n \"\"\"Test that the webapp passthrough (Next.js server) is accessible in the sandbox.\n\n This test:\n 1. Creates a sandbox pod\n 2. Sets up a session workspace\n 3. Verifies the Next.js server is running and accessible within the pod\n 4. Verifies get_nextjs_url returns the correct cluster URL format\n 5. Cleans up by terminating the sandbox\n \"\"\"\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context (required for multi-tenant operations)\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n # Get the manager instance\n manager = KubernetesSandboxManager()\n\n sandbox_id = uuid4()\n session_id = uuid4()\n\n # Create a test LLM config\n llm_config = LLMProviderConfig(\n provider=\"openai\",\n model_name=\"gpt-4\",\n api_key=\"test-key\",\n api_base=None,\n )\n\n try:\n # Provision the sandbox\n sandbox_info = manager.provision(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n llm_config=llm_config,\n )\n\n assert sandbox_info.status == SandboxStatus.RUNNING\n\n # Verify health check passes before testing webapp\n is_healthy = False\n for _ in range(10):\n is_healthy = manager.health_check(sandbox_id)\n if is_healthy:\n break\n time.sleep(10)\n\n assert is_healthy, \"Sandbox should be healthy before testing webapp passthrough\"\n print(\"DEBUG: Sandbox is healthy\")\n\n # Set up session workspace\n manager.setup_session_workspace(\n sandbox_id=sandbox_id,\n session_id=session_id,\n llm_config=llm_config,\n nextjs_port=SANDBOX_NEXTJS_PORT_START,\n file_system_path=None,\n snapshot_path=None,\n user_name=\"Test User\",\n user_role=\"Test Role\",\n )\n\n # Get Kubernetes client for exec operations\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n\n # Wait for Next.js server to be ready (it may take a few seconds to start)\n # The session uses the first port in the configured range\n test_nextjs_port = SANDBOX_NEXTJS_PORT_START\n nextjs_ready = False\n for attempt in range(30):\n exec_command = [\n \"/bin/sh\",\n \"-c\",\n (\n f\"curl -s -o /dev/null -w '%{{http_code}}' http://localhost:{test_nextjs_port}/ 2>/dev/null || echo 'failed'\"\n ),\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n print(f\"DEBUG: Next.js health check attempt {attempt + 1}: {resp}\")\n if resp and resp.strip() in (\"200\", \"304\"):\n nextjs_ready = True\n break\n time.sleep(2)\n\n assert (\n nextjs_ready\n ), f\"Next.js server should be accessible at localhost:{SANDBOX_NEXTJS_PORT_START}\"\n print(\"DEBUG: Next.js server is ready\")\n\n # Verify we can fetch actual content from the Next.js server\n exec_command = [\n \"/bin/sh\",\n \"-c\",\n f\"curl -s http://localhost:{SANDBOX_NEXTJS_PORT_START}/ | head -c 500\",\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None, \"Should receive content from Next.js server\"\n assert len(resp) > 0, \"Next.js server response should not be empty\"\n # Basic check that it looks like HTML\n assert (\n \"<\" in resp or \"html\" in resp.lower() or \"/dev/null || echo 'failed'\",\n ]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n print(f\"DEBUG: Cluster URL health check response: {resp}\")\n assert resp and resp.strip() in (\n \"200\",\n \"304\",\n ), f\"Next.js server should be accessible via cluster URL {nextjs_url}. Got response: {resp}\"\n\n finally:\n # Clean up: terminate the sandbox\n if sandbox_id:\n manager.terminate(sandbox_id)\n\n # Verify Kubernetes resources are cleaned up\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n\n # Give K8s a moment to delete resources\n time.sleep(2)\n\n # Verify pod is deleted (or being deleted)\n try:\n pod = k8s_client.read_namespaced_pod(\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n )\n # Pod might still exist but be terminating\n assert pod.metadata.deletion_timestamp is not None\n except ApiException as e:\n # 404 means pod was successfully deleted\n assert e.status == 404\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_kubernetes_sandbox_file_sync() -> None:\n \"\"\"Test that sync_files() triggers S3 sync in the file-sync sidecar.\n\n This test:\n 1. Creates a sandbox pod (which now has file-sync as sidecar)\n 2. Verifies the file-sync sidecar is running\n 3. Calls sync_files() to trigger S3 sync\n 4. Verifies the sync command executes successfully\n 5. Cleans up by terminating the sandbox\n \"\"\"\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context (required for multi-tenant operations)\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n # Get the manager instance\n manager = KubernetesSandboxManager()\n\n sandbox_id = uuid4()\n\n # Create a test LLM config\n llm_config = LLMProviderConfig(\n provider=\"openai\",\n model_name=\"gpt-4\",\n api_key=\"test-key\",\n api_base=None,\n )\n\n try:\n # Provision the sandbox\n sandbox_info = manager.provision(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n llm_config=llm_config,\n )\n\n assert sandbox_info.status == SandboxStatus.RUNNING\n\n # Verify the pod is running\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n pod = k8s_client.read_namespaced_pod(\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n )\n assert pod is not None\n assert pod.status.phase == \"Running\"\n\n # Verify file-sync sidecar container is running\n # With sidecar model, file-sync should be a regular container (not init)\n container_statuses = pod.status.container_statuses or []\n file_sync_status = next(\n (c for c in container_statuses if c.name == \"file-sync\"),\n None,\n )\n assert file_sync_status is not None, \"file-sync sidecar container should exist\"\n assert file_sync_status.ready, \"file-sync sidecar container should be ready\"\n print(f\"DEBUG: file-sync container status: {file_sync_status}\")\n\n # Wipe the /workspace/files directory to ensure files we find are from the sync\n exec_command = [\"/bin/sh\", \"-c\", \"rm -rf /workspace/files/*\"]\n k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"file-sync\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n print(\"DEBUG: Wiped /workspace/files directory\")\n\n # Verify the directory is empty\n exec_command = [\"/bin/sh\", \"-c\", \"find /workspace/files -type f | wc -l\"]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n file_count = int(resp.strip()) if resp else 0\n assert (\n file_count == 0\n ), f\"/workspace/files should be empty before sync, found {file_count} files\"\n print(\"DEBUG: Verified /workspace/files is empty\")\n\n # Call sync_files() to trigger S3 sync\n result = manager.sync_files(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n )\n assert result is True, \"sync_files() should return True on success\"\n print(\"DEBUG: sync_files() completed successfully\")\n\n # Verify /workspace/files exists and has files synced from S3\n # (verifies the shared volume is working and sync actually transferred files)\n exec_command = [\"/bin/sh\", \"-c\", \"find /workspace/files -type f | wc -l\"]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"sandbox\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None, \"/workspace/files should be accessible from sandbox\"\n file_count = int(resp.strip()) if resp else 0\n assert (\n file_count > 0\n ), f\"sync_files() should have synced files, but found {file_count} files\"\n print(f\"DEBUG: sync_files() synced {file_count} files to /workspace/files\")\n\n # Also verify we can exec into file-sync sidecar directly\n exec_command = [\"/bin/sh\", \"-c\", \"ls -la /workspace/files\"]\n resp = k8s_stream(\n k8s_client.connect_get_namespaced_pod_exec,\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n container=\"file-sync\",\n command=exec_command,\n stderr=True,\n stdin=False,\n stdout=True,\n tty=False,\n )\n assert resp is not None, \"/workspace/files should be accessible from file-sync\"\n print(f\"DEBUG: Contents of /workspace/files (from file-sync sidecar):\\n{resp}\")\n\n finally:\n # Clean up: terminate the sandbox\n if sandbox_id:\n manager.terminate(sandbox_id)\n\n # Verify Kubernetes resources are cleaned up\n k8s_client = _get_kubernetes_client()\n pod_name = f\"sandbox-{str(sandbox_id)[:8]}\"\n\n # Give K8s a moment to delete resources\n time.sleep(2)\n\n # Verify pod is deleted (or being deleted)\n try:\n pod = k8s_client.read_namespaced_pod(\n name=pod_name,\n namespace=SANDBOX_NAMESPACE,\n )\n # Pod might still exist but be terminating\n assert pod.metadata.deletion_timestamp is not None\n except ApiException as e:\n # 404 means pod was successfully deleted\n assert e.status == 404\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_health_check_returns_true_for_running_pod() -> None:\n \"\"\"Test that health_check() returns True for a healthy, running pod.\n\n This test:\n 1. Creates a sandbox pod\n 2. Calls health_check() and verifies it returns True\n 3. Cleans up by terminating the sandbox\n \"\"\"\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n manager = KubernetesSandboxManager()\n sandbox_id = uuid4()\n\n llm_config = LLMProviderConfig(\n provider=\"openai\",\n model_name=\"gpt-4\",\n api_key=\"test-key\",\n api_base=None,\n )\n\n try:\n # Provision the sandbox\n sandbox_info = manager.provision(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n llm_config=llm_config,\n )\n\n assert sandbox_info.status == SandboxStatus.RUNNING\n\n # Wait for pod to be fully healthy (it may take a few seconds)\n is_healthy = False\n for _ in range(10):\n is_healthy = manager.health_check(sandbox_id, timeout=5.0)\n if is_healthy:\n break\n time.sleep(2)\n\n assert (\n is_healthy\n ), \"health_check() should return True for a running, healthy pod\"\n\n finally:\n if sandbox_id:\n manager.terminate(sandbox_id)\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_health_check_returns_false_for_missing_pod() -> None:\n \"\"\"Test that health_check() returns False when the pod doesn't exist.\n\n This test:\n 1. Uses a random UUID that has no corresponding pod\n 2. Calls health_check() and verifies it returns False\n \"\"\"\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n manager = KubernetesSandboxManager()\n\n # Use a random UUID that definitely has no pod\n nonexistent_sandbox_id = uuid4()\n\n # health_check should return False for non-existent pod\n is_healthy = manager.health_check(nonexistent_sandbox_id, timeout=5.0)\n\n assert not is_healthy, \"health_check() should return False for a non-existent pod\"\n\n\n@pytest.mark.skipif(\n SANDBOX_BACKEND != SandboxBackend.KUBERNETES,\n reason=\"SANDBOX_BACKEND must be 'kubernetes' to run this test\",\n)\ndef test_health_check_returns_false_after_termination() -> None:\n \"\"\"Test that health_check() returns False after a pod has been terminated.\n\n This test:\n 1. Creates a sandbox pod\n 2. Verifies health_check() returns True\n 3. Terminates the sandbox\n 4. Verifies health_check() returns False\n \"\"\"\n _is_kubernetes_available()\n\n # Initialize the database engine\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Set up tenant context\n CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n\n manager = KubernetesSandboxManager()\n sandbox_id = uuid4()\n\n llm_config = LLMProviderConfig(\n provider=\"openai\",\n model_name=\"gpt-4\",\n api_key=\"test-key\",\n api_base=None,\n )\n\n # Provision the sandbox\n sandbox_info = manager.provision(\n sandbox_id=sandbox_id,\n user_id=TEST_USER_ID,\n tenant_id=TEST_TENANT_ID,\n llm_config=llm_config,\n )\n\n assert sandbox_info.status == SandboxStatus.RUNNING\n\n # Wait for pod to be fully healthy\n is_healthy = False\n for _ in range(10):\n is_healthy = manager.health_check(sandbox_id, timeout=5.0)\n if is_healthy:\n break\n time.sleep(2)\n\n assert is_healthy, \"Pod should be healthy before termination\"\n\n # Terminate the sandbox\n manager.terminate(sandbox_id)\n\n # Wait for pod to be deleted\n time.sleep(3)\n\n # health_check should now return False\n is_healthy_after = manager.health_check(sandbox_id, timeout=5.0)\n\n assert (\n not is_healthy_after\n ), \"health_check() should return False after pod has been terminated\"\n" }, { "path": "backend/tests/external_dependency_unit/craft/test_persistent_document_writer.py", "content": "\"\"\"\nTests for PersistentDocumentWriter (local) and S3PersistentDocumentWriter.\n\nRun with:\n python -m dotenv -f .vscode/.env run -- \\\n pytest backend/tests/external_dependency_unit/craft/test_persistent_document_writer.py -v\n\"\"\"\n\nimport json\nimport os\nimport tempfile\nfrom datetime import datetime\nfrom datetime import timezone\nfrom uuid import uuid4\n\nimport boto3\nimport pytest\nfrom botocore.exceptions import ClientError\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import TextSection\nfrom onyx.server.features.build.configs import SANDBOX_S3_BUCKET\nfrom onyx.server.features.build.indexing.persistent_document_writer import (\n PersistentDocumentWriter,\n)\nfrom onyx.server.features.build.indexing.persistent_document_writer import (\n S3PersistentDocumentWriter,\n)\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\ndef _create_test_document(doc_id: str, name: str) -> Document:\n \"\"\"Helper to create a test document.\"\"\"\n return Document(\n id=doc_id,\n semantic_identifier=name,\n title=name,\n source=DocumentSource.WEB,\n sections=[TextSection(text=\"Test content\", link=\"https://example.com\")],\n metadata={},\n doc_metadata={\"hierarchy\": {\"source_path\": [\"Folder\"]}},\n doc_updated_at=datetime.now(timezone.utc),\n primary_owners=[],\n secondary_owners=[],\n )\n\n\ndef test_local_persistent_document_writer() -> None:\n \"\"\"Test writing documents to local filesystem.\"\"\"\n with tempfile.TemporaryDirectory() as temp_dir:\n tenant_id = TEST_TENANT_ID\n user_id = str(uuid4())\n writer = PersistentDocumentWriter(\n base_path=temp_dir, tenant_id=tenant_id, user_id=user_id\n )\n\n doc = _create_test_document(\"doc-001\", \"Test Document\")\n written_paths = writer.write_documents([doc])\n\n assert len(written_paths) == 1\n assert written_paths[0] == os.path.join(\n temp_dir,\n tenant_id,\n \"knowledge\",\n user_id,\n \"web\",\n \"Folder\",\n \"Test_Document.json\",\n )\n assert os.path.exists(written_paths[0])\n\n with open(written_paths[0]) as f:\n content = json.load(f)\n assert content[\"id\"] == \"doc-001\"\n assert content[\"semantic_identifier\"] == \"Test Document\"\n\n\ndef _is_s3_available() -> bool:\n \"\"\"Check if S3 is available for testing.\"\"\"\n try:\n s3_client = boto3.client(\"s3\")\n s3_client.head_bucket(Bucket=SANDBOX_S3_BUCKET)\n return True\n except (ClientError, Exception):\n return False\n\n\n@pytest.mark.skipif(\n not _is_s3_available(),\n reason=f\"S3 bucket '{SANDBOX_S3_BUCKET}' not available\",\n)\ndef test_s3_persistent_document_writer() -> None:\n \"\"\"Test writing documents to S3.\"\"\"\n user_id = str(uuid4())\n writer = S3PersistentDocumentWriter(tenant_id=TEST_TENANT_ID, user_id=user_id)\n\n doc = _create_test_document(\"s3-doc-001\", \"S3 Test Doc\")\n written_keys = writer.write_documents([doc])\n\n try:\n assert len(written_keys) == 1\n assert f\"{TEST_TENANT_ID}/knowledge/{user_id}\" in written_keys[0]\n\n # Verify the object exists in S3\n s3_client = boto3.client(\"s3\")\n response = s3_client.get_object(Bucket=SANDBOX_S3_BUCKET, Key=written_keys[0])\n content = json.loads(response[\"Body\"].read().decode(\"utf-8\"))\n\n assert content[\"id\"] == \"s3-doc-001\"\n assert content[\"semantic_identifier\"] == \"S3 Test Doc\"\n finally:\n # Cleanup\n s3_client = boto3.client(\"s3\")\n try:\n s3_client.delete_object(Bucket=SANDBOX_S3_BUCKET, Key=written_keys[0])\n except Exception:\n pass\n" }, { "path": "backend/tests/external_dependency_unit/db/__init__.py", "content": "" }, { "path": "backend/tests/external_dependency_unit/db/conftest.py", "content": "\"\"\"Fixtures for testing DAL classes against a real PostgreSQL database.\n\nThese fixtures build on the db_session and tenant_context fixtures from\nthe parent conftest (tests/external_dependency_unit/conftest.py).\n\nRequires a running Postgres instance. Run with::\n\n python -m dotenv -f .vscode/.env run -- pytest tests/external_dependency_unit/db/\n\"\"\"\n\nfrom collections.abc import Callable\nfrom collections.abc import Generator\nfrom uuid import UUID\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.scim import ScimDAL\nfrom onyx.db.models import ScimToken\nfrom onyx.db.models import UserGroup\n\n\n@pytest.fixture\ndef scim_dal(db_session: Session) -> ScimDAL:\n \"\"\"A ScimDAL backed by the real test database session.\"\"\"\n return ScimDAL(db_session)\n\n\n@pytest.fixture\ndef scim_token_factory(\n db_session: Session,\n) -> Generator[Callable[..., ScimToken], None, None]:\n \"\"\"Factory that creates ScimToken rows and cleans them up after the test.\"\"\"\n created_ids: list[int] = []\n\n def _create(\n name: str = \"test-token\",\n hashed_token: str | None = None,\n token_display: str = \"onyx_scim_****test\",\n created_by_id: UUID | None = None,\n ) -> ScimToken:\n token = ScimToken(\n name=name,\n hashed_token=hashed_token or uuid4().hex,\n token_display=token_display,\n created_by_id=created_by_id or uuid4(),\n )\n db_session.add(token)\n db_session.flush()\n created_ids.append(token.id)\n return token\n\n yield _create\n\n for token_id in created_ids:\n obj = db_session.get(ScimToken, token_id)\n if obj:\n db_session.delete(obj)\n db_session.commit()\n\n\n@pytest.fixture\ndef user_group_factory(\n db_session: Session,\n) -> Generator[Callable[..., UserGroup], None, None]:\n \"\"\"Factory that creates UserGroup rows for testing group mappings.\"\"\"\n created_ids: list[int] = []\n\n def _create(name: str | None = None) -> UserGroup:\n group = UserGroup(name=name or f\"test-group-{uuid4().hex[:8]}\")\n db_session.add(group)\n db_session.flush()\n created_ids.append(group.id)\n return group\n\n yield _create\n\n for group_id in created_ids:\n obj = db_session.get(UserGroup, group_id)\n if obj:\n db_session.delete(obj)\n db_session.commit()\n" }, { "path": "backend/tests/external_dependency_unit/db/test_chat_session_eager_load.py", "content": "from sqlalchemy import inspect\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.chat import create_chat_session\nfrom onyx.db.chat import get_chat_session_by_id\nfrom onyx.db.models import Persona\nfrom onyx.db.models import UserProject\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef test_eager_load_persona_loads_relationships(db_session: Session) -> None:\n \"\"\"Verify that eager_load_persona pre-loads persona, its collections, and project.\"\"\"\n user = create_test_user(db_session, \"eager-load\")\n persona = Persona(name=\"eager-load-test\", description=\"test\")\n project = UserProject(name=\"eager-load-project\", user_id=user.id)\n db_session.add_all([persona, project])\n db_session.flush()\n\n chat_session = create_chat_session(\n db_session=db_session,\n description=\"test\",\n user_id=None,\n persona_id=persona.id,\n project_id=project.id,\n )\n\n loaded = get_chat_session_by_id(\n chat_session_id=chat_session.id,\n user_id=None,\n db_session=db_session,\n eager_load_persona=True,\n )\n\n try:\n tmp = inspect(loaded)\n assert tmp is not None\n unloaded = tmp.unloaded\n assert \"persona\" not in unloaded\n assert \"project\" not in unloaded\n\n tmp = inspect(loaded.persona)\n assert tmp is not None\n persona_unloaded = tmp.unloaded\n assert \"tools\" not in persona_unloaded\n assert \"user_files\" not in persona_unloaded\n assert \"document_sets\" not in persona_unloaded\n assert \"attached_documents\" not in persona_unloaded\n assert \"hierarchy_nodes\" not in persona_unloaded\n finally:\n db_session.rollback()\n" }, { "path": "backend/tests/external_dependency_unit/db/test_credential_sensitive_value.py", "content": "\"\"\"Test that Credential with nested JSON round-trips through SensitiveValue correctly.\n\nExercises the full encrypt → store → read → decrypt → SensitiveValue path\nwith realistic nested OAuth credential data, and verifies SQLAlchemy dirty\ntracking works with nested dict comparison.\n\nRequires a running Postgres instance.\n\"\"\"\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import Credential\nfrom onyx.utils.sensitive import SensitiveValue\n\n# NOTE: this is not the real shape of a Drive credential,\n# but it is intended to test nested JSON credential handling\n\n_NESTED_CRED_JSON = {\n \"oauth_tokens\": {\n \"access_token\": \"ya29.abc123\",\n \"refresh_token\": \"1//xEg-def456\",\n },\n \"scopes\": [\"read\", \"write\", \"admin\"],\n \"client_config\": {\n \"client_id\": \"123.apps.googleusercontent.com\",\n \"client_secret\": \"GOCSPX-secret\",\n },\n}\n\n\ndef test_nested_credential_json_round_trip(db_session: Session) -> None:\n \"\"\"Nested OAuth credential survives encrypt → store → read → decrypt.\"\"\"\n credential = Credential(\n source=DocumentSource.GOOGLE_DRIVE,\n credential_json=_NESTED_CRED_JSON,\n )\n db_session.add(credential)\n db_session.flush()\n\n # Immediate read (no DB round-trip) — tests the set event wrapping\n assert isinstance(credential.credential_json, SensitiveValue)\n assert credential.credential_json.get_value(apply_mask=False) == _NESTED_CRED_JSON\n\n # DB round-trip — tests process_result_value\n db_session.expire(credential)\n reloaded = credential.credential_json\n assert isinstance(reloaded, SensitiveValue)\n assert reloaded.get_value(apply_mask=False) == _NESTED_CRED_JSON\n\n db_session.rollback()\n\n\ndef test_reassign_same_nested_json_not_dirty(db_session: Session) -> None:\n \"\"\"Re-assigning the same nested dict should not mark the session dirty.\"\"\"\n credential = Credential(\n source=DocumentSource.GOOGLE_DRIVE,\n credential_json=_NESTED_CRED_JSON,\n )\n db_session.add(credential)\n db_session.flush()\n\n # Clear dirty state from the insert\n db_session.expire(credential)\n _ = credential.credential_json # force reload\n\n # Re-assign identical value\n credential.credential_json = _NESTED_CRED_JSON # type: ignore[assignment]\n assert not db_session.is_modified(credential)\n\n db_session.rollback()\n\n\ndef test_assign_different_nested_json_is_dirty(db_session: Session) -> None:\n \"\"\"Assigning a different nested dict should mark the session dirty.\"\"\"\n credential = Credential(\n source=DocumentSource.GOOGLE_DRIVE,\n credential_json=_NESTED_CRED_JSON,\n )\n db_session.add(credential)\n db_session.flush()\n\n db_session.expire(credential)\n _ = credential.credential_json # force reload\n\n modified_cred = {**_NESTED_CRED_JSON, \"scopes\": [\"read\"]}\n credential.credential_json = modified_cred # type: ignore[assignment]\n assert db_session.is_modified(credential)\n\n db_session.rollback()\n" }, { "path": "backend/tests/external_dependency_unit/db/test_rotate_encryption_key.py", "content": "\"\"\"Tests for rotate_encryption_key against real Postgres.\n\nUses real ORM models (Credential, InternetSearchProvider) and the actual\nPostgres database. Discovery is mocked in rotation tests to scope mutations\nto only the test rows — the real _discover_encrypted_columns walk is tested\nseparately in TestDiscoverEncryptedColumns.\n\nRequires a running Postgres instance. Run with::\n\n python -m dotenv -f .vscode/.env run -- pytest tests/external_dependency_unit/db/test_rotate_encryption_key.py\n\"\"\"\n\nimport json\nfrom collections.abc import Generator\nfrom unittest.mock import patch\n\nimport pytest\nfrom sqlalchemy import LargeBinary\nfrom sqlalchemy import select\nfrom sqlalchemy import text\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.utils.encryption import _decrypt_bytes\nfrom ee.onyx.utils.encryption import _encrypt_string\nfrom ee.onyx.utils.encryption import _get_trimmed_key\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.models import Credential\nfrom onyx.db.models import EncryptedJson\nfrom onyx.db.models import EncryptedString\nfrom onyx.db.models import InternetSearchProvider\nfrom onyx.db.rotate_encryption_key import _discover_encrypted_columns\nfrom onyx.db.rotate_encryption_key import rotate_encryption_key\nfrom onyx.utils.variable_functionality import fetch_versioned_implementation\nfrom onyx.utils.variable_functionality import global_version\n\nEE_MODULE = \"ee.onyx.utils.encryption\"\nROTATE_MODULE = \"onyx.db.rotate_encryption_key\"\n\nOLD_KEY = \"o\" * 16\nNEW_KEY = \"n\" * 16\n\n\n@pytest.fixture(autouse=True)\ndef _enable_ee() -> Generator[None, None, None]:\n prev = global_version._is_ee\n global_version.set_ee()\n fetch_versioned_implementation.cache_clear()\n yield\n global_version._is_ee = prev\n fetch_versioned_implementation.cache_clear()\n\n\n@pytest.fixture(autouse=True)\ndef _clear_key_cache() -> None:\n _get_trimmed_key.cache_clear()\n\n\ndef _raw_credential_bytes(db_session: Session, credential_id: int) -> bytes | None:\n \"\"\"Read raw bytes from credential_json, bypassing the TypeDecorator.\"\"\"\n col = Credential.__table__.c.credential_json\n stmt = select(col.cast(LargeBinary)).where(\n Credential.__table__.c.id == credential_id\n )\n return db_session.execute(stmt).scalar()\n\n\ndef _raw_isp_bytes(db_session: Session, isp_id: int) -> bytes | None:\n \"\"\"Read raw bytes from InternetSearchProvider.api_key.\"\"\"\n col = InternetSearchProvider.__table__.c.api_key\n stmt = select(col.cast(LargeBinary)).where(\n InternetSearchProvider.__table__.c.id == isp_id\n )\n return db_session.execute(stmt).scalar()\n\n\nclass TestDiscoverEncryptedColumns:\n \"\"\"Verify _discover_encrypted_columns finds real production models.\"\"\"\n\n def test_discovers_credential_json(self) -> None:\n results = _discover_encrypted_columns()\n found = {\n (model_cls.__tablename__, col_name, is_json) # type: ignore[attr-defined]\n for model_cls, col_name, _, is_json in results\n }\n assert (\"credential\", \"credential_json\", True) in found\n\n def test_discovers_internet_search_provider_api_key(self) -> None:\n results = _discover_encrypted_columns()\n found = {\n (model_cls.__tablename__, col_name, is_json) # type: ignore[attr-defined]\n for model_cls, col_name, _, is_json in results\n }\n assert (\"internet_search_provider\", \"api_key\", False) in found\n\n def test_all_encrypted_string_columns_are_not_json(self) -> None:\n results = _discover_encrypted_columns()\n for model_cls, col_name, _, is_json in results:\n col = getattr(model_cls, col_name).property.columns[0]\n if isinstance(col.type, EncryptedString):\n assert not is_json, (\n f\"{model_cls.__tablename__}.{col_name} is EncryptedString \" # type: ignore[attr-defined]\n f\"but is_json={is_json}\"\n )\n\n def test_all_encrypted_json_columns_are_json(self) -> None:\n results = _discover_encrypted_columns()\n for model_cls, col_name, _, is_json in results:\n col = getattr(model_cls, col_name).property.columns[0]\n if isinstance(col.type, EncryptedJson):\n assert is_json, (\n f\"{model_cls.__tablename__}.{col_name} is EncryptedJson \" # type: ignore[attr-defined]\n f\"but is_json={is_json}\"\n )\n\n\nclass TestRotateCredential:\n \"\"\"Test rotation against the real Credential table (EncryptedJson).\n\n Discovery is scoped to only the Credential model to avoid mutating\n other tables in the test database.\n \"\"\"\n\n @pytest.fixture(autouse=True)\n def _limit_discovery(self) -> Generator[None, None, None]:\n with patch(\n f\"{ROTATE_MODULE}._discover_encrypted_columns\",\n return_value=[(Credential, \"credential_json\", [\"id\"], True)],\n ):\n yield\n\n @pytest.fixture()\n def credential_id(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> Generator[int, None, None]:\n \"\"\"Insert a Credential row with raw encrypted bytes, clean up after.\"\"\"\n config = {\"api_key\": \"sk-test-1234\", \"endpoint\": \"https://example.com\"}\n encrypted = _encrypt_string(json.dumps(config), key=OLD_KEY)\n\n result = db_session.execute(\n text(\n \"INSERT INTO credential \"\n \"(source, credential_json, admin_public, curator_public) \"\n \"VALUES (:source, :cred_json, true, false) \"\n \"RETURNING id\"\n ),\n {\"source\": DocumentSource.INGESTION_API.value, \"cred_json\": encrypted},\n )\n cred_id = result.scalar_one()\n db_session.commit()\n\n yield cred_id\n\n db_session.execute(\n text(\"DELETE FROM credential WHERE id = :id\"), {\"id\": cred_id}\n )\n db_session.commit()\n\n def test_rotates_credential_json(\n self, db_session: Session, credential_id: int\n ) -> None:\n with (\n patch(f\"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n patch(f\"{EE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n ):\n totals = rotate_encryption_key(db_session, old_key=OLD_KEY)\n\n assert totals.get(\"credential.credential_json\", 0) >= 1\n\n raw = _raw_credential_bytes(db_session, credential_id)\n assert raw is not None\n decrypted = json.loads(_decrypt_bytes(raw, key=NEW_KEY))\n assert decrypted[\"api_key\"] == \"sk-test-1234\"\n assert decrypted[\"endpoint\"] == \"https://example.com\"\n\n def test_skips_already_rotated(\n self, db_session: Session, credential_id: int\n ) -> None:\n with (\n patch(f\"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n patch(f\"{EE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n ):\n rotate_encryption_key(db_session, old_key=OLD_KEY)\n _ = rotate_encryption_key(db_session, old_key=OLD_KEY)\n\n raw = _raw_credential_bytes(db_session, credential_id)\n assert raw is not None\n decrypted = json.loads(_decrypt_bytes(raw, key=NEW_KEY))\n assert decrypted[\"api_key\"] == \"sk-test-1234\"\n\n def test_dry_run_does_not_modify(\n self, db_session: Session, credential_id: int\n ) -> None:\n original = _raw_credential_bytes(db_session, credential_id)\n\n with (\n patch(f\"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n patch(f\"{EE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n ):\n totals = rotate_encryption_key(db_session, old_key=OLD_KEY, dry_run=True)\n\n assert totals.get(\"credential.credential_json\", 0) >= 1\n\n raw_after = _raw_credential_bytes(db_session, credential_id)\n assert raw_after == original\n\n\nclass TestRotateInternetSearchProvider:\n \"\"\"Test rotation against the real InternetSearchProvider table (EncryptedString).\n\n Discovery is scoped to only the InternetSearchProvider model to avoid\n mutating other tables in the test database.\n \"\"\"\n\n @pytest.fixture(autouse=True)\n def _limit_discovery(self) -> Generator[None, None, None]:\n with patch(\n f\"{ROTATE_MODULE}._discover_encrypted_columns\",\n return_value=[\n (InternetSearchProvider, \"api_key\", [\"id\"], False),\n ],\n ):\n yield\n\n @pytest.fixture()\n def isp_id(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> Generator[int, None, None]:\n \"\"\"Insert an InternetSearchProvider row with raw encrypted bytes.\"\"\"\n encrypted = _encrypt_string(\"sk-secret-api-key\", key=OLD_KEY)\n\n result = db_session.execute(\n text(\n \"INSERT INTO internet_search_provider \"\n \"(name, provider_type, api_key, is_active) \"\n \"VALUES (:name, :ptype, :api_key, false) \"\n \"RETURNING id\"\n ),\n {\n \"name\": f\"test-rotation-{id(self)}\",\n \"ptype\": \"test\",\n \"api_key\": encrypted,\n },\n )\n isp_id = result.scalar_one()\n db_session.commit()\n\n yield isp_id\n\n db_session.execute(\n text(\"DELETE FROM internet_search_provider WHERE id = :id\"),\n {\"id\": isp_id},\n )\n db_session.commit()\n\n def test_rotates_api_key(self, db_session: Session, isp_id: int) -> None:\n with (\n patch(f\"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n patch(f\"{EE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n ):\n totals = rotate_encryption_key(db_session, old_key=OLD_KEY)\n\n assert totals.get(\"internet_search_provider.api_key\", 0) >= 1\n\n raw = _raw_isp_bytes(db_session, isp_id)\n assert raw is not None\n assert _decrypt_bytes(raw, key=NEW_KEY) == \"sk-secret-api-key\"\n\n def test_rotates_from_unencrypted(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Test rotating data that was stored without any encryption key.\"\"\"\n result = db_session.execute(\n text(\n \"INSERT INTO internet_search_provider \"\n \"(name, provider_type, api_key, is_active) \"\n \"VALUES (:name, :ptype, :api_key, false) \"\n \"RETURNING id\"\n ),\n {\n \"name\": f\"test-raw-{id(self)}\",\n \"ptype\": \"test\",\n \"api_key\": b\"raw-api-key\",\n },\n )\n isp_id = result.scalar_one()\n db_session.commit()\n\n try:\n with (\n patch(f\"{ROTATE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n patch(f\"{EE_MODULE}.ENCRYPTION_KEY_SECRET\", NEW_KEY),\n ):\n totals = rotate_encryption_key(db_session, old_key=None)\n\n assert totals.get(\"internet_search_provider.api_key\", 0) >= 1\n\n raw = _raw_isp_bytes(db_session, isp_id)\n assert raw is not None\n assert _decrypt_bytes(raw, key=NEW_KEY) == \"raw-api-key\"\n finally:\n db_session.execute(\n text(\"DELETE FROM internet_search_provider WHERE id = :id\"),\n {\"id\": isp_id},\n )\n db_session.commit()\n" }, { "path": "backend/tests/external_dependency_unit/db/test_tag_race_condition.py", "content": "\"\"\"\nTest suite for tag creation race condition handling.\n\nTests that concurrent tag creation operations don't fail due to\nUniqueViolation errors, which would occur if the upsert logic\nisn't properly implemented.\n\"\"\"\n\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import Future\nfrom concurrent.futures import ThreadPoolExecutor\nfrom typing import Union\nfrom uuid import uuid4\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import Document\nfrom onyx.db.models import Tag\nfrom onyx.db.tag import create_or_add_document_tag\nfrom onyx.db.tag import create_or_add_document_tag_list\n\n\ndef _create_test_document(db_session: Session, doc_id: str) -> Document:\n \"\"\"Create a minimal test document.\"\"\"\n document = Document(\n id=doc_id,\n semantic_id=f\"semantic_{doc_id}\",\n boost=0,\n hidden=False,\n from_ingestion_api=False,\n )\n db_session.add(document)\n db_session.commit()\n return document\n\n\nclass TestTagRaceCondition:\n \"\"\"Tests for tag creation race condition handling.\"\"\"\n\n def test_concurrent_tag_creation_single_tag(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Test that multiple concurrent calls to create_or_add_document_tag\n with the same tag key/value all succeed without UniqueViolation errors.\n\n This simulates the race condition that occurs when multiple workers\n try to create the same tag simultaneously during document indexing.\n \"\"\"\n # Create multiple test documents that will all get the same tag\n num_documents = 20\n doc_ids = [f\"test_doc_race_{uuid4().hex[:8]}\" for _ in range(num_documents)]\n\n for doc_id in doc_ids:\n _create_test_document(db_session, doc_id)\n\n # Use a unique tag key/value for this test run to avoid interference\n test_tag_key = f\"test_key_{uuid4().hex[:8]}\"\n test_tag_value = f\"test_value_{uuid4().hex[:8]}\"\n test_source = DocumentSource.FILE\n\n errors: list[Exception] = []\n results: list[Tag | None] = []\n\n def create_tag_for_document(doc_id: str) -> Tag | None:\n \"\"\"Worker function that creates a tag for a document using its own session.\"\"\"\n with get_session_with_current_tenant() as session:\n return create_or_add_document_tag(\n tag_key=test_tag_key,\n tag_value=test_tag_value,\n source=test_source,\n document_id=doc_id,\n db_session=session,\n )\n\n # Run all tag creations concurrently with high parallelism\n with ThreadPoolExecutor(max_workers=num_documents) as executor:\n futures = {\n executor.submit(create_tag_for_document, doc_id): doc_id\n for doc_id in doc_ids\n }\n\n for future in as_completed(futures):\n doc_id = futures[future]\n try:\n result = future.result()\n results.append(result)\n except Exception as e:\n errors.append(e)\n\n # All operations should succeed without errors\n assert len(errors) == 0, f\"Got {len(errors)} errors: {errors}\"\n assert len(results) == num_documents\n\n # All results should be valid Tag objects\n for result in results:\n assert result is not None\n assert result.tag_key == test_tag_key\n assert result.tag_value == test_tag_value\n assert result.source == test_source\n\n # Verify only ONE tag was created in the database (not num_documents tags)\n with get_session_with_current_tenant() as session:\n tag_count = (\n session.execute(\n select(Tag).where(\n Tag.tag_key == test_tag_key,\n Tag.tag_value == test_tag_value,\n Tag.source == test_source,\n )\n )\n .scalars()\n .all()\n )\n\n assert len(tag_count) == 1, f\"Expected 1 tag, found {len(tag_count)}\"\n\n def test_concurrent_tag_list_creation(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Test that multiple concurrent calls to create_or_add_document_tag_list\n with the same tag values all succeed without UniqueViolation errors.\n \"\"\"\n # Create multiple test documents\n num_documents = 20\n doc_ids = [\n f\"test_doc_list_race_{uuid4().hex[:8]}\" for _ in range(num_documents)\n ]\n\n for doc_id in doc_ids:\n _create_test_document(db_session, doc_id)\n\n # Use unique tag key/values for this test run\n test_tag_key = f\"test_list_key_{uuid4().hex[:8]}\"\n test_tag_values = [f\"value_{i}_{uuid4().hex[:4]}\" for i in range(5)]\n test_source = DocumentSource.FILE\n\n errors: list[Exception] = []\n results: list[list[Tag]] = []\n\n def create_tag_list_for_document(doc_id: str) -> list[Tag]:\n \"\"\"Worker function that creates tag list for a document using its own session.\"\"\"\n with get_session_with_current_tenant() as session:\n return create_or_add_document_tag_list(\n tag_key=test_tag_key,\n tag_values=test_tag_values,\n source=test_source,\n document_id=doc_id,\n db_session=session,\n )\n\n # Run all tag creations concurrently\n with ThreadPoolExecutor(max_workers=num_documents) as executor:\n futures = {\n executor.submit(create_tag_list_for_document, doc_id): doc_id\n for doc_id in doc_ids\n }\n\n for future in as_completed(futures):\n doc_id = futures[future]\n try:\n result = future.result()\n results.append(result)\n except Exception as e:\n errors.append(e)\n\n # All operations should succeed without errors\n assert len(errors) == 0, f\"Got {len(errors)} errors: {errors}\"\n assert len(results) == num_documents\n\n # Each result should have all the expected tags\n for result in results:\n assert len(result) == len(test_tag_values)\n result_values = {tag.tag_value for tag in result}\n assert result_values == set(test_tag_values)\n\n # Verify exactly len(test_tag_values) tags were created (one per value)\n with get_session_with_current_tenant() as session:\n tags = (\n session.execute(\n select(Tag).where(\n Tag.tag_key == test_tag_key,\n Tag.tag_value.in_(test_tag_values),\n Tag.source == test_source,\n )\n )\n .scalars()\n .all()\n )\n\n assert len(tags) == len(\n test_tag_values\n ), f\"Expected {len(test_tag_values)} tags, found {len(tags)}\"\n\n def test_concurrent_mixed_tag_operations(\n self,\n db_session: Session,\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Test that concurrent single tag and tag list operations on the same\n tag key/value don't interfere with each other.\n\n This is a more realistic scenario where different documents might\n have the same metadata key but different value types (single vs list).\n \"\"\"\n num_documents = 10\n doc_ids_single = [\n f\"test_doc_single_{uuid4().hex[:8]}\" for _ in range(num_documents)\n ]\n doc_ids_list = [\n f\"test_doc_list_{uuid4().hex[:8]}\" for _ in range(num_documents)\n ]\n\n for doc_id in doc_ids_single + doc_ids_list:\n _create_test_document(db_session, doc_id)\n\n # Same key but used as both single value and list value\n test_tag_key = f\"mixed_key_{uuid4().hex[:8]}\"\n test_single_value = f\"single_value_{uuid4().hex[:8]}\"\n test_list_values = [test_single_value] # Same value but as list\n test_source = DocumentSource.FILE\n\n errors: list[Exception] = []\n\n def create_single_tag(doc_id: str) -> Tag | None:\n with get_session_with_current_tenant() as session:\n return create_or_add_document_tag(\n tag_key=test_tag_key,\n tag_value=test_single_value,\n source=test_source,\n document_id=doc_id,\n db_session=session,\n )\n\n def create_list_tag(doc_id: str) -> list[Tag]:\n with get_session_with_current_tenant() as session:\n return create_or_add_document_tag_list(\n tag_key=test_tag_key,\n tag_values=test_list_values,\n source=test_source,\n document_id=doc_id,\n db_session=session,\n )\n\n # Run both types of operations concurrently\n with ThreadPoolExecutor(max_workers=num_documents * 2) as executor:\n futures: list[Future[Union[Tag | None] | list[Tag]]] = []\n for doc_id in doc_ids_single:\n futures.append(executor.submit(create_single_tag, doc_id))\n for doc_id in doc_ids_list:\n futures.append(executor.submit(create_list_tag, doc_id))\n\n for future in as_completed(futures):\n try:\n future.result()\n except Exception as e:\n errors.append(e)\n\n # All operations should succeed\n assert len(errors) == 0, f\"Got {len(errors)} errors: {errors}\"\n\n # Should have exactly 2 tags: one with is_list=False, one with is_list=True\n with get_session_with_current_tenant() as session:\n tags = (\n session.execute(\n select(Tag).where(\n Tag.tag_key == test_tag_key,\n Tag.tag_value == test_single_value,\n Tag.source == test_source,\n )\n )\n .scalars()\n .all()\n )\n\n assert (\n len(tags) == 2\n ), f\"Expected 2 tags (is_list=True and False), found {len(tags)}\"\n is_list_values = {tag.is_list for tag in tags}\n assert is_list_values == {True, False}\n" }, { "path": "backend/tests/external_dependency_unit/db/test_user_account_type.py", "content": "\"\"\"\nTests that account_type is correctly set when creating users through\nthe internal DB functions: add_slack_user_if_not_exists and\nbatch_add_ext_perm_user_if_not_exists.\n\nThese functions are called by background workers (Slack bot, permission sync)\nand are not exposed via API endpoints, so they must be tested directly.\n\"\"\"\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import AccountType\nfrom onyx.db.models import UserRole\nfrom onyx.db.users import add_slack_user_if_not_exists\nfrom onyx.db.users import batch_add_ext_perm_user_if_not_exists\n\n\ndef test_slack_user_creation_sets_account_type_bot(db_session: Session) -> None:\n \"\"\"add_slack_user_if_not_exists sets account_type=BOT and role=SLACK_USER.\"\"\"\n user = add_slack_user_if_not_exists(db_session, \"slack_acct_type@test.com\")\n\n assert user.role == UserRole.SLACK_USER\n assert user.account_type == AccountType.BOT\n\n\ndef test_ext_perm_user_creation_sets_account_type(db_session: Session) -> None:\n \"\"\"batch_add_ext_perm_user_if_not_exists sets account_type=EXT_PERM_USER.\"\"\"\n users = batch_add_ext_perm_user_if_not_exists(\n db_session, [\"extperm_acct_type@test.com\"]\n )\n\n assert len(users) == 1\n user = users[0]\n assert user.role == UserRole.EXT_PERM_USER\n assert user.account_type == AccountType.EXT_PERM_USER\n\n\ndef test_ext_perm_to_slack_upgrade_updates_role_and_account_type(\n db_session: Session,\n) -> None:\n \"\"\"When an EXT_PERM_USER is upgraded to slack, both role and account_type update.\"\"\"\n email = \"ext_to_slack_acct_type@test.com\"\n\n # Create as ext_perm user first\n batch_add_ext_perm_user_if_not_exists(db_session, [email])\n\n # Now \"upgrade\" via slack path\n user = add_slack_user_if_not_exists(db_session, email)\n\n assert user.role == UserRole.SLACK_USER\n assert user.account_type == AccountType.BOT\n" }, { "path": "backend/tests/external_dependency_unit/discord_bot/conftest.py", "content": "\"\"\"Fixtures for Discord bot external dependency tests.\"\"\"\n\nfrom collections.abc import Generator\nfrom unittest.mock import AsyncMock\nfrom unittest.mock import MagicMock\n\nimport discord\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n\nTEST_TENANT_ID: str = \"public\"\n\n\n@pytest.fixture(scope=\"function\")\ndef db_session() -> Generator[Session, None, None]:\n \"\"\"Create a database session for testing.\"\"\"\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n with get_session_with_current_tenant() as session:\n yield session\n\n\n@pytest.fixture(scope=\"function\")\ndef tenant_context() -> Generator[None, None, None]:\n \"\"\"Set up tenant context for testing.\"\"\"\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n try:\n yield\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\n@pytest.fixture\ndef mock_cache_manager() -> MagicMock:\n \"\"\"Mock DiscordCacheManager.\"\"\"\n cache = MagicMock()\n cache.get_tenant.return_value = TEST_TENANT_ID\n cache.get_api_key.return_value = \"test_api_key\"\n cache.refresh_all = AsyncMock()\n cache.refresh_guild = AsyncMock()\n cache.is_initialized = True\n return cache\n\n\n@pytest.fixture\ndef mock_api_client() -> MagicMock:\n \"\"\"Mock OnyxAPIClient.\"\"\"\n client = MagicMock()\n client.initialize = AsyncMock()\n client.close = AsyncMock()\n client.is_initialized = True\n\n # Mock successful response\n mock_response = MagicMock()\n mock_response.answer = \"Test response from bot\"\n mock_response.citation_info = None\n mock_response.top_documents = None\n mock_response.error_msg = None\n\n client.send_chat_message = AsyncMock(return_value=mock_response)\n client.health_check = AsyncMock(return_value=True)\n return client\n\n\n@pytest.fixture\ndef mock_discord_guild() -> MagicMock:\n \"\"\"Mock Discord guild with channels.\"\"\"\n guild = MagicMock(spec=discord.Guild)\n guild.id = 123456789\n guild.name = \"Test Server\"\n guild.default_role = MagicMock()\n\n # Create some mock channels\n text_channel = MagicMock(spec=discord.TextChannel)\n text_channel.id = 111111111\n text_channel.name = \"general\"\n text_channel.type = discord.ChannelType.text\n perms = MagicMock()\n perms.view_channel = True\n text_channel.permissions_for.return_value = perms\n\n forum_channel = MagicMock(spec=discord.ForumChannel)\n forum_channel.id = 222222222\n forum_channel.name = \"forum\"\n forum_channel.type = discord.ChannelType.forum\n forum_channel.permissions_for.return_value = perms\n\n private_channel = MagicMock(spec=discord.TextChannel)\n private_channel.id = 333333333\n private_channel.name = \"private\"\n private_channel.type = discord.ChannelType.text\n private_perms = MagicMock()\n private_perms.view_channel = False\n private_channel.permissions_for.return_value = private_perms\n\n guild.channels = [text_channel, forum_channel, private_channel]\n guild.text_channels = [text_channel, private_channel]\n guild.forum_channels = [forum_channel]\n\n return guild\n\n\n@pytest.fixture\ndef mock_discord_message(mock_discord_guild: MagicMock) -> MagicMock:\n \"\"\"Mock Discord message for testing.\"\"\"\n msg = MagicMock(spec=discord.Message)\n msg.id = 555555555\n msg.author = MagicMock(spec=discord.Member)\n msg.author.id = 444444444\n msg.author.bot = False\n msg.author.display_name = \"TestUser\"\n msg.author.guild_permissions = MagicMock()\n msg.author.guild_permissions.administrator = True\n msg.author.guild_permissions.manage_guild = True\n msg.content = \"Hello bot\"\n msg.guild = mock_discord_guild\n msg.channel = MagicMock()\n msg.channel.id = 111111111\n msg.channel.name = \"general\"\n msg.channel.send = AsyncMock()\n msg.type = discord.MessageType.default\n msg.mentions = []\n msg.role_mentions = []\n msg.channel_mentions = []\n msg.reference = None\n msg.add_reaction = AsyncMock()\n msg.remove_reaction = AsyncMock()\n msg.reply = AsyncMock()\n msg.create_thread = AsyncMock()\n return msg\n\n\n@pytest.fixture\ndef mock_bot_user() -> MagicMock:\n \"\"\"Mock Discord bot user.\"\"\"\n user = MagicMock(spec=discord.ClientUser)\n user.id = 987654321\n user.display_name = \"OnyxBot\"\n user.bot = True\n return user\n\n\n@pytest.fixture\ndef mock_discord_bot(\n mock_cache_manager: MagicMock,\n mock_api_client: MagicMock,\n mock_bot_user: MagicMock,\n) -> MagicMock:\n \"\"\"Mock OnyxDiscordClient.\"\"\"\n bot = MagicMock()\n bot.user = mock_bot_user\n bot.cache = mock_cache_manager\n bot.api_client = mock_api_client\n bot.ready = True\n bot.loop = MagicMock()\n bot.is_closed.return_value = False\n bot.guilds = []\n return bot\n" }, { "path": "backend/tests/external_dependency_unit/discord_bot/test_discord_events.py", "content": "\"\"\"Tests for Discord bot event handling with mocked Discord API.\n\nThese tests mock the Discord API to test event handling logic.\n\"\"\"\n\nfrom unittest.mock import AsyncMock\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport discord\nimport pytest\n\nfrom onyx.onyxbot.discord.handle_commands import get_text_channels\nfrom onyx.onyxbot.discord.handle_commands import handle_dm\nfrom onyx.onyxbot.discord.handle_commands import handle_registration_command\nfrom onyx.onyxbot.discord.handle_commands import handle_sync_channels_command\nfrom onyx.onyxbot.discord.handle_message import process_chat_message\nfrom onyx.onyxbot.discord.handle_message import send_error_response\nfrom onyx.onyxbot.discord.handle_message import send_response\n\n\nclass TestGuildRegistrationCommand:\n \"\"\"Tests for !register command handling.\"\"\"\n\n @pytest.mark.asyncio\n async def test_register_guild_success(\n self,\n mock_discord_message: MagicMock,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"Valid registration key with admin perms succeeds.\"\"\"\n mock_discord_message.content = \"!register discord_public.valid_token\"\n\n with (\n patch(\n \"onyx.onyxbot.discord.handle_commands.parse_discord_registration_key\",\n return_value=\"public\",\n ),\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_session_with_tenant\"\n ) as mock_session,\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_guild_config_by_registration_key\"\n ) as mock_get_config,\n patch(\"onyx.onyxbot.discord.handle_commands.bulk_create_channel_configs\"),\n ):\n mock_db = MagicMock()\n mock_session.return_value.__enter__ = MagicMock(return_value=mock_db)\n mock_session.return_value.__exit__ = MagicMock()\n\n mock_config = MagicMock()\n mock_config.id = 1\n mock_config.guild_id = None # Not yet registered\n mock_get_config.return_value = mock_config\n\n mock_cache_manager.get_tenant.return_value = None # Not in cache yet\n\n result = await handle_registration_command(\n mock_discord_message, mock_cache_manager\n )\n\n assert result is True\n mock_discord_message.reply.assert_called()\n # Check that success message was sent\n call_args = mock_discord_message.reply.call_args\n assert \"Successfully registered\" in str(call_args)\n\n @pytest.mark.asyncio\n async def test_register_invalid_key_format(\n self,\n mock_discord_message: MagicMock,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"Malformed key DMs user and deletes message.\"\"\"\n mock_discord_message.content = \"!register abc\" # Malformed\n\n with patch(\n \"onyx.onyxbot.discord.handle_commands.parse_discord_registration_key\",\n return_value=None, # Invalid format\n ):\n result = await handle_registration_command(\n mock_discord_message, mock_cache_manager\n )\n\n assert result is True\n # On failure: DM the author and delete the message\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"Invalid\" in str(call_args)\n mock_discord_message.delete.assert_called()\n\n @pytest.mark.asyncio\n async def test_register_key_not_found(\n self,\n mock_discord_message: MagicMock,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"Key not in database DMs user and deletes message.\"\"\"\n mock_discord_message.content = \"!register discord_public.notexist\"\n\n with (\n patch(\n \"onyx.onyxbot.discord.handle_commands.parse_discord_registration_key\",\n return_value=\"public\",\n ),\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_session_with_tenant\"\n ) as mock_session,\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_guild_config_by_registration_key\",\n return_value=None, # Not found\n ),\n ):\n mock_db = MagicMock()\n mock_session.return_value.__enter__ = MagicMock(return_value=mock_db)\n # Must return False so exceptions are not suppressed\n mock_session.return_value.__exit__ = MagicMock(return_value=False)\n mock_cache_manager.get_tenant.return_value = None\n\n result = await handle_registration_command(\n mock_discord_message, mock_cache_manager\n )\n\n assert result is True\n # On failure: DM the author and delete the message\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"not found\" in str(call_args).lower()\n mock_discord_message.delete.assert_called()\n\n @pytest.mark.asyncio\n async def test_register_key_already_used(\n self,\n mock_discord_message: MagicMock,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"Previously used key DMs user and deletes message.\"\"\"\n mock_discord_message.content = \"!register discord_public.used_key\"\n\n with (\n patch(\n \"onyx.onyxbot.discord.handle_commands.parse_discord_registration_key\",\n return_value=\"public\",\n ),\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_session_with_tenant\"\n ) as mock_session,\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_guild_config_by_registration_key\"\n ) as mock_get_config,\n ):\n mock_db = MagicMock()\n mock_session.return_value.__enter__ = MagicMock(return_value=mock_db)\n # Must return False so exceptions are not suppressed\n mock_session.return_value.__exit__ = MagicMock(return_value=False)\n\n mock_config = MagicMock()\n mock_config.guild_id = 999999 # Already registered!\n mock_get_config.return_value = mock_config\n\n mock_cache_manager.get_tenant.return_value = None\n\n result = await handle_registration_command(\n mock_discord_message, mock_cache_manager\n )\n\n assert result is True\n # On failure: DM the author and delete the message\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"already\" in str(call_args).lower()\n mock_discord_message.delete.assert_called()\n\n @pytest.mark.asyncio\n async def test_register_guild_already_registered(\n self,\n mock_discord_message: MagicMock,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"Guild already in cache DMs user and deletes message.\"\"\"\n mock_discord_message.content = \"!register discord_public.valid_token\"\n\n with patch(\n \"onyx.onyxbot.discord.handle_commands.parse_discord_registration_key\",\n return_value=\"public\",\n ):\n # Guild already in cache\n mock_cache_manager.get_tenant.return_value = \"existing_tenant\"\n\n result = await handle_registration_command(\n mock_discord_message, mock_cache_manager\n )\n\n assert result is True\n # On failure: DM the author and delete the message\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"already registered\" in str(call_args).lower()\n mock_discord_message.delete.assert_called()\n\n @pytest.mark.asyncio\n async def test_register_no_permission(\n self,\n mock_discord_message: MagicMock,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"User without admin perms gets DM and message deleted.\"\"\"\n mock_discord_message.content = \"!register discord_public.valid_token\"\n mock_discord_message.author.guild_permissions.administrator = False\n mock_discord_message.author.guild_permissions.manage_guild = False\n\n result = await handle_registration_command(\n mock_discord_message, mock_cache_manager\n )\n\n assert result is True\n # On failure: DM the author and delete the message\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"permission\" in str(call_args).lower()\n mock_discord_message.delete.assert_called()\n\n @pytest.mark.asyncio\n async def test_register_in_dm(\n self,\n mock_cache_manager: MagicMock,\n ) -> None:\n \"\"\"Registration in DM sends DM and returns True.\"\"\"\n msg = MagicMock(spec=discord.Message)\n msg.guild = None # DM\n msg.content = \"!register discord_public.token\"\n msg.author = MagicMock()\n msg.author.send = AsyncMock()\n\n result = await handle_registration_command(msg, mock_cache_manager)\n\n assert result is True\n msg.author.send.assert_called()\n call_args = msg.author.send.call_args\n assert \"server\" in str(call_args).lower()\n\n @pytest.mark.asyncio\n async def test_register_syncs_forum_channels(\n self,\n mock_discord_message: MagicMock, # noqa: ARG002\n mock_discord_guild: MagicMock,\n ) -> None:\n \"\"\"Forum channels are included in sync.\"\"\"\n channels = get_text_channels(mock_discord_guild)\n\n channel_types = [c.channel_type for c in channels]\n assert \"forum\" in channel_types\n\n @pytest.mark.asyncio\n async def test_register_private_channel_detection(\n self,\n mock_discord_message: MagicMock, # noqa: ARG002\n mock_discord_guild: MagicMock,\n ) -> None:\n \"\"\"Private channels are marked correctly.\"\"\"\n channels = get_text_channels(mock_discord_guild)\n\n private_channels = [c for c in channels if c.is_private]\n assert len(private_channels) >= 1\n\n\nclass TestSyncChannelsCommand:\n \"\"\"Tests for !sync-channels command handling.\"\"\"\n\n @pytest.mark.asyncio\n async def test_sync_channels_adds_new(\n self,\n mock_discord_message: MagicMock,\n mock_discord_bot: MagicMock,\n ) -> None:\n \"\"\"New channel in Discord creates channel config.\"\"\"\n mock_discord_message.content = \"!sync-channels\"\n\n with (\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_session_with_tenant\"\n ) as mock_session,\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_guild_config_by_discord_id\"\n ) as mock_get_guild,\n patch(\n \"onyx.onyxbot.discord.handle_commands.get_guild_config_by_internal_id\"\n ) as mock_get_guild_internal,\n patch(\n \"onyx.onyxbot.discord.handle_commands.sync_channel_configs\"\n ) as mock_sync,\n ):\n mock_db = MagicMock()\n mock_session.return_value.__enter__ = MagicMock(return_value=mock_db)\n mock_session.return_value.__exit__ = MagicMock()\n\n mock_config = MagicMock()\n mock_config.id = 1\n mock_config.guild_id = 123456789\n mock_get_guild.return_value = mock_config\n mock_get_guild_internal.return_value = mock_config\n\n mock_sync.return_value = (1, 0, 0) # 1 added, 0 removed, 0 updated\n\n mock_discord_bot.get_guild.return_value = mock_discord_message.guild\n\n result = await handle_sync_channels_command(\n mock_discord_message, \"public\", mock_discord_bot\n )\n\n assert result is True\n mock_discord_message.reply.assert_called()\n\n @pytest.mark.asyncio\n async def test_sync_channels_no_permission(\n self,\n mock_discord_message: MagicMock,\n mock_discord_bot: MagicMock,\n ) -> None:\n \"\"\"User without admin perms gets DM and reaction.\"\"\"\n mock_discord_message.content = \"!sync-channels\"\n mock_discord_message.author.guild_permissions.administrator = False\n mock_discord_message.author.guild_permissions.manage_guild = False\n\n result = await handle_sync_channels_command(\n mock_discord_message, \"public\", mock_discord_bot\n )\n\n assert result is True\n # On failure: DM the author and react with ❌\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"permission\" in str(call_args).lower()\n mock_discord_message.add_reaction.assert_called_with(\"❌\")\n\n @pytest.mark.asyncio\n async def test_sync_channels_unregistered_guild(\n self,\n mock_discord_message: MagicMock,\n mock_discord_bot: MagicMock,\n ) -> None:\n \"\"\"Sync in unregistered guild gets DM and reaction.\"\"\"\n mock_discord_message.content = \"!sync-channels\"\n\n # tenant_id is None = not registered\n result = await handle_sync_channels_command(\n mock_discord_message, None, mock_discord_bot\n )\n\n assert result is True\n # On failure: DM the author and react with ❌\n mock_discord_message.author.send.assert_called()\n call_args = mock_discord_message.author.send.call_args\n assert \"not registered\" in str(call_args).lower()\n mock_discord_message.add_reaction.assert_called_with(\"❌\")\n\n\nclass TestMessageHandling:\n \"\"\"Tests for message handling behavior.\"\"\"\n\n @pytest.mark.asyncio\n async def test_message_adds_thinking_emoji(\n self,\n mock_discord_message: MagicMock,\n mock_api_client: MagicMock,\n mock_bot_user: MagicMock,\n ) -> None:\n \"\"\"Thinking emoji is added during processing.\"\"\"\n await process_chat_message(\n message=mock_discord_message,\n api_key=\"test_key\",\n persona_id=None,\n thread_only_mode=False,\n api_client=mock_api_client,\n bot_user=mock_bot_user,\n )\n\n mock_discord_message.add_reaction.assert_called()\n\n @pytest.mark.asyncio\n async def test_message_removes_thinking_emoji(\n self,\n mock_discord_message: MagicMock,\n mock_api_client: MagicMock,\n mock_bot_user: MagicMock,\n ) -> None:\n \"\"\"Thinking emoji is removed after response.\"\"\"\n await process_chat_message(\n message=mock_discord_message,\n api_key=\"test_key\",\n persona_id=None,\n thread_only_mode=False,\n api_client=mock_api_client,\n bot_user=mock_bot_user,\n )\n\n mock_discord_message.remove_reaction.assert_called()\n\n @pytest.mark.asyncio\n async def test_message_reaction_failure_non_blocking(\n self,\n mock_discord_message: MagicMock,\n mock_api_client: MagicMock,\n mock_bot_user: MagicMock,\n ) -> None:\n \"\"\"add_reaction failure doesn't block processing.\"\"\"\n mock_discord_message.add_reaction = AsyncMock(\n side_effect=discord.DiscordException(\"Cannot add reaction\")\n )\n\n # Should not raise - just log warning and continue\n await process_chat_message(\n message=mock_discord_message,\n api_key=\"test_key\",\n persona_id=None,\n thread_only_mode=False,\n api_client=mock_api_client,\n bot_user=mock_bot_user,\n )\n\n # Should still complete and send reply\n mock_discord_message.reply.assert_called()\n\n @pytest.mark.asyncio\n async def test_dm_response(self) -> None:\n \"\"\"DM to bot sends redirect message.\"\"\"\n msg = MagicMock(spec=discord.Message)\n msg.channel = MagicMock(spec=discord.DMChannel)\n msg.channel.send = AsyncMock()\n\n await handle_dm(msg)\n\n msg.channel.send.assert_called_once()\n call_args = msg.channel.send.call_args\n assert \"DM\" in str(call_args) or \"server\" in str(call_args).lower()\n\n\nclass TestThreadCreationAndResponseRouting:\n \"\"\"Tests for thread creation and response routing.\"\"\"\n\n @pytest.mark.asyncio\n async def test_response_in_existing_thread(\n self,\n mock_bot_user: MagicMock, # noqa: ARG002\n ) -> None:\n \"\"\"Message in thread - response appended to thread.\"\"\"\n thread = MagicMock(spec=discord.Thread)\n thread.send = AsyncMock()\n\n msg = MagicMock(spec=discord.Message)\n msg.channel = thread\n msg.reply = AsyncMock()\n msg.create_thread = AsyncMock()\n\n await send_response(msg, \"Test response\", thread_only_mode=False)\n\n # Should send to thread, not create new thread\n thread.send.assert_called()\n msg.create_thread.assert_not_called()\n\n @pytest.mark.asyncio\n async def test_response_creates_thread_thread_only_mode(\n self,\n mock_discord_message: MagicMock,\n mock_bot_user: MagicMock, # noqa: ARG002\n ) -> None:\n \"\"\"thread_only_mode=true creates new thread for response.\"\"\"\n mock_thread = MagicMock()\n mock_thread.send = AsyncMock()\n mock_discord_message.create_thread = AsyncMock(return_value=mock_thread)\n\n # Make sure it's not a thread\n mock_discord_message.channel = MagicMock(spec=discord.TextChannel)\n\n await send_response(\n mock_discord_message, \"Test response\", thread_only_mode=True\n )\n\n mock_discord_message.create_thread.assert_called()\n mock_thread.send.assert_called()\n\n @pytest.mark.asyncio\n async def test_response_replies_inline(\n self,\n mock_discord_message: MagicMock,\n mock_bot_user: MagicMock, # noqa: ARG002\n ) -> None:\n \"\"\"thread_only_mode=false uses message.reply().\"\"\"\n # Make sure it's not a thread\n mock_discord_message.channel = MagicMock(spec=discord.TextChannel)\n\n await send_response(\n mock_discord_message, \"Test response\", thread_only_mode=False\n )\n\n mock_discord_message.reply.assert_called()\n\n @pytest.mark.asyncio\n async def test_thread_name_truncation(\n self,\n mock_bot_user: MagicMock, # noqa: ARG002\n ) -> None:\n \"\"\"Thread name is truncated to 100 chars.\"\"\"\n msg = MagicMock(spec=discord.Message)\n msg.channel = MagicMock(spec=discord.TextChannel)\n msg.author = MagicMock()\n msg.author.display_name = \"A\" * 200 # Very long name\n\n mock_thread = MagicMock()\n mock_thread.send = AsyncMock()\n msg.create_thread = AsyncMock(return_value=mock_thread)\n\n await send_response(msg, \"Test\", thread_only_mode=True)\n\n call_args = msg.create_thread.call_args\n thread_name = call_args.kwargs.get(\"name\") or call_args[1].get(\"name\")\n assert len(thread_name) <= 100\n\n @pytest.mark.asyncio\n async def test_error_response_creates_thread(\n self,\n mock_discord_message: MagicMock,\n mock_bot_user: MagicMock,\n ) -> None:\n \"\"\"Error response in channel creates thread.\"\"\"\n mock_discord_message.channel = MagicMock(spec=discord.TextChannel)\n mock_thread = MagicMock()\n mock_thread.send = AsyncMock()\n mock_discord_message.create_thread = AsyncMock(return_value=mock_thread)\n\n await send_error_response(mock_discord_message, mock_bot_user)\n\n mock_discord_message.create_thread.assert_called()\n\n\nclass TestBotLifecycle:\n \"\"\"Tests for bot lifecycle management.\"\"\"\n\n @pytest.mark.asyncio\n async def test_setup_hook_initializes_cache(\n self,\n mock_cache_manager: MagicMock,\n mock_api_client: MagicMock,\n ) -> None:\n \"\"\"setup_hook calls cache.refresh_all().\"\"\"\n from onyx.onyxbot.discord.client import OnyxDiscordClient\n\n with (\n patch.object(\n OnyxDiscordClient,\n \"__init__\",\n lambda self: None, # noqa: ARG005\n ),\n patch(\n \"onyx.onyxbot.discord.client.DiscordCacheManager\",\n return_value=mock_cache_manager,\n ),\n patch(\n \"onyx.onyxbot.discord.client.OnyxAPIClient\",\n return_value=mock_api_client,\n ),\n ):\n bot = OnyxDiscordClient()\n bot.cache = mock_cache_manager\n bot.api_client = mock_api_client\n bot.loop = MagicMock()\n bot.loop.create_task = MagicMock()\n\n await bot.setup_hook()\n\n mock_cache_manager.refresh_all.assert_called()\n\n @pytest.mark.asyncio\n async def test_setup_hook_initializes_api_client(\n self,\n mock_cache_manager: MagicMock,\n mock_api_client: MagicMock,\n ) -> None:\n \"\"\"setup_hook calls api_client.initialize().\"\"\"\n from onyx.onyxbot.discord.client import OnyxDiscordClient\n\n with (\n patch.object(\n OnyxDiscordClient,\n \"__init__\",\n lambda self: None, # noqa: ARG005\n ),\n ):\n bot = OnyxDiscordClient()\n bot.cache = mock_cache_manager\n bot.api_client = mock_api_client\n bot.loop = MagicMock()\n bot.loop.create_task = MagicMock()\n\n await bot.setup_hook()\n\n mock_api_client.initialize.assert_called()\n\n @pytest.mark.asyncio\n async def test_close_closes_api_client(\n self,\n mock_cache_manager: MagicMock,\n mock_api_client: MagicMock,\n ) -> None:\n \"\"\"close() calls api_client.close().\"\"\"\n from onyx.onyxbot.discord.client import OnyxDiscordClient\n\n with (\n patch.object(\n OnyxDiscordClient,\n \"__init__\",\n lambda self: None, # noqa: ARG005\n ),\n patch.object(OnyxDiscordClient, \"is_closed\", return_value=True),\n ):\n bot = OnyxDiscordClient()\n bot.cache = mock_cache_manager\n bot.api_client = mock_api_client\n bot._cache_refresh_task = None\n bot.ready = True\n\n # Mock parent close\n async def mock_super_close() -> None:\n pass\n\n with patch(\"discord.ext.commands.Bot.close\", mock_super_close):\n await bot.close()\n\n mock_api_client.close.assert_called()\n mock_cache_manager.clear.assert_called()\n" }, { "path": "backend/tests/external_dependency_unit/document_index/conftest.py", "content": "\"\"\"Shared fixtures for document_index external dependency tests.\n\nProvides Vespa and OpenSearch index setup, tenant context, and chunk helpers.\n\"\"\"\n\nimport os\nimport time\nimport uuid\nfrom collections.abc import Generator\nfrom unittest.mock import patch\n\nimport httpx\nimport pytest\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import Document\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.document_index.interfaces_new import IndexingMetadata\nfrom onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchOldDocumentIndex,\n)\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client\nfrom onyx.document_index.vespa.shared_utils.utils import wait_for_vespa_with_timeout\nfrom onyx.indexing.models import ChunkEmbedding\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\nEMBEDDING_DIM = 128\n\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef make_chunk(\n doc_id: str,\n chunk_id: int = 0,\n content: str = \"test content\",\n) -> DocMetadataAwareIndexChunk:\n \"\"\"Create a chunk suitable for external dependency testing (128-dim embeddings).\"\"\"\n tenant_id = get_current_tenant_id()\n access = DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=True,\n )\n embeddings = ChunkEmbedding(\n full_embedding=[1.0] + [0.0] * (EMBEDDING_DIM - 1),\n mini_chunk_embeddings=[],\n )\n source_document = Document(\n id=doc_id,\n semantic_identifier=\"test_doc\",\n source=DocumentSource.FILE,\n sections=[],\n metadata={},\n title=\"test title\",\n )\n return DocMetadataAwareIndexChunk(\n tenant_id=tenant_id,\n access=access,\n document_sets=set(),\n user_project=[],\n personas=[],\n boost=0,\n aggregated_chunk_boost_factor=0,\n ancestor_hierarchy_node_ids=[],\n embeddings=embeddings,\n title_embedding=[1.0] + [0.0] * (EMBEDDING_DIM - 1),\n source_document=source_document,\n title_prefix=\"\",\n metadata_suffix_keyword=\"\",\n metadata_suffix_semantic=\"\",\n contextual_rag_reserved_tokens=0,\n doc_summary=\"\",\n chunk_context=\"\",\n mini_chunk_texts=None,\n large_chunk_id=None,\n chunk_id=chunk_id,\n blurb=content[:50],\n content=content,\n source_links={0: \"\"},\n image_file_id=None,\n section_continuation=False,\n )\n\n\ndef make_indexing_metadata(\n doc_ids: list[str],\n old_counts: list[int],\n new_counts: list[int],\n) -> IndexingMetadata:\n return IndexingMetadata(\n doc_id_to_chunk_cnt_diff={\n doc_id: IndexingMetadata.ChunkCounts(\n old_chunk_cnt=old,\n new_chunk_cnt=new,\n )\n for doc_id, old, new in zip(doc_ids, old_counts, new_counts)\n }\n )\n\n\n# ---------------------------------------------------------------------------\n# Fixtures\n# ---------------------------------------------------------------------------\n\n\n@pytest.fixture(scope=\"module\")\ndef tenant_context() -> Generator[None, None, None]:\n \"\"\"Sets up tenant context for testing.\"\"\"\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n try:\n yield\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n\n\n@pytest.fixture(scope=\"module\")\ndef test_index_name() -> Generator[str, None, None]:\n yield f\"test_index_{uuid.uuid4().hex[:8]}\"\n\n\n@pytest.fixture(scope=\"module\")\ndef httpx_client() -> Generator[httpx.Client, None, None]:\n client = get_vespa_http_client()\n try:\n yield client\n finally:\n client.close()\n\n\n@pytest.fixture(scope=\"module\")\ndef vespa_index(\n httpx_client: httpx.Client,\n tenant_context: None, # noqa: ARG001\n test_index_name: str,\n) -> Generator[VespaIndex, None, None]:\n \"\"\"Create a Vespa index, wait for schema readiness, and yield it.\"\"\"\n vespa_idx = VespaIndex(\n index_name=test_index_name,\n secondary_index_name=None,\n large_chunks_enabled=False,\n secondary_large_chunks_enabled=None,\n multitenant=MULTI_TENANT,\n httpx_client=httpx_client,\n )\n backend_dir = os.path.abspath(\n os.path.join(os.path.dirname(__file__), \"..\", \"..\", \"..\")\n )\n with patch(\"os.getcwd\", return_value=backend_dir):\n vespa_idx.ensure_indices_exist(\n primary_embedding_dim=EMBEDDING_DIM,\n primary_embedding_precision=EmbeddingPrecision.FLOAT,\n secondary_index_embedding_dim=None,\n secondary_index_embedding_precision=None,\n )\n if not wait_for_vespa_with_timeout(wait_limit=90):\n pytest.fail(\"Vespa is not available.\")\n\n # Wait until the schema is actually ready for writes on content nodes. We\n # probe by attempting a PUT; 200 means the schema is live, 400 means not\n # yet. This is only temporary until we entirely move off of Vespa.\n probe_doc = {\n \"fields\": {\n \"document_id\": \"__probe__\",\n \"chunk_id\": 0,\n \"blurb\": \"\",\n \"title\": \"\",\n \"skip_title\": True,\n \"content\": \"\",\n \"content_summary\": \"\",\n \"source_type\": \"file\",\n \"source_links\": \"null\",\n \"semantic_identifier\": \"\",\n \"section_continuation\": False,\n \"large_chunk_reference_ids\": [],\n \"metadata\": \"{}\",\n \"metadata_list\": [],\n \"metadata_suffix\": \"\",\n \"chunk_context\": \"\",\n \"doc_summary\": \"\",\n \"embeddings\": {\"full_chunk\": [1.0] + [0.0] * (EMBEDDING_DIM - 1)},\n \"access_control_list\": {},\n \"document_sets\": {},\n \"image_file_name\": None,\n \"user_project\": [],\n \"personas\": [],\n \"boost\": 0.0,\n \"aggregated_chunk_boost_factor\": 0.0,\n \"primary_owners\": [],\n \"secondary_owners\": [],\n }\n }\n probe_url = (\n f\"http://localhost:8081/document/v1/default/{test_index_name}/docid/__probe__\"\n )\n schema_ready = False\n for _ in range(60):\n resp = httpx_client.post(probe_url, json=probe_doc)\n if resp.status_code == 200:\n schema_ready = True\n httpx_client.delete(probe_url)\n break\n time.sleep(1)\n if not schema_ready:\n pytest.fail(f\"Vespa schema '{test_index_name}' did not become ready in time.\")\n\n yield vespa_idx\n\n\n@pytest.fixture(scope=\"module\")\ndef opensearch_old_index(\n tenant_context: None, # noqa: ARG001\n test_index_name: str,\n) -> Generator[OpenSearchOldDocumentIndex, None, None]:\n \"\"\"Create an OpenSearch index via the old adapter and yield it.\"\"\"\n if not wait_for_opensearch_with_timeout():\n pytest.fail(\"OpenSearch is not available.\")\n\n opensearch_idx = OpenSearchOldDocumentIndex(\n index_name=test_index_name,\n embedding_dim=EMBEDDING_DIM,\n embedding_precision=EmbeddingPrecision.FLOAT,\n secondary_index_name=None,\n secondary_embedding_dim=None,\n secondary_embedding_precision=None,\n large_chunks_enabled=False,\n secondary_large_chunks_enabled=None,\n multitenant=MULTI_TENANT,\n )\n opensearch_idx.ensure_indices_exist(\n primary_embedding_dim=EMBEDDING_DIM,\n primary_embedding_precision=EmbeddingPrecision.FLOAT,\n secondary_index_embedding_dim=None,\n secondary_index_embedding_precision=None,\n )\n\n yield opensearch_idx\n" }, { "path": "backend/tests/external_dependency_unit/document_index/test_document_index.py", "content": "\"\"\"External dependency tests for the new DocumentIndex interface.\n\nThese tests assume Vespa and OpenSearch are running.\n\"\"\"\n\nimport time\nimport uuid\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\n\nimport httpx\nimport pytest\n\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.document_index.interfaces_new import DocumentIndex as DocumentIndexNew\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchDocumentIndex,\n)\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchOldDocumentIndex,\n)\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\nfrom tests.external_dependency_unit.document_index.conftest import EMBEDDING_DIM\nfrom tests.external_dependency_unit.document_index.conftest import make_chunk\nfrom tests.external_dependency_unit.document_index.conftest import (\n make_indexing_metadata,\n)\n\n\n# ---------------------------------------------------------------------------\n# Fixtures\n# ---------------------------------------------------------------------------\n\n\n@pytest.fixture(scope=\"module\")\ndef vespa_document_index(\n vespa_index: VespaIndex, # noqa: ARG001 — ensures schema exists\n httpx_client: httpx.Client,\n test_index_name: str,\n) -> Generator[VespaDocumentIndex, None, None]:\n yield VespaDocumentIndex(\n index_name=test_index_name,\n tenant_state=TenantState(tenant_id=TEST_TENANT_ID, multitenant=False),\n large_chunks_enabled=False,\n httpx_client=httpx_client,\n )\n\n\n@pytest.fixture(scope=\"module\")\ndef opensearch_document_index(\n opensearch_old_index: OpenSearchOldDocumentIndex, # noqa: ARG001 — ensures index exists\n test_index_name: str,\n) -> Generator[OpenSearchDocumentIndex, None, None]:\n yield OpenSearchDocumentIndex(\n tenant_state=TenantState(tenant_id=TEST_TENANT_ID, multitenant=False),\n index_name=test_index_name,\n embedding_dim=EMBEDDING_DIM,\n embedding_precision=EmbeddingPrecision.FLOAT,\n )\n\n\n@pytest.fixture(scope=\"module\")\ndef document_indices(\n vespa_document_index: VespaDocumentIndex,\n opensearch_document_index: OpenSearchDocumentIndex,\n) -> Generator[list[DocumentIndexNew], None, None]:\n yield [opensearch_document_index, vespa_document_index]\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\nclass TestDocumentIndexNew:\n \"\"\"Tests the new DocumentIndex interface against real Vespa and OpenSearch.\"\"\"\n\n def test_index_single_new_doc(\n self,\n document_indices: list[DocumentIndexNew],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Indexing a single new document returns one record with already_existed=False.\"\"\"\n for document_index in document_indices:\n doc_id = f\"test_single_new_{uuid.uuid4().hex[:8]}\"\n chunk = make_chunk(doc_id)\n metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[1])\n\n results = document_index.index(chunks=[chunk], indexing_metadata=metadata)\n\n assert len(results) == 1\n assert results[0].document_id == doc_id\n assert results[0].already_existed is False\n\n def test_index_existing_doc_already_existed_true(\n self,\n document_indices: list[DocumentIndexNew],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Re-indexing a doc with previous chunks returns already_existed=True.\"\"\"\n for document_index in document_indices:\n doc_id = f\"test_existing_{uuid.uuid4().hex[:8]}\"\n chunk = make_chunk(doc_id)\n\n # First index — brand new document.\n metadata_first = make_indexing_metadata(\n [doc_id], old_counts=[0], new_counts=[1]\n )\n document_index.index(chunks=[chunk], indexing_metadata=metadata_first)\n\n # Allow near-real-time indexing to settle (needed for Vespa).\n time.sleep(1)\n\n # Re-index — old_chunk_cnt=1 signals the document already existed.\n metadata_second = make_indexing_metadata(\n [doc_id], old_counts=[1], new_counts=[1]\n )\n results = document_index.index(\n chunks=[chunk], indexing_metadata=metadata_second\n )\n\n assert len(results) == 1\n assert results[0].already_existed is True\n\n def test_index_multiple_docs(\n self,\n document_indices: list[DocumentIndexNew],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Indexing multiple documents returns one record per unique document.\"\"\"\n for document_index in document_indices:\n doc1 = f\"test_multi_1_{uuid.uuid4().hex[:8]}\"\n doc2 = f\"test_multi_2_{uuid.uuid4().hex[:8]}\"\n chunks = [\n make_chunk(doc1, chunk_id=0),\n make_chunk(doc1, chunk_id=1),\n make_chunk(doc2, chunk_id=0),\n ]\n metadata = make_indexing_metadata(\n [doc1, doc2], old_counts=[0, 0], new_counts=[2, 1]\n )\n\n results = document_index.index(chunks=chunks, indexing_metadata=metadata)\n\n result_map = {r.document_id: r.already_existed for r in results}\n assert len(result_map) == 2\n assert result_map[doc1] is False\n assert result_map[doc2] is False\n\n def test_index_deduplicates_doc_ids_in_results(\n self,\n document_indices: list[DocumentIndexNew],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"Multiple chunks from the same document produce only one\n DocumentInsertionRecord.\"\"\"\n for document_index in document_indices:\n doc_id = f\"test_dedup_{uuid.uuid4().hex[:8]}\"\n chunks = [make_chunk(doc_id, chunk_id=i) for i in range(5)]\n metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[5])\n\n results = document_index.index(chunks=chunks, indexing_metadata=metadata)\n\n assert len(results) == 1\n assert results[0].document_id == doc_id\n\n def test_index_mixed_new_and_existing_docs(\n self,\n document_indices: list[DocumentIndexNew],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"A batch with both new and existing documents returns the correct\n already_existed flag for each.\"\"\"\n for document_index in document_indices:\n existing_doc = f\"test_mixed_exist_{uuid.uuid4().hex[:8]}\"\n new_doc = f\"test_mixed_new_{uuid.uuid4().hex[:8]}\"\n\n # Pre-index the existing document.\n pre_chunk = make_chunk(existing_doc)\n pre_metadata = make_indexing_metadata(\n [existing_doc], old_counts=[0], new_counts=[1]\n )\n document_index.index(chunks=[pre_chunk], indexing_metadata=pre_metadata)\n\n time.sleep(1)\n\n # Now index a batch with the existing doc and a new doc.\n chunks = [\n make_chunk(existing_doc, chunk_id=0),\n make_chunk(new_doc, chunk_id=0),\n ]\n metadata = make_indexing_metadata(\n [existing_doc, new_doc], old_counts=[1, 0], new_counts=[1, 1]\n )\n\n results = document_index.index(chunks=chunks, indexing_metadata=metadata)\n\n result_map = {r.document_id: r.already_existed for r in results}\n assert len(result_map) == 2\n assert result_map[existing_doc] is True\n assert result_map[new_doc] is False\n\n def test_index_accepts_generator(\n self,\n document_indices: list[DocumentIndexNew],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"index() accepts a generator (any iterable), not just a list.\"\"\"\n for document_index in document_indices:\n doc_id = f\"test_gen_{uuid.uuid4().hex[:8]}\"\n metadata = make_indexing_metadata([doc_id], old_counts=[0], new_counts=[3])\n\n def chunk_gen() -> Iterator[DocMetadataAwareIndexChunk]:\n for i in range(3):\n yield make_chunk(doc_id, chunk_id=i)\n\n results = document_index.index(\n chunks=chunk_gen(), indexing_metadata=metadata\n )\n\n assert len(results) == 1\n assert results[0].document_id == doc_id\n assert results[0].already_existed is False\n" }, { "path": "backend/tests/external_dependency_unit/document_index/test_document_index_old.py", "content": "\"\"\"External dependency tests for the old DocumentIndex interface.\n\nThese tests assume Vespa and OpenSearch are running.\n\"\"\"\n\nimport time\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\n\nimport pytest\n\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.document_index.interfaces import IndexBatchParams\nfrom onyx.document_index.interfaces import VespaChunkRequest\nfrom onyx.document_index.interfaces import VespaDocumentUserFields\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n OpenSearchOldDocumentIndex,\n)\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.indexing.models import DocMetadataAwareIndexChunk\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom tests.external_dependency_unit.document_index.conftest import make_chunk\n\n\n@pytest.fixture(scope=\"module\")\ndef document_indices(\n vespa_index: VespaIndex,\n opensearch_old_index: OpenSearchOldDocumentIndex,\n) -> Generator[list[DocumentIndex], None, None]:\n # Ideally these are parametrized; doing so with pytest fixtures is tricky.\n yield [opensearch_old_index, vespa_index]\n\n\n@pytest.fixture(scope=\"function\")\ndef chunks(\n tenant_context: None, # noqa: ARG001\n) -> Generator[list[DocMetadataAwareIndexChunk], None, None]:\n yield [make_chunk(\"test_doc\", chunk_id=i) for i in range(5)]\n\n\n@pytest.fixture(scope=\"function\")\ndef index_batch_params(\n tenant_context: None, # noqa: ARG001\n) -> Generator[IndexBatchParams, None, None]:\n # WARNING: doc_id_to_previous_chunk_cnt={\"test_doc\": 0} is hardcoded to 0,\n # which is only correct on the very first index call. The document_indices\n # fixture is scope=\"module\", meaning the same OpenSearch and Vespa backends\n # persist across all test functions in this module. When a second test\n # function uses this fixture and calls document_index.index(...), the\n # backend already has 5 chunks for \"test_doc\" from the previous test run,\n # but the batch params still claim 0 prior chunks exist. This can lead to\n # orphaned/duplicate chunks that make subsequent assertions incorrect.\n # TODO: Whenever adding a second test, either change this or cleanup the\n # index between test cases.\n yield IndexBatchParams(\n doc_id_to_previous_chunk_cnt={\"test_doc\": 0},\n doc_id_to_new_chunk_cnt={\"test_doc\": 5},\n tenant_id=get_current_tenant_id(),\n large_chunks_enabled=False,\n )\n\n\nclass TestDocumentIndexOld:\n \"\"\"Tests the old DocumentIndex interface.\"\"\"\n\n # TODO(ENG-3864)(andrei): Re-enable this test.\n @pytest.mark.xfail(\n reason=\"Flaky test: Retrieved chunks vary non-deterministically before and after changing user projects and personas. Likely a timing issue with the index being updated.\"\n )\n def test_update_single_can_clear_user_projects_and_personas(\n self,\n document_indices: list[DocumentIndex],\n # This test case assumes all these chunks correspond to one document.\n chunks: list[DocMetadataAwareIndexChunk],\n index_batch_params: IndexBatchParams,\n ) -> None:\n \"\"\"\n Tests that update_single can clear user_projects and personas.\n \"\"\"\n for document_index in document_indices:\n # Precondition.\n # Ensure there is some non-empty value for user project and\n # personas.\n for chunk in chunks:\n chunk.user_project = [1]\n chunk.personas = [2]\n document_index.index(chunks, index_batch_params)\n\n # Ensure that we can get chunks as expected with filters.\n doc_id = chunks[0].source_document.id\n chunk_count = len(chunks)\n tenant_id = get_current_tenant_id()\n # We need to specify the chunk index range and specify\n # batch_retrieval=True below to trigger the codepath for Vespa's\n # search API, which uses the expected additive filtering for\n # project_id and persona_id. Otherwise we would use the codepath for\n # the visit API, which does not have this kind of filtering\n # implemented.\n chunk_request = VespaChunkRequest(\n document_id=doc_id, min_chunk_ind=0, max_chunk_ind=chunk_count - 1\n )\n project_persona_filters = IndexFilters(\n access_control_list=None,\n tenant_id=tenant_id,\n project_id_filter=1,\n persona_id_filter=2,\n # We need this even though none of the chunks belong to a\n # document set because project_id and persona_id are only\n # additive filters in the event the agent has knowledge scope;\n # if the agent does not, it is implied that it can see\n # everything it is allowed to.\n document_set=[\"1\"],\n )\n # Not best practice here but the API for refreshing the index to\n # ensure that the latest data is present is not exposed in this\n # class and is not the same for Vespa and OpenSearch, so we just\n # tolerate a sleep for now. As a consequence the number of tests in\n # this suite should be small. We only need to tolerate this for as\n # long as we continue to use Vespa, we can consider exposing\n # something for OpenSearch later.\n time.sleep(1)\n inference_chunks = document_index.id_based_retrieval(\n chunk_requests=[chunk_request],\n filters=project_persona_filters,\n batch_retrieval=True,\n )\n assert len(inference_chunks) == chunk_count\n # Sort by chunk id to easily test if we have all chunks.\n for i, inference_chunk in enumerate(\n sorted(inference_chunks, key=lambda x: x.chunk_id)\n ):\n assert inference_chunk.chunk_id == i\n assert inference_chunk.document_id == doc_id\n\n # Under test.\n # Explicitly set empty fields here.\n user_fields = VespaDocumentUserFields(user_projects=[], personas=[])\n document_index.update_single(\n doc_id=doc_id,\n chunk_count=chunk_count,\n tenant_id=tenant_id,\n fields=None,\n user_fields=user_fields,\n )\n\n # Postcondition.\n filters = IndexFilters(access_control_list=None, tenant_id=tenant_id)\n # We should expect to get back all expected chunks with no filters.\n # Again, not best practice here.\n time.sleep(1)\n inference_chunks = document_index.id_based_retrieval(\n chunk_requests=[chunk_request], filters=filters, batch_retrieval=True\n )\n assert len(inference_chunks) == chunk_count\n # Sort by chunk id to easily test if we have all chunks.\n for i, inference_chunk in enumerate(\n sorted(inference_chunks, key=lambda x: x.chunk_id)\n ):\n assert inference_chunk.chunk_id == i\n assert inference_chunk.document_id == doc_id\n # Now, we should expect to not get any chunks if we specify the user\n # project and personas filters.\n inference_chunks = document_index.id_based_retrieval(\n chunk_requests=[chunk_request],\n filters=project_persona_filters,\n batch_retrieval=True,\n )\n assert len(inference_chunks) == 0\n\n def test_index_accepts_generator(\n self,\n document_indices: list[DocumentIndex],\n tenant_context: None, # noqa: ARG002\n ) -> None:\n \"\"\"index() accepts a generator (any iterable), not just a list.\"\"\"\n for document_index in document_indices:\n\n def chunk_gen() -> Iterator[DocMetadataAwareIndexChunk]:\n for i in range(3):\n yield make_chunk(\"test_doc_gen\", chunk_id=i)\n\n index_batch_params = IndexBatchParams(\n doc_id_to_previous_chunk_cnt={\"test_doc_gen\": 0},\n doc_id_to_new_chunk_cnt={\"test_doc_gen\": 3},\n tenant_id=get_current_tenant_id(),\n large_chunks_enabled=False,\n )\n\n results = document_index.index(chunk_gen(), index_batch_params)\n\n assert len(results) == 1\n record = results.pop()\n assert record.document_id == \"test_doc_gen\"\n assert record.already_existed is False\n" }, { "path": "backend/tests/external_dependency_unit/feature_flags/__init__.py", "content": "# External dependency unit tests for feature flag service\n" }, { "path": "backend/tests/external_dependency_unit/feature_flags/test_feature_flag_provider_factory.py", "content": "\"\"\"\nExternal dependency unit tests for the feature flag service.\n\nThese tests verify the feature flag service implementation with real\nPostHog integration when available, and fallback behavior otherwise.\n\"\"\"\n\nfrom uuid import UUID\n\nfrom ee.onyx.feature_flags.posthog_provider import PostHogFeatureFlagProvider\nfrom onyx.feature_flags.factory import get_default_feature_flag_provider\nfrom onyx.feature_flags.interface import FeatureFlagProvider\nfrom onyx.feature_flags.interface import NoOpFeatureFlagProvider\n\n\nclass TestNoOpFeatureFlagProvider:\n \"\"\"Tests for the no-op feature flag provider.\"\"\"\n\n def test_always_returns_false(self) -> None:\n \"\"\"No-op provider should always return False.\"\"\"\n provider = NoOpFeatureFlagProvider()\n\n my_uuid = UUID(\"79a75f76-6b63-43ee-b04c-a0c6806900bd\")\n assert provider.feature_enabled(\"another-flag\", my_uuid) is False\n\n\nclass TestFeatureFlagFactory:\n \"\"\"Tests for the feature flag factory function.\"\"\"\n\n def test_factory_returns_provider(self) -> None:\n \"\"\"Factory should return a FeatureFlagProvider instance.\"\"\"\n provider = get_default_feature_flag_provider()\n assert isinstance(provider, FeatureFlagProvider)\n\n def test_posthog_provider(self) -> None:\n \"\"\"Posthog provider should return True if the feature is enabled.\"\"\"\n provider = PostHogFeatureFlagProvider()\n assert isinstance(provider, FeatureFlagProvider)\n" }, { "path": "backend/tests/external_dependency_unit/file_store/test_file_store_non_mocked.py", "content": "import os\nimport time\nimport uuid\nfrom collections.abc import Generator\nfrom concurrent.futures import as_completed\nfrom concurrent.futures import ThreadPoolExecutor\nfrom io import BytesIO\nfrom typing import Any\nfrom typing import cast\nfrom typing import Dict\nfrom typing import List\nfrom typing import Tuple\nfrom typing import TypedDict\nfrom unittest.mock import patch\n\nimport pytest\nfrom botocore.exceptions import ClientError\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.file_store.file_store import S3BackedFileStore\nfrom onyx.utils.logger import setup_logger\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\nlogger = setup_logger()\n\n\nTEST_BUCKET_NAME: str = \"onyx-file-store-tests\"\nTEST_FILE_PREFIX: str = \"test-files\"\n\n\n# Type definitions for test data\nclass BackendConfig(TypedDict):\n endpoint_url: str | None\n access_key: str\n secret_key: str\n region: str\n verify_ssl: bool\n backend_name: str\n\n\nclass FileTestData(TypedDict):\n name: str\n display_name: str\n content: str\n type: str\n origin: FileOrigin\n\n\nclass WorkerResult(TypedDict):\n worker_id: int\n file_name: str\n content: str\n\n\ndef _get_all_backend_configs() -> List[BackendConfig]:\n \"\"\"Get configurations for all available backends\"\"\"\n from onyx.configs.app_configs import (\n S3_ENDPOINT_URL,\n AWS_REGION_NAME,\n )\n\n s3_aws_access_key_id = os.environ.get(\"S3_AWS_ACCESS_KEY_ID_FOR_TEST\")\n s3_aws_secret_access_key = os.environ.get(\"S3_AWS_SECRET_ACCESS_KEY_FOR_TEST\")\n\n configs: List[BackendConfig] = []\n\n # MinIO configuration (if endpoint is configured)\n if S3_ENDPOINT_URL:\n minio_access_key = \"minioadmin\"\n minio_secret_key = \"minioadmin\"\n configs.append(\n {\n \"endpoint_url\": S3_ENDPOINT_URL,\n \"access_key\": minio_access_key,\n \"secret_key\": minio_secret_key,\n \"region\": \"us-east-1\",\n \"verify_ssl\": False,\n \"backend_name\": \"MinIO\",\n }\n )\n\n # AWS S3 configuration (if credentials are available)\n if s3_aws_access_key_id and s3_aws_secret_access_key:\n configs.append(\n {\n \"endpoint_url\": None,\n \"access_key\": s3_aws_access_key_id,\n \"secret_key\": s3_aws_secret_access_key,\n \"region\": AWS_REGION_NAME or \"us-east-2\",\n \"verify_ssl\": True,\n \"backend_name\": \"AWS S3\",\n }\n )\n\n if not configs:\n pytest.skip(\n \"No backend configurations available - set MinIO or AWS S3 credentials\"\n )\n\n return configs\n\n\n@pytest.fixture(\n scope=\"function\",\n params=_get_all_backend_configs(),\n ids=lambda config: config[\"backend_name\"],\n)\ndef file_store(\n request: pytest.FixtureRequest,\n db_session: Session, # noqa: ARG001\n tenant_context: None, # noqa: ARG001\n) -> Generator[S3BackedFileStore, None, None]:\n \"\"\"Create an S3BackedFileStore instance for testing with parametrized backend\"\"\"\n backend_config: BackendConfig = request.param\n\n # Create S3BackedFileStore with backend-specific configuration\n store = S3BackedFileStore(\n bucket_name=TEST_BUCKET_NAME,\n aws_access_key_id=backend_config[\"access_key\"],\n aws_secret_access_key=backend_config[\"secret_key\"],\n aws_region_name=backend_config[\"region\"],\n s3_endpoint_url=backend_config[\"endpoint_url\"],\n s3_prefix=f\"{TEST_FILE_PREFIX}-{uuid.uuid4()}\",\n s3_verify_ssl=backend_config[\"verify_ssl\"],\n )\n\n # Initialize the store and ensure bucket exists\n store.initialize()\n logger.info(\n f\"Successfully initialized {backend_config['backend_name']} file store with bucket {TEST_BUCKET_NAME}\"\n )\n\n yield store\n\n # Cleanup: Remove all test files from the bucket (including tenant-prefixed files)\n try:\n s3_client = store._get_s3_client()\n actual_bucket_name = store._get_bucket_name()\n\n # List and delete all objects in the test prefix (including tenant subdirectories)\n response = s3_client.list_objects_v2(\n Bucket=actual_bucket_name, Prefix=f\"{store._s3_prefix}/\"\n )\n\n if \"Contents\" in response:\n objects_to_delete = [{\"Key\": obj[\"Key\"]} for obj in response[\"Contents\"]]\n s3_client.delete_objects(\n Bucket=actual_bucket_name,\n Delete={\"Objects\": objects_to_delete}, # type: ignore[typeddict-item]\n )\n logger.info(\n f\"Cleaned up {len(objects_to_delete)} test objects from {backend_config['backend_name']}\"\n )\n except Exception as e:\n logger.warning(f\"Failed to cleanup test objects: {e}\")\n\n\nclass TestS3BackedFileStore:\n \"\"\"Test suite for S3BackedFileStore using real S3-compatible storage (MinIO or AWS S3)\"\"\"\n\n def test_store_initialization(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test that the file store initializes properly\"\"\"\n # The fixture already calls initialize(), so we just verify it worked\n bucket_name = file_store._get_bucket_name()\n assert bucket_name.startswith(TEST_BUCKET_NAME) # Should be backend-specific\n\n # Verify bucket exists by trying to list objects\n s3_client = file_store._get_s3_client()\n\n # This should not raise an exception\n s3_client.list_objects_v2(Bucket=bucket_name, MaxKeys=1)\n\n def test_save_and_read_text_file(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test saving and reading a text file\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Text File\"\n content = \"This is a test text file content.\\nWith multiple lines.\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Save the file\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Read the file back\n read_content_io = file_store.read_file(file_id)\n read_content = read_content_io.read().decode(\"utf-8\")\n\n assert read_content == content\n\n # Verify file record in database\n file_record = file_store.read_file_record(file_id)\n assert file_record.file_id == file_id\n assert file_record.display_name == display_name\n assert file_record.file_origin == file_origin\n assert file_record.file_type == file_type\n assert (\n file_record.bucket_name == file_store._get_bucket_name()\n ) # Use actual bucket name\n # The object key should include the tenant ID\n expected_object_key = f\"{file_store._s3_prefix}/{TEST_TENANT_ID}/{file_id}\"\n assert file_record.object_key == expected_object_key\n\n def test_save_and_read_binary_file(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test saving and reading a binary file\"\"\"\n file_id = f\"{uuid.uuid4()}.bin\"\n display_name = \"Test Binary File\"\n # Create some binary content\n content = bytes(range(256)) # 0-255 bytes\n file_type = \"application/octet-stream\"\n file_origin = FileOrigin.CONNECTOR\n\n # Save the file\n content_io = BytesIO(content)\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Read the file back\n read_content_io = file_store.read_file(file_id)\n read_content = read_content_io.read()\n\n assert read_content == content\n\n def test_save_with_metadata(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test saving a file with metadata\"\"\"\n file_id = f\"{uuid.uuid4()}.json\"\n display_name = \"Test Metadata File\"\n content = '{\"key\": \"value\", \"number\": 42}'\n file_type = \"application/json\"\n file_origin = FileOrigin.CHAT_UPLOAD\n metadata: Dict[str, Any] = {\n \"source\": \"test_suite\",\n \"version\": \"1.0\",\n \"tags\": [\"test\", \"json\"],\n \"size\": len(content),\n }\n\n # Save the file with metadata\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_metadata=metadata,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Verify metadata is stored in database\n file_record = file_store.read_file_record(file_id)\n assert file_record.file_metadata == metadata\n\n def test_has_file(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test the has_file method\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Has File\"\n content = \"Content for has_file test\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Initially, file should not exist\n assert not file_store.has_file(\n file_id=file_id,\n file_origin=file_origin,\n file_type=file_type,\n )\n\n # Save the file\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Now file should exist\n assert file_store.has_file(\n file_id=file_id,\n file_origin=file_origin,\n file_type=file_type,\n )\n\n # Test with wrong parameters\n assert not file_store.has_file(\n file_id=file_id,\n file_origin=FileOrigin.CONNECTOR, # Wrong origin\n file_type=file_type,\n )\n\n assert not file_store.has_file(\n file_id=file_id,\n file_origin=file_origin,\n file_type=\"application/pdf\", # Wrong type\n )\n\n def test_read_file_with_tempfile(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test reading a file using temporary file\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Temp File\"\n content = \"Content for temporary file test\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Save the file\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Read using temporary file\n temp_file = file_store.read_file(file_id, use_tempfile=True)\n\n # Read content from temp file\n temp_file.seek(0)\n read_content_bytes = temp_file.read()\n if isinstance(read_content_bytes, bytes):\n read_content_str = read_content_bytes.decode(\"utf-8\")\n else:\n read_content_str = str(read_content_bytes)\n\n assert read_content_str == content\n\n # Clean up the temp file\n temp_file.close()\n if hasattr(temp_file, \"name\"):\n try:\n os.unlink(temp_file.name)\n except (OSError, AttributeError):\n pass\n\n def test_delete_file(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test deleting a file\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Delete File\"\n content = \"Content for delete test\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Save the file\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Verify file exists\n assert file_store.has_file(\n file_id=file_id,\n file_origin=file_origin,\n file_type=file_type,\n )\n\n # Delete the file\n file_store.delete_file(file_id)\n\n # Verify file no longer exists\n assert not file_store.has_file(\n file_id=file_id,\n file_origin=file_origin,\n file_type=file_type,\n )\n\n # Verify trying to read deleted file raises exception\n with pytest.raises(RuntimeError, match=\"does not exist or was deleted\"):\n file_store.read_file(file_id)\n\n def test_get_file_with_mime_type(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test getting file with mime type detection\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test MIME Type\"\n content = \"This is a plain text file\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Save the file\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Get file with mime type\n file_with_mime = file_store.get_file_with_mime_type(file_id)\n\n assert file_with_mime is not None\n assert file_with_mime.data.decode(\"utf-8\") == content\n # The detected mime type might be different from what we stored\n assert file_with_mime.mime_type is not None\n\n def test_file_overwrite(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test overwriting an existing file\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Overwrite\"\n original_content = \"Original content\"\n new_content = \"New content after overwrite\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Save original file\n content_io = BytesIO(original_content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Verify original content\n read_content_io = file_store.read_file(file_id)\n assert read_content_io.read().decode(\"utf-8\") == original_content\n\n # Overwrite with new content\n new_content_io = BytesIO(new_content.encode(\"utf-8\"))\n returned_file_id_2 = file_store.save_file(\n content=new_content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id_2 == file_id\n\n # Verify new content\n read_content_io = file_store.read_file(file_id)\n assert read_content_io.read().decode(\"utf-8\") == new_content\n\n def test_large_file_handling(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test handling of larger files\"\"\"\n file_id = f\"{uuid.uuid4()}.bin\"\n display_name = \"Test Large File\"\n # Create a 1MB file\n content_size = 1024 * 1024 # 1MB\n content = b\"A\" * content_size\n file_type = \"application/octet-stream\"\n file_origin = FileOrigin.CONNECTOR\n\n # Save the large file\n content_io = BytesIO(content)\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Read the file back\n read_content_io = file_store.read_file(file_id)\n read_content = read_content_io.read()\n\n assert len(read_content) == content_size\n assert read_content == content\n\n def test_error_handling_nonexistent_file(\n self, file_store: S3BackedFileStore\n ) -> None:\n \"\"\"Test error handling when trying to read a non-existent file\"\"\"\n nonexistent_file_id = f\"{uuid.uuid4()}.txt\"\n\n with pytest.raises(RuntimeError, match=\"does not exist or was deleted\"):\n file_store.read_file(nonexistent_file_id)\n\n with pytest.raises(RuntimeError, match=\"does not exist or was deleted\"):\n file_store.read_file_record(nonexistent_file_id)\n\n # get_file_with_mime_type should return None for non-existent files\n result = file_store.get_file_with_mime_type(nonexistent_file_id)\n assert result is None\n\n def test_error_handling_delete_nonexistent_file(\n self, file_store: S3BackedFileStore\n ) -> None:\n \"\"\"Test error handling when trying to delete a non-existent file\"\"\"\n nonexistent_file_id = f\"{uuid.uuid4()}.txt\"\n\n # Should raise an exception when trying to delete non-existent file\n with pytest.raises(RuntimeError, match=\"does not exist or was deleted\"):\n file_store.delete_file(nonexistent_file_id)\n\n def test_multiple_files_different_origins(\n self, file_store: S3BackedFileStore\n ) -> None:\n \"\"\"Test storing multiple files with different origins and types\"\"\"\n files_data: List[FileTestData] = [\n {\n \"name\": f\"{uuid.uuid4()}.txt\",\n \"display_name\": \"Chat Upload File\",\n \"content\": \"Content from chat upload\",\n \"type\": \"text/plain\",\n \"origin\": FileOrigin.CHAT_UPLOAD,\n },\n {\n \"name\": f\"{uuid.uuid4()}.json\",\n \"display_name\": \"Connector File\",\n \"content\": '{\"from\": \"connector\"}',\n \"type\": \"application/json\",\n \"origin\": FileOrigin.CONNECTOR,\n },\n {\n \"name\": f\"{uuid.uuid4()}.csv\",\n \"display_name\": \"Generated Report\",\n \"content\": \"col1,col2\\nval1,val2\",\n \"type\": \"text/csv\",\n \"origin\": FileOrigin.GENERATED_REPORT,\n },\n ]\n\n # Save all files\n for file_data in files_data:\n content_io = BytesIO(file_data[\"content\"].encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=file_data[\"display_name\"],\n file_origin=file_data[\"origin\"],\n file_type=file_data[\"type\"],\n file_id=file_data[\"name\"],\n )\n assert returned_file_id == file_data[\"name\"]\n\n # Verify all files exist and have correct properties\n for file_data in files_data:\n assert file_store.has_file(\n file_id=file_data[\"name\"],\n file_origin=file_data[\"origin\"],\n file_type=file_data[\"type\"],\n )\n\n # Read and verify content\n read_content_io = file_store.read_file(file_data[\"name\"])\n read_content = read_content_io.read().decode(\"utf-8\")\n assert read_content == file_data[\"content\"]\n\n # Verify record\n file_record = file_store.read_file_record(file_data[\"name\"])\n assert file_record.file_origin == file_data[\"origin\"]\n assert file_record.file_type == file_data[\"type\"]\n\n def test_special_characters_in_filenames(\n self, file_store: S3BackedFileStore\n ) -> None:\n \"\"\"Test handling of special characters in filenames\"\"\"\n # Note: S3 keys have some restrictions, so we test reasonable special characters\n special_files: List[str] = [\n f\"{uuid.uuid4()} with spaces.txt\",\n f\"{uuid.uuid4()}-with-dashes.txt\",\n f\"{uuid.uuid4()}_with_underscores.txt\",\n f\"{uuid.uuid4()}.with.dots.txt\",\n f\"{uuid.uuid4()}(with)parentheses.txt\",\n ]\n\n for file_id in special_files:\n content = f\"Content for {file_id}\"\n content_io = BytesIO(content.encode(\"utf-8\"))\n\n # Save the file\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=f\"Display: {file_id}\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Read and verify\n read_content_io = file_store.read_file(file_id)\n read_content = read_content_io.read().decode(\"utf-8\")\n assert read_content == content\n\n @pytest.mark.skipif(\n not os.environ.get(\"TEST_S3_NETWORK_ERRORS\"),\n reason=\"Network error tests require TEST_S3_NETWORK_ERRORS environment variable\",\n )\n def test_network_error_handling(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test handling of network errors (requires special setup)\"\"\"\n # This test requires specific network configuration to simulate failures\n # It's marked as skip by default and only runs when explicitly enabled\n\n # Mock a network error during file operations\n with patch.object(file_store, \"_get_s3_client\") as mock_client:\n mock_s3 = mock_client.return_value\n mock_s3.put_object.side_effect = ClientError(\n error_response={\n \"Error\": {\n \"Code\": \"NetworkingError\",\n \"Message\": \"Connection timeout\",\n }\n },\n operation_name=\"PutObject\",\n )\n\n content_io = BytesIO(b\"test content\")\n\n with pytest.raises(ClientError):\n file_store.save_file(\n content=content_io,\n display_name=\"Network Error Test\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=f\"{uuid.uuid4()}.txt\",\n )\n\n def test_database_transaction_rollback(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test database transaction rollback behavior with PostgreSQL\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Rollback\"\n content = \"Content for rollback test\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Mock S3 to fail after database write but before commit\n with patch.object(file_store, \"_get_s3_client\") as mock_client:\n mock_s3 = mock_client.return_value\n mock_s3.put_object.side_effect = ClientError(\n error_response={\n \"Error\": {\"Code\": \"InternalError\", \"Message\": \"S3 internal error\"}\n },\n operation_name=\"PutObject\",\n )\n\n content_io = BytesIO(content.encode(\"utf-8\"))\n\n # This should fail and rollback the database transaction\n with pytest.raises(ClientError):\n file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n # Verify that the database record was not created due to rollback\n with pytest.raises(RuntimeError, match=\"does not exist or was deleted\"):\n file_store.read_file_record(file_id)\n\n def test_complex_jsonb_metadata(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test PostgreSQL JSONB metadata handling with complex data structures\"\"\"\n file_id = f\"{uuid.uuid4()}.json\"\n display_name = \"Test Complex Metadata\"\n content = '{\"data\": \"test\"}'\n file_type = \"application/json\"\n file_origin = FileOrigin.CONNECTOR\n\n # Complex metadata that tests PostgreSQL JSONB capabilities\n complex_metadata: Dict[str, Any] = {\n \"nested\": {\n \"array\": [1, 2, 3, {\"inner\": \"value\"}],\n \"boolean\": True,\n \"null_value\": None,\n \"number\": 42.5,\n },\n \"unicode\": \"测试数据 🚀\",\n \"special_chars\": \"Line 1\\nLine 2\\t\\r\\nSpecial: !@#$%^&*()\",\n \"large_text\": \"x\" * 1000, # Test large text in JSONB\n \"timestamps\": {\n \"created\": \"2024-01-01T00:00:00Z\",\n \"updated\": \"2024-01-02T12:30:45Z\",\n },\n }\n\n # Save file with complex metadata\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_metadata=complex_metadata,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Retrieve and verify the metadata was stored correctly\n file_record = file_store.read_file_record(file_id)\n stored_metadata = file_record.file_metadata\n\n # Verify all metadata fields were preserved\n assert stored_metadata == complex_metadata\n\n # Type casting for complex metadata access\n stored_metadata_dict = cast(Dict[str, Any], stored_metadata)\n nested_data = cast(Dict[str, Any], stored_metadata_dict[\"nested\"])\n array_data = cast(List[Any], nested_data[\"array\"])\n inner_obj = cast(Dict[str, Any], array_data[3])\n\n assert inner_obj[\"inner\"] == \"value\"\n assert stored_metadata_dict[\"unicode\"] == \"测试数据 🚀\"\n assert nested_data[\"boolean\"] is True\n assert nested_data[\"null_value\"] is None\n assert len(cast(str, stored_metadata_dict[\"large_text\"])) == 1000\n\n def test_database_consistency_after_s3_failure(\n self, file_store: S3BackedFileStore\n ) -> None:\n \"\"\"Test that database stays consistent when S3 operations fail\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test Consistency\"\n content = \"Initial content\"\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # First, save a file successfully\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Verify initial state\n assert file_store.has_file(file_id, file_origin, file_type)\n initial_record = file_store.read_file_record(file_id)\n\n # Now try to update but fail on S3 side\n with patch.object(file_store, \"_get_s3_client\") as mock_client:\n mock_s3 = mock_client.return_value\n # Let the first call (for reading/checking) succeed, but fail on put_object\n mock_s3.put_object.side_effect = ClientError(\n error_response={\n \"Error\": {\n \"Code\": \"ServiceUnavailable\",\n \"Message\": \"Service temporarily unavailable\",\n }\n },\n operation_name=\"PutObject\",\n )\n\n new_content = \"Updated content that should fail\"\n new_content_io = BytesIO(new_content.encode(\"utf-8\"))\n\n # This should fail and rollback\n with pytest.raises(ClientError):\n file_store.save_file(\n content=new_content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n # Verify the database record is unchanged (not updated)\n current_record = file_store.read_file_record(file_id)\n assert current_record.file_id == initial_record.file_id\n assert current_record.display_name == initial_record.display_name\n assert current_record.bucket_name == initial_record.bucket_name\n assert current_record.object_key == initial_record.object_key\n\n # Verify we can still read the original file content\n read_content_io = file_store.read_file(file_id)\n read_content = read_content_io.read().decode(\"utf-8\")\n assert read_content == content # Original content, not the failed update\n\n def test_concurrent_file_operations(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test handling of concurrent file operations on the same file\"\"\"\n base_file_name: str = str(uuid.uuid4())\n file_type: str = \"text/plain\"\n file_origin: FileOrigin = FileOrigin.OTHER\n\n # Get current file store configuration to replicate in workers\n current_bucket_name = file_store._get_bucket_name()\n current_access_key = file_store._aws_access_key_id\n current_secret_key = file_store._aws_secret_access_key\n current_region = file_store._aws_region_name\n current_endpoint_url = file_store._s3_endpoint_url\n current_verify_ssl = file_store._s3_verify_ssl\n\n results: List[Tuple[str, str]] = []\n errors: List[Tuple[int, str]] = []\n\n def save_file_worker(worker_id: int) -> bool:\n \"\"\"Worker function to save a file with its own database session\"\"\"\n try:\n # Set up tenant context for this worker\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(TEST_TENANT_ID)\n try:\n # Create a new database session for each worker to avoid conflicts\n with get_session_with_current_tenant() as worker_session:\n worker_file_store = S3BackedFileStore(\n bucket_name=current_bucket_name,\n aws_access_key_id=current_access_key,\n aws_secret_access_key=current_secret_key,\n aws_region_name=current_region,\n s3_endpoint_url=current_endpoint_url,\n s3_prefix=TEST_FILE_PREFIX,\n s3_verify_ssl=current_verify_ssl,\n )\n\n file_name: str = f\"{base_file_name}_{worker_id}.txt\"\n content: str = (\n f\"Content from worker {worker_id} at {time.time()}\"\n )\n content_io: BytesIO = BytesIO(content.encode(\"utf-8\"))\n\n worker_file_store.save_file(\n file_id=file_name,\n content=content_io,\n display_name=f\"Worker {worker_id} File\",\n file_origin=file_origin,\n file_type=file_type,\n db_session=worker_session,\n )\n results.append((file_name, content))\n return True\n finally:\n # Reset the tenant context after the worker completes\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n except Exception as e:\n errors.append((worker_id, str(e)))\n return False\n\n # Run multiple concurrent file save operations\n with ThreadPoolExecutor(max_workers=5) as executor:\n futures = [executor.submit(save_file_worker, i) for i in range(10)]\n\n for future in as_completed(futures):\n future.result() # Wait for completion\n\n # Verify all operations completed successfully\n assert len(errors) == 0, f\"Concurrent operations had errors: {errors}\"\n assert (\n len(results) == 10\n ), f\"Expected 10 successful operations, got {len(results)}\"\n\n # Verify all files were saved correctly\n for file_id, expected_content in results:\n # Check file exists\n assert file_store.has_file(\n file_id=file_id,\n file_origin=file_origin,\n file_type=file_type,\n )\n\n # Check content is correct\n read_content_io = file_store.read_file(file_id)\n actual_content: str = read_content_io.read().decode(\"utf-8\")\n assert actual_content == expected_content\n\n def test_list_files_by_prefix(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test listing files by prefix returns only correctly prefixed files\"\"\"\n test_prefix = \"documents-batch-\"\n\n # Files that should be returned (start with the prefix)\n prefixed_files: List[str] = [\n f\"{test_prefix}001.txt\",\n f\"{test_prefix}002.json\",\n f\"{test_prefix}abc.pdf\",\n f\"{test_prefix}xyz-final.docx\",\n ]\n\n # Files that should NOT be returned (don't start with prefix, even if they contain it)\n non_prefixed_files: List[str] = [\n f\"other-{test_prefix}001.txt\", # Contains prefix but doesn't start with it\n f\"backup-{test_prefix}data.txt\", # Contains prefix but doesn't start with it\n f\"{uuid.uuid4()}.txt\", # Random file without prefix\n \"reports-001.pdf\", # Different prefix\n f\"my-{test_prefix[:-1]}.txt\", # Similar but not exact prefix\n ]\n\n all_files = prefixed_files + non_prefixed_files\n saved_file_ids: List[str] = []\n\n # Save all test files\n for file_name in all_files:\n content = f\"Content for {file_name}\"\n content_io = BytesIO(content.encode(\"utf-8\"))\n\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=f\"Display: {file_name}\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_name,\n )\n saved_file_ids.append(returned_file_id)\n\n # Verify file was saved\n assert returned_file_id == file_name\n\n # Test the list_files_by_prefix functionality\n prefix_results = file_store.list_files_by_prefix(test_prefix)\n\n # Extract file IDs from results\n returned_file_ids = [record.file_id for record in prefix_results]\n\n # Verify correct number of files returned\n assert len(returned_file_ids) == len(prefixed_files), (\n f\"Expected {len(prefixed_files)} files with prefix '{test_prefix}', \"\n f\"but got {len(returned_file_ids)}: {returned_file_ids}\"\n )\n\n # Verify all prefixed files are returned\n for expected_file_id in prefixed_files:\n assert (\n expected_file_id in returned_file_ids\n ), f\"File '{expected_file_id}' should be in results but was not found. Returned files: {returned_file_ids}\"\n\n # Verify no non-prefixed files are returned\n for unexpected_file_id in non_prefixed_files:\n assert (\n unexpected_file_id not in returned_file_ids\n ), f\"File '{unexpected_file_id}' should NOT be in results but was found. Returned files: {returned_file_ids}\"\n\n # Verify the returned records have correct properties\n for record in prefix_results:\n assert record.file_id.startswith(test_prefix)\n assert record.display_name == f\"Display: {record.file_id}\"\n assert record.file_origin == FileOrigin.OTHER\n assert record.file_type == \"text/plain\"\n assert record.bucket_name == file_store._get_bucket_name()\n\n # Test with empty prefix (should return all files we created)\n all_results = file_store.list_files_by_prefix(\"\")\n all_returned_ids = [record.file_id for record in all_results]\n\n # Should include all our test files\n for file_id in saved_file_ids:\n assert (\n file_id in all_returned_ids\n ), f\"File '{file_id}' should be in results for empty prefix\"\n\n # Test with non-existent prefix\n nonexistent_results = file_store.list_files_by_prefix(\"nonexistent-prefix-\")\n assert (\n len(nonexistent_results) == 0\n ), \"Should return empty list for non-existent prefix\"\n\n def test_get_file_size(self, file_store: S3BackedFileStore) -> None:\n \"\"\"Test getting file size from S3\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n display_name = \"Test File Size\"\n content = \"This is test content for file size check.\"\n expected_size = len(content.encode(\"utf-8\"))\n file_type = \"text/plain\"\n file_origin = FileOrigin.OTHER\n\n # Save the file\n content_io = BytesIO(content.encode(\"utf-8\"))\n returned_file_id = file_store.save_file(\n content=content_io,\n display_name=display_name,\n file_origin=file_origin,\n file_type=file_type,\n file_id=file_id,\n )\n\n assert returned_file_id == file_id\n\n # Get file size\n file_size = file_store.get_file_size(file_id)\n\n assert file_size is not None\n assert file_size == expected_size\n\n def test_get_file_size_nonexistent_file(\n self, file_store: S3BackedFileStore\n ) -> None:\n \"\"\"Test getting file size for a non-existent file returns None\"\"\"\n nonexistent_file_id = f\"{uuid.uuid4()}.txt\"\n\n file_size = file_store.get_file_size(nonexistent_file_id)\n\n assert file_size is None\n" }, { "path": "backend/tests/external_dependency_unit/file_store/test_postgres_file_store_non_mocked.py", "content": "\"\"\"External dependency tests for PostgresBackedFileStore.\n\nThese tests interact with a real PostgreSQL database — no mocking.\nThey exercise Large Object creation, reading, streaming, deletion,\nand verify consistency between the file_record / file_content tables\nand the underlying pg_largeobject storage.\n\"\"\"\n\nimport uuid\nfrom collections.abc import Generator\nfrom io import BytesIO\nfrom io import StringIO\nfrom typing import Any\nfrom typing import Dict\nfrom typing import List\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import FileOrigin\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.file_content import get_file_content_by_file_id\nfrom onyx.db.file_content import get_file_content_by_file_id_optional\nfrom onyx.file_store.postgres_file_store import _get_raw_connection\nfrom onyx.file_store.postgres_file_store import _read_large_object\nfrom onyx.file_store.postgres_file_store import POSTGRES_BUCKET_SENTINEL\nfrom onyx.file_store.postgres_file_store import PostgresBackedFileStore\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\n# ------------------------------------------------------------------ fixtures --\n\n\n@pytest.fixture(scope=\"function\")\ndef pg_file_store(\n db_session: Session, # noqa: ARG001 — ensures engine is ready\n tenant_context: None, # noqa: ARG001\n) -> Generator[PostgresBackedFileStore, None, None]:\n \"\"\"Provide a PostgresBackedFileStore wired to the real test database.\"\"\"\n store = PostgresBackedFileStore()\n store.initialize()\n\n # Track file IDs so we can clean up after each test\n created_ids: list[str] = []\n original_save = store.save_file\n\n def _tracking_save(*args: Any, **kwargs: Any) -> str:\n file_id = original_save(*args, **kwargs)\n created_ids.append(file_id)\n return file_id\n\n store.save_file = _tracking_save # type: ignore[method-assign]\n\n yield store\n\n # Cleanup: delete every file we created (including Large Objects)\n for fid in created_ids:\n try:\n store.delete_file(fid)\n except Exception:\n pass\n\n\n# -------------------------------------------------------------------- tests --\n\n\nclass TestPostgresBackedFileStore:\n \"\"\"Full integration tests against a real PostgreSQL instance.\"\"\"\n\n # ── basic save / read ──────────────────────────────────────────\n\n def test_save_and_read_text_file(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n content = \"Hello, Postgres Large Objects!\"\n\n returned_id = pg_file_store.save_file(\n content=BytesIO(content.encode()),\n display_name=\"greeting.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert returned_id == file_id\n\n result = pg_file_store.read_file(file_id)\n assert result.read().decode() == content\n\n def test_save_and_read_binary_file(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n file_id = f\"{uuid.uuid4()}.bin\"\n content = bytes(range(256))\n\n pg_file_store.save_file(\n content=BytesIO(content),\n display_name=\"binary.bin\",\n file_origin=FileOrigin.CONNECTOR,\n file_type=\"application/octet-stream\",\n file_id=file_id,\n )\n\n assert pg_file_store.read_file(file_id).read() == content\n\n def test_save_string_io(self, pg_file_store: PostgresBackedFileStore) -> None:\n \"\"\"StringIO content should be transparently UTF-8 encoded.\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n text = \"StringIO content — including unicode: 测试 🚀\"\n\n pg_file_store.save_file(\n content=StringIO(text),\n display_name=\"stringio.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert pg_file_store.read_file(file_id).read().decode() == text\n\n def test_auto_generated_file_id(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n \"\"\"When no file_id is supplied, a UUID should be generated.\"\"\"\n returned_id = pg_file_store.save_file(\n content=BytesIO(b\"auto-id\"),\n display_name=\"auto.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n )\n\n # Should be a valid UUID\n uuid.UUID(returned_id)\n assert pg_file_store.read_file(returned_id).read() == b\"auto-id\"\n\n # ── read with tempfile (streaming) ─────────────────────────────\n\n def test_read_file_with_tempfile(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n content = \"Streamed via tempfile\"\n\n pg_file_store.save_file(\n content=BytesIO(content.encode()),\n display_name=\"streamed.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n tmp = pg_file_store.read_file(file_id, use_tempfile=True)\n try:\n tmp.seek(0)\n assert tmp.read().decode() == content\n finally:\n tmp.close()\n\n # ── file record metadata ───────────────────────────────────────\n\n def test_file_record_fields(self, pg_file_store: PostgresBackedFileStore) -> None:\n file_id = f\"{uuid.uuid4()}.json\"\n metadata: Dict[str, Any] = {\"source\": \"test\", \"version\": 1}\n\n pg_file_store.save_file(\n content=BytesIO(b'{\"k\":\"v\"}'),\n display_name=\"meta.json\",\n file_origin=FileOrigin.CHAT_UPLOAD,\n file_type=\"application/json\",\n file_metadata=metadata,\n file_id=file_id,\n )\n\n record = pg_file_store.read_file_record(file_id)\n assert record.file_id == file_id\n assert record.display_name == \"meta.json\"\n assert record.file_origin == FileOrigin.CHAT_UPLOAD\n assert record.file_type == \"application/json\"\n assert record.file_metadata == metadata\n assert record.bucket_name == POSTGRES_BUCKET_SENTINEL\n\n # object_key should be the stringified Large Object OID\n oid = int(record.object_key)\n assert oid > 0\n\n def test_file_content_record(self, pg_file_store: PostgresBackedFileStore) -> None:\n \"\"\"file_content row should track the OID and byte-size.\"\"\"\n file_id = f\"{uuid.uuid4()}.txt\"\n payload = b\"measure my size\"\n\n pg_file_store.save_file(\n content=BytesIO(payload),\n display_name=\"sized.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n with get_session_with_current_tenant() as session:\n fc = get_file_content_by_file_id(file_id, session)\n assert fc.file_size == len(payload)\n assert fc.lobj_oid > 0\n\n # ── has_file ───────────────────────────────────────────────────\n\n def test_has_file(self, pg_file_store: PostgresBackedFileStore) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n\n assert not pg_file_store.has_file(file_id, FileOrigin.OTHER, \"text/plain\")\n\n pg_file_store.save_file(\n content=BytesIO(b\"exists\"),\n display_name=\"exists.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert pg_file_store.has_file(file_id, FileOrigin.OTHER, \"text/plain\")\n # Wrong origin / type → False\n assert not pg_file_store.has_file(file_id, FileOrigin.CONNECTOR, \"text/plain\")\n assert not pg_file_store.has_file(file_id, FileOrigin.OTHER, \"image/png\")\n\n # ── get_file_size ──────────────────────────────────────────────\n\n def test_get_file_size(self, pg_file_store: PostgresBackedFileStore) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n payload = b\"exactly 24 bytes long!?!\"\n\n pg_file_store.save_file(\n content=BytesIO(payload),\n display_name=\"sized.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert pg_file_store.get_file_size(file_id) == len(payload)\n\n def test_get_file_size_nonexistent(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n assert pg_file_store.get_file_size(f\"{uuid.uuid4()}\") is None\n\n # ── delete ─────────────────────────────────────────────────────\n\n def test_delete_file(self, pg_file_store: PostgresBackedFileStore) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n\n pg_file_store.save_file(\n content=BytesIO(b\"delete me\"),\n display_name=\"doomed.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n pg_file_store.delete_file(file_id)\n\n assert not pg_file_store.has_file(file_id, FileOrigin.OTHER, \"text/plain\")\n\n with pytest.raises(RuntimeError, match=\"does not exist\"):\n pg_file_store.read_file(file_id)\n\n # file_content row should also be gone\n with get_session_with_current_tenant() as session:\n assert get_file_content_by_file_id_optional(file_id, session) is None\n\n def test_delete_nonexistent_raises(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n with pytest.raises(RuntimeError, match=\"does not exist\"):\n pg_file_store.delete_file(f\"{uuid.uuid4()}\")\n\n # ── overwrite (upsert) ─────────────────────────────────────────\n\n def test_overwrite_file(self, pg_file_store: PostgresBackedFileStore) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n\n pg_file_store.save_file(\n content=BytesIO(b\"original\"),\n display_name=\"v1.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert pg_file_store.read_file(file_id).read() == b\"original\"\n\n # Capture the OID of the original Large Object\n with get_session_with_current_tenant() as session:\n old_oid = get_file_content_by_file_id(file_id, session).lobj_oid\n\n pg_file_store.save_file(\n content=BytesIO(b\"overwritten\"),\n display_name=\"v2.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n assert pg_file_store.read_file(file_id).read() == b\"overwritten\"\n\n # The old Large Object should have been unlinked\n with get_session_with_current_tenant() as session:\n new_oid = get_file_content_by_file_id(file_id, session).lobj_oid\n assert new_oid != old_oid\n\n raw_conn = _get_raw_connection(session)\n with pytest.raises(Exception):\n _read_large_object(raw_conn, old_oid)\n\n # ── change_file_id ─────────────────────────────────────────────\n\n def test_change_file_id(self, pg_file_store: PostgresBackedFileStore) -> None:\n old_id = f\"{uuid.uuid4()}.txt\"\n new_id = f\"{uuid.uuid4()}.txt\"\n content = b\"portable content\"\n\n pg_file_store.save_file(\n content=BytesIO(content),\n display_name=\"rename.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=old_id,\n )\n\n pg_file_store.change_file_id(old_id, new_id)\n\n # Old ID should be gone\n assert not pg_file_store.has_file(old_id, FileOrigin.OTHER, \"text/plain\")\n\n # New ID should serve the same content\n assert pg_file_store.read_file(new_id).read() == content\n assert pg_file_store.get_file_size(new_id) == len(content)\n\n # Clean up the renamed file (fixture only tracks save_file calls)\n pg_file_store.delete_file(new_id)\n\n # ── list_files_by_prefix ───────────────────────────────────────\n\n def test_list_files_by_prefix(self, pg_file_store: PostgresBackedFileStore) -> None:\n prefix = f\"batch-{uuid.uuid4().hex[:8]}-\"\n\n # Create files with and without the prefix\n for i in range(3):\n pg_file_store.save_file(\n content=BytesIO(f\"prefixed-{i}\".encode()),\n display_name=f\"p{i}.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=f\"{prefix}{i}.txt\",\n )\n\n pg_file_store.save_file(\n content=BytesIO(b\"unrelated\"),\n display_name=\"other.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=f\"other-{uuid.uuid4()}.txt\",\n )\n\n results = pg_file_store.list_files_by_prefix(prefix)\n returned_ids = [r.file_id for r in results]\n\n assert len(returned_ids) == 3\n for i in range(3):\n assert f\"{prefix}{i}.txt\" in returned_ids\n\n # ── get_file_with_mime_type ────────────────────────────────────\n\n def test_get_file_with_mime_type(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n file_id = f\"{uuid.uuid4()}.txt\"\n\n pg_file_store.save_file(\n content=BytesIO(b\"plain text\"),\n display_name=\"mime.txt\",\n file_origin=FileOrigin.OTHER,\n file_type=\"text/plain\",\n file_id=file_id,\n )\n\n result = pg_file_store.get_file_with_mime_type(file_id)\n assert result is not None\n assert result.data == b\"plain text\"\n assert result.mime_type is not None\n\n def test_get_file_with_mime_type_nonexistent(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n assert pg_file_store.get_file_with_mime_type(f\"{uuid.uuid4()}\") is None\n\n # ── error handling ─────────────────────────────────────────────\n\n def test_read_nonexistent_raises(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n with pytest.raises(RuntimeError, match=\"does not exist\"):\n pg_file_store.read_file(f\"{uuid.uuid4()}\")\n\n def test_read_file_record_nonexistent_raises(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n with pytest.raises(RuntimeError, match=\"does not exist\"):\n pg_file_store.read_file_record(f\"{uuid.uuid4()}\")\n\n # ── large file ─────────────────────────────────────────────────\n\n def test_large_file_roundtrip(self, pg_file_store: PostgresBackedFileStore) -> None:\n \"\"\"Verify a 1 MB payload survives a full save / read cycle.\"\"\"\n file_id = f\"{uuid.uuid4()}.bin\"\n content = b\"X\" * (1024 * 1024)\n\n pg_file_store.save_file(\n content=BytesIO(content),\n display_name=\"big.bin\",\n file_origin=FileOrigin.CONNECTOR,\n file_type=\"application/octet-stream\",\n file_id=file_id,\n )\n\n assert pg_file_store.read_file(file_id).read() == content\n assert pg_file_store.get_file_size(file_id) == len(content)\n\n # ── multiple files with different origins ──────────────────────\n\n def test_multiple_files_different_origins(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n files: List[Dict[str, Any]] = [\n {\n \"id\": f\"{uuid.uuid4()}.txt\",\n \"content\": b\"chat upload\",\n \"origin\": FileOrigin.CHAT_UPLOAD,\n \"type\": \"text/plain\",\n },\n {\n \"id\": f\"{uuid.uuid4()}.json\",\n \"content\": b'{\"from\":\"connector\"}',\n \"origin\": FileOrigin.CONNECTOR,\n \"type\": \"application/json\",\n },\n {\n \"id\": f\"{uuid.uuid4()}.csv\",\n \"content\": b\"a,b\\n1,2\",\n \"origin\": FileOrigin.GENERATED_REPORT,\n \"type\": \"text/csv\",\n },\n ]\n\n for f in files:\n pg_file_store.save_file(\n content=BytesIO(f[\"content\"]),\n display_name=f[\"id\"],\n file_origin=f[\"origin\"],\n file_type=f[\"type\"],\n file_id=f[\"id\"],\n )\n\n for f in files:\n assert pg_file_store.has_file(f[\"id\"], f[\"origin\"], f[\"type\"])\n assert pg_file_store.read_file(f[\"id\"]).read() == f[\"content\"]\n\n # ── complex JSONB metadata ─────────────────────────────────────\n\n def test_complex_jsonb_metadata(\n self, pg_file_store: PostgresBackedFileStore\n ) -> None:\n file_id = f\"{uuid.uuid4()}.json\"\n metadata: Dict[str, Any] = {\n \"nested\": {\"array\": [1, 2, {\"inner\": True}], \"null_val\": None},\n \"unicode\": \"测试 🚀\",\n \"large_text\": \"z\" * 1000,\n }\n\n pg_file_store.save_file(\n content=BytesIO(b\"{}\"),\n display_name=\"meta.json\",\n file_origin=FileOrigin.OTHER,\n file_type=\"application/json\",\n file_metadata=metadata,\n file_id=file_id,\n )\n\n record = pg_file_store.read_file_record(file_id)\n assert record.file_metadata == metadata\n" }, { "path": "backend/tests/external_dependency_unit/full_setup.py", "content": "from __future__ import annotations\n\nimport os\nfrom pathlib import Path\nfrom typing import Optional\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.search_settings import get_active_search_settings\nfrom onyx.document_index.factory import get_all_document_indices\nfrom onyx.document_index.factory import get_default_document_index\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.indexing.models import IndexingSetting\nfrom onyx.setup import setup_document_indices\nfrom onyx.setup import setup_postgres\nfrom shared_configs import configs as shared_configs_module\nfrom shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\nfrom tests.external_dependency_unit.constants import TEST_TENANT_ID\n\n\n_SETUP_COMPLETE: bool = False\n\n\ndef ensure_full_deployment_setup(\n tenant_id: Optional[str] = None,\n opensearch_available: bool = False,\n) -> None:\n \"\"\"Initialize test environment to mirror a real deployment, on demand.\n\n - Initializes DB engine and sets tenant context\n - Skips model warm-ups during setup\n - Runs setup_onyx (Postgres defaults, Vespa indices)\n - Initializes file store (best-effort)\n - Ensures Vespa indices exist\n \"\"\"\n global _SETUP_COMPLETE\n if _SETUP_COMPLETE:\n return\n\n if os.environ.get(\"SKIP_EXTERNAL_DEPENDENCY_UNIT_SETUP\", \"\").lower() == \"true\":\n return\n\n tenant = tenant_id or TEST_TENANT_ID\n\n # Initialize engine (noop if already initialized)\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n # Avoid warm-up network calls during setup\n shared_configs_module.SKIP_WARM_UP = True\n\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant)\n original_cwd = os.getcwd()\n backend_dir = Path(__file__).resolve().parents[2] # points to 'backend'\n os.chdir(str(backend_dir))\n\n try:\n with get_session_with_current_tenant() as db_session:\n setup_postgres(db_session)\n\n # Initialize file store; ignore if not configured\n try:\n get_default_file_store().initialize()\n except Exception:\n pass\n\n # Also ensure indices exist explicitly (no-op if already created)\n with get_session_with_current_tenant() as db_session:\n active = get_active_search_settings(db_session)\n if opensearch_available:\n # We use this special bool here instead of just relying on\n # ENABLE_OPENSEARCH_INDEXING_FOR_ONYX because not all testing\n # infra is configured for OpenSearch.\n document_indices = get_all_document_indices(\n active.primary, active.secondary\n )\n else:\n document_indices = [\n get_default_document_index(\n active.primary, active.secondary, db_session\n )\n ]\n ok = setup_document_indices(\n document_indices=document_indices,\n index_setting=IndexingSetting.from_db_model(active.primary),\n secondary_index_setting=(\n IndexingSetting.from_db_model(active.secondary)\n if active.secondary\n else None\n ),\n )\n if not ok:\n raise RuntimeError(\n \"Vespa did not initialize within the specified timeout.\"\n )\n\n _SETUP_COMPLETE = True\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n os.chdir(original_cwd)\n" }, { "path": "backend/tests/external_dependency_unit/hierarchy/__init__.py", "content": "" }, { "path": "backend/tests/external_dependency_unit/hierarchy/test_hierarchy_access_filter.py", "content": "\"\"\"Tests for hierarchy node access filtering.\n\nValidates that the overlap operator on external_user_group_ids works correctly\nwith PostgreSQL's VARCHAR[] column type. This specifically tests the fix for\nthe `character varying[] && text[]` type mismatch error.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.db.hierarchy import _get_accessible_hierarchy_nodes_for_source\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.enums import HierarchyNodeType\nfrom onyx.db.models import HierarchyNode\n\n\ndef _make_node(\n raw_node_id: str,\n display_name: str,\n *,\n is_public: bool = False,\n external_user_emails: list[str] | None = None,\n external_user_group_ids: list[str] | None = None,\n) -> HierarchyNode:\n return HierarchyNode(\n raw_node_id=raw_node_id,\n display_name=display_name,\n source=DocumentSource.GOOGLE_DRIVE,\n node_type=HierarchyNodeType.FOLDER,\n is_public=is_public,\n external_user_emails=external_user_emails,\n external_user_group_ids=external_user_group_ids,\n )\n\n\n@pytest.fixture()\ndef seeded_nodes(db_session: Session) -> Generator[list[HierarchyNode], None, None]:\n \"\"\"Seed hierarchy nodes with various permission configurations.\"\"\"\n tag = uuid4().hex[:8]\n nodes = [\n _make_node(\n f\"public_{tag}\",\n f\"Public Folder {tag}\",\n is_public=True,\n ),\n _make_node(\n f\"email_only_{tag}\",\n f\"Email-Only Folder {tag}\",\n external_user_emails=[\"alice@example.com\"],\n ),\n _make_node(\n f\"group_only_{tag}\",\n f\"Group-Only Folder {tag}\",\n external_user_group_ids=[\"group_engineering\", \"group_design\"],\n ),\n _make_node(\n f\"private_{tag}\",\n f\"Private Folder {tag}\",\n ),\n ]\n for node in nodes:\n db_session.add(node)\n db_session.flush()\n\n yield nodes\n\n # Cleanup\n for node in nodes:\n db_session.delete(node)\n db_session.commit()\n\n\ndef test_group_overlap_filter(\n db_session: Session,\n seeded_nodes: list[HierarchyNode],\n) -> None:\n \"\"\"The overlap (&&) operator must work on the VARCHAR[] column.\n\n This is the core regression test: before the cast fix, PostgreSQL raised\n `operator does not exist: character varying[] && text[]`.\n \"\"\"\n results = _get_accessible_hierarchy_nodes_for_source(\n db_session,\n source=DocumentSource.GOOGLE_DRIVE,\n user_email=\"\",\n external_group_ids=[\"group_engineering\"],\n )\n result_ids = {n.raw_node_id for n in results}\n\n public_node, _, group_node, private_node = seeded_nodes\n assert public_node.raw_node_id in result_ids\n assert group_node.raw_node_id in result_ids\n assert private_node.raw_node_id not in result_ids\n\n\ndef test_email_filter(\n db_session: Session,\n seeded_nodes: list[HierarchyNode],\n) -> None:\n \"\"\"User email matching should return the email-permissioned node.\"\"\"\n results = _get_accessible_hierarchy_nodes_for_source(\n db_session,\n source=DocumentSource.GOOGLE_DRIVE,\n user_email=\"alice@example.com\",\n external_group_ids=[],\n )\n result_ids = {n.raw_node_id for n in results}\n\n public_node, email_node, group_node, private_node = seeded_nodes\n assert public_node.raw_node_id in result_ids\n assert email_node.raw_node_id in result_ids\n assert group_node.raw_node_id not in result_ids\n assert private_node.raw_node_id not in result_ids\n\n\ndef test_no_credentials_returns_only_public(\n db_session: Session,\n seeded_nodes: list[HierarchyNode],\n) -> None:\n \"\"\"With no email and no groups, only public nodes should be returned.\"\"\"\n results = _get_accessible_hierarchy_nodes_for_source(\n db_session,\n source=DocumentSource.GOOGLE_DRIVE,\n user_email=\"\",\n external_group_ids=[],\n )\n result_ids = {n.raw_node_id for n in results}\n\n public_node, email_node, group_node, private_node = seeded_nodes\n assert public_node.raw_node_id in result_ids\n assert email_node.raw_node_id not in result_ids\n assert group_node.raw_node_id not in result_ids\n assert private_node.raw_node_id not in result_ids\n\n\ndef test_combined_email_and_group(\n db_session: Session,\n seeded_nodes: list[HierarchyNode],\n) -> None:\n \"\"\"Both email and group filters should apply together via OR.\"\"\"\n results = _get_accessible_hierarchy_nodes_for_source(\n db_session,\n source=DocumentSource.GOOGLE_DRIVE,\n user_email=\"alice@example.com\",\n external_group_ids=[\"group_design\"],\n )\n result_ids = {n.raw_node_id for n in results}\n\n public_node, email_node, group_node, private_node = seeded_nodes\n assert public_node.raw_node_id in result_ids\n assert email_node.raw_node_id in result_ids\n assert group_node.raw_node_id in result_ids\n assert private_node.raw_node_id not in result_ids\n" }, { "path": "backend/tests/external_dependency_unit/llm/test_llm_provider.py", "content": "\"\"\"\nTests for the test_llm_configuration endpoint (/admin/llm/test).\n\nThis tests the LLM configuration testing functionality which verifies\nthat LLM credentials are valid before saving them.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import LLMModelFlowType\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import remove_llm_provider\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import UserRole\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LLM\nfrom onyx.server.manage.llm.api import (\n test_default_provider as run_test_default_provider,\n)\nfrom onyx.server.manage.llm.api import (\n test_llm_configuration as run_test_llm_configuration,\n)\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom onyx.server.manage.llm.models import TestLLMRequest as LLMTestRequest\n\n\ndef _create_mock_admin() -> MagicMock:\n \"\"\"Create a mock admin user for testing.\"\"\"\n mock_admin = MagicMock()\n mock_admin.role = UserRole.ADMIN\n return mock_admin\n\n\ndef _create_test_provider(\n db_session: Session,\n name: str,\n api_key: str = \"sk-test-key-00000000000000000000000000000000000\",\n) -> LLMProviderView:\n \"\"\"Helper to create a test LLM provider in the database.\"\"\"\n return upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=name,\n provider=LlmProviderNames.OPENAI,\n api_key=api_key,\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(name=\"gpt-4o-mini\", is_visible=True)\n ],\n ),\n db_session=db_session,\n )\n\n\ndef _cleanup_provider(db_session: Session, name: str) -> None:\n \"\"\"Helper to clean up a test provider by name.\"\"\"\n provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n if provider:\n remove_llm_provider(db_session, provider.id)\n\n\n@pytest.fixture\ndef provider_name() -> Generator[str, None, None]:\n \"\"\"Generate a unique provider name for each test.\"\"\"\n yield f\"test-provider-{uuid4().hex[:8]}\"\n\n\nclass TestLLMConfigurationEndpoint:\n \"\"\"Tests for the test_llm_configuration endpoint.\"\"\"\n\n def test_successful_llm_test_with_new_provider(\n self,\n db_session: Session,\n provider_name: str, # noqa: ARG002\n ) -> None:\n \"\"\"\n Test that a successful LLM test returns normally (no exception).\n\n When test_llm returns None (success), the endpoint should complete\n without raising an exception.\n \"\"\"\n captured_llms: list[LLM] = []\n\n def mock_test_llm_success(llm: LLM) -> str | None:\n \"\"\"Mock test_llm that always succeeds.\"\"\"\n captured_llms.append(llm)\n return None # Success\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_success\n ):\n # This should complete without exception\n run_test_llm_configuration(\n test_llm_request=LLMTestRequest(\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-new-test-key-0000000000000000000000000000\",\n api_key_changed=True,\n custom_config_changed=False,\n model=\"gpt-4o-mini\",\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify test_llm was called\n assert len(captured_llms) == 1, \"test_llm should have been called once\"\n\n # Verify the LLM was configured with the correct model\n assert captured_llms[0].config.model_name == \"gpt-4o-mini\"\n assert captured_llms[0].config.model_provider == LlmProviderNames.OPENAI\n\n finally:\n db_session.rollback()\n\n def test_failed_llm_test_raises_onyx_error(\n self,\n db_session: Session,\n provider_name: str, # noqa: ARG002\n ) -> None:\n \"\"\"\n Test that a failed LLM test raises an OnyxError with VALIDATION_ERROR.\n\n When test_llm returns an error message, the endpoint should raise\n an OnyxError with the error details.\n \"\"\"\n error_message = \"Invalid API key: Authentication failed\"\n\n def mock_test_llm_failure(llm: LLM) -> str | None: # noqa: ARG001\n \"\"\"Mock test_llm that always fails.\"\"\"\n return error_message\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_failure\n ):\n with pytest.raises(OnyxError) as exc_info:\n run_test_llm_configuration(\n test_llm_request=LLMTestRequest(\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-invalid-key-00000000000000000000000000\",\n api_key_changed=True,\n custom_config_changed=False,\n model=\"gpt-4o-mini\",\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert exc_info.value.detail == error_message\n\n finally:\n db_session.rollback()\n\n def test_uses_existing_provider_api_key_when_not_changed(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that when testing an existing provider without changing the API key,\n the stored API key from the database is used.\n \"\"\"\n original_api_key = \"sk-original-stored-key-00000000000000000000\"\n captured_llms: list[LLM] = []\n\n def mock_test_llm_capture(llm: LLM) -> str | None:\n \"\"\"Mock test_llm that captures the LLM for inspection.\"\"\"\n captured_llms.append(llm)\n return None\n\n try:\n # First, create the provider in the database\n provider = _create_test_provider(\n db_session, provider_name, api_key=original_api_key\n )\n\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n # Test with api_key_changed=False - should use stored key\n run_test_llm_configuration(\n test_llm_request=LLMTestRequest(\n id=provider.id,\n provider=LlmProviderNames.OPENAI,\n api_key=None, # Not providing a new key\n api_key_changed=False, # Using existing key\n custom_config_changed=False,\n model=\"gpt-4o-mini\",\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify test_llm was called with the original API key\n assert len(captured_llms) == 1\n assert captured_llms[0].config.api_key == original_api_key\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_uses_new_api_key_when_changed(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that when testing an existing provider with a new API key,\n the new API key is used instead of the stored one.\n \"\"\"\n original_api_key = \"sk-original-stored-key-00000000000000000000\"\n new_api_key = \"sk-new-updated-key-000000000000000000000000\"\n captured_llms: list[LLM] = []\n\n def mock_test_llm_capture(llm: LLM) -> str | None:\n \"\"\"Mock test_llm that captures the LLM for inspection.\"\"\"\n captured_llms.append(llm)\n return None\n\n try:\n # First, create the provider in the database\n provider = _create_test_provider(\n db_session, provider_name, api_key=original_api_key\n )\n\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n # Test with api_key_changed=True - should use new key\n run_test_llm_configuration(\n test_llm_request=LLMTestRequest(\n id=provider.id,\n provider=LlmProviderNames.OPENAI,\n api_key=new_api_key, # Providing a new key\n api_key_changed=True, # Key is being changed\n custom_config_changed=False,\n model=\"gpt-4o-mini\",\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify test_llm was called with the new API key\n assert len(captured_llms) == 1\n assert captured_llms[0].config.api_key == new_api_key\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_uses_existing_custom_config_when_not_changed(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that when testing an existing provider without changing custom_config,\n the stored custom_config from the database is used.\n \"\"\"\n original_custom_config = {\"custom_key\": \"original_value\"}\n captured_llms: list[LLM] = []\n\n def mock_test_llm_capture(llm: LLM) -> str | None:\n \"\"\"Mock test_llm that captures the LLM for inspection.\"\"\"\n captured_llms.append(llm)\n return None\n\n try:\n # First, create the provider in the database with custom_config\n provider = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n custom_config=original_custom_config,\n custom_config_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n )\n ],\n ),\n db_session=db_session,\n )\n\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n # Test with custom_config_changed=False - should use stored config\n run_test_llm_configuration(\n test_llm_request=LLMTestRequest(\n id=provider.id,\n provider=LlmProviderNames.OPENAI,\n api_key=None,\n api_key_changed=False,\n custom_config=None, # Not providing new config\n custom_config_changed=False, # Using existing config\n model=\"gpt-4o-mini\",\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify test_llm was called with the original custom_config\n assert len(captured_llms) == 1\n assert captured_llms[0].config.custom_config == original_custom_config\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_different_model_names(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test that the endpoint correctly passes different model names to the LLM.\n \"\"\"\n captured_llms: list[LLM] = []\n\n def mock_test_llm_capture(llm: LLM) -> str | None:\n captured_llms.append(llm)\n return None\n\n test_models = [\"gpt-4\", \"gpt-4o\", \"gpt-4o-mini\", \"gpt-3.5-turbo\"]\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n for model_name in test_models:\n run_test_llm_configuration(\n test_llm_request=LLMTestRequest(\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n custom_config_changed=False,\n model=model_name,\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify all models were tested\n assert len(captured_llms) == len(test_models)\n\n for i, llm in enumerate(captured_llms):\n assert (\n llm.config.model_name == test_models[i]\n ), f\"Expected model {test_models[i]}, got {llm.config.model_name}\"\n\n finally:\n db_session.rollback()\n\n\nclass TestDefaultProviderEndpoint:\n \"\"\"Tests for the test_default_provider endpoint (/admin/llm/test/default).\"\"\"\n\n def test_default_provider_switching(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test that run_test_default_provider correctly uses the default provider\n and responds to changes in default model and default provider.\n\n Steps:\n 1. Upload provider 1 with models, set as default\n 2. Call run_test_default_provider - should use provider 1's default model\n 3. Upload provider 2 with models (not default)\n 4. Call run_test_default_provider - should still use provider 1\n 5. Change the default model on provider 1\n 6. Call run_test_default_provider - should use new model on provider 1\n 7. Change the default provider to provider 2\n 8. Call run_test_default_provider - should use provider 2\n \"\"\"\n provider_1_name = f\"test-provider-1-{uuid4().hex[:8]}\"\n provider_2_name = f\"test-provider-2-{uuid4().hex[:8]}\"\n\n provider_1_api_key = \"sk-provider1-key-000000000000000000000000000\"\n provider_2_api_key = \"sk-provider2-key-000000000000000000000000000\"\n\n provider_1_initial_model = \"gpt-4\"\n provider_1_updated_model = \"gpt-4o\"\n provider_2_default_model = \"gpt-4o-mini\"\n\n captured_llms: list[LLM] = []\n\n def mock_test_llm_capture(llm: LLM) -> str | None:\n \"\"\"Mock test_llm that captures the LLM for inspection.\"\"\"\n captured_llms.append(llm)\n return None\n\n try:\n # Step 1: Create provider 1 with models, it becomes default (first provider)\n provider_1 = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=provider_1_name,\n provider=LlmProviderNames.OPENAI,\n api_key=provider_1_api_key,\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ],\n ),\n db_session=db_session,\n )\n\n # Set provider 1 as the default provider explicitly\n update_default_provider(provider_1.id, provider_1_initial_model, db_session)\n\n # Step 2: Call run_test_default_provider - should use provider 1's default model\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n run_test_default_provider(_=_create_mock_admin())\n\n assert len(captured_llms) == 1\n assert captured_llms[0].config.model_name == provider_1_initial_model\n assert captured_llms[0].config.api_key == provider_1_api_key\n captured_llms.clear()\n\n # Step 3: Create provider 2 (not default)\n provider_2 = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=provider_2_name,\n provider=LlmProviderNames.OPENAI,\n api_key=provider_2_api_key,\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-3.5-turbo\", is_visible=True\n ),\n ],\n ),\n db_session=db_session,\n )\n\n # Step 4: Call run_test_default_provider - should still use provider 1\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n run_test_default_provider(_=_create_mock_admin())\n\n assert len(captured_llms) == 1\n assert captured_llms[0].config.model_name == provider_1_initial_model\n assert captured_llms[0].config.api_key == provider_1_api_key\n captured_llms.clear()\n\n # Step 5: Update provider 1's default model\n upsert_llm_provider(\n LLMProviderUpsertRequest(\n id=provider_1.id,\n name=provider_1_name,\n provider=LlmProviderNames.OPENAI,\n api_key=provider_1_api_key,\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ],\n ),\n db_session=db_session,\n )\n\n # Set provider 1's default model to the updated model\n update_default_provider(provider_1.id, provider_1_updated_model, db_session)\n\n # Step 6: Call run_test_default_provider - should use new model on provider 1\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n run_test_default_provider(_=_create_mock_admin())\n\n assert len(captured_llms) == 1\n assert captured_llms[0].config.model_name == provider_1_updated_model\n assert captured_llms[0].config.api_key == provider_1_api_key\n captured_llms.clear()\n\n # Step 7: Change the default provider to provider 2\n update_default_provider(provider_2.id, provider_2_default_model, db_session)\n\n # Step 8: Call run_test_default_provider - should use provider 2\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n run_test_default_provider(_=_create_mock_admin())\n\n assert len(captured_llms) == 1\n assert captured_llms[0].config.model_name == provider_2_default_model\n assert captured_llms[0].config.api_key == provider_2_api_key\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_1_name)\n _cleanup_provider(db_session, provider_2_name)\n\n def test_no_default_provider_raises_exception(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test that when no default provider exists, the endpoint raises an exception.\n \"\"\"\n # Clear any existing providers to ensure no default exists\n from onyx.db.llm import fetch_existing_llm_providers\n\n try:\n existing_providers = fetch_existing_llm_providers(\n db_session, flow_type_filter=[LLMModelFlowType.CHAT]\n )\n provider_names_to_restore: list[str] = []\n\n for provider in existing_providers:\n provider_names_to_restore.append(provider.name)\n\n # Remove all providers temporarily\n for provider in existing_providers:\n remove_llm_provider(db_session, provider.id)\n\n # Now run_test_default_provider should fail\n with pytest.raises(OnyxError) as exc_info:\n run_test_default_provider(_=_create_mock_admin())\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert \"No LLM Provider setup\" in exc_info.value.detail\n\n finally:\n db_session.rollback()\n\n def test_default_provider_test_failure(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test that when the default provider's LLM test fails, an exception is raised.\n \"\"\"\n provider_name = f\"test-provider-{uuid4().hex[:8]}\"\n error_message = \"Connection to LLM provider failed\"\n\n def mock_test_llm_failure(llm: LLM) -> str | None: # noqa: ARG001\n \"\"\"Mock test_llm that always fails.\"\"\"\n return error_message\n\n try:\n # Create a provider and set it as default\n provider = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ],\n ),\n db_session=db_session,\n )\n update_default_provider(provider.id, \"gpt-4o-mini\", db_session)\n\n # Test should fail\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_failure\n ):\n with pytest.raises(OnyxError) as exc_info:\n run_test_default_provider(_=_create_mock_admin())\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert exc_info.value.detail == error_message\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n" }, { "path": "backend/tests/external_dependency_unit/llm/test_llm_provider_api_base.py", "content": "\"\"\"\nTests for LLM provider api_base and custom_config change restrictions.\n\nThis ensures we don't have a vulnerability where an admin could change the api_base\nor custom_config of an LLM provider without changing the API key, allowing them to\nredirect API requests (containing the real API key in headers) to an attacker-controlled\nserver.\n\nThese are external dependency unit tests because they need a real database but\nalso need to control the MULTI_TENANT setting via patching.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import remove_llm_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import UserRole\nfrom onyx.error_handling.error_codes import OnyxErrorCode\nfrom onyx.error_handling.exceptions import OnyxError\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.server.manage.llm.api import _mask_string\nfrom onyx.server.manage.llm.api import put_llm_provider\nfrom onyx.server.manage.llm.api import test_llm_configuration as run_llm_config_test\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom onyx.server.manage.llm.models import TestLLMRequest as LLMTestRequest\nfrom tests.external_dependency_unit.mock_llm import LLM\n\n\ndef _create_test_provider(\n db_session: Session,\n name: str,\n api_base: str | None = None,\n custom_config: dict[str, str] | None = None,\n) -> LLMProviderView:\n \"\"\"Helper to create a test LLM provider.\"\"\"\n return upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n api_base=api_base,\n custom_config=custom_config,\n model_configurations=[\n ModelConfigurationUpsertRequest(name=\"gpt-4o-mini\", is_visible=True)\n ],\n ),\n db_session=db_session,\n )\n\n\ndef _cleanup_provider(db_session: Session, name: str) -> None:\n \"\"\"Helper to clean up a test provider by name.\"\"\"\n provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n if provider:\n remove_llm_provider(db_session, provider.id)\n\n\ndef _create_mock_admin() -> MagicMock:\n \"\"\"Create a mock admin user for testing.\"\"\"\n mock_admin = MagicMock()\n mock_admin.role = UserRole.ADMIN\n return mock_admin\n\n\n@pytest.fixture\ndef provider_name() -> Generator[str, None, None]:\n \"\"\"Generate a unique provider name for each test.\"\"\"\n yield f\"test-provider-{uuid4().hex[:8]}\"\n\n\nclass TestLLMProviderChanges:\n \"\"\"Tests for api_base change restrictions when updating LLM providers.\"\"\"\n\n def test_blocks_api_base_change_without_key_change__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n In multi-tenant mode, changing api_base without also changing\n the API key should be blocked.\n \"\"\"\n try:\n provider = _create_test_provider(db_session, provider_name)\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_base=\"https://attacker.example.com\",\n )\n\n with pytest.raises(OnyxError) as exc_info:\n put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert \"cannot be changed without changing the API key\" in str(\n exc_info.value.detail\n )\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_api_base_change_with_key_change__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Changing api_base IS allowed when the API key is also being changed.\n \"\"\"\n try:\n provider = _create_test_provider(db_session, provider_name)\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-new-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n api_base=\"https://custom-endpoint.example.com/v1\",\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.api_base == \"https://custom-endpoint.example.com/v1\"\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_same_api_base__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Keeping the same api_base (no change) is allowed without changing the API key.\n \"\"\"\n original_api_base = \"https://original.example.com/v1\"\n\n try:\n provider = _create_test_provider(\n db_session, provider_name, api_base=original_api_base\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_base=original_api_base,\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.api_base == original_api_base\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_empty_string_api_base_when_existing_is_none__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Treat empty-string api_base from clients as unset when comparing provider\n changes. This allows model-only updates when provider has no custom base URL.\n \"\"\"\n try:\n view = _create_test_provider(db_session, provider_name, api_base=None)\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=view.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_base=\"\",\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.api_base is None\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_blocks_clearing_api_base__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Clearing api_base (setting to None when it was previously set)\n is also blocked without changing the API key.\n \"\"\"\n original_api_base = \"https://original.example.com/v1\"\n\n try:\n provider = _create_test_provider(\n db_session, provider_name, api_base=original_api_base\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_base=None,\n )\n\n with pytest.raises(OnyxError) as exc_info:\n put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert \"cannot be changed without changing the API key\" in str(\n exc_info.value.detail\n )\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_api_base_change__single_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n In single-tenant mode (MULTI_TENANT=False), changing api_base without\n changing the API key IS allowed. This is by design since single-tenant\n users have full control over their deployment.\n \"\"\"\n try:\n provider = _create_test_provider(db_session, provider_name)\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", False):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_base=\"https://custom.example.com/v1\",\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.api_base == \"https://custom.example.com/v1\"\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_new_provider_creation_not_affected__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Creating a new provider with an api_base should work regardless of\n api_key_changed (since there's no existing key to protect).\n \"\"\"\n try:\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n create_request = LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-new-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n api_base=\"https://custom.example.com/v1\",\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=create_request,\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.api_base == \"https://custom.example.com/v1\"\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_blocks_custom_config_change_without_key_change__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n In multi-tenant mode, changing custom_config without also changing\n the API key should be blocked (custom_config can set env vars that\n redirect LLM API requests).\n \"\"\"\n try:\n provider = _create_test_provider(\n db_session,\n provider_name,\n custom_config={\"SOME_CONFIG\": \"original_value\"},\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n custom_config={\"OPENAI_API_BASE\": \"https://attacker.example.com\"},\n custom_config_changed=True,\n )\n\n with pytest.raises(OnyxError) as exc_info:\n put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert \"cannot be changed without changing the API key\" in str(\n exc_info.value.detail\n )\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_blocks_adding_custom_config_without_key_change__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Adding custom_config when none existed should also be blocked\n without changing the API key.\n \"\"\"\n try:\n provider = _create_test_provider(db_session, provider_name)\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n custom_config={\"OPENAI_API_BASE\": \"https://attacker.example.com\"},\n custom_config_changed=True,\n )\n\n with pytest.raises(OnyxError) as exc_info:\n put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert exc_info.value.error_code == OnyxErrorCode.VALIDATION_ERROR\n assert \"cannot be changed without changing the API key\" in str(\n exc_info.value.detail\n )\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_custom_config_change_with_key_change__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Changing custom_config IS allowed when the API key is also being changed.\n \"\"\"\n new_config = {\"AWS_REGION_NAME\": \"us-west-2\"}\n\n try:\n provider = _create_test_provider(\n db_session,\n provider_name,\n custom_config={\"AWS_REGION_NAME\": \"us-east-1\"},\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-new-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n custom_config_changed=True,\n custom_config=new_config,\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.custom_config == new_config\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_same_custom_config__multi_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Keeping the same custom_config (no change) is allowed without changing the API key.\n \"\"\"\n original_config = {\"AWS_REGION_NAME\": \"us-east-1\"}\n\n try:\n provider = _create_test_provider(\n db_session, provider_name, custom_config=original_config\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", True):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n custom_config=original_config,\n custom_config_changed=True,\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.custom_config == original_config\n finally:\n _cleanup_provider(db_session, provider_name)\n\n def test_allows_custom_config_change__single_tenant(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n In single-tenant mode, changing custom_config without changing\n the API key IS allowed.\n \"\"\"\n new_config = {\"AWS_REGION_NAME\": \"eu-west-1\"}\n\n try:\n provider = _create_test_provider(\n db_session,\n provider_name,\n custom_config={\"AWS_REGION_NAME\": \"us-east-1\"},\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", False):\n update_request = LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n custom_config=new_config,\n custom_config_changed=True,\n )\n\n result = put_llm_provider(\n llm_provider_upsert_request=update_request,\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert result.custom_config == new_config\n finally:\n _cleanup_provider(db_session, provider_name)\n\n\ndef test_upload_with_custom_config_then_change(\n db_session: Session,\n) -> None:\n \"\"\"\n Run test + upload with a custom config (vertex).\n Edit attributes of provider that are not custom config or api key.\n Check that the test and update maintain the same values.\n \"\"\"\n custom_config = {\n \"vertex_credentials\": \"1234\",\n \"vertex_location\": \"us-east-1\",\n }\n name = \"test-provider-vertex-ai\"\n provider_name = LlmProviderNames.VERTEX_AI.value\n default_model_name = \"gemini-2.5-pro\"\n\n # List to capture LLM inputs passed to test_llm\n captured_llms: list = []\n\n def capture_test_llm(llm: LLM) -> str:\n \"\"\"Captures the LLM input and returns None (success).\"\"\"\n captured_llms.append(llm)\n return \"\"\n\n try:\n # Patch the test_llm method\n with patch(\"onyx.server.manage.llm.api.test_llm\", side_effect=capture_test_llm):\n run_llm_config_test(\n LLMTestRequest(\n provider=provider_name,\n model=default_model_name,\n api_key_changed=False,\n custom_config_changed=True,\n custom_config=custom_config,\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n provider = put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=name,\n provider=provider_name,\n custom_config=custom_config,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name, is_visible=True\n )\n ],\n api_key_changed=False,\n custom_config_changed=True,\n is_auto_mode=False,\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Turn auto mode off\n run_llm_config_test(\n LLMTestRequest(\n id=provider.id,\n provider=provider_name,\n model=default_model_name,\n api_key_changed=False,\n custom_config_changed=False,\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n id=provider.id,\n name=name,\n provider=provider_name,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name, is_visible=True\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ],\n api_key_changed=False,\n custom_config_changed=False,\n is_auto_mode=False,\n ),\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify that test_llm was called and custom_config matches\n assert len(captured_llms) == 2, \"test_llm should have been called 2 times\"\n\n for llm in captured_llms:\n assert (\n llm.config.custom_config == custom_config\n ), f\"Expected custom_config {custom_config}, but got {llm.config.custom_config}\"\n\n # Check inside the database and check that custom_config is the same as the original\n db_provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n if not db_provider:\n assert False, \"Provider not found in the database\"\n\n assert (\n db_provider.custom_config == custom_config\n ), f\"Expected custom_config {custom_config}, but got {db_provider.custom_config}\"\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, name)\n\n\ndef test_preserves_masked_sensitive_custom_config_on_provider_update(\n db_session: Session,\n) -> None:\n \"\"\"Masked sensitive values from the UI should not overwrite stored secrets.\"\"\"\n name = f\"test-provider-vertex-update-{uuid4().hex[:8]}\"\n provider = LlmProviderNames.VERTEX_AI.value\n default_model_name = \"gemini-2.5-pro\"\n original_custom_config = {\n \"vertex_credentials\": '{\"type\":\"service_account\",\"private_key\":\"REAL_PRIVATE_KEY\"}',\n \"vertex_location\": \"global\",\n }\n\n try:\n view = put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=name,\n provider=provider,\n custom_config=original_custom_config,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name, is_visible=True\n )\n ],\n api_key_changed=False,\n custom_config_changed=True,\n is_auto_mode=False,\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n with patch(\"onyx.server.manage.llm.api.MULTI_TENANT\", False):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n id=view.id,\n name=name,\n provider=provider,\n custom_config={\n \"vertex_credentials\": _mask_string(\n original_custom_config[\"vertex_credentials\"]\n ),\n \"vertex_location\": \"us-central1\",\n },\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name, is_visible=True\n )\n ],\n api_key_changed=False,\n custom_config_changed=True,\n is_auto_mode=False,\n ),\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n updated_provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n assert updated_provider is not None\n assert updated_provider.custom_config is not None\n assert (\n updated_provider.custom_config[\"vertex_credentials\"]\n == original_custom_config[\"vertex_credentials\"]\n )\n assert updated_provider.custom_config[\"vertex_location\"] == \"us-central1\"\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, name)\n\n\ndef test_preserves_masked_sensitive_custom_config_on_test_request(\n db_session: Session,\n) -> None:\n \"\"\"LLM test should restore masked sensitive custom config values before invocation.\"\"\"\n name = f\"test-provider-vertex-test-{uuid4().hex[:8]}\"\n provider_name = LlmProviderNames.VERTEX_AI.value\n default_model_name = \"gemini-2.5-pro\"\n original_custom_config = {\n \"vertex_credentials\": '{\"type\":\"service_account\",\"private_key\":\"REAL_PRIVATE_KEY\"}',\n \"vertex_location\": \"global\",\n }\n captured_llms: list[LLM] = []\n\n def capture_test_llm(llm: LLM) -> str:\n captured_llms.append(llm)\n return \"\"\n\n try:\n provider = put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=name,\n provider=provider_name,\n custom_config=original_custom_config,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name, is_visible=True\n )\n ],\n api_key_changed=False,\n custom_config_changed=True,\n is_auto_mode=False,\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n with patch(\"onyx.server.manage.llm.api.test_llm\", side_effect=capture_test_llm):\n run_llm_config_test(\n LLMTestRequest(\n id=provider.id,\n provider=provider_name,\n model=default_model_name,\n api_key_changed=False,\n custom_config_changed=True,\n custom_config={\n \"vertex_credentials\": _mask_string(\n original_custom_config[\"vertex_credentials\"]\n ),\n \"vertex_location\": \"us-central1\",\n },\n ),\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n assert len(captured_llms) == 1\n assert captured_llms[0].config.custom_config is not None\n assert (\n captured_llms[0].config.custom_config[\"vertex_credentials\"]\n == original_custom_config[\"vertex_credentials\"]\n )\n assert captured_llms[0].config.custom_config[\"vertex_location\"] == \"us-central1\"\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, name)\n" }, { "path": "backend/tests/external_dependency_unit/llm/test_llm_provider_auto_mode.py", "content": "\"\"\"\nTests for the LLM Provider Auto Mode feature.\n\nThis tests the automatic model syncing from GitHub config when a provider\nis uploaded with is_auto_mode=True.\n\"\"\"\n\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import LLMModelFlowType\nfrom onyx.db.llm import fetch_auto_mode_providers\nfrom onyx.db.llm import fetch_default_llm_model\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import fetch_existing_llm_providers\nfrom onyx.db.llm import fetch_llm_provider_view\nfrom onyx.db.llm import remove_llm_provider\nfrom onyx.db.llm import sync_auto_mode_models\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.models import UserRole\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.well_known_providers.auto_update_models import LLMProviderRecommendation\nfrom onyx.llm.well_known_providers.auto_update_models import LLMRecommendations\nfrom onyx.llm.well_known_providers.models import SimpleKnownModel\nfrom onyx.server.manage.llm.api import put_llm_provider\nfrom onyx.server.manage.llm.api import (\n test_default_provider as run_test_default_provider,\n)\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\n\n\ndef _create_mock_admin() -> MagicMock:\n \"\"\"Create a mock admin user for testing.\"\"\"\n mock_admin = MagicMock()\n mock_admin.role = UserRole.ADMIN\n return mock_admin\n\n\ndef _cleanup_provider(db_session: Session, name: str) -> None:\n \"\"\"Helper to clean up a test provider by name.\"\"\"\n provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n if provider:\n remove_llm_provider(db_session, provider.id)\n\n\ndef _create_mock_llm_recommendations(\n provider: str,\n default_model_name: str,\n additional_models: list[str],\n) -> LLMRecommendations:\n \"\"\"Create a mock LLMRecommendations object for testing.\n\n Args:\n provider: The provider name (e.g., \"openai\")\n default_model_name: The name of the default model\n additional_models: List of additional visible model names\n\n Returns:\n LLMRecommendations object with the specified configuration\n \"\"\"\n return LLMRecommendations(\n version=\"1.0.0\",\n updated_at=datetime.now(),\n providers={\n provider: LLMProviderRecommendation(\n default_model=SimpleKnownModel(\n name=default_model_name,\n display_name=default_model_name.upper(),\n ),\n additional_visible_models=[\n SimpleKnownModel(name=model, display_name=model.upper())\n for model in additional_models\n ],\n )\n },\n )\n\n\n@pytest.fixture\ndef provider_name() -> Generator[str, None, None]:\n \"\"\"Generate a unique provider name for each test.\"\"\"\n yield f\"test-auto-provider-{uuid4().hex[:8]}\"\n\n\nclass TestAutoModeSyncFeature:\n \"\"\"Tests for the Auto Mode model syncing feature.\"\"\"\n\n def test_auto_mode_syncs_models_from_github_config(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that when a provider is uploaded with auto mode enabled and no model\n configurations, the models from fetch_llm_recommendations_from_github()\n are synced to the provider.\n\n Steps:\n 1. Mock fetch_llm_recommendations_from_github to return a known config\n 2. Upload provider with is_auto_mode=True and no model_configurations\n 3. Fetch the provider and verify all recommended models are present\n 4. Set the provider as default\n 5. Fetch the default provider and verify the default model matches the config\n \"\"\"\n # Define the expected models from the mock GitHub config\n expected_default_model = \"gpt-4o\"\n expected_additional_models = [\"gpt-4o-mini\", \"gpt-4-turbo\"]\n all_expected_models = [expected_default_model] + expected_additional_models\n\n # Create the mock LLMRecommendations\n mock_recommendations = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=expected_default_model,\n additional_models=expected_additional_models,\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=mock_recommendations,\n ):\n # Step 1-2: Upload provider with auto mode on and no model configs\n # NOTE: We need to provide a default_model_name for the initial upsert,\n # but auto mode will override it with the GitHub config's default\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[], # No model configs provided\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 3: Verify all models from the GitHub config are present\n # Fetch the provider fresh from the database\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None, \"Provider should exist\"\n assert provider.is_auto_mode is True, \"Provider should be in auto mode\"\n\n # Check that all expected models are present and visible\n model_names = {mc.name for mc in provider.model_configurations}\n for expected_model in all_expected_models:\n assert (\n expected_model in model_names\n ), f\"Expected model '{expected_model}' not found in provider models\"\n\n # Verify visibility of all synced models\n for mc in provider.model_configurations:\n if mc.name in all_expected_models:\n assert mc.is_visible is True, f\"Model '{mc.name}' should be visible\"\n\n # Step 4: Set the provider as default\n update_default_provider(provider.id, expected_default_model, db_session)\n\n # Step 5: Fetch the default provider and verify\n default_model = fetch_default_llm_model(db_session)\n assert default_model is not None, \"Default provider should exist\"\n assert (\n default_model.llm_provider.name == provider_name\n ), \"Default provider should be our test provider\"\n assert (\n default_model.name == expected_default_model\n ), f\"Default provider's default model should be '{expected_default_model}'\"\n assert (\n default_model.llm_provider.is_auto_mode is True\n ), \"Default provider should be in auto mode\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_auto_mode_with_multiple_providers_in_config(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that auto mode only syncs models for the matching provider type,\n ignoring models from other providers in the config.\n \"\"\"\n # Create recommendations with multiple providers\n mock_recommendations = LLMRecommendations(\n version=\"1.0.0\",\n updated_at=datetime.now(),\n providers={\n LlmProviderNames.OPENAI: LLMProviderRecommendation(\n default_model=SimpleKnownModel(\n name=\"gpt-4o\", display_name=\"GPT-4o\"\n ),\n additional_visible_models=[\n SimpleKnownModel(name=\"gpt-4o-mini\", display_name=\"GPT-4o Mini\")\n ],\n ),\n LlmProviderNames.ANTHROPIC: LLMProviderRecommendation(\n default_model=SimpleKnownModel(\n name=\"claude-3-5-sonnet-latest\",\n display_name=\"Claude 3.5 Sonnet\",\n ),\n additional_visible_models=[\n SimpleKnownModel(\n name=\"claude-haiku-4-5\",\n display_name=\"Claude Haiku 4.5\",\n )\n ],\n ),\n },\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=mock_recommendations,\n ):\n # Upload an OpenAI provider with auto mode\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify only OpenAI models are synced, not Anthropic models\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n model_names = {mc.name for mc in provider.model_configurations}\n\n # OpenAI models should be present\n assert \"gpt-4o\" in model_names\n assert \"gpt-4o-mini\" in model_names\n\n # Anthropic models should NOT be present\n assert \"claude-3-5-sonnet-latest\" not in model_names\n assert \"claude-haiku-4-5\" not in model_names\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_existing_provider_transition_to_auto_mode(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that when an existing provider with visible models transitions to auto mode,\n models from the auto mode config become visible, and models not in the config\n become not visible.\n\n Steps:\n 1. Upload a provider with some visible model configurations (not in auto mode)\n 2. Update the provider to enable auto mode\n 3. Verify:\n - Models in the auto mode config are now visible\n - Models NOT in the auto mode config are now NOT visible\n \"\"\"\n # Initial models on the provider (all visible initially)\n initial_models = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True\n ), # Will NOT be in auto config\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True\n ), # Will be in auto config\n ModelConfigurationUpsertRequest(\n name=\"gpt-3.5-turbo\", is_visible=True\n ), # Will NOT be in auto config\n ]\n\n # Auto mode config: gpt-4o (default) + gpt-4o-mini (additional)\n # Note: gpt-4 and gpt-3.5-turbo are NOT in this config\n auto_mode_default = \"gpt-4o\"\n auto_mode_additional = [\"gpt-4o-mini\"]\n all_auto_mode_models = [auto_mode_default] + auto_mode_additional\n\n mock_recommendations = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=auto_mode_default,\n additional_models=auto_mode_additional,\n )\n\n try:\n # Step 1: Upload provider WITHOUT auto mode, with initial models\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=False, # Not in auto mode initially\n model_configurations=initial_models,\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Verify initial state: all models are visible\n initial_provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert initial_provider is not None\n assert initial_provider.is_auto_mode is False\n\n for mc in initial_provider.model_configurations:\n assert (\n mc.is_visible is True\n ), f\"Initial model '{mc.name}' should be visible\"\n\n # Step 2: Update provider to enable auto mode\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=mock_recommendations,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n id=initial_provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=None, # Not changing API key\n api_key_changed=False,\n is_auto_mode=True, # Now enabling auto mode\n model_configurations=[], # Auto mode will sync from config\n ),\n is_creation=False, # This is an update\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 3: Verify model visibility after auto mode transition\n # Expire session cache to force fresh fetch after sync_auto_mode_models committed\n db_session.expire_all()\n provider_view = fetch_llm_provider_view(\n provider_name=provider_name, db_session=db_session\n )\n assert provider_view is not None\n assert provider_view.is_auto_mode is True\n\n # Build a map of model name -> visibility\n model_visibility = {\n mc.name: mc.is_visible for mc in provider_view.model_configurations\n }\n\n # Models in auto mode config should be visible\n for model_name in all_auto_mode_models:\n assert (\n model_name in model_visibility\n ), f\"Auto mode model '{model_name}' should exist\"\n assert (\n model_visibility[model_name] is True\n ), f\"Auto mode model '{model_name}' should be visible\"\n\n # Models NOT in auto mode config should NOT be visible\n models_not_in_config = [\"gpt-4\", \"gpt-3.5-turbo\"]\n for model_name in models_not_in_config:\n if model_name in model_visibility:\n assert (\n model_visibility[model_name] is False\n ), f\"Model '{model_name}' not in auto config should NOT be visible\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_auto_mode_provider_not_in_config(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Test that when the provider type is not in the GitHub config,\n no model syncing occurs.\n \"\"\"\n # Create recommendations that don't include OpenAI\n mock_recommendations = LLMRecommendations(\n version=\"1.0.0\",\n updated_at=datetime.now(),\n providers={\n LlmProviderNames.ANTHROPIC: LLMProviderRecommendation(\n default_model=SimpleKnownModel(\n name=\"claude-3-5-sonnet-latest\",\n display_name=\"Claude 3.5 Sonnet\",\n ),\n additional_visible_models=[],\n ),\n },\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=mock_recommendations,\n ):\n # Upload an OpenAI provider (not in config)\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n )\n ],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Provider should be created but without synced models from config\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n assert provider.is_auto_mode is True\n\n # Only the default model provided in the request should exist\n model_names = {mc.name for mc in provider.model_configurations}\n assert \"gpt-4o\" in model_names\n # Anthropic models should NOT be synced\n assert \"claude-3-5-sonnet-latest\" not in model_names\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_switching_default_between_auto_mode_providers(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"\n Test switching the default provider between two auto mode providers\n and verifying test_default_provider uses the correct default model.\n\n Steps:\n 1. Create provider 1 (OpenAI) with auto mode, set as default\n 2. Create provider 2 (Anthropic) with auto mode\n 3. Verify provider 1 is the default\n 4. Change default to provider 2\n 5. Verify provider 2 is the default\n 6. Run test_default_provider and verify it uses provider 2's default model\n \"\"\"\n provider_1_name = f\"test-auto-openai-{uuid4().hex[:8]}\"\n provider_2_name = f\"test-auto-anthropic-{uuid4().hex[:8]}\"\n\n provider_1_api_key = \"sk-provider1-key-000000000000000000000000000\"\n provider_2_api_key = \"sk-ant-provider2-key-0000000000000000000000\"\n\n # Provider 1 (OpenAI) config\n provider_1_default_model = \"gpt-4o\"\n provider_1_additional_models = [\"gpt-4o-mini\"]\n\n # Provider 2 (Anthropic) config\n provider_2_default_model = \"claude-3-5-sonnet-latest\"\n provider_2_additional_models = [\"claude-haiku-4-5\"]\n\n # Create mock recommendations with both providers\n mock_recommendations = LLMRecommendations(\n version=\"1.0.0\",\n updated_at=datetime.now(),\n providers={\n LlmProviderNames.OPENAI: LLMProviderRecommendation(\n default_model=SimpleKnownModel(\n name=provider_1_default_model,\n display_name=\"GPT-4o\",\n ),\n additional_visible_models=[\n SimpleKnownModel(name=m, display_name=m.upper())\n for m in provider_1_additional_models\n ],\n ),\n LlmProviderNames.ANTHROPIC: LLMProviderRecommendation(\n default_model=SimpleKnownModel(\n name=provider_2_default_model,\n display_name=\"Claude 3.5 Sonnet\",\n ),\n additional_visible_models=[\n SimpleKnownModel(name=m, display_name=m.upper())\n for m in provider_2_additional_models\n ],\n ),\n },\n )\n\n captured_llms: list[LLM] = []\n\n def mock_test_llm_capture(llm: LLM) -> str | None:\n \"\"\"Mock test_llm that captures the LLM for inspection.\"\"\"\n captured_llms.append(llm)\n return None # Success\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=mock_recommendations,\n ):\n # Step 1: Create provider 1 (OpenAI) with auto mode\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_1_name,\n provider=LlmProviderNames.OPENAI,\n api_key=provider_1_api_key,\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Set provider 1 as the default\n db_session.expire_all()\n provider_1 = fetch_existing_llm_provider(\n name=provider_1_name, db_session=db_session\n )\n assert provider_1 is not None\n update_default_provider(provider_1.id, provider_1_default_model, db_session)\n\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=mock_recommendations,\n ):\n # Step 2: Create provider 2 (Anthropic) with auto mode\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_2_name,\n provider=LlmProviderNames.ANTHROPIC,\n api_key=provider_2_api_key,\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 3: Verify provider 1 is still the default\n db_session.expire_all()\n default_model = fetch_default_llm_model(db_session)\n assert default_model is not None\n assert default_model.llm_provider.name == provider_1_name\n assert default_model.name == provider_1_default_model\n assert default_model.llm_provider.is_auto_mode is True\n\n # Step 4: Change the default to provider 2\n provider_2 = fetch_existing_llm_provider(\n name=provider_2_name, db_session=db_session\n )\n assert provider_2 is not None\n update_default_provider(provider_2.id, provider_2_default_model, db_session)\n\n # Step 5: Verify provider 2 is now the default\n db_session.expire_all()\n default_model = fetch_default_llm_model(db_session)\n assert default_model is not None\n assert default_model.llm_provider.name == provider_2_name\n assert default_model.name == provider_2_default_model\n assert default_model.llm_provider.is_auto_mode is True\n\n # Step 6: Run test_default_provider and verify it uses provider 2's model\n with patch(\n \"onyx.server.manage.llm.api.test_llm\", side_effect=mock_test_llm_capture\n ):\n run_test_default_provider(_=_create_mock_admin())\n\n # Verify test_llm was called with provider 2's default model\n assert len(captured_llms) == 1\n assert captured_llms[0].config.model_name == provider_2_default_model\n assert captured_llms[0].config.api_key == provider_2_api_key\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_1_name)\n _cleanup_provider(db_session, provider_2_name)\n\n\nclass TestAutoModeMissingFlows:\n \"\"\"Regression test: sync_auto_mode_models must create LLMModelFlow rows\n for every ModelConfiguration it inserts, otherwise the provider vanishes\n from listing queries that join through LLMModelFlow.\"\"\"\n\n def test_sync_auto_mode_creates_flow_rows(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"\n Steps:\n 1. Create a provider with no model configs (empty shell).\n 2. Call sync_auto_mode_models to add models from a mock config.\n 3. Assert every new ModelConfiguration has at least one LLMModelFlow.\n 4. Assert fetch_existing_llm_providers (which joins through\n LLMModelFlow) returns the provider.\n \"\"\"\n mock_recommendations = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\"],\n )\n\n try:\n # Step 1: Create provider with no model configs\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 2: Run sync_auto_mode_models (simulating the periodic sync)\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=mock_recommendations,\n )\n\n # Step 3: Every ModelConfiguration must have at least one LLMModelFlow\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n synced_model_names = {mc.name for mc in provider.model_configurations}\n assert \"gpt-4o\" in synced_model_names\n assert \"gpt-4o-mini\" in synced_model_names\n\n for mc in provider.model_configurations:\n assert len(mc.llm_model_flows) > 0, (\n f\"ModelConfiguration '{mc.name}' (id={mc.id}) has no \"\n f\"LLMModelFlow rows — it will be invisible to listing queries\"\n )\n\n flow_types = {f.llm_model_flow_type for f in mc.llm_model_flows}\n assert (\n LLMModelFlowType.CHAT in flow_types\n ), f\"ModelConfiguration '{mc.name}' is missing a CHAT flow\"\n\n # Step 4: The provider must appear in fetch_existing_llm_providers\n listed_providers = fetch_existing_llm_providers(\n db_session=db_session,\n flow_type_filter=[LLMModelFlowType.CHAT],\n )\n listed_provider_names = {p.name for p in listed_providers}\n assert (\n provider_name in listed_provider_names\n ), f\"Provider '{provider_name}' not returned by fetch_existing_llm_providers — models are missing flow rows\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n\nclass TestAutoModeTransitionsAndResync:\n \"\"\"Tests for auto/manual transitions, config evolution, and sync idempotency.\"\"\"\n\n def test_transition_to_auto_mode_preserves_default(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"When the default provider transitions from manual to auto mode,\n the global default should be preserved (set to the recommended model).\n\n Steps:\n 1. Create a manual-mode provider with models, set it as global default.\n 2. Transition to auto mode (model_configurations=[] triggers cascade\n delete of old ModelConfigurations and their LLMModelFlow rows).\n 3. Verify the provider is still the global default, now using the\n recommended default model from the GitHub config.\n \"\"\"\n initial_models = [\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4o-mini\", is_visible=True),\n ]\n\n auto_config = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o-mini\",\n additional_models=[\"gpt-4o\"],\n )\n\n try:\n # Step 1: Create manual-mode provider and set as default\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=False,\n model_configurations=initial_models,\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n default_before = fetch_default_llm_model(db_session)\n assert default_before is not None\n assert default_before.name == \"gpt-4o\"\n assert default_before.llm_provider_id == provider.id\n\n # Step 2: Transition to auto mode\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=auto_config,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=None,\n api_key_changed=False,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 3: Default should be preserved on this provider\n db_session.expire_all()\n default_after = fetch_default_llm_model(db_session)\n assert default_after is not None, (\n \"Default model should not be None after transitioning to auto mode — \"\n \"the provider was the default before and should remain so\"\n )\n assert (\n default_after.llm_provider_id == provider.id\n ), \"Default should still belong to the same provider after transition\"\n assert (\n default_after.name == \"gpt-4o-mini\"\n ), f\"Default should be updated to the recommended model 'gpt-4o-mini', got '{default_after.name}'\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_auto_to_manual_mode_preserves_models_and_stops_syncing(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Disabling auto mode should preserve the current model list and\n prevent future syncs from altering visibility.\n\n Steps:\n 1. Create provider in auto mode — models synced from config.\n 2. Update provider to manual mode (is_auto_mode=False).\n 3. Verify all models remain with unchanged visibility.\n 4. Call sync_auto_mode_models with a *different* config.\n 5. Verify fetch_auto_mode_providers excludes this provider, so the\n periodic task would never call sync on it.\n \"\"\"\n initial_config = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\"],\n )\n\n try:\n # Step 1: Create in auto mode\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=initial_config,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n visibility_before = {\n mc.name: mc.is_visible for mc in provider.model_configurations\n }\n assert visibility_before == {\"gpt-4o\": True, \"gpt-4o-mini\": True}\n\n # Step 2: Switch to manual mode\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=None,\n api_key_changed=False,\n is_auto_mode=False,\n model_configurations=[\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ],\n ),\n is_creation=False,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 3: Models unchanged\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n assert provider.is_auto_mode is False\n visibility_after = {\n mc.name: mc.is_visible for mc in provider.model_configurations\n }\n assert visibility_after == visibility_before\n\n # Step 4-5: Provider excluded from auto mode queries\n auto_providers = fetch_auto_mode_providers(db_session)\n auto_provider_ids = {p.id for p in auto_providers}\n assert provider.id not in auto_provider_ids\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_resync_adds_new_and_hides_removed_models(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"When the GitHub config changes between syncs, a subsequent sync\n should add newly listed models and hide models that were removed.\n\n Steps:\n 1. Create provider in auto mode with config v1: [gpt-4o, gpt-4o-mini].\n 2. Sync with config v2: [gpt-4o, gpt-4-turbo] (gpt-4o-mini removed,\n gpt-4-turbo added).\n 3. Verify gpt-4o still visible, gpt-4o-mini hidden, gpt-4-turbo added\n and visible.\n \"\"\"\n config_v1 = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\"],\n )\n config_v2 = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4-turbo\"],\n )\n\n try:\n # Step 1: Create with config v1\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=config_v1,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 2: Re-sync with config v2\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n changes = sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config_v2,\n )\n assert changes > 0\n\n # Step 3: Verify\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n visibility = {\n mc.name: mc.is_visible for mc in provider.model_configurations\n }\n\n # gpt-4o: still in config -> visible\n assert visibility[\"gpt-4o\"] is True\n # gpt-4o-mini: removed from config -> hidden (not deleted)\n assert \"gpt-4o-mini\" in visibility, \"Removed model should still exist in DB\"\n assert visibility[\"gpt-4o-mini\"] is False\n # gpt-4-turbo: newly added -> visible\n assert visibility[\"gpt-4-turbo\"] is True\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_sync_is_idempotent(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Running sync twice with the same config should produce zero\n changes on the second call.\"\"\"\n config = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\", \"gpt-4-turbo\"],\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=config,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n # First explicit sync (may report changes if creation already synced)\n sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config,\n )\n\n # Snapshot state after first sync\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n snapshot = {\n mc.name: (mc.is_visible, mc.display_name)\n for mc in provider.model_configurations\n }\n\n # Second sync — should be a no-op\n changes = sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config,\n )\n assert (\n changes == 0\n ), f\"Expected 0 changes on idempotent re-sync, got {changes}\"\n\n # State should be identical\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n current = {\n mc.name: (mc.is_visible, mc.display_name)\n for mc in provider.model_configurations\n }\n assert current == snapshot\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_default_model_hidden_when_removed_from_config(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"When the current default model is removed from the config, sync\n should hide it. The default model flow row should still exist (it\n points at the ModelConfiguration), but the model is no longer visible.\n\n Steps:\n 1. Create provider with config: default=gpt-4o, additional=[gpt-4o-mini].\n 2. Set gpt-4o as the global default.\n 3. Re-sync with config: default=gpt-4o-mini (gpt-4o removed entirely).\n 4. Verify gpt-4o is hidden, gpt-4o-mini is visible, and\n fetch_default_llm_model still returns a result (the flow row persists).\n \"\"\"\n config_v1 = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\"],\n )\n config_v2 = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o-mini\",\n additional_models=[],\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=config_v1,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Step 2: Set gpt-4o as global default\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n default_before = fetch_default_llm_model(db_session)\n assert default_before is not None\n assert default_before.name == \"gpt-4o\"\n\n # Step 3: Re-sync with config v2 (gpt-4o removed)\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n changes = sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config_v2,\n )\n assert changes > 0\n\n # Step 4: Verify visibility\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n visibility = {\n mc.name: mc.is_visible for mc in provider.model_configurations\n }\n assert visibility[\"gpt-4o\"] is False, \"Removed default should be hidden\"\n assert visibility[\"gpt-4o-mini\"] is True, \"New default should be visible\"\n\n # The old default (gpt-4o) is now hidden. sync_auto_mode_models\n # should update the global default to the new recommended default\n # (gpt-4o-mini) so that it is not silently lost.\n db_session.expire_all()\n default_after = fetch_default_llm_model(db_session)\n assert (\n default_after is not None\n ), \"Default model should not be None — sync should set the new recommended default when the old one is hidden\"\n assert (\n default_after.name == \"gpt-4o-mini\"\n ), f\"Default should be updated to the new recommended model 'gpt-4o-mini', but got '{default_after.name}'\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_sync_updates_default_when_recommended_default_changes(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"When the provider owns the CHAT default and a sync arrives with a\n different recommended default model (both models still in config),\n the global default should be updated to the new recommendation.\n\n Steps:\n 1. Create auto-mode provider with config v1: default=gpt-4o.\n 2. Set gpt-4o as the global CHAT default.\n 3. Re-sync with config v2: default=gpt-4o-mini (gpt-4o still present).\n 4. Verify the CHAT default switched to gpt-4o-mini and both models\n remain visible.\n \"\"\"\n config_v1 = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\"],\n )\n config_v2 = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o-mini\",\n additional_models=[\"gpt-4o\"],\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=config_v1,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Set gpt-4o as the global CHAT default\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n default_before = fetch_default_llm_model(db_session)\n assert default_before is not None\n assert default_before.name == \"gpt-4o\"\n\n # Re-sync with config v2 (recommended default changed)\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n\n changes = sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config_v2,\n )\n assert changes > 0, \"Sync should report changes when default switches\"\n\n # Both models should remain visible\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n visibility = {\n mc.name: mc.is_visible for mc in provider.model_configurations\n }\n assert visibility[\"gpt-4o\"] is True\n assert visibility[\"gpt-4o-mini\"] is True\n\n # The CHAT default should now be gpt-4o-mini\n default_after = fetch_default_llm_model(db_session)\n assert default_after is not None\n assert (\n default_after.name == \"gpt-4o-mini\"\n ), f\"Default should be updated to 'gpt-4o-mini', got '{default_after.name}'\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n\n def test_sync_idempotent_when_default_already_matches(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"When the provider owns the CHAT default and it already matches the\n recommended default, re-syncing should report zero changes.\n\n This is a regression test for the bug where changes was unconditionally\n incremented even when the default was already correct.\n \"\"\"\n config = _create_mock_llm_recommendations(\n provider=LlmProviderNames.OPENAI,\n default_model_name=\"gpt-4o\",\n additional_models=[\"gpt-4o-mini\"],\n )\n\n try:\n with patch(\n \"onyx.server.manage.llm.api.fetch_llm_recommendations_from_github\",\n return_value=config,\n ):\n put_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n is_auto_mode=True,\n model_configurations=[],\n ),\n is_creation=True,\n _=_create_mock_admin(),\n db_session=db_session,\n )\n\n # Set gpt-4o (the recommended default) as global CHAT default\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n # First sync to stabilize state\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config,\n )\n\n # Second sync — default already matches, should be a no-op\n db_session.expire_all()\n provider = fetch_existing_llm_provider(\n name=provider_name, db_session=db_session\n )\n assert provider is not None\n changes = sync_auto_mode_models(\n db_session=db_session,\n provider=provider,\n llm_recommendations=config,\n )\n assert (\n changes == 0\n ), f\"Expected 0 changes when default already matches recommended, got {changes}\"\n\n # Default should still be gpt-4o\n default_model = fetch_default_llm_model(db_session)\n assert default_model is not None\n assert default_model.name == \"gpt-4o\"\n\n finally:\n db_session.rollback()\n _cleanup_provider(db_session, provider_name)\n" }, { "path": "backend/tests/external_dependency_unit/llm/test_llm_provider_called.py", "content": "from collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom fastapi_users.password import PasswordHelper\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.enums import AccountType\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import remove_llm_provider\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import User\nfrom onyx.db.models import UserRole\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom onyx.server.query_and_chat.chat_backend import create_new_chat_session\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.models import MessageResponseIDInfo\nfrom tests.external_dependency_unit.answer.stream_test_assertions import (\n assert_answer_stream_part_correct,\n)\nfrom tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder\nfrom tests.external_dependency_unit.answer.stream_test_utils import submit_query\nfrom tests.external_dependency_unit.answer.stream_test_utils import tokenise\nfrom tests.external_dependency_unit.mock_llm import LLMAnswerResponse\nfrom tests.external_dependency_unit.mock_llm import MockLLM\n\n\ndef _create_admin(db_session: Session) -> User:\n \"\"\"Create a mock admin user for testing.\"\"\"\n unique_email = f\"admin_{uuid4().hex[:8]}@example.com\"\n password_helper = PasswordHelper()\n password = password_helper.generate()\n hashed_password = password_helper.hash(password)\n\n user = User(\n id=uuid4(),\n email=unique_email,\n hashed_password=hashed_password,\n is_active=True,\n is_superuser=True,\n is_verified=True,\n role=UserRole.ADMIN,\n account_type=AccountType.STANDARD,\n )\n db_session.add(user)\n db_session.commit()\n db_session.refresh(user)\n return user\n\n\ndef _create_provider(\n db_session: Session,\n provider: LlmProviderNames,\n name: str,\n is_public: bool,\n) -> int:\n result = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=name,\n provider=provider,\n api_key=\"sk-ant-api03-...\",\n is_public=is_public,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"claude-3-5-sonnet-20240620\",\n is_visible=True,\n ),\n ],\n ),\n db_session=db_session,\n )\n return result.id\n\n\n@contextmanager\ndef use_mock_llm() -> (\n Generator[tuple[MockLLM, dict[str, bool | str | None]], None, None]\n):\n \"\"\"Context manager that patches LLM factory functions and tracks which ones are called.\"\"\"\n mock_llm = MockLLM()\n\n call_tracker: dict[str, bool | str | None] = {\n \"get_default_llm_called\": False,\n \"get_llm_called\": False,\n \"provider\": None,\n }\n\n def mock_get_default_llm(*_args: Any, **_kwargs: Any) -> MockLLM:\n call_tracker[\"get_default_llm_called\"] = True\n return mock_llm\n\n def mock_get_llm(provider: str, *_args: Any, **_kwargs: Any) -> MockLLM:\n call_tracker[\"get_llm_called\"] = True\n call_tracker[\"provider\"] = provider\n return mock_llm\n\n with (\n patch(\n \"onyx.llm.factory.get_default_llm\",\n side_effect=mock_get_default_llm,\n ),\n patch(\n \"onyx.llm.factory.get_llm\",\n side_effect=mock_get_llm,\n ),\n ):\n yield mock_llm, call_tracker\n\n\ndef _cleanup_provider(db_session: Session, name: str) -> None:\n \"\"\"Helper to clean up a test provider by name.\"\"\"\n provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n if provider:\n remove_llm_provider(db_session, provider.id)\n\n\ndef _assert_llm_calls(\n call_tracker: dict[str, bool | str | None], expected_provider: str\n) -> None:\n \"\"\"Assert that get_llm was called with expected provider and get_default_llm was not called.\"\"\"\n assert not call_tracker[\n \"get_default_llm_called\"\n ], \"get_default_llm should not be called when using private provider\"\n assert call_tracker[\n \"get_llm_called\"\n ], \"get_llm should be called when using private provider\"\n assert (\n call_tracker[\"provider\"] == expected_provider\n ), f\"Expected provider '{expected_provider}', got '{call_tracker['provider']}'\"\n\n\ndef _reset_call_tracker(call_tracker: dict[str, bool | str | None]) -> None:\n \"\"\"Reset the call tracker for the next test iteration.\"\"\"\n call_tracker[\"get_default_llm_called\"] = False\n call_tracker[\"get_llm_called\"] = False\n call_tracker[\"provider\"] = None\n\n\ndef test_user_sends_message_to_private_provider(\n db_session: Session,\n) -> None:\n \"\"\"Test that messages sent to a private provider use get_llm instead of get_default_llm.\"\"\"\n admin_user = _create_admin(db_session)\n\n # Create providers\n public_provider_id = _create_provider(\n db_session, LlmProviderNames.ANTHROPIC, \"public-provider\", True\n )\n _create_provider(db_session, LlmProviderNames.GOOGLE, \"private-provider\", False)\n\n update_default_provider(\n public_provider_id, \"claude-3-5-sonnet-20240620\", db_session\n )\n\n try:\n # Create chat session\n chat_session = create_new_chat_session(\n ChatSessionCreationRequest(),\n user=admin_user,\n db_session=db_session,\n )\n\n chat_session_id = chat_session.chat_session_id\n answer_tokens_1 = tokenise(\"Hello, how are you?\")\n answer_tokens_2 = tokenise(\"I'm good, thank you!\")\n\n with use_mock_llm() as (mock_llm, call_tracker):\n handler = StreamTestBuilder(llm_controller=mock_llm)\n\n # First message\n handler.add_response(LLMAnswerResponse(answer_tokens=answer_tokens_1))\n answer_stream = submit_query(\n query=\"Hello, how are you?\",\n chat_session_id=chat_session_id,\n db_session=db_session,\n user=admin_user,\n llm_override=LLMOverride(\n model_provider=\"private-provider\",\n model_version=\"claude-3-5-sonnet-20240620\",\n ),\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=1,\n reserved_assistant_message_id=1,\n ),\n )\n\n handler.expect_agent_response(\n answer_tokens=answer_tokens_1,\n turn_index=0,\n ).run_and_validate(stream=answer_stream)\n\n with pytest.raises(StopIteration):\n next(answer_stream)\n\n _assert_llm_calls(call_tracker, \"google\")\n _reset_call_tracker(call_tracker)\n\n # Second message\n handler.add_response(LLMAnswerResponse(answer_tokens=answer_tokens_2))\n answer_stream = submit_query(\n query=\"I'm good, thank you!\",\n chat_session_id=chat_session_id,\n db_session=db_session,\n user=admin_user,\n llm_override=LLMOverride(\n model_provider=\"private-provider\",\n model_version=\"claude-3-5-sonnet-20240620\",\n ),\n )\n\n assert_answer_stream_part_correct(\n received=next(answer_stream),\n expected=MessageResponseIDInfo(\n user_message_id=2,\n reserved_assistant_message_id=2,\n ),\n )\n\n handler.expect_agent_response(\n answer_tokens=answer_tokens_2,\n turn_index=0,\n ).run_and_validate(stream=answer_stream)\n\n with pytest.raises(StopIteration):\n next(answer_stream)\n\n _assert_llm_calls(call_tracker, \"google\")\n\n finally:\n _cleanup_provider(db_session, \"public-provider\")\n _cleanup_provider(db_session, \"private-provider\")\n" }, { "path": "backend/tests/external_dependency_unit/llm/test_llm_provider_default_model_protection.py", "content": "\"\"\"\nThis should act as the main point of reference for testing that default model\nlogic is consisten.\n\n -\n\"\"\"\n\nfrom collections.abc import Generator\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import remove_llm_provider\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import update_default_vision_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\n\n\ndef _create_test_provider(\n db_session: Session,\n name: str,\n models: list[ModelConfigurationUpsertRequest] | None = None,\n) -> LLMProviderView:\n \"\"\"Helper to create a test LLM provider with multiple models.\"\"\"\n if models is None:\n models = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True, supports_image_input=True\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True, supports_image_input=False\n ),\n ]\n return upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=models,\n ),\n db_session=db_session,\n )\n\n\ndef _cleanup_provider(db_session: Session, name: str) -> None:\n \"\"\"Helper to clean up a test provider by name.\"\"\"\n provider = fetch_existing_llm_provider(name=name, db_session=db_session)\n if provider:\n remove_llm_provider(db_session, provider.id)\n\n\n@pytest.fixture\ndef provider_name(db_session: Session) -> Generator[str, None, None]:\n \"\"\"Generate a unique provider name for each test, with automatic cleanup.\"\"\"\n name = f\"test-provider-{uuid4().hex[:8]}\"\n yield name\n db_session.rollback()\n _cleanup_provider(db_session, name)\n\n\nclass TestDefaultModelProtection:\n \"\"\"Tests that the default model cannot be removed or hidden.\"\"\"\n\n def test_cannot_remove_default_text_model(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Removing the default text model from a provider should raise ValueError.\"\"\"\n provider = _create_test_provider(db_session, provider_name)\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n # Try to update the provider without the default model\n with pytest.raises(ValueError, match=\"Cannot remove the default model\"):\n upsert_llm_provider(\n LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ],\n ),\n db_session=db_session,\n )\n\n def test_cannot_hide_default_text_model(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Setting is_visible=False on the default text model should raise ValueError.\"\"\"\n provider = _create_test_provider(db_session, provider_name)\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n # Try to hide the default model\n with pytest.raises(ValueError, match=\"Cannot hide the default model\"):\n upsert_llm_provider(\n LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=False\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ],\n ),\n db_session=db_session,\n )\n\n def test_cannot_remove_default_vision_model(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Removing the default vision model from a provider should raise ValueError.\"\"\"\n provider = _create_test_provider(db_session, provider_name)\n # Set gpt-4o as both the text and vision default\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n update_default_vision_provider(provider.id, \"gpt-4o\", db_session)\n\n # Try to remove the default vision model\n with pytest.raises(ValueError, match=\"Cannot remove the default model\"):\n upsert_llm_provider(\n LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ),\n ],\n ),\n db_session=db_session,\n )\n\n def test_can_remove_non_default_model(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Removing a non-default model should succeed.\"\"\"\n provider = _create_test_provider(db_session, provider_name)\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n # Remove gpt-4o-mini (not default) — should succeed\n updated = upsert_llm_provider(\n LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True, supports_image_input=True\n ),\n ],\n ),\n db_session=db_session,\n )\n\n model_names = {mc.name for mc in updated.model_configurations}\n assert \"gpt-4o\" in model_names\n assert \"gpt-4o-mini\" not in model_names\n\n def test_can_hide_non_default_model(\n self,\n db_session: Session,\n provider_name: str,\n ) -> None:\n \"\"\"Hiding a non-default model should succeed.\"\"\"\n provider = _create_test_provider(db_session, provider_name)\n update_default_provider(provider.id, \"gpt-4o\", db_session)\n\n # Hide gpt-4o-mini (not default) — should succeed\n updated = upsert_llm_provider(\n LLMProviderUpsertRequest(\n id=provider.id,\n name=provider_name,\n provider=LlmProviderNames.OPENAI,\n api_key=\"sk-test-key-00000000000000000000000000000000000\",\n api_key_changed=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True, supports_image_input=True\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=False\n ),\n ],\n ),\n db_session=db_session,\n )\n\n model_visibility = {\n mc.name: mc.is_visible for mc in updated.model_configurations\n }\n assert model_visibility[\"gpt-4o\"] is True\n assert model_visibility[\"gpt-4o-mini\"] is False\n" }, { "path": "backend/tests/external_dependency_unit/llm/test_prompt_caching.py", "content": "\"\"\"External dependency unit tests for prompt caching functionality.\n\nThese tests call LLM providers directly and use litellm's completion_cost() to verify\nthat prompt caching reduces costs.\n\"\"\"\n\nimport json\nimport os\nimport tempfile\nimport time\nfrom pathlib import Path\nfrom typing import Any\n\nimport pytest\nfrom litellm import completion_cost\nfrom sqlalchemy.orm import Session\n\nfrom onyx.llm.model_response import Usage\nfrom onyx.llm.models import AssistantMessage\nfrom onyx.llm.models import ChatCompletionMessage\nfrom onyx.llm.models import SystemMessage\nfrom onyx.llm.models import UserMessage\nfrom onyx.llm.multi_llm import LitellmLLM\nfrom onyx.llm.prompt_cache.processor import process_with_prompt_cache\n\n\nVERTEX_CREDENTIALS_ENV = \"VERTEX_CREDENTIALS\"\nVERTEX_LOCATION_ENV = \"VERTEX_LOCATION\"\nVERTEX_MODEL_ENV = \"VERTEX_MODEL_NAME\"\nDEFAULT_VERTEX_MODEL = \"gemini-2.5-flash\"\n\n\ndef _extract_cached_tokens(usage: Usage | None) -> int:\n \"\"\"Helper to extract cached_tokens from usage (dict or object).\"\"\"\n if not usage:\n print(\"Usage is None\")\n return 0\n\n cached_tokens = usage.cache_creation_input_tokens\n\n return cached_tokens\n\n\ndef _extract_prompt_tokens(usage: Usage | None) -> int:\n \"\"\"Helper to extract prompt_tokens from usage (dict or object).\"\"\"\n if not usage:\n print(\"Usage is None\")\n return 0\n\n return usage.prompt_tokens\n\n\ndef _extract_cache_read_tokens(usage: Usage | None) -> int:\n \"\"\"Extract cache read metrics from usage (dict or object).\"\"\"\n print(f\"usage: {usage}\")\n if not usage:\n print(\"Usage is None\")\n return 0\n\n return usage.cache_read_input_tokens\n\n\ndef _get_usage_value(usage: Any, key: str) -> int:\n \"\"\"Retrieve a numeric field from usage objects or dictionaries.\"\"\"\n if isinstance(usage, dict):\n value = usage.get(key)\n else:\n value = getattr(usage, key, None)\n return int(value or 0)\n\n\ndef _resolve_vertex_credentials() -> tuple[Path, bool]:\n \"\"\"Return a path to credentials; support inline JSON or filesystem path.\"\"\"\n raw_value = os.environ.get(VERTEX_CREDENTIALS_ENV)\n if not raw_value:\n raise FileNotFoundError(\"Vertex credentials environment variable not set.\")\n\n raw_value = raw_value.strip()\n candidate_path = Path(raw_value)\n if len(raw_value) < 100 and candidate_path.exists():\n return candidate_path, False\n\n try:\n json.loads(raw_value)\n except json.JSONDecodeError as exc:\n raise ValueError(\n \"Vertex credentials must be a valid JSON string or file path.\"\n ) from exc\n\n temp_file = tempfile.NamedTemporaryFile(\n mode=\"w\", suffix=\".json\", delete=False, encoding=\"utf-8\"\n )\n try:\n temp_file.write(raw_value)\n temp_file.flush()\n finally:\n temp_file.close()\n return Path(temp_file.name), True\n\n\ndef _validate_vertex_credentials_file(credentials_path: Path) -> None:\n \"\"\"Validate that the credentials file contains a usable service account.\"\"\"\n try:\n content = credentials_path.read_text(encoding=\"utf-8\")\n except OSError as exc:\n raise ValueError(f\"Failed to read credentials file: {exc}\") from exc\n\n try:\n data = json.loads(content)\n except json.JSONDecodeError as exc:\n raise ValueError(\"Credentials file does not contain valid JSON.\") from exc\n\n if not isinstance(data, dict):\n raise ValueError(\"Credentials JSON must be an object.\")\n\n cred_type = data.get(\"type\")\n if cred_type != \"service_account\":\n raise ValueError(\n f\"Unsupported credential type '{cred_type}'. Provide a service_account JSON blob.\"\n )\n\n missing_fields = [\n field\n for field in (\"project_id\", \"client_email\", \"private_key\")\n if not data.get(field)\n ]\n if missing_fields:\n raise ValueError(\n \"Missing required service account fields: \"\n + \", \".join(sorted(missing_fields))\n )\n\n try:\n from google.oauth2 import service_account\n\n service_account.Credentials.from_service_account_info(\n data,\n scopes=[\"https://www.googleapis.com/auth/cloud-platform\"],\n )\n except (\n Exception\n ) as exc: # pragma: no cover - depends on google SDK validation paths\n raise ValueError(\n f\"Failed to construct service account credentials: {exc}\"\n ) from exc\n\n\n@pytest.mark.skip(reason=\"OpenAI prompt caching is unreliable\")\n@pytest.mark.skipif(\n not os.environ.get(\"OPENAI_API_KEY\"),\n reason=\"OpenAI API key not available\",\n)\ndef test_openai_prompt_caching_reduces_costs(\n db_session: Session, # noqa: ARG001\n) -> None:\n \"\"\"Test that OpenAI prompt caching reduces costs on subsequent calls.\n\n OpenAI uses implicit caching for prompts >1024 tokens.\n \"\"\"\n attempts = 8\n successes = 0\n for _ in range(attempts):\n # Create OpenAI LLM\n llm = LitellmLLM(\n api_key=os.environ[\"OPENAI_API_KEY\"],\n model_provider=\"openai\",\n model_name=\"gpt-4o\",\n max_input_tokens=128000,\n )\n import random\n import string\n\n # Insert 32 random lowercase characters at the start of long_context\n # to prevent holdover cache from previous tests\n random_prefix = \"\".join(random.choices(string.ascii_lowercase, k=32))\n # Create a long context message to ensure caching threshold is met (>1024 tokens)\n long_context = (\n random_prefix\n + \"This is a comprehensive document about artificial intelligence and machine learning. \"\n + \" \".join(\n [\n f\"Section {i}: This section discusses various aspects of AI technology, \"\n f\"including neural networks, deep learning, natural language processing, \"\n f\"computer vision, and reinforcement learning. These technologies are \"\n f\"revolutionizing how we interact with computers and process information.\"\n for i in range(50)\n ]\n )\n )\n\n # Split into cacheable prefix (the long context) and suffix (the question)\n cacheable_prefix: list[ChatCompletionMessage] = [\n UserMessage(role=\"user\", content=long_context)\n ]\n\n # First call - creates cache\n print(\"\\n=== First call (cache creation) ===\")\n question1: list[ChatCompletionMessage] = [\n UserMessage(role=\"user\", content=\"What are the main topics discussed?\")\n ]\n\n # Apply prompt caching (for OpenAI, this is mostly a no-op but should still work)\n processed_messages1, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=question1,\n continuation=False,\n )\n # print(f\"Processed messages 1: {processed_messages1}\")\n # print(f\"Metadata 1: {metadata1}\")\n # print(f\"Cache key 1: {metadata1.cache_key if metadata1 else None}\")\n\n # Call litellm directly so we can get the raw response\n response1 = llm.invoke(prompt=processed_messages1)\n cost1 = completion_cost(\n completion_response=response1.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage1 = response1.usage\n cached_tokens_1 = _extract_cached_tokens(usage1)\n prompt_tokens_1 = _extract_prompt_tokens(usage1)\n # print(f\"Response 1 usage: {usage1}\")\n # print(f\"Cost 1: ${cost1:.10f}\")\n\n # Wait to ensure cache is available\n time.sleep(5)\n\n # Second call with same context - should use cache\n print(\"\\n=== Second call (cache read) ===\")\n question2: list[ChatCompletionMessage] = [\n UserMessage(role=\"user\", content=\"Can you elaborate on neural networks?\")\n ]\n\n # Apply prompt caching (same cacheable prefix)\n processed_messages2, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=question2,\n continuation=False,\n )\n # print(f\"Processed messages 2: {processed_messages2}\")\n response2 = llm.invoke(prompt=processed_messages2)\n cost2 = completion_cost(\n completion_response=response2.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage2 = response2.usage\n cached_tokens_2 = _extract_cache_read_tokens(usage2)\n prompt_tokens_2 = _extract_prompt_tokens(usage2)\n # print(f\"Response 2 usage: {usage2}\")\n # print(f\"Cost 2: ${cost2:.10f}\")\n\n # Verify caching occurred – OpenAI reports cached work via prompt_tokens_details.cached_tokens\n print(f\"\\nCached tokens call 1: {cached_tokens_1}, call 2: {cached_tokens_2}\")\n print(f\"Prompt tokens call 1: {prompt_tokens_1}, call 2: {prompt_tokens_2}\")\n print(f\"Cost delta (1 -> 2): ${cost1 - cost2:.10f}\")\n\n # The first call is expected to *create* cache (cached_tokens may be 0).\n # The second call should show cached tokens being used.\n if cached_tokens_2 > 0:\n successes += 1\n break\n\n # empirically there's a 60% chance of success per attempt, so we expect at least one success in 8 attempts\n # (99.94% probability). we can bump this number if the test is too flaky.\n assert (\n successes > 0\n ), f\"Expected at least one success. 0 of {attempts} attempts used prompt caching.\"\n\n\n@pytest.mark.skipif(\n not os.environ.get(\"ANTHROPIC_API_KEY\"),\n reason=\"Anthropic API key not available\",\n)\ndef test_anthropic_prompt_caching_reduces_costs(\n db_session: Session, # noqa: ARG001\n) -> None:\n \"\"\"Test that Anthropic prompt caching reduces costs on subsequent calls.\n\n Anthropic requires explicit cache_control parameters.\n \"\"\"\n # Prompt caching support is model/account specific.\n # Allow override via env var and otherwise try a few non-retired candidates.\n anthropic_prompt_cache_models_env = os.environ.get(\"ANTHROPIC_PROMPT_CACHE_MODELS\")\n if anthropic_prompt_cache_models_env:\n candidate_models = [\n model.strip()\n for model in anthropic_prompt_cache_models_env.split(\",\")\n if model.strip()\n ]\n else:\n candidate_models = [\n \"claude-haiku-4-5-20251001\",\n \"claude-sonnet-4-5-20250929\",\n \"claude-3-5-sonnet-20241022\",\n \"claude-3-5-sonnet-latest\",\n ]\n\n import random\n import string\n\n # Create a long context message.\n # Add a random prefix to avoid reusing an existing ephemeral cache from prior test runs.\n random_prefix = \"\".join(random.choices(string.ascii_lowercase, k=32))\n long_context = (\n random_prefix + \" \"\n \"This is a comprehensive document about artificial intelligence and machine learning. \"\n + \" \".join(\n [\n f\"Section {i}: This section discusses various aspects of AI technology, \"\n f\"including neural networks, deep learning, natural language processing, \"\n f\"computer vision, and reinforcement learning. These technologies are \"\n f\"revolutionizing how we interact with computers and process information.\"\n for i in range(50)\n ]\n )\n )\n\n base_messages: list[ChatCompletionMessage] = [\n UserMessage(role=\"user\", content=long_context)\n ]\n\n unavailable_models: list[str] = []\n non_caching_models: list[str] = []\n\n for model_name in candidate_models:\n llm = LitellmLLM(\n api_key=os.environ[\"ANTHROPIC_API_KEY\"],\n model_provider=\"anthropic\",\n model_name=model_name,\n max_input_tokens=200000,\n )\n\n # First call - creates cache\n print(f\"\\n=== First call (cache creation) model={model_name} ===\")\n question1: list[ChatCompletionMessage] = [\n UserMessage(\n role=\"user\",\n content=\"Reply with exactly one lowercase word: topics\",\n )\n ]\n\n processed_messages1, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=base_messages,\n suffix=question1,\n continuation=False,\n )\n\n try:\n response1 = llm.invoke(prompt=processed_messages1, max_tokens=8)\n except Exception as e:\n error_str = str(e).lower()\n if (\n \"not_found_error\" in error_str\n or \"model_not_found\" in error_str\n or ('\"type\":\"not_found_error\"' in error_str and \"model:\" in error_str)\n ):\n unavailable_models.append(model_name)\n continue\n raise\n\n cost1 = completion_cost(\n completion_response=response1.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage1 = response1.usage\n print(f\"Response 1 usage: {usage1}\")\n print(f\"Cost 1: ${cost1:.10f}\")\n\n # Wait to ensure cache is available\n time.sleep(2)\n\n # Second call with same context - should use cache\n print(f\"\\n=== Second call (cache read) model={model_name} ===\")\n question2: list[ChatCompletionMessage] = [\n UserMessage(\n role=\"user\",\n content=\"Reply with exactly one lowercase word: neural\",\n )\n ]\n\n processed_messages2, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=base_messages,\n suffix=question2,\n continuation=False,\n )\n\n response2 = llm.invoke(prompt=processed_messages2, max_tokens=8)\n cost2 = completion_cost(\n completion_response=response2.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage2 = response2.usage\n print(f\"Response 2 usage: {usage2}\")\n print(f\"Cost 2: ${cost2:.10f}\")\n\n cache_creation_tokens = _get_usage_value(usage1, \"cache_creation_input_tokens\")\n cache_read_tokens = _get_usage_value(usage2, \"cache_read_input_tokens\")\n\n print(f\"\\nCache creation tokens (call 1): {cache_creation_tokens}\")\n print(f\"Cache read tokens (call 2): {cache_read_tokens}\")\n print(f\"Cost reduction: ${cost1 - cost2:.10f}\")\n\n # Model is available but does not expose Anthropic cache usage metrics\n if cache_creation_tokens <= 0 or cache_read_tokens <= 0:\n non_caching_models.append(model_name)\n continue\n\n # Cost should be lower on second call\n assert (\n cost2 < cost1\n ), f\"Expected lower cost on cached call. Cost 1: ${cost1:.10f}, Cost 2: ${cost2:.10f}\"\n return\n\n pytest.skip(\n \"No Anthropic model available with observable prompt-cache metrics. \"\n f\"Tried models={candidate_models}, unavailable={unavailable_models}, non_caching={non_caching_models}\"\n )\n\n\n@pytest.mark.skipif(\n not os.environ.get(VERTEX_CREDENTIALS_ENV),\n reason=\"Vertex AI credentials file not available\",\n)\n@pytest.mark.skipif(\n not os.environ.get(VERTEX_LOCATION_ENV),\n reason=\"VERTEX_LOCATION required for Vertex AI context caching (e.g., 'us-central1')\",\n)\n@pytest.mark.skip(reason=\"Vertex AI prompt caching is disabled for now\")\ndef test_google_genai_prompt_caching_reduces_costs(\n db_session: Session, # noqa: ARG001\n) -> None:\n \"\"\"Test that Litellm Gemini prompt caching reduces costs on subsequent calls.\n\n Vertex AI requires explicit context caching via the Context Caching API,\n which needs both credentials and a valid location (e.g., us-central1).\n \"\"\"\n import random\n import string\n from litellm import exceptions as litellm_exceptions\n\n try:\n credentials_path, should_cleanup = _resolve_vertex_credentials()\n except FileNotFoundError:\n pytest.skip(\"Vertex credentials not available for test.\")\n except ValueError as exc:\n pytest.skip(str(exc))\n\n vertex_location = os.environ.get(VERTEX_LOCATION_ENV)\n if not vertex_location:\n pytest.skip(\"VERTEX_LOCATION required for Vertex AI context caching\")\n model_name = os.environ.get(VERTEX_MODEL_ENV, DEFAULT_VERTEX_MODEL)\n\n try:\n _validate_vertex_credentials_file(credentials_path)\n os.environ.setdefault(\"GOOGLE_APPLICATION_CREDENTIALS\", str(credentials_path))\n\n custom_config: dict[str, str] = {\"vertex_credentials\": str(credentials_path)}\n if vertex_location:\n custom_config[\"vertex_location\"] = vertex_location\n\n llm = LitellmLLM(\n api_key=None,\n model_provider=\"vertex_ai\",\n model_name=model_name,\n max_input_tokens=1_000_000,\n custom_config=custom_config,\n )\n\n attempts = 4\n success = False\n last_metrics: dict[str, Any] = {}\n\n for attempt in range(attempts):\n random_prefix = \"\".join(random.choices(string.ascii_lowercase, k=32))\n long_context = (\n random_prefix\n + \"This is a comprehensive document about artificial intelligence and machine learning. \"\n + \" \".join(\n [\n f\"Section {i}: This section discusses various aspects of AI technology, \"\n f\"including neural networks, deep learning, natural language processing, \"\n f\"computer vision, and reinforcement learning. These technologies are \"\n f\"revolutionizing how we interact with computers and process information.\"\n for i in range(50)\n ]\n )\n )\n\n cacheable_prefix: list[ChatCompletionMessage] = [\n SystemMessage(role=\"system\", content=long_context)\n ]\n\n print(f\"\\n=== Vertex attempt {attempt + 1} (cache creation) ===\")\n question1: list[ChatCompletionMessage] = [\n UserMessage(role=\"user\", content=\"What are the main topics discussed?\")\n ]\n\n processed_messages1, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=question1,\n continuation=False,\n )\n # Debug: print processed messages structure\n first_msg = (\n processed_messages1[0]\n if isinstance(processed_messages1, list) and processed_messages1\n else processed_messages1\n )\n print(f\"Processed messages structure (first msg): {first_msg}\")\n\n response1 = llm.invoke(prompt=processed_messages1)\n cost1 = completion_cost(\n completion_response=response1.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n usage1 = response1.usage\n cache_creation_tokens = _get_usage_value(\n usage1, \"cache_creation_input_tokens\"\n )\n cached_tokens_1 = _extract_cached_tokens(usage1)\n cache_read_tokens_1 = _extract_cache_read_tokens(usage1)\n\n print(f\"Vertex response 1 usage: {usage1}\")\n print(f\"Vertex cost 1: ${cost1:.10f}\")\n\n time.sleep(5)\n\n print(f\"\\n=== Vertex attempt {attempt + 1} (cache read) ===\")\n question2: list[ChatCompletionMessage] = [\n UserMessage(\n role=\"user\", content=\"Can you elaborate on neural networks?\"\n )\n ]\n\n processed_messages2, _ = process_with_prompt_cache(\n llm_config=llm.config,\n cacheable_prefix=cacheable_prefix,\n suffix=question2,\n continuation=False,\n )\n\n response2 = llm.invoke(prompt=processed_messages2)\n cost2 = completion_cost(\n completion_response=response2.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n usage2 = response2.usage\n cache_read_tokens_2 = _extract_cache_read_tokens(usage2)\n cached_tokens_2 = _extract_cached_tokens(usage2)\n\n print(f\"Vertex response 2 usage: {usage2}\")\n print(f\"Vertex cost 2: ${cost2:.10f}\")\n print(\n f\"Vertex cache metrics - creation: {cache_creation_tokens}, \"\n f\"call1 cached tokens: {cached_tokens_1}, \"\n f\"call1 cache read tokens: {cache_read_tokens_1}, \"\n f\"call2 cached tokens: {cached_tokens_2}, \"\n f\"call2 cache read tokens: {cache_read_tokens_2}\"\n )\n print(f\"Vertex cost delta (1 -> 2): ${cost1 - cost2:.10f}\")\n\n last_metrics = {\n \"cache_creation_tokens\": cache_creation_tokens,\n \"cached_tokens_1\": cached_tokens_1,\n \"cache_read_tokens_1\": cache_read_tokens_1,\n \"cached_tokens_2\": cached_tokens_2,\n \"cache_read_tokens_2\": cache_read_tokens_2,\n \"cost_delta\": cost1 - cost2,\n }\n\n if cache_read_tokens_2 > 0 or cached_tokens_2 > 0 or (cost1 - cost2) > 0:\n success = True\n break\n except ValueError as exc:\n pytest.fail(f\"Invalid Vertex credentials: {exc}\")\n except litellm_exceptions.APIConnectionError as exc:\n creds_details = json.loads(credentials_path.read_text(encoding=\"utf-8\"))\n pytest.fail(\n \"Vertex credentials appeared well-formed but failed to mint an access token. \"\n \"This typically means the service account lacks the required Vertex AI permissions \"\n \"or the key was revoked.\\n\"\n f\"project_id={creds_details.get('project_id')!r}, \"\n f\"client_email={creds_details.get('client_email')!r}\\n\"\n f\"Original error: {exc}\"\n )\n finally:\n if should_cleanup:\n try:\n credentials_path.unlink(missing_ok=True)\n except OSError:\n pass\n\n assert (\n success\n ), f\"Expected Gemini prompt caching evidence across attempts. Last observed metrics: {last_metrics}\"\n\n\n@pytest.mark.skipif(\n not os.environ.get(\"OPENAI_API_KEY\"),\n reason=\"OpenAI API key not available\",\n)\ndef test_prompt_caching_with_conversation_history(\n db_session: Session, # noqa: ARG001\n) -> None:\n \"\"\"Test that prompt caching works with multi-turn conversations.\n\n System message and history should be cached, only new user message is uncached.\n \"\"\"\n # Create OpenAI LLM\n llm = LitellmLLM(\n api_key=os.environ[\"OPENAI_API_KEY\"],\n model_provider=\"openai\",\n model_name=\"gpt-4o-mini\",\n max_input_tokens=128000,\n )\n\n # Create a long system message and context\n system_message: SystemMessage = SystemMessage(\n role=\"system\",\n content=(\n \"You are an AI assistant specialized in technology. \"\n + \" \".join(\n [\n f\"You have knowledge about topic {i} including detailed information. \"\n for i in range(50)\n ]\n )\n ),\n )\n\n long_context = \"This is a comprehensive document. \" + \" \".join(\n [f\"Section {i}: Details about topic {i}. \" * 20 for i in range(30)]\n )\n\n # Turn 1\n print(\"\\n=== Turn 1 ===\")\n messages_turn1: list[ChatCompletionMessage] = [\n system_message,\n UserMessage(role=\"user\", content=long_context + \"\\n\\nWhat is this about?\"),\n ]\n\n response1 = llm.invoke(prompt=messages_turn1)\n cost1 = completion_cost(\n completion_response=response1.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage1 = response1.usage\n print(f\"Turn 1 usage: {usage1}\")\n print(f\"Turn 1 cost: ${cost1:.10f}\")\n\n # Wait for cache\n time.sleep(2)\n\n # Turn 2 - add assistant response and new user message\n print(\"\\n=== Turn 2 (with cached history) ===\")\n messages_turn2: list[ChatCompletionMessage] = messages_turn1 + [\n AssistantMessage(\n role=\"assistant\", content=\"This document discusses various topics.\"\n ),\n UserMessage(role=\"user\", content=\"Tell me about the first topic.\"),\n ]\n\n response2 = llm.invoke(prompt=messages_turn2)\n cost2 = completion_cost(\n completion_response=response2.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage2 = response2.usage\n print(f\"Turn 2 usage: {usage2}\")\n print(f\"Turn 2 cost: ${cost2:.10f}\")\n\n # Turn 3 - continue conversation\n print(\"\\n=== Turn 3 (with even more cached history) ===\")\n messages_turn3: list[ChatCompletionMessage] = messages_turn2 + [\n AssistantMessage(role=\"assistant\", content=\"The first topic covers...\"),\n UserMessage(role=\"user\", content=\"What about the second topic?\"),\n ]\n\n response3 = llm.invoke(prompt=messages_turn3)\n cost3 = completion_cost(\n completion_response=response3.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage3 = response3.usage\n print(f\"Turn 3 usage: {usage3}\")\n print(f\"Turn 3 cost: ${cost3:.10f}\")\n\n # Verify caching in subsequent turns\n cache_tokens_2 = _get_usage_value(usage2, \"cache_read_input_tokens\")\n cache_tokens_3 = _get_usage_value(usage3, \"cache_read_input_tokens\")\n\n prompt_tokens_1 = _get_usage_value(usage1, \"prompt_tokens\")\n prompt_tokens_2 = _get_usage_value(usage2, \"prompt_tokens\")\n prompt_tokens_3 = _get_usage_value(usage3, \"prompt_tokens\")\n\n print(f\"\\nCache tokens - Turn 2: {cache_tokens_2}, Turn 3: {cache_tokens_3}\")\n print(\n f\"Prompt tokens - Turn 1: {prompt_tokens_1}, Turn 2: {prompt_tokens_2}, Turn 3: {prompt_tokens_3}\"\n )\n\n # Either cache tokens should increase or prompt tokens should be relatively stable\n # (not growing linearly with conversation length)\n assert (\n cache_tokens_2 > 0\n or cache_tokens_3 > 0\n or prompt_tokens_2 < prompt_tokens_1 * 1.5\n ), \"Expected caching benefits in multi-turn conversation\"\n\n\n@pytest.mark.skipif(\n not os.environ.get(\"OPENAI_API_KEY\"),\n reason=\"OpenAI API key not available\",\n)\ndef test_no_caching_without_process_with_prompt_cache(\n db_session: Session, # noqa: ARG001\n) -> None:\n \"\"\"Test baseline: without using process_with_prompt_cache, no special caching occurs.\n\n This establishes a baseline to compare against the caching tests.\n \"\"\"\n # Create OpenAI LLM\n llm = LitellmLLM(\n api_key=os.environ[\"OPENAI_API_KEY\"],\n model_provider=\"openai\",\n model_name=\"gpt-4o-mini\",\n max_input_tokens=128000,\n )\n\n # Create a long context\n long_context = \"This is a comprehensive document. \" + \" \".join(\n [f\"Section {i}: Details about technology topic {i}. \" * 10 for i in range(50)]\n )\n\n # First call - no explicit caching\n print(\"\\n=== First call (no explicit caching) ===\")\n messages1: list[ChatCompletionMessage] = [\n UserMessage(role=\"user\", content=long_context + \"\\n\\nSummarize this.\")\n ]\n\n response1 = llm.invoke(prompt=messages1)\n cost1 = completion_cost(\n completion_response=response1.model_dump(),\n model=f\"{llm._model_provider}/{llm._model_version}\",\n )\n\n usage1 = response1.usage\n print(f\"Response 1 usage: {usage1}\")\n print(f\"Cost 1: ${cost1:.10f}\")\n\n # This test just verifies the LLM works and we can calculate costs\n # It serves as a baseline comparison for the caching tests\n assert cost1 > 0, \"Should have non-zero cost\"\n assert usage1, \"Should have usage data\"\n\n print(\"\\nBaseline test passed - ready to compare with caching tests\")\n" }, { "path": "backend/tests/external_dependency_unit/mock_content_provider.py", "content": "import abc\nfrom collections.abc import Generator\nfrom collections.abc import Sequence\nfrom contextlib import contextmanager\nfrom unittest.mock import patch\n\nfrom pydantic import BaseModel\n\nfrom onyx.tools.tool_implementations.open_url.models import WebContent\nfrom onyx.tools.tool_implementations.open_url.models import WebContentProvider\n\n\nclass MockWebContent(BaseModel):\n title: str\n url: str\n content: str\n\n def to_web_content(self) -> WebContent:\n return WebContent(\n title=self.title,\n link=self.url,\n full_content=self.content,\n published_date=None,\n scrape_successful=True,\n )\n\n\nclass ContentProviderController(abc.ABC):\n @abc.abstractmethod\n def add_content(self, content: MockWebContent) -> None:\n raise NotImplementedError\n\n\nclass MockContentProvider(WebContentProvider, ContentProviderController):\n def __init__(self) -> None:\n self._contents: list[MockWebContent] = []\n\n def add_content(self, web_content: MockWebContent) -> None:\n self._contents.append(web_content)\n\n def contents(self, urls: Sequence[str]) -> list[WebContent]:\n filtered_contents = list(\n filter(lambda web_content: web_content.url in urls, self._contents)\n )\n\n return list(\n map(lambda web_content: web_content.to_web_content(), filtered_contents)\n )\n\n\n@contextmanager\ndef use_mock_content_provider() -> Generator[ContentProviderController, None, None]:\n content_provider = MockContentProvider()\n\n with patch(\n \"onyx.tools.tool_implementations.open_url.open_url_tool.get_default_content_provider\",\n return_value=content_provider,\n ):\n yield content_provider\n" }, { "path": "backend/tests/external_dependency_unit/mock_image_provider.py", "content": "import abc\nimport asyncio\nimport concurrent.futures\nimport time\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom datetime import datetime\nfrom typing import Any\nfrom unittest.mock import patch\n\nfrom litellm.types.utils import ImageObject\nfrom litellm.types.utils import ImageResponse\n\nfrom onyx.image_gen.interfaces import ImageGenerationProvider\nfrom onyx.image_gen.interfaces import ImageGenerationProviderCredentials\nfrom onyx.image_gen.interfaces import ReferenceImage\nfrom onyx.llm.interfaces import LLMConfig\n\n\nclass ImageGenerationProviderController(abc.ABC):\n @abc.abstractmethod\n def add_image(\n self,\n data: str,\n delay: float = 0.0,\n ) -> None:\n raise NotImplementedError\n\n\nclass MockImageGenerationProvider(\n ImageGenerationProvider, ImageGenerationProviderController\n):\n def __init__(self) -> None:\n self._images: list[str] = []\n self._delays: list[float] = []\n\n def add_image(\n self,\n data: str,\n delay: float = 0.0,\n ) -> None:\n self._images.append(data)\n self._delays.append(delay)\n\n @classmethod\n def validate_credentials(\n cls,\n credentials: ImageGenerationProviderCredentials, # noqa: ARG003\n ) -> bool:\n return True\n\n @classmethod\n def _build_from_credentials(\n cls,\n _: ImageGenerationProviderCredentials,\n ) -> ImageGenerationProvider:\n return cls()\n\n def generate_image(\n self,\n prompt: str,\n model: str, # noqa: ARG002\n size: str, # noqa: ARG002\n n: int, # noqa: ARG002\n quality: str | None = None, # noqa: ARG002\n reference_images: list[ReferenceImage] | None = None, # noqa: ARG002\n **kwargs: Any, # noqa: ARG002\n ) -> ImageResponse:\n image_data = self._images.pop(0)\n delay = self._delays.pop(0)\n\n if delay > 0.0:\n try:\n asyncio.get_running_loop()\n # Event loop is running - run sleep in executor to avoid blocking the event loop\n with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:\n future = executor.submit(time.sleep, delay)\n future.result()\n except RuntimeError:\n # No running event loop, use regular thread sleep\n time.sleep(delay)\n\n return ImageResponse(\n created=int(datetime.now().timestamp()),\n data=[\n ImageObject(\n b64_json=image_data,\n revised_prompt=prompt,\n )\n ],\n )\n\n\ndef _create_mock_image_generation_llm_config() -> LLMConfig:\n \"\"\"Create a mock LLMConfig for image generation.\"\"\"\n return LLMConfig(\n model_provider=\"openai\",\n model_name=\"gpt-image-1\",\n temperature=0.0,\n api_key=\"mock-api-key\",\n api_base=None,\n api_version=None,\n deployment_name=None,\n max_input_tokens=100000,\n custom_config=None,\n )\n\n\n@contextmanager\ndef use_mock_image_generation_provider() -> (\n Generator[ImageGenerationProviderController, None, None]\n):\n image_gen_provider = MockImageGenerationProvider()\n\n with (\n # Mock the image generation provider factory\n patch(\n \"onyx.tools.tool_implementations.images.image_generation_tool.get_image_generation_provider\",\n return_value=image_gen_provider,\n ),\n # Mock is_available to return True so the tool is registered\n patch(\n \"onyx.tools.tool_implementations.images.image_generation_tool.ImageGenerationTool.is_available\",\n return_value=True,\n ),\n # Mock the config lookup in tool_constructor to return a valid LLMConfig\n patch(\n \"onyx.tools.tool_constructor._get_image_generation_config\",\n return_value=_create_mock_image_generation_llm_config(),\n ),\n ):\n yield image_gen_provider\n" }, { "path": "backend/tests/external_dependency_unit/mock_llm.py", "content": "from __future__ import annotations\n\nimport abc\nimport threading\nimport time\nfrom collections.abc import Generator\nfrom collections.abc import Iterator\nfrom contextlib import contextmanager\nfrom enum import Enum\nfrom typing import Any\nfrom typing import cast\nfrom typing import Generic\nfrom typing import Literal\nfrom typing import TypeVar\nfrom unittest.mock import patch\n\nfrom pydantic import BaseModel\n\nfrom onyx.llm.interfaces import LanguageModelInput\nfrom onyx.llm.interfaces import LLM\nfrom onyx.llm.interfaces import LLMConfig\nfrom onyx.llm.interfaces import LLMUserIdentity\nfrom onyx.llm.interfaces import ReasoningEffort\nfrom onyx.llm.interfaces import ToolChoiceOptions\nfrom onyx.llm.model_response import ChatCompletionDeltaToolCall\nfrom onyx.llm.model_response import Delta\nfrom onyx.llm.model_response import FunctionCall\nfrom onyx.llm.model_response import ModelResponse\nfrom onyx.llm.model_response import ModelResponseStream\nfrom onyx.llm.model_response import StreamingChoice\n\nT = TypeVar(\"T\")\n\n\nclass LLMResponseType(str, Enum):\n REASONING = \"reasoning\"\n ANSWER = \"answer\"\n TOOL_CALL = \"tool_call\"\n\n\nclass LLMResponse(abc.ABC, BaseModel):\n type: str = \"\"\n\n @abc.abstractmethod\n def num_tokens(self) -> int:\n raise NotImplementedError\n\n\nclass LLMReasoningResponse(LLMResponse):\n type: Literal[\"reasoning\"] = LLMResponseType.REASONING.value\n reasoning_tokens: list[str]\n\n def num_tokens(self) -> int:\n return len(self.reasoning_tokens)\n\n\nclass LLMAnswerResponse(LLMResponse):\n type: Literal[\"answer\"] = LLMResponseType.ANSWER.value\n answer_tokens: list[str]\n\n def num_tokens(self) -> int:\n return len(self.answer_tokens)\n\n\nclass LLMToolCallResponse(LLMResponse):\n type: Literal[\"tool_call\"] = LLMResponseType.TOOL_CALL.value\n tool_name: str\n tool_call_id: str\n tool_call_argument_tokens: list[str]\n\n def num_tokens(self) -> int:\n return (\n len(self.tool_call_argument_tokens) + 1\n ) # +1 for the tool_call_id and tool_name\n\n\nclass StreamItem(BaseModel):\n \"\"\"Represents a single item in the mock LLM stream with its type.\"\"\"\n\n response_type: LLMResponseType\n data: Any\n\n\ndef _response_to_stream_items(response: LLMResponse) -> list[StreamItem]:\n match LLMResponseType(response.type):\n case LLMResponseType.REASONING:\n response = cast(LLMReasoningResponse, response)\n return [\n StreamItem(\n response_type=LLMResponseType.REASONING,\n data=token,\n )\n for token in response.reasoning_tokens\n ]\n case LLMResponseType.ANSWER:\n response = cast(LLMAnswerResponse, response)\n return [\n StreamItem(\n response_type=LLMResponseType.ANSWER,\n data=token,\n )\n for token in response.answer_tokens\n ]\n case LLMResponseType.TOOL_CALL:\n response = cast(LLMToolCallResponse, response)\n return [\n StreamItem(\n response_type=LLMResponseType.TOOL_CALL,\n data={\n \"tool_call_id\": response.tool_call_id,\n \"tool_name\": response.tool_name,\n \"arguments\": None,\n },\n )\n ] + [\n StreamItem(\n response_type=LLMResponseType.TOOL_CALL,\n data={\n \"tool_call_id\": None,\n \"tool_name\": None,\n \"arguments\": token,\n },\n )\n for token in response.tool_call_argument_tokens\n ]\n case _:\n raise ValueError(f\"Unknown response type: {response.type}\")\n\n\ndef create_delta_from_stream_item(item: StreamItem) -> Delta:\n response_type = item.response_type\n data = item.data\n if response_type == LLMResponseType.REASONING:\n return Delta(reasoning_content=data)\n elif response_type == LLMResponseType.ANSWER:\n return Delta(content=data)\n elif response_type == LLMResponseType.TOOL_CALL:\n # Handle grouped tool calls (list) vs single tool call (dict)\n if isinstance(data, list):\n # Multiple tool calls emitted together in the same tick\n tool_calls = []\n for tc_data in data:\n if tc_data[\"tool_call_id\"] is not None:\n tool_calls.append(\n ChatCompletionDeltaToolCall(\n id=tc_data[\"tool_call_id\"],\n index=tc_data[\"index\"],\n function=FunctionCall(\n arguments=\"\",\n name=tc_data[\"tool_name\"],\n ),\n )\n )\n else:\n tool_calls.append(\n ChatCompletionDeltaToolCall(\n index=tc_data[\"index\"],\n id=None,\n function=FunctionCall(\n arguments=tc_data[\"arguments\"],\n name=None,\n ),\n )\n )\n return Delta(tool_calls=tool_calls)\n else:\n # Single tool call (original behavior)\n # First tick has tool_call_id and tool_name, subsequent ticks have arguments\n if data[\"tool_call_id\"] is not None:\n return Delta(\n tool_calls=[\n ChatCompletionDeltaToolCall(\n id=data[\"tool_call_id\"],\n function=FunctionCall(\n name=data[\"tool_name\"],\n arguments=\"\",\n ),\n )\n ]\n )\n else:\n return Delta(\n tool_calls=[\n ChatCompletionDeltaToolCall(\n id=None,\n function=FunctionCall(\n name=None,\n arguments=data[\"arguments\"],\n ),\n )\n ]\n )\n else:\n raise ValueError(f\"Unknown response type: {response_type}\")\n\n\nclass MockLLMController(abc.ABC):\n @abc.abstractmethod\n def add_response(self, response: LLMResponse) -> None:\n \"\"\"Add a response to the current stream.\"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def add_responses_together(self, *responses: LLMResponse) -> None:\n \"\"\"Add multiple responses that should be emitted together in the same tick.\"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def forward(self, n: int) -> None:\n \"\"\"Forward the stream by n tokens.\"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def forward_till_end(self) -> None:\n \"\"\"Forward the stream until the end.\"\"\"\n raise NotImplementedError\n\n @abc.abstractmethod\n def set_max_timeout(self, timeout: float = 5.0) -> None:\n raise NotImplementedError\n\n\nclass MockLLM(LLM, MockLLMController):\n def __init__(self) -> None:\n self.stream_controller = SyncStreamController[StreamItem]()\n\n def add_response(self, response: LLMResponse) -> None:\n items = _response_to_stream_items(response)\n self.stream_controller.queue_items(items)\n\n def add_responses_together(self, *responses: LLMResponse) -> None:\n \"\"\"Add multiple responses that should be emitted together in the same tick.\n\n Currently only supports multiple tool call responses being grouped together.\n The initial tool call info (id, name) for all tool calls will be emitted\n in a single delta, followed by argument tokens for each tool call.\n \"\"\"\n tool_calls = [r for r in responses if r.type == LLMResponseType.TOOL_CALL]\n\n if len(tool_calls) != len(responses):\n raise ValueError(\n \"add_responses_together currently only supports multiple tool call responses\"\n )\n\n # Create combined first item with all tool call initial info\n combined_data = [\n {\n \"index\": idx,\n \"tool_call_id\": cast(LLMToolCallResponse, tc).tool_call_id,\n \"tool_name\": cast(LLMToolCallResponse, tc).tool_name,\n \"arguments\": None,\n }\n for idx, tc in enumerate(tool_calls)\n ]\n combined_item = StreamItem(\n response_type=LLMResponseType.TOOL_CALL,\n data=combined_data,\n )\n self.stream_controller.queue_items([combined_item])\n\n # Add argument tokens for each tool call with their index\n for idx, tc in enumerate(tool_calls):\n tc = cast(LLMToolCallResponse, tc)\n for token in tc.tool_call_argument_tokens:\n item = StreamItem(\n response_type=LLMResponseType.TOOL_CALL,\n data=[\n {\n \"index\": idx,\n \"tool_call_id\": None,\n \"tool_name\": None,\n \"arguments\": token,\n }\n ],\n )\n self.stream_controller.queue_items([item])\n\n def forward(self, n: int) -> None:\n if self.stream_controller:\n self.stream_controller.forward(n)\n else:\n raise ValueError(\"No response set\")\n\n def forward_till_end(self) -> None:\n if self.stream_controller:\n self.stream_controller.forward_till_end()\n else:\n raise ValueError(\"No response set\")\n\n def set_max_timeout(self, timeout: float = 5.0) -> None:\n self.stream_controller.timeout = timeout\n\n @property\n def config(self) -> LLMConfig:\n return LLMConfig(\n model_provider=\"mock\",\n model_name=\"mock\",\n temperature=1.0,\n max_input_tokens=1000000000,\n )\n\n def invoke(\n self,\n prompt: LanguageModelInput,\n tools: list[dict] | None = None,\n tool_choice: ToolChoiceOptions | None = None,\n structured_response_format: dict | None = None,\n timeout_override: int | None = None,\n max_tokens: int | None = None,\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO,\n user_identity: LLMUserIdentity | None = None,\n ) -> ModelResponse:\n raise NotImplementedError(\"We only care about streaming atm\")\n\n def stream(\n self,\n prompt: LanguageModelInput, # noqa: ARG002\n tools: list[dict] | None = None, # noqa: ARG002\n tool_choice: ToolChoiceOptions | None = None, # noqa: ARG002\n structured_response_format: dict | None = None, # noqa: ARG002\n timeout_override: int | None = None, # noqa: ARG002\n max_tokens: int | None = None, # noqa: ARG002\n reasoning_effort: ReasoningEffort = ReasoningEffort.AUTO, # noqa: ARG002\n user_identity: LLMUserIdentity | None = None, # noqa: ARG002\n ) -> Iterator[ModelResponseStream]:\n if not self.stream_controller:\n return\n\n for idx, item in enumerate(self.stream_controller):\n yield ModelResponseStream(\n id=\"chatcmp-123\",\n created=\"1\",\n choice=StreamingChoice(\n finish_reason=None,\n index=0, # Choice index should stay at 0 for all items in the same stream\n delta=create_delta_from_stream_item(item),\n ),\n usage=None,\n )\n\n\nclass StreamTimeoutError(Exception):\n \"\"\"Raised when the stream controller times out waiting for tokens.\"\"\"\n\n\nclass SyncStreamController(Generic[T]):\n def __init__(self, items: list[T] | None = None, timeout: float = 5.0) -> None:\n self.items = items if items is not None else []\n self.position = 0\n self.pending: list[int] = [] # The indices of the tokens that are pending\n self.timeout = timeout # Maximum time to wait for tokens before failing\n\n self._has_pending = threading.Event()\n\n def queue_items(self, new_items: list[T]) -> None:\n \"\"\"Queue additional tokens to the stream (for chaining responses like reasoning + tool calls).\"\"\"\n self.items.extend(new_items)\n\n def forward(self, n: int) -> None:\n \"\"\"Queue the next n tokens to be yielded\"\"\"\n end = min(self.position + n, len(self.items))\n self.pending.extend(range(self.position, end))\n self.position = end\n\n if self.pending:\n self._has_pending.set()\n\n def forward_till_end(self) -> None:\n self.forward(len(self.items) - self.position)\n\n @property\n def is_done(self) -> bool:\n return self.position >= len(self.items) and not self.pending\n\n def __iter__(self) -> SyncStreamController[T]:\n return self\n\n def __next__(self) -> T:\n start_time = time.monotonic()\n while not self.is_done:\n if self.pending:\n item_idx = self.pending.pop(0)\n if not self.pending:\n self._has_pending.clear()\n return self.items[item_idx]\n\n elapsed = time.monotonic() - start_time\n if elapsed >= self.timeout:\n raise StreamTimeoutError(\n f\"Stream controller timed out after {self.timeout}s waiting for tokens. \"\n f\"Position: {self.position}/{len(self.items)}, Pending: {len(self.pending)}\"\n )\n\n self._has_pending.wait(timeout=0.1)\n\n raise StopIteration\n\n\n@contextmanager\ndef use_mock_llm() -> Generator[MockLLMController, None, None]:\n mock_llm = MockLLM()\n\n with patch(\"onyx.chat.process_message.get_llm_for_persona\", return_value=mock_llm):\n yield mock_llm\n" }, { "path": "backend/tests/external_dependency_unit/mock_search_pipeline.py", "content": "from collections.abc import Callable\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.context.search.models import ChunkSearchRequest\nfrom onyx.context.search.models import InferenceChunk\nfrom onyx.context.search.models import PersonaSearchInfo\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.models import SearchSettings\nfrom onyx.db.models import User\nfrom onyx.document_index.interfaces import DocumentIndex\nfrom onyx.federated_connectors.federated_retrieval import FederatedRetrievalInfo\nfrom onyx.llm.interfaces import LLM\nfrom onyx.natural_language_processing.search_nlp_models import EmbeddingModel\nfrom onyx.tools.tool_implementations.search.search_tool import SearchTool\n\n\ndef run_functions_tuples_sequential(\n functions_with_args: list[tuple[Callable, tuple]],\n allow_failures: bool = False,\n max_workers: int | None = None, # noqa: ARG001\n timeout: float | None = None, # noqa: ARG001\n timeout_callback: Callable | None = None, # noqa: ARG001\n) -> list[Any]:\n \"\"\"\n A sequential replacement for run_functions_tuples_in_parallel.\n Useful in tests to make parallel tool calls deterministic.\n \"\"\"\n results = []\n for func, args in functions_with_args:\n try:\n results.append(func(*args))\n except Exception:\n if allow_failures:\n results.append(None)\n else:\n raise\n return results\n\n\nclass MockInternalSearchResult(BaseModel):\n document_id: str\n source_type: DocumentSource\n semantic_identifier: str\n chunk_ind: int\n\n def to_inference_chunk(self) -> InferenceChunk:\n return InferenceChunk(\n document_id=f\"{self.source_type.value.upper()}_{self.document_id}\",\n source_type=self.source_type,\n semantic_identifier=self.semantic_identifier,\n title=self.semantic_identifier,\n chunk_id=self.chunk_ind,\n blurb=\"\",\n content=\"\",\n source_links=None,\n image_file_id=None,\n section_continuation=False,\n boost=0,\n score=1.0,\n hidden=False,\n metadata={},\n match_highlights=[],\n doc_summary=\"\",\n chunk_context=\"\",\n updated_at=None,\n )\n\n def to_search_doc(self) -> SearchDoc:\n return SearchDoc(\n document_id=f\"{self.source_type.value.upper()}_{self.document_id}\",\n chunk_ind=self.chunk_ind,\n semantic_identifier=self.semantic_identifier,\n link=None,\n blurb=\"\",\n source_type=self.source_type,\n boost=0,\n hidden=False,\n metadata={},\n score=1.0,\n match_highlights=[],\n updated_at=None,\n )\n\n\nclass SearchPipelineController:\n def __init__(self) -> None:\n self.search_results: dict[str, list[MockInternalSearchResult]] = {}\n\n def add_search_results(\n self, query: str, results: list[MockInternalSearchResult]\n ) -> None:\n self.search_results[query] = results\n\n def get_search_results(self, query: str) -> list[InferenceChunk]:\n return [\n result.to_inference_chunk() for result in self.search_results.get(query, [])\n ]\n\n\n@contextmanager\ndef use_mock_search_pipeline(\n connectors: list[DocumentSource],\n) -> Generator[SearchPipelineController, None, None]:\n \"\"\"Mock the search pipeline and connector availability.\n\n Args:\n connectors: List of DocumentSource types to pretend are available.\n Pass an empty list to simulate no connectors.\n \"\"\"\n controller = SearchPipelineController()\n\n def mock_check_connectors_exist(db_session: Session) -> bool: # noqa: ARG001\n return len(connectors) > 0\n\n def mock_check_federated_connectors_exist(\n db_session: Session, # noqa: ARG001\n ) -> bool:\n # For now, federated connectors are not mocked as available\n return False\n\n def mock_check_user_files_exist(db_session: Session) -> bool: # noqa: ARG001\n # For now, user files are not mocked as available\n return False\n\n def mock_fetch_unique_document_sources(\n db_session: Session, # noqa: ARG001\n ) -> list[DocumentSource]:\n return connectors\n\n def override_search_pipeline(\n chunk_search_request: ChunkSearchRequest,\n document_index: DocumentIndex, # noqa: ARG001\n user: User | None, # noqa: ARG001\n persona_search_info: PersonaSearchInfo | None, # noqa: ARG001\n db_session: Session | None = None, # noqa: ARG001\n auto_detect_filters: bool = False, # noqa: ARG001\n llm: LLM | None = None, # noqa: ARG001\n project_id_filter: int | None = None, # noqa: ARG001\n persona_id_filter: int | None = None, # noqa: ARG001\n # Pre-fetched data (used by SearchTool to avoid DB access in parallel calls)\n acl_filters: list[str] | None = None, # noqa: ARG001\n embedding_model: EmbeddingModel | None = None, # noqa: ARG001\n prefetched_federated_retrieval_infos: ( # noqa: ARG001\n list[FederatedRetrievalInfo] | None\n ) = None,\n ) -> list[InferenceChunk]:\n return controller.get_search_results(chunk_search_request.query)\n\n # Mock the pre-fetch session and DB queries in SearchTool.run() so\n # tests don't need a fully initialised DB with search settings.\n @contextmanager\n def mock_get_session() -> Generator[MagicMock, None, None]:\n yield MagicMock(spec=Session)\n\n with (\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.search_pipeline\",\n new=override_search_pipeline,\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.check_connectors_exist\",\n new=mock_check_connectors_exist,\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.check_federated_connectors_exist\",\n new=mock_check_federated_connectors_exist,\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.semantic_query_rephrase\",\n return_value=\"\",\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.keyword_query_expansion\",\n return_value=[],\n ),\n patch(\n \"onyx.tools.tool_runner.run_functions_tuples_in_parallel\",\n new=run_functions_tuples_sequential,\n ),\n patch(\n \"onyx.db.connector.check_connectors_exist\",\n new=mock_check_connectors_exist,\n ),\n patch(\n \"onyx.db.connector.check_federated_connectors_exist\",\n new=mock_check_federated_connectors_exist,\n ),\n patch(\n \"onyx.db.connector.check_user_files_exist\",\n new=mock_check_user_files_exist,\n ),\n patch(\n \"onyx.db.connector.fetch_unique_document_sources\",\n new=mock_fetch_unique_document_sources,\n ),\n # Mock the pre-fetch phase of SearchTool.run()\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.get_session_with_current_tenant\",\n new=mock_get_session,\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.build_access_filters_for_user\",\n return_value=[],\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.get_current_search_settings\",\n return_value=MagicMock(spec=SearchSettings),\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.EmbeddingModel.from_db_model\",\n return_value=MagicMock(spec=EmbeddingModel),\n ),\n patch(\n \"onyx.tools.tool_implementations.search.search_tool.get_federated_retrieval_functions\",\n return_value=[],\n ),\n patch.object(\n SearchTool,\n \"_prefetch_slack_data\",\n return_value=(None, None, {}),\n ),\n ):\n yield controller\n" }, { "path": "backend/tests/external_dependency_unit/mock_search_provider.py", "content": "import abc\nfrom collections import defaultdict\nfrom collections.abc import Generator\nfrom collections.abc import Sequence\nfrom contextlib import contextmanager\nfrom unittest.mock import patch\n\nfrom pydantic import BaseModel\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import InternetSearchProvider\nfrom onyx.db.web_search import fetch_web_search_provider_by_name\nfrom onyx.tools.tool_implementations.web_search.models import WebSearchProvider\nfrom onyx.tools.tool_implementations.web_search.models import WebSearchResult\nfrom shared_configs.enums import WebSearchProviderType\n\n\nclass MockWebSearchResult(BaseModel):\n title: str\n link: str\n snippet: str\n\n def to_web_search_result(self) -> WebSearchResult:\n return WebSearchResult(\n title=self.title,\n link=self.link,\n snippet=self.snippet,\n author=None,\n published_date=None,\n )\n\n\nclass WebProviderController(abc.ABC):\n @abc.abstractmethod\n def add_results(self, query: str, results: list[MockWebSearchResult]) -> None:\n raise NotImplementedError\n\n\nclass MockWebProvider(WebSearchProvider, WebProviderController):\n def __init__(self) -> None:\n self._results: dict[str, list[MockWebSearchResult]] = defaultdict(list)\n\n def add_results(self, query: str, results: list[MockWebSearchResult]) -> None:\n self._results[query] = results\n\n def search(self, query: str) -> Sequence[WebSearchResult]:\n return list(\n map(lambda result: result.to_web_search_result(), self._results[query])\n )\n\n def test_connection(self) -> dict[str, str]:\n return {}\n\n\ndef add_web_provider_to_db(db_session: Session) -> None:\n # Write a provider to the database\n if fetch_web_search_provider_by_name(name=\"Test Provider 2\", db_session=db_session):\n return\n\n provider = InternetSearchProvider(\n name=\"Test Provider 2\",\n provider_type=WebSearchProviderType.EXA.value,\n api_key=\"test-api-key\",\n config={},\n is_active=True,\n )\n\n db_session.add(provider)\n db_session.commit()\n\n\ndef delete_web_provider_from_db(db_session: Session) -> None:\n provider = fetch_web_search_provider_by_name(\n name=\"Test Provider 2\", db_session=db_session\n )\n if provider is not None:\n db_session.delete(provider)\n db_session.commit()\n\n\n@contextmanager\ndef use_mock_web_provider(\n db_session: Session,\n) -> Generator[WebProviderController, None, None]:\n web_provider = MockWebProvider()\n\n # Write the tool to the database\n add_web_provider_to_db(db_session)\n\n # override the build function\n with patch(\n \"onyx.tools.tool_implementations.web_search.web_search_tool.build_search_provider_from_config\",\n return_value=web_provider,\n ):\n yield web_provider\n\n delete_web_provider_from_db(db_session)\n" }, { "path": "backend/tests/external_dependency_unit/opensearch/test_assistant_knowledge_filter.py", "content": "\"\"\"Tests for OpenSearch assistant knowledge filter construction.\n\nThese tests verify that when an assistant (persona) has knowledge attached,\nthe search filter includes the appropriate scope filters with OR logic (not AND),\nensuring documents are discoverable across knowledge types like attached documents,\nhierarchy nodes, document sets, and persona/project user files.\n\"\"\"\n\nfrom typing import Any\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.schema import ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DOCUMENT_ID_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DOCUMENT_SETS_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import PERSONAS_FIELD_NAME\nfrom onyx.document_index.opensearch.search import DocumentQuery\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\nATTACHED_DOCUMENT_ID = \"https://docs.google.com/document/d/test-doc-id\"\nHIERARCHY_NODE_ID = 42\nPERSONA_ID = 7\nKNOWLEDGE_FILTER_SCHEMA_FIELDS = {\n DOCUMENT_ID_FIELD_NAME,\n ANCESTOR_HIERARCHY_NODE_IDS_FIELD_NAME,\n DOCUMENT_SETS_FIELD_NAME,\n PERSONAS_FIELD_NAME,\n}\n\n\ndef _get_search_filters(\n source_types: list[DocumentSource],\n attached_document_ids: list[str] | None,\n hierarchy_node_ids: list[int] | None,\n persona_id_filter: int | None = None,\n document_sets: list[str] | None = None,\n) -> list[dict[str, Any]]:\n return DocumentQuery._get_search_filters(\n tenant_state=TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False),\n include_hidden=False,\n access_control_list=[\"user_email:test@example.com\"],\n source_types=source_types,\n tags=[],\n document_sets=document_sets or [],\n project_id_filter=None,\n persona_id_filter=persona_id_filter,\n time_cutoff=None,\n min_chunk_index=None,\n max_chunk_index=None,\n max_chunk_size=None,\n document_id=None,\n attached_document_ids=attached_document_ids,\n hierarchy_node_ids=hierarchy_node_ids,\n )\n\n\nclass TestAssistantKnowledgeFilter:\n \"\"\"Tests for assistant knowledge filter construction in OpenSearch queries.\"\"\"\n\n def test_persona_id_filter_added_when_knowledge_scope_exists(self) -> None:\n \"\"\"persona_id_filter should be OR'd into the knowledge scope filter\n when explicit knowledge attachments (attached_document_ids,\n hierarchy_node_ids, document_sets) are present.\"\"\"\n filter_clauses = _get_search_filters(\n source_types=[DocumentSource.FILE],\n attached_document_ids=[ATTACHED_DOCUMENT_ID],\n hierarchy_node_ids=[HIERARCHY_NODE_ID],\n persona_id_filter=PERSONA_ID,\n )\n\n knowledge_filter = None\n for clause in filter_clauses:\n if \"bool\" in clause and \"should\" in clause[\"bool\"]:\n if (\n clause[\"bool\"].get(\"minimum_should_match\") == 1\n and len(clause[\"bool\"][\"should\"]) > 0\n and (\n (\n clause[\"bool\"][\"should\"][0].get(\"term\", {}).keys()\n and list(\n clause[\"bool\"][\"should\"][0].get(\"term\", {}).keys()\n )[0]\n in KNOWLEDGE_FILTER_SCHEMA_FIELDS\n )\n or (\n clause[\"bool\"][\"should\"][0].get(\"terms\", {}).keys()\n and list(\n clause[\"bool\"][\"should\"][0].get(\"terms\", {}).keys()\n )[0]\n in KNOWLEDGE_FILTER_SCHEMA_FIELDS\n )\n )\n ):\n knowledge_filter = clause\n break\n\n assert knowledge_filter is not None, (\n \"Expected to find an assistant knowledge filter with \"\n \"'minimum_should_match: 1'\"\n )\n\n should_clauses = knowledge_filter[\"bool\"][\"should\"]\n persona_found = any(\n clause.get(\"term\", {}).get(PERSONAS_FIELD_NAME, {}).get(\"value\")\n == PERSONA_ID\n for clause in should_clauses\n )\n assert persona_found, (\n f\"Expected persona_id={PERSONA_ID} filter on {PERSONAS_FIELD_NAME} \"\n f\"in should clauses. Got: {should_clauses}\"\n )\n\n def test_persona_id_filter_alone_creates_knowledge_scope(self) -> None:\n \"\"\"persona_id_filter IS a primary knowledge scope trigger — a persona\n with user files is explicit knowledge, so it should restrict\n search on its own.\"\"\"\n filter_clauses = _get_search_filters(\n source_types=[],\n attached_document_ids=None,\n hierarchy_node_ids=None,\n persona_id_filter=PERSONA_ID,\n )\n\n knowledge_filter = None\n for clause in filter_clauses:\n if \"bool\" in clause and \"should\" in clause[\"bool\"]:\n if (\n clause[\"bool\"].get(\"minimum_should_match\") == 1\n and len(clause[\"bool\"][\"should\"]) > 0\n and (\n (\n clause[\"bool\"][\"should\"][0].get(\"term\", {}).keys()\n and list(\n clause[\"bool\"][\"should\"][0].get(\"term\", {}).keys()\n )[0]\n in KNOWLEDGE_FILTER_SCHEMA_FIELDS\n )\n or (\n clause[\"bool\"][\"should\"][0].get(\"terms\", {}).keys()\n and list(\n clause[\"bool\"][\"should\"][0].get(\"terms\", {}).keys()\n )[0]\n in KNOWLEDGE_FILTER_SCHEMA_FIELDS\n )\n )\n ):\n knowledge_filter = clause\n break\n\n assert (\n knowledge_filter is not None\n ), \"Expected persona_id_filter alone to create a knowledge scope filter\"\n persona_found = any(\n clause.get(\"term\", {}).get(PERSONAS_FIELD_NAME, {}).get(\"value\")\n == PERSONA_ID\n for clause in knowledge_filter[\"bool\"][\"should\"]\n )\n assert persona_found, (\n f\"Expected persona_id={PERSONA_ID} filter in knowledge scope. \"\n f\"Got: {knowledge_filter}\"\n )\n\n def test_knowledge_filter_with_document_sets_and_persona_filter(self) -> None:\n \"\"\"document_sets and persona_id_filter should be OR'd together in\n the knowledge scope filter.\"\"\"\n filter_clauses = _get_search_filters(\n source_types=[],\n attached_document_ids=None,\n hierarchy_node_ids=None,\n persona_id_filter=PERSONA_ID,\n document_sets=[\"engineering\"],\n )\n\n knowledge_filter = None\n for clause in filter_clauses:\n if \"bool\" in clause and \"should\" in clause[\"bool\"]:\n if (\n clause[\"bool\"].get(\"minimum_should_match\") == 1\n and len(clause[\"bool\"][\"should\"]) > 0\n and (\n (\n clause[\"bool\"][\"should\"][0].get(\"term\", {}).keys()\n and list(\n clause[\"bool\"][\"should\"][0].get(\"term\", {}).keys()\n )[0]\n in KNOWLEDGE_FILTER_SCHEMA_FIELDS\n )\n or (\n clause[\"bool\"][\"should\"][0].get(\"terms\", {}).keys()\n and list(\n clause[\"bool\"][\"should\"][0].get(\"terms\", {}).keys()\n )[0]\n in KNOWLEDGE_FILTER_SCHEMA_FIELDS\n )\n )\n ):\n knowledge_filter = clause\n break\n\n assert (\n knowledge_filter is not None\n ), \"Expected knowledge filter when document_sets is provided\"\n\n filter_str = str(knowledge_filter)\n assert (\n \"engineering\" in filter_str\n ), \"Expected document_set 'engineering' in knowledge filter\"\n assert (\n str(PERSONA_ID) in filter_str\n ), f\"Expected persona_id_filter {PERSONA_ID} in knowledge filter\"\n" }, { "path": "backend/tests/external_dependency_unit/opensearch/test_opensearch_client.py", "content": "\"\"\"External dependency unit tests for OpenSearchIndexClient.\n\nThese tests assume OpenSearch is running and test all implemented methods\nusing real schemas, pipelines, and search queries from the codebase.\n\"\"\"\n\nimport re\nimport uuid\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nimport pytest\nfrom opensearchpy import NotFoundError\n\nfrom onyx.access.models import DocumentAccess\nfrom onyx.access.utils import prefix_user_email\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.context.search.models import IndexFilters\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.client import OpenSearchIndexClient\nfrom onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout\nfrom onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE\nfrom onyx.document_index.opensearch.constants import HybridSearchNormalizationPipeline\nfrom onyx.document_index.opensearch.constants import HybridSearchSubqueryConfiguration\nfrom onyx.document_index.opensearch.opensearch_document_index import (\n generate_opensearch_filtered_access_control_list,\n)\nfrom onyx.document_index.opensearch.schema import CONTENT_FIELD_NAME\nfrom onyx.document_index.opensearch.schema import DocumentChunk\nfrom onyx.document_index.opensearch.schema import DocumentChunkWithoutVectors\nfrom onyx.document_index.opensearch.schema import DocumentSchema\nfrom onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id\nfrom onyx.document_index.opensearch.search import DocumentQuery\nfrom onyx.document_index.opensearch.search import (\n get_min_max_normalization_pipeline_name_and_config,\n)\nfrom onyx.document_index.opensearch.search import (\n get_normalization_pipeline_name_and_config,\n)\nfrom onyx.document_index.opensearch.search import (\n get_zscore_normalization_pipeline_name_and_config,\n)\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\n\ndef _patch_global_tenant_state(monkeypatch: pytest.MonkeyPatch, state: bool) -> None:\n \"\"\"Patches MULTI_TENANT wherever necessary for this test file.\n\n Args:\n monkeypatch: The test instance's monkeypatch instance, used for\n patching.\n state: The intended state of MULTI_TENANT.\n \"\"\"\n monkeypatch.setattr(\"shared_configs.configs.MULTI_TENANT\", state)\n monkeypatch.setattr(\"onyx.document_index.opensearch.schema.MULTI_TENANT\", state)\n\n\ndef _patch_hybrid_search_subquery_configuration(\n monkeypatch: pytest.MonkeyPatch, configuration: HybridSearchSubqueryConfiguration\n) -> None:\n \"\"\"\n Patches HYBRID_SEARCH_SUBQUERY_CONFIGURATION wherever necessary for this\n test file.\n\n Args:\n monkeypatch: The test instance's monkeypatch instance, used for\n patching.\n configuration: The intended state of\n HYBRID_SEARCH_SUBQUERY_CONFIGURATION.\n \"\"\"\n monkeypatch.setattr(\n \"onyx.document_index.opensearch.constants.HYBRID_SEARCH_SUBQUERY_CONFIGURATION\",\n configuration,\n )\n monkeypatch.setattr(\n \"onyx.document_index.opensearch.search.HYBRID_SEARCH_SUBQUERY_CONFIGURATION\",\n configuration,\n )\n\n\ndef _patch_hybrid_search_normalization_pipeline(\n monkeypatch: pytest.MonkeyPatch, pipeline: HybridSearchNormalizationPipeline\n) -> None:\n \"\"\"\n Patches HYBRID_SEARCH_NORMALIZATION_PIPELINE wherever necessary for this\n test file.\n \"\"\"\n monkeypatch.setattr(\n \"onyx.document_index.opensearch.constants.HYBRID_SEARCH_NORMALIZATION_PIPELINE\",\n pipeline,\n )\n monkeypatch.setattr(\n \"onyx.document_index.opensearch.search.HYBRID_SEARCH_NORMALIZATION_PIPELINE\",\n pipeline,\n )\n\n\ndef _patch_opensearch_match_highlights_disabled(\n monkeypatch: pytest.MonkeyPatch, disabled: bool\n) -> None:\n \"\"\"\n Patches OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED wherever necessary for this\n test file.\n \"\"\"\n monkeypatch.setattr(\n \"onyx.configs.app_configs.OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED\",\n disabled,\n )\n monkeypatch.setattr(\n \"onyx.document_index.opensearch.search.OPENSEARCH_MATCH_HIGHLIGHTS_DISABLED\",\n disabled,\n )\n\n\ndef _create_test_document_chunk(\n document_id: str,\n content: str,\n tenant_state: TenantState,\n chunk_index: int = 0,\n content_vector: list[float] | None = None,\n title: str | None = None,\n title_vector: list[float] | None = None,\n hidden: bool = False,\n document_access: DocumentAccess = DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=True,\n ),\n source_type: DocumentSource = DocumentSource.FILE,\n last_updated: datetime | None = None,\n) -> DocumentChunk:\n if content_vector is None:\n # Generate dummy vector - 128 dimensions for fast testing.\n content_vector = [0.1] * 128\n\n # If title is provided but no vector, generate one.\n if title is not None and title_vector is None:\n title_vector = [0.2] * 128\n\n return DocumentChunk(\n document_id=document_id,\n chunk_index=chunk_index,\n title=title,\n title_vector=title_vector,\n content=content,\n content_vector=content_vector,\n source_type=source_type.value,\n metadata_list=None,\n last_updated=last_updated,\n public=document_access.is_public,\n access_control_list=generate_opensearch_filtered_access_control_list(\n document_access\n ),\n hidden=hidden,\n global_boost=0,\n semantic_identifier=\"Test semantic identifier\",\n image_file_id=None,\n source_links=None,\n blurb=\"Test blurb\",\n doc_summary=\"Test doc summary\",\n chunk_context=\"Test chunk context\",\n document_sets=None,\n user_projects=None,\n primary_owners=None,\n secondary_owners=None,\n tenant_id=tenant_state,\n )\n\n\ndef _generate_test_vector(base_value: float = 0.1, dimension: int = 128) -> list[float]:\n \"\"\"Generates a test vector with slight variations.\n\n We round to eliminate floating point precision errors when comparing chunks\n for equality.\n \"\"\"\n return [round(base_value + (i * 0.001), 5) for i in range(dimension)]\n\n\n@pytest.fixture(scope=\"module\")\ndef opensearch_available() -> None:\n \"\"\"Verifies OpenSearch is running, skips all tests if not.\"\"\"\n if not wait_for_opensearch_with_timeout():\n pytest.fail(\"OpenSearch is not available.\")\n\n\n@pytest.fixture(scope=\"function\")\ndef test_client(\n opensearch_available: None, # noqa: ARG001\n) -> Generator[OpenSearchIndexClient, None, None]:\n \"\"\"Creates an OpenSearch client for testing with automatic cleanup.\"\"\"\n test_index_name = f\"test_index_{uuid.uuid4().hex[:8]}\"\n client = OpenSearchIndexClient(index_name=test_index_name)\n\n yield client # Test runs here.\n\n # Cleanup after test completes.\n try:\n client.delete_index()\n except Exception:\n pass\n finally:\n client.close()\n\n\n@pytest.fixture(scope=\"function\")\ndef search_pipeline(test_client: OpenSearchIndexClient) -> Generator[None, None, None]:\n \"\"\"Creates a search pipeline for testing with automatic cleanup.\"\"\"\n min_max_normalization_pipeline_name, min_max_normalization_pipeline_config = (\n get_min_max_normalization_pipeline_name_and_config()\n )\n zscore_normalization_pipeline_name, zscore_normalization_pipeline_config = (\n get_zscore_normalization_pipeline_name_and_config()\n )\n test_client.create_search_pipeline(\n pipeline_id=min_max_normalization_pipeline_name,\n pipeline_body=min_max_normalization_pipeline_config,\n )\n test_client.create_search_pipeline(\n pipeline_id=zscore_normalization_pipeline_name,\n pipeline_body=zscore_normalization_pipeline_config,\n )\n yield # Test runs here.\n try:\n test_client.delete_search_pipeline(\n pipeline_id=min_max_normalization_pipeline_name,\n )\n test_client.delete_search_pipeline(\n pipeline_id=zscore_normalization_pipeline_name,\n )\n except Exception:\n pass\n\n\nclass TestOpenSearchClient:\n \"\"\"Tests for OpenSearchIndexClient.\"\"\"\n\n def test_create_index(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests creating an index with a real schema.\"\"\"\n # Precondition.\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n\n # Under test.\n # Should not raise.\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Postcondition.\n # Verify index exists.\n assert test_client.validate_index(expected_mappings=mappings) is True\n\n def test_delete_existing_index(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests deleting an existing index returns True.\"\"\"\n # Precondition.\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test.\n # Delete should return True.\n result = test_client.delete_index()\n\n # Postcondition.\n assert result is True\n assert test_client.validate_index(expected_mappings=mappings) is False\n\n def test_delete_nonexistent_index(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests deleting a nonexistent index returns False.\"\"\"\n # Under test.\n # Don't create index, just try to delete.\n result = test_client.delete_index()\n\n # Postcondition.\n assert result is False\n\n def test_index_exists(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests checking if an index exists.\"\"\"\n # Precondition.\n # Index should not exist before creation.\n assert test_client.index_exists() is False\n\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test and postcondition.\n # Index should exist after creation.\n assert test_client.index_exists() is True\n\n def test_validate_index(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests validating an index.\"\"\"\n # Precondition.\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n\n # Under test and postcondition.\n # Should return False before creation.\n assert test_client.validate_index(expected_mappings=mappings) is False\n\n # Precondition.\n # Create index.\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test and postcondition.\n # Should return True after creation.\n assert test_client.validate_index(expected_mappings=mappings) is True\n\n def test_put_mapping_idempotent(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests put_mapping with same schema is idempotent.\"\"\"\n # Precondition.\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test.\n # Applying the same mappings again should succeed.\n test_client.put_mapping(mappings)\n\n # Postcondition.\n # Index should still be valid.\n assert test_client.validate_index(expected_mappings=mappings)\n\n def test_put_mapping_adds_new_field(\n self, test_client: OpenSearchIndexClient\n ) -> None:\n \"\"\"Tests put_mapping successfully adds new fields to existing index.\"\"\"\n # Precondition.\n # Create index with minimal schema (just required fields).\n initial_mappings = {\n \"dynamic\": \"strict\",\n \"properties\": {\n \"document_id\": {\"type\": \"keyword\"},\n \"chunk_index\": {\"type\": \"integer\"},\n \"content\": {\"type\": \"text\"},\n \"content_vector\": {\n \"type\": \"knn_vector\",\n \"dimension\": 128,\n \"method\": {\n \"name\": \"hnsw\",\n \"space_type\": \"cosinesimil\",\n \"engine\": \"lucene\",\n \"parameters\": {\"ef_construction\": 512, \"m\": 16},\n },\n },\n },\n }\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=initial_mappings, settings=settings)\n\n # Under test.\n # Add a new field using put_mapping.\n updated_mappings = {\n \"properties\": {\n \"document_id\": {\"type\": \"keyword\"},\n \"chunk_index\": {\"type\": \"integer\"},\n \"content\": {\"type\": \"text\"},\n \"content_vector\": {\n \"type\": \"knn_vector\",\n \"dimension\": 128,\n \"method\": {\n \"name\": \"hnsw\",\n \"space_type\": \"cosinesimil\",\n \"engine\": \"lucene\",\n \"parameters\": {\"ef_construction\": 512, \"m\": 16},\n },\n },\n # New field\n \"new_test_field\": {\"type\": \"keyword\"},\n },\n }\n # Should not raise.\n test_client.put_mapping(updated_mappings)\n\n # Postcondition.\n # Validate the new schema includes the new field.\n assert test_client.validate_index(expected_mappings=updated_mappings)\n\n def test_put_mapping_fails_on_type_change(\n self, test_client: OpenSearchIndexClient\n ) -> None:\n \"\"\"Tests put_mapping fails when trying to change existing field type.\"\"\"\n # Precondition.\n initial_mappings = {\n \"dynamic\": \"strict\",\n \"properties\": {\n \"document_id\": {\"type\": \"keyword\"},\n \"test_field\": {\"type\": \"keyword\"},\n },\n }\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=initial_mappings, settings=settings)\n\n # Under test and postcondition.\n # Try to change test_field type from keyword to text.\n conflicting_mappings = {\n \"properties\": {\n \"document_id\": {\"type\": \"keyword\"},\n \"test_field\": {\"type\": \"text\"}, # Changed from keyword to text\n },\n }\n # Should raise because field type cannot be changed.\n with pytest.raises(Exception, match=\"mapper|illegal_argument_exception\"):\n test_client.put_mapping(conflicting_mappings)\n\n def test_put_mapping_on_nonexistent_index(\n self, test_client: OpenSearchIndexClient\n ) -> None:\n \"\"\"Tests put_mapping on non-existent index raises an error.\"\"\"\n # Precondition.\n # Index does not exist yet.\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n\n # Under test and postcondition.\n with pytest.raises(Exception, match=\"index_not_found_exception|404\"):\n test_client.put_mapping(mappings)\n\n def test_create_duplicate_index(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests creating an index twice raises an error.\"\"\"\n # Precondition.\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=True\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n # Create once - should succeed.\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test and postcondition.\n # Create again - should raise.\n with pytest.raises(Exception, match=\"already exists\"):\n test_client.create_index(mappings=mappings, settings=settings)\n\n def test_update_settings(self, test_client: OpenSearchIndexClient) -> None:\n \"\"\"Tests that update_settings raises NotImplementedError.\"\"\"\n # Under test and postcondition.\n with pytest.raises(NotImplementedError):\n test_client.update_settings(settings={})\n\n def test_create_and_delete_search_pipeline(\n self, test_client: OpenSearchIndexClient\n ) -> None:\n \"\"\"Tests creating and deleting a search pipeline.\"\"\"\n # Precondition.\n pipeline_name, pipeline_config = get_normalization_pipeline_name_and_config()\n\n # Under test and postcondition.\n # Should not raise.\n test_client.create_search_pipeline(\n pipeline_id=pipeline_name,\n pipeline_body=pipeline_config,\n )\n\n # Under test and postcondition.\n # Should not raise.\n test_client.delete_search_pipeline(pipeline_id=pipeline_name)\n\n def test_index_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests indexing a document.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n doc = _create_test_document_chunk(\n document_id=\"test-doc-1\",\n chunk_index=0,\n content=\"Test content for indexing\",\n tenant_state=tenant_state,\n )\n\n # Under test and postcondition.\n # Should not raise.\n test_client.index_document(document=doc, tenant_state=tenant_state)\n # Should not raise if we supply update_if_exists.\n test_client.index_document(\n document=doc, tenant_state=tenant_state, update_if_exists=True\n )\n\n def test_bulk_index_documents(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests bulk indexing documents.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n docs = [\n _create_test_document_chunk(\n document_id=f\"test-doc-{i}\",\n chunk_index=i,\n content=f\"Test content for indexing {i}\",\n tenant_state=tenant_state,\n )\n for i in range(500)\n ]\n\n # Under test and postcondition.\n # Should not raise.\n test_client.bulk_index_documents(documents=docs, tenant_state=tenant_state)\n # Should not raise if we supply update_if_exists.\n test_client.bulk_index_documents(\n documents=docs, tenant_state=tenant_state, update_if_exists=True\n )\n\n def test_index_duplicate_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests indexing a duplicate document raises an error.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n doc = _create_test_document_chunk(\n document_id=\"test-doc-duplicate\",\n chunk_index=0,\n content=\"Duplicate test\",\n tenant_state=tenant_state,\n )\n\n # Index once - should succeed.\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n # Under test and postcondition.\n # Index again - should raise.\n with pytest.raises(Exception, match=\"already exists\"):\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n def test_get_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests getting a document.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n original_doc = _create_test_document_chunk(\n document_id=\"test-doc-get\",\n chunk_index=0,\n content=\"Content to retrieve\",\n tenant_state=tenant_state,\n # We only store second precision, so to make sure asserts work in\n # this test we'll deliberately lose some precision.\n last_updated=datetime.now(timezone.utc).replace(microsecond=0),\n )\n test_client.index_document(document=original_doc, tenant_state=tenant_state)\n\n # Under test.\n doc_chunk_id = get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=original_doc.document_id,\n chunk_index=original_doc.chunk_index,\n max_chunk_size=original_doc.max_chunk_size,\n )\n retrieved_doc = test_client.get_document(document_chunk_id=doc_chunk_id)\n\n # Postcondition.\n assert retrieved_doc == original_doc\n\n def test_get_nonexistent_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests getting a nonexistent document raises an error.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=False\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test and postcondition.\n with pytest.raises(Exception, match=\"404\"):\n test_client.get_document(\n document_chunk_id=\"test_source__nonexistent__512__0\"\n )\n\n def test_delete_existing_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests deleting an existing document returns True.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n doc = _create_test_document_chunk(\n document_id=\"test-doc-delete\",\n chunk_index=0,\n content=\"Content to delete\",\n tenant_state=tenant_state,\n )\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n # Under test.\n doc_chunk_id = get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=doc.document_id,\n chunk_index=doc.chunk_index,\n max_chunk_size=doc.max_chunk_size,\n )\n result = test_client.delete_document(document_chunk_id=doc_chunk_id)\n\n # Postcondition.\n assert result is True\n # Verify the document is gone.\n with pytest.raises(NotFoundError, match=\"404\"):\n test_client.get_document(document_chunk_id=doc_chunk_id)\n\n def test_delete_nonexistent_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests deleting a nonexistent document returns False.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test.\n result = test_client.delete_document(\n document_chunk_id=\"test_source__nonexistent__512__0\"\n )\n\n # Postcondition.\n assert result is False\n\n def test_delete_by_query(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests deleting documents by query.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index multiple documents.\n docs_to_delete = [\n _create_test_document_chunk(\n document_id=\"delete-me\",\n chunk_index=i,\n content=f\"Delete this {i}\",\n tenant_state=tenant_state,\n )\n for i in range(3)\n ]\n docs_to_keep = [\n _create_test_document_chunk(\n document_id=\"keep-me\",\n chunk_index=0,\n content=\"Keep this\",\n tenant_state=tenant_state,\n )\n ]\n\n for doc in docs_to_delete + docs_to_keep:\n test_client.index_document(document=doc, tenant_state=tenant_state)\n test_client.refresh_index()\n\n query_body = DocumentQuery.delete_from_document_id_query(\n document_id=\"delete-me\",\n tenant_state=tenant_state,\n )\n\n # Under test.\n num_deleted = test_client.delete_by_query(query_body=query_body)\n\n # Postcondition.\n assert num_deleted == 3\n\n # Verify deletion - the deleted documents should no longer exist.\n test_client.refresh_index()\n search_query = DocumentQuery.get_from_document_id_query(\n document_id=\"delete-me\",\n tenant_state=tenant_state,\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n min_chunk_index=None,\n max_chunk_index=None,\n get_full_document=False,\n )\n remaining_ids = test_client.search_for_document_ids(body=search_query)\n assert len(remaining_ids) == 0\n\n # Verify other documents still exist.\n keep_query = DocumentQuery.get_from_document_id_query(\n document_id=\"keep-me\",\n tenant_state=tenant_state,\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n min_chunk_index=None,\n max_chunk_index=None,\n get_full_document=False,\n )\n keep_ids = test_client.search_for_document_ids(body=keep_query)\n assert len(keep_ids) == 1\n\n def test_update_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests updating a document's properties.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Create a document to update.\n doc = _create_test_document_chunk(\n document_id=\"test-doc-update\",\n chunk_index=0,\n content=\"Original content\",\n tenant_state=tenant_state,\n hidden=False,\n )\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n # Under test.\n doc_chunk_id = get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=doc.document_id,\n chunk_index=doc.chunk_index,\n max_chunk_size=doc.max_chunk_size,\n )\n properties_to_update = {\n \"hidden\": True,\n \"global_boost\": 5,\n }\n test_client.update_document(\n document_chunk_id=doc_chunk_id,\n properties_to_update=properties_to_update,\n )\n\n # Postcondition.\n # Retrieve the document and verify updates were applied.\n updated_doc = test_client.get_document(document_chunk_id=doc_chunk_id)\n assert updated_doc.hidden is True\n assert updated_doc.global_boost == 5\n # Other properties should remain unchanged.\n assert updated_doc.document_id == doc.document_id\n assert updated_doc.content == doc.content\n assert updated_doc.public == doc.public\n\n def test_update_nonexistent_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests updating a nonexistent document raises an error.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Under test and postcondition.\n # Try to update a document that doesn't exist.\n with pytest.raises(NotFoundError, match=\"404\"):\n test_client.update_document(\n document_chunk_id=\"test_source__nonexistent__512__0\",\n properties_to_update={\"hidden\": True},\n )\n\n def test_hybrid_search_configurations_and_pipelines(\n self,\n test_client: OpenSearchIndexClient,\n search_pipeline: None, # noqa: ARG002\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"Tests all hybrid search configurations and pipelines.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n _patch_opensearch_match_highlights_disabled(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n # Index documents.\n docs = {\n \"doc-1\": _create_test_document_chunk(\n document_id=\"doc-1\",\n chunk_index=0,\n content=\"Python programming language tutorial\",\n content_vector=_generate_test_vector(0.1),\n tenant_state=tenant_state,\n ),\n \"doc-2\": _create_test_document_chunk(\n document_id=\"doc-2\",\n chunk_index=0,\n content=\"How to make cheese\",\n content_vector=_generate_test_vector(0.2),\n tenant_state=tenant_state,\n ),\n \"doc-3\": _create_test_document_chunk(\n document_id=\"doc-3\",\n chunk_index=0,\n content=\"C++ for newborns\",\n content_vector=_generate_test_vector(0.15),\n tenant_state=tenant_state,\n ),\n }\n for doc in docs.values():\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n for configuration in HybridSearchSubqueryConfiguration:\n _patch_hybrid_search_subquery_configuration(monkeypatch, configuration)\n for pipeline in HybridSearchNormalizationPipeline:\n _patch_hybrid_search_normalization_pipeline(monkeypatch, pipeline)\n pipeline_name, pipeline_config = (\n get_normalization_pipeline_name_and_config()\n )\n test_client.create_search_pipeline(\n pipeline_id=pipeline_name,\n pipeline_body=pipeline_config,\n )\n\n # Search query.\n query_text = \"Python programming\"\n query_vector = _generate_test_vector(0.12)\n search_body = DocumentQuery.get_hybrid_search_query(\n query_text=query_text,\n query_vector=query_vector,\n num_hits=5,\n tenant_state=tenant_state,\n # We're not worried about filtering here. tenant_id in this object\n # is not relevant.\n index_filters=IndexFilters(\n access_control_list=None, tenant_id=None\n ),\n include_hidden=False,\n )\n\n # Under test.\n results = test_client.search(\n body=search_body, search_pipeline_id=pipeline_name\n )\n\n # Postcondition.\n assert len(results) == len(docs)\n # Assert that all the chunks above are present.\n assert all(\n chunk.document_chunk.document_id in docs.keys() for chunk in results\n )\n # Make sure the chunk contents are preserved.\n for i, chunk in enumerate(results):\n expected = docs[chunk.document_chunk.document_id]\n assert chunk.document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(expected, k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n # Make sure score reporting seems reasonable (it should not be None\n # or 0).\n assert chunk.score\n # Make sure there is some kind of match highlight only for the first\n # result. The other results are so bad they're not expected to have\n # match highlights.\n if i == 0:\n assert chunk.match_highlights.get(CONTENT_FIELD_NAME, [])\n\n def test_search_empty_index(\n self,\n test_client: OpenSearchIndexClient,\n search_pipeline: None, # noqa: ARG002\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"Tests search on an empty index returns an empty list.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n # Note no documents were indexed.\n\n # Search query.\n query_text = \"test query\"\n query_vector = _generate_test_vector(0.5)\n search_body = DocumentQuery.get_hybrid_search_query(\n query_text=query_text,\n query_vector=query_vector,\n num_hits=5,\n tenant_state=tenant_state,\n # We're not worried about filtering here. tenant_id in this object\n # is not relevant.\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n )\n pipeline_name, _ = get_normalization_pipeline_name_and_config()\n\n # Under test.\n results = test_client.search(body=search_body, search_pipeline_id=pipeline_name)\n\n # Postcondition.\n assert len(results) == 0\n\n def test_hybrid_search_with_pipeline_and_filters(\n self,\n test_client: OpenSearchIndexClient,\n search_pipeline: None, # noqa: ARG002\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"\n Tests search filters for ACL, hidden documents, and tenant isolation.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, True)\n _patch_opensearch_match_highlights_disabled(monkeypatch, False)\n tenant_x = TenantState(tenant_id=\"tenant-x\", multitenant=True)\n tenant_y = TenantState(tenant_id=\"tenant-y\", multitenant=True)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_x.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index documents with different public/hidden, ACL, and tenant states.\n docs = {\n \"public-doc\": _create_test_document_chunk(\n document_id=\"public-doc\",\n chunk_index=0,\n content=\"Public document content\",\n hidden=False,\n tenant_state=tenant_x,\n ),\n \"hidden-doc\": _create_test_document_chunk(\n document_id=\"hidden-doc\",\n chunk_index=0,\n content=\"Hidden document content, spooky\",\n hidden=True,\n tenant_state=tenant_x,\n ),\n \"private-doc-user-a\": _create_test_document_chunk(\n document_id=\"private-doc-user-a\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 123-45-6789\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-a@example.com\", \"user-b@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n \"private-doc-user-b\": _create_test_document_chunk(\n document_id=\"private-doc-user-b\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 987-65-4321\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-b@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n \"should-not-exist-from-tenant-x-pov\": _create_test_document_chunk(\n document_id=\"should-not-exist-from-tenant-x-pov\",\n chunk_index=0,\n content=\"This is an entirely different tenant, x should never see this\",\n # Make this as permissive as possible to exercise tenant\n # isolation.\n hidden=False,\n tenant_state=tenant_y,\n ),\n }\n for doc in docs.values():\n test_client.index_document(document=doc, tenant_state=doc.tenant_id)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n query_text = \"document content\"\n query_vector = _generate_test_vector(0.6)\n search_body = DocumentQuery.get_hybrid_search_query(\n query_text=query_text,\n query_vector=query_vector,\n num_hits=5,\n tenant_state=tenant_x,\n # The user should only be able to see their private docs. tenant_id\n # in this object is not relevant.\n index_filters=IndexFilters(\n access_control_list=[\n prefix_user_email(\"user-a@example.com\"),\n prefix_user_email(\"user-c@example.com\"),\n ],\n tenant_id=None,\n ),\n include_hidden=False,\n )\n pipeline_name, _ = get_normalization_pipeline_name_and_config()\n\n # Under test.\n results = test_client.search(body=search_body, search_pipeline_id=pipeline_name)\n\n # Postcondition.\n # Should only get the public, non-hidden document, and the private\n # document for which the user has access.\n assert len(results) == 2\n # NOTE: This test is not explicitly testing for how well results are\n # ordered; we're just assuming which doc will be the first result here.\n assert results[0].document_chunk.document_id == \"public-doc\"\n # Make sure the chunk contents are preserved.\n assert results[0].document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(docs[\"public-doc\"], k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n # Make sure score reporting seems reasonable (it should not be None\n # or 0).\n assert results[0].score\n # Make sure there is some kind of match highlight.\n assert results[0].match_highlights.get(CONTENT_FIELD_NAME, [])\n # Same for the second result.\n assert results[1].document_chunk.document_id == \"private-doc-user-a\"\n assert results[1].document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(docs[\"private-doc-user-a\"], k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n assert results[1].score\n assert results[1].match_highlights.get(CONTENT_FIELD_NAME, [])\n\n def test_hybrid_search_with_pipeline_and_filters_returns_chunks_with_related_content_first(\n self,\n test_client: OpenSearchIndexClient,\n search_pipeline: None, # noqa: ARG002\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"\n Tests search with a normalization pipeline and filters returns chunks\n with related content first.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, True)\n _patch_opensearch_match_highlights_disabled(monkeypatch, False)\n tenant_x = TenantState(tenant_id=\"tenant-x\", multitenant=True)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_x.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index documents with varying relevance to the query.\n # Vectors closer to query_vector (0.1) should rank higher.\n docs = [\n _create_test_document_chunk(\n document_id=\"highly-relevant\",\n chunk_index=0,\n content=\"Artificial intelligence and machine learning transform technology\",\n content_vector=_generate_test_vector(\n 0.1\n ), # Very close to query vector.\n hidden=False,\n tenant_state=tenant_x,\n ),\n _create_test_document_chunk(\n document_id=\"somewhat-relevant\",\n chunk_index=0,\n content=\"Computer programming with various languages\",\n content_vector=_generate_test_vector(0.5), # Far from query vector.\n hidden=False,\n tenant_state=tenant_x,\n ),\n _create_test_document_chunk(\n document_id=\"not-very-relevant\",\n chunk_index=0,\n content=\"Cooking recipes for delicious meals\",\n content_vector=_generate_test_vector(\n 0.9\n ), # Very far from query vector.\n hidden=False,\n tenant_state=tenant_x,\n ),\n # These should be filtered out by public/hidden filters.\n _create_test_document_chunk(\n document_id=\"hidden-but-relevant\",\n chunk_index=0,\n content=\"Artificial intelligence research papers\",\n content_vector=_generate_test_vector(0.05), # Very close but hidden.\n hidden=True,\n tenant_state=tenant_x,\n ),\n _create_test_document_chunk(\n document_id=\"private-but-relevant\",\n chunk_index=0,\n content=\"Artificial intelligence industry analysis\",\n content_vector=_generate_test_vector(0.08), # Very close but private.\n document_access=DocumentAccess.build(\n user_emails=[],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n hidden=False,\n tenant_state=tenant_x,\n ),\n ]\n for doc in docs:\n test_client.index_document(document=doc, tenant_state=tenant_x)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n # Search query matching \"highly-relevant\" most closely.\n query_text = \"artificial intelligence\"\n query_vector = _generate_test_vector(0.1)\n search_body = DocumentQuery.get_hybrid_search_query(\n query_text=query_text,\n query_vector=query_vector,\n num_hits=5,\n tenant_state=tenant_x,\n # Explicitly pass in an empty list to enforce private doc filtering.\n index_filters=IndexFilters(access_control_list=[], tenant_id=None),\n include_hidden=False,\n )\n pipeline_name, _ = get_normalization_pipeline_name_and_config()\n\n # Under test.\n results = test_client.search(body=search_body, search_pipeline_id=pipeline_name)\n\n # Postcondition.\n # Should only get public, non-hidden documents (3 out of 5).\n assert len(results) == 3\n result_ids = [chunk.document_chunk.document_id for chunk in results]\n assert \"highly-relevant\" in result_ids\n assert \"somewhat-relevant\" in result_ids\n assert \"not-very-relevant\" in result_ids\n # Filtered out by public/hidden constraints.\n assert \"hidden-but-relevant\" not in result_ids\n assert \"private-but-relevant\" not in result_ids\n\n # Most relevant document should be first.\n assert results[0].document_chunk.document_id == \"highly-relevant\"\n\n # Make sure there is some kind of match highlight for the most relevant\n # result.\n match_highlights = results[0].match_highlights.get(CONTENT_FIELD_NAME, [])\n assert len(match_highlights) == 1\n # We expect the terms \"Artificial\" and \"intelligence\" to be matched.\n highlight_split = re.findall(r\"(.*?)\", match_highlights[0])\n assert len(highlight_split) == 2\n assert highlight_split[0] == \"Artificial\"\n assert highlight_split[1] == \"intelligence\"\n\n # Returned documents should be ordered by descending score.\n previous_score = float(\"inf\")\n for result in results:\n current_score = result.score\n assert current_score\n assert current_score < previous_score\n previous_score = current_score\n\n def test_delete_by_query_multitenant_isolation(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"\n Tests delete_by_query respects tenant boundaries in multi-tenant mode.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, True)\n tenant_x = TenantState(tenant_id=\"tenant-x\", multitenant=True)\n tenant_y = TenantState(tenant_id=\"tenant-y\", multitenant=True)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_x.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Although very unlikely in practice, let's use the same doc ID just to\n # make sure that doesn't break the index.\n tenant_x_chunks = [\n _create_test_document_chunk(\n document_id=\"doc\",\n chunk_index=i,\n content=f\"Tenant A Chunk {i}\",\n tenant_state=tenant_x,\n )\n for i in range(3)\n ]\n\n tenant_y_chunks = [\n _create_test_document_chunk(\n document_id=\"doc\",\n chunk_index=i,\n content=f\"Tenant B Chunk {i}\",\n tenant_state=tenant_y,\n )\n for i in range(2)\n ]\n\n for chunk in tenant_x_chunks + tenant_y_chunks:\n test_client.index_document(document=chunk, tenant_state=chunk.tenant_id)\n test_client.refresh_index()\n\n # Build deletion query for tenant-x only.\n query_body = DocumentQuery.delete_from_document_id_query(\n document_id=\"doc\",\n tenant_state=tenant_x,\n )\n\n # Under test.\n # Delete tenant-x chunks using delete_by_query.\n num_deleted = test_client.delete_by_query(query_body=query_body)\n\n # Postcondition.\n assert num_deleted == 3\n\n # Verify tenant-x chunks are deleted.\n test_client.refresh_index()\n verify_query_x = DocumentQuery.get_from_document_id_query(\n document_id=\"doc\",\n tenant_state=tenant_x,\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n min_chunk_index=None,\n max_chunk_index=None,\n get_full_document=False,\n )\n remaining_a_ids = test_client.search_for_document_ids(body=verify_query_x)\n assert len(remaining_a_ids) == 0\n\n # Verify tenant-y chunks still exist.\n verify_query_y = DocumentQuery.get_from_document_id_query(\n document_id=\"doc\",\n tenant_state=tenant_y,\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n min_chunk_index=None,\n max_chunk_index=None,\n get_full_document=False,\n )\n remaining_y_ids = test_client.search_for_document_ids(body=verify_query_y)\n assert len(remaining_y_ids) == 2\n expected_y_ids = {\n get_opensearch_doc_chunk_id(\n tenant_state=tenant_y,\n document_id=chunk.document_id,\n chunk_index=chunk.chunk_index,\n max_chunk_size=chunk.max_chunk_size,\n )\n for chunk in tenant_y_chunks\n }\n assert set(remaining_y_ids) == expected_y_ids\n\n def test_delete_by_query_nonexistent_document(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"\n Tests delete_by_query for non-existent document returns 0 deleted.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Don't index any documents.\n\n # Build deletion query.\n query_body = DocumentQuery.delete_from_document_id_query(\n document_id=\"nonexistent-doc\",\n tenant_state=tenant_state,\n )\n\n # Under test.\n num_deleted = test_client.delete_by_query(query_body=query_body)\n\n # Postcondition.\n assert num_deleted == 0\n\n def test_search_for_document_ids(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests search_for_document_ids method returns correct chunk IDs.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index chunks for two different documents.\n doc1_chunks = [\n _create_test_document_chunk(\n document_id=\"doc-1\",\n chunk_index=i,\n content=f\"Doc 1 Chunk {i}\",\n tenant_state=tenant_state,\n )\n for i in range(3)\n ]\n doc2_chunks = [\n _create_test_document_chunk(\n document_id=\"doc-2\",\n chunk_index=i,\n content=f\"Doc 2 Chunk {i}\",\n tenant_state=tenant_state,\n )\n for i in range(2)\n ]\n\n for chunk in doc1_chunks + doc2_chunks:\n test_client.index_document(document=chunk, tenant_state=tenant_state)\n test_client.refresh_index()\n\n # Build query for doc-1.\n query_body = DocumentQuery.get_from_document_id_query(\n document_id=\"doc-1\",\n tenant_state=tenant_state,\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n min_chunk_index=None,\n max_chunk_index=None,\n get_full_document=False,\n )\n\n # Under test.\n chunk_ids = test_client.search_for_document_ids(body=query_body)\n\n # Postcondition.\n assert len(chunk_ids) == 3\n expected_ids = {\n get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=chunk.document_id,\n chunk_index=chunk.chunk_index,\n max_chunk_size=chunk.max_chunk_size,\n )\n for chunk in doc1_chunks\n }\n assert set(chunk_ids) == expected_ids\n\n def test_search_with_no_document_access_can_retrieve_all_documents(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"\n Tests search with no document access can retrieve all documents, even\n private ones.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index documents with different public/hidden and tenant states.\n docs = {\n \"public-doc\": _create_test_document_chunk(\n document_id=\"public-doc\",\n chunk_index=0,\n content=\"Public document content\",\n hidden=False,\n tenant_state=tenant_state,\n ),\n \"hidden-doc\": _create_test_document_chunk(\n document_id=\"hidden-doc\",\n chunk_index=0,\n content=\"Hidden document content, spooky\",\n hidden=True,\n tenant_state=tenant_state,\n ),\n \"private-doc-user-a\": _create_test_document_chunk(\n document_id=\"private-doc-user-a\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 123-45-6789\",\n hidden=False,\n tenant_state=tenant_state,\n document_access=DocumentAccess.build(\n user_emails=[\"user-a@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n }\n for doc in docs.values():\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n # Build query for all documents.\n query_body = DocumentQuery.get_from_document_id_query(\n document_id=\"private-doc-user-a\",\n tenant_state=tenant_state,\n # This is the input under test, notice None for acl.\n index_filters=IndexFilters(access_control_list=None, tenant_id=None),\n include_hidden=False,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n min_chunk_index=None,\n max_chunk_index=None,\n get_full_document=False,\n )\n\n # Under test.\n chunk_ids = test_client.search_for_document_ids(body=query_body)\n\n # Postcondition.\n # Even though this doc is private, because we supplied None for acl we\n # were able to retrieve it.\n assert len(chunk_ids) == 1\n # Since this is a chunk ID, it will have the doc ID in it plus other\n # stuff we don't care about in this test.\n assert chunk_ids[0].startswith(\"private-doc-user-a\")\n\n def test_time_cutoff_filter(\n self,\n test_client: OpenSearchIndexClient,\n search_pipeline: None, # noqa: ARG002\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"Tests the time cutoff filter works.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index docs with various ages.\n one_day_ago = datetime.now(timezone.utc) - timedelta(days=1)\n one_week_ago = datetime.now(timezone.utc) - timedelta(days=7)\n six_months_ago = datetime.now(timezone.utc) - timedelta(days=180)\n one_year_ago = datetime.now(timezone.utc) - timedelta(days=365)\n docs = [\n _create_test_document_chunk(\n document_id=\"one-day-ago\",\n content=\"Good match\",\n last_updated=one_day_ago,\n tenant_state=tenant_state,\n ),\n _create_test_document_chunk(\n document_id=\"one-year-ago\",\n content=\"Good match\",\n last_updated=one_year_ago,\n tenant_state=tenant_state,\n ),\n _create_test_document_chunk(\n document_id=\"no-last-updated\",\n # Since we test for result ordering in the postconditions, let's\n # just make this content slightly less of a match with the query\n # so this test is not flaky from the ordering of the results.\n content=\"Still an ok match\",\n last_updated=None,\n tenant_state=tenant_state,\n ),\n ]\n for doc in docs:\n test_client.index_document(document=doc, tenant_state=tenant_state)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n # Build query for documents updated in the last week.\n last_week_search_body = DocumentQuery.get_hybrid_search_query(\n query_text=\"Good match\",\n query_vector=_generate_test_vector(0.1),\n num_hits=5,\n tenant_state=tenant_state,\n index_filters=IndexFilters(\n access_control_list=None, tenant_id=None, time_cutoff=one_week_ago\n ),\n include_hidden=False,\n )\n last_six_months_search_body = DocumentQuery.get_hybrid_search_query(\n query_text=\"Good match\",\n query_vector=_generate_test_vector(0.1),\n num_hits=5,\n tenant_state=tenant_state,\n index_filters=IndexFilters(\n access_control_list=None, tenant_id=None, time_cutoff=six_months_ago\n ),\n include_hidden=False,\n )\n pipeline_name, _ = get_normalization_pipeline_name_and_config()\n\n # Under test.\n last_week_results = test_client.search(\n body=last_week_search_body,\n search_pipeline_id=pipeline_name,\n )\n last_six_months_results = test_client.search(\n body=last_six_months_search_body,\n search_pipeline_id=pipeline_name,\n )\n\n # Postcondition.\n # We expect to only get one-day-ago.\n assert len(last_week_results) == 1\n assert last_week_results[0].document_chunk.document_id == \"one-day-ago\"\n # We expect to get one-day-ago and no-last-updated since six months >\n # ASSUMED_DOCUMENT_AGE_DAYS.\n assert len(last_six_months_results) == 2\n assert last_six_months_results[0].document_chunk.document_id == \"one-day-ago\"\n assert (\n last_six_months_results[1].document_chunk.document_id == \"no-last-updated\"\n )\n\n def test_random_search(\n self, test_client: OpenSearchIndexClient, monkeypatch: pytest.MonkeyPatch\n ) -> None:\n \"\"\"Tests the random search query works.\"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, False)\n tenant_state = TenantState(tenant_id=POSTGRES_DEFAULT_SCHEMA, multitenant=False)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_state.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index chunks for two different documents, one hidden one not.\n doc1_chunks = [\n _create_test_document_chunk(\n document_id=\"doc-1\",\n chunk_index=i,\n content=f\"Doc 1 Chunk {i}\",\n tenant_state=tenant_state,\n hidden=False,\n )\n for i in range(3)\n ]\n doc2_chunks = [\n _create_test_document_chunk(\n document_id=\"doc-2\",\n chunk_index=i,\n content=f\"Doc 2 Chunk {i}\",\n tenant_state=tenant_state,\n hidden=True,\n )\n for i in range(2)\n ]\n\n for chunk in doc1_chunks + doc2_chunks:\n test_client.index_document(document=chunk, tenant_state=tenant_state)\n test_client.refresh_index()\n\n # Build query.\n query_body = DocumentQuery.get_random_search_query(\n tenant_state=tenant_state,\n index_filters=IndexFilters(\n access_control_list=None, tenant_id=tenant_state.tenant_id\n ),\n num_to_retrieve=3,\n )\n\n # Under test.\n results = test_client.search(body=query_body, search_pipeline_id=None)\n\n # Postcondition.\n assert len(results) == 3\n assert set(result.document_chunk.chunk_index for result in results) == set(\n [0, 1, 2]\n )\n for result in results:\n # Note each result must be from doc 1, which is not hidden.\n expected_result = doc1_chunks[result.document_chunk.chunk_index]\n assert result.document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(expected_result, k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n\n def test_keyword_search(\n self,\n test_client: OpenSearchIndexClient,\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"\n Tests keyword search with filters for ACL, hidden documents, and tenant\n isolation.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, True)\n _patch_opensearch_match_highlights_disabled(monkeypatch, False)\n tenant_x = TenantState(tenant_id=\"tenant-x\", multitenant=True)\n tenant_y = TenantState(tenant_id=\"tenant-y\", multitenant=True)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_x.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index documents with different public/hidden, ACL, and tenant states.\n docs = {\n \"public-doc\": _create_test_document_chunk(\n document_id=\"public-doc\",\n chunk_index=0,\n content=\"Public document content\",\n hidden=False,\n tenant_state=tenant_x,\n ),\n \"hidden-doc\": _create_test_document_chunk(\n document_id=\"hidden-doc\",\n chunk_index=0,\n content=\"Hidden document content, spooky\",\n hidden=True,\n tenant_state=tenant_x,\n ),\n \"private-doc-user-a\": _create_test_document_chunk(\n document_id=\"private-doc-user-a\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 123-45-6789\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-a@example.com\", \"user-b@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n # Tests that we don't return documents that don't match keywords at\n # all, even if they match filters.\n \"private-but-not-relevant-doc-user-a\": _create_test_document_chunk(\n document_id=\"private-but-not-relevant-doc-user-a\",\n chunk_index=0,\n content=\"This text should not match the query at all\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-a@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n \"private-doc-user-b\": _create_test_document_chunk(\n document_id=\"private-doc-user-b\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 987-65-4321\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-b@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n \"should-not-exist-from-tenant-x-pov\": _create_test_document_chunk(\n document_id=\"should-not-exist-from-tenant-x-pov\",\n chunk_index=0,\n content=\"This is an entirely different tenant, x should never see this\",\n # Make this as permissive as possible to exercise tenant\n # isolation.\n hidden=False,\n tenant_state=tenant_y,\n ),\n }\n for doc in docs.values():\n test_client.index_document(document=doc, tenant_state=doc.tenant_id)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n # Should not match private-but-not-relevant-doc-user-a.\n query_text = \"document content\"\n search_body = DocumentQuery.get_keyword_search_query(\n query_text=query_text,\n num_hits=5,\n tenant_state=tenant_x,\n # The user should only be able to see their private docs. tenant_id\n # in this object is not relevant.\n index_filters=IndexFilters(\n access_control_list=[\n prefix_user_email(\"user-a@example.com\"),\n prefix_user_email(\"user-c@example.com\"),\n ],\n tenant_id=None,\n ),\n include_hidden=False,\n )\n\n # Under test.\n results = test_client.search(body=search_body, search_pipeline_id=None)\n\n # Postcondition.\n # Should only get the public, non-hidden document, and the private\n # document for which the user has access.\n assert len(results) == 2\n # This should be the highest-ranked result, as a higher percentage of\n # the content matches the query.\n assert results[0].document_chunk.document_id == \"public-doc\"\n # Make sure the chunk contents are preserved.\n assert results[0].document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(docs[\"public-doc\"], k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n # Make sure score reporting seems reasonable (it should not be None\n # or 0).\n assert results[0].score\n # Make sure there is some kind of match highlight.\n assert results[0].match_highlights.get(CONTENT_FIELD_NAME, [])\n # Same for the second result.\n assert results[1].document_chunk.document_id == \"private-doc-user-a\"\n assert results[1].document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(docs[\"private-doc-user-a\"], k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n assert results[1].score\n assert results[1].match_highlights.get(CONTENT_FIELD_NAME, [])\n assert results[1].score < results[0].score\n\n def test_semantic_search(\n self,\n test_client: OpenSearchIndexClient,\n monkeypatch: pytest.MonkeyPatch,\n ) -> None:\n \"\"\"\n Tests semantic search with filters for ACL, hidden documents, and tenant\n isolation.\n \"\"\"\n # Precondition.\n _patch_global_tenant_state(monkeypatch, True)\n tenant_x = TenantState(tenant_id=\"tenant-x\", multitenant=True)\n tenant_y = TenantState(tenant_id=\"tenant-y\", multitenant=True)\n mappings = DocumentSchema.get_document_schema(\n vector_dimension=128, multitenant=tenant_x.multitenant\n )\n settings = DocumentSchema.get_index_settings_based_on_environment()\n test_client.create_index(mappings=mappings, settings=settings)\n\n # Index documents with different public/hidden, ACL, and tenant states.\n docs = {\n \"public-doc\": _create_test_document_chunk(\n document_id=\"public-doc\",\n chunk_index=0,\n content=\"Public document content\",\n hidden=False,\n tenant_state=tenant_x,\n # Make this identical to the query vector to test that this\n # result is returned first.\n content_vector=_generate_test_vector(0.6),\n ),\n \"hidden-doc\": _create_test_document_chunk(\n document_id=\"hidden-doc\",\n chunk_index=0,\n content=\"Hidden document content, spooky\",\n hidden=True,\n tenant_state=tenant_x,\n ),\n \"private-doc-user-a\": _create_test_document_chunk(\n document_id=\"private-doc-user-a\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 123-45-6789\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-a@example.com\", \"user-b@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n # Make this different from the query vector to test that this\n # result is returned second.\n content_vector=_generate_test_vector(0.5),\n ),\n \"private-doc-user-b\": _create_test_document_chunk(\n document_id=\"private-doc-user-b\",\n chunk_index=0,\n content=\"Private document content, btw my SSN is 987-65-4321\",\n hidden=False,\n tenant_state=tenant_x,\n document_access=DocumentAccess.build(\n user_emails=[\"user-b@example.com\"],\n user_groups=[],\n external_user_emails=[],\n external_user_group_ids=[],\n is_public=False,\n ),\n ),\n \"should-not-exist-from-tenant-x-pov\": _create_test_document_chunk(\n document_id=\"should-not-exist-from-tenant-x-pov\",\n chunk_index=0,\n content=\"This is an entirely different tenant, x should never see this\",\n # Make this as permissive as possible to exercise tenant\n # isolation.\n hidden=False,\n tenant_state=tenant_y,\n ),\n }\n for doc in docs.values():\n test_client.index_document(document=doc, tenant_state=doc.tenant_id)\n\n # Refresh index to make documents searchable.\n test_client.refresh_index()\n\n query_vector = _generate_test_vector(0.6)\n search_body = DocumentQuery.get_semantic_search_query(\n query_embedding=query_vector,\n num_hits=5,\n tenant_state=tenant_x,\n # The user should only be able to see their private docs. tenant_id\n # in this object is not relevant.\n index_filters=IndexFilters(\n access_control_list=[\n prefix_user_email(\"user-a@example.com\"),\n prefix_user_email(\"user-c@example.com\"),\n ],\n tenant_id=None,\n ),\n include_hidden=False,\n )\n\n # Under test.\n results = test_client.search(body=search_body, search_pipeline_id=None)\n\n # Postcondition.\n # Should only get the public, non-hidden document, and the private\n # document for which the user has access.\n assert len(results) == 2\n # We explicitly expect this to be the highest-ranked result.\n assert results[0].document_chunk.document_id == \"public-doc\"\n # Make sure the chunk contents are preserved.\n assert results[0].document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(docs[\"public-doc\"], k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n assert results[0].score == 1.0\n # Same for the second result.\n assert results[1].document_chunk.document_id == \"private-doc-user-a\"\n assert results[1].document_chunk == DocumentChunkWithoutVectors(\n **{\n k: getattr(docs[\"private-doc-user-a\"], k)\n for k in DocumentChunkWithoutVectors.model_fields\n }\n )\n assert results[1].score\n assert 0.0 < results[1].score < 1.0\n" }, { "path": "backend/tests/external_dependency_unit/opensearch_migration/test_opensearch_migration_tasks.py", "content": "\"\"\"External dependency tests for OpenSearch migration celery tasks.\n\nThese tests require Postgres, Redis, Vespa, and OpenSearch to be running.\n\nWARNING: As with all external dependency tests, do not run them against a\ndatabase with data you care about. Your data will be destroyed.\n\"\"\"\n\nimport json\nfrom collections.abc import Generator\nfrom copy import deepcopy\nfrom datetime import datetime\nfrom typing import Any\nfrom unittest.mock import Mock\nfrom unittest.mock import patch\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.background.celery.tasks.opensearch_migration.constants import (\n GET_VESPA_CHUNKS_SLICE_COUNT,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.tasks import (\n is_continuation_token_done_for_all_slices,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.tasks import (\n migrate_chunks_from_vespa_to_opensearch_task,\n)\nfrom onyx.background.celery.tasks.opensearch_migration.transformer import (\n transform_vespa_chunks_to_opensearch_chunks,\n)\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\nfrom onyx.configs.constants import SOURCE_TYPE\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import Document\nfrom onyx.db.models import OpenSearchDocumentMigrationRecord\nfrom onyx.db.models import OpenSearchTenantMigrationRecord\nfrom onyx.db.opensearch_migration import build_sanitized_to_original_doc_id_mapping\nfrom onyx.db.search_settings import get_active_search_settings\nfrom onyx.document_index.interfaces_new import TenantState\nfrom onyx.document_index.opensearch.client import OpenSearchClient\nfrom onyx.document_index.opensearch.client import OpenSearchIndexClient\nfrom onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout\nfrom onyx.document_index.opensearch.constants import DEFAULT_MAX_CHUNK_SIZE\nfrom onyx.document_index.opensearch.schema import DocumentChunk\nfrom onyx.document_index.opensearch.schema import get_opensearch_doc_chunk_id\nfrom onyx.document_index.opensearch.search import DocumentQuery\nfrom onyx.document_index.vespa.shared_utils.utils import wait_for_vespa_with_timeout\nfrom onyx.document_index.vespa.vespa_document_index import VespaDocumentIndex\nfrom onyx.document_index.vespa_constants import ACCESS_CONTROL_LIST\nfrom onyx.document_index.vespa_constants import BLURB\nfrom onyx.document_index.vespa_constants import BOOST\nfrom onyx.document_index.vespa_constants import CHUNK_CONTEXT\nfrom onyx.document_index.vespa_constants import CHUNK_ID\nfrom onyx.document_index.vespa_constants import CONTENT\nfrom onyx.document_index.vespa_constants import DOC_SUMMARY\nfrom onyx.document_index.vespa_constants import DOC_UPDATED_AT\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID\nfrom onyx.document_index.vespa_constants import DOCUMENT_SETS\nfrom onyx.document_index.vespa_constants import EMBEDDINGS\nfrom onyx.document_index.vespa_constants import FULL_CHUNK_EMBEDDING_KEY\nfrom onyx.document_index.vespa_constants import HIDDEN\nfrom onyx.document_index.vespa_constants import IMAGE_FILE_NAME\nfrom onyx.document_index.vespa_constants import METADATA_LIST\nfrom onyx.document_index.vespa_constants import METADATA_SUFFIX\nfrom onyx.document_index.vespa_constants import PRIMARY_OWNERS\nfrom onyx.document_index.vespa_constants import SECONDARY_OWNERS\nfrom onyx.document_index.vespa_constants import SEMANTIC_IDENTIFIER\nfrom onyx.document_index.vespa_constants import SOURCE_LINKS\nfrom onyx.document_index.vespa_constants import TITLE\nfrom onyx.document_index.vespa_constants import TITLE_EMBEDDING\nfrom onyx.document_index.vespa_constants import USER_PROJECT\nfrom shared_configs.configs import MULTI_TENANT\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom tests.external_dependency_unit.full_setup import ensure_full_deployment_setup\n\n\nCHUNK_COUNT = 5\n\n\ndef _get_document_chunks_from_opensearch(\n opensearch_client: OpenSearchIndexClient,\n document_id: str,\n tenant_state: TenantState,\n) -> list[DocumentChunk]:\n opensearch_client.refresh_index()\n results: list[DocumentChunk] = []\n for i in range(CHUNK_COUNT):\n document_chunk_id: str = get_opensearch_doc_chunk_id(\n tenant_state=tenant_state,\n document_id=document_id,\n chunk_index=i,\n max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,\n )\n result = opensearch_client.get_document(document_chunk_id)\n results.append(result)\n return results\n\n\ndef _delete_document_chunks_from_opensearch(\n opensearch_client: OpenSearchIndexClient, document_id: str, current_tenant_id: str\n) -> None:\n opensearch_client.refresh_index()\n query_body = DocumentQuery.delete_from_document_id_query(\n document_id=document_id,\n tenant_state=TenantState(tenant_id=current_tenant_id, multitenant=False),\n )\n opensearch_client.delete_by_query(query_body)\n\n\ndef _generate_test_vector(dim: int) -> list[float]:\n \"\"\"Generate a deterministic test embedding vector.\"\"\"\n return [0.1 + (i * 0.001) for i in range(dim)]\n\n\ndef _insert_test_documents_with_commit(\n db_session: Session,\n document_ids: list[str],\n) -> list[Document]:\n \"\"\"Creates test Document records in Postgres.\"\"\"\n documents = [\n Document(\n id=document_id,\n semantic_id=document_id,\n chunk_count=CHUNK_COUNT,\n )\n for document_id in document_ids\n ]\n db_session.add_all(documents)\n db_session.commit()\n return documents\n\n\ndef _delete_test_documents_with_commit(\n db_session: Session,\n documents: list[Document],\n) -> None:\n \"\"\"Deletes test Document records from Postgres.\"\"\"\n for document in documents:\n db_session.delete(document)\n db_session.commit()\n\n\ndef _insert_test_migration_records_with_commit(\n db_session: Session,\n migration_records: list[OpenSearchDocumentMigrationRecord],\n) -> None:\n db_session.add_all(migration_records)\n db_session.commit()\n\n\ndef _create_raw_document_chunk(\n document_id: str,\n chunk_index: int,\n content: str,\n embedding: list[float],\n now: datetime,\n title: str | None = None,\n title_embedding: list[float] | None = None,\n) -> dict[str, Any]:\n return {\n DOCUMENT_ID: document_id,\n CHUNK_ID: chunk_index,\n CONTENT: content,\n EMBEDDINGS: {FULL_CHUNK_EMBEDDING_KEY: embedding},\n TITLE: title,\n TITLE_EMBEDDING: title_embedding,\n SOURCE_TYPE: \"test source type\",\n METADATA_LIST: [\"stuff=things\"],\n DOC_UPDATED_AT: int(now.timestamp()),\n HIDDEN: False,\n BOOST: 1,\n SEMANTIC_IDENTIFIER: \"test semantic identifier\",\n IMAGE_FILE_NAME: \"test.png\",\n SOURCE_LINKS: \"https://test.com\",\n BLURB: \"test blurb\",\n DOC_SUMMARY: \"test doc summary\",\n CHUNK_CONTEXT: \"test chunk context\",\n METADATA_SUFFIX: \"test metadata suffix\",\n DOCUMENT_SETS: {\"test document set\": 1},\n USER_PROJECT: [1],\n PRIMARY_OWNERS: [\"test primary owner\"],\n SECONDARY_OWNERS: [\"test secondary owner\"],\n ACCESS_CONTROL_LIST: {PUBLIC_DOC_PAT: 1, \"test user\": 1},\n }\n\n\ndef _assert_chunk_matches_vespa_chunk(\n opensearch_chunk: DocumentChunk,\n vespa_chunk: dict[str, Any],\n) -> None:\n assert opensearch_chunk.document_id == vespa_chunk[DOCUMENT_ID]\n assert opensearch_chunk.chunk_index == vespa_chunk[CHUNK_ID]\n assert opensearch_chunk.content == vespa_chunk[CONTENT]\n assert opensearch_chunk.content_vector == pytest.approx(\n vespa_chunk[EMBEDDINGS][FULL_CHUNK_EMBEDDING_KEY]\n )\n assert opensearch_chunk.title == vespa_chunk[TITLE]\n assert opensearch_chunk.title_vector == pytest.approx(vespa_chunk[TITLE_EMBEDDING])\n assert opensearch_chunk.source_type == vespa_chunk[SOURCE_TYPE]\n assert opensearch_chunk.metadata_list == vespa_chunk[METADATA_LIST]\n assert (\n opensearch_chunk.last_updated is not None\n and int(opensearch_chunk.last_updated.timestamp())\n == vespa_chunk[DOC_UPDATED_AT]\n )\n assert opensearch_chunk.public == vespa_chunk[ACCESS_CONTROL_LIST][PUBLIC_DOC_PAT]\n assert opensearch_chunk.access_control_list == [\n access_control\n for access_control in vespa_chunk[ACCESS_CONTROL_LIST]\n if access_control != PUBLIC_DOC_PAT\n ]\n assert opensearch_chunk.hidden == vespa_chunk[HIDDEN]\n assert opensearch_chunk.global_boost == vespa_chunk[BOOST]\n assert opensearch_chunk.semantic_identifier == vespa_chunk[SEMANTIC_IDENTIFIER]\n assert opensearch_chunk.image_file_id == vespa_chunk[IMAGE_FILE_NAME]\n assert opensearch_chunk.source_links == vespa_chunk[SOURCE_LINKS]\n assert opensearch_chunk.blurb == vespa_chunk[BLURB]\n assert opensearch_chunk.doc_summary == vespa_chunk[DOC_SUMMARY]\n assert opensearch_chunk.chunk_context == vespa_chunk[CHUNK_CONTEXT]\n assert opensearch_chunk.metadata_suffix == vespa_chunk[METADATA_SUFFIX]\n assert opensearch_chunk.document_sets == [\n doc_set for doc_set in vespa_chunk[DOCUMENT_SETS]\n ]\n assert opensearch_chunk.user_projects == vespa_chunk[USER_PROJECT]\n assert opensearch_chunk.primary_owners == vespa_chunk[PRIMARY_OWNERS]\n assert opensearch_chunk.secondary_owners == vespa_chunk[SECONDARY_OWNERS]\n\n\n@pytest.fixture(scope=\"module\")\ndef full_deployment_setup() -> Generator[None, None, None]:\n \"\"\"Optional fixture to perform full deployment-like setup on demand.\n\n Imports and calls\n tests.external_dependency_unit.startup.full_setup.ensure_full_deployment_setup\n to initialize Postgres defaults, Vespa indices, and seed initial docs.\n\n NOTE: We deliberately duplicate this logic from\n backend/tests/external_dependency_unit/conftest.py because we need to set\n opensearch_available just for this module, not the entire test session.\n\n TODO(ENG-3764)(andrei): Consolidate some of these test fixtures.\n \"\"\"\n # Patch ENABLE_OPENSEARCH_INDEXING_FOR_ONYX just for this test because we\n # don't yet want that enabled for all tests.\n # TODO(andrei): Remove this once CI enables OpenSearch for all tests.\n with (\n patch(\n \"onyx.configs.app_configs.ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\",\n True,\n ),\n patch(\"onyx.document_index.factory.ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\", True),\n ):\n ensure_full_deployment_setup(opensearch_available=True)\n yield # Test runs here.\n\n\n@pytest.fixture(scope=\"module\")\ndef db_session(\n full_deployment_setup: None, # noqa: ARG001\n) -> Generator[Session, None, None]:\n \"\"\"\n NOTE: We deliberately duplicate this logic from\n backend/tests/external_dependency_unit/conftest.py because we need a\n module-level fixture whereas the fixture in that file is function-level. I\n don't want to change it in this change to not risk inadvertently breaking\n things.\n \"\"\"\n with get_session_with_current_tenant() as session:\n yield session # Test runs here.\n\n\n@pytest.fixture(scope=\"module\")\ndef vespa_document_index(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n) -> Generator[VespaDocumentIndex, None, None]:\n \"\"\"Creates a Vespa document index for the test tenant.\"\"\"\n active = get_active_search_settings(db_session)\n yield VespaDocumentIndex(\n index_name=active.primary.index_name,\n tenant_state=TenantState(tenant_id=get_current_tenant_id(), multitenant=False),\n large_chunks_enabled=False,\n ) # Test runs here.\n\n\n@pytest.fixture(scope=\"module\")\ndef opensearch_client(\n db_session: Session,\n full_deployment_setup: None, # noqa: ARG001\n) -> Generator[OpenSearchIndexClient, None, None]:\n \"\"\"Creates an OpenSearch client for the test tenant.\"\"\"\n active = get_active_search_settings(db_session)\n yield OpenSearchIndexClient(index_name=active.primary.index_name) # Test runs here.\n\n\n@pytest.fixture(scope=\"module\")\ndef opensearch_available(\n opensearch_client: OpenSearchClient,\n) -> Generator[None, None, None]:\n \"\"\"Verifies OpenSearch is running, fails the test if not.\"\"\"\n if not wait_for_opensearch_with_timeout(client=opensearch_client):\n pytest.fail(\"OpenSearch is not available.\")\n yield # Test runs here.\n\n\n@pytest.fixture(scope=\"module\")\ndef vespa_available(\n full_deployment_setup: None, # noqa: ARG001\n) -> Generator[None, None, None]:\n \"\"\"Verifies Vespa is running, fails the test if not.\"\"\"\n # Try 90 seconds for testing in CI.\n if not wait_for_vespa_with_timeout(wait_limit=90):\n pytest.fail(\"Vespa is not available.\")\n yield # Test runs here.\n\n\n@pytest.fixture(scope=\"module\")\ndef test_embedding_dimension(db_session: Session) -> Generator[int, None, None]:\n active = get_active_search_settings(db_session)\n yield active.primary.model_dim # Test runs here.\n\n\n@pytest.fixture(scope=\"function\")\ndef patch_get_vespa_chunks_page_size() -> Generator[int, None, None]:\n test_page_size = 5\n with (\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.GET_VESPA_CHUNKS_PAGE_SIZE\",\n test_page_size,\n ),\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.constants.GET_VESPA_CHUNKS_PAGE_SIZE\",\n test_page_size,\n ),\n ):\n yield test_page_size # Test runs here.\n\n\n@pytest.fixture(scope=\"function\")\ndef test_documents(\n db_session: Session,\n vespa_document_index: VespaDocumentIndex,\n opensearch_client: OpenSearchIndexClient,\n patch_get_vespa_chunks_page_size: int,\n) -> Generator[list[Document], None, None]:\n \"\"\"\n Creates and cleans test Document records in Postgres and the document\n indices.\n \"\"\"\n # We use a large number of documents >\n # get_all_raw_document_chunks_paginated's page_size argument in the task.\n documents_to_create = patch_get_vespa_chunks_page_size * 2\n doc_ids = [f\"test_doc_{i}\" for i in range(documents_to_create)]\n documents = _insert_test_documents_with_commit(db_session, doc_ids)\n\n # NOTE: chunk_count must be passed because index_raw_chunks uses the \"new\"\n # chunk ID system (get_uuid_from_chunk_info). Without chunk_count, delete()\n # falls back to the \"old\" system (get_uuid_from_chunk_info_old) and won't\n # find/delete the chunks.\n for document in documents:\n vespa_document_index.delete(document.id, chunk_count=CHUNK_COUNT)\n\n for document in documents:\n _delete_document_chunks_from_opensearch(\n opensearch_client, document.id, get_current_tenant_id()\n )\n\n yield documents # Test runs here.\n\n # Cleanup.\n for document in documents:\n _delete_document_chunks_from_opensearch(\n opensearch_client, document.id, get_current_tenant_id()\n )\n\n for document in documents:\n vespa_document_index.delete(document.id, chunk_count=CHUNK_COUNT)\n\n _delete_test_documents_with_commit(db_session, documents)\n\n\n@pytest.fixture(scope=\"function\")\ndef clean_migration_tables(db_session: Session) -> Generator[None, None, None]:\n \"\"\"Cleans up migration-related tables before and after each test.\"\"\"\n # Clean before test.\n db_session.query(OpenSearchDocumentMigrationRecord).delete()\n db_session.query(OpenSearchTenantMigrationRecord).delete()\n db_session.commit()\n\n yield # Test runs here.\n\n # Clean after test.\n db_session.query(OpenSearchDocumentMigrationRecord).delete()\n db_session.query(OpenSearchTenantMigrationRecord).delete()\n db_session.commit()\n\n\n@pytest.fixture(scope=\"function\")\ndef enable_opensearch_indexing_for_onyx() -> Generator[None, None, None]:\n with patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\",\n True,\n ):\n yield # Test runs here.\n\n\n@pytest.fixture(scope=\"function\")\ndef disable_opensearch_indexing_for_onyx() -> Generator[None, None, None]:\n with patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.ENABLE_OPENSEARCH_INDEXING_FOR_ONYX\",\n False,\n ):\n yield # Test runs here.\n\n\nclass TestMigrateChunksFromVespaToOpenSearchTask:\n \"\"\"Tests migrate_chunks_from_vespa_to_opensearch_task.\"\"\"\n\n def test_chunk_migration_completes_successfully(\n self,\n db_session: Session,\n test_documents: list[Document],\n vespa_document_index: VespaDocumentIndex,\n opensearch_client: OpenSearchIndexClient,\n test_embedding_dimension: int,\n clean_migration_tables: None, # noqa: ARG002\n enable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Tests that all chunks are migrated from Vespa to OpenSearch.\n \"\"\"\n # Precondition.\n # Index chunks into Vespa.\n document_chunks: dict[str, list[dict[str, Any]]] = {\n document.id: [\n _create_raw_document_chunk(\n document_id=document.id,\n chunk_index=i,\n content=f\"Test content {i} for {document.id}\",\n embedding=_generate_test_vector(test_embedding_dimension),\n now=datetime.now(),\n title=f\"Test title {document.id}\",\n title_embedding=_generate_test_vector(test_embedding_dimension),\n )\n for i in range(CHUNK_COUNT)\n ]\n for document in test_documents\n }\n all_chunks: list[dict[str, Any]] = []\n for chunks in document_chunks.values():\n all_chunks.extend(chunks)\n vespa_document_index.index_raw_chunks(all_chunks)\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n\n # Under test.\n result = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n # Verify tenant migration record was updated.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated == len(all_chunks)\n # Visit is complete so continuation token should be None.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == len(all_chunks)\n\n # Verify chunks were indexed in OpenSearch.\n for document in test_documents:\n opensearch_chunks = _get_document_chunks_from_opensearch(\n opensearch_client, document.id, tenant_state\n )\n assert len(opensearch_chunks) == CHUNK_COUNT\n opensearch_chunks.sort(key=lambda x: x.chunk_index)\n for opensearch_chunk in opensearch_chunks:\n _assert_chunk_matches_vespa_chunk(\n opensearch_chunk,\n document_chunks[document.id][opensearch_chunk.chunk_index],\n )\n\n def test_chunk_migration_resumes_from_continuation_token(\n self,\n db_session: Session,\n test_documents: list[Document],\n vespa_document_index: VespaDocumentIndex,\n opensearch_client: OpenSearchIndexClient,\n test_embedding_dimension: int,\n clean_migration_tables: None, # noqa: ARG002\n enable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"Tests that chunk migration resumes from a saved continuation token.\n\n Simulates task time running out my mocking the locking behavior.\n \"\"\"\n # Precondition.\n # Index chunks into Vespa.\n document_chunks: dict[str, list[dict[str, Any]]] = {\n document.id: [\n _create_raw_document_chunk(\n document_id=document.id,\n chunk_index=i,\n content=f\"Test content {i} for {document.id}\",\n embedding=_generate_test_vector(test_embedding_dimension),\n now=datetime.now(),\n title=f\"Test title {document.id}\",\n title_embedding=_generate_test_vector(test_embedding_dimension),\n )\n for i in range(CHUNK_COUNT)\n ]\n for document in test_documents\n }\n all_chunks: list[dict[str, Any]] = []\n for chunks in document_chunks.values():\n all_chunks.extend(chunks)\n vespa_document_index.index_raw_chunks(all_chunks)\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n\n # Run the initial batch. To simulate partial progress we will mock the\n # redis lock to return True for the first invocation of .owned() and\n # False subsequently.\n mock_redis_client = Mock()\n mock_lock = Mock()\n mock_lock.owned.side_effect = [True, False, False]\n mock_lock.acquire.return_value = True\n mock_redis_client.lock.return_value = mock_lock\n with patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.get_redis_client\",\n return_value=mock_redis_client,\n ):\n result_1 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n assert result_1 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n\n # Verify partial progress was saved.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n partial_chunks_migrated = tenant_record.total_chunks_migrated\n assert partial_chunks_migrated > 0\n assert tenant_record.vespa_visit_continuation_token is not None\n # Slices are not necessarily evenly distributed across all document\n # chunks so we can't test that every token is non-None, but certainly at\n # least one must be.\n assert any(json.loads(tenant_record.vespa_visit_continuation_token).values())\n assert tenant_record.migration_completed_at is None\n assert tenant_record.approx_chunk_count_in_vespa is not None\n\n # Under test.\n # Run the remainder of the migration.\n result_2 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result_2 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n\n # Verify completion.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated > partial_chunks_migrated\n assert tenant_record.total_chunks_migrated == len(all_chunks)\n # Visit is complete so continuation token should be None.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == len(all_chunks)\n\n # Verify chunks were indexed in OpenSearch.\n for document in test_documents:\n opensearch_chunks = _get_document_chunks_from_opensearch(\n opensearch_client, document.id, tenant_state\n )\n assert len(opensearch_chunks) == CHUNK_COUNT\n opensearch_chunks.sort(key=lambda x: x.chunk_index)\n for opensearch_chunk in opensearch_chunks:\n _assert_chunk_matches_vespa_chunk(\n opensearch_chunk,\n document_chunks[document.id][opensearch_chunk.chunk_index],\n )\n\n def test_chunk_migration_visits_all_chunks_even_when_batch_size_varies(\n self,\n db_session: Session,\n test_documents: list[Document],\n vespa_document_index: VespaDocumentIndex,\n opensearch_client: OpenSearchIndexClient,\n test_embedding_dimension: int,\n clean_migration_tables: None, # noqa: ARG002\n enable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Tests that chunk migration works correctly even when the batch size\n changes halfway through a migration.\n\n Simulates task time running out my mocking the locking behavior.\n \"\"\"\n # Precondition.\n # Index chunks into Vespa.\n document_chunks: dict[str, list[dict[str, Any]]] = {\n document.id: [\n _create_raw_document_chunk(\n document_id=document.id,\n chunk_index=i,\n content=f\"Test content {i} for {document.id}\",\n embedding=_generate_test_vector(test_embedding_dimension),\n now=datetime.now(),\n title=f\"Test title {document.id}\",\n title_embedding=_generate_test_vector(test_embedding_dimension),\n )\n for i in range(CHUNK_COUNT)\n ]\n for document in test_documents\n }\n all_chunks: list[dict[str, Any]] = []\n for chunks in document_chunks.values():\n all_chunks.extend(chunks)\n vespa_document_index.index_raw_chunks(all_chunks)\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n\n # Run the initial batch. To simulate partial progress we will mock the\n # redis lock to return True for the first invocation of .owned() and\n # False subsequently.\n # NOTE: The batch size is currently set to 5 in\n # patch_get_vespa_chunks_page_size.\n mock_redis_client = Mock()\n mock_lock = Mock()\n mock_lock.owned.side_effect = [True, False, False]\n mock_lock.acquire.return_value = True\n mock_redis_client.lock.return_value = mock_lock\n with patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.get_redis_client\",\n return_value=mock_redis_client,\n ):\n result_1 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n assert result_1 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n\n # Verify partial progress was saved.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n partial_chunks_migrated = tenant_record.total_chunks_migrated\n assert partial_chunks_migrated > 0\n # page_size applies per slice, so one iteration can fetch up to\n # page_size * GET_VESPA_CHUNKS_SLICE_COUNT chunks total.\n assert partial_chunks_migrated <= 5 * GET_VESPA_CHUNKS_SLICE_COUNT\n assert tenant_record.vespa_visit_continuation_token is not None\n # Slices are not necessarily evenly distributed across all document\n # chunks so we can't test that every token is non-None, but certainly at\n # least one must be.\n assert any(json.loads(tenant_record.vespa_visit_continuation_token).values())\n assert tenant_record.migration_completed_at is None\n assert tenant_record.approx_chunk_count_in_vespa is not None\n\n # Under test.\n # Now patch the batch size to be some other number, like 2.\n mock_redis_client = Mock()\n mock_lock = Mock()\n mock_lock.owned.side_effect = [True, False, False]\n mock_lock.acquire.return_value = True\n mock_redis_client.lock.return_value = mock_lock\n with (\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.GET_VESPA_CHUNKS_PAGE_SIZE\",\n 2,\n ),\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.constants.GET_VESPA_CHUNKS_PAGE_SIZE\",\n 2,\n ),\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.get_redis_client\",\n return_value=mock_redis_client,\n ),\n ):\n result_2 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result_2 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n\n # Verify next partial progress was saved.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n new_partial_chunks_migrated = tenant_record.total_chunks_migrated\n assert new_partial_chunks_migrated > partial_chunks_migrated\n # page_size applies per slice, so one iteration can fetch up to\n # page_size * GET_VESPA_CHUNKS_SLICE_COUNT chunks total.\n assert new_partial_chunks_migrated <= (5 + 2) * GET_VESPA_CHUNKS_SLICE_COUNT\n assert tenant_record.vespa_visit_continuation_token is not None\n # Slices are not necessarily evenly distributed across all document\n # chunks so we can't test that every token is non-None, but certainly at\n # least one must be.\n assert any(json.loads(tenant_record.vespa_visit_continuation_token).values())\n assert tenant_record.migration_completed_at is None\n assert tenant_record.approx_chunk_count_in_vespa is not None\n\n # Under test.\n # Run the remainder of the migration.\n with (\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.tasks.GET_VESPA_CHUNKS_PAGE_SIZE\",\n 2,\n ),\n patch(\n \"onyx.background.celery.tasks.opensearch_migration.constants.GET_VESPA_CHUNKS_PAGE_SIZE\",\n 2,\n ),\n ):\n result_3 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result_3 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n\n # Verify completion.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated > new_partial_chunks_migrated\n assert tenant_record.total_chunks_migrated == len(all_chunks)\n # Visit is complete so continuation token should be None.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == len(all_chunks)\n\n # Verify chunks were indexed in OpenSearch.\n for document in test_documents:\n opensearch_chunks = _get_document_chunks_from_opensearch(\n opensearch_client, document.id, tenant_state\n )\n assert len(opensearch_chunks) == CHUNK_COUNT\n opensearch_chunks.sort(key=lambda x: x.chunk_index)\n for opensearch_chunk in opensearch_chunks:\n _assert_chunk_matches_vespa_chunk(\n opensearch_chunk,\n document_chunks[document.id][opensearch_chunk.chunk_index],\n )\n\n def test_chunk_migration_empty_vespa(\n self,\n db_session: Session,\n # Get this just to ensure Vespa is clean from previous test runs.\n test_documents: list[Document], # noqa: ARG002\n clean_migration_tables: None, # noqa: ARG002\n enable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Tests that chunk migration completes without error when Vespa is empty.\n \"\"\"\n # Under test.\n # No chunks in Vespa.\n result = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=get_current_tenant_id()\n )\n\n # Postcondition.\n assert result is True\n db_session.expire_all()\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated == 0\n # Visit is complete so continuation token should be marked as done for all slices.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n # Mark migration as completed even for empty Vespa.\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == 0\n\n def test_chunk_migration_updates_existing_chunks(\n self,\n db_session: Session,\n test_documents: list[Document],\n vespa_document_index: VespaDocumentIndex,\n opensearch_client: OpenSearchIndexClient,\n test_embedding_dimension: int,\n clean_migration_tables: None, # noqa: ARG002\n enable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Tests that the migration task updates existing chunks in OpenSearch if\n they already exist.\n\n Chunks existing in the index is not a failure mode as the document may\n have been dual indexed. Since dual indexing indexes into Vespa first, we\n can assume that the state of the chunk we want to migrate is the most\n up-to-date.\n \"\"\"\n # Precondition.\n # Index chunks into Vespa.\n document_chunks: dict[str, list[dict[str, Any]]] = {\n document.id: [\n _create_raw_document_chunk(\n document_id=document.id,\n chunk_index=i,\n content=f\"Test content {i} for {document.id}\",\n embedding=_generate_test_vector(test_embedding_dimension),\n now=datetime.now(),\n title=f\"Test title {document.id}\",\n title_embedding=_generate_test_vector(test_embedding_dimension),\n )\n for i in range(CHUNK_COUNT)\n ]\n for document in test_documents\n }\n all_chunks: list[dict[str, Any]] = []\n for chunks in document_chunks.values():\n all_chunks.extend(chunks)\n vespa_document_index.index_raw_chunks(all_chunks)\n # Index the first document into OpenSearch with some different content.\n document_in_opensearch = deepcopy(document_chunks[test_documents[0].id])\n for chunk in document_in_opensearch:\n chunk[\"content\"] = (\n f\"Different content {chunk[CHUNK_ID]} for {test_documents[0].id}\"\n )\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n chunks_for_document_in_opensearch, _ = (\n transform_vespa_chunks_to_opensearch_chunks(\n document_in_opensearch,\n tenant_state,\n {},\n )\n )\n opensearch_client.bulk_index_documents(\n documents=chunks_for_document_in_opensearch,\n tenant_state=tenant_state,\n update_if_exists=True,\n )\n\n # Under test.\n result = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated == len(all_chunks)\n # Visit is complete so continuation token should be None.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == len(all_chunks)\n\n # Verify chunks were indexed in OpenSearch.\n for document in test_documents:\n opensearch_chunks = _get_document_chunks_from_opensearch(\n opensearch_client, document.id, tenant_state\n )\n assert len(opensearch_chunks) == CHUNK_COUNT\n opensearch_chunks.sort(key=lambda x: x.chunk_index)\n for opensearch_chunk in opensearch_chunks:\n _assert_chunk_matches_vespa_chunk(\n opensearch_chunk,\n document_chunks[document.id][opensearch_chunk.chunk_index],\n )\n\n def test_chunk_migration_noops_when_migration_is_complete(\n self,\n db_session: Session,\n test_documents: list[Document],\n vespa_document_index: VespaDocumentIndex,\n opensearch_client: OpenSearchIndexClient,\n test_embedding_dimension: int,\n clean_migration_tables: None, # noqa: ARG002\n enable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"\n Tests that the migration task no-ops when the migration is complete.\n \"\"\"\n # Precondition.\n # Index chunks into Vespa.\n document_chunks: dict[str, list[dict[str, Any]]] = {\n document.id: [\n _create_raw_document_chunk(\n document_id=document.id,\n chunk_index=i,\n content=f\"Test content {i} for {document.id}\",\n embedding=_generate_test_vector(test_embedding_dimension),\n now=datetime.now(),\n title=f\"Test title {document.id}\",\n title_embedding=_generate_test_vector(test_embedding_dimension),\n )\n for i in range(CHUNK_COUNT)\n ]\n for document in test_documents\n }\n all_chunks: list[dict[str, Any]] = []\n for chunks in document_chunks.values():\n all_chunks.extend(chunks)\n vespa_document_index.index_raw_chunks(all_chunks)\n tenant_state = TenantState(\n tenant_id=get_current_tenant_id(), multitenant=MULTI_TENANT\n )\n\n # Under test.\n # First run.\n result_1 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result_1 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated == len(all_chunks)\n # Visit is complete so continuation token should be None.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == len(all_chunks)\n\n # Verify chunks were indexed in OpenSearch.\n for document in test_documents:\n opensearch_chunks = _get_document_chunks_from_opensearch(\n opensearch_client, document.id, tenant_state\n )\n assert len(opensearch_chunks) == CHUNK_COUNT\n opensearch_chunks.sort(key=lambda x: x.chunk_index)\n for opensearch_chunk in opensearch_chunks:\n _assert_chunk_matches_vespa_chunk(\n opensearch_chunk,\n document_chunks[document.id][opensearch_chunk.chunk_index],\n )\n\n # Under test.\n # Second run.\n result_2 = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=tenant_state.tenant_id\n )\n\n # Postcondition.\n assert result_2 is True\n # Expire the session cache to see the committed changes from the task.\n db_session.expire_all()\n # This all should be unchanged.\n tenant_record = db_session.query(OpenSearchTenantMigrationRecord).first()\n assert tenant_record is not None\n assert tenant_record.total_chunks_migrated == len(all_chunks)\n # Visit is complete so continuation token should be None.\n assert tenant_record.vespa_visit_continuation_token is not None\n assert is_continuation_token_done_for_all_slices(\n json.loads(tenant_record.vespa_visit_continuation_token)\n )\n assert tenant_record.migration_completed_at is not None\n assert tenant_record.approx_chunk_count_in_vespa == len(all_chunks)\n\n # Verify chunks were indexed in OpenSearch.\n for document in test_documents:\n opensearch_chunks = _get_document_chunks_from_opensearch(\n opensearch_client, document.id, tenant_state\n )\n assert len(opensearch_chunks) == CHUNK_COUNT\n opensearch_chunks.sort(key=lambda x: x.chunk_index)\n for opensearch_chunk in opensearch_chunks:\n _assert_chunk_matches_vespa_chunk(\n opensearch_chunk,\n document_chunks[document.id][opensearch_chunk.chunk_index],\n )\n\n def test_returns_none_when_feature_disabled(\n self,\n disable_opensearch_indexing_for_onyx: None, # noqa: ARG002\n ) -> None:\n \"\"\"Tests that task returns None when feature is disabled.\"\"\"\n # Under test.\n result = migrate_chunks_from_vespa_to_opensearch_task(\n tenant_id=get_current_tenant_id()\n )\n\n # Postcondition.\n assert result is None\n\n def test_vespa_get_chunk_count(\n self,\n vespa_document_index: VespaDocumentIndex,\n test_embedding_dimension: int,\n ) -> None:\n \"\"\"\n Tests that the VespaDocumentIndex.get_chunk_count() method returns the\n correct number of chunks.\n \"\"\"\n # Precondition.\n # Index chunks into Vespa.\n all_chunks = [\n _create_raw_document_chunk(\n document_id=\"test_doc_1\",\n chunk_index=i,\n content=f\"Test content {i} for test_doc_1\",\n embedding=_generate_test_vector(test_embedding_dimension),\n now=datetime.now(),\n title=f\"Test title {i}\",\n title_embedding=_generate_test_vector(test_embedding_dimension),\n )\n for i in range(500)\n ]\n vespa_document_index.index_raw_chunks(all_chunks)\n\n # Under test.\n chunk_count = vespa_document_index.get_chunk_count()\n\n # Postcondition.\n assert chunk_count == len(all_chunks)\n\n\nclass TestSanitizedDocIdResolution:\n \"\"\"Tests document ID resolution functions.\"\"\"\n\n def test_resolve_sanitized_document_ids_batch_normal(\n self,\n db_session: Session,\n test_documents: list[Document], # noqa: ARG002\n ) -> None:\n \"\"\"\n Tests batch resolution for normal document IDs (no sanitization needed).\n \"\"\"\n # Under test.\n result = build_sanitized_to_original_doc_id_mapping(db_session)\n\n # Postcondition.\n # Since we expect no IDs in test_documents to need sanitization, the\n # result should be empty.\n assert not result\n\n def test_resolve_sanitized_document_ids_batch_with_quotes(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"Tests batch resolution for a document ID containing single quotes.\"\"\"\n # Precondition.\n # Create a document with a single quote in its ID.\n original_id = \"test_doc_with'quote\"\n sanitized_id = \"test_doc_with_quote\"\n document = Document(\n id=original_id,\n semantic_id=original_id,\n chunk_count=1,\n )\n try:\n db_session.add(document)\n db_session.commit()\n\n # Under test.\n result = build_sanitized_to_original_doc_id_mapping(db_session)\n\n # Postcondition.\n assert len(result) == 1\n # The sanitized version should map to the original.\n assert sanitized_id in result\n assert result[sanitized_id] == original_id\n\n finally:\n _delete_test_documents_with_commit(db_session, [document])\n\n def test_raises_when_sanitized_id_matches_another_document(\n self,\n db_session: Session,\n ) -> None:\n \"\"\"\n Tests that the function raises when a sanitized ID matches another\n document's original ID.\n \"\"\"\n # Precondition.\n # Create a document with a single quote in its ID, and another document\n # with that string as its ID.\n original_id = \"test_doc_with'quote\"\n sanitized_id = \"test_doc_with_quote\"\n document_bad = Document(\n id=original_id,\n semantic_id=original_id,\n chunk_count=1,\n )\n document_fine = Document(\n id=sanitized_id,\n semantic_id=sanitized_id,\n chunk_count=1,\n )\n try:\n db_session.add(document_bad)\n db_session.add(document_fine)\n db_session.commit()\n\n # Under test.\n with pytest.raises(RuntimeError):\n build_sanitized_to_original_doc_id_mapping(db_session)\n\n finally:\n _delete_test_documents_with_commit(\n db_session, [document_bad, document_fine]\n )\n" }, { "path": "backend/tests/external_dependency_unit/permission_sync/test_doc_permission_sync_attempt.py", "content": "\"\"\"\nTest suite for DocPermissionSyncAttempt CRUD operations.\n\nTests the basic CRUD operations for document permission sync attempts,\nincluding creation, status updates, progress tracking, and querying.\n\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.permission_sync_attempt import complete_doc_permission_sync_attempt\nfrom onyx.db.permission_sync_attempt import create_doc_permission_sync_attempt\nfrom onyx.db.permission_sync_attempt import get_doc_permission_sync_attempt\nfrom onyx.db.permission_sync_attempt import (\n get_recent_doc_permission_sync_attempts_for_cc_pair,\n)\nfrom onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_failed\nfrom onyx.db.permission_sync_attempt import (\n mark_doc_permission_sync_attempt_in_progress,\n)\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef _create_test_connector_credential_pair(\n db_session: Session, source: DocumentSource = DocumentSource.GOOGLE_DRIVE\n) -> ConnectorCredentialPair:\n \"\"\"Create a test connector credential pair for testing.\"\"\"\n user = create_test_user(db_session, \"test_user\")\n\n connector = Connector(\n name=f\"Test {source.value} Connector\",\n source=source,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={},\n refresh_freq=None,\n prune_freq=None,\n indexing_start=datetime.now(timezone.utc),\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n credential_json={},\n user_id=user.id,\n admin_public=True,\n )\n db_session.add(credential)\n db_session.flush()\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.PUBLIC,\n )\n db_session.add(cc_pair)\n db_session.commit()\n\n return cc_pair\n\n\nclass TestDocPermissionSyncAttempt:\n def test_create_doc_permission_sync_attempt(self, db_session: Session) -> None:\n \"\"\"Test creating a new doc permission sync attempt.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n attempt_id = create_doc_permission_sync_attempt(\n connector_credential_pair_id=cc_pair.id,\n db_session=db_session,\n )\n\n assert attempt_id is not None\n assert isinstance(attempt_id, int)\n\n # Verify the attempt was created with correct defaults\n attempt = get_doc_permission_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.connector_credential_pair_id == cc_pair.id\n assert attempt.status == PermissionSyncStatus.NOT_STARTED\n assert attempt.total_docs_synced == 0\n assert attempt.docs_with_permission_errors == 0\n assert attempt.time_started is None\n assert attempt.time_finished is None\n assert attempt.time_created is not None\n\n def test_get_doc_permission_sync_attempt(self, db_session: Session) -> None:\n \"\"\"Test retrieving a doc permission sync attempt by ID.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Test basic retrieval\n attempt = get_doc_permission_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.id == attempt_id\n\n # Test with eager loading\n attempt_with_connector = get_doc_permission_sync_attempt(\n db_session, attempt_id, eager_load_connector=True\n )\n assert attempt_with_connector is not None\n assert attempt_with_connector.connector_credential_pair is not None\n assert attempt_with_connector.connector_credential_pair.id == cc_pair.id\n\n # Test non-existent ID\n non_existent_attempt = get_doc_permission_sync_attempt(db_session, 99999)\n assert non_existent_attempt is None\n\n def test_mark_doc_permission_sync_attempt_in_progress(\n self, db_session: Session\n ) -> None:\n \"\"\"Test marking a doc permission sync attempt as in progress.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress\n updated_attempt = mark_doc_permission_sync_attempt_in_progress(\n attempt_id, db_session\n )\n\n assert updated_attempt.status == PermissionSyncStatus.IN_PROGRESS\n assert updated_attempt.time_started is not None\n assert updated_attempt.time_finished is None\n\n # Verify it fails if already in progress\n with pytest.raises(RuntimeError, match=\"not in NOT_STARTED status\"):\n mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)\n\n def test_mark_doc_permission_sync_attempt_failed(self, db_session: Session) -> None:\n \"\"\"Test marking a doc permission sync attempt as failed.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Mark as failed with error message (should work even without starting)\n error_msg = \"Sync process crashed unexpectedly\"\n mark_doc_permission_sync_attempt_failed(\n attempt_id, db_session, error_message=error_msg\n )\n\n # Verify the status and timestamps\n attempt = get_doc_permission_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.status == PermissionSyncStatus.FAILED\n assert attempt.time_started is not None\n assert attempt.time_finished is not None\n assert attempt.error_message == error_msg\n\n def test_get_recent_doc_permission_sync_attempts_for_cc_pair(\n self, db_session: Session\n ) -> None:\n \"\"\"Test retrieving recent doc permission sync attempts for a connector credential pair.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Create multiple attempts\n attempt_ids = []\n for i in range(5):\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n attempt_ids.append(attempt_id)\n\n # Get recent attempts\n recent_attempts = get_recent_doc_permission_sync_attempts_for_cc_pair(\n cc_pair_id=cc_pair.id,\n limit=3,\n db_session=db_session,\n )\n\n assert len(recent_attempts) == 3\n\n # Verify they are ordered by time_created descending (most recent first)\n for i in range(len(recent_attempts) - 1):\n assert (\n recent_attempts[i].time_created >= recent_attempts[i + 1].time_created\n )\n\n # Verify they all belong to the correct cc_pair\n for attempt in recent_attempts:\n assert attempt.connector_credential_pair_id == cc_pair.id\n\n # Test with different cc_pair (should return empty)\n other_cc_pair = _create_test_connector_credential_pair(\n db_session, source=DocumentSource.SLACK\n )\n other_attempts = get_recent_doc_permission_sync_attempts_for_cc_pair(\n cc_pair_id=other_cc_pair.id,\n limit=10,\n db_session=db_session,\n )\n assert len(other_attempts) == 0\n\n def test_status_enum_methods(self, db_session: Session) -> None:\n \"\"\"Test the status enum helper methods.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Test NOT_STARTED status\n attempt = get_doc_permission_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert not attempt.status.is_terminal()\n assert not attempt.status.is_successful()\n\n # Test IN_PROGRESS status\n mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)\n attempt = get_doc_permission_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert not attempt.status.is_terminal()\n assert not attempt.status.is_successful()\n\n # Test SUCCESS status via complete function\n complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_docs_synced=100,\n docs_with_permission_errors=0,\n )\n attempt = get_doc_permission_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.status.is_terminal()\n assert attempt.status.is_successful()\n\n # Test FAILED status (create new attempt)\n failed_attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n mark_doc_permission_sync_attempt_failed(\n failed_attempt_id, db_session, error_message=\"Test failure\"\n )\n failed_attempt = get_doc_permission_sync_attempt(db_session, failed_attempt_id)\n assert failed_attempt is not None\n assert failed_attempt.status.is_terminal()\n assert not failed_attempt.status.is_successful()\n\n # Test COMPLETED_WITH_ERRORS status via complete function (create new attempt)\n error_attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n mark_doc_permission_sync_attempt_in_progress(error_attempt_id, db_session)\n complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=error_attempt_id,\n total_docs_synced=100,\n docs_with_permission_errors=10,\n )\n error_attempt = get_doc_permission_sync_attempt(db_session, error_attempt_id)\n assert error_attempt is not None\n assert error_attempt.status.is_terminal()\n assert (\n error_attempt.status.is_successful()\n ) # Completed with errors is still \"successful\"\n\n def test_complete_doc_permission_sync_attempt_success(\n self, db_session: Session\n ) -> None:\n \"\"\"Test completing a doc permission sync attempt without errors.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress first\n mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)\n\n # Complete without errors\n completed_attempt = complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_docs_synced=100,\n docs_with_permission_errors=0,\n )\n\n assert completed_attempt.status == PermissionSyncStatus.SUCCESS\n assert completed_attempt.total_docs_synced == 100\n assert completed_attempt.docs_with_permission_errors == 0\n assert completed_attempt.time_finished is not None\n\n def test_complete_doc_permission_sync_attempt_with_errors(\n self, db_session: Session\n ) -> None:\n \"\"\"Test completing a doc permission sync attempt with errors.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress first\n mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)\n\n # Complete with errors\n completed_attempt = complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_docs_synced=100,\n docs_with_permission_errors=15,\n )\n\n assert completed_attempt.status == PermissionSyncStatus.COMPLETED_WITH_ERRORS\n assert completed_attempt.total_docs_synced == 100\n assert completed_attempt.docs_with_permission_errors == 15\n assert completed_attempt.time_finished is not None\n\n def test_complete_doc_permission_sync_attempt_can_be_called_multiple_times(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that complete can be called multiple times if needed (accumulates correctly).\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_doc_permission_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress\n mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)\n\n # Complete once\n first_complete = complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_docs_synced=50,\n docs_with_permission_errors=5,\n )\n\n # Verify first completion\n assert first_complete.status == PermissionSyncStatus.COMPLETED_WITH_ERRORS\n assert first_complete.total_docs_synced == 50\n assert first_complete.docs_with_permission_errors == 5\n assert first_complete.time_finished is not None\n\n # Call complete again (simulating additional batch processing)\n second_complete = complete_doc_permission_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_docs_synced=50,\n docs_with_permission_errors=10,\n )\n\n # Should accumulate progress from both calls\n assert second_complete.status == PermissionSyncStatus.COMPLETED_WITH_ERRORS\n assert second_complete.total_docs_synced == 100\n assert second_complete.docs_with_permission_errors == 15\n assert second_complete.time_finished is not None\n" }, { "path": "backend/tests/external_dependency_unit/permission_sync/test_external_group_permission_sync_attempt.py", "content": "\"\"\"\nTest suite for ExternalGroupPermissionSyncAttempt CRUD operations.\n\nTests the basic CRUD operations for external group permission sync attempts,\nincluding creation, status updates, progress tracking, and querying.\nSupports both connector-specific and global group sync attempts.\n\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.models import Connector\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import Credential\nfrom onyx.db.models import ExternalGroupPermissionSyncAttempt\nfrom onyx.db.permission_sync_attempt import (\n complete_external_group_sync_attempt,\n)\nfrom onyx.db.permission_sync_attempt import (\n create_external_group_sync_attempt,\n)\nfrom onyx.db.permission_sync_attempt import (\n get_external_group_sync_attempt,\n)\nfrom onyx.db.permission_sync_attempt import (\n get_recent_external_group_sync_attempts_for_cc_pair,\n)\nfrom onyx.db.permission_sync_attempt import (\n mark_external_group_sync_attempt_failed,\n)\nfrom onyx.db.permission_sync_attempt import (\n mark_external_group_sync_attempt_in_progress,\n)\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef _create_test_connector_credential_pair(\n db_session: Session, source: DocumentSource = DocumentSource.GOOGLE_DRIVE\n) -> ConnectorCredentialPair:\n \"\"\"Create a test connector credential pair for testing.\"\"\"\n user = create_test_user(db_session, \"test_user\")\n\n connector = Connector(\n name=f\"Test {source.value} Connector\",\n source=source,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={},\n refresh_freq=None,\n prune_freq=None,\n indexing_start=datetime.now(timezone.utc),\n )\n db_session.add(connector)\n db_session.flush()\n\n credential = Credential(\n credential_json={},\n user_id=user.id,\n admin_public=True,\n )\n db_session.add(credential)\n db_session.flush()\n # Expire the credential so it reloads from DB with SensitiveValue wrapper\n db_session.expire(credential)\n\n cc_pair = ConnectorCredentialPair(\n connector_id=connector.id,\n credential_id=credential.id,\n name=\"Test CC Pair\",\n status=ConnectorCredentialPairStatus.ACTIVE,\n access_type=AccessType.PUBLIC,\n )\n db_session.add(cc_pair)\n db_session.commit()\n\n return cc_pair\n\n\ndef _cleanup_global_external_group_sync_attempts(db_session: Session) -> None:\n \"\"\"Clean up any existing global external group sync attempts from previous test runs.\"\"\"\n # Delete all global attempts (where connector_credential_pair_id is None)\n db_session.query(ExternalGroupPermissionSyncAttempt).filter(\n ExternalGroupPermissionSyncAttempt.connector_credential_pair_id.is_(None)\n ).delete()\n db_session.commit()\n\n\nclass TestExternalGroupPermissionSyncAttempt:\n def test_create_external_group_sync_attempt_with_cc_pair(\n self, db_session: Session\n ) -> None:\n \"\"\"Test creating a new external group sync attempt for a specific connector.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n attempt_id = create_external_group_sync_attempt(\n connector_credential_pair_id=cc_pair.id,\n db_session=db_session,\n )\n\n assert attempt_id is not None\n assert isinstance(attempt_id, int)\n\n # Verify the attempt was created with correct defaults\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.connector_credential_pair_id == cc_pair.id\n assert attempt.status == PermissionSyncStatus.NOT_STARTED\n assert attempt.total_users_processed == 0\n assert attempt.total_groups_processed == 0\n assert attempt.total_group_memberships_synced == 0\n assert attempt.time_started is None\n assert attempt.time_finished is None\n assert attempt.time_created is not None\n\n def test_create_global_external_group_sync_attempt(\n self, db_session: Session\n ) -> None:\n \"\"\"Test creating a new global external group sync attempt.\"\"\"\n attempt_id = create_external_group_sync_attempt(\n connector_credential_pair_id=None, # Global sync\n db_session=db_session,\n )\n\n assert attempt_id is not None\n assert isinstance(attempt_id, int)\n\n # Verify the attempt was created as global\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.connector_credential_pair_id is None\n assert attempt.status == PermissionSyncStatus.NOT_STARTED\n\n def test_get_external_group_sync_attempt(self, db_session: Session) -> None:\n \"\"\"Test retrieving an external group sync attempt by ID.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Test basic retrieval\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.id == attempt_id\n\n # Test with eager loading\n attempt_with_connector = get_external_group_sync_attempt(\n db_session, attempt_id, eager_load_connector=True\n )\n assert attempt_with_connector is not None\n assert attempt_with_connector.connector_credential_pair is not None\n assert attempt_with_connector.connector_credential_pair.id == cc_pair.id\n\n # Test non-existent ID\n non_existent_attempt = get_external_group_sync_attempt(db_session, 99999)\n assert non_existent_attempt is None\n\n def test_mark_external_group_sync_attempt_in_progress(\n self, db_session: Session\n ) -> None:\n \"\"\"Test marking an external group sync attempt as in progress.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress\n updated_attempt = mark_external_group_sync_attempt_in_progress(\n attempt_id, db_session\n )\n\n assert updated_attempt.status == PermissionSyncStatus.IN_PROGRESS\n assert updated_attempt.time_started is not None\n assert updated_attempt.time_finished is None\n\n # Verify it fails if already in progress\n with pytest.raises(RuntimeError, match=\"not in NOT_STARTED status\"):\n mark_external_group_sync_attempt_in_progress(attempt_id, db_session)\n\n def test_mark_external_group_sync_attempt_failed(self, db_session: Session) -> None:\n \"\"\"Test marking an external group sync attempt as failed.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Mark as failed with error message (should work even without starting)\n error_msg_1 = \"External group sync service unavailable\"\n mark_external_group_sync_attempt_failed(\n attempt_id, db_session, error_message=error_msg_1\n )\n\n # Verify the status and timestamps\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.status == PermissionSyncStatus.FAILED\n assert attempt.time_started is not None # Should be set if not already set\n assert attempt.time_finished is not None\n assert attempt.error_message == error_msg_1\n\n # Test with error message\n attempt_id_2 = create_external_group_sync_attempt(cc_pair.id, db_session)\n error_msg = \"Connection timeout to external service\"\n mark_external_group_sync_attempt_failed(\n attempt_id_2, db_session, error_message=error_msg\n )\n\n # Verify the error message was stored\n attempt_2 = get_external_group_sync_attempt(db_session, attempt_id_2)\n assert attempt_2 is not None\n assert attempt_2.status == PermissionSyncStatus.FAILED\n assert attempt_2.error_message == error_msg\n\n def test_get_recent_external_group_sync_attempts_for_cc_pair(\n self, db_session: Session\n ) -> None:\n \"\"\"Test retrieving recent external group sync attempts for a connector credential pair.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Create multiple attempts for the cc_pair\n attempt_ids = []\n for i in range(5):\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n attempt_ids.append(attempt_id)\n\n # Get recent attempts\n recent_attempts = get_recent_external_group_sync_attempts_for_cc_pair(\n cc_pair_id=cc_pair.id,\n limit=3,\n db_session=db_session,\n )\n\n assert len(recent_attempts) == 3\n\n # Verify they are ordered by time_created descending (most recent first)\n for i in range(len(recent_attempts) - 1):\n assert (\n recent_attempts[i].time_created >= recent_attempts[i + 1].time_created\n )\n\n # Verify they all belong to the correct cc_pair\n for attempt in recent_attempts:\n assert attempt.connector_credential_pair_id == cc_pair.id\n\n # Test with different cc_pair (should return empty)\n other_cc_pair = _create_test_connector_credential_pair(\n db_session, source=DocumentSource.SLACK\n )\n other_attempts = get_recent_external_group_sync_attempts_for_cc_pair(\n cc_pair_id=other_cc_pair.id,\n limit=10,\n db_session=db_session,\n )\n assert len(other_attempts) == 0\n\n def test_get_recent_global_external_group_sync_attempts(\n self, db_session: Session\n ) -> None:\n \"\"\"Test retrieving recent global external group sync attempts.\"\"\"\n # Clean up any existing global attempts from previous test runs\n _cleanup_global_external_group_sync_attempts(db_session)\n\n # Create a cc_pair specific attempt\n cc_pair = _create_test_connector_credential_pair(db_session)\n create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Create multiple global attempts\n global_attempt_ids = []\n for i in range(3):\n attempt_id = create_external_group_sync_attempt(None, db_session) # Global\n global_attempt_ids.append(attempt_id)\n\n # Get recent global attempts\n recent_global_attempts = get_recent_external_group_sync_attempts_for_cc_pair(\n cc_pair_id=None, # Global\n limit=5,\n db_session=db_session,\n )\n\n assert len(recent_global_attempts) == 3\n\n # Verify they are all global (cc_pair_id is None)\n for attempt in recent_global_attempts:\n assert attempt.connector_credential_pair_id is None\n\n # Verify they are ordered by time_created descending\n for i in range(len(recent_global_attempts) - 1):\n assert (\n recent_global_attempts[i].time_created\n >= recent_global_attempts[i + 1].time_created\n )\n\n def test_status_enum_methods(self, db_session: Session) -> None:\n \"\"\"Test the status enum helper methods.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Test NOT_STARTED status\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert not attempt.status.is_terminal()\n assert not attempt.status.is_successful()\n\n # Test IN_PROGRESS status\n mark_external_group_sync_attempt_in_progress(attempt_id, db_session)\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert not attempt.status.is_terminal()\n assert not attempt.status.is_successful()\n\n # Test SUCCESS status via complete function\n complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_users_processed=100,\n total_groups_processed=10,\n total_group_memberships_synced=500,\n errors_encountered=0,\n )\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.status.is_terminal()\n assert attempt.status.is_successful()\n\n # Test FAILED status (create new attempt)\n failed_attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n mark_external_group_sync_attempt_failed(\n failed_attempt_id, db_session, error_message=\"Test failure\"\n )\n failed_attempt = get_external_group_sync_attempt(db_session, failed_attempt_id)\n assert failed_attempt is not None\n assert failed_attempt.status.is_terminal()\n assert not failed_attempt.status.is_successful()\n\n # Test COMPLETED_WITH_ERRORS status via complete function (create new attempt)\n error_attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n mark_external_group_sync_attempt_in_progress(error_attempt_id, db_session)\n complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=error_attempt_id,\n total_users_processed=100,\n total_groups_processed=10,\n total_group_memberships_synced=500,\n errors_encountered=5,\n )\n error_attempt = get_external_group_sync_attempt(db_session, error_attempt_id)\n assert error_attempt is not None\n assert error_attempt.status.is_terminal()\n assert (\n error_attempt.status.is_successful()\n ) # Completed with errors is still \"successful\"\n\n def test_complete_external_group_sync_attempt_success(\n self, db_session: Session\n ) -> None:\n \"\"\"Test completing an external group sync attempt without errors.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress first\n mark_external_group_sync_attempt_in_progress(attempt_id, db_session)\n\n # Complete without errors\n completed_attempt = complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_users_processed=500,\n total_groups_processed=25,\n total_group_memberships_synced=1200,\n errors_encountered=0,\n )\n\n assert completed_attempt.status == PermissionSyncStatus.SUCCESS\n assert completed_attempt.total_users_processed == 500\n assert completed_attempt.total_groups_processed == 25\n assert completed_attempt.total_group_memberships_synced == 1200\n assert completed_attempt.time_finished is not None\n\n def test_complete_external_group_sync_attempt_with_errors(\n self, db_session: Session\n ) -> None:\n \"\"\"Test completing an external group sync attempt with errors.\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress first\n mark_external_group_sync_attempt_in_progress(attempt_id, db_session)\n\n # Complete with errors\n completed_attempt = complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_users_processed=500,\n total_groups_processed=25,\n total_group_memberships_synced=1200,\n errors_encountered=10,\n )\n\n assert completed_attempt.status == PermissionSyncStatus.COMPLETED_WITH_ERRORS\n assert completed_attempt.total_users_processed == 500\n assert completed_attempt.total_groups_processed == 25\n assert completed_attempt.total_group_memberships_synced == 1200\n assert completed_attempt.time_finished is not None\n\n def test_complete_external_group_sync_attempt_can_be_called_multiple_times(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that complete can be called multiple times if needed (accumulates correctly).\"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Mark as in progress\n mark_external_group_sync_attempt_in_progress(attempt_id, db_session)\n\n # Complete once\n first_complete = complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_users_processed=200,\n total_groups_processed=10,\n total_group_memberships_synced=600,\n errors_encountered=0,\n )\n\n # Verify first completion\n assert first_complete.status == PermissionSyncStatus.SUCCESS\n assert first_complete.total_users_processed == 200\n assert first_complete.total_groups_processed == 10\n assert first_complete.total_group_memberships_synced == 600\n assert first_complete.time_finished is not None\n\n # Call complete again (simulating additional batch processing)\n second_complete = complete_external_group_sync_attempt(\n db_session=db_session,\n attempt_id=attempt_id,\n total_users_processed=300,\n total_groups_processed=15,\n total_group_memberships_synced=600,\n errors_encountered=5,\n )\n\n # Should accumulate progress from both calls and update status\n assert second_complete.status == PermissionSyncStatus.COMPLETED_WITH_ERRORS\n assert second_complete.total_users_processed == 500\n assert second_complete.total_groups_processed == 25\n assert second_complete.total_group_memberships_synced == 1200\n assert second_complete.time_finished is not None\n\n def test_global_vs_connector_specific_attempts(self, db_session: Session) -> None:\n \"\"\"Test that global and connector-specific attempts are properly separated.\"\"\"\n # Clean up any existing global attempts from previous test runs\n _cleanup_global_external_group_sync_attempts(db_session)\n\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Create connector-specific attempts\n cc_attempt_1 = create_external_group_sync_attempt(cc_pair.id, db_session)\n cc_attempt_2 = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Create global attempts\n global_attempt_1 = create_external_group_sync_attempt(None, db_session)\n global_attempt_2 = create_external_group_sync_attempt(None, db_session)\n\n # Verify connector-specific attempts\n cc_attempts = get_recent_external_group_sync_attempts_for_cc_pair(\n cc_pair_id=cc_pair.id, limit=10, db_session=db_session\n )\n assert len(cc_attempts) == 2\n cc_attempt_ids = {attempt.id for attempt in cc_attempts}\n assert cc_attempt_ids == {cc_attempt_1, cc_attempt_2}\n\n # Verify global attempts\n global_attempts = get_recent_external_group_sync_attempts_for_cc_pair(\n cc_pair_id=None, limit=10, db_session=db_session\n )\n assert len(global_attempts) == 2\n global_attempt_ids = {attempt.id for attempt in global_attempts}\n assert global_attempt_ids == {global_attempt_1, global_attempt_2}\n\n def test_external_group_sync_attempt_not_stuck_on_early_failure(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that attempts transition to FAILED on early validation failures.\n\n This tests the bug fix where attempts could get stuck in NOT_STARTED status\n if validation checks failed after the attempt was created but before it was\n marked as IN_PROGRESS.\n \"\"\"\n cc_pair = _create_test_connector_credential_pair(db_session)\n\n # Create an attempt (simulating the start of a sync task)\n attempt_id = create_external_group_sync_attempt(cc_pair.id, db_session)\n\n # Verify it starts in NOT_STARTED\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.status == PermissionSyncStatus.NOT_STARTED\n assert attempt.error_message is None\n\n # Simulate an early validation failure (e.g., missing sync config)\n # In the actual code, this would be called by _fail_external_group_sync_attempt()\n error_msg = \"No group sync config found for source\"\n mark_external_group_sync_attempt_failed(\n attempt_id, db_session, error_message=error_msg\n )\n\n # Verify the attempt transitions to FAILED (not stuck in NOT_STARTED)\n attempt = get_external_group_sync_attempt(db_session, attempt_id)\n assert attempt is not None\n assert attempt.status == PermissionSyncStatus.FAILED\n assert attempt.error_message == error_msg\n assert attempt.time_started is not None # Should be set even on early failure\n assert attempt.time_finished is not None\n assert attempt.status.is_terminal()\n assert not attempt.status.is_successful()\n" }, { "path": "backend/tests/external_dependency_unit/search_settings/test_search_settings.py", "content": "\"\"\"Tests that search settings with contextual RAG are properly propagated\nto the indexing pipeline's LLM configuration.\"\"\"\n\nfrom unittest.mock import MagicMock\nfrom unittest.mock import patch\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.context.search.models import SavedSearchSettings\nfrom onyx.context.search.models import SearchSettingsCreationRequest\nfrom onyx.db.enums import EmbeddingPrecision\nfrom onyx.db.llm import fetch_default_contextual_rag_model\nfrom onyx.db.llm import fetch_existing_llm_provider\nfrom onyx.db.llm import update_default_contextual_model\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import IndexModelStatus\nfrom onyx.db.search_settings import create_search_settings\nfrom onyx.db.swap_index import check_and_perform_index_swap\nfrom onyx.indexing.indexing_pipeline import IndexingPipelineResult\nfrom onyx.indexing.indexing_pipeline import run_indexing_pipeline\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom onyx.server.manage.search_settings import set_new_search_settings\nfrom onyx.server.manage.search_settings import update_saved_search_settings\n\n\nTEST_CONTEXTUAL_RAG_LLM_NAME = \"test-contextual-model\"\nTEST_CONTEXTUAL_RAG_LLM_PROVIDER = \"test-contextual-provider\"\n\nUPDATED_CONTEXTUAL_RAG_LLM_NAME = \"updated-contextual-model\"\nUPDATED_CONTEXTUAL_RAG_LLM_PROVIDER = \"updated-contextual-provider\"\n\n\ndef _create_llm_provider_and_model(\n db_session: Session,\n provider_name: str,\n model_name: str,\n) -> None:\n \"\"\"Insert an LLM provider with a single visible model configuration.\"\"\"\n if fetch_existing_llm_provider(name=provider_name, db_session=db_session):\n return\n upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=provider_name,\n provider=\"openai\",\n api_key=\"test-api-key\",\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=model_name,\n is_visible=True,\n max_input_tokens=4096,\n )\n ],\n ),\n db_session=db_session,\n )\n\n\ndef _make_creation_request(\n llm_name: str = TEST_CONTEXTUAL_RAG_LLM_NAME,\n llm_provider: str = TEST_CONTEXTUAL_RAG_LLM_PROVIDER,\n enable_contextual_rag: bool = True,\n) -> SearchSettingsCreationRequest:\n return SearchSettingsCreationRequest(\n model_name=\"test-embedding-model\",\n model_dim=768,\n normalize=True,\n query_prefix=\"\",\n passage_prefix=\"\",\n provider_type=None,\n index_name=None,\n multipass_indexing=False,\n embedding_precision=EmbeddingPrecision.FLOAT,\n reduced_dimension=None,\n enable_contextual_rag=enable_contextual_rag,\n contextual_rag_llm_name=llm_name,\n contextual_rag_llm_provider=llm_provider,\n )\n\n\ndef _make_saved_search_settings(\n llm_name: str = TEST_CONTEXTUAL_RAG_LLM_NAME,\n llm_provider: str = TEST_CONTEXTUAL_RAG_LLM_PROVIDER,\n enable_contextual_rag: bool = True,\n) -> SavedSearchSettings:\n return SavedSearchSettings(\n model_name=\"test-embedding-model\",\n model_dim=768,\n normalize=True,\n query_prefix=\"\",\n passage_prefix=\"\",\n provider_type=None,\n index_name=\"test_index\",\n multipass_indexing=False,\n embedding_precision=EmbeddingPrecision.FLOAT,\n reduced_dimension=None,\n enable_contextual_rag=enable_contextual_rag,\n contextual_rag_llm_name=llm_name,\n contextual_rag_llm_provider=llm_provider,\n )\n\n\ndef _run_indexing_pipeline_with_mocks(\n mock_get_llm: MagicMock,\n mock_index_handler: MagicMock,\n db_session: Session,\n) -> None:\n \"\"\"Call run_indexing_pipeline with all heavy dependencies mocked out.\"\"\"\n mock_get_llm.return_value = MagicMock()\n mock_index_handler.return_value = IndexingPipelineResult(\n new_docs=0,\n total_docs=0,\n total_chunks=0,\n failures=[],\n )\n\n run_indexing_pipeline(\n document_batch=[],\n request_id=None,\n embedder=MagicMock(),\n document_indices=[],\n db_session=db_session,\n tenant_id=\"public\",\n adapter=MagicMock(),\n chunker=MagicMock(chunk_token_limit=512),\n )\n\n\n@pytest.fixture()\ndef baseline_search_settings(\n tenant_context: None, # noqa: ARG001\n db_session: Session,\n) -> None:\n \"\"\"Ensure a baseline PRESENT search settings row exists in the DB,\n which is required before set_new_search_settings can be called.\"\"\"\n baseline = _make_saved_search_settings(enable_contextual_rag=False)\n create_search_settings(\n search_settings=baseline,\n db_session=db_session,\n status=IndexModelStatus.PRESENT,\n )\n # Sync default contextual model to match PRESENT (clears any leftover state)\n update_default_contextual_model(\n db_session=db_session,\n enable_contextual_rag=baseline.enable_contextual_rag,\n contextual_rag_llm_provider=baseline.contextual_rag_llm_provider,\n contextual_rag_llm_name=baseline.contextual_rag_llm_name,\n )\n\n\n@patch(\"onyx.db.swap_index.get_all_document_indices\")\n@patch(\"onyx.server.manage.search_settings.get_all_document_indices\")\n@patch(\"onyx.server.manage.search_settings.get_default_document_index\")\n@patch(\"onyx.indexing.indexing_pipeline.get_llm_for_contextual_rag\")\n@patch(\"onyx.indexing.indexing_pipeline.index_doc_batch_with_handler\")\ndef test_indexing_pipeline_uses_contextual_rag_settings_from_create(\n mock_index_handler: MagicMock,\n mock_get_llm: MagicMock,\n mock_get_doc_index: MagicMock, # noqa: ARG001\n mock_get_all_doc_indices_search_settings: MagicMock, # noqa: ARG001\n mock_get_all_doc_indices: MagicMock,\n baseline_search_settings: None, # noqa: ARG001\n db_session: Session,\n) -> None:\n \"\"\"After creating FUTURE settings and swapping to PRESENT,\n fetch_default_contextual_rag_model should match the PRESENT settings\n and run_indexing_pipeline should call get_llm_for_contextual_rag.\"\"\"\n _create_llm_provider_and_model(\n db_session=db_session,\n provider_name=TEST_CONTEXTUAL_RAG_LLM_PROVIDER,\n model_name=TEST_CONTEXTUAL_RAG_LLM_NAME,\n )\n\n set_new_search_settings(\n search_settings_new=_make_creation_request(),\n _=MagicMock(),\n db_session=db_session,\n )\n\n # PRESENT still has contextual RAG disabled, so default should be None\n default_model = fetch_default_contextual_rag_model(db_session)\n assert default_model is None\n\n # Swap FUTURE → PRESENT (with 0 cc-pairs, REINDEX swaps immediately)\n mock_get_all_doc_indices.return_value = []\n old_settings = check_and_perform_index_swap(db_session)\n assert old_settings is not None, \"Swap should have occurred\"\n\n # Now PRESENT has contextual RAG enabled, default should match\n default_model = fetch_default_contextual_rag_model(db_session)\n assert default_model is not None\n assert default_model.name == TEST_CONTEXTUAL_RAG_LLM_NAME\n\n _run_indexing_pipeline_with_mocks(mock_get_llm, mock_index_handler, db_session)\n\n mock_get_llm.assert_called_once_with(\n TEST_CONTEXTUAL_RAG_LLM_NAME,\n TEST_CONTEXTUAL_RAG_LLM_PROVIDER,\n )\n\n\n@patch(\"onyx.db.swap_index.get_all_document_indices\")\n@patch(\"onyx.server.manage.search_settings.get_all_document_indices\")\n@patch(\"onyx.server.manage.search_settings.get_default_document_index\")\n@patch(\"onyx.indexing.indexing_pipeline.get_llm_for_contextual_rag\")\n@patch(\"onyx.indexing.indexing_pipeline.index_doc_batch_with_handler\")\ndef test_indexing_pipeline_uses_updated_contextual_rag_settings(\n mock_index_handler: MagicMock,\n mock_get_llm: MagicMock,\n mock_get_doc_index: MagicMock, # noqa: ARG001\n mock_get_all_doc_indices_search_settings: MagicMock, # noqa: ARG001\n mock_get_all_doc_indices: MagicMock,\n baseline_search_settings: None, # noqa: ARG001\n db_session: Session,\n) -> None:\n \"\"\"After creating FUTURE settings, swapping to PRESENT, then updating\n via update_saved_search_settings, run_indexing_pipeline should use\n the updated LLM names.\"\"\"\n _create_llm_provider_and_model(\n db_session=db_session,\n provider_name=TEST_CONTEXTUAL_RAG_LLM_PROVIDER,\n model_name=TEST_CONTEXTUAL_RAG_LLM_NAME,\n )\n _create_llm_provider_and_model(\n db_session=db_session,\n provider_name=UPDATED_CONTEXTUAL_RAG_LLM_PROVIDER,\n model_name=UPDATED_CONTEXTUAL_RAG_LLM_NAME,\n )\n\n # Create FUTURE settings with contextual RAG enabled\n set_new_search_settings(\n search_settings_new=_make_creation_request(),\n _=MagicMock(),\n db_session=db_session,\n )\n\n # PRESENT still has contextual RAG disabled, so default should be None\n default_model = fetch_default_contextual_rag_model(db_session)\n assert default_model is None\n\n # Swap FUTURE → PRESENT (with 0 cc-pairs, REINDEX swaps immediately)\n mock_get_all_doc_indices.return_value = []\n old_settings = check_and_perform_index_swap(db_session)\n assert old_settings is not None, \"Swap should have occurred\"\n\n # Now PRESENT has contextual RAG enabled, default should match\n default_model = fetch_default_contextual_rag_model(db_session)\n assert default_model is not None\n assert default_model.name == TEST_CONTEXTUAL_RAG_LLM_NAME\n\n # Update the PRESENT LLM names\n update_saved_search_settings(\n search_settings=_make_saved_search_settings(\n llm_name=UPDATED_CONTEXTUAL_RAG_LLM_NAME,\n llm_provider=UPDATED_CONTEXTUAL_RAG_LLM_PROVIDER,\n ),\n _=MagicMock(),\n db_session=db_session,\n )\n\n default_model = fetch_default_contextual_rag_model(db_session)\n assert default_model is not None\n assert default_model.name == UPDATED_CONTEXTUAL_RAG_LLM_NAME\n\n _run_indexing_pipeline_with_mocks(mock_get_llm, mock_index_handler, db_session)\n\n mock_get_llm.assert_called_once_with(\n UPDATED_CONTEXTUAL_RAG_LLM_NAME,\n UPDATED_CONTEXTUAL_RAG_LLM_PROVIDER,\n )\n\n\n@patch(\"onyx.server.manage.search_settings.get_all_document_indices\")\n@patch(\"onyx.server.manage.search_settings.get_default_document_index\")\n@patch(\"onyx.indexing.indexing_pipeline.get_llm_for_contextual_rag\")\n@patch(\"onyx.indexing.indexing_pipeline.index_doc_batch_with_handler\")\ndef test_indexing_pipeline_skips_llm_when_contextual_rag_disabled(\n mock_index_handler: MagicMock,\n mock_get_llm: MagicMock,\n mock_get_doc_index: MagicMock, # noqa: ARG001\n mock_get_all_doc_indices_search_settings: MagicMock, # noqa: ARG001\n baseline_search_settings: None, # noqa: ARG001\n db_session: Session,\n) -> None:\n \"\"\"When contextual RAG is disabled in search settings,\n get_llm_for_contextual_rag should not be called.\"\"\"\n _create_llm_provider_and_model(\n db_session=db_session,\n provider_name=TEST_CONTEXTUAL_RAG_LLM_PROVIDER,\n model_name=TEST_CONTEXTUAL_RAG_LLM_NAME,\n )\n\n set_new_search_settings(\n search_settings_new=_make_creation_request(enable_contextual_rag=False),\n _=MagicMock(),\n db_session=db_session,\n )\n\n # PRESENT has contextual RAG disabled, so default should be None\n default_model = fetch_default_contextual_rag_model(db_session)\n assert default_model is None\n\n _run_indexing_pipeline_with_mocks(mock_get_llm, mock_index_handler, db_session)\n\n mock_get_llm.assert_not_called()\n" }, { "path": "backend/tests/external_dependency_unit/slack_bot/__init__.py", "content": "" }, { "path": "backend/tests/external_dependency_unit/slack_bot/test_slack_bot_crud.py", "content": "\"\"\"Tests that SlackBot CRUD operations return properly typed SensitiveValue fields.\n\nRegression test for the bug where insert_slack_bot/update_slack_bot returned\nobjects with raw string tokens instead of SensitiveValue wrappers, causing\n'str object has no attribute get_value' errors in SlackBot.from_model().\n\"\"\"\n\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.slack_bot import insert_slack_bot\nfrom onyx.db.slack_bot import update_slack_bot\nfrom onyx.server.manage.models import SlackBot\nfrom onyx.utils.sensitive import SensitiveValue\n\n\ndef _unique(prefix: str) -> str:\n return f\"{prefix}-{uuid4().hex[:8]}\"\n\n\ndef test_insert_slack_bot_returns_sensitive_values(db_session: Session) -> None:\n bot_token = _unique(\"xoxb-insert\")\n app_token = _unique(\"xapp-insert\")\n user_token = _unique(\"xoxp-insert\")\n\n slack_bot = insert_slack_bot(\n db_session=db_session,\n name=_unique(\"test-bot-insert\"),\n enabled=True,\n bot_token=bot_token,\n app_token=app_token,\n user_token=user_token,\n )\n\n assert isinstance(slack_bot.bot_token, SensitiveValue)\n assert isinstance(slack_bot.app_token, SensitiveValue)\n assert isinstance(slack_bot.user_token, SensitiveValue)\n\n assert slack_bot.bot_token.get_value(apply_mask=False) == bot_token\n assert slack_bot.app_token.get_value(apply_mask=False) == app_token\n assert slack_bot.user_token.get_value(apply_mask=False) == user_token\n\n # Verify from_model works without error\n pydantic_bot = SlackBot.from_model(slack_bot)\n assert pydantic_bot.bot_token # masked, but not empty\n assert pydantic_bot.app_token\n\n\ndef test_update_slack_bot_returns_sensitive_values(db_session: Session) -> None:\n slack_bot = insert_slack_bot(\n db_session=db_session,\n name=_unique(\"test-bot-update\"),\n enabled=True,\n bot_token=_unique(\"xoxb-update\"),\n app_token=_unique(\"xapp-update\"),\n )\n\n new_bot_token = _unique(\"xoxb-update-new\")\n new_app_token = _unique(\"xapp-update-new\")\n new_user_token = _unique(\"xoxp-update-new\")\n\n updated = update_slack_bot(\n db_session=db_session,\n slack_bot_id=slack_bot.id,\n name=_unique(\"test-bot-updated\"),\n enabled=False,\n bot_token=new_bot_token,\n app_token=new_app_token,\n user_token=new_user_token,\n )\n\n assert isinstance(updated.bot_token, SensitiveValue)\n assert isinstance(updated.app_token, SensitiveValue)\n assert isinstance(updated.user_token, SensitiveValue)\n\n assert updated.bot_token.get_value(apply_mask=False) == new_bot_token\n assert updated.app_token.get_value(apply_mask=False) == new_app_token\n assert updated.user_token.get_value(apply_mask=False) == new_user_token\n\n # Verify from_model works without error\n pydantic_bot = SlackBot.from_model(updated)\n assert pydantic_bot.bot_token\n assert pydantic_bot.app_token\n assert pydantic_bot.user_token is not None\n" }, { "path": "backend/tests/external_dependency_unit/slack_bot/test_slack_bot_federated_search.py", "content": "# NOTE: ruff and black disagree after applying this noqa, so we just set file-level.\n# ruff: noqa: ARG005\nimport os\nfrom typing import Any\nfrom unittest.mock import MagicMock\nfrom unittest.mock import Mock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\n\n# Set environment variables to disable model server for testing\nos.environ[\"DISABLE_MODEL_SERVER\"] = \"true\"\nos.environ[\"MODEL_SERVER_HOST\"] = \"disabled\"\nos.environ[\"MODEL_SERVER_PORT\"] = \"9000\"\n\nfrom sqlalchemy import inspect\nfrom sqlalchemy.orm import Session\nfrom slack_sdk.errors import SlackApiError\n\nfrom onyx.configs.constants import FederatedConnectorSource\nfrom onyx.context.search.federated.slack_search import fetch_and_cache_channel_metadata\nfrom onyx.db.models import DocumentSet\nfrom onyx.db.models import FederatedConnector\nfrom onyx.db.models import FederatedConnector__DocumentSet\nfrom onyx.db.models import LLMProvider\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Persona__DocumentSet\nfrom onyx.db.models import Persona__Tool\nfrom onyx.db.models import SlackBot\nfrom onyx.db.models import SlackChannelConfig\nfrom onyx.db.models import User\nfrom onyx.onyxbot.slack.listener import process_message\nfrom onyx.onyxbot.slack.models import ChannelType\nfrom onyx.db.tools import get_builtin_tool\nfrom onyx.tools.built_in_tools import SearchTool\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom onyx.llm.constants import LlmProviderNames\n\n\ndef _create_test_persona_with_slack_config(db_session: Session) -> Persona | None:\n \"\"\"Helper to create a test persona configured for Slack federated search\"\"\"\n unique_id = str(uuid4())[:8]\n document_set = DocumentSet(\n name=f\"test_slack_docs_{unique_id}\",\n description=\"Test document set for Slack federated search\",\n )\n db_session.add(document_set)\n db_session.flush()\n\n persona = Persona(\n name=f\"test_slack_persona_{unique_id}\",\n description=\"Test persona for Slack federated search\",\n system_prompt=\"You are a helpful assistant.\",\n task_prompt=\"Answer the user's question based on the provided context.\",\n )\n db_session.add(persona)\n db_session.flush()\n\n persona_doc_set = Persona__DocumentSet(\n persona_id=persona.id,\n document_set_id=document_set.id,\n )\n db_session.add(persona_doc_set)\n db_session.commit()\n\n # Built-in tools are automatically seeded by migrations\n\n try:\n search_tool = get_builtin_tool(db_session=db_session, tool_type=SearchTool)\n if search_tool:\n persona_tool = Persona__Tool(persona_id=persona.id, tool_id=search_tool.id)\n db_session.add(persona_tool)\n except RuntimeError:\n # SearchTool not found, skip adding it\n pass\n\n db_session.commit()\n\n # Prompts are now directly on the persona table, no need for joinedload\n return persona\n\n\ndef _create_mock_slack_request(\n text: str, channel_id: str = \"C1234567890\", slack_bot_id: int = 12345\n) -> Mock:\n \"\"\"Create a mock Slack request\"\"\"\n mock_req = Mock()\n mock_req.type = \"events_api\"\n mock_req.envelope_id = \"test_envelope_id\"\n mock_req.payload = {\n \"event\": {\n \"type\": \"app_mention\",\n \"text\": f\"<@U1234567890> {text}\",\n \"channel\": channel_id,\n \"user\": \"U9876543210\",\n \"ts\": \"1234567890.123456\",\n }\n }\n mock_req.slack_bot_id = slack_bot_id\n return mock_req\n\n\ndef _create_mock_slack_client(\n channel_id: str = \"C1234567890\", # noqa: ARG001\n slack_bot_id: int = 12345,\n) -> Mock:\n \"\"\"Create a mock Slack client\"\"\"\n mock_client = Mock()\n mock_client.slack_bot_id = slack_bot_id\n mock_client.web_client = Mock()\n\n mock_post_message_response = {\"ok\": True, \"message_ts\": \"1234567890.123456\"}\n mock_client.web_client.chat_postMessage = Mock(\n return_value=mock_post_message_response\n )\n\n mock_users_info_response = Mock()\n mock_users_info_response.__getitem__ = Mock(\n side_effect=lambda key: {\"ok\": True}[key]\n )\n mock_users_info_response.data = {\n \"user\": {\n \"id\": \"U9876543210\",\n \"name\": \"testuser\",\n \"real_name\": \"Test User\",\n \"profile\": {\n \"display_name\": \"Test User\",\n \"first_name\": \"Test\",\n \"last_name\": \"User\",\n \"email\": \"test@example.com\",\n },\n }\n }\n mock_client.web_client.users_info = Mock(return_value=mock_users_info_response)\n\n mock_auth_test_response = {\n \"ok\": True,\n \"user_id\": \"U1234567890\",\n \"bot_id\": \"B1234567890\",\n }\n mock_client.web_client.auth_test = Mock(return_value=mock_auth_test_response)\n\n def mock_conversations_info_response(channel: str) -> Mock:\n channel_id = channel\n if channel_id == \"C1234567890\": # general - public\n mock_response = Mock()\n mock_response.validate.return_value = None\n mock_response.data = {\n \"channel\": {\n \"id\": \"C1234567890\",\n \"name\": \"general\",\n \"is_channel\": True,\n \"is_private\": False,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": False,\n }\n }\n mock_response.__getitem__ = lambda self, key: mock_response.data[key]\n return mock_response\n elif channel_id == \"C1111111111\": # support - public\n mock_response = Mock()\n mock_response.validate.return_value = None\n mock_response.data = {\n \"channel\": {\n \"id\": \"C1111111111\",\n \"name\": \"support\",\n \"is_channel\": True,\n \"is_private\": False,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": False,\n }\n }\n mock_response.__getitem__ = lambda self, key: mock_response.data[key]\n return mock_response\n elif channel_id == \"C9999999999\": # dev-team - private\n mock_response = Mock()\n mock_response.validate.return_value = None\n mock_response.data = {\n \"channel\": {\n \"id\": \"C9999999999\",\n \"name\": \"dev-team\",\n \"is_channel\": True,\n \"is_private\": True,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": False,\n }\n }\n mock_response.__getitem__ = lambda self, key: mock_response.data[key]\n return mock_response\n elif channel_id == \"D1234567890\": # DM\n mock_response = Mock()\n mock_response.validate.return_value = None\n mock_response.data = {\n \"channel\": {\n \"id\": \"D1234567890\",\n \"name\": \"directmessage\",\n \"is_channel\": False,\n \"is_private\": False,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": True,\n }\n }\n mock_response.__getitem__ = lambda self, key: mock_response.data[key]\n return mock_response\n else:\n mock_response = Mock()\n mock_response.validate.side_effect = Exception(\"channel_not_found\")\n return mock_response\n\n mock_client.web_client.conversations_info = Mock(\n side_effect=mock_conversations_info_response\n )\n\n mock_client.web_client.conversations_members = Mock(\n return_value={\"ok\": True, \"members\": [\"U9876543210\", \"U1234567890\"]}\n )\n\n mock_client.web_client.conversations_replies = Mock(\n return_value={\"ok\": True, \"messages\": []}\n )\n\n return mock_client\n\n\nclass TestSlackBotFederatedSearch:\n \"\"\"Test Slack bot federated search functionality\"\"\"\n\n def _setup_test_environment(\n self, db_session: Session\n ) -> tuple[User, Persona, FederatedConnector, SlackBot, SlackChannelConfig]:\n \"\"\"Setup test environment with user, persona, and federated connector\"\"\"\n user = create_test_user(db_session, \"slack_bot_test\")\n\n persona = _create_test_persona_with_slack_config(db_session)\n if persona is None:\n raise ValueError(\"Failed to create test persona\")\n\n federated_connector = FederatedConnector(\n source=FederatedConnectorSource.FEDERATED_SLACK,\n credentials={\"workspace_url\": \"https://test.slack.com\"},\n )\n db_session.add(federated_connector)\n db_session.flush()\n # Expire to ensure credentials is reloaded as SensitiveValue from DB\n db_session.expire(federated_connector)\n\n # Associate the federated connector with the persona's document sets\n # This is required for Slack federated search to be enabled\n for doc_set in persona.document_sets:\n federated_doc_set_mapping = FederatedConnector__DocumentSet(\n federated_connector_id=federated_connector.id,\n document_set_id=doc_set.id,\n entities={}, # Empty entities for test\n )\n db_session.add(federated_doc_set_mapping)\n db_session.flush()\n\n unique_id = str(uuid4())[:8]\n slack_bot = SlackBot(\n name=f\"Test Slack Bot {unique_id}\",\n bot_token=f\"xoxb-test-token-{unique_id}\",\n app_token=f\"xapp-test-token-{unique_id}\",\n user_token=f\"xoxp-test-user-token-{unique_id}\",\n enabled=True,\n )\n db_session.add(slack_bot)\n db_session.flush()\n # Expire to ensure tokens are reloaded as SensitiveValue from DB\n db_session.expire(slack_bot)\n\n slack_channel_config = SlackChannelConfig(\n slack_bot_id=slack_bot.id,\n persona_id=persona.id,\n channel_config={\"channel_name\": \"general\", \"disabled\": False},\n enable_auto_filters=True,\n is_default=True,\n )\n db_session.add(slack_channel_config)\n db_session.commit()\n\n return user, persona, federated_connector, slack_bot, slack_channel_config\n\n def _setup_slack_mocks(self, channel_name: str) -> tuple[list, list]:\n \"\"\"Setup only Slack API mocks - everything else runs live\"\"\"\n patches = [\n patch(\"slack_sdk.WebClient.search_messages\"),\n patch(\"onyx.context.search.federated.slack_search.query_slack\"),\n patch(\"onyx.onyxbot.slack.listener.get_channel_type_from_id\"),\n patch(\"onyx.context.search.utils.get_query_embeddings\"),\n ]\n\n started_patches = [p.start() for p in patches]\n\n self._setup_slack_api_mocks(started_patches[0], started_patches[0])\n\n self._setup_query_slack_mock(started_patches[1], channel_name)\n\n self._setup_channel_type_mock(started_patches[2], channel_name)\n\n self._setup_embedding_mock(started_patches[3])\n\n return patches, started_patches\n\n def _setup_embedding_mock(self, mock_get_query_embeddings: Mock) -> None:\n \"\"\"Mock embedding calls to avoid model server dependency\"\"\"\n # Return a dummy embedding vector for any query\n mock_get_query_embeddings.return_value = [[0.1] * 768] # 768-dimensional vector\n\n def _setup_slack_api_mocks(\n self,\n mock_search_messages: Mock,\n mock_conversations_info: Mock, # noqa: ARG002\n ) -> None:\n \"\"\"Setup Slack API mocks to return controlled data for testing filtering\"\"\"\n mock_search_response = Mock()\n mock_search_response.validate.return_value = None\n mock_search_response.get.return_value = {\n \"matches\": [\n {\n \"text\": \"Performance issue in API\",\n \"permalink\": \"https://test.slack.com/archives/C1234567890/p1234567890\",\n \"ts\": \"1234567890.123456\",\n \"channel\": {\"id\": \"C1234567890\", \"name\": \"general\"},\n \"username\": \"user1\",\n \"score\": 0.9,\n },\n {\n \"text\": \"Performance issue in dashboard\",\n \"permalink\": \"https://test.slack.com/archives/C1111111111/p1234567891\",\n \"ts\": \"1234567891.123456\",\n \"channel\": {\"id\": \"C1111111111\", \"name\": \"support\"},\n \"username\": \"user2\",\n \"score\": 0.8,\n },\n {\n \"text\": \"Performance issue in private channel\",\n \"permalink\": \"https://test.slack.com/archives/C9999999999/p1234567892\",\n \"ts\": \"1234567892.123456\",\n \"channel\": {\"id\": \"C9999999999\", \"name\": \"dev-team\"},\n \"username\": \"user3\",\n \"score\": 0.7,\n },\n {\n \"text\": \"Performance issue in DM\",\n \"permalink\": \"https://test.slack.com/archives/D1234567890/p1234567893\",\n \"ts\": \"1234567893.123456\",\n \"channel\": {\"id\": \"D1234567890\", \"name\": \"directmessage\"},\n \"username\": \"user4\",\n \"score\": 0.6,\n },\n ]\n }\n mock_search_messages.return_value = mock_search_response\n\n def _setup_query_slack_mock(\n self, mock_query_slack: Mock, channel_name: str\n ) -> None:\n \"\"\"Setup query_slack mock to capture filtering parameters\"\"\"\n from onyx.context.search.federated.slack_search import SlackQueryResult\n\n def mock_query_slack_capture_params(\n query_string: str, # noqa: ARG001\n access_token: str, # noqa: ARG001\n limit: int | None = None, # noqa: ARG001\n allowed_private_channel: str | None = None,\n bot_token: str | None = None, # noqa: ARG001\n include_dm: bool = False,\n entities: dict | None = None, # noqa: ARG001\n available_channels: list | None = None, # noqa: ARG001\n channel_metadata_dict: dict | None = None, # noqa: ARG001\n ) -> SlackQueryResult:\n self._captured_filtering_params = {\n \"allowed_private_channel\": allowed_private_channel,\n \"include_dm\": include_dm,\n \"channel_name\": channel_name,\n }\n\n return SlackQueryResult(messages=[], filtered_channels=[])\n\n mock_query_slack.side_effect = mock_query_slack_capture_params\n\n def _setup_channel_type_mock(\n self,\n mock_get_channel_type_from_id: Mock,\n channel_name: str, # noqa: ARG002\n ) -> None:\n \"\"\"Setup get_channel_type_from_id mock to return correct channel types\"\"\"\n\n def mock_channel_type_response(\n web_client: Mock, # noqa: ARG001\n channel_id: str,\n ) -> ChannelType:\n if channel_id == \"C1234567890\": # general - public\n return ChannelType.PUBLIC_CHANNEL\n elif channel_id == \"C1111111111\": # support - public\n return ChannelType.PUBLIC_CHANNEL\n elif channel_id == \"C9999999999\": # dev-team - private\n return ChannelType.PRIVATE_CHANNEL\n elif channel_id == \"D1234567890\": # DM\n return ChannelType.IM\n else:\n return ChannelType.PUBLIC_CHANNEL # default\n\n mock_get_channel_type_from_id.side_effect = mock_channel_type_response\n\n def _setup_llm_provider(self, db_session: Session) -> None:\n \"\"\"Create a default LLM provider in the database for testing with real API key\"\"\"\n # Delete any existing default LLM provider to ensure clean state\n # Use SQL-level delete to properly trigger ON DELETE CASCADE\n # (ORM-level delete tries to set foreign keys to NULL instead)\n from sqlalchemy import delete\n\n existing_providers = db_session.query(LLMProvider).all()\n for provider in existing_providers:\n db_session.execute(delete(LLMProvider).where(LLMProvider.id == provider.id))\n db_session.commit()\n\n api_key = os.getenv(\"OPENAI_API_KEY\")\n if not api_key:\n raise ValueError(\n \"OPENAI_API_KEY environment variable not set - test requires real API key\"\n )\n\n provider_view = upsert_llm_provider(\n LLMProviderUpsertRequest(\n name=f\"test-llm-provider-{uuid4().hex[:8]}\",\n provider=LlmProviderNames.OPENAI,\n api_key=api_key,\n is_public=True,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n max_input_tokens=None,\n display_name=\"gpt-4o\",\n ),\n ],\n ),\n db_session=db_session,\n )\n\n update_default_provider(provider_view.id, \"gpt-4o\", db_session)\n\n def _teardown_common_mocks(self, patches: list) -> None:\n \"\"\"Stop all patches\"\"\"\n for p in patches:\n p.stop()\n\n @patch(\"onyx.utils.gpu_utils.fast_gpu_status_request\", return_value=False)\n @patch(\n \"onyx.document_index.vespa.index.VespaIndex.hybrid_retrieval\", return_value=[]\n )\n def test_slack_bot_public_channel_filtering(\n self,\n mock_vespa: Mock, # noqa: ARG002\n mock_gpu_status: Mock, # noqa: ARG002\n db_session: Session,\n ) -> None:\n \"\"\"Test that slack bot in public channel sees only public channel messages\"\"\"\n self._setup_llm_provider(db_session)\n\n user, persona, federated_connector, slack_bot, slack_channel_config = (\n self._setup_test_environment(db_session)\n )\n\n channel_id = \"C1234567890\" # #general (public)\n channel_name = \"general\"\n\n patches, started_patches = self._setup_slack_mocks(channel_name)\n\n try:\n mock_req = _create_mock_slack_request(\n \"search for performance issues\", channel_id, slack_bot.id\n )\n mock_client = _create_mock_slack_client(channel_id, slack_bot.id)\n\n process_message(mock_req, mock_client)\n\n mock_client.web_client.chat_postMessage.assert_called()\n post_message_calls = mock_client.web_client.chat_postMessage.call_args_list\n last_call = post_message_calls[-1]\n assert (\n last_call[1][\"channel\"] == channel_id\n ), f\"Response should be sent to {channel_id}\"\n\n response_text = last_call[1].get(\"text\", \"\")\n assert len(response_text) > 0, \"Bot should have sent a non-empty response\"\n\n assert hasattr(\n self, \"_captured_filtering_params\"\n ), \"query_slack should have been called\"\n params = self._captured_filtering_params\n\n assert (\n params[\"allowed_private_channel\"] is None\n ), \"Public channels should not have private channel access\"\n assert (\n params[\"include_dm\"] is False\n ), \"Public channels should not include DMs\"\n assert (\n params[\"channel_name\"] == \"general\"\n ), \"Should be testing general channel\"\n\n finally:\n self._teardown_common_mocks(patches)\n\n @patch(\"onyx.utils.gpu_utils.fast_gpu_status_request\", return_value=False)\n @patch(\n \"onyx.document_index.vespa.index.VespaIndex.hybrid_retrieval\", return_value=[]\n )\n def test_slack_bot_private_channel_filtering(\n self,\n mock_vespa: Mock, # noqa: ARG002\n mock_gpu_status: Mock, # noqa: ARG002\n db_session: Session,\n ) -> None:\n \"\"\"Test that slack bot in private channel sees private + public channel messages\"\"\"\n self._setup_llm_provider(db_session)\n\n user, persona, federated_connector, slack_bot, slack_channel_config = (\n self._setup_test_environment(db_session)\n )\n\n channel_id = \"C9999999999\" # #dev-team (private)\n channel_name = \"dev-team\"\n\n patches, started_patches = self._setup_slack_mocks(channel_name)\n\n try:\n mock_req = _create_mock_slack_request(\n \"search for performance issues\", channel_id, slack_bot.id\n )\n mock_client = _create_mock_slack_client(channel_id, slack_bot.id)\n\n process_message(mock_req, mock_client)\n\n mock_client.web_client.chat_postMessage.assert_called()\n post_message_calls = mock_client.web_client.chat_postMessage.call_args_list\n last_call = post_message_calls[-1]\n assert (\n last_call[1][\"channel\"] == channel_id\n ), f\"Response should be sent to {channel_id}\"\n\n response_text = last_call[1].get(\"text\", \"\")\n assert len(response_text) > 0, \"Bot should have sent a non-empty response\"\n\n assert hasattr(\n self, \"_captured_filtering_params\"\n ), \"query_slack should have been called\"\n params = self._captured_filtering_params\n\n assert (\n params[\"allowed_private_channel\"] == \"C9999999999\"\n ), \"Private channels should have access to their specific private channel\"\n assert (\n params[\"include_dm\"] is False\n ), \"Private channels should not include DMs\"\n assert (\n params[\"channel_name\"] == \"dev-team\"\n ), \"Should be testing dev-team channel\"\n\n finally:\n self._teardown_common_mocks(patches)\n\n @patch(\"onyx.utils.gpu_utils.fast_gpu_status_request\", return_value=False)\n @patch(\n \"onyx.document_index.vespa.index.VespaIndex.hybrid_retrieval\", return_value=[]\n )\n def test_slack_bot_dm_filtering(\n self,\n mock_vespa: Mock, # noqa: ARG002\n mock_gpu_status: Mock, # noqa: ARG002\n db_session: Session,\n ) -> None:\n \"\"\"Test that slack bot in DM sees all messages (no filtering)\"\"\"\n self._setup_llm_provider(db_session)\n\n user, persona, federated_connector, slack_bot, slack_channel_config = (\n self._setup_test_environment(db_session)\n )\n\n channel_id = \"D1234567890\" # DM\n channel_name = \"directmessage\"\n\n patches, started_patches = self._setup_slack_mocks(channel_name)\n\n try:\n mock_req = _create_mock_slack_request(\n \"search for performance issues\", channel_id, slack_bot.id\n )\n mock_client = _create_mock_slack_client(channel_id, slack_bot.id)\n\n process_message(mock_req, mock_client)\n\n mock_client.web_client.chat_postMessage.assert_called()\n post_message_calls = mock_client.web_client.chat_postMessage.call_args_list\n last_call = post_message_calls[-1]\n assert (\n last_call[1][\"channel\"] == channel_id\n ), f\"Response should be sent to {channel_id}\"\n\n response_text = last_call[1].get(\"text\", \"\")\n assert len(response_text) > 0, \"Bot should have sent a non-empty response\"\n\n assert hasattr(\n self, \"_captured_filtering_params\"\n ), \"query_slack should have been called\"\n params = self._captured_filtering_params\n\n assert (\n params[\"allowed_private_channel\"] is None\n ), \"DMs should not have private channel access\"\n assert params[\"include_dm\"] is True, \"DMs should include DM messages\"\n assert (\n params[\"channel_name\"] == \"directmessage\"\n ), \"Should be testing directmessage channel\"\n\n finally:\n self._teardown_common_mocks(patches)\n\n\n@patch(\"onyx.context.search.federated.slack_search.get_redis_client\")\n@patch(\"onyx.context.search.federated.slack_search.WebClient\")\ndef test_missing_scope_resilience(\n mock_web_client: Mock, mock_redis_client: Mock\n) -> None:\n \"\"\"Test that missing scopes are handled gracefully\"\"\"\n # Setup mock Redis client\n mock_redis = MagicMock()\n mock_redis.get.return_value = None # Cache miss\n mock_redis_client.return_value = mock_redis\n\n # Setup mock Slack client that simulates missing_scope error\n mock_client_instance = MagicMock()\n mock_web_client.return_value = mock_client_instance\n\n # Track which channel types were attempted\n attempted_types: list[str] = []\n\n def mock_conversations_list(\n types: str | None = None,\n **kwargs: Any, # noqa: ARG001\n ) -> MagicMock:\n if types:\n attempted_types.append(types)\n\n # First call: all types including mpim -> missing_scope error\n if types and \"mpim\" in types:\n error_response = {\n \"ok\": False,\n \"error\": \"missing_scope\",\n \"needed\": \"mpim:read\",\n \"provided\": \"identify,channels:history,channels:read,groups:read,im:read,search:read\",\n }\n raise SlackApiError(\"missing_scope\", error_response)\n\n # Second call: without mpim -> success\n mock_response = MagicMock()\n mock_response.validate.return_value = None\n mock_response.data = {\n \"channels\": [\n {\n \"id\": \"C1234567890\",\n \"name\": \"general\",\n \"is_channel\": True,\n \"is_private\": False,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": False,\n \"is_member\": True,\n },\n {\n \"id\": \"D9876543210\",\n \"name\": \"\",\n \"is_channel\": False,\n \"is_private\": False,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": True,\n \"is_member\": True,\n },\n ],\n \"response_metadata\": {},\n }\n return mock_response\n\n mock_client_instance.conversations_list.side_effect = mock_conversations_list\n\n # Call the function\n result = fetch_and_cache_channel_metadata(\n access_token=\"xoxp-test-token\",\n team_id=\"T1234567890\",\n include_private=True,\n )\n\n # Assertions\n # Should have attempted twice: once with mpim, once without\n assert len(attempted_types) == 2, f\"Expected 2 attempts, got {len(attempted_types)}\"\n assert \"mpim\" in attempted_types[0], \"First attempt should include mpim\"\n assert \"mpim\" not in attempted_types[1], \"Second attempt should not include mpim\"\n\n # Should have successfully returned channels despite missing scope\n assert len(result) == 2, f\"Expected 2 channels, got {len(result)}\"\n assert \"C1234567890\" in result, \"Should have public channel\"\n assert \"D9876543210\" in result, \"Should have DM channel\"\n\n # Verify channel metadata structure\n assert result[\"C1234567890\"][\"name\"] == \"general\"\n assert result[\"C1234567890\"][\"type\"] == \"public_channel\"\n assert result[\"D9876543210\"][\"type\"] == \"im\"\n\n\n@patch(\"onyx.context.search.federated.slack_search.get_redis_client\")\n@patch(\"onyx.context.search.federated.slack_search.WebClient\")\ndef test_multiple_missing_scopes_resilience(\n mock_web_client: Mock, mock_redis_client: Mock\n) -> None:\n \"\"\"Test handling multiple missing scopes gracefully\"\"\"\n # Setup mock Redis client\n mock_redis = MagicMock()\n mock_redis.get.return_value = None # Cache miss\n mock_redis_client.return_value = mock_redis\n\n # Setup mock Slack client\n mock_client_instance = MagicMock()\n mock_web_client.return_value = mock_client_instance\n\n # Track attempts\n attempted_types: list[str] = []\n\n def mock_conversations_list(\n types: str | None = None,\n **kwargs: Any, # noqa: ARG001\n ) -> MagicMock:\n if types:\n attempted_types.append(types)\n\n # First: mpim missing\n if types and \"mpim\" in types:\n error_response = {\n \"ok\": False,\n \"error\": \"missing_scope\",\n \"needed\": \"mpim:read\",\n \"provided\": \"identify,channels:history,channels:read,groups:read\",\n }\n raise SlackApiError(\"missing_scope\", error_response)\n\n # Second: im missing\n if types and \"im\" in types:\n error_response = {\n \"ok\": False,\n \"error\": \"missing_scope\",\n \"needed\": \"im:read\",\n \"provided\": \"identify,channels:history,channels:read,groups:read\",\n }\n raise SlackApiError(\"missing_scope\", error_response)\n\n # Third: success with only public and private channels\n mock_response = MagicMock()\n mock_response.validate.return_value = None\n mock_response.data = {\n \"channels\": [\n {\n \"id\": \"C1234567890\",\n \"name\": \"general\",\n \"is_channel\": True,\n \"is_private\": False,\n \"is_group\": False,\n \"is_mpim\": False,\n \"is_im\": False,\n \"is_member\": True,\n }\n ],\n \"response_metadata\": {},\n }\n return mock_response\n\n mock_client_instance.conversations_list.side_effect = mock_conversations_list\n\n # Call the function\n result = fetch_and_cache_channel_metadata(\n access_token=\"xoxp-test-token\",\n team_id=\"T1234567890\",\n include_private=True,\n )\n\n # Should gracefully handle multiple missing scopes\n assert len(attempted_types) == 3, f\"Expected 3 attempts, got {len(attempted_types)}\"\n assert \"mpim\" in attempted_types[0], \"First attempt should include mpim\"\n assert \"mpim\" not in attempted_types[1], \"Second attempt should not include mpim\"\n assert \"im\" in attempted_types[1], \"Second attempt should include im\"\n assert \"im\" not in attempted_types[2], \"Third attempt should not include im\"\n\n # Should still return available channels\n assert len(result) == 1, f\"Expected 1 channel, got {len(result)}\"\n assert result[\"C1234567890\"][\"name\"] == \"general\"\n\n\ndef test_slack_channel_config_eager_loads_persona(db_session: Session) -> None:\n \"\"\"Test that fetch_slack_channel_config_for_channel_or_default eagerly loads persona.\n\n This prevents lazy loading failures when the session context changes later\n in the request handling flow (e.g., in handle_regular_answer).\n \"\"\"\n from onyx.db.slack_channel_config import (\n fetch_slack_channel_config_for_channel_or_default,\n )\n\n unique_id = str(uuid4())[:8]\n\n # Create a persona (using same fields as _create_test_persona_with_slack_config)\n persona = Persona(\n name=f\"test_eager_load_persona_{unique_id}\",\n description=\"Test persona for eager loading test\",\n system_prompt=\"You are a helpful assistant.\",\n task_prompt=\"Answer the user's question.\",\n )\n db_session.add(persona)\n db_session.flush()\n\n # Create a slack bot\n slack_bot = SlackBot(\n name=f\"Test Bot {unique_id}\",\n bot_token=f\"xoxb-test-{unique_id}\",\n app_token=f\"xapp-test-{unique_id}\",\n enabled=True,\n )\n db_session.add(slack_bot)\n db_session.flush()\n\n # Create slack channel config with persona\n channel_name = f\"test-channel-{unique_id}\"\n slack_channel_config = SlackChannelConfig(\n slack_bot_id=slack_bot.id,\n persona_id=persona.id,\n channel_config={\"channel_name\": channel_name, \"disabled\": False},\n enable_auto_filters=False,\n is_default=False,\n )\n db_session.add(slack_channel_config)\n db_session.commit()\n\n # Fetch the config using the function under test\n fetched_config = fetch_slack_channel_config_for_channel_or_default(\n db_session=db_session,\n slack_bot_id=slack_bot.id,\n channel_name=channel_name,\n )\n\n assert fetched_config is not None, \"Should find the channel config\"\n\n # Check that persona relationship is already loaded (not pending lazy load)\n insp = inspect(fetched_config)\n assert insp is not None, \"Should be able to inspect the config\"\n assert \"persona\" not in insp.unloaded, (\n \"Persona should be eagerly loaded, not pending lazy load. \"\n \"This is required to prevent fallback to default persona when \"\n \"session context changes in handle_regular_answer.\"\n )\n\n # Verify the persona is correct\n assert fetched_config.persona is not None, \"Persona should not be None\"\n assert fetched_config.persona.id == persona.id, \"Should load the correct persona\"\n assert fetched_config.persona.name == persona.name\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_image_generation_tool.py", "content": "# TODO re-enable this test\n# import os\n# import time\n# from typing import Any\n# from unittest.mock import patch\n\n# import pytest\n\n# from onyx.tools.models import ToolResponse\n# from onyx.tools.tool_implementations.images.image_generation_tool import (\n# IMAGE_GENERATION_HEARTBEAT_ID,\n# )\n# from onyx.tools.tool_implementations.images.image_generation_tool import (\n# IMAGE_GENERATION_RESPONSE_ID,\n# )\n# from onyx.tools.tool_implementations.images.image_generation_tool import (\n# ImageGenerationResponse,\n# )\n# from onyx.tools.tool_implementations.images.image_generation_tool import (\n# ImageGenerationTool,\n# )\n# from onyx.tools.tool_implementations.images.image_generation_tool import ImageShape\n\n\n# @pytest.fixture\n# def dalle3_tool() -> ImageGenerationTool:\n# \"\"\"Fixture for DALL-E 3 tool with API key from environment.\"\"\"\n# api_key = os.environ[\"OPENAI_API_KEY\"]\n# return ImageGenerationTool(\n# tool_id=0,\n# api_key=api_key,\n# api_base=None,\n# api_version=None,\n# model=\"dall-e-3\",\n# num_imgs=1,\n# )\n\n\n# def test_image_generation_with_heartbeats(dalle3_tool: ImageGenerationTool) -> None:\n# \"\"\"Test that heartbeat packets are yielded during image generation.\"\"\"\n# responses = []\n# heartbeat_count = 0\n# image_response_count = 0\n\n# # Collect all responses\n# for response in dalle3_tool.run(prompt=\"A simple red circle on white background\"):\n# responses.append(response)\n# if response.id == IMAGE_GENERATION_HEARTBEAT_ID:\n# heartbeat_count += 1\n# elif response.id == IMAGE_GENERATION_RESPONSE_ID:\n# image_response_count += 1\n\n# # Should have at least one heartbeat (depending on generation speed)\n# # and exactly one image response\n# assert image_response_count == 1\n# # May have 0 or more heartbeats depending on API speed\n# print(f\"Received {heartbeat_count} heartbeat packets\")\n\n# # Verify the final image response\n# final_response = responses[-1]\n# assert final_response.id == IMAGE_GENERATION_RESPONSE_ID\n# assert isinstance(final_response.response, list)\n# assert len(final_response.response) == 1\n\n# image = final_response.response[0]\n# assert isinstance(image, ImageGenerationResponse)\n# assert image.image_data is not None\n# assert len(image.image_data) > 100 # Base64 data should be substantial\n# assert image.revised_prompt is not None\n\n\n# def test_heartbeat_timing_with_mock() -> None:\n# \"\"\"Test that heartbeats are sent at correct intervals using mocked generation.\"\"\"\n# api_key = os.getenv(\"OPENAI_API_KEY\", \"mock-key-for-testing\")\n\n# tool = ImageGenerationTool(\n# tool_id=0,\n# api_key=api_key,\n# api_base=None,\n# api_version=None,\n# model=\"dall-e-3\",\n# num_imgs=1,\n# )\n\n# # Mock the _generate_image method to simulate slow generation\n# def slow_generate(*args: Any, **kwargs: Any) -> ImageGenerationResponse:\n# time.sleep(5) # Simulate 5 second generation time\n# return ImageGenerationResponse(\n# revised_prompt=\"Test prompt\",\n# image_data=\"base64encodedimagedata\",\n# )\n\n# with patch.object(tool, \"_generate_image\", side_effect=slow_generate):\n# start_time = time.time()\n# responses = list(tool.run(prompt=\"Test prompt\"))\n# time.time() - start_time\n\n# # Count heartbeats\n# heartbeat_count = sum(\n# 1 for r in responses if r.id == IMAGE_GENERATION_HEARTBEAT_ID\n# )\n\n# # With 5 second generation and 2 second intervals,\n# # we should get approximately 2 heartbeats\n# assert heartbeat_count >= 1\n# assert heartbeat_count <= 3 # Allow some timing variance\n\n# # Verify we still get the final result\n# image_responses = [r for r in responses if r.id == IMAGE_GENERATION_RESPONSE_ID]\n# assert len(image_responses) == 1\n# assert image_responses[0].response[0].image_data == \"base64encodedimagedata\"\n\n\n# def test_error_handling_with_heartbeats() -> None:\n# \"\"\"Test that errors are properly propagated even with heartbeat mechanism.\"\"\"\n# api_key = os.getenv(\"OPENAI_API_KEY\", \"mock-key-for-testing\")\n\n# tool = ImageGenerationTool(\n# tool_id=0,\n# api_key=api_key,\n# api_base=None,\n# api_version=None,\n# model=\"dall-e-3\",\n# num_imgs=1,\n# )\n\n# # Mock the _generate_image method to raise an error after delay\n# def error_generate(*args: Any, **kwargs: Any) -> None:\n# time.sleep(1) # Small delay to ensure at least one heartbeat\n# raise ValueError(\"Test error during generation\")\n\n# with patch.object(tool, \"_generate_image\", side_effect=error_generate):\n# with pytest.raises(ValueError, match=\"Test error during generation\"):\n# # Consume the generator to trigger the error\n# list(tool.run(prompt=\"Test prompt\"))\n\n\n# def test_tool_message_content_filters_heartbeats() -> None:\n# \"\"\"Test that get_llm_tool_response correctly filters heartbeats.\"\"\"\n# api_key = os.getenv(\"OPENAI_API_KEY\", \"mock-key-for-testing\")\n\n# tool = ImageGenerationTool(\n# tool_id=0,\n# api_key=api_key,\n# api_base=None,\n# api_version=None,\n# model=\"dall-e-3\",\n# num_imgs=1,\n# )\n\n# # Create mock responses\n# heartbeat1 = ToolResponse(\n# id=IMAGE_GENERATION_HEARTBEAT_ID,\n# response={\"status\": \"generating\", \"heartbeat\": 0},\n# )\n# heartbeat2 = ToolResponse(\n# id=IMAGE_GENERATION_HEARTBEAT_ID,\n# response={\"status\": \"generating\", \"heartbeat\": 1},\n# )\n# image_response = ToolResponse(\n# id=IMAGE_GENERATION_RESPONSE_ID,\n# response=[\n# ImageGenerationResponse(\n# revised_prompt=\"Test\",\n# image_data=\"base64encodedimagedata\",\n# )\n# ],\n# )\n\n# # Test that heartbeats are filtered out\n# result = tool.get_llm_tool_response(heartbeat1, heartbeat2, image_response)\n\n# # Should return JSON with image info, not heartbeats\n# assert isinstance(result, str)\n# assert \"Test\" in result\n# assert \"heartbeat\" not in result\n\n\n# def test_final_result_filters_heartbeats() -> None:\n# \"\"\"Test that final_result correctly filters heartbeats.\"\"\"\n# api_key = os.getenv(\"OPENAI_API_KEY\", \"mock-key-for-testing\")\n\n# tool = ImageGenerationTool(\n# tool_id=0,\n# api_key=api_key,\n# api_base=None,\n# api_version=None,\n# model=\"dall-e-3\",\n# num_imgs=1,\n# )\n\n# # Create mock responses\n# heartbeat = ToolResponse(\n# id=IMAGE_GENERATION_HEARTBEAT_ID,\n# response={\"status\": \"generating\", \"heartbeat\": 0},\n# )\n# image_response = ToolResponse(\n# id=IMAGE_GENERATION_RESPONSE_ID,\n# response=[\n# ImageGenerationResponse(\n# revised_prompt=\"Test prompt\",\n# image_data=\"base64encodedimagedata\",\n# )\n# ],\n# )\n\n# # Test that final_result returns only image data\n# result = tool.get_final_result(heartbeat, image_response)\n\n# assert isinstance(result, list)\n# assert len(result) == 1\n# assert result[0][\"revised_prompt\"] == \"Test prompt\"\n# assert result[0][\"image_data\"] == \"base64encodedimagedata\"\n\n\n# def test_different_image_shapes(dalle3_tool: ImageGenerationTool) -> None:\n# \"\"\"Test image generation with different shape parameters.\"\"\"\n# shapes_to_test = [\n# (ImageShape.SQUARE, \"A red square\"),\n# (ImageShape.PORTRAIT, \"A tall building\"),\n# (ImageShape.LANDSCAPE, \"A wide landscape\"),\n# ]\n\n# for shape, prompt in shapes_to_test:\n# responses = list(dalle3_tool.run(prompt=prompt, shape=shape.value))\n\n# # Find the image response\n# image_response = None\n# for response in responses:\n# if response.id == IMAGE_GENERATION_RESPONSE_ID:\n# image_response = response\n# break\n\n# assert image_response is not None\n# assert len(image_response.response) == 1\n# image = image_response.response[0]\n# assert image.image_data is not None\n# assert len(image.image_data) > 100 # Base64 data should be substantial\n# print(f\"Generated {shape.value} image (base64, {len(image.image_data)} chars)\")\n\n\n# def test_image_generation_response_format() -> None:\n# \"\"\"Test that image generation returns data in at least one format (URL or base64).\"\"\"\n# api_key = os.getenv(\"OPENAI_API_KEY\")\n# if not api_key:\n# pytest.skip(\"OPENAI_API_KEY environment variable not set\")\n\n# tool = ImageGenerationTool(\n# tool_id=0,\n# api_key=api_key,\n# api_base=None,\n# api_version=None,\n# model=\"dall-e-3\",\n# num_imgs=1,\n# )\n\n# responses = list(tool.run(prompt=\"A simple blue circle\"))\n\n# # Find the image response\n# image_response = None\n# for response in responses:\n# if response.id == IMAGE_GENERATION_RESPONSE_ID:\n# image_response = response\n# break\n\n# assert image_response is not None\n# assert len(image_response.response) == 1\n# image = image_response.response[0]\n# # Should always have base64 data\n# assert image.image_data is not None\n# assert len(image.image_data) > 100 # Base64 data should be substantial\n\n\n# if __name__ == \"__main__\":\n# # Run with: python -m pytest tests/external_dependency_unit/tools/test_image_generation_tool.py -v\n# pytest.main([__file__, \"-v\"])\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_mcp_passthrough_oauth.py", "content": "\"\"\"\nTest suite for MCP Pass-Through OAuth (PT_OAUTH) integration.\n\nTests the pass-through OAuth flow where Onyx forwards the user's login OAuth token\nto an MCP server for authentication.\n\nThis test:\n1. Creates a test user with an OAuthAccount (simulating Google OAuth login)\n2. Creates an MCP server with PT_OAUTH auth type\n3. Creates MCP tools for that server\n4. Verifies the user's OAuth token is correctly passed to MCPTool\n\nAll external HTTP calls are mocked, but Postgres and Redis are running.\n\"\"\"\n\nimport queue\nfrom typing import Any\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.emitter import Emitter\nfrom onyx.db.enums import MCPAuthenticationPerformer\nfrom onyx.db.enums import MCPAuthenticationType\nfrom onyx.db.enums import MCPTransport\nfrom onyx.db.mcp import create_mcp_server__no_commit\nfrom onyx.db.models import OAuthAccount\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Tool\nfrom onyx.db.models import User\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.server.query_and_chat.placement import Placement\nfrom onyx.tools.models import CustomToolCallSummary\nfrom onyx.tools.tool_constructor import construct_tools\nfrom onyx.tools.tool_constructor import SearchToolConfig\nfrom onyx.tools.tool_implementations.mcp.mcp_tool import MCPTool\nfrom tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef _create_test_persona_with_mcp_tool(\n db_session: Session, user: User, tools: list[Tool]\n) -> Persona:\n \"\"\"Helper to create a test persona with MCP tools\"\"\"\n persona = Persona(\n name=f\"Test MCP Persona {uuid4().hex[:8]}\",\n description=\"Test persona with MCP tools\",\n system_prompt=\"You are a helpful assistant\",\n task_prompt=\"Answer the user's question\",\n tools=tools,\n document_sets=[],\n users=[user],\n groups=[],\n is_listed=True,\n is_public=True,\n display_priority=None,\n starter_messages=None,\n deleted=False,\n )\n db_session.add(persona)\n db_session.commit()\n db_session.refresh(persona)\n return persona\n\n\nclass TestMCPPassThroughOAuth:\n \"\"\"Tests for MCP Pass-Through OAuth (PT_OAUTH) flow\"\"\"\n\n @pytest.fixture(autouse=True)\n def setup_llm_provider(self, db_session: Session) -> None:\n \"\"\"Ensure default LLM provider is set up for each test.\"\"\"\n ensure_default_llm_provider(db_session)\n\n def test_pt_oauth_passes_user_login_token(self, db_session: Session) -> None:\n \"\"\"\n Test that PT_OAUTH correctly passes the user's login OAuth token to MCPTool.\n\n This simulates a user who logged into Onyx with Google OAuth and is using\n an MCP server that requires their Google token for authentication.\n \"\"\"\n # Create user with login OAuth token (simulating Google OAuth login)\n user = create_test_user(db_session, \"pt_oauth_user\")\n user_oauth_token = \"google_oauth_token_abc123\"\n\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"google\",\n account_id=\"google_user_12345\",\n account_email=user.email,\n access_token=user_oauth_token,\n refresh_token=\"google_refresh_token\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n # Refresh user to load oauth_accounts relationship\n db_session.refresh(user)\n\n # Create MCP server with PT_OAUTH auth type\n mcp_server = create_mcp_server__no_commit(\n owner_email=user.email,\n name=f\"PT_OAUTH Test Server {uuid4().hex[:8]}\",\n description=\"MCP server for pass-through OAuth testing\",\n server_url=\"http://test-mcp-server.example.com/mcp\",\n auth_type=MCPAuthenticationType.PT_OAUTH,\n transport=MCPTransport.STREAMABLE_HTTP,\n auth_performer=MCPAuthenticationPerformer.ADMIN, # Not used for PT_OAUTH\n db_session=db_session,\n )\n db_session.commit()\n\n # Create MCP tool associated with this server\n mcp_tool_db = Tool(\n name=\"test_mcp_tool\",\n display_name=\"Test MCP Tool\",\n description=\"Test MCP tool for PT_OAUTH\",\n mcp_server_id=mcp_server.id,\n mcp_input_schema={\n \"type\": \"object\",\n \"properties\": {\n \"message\": {\"type\": \"string\", \"description\": \"Test message\"}\n },\n },\n user_id=user.id,\n )\n db_session.add(mcp_tool_db)\n db_session.commit()\n db_session.refresh(mcp_tool_db)\n\n # Create persona with the MCP tool\n persona = _create_test_persona_with_mcp_tool(db_session, user, [mcp_tool_db])\n llm = get_default_llm()\n\n # Construct tools\n search_tool_config = SearchToolConfig()\n\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=search_tool_config,\n )\n\n # Verify MCP tool was constructed\n assert mcp_tool_db.id in tool_dict\n constructed_tools = tool_dict[mcp_tool_db.id]\n assert len(constructed_tools) == 1\n mcp_tool = constructed_tools[0]\n assert isinstance(mcp_tool, MCPTool)\n\n # Verify the user's OAuth token was passed to the MCPTool\n assert mcp_tool._user_oauth_token == user_oauth_token\n\n def test_pt_oauth_without_user_oauth_account(self, db_session: Session) -> None:\n \"\"\"\n Test PT_OAUTH behavior when user doesn't have an OAuth account.\n\n The user logged in with basic auth (no OAuth token), so the MCP tool\n should have no OAuth token to pass through.\n \"\"\"\n # Create user WITHOUT OAuth account (basic auth login)\n user = create_test_user(db_session, \"basic_auth_user\")\n # No OAuthAccount created\n\n # Create MCP server with PT_OAUTH auth type\n mcp_server = create_mcp_server__no_commit(\n owner_email=user.email,\n name=f\"PT_OAUTH No Token Server {uuid4().hex[:8]}\",\n description=\"MCP server for testing missing OAuth token\",\n server_url=\"http://test-mcp-server.example.com/mcp\",\n auth_type=MCPAuthenticationType.PT_OAUTH,\n transport=MCPTransport.STREAMABLE_HTTP,\n auth_performer=MCPAuthenticationPerformer.ADMIN,\n db_session=db_session,\n )\n db_session.commit()\n\n # Create MCP tool\n mcp_tool_db = Tool(\n name=\"test_mcp_tool_no_token\",\n display_name=\"Test MCP Tool No Token\",\n description=\"Test MCP tool without OAuth token\",\n mcp_server_id=mcp_server.id,\n mcp_input_schema={\n \"type\": \"object\",\n \"properties\": {\"query\": {\"type\": \"string\"}},\n },\n user_id=user.id,\n )\n db_session.add(mcp_tool_db)\n db_session.commit()\n db_session.refresh(mcp_tool_db)\n\n # Create persona\n persona = _create_test_persona_with_mcp_tool(db_session, user, [mcp_tool_db])\n llm = get_default_llm()\n\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=SearchToolConfig(),\n )\n\n # Verify MCP tool was constructed\n assert mcp_tool_db.id in tool_dict\n constructed_tools = tool_dict[mcp_tool_db.id]\n assert len(constructed_tools) == 1\n mcp_tool = constructed_tools[0]\n assert isinstance(mcp_tool, MCPTool)\n\n # Verify NO OAuth token was passed (user has no OAuth account)\n assert mcp_tool._user_oauth_token is None\n\n def test_pt_oauth_vs_api_token_auth(self, db_session: Session) -> None:\n \"\"\"\n Test that PT_OAUTH and API_TOKEN auth types behave differently.\n\n PT_OAUTH should use the user's login token, while API_TOKEN should\n NOT use the user's login token (it uses the connection config instead).\n \"\"\"\n # Create user with OAuth account\n user = create_test_user(db_session, \"auth_type_test_user\")\n user_oauth_token = \"user_login_token_xyz789\"\n\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"google\",\n account_id=\"google_user_xyz\",\n account_email=user.email,\n access_token=user_oauth_token,\n refresh_token=\"\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n db_session.refresh(user)\n\n # Create MCP server with API_TOKEN auth type (not PT_OAUTH)\n mcp_server = create_mcp_server__no_commit(\n owner_email=user.email,\n name=f\"API Token Server {uuid4().hex[:8]}\",\n description=\"MCP server with API token auth\",\n server_url=\"http://api-token-server.example.com/mcp\",\n auth_type=MCPAuthenticationType.API_TOKEN, # Not PT_OAUTH\n transport=MCPTransport.STREAMABLE_HTTP,\n auth_performer=MCPAuthenticationPerformer.ADMIN,\n db_session=db_session,\n )\n db_session.commit()\n\n # Create MCP tool\n mcp_tool_db = Tool(\n name=\"api_token_tool\",\n display_name=\"API Token Tool\",\n description=\"Tool with API token auth\",\n mcp_server_id=mcp_server.id,\n mcp_input_schema={\n \"type\": \"object\",\n \"properties\": {\"data\": {\"type\": \"string\"}},\n },\n user_id=user.id,\n )\n db_session.add(mcp_tool_db)\n db_session.commit()\n db_session.refresh(mcp_tool_db)\n\n # Create persona\n persona = _create_test_persona_with_mcp_tool(db_session, user, [mcp_tool_db])\n llm = get_default_llm()\n\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=SearchToolConfig(),\n )\n # Verify MCP tool was constructed\n assert mcp_tool_db.id in tool_dict\n constructed_tools = tool_dict[mcp_tool_db.id]\n assert len(constructed_tools) == 1\n mcp_tool = constructed_tools[0]\n assert isinstance(mcp_tool, MCPTool)\n\n # Verify the user's OAuth token was NOT passed (API_TOKEN auth type)\n # API_TOKEN auth should use connection config, not user's login token\n assert mcp_tool._user_oauth_token is None\n\n def test_mcp_tool_run_sets_authorization_header_for_pt_oauth(\n self, db_session: Session\n ) -> None:\n \"\"\"\n Test that MCPTool.run() correctly sets the Authorization header\n when PT_OAUTH is configured.\n \"\"\"\n # Create user with OAuth token\n user = create_test_user(db_session, \"pt_oauth_header_user\")\n user_oauth_token = \"bearer_token_for_mcp_server\"\n\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"google\",\n account_id=\"google_header_user\",\n account_email=user.email,\n access_token=user_oauth_token,\n refresh_token=\"\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n db_session.refresh(user)\n\n # Create MCP server with PT_OAUTH\n mcp_server = create_mcp_server__no_commit(\n owner_email=user.email,\n name=f\"Header Test Server {uuid4().hex[:8]}\",\n description=\"Server for testing Authorization header\",\n server_url=\"http://header-test-server.example.com/mcp\",\n auth_type=MCPAuthenticationType.PT_OAUTH,\n transport=MCPTransport.STREAMABLE_HTTP,\n auth_performer=MCPAuthenticationPerformer.ADMIN,\n db_session=db_session,\n )\n db_session.commit()\n\n # Create MCP tool\n mcp_tool_db = Tool(\n name=\"header_test_tool\",\n display_name=\"Header Test Tool\",\n description=\"Tool to test Authorization header\",\n mcp_server_id=mcp_server.id,\n mcp_input_schema={\n \"type\": \"object\",\n \"properties\": {\"input\": {\"type\": \"string\"}},\n },\n user_id=user.id,\n )\n db_session.add(mcp_tool_db)\n db_session.commit()\n db_session.refresh(mcp_tool_db)\n\n # Create persona\n persona = _create_test_persona_with_mcp_tool(db_session, user, [mcp_tool_db])\n llm = get_default_llm()\n\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=SearchToolConfig(),\n )\n\n # Get the constructed MCPTool\n mcp_tool = tool_dict[mcp_tool_db.id][0]\n assert isinstance(mcp_tool, MCPTool)\n\n # Mock the call_mcp_tool function to capture the headers\n captured_headers: dict[str, str] = {}\n\n mocked_response = {\"result\": \"mocked_response\"}\n\n def mock_call_mcp_tool(\n server_url: str, # noqa: ARG001\n tool_name: str, # noqa: ARG001\n arguments: dict[str, Any], # noqa: ARG001\n connection_headers: dict[str, str],\n transport: MCPTransport, # noqa: ARG001\n auth: Any = None, # noqa: ARG001\n ) -> dict[str, Any]:\n captured_headers.update(connection_headers)\n return mocked_response\n\n with patch(\n \"onyx.tools.tool_implementations.mcp.mcp_tool.call_mcp_tool\",\n side_effect=mock_call_mcp_tool,\n ):\n # Run the tool\n response = mcp_tool.run(\n placement=Placement(turn_index=0, tab_index=0),\n override_kwargs=None,\n input=\"test\",\n )\n print(response.rich_response)\n assert isinstance(response.rich_response, CustomToolCallSummary)\n print(response.rich_response.tool_result)\n assert response.rich_response.tool_result[\"tool_result\"] == mocked_response\n\n # Verify Authorization header was set with the user's OAuth token\n assert \"Authorization\" in captured_headers\n assert captured_headers[\"Authorization\"] == f\"Bearer {user_oauth_token}\"\n\n def test_pt_oauth_works_with_oidc_provider(self, db_session: Session) -> None:\n \"\"\"\n Test that PT_OAUTH works correctly when user logged in via OIDC (not Google).\n\n This is important because OIDC providers (Okta, Auth0, Keycloak, etc.)\n use oauth_name='openid' while Google uses oauth_name='google'.\n The PT_OAUTH code should work with any OAuth provider.\n \"\"\"\n # Create user with OIDC OAuth token (simulating Okta/Auth0/Keycloak login)\n user = create_test_user(db_session, \"oidc_user\")\n # Use a random test token (not a real JWT to avoid pre-commit false positives)\n oidc_access_token = \"oidc_test_token_abc123_not_a_real_jwt_xyz789\"\n\n # OIDC providers use oauth_name='openid' by default\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"openid\", # This is the key difference from Google OAuth\n account_id=\"oidc_user_sub_12345\",\n account_email=user.email,\n access_token=oidc_access_token,\n refresh_token=\"oidc_refresh_token\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n db_session.refresh(user)\n\n # Create MCP server with PT_OAUTH auth type\n mcp_server = create_mcp_server__no_commit(\n owner_email=user.email,\n name=f\"PT_OAUTH OIDC Server {uuid4().hex[:8]}\",\n description=\"MCP server for OIDC pass-through OAuth testing\",\n server_url=\"http://oidc-mcp-server.example.com/mcp\",\n auth_type=MCPAuthenticationType.PT_OAUTH,\n transport=MCPTransport.STREAMABLE_HTTP,\n auth_performer=MCPAuthenticationPerformer.ADMIN,\n db_session=db_session,\n )\n db_session.commit()\n\n # Create MCP tool\n mcp_tool_db = Tool(\n name=\"oidc_mcp_tool\",\n display_name=\"OIDC MCP Tool\",\n description=\"Test MCP tool for OIDC PT_OAUTH\",\n mcp_server_id=mcp_server.id,\n mcp_input_schema={\n \"type\": \"object\",\n \"properties\": {\"query\": {\"type\": \"string\"}},\n },\n user_id=user.id,\n )\n db_session.add(mcp_tool_db)\n db_session.commit()\n db_session.refresh(mcp_tool_db)\n\n # Create persona\n persona = _create_test_persona_with_mcp_tool(db_session, user, [mcp_tool_db])\n llm = get_default_llm()\n\n # Construct tools\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=SearchToolConfig(),\n )\n # Verify MCP tool was constructed\n assert mcp_tool_db.id in tool_dict\n constructed_tools = tool_dict[mcp_tool_db.id]\n assert len(constructed_tools) == 1\n mcp_tool = constructed_tools[0]\n assert isinstance(mcp_tool, MCPTool)\n\n # Verify the OIDC token was passed to the MCPTool\n # (code should work identically for Google OAuth and OIDC)\n assert mcp_tool._user_oauth_token == oidc_access_token\n\n def test_pt_oauth_uses_first_oauth_account(self, db_session: Session) -> None:\n \"\"\"\n Test that PT_OAUTH uses the first OAuth account when user has multiple.\n\n Users might have OAuth accounts from multiple providers (unlikely but possible).\n The code should consistently use the first one.\n \"\"\"\n user = create_test_user(db_session, \"multi_oauth_user\")\n first_token = \"first_oauth_token_123\"\n second_token = \"second_oauth_token_456\"\n\n # Add first OAuth account (Google)\n oauth_account_1 = OAuthAccount(\n user_id=user.id,\n oauth_name=\"google\",\n account_id=\"google_user_123\",\n account_email=user.email,\n access_token=first_token,\n refresh_token=\"\",\n )\n db_session.add(oauth_account_1)\n db_session.commit()\n\n # Add second OAuth account (OIDC)\n oauth_account_2 = OAuthAccount(\n user_id=user.id,\n oauth_name=\"openid\",\n account_id=\"oidc_user_456\",\n account_email=user.email,\n access_token=second_token,\n refresh_token=\"\",\n )\n db_session.add(oauth_account_2)\n db_session.commit()\n db_session.refresh(user)\n\n # Create MCP server and tool\n mcp_server = create_mcp_server__no_commit(\n owner_email=user.email,\n name=f\"Multi OAuth Server {uuid4().hex[:8]}\",\n description=\"MCP server for multi-OAuth testing\",\n server_url=\"http://multi-oauth-server.example.com/mcp\",\n auth_type=MCPAuthenticationType.PT_OAUTH,\n transport=MCPTransport.STREAMABLE_HTTP,\n auth_performer=MCPAuthenticationPerformer.ADMIN,\n db_session=db_session,\n )\n db_session.commit()\n\n mcp_tool_db = Tool(\n name=\"multi_oauth_tool\",\n display_name=\"Multi OAuth Tool\",\n description=\"Test tool\",\n mcp_server_id=mcp_server.id,\n mcp_input_schema={\"type\": \"object\", \"properties\": {}},\n user_id=user.id,\n )\n db_session.add(mcp_tool_db)\n db_session.commit()\n db_session.refresh(mcp_tool_db)\n\n persona = _create_test_persona_with_mcp_tool(db_session, user, [mcp_tool_db])\n llm = get_default_llm()\n\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=SearchToolConfig(),\n )\n\n mcp_tool = tool_dict[mcp_tool_db.id][0]\n assert isinstance(mcp_tool, MCPTool)\n\n # Should use the first OAuth account's token\n assert mcp_tool._user_oauth_token == first_token\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_memory_tool_integration.py", "content": "\"\"\"Tests for MemoryTool integration: registration, construction, and DB persistence.\"\"\"\n\nimport pytest\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.memory import add_memory\nfrom onyx.db.memory import get_memories\nfrom onyx.db.memory import MAX_MEMORIES_PER_USER\nfrom onyx.db.memory import update_memory_at_index\nfrom onyx.db.models import Memory\nfrom onyx.db.models import User\nfrom onyx.tools.tool_implementations.memory.models import MemoryToolResponse\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\n@pytest.fixture()\ndef test_user(db_session: Session): # type: ignore\n \"\"\"Create a test user with use_memories enabled.\"\"\"\n user = create_test_user(db_session, \"memory_test\")\n user.use_memories = True\n db_session.commit()\n db_session.refresh(user)\n return user\n\n\n@pytest.fixture()\ndef test_user_no_memories(db_session: Session): # type: ignore\n \"\"\"Create a test user with use_memories disabled.\"\"\"\n user = create_test_user(db_session, \"memory_test_off\")\n user.use_memories = False\n db_session.commit()\n db_session.refresh(user)\n return user\n\n\nclass TestAddMemory:\n def test_add_memory_creates_row(self, db_session: Session, test_user: User) -> None:\n \"\"\"Verify that add_memory inserts a new Memory row.\"\"\"\n user_id = test_user.id\n memory = add_memory(\n user_id=user_id,\n memory_text=\"User prefers dark mode\",\n db_session=db_session,\n )\n\n assert memory.id is not None\n assert memory.user_id == user_id\n assert memory.memory_text == \"User prefers dark mode\"\n\n # Verify it persists\n fetched = db_session.get(Memory, memory.id)\n assert fetched is not None\n assert fetched.memory_text == \"User prefers dark mode\"\n\n def test_add_multiple_memories(self, db_session: Session, test_user: User) -> None:\n \"\"\"Verify that multiple memories can be added for the same user.\"\"\"\n user_id = test_user.id\n m1 = add_memory(\n user_id=user_id,\n memory_text=\"Favorite color is blue\",\n db_session=db_session,\n )\n m2 = add_memory(\n user_id=user_id,\n memory_text=\"Works in engineering\",\n db_session=db_session,\n )\n\n assert m1.id != m2.id\n assert m1.memory_text == \"Favorite color is blue\"\n assert m2.memory_text == \"Works in engineering\"\n\n\nclass TestUpdateMemoryAtIndex:\n def test_update_memory_at_valid_index(\n self, db_session: Session, test_user: User\n ) -> None:\n \"\"\"Verify that update_memory_at_index updates the correct row.\"\"\"\n user_id = test_user.id\n add_memory(user_id=user_id, memory_text=\"Memory 0\", db_session=db_session)\n add_memory(user_id=user_id, memory_text=\"Memory 1\", db_session=db_session)\n add_memory(user_id=user_id, memory_text=\"Memory 2\", db_session=db_session)\n\n updated = update_memory_at_index(\n user_id=user_id,\n index=1,\n new_text=\"Updated Memory 1\",\n db_session=db_session,\n )\n\n assert updated is not None\n assert updated.memory_text == \"Updated Memory 1\"\n\n def test_update_memory_at_out_of_range_index(\n self, db_session: Session, test_user: User\n ) -> None:\n \"\"\"Verify that out-of-range index returns None.\"\"\"\n user_id = test_user.id\n add_memory(user_id=user_id, memory_text=\"Only memory\", db_session=db_session)\n\n result = update_memory_at_index(\n user_id=user_id,\n index=5,\n new_text=\"Should not update\",\n db_session=db_session,\n )\n\n assert result is None\n\n def test_update_memory_at_negative_index(\n self, db_session: Session, test_user: User\n ) -> None:\n \"\"\"Verify that negative index returns None.\"\"\"\n user_id = test_user.id\n add_memory(user_id=user_id, memory_text=\"Only memory\", db_session=db_session)\n\n result = update_memory_at_index(\n user_id=user_id,\n index=-1,\n new_text=\"Should not update\",\n db_session=db_session,\n )\n\n assert result is None\n\n\nclass TestMemoryToolResponse:\n def test_response_with_add(self) -> None:\n \"\"\"Verify MemoryToolResponse correctly carries add (index_to_replace=None).\"\"\"\n response = MemoryToolResponse(\n memory_text=\"User likes Python\",\n index_to_replace=None,\n )\n assert response.memory_text == \"User likes Python\"\n assert response.index_to_replace is None\n\n def test_response_with_update(self) -> None:\n \"\"\"Verify MemoryToolResponse correctly carries update (index_to_replace=int).\"\"\"\n response = MemoryToolResponse(\n memory_text=\"User likes TypeScript\",\n index_to_replace=2,\n )\n assert response.memory_text == \"User likes TypeScript\"\n assert response.index_to_replace == 2\n\n\nclass TestMemoryCap:\n def test_add_memory_evicts_oldest_when_at_cap(\n self, db_session: Session, test_user: User\n ) -> None:\n \"\"\"When the user has MAX_MEMORIES_PER_USER memories, adding a new one\n should delete the oldest (lowest id) and keep the total at the cap.\"\"\"\n user_id = test_user.id\n\n # Fill up to the cap\n for i in range(MAX_MEMORIES_PER_USER):\n add_memory(\n user_id=user_id,\n memory_text=f\"Memory {i}\",\n db_session=db_session,\n )\n\n rows_before = db_session.scalars(\n Memory.__table__.select().where(Memory.user_id == user_id)\n ).all()\n assert len(rows_before) == MAX_MEMORIES_PER_USER\n\n # Add one more — should evict the oldest\n new_memory = add_memory(\n user_id=user_id,\n memory_text=\"New memory after cap\",\n db_session=db_session,\n )\n\n rows_after = db_session.scalars(\n select(Memory).where(Memory.user_id == user_id).order_by(Memory.id.asc())\n ).all()\n\n assert len(rows_after) == MAX_MEMORIES_PER_USER\n # Oldest (\"Memory 0\") should be gone; \"Memory 1\" is now the oldest\n assert rows_after[0].memory_text == \"Memory 1\"\n # Newest should be the one we just added\n assert rows_after[-1].id == new_memory.id\n assert rows_after[-1].memory_text == \"New memory after cap\"\n\n\nclass TestGetMemoriesWithUserId:\n def test_get_memories_populates_user_id(\n self, db_session: Session, test_user: User\n ) -> None:\n \"\"\"Verify that get_memories populates user_id on the returned context.\"\"\"\n context = get_memories(test_user, db_session)\n assert context.user_id == test_user.id\n\n def test_get_memories_disabled_still_populates_user_id(\n self, db_session: Session, test_user_no_memories: User\n ) -> None:\n \"\"\"Verify that get_memories with use_memories=False still returns a\n fully populated context (user_id, user_info, memories). The\n use_memories flag only controls whether memories are injected into\n the system prompt, not whether the context is fetched.\"\"\"\n # Add a memory for this user so we can verify it's fetched\n add_memory(\n user_id=test_user_no_memories.id,\n memory_text=\"Should still be fetched\",\n db_session=db_session,\n )\n\n context = get_memories(test_user_no_memories, db_session)\n assert context.user_id == test_user_no_memories.id\n assert context.user_info.email == test_user_no_memories.email\n assert len(context.memories) == 1\n assert context.memories[0] == \"Should still be fetched\"\n\n def test_get_memories_disabled_persistence_works(\n self, db_session: Session, test_user_no_memories: User\n ) -> None:\n \"\"\"Verify that add_memory and update_memory_at_index work correctly\n when use_memories=False, since the memory tool should still persist.\"\"\"\n user_id = test_user_no_memories.id\n\n # Add a memory\n memory = add_memory(\n user_id=user_id,\n memory_text=\"Memory with use_memories off\",\n db_session=db_session,\n )\n assert memory.memory_text == \"Memory with use_memories off\"\n\n # Update that memory\n updated = update_memory_at_index(\n user_id=user_id,\n index=0,\n new_text=\"Updated memory with use_memories off\",\n db_session=db_session,\n )\n assert updated is not None\n assert updated.memory_text == \"Updated memory with use_memories off\"\n\n # Verify get_memories returns the updated memory\n context = get_memories(test_user_no_memories, db_session)\n assert len(context.memories) == 1\n assert context.memories[0] == \"Updated memory with use_memories off\"\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_oauth_config_crud.py", "content": "\"\"\"\nTest suite for OAuth Config CRUD operations.\n\nTests the basic CRUD operations for OAuth configurations and user tokens,\nincluding creation, retrieval, updates, deletion, and token management.\n\"\"\"\n\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.models import OAuthConfig\nfrom onyx.db.models import Tool\nfrom onyx.db.oauth_config import create_oauth_config\nfrom onyx.db.oauth_config import delete_oauth_config\nfrom onyx.db.oauth_config import delete_user_oauth_token\nfrom onyx.db.oauth_config import get_oauth_config\nfrom onyx.db.oauth_config import get_oauth_configs\nfrom onyx.db.oauth_config import get_tools_by_oauth_config\nfrom onyx.db.oauth_config import get_user_oauth_token\nfrom onyx.db.oauth_config import update_oauth_config\nfrom onyx.db.oauth_config import upsert_user_oauth_token\nfrom onyx.db.tools import delete_tool__no_commit\nfrom onyx.db.tools import update_tool\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef _create_test_oauth_config(\n db_session: Session,\n name: str | None = None,\n) -> OAuthConfig:\n \"\"\"Helper to create a test OAuth config with unique name\"\"\"\n unique_name = name or f\"Test OAuth Config {uuid4().hex[:8]}\"\n return create_oauth_config(\n name=unique_name,\n authorization_url=\"https://github.com/login/oauth/authorize\",\n token_url=\"https://github.com/login/oauth/access_token\",\n client_id=\"test_client_id\",\n client_secret=\"test_client_secret\",\n scopes=[\"repo\", \"user\"],\n additional_params={\"test_param\": \"test_value\"},\n db_session=db_session,\n )\n\n\ndef _create_test_tool_with_oauth(\n db_session: Session, oauth_config: OAuthConfig\n) -> Tool:\n \"\"\"Helper to create a test tool with OAuth config\"\"\"\n user = create_test_user(db_session, \"tool_owner\")\n tool = Tool(\n name=\"Test Tool\",\n description=\"Test tool with OAuth\",\n openapi_schema={\"openapi\": \"3.0.0\"},\n user_id=user.id,\n oauth_config_id=oauth_config.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n return tool\n\n\nclass TestOAuthConfigCRUD:\n \"\"\"Tests for OAuth configuration CRUD operations\"\"\"\n\n def test_create_oauth_config(self, db_session: Session) -> None:\n \"\"\"Test creating a new OAuth configuration\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n\n assert oauth_config.id is not None\n assert oauth_config.name.startswith(\"Test OAuth Config\")\n assert (\n oauth_config.authorization_url == \"https://github.com/login/oauth/authorize\"\n )\n assert oauth_config.token_url == \"https://github.com/login/oauth/access_token\"\n assert oauth_config.scopes == [\"repo\", \"user\"]\n assert oauth_config.additional_params == {\"test_param\": \"test_value\"}\n assert oauth_config.created_at is not None\n assert oauth_config.updated_at is not None\n\n # Verify encrypted fields are stored (we can't decrypt in tests, but we can check they exist)\n assert oauth_config.client_id is not None\n assert oauth_config.client_secret is not None\n\n def test_get_oauth_config(self, db_session: Session) -> None:\n \"\"\"Test retrieving an OAuth config by ID\"\"\"\n created_config = _create_test_oauth_config(db_session)\n\n retrieved_config = get_oauth_config(created_config.id, db_session)\n\n assert retrieved_config is not None\n assert retrieved_config.id == created_config.id\n assert retrieved_config.name == created_config.name\n\n def test_get_oauth_config_not_found(self, db_session: Session) -> None:\n \"\"\"Test retrieving a non-existent OAuth config returns None\"\"\"\n config = get_oauth_config(99999, db_session)\n assert config is None\n\n def test_get_oauth_configs(self, db_session: Session) -> None:\n \"\"\"Test retrieving all OAuth configurations\"\"\"\n # Create multiple configs with unique names\n config1 = _create_test_oauth_config(db_session)\n config2 = _create_test_oauth_config(db_session)\n\n configs = get_oauth_configs(db_session)\n\n assert len(configs) >= 2\n config_ids = [c.id for c in configs]\n assert config1.id in config_ids\n assert config2.id in config_ids\n\n def test_update_oauth_config(self, db_session: Session) -> None:\n \"\"\"Test updating an OAuth configuration\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n original_name = oauth_config.name\n\n # Update the config with unique name\n new_name = f\"Updated GitHub OAuth {uuid4().hex[:8]}\"\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n name=new_name,\n scopes=[\"repo\", \"user\", \"admin\"],\n )\n\n assert updated_config.id == oauth_config.id\n assert updated_config.name == new_name\n assert updated_config.name != original_name\n assert updated_config.scopes == [\"repo\", \"user\", \"admin\"]\n\n def test_update_oauth_config_preserves_secrets(self, db_session: Session) -> None:\n \"\"\"Test that updating config without providing secrets preserves existing values\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n original_client_id = oauth_config.client_id\n original_client_secret = oauth_config.client_secret\n\n # Update config without providing client_id or client_secret\n new_name = f\"Updated Name {uuid4().hex[:8]}\"\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n name=new_name,\n client_id=None,\n client_secret=None,\n )\n\n # Secrets should be preserved\n assert updated_config.client_id is not None\n assert original_client_id is not None\n assert updated_config.client_id.get_value(\n apply_mask=False\n ) == original_client_id.get_value(apply_mask=False)\n assert updated_config.client_secret is not None\n assert original_client_secret is not None\n assert updated_config.client_secret.get_value(\n apply_mask=False\n ) == original_client_secret.get_value(apply_mask=False)\n # But name should be updated\n assert updated_config.name == new_name\n\n def test_update_oauth_config_not_found(self, db_session: Session) -> None:\n \"\"\"Test updating a non-existent OAuth config raises error\"\"\"\n with pytest.raises(\n ValueError, match=\"OAuth config with id 99999 does not exist\"\n ):\n update_oauth_config(99999, db_session, name=\"New Name\")\n\n def test_update_oauth_config_clear_client_id(self, db_session: Session) -> None:\n \"\"\"Test clearing client_id while preserving client_secret\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n original_client_secret = oauth_config.client_secret\n\n # Clear client_id\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n clear_client_id=True,\n )\n\n # client_id should be cleared (empty string)\n assert updated_config.client_id is not None\n assert updated_config.client_id.get_value(apply_mask=False) == \"\"\n # client_secret should be preserved\n assert updated_config.client_secret is not None\n assert original_client_secret is not None\n assert updated_config.client_secret.get_value(\n apply_mask=False\n ) == original_client_secret.get_value(apply_mask=False)\n\n def test_update_oauth_config_clear_client_secret(self, db_session: Session) -> None:\n \"\"\"Test clearing client_secret while preserving client_id\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n original_client_id = oauth_config.client_id\n\n # Clear client_secret\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n clear_client_secret=True,\n )\n\n # client_secret should be cleared (empty string)\n assert updated_config.client_secret is not None\n assert updated_config.client_secret.get_value(apply_mask=False) == \"\"\n # client_id should be preserved\n assert updated_config.client_id is not None\n assert original_client_id is not None\n assert updated_config.client_id.get_value(\n apply_mask=False\n ) == original_client_id.get_value(apply_mask=False)\n\n def test_update_oauth_config_clear_both_secrets(self, db_session: Session) -> None:\n \"\"\"Test clearing both client_id and client_secret\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n\n # Clear both secrets\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n clear_client_id=True,\n clear_client_secret=True,\n )\n\n # Both should be cleared (empty strings)\n assert updated_config.client_id is not None\n assert updated_config.client_id.get_value(apply_mask=False) == \"\"\n assert updated_config.client_secret is not None\n assert updated_config.client_secret.get_value(apply_mask=False) == \"\"\n\n def test_update_oauth_config_authorization_url(self, db_session: Session) -> None:\n \"\"\"Test updating authorization_url\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n new_auth_url = \"https://example.com/oauth/authorize\"\n\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n authorization_url=new_auth_url,\n )\n\n assert updated_config.authorization_url == new_auth_url\n\n def test_update_oauth_config_token_url(self, db_session: Session) -> None:\n \"\"\"Test updating token_url\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n new_token_url = \"https://example.com/oauth/token\"\n\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n token_url=new_token_url,\n )\n\n assert updated_config.token_url == new_token_url\n\n def test_update_oauth_config_additional_params(self, db_session: Session) -> None:\n \"\"\"Test updating additional_params\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n new_params = {\"access_type\": \"offline\", \"prompt\": \"consent\"}\n\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n additional_params=new_params,\n )\n\n assert updated_config.additional_params == new_params\n\n def test_update_oauth_config_multiple_fields(self, db_session: Session) -> None:\n \"\"\"Test updating multiple fields at once\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n new_name = f\"Updated Config {uuid4().hex[:8]}\"\n new_auth_url = \"https://example.com/oauth/authorize\"\n new_token_url = \"https://example.com/oauth/token\"\n new_scopes = [\"read\", \"write\", \"admin\"]\n new_params = {\"access_type\": \"offline\"}\n new_client_id = \"new_client_id\"\n\n updated_config = update_oauth_config(\n oauth_config.id,\n db_session,\n name=new_name,\n authorization_url=new_auth_url,\n token_url=new_token_url,\n scopes=new_scopes,\n additional_params=new_params,\n client_id=new_client_id,\n )\n\n assert updated_config.name == new_name\n assert updated_config.authorization_url == new_auth_url\n assert updated_config.token_url == new_token_url\n assert updated_config.scopes == new_scopes\n assert updated_config.additional_params == new_params\n assert updated_config.client_id is not None\n assert updated_config.client_id.get_value(apply_mask=False) == new_client_id\n\n def test_delete_oauth_config(self, db_session: Session) -> None:\n \"\"\"Test deleting an OAuth configuration\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n config_id = oauth_config.id\n\n # Delete the config\n delete_oauth_config(config_id, db_session)\n\n # Verify it's deleted\n deleted_config = get_oauth_config(config_id, db_session)\n assert deleted_config is None\n\n def test_delete_oauth_config_not_found(self, db_session: Session) -> None:\n \"\"\"Test deleting a non-existent OAuth config raises error\"\"\"\n with pytest.raises(\n ValueError, match=\"OAuth config with id 99999 does not exist\"\n ):\n delete_oauth_config(99999, db_session)\n\n def test_delete_oauth_config_sets_tool_reference_to_null(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that deleting OAuth config sets tool's oauth_config_id to NULL\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n tool = _create_test_tool_with_oauth(db_session, oauth_config)\n\n assert tool.oauth_config_id == oauth_config.id\n\n # Delete the OAuth config\n delete_oauth_config(oauth_config.id, db_session)\n\n # Refresh tool from database\n db_session.refresh(tool)\n\n # Tool should still exist but oauth_config_id should be NULL\n assert tool.oauth_config_id is None\n\n def test_update_tool_cleans_up_orphaned_oauth_config(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that changing a tool's oauth_config_id deletes the old config if no other tool uses it.\"\"\"\n old_config = _create_test_oauth_config(db_session)\n new_config = _create_test_oauth_config(db_session)\n tool = _create_test_tool_with_oauth(db_session, old_config)\n old_config_id = old_config.id\n\n update_tool(\n tool_id=tool.id,\n name=None,\n description=None,\n openapi_schema=None,\n custom_headers=None,\n user_id=None,\n db_session=db_session,\n passthrough_auth=None,\n oauth_config_id=new_config.id,\n )\n\n assert tool.oauth_config_id == new_config.id\n assert get_oauth_config(old_config_id, db_session) is None\n\n def test_delete_tool_cleans_up_orphaned_oauth_config(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that deleting the last tool referencing an OAuthConfig also deletes the config.\"\"\"\n config = _create_test_oauth_config(db_session)\n tool = _create_test_tool_with_oauth(db_session, config)\n config_id = config.id\n\n delete_tool__no_commit(tool.id, db_session)\n db_session.commit()\n\n assert get_oauth_config(config_id, db_session) is None\n\n def test_update_tool_preserves_shared_oauth_config(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that updating one tool's oauth_config_id preserves the config when another tool still uses it.\"\"\"\n shared_config = _create_test_oauth_config(db_session)\n new_config = _create_test_oauth_config(db_session)\n tool_a = _create_test_tool_with_oauth(db_session, shared_config)\n tool_b = _create_test_tool_with_oauth(db_session, shared_config)\n shared_config_id = shared_config.id\n\n # Move tool_a to a new config; tool_b still references shared_config\n update_tool(\n tool_id=tool_a.id,\n name=None,\n description=None,\n openapi_schema=None,\n custom_headers=None,\n user_id=None,\n db_session=db_session,\n passthrough_auth=None,\n oauth_config_id=new_config.id,\n )\n\n assert tool_a.oauth_config_id == new_config.id\n assert tool_b.oauth_config_id == shared_config_id\n assert get_oauth_config(shared_config_id, db_session) is not None\n\n def test_delete_tool_preserves_shared_oauth_config(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that deleting one tool preserves the config when another tool still uses it.\"\"\"\n shared_config = _create_test_oauth_config(db_session)\n tool_a = _create_test_tool_with_oauth(db_session, shared_config)\n tool_b = _create_test_tool_with_oauth(db_session, shared_config)\n shared_config_id = shared_config.id\n\n delete_tool__no_commit(tool_a.id, db_session)\n db_session.commit()\n\n assert tool_b.oauth_config_id == shared_config_id\n assert get_oauth_config(shared_config_id, db_session) is not None\n\n\nclass TestOAuthUserTokenCRUD:\n \"\"\"Tests for OAuth user token CRUD operations\"\"\"\n\n def test_upsert_user_oauth_token_create(self, db_session: Session) -> None:\n \"\"\"Test creating a new user OAuth token\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n token_data = {\n \"access_token\": \"test_access_token\",\n \"refresh_token\": \"test_refresh_token\",\n \"token_type\": \"Bearer\",\n \"expires_at\": 1234567890,\n }\n\n user_token = upsert_user_oauth_token(\n oauth_config.id, user.id, token_data, db_session\n )\n\n assert user_token.id is not None\n assert user_token.oauth_config_id == oauth_config.id\n assert user_token.user_id == user.id\n assert user_token.token_data is not None\n assert user_token.token_data.get_value(apply_mask=False) == token_data\n assert user_token.created_at is not None\n assert user_token.updated_at is not None\n\n def test_upsert_user_oauth_token_update(self, db_session: Session) -> None:\n \"\"\"Test updating an existing user OAuth token\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create initial token\n initial_token_data = {\n \"access_token\": \"initial_token\",\n \"expires_at\": 1234567890,\n }\n initial_token = upsert_user_oauth_token(\n oauth_config.id, user.id, initial_token_data, db_session\n )\n initial_token_id = initial_token.id\n\n # Update with new token data\n updated_token_data = {\n \"access_token\": \"updated_token\",\n \"expires_at\": 9876543210,\n }\n updated_token = upsert_user_oauth_token(\n oauth_config.id, user.id, updated_token_data, db_session\n )\n\n # Should be the same token record (updated, not inserted)\n assert updated_token.id == initial_token_id\n assert updated_token.token_data is not None\n assert (\n updated_token.token_data.get_value(apply_mask=False) == updated_token_data\n )\n assert (\n updated_token.token_data.get_value(apply_mask=False) != initial_token_data\n )\n\n def test_get_user_oauth_token(self, db_session: Session) -> None:\n \"\"\"Test retrieving a user's OAuth token\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n token_data = {\"access_token\": \"test_token\"}\n created_token = upsert_user_oauth_token(\n oauth_config.id, user.id, token_data, db_session\n )\n\n retrieved_token = get_user_oauth_token(oauth_config.id, user.id, db_session)\n\n assert retrieved_token is not None\n assert retrieved_token.id == created_token.id\n assert retrieved_token.token_data is not None\n assert retrieved_token.token_data.get_value(apply_mask=False) == token_data\n\n def test_get_user_oauth_token_not_found(self, db_session: Session) -> None:\n \"\"\"Test retrieving a non-existent user token returns None\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n token = get_user_oauth_token(oauth_config.id, user.id, db_session)\n assert token is None\n\n def test_delete_user_oauth_token(self, db_session: Session) -> None:\n \"\"\"Test deleting a user's OAuth token\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n token_data = {\"access_token\": \"test_token\"}\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n # Delete the token\n delete_user_oauth_token(oauth_config.id, user.id, db_session)\n\n # Verify it's deleted\n deleted_token = get_user_oauth_token(oauth_config.id, user.id, db_session)\n assert deleted_token is None\n\n def test_delete_user_oauth_token_not_found(self, db_session: Session) -> None:\n \"\"\"Test deleting a non-existent user token raises error\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n with pytest.raises(\n ValueError,\n match=f\"OAuth token for user {user.id} and config {oauth_config.id} does not exist\",\n ):\n delete_user_oauth_token(oauth_config.id, user.id, db_session)\n\n def test_unique_constraint_on_user_config(self, db_session: Session) -> None:\n \"\"\"Test that unique constraint prevents duplicate tokens per user per config\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create first token\n token_data1 = {\"access_token\": \"token1\"}\n upsert_user_oauth_token(oauth_config.id, user.id, token_data1, db_session)\n\n # Try to manually insert a duplicate (should fail at DB level)\n # But upsert should work fine (updates instead of inserting)\n token_data2 = {\"access_token\": \"token2\"}\n updated_token = upsert_user_oauth_token(\n oauth_config.id, user.id, token_data2, db_session\n )\n\n # Should only be one token\n retrieved_token = get_user_oauth_token(oauth_config.id, user.id, db_session)\n assert retrieved_token is not None\n assert retrieved_token.id == updated_token.id\n assert retrieved_token.token_data is not None\n assert retrieved_token.token_data.get_value(apply_mask=False) == token_data2\n\n def test_cascade_delete_user_tokens_on_config_deletion(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that deleting OAuth config cascades to user tokens\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user1 = create_test_user(db_session, \"user1\")\n user2 = create_test_user(db_session, \"user2\")\n\n # Create tokens for both users\n upsert_user_oauth_token(\n oauth_config.id, user1.id, {\"access_token\": \"token1\"}, db_session\n )\n upsert_user_oauth_token(\n oauth_config.id, user2.id, {\"access_token\": \"token2\"}, db_session\n )\n\n # Delete the OAuth config\n delete_oauth_config(oauth_config.id, db_session)\n\n # User tokens should be deleted\n token1 = get_user_oauth_token(oauth_config.id, user1.id, db_session)\n token2 = get_user_oauth_token(oauth_config.id, user2.id, db_session)\n assert token1 is None\n assert token2 is None\n\n\nclass TestOAuthHelperOperations:\n \"\"\"Tests for OAuth helper operations\"\"\"\n\n def test_get_tools_by_oauth_config(self, db_session: Session) -> None:\n \"\"\"Test retrieving tools that use a specific OAuth config\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n\n # Create multiple tools using this config\n tool1 = _create_test_tool_with_oauth(db_session, oauth_config)\n tool2 = _create_test_tool_with_oauth(db_session, oauth_config)\n\n # Create another tool without OAuth\n user = create_test_user(db_session, \"other_user\")\n tool3 = Tool(\n name=\"Tool without OAuth\",\n description=\"No OAuth config\",\n openapi_schema={\"openapi\": \"3.0.0\"},\n user_id=user.id,\n )\n db_session.add(tool3)\n db_session.commit()\n\n # Get tools by OAuth config\n tools = get_tools_by_oauth_config(oauth_config.id, db_session)\n\n assert len(tools) == 2\n tool_ids = [t.id for t in tools]\n assert tool1.id in tool_ids\n assert tool2.id in tool_ids\n assert tool3.id not in tool_ids\n\n def test_get_tools_by_oauth_config_empty(self, db_session: Session) -> None:\n \"\"\"Test retrieving tools for config with no associated tools\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n\n tools = get_tools_by_oauth_config(oauth_config.id, db_session)\n\n assert len(tools) == 0\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_oauth_token_manager.py", "content": "\"\"\"\nTest suite for OAuthTokenManager.\n\nTests the OAuth token management functionality including token validation,\nrefresh, expiration checking, and authorization URL building.\nAll HTTP requests to external OAuth providers are mocked.\n\"\"\"\n\nimport time\nfrom unittest.mock import Mock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom requests import HTTPError\nfrom requests import Response\nfrom sqlalchemy.orm import Session\n\nfrom onyx.auth.oauth_token_manager import OAuthTokenManager\nfrom onyx.db.models import OAuthConfig\nfrom onyx.db.oauth_config import create_oauth_config\nfrom onyx.db.oauth_config import upsert_user_oauth_token\nfrom onyx.utils.sensitive import SensitiveValue\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\ndef _create_test_oauth_config(db_session: Session) -> OAuthConfig:\n \"\"\"Helper to create a test OAuth config\"\"\"\n return create_oauth_config(\n name=f\"Test OAuth Config {uuid4().hex[:8]}\",\n authorization_url=\"https://github.com/login/oauth/authorize\",\n token_url=\"https://github.com/login/oauth/access_token\",\n client_id=\"test_client_id\",\n client_secret=\"test_client_secret\",\n scopes=[\"repo\", \"user\"],\n additional_params=None,\n db_session=db_session,\n )\n\n\nclass TestOAuthTokenManagerValidation:\n \"\"\"Tests for token validation and retrieval\"\"\"\n\n def test_get_valid_access_token_with_valid_token(self, db_session: Session) -> None:\n \"\"\"Test getting a valid access token that hasn't expired\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create a non-expired token\n future_timestamp = int(time.time()) + 3600 # Expires in 1 hour\n token_data = {\n \"access_token\": \"valid_token\",\n \"refresh_token\": \"refresh_token\",\n \"expires_at\": future_timestamp,\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n # Get the token\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n access_token = manager.get_valid_access_token()\n\n assert access_token == \"valid_token\"\n\n def test_get_valid_access_token_no_token_exists(self, db_session: Session) -> None:\n \"\"\"Test getting access token when no token exists returns None\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n access_token = manager.get_valid_access_token()\n\n assert access_token is None\n\n def test_get_valid_access_token_no_expiration(self, db_session: Session) -> None:\n \"\"\"Test getting access token without expiration data (assumes valid)\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create token without expiration\n token_data = {\n \"access_token\": \"token_without_expiry\",\n \"token_type\": \"Bearer\",\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n access_token = manager.get_valid_access_token()\n\n assert access_token == \"token_without_expiry\"\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_get_valid_access_token_with_expired_token_refreshes(\n self, mock_post: Mock, db_session: Session\n ) -> None:\n \"\"\"Test that expired token triggers automatic refresh\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create an expired token\n past_timestamp = int(time.time()) - 100 # Expired 100 seconds ago\n token_data = {\n \"access_token\": \"expired_token\",\n \"refresh_token\": \"refresh_token\",\n \"expires_at\": past_timestamp,\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n # Mock the refresh token response\n mock_response = Mock(spec=Response)\n mock_response.json.return_value = {\n \"access_token\": \"new_access_token\",\n \"refresh_token\": \"new_refresh_token\",\n \"expires_in\": 3600,\n }\n mock_response.raise_for_status = Mock()\n mock_post.return_value = mock_response\n\n # Get the token (should trigger refresh)\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n access_token = manager.get_valid_access_token()\n\n assert access_token == \"new_access_token\"\n # Verify refresh endpoint was called\n mock_post.assert_called_once()\n call_args = mock_post.call_args\n assert call_args[0][0] == oauth_config.token_url\n assert call_args[1][\"data\"][\"grant_type\"] == \"refresh_token\"\n assert call_args[1][\"data\"][\"refresh_token\"] == \"refresh_token\"\n\n def test_get_valid_access_token_expired_no_refresh_token(\n self, db_session: Session\n ) -> None:\n \"\"\"Test that expired token without refresh_token returns None\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create an expired token without refresh_token\n past_timestamp = int(time.time()) - 100\n token_data = {\n \"access_token\": \"expired_token\",\n \"expires_at\": past_timestamp,\n # No refresh_token\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n access_token = manager.get_valid_access_token()\n\n assert access_token is None\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_get_valid_access_token_refresh_fails(\n self, mock_post: Mock, db_session: Session\n ) -> None:\n \"\"\"Test that failed refresh returns None\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create an expired token\n past_timestamp = int(time.time()) - 100\n token_data = {\n \"access_token\": \"expired_token\",\n \"refresh_token\": \"refresh_token\",\n \"expires_at\": past_timestamp,\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n # Mock the refresh to fail\n mock_post.side_effect = HTTPError(\"Token refresh failed\")\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n access_token = manager.get_valid_access_token()\n\n assert access_token is None\n\n\nclass TestOAuthTokenManagerRefresh:\n \"\"\"Tests for token refresh functionality\"\"\"\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_refresh_token_success(self, mock_post: Mock, db_session: Session) -> None:\n \"\"\"Test successful token refresh\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create initial token\n token_data = {\n \"access_token\": \"old_token\",\n \"refresh_token\": \"old_refresh\",\n \"expires_at\": int(time.time()) - 100,\n }\n user_token = upsert_user_oauth_token(\n oauth_config.id, user.id, token_data, db_session\n )\n\n # Mock successful refresh\n new_expires_in = 3600\n mock_response = Mock(spec=Response)\n mock_response.json.return_value = {\n \"access_token\": \"new_token\",\n \"refresh_token\": \"new_refresh\",\n \"expires_in\": new_expires_in,\n }\n mock_response.raise_for_status = Mock()\n mock_post.return_value = mock_response\n\n # Refresh the token\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n new_access_token = manager.refresh_token(user_token)\n\n assert new_access_token == \"new_token\"\n\n # Verify token was updated in DB\n db_session.refresh(user_token)\n assert user_token.token_data is not None\n token_data = user_token.token_data.get_value(apply_mask=False)\n assert token_data[\"access_token\"] == \"new_token\"\n assert token_data[\"refresh_token\"] == \"new_refresh\"\n assert \"expires_at\" in token_data\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_refresh_token_preserves_refresh_token(\n self, mock_post: Mock, db_session: Session\n ) -> None:\n \"\"\"Test that refresh preserves old refresh_token if provider doesn't return new one\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create initial token\n token_data = {\n \"access_token\": \"old_token\",\n \"refresh_token\": \"old_refresh\",\n \"expires_at\": int(time.time()) - 100,\n }\n user_token = upsert_user_oauth_token(\n oauth_config.id, user.id, token_data, db_session\n )\n\n # Mock refresh response WITHOUT refresh_token\n mock_response = Mock(spec=Response)\n mock_response.json.return_value = {\n \"access_token\": \"new_token\",\n \"expires_in\": 3600,\n # No refresh_token returned\n }\n mock_response.raise_for_status = Mock()\n mock_post.return_value = mock_response\n\n # Refresh the token\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n manager.refresh_token(user_token)\n\n # Verify old refresh_token was preserved\n db_session.refresh(user_token)\n assert user_token.token_data is not None\n token_data = user_token.token_data.get_value(apply_mask=False)\n assert token_data[\"refresh_token\"] == \"old_refresh\"\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_refresh_token_http_error(\n self, mock_post: Mock, db_session: Session\n ) -> None:\n \"\"\"Test that HTTP error during refresh is raised\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n token_data = {\n \"access_token\": \"old_token\",\n \"refresh_token\": \"old_refresh\",\n \"expires_at\": int(time.time()) - 100,\n }\n user_token = upsert_user_oauth_token(\n oauth_config.id, user.id, token_data, db_session\n )\n\n # Mock HTTP error\n mock_response = Mock(spec=Response)\n mock_response.raise_for_status.side_effect = HTTPError(\"Invalid refresh token\")\n mock_post.return_value = mock_response\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n\n with pytest.raises(HTTPError):\n manager.refresh_token(user_token)\n\n\nclass TestOAuthTokenManagerExpiration:\n \"\"\"Tests for token expiration checking\"\"\"\n\n def test_is_token_expired_with_valid_token(self, db_session: Session) -> None:\n \"\"\"Test that non-expired token is detected as valid\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n\n # Token expires in 2 hours (well beyond 60 second buffer)\n token_data = {\"expires_at\": int(time.time()) + 7200}\n\n assert manager.is_token_expired(token_data) is False\n\n def test_is_token_expired_with_expired_token(self, db_session: Session) -> None:\n \"\"\"Test that expired token is detected\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n\n # Token expired 1 hour ago\n token_data = {\"expires_at\": int(time.time()) - 3600}\n\n assert manager.is_token_expired(token_data) is True\n\n def test_is_token_expired_with_buffer_zone(self, db_session: Session) -> None:\n \"\"\"Test that token within 60 second buffer is considered expired\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n\n # Token expires in 30 seconds (within 60 second buffer)\n token_data = {\"expires_at\": int(time.time()) + 30}\n\n assert manager.is_token_expired(token_data) is True\n\n def test_is_token_expired_no_expiration_data(self, db_session: Session) -> None:\n \"\"\"Test that token without expiration is considered valid\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n\n # Token without expires_at\n token_data = {\"access_token\": \"some_token\"}\n\n assert manager.is_token_expired(token_data) is False\n\n\nclass TestOAuthTokenManagerCodeExchange:\n \"\"\"Tests for authorization code exchange\"\"\"\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_exchange_code_for_token_success(\n self, mock_post: Mock, db_session: Session\n ) -> None:\n \"\"\"Test successful code exchange\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Mock successful token exchange\n mock_response = Mock(spec=Response)\n mock_response.json.return_value = {\n \"access_token\": \"new_access_token\",\n \"refresh_token\": \"new_refresh_token\",\n \"token_type\": \"Bearer\",\n \"expires_in\": 3600,\n \"scope\": \"repo user\",\n }\n mock_response.raise_for_status = Mock()\n mock_post.return_value = mock_response\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n token_data = manager.exchange_code_for_token(\n code=\"auth_code_123\", redirect_uri=\"https://example.com/callback\"\n )\n\n assert token_data[\"access_token\"] == \"new_access_token\"\n assert token_data[\"refresh_token\"] == \"new_refresh_token\"\n assert \"expires_at\" in token_data\n\n # Verify correct parameters were sent\n mock_post.assert_called_once()\n call_args = mock_post.call_args\n assert call_args[0][0] == oauth_config.token_url\n assert call_args[1][\"data\"][\"grant_type\"] == \"authorization_code\"\n assert call_args[1][\"data\"][\"code\"] == \"auth_code_123\"\n assert oauth_config.client_id is not None\n assert oauth_config.client_secret is not None\n assert call_args[1][\"data\"][\"client_id\"] == oauth_config.client_id.get_value(\n apply_mask=False\n )\n assert call_args[1][\"data\"][\n \"client_secret\"\n ] == oauth_config.client_secret.get_value(apply_mask=False)\n assert call_args[1][\"data\"][\"redirect_uri\"] == \"https://example.com/callback\"\n\n @patch(\"onyx.auth.oauth_token_manager.requests.post\")\n def test_exchange_code_for_token_http_error(\n self, mock_post: Mock, db_session: Session\n ) -> None:\n \"\"\"Test that HTTP error during code exchange is raised\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Mock HTTP error\n mock_response = Mock(spec=Response)\n mock_response.raise_for_status.side_effect = HTTPError(\"Invalid code\")\n mock_post.return_value = mock_response\n\n manager = OAuthTokenManager(oauth_config, user.id, db_session)\n\n with pytest.raises(HTTPError):\n manager.exchange_code_for_token(\n code=\"invalid_code\", redirect_uri=\"https://example.com/callback\"\n )\n\n\nclass TestOAuthTokenManagerURLBuilding:\n \"\"\"Tests for authorization URL building\"\"\"\n\n def test_build_authorization_url_basic(self, db_session: Session) -> None:\n \"\"\"Test building basic authorization URL\"\"\"\n oauth_config = _create_test_oauth_config(db_session)\n\n url = OAuthTokenManager.build_authorization_url(\n oauth_config=oauth_config,\n redirect_uri=\"https://example.com/callback\",\n state=\"random_state_123\",\n )\n\n assert url.startswith(oauth_config.authorization_url)\n assert \"client_id=test_client_id\" in url\n assert \"redirect_uri=https%3A%2F%2Fexample.com%2Fcallback\" in url\n assert \"response_type=code\" in url\n assert \"state=random_state_123\" in url\n # Check scopes are included\n assert \"scope=repo+user\" in url\n\n def test_build_authorization_url_with_additional_params(\n self, db_session: Session\n ) -> None:\n \"\"\"Test building URL with additional provider-specific parameters\"\"\"\n oauth_config = create_oauth_config(\n name=f\"Test OAuth {uuid4().hex[:8]}\",\n authorization_url=\"https://accounts.google.com/o/oauth2/v2/auth\",\n token_url=\"https://oauth2.googleapis.com/token\",\n client_id=\"google_client_id\",\n client_secret=\"google_client_secret\",\n scopes=[\"email\", \"profile\"],\n additional_params={\"access_type\": \"offline\", \"prompt\": \"consent\"},\n db_session=db_session,\n )\n\n url = OAuthTokenManager.build_authorization_url(\n oauth_config=oauth_config,\n redirect_uri=\"https://example.com/callback\",\n state=\"state_456\",\n )\n\n assert \"access_type=offline\" in url\n assert \"prompt=consent\" in url\n assert \"scope=email+profile\" in url\n\n def test_build_authorization_url_no_scopes(self, db_session: Session) -> None:\n \"\"\"Test building URL when no scopes are configured\"\"\"\n oauth_config = create_oauth_config(\n name=f\"Test OAuth {uuid4().hex[:8]}\",\n authorization_url=\"https://oauth.example.com/authorize\",\n token_url=\"https://oauth.example.com/token\",\n client_id=\"simple_client_id\",\n client_secret=\"simple_client_secret\",\n scopes=None, # No scopes\n additional_params=None,\n db_session=db_session,\n )\n\n url = OAuthTokenManager.build_authorization_url(\n oauth_config=oauth_config,\n redirect_uri=\"https://example.com/callback\",\n state=\"state_789\",\n )\n\n # Should not include scope parameter\n assert \"scope=\" not in url\n assert \"client_id=simple_client_id\" in url\n\n def test_build_authorization_url_with_existing_query_params(\n self, db_session: Session\n ) -> None:\n \"\"\"Test building URL when authorization_url already has query parameters\"\"\"\n oauth_config = create_oauth_config(\n name=f\"Test OAuth {uuid4().hex[:8]}\",\n authorization_url=\"https://oauth.example.com/authorize?foo=bar\",\n token_url=\"https://oauth.example.com/token\",\n client_id=\"custom_client_id\",\n client_secret=\"custom_client_secret\",\n scopes=[\"read\"],\n additional_params=None,\n db_session=db_session,\n )\n\n url = OAuthTokenManager.build_authorization_url(\n oauth_config=oauth_config,\n redirect_uri=\"https://example.com/callback\",\n state=\"state_xyz\",\n )\n\n # Should use & instead of ? since URL already has query params\n assert \"foo=bar&\" in url or \"?foo=bar\" in url\n assert \"client_id=custom_client_id\" in url\n\n\nclass TestUnwrapSensitiveStr:\n \"\"\"Tests for _unwrap_sensitive_str static method\"\"\"\n\n def test_unwrap_sensitive_str(self) -> None:\n \"\"\"Test that both SensitiveValue and plain str inputs are handled\"\"\"\n # SensitiveValue input\n sensitive = SensitiveValue[str](\n encrypted_bytes=b\"test_client_id\",\n decrypt_fn=lambda b: b.decode(),\n )\n assert OAuthTokenManager._unwrap_sensitive_str(sensitive) == \"test_client_id\"\n\n # Plain str input\n assert OAuthTokenManager._unwrap_sensitive_str(\"plain_string\") == \"plain_string\"\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_oauth_tool_integration.py", "content": "\"\"\"\nTest suite for OAuth integration in tool_constructor.\n\nTests the priority logic for OAuth tokens when constructing custom tools:\n1. Priority 1: OAuth config (per-tool OAuth)\n2. Priority 2: Passthrough auth (user's login OAuth token)\n\nAll external HTTP calls are mocked, but Postgres and Redis are running.\n\"\"\"\n\nimport queue\nfrom typing import Any\nfrom unittest.mock import Mock\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.chat.emitter import Emitter\nfrom onyx.db.models import OAuthAccount\nfrom onyx.db.models import OAuthConfig\nfrom onyx.db.models import Persona\nfrom onyx.db.models import Tool\nfrom onyx.db.models import User\nfrom onyx.db.oauth_config import create_oauth_config\nfrom onyx.db.oauth_config import upsert_user_oauth_token\nfrom onyx.llm.factory import get_default_llm\nfrom onyx.tools.tool_constructor import construct_tools\nfrom onyx.tools.tool_constructor import SearchToolConfig\nfrom onyx.tools.tool_implementations.custom.custom_tool import CustomTool\nfrom tests.external_dependency_unit.answer.conftest import ensure_default_llm_provider\nfrom tests.external_dependency_unit.conftest import create_test_user\n\n\n# Simple OpenAPI schema for testing\nSIMPLE_OPENAPI_SCHEMA: dict[str, Any] = {\n \"openapi\": \"3.0.0\",\n \"info\": {\"title\": \"Test API\", \"version\": \"1.0.0\"},\n \"servers\": [{\"url\": \"https://api.example.com\"}],\n \"paths\": {\n \"/test\": {\n \"get\": {\n \"operationId\": \"test_operation\",\n \"summary\": \"Test operation\",\n \"description\": \"A test operation\",\n \"responses\": {\"200\": {\"description\": \"Success\"}},\n }\n }\n },\n}\n\n\ndef _create_test_persona(db_session: Session, user: User, tools: list[Tool]) -> Persona:\n \"\"\"Helper to create a test persona with the given tools\"\"\"\n # Create persona with prompts directly on it\n persona = Persona(\n name=f\"Test Persona {uuid4().hex[:8]}\",\n description=\"Test persona\",\n system_prompt=\"You are a helpful assistant\",\n task_prompt=\"Answer the user's question\",\n tools=tools,\n document_sets=[],\n users=[user],\n groups=[],\n is_listed=True,\n is_public=True,\n display_priority=None,\n starter_messages=None,\n deleted=False,\n )\n db_session.add(persona)\n db_session.commit()\n db_session.refresh(persona)\n return persona\n\n\ndef _create_test_oauth_config(\n db_session: Session, name: str | None = None\n) -> OAuthConfig:\n \"\"\"Helper to create a test OAuth config\"\"\"\n return create_oauth_config(\n name=name or f\"Test OAuth Config {uuid4().hex[:8]}\",\n authorization_url=\"https://github.com/login/oauth/authorize\",\n token_url=\"https://github.com/login/oauth/access_token\",\n client_id=\"test_client_id\",\n client_secret=\"test_client_secret\",\n scopes=[\"repo\", \"user\"],\n additional_params=None,\n db_session=db_session,\n )\n\n\ndef _get_authorization_header(headers: dict[str, str]) -> str | None:\n \"\"\"\n Helper to extract authorization header from headers dict.\n Checks both 'authorization' and 'Authorization' keys.\n\n Returns:\n The authorization header value, or None if not present.\n \"\"\"\n return headers.get(\"authorization\") or headers.get(\"Authorization\")\n\n\ndef _assert_has_authorization_header(headers: dict[str, str]) -> None:\n \"\"\"Assert that headers contain an authorization header (any case).\"\"\"\n assert (\n \"authorization\" in headers or \"Authorization\" in headers\n ), \"Expected authorization header to be present\"\n\n\ndef _assert_no_authorization_header(headers: dict[str, str]) -> None:\n \"\"\"Assert that headers do NOT contain an authorization header.\"\"\"\n assert (\n \"authorization\" not in headers and \"Authorization\" not in headers\n ), \"Expected no authorization header\"\n\n\nclass TestOAuthToolIntegrationPriority:\n \"\"\"Tests for OAuth token priority logic in tool_constructor\"\"\"\n\n @pytest.fixture(autouse=True)\n def setup_llm_provider(self, db_session: Session) -> None:\n \"\"\"Ensure default LLM provider is set up for each test.\"\"\"\n ensure_default_llm_provider(db_session)\n\n def test_oauth_config_priority_over_passthrough(self, db_session: Session) -> None:\n \"\"\"\n Test that oauth_config_id takes priority over passthrough_auth.\n When both are set, the tool should use the OAuth config token.\n \"\"\"\n # Create user with login OAuth token\n user = create_test_user(db_session, \"oauth_user\")\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"github\",\n account_id=\"github_user_123\",\n account_email=user.email,\n access_token=\"user_login_token_12345\",\n refresh_token=\"\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n # Refresh user to load oauth_accounts relationship\n db_session.refresh(user)\n\n # Create OAuth config with a valid token\n oauth_config = _create_test_oauth_config(db_session)\n token_data = {\n \"access_token\": \"oauth_config_token_67890\",\n \"token_type\": \"Bearer\",\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n # Create tool with BOTH oauth_config_id and passthrough_auth set\n tool = Tool(\n name=\"test_tool\",\n description=\"Test tool\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=oauth_config.id, # Priority 1\n passthrough_auth=True, # Priority 2 - should be ignored\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona and chat session\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Construct tools\n search_tool_config = SearchToolConfig()\n\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n search_tool_config=search_tool_config,\n )\n\n # Verify tool was constructed\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify the OAuth config token is used (Priority 1), NOT passthrough token\n _assert_has_authorization_header(custom_tool.headers)\n auth_header = _get_authorization_header(custom_tool.headers)\n assert auth_header == \"Bearer oauth_config_token_67890\"\n\n def test_passthrough_auth_when_no_oauth_config(self, db_session: Session) -> None:\n \"\"\"\n Test that passthrough_auth works when oauth_config_id is not set.\n \"\"\"\n # Create user with login OAuth token\n user = create_test_user(db_session, \"oauth_user\")\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"google\",\n account_id=\"google_user_456\",\n account_email=user.email,\n access_token=\"user_passthrough_token_99999\",\n refresh_token=\"\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n # Refresh user to load oauth_accounts relationship\n db_session.refresh(user)\n\n # Create tool with only passthrough_auth set (no oauth_config_id)\n tool = Tool(\n name=\"test_tool_passthrough\",\n description=\"Test tool with passthrough\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=None, # No OAuth config\n passthrough_auth=True, # Should use user's login token\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Construct tools\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n )\n\n # Verify tool was constructed\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify the passthrough token is used\n _assert_has_authorization_header(custom_tool.headers)\n auth_header = _get_authorization_header(custom_tool.headers)\n assert auth_header == \"Bearer user_passthrough_token_99999\"\n\n def test_oauth_config_without_valid_token_logs_warning(\n self, db_session: Session, caplog: pytest.LogCaptureFixture\n ) -> None:\n \"\"\"\n Test that when oauth_config_id is set but no valid token exists,\n a warning is logged and the tool has no auth header.\n \"\"\"\n # Create user (no OAuth account)\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create OAuth config but DO NOT create a token for the user\n oauth_config = _create_test_oauth_config(db_session)\n\n # Create tool with oauth_config_id but user has no token\n tool = Tool(\n name=\"test_tool_no_token\",\n description=\"Test tool without token\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=oauth_config.id,\n passthrough_auth=False,\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Construct tools\n with caplog.at_level(\"WARNING\"):\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n )\n\n # Verify warning was logged\n assert any(\n \"No valid OAuth token found for tool\" in record.message\n for record in caplog.records\n )\n assert any(str(oauth_config.id) in record.message for record in caplog.records)\n\n # Verify tool was constructed but has no authorization header\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify NO authorization header is present\n _assert_no_authorization_header(custom_tool.headers)\n\n def test_no_auth_when_both_disabled(self, db_session: Session) -> None:\n \"\"\"\n Test that when neither oauth_config_id nor passthrough_auth is set,\n the tool has no authorization header.\n \"\"\"\n # Create user with OAuth account (but tool won't use it)\n user = create_test_user(db_session, \"oauth_user\")\n oauth_account = OAuthAccount(\n user_id=user.id,\n oauth_name=\"github\",\n account_id=\"github_user_789\",\n account_email=user.email,\n access_token=\"unused_token\",\n refresh_token=\"\",\n )\n db_session.add(oauth_account)\n db_session.commit()\n\n # Create tool with neither oauth_config_id nor passthrough_auth\n tool = Tool(\n name=\"test_tool_no_auth\",\n description=\"Test tool without auth\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=None,\n passthrough_auth=False,\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Construct tools\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n )\n\n # Verify tool was constructed\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify NO authorization header\n _assert_no_authorization_header(custom_tool.headers)\n\n def test_oauth_config_with_expired_token_refreshes(\n self, db_session: Session\n ) -> None:\n \"\"\"\n Test that expired OAuth config tokens are automatically refreshed.\n \"\"\"\n import time\n\n # Create user\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create OAuth config with expired token\n oauth_config = _create_test_oauth_config(db_session)\n expired_token_data = {\n \"access_token\": \"expired_token\",\n \"refresh_token\": \"refresh_token_12345\",\n \"expires_at\": int(time.time()) - 100, # Expired 100 seconds ago\n }\n upsert_user_oauth_token(\n oauth_config.id, user.id, expired_token_data, db_session\n )\n\n # Create tool with oauth_config_id\n tool = Tool(\n name=\"test_tool_refresh\",\n description=\"Test tool with token refresh\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=oauth_config.id,\n passthrough_auth=False,\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Mock the token refresh response\n mock_response = Mock()\n mock_response.json.return_value = {\n \"access_token\": \"refreshed_token_67890\",\n \"refresh_token\": \"refresh_token_12345\",\n \"expires_in\": 3600,\n \"token_type\": \"Bearer\",\n }\n mock_response.raise_for_status = Mock()\n\n with patch(\"onyx.auth.oauth_token_manager.requests.post\") as mock_post:\n mock_post.return_value = mock_response\n\n # Construct tools\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n )\n\n # Verify token refresh was called\n mock_post.assert_called_once()\n call_args = mock_post.call_args\n assert call_args[0][0] == oauth_config.token_url\n assert call_args[1][\"data\"][\"grant_type\"] == \"refresh_token\"\n assert call_args[1][\"data\"][\"refresh_token\"] == \"refresh_token_12345\"\n\n # Verify tool was constructed with refreshed token\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify the refreshed token is used\n _assert_has_authorization_header(custom_tool.headers)\n auth_header = _get_authorization_header(custom_tool.headers)\n assert auth_header == \"Bearer refreshed_token_67890\"\n\n def test_custom_headers_combined_with_oauth_token(\n self, db_session: Session\n ) -> None:\n \"\"\"\n Test that custom headers are properly combined with OAuth token.\n The OAuth Authorization header should be added to existing custom headers.\n \"\"\"\n # Create user\n user = create_test_user(db_session, \"oauth_user\")\n\n # Create OAuth config with token\n oauth_config = _create_test_oauth_config(db_session)\n token_data = {\n \"access_token\": \"oauth_token_abc123\",\n \"token_type\": \"Bearer\",\n }\n upsert_user_oauth_token(oauth_config.id, user.id, token_data, db_session)\n\n # Create tool with oauth_config_id AND custom headers\n tool = Tool(\n name=\"test_tool_combined\",\n description=\"Test tool with custom headers and OAuth\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=oauth_config.id,\n custom_headers=[\n {\"key\": \"X-Custom-Header\", \"value\": \"custom-value\"},\n {\"key\": \"X-API-Key\", \"value\": \"api-key-123\"},\n ],\n passthrough_auth=False,\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Construct tools\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n )\n\n # Verify tool was constructed\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify both OAuth token AND custom headers are present\n _assert_has_authorization_header(custom_tool.headers)\n auth_header = _get_authorization_header(custom_tool.headers)\n assert auth_header == \"Bearer oauth_token_abc123\"\n\n # Headers are capitalized by the tool\n assert \"X-Custom-Header\" in custom_tool.headers\n assert custom_tool.headers[\"X-Custom-Header\"] == \"custom-value\"\n assert \"X-API-Key\" in custom_tool.headers\n assert custom_tool.headers[\"X-API-Key\"] == \"api-key-123\"\n\n def test_passthrough_auth_without_user_oauth_account(\n self, db_session: Session\n ) -> None:\n \"\"\"\n Test that passthrough_auth handles gracefully when user has no OAuth account.\n \"\"\"\n # Create user WITHOUT OAuth account\n user = create_test_user(db_session, \"no_oauth_user\")\n\n # Create tool with passthrough_auth\n tool = Tool(\n name=\"test_tool_no_account\",\n description=\"Test tool passthrough without account\",\n openapi_schema=SIMPLE_OPENAPI_SCHEMA,\n oauth_config_id=None,\n passthrough_auth=True,\n user_id=user.id,\n )\n db_session.add(tool)\n db_session.commit()\n db_session.refresh(tool)\n\n # Create persona\n persona = _create_test_persona(db_session, user, [tool])\n llm = get_default_llm()\n\n # Construct tools\n tool_dict = construct_tools(\n persona=persona,\n db_session=db_session,\n emitter=Emitter(merged_queue=queue.Queue()),\n user=user,\n llm=llm,\n )\n\n # Verify tool was constructed\n assert tool.id in tool_dict\n custom_tools = tool_dict[tool.id]\n assert len(custom_tools) == 1\n custom_tool = custom_tools[0]\n assert isinstance(custom_tool, CustomTool)\n\n # Verify NO authorization header (user has no OAuth account)\n _assert_no_authorization_header(custom_tool.headers)\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_python_tool.py", "content": "# \"\"\"\n# External dependency unit tests for Python tool.\n\n# These tests run against a real Code Interpreter service (no mocking of the service).\n# They verify code execution, error handling, timeout behavior, and file generation.\n\n# Requirements:\n# - CODE_INTERPRETER_BASE_URL must be configured and point to a running service\n# - Tests use minimal mocking - only mock run_context infrastructure and db lookups\n# - File store operations execute for real (files are saved and read back)\n# \"\"\"\n\n# import asyncio\n# import io\n# import json\n# from unittest.mock import Mock\n# from unittest.mock import patch\n\n# import pytest\n# from agents import RunContextWrapper\n# from openpyxl import load_workbook\n# from pydantic import TypeAdapter\n# from sqlalchemy.orm import Session\n\n# from onyx.chat.turn.models import ChatTurnContext\n# from onyx.configs.app_configs import CODE_INTERPRETER_BASE_URL\n# from onyx.file_store.models import ChatFileType\n# from onyx.file_store.models import InMemoryChatFile\n# from onyx.file_store.utils import get_default_file_store\n# from onyx.server.query_and_chat.streaming_models import Packet\n# from onyx.server.query_and_chat.streaming_models import PythonToolDelta\n# from onyx.server.query_and_chat.streaming_models import PythonToolStart\n# from onyx.tools.tool_implementations.python.python_tool import PythonTool\n# from onyx.tools.tool_implementations_v2.code_interpreter_client import (\n# CodeInterpreterClient,\n# )\n# from onyx.tools.tool_implementations_v2.python import _python_execution_core\n# from onyx.tools.tool_implementations_v2.python import python\n# from onyx.tools.tool_implementations_v2.tool_result_models import (\n# LlmPythonExecutionResult,\n# )\n\n\n# # Apply initialize_file_store fixture to all tests in this module\n# pytestmark = pytest.mark.usefixtures(\"initialize_file_store\")\n\n\n# @pytest.fixture\n# def mock_run_context() -> RunContextWrapper[ChatTurnContext]:\n# \"\"\"Create a mock run context for testing.\"\"\"\n# # Create mock emitter\n# mock_emitter = Mock()\n# mock_emitter.emit = Mock()\n\n# # Create mock run dependencies\n# mock_dependencies = Mock()\n# mock_dependencies.emitter = mock_emitter\n# mock_dependencies.db_session = Mock()\n\n# # Create mock context\n# mock_context = Mock(spec=ChatTurnContext)\n# mock_context.current_run_step = 0\n# mock_context.run_dependencies = mock_dependencies\n# mock_context.iteration_instructions = []\n# mock_context.global_iteration_responses = []\n# mock_context.chat_files = []\n\n# # Create run context wrapper\n# run_context = Mock(spec=RunContextWrapper)\n# run_context.context = mock_context\n\n# return run_context\n\n\n# @pytest.fixture\n# def code_interpreter_client() -> CodeInterpreterClient:\n# \"\"\"Create a real Code Interpreter client for testing.\"\"\"\n# if not CODE_INTERPRETER_BASE_URL:\n# pytest.skip(\"CODE_INTERPRETER_BASE_URL not configured\")\n# return CodeInterpreterClient()\n\n\n# def test_python_execution_basic(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# ) -> None:\n# \"\"\"Test basic Python execution with simple code.\"\"\"\n# code = 'print(\"Hello, World!\")'\n\n# # Mock get_tool_by_name\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert \"Hello, World!\" in result.stdout\n# assert result.stderr == \"\"\n# assert result.exit_code == 0\n# assert not result.timed_out\n# assert len(result.generated_files) == 0\n\n# # Verify context was updated\n# # Note: @tool_accounting increments current_run_step from 0 to 1 before execution\n# assert len(mock_run_context.context.iteration_instructions) == 1\n# instruction = mock_run_context.context.iteration_instructions[0]\n# assert instruction.iteration_nr == 1\n# assert instruction.plan and \"Python\" in instruction.plan\n\n# assert len(mock_run_context.context.global_iteration_responses) == 1\n# answer = mock_run_context.context.global_iteration_responses[0]\n# assert answer.tool == \"PythonTool\"\n# assert \"Hello, World!\" in answer.answer\n\n# # Verify streaming packets were emitted\n# mock_emitter = mock_run_context.context.run_dependencies.emitter\n# emitter_calls = mock_emitter.emit.call_args_list # type: ignore\n# assert len(emitter_calls) >= 2 # At least start and delta\n\n# # Check for PythonToolStart packet\n# start_packets = [\n# call[0][0]\n# for call in emitter_calls\n# if isinstance(call[0][0].obj, PythonToolStart)\n# ]\n# assert len(start_packets) == 1\n\n# # Check for PythonToolDelta packet\n# delta_packets = [\n# call[0][0]\n# for call in emitter_calls\n# if isinstance(call[0][0].obj, PythonToolDelta)\n# ]\n# assert len(delta_packets) >= 1\n# assert \"Hello, World!\" in delta_packets[0].obj.stdout\n\n\n# def test_python_execution_with_syntax_error(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# ) -> None:\n# \"\"\"Test Python execution with syntax error.\"\"\"\n# code = \"print('missing closing quote\"\n\n# # Mock get_tool_by_name\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify error result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.stdout == \"\"\n# assert len(result.stderr) > 0\n# assert \"SyntaxError\" in result.stderr or \"unterminated\" in result.stderr.lower()\n# assert result.exit_code != 0\n# assert not result.timed_out\n# assert result.error is not None or len(result.stderr) > 0\n# assert len(result.generated_files) == 0\n\n\n# def test_python_execution_with_runtime_error(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# ) -> None:\n# \"\"\"Test Python execution with runtime error.\"\"\"\n# code = \"\"\"\n# x = 10\n# y = 0\n# result = x / y # Division by zero\n# print(result)\n# \"\"\"\n\n# # Mock get_tool_by_name\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify error result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code != 0\n# assert \"ZeroDivisionError\" in result.stderr or \"division\" in result.stderr.lower()\n# assert result.error is not None or len(result.stderr) > 0\n\n\n# def test_python_execution_timeout(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# ) -> None:\n# \"\"\"Test execution timeout handling.\"\"\"\n# # Code that will run longer than the timeout\n# code = \"\"\"\n# import time\n# time.sleep(10)\n# print(\"Should not reach here\")\n# \"\"\"\n\n# # Create client with short timeout (override via execute method)\n# if not CODE_INTERPRETER_BASE_URL:\n# pytest.skip(\"CODE_INTERPRETER_BASE_URL not configured\")\n\n# client = CodeInterpreterClient()\n\n# # Mock get_tool_by_name\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Mock the config to use a short timeout\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.CODE_INTERPRETER_DEFAULT_TIMEOUT_MS\",\n# 1000,\n# ):\n# # Execute code\n# result = _python_execution_core(mock_run_context, code, client)\n\n# # Verify timeout result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.timed_out\n\n\n# def test_python_execution_file_generation(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# db_session: Session, # Needed to initialize DB engine for file_store\n# ) -> None:\n# \"\"\"Test file generation and retrieval.\"\"\"\n# code = \"\"\"\n# import csv\n\n# # Create a CSV file\n# with open('test_output.csv', 'w', newline='') as f:\n# writer = csv.writer(f)\n# writer.writerow(['Name', 'Age', 'City'])\n# writer.writerow(['Alice', '30', 'New York'])\n# writer.writerow(['Bob', '25', 'San Francisco'])\n\n# print(\"CSV file created successfully\")\n# \"\"\"\n\n# # Mock only get_tool_by_name (database lookup)\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code - file store operations happen for real\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code == 0\n# assert \"CSV file created successfully\" in result.stdout\n# assert len(result.generated_files) == 1\n\n# # Verify file metadata\n# generated_file = result.generated_files[0]\n# assert generated_file.filename == \"test_output.csv\"\n# assert generated_file.file_link # File link exists\n# assert generated_file.file_link.startswith(\"http://localhost:3000/api/chat/file/\")\n\n# # Extract file_id from file_link\n# file_id = generated_file.file_link.split(\"/\")[-1]\n\n# # Verify we can read the file back from the file store\n# file_store = get_default_file_store()\n# file_io = file_store.read_file(file_id)\n# file_content = file_io.read()\n\n# # Verify file content\n# assert b\"Name,Age,City\" in file_content\n# assert b\"Alice,30,New York\" in file_content\n# assert b\"Bob,25,San Francisco\" in file_content\n\n# # Verify iteration answer includes file_ids\n# assert len(mock_run_context.context.global_iteration_responses) == 1\n# answer = mock_run_context.context.global_iteration_responses[0]\n# assert answer.file_ids == [file_id]\n\n\n# def test_python_execution_with_matplotlib(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# db_session: Session, # Needed to initialize DB engine for file_store\n# ) -> None:\n# \"\"\"Test matplotlib plot generation.\"\"\"\n# code = \"\"\"\n# import matplotlib\n# matplotlib.use('Agg') # Use non-interactive backend\n# import matplotlib.pyplot as plt\n# import numpy as np\n\n# # Generate data\n# x = np.linspace(0, 10, 100)\n# y = np.sin(x)\n\n# # Create plot\n# plt.figure(figsize=(10, 6))\n# plt.plot(x, y)\n# plt.title('Sine Wave')\n# plt.xlabel('x')\n# plt.ylabel('sin(x)')\n# plt.grid(True)\n\n# # Save plot\n# plt.savefig('sine_wave.png')\n# print(\"Plot saved successfully\")\n# \"\"\"\n\n# # Mock only get_tool_by_name (database lookup)\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code - file store operations happen for real\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code == 0\n# assert \"Plot saved successfully\" in result.stdout\n# assert len(result.generated_files) == 1\n\n# # Verify file metadata\n# generated_file = result.generated_files[0]\n# assert generated_file.filename == \"sine_wave.png\"\n# assert \".png\" in generated_file.filename\n\n# # Extract file_id from file_link\n# file_id = generated_file.file_link.split(\"/\")[-1]\n\n# # Verify we can read the file back from the file store\n# file_store = get_default_file_store()\n# file_io = file_store.read_file(file_id)\n# file_content = file_io.read()\n\n# # Verify the file is a valid PNG (check PNG magic bytes)\n# # PNG magic bytes: 89 50 4E 47 0D 0A 1A 0A\n# assert file_content[:8] == b\"\\x89PNG\\r\\n\\x1a\\n\"\n# assert len(file_content) > 1000 # PNG should be substantial\n\n\n# def test_python_execution_context_updates(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# ) -> None:\n# \"\"\"Test that run_context is properly updated.\"\"\"\n# code = 'print(\"Context update test\")'\n\n# # Mock get_tool_by_name\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 42\n# mock_get_tool.return_value = mock_tool\n\n# # Set specific run step - will be incremented to 6 by @tool_accounting\n# mock_run_context.context.current_run_step = 5\n\n# # Execute code\n# _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify iteration_instructions was updated\n# # Note: @tool_accounting increments from 5 to 6\n# assert len(mock_run_context.context.iteration_instructions) == 1\n# instruction = mock_run_context.context.iteration_instructions[0]\n# assert instruction.iteration_nr == 6\n# assert instruction.plan == \"Executing Python code\"\n# assert instruction.purpose == \"Running Python code\"\n# assert \"secure environment\" in instruction.reasoning\n\n# # Verify global_iteration_responses was updated\n# assert len(mock_run_context.context.global_iteration_responses) == 1\n# answer = mock_run_context.context.global_iteration_responses[0]\n# assert answer.tool == \"PythonTool\"\n# assert answer.tool_id == 42\n# assert answer.iteration_nr == 6\n# assert answer.parallelization_nr == 0\n# assert answer.question == \"Execute Python code\"\n# assert answer.reasoning and \"secure environment\" in answer.reasoning\n# assert \"Context update test\" in answer.answer\n# assert answer.cited_documents == {}\n\n# # Verify packets were emitted with correct index\n# mock_emitter = mock_run_context.context.run_dependencies.emitter\n# emitter_calls = mock_emitter.emit.call_args_list # type: ignore\n# for call in emitter_calls:\n# packet = call[0][0]\n# assert isinstance(packet, Packet)\n# assert packet.ind == 6\n\n\n# def test_python_tool_availability_with_url_set(db_session: Session) -> None:\n# \"\"\"Test PythonTool.is_available() returns True when URL is configured.\"\"\"\n# with patch(\n# \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n# \"http://localhost:8000\",\n# ):\n# assert PythonTool.is_available(db_session) is True\n\n\n# def test_python_tool_availability_without_url(db_session: Session) -> None:\n# \"\"\"Test PythonTool.is_available() returns False when URL is not configured.\"\"\"\n# with patch(\n# \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n# None,\n# ):\n# assert PythonTool.is_available(db_session) is False\n\n# with patch(\n# \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n# \"\",\n# ):\n# assert PythonTool.is_available(db_session) is False\n\n\n# def test_python_function_tool_wrapper(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# ) -> None:\n# \"\"\"Test the @function_tool decorated python() wrapper function.\"\"\"\n# code = 'print(\"Testing function tool wrapper\")'\n\n# # Mock get_tool_by_name and patch CodeInterpreterClient to use our fixture\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.CodeInterpreterClient\"\n# ) as mock_client_class:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n# mock_client_class.return_value = code_interpreter_client\n\n# # Call the function tool wrapper\n# result_coro = python.on_invoke_tool(mock_run_context, json.dumps({\"code\": code})) # type: ignore\n# result_json: str = asyncio.run(result_coro) # type: ignore\n\n# # Verify result is JSON string\n# assert isinstance(result_json, str)\n\n# # Parse and verify result\n# adapter = TypeAdapter(LlmPythonExecutionResult)\n# result = adapter.validate_json(result_json)\n\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert \"Testing function tool wrapper\" in result.stdout\n# assert result.exit_code == 0\n\n\n# def test_python_execution_output_truncation(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# ) -> None:\n# \"\"\"Test that large outputs are properly truncated.\"\"\"\n# # Generate code that produces output larger than truncation limit\n# code = \"\"\"\n# for i in range(10000):\n# print(f\"Line {i}: \" + \"x\" * 100)\n# \"\"\"\n\n# # Mock get_tool_by_name\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# # Set a small truncation limit for testing\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.CODE_INTERPRETER_MAX_OUTPUT_LENGTH\",\n# 5000,\n# ):\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code\n# result = _python_execution_core(\n# mock_run_context, code, code_interpreter_client\n# )\n\n# # Verify output was truncated\n# assert len(result.stdout) <= 5000 + 200 # Allow for truncation message\n# assert \"output truncated\" in result.stdout\n# assert \"characters omitted\" in result.stdout\n\n\n# def test_python_execution_multiple_files(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# db_session: Session, # Needed to initialize DB engine for file_store\n# ) -> None:\n# \"\"\"Test generation of multiple files.\"\"\"\n# code = \"\"\"\n# # Create multiple files\n# with open('file1.txt', 'w') as f:\n# f.write('Content of file 1')\n\n# with open('file2.txt', 'w') as f:\n# f.write('Content of file 2')\n\n# with open('file3.txt', 'w') as f:\n# f.write('Content of file 3')\n\n# print(\"Created 3 files\")\n# \"\"\"\n\n# # Mock only get_tool_by_name (database lookup)\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code - file store operations happen for real\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code == 0\n# assert \"Created 3 files\" in result.stdout\n# assert len(result.generated_files) == 3\n\n# # Verify all files have unique IDs and proper metadata\n# file_ids_result = [f.file_link.split(\"/\")[-1] for f in result.generated_files]\n# assert len(set(file_ids_result)) == 3 # All unique\n\n# # Verify filenames\n# filenames = [f.filename for f in result.generated_files]\n# assert \"file1.txt\" in filenames\n# assert \"file2.txt\" in filenames\n# assert \"file3.txt\" in filenames\n\n# # Verify we can read all files back from the file store\n# file_store = get_default_file_store()\n\n# # Create a mapping of filename to generated file for easier verification\n# files_by_name = {f.filename: f for f in result.generated_files}\n\n# # Verify each expected file\n# for i in range(1, 4):\n# filename = f\"file{i}.txt\"\n# assert filename in files_by_name, f\"Expected file {filename} not found\"\n\n# generated_file = files_by_name[filename]\n# file_id = generated_file.file_link.split(\"/\")[-1]\n# file_io = file_store.read_file(file_id)\n# file_content = file_io.read()\n# expected_content = f\"Content of file {i}\".encode()\n# assert (\n# expected_content in file_content\n# ), f\"Expected content not found in {filename}\"\n\n\n# def test_python_execution_client_error_handling(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# ) -> None:\n# \"\"\"Test error handling when Code Interpreter service fails.\"\"\"\n# code = 'print(\"Test\")'\n\n# # Create a client that will fail\n# if not CODE_INTERPRETER_BASE_URL:\n# pytest.skip(\"CODE_INTERPRETER_BASE_URL not configured\")\n\n# client = CodeInterpreterClient()\n\n# # Mock the execute method to raise an exception\n# with patch.object(client, \"execute\", side_effect=Exception(\"Service unavailable\")):\n# # Execute code\n# result = _python_execution_core(mock_run_context, code, client)\n\n# # Verify error result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code == -1\n# error_msg = result.error or \"\"\n# assert \"Service unavailable\" in result.stderr or \"Service unavailable\" in error_msg\n# assert not result.timed_out\n# assert len(result.generated_files) == 0\n\n# # Verify error delta was emitted\n# mock_emitter = mock_run_context.context.run_dependencies.emitter\n# emitter_calls = mock_emitter.emit.call_args_list # type: ignore\n# delta_packets = [\n# call[0][0]\n# for call in emitter_calls\n# if isinstance(call[0][0].obj, PythonToolDelta)\n# ]\n# assert len(delta_packets) >= 1\n# assert \"Service unavailable\" in delta_packets[-1].obj.stderr\n\n\n# def test_python_execution_with_excel_file(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# db_session: Session, # Needed to initialize DB engine for file_store\n# ) -> None:\n# \"\"\"Test Excel file generation with financial data.\"\"\"\n# code = \"\"\"\n# import pandas as pd\n\n# # Create financial sample data\n# data = {\n# 'Segment': ['Government', 'Government', 'Midmarket', 'Midmarket', 'Enterprise'],\n# 'Country': ['Canada', 'Germany', 'France', 'Germany', 'Canada'],\n# 'Product': ['Carretera', 'Carretera', 'Carretera', 'Carretera', 'Amarilla'],\n# 'Units Sold': [1618.5, 1321, 2178, 888, 2470],\n# 'Manufacturing Price': [3, 3, 3, 3, 260],\n# 'Sale Price': [20, 20, 20, 20, 300],\n# 'Gross Sales': [32370, 26420, 43560, 17760, 741000],\n# 'Discounts': [0, 0, 0, 0, 0],\n# 'Sales': [32370, 26420, 43560, 17760, 741000],\n# 'COGS': [16850, 13940, 22800, 9390, 642000],\n# 'Profit': [15520, 12480, 20760, 8370, 99000],\n# 'Month': ['January', 'January', 'June', 'April', 'September']\n# }\n\n# # Create DataFrame\n# df = pd.DataFrame(data)\n\n# # Write to Excel\n# df.to_excel('financial_report.xlsx', index=False, sheet_name='Financial Data')\n\n# print(f\"Excel file created with {len(df)} rows\")\n# \"\"\"\n\n# # Mock only get_tool_by_name (database lookup)\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code - file store operations happen for real\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code == 0\n# assert \"Excel file created with 5 rows\" in result.stdout\n# assert len(result.generated_files) == 1\n\n# # Verify file metadata\n# generated_file = result.generated_files[0]\n# assert generated_file.filename == \"financial_report.xlsx\"\n# assert \".xlsx\" in generated_file.filename\n\n# # Extract file_id from file_link\n# file_id = generated_file.file_link.split(\"/\")[-1]\n\n# # Verify we can read the file back from the file store\n# file_store = get_default_file_store()\n# file_io = file_store.read_file(file_id)\n# file_content = file_io.read()\n\n# # Verify the file is a valid Excel file (check ZIP magic bytes - xlsx is a ZIP archive)\n# # ZIP magic bytes: 50 4B 03 04\n# assert file_content[:4] == b\"PK\\x03\\x04\"\n# assert len(file_content) > 1000 # Excel file should be substantial\n\n# # Verify we can parse the Excel file with openpyxl directly\n# file_io = io.BytesIO(file_content)\n# workbook = load_workbook(file_io)\n# sheet = workbook[\"Financial Data\"]\n\n# # Verify data structure - get headers from first row\n# first_row = list(sheet.iter_rows(min_row=1, max_row=1, values_only=True))[0]\n# headers = list(first_row) if first_row else []\n# expected_columns = [\n# \"Segment\",\n# \"Country\",\n# \"Product\",\n# \"Units Sold\",\n# \"Manufacturing Price\",\n# \"Sale Price\",\n# \"Gross Sales\",\n# \"Discounts\",\n# \"Sales\",\n# \"COGS\",\n# \"Profit\",\n# \"Month\",\n# ]\n# assert headers == expected_columns\n\n# # Verify row count (excluding header)\n# assert sheet.max_row == 6 # 1 header + 5 data rows\n\n# # Read data rows\n# rows = []\n# for row in sheet.iter_rows(min_row=2, values_only=True):\n# rows.append(row)\n\n# assert len(rows) == 5\n\n# # Verify some sample data\n# segments = [row[0] for row in rows]\n# countries = [row[1] for row in rows]\n# units_sold = [float(row[3]) if row[3] is not None else 0.0 for row in rows] # type: ignore\n# profits = [float(row[10]) if row[10] is not None else 0.0 for row in rows] # type: ignore\n\n# assert \"Government\" in segments\n# assert \"Canada\" in countries\n# assert sum(units_sold) > 8000 # Total units sold\n# assert sum(profits) > 155000 # Total profit\n\n\n# def test_python_execution_with_excel_file_input(\n# mock_run_context: RunContextWrapper[ChatTurnContext],\n# code_interpreter_client: CodeInterpreterClient,\n# db_session: Session, # Needed to initialize DB engine for file_store\n# ) -> None:\n# \"\"\"Test processing an uploaded Excel file - reading and analyzing it.\"\"\"\n# # Load the sample Excel file\n# import os\n\n# test_file_path = os.path.join(\n# os.path.dirname(__file__), \"data\", \"financial-sample.xlsx\"\n# )\n\n# with open(test_file_path, \"rb\") as f:\n# file_content = f.read()\n\n# # Create InMemoryChatFile with the Excel file\n# chat_file = InMemoryChatFile(\n# file_id=\"test-financial-sample\",\n# content=file_content,\n# file_type=ChatFileType.DOC,\n# filename=\"financial-sample.xlsx\",\n# )\n\n# # Add the file to the mock context's chat_files\n# mock_run_context.context.chat_files = [chat_file]\n\n# # Code to analyze the uploaded Excel file\n# code = \"\"\"\n# import pandas as pd\n# import matplotlib\n# matplotlib.use('Agg')\n# import matplotlib.pyplot as plt\n# from openpyxl import load_workbook\n\n# # Read the uploaded Excel file using openpyxl directly\n# workbook = load_workbook('financial-sample.xlsx')\n# sheet = workbook.active\n\n# # Convert to pandas DataFrame\n# data = []\n# headers = [cell.value for cell in sheet[1]]\n# for row in sheet.iter_rows(min_row=2, values_only=True):\n# data.append(row)\n\n# df = pd.DataFrame(data, columns=headers)\n\n# print(f\"Loaded Excel file with {len(df)} rows and {len(df.columns)} columns\")\n# print(f\"\\\\nColumns: {', '.join(df.columns.tolist())}\")\n\n# # Perform analysis\n# print(f\"\\\\n=== Analysis ===\")\n\n# # Group by segment and calculate total sales and profit\n# segment_summary = df.groupby('Segment').agg({\n# ' Sales': 'sum',\n# 'Profit': 'sum',\n# 'Units Sold': 'sum'\n# }).round(2)\n\n# print(f\"\\\\nSales by Segment:\")\n# print(segment_summary)\n\n# # Find top 5 products by profit\n# top_products = df.groupby('Product')['Profit'].sum().sort_values(ascending=False).head(5)\n# print(f\"\\\\nTop 5 Products by Profit:\")\n# print(top_products)\n\n# # Calculate profit margin\n# total_sales = df[' Sales'].sum()\n# total_profit = df['Profit'].sum()\n# profit_margin = (total_profit / total_sales * 100) if total_sales > 0 else 0\n# print(f\"\\\\nOverall Profit Margin: {profit_margin:.2f}%\")\n\n# # Create a visualization\n# fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n\n# # Sales by Segment\n# segment_summary[' Sales'].plot(kind='bar', ax=ax1, color='steelblue')\n# ax1.set_title('Total Sales by Segment')\n# ax1.set_xlabel('Segment')\n# ax1.set_ylabel('Sales ($)')\n# ax1.tick_params(axis='x', rotation=45)\n\n# # Top 5 Products by Profit\n# top_products.plot(kind='barh', ax=ax2, color='seagreen')\n# ax2.set_title('Top 5 Products by Profit')\n# ax2.set_xlabel('Profit ($)')\n# ax2.set_ylabel('Product')\n\n# plt.tight_layout()\n# plt.savefig('financial_analysis.png', dpi=100, bbox_inches='tight')\n# print(f\"\\\\nVisualization saved as financial_analysis.png\")\n\n# # Create summary report Excel file\n# summary_data = {\n# 'Metric': ['Total Sales', 'Total Profit', 'Profit Margin %', 'Total Units Sold', 'Number of Records'],\n# 'Value': [\n# f\"${total_sales:,.2f}\",\n# f\"${total_profit:,.2f}\",\n# f\"{profit_margin:.2f}%\",\n# f\"{df['Units Sold'].sum():,.0f}\",\n# len(df)\n# ]\n# }\n# summary_df = pd.DataFrame(summary_data)\n\n# with pd.ExcelWriter('financial_summary.xlsx') as writer:\n# summary_df.to_excel(writer, sheet_name='Summary', index=False)\n# segment_summary.to_excel(writer, sheet_name='By Segment')\n\n# print(f\"Summary report saved as financial_summary.xlsx\")\n# \"\"\"\n\n# # Mock only get_tool_by_name (database lookup)\n# with patch(\n# \"onyx.tools.tool_implementations_v2.python.get_tool_by_name\"\n# ) as mock_get_tool:\n# mock_tool = Mock()\n# mock_tool.id = 1\n# mock_get_tool.return_value = mock_tool\n\n# # Execute code - file store operations happen for real\n# result = _python_execution_core(mock_run_context, code, code_interpreter_client)\n\n# # Verify result\n# assert isinstance(result, LlmPythonExecutionResult)\n# assert result.exit_code == 0\n# assert \"Loaded Excel file\" in result.stdout\n# assert \"Analysis\" in result.stdout\n# assert \"Sales by Segment\" in result.stdout\n# assert \"Top 5 Products by Profit\" in result.stdout\n# assert \"Profit Margin\" in result.stdout\n\n# # Should generate 2 files: PNG visualization and Excel summary\n# assert len(result.generated_files) == 2\n\n# # Verify generated files\n# filenames = [f.filename for f in result.generated_files]\n# assert \"financial_analysis.png\" in filenames\n# assert \"financial_summary.xlsx\" in filenames\n\n# # Verify we can read and validate the generated files\n# file_store = get_default_file_store()\n\n# # Check the PNG file\n# png_file = next(\n# f for f in result.generated_files if f.filename == \"financial_analysis.png\"\n# )\n# png_file_id = png_file.file_link.split(\"/\")[-1]\n# png_io = file_store.read_file(png_file_id)\n# png_content = png_io.read()\n# assert png_content[:8] == b\"\\x89PNG\\r\\n\\x1a\\n\" # PNG magic bytes\n# assert len(png_content) > 5000 # Should be substantial\n\n# # Check the Excel summary file\n# xlsx_file = next(\n# f for f in result.generated_files if f.filename == \"financial_summary.xlsx\"\n# )\n# xlsx_file_id = xlsx_file.file_link.split(\"/\")[-1]\n# xlsx_io = file_store.read_file(xlsx_file_id)\n# xlsx_content = xlsx_io.read()\n# assert xlsx_content[:4] == b\"PK\\x03\\x04\" # ZIP/Excel magic bytes\n\n# # Parse and verify the summary Excel file using openpyxl directly\n# xlsx_io_obj = io.BytesIO(xlsx_content)\n# workbook = load_workbook(xlsx_io_obj)\n# sheet = workbook[\"Summary\"]\n\n# # Read headers from first row\n# first_row = list(sheet.iter_rows(min_row=1, max_row=1, values_only=True))[0]\n# headers = list(first_row) if first_row else []\n# assert \"Metric\" in headers\n# assert \"Value\" in headers\n\n# # Read all rows and extract metrics\n# metrics = []\n# for row in sheet.iter_rows(min_row=2, values_only=True):\n# if row[0]: # Metric column\n# metrics.append(row[0])\n\n# assert \"Total Sales\" in metrics\n# assert \"Total Profit\" in metrics\n# assert \"Profit Margin %\" in metrics\n\n\n# if __name__ == \"__main__\":\n# # Run with: python -m pytest tests/external_dependency_unit/tools/test_python_tool.py -v\n# pytest.main([__file__, \"-v\"])\n\n\nfrom __future__ import annotations\n\nimport io\nimport json\nimport threading\nfrom collections.abc import Generator\nfrom http.server import BaseHTTPRequestHandler\nfrom http.server import HTTPServer\nfrom typing import Any\nfrom unittest.mock import patch\n\nimport pytest\nfrom fastapi import UploadFile\nfrom fastapi.background import BackgroundTasks\nfrom sqlalchemy.orm import Session\nfrom starlette.datastructures import Headers\n\nimport onyx.tools.tool_implementations.python.code_interpreter_client as ci_mod\nfrom onyx.chat.process_message import handle_stream_message_objects\nfrom onyx.db.models import Persona\nfrom onyx.db.tools import get_builtin_tool\nfrom onyx.file_store.models import ChatFileType\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.server.features.projects.api import upload_user_files\nfrom onyx.server.query_and_chat.chat_backend import get_chat_session\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.streaming_models import Packet\nfrom onyx.server.query_and_chat.streaming_models import PythonToolDelta\nfrom onyx.server.query_and_chat.streaming_models import PythonToolStart\nfrom onyx.server.query_and_chat.streaming_models import SectionEnd\nfrom onyx.server.query_and_chat.streaming_models import ToolCallArgumentDelta\nfrom onyx.tools.tool_implementations.python.python_tool import PythonTool\nfrom tests.external_dependency_unit.answer.stream_test_builder import StreamTestBuilder\nfrom tests.external_dependency_unit.answer.stream_test_utils import create_chat_session\nfrom tests.external_dependency_unit.answer.stream_test_utils import create_placement\nfrom tests.external_dependency_unit.conftest import create_test_user\nfrom tests.external_dependency_unit.mock_llm import LLMAnswerResponse\nfrom tests.external_dependency_unit.mock_llm import LLMToolCallResponse\nfrom tests.external_dependency_unit.mock_llm import use_mock_llm\n\n\n# ---------------------------------------------------------------------------\n# Mock Code Interpreter Server\n# ---------------------------------------------------------------------------\n\n\nclass CapturedRequest:\n \"\"\"A single HTTP request captured by the mock server.\"\"\"\n\n def __init__(self, method: str, path: str, body: bytes) -> None:\n self.method = method\n self.path = path\n self.body = body\n\n def json_body(self) -> dict[str, Any]:\n return json.loads(self.body)\n\n\nclass _MockCIHandler(BaseHTTPRequestHandler):\n \"\"\"HTTP handler that records every request and returns canned responses.\"\"\"\n\n server: MockCodeInterpreterServer\n\n def do_POST(self) -> None:\n body = self._read_body()\n self._capture(\"POST\", body)\n\n if self.path == \"/v1/files\":\n self.server._file_counter += 1\n self._respond_json(\n 200, {\"file_id\": f\"mock-ci-file-{self.server._file_counter}\"}\n )\n elif self.path == \"/v1/execute/stream\":\n if self.server.streaming_enabled:\n self._respond_sse(\n [\n (\n \"output\",\n {\"stream\": \"stdout\", \"data\": \"mock output\\n\"},\n ),\n (\n \"result\",\n {\n \"exit_code\": 0,\n \"timed_out\": False,\n \"duration_ms\": 50,\n \"files\": [],\n },\n ),\n ]\n )\n else:\n self._respond_json(404, {\"error\": \"not found\"})\n elif self.path == \"/v1/execute\":\n self._respond_json(\n 200,\n {\n \"stdout\": \"mock output\\n\",\n \"stderr\": \"\",\n \"exit_code\": 0,\n \"timed_out\": False,\n \"duration_ms\": 50,\n \"files\": [],\n },\n )\n else:\n self._respond_json(404, {\"error\": \"not found\"})\n\n def do_GET(self) -> None:\n self._capture(\"GET\", b\"\")\n if self.path == \"/health\":\n self._respond_json(200, {\"status\": \"ok\"})\n else:\n self._respond_json(404, {\"error\": \"not found\"})\n\n def do_DELETE(self) -> None:\n self._capture(\"DELETE\", b\"\")\n self.send_response(200)\n self.end_headers()\n\n def _read_body(self) -> bytes:\n length = int(self.headers.get(\"Content-Length\", 0))\n return self.rfile.read(length) if length else b\"\"\n\n def _capture(self, method: str, body: bytes) -> None:\n self.server.captured_requests.append(\n CapturedRequest(method=method, path=self.path, body=body)\n )\n\n def _respond_json(self, status: int, data: dict[str, Any]) -> None:\n payload = json.dumps(data).encode()\n self.send_response(status)\n self.send_header(\"Content-Type\", \"application/json\")\n self.send_header(\"Content-Length\", str(len(payload)))\n self.end_headers()\n self.wfile.write(payload)\n\n def _respond_sse(self, events: list[tuple[str, dict[str, Any]]]) -> None:\n frames = []\n for event_type, data in events:\n frames.append(f\"event: {event_type}\\ndata: {json.dumps(data)}\\n\\n\")\n payload = \"\".join(frames).encode()\n self.send_response(200)\n self.send_header(\"Content-Type\", \"text/event-stream\")\n self.send_header(\"Content-Length\", str(len(payload)))\n self.end_headers()\n self.wfile.write(payload)\n\n def log_message(self, format: str, *args: Any) -> None: # noqa: A002\n pass\n\n\nclass MockCodeInterpreterServer(HTTPServer):\n \"\"\"HTTPServer wrapper that records requests for assertions.\"\"\"\n\n def __init__(self) -> None:\n super().__init__((\"localhost\", 0), _MockCIHandler)\n self.captured_requests: list[CapturedRequest] = []\n self._file_counter = 0\n self.streaming_enabled: bool = True\n\n @property\n def url(self) -> str:\n host, port = self.server_address\n return f\"http://{host!s}:{port}\"\n\n def start(self) -> None:\n threading.Thread(target=self.serve_forever, daemon=True).start()\n\n def get_requests(\n self,\n method: str | None = None,\n path: str | None = None,\n ) -> list[CapturedRequest]:\n results = self.captured_requests\n if method:\n results = [r for r in results if r.method == method]\n if path:\n results = [r for r in results if r.path == path]\n return results\n\n\n# ---------------------------------------------------------------------------\n# Fixtures\n# ---------------------------------------------------------------------------\n\n\n@pytest.fixture(scope=\"module\")\ndef mock_ci_server() -> Generator[MockCodeInterpreterServer, None, None]:\n server = MockCodeInterpreterServer()\n server.start()\n yield server\n server.shutdown()\n\n\n@pytest.fixture(autouse=True)\ndef _clear_health_cache() -> None:\n \"\"\"Reset the health check cache before every test.\"\"\"\n import onyx.tools.tool_implementations.python.code_interpreter_client as mod\n\n mod._health_cache = {}\n\n\n@pytest.fixture()\ndef _attach_python_tool_to_default_persona(db_session: Session) -> None:\n \"\"\"Ensure the default persona (id=0) has the PythonTool attached.\"\"\"\n python_tool_db = get_builtin_tool(db_session, PythonTool)\n persona = db_session.get(Persona, 0)\n assert persona is not None, \"Default persona (id=0) not found\"\n\n if python_tool_db not in persona.tools:\n persona.tools.append(python_tool_db)\n db_session.commit()\n\n\n# ---------------------------------------------------------------------------\n# Test\n# ---------------------------------------------------------------------------\n\n\ndef test_code_interpreter_receives_chat_files(\n db_session: Session,\n mock_ci_server: MockCodeInterpreterServer,\n _attach_python_tool_to_default_persona: None,\n initialize_file_store: None, # noqa: ARG001\n) -> None:\n mock_ci_server.captured_requests.clear()\n mock_ci_server._file_counter = 0\n mock_url = mock_ci_server.url\n\n user = create_test_user(db_session, \"ci_test_admin\")\n chat_session = create_chat_session(db_session=db_session, user=user)\n\n # Upload a test CSV\n csv_content = b\"name,age,city\\nAlice,30,NYC\\nBob,25,SF\\n\"\n result = upload_user_files(\n bg_tasks=BackgroundTasks(),\n files=[\n UploadFile(\n file=io.BytesIO(csv_content),\n filename=\"data.csv\",\n size=len(csv_content),\n headers=Headers({\"content-type\": \"text/csv\"}),\n )\n ],\n project_id=None,\n temp_id_map=json.dumps({\"0|data.csv\": \"data.csv\"}),\n user=user,\n db_session=db_session,\n )\n assert len(result.user_files) == 1\n user_file = result.user_files[0]\n\n file_descriptor: FileDescriptor = {\n \"id\": user_file.file_id,\n \"type\": ChatFileType.TABULAR,\n \"name\": \"data.csv\",\n \"user_file_id\": str(user_file.id),\n }\n\n code = \"import pandas as pd\\ndf = pd.read_csv('data.csv')\\nprint(df)\"\n msg_req = SendMessageRequest(\n message=\"Read the CSV and print it.\",\n chat_session_id=chat_session.id,\n file_descriptors=[file_descriptor],\n stream=True,\n )\n\n original_defaults = ci_mod.CodeInterpreterClient.__init__.__defaults__\n with (\n use_mock_llm() as mock_llm,\n patch(\n \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n mock_url,\n ),\n patch(\n \"onyx.tools.tool_implementations.python.code_interpreter_client.CODE_INTERPRETER_BASE_URL\",\n mock_url,\n ),\n ):\n mock_llm.add_response(\n LLMToolCallResponse(\n tool_name=\"python\",\n tool_call_id=\"call_test_1\",\n tool_call_argument_tokens=[json.dumps({\"code\": code})],\n )\n )\n mock_llm.forward_till_end()\n\n ci_mod.CodeInterpreterClient.__init__.__defaults__ = (mock_url,)\n try:\n list(\n handle_stream_message_objects(\n new_msg_req=msg_req, user=user, db_session=db_session\n )\n )\n finally:\n ci_mod.CodeInterpreterClient.__init__.__defaults__ = original_defaults\n\n # Verify: file uploaded and code executed via streaming.\n assert len(mock_ci_server.get_requests(method=\"POST\", path=\"/v1/files\")) == 1\n assert (\n len(mock_ci_server.get_requests(method=\"POST\", path=\"/v1/execute/stream\")) == 1\n )\n\n # Staged input files are intentionally NOT deleted — PythonTool caches their\n # file IDs across agent-loop iterations to avoid re-uploading on every call.\n # The code interpreter cleans them up via its own TTL.\n assert len(mock_ci_server.get_requests(method=\"DELETE\")) == 0\n\n execute_body = mock_ci_server.get_requests(\n method=\"POST\", path=\"/v1/execute/stream\"\n )[0].json_body()\n assert execute_body[\"code\"] == code\n assert len(execute_body[\"files\"]) == 1\n assert execute_body[\"files\"][0][\"path\"] == \"data.csv\"\n\n\ndef test_code_interpreter_replay_packets_include_code_and_output(\n db_session: Session,\n mock_ci_server: MockCodeInterpreterServer,\n _attach_python_tool_to_default_persona: None,\n initialize_file_store: None, # noqa: ARG001\n) -> None:\n \"\"\"After a code interpreter message completes, retrieving the message\n via translate_assistant_message_to_packets should emit PythonToolStart\n (containing the executed code) and PythonToolDelta (containing\n stdout/stderr), not generic CustomTool packets.\"\"\"\n mock_ci_server.captured_requests.clear()\n mock_ci_server._file_counter = 0\n mock_url = mock_ci_server.url\n\n user = create_test_user(db_session, \"ci_replay_test\")\n chat_session = create_chat_session(db_session=db_session, user=user)\n\n code = 'x = 2 + 2\\nprint(f\"Result: {x}\")'\n msg_req = SendMessageRequest(\n message=\"Calculate 2 + 2\",\n chat_session_id=chat_session.id,\n stream=True,\n )\n\n original_defaults = ci_mod.CodeInterpreterClient.__init__.__defaults__\n with (\n use_mock_llm() as mock_llm,\n patch(\n \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n mock_url,\n ),\n patch(\n \"onyx.tools.tool_implementations.python.code_interpreter_client.CODE_INTERPRETER_BASE_URL\",\n mock_url,\n ),\n ):\n answer_tokens = [\"The \", \"result \", \"is \", \"4.\"]\n\n ci_mod.CodeInterpreterClient.__init__.__defaults__ = (mock_url,)\n try:\n handler = StreamTestBuilder(llm_controller=mock_llm)\n\n stream = handle_stream_message_objects(\n new_msg_req=msg_req, user=user, db_session=db_session\n )\n # First packet is always MessageResponseIDInfo\n next(stream)\n\n # Phase 1: LLM requests python tool execution.\n handler.add_response(\n LLMToolCallResponse(\n tool_name=\"python\",\n tool_call_id=\"call_replay_test\",\n tool_call_argument_tokens=[json.dumps({\"code\": code})],\n )\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=ToolCallArgumentDelta(\n tool_type=\"python\",\n argument_deltas={\"code\": code},\n ),\n ),\n forward=2,\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=PythonToolStart(code=code),\n ),\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=PythonToolDelta(stdout=\"mock output\\n\", stderr=\"\", file_ids=[]),\n ),\n forward=False,\n ).expect(\n Packet(\n placement=create_placement(0),\n obj=SectionEnd(),\n ),\n forward=False,\n ).run_and_validate(\n stream=stream\n )\n\n # Phase 2: LLM produces a final answer after tool execution.\n handler.add_response(\n LLMAnswerResponse(answer_tokens=answer_tokens)\n ).expect_agent_response(\n answer_tokens=answer_tokens,\n turn_index=1,\n ).run_and_validate(\n stream=stream\n )\n\n with pytest.raises(StopIteration):\n next(stream)\n\n finally:\n ci_mod.CodeInterpreterClient.__init__.__defaults__ = original_defaults\n\n # Retrieve the chat session through the same endpoint the frontend uses\n chat_detail = get_chat_session(\n session_id=chat_session.id,\n user=user,\n db_session=db_session,\n )\n\n assert (\n len(mock_ci_server.get_requests(method=\"POST\", path=\"/v1/execute/stream\")) == 1\n )\n\n # The response contains `packets` — a list of packet-lists, one per\n # assistant message. We should have exactly one assistant message.\n assert (\n len(chat_detail.packets) == 1\n ), f\"Expected 1 assistant packet list, got {len(chat_detail.packets)}\"\n packets = chat_detail.packets[0]\n\n # Extract PythonToolStart packets – these must contain the code\n start_packets = [p for p in packets if isinstance(p.obj, PythonToolStart)]\n assert (\n len(start_packets) == 1\n ), f\"Expected 1 PythonToolStart packet, got {len(start_packets)}. Packet types: {[type(p.obj).__name__ for p in packets]}\"\n start_obj = start_packets[0].obj\n assert isinstance(start_obj, PythonToolStart)\n assert start_obj.code == code\n\n # Extract PythonToolDelta packets – these must contain stdout/stderr\n delta_packets = [p for p in packets if isinstance(p.obj, PythonToolDelta)]\n assert len(delta_packets) >= 1, (\n f\"Expected at least 1 PythonToolDelta packet, got {len(delta_packets)}. \"\n f\"Packet types: {[type(p.obj).__name__ for p in packets]}\"\n )\n # The mock CI server returns \"mock output\\n\" as stdout\n delta_obj = delta_packets[0].obj\n assert isinstance(delta_obj, PythonToolDelta)\n assert \"mock output\" in delta_obj.stdout\n\n\ndef test_code_interpreter_streaming_fallback_to_batch(\n db_session: Session,\n mock_ci_server: MockCodeInterpreterServer,\n _attach_python_tool_to_default_persona: None,\n initialize_file_store: None, # noqa: ARG001\n) -> None:\n \"\"\"When the streaming endpoint is not available (older code-interpreter),\n execute_streaming should fall back to the batch /v1/execute endpoint.\"\"\"\n mock_ci_server.captured_requests.clear()\n mock_ci_server._file_counter = 0\n mock_ci_server.streaming_enabled = False\n mock_url = mock_ci_server.url\n\n user = create_test_user(db_session, \"ci_fallback_test\")\n chat_session = create_chat_session(db_session=db_session, user=user)\n\n code = 'print(\"fallback test\")'\n msg_req = SendMessageRequest(\n message=\"Print fallback test\",\n chat_session_id=chat_session.id,\n stream=True,\n )\n\n original_defaults = ci_mod.CodeInterpreterClient.__init__.__defaults__\n with (\n use_mock_llm() as mock_llm,\n patch(\n \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n mock_url,\n ),\n patch(\n \"onyx.tools.tool_implementations.python.code_interpreter_client.CODE_INTERPRETER_BASE_URL\",\n mock_url,\n ),\n ):\n mock_llm.add_response(\n LLMToolCallResponse(\n tool_name=\"python\",\n tool_call_id=\"call_fallback\",\n tool_call_argument_tokens=[json.dumps({\"code\": code})],\n )\n )\n mock_llm.forward_till_end()\n\n ci_mod.CodeInterpreterClient.__init__.__defaults__ = (mock_url,)\n try:\n packets = list(\n handle_stream_message_objects(\n new_msg_req=msg_req, user=user, db_session=db_session\n )\n )\n finally:\n ci_mod.CodeInterpreterClient.__init__.__defaults__ = original_defaults\n mock_ci_server.streaming_enabled = True\n\n # Streaming was attempted first (returned 404), then fell back to batch\n assert (\n len(mock_ci_server.get_requests(method=\"POST\", path=\"/v1/execute/stream\")) == 1\n )\n assert len(mock_ci_server.get_requests(method=\"POST\", path=\"/v1/execute\")) == 1\n\n # Verify output still made it through\n delta_packets = [\n p\n for p in packets\n if isinstance(p, Packet) and isinstance(p.obj, PythonToolDelta)\n ]\n assert len(delta_packets) >= 1\n first_delta = delta_packets[0].obj\n assert isinstance(first_delta, PythonToolDelta)\n assert \"mock output\" in first_delta.stdout\n" }, { "path": "backend/tests/external_dependency_unit/tools/test_python_tool_server_enabled.py", "content": "\"\"\"Tests that PythonTool.is_available() respects the server_enabled DB flag.\n\nUses a real DB session with CODE_INTERPRETER_BASE_URL mocked so the\nenvironment-variable check passes and the DB flag is the deciding factor.\n\"\"\"\n\nfrom unittest.mock import patch\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.code_interpreter import fetch_code_interpreter_server\nfrom onyx.db.code_interpreter import update_code_interpreter_server_enabled\nfrom onyx.tools.tool_implementations.python.python_tool import PythonTool\n\n\ndef test_python_tool_unavailable_when_server_disabled(\n db_session: Session,\n) -> None:\n \"\"\"With a valid base URL, the tool should be unavailable when\n server_enabled is False in the DB.\"\"\"\n server = fetch_code_interpreter_server(db_session)\n initial_enabled = server.server_enabled\n\n try:\n update_code_interpreter_server_enabled(db_session, enabled=False)\n\n with patch(\n \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n \"http://fake:8888\",\n ):\n assert PythonTool.is_available(db_session) is False\n finally:\n update_code_interpreter_server_enabled(db_session, enabled=initial_enabled)\n\n\ndef test_python_tool_available_when_server_enabled(\n db_session: Session,\n) -> None:\n \"\"\"With a valid base URL, the tool should be available when\n server_enabled is True in the DB.\"\"\"\n server = fetch_code_interpreter_server(db_session)\n initial_enabled = server.server_enabled\n\n try:\n update_code_interpreter_server_enabled(db_session, enabled=True)\n\n with patch(\n \"onyx.tools.tool_implementations.python.python_tool.CODE_INTERPRETER_BASE_URL\",\n \"http://fake:8888\",\n ):\n assert PythonTool.is_available(db_session) is True\n finally:\n update_code_interpreter_server_enabled(db_session, enabled=initial_enabled)\n" }, { "path": "backend/tests/external_dependency_unit/tracing/__init__.py", "content": "" }, { "path": "backend/tests/external_dependency_unit/tracing/test_llm_span_recording.py", "content": "\"\"\"Tests for LLM span recording utilities.\"\"\"\n\nfrom typing import Any\nfrom unittest.mock import MagicMock\n\nimport pytest\n\nfrom onyx.llm.model_response import ChatCompletionMessageToolCall\nfrom onyx.llm.model_response import Choice\nfrom onyx.llm.model_response import FunctionCall as ModelResponseFunctionCall\nfrom onyx.llm.model_response import Message\nfrom onyx.llm.model_response import ModelResponse\nfrom onyx.llm.model_response import Usage\nfrom onyx.llm.models import FunctionCall\nfrom onyx.llm.models import ToolCall\nfrom onyx.tracing.framework.span_data import GenerationSpanData\nfrom onyx.tracing.llm_utils import record_llm_response\nfrom onyx.tracing.llm_utils import record_llm_span_output\n\n\n@pytest.fixture\ndef mock_span() -> MagicMock:\n \"\"\"Create a mock span with GenerationSpanData.\"\"\"\n span = MagicMock()\n span.span_data = GenerationSpanData()\n return span\n\n\nclass TestRecordLlmResponse:\n \"\"\"Tests for record_llm_response function.\"\"\"\n\n def test_records_content_from_response(self, mock_span: MagicMock) -> None:\n \"\"\"Test that content is correctly extracted and recorded.\"\"\"\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(content=\"Hello, world!\", role=\"assistant\"),\n ),\n )\n\n record_llm_response(mock_span, response)\n\n assert mock_span.span_data.output == [\n {\"role\": \"assistant\", \"content\": \"Hello, world!\"}\n ]\n\n def test_records_reasoning_from_response(self, mock_span: MagicMock) -> None:\n \"\"\"Test that reasoning/extended thinking is recorded.\"\"\"\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(\n content=\"The answer is 42.\",\n role=\"assistant\",\n reasoning_content=\"Let me think step by step...\",\n ),\n ),\n )\n\n record_llm_response(mock_span, response)\n\n assert mock_span.span_data.output == [\n {\"role\": \"assistant\", \"content\": \"The answer is 42.\"}\n ]\n assert mock_span.span_data.reasoning == \"Let me think step by step...\"\n\n def test_records_tool_calls_from_response(self, mock_span: MagicMock) -> None:\n \"\"\"Test that tool calls are correctly extracted and recorded.\"\"\"\n tool_call = ChatCompletionMessageToolCall(\n id=\"call-123\",\n type=\"function\",\n function=ModelResponseFunctionCall(\n name=\"search_documents\",\n arguments='{\"query\": \"test query\"}',\n ),\n )\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(\n content=None,\n role=\"assistant\",\n tool_calls=[tool_call],\n ),\n ),\n )\n\n record_llm_response(mock_span, response)\n\n output = mock_span.span_data.output\n assert len(output) == 1\n assert output[0][\"role\"] == \"assistant\"\n assert \"tool_calls\" in output[0]\n assert len(output[0][\"tool_calls\"]) == 1\n assert output[0][\"tool_calls\"][0][\"id\"] == \"call-123\"\n assert output[0][\"tool_calls\"][0][\"function\"][\"name\"] == \"search_documents\"\n\n def test_records_usage_from_response(self, mock_span: MagicMock) -> None:\n \"\"\"Test that usage metrics are correctly recorded.\"\"\"\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(content=\"Test\", role=\"assistant\"),\n ),\n usage=Usage(\n prompt_tokens=100,\n completion_tokens=50,\n total_tokens=150,\n cache_creation_input_tokens=10,\n cache_read_input_tokens=20,\n ),\n )\n\n record_llm_response(mock_span, response)\n\n assert mock_span.span_data.usage is not None\n assert mock_span.span_data.usage[\"input_tokens\"] == 100\n assert mock_span.span_data.usage[\"output_tokens\"] == 50\n assert mock_span.span_data.usage[\"total_tokens\"] == 150\n assert mock_span.span_data.usage[\"cache_read_input_tokens\"] == 20\n assert mock_span.span_data.usage[\"cache_creation_input_tokens\"] == 10\n\n def test_handles_none_content(self, mock_span: MagicMock) -> None:\n \"\"\"Test that None content is handled (e.g., tool-only response).\"\"\"\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(content=None, role=\"assistant\"),\n ),\n )\n\n record_llm_response(mock_span, response)\n\n # Content should not be in output dict when None\n assert mock_span.span_data.output == [{\"role\": \"assistant\"}]\n\n def test_handles_no_usage(self, mock_span: MagicMock) -> None:\n \"\"\"Test that missing usage is handled gracefully.\"\"\"\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(content=\"Test\", role=\"assistant\"),\n ),\n usage=None,\n )\n\n record_llm_response(mock_span, response)\n\n # Usage should remain None/unset\n assert mock_span.span_data.usage is None\n\n def test_records_all_fields_together(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording a response with all fields present.\"\"\"\n tool_call = ChatCompletionMessageToolCall(\n id=\"call-456\",\n type=\"function\",\n function=ModelResponseFunctionCall(\n name=\"analyze\",\n arguments='{\"text\": \"sample\"}',\n ),\n )\n response = ModelResponse(\n id=\"test-id\",\n created=\"2024-01-01\",\n choice=Choice(\n message=Message(\n content=\"Here's my analysis:\",\n role=\"assistant\",\n reasoning_content=\"I need to think about this carefully...\",\n tool_calls=[tool_call],\n ),\n ),\n usage=Usage(\n prompt_tokens=200,\n completion_tokens=100,\n total_tokens=300,\n cache_creation_input_tokens=0,\n cache_read_input_tokens=50,\n ),\n )\n\n record_llm_response(mock_span, response)\n\n # Check output\n output = mock_span.span_data.output\n assert len(output) == 1\n assert output[0][\"role\"] == \"assistant\"\n assert output[0][\"content\"] == \"Here's my analysis:\"\n assert len(output[0][\"tool_calls\"]) == 1\n\n # Check reasoning\n assert (\n mock_span.span_data.reasoning == \"I need to think about this carefully...\"\n )\n\n # Check usage\n assert mock_span.span_data.usage[\"input_tokens\"] == 200\n assert mock_span.span_data.usage[\"output_tokens\"] == 100\n\n\nclass TestRecordLlmSpanOutput:\n \"\"\"Tests for record_llm_span_output function (streaming scenarios).\"\"\"\n\n def test_records_string_output(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording a simple string output.\"\"\"\n record_llm_span_output(mock_span, \"Hello, world!\")\n\n assert mock_span.span_data.output == [\n {\"role\": \"assistant\", \"content\": \"Hello, world!\"}\n ]\n\n def test_records_none_output(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording None output.\"\"\"\n record_llm_span_output(mock_span, None)\n\n assert mock_span.span_data.output == [{\"role\": \"assistant\", \"content\": None}]\n\n def test_records_sequence_output(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording a sequence of message dicts.\"\"\"\n messages: list[dict[str, Any]] = [\n {\"role\": \"assistant\", \"content\": \"Part 1\"},\n {\"role\": \"assistant\", \"content\": \"Part 2\"},\n ]\n\n record_llm_span_output(mock_span, messages)\n\n assert mock_span.span_data.output == messages\n\n def test_records_usage(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording usage information.\"\"\"\n usage = MagicMock()\n usage.prompt_tokens = 50\n usage.completion_tokens = 25\n usage.total_tokens = 75\n usage.cache_read_input_tokens = 10\n usage.cache_creation_input_tokens = 5\n\n record_llm_span_output(mock_span, \"Test output\", usage=usage)\n\n assert mock_span.span_data.usage is not None\n assert mock_span.span_data.usage[\"input_tokens\"] == 50\n assert mock_span.span_data.usage[\"output_tokens\"] == 25\n\n def test_records_reasoning(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording reasoning content.\"\"\"\n record_llm_span_output(\n mock_span, \"Final answer\", reasoning=\"Step by step thinking...\"\n )\n\n assert mock_span.span_data.reasoning == \"Step by step thinking...\"\n\n def test_records_tool_calls(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording tool calls in streaming scenario.\"\"\"\n tool_calls = [\n ToolCall(\n id=\"call-789\",\n type=\"function\",\n function=FunctionCall(\n name=\"get_weather\",\n arguments='{\"location\": \"NYC\"}',\n ),\n )\n ]\n\n record_llm_span_output(mock_span, \"Checking weather...\", tool_calls=tool_calls)\n\n output = mock_span.span_data.output\n assert len(output) == 1\n assert output[0][\"content\"] == \"Checking weather...\"\n assert \"tool_calls\" in output[0]\n assert len(output[0][\"tool_calls\"]) == 1\n assert output[0][\"tool_calls\"][0][\"id\"] == \"call-789\"\n\n def test_records_tool_calls_with_none_output(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording tool calls when output is None.\"\"\"\n tool_calls = [\n ToolCall(\n id=\"call-abc\",\n type=\"function\",\n function=FunctionCall(\n name=\"search\",\n arguments='{\"q\": \"test\"}',\n ),\n )\n ]\n\n record_llm_span_output(mock_span, None, tool_calls=tool_calls)\n\n output = mock_span.span_data.output\n assert len(output) == 1\n assert output[0][\"content\"] is None\n assert len(output[0][\"tool_calls\"]) == 1\n\n def test_records_all_streaming_fields(self, mock_span: MagicMock) -> None:\n \"\"\"Test recording all fields in streaming scenario.\"\"\"\n usage = MagicMock()\n usage.prompt_tokens = 100\n usage.completion_tokens = 50\n usage.total_tokens = 150\n usage.cache_read_input_tokens = 0\n usage.cache_creation_input_tokens = 0\n\n tool_calls = [\n ToolCall(\n id=\"call-xyz\",\n type=\"function\",\n function=FunctionCall(\n name=\"calculator\",\n arguments='{\"expr\": \"2+2\"}',\n ),\n )\n ]\n\n record_llm_span_output(\n mock_span,\n output=\"Computing...\",\n usage=usage,\n reasoning=\"Let me calculate this.\",\n tool_calls=tool_calls,\n )\n\n # Check all fields\n output = mock_span.span_data.output\n assert output[0][\"content\"] == \"Computing...\"\n assert len(output[0][\"tool_calls\"]) == 1\n assert mock_span.span_data.reasoning == \"Let me calculate this.\"\n assert mock_span.span_data.usage[\"input_tokens\"] == 100\n" }, { "path": "backend/tests/integration/Dockerfile", "content": "# syntax=docker/dockerfile:1.6\n# This image is only for running integration tests. It layers test-specific\n# files and dependencies on top of the backend image.\nFROM base AS integration-base\n\nWORKDIR /app\n\n# Integration test stuff\nCOPY ./requirements/dev.txt /tmp/dev-requirements.txt\nRUN uv pip install --system --no-cache-dir --upgrade -r /tmp/dev-requirements.txt && \\\n rm -rf ~/.cache/uv /tmp/*.txt\n\nCOPY ./pytest.ini /app/pytest.ini\nCOPY ./tests/integration /app/tests/integration\n# copies all files, but not folders, in the tests directory\nCOPY ./tests/* /app/tests/\n\nFROM base AS openapi-schema\nCOPY ./scripts/onyx_openapi_schema.py /app/scripts/onyx_openapi_schema.py\n# TODO(Nik): https://linear.app/onyx-app/issue/ENG-1/update-test-infra-to-use-test-license\nRUN LICENSE_ENFORCEMENT_ENABLED=false python scripts/onyx_openapi_schema.py --filename openapi.json\n\nFROM openapitools/openapi-generator-cli:latest AS openapi-client\nWORKDIR /local\nCOPY --from=openapi-schema /app/openapi.json /local/openapi.json\nRUN openapi-generator-cli generate \\\n -i /local/openapi.json \\\n -g python \\\n -o /local/onyx_openapi_client \\\n --package-name onyx_openapi_client \\\n --skip-validate-spec \\\n --openapi-normalizer \"SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true\"\n\nFROM integration-base AS integration\nCOPY --from=openapi-schema /app/openapi.json /app/generated/openapi.json\nCOPY --from=openapi-client /local/onyx_openapi_client /app/generated/onyx_openapi_client\n\n\nENV PYTHONPATH=/app\n\nENTRYPOINT [\"pytest\", \"-s\", \"-rs\"]\nCMD [\"/app/tests/integration\", \"--ignore=/app/tests/integration/multitenant_tests\"]\n" }, { "path": "backend/tests/integration/README.md", "content": "# Integration Tests\n\n## General Testing Overview\n\nThe integration tests are designed with a \"manager\" class and a \"test\" class for each type of object being manipulated (e.g., user, persona, credential):\n\n- **Manager Class**: Contains methods for each type of API call. Responsible for creating, deleting, and verifying the existence of an entity.\n- **Test Class**: Stores data for each entity being tested. This is our \"expected state\" of the object.\n\nThe idea is that each test can use the manager class to create (.create()) a \"test*\" object. It can then perform an operation on the object (e.g., send a request to the API) and then check if the \"test*\" object is in the expected state by using the manager class (.verify()) function.\n\n## Instructions for Running Integration Tests Locally\n0. Generate dependencies\nFirst install openap-generator\n```sh\nbrew install openapi-generator\n```\n\nThen, using the VSCode/Cursor debugger, run the `Onyx OpenAPI Schema Generator` task (see `CONTRIBUTING_VSCODE.md` for `launch.json` setup instructions).\nThe task automatically generates the Python client needed for integration tests.\n\nIf the client generation fails, try running this command manually:\n```sh\nopenapi-generator generate -i backend/generated/openapi.json -g python -o backend/generated/onyx_openapi_client --package-name onyx_openapi_client --skip-validate-spec --openapi-normalizer \"SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true\"\n```\n\n1. Launch onyx (using Docker or running with a debugger), ensuring the API server is running on port 8080.\n - If you'd like to set environment variables, you can do so by creating a `.env` file in the onyx/backend/tests/integration/ directory.\n - Onyx MUST be launched with AUTH_TYPE=basic and ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true\n - Tests that use `mock_llm_response` (e.g. llm workflow tool call tests) also require `INTEGRATION_TESTS_MODE=true` on the API server process.\n2. Navigate to `onyx/backend`.\n3. Run the following command in the terminal:\n ```sh\n python -m dotenv -f .env run -- pytest -s tests/integration/tests/\n ```\n or to run all tests in a file:\n ```sh\n python -m dotenv -f .env run -- pytest -s tests/integration/tests/path_to/test_file.py\n ```\n or to run a single test:\n ```sh\n python -m dotenv -f .env run -- pytest -s tests/integration/tests/path_to/test_file.py::test_function_name\n ```\n\nRunning some single tests require the `mock_connector_server` container to be running. If the above doesn't work, \nnavigate to `backend/tests/integration/mock_services` and run\n```sh\ndocker compose -f docker-compose.mock-it-services.yml -p mock-it-services-stack up -d\n```\nYou will have to modify the networks section of the docker-compose file to `_default` if you brought up the standard\nonyx services with a name different from the default `onyx`.\n\n## Guidelines for Writing Integration Tests\n\n- As authentication is currently required for all tests, each test should start by creating a user.\n- Each test should ideally focus on a single API flow.\n- The test writer should try to consider failure cases and edge cases for the flow and write the tests to check for these cases.\n- Every step of the test should be commented describing what is being done and what the expected behavior is.\n- A summary of the test should be given at the top of the test function as well!\n- When writing new tests, manager classes, manager functions, and test classes, try to copy the style of the other ones that have already been written.\n- Be careful for scope creep!\n - No need to overcomplicate every test by verifying after every single API call so long as the case you would be verifying is covered elsewhere (ideally in a test focused on covering that case).\n - An example of this is: Creating an admin user is done at the beginning of nearly every test, but we only need to verify that the user is actually an admin in the test focused on checking admin permissions. For every other test, we can just create the admin user and assume that the permissions are working as expected.\n\n## Current Testing Limitations\n\n### Test coverage\n\n- All tests are probably not as high coverage as they could be.\n- The \"connector\" tests in particular are super bare bones because we will be reworking connector/cc_pair sometime soon.\n- Global Curator role is not thoroughly tested.\n- No auth is not tested at all.\n\n### Failure checking\n\n- While we test expected auth failures, we only check that it failed at all.\n- We dont check that the return codes are what we expect.\n- This means that a test could be failing for a different reason than expected.\n- We should ensure that the proper codes are being returned for each failure case.\n- We should also query the db after each failure to ensure that the db is in the expected state.\n\n### Scope/focus\n\n- The tests may be scoped sub-optimally.\n- The scoping of each test may be overlapping.\n\n## Current Testing Coverage\n\nThe current testing coverage should be checked by reading the comments at the top of each test file.\n\n## TODO: Testing Coverage\n\n- Persona permissions testing\n- Read only (and/or basic) user permissions\n - Ensuring proper permission enforcement using the chat/doc_search endpoints\n- No auth\n\n## Ideas for integration testing design\n\n### Combine the \"test\" and \"manager\" classes\n\nThis 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.\n\n### Rework VespaClient\n\nRight now, its used a fixture and has to be passed around between manager classes.\nCould just be built where its used\n" }, { "path": "backend/tests/integration/__init__.py", "content": "" }, { "path": "backend/tests/integration/common_utils/chat.py", "content": "import requests\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import User\n\n\ndef test_create_chat_session_and_send_messages() -> None:\n # Create a test user\n with get_session_with_current_tenant() as db_session:\n test_user = User(email=\"test@example.com\", hashed_password=\"dummy_hash\")\n db_session.add(test_user)\n db_session.commit()\n\n base_url = \"http://localhost:8080\" # Adjust this to your API's base URL\n headers = {\"Authorization\": f\"Bearer {test_user.id}\"}\n\n # Create a new chat session\n create_session_response = requests.post(\n f\"{base_url}/chat/create-chat-session\",\n json={\n \"description\": \"Test Chat\",\n \"persona_id\": 1,\n }, # Assuming persona_id 1 exists\n headers=headers,\n )\n assert create_session_response.status_code == 200\n chat_session_id = create_session_response.json()[\"chat_session_id\"]\n\n # Send first message\n first_message = \"Hello, this is a test message.\"\n send_message_response = requests.post(\n f\"{base_url}/chat/send-chat-message\",\n json={\n \"chat_session_id\": chat_session_id,\n \"message\": first_message,\n \"retrieval_options\": {\"top_k\": 3},\n \"stream_response\": False,\n },\n headers=headers,\n )\n assert send_message_response.status_code == 200\n\n # Send second message\n second_message = \"Can you provide more information?\"\n send_message_response = requests.post(\n f\"{base_url}/chat/send-chat-message\",\n json={\n \"chat_session_id\": chat_session_id,\n \"message\": second_message,\n \"retrieval_options\": {\"top_k\": 3},\n \"stream_response\": False,\n },\n headers=headers,\n )\n assert send_message_response.status_code == 200\n\n # Verify chat session details\n get_session_response = requests.get(\n f\"{base_url}/chat/get-chat-session/{chat_session_id}\", headers=headers\n )\n assert get_session_response.status_code == 200\n session_details = get_session_response.json()\n assert session_details[\"chat_session_id\"] == chat_session_id\n assert session_details[\"description\"] == \"Test Chat\"\n assert len(session_details[\"messages\"]) == 4 # 2 user messages + 2 AI responses\n" }, { "path": "backend/tests/integration/common_utils/config.py", "content": "import generated.onyx_openapi_client.onyx_openapi_client as onyx_api # type: ignore[import-untyped,unused-ignore]\nfrom tests.integration.common_utils.constants import API_SERVER_URL\n\napi_config = onyx_api.Configuration(host=API_SERVER_URL)\n" }, { "path": "backend/tests/integration/common_utils/constants.py", "content": "import os\n\nADMIN_USER_NAME = \"admin_user\"\n\nAPI_SERVER_PROTOCOL = os.getenv(\"API_SERVER_PROTOCOL\") or \"http\"\nAPI_SERVER_HOST = os.getenv(\"API_SERVER_HOST\") or \"127.0.0.1\"\nAPI_SERVER_PORT = os.getenv(\"API_SERVER_PORT\") or \"8080\"\nAPI_SERVER_URL = f\"{API_SERVER_PROTOCOL}://{API_SERVER_HOST}:{API_SERVER_PORT}\"\nMAX_DELAY = 300\n\nMCP_SERVER_HOST = os.getenv(\"MCP_SERVER_HOST\") or \"127.0.0.1\"\nMCP_SERVER_PORT = os.getenv(\"MCP_SERVER_PORT\") or \"8090\"\nMCP_SERVER_URL = f\"{API_SERVER_PROTOCOL}://{MCP_SERVER_HOST}:{MCP_SERVER_PORT}\"\n\nGENERAL_HEADERS = {\"Content-Type\": \"application/json\"}\n\nNUM_DOCS = 5\n\nMOCK_CONNECTOR_SERVER_HOST = os.getenv(\"MOCK_CONNECTOR_SERVER_HOST\") or \"localhost\"\nMOCK_CONNECTOR_SERVER_PORT = os.getenv(\"MOCK_CONNECTOR_SERVER_PORT\") or 8001\n" }, { "path": "backend/tests/integration/common_utils/document_acl.py", "content": "\"\"\"\nUtilities for testing document access control lists (ACLs) and permissions.\n\"\"\"\n\nfrom typing import List\nfrom uuid import UUID\n\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom ee.onyx.access.access import _get_access_for_documents\nfrom ee.onyx.db.external_perm import fetch_external_groups_for_user\nfrom onyx.access.utils import prefix_external_group\nfrom onyx.access.utils import prefix_user_email\nfrom onyx.configs.constants import PUBLIC_DOC_PAT\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom onyx.db.models import User\nfrom onyx.db.users import fetch_user_by_id\nfrom onyx.utils.logger import setup_logger\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestUser\n\nlogger = setup_logger()\n\n\ndef get_user_acl(user: User, db_session: Session) -> set[str]:\n \"\"\"\n Get the ACL entries for a user, including their external groups, email, and public doc pattern.\n\n Args:\n user: The user object\n db_session: Database session\n\n Returns:\n Set of ACL entries for the user\n \"\"\"\n db_external_groups = (\n fetch_external_groups_for_user(db_session, user.id) if user else []\n )\n prefixed_external_groups = [\n prefix_external_group(db_external_group.external_user_group_id)\n for db_external_group in db_external_groups\n ]\n\n user_acl = set(prefixed_external_groups)\n user_acl.update({prefix_user_email(user.email), PUBLIC_DOC_PAT})\n return user_acl\n\n\ndef get_user_document_access_via_acl(\n test_user: DATestUser, document_ids: List[str], db_session: Session\n) -> List[str]:\n \"\"\"\n Determine which documents a user can access by comparing user ACL with document ACLs.\n\n This is a more reliable method than search-based verification as it directly checks\n permission logic without depending on search relevance or ranking.\n\n Args:\n test_user: The test user to check access for\n document_ids: List of document IDs to check\n db_session: Database session\n\n Returns:\n List of document IDs that the user can access\n \"\"\"\n # Get the actual User object from the database\n user = fetch_user_by_id(db_session, UUID(test_user.id))\n if not user:\n logger.error(f\"Could not find user with ID {test_user.id}\")\n return []\n\n user_acl = get_user_acl(user, db_session)\n logger.info(f\"User {user.email} ACL entries: {user_acl}\")\n\n # Get document access information\n doc_access_map = _get_access_for_documents(document_ids, db_session)\n logger.info(f\"Found access info for {len(doc_access_map)} documents\")\n\n accessible_docs = []\n for doc_id, doc_access in doc_access_map.items():\n doc_acl = doc_access.to_acl()\n logger.info(f\"Document {doc_id} ACL: {doc_acl}\")\n\n # Check if user has any matching ACL entry\n if user_acl.intersection(doc_acl):\n accessible_docs.append(doc_id)\n logger.info(f\"User {user.email} has access to document {doc_id}\")\n else:\n logger.info(f\"User {user.email} does NOT have access to document {doc_id}\")\n\n return accessible_docs\n\n\ndef get_all_connector_documents(\n cc_pair: DATestCCPair, db_session: Session\n) -> List[str]:\n \"\"\"\n Get all document IDs for a given connector/credential pair.\n\n Args:\n cc_pair: The connector-credential pair\n db_session: Database session\n\n Returns:\n List of document IDs\n \"\"\"\n stmt = select(DocumentByConnectorCredentialPair.id).where(\n DocumentByConnectorCredentialPair.connector_id == cc_pair.connector_id,\n DocumentByConnectorCredentialPair.credential_id == cc_pair.credential_id,\n )\n\n result = db_session.execute(stmt)\n document_ids = [row[0] for row in result.fetchall()]\n logger.info(\n f\"Found {len(document_ids)} documents for connector {cc_pair.connector_id}\"\n )\n\n return document_ids\n\n\ndef get_documents_by_permission_type(\n document_ids: List[str], db_session: Session\n) -> List[str]:\n \"\"\"\n Categorize documents by their permission types and return public documents.\n\n Args:\n document_ids: List of document IDs to check\n db_session: Database session\n\n Returns:\n List of document IDs that are public\n \"\"\"\n doc_access_map = _get_access_for_documents(document_ids, db_session)\n\n public_docs = []\n\n for doc_id, doc_access in doc_access_map.items():\n if doc_access.is_public:\n public_docs.append(doc_id)\n\n return public_docs\n" }, { "path": "backend/tests/integration/common_utils/managers/api_key.py", "content": "from uuid import uuid4\n\nimport requests\n\nfrom onyx.db.models import UserRole\nfrom onyx.server.api_key.models import APIKeyArgs\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass APIKeyManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n name: str | None = None,\n api_key_role: UserRole = UserRole.ADMIN,\n ) -> DATestAPIKey:\n name = f\"{name}-api-key\" if name else f\"test-api-key-{uuid4()}\"\n api_key_request = APIKeyArgs(\n name=name,\n role=api_key_role,\n )\n api_key_response = requests.post(\n f\"{API_SERVER_URL}/admin/api-key\",\n json=api_key_request.model_dump(),\n headers=user_performing_action.headers,\n )\n api_key_response.raise_for_status()\n api_key = api_key_response.json()\n result_api_key = DATestAPIKey(\n api_key_id=api_key[\"api_key_id\"],\n api_key_display=api_key[\"api_key_display\"],\n api_key=api_key[\"api_key\"],\n api_key_name=name,\n api_key_role=api_key_role,\n user_id=api_key[\"user_id\"],\n headers=GENERAL_HEADERS,\n )\n result_api_key.headers[\"Authorization\"] = f\"Bearer {result_api_key.api_key}\"\n return result_api_key\n\n @staticmethod\n def delete(\n api_key: DATestAPIKey,\n user_performing_action: DATestUser,\n ) -> None:\n api_key_response = requests.delete(\n f\"{API_SERVER_URL}/admin/api-key/{api_key.api_key_id}\",\n headers=user_performing_action.headers,\n )\n api_key_response.raise_for_status()\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[DATestAPIKey]:\n api_key_response = requests.get(\n f\"{API_SERVER_URL}/admin/api-key\",\n headers=user_performing_action.headers,\n )\n api_key_response.raise_for_status()\n return [DATestAPIKey(**api_key) for api_key in api_key_response.json()]\n\n @staticmethod\n def verify(\n api_key: DATestAPIKey,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n retrieved_keys = APIKeyManager.get_all(\n user_performing_action=user_performing_action\n )\n for key in retrieved_keys:\n if key.api_key_id == api_key.api_key_id:\n if verify_deleted:\n raise ValueError(\"API Key found when it should have been deleted\")\n if (\n key.api_key_name == api_key.api_key_name\n and key.api_key_role == api_key.api_key_role\n ):\n return\n\n if not verify_deleted:\n raise Exception(\"API Key not found\")\n" }, { "path": "backend/tests/integration/common_utils/managers/cc_pair.py", "content": "import time\nfrom datetime import datetime\nfrom typing import Any\nfrom uuid import uuid4\n\nimport requests\n\nimport generated.onyx_openapi_client.onyx_openapi_client as api # type: ignore[import-untyped,unused-ignore]\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import ConnectorCredentialPairStatus\nfrom onyx.server.documents.models import CCPairFullInfo\nfrom onyx.server.documents.models import ConnectorCredentialPairIdentifier\nfrom onyx.server.documents.models import ConnectorIndexingStatusLite\nfrom onyx.server.documents.models import ConnectorStatus\nfrom onyx.server.documents.models import DocumentSource\nfrom onyx.server.documents.models import DocumentSyncStatus\nfrom tests.integration.common_utils.config import api_config\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import MAX_DELAY\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _cc_pair_creator(\n connector_id: int,\n credential_id: int,\n user_performing_action: DATestUser,\n name: str | None = None,\n access_type: AccessType = AccessType.PUBLIC,\n groups: list[int] | None = None,\n) -> DATestCCPair:\n name = f\"{name}-cc-pair\" if name else f\"test-cc-pair-{uuid4()}\"\n\n with api.ApiClient(api_config) as api_client:\n api_instance = api.DefaultApi(api_client)\n connector_credential_pair_metadata = api.ConnectorCredentialPairMetadata(\n name=name, access_type=access_type, groups=groups or []\n )\n api_response: api.StatusResponseInt = (\n api_instance.associate_credential_to_connector(\n connector_id,\n credential_id,\n connector_credential_pair_metadata,\n _headers=user_performing_action.headers,\n )\n )\n\n return DATestCCPair(\n id=int(api_response.data),\n name=name,\n connector_id=connector_id,\n credential_id=credential_id,\n access_type=access_type,\n groups=groups or [],\n )\n\n\nclass CCPairManager:\n @staticmethod\n def create_from_scratch(\n user_performing_action: DATestUser,\n name: str | None = None,\n access_type: AccessType = AccessType.PUBLIC,\n groups: list[int] | None = None,\n source: DocumentSource = DocumentSource.FILE,\n input_type: InputType = InputType.LOAD_STATE,\n connector_specific_config: dict[str, Any] | None = None,\n credential_json: dict[str, Any] | None = None,\n refresh_freq: int | None = None,\n ) -> DATestCCPair:\n connector = ConnectorManager.create(\n user_performing_action=user_performing_action,\n name=name,\n source=source,\n input_type=input_type,\n connector_specific_config=connector_specific_config,\n access_type=access_type,\n groups=groups,\n refresh_freq=refresh_freq,\n )\n credential = CredentialManager.create(\n user_performing_action=user_performing_action,\n credential_json=credential_json,\n name=name,\n source=source,\n curator_public=(access_type == AccessType.PUBLIC),\n groups=groups,\n )\n cc_pair = _cc_pair_creator(\n connector_id=connector.id,\n credential_id=credential.id,\n name=name,\n access_type=access_type,\n groups=groups,\n user_performing_action=user_performing_action,\n )\n return cc_pair\n\n @staticmethod\n def create(\n connector_id: int,\n credential_id: int,\n user_performing_action: DATestUser,\n name: str | None = None,\n access_type: AccessType = AccessType.PUBLIC,\n groups: list[int] | None = None,\n ) -> DATestCCPair:\n cc_pair = _cc_pair_creator(\n connector_id=connector_id,\n credential_id=credential_id,\n name=name,\n access_type=access_type,\n groups=groups,\n user_performing_action=user_performing_action,\n )\n return cc_pair\n\n @staticmethod\n def pause_cc_pair(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> None:\n result = requests.put(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/status\",\n json={\"status\": \"PAUSED\"},\n headers=user_performing_action.headers,\n )\n result.raise_for_status()\n\n @staticmethod\n def unpause_cc_pair(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> None:\n result = requests.put(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/status\",\n json={\"status\": \"ACTIVE\"},\n headers=user_performing_action.headers,\n )\n result.raise_for_status()\n\n @staticmethod\n def delete(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> None:\n cc_pair_identifier = ConnectorCredentialPairIdentifier(\n connector_id=cc_pair.connector_id,\n credential_id=cc_pair.credential_id,\n )\n result = requests.post(\n url=f\"{API_SERVER_URL}/manage/admin/deletion-attempt\",\n json=cc_pair_identifier.model_dump(),\n headers=user_performing_action.headers,\n )\n result.raise_for_status()\n\n @staticmethod\n def get_single(\n cc_pair_id: int,\n user_performing_action: DATestUser,\n ) -> CCPairFullInfo | None:\n response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n cc_pair_json = response.json()\n return CCPairFullInfo(**cc_pair_json)\n\n @staticmethod\n def get_indexing_status_by_id(\n cc_pair_id: int,\n user_performing_action: DATestUser,\n ) -> ConnectorIndexingStatusLite | None:\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/connector/indexing-status\",\n headers=user_performing_action.headers,\n json={\"get_all_connectors\": True},\n )\n response.raise_for_status()\n indexing_status_response = response.json()\n for connectors_by_source in indexing_status_response:\n connectors = connectors_by_source[\"indexing_statuses\"]\n for connector in connectors:\n if connector[\"cc_pair_id\"] == cc_pair_id:\n return ConnectorIndexingStatusLite(**connector)\n\n return None\n\n @staticmethod\n def get_indexing_statuses(\n user_performing_action: DATestUser,\n ) -> list[ConnectorIndexingStatusLite]:\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/connector/indexing-status\",\n headers=user_performing_action.headers,\n json={\"get_all_connectors\": True},\n )\n response.raise_for_status()\n indexing_status_response = response.json()\n indexing_statuses = []\n for connectors_by_source in indexing_status_response:\n connectors = connectors_by_source[\"indexing_statuses\"]\n for connector in connectors:\n indexing_statuses.append(ConnectorIndexingStatusLite(**connector))\n return indexing_statuses\n\n @staticmethod\n def get_connector_statuses(\n user_performing_action: DATestUser,\n ) -> list[ConnectorStatus]:\n response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/connector/status\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [ConnectorStatus(**status) for status in response.json()]\n\n @staticmethod\n def verify(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n all_cc_pairs = CCPairManager.get_connector_statuses(user_performing_action)\n for retrieved_cc_pair in all_cc_pairs:\n if retrieved_cc_pair.cc_pair_id == cc_pair.id:\n if verify_deleted:\n # We assume that this check will be performed after the deletion is\n # already waited for\n raise ValueError(\n f\"CC pair {cc_pair.id} found but should be deleted\"\n )\n if (\n retrieved_cc_pair.name == cc_pair.name\n and retrieved_cc_pair.connector.id == cc_pair.connector_id\n and retrieved_cc_pair.credential.id == cc_pair.credential_id\n and retrieved_cc_pair.access_type == cc_pair.access_type\n and set(retrieved_cc_pair.groups) == set(cc_pair.groups)\n ):\n return\n\n if not verify_deleted:\n raise ValueError(f\"CC pair {cc_pair.id} not found\")\n\n @staticmethod\n def run_once(\n cc_pair: DATestCCPair,\n from_beginning: bool,\n user_performing_action: DATestUser,\n ) -> None:\n body = {\n \"connector_id\": cc_pair.connector_id,\n \"credential_ids\": [cc_pair.credential_id],\n \"from_beginning\": from_beginning,\n }\n result = requests.post(\n url=f\"{API_SERVER_URL}/manage/admin/connector/run-once\",\n json=body,\n headers=user_performing_action.headers,\n )\n result.raise_for_status()\n\n @staticmethod\n def wait_for_indexing_inactive(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n timeout: float = MAX_DELAY,\n ) -> None:\n \"\"\"wait for the number of docs to be indexed on the connector.\n This is used to test pausing a connector in the middle of indexing and\n terminating that indexing.\"\"\"\n print(f\"Indexing wait for inactive starting: cc_pair={cc_pair.id}\")\n start = time.monotonic()\n while True:\n fetched_cc_pairs = CCPairManager.get_indexing_statuses(\n user_performing_action\n )\n for fetched_cc_pair in fetched_cc_pairs:\n if fetched_cc_pair.cc_pair_id != cc_pair.id:\n continue\n\n if fetched_cc_pair.in_progress:\n continue\n\n print(f\"Indexing is inactive: cc_pair={cc_pair.id}\")\n return\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"Indexing wait for inactive timed out: cc_pair={cc_pair.id} timeout={timeout}s\"\n )\n\n print(\n f\"Indexing wait for inactive still waiting: cc_pair={cc_pair.id} elapsed={elapsed:.2f} timeout={timeout}s\"\n )\n time.sleep(5)\n\n @staticmethod\n def wait_for_indexing_in_progress(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n timeout: float = MAX_DELAY,\n num_docs: int = 16,\n ) -> None:\n \"\"\"wait for the number of docs to be indexed on the connector.\n This is used to test pausing a connector in the middle of indexing and\n terminating that indexing.\"\"\"\n start = time.monotonic()\n while True:\n fetched_cc_pairs = CCPairManager.get_indexing_statuses(\n user_performing_action\n )\n for fetched_cc_pair in fetched_cc_pairs:\n if fetched_cc_pair.cc_pair_id != cc_pair.id:\n continue\n\n if not fetched_cc_pair.in_progress:\n continue\n\n if fetched_cc_pair.docs_indexed < num_docs:\n print(\n f\"Indexing in progress: cc_pair={cc_pair.id} \"\n f\"docs_indexed={fetched_cc_pair.docs_indexed} num_docs={num_docs}\"\n )\n continue\n\n if fetched_cc_pair.docs_indexed >= num_docs:\n print(\n \"Indexed at least the requested number of docs: \"\n f\"cc_pair={cc_pair.id} \"\n f\"docs_indexed={fetched_cc_pair.docs_indexed} \"\n f\"num_docs={num_docs}\"\n )\n return\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"Indexing in progress wait timed out: cc_pair={cc_pair.id} timeout={timeout}s\"\n )\n\n print(\n f\"Indexing in progress waiting: cc_pair={cc_pair.id} elapsed={elapsed:.2f} timeout={timeout}s\"\n )\n time.sleep(5)\n\n @staticmethod\n def wait_for_indexing_completion(\n cc_pair: DATestCCPair,\n after: datetime,\n user_performing_action: DATestUser,\n timeout: float = MAX_DELAY,\n ) -> None:\n \"\"\"after: Wait for an indexing success time after this time\"\"\"\n start = time.monotonic()\n while True:\n fetched_cc_pairs = CCPairManager.get_indexing_statuses(\n user_performing_action\n )\n for fetched_cc_pair in fetched_cc_pairs:\n if fetched_cc_pair.cc_pair_id != cc_pair.id:\n continue\n\n if fetched_cc_pair.in_progress:\n continue\n\n if (\n fetched_cc_pair.last_success\n and fetched_cc_pair.last_success > after\n ):\n print(f\"Indexing complete: cc_pair={cc_pair.id}\")\n return\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"Indexing wait timed out: cc_pair={cc_pair.id} timeout={timeout}s\"\n )\n\n print(\n f\"Indexing wait for completion: cc_pair={cc_pair.id} elapsed={elapsed:.2f} timeout={timeout}s\"\n )\n time.sleep(5)\n\n @staticmethod\n def prune(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> None:\n result = requests.post(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/prune\",\n headers=user_performing_action.headers,\n )\n result.raise_for_status()\n\n @staticmethod\n def last_pruned(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> datetime | None:\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/last_pruned\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n response_str = response.json()\n\n # If the response itself is a datetime string, parse it\n if not isinstance(response_str, str):\n return None\n\n try:\n return datetime.fromisoformat(response_str)\n except ValueError:\n return None\n\n @staticmethod\n def wait_for_prune(\n cc_pair: DATestCCPair,\n after: datetime,\n user_performing_action: DATestUser,\n timeout: float = MAX_DELAY,\n ) -> None:\n \"\"\"after: The task register time must be after this time.\"\"\"\n start = time.monotonic()\n while True:\n last_pruned = CCPairManager.last_pruned(cc_pair, user_performing_action)\n if last_pruned and last_pruned > after:\n print(f\"Pruning complete: cc_pair={cc_pair.id}\")\n break\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"CC pair pruning was not completed within {timeout} seconds\"\n )\n\n print(\n f\"Waiting for CC pruning to complete. elapsed={elapsed:.2f} timeout={timeout}\"\n )\n time.sleep(5)\n\n @staticmethod\n def sync(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> None:\n \"\"\"This function triggers a permission sync.\n Naming / intent of this function probably could use improvement, but currently it's letting\n 409 Conflict pass through since if it's running that's what we were trying to do anyway.\n \"\"\"\n result = requests.post(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/sync-permissions\",\n headers=user_performing_action.headers,\n )\n if result.status_code != 409:\n result.raise_for_status()\n\n group_sync_result = requests.post(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/sync-groups\",\n headers=user_performing_action.headers,\n )\n if group_sync_result.status_code != 409:\n group_sync_result.raise_for_status()\n time.sleep(2)\n\n @staticmethod\n def get_doc_sync_task(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> datetime | None:\n doc_sync_response = requests.get(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/sync-permissions\",\n headers=user_performing_action.headers,\n )\n doc_sync_response.raise_for_status()\n doc_sync_response_str = doc_sync_response.json()\n\n # If the response itself is a datetime string, parse it\n if not isinstance(doc_sync_response_str, str):\n return None\n\n try:\n return datetime.fromisoformat(doc_sync_response_str)\n except ValueError:\n return None\n\n @staticmethod\n def get_group_sync_task(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> datetime | None:\n group_sync_response = requests.get(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/sync-groups\",\n headers=user_performing_action.headers,\n )\n group_sync_response.raise_for_status()\n group_sync_response_str = group_sync_response.json()\n\n # If the response itself is a datetime string, parse it\n if not isinstance(group_sync_response_str, str):\n return None\n\n try:\n return datetime.fromisoformat(group_sync_response_str)\n except ValueError:\n return None\n\n @staticmethod\n def get_doc_sync_statuses(\n cc_pair: DATestCCPair,\n user_performing_action: DATestUser,\n ) -> list[DocumentSyncStatus]:\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/get-docs-sync-status\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n doc_sync_statuses: list[DocumentSyncStatus] = []\n for doc_sync_status in response.json():\n last_synced = doc_sync_status.get(\"last_synced\")\n if last_synced:\n last_synced = datetime.fromisoformat(last_synced)\n\n last_modified = doc_sync_status.get(\"last_modified\")\n if last_modified:\n last_modified = datetime.fromisoformat(last_modified)\n\n doc_sync_statuses.append(\n DocumentSyncStatus(\n doc_id=doc_sync_status[\"doc_id\"],\n last_synced=last_synced,\n last_modified=last_modified,\n )\n )\n\n return doc_sync_statuses\n\n @staticmethod\n def wait_for_sync(\n cc_pair: DATestCCPair,\n after: datetime,\n user_performing_action: DATestUser,\n timeout: float = MAX_DELAY,\n number_of_updated_docs: int = 0,\n # Sometimes waiting for a group sync is not necessary\n should_wait_for_group_sync: bool = True,\n # Sometimes waiting for a vespa sync is not necessary\n should_wait_for_vespa_sync: bool = True,\n ) -> None:\n \"\"\"after: The task register time must be after this time.\"\"\"\n doc_synced = False\n group_synced = False\n start = time.monotonic()\n while True:\n # We are treating both syncs as part of one larger permission sync job\n doc_last_synced = CCPairManager.get_doc_sync_task(\n cc_pair, user_performing_action\n )\n group_last_synced = CCPairManager.get_group_sync_task(\n cc_pair, user_performing_action\n )\n\n if not doc_synced and doc_last_synced and doc_last_synced > after:\n print(f\"doc_last_synced: {doc_last_synced}\")\n print(f\"sync command start time: {after}\")\n print(f\"permission sync complete: cc_pair={cc_pair.id}\")\n doc_synced = True\n\n if not group_synced and group_last_synced and group_last_synced > after:\n print(f\"group_last_synced: {group_last_synced}\")\n print(f\"sync command start time: {after}\")\n print(f\"group sync complete: cc_pair={cc_pair.id}\")\n group_synced = True\n\n if doc_synced and (group_synced or not should_wait_for_group_sync):\n break\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"Permission sync was not completed within {timeout} seconds\"\n )\n\n print(\n f\"Waiting for CC sync to complete. elapsed={elapsed:.2f} timeout={timeout}\"\n )\n time.sleep(5)\n\n # TODO: remove this sleep,\n # this shouldnt be necessary but something is off with the timing for the sync jobs\n time.sleep(5)\n\n if not should_wait_for_vespa_sync:\n return\n\n print(\"waiting for vespa sync\")\n # wait for the vespa sync to complete once the permission sync is complete\n start = time.monotonic()\n while True:\n doc_sync_statuses = CCPairManager.get_doc_sync_statuses(\n cc_pair=cc_pair,\n user_performing_action=user_performing_action,\n )\n synced_docs = 0\n for doc_sync_status in doc_sync_statuses:\n if (\n doc_sync_status.last_synced is not None\n and doc_sync_status.last_modified is not None\n and doc_sync_status.last_synced >= doc_sync_status.last_modified\n and doc_sync_status.last_synced >= after\n and doc_sync_status.last_modified >= after\n ):\n synced_docs += 1\n\n if synced_docs >= number_of_updated_docs:\n print(f\"all docs synced: cc_pair={cc_pair.id}\")\n break\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"Vespa sync was not completed within {timeout} seconds\"\n )\n\n print(\n f\"Waiting for vespa sync to complete. elapsed={elapsed:.2f} timeout={timeout}\"\n )\n time.sleep(5)\n\n @staticmethod\n def wait_for_deletion_completion(\n user_performing_action: DATestUser,\n cc_pair_id: int | None = None,\n ) -> None:\n \"\"\"if cc_pair_id is not specified, just waits until no connectors are in the deleting state.\n if cc_pair_id is specified, checks to ensure the specific cc_pair_id is gone.\n We had a bug where the connector was paused in the middle of deleting, so specifying the\n cc_pair_id is good to do.\"\"\"\n start = time.monotonic()\n while True:\n cc_pairs = CCPairManager.get_indexing_statuses(user_performing_action)\n if cc_pair_id:\n found = False\n for cc_pair in cc_pairs:\n if cc_pair.cc_pair_id == cc_pair_id:\n found = True\n break\n\n if not found:\n return\n else:\n if all(\n cc_pair.cc_pair_status != ConnectorCredentialPairStatus.DELETING\n for cc_pair in cc_pairs\n ):\n return\n\n if time.monotonic() - start > MAX_DELAY:\n raise TimeoutError(\n f\"CC pairs deletion was not completed within the {MAX_DELAY} seconds\"\n )\n else:\n print(\"Some CC pairs are still being deleted, waiting...\")\n time.sleep(2)\n" }, { "path": "backend/tests/integration/common_utils/managers/chat.py", "content": "import json\nfrom typing import Any\nfrom typing import cast\nfrom typing import Literal\nfrom typing import TypedDict\nfrom uuid import UUID\n\nimport requests\nfrom requests.models import Response\n\nfrom onyx.context.search.models import SavedSearchDoc\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.llm.override_models import LLMOverride\nfrom onyx.server.query_and_chat.models import AUTO_PLACE_AFTER_LATEST_MESSAGE\nfrom onyx.server.query_and_chat.models import ChatSessionCreationRequest\nfrom onyx.server.query_and_chat.models import SendMessageRequest\nfrom onyx.server.query_and_chat.streaming_models import StreamingType\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestChatMessage\nfrom tests.integration.common_utils.test_models import DATestChatSession\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import ErrorResponse\nfrom tests.integration.common_utils.test_models import StreamedResponse\nfrom tests.integration.common_utils.test_models import ToolCallDebug\nfrom tests.integration.common_utils.test_models import ToolName\nfrom tests.integration.common_utils.test_models import ToolResult\n\n\nclass StreamPacketObj(TypedDict, total=False):\n \"\"\"Base structure for streaming packet objects.\"\"\"\n\n type: Literal[\n \"message_start\",\n \"message_delta\",\n \"search_tool_start\",\n \"search_tool_queries_delta\",\n \"search_tool_documents_delta\",\n \"image_generation_start\",\n \"image_generation_heartbeat\",\n \"image_generation_final\",\n \"tool_call_debug\",\n ]\n content: str\n final_documents: list[dict[str, Any]]\n is_internet_search: bool\n images: list[dict[str, Any]]\n queries: list[str]\n documents: list[dict[str, Any]]\n tool_call_id: str\n tool_name: str\n tool_args: dict[str, Any]\n\n\nclass PlacementData(TypedDict, total=False):\n \"\"\"Structure for packet placement information.\"\"\"\n\n turn_index: int\n tab_index: int\n sub_turn_index: int | None\n\n\nclass StreamPacketData(TypedDict, total=False):\n \"\"\"Structure for streaming response packets.\"\"\"\n\n reserved_assistant_message_id: int\n error: str\n stack_trace: str\n obj: StreamPacketObj\n placement: PlacementData\n\n\nclass ChatSessionManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n persona_id: int = 0,\n description: str = \"Test chat session\",\n project_id: int | None = None,\n ) -> DATestChatSession:\n chat_session_creation_req = ChatSessionCreationRequest(\n persona_id=persona_id,\n description=description,\n project_id=project_id,\n )\n response = requests.post(\n f\"{API_SERVER_URL}/chat/create-chat-session\",\n json=chat_session_creation_req.model_dump(),\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n chat_session_id = response.json()[\"chat_session_id\"]\n return DATestChatSession(\n id=chat_session_id, persona_id=persona_id, description=description\n )\n\n @staticmethod\n def send_message(\n chat_session_id: UUID,\n message: str,\n user_performing_action: DATestUser,\n parent_message_id: int | None = None,\n file_descriptors: list[FileDescriptor] | None = None,\n allowed_tool_ids: list[int] | None = None,\n forced_tool_ids: list[int] | None = None,\n chat_session: DATestChatSession | None = None,\n mock_llm_response: str | None = None,\n deep_research: bool = False,\n llm_override: LLMOverride | None = None,\n ) -> StreamedResponse:\n chat_message_req = SendMessageRequest(\n message=message,\n chat_session_id=chat_session_id,\n parent_message_id=(\n parent_message_id\n if parent_message_id is not None\n else AUTO_PLACE_AFTER_LATEST_MESSAGE\n ),\n file_descriptors=file_descriptors or [],\n allowed_tool_ids=allowed_tool_ids,\n forced_tool_id=forced_tool_ids[0] if forced_tool_ids else None,\n mock_llm_response=mock_llm_response,\n deep_research=deep_research,\n llm_override=llm_override,\n )\n\n response = requests.post(\n f\"{API_SERVER_URL}/chat/send-chat-message\",\n json=chat_message_req.model_dump(mode=\"json\"),\n headers=user_performing_action.headers,\n stream=True,\n cookies=user_performing_action.cookies,\n )\n\n streamed_response = ChatSessionManager.analyze_response(response)\n\n if not chat_session:\n return streamed_response\n\n # TODO: ideally we would get the research answer purpose from the chat history\n # but atm the field needed would not be used outside of testing, so we're not adding it.\n # chat_history = ChatSessionManager.get_chat_history(\n # chat_session=chat_session,\n # user_performing_action=user_performing_action,\n # )\n\n # for message_obj in chat_history:\n # if message_obj.message_type == MessageType.ASSISTANT:\n # streamed_response.research_answer_purpose = (\n # message_obj.research_answer_purpose\n # )\n # streamed_response.assistant_message_id = message_obj.id\n # break\n\n return streamed_response\n\n @staticmethod\n def send_message_with_disconnect(\n chat_session_id: UUID,\n message: str,\n user_performing_action: DATestUser,\n disconnect_after_packets: int = 0,\n parent_message_id: int | None = None,\n file_descriptors: list[FileDescriptor] | None = None,\n allowed_tool_ids: list[int] | None = None,\n forced_tool_ids: list[int] | None = None,\n mock_llm_response: str | None = None,\n deep_research: bool = False,\n llm_override: LLMOverride | None = None,\n ) -> None:\n \"\"\"\n Send a message and simulate client disconnect before stream completes.\n\n This is useful for testing how the server handles client disconnections\n during streaming responses.\n\n Args:\n chat_session_id: The chat session ID\n message: The message to send\n disconnect_after_packets: Disconnect after receiving this many packets.\n ... (other standard message parameters)\n\n Returns:\n None. Caller can verify server-side cleanup via get_chat_history etc.\n \"\"\"\n chat_message_req = SendMessageRequest(\n message=message,\n chat_session_id=chat_session_id,\n parent_message_id=(\n parent_message_id\n if parent_message_id is not None\n else AUTO_PLACE_AFTER_LATEST_MESSAGE\n ),\n file_descriptors=file_descriptors or [],\n allowed_tool_ids=allowed_tool_ids,\n forced_tool_id=forced_tool_ids[0] if forced_tool_ids else None,\n mock_llm_response=mock_llm_response,\n deep_research=deep_research,\n llm_override=llm_override,\n )\n\n packets_received = 0\n\n with requests.post(\n f\"{API_SERVER_URL}/chat/send-chat-message\",\n json=chat_message_req.model_dump(mode=\"json\"),\n headers=user_performing_action.headers,\n stream=True,\n cookies=user_performing_action.cookies,\n ) as response:\n for line in response.iter_lines():\n if not line:\n continue\n\n packets_received += 1\n if packets_received > disconnect_after_packets:\n break\n\n return None\n\n @staticmethod\n def analyze_response(response: Response) -> StreamedResponse:\n response_data = cast(\n list[StreamPacketData],\n [\n json.loads(line.decode(\"utf-8\"))\n for line in response.iter_lines()\n if line\n ],\n )\n ind_to_tool_use: dict[int, ToolResult] = {}\n tool_call_debug: list[ToolCallDebug] = []\n top_documents: list[SearchDoc] = []\n heartbeat_packets: list[StreamPacketData] = []\n full_message = \"\"\n assistant_message_id: int | None = None\n error = None\n ind: int\n for data in response_data:\n if reserved_id := data.get(\"reserved_assistant_message_id\"):\n assistant_message_id = reserved_id\n elif data.get(\"error\"):\n error = ErrorResponse(\n error=str(data[\"error\"]),\n stack_trace=str(data.get(\"stack_trace\") or \"\"),\n )\n elif (error_obj := cast(dict[str, Any], data.get(\"obj\") or {})) and (\n error_obj.get(\"error\")\n or error_obj.get(\"type\") == StreamingType.ERROR.value\n ):\n error = ErrorResponse(\n error=str(error_obj.get(\"error\") or \"Streaming error\"),\n stack_trace=str(\n error_obj.get(\"stack_trace\") or data.get(\"stack_trace\") or \"\"\n ),\n )\n elif (\n (data_obj := data.get(\"obj\"))\n and (packet_type := data_obj.get(\"type\"))\n and (\n ind := cast(\n int,\n (\n data.get(\"ind\")\n if data.get(\"ind\") is not None\n else data.get(\"placement\", {}).get(\"turn_index\")\n ),\n )\n )\n is not None\n ):\n packet_type_str = str(packet_type)\n if packet_type_str == StreamingType.MESSAGE_START.value:\n final_docs = data_obj.get(\"final_documents\")\n if isinstance(final_docs, list):\n top_documents = [SearchDoc(**doc) for doc in final_docs]\n full_message += data_obj.get(\"content\", \"\")\n elif packet_type_str == StreamingType.MESSAGE_DELTA.value:\n full_message += data_obj[\"content\"]\n elif packet_type_str == StreamingType.SEARCH_TOOL_START.value:\n tool_name = (\n ToolName.INTERNET_SEARCH\n if data_obj.get(\"is_internet_search\", False)\n else ToolName.INTERNAL_SEARCH\n )\n ind_to_tool_use[ind] = ToolResult(\n tool_name=tool_name,\n )\n elif packet_type_str == StreamingType.IMAGE_GENERATION_START.value:\n ind_to_tool_use[ind] = ToolResult(\n tool_name=ToolName.IMAGE_GENERATION,\n )\n elif packet_type_str == StreamingType.IMAGE_GENERATION_HEARTBEAT.value:\n # Track heartbeat packets for debugging/testing\n heartbeat_packets.append(data)\n elif packet_type_str == StreamingType.IMAGE_GENERATION_FINAL.value:\n from tests.integration.common_utils.test_models import (\n GeneratedImage,\n )\n\n images = data_obj.get(\"images\", [])\n ind_to_tool_use[ind].images.extend(\n [GeneratedImage(**img) for img in images]\n )\n elif packet_type_str == StreamingType.SEARCH_TOOL_QUERIES_DELTA.value:\n ind_to_tool_use[ind].queries.extend(data_obj.get(\"queries\", []))\n elif packet_type_str == StreamingType.SEARCH_TOOL_DOCUMENTS_DELTA.value:\n docs = []\n for doc in data_obj.get(\"documents\", []):\n if \"db_doc_id\" in doc:\n # Already a SavedSearchDoc format\n docs.append(SavedSearchDoc(**doc))\n else:\n # SearchDoc format - Convert to SavedSearchDoc\n search_doc = SearchDoc(**doc)\n docs.append(\n SavedSearchDoc.from_search_doc(search_doc, db_doc_id=0)\n )\n ind_to_tool_use[ind].documents.extend(docs)\n elif packet_type_str == StreamingType.TOOL_CALL_DEBUG.value:\n tool_call_debug.append(\n ToolCallDebug(\n tool_call_id=str(data_obj.get(\"tool_call_id\", \"\")),\n tool_name=str(data_obj.get(\"tool_name\", \"\")),\n tool_args=cast(\n dict[str, Any], data_obj.get(\"tool_args\") or {}\n ),\n )\n )\n # If there's an error, assistant_message_id might not be present\n if not assistant_message_id and not error:\n raise ValueError(\"Assistant message id not found\")\n return StreamedResponse(\n full_message=full_message,\n assistant_message_id=assistant_message_id or -1, # Use -1 for error cases\n top_documents=top_documents,\n used_tools=list(ind_to_tool_use.values()),\n tool_call_debug=tool_call_debug,\n heartbeat_packets=[dict(packet) for packet in heartbeat_packets],\n error=error,\n )\n\n @staticmethod\n def get_chat_history(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> list[DATestChatMessage]:\n response = requests.get(\n f\"{API_SERVER_URL}/chat/get-chat-session/{chat_session.id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n return [\n DATestChatMessage(\n id=msg[\"message_id\"],\n chat_session_id=chat_session.id,\n parent_message_id=msg.get(\"parent_message\"),\n message=msg[\"message\"],\n message_type=msg.get(\"message_type\"),\n files=msg.get(\"files\"),\n )\n for msg in response.json()[\"messages\"]\n ]\n\n @staticmethod\n def create_chat_message_feedback(\n message_id: int,\n is_positive: bool,\n user_performing_action: DATestUser,\n feedback_text: str | None = None,\n predefined_feedback: str | None = None,\n ) -> None:\n response = requests.post(\n url=f\"{API_SERVER_URL}/chat/create-chat-message-feedback\",\n json={\n \"chat_message_id\": message_id,\n \"is_positive\": is_positive,\n \"feedback_text\": feedback_text,\n \"predefined_feedback\": predefined_feedback,\n },\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def delete(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"\n Delete a chat session and all its related records (messages, agent data, etc.)\n Uses the default deletion method configured on the server.\n\n Returns True if deletion was successful, False otherwise.\n \"\"\"\n response = requests.delete(\n f\"{API_SERVER_URL}/chat/delete-chat-session/{chat_session.id}\",\n headers=user_performing_action.headers,\n )\n return response.ok\n\n @staticmethod\n def soft_delete(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"\n Soft delete a chat session (marks as deleted but keeps in database).\n\n Returns True if deletion was successful, False otherwise.\n \"\"\"\n # Since there's no direct API for soft delete, we'll use a query parameter approach\n # or make a direct call with hard_delete=False parameter via a new endpoint\n response = requests.delete(\n f\"{API_SERVER_URL}/chat/delete-chat-session/{chat_session.id}?hard_delete=false\",\n headers=user_performing_action.headers,\n )\n return response.ok\n\n @staticmethod\n def hard_delete(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"\n Hard delete a chat session (completely removes from database).\n\n Returns True if deletion was successful, False otherwise.\n \"\"\"\n response = requests.delete(\n f\"{API_SERVER_URL}/chat/delete-chat-session/{chat_session.id}?hard_delete=true\",\n headers=user_performing_action.headers,\n )\n return response.ok\n\n @staticmethod\n def verify_deleted(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"\n Verify that a chat session has been deleted by attempting to retrieve it.\n\n Returns True if the chat session is confirmed deleted, False if it still exists.\n \"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/chat/get-chat-session/{chat_session.id}\",\n headers=user_performing_action.headers,\n )\n # Chat session should return 404 if it doesn't exist or is deleted\n return response.status_code == 404\n\n @staticmethod\n def verify_soft_deleted(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"\n Verify that a chat session has been soft deleted (marked as deleted but still in DB).\n\n Returns True if the chat session is soft deleted, False otherwise.\n \"\"\"\n # Try to get the chat session with include_deleted=true\n response = requests.get(\n f\"{API_SERVER_URL}/chat/get-chat-session/{chat_session.id}?include_deleted=true\",\n headers=user_performing_action.headers,\n )\n\n if response.status_code == 200:\n # Chat exists, check if it's marked as deleted\n chat_data = response.json()\n return chat_data.get(\"deleted\", False) is True\n return False\n\n @staticmethod\n def verify_hard_deleted(\n chat_session: DATestChatSession,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"\n Verify that a chat session has been hard deleted (completely removed from DB).\n\n Returns True if the chat session is hard deleted, False otherwise.\n \"\"\"\n # Try to get the chat session with include_deleted=true\n response = requests.get(\n f\"{API_SERVER_URL}/chat/get-chat-session/{chat_session.id}?include_deleted=true\",\n headers=user_performing_action.headers,\n )\n\n # For hard delete, even with include_deleted=true, the record should not exist\n return response.status_code != 200\n" }, { "path": "backend/tests/integration/common_utils/managers/connector.py", "content": "from typing import Any\nfrom uuid import uuid4\n\nimport requests\n\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import ConnectorUpdateRequest\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass ConnectorManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n name: str | None = None,\n source: DocumentSource = DocumentSource.FILE,\n input_type: InputType = InputType.LOAD_STATE,\n connector_specific_config: dict[str, Any] | None = None,\n access_type: AccessType = AccessType.PUBLIC,\n groups: list[int] | None = None,\n refresh_freq: int | None = None,\n ) -> DATestConnector:\n name = f\"{name}-connector\" if name else f\"test-connector-{uuid4()}\"\n\n connector_update_request = ConnectorUpdateRequest(\n name=name,\n source=source,\n input_type=input_type,\n connector_specific_config=(\n connector_specific_config\n or (\n {\n \"file_locations\": [],\n \"file_names\": [],\n \"zip_metadata_file_id\": None,\n }\n if source == DocumentSource.FILE\n else {}\n )\n ),\n access_type=access_type,\n groups=groups or [],\n refresh_freq=refresh_freq,\n )\n\n response = requests.post(\n url=f\"{API_SERVER_URL}/manage/admin/connector\",\n json=connector_update_request.model_dump(),\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n response_data = response.json()\n return DATestConnector(\n id=response_data.get(\"id\"),\n name=name,\n source=source,\n input_type=input_type,\n connector_specific_config=connector_specific_config or {},\n groups=groups,\n access_type=access_type,\n )\n\n @staticmethod\n def edit(\n connector: DATestConnector,\n user_performing_action: DATestUser,\n ) -> None:\n response = requests.patch(\n url=f\"{API_SERVER_URL}/manage/admin/connector/{connector.id}\",\n json=connector.model_dump(exclude={\"id\"}),\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def delete(\n connector: DATestConnector,\n user_performing_action: DATestUser,\n ) -> None:\n response = requests.delete(\n url=f\"{API_SERVER_URL}/manage/admin/connector/{connector.id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[DATestConnector]:\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/connector\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [\n DATestConnector(\n id=conn.get(\"id\"),\n name=conn.get(\"name\", \"\"),\n source=conn.get(\"source\", DocumentSource.FILE),\n input_type=conn.get(\"input_type\", InputType.LOAD_STATE),\n connector_specific_config=conn.get(\"connector_specific_config\", {}),\n )\n for conn in response.json()\n ]\n\n @staticmethod\n def get(\n connector_id: int,\n user_performing_action: DATestUser,\n ) -> DATestConnector:\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/connector/{connector_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n conn = response.json()\n return DATestConnector(\n id=conn.get(\"id\"),\n name=conn.get(\"name\", \"\"),\n source=conn.get(\"source\", DocumentSource.FILE),\n input_type=conn.get(\"input_type\", InputType.LOAD_STATE),\n connector_specific_config=conn.get(\"connector_specific_config\", {}),\n )\n" }, { "path": "backend/tests/integration/common_utils/managers/credential.py", "content": "from typing import Any\nfrom uuid import uuid4\n\nimport requests\n\nfrom onyx.server.documents.models import CredentialSnapshot\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass CredentialManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n credential_json: dict[str, Any] | None = None,\n admin_public: bool = True,\n name: str | None = None,\n source: DocumentSource = DocumentSource.FILE,\n curator_public: bool = True,\n groups: list[int] | None = None,\n ) -> DATestCredential:\n name = f\"{name}-credential\" if name else f\"test-credential-{uuid4()}\"\n\n credential_request = {\n \"name\": name,\n \"credential_json\": credential_json or {},\n \"admin_public\": admin_public,\n \"source\": source,\n \"curator_public\": curator_public,\n \"groups\": groups or [],\n }\n\n response = requests.post(\n url=f\"{API_SERVER_URL}/manage/credential\",\n json=credential_request,\n headers=user_performing_action.headers,\n )\n\n response.raise_for_status()\n return DATestCredential(\n id=response.json()[\"id\"],\n name=name,\n credential_json=credential_json or {},\n admin_public=admin_public,\n source=source,\n curator_public=curator_public,\n groups=groups or [],\n )\n\n @staticmethod\n def edit(\n credential: DATestCredential,\n user_performing_action: DATestUser,\n ) -> None:\n request = credential.model_dump(include={\"name\", \"credential_json\"})\n response = requests.put(\n url=f\"{API_SERVER_URL}/manage/admin/credential/{credential.id}\",\n json=request,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def delete(\n credential: DATestCredential,\n user_performing_action: DATestUser,\n ) -> None:\n response = requests.delete(\n url=f\"{API_SERVER_URL}/manage/credential/{credential.id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def get(\n credential_id: int,\n user_performing_action: DATestUser,\n ) -> CredentialSnapshot:\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/credential/{credential_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return CredentialSnapshot(**response.json())\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[CredentialSnapshot]:\n response = requests.get(\n f\"{API_SERVER_URL}/manage/credential\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [CredentialSnapshot(**cred) for cred in response.json()]\n\n @staticmethod\n def verify(\n credential: DATestCredential,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n all_credentials = CredentialManager.get_all(user_performing_action)\n for fetched_credential in all_credentials:\n if credential.id == fetched_credential.id:\n if verify_deleted:\n raise ValueError(\n f\"Credential {credential.id} found but should be deleted\"\n )\n if (\n credential.name == fetched_credential.name\n and credential.admin_public == fetched_credential.admin_public\n and credential.source == fetched_credential.source\n and credential.curator_public == fetched_credential.curator_public\n ):\n return\n if not verify_deleted:\n raise ValueError(f\"Credential {credential.id} not found\")\n" }, { "path": "backend/tests/integration/common_utils/managers/discord_bot.py", "content": "\"\"\"Manager for Discord bot API integration tests.\"\"\"\n\nimport requests\n\nfrom onyx.db.discord_bot import create_channel_config\nfrom onyx.db.discord_bot import create_guild_config\nfrom onyx.db.discord_bot import register_guild\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.utils import DiscordChannelView\nfrom onyx.server.manage.discord_bot.utils import generate_discord_registration_key\nfrom shared_configs.contextvars import get_current_tenant_id\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestDiscordChannelConfig\nfrom tests.integration.common_utils.test_models import DATestDiscordGuildConfig\nfrom tests.integration.common_utils.test_models import DATestUser\n\nDISCORD_BOT_API_URL = f\"{API_SERVER_URL}/manage/admin/discord-bot\"\n\n\nclass DiscordBotManager:\n \"\"\"Manager for Discord bot API operations.\"\"\"\n\n # === Bot Config ===\n\n @staticmethod\n def get_bot_config(\n user_performing_action: DATestUser,\n ) -> dict:\n \"\"\"Get Discord bot config.\"\"\"\n response = requests.get(\n url=f\"{DISCORD_BOT_API_URL}/config\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def create_bot_config(\n bot_token: str,\n user_performing_action: DATestUser,\n ) -> dict:\n \"\"\"Create Discord bot config.\"\"\"\n response = requests.post(\n url=f\"{DISCORD_BOT_API_URL}/config\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n json={\"bot_token\": bot_token},\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def delete_bot_config(\n user_performing_action: DATestUser,\n ) -> dict:\n \"\"\"Delete Discord bot config.\"\"\"\n response = requests.delete(\n url=f\"{DISCORD_BOT_API_URL}/config\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n return response.json()\n\n # === Guild Config ===\n\n @staticmethod\n def list_guilds(\n user_performing_action: DATestUser,\n ) -> list[dict]:\n \"\"\"List all guild configs.\"\"\"\n response = requests.get(\n url=f\"{DISCORD_BOT_API_URL}/guilds\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def create_guild(\n user_performing_action: DATestUser,\n ) -> DATestDiscordGuildConfig:\n \"\"\"Create a new guild config with registration key.\"\"\"\n response = requests.post(\n url=f\"{DISCORD_BOT_API_URL}/guilds\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n data = response.json()\n return DATestDiscordGuildConfig(\n id=data[\"id\"],\n registration_key=data[\"registration_key\"],\n )\n\n @staticmethod\n def get_guild(\n config_id: int,\n user_performing_action: DATestUser,\n ) -> dict:\n \"\"\"Get a specific guild config.\"\"\"\n response = requests.get(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{config_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def update_guild(\n config_id: int,\n user_performing_action: DATestUser,\n enabled: bool | None = None,\n default_persona_id: int | None = None,\n ) -> dict:\n \"\"\"Update a guild config.\"\"\"\n # Fetch current guild config to get existing values\n current_guild = DiscordBotManager.get_guild(config_id, user_performing_action)\n\n # Build request body with required fields\n body: dict = {\n \"enabled\": enabled if enabled is not None else current_guild[\"enabled\"],\n \"default_persona_id\": (\n default_persona_id\n if default_persona_id is not None\n else current_guild.get(\"default_persona_id\")\n ),\n }\n\n response = requests.patch(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{config_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n json=body,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def delete_guild(\n config_id: int,\n user_performing_action: DATestUser,\n ) -> dict:\n \"\"\"Delete a guild config.\"\"\"\n response = requests.delete(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{config_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n return response.json()\n\n # === Channel Config ===\n\n @staticmethod\n def list_channels(\n guild_config_id: int,\n user_performing_action: DATestUser,\n ) -> list[DATestDiscordChannelConfig]:\n \"\"\"List all channel configs for a guild.\"\"\"\n response = requests.get(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{guild_config_id}/channels\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n response.raise_for_status()\n return [DATestDiscordChannelConfig(**c) for c in response.json()]\n\n @staticmethod\n def update_channel(\n guild_config_id: int,\n channel_config_id: int,\n user_performing_action: DATestUser,\n enabled: bool = False,\n thread_only_mode: bool = False,\n require_bot_invocation: bool = True,\n persona_override_id: int | None = None,\n ) -> DATestDiscordChannelConfig:\n \"\"\"Update a channel config.\n\n All fields are required by the API. Default values match the channel\n config defaults from create_channel_config.\n \"\"\"\n body: dict = {\n \"enabled\": enabled,\n \"thread_only_mode\": thread_only_mode,\n \"require_bot_invocation\": require_bot_invocation,\n \"persona_override_id\": persona_override_id,\n }\n\n response = requests.patch(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{guild_config_id}/channels/{channel_config_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n json=body,\n )\n response.raise_for_status()\n return DATestDiscordChannelConfig(**response.json())\n\n # === Utility methods for testing ===\n\n @staticmethod\n def create_registered_guild_in_db(\n guild_id: int,\n guild_name: str,\n ) -> DATestDiscordGuildConfig:\n \"\"\"Create a registered guild config directly in the database.\n\n This creates a guild that has already completed registration,\n with guild_id and guild_name set. Use this for testing channel\n endpoints which require a registered guild.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n tenant_id = get_current_tenant_id()\n registration_key = generate_discord_registration_key(tenant_id)\n config = create_guild_config(db_session, registration_key)\n config = register_guild(db_session, config, guild_id, guild_name)\n db_session.commit()\n\n return DATestDiscordGuildConfig(\n id=config.id,\n registration_key=registration_key,\n )\n\n @staticmethod\n def get_guild_or_none(\n config_id: int,\n user_performing_action: DATestUser,\n ) -> dict | None:\n \"\"\"Get a guild config, returning None if not found.\"\"\"\n response = requests.get(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{config_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n if response.status_code == 404:\n return None\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def delete_guild_if_exists(\n config_id: int,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"Delete a guild config if it exists. Returns True if deleted.\"\"\"\n response = requests.delete(\n url=f\"{DISCORD_BOT_API_URL}/guilds/{config_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n if response.status_code == 404:\n return False\n response.raise_for_status()\n return True\n\n @staticmethod\n def delete_bot_config_if_exists(\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"Delete bot config if it exists. Returns True if deleted.\"\"\"\n response = requests.delete(\n url=f\"{DISCORD_BOT_API_URL}/config\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n )\n if response.status_code == 404:\n return False\n response.raise_for_status()\n return True\n\n @staticmethod\n def create_test_channel_in_db(\n guild_config_id: int,\n channel_id: int,\n channel_name: str,\n channel_type: str = \"text\",\n is_private: bool = False,\n ) -> DATestDiscordChannelConfig:\n \"\"\"Create a test channel config directly in the database.\n\n This is needed because channels are normally synced from Discord,\n not created via API. For testing the channel API endpoints,\n we need to populate test data directly.\n \"\"\"\n with get_session_with_current_tenant() as db_session:\n channel_view = DiscordChannelView(\n channel_id=channel_id,\n channel_name=channel_name,\n channel_type=channel_type,\n is_private=is_private,\n )\n config = create_channel_config(db_session, guild_config_id, channel_view)\n db_session.commit()\n\n return DATestDiscordChannelConfig(\n id=config.id,\n guild_config_id=config.guild_config_id,\n channel_id=config.channel_id,\n channel_name=config.channel_name,\n channel_type=config.channel_type,\n is_private=config.is_private,\n enabled=config.enabled,\n thread_only_mode=config.thread_only_mode,\n require_bot_invocation=config.require_bot_invocation,\n persona_override_id=config.persona_override_id,\n )\n" }, { "path": "backend/tests/integration/common_utils/managers/document.py", "content": "from uuid import uuid4\n\nimport requests\nfrom sqlalchemy import and_\nfrom sqlalchemy import select\nfrom sqlalchemy.orm import Session\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.enums import AccessType\nfrom onyx.db.models import ConnectorCredentialPair\nfrom onyx.db.models import DocumentByConnectorCredentialPair\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import NUM_DOCS\nfrom tests.integration.common_utils.managers.api_key import DATestAPIKey\nfrom tests.integration.common_utils.managers.cc_pair import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import SimpleTestDocument\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef _verify_document_permissions(\n retrieved_doc: dict,\n cc_pair: DATestCCPair,\n doc_creating_user: DATestUser,\n doc_set_names: list[str] | None = None,\n group_names: list[str] | None = None,\n) -> None:\n acl_keys = set(retrieved_doc.get(\"access_control_list\", {}).keys())\n print(f\"ACL keys: {acl_keys}\")\n\n if cc_pair.access_type == AccessType.PUBLIC:\n if \"PUBLIC\" not in acl_keys:\n raise ValueError(\n f\"Document {retrieved_doc['document_id']} is public but does not have the PUBLIC ACL key\"\n )\n\n if f\"user_email:{doc_creating_user.email}\" not in acl_keys:\n raise ValueError(\n f\"Document {retrieved_doc['document_id']} was created by user\"\n f\" {doc_creating_user.email} but does not have the user_email:{doc_creating_user.email} ACL key\"\n )\n\n if group_names is not None:\n expected_group_keys = {f\"group:{group_name}\" for group_name in group_names}\n found_group_keys = {key for key in acl_keys if key.startswith(\"group:\")}\n if found_group_keys != expected_group_keys:\n raise ValueError(\n f\"Document {retrieved_doc['document_id']} has incorrect group ACL keys. \"\n f\"Expected: {expected_group_keys} Found: {found_group_keys}\\n\"\n f\"All ACL keys: {acl_keys}\"\n )\n\n if doc_set_names is not None:\n found_doc_set_names = set(retrieved_doc.get(\"document_sets\", {}).keys())\n if found_doc_set_names != set(doc_set_names):\n raise ValueError(\n f\"Document set names mismatch. \\nFound: {found_doc_set_names}, \\nExpected: {set(doc_set_names)}\"\n )\n\n\ndef _generate_dummy_document(\n document_id: str,\n cc_pair_id: int,\n content: str | None = None,\n extra_metadata: dict | None = None,\n) -> dict:\n text = content if content else f\"This is test document {document_id}\"\n\n metadata: dict = {\"document_id\": document_id}\n if extra_metadata:\n metadata.update(extra_metadata)\n\n return {\n \"document\": {\n \"id\": document_id,\n \"sections\": [\n {\n \"text\": text,\n \"link\": f\"{document_id}\",\n }\n ],\n \"source\": DocumentSource.NOT_APPLICABLE,\n \"metadata\": metadata,\n \"semantic_identifier\": f\"Test Document {document_id}\",\n \"from_ingestion_api\": True,\n },\n \"cc_pair_id\": cc_pair_id,\n }\n\n\nclass DocumentManager:\n \"\"\"\n Manager for seeding documents via the ingestion API.\n Used to test various connector features.\n \"\"\"\n\n @staticmethod\n def seed_dummy_docs(\n cc_pair: DATestCCPair,\n api_key: DATestAPIKey,\n num_docs: int = NUM_DOCS,\n document_ids: list[str] | None = None,\n ) -> list[SimpleTestDocument]:\n # Use provided document_ids if available, otherwise generate random UUIDs\n if document_ids is None:\n document_ids = [f\"test-doc-{uuid4()}\" for _ in range(num_docs)]\n else:\n num_docs = len(document_ids)\n # Create and ingest some documents\n documents: list[dict] = []\n for document_id in document_ids:\n document = _generate_dummy_document(document_id, cc_pair.id)\n documents.append(document)\n response = requests.post(\n f\"{API_SERVER_URL}/onyx-api/ingestion\",\n json=document,\n headers=api_key.headers,\n )\n response.raise_for_status()\n\n print(\n f\"Seeding docs for api_key_id={api_key.api_key_id} completed successfully.\"\n )\n return [\n SimpleTestDocument(\n id=document[\"document\"][\"id\"],\n content=document[\"document\"][\"sections\"][0][\"text\"],\n )\n for document in documents\n ]\n\n @staticmethod\n def seed_doc_with_content(\n cc_pair: DATestCCPair,\n content: str,\n api_key: DATestAPIKey,\n document_id: str | None = None,\n metadata: dict | None = None,\n ) -> SimpleTestDocument:\n # Use provided document_ids if available, otherwise generate random UUIDs\n if document_id is None:\n document_id = f\"test-doc-{uuid4()}\"\n # Create and ingest some documents\n document: dict = _generate_dummy_document(\n document_id,\n cc_pair.id,\n content,\n extra_metadata=metadata,\n )\n response = requests.post(\n f\"{API_SERVER_URL}/onyx-api/ingestion\",\n json=document,\n headers=api_key.headers,\n )\n response.raise_for_status()\n\n print(\n f\"Seeding doc for api_key_id={api_key.api_key_id} completed successfully.\"\n )\n\n return SimpleTestDocument(\n id=document[\"document\"][\"id\"],\n content=document[\"document\"][\"sections\"][0][\"text\"],\n )\n\n @staticmethod\n def verify(\n vespa_client: vespa_fixture,\n cc_pair: DATestCCPair,\n doc_creating_user: DATestUser,\n # If None, will not check doc sets or groups\n # If empty list, will check for empty doc sets or groups\n doc_set_names: list[str] | None = None,\n group_names: list[str] | None = None,\n verify_deleted: bool = False,\n ) -> None:\n doc_ids = [document.id for document in cc_pair.documents]\n retrieved_docs_dict = vespa_client.get_documents_by_id(doc_ids)[\"documents\"]\n\n retrieved_docs = {\n doc[\"fields\"][\"document_id\"]: doc[\"fields\"] for doc in retrieved_docs_dict\n }\n\n # NOTE(rkuo): too much log spam\n # Left this here for debugging purposes.\n # import json\n\n # print(\"DEBUGGING DOCUMENTS\")\n # print(retrieved_docs)\n # for doc in retrieved_docs.values():\n # printable_doc = doc.copy()\n # print(printable_doc.keys())\n # printable_doc.pop(\"embeddings\")\n # printable_doc.pop(\"title_embedding\")\n # print(json.dumps(printable_doc, indent=2))\n\n for document in cc_pair.documents:\n retrieved_doc = retrieved_docs.get(document.id)\n if not retrieved_doc:\n if not verify_deleted:\n print(f\"Document not found: {document.id}\")\n print(retrieved_docs.keys())\n print(retrieved_docs.values())\n raise ValueError(f\"Document not found: {document.id}\")\n continue\n if verify_deleted:\n raise ValueError(\n f\"Document found when it should be deleted: {document.id}\"\n )\n _verify_document_permissions(\n retrieved_doc,\n cc_pair,\n doc_creating_user,\n doc_set_names,\n group_names,\n )\n\n @staticmethod\n def fetch_documents_for_cc_pair(\n cc_pair_id: int,\n db_session: Session,\n vespa_client: vespa_fixture,\n ) -> list[SimpleTestDocument]:\n stmt = (\n select(DocumentByConnectorCredentialPair)\n .join(\n ConnectorCredentialPair,\n and_(\n DocumentByConnectorCredentialPair.connector_id\n == ConnectorCredentialPair.connector_id,\n DocumentByConnectorCredentialPair.credential_id\n == ConnectorCredentialPair.credential_id,\n ),\n )\n .where(ConnectorCredentialPair.id == cc_pair_id)\n )\n documents = db_session.execute(stmt).scalars().all()\n if not documents:\n return []\n\n doc_ids = [document.id for document in documents]\n retrieved_docs_dict = vespa_client.get_documents_by_id(doc_ids)[\"documents\"]\n\n final_docs: list[SimpleTestDocument] = []\n # NOTE: they are really chunks, but we're assuming that for these tests\n # we only have one chunk per document for now\n for doc_dict in retrieved_docs_dict:\n doc_id = doc_dict[\"fields\"][\"document_id\"]\n doc_content = doc_dict[\"fields\"][\"content\"]\n # still called `image_file_name` in Vespa for backwards compatibility\n image_file_id = doc_dict[\"fields\"].get(\"image_file_name\", None)\n final_docs.append(\n SimpleTestDocument(\n id=doc_id, content=doc_content, image_file_id=image_file_id\n )\n )\n\n return final_docs\n\n\nclass IngestionManager(DocumentManager):\n \"\"\"\n Manager for additional ingestion API endpoints not covered by DocumentManager.\n Used specifically to test the ingestion API.\n \"\"\"\n\n @staticmethod\n def list_all_ingestion_docs(\n api_key: DATestAPIKey,\n ) -> list[dict]:\n response = requests.get(\n f\"{API_SERVER_URL}/onyx-api/ingestion\",\n headers=api_key.headers,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def delete(\n document_id: str,\n api_key: DATestAPIKey,\n ) -> None:\n response = requests.delete(\n f\"{API_SERVER_URL}/onyx-api/ingestion/{document_id}\",\n headers=api_key.headers,\n )\n response.raise_for_status()\n print(f\"Deleted document {document_id} successfully.\")\n" }, { "path": "backend/tests/integration/common_utils/managers/document_search.py", "content": "import requests\n\nfrom ee.onyx.server.query_and_chat.models import SearchFullResponse\nfrom ee.onyx.server.query_and_chat.models import SendSearchQueryRequest\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass DocumentSearchManager:\n @staticmethod\n def search_documents(\n query: str,\n user_performing_action: DATestUser,\n ) -> list[str]:\n \"\"\"\n Search for documents using the EE search API.\n\n Args:\n query: The search query string\n user_performing_action: The user performing the search (for auth)\n\n Returns:\n A list of document content strings (blurbs) from the search results\n \"\"\"\n search_request = SendSearchQueryRequest(\n search_query=query,\n filters=None,\n stream=False,\n )\n result = requests.post(\n url=f\"{API_SERVER_URL}/search/send-search-message\",\n json=search_request.model_dump(),\n headers=user_performing_action.headers,\n )\n result.raise_for_status()\n result_json = result.json()\n search_response = SearchFullResponse(**result_json)\n\n # Return the blurbs as the document content\n # For small documents (like test docs), the blurb should contain the full content\n document_content_list: list[str] = [\n doc.blurb for doc in search_response.search_docs\n ]\n return document_content_list\n" }, { "path": "backend/tests/integration/common_utils/managers/document_set.py", "content": "import time\nfrom typing import Any\nfrom uuid import UUID\nfrom uuid import uuid4\n\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import MAX_DELAY\nfrom tests.integration.common_utils.test_models import DATestDocumentSet\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass DocumentSetManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n name: str | None = None,\n description: str | None = None,\n cc_pair_ids: list[int] | None = None,\n is_public: bool = True,\n users: list[str] | None = None,\n groups: list[int] | None = None,\n federated_connectors: list[dict[str, Any]] | None = None,\n ) -> DATestDocumentSet:\n if name is None:\n name = f\"test_doc_set_{str(uuid4())}\"\n\n doc_set_creation_request = {\n \"name\": name,\n \"description\": description or name,\n \"cc_pair_ids\": cc_pair_ids or [],\n \"is_public\": is_public,\n \"users\": [str(UUID(user_id)) for user_id in (users or [])],\n \"groups\": groups or [],\n \"federated_connectors\": federated_connectors or [],\n }\n\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/document-set\",\n json=doc_set_creation_request,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n return DATestDocumentSet(\n id=int(response.json()),\n name=name,\n description=description or name,\n cc_pair_ids=cc_pair_ids or [],\n is_public=is_public,\n is_up_to_date=True,\n users=users or [],\n groups=groups or [],\n federated_connectors=federated_connectors or [],\n )\n\n @staticmethod\n def edit(\n document_set: DATestDocumentSet,\n user_performing_action: DATestUser,\n ) -> bool:\n doc_set_update_request = {\n \"id\": document_set.id,\n \"description\": document_set.description,\n \"cc_pair_ids\": document_set.cc_pair_ids,\n \"is_public\": document_set.is_public,\n \"users\": [str(UUID(user_id)) for user_id in document_set.users],\n \"groups\": document_set.groups,\n \"federated_connectors\": document_set.federated_connectors,\n }\n response = requests.patch(\n f\"{API_SERVER_URL}/manage/admin/document-set\",\n json=doc_set_update_request,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return True\n\n @staticmethod\n def delete(\n document_set: DATestDocumentSet,\n user_performing_action: DATestUser,\n ) -> bool:\n response = requests.delete(\n f\"{API_SERVER_URL}/manage/admin/document-set/{document_set.id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return True\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[DATestDocumentSet]:\n response = requests.get(\n f\"{API_SERVER_URL}/manage/document-set\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [\n DATestDocumentSet(\n id=doc_set[\"id\"],\n name=doc_set[\"name\"],\n description=doc_set[\"description\"],\n cc_pair_ids=[cc_pair[\"id\"] for cc_pair in doc_set[\"cc_pair_summaries\"]],\n is_public=doc_set[\"is_public\"],\n is_up_to_date=doc_set[\"is_up_to_date\"],\n users=[str(user_id) for user_id in doc_set[\"users\"]],\n groups=doc_set[\"groups\"],\n federated_connectors=doc_set[\"federated_connector_summaries\"],\n )\n for doc_set in response.json()\n ]\n\n @staticmethod\n def wait_for_sync(\n user_performing_action: DATestUser,\n document_sets_to_check: list[DATestDocumentSet] | None = None,\n ) -> None:\n # wait for document sets to be synced\n start = time.time()\n while True:\n doc_sets = DocumentSetManager.get_all(user_performing_action)\n if document_sets_to_check:\n check_ids = {doc_set.id for doc_set in document_sets_to_check}\n doc_set_ids = {doc_set.id for doc_set in doc_sets}\n if not check_ids.issubset(doc_set_ids):\n raise RuntimeError(\"Document set not found\")\n doc_sets = [doc_set for doc_set in doc_sets if doc_set.id in check_ids]\n all_up_to_date = all(doc_set.is_up_to_date for doc_set in doc_sets)\n\n if all_up_to_date:\n print(\"Document sets synced successfully.\")\n break\n\n if time.time() - start > MAX_DELAY:\n not_synced_doc_sets = [\n doc_set for doc_set in doc_sets if not doc_set.is_up_to_date\n ]\n raise TimeoutError(\n f\"Document sets were not synced within the {MAX_DELAY} seconds. \"\n f\"Remaining unsynced document sets: {len(not_synced_doc_sets)}. \"\n f\"IDs: {[doc_set.id for doc_set in not_synced_doc_sets]}\"\n )\n else:\n not_synced_doc_sets = [\n doc_set for doc_set in doc_sets if not doc_set.is_up_to_date\n ]\n print(\n f\"Document sets were not synced yet, waiting... \"\n f\"{len(not_synced_doc_sets)}/{len(doc_sets)} document sets still syncing. \"\n f\"IDs: {[doc_set.id for doc_set in not_synced_doc_sets]}\"\n )\n\n time.sleep(2)\n\n @staticmethod\n def verify(\n document_set: DATestDocumentSet,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n doc_sets = DocumentSetManager.get_all(user_performing_action)\n for doc_set in doc_sets:\n if doc_set.id == document_set.id:\n if verify_deleted:\n raise ValueError(\n f\"Document set {document_set.id} found but should have been deleted\"\n )\n if (\n doc_set.name == document_set.name\n and set(doc_set.cc_pair_ids) == set(document_set.cc_pair_ids)\n and doc_set.is_public == document_set.is_public\n and set(doc_set.users) == set(document_set.users)\n and set(doc_set.groups) == set(document_set.groups)\n and doc_set.federated_connectors\n == document_set.federated_connectors\n ):\n return\n if not verify_deleted:\n raise ValueError(f\"Document set {document_set.id} not found\")\n" }, { "path": "backend/tests/integration/common_utils/managers/file.py", "content": "import io\nimport mimetypes\nfrom typing import cast\nfrom typing import IO\nfrom typing import List\nfrom typing import Tuple\n\nimport requests\n\nfrom onyx.file_store.models import FileDescriptor\nfrom onyx.server.documents.models import FileUploadResponse\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass FileManager:\n @staticmethod\n def upload_files(\n files: List[Tuple[str, IO]],\n user_performing_action: DATestUser,\n ) -> Tuple[List[FileDescriptor], str]:\n headers = user_performing_action.headers\n headers.pop(\"Content-Type\", None)\n\n files_param = []\n for filename, file_obj in files:\n mime_type, _ = mimetypes.guess_type(filename)\n if mime_type is None:\n mime_type = \"application/octet-stream\"\n files_param.append((\"files\", (filename, file_obj, mime_type)))\n\n response = requests.post(\n f\"{API_SERVER_URL}/user/projects/file/upload\",\n files=files_param,\n headers=headers,\n )\n\n if not response.ok:\n try:\n detail = response.json().get(\"detail\", response.text)\n except Exception:\n detail = response.text\n return (\n cast(List[FileDescriptor], []),\n f\"Failed to upload files - {detail}\",\n )\n\n response_json = response.json()\n # Convert UserFileSnapshot to FileDescriptor format\n file_descriptors: List[FileDescriptor] = []\n for user_file in response_json.get(\"user_files\", []):\n file_descriptors.append(\n {\n \"id\": user_file[\"file_id\"],\n \"type\": user_file[\"chat_file_type\"],\n \"name\": user_file[\"name\"],\n \"user_file_id\": str(user_file[\"id\"]),\n }\n )\n return file_descriptors, \"\"\n\n @staticmethod\n def fetch_uploaded_file(\n file_id: str,\n user_performing_action: DATestUser,\n ) -> bytes:\n response = requests.get(\n f\"{API_SERVER_URL}/chat/file/{file_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return response.content\n\n @staticmethod\n def upload_file_for_connector(\n file_path: str,\n file_name: str,\n user_performing_action: DATestUser,\n content_type: str = \"application/octet-stream\",\n ) -> FileUploadResponse:\n # Read the file content\n with open(file_path, \"rb\") as f:\n file_content = f.read()\n\n # Create a file-like object\n file_obj = io.BytesIO(file_content)\n\n # The 'files' form field expects a list of files\n files = [(\"files\", (file_name, file_obj, content_type))]\n\n # Use the user's headers but without Content-Type\n # as requests will set the correct multipart/form-data Content-Type for us\n headers = user_performing_action.headers.copy()\n if \"Content-Type\" in headers:\n del headers[\"Content-Type\"]\n\n # Make the request\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/connector/file/upload\",\n files=files,\n headers=headers,\n )\n\n if not response.ok:\n try:\n error_detail = response.json().get(\"detail\", \"Unknown error\")\n except Exception:\n error_detail = response.text\n\n raise Exception(\n f\"Unable to upload files - {error_detail} (Status code: {response.status_code})\"\n )\n\n response_json = response.json()\n return FileUploadResponse(**response_json)\n" }, { "path": "backend/tests/integration/common_utils/managers/image_generation.py", "content": "import json\nimport os\nfrom typing import Any\nfrom uuid import uuid4\n\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestImageGenerationConfig\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _serialize_custom_config(\n custom_config: dict[str, Any] | None,\n) -> dict[str, str] | None:\n \"\"\"Convert custom_config values to strings (API expects dict[str, str]).\"\"\"\n if custom_config is None:\n return None\n return {\n key: json.dumps(value) if not isinstance(value, str) else value\n for key, value in custom_config.items()\n }\n\n\nclass ImageGenerationConfigManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n image_provider_id: str | None = None,\n model_name: str = \"gpt-image-1\",\n provider: str = \"openai\",\n api_key: str | None = None,\n api_base: str | None = None,\n api_version: str | None = None,\n deployment_name: str | None = None,\n custom_config: dict[str, Any] | None = None,\n is_default: bool = False,\n ) -> DATestImageGenerationConfig:\n \"\"\"Create a new image generation config with new credentials.\"\"\"\n image_provider_id = image_provider_id or f\"test-provider-{uuid4()}\"\n\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n json={\n \"image_provider_id\": image_provider_id,\n \"model_name\": model_name,\n \"provider\": provider,\n \"api_key\": api_key or os.environ[\"OPENAI_API_KEY\"],\n \"api_base\": api_base,\n \"api_version\": api_version,\n \"deployment_name\": deployment_name,\n \"custom_config\": _serialize_custom_config(custom_config),\n \"is_default\": is_default,\n },\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n data = response.json()\n\n return DATestImageGenerationConfig(\n image_provider_id=data[\"image_provider_id\"],\n model_configuration_id=data[\"model_configuration_id\"],\n model_name=data[\"model_name\"],\n llm_provider_id=data[\"llm_provider_id\"],\n llm_provider_name=data[\"llm_provider_name\"],\n is_default=data[\"is_default\"],\n )\n\n @staticmethod\n def create_from_provider(\n source_llm_provider_id: int,\n user_performing_action: DATestUser,\n image_provider_id: str | None = None,\n model_name: str = \"gpt-image-1\",\n api_base: str | None = None,\n api_version: str | None = None,\n deployment_name: str | None = None,\n is_default: bool = False,\n ) -> DATestImageGenerationConfig:\n \"\"\"Create a new image generation config by cloning from an existing LLM provider.\"\"\"\n image_provider_id = image_provider_id or f\"test-provider-{uuid4()}\"\n\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n json={\n \"image_provider_id\": image_provider_id,\n \"model_name\": model_name,\n \"source_llm_provider_id\": source_llm_provider_id,\n \"api_base\": api_base,\n \"api_version\": api_version,\n \"deployment_name\": deployment_name,\n \"is_default\": is_default,\n },\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n data = response.json()\n\n return DATestImageGenerationConfig(\n image_provider_id=data[\"image_provider_id\"],\n model_configuration_id=data[\"model_configuration_id\"],\n model_name=data[\"model_name\"],\n llm_provider_id=data[\"llm_provider_id\"],\n llm_provider_name=data[\"llm_provider_name\"],\n is_default=data[\"is_default\"],\n )\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[DATestImageGenerationConfig]:\n \"\"\"Get all image generation configs.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [DATestImageGenerationConfig(**config) for config in response.json()]\n\n @staticmethod\n def get_credentials(\n image_provider_id: str,\n user_performing_action: DATestUser,\n ) -> dict:\n \"\"\"Get credentials for an image generation config.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/image-generation/config/{image_provider_id}/credentials\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def update(\n image_provider_id: str,\n model_name: str,\n user_performing_action: DATestUser,\n provider: str | None = None,\n api_key: str | None = None,\n source_llm_provider_id: int | None = None,\n api_base: str | None = None,\n api_version: str | None = None,\n deployment_name: str | None = None,\n ) -> DATestImageGenerationConfig:\n \"\"\"Update an existing image generation config.\"\"\"\n payload: dict = {\n \"model_name\": model_name,\n \"api_base\": api_base,\n \"api_version\": api_version,\n \"deployment_name\": deployment_name,\n }\n\n if source_llm_provider_id is not None:\n payload[\"source_llm_provider_id\"] = source_llm_provider_id\n elif api_key is not None and provider is not None:\n payload[\"provider\"] = provider\n payload[\"api_key\"] = api_key\n else:\n raise ValueError(\n f\"Either source_llm_provider_id or (api_key + provider) must be provided. \"\n f\"Got: source_llm_provider_id={source_llm_provider_id}, provider={provider}, api_key={'***' if api_key else None}\"\n )\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/image-generation/config/{image_provider_id}\",\n json=payload,\n headers=user_performing_action.headers,\n )\n if not response.ok:\n print(f\"Update failed with status {response.status_code}: {response.text}\")\n response.raise_for_status()\n data = response.json()\n\n return DATestImageGenerationConfig(\n image_provider_id=data[\"image_provider_id\"],\n model_configuration_id=data[\"model_configuration_id\"],\n model_name=data[\"model_name\"],\n llm_provider_id=data[\"llm_provider_id\"],\n llm_provider_name=data[\"llm_provider_name\"],\n is_default=data[\"is_default\"],\n )\n\n @staticmethod\n def delete(\n image_provider_id: str,\n user_performing_action: DATestUser,\n ) -> None:\n \"\"\"Delete an image generation config.\"\"\"\n response = requests.delete(\n f\"{API_SERVER_URL}/admin/image-generation/config/{image_provider_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def set_default(\n image_provider_id: str,\n user_performing_action: DATestUser,\n ) -> None:\n \"\"\"Set an image generation config as the default.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config/{image_provider_id}/default\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def verify(\n config: DATestImageGenerationConfig,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n \"\"\"Verify that a config exists (or doesn't exist if verify_deleted=True).\"\"\"\n all_configs = ImageGenerationConfigManager.get_all(user_performing_action)\n\n for fetched_config in all_configs:\n if fetched_config.image_provider_id == config.image_provider_id:\n if verify_deleted:\n raise ValueError(\n f\"ImageGenerationConfig {config.image_provider_id} found but should be deleted\"\n )\n # Verify the config matches\n if (\n fetched_config.model_name == config.model_name\n and fetched_config.is_default == config.is_default\n ):\n return\n\n if not verify_deleted:\n raise ValueError(\n f\"ImageGenerationConfig {config.image_provider_id} not found\"\n )\n" }, { "path": "backend/tests/integration/common_utils/managers/index_attempt.py", "content": "import time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom urllib.parse import urlencode\n\nimport requests\n\nfrom onyx.background.indexing.models import IndexAttemptErrorPydantic\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import IndexModelStatus\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.models import IndexingStatus\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.server.documents.models import IndexAttemptSnapshot\nfrom onyx.server.documents.models import PaginatedReturn\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import MAX_DELAY\nfrom tests.integration.common_utils.test_models import DATestIndexAttempt\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass IndexAttemptManager:\n @staticmethod\n def create_test_index_attempts(\n num_attempts: int,\n cc_pair_id: int,\n from_beginning: bool = False,\n status: IndexingStatus = IndexingStatus.SUCCESS,\n new_docs_indexed: int = 10,\n total_docs_indexed: int = 10,\n docs_removed_from_index: int = 0,\n error_msg: str | None = None,\n base_time: datetime | None = None,\n ) -> list[DATestIndexAttempt]:\n if base_time is None:\n base_time = datetime.now()\n\n attempts = []\n with get_session_with_current_tenant() as db_session:\n # Get the current search settings\n search_settings = get_current_search_settings(db_session)\n if (\n not search_settings\n or search_settings.status != IndexModelStatus.PRESENT\n ):\n raise ValueError(\"No current search settings found with PRESENT status\")\n\n for i in range(num_attempts):\n time_created = base_time - timedelta(hours=i)\n\n index_attempt = IndexAttempt(\n connector_credential_pair_id=cc_pair_id,\n from_beginning=from_beginning,\n status=status,\n new_docs_indexed=new_docs_indexed,\n total_docs_indexed=total_docs_indexed,\n docs_removed_from_index=docs_removed_from_index,\n error_msg=error_msg,\n time_created=time_created,\n time_started=time_created,\n time_updated=time_created,\n search_settings_id=search_settings.id,\n )\n\n db_session.add(index_attempt)\n db_session.flush() # To get the ID\n\n attempts.append(\n DATestIndexAttempt(\n id=index_attempt.id,\n status=index_attempt.status,\n new_docs_indexed=index_attempt.new_docs_indexed,\n total_docs_indexed=index_attempt.total_docs_indexed,\n docs_removed_from_index=index_attempt.docs_removed_from_index,\n error_msg=index_attempt.error_msg,\n time_started=index_attempt.time_started,\n time_updated=index_attempt.time_updated,\n )\n )\n\n db_session.commit()\n\n return attempts\n\n @staticmethod\n def get_index_attempt_page(\n cc_pair_id: int,\n user_performing_action: DATestUser,\n page: int = 0,\n page_size: int = 10,\n ) -> PaginatedReturn[IndexAttemptSnapshot]:\n query_params: dict[str, str | int] = {\n \"page_num\": page,\n \"page_size\": page_size,\n }\n\n url = f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}/index-attempts?{urlencode(query_params, doseq=True)}\"\n response = requests.get(\n url=url,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n data = response.json()\n return PaginatedReturn(\n items=[IndexAttemptSnapshot(**item) for item in data[\"items\"]],\n total_items=data[\"total_items\"],\n )\n\n @staticmethod\n def get_latest_index_attempt_for_cc_pair(\n cc_pair_id: int,\n user_performing_action: DATestUser,\n ) -> IndexAttemptSnapshot | None:\n \"\"\"Get an IndexAttempt by ID\"\"\"\n index_attempts = IndexAttemptManager.get_index_attempt_page(\n cc_pair_id, user_performing_action=user_performing_action\n ).items\n if not index_attempts:\n return None\n\n index_attempts = sorted(\n index_attempts, key=lambda x: x.time_started or \"0\", reverse=True\n )\n return index_attempts[0]\n\n @staticmethod\n def wait_for_index_attempt_start(\n cc_pair_id: int,\n user_performing_action: DATestUser,\n index_attempts_to_ignore: list[int] | None = None,\n timeout: float = MAX_DELAY,\n ) -> IndexAttemptSnapshot:\n \"\"\"Wait for an IndexAttempt to start\"\"\"\n start = datetime.now()\n index_attempts_to_ignore = index_attempts_to_ignore or []\n\n while True:\n index_attempt = IndexAttemptManager.get_latest_index_attempt_for_cc_pair(\n cc_pair_id=cc_pair_id,\n user_performing_action=user_performing_action,\n )\n if (\n index_attempt\n and index_attempt.time_started\n and index_attempt.id not in index_attempts_to_ignore\n ):\n return index_attempt\n\n elapsed = (datetime.now() - start).total_seconds()\n if elapsed > timeout:\n raise TimeoutError(\n f\"IndexAttempt for CC Pair {cc_pair_id} did not start within {timeout} seconds\"\n )\n\n @staticmethod\n def get_index_attempt_by_id(\n index_attempt_id: int,\n cc_pair_id: int,\n user_performing_action: DATestUser,\n ) -> IndexAttemptSnapshot:\n page_num = 0\n page_size = 10\n while True:\n page = IndexAttemptManager.get_index_attempt_page(\n cc_pair_id=cc_pair_id,\n page=page_num,\n page_size=page_size,\n user_performing_action=user_performing_action,\n )\n for attempt in page.items:\n if attempt.id == index_attempt_id:\n return attempt\n\n if len(page.items) < page_size:\n break\n\n page_num += 1\n\n raise ValueError(f\"IndexAttempt {index_attempt_id} not found\")\n\n @staticmethod\n def wait_for_index_attempt_completion(\n index_attempt_id: int,\n cc_pair_id: int,\n user_performing_action: DATestUser,\n timeout: float = MAX_DELAY,\n ) -> None:\n \"\"\"Wait for an IndexAttempt to complete\"\"\"\n start = time.monotonic()\n while True:\n index_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=index_attempt_id,\n cc_pair_id=cc_pair_id,\n user_performing_action=user_performing_action,\n )\n\n if index_attempt.status and index_attempt.status.is_terminal():\n print(\n f\"IndexAttempt {index_attempt_id} completed with status {index_attempt.status}\"\n )\n return\n\n elapsed = time.monotonic() - start\n if elapsed > timeout:\n raise TimeoutError(\n f\"IndexAttempt {index_attempt_id} did not complete within {timeout} seconds\"\n )\n\n print(\n f\"Waiting for IndexAttempt {index_attempt_id} to complete. elapsed={elapsed:.2f} timeout={timeout}\"\n )\n time.sleep(5)\n\n @staticmethod\n def get_index_attempt_errors_for_cc_pair(\n cc_pair_id: int,\n user_performing_action: DATestUser,\n include_resolved: bool = True,\n ) -> list[IndexAttemptErrorPydantic]:\n url = f\"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair_id}/errors?page_size=100\"\n if include_resolved:\n url += \"&include_resolved=true\"\n response = requests.get(\n url=url,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n data = response.json()\n return [IndexAttemptErrorPydantic(**item) for item in data[\"items\"]]\n" }, { "path": "backend/tests/integration/common_utils/managers/llm_provider.py", "content": "import os\nfrom uuid import uuid4\n\nimport requests\n\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.server.manage.llm.models import DefaultModel\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import LLMProviderView\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass LLMProviderManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n name: str | None = None,\n provider: str | None = None,\n api_key: str | None = None,\n default_model_name: str | None = None,\n api_base: str | None = None,\n api_version: str | None = None,\n groups: list[int] | None = None,\n personas: list[int] | None = None,\n is_public: bool | None = None,\n set_as_default: bool = True,\n ) -> DATestLLMProvider:\n print(f\"Seeding LLM Providers for {user_performing_action.email}...\")\n\n llm_provider = LLMProviderUpsertRequest(\n name=name or f\"test-provider-{uuid4()}\",\n provider=provider or LlmProviderNames.OPENAI,\n api_key=api_key or os.environ[\"OPENAI_API_KEY\"],\n api_base=api_base,\n api_version=api_version,\n custom_config=None,\n is_public=True if is_public is None else is_public,\n groups=groups or [],\n personas=personas or [],\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name or \"gpt-4o-mini\",\n is_visible=True,\n max_input_tokens=None,\n display_name=default_model_name or \"gpt-4o-mini\",\n supports_image_input=True,\n )\n ],\n api_key_changed=True,\n )\n\n llm_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n json=llm_provider.model_dump(),\n headers=user_performing_action.headers,\n )\n llm_response.raise_for_status()\n response_data = llm_response.json()\n\n result_llm = DATestLLMProvider(\n id=response_data[\"id\"],\n name=response_data[\"name\"],\n provider=response_data[\"provider\"],\n api_key=response_data[\"api_key\"],\n default_model_name=default_model_name or \"gpt-4o-mini\",\n is_public=response_data[\"is_public\"],\n is_auto_mode=response_data.get(\"is_auto_mode\", False),\n groups=response_data[\"groups\"],\n personas=response_data.get(\"personas\", []),\n api_base=response_data[\"api_base\"],\n api_version=response_data[\"api_version\"],\n )\n\n if set_as_default:\n if default_model_name is None:\n default_model_name = \"gpt-4o-mini\"\n set_default_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n json={\n \"provider_id\": response_data[\"id\"],\n \"model_name\": default_model_name,\n },\n headers=(\n user_performing_action.headers\n if user_performing_action\n else GENERAL_HEADERS\n ),\n )\n set_default_response.raise_for_status()\n\n return result_llm\n\n @staticmethod\n def delete(\n llm_provider: DATestLLMProvider,\n user_performing_action: DATestUser,\n ) -> bool:\n response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{llm_provider.id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return True\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[LLMProviderView]:\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [LLMProviderView(**p) for p in response.json()[\"providers\"]]\n\n @staticmethod\n def verify(\n llm_provider: DATestLLMProvider,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n all_llm_providers = LLMProviderManager.get_all(user_performing_action)\n default_model = LLMProviderManager.get_default_model(user_performing_action)\n for fetched_llm_provider in all_llm_providers:\n model_names = [\n model.name for model in fetched_llm_provider.model_configurations\n ]\n if llm_provider.id == fetched_llm_provider.id:\n if verify_deleted:\n raise ValueError(\n f\"LLM Provider {llm_provider.id} found but should be deleted\"\n )\n fetched_llm_groups = set(fetched_llm_provider.groups)\n llm_provider_groups = set(llm_provider.groups)\n\n # NOTE: returned api keys are sanitized and should not match\n if (\n fetched_llm_groups == llm_provider_groups\n and llm_provider.provider == fetched_llm_provider.provider\n and (\n default_model is None or default_model.model_name in model_names\n )\n and llm_provider.is_public == fetched_llm_provider.is_public\n and set(fetched_llm_provider.personas) == set(llm_provider.personas)\n ):\n return\n if not verify_deleted:\n raise ValueError(f\"LLM Provider {llm_provider.id} not found\")\n\n @staticmethod\n def get_default_model(\n user_performing_action: DATestUser | None = None,\n ) -> DefaultModel | None:\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=(\n user_performing_action.headers\n if user_performing_action\n else GENERAL_HEADERS\n ),\n )\n response.raise_for_status()\n default_text = response.json().get(\"default_text\")\n if default_text is None:\n return None\n return DefaultModel(**default_text)\n" }, { "path": "backend/tests/integration/common_utils/managers/pat.py", "content": "\"\"\"Helper for managing Personal Access Tokens in integration tests.\"\"\"\n\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestPAT\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass PATManager:\n \"\"\"Manager for creating and managing Personal Access Tokens in tests.\"\"\"\n\n @staticmethod\n def create(\n name: str,\n expiration_days: int | None,\n user_performing_action: DATestUser,\n ) -> DATestPAT:\n \"\"\"Create a Personal Access Token for a user.\n\n Args:\n name: Name of the token\n expiration_days: Number of days until expiration (None for never)\n user_performing_action: User creating the token\n\n Returns:\n DATestPAT with PAT data including the raw token\n \"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/user/pats\",\n json={\"name\": name, \"expiration_days\": expiration_days},\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n timeout=60,\n )\n response.raise_for_status()\n return DATestPAT(**response.json())\n\n @staticmethod\n def list(user_performing_action: DATestUser) -> list[DATestPAT]:\n \"\"\"List all PATs for a user.\n\n Args:\n user_performing_action: User listing their tokens\n\n Returns:\n List of DATestPAT (without raw tokens)\n \"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/user/pats\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n timeout=60,\n )\n response.raise_for_status()\n return [DATestPAT(**pat_data) for pat_data in response.json()]\n\n @staticmethod\n def revoke(token_id: int, user_performing_action: DATestUser) -> None:\n \"\"\"Revoke a Personal Access Token.\n\n Args:\n token_id: ID of the token to revoke\n user_performing_action: User revoking the token\n \"\"\"\n response = requests.delete(\n f\"{API_SERVER_URL}/user/pats/{token_id}\",\n headers=user_performing_action.headers,\n cookies=user_performing_action.cookies,\n timeout=60,\n )\n response.raise_for_status()\n\n @staticmethod\n def authenticate(token: str) -> requests.Response:\n \"\"\"Authenticate using a PAT token and get user info.\n\n Args:\n token: The raw PAT token\n\n Returns:\n Response from /me endpoint\n \"\"\"\n return requests.get(\n f\"{API_SERVER_URL}/me\",\n headers={\"Authorization\": f\"Bearer {token}\"},\n timeout=60,\n )\n\n @staticmethod\n def get_auth_headers(token: str) -> dict[str, str]:\n \"\"\"Get authorization headers for a PAT token.\n\n Args:\n token: The raw PAT token\n\n Returns:\n Headers dict with Authorization bearer token\n \"\"\"\n return {\"Authorization\": f\"Bearer {token}\"}\n" }, { "path": "backend/tests/integration/common_utils/managers/persona.py", "content": "from uuid import UUID\nfrom uuid import uuid4\n\nimport requests\n\nfrom onyx.server.features.persona.models import FullPersonaSnapshot\nfrom onyx.server.features.persona.models import PersonaUpsertRequest\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestPersona\nfrom tests.integration.common_utils.test_models import DATestPersonaLabel\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass PersonaManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n name: str | None = None,\n description: str | None = None,\n system_prompt: str | None = None,\n task_prompt: str | None = None,\n is_public: bool = True,\n datetime_aware: bool = False,\n document_set_ids: list[int] | None = None,\n tool_ids: list[int] | None = None,\n llm_model_provider_override: str | None = None,\n llm_model_version_override: str | None = None,\n users: list[str] | None = None,\n groups: list[int] | None = None,\n label_ids: list[int] | None = None,\n user_file_ids: list[str] | None = None,\n display_priority: int | None = None,\n featured: bool = False,\n ) -> DATestPersona:\n name = name or f\"test-persona-{uuid4()}\"\n description = description or f\"Description for {name}\"\n system_prompt = system_prompt or f\"System prompt for {name}\"\n task_prompt = task_prompt or f\"Task prompt for {name}\"\n\n persona_creation_request = PersonaUpsertRequest(\n name=name,\n description=description,\n system_prompt=system_prompt,\n task_prompt=task_prompt,\n datetime_aware=datetime_aware,\n is_public=is_public,\n document_set_ids=document_set_ids or [],\n tool_ids=tool_ids or [],\n llm_model_provider_override=llm_model_provider_override,\n llm_model_version_override=llm_model_version_override,\n users=[UUID(user) for user in (users or [])],\n groups=groups or [],\n label_ids=label_ids or [],\n user_file_ids=user_file_ids or [],\n display_priority=display_priority,\n is_featured=featured,\n )\n\n response = requests.post(\n f\"{API_SERVER_URL}/persona\",\n json=persona_creation_request.model_dump(mode=\"json\"),\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n persona_data = response.json()\n\n return DATestPersona(\n id=persona_data[\"id\"],\n name=name,\n description=description,\n is_public=is_public,\n system_prompt=system_prompt,\n task_prompt=task_prompt,\n datetime_aware=datetime_aware,\n document_set_ids=document_set_ids or [],\n tool_ids=tool_ids or [],\n llm_model_provider_override=llm_model_provider_override,\n llm_model_version_override=llm_model_version_override,\n users=users or [],\n groups=groups or [],\n label_ids=label_ids or [],\n is_featured=featured,\n )\n\n @staticmethod\n def edit(\n persona: DATestPersona,\n user_performing_action: DATestUser,\n name: str | None = None,\n description: str | None = None,\n system_prompt: str | None = None,\n task_prompt: str | None = None,\n is_public: bool | None = None,\n datetime_aware: bool = False,\n document_set_ids: list[int] | None = None,\n tool_ids: list[int] | None = None,\n llm_model_provider_override: str | None = None,\n llm_model_version_override: str | None = None,\n users: list[str] | None = None,\n groups: list[int] | None = None,\n label_ids: list[int] | None = None,\n featured: bool | None = None,\n ) -> DATestPersona:\n system_prompt = system_prompt or f\"System prompt for {persona.name}\"\n task_prompt = task_prompt or f\"Task prompt for {persona.name}\"\n\n persona_update_request = PersonaUpsertRequest(\n name=name or persona.name,\n description=description or persona.description,\n system_prompt=system_prompt,\n task_prompt=task_prompt,\n datetime_aware=datetime_aware,\n is_public=persona.is_public if is_public is None else is_public,\n document_set_ids=document_set_ids or persona.document_set_ids,\n tool_ids=tool_ids or persona.tool_ids,\n llm_model_provider_override=(\n llm_model_provider_override or persona.llm_model_provider_override\n ),\n llm_model_version_override=(\n llm_model_version_override or persona.llm_model_version_override\n ),\n users=[UUID(user) for user in (users or persona.users)],\n groups=groups or persona.groups,\n label_ids=label_ids or persona.label_ids,\n is_featured=featured if featured is not None else persona.is_featured,\n )\n\n response = requests.patch(\n f\"{API_SERVER_URL}/persona/{persona.id}\",\n json=persona_update_request.model_dump(mode=\"json\"),\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n updated_persona_data = response.json()\n\n return DATestPersona(\n id=updated_persona_data[\"id\"],\n name=updated_persona_data[\"name\"],\n description=updated_persona_data[\"description\"],\n is_public=updated_persona_data[\"is_public\"],\n system_prompt=system_prompt,\n task_prompt=task_prompt,\n datetime_aware=datetime_aware,\n document_set_ids=[ds[\"id\"] for ds in updated_persona_data[\"document_sets\"]],\n tool_ids=[t[\"id\"] for t in updated_persona_data[\"tools\"]],\n llm_model_provider_override=updated_persona_data[\n \"llm_model_provider_override\"\n ],\n llm_model_version_override=updated_persona_data[\n \"llm_model_version_override\"\n ],\n users=[user[\"email\"] for user in updated_persona_data[\"users\"]],\n groups=updated_persona_data[\"groups\"],\n label_ids=[label[\"id\"] for label in updated_persona_data[\"labels\"]],\n is_featured=updated_persona_data[\"is_featured\"],\n )\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[FullPersonaSnapshot]:\n response = requests.get(\n f\"{API_SERVER_URL}/admin/persona\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [FullPersonaSnapshot(**persona) for persona in response.json()]\n\n @staticmethod\n def get_one(\n persona_id: int,\n user_performing_action: DATestUser,\n ) -> list[FullPersonaSnapshot]:\n response = requests.get(\n f\"{API_SERVER_URL}/persona/{persona_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [FullPersonaSnapshot(**response.json())]\n\n @staticmethod\n def verify(\n persona: DATestPersona,\n user_performing_action: DATestUser,\n ) -> bool:\n all_personas = PersonaManager.get_one(\n persona_id=persona.id,\n user_performing_action=user_performing_action,\n )\n for fetched_persona in all_personas:\n if fetched_persona.id == persona.id:\n mismatches: list[tuple[str, object, object]] = []\n\n if fetched_persona.name != persona.name:\n mismatches.append((\"name\", persona.name, fetched_persona.name))\n if fetched_persona.description != persona.description:\n mismatches.append(\n (\n \"description\",\n persona.description,\n fetched_persona.description,\n )\n )\n if fetched_persona.is_public != persona.is_public:\n mismatches.append(\n (\"is_public\", persona.is_public, fetched_persona.is_public)\n )\n if fetched_persona.is_featured != persona.is_featured:\n mismatches.append(\n (\n \"is_featured\",\n persona.is_featured,\n fetched_persona.is_featured,\n )\n )\n if (\n fetched_persona.llm_model_provider_override\n != persona.llm_model_provider_override\n ):\n mismatches.append(\n (\n \"llm_model_provider_override\",\n persona.llm_model_provider_override,\n fetched_persona.llm_model_provider_override,\n )\n )\n if (\n fetched_persona.llm_model_version_override\n != persona.llm_model_version_override\n ):\n mismatches.append(\n (\n \"llm_model_version_override\",\n persona.llm_model_version_override,\n fetched_persona.llm_model_version_override,\n )\n )\n if fetched_persona.system_prompt != persona.system_prompt:\n mismatches.append(\n (\n \"system_prompt\",\n persona.system_prompt,\n fetched_persona.system_prompt,\n )\n )\n if fetched_persona.task_prompt != persona.task_prompt:\n mismatches.append(\n (\n \"task_prompt\",\n persona.task_prompt,\n fetched_persona.task_prompt,\n )\n )\n if fetched_persona.datetime_aware != persona.datetime_aware:\n mismatches.append(\n (\n \"datetime_aware\",\n persona.datetime_aware,\n fetched_persona.datetime_aware,\n )\n )\n\n fetched_document_set_ids = {\n document_set.id for document_set in fetched_persona.document_sets\n }\n expected_document_set_ids = set(persona.document_set_ids)\n if fetched_document_set_ids != expected_document_set_ids:\n mismatches.append(\n (\n \"document_set_ids\",\n sorted(expected_document_set_ids),\n sorted(fetched_document_set_ids),\n )\n )\n\n fetched_tool_ids = {tool.id for tool in fetched_persona.tools}\n expected_tool_ids = set(persona.tool_ids)\n if fetched_tool_ids != expected_tool_ids:\n mismatches.append(\n (\n \"tool_ids\",\n sorted(expected_tool_ids),\n sorted(fetched_tool_ids),\n )\n )\n\n fetched_user_emails = {user.email for user in fetched_persona.users}\n expected_user_emails = set(persona.users)\n if fetched_user_emails != expected_user_emails:\n mismatches.append(\n (\n \"users\",\n sorted(expected_user_emails),\n sorted(fetched_user_emails),\n )\n )\n\n fetched_group_ids = set(fetched_persona.groups)\n expected_group_ids = set(persona.groups)\n if fetched_group_ids != expected_group_ids:\n mismatches.append(\n (\n \"groups\",\n sorted(expected_group_ids),\n sorted(fetched_group_ids),\n )\n )\n\n fetched_label_ids = {label.id for label in fetched_persona.labels}\n expected_label_ids = set(persona.label_ids)\n if fetched_label_ids != expected_label_ids:\n mismatches.append(\n (\n \"label_ids\",\n sorted(expected_label_ids),\n sorted(fetched_label_ids),\n )\n )\n\n if mismatches:\n print(\n f\"Persona verification failed for id={persona.id}. Fields mismatched:\"\n )\n for field_name, expected_value, actual_value in mismatches:\n print(\n f\" - {field_name}: expected {expected_value!r}, got {actual_value!r}\"\n )\n return False\n return True\n print(\n f\"Persona verification failed: persona with id={persona.id} not found in fetched results.\"\n )\n return False\n\n @staticmethod\n def delete(\n persona: DATestPersona,\n user_performing_action: DATestUser,\n ) -> bool:\n response = requests.delete(\n f\"{API_SERVER_URL}/persona/{persona.id}\",\n headers=user_performing_action.headers,\n )\n return response.ok\n\n\nclass PersonaLabelManager:\n @staticmethod\n def create(\n label: DATestPersonaLabel,\n user_performing_action: DATestUser,\n ) -> DATestPersonaLabel:\n response = requests.post(\n f\"{API_SERVER_URL}/persona/labels\",\n json={\n \"name\": label.name,\n },\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n response_data = response.json()\n label.id = response_data[\"id\"]\n return label\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> list[DATestPersonaLabel]:\n response = requests.get(\n f\"{API_SERVER_URL}/persona/labels\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [DATestPersonaLabel(**label) for label in response.json()]\n\n @staticmethod\n def update(\n label: DATestPersonaLabel,\n user_performing_action: DATestUser,\n ) -> DATestPersonaLabel:\n response = requests.patch(\n f\"{API_SERVER_URL}/admin/persona/label/{label.id}\",\n json={\n \"label_name\": label.name,\n },\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return label\n\n @staticmethod\n def delete(\n label: DATestPersonaLabel,\n user_performing_action: DATestUser,\n ) -> bool:\n response = requests.delete(\n f\"{API_SERVER_URL}/admin/persona/label/{label.id}\",\n headers=user_performing_action.headers,\n )\n return response.ok\n\n @staticmethod\n def verify(\n label: DATestPersonaLabel,\n user_performing_action: DATestUser,\n ) -> bool:\n all_labels = PersonaLabelManager.get_all(user_performing_action)\n for fetched_label in all_labels:\n if fetched_label.id == label.id:\n return fetched_label.name == label.name\n return False\n" }, { "path": "backend/tests/integration/common_utils/managers/project.py", "content": "from typing import List\n\nimport requests\n\nfrom onyx.server.features.projects.models import CategorizedFilesSnapshot\nfrom onyx.server.features.projects.models import UserFileSnapshot\nfrom onyx.server.features.projects.models import UserProjectSnapshot\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass ProjectManager:\n @staticmethod\n def create(\n name: str,\n user_performing_action: DATestUser,\n ) -> UserProjectSnapshot:\n \"\"\"Create a new project via API.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/user/projects/create\",\n params={\"name\": name},\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return UserProjectSnapshot.model_validate(response.json())\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n ) -> List[UserProjectSnapshot]:\n \"\"\"Get all projects for a user via API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/user/projects\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [UserProjectSnapshot.model_validate(obj) for obj in response.json()]\n\n @staticmethod\n def delete(\n project_id: int,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"Delete a project via API.\"\"\"\n response = requests.delete(\n f\"{API_SERVER_URL}/user/projects/{project_id}\",\n headers=user_performing_action.headers,\n )\n return response.status_code == 204\n\n @staticmethod\n def verify_deleted(\n project_id: int,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"Verify that a project has been deleted by ensuring it's not in list.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/user/projects\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n projects = [UserProjectSnapshot.model_validate(obj) for obj in response.json()]\n return all(p.id != project_id for p in projects)\n\n @staticmethod\n def verify_files_unlinked(\n project_id: int,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"Verify that all files have been unlinked from the project via API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/user/projects/files/{project_id}\",\n headers=user_performing_action.headers,\n )\n if response.status_code == 404:\n return True\n if not response.ok:\n return False\n files = [UserFileSnapshot.model_validate(obj) for obj in response.json()]\n return len(files) == 0\n\n @staticmethod\n def verify_chat_sessions_unlinked(\n project_id: int,\n user_performing_action: DATestUser,\n ) -> bool:\n \"\"\"Verify that all chat sessions have been unlinked from the project via API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/user/projects/{project_id}\",\n headers=user_performing_action.headers,\n )\n if response.status_code == 404:\n return True\n if not response.ok:\n return False\n try:\n project = UserProjectSnapshot.model_validate(response.json())\n chat_sessions = getattr(project, \"chat_sessions\", [])\n return len(chat_sessions or []) == 0\n except Exception:\n # If response doesn't include chat_sessions, assume unlinked\n return True\n\n @staticmethod\n def upload_files(\n project_id: int,\n files: List[tuple[str, bytes]], # List of (filename, content) tuples\n user_performing_action: DATestUser,\n ) -> CategorizedFilesSnapshot:\n \"\"\"Upload files to a project via API.\"\"\"\n # Build multipart form-data\n files_payload = [\n (\n \"files\",\n (filename, content, \"text/plain\"),\n )\n for filename, content in files\n ]\n\n data = {\"project_id\": str(project_id)} if project_id is not None else {}\n\n # Let requests set Content-Type boundary by not overriding header\n headers = dict(user_performing_action.headers or {})\n headers.pop(\"Content-Type\", None)\n\n response = requests.post(\n f\"{API_SERVER_URL}/user/projects/file/upload\",\n data=data,\n files=files_payload,\n headers=headers,\n )\n response.raise_for_status()\n return CategorizedFilesSnapshot.model_validate(response.json())\n\n @staticmethod\n def get_project_files(\n project_id: int,\n user_performing_action: DATestUser,\n ) -> List[UserFileSnapshot]:\n \"\"\"Get all files associated with a project via API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/user/projects/files/{project_id}\",\n headers=user_performing_action.headers,\n )\n if response.status_code == 404:\n return []\n response.raise_for_status()\n return [UserFileSnapshot.model_validate(obj) for obj in response.json()]\n\n @staticmethod\n def set_instructions(\n project_id: int,\n instructions: str,\n user_performing_action: DATestUser,\n ) -> str:\n \"\"\"Set project instructions via API.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/user/projects/{project_id}/instructions\",\n json={\"instructions\": instructions},\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return (response.json() or {}).get(\"instructions\") or \"\"\n" }, { "path": "backend/tests/integration/common_utils/managers/query_history.py", "content": "import time\nfrom datetime import datetime\nfrom urllib.parse import urlencode\nfrom uuid import UUID\n\nimport requests\nfrom requests.models import CaseInsensitiveDict\n\nfrom ee.onyx.server.query_history.models import ChatSessionMinimal\nfrom ee.onyx.server.query_history.models import ChatSessionSnapshot\nfrom onyx.configs.constants import QAFeedbackType\nfrom onyx.db.enums import TaskStatus\nfrom onyx.server.documents.models import PaginatedReturn\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import MAX_DELAY\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass QueryHistoryManager:\n @staticmethod\n def get_query_history_page(\n user_performing_action: DATestUser,\n page_num: int = 0,\n page_size: int = 10,\n feedback_type: QAFeedbackType | None = None,\n start_time: datetime | None = None,\n end_time: datetime | None = None,\n ) -> PaginatedReturn[ChatSessionMinimal]:\n query_params: dict[str, str | int] = {\n \"page_num\": page_num,\n \"page_size\": page_size,\n }\n if feedback_type:\n query_params[\"feedback_type\"] = feedback_type.value\n if start_time:\n query_params[\"start_time\"] = start_time.isoformat()\n if end_time:\n query_params[\"end_time\"] = end_time.isoformat()\n\n response = requests.get(\n url=f\"{API_SERVER_URL}/admin/chat-session-history?{urlencode(query_params, doseq=True)}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n data = response.json()\n return PaginatedReturn(\n items=[ChatSessionMinimal(**item) for item in data[\"items\"]],\n total_items=data[\"total_items\"],\n )\n\n @staticmethod\n def get_chat_session_admin(\n chat_session_id: UUID | str,\n user_performing_action: DATestUser,\n ) -> ChatSessionSnapshot:\n response = requests.get(\n url=f\"{API_SERVER_URL}/admin/chat-session-history/{chat_session_id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return ChatSessionSnapshot(**response.json())\n\n @staticmethod\n def get_query_history_as_csv(\n user_performing_action: DATestUser,\n start_time: datetime | None = None,\n end_time: datetime | None = None,\n ) -> tuple[CaseInsensitiveDict[str], str]:\n query_params: dict[str, str | int] = {}\n if start_time:\n query_params[\"start\"] = start_time.isoformat()\n if end_time:\n query_params[\"end\"] = end_time.isoformat()\n\n start_response = requests.post(\n url=f\"{API_SERVER_URL}/admin/query-history/start-export?{urlencode(query_params, doseq=True)}\",\n headers=user_performing_action.headers,\n )\n start_response.raise_for_status()\n request_id = start_response.json()[\"request_id\"]\n\n deadline = time.time() + MAX_DELAY\n while time.time() < deadline:\n status_response = requests.get(\n url=f\"{API_SERVER_URL}/admin/query-history/export-status\",\n params={\"request_id\": request_id},\n headers=user_performing_action.headers,\n )\n status_response.raise_for_status()\n status = status_response.json()[\"status\"]\n if status == TaskStatus.SUCCESS:\n break\n if status == TaskStatus.FAILURE:\n raise RuntimeError(\"Query history export task failed\")\n time.sleep(2)\n else:\n raise TimeoutError(\n f\"Query history export not completed within {MAX_DELAY} seconds\"\n )\n\n download_response = requests.get(\n url=f\"{API_SERVER_URL}/admin/query-history/download\",\n params={\"request_id\": request_id},\n headers=user_performing_action.headers,\n )\n download_response.raise_for_status()\n\n if not download_response.content:\n raise RuntimeError(\n \"Query history CSV download returned zero-length content\"\n )\n\n return download_response.headers, download_response.content.decode()\n" }, { "path": "backend/tests/integration/common_utils/managers/scim_client.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\n\n\nclass ScimClient:\n \"\"\"HTTP client for making authenticated SCIM v2 requests.\"\"\"\n\n @staticmethod\n def _headers(raw_token: str) -> dict[str, str]:\n return {\n **GENERAL_HEADERS,\n \"Authorization\": f\"Bearer {raw_token}\",\n }\n\n @staticmethod\n def get(path: str, raw_token: str) -> requests.Response:\n return requests.get(\n f\"{API_SERVER_URL}/scim/v2{path}\",\n headers=ScimClient._headers(raw_token),\n timeout=60,\n )\n\n @staticmethod\n def post(path: str, raw_token: str, json: dict) -> requests.Response:\n return requests.post(\n f\"{API_SERVER_URL}/scim/v2{path}\",\n json=json,\n headers=ScimClient._headers(raw_token),\n timeout=60,\n )\n\n @staticmethod\n def put(path: str, raw_token: str, json: dict) -> requests.Response:\n return requests.put(\n f\"{API_SERVER_URL}/scim/v2{path}\",\n json=json,\n headers=ScimClient._headers(raw_token),\n timeout=60,\n )\n\n @staticmethod\n def patch(path: str, raw_token: str, json: dict) -> requests.Response:\n return requests.patch(\n f\"{API_SERVER_URL}/scim/v2{path}\",\n json=json,\n headers=ScimClient._headers(raw_token),\n timeout=60,\n )\n\n @staticmethod\n def delete(path: str, raw_token: str) -> requests.Response:\n return requests.delete(\n f\"{API_SERVER_URL}/scim/v2{path}\",\n headers=ScimClient._headers(raw_token),\n timeout=60,\n )\n\n @staticmethod\n def get_no_auth(path: str) -> requests.Response:\n return requests.get(\n f\"{API_SERVER_URL}/scim/v2{path}\",\n headers=GENERAL_HEADERS,\n timeout=60,\n )\n" }, { "path": "backend/tests/integration/common_utils/managers/scim_token.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestScimToken\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass ScimTokenManager:\n @staticmethod\n def create(\n name: str,\n user_performing_action: DATestUser,\n ) -> DATestScimToken:\n response = requests.post(\n f\"{API_SERVER_URL}/admin/enterprise-settings/scim/token\",\n json={\"name\": name},\n headers=user_performing_action.headers,\n timeout=60,\n )\n response.raise_for_status()\n data = response.json()\n return DATestScimToken(\n id=data[\"id\"],\n name=data[\"name\"],\n token_display=data[\"token_display\"],\n is_active=data[\"is_active\"],\n created_at=data[\"created_at\"],\n last_used_at=data.get(\"last_used_at\"),\n raw_token=data[\"raw_token\"],\n )\n\n @staticmethod\n def get_active(\n user_performing_action: DATestUser,\n ) -> DATestScimToken | None:\n response = requests.get(\n f\"{API_SERVER_URL}/admin/enterprise-settings/scim/token\",\n headers=user_performing_action.headers,\n timeout=60,\n )\n if response.status_code == 404:\n return None\n response.raise_for_status()\n data = response.json()\n return DATestScimToken(\n id=data[\"id\"],\n name=data[\"name\"],\n token_display=data[\"token_display\"],\n is_active=data[\"is_active\"],\n created_at=data[\"created_at\"],\n last_used_at=data.get(\"last_used_at\"),\n )\n" }, { "path": "backend/tests/integration/common_utils/managers/settings.py", "content": "from typing import Any\nfrom typing import Dict\nfrom typing import Optional\n\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestSettings\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass SettingsManager:\n @staticmethod\n def get_settings(\n user_performing_action: DATestUser,\n ) -> tuple[Dict[str, Any], str]:\n headers = user_performing_action.headers\n headers.pop(\"Content-Type\", None)\n\n response = requests.get(\n f\"{API_SERVER_URL}/admin/settings\",\n headers=headers,\n )\n\n if not response.ok:\n return (\n {},\n f\"Failed to get settings - {response.json().get('detail', 'Unknown error')}\",\n )\n\n return response.json(), \"\"\n\n @staticmethod\n def update_settings(\n settings: DATestSettings,\n user_performing_action: DATestUser,\n ) -> tuple[Dict[str, Any], str]:\n headers = user_performing_action.headers\n headers.pop(\"Content-Type\", None)\n\n payload = settings.model_dump()\n response = requests.put(\n f\"{API_SERVER_URL}/admin/settings\",\n json=payload,\n headers=headers,\n )\n\n if not response.ok:\n return (\n {},\n f\"Failed to update settings - {response.json().get('detail', 'Unknown error')}\",\n )\n\n return response.json(), \"\"\n\n @staticmethod\n def get_setting(\n key: str,\n user_performing_action: DATestUser,\n ) -> Optional[Any]:\n settings, error = SettingsManager.get_settings(user_performing_action)\n if error:\n return None\n return settings.get(key)\n" }, { "path": "backend/tests/integration/common_utils/managers/tenant.py", "content": "from datetime import datetime\nfrom datetime import timedelta\n\nimport jwt\nimport requests\n\nfrom onyx.server.manage.models import AllUsersResponse\nfrom onyx.server.models import FullUserSnapshot\nfrom onyx.server.models import InvitedUserSnapshot\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef generate_auth_token() -> str:\n payload = {\n \"iss\": \"control_plane\",\n \"exp\": datetime.utcnow() + timedelta(minutes=5),\n \"iat\": datetime.utcnow(),\n \"scope\": \"tenant:create\",\n }\n token = jwt.encode(payload, \"\", algorithm=\"HS256\")\n return token\n\n\nclass TenantManager:\n @staticmethod\n def get_all_users(\n user_performing_action: DATestUser,\n ) -> AllUsersResponse:\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/users\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n data = response.json()\n return AllUsersResponse(\n accepted=[FullUserSnapshot(**user) for user in data[\"accepted\"]],\n invited=[InvitedUserSnapshot(**user) for user in data[\"invited\"]],\n slack_users=[FullUserSnapshot(**user) for user in data[\"slack_users\"]],\n accepted_pages=data[\"accepted_pages\"],\n invited_pages=data[\"invited_pages\"],\n slack_users_pages=data[\"slack_users_pages\"],\n )\n\n @staticmethod\n def verify_user_in_tenant(\n user: DATestUser,\n user_performing_action: DATestUser,\n ) -> None:\n all_users = TenantManager.get_all_users(user_performing_action)\n for accepted_user in all_users.accepted:\n if accepted_user.email == user.email and accepted_user.id == user.id:\n return\n raise ValueError(f\"User {user.email} not found in tenant\")\n" }, { "path": "backend/tests/integration/common_utils/managers/tool.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestTool\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass ToolManager:\n @staticmethod\n def list_tools(\n user_performing_action: DATestUser,\n ) -> list[DATestTool]:\n response = requests.get(\n url=f\"{API_SERVER_URL}/tool\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return [\n DATestTool(\n id=tool.get(\"id\"),\n name=tool.get(\"name\"),\n description=tool.get(\"description\"),\n display_name=tool.get(\"display_name\"),\n in_code_tool_id=tool.get(\"in_code_tool_id\"),\n )\n for tool in response.json()\n ]\n" }, { "path": "backend/tests/integration/common_utils/managers/user.py", "content": "from copy import deepcopy\nfrom urllib.parse import urlencode\nfrom uuid import uuid4\n\nimport pytest\nimport requests\nfrom requests import HTTPError\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import ANONYMOUS_USER_EMAIL\nfrom onyx.configs.constants import ANONYMOUS_USER_UUID\nfrom onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME\nfrom onyx.server.documents.models import PaginatedReturn\nfrom onyx.server.manage.models import UserInfo\nfrom onyx.server.models import FullUserSnapshot\nfrom onyx.server.models import InvitedUserSnapshot\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\nfrom tests.integration.common_utils.test_models import DATestUser\n\nDOMAIN = \"example.com\"\nDEFAULT_PASSWORD = \"TestPassword123!\"\n\n\ndef build_email(name: str) -> str:\n return f\"{name}@example.com\"\n\n\nclass UserManager:\n @staticmethod\n def get_anonymous_user() -> DATestUser:\n \"\"\"Get a DATestUser representing the anonymous user.\n\n Anonymous users are real users in the database with LIMITED role.\n They don't have login cookies - requests are made with GENERAL_HEADERS.\n The anonymous_user_enabled setting must be True for these requests to work.\n \"\"\"\n return DATestUser(\n id=ANONYMOUS_USER_UUID,\n email=ANONYMOUS_USER_EMAIL,\n password=\"\",\n headers=GENERAL_HEADERS,\n role=UserRole.LIMITED,\n is_active=True,\n )\n\n @staticmethod\n def create(\n name: str | None = None,\n email: str | None = None,\n ) -> DATestUser:\n if name is None:\n name = f\"test{str(uuid4())}\"\n\n if email is None:\n email = build_email(name)\n\n password = DEFAULT_PASSWORD\n\n body = {\n \"email\": email,\n \"username\": email,\n \"password\": password,\n }\n response = requests.post(\n url=f\"{API_SERVER_URL}/auth/register\",\n json=body,\n headers=GENERAL_HEADERS,\n )\n response.raise_for_status()\n\n test_user = DATestUser(\n id=response.json()[\"id\"],\n email=email,\n password=password,\n headers=deepcopy(GENERAL_HEADERS),\n # fill as basic for now, the `login_as_user` call will\n # fill it in correctly\n role=UserRole.BASIC,\n is_active=True,\n )\n print(f\"Created user {test_user.email}\")\n\n return UserManager.login_as_user(test_user)\n\n @staticmethod\n def login_as_user(test_user: DATestUser) -> DATestUser:\n data = urlencode(\n {\n \"username\": test_user.email,\n \"password\": test_user.password,\n }\n )\n headers = test_user.headers.copy()\n headers[\"Content-Type\"] = \"application/x-www-form-urlencoded\"\n\n response = requests.post(\n url=f\"{API_SERVER_URL}/auth/login\",\n data=data,\n headers=headers,\n )\n\n response.raise_for_status()\n\n cookies = response.cookies.get_dict()\n session_cookie = cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)\n\n if not session_cookie:\n raise Exception(\"Failed to login\")\n\n # Set cookies in the headers\n test_user.headers[\"Cookie\"] = f\"fastapiusersauth={session_cookie}; \"\n test_user.cookies = {\"fastapiusersauth\": session_cookie}\n\n # Get user role from /me endpoint\n me_response = requests.get(\n url=f\"{API_SERVER_URL}/me\",\n headers=test_user.headers,\n cookies=test_user.cookies,\n )\n me_response.raise_for_status()\n me_response_json = me_response.json()\n test_user.id = me_response_json[\"id\"]\n role = UserRole(me_response_json[\"role\"])\n test_user.role = role\n\n return test_user\n\n @staticmethod\n def get_permissions(user: DATestUser) -> list[str]:\n response = requests.get(\n url=f\"{API_SERVER_URL}/me/permissions\",\n headers=user.headers,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def is_role(\n user_to_verify: DATestUser,\n target_role: UserRole,\n ) -> bool:\n response = requests.get(\n url=f\"{API_SERVER_URL}/me\",\n headers=user_to_verify.headers,\n cookies=user_to_verify.cookies,\n )\n\n if user_to_verify.is_active is False:\n with pytest.raises(HTTPError):\n response.raise_for_status()\n return user_to_verify.role == target_role\n else:\n response.raise_for_status()\n\n role_from_response = response.json().get(\"role\", None)\n\n if role_from_response is None:\n return user_to_verify.role == target_role\n\n return target_role == UserRole(role_from_response)\n\n @staticmethod\n def set_role(\n user_to_set: DATestUser,\n target_role: UserRole,\n user_performing_action: DATestUser,\n explicit_override: bool = False,\n ) -> DATestUser:\n response = requests.patch(\n url=f\"{API_SERVER_URL}/manage/set-user-role\",\n json={\n \"user_email\": user_to_set.email,\n \"new_role\": target_role.value,\n \"explicit_override\": explicit_override,\n },\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n new_user_updated_role = DATestUser(\n id=user_to_set.id,\n email=user_to_set.email,\n password=user_to_set.password,\n headers=user_to_set.headers,\n role=target_role,\n is_active=user_to_set.is_active,\n )\n return new_user_updated_role\n\n # TODO: Add a way to check invited status\n @staticmethod\n def is_status(user_to_verify: DATestUser, target_status: bool) -> bool:\n response = requests.get(\n url=f\"{API_SERVER_URL}/me\",\n headers=user_to_verify.headers,\n )\n\n if target_status is False:\n with pytest.raises(HTTPError):\n response.raise_for_status()\n else:\n response.raise_for_status()\n\n is_active = response.json().get(\"is_active\", None)\n if is_active is None:\n return user_to_verify.is_active == target_status\n return target_status == is_active\n\n @staticmethod\n def set_status(\n user_to_set: DATestUser,\n target_status: bool,\n user_performing_action: DATestUser,\n ) -> DATestUser:\n url_substring: str\n if target_status is True:\n url_substring = \"activate\"\n elif target_status is False:\n url_substring = \"deactivate\"\n response = requests.patch(\n url=f\"{API_SERVER_URL}/manage/admin/{url_substring}-user\",\n json={\"user_email\": user_to_set.email},\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n new_user_updated_status = DATestUser(\n id=user_to_set.id,\n email=user_to_set.email,\n password=user_to_set.password,\n headers=user_to_set.headers,\n role=user_to_set.role,\n is_active=target_status,\n )\n return new_user_updated_status\n\n @staticmethod\n def create_test_users(\n user_performing_action: DATestUser,\n user_name_prefix: str,\n count: int,\n role: UserRole = UserRole.BASIC,\n is_active: bool | None = None,\n ) -> list[DATestUser]:\n users_list = []\n for i in range(1, count + 1):\n user = UserManager.create(name=f\"{user_name_prefix}_{i}\")\n if role != UserRole.BASIC:\n user = UserManager.set_role(user, role, user_performing_action)\n if is_active is not None:\n user = UserManager.set_status(user, is_active, user_performing_action)\n users_list.append(user)\n return users_list\n\n @staticmethod\n def get_user_page(\n user_performing_action: DATestUser,\n page_num: int = 0,\n page_size: int = 10,\n search_query: str | None = None,\n role_filter: list[UserRole] | None = None,\n is_active_filter: bool | None = None,\n ) -> PaginatedReturn[FullUserSnapshot]:\n query_params: dict[str, str | list[str] | int] = {\n \"page_num\": page_num,\n \"page_size\": page_size,\n }\n if search_query:\n query_params[\"q\"] = search_query\n if role_filter:\n query_params[\"roles\"] = [role.value for role in role_filter]\n if is_active_filter is not None:\n query_params[\"is_active\"] = is_active_filter\n\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/users/accepted?{urlencode(query_params, doseq=True)}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n data = response.json()\n paginated_result = PaginatedReturn(\n items=[FullUserSnapshot(**user) for user in data[\"items\"]],\n total_items=data[\"total_items\"],\n )\n return paginated_result\n\n @staticmethod\n def invite_user(\n user_to_invite_email: str, user_performing_action: DATestUser\n ) -> None:\n \"\"\"Invite a user by email to join the organization.\n\n Args:\n user_to_invite_email: Email of the user to invite\n user_performing_action: User with admin permissions performing the invitation\n \"\"\"\n response = requests.put(\n url=f\"{API_SERVER_URL}/manage/admin/users\",\n headers=user_performing_action.headers,\n json={\"emails\": [user_to_invite_email]},\n )\n response.raise_for_status()\n\n @staticmethod\n def accept_invitation(tenant_id: str, user_performing_action: DATestUser) -> None:\n \"\"\"Accept an invitation to join the organization.\n\n Args:\n tenant_id: ID of the tenant/organization to accept invitation for\n user_performing_action: User accepting the invitation\n \"\"\"\n response = requests.post(\n url=f\"{API_SERVER_URL}/tenants/users/invite/accept\",\n headers=user_performing_action.headers,\n json={\"tenant_id\": tenant_id},\n )\n response.raise_for_status()\n\n @staticmethod\n def get_invited_users(\n user_performing_action: DATestUser,\n ) -> list[InvitedUserSnapshot]:\n \"\"\"Get a list of all invited users.\n\n Args:\n user_performing_action: User with admin permissions performing the action\n\n Returns:\n List of invited user snapshots\n \"\"\"\n response = requests.get(\n url=f\"{API_SERVER_URL}/manage/users/invited\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n return [InvitedUserSnapshot(**user) for user in response.json()]\n\n @staticmethod\n def get_user_info(user_performing_action: DATestUser) -> UserInfo:\n \"\"\"Get user info for the current user.\n\n Args:\n user_performing_action: User performing the action\n \"\"\"\n response = requests.get(\n url=f\"{API_SERVER_URL}/me\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return UserInfo(**response.json())\n" }, { "path": "backend/tests/integration/common_utils/managers/user_group.py", "content": "import time\nfrom uuid import uuid4\n\nimport requests\n\nfrom ee.onyx.server.user_group.models import UserGroup\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import MAX_DELAY\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import DATestUserGroup\n\n\nclass UserGroupManager:\n @staticmethod\n def create(\n user_performing_action: DATestUser,\n name: str | None = None,\n user_ids: list[str] | None = None,\n cc_pair_ids: list[int] | None = None,\n ) -> DATestUserGroup:\n name = f\"{name}-user-group\" if name else f\"test-user-group-{uuid4()}\"\n\n request = {\n \"name\": name,\n \"user_ids\": user_ids or [],\n \"cc_pair_ids\": cc_pair_ids or [],\n }\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/user-group\",\n json=request,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n test_user_group = DATestUserGroup(\n id=response.json()[\"id\"],\n name=response.json()[\"name\"],\n user_ids=[user[\"id\"] for user in response.json()[\"users\"]],\n cc_pair_ids=[cc_pair[\"id\"] for cc_pair in response.json()[\"cc_pairs\"]],\n )\n return test_user_group\n\n @staticmethod\n def edit(\n user_group: DATestUserGroup,\n user_performing_action: DATestUser,\n ) -> None:\n response = requests.patch(\n f\"{API_SERVER_URL}/manage/admin/user-group/{user_group.id}\",\n json=user_group.model_dump(),\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def delete(\n user_group: DATestUserGroup,\n user_performing_action: DATestUser,\n ) -> None:\n response = requests.delete(\n f\"{API_SERVER_URL}/manage/admin/user-group/{user_group.id}\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def add_users(\n user_group: DATestUserGroup,\n user_ids: list[str],\n user_performing_action: DATestUser,\n ) -> DATestUserGroup:\n request = {\n \"user_ids\": user_ids,\n }\n\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/user-group/{user_group.id}/add-users\",\n json=request,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n user_group.user_ids = [user[\"id\"] for user in response.json()[\"users\"]]\n user_group.cc_pair_ids = [\n cc_pair[\"id\"] for cc_pair in response.json()[\"cc_pairs\"]\n ]\n user_group.name = response.json()[\"name\"]\n return user_group\n\n @staticmethod\n def set_curator_status(\n test_user_group: DATestUserGroup,\n user_to_set_as_curator: DATestUser,\n user_performing_action: DATestUser,\n is_curator: bool = True,\n ) -> None:\n set_curator_request = {\n \"user_id\": user_to_set_as_curator.id,\n \"is_curator\": is_curator,\n }\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/user-group/{test_user_group.id}/set-curator\",\n json=set_curator_request,\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n\n @staticmethod\n def get_permissions(\n user_group: DATestUserGroup,\n user_performing_action: DATestUser,\n ) -> list[str]:\n response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/user-group/{user_group.id}/permissions\",\n headers=user_performing_action.headers,\n )\n response.raise_for_status()\n return response.json()\n\n @staticmethod\n def get_all(\n user_performing_action: DATestUser,\n include_default: bool = False,\n ) -> list[UserGroup]:\n params: dict[str, str] = {}\n if include_default:\n params[\"include_default\"] = \"true\"\n response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/user-group\",\n headers=user_performing_action.headers,\n params=params,\n )\n response.raise_for_status()\n return [UserGroup(**ug) for ug in response.json()]\n\n @staticmethod\n def verify(\n user_group: DATestUserGroup,\n user_performing_action: DATestUser,\n verify_deleted: bool = False,\n ) -> None:\n all_user_groups = UserGroupManager.get_all(user_performing_action)\n for fetched_user_group in all_user_groups:\n if user_group.id == fetched_user_group.id:\n if verify_deleted:\n raise ValueError(\n f\"User group {user_group.id} found but should be deleted\"\n )\n fetched_cc_ids = {cc_pair.id for cc_pair in fetched_user_group.cc_pairs}\n fetched_user_ids = {user.id for user in fetched_user_group.users}\n user_group_cc_ids = set(user_group.cc_pair_ids)\n user_group_user_ids = set(user_group.user_ids)\n if (\n fetched_cc_ids == user_group_cc_ids\n and fetched_user_ids == user_group_user_ids\n ):\n return\n if not verify_deleted:\n raise ValueError(f\"User group {user_group.id} not found\")\n\n @staticmethod\n def wait_for_sync(\n user_performing_action: DATestUser,\n user_groups_to_check: list[DATestUserGroup] | None = None,\n ) -> None:\n start = time.time()\n while True:\n user_groups = UserGroupManager.get_all(user_performing_action)\n if user_groups_to_check:\n check_ids = {user_group.id for user_group in user_groups_to_check}\n user_group_ids = {user_group.id for user_group in user_groups}\n if not check_ids.issubset(user_group_ids):\n raise RuntimeError(\"User group not found\")\n user_groups = [\n user_group\n for user_group in user_groups\n if user_group.id in check_ids\n ]\n if all(ug.is_up_to_date for ug in user_groups):\n print(\"User groups synced successfully.\")\n return\n\n if time.time() - start > MAX_DELAY:\n raise TimeoutError(\n f\"User groups were not synced within the {MAX_DELAY} seconds\"\n )\n else:\n print(\"User groups were not synced yet, waiting...\")\n time.sleep(2)\n\n @staticmethod\n def wait_for_deletion_completion(\n user_groups_to_check: list[DATestUserGroup],\n user_performing_action: DATestUser,\n ) -> None:\n start = time.time()\n user_group_ids_to_check = {user_group.id for user_group in user_groups_to_check}\n while True:\n fetched_user_groups = UserGroupManager.get_all(user_performing_action)\n fetched_user_group_ids = {\n user_group.id for user_group in fetched_user_groups\n }\n if not user_group_ids_to_check.intersection(fetched_user_group_ids):\n return\n\n if time.time() - start > MAX_DELAY:\n raise TimeoutError(\n f\"User groups deletion was not completed within the {MAX_DELAY} seconds\"\n )\n else:\n print(\"Some user groups are still being deleted, waiting...\")\n time.sleep(2)\n" }, { "path": "backend/tests/integration/common_utils/reset.py", "content": "import logging\nimport os\nimport time\nfrom types import SimpleNamespace\n\nimport psycopg2\nimport requests\n\nfrom alembic import command\nfrom alembic.config import Config\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PASSWORD\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.engine.sql_engine import SYNC_DB_API\nfrom onyx.db.engine.tenant_utils import get_all_tenant_ids\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.db.swap_index import check_and_perform_index_swap\nfrom onyx.document_index.document_index_utils import get_multipass_config\nfrom onyx.document_index.vespa.index import DOCUMENT_ID_ENDPOINT\nfrom onyx.document_index.vespa.index import VespaIndex\nfrom onyx.file_store.file_store import get_default_file_store\nfrom onyx.indexing.models import IndexingSetting\nfrom onyx.setup import setup_document_indices\nfrom onyx.setup import setup_postgres\nfrom onyx.utils.logger import setup_logger\nfrom tests.integration.common_utils.timeout import run_with_timeout_multiproc\n\nlogger = setup_logger()\n\n\ndef _run_migrations(\n database_url: str,\n config_name: str,\n direction: str = \"upgrade\",\n revision: str = \"head\",\n schema: str = \"public\",\n) -> None:\n # hide info logs emitted during migration\n logging.getLogger(\"alembic\").setLevel(logging.CRITICAL)\n\n # Create an Alembic configuration object\n alembic_cfg = Config(\"alembic.ini\")\n alembic_cfg.set_section_option(\"logger_alembic\", \"level\", \"WARN\")\n alembic_cfg.attributes[\"configure_logger\"] = False\n alembic_cfg.config_ini_section = config_name\n\n alembic_cfg.cmd_opts = SimpleNamespace() # type: ignore\n alembic_cfg.cmd_opts.x = [f\"schema={schema}\"] # type: ignore\n\n # Set the SQLAlchemy URL in the Alembic configuration\n alembic_cfg.set_main_option(\"sqlalchemy.url\", database_url)\n\n # Run the migration\n if direction == \"upgrade\":\n command.upgrade(alembic_cfg, revision)\n elif direction == \"downgrade\":\n command.downgrade(alembic_cfg, revision)\n else:\n raise ValueError(\n f\"Invalid direction: {direction}. Must be 'upgrade' or 'downgrade'.\"\n )\n\n logging.getLogger(\"alembic\").setLevel(logging.INFO)\n\n\ndef downgrade_postgres(\n database: str = \"postgres\",\n schema: str = \"public\",\n config_name: str = \"alembic\",\n revision: str = \"base\",\n clear_data: bool = False,\n) -> None:\n \"\"\"Downgrade Postgres database to base state.\"\"\"\n if clear_data:\n if revision != \"base\":\n raise ValueError(\"Clearing data without rolling back to base state\")\n\n conn = psycopg2.connect(\n dbname=database,\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n application_name=\"downgrade_postgres\",\n )\n conn.autocommit = True # Need autocommit for dropping schema\n cur = conn.cursor()\n\n # Close any existing connections to the schema before dropping\n cur.execute(\n f\"\"\"\n SELECT pg_terminate_backend(pg_stat_activity.pid)\n FROM pg_stat_activity\n WHERE pg_stat_activity.datname = '{database}'\n AND pg_stat_activity.state = 'idle in transaction'\n AND pid <> pg_backend_pid();\n \"\"\"\n )\n\n # Drop and recreate the public schema - this removes ALL objects\n cur.execute(f\"DROP SCHEMA {schema} CASCADE;\")\n cur.execute(f\"CREATE SCHEMA {schema};\")\n\n # Restore default privileges\n cur.execute(f\"GRANT ALL ON SCHEMA {schema} TO postgres;\")\n cur.execute(f\"GRANT ALL ON SCHEMA {schema} TO public;\")\n\n cur.close()\n conn.close()\n\n return\n\n # Downgrade to base\n conn_str = build_connection_string(\n db=database,\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n db_api=SYNC_DB_API,\n )\n _run_migrations(\n conn_str,\n config_name,\n direction=\"downgrade\",\n revision=revision,\n )\n\n\ndef upgrade_postgres(\n database: str = \"postgres\", config_name: str = \"alembic\", revision: str = \"head\"\n) -> None:\n \"\"\"Upgrade Postgres database to latest version.\"\"\"\n conn_str = build_connection_string(\n db=database,\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n db_api=SYNC_DB_API,\n app_name=\"upgrade_postgres\",\n )\n _run_migrations(\n conn_str,\n config_name,\n direction=\"upgrade\",\n revision=revision,\n )\n\n\ndef drop_multitenant_postgres(\n database: str = \"postgres\",\n) -> None:\n \"\"\"Reset the Postgres database.\"\"\"\n # this seems to hang due to locking issues, so run with a timeout with a few retries\n NUM_TRIES = 10\n TIMEOUT = 40\n success = False\n for _ in range(NUM_TRIES):\n logger.info(f\"drop_multitenant_postgres_task starting... ({_ + 1}/{NUM_TRIES})\")\n try:\n run_with_timeout_multiproc(\n drop_multitenant_postgres_task,\n TIMEOUT,\n kwargs={\n \"dbname\": database,\n },\n )\n success = True\n break\n except TimeoutError:\n logger.warning(\n f\"drop_multitenant_postgres_task timed out, retrying... ({_ + 1}/{NUM_TRIES})\"\n )\n except RuntimeError:\n logger.warning(\n f\"drop_multitenant_postgres_task exceptioned, retrying... ({_ + 1}/{NUM_TRIES})\"\n )\n\n if not success:\n raise RuntimeError(\"drop_multitenant_postgres_task failed after 10 timeouts.\")\n\n\ndef drop_multitenant_postgres_task(dbname: str) -> None:\n conn = psycopg2.connect(\n dbname=dbname,\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n connect_timeout=10,\n application_name=\"drop_multitenant_postgres_task\",\n )\n\n conn.autocommit = True\n cur = conn.cursor()\n\n logger.info(\"Selecting tenant schemas.\")\n # Get all tenant schemas\n cur.execute(\n \"\"\"\n SELECT schema_name\n FROM information_schema.schemata\n WHERE schema_name LIKE 'tenant_%'\n \"\"\"\n )\n tenant_schemas = cur.fetchall()\n\n # Drop all tenant schemas\n logger.info(\"Dropping all tenant schemas.\")\n for schema in tenant_schemas:\n # Close any existing connections to the schema before dropping\n cur.execute(\n \"\"\"\n SELECT pg_terminate_backend(pg_stat_activity.pid)\n FROM pg_stat_activity\n WHERE pg_stat_activity.datname = 'postgres'\n AND pg_stat_activity.state = 'idle in transaction'\n AND pid <> pg_backend_pid();\n \"\"\"\n )\n\n schema_name = schema[0]\n cur.execute(f'DROP SCHEMA \"{schema_name}\" CASCADE')\n\n # Drop tables in the public schema\n logger.info(\"Selecting public schema tables.\")\n cur.execute(\n \"\"\"\n SELECT tablename FROM pg_tables\n WHERE schemaname = 'public'\n \"\"\"\n )\n public_tables = cur.fetchall()\n\n logger.info(\"Dropping public schema tables.\")\n for table in public_tables:\n table_name = table[0]\n cur.execute(f'DROP TABLE IF EXISTS public.\"{table_name}\" CASCADE')\n\n cur.close()\n conn.close()\n\n\ndef reset_postgres(\n database: str = \"postgres\",\n config_name: str = \"alembic\",\n setup_onyx: bool = True,\n) -> None:\n \"\"\"Reset the Postgres database.\"\"\"\n # this seems to hang due to locking issues, so run with a timeout with a few retries\n NUM_TRIES = 10\n TIMEOUT = 40\n success = False\n for _ in range(NUM_TRIES):\n logger.info(f\"Downgrading Postgres... ({_ + 1}/{NUM_TRIES})\")\n try:\n run_with_timeout_multiproc(\n downgrade_postgres,\n TIMEOUT,\n kwargs={\n \"database\": database,\n \"config_name\": config_name,\n \"revision\": \"base\",\n \"clear_data\": True,\n },\n )\n success = True\n break\n except TimeoutError:\n logger.warning(\n f\"Postgres downgrade timed out, retrying... ({_ + 1}/{NUM_TRIES})\"\n )\n except RuntimeError:\n logger.warning(\n f\"Postgres downgrade exceptioned, retrying... ({_ + 1}/{NUM_TRIES})\"\n )\n\n if not success:\n raise RuntimeError(\"Postgres downgrade failed after 10 timeouts.\")\n\n logger.info(\"Upgrading Postgres...\")\n upgrade_postgres(database=database, config_name=config_name, revision=\"head\")\n if setup_onyx:\n logger.info(\"Setting up Postgres...\")\n with get_session_with_current_tenant() as db_session:\n setup_postgres(db_session)\n\n\ndef reset_vespa() -> None:\n \"\"\"Wipe all data from the Vespa index.\"\"\"\n\n with get_session_with_current_tenant() as db_session:\n # swap to the correct default model\n check_and_perform_index_swap(db_session)\n\n search_settings = get_current_search_settings(db_session)\n multipass_config = get_multipass_config(search_settings)\n index_name = search_settings.index_name\n\n success = setup_document_indices(\n document_indices=[\n VespaIndex(\n index_name=index_name,\n secondary_index_name=None,\n large_chunks_enabled=multipass_config.enable_large_chunks,\n secondary_large_chunks_enabled=None,\n )\n ],\n index_setting=IndexingSetting.from_db_model(search_settings),\n secondary_index_setting=None,\n )\n if not success:\n raise RuntimeError(\"Could not connect to Vespa within the specified timeout.\")\n\n for _ in range(5):\n try:\n continuation = None\n should_continue = True\n while should_continue:\n params = {\"selection\": \"true\", \"cluster\": \"danswer_index\"}\n if continuation:\n params = {**params, \"continuation\": continuation}\n response = requests.delete(\n DOCUMENT_ID_ENDPOINT.format(index_name=index_name), params=params\n )\n response.raise_for_status()\n\n response_json = response.json()\n\n continuation = response_json.get(\"continuation\")\n should_continue = bool(continuation)\n\n break\n except Exception as e:\n print(f\"Error deleting documents: {e}\")\n time.sleep(5)\n\n\ndef reset_postgres_multitenant() -> None:\n \"\"\"Reset the Postgres database for all tenants in a multitenant setup.\"\"\"\n\n drop_multitenant_postgres()\n reset_postgres(config_name=\"schema_private\", setup_onyx=False)\n\n\ndef reset_vespa_multitenant() -> None:\n \"\"\"Wipe all data from the Vespa index for all tenants.\"\"\"\n\n for tenant_id in get_all_tenant_ids():\n with get_session_with_tenant(tenant_id=tenant_id) as db_session:\n # swap to the correct default model for each tenant\n check_and_perform_index_swap(db_session)\n\n search_settings = get_current_search_settings(db_session)\n multipass_config = get_multipass_config(search_settings)\n index_name = search_settings.index_name\n\n success = setup_document_indices(\n document_indices=[\n VespaIndex(\n index_name=index_name,\n secondary_index_name=None,\n large_chunks_enabled=multipass_config.enable_large_chunks,\n secondary_large_chunks_enabled=None,\n )\n ],\n index_setting=IndexingSetting.from_db_model(search_settings),\n secondary_index_setting=None,\n )\n\n if not success:\n raise RuntimeError(\n f\"Could not connect to Vespa for tenant {tenant_id} within the specified timeout.\"\n )\n\n for _ in range(5):\n try:\n continuation = None\n should_continue = True\n while should_continue:\n params = {\"selection\": \"true\", \"cluster\": \"danswer_index\"}\n if continuation:\n params = {**params, \"continuation\": continuation}\n response = requests.delete(\n DOCUMENT_ID_ENDPOINT.format(index_name=index_name),\n params=params,\n )\n response.raise_for_status()\n\n response_json = response.json()\n\n continuation = response_json.get(\"continuation\")\n should_continue = bool(continuation)\n\n break\n except Exception as e:\n print(f\"Error deleting documents for tenant {tenant_id}: {e}\")\n time.sleep(5)\n\n\ndef reset_file_store() -> None:\n \"\"\"Reset the FileStore.\"\"\"\n filestore = get_default_file_store()\n for file_record in filestore.list_files_by_prefix(\"\"):\n filestore.delete_file(file_record.file_id)\n\n\ndef reset_all() -> None:\n if os.environ.get(\"SKIP_RESET\", \"\").lower() == \"true\":\n logger.info(\"Skipping reset.\")\n return\n\n logger.info(\"Resetting Postgres...\")\n reset_postgres()\n logger.info(\"Resetting Vespa...\")\n reset_vespa()\n logger.info(\"Resetting FileStore...\")\n reset_file_store()\n\n\ndef reset_all_multitenant() -> None:\n \"\"\"Reset both Postgres and Vespa for all tenants.\n\n Honors SKIP_RESET env var to allow callers (e.g., CI) to disable\n heavy resets entirely for faster end-to-end runs.\n \"\"\"\n if os.environ.get(\"SKIP_RESET\", \"\").lower() == \"true\":\n logger.info(\"SKIPPING multitenant reset due to SKIP_RESET=true\")\n return\n\n logger.info(\"Resetting Postgres for all tenants...\")\n reset_postgres_multitenant()\n logger.info(\"Resetting Vespa for all tenants...\")\n reset_vespa_multitenant()\n logger.info(\"Finished resetting all.\")\n" }, { "path": "backend/tests/integration/common_utils/test_document_utils.py", "content": "import uuid\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.connectors.models import TextSection\n\n\ndef create_test_document(\n doc_id: str | None = None,\n text: str = \"Test content\",\n link: str = \"http://example.com\",\n source: DocumentSource = DocumentSource.MOCK_CONNECTOR,\n metadata: dict | None = None,\n) -> Document:\n \"\"\"Create a test document with the given parameters.\n\n Args:\n doc_id: Optional document ID. If not provided, a random UUID will be generated.\n text: The text content of the document. Defaults to \"Test content\".\n link: The link for the document section. Defaults to \"http://example.com\".\n source: The document source. Defaults to MOCK_CONNECTOR.\n metadata: Optional metadata dictionary. Defaults to empty dict.\n \"\"\"\n doc_id = doc_id or f\"test-doc-{uuid.uuid4()}\"\n return Document(\n id=doc_id,\n sections=[TextSection(text=text, link=link)],\n source=source,\n semantic_identifier=doc_id,\n doc_updated_at=datetime.now(timezone.utc),\n metadata=metadata or {},\n )\n\n\ndef create_test_document_failure(\n doc_id: str,\n failure_message: str = \"Simulated failure\",\n document_link: str | None = None,\n) -> ConnectorFailure:\n \"\"\"Create a test document failure with the given parameters.\n\n Args:\n doc_id: The ID of the document that failed.\n failure_message: The failure message. Defaults to \"Simulated failure\".\n document_link: Optional link to the failed document.\n \"\"\"\n return ConnectorFailure(\n failed_document=DocumentFailure(\n document_id=doc_id,\n document_link=document_link,\n ),\n failure_message=failure_message,\n )\n" }, { "path": "backend/tests/integration/common_utils/test_file_utils.py", "content": "import io\n\nfrom PIL import Image\n\n\ndef create_test_image(\n width: int = 1,\n height: int = 1,\n color: str = \"white\",\n format: str = \"PNG\",\n) -> io.BytesIO:\n \"\"\"Create a test image file in memory for file attachment testing.\n\n Args:\n width: Width of the image in pixels. Defaults to 1.\n height: Height of the image in pixels. Defaults to 1.\n color: Color of the image. Defaults to \"white\".\n format: Image format (PNG, JPEG, etc.). Defaults to \"PNG\".\n\n Returns:\n A BytesIO object containing the image data, positioned at the start.\n \"\"\"\n image = Image.new(\"RGB\", (width, height), color=color)\n image_file = io.BytesIO()\n image.save(image_file, format=format)\n image_file.seek(0)\n return image_file\n\n\ndef create_test_text_file(content: str | bytes) -> io.BytesIO:\n \"\"\"Create a test text file in memory for file attachment testing.\n\n Args:\n content: The text content of the file. Can be string or bytes.\n\n Returns:\n A BytesIO object containing the text data, positioned at the start.\n \"\"\"\n if isinstance(content, str):\n content = content.encode(\"utf-8\")\n text_file = io.BytesIO(content)\n text_file.seek(0)\n return text_file\n" }, { "path": "backend/tests/integration/common_utils/test_models.py", "content": "from dataclasses import dataclass\nfrom datetime import datetime\nfrom enum import Enum\nfrom typing import Any\nfrom uuid import UUID\n\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import MessageType\nfrom onyx.configs.constants import QAFeedbackType\nfrom onyx.context.search.models import SavedSearchDoc\nfrom onyx.context.search.models import SearchDoc\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom onyx.server.documents.models import IndexAttemptSnapshot\nfrom onyx.server.documents.models import IndexingStatus\nfrom onyx.server.documents.models import InputType\nfrom onyx.server.query_and_chat.streaming_models import GeneratedImage\n\n\"\"\"\nThese data models are used to represent the data on the testing side of things.\nThis means the flow is:\n1. Make request that changes data in db\n2. Make a change to the testing model\n3. Retrieve data from db\n4. Compare db data with testing model to verify\n\"\"\"\n\n\nclass DATestPAT(BaseModel):\n \"\"\"Personal Access Token model for testing.\"\"\"\n\n id: int\n name: str\n token: str | None = None # Raw token - only present on initial creation\n token_display: str\n created_at: str\n expires_at: str | None = None\n last_used_at: str | None = None\n\n\nclass DATestScimToken(BaseModel):\n \"\"\"SCIM bearer token model for testing.\"\"\"\n\n id: int\n name: str\n raw_token: str | None = None # Only present on initial creation\n token_display: str\n is_active: bool\n created_at: str\n last_used_at: str | None = None\n\n\nclass DATestAPIKey(BaseModel):\n api_key_id: int\n api_key_display: str\n api_key: str | None = None # only present on initial creation\n api_key_name: str | None = None\n api_key_role: UserRole\n\n user_id: UUID\n headers: dict\n\n\nclass DATestUser(BaseModel):\n id: str\n email: str\n password: str\n headers: dict\n role: UserRole\n is_active: bool\n cookies: dict = {}\n\n\nclass DATestPersonaLabel(BaseModel):\n id: int | None = None\n name: str\n\n\nclass DATestCredential(BaseModel):\n id: int\n name: str\n credential_json: dict[str, Any]\n admin_public: bool\n source: DocumentSource\n curator_public: bool\n groups: list[int]\n\n\nclass DATestConnector(BaseModel):\n id: int\n name: str\n source: DocumentSource\n input_type: InputType\n connector_specific_config: dict[str, Any]\n groups: list[int] | None = None\n access_type: AccessType | None = None\n\n\nclass SimpleTestDocument(BaseModel):\n id: str\n content: str\n image_file_id: str | None = None\n\n\nclass DATestCCPair(BaseModel):\n id: int\n name: str\n connector_id: int\n credential_id: int\n access_type: AccessType\n groups: list[int]\n documents: list[SimpleTestDocument] = Field(default_factory=list)\n\n\nclass DATestUserGroup(BaseModel):\n id: int\n name: str\n user_ids: list[str]\n cc_pair_ids: list[int]\n\n\nclass DATestLLMProvider(BaseModel):\n id: int\n name: str\n provider: str\n api_key: str\n default_model_name: str | None = None\n is_public: bool\n is_auto_mode: bool = False\n groups: list[int]\n personas: list[int]\n api_base: str | None = None\n api_version: str | None = None\n\n\nclass DATestImageGenerationConfig(BaseModel):\n image_provider_id: str\n model_configuration_id: int\n model_name: str\n llm_provider_id: int\n llm_provider_name: str\n is_default: bool\n\n\nclass DATestDocumentSet(BaseModel):\n id: int\n name: str\n description: str\n cc_pair_ids: list[int] = Field(default_factory=list)\n is_public: bool\n is_up_to_date: bool\n users: list[str] = Field(default_factory=list)\n groups: list[int] = Field(default_factory=list)\n federated_connectors: list[dict[str, Any]] = Field(default_factory=list)\n\n\nclass DATestPersona(BaseModel):\n id: int\n name: str\n description: str\n is_public: bool\n document_set_ids: list[int]\n tool_ids: list[int]\n llm_model_provider_override: str | None\n llm_model_version_override: str | None\n users: list[str]\n groups: list[int]\n label_ids: list[int]\n is_featured: bool = False\n\n # Embedded prompt fields (no longer separate prompt_ids)\n system_prompt: str | None = None\n task_prompt: str | None = None\n datetime_aware: bool = True\n\n\nclass DATestChatMessage(BaseModel):\n id: int\n chat_session_id: UUID\n parent_message_id: int | None\n message: str\n message_type: MessageType | None = None\n files: list | None = None\n\n\nclass DATestChatSession(BaseModel):\n id: UUID\n persona_id: int\n description: str\n\n\nclass DAQueryHistoryEntry(DATestChatSession):\n feedback_type: QAFeedbackType | None\n\n\nclass ToolName(str, Enum):\n INTERNET_SEARCH = \"internet_search\"\n INTERNAL_SEARCH = \"run_search\"\n IMAGE_GENERATION = \"generate_image\"\n\n\nclass ToolResult(BaseModel):\n tool_name: ToolName\n\n queries: list[str] = Field(default_factory=list)\n documents: list[SavedSearchDoc] = Field(default_factory=list)\n images: list[GeneratedImage] = Field(default_factory=list)\n\n\nclass ToolCallDebug(BaseModel):\n tool_call_id: str\n tool_name: str\n tool_args: dict[str, Any]\n\n\nclass ErrorResponse(BaseModel):\n error: str\n stack_trace: str\n\n\nclass StreamedResponse(BaseModel):\n full_message: str\n assistant_message_id: int\n top_documents: list[SearchDoc]\n used_tools: list[ToolResult]\n tool_call_debug: list[ToolCallDebug] = Field(default_factory=list)\n error: ErrorResponse | None = None\n\n # Track heartbeat packets for image generation and other tools\n heartbeat_packets: list[dict[str, Any]]\n\n\nclass DATestGatingType(str, Enum):\n FULL = \"full\"\n PARTIAL = \"partial\"\n NONE = \"none\"\n\n\nclass DATestSettings(BaseModel):\n \"\"\"General settings\"\"\"\n\n # is float to allow for fractional days for easier automated testing\n maximum_chat_retention_days: float | None = None\n gpu_enabled: bool | None = None\n product_gating: DATestGatingType = DATestGatingType.NONE\n anonymous_user_enabled: bool | None = None\n image_extraction_and_analysis_enabled: bool | None = False\n search_time_image_analysis_enabled: bool | None = False\n\n\n@dataclass\nclass DATestIndexAttempt:\n id: int\n status: IndexingStatus | None\n new_docs_indexed: int | None\n total_docs_indexed: int | None\n docs_removed_from_index: int | None\n error_msg: str | None\n time_started: datetime | None\n time_updated: datetime | None\n\n @classmethod\n def from_index_attempt_snapshot(\n cls, index_attempt: IndexAttemptSnapshot\n ) -> \"DATestIndexAttempt\":\n return cls(\n id=index_attempt.id,\n status=index_attempt.status,\n new_docs_indexed=index_attempt.new_docs_indexed,\n total_docs_indexed=index_attempt.total_docs_indexed,\n docs_removed_from_index=index_attempt.docs_removed_from_index,\n error_msg=index_attempt.error_msg,\n time_started=(\n datetime.fromisoformat(index_attempt.time_started)\n if index_attempt.time_started\n else None\n ),\n time_updated=datetime.fromisoformat(index_attempt.time_updated),\n )\n\n\nclass DATestTool(BaseModel):\n id: int\n name: str\n description: str\n display_name: str\n in_code_tool_id: str | None\n\n\n# Discord Bot Models\nclass DATestDiscordGuildConfig(BaseModel):\n \"\"\"Discord guild config model for testing.\"\"\"\n\n id: int\n registration_key: str | None = None # Only present on creation\n guild_id: int | None = None\n guild_name: str | None = None\n enabled: bool = True\n default_persona_id: int | None = None\n\n\nclass DATestDiscordChannelConfig(BaseModel):\n \"\"\"Discord channel config model for testing.\"\"\"\n\n id: int\n guild_config_id: int\n channel_id: int\n channel_name: str\n channel_type: str\n is_private: bool\n enabled: bool = False\n thread_only_mode: bool = False\n require_bot_invocation: bool = True\n persona_override_id: int | None = None\n" }, { "path": "backend/tests/integration/common_utils/timeout.py", "content": "# import multiprocessing\n# from collections.abc import Callable\n# from typing import Any\n# from typing import TypeVar\n\n# T = TypeVar(\"T\")\n\n\n# def run_with_timeout_multiproc(\n# task: Callable[..., T], timeout: int, kwargs: dict[str, Any]\n# ) -> T:\n# # Use multiprocessing to prevent a thread from blocking the main thread\n# with multiprocessing.Pool(processes=1) as pool:\n# async_result = pool.apply_async(task, kwds=kwargs)\n# try:\n# # Wait at most timeout seconds for the function to complete\n# result = async_result.get(timeout=timeout)\n# return result\n# except multiprocessing.TimeoutError:\n# raise TimeoutError(f\"Function timed out after {timeout} seconds\")\n\n\nimport multiprocessing\nimport traceback\nfrom collections.abc import Callable\nfrom multiprocessing import Queue\nfrom typing import Any\nfrom typing import TypeVar\n\nT = TypeVar(\"T\")\n\n\ndef _multiproc_wrapper(\n task: Callable[..., T], kwargs: dict[str, Any], q: Queue\n) -> None:\n try:\n result = task(**kwargs)\n q.put((\"success\", result))\n except Exception:\n q.put((\"error\", traceback.format_exc()))\n\n\ndef run_with_timeout_multiproc(\n task: Callable[..., T], timeout: int, kwargs: dict[str, Any]\n) -> T:\n ctx = multiprocessing.get_context(\"spawn\")\n q: Queue = ctx.Queue()\n p = ctx.Process(\n target=_multiproc_wrapper,\n args=(\n task,\n kwargs,\n q,\n ),\n )\n p.start()\n p.join(timeout)\n\n if p.is_alive():\n p.terminate()\n raise TimeoutError(f\"{task.__name__} timed out after {timeout} seconds\")\n\n if not q.empty():\n status, result = q.get()\n if status == \"success\":\n return result\n else:\n raise RuntimeError(f\"{task.__name__} failed:\\n{result}\")\n else:\n raise RuntimeError(f\"{task.__name__} returned no result\")\n" }, { "path": "backend/tests/integration/common_utils/vespa.py", "content": "import requests\n\nfrom onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT\n\n\nclass vespa_fixture:\n def __init__(self, index_name: str):\n self.index_name = index_name\n self.vespa_document_url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)\n\n def get_documents_by_id(\n self, document_ids: list[str], wanted_doc_count: int = 1_000\n ) -> dict:\n selection = \" or \".join(\n f\"{self.index_name}.document_id=='{document_id}'\"\n for document_id in document_ids\n )\n params = {\n \"selection\": selection,\n \"wantedDocumentCount\": wanted_doc_count,\n }\n response = requests.get(\n self.vespa_document_url,\n params=params, # type: ignore\n )\n response.raise_for_status()\n return response.json()\n" }, { "path": "backend/tests/integration/conftest.py", "content": "import os\nfrom collections.abc import Callable\n\nimport pytest\n\n# Integration tests rely on this mode to enable mock_llm_response paths.\nos.environ[\"INTEGRATION_TESTS_MODE\"] = \"true\"\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.search_settings import get_current_search_settings\nfrom tests.integration.common_utils.constants import ADMIN_USER_NAME\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.image_generation import (\n ImageGenerationConfigManager,\n)\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import build_email\nfrom tests.integration.common_utils.managers.user import DEFAULT_PASSWORD\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.reset import reset_all_multitenant\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestImageGenerationConfig\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import SimpleTestDocument\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\nBASIC_USER_NAME = \"basic_user\"\n\nDocumentBuilderType = Callable[[list[str]], list[SimpleTestDocument]]\n\n\n@pytest.fixture(scope=\"session\", autouse=True)\ndef initialize_db() -> None:\n # Make sure that the db engine is initialized before any tests are run\n SqlEngine.init_engine(\n pool_size=10,\n max_overflow=5,\n )\n\n\ndef load_env_vars(env_file: str = \".env\") -> None:\n current_dir = os.path.dirname(os.path.abspath(__file__))\n env_path = os.path.join(current_dir, env_file)\n try:\n with open(env_path, \"r\") as f:\n for line in f:\n line = line.strip()\n if line and not line.startswith(\"#\"):\n key, value = line.split(\"=\", 1)\n # Preserve explicitly pre-set vars (e.g. INTEGRATION_TESTS_MODE).\n os.environ.setdefault(key, value.strip())\n print(\"Successfully loaded environment variables\")\n except FileNotFoundError:\n print(f\"File {env_file} not found\")\n\n\n# Load environment variables at the module level\nload_env_vars()\n\n\n\"\"\"NOTE: for some reason using this seems to lead to misc\n`sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) server closed the connection unexpectedly`\nerrors.\n\nCommenting out till we can get to the bottom of it. For now, just using\ninstantiate the session directly within the test.\n\"\"\"\n\n\n@pytest.fixture\ndef vespa_client() -> vespa_fixture:\n with get_session_with_current_tenant() as db_session:\n search_settings = get_current_search_settings(db_session)\n return vespa_fixture(index_name=search_settings.index_name)\n\n\n@pytest.fixture\ndef reset() -> None:\n reset_all()\n\n\n@pytest.fixture\ndef new_admin_user(reset: None) -> DATestUser: # noqa: ARG001\n return UserManager.create(name=ADMIN_USER_NAME)\n\n\n@pytest.fixture\ndef admin_user() -> DATestUser:\n try:\n user = UserManager.create(name=ADMIN_USER_NAME)\n\n # if there are other users for some reason, reset and try again\n if not UserManager.is_role(user, UserRole.ADMIN):\n print(\"Trying to reset\")\n reset_all()\n user = UserManager.create(name=ADMIN_USER_NAME)\n return user\n except Exception as e:\n print(f\"Failed to create admin user: {e}\")\n\n try:\n user = UserManager.login_as_user(\n DATestUser(\n id=\"\",\n email=build_email(\"admin_user\"),\n password=DEFAULT_PASSWORD,\n headers=GENERAL_HEADERS,\n role=UserRole.ADMIN,\n is_active=True,\n )\n )\n if not UserManager.is_role(user, UserRole.ADMIN):\n reset_all()\n user = UserManager.create(name=ADMIN_USER_NAME)\n return user\n\n return user\n except Exception as e:\n print(f\"Failed to create or login as admin user: {e}\")\n\n raise RuntimeError(\"Failed to create or login as admin user\")\n\n\n@pytest.fixture\ndef basic_user(\n # make sure the admin user exists first to ensure this new user\n # gets the BASIC role\n admin_user: DATestUser, # noqa: ARG001\n) -> DATestUser:\n try:\n user = UserManager.create(name=BASIC_USER_NAME)\n\n # Validate that the user has the BASIC role\n if user.role != UserRole.BASIC:\n raise RuntimeError(\n f\"Created user {BASIC_USER_NAME} does not have BASIC role\"\n )\n\n return user\n except Exception as e:\n print(f\"Failed to create basic user, trying to login as existing user: {e}\")\n\n # Try to login as existing basic user\n user = UserManager.login_as_user(\n DATestUser(\n id=\"\",\n email=build_email(BASIC_USER_NAME),\n password=DEFAULT_PASSWORD,\n headers=GENERAL_HEADERS,\n role=UserRole.BASIC,\n is_active=True,\n )\n )\n\n # Validate that the logged-in user has the BASIC role\n if not UserManager.is_role(user, UserRole.BASIC):\n raise RuntimeError(f\"User {BASIC_USER_NAME} does not have BASIC role\")\n\n return user\n\n\n@pytest.fixture(scope=\"session\")\ndef reset_multitenant() -> None:\n \"\"\"Initialize multi-tenant state once per test session.\n\n Intentionally avoid per-test resets to speed up the multitenant suite.\n The underlying reset function honors SKIP_RESET to allow CI to disable\n heavy resets entirely.\n \"\"\"\n reset_all_multitenant()\n\n\n@pytest.fixture\ndef llm_provider(admin_user: DATestUser) -> DATestLLMProvider:\n return LLMProviderManager.create(user_performing_action=admin_user)\n\n\n@pytest.fixture\ndef image_generation_config(\n admin_user: DATestUser,\n) -> DATestImageGenerationConfig:\n \"\"\"Create a default image generation config for tests.\"\"\"\n return ImageGenerationConfigManager.create(\n user_performing_action=admin_user,\n is_default=True,\n )\n\n\n@pytest.fixture\ndef document_builder(admin_user: DATestUser) -> DocumentBuilderType:\n api_key: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user,\n )\n\n # create connector\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n def _document_builder(contents: list[str]) -> list[SimpleTestDocument]:\n # seed documents\n docs: list[SimpleTestDocument] = [\n DocumentManager.seed_doc_with_content(\n cc_pair=cc_pair_1,\n content=content,\n api_key=api_key,\n )\n for content in contents\n ]\n\n return docs\n\n return _document_builder\n\n\ndef pytest_runtest_logstart(\n nodeid: str,\n location: tuple[str, int | None, str], # noqa: ARG001\n) -> None:\n print(f\"\\nTest start: {nodeid}\")\n\n\ndef pytest_runtest_logfinish(\n nodeid: str,\n location: tuple[str, int | None, str], # noqa: ARG001\n) -> None:\n print(f\"\\nTest end: {nodeid}\")\n" }, { "path": "backend/tests/integration/connector_job_tests/github/conftest.py", "content": "import os\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nGitHubTestEnvSetupTuple = tuple[\n DATestUser, # admin_user\n DATestUser, # test_user_1\n DATestUser, # test_user_2\n DATestCredential, # github_credential\n DATestConnector, # github_connector\n DATestCCPair, # github_cc_pair\n]\n\n\ndef _get_github_test_tokens() -> list[str]:\n \"\"\"\n Returns a list of GitHub tokens to run the GitHub connector suite against.\n\n Minimal setup:\n - Set ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN (token1)\n Optional:\n - Set ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC (token2 / classic)\n\n If the classic token is provided, the GitHub suite will run twice (once per token).\n \"\"\"\n token_1 = os.environ.get(\"ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN\")\n # Prefer the new \"classic\" name, but keep backward compatibility.\n token_2 = os.environ.get(\"ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC\")\n\n tokens: list[str] = []\n if token_1:\n tokens.append(token_1)\n if token_2:\n tokens.append(token_2)\n return tokens\n\n\n@pytest.fixture(scope=\"module\", params=_get_github_test_tokens())\ndef github_access_token(request: pytest.FixtureRequest) -> str:\n tokens = _get_github_test_tokens()\n if not tokens:\n pytest.skip(\n \"Skipping GitHub tests due to missing env vars \"\n \"ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN and \"\n \"ONYX_GITHUB_PERMISSION_SYNC_TEST_ACCESS_TOKEN_CLASSIC\"\n )\n return request.param\n\n\n@pytest.fixture(scope=\"module\")\ndef github_test_env_setup(\n github_access_token: str,\n) -> Generator[GitHubTestEnvSetupTuple]:\n \"\"\"\n Create a complete GitHub test environment with:\n - 3 users with email IDs from environment variables\n - GitHub credentials using ACCESS_TOKEN_GITHUB from environment\n - GitHub connector configured for testing\n - Connector-Credential pair linking them together\n\n Returns:\n Tuple containing: (admin_user, test_user_1, test_user_2, github_credential, github_connector, github_cc_pair)\n \"\"\"\n # Reset all resources before setting up the test environment\n reset_all()\n\n # Get user emails from environment (with fallbacks)\n admin_email = os.environ.get(\"ONYX_GITHUB_ADMIN_EMAIL\", \"admin@example.com\")\n test_user_1_email = os.environ.get(\n \"ONYX_GITHUB_TEST_USER_1_EMAIL\", \"subash@onyx.app\"\n )\n test_user_2_email = os.environ.get(\n \"ONYX_GITHUB_TEST_USER_2_EMAIL\", \"msubash203@gmail.com\"\n )\n\n if not admin_email or not test_user_1_email or not test_user_2_email:\n pytest.skip(\n \"Skipping GitHub test environment setup due to missing environment variables\"\n )\n\n # Create users\n admin_user: DATestUser = UserManager.create(email=admin_email)\n test_user_1: DATestUser = UserManager.create(email=test_user_1_email)\n test_user_2: DATestUser = UserManager.create(email=test_user_2_email)\n\n # Create LLM provider - required for document search to work\n LLMProviderManager.create(user_performing_action=admin_user)\n\n # Create GitHub credentials\n github_credentials = {\n \"github_access_token\": github_access_token,\n }\n\n github_credential: DATestCredential = CredentialManager.create(\n source=DocumentSource.GITHUB,\n credential_json=github_credentials,\n user_performing_action=admin_user,\n )\n\n # Create GitHub connector\n github_connector: DATestConnector = ConnectorManager.create(\n name=\"GitHub Test Connector\",\n input_type=InputType.POLL,\n source=DocumentSource.GITHUB,\n connector_specific_config={\n \"repo_owner\": \"permission-sync-test\",\n \"include_prs\": True,\n \"repositories\": \"perm-sync-test-minimal\",\n \"include_issues\": True,\n },\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n\n # Create CC pair linking connector and credential\n github_cc_pair: DATestCCPair = CCPairManager.create(\n credential_id=github_credential.id,\n connector_id=github_connector.id,\n name=\"GitHub Test CC Pair\",\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n\n # Wait for initial indexing to complete\n # GitHub API operations can be slow due to rate limiting and network latency\n # Use a longer timeout for initial indexing to avoid flaky test failures\n before = datetime.now(tz=timezone.utc)\n CCPairManager.wait_for_indexing_completion(\n cc_pair=github_cc_pair,\n after=before,\n user_performing_action=admin_user,\n timeout=900,\n )\n\n yield admin_user, test_user_1, test_user_2, github_credential, github_connector, github_cc_pair\n" }, { "path": "backend/tests/integration/connector_job_tests/github/test_github_permission_sync.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\nfrom github import Github\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.utils.logger import setup_logger\nfrom tests.integration.common_utils.document_acl import (\n get_all_connector_documents,\n)\nfrom tests.integration.common_utils.document_acl import (\n get_user_document_access_via_acl,\n)\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.connector_job_tests.github.conftest import (\n GitHubTestEnvSetupTuple,\n)\nfrom tests.integration.connector_job_tests.github.utils import GitHubManager\n\nlogger = setup_logger()\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_github_private_repo_permission_sync(\n github_test_env_setup: GitHubTestEnvSetupTuple,\n) -> None:\n\n (\n admin_user,\n test_user_1,\n test_user_2,\n github_credential,\n github_connector,\n github_cc_pair,\n ) = github_test_env_setup\n\n # Create GitHub client from credential\n # Note: github_credential is a DATestCredential (Pydantic model), not a SQLAlchemy model\n # so credential_json is already a plain dict\n github_access_token = github_credential.credential_json[\"github_access_token\"]\n github_client = Github(github_access_token)\n github_manager = GitHubManager(github_client)\n\n # Get repository configuration from connector\n repo_owner = github_connector.connector_specific_config[\"repo_owner\"]\n repo_name = github_connector.connector_specific_config[\"repositories\"]\n\n success = github_manager.change_repository_visibility(\n repo_owner=repo_owner, repo_name=repo_name, visibility=\"private\"\n )\n\n if not success:\n pytest.fail(f\"Failed to change repository {repo_owner}/{repo_name} to private\")\n\n # Add test-team to repository at the start\n logger.info(f\"Adding test-team to repository {repo_owner}/{repo_name}\")\n team_added = github_manager.add_team_to_repository(\n repo_owner=repo_owner,\n repo_name=repo_name,\n team_slug=\"test-team\",\n permission=\"pull\",\n )\n\n if not team_added:\n logger.warning(\n f\"Failed to add test-team to repository {repo_owner}/{repo_name}\"\n )\n\n try:\n after = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=github_cc_pair,\n user_performing_action=admin_user,\n )\n\n # Use a longer timeout for GitHub permission sync operations\n # GitHub API operations can be slow, especially with rate limiting\n # This accounts for document sync, group sync, and vespa sync operations\n CCPairManager.wait_for_sync(\n cc_pair=github_cc_pair,\n user_performing_action=admin_user,\n after=after,\n should_wait_for_group_sync=True,\n timeout=900,\n )\n\n # ACL-based verification\n with get_session_with_current_tenant() as db_session:\n # Get all documents for this connector\n all_document_ids = get_all_connector_documents(github_cc_pair, db_session)\n\n # Test access for both users using ACL verification\n accessible_docs_user1 = get_user_document_access_via_acl(\n test_user=test_user_1,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n accessible_docs_user2 = get_user_document_access_via_acl(\n test_user=test_user_2,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n logger.info(\n f\"test_user_1 has access to {len(accessible_docs_user1)} documents\"\n )\n logger.info(\n f\"test_user_2 has access to {len(accessible_docs_user2)} documents\"\n )\n\n # test_user_1 (part of test-team) should have access\n # test_user_2 (not part of test-team) should NOT have access\n assert len(accessible_docs_user1) > 0, (\n f\"test_user_1 should have access to private repository documents. \"\n f\"Found {len(accessible_docs_user1)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n assert len(accessible_docs_user2) == 0, (\n f\"test_user_2 should NOT have access to private repository documents. \"\n f\"Found {len(accessible_docs_user2)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n\n finally:\n # Remove test-team from repository at the end\n logger.info(f\"Removing test-team from repository {repo_owner}/{repo_name}\")\n team_removed = github_manager.remove_team_from_repository(\n repo_owner=repo_owner, repo_name=repo_name, team_slug=\"test-team\"\n )\n\n if not team_removed:\n logger.warning(\n f\"Failed to remove test-team from repository {repo_owner}/{repo_name}\"\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_github_public_repo_permission_sync(\n github_test_env_setup: GitHubTestEnvSetupTuple,\n) -> None:\n \"\"\"\n Test that when a repository is changed to public, both users can access the documents.\n \"\"\"\n (\n admin_user,\n test_user_1,\n test_user_2,\n github_credential,\n github_connector,\n github_cc_pair,\n ) = github_test_env_setup\n\n # Create GitHub client from credential\n # Note: github_credential is a DATestCredential (Pydantic model), not a SQLAlchemy model\n # so credential_json is already a plain dict\n github_access_token = github_credential.credential_json[\"github_access_token\"]\n github_client = Github(github_access_token)\n github_manager = GitHubManager(github_client)\n\n # Get repository configuration from connector\n repo_owner = github_connector.connector_specific_config[\"repo_owner\"]\n repo_name = github_connector.connector_specific_config[\"repositories\"]\n\n # Change repository to public\n logger.info(f\"Changing repository {repo_owner}/{repo_name} to public\")\n success = github_manager.change_repository_visibility(\n repo_owner=repo_owner, repo_name=repo_name, visibility=\"public\"\n )\n\n if not success:\n pytest.fail(f\"Failed to change repository {repo_owner}/{repo_name} to public\")\n\n # Verify repository is now public\n current_visibility = github_manager.get_repository_visibility(\n repo_owner=repo_owner, repo_name=repo_name\n )\n logger.info(f\"Repository {repo_owner}/{repo_name} visibility: {current_visibility}\")\n assert (\n current_visibility == \"public\"\n ), f\"Repository should be public, but is {current_visibility}\"\n\n # Trigger sync to update permissions\n after = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=github_cc_pair,\n user_performing_action=admin_user,\n )\n\n # Wait for sync to complete with group sync\n # Public repositories should be accessible to all users\n CCPairManager.wait_for_sync(\n cc_pair=github_cc_pair,\n user_performing_action=admin_user,\n after=after,\n should_wait_for_group_sync=True,\n timeout=900,\n )\n\n # ACL-based verification\n with get_session_with_current_tenant() as db_session:\n # Get all documents for this connector\n all_document_ids = get_all_connector_documents(github_cc_pair, db_session)\n\n # Test access for both users using ACL verification\n accessible_docs_user1 = get_user_document_access_via_acl(\n test_user=test_user_1,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n accessible_docs_user2 = get_user_document_access_via_acl(\n test_user=test_user_2,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n logger.info(f\"test_user_1 has access to {len(accessible_docs_user1)} documents\")\n logger.info(f\"test_user_2 has access to {len(accessible_docs_user2)} documents\")\n\n # Both users should have access to the public repository documents\n assert len(accessible_docs_user1) > 0, (\n f\"test_user_1 should have access to public repository documents. \"\n f\"Found {len(accessible_docs_user1)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n assert len(accessible_docs_user2) > 0, (\n f\"test_user_2 should have access to public repository documents. \"\n f\"Found {len(accessible_docs_user2)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n\n # Verify that both users get the same results (since repo is public)\n assert len(accessible_docs_user1) == len(accessible_docs_user2), (\n f\"Both users should see the same documents from public repository. \"\n f\"User1: {len(accessible_docs_user1)}, User2: {len(accessible_docs_user2)}\"\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_github_internal_repo_permission_sync(\n github_test_env_setup: GitHubTestEnvSetupTuple,\n) -> None:\n \"\"\"\n Test that when a repository is changed to internal, test_user_1 has access but test_user_2 doesn't.\n Internal repositories are accessible only to organization members.\n \"\"\"\n (\n admin_user,\n test_user_1,\n test_user_2,\n github_credential,\n github_connector,\n github_cc_pair,\n ) = github_test_env_setup\n\n # Create GitHub client from credential\n # Note: github_credential is a DATestCredential (Pydantic model), not a SQLAlchemy model\n # so credential_json is already a plain dict\n github_access_token = github_credential.credential_json[\"github_access_token\"]\n github_client = Github(github_access_token)\n github_manager = GitHubManager(github_client)\n\n # Get repository configuration from connector\n repo_owner = github_connector.connector_specific_config[\"repo_owner\"]\n repo_name = github_connector.connector_specific_config[\"repositories\"]\n\n # Change repository to internal\n logger.info(f\"Changing repository {repo_owner}/{repo_name} to internal\")\n success = github_manager.change_repository_visibility(\n repo_owner=repo_owner, repo_name=repo_name, visibility=\"internal\"\n )\n\n if not success:\n pytest.fail(f\"Failed to change repository {repo_owner}/{repo_name} to internal\")\n\n # Verify repository is now internal\n current_visibility = github_manager.get_repository_visibility(\n repo_owner=repo_owner, repo_name=repo_name\n )\n logger.info(f\"Repository {repo_owner}/{repo_name} visibility: {current_visibility}\")\n assert (\n current_visibility == \"internal\"\n ), f\"Repository should be internal, but is {current_visibility}\"\n\n # Trigger sync to update permissions\n after = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=github_cc_pair,\n user_performing_action=admin_user,\n )\n\n # Wait for sync to complete with group sync\n # Internal repositories should be accessible only to organization members\n CCPairManager.wait_for_sync(\n cc_pair=github_cc_pair,\n user_performing_action=admin_user,\n after=after,\n should_wait_for_group_sync=True,\n timeout=900,\n )\n\n # ACL-based verification\n with get_session_with_current_tenant() as db_session:\n # Get all documents for this connector\n all_document_ids = get_all_connector_documents(github_cc_pair, db_session)\n\n # Test access for both users using ACL verification\n accessible_docs_user1 = get_user_document_access_via_acl(\n test_user=test_user_1,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n accessible_docs_user2 = get_user_document_access_via_acl(\n test_user=test_user_2,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n logger.info(f\"test_user_1 has access to {len(accessible_docs_user1)} documents\")\n logger.info(f\"test_user_2 has access to {len(accessible_docs_user2)} documents\")\n\n # For internal repositories:\n # - test_user_1 should have access (assuming they're part of the organization)\n # - test_user_2 should NOT have access (assuming they're not part of the organization)\n assert len(accessible_docs_user1) > 0, (\n f\"test_user_1 should have access to internal repository documents (organization member). \"\n f\"Found {len(accessible_docs_user1)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n assert len(accessible_docs_user2) == 0, (\n f\"test_user_2 should NOT have access to internal repository documents (not organization member). \"\n f\"Found {len(accessible_docs_user2)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n" }, { "path": "backend/tests/integration/connector_job_tests/github/utils.py", "content": "from typing import Optional\n\nfrom github import Github\nfrom github.GithubException import GithubException\n\nfrom onyx.utils.logger import setup_logger\n\nlogger = setup_logger()\n\n\nclass GitHubManager:\n \"\"\"\n Manager class for GitHub operations used in testing.\n Provides methods to change repository visibility, check repository visibility, and manage teams.\n \"\"\"\n\n def __init__(self, github_client: Github):\n \"\"\"\n Initialize the GitHub manager with a GitHub client.\n\n Args:\n github_client: Authenticated GitHub client instance\n \"\"\"\n self.github_client = github_client\n\n def change_repository_visibility(\n self, repo_owner: str, repo_name: str, visibility: str\n ) -> bool:\n \"\"\"\n Change the visibility of a repository.\n\n Args:\n repo_owner: Repository owner (organization or username)\n repo_name: Repository name\n visibility: New visibility ('public', 'private', or 'internal')\n\n Returns:\n bool: True if successful, False otherwise\n\n Raises:\n ValueError: If visibility is not valid\n GithubException: If GitHub API call fails\n \"\"\"\n if visibility not in [\"public\", \"private\", \"internal\"]:\n raise ValueError(\n f\"Invalid visibility: {visibility}. Must be 'public', 'private', or 'internal'\"\n )\n\n try:\n repo = self.github_client.get_repo(f\"{repo_owner}/{repo_name}\")\n\n # Check if we have admin permissions\n if not repo.permissions.admin:\n logger.error(\n f\"No admin permissions for repository {repo_owner}/{repo_name}\"\n )\n return False\n\n # Note: Internal repositories are only available for GitHub Enterprise\n try:\n repo.edit(visibility=visibility)\n except GithubException as e:\n logger.warning(f\"Could not set repository to {visibility}: {e}\")\n return False\n\n logger.info(\n f\"Successfully changed {repo_owner}/{repo_name} visibility to {visibility}\"\n )\n return True\n\n except GithubException as e:\n logger.error(f\"Failed to change repository visibility: {e}\")\n return False\n\n def add_team_to_repository(\n self, repo_owner: str, repo_name: str, team_slug: str, permission: str = \"push\"\n ) -> bool:\n \"\"\"\n Add a team to a repository with specified permissions.\n\n Args:\n repo_owner: Repository owner (organization)\n repo_name: Repository name\n team_slug: Team slug (not team name)\n permission: Permission level ('pull', 'push', 'admin', 'maintain', 'triage')\n\n Returns:\n bool: True if successful, False otherwise\n\n Raises:\n GithubException: If GitHub API call fails\n \"\"\"\n valid_permissions = [\"pull\", \"push\", \"admin\", \"maintain\", \"triage\"]\n if permission not in valid_permissions:\n raise ValueError(\n f\"Invalid permission: {permission}. Must be one of {valid_permissions}\"\n )\n\n try:\n repo = self.github_client.get_repo(f\"{repo_owner}/{repo_name}\")\n org = self.github_client.get_organization(repo_owner)\n team = org.get_team_by_slug(team_slug)\n\n # Add team to repository\n team.add_to_repos(repo)\n\n # Set team permissions on the repository\n team.set_repo_permission(repo, permission)\n\n logger.info(\n f\"Successfully added team {team_slug} to {repo_owner}/{repo_name} with {permission} permissions\"\n )\n return True\n\n except GithubException as e:\n logger.error(f\"Failed to add team to repository: {e}\")\n return False\n\n def remove_team_from_repository(\n self, repo_owner: str, repo_name: str, team_slug: str\n ) -> bool:\n \"\"\"\n Remove a team from a repository.\n\n Args:\n repo_owner: Repository owner (organization)\n repo_name: Repository name\n team_slug: Team slug (not team name)\n\n Returns:\n bool: True if successful, False otherwise\n\n Raises:\n GithubException: If GitHub API call fails\n \"\"\"\n try:\n repo = self.github_client.get_repo(f\"{repo_owner}/{repo_name}\")\n org = self.github_client.get_organization(repo_owner)\n team = org.get_team_by_slug(team_slug)\n\n # Remove team from repository\n team.remove_from_repos(repo)\n\n logger.info(\n f\"Successfully removed team {team_slug} from {repo_owner}/{repo_name}\"\n )\n return True\n\n except GithubException as e:\n logger.error(f\"Failed to remove team from repository: {e}\")\n return False\n\n def get_repository_visibility(\n self, repo_owner: str, repo_name: str\n ) -> Optional[str]:\n \"\"\"\n Get the current visibility of a repository.\n\n Args:\n repo_owner: Repository owner\n repo_name: Repository name\n\n Returns:\n Optional[str]: Repository visibility ('public', 'private', 'internal') or None if failed\n \"\"\"\n try:\n repo = self.github_client.get_repo(f\"{repo_owner}/{repo_name}\")\n\n if hasattr(repo, \"visibility\"):\n return repo.visibility\n else:\n # Fallback for older GitHub API versions\n return \"private\" if repo.private else \"public\"\n\n except GithubException as e:\n logger.error(f\"Failed to get repository visibility: {e}\")\n return None\n" }, { "path": "backend/tests/integration/connector_job_tests/google/google_drive_api_utils.py", "content": "from typing import Any\nfrom uuid import uuid4\n\nfrom google.oauth2.service_account import Credentials\n\nfrom onyx.connectors.google_utils.resources import get_drive_service\nfrom onyx.connectors.google_utils.resources import get_google_docs_service\nfrom onyx.connectors.google_utils.resources import GoogleDocsService\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\n\n\nGOOGLE_SCOPES = {\n \"google_drive\": [\n \"https://www.googleapis.com/auth/drive\",\n \"https://www.googleapis.com/auth/admin.directory.group\",\n \"https://www.googleapis.com/auth/admin.directory.user\",\n ],\n}\n\n\ndef _create_doc_service(drive_service: GoogleDriveService) -> GoogleDocsService:\n docs_service = get_google_docs_service(\n creds=drive_service._http.credentials,\n user_email=drive_service._http.credentials._subject,\n )\n return docs_service\n\n\nclass GoogleDriveManager:\n @staticmethod\n def create_impersonated_drive_service(\n service_account_key: dict, impersonated_user_email: str\n ) -> GoogleDriveService:\n \"\"\"Gets a drive service that impersonates a specific user\"\"\"\n credentials = Credentials.from_service_account_info(\n service_account_key,\n scopes=GOOGLE_SCOPES[\"google_drive\"],\n subject=impersonated_user_email,\n )\n\n service = get_drive_service(credentials, impersonated_user_email)\n\n # Verify impersonation\n about = service.about().get(fields=\"user\").execute()\n if about.get(\"user\", {}).get(\"emailAddress\") != impersonated_user_email:\n raise ValueError(\n f\"Failed to impersonate {impersonated_user_email}. Instead got {about.get('user', {}).get('emailAddress')}\"\n )\n return service\n\n @staticmethod\n def create_shared_drive(\n drive_service: GoogleDriveService, admin_email: str, test_id: str\n ) -> str:\n \"\"\"\n Creates a shared drive and returns the drive's ID\n \"\"\"\n try:\n about = drive_service.about().get(fields=\"user\").execute()\n creating_user = about[\"user\"][\"emailAddress\"]\n\n # Verify we're still impersonating the admin\n if creating_user != admin_email:\n raise ValueError(\n f\"Expected to create drive as {admin_email}, but instead created drive as {creating_user}\"\n )\n\n drive_metadata = {\"name\": f\"perm_sync_drive_{test_id}\"}\n\n request_id = str(uuid4())\n drive = (\n drive_service.drives()\n .create(\n body=drive_metadata,\n requestId=request_id,\n fields=\"id,name,capabilities\",\n )\n .execute()\n )\n\n return drive[\"id\"]\n except Exception as e:\n print(f\"Error creating shared drive: {str(e)}\")\n raise\n\n @staticmethod\n def create_empty_doc(\n drive_service: Any,\n drive_id: str,\n ) -> str:\n \"\"\"\n Creates an empty document in the given drive and returns the document's ID\n \"\"\"\n file_metadata = {\n \"name\": f\"perm_sync_doc_{drive_id}_{str(uuid4())}\",\n \"mimeType\": \"application/vnd.google-apps.document\",\n \"parents\": [drive_id],\n }\n file = (\n drive_service.files()\n .create(body=file_metadata, supportsAllDrives=True)\n .execute()\n )\n\n return file[\"id\"]\n\n @staticmethod\n def append_text_to_doc(\n drive_service: GoogleDriveService, doc_id: str, text: str\n ) -> None:\n docs_service = _create_doc_service(drive_service)\n\n docs_service.documents().batchUpdate(\n documentId=doc_id,\n body={\n \"requests\": [{\"insertText\": {\"location\": {\"index\": 1}, \"text\": text}}]\n },\n ).execute()\n\n @staticmethod\n def update_file_permissions(\n drive_service: Any, file_id: str, email: str, role: str = \"reader\"\n ) -> None:\n permission = {\"type\": \"user\", \"role\": role, \"emailAddress\": email}\n drive_service.permissions().create(\n fileId=file_id,\n body=permission,\n supportsAllDrives=True,\n sendNotificationEmail=False,\n ).execute()\n\n @staticmethod\n def remove_file_permissions(\n drive_service: Any,\n file_id: str,\n email: str, # noqa: ARG004\n ) -> None:\n permissions = (\n drive_service.permissions()\n .list(fileId=file_id, supportsAllDrives=True)\n .execute()\n )\n # TODO: This is a hacky way to remove permissions. Removes anyone with reader role.\n # Need to find a way to map a user's email to a permission id.\n # The permissions.get returns a permissionID but email field is None,\n # something to do with it being a group or domain wide delegation.\n for permission in permissions.get(\"permissions\", []):\n if permission.get(\"role\") == \"reader\":\n drive_service.permissions().delete(\n fileId=file_id,\n permissionId=permission[\"id\"],\n supportsAllDrives=True,\n ).execute()\n break\n\n @staticmethod\n def make_file_public(drive_service: Any, file_id: str) -> None:\n permission = {\"type\": \"anyone\", \"role\": \"reader\"}\n drive_service.permissions().create(\n fileId=file_id, body=permission, supportsAllDrives=True\n ).execute()\n\n @staticmethod\n def cleanup_drive(drive_service: Any, drive_id: str) -> None:\n try:\n # Delete up to 2 files that match our pattern\n file_name_prefix = f\"perm_sync_doc_{drive_id}\"\n files = (\n drive_service.files()\n .list(\n q=f\"name contains '{file_name_prefix}'\",\n driveId=drive_id,\n includeItemsFromAllDrives=True,\n supportsAllDrives=True,\n corpora=\"drive\",\n fields=\"files(id)\",\n )\n .execute()\n )\n\n for file in files.get(\"files\", []):\n drive_service.files().delete(\n fileId=file[\"id\"], supportsAllDrives=True\n ).execute()\n\n # Then delete the drive\n drive_service.drives().delete(driveId=drive_id).execute()\n except Exception as e:\n print(f\"Error cleaning up drive {drive_id}: {e}\")\n" }, { "path": "backend/tests/integration/connector_job_tests/google/test_google_drive_permission_sync.py", "content": "import json\nimport os\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\nfrom uuid import uuid4\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.google_utils.resources import GoogleDriveService\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,\n)\nfrom onyx.connectors.google_utils.shared_constants import (\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY,\n)\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.document_search import (\n DocumentSearchManager,\n)\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\nfrom tests.integration.connector_job_tests.google.google_drive_api_utils import (\n GoogleDriveManager,\n)\n\n\n@pytest.fixture()\ndef google_drive_test_env_setup() -> Generator[\n tuple[GoogleDriveService, str, DATestCCPair, DATestUser, DATestUser, DATestUser],\n None,\n None,\n]:\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n # Creating a non-admin user\n test_user_1: DATestUser = UserManager.create(email=\"test_user_1@example.com\")\n # Creating a non-admin user\n test_user_2: DATestUser = UserManager.create(email=\"test_user_2@example.com\")\n\n service_account_key = os.environ[\"FULL_CONTROL_DRIVE_SERVICE_ACCOUNT\"]\n drive_id: str | None = None\n drive_service: GoogleDriveService | None = None\n\n try:\n credentials = {\n DB_CREDENTIALS_PRIMARY_ADMIN_KEY: admin_user.email,\n DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: service_account_key,\n }\n\n # Setup Google Drive\n drive_service = GoogleDriveManager.create_impersonated_drive_service(\n json.loads(service_account_key), admin_user.email\n )\n test_id = str(uuid4())\n drive_id = GoogleDriveManager.create_shared_drive(\n drive_service, admin_user.email, test_id\n )\n\n # Setup Onyx infrastructure\n LLMProviderManager.create(user_performing_action=admin_user)\n\n before = datetime.now(timezone.utc)\n credential: DATestCredential = CredentialManager.create(\n source=DocumentSource.GOOGLE_DRIVE,\n credential_json=credentials,\n user_performing_action=admin_user,\n )\n connector: DATestConnector = ConnectorManager.create(\n name=\"Google Drive Test\",\n input_type=InputType.POLL,\n source=DocumentSource.GOOGLE_DRIVE,\n connector_specific_config={\n \"shared_drive_urls\": f\"https://drive.google.com/drive/folders/{drive_id}\"\n },\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n cc_pair: DATestCCPair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair, after=before, user_performing_action=admin_user\n )\n\n yield drive_service, drive_id, cc_pair, admin_user, test_user_1, test_user_2\n\n except json.JSONDecodeError:\n pytest.skip(\"FULL_CONTROL_DRIVE_SERVICE_ACCOUNT is not valid JSON\")\n finally:\n # Cleanup drive and file\n if drive_id is not None:\n GoogleDriveManager.cleanup_drive(drive_service, drive_id)\n\n\n@pytest.mark.xfail(reason=\"Needs to be tested for flakiness\")\ndef test_google_permission_sync(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture, # noqa: ARG001\n google_drive_test_env_setup: tuple[\n GoogleDriveService, str, DATestCCPair, DATestUser, DATestUser, DATestUser\n ],\n) -> None:\n (\n drive_service,\n drive_id,\n cc_pair,\n admin_user,\n test_user_1,\n test_user_2,\n ) = google_drive_test_env_setup\n\n # ----------------------BASELINE TEST----------------------\n before = datetime.now(timezone.utc)\n\n # Create empty test doc in drive\n doc_id_1 = GoogleDriveManager.create_empty_doc(drive_service, drive_id)\n\n # Append text to doc\n doc_text_1 = \"The secret number is 12345\"\n GoogleDriveManager.append_text_to_doc(drive_service, doc_id_1, doc_text_1)\n\n # run indexing\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair, after=before, user_performing_action=admin_user\n )\n\n # run permission sync\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n )\n\n # Verify admin has access to document\n admin_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=admin_user\n )\n assert doc_text_1 in [result.strip(\"\\ufeff\") for result in admin_results]\n\n # Verify test_user_1 cannot access document\n user1_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=test_user_1\n )\n assert doc_text_1 not in [result.strip(\"\\ufeff\") for result in user1_results]\n\n # ----------------------GRANT USER 1 DOC PERMISSIONS TEST--------------------------\n before = datetime.now(timezone.utc)\n\n # Grant user 1 access to document 1\n GoogleDriveManager.update_file_permissions(\n drive_service=drive_service,\n file_id=doc_id_1,\n email=test_user_1.email,\n role=\"reader\",\n )\n\n # Create a second doc in the drive which user 1 should not have access to\n doc_id_2 = GoogleDriveManager.create_empty_doc(drive_service, drive_id)\n doc_text_2 = \"The secret number is 67890\"\n GoogleDriveManager.append_text_to_doc(drive_service, doc_id_2, doc_text_2)\n\n # Run indexing\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # Run permission sync\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n )\n\n # Verify admin can access both documents\n admin_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=admin_user\n )\n assert {doc_text_1, doc_text_2} == {\n result.strip(\"\\ufeff\") for result in admin_results\n }\n\n # Verify user 1 can access document 1\n user1_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=test_user_1\n )\n assert doc_text_1 in [result.strip(\"\\ufeff\") for result in user1_results]\n\n # Verify user 1 cannot access document 2\n user1_results_2 = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=test_user_1\n )\n assert doc_text_2 not in [result.strip(\"\\ufeff\") for result in user1_results_2]\n\n # ----------------------REMOVE USER 1 DOC PERMISSIONS TEST--------------------------\n before = datetime.now(timezone.utc)\n\n # Remove user 1 access to document 1\n GoogleDriveManager.remove_file_permissions(\n drive_service=drive_service, file_id=doc_id_1, email=test_user_1.email\n )\n # Run permission sync\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n )\n\n # Verify admin can access both documents\n admin_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=admin_user\n )\n assert {doc_text_1, doc_text_2} == {\n result.strip(\"\\ufeff\") for result in admin_results\n }\n\n # Verify user 1 cannot access either document\n user1_results = DocumentSearchManager.search_documents(\n query=\"secret numbers\", user_performing_action=test_user_1\n )\n assert {result.strip(\"\\ufeff\") for result in user1_results} == set()\n\n # ----------------------GRANT USER 1 DRIVE PERMISSIONS TEST--------------------------\n before = datetime.now(timezone.utc)\n\n # Grant user 1 access to drive\n GoogleDriveManager.update_file_permissions(\n drive_service=drive_service,\n file_id=drive_id,\n email=test_user_1.email,\n role=\"reader\",\n )\n\n # Run permission sync\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=2,\n user_performing_action=admin_user,\n # if we are only updating the group definition for this test we use this varaiable,\n # since it doesn't result in a vespa sync so we don't want to wait for it\n should_wait_for_vespa_sync=False,\n )\n\n # Verify user 1 can access both documents\n user1_results = DocumentSearchManager.search_documents(\n query=\"secret numbers\", user_performing_action=test_user_1\n )\n assert {doc_text_1, doc_text_2} == {\n result.strip(\"\\ufeff\") for result in user1_results\n }\n\n # ----------------------MAKE DRIVE PUBLIC TEST--------------------------\n before = datetime.now(timezone.utc)\n\n # Unable to make drive itself public as Google's security policies prevent this, so we make the documents public instead\n GoogleDriveManager.make_file_public(drive_service, doc_id_1)\n GoogleDriveManager.make_file_public(drive_service, doc_id_2)\n\n # Run permission sync\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=2,\n user_performing_action=admin_user,\n )\n\n # Verify all users can access both documents\n admin_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=admin_user\n )\n assert {doc_text_1, doc_text_2} == {\n result.strip(\"\\ufeff\") for result in admin_results\n }\n\n user1_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=test_user_1\n )\n assert {doc_text_1, doc_text_2} == {\n result.strip(\"\\ufeff\") for result in user1_results\n }\n\n user2_results = DocumentSearchManager.search_documents(\n query=\"secret number\", user_performing_action=test_user_2\n )\n assert {doc_text_1, doc_text_2} == {\n result.strip(\"\\ufeff\") for result in user2_results\n }\n" }, { "path": "backend/tests/integration/connector_job_tests/jira/conftest.py", "content": "import os\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import AccessType\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nJiraTestEnvSetupTuple = tuple[\n DATestUser,\n DATestCredential,\n DATestConnector,\n DATestCCPair,\n]\n\n\n@pytest.fixture()\ndef jira_test_env_setup() -> Generator[JiraTestEnvSetupTuple]:\n jira_base_url = os.environ[\"JIRA_BASE_URL\"]\n jira_user_email = os.environ[\"JIRA_USER_EMAIL\"]\n jira_api_token = os.environ[\"JIRA_API_TOKEN\"]\n\n credentials = {\n \"jira_user_email\": jira_user_email,\n \"jira_api_token\": jira_api_token,\n }\n\n admin_user: DATestUser = UserManager.create(email=jira_user_email)\n credential: DATestCredential = CredentialManager.create(\n source=DocumentSource.JIRA,\n credential_json=credentials,\n user_performing_action=admin_user,\n )\n connector: DATestConnector = ConnectorManager.create(\n name=\"Jira Test\",\n input_type=InputType.POLL,\n source=DocumentSource.JIRA,\n connector_specific_config={\n \"jira_base_url\": jira_base_url,\n },\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n cc_pair: DATestCCPair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n before = datetime.now(tz=timezone.utc)\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair, after=before, user_performing_action=admin_user\n )\n\n yield admin_user, credential, connector, cc_pair\n" }, { "path": "backend/tests/integration/connector_job_tests/jira/test_jira_permission_sync_full.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.connector_job_tests.jira.conftest import JiraTestEnvSetupTuple\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Jira permission sync is enterprise only\",\n)\n@pytest.mark.xfail(reason=\"Needs to be tested for flakiness\")\ndef test_jira_permission_sync_full(\n reset: None, # noqa: ARG001\n jira_test_env_setup: JiraTestEnvSetupTuple,\n) -> None:\n (\n admin_user,\n credential,\n connector,\n cc_pair,\n ) = jira_test_env_setup\n\n before = datetime.now(tz=timezone.utc)\n\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n timeout=float(\"inf\"),\n )\n" }, { "path": "backend/tests/integration/connector_job_tests/sharepoint/conftest.py", "content": "import os\nfrom collections.abc import Generator\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.connectors.sharepoint.connector import SharepointAuthMethod\nfrom onyx.db.enums import AccessType\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\n\nSharepointTestEnvSetupTuple = tuple[\n DATestUser, # admin_user\n DATestUser, # regular_user_1\n DATestUser, # regular_user_2\n DATestCredential,\n DATestConnector,\n DATestCCPair,\n]\n\n\n@pytest.fixture(scope=\"module\")\ndef sharepoint_test_env_setup() -> Generator[SharepointTestEnvSetupTuple]:\n # Reset all data before running the test\n reset_all()\n # Required environment variables for SharePoint certificate authentication\n sp_client_id = os.environ.get(\"PERM_SYNC_SHAREPOINT_CLIENT_ID\")\n sp_private_key = os.environ.get(\"PERM_SYNC_SHAREPOINT_PRIVATE_KEY\")\n sp_certificate_password = os.environ.get(\n \"PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD\"\n )\n sp_directory_id = os.environ.get(\"PERM_SYNC_SHAREPOINT_DIRECTORY_ID\")\n sharepoint_sites = \"https://danswerai.sharepoint.com/sites/Permisisonsync\"\n admin_email = \"admin@onyx.app\"\n user1_email = \"subash@onyx.app\"\n user2_email = \"raunak@onyx.app\"\n\n if not sp_private_key or not sp_certificate_password or not sp_directory_id:\n pytest.skip(\"Skipping test because required environment variables are not set\")\n\n # Certificate-based credentials\n credentials = {\n \"authentication_method\": SharepointAuthMethod.CERTIFICATE.value,\n \"sp_client_id\": sp_client_id,\n \"sp_private_key\": sp_private_key,\n \"sp_certificate_password\": sp_certificate_password,\n \"sp_directory_id\": sp_directory_id,\n }\n\n # Create users\n admin_user: DATestUser = UserManager.create(email=admin_email)\n regular_user_1: DATestUser = UserManager.create(email=user1_email)\n regular_user_2: DATestUser = UserManager.create(email=user2_email)\n\n # Create LLM provider for search functionality\n LLMProviderManager.create(user_performing_action=admin_user)\n\n # Create credential\n credential: DATestCredential = CredentialManager.create(\n source=DocumentSource.SHAREPOINT,\n credential_json=credentials,\n user_performing_action=admin_user,\n )\n\n # Create connector with SharePoint-specific configuration\n connector: DATestConnector = ConnectorManager.create(\n name=\"SharePoint Test\",\n input_type=InputType.POLL,\n source=DocumentSource.SHAREPOINT,\n connector_specific_config={\n \"sites\": sharepoint_sites.split(\",\"),\n \"treat_sharing_link_as_public\": True,\n },\n access_type=AccessType.SYNC, # Enable permission sync\n user_performing_action=admin_user,\n )\n\n # Create CC pair with permission sync enabled\n cc_pair: DATestCCPair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.SYNC, # Enable permission sync\n user_performing_action=admin_user,\n )\n\n # Wait for both indexing and permission sync to complete\n before = datetime.now(tz=timezone.utc)\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n timeout=float(\"inf\"),\n )\n\n # Wait for permission sync completion specifically\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n timeout=float(\"inf\"),\n )\n\n yield admin_user, regular_user_1, regular_user_2, credential, connector, cc_pair\n" }, { "path": "backend/tests/integration/connector_job_tests/sharepoint/test_sharepoint_permissions.py", "content": "import os\n\nimport pytest\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.utils.logger import setup_logger\nfrom tests.integration.common_utils.document_acl import (\n get_all_connector_documents,\n)\nfrom tests.integration.common_utils.document_acl import (\n get_documents_by_permission_type,\n)\nfrom tests.integration.common_utils.document_acl import (\n get_user_document_access_via_acl,\n)\nfrom tests.integration.connector_job_tests.sharepoint.conftest import (\n SharepointTestEnvSetupTuple,\n)\n\nlogger = setup_logger()\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_public_documents_accessible_by_all_users(\n sharepoint_test_env_setup: SharepointTestEnvSetupTuple,\n) -> None:\n \"\"\"Test that public documents are accessible by both test users using ACL verification\"\"\"\n (\n admin_user,\n regular_user_1,\n regular_user_2,\n credential,\n connector,\n cc_pair,\n ) = sharepoint_test_env_setup\n\n with get_session_with_current_tenant() as db_session:\n # Get all documents for this connector\n all_document_ids = get_all_connector_documents(cc_pair, db_session)\n\n # Test that regular_user_1 can access documents\n accessible_docs_user1 = get_user_document_access_via_acl(\n test_user=regular_user_1,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n # Test that regular_user_2 can access documents\n accessible_docs_user2 = get_user_document_access_via_acl(\n test_user=regular_user_2,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n logger.info(f\"User 1 has access to {len(accessible_docs_user1)} documents\")\n logger.info(f\"User 2 has access to {len(accessible_docs_user2)} documents\")\n\n # For public documents, both users should have access to at least some docs\n assert len(accessible_docs_user1) == 8, (\n f\"User 1 should have access to documents. Found \"\n f\"{len(accessible_docs_user1)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n assert len(accessible_docs_user2) == 1, (\n f\"User 2 should have access to documents. Found \"\n f\"{len(accessible_docs_user2)} accessible docs out of \"\n f\"{len(all_document_ids)} total\"\n )\n\n logger.info(\n \"Successfully verified public documents are accessible by users via ACL\"\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_group_based_permissions(\n sharepoint_test_env_setup: SharepointTestEnvSetupTuple,\n) -> None:\n \"\"\"Test that documents with group permissions are accessible only by users in that group using ACL verification\"\"\"\n (\n admin_user,\n regular_user_1,\n regular_user_2,\n credential,\n connector,\n cc_pair,\n ) = sharepoint_test_env_setup\n\n with get_session_with_current_tenant() as db_session:\n # Get all documents for this connector\n all_document_ids = get_all_connector_documents(cc_pair, db_session)\n\n if not all_document_ids:\n pytest.skip(\"No documents found for connector - skipping test\")\n\n # Test access for both users\n accessible_docs_user1 = get_user_document_access_via_acl(\n test_user=regular_user_1,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n accessible_docs_user2 = get_user_document_access_via_acl(\n test_user=regular_user_2,\n document_ids=all_document_ids,\n db_session=db_session,\n )\n\n logger.info(f\"User 1 has access to {len(accessible_docs_user1)} documents\")\n logger.info(f\"User 2 has access to {len(accessible_docs_user2)} documents\")\n\n public_docs = get_documents_by_permission_type(all_document_ids, db_session)\n\n # Check if user 2 has access to any non-public documents\n non_public_access_user2 = [\n doc for doc in accessible_docs_user2 if doc not in public_docs\n ]\n\n assert (\n len(non_public_access_user2) == 0\n ), f\"User 2 should only have access to public documents. Found access to non-public docs: {non_public_access_user2}\"\n" }, { "path": "backend/tests/integration/connector_job_tests/slack/conftest.py", "content": "import os\nfrom collections.abc import Generator\n\nimport pytest\n\nfrom onyx.connectors.slack.models import ChannelType\nfrom tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager\n\nSLACK_ADMIN_EMAIL = os.environ.get(\"SLACK_ADMIN_EMAIL\", \"evan@onyx.app\")\nSLACK_TEST_USER_1_EMAIL = os.environ.get(\"SLACK_TEST_USER_1_EMAIL\", \"evan+1@onyx.app\")\nSLACK_TEST_USER_2_EMAIL = os.environ.get(\"SLACK_TEST_USER_2_EMAIL\", \"justin@onyx.app\")\n\n\ndef _provision_slack_channels(\n bot_token: str,\n) -> Generator[tuple[ChannelType, ChannelType], None, None]:\n slack_client = SlackManager.get_slack_client(bot_token)\n\n auth_info = slack_client.auth_test()\n print(f\"\\nSlack workspace: {auth_info.get('team')} ({auth_info.get('url')})\")\n\n user_map = SlackManager.build_slack_user_email_id_map(slack_client)\n if SLACK_ADMIN_EMAIL not in user_map:\n raise KeyError(\n f\"'{SLACK_ADMIN_EMAIL}' not found in Slack workspace. Available emails: {sorted(user_map.keys())}\"\n )\n admin_user_id = user_map[SLACK_ADMIN_EMAIL]\n\n (\n public_channel,\n private_channel,\n run_id,\n ) = SlackManager.get_and_provision_available_slack_channels(\n slack_client=slack_client, admin_user_id=admin_user_id\n )\n\n yield public_channel, private_channel\n\n SlackManager.cleanup_after_test(slack_client=slack_client, test_id=run_id)\n\n\n@pytest.fixture()\ndef slack_test_setup() -> Generator[tuple[ChannelType, ChannelType], None, None]:\n yield from _provision_slack_channels(os.environ[\"SLACK_BOT_TOKEN\"])\n\n\n@pytest.fixture()\ndef slack_perm_sync_test_setup() -> (\n Generator[tuple[ChannelType, ChannelType], None, None]\n):\n yield from _provision_slack_channels(os.environ[\"SLACK_BOT_TOKEN_TEST_SPACE\"])\n" }, { "path": "backend/tests/integration/connector_job_tests/slack/slack_api_utils.py", "content": "\"\"\"\nAssumptions:\n- The test users have already been created\n- General is empty of messages\n- In addition to the normal slack oauth permissions, the following scopes are needed:\n - channels:manage\n - groups:write\n - chat:write\n - chat:write.public\n\"\"\"\n\nfrom typing import Any\nfrom typing import cast\nfrom uuid import uuid4\n\nfrom slack_sdk import WebClient\nfrom slack_sdk.errors import SlackApiError\n\nfrom onyx.connectors.slack.connector import get_channel_messages\nfrom onyx.connectors.slack.models import ChannelType\nfrom onyx.connectors.slack.utils import make_paginated_slack_api_call\n\n\ndef _get_slack_channel_id(channel: ChannelType) -> str:\n if not (channel_id := channel.get(\"id\")):\n raise ValueError(\"Channel ID is missing\")\n return channel_id\n\n\ndef _get_non_general_channels(\n slack_client: WebClient,\n get_private: bool,\n get_public: bool,\n only_get_done: bool = False,\n) -> list[ChannelType]:\n channel_types = []\n if get_private:\n channel_types.append(\"private_channel\")\n if get_public:\n channel_types.append(\"public_channel\")\n\n conversations: list[dict[str, Any]] = []\n for result in make_paginated_slack_api_call(\n slack_client.conversations_list,\n exclude_archived=False,\n types=channel_types,\n ):\n conversations.extend(result[\"channels\"])\n\n filtered_conversations = []\n for conversation in conversations:\n if conversation.get(\"is_general\", False):\n continue\n if only_get_done and \"done\" not in conversation.get(\"name\", \"\"):\n continue\n filtered_conversations.append(conversation)\n return cast(list[ChannelType], filtered_conversations)\n\n\ndef _clear_slack_conversation_members(\n slack_client: WebClient,\n admin_user_id: str,\n channel: ChannelType,\n) -> None:\n channel_id = _get_slack_channel_id(channel)\n member_ids: list[str] = []\n for result in make_paginated_slack_api_call(\n slack_client.conversations_members,\n channel=channel_id,\n ):\n member_ids.extend(result[\"members\"])\n\n for member_id in member_ids:\n if member_id == admin_user_id:\n continue\n try:\n slack_client.conversations_kick(channel=channel_id, user=member_id)\n print(f\"Kicked member: {member_id}\")\n except Exception as e:\n if \"cant_kick_self\" in str(e):\n continue\n print(f\"Error kicking member: {e}\")\n print(member_id)\n try:\n slack_client.conversations_unarchive(channel=channel_id)\n channel[\"is_archived\"] = False\n except Exception:\n # Channel is already unarchived\n pass\n\n\ndef _add_slack_conversation_members(\n slack_client: WebClient, channel: ChannelType, member_ids: list[str]\n) -> None:\n channel_id = _get_slack_channel_id(channel)\n for user_id in member_ids:\n try:\n slack_client.conversations_invite(channel=channel_id, users=user_id)\n except Exception as e:\n if \"already_in_channel\" in str(e):\n continue\n print(f\"Error inviting member: {e}\")\n print(user_id)\n\n\ndef _delete_slack_conversation_messages(\n slack_client: WebClient,\n channel: ChannelType,\n message_to_delete: str | None = None,\n) -> None:\n \"\"\"deletes all messages from a channel if message_to_delete is None\"\"\"\n channel_id = _get_slack_channel_id(channel)\n for message_batch in get_channel_messages(slack_client, channel):\n for message in message_batch:\n if message_to_delete and message.get(\"text\") != message_to_delete:\n continue\n print(\" removing message: \", message.get(\"text\"))\n\n try:\n if not (ts := message.get(\"ts\")):\n raise ValueError(\"Message timestamp is missing\")\n slack_client.chat_delete(channel=channel_id, ts=ts)\n except Exception as e:\n print(f\"Error deleting message: {e}\")\n print(message)\n\n\ndef _build_slack_channel_from_name(\n slack_client: WebClient,\n admin_user_id: str,\n suffix: str,\n is_private: bool,\n channel: ChannelType | None,\n) -> ChannelType:\n base = \"public_channel\" if not is_private else \"private_channel\"\n channel_name = f\"{base}-{suffix}\"\n if channel:\n # If channel is provided, we rename it\n channel_id = _get_slack_channel_id(channel)\n channel_response = slack_client.conversations_rename(\n channel=channel_id,\n name=channel_name,\n )\n else:\n # Otherwise, we create a new channel\n channel_response = slack_client.conversations_create(\n name=channel_name,\n is_private=is_private,\n )\n\n try:\n slack_client.conversations_unarchive(channel=channel_response[\"channel\"][\"id\"])\n except Exception:\n # Channel is already unarchived\n pass\n try:\n slack_client.conversations_invite(\n channel=channel_response[\"channel\"][\"id\"],\n users=[admin_user_id],\n )\n except Exception:\n pass\n\n final_channel = channel_response[\"channel\"] if channel_response else {}\n return cast(ChannelType, final_channel)\n\n\nclass SlackManager:\n @staticmethod\n def get_slack_client(token: str) -> WebClient:\n return WebClient(token=token)\n\n @staticmethod\n def get_and_provision_available_slack_channels(\n slack_client: WebClient, admin_user_id: str\n ) -> tuple[ChannelType, ChannelType, str]:\n run_id = str(uuid4())\n public_channels = _get_non_general_channels(\n slack_client, get_private=False, get_public=True, only_get_done=True\n )\n\n first_available_channel = (\n None if len(public_channels) < 1 else public_channels[0]\n )\n public_channel = _build_slack_channel_from_name(\n slack_client=slack_client,\n admin_user_id=admin_user_id,\n suffix=run_id,\n is_private=False,\n channel=first_available_channel,\n )\n _delete_slack_conversation_messages(\n slack_client=slack_client, channel=public_channel\n )\n\n private_channels = _get_non_general_channels(\n slack_client, get_private=True, get_public=False, only_get_done=True\n )\n second_available_channel = (\n None if len(private_channels) < 1 else private_channels[0]\n )\n private_channel = _build_slack_channel_from_name(\n slack_client=slack_client,\n admin_user_id=admin_user_id,\n suffix=run_id,\n is_private=True,\n channel=second_available_channel,\n )\n _delete_slack_conversation_messages(\n slack_client=slack_client, channel=private_channel\n )\n\n return public_channel, private_channel, run_id\n\n @staticmethod\n def build_slack_user_email_id_map(slack_client: WebClient) -> dict[str, str]:\n users: list[dict[str, Any]] = []\n\n for users_results in make_paginated_slack_api_call(\n slack_client.users_list,\n ):\n users.extend(users_results.get(\"members\", []))\n\n user_email_id_map = {}\n for user in users:\n if not (email := user.get(\"profile\", {}).get(\"email\")):\n continue\n if not (user_id := user.get(\"id\")):\n raise ValueError(\"User ID is missing\")\n user_email_id_map[email] = user_id\n return user_email_id_map\n\n @staticmethod\n def set_channel_members(\n slack_client: WebClient,\n admin_user_id: str,\n channel: ChannelType,\n user_ids: list[str],\n ) -> None:\n _clear_slack_conversation_members(\n slack_client=slack_client,\n channel=channel,\n admin_user_id=admin_user_id,\n )\n _add_slack_conversation_members(\n slack_client=slack_client, channel=channel, member_ids=user_ids\n )\n\n @staticmethod\n def add_message_to_channel(\n slack_client: WebClient, channel: ChannelType, message: str\n ) -> None:\n channel_id = _get_slack_channel_id(channel)\n slack_client.chat_postMessage(\n channel=channel_id,\n text=message,\n )\n\n @staticmethod\n def remove_message_from_channel(\n slack_client: WebClient, channel: ChannelType, message: str\n ) -> None:\n _delete_slack_conversation_messages(\n slack_client=slack_client, channel=channel, message_to_delete=message\n )\n\n @staticmethod\n def cleanup_after_test(\n slack_client: WebClient,\n test_id: str,\n ) -> None:\n channel_types = [\"private_channel\", \"public_channel\"]\n channels: list[ChannelType] = []\n for result in make_paginated_slack_api_call(\n slack_client.conversations_list,\n exclude_archived=False,\n types=channel_types,\n ):\n channels.extend(result[\"channels\"])\n\n for channel in channels:\n if test_id not in channel.get(\"name\", \"\"):\n continue\n # \"done\" in the channel name indicates that this channel is free to be used for a new test\n new_name = f\"done_{str(uuid4())}\"\n try:\n slack_client.conversations_rename(channel=channel[\"id\"], name=new_name)\n except SlackApiError as e:\n print(f\"Error renaming channel {channel['id']}: {e}\")\n" }, { "path": "backend/tests/integration/connector_job_tests/slack/test_permission_sync.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.connectors.models import InputType\nfrom onyx.connectors.slack.models import ChannelType\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.document_search import (\n DocumentSearchManager,\n)\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\nfrom tests.integration.connector_job_tests.slack.conftest import SLACK_ADMIN_EMAIL\nfrom tests.integration.connector_job_tests.slack.conftest import SLACK_TEST_USER_1_EMAIL\nfrom tests.integration.connector_job_tests.slack.conftest import SLACK_TEST_USER_2_EMAIL\nfrom tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager\n\n\n# NOTE(rkuo): it isn't yet clear if the reason these were previously xfail'd\n# still exists. May need to xfail again if flaky (DAN-789)\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_slack_permission_sync(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture, # noqa: ARG001\n slack_perm_sync_test_setup: tuple[ChannelType, ChannelType],\n) -> None:\n public_channel, private_channel = slack_perm_sync_test_setup\n\n admin_user: DATestUser = UserManager.create(\n email=SLACK_ADMIN_EMAIL,\n )\n\n test_user_1: DATestUser = UserManager.create(\n email=SLACK_TEST_USER_1_EMAIL,\n )\n\n test_user_2: DATestUser = UserManager.create(\n email=SLACK_TEST_USER_2_EMAIL,\n )\n\n bot_token = os.environ[\"SLACK_BOT_TOKEN_TEST_SPACE\"]\n slack_client = SlackManager.get_slack_client(bot_token)\n email_id_map = SlackManager.build_slack_user_email_id_map(slack_client)\n admin_user_id = email_id_map[admin_user.email]\n\n LLMProviderManager.create(user_performing_action=admin_user)\n\n before = datetime.now(timezone.utc)\n credential: DATestCredential = CredentialManager.create(\n source=DocumentSource.SLACK,\n credential_json={\n \"slack_bot_token\": bot_token,\n },\n user_performing_action=admin_user,\n )\n connector: DATestConnector = ConnectorManager.create(\n name=\"Slack\",\n input_type=InputType.POLL,\n source=DocumentSource.SLACK,\n connector_specific_config={\n \"channels\": [public_channel[\"name\"], private_channel[\"name\"]],\n \"include_bot_messages\": True,\n },\n access_type=AccessType.SYNC,\n groups=[],\n user_performing_action=admin_user,\n )\n cc_pair: DATestCCPair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # Add test_user_1 and admin_user to the private channel\n desired_channel_members = [admin_user, test_user_1]\n SlackManager.set_channel_members(\n slack_client=slack_client,\n admin_user_id=admin_user_id,\n channel=private_channel,\n user_ids=[email_id_map[user.email] for user in desired_channel_members],\n )\n\n public_message = \"Steve's favorite number is 809752\"\n private_message = \"Sara's favorite number is 346794\"\n\n SlackManager.add_message_to_channel(\n slack_client=slack_client,\n channel=public_channel,\n message=public_message,\n )\n SlackManager.add_message_to_channel(\n slack_client=slack_client,\n channel=private_channel,\n message=private_message,\n )\n\n # Run indexing\n before = datetime.now(timezone.utc)\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # Run permission sync. Since initial_index_should_sync=True for Slack,\n # permissions were already set during indexing above — the explicit sync\n # should find no changes to apply.\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=0,\n user_performing_action=admin_user,\n should_wait_for_group_sync=False,\n should_wait_for_vespa_sync=False,\n )\n\n # Verify admin can see messages from both channels\n admin_docs = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=admin_user,\n )\n assert public_message in admin_docs\n assert private_message in admin_docs\n\n # Verify test_user_2 can only see public channel messages\n user_2_docs = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=test_user_2,\n )\n assert public_message in user_2_docs\n assert private_message not in user_2_docs\n\n # Verify test_user_1 can see both channels (member of private channel)\n user_1_docs = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=test_user_1,\n )\n assert public_message in user_1_docs\n assert private_message in user_1_docs\n\n # Remove test_user_1 from the private channel\n before = datetime.now(timezone.utc)\n desired_channel_members = [admin_user]\n SlackManager.set_channel_members(\n slack_client=slack_client,\n admin_user_id=admin_user_id,\n channel=private_channel,\n user_ids=[email_id_map[user.email] for user in desired_channel_members],\n )\n\n # Run permission sync\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n should_wait_for_group_sync=False,\n )\n\n # Verify test_user_1 can no longer see private channel after removal\n user_1_docs = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=test_user_1,\n )\n assert public_message in user_1_docs\n assert private_message not in user_1_docs\n\n\n# NOTE(rkuo): it isn't yet clear if the reason these were previously xfail'd\n# still exists. May need to xfail again if flaky (DAN-789)\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission tests are enterprise only\",\n)\ndef test_slack_group_permission_sync(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture, # noqa: ARG001\n slack_perm_sync_test_setup: tuple[ChannelType, ChannelType],\n) -> None:\n \"\"\"\n This test ensures that permission sync overrides onyx group access.\n \"\"\"\n public_channel, private_channel = slack_perm_sync_test_setup\n\n admin_user: DATestUser = UserManager.create(\n email=SLACK_ADMIN_EMAIL,\n )\n\n test_user_1: DATestUser = UserManager.create(\n email=SLACK_TEST_USER_1_EMAIL,\n )\n\n # Create a user group and adding the non-admin user to it\n user_group = UserGroupManager.create(\n name=\"test_group\",\n user_ids=[test_user_1.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group],\n user_performing_action=admin_user,\n )\n\n bot_token = os.environ[\"SLACK_BOT_TOKEN_TEST_SPACE\"]\n slack_client = SlackManager.get_slack_client(bot_token)\n email_id_map = SlackManager.build_slack_user_email_id_map(slack_client)\n admin_user_id = email_id_map[admin_user.email]\n\n LLMProviderManager.create(user_performing_action=admin_user)\n\n # Add only admin to the private channel\n SlackManager.set_channel_members(\n slack_client=slack_client,\n admin_user_id=admin_user_id,\n channel=private_channel,\n user_ids=[admin_user_id],\n )\n\n before = datetime.now(timezone.utc)\n credential = CredentialManager.create(\n source=DocumentSource.SLACK,\n credential_json={\n \"slack_bot_token\": bot_token,\n },\n user_performing_action=admin_user,\n )\n\n # Create connector with sync access and assign it to the user group\n connector = ConnectorManager.create(\n name=\"Slack\",\n input_type=InputType.POLL,\n source=DocumentSource.SLACK,\n connector_specific_config={\n \"channels\": [private_channel[\"name\"]],\n \"include_bot_messages\": True,\n },\n access_type=AccessType.SYNC,\n groups=[user_group.id],\n user_performing_action=admin_user,\n )\n\n cc_pair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n groups=[user_group.id],\n )\n\n # Add a test message to the private channel\n private_message = \"This is a secret message: 987654\"\n SlackManager.add_message_to_channel(\n slack_client=slack_client,\n channel=private_channel,\n message=private_message,\n )\n\n # Run indexing\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # Run permission sync. Since initial_index_should_sync=True for Slack,\n # permissions were already set during indexing — no changes expected.\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=0,\n user_performing_action=admin_user,\n should_wait_for_group_sync=False,\n should_wait_for_vespa_sync=False,\n )\n\n # Verify admin can see the message\n admin_docs = DocumentSearchManager.search_documents(\n query=\"secret message\",\n user_performing_action=admin_user,\n )\n assert private_message in admin_docs\n\n # Verify test_user_1 cannot see the message despite being in the group\n # (Slack permissions should take precedence)\n user_1_docs = DocumentSearchManager.search_documents(\n query=\"secret message\",\n user_performing_action=test_user_1,\n )\n assert private_message not in user_1_docs\n" }, { "path": "backend/tests/integration/connector_job_tests/slack/test_prune.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.connectors.models import InputType\nfrom onyx.connectors.slack.models import ChannelType\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.document_search import (\n DocumentSearchManager,\n)\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestConnector\nfrom tests.integration.common_utils.test_models import DATestCredential\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\nfrom tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager\n\n\n@pytest.mark.xfail(reason=\"flaky - see DAN-986 for details\", strict=False)\ndef test_slack_prune(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture, # noqa: ARG001\n slack_test_setup: tuple[ChannelType, ChannelType],\n) -> None:\n public_channel, private_channel = slack_test_setup\n\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(\n email=\"admin@example.com\",\n )\n\n # Creating a non-admin user\n test_user_1: DATestUser = UserManager.create(\n email=\"test_user_1@example.com\",\n )\n\n slack_client = SlackManager.get_slack_client(os.environ[\"SLACK_BOT_TOKEN\"])\n email_id_map = SlackManager.build_slack_user_email_id_map(slack_client)\n admin_user_id = email_id_map[admin_user.email]\n\n LLMProviderManager.create(user_performing_action=admin_user)\n\n before = datetime.now(timezone.utc)\n credential: DATestCredential = CredentialManager.create(\n source=DocumentSource.SLACK,\n credential_json={\n \"slack_bot_token\": os.environ[\"SLACK_BOT_TOKEN\"],\n },\n user_performing_action=admin_user,\n )\n connector: DATestConnector = ConnectorManager.create(\n name=\"Slack\",\n input_type=InputType.POLL,\n source=DocumentSource.SLACK,\n connector_specific_config={\n \"channels\": [public_channel[\"name\"], private_channel[\"name\"]],\n },\n access_type=AccessType.PUBLIC,\n groups=[],\n user_performing_action=admin_user,\n )\n cc_pair: DATestCCPair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # ----------------------SETUP INITIAL SLACK STATE--------------------------\n # Add test_user_1 and admin_user to the private channel\n desired_channel_members = [admin_user, test_user_1]\n SlackManager.set_channel_members(\n slack_client=slack_client,\n admin_user_id=admin_user_id,\n channel=private_channel,\n user_ids=[email_id_map[user.email] for user in desired_channel_members],\n )\n\n public_message = \"Steve's favorite number is 809752\"\n private_message = \"Sara's favorite number is 346794\"\n message_to_delete = \"Rebecca's favorite number is 753468\"\n\n SlackManager.add_message_to_channel(\n slack_client=slack_client,\n channel=public_channel,\n message=public_message,\n )\n SlackManager.add_message_to_channel(\n slack_client=slack_client,\n channel=private_channel,\n message=private_message,\n )\n SlackManager.add_message_to_channel(\n slack_client=slack_client,\n channel=private_channel,\n message=message_to_delete,\n )\n\n # Run indexing\n before = datetime.now(timezone.utc)\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # Run permission sync\n before = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=admin_user,\n )\n\n # ----------------------TEST THE SETUP--------------------------\n # Search as admin with access to both channels\n onyx_doc_message_strings = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=admin_user,\n )\n print(\n \"\\ntop_documents content before deleting for admin: \",\n onyx_doc_message_strings,\n )\n\n # Ensure admin user can see all messages\n assert public_message in onyx_doc_message_strings\n assert private_message in onyx_doc_message_strings\n assert message_to_delete in onyx_doc_message_strings\n\n # Search as test_user_1 with access to both channels\n onyx_doc_message_strings = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=test_user_1,\n )\n print(\n \"\\ntop_documents content before deleting for test_user_1: \",\n onyx_doc_message_strings,\n )\n\n # Ensure test_user_1 can see all messages\n assert public_message in onyx_doc_message_strings\n assert private_message in onyx_doc_message_strings\n assert message_to_delete in onyx_doc_message_strings\n\n # ----------------------MAKE THE CHANGES--------------------------\n # Delete messages\n print(\"\\nDeleting message: \", message_to_delete)\n SlackManager.remove_message_from_channel(\n slack_client=slack_client,\n channel=private_channel,\n message=message_to_delete,\n )\n\n # Prune the cc_pair\n now = datetime.now(timezone.utc)\n CCPairManager.prune(cc_pair, user_performing_action=admin_user)\n CCPairManager.wait_for_prune(cc_pair, now, user_performing_action=admin_user)\n\n # ----------------------------VERIFY THE CHANGES---------------------------\n # Ensure admin user can't see deleted messages\n # Search as admin user with access to only the public channel\n onyx_doc_message_strings = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=admin_user,\n )\n print(\n \"\\ntop_documents content after deleting for admin: \",\n onyx_doc_message_strings,\n )\n\n # Ensure admin can't see deleted messages\n assert public_message in onyx_doc_message_strings\n assert private_message in onyx_doc_message_strings\n assert message_to_delete not in onyx_doc_message_strings\n\n # Ensure test_user_1 can't see deleted messages\n # Search as test_user_1 with access to only the public channel\n onyx_doc_message_strings = DocumentSearchManager.search_documents(\n query=\"favorite number\",\n user_performing_action=test_user_1,\n )\n print(\n \"\\ntop_documents content after prune for test_user_1: \",\n onyx_doc_message_strings,\n )\n\n # Ensure test_user_1 can't see deleted messages\n assert public_message in onyx_doc_message_strings\n assert private_message in onyx_doc_message_strings\n assert message_to_delete not in onyx_doc_message_strings\n" }, { "path": "backend/tests/integration/mock_services/docker-compose.mock-it-services.yml", "content": "version: '3.8'\n\nservices:\n mock_connector_server:\n build:\n context: ./mock_connector_server\n dockerfile: Dockerfile\n ports:\n - \"8001:8001\"\n healthcheck:\n test: [\"CMD\", \"curl\", \"-f\", \"http://localhost:8001/health\"]\n interval: 10s\n timeout: 5s\n retries: 5\n networks:\n - onyx_default\nnetworks:\n onyx_default:\n name: onyx_default\n external: true\n" }, { "path": "backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_api_key.py", "content": "import sys\n\nimport uvicorn\nfrom fastapi import FastAPI\nfrom fastapi.responses import PlainTextResponse\nfrom fastmcp import FastMCP\nfrom fastmcp.server.auth import StaticTokenVerifier\n\n\ndef make_many_tools(mcp: FastMCP) -> None:\n def make_tool(i: int) -> None:\n @mcp.tool(name=f\"tool_{i}\", description=f\"Get secret value {i}\")\n def tool_name(name: str) -> str: # noqa: ARG001\n \"\"\"Get secret value.\"\"\"\n return f\"Secret value {200 - i}!\"\n\n for i in range(100):\n make_tool(i)\n\n\nif __name__ == \"__main__\":\n # Accept only these tokens (treat them like API keys) and require a scope\n if len(sys.argv) > 1:\n api_key = sys.argv[1]\n else:\n api_key = \"dev-api-key-123\"\n\n if len(sys.argv) > 2:\n port = int(sys.argv[2])\n else:\n port = 8001\n\n auth = StaticTokenVerifier(\n tokens={\n api_key: {\"client_id\": \"evan\", \"scopes\": [\"mcp:use\"]},\n },\n required_scopes=[\"mcp:use\"],\n )\n\n # Create FastMCP instance - it will handle /mcp path internally\n mcp = FastMCP(\"My HTTP MCP\", auth=auth)\n make_many_tools(mcp)\n\n # Get the MCP HTTP app (configured to serve at /mcp)\n mcp_app = mcp.http_app()\n\n # Create wrapper FastAPI app with the MCP app's lifespan\n app = FastAPI(title=\"MCP API Key Test Server\", lifespan=mcp_app.lifespan)\n\n # Health check (unprotected)\n @app.get(\"/healthz\")\n def health() -> PlainTextResponse:\n return PlainTextResponse(\"ok\")\n\n # Mount MCP app at root - it handles /mcp internally\n app.mount(\"/\", mcp_app)\n\n # Run the server\n uvicorn.run(app, host=\"0.0.0.0\", port=port)\n" }, { "path": "backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_google_oauth.py", "content": "\"\"\"\nMCP Test Server for Google OAuth Pass-Through Authentication\n\nThis server validates Google OAuth access tokens that are passed through from\nOnyx. When users log into Onyx with Google OAuth, their access token is stored\nand can be passed to MCP servers that require authentication.\n\nThis server validates those tokens by calling Google's tokeninfo endpoint.\n\nUsage:\n python run_mcp_server_google_oauth.py [port]\n\nEnvironment Variables:\n MCP_SERVER_HOST: Host to bind to (default: 127.0.0.1)\n MCP_SERVER_PUBLIC_HOST: Public hostname for the server\n MCP_SERVER_PUBLIC_URL: Public URL for the server (e.g., for proxied setups)\n\"\"\"\n\nimport os\nimport sys\nfrom typing import Any\n\nimport httpx\nimport uvicorn\nfrom fastapi import FastAPI\nfrom fastapi.responses import PlainTextResponse\nfrom fastmcp import FastMCP\nfrom fastmcp.server.auth import AccessToken\nfrom fastmcp.server.auth import TokenVerifier\nfrom fastmcp.server.dependencies import get_access_token\n\n# Google's tokeninfo endpoint for validating access tokens\nGOOGLE_TOKENINFO_URL = \"https://oauth2.googleapis.com/tokeninfo\"\n\n\nclass GoogleOAuthTokenVerifier(TokenVerifier):\n \"\"\"\n Token verifier that validates Google OAuth access tokens.\n\n Google access tokens are opaque tokens (not JWTs), so they need to be\n validated by calling Google's tokeninfo endpoint. This verifier makes\n an HTTP request to Google to validate the token and extract user info.\n\n This is useful for testing pass-through OAuth scenarios where Onyx\n forwards the user's Google OAuth token to an MCP server.\n \"\"\"\n\n def __init__(\n self,\n required_scopes: list[str] | None = None,\n base_url: str | None = None,\n ):\n \"\"\"\n Initialize the Google OAuth token verifier.\n\n Args:\n required_scopes: Optional list of scopes that must be present in the token.\n Google tokens have scopes like 'openid', 'email', 'profile'.\n base_url: URL of this resource server (for RFC 8707)\n \"\"\"\n super().__init__(\n base_url=base_url,\n required_scopes=required_scopes,\n )\n self._http_client: httpx.AsyncClient | None = None\n\n async def _get_http_client(self) -> httpx.AsyncClient:\n \"\"\"Get or create the HTTP client for token validation.\"\"\"\n if self._http_client is None or self._http_client.is_closed:\n self._http_client = httpx.AsyncClient(timeout=10.0)\n return self._http_client\n\n async def verify_token(self, token: str) -> AccessToken | None:\n \"\"\"\n Verify a Google OAuth access token by calling Google's tokeninfo endpoint.\n\n Args:\n token: The Google OAuth access token to validate\n\n Returns:\n AccessToken object if valid, None if invalid or expired\n \"\"\"\n try:\n client = await self._get_http_client()\n\n # Call Google's tokeninfo endpoint\n response = await client.get(\n GOOGLE_TOKENINFO_URL,\n params={\"access_token\": token},\n )\n\n if response.status_code != 200:\n # Token is invalid or expired\n return None\n\n token_info = response.json()\n\n # Check if token has an error (Google returns 200 with error field sometimes)\n if \"error\" in token_info:\n return None\n\n # Extract scopes from the token\n scopes_str = token_info.get(\"scope\", \"\")\n scopes = scopes_str.split() if scopes_str else []\n\n # Check required scopes if configured\n if self.required_scopes:\n token_scopes = set(scopes)\n required = set(self.required_scopes)\n if not required.issubset(token_scopes):\n return None\n\n # Extract client/user ID - prefer email for user identification\n client_id = (\n token_info.get(\"email\")\n or token_info.get(\"sub\")\n or token_info.get(\"user_id\")\n or \"unknown\"\n )\n\n # Extract expiration time\n expires_in = token_info.get(\"expires_in\")\n expires_at = None\n if expires_in:\n import time\n\n expires_at = int(time.time()) + int(expires_in)\n\n return AccessToken(\n token=token,\n client_id=client_id,\n scopes=scopes,\n expires_at=expires_at,\n claims=token_info,\n )\n\n except httpx.HTTPError:\n # Network error or timeout\n return None\n except Exception:\n # Any other error during validation\n return None\n\n async def close(self) -> None:\n \"\"\"Close the HTTP client.\"\"\"\n if self._http_client and not self._http_client.is_closed:\n await self._http_client.aclose()\n\n\ndef make_tools(mcp: FastMCP) -> None:\n \"\"\"Create test tools for the MCP server.\"\"\"\n\n @mcp.tool(name=\"echo\", description=\"Echo back the input message\")\n def echo(message: str) -> str:\n \"\"\"Echo the message back to the caller.\"\"\"\n return f\"You said: {message}\"\n\n @mcp.tool(name=\"get_secret\", description=\"Get a secret value (requires auth)\")\n def get_secret(secret_name: str) -> str:\n \"\"\"Get a secret value. This proves the token was validated.\"\"\"\n return f\"Secret value for '{secret_name}': super-secret-value-12345\"\n\n @mcp.tool(name=\"whoami\", description=\"Get information about the authenticated user\")\n async def whoami() -> dict[str, Any]:\n \"\"\"Get information about the authenticated user from their Google token.\"\"\"\n tok = get_access_token()\n if not tok:\n return {\"error\": \"Not authenticated\"}\n\n return {\n \"client_id\": tok.client_id,\n \"scopes\": tok.scopes,\n \"email\": tok.claims.get(\"email\"),\n \"email_verified\": tok.claims.get(\"email_verified\"),\n \"expires_in\": tok.claims.get(\"expires_in\"),\n \"access_type\": tok.claims.get(\"access_type\"),\n }\n\n for i in range(5):\n\n @mcp.tool(name=f\"oauth_tool_{i}\", description=f\"Test tool number {i}\")\n def numbered_tool(name: str, _i: int = i) -> str:\n \"\"\"A numbered test tool.\"\"\"\n return f\"Tool {_i} says hello to {name}!\"\n\n\nif __name__ == \"__main__\":\n port = int(sys.argv[1] if len(sys.argv) > 1 else \"8006\")\n\n # Get configuration from environment\n bind_host = os.getenv(\"MCP_SERVER_HOST\", \"127.0.0.1\")\n public_host = os.getenv(\"MCP_SERVER_PUBLIC_HOST\", bind_host)\n public_url = os.getenv(\"MCP_SERVER_PUBLIC_URL\")\n\n # Optional: require specific scopes (Google tokens have scopes like 'email', 'profile')\n # Leave empty to accept any valid Google token\n required_scopes_str = os.getenv(\"MCP_GOOGLE_REQUIRED_SCOPES\", \"\")\n required_scopes = (\n required_scopes_str.split(\",\") if required_scopes_str.strip() else None\n )\n\n print(f\"Starting Google OAuth MCP Test Server on port {port}\")\n print(f\"Bind host: {bind_host}\")\n print(f\"Public host: {public_host}\")\n if public_url:\n print(f\"Public URL: {public_url}\")\n if required_scopes:\n print(f\"Required scopes: {required_scopes}\")\n else:\n print(\"No specific scopes required - any valid Google token accepted\")\n\n # Create the auth verifier\n auth = GoogleOAuthTokenVerifier(required_scopes=required_scopes)\n\n # Create FastMCP instance with auth\n mcp = FastMCP(\"Google OAuth Test MCP Server\", auth=auth)\n make_tools(mcp)\n\n # Get the MCP HTTP app\n mcp_app = mcp.http_app()\n\n # Create wrapper FastAPI app\n app = FastAPI(\n title=\"MCP Google OAuth Test Server\",\n description=\"MCP server that authenticates using Google OAuth tokens passed through from Onyx\",\n lifespan=mcp_app.lifespan,\n )\n\n # Health check (unprotected)\n @app.get(\"/healthz\")\n def health() -> PlainTextResponse:\n return PlainTextResponse(\"ok\")\n\n # Info endpoint (unprotected) - useful for debugging\n @app.get(\"/info\")\n def info() -> dict[str, Any]:\n return {\n \"server\": \"Google OAuth MCP Test Server\",\n \"auth_type\": \"google_oauth_pass_through\",\n \"description\": \"Validates Google OAuth tokens passed from Onyx\",\n \"tokeninfo_endpoint\": GOOGLE_TOKENINFO_URL,\n \"required_scopes\": required_scopes,\n }\n\n # Mount MCP app at root\n app.mount(\"/\", mcp_app)\n\n # Run the server\n uvicorn.run(app, host=bind_host, port=port)\n" }, { "path": "backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_no_auth.py", "content": "import os\nimport sys\n\nfrom fastmcp import FastMCP\n\nmcp = FastMCP(\"My HTTP MCP\")\n\n\n@mcp.tool\ndef hello(name: str) -> str:\n \"\"\"Say hi.\"\"\"\n return f\"Hello, {name}!\"\n\n\ndef make_many_tools() -> None:\n def make_tool(i: int) -> None:\n @mcp.tool(name=f\"tool_{i}\", description=f\"Get secret value {i}\")\n def tool_name(name: str) -> str: # noqa: ARG001\n \"\"\"Get secret value.\"\"\"\n return f\"Secret value {100 - i}!\"\n\n for i in range(100):\n make_tool(i)\n\n\nif __name__ == \"__main__\":\n # Get port from command-line argument first (passed by test)\n port_from_arg = int(sys.argv[1]) if len(sys.argv) > 1 else None\n # Streamable HTTP transport (recommended)\n make_many_tools()\n host = os.getenv(\"MCP_SERVER_BIND_HOST\", \"0.0.0.0\")\n # Use MOCK_MCP_SERVER_PORT to avoid conflicts with the real Onyx MCP server port (8090)\n # Priority: command-line arg > MOCK_MCP_SERVER_PORT > MCP_SERVER_PORT > default 8000\n if port_from_arg is not None:\n port = port_from_arg\n else:\n port = int(\n os.getenv(\"MOCK_MCP_SERVER_PORT\") or os.getenv(\"MCP_SERVER_PORT\") or \"8000\"\n )\n path = os.getenv(\"MCP_SERVER_PATH\", \"/mcp\")\n mcp.run(transport=\"http\", host=host, port=port, path=path)\n" }, { "path": "backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_oauth.py", "content": "import os\nfrom collections.abc import Awaitable\nfrom collections.abc import Callable\nfrom collections.abc import Iterable\nfrom typing import Any\nfrom urllib.parse import urlsplit\nfrom urllib.parse import urlunsplit\n\nimport uvicorn\nfrom fastapi import FastAPI\nfrom fastapi import Request\nfrom fastapi.responses import JSONResponse\nfrom fastapi.responses import PlainTextResponse\nfrom fastapi.responses import Response\nfrom fastmcp import FastMCP\nfrom fastmcp.server.auth.providers.jwt import JWTVerifier\nfrom fastmcp.server.dependencies import get_access_token\nfrom starlette.middleware.base import BaseHTTPMiddleware\n\n# uncomment for debug logs\n# logging.basicConfig(level=logging.DEBUG)\n\n\"\"\"\nSetup Okta:\n1. Create an authorization Server (Admin Console → Security →\nAPI → Authorization Servers), and get the Issuer, JWKS uri,\naudience (i.e. api://mcp). Add the mcp:use scope.\nGrant types should be Authorization Code and Refresh Token.\npolicy should allow your client (or All Clients) to grant oidc default scopes + mcp:use\nWARNING: due to the order of discovery urls, you actually need to use the default authorization server\nuntil Okta updates where their discovery urls are or the client library stops trying\nto go to /.well-known/oauth-authorization-server before trying the fallback\n\n2. Create a client (Admin Console → Applications → Create App Integration)\nEnable authorization code and store the client id and secret.\n\"\"\"\n\n\ndef make_many_tools(mcp: FastMCP) -> None:\n def make_tool(i: int) -> None:\n @mcp.tool(name=f\"tool_{i}\", description=f\"Get secret value {i}\")\n def tool_name(name: str) -> str: # noqa: ARG001\n \"\"\"Get secret value.\"\"\"\n return f\"Secret value {500 - i}!\"\n\n for i in range(100):\n make_tool(i)\n\n @mcp.tool\n async def whoami() -> dict[str, Any]:\n tok = get_access_token() # None if unauthenticated\n return {\n \"client_id\": tok.client_id if tok else None,\n \"scopes\": tok.scopes if tok else [],\n \"claims\": tok.claims if tok else {},\n }\n\n\n# ---------- FASTAPI APP ----------\n\n\ndef init_app(\n app: FastAPI,\n mcp_resource_url: str,\n authorization_servers: list[str],\n scopes_supported: list[str],\n) -> None:\n # 1) Protected Resource Metadata (RFC 9728) at well-known URL.\n # We accept both with and without the trailing resource suffix to be lenient in dev.\n @app.get(\"/.well-known/oauth-protected-resource\")\n @app.get(\"/.well-known/oauth-protected-resource/{_suffix:path}\")\n def oauth_protected_resource(_suffix: str = \"\") -> JSONResponse:\n \"\"\"\n Return PRM document. The 'resource' MUST equal the MCP resource identifier (the URL clients use),\n and should be validated by clients per RFC 9728 §3.3.\n \"\"\"\n return JSONResponse(\n {\n \"resource\": mcp_resource_url,\n \"authorization_servers\": authorization_servers,\n \"bearer_methods_supported\": [\"header\"],\n \"scopes_supported\": scopes_supported,\n # (Optional extras: jwks_uri, resource_signing_alg_values_supported, etc.)\n }\n )\n\n # Health check (unprotected)\n @app.get(\"/healthz\")\n def health() -> PlainTextResponse:\n return PlainTextResponse(\"ok\")\n\n\ndef metadata_url_for_resource(resource_url: str) -> str:\n \"\"\"\n RFC 9728: insert '/.well-known/oauth-protected-resource' between host and path.\n If the resource has a path (e.g., '/mcp'), append it after the well-known suffix.\n \"\"\"\n u = urlsplit(resource_url)\n path = u.path.lstrip(\"/\")\n suffix = \"/.well-known/oauth-protected-resource\"\n if path:\n suffix += f\"/{path}\"\n return urlunsplit((u.scheme, u.netloc, suffix, \"\", \"\"))\n\n\nPRM_URL = \"replace me\"\n\n\n# 2) Middleware that ensures 401s include a proper WWW-Authenticate challenge\n# pointing clients to our PRM URL (RFC 9728 §5.1), and includes RFC 6750 error info.\nclass WWWAuthenticateMiddleware(BaseHTTPMiddleware):\n def __init__(self, app: FastAPI, protected_prefixes: Iterable[str]) -> None:\n super().__init__(app)\n self.protected_prefixes = tuple(protected_prefixes)\n\n async def dispatch(\n self, request: Request, call_next: Callable[[Request], Awaitable[Response]]\n ) -> Response:\n # Only guard MCP endpoints (both Streamable HTTP and SSE)\n if not request.url.path.startswith(self.protected_prefixes):\n return await call_next(request)\n\n # Let FastMCP/verifier run first\n response = await call_next(request)\n\n # If unauthenticated or invalid token, attach RFC-compliant challenge header\n if response.status_code == 401:\n # RFC 9728: include resource_metadata param pointing to PRM URL.\n # RFC 6750: include error + error_description when appropriate.\n challenge = f'Bearer resource_metadata=\"{PRM_URL}\", error=\"invalid_token\", error_description=\"Authentication required\"'\n # Don't clobber if already present; append or set.\n if \"www-authenticate\" in response.headers:\n response.headers[\"www-authenticate\"] += \", \" + challenge\n else:\n response.headers[\"www-authenticate\"] = challenge\n # Helpful cache headers\n response.headers.setdefault(\"cache-control\", \"no-store\")\n response.headers.setdefault(\"pragma\", \"no-cache\")\n return response\n\n\nif __name__ == \"__main__\":\n import sys\n\n port = int(sys.argv[1] if len(sys.argv) > 1 else \"8004\")\n\n audience = os.getenv(\"MCP_OAUTH_AUDIENCE\", \"api://mcp\")\n issuer = os.getenv(\n \"MCP_OAUTH_ISSUER\",\n \"https://test-domain.okta.com/oauth2/default?well_known_override=https://test-domain.okta.com/oauth2//.well-known/oauth-authorization-server\",\n ) # NOTE: the mcp client library currently tries the root discovery url before\n # falling back to the one actually used by Okta. Our client code lets you specify this well_known_override\n # for Okta and other Idps that use these discovery urls.\n\n # issuer = os.getenv(\"MCP_OAUTH_ISSUER\", \"https://test-domain.okta.com/.well-known/oauth-authorization-server?issuer=https://test-domain.okta.com/oauth2/\")\n jwks_uri = os.getenv(\n \"MCP_OAUTH_JWKS_URI\", \"https://test-domain.okta.com/oauth2/default/v1/keys\"\n )\n required_scopes = os.getenv(\"MCP_OAUTH_REQUIRED_SCOPES\", \"mcp:use\")\n print(f\"Required scopes: {required_scopes}\")\n print(f\"Audience: {audience}\")\n print(f\"Issuer: {issuer}\")\n print(f\"JWKS URI: {jwks_uri}\")\n\n verifier = JWTVerifier(\n issuer=issuer.split(\"?\")[0], # ignore auth url override if present\n audience=audience, # exactly what you set on the AS\n jwks_uri=jwks_uri,\n required_scopes=required_scopes.split(\n \",\"\n ), # must be present in the token's `scp`\n )\n\n bind_host = os.getenv(\"MCP_SERVER_HOST\", \"127.0.0.1\")\n public_host = os.getenv(\"MCP_SERVER_PUBLIC_HOST\", bind_host)\n public_url = os.getenv(\"MCP_SERVER_PUBLIC_URL\")\n\n mcp = FastMCP(\"My HTTP MCP\", auth=verifier)\n make_many_tools(mcp)\n mcp_app = mcp.http_app()\n\n app = FastAPI(title=\"MCP over HTTP/SSE with OAuth\", lifespan=mcp_app.lifespan)\n\n if public_url:\n normalized_public_url = public_url.rstrip(\"/\")\n if not normalized_public_url.endswith(\"/mcp\"):\n normalized_public_url = f\"{normalized_public_url}/mcp\"\n mcp_resource_url = f\"{normalized_public_url}/\"\n else:\n mcp_resource_url = f\"http://{public_host}:{port}/mcp/\"\n authorization_servers = [issuer]\n scopes_supported = [\"mcp:use\"]\n\n init_app(app, mcp_resource_url, authorization_servers, scopes_supported)\n PRM_URL = metadata_url_for_resource(mcp_resource_url)\n print(f\"PRM URL: {PRM_URL}\")\n print(f\"MCP Resource URL: {mcp_resource_url}\")\n print(f\"Authorization Servers: {authorization_servers}\")\n print(f\"Scopes Supported: {scopes_supported}\")\n\n # Apply middleware at the parent app so it wraps mounted sub-apps too\n app.add_middleware(WWWAuthenticateMiddleware, protected_prefixes=[\"/mcp\", \"/sse\"])\n\n # 3) Mount MCP apps\n # Streamable HTTP transport (recommended for modern MCP clients)\n app.mount(\"/\", mcp_app)\n # SSE transport (some clients still use this)\n # app.mount(\"/sse\", mcp.sse_app()) # TODO: v2\n\n uvicorn.run(app, host=bind_host, port=port)\n" }, { "path": "backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py", "content": "import sys\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\nfrom typing import Dict\nfrom typing import Optional\n\nimport bcrypt\nfrom fastmcp import FastMCP\nfrom fastmcp.server.auth.auth import AccessToken\nfrom fastmcp.server.auth.auth import TokenVerifier\nfrom fastmcp.server.dependencies import get_access_token\n\n# pip install fastmcp bcrypt\n\n\n# ---- pretend database --------------------------------------------------------\n# Keys look like: \"mcp_live__\"\ndef _hash(secret: str) -> bytes:\n return bcrypt.hashpw(secret.encode(), bcrypt.gensalt(rounds=12))\n\n\nAPI_KEY_RECORDS: Dict[str, Dict[str, Any]] = {\n # key_id -> record\n \"kid_alice_001\": {\n \"user_id\": \"alice\",\n \"hashed_secret\": _hash(\"S3cr3tAlice\"),\n \"scopes\": [\"mcp:use\"],\n \"revoked\": False,\n \"expires_at\": None, # or datetime(...)\n \"metadata\": {\"plan\": \"pro\"},\n },\n \"kid_bob_001\": {\n \"user_id\": \"bob\",\n \"hashed_secret\": _hash(\"S3cr3tBob\"),\n \"scopes\": [\"mcp:use\"],\n \"revoked\": False,\n \"expires_at\": None,\n \"metadata\": {\"plan\": \"free\"},\n },\n}\n\n# These are inferrable from the file anyways, no need to obfuscate.\n# use them to test your auth with this server\n#\n# mcp_live-kid_alice_001-S3cr3tAlice\n# mcp_live-kid_bob_001-S3cr3tBob\n\n\n# ---- verifier ---------------------------------------------------------------\nclass ApiKeyVerifier(TokenVerifier):\n \"\"\"\n Accepts API keys in Authorization: Bearer mcp_live__\n Looks up in storage, bcrypt-verifies , returns AccessToken.\n \"\"\"\n\n def __init__(self, api_key_dict: dict[str, Any]):\n super().__init__()\n self.api_key_dict = api_key_dict\n\n async def verify_token(self, token: str) -> Optional[AccessToken]:\n # print(f\"Verifying token: {token}\")\n try:\n prefix, key_id, secret = token.split(\"-\")\n # print(f\"Prefix: {prefix}, Key ID: {key_id}, Secret: {secret}\")\n if prefix not in (\"mcp_live\", \"mcp_test\"):\n return None\n except ValueError:\n return None\n\n rec = self.api_key_dict.get(key_id)\n if not rec or rec.get(\"revoked\"):\n return None\n if rec.get(\"expires_at\") and rec[\"expires_at\"] < datetime.now(timezone.utc):\n return None\n\n # constant-time bcrypt verification\n if not bcrypt.checkpw(secret.encode(), rec[\"hashed_secret\"]):\n return None\n\n # Build an AccessToken with claims FastMCP can pass to your tools\n return AccessToken(\n token=token,\n client_id=rec[\"user_id\"],\n scopes=rec.get(\"scopes\", []),\n expires_at=rec.get(\"expires_at\"),\n resource=None,\n claims={\"key_id\": key_id, **rec.get(\"metadata\", {})},\n )\n\n\n# ---- server -----------------------------------------------------------------\n\n\ndef make_many_tools(mcp: FastMCP) -> None:\n def make_tool(i: int) -> None:\n @mcp.tool(name=f\"tool_{i}\", description=f\"Get secret value {i}\")\n def tool_name(name: str) -> str: # noqa: ARG001\n \"\"\"Get secret value.\"\"\"\n return f\"Secret value {400 - i}!\"\n\n for i in range(100):\n make_tool(i)\n\n\nif __name__ == \"__main__\":\n if len(sys.argv) > 1:\n port = int(sys.argv[1])\n else:\n port = 8003\n\n mcp = FastMCP(\"My HTTP MCP\", auth=ApiKeyVerifier(API_KEY_RECORDS))\n\n @mcp.tool\n def whoami() -> dict:\n \"\"\"Return authenticated identity info (for demo).\"\"\"\n # FastMCP exposes the verified AccessToken to tools; see docs for helpers\n tok = get_access_token()\n return {\n \"user\": tok.client_id if tok else None,\n \"scopes\": tok.scopes if tok else [],\n }\n\n make_many_tools(mcp)\n mcp.run(transport=\"http\", host=\"127.0.0.1\", port=port, path=\"/mcp\")\n" }, { "path": "backend/tests/integration/mock_services/mock_connector_server/Dockerfile", "content": "FROM python:3.11.7-slim-bookworm\n\nWORKDIR /app\n\nRUN pip install --no-cache-dir \"pydantic-core>=2.28.0\" fastapi uvicorn\n\nCOPY ./main.py /app/main.py\n\nCMD [\"uvicorn\", \"main:app\", \"--host\", \"0.0.0.0\", \"--port\", \"8001\"]\n" }, { "path": "backend/tests/integration/mock_services/mock_connector_server/main.py", "content": "from fastapi import FastAPI\nfrom fastapi import HTTPException\nfrom pydantic import BaseModel\nfrom pydantic import Field\n\n# We would like to import these, but it makes building this so much harder/slower\n# from onyx.connectors.mock_connector.connector import SingleConnectorYield\n# from onyx.connectors.models import ConnectorCheckpoint\n\napp = FastAPI()\n\n\n# Global state to store connector behavior configuration\nclass ConnectorBehavior(BaseModel):\n connector_yields: list[dict] = Field(\n default_factory=list\n ) # really list[SingleConnectorYield]\n called_with_checkpoints: list[dict] = Field(\n default_factory=list\n ) # really list[ConnectorCheckpoint]\n\n\ncurrent_behavior: ConnectorBehavior = ConnectorBehavior()\n\n\n@app.post(\"/set-behavior\")\nasync def set_behavior(behavior: list[dict]) -> None:\n \"\"\"Set the behavior for the next connector run\"\"\"\n global current_behavior\n current_behavior = ConnectorBehavior(connector_yields=behavior)\n\n\n@app.get(\"/get-documents\")\nasync def get_documents() -> list[dict]:\n \"\"\"Get the next batch of documents and update the checkpoint\"\"\"\n global current_behavior\n\n if not current_behavior.connector_yields:\n raise HTTPException(\n status_code=400, detail=\"No documents or failures configured\"\n )\n\n connector_yields = current_behavior.connector_yields\n\n # Clear the current behavior after returning it\n current_behavior = ConnectorBehavior()\n\n return connector_yields\n\n\n@app.post(\"/add-checkpoint\")\nasync def add_checkpoint(checkpoint: dict) -> None:\n \"\"\"Add a checkpoint to the list of checkpoints. Called by the MockConnector.\"\"\"\n global current_behavior\n current_behavior.called_with_checkpoints.append(checkpoint)\n\n\n@app.get(\"/get-checkpoints\")\nasync def get_checkpoints() -> list[dict]:\n \"\"\"Get the list of checkpoints. Used by the test to verify the\n proper checkpoint ordering.\"\"\"\n global current_behavior\n return current_behavior.called_with_checkpoints\n\n\n@app.post(\"/reset\")\nasync def reset() -> None:\n \"\"\"Reset the connector behavior to default\"\"\"\n global current_behavior\n current_behavior = ConnectorBehavior()\n\n\n@app.get(\"/health\")\nasync def health_check() -> dict[str, str]:\n \"\"\"Health check endpoint\"\"\"\n return {\"status\": \"healthy\"}\n" }, { "path": "backend/tests/integration/multitenant_tests/discord_bot/test_discord_bot_multitenant.py", "content": "\"\"\"Multi-tenant isolation tests for Discord bot.\n\nThese tests ensure tenant isolation and prevent data leakage between tenants.\nTests follow the multi-tenant integration test pattern using API requests.\n\"\"\"\n\nfrom unittest.mock import patch\nfrom uuid import uuid4\n\nimport pytest\nimport requests\n\nfrom onyx.configs.constants import AuthType\nfrom onyx.db.discord_bot import get_guild_config_by_registration_key\nfrom onyx.db.discord_bot import register_guild\nfrom onyx.db.engine.sql_engine import get_session_with_tenant\nfrom onyx.db.models import UserRole\nfrom onyx.onyxbot.discord.cache import DiscordCacheManager\nfrom onyx.server.manage.discord_bot.utils import generate_discord_registration_key\nfrom onyx.server.manage.discord_bot.utils import parse_discord_registration_key\nfrom onyx.server.manage.discord_bot.utils import REGISTRATION_KEY_PREFIX\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass TestBotConfigIsolationCloudMode:\n \"\"\"Tests for bot config isolation in cloud mode.\"\"\"\n\n def test_cannot_create_bot_config_in_cloud_mode(self) -> None:\n \"\"\"Bot config creation is blocked in cloud mode.\"\"\"\n with patch(\"onyx.configs.app_configs.AUTH_TYPE\", AuthType.CLOUD):\n from fastapi import HTTPException\n\n from onyx.server.manage.discord_bot.api import _check_bot_config_api_access\n\n with pytest.raises(HTTPException) as exc_info:\n _check_bot_config_api_access()\n\n assert exc_info.value.status_code == 403\n assert \"Cloud\" in str(exc_info.value.detail)\n\n def test_bot_token_from_env_only_in_cloud(self) -> None:\n \"\"\"Bot token comes from env var in cloud mode, ignores DB.\"\"\"\n from onyx.onyxbot.discord.utils import get_bot_token\n\n with (\n patch(\"onyx.onyxbot.discord.utils.DISCORD_BOT_TOKEN\", \"env_token\"),\n patch(\"onyx.onyxbot.discord.utils.AUTH_TYPE\", AuthType.CLOUD),\n ):\n result = get_bot_token()\n\n assert result == \"env_token\"\n\n\nclass TestGuildRegistrationIsolation:\n \"\"\"Tests for guild registration isolation between tenants.\"\"\"\n\n def test_guild_can_only_register_to_one_tenant(self) -> None:\n \"\"\"Guild registered to tenant 1 cannot be registered to tenant 2.\"\"\"\n cache = DiscordCacheManager()\n\n # Register guild to tenant 1\n cache._guild_tenants[123456789] = \"tenant1\"\n\n # Check if guild is already registered\n existing = cache.get_tenant(123456789)\n\n assert existing is not None\n assert existing == \"tenant1\"\n\n def test_registration_key_tenant_mismatch(self) -> None:\n \"\"\"Key created in tenant 1 cannot be used in tenant 2 context.\"\"\"\n key = generate_discord_registration_key(\"tenant1\")\n\n # Parse the key to get tenant\n parsed_tenant = parse_discord_registration_key(key)\n\n assert parsed_tenant == \"tenant1\"\n assert parsed_tenant != \"tenant2\"\n\n def test_registration_key_encodes_correct_tenant(self) -> None:\n \"\"\"Key format discord_. encodes correct tenant.\"\"\"\n tenant_id = \"my_tenant_123\"\n key = generate_discord_registration_key(tenant_id)\n\n assert key.startswith(REGISTRATION_KEY_PREFIX)\n assert \"my_tenant_123\" in key or \"my%5Ftenant%5F123\" in key\n\n parsed = parse_discord_registration_key(key)\n assert parsed == tenant_id\n\n\nclass TestGuildDataIsolation:\n \"\"\"Tests for guild data isolation between tenants via API.\"\"\"\n\n def test_tenant_cannot_see_other_tenant_guilds(\n self,\n reset_multitenant: None, # noqa: ARG002\n ) -> None:\n \"\"\"Guilds created in tenant 1 are not visible from tenant 2.\n\n Creates guilds via API in tenant 1, then queries from tenant 2\n context to verify the guilds are not visible.\n \"\"\"\n unique = uuid4().hex\n\n # Create admin user for tenant 1\n admin_user1: DATestUser = UserManager.create(\n email=f\"discord_admin1_{unique}@example.com\",\n )\n assert UserManager.is_role(admin_user1, UserRole.ADMIN)\n\n # Create admin user for tenant 2\n admin_user2: DATestUser = UserManager.create(\n email=f\"discord_admin2_{unique}@example.com\",\n )\n assert UserManager.is_role(admin_user2, UserRole.ADMIN)\n\n # Create a guild registration key in tenant 1\n response1 = requests.post(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user1.headers,\n )\n\n # If Discord bot feature is not enabled, skip the test\n if response1.status_code == 404:\n pytest.skip(\"Discord bot feature not enabled\")\n\n assert response1.ok, f\"Failed to create guild in tenant 1: {response1.text}\"\n guild1_data = response1.json()\n guild1_id = guild1_data[\"id\"]\n\n try:\n # List guilds from tenant 1 - should see the guild\n list_response1 = requests.get(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user1.headers,\n )\n assert list_response1.ok\n tenant1_guilds = list_response1.json()\n tenant1_guild_ids = [g[\"id\"] for g in tenant1_guilds]\n assert guild1_id in tenant1_guild_ids\n\n # List guilds from tenant 2 - should NOT see tenant 1's guild\n list_response2 = requests.get(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user2.headers,\n )\n assert list_response2.ok\n tenant2_guilds = list_response2.json()\n tenant2_guild_ids = [g[\"id\"] for g in tenant2_guilds]\n assert guild1_id not in tenant2_guild_ids\n\n finally:\n # Cleanup - delete guild from tenant 1\n requests.delete(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds/{guild1_id}\",\n headers=admin_user1.headers,\n )\n\n def test_guild_list_returns_only_own_tenant(\n self,\n reset_multitenant: None, # noqa: ARG002\n ) -> None:\n \"\"\"List guilds returns exactly the guilds for that tenant.\n\n Creates 1 guild in each tenant, registers them with different data,\n and verifies each tenant only sees their own guild.\n \"\"\"\n unique = uuid4().hex\n\n # Create admin users for two tenants\n admin_user1: DATestUser = UserManager.create(\n email=f\"discord_list1_{unique}@example.com\",\n )\n admin_user2: DATestUser = UserManager.create(\n email=f\"discord_list2_{unique}@example.com\",\n )\n\n # Create 1 guild in tenant 1\n response1 = requests.post(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user1.headers,\n )\n if response1.status_code == 404:\n pytest.skip(\"Discord bot feature not enabled\")\n assert response1.ok, f\"Failed to create guild in tenant 1: {response1.text}\"\n guild1_data = response1.json()\n guild1_id = guild1_data[\"id\"]\n registration_key1 = guild1_data[\"registration_key\"]\n tenant1_id = parse_discord_registration_key(registration_key1)\n assert (\n tenant1_id is not None\n ), \"Failed to parse tenant ID from registration key 1\"\n\n # Create 1 guild in tenant 2\n response2 = requests.post(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user2.headers,\n )\n assert response2.ok, f\"Failed to create guild in tenant 2: {response2.text}\"\n guild2_data = response2.json()\n guild2_id = guild2_data[\"id\"]\n registration_key2 = guild2_data[\"registration_key\"]\n tenant2_id = parse_discord_registration_key(registration_key2)\n assert (\n tenant2_id is not None\n ), \"Failed to parse tenant ID from registration key 2\"\n\n # Verify tenant IDs are different\n assert (\n tenant1_id != tenant2_id\n ), \"Tenant 1 and tenant 2 should have different tenant IDs\"\n\n # Register guild 1 with tenant 1's context - populate with different data\n with get_session_with_tenant(tenant_id=tenant1_id) as db_session:\n config1 = get_guild_config_by_registration_key(\n db_session, registration_key1\n )\n assert config1 is not None, \"Guild config 1 should exist\"\n register_guild(\n db_session=db_session,\n config=config1,\n guild_id=111111111111111111, # Different Discord guild ID\n guild_name=\"Tenant 1 Server\", # Different guild name\n )\n db_session.commit()\n\n # Register guild 2 with tenant 2's context - populate with different data\n with get_session_with_tenant(tenant_id=tenant2_id) as db_session:\n config2 = get_guild_config_by_registration_key(\n db_session, registration_key2\n )\n assert config2 is not None, \"Guild config 2 should exist\"\n register_guild(\n db_session=db_session,\n config=config2,\n guild_id=222222222222222222, # Different Discord guild ID\n guild_name=\"Tenant 2 Server\", # Different guild name\n )\n db_session.commit()\n\n try:\n # Verify tenant 1 sees only their guild\n list_response1 = requests.get(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user1.headers,\n )\n assert list_response1.ok\n tenant1_guilds = list_response1.json()\n\n # Tenant 1 should see exactly 1 guild\n assert (\n len(tenant1_guilds) == 1\n ), f\"Tenant 1 should see 1 guild, got {len(tenant1_guilds)}\"\n\n # Verify tenant 1's guild has the correct data\n tenant1_guild = tenant1_guilds[0]\n assert (\n tenant1_guild[\"id\"] == guild1_id\n ), \"Tenant 1 should see their own guild\"\n assert (\n tenant1_guild[\"guild_id\"] == 111111111111111111\n ), f\"Tenant 1's guild should have guild_id 111111111111111111, got {tenant1_guild['guild_id']}\"\n assert (\n tenant1_guild[\"guild_name\"] == \"Tenant 1 Server\"\n ), f\"Tenant 1's guild should have name 'Tenant 1 Server', got {tenant1_guild['guild_name']}\"\n assert (\n tenant1_guild[\"registered_at\"] is not None\n ), \"Tenant 1's guild should be registered\"\n\n # Tenant 1 should NOT see tenant 2's guild\n assert (\n tenant1_guild[\"guild_id\"] != 222222222222222222\n ), \"Tenant 1 should not see tenant 2's guild_id\"\n assert (\n tenant1_guild[\"guild_name\"] != \"Tenant 2 Server\"\n ), \"Tenant 1 should not see tenant 2's guild_name\"\n\n # Verify tenant 2 sees only their guild\n list_response2 = requests.get(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user2.headers,\n )\n assert list_response2.ok\n tenant2_guilds = list_response2.json()\n\n # Tenant 2 should see exactly 1 guild\n assert (\n len(tenant2_guilds) == 1\n ), f\"Tenant 2 should see 1 guild, got {len(tenant2_guilds)}\"\n\n # Verify tenant 2's guild has the correct data\n tenant2_guild = tenant2_guilds[0]\n assert (\n tenant2_guild[\"id\"] == guild2_id\n ), \"Tenant 2 should see their own guild\"\n assert (\n tenant2_guild[\"guild_id\"] == 222222222222222222\n ), f\"Tenant 2's guild should have guild_id 222222222222222222, got {tenant2_guild['guild_id']}\"\n assert (\n tenant2_guild[\"guild_name\"] == \"Tenant 2 Server\"\n ), f\"Tenant 2's guild should have name 'Tenant 2 Server', got {tenant2_guild['guild_name']}\"\n assert (\n tenant2_guild[\"registered_at\"] is not None\n ), \"Tenant 2's guild should be registered\"\n\n # Tenant 2 should NOT see tenant 1's guild\n assert (\n tenant2_guild[\"guild_id\"] != 111111111111111111\n ), \"Tenant 2 should not see tenant 1's guild_id\"\n assert (\n tenant2_guild[\"guild_name\"] != \"Tenant 1 Server\"\n ), \"Tenant 2 should not see tenant 1's guild_name\"\n\n # Verify the guilds are different (different data)\n assert (\n tenant1_guild[\"guild_id\"] != tenant2_guild[\"guild_id\"]\n ), \"Guilds should have different Discord guild IDs\"\n assert (\n tenant1_guild[\"guild_name\"] != tenant2_guild[\"guild_name\"]\n ), \"Guilds should have different names\"\n\n finally:\n # Cleanup\n requests.delete(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds/{guild1_id}\",\n headers=admin_user1.headers,\n )\n requests.delete(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds/{guild2_id}\",\n headers=admin_user2.headers,\n )\n\n\nclass TestGuildAccessIsolation:\n \"\"\"Tests for guild access isolation between tenants.\"\"\"\n\n def test_tenant_cannot_access_other_tenant_guild(\n self,\n reset_multitenant: None, # noqa: ARG002\n ) -> None:\n \"\"\"Tenant 2 cannot access or modify tenant 1's guild by ID.\n\n Creates a guild in tenant 1, then attempts to access it from tenant 2.\n \"\"\"\n unique = uuid4().hex\n\n # Create admin users for two tenants\n admin_user1: DATestUser = UserManager.create(\n email=f\"discord_access1_{unique}@example.com\",\n )\n admin_user2: DATestUser = UserManager.create(\n email=f\"discord_access2_{unique}@example.com\",\n )\n\n # Create a guild in tenant 1\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds\",\n headers=admin_user1.headers,\n )\n if response.status_code == 404:\n pytest.skip(\"Discord bot feature not enabled\")\n assert response.ok\n guild1_id = response.json()[\"id\"]\n\n try:\n # Tenant 2 tries to get the guild - should fail (404 or 403)\n get_response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds/{guild1_id}\",\n headers=admin_user2.headers,\n )\n # Should either return 404 (not found) or 403 (forbidden)\n assert get_response.status_code in [\n 403,\n 404,\n ], f\"Expected 403 or 404, got {get_response.status_code}\"\n\n # Tenant 2 tries to delete the guild - should fail\n delete_response = requests.delete(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds/{guild1_id}\",\n headers=admin_user2.headers,\n )\n assert delete_response.status_code in [403, 404]\n\n finally:\n # Cleanup - delete from tenant 1\n requests.delete(\n f\"{API_SERVER_URL}/manage/admin/discord-bot/guilds/{guild1_id}\",\n headers=admin_user1.headers,\n )\n\n\nclass TestCacheManagerIsolation:\n \"\"\"Tests for cache manager tenant isolation.\"\"\"\n\n def test_cache_maps_guild_to_correct_tenant(self) -> None:\n \"\"\"Cache correctly maps guild_id to tenant_id.\"\"\"\n cache = DiscordCacheManager()\n\n # Set up mappings\n cache._guild_tenants[111] = \"tenant1\"\n cache._guild_tenants[222] = \"tenant2\"\n cache._guild_tenants[333] = \"tenant1\"\n\n assert cache.get_tenant(111) == \"tenant1\"\n assert cache.get_tenant(222) == \"tenant2\"\n assert cache.get_tenant(333) == \"tenant1\"\n assert cache.get_tenant(444) is None\n\n def test_api_key_per_tenant_isolation(self) -> None:\n \"\"\"Each tenant has unique API key.\"\"\"\n cache = DiscordCacheManager()\n\n cache._api_keys[\"tenant1\"] = \"key_for_tenant1\"\n cache._api_keys[\"tenant2\"] = \"key_for_tenant2\"\n\n assert cache.get_api_key(\"tenant1\") == \"key_for_tenant1\"\n assert cache.get_api_key(\"tenant2\") == \"key_for_tenant2\"\n assert cache.get_api_key(\"tenant1\") != cache.get_api_key(\"tenant2\")\n\n\nclass TestAPIRequestIsolation:\n \"\"\"Tests for API request isolation between tenants.\"\"\"\n\n @pytest.mark.asyncio\n async def test_discord_bot_uses_tenant_specific_api_key(self) -> None:\n \"\"\"Message from guild in tenant 1 uses tenant 1's API key.\"\"\"\n cache = DiscordCacheManager()\n cache._guild_tenants[123456] = \"tenant1\"\n cache._api_keys[\"tenant1\"] = \"tenant1_api_key\"\n cache._api_keys[\"tenant2\"] = \"tenant2_api_key\"\n\n # When processing message from guild 123456\n tenant = cache.get_tenant(123456)\n assert tenant is not None\n api_key = cache.get_api_key(tenant)\n\n assert tenant == \"tenant1\"\n assert api_key == \"tenant1_api_key\"\n assert api_key != \"tenant2_api_key\"\n\n @pytest.mark.asyncio\n async def test_guild_message_routes_to_correct_tenant(self) -> None:\n \"\"\"Message from registered guild routes to correct tenant context.\"\"\"\n cache = DiscordCacheManager()\n cache._guild_tenants[999] = \"target_tenant\"\n cache._api_keys[\"target_tenant\"] = \"target_key\"\n\n # Simulate message routing\n guild_id = 999\n tenant = cache.get_tenant(guild_id)\n api_key = cache.get_api_key(tenant) if tenant else None\n\n assert tenant == \"target_tenant\"\n assert api_key == \"target_key\"\n" }, { "path": "backend/tests/integration/multitenant_tests/invitation/test_user_invitation.py", "content": "from uuid import uuid4\n\nfrom onyx.db.models import UserRole\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\nINVITED_BASIC_USER = \"basic_user\"\nINVITED_BASIC_USER_EMAIL = \"basic_user@example.com\"\n\n\ndef test_admin_can_invite_users(reset_multitenant: None) -> None: # noqa: ARG001\n \"\"\"Test that an admin can invite both registered and non-registered users.\"\"\"\n # Create first user (admin)\n unique = uuid4().hex\n admin_user: DATestUser = UserManager.create(name=f\"admin_{unique}\")\n assert UserManager.is_role(admin_user, UserRole.ADMIN)\n\n # Create second user\n invited_user: DATestUser = UserManager.create(name=f\"admin_invited_{unique}\")\n assert UserManager.is_role(invited_user, UserRole.ADMIN)\n\n # Admin user invites the previously registered and non-registered user\n UserManager.invite_user(invited_user.email, admin_user)\n UserManager.invite_user(f\"{INVITED_BASIC_USER}_{unique}@example.com\", admin_user)\n\n # Verify users are in the invited users list\n invited_users = UserManager.get_invited_users(admin_user)\n assert invited_user.email in [\n user.email for user in invited_users\n ], f\"User {invited_user.email} not found in invited users list\"\n\n\ndef test_non_registered_user_gets_basic_role(\n reset_multitenant: None, # noqa: ARG001\n) -> None:\n \"\"\"Test that a non-registered user gets a BASIC role when they register after being invited.\"\"\"\n # Create admin user\n unique = uuid4().hex\n admin_user: DATestUser = UserManager.create(name=f\"admin_{unique}\")\n assert UserManager.is_role(admin_user, UserRole.ADMIN)\n\n # Admin user invites a non-registered user\n invited_email = f\"{INVITED_BASIC_USER}_{unique}@example.com\"\n UserManager.invite_user(invited_email, admin_user)\n\n # Non-registered user registers\n invited_basic_user: DATestUser = UserManager.create(\n name=f\"{INVITED_BASIC_USER}_{unique}\", email=invited_email\n )\n assert UserManager.is_role(invited_basic_user, UserRole.BASIC)\n\n\ndef test_user_can_accept_invitation(\n reset_multitenant: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Test that a user can accept an invitation and join the organization with BASIC role.\"\"\"\n # Create admin user\n unique = uuid4().hex\n admin_user: DATestUser = UserManager.create(name=f\"admin_{unique}\")\n assert UserManager.is_role(admin_user, UserRole.ADMIN)\n\n # Create a user to be invited\n invited_user_email = f\"invited_user_{unique}@example.com\"\n\n # User registers with the same email as the invitation\n invited_user: DATestUser = UserManager.create(\n name=f\"invited_user_{unique}\", email=invited_user_email\n )\n # Admin user invites the user\n UserManager.invite_user(invited_user_email, admin_user)\n\n # Get user info to check tenant information\n user_info = UserManager.get_user_info(invited_user)\n\n # Extract the tenant_id from the invitation\n invited_tenant_id = (\n user_info.tenant_info.invitation.tenant_id\n if user_info.tenant_info and user_info.tenant_info.invitation\n else None\n )\n assert invited_tenant_id is not None, \"Expected to find an invitation tenant_id\"\n\n # User accepts invitation\n UserManager.accept_invitation(invited_tenant_id, invited_user)\n\n # User needs to reauthenticate after accepting invitation\n # Simulate this by creating a new user instance with the same credentials\n authenticated_user: DATestUser = UserManager.create(\n name=\"invited_user\", email=invited_user_email\n )\n\n # Get updated user info after accepting invitation and reauthenticating\n updated_user_info = UserManager.get_user_info(authenticated_user)\n\n # Verify the user has BASIC role in the organization\n assert (\n updated_user_info.role == UserRole.BASIC\n ), f\"Expected user to have BASIC role, but got {updated_user_info.role}\"\n\n # Verify user is in the organization\n user_page = UserManager.get_user_page(\n user_performing_action=admin_user, role_filter=[UserRole.BASIC]\n )\n\n # Check if the invited user is in the list of users with BASIC role\n invited_user_emails = [user.email for user in user_page.items]\n assert invited_user_email in invited_user_emails, (\n f\"User {invited_user_email} not found in the list of basic users \"\n f\"in the organization. Available users: {invited_user_emails}\"\n )\n\n invited_users = UserManager.get_invited_users(admin_user)\n assert invited_user.email not in [\n user.email for user in invited_users\n ], f\"User {invited_user.email} should not be found in invited users list after accepting invitation\"\n" }, { "path": "backend/tests/integration/multitenant_tests/migrations/test_run_multitenant_migrations.py", "content": "\"\"\"\nBlack-box integration tests for the parallel alembic migration runner\n(backend/alembic/run_multitenant_migrations.py).\n\nThe script is invoked as a subprocess — the same way it would be used in\nproduction. Tests verify exit codes and stdout messages.\n\nUsage:\n pytest tests/integration/tests/migrations/test_run_multitenant_migrations.py -v\n\"\"\"\n\nfrom __future__ import annotations\n\nimport os\nimport subprocess\nimport sys\nimport time\nimport uuid\nfrom collections.abc import Generator\n\nimport pytest\nfrom sqlalchemy import text\nfrom sqlalchemy.engine import Engine\n\nfrom onyx.db.engine.sql_engine import SqlEngine\n\n# Resolve the backend/ directory once so every helper can use it as cwd.\n_BACKEND_DIR = os.path.normpath(\n os.path.join(os.path.dirname(__file__), \"..\", \"..\", \"..\", \"..\")\n)\n\n_DROP_SCHEMA_MAX_RETRIES = 3\n_DROP_SCHEMA_RETRY_DELAY_SEC = 2\n\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\n\ndef _run_script(\n *extra_args: str,\n env_override: dict[str, str] | None = None,\n) -> subprocess.CompletedProcess[str]:\n \"\"\"Run ``python alembic/run_multitenant_migrations.py`` from the backend/ directory.\"\"\"\n env = {**os.environ, **(env_override or {})}\n return subprocess.run(\n [sys.executable, \"alembic/run_multitenant_migrations.py\", *extra_args],\n cwd=_BACKEND_DIR,\n stdout=subprocess.PIPE,\n stderr=subprocess.STDOUT,\n text=True,\n env=env,\n )\n\n\ndef _force_drop_schema(engine: Engine, schema: str) -> None:\n \"\"\"Terminate backends using *schema* then drop it, retrying on deadlock.\n\n Background Celery workers may discover test schemas (they match the\n ``tenant_`` prefix) and hold locks on tables inside them. A bare\n ``DROP SCHEMA … CASCADE`` can deadlock with those workers, so we\n first kill their connections and retry if we still hit a deadlock.\n \"\"\"\n for attempt in range(_DROP_SCHEMA_MAX_RETRIES):\n try:\n with engine.connect() as conn:\n conn.execute(\n text(\n \"\"\"\n SELECT pg_terminate_backend(l.pid)\n FROM pg_locks l\n JOIN pg_class c ON c.oid = l.relation\n JOIN pg_namespace n ON n.oid = c.relnamespace\n WHERE n.nspname = :schema\n AND l.pid != pg_backend_pid()\n \"\"\"\n ),\n {\"schema\": schema},\n )\n conn.execute(text(f'DROP SCHEMA IF EXISTS \"{schema}\" CASCADE'))\n conn.commit()\n return\n except Exception:\n if attempt == _DROP_SCHEMA_MAX_RETRIES - 1:\n raise\n time.sleep(_DROP_SCHEMA_RETRY_DELAY_SEC)\n\n\n# ---------------------------------------------------------------------------\n# Fixtures\n# ---------------------------------------------------------------------------\n\n\n@pytest.fixture\ndef engine() -> Engine:\n return SqlEngine.get_engine()\n\n\n@pytest.fixture\ndef current_head_rev() -> str:\n \"\"\"Get the head revision from the alembic script directory.\n\n Runs ``alembic heads`` as a subprocess — the same source of truth that\n ``run_multitenant_migrations.py`` uses internally.\n \"\"\"\n result = subprocess.run(\n [\"alembic\", \"heads\", \"--resolve-dependencies\"],\n cwd=_BACKEND_DIR,\n stdout=subprocess.PIPE,\n stderr=subprocess.STDOUT,\n text=True,\n )\n assert (\n result.returncode == 0\n ), f\"alembic heads failed (exit {result.returncode}):\\n{result.stdout}\"\n # Output looks like \"d5c86e2c6dc6 (head)\\n\"\n rev = result.stdout.strip().split()[0]\n assert len(rev) > 0\n return rev\n\n\n@pytest.fixture\ndef tenant_schema_at_head(\n engine: Engine, current_head_rev: str\n) -> Generator[str, None, None]:\n \"\"\"Create a temporary tenant schema whose alembic_version is at head.\"\"\"\n schema = f\"tenant_test_{uuid.uuid4().hex[:12]}\"\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA \"{schema}\"'))\n conn.execute(\n text(\n f'CREATE TABLE \"{schema}\".alembic_version (version_num VARCHAR(32) NOT NULL)'\n )\n )\n conn.execute(\n text(f'INSERT INTO \"{schema}\".alembic_version (version_num) VALUES (:rev)'),\n {\"rev\": current_head_rev},\n )\n conn.commit()\n\n yield schema\n\n _force_drop_schema(engine, schema)\n\n\n@pytest.fixture\ndef tenant_schema_empty(engine: Engine) -> Generator[str, None, None]:\n \"\"\"Create a temporary tenant schema with no tables at all.\n\n Alembic will treat it as a fresh schema and run every migration from base\n to head.\n \"\"\"\n schema = f\"tenant_test_{uuid.uuid4().hex[:12]}\"\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA \"{schema}\"'))\n conn.commit()\n\n yield schema\n\n _force_drop_schema(engine, schema)\n\n\n@pytest.fixture\ndef tenant_schema_bad_rev(engine: Engine) -> Generator[str, None, None]:\n \"\"\"Create a tenant schema whose alembic_version points to a non-existent\n revision. Alembic cannot find a migration path from this revision, so\n it will fail.\"\"\"\n schema = f\"tenant_test_{uuid.uuid4().hex[:12]}\"\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA \"{schema}\"'))\n conn.execute(\n text(\n f'CREATE TABLE \"{schema}\".alembic_version (version_num VARCHAR(32) NOT NULL)'\n )\n )\n conn.execute(\n text(\n f\"INSERT INTO \\\"{schema}\\\".alembic_version (version_num) VALUES ('00000bad0000')\"\n )\n )\n conn.commit()\n\n yield schema\n\n _force_drop_schema(engine, schema)\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\ndef test_no_tenant_schemas_exits_nonzero() -> None:\n \"\"\"In non-multi-tenant mode there are no tenant_ schemas, so the script\n should print a hint and exit 1.\"\"\"\n result = _run_script(env_override={\"MULTI_TENANT\": \"false\"})\n assert result.returncode == 1\n assert \"No tenant schemas found\" in result.stdout\n assert \"MULTI_TENANT\" in result.stdout\n\n\ndef test_at_head_schema_is_skipped(tenant_schema_at_head: str) -> None:\n \"\"\"A tenant schema already at head should not be targeted for migration.\"\"\"\n result = _run_script(\n \"--jobs\",\n \"1\",\n \"--batch-size\",\n \"50\",\n env_override={\"MULTI_TENANT\": \"true\"},\n )\n assert result.returncode == 0\n # Our at-head schema should not appear in any batch \"started\" lines.\n batch_start_lines = [\n line\n for line in result.stdout.splitlines()\n if \"Batch\" in line and \"started\" in line\n ]\n for line in batch_start_lines:\n assert tenant_schema_at_head not in line\n\n\ndef test_detects_schemas_needing_migration(\n tenant_schema_at_head: str,\n tenant_schema_empty: str,\n) -> None:\n \"\"\"When some schemas are behind, the script should report how many need\n migration, upgrade them, and succeed.\"\"\"\n result = _run_script(\n \"--jobs\",\n \"1\",\n \"--batch-size\",\n \"50\",\n env_override={\"MULTI_TENANT\": \"true\"},\n )\n assert result.returncode == 0, f\"Script failed:\\n{result.stdout}\"\n assert \"tenants need migration\" in result.stdout\n assert \"All migrations successful\" in result.stdout\n\n # The empty schema should appear in the batch that was started.\n assert tenant_schema_empty in result.stdout\n\n # The at-head schema should NOT appear in any batch \"started\" lines\n # (it was filtered out by get_schemas_needing_migration).\n batch_start_lines = [\n line\n for line in result.stdout.splitlines()\n if \"Batch\" in line and \"started\" in line\n ]\n for line in batch_start_lines:\n assert tenant_schema_at_head not in line\n\n\ndef test_failed_migration(\n tenant_schema_at_head: str,\n tenant_schema_empty: str,\n tenant_schema_bad_rev: str,\n) -> None:\n \"\"\"A schema with a bogus alembic revision causes alembic to fail.\n\n The script should:\n - Exit non-zero (some migrations failed).\n - Still skip the at-head schema.\n - Still attempt the other schemas via the ``continue=true`` retry.\n \"\"\"\n result = _run_script(\n \"--jobs\",\n \"1\",\n \"--batch-size\",\n \"50\",\n env_override={\"MULTI_TENANT\": \"true\"},\n )\n assert result.returncode == 1, f\"Expected failure but got:\\n{result.stdout}\"\n assert \"Some migrations failed\" in result.stdout\n\n # The bad-rev schema should appear in the batch (it needs migration).\n assert tenant_schema_bad_rev in result.stdout\n\n # The empty schema should also appear (it was attempted via continue=true retry).\n assert tenant_schema_empty in result.stdout\n\n # The at-head schema should still be skipped.\n batch_start_lines = [\n line\n for line in result.stdout.splitlines()\n if \"Batch\" in line and \"started\" in line\n ]\n for line in batch_start_lines:\n assert tenant_schema_at_head not in line\n" }, { "path": "backend/tests/integration/multitenant_tests/syncing/test_search_permissions.py", "content": "from typing import Any\nfrom uuid import uuid4\n\nfrom onyx.db.models import UserRole\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestChatSession\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import ToolName\n\n\ndef setup_test_tenants(reset_multitenant: None) -> dict[str, Any]: # noqa: ARG001\n \"\"\"Helper function to set up test tenants with documents and users.\"\"\"\n unique = uuid4().hex\n # Creating an admin user for Tenant 1\n admin_user1: DATestUser = UserManager.create(\n email=f\"admin_{unique}@example.com\",\n )\n assert UserManager.is_role(admin_user1, UserRole.ADMIN)\n\n # Create Tenant 2 and its Admin User\n admin_user2: DATestUser = UserManager.create(\n email=f\"admin2_{unique}@example.com\",\n )\n assert UserManager.is_role(admin_user2, UserRole.ADMIN)\n\n # Create connectors for Tenant 1\n cc_pair_1: DATestCCPair = CCPairManager.create_from_scratch(\n user_performing_action=admin_user1,\n )\n api_key_1: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user1,\n )\n api_key_1.headers.update(admin_user1.headers)\n LLMProviderManager.create(user_performing_action=admin_user1)\n\n # Create connectors for Tenant 2\n cc_pair_2: DATestCCPair = CCPairManager.create_from_scratch(\n user_performing_action=admin_user2,\n )\n api_key_2: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user2,\n )\n api_key_2.headers.update(admin_user2.headers)\n LLMProviderManager.create(user_performing_action=admin_user2)\n\n # Seed documents for Tenant 1\n cc_pair_1.documents = []\n doc1_tenant1 = DocumentManager.seed_doc_with_content(\n cc_pair=cc_pair_1,\n content=\"Tenant 1 Document Content\",\n api_key=api_key_1,\n )\n doc2_tenant1 = DocumentManager.seed_doc_with_content(\n cc_pair=cc_pair_1,\n content=\"Tenant 1 Document Content\",\n api_key=api_key_1,\n )\n cc_pair_1.documents.extend([doc1_tenant1, doc2_tenant1])\n\n # Seed documents for Tenant 2\n cc_pair_2.documents = []\n doc1_tenant2 = DocumentManager.seed_doc_with_content(\n cc_pair=cc_pair_2,\n content=\"Tenant 2 Document Content\",\n api_key=api_key_2,\n )\n doc2_tenant2 = DocumentManager.seed_doc_with_content(\n cc_pair=cc_pair_2,\n content=\"Tenant 2 Document Content\",\n api_key=api_key_2,\n )\n cc_pair_2.documents.extend([doc1_tenant2, doc2_tenant2])\n\n tenant1_doc_ids = {doc1_tenant1.id, doc2_tenant1.id}\n tenant2_doc_ids = {doc1_tenant2.id, doc2_tenant2.id}\n\n # Create chat sessions for each user\n chat_session1: DATestChatSession = ChatSessionManager.create(\n user_performing_action=admin_user1\n )\n chat_session2: DATestChatSession = ChatSessionManager.create(\n user_performing_action=admin_user2\n )\n\n return {\n \"admin_user1\": admin_user1,\n \"admin_user2\": admin_user2,\n \"chat_session1\": chat_session1,\n \"chat_session2\": chat_session2,\n \"tenant1_doc_ids\": tenant1_doc_ids,\n \"tenant2_doc_ids\": tenant2_doc_ids,\n }\n\n\ndef test_tenant1_can_access_own_documents(reset_multitenant: None) -> None:\n \"\"\"Test that Tenant 1 can access its own documents but not Tenant 2's.\"\"\"\n test_data = setup_test_tenants(reset_multitenant)\n\n # User 1 sends a message and gets a response\n response1 = ChatSessionManager.send_message(\n chat_session_id=test_data[\"chat_session1\"].id,\n message=\"What is in Tenant 1's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user1\"],\n )\n\n assert response1.error is None, \"Chat response should not have an error\"\n\n # Assert that only the internal search tool was used\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response1.used_tools\n )\n\n response_doc_ids = {doc.document_id for doc in response1.used_tools[0].documents}\n assert test_data[\"tenant1_doc_ids\"].issubset(\n response_doc_ids\n ), \"Not all Tenant 1 document IDs are in the response\"\n assert not response_doc_ids.intersection(\n test_data[\"tenant2_doc_ids\"]\n ), \"Tenant 2 document IDs should not be in the response\"\n\n # Assert that the contents are correct\n assert any(\n doc.blurb == \"Tenant 1 Document Content\"\n for doc in response1.used_tools[0].documents\n ), \"Tenant 1 Document Content not found in any document\"\n\n\ndef test_tenant2_can_access_own_documents(reset_multitenant: None) -> None:\n \"\"\"Test that Tenant 2 can access its own documents but not Tenant 1's.\"\"\"\n test_data = setup_test_tenants(reset_multitenant)\n\n # User 2 sends a message and gets a response\n response2 = ChatSessionManager.send_message(\n chat_session_id=test_data[\"chat_session2\"].id,\n message=\"What is in Tenant 2's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user2\"],\n )\n\n assert response2.error is None, \"Chat response should not have an error\"\n\n # Assert that the search tool was used\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response2.used_tools\n )\n\n # Assert that the tool_result contains Tenant 2's documents\n response_doc_ids = {doc.document_id for doc in response2.used_tools[0].documents}\n assert test_data[\"tenant2_doc_ids\"].issubset(\n response_doc_ids\n ), \"Not all Tenant 2 document IDs are in the response\"\n assert not response_doc_ids.intersection(\n test_data[\"tenant1_doc_ids\"]\n ), \"Tenant 1 document IDs should not be in the response\"\n\n # Assert that the contents are correct\n assert any(\n doc.blurb == \"Tenant 2 Document Content\"\n for doc in response2.used_tools[0].documents\n ), \"Tenant 2 Document Content not found in any document\"\n\n\ndef test_tenant1_cannot_access_tenant2_documents(reset_multitenant: None) -> None:\n \"\"\"Test that Tenant 1 cannot access Tenant 2's documents.\"\"\"\n test_data = setup_test_tenants(reset_multitenant)\n\n # User 1 tries to access Tenant 2's documents\n response_cross = ChatSessionManager.send_message(\n chat_session_id=test_data[\"chat_session1\"].id,\n message=\"What is in Tenant 2's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user1\"],\n )\n\n assert response_cross.error is None, \"Chat response should not have an error\"\n\n # Assert that the search tool was used\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response_cross.used_tools\n )\n\n # Assert that the tool_result is empty or does not contain Tenant 2's documents\n response_doc_ids = {\n doc.document_id for doc in response_cross.used_tools[0].documents\n }\n\n # Ensure none of Tenant 2's document IDs are in the response\n assert not response_doc_ids.intersection(test_data[\"tenant2_doc_ids\"])\n\n\ndef test_tenant2_cannot_access_tenant1_documents(reset_multitenant: None) -> None:\n \"\"\"Test that Tenant 2 cannot access Tenant 1's documents.\"\"\"\n test_data = setup_test_tenants(reset_multitenant)\n\n # User 2 tries to access Tenant 1's documents\n response_cross2 = ChatSessionManager.send_message(\n chat_session_id=test_data[\"chat_session2\"].id,\n message=\"What is in Tenant 1's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user2\"],\n )\n\n assert response_cross2.error is None, \"Chat response should not have an error\"\n\n # Assert that the search tool was used\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH\n for tool in response_cross2.used_tools\n )\n\n # Assert that the tool_result is empty or does not contain Tenant 1's documents\n response_doc_ids = {\n doc.document_id for doc in response_cross2.used_tools[0].documents\n }\n\n # Ensure none of Tenant 1's document IDs are in the response\n assert not response_doc_ids.intersection(test_data[\"tenant1_doc_ids\"])\n\n\ndef test_multi_tenant_access_control(reset_multitenant: None) -> None:\n \"\"\"Legacy test for multi-tenant access control.\"\"\"\n test_data = setup_test_tenants(reset_multitenant)\n\n # User 1 sends a message and gets a response with only Tenant 1's documents\n response1 = ChatSessionManager.send_message(\n chat_session_id=test_data[\"chat_session1\"].id,\n message=\"What is in Tenant 1's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user1\"],\n )\n assert response1.error is None, \"Chat response should not have an error\"\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response1.used_tools\n )\n response_doc_ids = {doc.document_id for doc in response1.used_tools[0].documents}\n assert test_data[\"tenant1_doc_ids\"].issubset(response_doc_ids)\n assert not response_doc_ids.intersection(test_data[\"tenant2_doc_ids\"])\n\n # User 2 sends a message and gets a response with only Tenant 2's documents\n response2 = ChatSessionManager.send_message(\n chat_session_id=test_data[\"chat_session2\"].id,\n message=\"What is in Tenant 2's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user2\"],\n )\n assert response2.error is None, \"Chat response should not have an error\"\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response2.used_tools\n )\n response_doc_ids = {doc.document_id for doc in response2.used_tools[0].documents}\n assert test_data[\"tenant2_doc_ids\"].issubset(response_doc_ids)\n assert not response_doc_ids.intersection(test_data[\"tenant1_doc_ids\"])\n\n # User 1 tries to access Tenant 2's documents and fails\n user1_second_chat_session = ChatSessionManager.create(\n user_performing_action=test_data[\"admin_user1\"]\n )\n response_cross = ChatSessionManager.send_message(\n chat_session_id=user1_second_chat_session.id,\n message=\"What is in Tenant 2's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user1\"],\n )\n assert response_cross.error is None, \"Chat response should not have an error\"\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response_cross.used_tools\n )\n response_doc_ids = {\n doc.document_id for doc in response_cross.used_tools[0].documents\n }\n assert not response_doc_ids.intersection(test_data[\"tenant2_doc_ids\"])\n\n # User 2 tries to access Tenant 1's documents and fails\n user2_second_chat_session = ChatSessionManager.create(\n user_performing_action=test_data[\"admin_user2\"]\n )\n response_cross2 = ChatSessionManager.send_message(\n chat_session_id=user2_second_chat_session.id,\n message=\"What is in Tenant 1's documents? Run an internal search.\",\n user_performing_action=test_data[\"admin_user2\"],\n )\n assert response_cross2.error is None, \"Chat response should not have an error\"\n assert all(\n tool.tool_name == ToolName.INTERNAL_SEARCH\n for tool in response_cross2.used_tools\n )\n response_doc_ids = {\n doc.document_id for doc in response_cross2.used_tools[0].documents\n }\n assert not response_doc_ids.intersection(test_data[\"tenant1_doc_ids\"])\n" }, { "path": "backend/tests/integration/multitenant_tests/tenants/test_tenant_creation.py", "content": "from http import HTTPStatus\nfrom uuid import uuid4\n\nimport requests\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.db.enums import AccessType\nfrom onyx.db.models import UserRole\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.image_generation import (\n ImageGenerationConfigManager,\n)\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_first_user_is_admin(reset_multitenant: None) -> None: # noqa: ARG001\n \"\"\"Test that the first user of a tenant is automatically assigned ADMIN role.\"\"\"\n unique = uuid4().hex\n test_user: DATestUser = UserManager.create(\n name=f\"test_{unique}\", email=f\"test_{unique}@example.com\"\n )\n assert UserManager.is_role(test_user, UserRole.ADMIN)\n\n\ndef test_admin_can_create_credential(\n reset_multitenant: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Test that an admin user can create a credential in their tenant.\"\"\"\n # Create admin user\n unique = uuid4().hex\n test_user: DATestUser = UserManager.create(\n name=f\"test_{unique}\", email=f\"test_{unique}@example.com\"\n )\n assert UserManager.is_role(test_user, UserRole.ADMIN)\n\n # Create credential\n test_credential = CredentialManager.create(\n name=\"admin_test_credential\",\n source=DocumentSource.FILE,\n curator_public=False,\n user_performing_action=test_user,\n )\n assert test_credential is not None\n\n\ndef test_admin_can_create_connector(\n reset_multitenant: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Test that an admin user can create a connector in their tenant.\"\"\"\n # Create admin user\n unique = uuid4().hex\n test_user: DATestUser = UserManager.create(\n name=f\"test_{unique}\", email=f\"test_{unique}@example.com\"\n )\n assert UserManager.is_role(test_user, UserRole.ADMIN)\n\n # Create connector\n test_connector = ConnectorManager.create(\n name=\"admin_test_connector\",\n source=DocumentSource.FILE,\n access_type=AccessType.PRIVATE,\n user_performing_action=test_user,\n )\n assert test_connector is not None\n\n\ndef test_admin_can_create_and_verify_cc_pair(\n reset_multitenant: None, # noqa: ARG001\n) -> None:\n \"\"\"Test that an admin user can create and verify a connector-credential pair in their tenant.\"\"\"\n # Create admin user\n unique = uuid4().hex\n test_user: DATestUser = UserManager.create(\n name=f\"test_{unique}\", email=f\"test_{unique}@example.com\"\n )\n assert UserManager.is_role(test_user, UserRole.ADMIN)\n\n # Create credential\n test_credential = CredentialManager.create(\n name=\"admin_test_credential\",\n source=DocumentSource.FILE,\n curator_public=False,\n user_performing_action=test_user,\n )\n\n # Create connector\n test_connector = ConnectorManager.create(\n name=\"admin_test_connector\",\n source=DocumentSource.FILE,\n access_type=AccessType.PRIVATE,\n user_performing_action=test_user,\n )\n\n # Create cc_pair\n test_cc_pair = CCPairManager.create(\n connector_id=test_connector.id,\n credential_id=test_credential.id,\n name=\"admin_test_cc_pair\",\n access_type=AccessType.PRIVATE,\n user_performing_action=test_user,\n )\n assert test_cc_pair is not None\n\n # Verify cc_pair\n CCPairManager.verify(cc_pair=test_cc_pair, user_performing_action=test_user)\n\n\ndef test_settings_access() -> None:\n \"\"\"Calls to the enterprise settings endpoint without authentication should fail with\n 403 (and not 500, which will lock the web UI into a \"maintenance mode\" page)\"\"\"\n\n response = requests.get(url=f\"{API_SERVER_URL}/enterprise-settings\")\n assert response.status_code == HTTPStatus.FORBIDDEN\n\n\ndef test_image_gen_config_created_on_tenant_provision(\n reset_multitenant: None, # noqa: ARG001\n) -> None:\n \"\"\"Test that image generation config is automatically created when a tenant is provisioned.\"\"\"\n unique = uuid4().hex\n test_user: DATestUser = UserManager.create(\n name=f\"test_{unique}\", email=f\"test_{unique}@example.com\"\n )\n assert UserManager.is_role(test_user, UserRole.ADMIN)\n\n # Check if image gen config was created during tenant provisioning\n all_configs = ImageGenerationConfigManager.get_all(user_performing_action=test_user)\n\n # Should have at least one config created during provisioning\n assert (\n len(all_configs) > 0\n ), \"Image generation config should be created during tenant provisioning\"\n\n # Verify a default config exists\n default_configs = [c for c in all_configs if c.is_default]\n assert (\n len(default_configs) == 1\n ), \"Exactly one default image generation config should exist\"\n\n # Verify expected properties\n default_config = default_configs[0]\n assert default_config.image_provider_id == \"openai_gpt_image_1\"\n assert default_config.model_name == \"gpt-image-1\"\n" }, { "path": "backend/tests/integration/multitenant_tests/tenants/test_tenant_provisioning_rollback.py", "content": "\"\"\"\nIntegration tests for tenant provisioning rollback behavior.\n\nTests the fix for the drop_schema bug where:\n1. isidentifier() rejected valid UUID tenant IDs (with hyphens)\n2. SQL syntax was broken (%(schema_name)s instead of proper identifier handling)\n\nThis test verifies the full flow: provisioning failure → rollback → schema cleanup.\n\"\"\"\n\nimport uuid\nfrom unittest.mock import patch\n\nfrom sqlalchemy import text\n\nfrom ee.onyx.server.tenants.schema_management import create_schema_if_not_exists\nfrom ee.onyx.server.tenants.schema_management import drop_schema\nfrom onyx.db.engine.sql_engine import get_session_with_shared_schema\nfrom shared_configs.configs import TENANT_ID_PREFIX\n\n\ndef _schema_exists(schema_name: str) -> bool:\n \"\"\"Check if a schema exists in the database.\"\"\"\n with get_session_with_shared_schema() as session:\n result = session.execute(\n text(\n \"SELECT 1 FROM information_schema.schemata WHERE schema_name = :schema\"\n ),\n {\"schema\": schema_name},\n ).fetchone()\n return result is not None\n\n\nclass TestTenantProvisioningRollback:\n \"\"\"Integration tests for provisioning failure and rollback.\"\"\"\n\n def test_failed_provisioning_cleans_up_schema(self) -> None:\n \"\"\"\n When setup_tenant fails after schema creation, rollback should\n clean up the orphaned schema.\n\n This is the actual bug scenario: pre_provision_tenant creates a schema,\n setup_tenant fails, rollback is called, but drop_schema was broken\n (isidentifier rejected UUIDs with hyphens), leaving orphaned schemas.\n \"\"\"\n from ee.onyx.background.celery.tasks.tenant_provisioning.tasks import (\n pre_provision_tenant,\n )\n\n # Track which tenant_id gets created\n created_tenant_id = None\n\n def track_schema_creation(tenant_id: str) -> bool:\n nonlocal created_tenant_id\n created_tenant_id = tenant_id\n return create_schema_if_not_exists(tenant_id)\n\n # Mock setup_tenant to fail after schema creation\n with patch(\n \"ee.onyx.background.celery.tasks.tenant_provisioning.tasks.setup_tenant\"\n ) as mock_setup:\n mock_setup.side_effect = Exception(\"Simulated provisioning failure\")\n\n with patch(\n \"ee.onyx.background.celery.tasks.tenant_provisioning.tasks.create_schema_if_not_exists\",\n side_effect=track_schema_creation,\n ):\n # Run pre-provisioning - it should fail and trigger rollback\n pre_provision_tenant()\n\n # Verify that the schema was created and then cleaned up\n assert created_tenant_id is not None, \"Schema should have been created\"\n assert created_tenant_id.startswith(\n TENANT_ID_PREFIX\n ), f\"Should have tenant prefix: {created_tenant_id}\"\n assert not _schema_exists(\n created_tenant_id\n ), f\"Schema {created_tenant_id} should have been rolled back\"\n\n def test_drop_schema_works_with_uuid_tenant_id(self) -> None:\n \"\"\"\n drop_schema should work with UUID-format tenant IDs.\n\n This directly tests the fix: UUID tenant IDs contain hyphens,\n which isidentifier() rejected. The new regex validation accepts them.\n \"\"\"\n tenant_id = f\"{TENANT_ID_PREFIX}{uuid.uuid4()}\"\n\n # Create schema\n create_schema_if_not_exists(tenant_id)\n assert _schema_exists(tenant_id), \"Schema should exist after creation\"\n\n # Drop schema\n drop_schema(tenant_id)\n assert not _schema_exists(tenant_id), \"Schema should be dropped\"\n" }, { "path": "backend/tests/integration/multitenant_tests/test_get_schemas_needing_migration.py", "content": "\"\"\"\nIntegration tests for onyx.db.engine.tenant_utils.get_schemas_needing_migration.\n\nThese tests require a live database and exercise the function directly,\nindependent of the alembic migration runner script.\n\nUsage:\n pytest tests/integration/multitenant_tests/test_get_schemas_needing_migration.py -v\n\"\"\"\n\nfrom __future__ import annotations\n\nimport subprocess\nimport uuid\nfrom collections.abc import Generator\n\nimport pytest\nfrom sqlalchemy import text\nfrom sqlalchemy.engine import Engine\n\nfrom onyx.db.engine.sql_engine import SqlEngine\nfrom onyx.db.engine.tenant_utils import get_schemas_needing_migration\n\n_BACKEND_DIR = __file__[: __file__.index(\"/tests/\")]\n\n\n# ---------------------------------------------------------------------------\n# Fixtures\n# ---------------------------------------------------------------------------\n\n\n@pytest.fixture\ndef engine() -> Engine:\n return SqlEngine.get_engine()\n\n\n@pytest.fixture\ndef current_head_rev() -> str:\n result = subprocess.run(\n [\"alembic\", \"heads\", \"--resolve-dependencies\"],\n cwd=_BACKEND_DIR,\n stdout=subprocess.PIPE,\n stderr=subprocess.STDOUT,\n text=True,\n )\n assert (\n result.returncode == 0\n ), f\"alembic heads failed (exit {result.returncode}):\\n{result.stdout}\"\n rev = result.stdout.strip().split()[0]\n assert len(rev) > 0\n return rev\n\n\n@pytest.fixture\ndef tenant_schema_at_head(\n engine: Engine, current_head_rev: str\n) -> Generator[str, None, None]:\n \"\"\"Tenant schema with alembic_version already at head — should be excluded.\"\"\"\n schema = f\"tenant_test_{uuid.uuid4().hex[:12]}\"\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA \"{schema}\"'))\n conn.execute(\n text(\n f'CREATE TABLE \"{schema}\".alembic_version (version_num VARCHAR(32) NOT NULL)'\n )\n )\n conn.execute(\n text(f'INSERT INTO \"{schema}\".alembic_version (version_num) VALUES (:rev)'),\n {\"rev\": current_head_rev},\n )\n conn.commit()\n\n yield schema\n\n with engine.connect() as conn:\n conn.execute(text(f'DROP SCHEMA IF EXISTS \"{schema}\" CASCADE'))\n conn.commit()\n\n\n@pytest.fixture\ndef tenant_schema_empty(engine: Engine) -> Generator[str, None, None]:\n \"\"\"Tenant schema with no tables — should be included (needs migration).\"\"\"\n schema = f\"tenant_test_{uuid.uuid4().hex[:12]}\"\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA \"{schema}\"'))\n conn.commit()\n\n yield schema\n\n with engine.connect() as conn:\n conn.execute(text(f'DROP SCHEMA IF EXISTS \"{schema}\" CASCADE'))\n conn.commit()\n\n\n@pytest.fixture\ndef tenant_schema_stale_rev(engine: Engine) -> Generator[str, None, None]:\n \"\"\"Tenant schema with a non-head revision — should be included (needs migration).\"\"\"\n schema = f\"tenant_test_{uuid.uuid4().hex[:12]}\"\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA \"{schema}\"'))\n conn.execute(\n text(\n f'CREATE TABLE \"{schema}\".alembic_version (version_num VARCHAR(32) NOT NULL)'\n )\n )\n conn.execute(\n text(\n f\"INSERT INTO \\\"{schema}\\\".alembic_version (version_num) VALUES ('stalerev000000000000')\"\n )\n )\n conn.commit()\n\n yield schema\n\n with engine.connect() as conn:\n conn.execute(text(f'DROP SCHEMA IF EXISTS \"{schema}\" CASCADE'))\n conn.commit()\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\ndef test_classifies_all_cases(\n current_head_rev: str,\n tenant_schema_at_head: str,\n tenant_schema_empty: str,\n tenant_schema_stale_rev: str,\n) -> None:\n \"\"\"Correctly classifies all three schema states:\n - at head → excluded\n - no table → included (needs migration)\n - stale rev → included (needs migration)\n \"\"\"\n all_schemas = [tenant_schema_at_head, tenant_schema_empty, tenant_schema_stale_rev]\n result = get_schemas_needing_migration(all_schemas, current_head_rev)\n\n assert tenant_schema_at_head not in result\n assert tenant_schema_empty in result\n assert tenant_schema_stale_rev in result\n\n\ndef test_idempotent(\n current_head_rev: str,\n tenant_schema_at_head: str,\n tenant_schema_empty: str,\n) -> None:\n \"\"\"Calling the function twice returns the same result.\n\n Verifies that the DROP TABLE IF EXISTS guards correctly clean up temp\n tables so a second call succeeds even if the first left state behind.\n \"\"\"\n schemas = [tenant_schema_at_head, tenant_schema_empty]\n\n first = get_schemas_needing_migration(schemas, current_head_rev)\n second = get_schemas_needing_migration(schemas, current_head_rev)\n\n assert first == second\n\n\ndef test_empty_input(current_head_rev: str) -> None:\n \"\"\"An empty input list returns immediately without touching the DB.\"\"\"\n assert get_schemas_needing_migration([], current_head_rev) == []\n" }, { "path": "backend/tests/integration/tests/anonymous_user/test_anonymous_user.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.settings import SettingsManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestSettings\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_me_endpoint_returns_anonymous_user_when_enabled(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Unauthenticated /me returns anonymous user info when anonymous access is enabled.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n SettingsManager.update_settings(\n DATestSettings(anonymous_user_enabled=True),\n user_performing_action=admin_user,\n )\n\n response = requests.get(f\"{API_SERVER_URL}/me\")\n\n assert response.status_code == 200\n data = response.json()\n assert data[\"is_anonymous_user\"] is True\n assert data[\"email\"] == \"anonymous@onyx.app\"\n assert data[\"role\"] == \"limited\"\n\n\ndef test_me_endpoint_returns_403_when_anonymous_disabled(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Unauthenticated /me returns 403 when anonymous access is disabled.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n SettingsManager.update_settings(\n DATestSettings(anonymous_user_enabled=False),\n user_performing_action=admin_user,\n )\n\n response = requests.get(f\"{API_SERVER_URL}/me\")\n\n # 403 is returned when user is not authenticated\n assert response.status_code == 403\n\n\ndef test_me_endpoint_returns_authenticated_user_info(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Authenticated /me returns the actual user's info.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n response = requests.get(\n f\"{API_SERVER_URL}/me\",\n headers=admin_user.headers,\n )\n\n assert response.status_code == 200\n data = response.json()\n assert data.get(\"is_anonymous_user\") is not True\n assert data[\"email\"] == admin_user.email\n assert data[\"role\"] == \"admin\"\n\n\ndef test_anonymous_user_can_access_persona_when_enabled(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Verify that anonymous users can access limited endpoints when enabled.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n SettingsManager.update_settings(\n DATestSettings(anonymous_user_enabled=True),\n user_performing_action=admin_user,\n )\n\n anon_user = UserManager.get_anonymous_user()\n\n response = requests.get(\n f\"{API_SERVER_URL}/persona\",\n headers=anon_user.headers,\n )\n assert response.status_code == 200\n\n\ndef test_anonymous_user_denied_persona_when_disabled(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Verify that anonymous users cannot access endpoints when disabled.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n SettingsManager.update_settings(\n DATestSettings(anonymous_user_enabled=False),\n user_performing_action=admin_user,\n )\n\n anon_user = UserManager.get_anonymous_user()\n\n response = requests.get(\n f\"{API_SERVER_URL}/persona\",\n headers=anon_user.headers,\n )\n # 403 is returned - BasicAuthenticationError uses HTTP 403 for all auth failures\n assert response.status_code == 403\n" }, { "path": "backend/tests/integration/tests/api_key/test_api_key.py", "content": "from uuid import UUID\n\nimport requests\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.db.enums import AccountType\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_limited(reset: None) -> None: # noqa: ARG001\n \"\"\"Verify that with a limited role key, limited endpoints are accessible and\n others are not.\"\"\"\n\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n api_key: DATestAPIKey = APIKeyManager.create(\n api_key_role=UserRole.LIMITED,\n user_performing_action=admin_user,\n )\n\n # test limited endpoint\n response = requests.get(\n f\"{API_SERVER_URL}/persona/0\",\n headers=api_key.headers,\n )\n assert response.status_code == 200\n\n # test admin endpoints\n response = requests.get(\n f\"{API_SERVER_URL}/admin/api-key\",\n headers=api_key.headers,\n )\n assert response.status_code == 403\n\n\ndef _get_service_account_account_type(\n admin_user: DATestUser,\n api_key_user_id: UUID,\n) -> AccountType:\n \"\"\"Fetch the account_type of a service account user via the user listing API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/manage/users\",\n headers=admin_user.headers,\n params={\"include_api_keys\": \"true\"},\n )\n response.raise_for_status()\n data = response.json()\n user_id_str = str(api_key_user_id)\n for user in data[\"accepted\"]:\n if user[\"id\"] == user_id_str:\n return AccountType(user[\"account_type\"])\n raise AssertionError(\n f\"Service account user {user_id_str} not found in user listing\"\n )\n\n\ndef _get_default_group_user_ids(\n admin_user: DATestUser,\n) -> tuple[set[str], set[str]]:\n \"\"\"Return (admin_group_user_ids, basic_group_user_ids) from default groups.\"\"\"\n all_groups = UserGroupManager.get_all(\n user_performing_action=admin_user,\n include_default=True,\n )\n admin_group = next(\n (g for g in all_groups if g.name == \"Admin\" and g.is_default), None\n )\n basic_group = next(\n (g for g in all_groups if g.name == \"Basic\" and g.is_default), None\n )\n assert admin_group is not None, \"Admin default group not found\"\n assert basic_group is not None, \"Basic default group not found\"\n\n admin_ids = {str(u.id) for u in admin_group.users}\n basic_ids = {str(u.id) for u in basic_group.users}\n return admin_ids, basic_ids\n\n\ndef test_api_key_limited_service_account(reset: None) -> None: # noqa: ARG001\n \"\"\"LIMITED role API key: account_type is SERVICE_ACCOUNT, no group membership.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n api_key: DATestAPIKey = APIKeyManager.create(\n api_key_role=UserRole.LIMITED,\n user_performing_action=admin_user,\n )\n\n # Verify account_type\n account_type = _get_service_account_account_type(admin_user, api_key.user_id)\n assert (\n account_type == AccountType.SERVICE_ACCOUNT\n ), f\"Expected account_type={AccountType.SERVICE_ACCOUNT}, got {account_type}\"\n\n # Verify no group membership\n admin_ids, basic_ids = _get_default_group_user_ids(admin_user)\n user_id_str = str(api_key.user_id)\n assert (\n user_id_str not in admin_ids\n ), \"LIMITED API key should NOT be in Admin default group\"\n assert (\n user_id_str not in basic_ids\n ), \"LIMITED API key should NOT be in Basic default group\"\n\n\ndef test_api_key_basic_service_account(reset: None) -> None: # noqa: ARG001\n \"\"\"BASIC role API key: account_type is SERVICE_ACCOUNT, in Basic group only.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n api_key: DATestAPIKey = APIKeyManager.create(\n api_key_role=UserRole.BASIC,\n user_performing_action=admin_user,\n )\n\n # Verify account_type\n account_type = _get_service_account_account_type(admin_user, api_key.user_id)\n assert (\n account_type == AccountType.SERVICE_ACCOUNT\n ), f\"Expected account_type={AccountType.SERVICE_ACCOUNT}, got {account_type}\"\n\n # Verify Basic group membership\n admin_ids, basic_ids = _get_default_group_user_ids(admin_user)\n user_id_str = str(api_key.user_id)\n assert user_id_str in basic_ids, \"BASIC API key should be in Basic default group\"\n assert (\n user_id_str not in admin_ids\n ), \"BASIC API key should NOT be in Admin default group\"\n\n\ndef test_api_key_admin_service_account(reset: None) -> None: # noqa: ARG001\n \"\"\"ADMIN role API key: account_type is SERVICE_ACCOUNT, in Admin group only.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n api_key: DATestAPIKey = APIKeyManager.create(\n api_key_role=UserRole.ADMIN,\n user_performing_action=admin_user,\n )\n\n # Verify account_type\n account_type = _get_service_account_account_type(admin_user, api_key.user_id)\n assert (\n account_type == AccountType.SERVICE_ACCOUNT\n ), f\"Expected account_type={AccountType.SERVICE_ACCOUNT}, got {account_type}\"\n\n # Verify Admin group membership\n admin_ids, basic_ids = _get_default_group_user_ids(admin_user)\n user_id_str = str(api_key.user_id)\n assert user_id_str in admin_ids, \"ADMIN API key should be in Admin default group\"\n assert (\n user_id_str not in basic_ids\n ), \"ADMIN API key should NOT be in Basic default group\"\n" }, { "path": "backend/tests/integration/tests/auth/test_saml_user_conversion.py", "content": "import os\n\nimport pytest\nimport requests\n\nfrom onyx.auth.schemas import UserRole\nfrom onyx.db.enums import AccountType\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _simulate_saml_login(email: str, admin_user: DATestUser) -> dict:\n \"\"\"Simulate a SAML login by calling the test upsert endpoint.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/manage/users/test-upsert-user\",\n json={\"email\": email},\n headers=admin_user.headers,\n )\n response.raise_for_status()\n return response.json()\n\n\ndef _get_basic_group_member_emails(admin_user: DATestUser) -> set[str]:\n \"\"\"Get the set of emails of all members in the Basic default group.\"\"\"\n all_groups = UserGroupManager.get_all(admin_user, include_default=True)\n basic_default = [g for g in all_groups if g.is_default and g.name == \"Basic\"]\n assert basic_default, \"Basic default group not found\"\n return {u.email for u in basic_default[0].users}\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"SAML tests are enterprise only\",\n)\ndef test_saml_user_conversion(reset: None) -> None: # noqa: ARG001\n \"\"\"\n Test that SAML login correctly converts users with non-authenticated roles\n (SLACK_USER or EXT_PERM_USER) to authenticated roles (BASIC).\n\n This test:\n 1. Creates an admin and a regular user\n 2. Changes the regular user's role to EXT_PERM_USER\n 3. Simulates a SAML login by calling the test endpoint\n 4. Verifies the user's role is converted to BASIC\n\n This tests the fix that ensures users with non-authenticated roles (SLACK_USER or EXT_PERM_USER)\n are properly converted to authenticated roles during SAML login.\n \"\"\"\n # Create an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n # Create a regular user that we'll convert to EXT_PERM_USER\n test_user_email = \"ext_perm_user@example.com\"\n test_user = UserManager.create(email=test_user_email)\n\n # Verify the user was created with BASIC role initially\n assert UserManager.is_role(test_user, UserRole.BASIC)\n\n # Change the user's role to EXT_PERM_USER using the UserManager\n UserManager.set_role(\n user_to_set=test_user,\n target_role=UserRole.EXT_PERM_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n\n # Verify the user has EXT_PERM_USER role now\n assert UserManager.is_role(test_user, UserRole.EXT_PERM_USER)\n\n # Simulate SAML login by calling the test endpoint\n user_data = _simulate_saml_login(test_user_email, admin_user)\n\n # Verify the response indicates the role changed to BASIC\n assert user_data[\"role\"] == UserRole.BASIC.value\n\n # Verify user role was changed in the database\n assert UserManager.is_role(test_user, UserRole.BASIC)\n\n # Do the same test with SLACK_USER\n slack_user_email = \"slack_user@example.com\"\n slack_user = UserManager.create(email=slack_user_email)\n\n # Verify the user was created with BASIC role initially\n assert UserManager.is_role(slack_user, UserRole.BASIC)\n\n # Change the user's role to SLACK_USER\n UserManager.set_role(\n user_to_set=slack_user,\n target_role=UserRole.SLACK_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n\n # Verify the user has SLACK_USER role\n assert UserManager.is_role(slack_user, UserRole.SLACK_USER)\n\n # Simulate SAML login again\n user_data = _simulate_saml_login(slack_user_email, admin_user)\n\n # Verify the response indicates the role changed to BASIC\n assert user_data[\"role\"] == UserRole.BASIC.value\n\n # Verify the user's role was changed in the database\n assert UserManager.is_role(slack_user, UserRole.BASIC)\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"SAML tests are enterprise only\",\n)\ndef test_saml_user_conversion_sets_account_type_and_group(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test that SAML login sets account_type to STANDARD when converting a\n non-web user (EXT_PERM_USER) and that the user receives the correct role\n (BASIC) after conversion.\n\n This validates the permissions-migration-phase2 changes which ensure that:\n 1. account_type is updated to 'standard' on SAML conversion\n 2. The converted user is assigned to the Basic default group\n \"\"\"\n # Create an admin user (first user is automatically admin)\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n # Create a user and set them as EXT_PERM_USER\n test_email = \"ext_convert@example.com\"\n test_user = UserManager.create(email=test_email)\n UserManager.set_role(\n user_to_set=test_user,\n target_role=UserRole.EXT_PERM_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert UserManager.is_role(test_user, UserRole.EXT_PERM_USER)\n\n # Simulate SAML login\n user_data = _simulate_saml_login(test_email, admin_user)\n\n # Verify account_type is set to standard after conversion\n assert (\n user_data[\"account_type\"] == AccountType.STANDARD.value\n ), f\"Expected account_type='{AccountType.STANDARD.value}', got '{user_data['account_type']}'\"\n\n # Verify role is BASIC after conversion\n assert user_data[\"role\"] == UserRole.BASIC.value\n\n # Verify the user was assigned to the Basic default group\n assert test_email in _get_basic_group_member_emails(\n admin_user\n ), f\"Converted user '{test_email}' not found in Basic default group\"\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"SAML tests are enterprise only\",\n)\ndef test_saml_normal_signin_assigns_group(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test that a brand-new user signing in via SAML for the first time\n is created with the correct role, account_type, and group membership.\n\n This validates that normal SAML sign-in (not an upgrade from\n SLACK_USER/EXT_PERM_USER) correctly:\n 1. Creates the user with role=BASIC and account_type=STANDARD\n 2. Assigns the user to the Basic default group\n \"\"\"\n # First user becomes admin\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n # New user signs in via SAML (no prior account)\n new_email = \"new_saml_user@example.com\"\n user_data = _simulate_saml_login(new_email, admin_user)\n\n # Verify role and account_type\n assert user_data[\"role\"] == UserRole.BASIC.value\n assert user_data[\"account_type\"] == AccountType.STANDARD.value\n\n # Verify user is in the Basic default group\n assert new_email in _get_basic_group_member_emails(\n admin_user\n ), f\"New SAML user '{new_email}' not found in Basic default group\"\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"SAML tests are enterprise only\",\n)\ndef test_saml_user_conversion_restores_group_membership(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test that SAML login restores Basic group membership when converting\n a non-authenticated user (EXT_PERM_USER or SLACK_USER) to BASIC.\n\n Group membership implies 'basic' permission (verified by\n test_new_group_gets_basic_permission).\n \"\"\"\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n # --- EXT_PERM_USER path ---\n ext_email = \"ext_perm_perms@example.com\"\n ext_user = UserManager.create(email=ext_email)\n assert ext_email in _get_basic_group_member_emails(admin_user)\n\n UserManager.set_role(\n user_to_set=ext_user,\n target_role=UserRole.EXT_PERM_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert ext_email not in _get_basic_group_member_emails(admin_user)\n\n user_data = _simulate_saml_login(ext_email, admin_user)\n assert user_data[\"role\"] == UserRole.BASIC.value\n assert ext_email in _get_basic_group_member_emails(\n admin_user\n ), \"EXT_PERM_USER should be back in Basic group after SAML conversion\"\n\n # --- SLACK_USER path ---\n slack_email = \"slack_perms@example.com\"\n slack_user = UserManager.create(email=slack_email)\n\n UserManager.set_role(\n user_to_set=slack_user,\n target_role=UserRole.SLACK_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert slack_email not in _get_basic_group_member_emails(admin_user)\n\n user_data = _simulate_saml_login(slack_email, admin_user)\n assert user_data[\"role\"] == UserRole.BASIC.value\n assert slack_email in _get_basic_group_member_emails(\n admin_user\n ), \"SLACK_USER should be back in Basic group after SAML conversion\"\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"SAML tests are enterprise only\",\n)\ndef test_saml_round_trip_group_lifecycle(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test the full round-trip: BASIC -> EXT_PERM -> SAML(BASIC) -> EXT_PERM -> SAML(BASIC).\n\n Verifies group membership is correctly removed and restored at each transition.\n \"\"\"\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n test_email = \"roundtrip@example.com\"\n test_user = UserManager.create(email=test_email)\n\n # Step 1: BASIC user is in Basic group\n assert test_email in _get_basic_group_member_emails(admin_user)\n\n # Step 2: Downgrade to EXT_PERM_USER — loses Basic group\n UserManager.set_role(\n user_to_set=test_user,\n target_role=UserRole.EXT_PERM_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert test_email not in _get_basic_group_member_emails(admin_user)\n\n # Step 3: SAML login — converts back to BASIC, regains Basic group\n _simulate_saml_login(test_email, admin_user)\n assert test_email in _get_basic_group_member_emails(\n admin_user\n ), \"Should be in Basic group after first SAML conversion\"\n\n # Step 4: Downgrade again\n UserManager.set_role(\n user_to_set=test_user,\n target_role=UserRole.EXT_PERM_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert test_email not in _get_basic_group_member_emails(admin_user)\n\n # Step 5: SAML login again — should still restore correctly\n _simulate_saml_login(test_email, admin_user)\n assert test_email in _get_basic_group_member_emails(\n admin_user\n ), \"Should be in Basic group after second SAML conversion\"\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"SAML tests are enterprise only\",\n)\ndef test_saml_slack_user_conversion_sets_account_type_and_group(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test that SAML login sets account_type to STANDARD and assigns Basic group\n when converting a SLACK_USER (BOT account_type).\n\n Mirrors test_saml_user_conversion_sets_account_type_and_group but for\n SLACK_USER instead of EXT_PERM_USER, and additionally verifies permissions.\n \"\"\"\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n test_email = \"slack_convert@example.com\"\n test_user = UserManager.create(email=test_email)\n\n UserManager.set_role(\n user_to_set=test_user,\n target_role=UserRole.SLACK_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert UserManager.is_role(test_user, UserRole.SLACK_USER)\n\n # SAML login\n user_data = _simulate_saml_login(test_email, admin_user)\n\n # Verify account_type and role\n assert (\n user_data[\"account_type\"] == AccountType.STANDARD.value\n ), f\"Expected STANDARD, got {user_data['account_type']}\"\n assert user_data[\"role\"] == UserRole.BASIC.value\n\n # Verify Basic group membership (implies 'basic' permission)\n assert test_email in _get_basic_group_member_emails(\n admin_user\n ), f\"Converted SLACK_USER '{test_email}' not found in Basic default group\"\n" }, { "path": "backend/tests/integration/tests/chat/test_chat_deletion.py", "content": "import pytest\n\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nMESSAGE = \"Hi\"\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef reset_for_module() -> None:\n \"\"\"Reset all data once before running any tests in this module.\"\"\"\n reset_all()\n\n\n@pytest.fixture\ndef llm_provider(admin_user: DATestUser) -> DATestLLMProvider:\n return LLMProviderManager.create(user_performing_action=admin_user)\n\n\ndef test_soft_delete_chat_session(\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"\n Test soft deletion of a chat session.\n Soft delete should mark the chat as deleted but keep it in the database.\n \"\"\"\n # Create a chat session\n test_chat_session = ChatSessionManager.create(\n persona_id=0, # Use default persona\n description=\"Test chat session for soft deletion\",\n user_performing_action=basic_user,\n )\n\n # Send a message to create some data\n response = ChatSessionManager.send_message(\n chat_session_id=test_chat_session.id,\n message=MESSAGE,\n user_performing_action=basic_user,\n )\n\n # Verify that the message was processed successfully\n assert response.error is None, \"Chat response should not have an error\"\n assert len(response.full_message) > 0, \"Chat response should not be empty\"\n\n # Verify that the chat session can be retrieved before deletion\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n )\n assert len(chat_history) > 0, \"Chat session should have messages\"\n\n # Test soft deletion of the chat session\n deletion_success = ChatSessionManager.soft_delete(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n )\n\n # Verify that the deletion was successful\n assert deletion_success, \"Chat session soft deletion should succeed\"\n\n # Verify that the chat session is soft deleted (marked as deleted but still in DB)\n assert ChatSessionManager.verify_soft_deleted(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n ), \"Chat session should be soft deleted\"\n\n # Verify that normal access is blocked\n assert ChatSessionManager.verify_deleted(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n ), \"Chat session should not be accessible normally after soft delete\"\n\n\ndef test_hard_delete_chat_session(\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"\n Test hard deletion of a chat session.\n Hard delete should completely remove the chat from the database.\n \"\"\"\n # Create a chat session\n test_chat_session = ChatSessionManager.create(\n persona_id=0, # Use default persona\n description=\"Test chat session for hard deletion\",\n user_performing_action=basic_user,\n )\n\n # Send a message to create some data\n response = ChatSessionManager.send_message(\n chat_session_id=test_chat_session.id,\n message=MESSAGE,\n user_performing_action=basic_user,\n )\n\n # Verify that the message was processed successfully\n assert response.error is None, \"Chat response should not have an error\"\n assert len(response.full_message) > 0, \"Chat response should not be empty\"\n\n # Verify that the chat session can be retrieved before deletion\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n )\n assert len(chat_history) > 0, \"Chat session should have messages\"\n\n # Test hard deletion of the chat session\n deletion_success = ChatSessionManager.hard_delete(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n )\n\n # Verify that the deletion was successful\n assert deletion_success, \"Chat session hard deletion should succeed\"\n\n # Verify that the chat session is hard deleted (completely removed from DB)\n assert ChatSessionManager.verify_hard_deleted(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n ), \"Chat session should be hard deleted\"\n\n # Verify that the chat session is not accessible at all\n assert ChatSessionManager.verify_deleted(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n ), \"Chat session should not be accessible after hard delete\"\n\n # Verify it's not soft deleted (since it doesn't exist at all)\n assert not ChatSessionManager.verify_soft_deleted(\n chat_session=test_chat_session,\n user_performing_action=basic_user,\n ), \"Hard deleted chat should not be found as soft deleted\"\n\n\ndef test_multiple_soft_deletions(\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"\n Test multiple chat session soft deletions to ensure proper handling\n when there are multiple related records.\n \"\"\"\n chat_sessions = []\n\n # Create multiple chat sessions with potential agent behavior\n for i in range(3):\n chat_session = ChatSessionManager.create(\n persona_id=0,\n description=f\"Test chat session {i} for multi-soft-deletion\",\n user_performing_action=basic_user,\n )\n\n # Send a message to create some data\n ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=f\"Tell me about topic {i} with detailed analysis\",\n user_performing_action=basic_user,\n )\n\n chat_sessions.append(chat_session)\n\n # Soft delete all chat sessions\n for chat_session in chat_sessions:\n deletion_success = ChatSessionManager.soft_delete(\n chat_session=chat_session,\n user_performing_action=basic_user,\n )\n assert deletion_success, f\"Failed to soft delete chat {chat_session.id}\"\n\n # Verify all chat sessions are soft deleted\n for chat_session in chat_sessions:\n assert ChatSessionManager.verify_soft_deleted(\n chat_session=chat_session,\n user_performing_action=basic_user,\n ), f\"Chat {chat_session.id} should be soft deleted\"\n\n assert ChatSessionManager.verify_deleted(\n chat_session=chat_session,\n user_performing_action=basic_user,\n ), f\"Chat {chat_session.id} should not be accessible normally\"\n\n\ndef test_multiple_hard_deletions_with_agent_data(\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"\n Test multiple chat session hard deletions to ensure CASCADE deletes work correctly\n when there are multiple related records.\n \"\"\"\n chat_sessions = []\n\n # Create multiple chat sessions with potential agent behavior\n for i in range(3):\n chat_session = ChatSessionManager.create(\n persona_id=0,\n description=f\"Test chat session {i} for multi-hard-deletion\",\n user_performing_action=basic_user,\n )\n\n # Send a message to create some data\n ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=f\"Tell me about topic {i} with detailed analysis\",\n user_performing_action=basic_user,\n )\n\n chat_sessions.append(chat_session)\n\n # Hard delete all chat sessions\n for chat_session in chat_sessions:\n deletion_success = ChatSessionManager.hard_delete(\n chat_session=chat_session,\n user_performing_action=basic_user,\n )\n assert deletion_success, f\"Failed to hard delete chat {chat_session.id}\"\n\n # Verify all chat sessions are hard deleted\n for chat_session in chat_sessions:\n assert ChatSessionManager.verify_hard_deleted(\n chat_session=chat_session,\n user_performing_action=basic_user,\n ), f\"Chat {chat_session.id} should be hard deleted\"\n\n assert ChatSessionManager.verify_deleted(\n chat_session=chat_session,\n user_performing_action=basic_user,\n ), f\"Chat {chat_session.id} should not be accessible\"\n\n\ndef test_soft_vs_hard_delete_edge_cases(\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"\n Test edge cases for both soft and hard deletion to ensure robustness.\n \"\"\"\n # Test 1: Soft delete a chat session with no messages\n empty_chat_session_soft = ChatSessionManager.create(\n persona_id=0,\n description=\"Empty chat session for soft delete\",\n user_performing_action=basic_user,\n )\n\n # Soft delete without sending any messages\n deletion_success = ChatSessionManager.soft_delete(\n chat_session=empty_chat_session_soft,\n user_performing_action=basic_user,\n )\n assert deletion_success, \"Empty chat session should be soft deletable\"\n assert ChatSessionManager.verify_soft_deleted(\n chat_session=empty_chat_session_soft,\n user_performing_action=basic_user,\n ), \"Empty chat session should be confirmed as soft deleted\"\n\n # Test 2: Hard delete a chat session with no messages\n empty_chat_session_hard = ChatSessionManager.create(\n persona_id=0,\n description=\"Empty chat session for hard delete\",\n user_performing_action=basic_user,\n )\n\n # Hard delete without sending any messages\n deletion_success = ChatSessionManager.hard_delete(\n chat_session=empty_chat_session_hard,\n user_performing_action=basic_user,\n )\n assert deletion_success, \"Empty chat session should be hard deletable\"\n assert ChatSessionManager.verify_hard_deleted(\n chat_session=empty_chat_session_hard,\n user_performing_action=basic_user,\n ), \"Empty chat session should be confirmed as hard deleted\"\n\n # Test 3: Soft delete a chat session with multiple messages\n multi_message_chat_soft = ChatSessionManager.create(\n persona_id=0,\n description=\"Multi-message chat session for soft delete\",\n user_performing_action=basic_user,\n )\n\n # Send multiple messages to create more complex data\n for i in range(3):\n ChatSessionManager.send_message(\n chat_session_id=multi_message_chat_soft.id,\n message=f\"Message {i}: Tell me about different aspects of this topic\",\n user_performing_action=basic_user,\n )\n\n # Verify messages exist\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=multi_message_chat_soft,\n user_performing_action=basic_user,\n )\n assert len(chat_history) >= 3, \"Chat should have multiple messages\"\n\n # Soft delete the chat with multiple messages\n deletion_success = ChatSessionManager.soft_delete(\n chat_session=multi_message_chat_soft,\n user_performing_action=basic_user,\n )\n assert deletion_success, \"Multi-message chat session should be soft deletable\"\n assert ChatSessionManager.verify_soft_deleted(\n chat_session=multi_message_chat_soft,\n user_performing_action=basic_user,\n ), \"Multi-message chat session should be confirmed as soft deleted\"\n\n # Test 4: Hard delete a chat session with multiple messages\n multi_message_chat_hard = ChatSessionManager.create(\n persona_id=0,\n description=\"Multi-message chat session for hard delete\",\n user_performing_action=basic_user,\n )\n\n # Send multiple messages to create more complex data\n for i in range(3):\n ChatSessionManager.send_message(\n chat_session_id=multi_message_chat_hard.id,\n message=f\"Message {i}: Tell me about different aspects of this topic\",\n user_performing_action=basic_user,\n )\n\n # Verify messages exist\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=multi_message_chat_hard,\n user_performing_action=basic_user,\n )\n assert len(chat_history) >= 3, \"Chat should have multiple messages\"\n\n # Hard delete the chat with multiple messages\n deletion_success = ChatSessionManager.hard_delete(\n chat_session=multi_message_chat_hard,\n user_performing_action=basic_user,\n )\n assert deletion_success, \"Multi-message chat session should be hard deletable\"\n assert ChatSessionManager.verify_hard_deleted(\n chat_session=multi_message_chat_hard,\n user_performing_action=basic_user,\n ), \"Multi-message chat session should be confirmed as hard deleted\"\n" }, { "path": "backend/tests/integration/tests/chat/test_chat_session_access.py", "content": "from uuid import uuid4\n\nimport pytest\nimport requests\nfrom requests import HTTPError\n\nfrom onyx.auth.schemas import UserRole\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.user import build_email\nfrom tests.integration.common_utils.managers.user import DEFAULT_PASSWORD\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef reset_for_module() -> None:\n \"\"\"Reset all data once before running any tests in this module.\"\"\"\n reset_all()\n\n\n@pytest.fixture\ndef second_user(admin_user: DATestUser) -> DATestUser: # noqa: ARG001\n # Ensure admin exists so this new user is created with BASIC role.\n try:\n return UserManager.create(name=\"second_basic_user\")\n except HTTPError as e:\n response = e.response\n if response is None:\n raise\n if response.status_code not in (400, 409):\n raise\n try:\n payload = response.json()\n except ValueError:\n raise\n detail = payload.get(\"detail\")\n if not _is_user_already_exists_detail(detail):\n raise\n print(\"Second basic user already exists; logging in instead.\")\n return UserManager.login_as_user(\n DATestUser(\n id=\"\",\n email=build_email(\"second_basic_user\"),\n password=DEFAULT_PASSWORD,\n headers=GENERAL_HEADERS,\n role=UserRole.BASIC,\n is_active=True,\n )\n )\n\n\ndef _is_user_already_exists_detail(detail: object) -> bool:\n if isinstance(detail, str):\n normalized = detail.lower()\n return (\n \"already exists\" in normalized\n or \"register_user_already_exists\" in normalized\n )\n if isinstance(detail, dict):\n code = detail.get(\"code\")\n if isinstance(code, str) and code.lower() == \"register_user_already_exists\":\n return True\n message = detail.get(\"message\")\n if isinstance(message, str) and \"already exists\" in message.lower():\n return True\n return False\n\n\ndef _get_chat_session(\n chat_session_id: str,\n user: DATestUser,\n is_shared: bool | None = None,\n include_deleted: bool | None = None,\n) -> requests.Response:\n params: dict[str, str] = {}\n if is_shared is not None:\n params[\"is_shared\"] = str(is_shared).lower()\n if include_deleted is not None:\n params[\"include_deleted\"] = str(include_deleted).lower()\n\n return requests.get(\n f\"{API_SERVER_URL}/chat/get-chat-session/{chat_session_id}\",\n params=params,\n headers=user.headers,\n cookies=user.cookies,\n )\n\n\ndef _set_sharing_status(\n chat_session_id: str, sharing_status: str, user: DATestUser\n) -> requests.Response:\n return requests.patch(\n f\"{API_SERVER_URL}/chat/chat-session/{chat_session_id}\",\n json={\"sharing_status\": sharing_status},\n headers=user.headers,\n cookies=user.cookies,\n )\n\n\ndef test_private_chat_session_access(\n basic_user: DATestUser, second_user: DATestUser\n) -> None:\n \"\"\"Verify private sessions are only accessible by the owner and never via share link.\"\"\"\n # Create a private chat session owned by basic_user.\n chat_session = ChatSessionManager.create(user_performing_action=basic_user)\n\n # Owner can access the private session normally.\n response = _get_chat_session(str(chat_session.id), basic_user)\n assert response.status_code == 200\n\n # Share link should be forbidden when the session is private.\n response = _get_chat_session(str(chat_session.id), basic_user, is_shared=True)\n assert response.status_code == 403\n\n # Other users cannot access private sessions directly.\n response = _get_chat_session(str(chat_session.id), second_user)\n assert response.status_code == 403\n\n # Other users also cannot access private sessions via share link.\n response = _get_chat_session(str(chat_session.id), second_user, is_shared=True)\n assert response.status_code == 403\n\n\ndef test_public_shared_chat_session_access(\n basic_user: DATestUser, second_user: DATestUser\n) -> None:\n \"\"\"Verify shared sessions are accessible only via share link for non-owners.\"\"\"\n # Create a private session, then mark it public.\n chat_session = ChatSessionManager.create(user_performing_action=basic_user)\n\n response = _set_sharing_status(str(chat_session.id), \"public\", basic_user)\n assert response.status_code == 200\n\n # Owner can access normally.\n response = _get_chat_session(str(chat_session.id), basic_user)\n assert response.status_code == 200\n\n # Owner can also access via share link.\n response = _get_chat_session(str(chat_session.id), basic_user, is_shared=True)\n assert response.status_code == 200\n\n # Non-owner cannot access without share link.\n response = _get_chat_session(str(chat_session.id), second_user)\n assert response.status_code == 403\n\n # Non-owner can access with share link for public sessions.\n response = _get_chat_session(str(chat_session.id), second_user, is_shared=True)\n assert response.status_code == 200\n\n\ndef test_deleted_chat_session_access(\n basic_user: DATestUser, second_user: DATestUser\n) -> None:\n \"\"\"Verify deleted sessions return 404, with include_deleted gated by access checks.\"\"\"\n # Create and soft-delete a session.\n chat_session = ChatSessionManager.create(user_performing_action=basic_user)\n\n deletion_success = ChatSessionManager.soft_delete(\n chat_session=chat_session, user_performing_action=basic_user\n )\n assert deletion_success is True\n\n # Deleted sessions are not accessible normally.\n response = _get_chat_session(str(chat_session.id), basic_user)\n assert response.status_code == 404\n\n # Owner can fetch deleted session only with include_deleted.\n response = _get_chat_session(str(chat_session.id), basic_user, include_deleted=True)\n assert response.status_code == 200\n assert response.json().get(\"deleted\") is True\n\n # Non-owner should be blocked even with include_deleted.\n response = _get_chat_session(\n str(chat_session.id), second_user, include_deleted=True\n )\n assert response.status_code == 403\n\n\ndef test_chat_session_not_found_returns_404(basic_user: DATestUser) -> None:\n \"\"\"Verify unknown IDs return 404.\"\"\"\n response = _get_chat_session(str(uuid4()), basic_user)\n assert response.status_code == 404\n" }, { "path": "backend/tests/integration/tests/chat_retention/test_chat_retention.py", "content": "import os\nimport time\n\nimport pytest\nimport requests\n\nfrom onyx.db.chat import delete_chat_session\nfrom onyx.db.chat import get_chat_sessions_older_than\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.settings import SettingsManager\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestSettings\nfrom tests.integration.common_utils.test_models import DATestUser\n\nRETENTION_SECONDS = 10\n\n\ndef _run_ttl_cleanup(retention_days: int) -> None:\n \"\"\"Directly execute TTL cleanup logic, bypassing Celery task infrastructure.\"\"\"\n with get_session_with_current_tenant() as db_session:\n old_chat_sessions = get_chat_sessions_older_than(retention_days, db_session)\n\n for user_id, session_id in old_chat_sessions:\n with get_session_with_current_tenant() as db_session:\n delete_chat_session(\n user_id,\n session_id,\n db_session,\n include_deleted=True,\n hard_delete=True,\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Chat retention tests are enterprise only\",\n)\ndef test_chat_retention(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"Test that chat sessions are deleted after the retention period expires.\"\"\"\n\n retention_days = RETENTION_SECONDS // 86400\n settings = DATestSettings(maximum_chat_retention_days=retention_days)\n SettingsManager.update_settings(settings, user_performing_action=admin_user)\n\n chat_session = ChatSessionManager.create(\n persona_id=0,\n description=\"Test chat retention\",\n user_performing_action=admin_user,\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"This message should be deleted soon\",\n user_performing_action=admin_user,\n )\n assert (\n response.error is None\n ), f\"Chat response should not have an error: {response.error}\"\n\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=chat_session,\n user_performing_action=admin_user,\n )\n assert len(chat_history) > 0, \"Chat session should have messages\"\n\n # Wait for the retention period to elapse, then directly run TTL cleanup\n time.sleep(RETENTION_SECONDS + 2)\n _run_ttl_cleanup(retention_days)\n\n # Verify the chat session was deleted\n session_deleted = False\n try:\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=chat_session,\n user_performing_action=admin_user,\n )\n session_deleted = len(chat_history) == 0\n except requests.exceptions.HTTPError as e:\n if e.response.status_code in (404, 400):\n session_deleted = True\n else:\n raise\n\n assert session_deleted, \"Chat session was not deleted after retention period\"\n" }, { "path": "backend/tests/integration/tests/code_interpreter/conftest.py", "content": "from collections.abc import Generator\n\nimport pytest\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\nCODE_INTERPRETER_URL = f\"{API_SERVER_URL}/admin/code-interpreter\"\n\n\n@pytest.fixture\ndef preserve_code_interpreter_state(\n admin_user: DATestUser,\n) -> Generator[None, None, None]:\n \"\"\"Capture the code interpreter enabled state before a test and restore it\n afterwards, so that tests that toggle the setting cannot leak state.\"\"\"\n response = requests.get(\n CODE_INTERPRETER_URL,\n headers=admin_user.headers,\n )\n response.raise_for_status()\n initial_enabled = response.json()[\"enabled\"]\n\n yield\n\n restore = requests.put(\n CODE_INTERPRETER_URL,\n json={\"enabled\": initial_enabled},\n headers=admin_user.headers,\n )\n restore.raise_for_status()\n" }, { "path": "backend/tests/integration/tests/code_interpreter/test_code_interpreter_api.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\nCODE_INTERPRETER_URL = f\"{API_SERVER_URL}/admin/code-interpreter\"\nCODE_INTERPRETER_HEALTH_URL = f\"{CODE_INTERPRETER_URL}/health\"\n\n\ndef test_get_code_interpreter_health_as_admin(\n admin_user: DATestUser,\n) -> None:\n \"\"\"Health endpoint should return a JSON object with a 'healthy' boolean.\"\"\"\n response = requests.get(\n CODE_INTERPRETER_HEALTH_URL,\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n data = response.json()\n assert \"healthy\" in data\n assert isinstance(data[\"healthy\"], bool)\n\n\ndef test_get_code_interpreter_status_as_admin(\n admin_user: DATestUser,\n) -> None:\n \"\"\"GET endpoint should return a JSON object with an 'enabled' boolean.\"\"\"\n response = requests.get(\n CODE_INTERPRETER_URL,\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n data = response.json()\n assert \"enabled\" in data\n assert isinstance(data[\"enabled\"], bool)\n\n\ndef test_update_code_interpreter_disable_and_enable(\n admin_user: DATestUser,\n preserve_code_interpreter_state: None, # noqa: ARG001\n) -> None:\n \"\"\"PUT endpoint should update the enabled flag and persist across reads.\"\"\"\n # Disable\n response = requests.put(\n CODE_INTERPRETER_URL,\n json={\"enabled\": False},\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n # Verify disabled\n response = requests.get(\n CODE_INTERPRETER_URL,\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n assert response.json()[\"enabled\"] is False\n\n # Re-enable\n response = requests.put(\n CODE_INTERPRETER_URL,\n json={\"enabled\": True},\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n # Verify enabled\n response = requests.get(\n CODE_INTERPRETER_URL,\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n assert response.json()[\"enabled\"] is True\n\n\ndef test_code_interpreter_endpoints_require_admin(\n basic_user: DATestUser,\n) -> None:\n \"\"\"All code interpreter endpoints should reject non-admin users.\"\"\"\n health_response = requests.get(\n CODE_INTERPRETER_HEALTH_URL,\n headers=basic_user.headers,\n )\n assert health_response.status_code == 403\n\n get_response = requests.get(\n CODE_INTERPRETER_URL,\n headers=basic_user.headers,\n )\n assert get_response.status_code == 403\n\n put_response = requests.put(\n CODE_INTERPRETER_URL,\n json={\"enabled\": True},\n headers=basic_user.headers,\n )\n assert put_response.status_code == 403\n" }, { "path": "backend/tests/integration/tests/connector/test_connector_creation.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\n\nfrom onyx.connectors.models import InputType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_connector_creation(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # create connectors\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n cc_pair_info = CCPairManager.get_single(\n cc_pair_1.id, user_performing_action=admin_user\n )\n assert cc_pair_info\n assert cc_pair_info.creator\n assert str(cc_pair_info.creator) == admin_user.id\n assert cc_pair_info.creator_email == admin_user.email\n\n\ndef test_overlapping_connector_creation(reset: None) -> None: # noqa: ARG001\n \"\"\"Tests that connectors indexing the same documents don't interfere with each other.\n A previous bug involved document by cc pair entries not being added for new connectors\n when the docs existed already via another connector and were up to date relative to the source.\n \"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n config = {\n \"wiki_base\": os.environ[\"CONFLUENCE_TEST_SPACE_URL\"],\n \"space\": \"DailyConne\",\n \"is_cloud\": True,\n }\n\n credential = {\n \"confluence_username\": os.environ[\"CONFLUENCE_USER_NAME\"],\n \"confluence_access_token\": os.environ[\"CONFLUENCE_ACCESS_TOKEN\"],\n }\n\n # store the time before we create the connector so that we know after\n # when the indexing should have started\n now = datetime.now(timezone.utc)\n\n # create connector\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.CONFLUENCE,\n connector_specific_config=config,\n credential_json=credential,\n user_performing_action=admin_user,\n input_type=InputType.POLL,\n )\n\n CCPairManager.wait_for_indexing_completion(\n cc_pair_1, now, timeout=300, user_performing_action=admin_user\n )\n\n now = datetime.now(timezone.utc)\n\n cc_pair_2 = CCPairManager.create_from_scratch(\n source=DocumentSource.CONFLUENCE,\n connector_specific_config=config,\n credential_json=credential,\n user_performing_action=admin_user,\n input_type=InputType.POLL,\n )\n\n CCPairManager.wait_for_indexing_completion(\n cc_pair_2, now, timeout=300, user_performing_action=admin_user\n )\n\n info_1 = CCPairManager.get_single(cc_pair_1.id, user_performing_action=admin_user)\n assert info_1\n\n info_2 = CCPairManager.get_single(cc_pair_2.id, user_performing_action=admin_user)\n assert info_2\n\n assert info_1.num_docs_indexed == info_2.num_docs_indexed\n\n\ndef test_connector_pause_while_indexing(reset: None) -> None: # noqa: ARG001\n \"\"\"Tests that we can pause a connector while indexing is in progress and that\n tasks end early or abort as a result.\n\n TODO: This does not specifically test for soft or hard termination code paths.\n Design specific tests for those use cases.\n \"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n config = {\n \"wiki_base\": os.environ[\"CONFLUENCE_TEST_SPACE_URL\"],\n \"is_cloud\": True,\n }\n\n credential = {\n \"confluence_username\": os.environ[\"CONFLUENCE_USER_NAME\"],\n \"confluence_access_token\": os.environ[\"CONFLUENCE_ACCESS_TOKEN\"],\n }\n\n # store the time before we create the connector so that we know after\n # when the indexing should have started\n datetime.now(timezone.utc)\n\n # create connector\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.CONFLUENCE,\n connector_specific_config=config,\n credential_json=credential,\n user_performing_action=admin_user,\n input_type=InputType.POLL,\n )\n\n # NOTE: A bit flaky in our CI due to varying indexing times. Empirically\n # 120s was not always enough to index 16 docs from Confluence so trying to\n # bump down the indexing progress to wait for to 4 docs from 16.\n CCPairManager.wait_for_indexing_in_progress(\n cc_pair_1, timeout=120, num_docs=4, user_performing_action=admin_user\n )\n\n CCPairManager.pause_cc_pair(cc_pair_1, user_performing_action=admin_user)\n\n CCPairManager.wait_for_indexing_inactive(\n cc_pair_1, timeout=60, user_performing_action=admin_user\n )\n return\n" }, { "path": "backend/tests/integration/tests/connector/test_connector_deletion.py", "content": "\"\"\"\nThis file contains tests for the following:\n- Ensuring deletion of a connector also:\n - deletes the documents in vespa for that connector\n - updates the document sets and user groups to remove the connector\n- Ensure that deleting a connector that is part of an overlapping document set and/or user group works as expected\n\"\"\"\n\nimport os\nfrom uuid import uuid4\n\nfrom sqlalchemy.orm import Session\n\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import DocumentFailure\nfrom onyx.db.engine.sql_engine import get_sqlalchemy_engine\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.index_attempt import create_index_attempt\nfrom onyx.db.index_attempt import create_index_attempt_error\nfrom onyx.db.models import IndexAttempt\nfrom onyx.db.search_settings import get_current_search_settings\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.constants import NUM_DOCS\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.document_set import DocumentSetManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import DATestUserGroup\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef test_connector_deletion(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture,\n) -> None:\n user_group_1: DATestUserGroup\n user_group_2: DATestUserGroup\n\n is_ee = (\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() == \"true\"\n )\n\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n # create api key\n api_key: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user,\n )\n\n # create connectors\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n cc_pair_2 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n # seed documents\n cc_pair_1.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_1,\n num_docs=NUM_DOCS,\n api_key=api_key,\n )\n cc_pair_2.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_2,\n num_docs=NUM_DOCS,\n api_key=api_key,\n )\n\n # create document sets\n doc_set_1 = DocumentSetManager.create(\n name=\"Test Document Set 1\",\n cc_pair_ids=[cc_pair_1.id],\n user_performing_action=admin_user,\n )\n doc_set_2 = DocumentSetManager.create(\n name=\"Test Document Set 2\",\n cc_pair_ids=[cc_pair_1.id, cc_pair_2.id],\n user_performing_action=admin_user,\n )\n\n # wait for document sets to be synced\n DocumentSetManager.wait_for_sync(user_performing_action=admin_user)\n\n print(\"Document sets created and synced\")\n\n if is_ee:\n # create user groups\n user_group_1 = UserGroupManager.create(\n cc_pair_ids=[cc_pair_1.id],\n user_performing_action=admin_user,\n )\n user_group_2 = UserGroupManager.create(\n cc_pair_ids=[cc_pair_1.id, cc_pair_2.id],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(user_performing_action=admin_user)\n\n # inject a finished index attempt and index attempt error (exercises foreign key errors)\n with Session(get_sqlalchemy_engine()) as db_session:\n primary_search_settings = get_current_search_settings(db_session)\n new_attempt = IndexAttempt(\n connector_credential_pair_id=cc_pair_1.id,\n search_settings_id=primary_search_settings.id,\n from_beginning=False,\n status=IndexingStatus.COMPLETED_WITH_ERRORS,\n )\n db_session.add(new_attempt)\n db_session.commit()\n\n create_index_attempt_error(\n index_attempt_id=new_attempt.id,\n connector_credential_pair_id=cc_pair_1.id,\n failure=ConnectorFailure(\n failure_message=\"Test error\",\n failed_document=DocumentFailure(\n document_id=cc_pair_1.documents[0].id,\n document_link=None,\n ),\n failed_entity=None,\n ),\n db_session=db_session,\n )\n\n # delete connector 1\n CCPairManager.pause_cc_pair(\n cc_pair=cc_pair_1,\n user_performing_action=admin_user,\n )\n CCPairManager.delete(\n cc_pair=cc_pair_1,\n user_performing_action=admin_user,\n )\n\n # inject an index attempt and index attempt error (exercises foreign key errors)\n with Session(get_sqlalchemy_engine()) as db_session:\n attempt_id = create_index_attempt(\n connector_credential_pair_id=cc_pair_1.id,\n search_settings_id=1,\n db_session=db_session,\n )\n create_index_attempt_error(\n index_attempt_id=attempt_id,\n connector_credential_pair_id=cc_pair_1.id,\n failure=ConnectorFailure(\n failure_message=\"Test error\",\n failed_document=DocumentFailure(\n document_id=cc_pair_1.documents[0].id,\n document_link=None,\n ),\n failed_entity=None,\n ),\n db_session=db_session,\n )\n\n # Update local records to match the database for later comparison\n doc_set_1.cc_pair_ids = []\n doc_set_2.cc_pair_ids = [cc_pair_2.id]\n cc_pair_1.groups = []\n if is_ee:\n cc_pair_2.groups = [user_group_2.id]\n else:\n cc_pair_2.groups = []\n\n CCPairManager.wait_for_deletion_completion(\n cc_pair_id=cc_pair_1.id, user_performing_action=admin_user\n )\n\n # validate vespa documents\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n doc_set_names=[],\n group_names=[],\n doc_creating_user=admin_user,\n verify_deleted=True,\n )\n\n cc_pair_2_group_name_expected = []\n if is_ee:\n cc_pair_2_group_name_expected = [user_group_2.name]\n\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n doc_set_names=[doc_set_2.name],\n group_names=cc_pair_2_group_name_expected,\n doc_creating_user=admin_user,\n verify_deleted=False,\n )\n\n # check that only connector 1 is deleted\n CCPairManager.verify(\n cc_pair=cc_pair_2,\n user_performing_action=admin_user,\n )\n\n # validate document sets\n DocumentSetManager.verify(\n document_set=doc_set_1,\n user_performing_action=admin_user,\n )\n DocumentSetManager.verify(\n document_set=doc_set_2,\n user_performing_action=admin_user,\n )\n\n if is_ee:\n user_group_1.cc_pair_ids = []\n user_group_2.cc_pair_ids = [cc_pair_2.id]\n\n # validate user groups\n UserGroupManager.verify(\n user_group=user_group_1,\n user_performing_action=admin_user,\n )\n UserGroupManager.verify(\n user_group=user_group_2,\n user_performing_action=admin_user,\n )\n\n\ndef test_connector_deletion_for_overlapping_connectors(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture,\n) -> None:\n \"\"\"Checks to make sure that connectors with overlapping documents work properly. Specifically, that the overlapping\n document (1) still exists and (2) has the right document set / group post-deletion of one of the connectors.\n \"\"\"\n user_group_1: DATestUserGroup\n user_group_2: DATestUserGroup\n\n is_ee = (\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() == \"true\"\n )\n\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n # create api key\n api_key: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user,\n )\n\n # create connectors\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n cc_pair_2 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n doc_ids = [str(uuid4())]\n cc_pair_1.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_1,\n document_ids=doc_ids,\n api_key=api_key,\n )\n cc_pair_2.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_2,\n document_ids=doc_ids,\n api_key=api_key,\n )\n\n # verify vespa document exists and that it is not in any document sets or groups\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n doc_set_names=[],\n group_names=[],\n doc_creating_user=admin_user,\n )\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n doc_set_names=[],\n group_names=[],\n doc_creating_user=admin_user,\n )\n\n # create document set\n doc_set_1 = DocumentSetManager.create(\n name=\"Test Document Set 1\",\n cc_pair_ids=[cc_pair_1.id],\n user_performing_action=admin_user,\n )\n DocumentSetManager.wait_for_sync(\n document_sets_to_check=[doc_set_1],\n user_performing_action=admin_user,\n )\n\n print(\"Document set 1 created and synced\")\n\n # verify vespa document is in the document set\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n doc_set_names=[doc_set_1.name],\n doc_creating_user=admin_user,\n )\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n doc_creating_user=admin_user,\n )\n\n if is_ee:\n # create a user group and attach it to connector 1\n user_group_1 = UserGroupManager.create(\n name=\"Test User Group 1\",\n cc_pair_ids=[cc_pair_1.id],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1],\n user_performing_action=admin_user,\n )\n cc_pair_1.groups = [user_group_1.id]\n\n print(\"User group 1 created and synced\")\n\n # create a user group and attach it to connector 2\n user_group_2 = UserGroupManager.create(\n name=\"Test User Group 2\",\n cc_pair_ids=[cc_pair_2.id],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_2],\n user_performing_action=admin_user,\n )\n cc_pair_2.groups = [user_group_2.id]\n\n print(\"User group 2 created and synced\")\n\n # verify vespa document is in the user group\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n group_names=[user_group_1.name, user_group_2.name],\n doc_creating_user=admin_user,\n )\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n group_names=[user_group_1.name, user_group_2.name],\n doc_creating_user=admin_user,\n )\n\n # delete connector 1\n CCPairManager.pause_cc_pair(\n cc_pair=cc_pair_1,\n user_performing_action=admin_user,\n )\n CCPairManager.delete(\n cc_pair=cc_pair_1,\n user_performing_action=admin_user,\n )\n\n # wait for deletion to finish\n CCPairManager.wait_for_deletion_completion(\n cc_pair_id=cc_pair_1.id, user_performing_action=admin_user\n )\n\n print(\"Connector 1 deleted\")\n\n # check that only connector 1 is deleted\n # TODO: check for the CC pair rather than the connector once the refactor is done\n CCPairManager.verify(\n cc_pair=cc_pair_1,\n verify_deleted=True,\n user_performing_action=admin_user,\n )\n CCPairManager.verify(\n cc_pair=cc_pair_2,\n user_performing_action=admin_user,\n )\n\n # verify the document is not in any document sets\n # verify the document is only in user group 2\n group_names_expected = []\n if is_ee:\n group_names_expected = [user_group_2.name]\n\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n doc_set_names=[],\n group_names=group_names_expected,\n doc_creating_user=admin_user,\n verify_deleted=False,\n )\n" }, { "path": "backend/tests/integration/tests/connector/test_last_indexed_time.py", "content": "\"\"\"\nIntegration tests for the \"Last Indexed\" time displayed on both the\nper-connector detail page and the all-connectors listing page.\n\nExpected behavior: \"Last Indexed\" = time_started of the most recent\nsuccessful index attempt for the cc pair, regardless of pagination.\n\nEdge cases:\n1. First page of index attempts is entirely errors — last_indexed should\n still reflect the older successful attempt beyond page 1.\n2. Credential swap — successful attempts, then failures after a\n \"credential change\"; last_indexed should reflect the most recent\n successful attempt.\n3. Mix of statuses — only the most recent successful attempt matters.\n4. COMPLETED_WITH_ERRORS counts as a success for last_indexed purposes.\n\"\"\"\n\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom onyx.db.models import IndexingStatus\nfrom onyx.server.documents.models import CCPairFullInfo\nfrom onyx.server.documents.models import ConnectorIndexingStatusLite\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.index_attempt import IndexAttemptManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _wait_for_real_success(\n cc_pair: DATestCCPair,\n admin: DATestUser,\n) -> None:\n \"\"\"Wait for the initial index attempt to complete successfully.\"\"\"\n CCPairManager.wait_for_indexing_completion(\n cc_pair,\n after=datetime(2000, 1, 1, tzinfo=timezone.utc),\n user_performing_action=admin,\n timeout=120,\n )\n\n\ndef _get_detail(cc_pair_id: int, admin: DATestUser) -> CCPairFullInfo:\n result = CCPairManager.get_single(cc_pair_id, admin)\n assert result is not None\n return result\n\n\ndef _get_listing(cc_pair_id: int, admin: DATestUser) -> ConnectorIndexingStatusLite:\n result = CCPairManager.get_indexing_status_by_id(cc_pair_id, admin)\n assert result is not None\n return result\n\n\ndef test_last_indexed_first_page_all_errors(reset: None) -> None: # noqa: ARG001\n \"\"\"When the first page of index attempts is entirely errors but an\n older successful attempt exists, both the detail page and the listing\n page should still show the time of that successful attempt.\n\n The detail page UI uses page size 8. We insert 10 failed attempts\n more recent than the initial success to push the success off page 1.\n \"\"\"\n admin = UserManager.create(name=\"admin_first_page_errors\")\n cc_pair = CCPairManager.create_from_scratch(user_performing_action=admin)\n _wait_for_real_success(cc_pair, admin)\n\n # Baseline: last_success should be set from the initial successful run\n listing_before = _get_listing(cc_pair.id, admin)\n assert listing_before.last_success is not None\n\n # 10 recent failures push the success off page 1\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=10,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.FAILED,\n error_msg=\"simulated failure\",\n base_time=datetime.now(tz=timezone.utc),\n )\n\n detail = _get_detail(cc_pair.id, admin)\n listing = _get_listing(cc_pair.id, admin)\n\n assert (\n detail.last_indexed is not None\n ), \"Detail page last_indexed is None even though a successful attempt exists\"\n assert (\n listing.last_success is not None\n ), \"Listing page last_success is None even though a successful attempt exists\"\n\n # Both surfaces must agree\n assert detail.last_indexed == listing.last_success, (\n f\"Detail last_indexed={detail.last_indexed} != \"\n f\"listing last_success={listing.last_success}\"\n )\n\n\ndef test_last_indexed_credential_swap_scenario(reset: None) -> None: # noqa: ARG001\n \"\"\"Perform an actual credential swap: create connector + cred1 (cc_pair_1),\n wait for success, then associate a new cred2 with the same connector\n (cc_pair_2), wait for that to succeed, and inject failures on cc_pair_2.\n\n cc_pair_2's last_indexed must reflect cc_pair_2's own success, not\n cc_pair_1's older one. Both the detail page and listing page must agree.\n \"\"\"\n admin = UserManager.create(name=\"admin_cred_swap\")\n\n connector = ConnectorManager.create(user_performing_action=admin)\n cred1 = CredentialManager.create(user_performing_action=admin)\n cc_pair_1 = CCPairManager.create(\n connector_id=connector.id,\n credential_id=cred1.id,\n user_performing_action=admin,\n )\n _wait_for_real_success(cc_pair_1, admin)\n\n cred2 = CredentialManager.create(user_performing_action=admin, name=\"swapped-cred\")\n cc_pair_2 = CCPairManager.create(\n connector_id=connector.id,\n credential_id=cred2.id,\n user_performing_action=admin,\n )\n _wait_for_real_success(cc_pair_2, admin)\n\n listing_after_swap = _get_listing(cc_pair_2.id, admin)\n assert listing_after_swap.last_success is not None\n\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=10,\n cc_pair_id=cc_pair_2.id,\n status=IndexingStatus.FAILED,\n error_msg=\"credential expired\",\n base_time=datetime.now(tz=timezone.utc),\n )\n\n detail = _get_detail(cc_pair_2.id, admin)\n listing = _get_listing(cc_pair_2.id, admin)\n\n assert detail.last_indexed is not None\n assert listing.last_success is not None\n\n assert detail.last_indexed == listing.last_success, (\n f\"Detail last_indexed={detail.last_indexed} != \"\n f\"listing last_success={listing.last_success}\"\n )\n\n\ndef test_last_indexed_mixed_statuses(reset: None) -> None: # noqa: ARG001\n \"\"\"Mix of in_progress, failed, and successful attempts. Only the most\n recent successful attempt's time matters.\"\"\"\n admin = UserManager.create(name=\"admin_mixed\")\n cc_pair = CCPairManager.create_from_scratch(user_performing_action=admin)\n _wait_for_real_success(cc_pair, admin)\n\n now = datetime.now(tz=timezone.utc)\n\n # Success 5 hours ago\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=1,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.SUCCESS,\n base_time=now - timedelta(hours=5),\n )\n\n # Failures 3 hours ago\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=3,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.FAILED,\n error_msg=\"transient failure\",\n base_time=now - timedelta(hours=3),\n )\n\n # In-progress 1 hour ago\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=1,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.IN_PROGRESS,\n base_time=now - timedelta(hours=1),\n )\n\n detail = _get_detail(cc_pair.id, admin)\n listing = _get_listing(cc_pair.id, admin)\n\n assert detail.last_indexed is not None\n assert listing.last_success is not None\n\n assert detail.last_indexed == listing.last_success, (\n f\"Detail last_indexed={detail.last_indexed} != \"\n f\"listing last_success={listing.last_success}\"\n )\n\n\ndef test_last_indexed_completed_with_errors(reset: None) -> None: # noqa: ARG001\n \"\"\"COMPLETED_WITH_ERRORS is treated as a successful attempt (matching\n IndexingStatus.is_successful()). When it is the most recent \"success\"\n and later attempts all failed, both surfaces should reflect its time.\"\"\"\n admin = UserManager.create(name=\"admin_completed_errors\")\n cc_pair = CCPairManager.create_from_scratch(user_performing_action=admin)\n _wait_for_real_success(cc_pair, admin)\n\n now = datetime.now(tz=timezone.utc)\n\n # COMPLETED_WITH_ERRORS 2 hours ago\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=1,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.COMPLETED_WITH_ERRORS,\n base_time=now - timedelta(hours=2),\n )\n\n # 10 failures after — push everything else off page 1\n IndexAttemptManager.create_test_index_attempts(\n num_attempts=10,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.FAILED,\n error_msg=\"post-partial failure\",\n base_time=now,\n )\n\n detail = _get_detail(cc_pair.id, admin)\n listing = _get_listing(cc_pair.id, admin)\n\n assert (\n detail.last_indexed is not None\n ), \"COMPLETED_WITH_ERRORS should count as a success for last_indexed\"\n assert (\n listing.last_success is not None\n ), \"COMPLETED_WITH_ERRORS should count as a success for last_success\"\n\n assert detail.last_indexed == listing.last_success, (\n f\"Detail last_indexed={detail.last_indexed} != \"\n f\"listing last_success={listing.last_success}\"\n )\n" }, { "path": "backend/tests/integration/tests/discord_bot/test_discord_bot_api.py", "content": "\"\"\"Integration tests for Discord bot API endpoints.\n\nThese tests hit actual API endpoints via HTTP requests.\n\"\"\"\n\nimport pytest\nimport requests\n\nfrom onyx.db.discord_bot import get_discord_service_api_key\nfrom onyx.db.discord_bot import get_or_create_discord_service_api_key\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom tests.integration.common_utils.managers.discord_bot import DiscordBotManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass TestBotConfigEndpoints:\n \"\"\"Tests for /manage/admin/discord-bot/config endpoints.\"\"\"\n\n def test_get_bot_config_not_configured(self, reset: None) -> None: # noqa: ARG002\n \"\"\"GET /config returns configured=False when no config exists.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Ensure no config exists\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n\n config = DiscordBotManager.get_bot_config(admin_user)\n\n assert config[\"configured\"] is False\n assert \"created_at\" not in config or config.get(\"created_at\") is None\n\n def test_create_bot_config(self, reset: None) -> None: # noqa: ARG002\n \"\"\"POST /config creates a new bot config.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Ensure no config exists\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n\n config = DiscordBotManager.create_bot_config(\n bot_token=\"test_token_123\",\n user_performing_action=admin_user,\n )\n\n assert config[\"configured\"] is True\n assert \"created_at\" in config\n\n # Cleanup\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n\n def test_create_bot_config_already_exists(\n self,\n reset: None, # noqa: ARG002\n ) -> None:\n \"\"\"POST /config returns 409 if config already exists.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Ensure no config exists, then create one\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n DiscordBotManager.create_bot_config(\n bot_token=\"token1\",\n user_performing_action=admin_user,\n )\n\n # Try to create another - should fail\n with pytest.raises(requests.HTTPError) as exc_info:\n DiscordBotManager.create_bot_config(\n bot_token=\"token2\",\n user_performing_action=admin_user,\n )\n\n assert exc_info.value.response.status_code == 409\n\n # Cleanup\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n\n def test_delete_bot_config(self, reset: None) -> None: # noqa: ARG002\n \"\"\"DELETE /config removes the bot config.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Ensure no config exists, then create one\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n DiscordBotManager.create_bot_config(\n bot_token=\"test_token\",\n user_performing_action=admin_user,\n )\n\n # Delete it\n result = DiscordBotManager.delete_bot_config(admin_user)\n assert result[\"deleted\"] is True\n\n # Verify it's gone\n config = DiscordBotManager.get_bot_config(admin_user)\n assert config[\"configured\"] is False\n\n def test_delete_bot_config_not_found(self, reset: None) -> None: # noqa: ARG002\n \"\"\"DELETE /config returns 404 if no config exists.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Ensure no config exists\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n\n # Try to delete - should fail\n with pytest.raises(requests.HTTPError) as exc_info:\n DiscordBotManager.delete_bot_config(admin_user)\n\n assert exc_info.value.response.status_code == 404\n\n\nclass TestGuildConfigEndpoints:\n \"\"\"Tests for /manage/admin/discord-bot/guilds endpoints.\"\"\"\n\n def test_create_guild_config(self, reset: None) -> None: # noqa: ARG002\n \"\"\"POST /guilds creates a new guild config with registration key.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n guild = DiscordBotManager.create_guild(admin_user)\n\n assert guild.id is not None\n assert guild.registration_key is not None\n assert guild.registration_key.startswith(\"discord_\")\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_list_guilds(self, reset: None) -> None: # noqa: ARG002\n \"\"\"GET /guilds returns all guild configs.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create some guilds\n guild1 = DiscordBotManager.create_guild(admin_user)\n guild2 = DiscordBotManager.create_guild(admin_user)\n\n guilds = DiscordBotManager.list_guilds(admin_user)\n\n guild_ids = [g[\"id\"] for g in guilds]\n assert guild1.id in guild_ids\n assert guild2.id in guild_ids\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild1.id, admin_user)\n DiscordBotManager.delete_guild_if_exists(guild2.id, admin_user)\n\n def test_get_guild_config(self, reset: None) -> None: # noqa: ARG002\n \"\"\"GET /guilds/{config_id} returns the specific guild config.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n guild = DiscordBotManager.create_guild(admin_user)\n\n fetched = DiscordBotManager.get_guild(guild.id, admin_user)\n\n assert fetched[\"id\"] == guild.id\n assert fetched[\"enabled\"] is True # Default\n assert fetched[\"guild_id\"] is None # Not registered yet\n assert fetched[\"guild_name\"] is None\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_get_guild_config_not_found(self, reset: None) -> None: # noqa: ARG002\n \"\"\"GET /guilds/{config_id} returns 404 for non-existent guild.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n result = DiscordBotManager.get_guild_or_none(999999, admin_user)\n assert result is None\n\n def test_update_guild_config(self, reset: None) -> None: # noqa: ARG002\n \"\"\"PATCH /guilds/{config_id} updates the guild config.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n guild = DiscordBotManager.create_guild(admin_user)\n\n # Update enabled status\n updated = DiscordBotManager.update_guild(\n guild.id,\n admin_user,\n enabled=False,\n )\n\n assert updated[\"enabled\"] is False\n\n # Verify persistence\n fetched = DiscordBotManager.get_guild(guild.id, admin_user)\n assert fetched[\"enabled\"] is False\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_delete_guild_config(self, reset: None) -> None: # noqa: ARG002\n \"\"\"DELETE /guilds/{config_id} removes the guild config.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n guild = DiscordBotManager.create_guild(admin_user)\n\n # Delete it\n result = DiscordBotManager.delete_guild(guild.id, admin_user)\n assert result[\"deleted\"] is True\n\n # Verify it's gone\n assert DiscordBotManager.get_guild_or_none(guild.id, admin_user) is None\n\n def test_delete_guild_config_not_found(self, reset: None) -> None: # noqa: ARG002\n \"\"\"DELETE /guilds/{config_id} returns 404 for non-existent guild.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n with pytest.raises(requests.HTTPError) as exc_info:\n DiscordBotManager.delete_guild(999999, admin_user)\n\n assert exc_info.value.response.status_code == 404\n\n def test_registration_key_format(self, reset: None) -> None: # noqa: ARG002\n \"\"\"Registration key has proper format with tenant encoded.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n guild = DiscordBotManager.create_guild(admin_user)\n\n # Key should be: discord_{encoded_tenant}.{random}\n key = guild.registration_key\n assert key is not None\n assert key.startswith(\"discord_\")\n\n # Should have two parts separated by dot\n key_body = key.removeprefix(\"discord_\")\n parts = key_body.split(\".\", 1)\n assert len(parts) == 2\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_each_registration_key_is_unique(self, reset: None) -> None: # noqa: ARG002\n \"\"\"Each created guild gets a unique registration key.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n guilds = [DiscordBotManager.create_guild(admin_user) for _ in range(5)]\n keys = [g.registration_key for g in guilds]\n\n assert len(set(keys)) == 5 # All unique\n\n # Cleanup\n for guild in guilds:\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n\nclass TestChannelConfigEndpoints:\n \"\"\"Tests for /manage/admin/discord-bot/guilds/{id}/channels endpoints.\"\"\"\n\n def test_list_channels_empty(self, reset: None) -> None: # noqa: ARG002\n \"\"\"GET /guilds/{id}/channels returns empty list when no channels exist.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create a registered guild (has guild_id set)\n guild = DiscordBotManager.create_registered_guild_in_db(\n guild_id=111111111,\n guild_name=\"Test Guild\",\n )\n\n channels = DiscordBotManager.list_channels(guild.id, admin_user)\n\n assert channels == []\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_list_channels_with_data(self, reset: None) -> None: # noqa: ARG002\n \"\"\"GET /guilds/{id}/channels returns channel configs.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create a registered guild (has guild_id set)\n guild = DiscordBotManager.create_registered_guild_in_db(\n guild_id=222222222,\n guild_name=\"Test Guild\",\n )\n\n # Create test channels directly in DB\n channel1 = DiscordBotManager.create_test_channel_in_db(\n guild_config_id=guild.id,\n channel_id=123456789,\n channel_name=\"general\",\n )\n channel2 = DiscordBotManager.create_test_channel_in_db(\n guild_config_id=guild.id,\n channel_id=987654321,\n channel_name=\"help\",\n channel_type=\"forum\",\n )\n\n channels = DiscordBotManager.list_channels(guild.id, admin_user)\n\n assert len(channels) == 2\n channel_ids = [c.id for c in channels]\n assert channel1.id in channel_ids\n assert channel2.id in channel_ids\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_update_channel_enabled(self, reset: None) -> None: # noqa: ARG002\n \"\"\"PATCH /guilds/{id}/channels/{id} updates enabled status.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create a registered guild (has guild_id set)\n guild = DiscordBotManager.create_registered_guild_in_db(\n guild_id=333333333,\n guild_name=\"Test Guild\",\n )\n channel = DiscordBotManager.create_test_channel_in_db(\n guild_config_id=guild.id,\n channel_id=123456789,\n channel_name=\"general\",\n )\n\n # Default is disabled\n assert channel.enabled is False\n\n # Enable the channel\n updated = DiscordBotManager.update_channel(\n guild.id,\n channel.id,\n admin_user,\n enabled=True,\n )\n\n assert updated.enabled is True\n\n # Verify persistence\n channels = DiscordBotManager.list_channels(guild.id, admin_user)\n found = next(c for c in channels if c.id == channel.id)\n assert found.enabled is True\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_update_channel_thread_only_mode(self, reset: None) -> None: # noqa: ARG002\n \"\"\"PATCH /guilds/{id}/channels/{id} updates thread_only_mode.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create a registered guild (has guild_id set)\n guild = DiscordBotManager.create_registered_guild_in_db(\n guild_id=444444444,\n guild_name=\"Test Guild\",\n )\n channel = DiscordBotManager.create_test_channel_in_db(\n guild_config_id=guild.id,\n channel_id=123456789,\n channel_name=\"general\",\n )\n\n # Default is False\n assert channel.thread_only_mode is False\n\n # Enable thread_only_mode\n updated = DiscordBotManager.update_channel(\n guild.id,\n channel.id,\n admin_user,\n thread_only_mode=True,\n )\n\n assert updated.thread_only_mode is True\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_update_channel_require_bot_invocation(\n self,\n reset: None, # noqa: ARG002\n ) -> None:\n \"\"\"PATCH /guilds/{id}/channels/{id} updates require_bot_invocation.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create a registered guild (has guild_id set)\n guild = DiscordBotManager.create_registered_guild_in_db(\n guild_id=555555555,\n guild_name=\"Test Guild\",\n )\n channel = DiscordBotManager.create_test_channel_in_db(\n guild_config_id=guild.id,\n channel_id=123456789,\n channel_name=\"general\",\n )\n\n # Default is True\n assert channel.require_bot_invocation is True\n\n # Disable require_bot_invocation\n updated = DiscordBotManager.update_channel(\n guild.id,\n channel.id,\n admin_user,\n require_bot_invocation=False,\n )\n\n assert updated.require_bot_invocation is False\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n def test_update_channel_not_found(self, reset: None) -> None: # noqa: ARG002\n \"\"\"PATCH /guilds/{id}/channels/{id} returns 404 for non-existent channel.\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create a registered guild (has guild_id set)\n guild = DiscordBotManager.create_registered_guild_in_db(\n guild_id=666666666,\n guild_name=\"Test Guild\",\n )\n\n with pytest.raises(requests.HTTPError) as exc_info:\n DiscordBotManager.update_channel(\n guild.id,\n 999999,\n admin_user,\n enabled=True,\n )\n\n assert exc_info.value.response.status_code == 404\n\n # Cleanup\n DiscordBotManager.delete_guild_if_exists(guild.id, admin_user)\n\n\nclass TestServiceApiKeyCleanup:\n \"\"\"Tests for service API key cleanup when bot/guild configs are deleted.\"\"\"\n\n def test_delete_bot_config_also_deletes_service_api_key(\n self,\n reset: None, # noqa: ARG002\n ) -> None:\n \"\"\"DELETE /config also deletes the service API key (self-hosted flow).\"\"\"\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Setup: create bot config via API\n DiscordBotManager.delete_bot_config_if_exists(admin_user)\n DiscordBotManager.create_bot_config(\n bot_token=\"test_token\",\n user_performing_action=admin_user,\n )\n\n # Create service API key directly in DB (simulating bot registration)\n with get_session_with_current_tenant() as db_session:\n get_or_create_discord_service_api_key(db_session, \"public\")\n db_session.commit()\n\n # Verify it exists\n assert get_discord_service_api_key(db_session) is not None\n\n # Delete bot config via API\n result = DiscordBotManager.delete_bot_config(admin_user)\n assert result[\"deleted\"] is True\n\n # Verify service API key was also deleted\n with get_session_with_current_tenant() as db_session:\n assert get_discord_service_api_key(db_session) is None\n" }, { "path": "backend/tests/integration/tests/discord_bot/test_discord_bot_db.py", "content": "\"\"\"Integration tests for Discord bot database operations.\n\nThese tests verify CRUD operations for Discord bot models.\n\"\"\"\n\nfrom collections.abc import Generator\n\nimport pytest\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.discord_bot import bulk_create_channel_configs\nfrom onyx.db.discord_bot import create_discord_bot_config\nfrom onyx.db.discord_bot import create_guild_config\nfrom onyx.db.discord_bot import delete_discord_bot_config\nfrom onyx.db.discord_bot import delete_discord_service_api_key\nfrom onyx.db.discord_bot import delete_guild_config\nfrom onyx.db.discord_bot import get_channel_configs\nfrom onyx.db.discord_bot import get_discord_bot_config\nfrom onyx.db.discord_bot import get_discord_service_api_key\nfrom onyx.db.discord_bot import get_guild_config_by_internal_id\nfrom onyx.db.discord_bot import get_guild_config_by_registration_key\nfrom onyx.db.discord_bot import get_guild_configs\nfrom onyx.db.discord_bot import get_or_create_discord_service_api_key\nfrom onyx.db.discord_bot import sync_channel_configs\nfrom onyx.db.discord_bot import update_discord_channel_config\nfrom onyx.db.discord_bot import update_guild_config\nfrom onyx.db.models import Persona\nfrom onyx.db.utils import DiscordChannelView\nfrom onyx.server.manage.discord_bot.utils import generate_discord_registration_key\n\n\ndef _create_test_persona(db_session: Session, persona_id: int, name: str) -> Persona:\n \"\"\"Create a minimal test persona.\"\"\"\n persona = Persona(\n id=persona_id,\n name=name,\n description=\"Test persona for Discord bot tests\",\n is_listed=True,\n is_featured=False,\n deleted=False,\n builtin_persona=False,\n )\n db_session.add(persona)\n db_session.flush()\n return persona\n\n\ndef _delete_test_persona(db_session: Session, persona_id: int) -> None:\n \"\"\"Delete a test persona.\"\"\"\n db_session.query(Persona).filter(Persona.id == persona_id).delete()\n db_session.flush()\n\n\nclass TestBotConfigAPI:\n \"\"\"Tests for bot config API operations.\"\"\"\n\n def test_create_bot_config(self, db_session: Session) -> None:\n \"\"\"Create bot config succeeds with valid token.\"\"\"\n # Clean up any existing config first\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n config = create_discord_bot_config(db_session, bot_token=\"test_token_123\")\n db_session.commit()\n\n assert config is not None\n assert config.bot_token is not None\n assert config.bot_token.get_value(apply_mask=False) == \"test_token_123\"\n\n # Cleanup\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n def test_create_bot_config_already_exists(self, db_session: Session) -> None:\n \"\"\"Creating config twice raises ValueError.\"\"\"\n # Clean up first\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n create_discord_bot_config(db_session, bot_token=\"token1\")\n db_session.commit()\n\n with pytest.raises(ValueError):\n create_discord_bot_config(db_session, bot_token=\"token2\")\n\n # Cleanup\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n def test_get_bot_config(self, db_session: Session) -> None:\n \"\"\"Get bot config returns config with masked token.\"\"\"\n # Clean up first\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n create_discord_bot_config(db_session, bot_token=\"my_secret_token\")\n db_session.commit()\n\n config = get_discord_bot_config(db_session)\n\n assert config is not None\n # Token should be stored (we don't mask in DB, only API response)\n assert config.bot_token is not None\n\n # Cleanup\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n def test_delete_bot_config(self, db_session: Session) -> None:\n \"\"\"Delete bot config removes it from DB.\"\"\"\n # Clean up first\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n create_discord_bot_config(db_session, bot_token=\"token\")\n db_session.commit()\n\n deleted = delete_discord_bot_config(db_session)\n db_session.commit()\n\n assert deleted is True\n assert get_discord_bot_config(db_session) is None\n\n def test_delete_bot_config_not_found(self, db_session: Session) -> None:\n \"\"\"Delete when no config exists returns False.\"\"\"\n # Ensure no config exists\n delete_discord_bot_config(db_session)\n db_session.commit()\n\n deleted = delete_discord_bot_config(db_session)\n assert deleted is False\n\n\nclass TestRegistrationKeyAPI:\n \"\"\"Tests for registration key API operations.\"\"\"\n\n def test_create_registration_key(self, db_session: Session) -> None:\n \"\"\"Create registration key with proper format.\"\"\"\n key = generate_discord_registration_key(\"test_tenant\")\n\n config = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n assert config is not None\n assert config.registration_key == key\n assert key.startswith(\"discord_\")\n assert \"test_tenant\" in key or \"test%5Ftenant\" in key\n\n # Cleanup\n delete_guild_config(db_session, config.id)\n db_session.commit()\n\n def test_registration_key_is_unique(\n self,\n db_session: Session, # noqa: ARG002\n ) -> None:\n \"\"\"Each generated key is unique.\"\"\"\n keys = [generate_discord_registration_key(\"tenant\") for _ in range(5)]\n assert len(set(keys)) == 5\n\n def test_delete_registration_key(self, db_session: Session) -> None:\n \"\"\"Deleted key can no longer be used.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n config = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n config_id = config.id\n\n # Delete\n deleted = delete_guild_config(db_session, config_id)\n db_session.commit()\n\n assert deleted is True\n\n # Should not find it anymore\n found = get_guild_config_by_registration_key(db_session, key)\n assert found is None\n\n\nclass TestGuildConfigAPI:\n \"\"\"Tests for guild config API operations.\"\"\"\n\n def test_list_guilds(self, db_session: Session) -> None:\n \"\"\"List guilds returns all guild configs.\"\"\"\n # Create some guild configs\n key1 = generate_discord_registration_key(\"t1\")\n key2 = generate_discord_registration_key(\"t2\")\n\n config1 = create_guild_config(db_session, registration_key=key1)\n config2 = create_guild_config(db_session, registration_key=key2)\n db_session.commit()\n\n configs = get_guild_configs(db_session)\n\n assert len(configs) >= 2\n\n # Cleanup\n delete_guild_config(db_session, config1.id)\n delete_guild_config(db_session, config2.id)\n db_session.commit()\n\n def test_get_guild_config(self, db_session: Session) -> None:\n \"\"\"Get specific guild config by ID.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n config = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n found = get_guild_config_by_internal_id(db_session, config.id)\n\n assert found is not None\n assert found.id == config.id\n assert found.registration_key == key\n\n # Cleanup\n delete_guild_config(db_session, config.id)\n db_session.commit()\n\n def test_update_guild_enabled(self, db_session: Session) -> None:\n \"\"\"Update guild enabled status.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n config = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n # Initially enabled is True by default\n assert config.enabled is True\n\n # Disable\n updated = update_guild_config(\n db_session, config, enabled=False, default_persona_id=None\n )\n db_session.commit()\n\n assert updated.enabled is False\n\n # Cleanup\n delete_guild_config(db_session, config.id)\n db_session.commit()\n\n def test_update_guild_persona(self, db_session: Session) -> None:\n \"\"\"Update guild default persona.\"\"\"\n # Create test persona first to satisfy foreign key constraint\n _create_test_persona(db_session, 5, \"Test Persona 5\")\n db_session.commit()\n\n key = generate_discord_registration_key(\"tenant\")\n config = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n # Set persona\n updated = update_guild_config(\n db_session, config, enabled=True, default_persona_id=5\n )\n db_session.commit()\n\n assert updated.default_persona_id == 5\n\n # Cleanup\n delete_guild_config(db_session, config.id)\n _delete_test_persona(db_session, 5)\n db_session.commit()\n\n\nclass TestChannelConfigAPI:\n \"\"\"Tests for channel config API operations.\"\"\"\n\n def test_list_channels_for_guild(self, db_session: Session) -> None:\n \"\"\"List channels returns all channel configs for guild.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n # Create some channels\n channels = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n DiscordChannelView(\n channel_id=222,\n channel_name=\"help\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n bulk_create_channel_configs(db_session, guild.id, channels)\n db_session.commit()\n\n channel_configs = get_channel_configs(db_session, guild.id)\n\n assert len(channel_configs) == 2\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n def test_update_channel_enabled(self, db_session: Session) -> None:\n \"\"\"Update channel enabled status.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n channels = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n created = bulk_create_channel_configs(db_session, guild.id, channels)\n db_session.commit()\n\n # Channels are disabled by default\n assert created[0].enabled is False\n\n # Enable\n updated = update_discord_channel_config(\n db_session,\n created[0],\n channel_name=\"general\",\n thread_only_mode=False,\n require_bot_invocation=True,\n enabled=True,\n )\n db_session.commit()\n\n assert updated.enabled is True\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n def test_update_channel_thread_only_mode(self, db_session: Session) -> None:\n \"\"\"Update channel thread_only_mode setting.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n channels = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n created = bulk_create_channel_configs(db_session, guild.id, channels)\n db_session.commit()\n\n # Update thread_only_mode\n updated = update_discord_channel_config(\n db_session,\n created[0],\n channel_name=\"general\",\n thread_only_mode=True,\n require_bot_invocation=True,\n enabled=True,\n )\n db_session.commit()\n\n assert updated.thread_only_mode is True\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n def test_sync_channels_adds_new(self, db_session: Session) -> None:\n \"\"\"Sync channels adds new channels.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n # Initial channels\n initial = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n bulk_create_channel_configs(db_session, guild.id, initial)\n db_session.commit()\n\n # Sync with new channel\n current = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n DiscordChannelView(\n channel_id=222,\n channel_name=\"new-channel\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n added, removed, updated = sync_channel_configs(db_session, guild.id, current)\n db_session.commit()\n\n assert added == 1\n assert removed == 0\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n def test_sync_channels_removes_deleted(self, db_session: Session) -> None:\n \"\"\"Sync channels removes deleted channels.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n # Initial channels\n initial = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n DiscordChannelView(\n channel_id=222,\n channel_name=\"old-channel\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n bulk_create_channel_configs(db_session, guild.id, initial)\n db_session.commit()\n\n # Sync with one channel removed\n current = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n added, removed, updated = sync_channel_configs(db_session, guild.id, current)\n db_session.commit()\n\n assert added == 0\n assert removed == 1\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n def test_sync_channels_updates_renamed(self, db_session: Session) -> None:\n \"\"\"Sync channels updates renamed channels.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n db_session.commit()\n\n # Initial channels\n initial = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"old-name\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n bulk_create_channel_configs(db_session, guild.id, initial)\n db_session.commit()\n\n # Sync with renamed channel\n current = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"new-name\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n added, removed, updated = sync_channel_configs(db_session, guild.id, current)\n db_session.commit()\n\n assert added == 0\n assert removed == 0\n assert updated == 1\n\n # Verify name was updated\n configs = get_channel_configs(db_session, guild.id)\n assert configs[0].channel_name == \"new-name\"\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n\nclass TestPersonaConfigurationAPI:\n \"\"\"Tests for persona configuration in API.\"\"\"\n\n def test_guild_persona_used_in_api_call(self, db_session: Session) -> None:\n \"\"\"Guild default_persona_id is used when no channel override.\"\"\"\n # Create test persona first\n _create_test_persona(db_session, 42, \"Test Persona 42\")\n db_session.commit()\n\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n update_guild_config(db_session, guild, enabled=True, default_persona_id=42)\n db_session.commit()\n\n # Verify persona is set\n config = get_guild_config_by_internal_id(db_session, guild.id)\n assert config is not None\n assert config.default_persona_id == 42\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n _delete_test_persona(db_session, 42)\n db_session.commit()\n\n def test_channel_persona_override_in_api_call(self, db_session: Session) -> None:\n \"\"\"Channel persona_override_id takes precedence over guild default.\"\"\"\n # Create test personas first\n _create_test_persona(db_session, 42, \"Test Persona 42\")\n _create_test_persona(db_session, 99, \"Test Persona 99\")\n db_session.commit()\n\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n update_guild_config(db_session, guild, enabled=True, default_persona_id=42)\n db_session.commit()\n\n channels = [\n DiscordChannelView(\n channel_id=111,\n channel_name=\"general\",\n channel_type=\"text\",\n is_private=False,\n ),\n ]\n created = bulk_create_channel_configs(db_session, guild.id, channels)\n db_session.commit()\n\n # Set channel persona override\n updated = update_discord_channel_config(\n db_session,\n created[0],\n channel_name=\"general\",\n thread_only_mode=False,\n require_bot_invocation=True,\n enabled=True,\n persona_override_id=99, # Override!\n )\n db_session.commit()\n\n assert updated.persona_override_id == 99\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n _delete_test_persona(db_session, 42)\n _delete_test_persona(db_session, 99)\n db_session.commit()\n\n def test_no_persona_uses_default(self, db_session: Session) -> None:\n \"\"\"Neither guild nor channel has persona - uses API default.\"\"\"\n key = generate_discord_registration_key(\"tenant\")\n guild = create_guild_config(db_session, registration_key=key)\n # No persona set\n db_session.commit()\n\n config = get_guild_config_by_internal_id(db_session, guild.id)\n assert config is not None\n assert config.default_persona_id is None\n\n # Cleanup\n delete_guild_config(db_session, guild.id)\n db_session.commit()\n\n\nclass TestServiceApiKeyAPI:\n \"\"\"Tests for Discord service API key operations.\"\"\"\n\n def test_create_service_api_key(self, db_session: Session) -> None:\n \"\"\"Create service API key returns valid key.\"\"\"\n # Clean up any existing key first\n delete_discord_service_api_key(db_session)\n db_session.commit()\n\n api_key = get_or_create_discord_service_api_key(db_session, \"public\")\n db_session.commit()\n\n assert api_key is not None\n assert len(api_key) > 0\n\n # Verify key was stored in database\n stored_key = get_discord_service_api_key(db_session)\n assert stored_key is not None\n\n # Cleanup\n delete_discord_service_api_key(db_session)\n db_session.commit()\n\n def test_get_or_create_returns_existing(self, db_session: Session) -> None:\n \"\"\"get_or_create_discord_service_api_key regenerates key if exists.\"\"\"\n # Clean up any existing key first\n delete_discord_service_api_key(db_session)\n db_session.commit()\n\n # Create first key\n key1 = get_or_create_discord_service_api_key(db_session, \"public\")\n db_session.commit()\n\n # Call again - should regenerate (per implementation, it regenerates to update cache)\n key2 = get_or_create_discord_service_api_key(db_session, \"public\")\n db_session.commit()\n\n # Keys should be different since it regenerates\n assert key1 != key2\n\n # But there should still be only one key in the database\n stored_key = get_discord_service_api_key(db_session)\n assert stored_key is not None\n\n # Cleanup\n delete_discord_service_api_key(db_session)\n db_session.commit()\n\n def test_delete_service_api_key(self, db_session: Session) -> None:\n \"\"\"Delete service API key removes it from DB.\"\"\"\n # Clean up any existing key first\n delete_discord_service_api_key(db_session)\n db_session.commit()\n\n # Create a key\n get_or_create_discord_service_api_key(db_session, \"public\")\n db_session.commit()\n\n # Delete it\n deleted = delete_discord_service_api_key(db_session)\n db_session.commit()\n\n assert deleted is True\n assert get_discord_service_api_key(db_session) is None\n\n def test_delete_service_api_key_not_found(self, db_session: Session) -> None:\n \"\"\"Delete when no key exists returns False.\"\"\"\n # Ensure no key exists\n delete_discord_service_api_key(db_session)\n db_session.commit()\n\n deleted = delete_discord_service_api_key(db_session)\n assert deleted is False\n\n\n# Pytest fixture for db_session\n@pytest.fixture\ndef db_session() -> Generator[Session, None, None]:\n \"\"\"Create database session for tests.\"\"\"\n from onyx.db.engine.sql_engine import get_session_with_current_tenant\n from onyx.db.engine.sql_engine import SqlEngine\n from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR\n\n SqlEngine.init_engine(pool_size=10, max_overflow=5)\n\n token = CURRENT_TENANT_ID_CONTEXTVAR.set(\"public\")\n try:\n with get_session_with_current_tenant() as session:\n yield session\n finally:\n CURRENT_TENANT_ID_CONTEXTVAR.reset(token)\n" }, { "path": "backend/tests/integration/tests/document_set/test_syncing.py", "content": "from onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.constants import NUM_DOCS\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.document_set import DocumentSetManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef test_multiple_document_sets_syncing_same_connnector(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture,\n) -> None:\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # create api key\n api_key: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user,\n )\n\n # create connector\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n # seed documents\n cc_pair_1.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_1,\n num_docs=NUM_DOCS,\n api_key=api_key,\n )\n\n # Create document sets\n doc_set_1 = DocumentSetManager.create(\n cc_pair_ids=[cc_pair_1.id],\n user_performing_action=admin_user,\n )\n doc_set_2 = DocumentSetManager.create(\n cc_pair_ids=[cc_pair_1.id],\n user_performing_action=admin_user,\n )\n\n DocumentSetManager.wait_for_sync(\n user_performing_action=admin_user,\n )\n\n DocumentSetManager.verify(\n document_set=doc_set_1,\n user_performing_action=admin_user,\n )\n DocumentSetManager.verify(\n document_set=doc_set_2,\n user_performing_action=admin_user,\n )\n\n # make sure documents are as expected\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n doc_set_names=[doc_set_1.name, doc_set_2.name],\n doc_creating_user=admin_user,\n )\n\n\ndef test_removing_connector(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture,\n) -> None:\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # create api key\n api_key: DATestAPIKey = APIKeyManager.create(\n user_performing_action=admin_user,\n )\n\n # create connectors\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n cc_pair_2 = CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n # seed documents\n cc_pair_1.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_1,\n num_docs=NUM_DOCS,\n api_key=api_key,\n )\n\n cc_pair_2.documents = DocumentManager.seed_dummy_docs(\n cc_pair=cc_pair_2,\n num_docs=NUM_DOCS,\n api_key=api_key,\n )\n\n # Create document sets\n doc_set_1 = DocumentSetManager.create(\n cc_pair_ids=[cc_pair_1.id, cc_pair_2.id],\n user_performing_action=admin_user,\n )\n\n DocumentSetManager.wait_for_sync(\n user_performing_action=admin_user,\n )\n\n DocumentSetManager.verify(\n document_set=doc_set_1,\n user_performing_action=admin_user,\n )\n\n # make sure cc_pair_1 docs are doc_set_1 only\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n doc_set_names=[doc_set_1.name],\n doc_creating_user=admin_user,\n )\n\n # make sure cc_pair_2 docs are doc_set_1 only\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n doc_set_names=[doc_set_1.name],\n doc_creating_user=admin_user,\n )\n\n # remove cc_pair_2 from document set\n doc_set_1.cc_pair_ids = [cc_pair_1.id]\n DocumentSetManager.edit(\n doc_set_1,\n user_performing_action=admin_user,\n )\n\n DocumentSetManager.wait_for_sync(\n user_performing_action=admin_user,\n )\n\n # make sure cc_pair_1 docs are doc_set_1 only\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_1,\n doc_set_names=[doc_set_1.name],\n doc_creating_user=admin_user,\n )\n\n # make sure cc_pair_2 docs have no doc set\n DocumentManager.verify(\n vespa_client=vespa_client,\n cc_pair=cc_pair_2,\n doc_set_names=[],\n doc_creating_user=admin_user,\n )\n" }, { "path": "backend/tests/integration/tests/image_generation/test_image_generation_config.py", "content": "\"\"\"Integration tests for image generation config endpoints.\n\nTests cover CRUD operations for /admin/image-generation/config endpoints.\nThe /admin/image-generation/test endpoint is not tested as it makes real API calls.\n\nUses module-scoped fixtures to reset DB and create users once per module for faster execution.\n\"\"\"\n\nimport pytest\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.image_generation import (\n ImageGenerationConfigManager,\n)\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n@pytest.fixture(scope=\"module\")\ndef setup_image_generation_tests() -> tuple[DATestUser, DATestLLMProvider]:\n \"\"\"Module-scoped fixture that runs once for all tests in this module.\n\n - Resets DB once at the start of the module\n - Creates admin user once\n - Creates LLM provider once (for clone-mode test)\n - Returns (admin_user, llm_provider) tuple for all tests to use\n \"\"\"\n reset_all()\n admin_user = UserManager.create(name=\"admin_user\")\n llm_provider = LLMProviderManager.create(user_performing_action=admin_user)\n return admin_user, llm_provider\n\n\ndef test_create_image_generation_config(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test creating an image generation config with new credentials.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"test-openai-dalle\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-test-key-12345\",\n is_default=False,\n user_performing_action=admin_user,\n )\n\n assert config.image_provider_id == \"test-openai-dalle\"\n assert config.model_name == \"dall-e-3\"\n assert config.is_default is False\n\n # Verify it exists in the list\n ImageGenerationConfigManager.verify(\n config=config,\n user_performing_action=admin_user,\n )\n\n\ndef test_create_image_generation_config_from_provider(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test creating an image generation config by cloning from an existing LLM provider.\"\"\"\n admin_user, llm_provider = setup_image_generation_tests\n\n # Create image generation config from the provider\n config = ImageGenerationConfigManager.create_from_provider(\n source_llm_provider_id=llm_provider.id,\n image_provider_id=\"test-from-provider\",\n model_name=\"gpt-image-1\",\n is_default=True,\n user_performing_action=admin_user,\n )\n\n assert config.image_provider_id == \"test-from-provider\"\n assert config.model_name == \"gpt-image-1\"\n assert config.is_default is True\n\n # Verify it exists\n ImageGenerationConfigManager.verify(\n config=config,\n user_performing_action=admin_user,\n )\n\n\ndef test_create_duplicate_config_fails(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test that creating a config with an existing image_provider_id fails.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create first config\n ImageGenerationConfigManager.create(\n image_provider_id=\"duplicate-test-id\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-test-key-1\",\n user_performing_action=admin_user,\n )\n\n # Try to create another with the same image_provider_id\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n json={\n \"image_provider_id\": \"duplicate-test-id\",\n \"model_name\": \"gpt-image-1\",\n \"provider\": \"openai\",\n \"api_key\": \"sk-test-key-2\",\n },\n headers=admin_user.headers,\n )\n\n assert response.status_code == 400\n assert \"already exists\" in response.json()[\"detail\"]\n\n\ndef test_get_all_configs(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test getting all image generation configs.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create multiple configs\n config1 = ImageGenerationConfigManager.create(\n image_provider_id=\"config-1\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-key-1\",\n user_performing_action=admin_user,\n )\n config2 = ImageGenerationConfigManager.create(\n image_provider_id=\"config-2\",\n model_name=\"gpt-image-1\",\n provider=\"openai\",\n api_key=\"sk-key-2\",\n user_performing_action=admin_user,\n )\n\n # Get all configs\n all_configs = ImageGenerationConfigManager.get_all(\n user_performing_action=admin_user\n )\n\n assert len(all_configs) >= 2\n config_ids = [c.image_provider_id for c in all_configs]\n assert config1.image_provider_id in config_ids\n assert config2.image_provider_id in config_ids\n\n\ndef test_get_config_credentials(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test getting credentials for an image generation config.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n test_api_key = \"sk-test-credentials-key-12345\"\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"credentials-test\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=test_api_key,\n user_performing_action=admin_user,\n )\n\n # Get credentials\n credentials = ImageGenerationConfigManager.get_credentials(\n image_provider_id=config.image_provider_id,\n user_performing_action=admin_user,\n )\n\n # Credentials should contain the masked API key (first 4 + **** + last 4)\n assert credentials[\"api_key\"] == \"sk-t****2345\"\n assert \"api_base\" in credentials\n assert \"api_version\" in credentials\n assert \"deployment_name\" in credentials\n\n\ndef test_get_credentials_not_found(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test getting credentials for a non-existent config returns 404.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n response = requests.get(\n f\"{API_SERVER_URL}/admin/image-generation/config/non-existent-id/credentials\",\n headers=admin_user.headers,\n )\n\n assert response.status_code == 404\n\n\ndef test_update_config_direct_key_entry(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test updating an image generation config with new direct credentials.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create initial config\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"update-direct-test\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-initial-key\",\n user_performing_action=admin_user,\n )\n\n assert config.model_name == \"dall-e-3\"\n\n # Update with new credentials and model\n new_api_key = \"sk-updated-key-12345\"\n updated_config = ImageGenerationConfigManager.update(\n image_provider_id=config.image_provider_id,\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=new_api_key,\n user_performing_action=admin_user,\n )\n\n assert updated_config.image_provider_id == config.image_provider_id\n assert updated_config.model_name == \"dall-e-3\"\n\n # Verify credentials were updated (masked: first 4 + **** + last 4)\n credentials = ImageGenerationConfigManager.get_credentials(\n image_provider_id=config.image_provider_id,\n user_performing_action=admin_user,\n )\n assert credentials[\"api_key\"] == \"sk-u****2345\"\n\n\ndef test_update_config_clone_mode(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test updating an image generation config by cloning from an LLM provider.\"\"\"\n admin_user, llm_provider = setup_image_generation_tests\n\n # Create initial config with direct credentials\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"update-clone-test\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-initial-direct-key\",\n user_performing_action=admin_user,\n )\n\n assert config.model_name == \"dall-e-3\"\n\n # Update by cloning from LLM provider\n updated_config = ImageGenerationConfigManager.update(\n image_provider_id=config.image_provider_id,\n model_name=\"gpt-image-1\",\n source_llm_provider_id=llm_provider.id,\n user_performing_action=admin_user,\n )\n\n assert updated_config.image_provider_id == config.image_provider_id\n assert updated_config.model_name == \"gpt-image-1\"\n\n # Verify config still exists and is accessible\n ImageGenerationConfigManager.verify(\n config=updated_config,\n user_performing_action=admin_user,\n )\n\n\ndef test_update_config_source_provider_not_found(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test that updating with non-existent source_llm_provider_id fails.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create initial config\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"update-bad-source-test\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-initial-key\",\n user_performing_action=admin_user,\n )\n\n # Try to update with non-existent source provider\n response = requests.put(\n f\"{API_SERVER_URL}/admin/image-generation/config/{config.image_provider_id}\",\n json={\n \"model_name\": \"gpt-image-1\",\n \"source_llm_provider_id\": 999999,\n },\n headers=admin_user.headers,\n )\n\n assert response.status_code == 404\n assert \"not found\" in response.json()[\"detail\"]\n\n\ndef test_delete_config(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test deleting an image generation config.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create a config\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"delete-test\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-delete-key\",\n user_performing_action=admin_user,\n )\n\n # Verify it exists\n ImageGenerationConfigManager.verify(\n config=config,\n user_performing_action=admin_user,\n )\n\n # Delete it\n ImageGenerationConfigManager.delete(\n image_provider_id=config.image_provider_id,\n user_performing_action=admin_user,\n )\n\n # Verify it's deleted\n ImageGenerationConfigManager.verify(\n config=config,\n verify_deleted=True,\n user_performing_action=admin_user,\n )\n\n\ndef test_delete_config_not_found(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test deleting a non-existent config returns 404.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n response = requests.delete(\n f\"{API_SERVER_URL}/admin/image-generation/config/non-existent-id\",\n headers=admin_user.headers,\n )\n\n assert response.status_code == 404\n\n\ndef test_set_default_config(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test setting a config as the default.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create a config that is not default\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"default-test\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-test-key\",\n is_default=False,\n user_performing_action=admin_user,\n )\n\n assert config.is_default is False\n\n # Set it as default\n ImageGenerationConfigManager.set_default(\n image_provider_id=config.image_provider_id,\n user_performing_action=admin_user,\n )\n\n # Verify it's now default\n all_configs = ImageGenerationConfigManager.get_all(\n user_performing_action=admin_user\n )\n updated_config = next(\n c for c in all_configs if c.image_provider_id == config.image_provider_id\n )\n assert updated_config.is_default is True\n\n\ndef test_set_default_clears_previous(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test that setting a new default clears the previous default.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Create first config as default\n config1 = ImageGenerationConfigManager.create(\n image_provider_id=\"first-default\",\n model_name=\"dall-e-3\",\n provider=\"openai\",\n api_key=\"sk-key-1\",\n is_default=True,\n user_performing_action=admin_user,\n )\n\n # Create second config not as default\n config2 = ImageGenerationConfigManager.create(\n image_provider_id=\"second-default\",\n model_name=\"gpt-image-1\",\n provider=\"openai\",\n api_key=\"sk-key-2\",\n is_default=False,\n user_performing_action=admin_user,\n )\n\n # Verify first is default\n all_configs = ImageGenerationConfigManager.get_all(\n user_performing_action=admin_user\n )\n first = next(\n c for c in all_configs if c.image_provider_id == config1.image_provider_id\n )\n second = next(\n c for c in all_configs if c.image_provider_id == config2.image_provider_id\n )\n assert first.is_default is True\n assert second.is_default is False\n\n # Set second as default\n ImageGenerationConfigManager.set_default(\n image_provider_id=config2.image_provider_id,\n user_performing_action=admin_user,\n )\n\n # Verify second is now default and first is not\n all_configs = ImageGenerationConfigManager.get_all(\n user_performing_action=admin_user\n )\n first = next(\n c for c in all_configs if c.image_provider_id == config1.image_provider_id\n )\n second = next(\n c for c in all_configs if c.image_provider_id == config2.image_provider_id\n )\n assert first.is_default is False\n assert second.is_default is True\n\n\ndef test_set_default_not_found(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test setting a non-existent config as default returns 404.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config/non-existent-id/default\",\n headers=admin_user.headers,\n )\n\n assert response.status_code == 404\n\n\ndef test_create_config_missing_credentials(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test that creating a config without credentials fails.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # Try to create without api_key/provider or source_llm_provider_id\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n json={\n \"image_provider_id\": \"no-creds-test\",\n \"model_name\": \"dall-e-3\",\n },\n headers=admin_user.headers,\n )\n\n assert response.status_code == 400\n assert \"No provider or source llm provided\" in response.json()[\"detail\"]\n\n\ndef test_create_config_source_provider_not_found(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"Test creating a config with non-existent source_llm_provider_id fails.\"\"\"\n admin_user, _ = setup_image_generation_tests\n\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n json={\n \"image_provider_id\": \"bad-source-test\",\n \"model_name\": \"dall-e-3\",\n \"source_llm_provider_id\": 999999, # Non-existent ID\n },\n headers=admin_user.headers,\n )\n\n assert response.status_code == 404\n assert \"not found\" in response.json()[\"detail\"]\n" }, { "path": "backend/tests/integration/tests/image_generation/test_image_generation_tool_visibility.py", "content": "\"\"\"Integration tests to check broader image generation config flow endpoints.\"\"\"\n\nimport pytest\n\nfrom onyx.tools.tool_implementations.images.image_generation_tool import (\n ImageGenerationTool,\n)\nfrom tests.integration.common_utils.managers.image_generation import (\n ImageGenerationConfigManager,\n)\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.tool import ToolManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\nIMAGE_GENERATION_TOOL_NAME = ImageGenerationTool.NAME\n\n\n@pytest.fixture(scope=\"module\")\ndef setup_image_generation_tests() -> tuple[DATestUser, DATestLLMProvider]:\n \"\"\"Module-scoped fixture that runs once for all tests in this module.\n\n - Resets DB once at the start of the module\n - Creates admin user once\n - Creates LLM provider once (for clone-mode test)\n - Returns (admin_user, llm_provider) tuple for all tests to use\n \"\"\"\n reset_all()\n admin_user = UserManager.create(name=\"admin_user\")\n llm_provider = LLMProviderManager.create(user_performing_action=admin_user)\n return admin_user, llm_provider\n\n\ndef test_vertex_creds_upload_image_tool_visibility(\n setup_image_generation_tests: tuple[DATestUser, DATestLLMProvider],\n) -> None:\n \"\"\"\n Tests the following scenario:\n 1. No image model added so tool not visible\n 2. Vertex AI creds uploaded\n 3. Image model added so tool visible\n \"\"\"\n admin_user, _ = setup_image_generation_tests\n\n # 1. Check the tools and check that image generation tool is not visible yet\n tools = ToolManager.list_tools(user_performing_action=admin_user)\n assert not any(tool.name == IMAGE_GENERATION_TOOL_NAME for tool in tools)\n\n # 2. Upload vertex ai credentials\n config = ImageGenerationConfigManager.create(\n image_provider_id=\"gemini-2.5-flash-image\",\n model_name=\"gemini-2.5-flash-image\",\n provider=\"vertex_ai\",\n custom_config={\n \"vertex_credentials\": {\n \"type\": \"service_account\",\n \"project_id\": \"test-project-id\",\n \"private_key_id\": \"test-private-key-id\",\n \"private_key\": \"test-private-key\",\n # ... Other random fields that we dont care about\n },\n \"vertex_location\": \"test-location\",\n },\n user_performing_action=admin_user,\n is_default=True,\n )\n\n assert config.image_provider_id == \"gemini-2.5-flash-image\"\n assert config.model_name == \"gemini-2.5-flash-image\"\n\n # 3. Check that the tool is visible\n tools = ToolManager.list_tools(user_performing_action=admin_user)\n assert any(tool.name == IMAGE_GENERATION_TOOL_NAME for tool in tools)\n" }, { "path": "backend/tests/integration/tests/image_indexing/test_indexing_images.py", "content": "import os\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.connectors.models import InputType\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.file import FileManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.settings import SettingsManager\nfrom tests.integration.common_utils.test_models import DATestSettings\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\nFILE_NAME = \"Sample.pdf\"\nFILE_PATH = \"tests/integration/common_utils/test_files\"\nDOCX_FILE_NAME = \"three_images.docx\"\n\n\ndef test_image_indexing(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n vespa_client: vespa_fixture,\n) -> None:\n os.makedirs(FILE_PATH, exist_ok=True)\n test_file_path = os.path.join(FILE_PATH, FILE_NAME)\n\n # Use FileManager to upload the test file\n upload_response = FileManager.upload_file_for_connector(\n file_path=test_file_path, file_name=FILE_NAME, user_performing_action=admin_user\n )\n\n LLMProviderManager.create(\n name=\"test_llm\",\n user_performing_action=admin_user,\n )\n\n SettingsManager.update_settings(\n DATestSettings(\n search_time_image_analysis_enabled=True,\n image_extraction_and_analysis_enabled=True,\n ),\n user_performing_action=admin_user,\n )\n\n file_paths = upload_response.file_paths\n\n if not file_paths:\n pytest.fail(\"File upload failed - no file paths returned\")\n\n # Create a dummy credential for the file connector\n credential = CredentialManager.create(\n source=DocumentSource.FILE,\n credential_json={},\n user_performing_action=admin_user,\n )\n\n # Create the connector\n connector_name = f\"FileConnector-{int(datetime.now().timestamp())}\"\n connector = ConnectorManager.create(\n name=connector_name,\n source=DocumentSource.FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={\n \"file_locations\": file_paths,\n \"file_names\": [FILE_NAME],\n \"zip_metadata_file_id\": None,\n },\n access_type=AccessType.PUBLIC,\n groups=[],\n user_performing_action=admin_user,\n )\n\n # Link the credential to the connector\n cc_pair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.PUBLIC,\n user_performing_action=admin_user,\n )\n\n # Explicitly run the connector to start indexing\n CCPairManager.run_once(\n cc_pair=cc_pair,\n from_beginning=True,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=datetime.now(timezone.utc),\n timeout=300,\n user_performing_action=admin_user,\n )\n\n with get_session_with_current_tenant() as db_session:\n # really gets the chunks from Vespa, which is why there are two;\n # one for the raw text and one for the summarized image.\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n\n assert len(documents) == 2\n for document in documents:\n if \"These are Johns dogs\" in document.content:\n assert document.image_file_id is None\n else:\n assert document.image_file_id is not None\n assert file_paths[0] in document.image_file_id\n\n\ndef test_docx_image_indexing(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n vespa_client: vespa_fixture,\n) -> None:\n \"\"\"Test that images from docx files are correctly extracted and indexed.\"\"\"\n os.makedirs(FILE_PATH, exist_ok=True)\n test_file_path = os.path.join(FILE_PATH, DOCX_FILE_NAME)\n\n # Use FileManager to upload the test file\n upload_response = FileManager.upload_file_for_connector(\n file_path=test_file_path,\n file_name=DOCX_FILE_NAME,\n user_performing_action=admin_user,\n )\n\n LLMProviderManager.create(\n name=\"test_llm_docx\",\n user_performing_action=admin_user,\n )\n\n SettingsManager.update_settings(\n DATestSettings(\n search_time_image_analysis_enabled=True,\n image_extraction_and_analysis_enabled=True,\n ),\n user_performing_action=admin_user,\n )\n\n file_paths = upload_response.file_paths\n\n if not file_paths:\n pytest.fail(\"File upload failed - no file paths returned\")\n\n # Create a dummy credential for the file connector\n credential = CredentialManager.create(\n source=DocumentSource.FILE,\n credential_json={},\n user_performing_action=admin_user,\n )\n\n # Create the connector\n connector_name = f\"DocxFileConnector-{int(datetime.now().timestamp())}\"\n connector = ConnectorManager.create(\n name=connector_name,\n source=DocumentSource.FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={\n \"file_locations\": file_paths,\n \"file_names\": [DOCX_FILE_NAME],\n \"zip_metadata_file_id\": None,\n },\n access_type=AccessType.PUBLIC,\n groups=[],\n user_performing_action=admin_user,\n )\n\n # Link the credential to the connector\n cc_pair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.PUBLIC,\n user_performing_action=admin_user,\n )\n\n # Explicitly run the connector to start indexing\n CCPairManager.run_once(\n cc_pair=cc_pair,\n from_beginning=True,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=datetime.now(timezone.utc),\n timeout=300,\n user_performing_action=admin_user,\n )\n\n with get_session_with_current_tenant() as db_session:\n # Fetch documents from Vespa - expect text content plus 3 images\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n\n # Should have documents for text content plus 3 images\n assert (\n len(documents) >= 3\n ), f\"Expected at least 3 documents (3 images), got {len(documents)}\"\n\n # Count documents with images\n image_documents = [doc for doc in documents if doc.image_file_id is not None]\n text_documents = [doc for doc in documents if doc.image_file_id is None]\n\n assert (\n len(image_documents) == 3\n ), f\"Expected exactly 3 image documents, got {len(image_documents)}\"\n assert (\n len(text_documents) >= 1\n ), f\"Expected at least 1 text document, got {len(text_documents)}\"\n\n # Verify each image document has a valid image_file_id pointing to our uploaded file\n for image_doc in image_documents:\n assert file_paths[0] in (\n image_doc.image_file_id or \"\"\n ), f\"Image document should reference uploaded file: {image_doc.image_file_id}\"\n" }, { "path": "backend/tests/integration/tests/index_attempt/test_index_attempt_pagination.py", "content": "import time\nfrom datetime import datetime\n\nfrom onyx.db.models import IndexingStatus\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.index_attempt import IndexAttemptManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _verify_index_attempt_pagination(\n cc_pair_id: int,\n index_attempt_ids: list[int],\n user_performing_action: DATestUser,\n page_size: int = 5,\n) -> None:\n retrieved_attempts: list[int] = []\n last_time_started = None # Track the last time_started seen\n\n for i in range(0, len(index_attempt_ids), page_size):\n paginated_result = IndexAttemptManager.get_index_attempt_page(\n cc_pair_id=cc_pair_id,\n page=(i // page_size),\n page_size=page_size,\n user_performing_action=user_performing_action,\n )\n\n # Verify that the total items is equal to the length of the index attempts list\n assert paginated_result.total_items == len(index_attempt_ids)\n # Verify that the number of items in the page is equal to the page size\n assert len(paginated_result.items) == min(page_size, len(index_attempt_ids) - i)\n\n # Verify time ordering within the page (descending order)\n for attempt in paginated_result.items:\n if last_time_started is not None:\n assert attempt.time_started is not None\n assert (\n attempt.time_started <= last_time_started\n ), \"Index attempts not in descending time order\"\n last_time_started = attempt.time_started\n\n # Add the retrieved index attempts to the list of retrieved attempts\n retrieved_attempts.extend([attempt.id for attempt in paginated_result.items])\n\n # Create a set of all the expected index attempt IDs\n all_expected_attempts = set(index_attempt_ids)\n # Create a set of all the retrieved index attempt IDs\n all_retrieved_attempts = set(retrieved_attempts)\n\n # Verify that the set of retrieved attempts is equal to the set of expected attempts\n assert all_expected_attempts == all_retrieved_attempts\n\n\ndef test_index_attempt_pagination(reset: None) -> None: # noqa: ARG001\n MAX_WAIT = 60\n all_attempt_ids: list[int] = []\n\n # Create an admin user to perform actions\n user_performing_action: DATestUser = UserManager.create(\n name=\"admin_performing_action\",\n )\n\n # Create a CC pair to attach index attempts to\n cc_pair = CCPairManager.create_from_scratch(\n user_performing_action=user_performing_action,\n )\n\n # Creating a CC pair will create an index attempt as well. wait for it.\n start = time.monotonic()\n while True:\n paginated_result = IndexAttemptManager.get_index_attempt_page(\n cc_pair_id=cc_pair.id,\n page=0,\n page_size=5,\n user_performing_action=user_performing_action,\n )\n\n if paginated_result.total_items == 1:\n all_attempt_ids.append(paginated_result.items[0].id)\n print(\"Initial index attempt from cc_pair creation detected. Continuing...\")\n break\n\n elapsed = time.monotonic() - start\n if elapsed > MAX_WAIT:\n raise TimeoutError(\n f\"Initial index attempt: Not detected within {MAX_WAIT} seconds.\"\n )\n\n print(\n f\"Waiting for initial index attempt: elapsed={elapsed:.2f} timeout={MAX_WAIT}\"\n )\n time.sleep(1)\n\n # Create 299 successful index attempts (for 300 total)\n base_time = datetime.now()\n generated_attempts = IndexAttemptManager.create_test_index_attempts(\n num_attempts=299,\n cc_pair_id=cc_pair.id,\n status=IndexingStatus.SUCCESS,\n base_time=base_time,\n )\n\n for attempt in generated_attempts:\n all_attempt_ids.append(attempt.id)\n\n # Verify basic pagination with different page sizes\n print(\"Verifying basic pagination with page size 5\")\n _verify_index_attempt_pagination(\n cc_pair_id=cc_pair.id,\n index_attempt_ids=all_attempt_ids,\n page_size=5,\n user_performing_action=user_performing_action,\n )\n\n # Test with a larger page size\n print(\"Verifying pagination with page size 100\")\n _verify_index_attempt_pagination(\n cc_pair_id=cc_pair.id,\n index_attempt_ids=all_attempt_ids,\n page_size=100,\n user_performing_action=user_performing_action,\n )\n" }, { "path": "backend/tests/integration/tests/indexing/conftest.py", "content": "import httpx\nimport pytest\n\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT\n\n\n@pytest.fixture\ndef mock_server_client() -> httpx.Client:\n print(\n f\"Initializing mock server client with host: {MOCK_CONNECTOR_SERVER_HOST} and port: {MOCK_CONNECTOR_SERVER_PORT}\"\n )\n return httpx.Client(\n base_url=f\"http://{MOCK_CONNECTOR_SERVER_HOST}:{MOCK_CONNECTOR_SERVER_PORT}\",\n timeout=5.0,\n )\n" }, { "path": "backend/tests/integration/tests/indexing/file_connector/test_file_connector_zip_metadata.py", "content": "import json\nimport os\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport pytest\n\nfrom onyx.connectors.models import InputType\nfrom onyx.db.document import get_documents_for_cc_pair\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.file import FileManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\n# This is a placeholder - you'll need to create this zip file with actual test files\nTEST_FILES_BASE = \"tests/integration/tests/indexing/file_connector/test_files\"\nTEST_META_ZIP_PATH = f\"{TEST_FILES_BASE}/with_meta.zip\"\nTEST_NO_META_ZIP_PATH = f\"{TEST_FILES_BASE}/without_meta.zip\"\nTEST_METADATA_FILE = f\"{TEST_FILES_BASE}/.onyx_metadata.json\"\n\n\n@pytest.mark.parametrize(\n \"zip_path, has_metadata\",\n [\n (TEST_META_ZIP_PATH, True),\n (TEST_NO_META_ZIP_PATH, False),\n ],\n)\ndef test_zip_metadata_handling(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture, # noqa: ARG001\n zip_path: str,\n has_metadata: bool,\n) -> None:\n before = datetime.now(timezone.utc)\n # Create an admin user\n admin_user: DATestUser = UserManager.create(\n email=\"admin@example.com\",\n )\n\n # Upload the test zip file (simulate this happening from frontend)\n upload_response = FileManager.upload_file_for_connector(\n file_path=zip_path,\n file_name=os.path.basename(zip_path),\n user_performing_action=admin_user,\n content_type=\"application/zip\",\n )\n\n file_paths = upload_response.file_paths\n assert file_paths, \"File upload failed - no file paths returned\"\n if has_metadata:\n zip_metadata_file_id = upload_response.zip_metadata_file_id\n assert zip_metadata_file_id, \"Metadata file ID should be present\"\n else:\n zip_metadata_file_id = None\n\n # Create a dummy credential for the file connector\n credential = CredentialManager.create(\n source=DocumentSource.FILE,\n credential_json={},\n user_performing_action=admin_user,\n )\n\n # Create the connector\n connector_name = f\"FileConnector-{int(datetime.now().timestamp())}\"\n connector = ConnectorManager.create(\n name=connector_name,\n source=DocumentSource.FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={\n \"file_locations\": file_paths,\n \"file_names\": [os.path.basename(file_path) for file_path in file_paths],\n \"zip_metadata_file_id\": zip_metadata_file_id,\n },\n access_type=AccessType.PUBLIC,\n groups=[],\n user_performing_action=admin_user,\n )\n\n # Link the credential to the connector\n cc_pair = CCPairManager.create(\n credential_id=credential.id,\n connector_id=connector.id,\n access_type=AccessType.PUBLIC,\n user_performing_action=admin_user,\n )\n\n # Run the connector to index the files\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair, after=before, user_performing_action=admin_user\n )\n\n # Get the indexed documents\n with get_session_with_current_tenant() as db_session:\n documents = get_documents_for_cc_pair(db_session, cc_pair.id)\n\n # Expected metadata from the .onyx_metadata.json file\n with open(TEST_METADATA_FILE, \"r\") as f:\n expected_metadata = json.load(f)\n\n # Verify each document has the correct metadata\n for doc in documents:\n filename = doc.semantic_id\n if filename in expected_metadata:\n expected = expected_metadata[filename]\n assert (\n doc.semantic_id == expected[\"display_name\"]\n ), f\"Display name mismatch for {filename}\"\n assert doc.link == expected[\"link\"], f\"Link mismatch for {filename}\"\n" }, { "path": "backend/tests/integration/tests/indexing/file_connector/test_files/.onyx_metadata.json", "content": "[\n {\n \"filename\": \"sample1.txt\",\n \"link\": \"https://www.google.com\",\n \"file_display_name\": \"Basically Google\",\n \"primary_owners\": [\"evan@onyx.app\"],\n \"status\": \"bingle bongle\"\n },\n {\n \"filename\": \"sample2.txt\",\n \"link\": \"https://www.youtube.com\",\n \"file_display_name\": \"Pretty much youtube\",\n \"primary_owners\": [\"chris@onyx.app\"],\n \"status\": \"not bingle bongle\"\n }\n]" }, { "path": "backend/tests/integration/tests/indexing/file_connector/test_files/sample1.txt", "content": "The following contains some excerpts from our docs.\n\nThe File Connector indexes user uploaded files. Currently supports .txt, .pdf, .docx, .pptx, .xlsx, .csv, .md, .mdx, .conf, .log, .json, .tsv, .xml, .yml, .yaml, .eml, and .epub files. \nYou can also upload a .zip containing these files - If there are other file types in the zip, the other file types are ignored. \nThere is also an optional metadata line that supports links, document owners, and time updated as metadata for Onyx’s retrieval and AI Answer.\n\nThe metadata line should be placed at the very top of the file and can take one of two formats:\n\n#ONYX_METADATA={\"link\": \"\"}\n\nWhere ONYX_METADATA= is followed by a json. The valid json keys are:\n\nlink\nprimary_owners\nsecondary_owners\ndoc_updated_at\nfile_display_name\nYou can also include arbitrary key/value pairs which will be understood as “tags”. \nThese tags can then be used in the UI as a filter if you want to constrain your search / conversation to only documents with certain tag(s) attached" }, { "path": "backend/tests/integration/tests/indexing/file_connector/test_files/sample2.txt", "content": "Hello, I hope you're having a wonderful day!" }, { "path": "backend/tests/integration/tests/indexing/test_checkpointing.py", "content": "import uuid\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nimport httpx\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.mock_connector.connector import MockConnectorCheckpoint\nfrom onyx.connectors.models import ConnectorFailure\nfrom onyx.connectors.models import EntityFailure\nfrom onyx.connectors.models import InputType\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import IndexingStatus\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.index_attempt import IndexAttemptManager\nfrom tests.integration.common_utils.test_document_utils import create_test_document\nfrom tests.integration.common_utils.test_document_utils import (\n create_test_document_failure,\n)\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef test_mock_connector_basic_flow(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture,\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that the mock connector can successfully process documents and failures\"\"\"\n # Set up mock server behavior\n doc_uuid = uuid.uuid4()\n test_doc = create_test_document(doc_id=f\"test-doc-{doc_uuid}\")\n\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [test_doc.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(\n mode=\"json\"\n ),\n \"failures\": [],\n }\n ],\n )\n assert response.status_code == 200\n\n # create CC Pair + index attempt\n cc_pair = CCPairManager.create_from_scratch(\n name=f\"mock-connector-{uuid.uuid4()}\",\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n user_performing_action=admin_user,\n )\n\n # wait for index attempt to start\n index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # wait for index attempt to finish\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # validate status\n finished_index_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_index_attempt.status == IndexingStatus.SUCCESS\n\n # Verify results\n with get_session_with_current_tenant() as db_session:\n chunks = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(chunks) == 1\n assert chunks[0].id == test_doc.id\n\n errors = IndexAttemptManager.get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert len(errors) == 0\n\n\ndef test_mock_connector_with_failures(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture,\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that the mock connector processes both successes and failures properly.\"\"\"\n doc1 = create_test_document()\n doc2 = create_test_document()\n doc2_failure = create_test_document_failure(doc_id=doc2.id)\n\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [doc1.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(\n mode=\"json\"\n ),\n \"failures\": [doc2_failure.model_dump(mode=\"json\")],\n }\n ],\n )\n assert response.status_code == 200\n\n # Create a CC Pair for the mock connector\n cc_pair = CCPairManager.create_from_scratch(\n name=f\"mock-connector-failure-{uuid.uuid4()}\",\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n user_performing_action=admin_user,\n )\n\n # Wait for the index attempt to start and then complete\n index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # validate status\n finished_index_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_index_attempt.status == IndexingStatus.COMPLETED_WITH_ERRORS\n\n # Verify results: doc1 should be indexed and doc2 should have an error entry\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(documents) == 1\n assert documents[0].id == doc1.id\n\n errors = IndexAttemptManager.get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert len(errors) == 1\n error = errors[0]\n assert error.failure_message == doc2_failure.failure_message\n assert error.document_id == doc2.id\n\n\ndef test_mock_connector_failure_recovery(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture,\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that a failed document can be successfully indexed in a subsequent attempt\n while maintaining previously successful documents.\"\"\"\n # Create test documents and failure\n doc1 = create_test_document()\n doc2 = create_test_document()\n doc2_failure = create_test_document_failure(doc_id=doc2.id)\n entity_id = \"test-entity-id\"\n entity_failure_msg = \"Simulated unhandled error\"\n\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [doc1.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(\n mode=\"json\"\n ),\n \"failures\": [\n doc2_failure.model_dump(mode=\"json\"),\n ConnectorFailure(\n failed_entity=EntityFailure(\n entity_id=entity_id,\n missed_time_range=(\n datetime.now(timezone.utc) - timedelta(days=1),\n datetime.now(timezone.utc),\n ),\n ),\n failure_message=entity_failure_msg,\n ).model_dump(mode=\"json\"),\n ],\n }\n ],\n )\n assert response.status_code == 200\n\n # Create CC Pair and run initial indexing attempt\n cc_pair = CCPairManager.create_from_scratch(\n name=f\"mock-connector-{uuid.uuid4()}\",\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n user_performing_action=admin_user,\n )\n\n # Wait for first index attempt to complete\n initial_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=initial_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # validate status\n finished_index_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=initial_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_index_attempt.status == IndexingStatus.COMPLETED_WITH_ERRORS\n\n # Verify initial state: doc1 indexed, doc2 failed\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(documents) == 1\n assert documents[0].id == doc1.id\n\n errors = IndexAttemptManager.get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert len(errors) == 2\n error_doc2 = next(error for error in errors if error.document_id == doc2.id)\n assert error_doc2.failure_message == doc2_failure.failure_message\n assert not error_doc2.is_resolved\n\n error_entity = next(error for error in errors if error.entity_id == entity_id)\n assert error_entity.failure_message == entity_failure_msg\n assert not error_entity.is_resolved\n\n # Update mock server to return success for both documents\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [\n doc1.model_dump(mode=\"json\"),\n doc2.model_dump(mode=\"json\"),\n ],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(\n mode=\"json\"\n ),\n \"failures\": [],\n }\n ],\n )\n assert response.status_code == 200\n\n # Trigger another indexing attempt\n # NOTE: must be from beginning to handle the entity failure\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n recovery_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n index_attempts_to_ignore=[initial_index_attempt.id],\n user_performing_action=admin_user,\n )\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=recovery_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n finished_second_index_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=recovery_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_second_index_attempt.status == IndexingStatus.SUCCESS\n\n # Verify both documents are now indexed\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(documents) == 2\n document_ids = {doc.id for doc in documents}\n assert doc2.id in document_ids\n assert doc1.id in document_ids\n\n # Verify original failures were marked as resolved\n errors = IndexAttemptManager.get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert len(errors) == 2\n error_doc2 = next(error for error in errors if error.document_id == doc2.id)\n error_entity = next(error for error in errors if error.entity_id == entity_id)\n\n assert error_doc2.is_resolved\n assert error_entity.is_resolved\n\n\ndef test_mock_connector_checkpoint_recovery(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture,\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that checkpointing works correctly when an unhandled exception occurs\n and that subsequent runs pick up from the last successful checkpoint.\"\"\"\n # Create test documents\n docs_batch_1 = [create_test_document() for _ in range(100)]\n doc2 = create_test_document()\n doc3 = create_test_document()\n\n # Set up mock server behavior for initial run:\n # - First yield: 100 docs with checkpoint1\n # - Second yield: doc2 with checkpoint2\n # - Third yield: unhandled exception\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [doc.model_dump(mode=\"json\") for doc in docs_batch_1],\n \"checkpoint\": MockConnectorCheckpoint(\n has_more=True, last_document_id=docs_batch_1[-1].id\n ).model_dump(mode=\"json\"),\n \"failures\": [],\n },\n {\n \"documents\": [doc2.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(\n has_more=True, last_document_id=doc2.id\n ).model_dump(mode=\"json\"),\n \"failures\": [],\n },\n {\n \"documents\": [],\n # should never hit this, unhandled exception happens first\n \"checkpoint\": MockConnectorCheckpoint(\n has_more=False, last_document_id=doc2.id\n ).model_dump(mode=\"json\"),\n \"failures\": [],\n \"unhandled_exception\": \"Simulated unhandled error\",\n },\n ],\n )\n assert response.status_code == 200\n\n # Create CC Pair and run initial indexing attempt\n # Note: Setting refresh_freq to allow manual retrigger after failure\n cc_pair = CCPairManager.create_from_scratch(\n name=f\"mock-connector-checkpoint-{uuid.uuid4()}\",\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n user_performing_action=admin_user,\n refresh_freq=60 * 60, # 1 hour\n )\n\n # Wait for first index attempt to complete\n initial_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=initial_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # validate status\n finished_index_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=initial_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_index_attempt.status == IndexingStatus.FAILED\n\n # Pause the connector immediately to prevent check_for_indexing from\n # creating automatic retry attempts while we reset the mock server.\n # Without this, the INITIAL_INDEXING status causes immediate retries\n # that would consume (or fail against) the mock server before we can\n # set up the recovery behavior.\n CCPairManager.pause_cc_pair(cc_pair, user_performing_action=admin_user)\n\n # Collect all index attempt IDs created so far (the initial one plus\n # any automatic retries that may have started before the pause took effect).\n all_prior_attempt_ids: list[int] = []\n index_attempts_page = IndexAttemptManager.get_index_attempt_page(\n cc_pair_id=cc_pair.id,\n page=0,\n page_size=100,\n user_performing_action=admin_user,\n )\n all_prior_attempt_ids = [ia.id for ia in index_attempts_page.items]\n\n # Verify initial state: both docs should be indexed\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n # This is no longer guaranteed because docfetching and docprocessing are decoupled!\n # Some batches may not be processed when docfetching fails, but they should still stick around\n # in the filestore and be ready for the next run.\n # assert len(documents) == 101 # 100 docs from first batch + doc2\n # document_ids = {doc.id for doc in documents}\n # assert doc2.id in document_ids\n # assert all(doc.id in document_ids for doc in docs_batch_1)\n\n # Get the checkpoints that were sent to the mock server\n response = mock_server_client.get(\"/get-checkpoints\")\n assert response.status_code == 200\n initial_checkpoints = response.json()\n\n # Verify we got the expected checkpoints in order\n assert len(initial_checkpoints) == 3\n assert initial_checkpoints[0] == {\n \"has_more\": True,\n \"last_document_id\": None,\n } # Initial empty checkpoint\n assert initial_checkpoints[1] == {\n \"has_more\": True,\n \"last_document_id\": docs_batch_1[-1].id,\n }\n assert initial_checkpoints[2] == {\"has_more\": True, \"last_document_id\": doc2.id}\n\n # Reset the mock server for the next run\n response = mock_server_client.post(\"/reset\")\n assert response.status_code == 200\n\n # Set up mock server behavior for recovery run - should succeed fully this time\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [doc3.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(\n has_more=False, last_document_id=doc3.id\n ).model_dump(mode=\"json\"),\n \"failures\": [],\n }\n ],\n )\n assert response.status_code == 200\n\n # Set the manual indexing trigger, then unpause to allow the recovery run.\n CCPairManager.run_once(\n cc_pair, from_beginning=False, user_performing_action=admin_user\n )\n CCPairManager.unpause_cc_pair(cc_pair, user_performing_action=admin_user)\n recovery_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n index_attempts_to_ignore=all_prior_attempt_ids,\n user_performing_action=admin_user,\n )\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=recovery_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # validate status\n finished_recovery_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=recovery_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_recovery_attempt.status == IndexingStatus.SUCCESS\n\n # Verify results\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(documents) == 102 # 100 docs from first batch + doc2 + doc3\n document_ids = {doc.id for doc in documents}\n assert doc3.id in document_ids\n assert doc2.id in document_ids\n assert all(doc.id in document_ids for doc in docs_batch_1)\n\n # Get the checkpoints from the recovery run\n response = mock_server_client.get(\"/get-checkpoints\")\n assert response.status_code == 200\n recovery_checkpoints = response.json()\n\n # Verify the recovery run started from the last successful checkpoint\n assert len(recovery_checkpoints) == 1\n assert recovery_checkpoints[0] == {\"has_more\": True, \"last_document_id\": doc2.id}\n" }, { "path": "backend/tests/integration/tests/indexing/test_initial_permission_sync.py", "content": "import os\nimport uuid\nfrom datetime import datetime\nfrom datetime import timezone\n\nimport httpx\nimport pytest\nfrom sqlalchemy import select\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.mock_connector.connector import EXTERNAL_USER_EMAILS\nfrom onyx.connectors.mock_connector.connector import EXTERNAL_USER_GROUP_IDS\nfrom onyx.connectors.mock_connector.connector import MockConnectorCheckpoint\nfrom onyx.connectors.models import Document\nfrom onyx.connectors.models import InputType\nfrom onyx.db.document import get_documents_by_ids\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import AccessType\nfrom onyx.db.enums import IndexingStatus\nfrom onyx.db.enums import PermissionSyncStatus\nfrom onyx.db.models import DocPermissionSyncAttempt\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.index_attempt import IndexAttemptManager\nfrom tests.integration.common_utils.test_document_utils import create_test_document\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef _setup_mock_connector(\n mock_server_client: httpx.Client,\n admin_user: DATestUser,\n) -> tuple[DATestCCPair, Document]:\n \"\"\"Common setup: create a test doc, configure mock server, create cc_pair, wait for indexing.\"\"\"\n doc_uuid = uuid.uuid4()\n test_doc = create_test_document(doc_id=f\"test-doc-{doc_uuid}\")\n\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[\n {\n \"documents\": [test_doc.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(\n mode=\"json\"\n ),\n \"failures\": [],\n }\n ],\n )\n assert response.status_code == 200\n\n cc_pair = CCPairManager.create_from_scratch(\n name=f\"mock-connector-{uuid.uuid4()}\",\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n access_type=AccessType.SYNC,\n user_performing_action=admin_user,\n )\n\n index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n finished = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished.status == IndexingStatus.SUCCESS\n return cc_pair, test_doc\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission sync is enterprise only\",\n)\ndef test_mock_connector_initial_permission_sync(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture,\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that the MockConnector fetches and sets permissions during initial indexing\n when AccessType.SYNC is used.\"\"\"\n\n cc_pair, test_doc = _setup_mock_connector(mock_server_client, admin_user)\n\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(documents) == 1\n assert documents[0].id == test_doc.id\n\n errors = IndexAttemptManager.get_index_attempt_errors_for_cc_pair(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert len(errors) == 0\n\n with get_session_with_current_tenant() as db_session:\n db_docs = get_documents_by_ids(\n db_session=db_session,\n document_ids=[test_doc.id],\n )\n assert len(db_docs) == 1\n db_doc = db_docs[0]\n\n assert db_doc.external_user_emails is not None\n assert db_doc.external_user_group_ids is not None\n assert set(db_doc.external_user_emails) == EXTERNAL_USER_EMAILS\n assert set(db_doc.external_user_group_ids) == EXTERNAL_USER_GROUP_IDS\n assert db_doc.is_public is False\n\n # After initial indexing, the beat task detects last_time_perm_sync is None\n # and triggers a doc permission sync. Explicitly trigger it to avoid\n # waiting for the 30s beat interval.\n before = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n should_wait_for_group_sync=False,\n should_wait_for_vespa_sync=False,\n )\n\n updated_cc_pair_info = CCPairManager.get_single(\n cc_pair.id, user_performing_action=admin_user\n )\n assert updated_cc_pair_info is not None\n assert updated_cc_pair_info.last_full_permission_sync is not None\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission sync attempt tracking is enterprise only\",\n)\ndef test_permission_sync_attempt_tracking_integration(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that permission sync attempts are properly tracked during real sync workflows.\"\"\"\n\n cc_pair, _test_doc = _setup_mock_connector(mock_server_client, admin_user)\n\n before = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n should_wait_for_group_sync=False,\n should_wait_for_vespa_sync=False,\n )\n\n with get_session_with_current_tenant() as db_session:\n attempt = db_session.execute(\n select(DocPermissionSyncAttempt).where(\n DocPermissionSyncAttempt.connector_credential_pair_id == cc_pair.id\n )\n ).scalar_one()\n\n assert attempt.status in [\n PermissionSyncStatus.SUCCESS,\n PermissionSyncStatus.COMPLETED_WITH_ERRORS,\n PermissionSyncStatus.FAILED,\n ]\n assert attempt.total_docs_synced is not None and attempt.total_docs_synced >= 0\n assert (\n attempt.docs_with_permission_errors is not None\n and attempt.docs_with_permission_errors >= 0\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission sync attempt tracking is enterprise only\",\n)\ndef test_permission_sync_attempt_status_success(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that permission sync attempts are marked as SUCCESS when sync completes without errors.\"\"\"\n\n cc_pair, _test_doc = _setup_mock_connector(mock_server_client, admin_user)\n\n before = datetime.now(timezone.utc)\n CCPairManager.sync(\n cc_pair=cc_pair,\n user_performing_action=admin_user,\n )\n\n CCPairManager.wait_for_sync(\n cc_pair=cc_pair,\n after=before,\n number_of_updated_docs=1,\n user_performing_action=admin_user,\n should_wait_for_group_sync=False,\n should_wait_for_vespa_sync=False,\n )\n\n with get_session_with_current_tenant() as db_session:\n attempt = db_session.execute(\n select(DocPermissionSyncAttempt).where(\n DocPermissionSyncAttempt.connector_credential_pair_id == cc_pair.id\n )\n ).scalar_one()\n\n assert attempt.status == PermissionSyncStatus.SUCCESS\n assert attempt.total_docs_synced is not None and attempt.total_docs_synced >= 0\n assert (\n attempt.docs_with_permission_errors is not None\n and attempt.docs_with_permission_errors == 0\n )\n" }, { "path": "backend/tests/integration/tests/indexing/test_polling.py", "content": "import uuid\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nimport httpx\n\nfrom onyx.configs.app_configs import POLL_CONNECTOR_OFFSET\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.mock_connector.connector import MockConnectorCheckpoint\nfrom onyx.connectors.models import InputType\nfrom onyx.db.enums import IndexingStatus\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.index_attempt import IndexAttemptManager\nfrom tests.integration.common_utils.test_document_utils import create_test_document\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _setup_mock_connector(\n mock_server_client: httpx.Client,\n admin_user: DATestUser, # noqa: ARG001\n) -> None:\n test_doc = create_test_document()\n successful_response = {\n \"documents\": [test_doc.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(mode=\"json\"),\n \"failures\": [],\n }\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[successful_response, successful_response], # For two attempts\n )\n assert response.status_code == 200\n\n\ndef test_poll_connector_time_ranges(\n mock_server_client: httpx.Client,\n admin_user: DATestUser,\n) -> None:\n \"\"\"\n Tests that poll connectors correctly set their poll_range_start and poll_range_end\n across multiple indexing attempts.\n \"\"\"\n # Set up mock server behavior - a simple successful response\n _setup_mock_connector(mock_server_client, admin_user)\n\n # Create a CC Pair for the mock connector with POLL input type\n cc_pair_name = f\"mock-poll-time-range-{uuid.uuid4()}\"\n cc_pair = CCPairManager.create_from_scratch(\n name=cc_pair_name,\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n user_performing_action=admin_user,\n refresh_freq=3, # refresh often to ensure the second attempt actually runs\n )\n\n # --- First Indexing Attempt ---\n time_before_first_attempt = datetime.now(timezone.utc)\n first_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=first_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n time_after_first_attempt = datetime.now(timezone.utc)\n\n # Fetch and validate the first attempt\n completed_first_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=first_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert completed_first_attempt.status == IndexingStatus.SUCCESS\n assert completed_first_attempt.poll_range_start is not None\n assert completed_first_attempt.poll_range_end is not None\n\n # For the first run (no prior successful attempts), poll_range_start should be epoch (0)\n expected_first_start = datetime.fromtimestamp(0, tz=timezone.utc)\n assert completed_first_attempt.poll_range_start == expected_first_start\n\n # `poll_range_end` should be sometime in between the time the attempt\n # started and the time it finished.\n # no way to have a more precise assertion here since the `poll_range_end`\n # can really be set anytime in that range and be \"correct\"\n assert (\n time_before_first_attempt\n <= completed_first_attempt.poll_range_end\n <= time_after_first_attempt\n )\n\n first_attempt_poll_end = completed_first_attempt.poll_range_end\n\n # --- Second Indexing Attempt ---\n # Trigger another run manually (since automatic refresh might be too slow for test)\n # Ensure there's a slight delay so the poll window moves\n # In a real scenario, the scheduler would wait for the refresh frequency.\n # Here we manually trigger a new run.\n _setup_mock_connector(mock_server_client, admin_user)\n CCPairManager.run_once(\n cc_pair, from_beginning=False, user_performing_action=admin_user\n )\n\n time_before_second_attempt = datetime.now(timezone.utc)\n second_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n index_attempts_to_ignore=[first_index_attempt.id],\n user_performing_action=admin_user,\n )\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=second_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n time_after_second_attempt = datetime.now(timezone.utc)\n\n # Fetch and validate the second attempt\n completed_second_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=second_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert completed_second_attempt.status == IndexingStatus.SUCCESS\n assert completed_second_attempt.poll_range_start is not None\n assert completed_second_attempt.poll_range_end is not None\n\n # For the second run, poll_range_start should be the previous successful attempt's\n # poll_range_end minus the POLL_CONNECTOR_OFFSET\n expected_second_start = first_attempt_poll_end - timedelta(\n minutes=POLL_CONNECTOR_OFFSET\n )\n assert completed_second_attempt.poll_range_start == expected_second_start\n\n # `poll_range_end` should be sometime in between the time the attempt\n # started and the time it finished.\n # again, no way to have a more precise assertion here since the `poll_range_end`\n # can really be set anytime in that range and be \"correct\"\n assert (\n time_before_second_attempt\n <= completed_second_attempt.poll_range_end\n <= time_after_second_attempt\n )\n" }, { "path": "backend/tests/integration/tests/indexing/test_repeated_error_state.py", "content": "import time\nimport uuid\n\nimport httpx\n\nfrom onyx.background.celery.tasks.docprocessing.utils import (\n NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE,\n)\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.mock_connector.connector import MockConnectorCheckpoint\nfrom onyx.connectors.models import InputType\nfrom onyx.db.connector_credential_pair import get_connector_credential_pair_from_id\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.enums import IndexingStatus\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_HOST\nfrom tests.integration.common_utils.constants import MOCK_CONNECTOR_SERVER_PORT\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.index_attempt import IndexAttemptManager\nfrom tests.integration.common_utils.test_document_utils import create_test_document\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef test_repeated_error_state_detection_and_recovery(\n mock_server_client: httpx.Client,\n vespa_client: vespa_fixture,\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that a connector is marked as in a repeated error state after\n NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE consecutive failures, and\n that it recovers after a successful indexing.\n\n This test ensures we properly wait for the required number of indexing attempts\n to fail before checking that the connector is in a repeated error state.\"\"\"\n\n # Create test document for successful response\n test_doc = create_test_document()\n\n # First, set up the mock server to consistently fail\n error_response = {\n \"documents\": [],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(mode=\"json\"),\n \"failures\": [],\n \"unhandled_exception\": \"Simulated unhandled error for testing repeated errors\",\n }\n\n # Create a list of failure responses with at least the same length\n # as NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE\n failure_behaviors = [error_response] * (\n 5 * NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE\n )\n\n response = mock_server_client.post(\n \"/set-behavior\",\n json=failure_behaviors,\n )\n assert response.status_code == 200\n\n # Create a new CC pair for testing\n cc_pair = CCPairManager.create_from_scratch(\n name=f\"mock-repeated-error-{uuid.uuid4()}\",\n source=DocumentSource.MOCK_CONNECTOR,\n input_type=InputType.POLL,\n connector_specific_config={\n \"mock_server_host\": MOCK_CONNECTOR_SERVER_HOST,\n \"mock_server_port\": MOCK_CONNECTOR_SERVER_PORT,\n },\n user_performing_action=admin_user,\n refresh_freq=60 * 60, # a very long time\n )\n\n # Wait for the required number of failed indexing attempts\n # This shouldn't take long, since we keep retrying while we haven't\n # succeeded yet\n start_time = time.monotonic()\n while True:\n index_attempts_page = IndexAttemptManager.get_index_attempt_page(\n cc_pair_id=cc_pair.id,\n page=0,\n page_size=100,\n user_performing_action=admin_user,\n )\n index_attempts = [\n ia\n for ia in index_attempts_page.items\n if ia.status and ia.status.is_terminal()\n ]\n if len(index_attempts) >= NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE:\n break\n\n if time.monotonic() - start_time > 180:\n raise TimeoutError(\n \"Did not get required number of failed attempts within 180 seconds\"\n )\n\n # make sure that we don't mark the connector as in repeated error state\n # before we have the required number of failed attempts\n with get_session_with_current_tenant() as db_session:\n cc_pair_obj = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair.id,\n )\n assert cc_pair_obj is not None\n assert not cc_pair_obj.in_repeated_error_state\n\n time.sleep(2)\n\n # Verify we have the correct number of failed attempts\n assert len(index_attempts) == NUM_REPEAT_ERRORS_BEFORE_REPEATED_ERROR_STATE\n for attempt in index_attempts:\n assert attempt.status == IndexingStatus.FAILED\n\n # Check if the connector is in a repeated error state\n start_time = time.monotonic()\n while True:\n with get_session_with_current_tenant() as db_session:\n cc_pair_obj = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair.id,\n )\n assert cc_pair_obj is not None\n if cc_pair_obj.in_repeated_error_state:\n # Pausing only happens for cloud deployments and the IT don't run with\n # that auth type :(\n # if AUTH_TYPE == AuthType.CLOUD:\n # assert cc_pair_obj.status == ConnectorCredentialPairStatus.PAUSED, (\n # f\"Expected status to be PAUSED when in repeated error state, \"\n # f\"but got {cc_pair_obj.status}\"\n # )\n break\n\n if time.monotonic() - start_time > 90:\n assert False, \"CC pair did not enter repeated error state within 90 seconds\"\n\n time.sleep(2)\n\n # Reset the mock server state\n response = mock_server_client.post(\"/reset\")\n assert response.status_code == 200\n\n # Now set up the mock server to succeed\n success_response = {\n \"documents\": [test_doc.model_dump(mode=\"json\")],\n \"checkpoint\": MockConnectorCheckpoint(has_more=False).model_dump(mode=\"json\"),\n \"failures\": [],\n }\n\n response = mock_server_client.post(\n \"/set-behavior\",\n json=[success_response],\n )\n assert response.status_code == 200\n\n # Set the manual indexing trigger first (while paused), then unpause.\n # This ensures the trigger is set before CHECK_FOR_INDEXING runs, which will\n # prevent the connector from being re-paused when repeated error state is detected.\n CCPairManager.run_once(\n cc_pair, from_beginning=True, user_performing_action=admin_user\n )\n CCPairManager.unpause_cc_pair(cc_pair, user_performing_action=admin_user)\n\n recovery_index_attempt = IndexAttemptManager.wait_for_index_attempt_start(\n cc_pair_id=cc_pair.id,\n index_attempts_to_ignore=[index_attempt.id for index_attempt in index_attempts],\n user_performing_action=admin_user,\n )\n\n IndexAttemptManager.wait_for_index_attempt_completion(\n index_attempt_id=recovery_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n\n # Validate the indexing succeeded\n finished_recovery_attempt = IndexAttemptManager.get_index_attempt_by_id(\n index_attempt_id=recovery_index_attempt.id,\n cc_pair_id=cc_pair.id,\n user_performing_action=admin_user,\n )\n assert finished_recovery_attempt.status == IndexingStatus.SUCCESS\n\n # Verify the document was indexed\n with get_session_with_current_tenant() as db_session:\n documents = DocumentManager.fetch_documents_for_cc_pair(\n cc_pair_id=cc_pair.id,\n db_session=db_session,\n vespa_client=vespa_client,\n )\n assert len(documents) == 1\n assert documents[0].id == test_doc.id\n\n # Verify the CC pair is no longer in a repeated error state\n start = time.monotonic()\n while True:\n with get_session_with_current_tenant() as db_session:\n cc_pair_obj = get_connector_credential_pair_from_id(\n db_session=db_session,\n cc_pair_id=cc_pair.id,\n )\n assert cc_pair_obj is not None\n if not cc_pair_obj.in_repeated_error_state:\n break\n\n elapsed = time.monotonic() - start\n if elapsed > 30:\n raise TimeoutError(\n \"CC pair did not exit repeated error state within 30 seconds\"\n )\n\n print(\n f\"Waiting for CC pair to exit repeated error state. elapsed={elapsed:.2f}\"\n )\n time.sleep(1)\n" }, { "path": "backend/tests/integration/tests/ingestion/test_ingestion_api.py", "content": "from onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import Document\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import IngestionManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\n\ndef test_ingestion_api_crud(\n reset: None, # noqa: ARG001\n vespa_client: vespa_fixture,\n) -> None:\n \"\"\"Test create, list, and delete via the ingestion API.\"\"\"\n admin_user: DATestUser = UserManager.create(email=\"admin@onyx.app\")\n cc_pair = CCPairManager.create_from_scratch(\n name=\"Ingestion-API-Test\",\n source=DocumentSource.FILE,\n input_type=InputType.LOAD_STATE,\n connector_specific_config={\n \"file_locations\": [],\n \"file_names\": [],\n \"zip_metadata_file_id\": None,\n },\n user_performing_action=admin_user,\n )\n api_key = APIKeyManager.create(user_performing_action=admin_user)\n api_key.headers.update(admin_user.headers)\n\n # CREATE\n doc = IngestionManager.seed_doc_with_content(\n cc_pair=cc_pair,\n content=\"Test document\",\n document_id=\"test-doc-1\",\n api_key=api_key,\n )\n\n with get_session_with_current_tenant() as db_session:\n doc_db = db_session.query(Document).filter(Document.id == doc.id).first()\n assert doc_db is not None\n assert doc_db.from_ingestion_api is True\n\n vespa_docs = vespa_client.get_documents_by_id([doc.id])[\"documents\"]\n assert len(vespa_docs) == 1\n\n # LIST\n docs_list = IngestionManager.list_all_ingestion_docs(api_key=api_key)\n assert any(d[\"document_id\"] == doc.id for d in docs_list)\n\n # DELETE\n IngestionManager.delete(document_id=doc.id, api_key=api_key)\n\n with get_session_with_current_tenant() as db_session:\n doc_db = db_session.query(Document).filter(Document.id == doc.id).first()\n assert doc_db is None\n\n vespa_docs = vespa_client.get_documents_by_id([doc.id])[\"documents\"]\n assert len(vespa_docs) == 0\n" }, { "path": "backend/tests/integration/tests/kg/test_kg_api.py", "content": "import json\nfrom datetime import datetime\nfrom http import HTTPStatus\n\nimport pytest\nimport requests\n\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.connectors.models import InputType\nfrom onyx.db.connector import create_connector\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.kg_config import get_kg_config_settings\nfrom onyx.db.kg_config import set_kg_config_settings\nfrom onyx.db.models import Connector\nfrom onyx.server.documents.models import ConnectorBase\nfrom onyx.server.kg.models import DisableKGConfigRequest\nfrom onyx.server.kg.models import EnableKGConfigRequest\nfrom onyx.server.kg.models import EntityType\nfrom onyx.server.kg.models import KGConfig as KGConfigAPIModel\nfrom onyx.server.kg.models import SourceAndEntityTypeView\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.reset import reset_all\n\n\n@pytest.fixture(autouse=True)\ndef reset_for_test() -> None:\n \"\"\"Reset all data before each test.\"\"\"\n reset_all()\n\n kg_config_settings = get_kg_config_settings()\n kg_config_settings.KG_EXPOSED = True\n set_kg_config_settings(kg_config_settings)\n\n\n@pytest.fixture()\ndef connectors() -> None:\n \"\"\"Set up connectors for tests.\"\"\"\n with get_session_with_current_tenant() as db_session:\n # Create Salesforce connector\n connector_data = ConnectorBase(\n name=\"Salesforce Test\",\n source=DocumentSource.SALESFORCE,\n input_type=InputType.POLL,\n connector_specific_config={},\n refresh_freq=None,\n indexing_start=None,\n prune_freq=None,\n )\n create_connector(db_session, connector_data)\n\n\ndef test_kg_enable_and_disable(connectors: None) -> None: # noqa: ARG001\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Enable KG\n # Need to `.model_dump_json()` and then `json.loads`.\n # Seems redundant, but this is because simply calling `json=data.model_dump()`\n # returns in a \"datetime cannot be JSON serialized error\".\n req1 = json.loads(\n EnableKGConfigRequest(\n vendor=\"Test\",\n vendor_domains=[\"test.app\", \"tester.ai\"],\n ignore_domains=[],\n coverage_start=datetime(1970, 1, 1, 0, 0),\n ).model_dump_json()\n )\n res1 = requests.put(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n json=req1,\n )\n assert (\n res1.status_code == HTTPStatus.OK\n ), f\"Error response: {res1.status_code} - {res1.text}\"\n\n # Check KG\n res2 = requests.get(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n )\n assert (\n res2.status_code == HTTPStatus.OK\n ), f\"Error response: {res2.status_code} - {res2.text}\"\n\n actual_config = KGConfigAPIModel.model_validate_json(res2.text)\n assert actual_config == KGConfigAPIModel(\n enabled=True,\n vendor=\"Test\",\n vendor_domains=[\"test.app\", \"tester.ai\"],\n ignore_domains=[],\n coverage_start=datetime(1970, 1, 1, 0, 0),\n )\n\n # Disable KG\n req3 = DisableKGConfigRequest().model_dump()\n res3 = requests.put(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n json=req3,\n )\n assert (\n res3.status_code == HTTPStatus.OK\n ), f\"Error response: {res3.status_code} - {res3.text}\"\n\n # Check KG\n res4 = requests.get(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n )\n assert (\n res4.status_code == HTTPStatus.OK\n ), f\"Error response: {res4.status_code} - {res4.text}\"\n\n actual_config = KGConfigAPIModel.model_validate_json(res4.text)\n assert actual_config == KGConfigAPIModel(\n enabled=False,\n vendor=\"Test\",\n vendor_domains=[\"test.app\", \"tester.ai\"],\n ignore_domains=[],\n coverage_start=datetime(1970, 1, 1, 0, 0),\n )\n\n\ndef test_kg_enable_with_missing_fields_should_fail() -> None:\n admin_user = UserManager.create(name=\"admin_user\")\n\n req = json.loads(\n EnableKGConfigRequest(\n vendor=\"Test\",\n vendor_domains=[],\n ignore_domains=[],\n coverage_start=datetime(1970, 1, 1, 0, 0),\n ).model_dump_json()\n )\n res = requests.put(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n json=req,\n )\n assert res.status_code == HTTPStatus.BAD_REQUEST\n\n\ndef test_update_kg_entity_types(connectors: None) -> None: # noqa: ARG001\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Enable kg and populate default entity types\n req1 = json.loads(\n EnableKGConfigRequest(\n vendor=\"Test\",\n vendor_domains=[\"test.app\", \"tester.ai\"],\n ignore_domains=[],\n coverage_start=datetime(1970, 1, 1, 0, 0),\n ).model_dump_json()\n )\n res1 = requests.put(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n json=req1,\n )\n assert (\n res1.status_code == HTTPStatus.OK\n ), f\"Error response: {res1.status_code} - {res1.text}\"\n\n # Get old entity types\n res2 = requests.get(\n f\"{API_SERVER_URL}/admin/kg/entity-types\",\n headers=admin_user.headers,\n )\n assert (\n res2.status_code == HTTPStatus.OK\n ), f\"Error response: {res2.status_code} - {res2.text}\"\n res2_parsed = SourceAndEntityTypeView.model_validate(res2.json())\n\n # Update entity types\n req3 = [\n EntityType(\n name=\"ACCOUNT\",\n description=\"Test.\",\n active=True,\n grounded_source_name=\"salesforce\",\n ).model_dump(),\n EntityType(\n name=\"OPPORTUNITY\",\n description=\"Test 2.\",\n active=False,\n ).model_dump(),\n ]\n res3 = requests.put(\n f\"{API_SERVER_URL}/admin/kg/entity-types\",\n headers=admin_user.headers,\n json=req3,\n )\n assert (\n res3.status_code == HTTPStatus.OK\n ), f\"Error response: {res3.status_code} - {res3.text}\"\n\n # Check connector kg_processing is enabled\n with get_session_with_current_tenant() as db_session:\n connector = (\n db_session.query(Connector)\n .filter(Connector.name == \"Salesforce Test\")\n .scalar()\n )\n assert connector.kg_processing_enabled\n\n # Check entity types looks correct\n res4 = requests.get(\n f\"{API_SERVER_URL}/admin/kg/entity-types\",\n headers=admin_user.headers,\n )\n assert (\n res4.status_code == HTTPStatus.OK\n ), f\"Error response: {res4.status_code} - {res4.text}\"\n res4_parsed = SourceAndEntityTypeView.model_validate(res4.json())\n\n def to_entity_type_map(map: dict[str, list[EntityType]]) -> dict[str, EntityType]:\n return {\n entity_type.name: entity_type\n for entity_types in map.values()\n for entity_type in entity_types\n }\n\n expected_entity_types = to_entity_type_map(map=res2_parsed.entity_types)\n new_entity_types = to_entity_type_map(map=res4_parsed.entity_types)\n\n # These are the updates.\n # We're just manually updating them.\n expected_entity_types[\"ACCOUNT\"].active = True\n expected_entity_types[\"ACCOUNT\"].description = \"Test.\"\n expected_entity_types[\"OPPORTUNITY\"].active = False\n expected_entity_types[\"OPPORTUNITY\"].description = \"Test 2.\"\n\n assert new_entity_types == expected_entity_types\n\n\ndef test_update_invalid_kg_entity_type_should_do_nothing(\n connectors: None, # noqa: ARG001\n) -> None:\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Enable kg and populate default entity types\n req1 = json.loads(\n EnableKGConfigRequest(\n vendor=\"Test\",\n vendor_domains=[\"test.app\", \"tester.ai\"],\n ignore_domains=[],\n coverage_start=datetime(1970, 1, 1, 0, 0),\n ).model_dump_json()\n )\n res1 = requests.put(\n f\"{API_SERVER_URL}/admin/kg/config\",\n headers=admin_user.headers,\n json=req1,\n )\n assert (\n res1.status_code == HTTPStatus.OK\n ), f\"Error response: {res1.status_code} - {res1.text}\"\n\n # Get old entity types\n res2 = requests.get(\n f\"{API_SERVER_URL}/admin/kg/entity-types\",\n headers=admin_user.headers,\n )\n assert (\n res2.status_code == HTTPStatus.OK\n ), f\"Error response: {res2.status_code} - {res2.text}\"\n\n # Update entity types with non-existent entity type\n req3 = [\n EntityType(name=\"NON-EXISTENT\", description=\"Test.\", active=False).model_dump(),\n ]\n res3 = requests.put(\n f\"{API_SERVER_URL}/admin/kg/entity-types\",\n headers=admin_user.headers,\n json=req3,\n )\n assert (\n res3.status_code == HTTPStatus.OK\n ), f\"Error response: {res3.status_code} - {res3.text}\"\n\n # Get entity types after the update attempt\n res4 = requests.get(\n f\"{API_SERVER_URL}/admin/kg/entity-types\",\n headers=admin_user.headers,\n )\n assert (\n res4.status_code == HTTPStatus.OK\n ), f\"Error response: {res4.status_code} - {res4.text}\"\n\n # Should be the same as before since non-existent entity type should be ignored\n assert res2.json() == res4.json()\n" }, { "path": "backend/tests/integration/tests/llm_auto_update/test_auto_llm_update.py", "content": "\"\"\"\nIntegration tests for Auto LLM model update feature.\n\nThese tests verify that LLM providers in Auto mode get their models\nautomatically synced from the GitHub config via the celery background task.\n\nEnvironment variables for testing:\n- AUTO_LLM_UPDATE_INTERVAL_SECONDS: Set to a low value (e.g., 5) for faster tests\n- AUTO_LLM_CONFIG_URL: Points to the config file to sync from\n\nThe celery beat scheduler will run the check_for_auto_llm_updates task\nat the configured interval.\n\"\"\"\n\nimport time\n\nimport pytest\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n# How long to wait for the celery task to run and sync models\n# This should be longer than AUTO_LLM_UPDATE_INTERVAL_SECONDS\nMAX_WAIT_TIME_SECONDS = 120\nPOLL_INTERVAL_SECONDS = 5\n\n\ndef _create_provider_with_api(\n admin_user: DATestUser,\n name: str,\n provider_type: str,\n default_model: str,\n is_auto_mode: bool,\n model_configurations: list[dict] | None = None,\n) -> dict:\n \"\"\"Create an LLM provider via the API.\"\"\"\n if model_configurations is None:\n model_configurations = [{\"name\": default_model, \"is_visible\": True}]\n\n llm_provider_data = {\n \"name\": name,\n \"provider\": provider_type,\n \"api_key\": \"test-api-key-for-auto-mode-testing\",\n \"api_base\": None,\n \"api_version\": None,\n \"custom_config\": None,\n \"is_public\": True,\n \"is_auto_mode\": is_auto_mode,\n \"groups\": [],\n \"personas\": [],\n \"model_configurations\": model_configurations,\n \"api_key_changed\": True,\n }\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n json=llm_provider_data,\n headers=admin_user.headers,\n )\n response.raise_for_status()\n return response.json()\n\n\ndef _get_provider_by_id(admin_user: DATestUser, provider_id: int) -> dict:\n \"\"\"Get a provider by ID via the API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n response.raise_for_status()\n for provider in response.json()[\"providers\"]:\n if provider[\"id\"] == provider_id:\n return provider\n raise ValueError(f\"Provider with id {provider_id} not found\")\n\n\ndef get_auto_config(admin_user: DATestUser) -> dict | None:\n \"\"\"Get the current auto config from the API.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/auto-config\",\n headers=admin_user.headers,\n )\n if response.status_code == 502:\n return None\n response.raise_for_status()\n return response.json()\n\n\ndef wait_for_model_sync(\n admin_user: DATestUser,\n provider_id: int,\n expected_model_names: set[str],\n max_wait_seconds: int = MAX_WAIT_TIME_SECONDS,\n) -> dict:\n \"\"\"\n Wait for the provider's models to match the expected set.\n\n Returns the provider data once models match, or raises an assertion error.\n \"\"\"\n start_time = time.time()\n last_provider: dict | None = None\n\n while time.time() - start_time < max_wait_seconds:\n provider = _get_provider_by_id(admin_user, provider_id)\n last_provider = provider\n current_models = {m[\"name\"] for m in provider[\"model_configurations\"]}\n\n # Check if we have all expected models\n if expected_model_names.issubset(current_models):\n return provider\n\n print(\n f\"Waiting for model sync... Current: {current_models}, Expected: {expected_model_names}\"\n )\n time.sleep(POLL_INTERVAL_SECONDS)\n\n # Timeout - return last state for debugging\n current_models = (\n {m[\"name\"] for m in last_provider[\"model_configurations\"]}\n if last_provider\n else set()\n )\n raise AssertionError(\n f\"Model sync timed out after {max_wait_seconds}s. Current models: {current_models}, Expected: {expected_model_names}\"\n )\n\n\ndef test_auto_mode_provider_gets_synced_from_github_config(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"\n Test that a provider in Auto mode gets its models synced from GitHub config.\n\n This test:\n 1. Fetches the current GitHub config to know what models to expect\n 2. Creates an OpenAI provider in Auto mode with outdated/minimal models\n 3. Waits for the celery task to sync models from GitHub\n 4. Verifies the models match the GitHub config\n \"\"\"\n # First, get the GitHub config to know what models we should expect\n github_config = get_auto_config(admin_user)\n if github_config is None:\n pytest.fail(\"GitHub config not found\")\n\n # Get expected models for OpenAI from the config\n if \"openai\" not in github_config.get(\"providers\", {}):\n pytest.fail(\"OpenAI not in GitHub config\")\n\n openai_config = github_config[\"providers\"][\"openai\"]\n\n # Build expected model names from default_model + additional_visible_models\n expected_models: set[str] = set()\n\n # Add default model\n default_model = openai_config.get(\"default_model\", {})\n if isinstance(default_model, dict):\n expected_models.add(default_model[\"name\"])\n elif isinstance(default_model, str):\n expected_models.add(default_model)\n\n # Add additional visible models\n for model in openai_config.get(\"additional_visible_models\", []):\n if isinstance(model, dict):\n expected_models.add(model[\"name\"])\n elif isinstance(model, str):\n expected_models.add(model)\n\n print(f\"Expected models from GitHub config: {expected_models}\")\n\n # Create an OpenAI provider in Auto mode with a single outdated model\n provider = _create_provider_with_api(\n admin_user=admin_user,\n name=\"test-auto-sync-openai\",\n provider_type=\"openai\",\n default_model=\"outdated-model-name\",\n is_auto_mode=True,\n model_configurations=[\n {\"name\": \"outdated-model-name\", \"is_visible\": True},\n ],\n )\n\n assert provider[\"is_auto_mode\"] is True\n print(f\"Created provider {provider['id']} in Auto mode\")\n\n # Wait for the celery task to sync models\n # The task runs at AUTO_LLM_UPDATE_INTERVAL_SECONDS interval\n synced_provider = wait_for_model_sync(\n admin_user=admin_user,\n provider_id=provider[\"id\"],\n expected_model_names=expected_models,\n )\n\n # Verify the models were synced\n synced_model_configs = synced_provider[\"model_configurations\"]\n synced_model_names = {m[\"name\"] for m in synced_model_configs}\n print(f\"Synced models: {synced_model_names}\")\n\n assert expected_models.issubset(\n synced_model_names\n ), f\"Expected models {expected_models} not found in synced models {synced_model_names}\"\n\n # Verify the outdated model still exists but is not visible\n # (Auto mode marks removed models as not visible, it doesn't delete them)\n outdated_model = next(\n (m for m in synced_model_configs if m[\"name\"] == \"outdated-model-name\"),\n None,\n )\n assert (\n outdated_model is not None\n ), \"Outdated model should still exist after sync (marked invisible, not deleted)\"\n assert not outdated_model[\n \"is_visible\"\n ], \"Outdated model should not be visible after sync\"\n\n\ndef test_manual_mode_provider_not_affected_by_auto_sync(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"\n Test that a provider in Manual mode is NOT affected by auto sync.\n\n This test:\n 1. Creates an OpenAI provider in Manual mode with custom models\n 2. Waits for a period longer than the sync interval\n 3. Verifies the models remain unchanged\n \"\"\"\n custom_model = \"my-custom-finetuned-model\"\n\n # Create a provider in Manual mode\n provider = _create_provider_with_api(\n admin_user=admin_user,\n name=\"test-manual-mode-unchanged\",\n provider_type=\"openai\",\n default_model=custom_model,\n is_auto_mode=False, # Manual mode\n model_configurations=[\n {\"name\": custom_model, \"is_visible\": True},\n {\"name\": \"another-custom-model\", \"is_visible\": True},\n ],\n )\n\n assert provider[\"is_auto_mode\"] is False\n initial_models = {m[\"name\"] for m in provider[\"model_configurations\"]}\n print(f\"Created manual mode provider with models: {initial_models}\")\n\n # Wait for longer than the sync interval\n wait_time = 15 # Should be longer than AUTO_LLM_UPDATE_INTERVAL_SECONDS\n print(f\"Waiting {wait_time}s to ensure sync task runs...\")\n time.sleep(wait_time)\n\n # Verify models are unchanged\n updated_provider = _get_provider_by_id(admin_user, provider[\"id\"])\n current_models = {m[\"name\"] for m in updated_provider[\"model_configurations\"]}\n\n assert (\n current_models == initial_models\n ), f\"Manual mode provider models should not change. Initial: {initial_models}, Current: {current_models}\"\n" }, { "path": "backend/tests/integration/tests/llm_provider/test_llm_provider.py", "content": "import uuid\nfrom typing import Any\n\nimport pytest\nimport requests\nfrom requests.models import Response\n\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.model_name_parser import parse_litellm_model_name\nfrom onyx.llm.utils import get_max_input_tokens\nfrom onyx.llm.utils import litellm_thinks_model_supports_image_input\nfrom onyx.llm.utils import model_is_reasoning_model\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _get_provider_by_id(admin_user: DATestUser, provider_id: str) -> dict | None:\n \"\"\"Utility function to fetch an LLM provider by ID\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n return next((p for p in providers if p[\"id\"] == provider_id), None)\n\n\ndef assert_response_is_equivalent(\n admin_user: DATestUser,\n response: Response,\n model_configurations: list[ModelConfigurationUpsertRequest],\n api_key: str | None = None,\n) -> None:\n assert response.status_code == 200\n created_provider = response.json()\n\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n\n assert provider_data[\"personas\"] == []\n\n def fill_max_input_tokens_and_supports_image_input(\n req: ModelConfigurationUpsertRequest,\n ) -> dict[str, Any]:\n provider_name = created_provider[\"provider\"]\n # Match how ModelConfigurationView.from_model builds the key for parsing\n model_key = req.name\n if provider_name and not model_key.startswith(f\"{provider_name}/\"):\n model_key = f\"{provider_name}/{model_key}\"\n parsed = parse_litellm_model_name(model_key)\n\n # Include region in display name for Bedrock cross-region models (matches from_model)\n display_name = (\n f\"{parsed.display_name} ({parsed.region})\"\n if parsed.region\n else parsed.display_name\n )\n\n filled_with_max_input_tokens = ModelConfigurationUpsertRequest(\n name=req.name,\n is_visible=req.is_visible,\n max_input_tokens=req.max_input_tokens\n or get_max_input_tokens(model_name=req.name, model_provider=provider_name),\n )\n return {\n **filled_with_max_input_tokens.model_dump(),\n \"supports_image_input\": litellm_thinks_model_supports_image_input(\n req.name, provider_name\n ),\n \"supports_reasoning\": model_is_reasoning_model(req.name, provider_name),\n \"display_name\": display_name,\n \"provider_display_name\": parsed.provider_display_name,\n \"vendor\": parsed.vendor,\n \"region\": parsed.region,\n \"version\": parsed.version,\n }\n\n # Compare model configurations by name (order-independent)\n actual_by_name = {\n config[\"name\"]: config for config in provider_data[\"model_configurations\"]\n }\n expected_by_name = {\n config.name: fill_max_input_tokens_and_supports_image_input(config)\n for config in model_configurations\n }\n\n assert set(actual_by_name.keys()) == set(\n expected_by_name.keys()\n ), f\"Model names don't match. Actual: {set(actual_by_name.keys())}, Expected: {set(expected_by_name.keys())}\"\n\n for name in actual_by_name:\n actual_config = actual_by_name[name]\n expected_config = expected_by_name[name]\n assert (\n actual_config == expected_config\n ), f\"Config mismatch for {name}:\\nActual: {actual_config}\\nExpected: {expected_config}\"\n\n # test that returned key is sanitized\n if api_key:\n assert provider_data[\"api_key\"] == api_key\n\n\n# Test creating an LLM Provider with some various model-configurations.\n@pytest.mark.parametrize(\n \"model_configurations, expected\",\n [\n # Test the case in which a basic model-configuration is passed.\n (\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n ),\n # Test the case in which multiple model-configuration are passed.\n (\n [\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ],\n [\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ],\n ),\n # Test the case in which duplicate model-configuration are passed.\n (\n [ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True)] * 4,\n [ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True)],\n ),\n ],\n)\ndef test_create_llm_provider(\n reset: None, # noqa: ARG001\n model_configurations: list[ModelConfigurationUpsertRequest],\n expected: list[ModelConfigurationUpsertRequest],\n) -> None:\n admin_user = UserManager.create(name=\"admin_user\")\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": str(uuid.uuid4()),\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n model_configuration.model_dump()\n for model_configuration in model_configurations\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n\n assert_response_is_equivalent(\n admin_user,\n response,\n expected,\n \"sk-0****0000\",\n )\n\n\n# Test creating a new LLM Provider with some given model-configurations, then performing some arbitrary update on it.\n@pytest.mark.parametrize(\n \"initial, initial_expected, updated, updated_expected\",\n [\n # Test the case in which a basic model-configuration is passed, but then it's updated to have *NO* max-input-tokens.\n (\n (\n \"gpt-4\",\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n ),\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n (\n \"gpt-4\",\n [ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True)],\n ),\n [ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True)],\n ),\n # Test the case where we insert 2 model-configurations, and then in the update the first,\n # we update one and delete the second.\n (\n (\n \"gpt-4\",\n [\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True, max_input_tokens=4096\n ),\n ],\n ),\n [\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True, max_input_tokens=4096\n ),\n ],\n (\n \"gpt-4\",\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n ),\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n ),\n ],\n)\ndef test_update_model_configurations(\n reset: None, # noqa: ARG001\n initial: tuple[str, list[ModelConfigurationUpsertRequest]],\n initial_expected: list[ModelConfigurationUpsertRequest],\n updated: tuple[str, list[ModelConfigurationUpsertRequest]],\n updated_expected: list[ModelConfigurationUpsertRequest],\n) -> None:\n admin_user = UserManager.create(name=\"admin_user\")\n\n default_model_name, model_configurations = initial\n updated_default_model_name, updated_model_configurations = updated\n\n name = str(uuid.uuid4())\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n model_configuration.dict()\n for model_configuration in model_configurations\n ],\n \"is_public\": True,\n \"groups\": [],\n \"api_key_changed\": True,\n },\n )\n created_provider = response.json()\n assert_response_is_equivalent(\n admin_user,\n response,\n initial_expected,\n )\n\n response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n headers=admin_user.headers,\n json={\n \"provider_id\": created_provider[\"id\"],\n \"model_name\": updated_default_model_name,\n },\n )\n assert response.status_code == 200\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n json={\n \"id\": created_provider[\"id\"],\n \"name\": name,\n \"provider\": created_provider[\"provider\"],\n \"api_key\": \"sk-000000000000000000000000000000000000000000000001\",\n \"model_configurations\": [\n model_configuration.dict()\n for model_configuration in updated_model_configurations\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert_response_is_equivalent(\n admin_user,\n response,\n updated_expected,\n \"sk-0****0000\",\n )\n\n response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n headers=admin_user.headers,\n json={\n \"provider_id\": created_provider[\"id\"],\n \"model_name\": updated_default_model_name,\n },\n )\n assert response.status_code == 200\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n json={\n \"id\": created_provider[\"id\"],\n \"name\": name,\n \"provider\": created_provider[\"provider\"],\n \"api_key\": \"sk-000000000000000000000000000000000000000000000001\",\n \"model_configurations\": [\n model_configuration.dict()\n for model_configuration in updated_model_configurations\n ],\n \"is_public\": True,\n \"groups\": [],\n \"api_key_changed\": True,\n },\n )\n assert_response_is_equivalent(\n admin_user,\n response,\n updated_expected,\n \"sk-0****0001\",\n )\n\n\n@pytest.mark.parametrize(\n \"model_configurations\",\n [\n [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\", is_visible=True, max_input_tokens=4096\n )\n ],\n [\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ],\n ],\n)\ndef test_delete_llm_provider(\n reset: None, # noqa: ARG001\n model_configurations: list[ModelConfigurationUpsertRequest],\n) -> None:\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a provider\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"test-provider-delete\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n model_configuration.dict()\n for model_configuration in model_configurations\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n created_provider = response.json()\n assert response.status_code == 200\n\n # Delete the provider\n response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{created_provider['id']}\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n # Verify provider is deleted by checking it's not in the list\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is None\n\n\ndef test_delete_default_llm_provider_rejected(\n reset: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Deleting the default LLM provider should return 400.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a provider\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"test-provider-default-delete\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert response.status_code == 200\n created_provider = response.json()\n\n # Set this provider as the default\n set_default_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n headers=admin_user.headers,\n json={\n \"provider_id\": created_provider[\"id\"],\n \"model_name\": \"gpt-4o-mini\",\n },\n )\n assert set_default_response.status_code == 200\n\n # Attempt to delete the default provider — should be rejected\n delete_response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{created_provider['id']}\",\n headers=admin_user.headers,\n )\n assert delete_response.status_code == 400\n assert \"Cannot delete the default LLM provider\" in delete_response.json()[\"detail\"]\n\n # Verify provider still exists\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n\n\ndef test_delete_non_default_llm_provider_with_default_set(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Deleting a non-default provider should succeed even when a default is set.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create two providers\n response_default = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"default-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert response_default.status_code == 200\n default_provider = response_default.json()\n\n response_other = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"other-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert response_other.status_code == 200\n other_provider = response_other.json()\n\n # Set the first provider as default\n set_default_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n headers=admin_user.headers,\n json={\n \"provider_id\": default_provider[\"id\"],\n \"model_name\": \"gpt-4o-mini\",\n },\n )\n assert set_default_response.status_code == 200\n\n # Delete the non-default provider — should succeed\n delete_response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{other_provider['id']}\",\n headers=admin_user.headers,\n )\n assert delete_response.status_code == 200\n\n # Verify the non-default provider is gone\n provider_data = _get_provider_by_id(admin_user, other_provider[\"id\"])\n assert provider_data is None\n\n # Verify the default provider still exists\n default_data = _get_provider_by_id(admin_user, default_provider[\"id\"])\n assert default_data is not None\n\n\ndef test_force_delete_default_llm_provider(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Force-deleting the default LLM provider should succeed.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a provider\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"test-provider-force-delete\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert response.status_code == 200\n created_provider = response.json()\n\n # Set this provider as the default\n set_default_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n headers=admin_user.headers,\n json={\n \"provider_id\": created_provider[\"id\"],\n \"model_name\": \"gpt-4o-mini\",\n },\n )\n assert set_default_response.status_code == 200\n\n # Attempt to delete without force — should be rejected\n delete_response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{created_provider['id']}\",\n headers=admin_user.headers,\n )\n assert delete_response.status_code == 400\n\n # Force delete — should succeed\n force_delete_response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{created_provider['id']}?force=true\",\n headers=admin_user.headers,\n )\n assert force_delete_response.status_code == 200\n\n # Verify provider is gone\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is None\n\n\ndef test_delete_default_vision_provider_clears_vision_default(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Deleting the default vision provider should succeed and clear the vision default.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a text provider and set it as default (so we have a default text provider)\n text_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"text-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000001\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert text_response.status_code == 200\n text_provider = text_response.json()\n _set_default_provider(admin_user, text_provider[\"id\"], \"gpt-4o-mini\")\n\n # Create a vision provider and set it as default vision\n vision_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"vision-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000002\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n supports_image_input=True,\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n },\n )\n assert vision_response.status_code == 200\n vision_provider = vision_response.json()\n _set_default_vision_provider(admin_user, vision_provider[\"id\"], \"gpt-4o\")\n\n # Verify vision default is set\n data = _get_providers_admin(admin_user)\n assert data is not None\n _, _, vision_default = _unpack_data(data)\n assert vision_default is not None\n assert vision_default[\"provider_id\"] == vision_provider[\"id\"]\n\n # Delete the vision provider — should succeed (only text default is protected)\n delete_response = requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{vision_provider['id']}\",\n headers=admin_user.headers,\n )\n assert delete_response.status_code == 200\n\n # Verify the vision provider is gone\n provider_data = _get_provider_by_id(admin_user, vision_provider[\"id\"])\n assert provider_data is None\n\n # Verify there is no default vision provider\n data = _get_providers_admin(admin_user)\n assert data is not None\n _, text_default, vision_default = _unpack_data(data)\n assert vision_default is None\n\n # Verify the text default is still intact\n assert text_default is not None\n assert text_default[\"provider_id\"] == text_provider[\"id\"]\n\n\ndef test_duplicate_provider_name_rejected(reset: None) -> None: # noqa: ARG001\n \"\"\"Creating a provider with a name that already exists should return 400.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n provider_name = f\"unique-provider-{uuid.uuid4()}\"\n\n base_payload = {\n \"name\": provider_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n }\n\n # First creation succeeds\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json=base_payload,\n )\n assert response.status_code == 200\n\n # Second creation with the same name is rejected\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json=base_payload,\n )\n assert response.status_code == 409\n assert \"already exists\" in response.json()[\"detail\"]\n\n\ndef test_rename_provider_rejected(reset: None) -> None: # noqa: ARG001\n \"\"\"Renaming a provider is not currently supported and should return 400.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n create_payload = {\n \"name\": f\"original-name-{uuid.uuid4()}\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\", is_visible=True\n ).model_dump()\n ],\n \"is_public\": True,\n \"groups\": [],\n }\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json=create_payload,\n )\n assert response.status_code == 200\n provider_id = response.json()[\"id\"]\n\n # Attempt to rename — should be rejected\n new_name = f\"renamed-provider-{uuid.uuid4()}\"\n update_payload = {**create_payload, \"id\": provider_id, \"name\": new_name}\n response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json=update_payload,\n )\n assert response.status_code == 400\n assert \"not currently supported\" in response.json()[\"detail\"]\n\n # Verify no duplicate was created — only the original provider should exist\n provider = _get_provider_by_id(admin_user, provider_id)\n assert provider is not None\n assert provider[\"name\"] == create_payload[\"name\"]\n\n all_response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n assert all_response.status_code == 200\n all_names = [p[\"name\"] for p in all_response.json()[\"providers\"]]\n assert new_name not in all_names\n\n\ndef test_model_visibility_preserved_on_edit(reset: None) -> None: # noqa: ARG001\n \"\"\"\n Test that model visibility flags are correctly preserved when editing an LLM provider.\n\n This test verifies the fix for the bug where editing a provider with specific visible models\n would incorrectly map visibility flags when the provider's model list differs from the\n descriptor's default model list.\n\n Scenario:\n 1. Create a provider with 3 models, 2 visible\n 2. Edit the provider to change visibility (make all 3 visible)\n 3. Verify all 3 models are now visible\n 4. Edit again to make only 1 visible\n 5. Verify only 1 is visible\n \"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Initial model configurations: 2 visible, 1 hidden\n model_configs = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\",\n is_visible=True,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4-turbo\",\n is_visible=False,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ]\n\n # Create the provider\n create_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": \"test-visibility-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [config.dict() for config in model_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response.status_code == 200\n created_provider = create_response.json()\n\n # Verify initial state: 2 visible models\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n visible_models = [\n model for model in provider_data[\"model_configurations\"] if model[\"is_visible\"]\n ]\n assert len(visible_models) == 2\n assert any(m[\"name\"] == \"gpt-4o\" for m in visible_models)\n assert any(m[\"name\"] == \"gpt-4o-mini\" for m in visible_models)\n\n # Edit 1: Make all 3 models visible\n edit_configs_all_visible = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\",\n is_visible=True,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4-turbo\",\n is_visible=True, # Now visible\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ]\n\n edit_response_1 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json={\n \"id\": created_provider[\"id\"],\n \"name\": \"test-visibility-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n config.dict() for config in edit_configs_all_visible\n ],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert edit_response_1.status_code == 200\n\n # Verify all 3 models are now visible\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n visible_models = [\n model for model in provider_data[\"model_configurations\"] if model[\"is_visible\"]\n ]\n assert len(visible_models) == 3\n\n # Edit 2: Make only 1 model visible\n edit_configs_one_visible = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True, # Only this one visible\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\",\n is_visible=False,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4-turbo\",\n is_visible=False,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ]\n\n edit_response_2 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json={\n \"id\": created_provider[\"id\"],\n \"name\": \"test-visibility-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n config.dict() for config in edit_configs_one_visible\n ],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert edit_response_2.status_code == 200\n\n # Verify only 1 model is visible\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n visible_models = [\n model for model in provider_data[\"model_configurations\"] if model[\"is_visible\"]\n ]\n assert len(visible_models) == 1\n assert visible_models[0][\"name\"] == \"gpt-4o\"\n\n # Make none visible\n edit_configs_none_visible = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=False,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o-mini\",\n is_visible=False,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4-turbo\",\n is_visible=False,\n max_input_tokens=None,\n supports_image_input=None,\n ),\n ]\n edit_response_3 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json={\n \"id\": created_provider[\"id\"],\n \"name\": \"test-visibility-provider\",\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [\n config.dict() for config in edit_configs_none_visible\n ],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert edit_response_3.status_code == 200\n\n # Verify no models are visible\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n visible_models = [\n model for model in provider_data[\"model_configurations\"] if model[\"is_visible\"]\n ]\n assert len(visible_models) == 0\n\n # Make gpt-4o the default\n _set_default_provider(admin_user, created_provider[\"id\"], \"gpt-4o\")\n\n # Verify gpt-4o is the default\n provider_data = _get_provider_by_id(admin_user, created_provider[\"id\"])\n assert provider_data is not None\n visible_models = [\n model for model in provider_data[\"model_configurations\"] if model[\"is_visible\"]\n ]\n assert len(visible_models) == 1\n assert visible_models[0][\"name\"] == \"gpt-4o\"\n\n\ndef _get_provider_by_name(providers: list[dict], provider_name: str) -> dict | None:\n return next((p for p in providers if p[\"name\"] == provider_name), None)\n\n\ndef _get_providers_admin(\n admin_user: DATestUser,\n) -> dict | None:\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n resp_json = response.json()\n\n return resp_json\n\n\ndef _unpack_data(data: dict) -> tuple[list[dict], dict | None, dict | None]:\n providers = data[\"providers\"]\n text_default = data.get(\"default_text\")\n vision_default = data.get(\"default_vision\")\n\n return providers, text_default, vision_default\n\n\ndef _get_providers_basic(\n user: DATestUser,\n) -> dict | None:\n response = requests.get(\n f\"{API_SERVER_URL}/llm/provider\",\n headers=user.headers,\n )\n assert response.status_code == 200\n resp_json = response.json()\n\n return resp_json\n\n\ndef _validate_default_model(\n default: dict | None,\n provider_id: int | None = None,\n model_name: str | None = None,\n) -> None:\n if default is None:\n assert provider_id is None and model_name is None\n return\n\n assert default[\"provider_id\"] == provider_id\n assert default[\"model_name\"] == model_name\n\n\ndef _get_provider_by_name_admin(\n admin_user: DATestUser, provider_name: str\n) -> dict | None:\n \"\"\"Utility function to fetch an LLM provider by name via admin endpoint.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n providers = response.json()\n return next((p for p in providers if p[\"name\"] == provider_name), None)\n\n\ndef _get_provider_by_name_basic(user: DATestUser, provider_name: str) -> dict | None:\n \"\"\"Utility function to fetch an LLM provider by name via basic (non-admin) endpoint.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/llm/provider\",\n headers=user.headers,\n )\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n return next((p for p in providers if p[\"name\"] == provider_name), None)\n\n\ndef _validate_model_configurations(\n actual_configs: list[dict],\n expected_model_names: list[str],\n expected_visible: dict[str, bool] | None = None,\n expected_image_support: dict[str, bool] | None = None,\n) -> None:\n \"\"\"\n Validate that model configurations match expectations.\n\n Args:\n actual_configs: List of model configuration dicts from the API response\n expected_model_names: List of expected model names\n expected_visible: Optional dict mapping model name to expected visibility\n expected_image_support: Optional dict mapping model name to expected supports_image_input\n \"\"\"\n actual_names = {config[\"name\"] for config in actual_configs}\n expected_names = set(expected_model_names)\n\n assert (\n actual_names == expected_names\n ), f\"Model names mismatch. Expected: {expected_names}, Actual: {actual_names}\"\n\n if expected_visible:\n for config in actual_configs:\n if config[\"name\"] in expected_visible:\n assert config[\"is_visible\"] == expected_visible[config[\"name\"]], (\n f\"Visibility mismatch for {config['name']}. \"\n f\"Expected: {expected_visible[config['name']]}, Actual: {config['is_visible']}\"\n )\n\n if expected_image_support:\n for config in actual_configs:\n if config[\"name\"] in expected_image_support:\n assert (\n config[\"supports_image_input\"]\n == expected_image_support[config[\"name\"]]\n ), (\n f\"supports_image_input mismatch for {config['name']}. \"\n f\"Expected: {expected_image_support[config['name']]}, \"\n f\"Actual: {config['supports_image_input']}\"\n )\n\n\ndef _validate_provider_data(\n provider_data: dict,\n expected_name: str,\n expected_provider: str,\n expected_model_names: list[str],\n expected_visible: dict[str, bool] | None = None,\n expected_is_public: bool | None = None,\n expected_image_support: dict[str, bool] | None = None,\n) -> None:\n \"\"\"\n Validate that provider data matches expectations.\n\n Args:\n provider_data: Provider dict from the API response\n expected_name: Expected provider name\n expected_provider: Expected provider type (e.g., 'openai')\n expected_model_names: List of expected model names in configurations\n expected_visible: Optional dict mapping model name to expected visibility\n expected_is_public: Optional expected is_public value (admin endpoint only)\n expected_image_support: Optional dict mapping model name to expected supports_image_input\n \"\"\"\n assert (\n provider_data[\"name\"] == expected_name\n ), f\"Provider name mismatch. Expected: {expected_name}, Actual: {provider_data['name']}\"\n assert (\n provider_data[\"provider\"] == expected_provider\n ), f\"Provider type mismatch. Expected: {expected_provider}, Actual: {provider_data['provider']}\"\n\n # Validate is_public if provided (only available in admin endpoint response)\n if expected_is_public is not None and \"is_public\" in provider_data:\n assert (\n provider_data[\"is_public\"] == expected_is_public\n ), f\"is_public mismatch. Expected: {expected_is_public}, Actual: {provider_data['is_public']}\"\n\n # Validate model configurations\n _validate_model_configurations(\n provider_data[\"model_configurations\"],\n expected_model_names,\n expected_visible,\n expected_image_support,\n )\n\n\ndef test_default_model_persistence_and_update(\n reset: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"\n Test that the default model is correctly set, persisted, and can be updated.\n\n This test verifies:\n 1. Admin creates a provider with a specific default model\n 2. Admin endpoint (/admin/llm/provider) shows correct default model\n 3. Basic endpoint (/llm/provider) shows correct default model for admin user\n 4. Non-admin user can see the same default model via basic endpoint\n 5. Admin updates the default model\n 6. Both admin and basic endpoints reflect the new default model\n 7. Non-admin user sees the updated default model\n \"\"\"\n from onyx.auth.schemas import UserRole\n\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a non-admin user\n basic_user = UserManager.create(name=\"basic_user\")\n # The first user is admin, subsequent users are basic by default\n assert basic_user.role == UserRole.BASIC or basic_user.role != UserRole.ADMIN\n\n provider_name = f\"test-default-model-{uuid.uuid4()}\"\n updated_default_model = \"gpt-4o\"\n\n # Model configurations including all models we'll use\n model_configs = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\",\n is_visible=True,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n ),\n ]\n\n # Expected model names and visibility\n expected_model_names = [\"gpt-4\", \"gpt-4o\"]\n expected_visible = {\"gpt-4\": True, \"gpt-4o\": True}\n\n # Step 1: Admin creates the provider with initial default model\n create_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": provider_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [config.model_dump() for config in model_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response.status_code == 200\n\n # Capture initial defaults (setup_postgres may have created a DevEnvPresetOpenAI default)\n initial_data = _get_providers_admin(admin_user)\n assert initial_data is not None\n _, initial_text_default, initial_vision_default = _unpack_data(initial_data)\n\n # Step 2: Verify via admin endpoint that all provider data is correct\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n # Defaults should be unchanged from initial state (new provider not set as default)\n assert text_default == initial_text_default\n assert vision_default == initial_vision_default\n\n admin_provider_data = _get_provider_by_name(providers, provider_name)\n assert admin_provider_data is not None\n\n _validate_provider_data(\n admin_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n expected_is_public=True,\n )\n\n # Step 3: Verify via basic endpoint (admin user) that all provider data is correct\n admin_basic_data = _get_providers_basic(admin_user)\n assert admin_basic_data is not None\n providers, text_default, vision_default = _unpack_data(admin_basic_data)\n assert text_default == initial_text_default\n assert vision_default == initial_vision_default\n\n admin_basic_provider_data = _get_provider_by_name(providers, provider_name)\n assert admin_basic_provider_data is not None\n _validate_provider_data(\n admin_basic_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n )\n\n # Step 4: Verify non-admin user sees the same provider data via basic endpoint\n basic_user_data = _get_providers_basic(basic_user)\n assert basic_user_data is not None\n providers, text_default, vision_default = _unpack_data(basic_user_data)\n assert text_default == initial_text_default\n assert vision_default == initial_vision_default\n\n basic_user_provider_data = _get_provider_by_name(providers, provider_name)\n assert basic_user_provider_data is not None\n _validate_provider_data(\n basic_user_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n )\n\n # Step 5: Admin updates the provider to change the default model\n update_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json={\n \"id\": create_response.json()[\"id\"],\n \"name\": provider_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [config.model_dump() for config in model_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert update_response.status_code == 200\n\n default_provider_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n json={\n \"provider_id\": update_response.json()[\"id\"],\n \"model_name\": updated_default_model,\n },\n headers=admin_user.headers,\n )\n assert default_provider_response.status_code == 200\n\n # Step 6a: Verify the updated provider data via admin endpoint\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default,\n provider_id=update_response.json()[\"id\"],\n model_name=updated_default_model,\n )\n _validate_default_model(vision_default) # None\n\n admin_provider_data = _get_provider_by_name(providers, provider_name)\n assert admin_provider_data is not None\n _validate_provider_data(\n admin_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n expected_is_public=True,\n )\n\n # Step 6b: Verify the updated provider data via basic endpoint (admin user)\n admin_basic_data = _get_providers_basic(admin_user)\n assert admin_basic_data is not None\n providers, text_default, vision_default = _unpack_data(admin_basic_data)\n _validate_default_model(\n text_default,\n provider_id=update_response.json()[\"id\"],\n model_name=updated_default_model,\n )\n _validate_default_model(vision_default) # None\n\n admin_basic_provider_data = _get_provider_by_name(providers, provider_name)\n assert admin_basic_provider_data is not None\n _validate_provider_data(\n admin_basic_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n )\n\n # Step 7: Verify non-admin user sees the updated provider data\n basic_user_data = _get_providers_basic(basic_user)\n assert basic_user_data is not None\n providers, text_default, vision_default = _unpack_data(basic_user_data)\n _validate_default_model(\n text_default,\n provider_id=update_response.json()[\"id\"],\n model_name=updated_default_model,\n )\n _validate_default_model(vision_default) # None\n\n basic_user_provider_data = _get_provider_by_name(providers, provider_name)\n assert basic_user_provider_data is not None\n _validate_provider_data(\n basic_user_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n )\n\n\ndef _get_all_providers_basic(user: DATestUser) -> list[dict]:\n \"\"\"Utility function to fetch all LLM providers via basic endpoint.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/llm/provider\",\n headers=user.headers,\n )\n assert response.status_code == 200\n return response.json()[\"providers\"]\n\n\ndef _get_all_providers_admin(admin_user: DATestUser) -> list[dict]:\n \"\"\"Utility function to fetch all LLM providers via admin endpoint.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n return response.json()[\"providers\"]\n\n\ndef _set_default_provider(\n admin_user: DATestUser, provider_id: int, model_name: str\n) -> None:\n \"\"\"Utility function to set a provider as the default.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n json={\n \"provider_id\": provider_id,\n \"model_name\": model_name,\n },\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n\ndef _set_default_vision_provider(\n admin_user: DATestUser, provider_id: int, vision_model: str | None = None\n) -> None:\n \"\"\"Utility function to set a provider as the default vision provider.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default-vision\",\n json={\n \"provider_id\": provider_id,\n \"model_name\": vision_model,\n },\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n\ndef test_multiple_providers_default_switching(\n reset: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"\n Test switching default providers and models across multiple LLM providers.\n\n This test verifies:\n 1. Admin creates multiple LLM providers\n 2. Admin sets one as the default provider with a specific default model\n 3. Both admin and basic_user query /provider and see the same default provider/model\n 4. Admin changes the default provider and model to something different\n 5. Both admin and basic_user verify they see the same updated default\n 6. Admin switches to a different provider that has a model with the same name\n 7. Both users should see the new provider as default with the same model name\n \"\"\"\n from onyx.auth.schemas import UserRole\n\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a non-admin user\n basic_user = UserManager.create(name=\"basic_user\")\n assert basic_user.role == UserRole.BASIC or basic_user.role != UserRole.ADMIN\n\n # We'll create two providers, both with a model named \"gpt-4\" to test the\n # scenario where different providers have models with the same name\n provider_1_name = f\"test-provider-1-{uuid.uuid4()}\"\n provider_2_name = f\"test-provider-2-{uuid.uuid4()}\"\n\n # Both providers will have \"gpt-4\" as a model\n shared_model_name = \"gpt-4\"\n provider_1_unique_model = \"gpt-4o\"\n provider_2_unique_model = \"gpt-4-turbo\"\n\n # Model configurations for provider 1\n provider_1_configs = [\n ModelConfigurationUpsertRequest(\n name=shared_model_name,\n is_visible=True,\n ),\n ModelConfigurationUpsertRequest(\n name=provider_1_unique_model,\n is_visible=True,\n ),\n ]\n\n # Model configurations for provider 2\n provider_2_configs = [\n ModelConfigurationUpsertRequest(\n name=shared_model_name,\n is_visible=True,\n ),\n ModelConfigurationUpsertRequest(\n name=provider_2_unique_model,\n is_visible=True,\n ),\n ]\n\n # Expected model names and visibility for each provider\n provider_1_model_names = [shared_model_name, provider_1_unique_model]\n provider_1_visible = {shared_model_name: True, provider_1_unique_model: True}\n provider_2_model_names = [shared_model_name, provider_2_unique_model]\n provider_2_visible = {shared_model_name: True, provider_2_unique_model: True}\n\n # Step 1: Create provider 1 with shared_model_name as default\n create_response_1 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": provider_1_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000001\",\n \"model_configurations\": [c.model_dump() for c in provider_1_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response_1.status_code == 200\n provider_1 = create_response_1.json()\n\n _set_default_provider(admin_user, provider_1[\"id\"], shared_model_name)\n\n # Create provider 2 with provider_2_unique_model as default initially\n create_response_2 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": provider_2_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000002\",\n \"model_configurations\": [c.model_dump() for c in provider_2_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response_2.status_code == 200\n provider_2 = create_response_2.json()\n\n # Step 2: Set provider 1 as the default provider\n _set_default_provider(admin_user, provider_1[\"id\"], shared_model_name)\n\n # Step 3: Both admin and basic_user query and verify they see the same default\n # Validate via admin endpoint\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default, provider_id=provider_1[\"id\"], model_name=shared_model_name\n )\n _validate_default_model(vision_default) # None\n admin_provider_data = _get_provider_by_name(providers, provider_1_name)\n assert admin_provider_data is not None\n _validate_provider_data(\n admin_provider_data,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n expected_is_public=True,\n )\n\n # Validate provider 2 via admin endpoint (should not be default)\n admin_provider_2 = _get_provider_by_name(providers, provider_2_name)\n assert admin_provider_2 is not None\n _validate_provider_data(\n admin_provider_2,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n expected_is_public=True,\n )\n\n # Validate via basic endpoint (basic_user)\n basic_data = _get_providers_basic(basic_user)\n assert basic_data is not None\n providers, text_default, vision_default = _unpack_data(basic_data)\n _validate_default_model(\n text_default, provider_id=provider_1[\"id\"], model_name=shared_model_name\n )\n _validate_default_model(vision_default) # None\n basic_provider_data = _get_provider_by_name(providers, provider_1_name)\n assert basic_provider_data is not None\n _validate_provider_data(\n basic_provider_data,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n )\n\n # Also verify admin sees the same via basic endpoint\n admin_basic_data = _get_providers_basic(admin_user)\n assert admin_basic_data is not None\n providers, text_default, vision_default = _unpack_data(admin_basic_data)\n _validate_default_model(\n text_default, provider_id=provider_1[\"id\"], model_name=shared_model_name\n )\n _validate_default_model(vision_default) # None\n admin_basic_provider_data = _get_provider_by_name(providers, provider_1_name)\n assert admin_basic_provider_data is not None\n _validate_provider_data(\n admin_basic_provider_data,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n )\n\n # Step 4: Admin changes the default provider to provider 2 and updates its default model\n # First update provider 2's default model to the unique model (it already is, but reconfirm)\n update_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json={\n \"id\": provider_2[\"id\"],\n \"name\": provider_2_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000002\",\n \"model_configurations\": [c.model_dump() for c in provider_2_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert update_response.status_code == 200\n\n # Now set provider 2 as the default\n _set_default_provider(admin_user, provider_2[\"id\"], provider_2_unique_model)\n\n # Step 5: Both admin and basic_user verify they see the updated default\n # Validate via admin endpoint\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default, provider_id=provider_2[\"id\"], model_name=provider_2_unique_model\n )\n _validate_default_model(vision_default) # None\n admin_provider_data = _get_provider_by_name(providers, provider_2_name)\n assert admin_provider_data is not None\n _validate_provider_data(\n admin_provider_data,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n expected_is_public=True,\n )\n\n # Validate provider 1 via admin endpoint (should no longer be default)\n admin_provider_1 = _get_provider_by_name(providers, provider_1_name)\n assert admin_provider_1 is not None\n _validate_provider_data(\n admin_provider_1,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n expected_is_public=True,\n )\n\n # Validate via basic endpoint (basic_user)\n basic_data = _get_providers_basic(basic_user)\n assert basic_data is not None\n providers, text_default, vision_default = _unpack_data(basic_data)\n _validate_default_model(\n text_default, provider_id=provider_2[\"id\"], model_name=provider_2_unique_model\n )\n _validate_default_model(vision_default) # None\n basic_provider_data = _get_provider_by_name(providers, provider_2_name)\n assert basic_provider_data is not None\n _validate_provider_data(\n basic_provider_data,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n )\n\n # Validate via basic endpoint (admin_user)\n admin_basic_data = _get_providers_basic(admin_user)\n assert admin_basic_data is not None\n providers, text_default, vision_default = _unpack_data(admin_basic_data)\n _validate_default_model(\n text_default, provider_id=provider_2[\"id\"], model_name=provider_2_unique_model\n )\n _validate_default_model(vision_default) # None\n admin_basic_provider_data = _get_provider_by_name(providers, provider_2_name)\n assert admin_basic_provider_data is not None\n _validate_provider_data(\n admin_basic_provider_data,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n )\n\n # Step 6: Admin changes provider 2's default model to the shared model name\n # (same model name as provider 1 had)\n update_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=false\",\n headers=admin_user.headers,\n json={\n \"id\": provider_2[\"id\"],\n \"name\": provider_2_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000002\",\n \"model_configurations\": [c.model_dump() for c in provider_2_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert update_response.status_code == 200\n\n _set_default_provider(\n admin_user, provider_2[\"id\"], shared_model_name\n ) # Same name as provider 1's model\n\n # Step 7: Both users verify they see provider 2 as default with the shared model name\n # Validate via admin endpoint\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default, provider_id=provider_2[\"id\"], model_name=shared_model_name\n )\n _validate_default_model(vision_default) # None\n admin_provider_data = _get_provider_by_name(providers, provider_2_name)\n assert admin_provider_data is not None\n _validate_provider_data(\n admin_provider_data,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n expected_is_public=True,\n )\n\n # Validate via basic endpoint (basic_user)\n basic_data = _get_providers_basic(basic_user)\n assert basic_data is not None\n providers, text_default, vision_default = _unpack_data(basic_data)\n _validate_default_model(\n text_default, provider_id=provider_2[\"id\"], model_name=shared_model_name\n )\n _validate_default_model(vision_default) # None\n basic_provider_data = _get_provider_by_name(providers, provider_2_name)\n assert basic_provider_data is not None\n _validate_provider_data(\n basic_provider_data,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n )\n\n # Validate via basic endpoint (admin_user)\n admin_basic_data = _get_providers_basic(admin_user)\n assert admin_basic_data is not None\n providers, text_default, vision_default = _unpack_data(admin_basic_data)\n _validate_default_model(\n text_default, provider_id=provider_2[\"id\"], model_name=shared_model_name\n )\n _validate_default_model(vision_default) # None\n admin_basic_provider_data = _get_provider_by_name(providers, provider_2_name)\n assert admin_basic_provider_data is not None\n _validate_provider_data(\n admin_basic_provider_data,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n )\n\n # Verify provider 1 is no longer the default and has correct data\n admin_provider_1 = _get_provider_by_name(providers, provider_1_name)\n assert admin_provider_1 is not None\n _validate_provider_data(\n admin_provider_1,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n expected_is_public=True,\n )\n\n basic_provider_1 = _get_provider_by_name(providers, provider_1_name)\n assert basic_provider_1 is not None\n _validate_provider_data(\n basic_provider_1,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n )\n\n\ndef test_default_provider_and_vision_provider_selection(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test setting separate default providers for regular LLM and vision capabilities.\n\n This test verifies:\n 1. Create provider 1 with mixed models (some with vision, some without)\n 2. Create provider 2 with only vision-capable models\n 3. Set a non-vision model from provider 1 as the general default\n 4. Set a vision model from provider 2 as the default vision model\n 5. Verify both admin and basic users see correct default provider and vision provider\n 6. Verify model configurations show correct image support capabilities\n \"\"\"\n from onyx.auth.schemas import UserRole\n\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a non-admin user\n basic_user = UserManager.create(name=\"basic_user\")\n assert basic_user.role == UserRole.BASIC or basic_user.role != UserRole.ADMIN\n\n provider_1_name = f\"test-mixed-models-{uuid.uuid4()}\"\n provider_2_name = f\"test-vision-only-{uuid.uuid4()}\"\n\n # Provider 1: Mixed models - some with vision support, some without\n # Using real model names that litellm recognizes for vision support\n provider_1_non_vision_model = \"gpt-4\" # No vision support\n provider_1_vision_model = \"gpt-4o\" # Has vision support\n\n # Provider 2: Only vision-capable models\n provider_2_vision_model_1 = \"gpt-4-vision-preview\" # Vision model\n provider_2_vision_model_2 = \"gpt-4o-mini\" # Also has vision support\n\n # Model configurations for provider 1 (mixed)\n provider_1_configs = [\n ModelConfigurationUpsertRequest(\n name=provider_1_non_vision_model,\n is_visible=True,\n ),\n ModelConfigurationUpsertRequest(\n name=provider_1_vision_model,\n is_visible=True,\n ),\n ]\n\n # Model configurations for provider 2 (vision only)\n provider_2_configs = [\n ModelConfigurationUpsertRequest(\n name=provider_2_vision_model_1,\n is_visible=True,\n supports_image_input=True,\n ),\n ModelConfigurationUpsertRequest(\n name=provider_2_vision_model_2,\n is_visible=True,\n supports_image_input=True,\n ),\n ]\n\n # Expected model names\n provider_1_model_names = [provider_1_non_vision_model, provider_1_vision_model]\n provider_1_visible = {\n provider_1_non_vision_model: True,\n provider_1_vision_model: True,\n }\n\n provider_2_model_names = [provider_2_vision_model_1, provider_2_vision_model_2]\n provider_2_visible = {\n provider_2_vision_model_1: True,\n provider_2_vision_model_2: True,\n }\n\n # Step 1: Create provider 1 with mixed models, set non-vision model as default\n create_response_1 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": provider_1_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000001\",\n \"model_configurations\": [c.model_dump() for c in provider_1_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response_1.status_code == 200\n provider_1 = create_response_1.json()\n\n # Step 2: Create provider 2 with vision-only models\n create_response_2 = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": provider_2_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000002\",\n \"model_configurations\": [c.model_dump() for c in provider_2_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response_2.status_code == 200\n provider_2 = create_response_2.json()\n\n # Step 3: Set provider 1 as the general default provider\n _set_default_provider(admin_user, provider_1[\"id\"], provider_1_non_vision_model)\n\n # Step 4: Set provider 2 with a specific vision model as the default vision provider\n _set_default_vision_provider(\n admin_user, provider_2[\"id\"], provider_2_vision_model_1\n )\n\n # Step 5: Verify via admin endpoint\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n\n # Find and validate the default provider (provider 1)\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default,\n provider_id=provider_1[\"id\"],\n model_name=provider_1_non_vision_model,\n )\n _validate_default_model(\n vision_default,\n provider_id=provider_2[\"id\"],\n model_name=provider_2_vision_model_1,\n )\n admin_default = _get_provider_by_name(providers, provider_1_name)\n assert admin_default is not None\n _validate_provider_data(\n admin_default,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n expected_is_public=True,\n )\n\n # Find and validate the default vision provider (provider 2)\n admin_vision_default = _get_provider_by_name(providers, provider_2_name)\n assert admin_vision_default is not None\n _validate_provider_data(\n admin_vision_default,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n expected_is_public=True,\n )\n\n # Step 6: Verify via basic endpoint (basic_user)\n # Find and validate the default provider (provider 1)\n basic_data = _get_providers_basic(basic_user)\n assert basic_data is not None\n providers, text_default, vision_default = _unpack_data(basic_data)\n _validate_default_model(\n text_default,\n provider_id=provider_1[\"id\"],\n model_name=provider_1_non_vision_model,\n )\n _validate_default_model(\n vision_default,\n provider_id=provider_2[\"id\"],\n model_name=provider_2_vision_model_1,\n )\n basic_default = _get_provider_by_name(providers, provider_1_name)\n assert basic_default is not None\n _validate_provider_data(\n basic_default,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n )\n\n # Find and validate the default vision provider (provider 2)\n basic_vision_default = _get_provider_by_name(providers, provider_2_name)\n assert basic_vision_default is not None\n _validate_provider_data(\n basic_vision_default,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n )\n\n # Step 7: Verify via basic endpoint (admin_user sees same as basic_user)\n admin_basic_data = _get_providers_basic(admin_user)\n assert admin_basic_data is not None\n providers, text_default, vision_default = _unpack_data(admin_basic_data)\n _validate_default_model(\n text_default,\n provider_id=provider_1[\"id\"],\n model_name=provider_1_non_vision_model,\n )\n _validate_default_model(\n vision_default,\n provider_id=provider_2[\"id\"],\n model_name=provider_2_vision_model_1,\n )\n admin_basic_default = _get_provider_by_name(providers, provider_1_name)\n assert admin_basic_default is not None\n _validate_provider_data(\n admin_basic_default,\n expected_name=provider_1_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_1_model_names,\n expected_visible=provider_1_visible,\n )\n\n admin_basic_vision_default = _get_provider_by_name(providers, provider_2_name)\n assert admin_basic_vision_default is not None\n _validate_provider_data(\n admin_basic_vision_default,\n expected_name=provider_2_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=provider_2_model_names,\n expected_visible=provider_2_visible,\n )\n\n # Verify that the providers are distinct (different providers for regular vs vision)\n assert (\n admin_default[\"name\"] != admin_vision_default[\"name\"]\n ), \"Default provider and vision provider should be different providers\"\n assert (\n basic_default[\"name\"] != basic_vision_default[\"name\"]\n ), \"Default provider and vision provider should be different providers (basic endpoint)\"\n\n\ndef test_default_provider_is_not_default_vision_provider(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"\n Test that setting a provider as the default provider does NOT make it\n the default vision provider.\n\n This test verifies:\n 1. Create a provider with some models\n 2. Set it as the default provider\n 3. Verify it is the default provider (is_default_provider=True)\n 4. Verify it is NOT the default vision provider (is_default_vision_provider should be None/False)\n \"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n provider_name = f\"test-default-not-vision-{uuid.uuid4()}\"\n\n # Model configurations\n model_configs = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4\",\n is_visible=True,\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\",\n is_visible=True,\n ),\n ]\n\n expected_model_names = [\"gpt-4\", \"gpt-4o\"]\n expected_visible = {\"gpt-4\": True, \"gpt-4o\": True}\n\n # Step 1: Create the provider\n create_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": provider_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000000\",\n \"model_configurations\": [c.model_dump() for c in model_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_response.status_code == 200\n created_provider = create_response.json()\n\n # Step 2: Set it as the default provider\n _set_default_provider(admin_user, created_provider[\"id\"], \"gpt-4\")\n\n # Step 3 & 4: Verify via admin endpoint\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default, provider_id=created_provider[\"id\"], model_name=\"gpt-4\"\n )\n _validate_default_model(vision_default) # None\n admin_provider_data = _get_provider_by_name(providers, provider_name)\n assert admin_provider_data is not None\n\n # Full validation of provider data\n _validate_provider_data(\n admin_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n expected_is_public=True,\n )\n\n # Also verify via basic endpoint\n basic_data = _get_providers_basic(admin_user)\n assert basic_data is not None\n providers, text_default, vision_default = _unpack_data(basic_data)\n _validate_default_model(\n text_default, provider_id=created_provider[\"id\"], model_name=\"gpt-4\"\n )\n _validate_default_model(vision_default) # None\n basic_provider_data = _get_provider_by_name(providers, provider_name)\n assert basic_provider_data is not None\n\n _validate_provider_data(\n basic_provider_data,\n expected_name=provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=expected_model_names,\n expected_visible=expected_visible,\n )\n\n\ndef _get_all_image_gen_configs(admin_user: DATestUser) -> list[dict]:\n \"\"\"Utility function to fetch all image generation configs.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n return response.json()\n\n\ndef _create_image_gen_config(\n admin_user: DATestUser,\n image_provider_id: str,\n model_name: str,\n source_llm_provider_id: int,\n is_default: bool = False,\n) -> dict:\n \"\"\"Utility function to create an image generation config using clone mode.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config\",\n headers=admin_user.headers,\n json={\n \"image_provider_id\": image_provider_id,\n \"model_name\": model_name,\n \"source_llm_provider_id\": source_llm_provider_id,\n \"is_default\": is_default,\n },\n )\n assert (\n response.status_code == 200\n ), f\"Failed to create image gen config: {response.text}\"\n return response.json()\n\n\ndef _set_image_gen_config_default(\n admin_user: DATestUser, image_provider_id: str\n) -> None:\n \"\"\"Utility function to set an image generation config as default.\"\"\"\n response = requests.post(\n f\"{API_SERVER_URL}/admin/image-generation/config/{image_provider_id}/default\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n\ndef _delete_image_gen_config(admin_user: DATestUser, image_provider_id: str) -> None:\n \"\"\"Utility function to delete an image generation config.\"\"\"\n response = requests.delete(\n f\"{API_SERVER_URL}/admin/image-generation/config/{image_provider_id}\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n\n\ndef test_all_three_provider_types_no_mixup(reset: None) -> None: # noqa: ARG001\n \"\"\"\n Test that regular LLM providers, vision providers, and image generation providers\n are all tracked separately with no mixup.\n\n This test verifies:\n 1. Create a regular LLM provider and set as default\n 2. Create a vision LLM provider and set as default vision\n 3. Create an image generation config (using clone mode from regular provider)\n 4. Set the image gen config as default\n 5. Verify all three are correctly identified:\n - Regular provider: is_default_provider=True, is_default_vision_provider=None\n - Vision provider: is_default_provider=None, is_default_vision_provider=True\n - Image gen config: is_default=True (separate from LLM provider defaults)\n 6. Verify image gen config doesn't appear in LLM provider lists\n 7. Verify LLM providers don't appear in image gen config list\n \"\"\"\n from onyx.auth.schemas import UserRole\n\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a non-admin user\n basic_user = UserManager.create(name=\"basic_user\")\n assert basic_user.role == UserRole.BASIC or basic_user.role != UserRole.ADMIN\n\n # Provider names\n regular_provider_name = f\"test-regular-provider-{uuid.uuid4()}\"\n vision_provider_name = f\"test-vision-provider-{uuid.uuid4()}\"\n image_gen_provider_id = f\"test-image-gen-{uuid.uuid4()}\"\n\n # Model configurations\n regular_model_configs = [\n ModelConfigurationUpsertRequest(name=\"gpt-4\", is_visible=True),\n ModelConfigurationUpsertRequest(name=\"gpt-4o\", is_visible=True),\n ]\n\n vision_model_configs = [\n ModelConfigurationUpsertRequest(\n name=\"gpt-4-vision-preview\", is_visible=True, supports_image_input=True\n ),\n ModelConfigurationUpsertRequest(\n name=\"gpt-4o\", is_visible=True, supports_image_input=True\n ),\n ]\n\n # Step 1: Create regular LLM provider\n create_regular_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": regular_provider_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000001\",\n \"model_configurations\": [c.model_dump() for c in regular_model_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_regular_response.status_code == 200\n regular_provider = create_regular_response.json()\n\n # Set as default provider\n _set_default_provider(admin_user, regular_provider[\"id\"], \"gpt-4\")\n\n # Step 2: Create vision LLM provider\n create_vision_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json={\n \"name\": vision_provider_name,\n \"provider\": LlmProviderNames.OPENAI,\n \"api_key\": \"sk-000000000000000000000000000000000000000000000002\",\n \"default_model_name\": \"gpt-4-vision-preview\",\n \"model_configurations\": [c.model_dump() for c in vision_model_configs],\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n },\n )\n assert create_vision_response.status_code == 200\n vision_provider = create_vision_response.json()\n\n # Set as default vision provider\n _set_default_vision_provider(\n admin_user, vision_provider[\"id\"], \"gpt-4-vision-preview\"\n )\n\n # Step 3: Create image generation config using clone mode from regular provider\n _create_image_gen_config(\n admin_user=admin_user,\n image_provider_id=image_gen_provider_id,\n model_name=\"dall-e-3\",\n source_llm_provider_id=regular_provider[\"id\"],\n is_default=True,\n )\n\n # Step 4: Verify all three types are correctly tracked\n\n # Get all LLM providers (via admin endpoint)\n admin_data = _get_providers_admin(admin_user)\n assert admin_data is not None\n providers, text_default, vision_default = _unpack_data(admin_data)\n _validate_default_model(\n text_default, provider_id=regular_provider[\"id\"], model_name=\"gpt-4\"\n )\n _validate_default_model(\n vision_default,\n provider_id=vision_provider[\"id\"],\n model_name=\"gpt-4-vision-preview\",\n )\n _validate_default_model(\n vision_default, vision_provider[\"id\"], \"gpt-4-vision-preview\"\n )\n _get_provider_by_name(providers, regular_provider_name)\n\n # Get all image generation configs\n image_gen_configs = _get_all_image_gen_configs(admin_user)\n\n # Verify the regular provider is the default provider\n admin_regular_provider_data = _get_provider_by_name(\n providers, regular_provider_name\n )\n assert admin_regular_provider_data is not None\n _validate_provider_data(\n admin_regular_provider_data,\n expected_name=regular_provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=[c.name for c in regular_model_configs],\n expected_visible={c.name: True for c in regular_model_configs},\n )\n admin_vision_provider_data = _get_provider_by_name(providers, vision_provider_name)\n assert admin_vision_provider_data is not None\n _validate_provider_data(\n admin_vision_provider_data,\n expected_name=vision_provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=[c.name for c in vision_model_configs],\n expected_visible={c.name: True for c in vision_model_configs},\n )\n\n # Verify the image gen config is the default image generation config\n image_gen_config_data = next(\n (\n c\n for c in image_gen_configs\n if c[\"image_provider_id\"] == image_gen_provider_id\n ),\n None,\n )\n assert image_gen_config_data is not None, \"Image gen config not found\"\n assert (\n image_gen_config_data[\"is_default\"] is True\n ), \"Image gen config should be the default\"\n assert (\n image_gen_config_data[\"model_name\"] == \"dall-e-3\"\n ), \"Image gen config should have correct model name\"\n\n # Step 5: Verify no mixup - image gen providers don't appear in LLM provider lists\n # Image gen provider should not appear in the list\n assert image_gen_provider_id not in [p[\"name\"] for p in providers]\n\n # Step 6: Verify via basic endpoint (non-admin user)\n basic_data = _get_providers_basic(basic_user)\n assert basic_data is not None\n providers, text_default, vision_default = _unpack_data(basic_data)\n _validate_default_model(\n text_default, provider_id=regular_provider[\"id\"], model_name=\"gpt-4\"\n )\n _validate_default_model(\n vision_default,\n provider_id=vision_provider[\"id\"],\n model_name=\"gpt-4-vision-preview\",\n )\n _validate_default_model(\n vision_default, vision_provider[\"id\"], \"gpt-4-vision-preview\"\n )\n basic_provider_data = _get_provider_by_name(providers, regular_provider_name)\n assert basic_provider_data is not None\n _validate_provider_data(\n basic_provider_data,\n expected_name=regular_provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=[c.name for c in regular_model_configs],\n expected_visible={c.name: True for c in regular_model_configs},\n )\n basic_vision_provider_data = _get_provider_by_name(providers, vision_provider_name)\n assert basic_vision_provider_data is not None\n _validate_provider_data(\n basic_vision_provider_data,\n expected_name=vision_provider_name,\n expected_provider=LlmProviderNames.OPENAI,\n expected_model_names=[c.name for c in vision_model_configs],\n expected_visible={c.name: True for c in vision_model_configs},\n )\n\n # Step 7: Verify the counts are as expected\n # We should have at least 2 user-created providers (setup_postgres may add more)\n assert len(providers) >= 2\n assert len(image_gen_configs) == 1\n\n # Clean up: Delete the image gen config (to clean up the internal LLM provider)\n _delete_image_gen_config(admin_user, image_gen_provider_id)\n" }, { "path": "backend/tests/integration/tests/llm_provider/test_llm_provider_access_control.py", "content": "import os\n\nimport pytest\nimport requests\nfrom sqlalchemy.orm import Session\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.llm import can_user_access_llm_provider\nfrom onyx.db.llm import fetch_user_group_ids\nfrom onyx.db.llm import update_default_provider\nfrom onyx.db.llm import upsert_llm_provider\nfrom onyx.db.models import LLMProvider as LLMProviderModel\nfrom onyx.db.models import LLMProvider__Persona\nfrom onyx.db.models import LLMProvider__UserGroup\nfrom onyx.db.models import Persona\nfrom onyx.db.models import User\nfrom onyx.db.models import User__UserGroup\nfrom onyx.db.models import UserGroup\nfrom onyx.llm.constants import LlmProviderNames\nfrom onyx.llm.factory import get_llm_for_persona\nfrom onyx.server.manage.llm.models import LLMProviderUpsertRequest\nfrom onyx.server.manage.llm.models import ModelConfigurationUpsertRequest\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\npytestmark = pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"LLM provider access control is enterprise only\",\n)\n\n\ndef _create_llm_provider(\n db_session: Session,\n *,\n name: str,\n default_model_name: str,\n is_public: bool,\n is_default: bool,\n) -> LLMProviderModel:\n _provider = upsert_llm_provider(\n llm_provider_upsert_request=LLMProviderUpsertRequest(\n name=name,\n provider=LlmProviderNames.OPENAI,\n api_key=None,\n api_base=None,\n api_version=None,\n custom_config=None,\n is_public=is_public,\n model_configurations=[\n ModelConfigurationUpsertRequest(\n name=default_model_name,\n is_visible=True,\n )\n ],\n ),\n db_session=db_session,\n )\n if is_default:\n update_default_provider(_provider.id, default_model_name, db_session)\n\n provider = db_session.get(LLMProviderModel, _provider.id)\n if not provider:\n raise ValueError(f\"Provider {name} not found\")\n return provider\n\n\ndef _create_persona(\n db_session: Session,\n *,\n name: str,\n provider_name: str,\n) -> Persona:\n persona = Persona(\n name=name,\n description=f\"{name} description\",\n llm_model_provider_override=provider_name,\n llm_model_version_override=\"gpt-4o-mini\",\n system_prompt=\"System prompt\",\n task_prompt=\"Task prompt\",\n datetime_aware=True,\n is_public=True,\n )\n db_session.add(persona)\n db_session.flush()\n return persona\n\n\n@pytest.fixture()\ndef users(reset: None) -> tuple[DATestUser, DATestUser]: # noqa: ARG001\n admin_user = UserManager.create(name=\"admin_user\")\n basic_user = UserManager.create(name=\"basic_user\")\n return admin_user, basic_user\n\n\ndef test_can_user_access_llm_provider_or_logic(\n users: tuple[DATestUser, DATestUser],\n) -> None:\n \"\"\"Test LLM provider access control with is_public flag and AND logic.\n\n Tests the new access control logic:\n - is_public=True providers are accessible to everyone\n - is_public=False with no restrictions locks the provider\n - When both groups AND personas are set, AND logic applies (must satisfy both)\n \"\"\"\n admin_user, basic_user = users\n\n with get_session_with_current_tenant() as db_session:\n # Public provider - accessible to everyone\n default_provider = _create_llm_provider(\n db_session,\n name=\"default-provider\",\n default_model_name=\"gpt-4o\",\n is_public=True,\n is_default=True,\n )\n # Locked provider - is_public=False with no restrictions\n locked_provider = _create_llm_provider(\n db_session,\n name=\"locked-provider\",\n default_model_name=\"gpt-4o\",\n is_public=False,\n is_default=False,\n )\n # Restricted provider - has both group AND persona restrictions (AND logic)\n restricted_provider = _create_llm_provider(\n db_session,\n name=\"restricted-provider\",\n default_model_name=\"gpt-4o-mini\",\n is_public=False,\n is_default=False,\n )\n\n allowed_persona = _create_persona(\n db_session,\n name=\"allowed-persona\",\n provider_name=restricted_provider.name,\n )\n blocked_persona = _create_persona(\n db_session,\n name=\"blocked-persona\",\n provider_name=restricted_provider.name,\n )\n\n access_group = UserGroup(name=\"access-group\")\n db_session.add(access_group)\n db_session.flush()\n\n # Add both group and persona restrictions to restricted_provider\n db_session.add(\n LLMProvider__UserGroup(\n llm_provider_id=restricted_provider.id,\n user_group_id=access_group.id,\n )\n )\n db_session.add(\n LLMProvider__Persona(\n llm_provider_id=restricted_provider.id,\n persona_id=allowed_persona.id,\n )\n )\n # Only admin_user is in the access_group\n db_session.add(\n User__UserGroup(\n user_group_id=access_group.id,\n user_id=admin_user.id,\n )\n )\n db_session.flush()\n\n db_session.refresh(restricted_provider)\n db_session.refresh(locked_provider)\n\n admin_model = db_session.get(User, admin_user.id)\n basic_model = db_session.get(User, basic_user.id)\n\n assert admin_model is not None\n assert basic_model is not None\n\n # Fetch user group IDs for both users\n admin_group_ids = fetch_user_group_ids(db_session, admin_model)\n basic_group_ids = fetch_user_group_ids(db_session, basic_model)\n\n # Test is_public flag\n assert default_provider.is_public\n assert not locked_provider.is_public\n assert not restricted_provider.is_public\n\n # Public provider - everyone can access\n assert can_user_access_llm_provider(\n default_provider,\n admin_group_ids,\n allowed_persona,\n )\n assert can_user_access_llm_provider(\n default_provider,\n basic_group_ids,\n blocked_persona,\n )\n\n # Locked provider (is_public=False, no restrictions) - nobody can access\n assert not can_user_access_llm_provider(\n locked_provider,\n admin_group_ids,\n allowed_persona,\n )\n assert not can_user_access_llm_provider(\n locked_provider,\n basic_group_ids,\n allowed_persona,\n )\n\n # Restricted provider with AND logic (both groups AND personas set)\n # admin_user in group + allowed_persona whitelisted → SUCCESS (both conditions met)\n assert can_user_access_llm_provider(\n restricted_provider,\n admin_group_ids,\n allowed_persona,\n )\n\n # admin_user in group + blocked_persona not whitelisted → FAIL (persona not allowed)\n assert not can_user_access_llm_provider(\n restricted_provider,\n admin_group_ids,\n blocked_persona,\n )\n\n # basic_user not in group + allowed_persona whitelisted → FAIL (user not in group)\n assert not can_user_access_llm_provider(\n restricted_provider,\n basic_group_ids,\n allowed_persona,\n )\n\n # basic_user not in group + blocked_persona not whitelisted → FAIL (neither condition met)\n assert not can_user_access_llm_provider(\n restricted_provider,\n basic_group_ids,\n blocked_persona,\n )\n\n\ndef test_public_provider_with_persona_restrictions(\n users: tuple[DATestUser, DATestUser],\n) -> None:\n \"\"\"Public providers should still enforce persona restrictions.\n\n Regression test for the bug where is_public=True caused\n can_user_access_llm_provider() to return True immediately,\n bypassing persona whitelist checks entirely.\n \"\"\"\n admin_user, _basic_user = users\n\n with get_session_with_current_tenant() as db_session:\n # Public provider with persona restrictions\n public_restricted = _create_llm_provider(\n db_session,\n name=\"public-persona-restricted\",\n default_model_name=\"gpt-4o\",\n is_public=True,\n is_default=True,\n )\n\n whitelisted_persona = _create_persona(\n db_session,\n name=\"whitelisted-persona\",\n provider_name=public_restricted.name,\n )\n non_whitelisted_persona = _create_persona(\n db_session,\n name=\"non-whitelisted-persona\",\n provider_name=public_restricted.name,\n )\n\n # Only whitelist one persona\n db_session.add(\n LLMProvider__Persona(\n llm_provider_id=public_restricted.id,\n persona_id=whitelisted_persona.id,\n )\n )\n db_session.flush()\n db_session.refresh(public_restricted)\n\n admin_model = db_session.get(User, admin_user.id)\n assert admin_model is not None\n admin_group_ids = fetch_user_group_ids(db_session, admin_model)\n\n # Whitelisted persona — should be allowed\n assert can_user_access_llm_provider(\n public_restricted,\n admin_group_ids,\n whitelisted_persona,\n )\n\n # Non-whitelisted persona — should be denied despite is_public=True\n assert not can_user_access_llm_provider(\n public_restricted,\n admin_group_ids,\n non_whitelisted_persona,\n )\n\n # No persona context (e.g. global provider list) — should be denied\n # because provider has persona restrictions set\n assert not can_user_access_llm_provider(\n public_restricted,\n admin_group_ids,\n persona=None,\n )\n\n\ndef test_public_provider_without_persona_restrictions(\n users: tuple[DATestUser, DATestUser],\n) -> None:\n \"\"\"Public providers with no persona restrictions remain accessible to all.\"\"\"\n admin_user, basic_user = users\n\n with get_session_with_current_tenant() as db_session:\n public_unrestricted = _create_llm_provider(\n db_session,\n name=\"public-unrestricted\",\n default_model_name=\"gpt-4o\",\n is_public=True,\n is_default=True,\n )\n\n any_persona = _create_persona(\n db_session,\n name=\"any-persona\",\n provider_name=public_unrestricted.name,\n )\n\n admin_model = db_session.get(User, admin_user.id)\n basic_model = db_session.get(User, basic_user.id)\n assert admin_model is not None\n assert basic_model is not None\n\n admin_group_ids = fetch_user_group_ids(db_session, admin_model)\n basic_group_ids = fetch_user_group_ids(db_session, basic_model)\n\n # Any user, any persona — all allowed\n assert can_user_access_llm_provider(\n public_unrestricted, admin_group_ids, any_persona\n )\n assert can_user_access_llm_provider(\n public_unrestricted, basic_group_ids, any_persona\n )\n assert can_user_access_llm_provider(\n public_unrestricted, admin_group_ids, persona=None\n )\n\n\ndef test_get_llm_for_persona_falls_back_when_access_denied(\n users: tuple[DATestUser, DATestUser],\n) -> None:\n admin_user, basic_user = users\n\n with get_session_with_current_tenant() as db_session:\n default_provider = _create_llm_provider(\n db_session,\n name=\"default-provider\",\n default_model_name=\"gpt-4o\",\n is_public=True,\n is_default=True,\n )\n restricted_provider = _create_llm_provider(\n db_session,\n name=\"restricted-provider\",\n default_model_name=\"gpt-4o-mini\",\n is_public=False,\n is_default=False,\n )\n\n persona = _create_persona(\n db_session,\n name=\"fallback-persona\",\n provider_name=restricted_provider.name,\n )\n\n access_group = UserGroup(name=\"persona-group\")\n db_session.add(access_group)\n db_session.flush()\n\n db_session.add(\n LLMProvider__UserGroup(\n llm_provider_id=restricted_provider.id,\n user_group_id=access_group.id,\n )\n )\n db_session.add(\n User__UserGroup(\n user_group_id=access_group.id,\n user_id=admin_user.id,\n )\n )\n db_session.flush()\n db_session.commit()\n\n db_session.refresh(default_provider)\n db_session.refresh(restricted_provider)\n db_session.refresh(persona)\n\n admin_model = db_session.get(User, admin_user.id)\n basic_model = db_session.get(User, basic_user.id)\n\n assert admin_model is not None\n assert basic_model is not None\n\n allowed_llm = get_llm_for_persona(\n persona=persona,\n user=admin_model,\n )\n assert (\n allowed_llm.config.model_name\n == restricted_provider.model_configurations[0].name\n )\n\n fallback_llm = get_llm_for_persona(\n persona=persona,\n user=basic_model,\n )\n assert (\n fallback_llm.config.model_name\n == default_provider.model_configurations[0].name\n )\n\n\ndef test_list_llm_provider_basics_excludes_non_public_unrestricted(\n users: tuple[DATestUser, DATestUser],\n) -> None:\n \"\"\"Test that the /llm/provider endpoint correctly excludes non-public providers\n with no group/persona restrictions.\n\n This tests the fix for the bug where non-public providers with no restrictions\n were incorrectly shown to all users instead of being admin-only.\n \"\"\"\n admin_user, basic_user = users\n\n # Create a public provider (should be visible to all)\n public_provider = LLMProviderManager.create(\n name=\"public-provider\",\n is_public=True,\n set_as_default=True,\n default_model_name=\"gpt-4o\",\n user_performing_action=admin_user,\n )\n\n # Create a non-public provider with no restrictions (should be admin-only)\n non_public_provider = LLMProviderManager.create(\n name=\"non-public-unrestricted\",\n is_public=False,\n groups=[],\n personas=[],\n set_as_default=False,\n user_performing_action=admin_user,\n )\n\n # Non-admin user calls the /llm/provider endpoint\n response = requests.get(\n f\"{API_SERVER_URL}/llm/provider\",\n headers=basic_user.headers,\n )\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n provider_names = [p[\"name\"] for p in providers]\n\n # Public provider should be visible\n assert public_provider.name in provider_names\n\n # Non-public provider with no restrictions should NOT be visible to non-admin\n assert non_public_provider.name not in provider_names\n\n # Admin user should see both providers\n admin_response = requests.get(\n f\"{API_SERVER_URL}/llm/provider\",\n headers=admin_user.headers,\n )\n assert admin_response.status_code == 200\n admin_providers = admin_response.json()[\"providers\"]\n admin_provider_names = [p[\"name\"] for p in admin_providers]\n\n assert public_provider.name in admin_provider_names\n assert non_public_provider.name in admin_provider_names\n\n\ndef test_provider_delete_clears_persona_references(\n reset: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Test that deleting a provider automatically clears persona references.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n\n # Create a default provider first so personas have something to fall back to\n LLMProviderManager.create(\n name=\"default-provider\",\n is_public=True,\n set_as_default=True,\n default_model_name=\"gpt-4o\",\n user_performing_action=admin_user,\n )\n\n provider = LLMProviderManager.create(\n is_public=False,\n set_as_default=False,\n user_performing_action=admin_user,\n )\n persona = PersonaManager.create(\n llm_model_provider_override=provider.name,\n user_performing_action=admin_user,\n )\n\n # Delete the provider - should succeed and automatically clear persona references\n assert LLMProviderManager.delete(\n provider,\n user_performing_action=admin_user,\n )\n\n # Verify the persona now falls back to default (llm_model_provider_override cleared)\n persona_response = requests.get(\n f\"{API_SERVER_URL}/persona/{persona.id}\",\n headers=admin_user.headers,\n )\n assert persona_response.status_code == 200\n updated_persona = persona_response.json()\n assert updated_persona[\"llm_model_provider_override\"] is None\n" }, { "path": "backend/tests/integration/tests/llm_provider/test_llm_provider_persona_access.py", "content": "\"\"\"\nIntegration tests for LLM Provider persona access authorization.\n\"\"\"\n\nimport os\n\nimport pytest\nimport requests\n\nfrom onyx.llm.constants import LlmProviderNames\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\npytestmark = pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"LLM provider persona access is enterprise only\",\n)\n\n\n@pytest.fixture()\ndef users_and_groups(\n reset: None, # noqa: ARG001\n) -> tuple[DATestUser, DATestUser, int, int]:\n \"\"\"Create admin, basic user, and two user groups.\"\"\"\n admin_user = UserManager.create(name=\"admin_user\")\n basic_user = UserManager.create(name=\"basic_user\")\n\n # Create two user groups\n group1 = UserGroupManager.create(\n user_performing_action=admin_user,\n name=\"test_group_1\",\n user_ids=[basic_user.id],\n )\n\n group2 = UserGroupManager.create(\n user_performing_action=admin_user,\n name=\"test_group_2\",\n user_ids=[], # basic_user is NOT in this group\n )\n\n return admin_user, basic_user, group1.id, group2.id\n\n\ndef test_unauthorized_persona_access_returns_403(\n users_and_groups: tuple[DATestUser, DATestUser, int, int],\n) -> None:\n \"\"\"Test that users cannot query providers for personas they don't have access to.\"\"\"\n admin_user, basic_user, group1_id, group2_id = users_and_groups\n\n # Create a persona restricted to group2 (which basic_user is NOT in)\n restricted_persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Restricted Persona\",\n description=\"Only accessible to group2\",\n is_public=False,\n groups=[group2_id],\n )\n\n # Try to query providers for the restricted persona as basic_user\n response = requests.get(\n f\"{API_SERVER_URL}/llm/persona/{restricted_persona.id}/providers\",\n headers=basic_user.headers,\n )\n\n # Should return 403 Forbidden\n assert response.status_code == 403\n assert \"don't have access to this assistant\" in response.json()[\"detail\"]\n\n\ndef test_authorized_persona_access_returns_filtered_providers(\n users_and_groups: tuple[DATestUser, DATestUser, int, int],\n) -> None:\n \"\"\"Test that users can query providers for personas they have access to.\"\"\"\n admin_user, basic_user, group1_id, group2_id = users_and_groups\n\n # Create a persona accessible to group1 (which basic_user IS in)\n accessible_persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Accessible Persona\",\n description=\"Accessible to group1\",\n is_public=False,\n groups=[group1_id],\n )\n\n # Create a restricted provider accessible only to the persona\n restricted_provider = LLMProviderManager.create(\n user_performing_action=admin_user,\n name=\"Restricted Provider\",\n provider=LlmProviderNames.OPENAI,\n api_key=\"test-key\",\n default_model_name=\"gpt-4o\",\n is_public=False,\n groups=[],\n personas=[accessible_persona.id],\n )\n\n # Query providers for the accessible persona as basic_user\n response = requests.get(\n f\"{API_SERVER_URL}/llm/persona/{accessible_persona.id}/providers\",\n headers=basic_user.headers,\n )\n\n # Should succeed\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n\n # Should include the restricted provider since basic_user can access the persona\n provider_names = [p[\"name\"] for p in providers]\n assert restricted_provider.name in provider_names\n\n\ndef test_persona_id_zero_applies_rbac(\n users_and_groups: tuple[DATestUser, DATestUser, int, int],\n) -> None:\n \"\"\"Test that persona_id=0 (default persona) properly applies RBAC.\"\"\"\n admin_user, basic_user, group1_id, group2_id = users_and_groups\n\n # Create a restricted provider accessible only to group2\n restricted_provider = LLMProviderManager.create(\n user_performing_action=admin_user,\n name=\"Group2 Only Provider\",\n provider=LlmProviderNames.OPENAI,\n api_key=\"test-key\",\n default_model_name=\"gpt-4o\",\n is_public=False,\n groups=[group2_id],\n personas=[],\n )\n\n # Query providers with persona_id=0 as basic_user\n response = requests.get(\n f\"{API_SERVER_URL}/llm/persona/0/providers\",\n headers=basic_user.headers,\n )\n\n # Should succeed (persona_id=0 refers to default persona, which is public)\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n\n # Should NOT include the restricted provider since basic_user is not in group2\n provider_names = [p[\"name\"] for p in providers]\n assert restricted_provider.name not in provider_names\n\n\ndef test_admin_can_query_any_persona(\n users_and_groups: tuple[DATestUser, DATestUser, int, int],\n) -> None:\n \"\"\"Test that admin users can query any persona's providers.\"\"\"\n admin_user, basic_user, group1_id, group2_id = users_and_groups\n\n # Create a persona restricted to group2 (admin is not explicitly in this group)\n restricted_persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Admin Test Persona\",\n description=\"Only accessible to group2\",\n is_public=False,\n groups=[group2_id],\n )\n\n # Create a restricted provider accessible only to the persona\n restricted_provider = LLMProviderManager.create(\n user_performing_action=admin_user,\n name=\"Admin Test Provider\",\n provider=LlmProviderNames.OPENAI,\n api_key=\"test-key\",\n default_model_name=\"gpt-4o\",\n is_public=False,\n groups=[],\n personas=[restricted_persona.id],\n )\n\n # Query providers for the restricted persona as admin_user\n response = requests.get(\n f\"{API_SERVER_URL}/llm/persona/{restricted_persona.id}/providers\",\n headers=admin_user.headers,\n )\n\n # Should succeed - admins can access any persona\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n\n # Should include the restricted provider\n provider_names = [p[\"name\"] for p in providers]\n assert restricted_provider.name in provider_names\n\n\ndef test_public_persona_accessible_to_all(\n users_and_groups: tuple[DATestUser, DATestUser, int, int],\n) -> None:\n \"\"\"Test that public personas are accessible to all users.\"\"\"\n admin_user, basic_user, group1_id, group2_id = users_and_groups\n\n # Create a public LLM provider so there's something to return\n public_provider = LLMProviderManager.create(\n user_performing_action=admin_user,\n name=\"Public Provider\",\n provider=LlmProviderNames.OPENAI,\n api_key=\"test-key\",\n default_model_name=\"gpt-4o\",\n is_public=True,\n set_as_default=True,\n )\n\n # Create a public persona\n public_persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Public Persona\",\n description=\"Accessible to everyone\",\n is_public=True,\n groups=[],\n )\n\n # Query providers for the public persona as basic_user\n response = requests.get(\n f\"{API_SERVER_URL}/llm/persona/{public_persona.id}/providers\",\n headers=basic_user.headers,\n )\n\n # Should succeed\n assert response.status_code == 200\n providers = response.json()[\"providers\"]\n\n # Should return the public provider\n assert len(providers) > 0\n provider_names = [p[\"name\"] for p in providers]\n assert public_provider.name in provider_names\n\n\ndef test_nonexistent_persona_returns_404(\n users_and_groups: tuple[DATestUser, DATestUser, int, int],\n) -> None:\n \"\"\"Test that querying a nonexistent persona returns 404.\"\"\"\n admin_user, basic_user, group1_id, group2_id = users_and_groups\n\n # Query providers for a nonexistent persona\n response = requests.get(\n f\"{API_SERVER_URL}/llm/persona/99999/providers\",\n headers=basic_user.headers,\n )\n\n # Should return 404\n assert response.status_code == 404\n assert \"Persona not found\" in response.json()[\"detail\"]\n" }, { "path": "backend/tests/integration/tests/llm_workflows/test_mock_llm_tool_calls.py", "content": "from onyx.configs import app_configs\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.tools.constants import SEARCH_TOOL_ID\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.tool import ToolManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n_DUMMY_OPENAI_API_KEY = \"sk-mock-llm-workflow-tests\"\n\n\ndef _get_internal_search_tool_id(admin_user: DATestUser) -> int:\n tools = ToolManager.list_tools(user_performing_action=admin_user)\n for tool in tools:\n if tool.in_code_tool_id == SEARCH_TOOL_ID:\n return tool.id\n raise AssertionError(\"SearchTool must exist for this test\")\n\n\ndef _assert_integration_mode_enabled() -> None:\n assert (\n app_configs.INTEGRATION_TESTS_MODE is True\n ), \"Integration tests require INTEGRATION_TESTS_MODE=true.\"\n\n\ndef _seed_connector_for_search_tool(admin_user: DATestUser) -> None:\n # SearchTool is only exposed when at least one non-default connector exists.\n CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n\ndef test_mock_llm_response_single_tool_call_debug(admin_user: DATestUser) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n\n LLMProviderManager.create(\n user_performing_action=admin_user,\n api_key=_DUMMY_OPENAI_API_KEY,\n )\n chat_session = ChatSessionManager.create(user_performing_action=admin_user)\n search_tool_id = _get_internal_search_tool_id(admin_user)\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"run the search tool\",\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n mock_llm_response='{\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"alpha\"]}}',\n )\n\n assert response.error is None, f\"Unexpected stream error: {response.error}\"\n assert len(response.tool_call_debug) == 1\n assert response.tool_call_debug[0].tool_name == \"internal_search\"\n assert response.tool_call_debug[0].tool_args == {\"queries\": [\"alpha\"]}\n\n\ndef test_mock_llm_response_parallel_tool_call_debug(admin_user: DATestUser) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n\n LLMProviderManager.create(\n user_performing_action=admin_user,\n api_key=_DUMMY_OPENAI_API_KEY,\n )\n chat_session = ChatSessionManager.create(user_performing_action=admin_user)\n search_tool_id = _get_internal_search_tool_id(admin_user)\n\n mock_response = \"\\n\".join(\n [\n '{\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"alpha\"]}}',\n '{\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"beta\"]}}',\n ]\n )\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"run the search tool twice\",\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n mock_llm_response=mock_response,\n )\n\n assert response.error is None, f\"Unexpected stream error: {response.error}\"\n assert len(response.tool_call_debug) == 2\n assert [entry.tool_name for entry in response.tool_call_debug] == [\n \"internal_search\",\n \"internal_search\",\n ]\n assert [entry.tool_args for entry in response.tool_call_debug] == [\n {\"queries\": [\"alpha\"]},\n {\"queries\": [\"beta\"]},\n ]\n\n\ndef test_mock_llm_response_embedded_json_fallback_tool_call_debug(\n admin_user: DATestUser,\n) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n\n LLMProviderManager.create(\n user_performing_action=admin_user,\n api_key=_DUMMY_OPENAI_API_KEY,\n )\n chat_session = ChatSessionManager.create(user_performing_action=admin_user)\n search_tool_id = _get_internal_search_tool_id(admin_user)\n\n # Validate fallback extraction when the model returns tool-call JSON embedded in\n # normal assistant text instead of structured tool_call objects.\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"use the search tool\",\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n mock_llm_response=(\n 'I will call a tool now. {\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"gamma\"]}}'\n ),\n )\n\n assert response.error is None, f\"Unexpected stream error: {response.error}\"\n assert len(response.tool_call_debug) == 1\n assert response.tool_call_debug[0].tool_name == \"internal_search\"\n assert response.tool_call_debug[0].tool_args == {\"queries\": [\"gamma\"]}\n" }, { "path": "backend/tests/integration/tests/llm_workflows/test_nightly_provider_chat_workflow.py", "content": "import json\nimport os\nimport time\nfrom uuid import uuid4\n\nimport pytest\nimport requests\nfrom pydantic import BaseModel\nfrom pydantic import ConfigDict\n\nfrom onyx.configs import app_configs\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.tools.constants import SEARCH_TOOL_ID\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.tool import ToolManager\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import ToolName\n\n\n_ENV_PROVIDER = \"NIGHTLY_LLM_PROVIDER\"\n_ENV_MODELS = \"NIGHTLY_LLM_MODELS\"\n_ENV_API_KEY = \"NIGHTLY_LLM_API_KEY\"\n_ENV_API_BASE = \"NIGHTLY_LLM_API_BASE\"\n_ENV_API_VERSION = \"NIGHTLY_LLM_API_VERSION\"\n_ENV_DEPLOYMENT_NAME = \"NIGHTLY_LLM_DEPLOYMENT_NAME\"\n_ENV_CUSTOM_CONFIG_JSON = \"NIGHTLY_LLM_CUSTOM_CONFIG_JSON\"\n_ENV_STRICT = \"NIGHTLY_LLM_STRICT\"\n\n\nclass NightlyProviderConfig(BaseModel):\n model_config = ConfigDict(frozen=True)\n\n provider: str\n model_names: list[str]\n api_key: str | None\n api_base: str | None\n api_version: str | None\n deployment_name: str | None\n custom_config: dict[str, str] | None\n strict: bool\n\n\ndef _stringify_custom_config_value(value: object) -> str:\n if isinstance(value, str):\n return value\n if isinstance(value, (dict, list)):\n return json.dumps(value)\n return str(value)\n\n\ndef _looks_like_vertex_credentials_payload(\n raw_custom_config: dict[object, object],\n) -> bool:\n normalized_keys = {str(key).strip().lower() for key in raw_custom_config}\n provider_specific_keys = {\n \"vertex_credentials\",\n \"credentials_file\",\n \"vertex_credentials_file\",\n \"google_application_credentials\",\n \"vertex_location\",\n \"location\",\n \"vertex_region\",\n \"region\",\n }\n if normalized_keys & provider_specific_keys:\n return False\n\n normalized_type = str(raw_custom_config.get(\"type\", \"\")).strip().lower()\n if normalized_type not in {\"service_account\", \"external_account\"}:\n return False\n\n # Service account JSON usually includes private_key/client_email, while external\n # account JSON includes credential_source. Either shape should be accepted.\n has_service_account_markers = any(\n key in normalized_keys for key in {\"private_key\", \"client_email\"}\n )\n has_external_account_markers = \"credential_source\" in normalized_keys\n return has_service_account_markers or has_external_account_markers\n\n\ndef _normalize_custom_config(\n provider: str, raw_custom_config: dict[object, object]\n) -> dict[str, str]:\n if provider == \"vertex_ai\" and _looks_like_vertex_credentials_payload(\n raw_custom_config\n ):\n return {\"vertex_credentials\": json.dumps(raw_custom_config)}\n\n normalized: dict[str, str] = {}\n for raw_key, raw_value in raw_custom_config.items():\n key = str(raw_key).strip()\n key_lower = key.lower()\n\n if provider == \"vertex_ai\":\n if key_lower in {\n \"vertex_credentials\",\n \"credentials_file\",\n \"vertex_credentials_file\",\n \"google_application_credentials\",\n }:\n key = \"vertex_credentials\"\n elif key_lower in {\n \"vertex_location\",\n \"location\",\n \"vertex_region\",\n \"region\",\n }:\n key = \"vertex_location\"\n\n normalized[key] = _stringify_custom_config_value(raw_value)\n\n return normalized\n\n\ndef _env_true(env_var: str, default: bool = False) -> bool:\n value = os.environ.get(env_var)\n if value is None:\n return default\n return value.strip().lower() in {\"1\", \"true\", \"yes\", \"on\"}\n\n\ndef _parse_models_env(env_var: str) -> list[str]:\n raw_value = os.environ.get(env_var, \"\").strip()\n if not raw_value:\n return []\n\n try:\n parsed_json = json.loads(raw_value)\n except json.JSONDecodeError:\n parsed_json = None\n\n if isinstance(parsed_json, list):\n return [str(model).strip() for model in parsed_json if str(model).strip()]\n\n return [part.strip() for part in raw_value.split(\",\") if part.strip()]\n\n\ndef _load_provider_config() -> NightlyProviderConfig:\n provider = os.environ.get(_ENV_PROVIDER, \"\").strip().lower()\n model_names = _parse_models_env(_ENV_MODELS)\n api_key = os.environ.get(_ENV_API_KEY) or None\n api_base = os.environ.get(_ENV_API_BASE) or None\n api_version = os.environ.get(_ENV_API_VERSION) or None\n deployment_name = os.environ.get(_ENV_DEPLOYMENT_NAME) or None\n strict = _env_true(_ENV_STRICT, default=False)\n\n custom_config: dict[str, str] | None = None\n custom_config_json = os.environ.get(_ENV_CUSTOM_CONFIG_JSON, \"\").strip()\n if custom_config_json:\n parsed = json.loads(custom_config_json)\n if not isinstance(parsed, dict):\n raise ValueError(f\"{_ENV_CUSTOM_CONFIG_JSON} must be a JSON object\")\n custom_config = _normalize_custom_config(\n provider=provider, raw_custom_config=parsed\n )\n\n if provider == \"ollama_chat\" and api_key and not custom_config:\n custom_config = {\"OLLAMA_API_KEY\": api_key}\n\n return NightlyProviderConfig(\n provider=provider,\n model_names=model_names,\n api_key=api_key,\n api_base=api_base,\n api_version=api_version,\n deployment_name=deployment_name,\n custom_config=custom_config,\n strict=strict,\n )\n\n\ndef _skip_or_fail(strict: bool, message: str) -> None:\n if strict:\n pytest.fail(message)\n pytest.skip(message)\n\n\ndef _validate_provider_config(config: NightlyProviderConfig) -> None:\n if not config.provider:\n _skip_or_fail(strict=config.strict, message=f\"{_ENV_PROVIDER} must be set\")\n\n if not config.model_names:\n _skip_or_fail(\n strict=config.strict,\n message=f\"{_ENV_MODELS} must include at least one model\",\n )\n\n if config.provider != \"ollama_chat\" and not (\n config.api_key or config.custom_config\n ):\n _skip_or_fail(\n strict=config.strict,\n message=(\n f\"{_ENV_API_KEY} or {_ENV_CUSTOM_CONFIG_JSON} is required for provider '{config.provider}'\"\n ),\n )\n\n if config.provider == \"ollama_chat\" and not (\n config.api_base or _default_api_base_for_provider(config.provider)\n ):\n _skip_or_fail(\n strict=config.strict,\n message=(f\"{_ENV_API_BASE} is required for provider '{config.provider}'\"),\n )\n\n if config.provider == \"azure\":\n if not config.api_base:\n _skip_or_fail(\n strict=config.strict,\n message=(\n f\"{_ENV_API_BASE} is required for provider '{config.provider}'\"\n ),\n )\n if not config.api_version:\n _skip_or_fail(\n strict=config.strict,\n message=(\n f\"{_ENV_API_VERSION} is required for provider '{config.provider}'\"\n ),\n )\n\n if config.provider == \"vertex_ai\":\n has_vertex_credentials = bool(\n config.custom_config and config.custom_config.get(\"vertex_credentials\")\n )\n if not has_vertex_credentials:\n configured_keys = (\n sorted(config.custom_config.keys()) if config.custom_config else []\n )\n _skip_or_fail(\n strict=config.strict,\n message=(\n f\"{_ENV_CUSTOM_CONFIG_JSON} must include 'vertex_credentials' \"\n f\"for provider '{config.provider}'. \"\n f\"Found keys: {configured_keys}\"\n ),\n )\n\n\ndef _assert_integration_mode_enabled() -> None:\n assert (\n app_configs.INTEGRATION_TESTS_MODE is True\n ), \"Integration tests require INTEGRATION_TESTS_MODE=true.\"\n\n\ndef _seed_connector_for_search_tool(admin_user: DATestUser) -> None:\n # SearchTool is only exposed when at least one non-default connector exists.\n CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n\ndef _get_internal_search_tool_id(admin_user: DATestUser) -> int:\n tools = ToolManager.list_tools(user_performing_action=admin_user)\n for tool in tools:\n if tool.in_code_tool_id == SEARCH_TOOL_ID:\n return tool.id\n raise AssertionError(\"SearchTool must exist for this test\")\n\n\ndef _default_api_base_for_provider(provider: str) -> str | None:\n if provider == \"openrouter\":\n return \"https://openrouter.ai/api/v1\"\n if provider == \"ollama_chat\":\n # host.docker.internal works when tests are running inside the integration test container.\n return \"http://host.docker.internal:11434\"\n return None\n\n\ndef _create_provider_payload(\n provider: str,\n provider_name: str,\n model_name: str,\n api_key: str | None,\n api_base: str | None,\n api_version: str | None,\n deployment_name: str | None,\n custom_config: dict[str, str] | None,\n) -> dict:\n return {\n \"name\": provider_name,\n \"provider\": provider,\n \"model\": model_name,\n \"api_key\": api_key,\n \"api_base\": api_base,\n \"api_version\": api_version,\n \"deployment_name\": deployment_name,\n \"custom_config\": custom_config,\n \"default_model_name\": model_name,\n \"is_public\": True,\n \"groups\": [],\n \"personas\": [],\n \"model_configurations\": [{\"name\": model_name, \"is_visible\": True}],\n \"api_key_changed\": bool(api_key),\n \"custom_config_changed\": bool(custom_config),\n }\n\n\ndef _ensure_provider_is_default(\n provider_id: int, model_name: str, admin_user: DATestUser\n) -> None:\n list_response = requests.get(\n f\"{API_SERVER_URL}/admin/llm/provider\",\n headers=admin_user.headers,\n )\n list_response.raise_for_status()\n default_text = list_response.json().get(\"default_text\")\n assert default_text is not None, \"Expected a default provider after setting default\"\n assert (\n default_text.get(\"provider_id\") == provider_id\n ), f\"Expected provider {provider_id} to be default, found {default_text.get('provider_id')}\"\n assert (\n default_text.get(\"model_name\") == model_name\n ), f\"Expected default model {model_name}, found {default_text.get('model_name')}\"\n\n\ndef _run_chat_assertions(\n admin_user: DATestUser,\n search_tool_id: int,\n provider: str,\n model_name: str,\n) -> None:\n last_error: str | None = None\n # Retry once to reduce transient nightly flakes due provider-side blips.\n for attempt in range(1, 3):\n chat_session = ChatSessionManager.create(user_performing_action=admin_user)\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=(\n \"Use internal_search to search for 'nightly-provider-regression-sentinel', \"\n \"then summarize the result in one short sentence.\"\n ),\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n )\n\n if response.error is None:\n used_internal_search = any(\n used_tool.tool_name == ToolName.INTERNAL_SEARCH\n for used_tool in response.used_tools\n )\n debug_has_internal_search = any(\n debug_tool_call.tool_name == \"internal_search\"\n for debug_tool_call in response.tool_call_debug\n )\n has_answer = bool(response.full_message.strip())\n\n if used_internal_search and debug_has_internal_search and has_answer:\n return\n\n last_error = (\n f\"attempt={attempt} provider={provider} model={model_name} \"\n f\"used_internal_search={used_internal_search} \"\n f\"debug_internal_search={debug_has_internal_search} \"\n f\"has_answer={has_answer} \"\n f\"tool_call_debug={response.tool_call_debug}\"\n )\n else:\n last_error = f\"attempt={attempt} provider={provider} model={model_name} stream_error={response.error.error}\"\n\n time.sleep(attempt)\n\n pytest.fail(f\"Chat/tool-call assertions failed: {last_error}\")\n\n\ndef _create_and_test_provider_for_model(\n admin_user: DATestUser,\n config: NightlyProviderConfig,\n model_name: str,\n search_tool_id: int,\n) -> None:\n provider_name = f\"nightly-{config.provider}-{uuid4().hex[:12]}\"\n resolved_api_base = config.api_base or _default_api_base_for_provider(\n config.provider\n )\n\n provider_payload = _create_provider_payload(\n provider=config.provider,\n provider_name=provider_name,\n model_name=model_name,\n api_key=config.api_key,\n api_base=resolved_api_base,\n api_version=config.api_version,\n deployment_name=config.deployment_name,\n custom_config=config.custom_config,\n )\n\n test_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/test\",\n headers=admin_user.headers,\n json=provider_payload,\n )\n assert test_response.status_code == 200, (\n f\"Provider test endpoint failed for provider={config.provider} \"\n f\"model={model_name}: {test_response.status_code} {test_response.text}\"\n )\n\n create_response = requests.put(\n f\"{API_SERVER_URL}/admin/llm/provider?is_creation=true\",\n headers=admin_user.headers,\n json=provider_payload,\n )\n assert create_response.status_code == 200, (\n f\"Provider creation failed for provider={config.provider} \"\n f\"model={model_name}: {create_response.status_code} {create_response.text}\"\n )\n provider_id = create_response.json()[\"id\"]\n\n try:\n set_default_response = requests.post(\n f\"{API_SERVER_URL}/admin/llm/default\",\n headers=admin_user.headers,\n json={\"provider_id\": provider_id, \"model_name\": model_name},\n )\n assert set_default_response.status_code == 200, (\n f\"Setting default provider failed for provider={config.provider} \"\n f\"model={model_name}: {set_default_response.status_code} \"\n f\"{set_default_response.text}\"\n )\n\n _ensure_provider_is_default(\n provider_id=provider_id, model_name=model_name, admin_user=admin_user\n )\n _run_chat_assertions(\n admin_user=admin_user,\n search_tool_id=search_tool_id,\n provider=config.provider,\n model_name=model_name,\n )\n finally:\n requests.delete(\n f\"{API_SERVER_URL}/admin/llm/provider/{provider_id}\",\n headers=admin_user.headers,\n )\n\n\ndef test_nightly_provider_chat_workflow(admin_user: DATestUser) -> None:\n \"\"\"Nightly regression test for provider setup + default selection + chat tool calls.\"\"\"\n _assert_integration_mode_enabled()\n config = _load_provider_config()\n _validate_provider_config(config)\n\n _seed_connector_for_search_tool(admin_user)\n search_tool_id = _get_internal_search_tool_id(admin_user)\n\n failures: list[str] = []\n for model_name in config.model_names:\n try:\n _create_and_test_provider_for_model(\n admin_user=admin_user,\n config=config,\n model_name=model_name,\n search_tool_id=search_tool_id,\n )\n except BaseException as exc:\n if isinstance(exc, (KeyboardInterrupt, SystemExit)):\n raise\n failures.append(\n f\"provider={config.provider} model={model_name} error={type(exc).__name__}: {exc}\"\n )\n\n if failures:\n pytest.fail(\"Nightly provider chat failures:\\n\" + \"\\n\".join(failures))\n" }, { "path": "backend/tests/integration/tests/llm_workflows/test_tool_policy_enforcement.py", "content": "from onyx.configs import app_configs\nfrom onyx.configs.constants import DocumentSource\nfrom onyx.tools.constants import SEARCH_TOOL_ID\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.tool import ToolManager\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.test_models import ToolName\n\n\n_DUMMY_OPENAI_API_KEY = \"sk-mock-tool-policy-tests\"\n\n\ndef _assert_integration_mode_enabled() -> None:\n assert (\n app_configs.INTEGRATION_TESTS_MODE is True\n ), \"Integration tests require INTEGRATION_TESTS_MODE=true.\"\n\n\ndef _seed_connector_for_search_tool(admin_user: DATestUser) -> None:\n # SearchTool is only exposed when at least one non-default connector exists.\n CCPairManager.create_from_scratch(\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n\ndef _get_internal_search_tool_id(admin_user: DATestUser) -> int:\n tools = ToolManager.list_tools(user_performing_action=admin_user)\n for tool in tools:\n if tool.in_code_tool_id == SEARCH_TOOL_ID:\n return tool.id\n raise AssertionError(\"SearchTool must exist for this test\")\n\n\ndef _ensure_llm_provider(admin_user: DATestUser) -> None:\n LLMProviderManager.create(\n user_performing_action=admin_user,\n api_key=_DUMMY_OPENAI_API_KEY,\n )\n\n\ndef test_forced_tool_executes_when_available(admin_user: DATestUser) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n _ensure_llm_provider(admin_user)\n\n search_tool_id = _get_internal_search_tool_id(admin_user)\n persona = PersonaManager.create(\n tool_ids=[search_tool_id], user_performing_action=admin_user\n )\n chat_session = ChatSessionManager.create(\n persona_id=persona.id, user_performing_action=admin_user\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"force the search tool\",\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n mock_llm_response='{\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"alpha\"]}}',\n )\n\n assert response.error is None, f\"Unexpected stream error: {response.error}\"\n assert any(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response.used_tools\n )\n assert len(response.tool_call_debug) == 1\n assert response.tool_call_debug[0].tool_name == \"internal_search\"\n assert response.tool_call_debug[0].tool_args == {\"queries\": [\"alpha\"]}\n\n\ndef test_forced_tool_rejected_when_not_in_persona_tools(\n admin_user: DATestUser,\n) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n _ensure_llm_provider(admin_user)\n\n search_tool_id = _get_internal_search_tool_id(admin_user)\n persona = PersonaManager.create(tool_ids=[], user_performing_action=admin_user)\n chat_session = ChatSessionManager.create(\n persona_id=persona.id, user_performing_action=admin_user\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"try forcing a missing tool\",\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n )\n\n assert response.error is not None\n assert response.error.error == f\"Forced tool {search_tool_id} not found in tools\"\n assert response.used_tools == []\n\n\ndef test_allowed_tool_ids_excludes_tools_outside_allowlist(\n admin_user: DATestUser,\n) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n _ensure_llm_provider(admin_user)\n\n search_tool_id = _get_internal_search_tool_id(admin_user)\n persona = PersonaManager.create(\n tool_ids=[search_tool_id], user_performing_action=admin_user\n )\n chat_session = ChatSessionManager.create(\n persona_id=persona.id, user_performing_action=admin_user\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"attempt tool use with empty allowlist\",\n user_performing_action=admin_user,\n allowed_tool_ids=[],\n mock_llm_response='{\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"beta\"]}}',\n )\n\n assert response.error is None, f\"Unexpected stream error: {response.error}\"\n assert response.used_tools == []\n assert response.tool_call_debug == []\n\n\ndef test_forced_and_allowlist_conflict_returns_validation_error(\n admin_user: DATestUser,\n) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n _ensure_llm_provider(admin_user)\n\n search_tool_id = _get_internal_search_tool_id(admin_user)\n persona = PersonaManager.create(\n tool_ids=[search_tool_id], user_performing_action=admin_user\n )\n chat_session = ChatSessionManager.create(\n persona_id=persona.id, user_performing_action=admin_user\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"force a tool blocked by allowlist\",\n user_performing_action=admin_user,\n allowed_tool_ids=[],\n forced_tool_ids=[search_tool_id],\n )\n\n assert response.error is not None\n assert response.error.error == f\"Forced tool {search_tool_id} not found in tools\"\n assert response.used_tools == []\n\n\ndef test_run_search_always_maps_to_forced_search_tool(admin_user: DATestUser) -> None:\n _assert_integration_mode_enabled()\n _seed_connector_for_search_tool(admin_user)\n _ensure_llm_provider(admin_user)\n\n search_tool_id = _get_internal_search_tool_id(admin_user)\n persona = PersonaManager.create(\n tool_ids=[search_tool_id], user_performing_action=admin_user\n )\n chat_session = ChatSessionManager.create(\n persona_id=persona.id, user_performing_action=admin_user\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"always run search\",\n user_performing_action=admin_user,\n forced_tool_ids=[search_tool_id],\n mock_llm_response='{\"name\":\"internal_search\",\"arguments\":{\"queries\":[\"gamma\"]}}',\n )\n\n assert response.error is None, f\"Unexpected stream error: {response.error}\"\n assert any(\n tool.tool_name == ToolName.INTERNAL_SEARCH for tool in response.used_tools\n )\n assert len(response.tool_call_debug) == 1\n assert response.tool_call_debug[0].tool_name == \"internal_search\"\n assert response.tool_call_debug[0].tool_args == {\"queries\": [\"gamma\"]}\n" }, { "path": "backend/tests/integration/tests/mcp/test_mcp_client_no_auth_flow.py", "content": "import os\nimport socket\nimport subprocess\nimport sys\nimport time\nfrom collections.abc import Generator\nfrom pathlib import Path\n\nimport pytest\nimport requests\n\nfrom onyx.db.enums import MCPAuthenticationPerformer\nfrom onyx.db.enums import MCPAuthenticationType\nfrom onyx.db.enums import MCPTransport\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n# TODO: update mcp client tests to use constants in common_utils/constants.py\n# NOTE: the tests for client should be independent of the Onyx MCP server\n# This means the port should probably stay to be 8010/not 8090 the Onyx MCP server port\n# Use MOCK_MCP_SERVER_PORT to avoid conflicts with the real Onyx MCP server port (8090)\nMCP_SERVER_HOST = os.getenv(\"TEST_WEB_HOSTNAME\", \"127.0.0.1\")\nMCP_SERVER_PORT = int(os.getenv(\"MOCK_MCP_SERVER_PORT\", \"8010\"))\nMCP_SERVER_URL = f\"http://{MCP_SERVER_HOST}:{MCP_SERVER_PORT}/mcp\"\nMCP_HELLO_TOOL = \"hello\"\n\nMCP_SERVER_SCRIPT = (\n Path(__file__).resolve().parents[2]\n / \"mock_services\"\n / \"mcp_test_server\"\n / \"run_mcp_server_no_auth.py\"\n)\n\n\ndef _wait_for_port(\n host: str,\n port: int,\n process: subprocess.Popen[bytes],\n timeout_seconds: float = 10.0,\n) -> None:\n start = time.monotonic()\n while time.monotonic() - start < timeout_seconds:\n if process.poll() is not None:\n raise RuntimeError(\"MCP server process exited unexpectedly during startup\")\n\n with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:\n sock.settimeout(0.5)\n try:\n sock.connect((host, port))\n return\n except OSError:\n time.sleep(0.1)\n\n raise TimeoutError(\"Timed out waiting for MCP server to accept connections\")\n\n\n@pytest.fixture(scope=\"module\")\ndef mcp_no_auth_server() -> Generator[None, None, None]:\n process = subprocess.Popen(\n [sys.executable, str(MCP_SERVER_SCRIPT), str(MCP_SERVER_PORT)],\n cwd=MCP_SERVER_SCRIPT.parent,\n )\n\n try:\n _wait_for_port(MCP_SERVER_HOST, MCP_SERVER_PORT, process)\n yield\n finally:\n process.terminate()\n try:\n process.wait(timeout=5)\n except subprocess.TimeoutExpired:\n process.kill()\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef ensure_mcp_server_exists() -> None:\n if not MCP_SERVER_SCRIPT.exists():\n raise FileNotFoundError(\n f\"Expected MCP server script at {MCP_SERVER_SCRIPT}, but it was not found\"\n )\n\n\ndef test_mcp_client_no_auth_flow(\n mcp_no_auth_server: None, # noqa: ARG001\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n # Step a) Create a no-auth MCP server via the admin API\n create_response = requests.post(\n f\"{API_SERVER_URL}/admin/mcp/servers/create\",\n json={\n \"name\": \"integration-mcp-no-auth\",\n \"description\": \"Integration test MCP server\",\n \"server_url\": MCP_SERVER_URL,\n \"transport\": MCPTransport.STREAMABLE_HTTP.value,\n \"auth_type\": MCPAuthenticationType.NONE.value,\n \"auth_performer\": MCPAuthenticationPerformer.ADMIN.value,\n },\n headers=admin_user.headers,\n cookies=admin_user.cookies,\n )\n create_response.raise_for_status()\n server_id = create_response.json()[\"server_id\"]\n\n # Step b) list the server's tools\n tools_response = requests.get(\n f\"{API_SERVER_URL}/admin/mcp/server/{server_id}/tools\",\n headers=admin_user.headers,\n cookies=admin_user.cookies,\n )\n tools_response.raise_for_status()\n tool_entries = tools_response.json()[\"tools\"]\n assert len(tool_entries) == 101\n\n # Update server status to CONNECTED\n status_response = requests.patch(\n f\"{API_SERVER_URL}/admin/mcp/server/{server_id}/status\",\n params={\"status\": \"CONNECTED\"},\n headers=admin_user.headers,\n cookies=admin_user.cookies,\n )\n status_response.raise_for_status()\n\n tools_response = requests.get(\n f\"{API_SERVER_URL}/admin/mcp/server/{server_id}/db-tools\",\n headers=admin_user.headers,\n cookies=admin_user.cookies,\n )\n tools_response.raise_for_status()\n tool_entries = tools_response.json()[\"tools\"]\n hello_tool_entry = next(\n tool for tool in tool_entries if tool[\"name\"] == MCP_HELLO_TOOL\n )\n tool_id = hello_tool_entry[\"id\"]\n\n # Step c) Create an assistant (persona) with the MCP tool attached\n persona = PersonaManager.create(\n name=\"integration-mcp-persona\",\n description=\"Persona for MCP integration test\",\n tool_ids=[tool_id],\n user_performing_action=admin_user,\n )\n persona_tools_response = requests.get(\n f\"{API_SERVER_URL}/persona\",\n headers=basic_user.headers,\n cookies=basic_user.cookies,\n )\n persona_tools_response.raise_for_status()\n persona_entries = persona_tools_response.json()\n persona_entry = next(\n entry for entry in persona_entries if entry[\"id\"] == persona.id\n )\n persona_tool_ids = {tool[\"id\"] for tool in persona_entry[\"tools\"]}\n assert tool_id in persona_tool_ids\n" }, { "path": "backend/tests/integration/tests/mcp/test_mcp_server_auth.py", "content": "\"\"\"Integration tests for MCP Server auth delegated to API /me.\"\"\"\n\nimport requests\n\nfrom tests.integration.common_utils.constants import MCP_SERVER_URL\nfrom tests.integration.common_utils.managers.pat import PATManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nSTREAMABLE_HTTP_URL = f\"{MCP_SERVER_URL.rstrip('/')}/?transportType=streamable-http\"\n\n\ndef test_mcp_server_health_check(reset: None) -> None: # noqa: ARG001\n \"\"\"Test MCP server health check endpoint.\"\"\"\n response = requests.get(f\"{MCP_SERVER_URL}/health\", timeout=10)\n assert response.status_code == 200\n assert response.json()[\"status\"] == \"healthy\"\n assert response.json()[\"service\"] == \"mcp_server\"\n\n\ndef test_mcp_server_auth_missing_token(reset: None) -> None: # noqa: ARG001\n \"\"\"Test MCP server rejects requests without credentials.\"\"\"\n response = requests.post(STREAMABLE_HTTP_URL)\n assert response.status_code == 401\n\n\ndef test_mcp_server_auth_invalid_token(reset: None) -> None: # noqa: ARG001\n \"\"\"Test MCP server rejects requests with an invalid bearer token.\"\"\"\n response = requests.post(\n STREAMABLE_HTTP_URL,\n headers={\"Authorization\": \"Bearer invalid-token\"},\n json={\"jsonrpc\": \"2.0\", \"method\": \"initialize\", \"id\": 1},\n )\n assert response.status_code == 401\n\n\ndef test_mcp_server_auth_valid_token(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test MCP server accepts requests with a valid bearer token.\"\"\"\n pat = PATManager.create(\n name=\"Test MCP Token\",\n expiration_days=7,\n user_performing_action=admin_user,\n )\n access_token = pat.token\n\n # Test connection with MCP protocol request\n response = requests.post(\n STREAMABLE_HTTP_URL,\n headers={\n \"Authorization\": f\"Bearer {access_token}\",\n \"Content-Type\": \"application/json\",\n \"Accept\": \"application/json\",\n \"MCP-Protocol-Version\": \"2025-03-26\",\n },\n json={\"jsonrpc\": \"2.0\", \"method\": \"initialize\", \"id\": 1},\n )\n\n # Should be authenticated (may return MCP protocol response or error)\n # 200 = valid MCP protocol response\n # 400 = valid protocol error (authenticated but bad request)\n assert response.status_code in [200, 400]\n" }, { "path": "backend/tests/integration/tests/mcp/test_mcp_server_search.py", "content": "\"\"\"Integration tests covering MCP document search flows.\"\"\"\n\nfrom __future__ import annotations\n\nimport asyncio\nimport json\nimport os\nfrom collections.abc import Awaitable\nfrom collections.abc import Callable\nfrom datetime import datetime\nfrom datetime import timezone\nfrom typing import Any\n\nimport pytest\nfrom mcp import ClientSession\nfrom mcp.client.streamable_http import streamablehttp_client\nfrom mcp.types import CallToolResult\nfrom mcp.types import TextContent\n\nfrom onyx.db.enums import AccessType\nfrom tests.integration.common_utils.constants import MCP_SERVER_URL\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document import DocumentManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.pat import PATManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestAPIKey\nfrom tests.integration.common_utils.test_models import DATestCCPair\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n# Constants\nMCP_SEARCH_TOOL = \"search_indexed_documents\"\nINDEXED_SOURCES_RESOURCE_URI = \"resource://indexed_sources\"\nDEFAULT_SEARCH_LIMIT = 5\nSTREAMABLE_HTTP_URL = f\"{MCP_SERVER_URL.rstrip('/')}/?transportType=streamable-http\"\n\n\ndef _run_with_mcp_session(\n headers: dict[str, str],\n action: Callable[[ClientSession], Awaitable[Any]],\n) -> Any:\n \"\"\"Run an async action with an MCP client session.\"\"\"\n\n async def _runner() -> Any:\n async with streamablehttp_client(STREAMABLE_HTTP_URL, headers=headers) as (\n read,\n write,\n _,\n ):\n async with ClientSession(read, write) as session:\n return await action(session)\n\n return asyncio.run(_runner())\n\n\ndef _extract_tool_payload(result: CallToolResult) -> dict[str, Any]:\n \"\"\"Extract JSON payload from MCP tool result.\"\"\"\n if result.isError:\n raise AssertionError(f\"MCP tool returned error: {result}\")\n\n text_blocks = [\n block.text\n for block in result.content\n if isinstance(block, TextContent) and block.text\n ]\n if not text_blocks:\n raise AssertionError(\"Expected textual content from MCP tool result\")\n\n return json.loads(text_blocks[-1])\n\n\ndef _call_search_tool(\n headers: dict[str, str], query: str, limit: int = DEFAULT_SEARCH_LIMIT\n) -> CallToolResult:\n \"\"\"Call the search_indexed_documents tool via MCP.\"\"\"\n\n async def _action(session: ClientSession) -> CallToolResult:\n await session.initialize()\n return await session.call_tool(\n MCP_SEARCH_TOOL,\n {\n \"query\": query,\n \"limit\": limit,\n },\n )\n\n return _run_with_mcp_session(headers, _action)\n\n\ndef _auth_headers(user: DATestUser, name: str) -> dict[str, str]:\n \"\"\"Create authorization headers with a PAT token.\"\"\"\n pat = PATManager.create(\n name=name,\n expiration_days=7,\n user_performing_action=user,\n )\n return {\"Authorization\": f\"Bearer {pat.token}\"}\n\n\ndef _seed_document_and_wait_for_indexing(\n cc_pair: DATestCCPair,\n content: str,\n api_key: DATestAPIKey,\n user_performing_action: DATestUser,\n) -> None:\n \"\"\"Seed a document and wait for indexing to complete.\"\"\"\n before = datetime.now(timezone.utc)\n DocumentManager.seed_doc_with_content(\n cc_pair=cc_pair,\n content=content,\n api_key=api_key,\n )\n CCPairManager.wait_for_indexing_completion(\n cc_pair=cc_pair,\n after=before,\n user_performing_action=user_performing_action,\n )\n\n\ndef test_mcp_document_search_flow(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test the complete MCP search flow: initialization, resources, tools, and search.\"\"\"\n # LLM provider is required for the document-search endpoint\n LLMProviderManager.create(user_performing_action=admin_user)\n\n api_key = APIKeyManager.create(user_performing_action=admin_user)\n cc_pair = CCPairManager.create_from_scratch(user_performing_action=admin_user)\n\n doc_text = \"MCP happy path search document\"\n _seed_document_and_wait_for_indexing(\n cc_pair=cc_pair,\n content=doc_text,\n api_key=api_key,\n user_performing_action=admin_user,\n )\n\n headers = _auth_headers(admin_user, name=\"mcp-search-flow\")\n\n async def _full_flow(session: ClientSession) -> Any:\n await session.initialize()\n resources = await session.list_resources()\n tools = await session.list_tools()\n search_result = await session.call_tool(\n MCP_SEARCH_TOOL,\n {\n \"query\": doc_text,\n \"limit\": DEFAULT_SEARCH_LIMIT,\n },\n )\n return resources, tools, search_result\n\n resources_result, tools_result, search_result = _run_with_mcp_session(\n headers, _full_flow\n )\n\n # Verify resources are available\n resource_uris = {str(resource.uri) for resource in resources_result.resources}\n assert INDEXED_SOURCES_RESOURCE_URI in resource_uris\n\n # Verify tools are available\n tool_names = {tool.name for tool in tools_result.tools}\n assert MCP_SEARCH_TOOL in tool_names\n\n # Verify search results\n payload = _extract_tool_payload(search_result)\n assert payload[\"query\"] == doc_text\n assert payload[\"total_results\"] >= 1\n assert isinstance(payload[\"documents\"], list)\n assert len(payload[\"documents\"]) > 0\n assert any(doc_text in (doc.get(\"content\") or \"\") for doc in payload[\"documents\"])\n\n # Verify document structure\n for doc in payload[\"documents\"]:\n assert isinstance(doc, dict)\n # Verify expected fields exist (may be None)\n assert \"content\" in doc\n assert \"semantic_identifier\" in doc\n assert \"source_type\" in doc\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"User group permissions are Enterprise-only\",\n)\ndef test_mcp_search_respects_acl_filters(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that search respects ACL filters - privileged users can access, others cannot.\"\"\"\n # LLM provider is required for the document-search endpoint\n LLMProviderManager.create(user_performing_action=admin_user)\n\n user_without_access = UserManager.create(name=\"mcp-acl-user-a\")\n privileged_user = UserManager.create(name=\"mcp-acl-user-b\")\n\n api_key = APIKeyManager.create(user_performing_action=admin_user)\n restricted_cc_pair = CCPairManager.create_from_scratch(\n access_type=AccessType.PRIVATE,\n user_performing_action=admin_user,\n )\n\n user_group = UserGroupManager.create(\n user_ids=[privileged_user.id],\n cc_pair_ids=[restricted_cc_pair.id],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_performing_action=admin_user, user_groups_to_check=[user_group]\n )\n\n restricted_doc_content = \"MCP restricted knowledge base document\"\n _seed_document_and_wait_for_indexing(\n cc_pair=restricted_cc_pair,\n content=restricted_doc_content,\n api_key=api_key,\n user_performing_action=admin_user,\n )\n\n privileged_headers = _auth_headers(privileged_user, \"mcp-acl-allowed\")\n restricted_headers = _auth_headers(user_without_access, \"mcp-acl-blocked\")\n\n # Privileged user should find the document\n allowed_result = _call_search_tool(privileged_headers, restricted_doc_content)\n allowed_payload = _extract_tool_payload(allowed_result)\n assert allowed_payload[\"total_results\"] >= 1\n assert any(\n restricted_doc_content in (doc.get(\"content\") or \"\")\n for doc in allowed_payload[\"documents\"]\n )\n\n # User without access should not find the document\n blocked_result = _call_search_tool(restricted_headers, restricted_doc_content)\n blocked_payload = _extract_tool_payload(blocked_result)\n assert blocked_payload[\"total_results\"] == 0\n assert blocked_payload[\"documents\"] == []\n" }, { "path": "backend/tests/integration/tests/migrations/conftest.py", "content": "\"\"\"\npytest-alembic configuration for testing Alembic migrations.\n\nThis module provides fixtures required by pytest-alembic to test the main\nschema migrations (alembic). For alembic_tenants, see test_alembic_tenants.py.\n\nUsage:\n Run all built-in pytest-alembic tests:\n pytest tests/integration/tests/migrations/test_alembic_main.py -v\n\nSee: https://pytest-alembic.readthedocs.io/en/latest/\n\"\"\"\n\nfrom collections.abc import Generator\nfrom typing import Any\n\nimport pytest\nfrom sqlalchemy import create_engine\nfrom sqlalchemy import text\nfrom sqlalchemy.engine import Engine\n\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PASSWORD\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.engine.sql_engine import SYNC_DB_API\nfrom shared_configs.configs import POSTGRES_DEFAULT_SCHEMA\n\n\ndef _create_sync_engine() -> Engine:\n \"\"\"Create a synchronous SQLAlchemy engine for pytest-alembic.\"\"\"\n conn_str = build_connection_string(\n db=\"postgres\",\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n db_api=SYNC_DB_API,\n )\n return create_engine(conn_str)\n\n\n@pytest.fixture\ndef alembic_config() -> dict[str, Any]:\n \"\"\"\n Configure pytest-alembic for the main schema migrations.\n\n Returns pytest-alembic configuration options.\n See: https://pytest-alembic.readthedocs.io/en/latest/setup.html\n \"\"\"\n return {\n \"file\": \"alembic.ini\",\n \"script_location\": \"alembic\",\n # Pass additional attributes to the alembic config\n # These will be available in env.py via context.config.attributes\n \"attributes\": {\n \"schema_name\": POSTGRES_DEFAULT_SCHEMA,\n },\n }\n\n\n@pytest.fixture\ndef alembic_engine() -> Generator[Engine, None, None]:\n \"\"\"\n Provide a synchronous SQLAlchemy engine for pytest-alembic.\n\n pytest-alembic requires a synchronous engine to run migrations.\n The engine is configured to connect to the test database.\n\n Note: pytest-alembic will internally perform commits, so ensure\n the database is in an appropriate state before running tests.\n \"\"\"\n engine = _create_sync_engine()\n\n # Ensure the default schema exists\n with engine.connect() as conn:\n conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS \"{POSTGRES_DEFAULT_SCHEMA}\"'))\n conn.commit()\n\n yield engine\n\n engine.dispose()\n" }, { "path": "backend/tests/integration/tests/migrations/test_alembic_main.py", "content": "\"\"\"\npytest-alembic tests for the main schema migrations.\n\nThese tests use pytest-alembic to verify that alembic migrations are correct.\nThe tests cover:\n- Single head revision (no diverged migration history)\n- Upgrade path from base to head\n- Up/down consistency (all downgrades succeed)\n\nUsage:\n pytest tests/integration/tests/migrations/test_alembic_main.py -v\n\nSee: https://github.com/schireson/pytest-alembic\n\"\"\"\n\nfrom pytest_alembic.tests import test_single_head_revision # type: ignore[import-not-found,unused-ignore]\nfrom pytest_alembic.tests import test_up_down_consistency # type: ignore[import-not-found,unused-ignore]\nfrom pytest_alembic.tests import test_upgrade # type: ignore[import-not-found,unused-ignore]\n\n__all__ = [\n \"test_single_head_revision\",\n \"test_up_down_consistency\",\n \"test_upgrade\",\n]\n" }, { "path": "backend/tests/integration/tests/migrations/test_alembic_tenants.py", "content": "\"\"\"\npytest-alembic tests for the tenants/public schema migrations.\n\nThese tests use pytest-alembic to verify that alembic_tenants migrations\nare correct. The alembic_tenants configuration handles migrations for\nthe public schema tables that are shared across tenants.\n\nUsage:\n pytest tests/integration/tests/migrations/test_alembic_tenants.py -v\n\nSee: https://github.com/schireson/pytest-alembic\n\"\"\"\n\nfrom collections.abc import Generator\nfrom typing import Any\n\nimport pytest\nfrom pytest_alembic import create_alembic_fixture # type: ignore[import-not-found,unused-ignore]\nfrom pytest_alembic.tests import test_single_head_revision # type: ignore[import-not-found,unused-ignore]\nfrom pytest_alembic.tests import test_up_down_consistency # type: ignore[import-not-found,unused-ignore]\nfrom pytest_alembic.tests import test_upgrade # type: ignore[import-not-found,unused-ignore]\nfrom sqlalchemy import create_engine\nfrom sqlalchemy.engine import Engine\n\nfrom onyx.configs.app_configs import POSTGRES_HOST\nfrom onyx.configs.app_configs import POSTGRES_PASSWORD\nfrom onyx.configs.app_configs import POSTGRES_PORT\nfrom onyx.configs.app_configs import POSTGRES_USER\nfrom onyx.db.engine.sql_engine import build_connection_string\nfrom onyx.db.engine.sql_engine import SYNC_DB_API\n\n\n@pytest.fixture\ndef alembic_config() -> dict[str, Any]:\n \"\"\"Override alembic_config for tenants configuration.\"\"\"\n return {\n \"file\": \"alembic.ini\",\n \"config_ini_section\": \"schema_private\",\n \"script_location\": \"alembic_tenants\",\n }\n\n\n@pytest.fixture\ndef alembic_engine() -> Generator[Engine, None, None]:\n \"\"\"Override alembic_engine for tenants configuration.\"\"\"\n conn_str = build_connection_string(\n db=\"postgres\",\n user=POSTGRES_USER,\n password=POSTGRES_PASSWORD,\n host=POSTGRES_HOST,\n port=POSTGRES_PORT,\n db_api=SYNC_DB_API,\n )\n engine = create_engine(conn_str)\n yield engine\n engine.dispose()\n\n\n# Create a custom alembic fixture for the tenants configuration\nalembic_runner = create_alembic_fixture()\n\n__all__ = [\n \"test_single_head_revision\",\n \"test_up_down_consistency\",\n \"test_upgrade\",\n]\n" }, { "path": "backend/tests/integration/tests/migrations/test_assistant_consolidation_migration.py", "content": "\"\"\"\nIntegration tests for the assistant consolidation migration.\n\nTests the migration from multiple default assistants (Search, General, Art, etc.)\nto a single default Assistant (ID 0) and the associated tool seeding.\n\"\"\"\n\nfrom sqlalchemy import text\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom tests.integration.common_utils.reset import downgrade_postgres\nfrom tests.integration.common_utils.reset import upgrade_postgres\n\n\ndef test_cold_startup_default_assistant() -> None:\n \"\"\"Test that cold startup creates only the default assistant.\"\"\"\n # Start fresh at the head revision\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"base\", clear_data=True\n )\n upgrade_postgres(database=\"postgres\", config_name=\"alembic\", revision=\"head\")\n\n with get_session_with_current_tenant() as db_session:\n # Check only default assistant exists\n result = db_session.execute(\n text(\n \"\"\"\n SELECT id, name, builtin_persona, is_featured, deleted\n FROM persona\n WHERE builtin_persona = true\n ORDER BY id\n \"\"\"\n )\n )\n assistants = result.fetchall()\n\n # Should have exactly one builtin assistant\n assert len(assistants) == 1, \"Should have exactly one builtin assistant\"\n default = assistants[0]\n assert default[0] == 0, \"Default assistant should have ID 0\"\n assert default[1] == \"Assistant\", \"Should be named 'Assistant'\"\n assert default[2] is True, \"Should be builtin\"\n assert default[3] is True, \"Should be is_featured\"\n assert default[4] is False, \"Should not be deleted\"\n\n # Check tools are properly associated\n result = db_session.execute(\n text(\n \"\"\"\n SELECT t.name, t.display_name\n FROM tool t\n JOIN persona__tool pt ON t.id = pt.tool_id\n WHERE pt.persona_id = 0\n ORDER BY t.name\n \"\"\"\n )\n )\n tool_associations = result.fetchall()\n tool_names = [row[0] for row in tool_associations]\n tool_display_names = [row[1] for row in tool_associations]\n\n # Verify all three main tools are attached\n assert (\n \"internal_search\" in tool_names\n ), \"Default assistant should have SearchTool attached\"\n assert (\n \"generate_image\" in tool_names\n ), \"Default assistant should have ImageGenerationTool attached\"\n assert (\n \"web_search\" in tool_names\n ), \"Default assistant should have WebSearchTool attached\"\n assert (\n \"read_file\" in tool_names\n ), \"Default assistant should have FileReaderTool attached\"\n assert (\n \"python\" in tool_names\n ), \"Default assistant should have PythonTool attached\"\n\n # Also verify by display names for clarity\n assert (\n \"Internal Search\" in tool_display_names\n ), \"Default assistant should have Internal Search tool\"\n assert (\n \"Image Generation\" in tool_display_names\n ), \"Default assistant should have Image Generation tool\"\n assert (\n \"Web Search\" in tool_display_names\n ), \"Default assistant should have Web Search tool\"\n assert (\n \"File Reader\" in tool_display_names\n ), \"Default assistant should have File Reader tool\"\n assert (\n \"Code Interpreter\" in tool_display_names\n ), \"Default assistant should have Code Interpreter tool\"\n\n # Should have exactly 6 tools\n assert (\n len(tool_associations) == 6\n ), f\"Default assistant should have exactly 6 tools attached, got {len(tool_associations)}\"\n" }, { "path": "backend/tests/integration/tests/migrations/test_migrations.py", "content": "# TODO(rkuo): All of the downgrade_postgres and upgrade_postgres operations here\n# are vulnerable to deadlocks. We could deal with them similar to reset_postgres\n# where we retry out of process\n\nimport json\n\nimport pytest\nfrom sqlalchemy import text\n\nfrom onyx.configs.constants import ANONYMOUS_USER_UUID\nfrom onyx.configs.constants import DEFAULT_BOOST\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom tests.integration.common_utils.reset import downgrade_postgres\nfrom tests.integration.common_utils.reset import upgrade_postgres\n\n\n@pytest.mark.skip(\n reason=\"Migration test no longer needed - migration has been applied to production\"\n)\ndef test_fix_capitalization_migration() -> None:\n \"\"\"Test that the be2ab2aa50ee migration correctly lowercases external_user_group_ids\"\"\"\n # Reset the database and run migrations up to the second to last migration\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"base\", clear_data=True\n )\n upgrade_postgres(\n database=\"postgres\",\n config_name=\"alembic\",\n # Upgrade it to the migration before the fix\n revision=\"369644546676\",\n )\n\n # Insert test data with mixed case group IDs\n test_data = [\n {\n \"id\": \"test_doc_1\",\n \"external_user_group_ids\": [\"Group1\", \"GROUP2\", \"group3\"],\n \"semantic_id\": \"test_doc_1\",\n \"boost\": DEFAULT_BOOST,\n \"hidden\": False,\n \"from_ingestion_api\": False,\n \"last_modified\": \"NOW()\",\n },\n {\n \"id\": \"test_doc_2\",\n \"external_user_group_ids\": [\"UPPER1\", \"upper2\", \"UPPER3\"],\n \"semantic_id\": \"test_doc_2\",\n \"boost\": DEFAULT_BOOST,\n \"hidden\": False,\n \"from_ingestion_api\": False,\n \"last_modified\": \"NOW()\",\n },\n ]\n\n # Insert the test data\n with get_session_with_current_tenant() as db_session:\n for doc in test_data:\n db_session.execute(\n text(\n \"\"\"\n INSERT INTO document (\n id,\n external_user_group_ids,\n semantic_id,\n boost,\n hidden,\n from_ingestion_api,\n last_modified\n )\n VALUES (\n :id,\n :group_ids,\n :semantic_id,\n :boost,\n :hidden,\n :from_ingestion_api,\n :last_modified\n )\n \"\"\"\n ),\n {\n \"id\": doc[\"id\"],\n \"group_ids\": doc[\"external_user_group_ids\"],\n \"semantic_id\": doc[\"semantic_id\"],\n \"boost\": doc[\"boost\"],\n \"hidden\": doc[\"hidden\"],\n \"from_ingestion_api\": doc[\"from_ingestion_api\"],\n \"last_modified\": doc[\"last_modified\"],\n },\n )\n db_session.commit()\n\n # Verify the data was inserted correctly\n with get_session_with_current_tenant() as db_session:\n results = db_session.execute(\n text(\n \"\"\"\n SELECT id, external_user_group_ids\n FROM document\n WHERE id IN ('test_doc_1', 'test_doc_2')\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n # Verify initial state\n assert len(results) == 2\n assert results[0].external_user_group_ids == [\"Group1\", \"GROUP2\", \"group3\"]\n assert results[1].external_user_group_ids == [\"UPPER1\", \"upper2\", \"UPPER3\"]\n\n # Run migrations again to apply the fix\n upgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"be2ab2aa50ee\"\n )\n\n # Verify the fix was applied\n with get_session_with_current_tenant() as db_session:\n results = db_session.execute(\n text(\n \"\"\"\n SELECT id, external_user_group_ids\n FROM document\n WHERE id IN ('test_doc_1', 'test_doc_2')\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n # Verify all group IDs are lowercase\n assert len(results) == 2\n assert results[0].external_user_group_ids == [\"group1\", \"group2\", \"group3\"]\n assert results[1].external_user_group_ids == [\"upper1\", \"upper2\", \"upper3\"]\n\n\ndef test_jira_connector_migration() -> None:\n \"\"\"Test that the da42808081e3 migration correctly updates Jira connector configurations\"\"\"\n # Reset the database and run migrations up to the migration before the Jira connector change\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"base\", clear_data=True\n )\n upgrade_postgres(\n database=\"postgres\",\n config_name=\"alembic\",\n # Upgrade it to the migration before the Jira connector change\n revision=\"f13db29f3101\",\n )\n\n # Insert test data with various Jira connector configurations\n test_data = [\n {\n \"id\": 1,\n \"name\": \"jira_connector_1\",\n \"source\": \"JIRA\",\n \"connector_specific_config\": {\n \"jira_project_url\": \"https://example.atlassian.net/projects/PROJ\",\n \"comment_email_blacklist\": [\"test@example.com\"],\n \"batch_size\": 100,\n \"labels_to_skip\": [\"skip-me\"],\n },\n },\n {\n \"id\": 2,\n \"name\": \"jira_connector_2\",\n \"source\": \"JIRA\",\n \"connector_specific_config\": {\n \"jira_project_url\": \"https://other.atlassian.net/projects/OTHER\"\n },\n },\n {\n \"id\": 3,\n \"name\": \"jira_connector_3\",\n \"source\": \"JIRA\",\n \"connector_specific_config\": {\n \"jira_project_url\": \"https://example.atlassian.net/projects/TEST\",\n \"batch_size\": 50,\n },\n },\n ]\n\n # Insert the test data\n with get_session_with_current_tenant() as db_session:\n for connector in test_data:\n db_session.execute(\n text(\n \"\"\"\n INSERT INTO connector (\n id,\n name,\n source,\n connector_specific_config\n )\n VALUES (\n :id,\n :name,\n :source,\n :config\n )\n \"\"\"\n ),\n {\n \"id\": connector[\"id\"],\n \"name\": connector[\"name\"],\n \"source\": connector[\"source\"],\n \"config\": json.dumps(connector[\"connector_specific_config\"]),\n },\n )\n db_session.commit()\n\n # Verify the data was inserted correctly\n with get_session_with_current_tenant() as db_session:\n results = db_session.execute(\n text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'JIRA'\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n # Verify initial state\n assert len(results) == 3\n assert (\n results[0].connector_specific_config\n == test_data[0][\"connector_specific_config\"]\n )\n assert (\n results[1].connector_specific_config\n == test_data[1][\"connector_specific_config\"]\n )\n assert (\n results[2].connector_specific_config\n == test_data[2][\"connector_specific_config\"]\n )\n\n # Run migrations again to apply the Jira connector change\n upgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"da42808081e3\"\n )\n # Verify the upgrade was applied correctly\n with get_session_with_current_tenant() as db_session:\n results = db_session.execute(\n text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'JIRA'\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n # Verify new format\n assert len(results) == 3\n\n # First connector - full config\n config_0 = results[0].connector_specific_config\n assert config_0[\"jira_base_url\"] == \"https://example.atlassian.net\"\n assert config_0[\"project_key\"] == \"PROJ\"\n assert config_0[\"comment_email_blacklist\"] == [\"test@example.com\"]\n assert config_0[\"batch_size\"] == 100\n assert config_0[\"labels_to_skip\"] == [\"skip-me\"]\n\n # Second connector - minimal config\n config_1 = results[1].connector_specific_config\n assert config_1[\"jira_base_url\"] == \"https://other.atlassian.net\"\n assert config_1[\"project_key\"] == \"OTHER\"\n assert \"comment_email_blacklist\" not in config_1\n assert \"batch_size\" not in config_1\n assert \"labels_to_skip\" not in config_1\n\n # Third connector - partial config\n config_2 = results[2].connector_specific_config\n assert config_2[\"jira_base_url\"] == \"https://example.atlassian.net\"\n assert config_2[\"project_key\"] == \"TEST\"\n assert config_2[\"batch_size\"] == 50\n assert \"comment_email_blacklist\" not in config_2\n assert \"labels_to_skip\" not in config_2\n\n # Test downgrade path\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"f13db29f3101\"\n )\n\n # Verify the downgrade was applied correctly\n with get_session_with_current_tenant() as db_session:\n results = db_session.execute(\n text(\n \"\"\"\n SELECT id, connector_specific_config\n FROM connector\n WHERE source = 'JIRA'\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n # Verify reverted to old format\n assert len(results) == 3\n\n # First connector - full config\n config_0 = results[0].connector_specific_config\n assert (\n config_0[\"jira_project_url\"]\n == \"https://example.atlassian.net/projects/PROJ\"\n )\n assert config_0[\"comment_email_blacklist\"] == [\"test@example.com\"]\n assert config_0[\"batch_size\"] == 100\n assert config_0[\"labels_to_skip\"] == [\"skip-me\"]\n\n # Second connector - minimal config\n config_1 = results[1].connector_specific_config\n assert (\n config_1[\"jira_project_url\"] == \"https://other.atlassian.net/projects/OTHER\"\n )\n\n # Third connector - partial config\n config_2 = results[2].connector_specific_config\n assert (\n config_2[\"jira_project_url\"]\n == \"https://example.atlassian.net/projects/TEST\"\n )\n assert config_2[\"batch_size\"] == 50\n\n\ndef test_anonymous_user_migration_dedupes_null_notifications() -> None:\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"base\", clear_data=True\n )\n upgrade_postgres(\n database=\"postgres\",\n config_name=\"alembic\",\n revision=\"f7ca3e2f45d9\",\n )\n\n with get_session_with_current_tenant() as db_session:\n db_session.execute(\n text(\n \"\"\"\n INSERT INTO notification (\n id,\n notif_type,\n user_id,\n dismissed,\n last_shown,\n first_shown,\n title,\n description,\n additional_data\n )\n VALUES\n (\n 1,\n 'RELEASE_NOTES',\n NULL,\n FALSE,\n NOW(),\n NOW(),\n 'Onyx v2.10.0 is available!',\n 'Check out what''s new in v2.10.0',\n '{\"version\":\"v2.10.0\",\"link\":\"https://docs.onyx.app/changelog#v2-10-0\"}'::jsonb\n ),\n (\n 2,\n 'RELEASE_NOTES',\n NULL,\n FALSE,\n NOW(),\n NOW(),\n 'Onyx v2.10.0 is available!',\n 'Check out what''s new in v2.10.0',\n '{\"version\":\"v2.10.0\",\"link\":\"https://docs.onyx.app/changelog#v2-10-0\"}'::jsonb\n )\n \"\"\"\n )\n )\n db_session.commit()\n\n upgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"e7f8a9b0c1d2\"\n )\n\n with get_session_with_current_tenant() as db_session:\n notifications = db_session.execute(\n text(\n \"\"\"\n SELECT id, user_id\n FROM notification\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n anonymous_user = db_session.execute(\n text(\n \"\"\"\n SELECT id, email, role\n FROM \"user\"\n WHERE id = :user_id\n \"\"\"\n ),\n {\"user_id\": ANONYMOUS_USER_UUID},\n ).fetchone()\n\n assert len(notifications) == 1\n assert notifications[0].id == 2 # Higher id wins when timestamps are equal\n assert str(notifications[0].user_id) == ANONYMOUS_USER_UUID\n assert anonymous_user is not None\n assert anonymous_user.email == \"anonymous@onyx.app\"\n assert anonymous_user.role == \"LIMITED\"\n\n\ndef test_anonymous_user_migration_collision_with_existing_anonymous_notification() -> (\n None\n):\n \"\"\"Test that a NULL-owned notification that collides with an already-existing\n anonymous-owned notification is removed during migration.\"\"\"\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"base\", clear_data=True\n )\n upgrade_postgres(\n database=\"postgres\",\n config_name=\"alembic\",\n revision=\"f7ca3e2f45d9\",\n )\n\n with get_session_with_current_tenant() as db_session:\n # Create the anonymous user early so we can insert a notification owned by it\n db_session.execute(\n text(\n \"\"\"\n INSERT INTO \"user\" (id, email, hashed_password, is_active, is_superuser, is_verified, role)\n VALUES (:id, 'anonymous@onyx.app', '', TRUE, FALSE, TRUE, 'LIMITED')\n ON CONFLICT (id) DO NOTHING\n \"\"\"\n ),\n {\"id\": ANONYMOUS_USER_UUID},\n )\n # Insert an anonymous-owned notification (already migrated in a prior partial run)\n db_session.execute(\n text(\n \"\"\"\n INSERT INTO notification (\n id, notif_type, user_id, dismissed, last_shown, first_shown,\n title, description, additional_data\n )\n VALUES\n (\n 1, 'RELEASE_NOTES', :user_id, FALSE, NOW(), NOW(),\n 'Onyx v2.10.0 is available!',\n 'Check out what''s new in v2.10.0',\n '{\"version\":\"v2.10.0\",\"link\":\"https://docs.onyx.app/changelog#v2-10-0\"}'::jsonb\n ),\n (\n 2, 'RELEASE_NOTES', NULL, FALSE, NOW(), NOW(),\n 'Onyx v2.10.0 is available!',\n 'Check out what''s new in v2.10.0',\n '{\"version\":\"v2.10.0\",\"link\":\"https://docs.onyx.app/changelog#v2-10-0\"}'::jsonb\n )\n \"\"\"\n ),\n {\"user_id\": ANONYMOUS_USER_UUID},\n )\n db_session.commit()\n\n upgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"e7f8a9b0c1d2\"\n )\n\n with get_session_with_current_tenant() as db_session:\n notifications = db_session.execute(\n text(\n \"\"\"\n SELECT id, user_id\n FROM notification\n ORDER BY id\n \"\"\"\n )\n ).fetchall()\n\n # Only the original anonymous-owned notification should remain;\n # the NULL-owned duplicate should have been deleted\n assert len(notifications) == 1\n assert notifications[0].id == 1\n assert str(notifications[0].user_id) == ANONYMOUS_USER_UUID\n" }, { "path": "backend/tests/integration/tests/migrations/test_tool_seeding.py", "content": "from pydantic import BaseModel\nfrom sqlalchemy import text\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom tests.integration.common_utils.reset import downgrade_postgres\nfrom tests.integration.common_utils.reset import upgrade_postgres\n\n\nclass ToolSeedingExpectedResult(BaseModel):\n name: str\n display_name: str\n in_code_tool_id: str\n user_id: str | None\n\n\nEXPECTED_TOOLS = {\n \"SearchTool\": ToolSeedingExpectedResult(\n name=\"internal_search\",\n display_name=\"Internal Search\",\n in_code_tool_id=\"SearchTool\",\n user_id=None,\n ),\n \"ImageGenerationTool\": ToolSeedingExpectedResult(\n name=\"generate_image\",\n display_name=\"Image Generation\",\n in_code_tool_id=\"ImageGenerationTool\",\n user_id=None,\n ),\n \"WebSearchTool\": ToolSeedingExpectedResult(\n name=\"web_search\",\n display_name=\"Web Search\",\n in_code_tool_id=\"WebSearchTool\",\n user_id=None,\n ),\n \"KnowledgeGraphTool\": ToolSeedingExpectedResult(\n name=\"run_kg_search\",\n display_name=\"Knowledge Graph Search\",\n in_code_tool_id=\"KnowledgeGraphTool\",\n user_id=None,\n ),\n \"PythonTool\": ToolSeedingExpectedResult(\n name=\"python\",\n display_name=\"Code Interpreter\",\n in_code_tool_id=\"PythonTool\",\n user_id=None,\n ),\n \"ResearchAgent\": ToolSeedingExpectedResult(\n name=\"research_agent\",\n display_name=\"Research Agent\",\n in_code_tool_id=\"ResearchAgent\",\n user_id=None,\n ),\n \"FileReaderTool\": ToolSeedingExpectedResult(\n name=\"read_file\",\n display_name=\"File Reader\",\n in_code_tool_id=\"FileReaderTool\",\n user_id=None,\n ),\n \"MemoryTool\": ToolSeedingExpectedResult(\n name=\"MemoryTool\",\n display_name=\"Add Memory\",\n in_code_tool_id=\"MemoryTool\",\n user_id=None,\n ),\n}\n\n\ndef test_tool_seeding_migration() -> None:\n \"\"\"Test that migration from base to head correctly seeds builtin tools.\"\"\"\n # Start from base and upgrade to just before tool seeding\n downgrade_postgres(\n database=\"postgres\", config_name=\"alembic\", revision=\"base\", clear_data=True\n )\n upgrade_postgres(\n database=\"postgres\",\n config_name=\"alembic\",\n revision=\"b7ec9b5b505f\", # Revision before tool seeding\n )\n\n # Verify no tools exist yet\n with get_session_with_current_tenant() as db_session:\n result = db_session.execute(text(\"SELECT COUNT(*) FROM tool\"))\n count = result.scalar()\n assert count == 0, \"No tools should exist before migration\"\n\n # Upgrade to head\n upgrade_postgres(\n database=\"postgres\",\n config_name=\"alembic\",\n revision=\"head\",\n )\n\n # Verify tools were created\n with get_session_with_current_tenant() as db_session:\n result = db_session.execute(\n text(\n \"\"\"\n SELECT id, name, display_name, description, in_code_tool_id,\n user_id\n FROM tool\n ORDER BY id\n \"\"\"\n )\n )\n tools = result.fetchall()\n\n # Should have all 9 builtin tools\n assert (\n len(tools) == 10\n ), f\"Should have created exactly 9 builtin tools, got {len(tools)}\"\n\n def validate_tool(expected: ToolSeedingExpectedResult) -> None:\n tool = next((t for t in tools if t[1] == expected.name), None)\n assert tool is not None, f\"{expected.name} should exist\"\n assert (\n tool[2] == expected.display_name\n ), f\"{expected.name} display name should be '{expected.display_name}'\"\n assert (\n tool[4] == expected.in_code_tool_id\n ), f\"{expected.name} in_code_tool_id should be '{expected.in_code_tool_id}'\"\n assert (\n tool[5] is None\n ), f\"{expected.name} should not have a user_id (builtin)\"\n\n # Check SearchTool\n validate_tool(EXPECTED_TOOLS[\"SearchTool\"])\n\n # Check ImageGenerationTool\n validate_tool(EXPECTED_TOOLS[\"ImageGenerationTool\"])\n\n # Check WebSearchTool\n validate_tool(EXPECTED_TOOLS[\"WebSearchTool\"])\n\n # Check KnowledgeGraphTool\n validate_tool(EXPECTED_TOOLS[\"KnowledgeGraphTool\"])\n\n # Check PythonTool\n validate_tool(EXPECTED_TOOLS[\"PythonTool\"])\n\n # Check ResearchAgent (Deep Research as a tool)\n validate_tool(EXPECTED_TOOLS[\"ResearchAgent\"])\n\n # Check FileReaderTool\n validate_tool(EXPECTED_TOOLS[\"FileReaderTool\"])\n\n # Check MemoryTool\n validate_tool(EXPECTED_TOOLS[\"MemoryTool\"])\n" }, { "path": "backend/tests/integration/tests/no_vectordb/conftest.py", "content": "\"\"\"Fixtures for no-vector-DB integration tests.\n\nThese tests are intended to run against an Onyx deployment started with\nDISABLE_VECTOR_DB=true. They are automatically **skipped** when the\nserver reports vector_db_enabled=true (i.e. when Vespa is available).\n\"\"\"\n\nimport pytest\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import GENERAL_HEADERS\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.reset import reset_file_store\nfrom tests.integration.common_utils.reset import reset_postgres\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _server_has_vector_db_disabled() -> bool:\n \"\"\"Query the running server to check whether DISABLE_VECTOR_DB is set.\"\"\"\n try:\n resp = requests.get(\n f\"{API_SERVER_URL}/settings\",\n headers=GENERAL_HEADERS,\n )\n if resp.ok:\n return resp.json().get(\"vector_db_enabled\") is False\n except Exception:\n pass\n return False\n\n\n# Skip the entire module when the server has vector DB enabled —\n# these tests only make sense in no-vector-DB deployments.\npytestmark = pytest.mark.skipif(\n not _server_has_vector_db_disabled(),\n reason=\"Server is running with vector DB enabled; skipping no-vectordb tests\",\n)\n\n\n@pytest.fixture()\ndef reset() -> None:\n \"\"\"Reset Postgres and the file store, but skip Vespa (not running).\"\"\"\n reset_postgres()\n reset_file_store()\n\n\n@pytest.fixture()\ndef llm_provider(admin_user: DATestUser) -> DATestLLMProvider:\n \"\"\"Ensure an LLM provider exists for the test session.\"\"\"\n return LLMProviderManager.create(user_performing_action=admin_user)\n" }, { "path": "backend/tests/integration/tests/no_vectordb/test_no_vectordb_chat.py", "content": "\"\"\"Integration tests for chat in no-vector-DB mode.\n\nCovers:\n- Uploading a file to a project, sending a chat message, and verifying the LLM\n receives the file content (small project — fits in context window).\n- Creating a persona with user_files and verifying chat works.\n- Verifying that persona creation with document_sets / hierarchy_nodes /\n document_ids is rejected with a 400.\n\"\"\"\n\nimport io\nimport time\n\nimport requests\n\nfrom onyx.db.enums import UserFileStatus\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.file import FileManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.project import ProjectManager\nfrom tests.integration.common_utils.managers.tool import ToolManager\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nFILE_READER_TOOL_ID = \"FileReaderTool\"\n\n\ndef _wait_for_file_processed(\n project_id: int,\n user: DATestUser,\n timeout: int = 30,\n) -> None:\n \"\"\"Poll until all files in the project reach COMPLETED status.\"\"\"\n deadline = time.time() + timeout\n while time.time() < deadline:\n files = ProjectManager.get_project_files(project_id, user)\n if files and all(f.status == UserFileStatus.COMPLETED for f in files):\n return\n time.sleep(1)\n raise TimeoutError(\n f\"Files in project {project_id} did not reach COMPLETED within {timeout}s\"\n )\n\n\n# ------------------------------------------------------------------\n# Small-project chat — file content loaded directly into context\n# ------------------------------------------------------------------\n\n\ndef test_chat_with_small_project_file(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"Upload a small text file to a project and send a chat message.\n\n The file is small enough to fit in the LLM context window, so the LLM\n should see the file content directly and be able to answer questions\n about it.\n \"\"\"\n project = ProjectManager.create(\n name=\"test-no-vectordb-small\", user_performing_action=admin_user\n )\n\n file_content = b\"The secret code is PINEAPPLE-42.\"\n ProjectManager.upload_files(\n project_id=project.id,\n files=[(\"secret.txt\", file_content)],\n user_performing_action=admin_user,\n )\n\n _wait_for_file_processed(project.id, admin_user)\n\n # Create a chat session associated with the project's default persona\n chat_session = ChatSessionManager.create(\n persona_id=0,\n description=\"no-vectordb small project test\",\n user_performing_action=admin_user,\n )\n\n # Link the chat session to the project\n resp = requests.post(\n f\"{API_SERVER_URL}/user/projects/{project.id}/move_chat_session\",\n json={\"chat_session_id\": str(chat_session.id)},\n headers=admin_user.headers,\n )\n resp.raise_for_status()\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"What is the secret code in the file?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None, f\"Chat returned an error: {response.error}\"\n assert (\n \"PINEAPPLE-42\" in response.full_message\n ), f\"Expected the LLM to reference the file content. Got: {response.full_message}\"\n\n\n# ------------------------------------------------------------------\n# Persona with user_files — should work in no-vector-DB mode\n# ------------------------------------------------------------------\n\n\ndef test_persona_with_user_files_chat(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"Create a persona with attached user files and verify chat works.\"\"\"\n # Upload a file first\n file_content = b\"Quarterly revenue was $42 million.\"\n file_obj = io.BytesIO(file_content)\n file_descriptors, error = FileManager.upload_files(\n files=[(\"revenue.txt\", file_obj)],\n user_performing_action=admin_user,\n )\n assert not error, f\"File upload failed: {error}\"\n assert len(file_descriptors) > 0\n\n user_file_id = file_descriptors[0].get(\"user_file_id\")\n assert user_file_id, \"Expected user_file_id in upload response\"\n\n # Wait for the file to be processed\n deadline = time.time() + 30\n while time.time() < deadline:\n time.sleep(1)\n # Check via file fetch — if it succeeds, the file is ready\n try:\n FileManager.fetch_uploaded_file(\n file_descriptors[0][\"id\"],\n admin_user,\n )\n break\n except Exception:\n continue\n\n # Find the FileReaderTool ID from available tools\n tools = ToolManager.list_tools(user_performing_action=admin_user)\n file_reader_tool = next(\n (t for t in tools if t.in_code_tool_id == FILE_READER_TOOL_ID), None\n )\n assert (\n file_reader_tool is not None\n ), \"FileReaderTool should be registered as a built-in tool\"\n\n # Create a persona with the user file attached\n persona = PersonaManager.create(\n name=\"no-vectordb-persona-test\",\n description=\"Test persona for no-vectordb mode\",\n system_prompt=\"You are a helpful assistant. Answer questions using the available tools and files.\",\n task_prompt=\"\",\n user_file_ids=[user_file_id],\n tool_ids=[file_reader_tool.id],\n user_performing_action=admin_user,\n )\n\n chat_session = ChatSessionManager.create(\n persona_id=persona.id,\n description=\"no-vectordb persona test\",\n user_performing_action=admin_user,\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"What was the quarterly revenue?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None, f\"Chat returned an error: {response.error}\"\n # The LLM should be able to answer about the revenue (either from direct\n # context injection or via the FileReaderTool)\n assert (\n \"$42 million\" in response.full_message or \"42\" in response.full_message\n ), f\"Expected the LLM to reference the file content. Got: {response.full_message}\"\n\n\n# ------------------------------------------------------------------\n# Persona validation — vector-DB knowledge types rejected\n# ------------------------------------------------------------------\n\n\ndef _base_persona_body(**overrides: object) -> dict:\n \"\"\"Build a valid PersonaUpsertRequest body with sensible defaults.\n\n Callers override only the fields under test so that Pydantic validation\n passes and the vector-DB guard (``_validate_vector_db_knowledge``) is\n the one that rejects the request.\n \"\"\"\n body: dict = {\n \"name\": \"should-fail\",\n \"description\": \"test\",\n \"system_prompt\": \"test\",\n \"task_prompt\": \"\",\n \"is_public\": True,\n \"datetime_aware\": False,\n \"document_set_ids\": [],\n \"tool_ids\": [],\n \"users\": [],\n \"groups\": [],\n }\n body.update(overrides)\n return body\n\n\ndef test_persona_rejects_document_sets_without_vector_db(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Creating a persona with document_set_ids should fail with 400.\"\"\"\n resp = requests.post(\n f\"{API_SERVER_URL}/persona\",\n json=_base_persona_body(document_set_ids=[1]),\n headers=admin_user.headers,\n )\n assert (\n resp.status_code == 400\n ), f\"Expected 400 for document_set_ids, got {resp.status_code}: {resp.text}\"\n\n\ndef test_persona_rejects_document_ids_without_vector_db(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Creating a persona with document_ids should fail with 400.\"\"\"\n resp = requests.post(\n f\"{API_SERVER_URL}/persona\",\n json=_base_persona_body(document_ids=[\"fake-doc-id\"]),\n headers=admin_user.headers,\n )\n assert (\n resp.status_code == 400\n ), f\"Expected 400 for document_ids, got {resp.status_code}: {resp.text}\"\n" }, { "path": "backend/tests/integration/tests/no_vectordb/test_no_vectordb_endpoints.py", "content": "\"\"\"Integration tests for endpoint gating when DISABLE_VECTOR_DB is set.\n\nVector-DB-dependent endpoints should return HTTP 501.\nNon-dependent endpoints (settings, document sets, chat, etc.) should work\nnormally.\n\"\"\"\n\nimport requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n# ------------------------------------------------------------------\n# Helper\n# ------------------------------------------------------------------\n\n\ndef _headers(user: DATestUser) -> dict[str, str]:\n return user.headers if user else {\"Content-Type\": \"application/json\"}\n\n\n# ------------------------------------------------------------------\n# Gated endpoints — should return 501\n# ------------------------------------------------------------------\n\n\ndef test_admin_search_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.post(\n f\"{API_SERVER_URL}/admin/search\",\n json={\"query\": \"test\", \"filters\": {}},\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501, f\"Expected 501, got {resp.status_code}\"\n\n\ndef test_document_size_info_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.get(\n f\"{API_SERVER_URL}/document/document-size-info\",\n params={\"document_id\": \"fake-doc\"},\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\ndef test_document_chunk_info_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.get(\n f\"{API_SERVER_URL}/document/chunk-info\",\n params={\"document_id\": \"fake-doc\"},\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\ndef test_set_new_search_settings_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.post(\n f\"{API_SERVER_URL}/search-settings/set-new-search-settings\",\n json={},\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\ndef test_cancel_new_embedding_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.post(\n f\"{API_SERVER_URL}/search-settings/cancel-new-embedding\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\ndef test_connector_router_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"The entire /manage router is gated — any connector endpoint should 501.\"\"\"\n resp = requests.get(\n f\"{API_SERVER_URL}/manage/connector\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\ndef test_ingestion_post_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.post(\n f\"{API_SERVER_URL}/onyx-api/ingestion\",\n json={\"document\": {}},\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\ndef test_ingestion_delete_returns_501(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.delete(\n f\"{API_SERVER_URL}/onyx-api/ingestion/fake-doc-id\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 501\n\n\n# ------------------------------------------------------------------\n# Non-gated endpoints — should work (2xx)\n# ------------------------------------------------------------------\n\n\ndef test_settings_endpoint_works(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.get(\n f\"{API_SERVER_URL}/settings\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 200\n data = resp.json()\n assert data[\"vector_db_enabled\"] is False\n\n\ndef test_document_set_list_works(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.get(\n f\"{API_SERVER_URL}/manage/document-set\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 200\n\n\ndef test_persona_list_works(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n resp = requests.get(\n f\"{API_SERVER_URL}/admin/persona\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 200\n\n\ndef test_tool_list_works(\n reset: None, admin_user: DATestUser # noqa: ARG001\n) -> None: # noqa: ARG001\n resp = requests.get(\n f\"{API_SERVER_URL}/tool\",\n headers=_headers(admin_user),\n )\n assert resp.status_code == 200\n tools = resp.json()\n tool_ids = {t[\"in_code_tool_id\"] for t in tools if t.get(\"in_code_tool_id\")}\n assert (\n \"FileReaderTool\" in tool_ids\n ), \"FileReaderTool should be registered as a built-in tool\"\n" }, { "path": "backend/tests/integration/tests/no_vectordb/test_no_vectordb_file_lifecycle.py", "content": "\"\"\"Integration test for the full user-file lifecycle in no-vector-DB mode.\n\nCovers: upload → COMPLETED → unlink from project → delete → gone.\n\nThe entire lifecycle is handled by FastAPI BackgroundTasks (no Celery workers\nneeded). The conftest-level ``pytestmark`` ensures these tests are skipped\nwhen the server is running with vector DB enabled.\n\"\"\"\n\nimport time\nfrom uuid import UUID\n\nimport requests\n\nfrom onyx.db.enums import UserFileStatus\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.project import ProjectManager\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\nPOLL_INTERVAL_SECONDS = 1\nPOLL_TIMEOUT_SECONDS = 30\n\n\ndef _poll_file_status(\n file_id: UUID,\n user: DATestUser,\n target_status: UserFileStatus,\n timeout: int = POLL_TIMEOUT_SECONDS,\n) -> None:\n \"\"\"Poll GET /user/projects/file/{file_id} until the file reaches *target_status*.\"\"\"\n deadline = time.time() + timeout\n while time.time() < deadline:\n resp = requests.get(\n f\"{API_SERVER_URL}/user/projects/file/{file_id}\",\n headers=user.headers,\n )\n if resp.ok:\n status = resp.json().get(\"status\")\n if status == target_status.value:\n return\n time.sleep(POLL_INTERVAL_SECONDS)\n raise TimeoutError(\n f\"File {file_id} did not reach {target_status.value} within {timeout}s\"\n )\n\n\ndef _file_is_gone(file_id: UUID, user: DATestUser, timeout: int = 15) -> None:\n \"\"\"Poll until GET /user/projects/file/{file_id} returns 404.\"\"\"\n deadline = time.time() + timeout\n while time.time() < deadline:\n resp = requests.get(\n f\"{API_SERVER_URL}/user/projects/file/{file_id}\",\n headers=user.headers,\n )\n if resp.status_code == 404:\n return\n time.sleep(POLL_INTERVAL_SECONDS)\n raise TimeoutError(\n f\"File {file_id} still accessible after {timeout}s (expected 404)\"\n )\n\n\ndef test_file_upload_process_delete_lifecycle(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"Full lifecycle: upload → COMPLETED → unlink → delete → 404.\n\n Validates that the API server handles all background processing\n (via FastAPI BackgroundTasks) without any Celery workers running.\n \"\"\"\n project = ProjectManager.create(\n name=\"lifecycle-test\", user_performing_action=admin_user\n )\n\n file_content = b\"Integration test file content for lifecycle verification.\"\n upload_result = ProjectManager.upload_files(\n project_id=project.id,\n files=[(\"lifecycle.txt\", file_content)],\n user_performing_action=admin_user,\n )\n assert upload_result.user_files, \"Expected at least one file in upload response\"\n\n user_file = upload_result.user_files[0]\n file_id = user_file.id\n\n _poll_file_status(file_id, admin_user, UserFileStatus.COMPLETED)\n\n project_files = ProjectManager.get_project_files(project.id, admin_user)\n assert any(\n f.id == file_id for f in project_files\n ), \"File should be listed in project files after processing\"\n\n # Unlink the file from the project so the delete endpoint will proceed\n unlink_resp = requests.delete(\n f\"{API_SERVER_URL}/user/projects/{project.id}/files/{file_id}\",\n headers=admin_user.headers,\n )\n assert (\n unlink_resp.status_code == 204\n ), f\"Expected 204 on unlink, got {unlink_resp.status_code}: {unlink_resp.text}\"\n\n delete_resp = requests.delete(\n f\"{API_SERVER_URL}/user/projects/file/{file_id}\",\n headers=admin_user.headers,\n )\n assert (\n delete_resp.ok\n ), f\"Delete request failed: {delete_resp.status_code} {delete_resp.text}\"\n body = delete_resp.json()\n assert (\n body[\"has_associations\"] is False\n ), f\"File still has associations after unlink: {body}\"\n\n _file_is_gone(file_id, admin_user)\n\n project_files_after = ProjectManager.get_project_files(project.id, admin_user)\n assert not any(\n f.id == file_id for f in project_files_after\n ), \"Deleted file should not appear in project files\"\n\n\ndef test_delete_blocked_while_associated(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"Deleting a file that still belongs to a project should return\n has_associations=True without actually deleting the file.\"\"\"\n project = ProjectManager.create(\n name=\"assoc-test\", user_performing_action=admin_user\n )\n\n upload_result = ProjectManager.upload_files(\n project_id=project.id,\n files=[(\"assoc.txt\", b\"associated file content\")],\n user_performing_action=admin_user,\n )\n file_id = upload_result.user_files[0].id\n\n _poll_file_status(file_id, admin_user, UserFileStatus.COMPLETED)\n\n # Attempt to delete while still linked\n delete_resp = requests.delete(\n f\"{API_SERVER_URL}/user/projects/file/{file_id}\",\n headers=admin_user.headers,\n )\n assert delete_resp.ok\n body = delete_resp.json()\n assert body[\"has_associations\"] is True, \"Should report existing associations\"\n assert project.name in body[\"project_names\"]\n\n # File should still be accessible\n get_resp = requests.get(\n f\"{API_SERVER_URL}/user/projects/file/{file_id}\",\n headers=admin_user.headers,\n )\n assert get_resp.status_code == 200, \"File should still exist after blocked delete\"\n" }, { "path": "backend/tests/integration/tests/opensearch_migration/test_opensearch_migration_api.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_migration_status_returns_defaults_when_no_record(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"When no migration record exists, status should return zeros/nulls.\"\"\"\n # Under test.\n response = requests.get(\n f\"{API_SERVER_URL}/admin/opensearch-migration/status\",\n headers=admin_user.headers,\n )\n\n # Postcondition.\n assert response.status_code == 200\n data = response.json()\n assert data[\"total_chunks_migrated\"] == 0\n assert data[\"created_at\"] is None\n assert data[\"migration_completed_at\"] is None\n assert data[\"approx_chunk_count_in_vespa\"] is None\n\n\ndef test_retrieval_status_returns_false_when_no_record(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"When no migration record exists, retrieval should default to disabled.\"\"\"\n # Under test.\n response = requests.get(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n headers=admin_user.headers,\n )\n\n # Postcondition.\n assert response.status_code == 200\n data = response.json()\n assert data[\"enable_opensearch_retrieval\"] is False\n\n\ndef test_set_and_get_retrieval_status(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Setting retrieval to True should persist and be readable.\"\"\"\n # Under test.\n # Enable retrieval.\n response = requests.put(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n json={\"enable_opensearch_retrieval\": True},\n headers=admin_user.headers,\n )\n\n # Postcondition.\n assert response.status_code == 200\n assert response.json()[\"enable_opensearch_retrieval\"] is True\n # Verify it persisted.\n response = requests.get(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n assert response.json()[\"enable_opensearch_retrieval\"] is True\n\n # Under test.\n # Disable retrieval.\n response = requests.put(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n json={\"enable_opensearch_retrieval\": False},\n headers=admin_user.headers,\n )\n\n # Postcondition.\n assert response.status_code == 200\n assert response.json()[\"enable_opensearch_retrieval\"] is False\n # Verify it persisted.\n response = requests.get(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n headers=admin_user.headers,\n )\n assert response.status_code == 200\n assert response.json()[\"enable_opensearch_retrieval\"] is False\n\n\ndef test_migration_status_after_record_created(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"After toggling retrieval (which creates the record), status should\n return a valid created_at timestamp.\"\"\"\n # Precondition.\n # Create the record by setting retrieval.\n requests.put(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n json={\"enable_opensearch_retrieval\": False},\n headers=admin_user.headers,\n )\n\n # Under test.\n response = requests.get(\n f\"{API_SERVER_URL}/admin/opensearch-migration/status\",\n headers=admin_user.headers,\n )\n\n # Postcondition.\n assert response.status_code == 200\n data = response.json()\n assert data[\"total_chunks_migrated\"] == 0\n assert data[\"created_at\"] is not None\n assert data[\"migration_completed_at\"] is None\n assert data[\"approx_chunk_count_in_vespa\"] is None\n\n\ndef test_endpoints_require_admin(\n reset: None, # noqa: ARG001\n admin_user: DATestUser, # noqa: ARG001\n) -> None:\n \"\"\"Endpoints should reject unauthenticated requests.\"\"\"\n for url in [\n f\"{API_SERVER_URL}/admin/opensearch-migration/status\",\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n ]:\n response = requests.get(url)\n assert response.status_code == 403\n\n response = requests.put(\n f\"{API_SERVER_URL}/admin/opensearch-migration/retrieval\",\n json={\"enable_opensearch_retrieval\": True},\n )\n assert response.status_code == 403\n" }, { "path": "backend/tests/integration/tests/pat/test_pat_api.py", "content": "\"\"\"\nIntegration tests for Personal Access Token (PAT) API.\n\nTest Suite:\n1. test_pat_lifecycle_happy_path - Complete PAT lifecycle (create, auth, revoke)\n2. test_pat_user_isolation_and_authentication - User authentication and multi-user isolation\n3. test_pat_expiration_flow - Expiration logic (end-of-day UTC, never-expiring)\n4. test_pat_validation_errors - Input validation and error handling\n5. test_pat_sorting_and_last_used - Sorting and last_used_at tracking\n6. test_pat_role_based_access_control - Admin vs Basic vs Curator permissions\n\"\"\"\n\nimport time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nimport requests\n\nfrom onyx.auth.schemas import UserRole\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.pat import PATManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_pat_lifecycle_happy_path(reset: None) -> None: # noqa: ARG001\n \"\"\"Complete PAT lifecycle: create, authenticate, revoke.\"\"\"\n user: DATestUser = UserManager.create(name=\"pat_user\")\n\n # Create PAT\n pat = PATManager.create(\n name=\"My Integration Token\",\n expiration_days=30,\n user_performing_action=user,\n )\n\n assert pat.id is not None\n assert pat.name == \"My Integration Token\"\n assert pat.token is not None # Raw token only returned on creation\n assert pat.token_display is not None\n assert pat.created_at is not None\n assert pat.expires_at is not None\n\n assert pat.token.startswith(\"onyx_pat_\")\n assert len(pat.token) > 20\n\n assert \"****\" in pat.token_display\n assert pat.token_display.startswith(\"onyx_pat_\")\n\n # List PATs\n tokens = PATManager.list(user)\n assert len(tokens) == 1\n assert tokens[0].id == pat.id\n assert tokens[0].name == \"My Integration Token\"\n assert tokens[0].token_display == pat.token_display\n assert tokens[0].token is None\n\n # Authenticate with PAT\n auth_response = PATManager.authenticate(pat.token)\n assert auth_response.status_code == 200\n me_data = auth_response.json()\n assert me_data[\"email\"] == user.email\n assert me_data[\"id\"] == user.id\n\n # Revoke PAT\n PATManager.revoke(pat.id, user)\n\n # Verify revoked token fails authentication\n revoked_auth_response = PATManager.authenticate(pat.token)\n assert revoked_auth_response.status_code == 403 # Revoked token returns 403\n\n # Verify token is no longer listed\n tokens_after_revoke = PATManager.list(user)\n assert len(tokens_after_revoke) == 0\n\n\ndef test_pat_user_isolation_and_authentication(\n reset: None, # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"\n PATs authenticate as real users, and users can only see/manage their own tokens.\n \"\"\"\n user_a: DATestUser = UserManager.create(name=\"user_a\")\n user_b: DATestUser = UserManager.create(name=\"user_b\")\n\n # Create tokens for both users\n user_a_pats = []\n for i in range(2):\n pat = PATManager.create(\n name=f\"User A Token {i + 1}\",\n expiration_days=30,\n user_performing_action=user_a,\n )\n user_a_pats.append(pat)\n\n user_b_pats = []\n for i in range(2):\n pat = PATManager.create(\n name=f\"User B Token {i + 1}\",\n expiration_days=30,\n user_performing_action=user_b,\n )\n user_b_pats.append(pat)\n\n # Verify PATs authenticate as the correct users\n for user, pat in [(user_a, user_a_pats[0]), (user_b, user_b_pats[0])]:\n assert pat.token is not None\n me_response = PATManager.authenticate(pat.token)\n assert me_response.status_code == 200\n me_data = me_response.json()\n assert me_data[\"email\"] == user.email\n assert me_data[\"id\"] == user.id\n\n # Verify each user only sees their own tokens\n user_a_list = PATManager.list(user_a)\n assert len(user_a_list) == 2\n\n user_b_list = PATManager.list(user_b)\n assert len(user_b_list) == 2\n\n # Verify user A cannot delete user B's token using their PAT\n assert user_a_pats[0].token is not None\n delete_response = requests.delete(\n f\"{API_SERVER_URL}/user/pats/{user_b_pats[0].id}\",\n headers=PATManager.get_auth_headers(user_a_pats[0].token),\n timeout=60,\n )\n assert delete_response.status_code == 404\n\n # Verify user B's token still exists\n user_b_list_after = PATManager.list(user_b)\n assert len(user_b_list_after) == 2\n\n # Verify deleting non-existent token returns 404\n delete_fake = requests.delete(\n f\"{API_SERVER_URL}/user/pats/999999\",\n headers=user_a.headers,\n timeout=60,\n )\n assert delete_fake.status_code == 404\n\n\ndef test_pat_expiration_flow(reset: None) -> None: # noqa: ARG001\n \"\"\"Expiration timestamp is end-of-day (23:59:59 UTC); never-expiring tokens work; revoked tokens fail.\"\"\"\n user: DATestUser = UserManager.create(name=\"expiration_user\")\n\n # Create expiring token\n pat = PATManager.create(\n name=\"Expiring Token\",\n expiration_days=7,\n user_performing_action=user,\n )\n\n assert pat.expires_at is not None\n expires_at = datetime.fromisoformat(pat.expires_at.replace(\"Z\", \"+00:00\"))\n\n # Verify end-of-day expiration\n assert expires_at.hour == 23\n assert expires_at.minute == 59\n assert expires_at.second == 59\n\n # Calculate expected end-of-day 7 days from now\n now = datetime.now(timezone.utc)\n expected_date = (now + timedelta(days=7)).date()\n expected_expiry = datetime.combine(expected_date, datetime.max.time()).replace(\n tzinfo=timezone.utc\n )\n # Allow for small timing differences (within a day)\n assert abs((expires_at - expected_expiry).total_seconds()) < 86400 # 1 day\n\n # Create never-expiring token\n never_expiring_pat = PATManager.create(\n name=\"Never Expiring Token\",\n expiration_days=None,\n user_performing_action=user,\n )\n assert never_expiring_pat.expires_at is None\n\n # Verify never-expiring token works\n assert never_expiring_pat.token is not None\n auth_response = PATManager.authenticate(never_expiring_pat.token)\n assert auth_response.status_code == 200\n\n # Revoke the never-expiring token\n PATManager.revoke(never_expiring_pat.id, user)\n\n # Verify revoked token fails (token var still holds the revoked value)\n revoked_auth_response = PATManager.authenticate(never_expiring_pat.token)\n assert revoked_auth_response.status_code == 403\n\n\ndef test_pat_validation_errors(reset: None) -> None: # noqa: ARG001\n \"\"\"Validate input errors: empty name, name too long, negative/zero expiration.\"\"\"\n user: DATestUser = UserManager.create(name=\"validation_user\")\n\n # Empty name should fail\n empty_name_response = requests.post(\n f\"{API_SERVER_URL}/user/pats\",\n json={\"name\": \"\", \"expiration_days\": 30},\n headers=user.headers,\n timeout=60,\n )\n assert empty_name_response.status_code == 422\n\n # Name too long should fail\n long_name = \"a\" * 101\n long_name_response = requests.post(\n f\"{API_SERVER_URL}/user/pats\",\n json={\"name\": long_name, \"expiration_days\": 30},\n headers=user.headers,\n timeout=60,\n )\n assert long_name_response.status_code == 422\n\n # Negative expiration should fail\n negative_exp_response = requests.post(\n f\"{API_SERVER_URL}/user/pats\",\n json={\"name\": \"Test Token\", \"expiration_days\": -1},\n headers=user.headers,\n timeout=60,\n )\n assert negative_exp_response.status_code == 422\n\n # Zero expiration should fail\n zero_exp_response = requests.post(\n f\"{API_SERVER_URL}/user/pats\",\n json={\"name\": \"Test Token\", \"expiration_days\": 0},\n headers=user.headers,\n timeout=60,\n )\n assert zero_exp_response.status_code == 422\n\n # Max length name (100 chars) should succeed\n valid_name = \"a\" * 100\n valid_pat = PATManager.create(\n name=valid_name,\n expiration_days=7,\n user_performing_action=user,\n )\n assert valid_pat.id is not None\n\n # Missing name should fail\n missing_name_response = requests.post(\n f\"{API_SERVER_URL}/user/pats\",\n json={\"expiration_days\": 30},\n headers=user.headers,\n timeout=60,\n )\n assert missing_name_response.status_code == 422\n\n\ndef test_pat_sorting_and_last_used(reset: None) -> None: # noqa: ARG001\n \"\"\"PATs are sorted by created_at DESC; last_used_at updates after authentication.\"\"\"\n user: DATestUser = UserManager.create(name=\"sorting_user\")\n\n # Create tokens with small delays to ensure different timestamps\n token1 = PATManager.create(\n name=\"First Token\",\n expiration_days=30,\n user_performing_action=user,\n )\n\n time.sleep(0.1)\n\n PATManager.create(\n name=\"Second Token\",\n expiration_days=30,\n user_performing_action=user,\n )\n\n time.sleep(0.1)\n\n PATManager.create(\n name=\"Third Token\",\n expiration_days=30,\n user_performing_action=user,\n )\n\n # Verify sorted by created_at DESC (newest first)\n tokens = PATManager.list(user)\n assert len(tokens) == 3\n\n assert tokens[0].name == \"Third Token\"\n assert tokens[1].name == \"Second Token\"\n assert tokens[2].name == \"First Token\"\n\n # Verify all tokens have no last_used_at initially\n for token in tokens:\n assert token.last_used_at is None\n\n # Use the first token (oldest)\n assert token1.token is not None\n auth_response = PATManager.authenticate(token1.token)\n assert auth_response.status_code == 200\n\n time.sleep(0.5)\n\n # Verify last_used_at is updated for the used token only\n tokens_after_use = PATManager.list(user)\n\n token1_after_use = next(t for t in tokens_after_use if t.name == \"First Token\")\n assert token1_after_use.last_used_at is not None\n\n token2_after_use = next(t for t in tokens_after_use if t.name == \"Second Token\")\n token3_after_use = next(t for t in tokens_after_use if t.name == \"Third Token\")\n assert token2_after_use.last_used_at is None\n assert token3_after_use.last_used_at is None\n\n\ndef test_pat_role_based_access_control(reset: None) -> None: # noqa: ARG001\n \"\"\"\n PATs inherit user roles and permissions:\n - Admin PAT: Full access to admin-only endpoints\n - Curator/Global Curator PATs: Access to management endpoints\n - Basic PAT: Denied access to admin and management endpoints\n \"\"\"\n # Create users with different roles\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n assert admin_user.role == UserRole.ADMIN\n\n basic_user: DATestUser = UserManager.create(name=\"basic_user\")\n assert basic_user.role == UserRole.BASIC\n\n curator_user: DATestUser = UserManager.create(name=\"curator_user\")\n curator_user = UserManager.set_role(\n user_to_set=curator_user,\n target_role=UserRole.CURATOR,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert curator_user.role == UserRole.CURATOR\n\n global_curator_user: DATestUser = UserManager.create(name=\"global_curator_user\")\n global_curator_user = UserManager.set_role(\n user_to_set=global_curator_user,\n target_role=UserRole.GLOBAL_CURATOR,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert global_curator_user.role == UserRole.GLOBAL_CURATOR\n\n # Create PATs for each user\n admin_pat = PATManager.create(\n name=\"Admin Token\",\n expiration_days=7,\n user_performing_action=admin_user,\n )\n\n basic_pat = PATManager.create(\n name=\"Basic Token\",\n expiration_days=7,\n user_performing_action=basic_user,\n )\n\n curator_pat = PATManager.create(\n name=\"Curator Token\",\n expiration_days=7,\n user_performing_action=curator_user,\n )\n\n global_curator_pat = PATManager.create(\n name=\"Global Curator Token\",\n expiration_days=7,\n user_performing_action=global_curator_user,\n )\n\n # Verify all tokens are present (type narrowing for mypy)\n assert admin_pat.token is not None\n assert basic_pat.token is not None\n assert curator_pat.token is not None\n assert global_curator_pat.token is not None\n\n # Test admin-only endpoint access\n print(\"\\n[Test] Admin PAT accessing admin-only endpoint...\")\n admin_endpoint_response = requests.get(\n f\"{API_SERVER_URL}/admin/api-key\",\n headers=PATManager.get_auth_headers(admin_pat.token),\n timeout=60,\n )\n assert admin_endpoint_response.status_code == 200\n print(\"[✓] Admin PAT successfully accessed /admin/api-key\")\n\n print(\"\\n[Test] Basic PAT accessing admin endpoint...\")\n basic_admin_response = requests.get(\n f\"{API_SERVER_URL}/admin/api-key\",\n headers=PATManager.get_auth_headers(basic_pat.token),\n timeout=60,\n )\n assert basic_admin_response.status_code == 403\n print(\"[✓] Basic PAT correctly denied access (403) to /admin/api-key\")\n\n print(\"\\n[Test] Curator PAT accessing admin-only endpoint...\")\n curator_admin_response = requests.get(\n f\"{API_SERVER_URL}/admin/api-key\",\n headers=PATManager.get_auth_headers(curator_pat.token),\n timeout=60,\n )\n assert curator_admin_response.status_code == 403\n print(\"[✓] Curator PAT correctly denied access (403) to /admin/api-key\")\n\n print(\"\\n[Test] Global Curator PAT accessing admin-only endpoint...\")\n global_curator_admin_response = requests.get(\n f\"{API_SERVER_URL}/admin/api-key\",\n headers=PATManager.get_auth_headers(global_curator_pat.token),\n timeout=60,\n )\n assert global_curator_admin_response.status_code == 403\n print(\"[✓] Global Curator PAT correctly denied access (403) to /admin/api-key\")\n\n # Test management endpoint access\n print(\"\\n[Test] Testing management endpoint access for curators...\")\n\n admin_manage_response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/connector\",\n headers=PATManager.get_auth_headers(admin_pat.token),\n timeout=60,\n )\n assert admin_manage_response.status_code == 200\n print(\"[✓] Admin PAT can access /manage/admin/connector\")\n\n curator_manage_response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/connector\",\n headers=PATManager.get_auth_headers(curator_pat.token),\n timeout=60,\n )\n assert curator_manage_response.status_code == 200\n print(\"[✓] Curator PAT can access /manage/admin/connector\")\n\n global_curator_manage_response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/connector\",\n headers=PATManager.get_auth_headers(global_curator_pat.token),\n timeout=60,\n )\n assert global_curator_manage_response.status_code == 200\n print(\"[✓] Global Curator PAT can access /manage/admin/connector\")\n\n basic_manage_response = requests.get(\n f\"{API_SERVER_URL}/manage/admin/connector\",\n headers=PATManager.get_auth_headers(basic_pat.token),\n timeout=60,\n )\n assert basic_manage_response.status_code in [403, 401]\n print(\n f\"[✓] Basic PAT correctly denied access ({basic_manage_response.status_code}) to /manage/admin/connector\"\n )\n\n # Verify PATs authenticate with correct identity and role\n print(\"\\n[Test] Verifying PATs authenticate as correct users with correct roles...\")\n\n admin_me = PATManager.authenticate(admin_pat.token)\n assert admin_me.status_code == 200\n assert admin_me.json()[\"email\"] == admin_user.email\n assert admin_me.json()[\"role\"] == UserRole.ADMIN.value\n\n basic_me = PATManager.authenticate(basic_pat.token)\n assert basic_me.status_code == 200\n assert basic_me.json()[\"email\"] == basic_user.email\n assert basic_me.json()[\"role\"] == UserRole.BASIC.value\n\n curator_me = PATManager.authenticate(curator_pat.token)\n assert curator_me.status_code == 200\n assert curator_me.json()[\"email\"] == curator_user.email\n assert curator_me.json()[\"role\"] == UserRole.CURATOR.value\n\n global_curator_me = PATManager.authenticate(global_curator_pat.token)\n assert global_curator_me.status_code == 200\n assert global_curator_me.json()[\"email\"] == global_curator_user.email\n assert global_curator_me.json()[\"role\"] == UserRole.GLOBAL_CURATOR.value\n\n print(\"[✓] All PATs authenticate with correct user identity and role\")\n\n # Verify all PATs can access basic endpoints\n print(\"\\n[Test] All PATs can access basic endpoints...\")\n for pat, user_name in [\n (admin_pat, \"Admin\"),\n (basic_pat, \"Basic\"),\n (curator_pat, \"Curator\"),\n (global_curator_pat, \"Global Curator\"),\n ]:\n assert pat.token is not None\n persona_response = requests.get(\n f\"{API_SERVER_URL}/persona\",\n headers=PATManager.get_auth_headers(pat.token),\n timeout=60,\n )\n assert persona_response.status_code == 200\n print(f\"[✓] {user_name} PAT can access /persona endpoint\")\n\n print(\"\\n[✓] All role-based access control tests passed!\")\n print(\"Summary:\")\n print(\n \" - Admin PAT: Full access to admin-only endpoints (/admin/*, /manage/admin/*)\"\n )\n print(\n \" - Curator PAT: Access to management endpoints (/manage/admin/*), denied on admin-only (/admin/*)\"\n )\n print(\n \" - Global Curator PAT: Access to management endpoints (/manage/admin/*), denied on admin-only (/admin/*)\"\n )\n print(\" - Basic PAT: Denied access to admin and management endpoints\")\n print(\" - All PATs: Can access basic endpoints (/persona, /me, etc.)\")\n print(\" - All PATs: Authenticate with correct user identity and role\")\n" }, { "path": "backend/tests/integration/tests/permissions/test_auth_permission_propagation.py", "content": "\"\"\"Integration tests for permission propagation across auth-triggered group changes.\n\nThese tests verify that effective permissions (via /me/permissions) actually\npropagate when users are added/removed from default groups through role changes.\nCustom permission grant tests will be added once the permission grant API is built.\n\"\"\"\n\nimport os\n\nimport pytest\n\nfrom onyx.auth.schemas import UserRole\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _get_basic_group_member_emails(admin_user: DATestUser) -> set[str]:\n all_groups = UserGroupManager.get_all(admin_user, include_default=True)\n basic_group = next(\n (g for g in all_groups if g.is_default and g.name == \"Basic\"), None\n )\n assert basic_group is not None, \"Basic default group not found\"\n return {u.email for u in basic_group.users}\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission propagation tests require enterprise features\",\n)\ndef test_basic_permission_granted_on_registration(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"New users should get 'basic' permission through default group assignment.\"\"\"\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n basic_user: DATestUser = UserManager.create(email=\"basic@example.com\")\n\n # Admin should have permissions from Admin group\n admin_perms = UserManager.get_permissions(admin_user)\n assert \"basic\" in admin_perms\n\n # Basic user should have 'basic' from Basic default group\n basic_perms = UserManager.get_permissions(basic_user)\n assert \"basic\" in basic_perms\n\n # Verify group membership matches\n assert basic_user.email in _get_basic_group_member_emails(admin_user)\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Permission propagation tests require enterprise features\",\n)\ndef test_role_downgrade_removes_basic_group_and_permission(\n reset: None, # noqa: ARG001\n) -> None:\n \"\"\"Downgrading to EXT_PERM_USER or SLACK_USER should remove from Basic group.\"\"\"\n admin_user: DATestUser = UserManager.create(email=\"admin@example.com\")\n\n # --- EXT_PERM_USER ---\n ext_user: DATestUser = UserManager.create(email=\"ext@example.com\")\n assert ext_user.email in _get_basic_group_member_emails(admin_user)\n\n UserManager.set_role(\n user_to_set=ext_user,\n target_role=UserRole.EXT_PERM_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert ext_user.email not in _get_basic_group_member_emails(admin_user)\n\n # --- SLACK_USER ---\n slack_user: DATestUser = UserManager.create(email=\"slack@example.com\")\n assert slack_user.email in _get_basic_group_member_emails(admin_user)\n\n UserManager.set_role(\n user_to_set=slack_user,\n target_role=UserRole.SLACK_USER,\n user_performing_action=admin_user,\n explicit_override=True,\n )\n assert slack_user.email not in _get_basic_group_member_emails(admin_user)\n" }, { "path": "backend/tests/integration/tests/permissions/test_cc_pair_permissions.py", "content": "\"\"\"\nThis file takes the happy path to adding a curator to a user group and then tests\nthe permissions of the curator manipulating connector-credential pairs.\n\"\"\"\n\nimport os\n\nimport pytest\nfrom onyx_openapi_client.exceptions import ApiException # type: ignore[import-untyped,unused-ignore,import-not-found]\n\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and User Group tests are enterprise only\",\n)\ndef test_cc_pair_permissions(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Creating a curator\n curator: DATestUser = UserManager.create(name=\"curator\")\n\n # Creating a user group\n user_group_1 = UserGroupManager.create(\n name=\"curated_user_group\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n # setting the user as a curator for the user group\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n\n # Creating another user group that the user is not a curator of\n user_group_2 = UserGroupManager.create(\n name=\"uncurated_user_group\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n connector_1 = ConnectorManager.create(\n name=\"admin_owned_connector\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n access_type=AccessType.PRIVATE,\n user_performing_action=admin_user,\n )\n # currently we dont enforce permissions at the connector level\n # pending cc_pair -> connector rework\n # connector_2 = ConnectorManager.create(\n # name=\"curator_visible_connector\",\n # source=DocumentSource.CONFLUENCE,\n # groups=[user_group_2.id],\n # is_public=False,\n # user_performing_action=admin_user,\n # )\n # Create a credentials that the curator is and is not curator of\n credential_1 = CredentialManager.create(\n name=\"curator_owned_credential\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n curator_public=False,\n user_performing_action=admin_user,\n )\n credential_2 = CredentialManager.create(\n name=\"curator_visible_credential\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_2.id],\n curator_public=False,\n user_performing_action=admin_user,\n )\n\n # END OF HAPPY PATH\n\n \"\"\"Tests for things Curators should not be able to do\"\"\"\n\n # Curators should not be able to create a cc\n # pair for a user group they are not a curator of\n with pytest.raises(ApiException):\n CCPairManager.create(\n connector_id=connector_1.id,\n credential_id=credential_1.id,\n name=\"invalid_cc_pair_2\",\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id, user_group_2.id],\n user_performing_action=curator,\n )\n\n # Curators should not be able to create a cc\n # pair without an attached user group\n with pytest.raises(ApiException):\n CCPairManager.create(\n connector_id=connector_1.id,\n credential_id=credential_1.id,\n name=\"invalid_cc_pair_2\",\n access_type=AccessType.PRIVATE,\n groups=[],\n user_performing_action=curator,\n )\n\n # # This test is currently disabled because permissions are\n # # not enforced at the connector level\n # # Curators should not be able to create a cc pair\n # # for a user group that the connector does not belong to (NOT WORKING)\n # with pytest.raises(HTTPError):\n # CCPairManager.create(\n # connector_id=connector_2.id,\n # credential_id=credential_1.id,\n # name=\"invalid_cc_pair_3\",\n # access_type=AccessType.PRIVATE,\n # groups=[user_group_1.id],\n # user_performing_action=curator,\n # )\n\n # Curators should not be able to create a cc\n # pair for a user group that the credential does not belong to\n with pytest.raises(ApiException):\n CCPairManager.create(\n connector_id=connector_1.id,\n credential_id=credential_2.id,\n name=\"invalid_cc_pair_4\",\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n \"\"\"Tests for things Curators should be able to do\"\"\"\n\n # Re-create connector since the credential_2 validation error above\n # triggers connector deletion in the exception handler\n connector_1 = ConnectorManager.create(\n name=\"admin_owned_connector_2\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n access_type=AccessType.PRIVATE,\n user_performing_action=admin_user,\n )\n\n # Curators should be able to create a private\n # cc pair for a user group they are a curator of\n valid_cc_pair = CCPairManager.create(\n name=\"valid_cc_pair\",\n connector_id=connector_1.id,\n credential_id=credential_1.id,\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n # Verify the created cc pair\n CCPairManager.verify(\n cc_pair=valid_cc_pair,\n user_performing_action=curator,\n )\n\n # Test pausing the cc pair\n CCPairManager.pause_cc_pair(valid_cc_pair, user_performing_action=curator)\n\n # Test deleting the cc pair\n CCPairManager.delete(valid_cc_pair, user_performing_action=curator)\n CCPairManager.wait_for_deletion_completion(\n cc_pair_id=valid_cc_pair.id, user_performing_action=curator\n )\n\n CCPairManager.verify(\n cc_pair=valid_cc_pair,\n verify_deleted=True,\n user_performing_action=curator,\n )\n" }, { "path": "backend/tests/integration/tests/permissions/test_connector_permissions.py", "content": "\"\"\"\nThis file takes the happy path to adding a curator to a user group and then tests\nthe permissions of the curator manipulating connectors.\n\"\"\"\n\nimport os\n\nimport pytest\nfrom requests.exceptions import HTTPError\n\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and user group tests are enterprise only\",\n)\ndef test_connector_permissions(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Creating a curator\n curator: DATestUser = UserManager.create(name=\"curator\")\n\n # Creating a user group\n user_group_1 = UserGroupManager.create(\n name=\"user_group_1\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n # setting the user as a curator for the user group\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n\n # Creating another user group that the user is not a curator of\n user_group_2 = UserGroupManager.create(\n name=\"user_group_2\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # END OF HAPPY PATH\n\n \"\"\"Tests for things Curators should not be able to do\"\"\"\n\n # Curators should not be able to create a connector for a\n # user group they are not a curator of\n with pytest.raises(HTTPError):\n ConnectorManager.create(\n name=\"invalid_connector_2\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id, user_group_2.id],\n access_type=AccessType.PRIVATE,\n user_performing_action=curator,\n )\n\n \"\"\"Tests for things Curators should be able to do\"\"\"\n\n # Curators should be able to create a private\n # connector for a user group they are a curator of\n valid_connector = ConnectorManager.create(\n name=\"valid_connector\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n access_type=AccessType.PRIVATE,\n user_performing_action=curator,\n )\n assert valid_connector.id is not None\n\n # Verify the created connector\n created_connector = ConnectorManager.get(\n valid_connector.id, user_performing_action=curator\n )\n assert created_connector.name == valid_connector.name\n assert created_connector.source == valid_connector.source\n\n # Verify that the connector can be found in the list of all connectors\n all_connectors = ConnectorManager.get_all(user_performing_action=curator)\n assert any(conn.id == valid_connector.id for conn in all_connectors)\n\n # Test editing the connector\n valid_connector.name = \"updated_valid_connector\"\n ConnectorManager.edit(valid_connector, user_performing_action=curator)\n\n # Verify the edit\n updated_connector = ConnectorManager.get(\n valid_connector.id, user_performing_action=curator\n )\n assert updated_connector.name == \"updated_valid_connector\"\n\n # Test deleting the connector\n ConnectorManager.delete(connector=valid_connector, user_performing_action=curator)\n\n # Verify the deletion\n all_connectors_after_delete = ConnectorManager.get_all(\n user_performing_action=curator\n )\n assert all(conn.id != valid_connector.id for conn in all_connectors_after_delete)\n\n # Test that curator cannot create a connector for a group they are not a curator of\n with pytest.raises(HTTPError):\n ConnectorManager.create(\n name=\"invalid_connector_3\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_2.id],\n access_type=AccessType.PRIVATE,\n user_performing_action=curator,\n )\n\n # Curators should be able to create a public connector\n public_connector = ConnectorManager.create(\n name=\"curator_public_connector\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n access_type=AccessType.PUBLIC,\n user_performing_action=curator,\n )\n assert public_connector.id is not None\n" }, { "path": "backend/tests/integration/tests/permissions/test_credential_permissions.py", "content": "\"\"\"\nThis file takes the happy path to adding a curator to a user group and then tests\nthe permissions of the curator manipulating credentials.\n\"\"\"\n\nimport os\n\nimport pytest\nfrom requests.exceptions import HTTPError\n\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and user group tests are enterprise only\",\n)\ndef test_credential_permissions(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Creating a curator\n curator: DATestUser = UserManager.create(name=\"curator\")\n\n # Creating a user group\n user_group_1 = UserGroupManager.create(\n name=\"user_group_1\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n # setting the user as a curator for the user group\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n\n # Creating another user group that the user is not a curator of\n user_group_2 = UserGroupManager.create(\n name=\"user_group_2\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # END OF HAPPY PATH\n\n \"\"\"Tests for things Curators should not be able to do\"\"\"\n\n # Curators should not be able to create a credential for a user group they are not a curator of\n with pytest.raises(HTTPError):\n CredentialManager.create(\n name=\"invalid_credential_2\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id, user_group_2.id],\n curator_public=False,\n user_performing_action=curator,\n )\n\n \"\"\"Tests for things Curators should be able to do\"\"\"\n # Curators should be able to create a private credential for a user group they are a curator of\n valid_credential = CredentialManager.create(\n name=\"valid_credential\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n curator_public=False,\n user_performing_action=curator,\n )\n\n # Verify the created credential\n CredentialManager.verify(\n credential=valid_credential,\n user_performing_action=curator,\n )\n\n # Test editing the credential\n valid_credential.name = \"updated_valid_credential\"\n CredentialManager.edit(valid_credential, user_performing_action=curator)\n\n # Verify the edit\n CredentialManager.verify(\n credential=valid_credential,\n user_performing_action=curator,\n )\n\n # Test deleting the credential\n CredentialManager.delete(valid_credential, user_performing_action=curator)\n\n # Verify the deletion\n CredentialManager.verify(\n credential=valid_credential,\n verify_deleted=True,\n user_performing_action=curator,\n )\n\n # Curators should be able to create a public credential\n public_credential = CredentialManager.create(\n name=\"curator_public_credential\",\n source=DocumentSource.CONFLUENCE,\n groups=[user_group_1.id],\n curator_public=True,\n user_performing_action=curator,\n )\n CredentialManager.verify(\n credential=public_credential,\n user_performing_action=curator,\n )\n" }, { "path": "backend/tests/integration/tests/permissions/test_doc_set_permissions.py", "content": "import os\n\nimport pytest\nfrom requests.exceptions import HTTPError\n\nfrom onyx.db.enums import AccessType\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.document_set import DocumentSetManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and user group tests are enterprise only\",\n)\ndef test_doc_set_permissions_setup(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Creating a second user (curator)\n curator: DATestUser = UserManager.create(name=\"curator\")\n\n # Creating the first user group\n user_group_1 = UserGroupManager.create(\n name=\"curated_user_group\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # Setting the curator as a curator for the first user group\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n\n # Creating a second user group\n user_group_2 = UserGroupManager.create(\n name=\"uncurated_user_group\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # Admin creates a cc_pair\n private_cc_pair = CCPairManager.create_from_scratch(\n access_type=AccessType.PRIVATE,\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n # Admin creates a public cc_pair\n public_cc_pair = CCPairManager.create_from_scratch(\n access_type=AccessType.PUBLIC,\n source=DocumentSource.INGESTION_API,\n user_performing_action=admin_user,\n )\n\n # END OF HAPPY PATH\n\n \"\"\"Tests for things Curators/Admins should not be able to do\"\"\"\n\n # Test that curator cannot create a non-public document set for the group they don't curate\n with pytest.raises(HTTPError):\n DocumentSetManager.create(\n name=\"Invalid Document Set 1\",\n is_public=False,\n groups=[user_group_2.id],\n cc_pair_ids=[public_cc_pair.id],\n user_performing_action=curator,\n )\n\n # Test that curator cannot create a document set attached to both groups\n with pytest.raises(HTTPError):\n DocumentSetManager.create(\n name=\"Invalid Document Set 2\",\n is_public=False,\n cc_pair_ids=[public_cc_pair.id],\n groups=[user_group_1.id, user_group_2.id],\n user_performing_action=curator,\n )\n\n # Test that curator cannot create a document set with no groups\n with pytest.raises(HTTPError):\n DocumentSetManager.create(\n name=\"Invalid Document Set 3\",\n is_public=False,\n cc_pair_ids=[public_cc_pair.id],\n groups=[],\n user_performing_action=curator,\n )\n\n # Test that curator cannot create a document set with no cc_pairs\n with pytest.raises(HTTPError):\n DocumentSetManager.create(\n name=\"Invalid Document Set 4\",\n is_public=False,\n cc_pair_ids=[],\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n # Test that admin cannot create a document set with no cc_pairs\n with pytest.raises(HTTPError):\n DocumentSetManager.create(\n name=\"Invalid Document Set 4\",\n is_public=False,\n cc_pair_ids=[],\n groups=[user_group_1.id],\n user_performing_action=admin_user,\n )\n\n \"\"\"Tests for things Curators should be able to do\"\"\"\n # Test that curator can create a document set for the group they curate\n valid_doc_set = DocumentSetManager.create(\n name=\"Valid Document Set\",\n is_public=False,\n cc_pair_ids=[public_cc_pair.id],\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n DocumentSetManager.wait_for_sync(\n document_sets_to_check=[valid_doc_set], user_performing_action=admin_user\n )\n\n # Verify that the valid document set was created\n DocumentSetManager.verify(\n document_set=valid_doc_set,\n user_performing_action=admin_user,\n )\n\n # Verify that only one document set exists\n all_doc_sets = DocumentSetManager.get_all(user_performing_action=admin_user)\n assert len(all_doc_sets) == 1\n\n # Add the private_cc_pair to the doc set on our end for later comparison\n valid_doc_set.cc_pair_ids.append(private_cc_pair.id)\n\n # Confirm the curator can't add the private_cc_pair to the doc set\n with pytest.raises(HTTPError):\n DocumentSetManager.edit(\n document_set=valid_doc_set,\n user_performing_action=curator,\n )\n # Confirm the admin can't add the private_cc_pair to the doc set\n with pytest.raises(HTTPError):\n DocumentSetManager.edit(\n document_set=valid_doc_set,\n user_performing_action=admin_user,\n )\n\n # Verify the document set has not been updated in the db\n with pytest.raises(ValueError):\n DocumentSetManager.verify(\n document_set=valid_doc_set,\n user_performing_action=admin_user,\n )\n\n # Add the private_cc_pair to the user group on our end for later comparison\n user_group_1.cc_pair_ids.append(private_cc_pair.id)\n\n # Admin adds the cc_pair to the group the curator curates\n UserGroupManager.edit(\n user_group=user_group_1,\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n UserGroupManager.verify(\n user_group=user_group_1,\n user_performing_action=admin_user,\n )\n\n # Confirm the curator can now add the cc_pair to the doc set\n DocumentSetManager.edit(\n document_set=valid_doc_set,\n user_performing_action=curator,\n )\n DocumentSetManager.wait_for_sync(\n document_sets_to_check=[valid_doc_set], user_performing_action=admin_user\n )\n # Verify the updated document set\n DocumentSetManager.verify(\n document_set=valid_doc_set,\n user_performing_action=admin_user,\n )\n" }, { "path": "backend/tests/integration/tests/permissions/test_file_connector_permissions.py", "content": "import io\nimport json\nimport os\n\nimport pytest\nimport requests\n\nfrom onyx.db.enums import AccessType\nfrom onyx.db.models import UserRole\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\ndef _upload_connector_file(\n *,\n user_performing_action: DATestUser,\n file_name: str,\n content: bytes,\n) -> tuple[str, str]:\n headers = user_performing_action.headers.copy()\n headers.pop(\"Content-Type\", None)\n\n response = requests.post(\n f\"{API_SERVER_URL}/manage/admin/connector/file/upload\",\n files=[(\"files\", (file_name, io.BytesIO(content), \"text/plain\"))],\n headers=headers,\n )\n response.raise_for_status()\n payload = response.json()\n return payload[\"file_paths\"][0], payload[\"file_names\"][0]\n\n\ndef _update_connector_files(\n *,\n connector_id: int,\n user_performing_action: DATestUser,\n file_ids_to_remove: list[str],\n new_file_name: str,\n new_file_content: bytes,\n) -> requests.Response:\n headers = user_performing_action.headers.copy()\n headers.pop(\"Content-Type\", None)\n\n return requests.post(\n f\"{API_SERVER_URL}/manage/admin/connector/{connector_id}/files/update\",\n data={\"file_ids_to_remove\": json.dumps(file_ids_to_remove)},\n files=[(\"files\", (new_file_name, io.BytesIO(new_file_content), \"text/plain\"))],\n headers=headers,\n )\n\n\ndef _list_connector_files(\n *,\n connector_id: int,\n user_performing_action: DATestUser,\n) -> requests.Response:\n return requests.get(\n f\"{API_SERVER_URL}/manage/admin/connector/{connector_id}/files\",\n headers=user_performing_action.headers,\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and user group tests are enterprise only\",\n)\n@pytest.mark.usefixtures(\"reset\")\ndef test_only_global_curator_can_update_public_file_connector_files() -> None:\n admin_user = UserManager.create(name=\"admin_user\")\n\n global_curator_creator = UserManager.create(name=\"global_curator_creator\")\n global_curator_creator = UserManager.set_role(\n user_to_set=global_curator_creator,\n target_role=UserRole.GLOBAL_CURATOR,\n user_performing_action=admin_user,\n )\n\n global_curator_editor = UserManager.create(name=\"global_curator_editor\")\n global_curator_editor = UserManager.set_role(\n user_to_set=global_curator_editor,\n target_role=UserRole.GLOBAL_CURATOR,\n user_performing_action=admin_user,\n )\n\n curator_user = UserManager.create(name=\"curator_user\")\n curator_group = UserGroupManager.create(\n name=\"curator_group\",\n user_ids=[curator_user.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[curator_group],\n user_performing_action=admin_user,\n )\n UserGroupManager.set_curator_status(\n test_user_group=curator_group,\n user_to_set_as_curator=curator_user,\n user_performing_action=admin_user,\n )\n\n initial_file_id, initial_file_name = _upload_connector_file(\n user_performing_action=global_curator_creator,\n file_name=\"initial-file.txt\",\n content=b\"initial file content\",\n )\n\n connector = ConnectorManager.create(\n user_performing_action=global_curator_creator,\n name=\"public_file_connector\",\n source=DocumentSource.FILE,\n connector_specific_config={\n \"file_locations\": [initial_file_id],\n \"file_names\": [initial_file_name],\n \"zip_metadata_file_id\": None,\n },\n access_type=AccessType.PUBLIC,\n groups=[],\n )\n credential = CredentialManager.create(\n user_performing_action=global_curator_creator,\n source=DocumentSource.FILE,\n curator_public=True,\n groups=[],\n name=\"public_file_connector_credential\",\n )\n CCPairManager.create(\n connector_id=connector.id,\n credential_id=credential.id,\n user_performing_action=global_curator_creator,\n access_type=AccessType.PUBLIC,\n groups=[],\n name=\"public_file_connector_cc_pair\",\n )\n\n curator_list_response = _list_connector_files(\n connector_id=connector.id,\n user_performing_action=curator_user,\n )\n curator_list_response.raise_for_status()\n curator_list_payload = curator_list_response.json()\n assert any(f[\"file_id\"] == initial_file_id for f in curator_list_payload[\"files\"])\n\n global_curator_list_response = _list_connector_files(\n connector_id=connector.id,\n user_performing_action=global_curator_editor,\n )\n global_curator_list_response.raise_for_status()\n global_curator_list_payload = global_curator_list_response.json()\n assert any(\n f[\"file_id\"] == initial_file_id for f in global_curator_list_payload[\"files\"]\n )\n\n denied_response = _update_connector_files(\n connector_id=connector.id,\n user_performing_action=curator_user,\n file_ids_to_remove=[initial_file_id],\n new_file_name=\"curator-file.txt\",\n new_file_content=b\"curator updated file\",\n )\n assert denied_response.status_code == 403\n\n allowed_response = _update_connector_files(\n connector_id=connector.id,\n user_performing_action=global_curator_editor,\n file_ids_to_remove=[initial_file_id],\n new_file_name=\"global-curator-file.txt\",\n new_file_content=b\"global curator updated file\",\n )\n allowed_response.raise_for_status()\n\n payload = allowed_response.json()\n assert initial_file_id not in payload[\"file_paths\"]\n assert \"global-curator-file.txt\" in payload[\"file_names\"]\n\n creator_group = UserGroupManager.create(\n name=\"creator_group\",\n user_ids=[global_curator_creator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[creator_group],\n user_performing_action=admin_user,\n )\n\n private_file_id, private_file_name = _upload_connector_file(\n user_performing_action=global_curator_creator,\n file_name=\"private-initial-file.txt\",\n content=b\"private initial file content\",\n )\n\n private_connector = ConnectorManager.create(\n user_performing_action=global_curator_creator,\n name=\"private_file_connector\",\n source=DocumentSource.FILE,\n connector_specific_config={\n \"file_locations\": [private_file_id],\n \"file_names\": [private_file_name],\n \"zip_metadata_file_id\": None,\n },\n access_type=AccessType.PRIVATE,\n groups=[creator_group.id],\n )\n private_credential = CredentialManager.create(\n user_performing_action=global_curator_creator,\n source=DocumentSource.FILE,\n curator_public=False,\n groups=[creator_group.id],\n name=\"private_file_connector_credential\",\n )\n CCPairManager.create(\n connector_id=private_connector.id,\n credential_id=private_credential.id,\n user_performing_action=global_curator_creator,\n access_type=AccessType.PRIVATE,\n groups=[creator_group.id],\n name=\"private_file_connector_cc_pair\",\n )\n\n private_denied_response = _update_connector_files(\n connector_id=private_connector.id,\n user_performing_action=global_curator_editor,\n file_ids_to_remove=[private_file_id],\n new_file_name=\"global-curator-private-file.txt\",\n new_file_content=b\"global curator private update\",\n )\n assert private_denied_response.status_code == 403\n" }, { "path": "backend/tests/integration/tests/permissions/test_persona_permissions.py", "content": "\"\"\"\nThis file tests the permissions for creating and editing personas for different user roles:\n- Basic users can create personas and edit their own\n- Curators can edit personas that belong exclusively to groups they curate\n- Admins can edit all personas\n\"\"\"\n\nimport os\n\nimport pytest\nfrom requests.exceptions import HTTPError\n\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and user group tests are enterprise only\",\n)\ndef test_persona_permissions(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Creating a curator user\n curator: DATestUser = UserManager.create(name=\"curator\")\n\n # Creating a basic user\n basic_user: DATestUser = UserManager.create(name=\"basic_user\")\n\n # Creating user groups\n user_group_1 = UserGroupManager.create(\n name=\"curated_user_group\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n # Setting the user as a curator for the user group\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n\n # Creating another user group that the user is not a curator of\n user_group_2 = UserGroupManager.create(\n name=\"uncurated_user_group\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_2], user_performing_action=admin_user\n )\n\n \"\"\"Test that any user can create a persona\"\"\"\n # Basic user creates a persona\n basic_user_persona = PersonaManager.create(\n name=\"basic_user_persona\",\n description=\"A persona created by basic user\",\n is_public=False,\n groups=[],\n users=[admin_user.id],\n user_performing_action=basic_user,\n )\n PersonaManager.verify(basic_user_persona, user_performing_action=basic_user)\n\n # Curator creates a persona\n curator_persona = PersonaManager.create(\n name=\"curator_persona\",\n description=\"A persona created by curator\",\n is_public=False,\n groups=[],\n user_performing_action=curator,\n )\n PersonaManager.verify(curator_persona, user_performing_action=curator)\n\n # Admin creates personas for different groups\n admin_persona_group_1 = PersonaManager.create(\n name=\"admin_persona_group_1\",\n description=\"A persona for group 1\",\n is_public=False,\n groups=[user_group_1.id],\n user_performing_action=admin_user,\n )\n admin_persona_group_2 = PersonaManager.create(\n name=\"admin_persona_group_2\",\n description=\"A persona for group 2\",\n is_public=False,\n groups=[user_group_2.id],\n user_performing_action=admin_user,\n )\n admin_persona_both_groups = PersonaManager.create(\n name=\"admin_persona_both_groups\",\n description=\"A persona for both groups\",\n is_public=False,\n groups=[user_group_1.id, user_group_2.id],\n user_performing_action=admin_user,\n )\n\n \"\"\"Test that users can edit their own personas\"\"\"\n # Basic user can edit their own persona\n PersonaManager.edit(\n persona=basic_user_persona,\n description=\"Updated description by basic user\",\n user_performing_action=basic_user,\n )\n PersonaManager.verify(basic_user_persona, user_performing_action=basic_user)\n\n # Basic user cannot edit other's personas\n with pytest.raises(HTTPError):\n PersonaManager.edit(\n persona=curator_persona,\n description=\"Invalid edit by basic user\",\n user_performing_action=basic_user,\n )\n\n \"\"\"Test curator permissions\"\"\"\n # Curator can edit personas that belong exclusively to groups they curate\n PersonaManager.edit(\n persona=admin_persona_group_1,\n description=\"Updated by curator\",\n user_performing_action=curator,\n )\n PersonaManager.verify(admin_persona_group_1, user_performing_action=curator)\n\n # Curator cannot edit personas in groups they don't curate\n with pytest.raises(HTTPError):\n PersonaManager.edit(\n persona=admin_persona_group_2,\n description=\"Invalid edit by curator\",\n user_performing_action=curator,\n )\n\n # Curator cannot edit personas that belong to multiple groups, even if they curate one\n with pytest.raises(HTTPError):\n PersonaManager.edit(\n persona=admin_persona_both_groups,\n description=\"Invalid edit by curator\",\n user_performing_action=curator,\n )\n\n \"\"\"Test admin permissions\"\"\"\n # Admin can edit any persona\n\n # the persona was shared with the admin user on creation\n # this edit call will simulate having the same user in the list twice.\n # The server side should dedupe and handle this correctly (prior bug)\n PersonaManager.edit(\n persona=basic_user_persona,\n description=\"Updated by admin 2\",\n users=[admin_user.id, admin_user.id],\n user_performing_action=admin_user,\n )\n PersonaManager.verify(basic_user_persona, user_performing_action=admin_user)\n\n PersonaManager.edit(\n persona=curator_persona,\n description=\"Updated by admin\",\n user_performing_action=admin_user,\n )\n PersonaManager.verify(curator_persona, user_performing_action=admin_user)\n\n PersonaManager.edit(\n persona=admin_persona_group_1,\n description=\"Updated by admin\",\n user_performing_action=admin_user,\n )\n PersonaManager.verify(admin_persona_group_1, user_performing_action=admin_user)\n\n PersonaManager.edit(\n persona=admin_persona_group_2,\n description=\"Updated by admin\",\n user_performing_action=admin_user,\n )\n PersonaManager.verify(admin_persona_group_2, user_performing_action=admin_user)\n\n PersonaManager.edit(\n persona=admin_persona_both_groups,\n description=\"Updated by admin\",\n user_performing_action=admin_user,\n )\n PersonaManager.verify(admin_persona_both_groups, user_performing_action=admin_user)\n" }, { "path": "backend/tests/integration/tests/permissions/test_user_file_permissions.py", "content": "\"\"\"\nThis file tests user file permissions in different scenarios:\n1. Public assistant with user files - files should be accessible to all users\n2. Direct file access - user files should NOT be accessible by users who don't own them\n\"\"\"\n\nimport io\nfrom typing import NamedTuple\n\nimport pytest\n\nfrom onyx.file_store.models import FileDescriptor\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.file import FileManager\nfrom tests.integration.common_utils.managers.llm_provider import LLMProviderManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestPersona\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\nclass UserFileTestSetup(NamedTuple):\n admin_user: DATestUser\n user1_file_owner: DATestUser\n user2_non_owner: DATestUser\n user1_file_descriptor: FileDescriptor\n user1_file_id: str\n public_assistant: DATestPersona\n\n\n@pytest.fixture\ndef user_file_setup(reset: None) -> UserFileTestSetup: # noqa: ARG001\n \"\"\"\n Common setup for user file permission tests.\n Creates users, files, and a public assistant with files.\n \"\"\"\n # Create an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # Create LLM provider for chat functionality\n LLMProviderManager.create(user_performing_action=admin_user)\n\n # Create user1 who will own the file\n user1: DATestUser = UserManager.create(name=\"user1_file_owner\")\n\n # Create user2 who will use the assistant but doesn't own the file\n user2: DATestUser = UserManager.create(name=\"user2_non_owner\")\n\n # Create a test file and upload as user1\n test_file_content = b\"This is test content for user file permission checking.\"\n test_file = (\"test_file.txt\", io.BytesIO(test_file_content))\n\n file_descriptors, error = FileManager.upload_files(\n files=[test_file],\n user_performing_action=user1,\n )\n\n assert not error, f\"Failed to upload file: {error}\"\n assert len(file_descriptors) == 1, \"Expected 1 file to be uploaded\"\n\n # Get the file descriptor and user_file_id\n user1_file_descriptor = file_descriptors[0]\n user_file_id = user1_file_descriptor.get(\"user_file_id\")\n\n assert user_file_id is not None, \"user_file_id should not be None\"\n\n # Create a public assistant with the user file attached\n public_assistant = PersonaManager.create(\n name=\"Public Assistant with Files\",\n description=\"A public assistant with user files for testing permissions\",\n is_public=True,\n user_file_ids=[user_file_id],\n user_performing_action=admin_user,\n )\n\n return UserFileTestSetup(\n admin_user=admin_user,\n user1_file_owner=user1,\n user2_non_owner=user2,\n user1_file_descriptor=user1_file_descriptor,\n user1_file_id=user_file_id,\n public_assistant=public_assistant,\n )\n\n\ndef test_public_assistant_with_user_files(\n user_file_setup: UserFileTestSetup,\n) -> None:\n \"\"\"\n Test that a public assistant with user files attached can be used by users\n who don't own those files without permission errors.\n \"\"\"\n # Create a chat session with the public assistant as user2\n chat_session = ChatSessionManager.create(\n persona_id=user_file_setup.public_assistant.id,\n description=\"Test chat session for user file permissions\",\n user_performing_action=user_file_setup.user2_non_owner,\n )\n\n # Send a message as user2 - this should not throw a permission error\n # even though user2 doesn't own the file attached to the assistant\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"Hello, can you help me?\",\n user_performing_action=user_file_setup.user2_non_owner,\n )\n\n # Verify the message was processed without errors\n assert (\n response.error is None\n ), f\"Expected no error when user2 uses public assistant with user1's files, but got error: {response.error}\"\n assert len(response.full_message) > 0, \"Expected a response from the assistant\"\n\n # Verify chat history is accessible\n chat_history = ChatSessionManager.get_chat_history(\n chat_session=chat_session,\n user_performing_action=user_file_setup.user2_non_owner,\n )\n assert (\n len(chat_history) >= 2\n ), \"Expected at least 2 messages (user message and assistant response)\"\n" }, { "path": "backend/tests/integration/tests/permissions/test_user_role_permissions.py", "content": "\"\"\"\nThis file tests the ability of different user types to set the role of other users.\n\"\"\"\n\nimport os\n\nimport pytest\nfrom requests.exceptions import HTTPError\n\nfrom onyx.db.models import UserRole\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator and user group tests are enterprise only\",\n)\ndef test_user_role_setting_permissions(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n assert UserManager.is_role(admin_user, UserRole.ADMIN)\n\n # Creating a basic user\n basic_user: DATestUser = UserManager.create(name=\"basic_user\")\n assert UserManager.is_role(basic_user, UserRole.BASIC)\n\n # Creating a curator\n curator: DATestUser = UserManager.create(name=\"curator\")\n assert UserManager.is_role(curator, UserRole.BASIC)\n\n # Creating a curator without adding to a group should not work\n with pytest.raises(HTTPError):\n UserManager.set_role(\n user_to_set=curator,\n target_role=UserRole.CURATOR,\n user_performing_action=admin_user,\n )\n\n global_curator: DATestUser = UserManager.create(name=\"global_curator\")\n assert UserManager.is_role(global_curator, UserRole.BASIC)\n\n # Setting the role of a global curator should not work for a basic user\n with pytest.raises(HTTPError):\n UserManager.set_role(\n user_to_set=global_curator,\n target_role=UserRole.GLOBAL_CURATOR,\n user_performing_action=basic_user,\n )\n\n # Setting the role of a global curator should work for an admin user\n UserManager.set_role(\n user_to_set=global_curator,\n target_role=UserRole.GLOBAL_CURATOR,\n user_performing_action=admin_user,\n )\n assert UserManager.is_role(global_curator, UserRole.GLOBAL_CURATOR)\n\n # Setting the role of a global curator should not work for an invalid curator\n with pytest.raises(HTTPError):\n UserManager.set_role(\n user_to_set=global_curator,\n target_role=UserRole.BASIC,\n user_performing_action=global_curator,\n )\n assert UserManager.is_role(global_curator, UserRole.GLOBAL_CURATOR)\n\n # Creating a user group\n user_group_1 = UserGroupManager.create(\n name=\"user_group_1\",\n user_ids=[],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # This should fail because the curator is not in the user group\n with pytest.raises(HTTPError):\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n\n # Adding the curator to the user group\n user_group_1.user_ids = [curator.id]\n UserGroupManager.edit(user_group=user_group_1, user_performing_action=admin_user)\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # This should work because the curator is in the user group\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n" }, { "path": "backend/tests/integration/tests/permissions/test_whole_curator_flow.py", "content": "\"\"\"\nThis test tests the happy path for curator permissions\n\"\"\"\n\nimport os\n\nimport pytest\n\nfrom onyx.db.enums import AccessType\nfrom onyx.db.models import UserRole\nfrom onyx.server.documents.models import DocumentSource\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.connector import ConnectorManager\nfrom tests.integration.common_utils.managers.credential import CredentialManager\nfrom tests.integration.common_utils.managers.user import DATestUser\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.managers.user_group import UserGroupManager\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator tests are enterprise only\",\n)\ndef test_whole_curator_flow(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n assert UserManager.is_role(admin_user, UserRole.ADMIN)\n\n # Creating a curator\n curator: DATestUser = UserManager.create(name=\"curator\")\n\n # Creating a user group\n user_group_1 = UserGroupManager.create(\n name=\"user_group_1\",\n user_ids=[curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n # Making curator a curator of user_group_1\n UserGroupManager.set_curator_status(\n test_user_group=user_group_1,\n user_to_set_as_curator=curator,\n user_performing_action=admin_user,\n )\n assert UserManager.is_role(curator, UserRole.CURATOR)\n\n # Creating a credential as curator\n test_credential = CredentialManager.create(\n name=\"curator_test_credential\",\n source=DocumentSource.FILE,\n curator_public=False,\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n # Creating a connector as curator\n test_connector = ConnectorManager.create(\n name=\"curator_test_connector\",\n source=DocumentSource.FILE,\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n # Test editing the connector\n test_connector.name = \"updated_test_connector\"\n ConnectorManager.edit(connector=test_connector, user_performing_action=curator)\n\n # Creating a CC pair as curator\n test_cc_pair = CCPairManager.create(\n connector_id=test_connector.id,\n credential_id=test_credential.id,\n name=\"curator_test_cc_pair\",\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id],\n user_performing_action=curator,\n )\n\n CCPairManager.verify(cc_pair=test_cc_pair, user_performing_action=admin_user)\n\n # Verify that the curator can pause and unpause the CC pair\n CCPairManager.pause_cc_pair(cc_pair=test_cc_pair, user_performing_action=curator)\n\n # Verify that the curator can delete the CC pair\n CCPairManager.delete(cc_pair=test_cc_pair, user_performing_action=curator)\n CCPairManager.wait_for_deletion_completion(\n cc_pair_id=test_cc_pair.id, user_performing_action=curator\n )\n\n # Verify that the CC pair has been deleted\n CCPairManager.verify(\n cc_pair=test_cc_pair,\n verify_deleted=True,\n user_performing_action=admin_user,\n )\n\n\n@pytest.mark.skipif(\n os.environ.get(\"ENABLE_PAID_ENTERPRISE_EDITION_FEATURES\", \"\").lower() != \"true\",\n reason=\"Curator tests are enterprise only\",\n)\ndef test_global_curator_flow(reset: None) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n assert UserManager.is_role(admin_user, UserRole.ADMIN)\n\n # Creating a user\n global_curator: DATestUser = UserManager.create(name=\"global_curator\")\n assert UserManager.is_role(global_curator, UserRole.BASIC)\n\n # Set the user to a global curator\n UserManager.set_role(\n user_to_set=global_curator,\n target_role=UserRole.GLOBAL_CURATOR,\n user_performing_action=admin_user,\n )\n assert UserManager.is_role(global_curator, UserRole.GLOBAL_CURATOR)\n\n # Creating a user group containing the global curator\n user_group_1 = UserGroupManager.create(\n name=\"user_group_1\",\n user_ids=[global_curator.id],\n cc_pair_ids=[],\n user_performing_action=admin_user,\n )\n UserGroupManager.wait_for_sync(\n user_groups_to_check=[user_group_1], user_performing_action=admin_user\n )\n\n # Creating a credential as global curator\n test_credential = CredentialManager.create(\n name=\"curator_test_credential\",\n source=DocumentSource.FILE,\n curator_public=False,\n groups=[user_group_1.id],\n user_performing_action=global_curator,\n )\n\n # Creating a connector as global curator\n test_connector = ConnectorManager.create(\n name=\"curator_test_connector\",\n source=DocumentSource.FILE,\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id],\n user_performing_action=global_curator,\n )\n\n # Test editing the connector\n test_connector.name = \"updated_test_connector\"\n ConnectorManager.edit(\n connector=test_connector, user_performing_action=global_curator\n )\n\n # Creating a CC pair as global curator\n test_cc_pair = CCPairManager.create(\n connector_id=test_connector.id,\n credential_id=test_credential.id,\n name=\"curator_test_cc_pair\",\n access_type=AccessType.PRIVATE,\n groups=[user_group_1.id],\n user_performing_action=global_curator,\n )\n\n CCPairManager.verify(cc_pair=test_cc_pair, user_performing_action=admin_user)\n\n # Verify that the curator can pause and unpause the CC pair\n CCPairManager.pause_cc_pair(\n cc_pair=test_cc_pair, user_performing_action=global_curator\n )\n\n # Verify that the curator can delete the CC pair\n CCPairManager.delete(cc_pair=test_cc_pair, user_performing_action=global_curator)\n CCPairManager.wait_for_deletion_completion(\n cc_pair_id=test_cc_pair.id, user_performing_action=global_curator\n )\n\n # Verify that the CC pair has been deleted\n CCPairManager.verify(\n cc_pair=test_cc_pair,\n verify_deleted=True,\n user_performing_action=admin_user,\n )\n" }, { "path": "backend/tests/integration/tests/personalization/test_personalization_flow.py", "content": "import requests\n\nfrom onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _get_auth_headers(user: DATestUser) -> tuple[dict, dict]:\n return user.headers, {\n FASTAPI_USERS_AUTH_COOKIE_NAME: user.cookies[FASTAPI_USERS_AUTH_COOKIE_NAME]\n }\n\n\ndef _get_me(headers: dict, cookies: dict) -> dict:\n response = requests.get(f\"{API_SERVER_URL}/me\", headers=headers, cookies=cookies)\n response.raise_for_status()\n return response.json()\n\n\ndef _patch_personalization(headers: dict, cookies: dict, payload: dict) -> None:\n response = requests.patch(\n f\"{API_SERVER_URL}/user/personalization\",\n json=payload,\n headers=headers,\n cookies=cookies,\n )\n response.raise_for_status()\n\n\ndef test_personalization_round_trip(reset: None) -> None: # noqa: ARG001\n user = UserManager.create()\n headers, cookies = _get_auth_headers(user)\n\n # baseline should have empty personalization\n me_initial = _get_me(headers, cookies)\n assert me_initial[\"personalization\"][\"name\"] == \"\"\n assert me_initial[\"personalization\"][\"role\"] == \"\"\n assert me_initial[\"personalization\"][\"use_memories\"] is True\n assert me_initial[\"personalization\"][\"enable_memory_tool\"] is True\n assert me_initial[\"personalization\"][\"memories\"] == []\n\n payload = {\n \"name\": \"Jane Doe\",\n \"role\": \"Developer advocate\",\n \"use_memories\": True,\n \"memories\": [\n {\"content\": \"Loves peanut butter\"},\n {\"content\": \"Prefers API docs\"},\n ],\n }\n\n _patch_personalization(headers, cookies, payload)\n\n me_after = _get_me(headers, cookies)\n personalization = me_after[\"personalization\"]\n\n assert personalization[\"name\"] == payload[\"name\"]\n assert personalization[\"role\"] == payload[\"role\"]\n assert personalization[\"use_memories\"] is True\n returned_memories = personalization[\"memories\"]\n assert len(returned_memories) == 2\n for mem in returned_memories:\n assert isinstance(mem[\"id\"], int)\n assert isinstance(mem[\"content\"], str)\n assert [m[\"content\"] for m in returned_memories] == [\n \"Prefers API docs\",\n \"Loves peanut butter\",\n ]\n\n # update memories to empty\n payload[\"memories\"] = []\n _patch_personalization(headers, cookies, payload)\n me_final = _get_me(headers, cookies)\n assert me_final[\"personalization\"][\"memories\"] == []\n\n\ndef test_enable_memory_tool_round_trip(reset: None) -> None: # noqa: ARG001\n user = UserManager.create()\n headers, cookies = _get_auth_headers(user)\n\n # default should be True\n me_initial = _get_me(headers, cookies)\n assert me_initial[\"personalization\"][\"enable_memory_tool\"] is True\n\n # disable enable_memory_tool\n _patch_personalization(headers, cookies, {\"enable_memory_tool\": False})\n me_after = _get_me(headers, cookies)\n assert me_after[\"personalization\"][\"enable_memory_tool\"] is False\n\n # re-enable\n _patch_personalization(headers, cookies, {\"enable_memory_tool\": True})\n me_reenabled = _get_me(headers, cookies)\n assert me_reenabled[\"personalization\"][\"enable_memory_tool\"] is True\n\n\ndef test_enable_memory_tool_independent_of_use_memories(\n reset: None, # noqa: ARG001\n) -> None:\n user = UserManager.create()\n headers, cookies = _get_auth_headers(user)\n\n # set use_memories=False and enable_memory_tool=True simultaneously\n _patch_personalization(\n headers, cookies, {\"use_memories\": False, \"enable_memory_tool\": True}\n )\n me = _get_me(headers, cookies)\n assert me[\"personalization\"][\"use_memories\"] is False\n assert me[\"personalization\"][\"enable_memory_tool\"] is True\n\n # reverse: use_memories=True and enable_memory_tool=False\n _patch_personalization(\n headers, cookies, {\"use_memories\": True, \"enable_memory_tool\": False}\n )\n me = _get_me(headers, cookies)\n assert me[\"personalization\"][\"use_memories\"] is True\n assert me[\"personalization\"][\"enable_memory_tool\"] is False\n" }, { "path": "backend/tests/integration/tests/personas/test_persona_categories.py", "content": "from uuid import uuid4\n\nimport pytest\nfrom requests.exceptions import HTTPError\n\nfrom tests.integration.common_utils.managers.persona import (\n PersonaLabelManager,\n)\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestPersonaLabel\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_persona_label_management(reset: None) -> None: # noqa: ARG001\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n persona_label = DATestPersonaLabel(\n id=None,\n name=f\"Test label {uuid4()}\",\n )\n persona_label = PersonaLabelManager.create(\n label=persona_label,\n user_performing_action=admin_user,\n )\n print(f\"Created persona label {persona_label.name} with id {persona_label.id}\")\n\n assert PersonaLabelManager.verify(\n label=persona_label,\n user_performing_action=admin_user,\n ), \"Persona label was not found after creation\"\n\n regular_user: DATestUser = UserManager.create(name=\"regular_user\")\n\n updated_persona_label = DATestPersonaLabel(\n id=persona_label.id,\n name=f\"Updated {persona_label.name}\",\n )\n with pytest.raises(HTTPError) as exc_info:\n PersonaLabelManager.update(\n label=updated_persona_label,\n user_performing_action=regular_user,\n )\n assert exc_info.value.response is not None\n assert exc_info.value.response.status_code == 403\n\n assert PersonaLabelManager.verify(\n label=persona_label,\n user_performing_action=admin_user,\n ), \"Persona label should not have been updated by non-admin user\"\n\n result = PersonaLabelManager.delete(\n label=persona_label,\n user_performing_action=regular_user,\n )\n assert (\n result is False\n ), \"Regular user should not be able to delete the persona label\"\n\n assert PersonaLabelManager.verify(\n label=persona_label,\n user_performing_action=admin_user,\n ), \"Persona label should not have been deleted by non-admin user\"\n\n updated_persona_label.name = f\"Updated {persona_label.name}\"\n updated_persona_label = PersonaLabelManager.update(\n label=updated_persona_label,\n user_performing_action=admin_user,\n )\n print(f\"Updated persona label to {updated_persona_label.name}\")\n\n assert PersonaLabelManager.verify(\n label=updated_persona_label,\n user_performing_action=admin_user,\n ), \"Persona label was not updated by admin\"\n\n success = PersonaLabelManager.delete(\n label=persona_label,\n user_performing_action=admin_user,\n )\n assert success, \"Admin user should be able to delete the persona label\"\n print(f\"Deleted persona label {persona_label.name} with id {persona_label.id}\")\n\n assert not PersonaLabelManager.verify(\n label=persona_label,\n user_performing_action=admin_user,\n ), \"Persona label should not exist after deletion by admin\"\n" }, { "path": "backend/tests/integration/tests/personas/test_persona_creation.py", "content": "import requests\n\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _list_minimal_personas(user: DATestUser) -> list[dict]:\n response = requests.get(\n f\"{API_SERVER_URL}/persona\",\n headers=user.headers,\n cookies=user.cookies,\n )\n response.raise_for_status()\n return response.json()\n\n\ndef _share_persona(\n persona_id: int, user_ids: list[str], acting_user: DATestUser\n) -> None:\n response = requests.patch(\n f\"{API_SERVER_URL}/persona/{persona_id}/share\",\n json={\"user_ids\": user_ids},\n headers=acting_user.headers,\n cookies=acting_user.cookies,\n )\n response.raise_for_status()\n\n\ndef test_persona_create_update_share_delete(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n basic_user: DATestUser,\n) -> None:\n # TODO: refactor `PersonaManager.verify`, not a good pattern\n # Create a persona as admin and verify it can be fetched\n expected_persona = PersonaManager.create(user_performing_action=admin_user)\n PersonaManager.verify(expected_persona, user_performing_action=admin_user)\n\n # Update the persona and verify changes\n updated_persona = PersonaManager.edit(\n expected_persona,\n name=f\"updated-{expected_persona.name}\",\n description=f\"updated-{expected_persona.description}\",\n is_public=False,\n user_performing_action=admin_user,\n )\n assert PersonaManager.verify(updated_persona, user_performing_action=admin_user)\n\n # Creator should see the persona in their minimal list\n creator_minimals = _list_minimal_personas(admin_user)\n assert any(p[\"id\"] == updated_persona.id for p in creator_minimals)\n\n # Regular user should not see a non-public, non-shared persona\n other_minimals_before = _list_minimal_personas(basic_user)\n assert all(p[\"id\"] != updated_persona.id for p in other_minimals_before)\n\n # Share persona with the regular user and verify visibility\n _share_persona(updated_persona.id, [basic_user.id], admin_user)\n other_minimals_after = _list_minimal_personas(basic_user)\n assert any(p[\"id\"] == updated_persona.id for p in other_minimals_after)\n\n # Delete persona and verify it no longer appears in lists\n assert PersonaManager.delete(updated_persona, user_performing_action=admin_user)\n\n # After deletion, list should not include it for either user\n creator_minimals_after_delete = _list_minimal_personas(admin_user)\n assert all(p[\"id\"] != updated_persona.id for p in creator_minimals_after_delete)\n\n regular_minimals_after_delete = _list_minimal_personas(basic_user)\n assert all(p[\"id\"] != updated_persona.id for p in regular_minimals_after_delete)\n" }, { "path": "backend/tests/integration/tests/personas/test_persona_file_context.py", "content": "\"\"\"\nIntegration tests for the unified persona file context flow.\n\nEnd-to-end tests that verify:\n1. Files can be uploaded and attached to a persona via API.\n2. The persona correctly reports its attached files.\n3. A chat session with a file-bearing persona processes without error.\n4. Precedence: custom persona files take priority over project files when\n the chat session is inside a project.\n\nThese tests run against a real Onyx deployment (all services running).\nFile processing is asynchronous, so we poll the file status endpoint\nuntil files reach COMPLETED before chatting.\n\"\"\"\n\nimport time\n\nimport requests\n\nfrom onyx.db.enums import UserFileStatus\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.constants import MAX_DELAY\nfrom tests.integration.common_utils.managers.chat import ChatSessionManager\nfrom tests.integration.common_utils.managers.file import FileManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.managers.project import ProjectManager\nfrom tests.integration.common_utils.test_file_utils import create_test_text_file\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n# ---------------------------------------------------------------------------\n# Helpers\n# ---------------------------------------------------------------------------\n\nFILE_PROCESSING_POLL_INTERVAL = 2\n\n\ndef _poll_file_statuses(\n user_file_ids: list[str],\n user: DATestUser,\n target_status: UserFileStatus = UserFileStatus.COMPLETED,\n timeout: int = MAX_DELAY,\n) -> None:\n \"\"\"Block until all files reach the target status or timeout expires.\"\"\"\n deadline = time.time() + timeout\n while time.time() < deadline:\n response = requests.post(\n f\"{API_SERVER_URL}/user/projects/file/statuses\",\n json={\"file_ids\": user_file_ids},\n headers=user.headers,\n )\n response.raise_for_status()\n statuses = response.json()\n if all(f[\"status\"] == target_status.value for f in statuses):\n return\n time.sleep(FILE_PROCESSING_POLL_INTERVAL)\n raise TimeoutError(\n f\"Files {user_file_ids} did not reach {target_status.value} within {timeout}s\"\n )\n\n\n# ---------------------------------------------------------------------------\n# Tests\n# ---------------------------------------------------------------------------\n\n\ndef test_persona_with_files_chat_no_error(\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"Upload files, attach them to a persona, wait for processing,\n then send a chat message. Verify no error is returned.\"\"\"\n\n # Upload files (creates UserFile records)\n text_file = create_test_text_file(\n \"The secret project codename is NIGHTINGALE. It was started in 2024 by the Advanced Research division.\"\n )\n file_descriptors, error = FileManager.upload_files(\n files=[(\"nightingale_brief.txt\", text_file)],\n user_performing_action=admin_user,\n )\n assert not error, f\"File upload failed: {error}\"\n assert len(file_descriptors) == 1\n\n user_file_id = file_descriptors[0][\"user_file_id\"]\n assert user_file_id is not None\n\n # Wait for file processing\n _poll_file_statuses([user_file_id], admin_user, timeout=120)\n\n # Create persona with the file attached\n persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Nightingale Agent\",\n description=\"Agent with secret file\",\n system_prompt=\"You are a helpful assistant with access to uploaded files.\",\n user_file_ids=[user_file_id],\n )\n\n # Verify persona has the file\n persona_snapshots = PersonaManager.get_one(persona.id, admin_user)\n assert len(persona_snapshots) == 1\n assert user_file_id in persona_snapshots[0].user_file_ids\n\n # Chat with the persona\n chat_session = ChatSessionManager.create(\n persona_id=persona.id,\n description=\"Test persona file context\",\n user_performing_action=admin_user,\n )\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"What is the secret project codename?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None, f\"Chat should succeed, got error: {response.error}\"\n assert len(response.full_message) > 0, \"Response should not be empty\"\n\n\ndef test_persona_without_files_still_works(\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"A persona with no attached files should still chat normally.\"\"\"\n persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Blank Agent\",\n description=\"No files attached\",\n system_prompt=\"You are a helpful assistant.\",\n )\n\n chat_session = ChatSessionManager.create(\n persona_id=persona.id,\n description=\"Test blank persona\",\n user_performing_action=admin_user,\n )\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"Hello, how are you?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None\n assert len(response.full_message) > 0\n\n\ndef test_persona_files_override_project_files(\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"When a custom persona (with its own files) is used inside a project,\n the persona's files take precedence — the project's files are invisible.\n\n We verify this by putting different content in project vs persona files\n and checking which content the model responds with.\"\"\"\n\n # Upload persona file\n persona_file = create_test_text_file(\"The persona's secret word is ALBATROSS.\")\n persona_fds, err1 = FileManager.upload_files(\n files=[(\"persona_secret.txt\", persona_file)],\n user_performing_action=admin_user,\n )\n assert not err1\n persona_user_file_id = persona_fds[0][\"user_file_id\"]\n assert persona_user_file_id is not None\n # Create a project and upload project files\n project = ProjectManager.create(\n name=\"Precedence Test Project\",\n user_performing_action=admin_user,\n )\n project_files = [\n (\"project_secret.txt\", b\"The project's secret word is FLAMINGO.\"),\n ]\n project_upload_result = ProjectManager.upload_files(\n project_id=project.id,\n files=project_files,\n user_performing_action=admin_user,\n )\n assert len(project_upload_result.user_files) == 1\n project_user_file_id = str(project_upload_result.user_files[0].id)\n\n # Wait for both persona and project file processing\n _poll_file_statuses([persona_user_file_id], admin_user, timeout=120)\n _poll_file_statuses([project_user_file_id], admin_user, timeout=120)\n\n # Create persona with persona file\n persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"Override Agent\",\n description=\"Persona with its own files\",\n system_prompt=\"You are a helpful assistant. Answer using the files.\",\n user_file_ids=[persona_user_file_id],\n )\n\n # Create chat session inside the project but using the custom persona\n chat_session = ChatSessionManager.create(\n persona_id=persona.id,\n project_id=project.id,\n user_performing_action=admin_user,\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"What is the secret word?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None, f\"Chat should succeed, got error: {response.error}\"\n # The persona's file should be what the model sees, not the project's\n message_lower = response.full_message.lower()\n assert (\n \"albatross\" in message_lower\n ), f\"Response should reference the persona file's secret word (ALBATROSS), but got: {response.full_message}\"\n\n\ndef test_default_persona_in_project_uses_project_files(\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"When the default persona (id=0) is used inside a project,\n the project's files should be used for context.\"\"\"\n project = ProjectManager.create(\n name=\"Default Persona Project\",\n user_performing_action=admin_user,\n )\n project_files = [\n (\"project_info.txt\", b\"The project mascot is a PANGOLIN.\"),\n ]\n upload_result = ProjectManager.upload_files(\n project_id=project.id,\n files=project_files,\n user_performing_action=admin_user,\n )\n assert len(upload_result.user_files) == 1\n\n # Wait for project file processing\n project_file_id = str(upload_result.user_files[0].id)\n _poll_file_statuses([project_file_id], admin_user, timeout=120)\n\n # Create chat session inside project using default persona (id=0)\n chat_session = ChatSessionManager.create(\n persona_id=0,\n project_id=project.id,\n user_performing_action=admin_user,\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"What is the project mascot?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None\n assert (\n \"pangolin\" in response.full_message.lower()\n ), f\"Response should reference the project file content (PANGOLIN), but got: {response.full_message}\"\n\n\ndef test_custom_persona_no_files_in_project_ignores_project(\n admin_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"A custom persona with NO files, used inside a project with files,\n should NOT see the project's files. The project is purely organizational.\n\n We verify by asking about content only in the project file and checking\n the model does NOT reference it.\"\"\"\n\n project = ProjectManager.create(\n name=\"Ignored Project\",\n user_performing_action=admin_user,\n )\n project_upload_result = ProjectManager.upload_files(\n project_id=project.id,\n files=[(\"project_only.txt\", b\"The project secret is CAPYBARA.\")],\n user_performing_action=admin_user,\n )\n assert len(project_upload_result.user_files) == 1\n project_user_file_id = str(project_upload_result.user_files[0].id)\n\n # Wait for project file processing\n _poll_file_statuses([project_user_file_id], admin_user, timeout=120)\n\n # Custom persona with no files\n persona = PersonaManager.create(\n user_performing_action=admin_user,\n name=\"No Files Agent\",\n description=\"No files, project is irrelevant\",\n system_prompt=(\n \"You are a helpful assistant. If you do not have information \"\n \"to answer a question, say 'I do not have that information.'\"\n ),\n )\n\n chat_session = ChatSessionManager.create(\n persona_id=persona.id,\n project_id=project.id,\n user_performing_action=admin_user,\n )\n\n response = ChatSessionManager.send_message(\n chat_session_id=chat_session.id,\n message=\"What is the project secret?\",\n user_performing_action=admin_user,\n )\n\n assert response.error is None\n assert len(response.full_message) > 0\n assert \"capybara\" not in response.full_message.lower(), (\n \"Response should NOT reference the project file content (CAPYBARA) \"\n \"because the custom persona has no files and should not inherit \"\n f\"project files, but got: {response.full_message}\"\n )\n" }, { "path": "backend/tests/integration/tests/personas/test_persona_label_updates.py", "content": "from uuid import uuid4\n\nimport requests\n\nfrom onyx.server.features.persona.models import PersonaUpsertRequest\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.persona import PersonaLabelManager\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.test_models import DATestPersonaLabel\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_update_persona_with_null_label_ids_preserves_labels(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n persona_label = PersonaLabelManager.create(\n label=DATestPersonaLabel(name=f\"Test label {uuid4()}\"),\n user_performing_action=admin_user,\n )\n assert persona_label.id is not None\n persona = PersonaManager.create(\n label_ids=[persona_label.id],\n user_performing_action=admin_user,\n )\n\n updated_description = f\"{persona.description}-updated\"\n update_request = PersonaUpsertRequest(\n name=persona.name,\n description=updated_description,\n system_prompt=persona.system_prompt or \"\",\n task_prompt=persona.task_prompt or \"\",\n datetime_aware=persona.datetime_aware,\n document_set_ids=persona.document_set_ids,\n is_public=persona.is_public,\n llm_model_provider_override=persona.llm_model_provider_override,\n llm_model_version_override=persona.llm_model_version_override,\n tool_ids=persona.tool_ids,\n users=[],\n groups=[],\n label_ids=None,\n )\n\n response = requests.patch(\n f\"{API_SERVER_URL}/persona/{persona.id}\",\n json=update_request.model_dump(mode=\"json\", exclude_none=False),\n headers=admin_user.headers,\n cookies=admin_user.cookies,\n )\n response.raise_for_status()\n\n fetched = requests.get(\n f\"{API_SERVER_URL}/persona/{persona.id}\",\n headers=admin_user.headers,\n cookies=admin_user.cookies,\n )\n fetched.raise_for_status()\n fetched_persona = fetched.json()\n\n assert fetched_persona[\"description\"] == updated_description\n fetched_label_ids = {label[\"id\"] for label in fetched_persona[\"labels\"]}\n assert persona_label.id in fetched_label_ids\n" }, { "path": "backend/tests/integration/tests/personas/test_persona_pagination.py", "content": "import requests\n\nfrom onyx.server.features.persona.constants import ADMIN_AGENTS_RESOURCE\nfrom onyx.server.features.persona.constants import AGENTS_RESOURCE\nfrom tests.integration.common_utils.constants import API_SERVER_URL\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef _get_agents_paginated(\n user: DATestUser,\n page_num: int,\n page_size: int,\n include_deleted: bool = False,\n get_editable: bool = False,\n include_default: bool = True,\n) -> tuple[dict, int]:\n \"\"\"Fetches a paginated page of agents, with status code.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}{AGENTS_RESOURCE}\",\n params={\n \"page_num\": page_num,\n \"page_size\": page_size,\n \"include_deleted\": include_deleted,\n \"get_editable\": get_editable,\n \"include_default\": include_default,\n },\n headers=user.headers,\n cookies=user.cookies,\n )\n return response.json(), response.status_code\n\n\ndef _get_agents_admin_paginated(\n user: DATestUser,\n page_num: int,\n page_size: int,\n include_deleted: bool = False,\n get_editable: bool = False,\n include_default: bool = True,\n) -> tuple[dict, int]:\n \"\"\"Fetches a paginated page of agents (admin endpoint) with status code.\"\"\"\n response = requests.get(\n f\"{API_SERVER_URL}{ADMIN_AGENTS_RESOURCE}\",\n params={\n \"page_num\": page_num,\n \"page_size\": page_size,\n \"include_deleted\": include_deleted,\n \"get_editable\": get_editable,\n \"include_default\": include_default,\n },\n headers=user.headers,\n cookies=user.cookies,\n )\n response.raise_for_status()\n return response.json(), response.status_code\n\n\ndef test_persona_pagination_basic(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test basic pagination - verify correct items and total count.\"\"\"\n # Preconditions\n personas_to_create = 25\n personas = []\n for i in range(personas_to_create):\n persona = PersonaManager.create(\n name=f\"Test Persona {i}\",\n user_performing_action=admin_user,\n )\n personas.append(persona)\n\n # Under test and postconditions\n # Test page 0 with size 10.\n page_0, _ = _get_agents_paginated(admin_user, page_num=0, page_size=10)\n assert \"items\" in page_0\n assert \"total_items\" in page_0\n assert len(page_0[\"items\"]) == 10\n assert (\n page_0[\"total_items\"] >= personas_to_create\n ) # At least personas_to_create (may have default personas)\n\n # Test page 2 with size 10 (should have 5+ items if only our test personas\n # exist).\n page_2, _ = _get_agents_paginated(admin_user, page_num=2, page_size=10)\n assert len(page_2[\"items\"]) >= 5\n assert page_2[\"total_items\"] >= personas_to_create\n\n # Test page beyond end (page 10 with size 10, offset 100).\n page_beyond, _ = _get_agents_paginated(admin_user, page_num=10, page_size=10)\n assert len(page_beyond[\"items\"]) == 0\n assert page_beyond[\"total_items\"] >= personas_to_create # Total doesn't change.\n\n\ndef test_persona_pagination_ordering(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test ordering - display_priority ASC nulls last, then ID ASC.\"\"\"\n # Preconditions\n # Create personas with specific display_priority values.\n persona_a = PersonaManager.create(\n name=\"Persona A\",\n description=\"This should be second\",\n user_performing_action=admin_user,\n display_priority=2,\n )\n persona_b = PersonaManager.create(\n name=\"Persona B\",\n description=\"This should be first\",\n user_performing_action=admin_user,\n display_priority=1,\n )\n persona_c = PersonaManager.create(\n name=\"Persona C\",\n description=\"This should be third\",\n user_performing_action=admin_user,\n display_priority=3,\n )\n persona_d = PersonaManager.create(\n name=\"Persona D\",\n description=\"This should be fourth\",\n user_performing_action=admin_user,\n display_priority=3, # Note the same prio as above, should sort by id\n )\n\n # Under test\n page_0, _ = _get_agents_paginated(admin_user, page_num=0, page_size=100)\n\n # Postconditions\n # Find our personas in the results.\n our_expected_ordered_persona_ids = [\n persona_b.id,\n persona_a.id,\n persona_c.id,\n persona_d.id,\n ]\n our_personas_in_results = [\n p for p in page_0[\"items\"] if p[\"id\"] in our_expected_ordered_persona_ids\n ]\n assert len(our_personas_in_results) == 4\n # Verify ordering.\n for i in range(len(our_expected_ordered_persona_ids)):\n assert our_expected_ordered_persona_ids[i] == our_personas_in_results[i][\"id\"]\n\n\ndef test_persona_pagination_admin_endpoint(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test admin paginated endpoint returns PersonaSnapshot format.\"\"\"\n # Preconditions\n personas_to_create = 5\n for i in range(personas_to_create):\n PersonaManager.create(\n name=f\"Admin Test Persona {i}\",\n user_performing_action=admin_user,\n )\n\n # Under test\n page_0, _ = _get_agents_admin_paginated(admin_user, page_num=0, page_size=10)\n\n # Postconditions\n assert \"items\" in page_0\n assert \"total_items\" in page_0\n assert len(page_0[\"items\"]) >= personas_to_create\n assert page_0[\"total_items\"] >= personas_to_create\n # Verify admin-specific fields are present (PersonaSnapshot has more\n # fields).\n first_persona = page_0[\"items\"][0]\n # PersonaSnapshot should have these fields that MinimalPersonaSnapshot\n # doesn't.\n assert \"users\" in first_persona\n assert \"groups\" in first_persona\n assert \"user_file_ids\" in first_persona\n\n\ndef test_persona_pagination_with_deleted(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test pagination with include_deleted parameter.\"\"\"\n # Preconditions\n # Create and delete a persona.\n persona = PersonaManager.create(\n name=\"To Be Deleted\",\n user_performing_action=admin_user,\n )\n PersonaManager.delete(persona, user_performing_action=admin_user)\n\n # Under test and postconditions\n # Without include_deleted, should not appear.\n page_without_deleted, _ = _get_agents_paginated(\n admin_user, page_num=0, page_size=100, include_deleted=False\n )\n persona_ids_without_deleted = [p[\"id\"] for p in page_without_deleted[\"items\"]]\n assert persona.id not in persona_ids_without_deleted\n\n # With include_deleted, should appear.\n page_with_deleted, _ = _get_agents_paginated(\n admin_user, page_num=0, page_size=100, include_deleted=True\n )\n persona_ids_with_deleted = [p[\"id\"] for p in page_with_deleted[\"items\"]]\n assert persona.id in persona_ids_with_deleted\n\n # Total counts should differ.\n assert page_with_deleted[\"total_items\"] > page_without_deleted[\"total_items\"]\n\n\ndef test_persona_pagination_page_size_limits(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test page_size parameter validation (max 1000).\"\"\"\n # Preconditions\n # Create a few personas.\n for i in range(5):\n PersonaManager.create(\n name=f\"Size Limit Test {i}\",\n user_performing_action=admin_user,\n )\n\n # Under test and postconditions\n # Valid page_size of 1\n data, _ = _get_agents_paginated(admin_user, page_num=0, page_size=1)\n assert len(data[\"items\"]) <= 1\n\n # Valid page_size of 1000\n data, _ = _get_agents_paginated(admin_user, page_num=0, page_size=1000)\n # We assume not that many default personas are made.\n assert len(data[\"items\"]) == data[\"total_items\"]\n\n # Invalid page_size of 1001 (exceeds max)\n _, status_code = _get_agents_paginated(admin_user, page_num=0, page_size=1001)\n assert status_code == 422 # Validation error\n\n # Invalid page_size of 0\n _, status_code = _get_agents_paginated(admin_user, page_num=0, page_size=0)\n assert status_code == 422 # Validation error\n\n\ndef test_persona_pagination_count_accuracy(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n) -> None:\n \"\"\"Test that total_items count is consistent across pages.\"\"\"\n # Preconditions\n # Create 15 personas.\n created_personas = []\n for i in range(15):\n persona = PersonaManager.create(\n name=f\"Count Test {i}\",\n user_performing_action=admin_user,\n )\n created_personas.append(persona)\n\n # Under test and postconditions\n # Fetch first page to get total count.\n page_0, _ = _get_agents_paginated(admin_user, page_num=0, page_size=5)\n total_items = page_0[\"total_items\"]\n assert total_items >= 15\n\n # Fetch all pages to cover all personas.\n all_ids_from_pages: set[int] = set()\n num_pages_needed = (total_items + 4) // 5 # Ceiling division\n for page_num in range(num_pages_needed):\n page, _ = _get_agents_paginated(admin_user, page_num=page_num, page_size=5)\n # All pages should report the same total.\n assert (\n page[\"total_items\"] == total_items\n ), f\"Page {page_num} has inconsistent total_items\"\n all_ids_from_pages.update(p[\"id\"] for p in page[\"items\"])\n\n # Our created personas should all appear.\n our_ids = {p.id for p in created_personas}\n assert our_ids.issubset(\n all_ids_from_pages\n ), \"All created personas should appear in paginated results\"\n\n\ndef test_persona_pagination_user_permissions(\n reset: None, # noqa: ARG001\n admin_user: DATestUser,\n basic_user: DATestUser,\n) -> None:\n \"\"\"Test that pagination respects user permissions.\"\"\"\n # Preconditions\n # Admin creates a private persona (not shared).\n private_persona = PersonaManager.create(\n name=\"Private Persona\",\n description=\"Not shared\",\n is_public=False,\n user_performing_action=admin_user,\n )\n # Admin creates a public persona.\n public_persona = PersonaManager.create(\n name=\"Public Persona\",\n description=\"Shared with all\",\n is_public=True,\n user_performing_action=admin_user,\n )\n\n # Under test and postconditions\n # Admin should see both in paginated results.\n admin_page, _ = _get_agents_paginated(admin_user, page_num=0, page_size=100)\n admin_ids = {p[\"id\"] for p in admin_page[\"items\"]}\n assert private_persona.id in admin_ids\n assert public_persona.id in admin_ids\n\n # Basic user should only see public persona.\n user_page, _ = _get_agents_paginated(basic_user, page_num=0, page_size=100)\n user_ids = {p[\"id\"] for p in user_page[\"items\"]}\n assert private_persona.id not in user_ids\n assert public_persona.id in user_ids\n\n # Totals should differ.\n assert admin_page[\"total_items\"] > user_page[\"total_items\"]\n" }, { "path": "backend/tests/integration/tests/personas/test_unified_assistant.py", "content": "\"\"\"Integration tests for the unified assistant.\"\"\"\n\nfrom tests.integration.common_utils.managers.persona import PersonaManager\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\ndef test_unified_assistant(\n reset: None, admin_user: DATestUser # noqa: ARG001\n) -> None: # noqa: ARG001\n \"\"\"Combined test verifying unified assistant existence, tools, and starter messages.\"\"\"\n # Fetch all personas\n personas = PersonaManager.get_all(admin_user)\n\n # Find the unified assistant (ID 0)\n unified_assistant = None\n for persona in personas:\n if persona.id == 0:\n unified_assistant = persona\n break\n\n # Assert that there are no other assistants (personas) besides the unified assistant\n # (ID 0)\n assert (\n len(personas) == 1\n ), f\"Expected only the unified assistant, found {len(personas)} personas\"\n\n # Verify the unified assistant exists\n assert unified_assistant is not None, \"Unified assistant (ID 0) not found\"\n\n # Verify basic properties\n assert unified_assistant.name == \"Assistant\"\n assert (\n \"search, web browsing, and image generation\"\n in unified_assistant.description.lower()\n )\n assert unified_assistant.is_featured is True\n assert unified_assistant.is_listed is True\n\n # Verify tools\n tools = unified_assistant.tools\n tool_names = [tool.name for tool in tools]\n assert \"internal_search\" in tool_names, \"SearchTool not found in unified assistant\"\n assert (\n \"generate_image\" in tool_names\n ), \"ImageGenerationTool not found in unified assistant\"\n assert \"web_search\" in tool_names, \"WebSearchTool not found in unified assistant\"\n\n # Verify no starter messages\n starter_messages = unified_assistant.starter_messages or []\n assert len(starter_messages) == 0, \"Starter messages found\"\n" }, { "path": "backend/tests/integration/tests/projects/test_projects.py", "content": "from typing import List\n\nimport pytest\n\nfrom onyx.db.engine.sql_engine import get_session_with_current_tenant\nfrom onyx.db.models import UserFile\nfrom onyx.server.features.projects.models import UserProjectSnapshot\nfrom tests.integration.common_utils.managers.project import ProjectManager\nfrom tests.integration.common_utils.reset import reset_all\nfrom tests.integration.common_utils.test_models import DATestLLMProvider\nfrom tests.integration.common_utils.test_models import DATestUser\n\n\n@pytest.fixture(scope=\"module\", autouse=True)\ndef reset_for_module() -> None:\n \"\"\"Reset all data once before running any tests in this module.\"\"\"\n reset_all()\n\n\ndef test_projects_flow(\n reset_for_module: None, # noqa: ARG001\n basic_user: DATestUser,\n llm_provider: DATestLLMProvider, # noqa: ARG001\n) -> None:\n \"\"\"End-to-end project flow covering creation, listing, files, instructions, deletion, and edge cases.\"\"\"\n # Case 1: Project creation and listing\n ProjectManager.create(\n name=\"Test Project 1\",\n user_performing_action=basic_user,\n )\n ProjectManager.create(\n name=\"Test Project 2\",\n user_performing_action=basic_user,\n )\n\n projects = ProjectManager.get_all(user_performing_action=basic_user)\n assert len(projects) >= 2\n project_names = {p.name for p in projects}\n assert \"Test Project 1\" in project_names\n assert \"Test Project 2\" in project_names\n assert all(str(p.user_id) == basic_user.id for p in projects)\n\n # Case 2: File upload and management\n file_project = ProjectManager.create(\n name=\"File Test Project\",\n user_performing_action=basic_user,\n )\n test_files = [\n (\"test1.txt\", b\"This is test file 1 content\"),\n (\"test2.txt\", b\"This is test file 2 content\"),\n ]\n upload_result = ProjectManager.upload_files(\n project_id=file_project.id,\n files=test_files,\n user_performing_action=basic_user,\n )\n assert len(upload_result.user_files) == 2\n assert len(upload_result.rejected_files) == 0\n project_files = ProjectManager.get_project_files(\n project_id=file_project.id,\n user_performing_action=basic_user,\n )\n assert len(project_files) == 2\n file_names = {f.name for f in project_files}\n assert \"test1.txt\" in file_names\n assert \"test2.txt\" in file_names\n\n # Case 3: Instructions set and update\n instructions_project = ProjectManager.create(\n name=\"Instructions Test Project\",\n user_performing_action=basic_user,\n )\n instructions = \"These are test project instructions\"\n result = ProjectManager.set_instructions(\n project_id=instructions_project.id,\n instructions=instructions,\n user_performing_action=basic_user,\n )\n assert result == instructions\n new_instructions = \"These are updated test project instructions\"\n result = ProjectManager.set_instructions(\n project_id=instructions_project.id,\n instructions=new_instructions,\n user_performing_action=basic_user,\n )\n assert result == new_instructions\n\n # Case 4: Deletion with files (unlink but do not delete files)\n delete_file_project = ProjectManager.create(\n name=\"Deletion Test Project\",\n user_performing_action=basic_user,\n )\n del_test_files = [\n (\"delete_test1.txt\", b\"This is test file 1 content\"),\n (\"delete_test2.txt\", b\"This is test file 2 content\"),\n ]\n ProjectManager.upload_files(\n project_id=delete_file_project.id,\n files=del_test_files,\n user_performing_action=basic_user,\n )\n del_project_files = ProjectManager.get_project_files(\n project_id=delete_file_project.id,\n user_performing_action=basic_user,\n )\n assert len(del_project_files) == 2\n deletion_success = ProjectManager.delete(\n project_id=delete_file_project.id,\n user_performing_action=basic_user,\n )\n assert deletion_success\n assert ProjectManager.verify_deleted(\n project_id=delete_file_project.id,\n user_performing_action=basic_user,\n )\n assert ProjectManager.verify_files_unlinked(\n project_id=delete_file_project.id,\n user_performing_action=basic_user,\n )\n with get_session_with_current_tenant() as db_session:\n file_ids = [f.id for f in del_project_files]\n remaining_files = (\n db_session.query(UserFile).filter(UserFile.id.in_(file_ids)).all()\n )\n assert len(remaining_files) == 2\n\n # Case 5: Deletion with chat sessions unlinked\n chat_project = ProjectManager.create(\n name=\"Chat Session Test Project\",\n user_performing_action=basic_user,\n )\n deletion_success = ProjectManager.delete(\n project_id=chat_project.id,\n user_performing_action=basic_user,\n )\n assert deletion_success\n assert ProjectManager.verify_chat_sessions_unlinked(\n project_id=chat_project.id,\n user_performing_action=basic_user,\n )\n\n # Case 6: Multiple project operations\n projects_group: List[UserProjectSnapshot] = []\n for i in range(3):\n proj = ProjectManager.create(\n name=f\"Multi-op Project {i}\",\n user_performing_action=basic_user,\n )\n projects_group.append(proj)\n\n for i, proj in enumerate(projects_group):\n tfiles = [\n (f\"multi_test{i}_1.txt\", b\"This is test file 1 content\"),\n (f\"multi_test{i}_2.txt\", b\"This is test file 2 content\"),\n ]\n ProjectManager.upload_files(\n project_id=proj.id,\n files=tfiles,\n user_performing_action=basic_user,\n )\n\n for i, proj in enumerate(projects_group):\n instr = f\"Instructions for project {i}\"\n res = ProjectManager.set_instructions(\n project_id=proj.id,\n instructions=instr,\n user_performing_action=basic_user,\n )\n assert res == instr\n\n for proj in projects_group:\n proj_files = ProjectManager.get_project_files(\n project_id=proj.id,\n user_performing_action=basic_user,\n )\n assert len(proj_files) == 2\n deletion_success = ProjectManager.delete(\n project_id=proj.id,\n user_performing_action=basic_user,\n )\n assert deletion_success\n assert ProjectManager.verify_deleted(\n project_id=proj.id,\n user_performing_action=basic_user,\n )\n assert ProjectManager.verify_files_unlinked(\n project_id=proj.id,\n user_performing_action=basic_user,\n )\n with get_session_with_current_tenant() as db_session:\n file_ids = [f.id for f in proj_files]\n remaining_files = (\n db_session.query(UserFile).filter(UserFile.id.in_(file_ids)).all()\n )\n assert len(remaining_files) == 2\n\n # Case 7: Edge cases\n with pytest.raises(Exception):\n ProjectManager.create(\n name=\"\",\n user_performing_action=basic_user,\n )\n\n non_existent_id = 99999\n deletion_success = ProjectManager.delete(\n project_id=non_existent_id,\n user_performing_action=basic_user,\n )\n assert not deletion_success\n\n with pytest.raises(Exception):\n ProjectManager.set_instructions(\n project_id=non_existent_id,\n instructions=\"Test instructions\",\n user_performing_action=basic_user,\n )\n\n with pytest.raises(Exception):\n ProjectManager.upload_files(\n project_id=non_existent_id,\n files=[(\"test.txt\", b\"content\")],\n user_performing_action=basic_user,\n )\n\n long_name = \"a\" * 1000\n with pytest.raises(Exception):\n ProjectManager.create(\n name=long_name,\n user_performing_action=basic_user,\n )\n\n long_instr_project = ProjectManager.create(\n name=\"Long Instructions Test\",\n user_performing_action=basic_user,\n )\n long_instructions = \"a\" * 10000\n result = ProjectManager.set_instructions(\n project_id=long_instr_project.id,\n instructions=long_instructions,\n user_performing_action=basic_user,\n )\n assert result == long_instructions\n" }, { "path": "backend/tests/integration/tests/pruning/test_pruning.py", "content": "import http.server\nimport os\nimport shutil\nimport tempfile\nimport threading\nfrom collections.abc import Generator\nfrom contextlib import contextmanager\nfrom datetime import datetime\nfrom datetime import timezone\nfrom time import sleep\nfrom typing import Any\n\nimport uvicorn\nfrom fastapi import FastAPI\nfrom fastapi.staticfiles import StaticFiles\n\nfrom onyx.server.documents.models import DocumentSource\nfrom onyx.utils.logger import setup_logger\nfrom tests.integration.common_utils.managers.api_key import APIKeyManager\nfrom tests.integration.common_utils.managers.cc_pair import CCPairManager\nfrom tests.integration.common_utils.managers.user import UserManager\nfrom tests.integration.common_utils.test_models import DATestUser\nfrom tests.integration.common_utils.vespa import vespa_fixture\n\nlogger = setup_logger()\n\n\n# FastAPI server for serving files\ndef create_fastapi_app(directory: str) -> FastAPI:\n app = FastAPI()\n\n # Mount the directory to serve static files\n app.mount(\"/\", StaticFiles(directory=directory, html=True), name=\"static\")\n\n return app\n\n\n# as far as we know, this doesn't hang when crawled. This is good.\n@contextmanager\ndef fastapi_server_context(\n directory: str, port: int = 8000\n) -> Generator[None, None, None]:\n app = create_fastapi_app(directory)\n\n config = uvicorn.Config(app=app, host=\"0.0.0.0\", port=port, log_level=\"info\")\n server = uvicorn.Server(config)\n\n # Create a thread to run the FastAPI server\n server_thread = threading.Thread(target=server.run)\n server_thread.daemon = (\n True # Ensures the thread will exit when the main program exits\n )\n\n try:\n # Start the server in the background\n server_thread.start()\n sleep(5) # Give it a few seconds to start\n yield # Yield control back to the calling function (context manager in use)\n finally:\n # Shutdown the server\n server.should_exit = True\n server_thread.join()\n\n\n# Leaving this here for posterity and experimentation, but the reason we're\n# not using this is python's web servers hang frequently when crawled\n# this is obviously not good for a unit test\n@contextmanager\ndef http_server_context(\n directory: str, port: int = 8000\n) -> Generator[http.server.ThreadingHTTPServer, None, None]:\n # Create a handler that serves files from the specified directory\n def handler_class(\n *args: Any, **kwargs: Any\n ) -> http.server.SimpleHTTPRequestHandler:\n return http.server.SimpleHTTPRequestHandler(\n *args, directory=directory, **kwargs\n )\n\n # Create an HTTPServer instance\n httpd = http.server.ThreadingHTTPServer((\"0.0.0.0\", port), handler_class)\n\n # Define a thread that runs the server in the background\n server_thread = threading.Thread(target=httpd.serve_forever)\n server_thread.daemon = (\n True # Ensures the thread will exit when the main program exits\n )\n\n try:\n # Start the server in the background\n server_thread.start()\n sleep(5) # give it a few seconds to start\n yield httpd\n finally:\n # Shutdown the server and wait for the thread to finish\n httpd.shutdown()\n httpd.server_close()\n server_thread.join()\n\n\ndef test_web_pruning(\n reset: None, vespa_client: vespa_fixture # noqa: ARG001\n) -> None: # noqa: ARG001\n # Creating an admin user (first user created is automatically an admin)\n admin_user: DATestUser = UserManager.create(name=\"admin_user\")\n\n # add api key to user\n APIKeyManager.create(\n user_performing_action=admin_user,\n )\n\n test_filename = os.path.realpath(__file__)\n test_directory = os.path.dirname(test_filename)\n with tempfile.TemporaryDirectory() as temp_dir:\n port = 8889\n\n website_src = os.path.join(test_directory, \"website\")\n website_tgt = os.path.join(temp_dir, \"website\")\n shutil.copytree(website_src, website_tgt)\n with fastapi_server_context(os.path.join(temp_dir, \"website\"), port):\n sleep(1) # sleep a tiny bit before starting everything\n\n hostname = os.getenv(\"TEST_WEB_HOSTNAME\", \"localhost\")\n config = {\n \"base_url\": f\"http://{hostname}:{port}/\",\n \"web_connector_type\": \"recursive\",\n }\n\n # store the time before we create the connector so that we know after\n # when the indexing should have started\n now = datetime.now(timezone.utc)\n\n # create connector\n cc_pair_1 = CCPairManager.create_from_scratch(\n source=DocumentSource.WEB,\n connector_specific_config=config,\n user_performing_action=admin_user,\n )\n\n CCPairManager.wait_for_indexing_completion(\n cc_pair_1, now, timeout=300, user_performing_action=admin_user\n )\n\n selected_cc_pair = CCPairManager.get_indexing_status_by_id(\n cc_pair_1.id, user_performing_action=admin_user\n )\n\n assert selected_cc_pair is not None, \"cc_pair not found after indexing!\"\n\n # used to be 15, but now\n # localhost:8889/ and localhost:8889/index.html are deduped\n assert selected_cc_pair.docs_indexed == 14\n\n logger.info(\"Removing about.html.\")\n os.remove(os.path.join(website_tgt, \"about.html\"))\n logger.info(\"Removing courses.html.\")\n os.remove(os.path.join(website_tgt, \"courses.html\"))\n\n now = datetime.now(timezone.utc)\n CCPairManager.prune(cc_pair_1, user_performing_action=admin_user)\n CCPairManager.wait_for_prune(\n cc_pair_1, now, timeout=300, user_performing_action=admin_user\n )\n\n selected_cc_pair = CCPairManager.get_indexing_status_by_id(\n cc_pair_1.id, user_performing_action=admin_user\n )\n assert selected_cc_pair is not None, \"cc_pair not found after pruning!\"\n assert selected_cc_pair.docs_indexed == 12\n\n # check vespa\n root_id = f\"http://{hostname}:{port}/\"\n index_id = f\"http://{hostname}:{port}/index.html\"\n about_id = f\"http://{hostname}:{port}/about.html\"\n courses_id = f\"http://{hostname}:{port}/courses.html\"\n\n doc_ids = [root_id, index_id, about_id, courses_id]\n retrieved_docs_dict = vespa_client.get_documents_by_id(doc_ids)[\"documents\"]\n retrieved_docs = {\n doc[\"fields\"][\"document_id\"]: doc[\"fields\"]\n for doc in retrieved_docs_dict\n }\n\n # verify root exists in Vespa\n retrieved_doc = retrieved_docs.get(root_id)\n assert retrieved_doc\n\n # verify index.html does not exist in Vespa since it is a duplicate of root\n retrieved_doc = retrieved_docs.get(index_id)\n assert not retrieved_doc\n\n # verify about and courses do not exist\n retrieved_doc = retrieved_docs.get(about_id)\n assert not retrieved_doc\n\n retrieved_doc = retrieved_docs.get(courses_id)\n assert not retrieved_doc\n" }, { "path": "backend/tests/integration/tests/pruning/website/about.html", "content": "\n\n \n \n Above Multi-purpose Free Bootstrap Responsive Template\n \n \n \n \n \n \n \n \n \n\n \n \n \n \n
    \n \n
    \n
    \n
    \n
    \n \n \n \n \n \n \"logo\"\n\n
    \n
    \n \n
    \n
    \n
    \n
    \n \n
    \n
    \n
    \n
    \n

    About Us

    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n

    We are awesome TEAM

    \n

    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et\n quasi architecto beatae vitae dicta sunt explicabo. Nemo\n enim ipsam voluptatem quia voluptas\n

    \n

    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et\n quasi architecto beatae vitae dicta sunt explicabo. Nemo\n enim ipsam voluptatem quia voluptas\n

    \n
    \n Read more\n
    \n
    \n
    \n\n
    \n
    \n \"\"\n
    \n
    \n
    \n

    \n Lorem ipsum dolor sit amet, cadipisicing sit amet, consectetur\n adipisicing elit. Atque sed, quidem quis praesentium, ut unde\n fuga error commodi architecto, laudantium culpa tenetur at id,\n beatae pet.\n

    \n

    \n Lorem ipsum dolor sit amet, consectetur adipisicing elit.\n adipisicing sit amet, consectetur adipisicing elit. Atque sed,\n quidem quis praesentium,m deserunt.\n

    \n
      \n
    • \n Lorem\n ipsum enimdolor sit amet\n
      • \n
    • \n \n Explicabo deleniti neque aliquid\n
      • \n
    • \n \n Consectetur adipisicing elit\n
      • \n
    • \n Lorem\n ipsum dolor sit amet\n
      • \n
    • \n Quo\n issimos molest quibusdam temporibus\n
      • \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n \n
    \n

    Why Choose Us?

    \n
    \n

    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et quasi\n architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam\n voluptatem quia voluptas sit aspernatur.

    Sed ut\n perspiciaatis iste natus error sit voluptatem probably haven't\n heard of them accusamus.\n

    \n
    \n
    \n
    \n

    Our Solution

    \n
    \n \n
    \n \n
    \n \n
    \n

    \n \n Accordion Heading\n Text Item # 1\n \n

    \n
    \n
    \n \n
    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et\n quasi architecto beatae vitae dicta sunt explicabo. Nemo\n enim ipsam voluptatem quia voluptas\n
    \n
    \n
    \n
    \n
    \n

    \n \n Accordion Heading\n Text Item # 2\n \n

    \n
    \n
    \n
    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et\n quasi architecto beatae vitae dicta sunt explicabo. Nemo\n enim ipsam voluptatem quia voluptas\n
    \n
    \n
    \n
    \n
    \n

    \n \n Accordion Heading\n Text Item # 3\n \n

    \n
    \n \n
    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et\n quasi architecto beatae vitae dicta sunt explicabo. Nemo\n enim ipsam voluptatem quia voluptas\n
    \n
    \n
    \n
    \n
    \n

    \n \n Accordion Heading\n Text Item # 4\n \n

    \n
    \n
    \n
    \n Sed ut perspiciaatis unde omnis iste natus error sit\n voluptatem accusantium doloremque laudantium, totam rem\n aperiam, eaque ipsa quae ab illo inventore veritatis et\n quasi architecto beatae vitae dicta sunt explicabo. Nemo\n enim ipsam voluptatem quia voluptas\n
    \n
    \n
    \n
    \n \n
    \n\n
    \n
    \n

    Our Expertise

    \n
    \n
    Web Development
    \n
    \n \n \n 40% Complete (success)\n
    \n
    \n
    Designing
    \n
    \n \n 40% Complete (success)\n
    \n
    \n
    User Experience
    \n
    \n \n 40% Complete (success)\n
    \n
    \n
    Development
    \n
    \n \n 40% Complete (success)\n
    \n
    \n \n \n\n
    \n
    \n \n\n \n
    \n

    Our Team

    \n
    \n
    \n\n \n\n
    \n
    \n
    \n \n
    \n \n \"\"\n \n

    Johne Doe

    \n Creative\n
    \n
    \n
    \n \n
    \n \n \"\"\n \n

    Jennifer

    \n Programmer\n
    \n
    \n
    \n \n
    \n \n \"\"\n \n

    Christean

    \n CEO\n
    \n
    \n
    \n \n
    \n \n \"\"\n \n

    Kerinele rase

    \n Manager\n
    \n
    \n
    \n
    \n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n" }, { "path": "backend/tests/integration/tests/pruning/website/contact.html", "content": "\n\n \n \n Above Multi-purpose Free Bootstrap Responsive Template\n \n \n \n \n \n \n \n \n \n\n \n \n \n \n
    \n \n
    \n
    \n
    \n
    \n \n \n \n \n \n \"logo\"\n\n
    \n
    \n \n
    \n
    \n
    \n
    \n \n
    \n
    \n
    \n
    \n

    Contact Us

    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n \n
    \n
    \n \n \n \n
    \n\n
    \n
    \n
    \n
    \n \n \n \n
    \n

    Settings

    \n
    \n\n
    \n
    \n
    GENERAL
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n \n
    \n
    \n Please enter a valid URL starting with http:// or https://\n
    \n
    \n
    \n\n \n
    \n
    \n
    \n\n \n \n\n" }, { "path": "desktop/src/titlebar.js", "content": "// Custom title bar for Onyx Desktop\n// This script injects a draggable title bar that matches Onyx design system\n\n(function () {\n const TITLEBAR_ID = \"onyx-desktop-titlebar\";\n const TITLEBAR_HEIGHT = 36;\n const STYLE_ID = \"onyx-desktop-titlebar-style\";\n const VIEWPORT_VAR = \"--onyx-desktop-viewport-height\";\n\n // Wait for DOM to be ready\n if (document.readyState === \"loading\") {\n document.addEventListener(\"DOMContentLoaded\", init);\n } else {\n init();\n }\n\n function getInvoke() {\n if (window.__TAURI__?.core?.invoke) return window.__TAURI__.core.invoke;\n if (window.__TAURI__?.invoke) return window.__TAURI__.invoke;\n if (window.__TAURI_INTERNALS__?.invoke)\n return window.__TAURI_INTERNALS__.invoke;\n return null;\n }\n\n async function startWindowDrag() {\n const invoke = getInvoke();\n\n if (invoke) {\n try {\n await invoke(\"start_drag_window\");\n return;\n } catch (err) {}\n }\n\n const appWindow =\n window.__TAURI__?.window?.getCurrent?.() ??\n window.__TAURI__?.window?.appWindow;\n\n if (appWindow?.startDragging) {\n try {\n await appWindow.startDragging();\n } catch (err) {}\n }\n }\n\n function injectStyles() {\n if (document.getElementById(STYLE_ID)) return;\n const style = document.createElement(\"style\");\n style.id = STYLE_ID;\n style.textContent = `\n :root {\n --onyx-desktop-titlebar-height: ${TITLEBAR_HEIGHT}px;\n --onyx-desktop-viewport-height: 100dvh;\n --onyx-desktop-safe-height: calc(var(--onyx-desktop-viewport-height) - var(--onyx-desktop-titlebar-height));\n }\n\n @supports not (height: 100dvh) {\n :root {\n --onyx-desktop-viewport-height: 100vh;\n }\n }\n\n html,\n body {\n height: var(--onyx-desktop-viewport-height);\n min-height: var(--onyx-desktop-viewport-height);\n margin: 0;\n padding: 0;\n overflow: hidden;\n }\n\n body {\n padding-top: var(--onyx-desktop-titlebar-height) !important;\n box-sizing: border-box;\n }\n\n body > div#__next,\n body > div#root,\n body > main {\n height: var(--onyx-desktop-safe-height);\n min-height: var(--onyx-desktop-safe-height);\n overflow: auto;\n }\n\n /* Override common Tailwind viewport helpers so content fits under the titlebar */\n .h-screen {\n height: var(--onyx-desktop-safe-height) !important;\n }\n\n .min-h-screen {\n min-height: var(--onyx-desktop-safe-height) !important;\n }\n\n .max-h-screen {\n max-height: var(--onyx-desktop-safe-height) !important;\n }\n\n #${TITLEBAR_ID} {\n cursor: default !important;\n -webkit-user-select: none !important;\n user-select: none !important;\n -webkit-app-region: drag;\n background: rgba(255, 255, 255, 0.85);\n height: var(--onyx-desktop-titlebar-height);\n }\n\n /* Dark mode support */\n .dark #${TITLEBAR_ID} {\n background: linear-gradient(180deg, rgba(18, 18, 18, 0.82) 0%, rgba(18, 18, 18, 0.72) 100%);\n border-bottom-color: rgba(255, 255, 255, 0.08);\n }\n `;\n document.head.appendChild(style);\n }\n\n function updateTitleBarTheme(isDark) {\n const titleBar = document.getElementById(TITLEBAR_ID);\n if (!titleBar) return;\n\n if (isDark) {\n titleBar.style.background =\n \"linear-gradient(180deg, rgba(18, 18, 18, 0.82) 0%, rgba(18, 18, 18, 0.72) 100%)\";\n titleBar.style.borderBottom = \"1px solid rgba(255, 255, 255, 0.08)\";\n titleBar.style.boxShadow = \"0 8px 28px rgba(0, 0, 0, 0.2)\";\n } else {\n titleBar.style.background =\n \"linear-gradient(180deg, rgba(255, 255, 255, 0.94) 0%, rgba(255, 255, 255, 0.78) 100%)\";\n titleBar.style.borderBottom = \"1px solid rgba(0, 0, 0, 0.06)\";\n titleBar.style.boxShadow = \"0 8px 28px rgba(0, 0, 0, 0.04)\";\n }\n }\n\n function buildTitleBar() {\n const titleBar = document.createElement(\"div\");\n titleBar.id = TITLEBAR_ID;\n titleBar.setAttribute(\"data-tauri-drag-region\", \"\");\n\n titleBar.addEventListener(\"mousedown\", (e) => {\n // Only start drag on left click and not on buttons/inputs\n const nonDraggable = [\n \"BUTTON\",\n \"INPUT\",\n \"TEXTAREA\",\n \"A\",\n \"SELECT\",\n \"OPTION\",\n ];\n if (e.button === 0 && !nonDraggable.includes(e.target.tagName)) {\n e.preventDefault();\n startWindowDrag();\n }\n });\n\n // Apply initial styles matching current theme\n const htmlHasDark = document.documentElement.classList.contains(\"dark\");\n const bodyHasDark = document.body?.classList.contains(\"dark\");\n const isDark = htmlHasDark || bodyHasDark;\n\n // Apply styles matching Onyx design system with translucent glass effect\n titleBar.style.cssText = `\n position: fixed;\n top: 0;\n left: 0;\n right: 0;\n height: ${TITLEBAR_HEIGHT}px;\n background: linear-gradient(180deg, rgba(255, 255, 255, 0.94) 0%, rgba(255, 255, 255, 0.78) 100%);\n border-bottom: 1px solid rgba(0, 0, 0, 0.06);\n box-shadow: 0 8px 28px rgba(0, 0, 0, 0.04);\n z-index: 999999;\n display: flex;\n align-items: center;\n justify-content: center;\n cursor: default;\n user-select: none;\n -webkit-user-select: none;\n font-family: 'Hanken Grotesk', -apple-system, BlinkMacSystemFont, sans-serif;\n backdrop-filter: blur(18px) saturate(180%);\n -webkit-backdrop-filter: blur(18px) saturate(180%);\n -webkit-app-region: drag;\n padding: 0 12px;\n transition: background 0.3s ease, border-bottom 0.3s ease, box-shadow 0.3s ease;\n `;\n\n // Apply correct theme\n updateTitleBarTheme(isDark);\n\n return titleBar;\n }\n\n function mountTitleBar() {\n if (!document.body) {\n return;\n }\n\n const existing = document.getElementById(TITLEBAR_ID);\n if (existing?.parentElement === document.body) {\n // Update theme on existing titlebar\n const htmlHasDark = document.documentElement.classList.contains(\"dark\");\n const bodyHasDark = document.body?.classList.contains(\"dark\");\n const isDark = htmlHasDark || bodyHasDark;\n updateTitleBarTheme(isDark);\n return;\n }\n\n if (existing) {\n existing.remove();\n }\n\n const titleBar = buildTitleBar();\n document.body.insertBefore(titleBar, document.body.firstChild);\n injectStyles();\n\n // Ensure theme is applied immediately after mount\n setTimeout(() => {\n const htmlHasDark = document.documentElement.classList.contains(\"dark\");\n const bodyHasDark = document.body?.classList.contains(\"dark\");\n const isDark = htmlHasDark || bodyHasDark;\n updateTitleBarTheme(isDark);\n }, 0);\n }\n\n function syncViewportHeight() {\n const viewportHeight =\n window.visualViewport?.height ??\n document.documentElement?.clientHeight ??\n window.innerHeight;\n\n if (viewportHeight) {\n document.documentElement.style.setProperty(\n VIEWPORT_VAR,\n `${viewportHeight}px`,\n );\n }\n }\n\n function observeThemeChanges() {\n let lastKnownTheme = null;\n\n function checkAndUpdateTheme() {\n // Check both html and body for dark class (some apps use body)\n const htmlHasDark = document.documentElement.classList.contains(\"dark\");\n const bodyHasDark = document.body?.classList.contains(\"dark\");\n const isDark = htmlHasDark || bodyHasDark;\n\n if (lastKnownTheme !== isDark) {\n lastKnownTheme = isDark;\n updateTitleBarTheme(isDark);\n }\n }\n\n // Immediate check on setup\n checkAndUpdateTheme();\n\n // Watch for theme changes on the HTML element\n const themeObserver = new MutationObserver(() => {\n checkAndUpdateTheme();\n });\n\n themeObserver.observe(document.documentElement, {\n attributes: true,\n attributeFilter: [\"class\"],\n });\n\n // Also observe body if it exists\n if (document.body) {\n const bodyObserver = new MutationObserver(() => {\n checkAndUpdateTheme();\n });\n bodyObserver.observe(document.body, {\n attributes: true,\n attributeFilter: [\"class\"],\n });\n }\n\n // Also check periodically in case classList is manipulated directly\n // or the theme loads asynchronously after page load\n const intervalId = setInterval(() => {\n checkAndUpdateTheme();\n }, 300);\n\n // Clean up after 30 seconds once theme should be stable\n setTimeout(() => {\n clearInterval(intervalId);\n // But keep checking every 2 seconds for manual theme changes\n setInterval(() => {\n checkAndUpdateTheme();\n }, 2000);\n }, 30000);\n }\n\n function init() {\n mountTitleBar();\n syncViewportHeight();\n observeThemeChanges();\n\n window.addEventListener(\"resize\", syncViewportHeight, { passive: true });\n window.visualViewport?.addEventListener(\"resize\", syncViewportHeight, {\n passive: true,\n });\n\n // Keep it around even if the app DOM re-renders\n const observer = new MutationObserver(() => {\n if (!document.getElementById(TITLEBAR_ID)) {\n mountTitleBar();\n }\n });\n\n observer.observe(document.documentElement, {\n childList: true,\n subtree: true,\n });\n\n // Fallback keep-alive check\n setInterval(() => {\n if (!document.getElementById(TITLEBAR_ID)) {\n mountTitleBar();\n }\n }, 1500);\n }\n})();\n" }, { "path": "desktop/src-tauri/Cargo.toml", "content": "[package]\nname = \"onyx\"\nversion = \"0.0.0-dev\"\ndescription = \"Lightweight desktop app for Onyx Cloud\"\nauthors = [\"you\"]\nedition = \"2021\"\n\n[build-dependencies]\ntauri-build = { version = \"2.5\", features = [] }\n\n[dependencies]\ntauri = { version = \"2.10\", features = [\"macos-private-api\", \"tray-icon\", \"image-png\"] }\ntauri-plugin-shell = \"2.3.5\"\ntauri-plugin-window-state = \"2.4.1\"\nserde = { version = \"1.0\", features = [\"derive\"] }\nserde_json = \"1.0\"\nuuid = { version = \"1.0\", features = [\"v4\"] }\ndirectories = \"5.0\"\ntokio = { version = \"1\", features = [\"time\"] }\nwindow-vibrancy = \"0.7.1\"\nurl = \"2.5\"\n\n[features]\ndefault = [\"custom-protocol\"]\ncustom-protocol = [\"tauri/custom-protocol\"]\ndevtools = [\"tauri/devtools\"]\n" }, { "path": "desktop/src-tauri/build.rs", "content": "fn main() {\n tauri_build::build()\n}\n" }, { "path": "desktop/src-tauri/gen/schemas/acl-manifests.json", "content": "{\"core\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default core plugins set.\",\"permissions\":[\"core:path:default\",\"core:event:default\",\"core:window:default\",\"core:webview:default\",\"core:app:default\",\"core:image:default\",\"core:resources:default\",\"core:menu:default\",\"core:tray:default\"]},\"permissions\":{},\"permission_sets\":{},\"global_scope_schema\":null},\"core:app\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin.\",\"permissions\":[\"allow-version\",\"allow-name\",\"allow-tauri-version\",\"allow-identifier\",\"allow-bundle-type\",\"allow-register-listener\",\"allow-remove-listener\"]},\"permissions\":{\"allow-app-hide\":{\"identifier\":\"allow-app-hide\",\"description\":\"Enables the app_hide command without any pre-configured scope.\",\"commands\":{\"allow\":[\"app_hide\"],\"deny\":[]}},\"allow-app-show\":{\"identifier\":\"allow-app-show\",\"description\":\"Enables the app_show command without any pre-configured scope.\",\"commands\":{\"allow\":[\"app_show\"],\"deny\":[]}},\"allow-bundle-type\":{\"identifier\":\"allow-bundle-type\",\"description\":\"Enables the bundle_type command without any pre-configured scope.\",\"commands\":{\"allow\":[\"bundle_type\"],\"deny\":[]}},\"allow-default-window-icon\":{\"identifier\":\"allow-default-window-icon\",\"description\":\"Enables the default_window_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[\"default_window_icon\"],\"deny\":[]}},\"allow-fetch-data-store-identifiers\":{\"identifier\":\"allow-fetch-data-store-identifiers\",\"description\":\"Enables the fetch_data_store_identifiers command without any pre-configured scope.\",\"commands\":{\"allow\":[\"fetch_data_store_identifiers\"],\"deny\":[]}},\"allow-identifier\":{\"identifier\":\"allow-identifier\",\"description\":\"Enables the identifier command without any pre-configured scope.\",\"commands\":{\"allow\":[\"identifier\"],\"deny\":[]}},\"allow-name\":{\"identifier\":\"allow-name\",\"description\":\"Enables the name command without any pre-configured scope.\",\"commands\":{\"allow\":[\"name\"],\"deny\":[]}},\"allow-register-listener\":{\"identifier\":\"allow-register-listener\",\"description\":\"Enables the register_listener command without any pre-configured scope.\",\"commands\":{\"allow\":[\"register_listener\"],\"deny\":[]}},\"allow-remove-data-store\":{\"identifier\":\"allow-remove-data-store\",\"description\":\"Enables the remove_data_store command without any pre-configured scope.\",\"commands\":{\"allow\":[\"remove_data_store\"],\"deny\":[]}},\"allow-remove-listener\":{\"identifier\":\"allow-remove-listener\",\"description\":\"Enables the remove_listener command without any pre-configured scope.\",\"commands\":{\"allow\":[\"remove_listener\"],\"deny\":[]}},\"allow-set-app-theme\":{\"identifier\":\"allow-set-app-theme\",\"description\":\"Enables the set_app_theme command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_app_theme\"],\"deny\":[]}},\"allow-set-dock-visibility\":{\"identifier\":\"allow-set-dock-visibility\",\"description\":\"Enables the set_dock_visibility command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_dock_visibility\"],\"deny\":[]}},\"allow-tauri-version\":{\"identifier\":\"allow-tauri-version\",\"description\":\"Enables the tauri_version command without any pre-configured scope.\",\"commands\":{\"allow\":[\"tauri_version\"],\"deny\":[]}},\"allow-version\":{\"identifier\":\"allow-version\",\"description\":\"Enables the version command without any pre-configured scope.\",\"commands\":{\"allow\":[\"version\"],\"deny\":[]}},\"deny-app-hide\":{\"identifier\":\"deny-app-hide\",\"description\":\"Denies the app_hide command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"app_hide\"]}},\"deny-app-show\":{\"identifier\":\"deny-app-show\",\"description\":\"Denies the app_show command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"app_show\"]}},\"deny-bundle-type\":{\"identifier\":\"deny-bundle-type\",\"description\":\"Denies the bundle_type command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"bundle_type\"]}},\"deny-default-window-icon\":{\"identifier\":\"deny-default-window-icon\",\"description\":\"Denies the default_window_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"default_window_icon\"]}},\"deny-fetch-data-store-identifiers\":{\"identifier\":\"deny-fetch-data-store-identifiers\",\"description\":\"Denies the fetch_data_store_identifiers command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"fetch_data_store_identifiers\"]}},\"deny-identifier\":{\"identifier\":\"deny-identifier\",\"description\":\"Denies the identifier command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"identifier\"]}},\"deny-name\":{\"identifier\":\"deny-name\",\"description\":\"Denies the name command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"name\"]}},\"deny-register-listener\":{\"identifier\":\"deny-register-listener\",\"description\":\"Denies the register_listener command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"register_listener\"]}},\"deny-remove-data-store\":{\"identifier\":\"deny-remove-data-store\",\"description\":\"Denies the remove_data_store command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"remove_data_store\"]}},\"deny-remove-listener\":{\"identifier\":\"deny-remove-listener\",\"description\":\"Denies the remove_listener command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"remove_listener\"]}},\"deny-set-app-theme\":{\"identifier\":\"deny-set-app-theme\",\"description\":\"Denies the set_app_theme command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_app_theme\"]}},\"deny-set-dock-visibility\":{\"identifier\":\"deny-set-dock-visibility\",\"description\":\"Denies the set_dock_visibility command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_dock_visibility\"]}},\"deny-tauri-version\":{\"identifier\":\"deny-tauri-version\",\"description\":\"Denies the tauri_version command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"tauri_version\"]}},\"deny-version\":{\"identifier\":\"deny-version\",\"description\":\"Denies the version command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"version\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:event\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin, which enables all commands.\",\"permissions\":[\"allow-listen\",\"allow-unlisten\",\"allow-emit\",\"allow-emit-to\"]},\"permissions\":{\"allow-emit\":{\"identifier\":\"allow-emit\",\"description\":\"Enables the emit command without any pre-configured scope.\",\"commands\":{\"allow\":[\"emit\"],\"deny\":[]}},\"allow-emit-to\":{\"identifier\":\"allow-emit-to\",\"description\":\"Enables the emit_to command without any pre-configured scope.\",\"commands\":{\"allow\":[\"emit_to\"],\"deny\":[]}},\"allow-listen\":{\"identifier\":\"allow-listen\",\"description\":\"Enables the listen command without any pre-configured scope.\",\"commands\":{\"allow\":[\"listen\"],\"deny\":[]}},\"allow-unlisten\":{\"identifier\":\"allow-unlisten\",\"description\":\"Enables the unlisten command without any pre-configured scope.\",\"commands\":{\"allow\":[\"unlisten\"],\"deny\":[]}},\"deny-emit\":{\"identifier\":\"deny-emit\",\"description\":\"Denies the emit command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"emit\"]}},\"deny-emit-to\":{\"identifier\":\"deny-emit-to\",\"description\":\"Denies the emit_to command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"emit_to\"]}},\"deny-listen\":{\"identifier\":\"deny-listen\",\"description\":\"Denies the listen command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"listen\"]}},\"deny-unlisten\":{\"identifier\":\"deny-unlisten\",\"description\":\"Denies the unlisten command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"unlisten\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:image\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin, which enables all commands.\",\"permissions\":[\"allow-new\",\"allow-from-bytes\",\"allow-from-path\",\"allow-rgba\",\"allow-size\"]},\"permissions\":{\"allow-from-bytes\":{\"identifier\":\"allow-from-bytes\",\"description\":\"Enables the from_bytes command without any pre-configured scope.\",\"commands\":{\"allow\":[\"from_bytes\"],\"deny\":[]}},\"allow-from-path\":{\"identifier\":\"allow-from-path\",\"description\":\"Enables the from_path command without any pre-configured scope.\",\"commands\":{\"allow\":[\"from_path\"],\"deny\":[]}},\"allow-new\":{\"identifier\":\"allow-new\",\"description\":\"Enables the new command without any pre-configured scope.\",\"commands\":{\"allow\":[\"new\"],\"deny\":[]}},\"allow-rgba\":{\"identifier\":\"allow-rgba\",\"description\":\"Enables the rgba command without any pre-configured scope.\",\"commands\":{\"allow\":[\"rgba\"],\"deny\":[]}},\"allow-size\":{\"identifier\":\"allow-size\",\"description\":\"Enables the size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"size\"],\"deny\":[]}},\"deny-from-bytes\":{\"identifier\":\"deny-from-bytes\",\"description\":\"Denies the from_bytes command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"from_bytes\"]}},\"deny-from-path\":{\"identifier\":\"deny-from-path\",\"description\":\"Denies the from_path command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"from_path\"]}},\"deny-new\":{\"identifier\":\"deny-new\",\"description\":\"Denies the new command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"new\"]}},\"deny-rgba\":{\"identifier\":\"deny-rgba\",\"description\":\"Denies the rgba command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"rgba\"]}},\"deny-size\":{\"identifier\":\"deny-size\",\"description\":\"Denies the size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"size\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:menu\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin, which enables all commands.\",\"permissions\":[\"allow-new\",\"allow-append\",\"allow-prepend\",\"allow-insert\",\"allow-remove\",\"allow-remove-at\",\"allow-items\",\"allow-get\",\"allow-popup\",\"allow-create-default\",\"allow-set-as-app-menu\",\"allow-set-as-window-menu\",\"allow-text\",\"allow-set-text\",\"allow-is-enabled\",\"allow-set-enabled\",\"allow-set-accelerator\",\"allow-set-as-windows-menu-for-nsapp\",\"allow-set-as-help-menu-for-nsapp\",\"allow-is-checked\",\"allow-set-checked\",\"allow-set-icon\"]},\"permissions\":{\"allow-append\":{\"identifier\":\"allow-append\",\"description\":\"Enables the append command without any pre-configured scope.\",\"commands\":{\"allow\":[\"append\"],\"deny\":[]}},\"allow-create-default\":{\"identifier\":\"allow-create-default\",\"description\":\"Enables the create_default command without any pre-configured scope.\",\"commands\":{\"allow\":[\"create_default\"],\"deny\":[]}},\"allow-get\":{\"identifier\":\"allow-get\",\"description\":\"Enables the get command without any pre-configured scope.\",\"commands\":{\"allow\":[\"get\"],\"deny\":[]}},\"allow-insert\":{\"identifier\":\"allow-insert\",\"description\":\"Enables the insert command without any pre-configured scope.\",\"commands\":{\"allow\":[\"insert\"],\"deny\":[]}},\"allow-is-checked\":{\"identifier\":\"allow-is-checked\",\"description\":\"Enables the is_checked command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_checked\"],\"deny\":[]}},\"allow-is-enabled\":{\"identifier\":\"allow-is-enabled\",\"description\":\"Enables the is_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_enabled\"],\"deny\":[]}},\"allow-items\":{\"identifier\":\"allow-items\",\"description\":\"Enables the items command without any pre-configured scope.\",\"commands\":{\"allow\":[\"items\"],\"deny\":[]}},\"allow-new\":{\"identifier\":\"allow-new\",\"description\":\"Enables the new command without any pre-configured scope.\",\"commands\":{\"allow\":[\"new\"],\"deny\":[]}},\"allow-popup\":{\"identifier\":\"allow-popup\",\"description\":\"Enables the popup command without any pre-configured scope.\",\"commands\":{\"allow\":[\"popup\"],\"deny\":[]}},\"allow-prepend\":{\"identifier\":\"allow-prepend\",\"description\":\"Enables the prepend command without any pre-configured scope.\",\"commands\":{\"allow\":[\"prepend\"],\"deny\":[]}},\"allow-remove\":{\"identifier\":\"allow-remove\",\"description\":\"Enables the remove command without any pre-configured scope.\",\"commands\":{\"allow\":[\"remove\"],\"deny\":[]}},\"allow-remove-at\":{\"identifier\":\"allow-remove-at\",\"description\":\"Enables the remove_at command without any pre-configured scope.\",\"commands\":{\"allow\":[\"remove_at\"],\"deny\":[]}},\"allow-set-accelerator\":{\"identifier\":\"allow-set-accelerator\",\"description\":\"Enables the set_accelerator command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_accelerator\"],\"deny\":[]}},\"allow-set-as-app-menu\":{\"identifier\":\"allow-set-as-app-menu\",\"description\":\"Enables the set_as_app_menu command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_as_app_menu\"],\"deny\":[]}},\"allow-set-as-help-menu-for-nsapp\":{\"identifier\":\"allow-set-as-help-menu-for-nsapp\",\"description\":\"Enables the set_as_help_menu_for_nsapp command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_as_help_menu_for_nsapp\"],\"deny\":[]}},\"allow-set-as-window-menu\":{\"identifier\":\"allow-set-as-window-menu\",\"description\":\"Enables the set_as_window_menu command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_as_window_menu\"],\"deny\":[]}},\"allow-set-as-windows-menu-for-nsapp\":{\"identifier\":\"allow-set-as-windows-menu-for-nsapp\",\"description\":\"Enables the set_as_windows_menu_for_nsapp command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_as_windows_menu_for_nsapp\"],\"deny\":[]}},\"allow-set-checked\":{\"identifier\":\"allow-set-checked\",\"description\":\"Enables the set_checked command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_checked\"],\"deny\":[]}},\"allow-set-enabled\":{\"identifier\":\"allow-set-enabled\",\"description\":\"Enables the set_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_enabled\"],\"deny\":[]}},\"allow-set-icon\":{\"identifier\":\"allow-set-icon\",\"description\":\"Enables the set_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_icon\"],\"deny\":[]}},\"allow-set-text\":{\"identifier\":\"allow-set-text\",\"description\":\"Enables the set_text command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_text\"],\"deny\":[]}},\"allow-text\":{\"identifier\":\"allow-text\",\"description\":\"Enables the text command without any pre-configured scope.\",\"commands\":{\"allow\":[\"text\"],\"deny\":[]}},\"deny-append\":{\"identifier\":\"deny-append\",\"description\":\"Denies the append command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"append\"]}},\"deny-create-default\":{\"identifier\":\"deny-create-default\",\"description\":\"Denies the create_default command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"create_default\"]}},\"deny-get\":{\"identifier\":\"deny-get\",\"description\":\"Denies the get command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"get\"]}},\"deny-insert\":{\"identifier\":\"deny-insert\",\"description\":\"Denies the insert command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"insert\"]}},\"deny-is-checked\":{\"identifier\":\"deny-is-checked\",\"description\":\"Denies the is_checked command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_checked\"]}},\"deny-is-enabled\":{\"identifier\":\"deny-is-enabled\",\"description\":\"Denies the is_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_enabled\"]}},\"deny-items\":{\"identifier\":\"deny-items\",\"description\":\"Denies the items command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"items\"]}},\"deny-new\":{\"identifier\":\"deny-new\",\"description\":\"Denies the new command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"new\"]}},\"deny-popup\":{\"identifier\":\"deny-popup\",\"description\":\"Denies the popup command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"popup\"]}},\"deny-prepend\":{\"identifier\":\"deny-prepend\",\"description\":\"Denies the prepend command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"prepend\"]}},\"deny-remove\":{\"identifier\":\"deny-remove\",\"description\":\"Denies the remove command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"remove\"]}},\"deny-remove-at\":{\"identifier\":\"deny-remove-at\",\"description\":\"Denies the remove_at command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"remove_at\"]}},\"deny-set-accelerator\":{\"identifier\":\"deny-set-accelerator\",\"description\":\"Denies the set_accelerator command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_accelerator\"]}},\"deny-set-as-app-menu\":{\"identifier\":\"deny-set-as-app-menu\",\"description\":\"Denies the set_as_app_menu command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_as_app_menu\"]}},\"deny-set-as-help-menu-for-nsapp\":{\"identifier\":\"deny-set-as-help-menu-for-nsapp\",\"description\":\"Denies the set_as_help_menu_for_nsapp command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_as_help_menu_for_nsapp\"]}},\"deny-set-as-window-menu\":{\"identifier\":\"deny-set-as-window-menu\",\"description\":\"Denies the set_as_window_menu command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_as_window_menu\"]}},\"deny-set-as-windows-menu-for-nsapp\":{\"identifier\":\"deny-set-as-windows-menu-for-nsapp\",\"description\":\"Denies the set_as_windows_menu_for_nsapp command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_as_windows_menu_for_nsapp\"]}},\"deny-set-checked\":{\"identifier\":\"deny-set-checked\",\"description\":\"Denies the set_checked command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_checked\"]}},\"deny-set-enabled\":{\"identifier\":\"deny-set-enabled\",\"description\":\"Denies the set_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_enabled\"]}},\"deny-set-icon\":{\"identifier\":\"deny-set-icon\",\"description\":\"Denies the set_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_icon\"]}},\"deny-set-text\":{\"identifier\":\"deny-set-text\",\"description\":\"Denies the set_text command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_text\"]}},\"deny-text\":{\"identifier\":\"deny-text\",\"description\":\"Denies the text command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"text\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:path\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin, which enables all commands.\",\"permissions\":[\"allow-resolve-directory\",\"allow-resolve\",\"allow-normalize\",\"allow-join\",\"allow-dirname\",\"allow-extname\",\"allow-basename\",\"allow-is-absolute\"]},\"permissions\":{\"allow-basename\":{\"identifier\":\"allow-basename\",\"description\":\"Enables the basename command without any pre-configured scope.\",\"commands\":{\"allow\":[\"basename\"],\"deny\":[]}},\"allow-dirname\":{\"identifier\":\"allow-dirname\",\"description\":\"Enables the dirname command without any pre-configured scope.\",\"commands\":{\"allow\":[\"dirname\"],\"deny\":[]}},\"allow-extname\":{\"identifier\":\"allow-extname\",\"description\":\"Enables the extname command without any pre-configured scope.\",\"commands\":{\"allow\":[\"extname\"],\"deny\":[]}},\"allow-is-absolute\":{\"identifier\":\"allow-is-absolute\",\"description\":\"Enables the is_absolute command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_absolute\"],\"deny\":[]}},\"allow-join\":{\"identifier\":\"allow-join\",\"description\":\"Enables the join command without any pre-configured scope.\",\"commands\":{\"allow\":[\"join\"],\"deny\":[]}},\"allow-normalize\":{\"identifier\":\"allow-normalize\",\"description\":\"Enables the normalize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"normalize\"],\"deny\":[]}},\"allow-resolve\":{\"identifier\":\"allow-resolve\",\"description\":\"Enables the resolve command without any pre-configured scope.\",\"commands\":{\"allow\":[\"resolve\"],\"deny\":[]}},\"allow-resolve-directory\":{\"identifier\":\"allow-resolve-directory\",\"description\":\"Enables the resolve_directory command without any pre-configured scope.\",\"commands\":{\"allow\":[\"resolve_directory\"],\"deny\":[]}},\"deny-basename\":{\"identifier\":\"deny-basename\",\"description\":\"Denies the basename command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"basename\"]}},\"deny-dirname\":{\"identifier\":\"deny-dirname\",\"description\":\"Denies the dirname command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"dirname\"]}},\"deny-extname\":{\"identifier\":\"deny-extname\",\"description\":\"Denies the extname command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"extname\"]}},\"deny-is-absolute\":{\"identifier\":\"deny-is-absolute\",\"description\":\"Denies the is_absolute command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_absolute\"]}},\"deny-join\":{\"identifier\":\"deny-join\",\"description\":\"Denies the join command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"join\"]}},\"deny-normalize\":{\"identifier\":\"deny-normalize\",\"description\":\"Denies the normalize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"normalize\"]}},\"deny-resolve\":{\"identifier\":\"deny-resolve\",\"description\":\"Denies the resolve command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"resolve\"]}},\"deny-resolve-directory\":{\"identifier\":\"deny-resolve-directory\",\"description\":\"Denies the resolve_directory command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"resolve_directory\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:resources\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin, which enables all commands.\",\"permissions\":[\"allow-close\"]},\"permissions\":{\"allow-close\":{\"identifier\":\"allow-close\",\"description\":\"Enables the close command without any pre-configured scope.\",\"commands\":{\"allow\":[\"close\"],\"deny\":[]}},\"deny-close\":{\"identifier\":\"deny-close\",\"description\":\"Denies the close command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"close\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:tray\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin, which enables all commands.\",\"permissions\":[\"allow-new\",\"allow-get-by-id\",\"allow-remove-by-id\",\"allow-set-icon\",\"allow-set-menu\",\"allow-set-tooltip\",\"allow-set-title\",\"allow-set-visible\",\"allow-set-temp-dir-path\",\"allow-set-icon-as-template\",\"allow-set-show-menu-on-left-click\"]},\"permissions\":{\"allow-get-by-id\":{\"identifier\":\"allow-get-by-id\",\"description\":\"Enables the get_by_id command without any pre-configured scope.\",\"commands\":{\"allow\":[\"get_by_id\"],\"deny\":[]}},\"allow-new\":{\"identifier\":\"allow-new\",\"description\":\"Enables the new command without any pre-configured scope.\",\"commands\":{\"allow\":[\"new\"],\"deny\":[]}},\"allow-remove-by-id\":{\"identifier\":\"allow-remove-by-id\",\"description\":\"Enables the remove_by_id command without any pre-configured scope.\",\"commands\":{\"allow\":[\"remove_by_id\"],\"deny\":[]}},\"allow-set-icon\":{\"identifier\":\"allow-set-icon\",\"description\":\"Enables the set_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_icon\"],\"deny\":[]}},\"allow-set-icon-as-template\":{\"identifier\":\"allow-set-icon-as-template\",\"description\":\"Enables the set_icon_as_template command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_icon_as_template\"],\"deny\":[]}},\"allow-set-menu\":{\"identifier\":\"allow-set-menu\",\"description\":\"Enables the set_menu command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_menu\"],\"deny\":[]}},\"allow-set-show-menu-on-left-click\":{\"identifier\":\"allow-set-show-menu-on-left-click\",\"description\":\"Enables the set_show_menu_on_left_click command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_show_menu_on_left_click\"],\"deny\":[]}},\"allow-set-temp-dir-path\":{\"identifier\":\"allow-set-temp-dir-path\",\"description\":\"Enables the set_temp_dir_path command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_temp_dir_path\"],\"deny\":[]}},\"allow-set-title\":{\"identifier\":\"allow-set-title\",\"description\":\"Enables the set_title command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_title\"],\"deny\":[]}},\"allow-set-tooltip\":{\"identifier\":\"allow-set-tooltip\",\"description\":\"Enables the set_tooltip command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_tooltip\"],\"deny\":[]}},\"allow-set-visible\":{\"identifier\":\"allow-set-visible\",\"description\":\"Enables the set_visible command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_visible\"],\"deny\":[]}},\"deny-get-by-id\":{\"identifier\":\"deny-get-by-id\",\"description\":\"Denies the get_by_id command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"get_by_id\"]}},\"deny-new\":{\"identifier\":\"deny-new\",\"description\":\"Denies the new command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"new\"]}},\"deny-remove-by-id\":{\"identifier\":\"deny-remove-by-id\",\"description\":\"Denies the remove_by_id command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"remove_by_id\"]}},\"deny-set-icon\":{\"identifier\":\"deny-set-icon\",\"description\":\"Denies the set_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_icon\"]}},\"deny-set-icon-as-template\":{\"identifier\":\"deny-set-icon-as-template\",\"description\":\"Denies the set_icon_as_template command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_icon_as_template\"]}},\"deny-set-menu\":{\"identifier\":\"deny-set-menu\",\"description\":\"Denies the set_menu command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_menu\"]}},\"deny-set-show-menu-on-left-click\":{\"identifier\":\"deny-set-show-menu-on-left-click\",\"description\":\"Denies the set_show_menu_on_left_click command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_show_menu_on_left_click\"]}},\"deny-set-temp-dir-path\":{\"identifier\":\"deny-set-temp-dir-path\",\"description\":\"Denies the set_temp_dir_path command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_temp_dir_path\"]}},\"deny-set-title\":{\"identifier\":\"deny-set-title\",\"description\":\"Denies the set_title command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_title\"]}},\"deny-set-tooltip\":{\"identifier\":\"deny-set-tooltip\",\"description\":\"Denies the set_tooltip command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_tooltip\"]}},\"deny-set-visible\":{\"identifier\":\"deny-set-visible\",\"description\":\"Denies the set_visible command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_visible\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:webview\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin.\",\"permissions\":[\"allow-get-all-webviews\",\"allow-webview-position\",\"allow-webview-size\",\"allow-internal-toggle-devtools\"]},\"permissions\":{\"allow-clear-all-browsing-data\":{\"identifier\":\"allow-clear-all-browsing-data\",\"description\":\"Enables the clear_all_browsing_data command without any pre-configured scope.\",\"commands\":{\"allow\":[\"clear_all_browsing_data\"],\"deny\":[]}},\"allow-create-webview\":{\"identifier\":\"allow-create-webview\",\"description\":\"Enables the create_webview command without any pre-configured scope.\",\"commands\":{\"allow\":[\"create_webview\"],\"deny\":[]}},\"allow-create-webview-window\":{\"identifier\":\"allow-create-webview-window\",\"description\":\"Enables the create_webview_window command without any pre-configured scope.\",\"commands\":{\"allow\":[\"create_webview_window\"],\"deny\":[]}},\"allow-get-all-webviews\":{\"identifier\":\"allow-get-all-webviews\",\"description\":\"Enables the get_all_webviews command without any pre-configured scope.\",\"commands\":{\"allow\":[\"get_all_webviews\"],\"deny\":[]}},\"allow-internal-toggle-devtools\":{\"identifier\":\"allow-internal-toggle-devtools\",\"description\":\"Enables the internal_toggle_devtools command without any pre-configured scope.\",\"commands\":{\"allow\":[\"internal_toggle_devtools\"],\"deny\":[]}},\"allow-print\":{\"identifier\":\"allow-print\",\"description\":\"Enables the print command without any pre-configured scope.\",\"commands\":{\"allow\":[\"print\"],\"deny\":[]}},\"allow-reparent\":{\"identifier\":\"allow-reparent\",\"description\":\"Enables the reparent command without any pre-configured scope.\",\"commands\":{\"allow\":[\"reparent\"],\"deny\":[]}},\"allow-set-webview-auto-resize\":{\"identifier\":\"allow-set-webview-auto-resize\",\"description\":\"Enables the set_webview_auto_resize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_webview_auto_resize\"],\"deny\":[]}},\"allow-set-webview-background-color\":{\"identifier\":\"allow-set-webview-background-color\",\"description\":\"Enables the set_webview_background_color command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_webview_background_color\"],\"deny\":[]}},\"allow-set-webview-focus\":{\"identifier\":\"allow-set-webview-focus\",\"description\":\"Enables the set_webview_focus command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_webview_focus\"],\"deny\":[]}},\"allow-set-webview-position\":{\"identifier\":\"allow-set-webview-position\",\"description\":\"Enables the set_webview_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_webview_position\"],\"deny\":[]}},\"allow-set-webview-size\":{\"identifier\":\"allow-set-webview-size\",\"description\":\"Enables the set_webview_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_webview_size\"],\"deny\":[]}},\"allow-set-webview-zoom\":{\"identifier\":\"allow-set-webview-zoom\",\"description\":\"Enables the set_webview_zoom command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_webview_zoom\"],\"deny\":[]}},\"allow-webview-close\":{\"identifier\":\"allow-webview-close\",\"description\":\"Enables the webview_close command without any pre-configured scope.\",\"commands\":{\"allow\":[\"webview_close\"],\"deny\":[]}},\"allow-webview-hide\":{\"identifier\":\"allow-webview-hide\",\"description\":\"Enables the webview_hide command without any pre-configured scope.\",\"commands\":{\"allow\":[\"webview_hide\"],\"deny\":[]}},\"allow-webview-position\":{\"identifier\":\"allow-webview-position\",\"description\":\"Enables the webview_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"webview_position\"],\"deny\":[]}},\"allow-webview-show\":{\"identifier\":\"allow-webview-show\",\"description\":\"Enables the webview_show command without any pre-configured scope.\",\"commands\":{\"allow\":[\"webview_show\"],\"deny\":[]}},\"allow-webview-size\":{\"identifier\":\"allow-webview-size\",\"description\":\"Enables the webview_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"webview_size\"],\"deny\":[]}},\"deny-clear-all-browsing-data\":{\"identifier\":\"deny-clear-all-browsing-data\",\"description\":\"Denies the clear_all_browsing_data command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"clear_all_browsing_data\"]}},\"deny-create-webview\":{\"identifier\":\"deny-create-webview\",\"description\":\"Denies the create_webview command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"create_webview\"]}},\"deny-create-webview-window\":{\"identifier\":\"deny-create-webview-window\",\"description\":\"Denies the create_webview_window command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"create_webview_window\"]}},\"deny-get-all-webviews\":{\"identifier\":\"deny-get-all-webviews\",\"description\":\"Denies the get_all_webviews command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"get_all_webviews\"]}},\"deny-internal-toggle-devtools\":{\"identifier\":\"deny-internal-toggle-devtools\",\"description\":\"Denies the internal_toggle_devtools command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"internal_toggle_devtools\"]}},\"deny-print\":{\"identifier\":\"deny-print\",\"description\":\"Denies the print command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"print\"]}},\"deny-reparent\":{\"identifier\":\"deny-reparent\",\"description\":\"Denies the reparent command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"reparent\"]}},\"deny-set-webview-auto-resize\":{\"identifier\":\"deny-set-webview-auto-resize\",\"description\":\"Denies the set_webview_auto_resize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_webview_auto_resize\"]}},\"deny-set-webview-background-color\":{\"identifier\":\"deny-set-webview-background-color\",\"description\":\"Denies the set_webview_background_color command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_webview_background_color\"]}},\"deny-set-webview-focus\":{\"identifier\":\"deny-set-webview-focus\",\"description\":\"Denies the set_webview_focus command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_webview_focus\"]}},\"deny-set-webview-position\":{\"identifier\":\"deny-set-webview-position\",\"description\":\"Denies the set_webview_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_webview_position\"]}},\"deny-set-webview-size\":{\"identifier\":\"deny-set-webview-size\",\"description\":\"Denies the set_webview_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_webview_size\"]}},\"deny-set-webview-zoom\":{\"identifier\":\"deny-set-webview-zoom\",\"description\":\"Denies the set_webview_zoom command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_webview_zoom\"]}},\"deny-webview-close\":{\"identifier\":\"deny-webview-close\",\"description\":\"Denies the webview_close command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"webview_close\"]}},\"deny-webview-hide\":{\"identifier\":\"deny-webview-hide\",\"description\":\"Denies the webview_hide command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"webview_hide\"]}},\"deny-webview-position\":{\"identifier\":\"deny-webview-position\",\"description\":\"Denies the webview_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"webview_position\"]}},\"deny-webview-show\":{\"identifier\":\"deny-webview-show\",\"description\":\"Denies the webview_show command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"webview_show\"]}},\"deny-webview-size\":{\"identifier\":\"deny-webview-size\",\"description\":\"Denies the webview_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"webview_size\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"core:window\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"Default permissions for the plugin.\",\"permissions\":[\"allow-get-all-windows\",\"allow-scale-factor\",\"allow-inner-position\",\"allow-outer-position\",\"allow-inner-size\",\"allow-outer-size\",\"allow-is-fullscreen\",\"allow-is-minimized\",\"allow-is-maximized\",\"allow-is-focused\",\"allow-is-decorated\",\"allow-is-resizable\",\"allow-is-maximizable\",\"allow-is-minimizable\",\"allow-is-closable\",\"allow-is-visible\",\"allow-is-enabled\",\"allow-title\",\"allow-current-monitor\",\"allow-primary-monitor\",\"allow-monitor-from-point\",\"allow-available-monitors\",\"allow-cursor-position\",\"allow-theme\",\"allow-is-always-on-top\",\"allow-internal-toggle-maximize\"]},\"permissions\":{\"allow-available-monitors\":{\"identifier\":\"allow-available-monitors\",\"description\":\"Enables the available_monitors command without any pre-configured scope.\",\"commands\":{\"allow\":[\"available_monitors\"],\"deny\":[]}},\"allow-center\":{\"identifier\":\"allow-center\",\"description\":\"Enables the center command without any pre-configured scope.\",\"commands\":{\"allow\":[\"center\"],\"deny\":[]}},\"allow-close\":{\"identifier\":\"allow-close\",\"description\":\"Enables the close command without any pre-configured scope.\",\"commands\":{\"allow\":[\"close\"],\"deny\":[]}},\"allow-create\":{\"identifier\":\"allow-create\",\"description\":\"Enables the create command without any pre-configured scope.\",\"commands\":{\"allow\":[\"create\"],\"deny\":[]}},\"allow-current-monitor\":{\"identifier\":\"allow-current-monitor\",\"description\":\"Enables the current_monitor command without any pre-configured scope.\",\"commands\":{\"allow\":[\"current_monitor\"],\"deny\":[]}},\"allow-cursor-position\":{\"identifier\":\"allow-cursor-position\",\"description\":\"Enables the cursor_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"cursor_position\"],\"deny\":[]}},\"allow-destroy\":{\"identifier\":\"allow-destroy\",\"description\":\"Enables the destroy command without any pre-configured scope.\",\"commands\":{\"allow\":[\"destroy\"],\"deny\":[]}},\"allow-get-all-windows\":{\"identifier\":\"allow-get-all-windows\",\"description\":\"Enables the get_all_windows command without any pre-configured scope.\",\"commands\":{\"allow\":[\"get_all_windows\"],\"deny\":[]}},\"allow-hide\":{\"identifier\":\"allow-hide\",\"description\":\"Enables the hide command without any pre-configured scope.\",\"commands\":{\"allow\":[\"hide\"],\"deny\":[]}},\"allow-inner-position\":{\"identifier\":\"allow-inner-position\",\"description\":\"Enables the inner_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"inner_position\"],\"deny\":[]}},\"allow-inner-size\":{\"identifier\":\"allow-inner-size\",\"description\":\"Enables the inner_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"inner_size\"],\"deny\":[]}},\"allow-internal-toggle-maximize\":{\"identifier\":\"allow-internal-toggle-maximize\",\"description\":\"Enables the internal_toggle_maximize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"internal_toggle_maximize\"],\"deny\":[]}},\"allow-is-always-on-top\":{\"identifier\":\"allow-is-always-on-top\",\"description\":\"Enables the is_always_on_top command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_always_on_top\"],\"deny\":[]}},\"allow-is-closable\":{\"identifier\":\"allow-is-closable\",\"description\":\"Enables the is_closable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_closable\"],\"deny\":[]}},\"allow-is-decorated\":{\"identifier\":\"allow-is-decorated\",\"description\":\"Enables the is_decorated command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_decorated\"],\"deny\":[]}},\"allow-is-enabled\":{\"identifier\":\"allow-is-enabled\",\"description\":\"Enables the is_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_enabled\"],\"deny\":[]}},\"allow-is-focused\":{\"identifier\":\"allow-is-focused\",\"description\":\"Enables the is_focused command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_focused\"],\"deny\":[]}},\"allow-is-fullscreen\":{\"identifier\":\"allow-is-fullscreen\",\"description\":\"Enables the is_fullscreen command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_fullscreen\"],\"deny\":[]}},\"allow-is-maximizable\":{\"identifier\":\"allow-is-maximizable\",\"description\":\"Enables the is_maximizable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_maximizable\"],\"deny\":[]}},\"allow-is-maximized\":{\"identifier\":\"allow-is-maximized\",\"description\":\"Enables the is_maximized command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_maximized\"],\"deny\":[]}},\"allow-is-minimizable\":{\"identifier\":\"allow-is-minimizable\",\"description\":\"Enables the is_minimizable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_minimizable\"],\"deny\":[]}},\"allow-is-minimized\":{\"identifier\":\"allow-is-minimized\",\"description\":\"Enables the is_minimized command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_minimized\"],\"deny\":[]}},\"allow-is-resizable\":{\"identifier\":\"allow-is-resizable\",\"description\":\"Enables the is_resizable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_resizable\"],\"deny\":[]}},\"allow-is-visible\":{\"identifier\":\"allow-is-visible\",\"description\":\"Enables the is_visible command without any pre-configured scope.\",\"commands\":{\"allow\":[\"is_visible\"],\"deny\":[]}},\"allow-maximize\":{\"identifier\":\"allow-maximize\",\"description\":\"Enables the maximize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"maximize\"],\"deny\":[]}},\"allow-minimize\":{\"identifier\":\"allow-minimize\",\"description\":\"Enables the minimize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"minimize\"],\"deny\":[]}},\"allow-monitor-from-point\":{\"identifier\":\"allow-monitor-from-point\",\"description\":\"Enables the monitor_from_point command without any pre-configured scope.\",\"commands\":{\"allow\":[\"monitor_from_point\"],\"deny\":[]}},\"allow-outer-position\":{\"identifier\":\"allow-outer-position\",\"description\":\"Enables the outer_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"outer_position\"],\"deny\":[]}},\"allow-outer-size\":{\"identifier\":\"allow-outer-size\",\"description\":\"Enables the outer_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"outer_size\"],\"deny\":[]}},\"allow-primary-monitor\":{\"identifier\":\"allow-primary-monitor\",\"description\":\"Enables the primary_monitor command without any pre-configured scope.\",\"commands\":{\"allow\":[\"primary_monitor\"],\"deny\":[]}},\"allow-request-user-attention\":{\"identifier\":\"allow-request-user-attention\",\"description\":\"Enables the request_user_attention command without any pre-configured scope.\",\"commands\":{\"allow\":[\"request_user_attention\"],\"deny\":[]}},\"allow-scale-factor\":{\"identifier\":\"allow-scale-factor\",\"description\":\"Enables the scale_factor command without any pre-configured scope.\",\"commands\":{\"allow\":[\"scale_factor\"],\"deny\":[]}},\"allow-set-always-on-bottom\":{\"identifier\":\"allow-set-always-on-bottom\",\"description\":\"Enables the set_always_on_bottom command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_always_on_bottom\"],\"deny\":[]}},\"allow-set-always-on-top\":{\"identifier\":\"allow-set-always-on-top\",\"description\":\"Enables the set_always_on_top command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_always_on_top\"],\"deny\":[]}},\"allow-set-background-color\":{\"identifier\":\"allow-set-background-color\",\"description\":\"Enables the set_background_color command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_background_color\"],\"deny\":[]}},\"allow-set-badge-count\":{\"identifier\":\"allow-set-badge-count\",\"description\":\"Enables the set_badge_count command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_badge_count\"],\"deny\":[]}},\"allow-set-badge-label\":{\"identifier\":\"allow-set-badge-label\",\"description\":\"Enables the set_badge_label command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_badge_label\"],\"deny\":[]}},\"allow-set-closable\":{\"identifier\":\"allow-set-closable\",\"description\":\"Enables the set_closable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_closable\"],\"deny\":[]}},\"allow-set-content-protected\":{\"identifier\":\"allow-set-content-protected\",\"description\":\"Enables the set_content_protected command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_content_protected\"],\"deny\":[]}},\"allow-set-cursor-grab\":{\"identifier\":\"allow-set-cursor-grab\",\"description\":\"Enables the set_cursor_grab command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_cursor_grab\"],\"deny\":[]}},\"allow-set-cursor-icon\":{\"identifier\":\"allow-set-cursor-icon\",\"description\":\"Enables the set_cursor_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_cursor_icon\"],\"deny\":[]}},\"allow-set-cursor-position\":{\"identifier\":\"allow-set-cursor-position\",\"description\":\"Enables the set_cursor_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_cursor_position\"],\"deny\":[]}},\"allow-set-cursor-visible\":{\"identifier\":\"allow-set-cursor-visible\",\"description\":\"Enables the set_cursor_visible command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_cursor_visible\"],\"deny\":[]}},\"allow-set-decorations\":{\"identifier\":\"allow-set-decorations\",\"description\":\"Enables the set_decorations command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_decorations\"],\"deny\":[]}},\"allow-set-effects\":{\"identifier\":\"allow-set-effects\",\"description\":\"Enables the set_effects command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_effects\"],\"deny\":[]}},\"allow-set-enabled\":{\"identifier\":\"allow-set-enabled\",\"description\":\"Enables the set_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_enabled\"],\"deny\":[]}},\"allow-set-focus\":{\"identifier\":\"allow-set-focus\",\"description\":\"Enables the set_focus command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_focus\"],\"deny\":[]}},\"allow-set-focusable\":{\"identifier\":\"allow-set-focusable\",\"description\":\"Enables the set_focusable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_focusable\"],\"deny\":[]}},\"allow-set-fullscreen\":{\"identifier\":\"allow-set-fullscreen\",\"description\":\"Enables the set_fullscreen command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_fullscreen\"],\"deny\":[]}},\"allow-set-icon\":{\"identifier\":\"allow-set-icon\",\"description\":\"Enables the set_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_icon\"],\"deny\":[]}},\"allow-set-ignore-cursor-events\":{\"identifier\":\"allow-set-ignore-cursor-events\",\"description\":\"Enables the set_ignore_cursor_events command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_ignore_cursor_events\"],\"deny\":[]}},\"allow-set-max-size\":{\"identifier\":\"allow-set-max-size\",\"description\":\"Enables the set_max_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_max_size\"],\"deny\":[]}},\"allow-set-maximizable\":{\"identifier\":\"allow-set-maximizable\",\"description\":\"Enables the set_maximizable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_maximizable\"],\"deny\":[]}},\"allow-set-min-size\":{\"identifier\":\"allow-set-min-size\",\"description\":\"Enables the set_min_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_min_size\"],\"deny\":[]}},\"allow-set-minimizable\":{\"identifier\":\"allow-set-minimizable\",\"description\":\"Enables the set_minimizable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_minimizable\"],\"deny\":[]}},\"allow-set-overlay-icon\":{\"identifier\":\"allow-set-overlay-icon\",\"description\":\"Enables the set_overlay_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_overlay_icon\"],\"deny\":[]}},\"allow-set-position\":{\"identifier\":\"allow-set-position\",\"description\":\"Enables the set_position command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_position\"],\"deny\":[]}},\"allow-set-progress-bar\":{\"identifier\":\"allow-set-progress-bar\",\"description\":\"Enables the set_progress_bar command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_progress_bar\"],\"deny\":[]}},\"allow-set-resizable\":{\"identifier\":\"allow-set-resizable\",\"description\":\"Enables the set_resizable command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_resizable\"],\"deny\":[]}},\"allow-set-shadow\":{\"identifier\":\"allow-set-shadow\",\"description\":\"Enables the set_shadow command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_shadow\"],\"deny\":[]}},\"allow-set-simple-fullscreen\":{\"identifier\":\"allow-set-simple-fullscreen\",\"description\":\"Enables the set_simple_fullscreen command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_simple_fullscreen\"],\"deny\":[]}},\"allow-set-size\":{\"identifier\":\"allow-set-size\",\"description\":\"Enables the set_size command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_size\"],\"deny\":[]}},\"allow-set-size-constraints\":{\"identifier\":\"allow-set-size-constraints\",\"description\":\"Enables the set_size_constraints command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_size_constraints\"],\"deny\":[]}},\"allow-set-skip-taskbar\":{\"identifier\":\"allow-set-skip-taskbar\",\"description\":\"Enables the set_skip_taskbar command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_skip_taskbar\"],\"deny\":[]}},\"allow-set-theme\":{\"identifier\":\"allow-set-theme\",\"description\":\"Enables the set_theme command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_theme\"],\"deny\":[]}},\"allow-set-title\":{\"identifier\":\"allow-set-title\",\"description\":\"Enables the set_title command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_title\"],\"deny\":[]}},\"allow-set-title-bar-style\":{\"identifier\":\"allow-set-title-bar-style\",\"description\":\"Enables the set_title_bar_style command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_title_bar_style\"],\"deny\":[]}},\"allow-set-visible-on-all-workspaces\":{\"identifier\":\"allow-set-visible-on-all-workspaces\",\"description\":\"Enables the set_visible_on_all_workspaces command without any pre-configured scope.\",\"commands\":{\"allow\":[\"set_visible_on_all_workspaces\"],\"deny\":[]}},\"allow-show\":{\"identifier\":\"allow-show\",\"description\":\"Enables the show command without any pre-configured scope.\",\"commands\":{\"allow\":[\"show\"],\"deny\":[]}},\"allow-start-dragging\":{\"identifier\":\"allow-start-dragging\",\"description\":\"Enables the start_dragging command without any pre-configured scope.\",\"commands\":{\"allow\":[\"start_dragging\"],\"deny\":[]}},\"allow-start-resize-dragging\":{\"identifier\":\"allow-start-resize-dragging\",\"description\":\"Enables the start_resize_dragging command without any pre-configured scope.\",\"commands\":{\"allow\":[\"start_resize_dragging\"],\"deny\":[]}},\"allow-theme\":{\"identifier\":\"allow-theme\",\"description\":\"Enables the theme command without any pre-configured scope.\",\"commands\":{\"allow\":[\"theme\"],\"deny\":[]}},\"allow-title\":{\"identifier\":\"allow-title\",\"description\":\"Enables the title command without any pre-configured scope.\",\"commands\":{\"allow\":[\"title\"],\"deny\":[]}},\"allow-toggle-maximize\":{\"identifier\":\"allow-toggle-maximize\",\"description\":\"Enables the toggle_maximize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"toggle_maximize\"],\"deny\":[]}},\"allow-unmaximize\":{\"identifier\":\"allow-unmaximize\",\"description\":\"Enables the unmaximize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"unmaximize\"],\"deny\":[]}},\"allow-unminimize\":{\"identifier\":\"allow-unminimize\",\"description\":\"Enables the unminimize command without any pre-configured scope.\",\"commands\":{\"allow\":[\"unminimize\"],\"deny\":[]}},\"deny-available-monitors\":{\"identifier\":\"deny-available-monitors\",\"description\":\"Denies the available_monitors command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"available_monitors\"]}},\"deny-center\":{\"identifier\":\"deny-center\",\"description\":\"Denies the center command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"center\"]}},\"deny-close\":{\"identifier\":\"deny-close\",\"description\":\"Denies the close command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"close\"]}},\"deny-create\":{\"identifier\":\"deny-create\",\"description\":\"Denies the create command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"create\"]}},\"deny-current-monitor\":{\"identifier\":\"deny-current-monitor\",\"description\":\"Denies the current_monitor command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"current_monitor\"]}},\"deny-cursor-position\":{\"identifier\":\"deny-cursor-position\",\"description\":\"Denies the cursor_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"cursor_position\"]}},\"deny-destroy\":{\"identifier\":\"deny-destroy\",\"description\":\"Denies the destroy command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"destroy\"]}},\"deny-get-all-windows\":{\"identifier\":\"deny-get-all-windows\",\"description\":\"Denies the get_all_windows command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"get_all_windows\"]}},\"deny-hide\":{\"identifier\":\"deny-hide\",\"description\":\"Denies the hide command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"hide\"]}},\"deny-inner-position\":{\"identifier\":\"deny-inner-position\",\"description\":\"Denies the inner_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"inner_position\"]}},\"deny-inner-size\":{\"identifier\":\"deny-inner-size\",\"description\":\"Denies the inner_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"inner_size\"]}},\"deny-internal-toggle-maximize\":{\"identifier\":\"deny-internal-toggle-maximize\",\"description\":\"Denies the internal_toggle_maximize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"internal_toggle_maximize\"]}},\"deny-is-always-on-top\":{\"identifier\":\"deny-is-always-on-top\",\"description\":\"Denies the is_always_on_top command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_always_on_top\"]}},\"deny-is-closable\":{\"identifier\":\"deny-is-closable\",\"description\":\"Denies the is_closable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_closable\"]}},\"deny-is-decorated\":{\"identifier\":\"deny-is-decorated\",\"description\":\"Denies the is_decorated command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_decorated\"]}},\"deny-is-enabled\":{\"identifier\":\"deny-is-enabled\",\"description\":\"Denies the is_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_enabled\"]}},\"deny-is-focused\":{\"identifier\":\"deny-is-focused\",\"description\":\"Denies the is_focused command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_focused\"]}},\"deny-is-fullscreen\":{\"identifier\":\"deny-is-fullscreen\",\"description\":\"Denies the is_fullscreen command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_fullscreen\"]}},\"deny-is-maximizable\":{\"identifier\":\"deny-is-maximizable\",\"description\":\"Denies the is_maximizable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_maximizable\"]}},\"deny-is-maximized\":{\"identifier\":\"deny-is-maximized\",\"description\":\"Denies the is_maximized command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_maximized\"]}},\"deny-is-minimizable\":{\"identifier\":\"deny-is-minimizable\",\"description\":\"Denies the is_minimizable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_minimizable\"]}},\"deny-is-minimized\":{\"identifier\":\"deny-is-minimized\",\"description\":\"Denies the is_minimized command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_minimized\"]}},\"deny-is-resizable\":{\"identifier\":\"deny-is-resizable\",\"description\":\"Denies the is_resizable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_resizable\"]}},\"deny-is-visible\":{\"identifier\":\"deny-is-visible\",\"description\":\"Denies the is_visible command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"is_visible\"]}},\"deny-maximize\":{\"identifier\":\"deny-maximize\",\"description\":\"Denies the maximize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"maximize\"]}},\"deny-minimize\":{\"identifier\":\"deny-minimize\",\"description\":\"Denies the minimize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"minimize\"]}},\"deny-monitor-from-point\":{\"identifier\":\"deny-monitor-from-point\",\"description\":\"Denies the monitor_from_point command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"monitor_from_point\"]}},\"deny-outer-position\":{\"identifier\":\"deny-outer-position\",\"description\":\"Denies the outer_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"outer_position\"]}},\"deny-outer-size\":{\"identifier\":\"deny-outer-size\",\"description\":\"Denies the outer_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"outer_size\"]}},\"deny-primary-monitor\":{\"identifier\":\"deny-primary-monitor\",\"description\":\"Denies the primary_monitor command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"primary_monitor\"]}},\"deny-request-user-attention\":{\"identifier\":\"deny-request-user-attention\",\"description\":\"Denies the request_user_attention command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"request_user_attention\"]}},\"deny-scale-factor\":{\"identifier\":\"deny-scale-factor\",\"description\":\"Denies the scale_factor command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"scale_factor\"]}},\"deny-set-always-on-bottom\":{\"identifier\":\"deny-set-always-on-bottom\",\"description\":\"Denies the set_always_on_bottom command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_always_on_bottom\"]}},\"deny-set-always-on-top\":{\"identifier\":\"deny-set-always-on-top\",\"description\":\"Denies the set_always_on_top command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_always_on_top\"]}},\"deny-set-background-color\":{\"identifier\":\"deny-set-background-color\",\"description\":\"Denies the set_background_color command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_background_color\"]}},\"deny-set-badge-count\":{\"identifier\":\"deny-set-badge-count\",\"description\":\"Denies the set_badge_count command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_badge_count\"]}},\"deny-set-badge-label\":{\"identifier\":\"deny-set-badge-label\",\"description\":\"Denies the set_badge_label command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_badge_label\"]}},\"deny-set-closable\":{\"identifier\":\"deny-set-closable\",\"description\":\"Denies the set_closable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_closable\"]}},\"deny-set-content-protected\":{\"identifier\":\"deny-set-content-protected\",\"description\":\"Denies the set_content_protected command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_content_protected\"]}},\"deny-set-cursor-grab\":{\"identifier\":\"deny-set-cursor-grab\",\"description\":\"Denies the set_cursor_grab command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_cursor_grab\"]}},\"deny-set-cursor-icon\":{\"identifier\":\"deny-set-cursor-icon\",\"description\":\"Denies the set_cursor_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_cursor_icon\"]}},\"deny-set-cursor-position\":{\"identifier\":\"deny-set-cursor-position\",\"description\":\"Denies the set_cursor_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_cursor_position\"]}},\"deny-set-cursor-visible\":{\"identifier\":\"deny-set-cursor-visible\",\"description\":\"Denies the set_cursor_visible command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_cursor_visible\"]}},\"deny-set-decorations\":{\"identifier\":\"deny-set-decorations\",\"description\":\"Denies the set_decorations command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_decorations\"]}},\"deny-set-effects\":{\"identifier\":\"deny-set-effects\",\"description\":\"Denies the set_effects command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_effects\"]}},\"deny-set-enabled\":{\"identifier\":\"deny-set-enabled\",\"description\":\"Denies the set_enabled command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_enabled\"]}},\"deny-set-focus\":{\"identifier\":\"deny-set-focus\",\"description\":\"Denies the set_focus command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_focus\"]}},\"deny-set-focusable\":{\"identifier\":\"deny-set-focusable\",\"description\":\"Denies the set_focusable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_focusable\"]}},\"deny-set-fullscreen\":{\"identifier\":\"deny-set-fullscreen\",\"description\":\"Denies the set_fullscreen command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_fullscreen\"]}},\"deny-set-icon\":{\"identifier\":\"deny-set-icon\",\"description\":\"Denies the set_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_icon\"]}},\"deny-set-ignore-cursor-events\":{\"identifier\":\"deny-set-ignore-cursor-events\",\"description\":\"Denies the set_ignore_cursor_events command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_ignore_cursor_events\"]}},\"deny-set-max-size\":{\"identifier\":\"deny-set-max-size\",\"description\":\"Denies the set_max_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_max_size\"]}},\"deny-set-maximizable\":{\"identifier\":\"deny-set-maximizable\",\"description\":\"Denies the set_maximizable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_maximizable\"]}},\"deny-set-min-size\":{\"identifier\":\"deny-set-min-size\",\"description\":\"Denies the set_min_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_min_size\"]}},\"deny-set-minimizable\":{\"identifier\":\"deny-set-minimizable\",\"description\":\"Denies the set_minimizable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_minimizable\"]}},\"deny-set-overlay-icon\":{\"identifier\":\"deny-set-overlay-icon\",\"description\":\"Denies the set_overlay_icon command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_overlay_icon\"]}},\"deny-set-position\":{\"identifier\":\"deny-set-position\",\"description\":\"Denies the set_position command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_position\"]}},\"deny-set-progress-bar\":{\"identifier\":\"deny-set-progress-bar\",\"description\":\"Denies the set_progress_bar command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_progress_bar\"]}},\"deny-set-resizable\":{\"identifier\":\"deny-set-resizable\",\"description\":\"Denies the set_resizable command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_resizable\"]}},\"deny-set-shadow\":{\"identifier\":\"deny-set-shadow\",\"description\":\"Denies the set_shadow command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_shadow\"]}},\"deny-set-simple-fullscreen\":{\"identifier\":\"deny-set-simple-fullscreen\",\"description\":\"Denies the set_simple_fullscreen command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_simple_fullscreen\"]}},\"deny-set-size\":{\"identifier\":\"deny-set-size\",\"description\":\"Denies the set_size command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_size\"]}},\"deny-set-size-constraints\":{\"identifier\":\"deny-set-size-constraints\",\"description\":\"Denies the set_size_constraints command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_size_constraints\"]}},\"deny-set-skip-taskbar\":{\"identifier\":\"deny-set-skip-taskbar\",\"description\":\"Denies the set_skip_taskbar command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_skip_taskbar\"]}},\"deny-set-theme\":{\"identifier\":\"deny-set-theme\",\"description\":\"Denies the set_theme command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_theme\"]}},\"deny-set-title\":{\"identifier\":\"deny-set-title\",\"description\":\"Denies the set_title command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_title\"]}},\"deny-set-title-bar-style\":{\"identifier\":\"deny-set-title-bar-style\",\"description\":\"Denies the set_title_bar_style command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_title_bar_style\"]}},\"deny-set-visible-on-all-workspaces\":{\"identifier\":\"deny-set-visible-on-all-workspaces\",\"description\":\"Denies the set_visible_on_all_workspaces command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"set_visible_on_all_workspaces\"]}},\"deny-show\":{\"identifier\":\"deny-show\",\"description\":\"Denies the show command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"show\"]}},\"deny-start-dragging\":{\"identifier\":\"deny-start-dragging\",\"description\":\"Denies the start_dragging command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"start_dragging\"]}},\"deny-start-resize-dragging\":{\"identifier\":\"deny-start-resize-dragging\",\"description\":\"Denies the start_resize_dragging command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"start_resize_dragging\"]}},\"deny-theme\":{\"identifier\":\"deny-theme\",\"description\":\"Denies the theme command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"theme\"]}},\"deny-title\":{\"identifier\":\"deny-title\",\"description\":\"Denies the title command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"title\"]}},\"deny-toggle-maximize\":{\"identifier\":\"deny-toggle-maximize\",\"description\":\"Denies the toggle_maximize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"toggle_maximize\"]}},\"deny-unmaximize\":{\"identifier\":\"deny-unmaximize\",\"description\":\"Denies the unmaximize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"unmaximize\"]}},\"deny-unminimize\":{\"identifier\":\"deny-unminimize\",\"description\":\"Denies the unminimize command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"unminimize\"]}}},\"permission_sets\":{},\"global_scope_schema\":null},\"shell\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\",\"permissions\":[\"allow-open\"]},\"permissions\":{\"allow-execute\":{\"identifier\":\"allow-execute\",\"description\":\"Enables the execute command without any pre-configured scope.\",\"commands\":{\"allow\":[\"execute\"],\"deny\":[]}},\"allow-kill\":{\"identifier\":\"allow-kill\",\"description\":\"Enables the kill command without any pre-configured scope.\",\"commands\":{\"allow\":[\"kill\"],\"deny\":[]}},\"allow-open\":{\"identifier\":\"allow-open\",\"description\":\"Enables the open command without any pre-configured scope.\",\"commands\":{\"allow\":[\"open\"],\"deny\":[]}},\"allow-spawn\":{\"identifier\":\"allow-spawn\",\"description\":\"Enables the spawn command without any pre-configured scope.\",\"commands\":{\"allow\":[\"spawn\"],\"deny\":[]}},\"allow-stdin-write\":{\"identifier\":\"allow-stdin-write\",\"description\":\"Enables the stdin_write command without any pre-configured scope.\",\"commands\":{\"allow\":[\"stdin_write\"],\"deny\":[]}},\"deny-execute\":{\"identifier\":\"deny-execute\",\"description\":\"Denies the execute command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"execute\"]}},\"deny-kill\":{\"identifier\":\"deny-kill\",\"description\":\"Denies the kill command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"kill\"]}},\"deny-open\":{\"identifier\":\"deny-open\",\"description\":\"Denies the open command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"open\"]}},\"deny-spawn\":{\"identifier\":\"deny-spawn\",\"description\":\"Denies the spawn command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"spawn\"]}},\"deny-stdin-write\":{\"identifier\":\"deny-stdin-write\",\"description\":\"Denies the stdin_write command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"stdin_write\"]}}},\"permission_sets\":{},\"global_scope_schema\":{\"$schema\":\"http://json-schema.org/draft-07/schema#\",\"anyOf\":[{\"additionalProperties\":false,\"properties\":{\"args\":{\"allOf\":[{\"$ref\":\"#/definitions/ShellScopeEntryAllowedArgs\"}],\"description\":\"The allowed arguments for the command execution.\"},\"cmd\":{\"description\":\"The command name. It can start with a variable that resolves to a system base directory. The variables are: `$AUDIO`, `$CACHE`, `$CONFIG`, `$DATA`, `$LOCALDATA`, `$DESKTOP`, `$DOCUMENT`, `$DOWNLOAD`, `$EXE`, `$FONT`, `$HOME`, `$PICTURE`, `$PUBLIC`, `$RUNTIME`, `$TEMPLATE`, `$VIDEO`, `$RESOURCE`, `$LOG`, `$TEMP`, `$APPCONFIG`, `$APPDATA`, `$APPLOCALDATA`, `$APPCACHE`, `$APPLOG`.\",\"type\":\"string\"},\"name\":{\"description\":\"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\"type\":\"string\"}},\"required\":[\"cmd\",\"name\"],\"type\":\"object\"},{\"additionalProperties\":false,\"properties\":{\"args\":{\"allOf\":[{\"$ref\":\"#/definitions/ShellScopeEntryAllowedArgs\"}],\"description\":\"The allowed arguments for the command execution.\"},\"name\":{\"description\":\"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\"type\":\"string\"},\"sidecar\":{\"description\":\"If this command is a sidecar command.\",\"type\":\"boolean\"}},\"required\":[\"name\",\"sidecar\"],\"type\":\"object\"}],\"definitions\":{\"ShellScopeEntryAllowedArg\":{\"anyOf\":[{\"description\":\"A non-configurable argument that is passed to the command in the order it was specified.\",\"type\":\"string\"},{\"additionalProperties\":false,\"description\":\"A variable that is set while calling the command from the webview API.\",\"properties\":{\"raw\":{\"default\":false,\"description\":\"Marks the validator as a raw regex, meaning the plugin should not make any modification at runtime.\\n\\nThis means the regex will not match on the entire string by default, which might be exploited if your regex allow unexpected input to be considered valid. When using this option, make sure your regex is correct.\",\"type\":\"boolean\"},\"validator\":{\"description\":\"[regex] validator to require passed values to conform to an expected input.\\n\\nThis will require the argument value passed to this variable to match the `validator` regex before it will be executed.\\n\\nThe regex string is by default surrounded by `^...$` to match the full string. For example the `https?://\\\\w+` regex would be registered as `^https?://\\\\w+$`.\\n\\n[regex]: \",\"type\":\"string\"}},\"required\":[\"validator\"],\"type\":\"object\"}],\"description\":\"A command argument allowed to be executed by the webview API.\"},\"ShellScopeEntryAllowedArgs\":{\"anyOf\":[{\"description\":\"Use a simple boolean to allow all or disable all arguments to this command configuration.\",\"type\":\"boolean\"},{\"description\":\"A specific set of [`ShellScopeEntryAllowedArg`] that are valid to call for the command configuration.\",\"items\":{\"$ref\":\"#/definitions/ShellScopeEntryAllowedArg\"},\"type\":\"array\"}],\"description\":\"A set of command arguments allowed to be executed by the webview API.\\n\\nA value of `true` will allow any arguments to be passed to the command. `false` will disable all arguments. A list of [`ShellScopeEntryAllowedArg`] will set those arguments as the only valid arguments to be passed to the attached command configuration.\"}},\"description\":\"Shell scope entry.\",\"title\":\"ShellScopeEntry\"}},\"window-state\":{\"default_permission\":{\"identifier\":\"default\",\"description\":\"This permission set configures what kind of\\noperations are available from the window state plugin.\\n\\n#### Granted Permissions\\n\\nAll operations are enabled by default.\\n\\n\",\"permissions\":[\"allow-filename\",\"allow-restore-state\",\"allow-save-window-state\"]},\"permissions\":{\"allow-filename\":{\"identifier\":\"allow-filename\",\"description\":\"Enables the filename command without any pre-configured scope.\",\"commands\":{\"allow\":[\"filename\"],\"deny\":[]}},\"allow-restore-state\":{\"identifier\":\"allow-restore-state\",\"description\":\"Enables the restore_state command without any pre-configured scope.\",\"commands\":{\"allow\":[\"restore_state\"],\"deny\":[]}},\"allow-save-window-state\":{\"identifier\":\"allow-save-window-state\",\"description\":\"Enables the save_window_state command without any pre-configured scope.\",\"commands\":{\"allow\":[\"save_window_state\"],\"deny\":[]}},\"deny-filename\":{\"identifier\":\"deny-filename\",\"description\":\"Denies the filename command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"filename\"]}},\"deny-restore-state\":{\"identifier\":\"deny-restore-state\",\"description\":\"Denies the restore_state command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"restore_state\"]}},\"deny-save-window-state\":{\"identifier\":\"deny-save-window-state\",\"description\":\"Denies the save_window_state command without any pre-configured scope.\",\"commands\":{\"allow\":[],\"deny\":[\"save_window_state\"]}}},\"permission_sets\":{},\"global_scope_schema\":null}}" }, { "path": "desktop/src-tauri/gen/schemas/capabilities.json", "content": "{}" }, { "path": "desktop/src-tauri/gen/schemas/desktop-schema.json", "content": "{\n \"$schema\": \"http://json-schema.org/draft-07/schema#\",\n \"title\": \"CapabilityFile\",\n \"description\": \"Capability formats accepted in a capability file.\",\n \"anyOf\": [\n {\n \"description\": \"A single capability.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Capability\"\n }\n ]\n },\n {\n \"description\": \"A list of capabilities.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/Capability\"\n }\n },\n {\n \"description\": \"A list of capabilities.\",\n \"type\": \"object\",\n \"required\": [\n \"capabilities\"\n ],\n \"properties\": {\n \"capabilities\": {\n \"description\": \"The list of capabilities.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/Capability\"\n }\n }\n }\n }\n ],\n \"definitions\": {\n \"Capability\": {\n \"description\": \"A grouping and boundary mechanism developers can use to isolate access to the IPC layer.\\n\\nIt controls application windows' and webviews' fine grained access to the Tauri core, application, or plugin commands. If a webview or its window is not matching any capability then it has no access to the IPC layer at all.\\n\\nThis can be done to create groups of windows, based on their required system access, which can reduce impact of frontend vulnerabilities in less privileged windows. Windows can be added to a capability by exact name (e.g. `main-window`) or glob patterns like `*` or `admin-*`. A Window can have none, one, or multiple associated capabilities.\\n\\n## Example\\n\\n```json { \\\"identifier\\\": \\\"main-user-files-write\\\", \\\"description\\\": \\\"This capability allows the `main` window on macOS and Windows access to `filesystem` write related commands and `dialog` commands to enable programmatic access to files selected by the user.\\\", \\\"windows\\\": [ \\\"main\\\" ], \\\"permissions\\\": [ \\\"core:default\\\", \\\"dialog:open\\\", { \\\"identifier\\\": \\\"fs:allow-write-text-file\\\", \\\"allow\\\": [{ \\\"path\\\": \\\"$HOME/test.txt\\\" }] }, ], \\\"platforms\\\": [\\\"macOS\\\",\\\"windows\\\"] } ```\",\n \"type\": \"object\",\n \"required\": [\n \"identifier\",\n \"permissions\"\n ],\n \"properties\": {\n \"identifier\": {\n \"description\": \"Identifier of the capability.\\n\\n## Example\\n\\n`main-user-files-write`\",\n \"type\": \"string\"\n },\n \"description\": {\n \"description\": \"Description of what the capability is intended to allow on associated windows.\\n\\nIt should contain a description of what the grouped permissions should allow.\\n\\n## Example\\n\\nThis capability allows the `main` window access to `filesystem` write related commands and `dialog` commands to enable programmatic access to files selected by the user.\",\n \"default\": \"\",\n \"type\": \"string\"\n },\n \"remote\": {\n \"description\": \"Configure remote URLs that can use the capability permissions.\\n\\nThis setting is optional and defaults to not being set, as our default use case is that the content is served from our local application.\\n\\n:::caution Make sure you understand the security implications of providing remote sources with local system access. :::\\n\\n## Example\\n\\n```json { \\\"urls\\\": [\\\"https://*.mydomain.dev\\\"] } ```\",\n \"anyOf\": [\n {\n \"$ref\": \"#/definitions/CapabilityRemote\"\n },\n {\n \"type\": \"null\"\n }\n ]\n },\n \"local\": {\n \"description\": \"Whether this capability is enabled for local app URLs or not. Defaults to `true`.\",\n \"default\": true,\n \"type\": \"boolean\"\n },\n \"windows\": {\n \"description\": \"List of windows that are affected by this capability. Can be a glob pattern.\\n\\nIf a window label matches any of the patterns in this list, the capability will be enabled on all the webviews of that window, regardless of the value of [`Self::webviews`].\\n\\nOn multiwebview windows, prefer specifying [`Self::webviews`] and omitting [`Self::windows`] for a fine grained access control.\\n\\n## Example\\n\\n`[\\\"main\\\"]`\",\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\"\n }\n },\n \"webviews\": {\n \"description\": \"List of webviews that are affected by this capability. Can be a glob pattern.\\n\\nThe capability will be enabled on all the webviews whose label matches any of the patterns in this list, regardless of whether the webview's window label matches a pattern in [`Self::windows`].\\n\\n## Example\\n\\n`[\\\"sub-webview-one\\\", \\\"sub-webview-two\\\"]`\",\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\"\n }\n },\n \"permissions\": {\n \"description\": \"List of permissions attached to this capability.\\n\\nMust include the plugin name as prefix in the form of `${plugin-name}:${permission-name}`. For commands directly implemented in the application itself only `${permission-name}` is required.\\n\\n## Example\\n\\n```json [ \\\"core:default\\\", \\\"shell:allow-open\\\", \\\"dialog:open\\\", { \\\"identifier\\\": \\\"fs:allow-write-text-file\\\", \\\"allow\\\": [{ \\\"path\\\": \\\"$HOME/test.txt\\\" }] } ] ```\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/PermissionEntry\"\n },\n \"uniqueItems\": true\n },\n \"platforms\": {\n \"description\": \"Limit which target platforms this capability applies to.\\n\\nBy default all platforms are targeted.\\n\\n## Example\\n\\n`[\\\"macOS\\\",\\\"windows\\\"]`\",\n \"type\": [\n \"array\",\n \"null\"\n ],\n \"items\": {\n \"$ref\": \"#/definitions/Target\"\n }\n }\n }\n },\n \"CapabilityRemote\": {\n \"description\": \"Configuration for remote URLs that are associated with the capability.\",\n \"type\": \"object\",\n \"required\": [\n \"urls\"\n ],\n \"properties\": {\n \"urls\": {\n \"description\": \"Remote domains this capability refers to using the [URLPattern standard](https://urlpattern.spec.whatwg.org/).\\n\\n## Examples\\n\\n- \\\"https://*.mydomain.dev\\\": allows subdomains of mydomain.dev - \\\"https://mydomain.dev/api/*\\\": allows any subpath of mydomain.dev/api\",\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\"\n }\n }\n }\n },\n \"PermissionEntry\": {\n \"description\": \"An entry for a permission value in a [`Capability`] can be either a raw permission [`Identifier`] or an object that references a permission and extends its scope.\",\n \"anyOf\": [\n {\n \"description\": \"Reference a permission or permission set by identifier.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Identifier\"\n }\n ]\n },\n {\n \"description\": \"Reference a permission or permission set by identifier and extends its scope.\",\n \"type\": \"object\",\n \"allOf\": [\n {\n \"if\": {\n \"properties\": {\n \"identifier\": {\n \"anyOf\": [\n {\n \"description\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\",\n \"type\": \"string\",\n \"const\": \"shell:default\",\n \"markdownDescription\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\"\n },\n {\n \"description\": \"Enables the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-execute\",\n \"markdownDescription\": \"Enables the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-kill\",\n \"markdownDescription\": \"Enables the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-open\",\n \"markdownDescription\": \"Enables the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-spawn\",\n \"markdownDescription\": \"Enables the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-stdin-write\",\n \"markdownDescription\": \"Enables the stdin_write command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-execute\",\n \"markdownDescription\": \"Denies the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-kill\",\n \"markdownDescription\": \"Denies the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-open\",\n \"markdownDescription\": \"Denies the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-spawn\",\n \"markdownDescription\": \"Denies the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-stdin-write\",\n \"markdownDescription\": \"Denies the stdin_write command without any pre-configured scope.\"\n }\n ]\n }\n }\n },\n \"then\": {\n \"properties\": {\n \"allow\": {\n \"items\": {\n \"title\": \"ShellScopeEntry\",\n \"description\": \"Shell scope entry.\",\n \"anyOf\": [\n {\n \"type\": \"object\",\n \"required\": [\n \"cmd\",\n \"name\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"cmd\": {\n \"description\": \"The command name. It can start with a variable that resolves to a system base directory. The variables are: `$AUDIO`, `$CACHE`, `$CONFIG`, `$DATA`, `$LOCALDATA`, `$DESKTOP`, `$DOCUMENT`, `$DOWNLOAD`, `$EXE`, `$FONT`, `$HOME`, `$PICTURE`, `$PUBLIC`, `$RUNTIME`, `$TEMPLATE`, `$VIDEO`, `$RESOURCE`, `$LOG`, `$TEMP`, `$APPCONFIG`, `$APPDATA`, `$APPLOCALDATA`, `$APPCACHE`, `$APPLOG`.\",\n \"type\": \"string\"\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n }\n },\n \"additionalProperties\": false\n },\n {\n \"type\": \"object\",\n \"required\": [\n \"name\",\n \"sidecar\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n },\n \"sidecar\": {\n \"description\": \"If this command is a sidecar command.\",\n \"type\": \"boolean\"\n }\n },\n \"additionalProperties\": false\n }\n ]\n }\n },\n \"deny\": {\n \"items\": {\n \"title\": \"ShellScopeEntry\",\n \"description\": \"Shell scope entry.\",\n \"anyOf\": [\n {\n \"type\": \"object\",\n \"required\": [\n \"cmd\",\n \"name\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"cmd\": {\n \"description\": \"The command name. It can start with a variable that resolves to a system base directory. The variables are: `$AUDIO`, `$CACHE`, `$CONFIG`, `$DATA`, `$LOCALDATA`, `$DESKTOP`, `$DOCUMENT`, `$DOWNLOAD`, `$EXE`, `$FONT`, `$HOME`, `$PICTURE`, `$PUBLIC`, `$RUNTIME`, `$TEMPLATE`, `$VIDEO`, `$RESOURCE`, `$LOG`, `$TEMP`, `$APPCONFIG`, `$APPDATA`, `$APPLOCALDATA`, `$APPCACHE`, `$APPLOG`.\",\n \"type\": \"string\"\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n }\n },\n \"additionalProperties\": false\n },\n {\n \"type\": \"object\",\n \"required\": [\n \"name\",\n \"sidecar\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n },\n \"sidecar\": {\n \"description\": \"If this command is a sidecar command.\",\n \"type\": \"boolean\"\n }\n },\n \"additionalProperties\": false\n }\n ]\n }\n }\n }\n },\n \"properties\": {\n \"identifier\": {\n \"description\": \"Identifier of the permission or permission set.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Identifier\"\n }\n ]\n }\n }\n },\n {\n \"properties\": {\n \"identifier\": {\n \"description\": \"Identifier of the permission or permission set.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Identifier\"\n }\n ]\n },\n \"allow\": {\n \"description\": \"Data that defines what is allowed by the scope.\",\n \"type\": [\n \"array\",\n \"null\"\n ],\n \"items\": {\n \"$ref\": \"#/definitions/Value\"\n }\n },\n \"deny\": {\n \"description\": \"Data that defines what is denied by the scope. This should be prioritized by validation logic.\",\n \"type\": [\n \"array\",\n \"null\"\n ],\n \"items\": {\n \"$ref\": \"#/definitions/Value\"\n }\n }\n }\n }\n ],\n \"required\": [\n \"identifier\"\n ]\n }\n ]\n },\n \"Identifier\": {\n \"description\": \"Permission identifier\",\n \"oneOf\": [\n {\n \"description\": \"Default core plugins set.\\n#### This default permission set includes:\\n\\n- `core:path:default`\\n- `core:event:default`\\n- `core:window:default`\\n- `core:webview:default`\\n- `core:app:default`\\n- `core:image:default`\\n- `core:resources:default`\\n- `core:menu:default`\\n- `core:tray:default`\",\n \"type\": \"string\",\n \"const\": \"core:default\",\n \"markdownDescription\": \"Default core plugins set.\\n#### This default permission set includes:\\n\\n- `core:path:default`\\n- `core:event:default`\\n- `core:window:default`\\n- `core:webview:default`\\n- `core:app:default`\\n- `core:image:default`\\n- `core:resources:default`\\n- `core:menu:default`\\n- `core:tray:default`\"\n },\n {\n \"description\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-version`\\n- `allow-name`\\n- `allow-tauri-version`\\n- `allow-identifier`\\n- `allow-bundle-type`\\n- `allow-register-listener`\\n- `allow-remove-listener`\",\n \"type\": \"string\",\n \"const\": \"core:app:default\",\n \"markdownDescription\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-version`\\n- `allow-name`\\n- `allow-tauri-version`\\n- `allow-identifier`\\n- `allow-bundle-type`\\n- `allow-register-listener`\\n- `allow-remove-listener`\"\n },\n {\n \"description\": \"Enables the app_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-app-hide\",\n \"markdownDescription\": \"Enables the app_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the app_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-app-show\",\n \"markdownDescription\": \"Enables the app_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the bundle_type command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-bundle-type\",\n \"markdownDescription\": \"Enables the bundle_type command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the default_window_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-default-window-icon\",\n \"markdownDescription\": \"Enables the default_window_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the fetch_data_store_identifiers command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-fetch-data-store-identifiers\",\n \"markdownDescription\": \"Enables the fetch_data_store_identifiers command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the identifier command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-identifier\",\n \"markdownDescription\": \"Enables the identifier command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the name command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-name\",\n \"markdownDescription\": \"Enables the name command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the register_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-register-listener\",\n \"markdownDescription\": \"Enables the register_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_data_store command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-remove-data-store\",\n \"markdownDescription\": \"Enables the remove_data_store command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-remove-listener\",\n \"markdownDescription\": \"Enables the remove_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_app_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-set-app-theme\",\n \"markdownDescription\": \"Enables the set_app_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_dock_visibility command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-set-dock-visibility\",\n \"markdownDescription\": \"Enables the set_dock_visibility command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the tauri_version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-tauri-version\",\n \"markdownDescription\": \"Enables the tauri_version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-version\",\n \"markdownDescription\": \"Enables the version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the app_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-app-hide\",\n \"markdownDescription\": \"Denies the app_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the app_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-app-show\",\n \"markdownDescription\": \"Denies the app_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the bundle_type command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-bundle-type\",\n \"markdownDescription\": \"Denies the bundle_type command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the default_window_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-default-window-icon\",\n \"markdownDescription\": \"Denies the default_window_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the fetch_data_store_identifiers command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-fetch-data-store-identifiers\",\n \"markdownDescription\": \"Denies the fetch_data_store_identifiers command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the identifier command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-identifier\",\n \"markdownDescription\": \"Denies the identifier command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the name command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-name\",\n \"markdownDescription\": \"Denies the name command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the register_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-register-listener\",\n \"markdownDescription\": \"Denies the register_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_data_store command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-remove-data-store\",\n \"markdownDescription\": \"Denies the remove_data_store command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-remove-listener\",\n \"markdownDescription\": \"Denies the remove_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_app_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-set-app-theme\",\n \"markdownDescription\": \"Denies the set_app_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_dock_visibility command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-set-dock-visibility\",\n \"markdownDescription\": \"Denies the set_dock_visibility command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the tauri_version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-tauri-version\",\n \"markdownDescription\": \"Denies the tauri_version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-version\",\n \"markdownDescription\": \"Denies the version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-listen`\\n- `allow-unlisten`\\n- `allow-emit`\\n- `allow-emit-to`\",\n \"type\": \"string\",\n \"const\": \"core:event:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-listen`\\n- `allow-unlisten`\\n- `allow-emit`\\n- `allow-emit-to`\"\n },\n {\n \"description\": \"Enables the emit command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-emit\",\n \"markdownDescription\": \"Enables the emit command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the emit_to command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-emit-to\",\n \"markdownDescription\": \"Enables the emit_to command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the listen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-listen\",\n \"markdownDescription\": \"Enables the listen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the unlisten command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-unlisten\",\n \"markdownDescription\": \"Enables the unlisten command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the emit command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-emit\",\n \"markdownDescription\": \"Denies the emit command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the emit_to command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-emit-to\",\n \"markdownDescription\": \"Denies the emit_to command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the listen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-listen\",\n \"markdownDescription\": \"Denies the listen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the unlisten command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-unlisten\",\n \"markdownDescription\": \"Denies the unlisten command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-from-bytes`\\n- `allow-from-path`\\n- `allow-rgba`\\n- `allow-size`\",\n \"type\": \"string\",\n \"const\": \"core:image:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-from-bytes`\\n- `allow-from-path`\\n- `allow-rgba`\\n- `allow-size`\"\n },\n {\n \"description\": \"Enables the from_bytes command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-from-bytes\",\n \"markdownDescription\": \"Enables the from_bytes command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the from_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-from-path\",\n \"markdownDescription\": \"Enables the from_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-new\",\n \"markdownDescription\": \"Enables the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the rgba command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-rgba\",\n \"markdownDescription\": \"Enables the rgba command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-size\",\n \"markdownDescription\": \"Enables the size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the from_bytes command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-from-bytes\",\n \"markdownDescription\": \"Denies the from_bytes command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the from_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-from-path\",\n \"markdownDescription\": \"Denies the from_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-new\",\n \"markdownDescription\": \"Denies the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the rgba command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-rgba\",\n \"markdownDescription\": \"Denies the rgba command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-size\",\n \"markdownDescription\": \"Denies the size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-append`\\n- `allow-prepend`\\n- `allow-insert`\\n- `allow-remove`\\n- `allow-remove-at`\\n- `allow-items`\\n- `allow-get`\\n- `allow-popup`\\n- `allow-create-default`\\n- `allow-set-as-app-menu`\\n- `allow-set-as-window-menu`\\n- `allow-text`\\n- `allow-set-text`\\n- `allow-is-enabled`\\n- `allow-set-enabled`\\n- `allow-set-accelerator`\\n- `allow-set-as-windows-menu-for-nsapp`\\n- `allow-set-as-help-menu-for-nsapp`\\n- `allow-is-checked`\\n- `allow-set-checked`\\n- `allow-set-icon`\",\n \"type\": \"string\",\n \"const\": \"core:menu:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-append`\\n- `allow-prepend`\\n- `allow-insert`\\n- `allow-remove`\\n- `allow-remove-at`\\n- `allow-items`\\n- `allow-get`\\n- `allow-popup`\\n- `allow-create-default`\\n- `allow-set-as-app-menu`\\n- `allow-set-as-window-menu`\\n- `allow-text`\\n- `allow-set-text`\\n- `allow-is-enabled`\\n- `allow-set-enabled`\\n- `allow-set-accelerator`\\n- `allow-set-as-windows-menu-for-nsapp`\\n- `allow-set-as-help-menu-for-nsapp`\\n- `allow-is-checked`\\n- `allow-set-checked`\\n- `allow-set-icon`\"\n },\n {\n \"description\": \"Enables the append command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-append\",\n \"markdownDescription\": \"Enables the append command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create_default command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-create-default\",\n \"markdownDescription\": \"Enables the create_default command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the get command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-get\",\n \"markdownDescription\": \"Enables the get command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the insert command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-insert\",\n \"markdownDescription\": \"Enables the insert command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-is-checked\",\n \"markdownDescription\": \"Enables the is_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-is-enabled\",\n \"markdownDescription\": \"Enables the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the items command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-items\",\n \"markdownDescription\": \"Enables the items command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-new\",\n \"markdownDescription\": \"Enables the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the popup command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-popup\",\n \"markdownDescription\": \"Enables the popup command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the prepend command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-prepend\",\n \"markdownDescription\": \"Enables the prepend command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-remove\",\n \"markdownDescription\": \"Enables the remove command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_at command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-remove-at\",\n \"markdownDescription\": \"Enables the remove_at command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_accelerator command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-accelerator\",\n \"markdownDescription\": \"Enables the set_accelerator command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_app_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-app-menu\",\n \"markdownDescription\": \"Enables the set_as_app_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_help_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-help-menu-for-nsapp\",\n \"markdownDescription\": \"Enables the set_as_help_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_window_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-window-menu\",\n \"markdownDescription\": \"Enables the set_as_window_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_windows_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-windows-menu-for-nsapp\",\n \"markdownDescription\": \"Enables the set_as_windows_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-checked\",\n \"markdownDescription\": \"Enables the set_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-enabled\",\n \"markdownDescription\": \"Enables the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-icon\",\n \"markdownDescription\": \"Enables the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-text\",\n \"markdownDescription\": \"Enables the set_text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-text\",\n \"markdownDescription\": \"Enables the text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the append command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-append\",\n \"markdownDescription\": \"Denies the append command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create_default command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-create-default\",\n \"markdownDescription\": \"Denies the create_default command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-get\",\n \"markdownDescription\": \"Denies the get command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the insert command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-insert\",\n \"markdownDescription\": \"Denies the insert command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-is-checked\",\n \"markdownDescription\": \"Denies the is_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-is-enabled\",\n \"markdownDescription\": \"Denies the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the items command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-items\",\n \"markdownDescription\": \"Denies the items command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-new\",\n \"markdownDescription\": \"Denies the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the popup command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-popup\",\n \"markdownDescription\": \"Denies the popup command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the prepend command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-prepend\",\n \"markdownDescription\": \"Denies the prepend command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-remove\",\n \"markdownDescription\": \"Denies the remove command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_at command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-remove-at\",\n \"markdownDescription\": \"Denies the remove_at command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_accelerator command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-accelerator\",\n \"markdownDescription\": \"Denies the set_accelerator command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_app_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-app-menu\",\n \"markdownDescription\": \"Denies the set_as_app_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_help_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-help-menu-for-nsapp\",\n \"markdownDescription\": \"Denies the set_as_help_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_window_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-window-menu\",\n \"markdownDescription\": \"Denies the set_as_window_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_windows_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-windows-menu-for-nsapp\",\n \"markdownDescription\": \"Denies the set_as_windows_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-checked\",\n \"markdownDescription\": \"Denies the set_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-enabled\",\n \"markdownDescription\": \"Denies the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-icon\",\n \"markdownDescription\": \"Denies the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-text\",\n \"markdownDescription\": \"Denies the set_text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-text\",\n \"markdownDescription\": \"Denies the text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-resolve-directory`\\n- `allow-resolve`\\n- `allow-normalize`\\n- `allow-join`\\n- `allow-dirname`\\n- `allow-extname`\\n- `allow-basename`\\n- `allow-is-absolute`\",\n \"type\": \"string\",\n \"const\": \"core:path:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-resolve-directory`\\n- `allow-resolve`\\n- `allow-normalize`\\n- `allow-join`\\n- `allow-dirname`\\n- `allow-extname`\\n- `allow-basename`\\n- `allow-is-absolute`\"\n },\n {\n \"description\": \"Enables the basename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-basename\",\n \"markdownDescription\": \"Enables the basename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the dirname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-dirname\",\n \"markdownDescription\": \"Enables the dirname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the extname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-extname\",\n \"markdownDescription\": \"Enables the extname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_absolute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-is-absolute\",\n \"markdownDescription\": \"Enables the is_absolute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the join command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-join\",\n \"markdownDescription\": \"Enables the join command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the normalize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-normalize\",\n \"markdownDescription\": \"Enables the normalize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the resolve command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-resolve\",\n \"markdownDescription\": \"Enables the resolve command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the resolve_directory command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-resolve-directory\",\n \"markdownDescription\": \"Enables the resolve_directory command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the basename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-basename\",\n \"markdownDescription\": \"Denies the basename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the dirname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-dirname\",\n \"markdownDescription\": \"Denies the dirname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the extname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-extname\",\n \"markdownDescription\": \"Denies the extname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_absolute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-is-absolute\",\n \"markdownDescription\": \"Denies the is_absolute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the join command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-join\",\n \"markdownDescription\": \"Denies the join command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the normalize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-normalize\",\n \"markdownDescription\": \"Denies the normalize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the resolve command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-resolve\",\n \"markdownDescription\": \"Denies the resolve command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the resolve_directory command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-resolve-directory\",\n \"markdownDescription\": \"Denies the resolve_directory command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-close`\",\n \"type\": \"string\",\n \"const\": \"core:resources:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-close`\"\n },\n {\n \"description\": \"Enables the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:resources:allow-close\",\n \"markdownDescription\": \"Enables the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:resources:deny-close\",\n \"markdownDescription\": \"Denies the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-get-by-id`\\n- `allow-remove-by-id`\\n- `allow-set-icon`\\n- `allow-set-menu`\\n- `allow-set-tooltip`\\n- `allow-set-title`\\n- `allow-set-visible`\\n- `allow-set-temp-dir-path`\\n- `allow-set-icon-as-template`\\n- `allow-set-show-menu-on-left-click`\",\n \"type\": \"string\",\n \"const\": \"core:tray:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-get-by-id`\\n- `allow-remove-by-id`\\n- `allow-set-icon`\\n- `allow-set-menu`\\n- `allow-set-tooltip`\\n- `allow-set-title`\\n- `allow-set-visible`\\n- `allow-set-temp-dir-path`\\n- `allow-set-icon-as-template`\\n- `allow-set-show-menu-on-left-click`\"\n },\n {\n \"description\": \"Enables the get_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-get-by-id\",\n \"markdownDescription\": \"Enables the get_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-new\",\n \"markdownDescription\": \"Enables the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-remove-by-id\",\n \"markdownDescription\": \"Enables the remove_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-icon\",\n \"markdownDescription\": \"Enables the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon_as_template command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-icon-as-template\",\n \"markdownDescription\": \"Enables the set_icon_as_template command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-menu\",\n \"markdownDescription\": \"Enables the set_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_show_menu_on_left_click command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-show-menu-on-left-click\",\n \"markdownDescription\": \"Enables the set_show_menu_on_left_click command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_temp_dir_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-temp-dir-path\",\n \"markdownDescription\": \"Enables the set_temp_dir_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-title\",\n \"markdownDescription\": \"Enables the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_tooltip command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-tooltip\",\n \"markdownDescription\": \"Enables the set_tooltip command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-visible\",\n \"markdownDescription\": \"Enables the set_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-get-by-id\",\n \"markdownDescription\": \"Denies the get_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-new\",\n \"markdownDescription\": \"Denies the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-remove-by-id\",\n \"markdownDescription\": \"Denies the remove_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-icon\",\n \"markdownDescription\": \"Denies the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon_as_template command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-icon-as-template\",\n \"markdownDescription\": \"Denies the set_icon_as_template command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-menu\",\n \"markdownDescription\": \"Denies the set_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_show_menu_on_left_click command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-show-menu-on-left-click\",\n \"markdownDescription\": \"Denies the set_show_menu_on_left_click command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_temp_dir_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-temp-dir-path\",\n \"markdownDescription\": \"Denies the set_temp_dir_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-title\",\n \"markdownDescription\": \"Denies the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_tooltip command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-tooltip\",\n \"markdownDescription\": \"Denies the set_tooltip command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-visible\",\n \"markdownDescription\": \"Denies the set_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-webviews`\\n- `allow-webview-position`\\n- `allow-webview-size`\\n- `allow-internal-toggle-devtools`\",\n \"type\": \"string\",\n \"const\": \"core:webview:default\",\n \"markdownDescription\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-webviews`\\n- `allow-webview-position`\\n- `allow-webview-size`\\n- `allow-internal-toggle-devtools`\"\n },\n {\n \"description\": \"Enables the clear_all_browsing_data command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-clear-all-browsing-data\",\n \"markdownDescription\": \"Enables the clear_all_browsing_data command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create_webview command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-create-webview\",\n \"markdownDescription\": \"Enables the create_webview command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create_webview_window command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-create-webview-window\",\n \"markdownDescription\": \"Enables the create_webview_window command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the get_all_webviews command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-get-all-webviews\",\n \"markdownDescription\": \"Enables the get_all_webviews command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the internal_toggle_devtools command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-internal-toggle-devtools\",\n \"markdownDescription\": \"Enables the internal_toggle_devtools command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the print command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-print\",\n \"markdownDescription\": \"Enables the print command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the reparent command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-reparent\",\n \"markdownDescription\": \"Enables the reparent command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_auto_resize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-auto-resize\",\n \"markdownDescription\": \"Enables the set_webview_auto_resize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-background-color\",\n \"markdownDescription\": \"Enables the set_webview_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-focus\",\n \"markdownDescription\": \"Enables the set_webview_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-position\",\n \"markdownDescription\": \"Enables the set_webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-size\",\n \"markdownDescription\": \"Enables the set_webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_zoom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-zoom\",\n \"markdownDescription\": \"Enables the set_webview_zoom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-close\",\n \"markdownDescription\": \"Enables the webview_close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-hide\",\n \"markdownDescription\": \"Enables the webview_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-position\",\n \"markdownDescription\": \"Enables the webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-show\",\n \"markdownDescription\": \"Enables the webview_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-size\",\n \"markdownDescription\": \"Enables the webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the clear_all_browsing_data command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-clear-all-browsing-data\",\n \"markdownDescription\": \"Denies the clear_all_browsing_data command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create_webview command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-create-webview\",\n \"markdownDescription\": \"Denies the create_webview command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create_webview_window command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-create-webview-window\",\n \"markdownDescription\": \"Denies the create_webview_window command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get_all_webviews command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-get-all-webviews\",\n \"markdownDescription\": \"Denies the get_all_webviews command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the internal_toggle_devtools command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-internal-toggle-devtools\",\n \"markdownDescription\": \"Denies the internal_toggle_devtools command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the print command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-print\",\n \"markdownDescription\": \"Denies the print command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the reparent command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-reparent\",\n \"markdownDescription\": \"Denies the reparent command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_auto_resize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-auto-resize\",\n \"markdownDescription\": \"Denies the set_webview_auto_resize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-background-color\",\n \"markdownDescription\": \"Denies the set_webview_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-focus\",\n \"markdownDescription\": \"Denies the set_webview_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-position\",\n \"markdownDescription\": \"Denies the set_webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-size\",\n \"markdownDescription\": \"Denies the set_webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_zoom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-zoom\",\n \"markdownDescription\": \"Denies the set_webview_zoom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-close\",\n \"markdownDescription\": \"Denies the webview_close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-hide\",\n \"markdownDescription\": \"Denies the webview_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-position\",\n \"markdownDescription\": \"Denies the webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-show\",\n \"markdownDescription\": \"Denies the webview_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-size\",\n \"markdownDescription\": \"Denies the webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-windows`\\n- `allow-scale-factor`\\n- `allow-inner-position`\\n- `allow-outer-position`\\n- `allow-inner-size`\\n- `allow-outer-size`\\n- `allow-is-fullscreen`\\n- `allow-is-minimized`\\n- `allow-is-maximized`\\n- `allow-is-focused`\\n- `allow-is-decorated`\\n- `allow-is-resizable`\\n- `allow-is-maximizable`\\n- `allow-is-minimizable`\\n- `allow-is-closable`\\n- `allow-is-visible`\\n- `allow-is-enabled`\\n- `allow-title`\\n- `allow-current-monitor`\\n- `allow-primary-monitor`\\n- `allow-monitor-from-point`\\n- `allow-available-monitors`\\n- `allow-cursor-position`\\n- `allow-theme`\\n- `allow-is-always-on-top`\\n- `allow-internal-toggle-maximize`\",\n \"type\": \"string\",\n \"const\": \"core:window:default\",\n \"markdownDescription\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-windows`\\n- `allow-scale-factor`\\n- `allow-inner-position`\\n- `allow-outer-position`\\n- `allow-inner-size`\\n- `allow-outer-size`\\n- `allow-is-fullscreen`\\n- `allow-is-minimized`\\n- `allow-is-maximized`\\n- `allow-is-focused`\\n- `allow-is-decorated`\\n- `allow-is-resizable`\\n- `allow-is-maximizable`\\n- `allow-is-minimizable`\\n- `allow-is-closable`\\n- `allow-is-visible`\\n- `allow-is-enabled`\\n- `allow-title`\\n- `allow-current-monitor`\\n- `allow-primary-monitor`\\n- `allow-monitor-from-point`\\n- `allow-available-monitors`\\n- `allow-cursor-position`\\n- `allow-theme`\\n- `allow-is-always-on-top`\\n- `allow-internal-toggle-maximize`\"\n },\n {\n \"description\": \"Enables the available_monitors command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-available-monitors\",\n \"markdownDescription\": \"Enables the available_monitors command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the center command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-center\",\n \"markdownDescription\": \"Enables the center command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-close\",\n \"markdownDescription\": \"Enables the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-create\",\n \"markdownDescription\": \"Enables the create command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the current_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-current-monitor\",\n \"markdownDescription\": \"Enables the current_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-cursor-position\",\n \"markdownDescription\": \"Enables the cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the destroy command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-destroy\",\n \"markdownDescription\": \"Enables the destroy command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the get_all_windows command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-get-all-windows\",\n \"markdownDescription\": \"Enables the get_all_windows command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-hide\",\n \"markdownDescription\": \"Enables the hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the inner_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-inner-position\",\n \"markdownDescription\": \"Enables the inner_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the inner_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-inner-size\",\n \"markdownDescription\": \"Enables the inner_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the internal_toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-internal-toggle-maximize\",\n \"markdownDescription\": \"Enables the internal_toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-always-on-top\",\n \"markdownDescription\": \"Enables the is_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-closable\",\n \"markdownDescription\": \"Enables the is_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_decorated command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-decorated\",\n \"markdownDescription\": \"Enables the is_decorated command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-enabled\",\n \"markdownDescription\": \"Enables the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_focused command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-focused\",\n \"markdownDescription\": \"Enables the is_focused command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-fullscreen\",\n \"markdownDescription\": \"Enables the is_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-maximizable\",\n \"markdownDescription\": \"Enables the is_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_maximized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-maximized\",\n \"markdownDescription\": \"Enables the is_maximized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-minimizable\",\n \"markdownDescription\": \"Enables the is_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_minimized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-minimized\",\n \"markdownDescription\": \"Enables the is_minimized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-resizable\",\n \"markdownDescription\": \"Enables the is_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-visible\",\n \"markdownDescription\": \"Enables the is_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-maximize\",\n \"markdownDescription\": \"Enables the maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the minimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-minimize\",\n \"markdownDescription\": \"Enables the minimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the monitor_from_point command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-monitor-from-point\",\n \"markdownDescription\": \"Enables the monitor_from_point command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the outer_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-outer-position\",\n \"markdownDescription\": \"Enables the outer_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the outer_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-outer-size\",\n \"markdownDescription\": \"Enables the outer_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the primary_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-primary-monitor\",\n \"markdownDescription\": \"Enables the primary_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the request_user_attention command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-request-user-attention\",\n \"markdownDescription\": \"Enables the request_user_attention command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the scale_factor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-scale-factor\",\n \"markdownDescription\": \"Enables the scale_factor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_always_on_bottom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-always-on-bottom\",\n \"markdownDescription\": \"Enables the set_always_on_bottom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-always-on-top\",\n \"markdownDescription\": \"Enables the set_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-background-color\",\n \"markdownDescription\": \"Enables the set_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_badge_count command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-badge-count\",\n \"markdownDescription\": \"Enables the set_badge_count command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_badge_label command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-badge-label\",\n \"markdownDescription\": \"Enables the set_badge_label command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-closable\",\n \"markdownDescription\": \"Enables the set_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_content_protected command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-content-protected\",\n \"markdownDescription\": \"Enables the set_content_protected command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_grab command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-grab\",\n \"markdownDescription\": \"Enables the set_cursor_grab command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-icon\",\n \"markdownDescription\": \"Enables the set_cursor_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-position\",\n \"markdownDescription\": \"Enables the set_cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-visible\",\n \"markdownDescription\": \"Enables the set_cursor_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_decorations command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-decorations\",\n \"markdownDescription\": \"Enables the set_decorations command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_effects command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-effects\",\n \"markdownDescription\": \"Enables the set_effects command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-enabled\",\n \"markdownDescription\": \"Enables the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-focus\",\n \"markdownDescription\": \"Enables the set_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_focusable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-focusable\",\n \"markdownDescription\": \"Enables the set_focusable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-fullscreen\",\n \"markdownDescription\": \"Enables the set_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-icon\",\n \"markdownDescription\": \"Enables the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_ignore_cursor_events command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-ignore-cursor-events\",\n \"markdownDescription\": \"Enables the set_ignore_cursor_events command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_max_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-max-size\",\n \"markdownDescription\": \"Enables the set_max_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-maximizable\",\n \"markdownDescription\": \"Enables the set_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_min_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-min-size\",\n \"markdownDescription\": \"Enables the set_min_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-minimizable\",\n \"markdownDescription\": \"Enables the set_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_overlay_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-overlay-icon\",\n \"markdownDescription\": \"Enables the set_overlay_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-position\",\n \"markdownDescription\": \"Enables the set_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_progress_bar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-progress-bar\",\n \"markdownDescription\": \"Enables the set_progress_bar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-resizable\",\n \"markdownDescription\": \"Enables the set_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_shadow command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-shadow\",\n \"markdownDescription\": \"Enables the set_shadow command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_simple_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-simple-fullscreen\",\n \"markdownDescription\": \"Enables the set_simple_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-size\",\n \"markdownDescription\": \"Enables the set_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_size_constraints command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-size-constraints\",\n \"markdownDescription\": \"Enables the set_size_constraints command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_skip_taskbar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-skip-taskbar\",\n \"markdownDescription\": \"Enables the set_skip_taskbar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-theme\",\n \"markdownDescription\": \"Enables the set_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-title\",\n \"markdownDescription\": \"Enables the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_title_bar_style command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-title-bar-style\",\n \"markdownDescription\": \"Enables the set_title_bar_style command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_visible_on_all_workspaces command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-visible-on-all-workspaces\",\n \"markdownDescription\": \"Enables the set_visible_on_all_workspaces command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-show\",\n \"markdownDescription\": \"Enables the show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the start_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-start-dragging\",\n \"markdownDescription\": \"Enables the start_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the start_resize_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-start-resize-dragging\",\n \"markdownDescription\": \"Enables the start_resize_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-theme\",\n \"markdownDescription\": \"Enables the theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-title\",\n \"markdownDescription\": \"Enables the title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-toggle-maximize\",\n \"markdownDescription\": \"Enables the toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the unmaximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-unmaximize\",\n \"markdownDescription\": \"Enables the unmaximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the unminimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-unminimize\",\n \"markdownDescription\": \"Enables the unminimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the available_monitors command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-available-monitors\",\n \"markdownDescription\": \"Denies the available_monitors command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the center command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-center\",\n \"markdownDescription\": \"Denies the center command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-close\",\n \"markdownDescription\": \"Denies the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-create\",\n \"markdownDescription\": \"Denies the create command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the current_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-current-monitor\",\n \"markdownDescription\": \"Denies the current_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-cursor-position\",\n \"markdownDescription\": \"Denies the cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the destroy command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-destroy\",\n \"markdownDescription\": \"Denies the destroy command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get_all_windows command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-get-all-windows\",\n \"markdownDescription\": \"Denies the get_all_windows command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-hide\",\n \"markdownDescription\": \"Denies the hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the inner_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-inner-position\",\n \"markdownDescription\": \"Denies the inner_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the inner_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-inner-size\",\n \"markdownDescription\": \"Denies the inner_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the internal_toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-internal-toggle-maximize\",\n \"markdownDescription\": \"Denies the internal_toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-always-on-top\",\n \"markdownDescription\": \"Denies the is_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-closable\",\n \"markdownDescription\": \"Denies the is_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_decorated command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-decorated\",\n \"markdownDescription\": \"Denies the is_decorated command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-enabled\",\n \"markdownDescription\": \"Denies the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_focused command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-focused\",\n \"markdownDescription\": \"Denies the is_focused command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-fullscreen\",\n \"markdownDescription\": \"Denies the is_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-maximizable\",\n \"markdownDescription\": \"Denies the is_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_maximized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-maximized\",\n \"markdownDescription\": \"Denies the is_maximized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-minimizable\",\n \"markdownDescription\": \"Denies the is_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_minimized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-minimized\",\n \"markdownDescription\": \"Denies the is_minimized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-resizable\",\n \"markdownDescription\": \"Denies the is_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-visible\",\n \"markdownDescription\": \"Denies the is_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-maximize\",\n \"markdownDescription\": \"Denies the maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the minimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-minimize\",\n \"markdownDescription\": \"Denies the minimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the monitor_from_point command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-monitor-from-point\",\n \"markdownDescription\": \"Denies the monitor_from_point command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the outer_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-outer-position\",\n \"markdownDescription\": \"Denies the outer_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the outer_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-outer-size\",\n \"markdownDescription\": \"Denies the outer_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the primary_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-primary-monitor\",\n \"markdownDescription\": \"Denies the primary_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the request_user_attention command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-request-user-attention\",\n \"markdownDescription\": \"Denies the request_user_attention command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the scale_factor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-scale-factor\",\n \"markdownDescription\": \"Denies the scale_factor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_always_on_bottom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-always-on-bottom\",\n \"markdownDescription\": \"Denies the set_always_on_bottom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-always-on-top\",\n \"markdownDescription\": \"Denies the set_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-background-color\",\n \"markdownDescription\": \"Denies the set_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_badge_count command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-badge-count\",\n \"markdownDescription\": \"Denies the set_badge_count command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_badge_label command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-badge-label\",\n \"markdownDescription\": \"Denies the set_badge_label command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-closable\",\n \"markdownDescription\": \"Denies the set_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_content_protected command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-content-protected\",\n \"markdownDescription\": \"Denies the set_content_protected command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_grab command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-grab\",\n \"markdownDescription\": \"Denies the set_cursor_grab command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-icon\",\n \"markdownDescription\": \"Denies the set_cursor_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-position\",\n \"markdownDescription\": \"Denies the set_cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-visible\",\n \"markdownDescription\": \"Denies the set_cursor_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_decorations command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-decorations\",\n \"markdownDescription\": \"Denies the set_decorations command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_effects command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-effects\",\n \"markdownDescription\": \"Denies the set_effects command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-enabled\",\n \"markdownDescription\": \"Denies the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-focus\",\n \"markdownDescription\": \"Denies the set_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_focusable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-focusable\",\n \"markdownDescription\": \"Denies the set_focusable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-fullscreen\",\n \"markdownDescription\": \"Denies the set_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-icon\",\n \"markdownDescription\": \"Denies the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_ignore_cursor_events command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-ignore-cursor-events\",\n \"markdownDescription\": \"Denies the set_ignore_cursor_events command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_max_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-max-size\",\n \"markdownDescription\": \"Denies the set_max_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-maximizable\",\n \"markdownDescription\": \"Denies the set_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_min_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-min-size\",\n \"markdownDescription\": \"Denies the set_min_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-minimizable\",\n \"markdownDescription\": \"Denies the set_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_overlay_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-overlay-icon\",\n \"markdownDescription\": \"Denies the set_overlay_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-position\",\n \"markdownDescription\": \"Denies the set_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_progress_bar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-progress-bar\",\n \"markdownDescription\": \"Denies the set_progress_bar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-resizable\",\n \"markdownDescription\": \"Denies the set_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_shadow command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-shadow\",\n \"markdownDescription\": \"Denies the set_shadow command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_simple_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-simple-fullscreen\",\n \"markdownDescription\": \"Denies the set_simple_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-size\",\n \"markdownDescription\": \"Denies the set_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_size_constraints command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-size-constraints\",\n \"markdownDescription\": \"Denies the set_size_constraints command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_skip_taskbar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-skip-taskbar\",\n \"markdownDescription\": \"Denies the set_skip_taskbar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-theme\",\n \"markdownDescription\": \"Denies the set_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-title\",\n \"markdownDescription\": \"Denies the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_title_bar_style command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-title-bar-style\",\n \"markdownDescription\": \"Denies the set_title_bar_style command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_visible_on_all_workspaces command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-visible-on-all-workspaces\",\n \"markdownDescription\": \"Denies the set_visible_on_all_workspaces command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-show\",\n \"markdownDescription\": \"Denies the show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the start_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-start-dragging\",\n \"markdownDescription\": \"Denies the start_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the start_resize_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-start-resize-dragging\",\n \"markdownDescription\": \"Denies the start_resize_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-theme\",\n \"markdownDescription\": \"Denies the theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-title\",\n \"markdownDescription\": \"Denies the title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-toggle-maximize\",\n \"markdownDescription\": \"Denies the toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the unmaximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-unmaximize\",\n \"markdownDescription\": \"Denies the unmaximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the unminimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-unminimize\",\n \"markdownDescription\": \"Denies the unminimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\",\n \"type\": \"string\",\n \"const\": \"shell:default\",\n \"markdownDescription\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\"\n },\n {\n \"description\": \"Enables the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-execute\",\n \"markdownDescription\": \"Enables the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-kill\",\n \"markdownDescription\": \"Enables the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-open\",\n \"markdownDescription\": \"Enables the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-spawn\",\n \"markdownDescription\": \"Enables the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-stdin-write\",\n \"markdownDescription\": \"Enables the stdin_write command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-execute\",\n \"markdownDescription\": \"Denies the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-kill\",\n \"markdownDescription\": \"Denies the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-open\",\n \"markdownDescription\": \"Denies the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-spawn\",\n \"markdownDescription\": \"Denies the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-stdin-write\",\n \"markdownDescription\": \"Denies the stdin_write command without any pre-configured scope.\"\n },\n {\n \"description\": \"This permission set configures what kind of\\noperations are available from the window state plugin.\\n\\n#### Granted Permissions\\n\\nAll operations are enabled by default.\\n\\n\\n#### This default permission set includes:\\n\\n- `allow-filename`\\n- `allow-restore-state`\\n- `allow-save-window-state`\",\n \"type\": \"string\",\n \"const\": \"window-state:default\",\n \"markdownDescription\": \"This permission set configures what kind of\\noperations are available from the window state plugin.\\n\\n#### Granted Permissions\\n\\nAll operations are enabled by default.\\n\\n\\n#### This default permission set includes:\\n\\n- `allow-filename`\\n- `allow-restore-state`\\n- `allow-save-window-state`\"\n },\n {\n \"description\": \"Enables the filename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:allow-filename\",\n \"markdownDescription\": \"Enables the filename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the restore_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:allow-restore-state\",\n \"markdownDescription\": \"Enables the restore_state command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the save_window_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:allow-save-window-state\",\n \"markdownDescription\": \"Enables the save_window_state command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the filename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:deny-filename\",\n \"markdownDescription\": \"Denies the filename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the restore_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:deny-restore-state\",\n \"markdownDescription\": \"Denies the restore_state command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the save_window_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:deny-save-window-state\",\n \"markdownDescription\": \"Denies the save_window_state command without any pre-configured scope.\"\n }\n ]\n },\n \"Value\": {\n \"description\": \"All supported ACL values.\",\n \"anyOf\": [\n {\n \"description\": \"Represents a null JSON value.\",\n \"type\": \"null\"\n },\n {\n \"description\": \"Represents a [`bool`].\",\n \"type\": \"boolean\"\n },\n {\n \"description\": \"Represents a valid ACL [`Number`].\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Number\"\n }\n ]\n },\n {\n \"description\": \"Represents a [`String`].\",\n \"type\": \"string\"\n },\n {\n \"description\": \"Represents a list of other [`Value`]s.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/Value\"\n }\n },\n {\n \"description\": \"Represents a map of [`String`] keys to [`Value`]s.\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"$ref\": \"#/definitions/Value\"\n }\n }\n ]\n },\n \"Number\": {\n \"description\": \"A valid ACL number.\",\n \"anyOf\": [\n {\n \"description\": \"Represents an [`i64`].\",\n \"type\": \"integer\",\n \"format\": \"int64\"\n },\n {\n \"description\": \"Represents a [`f64`].\",\n \"type\": \"number\",\n \"format\": \"double\"\n }\n ]\n },\n \"Target\": {\n \"description\": \"Platform target.\",\n \"oneOf\": [\n {\n \"description\": \"MacOS.\",\n \"type\": \"string\",\n \"enum\": [\n \"macOS\"\n ]\n },\n {\n \"description\": \"Windows.\",\n \"type\": \"string\",\n \"enum\": [\n \"windows\"\n ]\n },\n {\n \"description\": \"Linux.\",\n \"type\": \"string\",\n \"enum\": [\n \"linux\"\n ]\n },\n {\n \"description\": \"Android.\",\n \"type\": \"string\",\n \"enum\": [\n \"android\"\n ]\n },\n {\n \"description\": \"iOS.\",\n \"type\": \"string\",\n \"enum\": [\n \"iOS\"\n ]\n }\n ]\n },\n \"ShellScopeEntryAllowedArg\": {\n \"description\": \"A command argument allowed to be executed by the webview API.\",\n \"anyOf\": [\n {\n \"description\": \"A non-configurable argument that is passed to the command in the order it was specified.\",\n \"type\": \"string\"\n },\n {\n \"description\": \"A variable that is set while calling the command from the webview API.\",\n \"type\": \"object\",\n \"required\": [\n \"validator\"\n ],\n \"properties\": {\n \"raw\": {\n \"description\": \"Marks the validator as a raw regex, meaning the plugin should not make any modification at runtime.\\n\\nThis means the regex will not match on the entire string by default, which might be exploited if your regex allow unexpected input to be considered valid. When using this option, make sure your regex is correct.\",\n \"default\": false,\n \"type\": \"boolean\"\n },\n \"validator\": {\n \"description\": \"[regex] validator to require passed values to conform to an expected input.\\n\\nThis will require the argument value passed to this variable to match the `validator` regex before it will be executed.\\n\\nThe regex string is by default surrounded by `^...$` to match the full string. For example the `https?://\\\\w+` regex would be registered as `^https?://\\\\w+$`.\\n\\n[regex]: \",\n \"type\": \"string\"\n }\n },\n \"additionalProperties\": false\n }\n ]\n },\n \"ShellScopeEntryAllowedArgs\": {\n \"description\": \"A set of command arguments allowed to be executed by the webview API.\\n\\nA value of `true` will allow any arguments to be passed to the command. `false` will disable all arguments. A list of [`ShellScopeEntryAllowedArg`] will set those arguments as the only valid arguments to be passed to the attached command configuration.\",\n \"anyOf\": [\n {\n \"description\": \"Use a simple boolean to allow all or disable all arguments to this command configuration.\",\n \"type\": \"boolean\"\n },\n {\n \"description\": \"A specific set of [`ShellScopeEntryAllowedArg`] that are valid to call for the command configuration.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArg\"\n }\n }\n ]\n }\n }\n}" }, { "path": "desktop/src-tauri/gen/schemas/macOS-schema.json", "content": "{\n \"$schema\": \"http://json-schema.org/draft-07/schema#\",\n \"title\": \"CapabilityFile\",\n \"description\": \"Capability formats accepted in a capability file.\",\n \"anyOf\": [\n {\n \"description\": \"A single capability.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Capability\"\n }\n ]\n },\n {\n \"description\": \"A list of capabilities.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/Capability\"\n }\n },\n {\n \"description\": \"A list of capabilities.\",\n \"type\": \"object\",\n \"required\": [\n \"capabilities\"\n ],\n \"properties\": {\n \"capabilities\": {\n \"description\": \"The list of capabilities.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/Capability\"\n }\n }\n }\n }\n ],\n \"definitions\": {\n \"Capability\": {\n \"description\": \"A grouping and boundary mechanism developers can use to isolate access to the IPC layer.\\n\\nIt controls application windows' and webviews' fine grained access to the Tauri core, application, or plugin commands. If a webview or its window is not matching any capability then it has no access to the IPC layer at all.\\n\\nThis can be done to create groups of windows, based on their required system access, which can reduce impact of frontend vulnerabilities in less privileged windows. Windows can be added to a capability by exact name (e.g. `main-window`) or glob patterns like `*` or `admin-*`. A Window can have none, one, or multiple associated capabilities.\\n\\n## Example\\n\\n```json { \\\"identifier\\\": \\\"main-user-files-write\\\", \\\"description\\\": \\\"This capability allows the `main` window on macOS and Windows access to `filesystem` write related commands and `dialog` commands to enable programmatic access to files selected by the user.\\\", \\\"windows\\\": [ \\\"main\\\" ], \\\"permissions\\\": [ \\\"core:default\\\", \\\"dialog:open\\\", { \\\"identifier\\\": \\\"fs:allow-write-text-file\\\", \\\"allow\\\": [{ \\\"path\\\": \\\"$HOME/test.txt\\\" }] }, ], \\\"platforms\\\": [\\\"macOS\\\",\\\"windows\\\"] } ```\",\n \"type\": \"object\",\n \"required\": [\n \"identifier\",\n \"permissions\"\n ],\n \"properties\": {\n \"identifier\": {\n \"description\": \"Identifier of the capability.\\n\\n## Example\\n\\n`main-user-files-write`\",\n \"type\": \"string\"\n },\n \"description\": {\n \"description\": \"Description of what the capability is intended to allow on associated windows.\\n\\nIt should contain a description of what the grouped permissions should allow.\\n\\n## Example\\n\\nThis capability allows the `main` window access to `filesystem` write related commands and `dialog` commands to enable programmatic access to files selected by the user.\",\n \"default\": \"\",\n \"type\": \"string\"\n },\n \"remote\": {\n \"description\": \"Configure remote URLs that can use the capability permissions.\\n\\nThis setting is optional and defaults to not being set, as our default use case is that the content is served from our local application.\\n\\n:::caution Make sure you understand the security implications of providing remote sources with local system access. :::\\n\\n## Example\\n\\n```json { \\\"urls\\\": [\\\"https://*.mydomain.dev\\\"] } ```\",\n \"anyOf\": [\n {\n \"$ref\": \"#/definitions/CapabilityRemote\"\n },\n {\n \"type\": \"null\"\n }\n ]\n },\n \"local\": {\n \"description\": \"Whether this capability is enabled for local app URLs or not. Defaults to `true`.\",\n \"default\": true,\n \"type\": \"boolean\"\n },\n \"windows\": {\n \"description\": \"List of windows that are affected by this capability. Can be a glob pattern.\\n\\nIf a window label matches any of the patterns in this list, the capability will be enabled on all the webviews of that window, regardless of the value of [`Self::webviews`].\\n\\nOn multiwebview windows, prefer specifying [`Self::webviews`] and omitting [`Self::windows`] for a fine grained access control.\\n\\n## Example\\n\\n`[\\\"main\\\"]`\",\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\"\n }\n },\n \"webviews\": {\n \"description\": \"List of webviews that are affected by this capability. Can be a glob pattern.\\n\\nThe capability will be enabled on all the webviews whose label matches any of the patterns in this list, regardless of whether the webview's window label matches a pattern in [`Self::windows`].\\n\\n## Example\\n\\n`[\\\"sub-webview-one\\\", \\\"sub-webview-two\\\"]`\",\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\"\n }\n },\n \"permissions\": {\n \"description\": \"List of permissions attached to this capability.\\n\\nMust include the plugin name as prefix in the form of `${plugin-name}:${permission-name}`. For commands directly implemented in the application itself only `${permission-name}` is required.\\n\\n## Example\\n\\n```json [ \\\"core:default\\\", \\\"shell:allow-open\\\", \\\"dialog:open\\\", { \\\"identifier\\\": \\\"fs:allow-write-text-file\\\", \\\"allow\\\": [{ \\\"path\\\": \\\"$HOME/test.txt\\\" }] } ] ```\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/PermissionEntry\"\n },\n \"uniqueItems\": true\n },\n \"platforms\": {\n \"description\": \"Limit which target platforms this capability applies to.\\n\\nBy default all platforms are targeted.\\n\\n## Example\\n\\n`[\\\"macOS\\\",\\\"windows\\\"]`\",\n \"type\": [\n \"array\",\n \"null\"\n ],\n \"items\": {\n \"$ref\": \"#/definitions/Target\"\n }\n }\n }\n },\n \"CapabilityRemote\": {\n \"description\": \"Configuration for remote URLs that are associated with the capability.\",\n \"type\": \"object\",\n \"required\": [\n \"urls\"\n ],\n \"properties\": {\n \"urls\": {\n \"description\": \"Remote domains this capability refers to using the [URLPattern standard](https://urlpattern.spec.whatwg.org/).\\n\\n## Examples\\n\\n- \\\"https://*.mydomain.dev\\\": allows subdomains of mydomain.dev - \\\"https://mydomain.dev/api/*\\\": allows any subpath of mydomain.dev/api\",\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\"\n }\n }\n }\n },\n \"PermissionEntry\": {\n \"description\": \"An entry for a permission value in a [`Capability`] can be either a raw permission [`Identifier`] or an object that references a permission and extends its scope.\",\n \"anyOf\": [\n {\n \"description\": \"Reference a permission or permission set by identifier.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Identifier\"\n }\n ]\n },\n {\n \"description\": \"Reference a permission or permission set by identifier and extends its scope.\",\n \"type\": \"object\",\n \"allOf\": [\n {\n \"if\": {\n \"properties\": {\n \"identifier\": {\n \"anyOf\": [\n {\n \"description\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\",\n \"type\": \"string\",\n \"const\": \"shell:default\",\n \"markdownDescription\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\"\n },\n {\n \"description\": \"Enables the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-execute\",\n \"markdownDescription\": \"Enables the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-kill\",\n \"markdownDescription\": \"Enables the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-open\",\n \"markdownDescription\": \"Enables the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-spawn\",\n \"markdownDescription\": \"Enables the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-stdin-write\",\n \"markdownDescription\": \"Enables the stdin_write command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-execute\",\n \"markdownDescription\": \"Denies the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-kill\",\n \"markdownDescription\": \"Denies the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-open\",\n \"markdownDescription\": \"Denies the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-spawn\",\n \"markdownDescription\": \"Denies the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-stdin-write\",\n \"markdownDescription\": \"Denies the stdin_write command without any pre-configured scope.\"\n }\n ]\n }\n }\n },\n \"then\": {\n \"properties\": {\n \"allow\": {\n \"items\": {\n \"title\": \"ShellScopeEntry\",\n \"description\": \"Shell scope entry.\",\n \"anyOf\": [\n {\n \"type\": \"object\",\n \"required\": [\n \"cmd\",\n \"name\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"cmd\": {\n \"description\": \"The command name. It can start with a variable that resolves to a system base directory. The variables are: `$AUDIO`, `$CACHE`, `$CONFIG`, `$DATA`, `$LOCALDATA`, `$DESKTOP`, `$DOCUMENT`, `$DOWNLOAD`, `$EXE`, `$FONT`, `$HOME`, `$PICTURE`, `$PUBLIC`, `$RUNTIME`, `$TEMPLATE`, `$VIDEO`, `$RESOURCE`, `$LOG`, `$TEMP`, `$APPCONFIG`, `$APPDATA`, `$APPLOCALDATA`, `$APPCACHE`, `$APPLOG`.\",\n \"type\": \"string\"\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n }\n },\n \"additionalProperties\": false\n },\n {\n \"type\": \"object\",\n \"required\": [\n \"name\",\n \"sidecar\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n },\n \"sidecar\": {\n \"description\": \"If this command is a sidecar command.\",\n \"type\": \"boolean\"\n }\n },\n \"additionalProperties\": false\n }\n ]\n }\n },\n \"deny\": {\n \"items\": {\n \"title\": \"ShellScopeEntry\",\n \"description\": \"Shell scope entry.\",\n \"anyOf\": [\n {\n \"type\": \"object\",\n \"required\": [\n \"cmd\",\n \"name\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"cmd\": {\n \"description\": \"The command name. It can start with a variable that resolves to a system base directory. The variables are: `$AUDIO`, `$CACHE`, `$CONFIG`, `$DATA`, `$LOCALDATA`, `$DESKTOP`, `$DOCUMENT`, `$DOWNLOAD`, `$EXE`, `$FONT`, `$HOME`, `$PICTURE`, `$PUBLIC`, `$RUNTIME`, `$TEMPLATE`, `$VIDEO`, `$RESOURCE`, `$LOG`, `$TEMP`, `$APPCONFIG`, `$APPDATA`, `$APPLOCALDATA`, `$APPCACHE`, `$APPLOG`.\",\n \"type\": \"string\"\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n }\n },\n \"additionalProperties\": false\n },\n {\n \"type\": \"object\",\n \"required\": [\n \"name\",\n \"sidecar\"\n ],\n \"properties\": {\n \"args\": {\n \"description\": \"The allowed arguments for the command execution.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArgs\"\n }\n ]\n },\n \"name\": {\n \"description\": \"The name for this allowed shell command configuration.\\n\\nThis name will be used inside of the webview API to call this command along with any specified arguments.\",\n \"type\": \"string\"\n },\n \"sidecar\": {\n \"description\": \"If this command is a sidecar command.\",\n \"type\": \"boolean\"\n }\n },\n \"additionalProperties\": false\n }\n ]\n }\n }\n }\n },\n \"properties\": {\n \"identifier\": {\n \"description\": \"Identifier of the permission or permission set.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Identifier\"\n }\n ]\n }\n }\n },\n {\n \"properties\": {\n \"identifier\": {\n \"description\": \"Identifier of the permission or permission set.\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Identifier\"\n }\n ]\n },\n \"allow\": {\n \"description\": \"Data that defines what is allowed by the scope.\",\n \"type\": [\n \"array\",\n \"null\"\n ],\n \"items\": {\n \"$ref\": \"#/definitions/Value\"\n }\n },\n \"deny\": {\n \"description\": \"Data that defines what is denied by the scope. This should be prioritized by validation logic.\",\n \"type\": [\n \"array\",\n \"null\"\n ],\n \"items\": {\n \"$ref\": \"#/definitions/Value\"\n }\n }\n }\n }\n ],\n \"required\": [\n \"identifier\"\n ]\n }\n ]\n },\n \"Identifier\": {\n \"description\": \"Permission identifier\",\n \"oneOf\": [\n {\n \"description\": \"Default core plugins set.\\n#### This default permission set includes:\\n\\n- `core:path:default`\\n- `core:event:default`\\n- `core:window:default`\\n- `core:webview:default`\\n- `core:app:default`\\n- `core:image:default`\\n- `core:resources:default`\\n- `core:menu:default`\\n- `core:tray:default`\",\n \"type\": \"string\",\n \"const\": \"core:default\",\n \"markdownDescription\": \"Default core plugins set.\\n#### This default permission set includes:\\n\\n- `core:path:default`\\n- `core:event:default`\\n- `core:window:default`\\n- `core:webview:default`\\n- `core:app:default`\\n- `core:image:default`\\n- `core:resources:default`\\n- `core:menu:default`\\n- `core:tray:default`\"\n },\n {\n \"description\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-version`\\n- `allow-name`\\n- `allow-tauri-version`\\n- `allow-identifier`\\n- `allow-bundle-type`\\n- `allow-register-listener`\\n- `allow-remove-listener`\",\n \"type\": \"string\",\n \"const\": \"core:app:default\",\n \"markdownDescription\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-version`\\n- `allow-name`\\n- `allow-tauri-version`\\n- `allow-identifier`\\n- `allow-bundle-type`\\n- `allow-register-listener`\\n- `allow-remove-listener`\"\n },\n {\n \"description\": \"Enables the app_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-app-hide\",\n \"markdownDescription\": \"Enables the app_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the app_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-app-show\",\n \"markdownDescription\": \"Enables the app_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the bundle_type command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-bundle-type\",\n \"markdownDescription\": \"Enables the bundle_type command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the default_window_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-default-window-icon\",\n \"markdownDescription\": \"Enables the default_window_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the fetch_data_store_identifiers command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-fetch-data-store-identifiers\",\n \"markdownDescription\": \"Enables the fetch_data_store_identifiers command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the identifier command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-identifier\",\n \"markdownDescription\": \"Enables the identifier command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the name command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-name\",\n \"markdownDescription\": \"Enables the name command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the register_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-register-listener\",\n \"markdownDescription\": \"Enables the register_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_data_store command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-remove-data-store\",\n \"markdownDescription\": \"Enables the remove_data_store command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-remove-listener\",\n \"markdownDescription\": \"Enables the remove_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_app_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-set-app-theme\",\n \"markdownDescription\": \"Enables the set_app_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_dock_visibility command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-set-dock-visibility\",\n \"markdownDescription\": \"Enables the set_dock_visibility command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the tauri_version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-tauri-version\",\n \"markdownDescription\": \"Enables the tauri_version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:allow-version\",\n \"markdownDescription\": \"Enables the version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the app_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-app-hide\",\n \"markdownDescription\": \"Denies the app_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the app_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-app-show\",\n \"markdownDescription\": \"Denies the app_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the bundle_type command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-bundle-type\",\n \"markdownDescription\": \"Denies the bundle_type command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the default_window_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-default-window-icon\",\n \"markdownDescription\": \"Denies the default_window_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the fetch_data_store_identifiers command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-fetch-data-store-identifiers\",\n \"markdownDescription\": \"Denies the fetch_data_store_identifiers command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the identifier command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-identifier\",\n \"markdownDescription\": \"Denies the identifier command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the name command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-name\",\n \"markdownDescription\": \"Denies the name command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the register_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-register-listener\",\n \"markdownDescription\": \"Denies the register_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_data_store command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-remove-data-store\",\n \"markdownDescription\": \"Denies the remove_data_store command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_listener command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-remove-listener\",\n \"markdownDescription\": \"Denies the remove_listener command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_app_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-set-app-theme\",\n \"markdownDescription\": \"Denies the set_app_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_dock_visibility command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-set-dock-visibility\",\n \"markdownDescription\": \"Denies the set_dock_visibility command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the tauri_version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-tauri-version\",\n \"markdownDescription\": \"Denies the tauri_version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the version command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:app:deny-version\",\n \"markdownDescription\": \"Denies the version command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-listen`\\n- `allow-unlisten`\\n- `allow-emit`\\n- `allow-emit-to`\",\n \"type\": \"string\",\n \"const\": \"core:event:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-listen`\\n- `allow-unlisten`\\n- `allow-emit`\\n- `allow-emit-to`\"\n },\n {\n \"description\": \"Enables the emit command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-emit\",\n \"markdownDescription\": \"Enables the emit command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the emit_to command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-emit-to\",\n \"markdownDescription\": \"Enables the emit_to command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the listen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-listen\",\n \"markdownDescription\": \"Enables the listen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the unlisten command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:allow-unlisten\",\n \"markdownDescription\": \"Enables the unlisten command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the emit command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-emit\",\n \"markdownDescription\": \"Denies the emit command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the emit_to command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-emit-to\",\n \"markdownDescription\": \"Denies the emit_to command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the listen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-listen\",\n \"markdownDescription\": \"Denies the listen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the unlisten command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:event:deny-unlisten\",\n \"markdownDescription\": \"Denies the unlisten command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-from-bytes`\\n- `allow-from-path`\\n- `allow-rgba`\\n- `allow-size`\",\n \"type\": \"string\",\n \"const\": \"core:image:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-from-bytes`\\n- `allow-from-path`\\n- `allow-rgba`\\n- `allow-size`\"\n },\n {\n \"description\": \"Enables the from_bytes command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-from-bytes\",\n \"markdownDescription\": \"Enables the from_bytes command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the from_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-from-path\",\n \"markdownDescription\": \"Enables the from_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-new\",\n \"markdownDescription\": \"Enables the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the rgba command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-rgba\",\n \"markdownDescription\": \"Enables the rgba command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:allow-size\",\n \"markdownDescription\": \"Enables the size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the from_bytes command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-from-bytes\",\n \"markdownDescription\": \"Denies the from_bytes command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the from_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-from-path\",\n \"markdownDescription\": \"Denies the from_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-new\",\n \"markdownDescription\": \"Denies the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the rgba command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-rgba\",\n \"markdownDescription\": \"Denies the rgba command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:image:deny-size\",\n \"markdownDescription\": \"Denies the size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-append`\\n- `allow-prepend`\\n- `allow-insert`\\n- `allow-remove`\\n- `allow-remove-at`\\n- `allow-items`\\n- `allow-get`\\n- `allow-popup`\\n- `allow-create-default`\\n- `allow-set-as-app-menu`\\n- `allow-set-as-window-menu`\\n- `allow-text`\\n- `allow-set-text`\\n- `allow-is-enabled`\\n- `allow-set-enabled`\\n- `allow-set-accelerator`\\n- `allow-set-as-windows-menu-for-nsapp`\\n- `allow-set-as-help-menu-for-nsapp`\\n- `allow-is-checked`\\n- `allow-set-checked`\\n- `allow-set-icon`\",\n \"type\": \"string\",\n \"const\": \"core:menu:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-append`\\n- `allow-prepend`\\n- `allow-insert`\\n- `allow-remove`\\n- `allow-remove-at`\\n- `allow-items`\\n- `allow-get`\\n- `allow-popup`\\n- `allow-create-default`\\n- `allow-set-as-app-menu`\\n- `allow-set-as-window-menu`\\n- `allow-text`\\n- `allow-set-text`\\n- `allow-is-enabled`\\n- `allow-set-enabled`\\n- `allow-set-accelerator`\\n- `allow-set-as-windows-menu-for-nsapp`\\n- `allow-set-as-help-menu-for-nsapp`\\n- `allow-is-checked`\\n- `allow-set-checked`\\n- `allow-set-icon`\"\n },\n {\n \"description\": \"Enables the append command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-append\",\n \"markdownDescription\": \"Enables the append command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create_default command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-create-default\",\n \"markdownDescription\": \"Enables the create_default command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the get command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-get\",\n \"markdownDescription\": \"Enables the get command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the insert command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-insert\",\n \"markdownDescription\": \"Enables the insert command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-is-checked\",\n \"markdownDescription\": \"Enables the is_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-is-enabled\",\n \"markdownDescription\": \"Enables the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the items command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-items\",\n \"markdownDescription\": \"Enables the items command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-new\",\n \"markdownDescription\": \"Enables the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the popup command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-popup\",\n \"markdownDescription\": \"Enables the popup command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the prepend command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-prepend\",\n \"markdownDescription\": \"Enables the prepend command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-remove\",\n \"markdownDescription\": \"Enables the remove command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_at command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-remove-at\",\n \"markdownDescription\": \"Enables the remove_at command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_accelerator command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-accelerator\",\n \"markdownDescription\": \"Enables the set_accelerator command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_app_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-app-menu\",\n \"markdownDescription\": \"Enables the set_as_app_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_help_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-help-menu-for-nsapp\",\n \"markdownDescription\": \"Enables the set_as_help_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_window_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-window-menu\",\n \"markdownDescription\": \"Enables the set_as_window_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_as_windows_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-as-windows-menu-for-nsapp\",\n \"markdownDescription\": \"Enables the set_as_windows_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-checked\",\n \"markdownDescription\": \"Enables the set_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-enabled\",\n \"markdownDescription\": \"Enables the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-icon\",\n \"markdownDescription\": \"Enables the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-set-text\",\n \"markdownDescription\": \"Enables the set_text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:allow-text\",\n \"markdownDescription\": \"Enables the text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the append command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-append\",\n \"markdownDescription\": \"Denies the append command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create_default command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-create-default\",\n \"markdownDescription\": \"Denies the create_default command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-get\",\n \"markdownDescription\": \"Denies the get command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the insert command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-insert\",\n \"markdownDescription\": \"Denies the insert command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-is-checked\",\n \"markdownDescription\": \"Denies the is_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-is-enabled\",\n \"markdownDescription\": \"Denies the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the items command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-items\",\n \"markdownDescription\": \"Denies the items command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-new\",\n \"markdownDescription\": \"Denies the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the popup command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-popup\",\n \"markdownDescription\": \"Denies the popup command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the prepend command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-prepend\",\n \"markdownDescription\": \"Denies the prepend command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-remove\",\n \"markdownDescription\": \"Denies the remove command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_at command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-remove-at\",\n \"markdownDescription\": \"Denies the remove_at command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_accelerator command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-accelerator\",\n \"markdownDescription\": \"Denies the set_accelerator command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_app_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-app-menu\",\n \"markdownDescription\": \"Denies the set_as_app_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_help_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-help-menu-for-nsapp\",\n \"markdownDescription\": \"Denies the set_as_help_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_window_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-window-menu\",\n \"markdownDescription\": \"Denies the set_as_window_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_as_windows_menu_for_nsapp command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-as-windows-menu-for-nsapp\",\n \"markdownDescription\": \"Denies the set_as_windows_menu_for_nsapp command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_checked command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-checked\",\n \"markdownDescription\": \"Denies the set_checked command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-enabled\",\n \"markdownDescription\": \"Denies the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-icon\",\n \"markdownDescription\": \"Denies the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-set-text\",\n \"markdownDescription\": \"Denies the set_text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the text command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:menu:deny-text\",\n \"markdownDescription\": \"Denies the text command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-resolve-directory`\\n- `allow-resolve`\\n- `allow-normalize`\\n- `allow-join`\\n- `allow-dirname`\\n- `allow-extname`\\n- `allow-basename`\\n- `allow-is-absolute`\",\n \"type\": \"string\",\n \"const\": \"core:path:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-resolve-directory`\\n- `allow-resolve`\\n- `allow-normalize`\\n- `allow-join`\\n- `allow-dirname`\\n- `allow-extname`\\n- `allow-basename`\\n- `allow-is-absolute`\"\n },\n {\n \"description\": \"Enables the basename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-basename\",\n \"markdownDescription\": \"Enables the basename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the dirname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-dirname\",\n \"markdownDescription\": \"Enables the dirname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the extname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-extname\",\n \"markdownDescription\": \"Enables the extname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_absolute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-is-absolute\",\n \"markdownDescription\": \"Enables the is_absolute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the join command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-join\",\n \"markdownDescription\": \"Enables the join command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the normalize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-normalize\",\n \"markdownDescription\": \"Enables the normalize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the resolve command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-resolve\",\n \"markdownDescription\": \"Enables the resolve command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the resolve_directory command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:allow-resolve-directory\",\n \"markdownDescription\": \"Enables the resolve_directory command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the basename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-basename\",\n \"markdownDescription\": \"Denies the basename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the dirname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-dirname\",\n \"markdownDescription\": \"Denies the dirname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the extname command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-extname\",\n \"markdownDescription\": \"Denies the extname command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_absolute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-is-absolute\",\n \"markdownDescription\": \"Denies the is_absolute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the join command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-join\",\n \"markdownDescription\": \"Denies the join command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the normalize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-normalize\",\n \"markdownDescription\": \"Denies the normalize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the resolve command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-resolve\",\n \"markdownDescription\": \"Denies the resolve command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the resolve_directory command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:path:deny-resolve-directory\",\n \"markdownDescription\": \"Denies the resolve_directory command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-close`\",\n \"type\": \"string\",\n \"const\": \"core:resources:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-close`\"\n },\n {\n \"description\": \"Enables the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:resources:allow-close\",\n \"markdownDescription\": \"Enables the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:resources:deny-close\",\n \"markdownDescription\": \"Denies the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-get-by-id`\\n- `allow-remove-by-id`\\n- `allow-set-icon`\\n- `allow-set-menu`\\n- `allow-set-tooltip`\\n- `allow-set-title`\\n- `allow-set-visible`\\n- `allow-set-temp-dir-path`\\n- `allow-set-icon-as-template`\\n- `allow-set-show-menu-on-left-click`\",\n \"type\": \"string\",\n \"const\": \"core:tray:default\",\n \"markdownDescription\": \"Default permissions for the plugin, which enables all commands.\\n#### This default permission set includes:\\n\\n- `allow-new`\\n- `allow-get-by-id`\\n- `allow-remove-by-id`\\n- `allow-set-icon`\\n- `allow-set-menu`\\n- `allow-set-tooltip`\\n- `allow-set-title`\\n- `allow-set-visible`\\n- `allow-set-temp-dir-path`\\n- `allow-set-icon-as-template`\\n- `allow-set-show-menu-on-left-click`\"\n },\n {\n \"description\": \"Enables the get_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-get-by-id\",\n \"markdownDescription\": \"Enables the get_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-new\",\n \"markdownDescription\": \"Enables the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the remove_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-remove-by-id\",\n \"markdownDescription\": \"Enables the remove_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-icon\",\n \"markdownDescription\": \"Enables the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon_as_template command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-icon-as-template\",\n \"markdownDescription\": \"Enables the set_icon_as_template command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-menu\",\n \"markdownDescription\": \"Enables the set_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_show_menu_on_left_click command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-show-menu-on-left-click\",\n \"markdownDescription\": \"Enables the set_show_menu_on_left_click command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_temp_dir_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-temp-dir-path\",\n \"markdownDescription\": \"Enables the set_temp_dir_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-title\",\n \"markdownDescription\": \"Enables the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_tooltip command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-tooltip\",\n \"markdownDescription\": \"Enables the set_tooltip command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:allow-set-visible\",\n \"markdownDescription\": \"Enables the set_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-get-by-id\",\n \"markdownDescription\": \"Denies the get_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the new command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-new\",\n \"markdownDescription\": \"Denies the new command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the remove_by_id command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-remove-by-id\",\n \"markdownDescription\": \"Denies the remove_by_id command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-icon\",\n \"markdownDescription\": \"Denies the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon_as_template command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-icon-as-template\",\n \"markdownDescription\": \"Denies the set_icon_as_template command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_menu command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-menu\",\n \"markdownDescription\": \"Denies the set_menu command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_show_menu_on_left_click command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-show-menu-on-left-click\",\n \"markdownDescription\": \"Denies the set_show_menu_on_left_click command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_temp_dir_path command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-temp-dir-path\",\n \"markdownDescription\": \"Denies the set_temp_dir_path command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-title\",\n \"markdownDescription\": \"Denies the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_tooltip command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-tooltip\",\n \"markdownDescription\": \"Denies the set_tooltip command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:tray:deny-set-visible\",\n \"markdownDescription\": \"Denies the set_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-webviews`\\n- `allow-webview-position`\\n- `allow-webview-size`\\n- `allow-internal-toggle-devtools`\",\n \"type\": \"string\",\n \"const\": \"core:webview:default\",\n \"markdownDescription\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-webviews`\\n- `allow-webview-position`\\n- `allow-webview-size`\\n- `allow-internal-toggle-devtools`\"\n },\n {\n \"description\": \"Enables the clear_all_browsing_data command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-clear-all-browsing-data\",\n \"markdownDescription\": \"Enables the clear_all_browsing_data command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create_webview command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-create-webview\",\n \"markdownDescription\": \"Enables the create_webview command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create_webview_window command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-create-webview-window\",\n \"markdownDescription\": \"Enables the create_webview_window command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the get_all_webviews command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-get-all-webviews\",\n \"markdownDescription\": \"Enables the get_all_webviews command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the internal_toggle_devtools command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-internal-toggle-devtools\",\n \"markdownDescription\": \"Enables the internal_toggle_devtools command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the print command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-print\",\n \"markdownDescription\": \"Enables the print command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the reparent command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-reparent\",\n \"markdownDescription\": \"Enables the reparent command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_auto_resize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-auto-resize\",\n \"markdownDescription\": \"Enables the set_webview_auto_resize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-background-color\",\n \"markdownDescription\": \"Enables the set_webview_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-focus\",\n \"markdownDescription\": \"Enables the set_webview_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-position\",\n \"markdownDescription\": \"Enables the set_webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-size\",\n \"markdownDescription\": \"Enables the set_webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_webview_zoom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-set-webview-zoom\",\n \"markdownDescription\": \"Enables the set_webview_zoom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-close\",\n \"markdownDescription\": \"Enables the webview_close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-hide\",\n \"markdownDescription\": \"Enables the webview_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-position\",\n \"markdownDescription\": \"Enables the webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-show\",\n \"markdownDescription\": \"Enables the webview_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:allow-webview-size\",\n \"markdownDescription\": \"Enables the webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the clear_all_browsing_data command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-clear-all-browsing-data\",\n \"markdownDescription\": \"Denies the clear_all_browsing_data command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create_webview command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-create-webview\",\n \"markdownDescription\": \"Denies the create_webview command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create_webview_window command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-create-webview-window\",\n \"markdownDescription\": \"Denies the create_webview_window command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get_all_webviews command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-get-all-webviews\",\n \"markdownDescription\": \"Denies the get_all_webviews command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the internal_toggle_devtools command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-internal-toggle-devtools\",\n \"markdownDescription\": \"Denies the internal_toggle_devtools command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the print command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-print\",\n \"markdownDescription\": \"Denies the print command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the reparent command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-reparent\",\n \"markdownDescription\": \"Denies the reparent command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_auto_resize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-auto-resize\",\n \"markdownDescription\": \"Denies the set_webview_auto_resize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-background-color\",\n \"markdownDescription\": \"Denies the set_webview_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-focus\",\n \"markdownDescription\": \"Denies the set_webview_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-position\",\n \"markdownDescription\": \"Denies the set_webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-size\",\n \"markdownDescription\": \"Denies the set_webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_webview_zoom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-set-webview-zoom\",\n \"markdownDescription\": \"Denies the set_webview_zoom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-close\",\n \"markdownDescription\": \"Denies the webview_close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-hide\",\n \"markdownDescription\": \"Denies the webview_hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-position\",\n \"markdownDescription\": \"Denies the webview_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-show\",\n \"markdownDescription\": \"Denies the webview_show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the webview_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:webview:deny-webview-size\",\n \"markdownDescription\": \"Denies the webview_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-windows`\\n- `allow-scale-factor`\\n- `allow-inner-position`\\n- `allow-outer-position`\\n- `allow-inner-size`\\n- `allow-outer-size`\\n- `allow-is-fullscreen`\\n- `allow-is-minimized`\\n- `allow-is-maximized`\\n- `allow-is-focused`\\n- `allow-is-decorated`\\n- `allow-is-resizable`\\n- `allow-is-maximizable`\\n- `allow-is-minimizable`\\n- `allow-is-closable`\\n- `allow-is-visible`\\n- `allow-is-enabled`\\n- `allow-title`\\n- `allow-current-monitor`\\n- `allow-primary-monitor`\\n- `allow-monitor-from-point`\\n- `allow-available-monitors`\\n- `allow-cursor-position`\\n- `allow-theme`\\n- `allow-is-always-on-top`\\n- `allow-internal-toggle-maximize`\",\n \"type\": \"string\",\n \"const\": \"core:window:default\",\n \"markdownDescription\": \"Default permissions for the plugin.\\n#### This default permission set includes:\\n\\n- `allow-get-all-windows`\\n- `allow-scale-factor`\\n- `allow-inner-position`\\n- `allow-outer-position`\\n- `allow-inner-size`\\n- `allow-outer-size`\\n- `allow-is-fullscreen`\\n- `allow-is-minimized`\\n- `allow-is-maximized`\\n- `allow-is-focused`\\n- `allow-is-decorated`\\n- `allow-is-resizable`\\n- `allow-is-maximizable`\\n- `allow-is-minimizable`\\n- `allow-is-closable`\\n- `allow-is-visible`\\n- `allow-is-enabled`\\n- `allow-title`\\n- `allow-current-monitor`\\n- `allow-primary-monitor`\\n- `allow-monitor-from-point`\\n- `allow-available-monitors`\\n- `allow-cursor-position`\\n- `allow-theme`\\n- `allow-is-always-on-top`\\n- `allow-internal-toggle-maximize`\"\n },\n {\n \"description\": \"Enables the available_monitors command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-available-monitors\",\n \"markdownDescription\": \"Enables the available_monitors command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the center command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-center\",\n \"markdownDescription\": \"Enables the center command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-close\",\n \"markdownDescription\": \"Enables the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the create command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-create\",\n \"markdownDescription\": \"Enables the create command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the current_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-current-monitor\",\n \"markdownDescription\": \"Enables the current_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-cursor-position\",\n \"markdownDescription\": \"Enables the cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the destroy command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-destroy\",\n \"markdownDescription\": \"Enables the destroy command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the get_all_windows command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-get-all-windows\",\n \"markdownDescription\": \"Enables the get_all_windows command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-hide\",\n \"markdownDescription\": \"Enables the hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the inner_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-inner-position\",\n \"markdownDescription\": \"Enables the inner_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the inner_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-inner-size\",\n \"markdownDescription\": \"Enables the inner_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the internal_toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-internal-toggle-maximize\",\n \"markdownDescription\": \"Enables the internal_toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-always-on-top\",\n \"markdownDescription\": \"Enables the is_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-closable\",\n \"markdownDescription\": \"Enables the is_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_decorated command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-decorated\",\n \"markdownDescription\": \"Enables the is_decorated command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-enabled\",\n \"markdownDescription\": \"Enables the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_focused command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-focused\",\n \"markdownDescription\": \"Enables the is_focused command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-fullscreen\",\n \"markdownDescription\": \"Enables the is_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-maximizable\",\n \"markdownDescription\": \"Enables the is_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_maximized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-maximized\",\n \"markdownDescription\": \"Enables the is_maximized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-minimizable\",\n \"markdownDescription\": \"Enables the is_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_minimized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-minimized\",\n \"markdownDescription\": \"Enables the is_minimized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-resizable\",\n \"markdownDescription\": \"Enables the is_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the is_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-is-visible\",\n \"markdownDescription\": \"Enables the is_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-maximize\",\n \"markdownDescription\": \"Enables the maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the minimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-minimize\",\n \"markdownDescription\": \"Enables the minimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the monitor_from_point command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-monitor-from-point\",\n \"markdownDescription\": \"Enables the monitor_from_point command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the outer_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-outer-position\",\n \"markdownDescription\": \"Enables the outer_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the outer_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-outer-size\",\n \"markdownDescription\": \"Enables the outer_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the primary_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-primary-monitor\",\n \"markdownDescription\": \"Enables the primary_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the request_user_attention command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-request-user-attention\",\n \"markdownDescription\": \"Enables the request_user_attention command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the scale_factor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-scale-factor\",\n \"markdownDescription\": \"Enables the scale_factor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_always_on_bottom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-always-on-bottom\",\n \"markdownDescription\": \"Enables the set_always_on_bottom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-always-on-top\",\n \"markdownDescription\": \"Enables the set_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-background-color\",\n \"markdownDescription\": \"Enables the set_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_badge_count command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-badge-count\",\n \"markdownDescription\": \"Enables the set_badge_count command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_badge_label command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-badge-label\",\n \"markdownDescription\": \"Enables the set_badge_label command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-closable\",\n \"markdownDescription\": \"Enables the set_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_content_protected command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-content-protected\",\n \"markdownDescription\": \"Enables the set_content_protected command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_grab command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-grab\",\n \"markdownDescription\": \"Enables the set_cursor_grab command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-icon\",\n \"markdownDescription\": \"Enables the set_cursor_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-position\",\n \"markdownDescription\": \"Enables the set_cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_cursor_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-cursor-visible\",\n \"markdownDescription\": \"Enables the set_cursor_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_decorations command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-decorations\",\n \"markdownDescription\": \"Enables the set_decorations command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_effects command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-effects\",\n \"markdownDescription\": \"Enables the set_effects command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-enabled\",\n \"markdownDescription\": \"Enables the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-focus\",\n \"markdownDescription\": \"Enables the set_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_focusable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-focusable\",\n \"markdownDescription\": \"Enables the set_focusable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-fullscreen\",\n \"markdownDescription\": \"Enables the set_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-icon\",\n \"markdownDescription\": \"Enables the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_ignore_cursor_events command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-ignore-cursor-events\",\n \"markdownDescription\": \"Enables the set_ignore_cursor_events command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_max_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-max-size\",\n \"markdownDescription\": \"Enables the set_max_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-maximizable\",\n \"markdownDescription\": \"Enables the set_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_min_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-min-size\",\n \"markdownDescription\": \"Enables the set_min_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-minimizable\",\n \"markdownDescription\": \"Enables the set_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_overlay_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-overlay-icon\",\n \"markdownDescription\": \"Enables the set_overlay_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-position\",\n \"markdownDescription\": \"Enables the set_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_progress_bar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-progress-bar\",\n \"markdownDescription\": \"Enables the set_progress_bar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-resizable\",\n \"markdownDescription\": \"Enables the set_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_shadow command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-shadow\",\n \"markdownDescription\": \"Enables the set_shadow command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_simple_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-simple-fullscreen\",\n \"markdownDescription\": \"Enables the set_simple_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-size\",\n \"markdownDescription\": \"Enables the set_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_size_constraints command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-size-constraints\",\n \"markdownDescription\": \"Enables the set_size_constraints command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_skip_taskbar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-skip-taskbar\",\n \"markdownDescription\": \"Enables the set_skip_taskbar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-theme\",\n \"markdownDescription\": \"Enables the set_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-title\",\n \"markdownDescription\": \"Enables the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_title_bar_style command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-title-bar-style\",\n \"markdownDescription\": \"Enables the set_title_bar_style command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the set_visible_on_all_workspaces command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-set-visible-on-all-workspaces\",\n \"markdownDescription\": \"Enables the set_visible_on_all_workspaces command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-show\",\n \"markdownDescription\": \"Enables the show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the start_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-start-dragging\",\n \"markdownDescription\": \"Enables the start_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the start_resize_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-start-resize-dragging\",\n \"markdownDescription\": \"Enables the start_resize_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-theme\",\n \"markdownDescription\": \"Enables the theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-title\",\n \"markdownDescription\": \"Enables the title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-toggle-maximize\",\n \"markdownDescription\": \"Enables the toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the unmaximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-unmaximize\",\n \"markdownDescription\": \"Enables the unmaximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the unminimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:allow-unminimize\",\n \"markdownDescription\": \"Enables the unminimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the available_monitors command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-available-monitors\",\n \"markdownDescription\": \"Denies the available_monitors command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the center command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-center\",\n \"markdownDescription\": \"Denies the center command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the close command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-close\",\n \"markdownDescription\": \"Denies the close command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the create command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-create\",\n \"markdownDescription\": \"Denies the create command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the current_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-current-monitor\",\n \"markdownDescription\": \"Denies the current_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-cursor-position\",\n \"markdownDescription\": \"Denies the cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the destroy command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-destroy\",\n \"markdownDescription\": \"Denies the destroy command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the get_all_windows command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-get-all-windows\",\n \"markdownDescription\": \"Denies the get_all_windows command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the hide command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-hide\",\n \"markdownDescription\": \"Denies the hide command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the inner_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-inner-position\",\n \"markdownDescription\": \"Denies the inner_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the inner_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-inner-size\",\n \"markdownDescription\": \"Denies the inner_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the internal_toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-internal-toggle-maximize\",\n \"markdownDescription\": \"Denies the internal_toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-always-on-top\",\n \"markdownDescription\": \"Denies the is_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-closable\",\n \"markdownDescription\": \"Denies the is_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_decorated command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-decorated\",\n \"markdownDescription\": \"Denies the is_decorated command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-enabled\",\n \"markdownDescription\": \"Denies the is_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_focused command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-focused\",\n \"markdownDescription\": \"Denies the is_focused command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-fullscreen\",\n \"markdownDescription\": \"Denies the is_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-maximizable\",\n \"markdownDescription\": \"Denies the is_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_maximized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-maximized\",\n \"markdownDescription\": \"Denies the is_maximized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-minimizable\",\n \"markdownDescription\": \"Denies the is_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_minimized command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-minimized\",\n \"markdownDescription\": \"Denies the is_minimized command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-resizable\",\n \"markdownDescription\": \"Denies the is_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the is_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-is-visible\",\n \"markdownDescription\": \"Denies the is_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-maximize\",\n \"markdownDescription\": \"Denies the maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the minimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-minimize\",\n \"markdownDescription\": \"Denies the minimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the monitor_from_point command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-monitor-from-point\",\n \"markdownDescription\": \"Denies the monitor_from_point command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the outer_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-outer-position\",\n \"markdownDescription\": \"Denies the outer_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the outer_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-outer-size\",\n \"markdownDescription\": \"Denies the outer_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the primary_monitor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-primary-monitor\",\n \"markdownDescription\": \"Denies the primary_monitor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the request_user_attention command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-request-user-attention\",\n \"markdownDescription\": \"Denies the request_user_attention command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the scale_factor command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-scale-factor\",\n \"markdownDescription\": \"Denies the scale_factor command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_always_on_bottom command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-always-on-bottom\",\n \"markdownDescription\": \"Denies the set_always_on_bottom command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_always_on_top command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-always-on-top\",\n \"markdownDescription\": \"Denies the set_always_on_top command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_background_color command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-background-color\",\n \"markdownDescription\": \"Denies the set_background_color command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_badge_count command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-badge-count\",\n \"markdownDescription\": \"Denies the set_badge_count command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_badge_label command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-badge-label\",\n \"markdownDescription\": \"Denies the set_badge_label command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_closable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-closable\",\n \"markdownDescription\": \"Denies the set_closable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_content_protected command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-content-protected\",\n \"markdownDescription\": \"Denies the set_content_protected command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_grab command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-grab\",\n \"markdownDescription\": \"Denies the set_cursor_grab command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-icon\",\n \"markdownDescription\": \"Denies the set_cursor_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-position\",\n \"markdownDescription\": \"Denies the set_cursor_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_cursor_visible command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-cursor-visible\",\n \"markdownDescription\": \"Denies the set_cursor_visible command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_decorations command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-decorations\",\n \"markdownDescription\": \"Denies the set_decorations command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_effects command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-effects\",\n \"markdownDescription\": \"Denies the set_effects command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_enabled command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-enabled\",\n \"markdownDescription\": \"Denies the set_enabled command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_focus command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-focus\",\n \"markdownDescription\": \"Denies the set_focus command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_focusable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-focusable\",\n \"markdownDescription\": \"Denies the set_focusable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-fullscreen\",\n \"markdownDescription\": \"Denies the set_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-icon\",\n \"markdownDescription\": \"Denies the set_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_ignore_cursor_events command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-ignore-cursor-events\",\n \"markdownDescription\": \"Denies the set_ignore_cursor_events command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_max_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-max-size\",\n \"markdownDescription\": \"Denies the set_max_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_maximizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-maximizable\",\n \"markdownDescription\": \"Denies the set_maximizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_min_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-min-size\",\n \"markdownDescription\": \"Denies the set_min_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_minimizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-minimizable\",\n \"markdownDescription\": \"Denies the set_minimizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_overlay_icon command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-overlay-icon\",\n \"markdownDescription\": \"Denies the set_overlay_icon command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_position command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-position\",\n \"markdownDescription\": \"Denies the set_position command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_progress_bar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-progress-bar\",\n \"markdownDescription\": \"Denies the set_progress_bar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_resizable command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-resizable\",\n \"markdownDescription\": \"Denies the set_resizable command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_shadow command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-shadow\",\n \"markdownDescription\": \"Denies the set_shadow command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_simple_fullscreen command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-simple-fullscreen\",\n \"markdownDescription\": \"Denies the set_simple_fullscreen command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_size command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-size\",\n \"markdownDescription\": \"Denies the set_size command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_size_constraints command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-size-constraints\",\n \"markdownDescription\": \"Denies the set_size_constraints command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_skip_taskbar command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-skip-taskbar\",\n \"markdownDescription\": \"Denies the set_skip_taskbar command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-theme\",\n \"markdownDescription\": \"Denies the set_theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-title\",\n \"markdownDescription\": \"Denies the set_title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_title_bar_style command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-title-bar-style\",\n \"markdownDescription\": \"Denies the set_title_bar_style command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the set_visible_on_all_workspaces command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-set-visible-on-all-workspaces\",\n \"markdownDescription\": \"Denies the set_visible_on_all_workspaces command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the show command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-show\",\n \"markdownDescription\": \"Denies the show command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the start_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-start-dragging\",\n \"markdownDescription\": \"Denies the start_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the start_resize_dragging command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-start-resize-dragging\",\n \"markdownDescription\": \"Denies the start_resize_dragging command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the theme command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-theme\",\n \"markdownDescription\": \"Denies the theme command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the title command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-title\",\n \"markdownDescription\": \"Denies the title command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the toggle_maximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-toggle-maximize\",\n \"markdownDescription\": \"Denies the toggle_maximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the unmaximize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-unmaximize\",\n \"markdownDescription\": \"Denies the unmaximize command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the unminimize command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"core:window:deny-unminimize\",\n \"markdownDescription\": \"Denies the unminimize command without any pre-configured scope.\"\n },\n {\n \"description\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\",\n \"type\": \"string\",\n \"const\": \"shell:default\",\n \"markdownDescription\": \"This permission set configures which\\nshell functionality is exposed by default.\\n\\n#### Granted Permissions\\n\\nIt allows to use the `open` functionality with a reasonable\\nscope pre-configured. It will allow opening `http(s)://`,\\n`tel:` and `mailto:` links.\\n\\n#### This default permission set includes:\\n\\n- `allow-open`\"\n },\n {\n \"description\": \"Enables the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-execute\",\n \"markdownDescription\": \"Enables the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-kill\",\n \"markdownDescription\": \"Enables the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-open\",\n \"markdownDescription\": \"Enables the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-spawn\",\n \"markdownDescription\": \"Enables the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:allow-stdin-write\",\n \"markdownDescription\": \"Enables the stdin_write command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the execute command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-execute\",\n \"markdownDescription\": \"Denies the execute command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the kill command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-kill\",\n \"markdownDescription\": \"Denies the kill command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the open command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-open\",\n \"markdownDescription\": \"Denies the open command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the spawn command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-spawn\",\n \"markdownDescription\": \"Denies the spawn command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the stdin_write command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"shell:deny-stdin-write\",\n \"markdownDescription\": \"Denies the stdin_write command without any pre-configured scope.\"\n },\n {\n \"description\": \"This permission set configures what kind of\\noperations are available from the window state plugin.\\n\\n#### Granted Permissions\\n\\nAll operations are enabled by default.\\n\\n\\n#### This default permission set includes:\\n\\n- `allow-filename`\\n- `allow-restore-state`\\n- `allow-save-window-state`\",\n \"type\": \"string\",\n \"const\": \"window-state:default\",\n \"markdownDescription\": \"This permission set configures what kind of\\noperations are available from the window state plugin.\\n\\n#### Granted Permissions\\n\\nAll operations are enabled by default.\\n\\n\\n#### This default permission set includes:\\n\\n- `allow-filename`\\n- `allow-restore-state`\\n- `allow-save-window-state`\"\n },\n {\n \"description\": \"Enables the filename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:allow-filename\",\n \"markdownDescription\": \"Enables the filename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the restore_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:allow-restore-state\",\n \"markdownDescription\": \"Enables the restore_state command without any pre-configured scope.\"\n },\n {\n \"description\": \"Enables the save_window_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:allow-save-window-state\",\n \"markdownDescription\": \"Enables the save_window_state command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the filename command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:deny-filename\",\n \"markdownDescription\": \"Denies the filename command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the restore_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:deny-restore-state\",\n \"markdownDescription\": \"Denies the restore_state command without any pre-configured scope.\"\n },\n {\n \"description\": \"Denies the save_window_state command without any pre-configured scope.\",\n \"type\": \"string\",\n \"const\": \"window-state:deny-save-window-state\",\n \"markdownDescription\": \"Denies the save_window_state command without any pre-configured scope.\"\n }\n ]\n },\n \"Value\": {\n \"description\": \"All supported ACL values.\",\n \"anyOf\": [\n {\n \"description\": \"Represents a null JSON value.\",\n \"type\": \"null\"\n },\n {\n \"description\": \"Represents a [`bool`].\",\n \"type\": \"boolean\"\n },\n {\n \"description\": \"Represents a valid ACL [`Number`].\",\n \"allOf\": [\n {\n \"$ref\": \"#/definitions/Number\"\n }\n ]\n },\n {\n \"description\": \"Represents a [`String`].\",\n \"type\": \"string\"\n },\n {\n \"description\": \"Represents a list of other [`Value`]s.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/Value\"\n }\n },\n {\n \"description\": \"Represents a map of [`String`] keys to [`Value`]s.\",\n \"type\": \"object\",\n \"additionalProperties\": {\n \"$ref\": \"#/definitions/Value\"\n }\n }\n ]\n },\n \"Number\": {\n \"description\": \"A valid ACL number.\",\n \"anyOf\": [\n {\n \"description\": \"Represents an [`i64`].\",\n \"type\": \"integer\",\n \"format\": \"int64\"\n },\n {\n \"description\": \"Represents a [`f64`].\",\n \"type\": \"number\",\n \"format\": \"double\"\n }\n ]\n },\n \"Target\": {\n \"description\": \"Platform target.\",\n \"oneOf\": [\n {\n \"description\": \"MacOS.\",\n \"type\": \"string\",\n \"enum\": [\n \"macOS\"\n ]\n },\n {\n \"description\": \"Windows.\",\n \"type\": \"string\",\n \"enum\": [\n \"windows\"\n ]\n },\n {\n \"description\": \"Linux.\",\n \"type\": \"string\",\n \"enum\": [\n \"linux\"\n ]\n },\n {\n \"description\": \"Android.\",\n \"type\": \"string\",\n \"enum\": [\n \"android\"\n ]\n },\n {\n \"description\": \"iOS.\",\n \"type\": \"string\",\n \"enum\": [\n \"iOS\"\n ]\n }\n ]\n },\n \"ShellScopeEntryAllowedArg\": {\n \"description\": \"A command argument allowed to be executed by the webview API.\",\n \"anyOf\": [\n {\n \"description\": \"A non-configurable argument that is passed to the command in the order it was specified.\",\n \"type\": \"string\"\n },\n {\n \"description\": \"A variable that is set while calling the command from the webview API.\",\n \"type\": \"object\",\n \"required\": [\n \"validator\"\n ],\n \"properties\": {\n \"raw\": {\n \"description\": \"Marks the validator as a raw regex, meaning the plugin should not make any modification at runtime.\\n\\nThis means the regex will not match on the entire string by default, which might be exploited if your regex allow unexpected input to be considered valid. When using this option, make sure your regex is correct.\",\n \"default\": false,\n \"type\": \"boolean\"\n },\n \"validator\": {\n \"description\": \"[regex] validator to require passed values to conform to an expected input.\\n\\nThis will require the argument value passed to this variable to match the `validator` regex before it will be executed.\\n\\nThe regex string is by default surrounded by `^...$` to match the full string. For example the `https?://\\\\w+` regex would be registered as `^https?://\\\\w+$`.\\n\\n[regex]: \",\n \"type\": \"string\"\n }\n },\n \"additionalProperties\": false\n }\n ]\n },\n \"ShellScopeEntryAllowedArgs\": {\n \"description\": \"A set of command arguments allowed to be executed by the webview API.\\n\\nA value of `true` will allow any arguments to be passed to the command. `false` will disable all arguments. A list of [`ShellScopeEntryAllowedArg`] will set those arguments as the only valid arguments to be passed to the attached command configuration.\",\n \"anyOf\": [\n {\n \"description\": \"Use a simple boolean to allow all or disable all arguments to this command configuration.\",\n \"type\": \"boolean\"\n },\n {\n \"description\": \"A specific set of [`ShellScopeEntryAllowedArg`] that are valid to call for the command configuration.\",\n \"type\": \"array\",\n \"items\": {\n \"$ref\": \"#/definitions/ShellScopeEntryAllowedArg\"\n }\n }\n ]\n }\n }\n}" }, { "path": "desktop/src-tauri/icons/android/mipmap-anydpi-v26/ic_launcher.xml", "content": "\n\n \n \n" }, { "path": "desktop/src-tauri/icons/android/values/ic_launcher_background.xml", "content": "\n\n #fff\n" }, { "path": "desktop/src-tauri/src/main.rs", "content": "// Prevents additional console window on Windows in release\n#![cfg_attr(not(debug_assertions), windows_subsystem = \"windows\")]\n\nuse directories::ProjectDirs;\nuse serde::{Deserialize, Serialize};\nuse std::fs;\nuse std::path::PathBuf;\nuse std::process::Command;\nuse std::sync::{Mutex, RwLock};\nuse std::io::Write as IoWrite;\nuse std::time::SystemTime;\n#[cfg(target_os = \"macos\")]\nuse std::time::Duration;\nuse tauri::image::Image;\nuse tauri::menu::{\n CheckMenuItem, Menu, MenuBuilder, MenuItem, PredefinedMenuItem, SubmenuBuilder, HELP_SUBMENU_ID,\n};\nuse tauri::tray::{TrayIconBuilder, TrayIconEvent};\n#[cfg(target_os = \"macos\")]\nuse tauri::WebviewWindow;\nuse tauri::Wry;\nuse tauri::{\n webview::PageLoadPayload, AppHandle, Manager, Webview, WebviewUrl, WebviewWindowBuilder,\n};\n#[cfg(target_os = \"macos\")]\nuse tokio::time::sleep;\nuse url::Url;\n#[cfg(target_os = \"macos\")]\nuse window_vibrancy::{apply_vibrancy, NSVisualEffectMaterial};\n\n// ============================================================================\n// Configuration\n// ============================================================================\n\nconst DEFAULT_SERVER_URL: &str = \"https://cloud.onyx.app\";\nconst CONFIG_FILE_NAME: &str = \"config.json\";\n#[cfg(target_os = \"macos\")]\nconst TITLEBAR_SCRIPT: &str = include_str!(\"../../src/titlebar.js\");\nconst TRAY_ID: &str = \"onyx-tray\";\nconst TRAY_ICON_BYTES: &[u8] = include_bytes!(\"../icons/tray-icon.png\");\nconst TRAY_MENU_OPEN_APP_ID: &str = \"tray_open_app\";\nconst TRAY_MENU_OPEN_CHAT_ID: &str = \"tray_open_chat\";\nconst TRAY_MENU_SHOW_IN_BAR_ID: &str = \"tray_show_in_menu_bar\";\nconst TRAY_MENU_QUIT_ID: &str = \"tray_quit\";\nconst MENU_SHOW_MENU_BAR_ID: &str = \"show_menu_bar\";\nconst MENU_HIDE_DECORATIONS_ID: &str = \"hide_window_decorations\";\nconst CHAT_LINK_INTERCEPT_SCRIPT: &str = r##\"\n(() => {\n if (window.__ONYX_CHAT_LINK_INTERCEPT_INSTALLED__) {\n return;\n }\n\n window.__ONYX_CHAT_LINK_INTERCEPT_INSTALLED__ = true;\n\n function isChatSessionPage() {\n try {\n const currentUrl = new URL(window.location.href);\n return (\n currentUrl.pathname.startsWith(\"/app\") &&\n currentUrl.searchParams.has(\"chatId\")\n );\n } catch {\n return false;\n }\n }\n\n function getAllowedNavigationUrl(rawUrl) {\n try {\n const parsed = new URL(String(rawUrl), window.location.href);\n const scheme = parsed.protocol.toLowerCase();\n if (![\"http:\", \"https:\", \"mailto:\", \"tel:\"].includes(scheme)) {\n return null;\n }\n return parsed;\n } catch {\n return null;\n }\n }\n\n async function openWithTauri(url) {\n try {\n const invoke =\n window.__TAURI__?.core?.invoke || window.__TAURI_INTERNALS__?.invoke;\n if (typeof invoke !== \"function\") {\n return false;\n }\n\n await invoke(\"open_in_browser\", { url });\n return true;\n } catch {\n return false;\n }\n }\n\n function handleChatNavigation(rawUrl) {\n const parsedUrl = getAllowedNavigationUrl(rawUrl);\n if (!parsedUrl) {\n return false;\n }\n\n const safeUrl = parsedUrl.toString();\n const scheme = parsedUrl.protocol.toLowerCase();\n if (scheme === \"mailto:\" || scheme === \"tel:\") {\n void openWithTauri(safeUrl).then((opened) => {\n if (!opened) {\n window.location.assign(safeUrl);\n }\n });\n return true;\n }\n\n window.location.assign(safeUrl);\n return true;\n }\n\n document.addEventListener(\n \"click\",\n (event) => {\n if (!isChatSessionPage() || event.defaultPrevented) {\n return;\n }\n\n const element = event.target;\n if (!(element instanceof Element)) {\n return;\n }\n\n const anchor = element.closest(\"a\");\n if (!(anchor instanceof HTMLAnchorElement)) {\n return;\n }\n\n const target = (anchor.getAttribute(\"target\") || \"\").toLowerCase();\n if (target !== \"_blank\") {\n return;\n }\n\n const href = anchor.getAttribute(\"href\");\n if (!href || href.startsWith(\"#\")) {\n return;\n }\n\n if (!handleChatNavigation(href)) {\n return;\n }\n\n event.preventDefault();\n event.stopPropagation();\n },\n true\n );\n\n const nativeWindowOpen = window.open;\n window.open = function(url, target, features) {\n const resolvedTarget = typeof target === \"string\" ? target.toLowerCase() : \"\";\n const shouldNavigateInPlace = resolvedTarget === \"\" || resolvedTarget === \"_blank\";\n\n if (\n isChatSessionPage() &&\n shouldNavigateInPlace &&\n url != null &&\n String(url).length > 0\n ) {\n if (!handleChatNavigation(url)) {\n return null;\n }\n return null;\n }\n\n if (typeof nativeWindowOpen === \"function\") {\n return nativeWindowOpen.call(window, url, target, features);\n }\n return null;\n };\n})();\n\"##;\n\n#[cfg(not(target_os = \"macos\"))]\nconst MENU_KEY_HANDLER_SCRIPT: &str = r#\"\n(() => {\n if (window.__ONYX_MENU_KEY_HANDLER__) return;\n window.__ONYX_MENU_KEY_HANDLER__ = true;\n\n let altHeld = false;\n\n function invoke(cmd) {\n const fn_ =\n window.__TAURI__?.core?.invoke || window.__TAURI_INTERNALS__?.invoke;\n if (typeof fn_ === 'function') fn_(cmd);\n }\n\n function releaseAltAndHideMenu() {\n if (!altHeld) {\n return;\n }\n altHeld = false;\n invoke('hide_menu_bar_temporary');\n }\n\n document.addEventListener('keydown', (e) => {\n if (e.key === 'Alt') {\n if (!altHeld) {\n altHeld = true;\n invoke('show_menu_bar_temporarily');\n }\n return;\n }\n if (e.altKey && e.key === 'F1') {\n e.preventDefault();\n e.stopPropagation();\n altHeld = false;\n invoke('toggle_menu_bar');\n return;\n }\n }, true);\n\n document.addEventListener('keyup', (e) => {\n if (e.key === 'Alt' && altHeld) {\n releaseAltAndHideMenu();\n }\n }, true);\n\n window.addEventListener('blur', () => {\n releaseAltAndHideMenu();\n });\n\n document.addEventListener('visibilitychange', () => {\n if (document.hidden) {\n releaseAltAndHideMenu();\n }\n });\n})();\n\"#;\n\nconst CONSOLE_CAPTURE_SCRIPT: &str = r#\"\n(() => {\n if (window.__ONYX_CONSOLE_CAPTURE__) return;\n window.__ONYX_CONSOLE_CAPTURE__ = true;\n\n const levels = ['log', 'warn', 'error', 'info', 'debug'];\n const originals = {};\n\n levels.forEach(level => {\n originals[level] = console[level];\n console[level] = function(...args) {\n originals[level].apply(console, args);\n try {\n const invoke =\n window.__TAURI__?.core?.invoke || window.__TAURI_INTERNALS__?.invoke;\n if (typeof invoke === 'function') {\n const message = args.map(a => {\n try { return typeof a === 'string' ? a : JSON.stringify(a); }\n catch { return String(a); }\n }).join(' ');\n invoke('log_from_frontend', { level, message });\n }\n } catch {}\n };\n });\n\n window.addEventListener('error', (event) => {\n try {\n const invoke =\n window.__TAURI__?.core?.invoke || window.__TAURI_INTERNALS__?.invoke;\n if (typeof invoke === 'function') {\n invoke('log_from_frontend', {\n level: 'error',\n message: `[uncaught] ${event.message} at ${event.filename}:${event.lineno}:${event.colno}`\n });\n }\n } catch {}\n });\n\n window.addEventListener('unhandledrejection', (event) => {\n try {\n const invoke =\n window.__TAURI__?.core?.invoke || window.__TAURI_INTERNALS__?.invoke;\n if (typeof invoke === 'function') {\n invoke('log_from_frontend', {\n level: 'error',\n message: `[unhandled rejection] ${event.reason}`\n });\n }\n } catch {}\n });\n})();\n\"#;\n\nconst MENU_TOGGLE_DEVTOOLS_ID: &str = \"toggle_devtools\";\nconst MENU_OPEN_DEBUG_LOG_ID: &str = \"open_debug_log\";\n\n#[derive(Debug, Clone, Serialize, Deserialize)]\npub struct AppConfig {\n pub server_url: String,\n\n #[serde(default = \"default_window_title\")]\n pub window_title: String,\n\n #[serde(default = \"default_show_menu_bar\")]\n pub show_menu_bar: bool,\n\n #[serde(default)]\n pub hide_window_decorations: bool,\n}\n\nfn default_window_title() -> String {\n \"Onyx\".to_string()\n}\n\nfn default_show_menu_bar() -> bool {\n true\n}\n\nimpl Default for AppConfig {\n fn default() -> Self {\n Self {\n server_url: DEFAULT_SERVER_URL.to_string(),\n window_title: default_window_title(),\n show_menu_bar: true,\n hide_window_decorations: false,\n }\n }\n}\n\n/// Get the config directory path\nfn get_config_dir() -> Option {\n ProjectDirs::from(\"app\", \"onyx\", \"onyx-desktop\").map(|dirs| dirs.config_dir().to_path_buf())\n}\n\n/// Get the full config file path\nfn get_config_path() -> Option {\n get_config_dir().map(|dir| dir.join(CONFIG_FILE_NAME))\n}\n\n/// Load config from file, or create default if it doesn't exist\nfn load_config() -> (AppConfig, bool) {\n let config_path = match get_config_path() {\n Some(path) => path,\n None => {\n return (AppConfig::default(), false);\n }\n };\n\n if !config_path.exists() {\n return (AppConfig::default(), false);\n }\n\n match fs::read_to_string(&config_path) {\n Ok(contents) => match serde_json::from_str(&contents) {\n Ok(config) => (config, true),\n Err(_) => (AppConfig::default(), false),\n },\n Err(_) => (AppConfig::default(), false),\n }\n}\n\n/// Save config to file\nfn save_config(config: &AppConfig) -> Result<(), String> {\n let config_dir = get_config_dir().ok_or(\"Could not determine config directory\")?;\n let config_path = config_dir.join(CONFIG_FILE_NAME);\n\n // Ensure config directory exists\n fs::create_dir_all(&config_dir).map_err(|e| format!(\"Failed to create config dir: {}\", e))?;\n\n let json = serde_json::to_string_pretty(config)\n .map_err(|e| format!(\"Failed to serialize config: {}\", e))?;\n\n fs::write(&config_path, json).map_err(|e| format!(\"Failed to write config: {}\", e))?;\n\n Ok(())\n}\n\n// ============================================================================\n// Debug Mode\n// ============================================================================\n\nfn is_debug_mode() -> bool {\n std::env::args().any(|arg| arg == \"--debug\") || std::env::var(\"ONYX_DEBUG\").is_ok()\n}\n\nfn get_debug_log_path() -> Option {\n get_config_dir().map(|dir| dir.join(\"frontend_debug.log\"))\n}\n\nfn init_debug_log_file() -> Option {\n let log_path = get_debug_log_path()?;\n if let Some(parent) = log_path.parent() {\n let _ = fs::create_dir_all(parent);\n }\n fs::OpenOptions::new()\n .create(true)\n .append(true)\n .open(&log_path)\n .ok()\n}\n\nfn format_utc_timestamp() -> String {\n let now = SystemTime::now()\n .duration_since(SystemTime::UNIX_EPOCH)\n .unwrap_or_default();\n let total_secs = now.as_secs();\n let millis = now.subsec_millis();\n\n let days = total_secs / 86400;\n let secs_of_day = total_secs % 86400;\n let hours = secs_of_day / 3600;\n let mins = (secs_of_day % 3600) / 60;\n let secs = secs_of_day % 60;\n\n // Days since Unix epoch -> Y/M/D via civil calendar arithmetic\n let z = days as i64 + 719468;\n let era = z / 146097;\n let doe = z - era * 146097;\n let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;\n let y = yoe + era * 400;\n let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);\n let mp = (5 * doy + 2) / 153;\n let d = doy - (153 * mp + 2) / 5 + 1;\n let m = if mp < 10 { mp + 3 } else { mp - 9 };\n let y = if m <= 2 { y + 1 } else { y };\n\n format!(\n \"{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:03}Z\",\n y, m, d, hours, mins, secs, millis\n )\n}\n\nfn inject_console_capture(webview: &Webview) {\n let _ = webview.eval(CONSOLE_CAPTURE_SCRIPT);\n}\n\nfn maybe_open_devtools(app: &AppHandle, window: &tauri::WebviewWindow) {\n #[cfg(any(debug_assertions, feature = \"devtools\"))]\n {\n let state = app.state::();\n if state.debug_mode {\n window.open_devtools();\n }\n }\n #[cfg(not(any(debug_assertions, feature = \"devtools\")))]\n {\n let _ = (app, window);\n }\n}\n\n// Global config state\nstruct ConfigState {\n config: RwLock,\n config_initialized: RwLock,\n app_base_url: RwLock>,\n menu_temporarily_visible: RwLock,\n debug_mode: bool,\n debug_log_file: Mutex>,\n}\n\nfn focus_main_window(app: &AppHandle) {\n if let Some(window) = app.get_webview_window(\"main\") {\n let _ = window.unminimize();\n let _ = window.show();\n let _ = window.set_focus();\n } else {\n trigger_new_window(app);\n }\n}\n\nfn trigger_new_chat(app: &AppHandle) {\n let state = app.state::();\n let server_url = state.config.read().unwrap().server_url.clone();\n\n if let Some(window) = app.get_webview_window(\"main\") {\n let url = format!(\"{}/chat\", server_url);\n let _ = window.eval(&format!(\"window.location.href = '{}'\", url));\n }\n}\n\nfn trigger_new_window(app: &AppHandle) {\n let state = app.state::();\n let server_url = state.config.read().unwrap().server_url.clone();\n let handle = app.clone();\n\n tauri::async_runtime::spawn(async move {\n let window_label = format!(\"onyx-{}\", uuid::Uuid::new_v4());\n let builder = WebviewWindowBuilder::new(\n &handle,\n &window_label,\n WebviewUrl::External(server_url.parse().unwrap()),\n )\n .title(\"Onyx\")\n .inner_size(1200.0, 800.0)\n .min_inner_size(800.0, 600.0)\n .transparent(true);\n\n #[cfg(target_os = \"macos\")]\n let builder = builder\n .title_bar_style(tauri::TitleBarStyle::Overlay)\n .hidden_title(true);\n\n #[cfg(target_os = \"linux\")]\n let builder = builder.background_color(tauri::window::Color(0x1a, 0x1a, 0x2e, 0xff));\n\n if let Ok(window) = builder.build() {\n #[cfg(target_os = \"macos\")]\n {\n let _ = apply_vibrancy(&window, NSVisualEffectMaterial::Sidebar, None, None);\n inject_titlebar(window.clone());\n }\n\n apply_settings_to_window(&handle, &window);\n maybe_open_devtools(&handle, &window);\n let _ = window.set_focus();\n }\n });\n}\n\nfn open_docs() {\n let _ = open_in_default_browser(\"https://docs.onyx.app\");\n}\n\nfn open_settings(app: &AppHandle) {\n // Navigate main window to the settings page (index.html) with settings flag\n let state = app.state::();\n let settings_url = state\n .app_base_url\n .read()\n .unwrap()\n .as_ref()\n .cloned()\n .and_then(|mut url| {\n url.set_query(None);\n url.set_fragment(Some(\"settings\"));\n url.set_path(\"/\");\n Some(url)\n })\n .or_else(|| Url::parse(\"tauri://localhost/#settings\").ok());\n\n if let Some(window) = app.get_webview_window(\"main\") {\n if let Some(url) = settings_url {\n let _ = window.navigate(url);\n }\n }\n}\n\nfn same_origin(left: &Url, right: &Url) -> bool {\n left.scheme() == right.scheme()\n && left.host_str() == right.host_str()\n && left.port_or_known_default() == right.port_or_known_default()\n}\n\nfn is_chat_session_url(url: &Url) -> bool {\n url.path().starts_with(\"/app\") && url.query_pairs().any(|(key, _)| key == \"chatId\")\n}\n\nfn should_open_in_external_browser(current_url: &Url, destination_url: &Url) -> bool {\n if !is_chat_session_url(current_url) {\n return false;\n }\n\n match destination_url.scheme() {\n \"mailto\" | \"tel\" => true,\n \"http\" | \"https\" => !same_origin(current_url, destination_url),\n _ => false,\n }\n}\n\nfn open_in_default_browser(url: &str) -> bool {\n #[cfg(target_os = \"macos\")]\n {\n return Command::new(\"open\").arg(url).status().is_ok();\n }\n #[cfg(target_os = \"linux\")]\n {\n return Command::new(\"xdg-open\").arg(url).status().is_ok();\n }\n #[cfg(target_os = \"windows\")]\n {\n return Command::new(\"rundll32\")\n .arg(\"url.dll,FileProtocolHandler\")\n .arg(url)\n .status()\n .is_ok();\n }\n #[allow(unreachable_code)]\n false\n}\n\n#[tauri::command]\nfn open_in_browser(url: String) -> Result<(), String> {\n let parsed_url = Url::parse(&url).map_err(|_| \"Invalid URL\".to_string())?;\n match parsed_url.scheme() {\n \"http\" | \"https\" | \"mailto\" | \"tel\" => {}\n _ => return Err(\"Unsupported URL scheme\".to_string()),\n }\n\n if open_in_default_browser(parsed_url.as_str()) {\n Ok(())\n } else {\n Err(\"Failed to open URL in default browser\".to_string())\n }\n}\n\nfn inject_chat_link_intercept(webview: &Webview) {\n let _ = webview.eval(CHAT_LINK_INTERCEPT_SCRIPT);\n}\n\nfn handle_toggle_devtools(app: &AppHandle) {\n #[cfg(any(debug_assertions, feature = \"devtools\"))]\n {\n let windows: Vec<_> = app.webview_windows().into_values().collect();\n let any_open = windows.iter().any(|w| w.is_devtools_open());\n for window in &windows {\n if any_open {\n window.close_devtools();\n } else {\n window.open_devtools();\n }\n }\n }\n #[cfg(not(any(debug_assertions, feature = \"devtools\")))]\n {\n let _ = app;\n }\n}\n\nfn handle_open_debug_log() {\n let log_path = match get_debug_log_path() {\n Some(p) => p,\n None => return,\n };\n\n if !log_path.exists() {\n eprintln!(\"[ONYX DEBUG] Log file does not exist yet: {:?}\", log_path);\n return;\n }\n\n let url_path = log_path.to_string_lossy().replace('\\\\', \"/\");\n let _ = open_in_default_browser(&format!(\n \"file:///{}\",\n url_path.trim_start_matches('/')\n ));\n}\n\n// ============================================================================\n// Tauri Commands\n// ============================================================================\n\n#[tauri::command]\nfn log_from_frontend(level: String, message: String, state: tauri::State) {\n if !state.debug_mode {\n return;\n }\n let timestamp = format_utc_timestamp();\n let log_line = format!(\"[{}] [{}] {}\", timestamp, level.to_uppercase(), message);\n\n eprintln!(\"{}\", log_line);\n\n if let Ok(mut guard) = state.debug_log_file.lock() {\n if let Some(ref mut file) = *guard {\n let _ = writeln!(file, \"{}\", log_line);\n let _ = file.flush();\n }\n }\n}\n\n/// Get the current server URL\n#[tauri::command]\nfn get_server_url(state: tauri::State) -> String {\n state.config.read().unwrap().server_url.clone()\n}\n\n#[derive(Serialize)]\nstruct BootstrapState {\n server_url: String,\n config_exists: bool,\n}\n\n/// Get the server URL plus whether a config file exists\n#[tauri::command]\nfn get_bootstrap_state(state: tauri::State) -> BootstrapState {\n let server_url = state.config.read().unwrap().server_url.clone();\n let config_initialized = *state.config_initialized.read().unwrap();\n let config_exists =\n config_initialized && get_config_path().map(|path| path.exists()).unwrap_or(false);\n\n BootstrapState {\n server_url,\n config_exists,\n }\n}\n\n/// Set a new server URL and save to config\n#[tauri::command]\nfn set_server_url(state: tauri::State, url: String) -> Result {\n // Validate URL\n if !url.starts_with(\"http://\") && !url.starts_with(\"https://\") {\n return Err(\"URL must start with http:// or https://\".to_string());\n }\n\n let mut config = state.config.write().unwrap();\n config.server_url = url.trim_end_matches('/').to_string();\n save_config(&config)?;\n *state.config_initialized.write().unwrap() = true;\n\n Ok(config.server_url.clone())\n}\n\n/// Get the config file path (so users know where to edit)\n#[tauri::command]\nfn get_config_path_cmd() -> Result {\n get_config_path()\n .map(|p| p.to_string_lossy().to_string())\n .ok_or_else(|| \"Could not determine config path\".to_string())\n}\n\n/// Open the config file in the default editor\n#[tauri::command]\nfn open_config_file() -> Result<(), String> {\n let config_path = get_config_path().ok_or(\"Could not determine config path\")?;\n\n // Ensure config exists\n if !config_path.exists() {\n save_config(&AppConfig::default())?;\n }\n\n #[cfg(target_os = \"macos\")]\n {\n std::process::Command::new(\"open\")\n .arg(\"-t\")\n .arg(&config_path)\n .spawn()\n .map_err(|e| format!(\"Failed to open config: {}\", e))?;\n }\n\n #[cfg(target_os = \"linux\")]\n {\n std::process::Command::new(\"xdg-open\")\n .arg(&config_path)\n .spawn()\n .map_err(|e| format!(\"Failed to open config: {}\", e))?;\n }\n\n #[cfg(target_os = \"windows\")]\n {\n std::process::Command::new(\"notepad\")\n .arg(&config_path)\n .spawn()\n .map_err(|e| format!(\"Failed to open config: {}\", e))?;\n }\n\n Ok(())\n}\n\n/// Open the config directory in file manager\n#[tauri::command]\nfn open_config_directory() -> Result<(), String> {\n let config_dir = get_config_dir().ok_or(\"Could not determine config directory\")?;\n\n // Ensure directory exists\n fs::create_dir_all(&config_dir).map_err(|e| format!(\"Failed to create config dir: {}\", e))?;\n\n #[cfg(target_os = \"macos\")]\n {\n std::process::Command::new(\"open\")\n .arg(&config_dir)\n .spawn()\n .map_err(|e| format!(\"Failed to open directory: {}\", e))?;\n }\n\n #[cfg(target_os = \"linux\")]\n {\n std::process::Command::new(\"xdg-open\")\n .arg(&config_dir)\n .spawn()\n .map_err(|e| format!(\"Failed to open directory: {}\", e))?;\n }\n\n #[cfg(target_os = \"windows\")]\n {\n std::process::Command::new(\"explorer\")\n .arg(&config_dir)\n .spawn()\n .map_err(|e| format!(\"Failed to open directory: {}\", e))?;\n }\n\n Ok(())\n}\n\n/// Navigate to a specific path on the configured server\n#[tauri::command]\nfn navigate_to(window: tauri::WebviewWindow, state: tauri::State, path: &str) {\n let base_url = state.config.read().unwrap().server_url.clone();\n let url = format!(\"{}{}\", base_url, path);\n let _ = window.eval(&format!(\"window.location.href = '{}'\", url));\n}\n\n/// Reload the current page\n#[tauri::command]\nfn reload_page(window: tauri::WebviewWindow) {\n let _ = window.eval(\"window.location.reload()\");\n}\n\n/// Go back in history\n#[tauri::command]\nfn go_back(window: tauri::WebviewWindow) {\n let _ = window.eval(\"window.history.back()\");\n}\n\n/// Go forward in history\n#[tauri::command]\nfn go_forward(window: tauri::WebviewWindow) {\n let _ = window.eval(\"window.history.forward()\");\n}\n\n/// Open a new window\n#[tauri::command]\nasync fn new_window(app: AppHandle, state: tauri::State<'_, ConfigState>) -> Result<(), String> {\n let server_url = state.config.read().unwrap().server_url.clone();\n let window_label = format!(\"onyx-{}\", uuid::Uuid::new_v4());\n\n let builder = WebviewWindowBuilder::new(\n &app,\n &window_label,\n WebviewUrl::External(\n server_url\n .parse()\n .map_err(|e| format!(\"Invalid URL: {}\", e))?,\n ),\n )\n .title(\"Onyx\")\n .inner_size(1200.0, 800.0)\n .min_inner_size(800.0, 600.0)\n .transparent(true);\n\n #[cfg(target_os = \"macos\")]\n let builder = builder\n .title_bar_style(tauri::TitleBarStyle::Overlay)\n .hidden_title(true);\n\n #[cfg(target_os = \"linux\")]\n let builder = builder.background_color(tauri::window::Color(0x1a, 0x1a, 0x2e, 0xff));\n\n let window = builder.build().map_err(|e| e.to_string())?;\n\n #[cfg(target_os = \"macos\")]\n {\n let _ = apply_vibrancy(&window, NSVisualEffectMaterial::Sidebar, None, None);\n inject_titlebar(window.clone());\n }\n\n apply_settings_to_window(&app, &window);\n maybe_open_devtools(&app, &window);\n\n Ok(())\n}\n\n/// Reset config to defaults\n#[tauri::command]\nfn reset_config(state: tauri::State) -> Result<(), String> {\n let mut config = state.config.write().unwrap();\n *config = AppConfig::default();\n save_config(&config)?;\n *state.config_initialized.write().unwrap() = true;\n Ok(())\n}\n\n#[cfg(target_os = \"macos\")]\nfn inject_titlebar(window: WebviewWindow) {\n let script = TITLEBAR_SCRIPT.to_string();\n tauri::async_runtime::spawn(async move {\n // Keep trying for a few seconds to survive navigations and slow loads\n let delays = [0u64, 200, 600, 1200, 2000, 4000, 6000, 8000, 10000];\n for delay in delays {\n if delay > 0 {\n sleep(Duration::from_millis(delay)).await;\n }\n let _ = window.eval(&script);\n }\n });\n}\n\n/// Start dragging the window\n#[tauri::command]\nasync fn start_drag_window(window: tauri::Window) -> Result<(), String> {\n window.start_dragging().map_err(|e| e.to_string())\n}\n\n// ============================================================================\n// Window Settings\n// ============================================================================\n\nfn find_check_menu_item(\n app: &AppHandle,\n id: &str,\n) -> Option> {\n let menu = app.menu()?;\n for item in menu.items().ok()? {\n if let Some(submenu) = item.as_submenu() {\n for sub_item in submenu.items().ok()? {\n if let Some(check) = sub_item.as_check_menuitem() {\n if check.id().as_ref() == id {\n return Some(check.clone());\n }\n }\n }\n }\n }\n None\n}\n\nfn apply_settings_to_window(app: &AppHandle, window: &tauri::WebviewWindow) {\n if cfg!(target_os = \"macos\") {\n return;\n }\n let state = app.state::();\n let config = state.config.read().unwrap();\n let temp_visible = *state.menu_temporarily_visible.read().unwrap();\n if !config.show_menu_bar && !temp_visible {\n let _ = window.hide_menu();\n }\n if config.hide_window_decorations {\n let _ = window.set_decorations(false);\n }\n}\n\nfn handle_menu_bar_toggle(app: &AppHandle) {\n if cfg!(target_os = \"macos\") {\n return;\n }\n let state = app.state::();\n let show = {\n let mut config = state.config.write().unwrap();\n config.show_menu_bar = !config.show_menu_bar;\n let _ = save_config(&config);\n config.show_menu_bar\n };\n\n *state.menu_temporarily_visible.write().unwrap() = false;\n\n for (_, window) in app.webview_windows() {\n if show {\n let _ = window.show_menu();\n } else {\n let _ = window.hide_menu();\n }\n }\n}\n\nfn handle_decorations_toggle(app: &AppHandle) {\n if cfg!(target_os = \"macos\") {\n return;\n }\n let state = app.state::();\n let hide = {\n let mut config = state.config.write().unwrap();\n config.hide_window_decorations = !config.hide_window_decorations;\n let _ = save_config(&config);\n config.hide_window_decorations\n };\n\n for (_, window) in app.webview_windows() {\n let _ = window.set_decorations(!hide);\n }\n}\n\n#[tauri::command]\nfn toggle_menu_bar(app: AppHandle) {\n if cfg!(target_os = \"macos\") {\n return;\n }\n handle_menu_bar_toggle(&app);\n\n let state = app.state::();\n let checked = state.config.read().unwrap().show_menu_bar;\n if let Some(check) = find_check_menu_item(&app, MENU_SHOW_MENU_BAR_ID) {\n let _ = check.set_checked(checked);\n }\n}\n\n#[tauri::command]\nfn show_menu_bar_temporarily(app: AppHandle) {\n if cfg!(target_os = \"macos\") {\n return;\n }\n let state = app.state::();\n if state.config.read().unwrap().show_menu_bar {\n return;\n }\n\n let mut temp = state.menu_temporarily_visible.write().unwrap();\n if *temp {\n return;\n }\n *temp = true;\n drop(temp);\n\n for (_, window) in app.webview_windows() {\n let _ = window.show_menu();\n }\n}\n\n#[tauri::command]\nfn hide_menu_bar_temporary(app: AppHandle) {\n if cfg!(target_os = \"macos\") {\n return;\n }\n let state = app.state::();\n let mut temp = state.menu_temporarily_visible.write().unwrap();\n if !*temp {\n return;\n }\n *temp = false;\n drop(temp);\n\n if state.config.read().unwrap().show_menu_bar {\n return;\n }\n\n for (_, window) in app.webview_windows() {\n let _ = window.hide_menu();\n }\n}\n\n// ============================================================================\n// Menu Setup\n// ============================================================================\n\nfn setup_app_menu(app: &AppHandle) -> tauri::Result<()> {\n let menu = app.menu().unwrap_or(Menu::default(app)?);\n\n let new_chat_item = MenuItem::with_id(app, \"new_chat\", \"New Chat\", true, Some(\"CmdOrCtrl+N\"))?;\n let new_window_item = MenuItem::with_id(\n app,\n \"new_window\",\n \"New Window\",\n true,\n Some(\"CmdOrCtrl+Shift+N\"),\n )?;\n let settings_item = MenuItem::with_id(\n app,\n \"open_settings\",\n \"Settings...\",\n true,\n Some(\"CmdOrCtrl+Comma\"),\n )?;\n let docs_item = MenuItem::with_id(app, \"open_docs\", \"Onyx Documentation\", true, None::<&str>)?;\n\n if let Some(file_menu) = menu\n .items()?\n .into_iter()\n .filter_map(|item| item.as_submenu().cloned())\n .find(|submenu| submenu.text().ok().as_deref() == Some(\"File\"))\n {\n file_menu.insert_items(&[&new_chat_item, &new_window_item, &settings_item], 0)?;\n } else {\n let file_menu = SubmenuBuilder::new(app, \"File\")\n .items(&[\n &new_chat_item,\n &new_window_item,\n &settings_item,\n &PredefinedMenuItem::close_window(app, None)?,\n ])\n .build()?;\n menu.prepend(&file_menu)?;\n }\n\n #[cfg(not(target_os = \"macos\"))]\n {\n let config = app.state::();\n let config_guard = config.config.read().unwrap();\n\n let show_menu_bar_item = CheckMenuItem::with_id(\n app,\n MENU_SHOW_MENU_BAR_ID,\n \"Show Menu Bar\",\n true,\n config_guard.show_menu_bar,\n None::<&str>,\n )?;\n\n let hide_decorations_item = CheckMenuItem::with_id(\n app,\n MENU_HIDE_DECORATIONS_ID,\n \"Hide Window Decorations\",\n true,\n config_guard.hide_window_decorations,\n None::<&str>,\n )?;\n\n drop(config_guard);\n\n if let Some(window_menu) = menu\n .items()?\n .into_iter()\n .filter_map(|item| item.as_submenu().cloned())\n .find(|submenu| submenu.text().ok().as_deref() == Some(\"Window\"))\n {\n window_menu.append(&show_menu_bar_item)?;\n window_menu.append(&hide_decorations_item)?;\n } else {\n let window_menu = SubmenuBuilder::new(app, \"Window\")\n .item(&show_menu_bar_item)\n .item(&hide_decorations_item)\n .build()?;\n\n let items = menu.items()?;\n let help_idx = items\n .iter()\n .position(|item| {\n item.as_submenu()\n .and_then(|s| s.text().ok())\n .as_deref()\n == Some(\"Help\")\n })\n .unwrap_or(items.len());\n menu.insert(&window_menu, help_idx)?;\n }\n }\n\n if let Some(help_menu) = menu\n .get(HELP_SUBMENU_ID)\n .and_then(|item| item.as_submenu().cloned())\n {\n help_menu.append(&docs_item)?;\n } else {\n let help_menu = SubmenuBuilder::with_id(app, HELP_SUBMENU_ID, \"Help\")\n .item(&docs_item)\n .build()?;\n menu.append(&help_menu)?;\n }\n\n let state = app.state::();\n if state.debug_mode {\n let toggle_devtools_item = MenuItem::with_id(\n app,\n MENU_TOGGLE_DEVTOOLS_ID,\n \"Toggle DevTools\",\n true,\n Some(\"F12\"),\n )?;\n let open_log_item = MenuItem::with_id(\n app,\n MENU_OPEN_DEBUG_LOG_ID,\n \"Open Debug Log\",\n true,\n None::<&str>,\n )?;\n\n let debug_menu = SubmenuBuilder::new(app, \"Debug\")\n .item(&toggle_devtools_item)\n .item(&open_log_item)\n .build()?;\n menu.append(&debug_menu)?;\n }\n\n app.set_menu(menu)?;\n Ok(())\n}\n\nfn build_tray_menu(app: &AppHandle) -> tauri::Result> {\n let open_app = MenuItem::with_id(app, TRAY_MENU_OPEN_APP_ID, \"Open Onyx\", true, None::<&str>)?;\n let open_chat = MenuItem::with_id(\n app,\n TRAY_MENU_OPEN_CHAT_ID,\n \"Open Chat Window\",\n true,\n None::<&str>,\n )?;\n let show_in_menu_bar = CheckMenuItem::with_id(\n app,\n TRAY_MENU_SHOW_IN_BAR_ID,\n \"Show in Menu Bar\",\n true,\n true,\n None::<&str>,\n )?;\n // Keep it visible/pinned without letting users uncheck (avoids orphaning the tray)\n let _ = show_in_menu_bar.set_enabled(false);\n let quit = PredefinedMenuItem::quit(app, Some(\"Quit Onyx\"))?;\n\n MenuBuilder::new(app)\n .item(&open_app)\n .item(&open_chat)\n .separator()\n .item(&show_in_menu_bar)\n .separator()\n .item(&quit)\n .build()\n}\n\nfn handle_tray_menu_event(app: &AppHandle, id: &str) {\n match id {\n TRAY_MENU_OPEN_APP_ID => {\n focus_main_window(app);\n }\n TRAY_MENU_OPEN_CHAT_ID => {\n focus_main_window(app);\n trigger_new_chat(app);\n }\n TRAY_MENU_QUIT_ID => {\n app.exit(0);\n }\n TRAY_MENU_SHOW_IN_BAR_ID => {\n // No-op for now; the item stays checked/disabled to indicate it's pinned.\n }\n _ => {}\n }\n}\n\nfn setup_tray_icon(app: &AppHandle) -> tauri::Result<()> {\n let mut builder = TrayIconBuilder::with_id(TRAY_ID).tooltip(\"Onyx\");\n\n let tray_icon = Image::from_bytes(TRAY_ICON_BYTES)\n .ok()\n .or_else(|| app.default_window_icon().cloned());\n\n if let Some(icon) = tray_icon {\n builder = builder.icon(icon);\n\n #[cfg(target_os = \"macos\")]\n {\n builder = builder.icon_as_template(true);\n }\n }\n\n if let Ok(menu) = build_tray_menu(app) {\n builder = builder.menu(&menu);\n }\n\n builder\n .on_tray_icon_event(|tray, event| {\n if let TrayIconEvent::Click { .. } = event {\n focus_main_window(tray.app_handle());\n }\n })\n .on_menu_event(|app, event| handle_tray_menu_event(app, event.id().as_ref()))\n .build(app)?;\n\n Ok(())\n}\n\n// ============================================================================\n// Main\n// ============================================================================\n\nfn main() {\n let (config, config_initialized) = load_config();\n let debug_mode = is_debug_mode();\n\n let debug_log_file = if debug_mode {\n eprintln!(\"[ONYX DEBUG] Debug mode enabled\");\n if let Some(path) = get_debug_log_path() {\n eprintln!(\"[ONYX DEBUG] Frontend logs: {}\", path.display());\n }\n eprintln!(\"[ONYX DEBUG] DevTools will open automatically\");\n eprintln!(\"[ONYX DEBUG] Capturing console.log/warn/error/info/debug from webview\");\n init_debug_log_file()\n } else {\n None\n };\n\n tauri::Builder::default()\n .plugin(tauri_plugin_shell::init())\n .plugin(\n tauri::plugin::Builder::::new(\"chat-external-navigation-handler\")\n .on_navigation(|webview, destination_url| {\n let Ok(current_url) = webview.url() else {\n return true;\n };\n\n if should_open_in_external_browser(¤t_url, destination_url) {\n if !open_in_default_browser(destination_url.as_str()) {\n eprintln!(\n \"Failed to open external URL in default browser: {}\",\n destination_url\n );\n }\n return false;\n }\n\n true\n })\n .build(),\n )\n .plugin(tauri_plugin_window_state::Builder::default().build())\n .manage(ConfigState {\n config: RwLock::new(config),\n config_initialized: RwLock::new(config_initialized),\n app_base_url: RwLock::new(None),\n menu_temporarily_visible: RwLock::new(false),\n debug_mode,\n debug_log_file: Mutex::new(debug_log_file),\n })\n .invoke_handler(tauri::generate_handler![\n get_server_url,\n get_bootstrap_state,\n set_server_url,\n get_config_path_cmd,\n open_in_browser,\n open_config_file,\n open_config_directory,\n navigate_to,\n reload_page,\n go_back,\n go_forward,\n new_window,\n reset_config,\n start_drag_window,\n toggle_menu_bar,\n show_menu_bar_temporarily,\n hide_menu_bar_temporary,\n log_from_frontend\n ])\n .on_menu_event(|app, event| match event.id().as_ref() {\n \"open_docs\" => open_docs(),\n \"new_chat\" => trigger_new_chat(app),\n \"new_window\" => trigger_new_window(app),\n \"open_settings\" => open_settings(app),\n \"show_menu_bar\" => handle_menu_bar_toggle(app),\n \"hide_window_decorations\" => handle_decorations_toggle(app),\n MENU_TOGGLE_DEVTOOLS_ID => handle_toggle_devtools(app),\n MENU_OPEN_DEBUG_LOG_ID => handle_open_debug_log(),\n _ => {}\n })\n .setup(move |app| {\n let app_handle = app.handle();\n\n if let Err(e) = setup_app_menu(&app_handle) {\n eprintln!(\"Failed to setup menu: {}\", e);\n }\n\n if let Err(e) = setup_tray_icon(&app_handle) {\n eprintln!(\"Failed to setup tray icon: {}\", e);\n }\n\n // Setup main window with vibrancy effect\n if let Some(window) = app.get_webview_window(\"main\") {\n // Apply vibrancy effect for translucent glass look\n #[cfg(target_os = \"macos\")]\n {\n let _ = apply_vibrancy(&window, NSVisualEffectMaterial::Sidebar, None, None);\n }\n\n if let Ok(url) = window.url() {\n let mut base_url = url;\n base_url.set_query(None);\n base_url.set_fragment(None);\n base_url.set_path(\"/\");\n *app.state::().app_base_url.write().unwrap() = Some(base_url);\n }\n\n #[cfg(target_os = \"macos\")]\n inject_titlebar(window.clone());\n\n apply_settings_to_window(&app_handle, &window);\n maybe_open_devtools(&app_handle, &window);\n\n let _ = window.set_focus();\n }\n\n Ok(())\n })\n .on_page_load(|webview: &Webview, _payload: &PageLoadPayload| {\n inject_chat_link_intercept(webview);\n\n {\n let app = webview.app_handle();\n let state = app.state::();\n if state.debug_mode {\n inject_console_capture(webview);\n }\n }\n\n #[cfg(not(target_os = \"macos\"))]\n {\n let _ = webview.eval(MENU_KEY_HANDLER_SCRIPT);\n\n let app = webview.app_handle();\n let state = app.state::();\n let config = state.config.read().unwrap();\n let temp_visible = *state.menu_temporarily_visible.read().unwrap();\n let label = webview.label().to_string();\n if !config.show_menu_bar && !temp_visible {\n if let Some(win) = app.get_webview_window(&label) {\n let _ = win.hide_menu();\n }\n }\n if config.hide_window_decorations {\n if let Some(win) = app.get_webview_window(&label) {\n let _ = win.set_decorations(false);\n }\n }\n }\n\n #[cfg(target_os = \"macos\")]\n let _ = webview.eval(TITLEBAR_SCRIPT);\n })\n .run(tauri::generate_context!())\n .expect(\"error while running tauri application\");\n}\n" }, { "path": "desktop/src-tauri/tauri.conf.json", "content": "{\n \"$schema\": \"https://schema.tauri.app/config/2.0.0\",\n \"productName\": \"Onyx\",\n \"version\": \"0.0.0-dev\",\n \"identifier\": \"app.onyx.desktop\",\n \"build\": {\n \"beforeBuildCommand\": \"\",\n \"beforeDevCommand\": \"\",\n \"frontendDist\": \"../src\"\n },\n \"app\": {\n \"withGlobalTauri\": true,\n \"windows\": [\n {\n \"title\": \"Onyx\",\n \"label\": \"main\",\n \"url\": \"index.html\",\n \"width\": 1200,\n \"height\": 800,\n \"minWidth\": 800,\n \"minHeight\": 600,\n \"resizable\": true,\n \"fullscreen\": false,\n \"decorations\": true,\n \"transparent\": true,\n \"backgroundColor\": \"#1a1a2e\",\n \"titleBarStyle\": \"Overlay\",\n \"hiddenTitle\": true,\n \"acceptFirstMouse\": true,\n \"tabbingIdentifier\": \"onyx\"\n }\n ],\n \"security\": {\n \"csp\": null\n },\n \"macOSPrivateApi\": true\n },\n \"bundle\": {\n \"active\": true,\n \"targets\": \"all\",\n \"icon\": [\n \"icons/32x32.png\",\n \"icons/128x128.png\",\n \"icons/128x128@2x.png\",\n \"icons/icon.icns\",\n \"icons/icon.ico\"\n ],\n \"category\": \"Productivity\",\n \"shortDescription\": \"Onyx Cloud Desktop App\",\n \"longDescription\": \"A lightweight desktop wrapper for Onyx Cloud - your AI-powered knowledge assistant.\",\n \"macOS\": {\n \"entitlements\": null,\n \"exceptionDomain\": \"cloud.onyx.app\",\n \"minimumSystemVersion\": \"10.15\",\n \"signingIdentity\": null,\n \"dmg\": {\n \"windowSize\": {\n \"width\": 660,\n \"height\": 400\n }\n }\n }\n },\n \"plugins\": {\n \"shell\": {\n \"open\": true\n }\n }\n}\n" }, { "path": "docker-bake.hcl", "content": "group \"default\" {\n targets = [\"backend\", \"model-server\", \"web\"]\n}\n\nvariable \"BACKEND_REPOSITORY\" {\n default = \"onyxdotapp/onyx-backend\"\n}\n\nvariable \"WEB_SERVER_REPOSITORY\" {\n default = \"onyxdotapp/onyx-web-server\"\n}\n\nvariable \"MODEL_SERVER_REPOSITORY\" {\n default = \"onyxdotapp/onyx-model-server\"\n}\n\nvariable \"INTEGRATION_REPOSITORY\" {\n default = \"onyxdotapp/onyx-integration\"\n}\n\nvariable \"CLI_REPOSITORY\" {\n default = \"onyxdotapp/onyx-cli\"\n}\n\nvariable \"TAG\" {\n default = \"latest\"\n}\n\ntarget \"backend\" {\n context = \"backend\"\n dockerfile = \"Dockerfile\"\n\n cache-from = [\"type=registry,ref=${BACKEND_REPOSITORY}:latest\"]\n cache-to = [\"type=inline\"]\n\n tags = [\"${BACKEND_REPOSITORY}:${TAG}\"]\n}\n\ntarget \"web\" {\n context = \"web\"\n dockerfile = \"Dockerfile\"\n\n cache-from = [\"type=registry,ref=${WEB_SERVER_REPOSITORY}:latest\"]\n cache-to = [\"type=inline\"]\n\n tags = [\"${WEB_SERVER_REPOSITORY}:${TAG}\"]\n}\n\ntarget \"model-server\" {\n context = \"backend\"\n\n dockerfile = \"Dockerfile.model_server\"\n\n cache-from = [\"type=registry,ref=${MODEL_SERVER_REPOSITORY}:latest\"]\n cache-to = [\"type=inline\"]\n\n tags = [\"${MODEL_SERVER_REPOSITORY}:${TAG}\"]\n}\n\ntarget \"integration\" {\n context = \"backend\"\n dockerfile = \"tests/integration/Dockerfile\"\n\n // Provide the base image via build context from the backend target\n contexts = {\n base = \"target:backend\"\n }\n\n tags = [\"${INTEGRATION_REPOSITORY}:${TAG}\"]\n}\n\ntarget \"cli\" {\n context = \"cli\"\n dockerfile = \"Dockerfile\"\n\n cache-from = [\"type=registry,ref=${CLI_REPOSITORY}:latest\"]\n cache-to = [\"type=inline\"]\n\n tags = [\"${CLI_REPOSITORY}:${TAG}\"]\n}\n" }, { "path": "docs/METRICS.md", "content": "# Onyx Prometheus Metrics Reference\n\n## Adding New Metrics\n\nAll Prometheus metrics live in the `backend/onyx/server/metrics/` package. Follow these steps to add a new metric.\n\n### 1. Choose the right file (or create a new one)\n\n| File | Purpose |\n|------|---------|\n| `metrics/slow_requests.py` | Slow request counter + callback |\n| `metrics/postgres_connection_pool.py` | SQLAlchemy connection pool metrics |\n| `metrics/prometheus_setup.py` | FastAPI instrumentator config (orchestrator) |\n\nIf your metric is a standalone concern (e.g. cache hit rates, queue depths), create a new file under `metrics/` and keep one metric concept per file.\n\n### 2. Define the metric\n\nUse `prometheus_client` types directly at module level:\n\n```python\n# metrics/my_metric.py\nfrom prometheus_client import Counter\n\n_my_counter = Counter(\n \"onyx_my_counter_total\", # Always prefix with onyx_\n \"Human-readable description\",\n [\"label_a\", \"label_b\"], # Keep label cardinality low\n)\n```\n\n**Naming conventions:**\n- Prefix all metric names with `onyx_`\n- Counters: `_total` suffix (e.g. `onyx_api_slow_requests_total`)\n- Histograms: `_seconds` or `_bytes` suffix for durations/sizes\n- Gauges: no special suffix\n\n**Label cardinality:** Avoid high-cardinality labels (raw user IDs, UUIDs, raw paths). Use route templates like `/api/items/{item_id}` instead of `/api/items/abc-123`.\n\n### 3. Wire it into the instrumentator (if request-scoped)\n\nIf your metric needs to run on every HTTP request, write a callback and register it in `prometheus_setup.py`:\n\n```python\n# metrics/my_metric.py\nfrom prometheus_fastapi_instrumentator.metrics import Info\n\ndef my_metric_callback(info: Info) -> None:\n _my_counter.labels(label_a=info.method, label_b=info.modified_handler).inc()\n```\n\n```python\n# metrics/prometheus_setup.py\nfrom onyx.server.metrics.my_metric import my_metric_callback\n\n# Inside setup_prometheus_metrics():\ninstrumentator.add(my_metric_callback)\n```\n\n### 4. Wire it into setup_prometheus_metrics (if infrastructure-scoped)\n\nFor metrics that attach to engines, pools, or background systems, add a setup function and call it from `setup_prometheus_metrics()` in `metrics/prometheus_setup.py`:\n\n```python\n# metrics/my_metric.py\ndef setup_my_metrics(resource: SomeResource) -> None:\n # Register collectors, attach event listeners, etc.\n ...\n```\n\n```python\n# metrics/prometheus_setup.py — inside setup_prometheus_metrics()\nfrom onyx.server.metrics.my_metric import setup_my_metrics\n\ndef setup_prometheus_metrics(app, engines=None) -> None:\n setup_my_metrics(resource) # Add your call here\n ...\n```\n\nAll metrics initialization is funneled through the single `setup_prometheus_metrics()` call in `onyx/main.py:lifespan()`. Do not add separate setup calls to `main.py`.\n\n### 5. Write tests\n\nAdd tests in `backend/tests/unit/onyx/server/`. Use `unittest.mock.patch` to mock the prometheus objects — don't increment real global counters in tests.\n\n### 6. Document the metric\n\nAdd your metric to the reference tables below in this file. Include the metric name, type, labels, and description.\n\n### 7. Update Grafana dashboards\n\nAfter deploying, add panels to the relevant Grafana dashboard:\n\n1. Open Grafana and navigate to the Onyx dashboard (or create a new one)\n2. Add a new panel — choose the appropriate visualization:\n - **Counters** → use `rate()` in a time series panel (e.g. `rate(onyx_my_counter_total[5m])`)\n - **Histograms** → use `histogram_quantile()` for percentiles, or `_sum/_count` for averages\n - **Gauges** → display directly as a stat or gauge panel\n3. Add meaningful thresholds and alerts where appropriate\n4. Group related panels into rows (e.g. \"API Performance\", \"Database Pool\")\n\n---\n\n## API Server Metrics\n\nThese metrics are exposed at `GET /metrics` on the API server.\n\n### Built-in (via `prometheus-fastapi-instrumentator`)\n\n| Metric | Type | Labels | Description |\n|--------|------|--------|-------------|\n| `http_requests_total` | Counter | `method`, `status`, `handler` | Total request count |\n| `http_request_duration_highr_seconds` | Histogram | _(none)_ | High-resolution latency (many buckets, no labels) |\n| `http_request_duration_seconds` | Histogram | `method`, `handler` | Latency by handler (custom buckets for P95/P99) |\n| `http_request_size_bytes` | Summary | `handler` | Incoming request content length |\n| `http_response_size_bytes` | Summary | `handler` | Outgoing response content length |\n| `http_requests_inprogress` | Gauge | `method`, `handler` | Currently in-flight requests |\n\n### Custom (via `onyx.server.metrics`)\n\n| Metric | Type | Labels | Description |\n|--------|------|--------|-------------|\n| `onyx_api_slow_requests_total` | Counter | `method`, `handler`, `status` | Requests exceeding `SLOW_REQUEST_THRESHOLD_SECONDS` (default 1s) |\n\n### Configuration\n\n| Env Var | Default | Description |\n|---------|---------|-------------|\n| `SLOW_REQUEST_THRESHOLD_SECONDS` | `1.0` | Duration threshold for slow request counting |\n\n### Instrumentator Settings\n\n- `should_group_status_codes=False` — Reports exact HTTP status codes (e.g. 401, 403, 500)\n- `should_instrument_requests_inprogress=True` — Enables the in-progress request gauge\n- `inprogress_labels=True` — Breaks down in-progress gauge by `method` and `handler`\n- `excluded_handlers=[\"/health\", \"/metrics\", \"/openapi.json\"]` — Excludes noisy endpoints from metrics\n\n## Database Pool Metrics\n\nThese metrics provide visibility into SQLAlchemy connection pool state across all three engines (`sync`, `async`, `readonly`). Collected via `onyx.server.metrics.postgres_connection_pool`.\n\n### Pool State (via custom Prometheus collector — snapshot on each scrape)\n\n| Metric | Type | Labels | Description |\n|--------|------|--------|-------------|\n| `onyx_db_pool_checked_out` | Gauge | `engine` | Currently checked-out connections |\n| `onyx_db_pool_checked_in` | Gauge | `engine` | Idle connections available in the pool |\n| `onyx_db_pool_overflow` | Gauge | `engine` | Current overflow connections beyond `pool_size` |\n| `onyx_db_pool_size` | Gauge | `engine` | Configured pool size (constant) |\n\n### Pool Lifecycle (via SQLAlchemy pool event listeners)\n\n| Metric | Type | Labels | Description |\n|--------|------|--------|-------------|\n| `onyx_db_pool_checkout_total` | Counter | `engine` | Total connection checkouts from the pool |\n| `onyx_db_pool_checkin_total` | Counter | `engine` | Total connection checkins to the pool |\n| `onyx_db_pool_connections_created_total` | Counter | `engine` | Total new database connections created |\n| `onyx_db_pool_invalidations_total` | Counter | `engine` | Total connection invalidations |\n| `onyx_db_pool_checkout_timeout_total` | Counter | `engine` | Total connection checkout timeouts |\n\n### Per-Endpoint Attribution (via pool events + endpoint context middleware)\n\n| Metric | Type | Labels | Description |\n|--------|------|--------|-------------|\n| `onyx_db_connections_held_by_endpoint` | Gauge | `handler`, `engine` | DB connections currently held, by endpoint |\n| `onyx_db_connection_hold_seconds` | Histogram | `handler`, `engine` | Duration a DB connection is held by an endpoint |\n\nEngine label values: `sync` (main read-write), `async` (async sessions), `readonly` (read-only user).\n\nConnections from background tasks (Celery) or boot-time warmup appear as `handler=\"unknown\"`.\n\n## OpenSearch Search Metrics\n\nThese metrics track OpenSearch search latency and throughput. Collected via `onyx.server.metrics.opensearch_search`.\n\n| Metric | Type | Labels | Description |\n|--------|------|--------|-------------|\n| `onyx_opensearch_search_client_duration_seconds` | Histogram | `search_type` | Client-side end-to-end latency (network + serialization + server execution) |\n| `onyx_opensearch_search_server_duration_seconds` | Histogram | `search_type` | Server-side execution time from OpenSearch `took` field |\n| `onyx_opensearch_search_total` | Counter | `search_type` | Total search requests sent to OpenSearch |\n| `onyx_opensearch_searches_in_progress` | Gauge | `search_type` | Currently in-flight OpenSearch searches |\n\nSearch type label values: See `OpenSearchSearchType`.\n\n---\n\n## Example PromQL Queries\n\n### Which endpoints are saturated right now?\n\n```promql\n# Top 10 endpoints by in-progress requests\ntopk(10, http_requests_inprogress)\n```\n\n### What's the P99 latency per endpoint?\n\n```promql\n# P99 latency by handler over the last 5 minutes\nhistogram_quantile(0.99, sum by (handler, le) (rate(http_request_duration_seconds_bucket[5m])))\n```\n\n### Which endpoints have the highest request rate?\n\n```promql\n# Requests per second by handler, top 10\ntopk(10, sum by (handler) (rate(http_requests_total[5m])))\n```\n\n### Which endpoints are returning errors?\n\n```promql\n# 5xx error rate by handler\nsum by (handler) (rate(http_requests_total{status=~\"5..\"}[5m]))\n```\n\n### Slow request hotspots\n\n```promql\n# Slow requests per minute by handler\nsum by (handler) (rate(onyx_api_slow_requests_total[5m])) * 60\n```\n\n### Latency trending up?\n\n```promql\n# Compare P50 latency now vs 1 hour ago\nhistogram_quantile(0.5, sum by (le) (rate(http_request_duration_highr_seconds_bucket[5m])))\n -\nhistogram_quantile(0.5, sum by (le) (rate(http_request_duration_highr_seconds_bucket[5m] offset 1h)))\n```\n\n### Overall request throughput\n\n```promql\n# Total requests per second across all endpoints\nsum(rate(http_requests_total[5m]))\n```\n\n### Pool utilization (% of capacity in use)\n\n```promql\n# Sync pool utilization: checked-out / (pool_size + max_overflow)\n# NOTE: Replace 10 with your actual POSTGRES_API_SERVER_POOL_OVERFLOW value.\nonyx_db_pool_checked_out{engine=\"sync\"} / (onyx_db_pool_size{engine=\"sync\"} + 10) * 100\n```\n\n### Pool approaching exhaustion?\n\n```promql\n# Alert when checked-out connections exceed 80% of pool capacity\n# NOTE: Replace 10 with your actual POSTGRES_API_SERVER_POOL_OVERFLOW value.\nonyx_db_pool_checked_out{engine=\"sync\"} > 0.8 * (onyx_db_pool_size{engine=\"sync\"} + 10)\n```\n\n### Which endpoints are hogging DB connections?\n\n```promql\n# Top 10 endpoints by connections currently held\ntopk(10, onyx_db_connections_held_by_endpoint{engine=\"sync\"})\n```\n\n### Which endpoints hold connections the longest?\n\n```promql\n# P99 connection hold time by endpoint\nhistogram_quantile(0.99, sum by (handler, le) (rate(onyx_db_connection_hold_seconds_bucket{engine=\"sync\"}[5m])))\n```\n\n### Connection checkout/checkin rate\n\n```promql\n# Checkouts per second by engine\nsum by (engine) (rate(onyx_db_pool_checkout_total[5m]))\n```\n\n### OpenSearch P99 search latency by type\n\n```promql\n# P99 client-side latency by search type\nhistogram_quantile(0.99, sum by (search_type, le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\n```\n\n### OpenSearch search throughput\n\n```promql\n# Searches per second by type\nsum by (search_type) (rate(onyx_opensearch_search_total[5m]))\n```\n\n### OpenSearch concurrent searches\n\n```promql\n# Total in-flight searches across all instances\nsum(onyx_opensearch_searches_in_progress)\n```\n\n### OpenSearch network overhead\n\n```promql\n# Difference between client and server P50 reveals network/serialization cost.\nhistogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\n -\nhistogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_server_duration_seconds_bucket[5m])))\n```\n" }, { "path": "examples/assistants-api/topics_analyzer.py", "content": "import argparse\nimport os\nimport time\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom datetime import timezone\n\nfrom openai import OpenAI\n\n\nASSISTANT_NAME = \"Topic Analyzer\"\nSYSTEM_PROMPT = \"\"\"\nYou are a helpful assistant that analyzes topics by searching through available \\\ndocuments and providing insights. These available documents come from common \\\nworkplace tools like Slack, emails, Confluence, Google Drive, etc.\n\nWhen analyzing a topic:\n1. Search for relevant information using the search tool\n2. Synthesize the findings into clear insights\n3. Highlight key trends, patterns, or notable developments\n4. Maintain objectivity and cite sources where relevant\n\"\"\"\nUSER_PROMPT = \"\"\"\nPlease analyze and provide insights about this topic: {topic}.\n\nIMPORTANT: do not mention things that are not relevant to the specified topic. \\\nIf there is no relevant information, just say \"No relevant information found.\"\n\"\"\"\n\n\ndef wait_on_run(client: OpenAI, run, thread): # type: ignore\n while run.status == \"queued\" or run.status == \"in_progress\":\n run = client.beta.threads.runs.retrieve(\n thread_id=thread.id,\n run_id=run.id,\n )\n time.sleep(0.5)\n return run\n\n\ndef show_response(messages) -> None: # type: ignore\n # Get only the assistant's response text\n for message in messages.data[::-1]:\n if message.role == \"assistant\":\n for content in message.content:\n if content.type == \"text\":\n print(content.text)\n break\n\n\ndef analyze_topics(topics: list[str]) -> None:\n openai_api_key = os.environ.get(\n \"OPENAI_API_KEY\", \"\"\n )\n onyx_api_key = os.environ.get(\n \"DANSWER_API_KEY\", \"\"\n )\n client = OpenAI(\n api_key=openai_api_key,\n base_url=\"http://localhost:8080/openai-assistants\",\n default_headers={\n \"Authorization\": f\"Bearer {onyx_api_key}\",\n },\n )\n\n # Create an assistant if it doesn't exist\n try:\n assistants = client.beta.assistants.list(limit=100)\n # Find the Topic Analyzer assistant if it exists\n assistant = next((a for a in assistants.data if a.name == ASSISTANT_NAME))\n client.beta.assistants.delete(assistant.id)\n except Exception:\n pass\n\n assistant = client.beta.assistants.create(\n name=ASSISTANT_NAME,\n instructions=SYSTEM_PROMPT,\n tools=[{\"type\": \"SearchTool\"}], # type: ignore\n model=\"gpt-4o\",\n )\n\n # Process each topic individually\n for topic in topics:\n thread = client.beta.threads.create()\n message = client.beta.threads.messages.create(\n thread_id=thread.id,\n role=\"user\",\n content=USER_PROMPT.format(topic=topic),\n )\n\n run = client.beta.threads.runs.create(\n thread_id=thread.id,\n assistant_id=assistant.id,\n tools=[\n { # type: ignore\n \"type\": \"SearchTool\",\n \"retrieval_details\": {\n \"run_search\": \"always\",\n \"filters\": {\n \"time_cutoff\": str(\n datetime.now(timezone.utc) - timedelta(days=7)\n )\n },\n },\n }\n ],\n )\n\n run = wait_on_run(client, run, thread)\n messages = client.beta.threads.messages.list(\n thread_id=thread.id, order=\"asc\", after=message.id\n )\n print(f\"\\nAnalysis for topic: {topic}\")\n print(\"-\" * 40)\n show_response(messages)\n print()\n\n\n# Example usage\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(description=\"Analyze specific topics\")\n parser.add_argument(\"topics\", nargs=\"+\", help=\"Topics to analyze (one or more)\")\n\n args = parser.parse_args()\n analyze_topics(args.topics)\n" }, { "path": "examples/widget/.eslintrc.json", "content": "{\n \"extends\": \"next/core-web-vitals\"\n}\n" }, { "path": "examples/widget/.gitignore", "content": "# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.\n\n# dependencies\n/node_modules\n/.pnp\n.pnp.js\n.yarn/install-state.gz\n\n# testing\n/coverage\n\n# next.js\n/.next/\n/out/\n\n# production\n/build\n\n# misc\n.DS_Store\n*.pem\n\n# debug\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\n\n# local env files\n.env*.local\n\n# vercel\n.vercel\n\n# typescript\n*.tsbuildinfo\nnext-env.d.ts\n" }, { "path": "examples/widget/README.md", "content": "# Onyx Chat Bot Widget\n\nNote: The widget requires a Onyx API key, which is a paid (cloud/enterprise) feature.\n\nThis is a code example for how you can use Onyx's APIs to build a chat bot widget for a website! The main code to look at can be found in `src/app/widget/Widget.tsx`.\n\n## Getting Started\n\nTo get the widget working on your webpage, follow these steps:\n\n### 1. Install Dependencies\n\nEnsure you have the necessary dependencies installed. From the `examples/widget/README.md` file:\n\n```bash\nnpm i\n```\n\n### 2. Set Environment Variables\n\nMake sure to set the environment variables `NEXT_PUBLIC_API_URL` and `NEXT_PUBLIC_API_KEY` in a `.env` file at the root of your project:\n\n```bash\nNEXT_PUBLIC_API_URL=\nNEXT_PUBLIC_API_KEY=\n```\n\n### 3. Run the Development Server\n\nStart the development server to see the widget in action.\n\n```bash\nnpm run dev\n```\n\nOpen [http://localhost:3000](http://localhost:3000) with your browser to see the result.\n\n### 4. Integrate the Widget\n\nTo integrate the widget into your webpage, you can use the `ChatWidget` component. Here’s an example of how to include it in a page component:\n\n```jsx\nimport ChatWidget from \"path/to/ChatWidget\";\nfunction MyPage() {\n return (\n
    \n

    My Webpage

    \n \n
    \n );\n}\nexport default MyPage;\n```\n\n### 5. Deploy\n\nOnce you are satisfied with the widget, you can build and start the application for production:\n\n```bash\nnpm run build\nnpm run start\n```\n\n### Custom Styling and Configuration\n\nIf you need to customize the widget, you can modify the `ChatWidget` component in the `examples/widget/src/app/widget/Widget.tsx` file.\n\nBy following these steps, you should be able to get the chat widget working on your webpage.\n\nIf you want to get fancier, then take a peek at the Chat implementation within Onyx itself [here](https://github.com/onyx-dot-app/onyx/blob/main/web/src/app/chat/ChatPage.tsx#L82).\n" }, { "path": "examples/widget/next.config.mjs", "content": "/** @type {import('next').NextConfig} */\nconst nextConfig = {};\n\nexport default nextConfig;\n" }, { "path": "examples/widget/package.json", "content": "{\n \"name\": \"widget\",\n \"version\": \"0.1.0\",\n \"private\": true,\n \"scripts\": {\n \"dev\": \"next dev\",\n \"build\": \"next build\",\n \"start\": \"next start\",\n \"lint\": \"next lint\"\n },\n \"dependencies\": {\n \"next\": \"^16.1.7\",\n \"react\": \"^19\",\n \"react-dom\": \"^19\",\n \"react-markdown\": \"^10.1.0\"\n },\n \"devDependencies\": {\n \"@tailwindcss/postcss\": \"^4.1.18\",\n \"@types/node\": \"^25\",\n \"@types/react\": \"^19\",\n \"@types/react-dom\": \"^19\",\n \"autoprefixer\": \"^10.4.23\",\n \"eslint\": \"^9\",\n \"eslint-config-next\": \"16.1.2\",\n \"postcss\": \"^8.5.6\",\n \"tailwindcss\": \"^4.1.18\",\n \"typescript\": \"^5\"\n }\n}\n" }, { "path": "examples/widget/postcss.config.mjs", "content": "/** @type {import('postcss-load-config').Config} */\nconst config = {\n plugins: {\n \"@tailwindcss/postcss\": {},\n },\n};\n\nexport default config;\n" }, { "path": "examples/widget/src/app/globals.css", "content": "@import \"tailwindcss\";\n" }, { "path": "examples/widget/src/app/layout.tsx", "content": "import type { Metadata } from \"next\";\nimport { Inter } from \"next/font/google\";\n\nimport \"./globals.css\";\n\nconst inter = Inter({ subsets: [\"latin\"] });\n\nexport const metadata: Metadata = {\n title: \"Example Onyx Widget\",\n description: \"Example Onyx Widget\",\n};\n\nexport default function RootLayout({\n children,\n}: Readonly<{\n children: React.ReactNode;\n}>) {\n return (\n \n {children}\n \n );\n}\n" }, { "path": "examples/widget/src/app/page.tsx", "content": "import { ChatWidget } from \"./widget/Widget\";\n\nexport default function Home() {\n return (\n
    \n \n
    \n );\n}\n" }, { "path": "examples/widget/src/app/widget/Widget.tsx", "content": "\"use client\";\n\nimport React, { useState } from \"react\";\nimport ReactMarkdown from \"react-markdown\";\n\nconst API_URL = process.env.NEXT_PUBLIC_API_URL || \"http://localhost:8080\";\nconst API_KEY = process.env.NEXT_PUBLIC_API_KEY || \"\";\n\ntype NonEmptyObject = { [k: string]: any };\n\nconst processSingleChunk = (\n chunk: string,\n currPartialChunk: string | null,\n): [T | null, string | null] => {\n const completeChunk = (currPartialChunk || \"\") + chunk;\n try {\n // every complete chunk should be valid JSON\n const chunkJson = JSON.parse(completeChunk);\n return [chunkJson, null];\n } catch (err) {\n // if it's not valid JSON, then it's probably an incomplete chunk\n return [null, completeChunk];\n }\n};\n\nconst processRawChunkString = (\n rawChunkString: string,\n previousPartialChunk: string | null,\n): [T[], string | null] => {\n /* This is required because, in practice, we see that nginx does not send over\n each chunk one at a time even with buffering turned off. Instead,\n chunks are sometimes in batches or are sometimes incomplete */\n if (!rawChunkString) {\n return [[], null];\n }\n const chunkSections = rawChunkString\n .split(\"\\n\")\n .filter((chunk) => chunk.length > 0);\n let parsedChunkSections: T[] = [];\n let currPartialChunk = previousPartialChunk;\n chunkSections.forEach((chunk) => {\n const [processedChunk, partialChunk] = processSingleChunk(\n chunk,\n currPartialChunk,\n );\n if (processedChunk) {\n parsedChunkSections.push(processedChunk);\n currPartialChunk = null;\n } else {\n currPartialChunk = partialChunk;\n }\n });\n\n return [parsedChunkSections, currPartialChunk];\n};\n\nasync function* handleStream(\n streamingResponse: Response,\n): AsyncGenerator {\n const reader = streamingResponse.body?.getReader();\n const decoder = new TextDecoder(\"utf-8\");\n\n let previousPartialChunk: string | null = null;\n while (true) {\n const rawChunk = await reader?.read();\n if (!rawChunk) {\n throw new Error(\"Unable to process chunk\");\n }\n const { done, value } = rawChunk;\n if (done) {\n break;\n }\n\n const [completedChunks, partialChunk] = processRawChunkString(\n decoder.decode(value, { stream: true }),\n previousPartialChunk,\n );\n if (!completedChunks.length && !partialChunk) {\n break;\n }\n previousPartialChunk = partialChunk as string | null;\n\n yield await Promise.resolve(completedChunks);\n }\n}\n\nasync function* sendMessage({\n message,\n chatSessionId,\n parentMessageId,\n}: {\n message: string;\n chatSessionId?: number;\n parentMessageId?: number;\n}) {\n if (!chatSessionId || !parentMessageId) {\n // Create a new chat session if one doesn't exist\n const createSessionResponse = await fetch(\n `${API_URL}/chat/create-chat-session`,\n {\n method: \"POST\",\n headers: {\n \"Content-Type\": \"application/json\",\n Authorization: `Bearer ${API_KEY}`,\n },\n body: JSON.stringify({\n // or specify an assistant you have defined\n persona_id: 0,\n }),\n },\n );\n\n if (!createSessionResponse.ok) {\n const errorJson = await createSessionResponse.json();\n const errorMsg = errorJson.message || errorJson.detail || \"\";\n throw Error(`Failed to create chat session - ${errorMsg}`);\n }\n\n const sessionData = await createSessionResponse.json();\n chatSessionId = sessionData.chat_session_id;\n }\n\n const sendMessageResponse = await fetch(`${API_URL}/chat/send-message`, {\n method: \"POST\",\n headers: {\n \"Content-Type\": \"application/json\",\n Authorization: `Bearer ${API_KEY}`,\n },\n body: JSON.stringify({\n chat_session_id: chatSessionId,\n parent_message_id: parentMessageId || null,\n message: message,\n prompt_id: null,\n search_doc_ids: null,\n file_descriptors: [],\n // checkout https://github.com/onyx-dot-app/onyx/blob/main/backend/onyx/search/models.py#L105 for\n // all available options\n retrieval_options: {\n run_search: \"always\",\n filters: null,\n },\n query_override: null,\n }),\n });\n if (!sendMessageResponse.ok) {\n const errorJson = await sendMessageResponse.json();\n const errorMsg = errorJson.message || errorJson.detail || \"\";\n throw Error(`Failed to send message - ${errorMsg}`);\n }\n\n yield* handleStream(sendMessageResponse);\n}\n\nexport const ChatWidget = () => {\n const [messages, setMessages] = useState<{ text: string; isUser: boolean }[]>(\n [],\n );\n const [inputText, setInputText] = useState(\"\");\n const [isLoading, setIsLoading] = useState(false);\n\n const handleSubmit = async (e: React.FormEvent) => {\n e.preventDefault();\n if (inputText.trim()) {\n const initialPrevMessages = messages;\n setMessages([...initialPrevMessages, { text: inputText, isUser: true }]);\n setInputText(\"\");\n setIsLoading(true);\n\n try {\n const messageGenerator = sendMessage({\n message: inputText,\n chatSessionId: undefined,\n parentMessageId: undefined,\n });\n let fullResponse = \"\";\n\n for await (const chunks of messageGenerator) {\n for (const chunk of chunks) {\n if (\"answer_piece\" in chunk) {\n fullResponse += chunk.answer_piece;\n setMessages([\n ...initialPrevMessages,\n { text: inputText, isUser: true },\n { text: fullResponse, isUser: false },\n ]);\n }\n }\n }\n } catch (error) {\n console.error(\"Error sending message:\", error);\n setMessages((prevMessages) => [\n ...prevMessages,\n { text: \"An error occurred. Please try again.\", isUser: false },\n ]);\n } finally {\n setIsLoading(false);\n }\n }\n };\n\n return (\n \n \n Chat Support\n
    \n \n {messages.map((message, index) => (\n \n \n {message.text}\n
    \n
    \n ))}\n {isLoading && (\n
    \n
    \n
    \n
    \n
    \n
    \n
    \n )}\n
    \n \n
    \n setInputText(e.target.value)}\n placeholder=\"Type a message...\"\n className=\"\n w-full\n p-2\n pr-10\n border\n border-gray-300\n rounded-full\n focus:outline-none\n focus:ring-2\n focus:ring-blue-500\n focus:border-transparent\n \"\n disabled={isLoading}\n />\n \n \n \n \n \n
    \n \n
    \n );\n};\n" }, { "path": "examples/widget/tailwind.config.ts", "content": "import type { Config } from \"tailwindcss\";\n\nconst config: Config = {\n content: [\n \"./src/pages/**/*.{js,ts,jsx,tsx,mdx}\",\n \"./src/components/**/*.{js,ts,jsx,tsx,mdx}\",\n \"./src/app/**/*.{js,ts,jsx,tsx,mdx}\",\n ],\n theme: {\n extend: {\n backgroundImage: {\n \"gradient-radial\": \"radial-gradient(var(--tw-gradient-stops))\",\n \"gradient-conic\":\n \"conic-gradient(from 180deg at 50% 50%, var(--tw-gradient-stops))\",\n },\n },\n },\n plugins: [],\n};\nexport default config;\n" }, { "path": "examples/widget/tsconfig.json", "content": "{\n \"compilerOptions\": {\n \"lib\": [\n \"dom\",\n \"dom.iterable\",\n \"esnext\"\n ],\n \"allowJs\": true,\n \"skipLibCheck\": true,\n \"strict\": true,\n \"noEmit\": true,\n \"esModuleInterop\": true,\n \"module\": \"esnext\",\n \"moduleResolution\": \"bundler\",\n \"resolveJsonModule\": true,\n \"isolatedModules\": true,\n \"jsx\": \"preserve\",\n \"incremental\": true,\n \"plugins\": [\n {\n \"name\": \"next\"\n }\n ],\n \"paths\": {\n \"@/*\": [\n \"./src/*\"\n ]\n },\n \"target\": \"ES2017\"\n },\n \"include\": [\n \"next-env.d.ts\",\n \"**/*.ts\",\n \"**/*.tsx\",\n \".next/types/**/*.ts\",\n \".next/dev/types/**/*.ts\"\n ],\n \"exclude\": [\n \"node_modules\"\n ]\n}\n" }, { "path": "extensions/chrome/LICENSE", "content": "MIT License\n\nCopyright (c) 2025 DanswerAI, Inc.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "extensions/chrome/README.md", "content": "# Onyx Chrome Extension\n\nThe Onyx chrome extension lets you research, create, and automate with LLMs powered by your team's unique knowledge. Just hit Ctrl + O on Mac or Alt + O on Windows to instantly access Onyx in your browser:\n\n💡 Know what your company knows, instantly with the Onyx sidebar\n💬 Chat: Onyx provides a natural language chat interface as the main way of interacting with the features.\n🌎 Internal Search: Ask questions and get answers from all your team's knowledge, powered by Onyx's 50+ connectors to all the tools your team uses\n🚀 With a simple Ctrl + O on Mac or Alt + O on Windows - instantly summarize information from any work application\n\n⚡️ Get quick access to the work resources you need.\n🆕 Onyx new tab page puts all of your company’s knowledge at your fingertips\n🤖 Access custom AI Agents for unique use cases, and give them access to tools to take action.\n\n—\n\nOnyx connects with dozens of popular workplace apps like Google Drive, Jira, Confluence, Slack, and more. Use this extension if you have an account created by your team admin.\n\n## Installation\n\nFor Onyx Cloud Users, please visit the Chrome Plugin Store (pending approval still)\n\n## Development\n\n- Load unpacked extension in your browser\n- Modify files in `src` directory\n- Refresh extension in Chrome\n\n## Contributing\n\nSubmit issues or pull requests for improvements\n" }, { "path": "extensions/chrome/manifest.json", "content": "{\n \"manifest_version\": 3,\n \"name\": \"Onyx\",\n \"version\": \"1.1\",\n \"description\": \"Onyx lets you research, create, and automate with LLMs powered by your team's unique knowledge\",\n \"permissions\": [\n \"sidePanel\",\n \"storage\",\n \"activeTab\",\n \"tabs\"\n ],\n \"host_permissions\": [\"\"],\n \"background\": {\n \"service_worker\": \"service_worker.js\",\n \"type\": \"module\"\n },\n \"action\": {\n \"default_icon\": {\n \"16\": \"public/icon16.png\",\n \"48\": \"public/icon48.png\",\n \"128\": \"public/icon128.png\"\n },\n \"default_popup\": \"src/pages/popup.html\"\n },\n \"icons\": {\n \"16\": \"public/icon16.png\",\n \"48\": \"public/icon48.png\",\n \"128\": \"public/icon128.png\"\n },\n \"options_page\": \"src/pages/options.html\",\n \"chrome_url_overrides\": {\n \"newtab\": \"src/pages/onyx_home.html\"\n },\n \"commands\": {\n \"toggleNewTabOverride\": {\n \"suggested_key\": {\n \"default\": \"Ctrl+Shift+O\",\n \"mac\": \"Command+Shift+O\"\n },\n \"description\": \"Toggle Onyx New Tab Override\"\n },\n \"openSidePanel\": {\n \"suggested_key\": {\n \"default\": \"Ctrl+O\",\n \"windows\": \"Alt+O\",\n \"mac\": \"MacCtrl+O\"\n },\n \"description\": \"Open Onyx Side Panel\"\n }\n },\n \"side_panel\": {\n \"default_path\": \"src/pages/panel.html\"\n },\n \"omnibox\": {\n \"keyword\": \"onyx\"\n },\n \"content_scripts\": [\n {\n \"matches\": [\"\"],\n \"js\": [\"src/utils/selection-icon.js\"],\n \"css\": [\"src/styles/selection-icon.css\"]\n }\n ],\n \"web_accessible_resources\": [\n {\n \"resources\": [\"public/icon32.png\"],\n \"matches\": [\"\"]\n }\n ]\n}\n" }, { "path": "extensions/chrome/service_worker.js", "content": "import {\n DEFAULT_ONYX_DOMAIN,\n CHROME_SPECIFIC_STORAGE_KEYS,\n ACTIONS,\n SIDE_PANEL_PATH,\n} from \"./src/utils/constants.js\";\n\n// Track side panel state per window\nconst sidePanelOpenState = new Map();\n\n// Open welcome page on first install\nchrome.runtime.onInstalled.addListener((details) => {\n if (details.reason === \"install\") {\n chrome.storage.local.get(\n { [CHROME_SPECIFIC_STORAGE_KEYS.ONBOARDING_COMPLETE]: false },\n (result) => {\n if (!result[CHROME_SPECIFIC_STORAGE_KEYS.ONBOARDING_COMPLETE]) {\n chrome.tabs.create({ url: \"src/pages/welcome.html\" });\n }\n },\n );\n }\n});\n\nasync function setupSidePanel() {\n if (chrome.sidePanel) {\n try {\n // Don't auto-open side panel on action click since we have a popup menu\n await chrome.sidePanel.setPanelBehavior({\n openPanelOnActionClick: false,\n });\n } catch (error) {\n console.error(\"Error setting up side panel:\", error);\n }\n }\n}\n\nasync function openSidePanel(tabId) {\n try {\n await chrome.sidePanel.open({ tabId });\n } catch (error) {\n console.error(\"Error opening side panel:\", error);\n }\n}\n\nfunction encodeUserPrompt(text) {\n return encodeURIComponent(text).replace(/\\(/g, \"%28\").replace(/\\)/g, \"%29\");\n}\n\nasync function sendToOnyx(info, tab) {\n const selectedText = encodeUserPrompt(info.selectionText);\n const currentUrl = encodeURIComponent(tab.url);\n\n try {\n const result = await chrome.storage.local.get({\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN,\n });\n const url = `${\n result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]\n }${SIDE_PANEL_PATH}?user-prompt=${selectedText}`;\n\n await openSidePanel(tab.id);\n chrome.runtime.sendMessage({\n action: ACTIONS.OPEN_SIDE_PANEL_WITH_INPUT,\n url: url,\n pageUrl: tab.url,\n });\n } catch (error) {\n console.error(\"Error sending to Onyx:\", error);\n }\n}\n\nasync function toggleNewTabOverride() {\n try {\n const result = await chrome.storage.local.get(\n CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB,\n );\n const newValue =\n !result[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB];\n await chrome.storage.local.set({\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]: newValue,\n });\n\n chrome.notifications.create({\n type: \"basic\",\n iconUrl: \"icon.png\",\n title: \"Onyx New Tab\",\n message: `New Tab Override ${newValue ? \"enabled\" : \"disabled\"}`,\n });\n\n // Send a message to inform all tabs about the change\n chrome.tabs.query({}, (tabs) => {\n tabs.forEach((tab) => {\n chrome.tabs.sendMessage(tab.id, {\n action: \"newTabOverrideToggled\",\n value: newValue,\n });\n });\n });\n } catch (error) {\n console.error(\"Error toggling new tab override:\", error);\n }\n}\n\n// Note: This listener won't fire when a popup is defined in manifest.json\n// The popup will show instead. This is kept as a fallback if popup is removed.\nchrome.action.onClicked.addListener((tab) => {\n openSidePanel(tab.id);\n});\n\nchrome.commands.onCommand.addListener(async (command) => {\n if (command === ACTIONS.SEND_TO_ONYX) {\n try {\n const [tab] = await chrome.tabs.query({\n active: true,\n lastFocusedWindow: true,\n });\n if (tab) {\n const response = await chrome.tabs.sendMessage(tab.id, {\n action: ACTIONS.GET_SELECTED_TEXT,\n });\n const selectedText = response?.selectedText || \"\";\n sendToOnyx({ selectionText: selectedText }, tab);\n }\n } catch (error) {\n console.error(\"Error sending to Onyx:\", error);\n }\n } else if (command === ACTIONS.TOGGLE_NEW_TAB_OVERRIDE) {\n toggleNewTabOverride();\n } else if (command === ACTIONS.CLOSE_SIDE_PANEL) {\n try {\n await chrome.sidePanel.hide();\n } catch (error) {\n console.error(\"Error closing side panel via command:\", error);\n }\n } else if (command === ACTIONS.OPEN_SIDE_PANEL) {\n chrome.tabs.query({ active: true, lastFocusedWindow: true }, (tabs) => {\n if (tabs && tabs.length > 0) {\n const tab = tabs[0];\n const windowId = tab.windowId;\n const isOpen = sidePanelOpenState.get(windowId) || false;\n\n if (isOpen) {\n chrome.sidePanel.setOptions({ enabled: false }, () => {\n chrome.sidePanel.setOptions({ enabled: true });\n sidePanelOpenState.set(windowId, false);\n });\n } else {\n chrome.sidePanel.open({ tabId: tab.id });\n sidePanelOpenState.set(windowId, true);\n }\n }\n });\n return;\n } else {\n console.log(\"Unhandled command:\", command);\n }\n});\n\nasync function sendActiveTabUrlToPanel() {\n try {\n const [tab] = await chrome.tabs.query({\n active: true,\n lastFocusedWindow: true,\n });\n if (tab?.url) {\n chrome.runtime.sendMessage({\n action: ACTIONS.TAB_URL_UPDATED,\n url: tab.url,\n });\n }\n } catch (error) {\n console.error(\"[Onyx SW] Error sending tab URL:\", error);\n }\n}\n\nchrome.runtime.onMessage.addListener((request, sender, sendResponse) => {\n if (request.action === ACTIONS.GET_CURRENT_ONYX_DOMAIN) {\n chrome.storage.local.get(\n { [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN },\n (result) => {\n sendResponse({\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]:\n result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN],\n });\n },\n );\n return true;\n }\n if (request.action === ACTIONS.CLOSE_SIDE_PANEL) {\n closeSidePanel();\n chrome.storage.local.get(\n { [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN },\n (result) => {\n chrome.tabs.create({\n url: `${result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]}/auth/login`,\n active: true,\n });\n },\n );\n return true;\n }\n if (request.action === ACTIONS.OPEN_SIDE_PANEL_WITH_INPUT) {\n const { selectedText, pageUrl } = request;\n const tabId = sender.tab?.id;\n const windowId = sender.tab?.windowId;\n\n if (tabId && windowId) {\n chrome.storage.local.get(\n { [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN },\n (result) => {\n const encodedText = encodeUserPrompt(selectedText);\n const onyxDomain = result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n const url = `${onyxDomain}${SIDE_PANEL_PATH}?user-prompt=${encodedText}`;\n\n chrome.storage.session.set({\n pendingInput: {\n url: url,\n pageUrl: pageUrl,\n timestamp: Date.now(),\n },\n });\n\n chrome.sidePanel\n .open({ windowId })\n .then(() => {\n chrome.runtime.sendMessage({\n action: ACTIONS.OPEN_ONYX_WITH_INPUT,\n url: url,\n pageUrl: pageUrl,\n });\n })\n .catch((error) => {\n console.error(\n \"[Onyx SW] Error opening side panel with text:\",\n error,\n );\n });\n },\n );\n } else {\n console.error(\"[Onyx SW] Missing tabId or windowId\");\n }\n return true;\n }\n if (request.action === ACTIONS.TAB_READING_ENABLED) {\n chrome.storage.session.set({ tabReadingEnabled: true });\n sendActiveTabUrlToPanel();\n return false;\n }\n if (request.action === ACTIONS.TAB_READING_DISABLED) {\n chrome.storage.session.set({ tabReadingEnabled: false });\n return false;\n }\n});\n\nchrome.storage.onChanged.addListener((changes, namespace) => {\n if (\n namespace === \"local\" &&\n changes[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]\n ) {\n const newValue =\n changes[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]\n .newValue;\n\n if (newValue === false) {\n chrome.runtime.openOptionsPage();\n }\n }\n});\n\nchrome.windows.onRemoved.addListener((windowId) => {\n sidePanelOpenState.delete(windowId);\n});\n\nchrome.omnibox.setDefaultSuggestion({\n description: 'Search Onyx for \"%s\"',\n});\n\nchrome.omnibox.onInputEntered.addListener(async (text) => {\n try {\n const result = await chrome.storage.local.get({\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN,\n });\n\n const domain = result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n const searchUrl = `${domain}/chat?user-prompt=${encodeURIComponent(text)}`;\n\n chrome.tabs.update({ url: searchUrl });\n } catch (error) {\n console.error(\"Error handling omnibox search:\", error);\n }\n});\n\nchrome.omnibox.onInputChanged.addListener((text, suggest) => {\n if (text.trim()) {\n suggest([\n {\n content: text,\n description: `Search Onyx for \"${text}\"`,\n },\n ]);\n }\n});\n\nchrome.tabs.onActivated.addListener(async (activeInfo) => {\n const result = await chrome.storage.session.get({ tabReadingEnabled: false });\n if (!result.tabReadingEnabled) return;\n try {\n const tab = await chrome.tabs.get(activeInfo.tabId);\n if (tab.url) {\n chrome.runtime.sendMessage({\n action: ACTIONS.TAB_URL_UPDATED,\n url: tab.url,\n });\n }\n } catch (error) {\n console.error(\"[Onyx SW] Error on tab activated:\", error);\n }\n});\n\nchrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {\n if (!changeInfo.url) return;\n const result = await chrome.storage.session.get({ tabReadingEnabled: false });\n if (!result.tabReadingEnabled) return;\n try {\n const [activeTab] = await chrome.tabs.query({\n active: true,\n lastFocusedWindow: true,\n });\n if (activeTab?.id === tabId) {\n chrome.runtime.sendMessage({\n action: ACTIONS.TAB_URL_UPDATED,\n url: changeInfo.url,\n });\n }\n } catch (error) {\n console.error(\"[Onyx SW] Error on tab updated:\", error);\n }\n});\n\nsetupSidePanel();\n" }, { "path": "extensions/chrome/src/pages/onyx_home.html", "content": "\n\n \n \n \n \n Onyx Home\n \n \n \n\n \n
    \n
    \n \n
    \n \n \n\n" }, { "path": "extensions/chrome/src/pages/onyx_home.js", "content": "import {\n CHROME_MESSAGE,\n CHROME_SPECIFIC_STORAGE_KEYS,\n WEB_MESSAGE,\n} from \"../utils/constants.js\";\nimport {\n showErrorModal,\n hideErrorModal,\n initErrorModal,\n} from \"../utils/error-modal.js\";\nimport { getOnyxDomain } from \"../utils/storage.js\";\n\n(function () {\n let mainIframe = document.getElementById(\"onyx-iframe\");\n let preloadedIframe = null;\n const background = document.getElementById(\"background\");\n const content = document.getElementById(\"content\");\n const DEFAULT_LIGHT_BACKGROUND_IMAGE =\n \"https://images.unsplash.com/photo-1692520883599-d543cfe6d43d?q=80&w=2666&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D\";\n const DEFAULT_DARK_BACKGROUND_IMAGE =\n \"https://images.unsplash.com/photo-1692520883599-d543cfe6d43d?q=80&w=2666&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D\";\n\n let iframeLoadTimeout;\n let iframeLoaded = false;\n\n initErrorModal();\n\n async function preloadChatInterface() {\n preloadedIframe = document.createElement(\"iframe\");\n\n const domain = await getOnyxDomain();\n preloadedIframe.src = domain + \"/chat\";\n preloadedIframe.style.opacity = \"0\";\n preloadedIframe.style.visibility = \"hidden\";\n preloadedIframe.style.transition = \"opacity 0.3s ease-in\";\n preloadedIframe.style.border = \"none\";\n preloadedIframe.style.width = \"100%\";\n preloadedIframe.style.height = \"100%\";\n preloadedIframe.style.position = \"absolute\";\n preloadedIframe.style.top = \"0\";\n preloadedIframe.style.left = \"0\";\n preloadedIframe.style.zIndex = \"1\";\n content.appendChild(preloadedIframe);\n }\n\n function setIframeSrc(url) {\n mainIframe.src = url;\n startIframeLoadTimeout();\n iframeLoaded = false;\n }\n\n function startIframeLoadTimeout() {\n clearTimeout(iframeLoadTimeout);\n iframeLoadTimeout = setTimeout(() => {\n if (!iframeLoaded) {\n try {\n if (\n mainIframe.contentWindow.location.pathname.includes(\"/auth/login\")\n ) {\n showLoginPage();\n } else {\n showErrorModal(mainIframe.src);\n }\n } catch (error) {\n showErrorModal(mainIframe.src);\n }\n }\n }, 2500);\n }\n\n function showLoginPage() {\n background.style.opacity = \"0\";\n mainIframe.style.opacity = \"1\";\n mainIframe.style.visibility = \"visible\";\n content.style.opacity = \"1\";\n hideErrorModal();\n }\n\n function setTheme(theme, customBackgroundImage) {\n const imageUrl =\n customBackgroundImage ||\n (theme === \"dark\"\n ? DEFAULT_DARK_BACKGROUND_IMAGE\n : DEFAULT_LIGHT_BACKGROUND_IMAGE);\n background.style.backgroundImage = `url('${imageUrl}')`;\n }\n\n function fadeInContent() {\n content.style.transition = \"opacity 0.5s ease-in\";\n mainIframe.style.transition = \"opacity 0.5s ease-in\";\n content.style.opacity = \"0\";\n mainIframe.style.opacity = \"0\";\n mainIframe.style.visibility = \"visible\";\n\n requestAnimationFrame(() => {\n content.style.opacity = \"1\";\n mainIframe.style.opacity = \"1\";\n\n setTimeout(() => {\n background.style.transition = \"opacity 0.3s ease-out\";\n background.style.opacity = \"0\";\n }, 500);\n });\n }\n\n function checkOnyxPreference() {\n chrome.storage.local.get(\n [\n CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB,\n CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN,\n ],\n (items) => {\n let useOnyxAsDefaultNewTab =\n items[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB];\n\n if (useOnyxAsDefaultNewTab === undefined) {\n useOnyxAsDefaultNewTab = !!(\n localStorage.getItem(\n CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB,\n ) === \"1\"\n );\n chrome.storage.local.set({\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]:\n useOnyxAsDefaultNewTab,\n });\n }\n\n if (!useOnyxAsDefaultNewTab) {\n chrome.tabs.update({\n url: \"chrome://new-tab-page\",\n });\n return;\n }\n\n setIframeSrc(items[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN] + \"/nrf\");\n },\n );\n }\n\n function loadThemeAndBackground() {\n chrome.storage.local.get(\n [\n CHROME_SPECIFIC_STORAGE_KEYS.THEME,\n CHROME_SPECIFIC_STORAGE_KEYS.BACKGROUND_IMAGE,\n CHROME_SPECIFIC_STORAGE_KEYS.DARK_BG_URL,\n CHROME_SPECIFIC_STORAGE_KEYS.LIGHT_BG_URL,\n ],\n function (result) {\n const theme = result[CHROME_SPECIFIC_STORAGE_KEYS.THEME] || \"light\";\n const customBackgroundImage =\n result[CHROME_SPECIFIC_STORAGE_KEYS.BACKGROUND_IMAGE];\n const darkBgUrl = result[CHROME_SPECIFIC_STORAGE_KEYS.DARK_BG_URL];\n const lightBgUrl = result[CHROME_SPECIFIC_STORAGE_KEYS.LIGHT_BG_URL];\n\n let backgroundImage;\n if (customBackgroundImage) {\n backgroundImage = customBackgroundImage;\n } else if (theme === \"dark\" && darkBgUrl) {\n backgroundImage = darkBgUrl;\n } else if (theme === \"light\" && lightBgUrl) {\n backgroundImage = lightBgUrl;\n }\n\n setTheme(theme, backgroundImage);\n checkOnyxPreference();\n },\n );\n }\n\n function loadNewPage(newSrc) {\n if (preloadedIframe && preloadedIframe.contentWindow) {\n preloadedIframe.contentWindow.postMessage(\n { type: WEB_MESSAGE.PAGE_CHANGE, href: newSrc },\n \"*\",\n );\n } else {\n console.error(\"Preloaded iframe not available\");\n }\n }\n\n function completePendingPageLoad() {\n if (preloadedIframe) {\n preloadedIframe.style.visibility = \"visible\";\n preloadedIframe.style.opacity = \"1\";\n preloadedIframe.style.zIndex = \"1\";\n mainIframe.style.zIndex = \"2\";\n mainIframe.style.opacity = \"0\";\n\n setTimeout(() => {\n if (content.contains(mainIframe)) {\n content.removeChild(mainIframe);\n }\n\n mainIframe = preloadedIframe;\n mainIframe.id = \"onyx-iframe\";\n mainIframe.style.zIndex = \"\";\n iframeLoaded = true;\n clearTimeout(iframeLoadTimeout);\n }, 200);\n } else {\n console.warn(\"No preloaded iframe available\");\n }\n }\n\n chrome.storage.onChanged.addListener(function (changes, namespace) {\n if (namespace === \"local\" && changes.useOnyxAsDefaultNewTab) {\n checkOnyxPreference();\n }\n });\n\n window.addEventListener(\"message\", function (event) {\n if (event.data.type === CHROME_MESSAGE.SET_DEFAULT_NEW_TAB) {\n chrome.storage.local.set({ useOnyxAsDefaultNewTab: event.data.value });\n } else if (event.data.type === CHROME_MESSAGE.ONYX_APP_LOADED) {\n clearTimeout(iframeLoadTimeout);\n hideErrorModal();\n fadeInContent();\n iframeLoaded = true;\n } else if (event.data.type === CHROME_MESSAGE.PREFERENCES_UPDATED) {\n const { theme, backgroundUrl } = event.data.payload;\n chrome.storage.local.set(\n {\n [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: theme,\n [CHROME_SPECIFIC_STORAGE_KEYS.BACKGROUND_IMAGE]: backgroundUrl,\n },\n () => {},\n );\n } else if (event.data.type === CHROME_MESSAGE.LOAD_NEW_PAGE) {\n loadNewPage(event.data.href);\n } else if (event.data.type === CHROME_MESSAGE.LOAD_NEW_CHAT_PAGE) {\n completePendingPageLoad();\n }\n });\n\n mainIframe.onload = function () {\n clearTimeout(iframeLoadTimeout);\n startIframeLoadTimeout();\n };\n\n mainIframe.onerror = function (error) {\n showErrorModal(mainIframe.src);\n };\n\n loadThemeAndBackground();\n preloadChatInterface();\n})();\n" }, { "path": "extensions/chrome/src/pages/options.html", "content": "\n\n \n \n \n \n Onyx - Settings\n \n \n \n\n \n
    \n
    \n
    \n
    \n
    \n \"Onyx\"\n
    \n

    Settings

    \n
    \n \n \n \n \n \n \n
    \n\n
    \n \n
    \n
    General
    \n
    \n
    \n
    \n
    \n
    \n
    \n
    \n \n
    \n
    \n
    \n
    \n
    \n \n
    \n
    \n
    \n\n \n
    \n
    Search Engine
    \n
    \n
    \n
    \n \n
    \n Type onyx followed by a space in Chrome's address\n bar, then enter your search query and press Enter\n
    \n
    \n
    \n
    \n
    \n
    \n
    \n Searches will be directed to your configured Onyx instance\n at the Root Domain above\n
    \n
    \n
    \n
    \n
    \n\n \n
    \n

    \n \n
    \n
    \n
    \n
    \n \n \n\n" }, { "path": "extensions/chrome/src/pages/options.js", "content": "import {\n CHROME_SPECIFIC_STORAGE_KEYS,\n DEFAULT_ONYX_DOMAIN,\n} from \"../utils/constants.js\";\n\ndocument.addEventListener(\"DOMContentLoaded\", function () {\n const domainInput = document.getElementById(\"onyxDomain\");\n const useOnyxAsDefaultToggle = document.getElementById(\"useOnyxAsDefault\");\n const statusContainer = document.getElementById(\"statusContainer\");\n const statusElement = document.getElementById(\"status\");\n const newTabButton = document.getElementById(\"newTab\");\n const themeToggle = document.getElementById(\"themeToggle\");\n const themeIcon = document.getElementById(\"themeIcon\");\n\n let currentTheme = \"dark\";\n\n function updateThemeIcon(theme) {\n if (!themeIcon) return;\n\n if (theme === \"light\") {\n themeIcon.innerHTML = `\n \n \n `;\n } else {\n themeIcon.innerHTML = `\n \n `;\n }\n }\n\n function loadStoredValues() {\n chrome.storage.local.get(\n {\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN,\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]: false,\n [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: \"dark\",\n },\n (result) => {\n if (domainInput)\n domainInput.value = result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n if (useOnyxAsDefaultToggle)\n useOnyxAsDefaultToggle.checked =\n result[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB];\n\n currentTheme = result[CHROME_SPECIFIC_STORAGE_KEYS.THEME] || \"dark\";\n updateThemeIcon(currentTheme);\n\n document.body.className = currentTheme === \"light\" ? \"light-theme\" : \"\";\n },\n );\n }\n\n function saveSettings() {\n const domain = domainInput.value.trim();\n const useOnyxAsDefault = useOnyxAsDefaultToggle\n ? useOnyxAsDefaultToggle.checked\n : false;\n\n chrome.storage.local.set(\n {\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: domain,\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]:\n useOnyxAsDefault,\n [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: currentTheme,\n },\n () => {\n showStatusMessage(\n useOnyxAsDefault\n ? \"Settings updated. Open a new tab to test it out. Click on the extension icon to bring up Onyx from any page.\"\n : \"Settings updated.\",\n );\n },\n );\n }\n\n function showStatusMessage(message) {\n if (statusElement) {\n const useOnyxAsDefault = useOnyxAsDefaultToggle\n ? useOnyxAsDefaultToggle.checked\n : false;\n\n statusElement.textContent =\n message ||\n (useOnyxAsDefault\n ? \"Settings updated. Open a new tab to test it out. Click on the extension icon to bring up Onyx from any page.\"\n : \"Settings updated.\");\n\n if (newTabButton) {\n newTabButton.style.display = useOnyxAsDefault ? \"block\" : \"none\";\n }\n }\n\n if (statusContainer) {\n statusContainer.classList.add(\"show\");\n }\n\n setTimeout(hideStatusMessage, 5000);\n }\n\n function hideStatusMessage() {\n if (statusContainer) {\n statusContainer.classList.remove(\"show\");\n }\n }\n\n function toggleTheme() {\n currentTheme = currentTheme === \"light\" ? \"dark\" : \"light\";\n updateThemeIcon(currentTheme);\n\n document.body.className = currentTheme === \"light\" ? \"light-theme\" : \"\";\n\n chrome.storage.local.set({\n [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: currentTheme,\n });\n }\n\n function openNewTab() {\n chrome.tabs.create({});\n }\n\n if (domainInput) {\n domainInput.addEventListener(\"input\", () => {\n clearTimeout(domainInput.saveTimeout);\n domainInput.saveTimeout = setTimeout(saveSettings, 1000);\n });\n }\n\n if (useOnyxAsDefaultToggle) {\n useOnyxAsDefaultToggle.addEventListener(\"change\", saveSettings);\n }\n\n if (themeToggle) {\n themeToggle.addEventListener(\"click\", toggleTheme);\n }\n\n if (newTabButton) {\n newTabButton.addEventListener(\"click\", openNewTab);\n }\n\n loadStoredValues();\n});\n" }, { "path": "extensions/chrome/src/pages/panel.html", "content": "\n\n \n \n \n \n Onyx Panel\n \n \n \n\n \n
    \n
    \n
    Loading Onyx...
    \n
    \n \n \n \n \n\n" }, { "path": "extensions/chrome/src/pages/panel.js", "content": "import { showErrorModal, showAuthModal } from \"../utils/error-modal.js\";\nimport {\n ACTIONS,\n CHROME_MESSAGE,\n WEB_MESSAGE,\n CHROME_SPECIFIC_STORAGE_KEYS,\n SIDE_PANEL_PATH,\n} from \"../utils/constants.js\";\n(function () {\n const iframe = document.getElementById(\"onyx-panel-iframe\");\n const loadingScreen = document.getElementById(\"loading-screen\");\n\n let currentUrl = \"\";\n let iframeLoaded = false;\n let iframeLoadTimeout;\n let authRequired = false;\n\n // Returns the origin of the Onyx app loaded in the iframe.\n // We derive the origin from iframe.src so postMessage payloads\n // (including tab URLs) are only delivered to the expected page.\n // Throws if iframe.src is not a valid URL — this is intentional:\n // postMessage must never fall back to the unsafe wildcard \"*\".\n function getIframeOrigin() {\n return new URL(iframe.src).origin;\n }\n\n async function checkPendingInput() {\n try {\n const result = await chrome.storage.session.get(\"pendingInput\");\n if (result.pendingInput) {\n const { url, pageUrl, timestamp } = result.pendingInput;\n if (Date.now() - timestamp < 5000) {\n setIframeSrc(url, pageUrl);\n await chrome.storage.session.remove(\"pendingInput\");\n return true;\n }\n await chrome.storage.session.remove(\"pendingInput\");\n }\n } catch (error) {\n console.error(\"[Onyx Panel] Error checking pending input:\", error);\n }\n return false;\n }\n\n async function initializePanel() {\n loadingScreen.style.display = \"flex\";\n loadingScreen.style.opacity = \"1\";\n iframe.style.opacity = \"0\";\n\n // Check for pending input first (from selection icon click)\n const hasPendingInput = await checkPendingInput();\n if (!hasPendingInput) {\n loadOnyxDomain();\n }\n }\n\n function setIframeSrc(url, pageUrl) {\n iframe.src = url;\n currentUrl = pageUrl;\n }\n\n function sendWebsiteToIframe(pageUrl) {\n if (iframe.contentWindow && pageUrl !== currentUrl) {\n iframe.contentWindow.postMessage(\n {\n type: WEB_MESSAGE.PAGE_CHANGE,\n url: pageUrl,\n },\n getIframeOrigin(),\n );\n currentUrl = pageUrl;\n }\n }\n\n function startIframeLoadTimeout() {\n iframeLoadTimeout = setTimeout(() => {\n if (!iframeLoaded) {\n if (authRequired) {\n showAuthModal();\n } else {\n showErrorModal(iframe.src);\n }\n }\n }, 2500);\n }\n\n function handleMessage(event) {\n // Only trust messages from the Onyx app iframe.\n // Check both source identity and origin so that a cross-origin page\n // navigated to inside the iframe cannot send privileged extension\n // messages (e.g. TAB_READING_ENABLED) after iframe.src changes.\n // getIframeOrigin() throws if iframe.src is not yet a valid URL —\n // catching it here fails closed (message is rejected, not processed).\n if (event.source !== iframe.contentWindow) return;\n try {\n if (event.origin !== getIframeOrigin()) return;\n } catch {\n return;\n }\n if (event.data.type === CHROME_MESSAGE.ONYX_APP_LOADED) {\n clearTimeout(iframeLoadTimeout);\n iframeLoaded = true;\n showIframe();\n if (iframe.contentWindow) {\n iframe.contentWindow.postMessage(\n { type: \"PANEL_READY\" },\n getIframeOrigin(),\n );\n }\n } else if (event.data.type === CHROME_MESSAGE.AUTH_REQUIRED) {\n authRequired = true;\n } else if (event.data.type === CHROME_MESSAGE.TAB_READING_ENABLED) {\n chrome.runtime.sendMessage({ action: ACTIONS.TAB_READING_ENABLED });\n } else if (event.data.type === CHROME_MESSAGE.TAB_READING_DISABLED) {\n chrome.runtime.sendMessage({ action: ACTIONS.TAB_READING_DISABLED });\n }\n }\n\n function showIframe() {\n iframe.style.opacity = \"1\";\n loadingScreen.style.opacity = \"0\";\n setTimeout(() => {\n loadingScreen.style.display = \"none\";\n }, 500);\n }\n\n async function loadOnyxDomain() {\n const response = await chrome.runtime.sendMessage({\n action: ACTIONS.GET_CURRENT_ONYX_DOMAIN,\n });\n if (response && response[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]) {\n setIframeSrc(\n response[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN] + SIDE_PANEL_PATH,\n \"\",\n );\n } else {\n console.warn(\"Onyx domain not found, using default\");\n const domain = await getOnyxDomain();\n setIframeSrc(domain + SIDE_PANEL_PATH, \"\");\n }\n }\n\n chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {\n if (request.action === ACTIONS.OPEN_ONYX_WITH_INPUT) {\n setIframeSrc(request.url, request.pageUrl);\n } else if (request.action === ACTIONS.UPDATE_PAGE_URL) {\n sendWebsiteToIframe(request.pageUrl);\n } else if (request.action === ACTIONS.TAB_URL_UPDATED) {\n if (iframe.contentWindow) {\n iframe.contentWindow.postMessage(\n { type: CHROME_MESSAGE.TAB_URL_UPDATED, url: request.url },\n getIframeOrigin(),\n );\n }\n }\n });\n\n window.addEventListener(\"message\", handleMessage);\n\n initializePanel();\n startIframeLoadTimeout();\n})();\n" }, { "path": "extensions/chrome/src/pages/popup.html", "content": "\n\n \n \n \n \n Onyx\n \n \n \n \n
    \n
    \n
    \n \"Onyx\"\n
    \n

    Onyx

    \n
    \n\n
    \n
    \n \n \n
    \n
    \n\n
    \n \n\n \n
    \n
    \n \n \n\n" }, { "path": "extensions/chrome/src/pages/popup.js", "content": "import { CHROME_SPECIFIC_STORAGE_KEYS } from \"../utils/constants.js\";\n\ndocument.addEventListener(\"DOMContentLoaded\", async function () {\n const defaultNewTabToggle = document.getElementById(\"defaultNewTabToggle\");\n const openSidePanelButton = document.getElementById(\"openSidePanel\");\n const openOptionsButton = document.getElementById(\"openOptions\");\n\n async function loadSetting() {\n const result = await chrome.storage.local.get({\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]: false,\n });\n if (defaultNewTabToggle) {\n defaultNewTabToggle.checked =\n result[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB];\n }\n }\n\n async function toggleSetting() {\n const currentValue = defaultNewTabToggle.checked;\n await chrome.storage.local.set({\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]: currentValue,\n });\n }\n\n async function openSidePanel() {\n try {\n const [tab] = await chrome.tabs.query({\n active: true,\n currentWindow: true,\n });\n if (tab && chrome.sidePanel) {\n await chrome.sidePanel.open({ tabId: tab.id });\n window.close();\n }\n } catch (error) {\n console.error(\"Error opening side panel:\", error);\n }\n }\n\n function openOptions() {\n chrome.runtime.openOptionsPage();\n window.close();\n }\n\n await loadSetting();\n\n if (defaultNewTabToggle) {\n defaultNewTabToggle.addEventListener(\"change\", toggleSetting);\n }\n\n if (openSidePanelButton) {\n openSidePanelButton.addEventListener(\"click\", openSidePanel);\n }\n\n if (openOptionsButton) {\n openOptionsButton.addEventListener(\"click\", openOptions);\n }\n});\n" }, { "path": "extensions/chrome/src/pages/welcome.html", "content": "\n\n \n \n \n Welcome to Onyx\n \n \n \n \n \n \n\n \n
    \n
    \n
    \n
    \n
    \n \"Onyx\"\n
    \n

    Onyx

    \n
    \n \n \n \n \n \n
    \n\n
    \n
    \n
    \n
    \n
    \n\n \n
    \n

    Welcome to Onyx

    \n

    \n Enter your Onyx instance URL to get started. This is where your\n Onyx deployment is hosted.\n

    \n\n
    \n \n \n
    \n\n
    \n \n
    \n
    \n\n \n
    \n

    Customize Your Experience

    \n

    \n Set Onyx as your new tab page for quick access to your AI\n assistant.\n

    \n\n
    \n
    \n
    \n Use Onyx as new tab page\n Open Onyx every time you create a new tab\n
    \n \n
    \n
    \n\n
    \n \n \n
    \n
    \n
    \n
    \n
    \n \n \n\n" }, { "path": "extensions/chrome/src/pages/welcome.js", "content": "import {\n CHROME_SPECIFIC_STORAGE_KEYS,\n DEFAULT_ONYX_DOMAIN,\n} from \"../utils/constants.js\";\n\ndocument.addEventListener(\"DOMContentLoaded\", function () {\n const domainInput = document.getElementById(\"onyxDomain\");\n const useOnyxAsDefaultToggle = document.getElementById(\"useOnyxAsDefault\");\n const continueBtn = document.getElementById(\"continueBtn\");\n const backBtn = document.getElementById(\"backBtn\");\n const finishBtn = document.getElementById(\"finishBtn\");\n const themeToggle = document.getElementById(\"themeToggle\");\n const themeIcon = document.getElementById(\"themeIcon\");\n\n const step1 = document.getElementById(\"step1\");\n const step2 = document.getElementById(\"step2\");\n const stepDots = document.querySelectorAll(\".step-dot\");\n\n let currentStep = 1;\n let currentTheme = \"dark\";\n\n // Initialize theme based on system preference or stored value\n function initTheme() {\n chrome.storage.local.get(\n { [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: null },\n (result) => {\n const storedTheme = result[CHROME_SPECIFIC_STORAGE_KEYS.THEME];\n if (storedTheme) {\n currentTheme = storedTheme;\n } else {\n // Check system preference\n currentTheme = window.matchMedia(\"(prefers-color-scheme: light)\")\n .matches\n ? \"light\"\n : \"dark\";\n }\n applyTheme();\n },\n );\n }\n\n function applyTheme() {\n document.body.className = currentTheme === \"light\" ? \"light-theme\" : \"\";\n updateThemeIcon();\n }\n\n function updateThemeIcon() {\n if (!themeIcon) return;\n\n if (currentTheme === \"light\") {\n themeIcon.innerHTML = `\n \n \n `;\n } else {\n themeIcon.innerHTML = `\n \n `;\n }\n }\n\n function toggleTheme() {\n currentTheme = currentTheme === \"light\" ? \"dark\" : \"light\";\n applyTheme();\n chrome.storage.local.set({\n [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: currentTheme,\n });\n }\n\n function goToStep(step) {\n if (step === 1) {\n step2.classList.remove(\"active\");\n setTimeout(() => {\n step1.classList.add(\"active\");\n }, 50);\n } else if (step === 2) {\n step1.classList.remove(\"active\");\n setTimeout(() => {\n step2.classList.add(\"active\");\n }, 50);\n }\n\n stepDots.forEach((dot) => {\n const dotStep = parseInt(dot.dataset.step);\n if (dotStep === step) {\n dot.classList.add(\"active\");\n } else {\n dot.classList.remove(\"active\");\n }\n });\n\n currentStep = step;\n }\n\n // Validate domain input\n function validateDomain(domain) {\n if (!domain) return false;\n try {\n new URL(domain);\n return true;\n } catch {\n return false;\n }\n }\n\n function handleContinue() {\n const domain = domainInput.value.trim();\n\n if (domain && !validateDomain(domain)) {\n domainInput.style.borderColor = \"rgba(255, 100, 100, 0.5)\";\n domainInput.focus();\n return;\n }\n\n domainInput.style.borderColor = \"\";\n goToStep(2);\n }\n\n function handleBack() {\n goToStep(1);\n }\n\n function handleFinish() {\n const domain = domainInput.value.trim() || DEFAULT_ONYX_DOMAIN;\n const useOnyxAsDefault = useOnyxAsDefaultToggle.checked;\n\n chrome.storage.local.set(\n {\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: domain,\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]:\n useOnyxAsDefault,\n [CHROME_SPECIFIC_STORAGE_KEYS.THEME]: currentTheme,\n [CHROME_SPECIFIC_STORAGE_KEYS.ONBOARDING_COMPLETE]: true,\n },\n () => {\n // Open a new tab if they enabled the new tab feature, otherwise just close\n if (useOnyxAsDefault) {\n chrome.tabs.create({}, () => {\n window.close();\n });\n } else {\n window.close();\n }\n },\n );\n }\n\n // Load any existing values (in case user returns to this page)\n function loadStoredValues() {\n chrome.storage.local.get(\n {\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: \"\",\n [CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB]: true,\n },\n (result) => {\n if (result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]) {\n domainInput.value = result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n }\n useOnyxAsDefaultToggle.checked =\n result[CHROME_SPECIFIC_STORAGE_KEYS.USE_ONYX_AS_DEFAULT_NEW_TAB];\n },\n );\n }\n\n if (themeToggle) {\n themeToggle.addEventListener(\"click\", toggleTheme);\n }\n\n if (continueBtn) {\n continueBtn.addEventListener(\"click\", handleContinue);\n }\n\n if (backBtn) {\n backBtn.addEventListener(\"click\", handleBack);\n }\n\n if (finishBtn) {\n finishBtn.addEventListener(\"click\", handleFinish);\n }\n\n // Allow Enter key to proceed\n if (domainInput) {\n domainInput.addEventListener(\"keydown\", (e) => {\n if (e.key === \"Enter\") {\n handleContinue();\n }\n });\n }\n\n initTheme();\n loadStoredValues();\n});\n" }, { "path": "extensions/chrome/src/styles/selection-icon.css", "content": "#onyx-selection-icon {\n position: fixed;\n z-index: 2147483647;\n width: 32px;\n height: 32px;\n border-radius: 50%;\n background-color: #ffffff;\n border: 1px solid #e0e0e0;\n box-shadow: 0 2px 8px rgba(0, 0, 0, 0.15);\n cursor: pointer;\n display: flex;\n align-items: center;\n justify-content: center;\n opacity: 0;\n transform: scale(0.8);\n transition:\n opacity 0.15s ease,\n transform 0.15s ease,\n box-shadow 0.15s ease;\n pointer-events: none;\n}\n\n#onyx-selection-icon.visible {\n opacity: 1;\n transform: scale(1);\n pointer-events: auto;\n}\n\n#onyx-selection-icon:hover {\n box-shadow: 0 4px 12px rgba(0, 0, 0, 0.2);\n transform: scale(1.1);\n}\n\n#onyx-selection-icon:active {\n transform: scale(0.95);\n}\n\n#onyx-selection-icon img {\n width: 20px;\n height: 20px;\n pointer-events: none;\n}\n" }, { "path": "extensions/chrome/src/styles/shared.css", "content": "/* Import Hanken Grotesk font */\n@import url(\"https://fonts.googleapis.com/css2?family=Hanken+Grotesk:wght@300;400;500;600;700&display=swap\");\n\n:root {\n --primary-color: #4285f4;\n --primary-hover-color: #3367d6;\n --secondary-color: #f1f3f4;\n --secondary-hover-color: #e8eaed;\n --text-color: #333;\n --text-light-color: #666;\n --background-color: #f1f3f4;\n --card-background-color: #fff;\n --border-color: #ccc;\n --font-family: Arial, sans-serif;\n --font-hanken-grotesk: \"Hanken Grotesk\", sans-serif;\n}\n\nbody {\n font-family: var(--font-hanken-grotesk);\n margin: 0;\n padding: 0;\n}\n\n.container {\n max-width: 500px;\n width: 90%;\n margin: 0 auto;\n}\n\n.card {\n background-color: var(--card-background-color);\n padding: 25px;\n border-radius: 10px;\n box-shadow: 0 3px 5px rgba(0, 0, 0, 0.1);\n}\n\nh1 {\n color: var(--text-color);\n font-size: 24px;\n font-weight: 600;\n margin-top: 0;\n margin-bottom: 20px;\n}\n\n.option-group {\n margin-bottom: 20px;\n}\n\nlabel {\n display: block;\n margin-bottom: 5px;\n color: var(--text-light-color);\n font-weight: 400;\n font-size: 16px;\n}\n\ninput[type=\"text\"] {\n width: 100%;\n padding: 8px;\n border: 1px solid var(--border-color);\n border-radius: 4px;\n font-size: 14px;\n background-color: var(--card-background-color);\n color: var(--text-color);\n}\n\n.button {\n width: 100%;\n padding: 10px 20px;\n border-radius: 5px;\n border: none;\n cursor: pointer;\n font-size: 16px;\n font-weight: 500;\n transition: background-color 0.3s;\n}\n\n.button.primary {\n background-color: var(--primary-color);\n color: #fff;\n}\n\n.button.primary:hover {\n background-color: var(--primary-hover-color);\n}\n\n.button.secondary {\n background-color: var(--secondary-color);\n color: var(--text-color);\n}\n\n.button.secondary:hover {\n background-color: var(--secondary-hover-color);\n}\n\n.status-container {\n margin-top: 10px;\n margin-bottom: 15px;\n}\n\n.status-message {\n margin: 0 0 10px 0;\n color: var(--text-color);\n font-weight: 500;\n text-align: center;\n font-size: 16px;\n transition: opacity 0.5s ease-in-out;\n}\n\nkbd {\n background-color: var(--secondary-color);\n border: 1px solid var(--border-color);\n border-radius: 3px;\n padding: 2px 5px;\n font-family: monospace;\n font-weight: 500;\n color: var(--text-color);\n}\n\n.toggle-label {\n display: flex;\n justify-content: space-between;\n align-items: center;\n}\n\n.toggle-switch {\n position: relative;\n display: inline-block;\n width: 50px;\n height: 24px;\n}\n\n.toggle-switch input {\n opacity: 0;\n width: 0;\n height: 0;\n}\n\n.slider {\n position: absolute;\n cursor: pointer;\n top: 0;\n left: 0;\n right: 0;\n bottom: 0;\n background-color: var(--secondary-color);\n transition: 0.4s;\n border-radius: 24px;\n}\n\n.slider:before {\n position: absolute;\n content: \"\";\n height: 20px;\n width: 20px;\n left: 2px;\n bottom: 2px;\n background-color: white;\n transition: 0.4s;\n border-radius: 50%;\n}\n\ninput:checked + .slider {\n background-color: var(--primary-color);\n}\n\ninput:checked + .slider:before {\n transform: translateX(26px);\n}\n" }, { "path": "extensions/chrome/src/utils/constants.js", "content": "export const THEMES = {\n LIGHT: \"light\",\n DARK: \"dark\",\n};\n\nexport const DEFAULT_ONYX_DOMAIN = \"http://localhost:3000\";\n\nexport const SIDE_PANEL_PATH = \"/nrf/side-panel\";\n\nexport const ACTIONS = {\n GET_SELECTED_TEXT: \"getSelectedText\",\n GET_CURRENT_ONYX_DOMAIN: \"getCurrentOnyxDomain\",\n UPDATE_PAGE_URL: \"updatePageUrl\",\n SEND_TO_ONYX: \"sendToOnyx\",\n OPEN_SIDE_PANEL: \"openSidePanel\",\n TOGGLE_NEW_TAB_OVERRIDE: \"toggleNewTabOverride\",\n OPEN_SIDE_PANEL_WITH_INPUT: \"openSidePanelWithInput\",\n OPEN_ONYX_WITH_INPUT: \"openOnyxWithInput\",\n CLOSE_SIDE_PANEL: \"closeSidePanel\",\n TAB_URL_UPDATED: \"tabUrlUpdated\",\n TAB_READING_ENABLED: \"tabReadingEnabled\",\n TAB_READING_DISABLED: \"tabReadingDisabled\",\n};\n\nexport const CHROME_SPECIFIC_STORAGE_KEYS = {\n ONYX_DOMAIN: \"onyxExtensionDomain\",\n USE_ONYX_AS_DEFAULT_NEW_TAB: \"onyxExtensionDefaultNewTab\",\n THEME: \"onyxExtensionTheme\",\n BACKGROUND_IMAGE: \"onyxExtensionBackgroundImage\",\n DARK_BG_URL: \"onyxExtensionDarkBgUrl\",\n LIGHT_BG_URL: \"onyxExtensionLightBgUrl\",\n ONBOARDING_COMPLETE: \"onyxExtensionOnboardingComplete\",\n};\n\nexport const CHROME_MESSAGE = {\n PREFERENCES_UPDATED: \"PREFERENCES_UPDATED\",\n ONYX_APP_LOADED: \"ONYX_APP_LOADED\",\n SET_DEFAULT_NEW_TAB: \"SET_DEFAULT_NEW_TAB\",\n LOAD_NEW_CHAT_PAGE: \"LOAD_NEW_CHAT_PAGE\",\n LOAD_NEW_PAGE: \"LOAD_NEW_PAGE\",\n AUTH_REQUIRED: \"AUTH_REQUIRED\",\n TAB_READING_ENABLED: \"TAB_READING_ENABLED\",\n TAB_READING_DISABLED: \"TAB_READING_DISABLED\",\n TAB_URL_UPDATED: \"TAB_URL_UPDATED\",\n};\n\nexport const WEB_MESSAGE = {\n PAGE_CHANGE: \"PAGE_CHANGE\",\n};\n" }, { "path": "extensions/chrome/src/utils/content.js", "content": "let sidePanel = null;\n\nfunction createSidePanel() {\n sidePanel = document.createElement(\"div\");\n sidePanel.id = \"onyx-side-panel\";\n sidePanel.style.cssText = `\n position: fixed;\n top: 0;\n right: -400px;\n width: 400px;\n height: 100%;\n background-color: white;\n box-shadow: -2px 0 5px rgba(0,0,0,0.2);\n transition: right 0.3s ease-in-out;\n z-index: 9999;\n `;\n\n const iframe = document.createElement(\"iframe\");\n iframe.style.cssText = `\n width: 100%;\n height: 100%;\n border: none;\n `;\n\n chrome.runtime.sendMessage(\n { action: ACTIONS.GET_CURRENT_ONYX_DOMAIN },\n function (response) {\n iframe.src = response[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n },\n );\n\n sidePanel.appendChild(iframe);\n document.body.appendChild(sidePanel);\n}\n" }, { "path": "extensions/chrome/src/utils/error-modal.js", "content": "import {\n CHROME_SPECIFIC_STORAGE_KEYS,\n DEFAULT_ONYX_DOMAIN,\n ACTIONS,\n} from \"./constants.js\";\n\nconst errorModalHTML = `\n
    \n
    \n
    \n
    \n
    \n \n \n \n \n \n
    \n

    Configuration Error

    \n
    \n
    \n

    The Onyx configuration needs to be updated. Please check your settings or contact your Onyx administrator.

    \n
    \n Attempted to load:\n \n
    \n
    \n
    \n
    \n \n \n
    \n
    \n
    \n
    \n`;\n\nconst style = document.createElement(\"style\");\nstyle.textContent = `\n :root {\n --background-900: #0a0a0a;\n --background-800: #1a1a1a;\n --text-light-05: rgba(255, 255, 255, 0.95);\n --text-light-03: rgba(255, 255, 255, 0.6);\n --white-10: rgba(255, 255, 255, 0.1);\n --white-15: rgba(255, 255, 255, 0.15);\n --white-20: rgba(255, 255, 255, 0.2);\n --white-30: rgba(255, 255, 255, 0.3);\n }\n\n #error-modal {\n position: fixed;\n top: 0;\n left: 0;\n width: 100%;\n height: 100%;\n display: none;\n align-items: center;\n justify-content: center;\n z-index: 2000;\n font-family: var(--font-hanken-grotesk), 'Hanken Grotesk', sans-serif;\n }\n\n #error-modal .modal-backdrop {\n position: absolute;\n top: 0;\n left: 0;\n width: 100%;\n height: 100%;\n background: rgba(0, 0, 0, 0.7);\n backdrop-filter: blur(8px);\n }\n\n #error-modal .modal-content {\n position: relative;\n background: linear-gradient(to bottom, rgba(10, 10, 10, 0.95), rgba(26, 26, 26, 0.95));\n backdrop-filter: blur(24px);\n border-radius: 16px;\n border: 1px solid var(--white-10);\n max-width: 95%;\n width: 500px;\n box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4);\n overflow: hidden;\n }\n\n #error-modal .modal-header {\n padding: 24px;\n border-bottom: 1px solid var(--white-10);\n display: flex;\n align-items: center;\n gap: 12px;\n }\n\n #error-modal .modal-icon {\n width: 40px;\n height: 40px;\n border-radius: 12px;\n background: rgba(255, 87, 87, 0.15);\n display: flex;\n align-items: center;\n justify-content: center;\n flex-shrink: 0;\n }\n\n #error-modal .modal-icon svg {\n width: 24px;\n height: 24px;\n stroke: #ff5757;\n }\n\n #error-modal .modal-icon.auth-icon {\n background: rgba(66, 133, 244, 0.15);\n }\n\n #error-modal .modal-icon.auth-icon svg {\n stroke: #4285f4;\n }\n\n #error-modal h2 {\n margin: 0;\n color: var(--text-light-05);\n font-size: 20px;\n font-weight: 600;\n }\n\n #error-modal .modal-body {\n padding: 24px;\n }\n\n #error-modal .modal-description {\n color: var(--text-light-05);\n margin: 0 0 20px 0;\n font-size: 14px;\n line-height: 1.6;\n font-weight: 400;\n }\n\n #error-modal .url-display {\n background: rgba(255, 255, 255, 0.05);\n border-radius: 8px;\n padding: 12px;\n border: 1px solid var(--white-10);\n }\n\n #error-modal .url-label {\n display: block;\n font-size: 12px;\n color: var(--text-light-03);\n margin-bottom: 6px;\n font-weight: 500;\n text-transform: uppercase;\n letter-spacing: 0.05em;\n }\n\n #error-modal .url-value {\n display: block;\n font-size: 13px;\n color: var(--text-light-05);\n word-break: break-all;\n font-family: monospace;\n line-height: 1.5;\n }\n\n #error-modal .modal-footer {\n padding: 0 24px 24px 24px;\n }\n\n #error-modal .button-container {\n display: flex;\n flex-direction: column;\n gap: 10px;\n margin-bottom: 16px;\n }\n\n #error-modal .button {\n padding: 12px 20px;\n border-radius: 8px;\n border: none;\n cursor: pointer;\n font-size: 14px;\n font-weight: 500;\n transition: all 0.2s;\n font-family: var(--font-hanken-grotesk), 'Hanken Grotesk', sans-serif;\n }\n\n #error-modal .button.primary {\n background: rgba(255, 255, 255, 0.15);\n color: var(--text-light-05);\n border: 1px solid var(--white-10);\n }\n\n #error-modal .button.primary:hover {\n background: rgba(255, 255, 255, 0.2);\n border-color: var(--white-20);\n }\n\n #error-modal .button.secondary {\n background: rgba(255, 255, 255, 0.05);\n color: var(--text-light-05);\n border: 1px solid var(--white-10);\n }\n\n #error-modal .button.secondary:hover {\n background: rgba(255, 255, 255, 0.1);\n border-color: var(--white-15);\n }\n\n #error-modal kbd {\n background: rgba(255, 255, 255, 0.1);\n border: 1px solid var(--white-10);\n border-radius: 4px;\n padding: 2px 6px;\n font-family: monospace;\n font-weight: 500;\n color: var(--text-light-05);\n font-size: 11px;\n }\n\n @media (min-width: 768px) {\n #error-modal .button-container {\n flex-direction: row;\n }\n\n #error-modal .button {\n flex: 1;\n }\n }\n`;\n\nconst authModalHTML = `\n
    \n
    \n
    \n
    \n
    \n \n \n \n \n
    \n

    Authentication Required

    \n
    \n
    \n

    You need to log in to access Onyx. Click the button below to authenticate.

    \n
    \n
    \n
    \n \n
    \n
    \n
    \n
    \n`;\n\nlet errorModal, attemptedUrlSpan, openOptionsButton, disableOverrideButton;\n\nlet authModal, openAuthButton;\n\nexport function initErrorModal() {\n if (!document.getElementById(\"error-modal\")) {\n const link = document.createElement(\"link\");\n link.rel = \"stylesheet\";\n link.href = \"../styles/shared.css\";\n document.head.appendChild(link);\n\n document.body.insertAdjacentHTML(\"beforeend\", errorModalHTML);\n document.head.appendChild(style);\n\n errorModal = document.getElementById(\"error-modal\");\n authModal = document.getElementById(\"error-modal\");\n attemptedUrlSpan = document.getElementById(\"attempted-url\");\n openOptionsButton = document.getElementById(\"open-options\");\n disableOverrideButton = document.getElementById(\"disable-override\");\n\n openOptionsButton.addEventListener(\"click\", (e) => {\n e.preventDefault();\n chrome.runtime.openOptionsPage();\n });\n\n disableOverrideButton.addEventListener(\"click\", () => {\n chrome.storage.local.set({ useOnyxAsDefaultNewTab: false }, () => {\n chrome.tabs.update({ url: \"chrome://new-tab-page\" });\n });\n });\n }\n}\n\nexport function showErrorModal(url) {\n if (!errorModal) {\n initErrorModal();\n }\n if (errorModal) {\n errorModal.style.display = \"flex\";\n errorModal.style.zIndex = \"9999\";\n attemptedUrlSpan.textContent = url;\n document.body.style.overflow = \"hidden\";\n }\n}\n\nexport function hideErrorModal() {\n if (errorModal) {\n errorModal.style.display = \"none\";\n document.body.style.overflow = \"auto\";\n }\n}\n\nexport function checkModalVisibility() {\n return errorModal\n ? window.getComputedStyle(errorModal).display !== \"none\"\n : false;\n}\n\nexport function initAuthModal() {\n if (!document.getElementById(\"error-modal\")) {\n const link = document.createElement(\"link\");\n link.rel = \"stylesheet\";\n link.href = \"../styles/shared.css\";\n document.head.appendChild(link);\n\n document.body.insertAdjacentHTML(\"beforeend\", authModalHTML);\n document.head.appendChild(style);\n\n authModal = document.getElementById(\"error-modal\");\n openAuthButton = document.getElementById(\"open-auth\");\n\n openAuthButton.addEventListener(\"click\", (e) => {\n e.preventDefault();\n chrome.storage.local.get(\n { [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN },\n (result) => {\n const onyxDomain = result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n chrome.runtime.sendMessage(\n { action: ACTIONS.CLOSE_SIDE_PANEL },\n () => {\n if (chrome.runtime.lastError) {\n console.error(\n \"Error closing side panel:\",\n chrome.runtime.lastError,\n );\n }\n chrome.tabs.create(\n {\n url: `${onyxDomain}/auth/login`,\n active: true,\n },\n (_) => {\n if (chrome.runtime.lastError) {\n console.error(\n \"Error opening auth tab:\",\n chrome.runtime.lastError,\n );\n }\n },\n );\n },\n );\n },\n );\n });\n }\n}\n\nexport function showAuthModal() {\n if (!authModal) {\n initAuthModal();\n }\n if (authModal) {\n authModal.style.display = \"flex\";\n authModal.style.zIndex = \"9999\";\n document.body.style.overflow = \"hidden\";\n }\n}\n\nexport function hideAuthModal() {\n if (authModal) {\n authModal.style.display = \"none\";\n document.body.style.overflow = \"auto\";\n }\n}\n" }, { "path": "extensions/chrome/src/utils/selection-icon.js", "content": "(function () {\n const OPEN_SIDE_PANEL_WITH_INPUT = \"openSidePanelWithInput\";\n\n let selectionIcon = null;\n let currentSelectedText = \"\";\n\n function createSelectionIcon() {\n if (selectionIcon) return;\n\n selectionIcon = document.createElement(\"div\");\n selectionIcon.id = \"onyx-selection-icon\";\n\n const img = document.createElement(\"img\");\n img.src = chrome.runtime.getURL(\"public/icon32.png\");\n img.alt = \"Search with Onyx\";\n\n selectionIcon.appendChild(img);\n document.body.appendChild(selectionIcon);\n\n selectionIcon.addEventListener(\"mousedown\", handleIconClick);\n }\n\n function showIcon(text) {\n if (!selectionIcon) {\n createSelectionIcon();\n }\n\n currentSelectedText = text;\n\n const selection = window.getSelection();\n if (!selection.rangeCount) return;\n\n const range = selection.getRangeAt(0);\n const rect = range.getBoundingClientRect();\n\n const iconSize = 32;\n const offset = 4;\n\n let posX = rect.right + offset;\n let posY = rect.bottom + offset;\n\n if (posX + iconSize > window.innerWidth) {\n posX = rect.left - iconSize - offset;\n }\n if (posY + iconSize > window.innerHeight) {\n posY = rect.top - iconSize - offset;\n }\n\n posX = Math.max(\n offset,\n Math.min(posX, window.innerWidth - iconSize - offset),\n );\n posY = Math.max(\n offset,\n Math.min(posY, window.innerHeight - iconSize - offset),\n );\n\n selectionIcon.style.left = `${posX}px`;\n selectionIcon.style.top = `${posY}px`;\n selectionIcon.classList.add(\"visible\");\n }\n\n function hideIcon() {\n if (selectionIcon) {\n selectionIcon.classList.remove(\"visible\");\n }\n currentSelectedText = \"\";\n }\n\n function handleIconClick(e) {\n e.preventDefault();\n e.stopPropagation();\n\n const textToSend = currentSelectedText;\n\n if (textToSend) {\n chrome.runtime.sendMessage(\n {\n action: OPEN_SIDE_PANEL_WITH_INPUT,\n selectedText: textToSend,\n pageUrl: window.location.href,\n },\n (response) => {\n if (chrome.runtime.lastError) {\n console.error(\n \"[Onyx] Error sending message:\",\n chrome.runtime.lastError.message,\n );\n } else {\n }\n },\n );\n }\n\n hideIcon();\n }\n\n document.addEventListener(\"mouseup\", (e) => {\n if (\n e.target.id === \"onyx-selection-icon\" ||\n e.target.closest(\"#onyx-selection-icon\")\n ) {\n return;\n }\n\n setTimeout(() => {\n const selection = window.getSelection();\n const selectedText = selection.toString().trim();\n\n if (selectedText && selectedText.length > 0) {\n showIcon(selectedText);\n } else {\n hideIcon();\n }\n }, 10);\n });\n\n document.addEventListener(\"mousedown\", (e) => {\n if (\n e.target.id !== \"onyx-selection-icon\" &&\n !e.target.closest(\"#onyx-selection-icon\")\n ) {\n const selection = window.getSelection();\n const selectedText = selection.toString().trim();\n if (!selectedText) {\n hideIcon();\n }\n }\n });\n\n document.addEventListener(\n \"scroll\",\n () => {\n hideIcon();\n },\n true,\n );\n\n document.addEventListener(\"selectionchange\", () => {\n const selection = window.getSelection();\n const selectedText = selection.toString().trim();\n if (!selectedText) {\n hideIcon();\n }\n });\n\n if (document.readyState === \"loading\") {\n document.addEventListener(\"DOMContentLoaded\", createSelectionIcon);\n } else {\n createSelectionIcon();\n }\n})();\n" }, { "path": "extensions/chrome/src/utils/storage.js", "content": "import {\n DEFAULT_ONYX_DOMAIN,\n CHROME_SPECIFIC_STORAGE_KEYS,\n} from \"./constants.js\";\n\nexport async function getOnyxDomain() {\n const result = await chrome.storage.local.get({\n [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: DEFAULT_ONYX_DOMAIN,\n });\n return result[CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN];\n}\n\nexport function setOnyxDomain(domain, callback) {\n chrome.storage.local.set(\n { [CHROME_SPECIFIC_STORAGE_KEYS.ONYX_DOMAIN]: domain },\n callback,\n );\n}\n\nexport function getOnyxDomainSync() {\n return new Promise((resolve) => {\n getOnyxDomain(resolve);\n });\n}\n" }, { "path": "profiling/grafana/dashboards/onyx/opensearch-search-latency.json", "content": "{\n \"annotations\": {\n \"list\": [\n {\n \"builtIn\": 1,\n \"datasource\": { \"type\": \"grafana\", \"uid\": \"-- Grafana --\" },\n \"enable\": true,\n \"hide\": true,\n \"iconColor\": \"rgba(0, 211, 255, 1)\",\n \"name\": \"Annotations & Alerts\",\n \"type\": \"dashboard\"\n }\n ]\n },\n \"editable\": true,\n \"fiscalYearStartMonth\": 0,\n \"graphTooltip\": 1,\n \"id\": null,\n \"links\": [],\n \"liveNow\": true,\n \"panels\": [\n {\n \"title\": \"Client-Side Search Latency (P50 / P95 / P99)\",\n \"description\": \"End-to-end latency as measured by the Python client, including network round-trip and serialization overhead.\",\n \"type\": \"timeseries\",\n \"gridPos\": { \"h\": 10, \"w\": 12, \"x\": 0, \"y\": 0 },\n \"id\": 1,\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"fieldConfig\": {\n \"defaults\": {\n \"color\": { \"mode\": \"palette-classic\" },\n \"custom\": {\n \"axisBorderShow\": false,\n \"axisCenteredZero\": false,\n \"axisLabel\": \"seconds\",\n \"axisPlacement\": \"auto\",\n \"drawStyle\": \"line\",\n \"fillOpacity\": 0,\n \"gradientMode\": \"none\",\n \"lineInterpolation\": \"smooth\",\n \"lineWidth\": 2,\n \"pointSize\": 5,\n \"scaleDistribution\": { \"type\": \"linear\" },\n \"showPoints\": \"never\",\n \"spanNulls\": false,\n \"stacking\": { \"group\": \"A\", \"mode\": \"none\" },\n \"thresholdsStyle\": { \"mode\": \"dashed\" }\n },\n \"thresholds\": {\n \"mode\": \"absolute\",\n \"steps\": [\n { \"color\": \"green\", \"value\": null },\n { \"color\": \"yellow\", \"value\": 0.5 },\n { \"color\": \"red\", \"value\": 2.0 }\n ]\n },\n \"unit\": \"s\",\n \"min\": 0\n },\n \"overrides\": []\n },\n \"targets\": [\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"P50\",\n \"refId\": \"A\"\n },\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.95, sum by (le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"P95\",\n \"refId\": \"B\"\n },\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.99, sum by (le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"P99\",\n \"refId\": \"C\"\n }\n ]\n },\n {\n \"title\": \"Server-Side Search Latency (P50 / P95 / P99)\",\n \"description\": \"OpenSearch server-side execution time from the 'took' field in the response. Does not include network or client-side overhead.\",\n \"type\": \"timeseries\",\n \"gridPos\": { \"h\": 10, \"w\": 12, \"x\": 12, \"y\": 0 },\n \"id\": 2,\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"fieldConfig\": {\n \"defaults\": {\n \"color\": { \"mode\": \"palette-classic\" },\n \"custom\": {\n \"axisBorderShow\": false,\n \"axisCenteredZero\": false,\n \"axisLabel\": \"seconds\",\n \"axisPlacement\": \"auto\",\n \"drawStyle\": \"line\",\n \"fillOpacity\": 0,\n \"gradientMode\": \"none\",\n \"lineInterpolation\": \"smooth\",\n \"lineWidth\": 2,\n \"pointSize\": 5,\n \"scaleDistribution\": { \"type\": \"linear\" },\n \"showPoints\": \"never\",\n \"spanNulls\": false,\n \"stacking\": { \"group\": \"A\", \"mode\": \"none\" },\n \"thresholdsStyle\": { \"mode\": \"dashed\" }\n },\n \"thresholds\": {\n \"mode\": \"absolute\",\n \"steps\": [\n { \"color\": \"green\", \"value\": null },\n { \"color\": \"yellow\", \"value\": 0.5 },\n { \"color\": \"red\", \"value\": 2.0 }\n ]\n },\n \"unit\": \"s\",\n \"min\": 0\n },\n \"overrides\": []\n },\n \"targets\": [\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_server_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"P50\",\n \"refId\": \"A\"\n },\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.95, sum by (le) (rate(onyx_opensearch_search_server_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"P95\",\n \"refId\": \"B\"\n },\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.99, sum by (le) (rate(onyx_opensearch_search_server_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"P99\",\n \"refId\": \"C\"\n }\n ]\n },\n {\n \"title\": \"Client-Side Latency by Search Type (P95)\",\n \"description\": \"P95 client-side latency broken down by search type (hybrid, keyword, semantic, random, doc_id_retrieval).\",\n \"type\": \"timeseries\",\n \"gridPos\": { \"h\": 10, \"w\": 12, \"x\": 0, \"y\": 10 },\n \"id\": 3,\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"fieldConfig\": {\n \"defaults\": {\n \"color\": { \"mode\": \"palette-classic\" },\n \"custom\": {\n \"axisBorderShow\": false,\n \"axisCenteredZero\": false,\n \"axisLabel\": \"seconds\",\n \"axisPlacement\": \"auto\",\n \"drawStyle\": \"line\",\n \"fillOpacity\": 0,\n \"gradientMode\": \"none\",\n \"lineInterpolation\": \"smooth\",\n \"lineWidth\": 2,\n \"pointSize\": 5,\n \"scaleDistribution\": { \"type\": \"linear\" },\n \"showPoints\": \"never\",\n \"spanNulls\": false,\n \"stacking\": { \"group\": \"A\", \"mode\": \"none\" },\n \"thresholdsStyle\": { \"mode\": \"off\" }\n },\n \"unit\": \"s\",\n \"min\": 0\n },\n \"overrides\": []\n },\n \"targets\": [\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.95, sum by (search_type, le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"{{ search_type }}\",\n \"refId\": \"A\"\n }\n ]\n },\n {\n \"title\": \"Search Throughput by Type\",\n \"description\": \"Searches per second broken down by search type.\",\n \"type\": \"timeseries\",\n \"gridPos\": { \"h\": 10, \"w\": 12, \"x\": 12, \"y\": 10 },\n \"id\": 4,\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"fieldConfig\": {\n \"defaults\": {\n \"color\": { \"mode\": \"palette-classic\" },\n \"custom\": {\n \"axisBorderShow\": false,\n \"axisCenteredZero\": false,\n \"axisLabel\": \"searches/s\",\n \"axisPlacement\": \"auto\",\n \"drawStyle\": \"line\",\n \"fillOpacity\": 0,\n \"gradientMode\": \"none\",\n \"lineInterpolation\": \"smooth\",\n \"lineWidth\": 2,\n \"pointSize\": 5,\n \"scaleDistribution\": { \"type\": \"linear\" },\n \"showPoints\": \"never\",\n \"spanNulls\": false,\n \"stacking\": { \"group\": \"A\", \"mode\": \"normal\" },\n \"thresholdsStyle\": { \"mode\": \"off\" }\n },\n \"unit\": \"ops\",\n \"min\": 0\n },\n \"overrides\": []\n },\n \"targets\": [\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"sum by (search_type) (rate(onyx_opensearch_search_total[5m]))\",\n \"legendFormat\": \"{{ search_type }}\",\n \"refId\": \"A\"\n }\n ]\n },\n {\n \"title\": \"Concurrent Searches In Progress\",\n \"description\": \"Number of OpenSearch searches currently in flight, broken down by search type. Summed across all instances.\",\n \"type\": \"timeseries\",\n \"gridPos\": { \"h\": 10, \"w\": 12, \"x\": 0, \"y\": 20 },\n \"id\": 5,\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"fieldConfig\": {\n \"defaults\": {\n \"color\": { \"mode\": \"palette-classic\" },\n \"custom\": {\n \"axisBorderShow\": false,\n \"axisCenteredZero\": false,\n \"axisLabel\": \"searches\",\n \"axisPlacement\": \"auto\",\n \"drawStyle\": \"line\",\n \"fillOpacity\": 0,\n \"gradientMode\": \"none\",\n \"lineInterpolation\": \"smooth\",\n \"lineWidth\": 2,\n \"pointSize\": 5,\n \"scaleDistribution\": { \"type\": \"linear\" },\n \"showPoints\": \"never\",\n \"spanNulls\": false,\n \"stacking\": { \"group\": \"A\", \"mode\": \"normal\" },\n \"thresholdsStyle\": { \"mode\": \"off\" }\n },\n \"min\": 0\n },\n \"overrides\": []\n },\n \"targets\": [\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"sum by (search_type) (onyx_opensearch_searches_in_progress)\",\n \"legendFormat\": \"{{ search_type }}\",\n \"refId\": \"A\"\n }\n ]\n },\n {\n \"title\": \"Client vs Server Latency Overhead (P50)\",\n \"description\": \"Difference between client-side and server-side P50 latency. Reveals network, serialization, and untracked OpenSearch overhead.\",\n \"type\": \"timeseries\",\n \"gridPos\": { \"h\": 10, \"w\": 12, \"x\": 12, \"y\": 20 },\n \"id\": 6,\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"fieldConfig\": {\n \"defaults\": {\n \"color\": { \"mode\": \"palette-classic\" },\n \"custom\": {\n \"axisBorderShow\": false,\n \"axisCenteredZero\": false,\n \"axisLabel\": \"seconds\",\n \"axisPlacement\": \"auto\",\n \"drawStyle\": \"line\",\n \"fillOpacity\": 0,\n \"gradientMode\": \"none\",\n \"lineInterpolation\": \"smooth\",\n \"lineWidth\": 2,\n \"pointSize\": 5,\n \"scaleDistribution\": { \"type\": \"linear\" },\n \"showPoints\": \"never\",\n \"spanNulls\": false,\n \"stacking\": { \"group\": \"A\", \"mode\": \"none\" },\n \"thresholdsStyle\": { \"mode\": \"off\" }\n },\n \"unit\": \"s\",\n \"min\": 0\n },\n \"overrides\": []\n },\n \"targets\": [\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m]))) - histogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_server_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"Client - Server overhead (P50)\",\n \"refId\": \"A\"\n },\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_client_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"Client P50\",\n \"refId\": \"B\"\n },\n {\n \"datasource\": { \"type\": \"prometheus\", \"uid\": \"${DS_PROMETHEUS}\" },\n \"expr\": \"histogram_quantile(0.5, sum by (le) (rate(onyx_opensearch_search_server_duration_seconds_bucket[5m])))\",\n \"legendFormat\": \"Server P50\",\n \"refId\": \"C\"\n }\n ]\n }\n ],\n \"refresh\": \"5s\",\n \"schemaVersion\": 37,\n \"style\": \"dark\",\n \"tags\": [\"onyx\", \"opensearch\", \"search\", \"latency\"],\n \"templating\": {\n \"list\": [\n {\n \"current\": {\n \"text\": \"Prometheus\",\n \"value\": \"prometheus\"\n },\n \"includeAll\": false,\n \"name\": \"DS_PROMETHEUS\",\n \"options\": [],\n \"query\": \"prometheus\",\n \"refresh\": 1,\n \"type\": \"datasource\"\n }\n ]\n },\n \"time\": { \"from\": \"now-60m\", \"to\": \"now\" },\n \"timepicker\": {\n \"refresh_intervals\": [\"5s\", \"10s\", \"30s\", \"1m\"]\n },\n \"timezone\": \"\",\n \"title\": \"Onyx OpenSearch Search Latency\",\n \"uid\": \"onyx-opensearch-search-latency\",\n \"version\": 0,\n \"weekStart\": \"\"\n}\n" }, { "path": "pyproject.toml", "content": "[build-system]\nrequires = [\"setuptools>=61\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"onyx\"\nversion = \"0.0.0\"\nrequires-python = \">=3.11\"\n# Shared dependencies between backend and model_server\ndependencies = [\n \"aioboto3==15.1.0\",\n \"cohere==5.6.1\",\n \"fastapi==0.133.1\",\n \"google-genai==1.52.0\",\n \"litellm==1.81.6\",\n \"openai==2.14.0\",\n \"pydantic==2.11.7\",\n \"prometheus_client>=0.21.1\",\n \"prometheus_fastapi_instrumentator==7.1.0\",\n \"retry==0.9.2\", # This pulls in py which is in CVE-2022-42969, must remove py from image\n \"sentry-sdk==2.14.0\",\n \"uvicorn==0.35.0\",\n \"voyageai==0.2.3\",\n \"brotli>=1.2.0\",\n \"claude-agent-sdk>=0.1.19\",\n \"agent-client-protocol>=0.7.1\",\n \"discord-py==2.4.0\",\n \"kubernetes>=31.0.0\",\n]\n\n[project.optional-dependencies]\n# Main backend application dependencies\nbackend = [\n \"aiohttp==3.13.4\",\n \"alembic==1.10.4\",\n \"asyncpg==0.30.0\",\n \"atlassian-python-api==3.41.16\",\n \"azure-cognitiveservices-speech==1.38.0\",\n \"beautifulsoup4==4.12.3\",\n \"boto3==1.39.11\",\n \"boto3-stubs[s3]==1.39.11\",\n \"celery==5.5.1\",\n \"chardet==5.2.0\",\n \"chonkie==1.0.10\",\n \"dask==2026.1.1\",\n \"ddtrace==3.10.0\",\n \"discord.py==2.4.0\",\n \"distributed==2026.1.1\",\n \"fastapi-users==15.0.4\",\n \"fastapi-users-db-sqlalchemy==7.0.0\",\n \"fastapi-limiter==0.1.6\",\n \"fastmcp==3.2.0\",\n \"filelock==3.20.3\",\n \"google-api-python-client==2.86.0\",\n \"google-auth-httplib2==0.1.0\",\n \"google-auth-oauthlib==1.0.0\",\n # GPT4All library has issues running on Macs and python:3.11.4-slim-bookworm\n # will reintroduce this when library version catches up\n # \"gpt4all==2.0.2\",\n \"httpcore==1.0.9\",\n \"httpx[http2]==0.28.1\",\n \"httpx-oauth==0.15.1\",\n \"huggingface-hub==0.35.3\",\n \"inflection==0.5.1\",\n \"jira==3.10.5\",\n \"jsonref==1.1.0\",\n \"kubernetes==31.0.0\",\n \"trafilatura==1.12.2\",\n \"langchain-core==1.2.22\",\n \"lazy_imports==1.0.1\",\n \"lxml==5.3.0\",\n \"Mako==1.2.4\",\n \"markitdown[pdf, docx, pptx, xlsx, xls]==0.1.2\",\n \"mcp[cli]==1.26.0\",\n \"msal==1.34.0\",\n \"msoffcrypto-tool==5.4.2\",\n \"Office365-REST-Python-Client==2.6.2\",\n \"oauthlib==3.2.2\",\n # NOTE: This is frozen to avoid https://foss.heptapod.net/openpyxl/openpyxl/-/issues/2147\n \"openpyxl==3.0.10\",\n \"opensearch-py==3.0.0\",\n \"passlib==1.7.4\",\n \"playwright==1.55.0\",\n \"psutil==7.1.3\",\n \"psycopg2-binary==2.9.9\",\n \"puremagic==1.28\",\n \"pyairtable==3.0.1\",\n \"pycryptodome==3.19.1\",\n \"PyGithub==2.5.0\",\n \"pympler==1.1\",\n \"python-dateutil==2.8.2\",\n \"python-gitlab==5.6.0\",\n \"python-pptx==0.6.23\",\n \"pypandoc_binary==1.16.2\",\n \"pypdf==6.9.2\",\n \"pytest-mock==3.12.0\",\n \"pytest-playwright==0.7.0\",\n \"python-docx==1.1.2\",\n \"python-dotenv==1.1.1\",\n \"python-multipart==0.0.22\",\n \"pywikibot==9.0.0\",\n \"redis==5.0.8\",\n \"requests==2.33.0\",\n \"requests-oauthlib==1.3.1\",\n \"rfc3986==1.5.0\",\n \"simple-salesforce==1.12.6\",\n \"slack-sdk==3.20.2\",\n \"SQLAlchemy[mypy]==2.0.15\",\n \"starlette==0.49.3\",\n \"supervisor==4.3.0\",\n \"RapidFuzz==3.13.0\",\n \"tiktoken==0.7.0\",\n \"timeago==1.0.16\",\n \"types-openpyxl==3.0.4.7\",\n \"unstructured==0.18.27\",\n \"unstructured-client==0.42.6\",\n \"zulip==0.8.2\",\n \"hubspot-api-client==11.1.0\",\n \"asana==5.0.8\",\n \"dropbox==12.0.2\",\n \"shapely==2.0.6\",\n \"stripe==10.12.0\",\n \"urllib3==2.6.3\",\n \"mistune==3.2.0\",\n \"sendgrid==6.12.5\",\n \"exa_py==1.15.4\",\n \"braintrust==0.3.9\",\n \"langfuse==3.10.0\",\n \"nest_asyncio==1.6.0\",\n \"openinference-instrumentation==0.1.42\",\n \"opentelemetry-proto>=1.39.0\",\n \"python3-saml==1.15.0\",\n \"xmlsec==1.3.14\",\n]\n\n# Dev tools\ndev = [\n \"black==25.1.0\",\n \"celery-types==0.19.0\",\n \"faker==40.1.2\",\n \"hatchling==1.28.0\",\n \"ipykernel==6.29.5\",\n \"manygo==0.2.0\",\n \"matplotlib==3.10.8\",\n \"mypy-extensions==1.0.0\",\n \"mypy==1.13.0\",\n \"onyx-devtools==0.7.2\",\n \"openapi-generator-cli==7.17.0\",\n \"pandas-stubs~=2.3.3\",\n \"pre-commit==3.2.2\",\n \"pytest-alembic==0.12.1\",\n \"pytest-asyncio==1.3.0\",\n \"pytest-dotenv==0.5.2\",\n \"pytest-repeat==0.9.4\",\n \"pytest-xdist==3.8.0\",\n \"pytest==8.3.5\",\n \"release-tag==0.5.2\",\n \"reorder-python-imports-black==3.14.0\",\n \"ruff==0.12.0\",\n \"types-beautifulsoup4==4.12.0.3\",\n \"types-html5lib==1.1.11.13\",\n \"types-oauthlib==3.2.0.9\",\n \"types-passlib==1.7.7.20240106\",\n \"types-Pillow==10.2.0.20240822\",\n \"types-psutil==7.1.3.20251125\",\n \"types-psycopg2==2.9.21.10\",\n \"types-python-dateutil==2.8.19.13\",\n \"types-PyYAML==6.0.12.11\",\n \"types-pytz==2023.3.1.1\",\n \"types-regex==2023.3.23.1\",\n \"types-requests==2.32.0.20250328\",\n \"types-retry==0.9.9.3\",\n \"types-setuptools==68.0.0.3\",\n \"zizmor==1.18.0\",\n]\n\n# Enterprise Edition features\nee = [\n \"posthog==3.7.4\",\n]\n\n# Model server specific dependencies (ML packages)\nmodel_server = [\n \"accelerate==1.6.0\",\n \"einops==0.8.1\",\n \"numpy==2.4.1\",\n \"safetensors==0.5.3\",\n \"sentence-transformers==4.0.2\",\n \"torch==2.9.1\",\n \"transformers==4.53.0\",\n \"sentry-sdk[fastapi,celery,starlette]==2.14.0\",\n]\n\n[tool.mypy]\nplugins = \"sqlalchemy.ext.mypy.plugin\"\nmypy_path = \"backend\"\nexplicit_package_bases = true\ndisallow_untyped_defs = true\nwarn_unused_ignores = true\nenable_error_code = [\"possibly-undefined\"]\nstrict_equality = true\n# Patterns match paths whether mypy is run from backend/ (CI) or repo root (e.g. VS Code extension with target ./backend)\nexclude = [\n \"(?:^|/)generated/\",\n \"(?:^|/)\\\\.venv/\",\n \"(?:^|/)onyx/server/features/build/sandbox/kubernetes/docker/skills/\",\n \"(?:^|/)onyx/server/features/build/sandbox/kubernetes/docker/templates/\",\n]\n\n[[tool.mypy.overrides]]\nmodule = \"alembic.versions.*\"\ndisable_error_code = [\"var-annotated\"]\n\n[[tool.mypy.overrides]]\nmodule = \"alembic_tenants.versions.*\"\ndisable_error_code = [\"var-annotated\"]\n\n[[tool.mypy.overrides]]\nmodule = \"generated.*\"\nfollow_imports = \"silent\"\nignore_errors = true\n\n[[tool.mypy.overrides]]\nmodule = \"transformers.*\"\nfollow_imports = \"skip\"\nignore_errors = true\n\n[tool.uv.workspace]\nmembers = [\"backend\", \"tools/ods\"]\n\n[tool.basedpyright]\ninclude = [\"backend\"]\nexclude = [\"backend/generated\", \"backend/onyx/server/features/build/sandbox/kubernetes/docker/skills/pptx\", \"backend/onyx/server/features/build/sandbox/kubernetes/docker/templates/venv\"]\ntypeCheckingMode = \"off\"\n\n[tool.ruff]\nline-length = 130\ntarget-version = \"py311\"\n\n[tool.ruff.lint]\nignore = [\n \"E501\", # Long lines are handled by Black.\n]\nselect = [\n \"ARG\",\n \"E\",\n \"F\",\n \"S324\",\n \"W\",\n]\n\n[tool.setuptools.packages.find]\nwhere = [\"backend\"]\ninclude = [\"onyx*\", \"tests*\"]\n" }, { "path": "web/.dockerignore", "content": "node_modules\n.next\n/tests/\n\n# Explicitly include src/app/build (overrides .gitignore /build pattern)\n!src/app/build\n" }, { "path": "web/.eslintrc.json", "content": "{\n \"extends\": \"next/core-web-vitals\",\n \"plugins\": [\"unused-imports\"],\n \"rules\": {\n \"@next/next/no-img-element\": \"off\",\n \"react-hooks/exhaustive-deps\": \"off\",\n \"no-unused-vars\": \"off\",\n \"@typescript-eslint/no-unused-vars\": \"off\",\n \"unused-imports/no-unused-imports\": \"warn\",\n \"unused-imports/no-unused-vars\": [\n \"warn\",\n {\n \"vars\": \"all\",\n \"varsIgnorePattern\": \"^_\",\n \"args\": \"after-used\",\n \"argsIgnorePattern\": \"^_\",\n \"ignoreRestSiblings\": true\n }\n ]\n }\n}\n" }, { "path": "web/.gitignore", "content": "# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.\n.env.sentry-build-plugin\n\n# dependencies\nnode_modules\n/.pnp\n.pnp.js\n\n# testing\n/coverage\n\n# next.js\n/.next/\n/out/\n\n# production\n/build\n\n# misc\n.DS_Store\n*.pem\n\n# debug\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\n.pnpm-debug.log*\n\n# local env files\n.env*.local\n\n# vercel\n.vercel\n\n# typescript\n*.tsbuildinfo\nnext-env.d.ts\n\n# playwright testing temp files\n/admin*_auth.json\n/worker*_auth.json\n/user_auth.json\n/build-archive.log\n/test-results\n/output/\n\n# generated clients ... in particular, the API to the Onyx backend itself!\n/src/lib/generated\n.jest-cache\n\n# storybook\nstorybook-static\n" }, { "path": "web/.prettierignore", "content": "**/.git\n**/.svn\n**/.hg\n**/node_modules\n**/.next\n**/.vscode" }, { "path": "web/.prettierrc.json", "content": "{\n \"trailingComma\": \"es5\"\n}\n" }, { "path": "web/.storybook/Introduction.mdx", "content": "import { Meta } from \"@storybook/blocks\";\n\n\n\n# Onyx Storybook\n\nA living catalog for browsing, testing, and documenting Onyx UI components in isolation.\n\n---\n\n## What is this?\n\nThis Storybook contains interactive examples of every reusable UI component in the Onyx frontend. Each component has a dedicated page with:\n\n- **Live demos** you can interact with directly\n- **Controls** to tweak props and see how the component responds\n- **Auto-generated docs** showing the full props API\n- **Dark mode toggle** in the toolbar to preview both themes\n\n---\n\n## Navigating Storybook\n\n### Sidebar\n\nThe left sidebar organizes components by layer:\n\n- **opal/core** — Low-level primitives (`Interactive`, `Hoverable`)\n- **opal/components** — Design system atoms (`Button`, `OpenButton`, `Tag`)\n- **Layouts** — Structural layouts (`Content`, `ContentAction`, `IllustrationContent`)\n- **refresh-components** — App-level components (inputs, modals, tables, text, etc.)\n\nClick any component to see its stories. Click **Docs** to see the auto-generated props table.\n\n### Controls panel\n\nAt the bottom of each story, the **Controls** panel lets you change props in real time. Toggle booleans, pick from enums, type in strings — the preview updates instantly.\n\n### Theme toggle\n\nUse the paint roller icon in the top toolbar to switch between **light** and **dark** mode. All components use CSS variables that automatically adapt.\n\n---\n\n## Running locally\n\n```bash\ncd web\nnpm run storybook # dev server on :6006\nnpm run storybook:build # static build to storybook-static/\n```\n\n---\n\n## Adding a new story\n\nStories are **co-located** next to their component:\n\n```\nlib/opal/src/components/buttons/Button/\n├── components.tsx ← the component\n├── Button.stories.tsx ← the story\n├── styles.css\n└── README.md\n```\n\n### Minimal template\n\n```tsx\nimport type { Meta, StoryObj } from \"@storybook/react\";\nimport { MyComponent } from \"./MyComponent\";\n\nconst meta: Meta = {\n title: \"opal/components/MyComponent\", // sidebar path\n component: MyComponent,\n tags: [\"autodocs\"], // auto-generate docs page\n};\n\nexport default meta;\ntype Story = StoryObj;\n\nexport const Default: Story = {\n args: {\n title: \"Hello\",\n },\n};\n\nexport const WithCustomLayout: Story = {\n render: () => (\n
    \n \n \n
    \n ),\n};\n```\n\n### Conventions\n\n- **Title format:** `opal/core/Name`, `opal/components/Name`, `Layouts/Name`, or `refresh-components/Name`\n- **Tags:** Add `tags: [\"autodocs\"]` to auto-generate a docs page from props\n- **Decorators:** If your component needs `TooltipPrimitive.Provider` (anything with tooltips), add it as a decorator\n- **Layout:** Use `parameters: { layout: \"fullscreen\" }` for modals/popovers that use portals\n\n---\n\n## Deployment\n\nProduction builds deploy to [onyx-storybook.vercel.app](https://onyx-storybook.vercel.app) automatically when PRs touching component files merge to `main`.\n\nMonitored paths:\n\n- `web/lib/opal/**`\n- `web/src/refresh-components/**`\n- `web/.storybook/**`\n" }, { "path": "web/.storybook/README.md", "content": "# Onyx Storybook\n\nStorybook is an isolated development environment for UI components. It renders each component in a standalone \"story\" outside of the main app, so you can visually verify appearance, interact with props, and catch regressions without navigating through the full application.\n\nThe Onyx Storybook covers the full component library — from low-level `@opal/core` primitives up through `refresh-components` — giving designers and engineers a shared reference for every visual state.\n\n**Production:** [onyx-storybook.vercel.app](https://onyx-storybook.vercel.app)\n\n## Running Locally\n\n```bash\ncd web\nnpm run storybook # dev server on http://localhost:6006\nnpm run storybook:build # static build to storybook-static/\n```\n\nThe dev server hot-reloads when you edit a component or story file.\n\n## Writing Stories\n\nStories are **co-located** next to their component source:\n\n```\nlib/opal/src/core/interactive/\n├── components.tsx ← the component\n├── Interactive.stories.tsx ← the story\n└── styles.css\n\nsrc/refresh-components/buttons/\n├── Button.tsx\n└── Button.stories.tsx\n```\n\n### Minimal Template\n\n```tsx\nimport type { Meta, StoryObj } from \"@storybook/react\";\nimport { MyComponent } from \"./MyComponent\";\n\nconst meta: Meta = {\n title: \"Category/MyComponent\", // sidebar path\n component: MyComponent,\n tags: [\"autodocs\"], // generates a docs page from props\n};\n\nexport default meta;\ntype Story = StoryObj;\n\nexport const Default: Story = {\n args: { label: \"Hello\" },\n};\n```\n\n### Conventions\n\n- **Title format:** `Core/Name`, `Components/Name`, `Layouts/Name`, or `refresh-components/category/Name`\n- **Tags:** Add `tags: [\"autodocs\"]` to auto-generate a props docs page\n- **Decorators:** Components that use Radix tooltips need a `TooltipPrimitive.Provider` decorator\n- **Layout:** Use `parameters: { layout: \"fullscreen\" }` for modals/popovers that use portals\n\n## Dark Mode\n\nUse the theme toggle (paint roller icon) in the Storybook toolbar to switch between light and dark modes. This adds/removes the `dark` class on the preview body, matching the app's `darkMode: \"class\"` Tailwind config. All color tokens from `colors.css` adapt automatically.\n\n## Deployment\n\nThe production Storybook is deployed as a static site on Vercel. The build runs `npm run storybook:build` which outputs to `storybook-static/`, and Vercel serves that directory.\n\nDeploys are triggered on merges to `main` when files in `web/lib/opal/`, `web/src/refresh-components/`, or `web/.storybook/` change.\n\n## Component Layers\n\nThe sidebar organizes components by their layer in the design system:\n\n| Layer | Path | Examples |\n|-------|------|----------|\n| **Core** | `lib/opal/src/core/` | Interactive, Hoverable |\n| **Components** | `lib/opal/src/components/` | Button, OpenButton, Tag |\n| **Layouts** | `lib/opal/src/layouts/` | Content, ContentAction, IllustrationContent |\n| **refresh-components** | `src/refresh-components/` | Inputs, tables, modals, text, cards, tiles, etc. |\n" }, { "path": "web/.storybook/main.ts", "content": "import type { StorybookConfig } from \"@storybook/react-vite\";\nimport path from \"path\";\n\nconst config: StorybookConfig = {\n stories: [\n \"./*.mdx\",\n \"../lib/opal/src/**/*.stories.@(ts|tsx)\",\n \"../src/refresh-components/**/*.stories.@(ts|tsx)\",\n ],\n addons: [\"@storybook/addon-essentials\", \"@storybook/addon-themes\"],\n framework: {\n name: \"@storybook/react-vite\",\n options: {},\n },\n staticDirs: [\"../public\"],\n docs: {\n autodocs: \"tag\",\n },\n typescript: {\n reactDocgen: \"react-docgen-typescript\",\n },\n viteFinal: async (config) => {\n config.resolve = config.resolve ?? {};\n config.resolve.alias = {\n ...config.resolve.alias,\n \"@\": path.resolve(__dirname, \"../src\"),\n \"@opal\": path.resolve(__dirname, \"../lib/opal/src\"),\n \"@public\": path.resolve(__dirname, \"../public\"),\n // Next.js module stubs for Vite\n \"next/link\": path.resolve(__dirname, \"mocks/next-link.tsx\"),\n \"next/navigation\": path.resolve(__dirname, \"mocks/next-navigation.tsx\"),\n \"next/image\": path.resolve(__dirname, \"mocks/next-image.tsx\"),\n };\n\n // Process CSS with Tailwind via PostCSS\n config.css = config.css ?? {};\n config.css.postcss = path.resolve(__dirname, \"..\");\n\n return config;\n },\n};\n\nexport default config;\n" }, { "path": "web/.storybook/mocks/next-image.tsx", "content": "import React from \"react\";\n\ninterface ImageProps {\n src: string;\n alt: string;\n width?: number;\n height?: number;\n fill?: boolean;\n [key: string]: unknown;\n}\n\nfunction Image({ src, alt, width, height, fill, ...props }: ImageProps) {\n const fillStyle: React.CSSProperties = fill\n ? { position: \"absolute\", inset: 0, width: \"100%\", height: \"100%\" }\n : {};\n return (\n )}\n src={src}\n alt={alt}\n width={fill ? undefined : width}\n height={fill ? undefined : height}\n style={{ ...(props.style as React.CSSProperties), ...fillStyle }}\n />\n );\n}\n\nexport default Image;\n" }, { "path": "web/.storybook/mocks/next-link.tsx", "content": "import React from \"react\";\n\ninterface LinkProps {\n href: string;\n children: React.ReactNode;\n [key: string]: unknown;\n}\n\nfunction Link({\n href,\n children,\n prefetch: _prefetch,\n scroll: _scroll,\n shallow: _shallow,\n replace: _replace,\n passHref: _passHref,\n locale: _locale,\n legacyBehavior: _legacyBehavior,\n ...props\n}: LinkProps) {\n return (\n \n {children}\n \n );\n}\n\nexport default Link;\n" }, { "path": "web/.storybook/mocks/next-navigation.tsx", "content": "export function useRouter() {\n return {\n push: (_url: string) => {},\n replace: (_url: string) => {},\n back: () => {},\n forward: () => {},\n refresh: () => {},\n prefetch: (_url: string) => Promise.resolve(),\n };\n}\n\nexport function usePathname() {\n return \"/\";\n}\n\nexport function useSearchParams() {\n return new URLSearchParams() as ReadonlyURLSearchParams;\n}\n\nexport function useParams() {\n return {};\n}\n\nexport function redirect(_url: string): never {\n throw new Error(\"redirect() called in Storybook\");\n}\n\nexport function notFound(): never {\n throw new Error(\"notFound() called in Storybook\");\n}\n" }, { "path": "web/.storybook/preview-head.html", "content": "\n\n\n" }, { "path": "web/.storybook/preview.ts", "content": "import type { Preview } from \"@storybook/react\";\nimport { withThemeByClassName } from \"@storybook/addon-themes\";\nimport \"../src/app/globals.css\";\n\nconst preview: Preview = {\n parameters: {\n layout: \"centered\",\n backgrounds: { disable: true },\n controls: {\n matchers: {\n color: /(background|color)$/i,\n date: /Date$/i,\n },\n },\n },\n decorators: [\n withThemeByClassName({\n themes: {\n light: \"\",\n dark: \"dark\",\n },\n defaultTheme: \"light\",\n }),\n ],\n};\n\nexport default preview;\n" }, { "path": "web/@types/favicon-fetch.d.ts", "content": "declare module \"favicon-fetch\" {\n interface FaviconFetchOptions {\n uri: string;\n }\n\n function faviconFetch(options: FaviconFetchOptions): string | null;\n\n export default faviconFetch;\n}\n" }, { "path": "web/@types/images.d.ts", "content": "declare module \"*.png\" {\n const content: string;\n export default content;\n}\n\ndeclare module \"*.svg\" {\n const content: string;\n export default content;\n}\n\ndeclare module \"*.jpeg\" {\n const content: string;\n export default content;\n}\n\ndeclare module \"*.jpg\" {\n const content: string;\n export default content;\n}\n\ndeclare module \"*.gif\" {\n const content: string;\n export default content;\n}\n\ndeclare module \"*.webp\" {\n const content: string;\n export default content;\n}\n\ndeclare module \"*.ico\" {\n const content: string;\n export default content;\n}\n" }, { "path": "web/AGENTS.md", "content": "# Frontend Standards\n\nThis file is the single source of truth for frontend coding standards across all Onyx frontend\nprojects (including, but not limited to, `/web`, `/desktop`).\n\n# Components\n\nUI components are spread across several directories while the codebase migrates to Opal:\n\n- **`web/lib/opal/src/`** — The Opal design system. Preferred for all new components.\n- **`web/src/refresh-components/`** — Production components not yet migrated to Opal.\n- **`web/src/sections/`** — Feature-specific composite components (cards, modals, etc.).\n- **`web/src/layouts/`** — Page-level layout components (settings pages, etc.).\n\n**Do NOT use anything from `web/src/components/`** — this directory contains legacy components\nthat are being phased out. Always prefer Opal first; fall back to `refresh-components` only for\ncomponents not yet available in Opal.\n\n## Opal Layouts (`lib/opal/src/layouts/`)\n\nAll layout primitives are imported from `@opal/layouts`. They handle sizing, font selection, icon\nalignment, and optional inline editing.\n\n```typescript\nimport { Content, ContentAction, IllustrationContent } from \"@opal/layouts\";\n```\n\n### Content\n\n**Use this for any combination of icon + title + description.**\n\nA two-axis layout component that automatically routes to the correct internal layout\n(`ContentXl`, `ContentLg`, `ContentMd`, `ContentSm`) based on `sizePreset` and `variant`:\n\n| sizePreset | variant | Routes to | Layout |\n|---|---|---|---|\n| `headline` / `section` | `heading` | `ContentXl` | Icon on top (flex-col) |\n| `headline` / `section` | `section` | `ContentLg` | Icon inline (flex-row) |\n| `main-content` / `main-ui` / `secondary` | `section` / `heading` | `ContentMd` | Compact inline |\n| `main-content` / `main-ui` / `secondary` | `body` | `ContentSm` | Body text layout |\n\n```typescript\n\n```\n\n### ContentAction\n\n**Use this when a Content block needs right-side actions** (buttons, badges, icons, etc.).\n\nWraps `Content` and adds a `rightChildren` slot. Accepts all `Content` props plus:\n- `rightChildren`: `ReactNode` — actions rendered on the right\n- `paddingVariant`: `SizeVariant` — controls outer padding\n\n```typescript\nEdit}\n/>\n```\n\n### IllustrationContent\n\n**Use this for empty states, error pages, and informational placeholders.**\n\nA vertically-stacked, center-aligned layout that pairs a large illustration (7.5rem x 7.5rem)\nwith a title and optional description.\n\n```typescript\nimport SvgNoResult from \"@opal/illustrations/no-result\";\n\n\n```\n\nProps:\n- `illustration`: `IconFunctionComponent` — optional, from `@opal/illustrations`\n- `title`: `string` — required\n- `description`: `string` — optional\n\n## Settings Page Layout (`src/layouts/settings-layouts.tsx`)\n\n**Use this for all admin/settings pages.** Provides a standardized layout with scroll-aware\nsticky headers, centered content containers, and responsive behavior.\n\n```typescript\nimport SettingsLayouts from \"@/layouts/settings-layouts\";\n\nfunction MySettingsPage() {\n return (\n \n Save}\n >\n \n \n\n \n Settings content here\n \n \n );\n}\n```\n\nSub-components:\n- **`SettingsLayouts.Root`** — Wrapper with centered, scrollable container. Width options:\n `\"sm\"` (672px), `\"sm-md\"` (752px), `\"md\"` (872px, default), `\"lg\"` (992px), `\"full\"` (100%).\n- **`SettingsLayouts.Header`** — Sticky header with icon, title, description, optional\n `rightChildren` actions, optional `children` below (e.g., search/filter), optional `backButton`,\n and optional `separator`. Automatically shows a scroll shadow when scrolled.\n- **`SettingsLayouts.Body`** — Content container with consistent padding and vertical spacing.\n\n## Cards (`src/sections/cards/`)\n\n**When building a card that displays information about a specific entity (agent, document set,\nfile, connector, etc.), add it to `web/src/sections/cards/`.**\n\nEach card is a self-contained component focused on a single entity type. Cards typically include\nentity identification (name, avatar, icon), summary information, and quick actions.\n\n```typescript\nimport AgentCard from \"@/sections/cards/AgentCard\";\nimport DocumentSetCard from \"@/sections/cards/DocumentSetCard\";\nimport FileCard from \"@/sections/cards/FileCard\";\n```\n\nGuidelines:\n- One card per entity type — keep card-specific logic within the card component.\n- Cards should be reusable across different pages and contexts.\n- Use shared components from `@opal/components`, `@opal/layouts`, and `@/refresh-components`\n inside cards — do not duplicate layout or styling logic.\n\n## Button (`components/buttons/button/`)\n\n**Always use the Opal `Button`.** Do not use raw `\n\n// Icon-only button (omit children)\n\n \n )\n}\n\n// ❌ Bad\nfunction ContactForm() {\n return (\n
    \n \n